diff --git a/docs/language-reference.rst b/docs/language-reference.rst --- a/docs/language-reference.rst +++ b/docs/language-reference.rst @@ -34,12 +34,11 @@ Many things in Futhark are named. When we are defining something, we give it an unqualified name (`id`). When referencing something inside a module, we use a qualified name (`qualid`). The constructor names -of a sum type (:ref:`compounds`) are identifiers prefixed with ``#``, -with no space afterwards. The fields of a record are named with -`fieldid`. Note that a `fieldid` can be a decimal number. Futhark -has three distinct name spaces: terms, module types, and types. -Modules (including parametric modules) and values both share the term -namespace. +of a sum type are identifiers prefixed with ``#``, with no space +afterwards. The fields of a record are named with `fieldid`. Note +that a `fieldid` can be a decimal number. Futhark has three distinct +name spaces: terms, module types, and types. Modules (including +parametric modules) and values both share the term namespace. .. _primitives: @@ -94,8 +93,6 @@ hexdigit: `decdigit` | "a"..."f" | "A"..."F" bindigit: "0" | "1" -.. _compounds: - Compound Types and Values ~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -329,6 +326,9 @@ Any top-level function named ``main`` will always be considered an entry point, whether it is declared with ``entry`` or not. +The name of an entry point must not contain an apostrophe (``'``), +even though that is normally permitted in Futhark identifiers. + Value Declarations ~~~~~~~~~~~~~~~~~~ @@ -594,7 +594,7 @@ Return the element at the given position in the array. The index may be a comma-separated list of indexes instead of just a single index. If the number of indices given is less than the rank of the array, an -array is returned. +array is returned. The index may be of any unsigned integer type. The array ``a`` must be a variable name or a parenthesised expression. Furthermore, there *may not* be a space between ``a`` and the opening @@ -610,7 +610,8 @@ former inclusive and the latter exclusive, taking every ``s``-th 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). +(not inclusive). Slicing can be done only with expressions of type +``i32``. It is generally a bad idea for ``s`` to be non-constant. Slicing of multiple dimensions can be done by separating with commas, @@ -647,9 +648,9 @@ ``x..y...z`` ............ -Construct an integer array whose first element is ``x`` and which -proceeds stride of ``y-x`` until reaching ``z`` (inclusive). The -``..y`` part can be elided in which case a stride of 1 is used. A +Construct a signed integer array whose first element is ``x`` and +which proceeds stride of ``y-x`` until reaching ``z`` (inclusive). +The ``..y`` part can be elided in which case a stride of 1 is used. A run-time error occurs if ``z`` is lesser than ``x`` or ``y``, or if ``x`` and ``y`` are the same value. @@ -666,8 +667,8 @@ ``x..y..z`` ............... -Construct an integer array whose first elements is ``x``, and which -proceeds downwards with a stride of ``y`` until reaching ``z`` +Construct a signed integer array whose first elements is ``x``, and +which proceeds downwards with a stride of ``y`` until reaching ``z`` (exclusive). The ``..y`` part can be elided in which case a stride of -1 is used. A run-time error occurs if ``z`` is greater than ``x`` or ``y``, or if ``x`` and ``y`` are the same value. 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 @@ -62,6 +62,10 @@ The following options are accepted by executables generated by ``futhark c``. +-h, --help + + Print help text to standard output and exit. + -b, --binary-output Print the program result in the binary output format. The default 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 @@ -23,10 +23,14 @@ resulting program will otherwise behave exactly as one compiled with ``futhark c``. -``futhark cuda`` uses ``-lcuda -lnvrtc`` to link. If using +``futhark cuda`` uses ``-lcuda -lcudart -lnvrtc`` to link. If using ``--library``, you will need to do the same when linking the final binary. +The generated CUDA code can be called from multiple CPU threads, as it +brackets every API operation with ``cuCtxPushCurrent()`` and +``cuCtxPopCurrent()``. + OPTIONS ======= @@ -68,6 +72,10 @@ nomenclature ("group" instead of "thread block"). The following additional options are accepted. + +-h, --help + + Print help text to standard output and exit. --default-group-size=INT 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 @@ -73,6 +73,10 @@ The following additional options are accepted. +-h, --help + + Print help text to standard output and exit. + --build-option=OPT Add an additional build option to the string passed to @@ -154,6 +158,10 @@ --tuning=FILE Read size=value assignments from the given file. + +--list-devices + + List all OpenCL devices and platforms available on the system. SEE ALSO ======== diff --git a/docs/usage.rst b/docs/usage.rst --- a/docs/usage.rst +++ b/docs/usage.rst @@ -75,6 +75,10 @@ All generated executables support the following options. + ``-h/--help`` + + Print help text to standard output and exit. + ``-t FILE`` Print the time taken to execute the program to the indicated file, @@ -205,6 +209,10 @@ ``clBuildProgram()``. Refer to the OpenCL documentation for which options are supported. Be careful - some options can easily result in invalid results. + + ``--list-devices`` + + List all OpenCL devices and platforms available on the system. There is rarely a need to use both ``-p`` and ``-d``. For example, to run on the first available NVIDIA GPU, ``-p NVIDIA`` is sufficient, as diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -1,7 +1,7 @@ cabal-version: 2.4 name: futhark -version: 0.16.4 +version: 0.17.1 synopsis: An optimising compiler for a functional, array-oriented language. description: Futhark is a small programming language designed to be compiled to @@ -118,6 +118,7 @@ Futhark.CodeGen.ImpGen.Kernels.Transpose Futhark.CodeGen.ImpGen.OpenCL Futhark.CodeGen.ImpGen.Sequential + Futhark.CodeGen.ImpGen.Transpose Futhark.CodeGen.OpenCL.Heuristics Futhark.CodeGen.SetDefaultSpace Futhark.Compiler @@ -286,6 +287,7 @@ , process >=1.4.3.0 , process-extras >=0.7.2 , regex-tdfa >=1.2 + , sexp-grammar >= 2.2.1 , srcloc >=0.4 , template-haskell >=2.11.1 , temporary @@ -342,6 +344,7 @@ , megaparsec , mtl , parser-combinators + , sexp-grammar , tasty , tasty-hunit , tasty-quickcheck diff --git a/prelude/array.fut b/prelude/array.fut --- a/prelude/array.fut +++ b/prelude/array.fut @@ -24,13 +24,13 @@ let init [n] 't (x: [n]t) = x[0:n-1] -- | Take some number of elements from the head of the array. -let take [n] 't (i: i32) (x: [n]t): [i]t = x[0:i] +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. -let drop [n] 't (i: i32) (x: [n]t) = x[i:] +let drop [n] 't (i: i64) (x: [n]t) = x[i:] -- | Split an array at a given position. -let split [n] 't (i: i32) (xs: [n]t): ([i]t, []t) = +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. @@ -46,28 +46,28 @@ -- | Concatenation where the result has a predetermined size. If the -- provided size is wrong, the function will fail with a run-time -- error. -let concat_to [n] [m] 't (k: i32) (xs: [n]t) (ys: [m]t): *[k]t = xs ++ ys :> [k]t +let concat_to [n] [m] 't (k: i64) (xs: [n]t) (ys: [m]t): *[k]t = xs ++ ys :> [k]t -- | Rotate an array some number of elements to the left. A negative -- rotation amount is also supported. -- -- For example, if `b==rotate r a`, then `b[x+r] = a[x]`. -let rotate [n] 't (r: i32) (xs: [n]t): [n]t = intrinsics.rotate (r, xs) :> [n]t +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. -let iota (n: i32): *[n]i32 = - i32.iota n :> [n]i32 +let iota (n: i64): *[n]i64 = + 0..1.. [n]t +let replicate 't (n: i64) (x: t): *[n]t = + map (const x) (iota n) -- | Copy a value. The result will not alias anything. let copy 't (a: t): *t = @@ -79,7 +79,7 @@ -- | Like `flatten`@term, but where the final size is known. Fails at -- runtime if the provided size is wrong. -let flatten_to [n][m] 't (l: i32) (xs: [n][m]t): [l]t = +let flatten_to [n][m] 't (l: i64) (xs: [n][m]t): [l]t = flatten xs :> [l]t -- | Combines the outer three dimensions of an array. @@ -91,15 +91,15 @@ flatten (flatten_3d xs) -- | Splits the outer dimension of an array in two. -let unflatten [p] 't (n: i32) (m: i32) (xs: [p]t): [n][m]t = +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. -let unflatten_3d [p] 't (n: i32) (m: i32) (l: i32) (xs: [p]t): [n][m][l]t = +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. -let unflatten_4d [p] 't (n: i32) (m: i32) (l: i32) (k: i32) (xs: [p]t): [n][m][l][k]t = +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) let transpose [n] [m] 't (a: [n][m]t): [m][n]t = @@ -122,13 +122,13 @@ foldl (flip f) acc (reverse bs) -- | Create a value for each point in a one-dimensional index space. -let tabulate 'a (n: i32) (f: i32 -> a): *[n]a = +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. -let tabulate_2d 'a (n: i32) (m: i32) (f: i32 -> i32 -> a): *[n][m]a = +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. -let tabulate_3d 'a (n: i32) (m: i32) (o: i32) (f: i32 -> i32 -> i32 -> a): *[n][m][o]a = +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 @@ -2,8 +2,6 @@ import "soacs" -local let const 'a 'b (x: a) (_: b): a = x - -- | Describes types of values that can be created from the primitive -- numeric types (and bool). module type from_prim = { @@ -118,22 +116,11 @@ val ctz: t -> i32 } --- | An extension of `size`@mtype that further includes facilities for --- constructing arrays where the size is provided as a value of the --- given integral type. -module type size = { - include integral - - val iota: t -> *[]t - val replicate 'v: t -> v -> *[]v -} - -- | Numbers that model real numbers to some degree. module type real = { include numeric - val from_fraction: i32 -> i32 -> t - val to_i32: t -> i32 + val from_fraction: i64 -> i64 -> t val to_i64: t -> i64 val to_f64: t -> f64 @@ -239,7 +226,7 @@ let bool (x: bool) = x } -module i8: (size with t = i8) = { +module i8: (integral with t = i8) = { type t = i8 let (x: i8) + (y: i8) = intrinsics.add8 (x, y) @@ -305,16 +292,13 @@ let clz = intrinsics.clz8 let ctz = intrinsics.ctz8 - let iota (n: i8) = 0i8..1i8.. a -> a) (ne : a) (is : [n]i32) (as : [n]a) : *[m]a = +let reduce_by_index 'a [m] [n] (dest : *[m]a) (f : a -> a -> a) (ne : a) (is : [n]i64) (as : [n]a) : *[m]a = intrinsics.hist (1, dest, f, ne, is, as) :> *[m]a -- | Inclusive prefix scan. Has the same caveats with respect to @@ -163,7 +163,7 @@ -- | `reduce_stream op f as` splits `as` into chunks, applies `f` to each -- of these in parallel, and uses `op` (which must be associative) to --- combine the per-chunk results into a final result. The `i32` +-- combine the per-chunk results into a final result. The `i64` -- passed to `f` is the size of the chunk. This SOAC is useful when -- `f` can be given a particularly work-efficient sequential -- implementation. Operationally, we can imagine that `as` is divided @@ -176,7 +176,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(log(n))* -let reduce_stream [n] 'a 'b (op: b -> b -> b) (f: (k: i32) -> [k]a -> b) (as: [n]a): b = +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) -- | As `reduce_stream`@term, but the chunks do not necessarily @@ -186,7 +186,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(log(n))* -let reduce_stream_per [n] 'a 'b (op: b -> b -> b) (f: (k: i32) -> [k]a -> b) (as: [n]a): b = +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) -- | Similar to `reduce_stream`@term, except that each chunk must produce @@ -196,7 +196,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(1)* -let map_stream [n] 'a 'b (f: (k: i32) -> [k]a -> [k]b) (as: [n]a): *[n]b = +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 -- | Similar to `map_stream`@term, but the chunks do not necessarily @@ -206,7 +206,7 @@ -- **Work:** *O(n)* -- -- **Span:** *O(1)* -let map_stream_per [n] 'a 'b (f: (k: i32) -> [k]a -> [k]b) (as: [n]a): *[n]b = +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 @@ -252,5 +252,5 @@ -- **Work:** *O(n)* -- -- **Span:** *O(1)* -let scatter 't [m] [n] (dest: *[m]t) (is: [n]i32) (vs: [n]t): *[m]t = +let scatter 't [m] [n] (dest: *[m]t) (is: [n]i64) (vs: [n]t): *[m]t = intrinsics.scatter (dest, is, vs) :> *[m]t diff --git a/rts/c/lock.h b/rts/c/lock.h --- a/rts/c/lock.h +++ b/rts/c/lock.h @@ -11,7 +11,7 @@ typedef HANDLE lock_t; -static lock_t create_lock(lock_t *lock) { +static void create_lock(lock_t *lock) { *lock = CreateMutex(NULL, // Default security attributes. FALSE, // Initially unlocked. NULL); // Unnamed. diff --git a/rts/c/opencl.h b/rts/c/opencl.h --- a/rts/c/opencl.h +++ b/rts/c/opencl.h @@ -342,6 +342,33 @@ } // Returns 0 on success. +static int list_devices(struct opencl_config *cfg) { + struct opencl_device_option *devices; + size_t num_devices; + + opencl_all_device_options(&devices, &num_devices); + + const char *cur_platform = ""; + for (size_t i = 0; i < num_devices; i++) { + struct opencl_device_option device = devices[i]; + if (strcmp(cur_platform, device.platform_name) != 0) { + printf("Platform: %s\n", device.platform_name); + cur_platform = device.platform_name; + } + printf("[%d]: %s\n", (int)i, device.device_name); + } + + // Free all the platform and device names. + for (size_t j = 0; j < num_devices; j++) { + free(devices[j].platform_name); + free(devices[j].device_name); + } + free(devices); + + return 0; +} + +// Returns 0 on success. static int select_device_interactively(struct opencl_config *cfg) { struct opencl_device_option *devices; size_t num_devices; @@ -529,11 +556,18 @@ // Futhark reserves 4 bytes for bookkeeping information. max_local_memory -= 4; - // NVIDIA reserves some more bytes for who-knows-what. The number - // of bytes here has been experimentally determined, but the - // overhead seems to vary a bit depending on what the kernel does. + // The OpenCL implementation may reserve some local memory bytes for + // various purposes. In principle, we should use + // clGetKernelWorkGroupInfo() to figure out for each kernel how much + // is actually available, but our current code generator design + // makes this infeasible. Instead, we have this nasty hack where we + // arbitrarily subtract some bytes, based on empirical measurements + // (but which might be arbitrarily wrong). Fortunately, we rarely + // try to really push the local memory usage. if (strstr(device_option.platform_name, "NVIDIA CUDA") != NULL) { max_local_memory -= 12; + } else if (strstr(device_option.platform_name, "AMD") != NULL) { + max_local_memory -= 16; } // Make sure this function is defined. @@ -843,6 +877,7 @@ } static int opencl_alloc(struct opencl_context *ctx, size_t min_size, const char *tag, cl_mem *mem_out) { + (void)tag; if (min_size < sizeof(int)) { min_size = sizeof(int); } diff --git a/rts/python/opencl.py b/rts/python/opencl.py --- a/rts/python/opencl.py +++ b/rts/python/opencl.py @@ -115,12 +115,14 @@ # See comment in rts/c/opencl.h. if self.platform.name.find('NVIDIA CUDA') >= 0: self.max_local_memory -= 12 + elif self.platform.name.find('AMD') >= 0: + self.max_local_memory -= 16 self.free_list = {} self.global_failure = self.pool.allocate(np.int32().itemsize) cl.enqueue_fill_buffer(self.queue, self.global_failure, np.int32(-1), 0, np.int32().itemsize) - self.global_failure_args = self.pool.allocate(np.int32().itemsize * + self.global_failure_args = self.pool.allocate(np.int64().itemsize * (self.global_failure_args_max+1)) self.failure_is_an_option = np.int32(0) @@ -223,7 +225,7 @@ cl.enqueue_fill_buffer(self.queue, self.global_failure, np.int32(-1), 0, np.int32().itemsize) # Read failure args. - failure_args = np.empty(self.global_failure_args_max+1, dtype=np.int32) + failure_args = np.empty(self.global_failure_args_max+1, dtype=np.int64) cl.enqueue_copy(self.queue, failure_args, self.global_failure_args, is_blocking=True) raise Exception(self.failure_msgs[failure[0]].format(*failure_args)) diff --git a/src/Futhark/Actions.hs b/src/Futhark/Actions.hs --- a/src/Futhark/Actions.hs +++ b/src/Futhark/Actions.hs @@ -1,75 +1,106 @@ {-# LANGUAGE FlexibleContexts #-} + -- | All (almost) compiler pipelines end with an 'Action', which does -- something with the result of the pipeline. module Futhark.Actions - ( printAction - , impCodeGenAction - , kernelImpCodeGenAction - , metricsAction - , compileCAction - , compileOpenCLAction - , compileCUDAAction + ( printAction, + impCodeGenAction, + kernelImpCodeGenAction, + metricsAction, + compileCAction, + compileOpenCLAction, + compileCUDAAction, + sexpAction, ) where import Control.Monad import Control.Monad.IO.Class -import System.Exit -import System.FilePath -import qualified System.Info - -import Futhark.Compiler.CLI +import qualified Data.ByteString.Lazy.Char8 as ByteString import Futhark.Analysis.Alias +import Futhark.Analysis.Metrics +import qualified Futhark.CodeGen.Backends.CCUDA as CCUDA +import qualified Futhark.CodeGen.Backends.COpenCL as COpenCL +import qualified Futhark.CodeGen.Backends.SequentialC as SequentialC +import qualified Futhark.CodeGen.ImpGen.Kernels as ImpGenKernels +import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGenSequential +import Futhark.Compiler.CLI import Futhark.IR -import Futhark.IR.Prop.Aliases import Futhark.IR.KernelsMem (KernelsMem) +import Futhark.IR.Prop.Aliases import Futhark.IR.SeqMem (SeqMem) -import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGenSequential -import qualified Futhark.CodeGen.ImpGen.Kernels as ImpGenKernels -import qualified Futhark.CodeGen.Backends.SequentialC as SequentialC -import qualified Futhark.CodeGen.Backends.CCUDA as CCUDA -import qualified Futhark.CodeGen.Backends.COpenCL as COpenCL -import Futhark.Analysis.Metrics import Futhark.Util (runProgramWithExitCode) +import Language.SexpGrammar as Sexp +import System.Exit +import System.FilePath +import qualified System.Info -- | Print the result to stdout, with alias annotations. printAction :: (ASTLore lore, CanBeAliased (Op lore)) => Action lore printAction = - Action { actionName = "Prettyprint" - , actionDescription = "Prettyprint the resulting internal representation on standard output." - , actionProcedure = liftIO . putStrLn . pretty . aliasAnalysis - } + Action + { actionName = "Prettyprint", + actionDescription = "Prettyprint the resulting internal representation on standard output.", + actionProcedure = liftIO . putStrLn . pretty . aliasAnalysis + } -- | Print metrics about AST node counts to stdout. metricsAction :: OpMetrics (Op lore) => Action lore metricsAction = - Action { actionName = "Compute metrics" - , actionDescription = "Print metrics on the final AST." - , actionProcedure = liftIO . putStr . show . progMetrics - } + Action + { actionName = "Compute metrics", + actionDescription = "Print metrics on the final AST.", + actionProcedure = liftIO . putStr . show . progMetrics + } -- | Convert the program to sequential ImpCode and print it to stdout. impCodeGenAction :: Action SeqMem impCodeGenAction = - Action { actionName = "Compile imperative" - , actionDescription = "Translate program into imperative IL and write it on standard output." - , actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenSequential.compileProg - } + Action + { actionName = "Compile imperative", + actionDescription = "Translate program into imperative IL and write it on standard output.", + actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenSequential.compileProg + } -- | Convert the program to GPU ImpCode and print it to stdout. kernelImpCodeGenAction :: Action KernelsMem kernelImpCodeGenAction = - Action { actionName = "Compile imperative kernels" - , actionDescription = "Translate program into imperative IL with kernels and write it on standard output." - , actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenKernels.compileProgOpenCL - } + Action + { actionName = "Compile imperative kernels", + actionDescription = "Translate program into imperative IL with kernels and write it on standard output.", + actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenKernels.compileProgOpenCL + } +-- | Print metrics about AST node counts to stdout. +sexpAction :: ASTLore lore => Action lore +sexpAction = + Action + { actionName = "Print sexps", + actionDescription = "Print sexps on the final IR.", + actionProcedure = liftIO . helper + } + where + helper :: ASTLore lore => Prog lore -> IO () + helper prog = + case encodePretty prog of + Right prog' -> do + ByteString.putStrLn prog' + let prog'' = decode prog' + unless (prog'' == Right prog) $ + error $ + "S-exp not isomorph!\n" + ++ either show pretty prog'' + Left s -> + error $ "Couldn't encode program: " ++ s + -- | The @futhark c@ action. compileCAction :: FutharkConfig -> CompilerMode -> FilePath -> Action SeqMem compileCAction fcfg mode outpath = - Action { actionName = "Compile to OpenCL" - , actionDescription = "Compile to OpenCL" - , actionProcedure = helper } + Action + { actionName = "Compile to OpenCL", + actionDescription = "Compile to OpenCL", + actionProcedure = helper + } where helper prog = do cprog <- handleWarnings fcfg $ SequentialC.compileProg prog @@ -83,23 +114,32 @@ liftIO $ writeFile cpath impl ToExecutable -> do liftIO $ writeFile cpath $ SequentialC.asExecutable cprog - ret <- liftIO $ runProgramWithExitCode "gcc" - [cpath, "-O3", "-std=c99", "-lm", "-o", outpath] mempty + ret <- + liftIO $ + runProgramWithExitCode + "gcc" + [cpath, "-O3", "-std=c99", "-lm", "-o", outpath] + mempty case ret of Left err -> externalErrorS $ "Failed to run gcc: " ++ show err Right (ExitFailure code, _, gccerr) -> - externalErrorS $ "gcc failed with code " ++ - show code ++ ":\n" ++ gccerr + externalErrorS $ + "gcc failed with code " + ++ show code + ++ ":\n" + ++ gccerr Right (ExitSuccess, _, _) -> return () -- | The @futhark opencl@ action. compileOpenCLAction :: FutharkConfig -> CompilerMode -> FilePath -> Action KernelsMem compileOpenCLAction fcfg mode outpath = - Action { actionName = "Compile to OpenCL" - , actionDescription = "Compile to OpenCL" - , actionProcedure = helper } + Action + { actionName = "Compile to OpenCL", + actionDescription = "Compile to OpenCL", + actionProcedure = helper + } where helper prog = do cprog <- handleWarnings fcfg $ COpenCL.compileProg prog @@ -107,11 +147,11 @@ hpath = outpath `addExtension` "h" extra_options | System.Info.os == "darwin" = - ["-framework", "OpenCL"] + ["-framework", "OpenCL"] | System.Info.os == "mingw32" = - ["-lOpenCL64"] + ["-lOpenCL64"] | otherwise = - ["-lOpenCL"] + ["-lOpenCL"] case mode of ToLibrary -> do @@ -120,32 +160,42 @@ liftIO $ writeFile cpath impl ToExecutable -> do liftIO $ writeFile cpath $ COpenCL.asExecutable cprog - ret <- liftIO $ runProgramWithExitCode "gcc" - ([cpath, "-O", "-std=c99", "-lm", "-o", outpath] ++ extra_options) mempty + ret <- + liftIO $ + runProgramWithExitCode + "gcc" + ([cpath, "-O", "-std=c99", "-lm", "-o", outpath] ++ extra_options) + mempty case ret of Left err -> externalErrorS $ "Failed to run gcc: " ++ show err Right (ExitFailure code, _, gccerr) -> - externalErrorS $ "gcc failed with code " ++ - show code ++ ":\n" ++ gccerr + externalErrorS $ + "gcc failed with code " + ++ show code + ++ ":\n" + ++ gccerr Right (ExitSuccess, _, _) -> return () -- | The @futhark cuda@ action. compileCUDAAction :: FutharkConfig -> CompilerMode -> FilePath -> Action KernelsMem compileCUDAAction fcfg mode outpath = - Action { actionName = "Compile to CUDA" - , actionDescription = "Compile to CUDA" - , actionProcedure = helper } + Action + { actionName = "Compile to CUDA", + actionDescription = "Compile to CUDA", + actionProcedure = helper + } where helper prog = do cprog <- handleWarnings fcfg $ CCUDA.compileProg prog let cpath = outpath `addExtension` "c" hpath = outpath `addExtension` "h" - extra_options = [ "-lcuda" - , "-lcudart" - , "-lnvrtc" - ] + extra_options = + [ "-lcuda", + "-lcudart", + "-lnvrtc" + ] case mode of ToLibrary -> do let (header, impl) = CCUDA.asLibrary cprog @@ -153,14 +203,18 @@ liftIO $ writeFile cpath impl ToExecutable -> do liftIO $ writeFile cpath $ CCUDA.asExecutable cprog - let args = [cpath, "-O", "-std=c99", "-lm", "-o", outpath] - ++ extra_options + let args = + [cpath, "-O", "-std=c99", "-lm", "-o", outpath] + ++ extra_options ret <- liftIO $ runProgramWithExitCode "gcc" args mempty case ret of Left err -> externalErrorS $ "Failed to run gcc: " ++ show err Right (ExitFailure code, _, gccerr) -> - externalErrorS $ "gcc failed with code " ++ - show code ++ ":\n" ++ gccerr + externalErrorS $ + "gcc failed with code " + ++ show code + ++ ":\n" + ++ gccerr Right (ExitSuccess, _, _) -> return () 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 @@ -1,4 +1,5 @@ {-# LANGUAGE FlexibleContexts #-} + -- | Alias analysis of a full Futhark program. Takes as input a -- program with an arbitrary lore and produces one with aliases. This -- module does not implement the aliasing logic itself, and derives @@ -8,102 +9,126 @@ -- here will include transitive aliases (note that this is not what -- the building blocks do). module Futhark.Analysis.Alias - ( aliasAnalysis - -- * Ad-hoc utilities - , AliasTable - , analyseFun - , analyseStms - , analyseExp - , analyseBody - , analyseLambda - ) - where + ( aliasAnalysis, + -- * Ad-hoc utilities + AliasTable, + analyseFun, + analyseStms, + analyseExp, + analyseBody, + analyseLambda, + ) +where + import Data.List (foldl') import qualified Data.Map as M - -import Futhark.IR.Syntax import Futhark.IR.Aliases -- | Perform alias analysis on a Futhark program. -aliasAnalysis :: (ASTLore lore, CanBeAliased (Op lore)) => - Prog lore -> Prog (Aliases lore) +aliasAnalysis :: + (ASTLore lore, CanBeAliased (Op lore)) => + Prog lore -> + Prog (Aliases lore) aliasAnalysis (Prog consts funs) = Prog (fst (analyseStms mempty consts)) (map analyseFun funs) -analyseFun :: (ASTLore lore, CanBeAliased (Op lore)) => - FunDef lore -> FunDef (Aliases lore) +analyseFun :: + (ASTLore lore, CanBeAliased (Op lore)) => + FunDef lore -> + FunDef (Aliases lore) analyseFun (FunDef entry attrs fname restype params body) = FunDef entry attrs fname restype params body' - where body' = analyseBody mempty body + 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)) => - AliasTable -> Body lore -> Body (Aliases lore) +analyseBody :: + ( ASTLore lore, + CanBeAliased (Op lore) + ) => + AliasTable -> + Body lore -> + Body (Aliases lore) analyseBody atable (Body lore stms result) = let (stms', _atable') = analyseStms atable stms - in mkAliasedBody lore stms' result + in mkAliasedBody lore stms' result -analyseStms :: (ASTLore lore, CanBeAliased (Op lore)) => - AliasTable -> Stms lore -> (Stms (Aliases lore), AliasesAndConsumed) +analyseStms :: + (ASTLore lore, CanBeAliased (Op lore)) => + AliasTable -> + Stms lore -> + (Stms (Aliases lore), AliasesAndConsumed) analyseStms orig_aliases = foldl' f (mempty, (orig_aliases, mempty)) . stmsToList - where f (stms, aliases) stm = - let stm' = analyseStm (fst aliases) stm - atable' = trackAliases aliases stm' - in (stms<>oneStm stm', atable') + where + f (stms, aliases) stm = + let stm' = analyseStm (fst aliases) stm + atable' = trackAliases aliases stm' + in (stms <> oneStm stm', atable') -analyseStm :: (ASTLore lore, CanBeAliased (Op lore)) => - AliasTable -> Stm lore -> Stm (Aliases lore) +analyseStm :: + (ASTLore lore, CanBeAliased (Op lore)) => + AliasTable -> + Stm lore -> + Stm (Aliases lore) analyseStm aliases (Let pat (StmAux cs attrs dec) e) = let e' = analyseExp aliases e pat' = addAliasesToPattern pat e' lore' = (AliasDec $ consumedInExp e', dec) - in Let pat' (StmAux cs attrs lore') e' - -analyseExp :: (ASTLore lore, CanBeAliased (Op lore)) => - AliasTable -> Exp lore -> Exp (Aliases lore) + in Let pat' (StmAux cs attrs lore') e' +analyseExp :: + (ASTLore lore, CanBeAliased (Op lore)) => + AliasTable -> + Exp lore -> + Exp (Aliases lore) -- Would be better to put this in a BranchType annotation, but that -- requires a lot of other work. analyseExp aliases (If cond tb fb dec) = let Body ((tb_als, tb_cons), tb_dec) tb_stms tb_res = analyseBody aliases tb Body ((fb_als, fb_cons), fb_dec) fb_stms fb_res = analyseBody aliases fb cons = tb_cons <> fb_cons - isConsumed v = any (`nameIn` unAliases cons) $ - v : namesToList (M.findWithDefault mempty v aliases) - notConsumed = AliasDec . namesFromList . - filter (not . isConsumed) . - namesToList . unAliases + isConsumed v = + any (`nameIn` unAliases cons) $ + v : namesToList (M.findWithDefault mempty v aliases) + notConsumed = + AliasDec . namesFromList + . filter (not . isConsumed) + . namesToList + . unAliases tb_als' = map notConsumed tb_als fb_als' = map notConsumed fb_als tb' = Body ((tb_als', tb_cons), tb_dec) tb_stms tb_res fb' = Body ((fb_als', fb_cons), fb_dec) fb_stms fb_res - in If cond tb' fb' dec - + in If cond tb' fb' dec analyseExp aliases e = mapExp analyse e - where analyse = - Mapper { mapOnSubExp = return - , mapOnVName = return - , mapOnBody = const $ return . analyseBody aliases - , mapOnRetType = return - , mapOnBranchType = return - , mapOnFParam = return - , mapOnLParam = return - , mapOnOp = return . addOpAliases - } + where + analyse = + Mapper + { mapOnSubExp = return, + mapOnVName = return, + mapOnBody = const $ return . analyseBody aliases, + mapOnRetType = return, + mapOnBranchType = return, + mapOnFParam = return, + mapOnLParam = return, + mapOnOp = return . addOpAliases + } -analyseLambda :: (ASTLore lore, CanBeAliased (Op lore)) => - Lambda lore -> Lambda (Aliases lore) +analyseLambda :: + (ASTLore lore, CanBeAliased (Op lore)) => + 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 - in lam { lambdaBody = body - , lambdaParams = lambdaParams lam - } + in lam + { lambdaBody = body, + lambdaParams = lambdaParams lam + } diff --git a/src/Futhark/Analysis/CallGraph.hs b/src/Futhark/Analysis/CallGraph.hs --- a/src/Futhark/Analysis/CallGraph.hs +++ b/src/Futhark/Analysis/CallGraph.hs @@ -1,30 +1,31 @@ {-# LANGUAGE OverloadedStrings #-} + -- | This module exports functionality for generating a call graph of -- an Futhark program. module Futhark.Analysis.CallGraph - ( CallGraph - , buildCallGraph - , isFunInCallGraph - , calls - , calledByConsts - , allCalledBy - , findNoninlined + ( CallGraph, + buildCallGraph, + isFunInCallGraph, + calls, + calledByConsts, + allCalledBy, + findNoninlined, ) - where +where import Control.Monad.Writer.Strict +import Data.List (foldl') import qualified Data.Map.Strict as M -import qualified Data.Set as S import Data.Maybe (fromMaybe) -import Data.List (foldl') - +import qualified Data.Set as S import Futhark.IR.SOACS type FunctionTable = M.Map Name (FunDef SOACS) buildFunctionTable :: Prog SOACS -> FunctionTable buildFunctionTable = foldl expand M.empty . progFuns - where expand ftab f = M.insert (funDefName f) f ftab + where + expand ftab f = M.insert (funDefName f) f ftab type FunGraph = M.Map Name (S.Set Name) @@ -33,9 +34,10 @@ -- transitively!) by the function. -- -- We keep track separately of the functions called by constants. -data CallGraph = CallGraph { calledByFuns :: M.Map Name (S.Set Name) - , calledInConsts :: S.Set Name - } +data CallGraph = CallGraph + { calledByFuns :: M.Map Name (S.Set Name), + calledInConsts :: S.Set Name + } -- | Is the given function known to the call graph? isFunInCallGraph :: Name -> CallGraph -> Bool @@ -58,23 +60,24 @@ buildCallGraph :: Prog SOACS -> CallGraph buildCallGraph prog = CallGraph fg $ buildFGStms $ progConsts prog - where fg = foldl' (buildFGfun ftable) M.empty entry_points + where + fg = foldl' (buildFGfun ftable) M.empty entry_points - entry_points = map funDefName $ progFuns prog - ftable = buildFunctionTable prog + entry_points = map funDefName $ progFuns prog + ftable = buildFunctionTable prog -- | @buildCallGraph ftable fg fname@ updates @fg@ with the -- contributions of function @fname@. buildFGfun :: FunctionTable -> FunGraph -> Name -> FunGraph -buildFGfun ftable fg fname = +buildFGfun ftable fg fname = -- Check if function is a non-builtin that we have not already -- processed. case M.lookup fname ftable of Just f | Nothing <- M.lookup fname fg -> do - let callees = buildFGBody $ funDefBody f - fg' = M.insert fname callees fg - -- recursively build the callees - foldl' (buildFGfun ftable) fg' callees + let callees = buildFGBody $ funDefBody f + fg' = M.insert fname callees fg + -- recursively build the callees + foldl' (buildFGfun ftable) fg' callees _ -> fg buildFGStms :: Stms SOACS -> S.Set Name @@ -86,41 +89,51 @@ buildFGexp :: Exp -> S.Set Name buildFGexp (Apply fname _ _ _) = S.singleton fname buildFGexp (Op op) = execWriter $ mapSOACM folder op - where folder = identitySOACMapper { - mapOnSOACLambda = \lam -> do tell $ buildFGBody $ lambdaBody lam - return lam - } + where + folder = + identitySOACMapper + { mapOnSOACLambda = \lam -> do + tell $ buildFGBody $ lambdaBody lam + return lam + } buildFGexp e = execWriter $ mapExpM folder e - where folder = identityMapper { - mapOnBody = \_ body -> do tell $ buildFGBody body - return body - } + where + folder = + identityMapper + { mapOnBody = \_ body -> do + tell $ buildFGBody body + return body + } -- | The set of all functions that are called noinline somewhere, or -- have a noinline attribute on their definition. findNoninlined :: Prog SOACS -> S.Set Name findNoninlined prog = - foldMap noinlineDef (progFuns prog) <> - foldMap onStm (foldMap (bodyStms . funDefBody) (progFuns prog) <> progConsts prog) - where onStm :: Stm -> S.Set Name - onStm (Let _ aux (Apply fname _ _ _)) - | "noinline" `inAttrs` stmAuxAttrs aux = - S.singleton fname - onStm (Let _ _ e) = execWriter $ mapExpM folder e - where folder = - identityMapper - { mapOnBody = \_ body -> do tell $ foldMap onStm $ bodyStms body - return body - , mapOnOp = mapSOACM - identitySOACMapper - { mapOnSOACLambda = \lam -> do - tell $ foldMap onStm $ bodyStms $ lambdaBody lam - return lam - } - } + foldMap noinlineDef (progFuns prog) + <> foldMap onStm (foldMap (bodyStms . funDefBody) (progFuns prog) <> progConsts prog) + where + onStm :: Stm -> S.Set Name + onStm (Let _ aux (Apply fname _ _ _)) + | "noinline" `inAttrs` stmAuxAttrs aux = + S.singleton fname + onStm (Let _ _ e) = execWriter $ mapExpM folder e + where + folder = + identityMapper + { mapOnBody = \_ body -> do + tell $ foldMap onStm $ bodyStms body + return body, + mapOnOp = + mapSOACM + identitySOACMapper + { mapOnSOACLambda = \lam -> do + tell $ foldMap onStm $ bodyStms $ lambdaBody lam + return lam + } + } - noinlineDef fd - | "noinline" `inAttrs` funDefAttrs fd = - S.singleton $ funDefName fd - | otherwise = - mempty + noinlineDef fd + | "noinline" `inAttrs` funDefAttrs fd = + S.singleton $ funDefName fd + | otherwise = + mempty diff --git a/src/Futhark/Analysis/DataDependencies.hs b/src/Futhark/Analysis/DataDependencies.hs --- a/src/Futhark/Analysis/DataDependencies.hs +++ b/src/Futhark/Analysis/DataDependencies.hs @@ -1,14 +1,14 @@ {-# LANGUAGE FlexibleContexts #-} + -- | Facilities for inspecting the data dependencies of a program. module Futhark.Analysis.DataDependencies - ( Dependencies - , dataDependencies - , findNecessaryForReturned + ( Dependencies, + dataDependencies, + findNecessaryForReturned, ) - where +where import qualified Data.Map.Strict as M - import Futhark.IR -- | A mapping from a variable name @v@, to those variables on which @@ -21,31 +21,39 @@ dataDependencies :: ASTLore lore => Body lore -> Dependencies dataDependencies = dataDependencies' M.empty -dataDependencies' :: ASTLore lore => - Dependencies -> Body lore -> Dependencies +dataDependencies' :: + ASTLore lore => + Dependencies -> + Body lore -> + Dependencies dataDependencies' startdeps = foldl grow startdeps . bodyStms - where grow deps (Let pat _ (If c tb fb _)) = - let tdeps = dataDependencies' deps tb - fdeps = dataDependencies' deps fb - cdeps = depsOf deps c - comb (pe, tres, fres) = - (patElemName pe, - mconcat $ [freeIn pe, cdeps, depsOf tdeps tres, depsOf fdeps fres] ++ - map (depsOfVar deps) (namesToList $ freeIn pe)) - branchdeps = - M.fromList $ map comb $ zip3 (patternElements pat) - (bodyResult tb) - (bodyResult fb) - in M.unions [branchdeps, deps, tdeps, fdeps] - - grow deps (Let pat _ e) = - let free = freeIn pat <> freeIn e - freeDeps = mconcat $ map (depsOfVar deps) $ namesToList free - in M.fromList [ (name, freeDeps) | name <- patternNames pat ] `M.union` deps + where + grow deps (Let pat _ (If c tb fb _)) = + let tdeps = dataDependencies' deps tb + fdeps = dataDependencies' deps fb + cdeps = depsOf deps c + comb (pe, tres, fres) = + ( patElemName pe, + mconcat $ + [freeIn pe, cdeps, depsOf tdeps tres, depsOf fdeps fres] + ++ map (depsOfVar deps) (namesToList $ freeIn pe) + ) + branchdeps = + M.fromList $ + map comb $ + zip3 + (patternElements pat) + (bodyResult tb) + (bodyResult fb) + in M.unions [branchdeps, deps, tdeps, fdeps] + grow deps (Let pat _ e) = + let free = freeIn pat <> freeIn e + freeDeps = mconcat $ map (depsOfVar deps) $ namesToList free + in M.fromList [(name, freeDeps) | name <- patternNames pat] `M.union` deps depsOf :: Dependencies -> SubExp -> Names depsOf _ (Constant _) = mempty -depsOf deps (Var v) = depsOfVar deps v +depsOf deps (Var v) = depsOfVar deps v depsOfVar :: Dependencies -> VName -> Names depsOfVar deps name = oneName name <> M.findWithDefault mempty name deps @@ -56,21 +64,25 @@ -- of that parameter is live after the loop. @deps@ is the data -- dependencies of the loop body. This is computed by straightforward -- fixpoint iteration. -findNecessaryForReturned :: (Param dec -> Bool) -> [(Param dec, SubExp)] - -> M.Map VName Names - -> Names +findNecessaryForReturned :: + (Param dec -> Bool) -> + [(Param dec, SubExp)] -> + M.Map VName Names -> + Names findNecessaryForReturned usedAfterLoop merge_and_res allDependencies = - iterateNecessary mempty <> - namesFromList (map paramName $ filter usedAfterLoop $ map fst merge_and_res) - where iterateNecessary prev_necessary - | necessary == prev_necessary = necessary - | otherwise = iterateNecessary necessary - where necessary = mconcat $ map dependencies returnedResultSubExps - usedAfterLoopOrNecessary param = - usedAfterLoop param || paramName param `nameIn` prev_necessary - returnedResultSubExps = - map snd $ filter (usedAfterLoopOrNecessary . fst) merge_and_res - dependencies (Constant _) = - mempty - dependencies (Var v) = - M.findWithDefault (oneName v) v allDependencies + iterateNecessary mempty + <> namesFromList (map paramName $ filter usedAfterLoop $ map fst merge_and_res) + where + iterateNecessary prev_necessary + | necessary == prev_necessary = necessary + | otherwise = iterateNecessary necessary + where + necessary = mconcat $ map dependencies returnedResultSubExps + usedAfterLoopOrNecessary param = + usedAfterLoop param || paramName param `nameIn` prev_necessary + returnedResultSubExps = + map snd $ filter (usedAfterLoopOrNecessary . fst) merge_and_res + dependencies (Constant _) = + mempty + dependencies (Var v) = + M.findWithDefault (oneName v) v allDependencies diff --git a/src/Futhark/Analysis/HORep/MapNest.hs b/src/Futhark/Analysis/HORep/MapNest.hs --- a/src/Futhark/Analysis/HORep/MapNest.hs +++ b/src/Futhark/Analysis/HORep/MapNest.hs @@ -1,50 +1,50 @@ {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} + module Futhark.Analysis.HORep.MapNest - ( Nesting (..) - , MapNest (..) - , typeOf - , params - , inputs - , setInputs - , fromSOAC - , toSOAC + ( Nesting (..), + MapNest (..), + typeOf, + params, + inputs, + setInputs, + fromSOAC, + toSOAC, ) where import Data.List (find) -import Data.Maybe import qualified Data.Map.Strict as M - -import qualified Futhark.Analysis.HORep.SOAC as SOAC +import Data.Maybe import Futhark.Analysis.HORep.SOAC (SOAC) +import qualified Futhark.Analysis.HORep.SOAC as SOAC +import Futhark.Construct +import Futhark.IR hiding (typeOf) import qualified Futhark.IR.SOACS.SOAC as Futhark import Futhark.Transform.Substitute -import Futhark.IR hiding (typeOf) -import Futhark.MonadFreshNames -import Futhark.Construct -data Nesting lore = Nesting { - nestingParamNames :: [VName] - , nestingResult :: [VName] - , nestingReturnType :: [Type] - , nestingWidth :: SubExp - } deriving (Eq, Ord, Show) +data Nesting lore = Nesting + { nestingParamNames :: [VName], + nestingResult :: [VName], + nestingReturnType :: [Type], + nestingWidth :: SubExp + } + deriving (Eq, Ord, Show) data MapNest lore = MapNest SubExp (Lambda lore) [Nesting lore] [SOAC.Input] - deriving (Show) + deriving (Show) typeOf :: MapNest lore -> [Type] typeOf (MapNest w lam [] _) = map (`arrayOfRow` w) $ lambdaReturnType lam -typeOf (MapNest w _ (nest:_) _) = +typeOf (MapNest w _ (nest : _) _) = map (`arrayOfRow` w) $ nestingReturnType nest params :: MapNest lore -> [VName] -params (MapNest _ lam [] _) = +params (MapNest _ lam [] _) = map paramName $ lambdaParams lam -params (MapNest _ _ (nest:_) _) = +params (MapNest _ _ (nest : _) _) = nestingParamNames nest inputs :: MapNest lore -> [SOAC.Input] @@ -52,32 +52,41 @@ setInputs :: [SOAC.Input] -> MapNest lore -> MapNest lore setInputs [] (MapNest w body ns _) = MapNest w body ns [] -setInputs (inp:inps) (MapNest _ body ns _) = MapNest w body ns' (inp:inps) - where w = arraySize 0 $ SOAC.inputType inp - ws = drop 1 $ arrayDims $ SOAC.inputType inp - ns' = zipWith setDepth ns ws - setDepth n nw = n { nestingWidth = nw } +setInputs (inp : inps) (MapNest _ body ns _) = MapNest w body ns' (inp : inps) + where + w = arraySize 0 $ SOAC.inputType inp + ws = drop 1 $ arrayDims $ SOAC.inputType inp + ns' = zipWith setDepth ns ws + setDepth n nw = n {nestingWidth = nw} -fromSOAC :: (Bindable lore, MonadFreshNames m, - LocalScope lore m, - Op lore ~ Futhark.SOAC lore) => - SOAC lore -> m (Maybe (MapNest lore)) +fromSOAC :: + ( Bindable lore, + MonadFreshNames m, + LocalScope lore m, + Op lore ~ Futhark.SOAC lore + ) => + SOAC lore -> + m (Maybe (MapNest lore)) fromSOAC = fromSOAC' mempty -fromSOAC' :: (Bindable lore, MonadFreshNames m, - LocalScope lore m, - Op lore ~ Futhark.SOAC lore) => - [Ident] - -> SOAC lore - -> m (Maybe (MapNest lore)) - +fromSOAC' :: + ( Bindable lore, + MonadFreshNames m, + LocalScope lore m, + Op lore ~ Futhark.SOAC lore + ) => + [Ident] -> + SOAC lore -> + m (Maybe (MapNest lore)) fromSOAC' bound (SOAC.Screma w (SOAC.ScremaForm [] [] lam) inps) = do - maybenest <- case (stmsToList $ bodyStms $ lambdaBody lam, - bodyResult $ lambdaBody lam) of - ([Let pat _ e], res) | res == map Var (patternNames pat) -> - localScope (scopeOfLParams $ lambdaParams lam) $ - SOAC.fromExp e >>= - either (return . Left) (fmap (Right . fmap (pat,)) . fromSOAC' bound') + maybenest <- case ( stmsToList $ bodyStms $ lambdaBody lam, + bodyResult $ lambdaBody lam + ) of + ([Let pat _ e], res) + | res == map Var (patternNames pat) -> + localScope (scopeOfLParams $ lambdaParams lam) $ + SOAC.fromExp e + >>= either (return . Left) (fmap (Right . fmap (pat,)) . fromSOAC' bound') _ -> return $ Right Nothing @@ -85,72 +94,93 @@ -- Do we have a nested MapNest? Right (Just (pat, mn@(MapNest inner_w body' ns' inps'))) -> do (ps, inps'') <- - unzip <$> - fixInputs w (zip (map paramName $ lambdaParams lam) inps) - (zip (params mn) inps') - let n' = Nesting { nestingParamNames = ps - , nestingResult = patternNames pat - , nestingReturnType = typeOf mn - , nestingWidth = inner_w - } - return $ Just $ MapNest w body' (n':ns') inps'' + unzip + <$> fixInputs + w + (zip (map paramName $ lambdaParams lam) inps) + (zip (params mn) inps') + let n' = + Nesting + { nestingParamNames = ps, + nestingResult = patternNames pat, + nestingReturnType = typeOf mn, + nestingWidth = inner_w + } + return $ Just $ MapNest w body' (n' : ns') inps'' -- No nested MapNest it seems. _ -> do let isBound name - | Just param <- find ((name==) . identName) bound = + | Just param <- find ((name ==) . identName) bound = Just param | otherwise = Nothing boundUsedInBody = mapMaybe isBound $ namesToList $ freeIn lam - newParams <- mapM (newIdent' (++"_wasfree")) boundUsedInBody - let subst = M.fromList $ - zip (map identName boundUsedInBody) (map identName newParams) - inps' = inps ++ - map (SOAC.addTransform (SOAC.Replicate mempty $ Shape [w]) . SOAC.identInput) - boundUsedInBody + newParams <- mapM (newIdent' (++ "_wasfree")) boundUsedInBody + let subst = + M.fromList $ + zip (map identName boundUsedInBody) (map identName newParams) + inps' = + inps + ++ map + (SOAC.addTransform (SOAC.Replicate mempty $ Shape [w]) . SOAC.identInput) + boundUsedInBody lam' = - lam { lambdaBody = - substituteNames subst $ lambdaBody lam - , lambdaParams = - lambdaParams lam ++ [ Param name t - | Ident name t <- newParams ] - } + lam + { lambdaBody = + substituteNames subst $ lambdaBody lam, + lambdaParams = + lambdaParams lam + ++ [ Param name t + | Ident name t <- newParams + ] + } return $ Just $ MapNest w lam' [] inps' - where bound' = bound <> map paramIdent (lambdaParams lam) - + where + bound' = bound <> map paramIdent (lambdaParams lam) fromSOAC' _ _ = return Nothing -toSOAC :: (MonadFreshNames m, HasScope lore m, - Bindable lore, BinderOps lore, Op lore ~ Futhark.SOAC lore) => - MapNest lore -> m (SOAC lore) +toSOAC :: + ( MonadFreshNames m, + HasScope lore m, + Bindable lore, + BinderOps lore, + Op lore ~ Futhark.SOAC lore + ) => + MapNest lore -> + m (SOAC lore) toSOAC (MapNest w lam [] inps) = return $ SOAC.Screma w (Futhark.mapSOAC lam) inps -toSOAC (MapNest w lam (Nesting npnames nres nrettype nw:ns) inps) = do +toSOAC (MapNest w lam (Nesting npnames nres nrettype nw : ns) inps) = do let nparams = zipWith Param npnames $ map SOAC.inputRowType inps - body <- runBodyBinder $ localScope (scopeOfLParams nparams) $ do - letBindNames nres =<< SOAC.toExp =<< - toSOAC (MapNest nw lam ns $ map (SOAC.identInput . paramIdent) nparams) - return $ resultBody $ map Var nres - let outerlam = Lambda { lambdaParams = nparams - , lambdaBody = body - , lambdaReturnType = nrettype - } + body <- runBodyBinder $ + localScope (scopeOfLParams nparams) $ do + letBindNames nres =<< SOAC.toExp + =<< toSOAC (MapNest nw lam ns $ map (SOAC.identInput . paramIdent) nparams) + return $ resultBody $ map Var nres + let outerlam = + Lambda + { lambdaParams = nparams, + lambdaBody = body, + lambdaReturnType = nrettype + } return $ SOAC.Screma w (Futhark.mapSOAC outerlam) inps -fixInputs :: MonadFreshNames m => - SubExp -> [(VName, SOAC.Input)] -> [(VName, SOAC.Input)] - -> m [(VName, SOAC.Input)] +fixInputs :: + MonadFreshNames m => + SubExp -> + [(VName, SOAC.Input)] -> + [(VName, SOAC.Input)] -> + m [(VName, SOAC.Input)] fixInputs w ourInps = mapM inspect where isParam x (y, _) = x == y inspect (_, SOAC.Input ts v _) - | Just (p,pInp) <- find (isParam v) ourInps = do - let pInp' = SOAC.transformRows ts pInp - p' <- newNameFromString $ baseString p - return (p', pInp') - + | Just (p, pInp) <- find (isParam v) ourInps = do + let pInp' = SOAC.transformRows ts pInp + p' <- newNameFromString $ baseString p + return (p', pInp') inspect (param, SOAC.Input ts a t) = do param' <- newNameFromString (baseString param ++ "_rep") return (param', SOAC.Input (ts SOAC.|> SOAC.Replicate mempty (Shape [w])) a t) 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 @@ -1,5 +1,6 @@ -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} + -- | High-level representation of SOACs. When performing -- SOAC-transformations, operating on normal 'Exp' values is somewhat -- of a nuisance, as they can represent terms that are not proper @@ -21,80 +22,94 @@ -- import qualified Futhark.Analysis.HORep.SOAC as SOAC -- @ module Futhark.Analysis.HORep.SOAC - ( - -- * SOACs - SOAC (..) - , Futhark.ScremaForm(..) - , inputs - , setInputs - , lambda - , setLambda - , typeOf - , width - -- ** Converting to and from expressions - , NotSOAC (..) - , fromExp - , toExp - , toSOAC - -- * SOAC inputs - , Input (..) - , varInput - , identInput - , isVarInput - , isVarishInput - , addTransform - , addInitialTransforms - , inputArray - , inputRank - , inputType - , inputRowType - , transformRows - , transposeInput - -- ** Input transformations - , ArrayTransforms - , noTransforms - , nullTransforms - , (|>) - , (<|) - , viewf - , ViewF(..) - , viewl - , ViewL(..) - , ArrayTransform(..) - , transformFromExp - , soacToStream + ( -- * SOACs + SOAC (..), + Futhark.ScremaForm (..), + inputs, + setInputs, + lambda, + setLambda, + typeOf, + width, + + -- ** Converting to and from expressions + NotSOAC (..), + fromExp, + toExp, + toSOAC, + + -- * SOAC inputs + Input (..), + varInput, + identInput, + isVarInput, + isVarishInput, + addTransform, + addInitialTransforms, + inputArray, + inputRank, + inputType, + inputRowType, + transformRows, + transposeInput, + + -- ** Input transformations + ArrayTransforms, + noTransforms, + nullTransforms, + (|>), + (<|), + viewf, + ViewF (..), + viewl, + ViewL (..), + ArrayTransform (..), + transformFromExp, + soacToStream, ) - where +where import Data.Foldable as Foldable import Data.Maybe import qualified Data.Sequence as Seq - +import Futhark.Construct hiding (toExp) +import Futhark.IR hiding + ( Iota, + Rearrange, + Replicate, + Reshape, + Var, + typeOf, + ) import qualified Futhark.IR as Futhark import Futhark.IR.SOACS.SOAC - (StreamForm(..), ScremaForm(..), scremaType, getStreamAccums, HistOp(..), StreamOrd(..)) + ( HistOp (..), + ScremaForm (..), + StreamForm (..), + StreamOrd (..), + getStreamAccums, + scremaType, + ) import qualified Futhark.IR.SOACS.SOAC as Futhark -import Futhark.IR - hiding (Var, Iota, Rearrange, Reshape, Replicate, typeOf) -import Futhark.Transform.Substitute -import Futhark.Construct hiding (toExp) import Futhark.Transform.Rename (renameLambda) -import qualified Futhark.Util.Pretty as PP +import Futhark.Transform.Substitute import Futhark.Util.Pretty (ppr, text) +import qualified Futhark.Util.Pretty as PP -- | A single, simple transformation. If you want several, don't just -- create a list, use 'ArrayTransforms' instead. -data ArrayTransform = Rearrange Certificates [Int] - -- ^ A permutation of an otherwise valid input. - | Reshape Certificates (ShapeChange SubExp) - -- ^ A reshaping of an otherwise valid input. - | ReshapeOuter Certificates (ShapeChange SubExp) - -- ^ A reshaping of the outer dimension. - | ReshapeInner Certificates (ShapeChange SubExp) - -- ^ A reshaping of everything but the outer dimension. - | Replicate Certificates Shape - -- ^ Replicate the rows of the array a number of times. - deriving (Show, Eq, Ord) +data ArrayTransform + = -- | A permutation of an otherwise valid input. + Rearrange Certificates [Int] + | -- | A reshaping of an otherwise valid input. + Reshape Certificates (ShapeChange SubExp) + | -- | A reshaping of the outer dimension. + ReshapeOuter Certificates (ShapeChange SubExp) + | -- | A reshaping of everything but the outer dimension. + ReshapeInner Certificates (ShapeChange SubExp) + | -- | Replicate the rows of the array a number of times. + Replicate Certificates Shape + deriving (Show, Eq, Ord) instance Substitute ArrayTransform where substituteNames substs (Rearrange cs xs) = @@ -124,8 +139,8 @@ instance Semigroup ArrayTransforms where ts1 <> ts2 = case viewf ts2 of - t :< ts2' -> (ts1 |> t) <> ts2' - EmptyF -> ts1 + t :< ts2' -> (ts1 |> t) <> ts2' + EmptyF -> ts1 instance Monoid ArrayTransforms where mempty = noTransforms @@ -145,59 +160,67 @@ -- | Decompose the input-end of the transformation sequence. viewf :: ArrayTransforms -> ViewF viewf (ArrayTransforms s) = case Seq.viewl s of - t Seq.:< s' -> t :< ArrayTransforms s' - Seq.EmptyL -> EmptyF + t Seq.:< s' -> t :< ArrayTransforms s' + Seq.EmptyL -> EmptyF -- | A view of the first transformation to be applied. -data ViewF = EmptyF - | ArrayTransform :< ArrayTransforms +data ViewF + = EmptyF + | ArrayTransform :< ArrayTransforms -- | Decompose the output-end of the transformation sequence. viewl :: ArrayTransforms -> ViewL viewl (ArrayTransforms s) = case Seq.viewr s of - s' Seq.:> t -> ArrayTransforms s' :> t - Seq.EmptyR -> EmptyL + s' Seq.:> t -> ArrayTransforms s' :> t + Seq.EmptyR -> EmptyL -- | A view of the last transformation to be applied. -data ViewL = EmptyL - | ArrayTransforms :> ArrayTransform +data ViewL + = EmptyL + | ArrayTransforms :> ArrayTransform -- | Add a transform to the end of the transformation list. (|>) :: ArrayTransforms -> ArrayTransform -> ArrayTransforms (|>) = flip $ addTransform' extract add $ uncurry (flip (,)) - where extract ts' = case viewl ts' of - EmptyL -> Nothing - ts'' :> t' -> Just (t', ts'') - add t' (ArrayTransforms ts') = ArrayTransforms $ ts' Seq.|> t' + where + extract ts' = case viewl ts' of + EmptyL -> Nothing + ts'' :> t' -> Just (t', ts'') + add t' (ArrayTransforms ts') = ArrayTransforms $ ts' Seq.|> t' -- | Add a transform at the beginning of the transformation list. (<|) :: ArrayTransform -> ArrayTransforms -> ArrayTransforms (<|) = addTransform' extract add id - where extract ts' = case viewf ts' of - EmptyF -> Nothing - t' :< ts'' -> Just (t', ts'') - add t' (ArrayTransforms ts') = ArrayTransforms $ t' Seq.<| ts' + where + extract ts' = case viewf ts' of + EmptyF -> Nothing + t' :< ts'' -> Just (t', ts'') + add t' (ArrayTransforms ts') = ArrayTransforms $ t' Seq.<| ts' -addTransform' :: (ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)) - -> (ArrayTransform -> ArrayTransforms -> ArrayTransforms) - -> ((ArrayTransform,ArrayTransform) -> (ArrayTransform,ArrayTransform)) - -> ArrayTransform -> ArrayTransforms - -> ArrayTransforms +addTransform' :: + (ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)) -> + (ArrayTransform -> ArrayTransforms -> ArrayTransforms) -> + ((ArrayTransform, ArrayTransform) -> (ArrayTransform, ArrayTransform)) -> + ArrayTransform -> + ArrayTransforms -> + ArrayTransforms addTransform' extract add swap t ts = fromMaybe (t `add` ts) $ do (t', ts') <- extract ts combined <- uncurry combineTransforms $ swap (t', t) - Just $ if identityTransform combined then ts' - else addTransform' extract add swap combined ts' + Just $ + if identityTransform combined + then ts' + else addTransform' extract add swap combined ts' identityTransform :: ArrayTransform -> Bool identityTransform (Rearrange _ perm) = - Foldable.and $ zipWith (==) perm [0..] + Foldable.and $ zipWith (==) perm [0 ..] identityTransform _ = False combineTransforms :: ArrayTransform -> ArrayTransform -> Maybe ArrayTransform combineTransforms (Rearrange cs2 perm2) (Rearrange cs1 perm1) = - Just $ Rearrange (cs1<>cs2) $ perm2 `rearrangeCompose` perm1 + Just $ Rearrange (cs1 <> cs2) $ perm2 `rearrangeCompose` perm1 combineTransforms _ _ = Nothing -- | Given an expression, determine whether the expression represents @@ -219,17 +242,20 @@ -- plain expressions. The transforms are done left-to-right, that is, -- the first element of the 'ArrayTransform' list is applied first. data Input = Input ArrayTransforms VName Type - deriving (Show, Eq, Ord) + deriving (Show, Eq, Ord) instance Substitute Input where substituteNames substs (Input ts v t) = - Input (substituteNames substs ts) - (substituteNames substs v) (substituteNames substs t) + Input + (substituteNames substs ts) + (substituteNames substs v) + (substituteNames substs t) -- | Create a plain array variable input with no transformations. varInput :: HasScope t f => VName -> f Input varInput v = withType <$> lookupType v - where withType = Input (ArrayTransforms Seq.empty) v + where + withType = Input (ArrayTransforms Seq.empty) v -- | Create a plain array variable input with no transformations, from an 'Ident'. identInput :: Ident -> Input @@ -239,15 +265,16 @@ -- return the variable. isVarInput :: Input -> Maybe VName isVarInput (Input ts v _) | nullTransforms ts = Just v -isVarInput _ = Nothing +isVarInput _ = Nothing -- | If the given input is a plain variable input, with no non-vacuous transforms, -- return the variable. isVarishInput :: Input -> Maybe VName isVarishInput (Input ts v t) | nullTransforms ts = Just v - | Reshape cs [DimCoercion _] :< ts' <- viewf ts, cs == mempty = - isVarishInput $ Input ts' v t + | Reshape cs [DimCoercion _] :< ts' <- viewf ts, + cs == mempty = + isVarishInput $ Input ts' v t isVarishInput _ = Nothing -- | Add a transformation to the end of the transformation list. @@ -261,33 +288,32 @@ addInitialTransforms ts (Input ots a t) = Input (ts <> ots) a t -- | Convert SOAC inputs to the corresponding expressions. -inputsToSubExps :: (MonadBinder m) => - [Input] -> m [VName] +inputsToSubExps :: + (MonadBinder m) => + [Input] -> + m [VName] inputsToSubExps = mapM inputToExp' - where inputToExp' (Input (ArrayTransforms ts) a _) = - foldlM transform a ts - - transform ia (Replicate cs n) = - certifying cs $ - letExp "repeat" $ BasicOp $ Futhark.Replicate n (Futhark.Var ia) - - transform ia (Rearrange cs perm) = - certifying cs $ - letExp "rearrange" $ BasicOp $ Futhark.Rearrange perm ia - - transform ia (Reshape cs shape) = - certifying cs $ - letExp "reshape" $ BasicOp $ Futhark.Reshape shape ia - - transform ia (ReshapeOuter cs shape) = do - shape' <- reshapeOuter shape 1 . arrayShape <$> lookupType ia - certifying cs $ - letExp "reshape_outer" $ BasicOp $ Futhark.Reshape shape' ia + where + inputToExp' (Input (ArrayTransforms ts) a _) = + foldlM transform a ts - transform ia (ReshapeInner cs shape) = do - shape' <- reshapeInner shape 1 . arrayShape <$> lookupType ia - certifying cs $ - letExp "reshape_inner" $ BasicOp $ Futhark.Reshape shape' ia + transform ia (Replicate cs n) = + certifying cs $ + letExp "repeat" $ BasicOp $ Futhark.Replicate n (Futhark.Var ia) + transform ia (Rearrange cs perm) = + certifying cs $ + letExp "rearrange" $ BasicOp $ Futhark.Rearrange perm ia + transform ia (Reshape cs shape) = + certifying cs $ + letExp "reshape" $ BasicOp $ Futhark.Reshape shape ia + transform ia (ReshapeOuter cs shape) = do + shape' <- reshapeOuter shape 1 . arrayShape <$> lookupType ia + certifying cs $ + letExp "reshape_outer" $ BasicOp $ Futhark.Reshape shape' ia + transform ia (ReshapeInner cs shape) = do + shape' <- reshapeInner shape 1 . arrayShape <$> lookupType ia + certifying cs $ + letExp "reshape_inner" $ BasicOp $ Futhark.Reshape shape' ia -- | Return the array name of the input. inputArray :: Input -> VName @@ -297,18 +323,19 @@ inputType :: Input -> Type inputType (Input (ArrayTransforms ts) _ at) = Foldable.foldl transformType at ts - where transformType t (Replicate _ shape) = - arrayOfShape t shape - transformType t (Rearrange _ perm) = - rearrangeType perm t - transformType t (Reshape _ shape) = - t `setArrayShape` newShape shape - transformType t (ReshapeOuter _ shape) = - let Shape oldshape = arrayShape t - in t `setArrayShape` Shape (newDims shape ++ drop 1 oldshape) - transformType t (ReshapeInner _ shape) = - let Shape oldshape = arrayShape t - in t `setArrayShape` Shape (take 1 oldshape ++ newDims shape) + where + transformType t (Replicate _ shape) = + arrayOfShape t shape + transformType t (Rearrange _ perm) = + rearrangeType perm t + transformType t (Reshape _ shape) = + t `setArrayShape` newShape shape + transformType t (ReshapeOuter _ shape) = + let Shape oldshape = arrayShape t + in t `setArrayShape` Shape (newDims shape ++ drop 1 oldshape) + transformType t (ReshapeInner _ shape) = + let Shape oldshape = arrayShape t + in t `setArrayShape` Shape (take 1 oldshape ++ newDims shape) -- | Return the row type of an input. Just a convenient alias. inputRowType :: Input -> Type @@ -323,47 +350,51 @@ transformRows :: ArrayTransforms -> Input -> Input transformRows (ArrayTransforms ts) = flip (Foldable.foldl transformRows') ts - where transformRows' inp (Rearrange cs perm) = - addTransform (Rearrange cs (0:map (+1) perm)) inp - transformRows' inp (Reshape cs shape) = - addTransform (ReshapeInner cs shape) inp - transformRows' inp (Replicate cs n) - | inputRank inp == 1 = - Rearrange mempty [1,0] `addTransform` - (Replicate cs n `addTransform` inp) - | otherwise = - Rearrange mempty (2:0:1:[3..inputRank inp]) `addTransform` - (Replicate cs n `addTransform` - (Rearrange mempty (1:0:[2..inputRank inp-1]) `addTransform` inp)) - transformRows' inp nts = - error $ "transformRows: Cannot transform this yet:\n" ++ show nts ++ "\n" ++ show inp + where + transformRows' inp (Rearrange cs perm) = + addTransform (Rearrange cs (0 : map (+ 1) perm)) inp + transformRows' inp (Reshape cs shape) = + addTransform (ReshapeInner cs shape) inp + transformRows' inp (Replicate cs n) + | inputRank inp == 1 = + Rearrange mempty [1, 0] + `addTransform` (Replicate cs n `addTransform` inp) + | otherwise = + Rearrange mempty (2 : 0 : 1 : [3 .. inputRank inp]) + `addTransform` ( Replicate cs n + `addTransform` (Rearrange mempty (1 : 0 : [2 .. inputRank inp -1]) `addTransform` inp) + ) + transformRows' inp nts = + error $ "transformRows: Cannot transform this yet:\n" ++ show nts ++ "\n" ++ show inp -- | Add to the input a 'Rearrange' transform that performs an @(k,n)@ -- transposition. The new transform will be at the end of the current -- transformation list. transposeInput :: Int -> Int -> Input -> Input transposeInput k n inp = - addTransform (Rearrange mempty $ transposeIndex k n [0..inputRank inp-1]) inp + addTransform (Rearrange mempty $ transposeIndex k n [0 .. inputRank inp -1]) inp -- | 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)] - | Screma SubExp (ScremaForm lore) [Input] - | Hist SubExp [HistOp lore] (Lambda lore) [Input] - deriving (Eq, Show) +data SOAC lore + = Stream SubExp (StreamForm lore) (Lambda lore) [Input] + | Scatter SubExp (Lambda lore) [Input] [(SubExp, Int, VName)] + | Screma SubExp (ScremaForm lore) [Input] + | Hist SubExp [HistOp lore] (Lambda lore) [Input] + deriving (Eq, Show) instance PP.Pretty Input where ppr (Input (ArrayTransforms ts) arr _) = foldl f (ppr arr) ts - where f e (Rearrange cs perm) = - text "rearrange" <> ppr cs <> PP.apply [PP.apply (map ppr perm), e] - f e (Reshape cs shape) = - text "reshape" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] - f e (ReshapeOuter cs shape) = - text "reshape_outer" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] - f e (ReshapeInner cs shape) = - text "reshape_inner" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] - f e (Replicate cs ne) = - text "replicate" <> ppr cs <> PP.apply [ppr ne, e] + where + f e (Rearrange cs perm) = + text "rearrange" <> ppr cs <> PP.apply [PP.apply (map ppr perm), e] + f e (Reshape cs shape) = + text "reshape" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] + f e (ReshapeOuter cs shape) = + text "reshape_outer" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] + f e (ReshapeInner cs shape) = + text "reshape_inner" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] + f e (Replicate cs ne) = + text "replicate" <> ppr cs <> PP.apply [ppr ne, e] instance PrettyLore lore => PP.Pretty (SOAC lore) where ppr (Screma w form arrs) = Futhark.ppScrema w form arrs @@ -373,9 +404,9 @@ -- | Returns the inputs used in a SOAC. inputs :: SOAC lore -> [Input] -inputs (Stream _ _ _ arrs) = arrs -inputs (Scatter _len _lam ivs _as) = ivs -inputs (Screma _ _ 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. @@ -391,11 +422,11 @@ newWidth :: [Input] -> SubExp -> SubExp newWidth [] w = w -newWidth (inp:_) _ = arraySize 0 $ inputType inp +newWidth (inp : _) _ = arraySize 0 $ inputType inp -- | 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 @@ -414,16 +445,19 @@ -- | The return type of a SOAC. typeOf :: SOAC lore -> [Type] typeOf (Stream w form lam _) = - let nes = getStreamAccums form + let nes = getStreamAccums form accrtps = take (length nes) $ lambdaReturnType lam - arrtps = [ arrayOf (stripArray 1 t) (Shape [w]) NoUniqueness - | t <- drop (length nes) (lambdaReturnType lam) ] - in accrtps ++ arrtps + 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 - where lam_ts = lambdaReturnType lam - n = length lam_ts - (aws, _, _) = unzip3 dests + where + lam_ts = lambdaReturnType lam + n = length lam_ts + (aws, _, _) = unzip3 dests typeOf (Screma w form _) = scremaType w form typeOf (Hist _ ops _ _) = do @@ -439,13 +473,17 @@ width (Hist w _ _ _) = w -- | Convert a SOAC to the corresponding expression. -toExp :: (MonadBinder m, Op (Lore m) ~ Futhark.SOAC (Lore m)) => - SOAC (Lore m) -> m (Exp (Lore m)) +toExp :: + (MonadBinder m, Op (Lore m) ~ Futhark.SOAC (Lore m)) => + SOAC (Lore m) -> + m (Exp (Lore m)) toExp soac = Op <$> toSOAC soac -- | Convert a SOAC to a Futhark-level SOAC. -toSOAC :: MonadBinder m => - SOAC (Lore m) -> m (Futhark.SOAC (Lore m)) +toSOAC :: + MonadBinder m => + SOAC (Lore m) -> + m (Futhark.SOAC (Lore m)) toSOAC (Stream w form lam inps) = Futhark.Stream w form lam <$> inputsToSubExps inps toSOAC (Scatter len lam ivs dests) = do @@ -458,14 +496,18 @@ -- | The reason why some expression cannot be converted to a 'SOAC' -- value. -data NotSOAC = NotSOAC -- ^ The expression is not a (tuple-)SOAC at all. - deriving (Show) +data NotSOAC + = -- | The expression is not a (tuple-)SOAC at all. + NotSOAC + deriving (Show) -- | Either convert an expression to the normalised SOAC -- representation, or a reason why the expression does not have the -- valid form. -fromExp :: (Op lore ~ Futhark.SOAC lore, HasScope lore m) => - Exp lore -> m (Either NotSOAC (SOAC lore)) +fromExp :: + (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.Scatter len lam ivs as)) = @@ -479,159 +521,213 @@ -- | To-Stream translation of SOACs. -- Returns the Stream SOAC and the -- extra-accumulator body-result ident if any. -soacToStream :: (MonadFreshNames m, Bindable lore, Op lore ~ Futhark.SOAC lore) => - SOAC lore -> m (SOAC lore,[Ident]) +soacToStream :: + (MonadFreshNames m, Bindable lore, Op lore ~ Futhark.SOAC lore) => + SOAC lore -> + m (SOAC lore, [Ident]) soacToStream soac = do - chunk_param <- newParam "chunk" $ Prim int32 - let chvar= Futhark.Var $ paramName chunk_param + chunk_param <- newParam "chunk" $ Prim int64 + let chvar = Futhark.Var $ paramName chunk_param (lam, inps) = (lambda soac, inputs soac) w = width soac - lam' <- renameLambda lam - let arrrtps= mapType w lam + lam' <- renameLambda lam + let arrrtps = mapType w lam -- the chunked-outersize of the array result and input types - loutps = [ arrayOfRow t chvar | t <- map rowType arrrtps ] - lintps = [ arrayOfRow t chvar | t <- map inputRowType inps ] + loutps = [arrayOfRow t chvar | t <- map rowType arrrtps] + lintps = [arrayOfRow t chvar | t <- map inputRowType inps] strm_inpids <- mapM (newParam "inp") lintps -- Treat each SOAC case individually: case soac of Screma _ form _ | Just _ <- Futhark.isMapSOAC form -> do - -- Map(f,a) => is translated in strem's body to: - -- let strm_resids = map(f,a_ch) in strm_resids - -- - -- array result and input IDs of the stream's lambda - strm_resids <- mapM (newIdent "res") loutps - let insoac = Futhark.Screma chvar (Futhark.mapSOAC lam') $ map paramName strm_inpids - insbnd = mkLet [] strm_resids $ Op insoac - strmbdy= mkBody (oneStm insbnd) $ map (Futhark.Var . identName) strm_resids - strmpar= chunk_param:strm_inpids - 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, []) - + -- Map(f,a) => is translated in strem's body to: + -- let strm_resids = map(f,a_ch) in strm_resids + -- + -- array result and input IDs of the stream's lambda + strm_resids <- mapM (newIdent "res") loutps + let insoac = Futhark.Screma chvar (Futhark.mapSOAC lam') $ map paramName strm_inpids + insbnd = mkLet [] strm_resids $ Op insoac + strmbdy = mkBody (oneStm insbnd) $ map (Futhark.Var . identName) strm_resids + strmpar = chunk_param : strm_inpids + 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, []) | 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: - -- 1. let (scan0_ids,map_resids) = scanomap(scan_lam, nes, map_lam, a_ch) - -- 2. let strm_resids = map (acc `+`,nes, scan0_ids) - -- 3. let outerszm1id = sizeof(0,strm_resids) - 1 - -- 4. let lasteel_ids = if outerszm1id < 0 - -- then nes - -- else strm_resids[outerszm1id] - -- 5. let acc' = acc + lasteel_ids - -- {acc', strm_resids, map_resids} - -- the array and accumulator result types - let scan_arr_ts = map (`arrayOfRow` chvar) $ lambdaReturnType scan_lam - map_arr_ts = drop (length nes) loutps - accrtps = lambdaReturnType scan_lam - - -- array result and input IDs of the stream's lambda - strm_resids <- mapM (newIdent "res") scan_arr_ts - scan0_ids <- mapM (newIdent "resarr0") scan_arr_ts - map_resids <- mapM (newIdent "map_res") map_arr_ts + -- scanomap(scan_lam,nes,map_lam,a) => is translated in strem's body to: + -- 1. let (scan0_ids,map_resids) = scanomap(scan_lam, nes, map_lam, a_ch) + -- 2. let strm_resids = map (acc `+`,nes, scan0_ids) + -- 3. let outerszm1id = sizeof(0,strm_resids) - 1 + -- 4. let lasteel_ids = if outerszm1id < 0 + -- then nes + -- else strm_resids[outerszm1id] + -- 5. let acc' = acc + lasteel_ids + -- {acc', strm_resids, map_resids} + -- the array and accumulator result types + let scan_arr_ts = map (`arrayOfRow` chvar) $ lambdaReturnType scan_lam + map_arr_ts = drop (length nes) loutps + accrtps = lambdaReturnType scan_lam - lastel_ids <- mapM (newIdent "lstel") accrtps - lastel_tmp_ids <- mapM (newIdent "lstel_tmp") accrtps - empty_arr <- newIdent "empty_arr" $ Prim Bool - inpacc_ids <- mapM (newParam "inpacc") accrtps - outszm1id <- newIdent "szm1" $ Prim int32 - -- 1. let (scan0_ids,map_resids) = scanomap(scan_lam,nes,map_lam,a_ch) - let insbnd = mkLet [] (scan0_ids++map_resids) $ Op $ - Futhark.Screma chvar (Futhark.scanomapSOAC [Futhark.Scan scan_lam nes] lam') $ - map paramName strm_inpids - -- 2. let outerszm1id = chunksize - 1 - outszm1bnd = mkLet [] [outszm1id] $ BasicOp $ - BinOp (Sub Int32 OverflowUndef) - (Futhark.Var $ paramName chunk_param) - (constant (1::Int32)) - -- 3. let lasteel_ids = ... - empty_arr_bnd = mkLet [] [empty_arr] $ BasicOp $ CmpOp (CmpSlt Int32) - (Futhark.Var $ identName outszm1id) - (constant (0::Int32)) - leltmpbnds= zipWith (\ lid arrid -> mkLet [] [lid] $ BasicOp $ - Index (identName arrid) $ - fullSlice (identType arrid) - [DimFix $ Futhark.Var $ identName outszm1id] - ) lastel_tmp_ids scan0_ids - lelbnd = mkLet [] lastel_ids $ - If (Futhark.Var $ identName empty_arr) - (mkBody mempty nes) - (mkBody (stmsFromList leltmpbnds) $ - map (Futhark.Var . identName) lastel_tmp_ids) $ - ifCommon $ map identType lastel_tmp_ids - -- 4. let strm_resids = map (acc `+`,nes, scan0_ids) - maplam <- mkMapPlusAccLam (map (Futhark.Var . paramName) inpacc_ids) scan_lam - let mapbnd = mkLet [] strm_resids $ Op $ - Futhark.Screma chvar (Futhark.mapSOAC maplam) $ - map identName scan0_ids - -- 5. let acc' = acc + lasteel_ids - addlelbdy <- mkPlusBnds scan_lam $ map Futhark.Var $ - map paramName inpacc_ids++map identName lastel_ids - -- Finally, construct the stream - let (addlelbnd,addlelres) = (bodyStms addlelbdy, bodyResult addlelbdy) - strmbdy= mkBody (stmsFromList [insbnd,outszm1bnd,empty_arr_bnd,lelbnd,mapbnd]<>addlelbnd) $ - addlelres ++ map (Futhark.Var . identName) (strm_resids ++ map_resids) - strmpar= chunk_param:inpacc_ids++strm_inpids - strmlam= Lambda strmpar strmbdy (accrtps++loutps) - return (Stream w (Sequential nes) strmlam inps, - map paramIdent inpacc_ids) + -- array result and input IDs of the stream's lambda + strm_resids <- mapM (newIdent "res") scan_arr_ts + scan0_ids <- mapM (newIdent "resarr0") scan_arr_ts + map_resids <- mapM (newIdent "map_res") map_arr_ts + lastel_ids <- mapM (newIdent "lstel") accrtps + lastel_tmp_ids <- mapM (newIdent "lstel_tmp") accrtps + empty_arr <- newIdent "empty_arr" $ Prim Bool + inpacc_ids <- mapM (newParam "inpacc") accrtps + outszm1id <- newIdent "szm1" $ Prim int64 + -- 1. let (scan0_ids,map_resids) = scanomap(scan_lam,nes,map_lam,a_ch) + let insbnd = + mkLet [] (scan0_ids ++ map_resids) $ + Op $ + Futhark.Screma chvar (Futhark.scanomapSOAC [Futhark.Scan scan_lam nes] lam') $ + map paramName strm_inpids + -- 2. let outerszm1id = chunksize - 1 + outszm1bnd = + mkLet [] [outszm1id] $ + BasicOp $ + BinOp + (Sub Int64 OverflowUndef) + (Futhark.Var $ paramName chunk_param) + (constant (1 :: Int64)) + -- 3. let lasteel_ids = ... + empty_arr_bnd = + mkLet [] [empty_arr] $ + BasicOp $ + CmpOp + (CmpSlt Int64) + (Futhark.Var $ identName outszm1id) + (constant (0 :: Int64)) + leltmpbnds = + zipWith + ( \lid arrid -> + mkLet [] [lid] $ + BasicOp $ + Index (identName arrid) $ + fullSlice + (identType arrid) + [DimFix $ Futhark.Var $ identName outszm1id] + ) + lastel_tmp_ids + scan0_ids + lelbnd = + mkLet [] lastel_ids $ + If + (Futhark.Var $ identName empty_arr) + (mkBody mempty nes) + ( mkBody (stmsFromList leltmpbnds) $ + map (Futhark.Var . identName) lastel_tmp_ids + ) + $ ifCommon $ map identType lastel_tmp_ids + -- 4. let strm_resids = map (acc `+`,nes, scan0_ids) + maplam <- mkMapPlusAccLam (map (Futhark.Var . paramName) inpacc_ids) scan_lam + let mapbnd = + mkLet [] strm_resids $ + Op $ + Futhark.Screma chvar (Futhark.mapSOAC maplam) $ + map identName scan0_ids + -- 5. let acc' = acc + lasteel_ids + addlelbdy <- + mkPlusBnds scan_lam $ + map Futhark.Var $ + map paramName inpacc_ids ++ map identName lastel_ids + -- Finally, construct the stream + let (addlelbnd, addlelres) = (bodyStms addlelbdy, bodyResult addlelbdy) + strmbdy = + mkBody (stmsFromList [insbnd, outszm1bnd, empty_arr_bnd, lelbnd, mapbnd] <> addlelbnd) $ + addlelres ++ map (Futhark.Var . identName) (strm_resids ++ map_resids) + strmpar = chunk_param : inpacc_ids ++ strm_inpids + strmlam = Lambda strmpar strmbdy (accrtps ++ loutps) + return + ( Stream w (Sequential nes) strmlam inps, + map paramIdent inpacc_ids + ) | Just (reds, _) <- Futhark.isRedomapSOAC form, Futhark.Reduce comm lamin nes <- Futhark.singleReduce reds -> do - -- Redomap(+,lam,nes,a) => is translated in strem's body to: - -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in - -- 2. let acc' = acc + acc0_ids in - -- {acc', strm_resids} + -- Redomap(+,lam,nes,a) => is translated in strem's body to: + -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in + -- 2. let acc' = acc + acc0_ids in + -- {acc', strm_resids} - let accrtps= take (length nes) $ lambdaReturnType lam - -- the chunked-outersize of the array result and input types - loutps' = drop (length nes) loutps - -- the lambda with proper index - foldlam = lam' - -- array result and input IDs of the stream's lambda - strm_resids <- mapM (newIdent "res") loutps' - inpacc_ids <- mapM (newParam "inpacc") accrtps - acc0_ids <- mapM (newIdent "acc0" ) accrtps - -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in - let insoac = Futhark.Screma chvar - (Futhark.redomapSOAC [Futhark.Reduce comm lamin nes] foldlam) $ - map paramName strm_inpids - insbnd = mkLet [] (acc0_ids++strm_resids) $ Op insoac - -- 2. let acc' = acc + acc0_ids in - addaccbdy <- mkPlusBnds lamin $ map Futhark.Var $ - map paramName inpacc_ids++map identName acc0_ids - -- Construct the stream - let (addaccbnd,addaccres) = (bodyStms addaccbdy, bodyResult addaccbdy) - strmbdy= mkBody (oneStm insbnd <> addaccbnd) $ - addaccres ++ map (Futhark.Var . identName) strm_resids - 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, []) + let accrtps = take (length nes) $ lambdaReturnType lam + -- the chunked-outersize of the array result and input types + loutps' = drop (length nes) loutps + -- the lambda with proper index + foldlam = lam' + -- array result and input IDs of the stream's lambda + strm_resids <- mapM (newIdent "res") loutps' + inpacc_ids <- mapM (newParam "inpacc") accrtps + acc0_ids <- mapM (newIdent "acc0") accrtps + -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in + let insoac = + Futhark.Screma + chvar + (Futhark.redomapSOAC [Futhark.Reduce comm lamin nes] foldlam) + $ map paramName strm_inpids + insbnd = mkLet [] (acc0_ids ++ strm_resids) $ Op insoac + -- 2. let acc' = acc + acc0_ids in + addaccbdy <- + mkPlusBnds lamin $ + map Futhark.Var $ + map paramName inpacc_ids ++ map identName acc0_ids + -- Construct the stream + let (addaccbnd, addaccres) = (bodyStms addaccbdy, bodyResult addaccbdy) + strmbdy = + mkBody (oneStm insbnd <> addaccbnd) $ + addaccres ++ map (Futhark.Var . identName) strm_resids + 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, []) -- Otherwise it cannot become a stream. - _ -> return (soac,[]) - where mkMapPlusAccLam :: (MonadFreshNames m, Bindable lore) - => [SubExp] -> Lambda lore -> m (Lambda lore) - mkMapPlusAccLam accs plus = do - let (accpars, rempars) = splitAt (length accs) $ lambdaParams plus - parbnds = zipWith (\ par se -> mkLet [] [paramIdent par] - (BasicOp $ SubExp se) - ) accpars accs - plus_bdy = lambdaBody plus - newlambdy = Body (bodyDec plus_bdy) - (stmsFromList parbnds <> bodyStms plus_bdy) - (bodyResult plus_bdy) - renameLambda $ Lambda rempars newlambdy $ lambdaReturnType plus + _ -> return (soac, []) + where + mkMapPlusAccLam :: + (MonadFreshNames m, Bindable lore) => + [SubExp] -> + Lambda lore -> + m (Lambda lore) + mkMapPlusAccLam accs plus = do + let (accpars, rempars) = splitAt (length accs) $ lambdaParams plus + parbnds = + zipWith + ( \par se -> + mkLet + [] + [paramIdent par] + (BasicOp $ SubExp se) + ) + accpars + accs + plus_bdy = lambdaBody plus + newlambdy = + Body + (bodyDec plus_bdy) + (stmsFromList parbnds <> bodyStms plus_bdy) + (bodyResult plus_bdy) + renameLambda $ Lambda rempars newlambdy $ lambdaReturnType plus - mkPlusBnds :: (MonadFreshNames m, Bindable lore) - => Lambda lore -> [SubExp] -> m (Body lore) - mkPlusBnds plus accels = do - plus' <- renameLambda plus - let parbnds = zipWith (\ par se -> mkLet [] [paramIdent par] - (BasicOp $ SubExp se) - ) (lambdaParams plus') accels - body = lambdaBody plus' - return $ body { bodyStms = stmsFromList parbnds <> bodyStms body } + mkPlusBnds :: + (MonadFreshNames m, Bindable lore) => + Lambda lore -> + [SubExp] -> + m (Body lore) + mkPlusBnds plus accels = do + plus' <- renameLambda plus + let parbnds = + zipWith + ( \par se -> + mkLet + [] + [paramIdent par] + (BasicOp $ SubExp se) + ) + (lambdaParams plus') + accels + body = lambdaBody plus' + return $ body {bodyStms = stmsFromList parbnds <> bodyStms body} diff --git a/src/Futhark/Analysis/Metrics.hs b/src/Futhark/Analysis/Metrics.hs --- a/src/Futhark/Analysis/Metrics.hs +++ b/src/Futhark/Analysis/Metrics.hs @@ -1,28 +1,29 @@ {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Trustworthy #-} + -- | Abstract Syntax Tree metrics. This is used in the @futhark test@ -- program, for the @structure@ stanzas. module Futhark.Analysis.Metrics - ( AstMetrics(..) - , progMetrics + ( AstMetrics (..), + progMetrics, - -- * Extensibility - , OpMetrics(..) - , seen - , inside - , MetricsM - , stmMetrics - , lambdaMetrics - ) where + -- * Extensibility + OpMetrics (..), + seen, + inside, + MetricsM, + stmMetrics, + lambdaMetrics, + ) +where import Control.Monad.Writer -import Data.Text (Text) -import qualified Data.Text as T import Data.List (tails) import qualified Data.Map.Strict as M - +import Data.Text (Text) +import qualified Data.Text as T import Futhark.IR -- | AST metrics are simply a collection from identifiable node names @@ -31,15 +32,17 @@ instance Show AstMetrics where show (AstMetrics m) = unlines $ map metric $ M.toList m - where metric (k, v) = pretty k ++ " " ++ pretty v + where + metric (k, v) = pretty k ++ " " ++ pretty v instance Read AstMetrics where readsPrec _ s = maybe [] success $ mapM onLine $ lines s - where onLine l = case words l of - [k, x] | [(n, "")] <- reads x -> Just (T.pack k, n) - _ -> Nothing - success m = [(AstMetrics $ M.fromList m, "")] + where + onLine l = case words l of + [k, x] | [(n, "")] <- reads x -> Just (T.pack k, n) + _ -> Nothing + success m = [(AstMetrics $ M.fromList m, "")] -- | Compute the metrics for some operation. class OpMetrics op where @@ -59,15 +62,21 @@ actualMetrics :: CountMetrics -> AstMetrics actualMetrics (CountMetrics metrics) = AstMetrics $ M.fromListWith (+) $ concatMap expand metrics - where expand (ctx, k) = - [ (T.intercalate "/" (ctx' ++ [k]), 1) - | ctx' <- tails $ "" : ctx ] + where + expand (ctx, k) = + [ (T.intercalate "/" (ctx' ++ [k]), 1) + | ctx' <- tails $ "" : ctx + ] -- | This monad is used for computing metrics. It internally keeps -- track of what we've seen so far. Use 'seen' to add more stuff. -newtype MetricsM a = MetricsM { runMetricsM :: Writer CountMetrics a } - deriving (Monad, Applicative, Functor, - MonadWriter CountMetrics) +newtype MetricsM a = MetricsM {runMetricsM :: Writer CountMetrics a} + deriving + ( Monad, + Applicative, + Functor, + MonadWriter CountMetrics + ) -- | Add this node to the current tally. seen :: Text -> MetricsM () @@ -78,16 +87,19 @@ -- enclosed operation. inside :: Text -> MetricsM () -> MetricsM () inside what m = seen what >> censor addWhat m - where addWhat (CountMetrics metrics) = - CountMetrics (map addWhat' metrics) - addWhat' (ctx, k) = (what : ctx, k) + where + addWhat (CountMetrics metrics) = + CountMetrics (map addWhat' metrics) + addWhat' (ctx, k) = (what : ctx, k) -- | Compute the metrics for a program. progMetrics :: OpMetrics (Op lore) => Prog lore -> AstMetrics progMetrics prog = - actualMetrics $ execWriter $ runMetricsM $ do - mapM_ funDefMetrics $ progFuns prog - mapM_ stmMetrics $ progConsts prog + actualMetrics $ + execWriter $ + runMetricsM $ do + mapM_ funDefMetrics $ progFuns prog + mapM_ stmMetrics $ progConsts prog funDefMetrics :: OpMetrics (Op lore) => FunDef lore -> MetricsM () funDefMetrics = bodyMetrics . funDefBody @@ -102,15 +114,15 @@ expMetrics :: OpMetrics (Op lore) => Exp lore -> MetricsM () expMetrics (BasicOp op) = seen "BasicOp" >> primOpMetrics op -expMetrics (DoLoop _ _ ForLoop{} body) = +expMetrics (DoLoop _ _ ForLoop {} body) = inside "DoLoop" $ seen "ForLoop" >> bodyMetrics body -expMetrics (DoLoop _ _ WhileLoop{} body) = +expMetrics (DoLoop _ _ WhileLoop {} body) = inside "DoLoop" $ seen "WhileLoop" >> bodyMetrics body expMetrics (If _ tb fb _) = inside "If" $ do inside "True" $ bodyMetrics tb inside "False" $ bodyMetrics fb -expMetrics Apply{} = +expMetrics Apply {} = seen "Apply" expMetrics (Op op) = opMetrics op @@ -118,23 +130,23 @@ primOpMetrics :: BasicOp -> MetricsM () primOpMetrics (SubExp _) = seen "SubExp" primOpMetrics (Opaque _) = seen "Opaque" -primOpMetrics ArrayLit{} = seen "ArrayLit" -primOpMetrics BinOp{} = seen "BinOp" -primOpMetrics UnOp{} = seen "UnOp" -primOpMetrics ConvOp{} = seen "ConvOp" -primOpMetrics CmpOp{} = seen "ConvOp" -primOpMetrics Assert{} = seen "Assert" -primOpMetrics Index{} = seen "Index" -primOpMetrics Update{} = seen "Update" -primOpMetrics Concat{} = seen "Concat" -primOpMetrics Copy{} = seen "Copy" -primOpMetrics Manifest{} = seen "Manifest" -primOpMetrics Iota{} = seen "Iota" -primOpMetrics Replicate{} = seen "Replicate" -primOpMetrics Scratch{} = seen "Scratch" -primOpMetrics Reshape{} = seen "Reshape" -primOpMetrics Rearrange{} = seen "Rearrange" -primOpMetrics Rotate{} = seen "Rotate" +primOpMetrics ArrayLit {} = seen "ArrayLit" +primOpMetrics BinOp {} = seen "BinOp" +primOpMetrics UnOp {} = seen "UnOp" +primOpMetrics ConvOp {} = seen "ConvOp" +primOpMetrics CmpOp {} = seen "ConvOp" +primOpMetrics Assert {} = seen "Assert" +primOpMetrics Index {} = seen "Index" +primOpMetrics Update {} = seen "Update" +primOpMetrics Concat {} = seen "Concat" +primOpMetrics Copy {} = seen "Copy" +primOpMetrics Manifest {} = seen "Manifest" +primOpMetrics Iota {} = seen "Iota" +primOpMetrics Replicate {} = seen "Replicate" +primOpMetrics Scratch {} = seen "Scratch" +primOpMetrics Reshape {} = seen "Reshape" +primOpMetrics Rearrange {} = seen "Rearrange" +primOpMetrics Rotate {} = seen "Rotate" -- | Compute metrics for this lambda. lambdaMetrics :: OpMetrics (Op lore) => Lambda lore -> MetricsM () 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,31 +1,74 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} + -- | A primitive expression is an expression where the non-leaves are -- primitive operators. Our representation does not guarantee that -- the expression is type-correct. module Futhark.Analysis.PrimExp - ( PrimExp (..) - , evalPrimExp - , primExpType - , primExpSizeAtLeast - , coerceIntPrimExp - , leafExpTypes - , true - , false - , constFoldPrimExp + ( PrimExp (..), + TPrimExp (..), + isInt8, + isInt16, + isInt32, + isInt64, + isBool, + isF32, + isF64, + evalPrimExp, + primExpType, + primExpSizeAtLeast, + coerceIntPrimExp, + leafExpTypes, + true, + false, + constFoldPrimExp, - , module Futhark.IR.Primitive - , sExt, zExt - , (.&&.), (.||.), (.<.), (.<=.), (.>.), (.>=.), (.==.), (.&.), (.|.), (.^.) - ) where + -- * Construction + module Futhark.IR.Primitive, + NumExp (..), + IntExp, + FloatExp (..), + sExt, + zExt, + (.&&.), + (.||.), + (.<.), + (.<=.), + (.>.), + (.>=.), + (.==.), + (.&.), + (.|.), + (.^.), + bNot, + sMax32, + sMin32, + sMax64, + sMin64, + sExt32, + sExt64, + zExt32, + zExt64, + fMin64, + fMax64, + ) +where -import Control.Monad -import Data.Traversable +import Control.Category +import Control.Monad import qualified Data.Map as M import qualified Data.Set as S - -import Futhark.IR.Prop.Names -import Futhark.IR.Primitive -import Futhark.Util.IntegralExp -import Futhark.Util.Pretty +import qualified Data.Text as T +import Data.Traversable +import Futhark.IR.Primitive +import Futhark.IR.Prop.Names +import Futhark.Util.IntegralExp +import Futhark.Util.Pretty +import GHC.Generics (Generic) +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) -- | A primitive expression parametrised over the representation of -- free variables. Note that the 'Functor', 'Traversable', and 'Num' @@ -33,36 +76,28 @@ -- -- Note also that the 'Num' instance assumes 'OverflowUndef' -- semantics! -data PrimExp v = LeafExp v PrimType - | ValueExp PrimValue - | BinOpExp BinOp (PrimExp v) (PrimExp v) - | CmpOpExp CmpOp (PrimExp v) (PrimExp v) - | UnOpExp UnOp (PrimExp v) - | ConvOpExp ConvOp (PrimExp v) - | FunExp String [PrimExp v] PrimType - deriving (Ord, Show) +data PrimExp v + = LeafExp v PrimType + | ValueExp PrimValue + | BinOpExp BinOp (PrimExp v) (PrimExp v) + | CmpOpExp CmpOp (PrimExp v) (PrimExp v) + | UnOpExp UnOp (PrimExp v) + | ConvOpExp ConvOp (PrimExp v) + | FunExp String [PrimExp v] PrimType + deriving (Eq, Ord, Show, Generic) --- The Eq instance upcoerces all integer constants to their largest --- type before comparing for equality. This is technically not a good --- idea, but solves annoying problems related to the Num instance --- always producing Int64s. -instance Eq v => Eq (PrimExp v) where - LeafExp x xt == LeafExp y yt = x == y && xt == yt - ValueExp (IntValue x) == ValueExp (IntValue y) = - intToInt64 x == intToInt64 y - ValueExp x == ValueExp y = - x == y - BinOpExp xop x1 x2 == BinOpExp yop y1 y2 = - xop == yop && x1 == y1 && x2 == y2 - CmpOpExp xop x1 x2 == CmpOpExp yop y1 y2 = - xop == yop && x1 == y1 && x2 == y2 - UnOpExp xop x == UnOpExp yop y = - xop == yop && x == y - ConvOpExp xop x == ConvOpExp yop y = - xop == yop && x == y - FunExp xf xargs _ == FunExp yf yargs _ = - xf == yf && xargs == yargs - _ == _ = False +instance SexpIso v => SexpIso (PrimExp v) where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "leaf") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "value") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "bin-op") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "cmp-op") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "un-op") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "conv-op") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "fun") >>> Sexp.el (iso T.unpack T.pack . sexpIso) >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End instance Functor PrimExp where fmap = fmapDefault @@ -89,48 +124,102 @@ instance FreeIn v => FreeIn (PrimExp v) where freeIn' = foldMap freeIn' +-- | A 'PrimExp' tagged with a phantom type used to provide type-safe +-- construction. Does not guarantee that the underlying expression is +-- actually type correct. +newtype TPrimExp t v = TPrimExp {untyped :: PrimExp v} + deriving (Eq, Ord, Show, Generic) + +instance SexpIso v => SexpIso (TPrimExp t v) where + sexpIso = with $ \e -> sexpIso >>> e + +instance Functor (TPrimExp t) where + fmap = fmapDefault + +instance Foldable (TPrimExp t) where + foldMap = foldMapDefault + +instance Traversable (TPrimExp t) where + traverse f (TPrimExp e) = TPrimExp <$> traverse f e + +instance FreeIn v => FreeIn (TPrimExp t v) where + freeIn' = freeIn' . untyped + +-- | This expression is of type 'Int8'. +isInt8 :: PrimExp v -> TPrimExp Int8 v +isInt8 = TPrimExp + +-- | This expression is of type 'Int16'. +isInt16 :: PrimExp v -> TPrimExp Int16 v +isInt16 = TPrimExp + +-- | This expression is of type 'Int32'. +isInt32 :: PrimExp v -> TPrimExp Int32 v +isInt32 = TPrimExp + +-- | This expression is of type 'Int64'. +isInt64 :: PrimExp v -> TPrimExp Int64 v +isInt64 = TPrimExp + +-- | This is a boolean expression. +isBool :: PrimExp v -> TPrimExp Bool v +isBool = TPrimExp + +-- | This expression is of type 'Float'. +isF32 :: PrimExp v -> TPrimExp Float v +isF32 = TPrimExp + +-- | This expression is of type 'Double'. +isF64 :: PrimExp v -> TPrimExp Double v +isF64 = TPrimExp + -- | True if the 'PrimExp' has at least this many nodes. This can be -- much more efficient than comparing with 'length' for large -- 'PrimExp's, as this function is lazy. primExpSizeAtLeast :: Int -> PrimExp v -> Bool -primExpSizeAtLeast k = maybe True (>=k) . descend 0 - where descend i _ - | i >= k = Nothing - descend i LeafExp{} = Just (i+1) - descend i ValueExp{} = Just (i+1) - descend i (BinOpExp _ x y) = do x' <- descend (i+1) x - descend x' y - descend i (CmpOpExp _ x y) = do x' <- descend (i+1) x - descend x' y - descend i (ConvOpExp _ x) = descend (i+1) x - descend i (UnOpExp _ x) = descend (i+1) x - descend i (FunExp _ args _) = foldM descend (i+1) args +primExpSizeAtLeast k = maybe True (>= k) . descend 0 + where + descend i _ + | i >= k = Nothing + descend i LeafExp {} = Just (i + 1) + descend i ValueExp {} = Just (i + 1) + descend i (BinOpExp _ x y) = do + x' <- descend (i + 1) x + descend x' y + descend i (CmpOpExp _ x y) = do + x' <- descend (i + 1) x + descend x' y + descend i (ConvOpExp _ x) = descend (i + 1) x + descend i (UnOpExp _ x) = descend (i + 1) x + descend i (FunExp _ args _) = foldM descend (i + 1) args -- | Perform quick and dirty constant folding on the top level of a -- PrimExp. This is necessary because we want to consider -- e.g. equality modulo constant folding. constFoldPrimExp :: PrimExp v -> PrimExp v -constFoldPrimExp (BinOpExp Add{} x y) +constFoldPrimExp (BinOpExp Add {} x y) | zeroIshExp x = y | zeroIshExp y = x -constFoldPrimExp (BinOpExp Sub{} x y) +constFoldPrimExp (BinOpExp Sub {} x y) | zeroIshExp y = x -constFoldPrimExp (BinOpExp Mul{} x y) +constFoldPrimExp (BinOpExp Mul {} x y) | oneIshExp x = y | oneIshExp y = x - | zeroIshExp x, IntType it <- primExpType y = - ValueExp $ IntValue $ intValue it (0::Int) - | zeroIshExp y, IntType it <- primExpType x = - ValueExp $ IntValue $ intValue it (0::Int) -constFoldPrimExp (BinOpExp SDiv{} x y) + | zeroIshExp x, + IntType it <- primExpType y = + ValueExp $ IntValue $ intValue it (0 :: Int) + | zeroIshExp y, + IntType it <- primExpType x = + ValueExp $ IntValue $ intValue it (0 :: Int) +constFoldPrimExp (BinOpExp SDiv {} x y) | oneIshExp y = x -constFoldPrimExp (BinOpExp SQuot{} x y) +constFoldPrimExp (BinOpExp SQuot {} x y) | oneIshExp y = x -constFoldPrimExp (BinOpExp UDiv{} x y) +constFoldPrimExp (BinOpExp UDiv {} x y) | oneIshExp y = x constFoldPrimExp (BinOpExp bop (ValueExp x) (ValueExp y)) | Just z <- doBinOp bop x y = - ValueExp z + ValueExp z constFoldPrimExp (BinOpExp LogAnd x y) | oneIshExp x = y | oneIshExp y = x @@ -143,124 +232,201 @@ | zeroIshExp y = x constFoldPrimExp e = e --- The Num instance performs a little bit of magic: whenever an --- expression and a constant is combined with a binary operator, the --- type of the constant may be changed to be the type of the --- expression, if they are not already the same. This permits us to --- write e.g. @x * 4@, where @x@ is an arbitrary PrimExp, and have the --- @4@ converted to the proper primitive type. We also support --- converting integers to floating point values, but not the other way --- around. All numeric instances assume unsigned integers for such --- conversions. --- --- We also perform simple constant folding, in particular to reduce --- expressions to constants so that the above works. However, it is --- still a bit of a hack. -instance Pretty v => Num (PrimExp v) where - x + y | Just z <- msum [asIntOp (`Add` OverflowUndef) x y, - asFloatOp FAdd x y] = constFoldPrimExp z - | otherwise = numBad "+" (x,y) +-- | The class of numeric types that can be used for constructing +-- 'TPrimExp's. +class NumExp t where + -- | Construct a typed expression from an integer. + fromInteger' :: Integer -> TPrimExp t v - x - y | Just z <- msum [asIntOp (`Sub` OverflowUndef) x y, - asFloatOp FSub x y] = constFoldPrimExp z - | otherwise = numBad "-" (x,y) +-- | The class of integer types that can be used for constructing +-- 'TPrimExp's. +class NumExp t => IntExp t - x * y | Just z <- msum [asIntOp (`Mul` OverflowUndef) x y, - asFloatOp FMul x y] = constFoldPrimExp z - | otherwise = numBad "*" (x,y) +instance NumExp Int8 where + fromInteger' = isInt8 . ValueExp . IntValue . Int8Value . fromInteger - abs x | IntType t <- primExpType x = UnOpExp (Abs t) x - | FloatType t <- primExpType x = UnOpExp (FAbs t) x - | otherwise = numBad "abs" x +instance IntExp Int8 - signum x | IntType t <- primExpType x = UnOpExp (SSignum t) x - | otherwise = numBad "signum" x +instance NumExp Int16 where + fromInteger' = isInt16 . ValueExp . IntValue . Int16Value . fromInteger - fromInteger = fromInt32 . fromInteger +instance IntExp Int16 -instance Pretty v => Fractional (PrimExp v) where - x / y | Just z <- msum [asFloatOp FDiv x y] = constFoldPrimExp z - | otherwise = numBad "/" (x,y) +instance NumExp Int32 where + fromInteger' = isInt32 . ValueExp . IntValue . Int32Value . fromInteger - fromRational = ValueExp . FloatValue . Float64Value . fromRational +instance IntExp Int32 -instance Pretty v => IntegralExp (PrimExp v) where - x `div` y | Just z <- msum [asIntOp (`SDiv` Unsafe) x y, - asFloatOp FDiv x y] = - constFoldPrimExp z - | otherwise = numBad "div" (x,y) +instance NumExp Int64 where + fromInteger' = isInt64 . ValueExp . IntValue . Int64Value . fromInteger - x `mod` y | Just z <- msum [asIntOp (`SMod` Unsafe) x y] = z - | otherwise = numBad "mod" (x,y) +instance IntExp Int64 - x `quot` y | oneIshExp y = x - | Just z <- msum [asIntOp (`SQuot` Unsafe) x y] = constFoldPrimExp z - | otherwise = numBad "quot" (x,y) +-- | The class of floating-point types that can be used for +-- constructing 'TPrimExp's. +class NumExp t => FloatExp t where + -- | Construct a typed expression from a rational. + fromRational' :: Rational -> TPrimExp t v - x `rem` y | Just z <- msum [asIntOp (`SRem` Unsafe) x y] = constFoldPrimExp z - | otherwise = numBad "rem" (x,y) +instance NumExp Float where + fromInteger' = TPrimExp . ValueExp . FloatValue . Float32Value . fromInteger - x `divUp` y | Just z <- msum [asIntOp (`SDivUp` Unsafe) x y] = - constFoldPrimExp z - | otherwise = numBad "divRoundingUp" (x,y) +instance NumExp Double where + fromInteger' = TPrimExp . ValueExp . FloatValue . Float64Value . fromInteger - sgn (ValueExp (IntValue i)) = Just $ signum $ valueIntegral i - sgn _ = Nothing +instance FloatExp Float where + fromRational' = TPrimExp . ValueExp . FloatValue . Float32Value . fromRational - fromInt8 = ValueExp . IntValue . Int8Value - fromInt16 = ValueExp . IntValue . Int16Value - fromInt32 = ValueExp . IntValue . Int32Value - fromInt64 = ValueExp . IntValue . Int64Value +instance FloatExp Double where + fromRational' = TPrimExp . ValueExp . FloatValue . Float64Value . fromRational +instance (NumExp t, Pretty v) => Num (TPrimExp t v) where + TPrimExp x + TPrimExp y + | Just z <- + msum + [ asIntOp (`Add` OverflowUndef) x y, + asFloatOp FAdd x y + ] = + TPrimExp $ constFoldPrimExp z + | otherwise = numBad "+" (x, y) + + TPrimExp x - TPrimExp y + | Just z <- + msum + [ asIntOp (`Sub` OverflowUndef) x y, + asFloatOp FSub x y + ] = + TPrimExp $ constFoldPrimExp z + | otherwise = numBad "-" (x, y) + + TPrimExp x * TPrimExp y + | Just z <- + msum + [ asIntOp (`Mul` OverflowUndef) x y, + asFloatOp FMul x y + ] = + TPrimExp $ constFoldPrimExp z + | otherwise = numBad "*" (x, y) + + abs (TPrimExp x) + | IntType t <- primExpType x = TPrimExp $ UnOpExp (Abs t) x + | FloatType t <- primExpType x = TPrimExp $ UnOpExp (FAbs t) x + | otherwise = numBad "abs" x + + signum (TPrimExp x) + | IntType t <- primExpType x = TPrimExp $ UnOpExp (SSignum t) x + | otherwise = numBad "signum" x + + fromInteger = fromInteger' + +instance (FloatExp t, Pretty v) => Fractional (TPrimExp t v) where + TPrimExp x / TPrimExp y + | Just z <- msum [asFloatOp FDiv x y] = TPrimExp $ constFoldPrimExp z + | otherwise = numBad "/" (x, y) + + fromRational = fromRational' + +instance (IntExp t, Pretty v) => IntegralExp (TPrimExp t v) where + TPrimExp x `div` TPrimExp y + | Just z <- + msum + [ asIntOp (`SDiv` Unsafe) x y, + asFloatOp FDiv x y + ] = + TPrimExp $ constFoldPrimExp z + | otherwise = numBad "div" (x, y) + + TPrimExp x `mod` TPrimExp y + | Just z <- msum [asIntOp (`SMod` Unsafe) x y] = + TPrimExp z + | otherwise = numBad "mod" (x, y) + + TPrimExp x `quot` TPrimExp y + | oneIshExp y = TPrimExp x + | Just z <- msum [asIntOp (`SQuot` Unsafe) x y] = + TPrimExp $ constFoldPrimExp z + | otherwise = numBad "quot" (x, y) + + TPrimExp x `rem` TPrimExp y + | Just z <- msum [asIntOp (`SRem` Unsafe) x y] = + TPrimExp $ constFoldPrimExp z + | otherwise = numBad "rem" (x, y) + + TPrimExp x `divUp` TPrimExp y + | Just z <- msum [asIntOp (`SDivUp` Unsafe) x y] = + TPrimExp $ constFoldPrimExp z + | otherwise = numBad "divRoundingUp" (x, y) + + sgn (TPrimExp (ValueExp (IntValue i))) = Just $ signum $ valueIntegral i + sgn _ = Nothing + -- | Lifted logical conjunction. -(.&&.) :: PrimExp v -> PrimExp v -> PrimExp v -x .&&. y = constFoldPrimExp $ BinOpExp LogAnd x y +(.&&.) :: TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v +TPrimExp x .&&. TPrimExp y = TPrimExp $ constFoldPrimExp $ BinOpExp LogAnd x y -- | Lifted logical conjunction. -(.||.) :: PrimExp v -> PrimExp v -> PrimExp v -x .||. y = constFoldPrimExp $ BinOpExp LogOr x y +(.||.) :: TPrimExp Bool v -> TPrimExp Bool v -> TPrimExp Bool v +TPrimExp x .||. TPrimExp y = TPrimExp $ constFoldPrimExp $ BinOpExp LogOr x y -- | Lifted relational operators; assuming signed numbers in case of -- integers. -(.<.), (.>.), (.<=.), (.>=.), (.==.) :: PrimExp v -> PrimExp v -> PrimExp v -x .<. y = constFoldPrimExp $ - CmpOpExp cmp x y where cmp = case primExpType x of - IntType t -> CmpSlt $ t `min` primExpIntType y - FloatType t -> FCmpLt t - _ -> CmpLlt -x .<=. y = constFoldPrimExp $ - CmpOpExp cmp x y where cmp = case primExpType x of - IntType t -> CmpSle $ t `min` primExpIntType y - FloatType t -> FCmpLe t - _ -> CmpLle -x .==. y = constFoldPrimExp $ - CmpOpExp (CmpEq $ primExpType x `min` primExpType y) x y +(.<.), (.>.), (.<=.), (.>=.), (.==.) :: TPrimExp t v -> TPrimExp t v -> TPrimExp Bool v +TPrimExp x .<. TPrimExp y = + TPrimExp $ + constFoldPrimExp $ + CmpOpExp cmp x y + where + cmp = case primExpType x of + IntType t -> CmpSlt t + FloatType t -> FCmpLt t + _ -> CmpLlt +TPrimExp x .<=. TPrimExp y = + TPrimExp $ + constFoldPrimExp $ + CmpOpExp cmp x y + where + cmp = case primExpType x of + IntType t -> CmpSle t + FloatType t -> FCmpLe t + _ -> CmpLle +TPrimExp x .==. TPrimExp y = + TPrimExp $ + constFoldPrimExp $ + CmpOpExp (CmpEq $ primExpType x `min` primExpType y) x y x .>. y = y .<. x x .>=. y = y .<=. x -- | Lifted bitwise operators. -(.&.), (.|.), (.^.) :: PrimExp v -> PrimExp v -> PrimExp v -x .&. y = constFoldPrimExp $ - BinOpExp (And $ primExpIntType x `min` primExpIntType y) x y -x .|. y = constFoldPrimExp $ - BinOpExp (Or $ primExpIntType x `min` primExpIntType y) x y -x .^. y = constFoldPrimExp $ - BinOpExp (Xor $ primExpIntType x `min` primExpIntType y) x y +(.&.), (.|.), (.^.) :: TPrimExp t v -> TPrimExp t v -> TPrimExp t v +TPrimExp x .&. TPrimExp y = + TPrimExp $ + constFoldPrimExp $ + BinOpExp (And $ primExpIntType x) x y +TPrimExp x .|. TPrimExp y = + TPrimExp $ + constFoldPrimExp $ + BinOpExp (Or $ primExpIntType x) x y +TPrimExp x .^. TPrimExp y = + TPrimExp $ + constFoldPrimExp $ + BinOpExp (Xor $ primExpIntType x) x y infix 4 .==., .<., .>., .<=., .>=. + infixr 3 .&&. + infixr 2 .||. --- | Smart constructor for sign extension that does a bit of constant --- folding. +-- | Untyped smart constructor for sign extension that does a bit of +-- constant folding. sExt :: IntType -> PrimExp v -> PrimExp v sExt it (ValueExp (IntValue v)) = ValueExp $ IntValue $ doSExt v it sExt it e | primExpIntType e == it = e | otherwise = ConvOpExp (SExt (primExpIntType e) it) e --- | Smart constructor for zero extension that does a bit of constant --- folding. +-- | Untyped smart constructor for zero extension that does a bit of +-- constant folding. zExt :: IntType -> PrimExp v -> PrimExp v zExt it (ValueExp (IntValue v)) = ValueExp $ IntValue $ doZExt v it zExt it e @@ -269,50 +435,14 @@ asIntOp :: (IntType -> BinOp) -> PrimExp v -> PrimExp v -> Maybe (PrimExp v) asIntOp f x y - -- If either of the operands is a constant, then we prefer the type - -- of the other operand. This lets us use literals via fromInteger - -- without imposing a specific type. - | ValueExp{} <- x, - IntType y_t <- primExpType y, - Just x' <- asIntExp y_t x = Just $ BinOpExp (f y_t) x' y - | ValueExp{} <- y, - IntType x_t <- primExpType x, - Just y' <- asIntExp x_t y = Just $ BinOpExp (f x_t) x y' - - -- Otherwise prefer the type of the leftmost operand. - | IntType t <- primExpType x, - Just y' <- asIntExp t y = Just $ BinOpExp (f t) x y' - | IntType t <- primExpType y, - Just x' <- asIntExp t x = Just $ BinOpExp (f t) x' y - + | IntType x_t <- primExpType x = Just $ BinOpExp (f x_t) x y | otherwise = Nothing -asIntExp :: IntType -> PrimExp v -> Maybe (PrimExp v) -asIntExp t e - | primExpType e == IntType t = Just e -asIntExp t (ValueExp (IntValue v)) = - Just $ ValueExp $ IntValue $ doSExt v t -asIntExp _ _ = - Nothing - asFloatOp :: (FloatType -> BinOp) -> PrimExp v -> PrimExp v -> Maybe (PrimExp v) asFloatOp f x y - | FloatType t <- primExpType x, - Just y' <- asFloatExp t y = Just $ BinOpExp (f t) x y' - | FloatType t <- primExpType y, - Just x' <- asFloatExp t x = Just $ BinOpExp (f t) x' y + | FloatType t <- primExpType x = Just $ BinOpExp (f t) x y | otherwise = Nothing -asFloatExp :: FloatType -> PrimExp v -> Maybe (PrimExp v) -asFloatExp t e - | primExpType e == FloatType t = Just e -asFloatExp t (ValueExp (FloatValue v)) = - Just $ ValueExp $ FloatValue $ doFPConv v t -asFloatExp t (ValueExp (IntValue v)) = - Just $ ValueExp $ FloatValue $ doSIToFP v t -asFloatExp _ _ = - Nothing - numBad :: Pretty a => String -> a -> b numBad s x = error $ "Invalid argument to PrimExp method " ++ s ++ ": " ++ pretty x @@ -325,11 +455,11 @@ evalPrimExp f (BinOpExp op x y) = do x' <- evalPrimExp f x y' <- evalPrimExp f y - maybe (evalBad op (x,y)) return $ doBinOp op x' y' + maybe (evalBad op (x, y)) return $ doBinOp op x' y' evalPrimExp f (CmpOpExp op x y) = do x' <- evalPrimExp f x y' <- evalPrimExp f y - maybe (evalBad op (x,y)) (return . BoolValue) $ doCmpOp op x' y' + maybe (evalBad op (x, y)) (return . BoolValue) $ doCmpOp op x' y' evalPrimExp f (UnOpExp op x) = do x' <- evalPrimExp f x maybe (evalBad op x) return $ doUnOp op x' @@ -338,60 +468,113 @@ maybe (evalBad op x) return $ doConvOp op x' evalPrimExp f (FunExp h args _) = do args' <- mapM (evalPrimExp f) args - maybe (evalBad h args) return $ do (_, _, fun) <- M.lookup h primFuns - fun args' + maybe (evalBad h args) return $ do + (_, _, fun) <- M.lookup h primFuns + fun args' evalBad :: (Pretty a, Pretty b, MonadFail m) => a -> b -> m c -evalBad op arg = fail $ "evalPrimExp: Type error when applying " ++ - pretty op ++ " to " ++ pretty arg +evalBad op arg = + fail $ + "evalPrimExp: Type error when applying " + ++ pretty op + ++ " to " + ++ pretty arg -- | The type of values returned by a 'PrimExp'. This function -- returning does not imply that the 'PrimExp' is type-correct. primExpType :: PrimExp v -> PrimType -primExpType (LeafExp _ t) = t -primExpType (ValueExp v) = primValueType v +primExpType (LeafExp _ t) = t +primExpType (ValueExp v) = primValueType v primExpType (BinOpExp op _ _) = binOpType op -primExpType CmpOpExp{} = Bool -primExpType (UnOpExp op _) = unOpType op -primExpType (ConvOpExp op _) = snd $ convOpType op -primExpType (FunExp _ _ t) = t +primExpType CmpOpExp {} = Bool +primExpType (UnOpExp op _) = unOpType op +primExpType (ConvOpExp op _) = snd $ convOpType op +primExpType (FunExp _ _ t) = t -- | Is the expression a constant zero of some sort? zeroIshExp :: PrimExp v -> Bool zeroIshExp (ValueExp v) = zeroIsh v -zeroIshExp _ = False +zeroIshExp _ = False -- | Is the expression a constant one of some sort? oneIshExp :: PrimExp v -> Bool oneIshExp (ValueExp v) = oneIsh v -oneIshExp _ = False +oneIshExp _ = False -- | If the given 'PrimExp' is a constant of the wrong integer type, -- coerce it to the given integer type. This is a workaround for an -- issue in the 'Num' instance. coerceIntPrimExp :: IntType -> PrimExp v -> PrimExp v coerceIntPrimExp t (ValueExp (IntValue v)) = ValueExp $ IntValue $ doSExt v t -coerceIntPrimExp _ e = e +coerceIntPrimExp _ e = e primExpIntType :: PrimExp v -> IntType -primExpIntType e = case primExpType e of IntType t -> t - _ -> Int64 +primExpIntType e = case primExpType e of + IntType t -> t + _ -> Int64 -- | Boolean-valued PrimExps. -true, false :: PrimExp v -true = ValueExp $ BoolValue True -false = ValueExp $ BoolValue False +true, false :: TPrimExp Bool v +true = TPrimExp $ ValueExp $ BoolValue True +false = TPrimExp $ ValueExp $ BoolValue False +-- | Boolean negation smart constructor. +bNot :: TPrimExp Bool v -> TPrimExp Bool v +bNot = TPrimExp . UnOpExp Not . untyped + +-- | SMax on 32-bit integers. +sMax32 :: TPrimExp Int32 v -> TPrimExp Int32 v -> TPrimExp Int32 v +sMax32 x y = TPrimExp $ BinOpExp (SMax Int32) (untyped x) (untyped y) + +-- | SMin on 32-bit integers. +sMin32 :: TPrimExp Int32 v -> TPrimExp Int32 v -> TPrimExp Int32 v +sMin32 x y = TPrimExp $ BinOpExp (SMin Int32) (untyped x) (untyped y) + +-- | SMax on 64-bit integers. +sMax64 :: TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v +sMax64 x y = TPrimExp $ BinOpExp (SMax Int64) (untyped x) (untyped y) + +-- | SMin on 64-bit integers. +sMin64 :: TPrimExp Int64 v -> TPrimExp Int64 v -> TPrimExp Int64 v +sMin64 x y = TPrimExp $ BinOpExp (SMin Int64) (untyped x) (untyped y) + +-- | Sign-extend to 32 bit integer. +sExt32 :: IntExp t => TPrimExp t v -> TPrimExp Int32 v +sExt32 = isInt32 . sExt Int32 . untyped + +-- | Sign-extend to 64 bit integer. +sExt64 :: IntExp t => TPrimExp t v -> TPrimExp Int64 v +sExt64 = isInt64 . sExt Int64 . untyped + +-- | Zero-extend to 32 bit integer. +zExt32 :: IntExp t => TPrimExp t v -> TPrimExp Int32 v +zExt32 = isInt32 . zExt Int32 . untyped + +-- | Zero-extend to 64 bit integer. +zExt64 :: IntExp t => TPrimExp t v -> TPrimExp Int64 v +zExt64 = isInt64 . zExt Int64 . untyped + +-- | 64-bit float minimum. +fMin64 :: TPrimExp Double v -> TPrimExp Double v -> TPrimExp Double v +fMin64 x y = TPrimExp $ BinOpExp (FMin Float64) (untyped x) (untyped y) + +-- | 64-bit float maximum. +fMax64 :: TPrimExp Double v -> TPrimExp Double v -> TPrimExp Double v +fMax64 x y = TPrimExp $ BinOpExp (FMax Float64) (untyped x) (untyped y) + -- Prettyprinting instances instance Pretty v => Pretty (PrimExp v) where - ppr (LeafExp v _) = ppr v - ppr (ValueExp v) = ppr v + ppr (LeafExp v _) = ppr v + ppr (ValueExp v) = ppr v ppr (BinOpExp op x y) = ppr op <+> parens (ppr x) <+> parens (ppr y) ppr (CmpOpExp op x y) = ppr op <+> parens (ppr x) <+> parens (ppr y) - ppr (ConvOpExp op x) = ppr op <+> parens (ppr x) - ppr (UnOpExp op x) = ppr op <+> parens (ppr x) + ppr (ConvOpExp op x) = ppr op <+> parens (ppr x) + ppr (UnOpExp op x) = ppr op <+> parens (ppr x) ppr (FunExp h args _) = text h <+> parens (commasep $ map ppr args) + +instance Pretty v => Pretty (TPrimExp t v) where + ppr = ppr . untyped -- | Produce a mapping from the leaves of the 'PrimExp' to their -- designated types. diff --git a/src/Futhark/Analysis/PrimExp/Convert.hs b/src/Futhark/Analysis/PrimExp/Convert.hs --- a/src/Futhark/Analysis/PrimExp/Convert.hs +++ b/src/Futhark/Analysis/PrimExp/Convert.hs @@ -1,29 +1,33 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} + -- | Converting back and forth between 'PrimExp's. Use the 'ToExp' -- instance to convert to Futhark expressions. module Futhark.Analysis.PrimExp.Convert - ( - primExpFromExp - , primExpFromSubExp - , primExpFromSubExpM - , replaceInPrimExp - , replaceInPrimExpM - , substituteInPrimExp - , primExpSlice - , subExpSlice + ( primExpFromExp, + primExpFromSubExp, + pe32, + le32, + pe64, + le64, + primExpFromSubExpM, + replaceInPrimExp, + replaceInPrimExpM, + substituteInPrimExp, + primExpSlice, + subExpSlice, -- * Module reexport - , module Futhark.Analysis.PrimExp - ) where + module Futhark.Analysis.PrimExp, + ) +where import qualified Control.Monad.Fail as Fail -import Control.Monad.Identity +import Control.Monad.Identity import qualified Data.Map.Strict as M -import Data.Maybe - -import Futhark.Analysis.PrimExp -import Futhark.Construct -import Futhark.IR +import Data.Maybe +import Futhark.Analysis.PrimExp +import Futhark.Construct +import Futhark.IR instance ToExp v => ToExp (PrimExp v) where toExp (BinOpExp op x y) = @@ -37,18 +41,25 @@ toExp (ValueExp v) = return $ BasicOp $ SubExp $ Constant v toExp (FunExp h args t) = - Apply (nameFromString h) <$> args' <*> pure [primRetType t] <*> - pure (Safe, mempty, []) - where args' = zip <$> mapM (toSubExp "apply_arg") args <*> pure (repeat Observe) + Apply (nameFromString h) <$> args' <*> pure [primRetType t] + <*> pure (Safe, mempty, []) + where + args' = zip <$> mapM (toSubExp "apply_arg") args <*> pure (repeat Observe) toExp (LeafExp v _) = toExp v +instance ToExp v => ToExp (TPrimExp t v) where + toExp = toExp . untyped + -- | Convert an expression to a 'PrimExp'. The provided function is -- used to convert expressions that are not trivially 'PrimExp's. -- This includes constants and variable names, which are passed as -- t'SubExp's. -primExpFromExp :: (Fail.MonadFail m, Decorations lore) => - (VName -> m (PrimExp v)) -> Exp lore -> m (PrimExp v) +primExpFromExp :: + (Fail.MonadFail m, Decorations lore) => + (VName -> m (PrimExp v)) -> + Exp lore -> + m (PrimExp v) primExpFromExp f (BasicOp (BinOp op x y)) = BinOpExp op <$> primExpFromSubExpM f x <*> primExpFromSubExpM f y primExpFromExp f (BasicOp (CmpOp op x y)) = @@ -60,8 +71,9 @@ primExpFromExp f (BasicOp (SubExp se)) = primExpFromSubExpM f se primExpFromExp f (Apply fname args ts _) - | isBuiltInFunction fname, [Prim t] <- map declExtTypeOf ts = - FunExp (nameToString fname) <$> mapM (primExpFromSubExpM f . fst) args <*> pure t + | isBuiltInFunction fname, + [Prim t] <- map declExtTypeOf ts = + FunExp (nameToString fname) <$> mapM (primExpFromSubExpM f . fst) args <*> pure t primExpFromExp _ _ = fail "Not a PrimExp" -- | Like 'primExpFromExp', but for a t'SubExp'. @@ -71,20 +83,38 @@ -- | Convert t'SubExp's of a given type. primExpFromSubExp :: PrimType -> SubExp -> PrimExp VName -primExpFromSubExp t (Var v) = LeafExp v t +primExpFromSubExp t (Var v) = LeafExp v t primExpFromSubExp _ (Constant v) = ValueExp v +-- | Shorthand for constructing a 'TPrimExp' of type 'Int32'. +pe32 :: SubExp -> TPrimExp Int32 VName +pe32 = isInt32 . primExpFromSubExp int32 + +-- | Shorthand for constructing a 'TPrimExp' of type 'Int32', from a leaf. +le32 :: a -> TPrimExp Int32 a +le32 = isInt32 . flip LeafExp int32 + +-- | Shorthand for constructing a 'TPrimExp' of type 'Int64'. +pe64 :: SubExp -> TPrimExp Int64 VName +pe64 = isInt64 . primExpFromSubExp int64 + +-- | Shorthand for constructing a 'TPrimExp' of type 'Int64', from a leaf. +le64 :: a -> TPrimExp Int64 a +le64 = isInt64 . flip LeafExp int64 + -- | Applying a monadic transformation to the leaves in a 'PrimExp'. -replaceInPrimExpM :: Monad m => - (a -> PrimType -> m (PrimExp b)) -> - PrimExp a -> m (PrimExp b) +replaceInPrimExpM :: + Monad m => + (a -> PrimType -> m (PrimExp b)) -> + PrimExp a -> + m (PrimExp b) replaceInPrimExpM f (LeafExp v pt) = f v pt replaceInPrimExpM _ (ValueExp v) = return $ ValueExp v replaceInPrimExpM f (BinOpExp bop pe1 pe2) = - constFoldPrimExp <$> - (BinOpExp bop <$> replaceInPrimExpM f pe1 <*> replaceInPrimExpM f pe2) + constFoldPrimExp + <$> (BinOpExp bop <$> replaceInPrimExpM f pe1 <*> replaceInPrimExpM f pe2) replaceInPrimExpM f (CmpOpExp cop pe1 pe2) = CmpOpExp cop <$> replaceInPrimExpM f pe1 <*> replaceInPrimExpM f pe2 replaceInPrimExpM f (UnOpExp uop pe) = @@ -95,21 +125,27 @@ FunExp h <$> mapM (replaceInPrimExpM f) args <*> pure t -- | As 'replaceInPrimExpM', but in the identity monad. -replaceInPrimExp :: (a -> PrimType -> PrimExp b) -> - PrimExp a -> PrimExp b +replaceInPrimExp :: + (a -> PrimType -> PrimExp b) -> + PrimExp a -> + PrimExp b replaceInPrimExp f e = runIdentity $ replaceInPrimExpM f' e - where f' x y = return $ f x y + where + f' x y = return $ f x y -- | Substituting names in a PrimExp with other PrimExps -substituteInPrimExp :: Ord v => M.Map v (PrimExp v) - -> PrimExp v -> PrimExp v +substituteInPrimExp :: + Ord v => + M.Map v (PrimExp v) -> + PrimExp v -> + PrimExp v substituteInPrimExp tab = replaceInPrimExp $ \v t -> fromMaybe (LeafExp v t) $ M.lookup v tab -- | Convert a 'SubExp' slice to a 'PrimExp' slice. -primExpSlice :: Slice SubExp -> Slice (PrimExp VName) -primExpSlice = map $ fmap $ primExpFromSubExp int32 +primExpSlice :: Slice SubExp -> Slice (TPrimExp Int64 VName) +primExpSlice = map $ fmap pe64 -- | Convert a 'PrimExp' slice to a 'SubExp' slice. -subExpSlice :: MonadBinder m => Slice (PrimExp VName) -> m (Slice SubExp) +subExpSlice :: MonadBinder m => Slice (TPrimExp Int64 VName) -> m (Slice SubExp) subExpSlice = mapM $ traverse $ toSubExp "slice" diff --git a/src/Futhark/Analysis/PrimExp/Generalize.hs b/src/Futhark/Analysis/PrimExp/Generalize.hs --- a/src/Futhark/Analysis/PrimExp/Generalize.hs +++ b/src/Futhark/Analysis/PrimExp/Generalize.hs @@ -1,69 +1,73 @@ -- | Generalization (anti-unification) of 'PrimExp's. module Futhark.Analysis.PrimExp.Generalize - ( - leastGeneralGeneralization - ) where - -import Data.List (elemIndex) + ( leastGeneralGeneralization, + ) +where -import Futhark.Analysis.PrimExp -import Futhark.IR.Syntax.Core (Ext(..)) +import Data.List (elemIndex) +import Futhark.Analysis.PrimExp +import Futhark.IR.Syntax.Core (Ext (..)) -- | Generalize two 'PrimExp's of the the same type. -leastGeneralGeneralization :: (Eq v) => [(PrimExp v, PrimExp v)] -> PrimExp v -> PrimExp v -> - (PrimExp (Ext v), [(PrimExp v, PrimExp v)]) +leastGeneralGeneralization :: + (Eq v) => + [(PrimExp v, PrimExp v)] -> + PrimExp v -> + PrimExp v -> + (PrimExp (Ext v), [(PrimExp v, PrimExp v)]) leastGeneralGeneralization m exp1@(LeafExp v1 t1) exp2@(LeafExp v2 _) = - if v1 == v2 then - (LeafExp (Free v1) t1, m) - else - generalize m exp1 exp2 + if v1 == v2 + then (LeafExp (Free v1) t1, m) + else generalize m exp1 exp2 leastGeneralGeneralization m exp1@(ValueExp v1) exp2@(ValueExp v2) = - if v1 == v2 then - (ValueExp v1, m) - else - generalize m exp1 exp2 + if v1 == v2 + then (ValueExp v1, m) + else generalize m exp1 exp2 leastGeneralGeneralization m exp1@(BinOpExp op1 e11 e12) exp2@(BinOpExp op2 e21 e22) = - if op1 == op2 then - let (e1, m1) = leastGeneralGeneralization m e11 e21 - (e2, m2) = leastGeneralGeneralization m1 e12 e22 - in (BinOpExp op1 e1 e2, m2) - else - generalize m exp1 exp2 + if op1 == op2 + then + let (e1, m1) = leastGeneralGeneralization m e11 e21 + (e2, m2) = leastGeneralGeneralization m1 e12 e22 + in (BinOpExp op1 e1 e2, m2) + else generalize m exp1 exp2 leastGeneralGeneralization m exp1@(CmpOpExp op1 e11 e12) exp2@(CmpOpExp op2 e21 e22) = - if op1 == op2 then - let (e1, m1) = leastGeneralGeneralization m e11 e21 - (e2, m2) = leastGeneralGeneralization m1 e12 e22 - in (CmpOpExp op1 e1 e2, m2) - else - generalize m exp1 exp2 + if op1 == op2 + then + let (e1, m1) = leastGeneralGeneralization m e11 e21 + (e2, m2) = leastGeneralGeneralization m1 e12 e22 + in (CmpOpExp op1 e1 e2, m2) + else generalize m exp1 exp2 leastGeneralGeneralization m exp1@(UnOpExp op1 e1) exp2@(UnOpExp op2 e2) = - if op1 == op2 then - let (e, m1) = leastGeneralGeneralization m e1 e2 - in (UnOpExp op1 e, m1) - else - generalize m exp1 exp2 + if op1 == op2 + then + let (e, m1) = leastGeneralGeneralization m e1 e2 + in (UnOpExp op1 e, m1) + else generalize m exp1 exp2 leastGeneralGeneralization m exp1@(ConvOpExp op1 e1) exp2@(ConvOpExp op2 e2) = - if op1 == op2 then - let (e, m1) = leastGeneralGeneralization m e1 e2 - in (ConvOpExp op1 e, m1) - else - generalize m exp1 exp2 + if op1 == op2 + then + let (e, m1) = leastGeneralGeneralization m e1 e2 + in (ConvOpExp op1 e, m1) + else generalize m exp1 exp2 leastGeneralGeneralization m exp1@(FunExp s1 args1 t1) exp2@(FunExp s2 args2 _) = - if s1 == s2 && length args1 == length args2 then - let (args, m') = - foldl (\(arg_acc, m_acc) (a1, a2) -> - let (a, m'') = leastGeneralGeneralization m_acc a1 a2 - in (a : arg_acc, m'') - ) ([], m) (zip args1 args2) - in (FunExp s1 (reverse args) t1, m') - else - generalize m exp1 exp2 + if s1 == s2 && length args1 == length args2 + then + let (args, m') = + foldl + ( \(arg_acc, m_acc) (a1, a2) -> + let (a, m'') = leastGeneralGeneralization m_acc a1 a2 + in (a : arg_acc, m'') + ) + ([], m) + (zip args1 args2) + in (FunExp s1 (reverse args) t1, m') + else generalize m exp1 exp2 leastGeneralGeneralization m exp1 exp2 = generalize m exp1 exp2 generalize :: Eq v => [(PrimExp v, PrimExp v)] -> PrimExp v -> PrimExp v -> (PrimExp (Ext v), [(PrimExp v, PrimExp v)]) generalize m exp1 exp2 = let t = primExpType exp1 - in case elemIndex (exp1, exp2) m of - Just i -> (LeafExp (Ext i) t, m) - Nothing -> (LeafExp (Ext $ length m) t, m ++ [(exp1, exp2)]) + in case elemIndex (exp1, exp2) m of + Just i -> (LeafExp (Ext i) t, m) + Nothing -> (LeafExp (Ext $ length m) t, m ++ [(exp1, exp2)]) diff --git a/src/Futhark/Analysis/PrimExp/Simplify.hs b/src/Futhark/Analysis/PrimExp/Simplify.hs --- a/src/Futhark/Analysis/PrimExp/Simplify.hs +++ b/src/Futhark/Analysis/PrimExp/Simplify.hs @@ -1,39 +1,46 @@ {-# LANGUAGE FlexibleContexts #-} + -- | Defines simplification functions for 'PrimExp's. -module Futhark.Analysis.PrimExp.Simplify - (simplifyPrimExp, simplifyExtPrimExp) -where +module Futhark.Analysis.PrimExp.Simplify (simplifyPrimExp, simplifyExtPrimExp) where -import Futhark.Analysis.PrimExp -import Futhark.Optimise.Simplify.Engine as Engine -import Futhark.IR +import Futhark.Analysis.PrimExp +import Futhark.IR +import Futhark.Optimise.Simplify.Engine as Engine -- | Simplify a 'PrimExp', including copy propagation. If a 'LeafExp' -- refers to a name that is a 'Constant', the node turns into a -- 'ValueExp'. -simplifyPrimExp :: SimplifiableLore lore => - PrimExp VName -> SimpleM lore (PrimExp VName) +simplifyPrimExp :: + SimplifiableLore lore => + PrimExp VName -> + SimpleM lore (PrimExp VName) simplifyPrimExp = simplifyAnyPrimExp onLeaf - where onLeaf v pt = do - se <- simplify $ Var v - case se of - Var v' -> return $ LeafExp v' pt - Constant pv -> return $ ValueExp pv + where + onLeaf v pt = do + se <- simplify $ Var v + case se of + Var v' -> return $ LeafExp v' pt + Constant pv -> return $ ValueExp pv -- | Like 'simplifyPrimExp', but where leaves may be 'Ext's. -simplifyExtPrimExp :: SimplifiableLore lore => - PrimExp (Ext VName) -> SimpleM lore (PrimExp (Ext VName)) +simplifyExtPrimExp :: + SimplifiableLore lore => + PrimExp (Ext VName) -> + SimpleM lore (PrimExp (Ext VName)) simplifyExtPrimExp = simplifyAnyPrimExp onLeaf - where onLeaf (Free v) pt = do - se <- simplify $ Var v - case se of - Var v' -> return $ LeafExp (Free v') pt - Constant pv -> return $ ValueExp pv - onLeaf (Ext i) pt = return $ LeafExp (Ext i) pt + where + onLeaf (Free v) pt = do + se <- simplify $ Var v + case se of + Var v' -> return $ LeafExp (Free v') pt + Constant pv -> return $ ValueExp pv + onLeaf (Ext i) pt = return $ LeafExp (Ext i) pt -simplifyAnyPrimExp :: SimplifiableLore lore => - (a -> PrimType -> SimpleM lore (PrimExp a)) - -> PrimExp a -> SimpleM lore (PrimExp a) +simplifyAnyPrimExp :: + SimplifiableLore lore => + (a -> PrimType -> SimpleM lore (PrimExp a)) -> + PrimExp a -> + SimpleM lore (PrimExp a) simplifyAnyPrimExp f (LeafExp v pt) = f v pt simplifyAnyPrimExp _ (ValueExp pv) = return $ ValueExp pv diff --git a/src/Futhark/Analysis/Rephrase.hs b/src/Futhark/Analysis/Rephrase.hs --- a/src/Futhark/Analysis/Rephrase.hs +++ b/src/Futhark/Analysis/Rephrase.hs @@ -1,18 +1,19 @@ -{-# LANGUAGE GADTs #-} {-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE GADTs #-} + -- | Facilities for changing the lore of some fragment, with no -- context. We call this "rephrasing", for no deep reason. module Futhark.Analysis.Rephrase - ( rephraseProg - , rephraseFunDef - , rephraseExp - , rephraseBody - , rephraseStm - , rephraseLambda - , rephrasePattern - , rephrasePatElem - , Rephraser (..) - ) + ( rephraseProg, + rephraseFunDef, + rephraseExp, + rephraseBody, + rephraseStm, + rephraseLambda, + rephrasePattern, + rephrasePatElem, + Rephraser (..), + ) where import Futhark.IR @@ -22,23 +23,23 @@ -- monad, we can conveniently do rephrasing that might fail. This is -- useful if you want to see if some IR in e.g. the @Kernels@ lore -- actually uses any @Kernels@-specific operations. -data Rephraser m from to - = Rephraser { rephraseExpLore :: ExpDec from -> m (ExpDec to) - , rephraseLetBoundLore :: LetDec from -> m (LetDec to) - , rephraseFParamLore :: FParamInfo from -> m (FParamInfo to) - , rephraseLParamLore :: LParamInfo from -> m (LParamInfo to) - , rephraseBodyLore :: BodyDec from -> m (BodyDec to) - , rephraseRetType :: RetType from -> m (RetType to) - , rephraseBranchType :: BranchType from -> m (BranchType to) - , rephraseOp :: Op from -> m (Op to) - } +data Rephraser m from to = Rephraser + { rephraseExpLore :: ExpDec from -> m (ExpDec to), + rephraseLetBoundLore :: LetDec from -> m (LetDec to), + rephraseFParamLore :: FParamInfo from -> m (FParamInfo to), + rephraseLParamLore :: LParamInfo from -> m (LParamInfo to), + rephraseBodyLore :: BodyDec from -> m (BodyDec to), + rephraseRetType :: RetType from -> m (RetType to), + rephraseBranchType :: BranchType from -> m (BranchType to), + rephraseOp :: Op from -> m (Op to) + } -- | Rephrase an entire program. rephraseProg :: Monad m => Rephraser m from to -> Prog from -> m (Prog to) rephraseProg rephraser (Prog consts funs) = Prog - <$> mapM (rephraseStm rephraser) consts - <*> mapM (rephraseFunDef rephraser) funs + <$> mapM (rephraseStm rephraser) consts + <*> mapM (rephraseFunDef rephraser) funs -- | Rephrase a function definition. rephraseFunDef :: Monad m => Rephraser m from to -> FunDef from -> m (FunDef to) @@ -46,7 +47,7 @@ body' <- rephraseBody rephraser $ funDefBody fundec params' <- mapM (rephraseParam $ rephraseFParamLore rephraser) $ funDefParams fundec rettype' <- mapM (rephraseRetType rephraser) $ funDefRetType fundec - return fundec { funDefBody = body', funDefParams = params', funDefRetType = rettype' } + return fundec {funDefBody = body', funDefParams = params', funDefRetType = rettype'} -- | Rephrase an expression. rephraseExp :: Monad m => Rephraser m from to -> Exp from -> m (Exp to) @@ -55,16 +56,17 @@ -- | Rephrase a statement. rephraseStm :: Monad m => Rephraser m from to -> Stm from -> m (Stm to) rephraseStm rephraser (Let pat (StmAux cs attrs dec) e) = - Let <$> - rephrasePattern (rephraseLetBoundLore rephraser) pat <*> - (StmAux cs attrs <$> rephraseExpLore rephraser dec) <*> - rephraseExp rephraser e + Let + <$> rephrasePattern (rephraseLetBoundLore rephraser) pat + <*> (StmAux cs attrs <$> rephraseExpLore rephraser dec) + <*> rephraseExp rephraser e -- | Rephrase a pattern. -rephrasePattern :: Monad m => - (from -> m to) - -> PatternT from - -> m (PatternT to) +rephrasePattern :: + Monad m => + (from -> m to) -> + PatternT from -> + m (PatternT to) rephrasePattern = traverse -- | Rephrase a pattern element. @@ -80,24 +82,25 @@ -- | Rephrase a body. rephraseBody :: Monad m => Rephraser m from to -> Body from -> m (Body to) rephraseBody rephraser (Body lore bnds res) = - Body <$> - rephraseBodyLore rephraser lore <*> - (stmsFromList <$> mapM (rephraseStm rephraser) (stmsToList bnds)) <*> - pure res + Body + <$> rephraseBodyLore rephraser lore + <*> (stmsFromList <$> mapM (rephraseStm rephraser) (stmsToList bnds)) + <*> pure res -- | Rephrase a lambda. rephraseLambda :: Monad m => Rephraser m from to -> Lambda from -> m (Lambda to) rephraseLambda rephraser lam = do body' <- rephraseBody rephraser $ lambdaBody lam params' <- mapM (rephraseParam $ rephraseLParamLore rephraser) $ lambdaParams lam - return lam { lambdaBody = body', lambdaParams = params' } + return lam {lambdaBody = body', lambdaParams = params'} mapper :: Monad m => Rephraser m from to -> Mapper from to m -mapper rephraser = identityMapper { - mapOnBody = const $ rephraseBody rephraser - , mapOnRetType = rephraseRetType rephraser - , mapOnBranchType = rephraseBranchType rephraser - , mapOnFParam = rephraseParam (rephraseFParamLore rephraser) - , mapOnLParam = rephraseParam (rephraseLParamLore rephraser) - , mapOnOp = rephraseOp rephraser - } +mapper rephraser = + identityMapper + { mapOnBody = const $ rephraseBody rephraser, + mapOnRetType = rephraseRetType rephraser, + mapOnBranchType = rephraseBranchType rephraser, + mapOnFParam = rephraseParam (rephraseFParamLore rephraser), + mapOnLParam = rephraseParam (rephraseLParamLore rephraser), + mapOnOp = rephraseOp rephraser + } diff --git a/src/Futhark/Analysis/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs --- a/src/Futhark/Analysis/SymbolTable.hs +++ b/src/Futhark/Analysis/SymbolTable.hs @@ -1,85 +1,87 @@ -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + module Futhark.Analysis.SymbolTable - ( SymbolTable (bindings, loopDepth, availableAtClosestLoop, simplifyMemory) - , empty - , fromScope - , toScope + ( SymbolTable (bindings, loopDepth, availableAtClosestLoop, simplifyMemory), + empty, + fromScope, + toScope, -- * Entries - , Entry - , deepen - , entryDepth - , entryLetBoundDec - , entryIsSize + Entry, + deepen, + entryDepth, + entryLetBoundDec, + entryIsSize, -- * Lookup - , elem - , lookup - , lookupStm - , lookupExp - , lookupBasicOp - , lookupType - , lookupSubExp - , lookupAliases - , lookupLoopVar - , available - , consume - , index - , index' - , Indexed(..) - , indexedAddCerts - , IndexOp(..) + elem, + lookup, + lookupStm, + lookupExp, + lookupBasicOp, + lookupType, + lookupSubExp, + lookupAliases, + lookupLoopVar, + lookupLoopParam, + available, + consume, + index, + index', + Indexed (..), + indexedAddCerts, + IndexOp (..), -- * Insertion - , insertStm - , insertStms - , insertFParams - , insertLParam - , insertLoopVar + insertStm, + insertStms, + insertFParams, + insertLParam, + insertLoopVar, + insertLoopMerge, -- * Misc - , hideCertified + hideCertified, ) - where +where import Control.Arrow ((&&&)) import Control.Monad -import Data.Ord -import Data.Maybe -import Data.List (foldl', elemIndex) +import Data.List (elemIndex, foldl') import qualified Data.Map.Strict as M - -import Prelude hiding (elem, lookup) - +import Data.Maybe +import Data.Ord import Futhark.Analysis.PrimExp.Convert import Futhark.IR hiding (FParam, lookupType) import qualified Futhark.IR as AST - import qualified Futhark.IR.Prop.Aliases as Aliases +import Prelude hiding (elem, lookup) -data SymbolTable lore = SymbolTable { - loopDepth :: Int - , bindings :: M.Map VName (Entry lore) - , availableAtClosestLoop :: Names - -- ^ Which names are available just before the most enclosing +data SymbolTable lore = SymbolTable + { loopDepth :: Int, + bindings :: M.Map VName (Entry lore), + -- | Which names are available just before the most enclosing -- loop? - , simplifyMemory :: Bool - -- ^ We are in a situation where we should + availableAtClosestLoop :: Names, + -- | We are in a situation where we should -- simplify/hoist/un-existentialise memory as much as possible - -- typically, inside a kernel. + simplifyMemory :: Bool } instance Semigroup (SymbolTable lore) where table1 <> table2 = - SymbolTable { loopDepth = max (loopDepth table1) (loopDepth table2) - , bindings = bindings table1 <> bindings table2 - , availableAtClosestLoop = availableAtClosestLoop table1 <> - availableAtClosestLoop table2 - , simplifyMemory = simplifyMemory table1 || simplifyMemory table2 - } + SymbolTable + { loopDepth = max (loopDepth table1) (loopDepth table2), + bindings = bindings table1 <> bindings table2, + availableAtClosestLoop = + availableAtClosestLoop table1 + <> availableAtClosestLoop table2, + simplifyMemory = simplifyMemory table1 || simplifyMemory table2 + } instance Monoid (SymbolTable lore) where mempty = empty @@ -89,95 +91,103 @@ fromScope :: ASTLore lore => Scope lore -> SymbolTable lore fromScope = M.foldlWithKey' insertFreeVar' empty - where insertFreeVar' m k dec = insertFreeVar k dec m + where + insertFreeVar' m k dec = insertFreeVar k dec m toScope :: SymbolTable lore -> Scope lore toScope = M.map entryInfo . bindings deepen :: SymbolTable lore -> SymbolTable lore -deepen vtable = vtable { loopDepth = loopDepth vtable + 1, - availableAtClosestLoop = namesFromList $ M.keys $ bindings vtable - } +deepen vtable = + vtable + { loopDepth = loopDepth vtable + 1, + availableAtClosestLoop = namesFromList $ M.keys $ bindings vtable + } -- | The result of indexing a delayed array. -data Indexed = Indexed Certificates (PrimExp VName) - -- ^ A PrimExp based on the indexes (that is, without - -- accessing any actual array). - | IndexedArray Certificates VName [PrimExp VName] - -- ^ The indexing corresponds to another (perhaps more - -- advantageous) array. +data Indexed + = -- | A PrimExp based on the indexes (that is, without + -- accessing any actual array). + Indexed Certificates (PrimExp VName) + | -- | The indexing corresponds to another (perhaps more + -- advantageous) array. + IndexedArray Certificates VName [TPrimExp Int64 VName] indexedAddCerts :: Certificates -> Indexed -> Indexed -indexedAddCerts cs1 (Indexed cs2 v) = Indexed (cs1<>cs2) v -indexedAddCerts cs1 (IndexedArray cs2 arr v) = IndexedArray (cs1<>cs2) arr v +indexedAddCerts cs1 (Indexed cs2 v) = Indexed (cs1 <> cs2) v +indexedAddCerts cs1 (IndexedArray cs2 arr v) = IndexedArray (cs1 <> cs2) arr v instance FreeIn Indexed where freeIn' (Indexed cs v) = freeIn' cs <> freeIn' v freeIn' (IndexedArray cs arr v) = freeIn' cs <> freeIn' arr <> freeIn' v -- | Indexing a delayed array if possible. -type IndexArray = [PrimExp VName] -> Maybe Indexed +type IndexArray = [TPrimExp Int64 VName] -> Maybe Indexed -data Entry lore = - Entry { entryConsumed :: Bool - -- ^ True if consumed. - , entryDepth :: Int - , entryIsSize :: Bool - -- ^ True if this name has been used as an array size, - -- implying that it is non-negative. - , entryType :: EntryType lore - } +data Entry lore = Entry + { -- | True if consumed. + entryConsumed :: Bool, + entryDepth :: Int, + -- | True if this name has been used as an array size, + -- implying that it is non-negative. + entryIsSize :: Bool, + entryType :: EntryType lore + } -data EntryType lore = LoopVar (LoopVarEntry lore) - | LetBound (LetBoundEntry lore) - | FParam (FParamEntry lore) - | LParam (LParamEntry lore) - | FreeVar (FreeVarEntry lore) +data EntryType lore + = LoopVar (LoopVarEntry lore) + | LetBound (LetBoundEntry lore) + | FParam (FParamEntry lore) + | LParam (LParamEntry lore) + | FreeVar (FreeVarEntry lore) -data LoopVarEntry lore = - LoopVarEntry { loopVarType :: IntType - , loopVarBound :: SubExp - } +data LoopVarEntry lore = LoopVarEntry + { loopVarType :: IntType, + loopVarBound :: SubExp + } -data LetBoundEntry lore = - LetBoundEntry { letBoundDec :: LetDec lore - , letBoundAliases :: Names - , letBoundStm :: Stm lore - , letBoundIndex :: Int -> IndexArray - -- ^ Index a delayed array, if possible. - } +data LetBoundEntry lore = LetBoundEntry + { letBoundDec :: LetDec lore, + letBoundAliases :: Names, + letBoundStm :: Stm lore, + -- | Index a delayed array, if possible. + letBoundIndex :: Int -> IndexArray + } -data FParamEntry lore = - FParamEntry { fparamDec :: FParamInfo lore - , fparamAliases :: Names - } +data FParamEntry lore = FParamEntry + { fparamDec :: FParamInfo lore, + fparamAliases :: Names, + -- | If a loop parameter, the initial value and the eventual + -- result. The result need not be in scope in the symbol table. + fparamMerge :: Maybe (SubExp, SubExp) + } -data LParamEntry lore = - LParamEntry { lparamDec :: LParamInfo lore - , lparamIndex :: IndexArray - } +data LParamEntry lore = LParamEntry + { lparamDec :: LParamInfo lore, + lparamIndex :: IndexArray + } -data FreeVarEntry lore = - FreeVarEntry { freeVarDec :: NameInfo lore - , freeVarIndex :: VName -> IndexArray - -- ^ Index a delayed array, if possible. - } +data FreeVarEntry lore = FreeVarEntry + { freeVarDec :: NameInfo lore, + -- | Index a delayed array, if possible. + freeVarIndex :: VName -> IndexArray + } instance ASTLore lore => Typed (Entry lore) where typeOf = typeOf . entryInfo entryInfo :: Entry lore -> NameInfo lore entryInfo e = case entryType e of - LetBound entry -> LetName $ letBoundDec entry - LoopVar entry -> IndexName $ loopVarType entry - FParam entry -> FParamName $ fparamDec entry - LParam entry -> LParamName $ lparamDec entry - FreeVar entry -> freeVarDec entry + LetBound entry -> LetName $ letBoundDec entry + LoopVar entry -> IndexName $ loopVarType entry + FParam entry -> FParamName $ fparamDec entry + LParam entry -> LParamName $ lparamDec entry + FreeVar entry -> freeVarDec entry isLetBound :: Entry lore -> Maybe (LetBoundEntry lore) isLetBound e = case entryType e of - LetBound entry -> Just entry - _ -> Nothing + LetBound entry -> Just entry + _ -> Nothing entryStm :: Entry lore -> Maybe (Stm lore) entryStm = fmap letBoundStm . isLetBound @@ -200,7 +210,7 @@ lookupBasicOp :: VName -> SymbolTable lore -> Maybe (BasicOp, Certificates) lookupBasicOp name vtable = case lookupExp name vtable of Just (BasicOp e, cs) -> Just (e, cs) - _ -> Nothing + _ -> Nothing lookupType :: ASTLore lore => VName -> SymbolTable lore -> Maybe Type lookupType name vtable = typeOf <$> lookup name vtable @@ -211,17 +221,17 @@ lookupSubExp :: VName -> SymbolTable lore -> Maybe (SubExp, Certificates) lookupSubExp name vtable = do - (e,cs) <- lookupExp name vtable + (e, cs) <- lookupExp name vtable case e of - BasicOp (SubExp se) -> Just (se,cs) - _ -> Nothing + BasicOp (SubExp se) -> Just (se, cs) + _ -> Nothing lookupAliases :: VName -> SymbolTable lore -> Names lookupAliases name vtable = case entryType <$> M.lookup name (bindings vtable) of Just (LetBound e) -> letBoundAliases e - Just (FParam e) -> fparamAliases e - _ -> mempty + Just (FParam e) -> fparamAliases e + _ -> mempty -- | If the given variable name is the name of a 'ForLoop' parameter, -- then return the bound of that loop. @@ -230,27 +240,42 @@ LoopVar e <- entryType <$> M.lookup name (bindings vtable) return $ loopVarBound e +lookupLoopParam :: VName -> SymbolTable lore -> Maybe (SubExp, SubExp) +lookupLoopParam name vtable = do + FParam e <- entryType <$> M.lookup name (bindings vtable) + fparamMerge e + -- | In symbol table and not consumed. available :: VName -> SymbolTable lore -> Bool available name = maybe False (not . entryConsumed) . M.lookup name . bindings -index :: ASTLore lore => VName -> [SubExp] -> SymbolTable lore - -> Maybe Indexed +index :: + ASTLore lore => + VName -> + [SubExp] -> + SymbolTable lore -> + Maybe Indexed index name is table = do is' <- mapM asPrimExp is index' name is' table - where asPrimExp i = do - Prim t <- lookupSubExpType i table - return $ primExpFromSubExp t i + where + asPrimExp i = do + Prim t <- lookupSubExpType i table + return $ TPrimExp $ primExpFromSubExp t i -index' :: VName -> [PrimExp VName] -> SymbolTable lore - -> Maybe Indexed +index' :: + VName -> + [TPrimExp Int64 VName] -> + SymbolTable lore -> + Maybe Indexed index' name is vtable = do entry <- lookup name vtable case entryType entry of - LetBound entry' | - Just k <- elemIndex name $ patternValueNames $ - stmPattern $ letBoundStm entry' -> + LetBound entry' + | Just k <- + elemIndex name $ + patternValueNames $ + stmPattern $ letBoundStm entry' -> letBoundIndex entry' k is FreeVar entry' -> freeVarIndex entry' name is @@ -258,73 +283,83 @@ _ -> Nothing class IndexOp op where - indexOp :: (ASTLore lore, IndexOp (Op lore)) => - SymbolTable lore -> Int -> op - -> [PrimExp VName] -> Maybe Indexed + indexOp :: + (ASTLore lore, IndexOp (Op lore)) => + SymbolTable lore -> + Int -> + op -> + [TPrimExp Int64 VName] -> + Maybe Indexed indexOp _ _ _ _ = Nothing -instance IndexOp () where - -indexExp :: (IndexOp (Op lore), ASTLore lore) => - SymbolTable lore -> Exp lore -> Int -> IndexArray +instance IndexOp () +indexExp :: + (IndexOp (Op lore), ASTLore lore) => + SymbolTable lore -> + Exp lore -> + Int -> + IndexArray indexExp vtable (Op op) k is = indexOp vtable k op is - indexExp _ (BasicOp (Iota _ x s to_it)) _ [i] = - Just $ Indexed mempty $ - sExt to_it i - * primExpFromSubExp (IntType to_it) s - + primExpFromSubExp (IntType to_it) x - + Just $ + Indexed mempty $ + ( sExt to_it (untyped i) + `mul` primExpFromSubExp (IntType to_it) s + ) + `add` primExpFromSubExp (IntType to_it) x + where + mul = BinOpExp (Mul to_it OverflowWrap) + add = BinOpExp (Add to_it OverflowWrap) indexExp table (BasicOp (Replicate (Shape ds) v)) _ is | length ds == length is, Just (Prim t) <- lookupSubExpType v table = - Just $ Indexed mempty $ primExpFromSubExp t v - -indexExp table (BasicOp (Replicate (Shape [_]) (Var v))) _ (_:is) = + Just $ Indexed mempty $ primExpFromSubExp t v +indexExp table (BasicOp (Replicate (Shape [_]) (Var v))) _ (_ : is) = index' v is table - indexExp table (BasicOp (Reshape newshape v)) _ is | Just oldshape <- arrayDims <$> lookupType v table = - let is' = - reshapeIndex (map (primExpFromSubExp int32) oldshape) - (map (primExpFromSubExp int32) $ newDims newshape) - is - in index' v is' table - + let is' = + reshapeIndex + (map pe64 oldshape) + (map pe64 $ newDims newshape) + is + in index' v is' table indexExp table (BasicOp (Index v slice)) _ is = index' v (adjust slice is) table - where adjust (DimFix j:js') is' = - pe j : adjust js' is' - adjust (DimSlice j _ s:js') (i:is') = - let i_t_s = i * pe s - j_p_i_t_s = pe j + i_t_s - in j_p_i_t_s : adjust js' is' - adjust _ _ = [] - - pe = primExpFromSubExp (IntType Int32) - + where + adjust (DimFix j : js') is' = + pe64 j : adjust js' is' + adjust (DimSlice j _ s : js') (i : is') = + let i_t_s = i * pe64 s + j_p_i_t_s = pe64 j + i_t_s + in j_p_i_t_s : adjust js' is' + adjust _ _ = [] indexExp _ _ _ _ = Nothing -defBndEntry :: (ASTLore lore, IndexOp (Op lore)) => - SymbolTable lore - -> PatElem lore - -> Names - -> Stm lore - -> LetBoundEntry lore +defBndEntry :: + (ASTLore lore, IndexOp (Op lore)) => + SymbolTable lore -> + PatElem lore -> + Names -> + Stm lore -> + LetBoundEntry lore defBndEntry vtable patElem als bnd = - LetBoundEntry { - letBoundDec = patElemDec patElem - , letBoundAliases = als - , letBoundStm = bnd - , letBoundIndex = \k -> fmap (indexedAddCerts (stmAuxCerts $ stmAux bnd)) . - indexExp vtable (stmExp bnd) k + LetBoundEntry + { letBoundDec = patElemDec patElem, + letBoundAliases = als, + letBoundStm = bnd, + letBoundIndex = \k -> + fmap (indexedAddCerts (stmAuxCerts $ stmAux bnd)) + . indexExp vtable (stmExp bnd) k } -bindingEntries :: (ASTLore lore, Aliases.Aliased lore, IndexOp (Op lore)) => - Stm lore -> SymbolTable lore - -> [LetBoundEntry lore] +bindingEntries :: + (ASTLore lore, Aliases.Aliased lore, IndexOp (Op lore)) => + Stm lore -> + SymbolTable lore -> + [LetBoundEntry lore] bindingEntries bnd@(Let pat _ _) vtable = do pat_elem <- patternElements pat return $ defBndEntry vtable pat_elem (Aliases.aliasesOf pat_elem) bnd @@ -332,111 +367,186 @@ adjustSeveral :: Ord k => (v -> v) -> [k] -> M.Map k v -> M.Map k v adjustSeveral f = flip $ foldl' $ flip $ M.adjust f -insertEntry :: ASTLore lore => - VName -> EntryType lore -> SymbolTable lore - -> SymbolTable lore +insertEntry :: + ASTLore lore => + VName -> + EntryType lore -> + SymbolTable lore -> + SymbolTable lore insertEntry name entry vtable = - let entry' = Entry { entryConsumed = False - , entryDepth = loopDepth vtable - , entryIsSize = False - , entryType = entry - } + let entry' = + Entry + { entryConsumed = False, + entryDepth = loopDepth vtable, + entryIsSize = False, + entryType = entry + } dims = mapMaybe subExpVar $ arrayDims $ typeOf entry' - isSize e = e { entryIsSize = True } - in vtable { bindings = adjustSeveral isSize dims $ - M.insert name entry' $ bindings vtable } + isSize e = e {entryIsSize = True} + in vtable + { bindings = + adjustSeveral isSize dims $ + M.insert name entry' $ bindings vtable + } -insertEntries :: ASTLore lore => - [(VName, EntryType lore)] -> SymbolTable lore - -> SymbolTable lore +insertEntries :: + ASTLore lore => + [(VName, EntryType lore)] -> + SymbolTable lore -> + SymbolTable lore insertEntries entries vtable = foldl' add vtable entries - where add vtable' (name, entry) = insertEntry name entry vtable' + where + add vtable' (name, entry) = insertEntry name entry vtable' -insertStm :: (ASTLore lore, IndexOp (Op lore), Aliases.Aliased lore) => - Stm lore - -> SymbolTable lore - -> SymbolTable lore +insertStm :: + (ASTLore lore, IndexOp (Op lore), Aliases.Aliased lore) => + Stm lore -> + SymbolTable lore -> + SymbolTable lore insertStm stm vtable = flip (foldl' $ flip consume) (namesToList stm_consumed) $ - flip (foldl' addRevAliases) (patternElements $ stmPattern stm) $ - insertEntries (zip names $ map LetBound $ bindingEntries stm vtable) vtable - where names = patternNames $ stmPattern stm - stm_consumed = expandAliases (Aliases.consumedInStm stm) vtable - addRevAliases vtable' pe = - vtable' { bindings = adjustSeveral update inedges $ bindings vtable' } - where inedges = namesToList $ expandAliases (Aliases.aliasesOf pe) vtable' - update e = e { entryType = update' $ entryType e } - update' (LetBound entry) = - LetBound entry - { letBoundAliases = oneName (patElemName pe) <> letBoundAliases entry } - update' (FParam entry) = - FParam entry - { fparamAliases = oneName (patElemName pe) <> fparamAliases entry } - update' e = e + flip (foldl' addRevAliases) (patternElements $ stmPattern stm) $ + insertEntries (zip names $ map LetBound $ bindingEntries stm vtable) vtable + where + names = patternNames $ stmPattern stm + stm_consumed = expandAliases (Aliases.consumedInStm stm) vtable + addRevAliases vtable' pe = + vtable' {bindings = adjustSeveral update inedges $ bindings vtable'} + where + inedges = namesToList $ expandAliases (Aliases.aliasesOf pe) vtable' + update e = e {entryType = update' $ entryType e} + update' (LetBound entry) = + LetBound + entry + { letBoundAliases = oneName (patElemName pe) <> letBoundAliases entry + } + update' (FParam entry) = + FParam + entry + { fparamAliases = oneName (patElemName pe) <> fparamAliases entry + } + update' e = e -insertStms :: (ASTLore lore, IndexOp (Op lore), Aliases.Aliased lore) => - Stms lore -> SymbolTable lore -> SymbolTable lore +insertStms :: + (ASTLore lore, IndexOp (Op lore), Aliases.Aliased lore) => + Stms lore -> + SymbolTable lore -> + SymbolTable lore insertStms stms vtable = foldl' (flip insertStm) vtable $ stmsToList stms expandAliases :: Names -> SymbolTable lore -> Names expandAliases names vtable = names <> aliasesOfAliases - where aliasesOfAliases = - mconcat . map (`lookupAliases` vtable) . namesToList $ names + where + aliasesOfAliases = + mconcat . map (`lookupAliases` vtable) . namesToList $ names -insertFParam :: ASTLore lore => - AST.FParam lore -> SymbolTable lore -> SymbolTable lore +insertFParam :: + ASTLore lore => + AST.FParam lore -> + SymbolTable lore -> + SymbolTable lore insertFParam fparam = insertEntry name entry - where name = AST.paramName fparam - entry = FParam FParamEntry { fparamDec = AST.paramDec fparam - , fparamAliases = mempty - } + where + name = AST.paramName fparam + entry = + FParam + FParamEntry + { fparamDec = AST.paramDec fparam, + fparamAliases = mempty, + fparamMerge = Nothing + } -insertFParams :: ASTLore lore => - [AST.FParam lore] -> SymbolTable lore -> SymbolTable lore +insertFParams :: + ASTLore lore => + [AST.FParam lore] -> + SymbolTable lore -> + SymbolTable lore insertFParams fparams symtable = foldl' (flip insertFParam) symtable fparams insertLParam :: ASTLore lore => LParam lore -> SymbolTable lore -> SymbolTable lore insertLParam param = insertEntry name bind - where bind = LParam LParamEntry { lparamDec = AST.paramDec param - , lparamIndex = const Nothing - } - name = AST.paramName param + where + bind = + LParam + LParamEntry + { lparamDec = AST.paramDec param, + lparamIndex = const Nothing + } + name = AST.paramName param +-- | Insert entries corresponding to the parameters of a loop (not +-- distinguishing contect and value part). Apart from the parameter +-- itself, we also insert the initial value and the subexpression +-- providing the final value. Note that the latter is likely not in +-- scope in the symbol at this point. This is OK, and can still be +-- used to help some loop optimisations detect invariant loop +-- parameters. +insertLoopMerge :: + ASTLore lore => + [(AST.FParam lore, SubExp, SubExp)] -> + SymbolTable lore -> + SymbolTable lore +insertLoopMerge = flip $ foldl' $ flip bind + where + bind (p, initial, res) = + insertEntry (paramName p) $ + FParam + FParamEntry + { fparamDec = AST.paramDec p, + fparamAliases = mempty, + fparamMerge = Just (initial, res) + } + insertLoopVar :: ASTLore lore => VName -> IntType -> SubExp -> SymbolTable lore -> SymbolTable lore insertLoopVar name it bound = insertEntry name bind - where bind = LoopVar LoopVarEntry { - loopVarType = it - , loopVarBound = bound + where + bind = + LoopVar + LoopVarEntry + { loopVarType = it, + loopVarBound = bound } insertFreeVar :: ASTLore lore => VName -> NameInfo lore -> SymbolTable lore -> SymbolTable lore insertFreeVar name dec = insertEntry name entry - where entry = FreeVar FreeVarEntry { - freeVarDec = dec - , freeVarIndex = \_ _ -> Nothing + where + entry = + FreeVar + FreeVarEntry + { freeVarDec = dec, + freeVarIndex = \_ _ -> Nothing } consume :: VName -> SymbolTable lore -> SymbolTable lore -consume consumee vtable = foldl' consume' vtable $ namesToList $ - expandAliases (oneName consumee) vtable - where consume' vtable' v = - vtable' { bindings = M.adjust consume'' v $ bindings vtable' } - consume'' e = e { entryConsumed = True } +consume consumee vtable = + foldl' consume' vtable $ + namesToList $ + expandAliases (oneName consumee) vtable + where + consume' vtable' v = + vtable' {bindings = M.adjust consume'' v $ bindings vtable'} + consume'' e = e {entryConsumed = True} -- | Hide definitions of those entries that satisfy some predicate. hideIf :: (Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore -hideIf hide vtable = vtable { bindings = M.map maybeHide $ bindings vtable } - where maybeHide entry - | hide entry = entry { entryType = - FreeVar FreeVarEntry { freeVarDec = entryInfo entry - , freeVarIndex = \_ _ -> Nothing - } - } - | otherwise = entry +hideIf hide vtable = vtable {bindings = M.map maybeHide $ bindings vtable} + where + maybeHide entry + | hide entry = + entry + { entryType = + FreeVar + FreeVarEntry + { freeVarDec = entryInfo entry, + freeVarIndex = \_ _ -> Nothing + } + } + | otherwise = entry -- | Hide these definitions, if they are protected by certificates in -- the set of names. hideCertified :: Names -> SymbolTable lore -> SymbolTable lore hideCertified to_hide = hideIf $ maybe False hide . entryStm - where hide = any (`nameIn` to_hide) . unCertificates . stmCerts + where + hide = any (`nameIn` to_hide) . unCertificates . stmCerts diff --git a/src/Futhark/Analysis/UsageTable.hs b/src/Futhark/Analysis/UsageTable.hs --- a/src/Futhark/Analysis/UsageTable.hs +++ b/src/Futhark/Analysis/UsageTable.hs @@ -1,40 +1,39 @@ {-# LANGUAGE Strict #-} + -- | A usage-table is sort of a bottom-up symbol table, describing how -- (and if) a variable is used. module Futhark.Analysis.UsageTable - ( UsageTable - , without - , lookup - , used - , expand - , isConsumed - , isInResult - , isUsedDirectly - , isSize - , usages - , usage - , consumedUsage - , inResultUsage - , sizeUsage - , sizeUsages - , Usages - , usageInStm + ( UsageTable, + without, + lookup, + used, + expand, + isConsumed, + isInResult, + isUsedDirectly, + isSize, + usages, + usage, + consumedUsage, + inResultUsage, + sizeUsage, + sizeUsages, + Usages, + usageInStm, ) - where +where import Data.Bits import qualified Data.Foldable as Foldable import qualified Data.IntMap.Strict as IM import Data.List (foldl') - -import Prelude hiding (lookup) - import Futhark.IR import Futhark.IR.Prop.Aliases +import Prelude hiding (lookup) -- | A usage table. newtype UsageTable = UsageTable (IM.IntMap Usages) - deriving (Eq, Show) + deriving (Eq, Show) instance Semigroup UsageTable where UsageTable table1 <> UsageTable table2 = @@ -62,9 +61,11 @@ -- | Expand the usage table based on aliasing information. expand :: (VName -> Names) -> UsageTable -> UsageTable expand look (UsageTable m) = UsageTable $ foldl' grow m $ IM.toList m - where grow m' (k, v) = foldl' (grow'' $ v `withoutU` presentU) m' $ - namesIntMap $ look $ VName (nameFromString "") k - grow'' v m'' k = IM.insertWith (<>) (baseTag k) v m'' + where + grow m' (k, v) = + foldl' (grow'' $ v `withoutU` presentU) m' $ + namesIntMap $ look $ VName (nameFromString "") k + grow'' v m'' k = IM.insertWith (<>) (baseTag k) v m'' is :: Usages -> VName -> UsageTable -> Bool is = lookupPred . matches @@ -144,30 +145,43 @@ -- a single statement. usageInStm :: (ASTLore lore, Aliased lore) => Stm lore -> UsageTable usageInStm (Let pat lore e) = - mconcat [usageInPat, - usageInExpLore, - usageInExp e, - usages (freeIn e)] - where usageInPat = - usages (mconcat (map freeIn $ patternElements pat) - `namesSubtract` - namesFromList (patternNames pat)) <> - sizeUsages (foldMap (freeIn . patElemType) (patternElements pat)) - usageInExpLore = - usages $ freeIn lore + mconcat + [ usageInPat, + usageInExpLore, + usageInExp e, + usages (freeIn e) + ] + where + usageInPat = + usages + ( mconcat (map freeIn $ patternElements pat) + `namesSubtract` namesFromList (patternNames pat) + ) + <> sizeUsages (foldMap (freeIn . patElemType) (patternElements pat)) + usageInExpLore = + usages $ freeIn lore usageInExp :: Aliased lore => Exp lore -> UsageTable usageInExp (Apply _ args _ _) = - mconcat [ mconcat $ map consumedUsage $ - namesToList $ subExpAliases arg - | (arg,d) <- args, d == Consume ] + mconcat + [ mconcat $ + map consumedUsage $ + namesToList $ subExpAliases arg + | (arg, d) <- args, + d == Consume + ] usageInExp (DoLoop _ merge _ _) = - mconcat [ mconcat $ map consumedUsage $ - namesToList $ subExpAliases se - | (v,se) <- merge, unique $ paramDeclType v ] + mconcat + [ mconcat $ + map consumedUsage $ + namesToList $ subExpAliases se + | (v, se) <- merge, + unique $ paramDeclType v + ] usageInExp (If _ tbranch fbranch _) = - foldMap consumedUsage $ namesToList $ - consumedInBody tbranch <> consumedInBody fbranch + foldMap consumedUsage $ + namesToList $ + consumedInBody tbranch <> consumedInBody fbranch usageInExp (BasicOp (Update src _ _)) = consumedUsage src usageInExp (Op op) = diff --git a/src/Futhark/Bench.hs b/src/Futhark/Bench.hs --- a/src/Futhark/Bench.hs +++ b/src/Futhark/Bench.hs @@ -1,61 +1,59 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} + -- | Facilities for handling Futhark benchmark results. A Futhark -- benchmark program is just like a Futhark test program. module Futhark.Bench - ( RunResult (..) - , DataResult(..) - , BenchResult(..) - , encodeBenchResults - , decodeBenchResults - - , binaryName - - , benchmarkDataset - , RunOptions(..) - - , prepareBenchmarkProgram - , CompileOptions(..) + ( RunResult (..), + DataResult (..), + BenchResult (..), + encodeBenchResults, + decodeBenchResults, + binaryName, + benchmarkDataset, + RunOptions (..), + prepareBenchmarkProgram, + CompileOptions (..), ) - where +where import Control.Applicative import Control.Concurrent (forkIO, killThread, threadDelay) import Control.Monad.Except +import qualified Data.Aeson as JSON import qualified Data.ByteString.Char8 as SBS import qualified Data.ByteString.Lazy.Char8 as LBS import qualified Data.HashMap.Strict as HM -import qualified Data.Aeson as JSON import qualified Data.Text as T -import qualified Data.Text.IO as T import qualified Data.Text.Encoding as T -import System.FilePath +import qualified Data.Text.IO as T +import Futhark.Test import System.Exit +import System.FilePath import System.IO import System.IO.Error import System.IO.Temp (withSystemTempFile) import System.Process.ByteString (readProcessWithExitCode) import System.Timeout (timeout) -import Futhark.Test - -- | The runtime of a single succesful run. -newtype RunResult = RunResult { runMicroseconds :: Int } - deriving (Eq, Show) +newtype RunResult = RunResult {runMicroseconds :: Int} + deriving (Eq, Show) -- | The results for a single named dataset is either an error -- message, or runtime measurements along the stderr that was -- produced. data DataResult = DataResult String (Either T.Text ([RunResult], T.Text)) - deriving (Eq, Show) + deriving (Eq, Show) -- | The results for all datasets for some benchmark program. data BenchResult = BenchResult FilePath [DataResult] - deriving (Eq, Show) + deriving (Eq, Show) -newtype DataResults = DataResults { unDataResults :: [DataResult] } -newtype BenchResults = BenchResults { unBenchResults :: [BenchResult] } +newtype DataResults = DataResults {unDataResults :: [DataResult]} +newtype BenchResults = BenchResults {unBenchResults :: [BenchResult]} + instance JSON.ToJSON RunResult where toJSON = JSON.toJSON . runMicroseconds @@ -71,24 +69,29 @@ instance JSON.FromJSON DataResults where parseJSON = JSON.withObject "datasets" $ \o -> DataResults <$> mapM datasetResult (HM.toList o) - where datasetResult (k, v) = - DataResult (T.unpack k) <$> - ((Right <$> success v) <|> (Left <$> JSON.parseJSON v)) - success = JSON.withObject "result" $ \o -> - (,) <$> o JSON..: "runtimes" <*> o JSON..: "stderr" + where + datasetResult (k, v) = + DataResult (T.unpack k) + <$> ((Right <$> success v) <|> (Left <$> JSON.parseJSON v)) + success = JSON.withObject "result" $ \o -> + (,) <$> o JSON..: "runtimes" <*> o JSON..: "stderr" dataResultJSON :: DataResult -> (T.Text, JSON.Value) dataResultJSON (DataResult desc (Left err)) = (T.pack desc, JSON.toJSON err) dataResultJSON (DataResult desc (Right (runtimes, progerr))) = - (T.pack desc, JSON.object - [("runtimes", JSON.toJSON $ map runMicroseconds runtimes), - ("stderr", JSON.toJSON progerr)]) + ( T.pack desc, + JSON.object + [ ("runtimes", JSON.toJSON $ map runMicroseconds runtimes), + ("stderr", JSON.toJSON progerr) + ] + ) benchResultJSON :: BenchResult -> (T.Text, JSON.Value) benchResultJSON (BenchResult prog r) = - (T.pack prog, - JSON.Object $ HM.singleton "datasets" (JSON.toJSON $ DataResults r)) + ( T.pack prog, + JSON.Object $ HM.singleton "datasets" (JSON.toJSON $ DataResults r) + ) instance JSON.ToJSON BenchResults where toJSON (BenchResults rs) = @@ -97,11 +100,12 @@ instance JSON.FromJSON BenchResults where parseJSON = JSON.withObject "benchmarks" $ \o -> BenchResults <$> mapM onBenchmark (HM.toList o) - where onBenchmark (k, v) = - BenchResult (T.unpack k) <$> - JSON.withObject "benchmark" onBenchmark' v - onBenchmark' o = - fmap unDataResults . JSON.parseJSON =<< o JSON..: "datasets" + where + onBenchmark (k, v) = + BenchResult (T.unpack k) + <$> JSON.withObject "benchmark" onBenchmark' v + onBenchmark' o = + fmap unDataResults . JSON.parseJSON =<< o JSON..: "datasets" -- | Transform benchmark results to a JSON bytestring. encodeBenchResults :: [BenchResult] -> LBS.ByteString @@ -116,18 +120,24 @@ readRuntime :: T.Text -> Maybe Int readRuntime s = case reads $ T.unpack s of [(runtime, _)] -> Just runtime - _ -> Nothing + _ -> Nothing didNotFail :: FilePath -> ExitCode -> T.Text -> ExceptT T.Text IO () didNotFail _ ExitSuccess _ = return () didNotFail program (ExitFailure code) stderr_s = - throwError $ T.pack $ program ++ " failed with error code " ++ show code ++ - " and output:\n" ++ T.unpack stderr_s + throwError $ + T.pack $ + program ++ " failed with error code " ++ show code + ++ " and output:\n" + ++ T.unpack stderr_s -compareResult :: (MonadError T.Text m, MonadIO m) => - FilePath -> (SBS.ByteString, [Value]) -> (SBS.ByteString, [Value]) - -> m () +compareResult :: + (MonadError T.Text m, MonadIO m) => + FilePath -> + (SBS.ByteString, [Value]) -> + (SBS.ByteString, [Value]) -> + m () compareResult program (expected_bs, expected_vs) (actual_bs, actual_vs) = case compareValues1 actual_vs expected_vs of Just mismatch -> do @@ -135,165 +145,194 @@ expectedf = program `replaceExtension` "expected" liftIO $ SBS.writeFile actualf actual_bs liftIO $ SBS.writeFile expectedf expected_bs - throwError $ T.pack actualf <> " and " <> T.pack expectedf <> - " do not match:\n" <> T.pack (show mismatch) + throwError $ + T.pack actualf <> " and " <> T.pack expectedf + <> " do not match:\n" + <> T.pack (show mismatch) Nothing -> return () -runResult :: (MonadError T.Text m, MonadIO m) => - FilePath - -> ExitCode - -> SBS.ByteString - -> SBS.ByteString - -> m (SBS.ByteString, [Value]) +runResult :: + (MonadError T.Text m, MonadIO m) => + FilePath -> + ExitCode -> + SBS.ByteString -> + SBS.ByteString -> + m (SBS.ByteString, [Value]) runResult program ExitSuccess stdout_s _ = case valuesFromByteString "stdout" $ LBS.fromStrict stdout_s of - Left e -> do + Left e -> do let actualf = program `replaceExtension` "actual" liftIO $ SBS.writeFile actualf stdout_s throwError $ T.pack $ show e <> "\n(See " <> actualf <> ")" Right vs -> return (stdout_s, vs) runResult program (ExitFailure code) _ stderr_s = - throwError $ T.pack $ binaryName program ++ " failed with error code " ++ show code ++ - " and output:\n" ++ T.unpack (T.decodeUtf8 stderr_s) + throwError $ + T.pack $ + binaryName program ++ " failed with error code " ++ show code + ++ " and output:\n" + ++ T.unpack (T.decodeUtf8 stderr_s) -- | How to run a benchmark. -data RunOptions = - RunOptions - { runRunner :: String - , runRuns :: Int - , runExtraOptions :: [String] - , runTimeout :: Int - , runVerbose :: Int - , runResultAction :: Maybe (Int -> IO ()) - -- ^ Invoked for every runtime measured during the run. Can be +data RunOptions = RunOptions + { runRunner :: String, + runRuns :: Int, + runExtraOptions :: [String], + runTimeout :: Int, + runVerbose :: Int, + -- | Invoked for every runtime measured during the run. Can be -- used to provide a progress bar. + runResultAction :: Maybe (Int -> IO ()) } - -- | Like @tail -f@, but running an arbitrary IO action per line. follow :: (String -> IO ()) -> FilePath -> IO () follow f fname = go 0 - where go i = do - i' <- withFile fname ReadMode $ \h -> do - hSeek h AbsoluteSeek i - goH h i - go i' + where + go i = do + i' <- withFile fname ReadMode $ \h -> do + hSeek h AbsoluteSeek i + goH h i + go i' - goH h i = do - res <- tryIOError $ hGetLine h - case res of - Left e | isEOFError e -> do - threadDelay followDelayMicroseconds - pure i - | otherwise -> ioError e - Right l -> do f l - goH h =<< hTell h + goH h i = do + res <- tryIOError $ hGetLine h + case res of + Left e + | isEOFError e -> do + threadDelay followDelayMicroseconds + pure i + | otherwise -> ioError e + Right l -> do + f l + goH h =<< hTell h - triesPerSecond = 10 - followDelayMicroseconds = 1000000 `div` triesPerSecond + triesPerSecond = 10 + followDelayMicroseconds = 1000000 `div` triesPerSecond -- | Run the benchmark program on the indicated dataset. -benchmarkDataset :: RunOptions -> FilePath -> T.Text - -> Values -> Maybe Success -> FilePath - -> IO (Either T.Text ([RunResult], T.Text)) +benchmarkDataset :: + RunOptions -> + FilePath -> + T.Text -> + Values -> + Maybe Success -> + FilePath -> + IO (Either T.Text ([RunResult], T.Text)) benchmarkDataset opts program entry input_spec expected_spec ref_out = -- We store the runtime in a temporary file. withSystemTempFile "futhark-bench" $ \tmpfile h -> do - hClose h -- We will be writing and reading this ourselves. - input <- getValuesBS dir input_spec - let getValuesAndBS (SuccessValues vs) = do - vs' <- getValues dir vs - bs <- getValuesBS dir vs - return (LBS.toStrict bs, vs') - getValuesAndBS SuccessGenerateValues = - getValuesAndBS $ SuccessValues $ InFile ref_out - maybe_expected <- maybe (return Nothing) (fmap Just . getValuesAndBS) expected_spec - let options = runExtraOptions opts ++ ["-e", T.unpack entry, - "-t", tmpfile, - "-r", show $ runRuns opts, - "-b"] + hClose h -- We will be writing and reading this ourselves. + input <- getValuesBS dir input_spec + let getValuesAndBS (SuccessValues vs) = do + vs' <- getValues dir vs + bs <- getValuesBS dir vs + return (LBS.toStrict bs, vs') + getValuesAndBS SuccessGenerateValues = + getValuesAndBS $ SuccessValues $ InFile ref_out + maybe_expected <- maybe (return Nothing) (fmap Just . getValuesAndBS) expected_spec + let options = + runExtraOptions opts + ++ [ "-e", + T.unpack entry, + "-t", + tmpfile, + "-r", + show $ runRuns opts, + "-b" + ] - -- Explicitly prefixing the current directory is necessary for - -- readProcessWithExitCode to find the binary when binOutputf has - -- no program component. - let (to_run, to_run_args) - | null $ runRunner opts = ("." binaryName program, options) - | otherwise = (runRunner opts, binaryName program : options) + -- Explicitly prefixing the current directory is necessary for + -- readProcessWithExitCode to find the binary when binOutputf has + -- no program component. + let (to_run, to_run_args) + | null $ runRunner opts = ("." binaryName program, options) + | otherwise = (runRunner opts, binaryName program : options) - when (runVerbose opts > 1) $ - putStrLn $ unwords ["Running executable", show to_run, - "with arguments", show to_run_args] + when (runVerbose opts > 1) $ + putStrLn $ + unwords + [ "Running executable", + show to_run, + "with arguments", + show to_run_args + ] - let onResult l - | Just f <- runResultAction opts, - [(x, "")] <- reads l = + let onResult l + | Just f <- runResultAction opts, + [(x, "")] <- reads l = f x - | otherwise = + | otherwise = pure () - watcher <- forkIO $ follow onResult tmpfile - - run_res <- - timeout (runTimeout opts * 1000000) $ - readProcessWithExitCode to_run to_run_args $ - LBS.toStrict input + watcher <- forkIO $ follow onResult tmpfile - killThread watcher + run_res <- + timeout (runTimeout opts * 1000000) $ + readProcessWithExitCode to_run to_run_args $ + LBS.toStrict input - runExceptT $ case run_res of - Just (progCode, output, progerr) -> do - case maybe_expected of - Nothing -> - didNotFail program progCode $ T.decodeUtf8 progerr - Just expected -> - compareResult program expected =<< - runResult program progCode output progerr - runtime_result <- liftIO $ T.readFile tmpfile - runtimes <- case mapM readRuntime $ T.lines runtime_result of - Just runtimes -> return $ map RunResult runtimes - Nothing -> throwError $ "Runtime file has invalid contents:\n" <> runtime_result + killThread watcher - return (runtimes, T.decodeUtf8 progerr) - Nothing -> - throwError $ T.pack $ "Execution exceeded " ++ show (runTimeout opts) ++ " seconds." + runExceptT $ case run_res of + Just (progCode, output, progerr) -> do + case maybe_expected of + Nothing -> + didNotFail program progCode $ T.decodeUtf8 progerr + Just expected -> + compareResult program expected + =<< runResult program progCode output progerr + runtime_result <- liftIO $ T.readFile tmpfile + runtimes <- case mapM readRuntime $ T.lines runtime_result of + Just runtimes -> return $ map RunResult runtimes + Nothing -> throwError $ "Runtime file has invalid contents:\n" <> runtime_result - where dir = takeDirectory program + return (runtimes, T.decodeUtf8 progerr) + Nothing -> + throwError $ T.pack $ "Execution exceeded " ++ show (runTimeout opts) ++ " seconds." + where + dir = takeDirectory program -- | How to compile a benchmark. -data CompileOptions = - CompileOptions - { compFuthark :: String - , compBackend :: String - , compOptions :: [String] +data CompileOptions = CompileOptions + { compFuthark :: String, + compBackend :: String, + compOptions :: [String] } progNotFound :: String -> String progNotFound s = s ++ ": command not found" -- | Compile and produce reference datasets. -prepareBenchmarkProgram :: MonadIO m => - Maybe Int - -> CompileOptions - -> FilePath - -> [InputOutputs] - -> m (Either (String, Maybe SBS.ByteString) ()) +prepareBenchmarkProgram :: + MonadIO m => + Maybe Int -> + CompileOptions -> + FilePath -> + [InputOutputs] -> + m (Either (String, Maybe SBS.ByteString) ()) prepareBenchmarkProgram concurrency opts program cases = do let futhark = compFuthark opts ref_res <- runExceptT $ ensureReferenceOutput concurrency futhark "c" program cases case ref_res of Left err -> - return $ Left ("Reference output generation for " ++ program ++ " failed:\n" ++ - unlines (map T.unpack err), - Nothing) - + return $ + Left + ( "Reference output generation for " ++ program ++ " failed:\n" + ++ unlines (map T.unpack err), + Nothing + ) Right () -> do - (futcode, _, futerr) <- liftIO $ readProcessWithExitCode futhark - ([compBackend opts, program, "-o", binaryName program] <> - compOptions opts) - "" + (futcode, _, futerr) <- + liftIO $ + readProcessWithExitCode + futhark + ( [compBackend opts, program, "-o", binaryName program] + <> compOptions opts + ) + "" case futcode of - ExitSuccess -> return $ Right () + ExitSuccess -> return $ Right () ExitFailure 127 -> return $ Left (progNotFound futhark, Nothing) - ExitFailure _ -> return $ Left ("Compilation of " ++ program ++ " failed:\n", Just futerr) + ExitFailure _ -> return $ Left ("Compilation of " ++ program ++ " failed:\n", Just futerr) diff --git a/src/Futhark/Binder.hs b/src/Futhark/Binder.hs --- a/src/Futhark/Binder.hs +++ b/src/Futhark/Binder.hs @@ -1,7 +1,13 @@ -{-# LANGUAGE FlexibleContexts, GeneralizedNewtypeDeriving, TypeFamilies, FlexibleInstances, MultiParamTypeClasses, UndecidableInstances #-} {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + -- | This module defines a convenience monad/typeclass for creating -- normalised programs. The fundamental building block is 'BinderT' -- and its execution functions, but it is usually easier to use @@ -10,27 +16,28 @@ -- See "Futhark.Construct" for a high-level description. module Futhark.Binder ( -- * A concrete @MonadBinder@ monad. - BinderT - , runBinderT, runBinderT_ - , runBinderT', runBinderT'_ - , BinderOps (..) - , Binder - , runBinder - , runBinder_ - , runBodyBinder + BinderT, + runBinderT, + runBinderT_, + runBinderT', + runBinderT'_, + BinderOps (..), + Binder, + runBinder, + runBinder_, + runBodyBinder, - -- * The 'MonadBinder' typeclass - , module Futhark.Binder.Class + -- * The 'MonadBinder' typeclass + module Futhark.Binder.Class, ) where import Control.Arrow (second) -import Control.Monad.Writer -import Control.Monad.State.Strict -import Control.Monad.Reader import Control.Monad.Error.Class +import Control.Monad.Reader +import Control.Monad.State.Strict +import Control.Monad.Writer import qualified Data.Map.Strict as M - import Futhark.Binder.Class import Futhark.IR @@ -38,23 +45,41 @@ -- 'MonadBinder' for lores that implement this type class, which -- contains methods for constructing statements. class ASTLore lore => BinderOps lore where - mkExpDecB :: (MonadBinder m, Lore m ~ lore) => - Pattern lore -> Exp lore -> m (ExpDec lore) - mkBodyB :: (MonadBinder m, Lore m ~ lore) => - Stms lore -> Result -> m (Body lore) - mkLetNamesB :: (MonadBinder m, Lore m ~ lore) => - [VName] -> Exp lore -> m (Stm lore) + mkExpDecB :: + (MonadBinder m, Lore m ~ lore) => + Pattern lore -> + Exp lore -> + m (ExpDec lore) + mkBodyB :: + (MonadBinder m, Lore m ~ lore) => + Stms lore -> + Result -> + m (Body lore) + mkLetNamesB :: + (MonadBinder m, Lore m ~ lore) => + [VName] -> + Exp lore -> + m (Stm lore) - default mkExpDecB :: (MonadBinder m, Bindable lore) => - Pattern lore -> Exp lore -> m (ExpDec lore) + default mkExpDecB :: + (MonadBinder m, Bindable lore) => + Pattern lore -> + Exp lore -> + m (ExpDec lore) mkExpDecB pat e = return $ mkExpDec pat e - default mkBodyB :: (MonadBinder m, Bindable lore) => - Stms lore -> Result -> m (Body lore) + default mkBodyB :: + (MonadBinder m, Bindable lore) => + Stms lore -> + Result -> + m (Body lore) mkBodyB stms res = return $ mkBody stms res - default mkLetNamesB :: (MonadBinder m, Lore m ~ lore, Bindable lore) => - [VName] -> Exp lore -> m (Stm lore) + default mkLetNamesB :: + (MonadBinder m, Lore m ~ lore, Bindable lore) => + [VName] -> + Exp lore -> + m (Stm lore) mkLetNamesB = mkLetNames -- | A monad transformer that tracks statements and provides a @@ -76,8 +101,10 @@ getNameSource = lift getNameSource putNameSource = lift . putNameSource -instance (ASTLore lore, Monad m) => - HasScope lore (BinderT lore m) where +instance + (ASTLore lore, Monad m) => + HasScope lore (BinderT lore m) + where lookupType name = do t <- BinderT $ gets $ M.lookup name . snd case t of @@ -85,24 +112,29 @@ Just t' -> return $ typeOf t' askScope = BinderT $ gets snd -instance (ASTLore lore, Monad m) => - LocalScope lore (BinderT lore m) where +instance + (ASTLore lore, Monad m) => + LocalScope lore (BinderT lore m) + where localScope types (BinderT m) = BinderT $ do modify $ second (M.union types) x <- m modify $ second (`M.difference` types) return x -instance (ASTLore lore, MonadFreshNames m, BinderOps lore) => - MonadBinder (BinderT lore m) where +instance + (ASTLore lore, MonadFreshNames m, BinderOps lore) => + MonadBinder (BinderT lore m) + where type Lore (BinderT lore m) = lore mkExpDecM = mkExpDecB mkBodyM = mkBodyB mkLetNamesM = mkLetNamesB addStms stms = - BinderT $ modify $ \(cur_stms,scope) -> - (cur_stms<>stms, scope `M.union` scopeOf stms) + BinderT $ + modify $ \(cur_stms, scope) -> + (cur_stms <> stms, scope `M.union` scopeOf stms) collectStms m = do (old_stms, old_scope) <- BinderT get @@ -114,67 +146,88 @@ -- | Run a binder action given an initial scope, returning a value and -- the statements added ('addStm') during the action. -runBinderT :: MonadFreshNames m => - BinderT lore m a - -> Scope lore - -> m (a, Stms lore) +runBinderT :: + MonadFreshNames m => + BinderT lore m a -> + Scope lore -> + m (a, Stms lore) runBinderT (BinderT m) scope = do (x, (stms, _)) <- runStateT m (mempty, scope) return (x, stms) -- | Like 'runBinderT', but return only the statements. -runBinderT_ :: MonadFreshNames m => - BinderT lore m () -> Scope lore -> m (Stms lore) +runBinderT_ :: + MonadFreshNames m => + BinderT lore m () -> + Scope lore -> + m (Stms lore) runBinderT_ m = fmap snd . runBinderT m -- | Like 'runBinderT', but get the initial scope from the current -- monad. -runBinderT' :: (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) => - BinderT lore m a - -> m (a, Stms lore) +runBinderT' :: + (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) => + BinderT lore m a -> + m (a, Stms lore) runBinderT' m = do scope <- askScope runBinderT m $ castScope scope -- | Like 'runBinderT_', but get the initial scope from the current -- monad. -runBinderT'_ :: (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) => - BinderT lore m a -> m (Stms lore) +runBinderT'_ :: + (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) => + BinderT lore m a -> + m (Stms lore) runBinderT'_ = fmap snd . runBinderT' -- | Run a binder action, returning a value and the statements added -- ('addStm') during the action. Assumes that the current monad -- provides initial scope and name source. -runBinder :: (MonadFreshNames m, - HasScope somelore m, SameScope somelore lore) => - Binder lore a - -> m (a, Stms lore) +runBinder :: + ( MonadFreshNames m, + HasScope somelore m, + SameScope somelore lore + ) => + Binder lore a -> + m (a, Stms lore) runBinder m = do types <- askScope modifyNameSource $ runState $ runBinderT m $ castScope types -- | Like 'runBinder', but throw away the result and just return the -- added statements. -runBinder_ :: (MonadFreshNames m, - HasScope somelore m, SameScope somelore lore) => - Binder lore a - -> m (Stms lore) +runBinder_ :: + ( MonadFreshNames m, + HasScope somelore m, + SameScope somelore lore + ) => + Binder lore a -> + m (Stms lore) runBinder_ = fmap snd . runBinder -- | Run a binder that produces a t'Body', and prefix that t'Body' by -- the statements produced during execution of the action. -runBodyBinder :: (Bindable lore, MonadFreshNames m, - HasScope somelore m, SameScope somelore lore) => - Binder lore (Body lore) -> m (Body lore) +runBodyBinder :: + ( Bindable lore, + MonadFreshNames m, + HasScope somelore m, + SameScope somelore lore + ) => + Binder lore (Body lore) -> + m (Body lore) runBodyBinder = fmap (uncurry $ flip insertStms) . runBinder -- Utility instance defintions for MTL classes. These require -- UndecidableInstances, but save on typing elsewhere. -mapInner :: Monad m => - (m (a, (Stms lore, Scope lore)) - -> m (b, (Stms lore, Scope lore))) - -> BinderT lore m a -> BinderT lore m b +mapInner :: + Monad m => + ( m (a, (Stms lore, Scope lore)) -> + m (b, (Stms lore, Scope lore)) + ) -> + BinderT lore m a -> + BinderT lore m b mapInner f (BinderT m) = BinderT $ do s <- get (x, s') <- lift $ f $ runStateT m s @@ -202,4 +255,5 @@ throwError = lift . throwError catchError (BinderT m) f = BinderT $ catchError m $ unBinder . f - where unBinder (BinderT m') = m' + where + unBinder (BinderT m') = m' diff --git a/src/Futhark/Binder/Class.hs b/src/Futhark/Binder/Class.hs --- a/src/Futhark/Binder/Class.hs +++ b/src/Futhark/Binder/Class.hs @@ -1,28 +1,28 @@ -{-# LANGUAGE FlexibleContexts, TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + -- | This module defines a convenience typeclass for creating -- normalised programs. -- -- See "Futhark.Construct" for a high-level description. module Futhark.Binder.Class - ( Bindable (..) - , mkLet - , mkLet' - , MonadBinder (..) - , insertStms - , insertStm - , letBind - , letBindNames - , collectStms_ - , bodyBind - , attributing - , auxing - - , module Futhark.MonadFreshNames + ( Bindable (..), + mkLet, + mkLet', + MonadBinder (..), + insertStms, + insertStm, + letBind, + letBindNames, + collectStms_, + bodyBind, + attributing, + auxing, + module Futhark.MonadFreshNames, ) where import qualified Data.Kind - import Futhark.IR import Futhark.MonadFreshNames @@ -32,18 +32,24 @@ -- often than you think, and the results thrown away. If used -- exclusively within a 'MonadBinder' instance, it is acceptable for -- them to create new bindings, however. -class (ASTLore lore, - FParamInfo lore ~ DeclType, - LParamInfo lore ~ Type, - RetType lore ~ DeclExtType, - BranchType lore ~ ExtType, - SetType (LetDec lore)) => - Bindable lore where +class + ( ASTLore lore, + FParamInfo lore ~ DeclType, + LParamInfo lore ~ Type, + RetType lore ~ DeclExtType, + BranchType lore ~ ExtType, + SetType (LetDec lore) + ) => + Bindable lore + where mkExpPat :: [Ident] -> [Ident] -> Exp lore -> Pattern lore mkExpDec :: Pattern lore -> Exp lore -> ExpDec lore mkBody :: Stms lore -> Result -> Body lore - mkLetNames :: (MonadFreshNames m, HasScope lore m) => - [VName] -> Exp lore -> m (Stm lore) + mkLetNames :: + (MonadFreshNames m, HasScope lore m) => + [VName] -> + Exp lore -> + m (Stm lore) -- | A monad that supports the creation of bindings from expressions -- and bodies from bindings, with a specific lore. This is the main @@ -56,21 +62,26 @@ -- effects! They may be called more often than you think, and the -- results thrown away. It is acceptable for them to create new -- bindings, however. -class (ASTLore (Lore m), - MonadFreshNames m, Applicative m, Monad m, - LocalScope (Lore m) m) => - MonadBinder m where +class + ( ASTLore (Lore m), + MonadFreshNames m, + Applicative m, + Monad m, + LocalScope (Lore m) m + ) => + MonadBinder m + where type Lore m :: Data.Kind.Type mkExpDecM :: Pattern (Lore m) -> Exp (Lore m) -> m (ExpDec (Lore m)) mkBodyM :: Stms (Lore m) -> Result -> m (Body (Lore m)) mkLetNamesM :: [VName] -> Exp (Lore m) -> m (Stm (Lore m)) -- | Add a statement to the 'Stms' under construction. - addStm :: Stm (Lore m) -> m () - addStm = addStms . oneStm + addStm :: Stm (Lore m) -> m () + addStm = addStms . oneStm -- | Add multiple statements to the 'Stms' under construction. - addStms :: Stms (Lore m) -> m () + addStms :: Stms (Lore m) -> m () -- | Obtain the statements constructed during a monadic action, -- instead of adding them to the state. @@ -82,9 +93,11 @@ certifying = censorStms . fmap . certify -- | Apply a function to the statements added by this action. -censorStms :: MonadBinder m => - (Stms (Lore m) -> Stms (Lore m)) - -> m a -> m a +censorStms :: + MonadBinder m => + (Stms (Lore m) -> Stms (Lore m)) -> + m a -> + m a censorStms f m = do (x, stms) <- collectStms m addStms $ f stms @@ -93,22 +106,30 @@ -- | Add the given attributes to any statements added by this action. attributing :: MonadBinder m => Attrs -> m a -> m a attributing attrs = censorStms $ fmap onStm - where onStm (Let pat aux e) = - Let pat aux { stmAuxAttrs = attrs <> stmAuxAttrs aux } e + where + onStm (Let pat aux e) = + Let pat aux {stmAuxAttrs = attrs <> stmAuxAttrs aux} e -- | Add the certificates and attributes to any statements added by -- this action. auxing :: MonadBinder m => StmAux anylore -> m a -> m a auxing (StmAux cs attrs _) = censorStms $ fmap onStm - where onStm (Let pat aux e) = - Let pat aux' e - where aux' = aux { stmAuxAttrs = attrs <> stmAuxAttrs aux - , stmAuxCerts = cs <> stmAuxCerts aux - } + where + onStm (Let pat aux e) = + Let pat aux' e + where + aux' = + aux + { stmAuxAttrs = attrs <> stmAuxAttrs aux, + stmAuxCerts = cs <> stmAuxCerts aux + } -- | Add a statement with the given pattern and expression. -letBind :: MonadBinder m => - Pattern (Lore m) -> Exp (Lore m) -> m () +letBind :: + MonadBinder m => + Pattern (Lore m) -> + Exp (Lore m) -> + m () letBind pat e = addStm =<< Let pat <$> (defAux <$> mkExpDecM pat e) <*> pure e @@ -118,7 +139,7 @@ mkLet ctx val e = let pat = mkExpPat ctx val e dec = mkExpDec pat e - in Let pat (defAux dec) e + in Let pat (defAux dec) e -- | Like mkLet, but also take attributes and certificates from the -- given 'StmAux'. @@ -126,7 +147,7 @@ mkLet' ctx val (StmAux cs attrs _) e = let pat = mkExpPat ctx val e dec = mkExpDec pat e - in Let pat (StmAux cs attrs dec) e + in Let pat (StmAux cs attrs dec) e -- | Add a statement with the given pattern element names and -- expression. @@ -145,7 +166,7 @@ -- | Add several bindings at the outermost level of a t'Body'. insertStms :: Bindable lore => Stms lore -> Body lore -> Body lore -insertStms stms1 (Body _ stms2 res) = mkBody (stms1<>stms2) res +insertStms stms1 (Body _ stms2 res) = mkBody (stms1 <> stms2) res -- | Add a single binding at the outermost level of a t'Body'. insertStm :: Bindable lore => Stm lore -> Body lore -> Body lore diff --git a/src/Futhark/CLI/Autotune.hs b/src/Futhark/CLI/Autotune.hs --- a/src/Futhark/CLI/Autotune.hs +++ b/src/Futhark/CLI/Autotune.hs @@ -1,40 +1,39 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} + -- | @futhark autotune@ module Futhark.CLI.Autotune (main) where import Control.Monad import qualified Data.ByteString.Char8 as SBS import Data.Function (on) -import Data.Tree -import Data.List (intersect, isPrefixOf, sort, sortOn, elemIndex, minimumBy) +import Data.List (elemIndex, intersect, isPrefixOf, minimumBy, sort, sortOn) import Data.Maybe -import Text.Read (readMaybe) -import Text.Regex.TDFA -import qualified Data.Text as T import qualified Data.Set as S - -import System.Environment (getExecutablePath) -import System.Exit -import System.Process -import System.FilePath -import System.Console.GetOpt - +import qualified Data.Text as T +import Data.Tree import Futhark.Bench import Futhark.Test import Futhark.Util (maxinum) import Futhark.Util.Options +import System.Console.GetOpt +import System.Environment (getExecutablePath) +import System.Exit +import System.FilePath +import System.Process +import Text.Read (readMaybe) +import Text.Regex.TDFA data AutotuneOptions = AutotuneOptions - { optBackend :: String - , optFuthark :: Maybe String - , optRuns :: Int - , optTuning :: Maybe String - , optExtraOptions :: [String] - , optVerbose :: Int - , optTimeout :: Int - , optDefaultThreshold :: Int - } + { optBackend :: String, + optFuthark :: Maybe String, + optRuns :: Int, + optTuning :: Maybe String, + optExtraOptions :: [String], + optVerbose :: Int, + optTimeout :: Int, + optDefaultThreshold :: Int + } initialAutotuneOptions :: AutotuneOptions initialAutotuneOptions = @@ -43,25 +42,30 @@ compileOptions :: AutotuneOptions -> IO CompileOptions compileOptions opts = do futhark <- maybe getExecutablePath return $ optFuthark opts - return $ CompileOptions { compFuthark = futhark - , compBackend = optBackend opts - , compOptions = mempty - } + return $ + CompileOptions + { compFuthark = futhark, + compBackend = optBackend opts, + compOptions = mempty + } runOptions :: Path -> Int -> AutotuneOptions -> RunOptions runOptions path timeout_s opts = - RunOptions { runRunner = "" - , runRuns = optRuns opts - , runExtraOptions = "--default-threshold" : - show (optDefaultThreshold opts) : - "-L" : - map opt path ++ - optExtraOptions opts - , runTimeout = timeout_s - , runVerbose = optVerbose opts - , runResultAction = Nothing - } - where opt (name, val) = "--size=" ++ name ++ "=" ++ show val + RunOptions + { runRunner = "", + runRuns = optRuns opts, + runExtraOptions = + "--default-threshold" : + show (optDefaultThreshold opts) : + "-L" : + map opt path + ++ optExtraOptions opts, + runTimeout = timeout_s, + runVerbose = optVerbose opts, + runResultAction = Nothing + } + where + opt (name, val) = "--size=" ++ name ++ "=" ++ show val type Path = [(String, Int)] @@ -71,13 +75,14 @@ matchM regex s :: Maybe (String, String, String, [String]) Just groups -comparisons :: String -> [(String,Int)] +comparisons :: String -> [(String, Int)] comparisons = mapMaybe isComparison . lines - where regex = makeRegex ("Compared ([^ ]+) <= (-?[0-9]+)" :: String) - isComparison l = do [thresh, val] <- regexGroups regex l - val' <- readMaybe val - return (thresh, val') - + where + regex = makeRegex ("Compared ([^ ]+) <= (-?[0-9]+)" :: String) + isComparison l = do + [thresh, val] <- regexGroups regex l + val' <- readMaybe val + return (thresh, val') type RunDataset = Int -> Path -> IO (Either String ([(String, Int)], Int)) @@ -91,18 +96,18 @@ truns <- case testAction spec of RunCases ios _ _ | not $ null ios -> do - when (optVerbose opts > 1) $ - putStrLn $ - unwords ("Entry points:" : map (T.unpack . iosEntryPoint) ios) + when (optVerbose opts > 1) $ + putStrLn $ + unwords ("Entry points:" : map (T.unpack . iosEntryPoint) ios) - res <- prepareBenchmarkProgram Nothing copts prog ios - case res of - Left (err, errstr) -> do - putStrLn err - maybe (return ()) SBS.putStrLn errstr - exitFailure - Right () -> - return ios + res <- prepareBenchmarkProgram Nothing copts prog ios + case res of + Left (err, errstr) -> do + putStrLn err + maybe (return ()) SBS.putStrLn errstr + exitFailure + Right () -> + return ios _ -> fail "Unsupported test spec." @@ -110,36 +115,44 @@ case runExpectedResult trun of Succeeds expected | null (runTags trun `intersect` ["notune", "disable"]) -> - Just (runDescription trun, run entry_point trun expected) - + Just (runDescription trun, run entry_point trun expected) _ -> Nothing - fmap concat $ forM truns $ \ios -> - forM (mapMaybe (runnableDataset $ iosEntryPoint ios) - (iosTestRuns ios)) $ - \(dataset, do_run) -> - return (dataset, do_run, iosEntryPoint ios) - - where run entry_point trun expected timeout path = do - let bestRuntime :: ([RunResult], T.Text) -> ([(String, Int)], Int) - bestRuntime (runres, errout) = - (comparisons (T.unpack errout), - minimum $ map runMicroseconds runres) + fmap concat $ + forM truns $ \ios -> + forM + ( mapMaybe + (runnableDataset $ iosEntryPoint ios) + (iosTestRuns ios) + ) + $ \(dataset, do_run) -> + return (dataset, do_run, iosEntryPoint ios) + where + run entry_point trun expected timeout path = do + let bestRuntime :: ([RunResult], T.Text) -> ([(String, Int)], Int) + bestRuntime (runres, errout) = + ( comparisons (T.unpack errout), + minimum $ map runMicroseconds runres + ) - ropts = runOptions path timeout opts + ropts = runOptions path timeout opts - when (optVerbose opts > 1) $ - putStrLn $ "Running with options: " ++ unwords (runExtraOptions ropts) + when (optVerbose opts > 1) $ + putStrLn $ "Running with options: " ++ unwords (runExtraOptions ropts) - either (Left . T.unpack) (Right . bestRuntime) <$> - benchmarkDataset ropts prog entry_point - (runInput trun) expected - (testRunReferenceOutput prog entry_point trun) + either (Left . T.unpack) (Right . bestRuntime) + <$> benchmarkDataset + ropts + prog + entry_point + (runInput trun) + expected + (testRunReferenceOutput prog entry_point trun) --- Benchmarking a program data DatasetResult = DatasetResult [(String, Int)] Double - deriving Show + deriving (Show) --- Finding initial comparisons. @@ -156,101 +169,120 @@ -- are used. tuningPaths :: ThresholdForest -> [(String, Path)] tuningPaths = concatMap (treePaths []) - where treePaths ancestors (Node (v, _) children) = - concatMap (onChild ancestors v) children ++ [(v, ancestors)] + where + treePaths ancestors (Node (v, _) children) = + concatMap (onChild ancestors v) children ++ [(v, ancestors)] - onChild ancestors v child@(Node (_, cmp) _) = - treePaths (ancestors++[(v, t cmp)]) child + onChild ancestors v child@(Node (_, cmp) _) = + treePaths (ancestors ++ [(v, t cmp)]) child - t False = thresholdMax - t True = thresholdMin + t False = thresholdMax + t True = thresholdMin thresholdForest :: FilePath -> IO ThresholdForest thresholdForest prog = do - thresholds <- getThresholds <$> - readProcess ("." dropExtension prog) ["--print-sizes"] "" + thresholds <- + getThresholds + <$> readProcess ("." dropExtension prog) ["--print-sizes"] "" let root (v, _) = ((v, False), []) - return $ unfoldForest (unfold thresholds) $ - map root $ filter (null . snd) thresholds - where getThresholds = mapMaybe findThreshold . lines - regex = makeRegex ("(.*)\\ \\(threshold\\ \\((.*)\\)\\)" :: String) - - findThreshold :: String -> Maybe (String, [(String, Bool)]) - findThreshold l = do [grp1, grp2] <- regexGroups regex l - return (grp1, - filter (not . null . fst) $ - map (\x -> if "!" `isPrefixOf` x - then (drop 1 x, False) - else (x, True)) $ - words grp2) + return $ + unfoldForest (unfold thresholds) $ + map root $ filter (null . snd) thresholds + where + getThresholds = mapMaybe findThreshold . lines + regex = makeRegex ("(.*)\\ \\(threshold\\ \\((.*)\\)\\)" :: String) - unfold thresholds ((parent, parent_cmp), ancestors) = - let ancestors' = parent : ancestors + findThreshold :: String -> Maybe (String, [(String, Bool)]) + findThreshold l = do + [grp1, grp2] <- regexGroups regex l + return + ( grp1, + filter (not . null . fst) $ + map + ( \x -> + if "!" `isPrefixOf` x + then (drop 1 x, False) + else (x, True) + ) + $ words grp2 + ) - isChild (v, v_ancestors) = do - cmp <- lookup parent v_ancestors - guard $ - sort (map fst v_ancestors) == - sort (parent : ancestors) - return ((v, cmp), ancestors') + unfold thresholds ((parent, parent_cmp), ancestors) = + let ancestors' = parent : ancestors - in ((parent, parent_cmp), mapMaybe isChild thresholds) + isChild (v, v_ancestors) = do + cmp <- lookup parent v_ancestors + guard $ + sort (map fst v_ancestors) + == sort (parent : ancestors) + return ((v, cmp), ancestors') + in ((parent, parent_cmp), mapMaybe isChild thresholds) --- Doing the atual tuning -tuneThreshold :: AutotuneOptions - -> [(DatasetName, RunDataset, T.Text)] - -> Path -> (String, Path) - -> IO Path +tuneThreshold :: + AutotuneOptions -> + [(DatasetName, RunDataset, T.Text)] -> + Path -> + (String, Path) -> + IO Path tuneThreshold opts datasets already_tuned (v, _v_path) = do (_, threshold) <- foldM tuneDataset (thresholdMin, thresholdMax) datasets return $ (v, threshold) : already_tuned - where tuneDataset :: (Int, Int) -> (DatasetName, RunDataset, T.Text) -> IO (Int, Int) tuneDataset (tMin, tMax) (dataset_name, run, entry_point) = - if not $ isPrefixOf (T.unpack entry_point ++ ".") v then do - when (optVerbose opts > 0) $ - putStrLn $ unwords [v, "is irrelevant for", T.unpack entry_point] - return (tMin, tMax) - else do - - putStrLn $ unwords ["Tuning", v, "on entry point", T.unpack entry_point, - "and dataset", dataset_name] + if not $ isPrefixOf (T.unpack entry_point ++ ".") v + then do + when (optVerbose opts > 0) $ + putStrLn $ unwords [v, "is irrelevant for", T.unpack entry_point] + return (tMin, tMax) + else do + putStrLn $ + unwords + [ "Tuning", + v, + "on entry point", + T.unpack entry_point, + "and dataset", + dataset_name + ] - sample_run <- run (optTimeout opts) ((v, tMax) : already_tuned) + sample_run <- run (optTimeout opts) ((v, tMax) : already_tuned) - case sample_run of - Left err -> do - -- If the sampling run fails, we treat it as zero information. - -- One of our ancestor thresholds will have be set such that - -- this path is never taken. - when (optVerbose opts > 0) $ putStrLn $ - "Sampling run failed:\n" ++ err - return (tMin, tMax) - Right (cmps, t) -> do - let ePars = S.toAscList $ - S.map snd $ + case sample_run of + Left err -> do + -- If the sampling run fails, we treat it as zero information. + -- One of our ancestor thresholds will have be set such that + -- this path is never taken. + when (optVerbose opts > 0) $ + putStrLn $ + "Sampling run failed:\n" ++ err + return (tMin, tMax) + Right (cmps, t) -> do + let ePars = + S.toAscList $ + S.map snd $ S.filter (candidateEPar (tMin, tMax)) $ - S.fromList cmps + S.fromList cmps - runner :: Int -> Int -> IO (Maybe Int) - runner timeout' threshold = do - res <- run timeout' ((v, threshold) : already_tuned) - case res of - Right (_, runTime) -> - return $ Just runTime - _ -> - return Nothing + runner :: Int -> Int -> IO (Maybe Int) + runner timeout' threshold = do + res <- run timeout' ((v, threshold) : already_tuned) + case res of + Right (_, runTime) -> + return $ Just runTime + _ -> + return Nothing - when (optVerbose opts > 1) $ - putStrLn $ unwords ("Got ePars: " : map show ePars) + when (optVerbose opts > 1) $ + putStrLn $ unwords ("Got ePars: " : map show ePars) - newMax <- binarySearch runner (t, tMax) ePars - let newMinIdx = pred <$> elemIndex newMax ePars - let newMin = maxinum $ catMaybes [Just tMin, newMinIdx] - return (newMin, newMax) + newMax <- binarySearch runner (t, tMax) ePars + let newMinIdx = pred <$> elemIndex newMax ePars + let newMin = maxinum $ catMaybes [Just tMin, newMinIdx] + return (newMin, newMax) bestPair :: [(Int, Int)] -> (Int, Int) bestPair = minimumBy (compare `on` fst) @@ -263,40 +295,49 @@ candidateEPar (tMin, tMax) (threshold, ePar) = ePar > tMin && ePar < tMax && threshold == v - binarySearch :: (Int -> Int -> IO (Maybe Int)) -> (Int, Int) -> [Int] -> IO Int binarySearch runner best@(best_t, best_e_par) xs = case splitAt (length xs `div` 2) xs of (lower, middle : middle' : upper) -> do when (optVerbose opts > 0) $ - putStrLn $ unwords ["Trying e_par", show middle, - "and", show middle'] + putStrLn $ + unwords + [ "Trying e_par", + show middle, + "and", + show middle' + ] candidate <- runner (timeout best_t) middle candidate' <- runner (timeout best_t) middle' case (candidate, candidate') of (Just new_t, Just new_t') -> - if new_t < new_t' then - -- recurse into lower half - binarySearch runner (bestPair [(new_t, middle), best]) lower - else - -- recurse into upper half - binarySearch runner (bestPair [(new_t', middle'), best]) upper + if new_t < new_t' + then -- recurse into lower half + binarySearch runner (bestPair [(new_t, middle), best]) lower + else -- recurse into upper half + binarySearch runner (bestPair [(new_t', middle'), best]) upper (Just new_t, Nothing) -> -- recurse into lower half binarySearch runner (bestPair [(new_t, middle), best]) lower (Nothing, Just new_t') -> - -- recurse into upper half - binarySearch runner (bestPair [(new_t', middle'), best]) upper + -- recurse into upper half + binarySearch runner (bestPair [(new_t', middle'), best]) upper (Nothing, Nothing) -> do when (optVerbose opts > 2) $ - putStrLn $ unwords ["Timing failed for candidates", - show middle, "and", show middle'] + putStrLn $ + unwords + [ "Timing failed for candidates", + show middle, + "and", + show middle' + ] return best_e_par (_, _) -> do when (optVerbose opts > 0) $ putStrLn $ unwords ["Trying e_pars", show xs] - candidates <- catMaybes . zipWith (fmap . flip (,)) xs <$> - mapM (runner $ timeout best_t) xs + candidates <- + catMaybes . zipWith (fmap . flip (,)) xs + <$> mapM (runner $ timeout best_t) xs return $ snd $ bestPair $ best : candidates --- CLI @@ -308,7 +349,7 @@ forest <- thresholdForest prog when (optVerbose opts > 0) $ - putStrLn $ ("Threshold forest:\n"++) $ drawForest $ map (fmap show) forest + putStrLn $ ("Threshold forest:\n" ++) $ drawForest $ map (fmap show) forest foldM (tuneThreshold opts datasets) [] $ tuningPaths forest @@ -316,8 +357,9 @@ runAutotuner opts prog = do best <- tune opts prog - let tuning = unlines $ do (s, n) <- sortOn fst best - return $ s ++ "=" ++ show n + let tuning = unlines $ do + (s, n) <- sortOn fst best + return $ s ++ "=" ++ show n case optTuning opts of Nothing -> return () @@ -328,54 +370,87 @@ putStrLn $ "Result of autotuning:\n" ++ tuning commandLineOptions :: [FunOptDescr AutotuneOptions] -commandLineOptions = [ - Option "r" ["runs"] - (ReqArg (\n -> +commandLineOptions = + [ Option + "r" + ["runs"] + ( ReqArg + ( \n -> case reads n of [(n', "")] | n' >= 0 -> Right $ \config -> - config { optRuns = n' - } + config + { optRuns = n' + } _ -> - Left $ error $ "'" ++ n ++ "' is not a non-negative integer.") - "RUNS") - "Run each test case this many times." - , Option [] ["backend"] - (ReqArg (\backend -> Right $ \config -> config { optBackend = backend }) - "BACKEND") - "The compiler used (defaults to 'opencl')." - , Option [] ["futhark"] - (ReqArg (\prog -> Right $ \config -> config { optFuthark = Just prog }) - "PROGRAM") - "The binary used for operations (defaults to 'futhark')." - , Option [] ["pass-option"] - (ReqArg (\opt -> - Right $ \config -> - config { optExtraOptions = opt : optExtraOptions config }) - "OPT") - "Pass this option to programs being run." - , Option [] ["tuning"] - (ReqArg (\s -> Right $ \config -> config { optTuning = Just s }) - "EXTENSION") - "Write tuning files with this extension (default: .tuning)." - , Option [] ["timeout"] - (ReqArg (\n -> - case reads n of - [(n', "")] -> - Right $ \config -> config { optTimeout = n' } - _ -> - Left $ error $ "'" ++ n ++ "' is not a non-negative integer.") - "SECONDS") - "Initial tuning timeout for each dataset. Later tuning runs are based off of the runtime of the first run." - , Option "v" ["verbose"] - (NoArg $ Right $ \config -> config { optVerbose = optVerbose config + 1 }) - "Enable logging. Pass multiple times for more." - ] + Left $ error $ "'" ++ n ++ "' is not a non-negative integer." + ) + "RUNS" + ) + "Run each test case this many times.", + Option + [] + ["backend"] + ( ReqArg + (\backend -> Right $ \config -> config {optBackend = backend}) + "BACKEND" + ) + "The compiler used (defaults to 'opencl').", + Option + [] + ["futhark"] + ( ReqArg + (\prog -> Right $ \config -> config {optFuthark = Just prog}) + "PROGRAM" + ) + "The binary used for operations (defaults to 'futhark').", + Option + [] + ["pass-option"] + ( ReqArg + ( \opt -> + Right $ \config -> + config {optExtraOptions = opt : optExtraOptions config} + ) + "OPT" + ) + "Pass this option to programs being run.", + Option + [] + ["tuning"] + ( ReqArg + (\s -> Right $ \config -> config {optTuning = Just s}) + "EXTENSION" + ) + "Write tuning files with this extension (default: .tuning).", + Option + [] + ["timeout"] + ( ReqArg + ( \n -> + case reads n of + [(n', "")] -> + Right $ \config -> config {optTimeout = n'} + _ -> + Left $ error $ "'" ++ n ++ "' is not a non-negative integer." + ) + "SECONDS" + ) + "Initial tuning timeout for each dataset. Later tuning runs are based off of the runtime of the first run.", + Option + "v" + ["verbose"] + (NoArg $ Right $ \config -> config {optVerbose = optVerbose config + 1}) + "Enable logging. Pass multiple times for more." + ] -- | Run @futhark autotune@ main :: String -> [String] -> IO () -main = mainWithOptions initialAutotuneOptions commandLineOptions - "options... program" $ - \progs config -> - case progs of [prog] -> Just $ runAutotuner config prog - _ -> Nothing +main = mainWithOptions + initialAutotuneOptions + commandLineOptions + "options... program" + $ \progs config -> + case progs of + [prog] -> Just $ runAutotuner config prog + _ -> Nothing 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 @@ -1,8 +1,9 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} + -- | @futhark bench@ -module Futhark.CLI.Bench ( main ) where +module Futhark.CLI.Bench (main) where import Control.Monad import Control.Monad.Except @@ -10,46 +11,60 @@ import qualified Data.ByteString.Lazy.Char8 as LBS import Data.Either import Data.IORef -import Data.Maybe import Data.List (foldl', sortBy) +import Data.Maybe import Data.Ord import qualified Data.Text as T +import Futhark.Bench +import Futhark.Test +import Futhark.Util (fancyTerminal, maxinum, maybeNth, pmapIO) +import Futhark.Util.Console +import Futhark.Util.Options import System.Console.ANSI (clearLine) import System.Console.GetOpt import System.Directory import System.Environment -import System.IO import System.Exit +import System.IO import Text.Printf import Text.Regex.TDFA -import Futhark.Bench -import Futhark.Test -import Futhark.Util (fancyTerminal, maybeNth, maxinum, pmapIO) -import Futhark.Util.Console -import Futhark.Util.Options - data BenchOptions = BenchOptions - { optBackend :: String - , optFuthark :: Maybe String - , optRunner :: String - , optRuns :: Int - , optExtraOptions :: [String] - , optCompilerOptions :: [String] - , optJSON :: Maybe FilePath - , optTimeout :: Int - , optSkipCompilation :: Bool - , optExcludeCase :: [String] - , optIgnoreFiles :: [Regex] - , optEntryPoint :: Maybe String - , optTuning :: Maybe String - , optConcurrency :: Maybe Int - , optVerbose :: Int - } + { optBackend :: String, + optFuthark :: Maybe String, + optRunner :: String, + optRuns :: Int, + optExtraOptions :: [String], + optCompilerOptions :: [String], + optJSON :: Maybe FilePath, + optTimeout :: Int, + optSkipCompilation :: Bool, + optExcludeCase :: [String], + optIgnoreFiles :: [Regex], + optEntryPoint :: Maybe String, + optTuning :: Maybe String, + optConcurrency :: Maybe Int, + optVerbose :: Int + } initialBenchOptions :: BenchOptions -initialBenchOptions = BenchOptions "c" Nothing "" 10 [] [] Nothing (-1) False - ["nobench", "disable"] [] Nothing (Just "tuning") Nothing 0 +initialBenchOptions = + BenchOptions + "c" + Nothing + "" + 10 + [] + [] + Nothing + (-1) + False + ["nobench", "disable"] + [] + Nothing + (Just "tuning") + Nothing + 0 runBenchmarks :: BenchOptions -> [FilePath] -> IO () runBenchmarks opts paths = do @@ -60,9 +75,10 @@ benchmarks <- filter (not . ignored . fst) <$> testSpecsFromPathsOrDie paths -- Try to avoid concurrency at both program and data set level. - let opts' = if length paths /= 1 - then opts { optConcurrency = Just 1} - else opts + let opts' = + if length paths /= 1 + then opts {optConcurrency = Just 1} + else opts (skipped_benchmarks, compiled_benchmarks) <- partitionEithers <$> pmapIO (optConcurrency opts) (compileBenchmark opts') benchmarks @@ -70,24 +86,28 @@ putStrLn $ "Reporting average runtime of " ++ show (optRuns opts) ++ " runs for each dataset." - results <- concat <$> mapM (runBenchmark opts) - (sortBy (comparing fst) compiled_benchmarks) + results <- + concat + <$> mapM + (runBenchmark opts) + (sortBy (comparing fst) compiled_benchmarks) case optJSON opts of Nothing -> return () Just file -> LBS.writeFile file $ encodeBenchResults results when (anyFailed results) exitFailure - - where ignored f = any (`match` f) $ optIgnoreFiles opts + where + ignored f = any (`match` f) $ optIgnoreFiles opts anyFailed :: [BenchResult] -> Bool anyFailed = any failedBenchResult - where failedBenchResult (BenchResult _ xs) = - any failedResult xs - failedResult (DataResult _ Left{}) = True - failedResult _ = False + where + failedBenchResult (BenchResult _ xs) = + any failedResult xs + failedResult (DataResult _ Left {}) = True + failedResult _ = False anyFailedToCompile :: [SkipReason] -> Bool -anyFailedToCompile = not . all (==Skipped) +anyFailedToCompile = not . all (== Skipped) data SkipReason = Skipped | FailedToCompile deriving (Eq) @@ -95,89 +115,108 @@ compileOptions :: BenchOptions -> IO CompileOptions compileOptions opts = do futhark <- maybe getExecutablePath return $ optFuthark opts - return $ CompileOptions { compFuthark = futhark - , compBackend = optBackend opts - , compOptions = optCompilerOptions opts - } + return $ + CompileOptions + { compFuthark = futhark, + compBackend = optBackend opts, + compOptions = optCompilerOptions opts + } -compileBenchmark :: BenchOptions -> (FilePath, ProgramTest) - -> IO (Either SkipReason (FilePath, [InputOutputs])) +compileBenchmark :: + BenchOptions -> + (FilePath, ProgramTest) -> + IO (Either SkipReason (FilePath, [InputOutputs])) compileBenchmark opts (program, spec) = case testAction spec of - RunCases cases _ _ | "nobench" `notElem` testTags spec, - "disable" `notElem` testTags spec, - any hasRuns cases -> - if optSkipCompilation opts - then do - exists <- doesFileExist $ binaryName program - if exists - then return $ Right (program, cases) - else do putStrLn $ binaryName program ++ " does not exist, but --skip-compilation passed." - return $ Left FailedToCompile - else do - - putStr $ "Compiling " ++ program ++ "...\n" - - compile_opts <- compileOptions opts + RunCases cases _ _ + | "nobench" `notElem` testTags spec, + "disable" `notElem` testTags spec, + any hasRuns cases -> + if optSkipCompilation opts + then do + exists <- doesFileExist $ binaryName program + if exists + then return $ Right (program, cases) + else do + putStrLn $ binaryName program ++ " does not exist, but --skip-compilation passed." + return $ Left FailedToCompile + else do + putStr $ "Compiling " ++ program ++ "...\n" - res <- prepareBenchmarkProgram (optConcurrency opts) compile_opts program cases + compile_opts <- compileOptions opts - case res of - Left (err, errstr) -> do - putStrLn $ inRed err - maybe (return ()) SBS.putStrLn errstr - return $ Left FailedToCompile - Right () -> - return $ Right (program, cases) + res <- prepareBenchmarkProgram (optConcurrency opts) compile_opts program cases + case res of + Left (err, errstr) -> do + putStrLn $ inRed err + maybe (return ()) SBS.putStrLn errstr + return $ Left FailedToCompile + Right () -> + return $ Right (program, cases) _ -> return $ Left Skipped - where hasRuns (InputOutputs _ runs) = not $ null runs + where + hasRuns (InputOutputs _ runs) = not $ null runs runBenchmark :: BenchOptions -> (FilePath, [InputOutputs]) -> IO [BenchResult] runBenchmark opts (program, cases) = mapM forInputOutputs $ filter relevant cases - where forInputOutputs (InputOutputs entry_name runs) = do - (tuning_opts, tuning_desc) <- determineTuning (optTuning opts) program + where + forInputOutputs (InputOutputs entry_name runs) = do + (tuning_opts, tuning_desc) <- determineTuning (optTuning opts) program - putStr $ inBold $ "\nResults for " ++ program' ++ tuning_desc ++ ":\n" - let opts' = opts { optExtraOptions = - optExtraOptions opts ++ tuning_opts } - BenchResult program' . catMaybes <$> - mapM (runBenchmarkCase opts' program entry_name pad_to) runs - where program' = if entry_name == "main" - then program - else program ++ ":" ++ T.unpack entry_name + putStr $ inBold $ "\nResults for " ++ program' ++ tuning_desc ++ ":\n" + let opts' = + opts + { optExtraOptions = + optExtraOptions opts ++ tuning_opts + } + BenchResult program' . catMaybes + <$> mapM (runBenchmarkCase opts' program entry_name pad_to) runs + where + program' = + if entry_name == "main" + then program + else program ++ ":" ++ T.unpack entry_name - relevant = maybe (const True) (==) (optEntryPoint opts) . T.unpack . iosEntryPoint + relevant = maybe (const True) (==) (optEntryPoint opts) . T.unpack . iosEntryPoint - pad_to = foldl max 0 $ concatMap (map (length . runDescription) . iosTestRuns) cases + pad_to = foldl max 0 $ concatMap (map (length . runDescription) . iosTestRuns) cases runOptions :: (Int -> IO ()) -> BenchOptions -> RunOptions runOptions f opts = - RunOptions { runRunner = optRunner opts - , runRuns = optRuns opts - , runExtraOptions = optExtraOptions opts - , runTimeout = optTimeout opts - , runVerbose = optVerbose opts - , runResultAction = Just f - } + RunOptions + { runRunner = optRunner opts, + runRuns = optRuns opts, + runExtraOptions = optExtraOptions opts, + runTimeout = optTimeout opts, + runVerbose = optVerbose opts, + runResultAction = Just f + } progressBar :: Int -> Int -> Int -> String progressBar cur bound steps = - "[" ++ map cell [1..steps] ++ "] " ++ show cur ++ "/" ++ show bound - where step_size :: Double - step_size = fromIntegral bound / fromIntegral steps - cur' = fromIntegral cur - chars = " ▏▎▍▍▌▋▊▉█" - char i = fromMaybe ' ' $ maybeNth (i::Int) chars - num_chars = fromIntegral $ length chars + "[" ++ map cell [1 .. steps] ++ "] " ++ show cur ++ "/" ++ show bound + where + step_size :: Double + step_size = fromIntegral bound / fromIntegral steps + cur' = fromIntegral cur + chars = " ▏▎▍▍▌▋▊▉█" + char i = fromMaybe ' ' $ maybeNth (i :: Int) chars + num_chars = fromIntegral $ length chars - cell :: Int -> Char - cell i - | i' * step_size <= cur' = char 9 - | otherwise = char (floor (((cur' - (i'-1) * step_size) * num_chars) - / step_size)) - where i' = fromIntegral i + cell :: Int -> Char + cell i + | i' * step_size <= cur' = char 9 + | otherwise = + char + ( floor + ( ((cur' - (i' -1) * step_size) * num_chars) + / step_size + ) + ) + where + i' = fromIntegral i descString :: String -> Int -> String descString desc pad_to = desc ++ ": " ++ replicate (pad_to - length desc) ' ' @@ -185,36 +224,45 @@ mkProgressPrompt :: Int -> Int -> String -> IO (Maybe Int -> IO ()) mkProgressPrompt runs pad_to dataset_desc | fancyTerminal = do - count <- newIORef (0::Int) - return $ \us -> do - putStr "\r" -- Go to start of line. - i <- readIORef count - let i' = if isJust us then i+1 else i - writeIORef count i' - putStr $ descString dataset_desc pad_to ++ progressBar i' runs 10 - putStr " " -- Just to move the cursor away from the progress bar. - hFlush stdout - - | otherwise = do - putStr $ descString dataset_desc pad_to + count <- newIORef (0 :: Int) + return $ \us -> do + putStr "\r" -- Go to start of line. + i <- readIORef count + let i' = if isJust us then i + 1 else i + writeIORef count i' + putStr $ descString dataset_desc pad_to ++ progressBar i' runs 10 + putStr " " -- Just to move the cursor away from the progress bar. hFlush stdout - return $ const $ return () + | otherwise = do + putStr $ descString dataset_desc pad_to + hFlush stdout + return $ const $ return () reportResult :: [RunResult] -> IO () reportResult results = do let runtimes = map (fromIntegral . runMicroseconds) results avg = sum runtimes / fromIntegral (length runtimes) rsd = stddevp runtimes / mean runtimes :: Double - putStrLn $ printf "%10.0fμs (RSD: %.3f; min: %3.0f%%; max: %+3.0f%%)" - avg rsd ((minimum runtimes / avg - 1) * 100) ((maxinum runtimes / avg - 1) * 100) + putStrLn $ + printf + "%10.0fμs (RSD: %.3f; min: %3.0f%%; max: %+3.0f%%)" + avg + rsd + ((minimum runtimes / avg - 1) * 100) + ((maxinum runtimes / avg - 1) * 100) -runBenchmarkCase :: BenchOptions -> FilePath -> T.Text -> Int -> TestRun - -> IO (Maybe DataResult) -runBenchmarkCase _ _ _ _ (TestRun _ _ RunTimeFailure{} _ _) = +runBenchmarkCase :: + BenchOptions -> + FilePath -> + T.Text -> + Int -> + TestRun -> + IO (Maybe DataResult) +runBenchmarkCase _ _ _ _ (TestRun _ _ RunTimeFailure {} _ _) = return Nothing -- Not our concern, we are not a testing tool. runBenchmarkCase opts _ _ _ (TestRun tags _ _ _ _) | any (`elem` tags) $ optExcludeCase opts = - return Nothing + return Nothing runBenchmarkCase opts program entry pad_to tr@(TestRun _ input_spec (Succeeds expected_spec) _ dataset_desc) = do prompt <- mkProgressPrompt (optRuns opts) pad_to dataset_desc @@ -222,9 +270,14 @@ -- error occurs it's easier to see where. prompt Nothing - res <- benchmarkDataset (runOptions (prompt . Just) opts) - program entry input_spec expected_spec - (testRunReferenceOutput program entry tr) + res <- + benchmarkDataset + (runOptions (prompt . Just) opts) + program + entry + input_spec + expected_spec + (testRunReferenceOutput program entry tr) when fancyTerminal $ do clearLine @@ -243,98 +296,169 @@ return $ Just $ DataResult dataset_desc $ Right (runtimes, errout) commandLineOptions :: [FunOptDescr BenchOptions] -commandLineOptions = [ - Option "r" ["runs"] - (ReqArg (\n -> +commandLineOptions = + [ Option + "r" + ["runs"] + ( ReqArg + ( \n -> case reads n of [(n', "")] | n' > 0 -> Right $ \config -> - config { optRuns = n' - } + config + { optRuns = n' + } _ -> - Left $ error $ "'" ++ n ++ "' is not a positive integer.") - "RUNS") - "Run each test case this many times." - , Option [] ["backend"] - (ReqArg (\backend -> Right $ \config -> config { optBackend = backend }) - "PROGRAM") - "The compiler used (defaults to 'futhark-c')." - , Option [] ["futhark"] - (ReqArg (\prog -> Right $ \config -> config { optFuthark = Just prog }) - "PROGRAM") - "The binary used for operations (defaults to same binary as 'futhark bench')." - , Option [] ["runner"] - (ReqArg (\prog -> Right $ \config -> config { optRunner = prog }) "PROGRAM") - "The program used to run the Futhark-generated programs (defaults to nothing)." - , Option "p" ["pass-option"] - (ReqArg (\opt -> - Right $ \config -> - config { optExtraOptions = opt : optExtraOptions config }) - "OPT") - "Pass this option to programs being run." - , Option [] ["pass-compiler-option"] - (ReqArg (\opt -> - Right $ \config -> - config { optCompilerOptions = opt : optCompilerOptions config }) - "OPT") - "Pass this option to the compiler (or typechecker if in -t mode)." - , Option [] ["json"] - (ReqArg (\file -> - Right $ \config -> config { optJSON = Just file}) - "FILE") - "Scatter results in JSON format here." - , Option [] ["timeout"] - (ReqArg (\n -> - case reads n of - [(n', "")] - | n' < max_timeout -> - Right $ \config -> config { optTimeout = fromIntegral n' } - _ -> - Left $ error $ "'" ++ n ++ - "' is not an integer smaller than" ++ show max_timeout ++ ".") - "SECONDS") - "Number of seconds before a dataset is aborted." - , Option [] ["skip-compilation"] - (NoArg $ Right $ \config -> config { optSkipCompilation = True }) - "Use already compiled program." - , Option [] ["exclude-case"] - (ReqArg (\s -> Right $ \config -> - config { optExcludeCase = s : optExcludeCase config }) - "TAG") - "Do not run test cases with this tag." - , Option [] ["ignore-files"] - (ReqArg (\s -> Right $ \config -> - config { optIgnoreFiles = makeRegex s : optIgnoreFiles config }) - "REGEX") - "Ignore files matching this regular expression." - , Option "e" ["entry-point"] - (ReqArg (\s -> Right $ \config -> - config { optEntryPoint = Just s }) - "NAME") - "Only run this entry point." - , Option [] ["tuning"] - (ReqArg (\s -> Right $ \config -> config { optTuning = Just s }) - "EXTENSION") - "Look for tuning files with this extension (defaults to .tuning)." - , Option [] ["no-tuning"] - (NoArg $ Right $ \config -> config { optTuning = Nothing }) - "Do not load tuning files." - , Option [] ["concurrency"] - (ReqArg (\n -> - case reads n of - [(n', "")] - | n' > 0 -> - Right $ \config -> config { optConcurrency = Just n' } - _ -> - Left $ error $ "'" ++ n ++ "' is not a positive integer.") - "NUM") - "Number of benchmarks to prepare (not run) concurrently." - , Option "v" ["verbose"] - (NoArg $ Right $ \config -> config { optVerbose = optVerbose config + 1 }) - "Enable logging. Pass multiple times for more." + Left $ error $ "'" ++ n ++ "' is not a positive integer." + ) + "RUNS" + ) + "Run each test case this many times.", + Option + [] + ["backend"] + ( ReqArg + (\backend -> Right $ \config -> config {optBackend = backend}) + "PROGRAM" + ) + "The compiler used (defaults to 'futhark-c').", + Option + [] + ["futhark"] + ( ReqArg + (\prog -> Right $ \config -> config {optFuthark = Just prog}) + "PROGRAM" + ) + "The binary used for operations (defaults to same binary as 'futhark bench').", + Option + [] + ["runner"] + (ReqArg (\prog -> Right $ \config -> config {optRunner = prog}) "PROGRAM") + "The program used to run the Futhark-generated programs (defaults to nothing).", + Option + "p" + ["pass-option"] + ( ReqArg + ( \opt -> + Right $ \config -> + config {optExtraOptions = opt : optExtraOptions config} + ) + "OPT" + ) + "Pass this option to programs being run.", + Option + [] + ["pass-compiler-option"] + ( ReqArg + ( \opt -> + Right $ \config -> + config {optCompilerOptions = opt : optCompilerOptions config} + ) + "OPT" + ) + "Pass this option to the compiler (or typechecker if in -t mode).", + Option + [] + ["json"] + ( ReqArg + ( \file -> + Right $ \config -> config {optJSON = Just file} + ) + "FILE" + ) + "Scatter results in JSON format here.", + Option + [] + ["timeout"] + ( ReqArg + ( \n -> + case reads n of + [(n', "")] + | n' < max_timeout -> + Right $ \config -> config {optTimeout = fromIntegral n'} + _ -> + Left $ + error $ + "'" ++ n + ++ "' is not an integer smaller than" + ++ show max_timeout + ++ "." + ) + "SECONDS" + ) + "Number of seconds before a dataset is aborted.", + Option + [] + ["skip-compilation"] + (NoArg $ Right $ \config -> config {optSkipCompilation = True}) + "Use already compiled program.", + Option + [] + ["exclude-case"] + ( ReqArg + ( \s -> Right $ \config -> + config {optExcludeCase = s : optExcludeCase config} + ) + "TAG" + ) + "Do not run test cases with this tag.", + Option + [] + ["ignore-files"] + ( ReqArg + ( \s -> Right $ \config -> + config {optIgnoreFiles = makeRegex s : optIgnoreFiles config} + ) + "REGEX" + ) + "Ignore files matching this regular expression.", + Option + "e" + ["entry-point"] + ( ReqArg + ( \s -> Right $ \config -> + config {optEntryPoint = Just s} + ) + "NAME" + ) + "Only run this entry point.", + Option + [] + ["tuning"] + ( ReqArg + (\s -> Right $ \config -> config {optTuning = Just s}) + "EXTENSION" + ) + "Look for tuning files with this extension (defaults to .tuning).", + Option + [] + ["no-tuning"] + (NoArg $ Right $ \config -> config {optTuning = Nothing}) + "Do not load tuning files.", + Option + [] + ["concurrency"] + ( ReqArg + ( \n -> + case reads n of + [(n', "")] + | n' > 0 -> + Right $ \config -> config {optConcurrency = Just n'} + _ -> + Left $ error $ "'" ++ n ++ "' is not a positive integer." + ) + "NUM" + ) + "Number of benchmarks to prepare (not run) concurrently.", + Option + "v" + ["verbose"] + (NoArg $ Right $ \config -> config {optVerbose = optVerbose config + 1}) + "Enable logging. Pass multiple times for more." ] - where max_timeout :: Int - max_timeout = maxBound `div` 1000000 + where + max_timeout :: Int + max_timeout = maxBound `div` 1000000 -- | Run @futhark bench@. main :: String -> [String] -> IO () @@ -346,7 +470,7 @@ -- | Numerically stable mean mean :: Floating a => [a] -> a -mean x = fst $ foldl' (\(!m, !n) x' -> (m+(x'-m)/(n+1),n+1)) (0,0) x +mean x = fst $ foldl' (\(!m, !n) x' -> (m + (x' - m) / (n + 1), n + 1)) (0, 0) x -- | Standard deviation of population stddevp :: (Floating a) => [a] -> a @@ -354,12 +478,12 @@ -- | Population variance pvar :: (Floating a) => [a] -> a -pvar xs = centralMoment xs (2::Int) +pvar xs = centralMoment xs (2 :: Int) -- | Central moments centralMoment :: (Floating b, Integral t) => [b] -> t -> b -centralMoment _ 1 = 0 -centralMoment xs r = sum (map (\x -> (x-m)^r) xs) / n - where - m = mean xs - n = fromIntegral $ length xs +centralMoment _ 1 = 0 +centralMoment xs r = sum (map (\x -> (x - m) ^ r) xs) / n + where + m = mean xs + n = fromIntegral $ length xs diff --git a/src/Futhark/CLI/C.hs b/src/Futhark/CLI/C.hs --- a/src/Futhark/CLI/C.hs +++ b/src/Futhark/CLI/C.hs @@ -1,14 +1,19 @@ {-# LANGUAGE FlexibleContexts #-} + -- | @futhark c@ module Futhark.CLI.C (main) where import Futhark.Actions (compileCAction) -import Futhark.Passes (sequentialCpuPipeline) import Futhark.Compiler.CLI +import Futhark.Passes (sequentialCpuPipeline) -- | Run @futhark c@ main :: String -> [String] -> IO () -main = compilerMain () [] - "Compile sequential C" "Generate sequential C code from optimised Futhark program." - sequentialCpuPipeline $ \fcfg () mode outpath prog -> - actionProcedure (compileCAction fcfg mode outpath) prog +main = compilerMain + () + [] + "Compile sequential C" + "Generate sequential C code from optimised Futhark program." + sequentialCpuPipeline + $ \fcfg () mode outpath prog -> + actionProcedure (compileCAction fcfg mode outpath) prog diff --git a/src/Futhark/CLI/CUDA.hs b/src/Futhark/CLI/CUDA.hs --- a/src/Futhark/CLI/CUDA.hs +++ b/src/Futhark/CLI/CUDA.hs @@ -1,14 +1,19 @@ {-# LANGUAGE FlexibleContexts #-} + -- | @futhark cuda@ module Futhark.CLI.CUDA (main) where import Futhark.Actions (compileCUDAAction) -import Futhark.Passes (gpuPipeline) import Futhark.Compiler.CLI +import Futhark.Passes (gpuPipeline) -- | Run @futhark cuda@. main :: String -> [String] -> IO () -main = compilerMain () [] - "Compile CUDA" "Generate CUDA/C code from optimised Futhark program." - gpuPipeline $ \fcfg () mode outpath prog -> - actionProcedure (compileCUDAAction fcfg mode outpath) prog +main = compilerMain + () + [] + "Compile CUDA" + "Generate CUDA/C code from optimised Futhark program." + gpuPipeline + $ \fcfg () mode outpath prog -> + actionProcedure (compileCUDAAction fcfg mode outpath) prog diff --git a/src/Futhark/CLI/Check.hs b/src/Futhark/CLI/Check.hs --- a/src/Futhark/CLI/Check.hs +++ b/src/Futhark/CLI/Check.hs @@ -3,21 +3,24 @@ import Control.Monad import Control.Monad.IO.Class -import System.Console.GetOpt -import System.IO - import Futhark.Compiler import Futhark.Util.Options +import System.Console.GetOpt +import System.IO -newtype CheckConfig = CheckConfig { checkWarn :: Bool } +newtype CheckConfig = CheckConfig {checkWarn :: Bool} newCheckConfig :: CheckConfig newCheckConfig = CheckConfig True options :: [FunOptDescr CheckConfig] -options = [Option "w" [] - (NoArg $ Right $ \cfg -> cfg { checkWarn = False }) - "Disable all warnings."] +options = + [ Option + "w" + [] + (NoArg $ Right $ \cfg -> cfg {checkWarn = False}) + "Disable all warnings." + ] -- | Run @futhark check@. main :: String -> [String] -> IO () diff --git a/src/Futhark/CLI/Datacmp.hs b/src/Futhark/CLI/Datacmp.hs --- a/src/Futhark/CLI/Datacmp.hs +++ b/src/Futhark/CLI/Datacmp.hs @@ -1,29 +1,30 @@ {-# LANGUAGE OverloadedStrings #-} + -- | @futhark datacmp@ module Futhark.CLI.Datacmp (main) where import qualified Data.ByteString.Lazy.Char8 as BS -import System.Exit - import Futhark.Test.Values import Futhark.Util.Options +import System.Exit -- | Run @futhark datacmp@ main :: String -> [String] -> IO () main = mainWithOptions () [] " " f - where f [file_a, file_b] () = Just $ do - vs_a_maybe <- readValues <$> BS.readFile file_a - vs_b_maybe <- readValues <$> BS.readFile file_b - case (vs_a_maybe, vs_b_maybe) of - (Nothing, _) -> - error $ "Error reading values from " ++ file_a - (_, Nothing) -> - error $ "Error reading values from " ++ file_b - (Just vs_a, Just vs_b) -> - case compareValues vs_a vs_b of - [] -> return () - es -> do mapM_ print es - exitWith $ ExitFailure 2 - - f _ _ = - Nothing + where + f [file_a, file_b] () = Just $ do + vs_a_maybe <- readValues <$> BS.readFile file_a + vs_b_maybe <- readValues <$> BS.readFile file_b + case (vs_a_maybe, vs_b_maybe) of + (Nothing, _) -> + error $ "Error reading values from " ++ file_a + (_, Nothing) -> + error $ "Error reading values from " ++ file_b + (Just vs_a, Just vs_b) -> + case compareValues vs_a vs_b of + [] -> return () + es -> do + mapM_ print es + exitWith $ ExitFailure 2 + f _ _ = + Nothing diff --git a/src/Futhark/CLI/Dataset.hs b/src/Futhark/CLI/Dataset.hs --- a/src/Futhark/CLI/Dataset.hs +++ b/src/Futhark/CLI/Dataset.hs @@ -1,5 +1,6 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Strict #-} + -- | @futhark dataset@ module Futhark.CLI.Dataset (main) where @@ -9,155 +10,193 @@ import qualified Data.ByteString.Lazy.Char8 as BS import qualified Data.Map.Strict as M import qualified Data.Text as T -import Data.Word -import qualified Data.Vector.Unboxed.Mutable as UMVec -import qualified Data.Vector.Unboxed as UVec import Data.Vector.Generic (freeze) +import qualified Data.Vector.Unboxed as UVec +import qualified Data.Vector.Unboxed.Mutable as UMVec +import Data.Word +import Futhark.Test.Values +import Futhark.Util.Options +import Language.Futhark.Parser +import Language.Futhark.Pretty () +import Language.Futhark.Prop (UncheckedTypeExp, namesToPrimTypes) +import Language.Futhark.Syntax hiding + ( FloatValue (..), + IntValue (..), + PrimValue (..), + Value, + ValueType, + ) import System.Console.GetOpt import System.Exit import System.IO -import System.Random.PCG (initialize, Variate, uniformR) - -import Language.Futhark.Syntax hiding - (Value, ValueType, PrimValue(..), IntValue(..), FloatValue(..)) -import Language.Futhark.Prop (UncheckedTypeExp, namesToPrimTypes) -import Language.Futhark.Parser -import Language.Futhark.Pretty () - -import Futhark.Test.Values -import Futhark.Util.Options +import System.Random.PCG (Variate, initialize, uniformR) -- | Run @futhark dataset@. main :: String -> [String] -> IO () main = mainWithOptions initialDataOptions commandLineOptions "options..." f - where f [] config - | null $ optOrders config = Just $ do - maybe_vs <- readValues <$> BS.getContents - case maybe_vs of - Nothing -> do hPutStrLn stderr "Malformed data on standard input." - exitFailure - Just vs -> - case format config of - Text -> mapM_ (putStrLn . pretty) vs - Binary -> mapM_ (BS.putStr . Bin.encode) vs - Type -> mapM_ (putStrLn . pretty . valueType) vs - | otherwise = - Just $ zipWithM_ ($) (optOrders config) - [fromIntegral (optSeed config)..] - f _ _ = - Nothing + where + f [] config + | null $ optOrders config = Just $ do + maybe_vs <- readValues <$> BS.getContents + case maybe_vs of + Nothing -> do + hPutStrLn stderr "Malformed data on standard input." + exitFailure + Just vs -> + case format config of + Text -> mapM_ (putStrLn . pretty) vs + Binary -> mapM_ (BS.putStr . Bin.encode) vs + Type -> mapM_ (putStrLn . pretty . valueType) vs + | otherwise = + Just $ + zipWithM_ + ($) + (optOrders config) + [fromIntegral (optSeed config) ..] + f _ _ = + Nothing -data OutputFormat = Text - | Binary - | Type - deriving (Eq, Ord, Show) +data OutputFormat + = Text + | Binary + | Type + deriving (Eq, Ord, Show) data DataOptions = DataOptions - { optSeed :: Int - , optRange :: RandomConfiguration - , optOrders :: [Word64 -> IO ()] - , format :: OutputFormat - } + { optSeed :: Int, + optRange :: RandomConfiguration, + optOrders :: [Word64 -> IO ()], + format :: OutputFormat + } initialDataOptions :: DataOptions initialDataOptions = DataOptions 1 initialRandomConfiguration [] Text commandLineOptions :: [FunOptDescr DataOptions] -commandLineOptions = [ - Option "s" ["seed"] - (ReqArg (\n -> +commandLineOptions = + [ Option + "s" + ["seed"] + ( ReqArg + ( \n -> case reads n of [(n', "")] -> - Right $ \config -> config { optSeed = n' } + Right $ \config -> config {optSeed = n'} _ -> - Left $ do hPutStrLn stderr $ "'" ++ n ++ "' is not an integer." - exitFailure) - "SEED") - "The seed to use when initialising the RNG." - , Option "g" ["generate"] - (ReqArg (\t -> + Left $ do + hPutStrLn stderr $ "'" ++ n ++ "' is not an integer." + exitFailure + ) + "SEED" + ) + "The seed to use when initialising the RNG.", + Option + "g" + ["generate"] + ( ReqArg + ( \t -> case tryMakeGenerator t of Right g -> Right $ \config -> - config { optOrders = - optOrders config ++ - [g (optRange config) (format config)] - } + config + { optOrders = + optOrders config + ++ [g (optRange config) (format config)] + } Left err -> - Left $ do hPutStrLn stderr err - exitFailure) - "TYPE") - "Generate a random value of this type." - , Option [] ["text"] - (NoArg $ Right $ \opts -> opts { format = Text }) - "Output data in text format (must precede --generate)." - , Option "b" ["binary"] - (NoArg $ Right $ \opts -> opts { format = Binary }) - "Output data in binary Futhark format (must precede --generate)." - , Option "t" ["type"] - (NoArg $ Right $ \opts -> opts { format = Type }) - "Output the type (textually) rather than the value (must precede --generate)." - , setRangeOption "i8" seti8Range - , setRangeOption "i16" seti16Range - , setRangeOption "i32" seti32Range - , setRangeOption "i64" seti64Range - , setRangeOption "u8" setu8Range - , setRangeOption "u16" setu16Range - , setRangeOption "u32" setu32Range - , setRangeOption "u64" setu64Range - , setRangeOption "f32" setf32Range - , setRangeOption "f64" setf64Range + Left $ do + hPutStrLn stderr err + exitFailure + ) + "TYPE" + ) + "Generate a random value of this type.", + Option + [] + ["text"] + (NoArg $ Right $ \opts -> opts {format = Text}) + "Output data in text format (must precede --generate).", + Option + "b" + ["binary"] + (NoArg $ Right $ \opts -> opts {format = Binary}) + "Output data in binary Futhark format (must precede --generate).", + Option + "t" + ["type"] + (NoArg $ Right $ \opts -> opts {format = Type}) + "Output the type (textually) rather than the value (must precede --generate).", + setRangeOption "i8" seti8Range, + setRangeOption "i16" seti16Range, + setRangeOption "i32" seti32Range, + setRangeOption "i64" seti64Range, + setRangeOption "u8" setu8Range, + setRangeOption "u16" setu16Range, + setRangeOption "u32" setu32Range, + setRangeOption "u64" setu64Range, + setRangeOption "f32" setf32Range, + setRangeOption "f64" setf64Range ] -setRangeOption :: Read a => String - -> (Range a -> RandomConfiguration -> RandomConfiguration) - -> FunOptDescr DataOptions +setRangeOption :: + Read a => + String -> + (Range a -> RandomConfiguration -> RandomConfiguration) -> + FunOptDescr DataOptions setRangeOption tname set = - Option "" [name] - (ReqArg (\b -> - let (lower,rest) = span (/=':') b + Option + "" + [name] + ( ReqArg + ( \b -> + let (lower, rest) = span (/= ':') b upper = drop 1 rest - in case (reads lower, reads upper) of - ([(lower', "")], [(upper', "")]) -> - Right $ \config -> - config { optRange = set (lower', upper') $ optRange config } - _ -> - Left $ do - hPutStrLn stderr $ "Invalid bounds for " ++ tname ++ ": " ++ b - exitFailure - ) - "MIN:MAX") $ - "Range of " ++ tname ++ " values." - where name = tname ++ "-bounds" + in case (reads lower, reads upper) of + ([(lower', "")], [(upper', "")]) -> + Right $ \config -> + config {optRange = set (lower', upper') $ optRange config} + _ -> + Left $ do + hPutStrLn stderr $ "Invalid bounds for " ++ tname ++ ": " ++ b + exitFailure + ) + "MIN:MAX" + ) + $ "Range of " ++ tname ++ " values." + where + name = tname ++ "-bounds" -tryMakeGenerator :: String - -> Either String (RandomConfiguration -> OutputFormat -> Word64 -> IO ()) +tryMakeGenerator :: + String -> + Either String (RandomConfiguration -> OutputFormat -> Word64 -> IO ()) tryMakeGenerator t | Just vs <- readValues $ BS.pack t = - return $ \_ fmt _ -> mapM_ (putValue fmt) vs + return $ \_ fmt _ -> mapM_ (putValue fmt) vs | otherwise = do - t' <- toValueType =<< either (Left . show) Right (parseType name (T.pack t)) - return $ \conf fmt seed -> do - let v = randomValue conf t' seed - putValue fmt v - where name = "option " ++ t - putValue Text = putStrLn . pretty - putValue Binary = BS.putStr . Bin.encode - putValue Type = putStrLn . pretty . valueType + t' <- toValueType =<< either (Left . show) Right (parseType name (T.pack t)) + return $ \conf fmt seed -> do + let v = randomValue conf t' seed + putValue fmt v + where + name = "option " ++ t + putValue Text = putStrLn . pretty + putValue Binary = BS.putStr . Bin.encode + putValue Type = putStrLn . pretty . valueType toValueType :: UncheckedTypeExp -> Either String ValueType -toValueType TETuple{} = Left "Cannot handle tuples yet." -toValueType TERecord{} = Left "Cannot handle records yet." -toValueType TEApply{} = Left "Cannot handle type applications yet." -toValueType TEArrow{} = Left "Cannot generate functions." -toValueType TESum{} = Left "Cannot handle sumtypes yet." +toValueType TETuple {} = Left "Cannot handle tuples yet." +toValueType TERecord {} = Left "Cannot handle records yet." +toValueType TEApply {} = Left "Cannot handle type applications yet." +toValueType TEArrow {} = Left "Cannot generate functions." +toValueType TESum {} = Left "Cannot handle sumtypes yet." toValueType (TEUnique t _) = toValueType t toValueType (TEArray t d _) = do d' <- constantDim d ValueType ds t' <- toValueType t - return $ ValueType (d':ds) t' - where constantDim (DimExpConst k _) = Right k - constantDim _ = Left "Array has non-constant dimension declaration." + return $ ValueType (d' : ds) t' + where + constantDim (DimExpConst k _) = Right k + constantDim _ = Left "Array has non-constant dimension declaration." toValueType (TEVar (QualName [] v) _) | Just t <- M.lookup v namesToPrimTypes = Right $ ValueType [] t toValueType (TEVar v _) = @@ -167,70 +206,88 @@ type Range a = (a, a) data RandomConfiguration = RandomConfiguration - { i8Range :: Range Int8 - , i16Range :: Range Int16 - , i32Range :: Range Int32 - , i64Range :: Range Int64 - , u8Range :: Range Word8 - , u16Range :: Range Word16 - , u32Range :: Range Word32 - , u64Range :: Range Word64 - , f32Range :: Range Float - , f64Range :: Range Double - } + { i8Range :: Range Int8, + i16Range :: Range Int16, + i32Range :: Range Int32, + i64Range :: Range Int64, + u8Range :: Range Word8, + u16Range :: Range Word16, + u32Range :: Range Word32, + u64Range :: Range Word64, + f32Range :: Range Float, + f64Range :: Range Double + } -- The following lines provide evidence about how Haskells record -- system sucks. seti8Range :: Range Int8 -> RandomConfiguration -> RandomConfiguration -seti8Range bounds config = config { i8Range = bounds } +seti8Range bounds config = config {i8Range = bounds} + seti16Range :: Range Int16 -> RandomConfiguration -> RandomConfiguration -seti16Range bounds config = config { i16Range = bounds } +seti16Range bounds config = config {i16Range = bounds} + seti32Range :: Range Int32 -> RandomConfiguration -> RandomConfiguration -seti32Range bounds config = config { i32Range = bounds } +seti32Range bounds config = config {i32Range = bounds} + seti64Range :: Range Int64 -> RandomConfiguration -> RandomConfiguration -seti64Range bounds config = config { i64Range = bounds } +seti64Range bounds config = config {i64Range = bounds} setu8Range :: Range Word8 -> RandomConfiguration -> RandomConfiguration -setu8Range bounds config = config { u8Range = bounds } +setu8Range bounds config = config {u8Range = bounds} + setu16Range :: Range Word16 -> RandomConfiguration -> RandomConfiguration -setu16Range bounds config = config { u16Range = bounds } +setu16Range bounds config = config {u16Range = bounds} + setu32Range :: Range Word32 -> RandomConfiguration -> RandomConfiguration -setu32Range bounds config = config { u32Range = bounds } +setu32Range bounds config = config {u32Range = bounds} + setu64Range :: Range Word64 -> RandomConfiguration -> RandomConfiguration -setu64Range bounds config = config { u64Range = bounds } +setu64Range bounds config = config {u64Range = bounds} setf32Range :: Range Float -> RandomConfiguration -> RandomConfiguration -setf32Range bounds config = config { f32Range = bounds } +setf32Range bounds config = config {f32Range = bounds} + setf64Range :: Range Double -> RandomConfiguration -> RandomConfiguration -setf64Range bounds config = config { f64Range = bounds } +setf64Range bounds config = config {f64Range = bounds} initialRandomConfiguration :: RandomConfiguration -initialRandomConfiguration = RandomConfiguration - (minBound, maxBound) (minBound, maxBound) (minBound, maxBound) (minBound, maxBound) - (minBound, maxBound) (minBound, maxBound) (minBound, maxBound) (minBound, maxBound) - (0.0, 1.0) (0.0, 1.0) +initialRandomConfiguration = + RandomConfiguration + (minBound, maxBound) + (minBound, maxBound) + (minBound, maxBound) + (minBound, maxBound) + (minBound, maxBound) + (minBound, maxBound) + (minBound, maxBound) + (minBound, maxBound) + (0.0, 1.0) + (0.0, 1.0) randomValue :: RandomConfiguration -> ValueType -> Word64 -> Value randomValue conf (ValueType ds t) seed = case t of - Signed Int8 -> gen i8Range Int8Value + Signed Int8 -> gen i8Range Int8Value Signed Int16 -> gen i16Range Int16Value Signed Int32 -> gen i32Range Int32Value Signed Int64 -> gen i64Range Int64Value - Unsigned Int8 -> gen u8Range Word8Value + Unsigned Int8 -> gen u8Range Word8Value Unsigned Int16 -> gen u16Range Word16Value Unsigned Int32 -> gen u32Range Word32Value Unsigned Int64 -> gen u64Range Word64Value FloatType Float32 -> gen f32Range Float32Value FloatType Float64 -> gen f64Range Float64Value - Bool -> gen (const (False,True)) BoolValue - where gen range final = randomVector (range conf) final ds seed + Bool -> gen (const (False, True)) BoolValue + where + gen range final = randomVector (range conf) final ds seed -randomVector :: (UMVec.Unbox v, Variate v) => - Range v - -> (UVec.Vector Int -> UVec.Vector v -> Value) - -> [Int] -> Word64 - -> Value +randomVector :: + (UMVec.Unbox v, Variate v) => + Range v -> + (UVec.Vector Int -> UVec.Vector v -> Value) -> + [Int] -> + Word64 -> + Value randomVector range final ds seed = runST $ do -- USe some nice impure computation where we can preallocate a -- vector of the desired size, populate it via the random number @@ -239,10 +296,11 @@ g <- initialize 6364136223846793006 seed let fill i | i < n = do - v <- uniformR range g - UMVec.write arr i v - fill $! i+1 + v <- uniformR range g + UMVec.write arr i v + fill $! i + 1 | otherwise = - final (UVec.fromList ds) <$> freeze arr + final (UVec.fromList ds) <$> freeze arr fill 0 - where n = product ds + where + n = product ds 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 @@ -1,96 +1,103 @@ +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeApplications #-} + -- | Futhark Compiler Driver module Futhark.CLI.Dev (main) where -import Data.Maybe -import Data.List (intersperse) import Control.Category (id) import Control.Monad import Control.Monad.State +import qualified Data.ByteString.Lazy as ByteString +import Data.List (intersperse) +import Data.Maybe import qualified Data.Text.IO as T -import System.IO -import System.Exit -import System.Console.GetOpt -import System.FilePath - -import Prelude hiding (id) - -import Futhark.Compiler.CLI -import Futhark.Pass import Futhark.Actions -import Language.Futhark.Parser (parseFuthark) -import Futhark.Util.Options -import qualified Futhark.IR.SOACS as SOACS +import Futhark.Analysis.Metrics (OpMetrics) +import Futhark.Compiler.CLI +import Futhark.IR (ASTLore, Op, Prog, pretty) import qualified Futhark.IR.Kernels as Kernels -import qualified Futhark.IR.Seq as Seq import qualified Futhark.IR.KernelsMem as KernelsMem +import Futhark.IR.Prop.Aliases (CanBeAliased) +import qualified Futhark.IR.SOACS as SOACS +import qualified Futhark.IR.Seq as Seq import qualified Futhark.IR.SeqMem as SeqMem -import Futhark.IR (Prog, pretty) -import Futhark.TypeCheck (Checkable) -import qualified Futhark.Util.Pretty as PP - +import Futhark.Internalise.Defunctionalise as Defunctionalise import Futhark.Internalise.Defunctorise as Defunctorise import Futhark.Internalise.Monomorphise as Monomorphise -import Futhark.Internalise.Defunctionalise as Defunctionalise -import Futhark.Optimise.InliningDeadFun import Futhark.Optimise.CSE +import Futhark.Optimise.DoubleBuffer import Futhark.Optimise.Fusion -import Futhark.Pass.FirstOrderTransform -import Futhark.Pass.Simplify import Futhark.Optimise.InPlaceLowering -import Futhark.Optimise.DoubleBuffer +import Futhark.Optimise.InliningDeadFun import Futhark.Optimise.Sink import Futhark.Optimise.TileLoops import Futhark.Optimise.Unstream -import Futhark.Pass.KernelBabysitting -import Futhark.Pass.ExtractKernels +import Futhark.Pass import Futhark.Pass.ExpandAllocations import qualified Futhark.Pass.ExplicitAllocations.Kernels as Kernels import qualified Futhark.Pass.ExplicitAllocations.Seq as Seq +import Futhark.Pass.ExtractKernels +import Futhark.Pass.FirstOrderTransform +import Futhark.Pass.KernelBabysitting +import Futhark.Pass.Simplify import Futhark.Passes +import Futhark.TypeCheck (Checkable) import Futhark.Util.Log +import Futhark.Util.Options +import qualified Futhark.Util.Pretty as PP +import Language.Futhark.Parser (parseFuthark) +import qualified Language.SexpGrammar as Sexp +import System.Console.GetOpt +import System.Exit +import System.FilePath +import System.IO +import Prelude hiding (id) -- | What to do with the program after it has been read. -data FutharkPipeline = PrettyPrint - -- ^ Just print it. - | TypeCheck - -- ^ Run the type checker; print type errors. - | Pipeline [UntypedPass] - -- ^ Run this pipeline. - | Defunctorise - -- ^ Partially evaluate away the module language. - | Monomorphise - -- ^ Defunctorise and monomorphise. - | Defunctionalise - -- ^ Defunctorise, monomorphise, and defunctionalise. - -data Config = Config { futharkConfig :: FutharkConfig - , futharkPipeline :: FutharkPipeline - -- ^ Nothing is distinct from a empty pipeline - - -- it means we don't even run the internaliser. - , futharkAction :: UntypedAction - , futharkPrintAST :: Bool - -- ^ If true, prints programs as raw ASTs instead - -- of their prettyprinted form. - } +data FutharkPipeline + = -- | Just print it. + PrettyPrint + | -- | Run the type checker; print type errors. + TypeCheck + | -- | Run this pipeline. + Pipeline [UntypedPass] + | -- | Partially evaluate away the module language. + Defunctorise + | -- | Defunctorise and monomorphise. + Monomorphise + | -- | Defunctorise, monomorphise, and defunctionalise. + Defunctionalise +data Config = Config + { futharkConfig :: FutharkConfig, + -- | Nothing is distinct from a empty pipeline - + -- it means we don't even run the internaliser. + futharkPipeline :: FutharkPipeline, + futharkAction :: UntypedAction, + -- | If true, prints programs as raw ASTs instead + -- of their prettyprinted form. + futharkPrintAST :: Bool + } -- | Get a Futhark pipeline from the configuration - an empty one if -- none exists. getFutharkPipeline :: Config -> [UntypedPass] getFutharkPipeline = toPipeline . futharkPipeline - where toPipeline (Pipeline p) = p - toPipeline _ = [] + where + toPipeline (Pipeline p) = p + toPipeline _ = [] -data UntypedPassState = SOACS (Prog SOACS.SOACS) - | Kernels (Prog Kernels.Kernels) - | Seq (Prog Seq.Seq) - | KernelsMem (Prog KernelsMem.KernelsMem) - | SeqMem (Prog SeqMem.SeqMem) +data UntypedPassState + = SOACS (Prog SOACS.SOACS) + | Kernels (Prog Kernels.Kernels) + | Seq (Prog Seq.Seq) + | KernelsMem (Prog KernelsMem.KernelsMem) + | SeqMem (Prog SeqMem.SeqMem) getSOACSProg :: UntypedPassState -> Maybe (Prog SOACS.SOACS) getSOACSProg (SOACS prog) = Just prog -getSOACSProg _ = Nothing +getSOACSProg _ = Nothing class Representation s where -- | A human-readable description of the representation expected or @@ -111,382 +118,512 @@ ppr (SeqMem prog) = PP.ppr prog ppr (KernelsMem prog) = PP.ppr prog -newtype UntypedPass = UntypedPass (UntypedPassState - -> PipelineConfig - -> FutharkM UntypedPassState) - -data AllActions = AllActions - { actionSOACS :: Action SOACS.SOACS - , actionKernels :: Action Kernels.Kernels - , actionSeq :: Action Seq.Seq - , actionKernelsMem :: Action KernelsMem.KernelsMem - , actionSeqMem :: Action SeqMem.SeqMem - } +newtype UntypedPass + = UntypedPass + ( UntypedPassState -> + PipelineConfig -> + FutharkM UntypedPassState + ) -data UntypedAction = SOACSAction (Action SOACS.SOACS) - | KernelsAction (Action Kernels.Kernels) - | KernelsMemAction (FilePath -> Action KernelsMem.KernelsMem) - | SeqMemAction (FilePath -> Action SeqMem.SeqMem) - | PolyAction AllActions +data UntypedAction + = SOACSAction (Action SOACS.SOACS) + | KernelsAction (Action Kernels.Kernels) + | KernelsMemAction (FilePath -> Action KernelsMem.KernelsMem) + | SeqMemAction (FilePath -> Action SeqMem.SeqMem) + | PolyAction + ( forall lore. + ( ASTLore lore, + (CanBeAliased (Op lore)), + (OpMetrics (Op lore)) + ) => + Action lore + ) untypedActionName :: UntypedAction -> String untypedActionName (SOACSAction a) = actionName a untypedActionName (KernelsAction a) = actionName a untypedActionName (SeqMemAction a) = actionName $ a "" untypedActionName (KernelsMemAction a) = actionName $ a "" -untypedActionName (PolyAction a) = actionName (actionSOACS a) +untypedActionName (PolyAction a) = actionName (a :: Action SOACS.SOACS) instance Representation UntypedAction where representation (SOACSAction _) = "SOACS" representation (KernelsAction _) = "Kernels" representation (KernelsMemAction _) = "KernelsMem" representation (SeqMemAction _) = "SeqMem" - representation PolyAction{} = "" + representation PolyAction {} = "" newConfig :: Config newConfig = Config newFutharkConfig (Pipeline []) action False - where action = PolyAction $ AllActions printAction printAction printAction printAction printAction + where + action = PolyAction printAction -changeFutharkConfig :: (FutharkConfig -> FutharkConfig) - -> Config -> Config -changeFutharkConfig f cfg = cfg { futharkConfig = f $ futharkConfig cfg } +changeFutharkConfig :: + (FutharkConfig -> FutharkConfig) -> + Config -> + Config +changeFutharkConfig f cfg = cfg {futharkConfig = f $ futharkConfig cfg} type FutharkOption = FunOptDescr Config passOption :: String -> UntypedPass -> String -> [String] -> FutharkOption passOption desc pass short long = - Option short long - (NoArg $ Right $ \cfg -> - cfg { futharkPipeline = Pipeline $ getFutharkPipeline cfg ++ [pass] }) - desc + Option + short + long + ( NoArg $ + Right $ \cfg -> + cfg {futharkPipeline = Pipeline $ getFutharkPipeline cfg ++ [pass]} + ) + desc -kernelsMemProg :: String -> UntypedPassState - -> FutharkM (Prog KernelsMem.KernelsMem) +kernelsMemProg :: + String -> + UntypedPassState -> + FutharkM (Prog KernelsMem.KernelsMem) kernelsMemProg _ (KernelsMem prog) = return prog kernelsMemProg name rep = - externalErrorS $ "Pass " ++ name ++ - " expects KernelsMem representation, but got " ++ representation rep + externalErrorS $ + "Pass " ++ name + ++ " expects KernelsMem representation, but got " + ++ representation rep soacsProg :: String -> UntypedPassState -> FutharkM (Prog SOACS.SOACS) soacsProg _ (SOACS prog) = return prog soacsProg name rep = - externalErrorS $ "Pass " ++ name ++ - " expects SOACS representation, but got " ++ representation rep + externalErrorS $ + "Pass " ++ name + ++ " expects SOACS representation, but got " + ++ representation rep kernelsProg :: String -> UntypedPassState -> FutharkM (Prog Kernels.Kernels) kernelsProg _ (Kernels prog) = return prog kernelsProg name rep = externalErrorS $ - "Pass " ++ name ++" expects Kernels representation, but got " ++ representation rep + "Pass " ++ name ++ " expects Kernels representation, but got " ++ representation rep -typedPassOption :: Checkable tolore => - (String -> UntypedPassState -> FutharkM (Prog fromlore)) - -> (Prog tolore -> UntypedPassState) - -> Pass fromlore tolore - -> String - -> FutharkOption +typedPassOption :: + Checkable tolore => + (String -> UntypedPassState -> FutharkM (Prog fromlore)) -> + (Prog tolore -> UntypedPassState) -> + Pass fromlore tolore -> + String -> + FutharkOption typedPassOption getProg putProg pass short = passOption (passDescription pass) (UntypedPass perform) short long - where perform s config = do - prog <- getProg (passName pass) s - putProg <$> runPipeline (onePass pass) config prog + where + perform s config = do + prog <- getProg (passName pass) s + putProg <$> runPipeline (onePass pass) config prog - long = [passLongOption pass] + long = [passLongOption pass] soacsPassOption :: Pass SOACS.SOACS SOACS.SOACS -> String -> FutharkOption soacsPassOption = typedPassOption soacsProg SOACS -kernelsPassOption :: Pass Kernels.Kernels Kernels.Kernels - -> String -> FutharkOption +kernelsPassOption :: + Pass Kernels.Kernels Kernels.Kernels -> + String -> + FutharkOption kernelsPassOption = typedPassOption kernelsProg Kernels -kernelsMemPassOption :: Pass KernelsMem.KernelsMem KernelsMem.KernelsMem - -> String -> FutharkOption +kernelsMemPassOption :: + Pass KernelsMem.KernelsMem KernelsMem.KernelsMem -> + String -> + FutharkOption kernelsMemPassOption = typedPassOption kernelsMemProg KernelsMem simplifyOption :: String -> FutharkOption simplifyOption short = passOption (passDescription pass) (UntypedPass perform) short long - where perform (SOACS prog) config = - SOACS <$> runPipeline (onePass simplifySOACS) config prog - perform (Kernels prog) config = - Kernels <$> runPipeline (onePass simplifyKernels) config prog - perform (Seq prog) config = - Seq <$> runPipeline (onePass simplifySeq) config prog - perform (SeqMem prog) config = - SeqMem <$> runPipeline (onePass simplifySeqMem) config prog - perform (KernelsMem prog) config = - KernelsMem <$> runPipeline (onePass simplifyKernelsMem) config prog + where + perform (SOACS prog) config = + SOACS <$> runPipeline (onePass simplifySOACS) config prog + perform (Kernels prog) config = + Kernels <$> runPipeline (onePass simplifyKernels) config prog + perform (Seq prog) config = + Seq <$> runPipeline (onePass simplifySeq) config prog + perform (SeqMem prog) config = + SeqMem <$> runPipeline (onePass simplifySeqMem) config prog + perform (KernelsMem prog) config = + KernelsMem <$> runPipeline (onePass simplifyKernelsMem) config prog - long = [passLongOption pass] - pass = simplifySOACS + long = [passLongOption pass] + pass = simplifySOACS allocateOption :: String -> FutharkOption allocateOption short = passOption (passDescription pass) (UntypedPass perform) short long - where perform (Kernels prog) config = - KernelsMem <$> - runPipeline (onePass Kernels.explicitAllocations) config prog - perform (Seq prog) config = - SeqMem <$> - runPipeline (onePass Seq.explicitAllocations) config prog - perform s _ = - externalErrorS $ - "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s + where + perform (Kernels prog) config = + KernelsMem + <$> runPipeline (onePass Kernels.explicitAllocations) config prog + perform (Seq prog) config = + SeqMem + <$> runPipeline (onePass Seq.explicitAllocations) config prog + perform s _ = + externalErrorS $ + "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s - long = [passLongOption pass] - pass = Seq.explicitAllocations + long = [passLongOption pass] + pass = Seq.explicitAllocations iplOption :: String -> FutharkOption iplOption short = passOption (passDescription pass) (UntypedPass perform) short long - where perform (Kernels prog) config = - Kernels <$> - runPipeline (onePass inPlaceLoweringKernels) config prog - perform (Seq prog) config = - Seq <$> - runPipeline (onePass inPlaceLoweringSeq) config prog - perform s _ = - externalErrorS $ - "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s + where + perform (Kernels prog) config = + Kernels + <$> runPipeline (onePass inPlaceLoweringKernels) config prog + perform (Seq prog) config = + Seq + <$> runPipeline (onePass inPlaceLoweringSeq) config prog + perform s _ = + externalErrorS $ + "Pass '" ++ passDescription pass ++ "' cannot operate on " ++ representation s - long = [passLongOption pass] - pass = inPlaceLoweringSeq + long = [passLongOption pass] + pass = inPlaceLoweringSeq cseOption :: String -> FutharkOption cseOption short = passOption (passDescription pass) (UntypedPass perform) short long - where perform (SOACS prog) config = - SOACS <$> runPipeline (onePass $ performCSE True) config prog - perform (Kernels prog) config = - Kernels <$> runPipeline (onePass $ performCSE True) config prog - perform (Seq prog) config = - Seq <$> runPipeline (onePass $ performCSE True) config prog - perform (SeqMem prog) config = - SeqMem <$> runPipeline (onePass $ performCSE False) config prog - perform (KernelsMem prog) config = - KernelsMem <$> runPipeline (onePass $ performCSE False) config prog + where + perform (SOACS prog) config = + SOACS <$> runPipeline (onePass $ performCSE True) config prog + perform (Kernels prog) config = + Kernels <$> runPipeline (onePass $ performCSE True) config prog + perform (Seq prog) config = + Seq <$> runPipeline (onePass $ performCSE True) config prog + perform (SeqMem prog) config = + SeqMem <$> runPipeline (onePass $ performCSE False) config prog + perform (KernelsMem prog) config = + KernelsMem <$> runPipeline (onePass $ performCSE False) config prog - long = [passLongOption pass] - pass = performCSE True :: Pass SOACS.SOACS SOACS.SOACS + long = [passLongOption pass] + pass = performCSE True :: Pass SOACS.SOACS SOACS.SOACS -pipelineOption :: (UntypedPassState -> Maybe (Prog fromlore)) - -> String - -> (Prog tolore -> UntypedPassState) - -> String - -> Pipeline fromlore tolore - -> String - -> [String] - -> FutharkOption +pipelineOption :: + (UntypedPassState -> Maybe (Prog fromlore)) -> + String -> + (Prog tolore -> UntypedPassState) -> + String -> + Pipeline fromlore tolore -> + String -> + [String] -> + FutharkOption pipelineOption getprog repdesc repf desc pipeline = passOption desc $ UntypedPass pipelinePass - where pipelinePass rep config = - case getprog rep of - Just prog -> - repf <$> runPipeline pipeline config prog - Nothing -> - externalErrorS $ "Expected " ++ repdesc ++ " representation, but got " ++ - representation rep + where + pipelinePass rep config = + case getprog rep of + Just prog -> + repf <$> runPipeline pipeline config prog + Nothing -> + externalErrorS $ + "Expected " ++ repdesc ++ " representation, but got " + ++ representation rep -soacsPipelineOption :: String -> Pipeline SOACS.SOACS SOACS.SOACS -> String -> [String] - -> FutharkOption +soacsPipelineOption :: + String -> + Pipeline SOACS.SOACS SOACS.SOACS -> + String -> + [String] -> + FutharkOption soacsPipelineOption = pipelineOption getSOACSProg "SOACS" SOACS commandLineOptions :: [FutharkOption] commandLineOptions = - [ Option "v" ["verbose"] - (OptArg (Right . changeFutharkConfig . incVerbosity) "FILE") - "Print verbose output on standard error; wrong program to FILE." - , Option [] ["Werror"] - (NoArg $ Right $ changeFutharkConfig $ \opts -> opts { futharkWerror = True }) - "Treat warnings as errors." - - , Option "t" ["type-check"] - (NoArg $ Right $ \opts -> - opts { futharkPipeline = TypeCheck }) - "Print on standard output the type-checked program." - - , Option [] ["pretty-print"] - (NoArg $ Right $ \opts -> - opts { futharkPipeline = PrettyPrint }) - "Parse and pretty-print the AST of the given program." - - , Option [] ["compile-imperative"] - (NoArg $ Right $ \opts -> - opts { futharkAction = SeqMemAction $ const impCodeGenAction }) - "Translate program into the imperative IL and write it on standard output." - , Option [] ["compile-imperative-kernels"] - (NoArg $ Right $ \opts -> - opts { futharkAction = KernelsMemAction $ const kernelImpCodeGenAction }) - "Translate program into the imperative IL with kernels and write it on standard output." - , Option [] ["compile-opencl"] - (NoArg $ Right $ \opts -> - opts { futharkAction = KernelsMemAction $ compileOpenCLAction newFutharkConfig ToExecutable }) - "Compile the program using the OpenCL backend." - , Option [] ["compile-c"] - (NoArg $ Right $ \opts -> - opts { futharkAction = SeqMemAction $ compileCAction newFutharkConfig ToExecutable }) - "Compile the program using the C backend." - , Option "p" ["print"] - (NoArg $ Right $ \opts -> - opts { futharkAction = PolyAction $ AllActions printAction printAction printAction printAction printAction }) - "Prettyprint the resulting internal representation on standard output (default action)." - , Option "m" ["metrics"] - (NoArg $ Right $ \opts -> - opts { futharkAction = PolyAction $ AllActions metricsAction metricsAction metricsAction metricsAction metricsAction }) - "Print AST metrics of the resulting internal representation on standard output." - , Option [] ["defunctorise"] - (NoArg $ Right $ \opts -> opts { futharkPipeline = Defunctorise }) - "Partially evaluate all module constructs and print the residual program." - , Option [] ["monomorphise"] - (NoArg $ Right $ \opts -> opts { futharkPipeline = Monomorphise }) - "Monomorphise the program." - , Option [] ["defunctionalise"] - (NoArg $ Right $ \opts -> opts { futharkPipeline = Defunctionalise }) - "Defunctionalise the program." - , Option [] ["ast"] - (NoArg $ Right $ \opts -> opts { futharkPrintAST = True }) - "Output ASTs instead of prettyprinted programs." - , Option [] ["safe"] - (NoArg $ Right $ changeFutharkConfig $ \opts -> opts { futharkSafe = True }) - "Ignore 'unsafe'." - , typedPassOption soacsProg Seq firstOrderTransform "f" - , soacsPassOption fuseSOACs "o" - , soacsPassOption inlineFunctions [] - , kernelsPassOption babysitKernels [] - , kernelsPassOption tileLoops [] - , kernelsPassOption unstream [] - , kernelsPassOption sink [] - , typedPassOption soacsProg Kernels extractKernels [] - - , iplOption [] - , allocateOption "a" - - , kernelsMemPassOption doubleBuffer [] - , kernelsMemPassOption expandAllocations [] - - , cseOption [] - , simplifyOption "e" - - , soacsPipelineOption "Run the default optimised pipeline" - standardPipeline "s" ["standard"] - , pipelineOption getSOACSProg "Kernels" Kernels - "Run the default optimised kernels pipeline" - kernelsPipeline [] ["kernels"] - , pipelineOption getSOACSProg "KernelsMem" KernelsMem - "Run the full GPU compilation pipeline" - gpuPipeline [] ["gpu"] - , pipelineOption getSOACSProg "KernelsMem" SeqMem - "Run the sequential CPU compilation pipeline" - sequentialCpuPipeline [] ["cpu"] + [ Option + "v" + ["verbose"] + (OptArg (Right . changeFutharkConfig . incVerbosity) "FILE") + "Print verbose output on standard error; wrong program to FILE.", + Option + [] + ["Werror"] + (NoArg $ Right $ changeFutharkConfig $ \opts -> opts {futharkWerror = True}) + "Treat warnings as errors.", + Option + "t" + ["type-check"] + ( NoArg $ + Right $ \opts -> + opts {futharkPipeline = TypeCheck} + ) + "Print on standard output the type-checked program.", + Option + [] + ["pretty-print"] + ( NoArg $ + Right $ \opts -> + opts {futharkPipeline = PrettyPrint} + ) + "Parse and pretty-print the AST of the given program.", + Option + [] + ["compile-imperative"] + ( NoArg $ + Right $ \opts -> + opts {futharkAction = SeqMemAction $ const impCodeGenAction} + ) + "Translate program into the imperative IL and write it on standard output.", + Option + [] + ["compile-imperative-kernels"] + ( NoArg $ + Right $ \opts -> + opts {futharkAction = KernelsMemAction $ const kernelImpCodeGenAction} + ) + "Translate program into the imperative IL with kernels and write it on standard output.", + Option + [] + ["compile-opencl"] + ( NoArg $ + Right $ \opts -> + opts {futharkAction = KernelsMemAction $ compileOpenCLAction newFutharkConfig ToExecutable} + ) + "Compile the program using the OpenCL backend.", + Option + [] + ["compile-c"] + ( NoArg $ + Right $ \opts -> + opts {futharkAction = SeqMemAction $ compileCAction newFutharkConfig ToExecutable} + ) + "Compile the program using the C backend.", + Option + "p" + ["print"] + (NoArg $ Right $ \opts -> opts {futharkAction = PolyAction printAction}) + "Prettyprint the resulting internal representation on standard output (default action).", + Option + "m" + ["metrics"] + (NoArg $ Right $ \opts -> opts {futharkAction = PolyAction metricsAction}) + "Print AST metrics of the resulting internal representation on standard output.", + Option + [] + ["sexp"] + (NoArg $ Right $ \opts -> opts {futharkAction = PolyAction sexpAction}) + "Print the resulting IR as S-expressions to standard output.", + Option + [] + ["defunctorise"] + (NoArg $ Right $ \opts -> opts {futharkPipeline = Defunctorise}) + "Partially evaluate all module constructs and print the residual program.", + Option + [] + ["monomorphise"] + (NoArg $ Right $ \opts -> opts {futharkPipeline = Monomorphise}) + "Monomorphise the program.", + Option + [] + ["defunctionalise"] + (NoArg $ Right $ \opts -> opts {futharkPipeline = Defunctionalise}) + "Defunctionalise the program.", + Option + [] + ["ast"] + (NoArg $ Right $ \opts -> opts {futharkPrintAST = True}) + "Output ASTs instead of prettyprinted programs.", + Option + [] + ["safe"] + (NoArg $ Right $ changeFutharkConfig $ \opts -> opts {futharkSafe = True}) + "Ignore 'unsafe'.", + typedPassOption soacsProg Seq firstOrderTransform "f", + soacsPassOption fuseSOACs "o", + soacsPassOption inlineFunctions [], + kernelsPassOption babysitKernels [], + kernelsPassOption tileLoops [], + kernelsPassOption unstream [], + kernelsPassOption sink [], + typedPassOption soacsProg Kernels extractKernels [], + iplOption [], + allocateOption "a", + kernelsMemPassOption doubleBuffer [], + kernelsMemPassOption expandAllocations [], + cseOption [], + simplifyOption "e", + soacsPipelineOption + "Run the default optimised pipeline" + standardPipeline + "s" + ["standard"], + pipelineOption + getSOACSProg + "Kernels" + Kernels + "Run the default optimised kernels pipeline" + kernelsPipeline + [] + ["kernels"], + pipelineOption + getSOACSProg + "KernelsMem" + KernelsMem + "Run the full GPU compilation pipeline" + gpuPipeline + [] + ["gpu"], + pipelineOption + getSOACSProg + "KernelsMem" + SeqMem + "Run the sequential CPU compilation pipeline" + sequentialCpuPipeline + [] + ["cpu"] ] incVerbosity :: Maybe FilePath -> FutharkConfig -> FutharkConfig incVerbosity file cfg = - cfg { futharkVerbose = (v, file `mplus` snd (futharkVerbose cfg)) } - where v = case fst $ futharkVerbose cfg of - NotVerbose -> Verbose - Verbose -> VeryVerbose - VeryVerbose -> VeryVerbose + cfg {futharkVerbose = (v, file `mplus` snd (futharkVerbose cfg))} + where + v = case fst $ futharkVerbose cfg of + NotVerbose -> Verbose + Verbose -> VeryVerbose + VeryVerbose -> VeryVerbose -- | Entry point. Non-interactive, except when reading interpreter -- input from standard input. main :: String -> [String] -> IO () main = mainWithOptions newConfig commandLineOptions "options... program" compile - where compile [file] config = - Just $ do - res <- runFutharkM (m file config) $ - fst $ futharkVerbose $ futharkConfig config - case res of - Left err -> do - dumpError (futharkConfig config) err - exitWith $ ExitFailure 2 - Right () -> return () - compile _ _ = - Nothing - m file config = do - let p :: (Show a, PP.Pretty a) => [a] -> IO () - p = mapM_ putStrLn . - intersperse "" . - map (if futharkPrintAST config then show else pretty) + where + compile [file] config = + Just $ do + res <- + runFutharkM (m file config) $ + fst $ futharkVerbose $ futharkConfig config + case res of + Left err -> do + dumpError (futharkConfig config) err + exitWith $ ExitFailure 2 + Right () -> return () + compile _ _ = + Nothing + m file config = do + let p :: (Show a, PP.Pretty a) => [a] -> IO () + p = + mapM_ putStrLn + . intersperse "" + . map (if futharkPrintAST config then show else pretty) - case futharkPipeline config of - PrettyPrint -> liftIO $ do - maybe_prog <- parseFuthark file <$> T.readFile file - case maybe_prog of - Left err -> fail $ show err - Right prog | futharkPrintAST config -> print prog - | otherwise -> putStrLn $ pretty prog - TypeCheck -> do - (_, imports, _) <- readProgram file - liftIO $ forM_ (map snd imports) $ \fm -> - putStrLn $ if futharkPrintAST config - then show $ fileProg fm - else pretty $ fileProg fm - Defunctorise -> do - (_, imports, src) <- readProgram file - liftIO $ p $ evalState (Defunctorise.transformProg imports) src - Monomorphise -> do - (_, imports, src) <- readProgram file - liftIO $ p $ - flip evalState src $ + case futharkPipeline config of + PrettyPrint -> liftIO $ do + maybe_prog <- parseFuthark file <$> T.readFile file + case maybe_prog of + Left err -> fail $ show err + Right prog + | futharkPrintAST config -> print prog + | otherwise -> putStrLn $ pretty prog + TypeCheck -> do + (_, imports, _) <- readProgram file + liftIO $ + forM_ (map snd imports) $ \fm -> + putStrLn $ + if futharkPrintAST config + then show $ fileProg fm + else pretty $ fileProg fm + Defunctorise -> do + (_, imports, src) <- readProgram file + liftIO $ p $ evalState (Defunctorise.transformProg imports) src + Monomorphise -> do + (_, imports, src) <- readProgram file + liftIO $ + p $ + flip evalState src $ Defunctorise.transformProg imports - >>= Monomorphise.transformProg - Defunctionalise -> do - (_, imports, src) <- readProgram file - liftIO $ p $ - flip evalState src $ + >>= Monomorphise.transformProg + Defunctionalise -> do + (_, imports, src) <- readProgram file + liftIO $ + p $ + flip evalState src $ Defunctorise.transformProg imports - >>= Monomorphise.transformProg - >>= Defunctionalise.transformProg - Pipeline{} -> do + >>= Monomorphise.transformProg + >>= Defunctionalise.transformProg + Pipeline {} -> + case splitExtensions file of + (base, ".fut") -> do prog <- runPipelineOnProgram (futharkConfig config) id file - runPolyPasses config (file `replaceExtension` "") (SOACS prog) + runPolyPasses config base (SOACS prog) + (base, ".sexp") -> do + input <- liftIO $ ByteString.readFile file + prog <- case Sexp.decode @(Prog SOACS.SOACS) input of + Right prog' -> return $ SOACS prog' + Left _ -> + case Sexp.decode @(Prog Kernels.Kernels) input of + Right prog' -> return $ Kernels prog' + Left _ -> + case Sexp.decode @(Prog Seq.Seq) input of + Right prog' -> return $ Seq prog' + Left _ -> + case Sexp.decode @(Prog KernelsMem.KernelsMem) input of + Right prog' -> return $ KernelsMem prog' + Left _ -> + case Sexp.decode @(Prog SeqMem.SeqMem) input of + Right prog' -> return $ SeqMem prog' + Left e -> externalErrorS $ "Couldn't parse sexp input: " ++ show e + runPolyPasses config base prog + (_, ext) -> + externalErrorS $ unwords ["Unsupported extension", show ext, ". Supported extensions: sexp, fut"] runPolyPasses :: Config -> FilePath -> UntypedPassState -> FutharkM () runPolyPasses config base initial_prog = do - end_prog <- foldM (runPolyPass pipeline_config) initial_prog - (getFutharkPipeline config) - logMsg $ "Running action " ++ untypedActionName (futharkAction config) - case (end_prog, futharkAction config) of - (SOACS prog, SOACSAction action) -> - actionProcedure action prog - (Kernels prog, KernelsAction action) -> - actionProcedure action prog - (SeqMem prog, SeqMemAction action) -> - actionProcedure (action base) prog - (KernelsMem prog, KernelsMemAction action) -> - actionProcedure (action base) prog - - (SOACS soacs_prog, PolyAction acs) -> - actionProcedure (actionSOACS acs) soacs_prog - (Kernels kernels_prog, PolyAction acs) -> - actionProcedure (actionKernels acs) kernels_prog - (Seq seq_prog, PolyAction acs) -> - actionProcedure (actionSeq acs) seq_prog - (KernelsMem mem_prog, PolyAction acs) -> - actionProcedure (actionKernelsMem acs) mem_prog - (SeqMem mem_prog, PolyAction acs) -> - actionProcedure (actionSeqMem acs) mem_prog - - (_, action) -> - externalErrorS $ "Action " <> - untypedActionName action <> - " expects " ++ representation action ++ " representation, but got " ++ - representation end_prog ++ "." - logMsg ("Done." :: String) - where pipeline_config = - PipelineConfig { pipelineVerbose = fst (futharkVerbose $ futharkConfig config) > NotVerbose - , pipelineValidate = True - } + end_prog <- + foldM + (runPolyPass pipeline_config) + initial_prog + (getFutharkPipeline config) + logMsg $ "Running action " ++ untypedActionName (futharkAction config) + case (end_prog, futharkAction config) of + (SOACS prog, SOACSAction action) -> + actionProcedure action prog + (Kernels prog, KernelsAction action) -> + actionProcedure action prog + (SeqMem prog, SeqMemAction action) -> + actionProcedure (action base) prog + (KernelsMem prog, KernelsMemAction action) -> + actionProcedure (action base) prog + (SOACS soacs_prog, PolyAction acs) -> + actionProcedure acs soacs_prog + (Kernels kernels_prog, PolyAction acs) -> + actionProcedure acs kernels_prog + (Seq seq_prog, PolyAction acs) -> + actionProcedure acs seq_prog + (KernelsMem mem_prog, PolyAction acs) -> + actionProcedure acs mem_prog + (SeqMem mem_prog, PolyAction acs) -> + actionProcedure acs mem_prog + (_, action) -> + externalErrorS $ + "Action " + <> untypedActionName action + <> " expects " + ++ representation action + ++ " representation, but got " + ++ representation end_prog + ++ "." + logMsg ("Done." :: String) + where + pipeline_config = + PipelineConfig + { pipelineVerbose = fst (futharkVerbose $ futharkConfig config) > NotVerbose, + pipelineValidate = True + } -runPolyPass :: PipelineConfig - -> UntypedPassState -> UntypedPass -> FutharkM UntypedPassState +runPolyPass :: + PipelineConfig -> + UntypedPassState -> + UntypedPass -> + FutharkM UntypedPassState runPolyPass pipeline_config s (UntypedPass f) = f s pipeline_config diff --git a/src/Futhark/CLI/Doc.hs b/src/Futhark/CLI/Doc.hs --- a/src/Futhark/CLI/Doc.hs +++ b/src/Futhark/CLI/Doc.hs @@ -1,105 +1,120 @@ -{-# LANGUAGE TupleSections #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} + -- | @futhark doc@ module Futhark.CLI.Doc (main) where import Control.Monad.State import Data.FileEmbed import Data.List (nubBy) -import System.FilePath -import System.Directory (createDirectoryIfMissing) -import System.Console.GetOpt -import System.IO -import System.Exit import qualified Data.Text.Lazy as T import qualified Data.Text.Lazy.IO as T -import Text.Blaze.Html.Renderer.Text - +import Futhark.Compiler (Imports, dumpError, fileProg, newFutharkConfig, readLibrary) import Futhark.Doc.Generator -import Futhark.Compiler (readLibrary, dumpError, newFutharkConfig, Imports, fileProg) -import Futhark.Pipeline (runFutharkM, FutharkM, Verbosity(..)) -import Language.Futhark.Syntax (progDoc, DocComment(..)) -import Futhark.Util.Options +import Futhark.Pipeline (FutharkM, Verbosity (..), runFutharkM) import Futhark.Util (directoryContents, trim) +import Futhark.Util.Options +import Language.Futhark.Syntax (DocComment (..), progDoc) +import System.Console.GetOpt +import System.Directory (createDirectoryIfMissing) +import System.Exit +import System.FilePath +import System.IO +import Text.Blaze.Html.Renderer.Text -- | Run @futhark doc@. main :: String -> [String] -> IO () main = mainWithOptions initialDocConfig commandLineOptions "options... -o outdir programs..." f - where f [dir] config = Just $ do - res <- runFutharkM (m config dir) Verbose - case res of - Left err -> liftIO $ do - dumpError newFutharkConfig err - exitWith $ ExitFailure 2 - Right () -> - return () - f _ _ = Nothing + where + f [dir] config = Just $ do + res <- runFutharkM (m config dir) Verbose + case res of + Left err -> liftIO $ do + dumpError newFutharkConfig err + exitWith $ ExitFailure 2 + Right () -> + return () + f _ _ = Nothing - m :: DocConfig -> FilePath -> FutharkM () - m config dir = - case docOutput config of - Nothing -> liftIO $ do - hPutStrLn stderr "Must specify output directory with -o." - exitWith $ ExitFailure 1 - Just outdir -> do - files <- liftIO $ futFiles dir - when (docVerbose config) $ liftIO $ do - mapM_ (hPutStrLn stderr . ("Found source file "<>)) files - hPutStrLn stderr "Reading files..." - (_w, imports, _vns) <- readLibrary files - liftIO $ printDecs config outdir files $ nubBy sameImport imports + m :: DocConfig -> FilePath -> FutharkM () + m config dir = + case docOutput config of + Nothing -> liftIO $ do + hPutStrLn stderr "Must specify output directory with -o." + exitWith $ ExitFailure 1 + Just outdir -> do + files <- liftIO $ futFiles dir + when (docVerbose config) $ + liftIO $ do + mapM_ (hPutStrLn stderr . ("Found source file " <>)) files + hPutStrLn stderr "Reading files..." + (_w, imports, _vns) <- readLibrary files + liftIO $ printDecs config outdir files $ nubBy sameImport imports - sameImport (x, _) (y, _) = x == y + sameImport (x, _) (y, _) = x == y futFiles :: FilePath -> IO [FilePath] futFiles dir = filter isFut <$> directoryContents dir - where isFut = (==".fut") . takeExtension + where + isFut = (== ".fut") . takeExtension printDecs :: DocConfig -> FilePath -> [FilePath] -> Imports -> IO () printDecs cfg dir files imports = do let direct_imports = map (normalise . dropExtension) files - (file_htmls, _warnings) = renderFiles direct_imports $ - filter (not . ignored) imports + (file_htmls, _warnings) = + renderFiles direct_imports $ + filter (not . ignored) imports mapM_ (write . fmap renderHtml) file_htmls write ("style.css", cssFile) - - where write :: (String, T.Text) -> IO () - write (name, content) = do let file = dir makeRelative "/" name - when (docVerbose cfg) $ - hPutStrLn stderr $ "Writing " <> file - createDirectoryIfMissing True $ takeDirectory file - T.writeFile file content + where + write :: (String, T.Text) -> IO () + write (name, content) = do + let file = dir makeRelative "/" name + when (docVerbose cfg) $ + hPutStrLn stderr $ "Writing " <> file + createDirectoryIfMissing True $ takeDirectory file + T.writeFile file content - -- Some files are not worth documenting; typically because - -- they contain tests. The current crude mechanism is to - -- recognise them by a file comment containing "ignore". - ignored (_, fm) = - case progDoc (fileProg fm) of - Just (DocComment s _) -> trim s == "ignore" - _ -> False + -- Some files are not worth documenting; typically because + -- they contain tests. The current crude mechanism is to + -- recognise them by a file comment containing "ignore". + ignored (_, fm) = + case progDoc (fileProg fm) of + Just (DocComment s _) -> trim s == "ignore" + _ -> False cssFile :: T.Text cssFile = $(embedStringFile "rts/futhark-doc/style.css") -data DocConfig = DocConfig { docOutput :: Maybe FilePath - , docVerbose :: Bool - } +data DocConfig = DocConfig + { docOutput :: Maybe FilePath, + docVerbose :: Bool + } initialDocConfig :: DocConfig -initialDocConfig = DocConfig { docOutput = Nothing - , docVerbose = False - } +initialDocConfig = + DocConfig + { docOutput = Nothing, + docVerbose = False + } type DocOption = OptDescr (Either (IO ()) (DocConfig -> DocConfig)) commandLineOptions :: [DocOption] -commandLineOptions = [ Option "o" ["output-directory"] - (ReqArg (\dirname -> Right $ \config -> config { docOutput = Just dirname }) - "DIR") - "Directory in which to put generated documentation." - , Option "v" ["verbose"] - (NoArg $ Right $ \config -> config { docVerbose = True }) - "Print status messages on stderr." - ] +commandLineOptions = + [ Option + "o" + ["output-directory"] + ( ReqArg + (\dirname -> Right $ \config -> config {docOutput = Just dirname}) + "DIR" + ) + "Directory in which to put generated documentation.", + Option + "v" + ["verbose"] + (NoArg $ Right $ \config -> config {docVerbose = True}) + "Print status messages on stderr." + ] diff --git a/src/Futhark/CLI/Misc.hs b/src/Futhark/CLI/Misc.hs --- a/src/Futhark/CLI/Misc.hs +++ b/src/Futhark/CLI/Misc.hs @@ -1,22 +1,22 @@ {-# LANGUAGE FlexibleContexts #-} + -- | Various small subcommands that are too simple to deserve their own file. module Futhark.CLI.Misc - ( mainImports - , mainDataget + ( mainImports, + mainDataget, ) where +import Control.Monad.State import qualified Data.ByteString.Lazy as BS -import Data.List (isPrefixOf, isInfixOf, nubBy) import Data.Function (on) -import Control.Monad.State -import System.FilePath -import System.IO -import System.Exit - +import Data.List (isInfixOf, isPrefixOf, nubBy) import Futhark.Compiler -import Futhark.Util.Options import Futhark.Test +import Futhark.Util.Options +import System.Exit +import System.FilePath +import System.IO -- | @futhark imports@ mainImports :: String -> [String] -> IO () @@ -24,9 +24,12 @@ case args of [file] -> Just $ do (_, prog_imports, _) <- readProgramOrDie file - liftIO $ putStr $ unlines $ map (++ ".fut") - $ filter (\f -> not ("prelude/" `isPrefixOf` f)) - $ map fst prog_imports + liftIO $ + putStr $ + unlines $ + map (++ ".fut") $ + filter (\f -> not ("prelude/" `isPrefixOf` f)) $ + map fst prog_imports _ -> Nothing -- | @futhark dataget@ @@ -35,27 +38,28 @@ case args of [file, dataset] -> Just $ dataget file dataset _ -> Nothing - where dataget prog dataset = do - let dir = takeDirectory prog - - runs <- testSpecRuns <$> testSpecFromFileOrDie prog - - let exact = filter ((dataset==) . runDescription) runs - infixes = filter ((dataset `isInfixOf`) . runDescription) runs + where + dataget prog dataset = do + let dir = takeDirectory prog - case nubBy ((==) `on` runDescription) $ - if null exact then infixes else exact of - [x] -> BS.putStr =<< getValuesBS dir (runInput x) + runs <- testSpecRuns <$> testSpecFromFileOrDie prog - [] -> do hPutStr stderr $ "No dataset '" ++ dataset ++ "'.\n" - hPutStr stderr "Available datasets:\n" - mapM_ (hPutStrLn stderr . (" "++) . runDescription) runs - exitFailure + let exact = filter ((dataset ==) . runDescription) runs + infixes = filter ((dataset `isInfixOf`) . runDescription) runs - runs' -> do hPutStr stderr $ "Dataset '" ++ dataset ++ "' ambiguous:\n" - mapM_ (hPutStrLn stderr . (" "++) . runDescription) runs' - exitFailure + case nubBy ((==) `on` runDescription) $ + if null exact then infixes else exact of + [x] -> BS.putStr =<< getValuesBS dir (runInput x) + [] -> do + hPutStr stderr $ "No dataset '" ++ dataset ++ "'.\n" + hPutStr stderr "Available datasets:\n" + mapM_ (hPutStrLn stderr . (" " ++) . runDescription) runs + exitFailure + runs' -> do + hPutStr stderr $ "Dataset '" ++ dataset ++ "' ambiguous:\n" + mapM_ (hPutStrLn stderr . (" " ++) . runDescription) runs' + exitFailure - testSpecRuns = testActionRuns . testAction - testActionRuns CompileTimeFailure{} = [] - testActionRuns (RunCases ios _ _) = concatMap iosTestRuns ios + testSpecRuns = testActionRuns . testAction + testActionRuns CompileTimeFailure {} = [] + testActionRuns (RunCases ios _ _) = concatMap iosTestRuns ios diff --git a/src/Futhark/CLI/OpenCL.hs b/src/Futhark/CLI/OpenCL.hs --- a/src/Futhark/CLI/OpenCL.hs +++ b/src/Futhark/CLI/OpenCL.hs @@ -1,14 +1,19 @@ {-# LANGUAGE FlexibleContexts #-} + -- | @futhark opencl@ module Futhark.CLI.OpenCL (main) where import Futhark.Actions (compileOpenCLAction) -import Futhark.Passes (gpuPipeline) import Futhark.Compiler.CLI +import Futhark.Passes (gpuPipeline) -- | Run @futhark opencl@ main :: String -> [String] -> IO () -main = compilerMain () [] - "Compile OpenCL" "Generate OpenCL/C code from optimised Futhark program." - gpuPipeline $ \fcfg () mode outpath prog -> - actionProcedure (compileOpenCLAction fcfg mode outpath) prog +main = compilerMain + () + [] + "Compile OpenCL" + "Generate OpenCL/C code from optimised Futhark program." + gpuPipeline + $ \fcfg () mode outpath prog -> + actionProcedure (compileOpenCLAction fcfg mode outpath) prog diff --git a/src/Futhark/CLI/Pkg.hs b/src/Futhark/CLI/Pkg.hs --- a/src/Futhark/CLI/Pkg.hs +++ b/src/Futhark/CLI/Pkg.hs @@ -1,37 +1,35 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} + -- | @futhark pkg@ module Futhark.CLI.Pkg (main) where +import qualified Codec.Archive.Zip as Zip import Control.Monad.IO.Class -import Control.Monad.State import Control.Monad.Reader -import Data.Maybe +import Control.Monad.State +import qualified Data.ByteString.Lazy as LBS +import Data.List (intercalate, isPrefixOf) import qualified Data.Map as M +import Data.Maybe +import Data.Monoid import qualified Data.Text as T import qualified Data.Text.IO as T -import qualified Data.ByteString.Lazy as LBS -import Data.List (isPrefixOf, intercalate) -import Data.Monoid +import Futhark.Pkg.Info +import Futhark.Pkg.Solve +import Futhark.Pkg.Types +import Futhark.Util (directoryContents, maxinum) +import Futhark.Util.Log +import Futhark.Util.Options +import System.Console.GetOpt import System.Directory -import System.FilePath -import qualified System.FilePath.Posix as Posix import System.Environment import System.Exit +import System.FilePath +import qualified System.FilePath.Posix as Posix import System.IO -import System.Console.GetOpt - -import qualified Codec.Archive.Zip as Zip - import Prelude -import Futhark.Util.Options -import Futhark.Pkg.Types -import Futhark.Pkg.Info -import Futhark.Pkg.Solve -import Futhark.Util (directoryContents, maxinum) -import Futhark.Util.Log - --- Installing packages installInDir :: BuildList -> FilePath -> PkgM () @@ -39,23 +37,25 @@ let putEntry from_dir pdir entry -- The archive may contain all kinds of other stuff that we don't want. | not (isInPkgDir from_dir $ Zip.eRelativePath entry) - || hasTrailingPathSeparator (Zip.eRelativePath entry) = return Nothing + || hasTrailingPathSeparator (Zip.eRelativePath entry) = + return Nothing | otherwise = do - -- Since we are writing to paths indicated in a zipfile we - -- downloaded from the wild Internet, we are going to be a - -- little bit paranoid. Specifically, we want to avoid - -- writing outside of the 'lib/' directory. We do this by - -- bailing out if the path contains any '..' components. We - -- have to use System.FilePath.Posix, because the zip library - -- claims to encode filepaths with '/' directory seperators no - -- matter the host OS. - when (".." `elem` Posix.splitPath (Zip.eRelativePath entry)) $ - fail $ "Zip archive for " <> pdir <> " contains suspicious path: " <> - Zip.eRelativePath entry - let f = pdir makeRelative from_dir (Zip.eRelativePath entry) - createDirectoryIfMissing True $ takeDirectory f - LBS.writeFile f $ Zip.fromEntry entry - return $ Just f + -- Since we are writing to paths indicated in a zipfile we + -- downloaded from the wild Internet, we are going to be a + -- little bit paranoid. Specifically, we want to avoid + -- writing outside of the 'lib/' directory. We do this by + -- bailing out if the path contains any '..' components. We + -- have to use System.FilePath.Posix, because the zip library + -- claims to encode filepaths with '/' directory seperators no + -- matter the host OS. + when (".." `elem` Posix.splitPath (Zip.eRelativePath entry)) $ + fail $ + "Zip archive for " <> pdir <> " contains suspicious path: " + <> Zip.eRelativePath entry + let f = pdir makeRelative from_dir (Zip.eRelativePath entry) + createDirectoryIfMissing True $ takeDirectory f + LBS.writeFile f $ Zip.fromEntry entry + return $ Just f isInPkgDir from_dir f = Posix.splitPath from_dir `isPrefixOf` Posix.splitPath f @@ -67,8 +67,11 @@ -- Compute the directory in the zipball that should contain the -- package files. - let noPkgDir = fail $ "futhark.pkg for " ++ T.unpack p ++ "-" ++ - T.unpack (prettySemVer v) ++ " does not define a package path." + let noPkgDir = + fail $ + "futhark.pkg for " ++ T.unpack p ++ "-" + ++ T.unpack (prettySemVer v) + ++ " does not define a package path." from_dir <- maybe noPkgDir (return . (pkgRevZipballDir info <>)) $ pkgDir m -- The directory in the local file system that will contain the @@ -86,8 +89,10 @@ catMaybes <$> liftIO (mapM (putEntry from_dir pdir) $ Zip.zEntries a) when (null written) $ - fail $ "Zip archive for package " ++ T.unpack p ++ - " does not contain any files in " ++ from_dir + fail $ + "Zip archive for package " ++ T.unpack p + ++ " does not contain any files in " + ++ from_dir libDir, libNewDir, libOldDir :: FilePath (libDir, libNewDir, libOldDir) = ("lib", "lib~new", "lib~old") @@ -123,16 +128,18 @@ libdir_exists <- liftIO $ doesDirectoryExist libDir -- 1 - liftIO $ do removePathForcibly libNewDir - createDirectoryIfMissing False libNewDir + liftIO $ do + removePathForcibly libNewDir + createDirectoryIfMissing False libNewDir -- 2 installInDir bl libNewDir -- 3 - when libdir_exists $ liftIO $ do - removePathForcibly libOldDir - renameDirectory libDir libOldDir + when libdir_exists $ + liftIO $ do + removePathForcibly libOldDir + renameDirectory libDir libOldDir -- 4 liftIO $ renameDirectory libNewDir libDir @@ -158,20 +165,23 @@ case (file_exists, dir_exists) of (True, _) -> liftIO $ parsePkgManifestFromFile futharkPkg - (_, True) -> fail $ futharkPkg <> - " exists, but it is a directory! What in Odin's beard..." - _ -> liftIO $ do T.putStrLn $ T.pack futharkPkg <> " not found - pretending it's empty." - return $ newPkgManifest Nothing + (_, True) -> + fail $ + futharkPkg + <> " exists, but it is a directory! What in Odin's beard..." + _ -> liftIO $ do + T.putStrLn $ T.pack futharkPkg <> " not found - pretending it's empty." + return $ newPkgManifest Nothing putPkgManifest :: PkgManifest -> PkgM () putPkgManifest = liftIO . T.writeFile futharkPkg . prettyPkgManifest --- The CLI -newtype PkgConfig = PkgConfig { pkgVerbose :: Bool } +newtype PkgConfig = PkgConfig {pkgVerbose :: Bool} -- | The monad in which futhark-pkg runs. -newtype PkgM a = PkgM { unPkgM :: ReaderT PkgConfig (StateT (PkgRegistry PkgM) IO) a } +newtype PkgM a = PkgM {unPkgM :: ReaderT PkgConfig (StateT (PkgRegistry PkgM) IO) a} deriving (Functor, Applicative, MonadIO, MonadReader PkgConfig) instance Monad PkgM where @@ -196,12 +206,21 @@ runPkgM :: PkgConfig -> PkgM a -> IO a runPkgM cfg (PkgM m) = evalStateT (runReaderT m cfg) mempty -cmdMain :: String -> ([String] -> PkgConfig -> Maybe (IO ())) - -> String -> [String] -> IO () +cmdMain :: + String -> + ([String] -> PkgConfig -> Maybe (IO ())) -> + String -> + [String] -> + IO () cmdMain = mainWithOptions (PkgConfig False) options - where options = [ Option "v" ["verbose"] - (NoArg $ Right $ \cfg -> cfg { pkgVerbose = True }) - "Write running diagnostics to stderr."] + where + options = + [ Option + "v" + ["verbose"] + (NoArg $ Right $ \cfg -> cfg {pkgVerbose = True}) + "Write running diagnostics to stderr." + ] doFmt :: String -> [String] -> IO () doFmt = mainWithOptions () [] "" $ \args () -> @@ -214,49 +233,56 @@ doCheck :: String -> [String] -> IO () doCheck = cmdMain "check" $ \args cfg -> case args of - [] -> Just $ runPkgM cfg $ do - m <- getPkgManifest - bl <- solveDeps $ pkgRevDeps m + [] -> Just $ + runPkgM cfg $ do + m <- getPkgManifest + bl <- solveDeps $ pkgRevDeps m - liftIO $ T.putStrLn "Dependencies chosen:" - liftIO $ T.putStr $ prettyBuildList bl + liftIO $ T.putStrLn "Dependencies chosen:" + liftIO $ T.putStr $ prettyBuildList bl - case commented $ manifestPkgPath m of - Nothing -> return () - Just p -> do - let pdir = "lib" T.unpack p + case commented $ manifestPkgPath m of + Nothing -> return () + Just p -> do + let pdir = "lib" T.unpack p - pdir_exists <- liftIO $ doesDirectoryExist pdir + pdir_exists <- liftIO $ doesDirectoryExist pdir - unless pdir_exists $ liftIO $ do - T.putStrLn $ "Problem: the directory " <> T.pack pdir <> " does not exist." - exitFailure + unless pdir_exists $ + liftIO $ do + T.putStrLn $ "Problem: the directory " <> T.pack pdir <> " does not exist." + exitFailure - anything <- liftIO $ any ((==".fut") . takeExtension) <$> - directoryContents ("lib" T.unpack p) - unless anything $ liftIO $ do - T.putStrLn $ "Problem: the directory " <> T.pack pdir <> " does not contain any .fut files." - exitFailure + anything <- + liftIO $ + any ((== ".fut") . takeExtension) + <$> directoryContents ("lib" T.unpack p) + unless anything $ + liftIO $ do + T.putStrLn $ "Problem: the directory " <> T.pack pdir <> " does not contain any .fut files." + exitFailure _ -> Nothing doSync :: String -> [String] -> IO () doSync = cmdMain "" $ \args cfg -> case args of - [] -> Just $ runPkgM cfg $ do - m <- getPkgManifest - bl <- solveDeps $ pkgRevDeps m - installBuildList (commented $ manifestPkgPath m) bl + [] -> Just $ + runPkgM cfg $ do + m <- getPkgManifest + bl <- solveDeps $ pkgRevDeps m + installBuildList (commented $ manifestPkgPath m) bl _ -> Nothing doAdd :: String -> [String] -> IO () doAdd = cmdMain "PKGPATH" $ \args cfg -> case args of [p, v] | Right v' <- parseVersion $ T.pack v -> Just $ runPkgM cfg $ doAdd' (T.pack p) v' - [p] -> Just $ runPkgM cfg $ - -- Look up the newest revision of the package. - doAdd' (T.pack p) =<< lookupNewestRev (T.pack p) + [p] -> + Just $ + runPkgM cfg $ + -- Look up the newest revision of the package. + doAdd' (T.pack p) =<< lookupNewestRev (T.pack p) _ -> Nothing - where doAdd' p v = do m <- getPkgManifest @@ -272,11 +298,11 @@ -- we add a new one. p_info <- lookupPackageRev p v let hash = case (_svMajor v, _svMinor v, _svPatch v) of - -- We do not perform hash-pinning for - -- (0,0,0)-versions, because these already embed a - -- specific revision ID into their version number. - (0, 0, 0) -> Nothing - _ -> Just $ pkgRevCommit p_info + -- We do not perform hash-pinning for + -- (0,0,0)-versions, because these already embed a + -- specific revision ID into their version number. + (0, 0, 0) -> Nothing + _ -> Just $ pkgRevCommit p_info req = Required p v hash (m', prev_r) = addRequiredToManifest req m @@ -285,8 +311,13 @@ | requiredPkgRev prev_r' == v -> liftIO $ T.putStrLn $ "Package already at version " <> prettySemVer v <> "; nothing to do." | otherwise -> - liftIO $ T.putStrLn $ "Replaced " <> p <> " " <> - prettySemVer (requiredPkgRev prev_r') <> " => " <> prettySemVer v <> "." + liftIO $ + T.putStrLn $ + "Replaced " <> p <> " " + <> prettySemVer (requiredPkgRev prev_r') + <> " => " + <> prettySemVer v + <> "." Nothing -> liftIO $ T.putStrLn $ "Added new required package " <> p <> " " <> prettySemVer v <> "." putPkgManifest m' @@ -316,9 +347,10 @@ where doCreate' p = do exists <- liftIO $ (||) <$> doesFileExist futharkPkg <*> doesDirectoryExist futharkPkg - when exists $ liftIO $ do - T.putStrLn $ T.pack futharkPkg <> " already exists." - exitFailure + when exists $ + liftIO $ do + T.putStrLn $ T.pack futharkPkg <> " already exists." + exitFailure liftIO $ createDirectoryIfMissing True $ "lib" T.unpack p liftIO $ T.putStrLn $ "Created directory " <> T.pack ("lib" T.unpack p) <> "." @@ -329,33 +361,44 @@ doUpgrade :: String -> [String] -> IO () doUpgrade = cmdMain "" $ \args cfg -> case args of - [] -> Just $ runPkgM cfg $ do - m <- getPkgManifest - rs <- traverse (mapM (traverse upgrade)) $ manifestRequire m - putPkgManifest m { manifestRequire = rs } - if rs == manifestRequire m - then liftIO $ T.putStrLn "Nothing to upgrade." - else liftIO $ T.putStrLn "Remember to run 'futhark pkg sync'." + [] -> Just $ + runPkgM cfg $ do + m <- getPkgManifest + rs <- traverse (mapM (traverse upgrade)) $ manifestRequire m + putPkgManifest m {manifestRequire = rs} + if rs == manifestRequire m + then liftIO $ T.putStrLn "Nothing to upgrade." + else liftIO $ T.putStrLn "Remember to run 'futhark pkg sync'." _ -> Nothing - where upgrade req = do - v <- lookupNewestRev $ requiredPkg req - h <- pkgRevCommit <$> lookupPackageRev (requiredPkg req) v + where + upgrade req = do + v <- lookupNewestRev $ requiredPkg req + h <- pkgRevCommit <$> lookupPackageRev (requiredPkg req) v - when (v /= requiredPkgRev req) $ - liftIO $ T.putStrLn $ "Upgraded " <> requiredPkg req <> " " <> - prettySemVer (requiredPkgRev req) <> " => " <> prettySemVer v <> "." + when (v /= requiredPkgRev req) $ + liftIO $ + T.putStrLn $ + "Upgraded " <> requiredPkg req <> " " + <> prettySemVer (requiredPkgRev req) + <> " => " + <> prettySemVer v + <> "." - return req { requiredPkgRev = v - , requiredHash = Just h } + return + req + { requiredPkgRev = v, + requiredHash = Just h + } doVersions :: String -> [String] -> IO () doVersions = cmdMain "PKGPATH" $ \args cfg -> case args of [p] -> Just $ runPkgM cfg $ doVersions' $ T.pack p _ -> Nothing - where doVersions' = - mapM_ (liftIO . T.putStrLn . prettySemVer) . M.keys . pkgVersions - <=< lookupPackage + where + doVersions' = + mapM_ (liftIO . T.putStrLn . prettySemVer) . M.keys . pkgVersions + <=< lookupPackage -- | Run @futhark pkg@. main :: String -> [String] -> IO () @@ -363,37 +406,49 @@ -- Avoid Git asking for credentials. We prefer failure. liftIO $ setEnv "GIT_TERMINAL_PROMPT" "0" - let commands = [ ("add", - (doAdd, "Add another required package to futhark.pkg.")) - , ("check", - (doCheck, "Check that futhark.pkg is satisfiable.")) - , ("init", - (doInit, "Create a new futhark.pkg and a lib/ skeleton.")) - , ("fmt", - (doFmt, "Reformat futhark.pkg.")) - , ("sync", - (doSync, "Populate lib/ as specified by futhark.pkg.")) - , ("remove", - (doRemove, "Remove a required package from futhark.pkg.")) - , ("upgrade", - (doUpgrade, "Upgrade all packages to newest versions.")) - , ("versions", - (doVersions, "List available versions for a package.")) - ] + let commands = + [ ( "add", + (doAdd, "Add another required package to futhark.pkg.") + ), + ( "check", + (doCheck, "Check that futhark.pkg is satisfiable.") + ), + ( "init", + (doInit, "Create a new futhark.pkg and a lib/ skeleton.") + ), + ( "fmt", + (doFmt, "Reformat futhark.pkg.") + ), + ( "sync", + (doSync, "Populate lib/ as specified by futhark.pkg.") + ), + ( "remove", + (doRemove, "Remove a required package from futhark.pkg.") + ), + ( "upgrade", + (doUpgrade, "Upgrade all packages to newest versions.") + ), + ( "versions", + (doVersions, "List available versions for a package.") + ) + ] usage = "options... <" <> intercalate "|" (map fst commands) <> ">" case args of - cmd : args' | Just (m, _) <- lookup cmd commands -> - m (unwords [prog, cmd]) args' + cmd : args' + | Just (m, _) <- lookup cmd commands -> + m (unwords [prog, cmd]) args' _ -> do let bad _ () = Just $ do let k = maxinum (map (length . fst) commands) + 3 - usageMsg $ T.unlines $ - [" ...:", "", "Commands:"] ++ - [ " " <> T.pack cmd <> T.pack (replicate (k - length cmd) ' ') <> desc - | (cmd, (_, desc)) <- commands ] + usageMsg $ + T.unlines $ + [" ...:", "", "Commands:"] + ++ [ " " <> T.pack cmd <> T.pack (replicate (k - length cmd) ' ') <> desc + | (cmd, (_, desc)) <- commands + ] mainWithOptions () [] usage bad prog args - - where usageMsg s = do - T.putStrLn $ "Usage: " <> T.pack prog <> " [--version] [--help] " <> s - exitFailure + where + usageMsg s = do + T.putStrLn $ "Usage: " <> T.pack prog <> " [--version] [--help] " <> s + exitFailure diff --git a/src/Futhark/CLI/PyOpenCL.hs b/src/Futhark/CLI/PyOpenCL.hs --- a/src/Futhark/CLI/PyOpenCL.hs +++ b/src/Futhark/CLI/PyOpenCL.hs @@ -1,29 +1,34 @@ {-# LANGUAGE FlexibleContexts #-} + -- | @futhark pyopencl@ module Futhark.CLI.PyOpenCL (main) where import Control.Monad.IO.Class -import System.FilePath -import System.Directory - -import Futhark.Passes import qualified Futhark.CodeGen.Backends.PyOpenCL as PyOpenCL import Futhark.Compiler.CLI +import Futhark.Passes +import System.Directory +import System.FilePath -- | Run @futhark pyopencl@. main :: String -> [String] -> IO () -main = compilerMain () [] - "Compile PyOpenCL" "Generate Python + OpenCL code from optimised Futhark program." - gpuPipeline $ \fcfg () mode outpath prog -> do - let class_name = - case mode of ToLibrary -> Just $ takeBaseName outpath - ToExecutable -> Nothing - pyprog <- handleWarnings fcfg $ PyOpenCL.compileProg class_name prog - +main = compilerMain + () + [] + "Compile PyOpenCL" + "Generate Python + OpenCL code from optimised Futhark program." + gpuPipeline + $ \fcfg () mode outpath prog -> do + let class_name = case mode of - ToLibrary -> - liftIO $ writeFile (outpath `addExtension` "py") pyprog - ToExecutable -> liftIO $ do - writeFile outpath pyprog - perms <- liftIO $ getPermissions outpath - setPermissions outpath $ setOwnerExecutable True perms + ToLibrary -> Just $ takeBaseName outpath + ToExecutable -> Nothing + pyprog <- handleWarnings fcfg $ PyOpenCL.compileProg class_name prog + + case mode of + ToLibrary -> + liftIO $ writeFile (outpath `addExtension` "py") pyprog + ToExecutable -> liftIO $ do + writeFile outpath pyprog + perms <- liftIO $ getPermissions outpath + setPermissions outpath $ setOwnerExecutable True perms diff --git a/src/Futhark/CLI/Python.hs b/src/Futhark/CLI/Python.hs --- a/src/Futhark/CLI/Python.hs +++ b/src/Futhark/CLI/Python.hs @@ -1,29 +1,34 @@ {-# LANGUAGE FlexibleContexts #-} + -- | @futhark py@ module Futhark.CLI.Python (main) where import Control.Monad.IO.Class -import System.FilePath -import System.Directory - -import Futhark.Passes import qualified Futhark.CodeGen.Backends.SequentialPython as SequentialPy import Futhark.Compiler.CLI +import Futhark.Passes +import System.Directory +import System.FilePath -- | Run @futhark py@ main :: String -> [String] -> IO () -main = compilerMain () [] - "Compile sequential Python" "Generate sequential Python code from optimised Futhark program." - sequentialCpuPipeline $ \fcfg () mode outpath prog -> do - let class_name = - case mode of ToLibrary -> Just $ takeBaseName outpath - ToExecutable -> Nothing - pyprog <- handleWarnings fcfg $ SequentialPy.compileProg class_name prog - +main = compilerMain + () + [] + "Compile sequential Python" + "Generate sequential Python code from optimised Futhark program." + sequentialCpuPipeline + $ \fcfg () mode outpath prog -> do + let class_name = case mode of - ToLibrary -> - liftIO $ writeFile (outpath `addExtension` "py") pyprog - ToExecutable -> liftIO $ do - writeFile outpath pyprog - perms <- liftIO $ getPermissions outpath - setPermissions outpath $ setOwnerExecutable True perms + ToLibrary -> Just $ takeBaseName outpath + ToExecutable -> Nothing + pyprog <- handleWarnings fcfg $ SequentialPy.compileProg class_name prog + + case mode of + ToLibrary -> + liftIO $ writeFile (outpath `addExtension` "py") pyprog + ToExecutable -> liftIO $ do + writeFile outpath pyprog + perms <- liftIO $ getPermissions outpath + setPermissions outpath $ setOwnerExecutable True perms diff --git a/src/Futhark/CLI/Query.hs b/src/Futhark/CLI/Query.hs --- a/src/Futhark/CLI/Query.hs +++ b/src/Futhark/CLI/Query.hs @@ -1,17 +1,16 @@ -- | @futhark query@ module Futhark.CLI.Query (main) where -import Text.Read (readMaybe) - import Futhark.Compiler -import Futhark.Util.Options import Futhark.Util.Loc +import Futhark.Util.Options import Language.Futhark.Query import Language.Futhark.Syntax +import Text.Read (readMaybe) -- | Run @futhark query@. main :: String -> [String] -> IO () -main = mainWithOptions () [] "program line:col" $ \args () -> +main = mainWithOptions () [] "program line col" $ \args () -> case args of [file, line, col] -> do line' <- readMaybe line diff --git a/src/Futhark/CLI/REPL.hs b/src/Futhark/CLI/REPL.hs --- a/src/Futhark/CLI/REPL.hs +++ b/src/Futhark/CLI/REPL.hs @@ -1,58 +1,58 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE ScopedTypeVariables #-} + -- | @futhark repl@ module Futhark.CLI.REPL (main) where -import Control.Monad.Free.Church import Control.Exception -import Data.Char -import Data.List (intercalate, intersperse) -import Data.Maybe -import Data.Version import Control.Monad -import Control.Monad.IO.Class -import Control.Monad.State import Control.Monad.Except +import Control.Monad.Free.Church +import Control.Monad.State +import Data.Char +import Data.List (intercalate, intersperse) import qualified Data.List.NonEmpty as NE +import Data.Maybe import qualified Data.Text as T import qualified Data.Text.IO as T -import NeatInterpolation (text) -import System.Directory -import System.FilePath -import System.Console.GetOpt -import Text.Read (readMaybe) -import qualified System.Console.Haskeline as Haskeline - -import Language.Futhark -import Language.Futhark.Parser hiding (EOF) -import qualified Language.Futhark.TypeChecker as T -import qualified Language.Futhark.Semantic as T -import Futhark.MonadFreshNames -import Futhark.Version +import Data.Version import Futhark.Compiler +import Futhark.MonadFreshNames import Futhark.Pipeline -import Futhark.Util.Options import Futhark.Util (toPOSIX) - +import Futhark.Util.Options +import Futhark.Version +import Language.Futhark import qualified Language.Futhark.Interpreter as I +import Language.Futhark.Parser hiding (EOF) +import qualified Language.Futhark.Semantic as T +import qualified Language.Futhark.TypeChecker as T +import NeatInterpolation (text) +import System.Console.GetOpt +import qualified System.Console.Haskeline as Haskeline +import System.Directory +import System.FilePath +import Text.Read (readMaybe) banner :: String -banner = unlines [ - "|// |\\ | |\\ |\\ /", - "|/ | \\ |\\ |\\ |/ /", - "| | \\ |/ | |\\ \\", - "| | \\ | | | \\ \\" - ] +banner = + unlines + [ "|// |\\ | |\\ |\\ /", + "|/ | \\ |\\ |\\ |/ /", + "| | \\ |/ | |\\ \\", + "| | \\ | | | \\ \\" + ] -- | Run @futhark repl@. main :: String -> [String] -> IO () main = mainWithOptions interpreterConfig options "options... [program.fut]" run - where run [] _ = Just $ repl Nothing - run [prog] _ = Just $ repl $ Just prog - run _ _ = Nothing + where + run [] _ = Just $ repl Nothing + run [prog] _ = Just $ repl $ Just prog + run _ _ = Nothing data StopReason = EOF | Stop | Exit | Load FilePath @@ -77,8 +77,9 @@ liftIO $ newFutharkiState (futharkiCount s) $ Just file case maybe_new_state of Right new_state -> toploop new_state - Left err -> do liftIO $ putStrLn err - toploop s' + Left err -> do + liftIO $ putStrLn err + toploop s' Right _ -> return () finish s = do @@ -99,92 +100,116 @@ Nothing -> return True -- EOF Just 'y' -> return True Just 'n' -> return False - _ -> confirmQuit + _ -> confirmQuit -newtype InterpreterConfig = InterpreterConfig { interpreterEntryPoint :: Name } +newtype InterpreterConfig = InterpreterConfig {interpreterEntryPoint :: Name} interpreterConfig :: InterpreterConfig interpreterConfig = InterpreterConfig defaultEntryPoint options :: [FunOptDescr InterpreterConfig] -options = [ Option "e" ["entry-point"] - (ReqArg (\entry -> Right $ \config -> - config { interpreterEntryPoint = nameFromString entry }) - "NAME") - "The entry point to execute." - ] +options = + [ Option + "e" + ["entry-point"] + ( ReqArg + ( \entry -> Right $ \config -> + config {interpreterEntryPoint = nameFromString entry} + ) + "NAME" + ) + "The entry point to execute." + ] -- | Representation of breaking at a breakpoint, to allow for -- navigating through the stack frames and such. -data Breaking = Breaking { breakingStack :: NE.NonEmpty I.StackFrame - , breakingAt :: Int - -- ^ Index of the current breakpoint (with - -- 0 being the outermost). - } +data Breaking = Breaking + { breakingStack :: NE.NonEmpty I.StackFrame, + -- | Index of the current breakpoint (with + -- 0 being the outermost). + breakingAt :: Int + } -data FutharkiState = - FutharkiState { futharkiImports :: Imports - , futharkiNameSource :: VNameSource - , futharkiCount :: Int - , futharkiEnv :: (T.Env, I.Ctx) - , futharkiBreaking :: Maybe Breaking - -- ^ Are we currently stopped at a breakpoint? - , futharkiSkipBreaks :: [Loc] - -- ^ Skip breakpoints at these locations. - , futharkiBreakOnNaN :: Bool - , futharkiLoaded :: Maybe FilePath - -- ^ The currently loaded file. - } +data FutharkiState = FutharkiState + { futharkiImports :: Imports, + futharkiNameSource :: VNameSource, + futharkiCount :: Int, + futharkiEnv :: (T.Env, I.Ctx), + -- | Are we currently stopped at a breakpoint? + futharkiBreaking :: Maybe Breaking, + -- | Skip breakpoints at these locations. + futharkiSkipBreaks :: [Loc], + futharkiBreakOnNaN :: Bool, + -- | The currently loaded file. + futharkiLoaded :: Maybe FilePath + } newFutharkiState :: Int -> Maybe FilePath -> IO (Either String FutharkiState) newFutharkiState count maybe_file = runExceptT $ do (imports, src, tenv, ienv) <- case maybe_file of - Nothing -> do -- Load the builtins through the type checker. (_, imports, src) <- badOnLeft show =<< runExceptT (readLibrary []) -- Then into the interpreter. - ienv <- foldM (\ctx -> badOnLeft show <=< runInterpreter' . I.interpretImport ctx) - I.initialCtx $ map (fmap fileProg) imports + ienv <- + foldM + (\ctx -> badOnLeft show <=< runInterpreter' . I.interpretImport ctx) + I.initialCtx + $ map (fmap fileProg) imports -- Then make the prelude available in the type checker. - (tenv, d, src') <- badOnLeft pretty $ T.checkDec imports src T.initialEnv - (T.mkInitialImport ".") $ mkOpen "/prelude/prelude" + (tenv, d, src') <- + badOnLeft pretty $ + T.checkDec + imports + src + T.initialEnv + (T.mkInitialImport ".") + $ mkOpen "/prelude/prelude" -- Then in the interpreter. ienv' <- badOnLeft show =<< runInterpreter' (I.interpretDec ienv d) return (imports, src', tenv, ienv') - Just file -> do (ws, imports, src) <- - badOnLeft show =<< - liftIO (runExceptT (readProgram file) - `catch` \(err::IOException) -> - return (externalErrorS (show err))) + badOnLeft show + =<< liftIO + ( runExceptT (readProgram file) + `catch` \(err :: IOException) -> + return (externalErrorS (show err)) + ) liftIO $ print ws let imp = T.mkInitialImport "." - ienv1 <- foldM (\ctx -> badOnLeft show <=< runInterpreter' . I.interpretImport ctx) I.initialCtx $ - map (fmap fileProg) imports - (tenv1, d1, src') <- badOnLeft pretty $ T.checkDec imports src T.initialEnv imp $ - mkOpen "/prelude/prelude" - (tenv2, d2, src'') <- badOnLeft pretty $ T.checkDec imports src' tenv1 imp $ - mkOpen $ toPOSIX $ dropExtension file + ienv1 <- + foldM (\ctx -> badOnLeft show <=< runInterpreter' . I.interpretImport ctx) I.initialCtx $ + map (fmap fileProg) imports + (tenv1, d1, src') <- + badOnLeft pretty $ + T.checkDec imports src T.initialEnv imp $ + mkOpen "/prelude/prelude" + (tenv2, d2, src'') <- + badOnLeft pretty $ + T.checkDec imports src' tenv1 imp $ + mkOpen $ toPOSIX $ dropExtension file ienv2 <- badOnLeft show =<< runInterpreter' (I.interpretDec ienv1 d1) ienv3 <- badOnLeft show =<< runInterpreter' (I.interpretDec ienv2 d2) return (imports, src'', tenv2, ienv3) - return FutharkiState { futharkiImports = imports - , futharkiNameSource = src - , futharkiCount = count - , futharkiEnv = (tenv, ienv) - , futharkiBreaking = Nothing - , futharkiSkipBreaks = mempty - , futharkiBreakOnNaN = False - , futharkiLoaded = maybe_file - } - where badOnLeft :: (err -> String) -> Either err a -> ExceptT String IO a - badOnLeft _ (Right x) = return x - badOnLeft p (Left err) = throwError $ p err + return + FutharkiState + { futharkiImports = imports, + futharkiNameSource = src, + futharkiCount = count, + futharkiEnv = (tenv, ienv), + futharkiBreaking = Nothing, + futharkiSkipBreaks = mempty, + futharkiBreakOnNaN = False, + futharkiLoaded = maybe_file + } + where + badOnLeft :: (err -> String) -> Either err a -> ExceptT String IO a + badOnLeft _ (Right x) = return x + badOnLeft p (Left err) = throwError $ p err getPrompt :: FutharkiM String getPrompt = do @@ -195,10 +220,15 @@ mkOpen f = OpenDec (ModImport f NoInfo mempty) mempty -- The ExceptT part is more of a continuation, really. -newtype FutharkiM a = - FutharkiM { runFutharkiM :: ExceptT StopReason (StateT FutharkiState (Haskeline.InputT IO)) a } - deriving (Functor, Applicative, Monad, - MonadState FutharkiState, MonadIO, MonadError StopReason) +newtype FutharkiM a = FutharkiM {runFutharkiM :: ExceptT StopReason (StateT FutharkiState (Haskeline.InputT IO)) a} + deriving + ( Functor, + Applicative, + Monad, + MonadState FutharkiState, + MonadIO, + MonadError StopReason + ) readEvalPrint :: FutharkiM () readEvalPrint = do @@ -209,16 +239,17 @@ Nothing | isJust breaking -> throwError Stop | otherwise -> return () - Just (':', command) -> do let (cmdname, rest) = T.break isSpace command arg = T.dropWhileEnd isSpace $ T.dropWhile isSpace rest case filter ((cmdname `T.isPrefixOf`) . fst) commands of [] -> liftIO $ T.putStrLn $ "Unknown command '" <> cmdname <> "'" [(_, (cmdf, _))] -> cmdf arg - matches -> liftIO $ T.putStrLn $ "Ambiguous command; could be one of " <> - mconcat (intersperse ", " (map fst matches)) - + matches -> + liftIO $ + T.putStrLn $ + "Ambiguous command; could be one of " + <> mconcat (intersperse ", " (map fst matches)) _ -> do -- Read a declaration or expression. maybe_dec_or_e <- parseDecOrExpIncrM (inputLine " ") prompt line @@ -227,12 +258,13 @@ Left err -> liftIO $ print err Right (Left d) -> onDec d Right (Right e) -> onExp e - modify $ \s -> s { futharkiCount = futharkiCount s + 1 } - where inputLine prompt = do - inp <- FutharkiM $ lift $ lift $ Haskeline.getInputLine prompt - case inp of - Just s -> return $ T.pack s - Nothing -> throwError EOF + modify $ \s -> s {futharkiCount = futharkiCount s + 1} + where + inputLine prompt = do + inp <- FutharkiM $ lift $ lift $ Haskeline.getInputLine prompt + case inp of + Just s -> return $ T.pack s + Nothing -> throwError EOF getIt :: FutharkiM (Imports, VNameSource, T.Env, I.Ctx) getIt = do @@ -256,7 +288,7 @@ case imp_r of Left e -> liftIO $ print e - Right (_, imports', src') -> + Right (_, imports', src') -> case T.checkDec imports' src' tenv cur_import d of Left e -> liftIO $ putStrLn $ pretty e Right (tenv', d', src'') -> do @@ -268,28 +300,30 @@ I.interpretDec ienv' d' case int_r of Left err -> liftIO $ print err - Right ienv' -> modify $ \s -> s { futharkiEnv = (tenv', ienv') - , futharkiImports = imports' - , futharkiNameSource = src'' - } + Right ienv' -> modify $ \s -> + s + { futharkiEnv = (tenv', ienv'), + futharkiImports = imports', + futharkiNameSource = src'' + } onExp :: UncheckedExp -> FutharkiM () onExp e = do (imports, src, tenv, ienv) <- getIt case either (Left . pretty) Right $ - T.checkExp imports src tenv e of + T.checkExp imports src tenv e of Left err -> liftIO $ putStrLn err Right (tparams, e') | null tparams -> do - r <- runInterpreter $ I.interpretExp ienv e' - case r of - Left err -> liftIO $ print err - Right v -> liftIO $ putStrLn $ pretty v - + r <- runInterpreter $ I.interpretExp ienv e' + case r of + Left err -> liftIO $ print err + Right v -> liftIO $ putStrLn $ pretty v | otherwise -> liftIO $ do - putStrLn $ "Inferred type of expression: " ++ pretty (typeOf e') - putStrLn $ "The following types are ambiguous: " ++ - intercalate ", " (map (prettyName . typeParamName) tparams) + putStrLn $ "Inferred type of expression: " ++ pretty (typeOf e') + putStrLn $ + "The following types are ambiguous: " + ++ intercalate ", " (map (prettyName . typeParamName) tparams) prettyBreaking :: Breaking -> String prettyBreaking b = @@ -302,8 +336,8 @@ breakForReason s _ I.BreakNaN | not $ futharkiBreakOnNaN s = False breakForReason s top _ = - isNothing (futharkiBreaking s) && - locOf top `notElem` futharkiSkipBreaks s + isNothing (futharkiBreaking s) + && locOf top `notElem` futharkiSkipBreaks s runInterpreter :: F I.ExtOp a -> FutharkiM (Either I.InterpreterError a) runInterpreter m = runF m (return . Right) intOp @@ -316,8 +350,9 @@ intOp (I.ExtOpBreak why callstack c) = do s <- get - let why' = case why of I.BreakPoint -> "Breakpoint" - I.BreakNaN -> "NaN produced" + let why' = case why of + I.BreakPoint -> "Breakpoint" + I.BreakNaN -> "NaN produced" top = NE.head callstack ctx = I.stackFrameCtx top tenv = I.typeCheckerEnv $ I.ctxEnv ctx @@ -332,33 +367,41 @@ -- Note the cleverness to preserve the Haskeline session (for -- line history and such). (stop, s') <- - FutharkiM $ lift $ lift $ - runStateT (runExceptT $ runFutharkiM $ forever readEvalPrint) - s { futharkiEnv = (tenv, ctx) - , futharkiCount = futharkiCount s + 1 - , futharkiBreaking = Just breaking - } + FutharkiM $ + lift $ + lift $ + runStateT + (runExceptT $ runFutharkiM $ forever readEvalPrint) + s + { futharkiEnv = (tenv, ctx), + futharkiCount = futharkiCount s + 1, + futharkiBreaking = Just breaking + } case stop of Left (Load file) -> throwError $ Load file - _ -> do liftIO $ putStrLn "Continuing..." - put s { futharkiCount = - futharkiCount s' - , futharkiSkipBreaks = - futharkiSkipBreaks s' <> futharkiSkipBreaks s - , futharkiBreakOnNaN = - futharkiBreakOnNaN s' - } + _ -> do + liftIO $ putStrLn "Continuing..." + put + s + { futharkiCount = + futharkiCount s', + futharkiSkipBreaks = + futharkiSkipBreaks s' <> futharkiSkipBreaks s, + futharkiBreakOnNaN = + futharkiBreakOnNaN s' + } c runInterpreter' :: MonadIO m => F I.ExtOp a -> m (Either I.InterpreterError a) runInterpreter' m = runF m (return . Right) intOp - where intOp (I.ExtOpError err) = return $ Left err - intOp (I.ExtOpTrace w v c) = do - liftIO $ putStrLn $ "Trace at " ++ locStr w ++ ": " ++ v - c - intOp (I.ExtOpBreak _ _ c) = c + where + intOp (I.ExtOpError err) = return $ Left err + intOp (I.ExtOpTrace w v c) = do + liftIO $ putStrLn $ "Trace at " ++ locStr w ++ ": " ++ v + c + intOp (I.ExtOpBreak _ _ c) = c type Command = T.Text -> FutharkiM () @@ -370,11 +413,12 @@ (True, Nothing) -> liftIO $ T.putStrLn "No file specified and no file previously loaded." (False, _) -> throwError $ Load $ T.unpack file -genTypeCommand :: Show err => - (String -> T.Text -> Either err a) - -> (Imports -> VNameSource -> T.Env -> a -> Either T.TypeError b) - -> (b -> String) - -> Command +genTypeCommand :: + Show err => + (String -> T.Text -> Either err a) -> + (Imports -> VNameSource -> T.Env -> a -> Either T.TypeError b) -> + (b -> String) -> + Command genTypeCommand f g h e = do prompt <- getPrompt case f prompt e of @@ -389,8 +433,9 @@ typeCommand :: Command typeCommand = genTypeCommand parseExp T.checkExp $ \(ps, e) -> - pretty e <> concatMap ((" "<>) . pretty) ps <> - " : " <> pretty (typeOf e) + pretty e <> concatMap ((" " <>) . pretty) ps + <> " : " + <> pretty (typeOf e) mtypeCommand :: Command mtypeCommand = genTypeCommand parseModExp T.checkModExp $ pretty . fst @@ -400,31 +445,35 @@ top <- gets $ fmap (NE.head . breakingStack) . futharkiBreaking case top of Nothing -> liftIO $ putStrLn "Not currently stopped at a breakpoint." - Just top' -> do modify $ \s -> s { futharkiSkipBreaks = locOf top' : futharkiSkipBreaks s } - throwError Stop + Just top' -> do + modify $ \s -> s {futharkiSkipBreaks = locOf top' : futharkiSkipBreaks s} + throwError Stop nanbreakCommand :: Command nanbreakCommand _ = do - modify $ \s -> s { futharkiBreakOnNaN = not $ futharkiBreakOnNaN s } + modify $ \s -> s {futharkiBreakOnNaN = not $ futharkiBreakOnNaN s} b <- gets futharkiBreakOnNaN - liftIO $ putStrLn $ - if b - then "Now treating NaNs as breakpoints." - else "No longer treating NaNs as breakpoints." + liftIO $ + putStrLn $ + if b + then "Now treating NaNs as breakpoints." + else "No longer treating NaNs as breakpoints." frameCommand :: Command frameCommand which = do maybe_stack <- gets $ fmap breakingStack . futharkiBreaking case (maybe_stack, readMaybe $ T.unpack which) of (Just stack, Just i) - | frame:_ <- NE.drop i stack -> do - let breaking = Breaking stack i - ctx = I.stackFrameCtx frame - tenv = I.typeCheckerEnv $ I.ctxEnv ctx - modify $ \s -> s { futharkiEnv = (tenv, ctx) - , futharkiBreaking = Just breaking - } - liftIO $ putStrLn $ prettyBreaking breaking + | frame : _ <- NE.drop i stack -> do + let breaking = Breaking stack i + ctx = I.stackFrameCtx frame + tenv = I.typeCheckerEnv $ I.ctxEnv ctx + modify $ \s -> + s + { futharkiEnv = (tenv, ctx), + futharkiBreaking = Just breaking + } + liftIO $ putStrLn $ prettyBreaking breaking (Just _, _) -> liftIO $ putStrLn $ "Invalid stack index: " ++ T.unpack which (Nothing, _) -> @@ -435,15 +484,17 @@ cdCommand :: Command cdCommand dir - | T.null dir = liftIO $ putStrLn "Usage: ':cd '." - | otherwise = - liftIO $ setCurrentDirectory (T.unpack dir) - `catch` \(err::IOException) -> print err + | T.null dir = liftIO $ putStrLn "Usage: ':cd '." + | otherwise = + liftIO $ + setCurrentDirectory (T.unpack dir) + `catch` \(err :: IOException) -> print err helpCommand :: Command -helpCommand _ = liftIO $ forM_ commands $ \(cmd, (_, desc)) -> do +helpCommand _ = liftIO $ + forM_ commands $ \(cmd, (_, desc)) -> do T.putStrLn $ ":" <> cmd - T.putStrLn $ T.replicate (1+T.length cmd) "-" + T.putStrLn $ T.replicate (1 + T.length cmd) "-" T.putStr desc T.putStrLn "" T.putStrLn "" @@ -452,7 +503,10 @@ quitCommand _ = throwError Exit commands :: [(T.Text, (Command, T.Text))] -commands = [("load", (loadCommand, [text| +commands = + [ ( "load", + ( loadCommand, + [text| Load a Futhark source file. Usage: > :load foo.fut @@ -464,34 +518,73 @@ second time will replace the previously loaded file. It will also replace any declarations entered at the REPL. -|])), - ("type", (typeCommand, [text| +|] + ) + ), + ( "type", + ( typeCommand, + [text| Show the type of an expression, which must fit on a single line. -|])), - ("mtype", (mtypeCommand, [text| +|] + ) + ), + ( "mtype", + ( mtypeCommand, + [text| Show the type of a module expression, which must fit on a single line. -|])), - ("unbreak", (unbreakCommand, [text| +|] + ) + ), + ( "unbreak", + ( unbreakCommand, + [text| Skip all future occurences of the current breakpoint. -|])), - ("nanbreak", (nanbreakCommand, [text| +|] + ) + ), + ( "nanbreak", + ( nanbreakCommand, + [text| Toggle treating operators that produce new NaNs as breakpoints. We consider a NaN to be "new" if none of the arguments to the operator in question is a NaN. -|])), - ("frame", (frameCommand, [text| +|] + ) + ), + ( "frame", + ( frameCommand, + [text| While at a break point, jump to another stack frame, whose variables can then be inspected. Resuming from the breakpoint will jump back to the innermost stack frame. -|])), - ("pwd", (pwdCommand, [text| +|] + ) + ), + ( "pwd", + ( pwdCommand, + [text| Print the current working directory. -|])), - ("cd", (cdCommand, [text| +|] + ) + ), + ( "cd", + ( cdCommand, + [text| Change the current working directory. -|])), - ("help", (helpCommand, [text| +|] + ) + ), + ( "help", + ( helpCommand, + [text| Print a list of commands and a description of their behaviour. -|])), - ("quit", (quitCommand, [text| +|] + ) + ), + ( "quit", + ( quitCommand, + [text| Exit REPL. -|]))] +|] + ) + ) + ] diff --git a/src/Futhark/CLI/Run.hs b/src/Futhark/CLI/Run.hs --- a/src/Futhark/CLI/Run.hs +++ b/src/Futhark/CLI/Run.hs @@ -1,47 +1,47 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} + -- | @futhark run@ module Futhark.CLI.Run (main) where -import Control.Monad.Free.Church import Control.Exception -import Data.Maybe -import qualified Data.Map as M import Control.Monad -import Control.Monad.IO.Class import Control.Monad.Except +import Control.Monad.Free.Church +import qualified Data.Map as M +import Data.Maybe import qualified Data.Text.IO as T -import System.FilePath -import System.Exit -import System.Console.GetOpt -import System.IO - -import Prelude - -import Language.Futhark -import Language.Futhark.Parser hiding (EOF) -import qualified Language.Futhark.TypeChecker as T -import qualified Language.Futhark.Semantic as T import Futhark.Compiler import Futhark.Pipeline -import Futhark.Util.Options import Futhark.Util (toPOSIX) - +import Futhark.Util.Options +import Language.Futhark import qualified Language.Futhark.Interpreter as I +import Language.Futhark.Parser hiding (EOF) +import qualified Language.Futhark.Semantic as T +import qualified Language.Futhark.TypeChecker as T +import System.Console.GetOpt +import System.Exit +import System.FilePath +import System.IO +import Prelude -- | Run @futhark run@. main :: String -> [String] -> IO () main = mainWithOptions interpreterConfig options "options... " run - where run [prog] config = Just $ interpret config prog - run _ _ = Nothing + where + run [prog] config = Just $ interpret config prog + run _ _ = Nothing interpret :: InterpreterConfig -> FilePath -> IO () interpret config fp = do pr <- newFutharkiState config fp - (tenv, ienv) <- case pr of Left err -> do hPutStrLn stderr err - exitFailure - Right env -> return env + (tenv, ienv) <- case pr of + Left err -> do + hPutStrLn stderr err + exitFailure + Right env -> return env let entry = interpreterEntryPoint config vr <- parseValues "stdin" <$> T.getContents @@ -58,18 +58,21 @@ case M.lookup (T.Term, entry) $ T.envNameMap tenv of Just fname | Just (T.BoundV _ t) <- M.lookup (qualLeaf fname) $ T.envVtable tenv -> - return (fname, toStructural $ snd $ unfoldFunType t) - _ -> do hPutStrLn stderr $ "Invalid entry point: " ++ pretty entry - exitFailure + return (fname, toStructural $ snd $ unfoldFunType t) + _ -> do + hPutStrLn stderr $ "Invalid entry point: " ++ pretty entry + exitFailure case I.interpretFunction ienv (qualLeaf fname) inps of - Left err -> do hPutStrLn stderr err - exitFailure + Left err -> do + hPutStrLn stderr err + exitFailure Right run -> do run' <- runInterpreter' run case run' of - Left err -> do hPrint stderr err - exitFailure + Left err -> do + hPrint stderr err + exitFailure Right res -> case (I.fromTuple res, isTupleRecord ret) of (Just vs, Just ts) -> zipWithM_ putValue vs ts @@ -80,57 +83,76 @@ | I.isEmptyArray v = putStrLn $ I.prettyEmptyArray t v | otherwise = putStrLn $ pretty v -data InterpreterConfig = - InterpreterConfig { interpreterEntryPoint :: Name - , interpreterPrintWarnings :: Bool - } +data InterpreterConfig = InterpreterConfig + { interpreterEntryPoint :: Name, + interpreterPrintWarnings :: Bool + } interpreterConfig :: InterpreterConfig interpreterConfig = InterpreterConfig defaultEntryPoint True options :: [FunOptDescr InterpreterConfig] -options = [ Option "e" ["entry-point"] - (ReqArg (\entry -> Right $ \config -> - config { interpreterEntryPoint = nameFromString entry }) - "NAME") - "The entry point to execute." - , Option "w" ["no-warnings"] - (NoArg $ Right $ \config -> config { interpreterPrintWarnings = False }) - "Do not print warnings." - ] +options = + [ Option + "e" + ["entry-point"] + ( ReqArg + ( \entry -> Right $ \config -> + config {interpreterEntryPoint = nameFromString entry} + ) + "NAME" + ) + "The entry point to execute.", + Option + "w" + ["no-warnings"] + (NoArg $ Right $ \config -> config {interpreterPrintWarnings = False}) + "Do not print warnings." + ] -newFutharkiState :: InterpreterConfig -> FilePath - -> IO (Either String (T.Env, I.Ctx)) +newFutharkiState :: + InterpreterConfig -> + FilePath -> + IO (Either String (T.Env, I.Ctx)) newFutharkiState cfg file = runExceptT $ do (ws, imports, src) <- - badOnLeft show =<< - liftIO (runExceptT (readProgram file) - `catch` \(err::IOException) -> - return (externalErrorS (show err))) + badOnLeft show + =<< liftIO + ( runExceptT (readProgram file) + `catch` \(err :: IOException) -> + return (externalErrorS (show err)) + ) when (interpreterPrintWarnings cfg) $ liftIO $ hPutStr stderr $ show ws let imp = T.mkInitialImport "." - ienv1 <- foldM (\ctx -> badOnLeft show <=< runInterpreter' . I.interpretImport ctx) I.initialCtx $ - map (fmap fileProg) imports - (tenv1, d1, src') <- badOnLeft pretty $ T.checkDec imports src T.initialEnv imp $ - mkOpen "/prelude/prelude" - (tenv2, d2, _) <- badOnLeft pretty $ T.checkDec imports src' tenv1 imp $ - mkOpen $ toPOSIX $ dropExtension file + ienv1 <- + foldM (\ctx -> badOnLeft show <=< runInterpreter' . I.interpretImport ctx) I.initialCtx $ + map (fmap fileProg) imports + (tenv1, d1, src') <- + badOnLeft pretty $ + T.checkDec imports src T.initialEnv imp $ + mkOpen "/prelude/prelude" + (tenv2, d2, _) <- + badOnLeft pretty $ + T.checkDec imports src' tenv1 imp $ + mkOpen $ toPOSIX $ dropExtension file ienv2 <- badOnLeft show =<< runInterpreter' (I.interpretDec ienv1 d1) ienv3 <- badOnLeft show =<< runInterpreter' (I.interpretDec ienv2 d2) return (tenv2, ienv3) - where badOnLeft :: (err -> String) -> Either err a -> ExceptT String IO a - badOnLeft _ (Right x) = return x - badOnLeft p (Left err) = throwError $ p err + where + badOnLeft :: (err -> String) -> Either err a -> ExceptT String IO a + badOnLeft _ (Right x) = return x + badOnLeft p (Left err) = throwError $ p err mkOpen :: FilePath -> UncheckedDec mkOpen f = OpenDec (ModImport f NoInfo mempty) mempty runInterpreter' :: MonadIO m => F I.ExtOp a -> m (Either I.InterpreterError a) runInterpreter' m = runF m (return . Right) intOp - where intOp (I.ExtOpError err) = return $ Left err - intOp (I.ExtOpTrace w v c) = do - liftIO $ putStrLn $ "Trace at " ++ locStr w ++ ": " ++ v - c - intOp (I.ExtOpBreak _ _ c) = c + where + intOp (I.ExtOpError err) = return $ Left err + intOp (I.ExtOpTrace w v c) = do + liftIO $ putStrLn $ "Trace at " ++ locStr w ++ ": " ++ v + c + intOp (I.ExtOpBreak _ _ c) = c diff --git a/src/Futhark/CLI/Test.hs b/src/Futhark/CLI/Test.hs --- a/src/Futhark/CLI/Test.hs +++ b/src/Futhark/CLI/Test.hs @@ -1,8 +1,11 @@ -{-# LANGUAGE OverloadedStrings, FlexibleContexts, LambdaCase #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} + -- | @futhark test@ module Futhark.CLI.Test (main) where -import Control.Applicative.Lift (runErrors, failure, Errors, Lift(..)) +import Control.Applicative.Lift (Errors, Lift (..), failure, runErrors) import Control.Concurrent import Control.Exception import Control.Monad @@ -15,23 +18,22 @@ import qualified Data.Text as T import qualified Data.Text.Encoding as T import qualified Data.Text.IO as T -import System.Console.ANSI -import System.Process.ByteString (readProcessWithExitCode) -import System.Environment -import System.Exit -import System.FilePath -import System.Console.GetOpt -import qualified System.Console.Terminal.Size as Terminal -import System.IO -import Text.Regex.TDFA - import Futhark.Analysis.Metrics import Futhark.Test import Futhark.Util (fancyTerminal) +import Futhark.Util.Console import Futhark.Util.Options import Futhark.Util.Pretty (prettyText) -import Futhark.Util.Console import Futhark.Util.Table +import System.Console.ANSI +import System.Console.GetOpt +import qualified System.Console.Terminal.Size as Terminal +import System.Environment +import System.Exit +import System.FilePath +import System.IO +import System.Process.ByteString (readProcessWithExitCode) +import Text.Regex.TDFA --- Test execution @@ -53,8 +55,8 @@ context :: T.Text -> TestM a -> TestM a context s = withExceptT $ \case - [] -> [] - (e:es') -> (s <> ":\n" <> e):es' + [] -> [] + (e : es') -> (s <> ":\n" <> e) : es' accErrors :: [TestM a] -> TestM [a] accErrors tests = do @@ -65,16 +67,18 @@ accErrors_ :: [TestM a] -> TestM () accErrors_ = void . accErrors -data TestResult = Success - | Failure [T.Text] - deriving (Eq, Show) +data TestResult + = Success + | Failure [T.Text] + deriving (Eq, Show) -data TestCase = TestCase { _testCaseMode :: TestMode - , testCaseProgram :: FilePath - , testCaseTest :: ProgramTest - , _testCasePrograms :: ProgConfig - } - deriving (Show) +data TestCase = TestCase + { _testCaseMode :: TestMode, + testCaseProgram :: FilePath, + testCaseTest :: ProgramTest, + _testCasePrograms :: ProgConfig + } + deriving (Show) instance Eq TestCase where x == y = testCaseProgram x == testCaseProgram y @@ -82,95 +86,119 @@ instance Ord TestCase where x `compare` y = testCaseProgram x `compare` testCaseProgram y -data RunResult = ErrorResult Int SBS.ByteString - | SuccessResult [Value] +data RunResult + = ErrorResult Int SBS.ByteString + | SuccessResult [Value] progNotFound :: T.Text -> T.Text progNotFound s = s <> ": command not found" optimisedProgramMetrics :: ProgConfig -> StructurePipeline -> FilePath -> TestM AstMetrics optimisedProgramMetrics programs pipeline program = - case pipeline of SOACSPipeline -> - check "-s" - KernelsPipeline -> - check "--kernels" - SequentialCpuPipeline -> - check "--cpu" - GpuPipeline -> - check "--gpu" - where check opt = do - futhark <- io $ maybe getExecutablePath return $ configFuthark programs - (code, output, err) <- - io $ readProcessWithExitCode futhark ["dev", opt, "--metrics", program] "" - let output' = T.decodeUtf8 output - case code of - ExitSuccess - | [(m, [])] <- reads $ T.unpack output' -> return m - | otherwise -> throwError $ "Could not read metrics output:\n" <> output' - ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark - ExitFailure _ -> throwError $ T.decodeUtf8 err + case pipeline of + SOACSPipeline -> + check "-s" + KernelsPipeline -> + check "--kernels" + SequentialCpuPipeline -> + check "--cpu" + GpuPipeline -> + check "--gpu" + where + check opt = do + futhark <- io $ maybe getExecutablePath return $ configFuthark programs + (code, output, err) <- + io $ readProcessWithExitCode futhark ["dev", opt, "--metrics", program] "" + let output' = T.decodeUtf8 output + case code of + ExitSuccess + | [(m, [])] <- reads $ T.unpack output' -> return m + | otherwise -> throwError $ "Could not read metrics output:\n" <> output' + ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark + ExitFailure _ -> throwError $ T.decodeUtf8 err testMetrics :: ProgConfig -> FilePath -> StructureTest -> TestM () testMetrics programs program (StructureTest pipeline (AstMetrics expected)) = context "Checking metrics" $ do actual <- optimisedProgramMetrics programs pipeline program accErrors_ $ map (ok actual) $ M.toList expected - where ok (AstMetrics metrics) (name, expected_occurences) = - case M.lookup name metrics of - Nothing - | expected_occurences > 0 -> - throwError $ name <> " should have occurred " <> T.pack (show expected_occurences) <> - " times, but did not occur at all in optimised program." - Just actual_occurences - | expected_occurences /= actual_occurences -> - throwError $ name <> " should have occurred " <> T.pack (show expected_occurences) <> - " times, but occurred " <> T.pack (show actual_occurences) <> " times." - _ -> return () + where + ok (AstMetrics metrics) (name, expected_occurences) = + case M.lookup name metrics of + Nothing + | expected_occurences > 0 -> + throwError $ + name <> " should have occurred " <> T.pack (show expected_occurences) + <> " times, but did not occur at all in optimised program." + Just actual_occurences + | expected_occurences /= actual_occurences -> + throwError $ + name <> " should have occurred " <> T.pack (show expected_occurences) + <> " times, but occurred " + <> T.pack (show actual_occurences) + <> " times." + _ -> return () testWarnings :: [WarningTest] -> SBS.ByteString -> TestM () testWarnings warnings futerr = accErrors_ $ map testWarning warnings - where testWarning (ExpectedWarning regex_s regex) - | not (match regex $ T.unpack $ T.decodeUtf8 futerr) = - throwError $ "Expected warning:\n " <> regex_s <> - "\nGot warnings:\n " <> T.decodeUtf8 futerr - | otherwise = return () + where + testWarning (ExpectedWarning regex_s regex) + | not (match regex $ T.unpack $ T.decodeUtf8 futerr) = + throwError $ + "Expected warning:\n " <> regex_s + <> "\nGot warnings:\n " + <> T.decodeUtf8 futerr + | otherwise = return () runTestCase :: TestCase -> TestM () runTestCase (TestCase mode program testcase progs) = do futhark <- io $ maybe getExecutablePath return $ configFuthark progs case testAction testcase of - CompileTimeFailure expected_error -> - context (mconcat ["Type-checking with '", T.pack futhark, - " check ", T.pack program, "'"]) $ do - (code, _, err) <- - io $ readProcessWithExitCode futhark ["check", program] "" - case code of - ExitSuccess -> throwError "Expected failure\n" - ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark - ExitFailure 1 -> throwError $ T.decodeUtf8 err - ExitFailure _ -> checkError expected_error err - - RunCases{} | mode == TypeCheck -> do + context + ( mconcat + [ "Type-checking with '", + T.pack futhark, + " check ", + T.pack program, + "'" + ] + ) + $ do + (code, _, err) <- + io $ readProcessWithExitCode futhark ["check", program] "" + case code of + ExitSuccess -> throwError "Expected failure\n" + ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark + ExitFailure 1 -> throwError $ T.decodeUtf8 err + ExitFailure _ -> checkError expected_error err + RunCases {} | mode == TypeCheck -> do let options = ["check", program] ++ configExtraCompilerOptions progs - context (mconcat ["Type-checking with '", T.pack futhark, - " check ", T.pack program, "'"]) $ do - (code, _, err) <- io $ readProcessWithExitCode futhark options "" - - case code of - ExitSuccess -> return () - ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark - ExitFailure _ -> throwError $ T.decodeUtf8 err + context + ( mconcat + [ "Type-checking with '", + T.pack futhark, + " check ", + T.pack program, + "'" + ] + ) + $ do + (code, _, err) <- io $ readProcessWithExitCode futhark options "" + case code of + ExitSuccess -> return () + ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark + ExitFailure _ -> throwError $ T.decodeUtf8 err RunCases ios structures warnings -> do -- Compile up-front and reuse same executable for several entry points. let backend = configBackend progs extra_options = configExtraCompilerOptions progs unless (mode == Compile) $ context "Generating reference outputs" $ - -- We probably get the concurrency at the test program level, - -- so force just one data set at a time here. - ensureReferenceOutput (Just 1) futhark "c" program ios + -- We probably get the concurrency at the test program level, + -- so force just one data set at a time here. + ensureReferenceOutput (Just 1) futhark "c" program ios unless (mode == Interpreted) $ context ("Compiling with --backend=" <> T.pack backend) $ do compileTestProgram extra_options futhark backend program warnings @@ -178,34 +206,40 @@ unless (mode == Compile) $ do (tuning_opts, _) <- liftIO $ determineTuning (configTuning progs) program - let progs' = progs { configExtraOptions = - tuning_opts ++ configExtraOptions progs } + let progs' = + progs + { configExtraOptions = + tuning_opts ++ configExtraOptions progs + } context "Running compiled program" $ accErrors_ $ map (runCompiledEntry program progs') ios unless (mode == Compile || mode == Compiled) $ context "Interpreting" $ accErrors_ $ map (runInterpretedEntry futhark program) ios -runInterpretedEntry :: String -> FilePath -> InputOutputs -> TestM() +runInterpretedEntry :: String -> FilePath -> InputOutputs -> TestM () runInterpretedEntry futhark program (InputOutputs entry run_cases) = let dir = takeDirectory program runInterpretedCase run@(TestRun _ inputValues _ index _) = unless ("compiled" `elem` runTags run) $ - context ("Entry point: " <> entry - <> "; dataset: " <> T.pack (runDescription run)) $ do - - input <- T.unlines . map prettyText <$> getValues dir inputValues - expectedResult' <- getExpectedResult program entry run - (code, output, err) <- - io $ readProcessWithExitCode futhark ["run", "-e", T.unpack entry, program] $ - T.encodeUtf8 input - case code of - ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark - - _ -> compareResult entry index program expectedResult' - =<< runResult program code output err - - in accErrors_ $ map runInterpretedCase run_cases + context + ( "Entry point: " <> entry + <> "; dataset: " + <> T.pack (runDescription run) + ) + $ do + input <- T.unlines . map prettyText <$> getValues dir inputValues + expectedResult' <- getExpectedResult program entry run + (code, output, err) <- + io $ + readProcessWithExitCode futhark ["run", "-e", T.unpack entry, program] $ + T.encodeUtf8 input + case code of + ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark + _ -> + compareResult entry index program expectedResult' + =<< runResult program code output err + in accErrors_ $ map runInterpretedCase run_cases runCompiledEntry :: FilePath -> ProgConfig -> InputOutputs -> TestM () runCompiledEntry program progs (InputOutputs entry run_cases) = @@ -220,29 +254,34 @@ extra_options = configExtraOptions progs runCompiledCase run@(TestRun _ inputValues _ index _) = - context ("Entry point: " <> entry - <> "; dataset: " <> T.pack (runDescription run)) $ do - expected <- getExpectedResult program entry run - (progCode, output, progerr) <- - runProgram runner extra_options program entry inputValues - compareResult entry index program expected - =<< runResult program progCode output progerr - - in context ("Running " <> T.pack (unwords $ binpath : entry_options ++ extra_options)) $ - accErrors_ $ map runCompiledCase run_cases + context + ( "Entry point: " <> entry + <> "; dataset: " + <> T.pack (runDescription run) + ) + $ do + expected <- getExpectedResult program entry run + (progCode, output, progerr) <- + runProgram runner extra_options program entry inputValues + compareResult entry index program expected + =<< runResult program progCode output progerr + in context ("Running " <> T.pack (unwords $ binpath : entry_options ++ extra_options)) $ + accErrors_ $ map runCompiledCase run_cases checkError :: ExpectedError -> SBS.ByteString -> TestM () checkError (ThisError regex_s regex) err | not (match regex $ T.unpack $ T.decodeUtf8 err) = - throwError $ "Expected error:\n " <> regex_s <> - "\nGot error:\n " <> T.decodeUtf8 err + throwError $ + "Expected error:\n " <> regex_s + <> "\nGot error:\n " + <> T.decodeUtf8 err checkError _ _ = return () runResult :: FilePath -> ExitCode -> SBS.ByteString -> SBS.ByteString -> TestM RunResult runResult program ExitSuccess stdout_s _ = case valuesFromByteString "stdout" $ LBS.fromStrict stdout_s of - Left e -> do + Left e -> do let actualf = program `addExtension` "actual" io $ SBS.writeFile actualf stdout_s throwError $ T.pack e <> "\n(See " <> T.pack actualf <> ")" @@ -255,28 +294,41 @@ (_, futerr) <- compileProgram extra_options futhark backend program testWarnings warnings futerr -compareResult :: T.Text -> Int -> FilePath -> ExpectedResult [Value] -> RunResult - -> TestM () -compareResult _ _ _ (Succeeds Nothing) SuccessResult{} = +compareResult :: + T.Text -> + Int -> + FilePath -> + ExpectedResult [Value] -> + RunResult -> + TestM () +compareResult _ _ _ (Succeeds Nothing) SuccessResult {} = return () compareResult entry index program (Succeeds (Just expectedResult)) (SuccessResult actualResult) = case compareValues1 actualResult expectedResult of Just mismatch -> do let actualf = program <.> T.unpack entry <.> show index <.> "actual" expectedf = program <.> T.unpack entry <.> show index <.> "expected" - io $ SBS.writeFile actualf $ - T.encodeUtf8 $ T.unlines $ map prettyText actualResult - io $ SBS.writeFile expectedf $ - T.encodeUtf8 $ T.unlines $ map prettyText expectedResult - throwError $ T.pack actualf <> " and " <> T.pack expectedf <> - " do not match:\n" <> T.pack (show mismatch) <> "\n" + io $ + SBS.writeFile actualf $ + T.encodeUtf8 $ T.unlines $ map prettyText actualResult + io $ + SBS.writeFile expectedf $ + T.encodeUtf8 $ T.unlines $ map prettyText expectedResult + throwError $ + T.pack actualf <> " and " <> T.pack expectedf + <> " do not match:\n" + <> T.pack (show mismatch) + <> "\n" Nothing -> return () compareResult _ _ _ (RunTimeFailure expectedError) (ErrorResult _ actualError) = checkError expectedError actualError compareResult _ _ _ (Succeeds _) (ErrorResult code err) = - throwError $ "Program failed with error code " <> - T.pack (show code) <> " and stderr:\n " <> T.decodeUtf8 err + throwError $ + "Program failed with error code " + <> T.pack (show code) + <> " and stderr:\n " + <> T.decodeUtf8 err compareResult _ _ _ (RunTimeFailure f) (SuccessResult _) = throwError $ "Program succeeded, but expected failure:\n " <> T.pack (show f) @@ -284,21 +336,23 @@ --- Test manager --- -data TestStatus = TestStatus { testStatusRemain :: [TestCase] - , testStatusRun :: [TestCase] - , testStatusTotal :: Int - , testStatusFail :: Int - , testStatusPass :: Int - , testStatusRuns :: Int - , testStatusRunsRemain :: Int - , testStatusRunPass :: Int - , testStatusRunFail :: Int - } +data TestStatus = TestStatus + { testStatusRemain :: [TestCase], + testStatusRun :: [TestCase], + testStatusTotal :: Int, + testStatusFail :: Int, + testStatusPass :: Int, + testStatusRuns :: Int, + testStatusRunsRemain :: Int, + testStatusRunPass :: Int, + testStatusRunFail :: Int + } catching :: IO TestResult -> IO TestResult catching m = m `catch` save - where save :: SomeException -> IO TestResult - save e = return $ Failure [T.pack $ show e] + where + save :: SomeException -> IO TestResult + save e = return $ Failure [T.pack $ show e] doTest :: TestCase -> IO TestResult doTest = catching . runTestM . runTestCase @@ -307,8 +361,9 @@ makeTestCase config mode (file, spec) = TestCase mode file spec $ configPrograms config -data ReportMsg = TestStarted TestCase - | TestDone TestCase TestResult +data ReportMsg + = TestStarted TestCase + | TestDone TestCase TestResult runTest :: MVar TestCase -> MVar ReportMsg -> IO () runTest testmvar resmvar = forever $ do @@ -324,39 +379,42 @@ -- | Exclude those test cases that have tags we do not wish to run. excludeCases :: TestConfig -> TestCase -> TestCase excludeCases config tcase = - tcase { testCaseTest = onTest $ testCaseTest tcase } - where onTest (ProgramTest desc tags action) = - ProgramTest desc tags $ onAction action - onAction (RunCases ios stest wtest) = - RunCases (map onIOs ios) stest wtest - onAction action = action - onIOs (InputOutputs entry runs) = - InputOutputs entry $ filter (not . any excluded . runTags) runs - excluded = (`elem` configExclude config) . T.pack + tcase {testCaseTest = onTest $ testCaseTest tcase} + where + onTest (ProgramTest desc tags action) = + ProgramTest desc tags $ onAction action + onAction (RunCases ios stest wtest) = + RunCases (map onIOs ios) stest wtest + onAction action = action + onIOs (InputOutputs entry runs) = + InputOutputs entry $ filter (not . any excluded . runTags) runs + excluded = (`elem` configExclude config) . T.pack statusTable :: TestStatus -> String statusTable ts = buildTable rows 1 - where rows = - [ [ mkEntry "", passed, failed, mkEntry "remaining"] - , map mkEntry ["programs", passedProgs, failedProgs, remainProgs'] - , map mkEntry ["runs", passedRuns, failedRuns, remainRuns'] - ] - passed = ("passed", [SetColor Foreground Vivid Green]) - failed = ("failed", [SetColor Foreground Vivid Red]) - passedProgs = show $ testStatusPass ts - failedProgs = show $ testStatusFail ts - totalProgs = show $ testStatusTotal ts - totalRuns = show $ testStatusRuns ts - passedRuns = show $ testStatusRunPass ts - failedRuns = show $ testStatusRunFail ts - remainProgs = show . length $ testStatusRemain ts - remainProgs' = remainProgs ++ "/" ++ totalProgs - remainRuns = show $ testStatusRunsRemain ts - remainRuns' = remainRuns ++ "/" ++ totalRuns + where + rows = + [ [mkEntry "", passed, failed, mkEntry "remaining"], + map mkEntry ["programs", passedProgs, failedProgs, remainProgs'], + map mkEntry ["runs", passedRuns, failedRuns, remainRuns'] + ] + passed = ("passed", [SetColor Foreground Vivid Green]) + failed = ("failed", [SetColor Foreground Vivid Red]) + passedProgs = show $ testStatusPass ts + failedProgs = show $ testStatusFail ts + totalProgs = show $ testStatusTotal ts + totalRuns = show $ testStatusRuns ts + passedRuns = show $ testStatusRunPass ts + failedRuns = show $ testStatusRunFail ts + remainProgs = show . length $ testStatusRemain ts + remainProgs' = remainProgs ++ "/" ++ totalProgs + remainRuns = show $ testStatusRunsRemain ts + remainRuns' = remainRuns ++ "/" ++ totalRuns tableLines :: Int tableLines = 1 + (length . lines $ blankTable) - where blankTable = statusTable $ TestStatus [] [] 0 0 0 0 0 0 0 + where + blankTable = statusTable $ TestStatus [] [] 0 0 0 0 0 0 0 spaceTable :: IO () spaceTable = putStr $ replicate tableLines '\n' @@ -367,23 +425,29 @@ putStrLn $ statusTable ts clearLine w <- maybe 80 Terminal.width <$> Terminal.size - putStrLn $ atMostChars (w-length labelstr) running - where running = labelstr ++ (unwords . reverse . map testCaseProgram . testStatusRun) ts - labelstr = "Now testing: " + putStrLn $ atMostChars (w - length labelstr) running + where + running = labelstr ++ (unwords . reverse . map testCaseProgram . testStatusRun) ts + labelstr = "Now testing: " moveCursorToTableTop :: IO () moveCursorToTableTop = cursorUpLine tableLines atMostChars :: Int -> String -> String -atMostChars n s | length s > n = take (n-3) s ++ "..." - | otherwise = s +atMostChars n s + | length s > n = take (n -3) s ++ "..." + | otherwise = s reportText :: TestStatus -> IO () reportText ts = - putStr $ "(" ++ show (testStatusFail ts) ++ " failed, " ++ - show (testStatusPass ts) ++ " passed, " ++ - show num_remain ++ " to go).\n" - where num_remain = length $ testStatusRemain ts + putStr $ + "(" ++ show (testStatusFail ts) ++ " failed, " + ++ show (testStatusPass ts) + ++ " passed, " + ++ show num_remain + ++ " to go).\n" + where + num_remain = length $ testStatusRemain ts runTests :: TestConfig -> [FilePath] -> IO () runTests config paths = do @@ -393,8 +457,9 @@ hSetBuffering stdout LineBuffering let mode = configTestMode config - all_tests <- map (makeTestCase config mode) <$> - testSpecsFromPathsOrDie paths + all_tests <- + map (makeTestCase config mode) + <$> testSpecsFromPathsOrDie paths testmvar <- newEmptyMVar reportmvar <- newEmptyMVar concurrency <- maybe getNumCapabilities pure $ configConcurrency config @@ -405,16 +470,20 @@ let fancy = not (configLineOutput config) && fancyTerminal - report | fancy = reportTable - | otherwise = reportText - clear | fancy = clearFromCursorToScreenEnd - | otherwise = putStr "\n" + report + | fancy = reportTable + | otherwise = reportText + clear + | fancy = clearFromCursorToScreenEnd + | otherwise = putStr "\n" numTestCases tc = case testAction $ testCaseTest tc of CompileTimeFailure _ -> 1 - RunCases ios sts wts -> (length . concat) (iosTestRuns <$> ios) - + length sts + length wts + RunCases ios sts wts -> + (length . concat) (iosTestRuns <$> ios) + + length sts + + length wts getResults ts | null (testStatusRemain ts) = report ts >> return ts @@ -427,174 +496,222 @@ putStr $ "Started testing " <> testCaseProgram test <> " " getResults $ ts {testStatusRun = test : testStatusRun ts} TestDone test res -> do - let ts' = ts { testStatusRemain = test `delete` testStatusRemain ts - , testStatusRun = test `delete` testStatusRun ts - , testStatusRunsRemain = testStatusRunsRemain ts - - numTestCases test - } + let ts' = + ts + { testStatusRemain = test `delete` testStatusRemain ts, + testStatusRun = test `delete` testStatusRun ts, + testStatusRunsRemain = + testStatusRunsRemain ts + - numTestCases test + } case res of Success -> do - let ts'' = ts' { testStatusRunPass = - testStatusRunPass ts' + numTestCases test - } + let ts'' = + ts' + { testStatusRunPass = + testStatusRunPass ts' + numTestCases test + } unless fancy $ putStr $ "Finished testing " <> testCaseProgram test <> " " - getResults $ ts'' { testStatusPass = testStatusPass ts + 1} + getResults $ ts'' {testStatusPass = testStatusPass ts + 1} Failure s -> do when fancy moveCursorToTableTop clear T.putStr $ (T.pack (inRed $ testCaseProgram test) <> ":\n") <> T.unlines s when fancy spaceTable - getResults $ ts' { testStatusFail = testStatusFail ts' + 1 - , testStatusRunPass = testStatusRunPass ts' - + numTestCases test - length s - - , testStatusRunFail = testStatusRunFail ts' - + length s - } + getResults $ + ts' + { testStatusFail = testStatusFail ts' + 1, + testStatusRunPass = + testStatusRunPass ts' + + numTestCases test - length s, + testStatusRunFail = + testStatusRunFail ts' + + length s + } when fancy spaceTable - ts <- getResults TestStatus { testStatusRemain = included - , testStatusRun = [] - , testStatusTotal = length included - , testStatusFail = 0 - , testStatusPass = 0 - , testStatusRuns = sum $ map numTestCases included - , testStatusRunsRemain = sum $ map numTestCases included - , testStatusRunPass = 0 - , testStatusRunFail = 0 - } + ts <- + getResults + TestStatus + { testStatusRemain = included, + testStatusRun = [], + testStatusTotal = length included, + testStatusFail = 0, + testStatusPass = 0, + testStatusRuns = sum $ map numTestCases included, + testStatusRunsRemain = sum $ map numTestCases included, + testStatusRunPass = 0, + testStatusRunFail = 0 + } -- Removes "Now testing" output. when fancy $ cursorUpLine 1 >> clearLine - let excluded_str | null excluded = "" - | otherwise = " (" ++ show (length excluded) ++ " program(s) excluded).\n" + let excluded_str + | null excluded = "" + | otherwise = " (" ++ show (length excluded) ++ " program(s) excluded).\n" putStr excluded_str - exitWith $ case testStatusFail ts of 0 -> ExitSuccess - _ -> ExitFailure 1 + exitWith $ case testStatusFail ts of + 0 -> ExitSuccess + _ -> ExitFailure 1 --- --- Configuration and command line parsing --- data TestConfig = TestConfig - { configTestMode :: TestMode - , configPrograms :: ProgConfig - , configExclude :: [T.Text] - , configLineOutput :: Bool - , configConcurrency :: Maybe Int - } + { configTestMode :: TestMode, + configPrograms :: ProgConfig, + configExclude :: [T.Text], + configLineOutput :: Bool, + configConcurrency :: Maybe Int + } defaultConfig :: TestConfig -defaultConfig = TestConfig { configTestMode = Everything - , configExclude = [ "disable" ] - , configPrograms = - ProgConfig - { configBackend = "c" - , configFuthark = Nothing - , configRunner = "" - , configExtraOptions = [] - , configExtraCompilerOptions = [] - , configTuning = Just "tuning" - } - , configLineOutput = False - , configConcurrency = Nothing - } +defaultConfig = + TestConfig + { configTestMode = Everything, + configExclude = ["disable"], + configPrograms = + ProgConfig + { configBackend = "c", + configFuthark = Nothing, + configRunner = "", + configExtraOptions = [], + configExtraCompilerOptions = [], + configTuning = Just "tuning" + }, + configLineOutput = False, + configConcurrency = Nothing + } data ProgConfig = ProgConfig - { configBackend :: String - , configFuthark :: Maybe FilePath - , configRunner :: FilePath - , configExtraCompilerOptions :: [String] - , configTuning :: Maybe String - , configExtraOptions :: [String] - -- ^ Extra options passed to the programs being run. - } - deriving (Show) + { configBackend :: String, + configFuthark :: Maybe FilePath, + configRunner :: FilePath, + configExtraCompilerOptions :: [String], + configTuning :: Maybe String, + -- | Extra options passed to the programs being run. + configExtraOptions :: [String] + } + deriving (Show) changeProgConfig :: (ProgConfig -> ProgConfig) -> TestConfig -> TestConfig -changeProgConfig f config = config { configPrograms = f $ configPrograms config } +changeProgConfig f config = config {configPrograms = f $ configPrograms config} setBackend :: FilePath -> ProgConfig -> ProgConfig setBackend backend config = - config { configBackend = backend } + config {configBackend = backend} setFuthark :: FilePath -> ProgConfig -> ProgConfig setFuthark futhark config = - config { configFuthark = Just futhark } + config {configFuthark = Just futhark} setRunner :: FilePath -> ProgConfig -> ProgConfig setRunner runner config = - config { configRunner = runner } + config {configRunner = runner} addCompilerOption :: String -> ProgConfig -> ProgConfig addCompilerOption option config = - config { configExtraCompilerOptions = configExtraCompilerOptions config ++ [option] } + config {configExtraCompilerOptions = configExtraCompilerOptions config ++ [option]} addOption :: String -> ProgConfig -> ProgConfig addOption option config = - config { configExtraOptions = configExtraOptions config ++ [option] } + config {configExtraOptions = configExtraOptions config ++ [option]} -data TestMode = TypeCheck - | Compile - | Compiled - | Interpreted - | Everything - deriving (Eq, Show) +data TestMode + = TypeCheck + | Compile + | Compiled + | Interpreted + | Everything + deriving (Eq, Show) commandLineOptions :: [FunOptDescr TestConfig] -commandLineOptions = [ - Option "t" ["typecheck"] - (NoArg $ Right $ \config -> config { configTestMode = TypeCheck }) - "Only perform type-checking" - , Option "i" ["interpreted"] - (NoArg $ Right $ \config -> config { configTestMode = Interpreted }) - "Only interpret" - , Option "c" ["compiled"] - (NoArg $ Right $ \config -> config { configTestMode = Compiled }) - "Only run compiled code" - , Option "C" ["compile"] - (NoArg $ Right $ \config -> config { configTestMode = Compile }) - "Only compile, do not run." - , Option [] ["no-terminal", "notty"] - (NoArg $ Right $ \config -> config { configLineOutput = True }) - "Provide simpler line-based output." - , Option [] ["backend"] - (ReqArg (Right . changeProgConfig . setBackend) "BACKEND") - "Backend used for compilation (defaults to 'c')." - , Option [] ["futhark"] - (ReqArg (Right . changeProgConfig . setFuthark) "PROGRAM") - "Program to run for subcommands (defaults to same binary as 'futhark test')." - , Option [] ["runner"] - (ReqArg (Right . changeProgConfig . setRunner) "PROGRAM") - "The program used to run the Futhark-generated programs (defaults to nothing)." - , Option [] ["exclude"] - (ReqArg (\tag -> - Right $ \config -> - config { configExclude = T.pack tag : configExclude config }) - "TAG") - "Exclude test programs that define this tag." - , Option "p" ["pass-option"] - (ReqArg (Right . changeProgConfig . addOption) "OPT") - "Pass this option to programs being run." - , Option [] ["pass-compiler-option"] - (ReqArg (Right . changeProgConfig . addCompilerOption) "OPT") - "Pass this option to the compiler (or typechecker if in -t mode)." - , Option [] ["no-tuning"] - (NoArg $ Right $ changeProgConfig $ \config -> config { configTuning = Nothing }) - "Do not load tuning files." - , Option [] ["concurrency"] - (ReqArg (\n -> - case reads n of - [(n', "")] - | n' > 0 -> - Right $ \config -> config { configConcurrency = Just n' } - _ -> - Left $ error $ "'" ++ n ++ "' is not a positive integer.") - "NUM") - "Number of tests to run concurrently." +commandLineOptions = + [ Option + "t" + ["typecheck"] + (NoArg $ Right $ \config -> config {configTestMode = TypeCheck}) + "Only perform type-checking", + Option + "i" + ["interpreted"] + (NoArg $ Right $ \config -> config {configTestMode = Interpreted}) + "Only interpret", + Option + "c" + ["compiled"] + (NoArg $ Right $ \config -> config {configTestMode = Compiled}) + "Only run compiled code", + Option + "C" + ["compile"] + (NoArg $ Right $ \config -> config {configTestMode = Compile}) + "Only compile, do not run.", + Option + [] + ["no-terminal", "notty"] + (NoArg $ Right $ \config -> config {configLineOutput = True}) + "Provide simpler line-based output.", + Option + [] + ["backend"] + (ReqArg (Right . changeProgConfig . setBackend) "BACKEND") + "Backend used for compilation (defaults to 'c').", + Option + [] + ["futhark"] + (ReqArg (Right . changeProgConfig . setFuthark) "PROGRAM") + "Program to run for subcommands (defaults to same binary as 'futhark test').", + Option + [] + ["runner"] + (ReqArg (Right . changeProgConfig . setRunner) "PROGRAM") + "The program used to run the Futhark-generated programs (defaults to nothing).", + Option + [] + ["exclude"] + ( ReqArg + ( \tag -> + Right $ \config -> + config {configExclude = T.pack tag : configExclude config} + ) + "TAG" + ) + "Exclude test programs that define this tag.", + Option + "p" + ["pass-option"] + (ReqArg (Right . changeProgConfig . addOption) "OPT") + "Pass this option to programs being run.", + Option + [] + ["pass-compiler-option"] + (ReqArg (Right . changeProgConfig . addCompilerOption) "OPT") + "Pass this option to the compiler (or typechecker if in -t mode).", + Option + [] + ["no-tuning"] + (NoArg $ Right $ changeProgConfig $ \config -> config {configTuning = Nothing}) + "Do not load tuning files.", + Option + [] + ["concurrency"] + ( ReqArg + ( \n -> + case reads n of + [(n', "")] + | n' > 0 -> + Right $ \config -> config {configConcurrency = Just n'} + _ -> + Left $ error $ "'" ++ n ++ "' is not a positive integer." + ) + "NUM" + ) + "Number of tests to run concurrently." ] -- | Run @futhark test@. 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 @@ -1,27 +1,31 @@ {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TupleSections #-} + -- | Code generation for CUDA. module Futhark.CodeGen.Backends.CCUDA - ( compileProg - , GC.CParts(..) - , GC.asLibrary - , GC.asExecutable - ) where + ( compileProg, + GC.CParts (..), + GC.asLibrary, + GC.asExecutable, + ) +where import Control.Monad import Data.List (intercalate) import Data.Maybe (catMaybes) -import qualified Language.C.Quote.OpenCL as C - +import Futhark.CodeGen.Backends.CCUDA.Boilerplate +import Futhark.CodeGen.Backends.COpenCL.Boilerplate (commonOptions) import qualified Futhark.CodeGen.Backends.GenericC as GC +import Futhark.CodeGen.Backends.GenericC.Options +import Futhark.CodeGen.ImpCode.OpenCL import qualified Futhark.CodeGen.ImpGen.CUDA as ImpGen -import Futhark.IR.KernelsMem - hiding (GetSize, CmpSizeLe, GetSizeMax) +import Futhark.IR.KernelsMem hiding + ( CmpSizeLe, + GetSize, + GetSizeMax, + ) import Futhark.MonadFreshNames -import Futhark.CodeGen.ImpCode.OpenCL -import Futhark.CodeGen.Backends.COpenCL.Boilerplate (commonOptions) -import Futhark.CodeGen.Backends.CCUDA.Boilerplate -import Futhark.CodeGen.Backends.GenericC.Options +import qualified Language.C.Quote.OpenCL as C -- | Compile the program to C with calls to CUDA. compileProg :: MonadFreshNames m => Prog KernelsMem -> m (ImpGen.Warnings, GC.CParts) @@ -29,73 +33,111 @@ (ws, Program cuda_code cuda_prelude kernels _ sizes failures prog') <- ImpGen.compileProg prog let cost_centres = - [copyDevToDev, copyDevToHost, copyHostToDev, - copyScalarToDev, copyScalarFromDev] - extra = generateBoilerplate cuda_code cuda_prelude - cost_centres kernels sizes failures - (ws,) <$> - GC.compileProg "cuda" operations extra cuda_includes - [Space "device", DefaultSpace] cliOptions prog' + [ copyDevToDev, + copyDevToHost, + copyHostToDev, + copyScalarToDev, + copyScalarFromDev + ] + extra = + generateBoilerplate + cuda_code + cuda_prelude + cost_centres + kernels + sizes + failures + (ws,) + <$> GC.compileProg + "cuda" + operations + extra + cuda_includes + [Space "device", DefaultSpace] + cliOptions + prog' where operations :: GC.Operations OpenCL () - operations = GC.defaultOperations - { GC.opsWriteScalar = writeCUDAScalar - , GC.opsReadScalar = readCUDAScalar - , GC.opsAllocate = allocateCUDABuffer - , GC.opsDeallocate = deallocateCUDABuffer - , GC.opsCopy = copyCUDAMemory - , GC.opsStaticArray = staticCUDAArray - , GC.opsMemoryType = cudaMemoryType - , GC.opsCompiler = callKernel - , GC.opsFatMemory = True - } - cuda_includes = unlines [ "#include " - , "#include " - , "#include " - ] + operations = + GC.defaultOperations + { GC.opsWriteScalar = writeCUDAScalar, + GC.opsReadScalar = readCUDAScalar, + GC.opsAllocate = allocateCUDABuffer, + GC.opsDeallocate = deallocateCUDABuffer, + GC.opsCopy = copyCUDAMemory, + GC.opsStaticArray = staticCUDAArray, + GC.opsMemoryType = cudaMemoryType, + GC.opsCompiler = callKernel, + GC.opsFatMemory = True, + GC.opsCritical = + ( [C.citems|CUDA_SUCCEED(cuCtxPushCurrent(ctx->cuda.cu_ctx));|], + [C.citems|CUDA_SUCCEED(cuCtxPopCurrent(&ctx->cuda.cu_ctx));|] + ) + } + cuda_includes = + unlines + [ "#include ", + "#include ", + "#include " + ] cliOptions :: [Option] cliOptions = - commonOptions ++ - [ Option { optionLongName = "dump-cuda" - , optionShortName = Nothing - , optionArgument = RequiredArgument "FILE" - , optionAction = [C.cstm|{futhark_context_config_dump_program_to(cfg, optarg); + commonOptions + ++ [ Option + { optionLongName = "dump-cuda", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Dump the embedded CUDA kernels to the indicated file.", + optionAction = + [C.cstm|{futhark_context_config_dump_program_to(cfg, optarg); entry_point = NULL;}|] - } - , Option { optionLongName = "load-cuda" - , optionShortName = Nothing - , optionArgument = RequiredArgument "FILE" - , optionAction = [C.cstm|futhark_context_config_load_program_from(cfg, optarg);|] - } - , Option { optionLongName = "dump-ptx" - , optionShortName = Nothing - , optionArgument = RequiredArgument "FILE" - , optionAction = [C.cstm|{futhark_context_config_dump_ptx_to(cfg, optarg); + }, + Option + { optionLongName = "load-cuda", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Instead of using the embedded CUDA kernels, load them from the indicated file.", + optionAction = [C.cstm|futhark_context_config_load_program_from(cfg, optarg);|] + }, + Option + { optionLongName = "dump-ptx", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Dump the PTX-compiled version of the embedded kernels to the indicated file.", + optionAction = + [C.cstm|{futhark_context_config_dump_ptx_to(cfg, optarg); entry_point = NULL;}|] - } - , Option { optionLongName = "load-ptx" - , optionShortName = Nothing - , optionArgument = RequiredArgument "FILE" - , optionAction = [C.cstm|futhark_context_config_load_ptx_from(cfg, optarg);|] - } - , Option { optionLongName = "nvrtc-option" - , optionShortName = Nothing - , optionArgument = RequiredArgument "OPT" - , optionAction = [C.cstm|futhark_context_config_add_nvrtc_option(cfg, optarg);|] - } - , Option { optionLongName = "profile" - , optionShortName = Just 'P' - , optionArgument = NoArgument - , optionAction = [C.cstm|futhark_context_config_set_profiling(cfg, 1);|] + }, + Option + { optionLongName = "load-ptx", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Load PTX code from the indicated file.", + optionAction = [C.cstm|futhark_context_config_load_ptx_from(cfg, optarg);|] + }, + Option + { optionLongName = "nvrtc-option", + optionShortName = Nothing, + optionArgument = RequiredArgument "OPT", + optionDescription = "Add an additional build option to the string passed to NVRTC.", + optionAction = [C.cstm|futhark_context_config_add_nvrtc_option(cfg, optarg);|] + }, + Option + { optionLongName = "profile", + optionShortName = Just 'P', + optionArgument = NoArgument, + optionDescription = "Gather profiling data while executing and print out a summary at the end.", + optionAction = [C.cstm|futhark_context_config_set_profiling(cfg, 1);|] } - ] + ] writeCUDAScalar :: GC.WriteScalar OpenCL () writeCUDAScalar mem idx t "device" _ val = do val' <- newVName "write_tmp" let (bef, aft) = profilingEnclosure copyScalarToDev - GC.item [C.citem|{$ty:t $id:val' = $exp:val; + GC.item + [C.citem|{$ty:t $id:val' = $exp:val; $items:bef CUDA_SUCCEED( cuMemcpyHtoD($exp:mem + $exp:idx * sizeof($ty:t), @@ -110,7 +152,8 @@ readCUDAScalar mem idx t "device" _ = do val <- newVName "read_res" let (bef, aft) = profilingEnclosure copyScalarFromDev - mapM_ GC.item + mapM_ + GC.item [C.citems| $ty:t $id:val; { @@ -143,7 +186,8 @@ copyCUDAMemory dstmem dstidx dstSpace srcmem srcidx srcSpace nbytes = do let (fn, prof) = memcpyFun dstSpace srcSpace (bef, aft) = profilingEnclosure prof - GC.item [C.citem|{ + GC.item + [C.citem|{ $items:bef CUDA_SUCCEED( $id:fn($exp:dstmem + $exp:dstidx, @@ -153,11 +197,15 @@ } |] where - memcpyFun DefaultSpace (Space "device") = ("cuMemcpyDtoH", copyDevToHost) - memcpyFun (Space "device") DefaultSpace = ("cuMemcpyHtoD", copyHostToDev) - memcpyFun (Space "device") (Space "device") = ("cuMemcpy", copyDevToDev) - memcpyFun _ _ = error $ "Cannot copy to '" ++ show dstSpace ++ - "' from '" ++ show srcSpace ++ "'." + memcpyFun DefaultSpace (Space "device") = ("cuMemcpyDtoH", copyDevToHost) + memcpyFun (Space "device") DefaultSpace = ("cuMemcpyHtoD", copyHostToDev) + memcpyFun (Space "device") (Space "device") = ("cuMemcpy", copyDevToDev) + memcpyFun _ _ = + error $ + "Cannot copy to '" ++ show dstSpace + ++ "' from '" + ++ show srcSpace + ++ "'." staticCUDAArray :: GC.StaticArray OpenCL () staticCUDAArray name "device" t vs = do @@ -174,7 +222,8 @@ -- Fake a memory block. GC.contextField (C.toIdent name mempty) [C.cty|struct memblock_device|] Nothing -- During startup, copy the data to where we need it. - GC.atInit [C.cstm|{ + GC.atInit + [C.cstm|{ ctx->$id:name.references = NULL; ctx->$id:name.size = 0; CUDA_SUCCEED(cuMemAlloc(&ctx->$id:name.mem, @@ -186,8 +235,9 @@ }|] GC.item [C.citem|struct memblock_device $id:name = ctx->$id:name;|] staticCUDAArray _ space _ _ = - error $ "CUDA backend cannot create static array in '" ++ space - ++ "' memory space" + error $ + "CUDA backend cannot create static array in '" ++ space + ++ "' memory space" cudaMemoryType :: GC.MemoryType OpenCL () cudaMemoryType "device" = return [C.cty|typename CUdeviceptr|] @@ -202,9 +252,9 @@ GC.stm [C.cstm|$id:v = ctx->sizes.$id:key <= $exp:x';|] callKernel (GetSizeMax v size_class) = let field = "max_" ++ cudaSizeClass size_class - in GC.stm [C.cstm|$id:v = ctx->cuda.$id:field;|] + in GC.stm [C.cstm|$id:v = ctx->cuda.$id:field;|] where - cudaSizeClass SizeThreshold{} = "threshold" + cudaSizeClass SizeThreshold {} = "threshold" cudaSizeClass SizeGroup = "block_size" cudaSizeClass SizeNumGroups = "grid_size" cudaSizeClass SizeTile = "tile_size" @@ -219,24 +269,35 @@ shared_offsets_sc = mkOffsets shared_sizes shared_args = zip shared_offsets shared_offsets_sc shared_tot = last shared_offsets_sc - forM_ shared_args $ \(arg,offset) -> + forM_ shared_args $ \(arg, offset) -> GC.decl [C.cdecl|unsigned int $id:arg = $exp:offset;|] (grid_x, grid_y, grid_z) <- mkDims <$> mapM GC.compileExp num_blocks (block_x, block_y, block_z) <- mkDims <$> mapM GC.compileExp block_size let perm_args - | length num_blocks == 3 = [ [C.cinit|&perm[0]|], [C.cinit|&perm[1]|], [C.cinit|&perm[2]|] ] + | length num_blocks == 3 = [[C.cinit|&perm[0]|], [C.cinit|&perm[1]|], [C.cinit|&perm[2]|]] | otherwise = [] - failure_args = take (numFailureParams safety) - [[C.cinit|&ctx->global_failure|], - [C.cinit|&ctx->failure_is_an_option|], - [C.cinit|&ctx->global_failure_args|]] - args'' = perm_args ++ failure_args ++ [ [C.cinit|&$id:a|] | a <- args' ] - sizes_nonzero = expsNotZero [grid_x, grid_y, grid_z, - block_x, block_y, block_z] + failure_args = + take + (numFailureParams safety) + [ [C.cinit|&ctx->global_failure|], + [C.cinit|&ctx->failure_is_an_option|], + [C.cinit|&ctx->global_failure_args|] + ] + args'' = perm_args ++ failure_args ++ [[C.cinit|&$id:a|] | a <- args'] + sizes_nonzero = + expsNotZero + [ grid_x, + grid_y, + grid_z, + block_x, + block_y, + block_z + ] (bef, aft) = profilingEnclosure kernel_name - GC.stm [C.cstm| + GC.stm + [C.cstm| if ($exp:sizes_nonzero) { int perm[3] = { 0, 1, 2 }; @@ -283,12 +344,11 @@ when (safety >= SafetyFull) $ GC.stm [C.cstm|ctx->failure_is_an_option = 1;|] - where - mkDims [] = ([C.cexp|0|] , [C.cexp|0|], [C.cexp|0|]) + mkDims [] = ([C.cexp|0|], [C.cexp|0|], [C.cexp|0|]) mkDims [x] = (x, [C.cexp|1|], [C.cexp|1|]) - mkDims [x,y] = (x, y, [C.cexp|1|]) - mkDims (x:y:z:_) = (x, y, z) + mkDims [x, y] = (x, y, [C.cexp|1|]) + mkDims (x : y : z : _) = (x, y, z) addExp x y = [C.cexp|$exp:x + $exp:y|] alignExp e = [C.cexp|$exp:e + ((8 - ($exp:e % 8)) % 8)|] mkOffsets = scanl (\a b -> a `addExp` alignExp b) [C.cexp|0|] 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 @@ -1,27 +1,33 @@ {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TemplateHaskell #-} + -- | Various boilerplate definitions for the CUDA backend. module Futhark.CodeGen.Backends.CCUDA.Boilerplate - ( - generateBoilerplate - , profilingEnclosure - , module Futhark.CodeGen.Backends.COpenCL.Boilerplate - ) where - -import qualified Language.C.Syntax as C -import qualified Language.C.Quote.OpenCL as C + ( generateBoilerplate, + profilingEnclosure, + module Futhark.CodeGen.Backends.COpenCL.Boilerplate, + ) +where +import Data.FileEmbed (embedStringFile) +import qualified Data.Map as M +import Data.Maybe +import Futhark.CodeGen.Backends.COpenCL.Boilerplate + ( copyDevToDev, + copyDevToHost, + copyHostToDev, + copyScalarFromDev, + copyScalarToDev, + costCentreReport, + failureSwitch, + kernelRuns, + kernelRuntime, + ) import qualified Futhark.CodeGen.Backends.GenericC as GC import Futhark.CodeGen.ImpCode.OpenCL -import Futhark.CodeGen.Backends.COpenCL.Boilerplate - (failureSwitch, kernelRuntime, kernelRuns, costCentreReport, - copyDevToDev, copyDevToHost, copyHostToDev, - copyScalarToDev, copyScalarFromDev) import Futhark.Util (chunk, zEncodeString) - -import qualified Data.Map as M -import Data.Maybe -import Data.FileEmbed (embedStringFile) +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C errorMsgNumArgs :: ErrorMsg a -> Int errorMsgNumArgs = length . errorMsgArgTypes @@ -29,7 +35,7 @@ -- | Block items to put before and after a thing to be profiled. profilingEnclosure :: Name -> ([C.BlockItem], [C.BlockItem]) profilingEnclosure name = - ([C.citems| + ( [C.citems| typename cudaEvent_t *pevents = NULL; if (ctx->profiling && !ctx->profiling_paused) { pevents = cuda_get_events(&ctx->cuda, @@ -38,20 +44,27 @@ CUDA_SUCCEED(cudaEventRecord(pevents[0], 0)); } |], - [C.citems| + [C.citems| if (pevents != NULL) { CUDA_SUCCEED(cudaEventRecord(pevents[1], 0)); } - |]) + |] + ) -- | Called after most code has been generated to generate the bulk of -- the boilerplate. -generateBoilerplate :: String -> String -> [Name] -> M.Map KernelName KernelSafety - -> M.Map Name SizeClass - -> [FailureMsg] - -> GC.CompilerM OpenCL () () +generateBoilerplate :: + String -> + String -> + [Name] -> + M.Map KernelName KernelSafety -> + M.Map Name SizeClass -> + [FailureMsg] -> + GC.CompilerM OpenCL () () generateBoilerplate cuda_program cuda_prelude cost_centres kernels sizes failures = do - mapM_ GC.earlyDecl [C.cunit| + mapM_ + GC.earlyDecl + [C.cunit| $esc:("#include ") $esc:("#include ") $esc:("typedef CUdeviceptr fl_mem_t;") @@ -66,12 +79,12 @@ GC.profileReport [C.citem|CUDA_SUCCEED(cuda_tally_profiling_records(&ctx->cuda));|] mapM_ GC.profileReport $ costCentreReport $ cost_centres ++ M.keys kernels - where cuda_h = $(embedStringFile "rts/c/cuda.h") free_list_h = $(embedStringFile "rts/c/free_list.h") - fragments = map (\s -> [C.cinit|$string:s|]) - $ chunk 2000 (cuda_prelude ++ cuda_program) + fragments = + map (\s -> [C.cinit|$string:s|]) $ + chunk 2000 (cuda_prelude ++ cuda_program) generateSizeFuns :: M.Map Name SizeClass -> GC.CompilerM OpenCL () () generateSizeFuns sizes = do @@ -85,22 +98,25 @@ 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) { + ( [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) { + ( [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) { + ( [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 @@ -108,20 +124,22 @@ num_sizes = M.size sizes GC.earlyDecl [C.cedecl|struct sizes { $sdecls:size_decls };|] cfg <- GC.publicDef "context_config" GC.InitDecl $ \s -> - ([C.cedecl|struct $id:s;|], - [C.cedecl|struct $id:s { struct cuda_config cu_cfg; + ( [C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { struct cuda_config cu_cfg; int profiling; size_t sizes[$int:num_sizes]; int num_nvrtc_opts; const char **nvrtc_opts; - };|]) + };|] + ) - let size_value_inits = zipWith sizeInit [0..M.size sizes-1] (M.elems sizes) + let size_value_inits = zipWith sizeInit [0 .. M.size sizes -1] (M.elems sizes) sizeInit i size = [C.cstm|cfg->sizes[$int:i] = $int:val;|] - where val = fromMaybe 0 $ sizeDefault size + where + val = fromMaybe 0 $ sizeDefault size GC.publicDef_ "context_config_new" GC.InitDecl $ \s -> - ([C.cedecl|struct $id:cfg* $id:s(void);|], - [C.cedecl|struct $id:cfg* $id:s(void) { + ( [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; @@ -136,102 +154,117 @@ size_names, size_vars, cfg->sizes, size_classes); 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) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg) { free(cfg->nvrtc_opts); free(cfg); - }|]) + }|] + ) GC.publicDef_ "context_config_add_nvrtc_option" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) { cfg->nvrtc_opts[cfg->num_nvrtc_opts] = opt; cfg->num_nvrtc_opts++; cfg->nvrtc_opts = (const char**) realloc(cfg->nvrtc_opts, (cfg->num_nvrtc_opts+1) * sizeof(const char*)); cfg->nvrtc_opts[cfg->num_nvrtc_opts] = NULL; - }|]) + }|] + ) 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 flag) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { cfg->cu_cfg.logging = cfg->cu_cfg.debugging = flag; - }|]) + }|] + ) GC.publicDef_ "context_config_set_profiling" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { cfg->profiling = flag; - }|]) + }|] + ) 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 flag) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { cfg->cu_cfg.logging = flag; - }|]) + }|] + ) GC.publicDef_ "context_config_set_device" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s) { set_preferred_device(&cfg->cu_cfg, s); - }|]) + }|] + ) GC.publicDef_ "context_config_dump_program_to" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { cfg->cu_cfg.dump_program_to = path; - }|]) + }|] + ) GC.publicDef_ "context_config_load_program_from" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { cfg->cu_cfg.load_program_from = path; - }|]) + }|] + ) GC.publicDef_ "context_config_dump_ptx_to" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { cfg->cu_cfg.dump_ptx_to = path; - }|]) + }|] + ) GC.publicDef_ "context_config_load_ptx_from" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { cfg->cu_cfg.load_ptx_from = path; - }|]) + }|] + ) GC.publicDef_ "context_config_set_default_group_size" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, int size);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int size);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { cfg->cu_cfg.default_block_size = size; cfg->cu_cfg.default_block_size_changed = 1; - }|]) + }|] + ) GC.publicDef_ "context_config_set_default_num_groups" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, int num) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int num) { cfg->cu_cfg.default_grid_size = num; cfg->cu_cfg.default_grid_size_changed = 1; - }|]) + }|] + ) GC.publicDef_ "context_config_set_default_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) { + ( [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_tile_size = size; cfg->cu_cfg.default_tile_size_changed = 1; - }|]) + }|] + ) 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) { + ( [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_threshold = size; - }|]) + }|] + ) 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) { + ( [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) { for (int i = 0; i < $int:num_sizes; i++) { if (strcmp(size_name, size_names[i]) == 0) { @@ -260,38 +293,46 @@ return 0; } return 1; - }|]) + }|] + ) return cfg -generateContextFuns :: String -> [Name] -> M.Map KernelName KernelSafety - -> M.Map Name SizeClass - -> [FailureMsg] - -> GC.CompilerM OpenCL () () +generateContextFuns :: + String -> + [Name] -> + M.Map KernelName KernelSafety -> + M.Map Name SizeClass -> + [FailureMsg] -> + GC.CompilerM OpenCL () () generateContextFuns cfg cost_centres kernels sizes failures = do final_inits <- GC.contextFinalInits (fields, init_fields) <- GC.contextContents let forCostCentre name = - [([C.csdecl|typename int64_t $id:(kernelRuntime name);|], - [C.cstm|ctx->$id:(kernelRuntime name) = 0;|]), - ([C.csdecl|int $id:(kernelRuns name);|], - [C.cstm|ctx->$id:(kernelRuns name) = 0;|])] + [ ( [C.csdecl|typename int64_t $id:(kernelRuntime name);|], + [C.cstm|ctx->$id:(kernelRuntime name) = 0;|] + ), + ( [C.csdecl|int $id:(kernelRuns name);|], + [C.cstm|ctx->$id:(kernelRuns name) = 0;|] + ) + ] forKernel name = - ([C.csdecl|typename CUfunction $id:name;|], - [C.cstm|CUDA_SUCCEED(cuModuleGetFunction( + ( [C.csdecl|typename CUfunction $id:name;|], + [C.cstm|CUDA_SUCCEED(cuModuleGetFunction( &ctx->$id:name, ctx->cuda.module, - $string:(pretty (C.toIdent name mempty))));|]) - : forCostCentre name + $string:(pretty (C.toIdent name mempty))));|] + ) : + forCostCentre name (kernel_fields, init_kernel_fields) = unzip $ - concatMap forKernel (M.keys kernels) ++ - concatMap forCostCentre cost_centres + concatMap forKernel (M.keys kernels) + ++ concatMap forCostCentre cost_centres ctx <- GC.publicDef "context" GC.InitDecl $ \s -> - ([C.cedecl|struct $id:s;|], - [C.cedecl|struct $id:s { + ( [C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { int detail_memory; int debugging; int profiling; @@ -309,16 +350,20 @@ int total_runs; long int total_runtime; - };|]) + };|] + ) - let set_sizes = zipWith (\i k -> [C.cstm|ctx->sizes.$id:k = cfg->sizes[$int:i];|]) - [(0::Int)..] $ M.keys sizes + let set_sizes = + zipWith + (\i k -> [C.cstm|ctx->sizes.$id:k = cfg->sizes[$int:i];|]) + [(0 :: Int) ..] + $ M.keys sizes max_failure_args = foldl max 0 $ map (errorMsgNumArgs . failureError) failures 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) { + ( [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; @@ -347,7 +392,7 @@ CUDA_SUCCEED(cuMemAlloc(&ctx->global_failure, sizeof(no_error))); CUDA_SUCCEED(cuMemcpyHtoD(ctx->global_failure, &no_error, sizeof(no_error))); // The +1 is to avoid zero-byte allocations. - CUDA_SUCCEED(cuMemAlloc(&ctx->global_failure_args, sizeof(int32_t)*($int:max_failure_args+1))); + CUDA_SUCCEED(cuMemAlloc(&ctx->global_failure_args, sizeof(int64_t)*($int:max_failure_args+1))); $stms:init_kernel_fields @@ -361,20 +406,23 @@ futhark_context_sync(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) { + ( [C.cedecl|void $id:s(struct $id:ctx* ctx);|], + [C.cedecl|void $id:s(struct $id:ctx* ctx) { free_constants(ctx); cuda_cleanup(&ctx->cuda); 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) { + ( [C.cedecl|int $id:s(struct $id:ctx* ctx);|], + [C.cedecl|int $id:s(struct $id:ctx* ctx) { + CUDA_SUCCEED(cuCtxPushCurrent(ctx->cuda.cu_ctx)); CUDA_SUCCEED(cuCtxSynchronize()); if (ctx->failure_is_an_option) { // Check for any delayed error. @@ -394,7 +442,7 @@ &no_failure, sizeof(int32_t))); - typename int32_t args[$int:max_failure_args+1]; + typename int64_t args[$int:max_failure_args+1]; CUDA_SUCCEED( cuMemcpyDtoH(&args, ctx->global_failure_args, @@ -405,13 +453,15 @@ return 1; } } + CUDA_SUCCEED(cuCtxPopCurrent(&ctx->cuda.cu_ctx)); return 0; - }|]) - + }|] + ) GC.publicDef_ "context_clear_caches" GC.MiscDecl $ \s -> - ([C.cedecl|int $id:s(struct $id:ctx* ctx);|], - [C.cedecl|int $id:s(struct $id:ctx* ctx) { + ( [C.cedecl|int $id:s(struct $id:ctx* ctx);|], + [C.cedecl|int $id:s(struct $id:ctx* ctx) { CUDA_SUCCEED(cuda_free_all(&ctx->cuda)); return 0; - }|]) + }|] + ) diff --git a/src/Futhark/CodeGen/Backends/COpenCL.hs b/src/Futhark/CodeGen/Backends/COpenCL.hs --- a/src/Futhark/CodeGen/Backends/COpenCL.hs +++ b/src/Futhark/CodeGen/Backends/COpenCL.hs @@ -1,105 +1,162 @@ -{-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TupleSections #-} + -- | Code generation for C with OpenCL. module Futhark.CodeGen.Backends.COpenCL - ( compileProg - , GC.CParts(..) - , GC.asLibrary - , GC.asExecutable - ) where + ( compileProg, + GC.CParts (..), + GC.asLibrary, + GC.asExecutable, + ) +where import Control.Monad hiding (mapM) import Data.List (intercalate) - -import qualified Language.C.Syntax as C -import qualified Language.C.Quote.OpenCL as C - -import Futhark.IR.KernelsMem - hiding (GetSize, CmpSizeLe, GetSizeMax) import Futhark.CodeGen.Backends.COpenCL.Boilerplate import qualified Futhark.CodeGen.Backends.GenericC as GC import Futhark.CodeGen.Backends.GenericC.Options import Futhark.CodeGen.ImpCode.OpenCL import qualified Futhark.CodeGen.ImpGen.OpenCL as ImpGen +import Futhark.IR.KernelsMem hiding + ( CmpSizeLe, + GetSize, + GetSizeMax, + ) import Futhark.MonadFreshNames +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C -- | Compile the program to C with calls to OpenCL. compileProg :: MonadFreshNames m => Prog KernelsMem -> m (ImpGen.Warnings, GC.CParts) compileProg prog = do - (ws, Program opencl_code opencl_prelude kernels - types sizes failures prog') <- ImpGen.compileProg prog + ( ws, + Program + opencl_code + opencl_prelude + kernels + types + sizes + failures + prog' + ) <- + ImpGen.compileProg prog let cost_centres = - [copyDevToDev, copyDevToHost, copyHostToDev, - copyScalarToDev, copyScalarFromDev] - (ws,) <$> - GC.compileProg "opencl" operations - (generateBoilerplate opencl_code opencl_prelude - cost_centres kernels types sizes failures) - include_opencl_h [Space "device", DefaultSpace] - cliOptions prog' - where operations :: GC.Operations OpenCL () - operations = GC.defaultOperations - { GC.opsCompiler = callKernel - , GC.opsWriteScalar = writeOpenCLScalar - , GC.opsReadScalar = readOpenCLScalar - , GC.opsAllocate = allocateOpenCLBuffer - , GC.opsDeallocate = deallocateOpenCLBuffer - , GC.opsCopy = copyOpenCLMemory - , GC.opsStaticArray = staticOpenCLArray - , GC.opsMemoryType = openclMemoryType - , GC.opsFatMemory = True - } - include_opencl_h = unlines ["#define CL_TARGET_OPENCL_VERSION 120", - "#define CL_USE_DEPRECATED_OPENCL_1_2_APIS", - "#ifdef __APPLE__", - "#define CL_SILENCE_DEPRECATION", - "#include ", - "#else", - "#include ", - "#endif"] + [ copyDevToDev, + copyDevToHost, + copyHostToDev, + copyScalarToDev, + copyScalarFromDev + ] + (ws,) + <$> GC.compileProg + "opencl" + operations + ( generateBoilerplate + opencl_code + opencl_prelude + cost_centres + kernels + types + sizes + failures + ) + include_opencl_h + [Space "device", DefaultSpace] + cliOptions + prog' + where + operations :: GC.Operations OpenCL () + operations = + GC.defaultOperations + { GC.opsCompiler = callKernel, + GC.opsWriteScalar = writeOpenCLScalar, + GC.opsReadScalar = readOpenCLScalar, + GC.opsAllocate = allocateOpenCLBuffer, + GC.opsDeallocate = deallocateOpenCLBuffer, + GC.opsCopy = copyOpenCLMemory, + GC.opsStaticArray = staticOpenCLArray, + GC.opsMemoryType = openclMemoryType, + GC.opsFatMemory = True + } + include_opencl_h = + unlines + [ "#define CL_TARGET_OPENCL_VERSION 120", + "#define CL_USE_DEPRECATED_OPENCL_1_2_APIS", + "#ifdef __APPLE__", + "#define CL_SILENCE_DEPRECATION", + "#include ", + "#else", + "#include ", + "#endif" + ] cliOptions :: [Option] cliOptions = - commonOptions ++ - [ Option { optionLongName = "platform" - , optionShortName = Just 'p' - , optionArgument = RequiredArgument "NAME" - , optionAction = [C.cstm|futhark_context_config_set_platform(cfg, optarg);|] - } - , Option { optionLongName = "dump-opencl" - , optionShortName = Nothing - , optionArgument = RequiredArgument "FILE" - , optionAction = [C.cstm|{futhark_context_config_dump_program_to(cfg, optarg); + commonOptions + ++ [ Option + { optionLongName = "platform", + optionShortName = Just 'p', + optionArgument = RequiredArgument "NAME", + optionDescription = "Use the first OpenCL platform whose name contains the given string.", + optionAction = [C.cstm|futhark_context_config_set_platform(cfg, optarg);|] + }, + Option + { optionLongName = "dump-opencl", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Dump the embedded OpenCL program to the indicated file.", + optionAction = + [C.cstm|{futhark_context_config_dump_program_to(cfg, optarg); entry_point = NULL;}|] - } - , Option { optionLongName = "load-opencl" - , optionShortName = Nothing - , optionArgument = RequiredArgument "FILE" - , optionAction = [C.cstm|futhark_context_config_load_program_from(cfg, optarg);|] - } - , Option { optionLongName = "dump-opencl-binary" - , optionShortName = Nothing - , optionArgument = RequiredArgument "FILE" - , optionAction = [C.cstm|{futhark_context_config_dump_binary_to(cfg, optarg); + }, + Option + { optionLongName = "load-opencl", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Instead of using the embedded OpenCL program, load it from the indicated file.", + optionAction = [C.cstm|futhark_context_config_load_program_from(cfg, optarg);|] + }, + Option + { optionLongName = "dump-opencl-binary", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Dump the compiled version of the embedded OpenCL program to the indicated file.", + optionAction = + [C.cstm|{futhark_context_config_dump_binary_to(cfg, optarg); entry_point = NULL;}|] - } - , Option { optionLongName = "load-opencl-binary" - , optionShortName = Nothing - , optionArgument = RequiredArgument "FILE" - , optionAction = [C.cstm|futhark_context_config_load_binary_from(cfg, optarg);|] - } - , Option { optionLongName = "build-option" - , optionShortName = Nothing - , optionArgument = RequiredArgument "OPT" - , optionAction = [C.cstm|futhark_context_config_add_build_option(cfg, optarg);|] - } - , Option { optionLongName = "profile" - , optionShortName = Just 'P' - , optionArgument = NoArgument - , optionAction = [C.cstm|futhark_context_config_set_profiling(cfg, 1);|] + }, + Option + { optionLongName = "load-opencl-binary", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Load an OpenCL binary from the indicated file.", + optionAction = [C.cstm|futhark_context_config_load_binary_from(cfg, optarg);|] + }, + Option + { optionLongName = "build-option", + optionShortName = Nothing, + optionArgument = RequiredArgument "OPT", + optionDescription = "Add an additional build option to the string passed to clBuildProgram().", + optionAction = [C.cstm|futhark_context_config_add_build_option(cfg, optarg);|] + }, + Option + { optionLongName = "profile", + optionShortName = Just 'P', + optionArgument = NoArgument, + optionDescription = "Gather profiling data while executing and print out a summary at the end.", + optionAction = [C.cstm|futhark_context_config_set_profiling(cfg, 1);|] + }, + Option + { optionLongName = "list-devices", + optionShortName = Nothing, + optionArgument = NoArgument, + optionDescription = "List all OpenCL devices and platforms available on the system.", + optionAction = + [C.cstm|{futhark_context_config_list_devices(cfg); + entry_point = NULL;}|] } - ] + ] -- We detect the special case of writing a constant and turn it into a -- non-blocking write. This may be slightly faster, as it prevents @@ -112,9 +169,10 @@ val' <- newVName "write_tmp" let (decl, blocking) = case val of - C.Const{} -> ([C.citem|static $ty:t $id:val' = $exp:val;|], [C.cexp|CL_FALSE|]) - _ -> ([C.citem|$ty:t $id:val' = $exp:val;|], [C.cexp|CL_TRUE|]) - GC.stm [C.cstm|{$item:decl + C.Const {} -> ([C.citem|static $ty:t $id:val' = $exp:val;|], [C.cexp|CL_FALSE|]) + _ -> ([C.citem|$ty:t $id:val' = $exp:val;|], [C.cexp|CL_TRUE|]) + GC.stm + [C.cstm|{$item:decl OPENCL_SUCCEED_OR_RETURN( clEnqueueWriteBuffer(ctx->opencl.queue, $exp:mem, $exp:blocking, $exp:i * sizeof($ty:t), sizeof($ty:t), @@ -132,14 +190,16 @@ readOpenCLScalar mem i t "device" _ = do val <- newVName "read_res" GC.decl [C.cdecl|$ty:t $id:val;|] - GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN( + GC.stm + [C.cstm|OPENCL_SUCCEED_OR_RETURN( clEnqueueReadBuffer(ctx->opencl.queue, $exp:mem, ctx->failure_is_an_option ? CL_FALSE : CL_TRUE, $exp:i * sizeof($ty:t), sizeof($ty:t), &$id:val, 0, NULL, $exp:(profilingEvent copyScalarFromDev))); |] - GC.stm [C.cstm|if (ctx->failure_is_an_option && + GC.stm + [C.cstm|if (ctx->failure_is_an_option && futhark_context_sync(ctx) != 0) { return 1; }|] return [C.cexp|$id:val|] readOpenCLScalar _ _ _ space _ = @@ -162,7 +222,8 @@ -- out of bounds, even if asked to read zero bytes. We protect with a -- branch to avoid this. copyOpenCLMemory destmem destidx DefaultSpace srcmem srcidx (Space "device") nbytes = - GC.stm [C.cstm| + GC.stm + [C.cstm| if ($exp:nbytes > 0) { OPENCL_SUCCEED_OR_RETURN( clEnqueueReadBuffer(ctx->opencl.queue, $exp:srcmem, @@ -175,7 +236,8 @@ } |] copyOpenCLMemory destmem destidx (Space "device") srcmem srcidx DefaultSpace nbytes = - GC.stm [C.cstm| + GC.stm + [C.cstm| if ($exp:nbytes > 0) { OPENCL_SUCCEED_OR_RETURN( clEnqueueWriteBuffer(ctx->opencl.queue, $exp:destmem, CL_TRUE, @@ -187,7 +249,8 @@ copyOpenCLMemory destmem destidx (Space "device") srcmem srcidx (Space "device") nbytes = -- Be aware that OpenCL swaps the usual order of operands for -- memcpy()-like functions. The order below is not a typo. - GC.stm [C.cstm|{ + GC.stm + [C.cstm|{ if ($exp:nbytes > 0) { OPENCL_SUCCEED_OR_RETURN( clEnqueueCopyBuffer(ctx->opencl.queue, @@ -225,7 +288,8 @@ -- Fake a memory block. GC.contextField (C.toIdent name mempty) [C.cty|struct memblock_device|] Nothing -- During startup, copy the data to where we need it. - GC.atInit [C.cstm|{ + GC.atInit + [C.cstm|{ typename cl_int success; ctx->$id:name.references = NULL; ctx->$id:name.size = 0; @@ -243,7 +307,6 @@ } }|] GC.item [C.citem|struct memblock_device $id:name = ctx->$id:name;|] - staticOpenCLArray _ space _ _ = error $ "OpenCL backend cannot create static array in memory space '" ++ space ++ "'" @@ -253,25 +316,25 @@ callKernel (CmpSizeLe v key x) = do x' <- GC.compileExp x GC.stm [C.cstm|$id:v = ctx->sizes.$id:key <= $exp:x';|] - GC.stm [C.cstm|if (ctx->logging) { + GC.stm + [C.cstm|if (ctx->logging) { fprintf(stderr, "Compared %s <= %d.\n", $string:(pretty key), $exp:x'); }|] callKernel (GetSizeMax v size_class) = let field = "max_" ++ pretty size_class - in GC.stm [C.cstm|$id:v = ctx->opencl.$id:field;|] - + in GC.stm [C.cstm|$id:v = ctx->opencl.$id:field;|] callKernel (LaunchKernel safety name args num_workgroups workgroup_size) = do - -- The other failure args are set automatically when the kernel is -- first created. when (safety == SafetyFull) $ - GC.stm [C.cstm| + GC.stm + [C.cstm| OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->$id:name, 1, sizeof(ctx->failure_is_an_option), &ctx->failure_is_an_option)); |] - zipWithM_ setKernelArg [numFailureParams safety..] args + zipWithM_ setKernelArg [numFailureParams safety ..] args num_workgroups' <- mapM GC.compileExp num_workgroups workgroup_size' <- mapM GC.compileExp workgroup_size local_bytes <- foldM localBytes [C.cexp|0|] args @@ -280,32 +343,38 @@ when (safety >= SafetyFull) $ GC.stm [C.cstm|ctx->failure_is_an_option = 1;|] - - where setKernelArg i (ValueKArg e bt) = do - v <- GC.compileExpToName "kernel_arg" bt e - GC.stm [C.cstm| + where + setKernelArg i (ValueKArg e bt) = do + v <- GC.compileExpToName "kernel_arg" bt e + GC.stm + [C.cstm| OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->$id:name, $int:i, sizeof($id:v), &$id:v)); |] - - setKernelArg i (MemKArg v) = do - v' <- GC.rawMem v - GC.stm [C.cstm| + setKernelArg i (MemKArg v) = do + v' <- GC.rawMem v + GC.stm + [C.cstm| OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->$id:name, $int:i, sizeof($exp:v'), &$exp:v')); |] - - setKernelArg i (SharedMemoryKArg num_bytes) = do - num_bytes' <- GC.compileExp $ unCount num_bytes - GC.stm [C.cstm| + setKernelArg i (SharedMemoryKArg num_bytes) = do + num_bytes' <- GC.compileExp $ unCount num_bytes + GC.stm + [C.cstm| OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->$id:name, $int:i, $exp:num_bytes', NULL)); |] - localBytes cur (SharedMemoryKArg num_bytes) = do - num_bytes' <- GC.compileExp $ unCount num_bytes - return [C.cexp|$exp:cur + $exp:num_bytes'|] - localBytes cur _ = return cur + localBytes cur (SharedMemoryKArg num_bytes) = do + num_bytes' <- GC.compileExp $ unCount num_bytes + return [C.cexp|$exp:cur + $exp:num_bytes'|] + localBytes cur _ = return cur -launchKernel :: C.ToExp a => - KernelName -> [a] -> [a] -> a -> GC.CompilerM op s () +launchKernel :: + C.ToExp a => + KernelName -> + [a] -> + [a] -> + a -> + GC.CompilerM op s () launchKernel kernel_name num_workgroups workgroup_dims local_bytes = do global_work_size <- newVName "global_work_size" time_start <- newVName "time_start" @@ -313,7 +382,8 @@ time_diff <- newVName "time_diff" local_work_size <- newVName "local_work_size" - GC.stm [C.cstm| + GC.stm + [C.cstm| if ($exp:total_elements != 0) { const size_t $id:global_work_size[$int:kernel_rank] = {$inits:kernel_dims'}; const size_t $id:local_work_size[$int:kernel_rank] = {$inits:workgroup_dims'}; @@ -338,18 +408,20 @@ $string:(pretty kernel_name), $id:time_diff); } }|] - where kernel_rank = length kernel_dims - kernel_dims = zipWith multExp num_workgroups workgroup_dims - kernel_dims' = map toInit kernel_dims - workgroup_dims' = map toInit workgroup_dims - total_elements = foldl multExp [C.cexp|1|] kernel_dims + where + kernel_rank = length kernel_dims + kernel_dims = zipWith multExp (map toSize num_workgroups) (map toSize workgroup_dims) + kernel_dims' = map toInit kernel_dims + workgroup_dims' = map toInit workgroup_dims + total_elements = foldl multExp [C.cexp|1|] kernel_dims - toInit e = [C.cinit|$exp:e|] - multExp x y = [C.cexp|$exp:x * $exp:y|] + toInit e = [C.cinit|$exp:e|] + multExp x y = [C.cexp|$exp:x * $exp:y|] + toSize e = [C.cexp|(size_t)$exp:e|] - printKernelSize :: VName -> [C.Stm] - printKernelSize work_size = - intercalate [[C.cstm|fprintf(stderr, ", ");|]] $ - map (printKernelDim work_size) [0..kernel_rank-1] - printKernelDim global_work_size i = - [[C.cstm|fprintf(stderr, "%zu", $id:global_work_size[$int:i]);|]] + printKernelSize :: VName -> [C.Stm] + printKernelSize work_size = + intercalate [[C.cstm|fprintf(stderr, ", ");|]] $ + map (printKernelDim work_size) [0 .. kernel_rank -1] + printKernelDim global_work_size i = + [[C.cstm|fprintf(stderr, "%zu", $id:global_work_size[$int:i]);|]] 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 @@ -1,49 +1,55 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE FlexibleContexts #-} -module Futhark.CodeGen.Backends.COpenCL.Boilerplate - ( generateBoilerplate - , profilingEvent - , copyDevToDev, copyDevToHost, copyHostToDev, copyScalarToDev, copyScalarFromDev - , commonOptions - , failureSwitch - , costCentreReport - , kernelRuntime - , kernelRuns - ) where +module Futhark.CodeGen.Backends.COpenCL.Boilerplate + ( generateBoilerplate, + profilingEvent, + copyDevToDev, + copyDevToHost, + copyHostToDev, + copyScalarToDev, + copyScalarFromDev, + commonOptions, + failureSwitch, + costCentreReport, + kernelRuntime, + kernelRuns, + ) +where import Control.Monad.State import Data.FileEmbed import qualified Data.Map as M import Data.Maybe -import qualified Language.C.Syntax as C -import qualified Language.C.Quote.OpenCL as C - -import Futhark.CodeGen.ImpCode.OpenCL import qualified Futhark.CodeGen.Backends.GenericC as GC import Futhark.CodeGen.Backends.GenericC.Options +import Futhark.CodeGen.ImpCode.OpenCL import Futhark.CodeGen.OpenCL.Heuristics import Futhark.Util (chunk, zEncodeString) +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C errorMsgNumArgs :: ErrorMsg a -> Int errorMsgNumArgs = length . errorMsgArgTypes failureSwitch :: [FailureMsg] -> C.Stm failureSwitch failures = - let printfEscape = let escapeChar '%' = "%%" - escapeChar c = [c] - in concatMap escapeChar + let printfEscape = + let escapeChar '%' = "%%" + escapeChar c = [c] + in concatMap escapeChar onPart (ErrorString s) = printfEscape s - onPart ErrorInt32{} = "%d" + onPart ErrorInt32 {} = "%lld" + onPart ErrorInt64 {} = "%lld" onFailure i (FailureMsg emsg@(ErrorMsg parts) backtrace) = - let msg = concatMap onPart parts ++ "\n" ++ printfEscape backtrace - msgargs = [ [C.cexp|args[$int:j]|] | j <- [0..errorMsgNumArgs emsg-1] ] + let msg = concatMap onPart parts ++ "\n" ++ printfEscape backtrace + msgargs = [[C.cexp|args[$int:j]|] | j <- [0 .. errorMsgNumArgs emsg -1]] in [C.cstm|case $int:i: {ctx->error = msgprintf($string:msg, $args:msgargs); break;}|] failure_cases = - zipWith onFailure [(0::Int)..] failures - in [C.cstm|switch (failure_idx) { $stms:failure_cases }|] + zipWith onFailure [(0 :: Int) ..] failures + in [C.cstm|switch (failure_idx) { $stms:failure_cases }|] copyDevToDev, copyDevToHost, copyHostToDev, copyScalarToDev, copyScalarFromDev :: Name copyDevToDev = "copy_dev_to_dev" @@ -61,12 +67,15 @@ -- | Called after most code has been generated to generate the bulk of -- the boilerplate. -generateBoilerplate :: String -> String -> [Name] - -> M.Map KernelName KernelSafety - -> [PrimType] - -> M.Map Name SizeClass - -> [FailureMsg] - -> GC.CompilerM OpenCL () () +generateBoilerplate :: + String -> + String -> + [Name] -> + M.Map KernelName KernelSafety -> + [PrimType] -> + M.Map Name SizeClass -> + [FailureMsg] -> + GC.CompilerM OpenCL () () generateBoilerplate opencl_code opencl_prelude cost_centres kernels types sizes failures = do final_inits <- GC.contextFinalInits @@ -85,39 +94,44 @@ 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) { + ( [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) { + ( [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) { + ( [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|size_t $id:k;|]) $ M.keys sizes GC.earlyDecl [C.cedecl|struct sizes { $sdecls:size_decls };|] cfg <- GC.publicDef "context_config" GC.InitDecl $ \s -> - ([C.cedecl|struct $id:s;|], - [C.cedecl|struct $id:s { struct opencl_config opencl; + ( [C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { struct opencl_config opencl; size_t sizes[$int:num_sizes]; int num_build_opts; const char **build_opts; - };|]) + };|] + ) - let size_value_inits = zipWith sizeInit [0..M.size sizes-1] (M.elems sizes) + let size_value_inits = zipWith sizeInit [0 .. M.size sizes -1] (M.elems sizes) sizeInit i size = [C.cstm|cfg->sizes[$int:i] = $int:val;|] - where val = fromMaybe 0 $ sizeDefault size + where + val = fromMaybe 0 $ sizeDefault size GC.publicDef_ "context_config_new" GC.InitDecl $ \s -> - ([C.cedecl|struct $id:cfg* $id:s(void);|], - [C.cedecl|struct $id:cfg* $id:s(void) { + ( [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; @@ -131,114 +145,137 @@ size_names, size_vars, cfg->sizes, size_classes); 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) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg) { free(cfg->build_opts); free(cfg); - }|]) + }|] + ) GC.publicDef_ "context_config_add_build_option" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) { cfg->build_opts[cfg->num_build_opts] = opt; cfg->num_build_opts++; cfg->build_opts = (const char**) realloc(cfg->build_opts, (cfg->num_build_opts+1) * sizeof(const char*)); cfg->build_opts[cfg->num_build_opts] = NULL; - }|]) + }|] + ) 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 flag) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { cfg->opencl.profiling = cfg->opencl.logging = cfg->opencl.debugging = flag; - }|]) + }|] + ) GC.publicDef_ "context_config_set_profiling" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { cfg->opencl.profiling = flag; - }|]) + }|] + ) 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 flag) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { cfg->opencl.logging = flag; - }|]) + }|] + ) GC.publicDef_ "context_config_set_device" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s) { set_preferred_device(&cfg->opencl, s); - }|]) + }|] + ) GC.publicDef_ "context_config_set_platform" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s) { set_preferred_platform(&cfg->opencl, s); - }|]) + }|] + ) GC.publicDef_ "context_config_select_device_interactively" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg) { select_device_interactively(&cfg->opencl); - }|]) + }|] + ) + GC.publicDef_ "context_config_list_devices" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:cfg* cfg);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg) { + list_devices(&cfg->opencl); + }|] + ) + GC.publicDef_ "context_config_dump_program_to" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { cfg->opencl.dump_program_to = path; - }|]) + }|] + ) GC.publicDef_ "context_config_load_program_from" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { cfg->opencl.load_program_from = path; - }|]) - + }|] + ) GC.publicDef_ "context_config_dump_binary_to" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { cfg->opencl.dump_binary_to = path; - }|]) + }|] + ) GC.publicDef_ "context_config_load_binary_from" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { cfg->opencl.load_binary_from = path; - }|]) + }|] + ) GC.publicDef_ "context_config_set_default_group_size" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, int size);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int size);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { cfg->opencl.default_group_size = size; cfg->opencl.default_group_size_changed = 1; - }|]) + }|] + ) GC.publicDef_ "context_config_set_default_num_groups" GC.InitDecl $ \s -> - ([C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, int num) { + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int num) { cfg->opencl.default_num_groups = num; - }|]) + }|] + ) GC.publicDef_ "context_config_set_default_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) { + ( [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_tile_size = size; cfg->opencl.default_tile_size_changed = 1; - }|]) + }|] + ) 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) { + ( [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_threshold = size; - }|]) + }|] + ) 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) { + ( [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) { for (int i = 0; i < $int:num_sizes; i++) { if (strcmp(size_name, size_names[i]) == 0) { @@ -268,12 +305,13 @@ } return 1; - }|]) + }|] + ) (fields, init_fields) <- GC.contextContents ctx <- GC.publicDef "context" GC.InitDecl $ \s -> - ([C.cedecl|struct $id:s;|], - [C.cedecl|struct $id:s { + ( [C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { int detail_memory; int debugging; int profiling; @@ -289,11 +327,13 @@ struct sizes sizes; // True if a potentially failing kernel has been enqueued. typename cl_int failure_is_an_option; - };|]) + };|] + ) mapM_ GC.earlyDecl later_top_decls - GC.earlyDecl [C.cedecl|static void init_context_early(struct $id:cfg *cfg, struct $id:ctx* ctx) { + GC.earlyDecl + [C.cedecl|static void init_context_early(struct $id:cfg *cfg, struct $id:ctx* ctx) { ctx->opencl.cfg = cfg->opencl; ctx->detail_memory = cfg->opencl.debugging; ctx->debugging = cfg->opencl.debugging; @@ -313,12 +353,16 @@ $stms:ctx_opencl_inits }|] - let set_sizes = zipWith (\i k -> [C.cstm|ctx->sizes.$id:k = cfg->sizes[$int:i];|]) - [(0::Int)..] $ M.keys sizes + let set_sizes = + zipWith + (\i k -> [C.cstm|ctx->sizes.$id:k = cfg->sizes[$int:i];|]) + [(0 :: Int) ..] + $ M.keys sizes max_failure_args = foldl max 0 $ map (errorMsgNumArgs . failureError) failures - GC.earlyDecl [C.cedecl|static int init_context_late(struct $id:cfg *cfg, struct $id:ctx* ctx, typename cl_program prog) { + GC.earlyDecl + [C.cedecl|static int init_context_late(struct $id:cfg *cfg, struct $id:ctx* ctx, typename cl_program prog) { typename cl_int error; typename cl_int no_error = -1; @@ -332,7 +376,7 @@ ctx->global_failure_args = clCreateBuffer(ctx->opencl.ctx, CL_MEM_READ_WRITE, - sizeof(cl_int)*($int:max_failure_args+1), NULL, &error); + sizeof(int64_t)*($int:max_failure_args+1), NULL, &error); OPENCL_SUCCEED_OR_RETURN(error); // Load all the kernels. @@ -351,12 +395,14 @@ return futhark_context_sync(ctx); }|] - let set_required_types = [ [C.cstm|required_types |= OPENCL_F64; |] - | FloatType Float64 `elem` types ] + let set_required_types = + [ [C.cstm|required_types |= OPENCL_F64; |] + | FloatType Float64 `elem` types + ] 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) { + ( [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; @@ -369,11 +415,12 @@ typename cl_program prog = setup_opencl(&ctx->opencl, opencl_program, required_types, cfg->build_opts); init_context_late(cfg, ctx, prog); return ctx; - }|]) + }|] + ) GC.publicDef_ "context_new_with_command_queue" GC.InitDecl $ \s -> - ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue);|], - [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue) { + ( [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue);|], + [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue) { struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx)); if (ctx == NULL) { return NULL; @@ -386,21 +433,23 @@ typename cl_program prog = setup_opencl_with_command_queue(&ctx->opencl, queue, opencl_program, required_types, cfg->build_opts); init_context_late(cfg, ctx, prog); 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) { + ( [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); $stms:(map releaseKernel (M.toList kernels)) teardown_opencl(&ctx->opencl); 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) { + ( [C.cedecl|int $id:s(struct $id:ctx* ctx);|], + [C.cedecl|int $id:s(struct $id:ctx* ctx) { // Check for any delayed error. typename cl_int failure_idx = -1; if (ctx->failure_is_an_option) { @@ -424,7 +473,7 @@ 0, sizeof(cl_int), &no_failure, 0, NULL, NULL)); - typename cl_int args[$int:max_failure_args+1]; + typename int64_t args[$int:max_failure_args+1]; OPENCL_SUCCEED_OR_RETURN( clEnqueueReadBuffer(ctx->opencl.queue, ctx->global_failure_args, @@ -437,73 +486,90 @@ return 1; } return 0; - }|]) + }|] + ) GC.publicDef_ "context_clear_caches" GC.MiscDecl $ \s -> - ([C.cedecl|int $id:s(struct $id:ctx* ctx);|], - [C.cedecl|int $id:s(struct $id:ctx* ctx) { + ( [C.cedecl|int $id:s(struct $id:ctx* ctx);|], + [C.cedecl|int $id:s(struct $id:ctx* ctx) { ctx->error = OPENCL_SUCCEED_NONFATAL(opencl_free_all(&ctx->opencl)); return ctx->error != NULL; - }|]) + }|] + ) GC.publicDef_ "context_get_command_queue" GC.InitDecl $ \s -> - ([C.cedecl|typename cl_command_queue $id:s(struct $id:ctx* ctx);|], - [C.cedecl|typename cl_command_queue $id:s(struct $id:ctx* ctx) { + ( [C.cedecl|typename cl_command_queue $id:s(struct $id:ctx* ctx);|], + [C.cedecl|typename cl_command_queue $id:s(struct $id:ctx* ctx) { return ctx->opencl.queue; - }|]) + }|] + ) GC.profileReport [C.citem|OPENCL_SUCCEED_FATAL(opencl_tally_profiling_records(&ctx->opencl));|] - mapM_ GC.profileReport $ costCentreReport $ - cost_centres ++ M.keys kernels + mapM_ GC.profileReport $ + costCentreReport $ + cost_centres ++ M.keys kernels -openClDecls :: [Name] -> M.Map KernelName KernelSafety -> String -> String - -> ([C.FieldGroup], [C.Stm], [C.Definition], [C.Definition]) +openClDecls :: + [Name] -> + M.Map KernelName KernelSafety -> + String -> + String -> + ([C.FieldGroup], [C.Stm], [C.Definition], [C.Definition]) openClDecls cost_centres kernels opencl_program opencl_prelude = (ctx_fields, ctx_inits, openCL_boilerplate, openCL_load) - where opencl_program_fragments = - -- Some C compilers limit the size of literal strings, so - -- chunk the entire program into small bits here, and - -- concatenate it again at runtime. - [ [C.cinit|$string:s|] | s <- chunk 2000 (opencl_prelude++opencl_program) ] + where + opencl_program_fragments = + -- Some C compilers limit the size of literal strings, so + -- chunk the entire program into small bits here, and + -- concatenate it again at runtime. + [[C.cinit|$string:s|] | s <- chunk 2000 (opencl_prelude ++ opencl_program)] - ctx_fields = - [ [C.csdecl|int total_runs;|], - [C.csdecl|long int total_runtime;|] ] ++ - [ [C.csdecl|typename cl_kernel $id:name;|] - | name <- M.keys kernels ] ++ - concat - [ [ [C.csdecl|typename int64_t $id:(kernelRuntime name);|] - , [C.csdecl|int $id:(kernelRuns name);|] + ctx_fields = + [ [C.csdecl|int total_runs;|], + [C.csdecl|long int total_runtime;|] + ] + ++ [ [C.csdecl|typename cl_kernel $id:name;|] + | name <- M.keys kernels + ] + ++ concat + [ [ [C.csdecl|typename int64_t $id:(kernelRuntime name);|], + [C.csdecl|int $id:(kernelRuns name);|] ] - | name <- cost_centres ++ M.keys kernels ] + | name <- cost_centres ++ M.keys kernels + ] - ctx_inits = - [ [C.cstm|ctx->total_runs = 0;|], - [C.cstm|ctx->total_runtime = 0;|] ] ++ - concat - [ [ [C.cstm|ctx->$id:(kernelRuntime name) = 0;|] - , [C.cstm|ctx->$id:(kernelRuns name) = 0;|] + ctx_inits = + [ [C.cstm|ctx->total_runs = 0;|], + [C.cstm|ctx->total_runtime = 0;|] + ] + ++ concat + [ [ [C.cstm|ctx->$id:(kernelRuntime name) = 0;|], + [C.cstm|ctx->$id:(kernelRuns name) = 0;|] ] - | name <- cost_centres ++ M.keys kernels ] + | name <- cost_centres ++ M.keys kernels + ] - openCL_load = [ - [C.cedecl| + openCL_load = + [ [C.cedecl| void post_opencl_setup(struct opencl_context *ctx, struct opencl_device_option *option) { $stms:(map sizeHeuristicsCode sizeHeuristicsTable) -}|]] +}|] + ] - free_list_h = $(embedStringFile "rts/c/free_list.h") - openCL_h = $(embedStringFile "rts/c/opencl.h") + free_list_h = $(embedStringFile "rts/c/free_list.h") + openCL_h = $(embedStringFile "rts/c/opencl.h") - program_fragments = opencl_program_fragments ++ [[C.cinit|NULL|]] - openCL_boilerplate = [C.cunit| + program_fragments = opencl_program_fragments ++ [[C.cinit|NULL|]] + openCL_boilerplate = + [C.cunit| $esc:("typedef cl_mem fl_mem_t;") $esc:free_list_h $esc:openCL_h static const char *opencl_program[] = {$inits:program_fragments};|] loadKernel :: (KernelName, KernelSafety) -> C.Stm -loadKernel (name, safety) = [C.cstm|{ +loadKernel (name, safety) = + [C.cstm|{ ctx->$id:name = clCreateKernel(prog, $string:(pretty (C.toIdent name mempty)), &error); OPENCL_SUCCEED_FATAL(error); $items:set_args @@ -511,126 +577,148 @@ fprintf(stderr, "Created kernel %s.\n", $string:(pretty name)); } }|] - where set_global_failure = - [C.citem|OPENCL_SUCCEED_FATAL( + where + set_global_failure = + [C.citem|OPENCL_SUCCEED_FATAL( clSetKernelArg(ctx->$id:name, 0, sizeof(typename cl_mem), &ctx->global_failure));|] - set_global_failure_args = - [C.citem|OPENCL_SUCCEED_FATAL( + set_global_failure_args = + [C.citem|OPENCL_SUCCEED_FATAL( clSetKernelArg(ctx->$id:name, 2, sizeof(typename cl_mem), &ctx->global_failure_args));|] - set_args = case safety of - SafetyNone -> [] - SafetyCheap -> [set_global_failure] - SafetyFull -> [set_global_failure, set_global_failure_args] + set_args = case safety of + SafetyNone -> [] + SafetyCheap -> [set_global_failure] + SafetyFull -> [set_global_failure, set_global_failure_args] releaseKernel :: (KernelName, KernelSafety) -> C.Stm releaseKernel (name, _) = [C.cstm|OPENCL_SUCCEED_FATAL(clReleaseKernel(ctx->$id:name));|] kernelRuntime :: KernelName -> Name -kernelRuntime = (<>"_total_runtime") +kernelRuntime = (<> "_total_runtime") kernelRuns :: KernelName -> Name -kernelRuns = (<>"_runs") +kernelRuns = (<> "_runs") costCentreReport :: [Name] -> [C.BlockItem] costCentreReport names = report_kernels ++ [report_total] - where longest_name = foldl max 0 $ map (length . pretty) names - report_kernels = concatMap reportKernel names - format_string name = - let padding = replicate (longest_name - length name) ' ' - in unwords [name ++ padding, - "ran %5d times; avg: %8ldus; total: %8ldus\n"] - reportKernel name = - let runs = kernelRuns name - total_runtime = kernelRuntime name - in [[C.citem| + where + longest_name = foldl max 0 $ map (length . pretty) names + report_kernels = concatMap reportKernel names + format_string name = + let padding = replicate (longest_name - length name) ' ' + in unwords + [ name ++ padding, + "ran %5d times; avg: %8ldus; total: %8ldus\n" + ] + reportKernel name = + let runs = kernelRuns name + total_runtime = kernelRuntime name + in [ [C.citem| str_builder(&builder, $string:(format_string (pretty name)), ctx->$id:runs, (long int) ctx->$id:total_runtime / (ctx->$id:runs != 0 ? ctx->$id:runs : 1), (long int) ctx->$id:total_runtime); |], - [C.citem|ctx->total_runtime += ctx->$id:total_runtime;|], - [C.citem|ctx->total_runs += ctx->$id:runs;|]] + [C.citem|ctx->total_runtime += ctx->$id:total_runtime;|], + [C.citem|ctx->total_runs += ctx->$id:runs;|] + ] - report_total = [C.citem| + report_total = + [C.citem| str_builder(&builder, "%d operations with cumulative runtime: %6ldus\n", ctx->total_runs, ctx->total_runtime); |] sizeHeuristicsCode :: SizeHeuristic -> C.Stm -sizeHeuristicsCode (SizeHeuristic platform_name device_type which what) = +sizeHeuristicsCode (SizeHeuristic platform_name device_type which (TPrimExp what)) = [C.cstm| if ($exp:which' == 0 && strstr(option->platform_name, $string:platform_name) != NULL && (option->device_type & $exp:(clDeviceType device_type)) == $exp:(clDeviceType device_type)) { $items:get_size }|] - where clDeviceType DeviceGPU = [C.cexp|CL_DEVICE_TYPE_GPU|] - clDeviceType DeviceCPU = [C.cexp|CL_DEVICE_TYPE_CPU|] + where + clDeviceType DeviceGPU = [C.cexp|CL_DEVICE_TYPE_GPU|] + clDeviceType DeviceCPU = [C.cexp|CL_DEVICE_TYPE_CPU|] - which' = case which of - LockstepWidth -> [C.cexp|ctx->lockstep_width|] - NumGroups -> [C.cexp|ctx->cfg.default_num_groups|] - GroupSize -> [C.cexp|ctx->cfg.default_group_size|] - TileSize -> [C.cexp|ctx->cfg.default_tile_size|] - Threshold -> [C.cexp|ctx->cfg.default_threshold|] + which' = case which of + LockstepWidth -> [C.cexp|ctx->lockstep_width|] + NumGroups -> [C.cexp|ctx->cfg.default_num_groups|] + GroupSize -> [C.cexp|ctx->cfg.default_group_size|] + TileSize -> [C.cexp|ctx->cfg.default_tile_size|] + Threshold -> [C.cexp|ctx->cfg.default_threshold|] - get_size = - let (e, m) = runState (GC.compilePrimExp onLeaf what) mempty - in concat (M.elems m) ++ [[C.citem|$exp:which' = $exp:e;|]] + get_size = + let (e, m) = runState (GC.compilePrimExp onLeaf what) mempty + in concat (M.elems m) ++ [[C.citem|$exp:which' = $exp:e;|]] - onLeaf (DeviceInfo s) = do - let s' = "CL_DEVICE_" ++ s - v = s ++ "_val" - m <- get - case M.lookup s m of - Nothing -> - -- Cheating with the type here; works for the infos we - -- currently use, but should be made more size-aware in - -- the future. - modify $ M.insert s' + onLeaf (DeviceInfo s) = do + let s' = "CL_DEVICE_" ++ s + v = s ++ "_val" + m <- get + case M.lookup s m of + Nothing -> + -- Cheating with the type here; works for the infos we + -- currently use, but should be made more size-aware in + -- the future. + modify $ + M.insert + s' [C.citems|size_t $id:v; clGetDeviceInfo(ctx->device, $id:s', sizeof($id:v), &$id:v, NULL);|] - Just _ -> return () + Just _ -> return () - return [C.cexp|$id:v|] + return [C.cexp|$id:v|] -- Options that are common to multiple GPU-like backends. commonOptions :: [Option] commonOptions = - [ Option { optionLongName = "device" - , optionShortName = Just 'd' - , optionArgument = RequiredArgument "NAME" - , optionAction = [C.cstm|futhark_context_config_set_device(cfg, optarg);|] - } - , Option { optionLongName = "default-group-size" - , optionShortName = Nothing - , optionArgument = RequiredArgument "INT" - , optionAction = [C.cstm|futhark_context_config_set_default_group_size(cfg, atoi(optarg));|] - } - , Option { optionLongName = "default-num-groups" - , optionShortName = Nothing - , optionArgument = RequiredArgument "INT" - , optionAction = [C.cstm|futhark_context_config_set_default_num_groups(cfg, atoi(optarg));|] - } - , Option { optionLongName = "default-tile-size" - , optionShortName = Nothing - , optionArgument = RequiredArgument "INT" - , optionAction = [C.cstm|futhark_context_config_set_default_tile_size(cfg, atoi(optarg));|] - } - , Option { optionLongName = "default-threshold" - , optionShortName = Nothing - , optionArgument = RequiredArgument "INT" - , optionAction = [C.cstm|futhark_context_config_set_default_threshold(cfg, atoi(optarg));|] - } - , Option { optionLongName = "print-sizes" - , optionShortName = Nothing - , optionArgument = NoArgument - , optionAction = [C.cstm|{ + [ Option + { optionLongName = "device", + optionShortName = Just 'd', + optionArgument = RequiredArgument "NAME", + optionDescription = "Use the first OpenCL device whose name contains the given string.", + optionAction = [C.cstm|futhark_context_config_set_device(cfg, optarg);|] + }, + Option + { optionLongName = "default-group-size", + optionShortName = Nothing, + optionArgument = RequiredArgument "INT", + optionDescription = "The default size of OpenCL workgroups that are launched.", + optionAction = [C.cstm|futhark_context_config_set_default_group_size(cfg, atoi(optarg));|] + }, + Option + { optionLongName = "default-num-groups", + optionShortName = Nothing, + optionArgument = RequiredArgument "INT", + optionDescription = "The default number of OpenCL workgroups that are launched.", + optionAction = [C.cstm|futhark_context_config_set_default_num_groups(cfg, atoi(optarg));|] + }, + Option + { optionLongName = "default-tile-size", + optionShortName = Nothing, + optionArgument = RequiredArgument "INT", + optionDescription = "The default tile size used when performing two-dimensional tiling.", + optionAction = [C.cstm|futhark_context_config_set_default_tile_size(cfg, atoi(optarg));|] + }, + Option + { optionLongName = "default-threshold", + 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), @@ -638,11 +726,14 @@ } exit(0); }|] - } - , Option { optionLongName = "size" - , optionShortName = Nothing - , optionArgument = RequiredArgument "NAME=INT" - , optionAction = [C.cstm|{ + }, + 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; @@ -655,15 +746,18 @@ } else { futhark_panic(1, "Invalid argument for size option: %s\n", optarg); }}|] - } - , Option { optionLongName = "tuning" - , optionShortName = Nothing - , optionArgument = RequiredArgument "FILE" - , optionAction = [C.cstm|{ + }, + 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/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs --- a/src/Futhark/CodeGen/Backends/GenericC.hs +++ b/src/Futhark/CodeGen/Backends/GenericC.hs @@ -1,2200 +1,2414 @@ -{-# LANGUAGE QuasiQuotes, GeneralizedNewtypeDeriving, FlexibleInstances #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE Trustworthy #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} --- | C code generator framework. -module Futhark.CodeGen.Backends.GenericC - ( compileProg - , CParts(..) - , asLibrary - , asExecutable - - -- * Pluggable compiler - , Operations (..) - , defaultOperations - , OpCompiler - , ErrorCompiler - , CallCompiler - - , PointerQuals - , MemoryType - , WriteScalar - , writeScalarPointerWithQuals - , ReadScalar - , readScalarPointerWithQuals - , Allocate - , Deallocate - , Copy - , StaticArray - - -- * Monadic compiler interface - , CompilerM - , CompilerState (compUserState) - , getUserState - , modifyUserState - , contextContents - , contextFinalInits - , runCompilerM - , cachingMemory - , blockScope - , compileFun - , compileCode - , compileExp - , compilePrimExp - , compilePrimValue - , compileExpToName - , rawMem - , item - , items - , stm - , stms - , decl - , atInit - , headerDecl - , publicDef - , publicDef_ - , profileReport - , HeaderSection(..) - , libDecl - , earlyDecl - , publicName - , contextType - , contextField - - -- * Building Blocks - , primTypeToCType - , copyMemoryDefaultSpace - ) where - -import Control.Monad.Identity -import Control.Monad.State -import Control.Monad.Reader -import Control.Monad.Writer -import Control.Monad.RWS -import Data.Bifunctor (first) -import Data.Bits (xor, shiftR) -import Data.Char (ord, isDigit, isAlphaNum) -import qualified Data.Map.Strict as M -import qualified Data.DList as DL -import Data.List (unzip4) -import Data.Loc -import Data.Maybe -import Data.FileEmbed -import Text.Printf - -import qualified Language.C.Syntax as C -import qualified Language.C.Quote.OpenCL as C - -import Futhark.CodeGen.ImpCode -import Futhark.MonadFreshNames -import Futhark.CodeGen.Backends.SimpleRep -import Futhark.CodeGen.Backends.GenericC.Options -import Futhark.Util (zEncodeString) -import Futhark.IR.Prop (isBuiltInFunction) - - -data CompilerState s = CompilerState { - compArrayStructs :: [((C.Type, Int), (C.Type, [C.Definition]))] - , compOpaqueStructs :: [(String, (C.Type, [C.Definition]))] - , compEarlyDecls :: DL.DList C.Definition - , compInit :: [C.Stm] - , compNameSrc :: VNameSource - , compUserState :: s - , compHeaderDecls :: M.Map HeaderSection (DL.DList C.Definition) - , compLibDecls :: DL.DList C.Definition - , compCtxFields :: DL.DList (C.Id, C.Type, Maybe C.Exp) - , compProfileItems :: DL.DList C.BlockItem - , compDeclaredMem :: [(VName,Space)] - } - -newCompilerState :: VNameSource -> s -> CompilerState s -newCompilerState src s = CompilerState { compArrayStructs = [] - , compOpaqueStructs = [] - , compEarlyDecls = mempty - , compInit = [] - , compNameSrc = src - , compUserState = s - , compHeaderDecls = mempty - , compLibDecls = mempty - , compCtxFields = mempty - , compProfileItems = mempty - , compDeclaredMem = mempty - } - --- | In which part of the header file we put the declaration. This is --- to ensure that the header file remains structured and readable. -data HeaderSection = ArrayDecl String - | OpaqueDecl String - | EntryDecl - | MiscDecl - | InitDecl - deriving (Eq, Ord) - --- | A substitute expression compiler, tried before the main --- compilation function. -type OpCompiler op s = op -> CompilerM op s () - -type ErrorCompiler op s = ErrorMsg Exp -> String -> CompilerM op s () - --- | The address space qualifiers for a pointer of the given type with --- the given annotation. -type PointerQuals op s = String -> CompilerM op s [C.TypeQual] - --- | The type of a memory block in the given memory space. -type MemoryType op s = SpaceId -> CompilerM op s C.Type - --- | Write a scalar to the given memory block with the given element --- index and in the given memory space. -type WriteScalar op s = - C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> C.Exp -> CompilerM op s () - --- | Read a scalar from the given memory block with the given element --- index and in the given memory space. -type ReadScalar op s = - C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> CompilerM op s C.Exp - --- | Allocate a memory block of the given size and with the given tag --- in the given memory space, saving a reference in the given variable --- name. -type Allocate op s = C.Exp -> C.Exp -> C.Exp -> SpaceId - -> CompilerM op s () - --- | De-allocate the given memory block with the given tag, which is --- in the given memory space. -type Deallocate op s = C.Exp -> C.Exp -> SpaceId -> CompilerM op s () - --- | Create a static array of values - initialised at load time. -type StaticArray op s = VName -> SpaceId -> PrimType -> ArrayContents -> CompilerM op s () - --- | Copy from one memory block to another. -type Copy op s = C.Exp -> C.Exp -> Space -> - C.Exp -> C.Exp -> Space -> - C.Exp -> - CompilerM op s () - --- | Call a function. -type CallCompiler op s = [VName] -> Name -> [C.Exp] -> CompilerM op s () - -data Operations op s = - Operations { opsWriteScalar :: WriteScalar op s - , opsReadScalar :: ReadScalar op s - , opsAllocate :: Allocate op s - , opsDeallocate :: Deallocate op s - , opsCopy :: Copy op s - , opsStaticArray :: StaticArray op s - - , opsMemoryType :: MemoryType op s - , opsCompiler :: OpCompiler op s - , opsError :: ErrorCompiler op s - , opsCall :: CallCompiler op s - - , opsFatMemory :: Bool - -- ^ If true, use reference counting. Otherwise, bare - -- pointers. - } - -defError :: ErrorCompiler op s -defError (ErrorMsg parts) stacktrace = do - free_all_mem <- collect $ mapM_ (uncurry unRefMem) =<< gets compDeclaredMem - let onPart (ErrorString s) = return ("%s", [C.cexp|$string:s|]) - onPart (ErrorInt32 x) = ("%d",) <$> compileExp x - (formatstrs, formatargs) <- unzip <$> mapM onPart parts - let formatstr = "Error: " ++ concat formatstrs ++ "\n\nBacktrace:\n%s" - items [C.citems|ctx->error = msgprintf($string:formatstr, $args:formatargs, $string:stacktrace); - $items:free_all_mem - return 1;|] - -defCall :: CallCompiler op s -defCall dests fname args = do - let out_args = [ [C.cexp|&$id:d|] | d <- dests ] - args' | isBuiltInFunction fname = args - | otherwise = [C.cexp|ctx|] : out_args ++ args - case dests of - [dest] | isBuiltInFunction fname -> - stm [C.cstm|$id:dest = $id:(funName fname)($args:args');|] - _ -> - item [C.citem|if ($id:(funName fname)($args:args') != 0) { err = 1; goto cleanup; }|] - --- | A set of operations that fail for every operation involving --- non-default memory spaces. Uses plain pointers and @malloc@ for --- memory management. -defaultOperations :: Operations op s -defaultOperations = Operations { opsWriteScalar = defWriteScalar - , opsReadScalar = defReadScalar - , opsAllocate = defAllocate - , opsDeallocate = defDeallocate - , opsCopy = defCopy - , opsStaticArray = defStaticArray - , opsMemoryType = defMemoryType - , opsCompiler = defCompiler - , opsFatMemory = True - , opsError = defError - , opsCall = defCall - } - where defWriteScalar _ _ _ _ _ = - error "Cannot write to non-default memory space because I am dumb" - defReadScalar _ _ _ _ = - error "Cannot read from non-default memory space" - defAllocate _ _ _ = - error "Cannot allocate in non-default memory space" - defDeallocate _ _ = - error "Cannot deallocate in non-default memory space" - defCopy destmem destoffset DefaultSpace srcmem srcoffset DefaultSpace size = - copyMemoryDefaultSpace destmem destoffset srcmem srcoffset size - defCopy _ _ _ _ _ _ _ = - error "Cannot copy to or from non-default memory space" - defStaticArray _ _ _ _ = - error "Cannot create static array in non-default memory space" - defMemoryType _ = - error "Has no type for non-default memory space" - defCompiler _ = - error "The default compiler cannot compile extended operations" - - -data CompilerEnv op s = CompilerEnv - { envOperations :: Operations op s - , envCachedMem :: M.Map C.Exp VName - -- ^ Mapping memory blocks to sizes. These memory blocks are CPU - -- memory that we know are used in particularly simple ways (no - -- reference counting necessary). To cut down on allocator - -- pressure, we keep these allocations around for a long time, and - -- record their sizes so we can reuse them if possible (and - -- realloc() when needed). - } - -newtype CompilerAcc op s = CompilerAcc { - accItems :: DL.DList C.BlockItem - } - -instance Semigroup (CompilerAcc op s) where - CompilerAcc items1 <> CompilerAcc items2 = - CompilerAcc (items1<>items2) - -instance Monoid (CompilerAcc op s) where - mempty = CompilerAcc mempty - -envOpCompiler :: CompilerEnv op s -> OpCompiler op s -envOpCompiler = opsCompiler . envOperations - -envMemoryType :: CompilerEnv op s -> MemoryType op s -envMemoryType = opsMemoryType . envOperations - -envReadScalar :: CompilerEnv op s -> ReadScalar op s -envReadScalar = opsReadScalar . envOperations - -envWriteScalar :: CompilerEnv op s -> WriteScalar op s -envWriteScalar = opsWriteScalar . envOperations - -envAllocate :: CompilerEnv op s -> Allocate op s -envAllocate = opsAllocate . envOperations - -envDeallocate :: CompilerEnv op s -> Deallocate op s -envDeallocate = opsDeallocate . envOperations - -envCopy :: CompilerEnv op s -> Copy op s -envCopy = opsCopy . envOperations - -envStaticArray :: CompilerEnv op s -> StaticArray op s -envStaticArray = opsStaticArray . envOperations - -envFatMemory :: CompilerEnv op s -> Bool -envFatMemory = opsFatMemory . envOperations - -arrayDefinitions, opaqueDefinitions :: CompilerState s -> [C.Definition] -arrayDefinitions = concatMap (snd . snd) . compArrayStructs -opaqueDefinitions = concatMap (snd . snd) . compOpaqueStructs - -initDecls, arrayDecls, opaqueDecls, entryDecls, miscDecls :: CompilerState s -> [C.Definition] -initDecls = concatMap (DL.toList . snd) . filter ((==InitDecl) . fst) . M.toList . compHeaderDecls -arrayDecls = concatMap (DL.toList . snd) . filter (isArrayDecl . fst) . M.toList . compHeaderDecls - where isArrayDecl ArrayDecl{} = True - isArrayDecl _ = False -opaqueDecls = concatMap (DL.toList . snd) . filter (isOpaqueDecl . fst) . M.toList . compHeaderDecls - where isOpaqueDecl OpaqueDecl{} = True - isOpaqueDecl _ = False -entryDecls = concatMap (DL.toList . snd) . filter ((==EntryDecl) . fst) . M.toList . compHeaderDecls -miscDecls = concatMap (DL.toList . snd) . filter ((==MiscDecl) . fst) . M.toList . compHeaderDecls - -contextContents :: CompilerM op s ([C.FieldGroup], [C.Stm]) -contextContents = do - (field_names, field_types, field_values) <- gets $ unzip3 . DL.toList . compCtxFields - let fields = [ [C.csdecl|$ty:ty $id:name;|] - | (name, ty) <- zip field_names field_types ] - init_fields = [ [C.cstm|ctx->$id:name = $exp:e;|] - | (name, Just e) <- zip field_names field_values ] - return (fields, init_fields) - -contextFinalInits :: CompilerM op s [C.Stm] -contextFinalInits = gets compInit - -newtype CompilerM op s a = CompilerM (RWS - (CompilerEnv op s) - (CompilerAcc op s) - (CompilerState s) a) - deriving (Functor, Applicative, Monad, - MonadState (CompilerState s), - MonadReader (CompilerEnv op s), - MonadWriter (CompilerAcc op s)) - -instance MonadFreshNames (CompilerM op s) where - getNameSource = gets compNameSrc - putNameSource src = modify $ \s -> s { compNameSrc = src } - -runCompilerM :: Operations op s -> VNameSource -> s - -> CompilerM op s a - -> (a, CompilerState s) -runCompilerM ops src userstate (CompilerM m) = - let (x, s, _) = runRWS m (CompilerEnv ops mempty) (newCompilerState src userstate) - in (x, s) - -getUserState :: CompilerM op s s -getUserState = gets compUserState - -modifyUserState :: (s -> s) -> CompilerM op s () -modifyUserState f = modify $ \compstate -> - compstate { compUserState = f $ compUserState compstate } - -atInit :: C.Stm -> CompilerM op s () -atInit x = modify $ \s -> - s { compInit = compInit s ++ [x] } - -collect :: CompilerM op s () -> CompilerM op s [C.BlockItem] -collect m = snd <$> collect' m - -collect' :: CompilerM op s a -> CompilerM op s (a, [C.BlockItem]) -collect' m = pass $ do - (x, w) <- listen m - return ((x, DL.toList $ accItems w), - const w { accItems = mempty}) - -item :: C.BlockItem -> CompilerM op s () -item x = tell $ mempty { accItems = DL.singleton x } - -items :: [C.BlockItem] -> CompilerM op s () -items = mapM_ item - -fatMemory :: Space -> CompilerM op s Bool -fatMemory ScalarSpace{} = return False -fatMemory _ = asks envFatMemory - -cacheMem :: C.ToExp a => a -> CompilerM op s (Maybe VName) -cacheMem a = asks $ M.lookup (C.toExp a noLoc) . envCachedMem - -instance C.ToIdent Name where - toIdent = C.toIdent . zEncodeString . nameToString - -instance C.ToIdent VName where - toIdent = C.toIdent . zEncodeString . pretty - -instance C.ToExp VName where - toExp v _ = [C.cexp|$id:v|] - -instance C.ToExp IntValue where - toExp (Int8Value v) = C.toExp v - toExp (Int16Value v) = C.toExp v - toExp (Int32Value v) = C.toExp v - toExp (Int64Value v) = C.toExp v - -instance C.ToExp FloatValue where - toExp (Float32Value v) = C.toExp v - toExp (Float64Value v) = C.toExp v - -instance C.ToExp PrimValue where - toExp (IntValue v) = C.toExp v - toExp (FloatValue v) = C.toExp v - toExp (BoolValue True) = C.toExp (1::Int8) - toExp (BoolValue False) = C.toExp (0::Int8) - toExp Checked = C.toExp (1::Int8) - -instance C.ToExp SubExp where - toExp (Var v) = C.toExp v - toExp (Constant c) = C.toExp c - --- | Construct a publicly visible definition using the specified name --- as the template. The first returned definition is put in the --- header file, and the second is the implementation. Returns the public --- name. -publicDef :: String -> HeaderSection -> (String -> (C.Definition, C.Definition)) - -> CompilerM op s String -publicDef s h f = do - s' <- publicName s - let (pub, priv) = f s' - headerDecl h pub - earlyDecl priv - return s' - --- | As 'publicDef', but ignores the public name. -publicDef_ :: String -> HeaderSection -> (String -> (C.Definition, C.Definition)) - -> CompilerM op s () -publicDef_ s h f = void $ publicDef s h f - -headerDecl :: HeaderSection -> C.Definition -> CompilerM op s () -headerDecl sec def = modify $ \s -> - s { compHeaderDecls = M.unionWith (<>) (compHeaderDecls s) - (M.singleton sec (DL.singleton def)) } - -libDecl :: C.Definition -> CompilerM op s () -libDecl def = modify $ \s -> - s { compLibDecls = compLibDecls s <> DL.singleton def } - -earlyDecl :: C.Definition -> CompilerM op s () -earlyDecl def = modify $ \s -> - s { compEarlyDecls = compEarlyDecls s <> DL.singleton def } - -contextField :: C.Id -> C.Type -> Maybe C.Exp -> CompilerM op s () -contextField name ty initial = modify $ \s -> - s { compCtxFields = compCtxFields s <> DL.singleton (name,ty,initial) } - -profileReport :: C.BlockItem -> CompilerM op s () -profileReport x = modify $ \s -> - s { compProfileItems = compProfileItems s <> DL.singleton x } - -stm :: C.Stm -> CompilerM op s () -stm s = item [C.citem|$stm:s|] - -stms :: [C.Stm] -> CompilerM op s () -stms = mapM_ stm - -decl :: C.InitGroup -> CompilerM op s () -decl x = item [C.citem|$decl:x;|] - -addrOf :: C.Exp -> C.Exp -addrOf e = [C.cexp|&$exp:e|] - --- | Public names must have a consitent prefix. -publicName :: String -> CompilerM op s String -publicName s = return $ "futhark_" ++ s - --- | The generated code must define a struct with this name. -contextType :: CompilerM op s C.Type -contextType = do - name <- publicName "context" - return [C.cty|struct $id:name|] - -memToCType :: VName -> Space -> CompilerM op s C.Type -memToCType v space = do - refcount <- fatMemory space - cached <- isJust <$> cacheMem v - if refcount && not cached - then return $ fatMemType space - else rawMemCType space - -rawMemCType :: Space -> CompilerM op s C.Type -rawMemCType DefaultSpace = return defaultMemBlockType -rawMemCType (Space sid) = join $ asks envMemoryType <*> pure sid -rawMemCType (ScalarSpace [] t) = - return [C.cty|$ty:(primTypeToCType t)[1]|] -rawMemCType (ScalarSpace ds t) = - return [C.cty|$ty:(primTypeToCType t)[$exp:(cproduct ds')]|] - where ds' = map (`C.toExp` noLoc) ds - -fatMemType :: Space -> C.Type -fatMemType space = - [C.cty|struct $id:name|] - where name = case space of - Space sid -> "memblock_" ++ sid - _ -> "memblock" - -fatMemSet :: Space -> String -fatMemSet (Space sid) = "memblock_set_" ++ sid -fatMemSet _ = "memblock_set" - -fatMemAlloc :: Space -> String -fatMemAlloc (Space sid) = "memblock_alloc_" ++ sid -fatMemAlloc _ = "memblock_alloc" - -fatMemUnRef :: Space -> String -fatMemUnRef (Space sid) = "memblock_unref_" ++ sid -fatMemUnRef _ = "memblock_unref" - -rawMem :: VName -> CompilerM op s C.Exp -rawMem v = rawMem' <$> fat <*> pure v - where fat = asks ((&&) . envFatMemory) <*> (isNothing <$> cacheMem v) - -rawMem' :: C.ToExp a => Bool -> a -> C.Exp -rawMem' True e = [C.cexp|$exp:e.mem|] -rawMem' False e = [C.cexp|$exp:e|] - -allocRawMem :: (C.ToExp a, C.ToExp b, C.ToExp c) => - a -> b -> Space -> c -> CompilerM op s () -allocRawMem dest size space desc = case space of - Space sid -> - join $ asks envAllocate <*> pure [C.cexp|$exp:dest|] <*> - pure [C.cexp|$exp:size|] <*> pure [C.cexp|$exp:desc|] <*> pure sid - _ -> - stm [C.cstm|$exp:dest = (char*) malloc($exp:size);|] - -freeRawMem :: (C.ToExp a, C.ToExp b) => - a -> Space -> b -> CompilerM op s () -freeRawMem mem space desc = - case space of - Space sid -> do free_mem <- asks envDeallocate - free_mem [C.cexp|$exp:mem|] [C.cexp|$exp:desc|] sid - _ -> item [C.citem|free($exp:mem);|] - -defineMemorySpace :: Space -> CompilerM op s (C.Definition, [C.Definition], C.BlockItem) -defineMemorySpace space = do - rm <- rawMemCType space - let structdef = - [C.cedecl|struct $id:sname { int *references; - $ty:rm mem; - typename int64_t size; - const char *desc; };|] - - contextField peakname [C.cty|typename int64_t|] $ Just [C.cexp|0|] - contextField usagename [C.cty|typename int64_t|] $ Just [C.cexp|0|] - - -- Unreferencing a memory block consists of decreasing its reference - -- count and freeing the corresponding memory if the count reaches - -- zero. - free <- collect $ freeRawMem [C.cexp|block->mem|] space [C.cexp|desc|] - ctx_ty <- contextType - let unrefdef = [C.cedecl|static int $id:(fatMemUnRef space) ($ty:ctx_ty *ctx, $ty:mty *block, const char *desc) { - if (block->references != NULL) { - *(block->references) -= 1; - if (ctx->detail_memory) { - fprintf(stderr, "Unreferencing block %s (allocated as %s) in %s: %d references remaining.\n", - desc, block->desc, $string:spacedesc, *(block->references)); - } - if (*(block->references) == 0) { - ctx->$id:usagename -= block->size; - $items:free - free(block->references); - if (ctx->detail_memory) { - fprintf(stderr, "%lld bytes freed (now allocated: %lld bytes)\n", - (long long) block->size, (long long) ctx->$id:usagename); - } - } - block->references = NULL; - } - return 0; -}|] - - -- When allocating a memory block we initialise the reference count to 1. - alloc <- collect $ - allocRawMem [C.cexp|block->mem|] [C.cexp|size|] space [C.cexp|desc|] - let allocdef = [C.cedecl|static int $id:(fatMemAlloc space) ($ty:ctx_ty *ctx, $ty:mty *block, typename int64_t size, const char *desc) { - if (size < 0) { - futhark_panic(1, "Negative allocation of %lld bytes attempted for %s in %s.\n", - (long long)size, desc, $string:spacedesc, ctx->$id:usagename); - } - int ret = $id:(fatMemUnRef space)(ctx, block, desc); - - ctx->$id:usagename += size; - if (ctx->detail_memory) { - fprintf(stderr, "Allocating %lld bytes for %s in %s (then allocated: %lld bytes)", - (long long) size, - desc, $string:spacedesc, - (long long) ctx->$id:usagename); - } - if (ctx->$id:usagename > ctx->$id:peakname) { - ctx->$id:peakname = ctx->$id:usagename; - if (ctx->detail_memory) { - fprintf(stderr, " (new peak).\n"); - } - } else if (ctx->detail_memory) { - fprintf(stderr, ".\n"); - } - - $items:alloc - block->references = (int*) malloc(sizeof(int)); - *(block->references) = 1; - block->size = size; - block->desc = desc; - return ret; - }|] - - -- Memory setting - unreference the destination and increase the - -- count of the source by one. - let setdef = [C.cedecl|static int $id:(fatMemSet space) ($ty:ctx_ty *ctx, $ty:mty *lhs, $ty:mty *rhs, const char *lhs_desc) { - int ret = $id:(fatMemUnRef space)(ctx, lhs, lhs_desc); - (*(rhs->references))++; - *lhs = *rhs; - return ret; -} -|] - - let peakmsg = "Peak memory usage for " ++ spacedesc ++ ": %lld bytes.\n" - return (structdef, - [unrefdef, allocdef, setdef], - -- Do not report memory usage for DefaultSpace (CPU memory), - -- because it would not be accurate anyway. This whole - -- tracking probably needs to be rethought. - if space == DefaultSpace - then [C.citem|{}|] - else [C.citem|str_builder(&builder, $string:peakmsg, (long long) ctx->$id:peakname);|]) - where mty = fatMemType space - (peakname, usagename, sname, spacedesc) = case space of - Space sid -> (C.toIdent ("peak_mem_usage_" ++ sid) noLoc, - C.toIdent ("cur_mem_usage_" ++ sid) noLoc, - C.toIdent ("memblock_" ++ sid) noLoc, - "space '" ++ sid ++ "'") - _ -> ("peak_mem_usage_default", - "cur_mem_usage_default", - "memblock", - "default space") - -declMem :: VName -> Space -> CompilerM op s () -declMem name space = do - cached <- isJust <$> cacheMem name - unless cached $ do - ty <- memToCType name space - decl [C.cdecl|$ty:ty $id:name;|] - resetMem name space - modify $ \s -> s { compDeclaredMem = (name, space) : compDeclaredMem s } - -resetMem :: C.ToExp a => a -> Space -> CompilerM op s () -resetMem mem space = do - refcount <- fatMemory space - cached <- isJust <$> cacheMem mem - if cached - then stm [C.cstm|$exp:mem = NULL;|] - else when refcount $ - stm [C.cstm|$exp:mem.references = NULL;|] - -setMem :: (C.ToExp a, C.ToExp b) => a -> b -> Space -> CompilerM op s () -setMem dest src space = do - refcount <- fatMemory space - let src_s = pretty $ C.toExp src noLoc - if refcount - then stm [C.cstm|if ($id:(fatMemSet space)(ctx, &$exp:dest, &$exp:src, - $string:src_s) != 0) { - return 1; - }|] - else case space of - ScalarSpace ds _ -> do - i' <- newVName "i" - let i = C.toIdent i' - it = primTypeToCType $ IntType Int32 - ds' = map (`C.toExp` noLoc) ds - bound = cproduct ds' - stm [C.cstm|for ($ty:it $id:i = 0; $id:i < $exp:bound; $id:i++) { - $exp:dest[$id:i] = $exp:src[$id:i]; - }|] - _ -> stm [C.cstm|$exp:dest = $exp:src;|] - -unRefMem :: C.ToExp a => a -> Space -> CompilerM op s () -unRefMem mem space = do - refcount <- fatMemory space - cached <- isJust <$> cacheMem mem - let mem_s = pretty $ C.toExp mem noLoc - when (refcount && not cached) $ - stm [C.cstm|if ($id:(fatMemUnRef space)(ctx, &$exp:mem, $string:mem_s) != 0) { - return 1; - }|] - -allocMem :: (C.ToExp a, C.ToExp b) => - a -> b -> Space -> C.Stm -> CompilerM op s () -allocMem mem size space on_failure = do - refcount <- fatMemory space - let mem_s = pretty $ C.toExp mem noLoc - if refcount - then stm [C.cstm|if ($id:(fatMemAlloc space)(ctx, &$exp:mem, $exp:size, - $string:mem_s)) { - $stm:on_failure - }|] - else do freeRawMem mem space mem_s - allocRawMem mem size space [C.cexp|desc|] - -primTypeInfo :: PrimType -> Signedness -> C.Exp -primTypeInfo (IntType it) t = case (it, t) of - (Int8, TypeUnsigned) -> [C.cexp|u8_info|] - (Int16, TypeUnsigned) -> [C.cexp|u16_info|] - (Int32, TypeUnsigned) -> [C.cexp|u32_info|] - (Int64, TypeUnsigned) -> [C.cexp|u64_info|] - (Int8, _) -> [C.cexp|i8_info|] - (Int16, _) -> [C.cexp|i16_info|] - (Int32, _) -> [C.cexp|i32_info|] - (Int64, _) -> [C.cexp|i64_info|] -primTypeInfo (FloatType Float32) _ = [C.cexp|f32_info|] -primTypeInfo (FloatType Float64) _ = [C.cexp|f64_info|] -primTypeInfo Bool _ = [C.cexp|bool_info|] -primTypeInfo Cert _ = [C.cexp|bool_info|] - -copyMemoryDefaultSpace :: C.Exp -> C.Exp -> C.Exp -> C.Exp -> C.Exp -> - CompilerM op s () -copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes = - stm [C.cstm|memmove($exp:destmem + $exp:destidx, - $exp:srcmem + $exp:srcidx, - $exp:nbytes);|] - ---- Entry points. - -arrayName :: PrimType -> Signedness -> Int -> String -arrayName pt signed rank = - prettySigned (signed==TypeUnsigned) pt ++ "_" ++ show rank ++ "d" - -opaqueName :: String -> [ValueDesc] -> String -opaqueName s _ - | valid = "opaque_" ++ s - where valid = head s /= '_' && - not (isDigit $ head s) && - all ok s - ok c = isAlphaNum c || c == '_' -opaqueName s vds = "opaque_" ++ hash (zipWith xor [0..] $ map ord (s ++ concatMap p vds)) - where p (ScalarValue pt signed _) = - show (pt, signed) - p (ArrayValue _ space pt signed dims) = - show (space, pt, signed, length dims) - - -- FIXME: a stupid hash algorithm; may have collisions. - hash = printf "%x" . foldl xor 0 . map (iter . (*0x45d9f3b) . - iter . (*0x45d9f3b) . - iter . fromIntegral) - iter x = ((x::Word32) `shiftR` 16) `xor` x - -criticalSection :: [C.BlockItem] -> [C.BlockItem] -criticalSection x = [C.citems| - lock_lock(&ctx->lock); - $items:x - lock_unlock(&ctx->lock); - |] - -arrayLibraryFunctions :: Space -> PrimType -> Signedness -> [DimSize] - -> CompilerM op s [C.Definition] -arrayLibraryFunctions space pt signed shape = do - let rank = length shape - pt' = signedPrimTypeToCType signed pt - name = arrayName pt signed rank - arr_name = "futhark_" ++ name - array_type = [C.cty|struct $id:arr_name|] - - new_array <- publicName $ "new_" ++ name - new_raw_array <- publicName $ "new_raw_" ++ name - free_array <- publicName $ "free_" ++ name - values_array <- publicName $ "values_" ++ name - values_raw_array <- publicName $ "values_raw_" ++ name - shape_array <- publicName $ "shape_" ++ name - - let shape_names = [ "dim"++show i | i <- [0..rank-1] ] - shape_params = [ [C.cparam|typename int64_t $id:k|] | k <- shape_names ] - arr_size = cproduct [ [C.cexp|$id:k|] | k <- shape_names ] - arr_size_array = cproduct [ [C.cexp|arr->shape[$int:i]|] | i <- [0..rank-1] ] - copy <- asks envCopy - - memty <- rawMemCType space - - let prepare_new = do - resetMem [C.cexp|arr->mem|] space - allocMem [C.cexp|arr->mem|] [C.cexp|((size_t)$exp:arr_size) * sizeof($ty:pt')|] space - [C.cstm|return NULL;|] - forM_ [0..rank-1] $ \i -> - let dim_s = "dim"++show i - in stm [C.cstm|arr->shape[$int:i] = $id:dim_s;|] - - new_body <- collect $ do - prepare_new - copy [C.cexp|arr->mem.mem|] [C.cexp|0|] space - [C.cexp|data|] [C.cexp|0|] DefaultSpace - [C.cexp|((size_t)$exp:arr_size) * sizeof($ty:pt')|] - - new_raw_body <- collect $ do - prepare_new - copy [C.cexp|arr->mem.mem|] [C.cexp|0|] space - [C.cexp|data|] [C.cexp|offset|] space - [C.cexp|((size_t)$exp:arr_size) * sizeof($ty:pt')|] - - free_body <- collect $ unRefMem [C.cexp|arr->mem|] space - - values_body <- collect $ - copy [C.cexp|data|] [C.cexp|0|] DefaultSpace - [C.cexp|arr->mem.mem|] [C.cexp|0|] space - [C.cexp|((size_t)$exp:arr_size_array) * sizeof($ty:pt')|] - - ctx_ty <- contextType - - headerDecl (ArrayDecl name) - [C.cedecl|struct $id:arr_name;|] - headerDecl (ArrayDecl name) - [C.cedecl|$ty:array_type* $id:new_array($ty:ctx_ty *ctx, $ty:pt' *data, $params:shape_params);|] - headerDecl (ArrayDecl name) - [C.cedecl|$ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, $ty:memty data, int offset, $params:shape_params);|] - headerDecl (ArrayDecl name) - [C.cedecl|int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] - headerDecl (ArrayDecl name) - [C.cedecl|int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data);|] - headerDecl (ArrayDecl name) - [C.cedecl|$ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] - headerDecl (ArrayDecl name) - [C.cedecl|const typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] - - return [C.cunit| - $ty:array_type* $id:new_array($ty:ctx_ty *ctx, $ty:pt' *data, $params:shape_params) { - $ty:array_type* bad = NULL; - $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type)); - if (arr == NULL) { - return bad; - } - $items:(criticalSection new_body) - return arr; - } - - $ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, $ty:memty data, int offset, - $params:shape_params) { - $ty:array_type* bad = NULL; - $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type)); - if (arr == NULL) { - return bad; - } - $items:(criticalSection new_raw_body) - return arr; - } - - int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr) { - $items:(criticalSection free_body) - free(arr); - return 0; - } - - int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data) { - $items:(criticalSection values_body) - return 0; - } - - $ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr) { - (void)ctx; - return arr->mem.mem; - } - - const typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr) { - (void)ctx; - return arr->shape; - } - |] - -opaqueLibraryFunctions :: String -> [ValueDesc] - -> CompilerM op s [C.Definition] -opaqueLibraryFunctions desc vds = do - name <- publicName $ opaqueName desc vds - free_opaque <- publicName $ "free_" ++ opaqueName desc vds - - let opaque_type = [C.cty|struct $id:name|] - - freeComponent _ ScalarValue{} = - return () - freeComponent i (ArrayValue _ _ pt signed shape) = do - let rank = length shape - free_array <- publicName $ "free_" ++ arrayName pt signed rank - stm [C.cstm|if ((tmp = $id:free_array(ctx, obj->$id:(tupleField i))) != 0) { - ret = tmp; - }|] - - ctx_ty <- contextType - - free_body <- collect $ zipWithM_ freeComponent [0..] vds - - headerDecl (OpaqueDecl desc) - [C.cedecl|int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj);|] - - return [C.cunit| - int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj) { - int ret = 0, tmp; - $items:free_body - free(obj); - return ret; - } - |] - -valueDescToCType :: ValueDesc -> CompilerM op s C.Type -valueDescToCType (ScalarValue pt signed _) = - return $ signedPrimTypeToCType signed pt -valueDescToCType (ArrayValue mem space pt signed shape) = do - let pt' = signedPrimTypeToCType signed pt - rank = length shape - exists <- gets $ lookup (pt',rank) . compArrayStructs - case exists of - Just (cty, _) -> return cty - Nothing -> do - memty <- memToCType mem space - name <- publicName $ arrayName pt signed rank - let struct = [C.cedecl|struct $id:name { $ty:memty mem; typename int64_t shape[$int:rank]; };|] - stype = [C.cty|struct $id:name|] - library <- arrayLibraryFunctions space pt signed shape - modify $ \s -> s { compArrayStructs = - ((pt', rank), (stype, struct : library)) : compArrayStructs s - } - return stype - -opaqueToCType :: String -> [ValueDesc] -> CompilerM op s C.Type -opaqueToCType desc vds = do - name <- publicName $ opaqueName desc vds - exists <- gets $ lookup name . compOpaqueStructs - case exists of - Just (ty, _) -> return ty - Nothing -> do - members <- zipWithM field vds [(0::Int)..] - let struct = [C.cedecl|struct $id:name { $sdecls:members };|] - stype = [C.cty|struct $id:name|] - headerDecl (OpaqueDecl desc) [C.cedecl|struct $id:name;|] - library <- opaqueLibraryFunctions desc vds - modify $ \s -> s { compOpaqueStructs = - (name, (stype, struct : library)) : - compOpaqueStructs s } - return stype - where field vd@ScalarValue{} i = do - ct <- valueDescToCType vd - return [C.csdecl|$ty:ct $id:(tupleField i);|] - field vd i = do - ct <- valueDescToCType vd - return [C.csdecl|$ty:ct *$id:(tupleField i);|] - -externalValueToCType :: ExternalValue -> CompilerM op s C.Type -externalValueToCType (TransparentValue vd) = valueDescToCType vd -externalValueToCType (OpaqueValue desc vds) = opaqueToCType desc vds - -prepareEntryInputs :: [ExternalValue] -> CompilerM op s [C.Param] -prepareEntryInputs = zipWithM prepare [(0::Int)..] - where prepare pno (TransparentValue vd) = do - let pname = "in" ++ show pno - ty <- prepareValue [C.cexp|$id:pname|] vd - return [C.cparam|const $ty:ty $id:pname|] - - prepare pno (OpaqueValue desc vds) = do - ty <- opaqueToCType desc vds - let pname = "in" ++ show pno - field i ScalarValue{} = [C.cexp|$id:pname->$id:(tupleField i)|] - field i ArrayValue{} = [C.cexp|$id:pname->$id:(tupleField i)|] - zipWithM_ prepareValue (zipWith field [0..] vds) vds - return [C.cparam|const $ty:ty *$id:pname|] - - prepareValue src (ScalarValue pt signed name) = do - let pt' = signedPrimTypeToCType signed pt - stm [C.cstm|$id:name = $exp:src;|] - return pt' - - prepareValue src vd@(ArrayValue mem _ _ _ shape) = do - ty <- valueDescToCType vd - - stm [C.cstm|$exp:mem = $exp:src->mem;|] - - let rank = length shape - maybeCopyDim (Var d) i = - Just [C.cstm|$id:d = $exp:src->shape[$int:i];|] - maybeCopyDim _ _ = Nothing - - stms $ catMaybes $ zipWith maybeCopyDim shape [0..rank-1] - - return [C.cty|$ty:ty*|] - -prepareEntryOutputs :: [ExternalValue] -> CompilerM op s [C.Param] -prepareEntryOutputs = zipWithM prepare [(0::Int)..] - where prepare pno (TransparentValue vd) = do - let pname = "out" ++ show pno - ty <- valueDescToCType vd - - case vd of - ArrayValue{} -> do - stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|] - prepareValue [C.cexp|*$id:pname|] vd - return [C.cparam|$ty:ty **$id:pname|] - ScalarValue{} -> do - prepareValue [C.cexp|*$id:pname|] vd - return [C.cparam|$ty:ty *$id:pname|] - - prepare pno (OpaqueValue desc vds) = do - let pname = "out" ++ show pno - ty <- opaqueToCType desc vds - vd_ts <- mapM valueDescToCType vds - - stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|] - - - forM_ (zip3 [0..] vd_ts vds) $ \(i,ct,vd) -> do - let field = [C.cexp|(*$id:pname)->$id:(tupleField i)|] - case vd of - ScalarValue{} -> return () - _ -> stm [C.cstm|assert(($exp:field = ($ty:ct*) malloc(sizeof($ty:ct))) != NULL);|] - prepareValue field vd - - return [C.cparam|$ty:ty **$id:pname|] - - prepareValue dest (ScalarValue _ _ name) = - stm [C.cstm|$exp:dest = $id:name;|] - - prepareValue dest (ArrayValue mem _ _ _ shape) = do - stm [C.cstm|$exp:dest->mem = $id:mem;|] - - let rank = length shape - maybeCopyDim (Constant x) i = - [C.cstm|$exp:dest->shape[$int:i] = $exp:x;|] - maybeCopyDim (Var d) i = - [C.cstm|$exp:dest->shape[$int:i] = $id:d;|] - stms $ zipWith maybeCopyDim shape [0..rank-1] - -onEntryPoint :: Name -> Function op - -> CompilerM op s (C.Definition, C.Definition, C.Initializer) -onEntryPoint fname function@(Function _ outputs inputs _ results args) = do - let out_args = map (\p -> [C.cexp|&$id:(paramName p)|]) outputs - in_args = map (\p -> [C.cexp|$id:(paramName p)|]) inputs - - inputdecls <- collect $ mapM_ stubParam inputs - outputdecls <- collect $ mapM_ stubParam outputs - - let entry_point_name = nameToString fname - entry_point_function_name <- publicName $ "entry_" ++ entry_point_name - - (entry_point_input_params, unpack_entry_inputs) <- - collect' $ prepareEntryInputs args - (entry_point_output_params, pack_entry_outputs) <- - collect' $ prepareEntryOutputs results - - (cli_entry_point, cli_init) <- cliEntryPoint fname function - - ctx_ty <- contextType - - headerDecl EntryDecl [C.cedecl|int $id:entry_point_function_name - ($ty:ctx_ty *ctx, - $params:entry_point_output_params, - $params:entry_point_input_params);|] - - return ([C.cedecl|int $id:entry_point_function_name - ($ty:ctx_ty *ctx, - $params:entry_point_output_params, - $params:entry_point_input_params) { - $items:inputdecls - $items:outputdecls - - lock_lock(&ctx->lock); - - $items:unpack_entry_inputs - - int ret = $id:(funName fname)(ctx, $args:out_args, $args:in_args); - - if (ret == 0) { - $items:pack_entry_outputs - } - - lock_unlock(&ctx->lock); - - return ret; -} - |], - cli_entry_point, - cli_init) - where stubParam (MemParam name space) = - declMem name space - stubParam (ScalarParam name ty) = do - let ty' = primTypeToCType ty - decl [C.cdecl|$ty:ty' $id:name;|] - ---- CLI interface --- --- Our strategy for CLI entry points is to parse everything into --- host memory ('DefaultSpace') and copy the result into host memory --- after the entry point has returned. We have some ad-hoc frobbery --- to copy the host-level memory blocks to another memory space if --- necessary. This will break if the Futhark entry point uses --- non-trivial index functions for its input or output. --- --- The idea here is to keep the nastyness in the wrapper, whilst not --- messing up anything else. - -printPrimStm :: (C.ToExp a, C.ToExp b) => a -> b -> PrimType -> Signedness -> C.Stm -printPrimStm dest val bt ept = - [C.cstm|write_scalar($exp:dest, binary_output, &$exp:(primTypeInfo bt ept), &$exp:val);|] - --- | Return a statement printing the given external value. -printStm :: ExternalValue -> C.Exp -> CompilerM op s C.Stm -printStm (OpaqueValue desc _) _ = - return [C.cstm|printf("#", $string:desc);|] -printStm (TransparentValue (ScalarValue bt ept _)) e = - return $ printPrimStm [C.cexp|stdout|] e bt ept -printStm (TransparentValue (ArrayValue _ _ bt ept shape)) e = do - values_array <- publicName $ "values_" ++ name - shape_array <- publicName $ "shape_" ++ name - let num_elems = cproduct [ [C.cexp|$id:shape_array(ctx, $exp:e)[$int:i]|] | i <- [0..rank-1] ] - return [C.cstm|{ - $ty:bt' *arr = calloc(sizeof($ty:bt'), $exp:num_elems); - assert(arr != NULL); - assert($id:values_array(ctx, $exp:e, arr) == 0); - write_array(stdout, binary_output, &$exp:(primTypeInfo bt ept), arr, - $id:shape_array(ctx, $exp:e), $int:rank); - free(arr); - }|] - where rank = length shape - bt' = primTypeToCType bt - name = arrayName bt ept rank - -readPrimStm :: C.ToExp a => a -> Int -> PrimType -> Signedness -> C.Stm -readPrimStm place i t ept = - [C.cstm|if (read_scalar(&$exp:(primTypeInfo t ept),&$exp:place) != 0) { - futhark_panic(1, "Error when reading input #%d of type %s (errno: %s).\n", - $int:i, - $exp:(primTypeInfo t ept).type_name, - strerror(errno)); - }|] - -readInputs :: [ExternalValue] -> CompilerM op s [(C.Stm, C.Stm, C.Stm, C.Exp)] -readInputs = zipWithM readInput [0..] - -readInput :: Int -> ExternalValue -> CompilerM op s (C.Stm, C.Stm, C.Stm, C.Exp) -readInput i (OpaqueValue desc _) = do - stm [C.cstm|futhark_panic(1, "Cannot read input #%d of type %s\n", $int:i, $string:desc);|] - return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|NULL|]) -readInput i (TransparentValue (ScalarValue t ept _)) = do - dest <- newVName "read_value" - item [C.citem|$ty:(primTypeToCType t) $id:dest;|] - stm $ readPrimStm dest i t ept - return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|$id:dest|]) -readInput i (TransparentValue vd@(ArrayValue _ _ t ept dims)) = do - dest <- newVName "read_value" - shape <- newVName "read_shape" - arr <- newVName "read_arr" - ty <- valueDescToCType vd - item [C.citem|$ty:ty *$id:dest;|] - - let t' = signedPrimTypeToCType ept t - rank = length dims - name = arrayName t ept rank - dims_exps = [ [C.cexp|$id:shape[$int:j]|] | j <- [0..rank-1] ] - dims_s = concat $ replicate rank "[]" - - new_array <- publicName $ "new_" ++ name - free_array <- publicName $ "free_" ++ name - - items [C.citems| - typename int64_t $id:shape[$int:rank]; - $ty:t' *$id:arr = NULL; - errno = 0; - if (read_array(&$exp:(primTypeInfo t ept), - (void**) &$id:arr, - $id:shape, - $int:(length dims)) - != 0) { - futhark_panic(1, "Cannot read input #%d of type %s%s (errno: %s).\n", - $int:i, - $string:dims_s, - $exp:(primTypeInfo t ept).type_name, - strerror(errno)); - }|] - - return ([C.cstm|assert(($exp:dest = $id:new_array(ctx, $id:arr, $args:dims_exps)) != 0);|], - [C.cstm|assert($id:free_array(ctx, $exp:dest) == 0);|], - [C.cstm|free($id:arr);|], - [C.cexp|$id:dest|]) - -prepareOutputs :: [ExternalValue] -> CompilerM op s [(C.Exp, C.Stm)] -prepareOutputs = mapM prepareResult - where prepareResult ev = do - ty <- externalValueToCType ev - result <- newVName "result" - - case ev of - TransparentValue ScalarValue{} -> do - item [C.citem|$ty:ty $id:result;|] - return ([C.cexp|$id:result|], [C.cstm|;|]) - TransparentValue (ArrayValue _ _ t ept dims) -> do - let name = arrayName t ept $ length dims - free_array <- publicName $ "free_" ++ name - item [C.citem|$ty:ty *$id:result;|] - return ([C.cexp|$id:result|], - [C.cstm|assert($id:free_array(ctx, $exp:result) == 0);|]) - OpaqueValue desc vds -> do - free_opaque <- publicName $ "free_" ++ opaqueName desc vds - item [C.citem|$ty:ty *$id:result;|] - return ([C.cexp|$id:result|], - [C.cstm|assert($id:free_opaque(ctx, $exp:result) == 0);|]) - -printResult :: [(ExternalValue,C.Exp)] -> CompilerM op s [C.Stm] -printResult vs = fmap concat $ forM vs $ \(v,e) -> do - p <- printStm v e - return [p, [C.cstm|printf("\n");|]] - -cliEntryPoint :: Name - -> FunctionT a - -> CompilerM op s (C.Definition, C.Initializer) -cliEntryPoint fname (Function _ _ _ _ results args) = do - ((pack_input, free_input, free_parsed, input_args), input_items) <- - collect' $ unzip4 <$> readInputs args - - ((output_vals, free_outputs), output_decls) <- - collect' $ unzip <$> prepareOutputs results - printstms <- printResult $ zip results output_vals - - ctx_ty <- contextType - sync_ctx <- publicName "context_sync" - error_ctx <- publicName "context_get_error" - - let entry_point_name = nameToString fname - cli_entry_point_function_name = "futrts_cli_entry_" ++ entry_point_name - entry_point_function_name <- publicName $ "entry_" ++ entry_point_name - - pause_profiling <- publicName "context_pause_profiling" - unpause_profiling <- publicName "context_unpause_profiling" - - let run_it = [C.citems| - int r; - // Run the program once. - $stms:pack_input - if ($id:sync_ctx(ctx) != 0) { - futhark_panic(1, "%s", $id:error_ctx(ctx)); - }; - // Only profile last run. - if (profile_run) { - $id:unpause_profiling(ctx); - } - t_start = get_wall_time(); - r = $id:entry_point_function_name(ctx, - $args:(map addrOf output_vals), - $args:input_args); - if (r != 0) { - futhark_panic(1, "%s", $id:error_ctx(ctx)); - } - if ($id:sync_ctx(ctx) != 0) { - futhark_panic(1, "%s", $id:error_ctx(ctx)); - }; - if (profile_run) { - $id:pause_profiling(ctx); - } - t_end = get_wall_time(); - long int elapsed_usec = t_end - t_start; - if (time_runs && runtime_file != NULL) { - fprintf(runtime_file, "%lld\n", (long long) elapsed_usec); - fflush(runtime_file); - } - $stms:free_input - |] - - return ([C.cedecl|static void $id:cli_entry_point_function_name($ty:ctx_ty *ctx) { - typename int64_t t_start, t_end; - int time_runs = 0, profile_run = 0; - - // We do not want to profile all the initialisation. - $id:pause_profiling(ctx); - - // Declare and read input. - set_binary_mode(stdin); - $items:input_items - - if (end_of_input() != 0) { - futhark_panic(1, "Expected EOF on stdin after reading input for %s.\n", $string:(quote (pretty fname))); - } - - $items:output_decls - - // Warmup run - if (perform_warmup) { - $items:run_it - $stms:free_outputs - } - time_runs = 1; - // Proper run. - for (int run = 0; run < num_runs; run++) { - // Only profile last run. - profile_run = run == num_runs -1; - $items:run_it - if (run < num_runs-1) { - $stms:free_outputs - } - } - - // Free the parsed input. - $stms:free_parsed - - // Print the final result. - if (binary_output) { - set_binary_mode(stdout); - } - $stms:printstms - - $stms:free_outputs - } - |], - [C.cinit|{ .name = $string:entry_point_name, - .fun = $id:cli_entry_point_function_name }|] - ) - -benchmarkOptions :: [Option] -benchmarkOptions = - [ Option { optionLongName = "write-runtime-to" - , optionShortName = Just 't' - , optionArgument = RequiredArgument "FILE" - , optionAction = set_runtime_file - } - , Option { optionLongName = "runs" - , optionShortName = Just 'r' - , optionArgument = RequiredArgument "INT" - , optionAction = set_num_runs - } - , Option { optionLongName = "debugging" - , optionShortName = Just 'D' - , optionArgument = NoArgument - , optionAction = [C.cstm|futhark_context_config_set_debugging(cfg, 1);|] - } - , Option { optionLongName = "log" - , optionShortName = Just 'L' - , optionArgument = NoArgument - , optionAction = [C.cstm|futhark_context_config_set_logging(cfg, 1);|] - } - , Option { optionLongName = "entry-point" - , optionShortName = Just 'e' - , optionArgument = RequiredArgument "NAME" - , optionAction = [C.cstm|if (entry_point != NULL) entry_point = optarg;|] - } - , Option { optionLongName = "binary-output" - , optionShortName = Just 'b' - , optionArgument = NoArgument - , optionAction = [C.cstm|binary_output = 1;|] - } - ] - where set_runtime_file = [C.cstm|{ - runtime_file = fopen(optarg, "w"); - if (runtime_file == NULL) { - futhark_panic(1, "Cannot open %s: %s\n", optarg, strerror(errno)); - } - }|] - set_num_runs = [C.cstm|{ - num_runs = atoi(optarg); - perform_warmup = 1; - if (num_runs <= 0) { - futhark_panic(1, "Need a positive number of runs, not %s\n", optarg); - } - }|] - --- | The result of compilation to C is four parts, which can be put --- together in various ways. The obvious way is to concatenate all of --- them, which yields a CLI program. Another is to compile the --- library part by itself, and use the header file to call into it. -data CParts = CParts { cHeader :: String - , cUtils :: String - -- ^ Utility definitions that must be visible - -- to both CLI and library parts. - , cCLI :: String - , cLib :: String - } - --- We may generate variables that are never used (e.g. for --- certificates) or functions that are never called (e.g. unused --- intrinsics), and generated code may have other cosmetic issues that --- compilers warn about. We disable these warnings to not clutter the --- compilation logs. -disableWarnings :: String -disableWarnings = pretty [C.cunit| -$esc:("#ifdef __GNUC__") -$esc:("#pragma GCC diagnostic ignored \"-Wunused-function\"") -$esc:("#pragma GCC diagnostic ignored \"-Wunused-variable\"") -$esc:("#pragma GCC diagnostic ignored \"-Wparentheses\"") -$esc:("#pragma GCC diagnostic ignored \"-Wunused-label\"") -$esc:("#endif") - -$esc:("#ifdef __clang__") -$esc:("#pragma clang diagnostic ignored \"-Wunused-function\"") -$esc:("#pragma clang diagnostic ignored \"-Wunused-variable\"") -$esc:("#pragma clang diagnostic ignored \"-Wparentheses\"") -$esc:("#pragma clang diagnostic ignored \"-Wunused-label\"") -$esc:("#endif") -|] - --- | Produce header and implementation files. -asLibrary :: CParts -> (String, String) -asLibrary parts = ("#pragma once\n\n" <> cHeader parts, - disableWarnings <> cHeader parts <> cUtils parts <> cLib parts) - --- | As executable with command-line interface. -asExecutable :: CParts -> String -asExecutable (CParts a b c d) = disableWarnings <> a <> b <> c <> d - --- | Compile imperative program to a C program. Always uses the --- function named "main" as entry point, so make sure it is defined. -compileProg :: MonadFreshNames m => - String - -> Operations op () - -> CompilerM op () () - -> String - -> [Space] - -> [Option] - -> Definitions op - -> m CParts -compileProg backend ops extra header_extra spaces options prog = do - src <- getNameSource - let ((prototypes, definitions, entry_points), endstate) = - runCompilerM ops src () compileProg' - (entry_point_decls, cli_entry_point_decls, entry_point_inits) = - unzip3 entry_points - option_parser = generateOptionParser "parse_options" $ benchmarkOptions++options - - let headerdefs = [C.cunit| -$esc:("// Headers\n") -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") -$esc:(header_extra) - -$esc:("\n// Initialisation\n") -$edecls:(initDecls endstate) - -$esc:("\n// Arrays\n") -$edecls:(arrayDecls endstate) - -$esc:("\n// Opaque values\n") -$edecls:(opaqueDecls endstate) - -$esc:("\n// Entry points\n") -$edecls:(entryDecls endstate) - -$esc:("\n// Miscellaneous\n") -$edecls:(miscDecls endstate) -$esc:("#define FUTHARK_BACKEND_"++backend) - |] - - let utildefs = [C.cunit| -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") -// If NDEBUG is set, the assert() macro will do nothing. Since Futhark -// (unfortunately) makes use of assert() for error detection (and even some -// side effects), we want to avoid that. -$esc:("#undef NDEBUG") -$esc:("#include ") -$esc:("#include ") - -$esc:util_h - -$esc:timing_h -|] - - let clidefs = [C.cunit| -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") - -$esc:values_h - -$esc:("#define __private") - -static int binary_output = 0; -static typename FILE *runtime_file; -static int perform_warmup = 0; -static int num_runs = 1; -// If the entry point is NULL, the program will terminate after doing initialisation and such. -static const char *entry_point = "main"; - -$esc:tuning_h - -$func:option_parser - -$edecls:cli_entry_point_decls - -typedef void entry_point_fun(struct futhark_context*); - -struct entry_point_entry { - const char *name; - entry_point_fun *fun; -}; - -int main(int argc, char** argv) { - fut_progname = argv[0]; - - struct entry_point_entry entry_points[] = { - $inits:entry_point_inits - }; - - struct futhark_context_config *cfg = futhark_context_config_new(); - assert(cfg != NULL); - - int parsed_options = parse_options(cfg, argc, argv); - argc -= parsed_options; - argv += parsed_options; - - if (argc != 0) { - futhark_panic(1, "Excess non-option: %s\n", argv[0]); - } - - struct futhark_context *ctx = futhark_context_new(cfg); - assert (ctx != NULL); - - char* error = futhark_context_get_error(ctx); - if (error != NULL) { - futhark_panic(1, "%s", error); - } - - if (entry_point != NULL) { - int num_entry_points = sizeof(entry_points) / sizeof(entry_points[0]); - entry_point_fun *entry_point_fun = NULL; - for (int i = 0; i < num_entry_points; i++) { - if (strcmp(entry_points[i].name, entry_point) == 0) { - entry_point_fun = entry_points[i].fun; - break; - } - } - - if (entry_point_fun == NULL) { - fprintf(stderr, "No entry point '%s'. Select another with --entry-point. Options are:\n", - entry_point); - for (int i = 0; i < num_entry_points; i++) { - fprintf(stderr, "%s\n", entry_points[i].name); - } - return 1; - } - - entry_point_fun(ctx); - - if (runtime_file != NULL) { - fclose(runtime_file); - } - - char *report = futhark_context_report(ctx); - fputs(report, stderr); - free(report); - } - - futhark_context_free(ctx); - futhark_context_config_free(cfg); - return 0; -} - |] - - let early_decls = DL.toList $ compEarlyDecls endstate - let lib_decls = DL.toList $ compLibDecls endstate - let libdefs = [C.cunit| -$esc:("#ifdef _MSC_VER\n#define inline __inline\n#endif") -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") -$esc:("#include ") - -$esc:(header_extra) - -$esc:lock_h - -$edecls:builtin - -$edecls:early_decls - -$edecls:prototypes - -$edecls:lib_decls - -$edecls:(map funcToDef definitions) - -$edecls:(arrayDefinitions endstate) - -$edecls:(opaqueDefinitions endstate) - -$edecls:entry_point_decls - |] - - return $ CParts (pretty headerdefs) (pretty utildefs) (pretty clidefs) (pretty libdefs) - where - compileProg' = do - let Definitions consts (Functions funs) = prog - - (memstructs, memfuns, memreport) <- unzip3 <$> mapM defineMemorySpace spaces - - get_consts <- compileConstants consts - - ctx_ty <- contextType - - (prototypes, definitions) <- - unzip <$> mapM (compileFun get_consts [[C.cparam|$ty:ctx_ty *ctx|]]) funs - - mapM_ earlyDecl memstructs - entry_points <- - mapM (uncurry onEntryPoint) $ filter (functionEntry . snd) funs - - extra - - mapM_ earlyDecl $ concat memfuns - - commonLibFuns memreport - - return (prototypes, definitions, entry_points) - - funcToDef func = C.FuncDef func loc - where loc = case func of - C.OldFunc _ _ _ _ _ _ l -> l - C.Func _ _ _ _ _ l -> l - - builtin = cIntOps ++ cFloat32Ops ++ cFloat64Ops ++ cFloatConvOps ++ - cFloat32Funs ++ cFloat64Funs - - util_h = $(embedStringFile "rts/c/util.h") - values_h = $(embedStringFile "rts/c/values.h") - timing_h = $(embedStringFile "rts/c/timing.h") - lock_h = $(embedStringFile "rts/c/lock.h") - tuning_h = $(embedStringFile "rts/c/tuning.h") - -commonLibFuns :: [C.BlockItem] -> CompilerM op s () -commonLibFuns memreport = do - ctx <- contextType - profilereport <- gets $ DL.toList . compProfileItems - - publicDef_ "context_report" MiscDecl $ \s -> - ([C.cedecl|char* $id:s($ty:ctx *ctx);|], - [C.cedecl|char* $id:s($ty:ctx *ctx) { - struct str_builder builder; - str_builder_init(&builder); - if (ctx->detail_memory || ctx->profiling) { - $items:memreport - } - if (ctx->profiling) { - $items:profilereport - } - return builder.str; - }|]) - - publicDef_ "context_get_error" MiscDecl $ \s -> - ([C.cedecl|char* $id:s($ty:ctx* ctx);|], - [C.cedecl|char* $id:s($ty:ctx* ctx) { - char* error = ctx->error; - ctx->error = NULL; - return error; - }|]) - - publicDef_ "context_pause_profiling" MiscDecl $ \s -> - ([C.cedecl|void $id:s($ty:ctx* ctx);|], - [C.cedecl|void $id:s($ty:ctx* ctx) { - ctx->profiling_paused = 1; - }|]) - - publicDef_ "context_unpause_profiling" MiscDecl $ \s -> - ([C.cedecl|void $id:s($ty:ctx* ctx);|], - [C.cedecl|void $id:s($ty:ctx* ctx) { - ctx->profiling_paused = 0; - }|]) - -compileConstants :: Constants op -> CompilerM op s [C.BlockItem] -compileConstants (Constants ps init_consts) = do - ctx_ty <- contextType - const_fields <- mapM constParamField ps - -- Avoid an empty struct, as that is apparently undefined behaviour. - let const_fields' | null const_fields = [[C.csdecl|int dummy;|]] - | otherwise = const_fields - contextField "constants" [C.cty|struct { $sdecls:const_fields' }|] Nothing - earlyDecl [C.cedecl|int init_constants($ty:ctx_ty*);|] - earlyDecl [C.cedecl|int free_constants($ty:ctx_ty*);|] - - -- We locally define macros for the constants, so that when we - -- generate assignments to local variables, we actually assign into - -- the constants struct. This is not needed for functions, because - -- they can only read constants, not write them. - let (defs, undefs) = unzip $ map constMacro ps - init_consts' <- blockScope $ do - mapM_ resetMemConst ps - compileCode init_consts - libDecl [C.cedecl|int init_constants($ty:ctx_ty *ctx) { - (void)ctx; - int err = 0; - $items:defs - $items:init_consts' - $items:undefs - cleanup: - return err; - }|] - - free_consts <- collect $ mapM_ freeConst ps - libDecl [C.cedecl|int free_constants($ty:ctx_ty *ctx) { - (void)ctx; - $items:free_consts - return 0; - }|] - - mapM getConst ps - - where constParamField (ScalarParam name bt) = do - let ctp = primTypeToCType bt - return [C.csdecl|$ty:ctp $id:name;|] - constParamField (MemParam name space) = do - ty <- memToCType name space - return [C.csdecl|$ty:ty $id:name;|] - - constMacro p = ([C.citem|$escstm:def|], [C.citem|$escstm:undef|]) - where p' = pretty (C.toIdent (paramName p) mempty) - def = "#define " ++ p' ++ " (" ++ "ctx->constants." ++ p' ++ ")" - undef = "#undef " ++ p' - - resetMemConst ScalarParam{} = return () - resetMemConst (MemParam name space) = resetMem name space - - freeConst ScalarParam{} = return () - freeConst (MemParam name space) = unRefMem [C.cexp|ctx->constants.$id:name|] space - - getConst (ScalarParam name bt) = do - let ctp = primTypeToCType bt - return [C.citem|$ty:ctp $id:name = ctx->constants.$id:name;|] - getConst (MemParam name space) = do - ty <- memToCType name space - return [C.citem|$ty:ty $id:name = ctx->constants.$id:name;|] - -cachingMemory :: M.Map VName Space - -> ([C.BlockItem] -> [C.Stm] -> CompilerM op s a) - -> CompilerM op s a -cachingMemory lexical f = do - -- We only consider lexical 'DefaultSpace' memory blocks to be - -- cached. This is not a deep technical restriction, but merely a - -- heuristic based on GPU memory usually involving larger - -- allocations, that do not suffer from the overhead of reference - -- counting. - let cached = M.keys $ M.filter (==DefaultSpace) lexical - - cached' <- forM cached $ \mem -> do - size <- newVName $ pretty mem <> "_cached_size" - return (mem, size) - - let lexMem env = - env { envCachedMem = - M.fromList (map (first (`C.toExp` noLoc)) cached') - <> envCachedMem env - } - - declCached (mem, size) = - [[C.citem|size_t $id:size = 0;|], - [C.citem|$ty:defaultMemBlockType $id:mem = NULL;|]] - - freeCached (mem, _) = - [C.cstm|free($id:mem);|] - - local lexMem $ f (concatMap declCached cached') (map freeCached cached') - -compileFun :: [C.BlockItem] -> [C.Param] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func) -compileFun get_constants extra (fname, func@(Function _ outputs inputs body _ _)) = do - (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs - inparams <- mapM compileInput inputs - - cachingMemory (lexicalMemoryUsage func) $ \decl_cached free_cached -> do - body' <- blockScope $ compileFunBody out_ptrs outputs body - - return ([C.cedecl|static int $id:(funName fname)($params:extra, $params:outparams, $params:inparams);|], - [C.cfun|static int $id:(funName fname)($params:extra, $params:outparams, $params:inparams) { - int err = 0; - $items:decl_cached - $items:get_constants - $items:body' - cleanup: - {} - $stms:free_cached - return err; - }|]) - - where compileInput (ScalarParam name bt) = do - let ctp = primTypeToCType bt - return [C.cparam|$ty:ctp $id:name|] - compileInput (MemParam name space) = do - ty <- memToCType name space - return [C.cparam|$ty:ty $id:name|] - - compileOutput (ScalarParam name bt) = do - let ctp = primTypeToCType bt - p_name <- newVName $ "out_" ++ baseString name - return ([C.cparam|$ty:ctp *$id:p_name|], [C.cexp|$id:p_name|]) - compileOutput (MemParam name space) = do - ty <- memToCType name space - p_name <- newVName $ baseString name ++ "_p" - return ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|]) - -compilePrimValue :: PrimValue -> C.Exp - -compilePrimValue (IntValue (Int8Value k)) = [C.cexp|$int:k|] -compilePrimValue (IntValue (Int16Value k)) = [C.cexp|$int:k|] -compilePrimValue (IntValue (Int32Value k)) = [C.cexp|$int:k|] -compilePrimValue (IntValue (Int64Value k)) = [C.cexp|$int:k|] - -compilePrimValue (FloatValue (Float64Value x)) - | isInfinite x = - if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|] - | isNaN x = - [C.cexp|NAN|] - | otherwise = - [C.cexp|$double:x|] -compilePrimValue (FloatValue (Float32Value x)) - | isInfinite x = - if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|] - | isNaN x = - [C.cexp|NAN|] - | otherwise = - [C.cexp|$float:x|] - -compilePrimValue (BoolValue b) = - [C.cexp|$int:b'|] - where b' :: Int - b' = if b then 1 else 0 - -compilePrimValue Checked = - [C.cexp|0|] - -derefPointer :: C.Exp -> C.Exp -> C.Type -> C.Exp -derefPointer ptr i res_t = - [C.cexp|(($ty:res_t)$exp:ptr)[$exp:i]|] - -volQuals :: Volatility -> [C.TypeQual] -volQuals Volatile = [C.ctyquals|volatile|] -volQuals Nonvolatile = [] - -writeScalarPointerWithQuals :: PointerQuals op s -> WriteScalar op s -writeScalarPointerWithQuals quals_f dest i elemtype space vol v = do - quals <- quals_f space - let quals' = volQuals vol ++ quals - deref = derefPointer dest i - [C.cty|$tyquals:quals' $ty:elemtype*|] - stm [C.cstm|$exp:deref = $exp:v;|] - -readScalarPointerWithQuals :: PointerQuals op s -> ReadScalar op s -readScalarPointerWithQuals quals_f dest i elemtype space vol = do - quals <- quals_f space - let quals' = volQuals vol ++ quals - return $ derefPointer dest i [C.cty|$tyquals:quals' $ty:elemtype*|] - -compileExpToName :: String -> PrimType -> Exp -> CompilerM op s VName -compileExpToName _ _ (LeafExp (ScalarVar v) _) = - return v -compileExpToName desc t e = do - desc' <- newVName desc - e' <- compileExp e - decl [C.cdecl|$ty:(primTypeToCType t) $id:desc' = $e';|] - return desc' - -compileExp :: Exp -> CompilerM op s C.Exp - -compileExp = compilePrimExp compileLeaf - where compileLeaf (ScalarVar src) = - return [C.cexp|$id:src|] - - compileLeaf (Index src (Count iexp) restype DefaultSpace vol) = do - src' <- rawMem src - derefPointer src' - <$> compileExp iexp - <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primTypeToCType restype)*|] - - compileLeaf (Index src (Count iexp) restype (Space space) vol) = - join $ asks envReadScalar - <*> rawMem src <*> compileExp iexp - <*> pure (primTypeToCType restype) <*> pure space <*> pure vol - - compileLeaf (Index src (Count iexp) _ ScalarSpace{} _) = do - iexp' <- compileExp iexp - return [C.cexp|$id:src[$exp:iexp']|] - - compileLeaf (SizeOf t) = - return [C.cexp|(typename int32_t)sizeof($ty:t')|] - where t' = primTypeToCType t - --- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you. -compilePrimExp :: Monad m => (v -> m C.Exp) -> PrimExp v -> m C.Exp - -compilePrimExp _ (ValueExp val) = - return $ compilePrimValue val - -compilePrimExp f (LeafExp v _) = - f v - -compilePrimExp f (UnOpExp Complement{} x) = do - x' <- compilePrimExp f x - return [C.cexp|~$exp:x'|] - -compilePrimExp f (UnOpExp Not{} x) = do - x' <- compilePrimExp f x - return [C.cexp|!$exp:x'|] - -compilePrimExp f (UnOpExp Abs{} x) = do - x' <- compilePrimExp f x - return [C.cexp|abs($exp:x')|] - -compilePrimExp f (UnOpExp (FAbs Float32) x) = do - x' <- compilePrimExp f x - return [C.cexp|(float)fabs($exp:x')|] - -compilePrimExp f (UnOpExp (FAbs Float64) x) = do - x' <- compilePrimExp f x - return [C.cexp|fabs($exp:x')|] - -compilePrimExp f (UnOpExp SSignum{} x) = do - x' <- compilePrimExp f x - return [C.cexp|($exp:x' > 0) - ($exp:x' < 0)|] - -compilePrimExp f (UnOpExp USignum{} x) = do - x' <- compilePrimExp f x - return [C.cexp|($exp:x' > 0) - ($exp:x' < 0) != 0|] - -compilePrimExp f (CmpOpExp cmp x y) = do - x' <- compilePrimExp f x - y' <- compilePrimExp f y - return $ case cmp of - CmpEq{} -> [C.cexp|$exp:x' == $exp:y'|] - - FCmpLt{} -> [C.cexp|$exp:x' < $exp:y'|] - FCmpLe{} -> [C.cexp|$exp:x' <= $exp:y'|] - - CmpLlt{} -> [C.cexp|$exp:x' < $exp:y'|] - CmpLle{} -> [C.cexp|$exp:x' <= $exp:y'|] - - _ -> [C.cexp|$id:(pretty cmp)($exp:x', $exp:y')|] - -compilePrimExp f (ConvOpExp conv x) = do - x' <- compilePrimExp f x - return [C.cexp|$id:(pretty conv)($exp:x')|] - -compilePrimExp f (BinOpExp bop x y) = do - x' <- compilePrimExp f x - y' <- compilePrimExp f y - -- Note that integer addition, subtraction, and multiplication with - -- OverflowWrap are not handled by explicit operators, but rather by - -- functions. This is because we want to implicitly convert them to - -- unsigned numbers, so we can do overflow without invoking - -- undefined behaviour. - return $ case bop of - Add _ OverflowUndef -> [C.cexp|$exp:x' + $exp:y'|] - Sub _ OverflowUndef -> [C.cexp|$exp:x' - $exp:y'|] - Mul _ OverflowUndef -> [C.cexp|$exp:x' * $exp:y'|] - FAdd{} -> [C.cexp|$exp:x' + $exp:y'|] - FSub{} -> [C.cexp|$exp:x' - $exp:y'|] - FMul{} -> [C.cexp|$exp:x' * $exp:y'|] - FDiv{} -> [C.cexp|$exp:x' / $exp:y'|] - Xor{} -> [C.cexp|$exp:x' ^ $exp:y'|] - And{} -> [C.cexp|$exp:x' & $exp:y'|] - Or{} -> [C.cexp|$exp:x' | $exp:y'|] - Shl{} -> [C.cexp|$exp:x' << $exp:y'|] - LogAnd{} -> [C.cexp|$exp:x' && $exp:y'|] - LogOr{} -> [C.cexp|$exp:x' || $exp:y'|] - _ -> [C.cexp|$id:(pretty bop)($exp:x', $exp:y')|] - -compilePrimExp f (FunExp h args _) = do - args' <- mapM (compilePrimExp f) args - return [C.cexp|$id:(funName (nameFromString h))($args:args')|] - -compileCode :: Code op -> CompilerM op s () - -compileCode (Op op) = - join $ asks envOpCompiler <*> pure op - -compileCode Skip = return () - -compileCode (Comment s code) = do - xs <- blockScope $ compileCode code - let comment = "// " ++ s - stm [C.cstm|$comment:comment - { $items:xs } - |] - -compileCode (DebugPrint s (Just e)) = do - e' <- compileExp e - stm [C.cstm|if (ctx->debugging) { - fprintf(stderr, $string:fmtstr, $exp:s, ($ty:ety)$exp:e', '\n'); - }|] - where (fmt, ety) = case primExpType e of - IntType _ -> ("llu", [C.cty|long long int|]) - FloatType _ -> ("f", [C.cty|double|]) - _ -> ("d", [C.cty|int|]) - fmtstr = "%s: %" ++ fmt ++ "%c" - -compileCode (DebugPrint s Nothing) = - stm [C.cstm|if (ctx->debugging) { - fprintf(stderr, "%s\n", $exp:s); - }|] - -compileCode c - | Just (name, vol, t, e, c') <- declareAndSet c = do - let ct = primTypeToCType t - e' <- compileExp e - item [C.citem|$tyquals:(volQuals vol) $ty:ct $id:name = $exp:e';|] - compileCode c' - -compileCode (c1 :>>: c2) = compileCode c1 >> compileCode c2 - -compileCode (Assert e msg (loc, locs)) = do - e' <- compileExp e - err <- collect $ join $ - asks (opsError . envOperations) <*> pure msg <*> pure stacktrace - stm [C.cstm|if (!$exp:e') { $items:err }|] - where stacktrace = prettyStacktrace 0 $ map locStr $ loc:locs - -compileCode (Allocate _ _ ScalarSpace{}) = - -- Handled by the declaration of the memory block, which is - -- translated to an actual array. - return () - -compileCode (Allocate name (Count e) space) = do - size <- compileExp e - cached <- cacheMem name - case cached of - Just cur_size -> - stm [C.cstm|if ($exp:cur_size < (size_t)$exp:size) { - $exp:name = realloc($exp:name, $exp:size); - $exp:cur_size = $exp:size; - }|] - _ -> - allocMem name size space [C.cstm|{err = 1; goto cleanup;}|] - -compileCode (Free name space) = do - cached <- isJust <$> cacheMem name - unless cached $ unRefMem name space - -compileCode (For i it bound body) = do - let i' = C.toIdent i - it' = primTypeToCType $ IntType it - bound' <- compileExp bound - body' <- blockScope $ compileCode body - stm [C.cstm|for ($ty:it' $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) { - $items:body' - }|] - -compileCode (While cond body) = do - cond' <- compileExp cond - body' <- blockScope $ compileCode body - stm [C.cstm|while ($exp:cond') { - $items:body' - }|] - -compileCode (If cond tbranch fbranch) = do - cond' <- compileExp cond - tbranch' <- blockScope $ compileCode tbranch - fbranch' <- blockScope $ compileCode fbranch - stm $ case (tbranch', fbranch') of - (_, []) -> - [C.cstm|if ($exp:cond') { $items:tbranch' }|] - ([], _) -> - [C.cstm|if (!($exp:cond')) { $items:fbranch' }|] - _ -> - [C.cstm|if ($exp:cond') { $items:tbranch' } else { $items:fbranch' }|] - -compileCode (Copy dest (Count destoffset) DefaultSpace src (Count srcoffset) DefaultSpace (Count size)) = - join $ copyMemoryDefaultSpace - <$> rawMem dest <*> compileExp destoffset - <*> rawMem src <*> compileExp srcoffset - <*> compileExp size - -compileCode (Copy dest (Count destoffset) destspace src (Count srcoffset) srcspace (Count size)) = do - copy <- asks envCopy - join $ copy - <$> rawMem dest <*> compileExp destoffset <*> pure destspace - <*> rawMem src <*> compileExp srcoffset <*> pure srcspace - <*> compileExp size - -compileCode (Write dest (Count idx) elemtype DefaultSpace vol elemexp) = do - dest' <- rawMem dest - deref <- derefPointer dest' - <$> compileExp idx - <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primTypeToCType elemtype)*|] - elemexp' <- compileExp elemexp - stm [C.cstm|$exp:deref = $exp:elemexp';|] - -compileCode (Write dest (Count idx) _ ScalarSpace{} _ elemexp) = do - idx' <- compileExp idx - elemexp' <- compileExp elemexp - stm [C.cstm|$id:dest[$exp:idx'] = $exp:elemexp';|] - -compileCode (Write dest (Count idx) elemtype (Space space) vol elemexp) = - join $ asks envWriteScalar - <*> rawMem dest - <*> compileExp idx - <*> pure (primTypeToCType elemtype) - <*> pure space - <*> pure vol - <*> compileExp elemexp - -compileCode (DeclareMem name space) = - declMem name space - -compileCode (DeclareScalar name vol t) = do - let ct = primTypeToCType t - decl [C.cdecl|$tyquals:(volQuals vol) $ty:ct $id:name;|] - -compileCode (DeclareArray name ScalarSpace{} _ _) = - error $ "Cannot declare array " ++ pretty name ++ " in scalar space." - -compileCode (DeclareArray name DefaultSpace t vs) = do - name_realtype <- newVName $ baseString name ++ "_realtype" - let ct = primTypeToCType t - case vs of - ArrayValues vs' -> do - let vs'' = [[C.cinit|$exp:(compilePrimValue v)|] | v <- vs'] - earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:(length vs')] = {$inits:vs''};|] - ArrayZeros n -> - earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:n];|] - -- Fake a memory block. - contextField (C.toIdent name noLoc) - [C.cty|struct memblock|] $ - Just [C.cexp|(struct memblock){NULL, (char*)$id:name_realtype, 0}|] - item [C.citem|struct memblock $id:name = ctx->$id:name;|] - -compileCode (DeclareArray name (Space space) t vs) = - join $ asks envStaticArray <*> - pure name <*> pure space <*> pure t <*> pure vs - --- For assignments of the form 'x = x OP e', we generate C assignment --- operators to make the resulting code slightly nicer. This has no --- effect on performance. -compileCode (SetScalar dest (BinOpExp op (LeafExp (ScalarVar x) _) y)) - | dest == x, Just f <- assignmentOperator op = do - y' <- compileExp y - stm [C.cstm|$exp:(f dest y');|] - -compileCode (SetScalar dest src) = do - src' <- compileExp src - stm [C.cstm|$id:dest = $exp:src';|] - -compileCode (SetMem dest src space) = - setMem dest src space - -compileCode (Call dests fname args) = - join $ asks (opsCall . envOperations) - <*> pure dests <*> pure fname <*> mapM compileArg args - where compileArg (MemArg m) = return [C.cexp|$exp:m|] - compileArg (ExpArg e) = compileExp e - -blockScope :: CompilerM op s () -> CompilerM op s [C.BlockItem] -blockScope = fmap snd . blockScope' - -blockScope' :: CompilerM op s a -> CompilerM op s (a, [C.BlockItem]) -blockScope' m = do - old_allocs <- gets compDeclaredMem - (x, xs) <- pass $ do - (x, w) <- listen m - let xs = DL.toList $ accItems w - return ((x, xs), const mempty) - new_allocs <- gets $ filter (`notElem` old_allocs) . compDeclaredMem - modify $ \s -> s { compDeclaredMem = old_allocs } - releases <- collect $ mapM_ (uncurry unRefMem) new_allocs - return (x, xs <> releases) - -compileFunBody :: [C.Exp] -> [Param] -> Code op -> CompilerM op s () -compileFunBody output_ptrs outputs code = do - mapM_ declareOutput outputs - compileCode code - zipWithM_ setRetVal' output_ptrs outputs - where declareOutput (MemParam name space) = - declMem name space - declareOutput (ScalarParam name pt) = do - let ctp = primTypeToCType pt - decl [C.cdecl|$ty:ctp $id:name;|] - - setRetVal' p (MemParam name space) = do - resetMem [C.cexp|*$exp:p|] space - setMem [C.cexp|*$exp:p|] name space - setRetVal' p (ScalarParam name _) = - stm [C.cstm|*$exp:p = $id:name;|] - -declareAndSet :: Code op -> Maybe (VName, Volatility, PrimType, Exp, Code op) -declareAndSet code = do - (DeclareScalar name vol t, code') <- nextCode code - (SetScalar dest e, code'') <- nextCode code' - guard $ name == dest - Just (name, vol, t, e, code'') - -nextCode :: Code op -> Maybe (Code op, Code op) -nextCode (x :>>: y) - | Just (x_a, x_b) <- nextCode x = - Just (x_a, x_b <> y) - | otherwise = - Just (x, y) -nextCode _ = Nothing - -assignmentOperator :: BinOp -> Maybe (VName -> C.Exp -> C.Exp) -assignmentOperator Add{} = Just $ \d e -> [C.cexp|$id:d += $exp:e|] -assignmentOperator Sub{} = Just $ \d e -> [C.cexp|$id:d -= $exp:e|] -assignmentOperator Mul{} = Just $ \d e -> [C.cexp|$id:d *= $exp:e|] -assignmentOperator _ = Nothing - --- | Return an expression multiplying together the given expressions. --- If an empty list is given, the expression @1@ is returned. -cproduct :: [C.Exp] -> C.Exp -cproduct [] = [C.cexp|1|] -cproduct (e:es) = foldl mult e es - where mult x y = [C.cexp|$exp:x * $exp:y|] +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TupleSections #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +-- | C code generator framework. +module Futhark.CodeGen.Backends.GenericC + ( compileProg, + CParts (..), + asLibrary, + asExecutable, + + -- * Pluggable compiler + Operations (..), + defaultOperations, + OpCompiler, + ErrorCompiler, + CallCompiler, + PointerQuals, + MemoryType, + WriteScalar, + writeScalarPointerWithQuals, + ReadScalar, + readScalarPointerWithQuals, + Allocate, + Deallocate, + Copy, + StaticArray, + + -- * Monadic compiler interface + CompilerM, + CompilerState (compUserState), + getUserState, + modifyUserState, + contextContents, + contextFinalInits, + runCompilerM, + cachingMemory, + blockScope, + compileFun, + compileCode, + compileExp, + compilePrimExp, + compilePrimValue, + compileExpToName, + rawMem, + item, + items, + stm, + stms, + decl, + atInit, + headerDecl, + publicDef, + publicDef_, + profileReport, + HeaderSection (..), + libDecl, + earlyDecl, + publicName, + contextType, + contextField, + + -- * Building Blocks + primTypeToCType, + copyMemoryDefaultSpace, + ) +where + +import Control.Monad.Identity +import Control.Monad.RWS +import Data.Bifunctor (first) +import Data.Bits (shiftR, xor) +import Data.Char (isAlphaNum, isDigit, ord) +import qualified Data.DList as DL +import Data.FileEmbed +import Data.List (unzip4) +import Data.Loc +import qualified Data.Map.Strict as M +import Data.Maybe +import Futhark.CodeGen.Backends.GenericC.Options +import Futhark.CodeGen.Backends.SimpleRep +import Futhark.CodeGen.ImpCode +import Futhark.IR.Prop (isBuiltInFunction) +import Futhark.MonadFreshNames +import Futhark.Util (zEncodeString) +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C +import Text.Printf + +data CompilerState s = CompilerState + { compArrayStructs :: [((C.Type, Int), (C.Type, [C.Definition]))], + compOpaqueStructs :: [(String, (C.Type, [C.Definition]))], + compEarlyDecls :: DL.DList C.Definition, + compInit :: [C.Stm], + compNameSrc :: VNameSource, + compUserState :: s, + compHeaderDecls :: M.Map HeaderSection (DL.DList C.Definition), + compLibDecls :: DL.DList C.Definition, + compCtxFields :: DL.DList (C.Id, C.Type, Maybe C.Exp), + compProfileItems :: DL.DList C.BlockItem, + compDeclaredMem :: [(VName, Space)] + } + +newCompilerState :: VNameSource -> s -> CompilerState s +newCompilerState src s = + CompilerState + { compArrayStructs = [], + compOpaqueStructs = [], + compEarlyDecls = mempty, + compInit = [], + compNameSrc = src, + compUserState = s, + compHeaderDecls = mempty, + compLibDecls = mempty, + compCtxFields = mempty, + compProfileItems = mempty, + compDeclaredMem = mempty + } + +-- | In which part of the header file we put the declaration. This is +-- to ensure that the header file remains structured and readable. +data HeaderSection + = ArrayDecl String + | OpaqueDecl String + | EntryDecl + | MiscDecl + | InitDecl + deriving (Eq, Ord) + +-- | A substitute expression compiler, tried before the main +-- compilation function. +type OpCompiler op s = op -> CompilerM op s () + +type ErrorCompiler op s = ErrorMsg Exp -> String -> CompilerM op s () + +-- | The address space qualifiers for a pointer of the given type with +-- the given annotation. +type PointerQuals op s = String -> CompilerM op s [C.TypeQual] + +-- | The type of a memory block in the given memory space. +type MemoryType op s = SpaceId -> CompilerM op s C.Type + +-- | Write a scalar to the given memory block with the given element +-- index and in the given memory space. +type WriteScalar op s = + C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> C.Exp -> CompilerM op s () + +-- | Read a scalar from the given memory block with the given element +-- index and in the given memory space. +type ReadScalar op s = + C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> CompilerM op s C.Exp + +-- | Allocate a memory block of the given size and with the given tag +-- in the given memory space, saving a reference in the given variable +-- name. +type Allocate op s = + C.Exp -> + C.Exp -> + C.Exp -> + SpaceId -> + CompilerM op s () + +-- | De-allocate the given memory block with the given tag, which is +-- in the given memory space. +type Deallocate op s = C.Exp -> C.Exp -> SpaceId -> CompilerM op s () + +-- | Create a static array of values - initialised at load time. +type StaticArray op s = VName -> SpaceId -> PrimType -> ArrayContents -> CompilerM op s () + +-- | Copy from one memory block to another. +type Copy op s = + C.Exp -> + C.Exp -> + Space -> + C.Exp -> + C.Exp -> + Space -> + C.Exp -> + CompilerM op s () + +-- | Call a function. +type CallCompiler op s = [VName] -> Name -> [C.Exp] -> CompilerM op s () + +data Operations op s = Operations + { opsWriteScalar :: WriteScalar op s, + opsReadScalar :: ReadScalar op s, + opsAllocate :: Allocate op s, + opsDeallocate :: Deallocate op s, + opsCopy :: Copy op s, + opsStaticArray :: StaticArray op s, + opsMemoryType :: MemoryType op s, + opsCompiler :: OpCompiler op s, + opsError :: ErrorCompiler op s, + opsCall :: CallCompiler op s, + -- | If true, use reference counting. Otherwise, bare + -- pointers. + opsFatMemory :: Bool, + -- | Code to bracket critical sections. + opsCritical :: ([C.BlockItem], [C.BlockItem]) + } + +defError :: ErrorCompiler op s +defError (ErrorMsg parts) stacktrace = do + free_all_mem <- collect $ mapM_ (uncurry unRefMem) =<< gets compDeclaredMem + let onPart (ErrorString s) = return ("%s", [C.cexp|$string:s|]) + onPart (ErrorInt32 x) = ("%d",) <$> compileExp x + onPart (ErrorInt64 x) = ("%lld",) <$> compileExp x + (formatstrs, formatargs) <- unzip <$> mapM onPart parts + let formatstr = "Error: " ++ concat formatstrs ++ "\n\nBacktrace:\n%s" + items + [C.citems|ctx->error = msgprintf($string:formatstr, $args:formatargs, $string:stacktrace); + $items:free_all_mem + return 1;|] + +defCall :: CallCompiler op s +defCall dests fname args = do + let out_args = [[C.cexp|&$id:d|] | d <- dests] + args' + | isBuiltInFunction fname = args + | otherwise = [C.cexp|ctx|] : out_args ++ args + case dests of + [dest] + | isBuiltInFunction fname -> + stm [C.cstm|$id:dest = $id:(funName fname)($args:args');|] + _ -> + item [C.citem|if ($id:(funName fname)($args:args') != 0) { err = 1; goto cleanup; }|] + +-- | A set of operations that fail for every operation involving +-- non-default memory spaces. Uses plain pointers and @malloc@ for +-- memory management. +defaultOperations :: Operations op s +defaultOperations = + Operations + { opsWriteScalar = defWriteScalar, + opsReadScalar = defReadScalar, + opsAllocate = defAllocate, + opsDeallocate = defDeallocate, + opsCopy = defCopy, + opsStaticArray = defStaticArray, + opsMemoryType = defMemoryType, + opsCompiler = defCompiler, + opsFatMemory = True, + opsError = defError, + opsCall = defCall, + opsCritical = mempty + } + where + defWriteScalar _ _ _ _ _ = + error "Cannot write to non-default memory space because I am dumb" + defReadScalar _ _ _ _ = + error "Cannot read from non-default memory space" + defAllocate _ _ _ = + error "Cannot allocate in non-default memory space" + defDeallocate _ _ = + error "Cannot deallocate in non-default memory space" + defCopy destmem destoffset DefaultSpace srcmem srcoffset DefaultSpace size = + copyMemoryDefaultSpace destmem destoffset srcmem srcoffset size + defCopy _ _ _ _ _ _ _ = + error "Cannot copy to or from non-default memory space" + defStaticArray _ _ _ _ = + error "Cannot create static array in non-default memory space" + defMemoryType _ = + error "Has no type for non-default memory space" + defCompiler _ = + error "The default compiler cannot compile extended operations" + +data CompilerEnv op s = CompilerEnv + { envOperations :: Operations op s, + -- | Mapping memory blocks to sizes. These memory blocks are CPU + -- memory that we know are used in particularly simple ways (no + -- reference counting necessary). To cut down on allocator + -- pressure, we keep these allocations around for a long time, and + -- record their sizes so we can reuse them if possible (and + -- realloc() when needed). + envCachedMem :: M.Map C.Exp VName + } + +newtype CompilerAcc op s = CompilerAcc + { accItems :: DL.DList C.BlockItem + } + +instance Semigroup (CompilerAcc op s) where + CompilerAcc items1 <> CompilerAcc items2 = + CompilerAcc (items1 <> items2) + +instance Monoid (CompilerAcc op s) where + mempty = CompilerAcc mempty + +envOpCompiler :: CompilerEnv op s -> OpCompiler op s +envOpCompiler = opsCompiler . envOperations + +envMemoryType :: CompilerEnv op s -> MemoryType op s +envMemoryType = opsMemoryType . envOperations + +envReadScalar :: CompilerEnv op s -> ReadScalar op s +envReadScalar = opsReadScalar . envOperations + +envWriteScalar :: CompilerEnv op s -> WriteScalar op s +envWriteScalar = opsWriteScalar . envOperations + +envAllocate :: CompilerEnv op s -> Allocate op s +envAllocate = opsAllocate . envOperations + +envDeallocate :: CompilerEnv op s -> Deallocate op s +envDeallocate = opsDeallocate . envOperations + +envCopy :: CompilerEnv op s -> Copy op s +envCopy = opsCopy . envOperations + +envStaticArray :: CompilerEnv op s -> StaticArray op s +envStaticArray = opsStaticArray . envOperations + +envFatMemory :: CompilerEnv op s -> Bool +envFatMemory = opsFatMemory . envOperations + +arrayDefinitions, opaqueDefinitions :: CompilerState s -> [C.Definition] +arrayDefinitions = concatMap (snd . snd) . compArrayStructs +opaqueDefinitions = concatMap (snd . snd) . compOpaqueStructs + +initDecls, arrayDecls, opaqueDecls, entryDecls, miscDecls :: CompilerState s -> [C.Definition] +initDecls = concatMap (DL.toList . snd) . filter ((== InitDecl) . fst) . M.toList . compHeaderDecls +arrayDecls = concatMap (DL.toList . snd) . filter (isArrayDecl . fst) . M.toList . compHeaderDecls + where + isArrayDecl ArrayDecl {} = True + isArrayDecl _ = False +opaqueDecls = concatMap (DL.toList . snd) . filter (isOpaqueDecl . fst) . M.toList . compHeaderDecls + where + isOpaqueDecl OpaqueDecl {} = True + isOpaqueDecl _ = False +entryDecls = concatMap (DL.toList . snd) . filter ((== EntryDecl) . fst) . M.toList . compHeaderDecls +miscDecls = concatMap (DL.toList . snd) . filter ((== MiscDecl) . fst) . M.toList . compHeaderDecls + +contextContents :: CompilerM op s ([C.FieldGroup], [C.Stm]) +contextContents = do + (field_names, field_types, field_values) <- gets $ unzip3 . DL.toList . compCtxFields + let fields = + [ [C.csdecl|$ty:ty $id:name;|] + | (name, ty) <- zip field_names field_types + ] + init_fields = + [ [C.cstm|ctx->$id:name = $exp:e;|] + | (name, Just e) <- zip field_names field_values + ] + return (fields, init_fields) + +contextFinalInits :: CompilerM op s [C.Stm] +contextFinalInits = gets compInit + +newtype CompilerM op s a + = CompilerM + ( RWS + (CompilerEnv op s) + (CompilerAcc op s) + (CompilerState s) + a + ) + deriving + ( Functor, + Applicative, + Monad, + MonadState (CompilerState s), + MonadReader (CompilerEnv op s), + MonadWriter (CompilerAcc op s) + ) + +instance MonadFreshNames (CompilerM op s) where + getNameSource = gets compNameSrc + putNameSource src = modify $ \s -> s {compNameSrc = src} + +runCompilerM :: + Operations op s -> + VNameSource -> + s -> + CompilerM op s a -> + (a, CompilerState s) +runCompilerM ops src userstate (CompilerM m) = + let (x, s, _) = runRWS m (CompilerEnv ops mempty) (newCompilerState src userstate) + in (x, s) + +getUserState :: CompilerM op s s +getUserState = gets compUserState + +modifyUserState :: (s -> s) -> CompilerM op s () +modifyUserState f = modify $ \compstate -> + compstate {compUserState = f $ compUserState compstate} + +atInit :: C.Stm -> CompilerM op s () +atInit x = modify $ \s -> + s {compInit = compInit s ++ [x]} + +collect :: CompilerM op s () -> CompilerM op s [C.BlockItem] +collect m = snd <$> collect' m + +collect' :: CompilerM op s a -> CompilerM op s (a, [C.BlockItem]) +collect' m = pass $ do + (x, w) <- listen m + return + ( (x, DL.toList $ accItems w), + const w {accItems = mempty} + ) + +item :: C.BlockItem -> CompilerM op s () +item x = tell $ mempty {accItems = DL.singleton x} + +items :: [C.BlockItem] -> CompilerM op s () +items = mapM_ item + +fatMemory :: Space -> CompilerM op s Bool +fatMemory ScalarSpace {} = return False +fatMemory _ = asks envFatMemory + +cacheMem :: C.ToExp a => a -> CompilerM op s (Maybe VName) +cacheMem a = asks $ M.lookup (C.toExp a noLoc) . envCachedMem + +instance C.ToIdent Name where + toIdent = C.toIdent . zEncodeString . nameToString + +instance C.ToIdent VName where + toIdent = C.toIdent . zEncodeString . pretty + +instance C.ToExp VName where + toExp v _ = [C.cexp|$id:v|] + +instance C.ToExp IntValue where + toExp (Int8Value v) = C.toExp v + toExp (Int16Value v) = C.toExp v + toExp (Int32Value v) = C.toExp v + toExp (Int64Value v) = C.toExp v + +instance C.ToExp FloatValue where + toExp (Float32Value v) = C.toExp v + toExp (Float64Value v) = C.toExp v + +instance C.ToExp PrimValue where + toExp (IntValue v) = C.toExp v + toExp (FloatValue v) = C.toExp v + toExp (BoolValue True) = C.toExp (1 :: Int8) + toExp (BoolValue False) = C.toExp (0 :: Int8) + toExp Checked = C.toExp (1 :: Int8) + +instance C.ToExp SubExp where + toExp (Var v) = C.toExp v + toExp (Constant c) = C.toExp c + +-- | Construct a publicly visible definition using the specified name +-- as the template. The first returned definition is put in the +-- header file, and the second is the implementation. Returns the public +-- name. +publicDef :: + String -> + HeaderSection -> + (String -> (C.Definition, C.Definition)) -> + CompilerM op s String +publicDef s h f = do + s' <- publicName s + let (pub, priv) = f s' + headerDecl h pub + earlyDecl priv + return s' + +-- | As 'publicDef', but ignores the public name. +publicDef_ :: + String -> + HeaderSection -> + (String -> (C.Definition, C.Definition)) -> + CompilerM op s () +publicDef_ s h f = void $ publicDef s h f + +headerDecl :: HeaderSection -> C.Definition -> CompilerM op s () +headerDecl sec def = modify $ \s -> + s + { compHeaderDecls = + M.unionWith + (<>) + (compHeaderDecls s) + (M.singleton sec (DL.singleton def)) + } + +libDecl :: C.Definition -> CompilerM op s () +libDecl def = modify $ \s -> + s {compLibDecls = compLibDecls s <> DL.singleton def} + +earlyDecl :: C.Definition -> CompilerM op s () +earlyDecl def = modify $ \s -> + s {compEarlyDecls = compEarlyDecls s <> DL.singleton def} + +contextField :: C.Id -> C.Type -> Maybe C.Exp -> CompilerM op s () +contextField name ty initial = modify $ \s -> + s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, initial)} + +profileReport :: C.BlockItem -> CompilerM op s () +profileReport x = modify $ \s -> + s {compProfileItems = compProfileItems s <> DL.singleton x} + +stm :: C.Stm -> CompilerM op s () +stm s = item [C.citem|$stm:s|] + +stms :: [C.Stm] -> CompilerM op s () +stms = mapM_ stm + +decl :: C.InitGroup -> CompilerM op s () +decl x = item [C.citem|$decl:x;|] + +addrOf :: C.Exp -> C.Exp +addrOf e = [C.cexp|&$exp:e|] + +-- | Public names must have a consitent prefix. +publicName :: String -> CompilerM op s String +publicName s = return $ "futhark_" ++ s + +-- | The generated code must define a struct with this name. +contextType :: CompilerM op s C.Type +contextType = do + name <- publicName "context" + return [C.cty|struct $id:name|] + +memToCType :: VName -> Space -> CompilerM op s C.Type +memToCType v space = do + refcount <- fatMemory space + cached <- isJust <$> cacheMem v + if refcount && not cached + then return $ fatMemType space + else rawMemCType space + +rawMemCType :: Space -> CompilerM op s C.Type +rawMemCType DefaultSpace = return defaultMemBlockType +rawMemCType (Space sid) = join $ asks envMemoryType <*> pure sid +rawMemCType (ScalarSpace [] t) = + return [C.cty|$ty:(primTypeToCType t)[1]|] +rawMemCType (ScalarSpace ds t) = + return [C.cty|$ty:(primTypeToCType t)[$exp:(cproduct ds')]|] + where + ds' = map (`C.toExp` noLoc) ds + +fatMemType :: Space -> C.Type +fatMemType space = + [C.cty|struct $id:name|] + where + name = case space of + Space sid -> "memblock_" ++ sid + _ -> "memblock" + +fatMemSet :: Space -> String +fatMemSet (Space sid) = "memblock_set_" ++ sid +fatMemSet _ = "memblock_set" + +fatMemAlloc :: Space -> String +fatMemAlloc (Space sid) = "memblock_alloc_" ++ sid +fatMemAlloc _ = "memblock_alloc" + +fatMemUnRef :: Space -> String +fatMemUnRef (Space sid) = "memblock_unref_" ++ sid +fatMemUnRef _ = "memblock_unref" + +rawMem :: VName -> CompilerM op s C.Exp +rawMem v = rawMem' <$> fat <*> pure v + where + fat = asks ((&&) . envFatMemory) <*> (isNothing <$> cacheMem v) + +rawMem' :: C.ToExp a => Bool -> a -> C.Exp +rawMem' True e = [C.cexp|$exp:e.mem|] +rawMem' False e = [C.cexp|$exp:e|] + +allocRawMem :: + (C.ToExp a, C.ToExp b, C.ToExp c) => + a -> + b -> + Space -> + c -> + CompilerM op s () +allocRawMem dest size space desc = case space of + Space sid -> + join $ + asks envAllocate <*> pure [C.cexp|$exp:dest|] + <*> pure [C.cexp|$exp:size|] + <*> pure [C.cexp|$exp:desc|] + <*> pure sid + _ -> + stm [C.cstm|$exp:dest = (char*) malloc($exp:size);|] + +freeRawMem :: + (C.ToExp a, C.ToExp b) => + a -> + Space -> + b -> + CompilerM op s () +freeRawMem mem space desc = + case space of + Space sid -> do + free_mem <- asks envDeallocate + free_mem [C.cexp|$exp:mem|] [C.cexp|$exp:desc|] sid + _ -> item [C.citem|free($exp:mem);|] + +defineMemorySpace :: Space -> CompilerM op s (C.Definition, [C.Definition], C.BlockItem) +defineMemorySpace space = do + rm <- rawMemCType space + let structdef = + [C.cedecl|struct $id:sname { int *references; + $ty:rm mem; + typename int64_t size; + const char *desc; };|] + + contextField peakname [C.cty|typename int64_t|] $ Just [C.cexp|0|] + contextField usagename [C.cty|typename int64_t|] $ Just [C.cexp|0|] + + -- Unreferencing a memory block consists of decreasing its reference + -- count and freeing the corresponding memory if the count reaches + -- zero. + free <- collect $ freeRawMem [C.cexp|block->mem|] space [C.cexp|desc|] + ctx_ty <- contextType + let unrefdef = + [C.cedecl|static int $id:(fatMemUnRef space) ($ty:ctx_ty *ctx, $ty:mty *block, const char *desc) { + if (block->references != NULL) { + *(block->references) -= 1; + if (ctx->detail_memory) { + fprintf(stderr, "Unreferencing block %s (allocated as %s) in %s: %d references remaining.\n", + desc, block->desc, $string:spacedesc, *(block->references)); + } + if (*(block->references) == 0) { + ctx->$id:usagename -= block->size; + $items:free + free(block->references); + if (ctx->detail_memory) { + fprintf(stderr, "%lld bytes freed (now allocated: %lld bytes)\n", + (long long) block->size, (long long) ctx->$id:usagename); + } + } + block->references = NULL; + } + return 0; +}|] + + -- When allocating a memory block we initialise the reference count to 1. + alloc <- + collect $ + allocRawMem [C.cexp|block->mem|] [C.cexp|size|] space [C.cexp|desc|] + let allocdef = + [C.cedecl|static int $id:(fatMemAlloc space) ($ty:ctx_ty *ctx, $ty:mty *block, typename int64_t size, const char *desc) { + if (size < 0) { + futhark_panic(1, "Negative allocation of %lld bytes attempted for %s in %s.\n", + (long long)size, desc, $string:spacedesc, ctx->$id:usagename); + } + int ret = $id:(fatMemUnRef space)(ctx, block, desc); + + ctx->$id:usagename += size; + if (ctx->detail_memory) { + fprintf(stderr, "Allocating %lld bytes for %s in %s (then allocated: %lld bytes)", + (long long) size, + desc, $string:spacedesc, + (long long) ctx->$id:usagename); + } + if (ctx->$id:usagename > ctx->$id:peakname) { + ctx->$id:peakname = ctx->$id:usagename; + if (ctx->detail_memory) { + fprintf(stderr, " (new peak).\n"); + } + } else if (ctx->detail_memory) { + fprintf(stderr, ".\n"); + } + + $items:alloc + block->references = (int*) malloc(sizeof(int)); + *(block->references) = 1; + block->size = size; + block->desc = desc; + return ret; + }|] + + -- Memory setting - unreference the destination and increase the + -- count of the source by one. + let setdef = + [C.cedecl|static int $id:(fatMemSet space) ($ty:ctx_ty *ctx, $ty:mty *lhs, $ty:mty *rhs, const char *lhs_desc) { + int ret = $id:(fatMemUnRef space)(ctx, lhs, lhs_desc); + (*(rhs->references))++; + *lhs = *rhs; + return ret; +} +|] + + let peakmsg = "Peak memory usage for " ++ spacedesc ++ ": %lld bytes.\n" + return + ( structdef, + [unrefdef, allocdef, setdef], + -- Do not report memory usage for DefaultSpace (CPU memory), + -- because it would not be accurate anyway. This whole + -- tracking probably needs to be rethought. + if space == DefaultSpace + then [C.citem|{}|] + else [C.citem|str_builder(&builder, $string:peakmsg, (long long) ctx->$id:peakname);|] + ) + where + mty = fatMemType space + (peakname, usagename, sname, spacedesc) = case space of + Space sid -> + ( C.toIdent ("peak_mem_usage_" ++ sid) noLoc, + C.toIdent ("cur_mem_usage_" ++ sid) noLoc, + C.toIdent ("memblock_" ++ sid) noLoc, + "space '" ++ sid ++ "'" + ) + _ -> + ( "peak_mem_usage_default", + "cur_mem_usage_default", + "memblock", + "default space" + ) + +declMem :: VName -> Space -> CompilerM op s () +declMem name space = do + cached <- isJust <$> cacheMem name + unless cached $ do + ty <- memToCType name space + decl [C.cdecl|$ty:ty $id:name;|] + resetMem name space + modify $ \s -> s {compDeclaredMem = (name, space) : compDeclaredMem s} + +resetMem :: C.ToExp a => a -> Space -> CompilerM op s () +resetMem mem space = do + refcount <- fatMemory space + cached <- isJust <$> cacheMem mem + if cached + then stm [C.cstm|$exp:mem = NULL;|] + else + when refcount $ + stm [C.cstm|$exp:mem.references = NULL;|] + +setMem :: (C.ToExp a, C.ToExp b) => a -> b -> Space -> CompilerM op s () +setMem dest src space = do + refcount <- fatMemory space + let src_s = pretty $ C.toExp src noLoc + if refcount + then + stm + [C.cstm|if ($id:(fatMemSet space)(ctx, &$exp:dest, &$exp:src, + $string:src_s) != 0) { + return 1; + }|] + else case space of + ScalarSpace ds _ -> do + i' <- newVName "i" + let i = C.toIdent i' + it = primTypeToCType $ IntType Int32 + ds' = map (`C.toExp` noLoc) ds + bound = cproduct ds' + stm + [C.cstm|for ($ty:it $id:i = 0; $id:i < $exp:bound; $id:i++) { + $exp:dest[$id:i] = $exp:src[$id:i]; + }|] + _ -> stm [C.cstm|$exp:dest = $exp:src;|] + +unRefMem :: C.ToExp a => a -> Space -> CompilerM op s () +unRefMem mem space = do + refcount <- fatMemory space + cached <- isJust <$> cacheMem mem + let mem_s = pretty $ C.toExp mem noLoc + when (refcount && not cached) $ + stm + [C.cstm|if ($id:(fatMemUnRef space)(ctx, &$exp:mem, $string:mem_s) != 0) { + return 1; + }|] + +allocMem :: + (C.ToExp a, C.ToExp b) => + a -> + b -> + Space -> + C.Stm -> + CompilerM op s () +allocMem mem size space on_failure = do + refcount <- fatMemory space + let mem_s = pretty $ C.toExp mem noLoc + if refcount + then + stm + [C.cstm|if ($id:(fatMemAlloc space)(ctx, &$exp:mem, $exp:size, + $string:mem_s)) { + $stm:on_failure + }|] + else do + freeRawMem mem space mem_s + allocRawMem mem size space [C.cexp|desc|] + +primTypeInfo :: PrimType -> Signedness -> C.Exp +primTypeInfo (IntType it) t = case (it, t) of + (Int8, TypeUnsigned) -> [C.cexp|u8_info|] + (Int16, TypeUnsigned) -> [C.cexp|u16_info|] + (Int32, TypeUnsigned) -> [C.cexp|u32_info|] + (Int64, TypeUnsigned) -> [C.cexp|u64_info|] + (Int8, _) -> [C.cexp|i8_info|] + (Int16, _) -> [C.cexp|i16_info|] + (Int32, _) -> [C.cexp|i32_info|] + (Int64, _) -> [C.cexp|i64_info|] +primTypeInfo (FloatType Float32) _ = [C.cexp|f32_info|] +primTypeInfo (FloatType Float64) _ = [C.cexp|f64_info|] +primTypeInfo Bool _ = [C.cexp|bool_info|] +primTypeInfo Cert _ = [C.cexp|bool_info|] + +copyMemoryDefaultSpace :: + C.Exp -> + C.Exp -> + C.Exp -> + C.Exp -> + C.Exp -> + CompilerM op s () +copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes = + stm + [C.cstm|memmove($exp:destmem + $exp:destidx, + $exp:srcmem + $exp:srcidx, + $exp:nbytes);|] + +--- Entry points. + +arrayName :: PrimType -> Signedness -> Int -> String +arrayName pt signed rank = + prettySigned (signed == TypeUnsigned) pt ++ "_" ++ show rank ++ "d" + +opaqueName :: String -> [ValueDesc] -> String +opaqueName s _ + | valid = "opaque_" ++ s + where + valid = + head s /= '_' + && not (isDigit $ head s) + && all ok s + ok c = isAlphaNum c || c == '_' +opaqueName s vds = "opaque_" ++ hash (zipWith xor [0 ..] $ map ord (s ++ concatMap p vds)) + where + p (ScalarValue pt signed _) = + show (pt, signed) + p (ArrayValue _ space pt signed dims) = + show (space, pt, signed, length dims) + + -- FIXME: a stupid hash algorithm; may have collisions. + hash = + printf "%x" . foldl xor 0 + . map + ( iter . (* 0x45d9f3b) + . iter + . (* 0x45d9f3b) + . iter + . fromIntegral + ) + iter x = ((x :: Word32) `shiftR` 16) `xor` x + +criticalSection :: Operations op s -> [C.BlockItem] -> [C.BlockItem] +criticalSection ops x = + [C.citems|lock_lock(&ctx->lock); + $items:(fst (opsCritical ops)) + $items:x + $items:(snd (opsCritical ops)) + lock_unlock(&ctx->lock); + |] + +arrayLibraryFunctions :: + Space -> + PrimType -> + Signedness -> + [DimSize] -> + CompilerM op s [C.Definition] +arrayLibraryFunctions space pt signed shape = do + let rank = length shape + pt' = signedPrimTypeToCType signed pt + name = arrayName pt signed rank + arr_name = "futhark_" ++ name + array_type = [C.cty|struct $id:arr_name|] + + new_array <- publicName $ "new_" ++ name + new_raw_array <- publicName $ "new_raw_" ++ name + free_array <- publicName $ "free_" ++ name + values_array <- publicName $ "values_" ++ name + values_raw_array <- publicName $ "values_raw_" ++ name + shape_array <- publicName $ "shape_" ++ name + + let shape_names = ["dim" ++ show i | i <- [0 .. rank -1]] + shape_params = [[C.cparam|typename int64_t $id:k|] | k <- shape_names] + arr_size = cproduct [[C.cexp|$id:k|] | k <- shape_names] + arr_size_array = cproduct [[C.cexp|arr->shape[$int:i]|] | i <- [0 .. rank -1]] + copy <- asks envCopy + + memty <- rawMemCType space + + let prepare_new = do + resetMem [C.cexp|arr->mem|] space + allocMem + [C.cexp|arr->mem|] + [C.cexp|((size_t)$exp:arr_size) * sizeof($ty:pt')|] + space + [C.cstm|return NULL;|] + forM_ [0 .. rank -1] $ \i -> + let dim_s = "dim" ++ show i + in stm [C.cstm|arr->shape[$int:i] = $id:dim_s;|] + + new_body <- collect $ do + prepare_new + copy + [C.cexp|arr->mem.mem|] + [C.cexp|0|] + space + [C.cexp|data|] + [C.cexp|0|] + DefaultSpace + [C.cexp|((size_t)$exp:arr_size) * sizeof($ty:pt')|] + + new_raw_body <- collect $ do + prepare_new + copy + [C.cexp|arr->mem.mem|] + [C.cexp|0|] + space + [C.cexp|data|] + [C.cexp|offset|] + space + [C.cexp|((size_t)$exp:arr_size) * sizeof($ty:pt')|] + + free_body <- collect $ unRefMem [C.cexp|arr->mem|] space + + values_body <- + collect $ + copy + [C.cexp|data|] + [C.cexp|0|] + DefaultSpace + [C.cexp|arr->mem.mem|] + [C.cexp|0|] + space + [C.cexp|((size_t)$exp:arr_size_array) * sizeof($ty:pt')|] + + ctx_ty <- contextType + ops <- asks envOperations + + headerDecl + (ArrayDecl name) + [C.cedecl|struct $id:arr_name;|] + headerDecl + (ArrayDecl name) + [C.cedecl|$ty:array_type* $id:new_array($ty:ctx_ty *ctx, const $ty:pt' *data, $params:shape_params);|] + headerDecl + (ArrayDecl name) + [C.cedecl|$ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, const $ty:memty data, int offset, $params:shape_params);|] + headerDecl + (ArrayDecl name) + [C.cedecl|int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] + headerDecl + (ArrayDecl name) + [C.cedecl|int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data);|] + headerDecl + (ArrayDecl name) + [C.cedecl|$ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] + headerDecl + (ArrayDecl name) + [C.cedecl|const typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] + + return + [C.cunit| + $ty:array_type* $id:new_array($ty:ctx_ty *ctx, const $ty:pt' *data, $params:shape_params) { + $ty:array_type* bad = NULL; + $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type)); + if (arr == NULL) { + return bad; + } + $items:(criticalSection ops new_body) + return arr; + } + + $ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, const $ty:memty data, int offset, + $params:shape_params) { + $ty:array_type* bad = NULL; + $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type)); + if (arr == NULL) { + return bad; + } + $items:(criticalSection ops new_raw_body) + return arr; + } + + int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr) { + $items:(criticalSection ops free_body) + free(arr); + return 0; + } + + int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data) { + $items:(criticalSection ops values_body) + return 0; + } + + $ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr) { + (void)ctx; + return arr->mem.mem; + } + + const typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr) { + (void)ctx; + return arr->shape; + } + |] + +opaqueLibraryFunctions :: + String -> + [ValueDesc] -> + CompilerM op s [C.Definition] +opaqueLibraryFunctions desc vds = do + name <- publicName $ opaqueName desc vds + free_opaque <- publicName $ "free_" ++ opaqueName desc vds + + let opaque_type = [C.cty|struct $id:name|] + + freeComponent _ ScalarValue {} = + return () + freeComponent i (ArrayValue _ _ pt signed shape) = do + let rank = length shape + free_array <- publicName $ "free_" ++ arrayName pt signed rank + stm + [C.cstm|if ((tmp = $id:free_array(ctx, obj->$id:(tupleField i))) != 0) { + ret = tmp; + }|] + + ctx_ty <- contextType + + free_body <- collect $ zipWithM_ freeComponent [0 ..] vds + + headerDecl + (OpaqueDecl desc) + [C.cedecl|int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj);|] + + return + [C.cunit| + int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj) { + int ret = 0, tmp; + $items:free_body + free(obj); + return ret; + } + |] + +valueDescToCType :: ValueDesc -> CompilerM op s C.Type +valueDescToCType (ScalarValue pt signed _) = + return $ signedPrimTypeToCType signed pt +valueDescToCType (ArrayValue mem space pt signed shape) = do + let pt' = signedPrimTypeToCType signed pt + rank = length shape + exists <- gets $ lookup (pt', rank) . compArrayStructs + case exists of + Just (cty, _) -> return cty + Nothing -> do + memty <- memToCType mem space + name <- publicName $ arrayName pt signed rank + let struct = [C.cedecl|struct $id:name { $ty:memty mem; typename int64_t shape[$int:rank]; };|] + stype = [C.cty|struct $id:name|] + library <- arrayLibraryFunctions space pt signed shape + modify $ \s -> + s + { compArrayStructs = + ((pt', rank), (stype, struct : library)) : compArrayStructs s + } + return stype + +opaqueToCType :: String -> [ValueDesc] -> CompilerM op s C.Type +opaqueToCType desc vds = do + name <- publicName $ opaqueName desc vds + exists <- gets $ lookup name . compOpaqueStructs + case exists of + Just (ty, _) -> return ty + Nothing -> do + members <- zipWithM field vds [(0 :: Int) ..] + let struct = [C.cedecl|struct $id:name { $sdecls:members };|] + stype = [C.cty|struct $id:name|] + headerDecl (OpaqueDecl desc) [C.cedecl|struct $id:name;|] + library <- opaqueLibraryFunctions desc vds + modify $ \s -> + s + { compOpaqueStructs = + (name, (stype, struct : library)) : + compOpaqueStructs s + } + return stype + where + field vd@ScalarValue {} i = do + ct <- valueDescToCType vd + return [C.csdecl|$ty:ct $id:(tupleField i);|] + field vd i = do + ct <- valueDescToCType vd + return [C.csdecl|$ty:ct *$id:(tupleField i);|] + +externalValueToCType :: ExternalValue -> CompilerM op s C.Type +externalValueToCType (TransparentValue vd) = valueDescToCType vd +externalValueToCType (OpaqueValue desc vds) = opaqueToCType desc vds + +prepareEntryInputs :: [ExternalValue] -> CompilerM op s [C.Param] +prepareEntryInputs = zipWithM prepare [(0 :: Int) ..] + where + prepare pno (TransparentValue vd) = do + let pname = "in" ++ show pno + ty <- prepareValue [C.cexp|$id:pname|] vd + return [C.cparam|const $ty:ty $id:pname|] + prepare pno (OpaqueValue desc vds) = do + ty <- opaqueToCType desc vds + let pname = "in" ++ show pno + field i ScalarValue {} = [C.cexp|$id:pname->$id:(tupleField i)|] + field i ArrayValue {} = [C.cexp|$id:pname->$id:(tupleField i)|] + zipWithM_ prepareValue (zipWith field [0 ..] vds) vds + return [C.cparam|const $ty:ty *$id:pname|] + + prepareValue src (ScalarValue pt signed name) = do + let pt' = signedPrimTypeToCType signed pt + stm [C.cstm|$id:name = $exp:src;|] + return pt' + prepareValue src vd@(ArrayValue mem _ _ _ shape) = do + ty <- valueDescToCType vd + + stm [C.cstm|$exp:mem = $exp:src->mem;|] + + let rank = length shape + maybeCopyDim (Var d) i = + Just [C.cstm|$id:d = $exp:src->shape[$int:i];|] + maybeCopyDim _ _ = Nothing + + stms $ catMaybes $ zipWith maybeCopyDim shape [0 .. rank -1] + + return [C.cty|$ty:ty*|] + +prepareEntryOutputs :: [ExternalValue] -> CompilerM op s [C.Param] +prepareEntryOutputs = zipWithM prepare [(0 :: Int) ..] + where + prepare pno (TransparentValue vd) = do + let pname = "out" ++ show pno + ty <- valueDescToCType vd + + case vd of + ArrayValue {} -> do + stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|] + prepareValue [C.cexp|*$id:pname|] vd + return [C.cparam|$ty:ty **$id:pname|] + ScalarValue {} -> do + prepareValue [C.cexp|*$id:pname|] vd + return [C.cparam|$ty:ty *$id:pname|] + prepare pno (OpaqueValue desc vds) = do + let pname = "out" ++ show pno + ty <- opaqueToCType desc vds + vd_ts <- mapM valueDescToCType vds + + stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|] + + forM_ (zip3 [0 ..] vd_ts vds) $ \(i, ct, vd) -> do + let field = [C.cexp|(*$id:pname)->$id:(tupleField i)|] + case vd of + ScalarValue {} -> return () + _ -> stm [C.cstm|assert(($exp:field = ($ty:ct*) malloc(sizeof($ty:ct))) != NULL);|] + prepareValue field vd + + return [C.cparam|$ty:ty **$id:pname|] + + prepareValue dest (ScalarValue _ _ name) = + stm [C.cstm|$exp:dest = $id:name;|] + prepareValue dest (ArrayValue mem _ _ _ shape) = do + stm [C.cstm|$exp:dest->mem = $id:mem;|] + + let rank = length shape + maybeCopyDim (Constant x) i = + [C.cstm|$exp:dest->shape[$int:i] = $exp:x;|] + maybeCopyDim (Var d) i = + [C.cstm|$exp:dest->shape[$int:i] = $id:d;|] + stms $ zipWith maybeCopyDim shape [0 .. rank -1] + +onEntryPoint :: + Name -> + Function op -> + CompilerM op s (C.Definition, C.Definition, C.Initializer) +onEntryPoint fname function@(Function _ outputs inputs _ results args) = do + let out_args = map (\p -> [C.cexp|&$id:(paramName p)|]) outputs + in_args = map (\p -> [C.cexp|$id:(paramName p)|]) inputs + + inputdecls <- collect $ mapM_ stubParam inputs + outputdecls <- collect $ mapM_ stubParam outputs + + let entry_point_name = nameToString fname + entry_point_function_name <- publicName $ "entry_" ++ entry_point_name + + (entry_point_input_params, unpack_entry_inputs) <- + collect' $ prepareEntryInputs args + (entry_point_output_params, pack_entry_outputs) <- + collect' $ prepareEntryOutputs results + + (cli_entry_point, cli_init) <- cliEntryPoint fname function + + ctx_ty <- contextType + + headerDecl + EntryDecl + [C.cedecl|int $id:entry_point_function_name + ($ty:ctx_ty *ctx, + $params:entry_point_output_params, + $params:entry_point_input_params);|] + + let critical = + [C.citems| + $items:unpack_entry_inputs + + int ret = $id:(funName fname)(ctx, $args:out_args, $args:in_args); + + if (ret == 0) { + $items:pack_entry_outputs + } + |] + + ops <- asks envOperations + + return + ( [C.cedecl| + int $id:entry_point_function_name + ($ty:ctx_ty *ctx, + $params:entry_point_output_params, + $params:entry_point_input_params) { + $items:inputdecls + $items:outputdecls + + $items:(criticalSection ops critical) + + return ret; + }|], + cli_entry_point, + cli_init + ) + where + stubParam (MemParam name space) = + declMem name space + stubParam (ScalarParam name ty) = do + let ty' = primTypeToCType ty + decl [C.cdecl|$ty:ty' $id:name;|] + +--- CLI interface +-- +-- Our strategy for CLI entry points is to parse everything into +-- host memory ('DefaultSpace') and copy the result into host memory +-- after the entry point has returned. We have some ad-hoc frobbery +-- to copy the host-level memory blocks to another memory space if +-- necessary. This will break if the Futhark entry point uses +-- non-trivial index functions for its input or output. +-- +-- The idea here is to keep the nastyness in the wrapper, whilst not +-- messing up anything else. + +printPrimStm :: (C.ToExp a, C.ToExp b) => a -> b -> PrimType -> Signedness -> C.Stm +printPrimStm dest val bt ept = + [C.cstm|write_scalar($exp:dest, binary_output, &$exp:(primTypeInfo bt ept), &$exp:val);|] + +-- | Return a statement printing the given external value. +printStm :: ExternalValue -> C.Exp -> CompilerM op s C.Stm +printStm (OpaqueValue desc _) _ = + return [C.cstm|printf("#", $string:desc);|] +printStm (TransparentValue (ScalarValue bt ept _)) e = + return $ printPrimStm [C.cexp|stdout|] e bt ept +printStm (TransparentValue (ArrayValue _ _ bt ept shape)) e = do + values_array <- publicName $ "values_" ++ name + shape_array <- publicName $ "shape_" ++ name + let num_elems = cproduct [[C.cexp|$id:shape_array(ctx, $exp:e)[$int:i]|] | i <- [0 .. rank -1]] + return + [C.cstm|{ + $ty:bt' *arr = calloc(sizeof($ty:bt'), $exp:num_elems); + assert(arr != NULL); + assert($id:values_array(ctx, $exp:e, arr) == 0); + write_array(stdout, binary_output, &$exp:(primTypeInfo bt ept), arr, + $id:shape_array(ctx, $exp:e), $int:rank); + free(arr); + }|] + where + rank = length shape + bt' = primTypeToCType bt + name = arrayName bt ept rank + +readPrimStm :: C.ToExp a => a -> Int -> PrimType -> Signedness -> C.Stm +readPrimStm place i t ept = + [C.cstm|if (read_scalar(&$exp:(primTypeInfo t ept),&$exp:place) != 0) { + futhark_panic(1, "Error when reading input #%d of type %s (errno: %s).\n", + $int:i, + $exp:(primTypeInfo t ept).type_name, + strerror(errno)); + }|] + +readInputs :: [ExternalValue] -> CompilerM op s [(C.Stm, C.Stm, C.Stm, C.Exp)] +readInputs = zipWithM readInput [0 ..] + +readInput :: Int -> ExternalValue -> CompilerM op s (C.Stm, C.Stm, C.Stm, C.Exp) +readInput i (OpaqueValue desc _) = do + stm [C.cstm|futhark_panic(1, "Cannot read input #%d of type %s\n", $int:i, $string:desc);|] + return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|NULL|]) +readInput i (TransparentValue (ScalarValue t ept _)) = do + dest <- newVName "read_value" + item [C.citem|$ty:(primTypeToCType t) $id:dest;|] + stm $ readPrimStm dest i t ept + return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|$id:dest|]) +readInput i (TransparentValue vd@(ArrayValue _ _ t ept dims)) = do + dest <- newVName "read_value" + shape <- newVName "read_shape" + arr <- newVName "read_arr" + ty <- valueDescToCType vd + item [C.citem|$ty:ty *$id:dest;|] + + let t' = signedPrimTypeToCType ept t + rank = length dims + name = arrayName t ept rank + dims_exps = [[C.cexp|$id:shape[$int:j]|] | j <- [0 .. rank -1]] + dims_s = concat $ replicate rank "[]" + + new_array <- publicName $ "new_" ++ name + free_array <- publicName $ "free_" ++ name + + items + [C.citems| + typename int64_t $id:shape[$int:rank]; + $ty:t' *$id:arr = NULL; + errno = 0; + if (read_array(&$exp:(primTypeInfo t ept), + (void**) &$id:arr, + $id:shape, + $int:(length dims)) + != 0) { + futhark_panic(1, "Cannot read input #%d of type %s%s (errno: %s).\n", + $int:i, + $string:dims_s, + $exp:(primTypeInfo t ept).type_name, + strerror(errno)); + }|] + + return + ( [C.cstm|assert(($exp:dest = $id:new_array(ctx, $id:arr, $args:dims_exps)) != 0);|], + [C.cstm|assert($id:free_array(ctx, $exp:dest) == 0);|], + [C.cstm|free($id:arr);|], + [C.cexp|$id:dest|] + ) + +prepareOutputs :: [ExternalValue] -> CompilerM op s [(C.Exp, C.Stm)] +prepareOutputs = mapM prepareResult + where + prepareResult ev = do + ty <- externalValueToCType ev + result <- newVName "result" + + case ev of + TransparentValue ScalarValue {} -> do + item [C.citem|$ty:ty $id:result;|] + return ([C.cexp|$id:result|], [C.cstm|;|]) + TransparentValue (ArrayValue _ _ t ept dims) -> do + let name = arrayName t ept $ length dims + free_array <- publicName $ "free_" ++ name + item [C.citem|$ty:ty *$id:result;|] + return + ( [C.cexp|$id:result|], + [C.cstm|assert($id:free_array(ctx, $exp:result) == 0);|] + ) + OpaqueValue desc vds -> do + free_opaque <- publicName $ "free_" ++ opaqueName desc vds + item [C.citem|$ty:ty *$id:result;|] + return + ( [C.cexp|$id:result|], + [C.cstm|assert($id:free_opaque(ctx, $exp:result) == 0);|] + ) + +printResult :: [(ExternalValue, C.Exp)] -> CompilerM op s [C.Stm] +printResult vs = fmap concat $ + forM vs $ \(v, e) -> do + p <- printStm v e + return [p, [C.cstm|printf("\n");|]] + +cliEntryPoint :: + Name -> + FunctionT a -> + CompilerM op s (C.Definition, C.Initializer) +cliEntryPoint fname (Function _ _ _ _ results args) = do + ((pack_input, free_input, free_parsed, input_args), input_items) <- + collect' $ unzip4 <$> readInputs args + + ((output_vals, free_outputs), output_decls) <- + collect' $ unzip <$> prepareOutputs results + printstms <- printResult $ zip results output_vals + + ctx_ty <- contextType + sync_ctx <- publicName "context_sync" + error_ctx <- publicName "context_get_error" + + let entry_point_name = nameToString fname + cli_entry_point_function_name = "futrts_cli_entry_" ++ entry_point_name + entry_point_function_name <- publicName $ "entry_" ++ entry_point_name + + pause_profiling <- publicName "context_pause_profiling" + unpause_profiling <- publicName "context_unpause_profiling" + + let run_it = + [C.citems| + int r; + // Run the program once. + $stms:pack_input + if ($id:sync_ctx(ctx) != 0) { + futhark_panic(1, "%s", $id:error_ctx(ctx)); + }; + // Only profile last run. + if (profile_run) { + $id:unpause_profiling(ctx); + } + t_start = get_wall_time(); + r = $id:entry_point_function_name(ctx, + $args:(map addrOf output_vals), + $args:input_args); + if (r != 0) { + futhark_panic(1, "%s", $id:error_ctx(ctx)); + } + if ($id:sync_ctx(ctx) != 0) { + futhark_panic(1, "%s", $id:error_ctx(ctx)); + }; + if (profile_run) { + $id:pause_profiling(ctx); + } + t_end = get_wall_time(); + long int elapsed_usec = t_end - t_start; + if (time_runs && runtime_file != NULL) { + fprintf(runtime_file, "%lld\n", (long long) elapsed_usec); + fflush(runtime_file); + } + $stms:free_input + |] + + return + ( [C.cedecl|static void $id:cli_entry_point_function_name($ty:ctx_ty *ctx) { + typename int64_t t_start, t_end; + int time_runs = 0, profile_run = 0; + + // We do not want to profile all the initialisation. + $id:pause_profiling(ctx); + + // Declare and read input. + set_binary_mode(stdin); + $items:input_items + + if (end_of_input() != 0) { + futhark_panic(1, "Expected EOF on stdin after reading input for %s.\n", $string:(quote (pretty fname))); + } + + $items:output_decls + + // Warmup run + if (perform_warmup) { + $items:run_it + $stms:free_outputs + } + time_runs = 1; + // Proper run. + for (int run = 0; run < num_runs; run++) { + // Only profile last run. + profile_run = run == num_runs -1; + $items:run_it + if (run < num_runs-1) { + $stms:free_outputs + } + } + + // Free the parsed input. + $stms:free_parsed + + // Print the final result. + if (binary_output) { + set_binary_mode(stdout); + } + $stms:printstms + + $stms:free_outputs + } + |], + [C.cinit|{ .name = $string:entry_point_name, + .fun = $id:cli_entry_point_function_name }|] + ) + +genericOptions :: [Option] +genericOptions = + [ Option + { optionLongName = "write-runtime-to", + optionShortName = Just 't', + optionArgument = RequiredArgument "FILE", + optionDescription = "Print the time taken to execute the program to the indicated file, an integral number of microseconds.", + optionAction = set_runtime_file + }, + Option + { optionLongName = "runs", + optionShortName = Just 'r', + optionArgument = RequiredArgument "INT", + optionDescription = "Perform NUM runs of the program.", + optionAction = set_num_runs + }, + Option + { optionLongName = "debugging", + optionShortName = Just 'D', + optionArgument = NoArgument, + optionDescription = "Perform possibly expensive internal correctness checks and verbose logging.", + optionAction = [C.cstm|futhark_context_config_set_debugging(cfg, 1);|] + }, + Option + { optionLongName = "log", + optionShortName = Just 'L', + optionArgument = NoArgument, + optionDescription = "Print various low-overhead logging information to stderr while running.", + optionAction = [C.cstm|futhark_context_config_set_logging(cfg, 1);|] + }, + Option + { optionLongName = "entry-point", + optionShortName = Just 'e', + optionArgument = RequiredArgument "NAME", + optionDescription = "The entry point to run. Defaults to main.", + optionAction = [C.cstm|if (entry_point != NULL) entry_point = optarg;|] + }, + Option + { optionLongName = "binary-output", + optionShortName = Just 'b', + optionArgument = NoArgument, + optionDescription = "Print the program result in the binary output format.", + optionAction = [C.cstm|binary_output = 1;|] + }, + Option + { optionLongName = "help", + optionShortName = Just 'h', + optionArgument = NoArgument, + optionDescription = "Print help information and exit.", + optionAction = + [C.cstm|{ + printf("Usage: %s [OPTION]...\nOptions:\n\n%s\nFor more information, consult the Futhark User's Guide or the man pages.\n", + fut_progname, option_descriptions); + exit(0); + }|] + } + ] + where + set_runtime_file = + [C.cstm|{ + runtime_file = fopen(optarg, "w"); + if (runtime_file == NULL) { + futhark_panic(1, "Cannot open %s: %s\n", optarg, strerror(errno)); + } + }|] + set_num_runs = + [C.cstm|{ + num_runs = atoi(optarg); + perform_warmup = 1; + if (num_runs <= 0) { + futhark_panic(1, "Need a positive number of runs, not %s\n", optarg); + } + }|] + +-- | The result of compilation to C is four parts, which can be put +-- together in various ways. The obvious way is to concatenate all of +-- them, which yields a CLI program. Another is to compile the +-- library part by itself, and use the header file to call into it. +data CParts = CParts + { cHeader :: String, + -- | Utility definitions that must be visible + -- to both CLI and library parts. + cUtils :: String, + cCLI :: String, + cLib :: String + } + +-- We may generate variables that are never used (e.g. for +-- certificates) or functions that are never called (e.g. unused +-- intrinsics), and generated code may have other cosmetic issues that +-- compilers warn about. We disable these warnings to not clutter the +-- compilation logs. +disableWarnings :: String +disableWarnings = + pretty + [C.cunit| +$esc:("#ifdef __GNUC__") +$esc:("#pragma GCC diagnostic ignored \"-Wunused-function\"") +$esc:("#pragma GCC diagnostic ignored \"-Wunused-variable\"") +$esc:("#pragma GCC diagnostic ignored \"-Wparentheses\"") +$esc:("#pragma GCC diagnostic ignored \"-Wunused-label\"") +$esc:("#endif") + +$esc:("#ifdef __clang__") +$esc:("#pragma clang diagnostic ignored \"-Wunused-function\"") +$esc:("#pragma clang diagnostic ignored \"-Wunused-variable\"") +$esc:("#pragma clang diagnostic ignored \"-Wparentheses\"") +$esc:("#pragma clang diagnostic ignored \"-Wunused-label\"") +$esc:("#endif") +|] + +-- | Produce header and implementation files. +asLibrary :: CParts -> (String, String) +asLibrary parts = + ( "#pragma once\n\n" <> cHeader parts, + disableWarnings <> cHeader parts <> cUtils parts <> cLib parts + ) + +-- | As executable with command-line interface. +asExecutable :: CParts -> String +asExecutable (CParts a b c d) = disableWarnings <> a <> b <> c <> d + +-- | Compile imperative program to a C program. Always uses the +-- function named "main" as entry point, so make sure it is defined. +compileProg :: + MonadFreshNames m => + String -> + Operations op () -> + CompilerM op () () -> + String -> + [Space] -> + [Option] -> + Definitions op -> + m CParts +compileProg backend ops extra header_extra spaces options prog = do + src <- getNameSource + let ((prototypes, definitions, entry_points), endstate) = + runCompilerM ops src () compileProg' + (entry_point_decls, cli_entry_point_decls, entry_point_inits) = + unzip3 entry_points + option_parser = generateOptionParser "parse_options" $ genericOptions ++ options + + let headerdefs = + [C.cunit| +$esc:("// Headers\n") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:(header_extra) + +$esc:("\n// Initialisation\n") +$edecls:(initDecls endstate) + +$esc:("\n// Arrays\n") +$edecls:(arrayDecls endstate) + +$esc:("\n// Opaque values\n") +$edecls:(opaqueDecls endstate) + +$esc:("\n// Entry points\n") +$edecls:(entryDecls endstate) + +$esc:("\n// Miscellaneous\n") +$edecls:(miscDecls endstate) +$esc:("#define FUTHARK_BACKEND_"++backend) + |] + + let utildefs = + [C.cunit| +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +// If NDEBUG is set, the assert() macro will do nothing. Since Futhark +// (unfortunately) makes use of assert() for error detection (and even some +// side effects), we want to avoid that. +$esc:("#undef NDEBUG") +$esc:("#include ") +$esc:("#include ") + +$esc:util_h + +$esc:timing_h +|] + + let clidefs = + [C.cunit| +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") + +$esc:values_h + +$esc:("#define __private") + +static int binary_output = 0; +static typename FILE *runtime_file; +static int perform_warmup = 0; +static int num_runs = 1; +// If the entry point is NULL, the program will terminate after doing initialisation and such. +static const char *entry_point = "main"; + +$esc:tuning_h + +$func:option_parser + +$edecls:cli_entry_point_decls + +typedef void entry_point_fun(struct futhark_context*); + +struct entry_point_entry { + const char *name; + entry_point_fun *fun; +}; + +int main(int argc, char** argv) { + fut_progname = argv[0]; + + struct entry_point_entry entry_points[] = { + $inits:entry_point_inits + }; + + struct futhark_context_config *cfg = futhark_context_config_new(); + assert(cfg != NULL); + + int parsed_options = parse_options(cfg, argc, argv); + argc -= parsed_options; + argv += parsed_options; + + if (argc != 0) { + futhark_panic(1, "Excess non-option: %s\n", argv[0]); + } + + struct futhark_context *ctx = futhark_context_new(cfg); + assert (ctx != NULL); + + char* error = futhark_context_get_error(ctx); + if (error != NULL) { + futhark_panic(1, "%s", error); + } + + if (entry_point != NULL) { + int num_entry_points = sizeof(entry_points) / sizeof(entry_points[0]); + entry_point_fun *entry_point_fun = NULL; + for (int i = 0; i < num_entry_points; i++) { + if (strcmp(entry_points[i].name, entry_point) == 0) { + entry_point_fun = entry_points[i].fun; + break; + } + } + + if (entry_point_fun == NULL) { + fprintf(stderr, "No entry point '%s'. Select another with --entry-point. Options are:\n", + entry_point); + for (int i = 0; i < num_entry_points; i++) { + fprintf(stderr, "%s\n", entry_points[i].name); + } + return 1; + } + + entry_point_fun(ctx); + + if (runtime_file != NULL) { + fclose(runtime_file); + } + + char *report = futhark_context_report(ctx); + fputs(report, stderr); + free(report); + } + + futhark_context_free(ctx); + futhark_context_config_free(cfg); + return 0; +} + |] + + let early_decls = DL.toList $ compEarlyDecls endstate + let lib_decls = DL.toList $ compLibDecls endstate + let libdefs = + [C.cunit| +$esc:("#ifdef _MSC_VER\n#define inline __inline\n#endif") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") + +$esc:(header_extra) + +$esc:lock_h + +$edecls:builtin + +$edecls:early_decls + +$edecls:prototypes + +$edecls:lib_decls + +$edecls:(map funcToDef definitions) + +$edecls:(arrayDefinitions endstate) + +$edecls:(opaqueDefinitions endstate) + +$edecls:entry_point_decls + |] + + return $ CParts (pretty headerdefs) (pretty utildefs) (pretty clidefs) (pretty libdefs) + where + compileProg' = do + let Definitions consts (Functions funs) = prog + + (memstructs, memfuns, memreport) <- unzip3 <$> mapM defineMemorySpace spaces + + get_consts <- compileConstants consts + + ctx_ty <- contextType + + (prototypes, definitions) <- + unzip <$> mapM (compileFun get_consts [[C.cparam|$ty:ctx_ty *ctx|]]) funs + + mapM_ earlyDecl memstructs + entry_points <- + mapM (uncurry onEntryPoint) $ filter (functionEntry . snd) funs + + extra + + mapM_ earlyDecl $ concat memfuns + + commonLibFuns memreport + + return (prototypes, definitions, entry_points) + + funcToDef func = C.FuncDef func loc + where + loc = case func of + C.OldFunc _ _ _ _ _ _ l -> l + C.Func _ _ _ _ _ l -> l + + builtin = + cIntOps ++ cFloat32Ops ++ cFloat64Ops ++ cFloatConvOps + ++ cFloat32Funs + ++ cFloat64Funs + + util_h = $(embedStringFile "rts/c/util.h") + values_h = $(embedStringFile "rts/c/values.h") + timing_h = $(embedStringFile "rts/c/timing.h") + lock_h = $(embedStringFile "rts/c/lock.h") + tuning_h = $(embedStringFile "rts/c/tuning.h") + +commonLibFuns :: [C.BlockItem] -> CompilerM op s () +commonLibFuns memreport = do + ctx <- contextType + profilereport <- gets $ DL.toList . compProfileItems + + publicDef_ "context_report" MiscDecl $ \s -> + ( [C.cedecl|char* $id:s($ty:ctx *ctx);|], + [C.cedecl|char* $id:s($ty:ctx *ctx) { + struct str_builder builder; + str_builder_init(&builder); + if (ctx->detail_memory || ctx->profiling) { + $items:memreport + } + if (ctx->profiling) { + $items:profilereport + } + return builder.str; + }|] + ) + + publicDef_ "context_get_error" MiscDecl $ \s -> + ( [C.cedecl|char* $id:s($ty:ctx* ctx);|], + [C.cedecl|char* $id:s($ty:ctx* ctx) { + char* error = ctx->error; + ctx->error = NULL; + return error; + }|] + ) + + publicDef_ "context_pause_profiling" MiscDecl $ \s -> + ( [C.cedecl|void $id:s($ty:ctx* ctx);|], + [C.cedecl|void $id:s($ty:ctx* ctx) { + ctx->profiling_paused = 1; + }|] + ) + + publicDef_ "context_unpause_profiling" MiscDecl $ \s -> + ( [C.cedecl|void $id:s($ty:ctx* ctx);|], + [C.cedecl|void $id:s($ty:ctx* ctx) { + ctx->profiling_paused = 0; + }|] + ) + +compileConstants :: Constants op -> CompilerM op s [C.BlockItem] +compileConstants (Constants ps init_consts) = do + ctx_ty <- contextType + const_fields <- mapM constParamField ps + -- Avoid an empty struct, as that is apparently undefined behaviour. + let const_fields' + | null const_fields = [[C.csdecl|int dummy;|]] + | otherwise = const_fields + contextField "constants" [C.cty|struct { $sdecls:const_fields' }|] Nothing + earlyDecl [C.cedecl|int init_constants($ty:ctx_ty*);|] + earlyDecl [C.cedecl|int free_constants($ty:ctx_ty*);|] + + -- We locally define macros for the constants, so that when we + -- generate assignments to local variables, we actually assign into + -- the constants struct. This is not needed for functions, because + -- they can only read constants, not write them. + let (defs, undefs) = unzip $ map constMacro ps + init_consts' <- blockScope $ do + mapM_ resetMemConst ps + compileCode init_consts + libDecl + [C.cedecl|int init_constants($ty:ctx_ty *ctx) { + (void)ctx; + int err = 0; + $items:defs + $items:init_consts' + $items:undefs + cleanup: + return err; + }|] + + free_consts <- collect $ mapM_ freeConst ps + libDecl + [C.cedecl|int free_constants($ty:ctx_ty *ctx) { + (void)ctx; + $items:free_consts + return 0; + }|] + + mapM getConst ps + where + constParamField (ScalarParam name bt) = do + let ctp = primTypeToCType bt + return [C.csdecl|$ty:ctp $id:name;|] + constParamField (MemParam name space) = do + ty <- memToCType name space + return [C.csdecl|$ty:ty $id:name;|] + + constMacro p = ([C.citem|$escstm:def|], [C.citem|$escstm:undef|]) + where + p' = pretty (C.toIdent (paramName p) mempty) + def = "#define " ++ p' ++ " (" ++ "ctx->constants." ++ p' ++ ")" + undef = "#undef " ++ p' + + resetMemConst ScalarParam {} = return () + resetMemConst (MemParam name space) = resetMem name space + + freeConst ScalarParam {} = return () + freeConst (MemParam name space) = unRefMem [C.cexp|ctx->constants.$id:name|] space + + getConst (ScalarParam name bt) = do + let ctp = primTypeToCType bt + return [C.citem|$ty:ctp $id:name = ctx->constants.$id:name;|] + getConst (MemParam name space) = do + ty <- memToCType name space + return [C.citem|$ty:ty $id:name = ctx->constants.$id:name;|] + +cachingMemory :: + M.Map VName Space -> + ([C.BlockItem] -> [C.Stm] -> CompilerM op s a) -> + CompilerM op s a +cachingMemory lexical f = do + -- We only consider lexical 'DefaultSpace' memory blocks to be + -- cached. This is not a deep technical restriction, but merely a + -- heuristic based on GPU memory usually involving larger + -- allocations, that do not suffer from the overhead of reference + -- counting. + let cached = M.keys $ M.filter (== DefaultSpace) lexical + + cached' <- forM cached $ \mem -> do + size <- newVName $ pretty mem <> "_cached_size" + return (mem, size) + + let lexMem env = + env + { envCachedMem = + M.fromList (map (first (`C.toExp` noLoc)) cached') + <> envCachedMem env + } + + declCached (mem, size) = + [ [C.citem|size_t $id:size = 0;|], + [C.citem|$ty:defaultMemBlockType $id:mem = NULL;|] + ] + + freeCached (mem, _) = + [C.cstm|free($id:mem);|] + + local lexMem $ f (concatMap declCached cached') (map freeCached cached') + +compileFun :: [C.BlockItem] -> [C.Param] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func) +compileFun get_constants extra (fname, func@(Function _ outputs inputs body _ _)) = do + (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs + inparams <- mapM compileInput inputs + + cachingMemory (lexicalMemoryUsage func) $ \decl_cached free_cached -> do + body' <- blockScope $ compileFunBody out_ptrs outputs body + + return + ( [C.cedecl|static int $id:(funName fname)($params:extra, $params:outparams, $params:inparams);|], + [C.cfun|static int $id:(funName fname)($params:extra, $params:outparams, $params:inparams) { + $stms:ignores + int err = 0; + $items:decl_cached + $items:get_constants + $items:body' + cleanup: + {} + $stms:free_cached + return err; + }|] + ) + where + -- Ignore all the boilerplate parameters, just in case we don't + -- actually need to use them. + ignores = [[C.cstm|(void)$id:p;|] | C.Param (Just p) _ _ _ <- extra] + + compileInput (ScalarParam name bt) = do + let ctp = primTypeToCType bt + return [C.cparam|$ty:ctp $id:name|] + compileInput (MemParam name space) = do + ty <- memToCType name space + return [C.cparam|$ty:ty $id:name|] + + compileOutput (ScalarParam name bt) = do + let ctp = primTypeToCType bt + p_name <- newVName $ "out_" ++ baseString name + return ([C.cparam|$ty:ctp *$id:p_name|], [C.cexp|$id:p_name|]) + compileOutput (MemParam name space) = do + ty <- memToCType name space + p_name <- newVName $ baseString name ++ "_p" + return ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|]) + +compilePrimValue :: PrimValue -> C.Exp +compilePrimValue (IntValue (Int8Value k)) = [C.cexp|$int:k|] +compilePrimValue (IntValue (Int16Value k)) = [C.cexp|$int:k|] +compilePrimValue (IntValue (Int32Value k)) = [C.cexp|$int:k|] +compilePrimValue (IntValue (Int64Value k)) = [C.cexp|$int:k|] +compilePrimValue (FloatValue (Float64Value x)) + | isInfinite x = + if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|] + | isNaN x = + [C.cexp|NAN|] + | otherwise = + [C.cexp|$double:x|] +compilePrimValue (FloatValue (Float32Value x)) + | isInfinite x = + if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|] + | isNaN x = + [C.cexp|NAN|] + | otherwise = + [C.cexp|$float:x|] +compilePrimValue (BoolValue b) = + [C.cexp|$int:b'|] + where + b' :: Int + b' = if b then 1 else 0 +compilePrimValue Checked = + [C.cexp|0|] + +derefPointer :: C.Exp -> C.Exp -> C.Type -> C.Exp +derefPointer ptr i res_t = + [C.cexp|(($ty:res_t)$exp:ptr)[$exp:i]|] + +volQuals :: Volatility -> [C.TypeQual] +volQuals Volatile = [C.ctyquals|volatile|] +volQuals Nonvolatile = [] + +writeScalarPointerWithQuals :: PointerQuals op s -> WriteScalar op s +writeScalarPointerWithQuals quals_f dest i elemtype space vol v = do + quals <- quals_f space + let quals' = volQuals vol ++ quals + deref = + derefPointer + dest + i + [C.cty|$tyquals:quals' $ty:elemtype*|] + stm [C.cstm|$exp:deref = $exp:v;|] + +readScalarPointerWithQuals :: PointerQuals op s -> ReadScalar op s +readScalarPointerWithQuals quals_f dest i elemtype space vol = do + quals <- quals_f space + let quals' = volQuals vol ++ quals + return $ derefPointer dest i [C.cty|$tyquals:quals' $ty:elemtype*|] + +compileExpToName :: String -> PrimType -> Exp -> CompilerM op s VName +compileExpToName _ _ (LeafExp (ScalarVar v) _) = + return v +compileExpToName desc t e = do + desc' <- newVName desc + e' <- compileExp e + decl [C.cdecl|$ty:(primTypeToCType t) $id:desc' = $e';|] + return desc' + +compileExp :: Exp -> CompilerM op s C.Exp +compileExp = compilePrimExp compileLeaf + where + compileLeaf (ScalarVar src) = + return [C.cexp|$id:src|] + compileLeaf (Index src (Count iexp) restype DefaultSpace vol) = do + src' <- rawMem src + derefPointer src' + <$> compileExp (untyped iexp) + <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primTypeToCType restype)*|] + compileLeaf (Index src (Count iexp) restype (Space space) vol) = + join $ + asks envReadScalar + <*> rawMem src + <*> compileExp (untyped iexp) + <*> pure (primTypeToCType restype) + <*> pure space + <*> pure vol + compileLeaf (Index src (Count iexp) _ ScalarSpace {} _) = do + iexp' <- compileExp $ untyped iexp + return [C.cexp|$id:src[$exp:iexp']|] + compileLeaf (SizeOf t) = + return [C.cexp|(typename int64_t)sizeof($ty:t')|] + where + t' = primTypeToCType t + +-- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you. +compilePrimExp :: Monad m => (v -> m C.Exp) -> PrimExp v -> m C.Exp +compilePrimExp _ (ValueExp val) = + return $ compilePrimValue val +compilePrimExp f (LeafExp v _) = + f v +compilePrimExp f (UnOpExp Complement {} x) = do + x' <- compilePrimExp f x + return [C.cexp|~$exp:x'|] +compilePrimExp f (UnOpExp Not {} x) = do + x' <- compilePrimExp f x + return [C.cexp|!$exp:x'|] +compilePrimExp f (UnOpExp Abs {} x) = do + x' <- compilePrimExp f x + return [C.cexp|abs($exp:x')|] +compilePrimExp f (UnOpExp (FAbs Float32) x) = do + x' <- compilePrimExp f x + return [C.cexp|(float)fabs($exp:x')|] +compilePrimExp f (UnOpExp (FAbs Float64) x) = do + x' <- compilePrimExp f x + return [C.cexp|fabs($exp:x')|] +compilePrimExp f (UnOpExp SSignum {} x) = do + x' <- compilePrimExp f x + return [C.cexp|($exp:x' > 0) - ($exp:x' < 0)|] +compilePrimExp f (UnOpExp USignum {} x) = do + x' <- compilePrimExp f x + return [C.cexp|($exp:x' > 0) - ($exp:x' < 0) != 0|] +compilePrimExp f (CmpOpExp cmp x y) = do + x' <- compilePrimExp f x + y' <- compilePrimExp f y + return $ case cmp of + CmpEq {} -> [C.cexp|$exp:x' == $exp:y'|] + FCmpLt {} -> [C.cexp|$exp:x' < $exp:y'|] + FCmpLe {} -> [C.cexp|$exp:x' <= $exp:y'|] + CmpLlt {} -> [C.cexp|$exp:x' < $exp:y'|] + CmpLle {} -> [C.cexp|$exp:x' <= $exp:y'|] + _ -> [C.cexp|$id:(pretty cmp)($exp:x', $exp:y')|] +compilePrimExp f (ConvOpExp conv x) = do + x' <- compilePrimExp f x + return [C.cexp|$id:(pretty conv)($exp:x')|] +compilePrimExp f (BinOpExp bop x y) = do + x' <- compilePrimExp f x + y' <- compilePrimExp f y + -- Note that integer addition, subtraction, and multiplication with + -- OverflowWrap are not handled by explicit operators, but rather by + -- functions. This is because we want to implicitly convert them to + -- unsigned numbers, so we can do overflow without invoking + -- undefined behaviour. + return $ case bop of + Add _ OverflowUndef -> [C.cexp|$exp:x' + $exp:y'|] + Sub _ OverflowUndef -> [C.cexp|$exp:x' - $exp:y'|] + Mul _ OverflowUndef -> [C.cexp|$exp:x' * $exp:y'|] + FAdd {} -> [C.cexp|$exp:x' + $exp:y'|] + FSub {} -> [C.cexp|$exp:x' - $exp:y'|] + FMul {} -> [C.cexp|$exp:x' * $exp:y'|] + FDiv {} -> [C.cexp|$exp:x' / $exp:y'|] + Xor {} -> [C.cexp|$exp:x' ^ $exp:y'|] + And {} -> [C.cexp|$exp:x' & $exp:y'|] + Or {} -> [C.cexp|$exp:x' | $exp:y'|] + Shl {} -> [C.cexp|$exp:x' << $exp:y'|] + LogAnd {} -> [C.cexp|$exp:x' && $exp:y'|] + LogOr {} -> [C.cexp|$exp:x' || $exp:y'|] + _ -> [C.cexp|$id:(pretty bop)($exp:x', $exp:y')|] +compilePrimExp f (FunExp h args _) = do + args' <- mapM (compilePrimExp f) args + return [C.cexp|$id:(funName (nameFromString h))($args:args')|] + +compileCode :: Code op -> CompilerM op s () +compileCode (Op op) = + join $ asks envOpCompiler <*> pure op +compileCode Skip = return () +compileCode (Comment s code) = do + xs <- blockScope $ compileCode code + let comment = "// " ++ s + stm + [C.cstm|$comment:comment + { $items:xs } + |] +compileCode (DebugPrint s (Just e)) = do + e' <- compileExp e + stm + [C.cstm|if (ctx->debugging) { + fprintf(stderr, $string:fmtstr, $exp:s, ($ty:ety)$exp:e', '\n'); + }|] + where + (fmt, ety) = case primExpType e of + IntType _ -> ("llu", [C.cty|long long int|]) + FloatType _ -> ("f", [C.cty|double|]) + _ -> ("d", [C.cty|int|]) + fmtstr = "%s: %" ++ fmt ++ "%c" +compileCode (DebugPrint s Nothing) = + stm + [C.cstm|if (ctx->debugging) { + fprintf(stderr, "%s\n", $exp:s); + }|] +compileCode c + | Just (name, vol, t, e, c') <- declareAndSet c = do + let ct = primTypeToCType t + e' <- compileExp e + item [C.citem|$tyquals:(volQuals vol) $ty:ct $id:name = $exp:e';|] + compileCode c' +compileCode (c1 :>>: c2) = compileCode c1 >> compileCode c2 +compileCode (Assert e msg (loc, locs)) = do + e' <- compileExp e + err <- + collect $ + join $ + asks (opsError . envOperations) <*> pure msg <*> pure stacktrace + stm [C.cstm|if (!$exp:e') { $items:err }|] + where + stacktrace = prettyStacktrace 0 $ map locStr $ loc : locs +compileCode (Allocate _ _ ScalarSpace {}) = + -- Handled by the declaration of the memory block, which is + -- translated to an actual array. + return () +compileCode (Allocate name (Count (TPrimExp e)) space) = do + size <- compileExp e + cached <- cacheMem name + case cached of + Just cur_size -> + stm + [C.cstm|if ($exp:cur_size < (size_t)$exp:size) { + $exp:name = realloc($exp:name, $exp:size); + $exp:cur_size = $exp:size; + }|] + _ -> + allocMem name size space [C.cstm|{err = 1; goto cleanup;}|] +compileCode (Free name space) = do + cached <- isJust <$> cacheMem name + unless cached $ unRefMem name space +compileCode (For i bound body) = do + let i' = C.toIdent i + t = primTypeToCType $ primExpType bound + bound' <- compileExp bound + body' <- blockScope $ compileCode body + stm + [C.cstm|for ($ty:t $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) { + $items:body' + }|] +compileCode (While cond body) = do + cond' <- compileExp $ untyped cond + body' <- blockScope $ compileCode body + stm + [C.cstm|while ($exp:cond') { + $items:body' + }|] +compileCode (If cond tbranch fbranch) = do + cond' <- compileExp $ untyped cond + tbranch' <- blockScope $ compileCode tbranch + fbranch' <- blockScope $ compileCode fbranch + stm $ case (tbranch', fbranch') of + (_, []) -> + [C.cstm|if ($exp:cond') { $items:tbranch' }|] + ([], _) -> + [C.cstm|if (!($exp:cond')) { $items:fbranch' }|] + _ -> + [C.cstm|if ($exp:cond') { $items:tbranch' } else { $items:fbranch' }|] +compileCode (Copy dest (Count destoffset) DefaultSpace src (Count srcoffset) DefaultSpace (Count size)) = + join $ + copyMemoryDefaultSpace + <$> rawMem dest + <*> compileExp (untyped destoffset) + <*> rawMem src + <*> compileExp (untyped srcoffset) + <*> compileExp (untyped size) +compileCode (Copy dest (Count destoffset) destspace src (Count srcoffset) srcspace (Count size)) = do + copy <- asks envCopy + join $ + copy + <$> rawMem dest + <*> compileExp (untyped destoffset) + <*> pure destspace + <*> rawMem src + <*> compileExp (untyped srcoffset) + <*> pure srcspace + <*> compileExp (untyped size) +compileCode (Write dest (Count idx) elemtype DefaultSpace vol elemexp) = do + dest' <- rawMem dest + deref <- + derefPointer dest' + <$> compileExp (untyped idx) + <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primTypeToCType elemtype)*|] + elemexp' <- compileExp elemexp + stm [C.cstm|$exp:deref = $exp:elemexp';|] +compileCode (Write dest (Count idx) _ ScalarSpace {} _ elemexp) = do + idx' <- compileExp (untyped idx) + elemexp' <- compileExp elemexp + stm [C.cstm|$id:dest[$exp:idx'] = $exp:elemexp';|] +compileCode (Write dest (Count idx) elemtype (Space space) vol elemexp) = + join $ + asks envWriteScalar + <*> rawMem dest + <*> compileExp (untyped idx) + <*> pure (primTypeToCType elemtype) + <*> pure space + <*> pure vol + <*> compileExp elemexp +compileCode (DeclareMem name space) = + declMem name space +compileCode (DeclareScalar name vol t) = do + let ct = primTypeToCType t + decl [C.cdecl|$tyquals:(volQuals vol) $ty:ct $id:name;|] +compileCode (DeclareArray name ScalarSpace {} _ _) = + error $ "Cannot declare array " ++ pretty name ++ " in scalar space." +compileCode (DeclareArray name DefaultSpace t vs) = do + name_realtype <- newVName $ baseString name ++ "_realtype" + let ct = primTypeToCType t + case vs of + ArrayValues vs' -> do + let vs'' = [[C.cinit|$exp:(compilePrimValue v)|] | v <- vs'] + earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:(length vs')] = {$inits:vs''};|] + ArrayZeros n -> + earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:n];|] + -- Fake a memory block. + contextField + (C.toIdent name noLoc) + [C.cty|struct memblock|] + $ Just [C.cexp|(struct memblock){NULL, (char*)$id:name_realtype, 0}|] + item [C.citem|struct memblock $id:name = ctx->$id:name;|] +compileCode (DeclareArray name (Space space) t vs) = + join $ + asks envStaticArray + <*> pure name + <*> pure space + <*> pure t + <*> pure vs +-- For assignments of the form 'x = x OP e', we generate C assignment +-- operators to make the resulting code slightly nicer. This has no +-- effect on performance. +compileCode (SetScalar dest (BinOpExp op (LeafExp (ScalarVar x) _) y)) + | dest == x, + Just f <- assignmentOperator op = do + y' <- compileExp y + stm [C.cstm|$exp:(f dest y');|] +compileCode (SetScalar dest src) = do + src' <- compileExp src + stm [C.cstm|$id:dest = $exp:src';|] +compileCode (SetMem dest src space) = + setMem dest src space +compileCode (Call dests fname args) = + join $ + asks (opsCall . envOperations) + <*> pure dests + <*> pure fname + <*> mapM compileArg args + where + compileArg (MemArg m) = return [C.cexp|$exp:m|] + compileArg (ExpArg e) = compileExp e + +blockScope :: CompilerM op s () -> CompilerM op s [C.BlockItem] +blockScope = fmap snd . blockScope' + +blockScope' :: CompilerM op s a -> CompilerM op s (a, [C.BlockItem]) +blockScope' m = do + old_allocs <- gets compDeclaredMem + (x, xs) <- pass $ do + (x, w) <- listen m + let xs = DL.toList $ accItems w + return ((x, xs), const mempty) + new_allocs <- gets $ filter (`notElem` old_allocs) . compDeclaredMem + modify $ \s -> s {compDeclaredMem = old_allocs} + releases <- collect $ mapM_ (uncurry unRefMem) new_allocs + return (x, xs <> releases) + +compileFunBody :: [C.Exp] -> [Param] -> Code op -> CompilerM op s () +compileFunBody output_ptrs outputs code = do + mapM_ declareOutput outputs + compileCode code + zipWithM_ setRetVal' output_ptrs outputs + where + declareOutput (MemParam name space) = + declMem name space + declareOutput (ScalarParam name pt) = do + let ctp = primTypeToCType pt + decl [C.cdecl|$ty:ctp $id:name;|] + + setRetVal' p (MemParam name space) = do + resetMem [C.cexp|*$exp:p|] space + setMem [C.cexp|*$exp:p|] name space + setRetVal' p (ScalarParam name _) = + stm [C.cstm|*$exp:p = $id:name;|] + +declareAndSet :: Code op -> Maybe (VName, Volatility, PrimType, Exp, Code op) +declareAndSet code = do + (DeclareScalar name vol t, code') <- nextCode code + (SetScalar dest e, code'') <- nextCode code' + guard $ name == dest + Just (name, vol, t, e, code'') + +nextCode :: Code op -> Maybe (Code op, Code op) +nextCode (x :>>: y) + | Just (x_a, x_b) <- nextCode x = + Just (x_a, x_b <> y) + | otherwise = + Just (x, y) +nextCode _ = Nothing + +assignmentOperator :: BinOp -> Maybe (VName -> C.Exp -> C.Exp) +assignmentOperator Add {} = Just $ \d e -> [C.cexp|$id:d += $exp:e|] +assignmentOperator Sub {} = Just $ \d e -> [C.cexp|$id:d -= $exp:e|] +assignmentOperator Mul {} = Just $ \d e -> [C.cexp|$id:d *= $exp:e|] +assignmentOperator _ = Nothing + +-- | Return an expression multiplying together the given expressions. +-- If an empty list is given, the expression @1@ is returned. +cproduct :: [C.Exp] -> C.Exp +cproduct [] = [C.cexp|1|] +cproduct (e : es) = foldl mult e es + where + mult x y = [C.cexp|$exp:x * $exp:y|] diff --git a/src/Futhark/CodeGen/Backends/GenericC/Options.hs b/src/Futhark/CodeGen/Backends/GenericC/Options.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Options.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Options.hs @@ -1,36 +1,42 @@ {-# LANGUAGE QuasiQuotes #-} + -- | This module defines a generator for @getopt_long@ based command -- line argument parsing. Each option is associated with arbitrary C -- code that will perform side effects, usually by setting some global -- variables. module Futhark.CodeGen.Backends.GenericC.Options - ( Option (..) - , OptionArgument (..) - , generateOptionParser - ) - where + ( Option (..), + OptionArgument (..), + generateOptionParser, + ) +where +import Data.Char (isSpace) +import Data.Function ((&)) +import Data.List (intercalate) import Data.Maybe - -import qualified Language.C.Syntax as C import qualified Language.C.Quote.C as C +import qualified Language.C.Syntax as C -- | Specification if a single command line option. The option must -- have a long name, and may also have a short name. -- -- In the action, the option argument (if any) is stored as in the -- @char*@-typed variable @optarg@. -data Option = Option { optionLongName :: String - , optionShortName :: Maybe Char - , optionArgument :: OptionArgument - , optionAction :: C.Stm - } +data Option = Option + { optionLongName :: String, + optionShortName :: Maybe Char, + optionArgument :: OptionArgument, + optionDescription :: String, + optionAction :: C.Stm + } -- | Whether an option accepts an argument. -data OptionArgument = NoArgument - | RequiredArgument String - -- ^ The 'String' becomes part of the help text. - | OptionalArgument +data OptionArgument + = NoArgument + | -- | The 'String' becomes part of the help text. + RequiredArgument String + | OptionalArgument -- | Generate an option parser as a function of the given name, that -- accepts the given command line options. The result is a function @@ -46,6 +52,8 @@ static struct option long_options[] = { $inits:option_fields, {0, 0, 0, 0} }; + static char* option_descriptions = $string:option_descriptions; + while (($id:chosen_option = getopt_long(argc, argv, $string:option_string, long_options, NULL)) != -1) { $stms:option_applications @@ -60,47 +68,76 @@ return optind; } |] - where chosen_option = "ch" - option_string = ':' : optionString options - option_applications = optionApplications chosen_option options - option_fields = optionFields options - option_descriptions = unwords $ map describeOption options + where + chosen_option = "ch" + option_string = ':' : optionString options + option_applications = optionApplications chosen_option options + option_fields = optionFields options + option_descriptions = describeOptions options -describeOption :: Option -> String -describeOption opt = - concat [ "[" - , maybe "" (\c -> "-" ++ [c] ++ "/") $ optionShortName opt - , "--" ++ optionLongName opt - , case optionArgument opt of - NoArgument -> "" - RequiredArgument what -> " " ++ what - OptionalArgument -> " [ARG]" - , "]" - ] +trim :: String -> String +trim = f . f + where + f = reverse . dropWhile isSpace +describeOptions :: [Option] -> String +describeOptions opts = + let + in unlines $ fmap extendDesc with_short_descs + where + with_short_descs = fmap (\opt -> (opt, shortDesc opt)) opts + max_short_desc_len = maximum $ fmap (length . snd) with_short_descs + extendDesc :: (Option, String) -> String + extendDesc (opt, short) = + take (max_short_desc_len + 1) (short ++ repeat ' ') + ++ ( optionDescription opt + & lines + & fmap trim + & intercalate ('\n' : replicate (max_short_desc_len + 1) ' ') + ) + shortDesc :: Option -> String + shortDesc opt = + concat + [ " ", + maybe "" (\c -> "-" ++ [c] ++ "/") $ optionShortName opt, + "--" ++ optionLongName opt, + case optionArgument opt of + NoArgument -> "" + RequiredArgument what -> " " ++ what + OptionalArgument -> " [ARG]" + ] + optionFields :: [Option] -> [C.Initializer] -optionFields = zipWith field [(1::Int)..] - where field i option = - [C.cinit| { $string:(optionLongName option), $id:arg, NULL, $int:i } |] - where arg = case optionArgument option of - NoArgument -> "no_argument" - RequiredArgument _ -> "required_argument" - OptionalArgument -> "optional_argument" +optionFields = zipWith field [(1 :: Int) ..] + where + field i option = + [C.cinit| { $string:(optionLongName option), $id:arg, NULL, $int:i } |] + where + arg = case optionArgument option of + NoArgument -> "no_argument" + RequiredArgument _ -> "required_argument" + OptionalArgument -> "optional_argument" optionApplications :: String -> [Option] -> [C.Stm] -optionApplications chosen_option = zipWith check [(1::Int)..] - where check i option = - [C.cstm|if ($exp:cond) $stm:(optionAction option)|] - where cond = case optionShortName option of - Nothing -> [C.cexp|$id:chosen_option == $int:i|] - Just c -> [C.cexp|($id:chosen_option == $int:i) || +optionApplications chosen_option = zipWith check [(1 :: Int) ..] + where + check i option = + [C.cstm|if ($exp:cond) $stm:(optionAction option)|] + where + cond = case optionShortName option of + Nothing -> [C.cexp|$id:chosen_option == $int:i|] + Just c -> + [C.cexp|($id:chosen_option == $int:i) || ($id:chosen_option == $char:c)|] + optionString :: [Option] -> String optionString = concat . mapMaybe optionStringChunk - where optionStringChunk option = do - short <- optionShortName option - return $ short : - case optionArgument option of - NoArgument -> "" - RequiredArgument _ -> ":" - OptionalArgument -> "::" + where + optionStringChunk option = do + short <- optionShortName option + return $ + short : + case optionArgument option of + NoArgument -> "" + RequiredArgument _ -> ":" + OptionalArgument -> "::" 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 @@ -1,1034 +1,1213 @@ -{-# LANGUAGE OverloadedStrings, GeneralizedNewtypeDeriving, LambdaCase #-} -{-# LANGUAGE TupleSections #-} --- | A generic Python code generator which is polymorphic in the type --- of the operations. Concretely, we use this to handle both --- sequential and PyOpenCL Python code. -module Futhark.CodeGen.Backends.GenericPython - ( compileProg - , Constructor (..) - , emptyConstructor - - , compileName - , compileVar - , compileDim - , compileExp - , compilePrimExp - , compileCode - , compilePrimValue - , compilePrimType - , compilePrimTypeExt - , compilePrimToNp - , compilePrimToExtNp - - , Operations (..) - , defaultOperations - - , unpackDim - - , CompilerM (..) - , OpCompiler - , WriteScalar - , ReadScalar - , Allocate - , Copy - , StaticArray - , EntryOutput - , EntryInput - - , CompilerEnv(..) - , CompilerState(..) - , stm - , atInit - , collect' - , collect - , simpleCall - - , copyMemoryDefaultSpace - ) where - -import Control.Monad.Identity -import Control.Monad.State -import Control.Monad.Reader -import Control.Monad.Writer -import Control.Monad.RWS -import Data.Maybe -import qualified Data.Map as M - -import Futhark.IR.Primitive hiding (Bool) -import Futhark.MonadFreshNames -import Futhark.IR.Syntax (Space(..)) -import qualified Futhark.CodeGen.ImpCode as Imp -import Futhark.CodeGen.Backends.GenericPython.AST -import Futhark.CodeGen.Backends.GenericPython.Options -import Futhark.CodeGen.Backends.GenericPython.Definitions -import Futhark.Util (zEncodeString) -import Futhark.IR.Prop (isBuiltInFunction) - --- | A substitute expression compiler, tried before the main --- compilation function. -type OpCompiler op s = op -> CompilerM op s () - --- | Write a scalar to the given memory block with the given index and --- in the given memory space. -type WriteScalar op s = PyExp -> PyExp -> PrimType -> Imp.SpaceId -> PyExp - -> CompilerM op s () - --- | Read a scalar from the given memory block with the given index and --- in the given memory space. -type ReadScalar op s = PyExp -> PyExp -> PrimType -> Imp.SpaceId - -> CompilerM op s PyExp - --- | Allocate a memory block of the given size in the given memory --- space, saving a reference in the given variable name. -type Allocate op s = PyExp -> PyExp -> Imp.SpaceId - -> CompilerM op s () - --- | Copy from one memory block to another. -type Copy op s = PyExp -> PyExp -> Imp.Space -> - PyExp -> PyExp -> Imp.Space -> - PyExp -> PrimType -> - CompilerM op s () - --- | Create a static array of values - initialised at load time. -type StaticArray op s = VName -> Imp.SpaceId -> PrimType -> Imp.ArrayContents -> CompilerM op s () - --- | Construct the Python array being returned from an entry point. -type EntryOutput op s = VName -> Imp.SpaceId -> - PrimType -> Imp.Signedness -> - [Imp.DimSize] -> - CompilerM op s PyExp - --- | Unpack the array being passed to an entry point. -type EntryInput op s = PyExp -> Imp.SpaceId -> - PrimType -> Imp.Signedness -> - [Imp.DimSize] -> - PyExp -> - CompilerM op s () - - -data Operations op s = Operations { opsWriteScalar :: WriteScalar op s - , opsReadScalar :: ReadScalar op s - , opsAllocate :: Allocate op s - , opsCopy :: Copy op s - , opsStaticArray :: StaticArray op s - , opsCompiler :: OpCompiler op s - , opsEntryOutput :: EntryOutput op s - , opsEntryInput :: EntryInput op s - } - --- | A set of operations that fail for every operation involving --- non-default memory spaces. Uses plain pointers and @malloc@ for --- memory management. -defaultOperations :: Operations op s -defaultOperations = Operations { opsWriteScalar = defWriteScalar - , opsReadScalar = defReadScalar - , opsAllocate = defAllocate - , opsCopy = defCopy - , opsStaticArray = defStaticArray - , opsCompiler = defCompiler - , opsEntryOutput = defEntryOutput - , opsEntryInput = defEntryInput - } - where defWriteScalar _ _ _ _ _ = - error "Cannot write to non-default memory space because I am dumb" - defReadScalar _ _ _ _ = - error "Cannot read from non-default memory space" - defAllocate _ _ _ = - error "Cannot allocate in non-default memory space" - defCopy _ _ _ _ _ _ _ _ = - error "Cannot copy to or from non-default memory space" - defStaticArray _ _ _ _ = - error "Cannot create static array in non-default memory space" - defCompiler _ = - error "The default compiler cannot compile extended operations" - defEntryOutput _ _ _ _ = - error "Cannot return array not in default memory space" - defEntryInput _ _ _ _ = - error "Cannot accept array not in default memory space" - -data CompilerEnv op s = CompilerEnv - { envOperations :: Operations op s - , envVarExp :: M.Map VName PyExp - } - -envOpCompiler :: CompilerEnv op s -> OpCompiler op s -envOpCompiler = opsCompiler . envOperations - -envReadScalar :: CompilerEnv op s -> ReadScalar op s -envReadScalar = opsReadScalar . envOperations - -envWriteScalar :: CompilerEnv op s -> WriteScalar op s -envWriteScalar = opsWriteScalar . envOperations - -envAllocate :: CompilerEnv op s -> Allocate op s -envAllocate = opsAllocate . envOperations - -envCopy :: CompilerEnv op s -> Copy op s -envCopy = opsCopy . envOperations - -envStaticArray :: CompilerEnv op s -> StaticArray op s -envStaticArray = opsStaticArray . envOperations - -envEntryOutput :: CompilerEnv op s -> EntryOutput op s -envEntryOutput = opsEntryOutput . envOperations - -envEntryInput :: CompilerEnv op s -> EntryInput op s -envEntryInput = opsEntryInput . envOperations - -newCompilerEnv :: Operations op s -> CompilerEnv op s -newCompilerEnv ops = CompilerEnv { envOperations = ops - , envVarExp = mempty } - -data CompilerState s = CompilerState { - compNameSrc :: VNameSource - , compInit :: [PyStmt] - , compUserState :: s -} - -newCompilerState :: VNameSource -> s -> CompilerState s -newCompilerState src s = CompilerState { compNameSrc = src - , compInit = [] - , compUserState = s } - -newtype CompilerM op s a = CompilerM (RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a) - deriving (Functor, Applicative, Monad, - MonadState (CompilerState s), - MonadReader (CompilerEnv op s), - MonadWriter [PyStmt]) - -instance MonadFreshNames (CompilerM op s) where - getNameSource = gets compNameSrc - putNameSource src = modify $ \s -> s { compNameSrc = src } - -collect :: CompilerM op s () -> CompilerM op s [PyStmt] -collect m = pass $ do - ((), w) <- listen m - return (w, const mempty) - -collect' :: CompilerM op s a -> CompilerM op s (a, [PyStmt]) -collect' m = pass $ do - (x, w) <- listen m - return ((x, w), const mempty) - -atInit :: PyStmt -> CompilerM op s () -atInit x = modify $ \s -> - s { compInit = compInit s ++ [x] } - -stm :: PyStmt -> CompilerM op s () -stm x = tell [x] - -futharkFun :: String -> String -futharkFun s = "futhark_" ++ zEncodeString s - -compileOutput :: [Imp.Param] -> [PyExp] -compileOutput = map (Var . compileName . Imp.paramName) - -runCompilerM :: Operations op s - -> VNameSource - -> s - -> CompilerM op s a - -> a -runCompilerM ops src userstate (CompilerM m) = - fst $ evalRWS m (newCompilerEnv ops) (newCompilerState src userstate) - -standardOptions :: [Option] -standardOptions = [ - Option { optionLongName = "write-runtime-to" - , optionShortName = Just 't' - , optionArgument = RequiredArgument "str" - , optionAction = - [ - If (Var "runtime_file") - [Exp $ simpleCall "runtime_file.close" []] [] - , Assign (Var "runtime_file") $ - simpleCall "open" [Var "optarg", String "w"] - ] - }, - Option { optionLongName = "runs" - , optionShortName = Just 'r' - , optionArgument = RequiredArgument "str" - , optionAction = - [ Assign (Var "num_runs") $ Var "optarg" - , Assign (Var "do_warmup_run") $ Bool True - ] - }, - Option { optionLongName = "entry-point" - , optionShortName = Just 'e' - , optionArgument = RequiredArgument "str" - , optionAction = - [ Assign (Var "entry_point") $ Var "optarg" ] - }, - Option { optionLongName = "binary-output" - , optionShortName = Just 'b' - , optionArgument = NoArgument - , optionAction = [Assign (Var "binary_output") $ Bool True] - }, - Option { optionLongName = "tuning" - , optionShortName = Nothing - , optionArgument = RequiredArgument "open" - , optionAction = [Exp $ simpleCall "read_tuning_file" [Var "sizes", Var "optarg"]] - } - ] - - --- | The class generated by the code generator must have a --- constructor, although it can be vacuous. -data Constructor = Constructor [String] [PyStmt] - --- | A constructor that takes no arguments and does nothing. -emptyConstructor :: Constructor -emptyConstructor = Constructor ["self"] [Pass] - -constructorToFunDef :: Constructor -> [PyStmt] -> PyFunDef -constructorToFunDef (Constructor params body) at_init = - Def "__init__" params $ body <> at_init - -compileProg :: MonadFreshNames m => - Maybe String - -> Constructor - -> [PyStmt] - -> [PyStmt] - -> Operations op s - -> s - -> [PyStmt] - -> [Option] - -> Imp.Definitions op - -> m String -compileProg module_name constructor imports defines ops userstate sync options prog = do - src <- getNameSource - let prog' = runCompilerM ops src userstate compileProg' - maybe_shebang = - case module_name of Nothing -> "#!/usr/bin/env python\n" - Just _ -> "" - return $ maybe_shebang ++ - pretty (PyProg $ imports ++ - [ Import "argparse" Nothing - , Assign (Var "sizes") $ Dict [] - ] ++ - defines ++ - [Escape pyUtility] ++ - prog') - where Imp.Definitions consts (Imp.Functions funs) = prog - compileProg' = withConstantSubsts consts $ do - - compileConstants consts - - definitions <- mapM compileFunc funs - at_inits <- gets compInit - - let constructor' = constructorToFunDef constructor at_inits - - case module_name of - Just name -> do - (entry_points, entry_point_types) <- - unzip <$> mapM (compileEntryFun sync) (filter (Imp.functionEntry . snd) funs) - return [ClassDef $ Class name $ - Assign (Var "entry_points") (Dict entry_point_types) : - map FunDef (constructor' : definitions ++ entry_points)] - Nothing -> do - let classinst = Assign (Var "self") $ simpleCall "internal" [] - (entry_point_defs, entry_point_names, entry_points) <- - unzip3 <$> mapM (callEntryFun sync) - (filter (Imp.functionEntry . snd) funs) - return (parse_options ++ - ClassDef (Class "internal" $ map FunDef $ - constructor' : definitions) : - classinst : - map FunDef entry_point_defs ++ - selectEntryPoint entry_point_names entry_points) - - parse_options = - Assign (Var "runtime_file") None : - Assign (Var "do_warmup_run") (Bool False) : - Assign (Var "num_runs") (Integer 1) : - Assign (Var "entry_point") (String "main") : - Assign (Var "binary_output") (Bool False) : - generateOptionParser (standardOptions ++ options) - - selectEntryPoint entry_point_names entry_points = - [ Assign (Var "entry_points") $ - Dict $ zip (map String entry_point_names) entry_points, - Assign (Var "entry_point_fun") $ - simpleCall "entry_points.get" [Var "entry_point"], - If (BinOp "==" (Var "entry_point_fun") None) - [Exp $ simpleCall "sys.exit" - [Call (Field - (String "No entry point '{}'. Select another with --entry point. Options are:\n{}") - "format") - [Arg $ Var "entry_point", - - Arg $ Call (Field (String "\n") "join") - [Arg $ simpleCall "entry_points.keys" []]]]] - [Exp $ simpleCall "entry_point_fun" []] - ] - -withConstantSubsts :: Imp.Constants op -> CompilerM op s a -> CompilerM op s a -withConstantSubsts (Imp.Constants ps _) = - local $ \env -> env { envVarExp = foldMap constExp ps } - where constExp p = - M.singleton (Imp.paramName p) $ Index (Var "self.constants") $ - IdxExp $ String $ pretty $ Imp.paramName p - -compileConstants :: Imp.Constants op -> CompilerM op s () -compileConstants (Imp.Constants _ init_consts) = do - atInit $ Assign (Var "self.constants") $ Dict [] - mapM_ atInit =<< collect (compileCode init_consts) - -compileFunc :: (Name, Imp.Function op) -> CompilerM op s PyFunDef -compileFunc (fname, Imp.Function _ outputs inputs body _ _) = do - body' <- collect $ compileCode body - let inputs' = map (compileName . Imp.paramName) inputs - let ret = Return $ tupleOrSingle $ compileOutput outputs - return $ Def (futharkFun . nameToString $ fname) ("self" : inputs') $ - body'++[ret] - -tupleOrSingle :: [PyExp] -> PyExp -tupleOrSingle [e] = e -tupleOrSingle es = Tuple es - --- | A 'Call' where the function is a variable and every argument is a --- simple 'Arg'. -simpleCall :: String -> [PyExp] -> PyExp -simpleCall fname = Call (Var fname) . map Arg - -compileName :: VName -> String -compileName = zEncodeString . pretty - -compileDim :: Imp.DimSize -> PyExp -compileDim (Imp.Constant v) = compilePrimValue v -compileDim (Imp.Var v) = Var $ compileName v - -unpackDim :: PyExp -> Imp.DimSize -> Int32 -> CompilerM op s () -unpackDim arr_name (Imp.Constant c) i = do - let shape_name = Field arr_name "shape" - let constant_c = compilePrimValue c - let constant_i = Integer $ toInteger i - stm $ Assert (BinOp "==" constant_c (Index shape_name $ IdxExp constant_i)) $ - String "constant dimension wrong" -unpackDim arr_name (Imp.Var var) i = do - let shape_name = Field arr_name "shape" - src = Index shape_name $ IdxExp $ Integer $ toInteger i - var' <- compileVar var - stm $ Assign var' $ simpleCall "np.int32" [src] - -entryPointOutput :: Imp.ExternalValue -> CompilerM op s PyExp -entryPointOutput (Imp.OpaqueValue desc vs) = - simpleCall "opaque" . (String (pretty desc):) <$> - mapM (entryPointOutput . Imp.TransparentValue) vs -entryPointOutput (Imp.TransparentValue (Imp.ScalarValue bt ept name)) = do - name' <- compileVar name - return $ simpleCall tf [name'] - where tf = compilePrimToExtNp bt ept -entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims)) = do - pack_output <- asks envEntryOutput - pack_output mem sid bt ept dims -entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ bt ept dims)) = do - mem' <- compileVar mem - let cast = Cast mem' (compilePrimTypeExt bt ept) - return $ simpleCall "createArray" [cast, Tuple $ map compileDim dims] - -badInput :: Int -> PyExp -> String -> PyStmt -badInput i e t = - Raise $ simpleCall "TypeError" - [Call (Field (String err_msg) "format") - [Arg (String t), Arg $ simpleCall "type" [e], Arg e]] - where err_msg = unlines [ "Argument #" ++ show i ++ " has invalid value" - , "Futhark type: {}" - , "Argument has Python type {} and value: {}"] - -badInputType :: Int -> PyExp -> String -> PyExp -> PyExp -> PyStmt -badInputType i e t de dg = - Raise $ simpleCall "TypeError" - [Call (Field (String err_msg) "format") - [Arg (String t), Arg $ simpleCall "type" [e], Arg e, Arg de, Arg dg]] - where err_msg = unlines [ "Argument #" ++ show i ++ " has invalid value" - , "Futhark type: {}" - , "Argument has Python type {} and value: {}" - , "Expected array with elements of dtype: {}" - , "The array given has elements of dtype: {}"] - -badInputDim :: Int -> PyExp -> String -> Int -> PyStmt -badInputDim i e typ dimf = - Raise $ simpleCall "TypeError" - [Call (Field (String err_msg) "format") - [Arg eft, Arg aft]] - where eft = String (concat (replicate dimf "[]") ++ typ) - aft = BinOp "+" (BinOp "*" (String "[]") (Field e "ndim")) (String typ) - err_msg = unlines [ "Argument #" ++ show i ++ " has invalid value" - , "Dimensionality mismatch" - , "Expected Futhark type: {}" - , "Bad Python value passed" - , "Actual Futhark type: {}"] - -entryPointInput :: (Int, Imp.ExternalValue, PyExp) -> CompilerM op s () -entryPointInput (i, Imp.OpaqueValue desc vs, e) = do - let type_is_ok = BinOp "and" (simpleCall "isinstance" [e, Var "opaque"]) - (BinOp "==" (Field e "desc") (String desc)) - stm $ If (UnOp "not" type_is_ok) [badInput i e desc] [] - mapM_ entryPointInput $ zip3 (repeat i) (map Imp.TransparentValue vs) $ - map (Index (Field e "data") . IdxExp . Integer) [0..] - -entryPointInput (i, Imp.TransparentValue (Imp.ScalarValue bt s name), e) = do - vname' <- compileVar name - let -- HACK: A Numpy int64 will signal an OverflowError if we pass - -- it a number bigger than 2**63. This does not happen if we - -- pass e.g. int8 a number bigger than 2**7. As a workaround, - -- we first go through the corresponding ctypes type, which does - -- not have this problem. - ctobject = compilePrimType bt - ctcall = simpleCall ctobject [e] - npobject = compilePrimToNp bt - npcall = simpleCall npobject [ctcall] - stm $ Try [Assign vname' npcall] - [Catch (Tuple [Var "TypeError", Var "AssertionError"]) - [badInput i e $ prettySigned (s==Imp.TypeUnsigned) bt]] - -entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims), e) = do - unpack_input <- asks envEntryInput - mem' <- compileVar mem - unpack <- collect $ unpack_input mem' sid bt ept dims e - stm $ Try unpack - [Catch (Tuple [Var "TypeError", Var "AssertionError"]) - [badInput i e $ concat (replicate (length dims) "[]") ++ - prettySigned (ept==Imp.TypeUnsigned) bt]] - -entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem _ t s dims), e) = do - let type_is_wrong = UnOp "not" $ BinOp "in" (simpleCall "type" [e]) $ List [Var "np.ndarray"] - let dtype_is_wrong = UnOp "not" $ BinOp "==" (Field e "dtype") $ Var $ compilePrimToExtNp t s - let dim_is_wrong = UnOp "not" $ BinOp "==" (Field e "ndim") $ Integer $ toInteger $ length dims - stm $ If type_is_wrong - [badInput i e $ concat (replicate (length dims) "[]") ++ - prettySigned (s==Imp.TypeUnsigned) t] - [] - stm $ If dtype_is_wrong - [badInputType i e - (concat (replicate (length dims) "[]") ++ prettySigned (s==Imp.TypeUnsigned) t) - (simpleCall "np.dtype" [Var (compilePrimToExtNp t s)]) - (Field e "dtype")] - [] - stm $ If dim_is_wrong - [badInputDim i e (prettySigned (s==Imp.TypeUnsigned) t ) (length dims)] - [] - - zipWithM_ (unpackDim e) dims [0..] - dest <- compileVar mem - let unwrap_call = simpleCall "unwrapArray" [e] - - stm $ Assign dest unwrap_call - -extValueDescName :: Imp.ExternalValue -> String -extValueDescName (Imp.TransparentValue v) = extName $ valueDescName v -extValueDescName (Imp.OpaqueValue desc []) = extName $ zEncodeString desc -extValueDescName (Imp.OpaqueValue desc (v:_)) = - extName $ zEncodeString desc ++ "_" ++ pretty (baseTag (valueDescVName v)) - -extName :: String -> String -extName = (++"_ext") - -valueDescName :: Imp.ValueDesc -> String -valueDescName = compileName . valueDescVName - -valueDescVName :: Imp.ValueDesc -> VName -valueDescVName (Imp.ScalarValue _ _ vname) = vname -valueDescVName (Imp.ArrayValue vname _ _ _ _) = vname - --- Key into the FUTHARK_PRIMTYPES dict. -readTypeEnum :: PrimType -> Imp.Signedness -> String -readTypeEnum (IntType Int8) Imp.TypeUnsigned = "u8" -readTypeEnum (IntType Int16) Imp.TypeUnsigned = "u16" -readTypeEnum (IntType Int32) Imp.TypeUnsigned = "u32" -readTypeEnum (IntType Int64) Imp.TypeUnsigned = "u64" -readTypeEnum (IntType Int8) Imp.TypeDirect = "i8" -readTypeEnum (IntType Int16) Imp.TypeDirect = "i16" -readTypeEnum (IntType Int32) Imp.TypeDirect = "i32" -readTypeEnum (IntType Int64) Imp.TypeDirect = "i64" -readTypeEnum (FloatType Float32) _ = "f32" -readTypeEnum (FloatType Float64) _ = "f64" -readTypeEnum Imp.Bool _ = "bool" -readTypeEnum Cert _ = error "readTypeEnum: cert" - -readInput :: Imp.ExternalValue -> PyStmt -readInput (Imp.OpaqueValue desc _) = - Raise $ simpleCall "Exception" - [String $ "Cannot read argument of type " ++ desc ++ "."] - -readInput decl@(Imp.TransparentValue (Imp.ScalarValue bt ept _)) = - let type_name = readTypeEnum bt ept - in Assign (Var $ extValueDescName decl) $ simpleCall "read_value" [String type_name] - -readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ bt ept dims)) = - let type_name = readTypeEnum bt ept - in Assign (Var $ extValueDescName decl) $ simpleCall "read_value" - [String $ concat (replicate (length dims) "[]") ++ type_name] - -printValue :: [(Imp.ExternalValue, PyExp)] -> CompilerM op s [PyStmt] -printValue = fmap concat . mapM (uncurry printValue') - -- We copy non-host arrays to the host before printing. This is - -- done in a hacky way - we assume the value has a .get()-method - -- that returns an equivalent Numpy array. This works for PyOpenCL, - -- but we will probably need yet another plugin mechanism here in - -- the future. - where printValue' (Imp.OpaqueValue desc _) _ = - return [Exp $ simpleCall "sys.stdout.write" - [String $ "#"]] - printValue' (Imp.TransparentValue (Imp.ArrayValue mem (Space _) bt ept shape)) e = - printValue' (Imp.TransparentValue (Imp.ArrayValue mem DefaultSpace bt ept shape)) $ - simpleCall (pretty e ++ ".get") [] - printValue' (Imp.TransparentValue _) e = - return [Exp $ Call (Var "write_value") - [Arg e, - ArgKeyword "binary" (Var "binary_output")], - Exp $ simpleCall "sys.stdout.write" [String "\n"]] - -prepareEntry :: (Name, Imp.Function op) -> CompilerM op s - (String, [String], [PyStmt], [PyStmt], [PyStmt], [PyStmt], - [(Imp.ExternalValue, PyExp)], [PyStmt]) -prepareEntry (fname, Imp.Function _ outputs inputs _ results args) = do - let output_paramNames = map (compileName . Imp.paramName) outputs - funTuple = tupleOrSingle $ fmap Var output_paramNames - - (argexps_mem_copies, prepare_run) <- collect' $ forM inputs $ \case - Imp.MemParam name space -> do - -- A program might write to its input parameters, so create a new memory - -- block and copy the source there. This way the program can be run more - -- than once. - name' <- newVName $ baseString name <> "_copy" - copy <- asks envCopy - allocate <- asks envAllocate - let size = Var (extName (compileName name) ++ ".nbytes") -- FIXME - dest = name' - src = name - offset = Integer 0 - case space of - Space sid -> - allocate (Var (compileName name')) size sid - _ -> - stm $ Assign (Var (compileName name')) - (simpleCall "allocateMem" [size]) -- FIXME - dest' <- compileVar dest - src' <- compileVar src - copy dest' offset space src' offset space size (IntType Int32) -- FIXME - return $ Just $ compileName name' - _ -> return Nothing - - prepareIn <- collect $ mapM_ entryPointInput $ zip3 [0..] args $ - map (Var . extValueDescName) args - (res, prepareOut) <- collect' $ mapM entryPointOutput results - - let argexps_lib = map (compileName . Imp.paramName) inputs - argexps_bin = zipWith fromMaybe argexps_lib argexps_mem_copies - fname' = "self." ++ futharkFun (nameToString fname) - call_lib = [Assign funTuple $ simpleCall fname' (fmap Var argexps_lib)] - call_bin = [Assign funTuple $ simpleCall fname' (fmap Var argexps_bin)] - - return (nameToString fname, map extValueDescName args, - prepareIn, call_lib, call_bin, prepareOut, - zip results res, prepare_run) - -copyMemoryDefaultSpace :: PyExp -> PyExp -> PyExp -> PyExp -> PyExp -> - CompilerM op s () -copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes = do - let offset_call1 = simpleCall "addressOffset" - [destmem, destidx, Var "ct.c_byte"] - let offset_call2 = simpleCall "addressOffset" - [srcmem, srcidx, Var "ct.c_byte"] - stm $ Exp $ simpleCall "ct.memmove" [offset_call1, offset_call2, nbytes] - -compileEntryFun :: [PyStmt] -> (Name, Imp.Function op) - -> CompilerM op s (PyFunDef, (PyExp, PyExp)) -compileEntryFun sync entry = do - (fname', params, prepareIn, body_lib, _, prepareOut, res, _) <- prepareEntry entry - let ret = Return $ tupleOrSingle $ map snd res - (pts, rts) = entryTypes $ snd entry - return (Def fname' ("self" : params) $ - prepareIn ++ body_lib ++ prepareOut ++ sync ++ [ret], - (String fname', Tuple [List (map String pts), List (map String rts)])) - -entryTypes :: Imp.Function op -> ([String], [String]) -entryTypes func = (map desc $ Imp.functionArgs func, - map desc $ Imp.functionResult func) - where desc (Imp.OpaqueValue d _) = d - desc (Imp.TransparentValue (Imp.ScalarValue pt s _)) = readTypeEnum pt s - desc (Imp.TransparentValue (Imp.ArrayValue _ _ pt s dims)) = - concat (replicate (length dims) "[]") ++ readTypeEnum pt s - -callEntryFun :: [PyStmt] -> (Name, Imp.Function op) - -> CompilerM op s (PyFunDef, String, PyExp) -callEntryFun pre_timing entry@(fname, Imp.Function _ _ _ _ _ decl_args) = do - (_, _, prepare_in, _, body_bin, _, res, prepare_run) <- prepareEntry entry - - let str_input = map readInput decl_args - end_of_input = [Exp $ simpleCall "end_of_input" [String $ pretty fname]] - - exitcall = [Exp $ simpleCall "sys.exit" [Field (String "Assertion.{} failed") "format(e)"]] - except' = Catch (Var "AssertionError") exitcall - do_run = body_bin ++ pre_timing - (do_run_with_timing, close_runtime_file) = addTiming do_run - - -- We ignore overflow errors and the like for executable entry - -- points. These are (somewhat) well-defined in Futhark. - ignore s = ArgKeyword s $ String "ignore" - errstate = Call (Var "np.errstate") $ map ignore ["divide", "over", "under", "invalid"] - - do_warmup_run = - If (Var "do_warmup_run") (prepare_run ++ do_run) [] - - do_num_runs = - For "i" (simpleCall "range" [simpleCall "int" [Var "num_runs"]]) - (prepare_run ++ do_run_with_timing) - - str_output <- printValue res - - let fname' = "entry_" ++ nameToString fname - - return (Def fname' [] $ - str_input ++ end_of_input ++ prepare_in ++ - [Try [With errstate [do_warmup_run, do_num_runs]] [except']] ++ - [close_runtime_file] ++ - str_output, - - nameToString fname, - - Var fname') - -addTiming :: [PyStmt] -> ([PyStmt], PyStmt) -addTiming statements = - ([ Assign (Var "time_start") $ simpleCall "time.time" [] ] ++ - statements ++ - [ Assign (Var "time_end") $ simpleCall "time.time" [] - , If (Var "runtime_file") print_runtime [] ], - - If (Var "runtime_file") [Exp $ simpleCall "runtime_file.close" []] []) - where print_runtime = - [Exp $ simpleCall "runtime_file.write" - [simpleCall "str" - [BinOp "-" - (toMicroseconds (Var "time_end")) - (toMicroseconds (Var "time_start"))]], - Exp $ simpleCall "runtime_file.write" [String "\n"], - Exp $ simpleCall "runtime_file.flush" []] - toMicroseconds x = - simpleCall "int" [BinOp "*" x $ Integer 1000000] - -compileUnOp :: Imp.UnOp -> String -compileUnOp op = - case op of - Not -> "not" - Complement{} -> "~" - Abs{} -> "abs" - FAbs{} -> "abs" - SSignum{} -> "ssignum" - USignum{} -> "usignum" - -compileBinOpLike :: Monad m => - (v -> m PyExp) - -> Imp.PrimExp v -> Imp.PrimExp v - -> m (PyExp, PyExp, String -> m PyExp) -compileBinOpLike f x y = do - x' <- compilePrimExp f x - y' <- compilePrimExp f y - let simple s = return $ BinOp s x' y' - return (x', y', simple) - --- | The ctypes type corresponding to a 'PrimType'. -compilePrimType :: PrimType -> String -compilePrimType t = - case t of - IntType Int8 -> "ct.c_int8" - IntType Int16 -> "ct.c_int16" - IntType Int32 -> "ct.c_int32" - IntType Int64 -> "ct.c_int64" - FloatType Float32 -> "ct.c_float" - FloatType Float64 -> "ct.c_double" - Imp.Bool -> "ct.c_bool" - Cert -> "ct.c_bool" - --- | The ctypes type corresponding to a 'PrimType', taking sign into account. -compilePrimTypeExt :: PrimType -> Imp.Signedness -> String -compilePrimTypeExt t ept = - case (t, ept) of - (IntType Int8, Imp.TypeUnsigned) -> "ct.c_uint8" - (IntType Int16, Imp.TypeUnsigned) -> "ct.c_uint16" - (IntType Int32, Imp.TypeUnsigned) -> "ct.c_uint32" - (IntType Int64, Imp.TypeUnsigned) -> "ct.c_uint64" - (IntType Int8, _) -> "ct.c_int8" - (IntType Int16, _) -> "ct.c_int16" - (IntType Int32, _) -> "ct.c_int32" - (IntType Int64, _) -> "ct.c_int64" - (FloatType Float32, _) -> "ct.c_float" - (FloatType Float64, _) -> "ct.c_double" - (Imp.Bool, _) -> "ct.c_bool" - (Cert, _) -> "ct.c_byte" - --- | The Numpy type corresponding to a 'PrimType'. -compilePrimToNp :: Imp.PrimType -> String -compilePrimToNp bt = - case bt of - IntType Int8 -> "np.int8" - IntType Int16 -> "np.int16" - IntType Int32 -> "np.int32" - IntType Int64 -> "np.int64" - FloatType Float32 -> "np.float32" - FloatType Float64 -> "np.float64" - Imp.Bool -> "np.byte" - Cert -> "np.byte" - --- | The Numpy type corresponding to a 'PrimType', taking sign into account. -compilePrimToExtNp :: Imp.PrimType -> Imp.Signedness -> String -compilePrimToExtNp bt ept = - case (bt,ept) of - (IntType Int8, Imp.TypeUnsigned) -> "np.uint8" - (IntType Int16, Imp.TypeUnsigned) -> "np.uint16" - (IntType Int32, Imp.TypeUnsigned) -> "np.uint32" - (IntType Int64, Imp.TypeUnsigned) -> "np.uint64" - (IntType Int8, _) -> "np.int8" - (IntType Int16, _) -> "np.int16" - (IntType Int32, _) -> "np.int32" - (IntType Int64, _) -> "np.int64" - (FloatType Float32, _) -> "np.float32" - (FloatType Float64, _) -> "np.float64" - (Imp.Bool, _) -> "np.bool_" - (Cert, _) -> "np.byte" - -compilePrimValue :: Imp.PrimValue -> PyExp -compilePrimValue (IntValue (Int8Value v)) = - simpleCall "np.int8" [Integer $ toInteger v] -compilePrimValue (IntValue (Int16Value v)) = - simpleCall "np.int16" [Integer $ toInteger v] -compilePrimValue (IntValue (Int32Value v)) = - simpleCall "np.int32" [Integer $ toInteger v] -compilePrimValue (IntValue (Int64Value v)) = - simpleCall "np.int64" [Integer $ toInteger v] -compilePrimValue (FloatValue (Float32Value v)) - | isInfinite v = - if v > 0 then Var "np.inf" else Var "-np.inf" - | isNaN v = - Var "np.nan" - | otherwise = simpleCall "np.float32" [Float $ fromRational $ toRational v] -compilePrimValue (FloatValue (Float64Value v)) - | isInfinite v = - if v > 0 then Var "np.inf" else Var "-np.inf" - | isNaN v = - Var "np.nan" - | otherwise = simpleCall "np.float64" [Float $ fromRational $ toRational v] -compilePrimValue (BoolValue v) = Bool v -compilePrimValue Checked = Var "True" - -compileVar :: VName -> CompilerM op s PyExp -compileVar v = - asks $ fromMaybe (Var $ compileName v) . M.lookup v . envVarExp - --- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you. -compilePrimExp :: Monad m => (v -> m PyExp) -> Imp.PrimExp v -> m PyExp - -compilePrimExp _ (Imp.ValueExp v) = return $ compilePrimValue v - -compilePrimExp f (Imp.LeafExp v _) = f v - -compilePrimExp f (Imp.BinOpExp op x y) = do - (x', y', simple) <- compileBinOpLike f x y - case op of - Add{} -> simple "+" - Sub{} -> simple "-" - Mul{} -> simple "*" - FAdd{} -> simple "+" - FSub{} -> simple "-" - FMul{} -> simple "*" - FDiv{} -> simple "/" - FMod{} -> simple "%" - Xor{} -> simple "^" - And{} -> simple "&" - Or{} -> simple "|" - Shl{} -> simple "<<" - LogAnd{} -> simple "and" - LogOr{} -> simple "or" - _ -> return $ simpleCall (pretty op) [x', y'] - -compilePrimExp f (Imp.ConvOpExp conv x) = do - x' <- compilePrimExp f x - return $ simpleCall (pretty conv) [x'] - -compilePrimExp f (Imp.CmpOpExp cmp x y) = do - (x', y', simple) <- compileBinOpLike f x y - case cmp of - CmpEq{} -> simple "==" - FCmpLt{} -> simple "<" - FCmpLe{} -> simple "<=" - CmpLlt -> simple "<" - CmpLle -> simple "<=" - _ -> return $ simpleCall (pretty cmp) [x', y'] - -compilePrimExp f (Imp.UnOpExp op exp1) = - UnOp (compileUnOp op) <$> compilePrimExp f exp1 - -compilePrimExp f (Imp.FunExp h args _) = - simpleCall (futharkFun (pretty h)) <$> mapM (compilePrimExp f) args - -compileExp :: Imp.Exp -> CompilerM op s PyExp -compileExp = compilePrimExp compileLeaf - where compileLeaf (Imp.ScalarVar vname) = - compileVar vname - - compileLeaf (Imp.SizeOf t) = - return $ simpleCall (compilePrimToNp $ IntType Int32) [Integer $ primByteSize t] - - compileLeaf (Imp.Index src (Imp.Count iexp) restype (Imp.Space space) _) = - join $ asks envReadScalar - <*> compileVar src <*> compileExp iexp - <*> pure restype <*> pure space - - compileLeaf (Imp.Index src (Imp.Count iexp) bt _ _) = do - iexp' <- compileExp iexp - let bt' = compilePrimType bt - nptype = compilePrimToNp bt - src' <- compileVar src - return $ simpleCall "indexArray" [src', iexp', Var bt', Var nptype] - -compileCode :: Imp.Code op -> CompilerM op s () - -compileCode Imp.DebugPrint{} = - return () - -compileCode (Imp.Op op) = - join $ asks envOpCompiler <*> pure op - -compileCode (Imp.If cond tb fb) = do - cond' <- compileExp cond - tb' <- collect $ compileCode tb - fb' <- collect $ compileCode fb - stm $ If cond' tb' fb' - -compileCode (c1 Imp.:>>: c2) = do - compileCode c1 - compileCode c2 - -compileCode (Imp.While cond body) = do - cond' <- compileExp cond - body' <- collect $ compileCode body - stm $ While cond' body' - -compileCode (Imp.For i it bound body) = do - bound' <- compileExp bound - let i' = compileName i - body' <- collect $ compileCode body - counter <- pretty <$> newVName "counter" - one <- pretty <$> newVName "one" - stm $ Assign (Var i') $ simpleCall (compilePrimToNp (IntType it)) [Integer 0] - stm $ Assign (Var one) $ simpleCall (compilePrimToNp (IntType it)) [Integer 1] - stm $ For counter (simpleCall "range" [bound']) $ - body' ++ [AssignOp "+" (Var i') (Var one)] - -compileCode (Imp.SetScalar name exp1) = - stm =<< Assign <$> compileVar name <*> compileExp exp1 - -compileCode Imp.DeclareMem{} = return () -compileCode (Imp.DeclareScalar v _ Cert) = do - v' <- compileVar v - stm $ Assign v' $ Var "True" -compileCode Imp.DeclareScalar{} = return () - -compileCode (Imp.DeclareArray name (Space space) t vs) = - join $ asks envStaticArray <*> - pure name <*> pure space <*> pure t <*> pure vs - -compileCode (Imp.DeclareArray name _ t vs) = do - let arr_name = compileName name <> "_arr" - -- It is important to store the Numpy array in a temporary variable - -- to prevent it from going "out-of-scope" before calling - -- unwrapArray (which internally uses the .ctype method); see - -- https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.ctypes.html - atInit $ Assign (Field (Var "self") arr_name) $ case vs of - Imp.ArrayValues vs' -> - Call (Var "np.array") - [Arg $ List $ map compilePrimValue vs', - ArgKeyword "dtype" $ Var $ compilePrimToNp t] - Imp.ArrayZeros n -> - Call (Var "np.zeros") - [Arg $ Integer $ fromIntegral n, - ArgKeyword "dtype" $ Var $ compilePrimToNp t] - atInit $ - Assign (Field (Var "self") (compileName name)) $ - simpleCall "unwrapArray" [Field (Var "self") arr_name] - name' <- compileVar name - stm $ Assign name' $ Field (Var "self") (compileName name) - -compileCode (Imp.Comment s code) = do - code' <- collect $ compileCode code - stm $ Comment s code' - -compileCode (Imp.Assert e (Imp.ErrorMsg parts) (loc,locs)) = do - e' <- compileExp e - let onPart (Imp.ErrorString s) = return ("%s", String s) - onPart (Imp.ErrorInt32 x) = ("%d",) <$> compileExp x - (formatstrs, formatargs) <- unzip <$> mapM onPart parts - stm $ Assert e' (BinOp "%" - (String $ "Error: " ++ concat formatstrs ++ "\n\nBacktrace:\n" ++ stacktrace) - (Tuple formatargs)) - where stacktrace = prettyStacktrace 0 $ map locStr $ loc:locs - -compileCode (Imp.Call dests fname args) = do - args' <- mapM compileArg args - dests' <- tupleOrSingle <$> mapM compileVar dests - let fname' - | isBuiltInFunction fname = futharkFun (pretty fname) - | otherwise = "self." ++ futharkFun (pretty fname) - call' = simpleCall fname' args' - -- If the function returns nothing (is called only for side - -- effects), take care not to assign to an empty tuple. - stm $ if null dests - then Exp call' - else Assign dests' call' - where compileArg (Imp.MemArg m) = compileVar m - compileArg (Imp.ExpArg e) = compileExp e - -compileCode (Imp.SetMem dest src _) = - stm =<< Assign <$> compileVar dest <*> compileVar src - -compileCode (Imp.Allocate name (Imp.Count e) (Imp.Space space)) = - join $ asks envAllocate - <*> compileVar name - <*> compileExp e - <*> pure space - -compileCode (Imp.Allocate name (Imp.Count e) _) = do - e' <- compileExp e - let allocate' = simpleCall "allocateMem" [e'] - stm =<< Assign <$> compileVar name <*> pure allocate' - -compileCode (Imp.Free name _) = - stm =<< Assign <$> compileVar name <*> pure None - -compileCode (Imp.Copy dest (Imp.Count destoffset) DefaultSpace src (Imp.Count srcoffset) DefaultSpace (Imp.Count size)) = do - destoffset' <- compileExp destoffset - srcoffset' <- compileExp srcoffset - dest' <- compileVar dest - src' <- compileVar src - size' <- compileExp size - let offset_call1 = simpleCall "addressOffset" [dest', destoffset', Var "ct.c_byte"] - let offset_call2 = simpleCall "addressOffset" [src', srcoffset', Var "ct.c_byte"] - stm $ Exp $ simpleCall "ct.memmove" [offset_call1, offset_call2, size'] - -compileCode (Imp.Copy dest (Imp.Count destoffset) destspace src (Imp.Count srcoffset) srcspace (Imp.Count size)) = do - copy <- asks envCopy - join $ copy - <$> compileVar dest <*> compileExp destoffset <*> pure destspace - <*> compileVar src <*> compileExp srcoffset <*> pure srcspace - <*> compileExp size <*> pure (IntType Int32) -- FIXME - -compileCode (Imp.Write dest (Imp.Count idx) elemtype (Imp.Space space) _ elemexp) = - join $ asks envWriteScalar - <*> compileVar dest - <*> compileExp idx - <*> pure elemtype - <*> pure space - <*> compileExp elemexp - -compileCode (Imp.Write dest (Imp.Count idx) elemtype _ _ elemexp) = do - idx' <- compileExp idx - elemexp' <- compileExp elemexp - dest' <- compileVar dest - let elemtype' = compilePrimType elemtype - ctype = simpleCall elemtype' [elemexp'] - stm $ Exp $ simpleCall "writeScalarArray" [dest', idx', ctype] - +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} + +-- | A generic Python code generator which is polymorphic in the type +-- of the operations. Concretely, we use this to handle both +-- sequential and PyOpenCL Python code. +module Futhark.CodeGen.Backends.GenericPython + ( compileProg, + Constructor (..), + emptyConstructor, + compileName, + compileVar, + compileDim, + compileExp, + compilePrimExp, + compileCode, + compilePrimValue, + compilePrimType, + compilePrimTypeExt, + compilePrimToNp, + compilePrimToExtNp, + Operations (..), + defaultOperations, + unpackDim, + CompilerM (..), + OpCompiler, + WriteScalar, + ReadScalar, + Allocate, + Copy, + StaticArray, + EntryOutput, + EntryInput, + CompilerEnv (..), + CompilerState (..), + stm, + atInit, + collect', + collect, + simpleCall, + copyMemoryDefaultSpace, + ) +where + +import Control.Monad.Identity +import Control.Monad.RWS +import qualified Data.Map as M +import Data.Maybe +import Futhark.CodeGen.Backends.GenericPython.AST +import Futhark.CodeGen.Backends.GenericPython.Definitions +import Futhark.CodeGen.Backends.GenericPython.Options +import qualified Futhark.CodeGen.ImpCode as Imp +import Futhark.IR.Primitive hiding (Bool) +import Futhark.IR.Prop (isBuiltInFunction) +import Futhark.IR.Syntax (Space (..)) +import Futhark.MonadFreshNames +import Futhark.Util (zEncodeString) + +-- | A substitute expression compiler, tried before the main +-- compilation function. +type OpCompiler op s = op -> CompilerM op s () + +-- | Write a scalar to the given memory block with the given index and +-- in the given memory space. +type WriteScalar op s = + PyExp -> + PyExp -> + PrimType -> + Imp.SpaceId -> + PyExp -> + CompilerM op s () + +-- | Read a scalar from the given memory block with the given index and +-- in the given memory space. +type ReadScalar op s = + PyExp -> + PyExp -> + PrimType -> + Imp.SpaceId -> + CompilerM op s PyExp + +-- | Allocate a memory block of the given size in the given memory +-- space, saving a reference in the given variable name. +type Allocate op s = + PyExp -> + PyExp -> + Imp.SpaceId -> + CompilerM op s () + +-- | Copy from one memory block to another. +type Copy op s = + PyExp -> + PyExp -> + Imp.Space -> + PyExp -> + PyExp -> + Imp.Space -> + PyExp -> + PrimType -> + CompilerM op s () + +-- | Create a static array of values - initialised at load time. +type StaticArray op s = VName -> Imp.SpaceId -> PrimType -> Imp.ArrayContents -> CompilerM op s () + +-- | Construct the Python array being returned from an entry point. +type EntryOutput op s = + VName -> + Imp.SpaceId -> + PrimType -> + Imp.Signedness -> + [Imp.DimSize] -> + CompilerM op s PyExp + +-- | Unpack the array being passed to an entry point. +type EntryInput op s = + PyExp -> + Imp.SpaceId -> + PrimType -> + Imp.Signedness -> + [Imp.DimSize] -> + PyExp -> + CompilerM op s () + +data Operations op s = Operations + { opsWriteScalar :: WriteScalar op s, + opsReadScalar :: ReadScalar op s, + opsAllocate :: Allocate op s, + opsCopy :: Copy op s, + opsStaticArray :: StaticArray op s, + opsCompiler :: OpCompiler op s, + opsEntryOutput :: EntryOutput op s, + opsEntryInput :: EntryInput op s + } + +-- | A set of operations that fail for every operation involving +-- non-default memory spaces. Uses plain pointers and @malloc@ for +-- memory management. +defaultOperations :: Operations op s +defaultOperations = + Operations + { opsWriteScalar = defWriteScalar, + opsReadScalar = defReadScalar, + opsAllocate = defAllocate, + opsCopy = defCopy, + opsStaticArray = defStaticArray, + opsCompiler = defCompiler, + opsEntryOutput = defEntryOutput, + opsEntryInput = defEntryInput + } + where + defWriteScalar _ _ _ _ _ = + error "Cannot write to non-default memory space because I am dumb" + defReadScalar _ _ _ _ = + error "Cannot read from non-default memory space" + defAllocate _ _ _ = + error "Cannot allocate in non-default memory space" + defCopy _ _ _ _ _ _ _ _ = + error "Cannot copy to or from non-default memory space" + defStaticArray _ _ _ _ = + error "Cannot create static array in non-default memory space" + defCompiler _ = + error "The default compiler cannot compile extended operations" + defEntryOutput _ _ _ _ = + error "Cannot return array not in default memory space" + defEntryInput _ _ _ _ = + error "Cannot accept array not in default memory space" + +data CompilerEnv op s = CompilerEnv + { envOperations :: Operations op s, + envVarExp :: M.Map VName PyExp + } + +envOpCompiler :: CompilerEnv op s -> OpCompiler op s +envOpCompiler = opsCompiler . envOperations + +envReadScalar :: CompilerEnv op s -> ReadScalar op s +envReadScalar = opsReadScalar . envOperations + +envWriteScalar :: CompilerEnv op s -> WriteScalar op s +envWriteScalar = opsWriteScalar . envOperations + +envAllocate :: CompilerEnv op s -> Allocate op s +envAllocate = opsAllocate . envOperations + +envCopy :: CompilerEnv op s -> Copy op s +envCopy = opsCopy . envOperations + +envStaticArray :: CompilerEnv op s -> StaticArray op s +envStaticArray = opsStaticArray . envOperations + +envEntryOutput :: CompilerEnv op s -> EntryOutput op s +envEntryOutput = opsEntryOutput . envOperations + +envEntryInput :: CompilerEnv op s -> EntryInput op s +envEntryInput = opsEntryInput . envOperations + +newCompilerEnv :: Operations op s -> CompilerEnv op s +newCompilerEnv ops = + CompilerEnv + { envOperations = ops, + envVarExp = mempty + } + +data CompilerState s = CompilerState + { compNameSrc :: VNameSource, + compInit :: [PyStmt], + compUserState :: s + } + +newCompilerState :: VNameSource -> s -> CompilerState s +newCompilerState src s = + CompilerState + { compNameSrc = src, + compInit = [], + compUserState = s + } + +newtype CompilerM op s a = CompilerM (RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a) + deriving + ( Functor, + Applicative, + Monad, + MonadState (CompilerState s), + MonadReader (CompilerEnv op s), + MonadWriter [PyStmt] + ) + +instance MonadFreshNames (CompilerM op s) where + getNameSource = gets compNameSrc + putNameSource src = modify $ \s -> s {compNameSrc = src} + +collect :: CompilerM op s () -> CompilerM op s [PyStmt] +collect m = pass $ do + ((), w) <- listen m + return (w, const mempty) + +collect' :: CompilerM op s a -> CompilerM op s (a, [PyStmt]) +collect' m = pass $ do + (x, w) <- listen m + return ((x, w), const mempty) + +atInit :: PyStmt -> CompilerM op s () +atInit x = modify $ \s -> + s {compInit = compInit s ++ [x]} + +stm :: PyStmt -> CompilerM op s () +stm x = tell [x] + +futharkFun :: String -> String +futharkFun s = "futhark_" ++ zEncodeString s + +compileOutput :: [Imp.Param] -> [PyExp] +compileOutput = map (Var . compileName . Imp.paramName) + +runCompilerM :: + Operations op s -> + VNameSource -> + s -> + CompilerM op s a -> + a +runCompilerM ops src userstate (CompilerM m) = + fst $ evalRWS m (newCompilerEnv ops) (newCompilerState src userstate) + +standardOptions :: [Option] +standardOptions = + [ Option + { optionLongName = "write-runtime-to", + optionShortName = Just 't', + optionArgument = RequiredArgument "str", + optionAction = + [ If + (Var "runtime_file") + [Exp $ simpleCall "runtime_file.close" []] + [], + Assign (Var "runtime_file") $ + simpleCall "open" [Var "optarg", String "w"] + ] + }, + Option + { optionLongName = "runs", + optionShortName = Just 'r', + optionArgument = RequiredArgument "str", + optionAction = + [ Assign (Var "num_runs") $ Var "optarg", + Assign (Var "do_warmup_run") $ Bool True + ] + }, + Option + { optionLongName = "entry-point", + optionShortName = Just 'e', + optionArgument = RequiredArgument "str", + optionAction = + [Assign (Var "entry_point") $ Var "optarg"] + }, + Option + { optionLongName = "binary-output", + optionShortName = Just 'b', + optionArgument = NoArgument, + optionAction = [Assign (Var "binary_output") $ Bool True] + }, + Option + { optionLongName = "tuning", + optionShortName = Nothing, + optionArgument = RequiredArgument "open", + optionAction = [Exp $ simpleCall "read_tuning_file" [Var "sizes", Var "optarg"]] + } + ] + +-- | The class generated by the code generator must have a +-- constructor, although it can be vacuous. +data Constructor = Constructor [String] [PyStmt] + +-- | A constructor that takes no arguments and does nothing. +emptyConstructor :: Constructor +emptyConstructor = Constructor ["self"] [Pass] + +constructorToFunDef :: Constructor -> [PyStmt] -> PyFunDef +constructorToFunDef (Constructor params body) at_init = + Def "__init__" params $ body <> at_init + +compileProg :: + MonadFreshNames m => + Maybe String -> + Constructor -> + [PyStmt] -> + [PyStmt] -> + Operations op s -> + s -> + [PyStmt] -> + [Option] -> + Imp.Definitions op -> + m String +compileProg module_name constructor imports defines ops userstate sync options prog = do + src <- getNameSource + let prog' = runCompilerM ops src userstate compileProg' + maybe_shebang = + case module_name of + Nothing -> "#!/usr/bin/env python\n" + Just _ -> "" + return $ + maybe_shebang + ++ pretty + ( PyProg $ + imports + ++ [ Import "argparse" Nothing, + Assign (Var "sizes") $ Dict [] + ] + ++ defines + ++ [Escape pyUtility] + ++ prog' + ) + where + Imp.Definitions consts (Imp.Functions funs) = prog + compileProg' = withConstantSubsts consts $ do + compileConstants consts + + definitions <- mapM compileFunc funs + at_inits <- gets compInit + + let constructor' = constructorToFunDef constructor at_inits + + case module_name of + Just name -> do + (entry_points, entry_point_types) <- + unzip <$> mapM (compileEntryFun sync) (filter (Imp.functionEntry . snd) funs) + return + [ ClassDef $ + Class name $ + Assign (Var "entry_points") (Dict entry_point_types) : + map FunDef (constructor' : definitions ++ entry_points) + ] + Nothing -> do + let classinst = Assign (Var "self") $ simpleCall "internal" [] + (entry_point_defs, entry_point_names, entry_points) <- + unzip3 + <$> mapM + (callEntryFun sync) + (filter (Imp.functionEntry . snd) funs) + return + ( parse_options + ++ ClassDef + ( Class "internal" $ + map FunDef $ + constructor' : definitions + ) : + classinst : + map FunDef entry_point_defs + ++ selectEntryPoint entry_point_names entry_points + ) + + parse_options = + Assign (Var "runtime_file") None : + Assign (Var "do_warmup_run") (Bool False) : + Assign (Var "num_runs") (Integer 1) : + Assign (Var "entry_point") (String "main") : + Assign (Var "binary_output") (Bool False) : + generateOptionParser (standardOptions ++ options) + + selectEntryPoint entry_point_names entry_points = + [ Assign (Var "entry_points") $ + Dict $ zip (map String entry_point_names) entry_points, + Assign (Var "entry_point_fun") $ + simpleCall "entry_points.get" [Var "entry_point"], + If + (BinOp "==" (Var "entry_point_fun") None) + [ Exp $ + simpleCall + "sys.exit" + [ Call + ( Field + (String "No entry point '{}'. Select another with --entry point. Options are:\n{}") + "format" + ) + [ Arg $ Var "entry_point", + Arg $ + Call + (Field (String "\n") "join") + [Arg $ simpleCall "entry_points.keys" []] + ] + ] + ] + [Exp $ simpleCall "entry_point_fun" []] + ] + +withConstantSubsts :: Imp.Constants op -> CompilerM op s a -> CompilerM op s a +withConstantSubsts (Imp.Constants ps _) = + local $ \env -> env {envVarExp = foldMap constExp ps} + where + constExp p = + M.singleton (Imp.paramName p) $ + Index (Var "self.constants") $ + IdxExp $ String $ pretty $ Imp.paramName p + +compileConstants :: Imp.Constants op -> CompilerM op s () +compileConstants (Imp.Constants _ init_consts) = do + atInit $ Assign (Var "self.constants") $ Dict [] + mapM_ atInit =<< collect (compileCode init_consts) + +compileFunc :: (Name, Imp.Function op) -> CompilerM op s PyFunDef +compileFunc (fname, Imp.Function _ outputs inputs body _ _) = do + body' <- collect $ compileCode body + let inputs' = map (compileName . Imp.paramName) inputs + let ret = Return $ tupleOrSingle $ compileOutput outputs + return $ + Def (futharkFun . nameToString $ fname) ("self" : inputs') $ + body' ++ [ret] + +tupleOrSingle :: [PyExp] -> PyExp +tupleOrSingle [e] = e +tupleOrSingle es = Tuple es + +-- | A 'Call' where the function is a variable and every argument is a +-- simple 'Arg'. +simpleCall :: String -> [PyExp] -> PyExp +simpleCall fname = Call (Var fname) . map Arg + +compileName :: VName -> String +compileName = zEncodeString . pretty + +compileDim :: Imp.DimSize -> PyExp +compileDim (Imp.Constant v) = compilePrimValue v +compileDim (Imp.Var v) = Var $ compileName v + +unpackDim :: PyExp -> Imp.DimSize -> Int32 -> CompilerM op s () +unpackDim arr_name (Imp.Constant c) i = do + let shape_name = Field arr_name "shape" + let constant_c = compilePrimValue c + let constant_i = Integer $ toInteger i + stm $ + Assert (BinOp "==" constant_c (Index shape_name $ IdxExp constant_i)) $ + String "constant dimension wrong" +unpackDim arr_name (Imp.Var var) i = do + let shape_name = Field arr_name "shape" + src = Index shape_name $ IdxExp $ Integer $ toInteger i + var' <- compileVar var + stm $ Assign var' $ simpleCall "np.int32" [src] + +entryPointOutput :: Imp.ExternalValue -> CompilerM op s PyExp +entryPointOutput (Imp.OpaqueValue desc vs) = + simpleCall "opaque" . (String (pretty desc) :) + <$> mapM (entryPointOutput . Imp.TransparentValue) vs +entryPointOutput (Imp.TransparentValue (Imp.ScalarValue bt ept name)) = do + name' <- compileVar name + return $ simpleCall tf [name'] + where + tf = compilePrimToExtNp bt ept +entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims)) = do + pack_output <- asks envEntryOutput + pack_output mem sid bt ept dims +entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ bt ept dims)) = do + mem' <- compileVar mem + let cast = Cast mem' (compilePrimTypeExt bt ept) + return $ simpleCall "createArray" [cast, Tuple $ map compileDim dims] + +badInput :: Int -> PyExp -> String -> PyStmt +badInput i e t = + Raise $ + simpleCall + "TypeError" + [ Call + (Field (String err_msg) "format") + [Arg (String t), Arg $ simpleCall "type" [e], Arg e] + ] + where + err_msg = + unlines + [ "Argument #" ++ show i ++ " has invalid value", + "Futhark type: {}", + "Argument has Python type {} and value: {}" + ] + +badInputType :: Int -> PyExp -> String -> PyExp -> PyExp -> PyStmt +badInputType i e t de dg = + Raise $ + simpleCall + "TypeError" + [ Call + (Field (String err_msg) "format") + [Arg (String t), Arg $ simpleCall "type" [e], Arg e, Arg de, Arg dg] + ] + where + err_msg = + unlines + [ "Argument #" ++ show i ++ " has invalid value", + "Futhark type: {}", + "Argument has Python type {} and value: {}", + "Expected array with elements of dtype: {}", + "The array given has elements of dtype: {}" + ] + +badInputDim :: Int -> PyExp -> String -> Int -> PyStmt +badInputDim i e typ dimf = + Raise $ + simpleCall + "TypeError" + [ Call + (Field (String err_msg) "format") + [Arg eft, Arg aft] + ] + where + eft = String (concat (replicate dimf "[]") ++ typ) + aft = BinOp "+" (BinOp "*" (String "[]") (Field e "ndim")) (String typ) + err_msg = + unlines + [ "Argument #" ++ show i ++ " has invalid value", + "Dimensionality mismatch", + "Expected Futhark type: {}", + "Bad Python value passed", + "Actual Futhark type: {}" + ] + +entryPointInput :: (Int, Imp.ExternalValue, PyExp) -> CompilerM op s () +entryPointInput (i, Imp.OpaqueValue desc vs, e) = do + let type_is_ok = + BinOp + "and" + (simpleCall "isinstance" [e, Var "opaque"]) + (BinOp "==" (Field e "desc") (String desc)) + stm $ If (UnOp "not" type_is_ok) [badInput i e desc] [] + mapM_ entryPointInput $ + zip3 (repeat i) (map Imp.TransparentValue vs) $ + map (Index (Field e "data") . IdxExp . Integer) [0 ..] +entryPointInput (i, Imp.TransparentValue (Imp.ScalarValue bt s name), e) = do + vname' <- compileVar name + let -- HACK: A Numpy int64 will signal an OverflowError if we pass + -- it a number bigger than 2**63. This does not happen if we + -- pass e.g. int8 a number bigger than 2**7. As a workaround, + -- we first go through the corresponding ctypes type, which does + -- not have this problem. + ctobject = compilePrimType bt + ctcall = simpleCall ctobject [e] + npobject = compilePrimToNp bt + npcall = simpleCall npobject [ctcall] + stm $ + Try + [Assign vname' npcall] + [ Catch + (Tuple [Var "TypeError", Var "AssertionError"]) + [badInput i e $ prettySigned (s == Imp.TypeUnsigned) bt] + ] +entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims), e) = do + unpack_input <- asks envEntryInput + mem' <- compileVar mem + unpack <- collect $ unpack_input mem' sid bt ept dims e + stm $ + Try + unpack + [ Catch + (Tuple [Var "TypeError", Var "AssertionError"]) + [ badInput i e $ + concat (replicate (length dims) "[]") + ++ prettySigned (ept == Imp.TypeUnsigned) bt + ] + ] +entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem _ t s dims), e) = do + let type_is_wrong = UnOp "not" $ BinOp "in" (simpleCall "type" [e]) $ List [Var "np.ndarray"] + let dtype_is_wrong = UnOp "not" $ BinOp "==" (Field e "dtype") $ Var $ compilePrimToExtNp t s + let dim_is_wrong = UnOp "not" $ BinOp "==" (Field e "ndim") $ Integer $ toInteger $ length dims + stm $ + If + type_is_wrong + [ badInput i e $ + concat (replicate (length dims) "[]") + ++ prettySigned (s == Imp.TypeUnsigned) t + ] + [] + stm $ + If + dtype_is_wrong + [ badInputType + i + e + (concat (replicate (length dims) "[]") ++ prettySigned (s == Imp.TypeUnsigned) t) + (simpleCall "np.dtype" [Var (compilePrimToExtNp t s)]) + (Field e "dtype") + ] + [] + stm $ + If + dim_is_wrong + [badInputDim i e (prettySigned (s == Imp.TypeUnsigned) t) (length dims)] + [] + + zipWithM_ (unpackDim e) dims [0 ..] + dest <- compileVar mem + let unwrap_call = simpleCall "unwrapArray" [e] + + stm $ Assign dest unwrap_call + +extValueDescName :: Imp.ExternalValue -> String +extValueDescName (Imp.TransparentValue v) = extName $ valueDescName v +extValueDescName (Imp.OpaqueValue desc []) = extName $ zEncodeString desc +extValueDescName (Imp.OpaqueValue desc (v : _)) = + extName $ zEncodeString desc ++ "_" ++ pretty (baseTag (valueDescVName v)) + +extName :: String -> String +extName = (++ "_ext") + +valueDescName :: Imp.ValueDesc -> String +valueDescName = compileName . valueDescVName + +valueDescVName :: Imp.ValueDesc -> VName +valueDescVName (Imp.ScalarValue _ _ vname) = vname +valueDescVName (Imp.ArrayValue vname _ _ _ _) = vname + +-- Key into the FUTHARK_PRIMTYPES dict. +readTypeEnum :: PrimType -> Imp.Signedness -> String +readTypeEnum (IntType Int8) Imp.TypeUnsigned = "u8" +readTypeEnum (IntType Int16) Imp.TypeUnsigned = "u16" +readTypeEnum (IntType Int32) Imp.TypeUnsigned = "u32" +readTypeEnum (IntType Int64) Imp.TypeUnsigned = "u64" +readTypeEnum (IntType Int8) Imp.TypeDirect = "i8" +readTypeEnum (IntType Int16) Imp.TypeDirect = "i16" +readTypeEnum (IntType Int32) Imp.TypeDirect = "i32" +readTypeEnum (IntType Int64) Imp.TypeDirect = "i64" +readTypeEnum (FloatType Float32) _ = "f32" +readTypeEnum (FloatType Float64) _ = "f64" +readTypeEnum Imp.Bool _ = "bool" +readTypeEnum Cert _ = error "readTypeEnum: cert" + +readInput :: Imp.ExternalValue -> PyStmt +readInput (Imp.OpaqueValue desc _) = + Raise $ + simpleCall + "Exception" + [String $ "Cannot read argument of type " ++ desc ++ "."] +readInput decl@(Imp.TransparentValue (Imp.ScalarValue bt ept _)) = + let type_name = readTypeEnum bt ept + in Assign (Var $ extValueDescName decl) $ simpleCall "read_value" [String type_name] +readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ bt ept dims)) = + let type_name = readTypeEnum bt ept + in Assign (Var $ extValueDescName decl) $ + simpleCall + "read_value" + [String $ concat (replicate (length dims) "[]") ++ type_name] + +printValue :: [(Imp.ExternalValue, PyExp)] -> CompilerM op s [PyStmt] +printValue = fmap concat . mapM (uncurry printValue') + where + -- We copy non-host arrays to the host before printing. This is + -- done in a hacky way - we assume the value has a .get()-method + -- that returns an equivalent Numpy array. This works for PyOpenCL, + -- but we will probably need yet another plugin mechanism here in + -- the future. + printValue' (Imp.OpaqueValue desc _) _ = + return + [ Exp $ + simpleCall + "sys.stdout.write" + [String $ "#"] + ] + printValue' (Imp.TransparentValue (Imp.ArrayValue mem (Space _) bt ept shape)) e = + printValue' (Imp.TransparentValue (Imp.ArrayValue mem DefaultSpace bt ept shape)) $ + simpleCall (pretty e ++ ".get") [] + printValue' (Imp.TransparentValue _) e = + return + [ Exp $ + Call + (Var "write_value") + [ Arg e, + ArgKeyword "binary" (Var "binary_output") + ], + Exp $ simpleCall "sys.stdout.write" [String "\n"] + ] + +prepareEntry :: + (Name, Imp.Function op) -> + CompilerM + op + s + ( String, + [String], + [PyStmt], + [PyStmt], + [PyStmt], + [PyStmt], + [(Imp.ExternalValue, PyExp)], + [PyStmt] + ) +prepareEntry (fname, Imp.Function _ outputs inputs _ results args) = do + let output_paramNames = map (compileName . Imp.paramName) outputs + funTuple = tupleOrSingle $ fmap Var output_paramNames + + (argexps_mem_copies, prepare_run) <- collect' $ + forM inputs $ \case + Imp.MemParam name space -> do + -- A program might write to its input parameters, so create a new memory + -- block and copy the source there. This way the program can be run more + -- than once. + name' <- newVName $ baseString name <> "_copy" + copy <- asks envCopy + allocate <- asks envAllocate + let size = Var (extName (compileName name) ++ ".nbytes") -- FIXME + dest = name' + src = name + offset = Integer 0 + case space of + Space sid -> + allocate (Var (compileName name')) size sid + _ -> + stm $ + Assign + (Var (compileName name')) + (simpleCall "allocateMem" [size]) -- FIXME + dest' <- compileVar dest + src' <- compileVar src + copy dest' offset space src' offset space size (IntType Int32) -- FIXME + return $ Just $ compileName name' + _ -> return Nothing + + prepareIn <- + collect $ + mapM_ entryPointInput $ + zip3 [0 ..] args $ + map (Var . extValueDescName) args + (res, prepareOut) <- collect' $ mapM entryPointOutput results + + let argexps_lib = map (compileName . Imp.paramName) inputs + argexps_bin = zipWith fromMaybe argexps_lib argexps_mem_copies + fname' = "self." ++ futharkFun (nameToString fname) + call_lib = [Assign funTuple $ simpleCall fname' (fmap Var argexps_lib)] + call_bin = [Assign funTuple $ simpleCall fname' (fmap Var argexps_bin)] + + return + ( nameToString fname, + map extValueDescName args, + prepareIn, + call_lib, + call_bin, + prepareOut, + zip results res, + prepare_run + ) + +copyMemoryDefaultSpace :: + PyExp -> + PyExp -> + PyExp -> + PyExp -> + PyExp -> + CompilerM op s () +copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes = do + let offset_call1 = + simpleCall + "addressOffset" + [destmem, destidx, Var "ct.c_byte"] + let offset_call2 = + simpleCall + "addressOffset" + [srcmem, srcidx, Var "ct.c_byte"] + stm $ Exp $ simpleCall "ct.memmove" [offset_call1, offset_call2, nbytes] + +compileEntryFun :: + [PyStmt] -> + (Name, Imp.Function op) -> + CompilerM op s (PyFunDef, (PyExp, PyExp)) +compileEntryFun sync entry = do + (fname', params, prepareIn, body_lib, _, prepareOut, res, _) <- prepareEntry entry + let ret = Return $ tupleOrSingle $ map snd res + (pts, rts) = entryTypes $ snd entry + return + ( Def fname' ("self" : params) $ + prepareIn ++ body_lib ++ prepareOut ++ sync ++ [ret], + (String fname', Tuple [List (map String pts), List (map String rts)]) + ) + +entryTypes :: Imp.Function op -> ([String], [String]) +entryTypes func = + ( map desc $ Imp.functionArgs func, + map desc $ Imp.functionResult func + ) + where + desc (Imp.OpaqueValue d _) = d + desc (Imp.TransparentValue (Imp.ScalarValue pt s _)) = readTypeEnum pt s + desc (Imp.TransparentValue (Imp.ArrayValue _ _ pt s dims)) = + concat (replicate (length dims) "[]") ++ readTypeEnum pt s + +callEntryFun :: + [PyStmt] -> + (Name, Imp.Function op) -> + CompilerM op s (PyFunDef, String, PyExp) +callEntryFun pre_timing entry@(fname, Imp.Function _ _ _ _ _ decl_args) = do + (_, _, prepare_in, _, body_bin, _, res, prepare_run) <- prepareEntry entry + + let str_input = map readInput decl_args + end_of_input = [Exp $ simpleCall "end_of_input" [String $ pretty fname]] + + exitcall = [Exp $ simpleCall "sys.exit" [Field (String "Assertion.{} failed") "format(e)"]] + except' = Catch (Var "AssertionError") exitcall + do_run = body_bin ++ pre_timing + (do_run_with_timing, close_runtime_file) = addTiming do_run + + -- We ignore overflow errors and the like for executable entry + -- points. These are (somewhat) well-defined in Futhark. + ignore s = ArgKeyword s $ String "ignore" + errstate = Call (Var "np.errstate") $ map ignore ["divide", "over", "under", "invalid"] + + do_warmup_run = + If (Var "do_warmup_run") (prepare_run ++ do_run) [] + + do_num_runs = + For + "i" + (simpleCall "range" [simpleCall "int" [Var "num_runs"]]) + (prepare_run ++ do_run_with_timing) + + str_output <- printValue res + + let fname' = "entry_" ++ nameToString fname + + return + ( Def fname' [] $ + str_input ++ end_of_input ++ prepare_in + ++ [Try [With errstate [do_warmup_run, do_num_runs]] [except']] + ++ [close_runtime_file] + ++ str_output, + nameToString fname, + Var fname' + ) + +addTiming :: [PyStmt] -> ([PyStmt], PyStmt) +addTiming statements = + ( [Assign (Var "time_start") $ simpleCall "time.time" []] + ++ statements + ++ [ Assign (Var "time_end") $ simpleCall "time.time" [], + If (Var "runtime_file") print_runtime [] + ], + If (Var "runtime_file") [Exp $ simpleCall "runtime_file.close" []] [] + ) + where + print_runtime = + [ Exp $ + simpleCall + "runtime_file.write" + [ simpleCall + "str" + [ BinOp + "-" + (toMicroseconds (Var "time_end")) + (toMicroseconds (Var "time_start")) + ] + ], + Exp $ simpleCall "runtime_file.write" [String "\n"], + Exp $ simpleCall "runtime_file.flush" [] + ] + toMicroseconds x = + simpleCall "int" [BinOp "*" x $ Integer 1000000] + +compileUnOp :: Imp.UnOp -> String +compileUnOp op = + case op of + Not -> "not" + Complement {} -> "~" + Abs {} -> "abs" + FAbs {} -> "abs" + SSignum {} -> "ssignum" + USignum {} -> "usignum" + +compileBinOpLike :: + Monad m => + (v -> m PyExp) -> + Imp.PrimExp v -> + Imp.PrimExp v -> + m (PyExp, PyExp, String -> m PyExp) +compileBinOpLike f x y = do + x' <- compilePrimExp f x + y' <- compilePrimExp f y + let simple s = return $ BinOp s x' y' + return (x', y', simple) + +-- | The ctypes type corresponding to a 'PrimType'. +compilePrimType :: PrimType -> String +compilePrimType t = + case t of + IntType Int8 -> "ct.c_int8" + IntType Int16 -> "ct.c_int16" + IntType Int32 -> "ct.c_int32" + IntType Int64 -> "ct.c_int64" + FloatType Float32 -> "ct.c_float" + FloatType Float64 -> "ct.c_double" + Imp.Bool -> "ct.c_bool" + Cert -> "ct.c_bool" + +-- | The ctypes type corresponding to a 'PrimType', taking sign into account. +compilePrimTypeExt :: PrimType -> Imp.Signedness -> String +compilePrimTypeExt t ept = + case (t, ept) of + (IntType Int8, Imp.TypeUnsigned) -> "ct.c_uint8" + (IntType Int16, Imp.TypeUnsigned) -> "ct.c_uint16" + (IntType Int32, Imp.TypeUnsigned) -> "ct.c_uint32" + (IntType Int64, Imp.TypeUnsigned) -> "ct.c_uint64" + (IntType Int8, _) -> "ct.c_int8" + (IntType Int16, _) -> "ct.c_int16" + (IntType Int32, _) -> "ct.c_int32" + (IntType Int64, _) -> "ct.c_int64" + (FloatType Float32, _) -> "ct.c_float" + (FloatType Float64, _) -> "ct.c_double" + (Imp.Bool, _) -> "ct.c_bool" + (Cert, _) -> "ct.c_byte" + +-- | The Numpy type corresponding to a 'PrimType'. +compilePrimToNp :: Imp.PrimType -> String +compilePrimToNp bt = + case bt of + IntType Int8 -> "np.int8" + IntType Int16 -> "np.int16" + IntType Int32 -> "np.int32" + IntType Int64 -> "np.int64" + FloatType Float32 -> "np.float32" + FloatType Float64 -> "np.float64" + Imp.Bool -> "np.byte" + Cert -> "np.byte" + +-- | The Numpy type corresponding to a 'PrimType', taking sign into account. +compilePrimToExtNp :: Imp.PrimType -> Imp.Signedness -> String +compilePrimToExtNp bt ept = + case (bt, ept) of + (IntType Int8, Imp.TypeUnsigned) -> "np.uint8" + (IntType Int16, Imp.TypeUnsigned) -> "np.uint16" + (IntType Int32, Imp.TypeUnsigned) -> "np.uint32" + (IntType Int64, Imp.TypeUnsigned) -> "np.uint64" + (IntType Int8, _) -> "np.int8" + (IntType Int16, _) -> "np.int16" + (IntType Int32, _) -> "np.int32" + (IntType Int64, _) -> "np.int64" + (FloatType Float32, _) -> "np.float32" + (FloatType Float64, _) -> "np.float64" + (Imp.Bool, _) -> "np.bool_" + (Cert, _) -> "np.byte" + +compilePrimValue :: Imp.PrimValue -> PyExp +compilePrimValue (IntValue (Int8Value v)) = + simpleCall "np.int8" [Integer $ toInteger v] +compilePrimValue (IntValue (Int16Value v)) = + simpleCall "np.int16" [Integer $ toInteger v] +compilePrimValue (IntValue (Int32Value v)) = + simpleCall "np.int32" [Integer $ toInteger v] +compilePrimValue (IntValue (Int64Value v)) = + simpleCall "np.int64" [Integer $ toInteger v] +compilePrimValue (FloatValue (Float32Value v)) + | isInfinite v = + if v > 0 then Var "np.inf" else Var "-np.inf" + | isNaN v = + Var "np.nan" + | otherwise = simpleCall "np.float32" [Float $ fromRational $ toRational v] +compilePrimValue (FloatValue (Float64Value v)) + | isInfinite v = + if v > 0 then Var "np.inf" else Var "-np.inf" + | isNaN v = + Var "np.nan" + | otherwise = simpleCall "np.float64" [Float $ fromRational $ toRational v] +compilePrimValue (BoolValue v) = Bool v +compilePrimValue Checked = Var "True" + +compileVar :: VName -> CompilerM op s PyExp +compileVar v = + asks $ fromMaybe (Var $ compileName v) . M.lookup v . envVarExp + +-- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you. +compilePrimExp :: Monad m => (v -> m PyExp) -> Imp.PrimExp v -> m PyExp +compilePrimExp _ (Imp.ValueExp v) = return $ compilePrimValue v +compilePrimExp f (Imp.LeafExp v _) = f v +compilePrimExp f (Imp.BinOpExp op x y) = do + (x', y', simple) <- compileBinOpLike f x y + case op of + Add {} -> simple "+" + Sub {} -> simple "-" + Mul {} -> simple "*" + FAdd {} -> simple "+" + FSub {} -> simple "-" + FMul {} -> simple "*" + FDiv {} -> simple "/" + FMod {} -> simple "%" + Xor {} -> simple "^" + And {} -> simple "&" + Or {} -> simple "|" + Shl {} -> simple "<<" + LogAnd {} -> simple "and" + LogOr {} -> simple "or" + _ -> return $ simpleCall (pretty op) [x', y'] +compilePrimExp f (Imp.ConvOpExp conv x) = do + x' <- compilePrimExp f x + return $ simpleCall (pretty conv) [x'] +compilePrimExp f (Imp.CmpOpExp cmp x y) = do + (x', y', simple) <- compileBinOpLike f x y + case cmp of + CmpEq {} -> simple "==" + FCmpLt {} -> simple "<" + FCmpLe {} -> simple "<=" + CmpLlt -> simple "<" + CmpLle -> simple "<=" + _ -> return $ simpleCall (pretty cmp) [x', y'] +compilePrimExp f (Imp.UnOpExp op exp1) = + UnOp (compileUnOp op) <$> compilePrimExp f exp1 +compilePrimExp f (Imp.FunExp h args _) = + simpleCall (futharkFun (pretty h)) <$> mapM (compilePrimExp f) args + +compileExp :: Imp.Exp -> CompilerM op s PyExp +compileExp = compilePrimExp compileLeaf + where + compileLeaf (Imp.ScalarVar vname) = + compileVar vname + compileLeaf (Imp.SizeOf t) = + return $ simpleCall (compilePrimToNp $ IntType Int32) [Integer $ primByteSize t] + compileLeaf (Imp.Index src (Imp.Count iexp) restype (Imp.Space space) _) = + join $ + asks envReadScalar + <*> compileVar src + <*> compileExp (Imp.untyped iexp) + <*> pure restype + <*> pure space + compileLeaf (Imp.Index src (Imp.Count iexp) bt _ _) = do + iexp' <- compileExp $ Imp.untyped iexp + let bt' = compilePrimType bt + nptype = compilePrimToNp bt + src' <- compileVar src + return $ simpleCall "indexArray" [src', iexp', Var bt', Var nptype] + +compileCode :: Imp.Code op -> CompilerM op s () +compileCode Imp.DebugPrint {} = + return () +compileCode (Imp.Op op) = + join $ asks envOpCompiler <*> pure op +compileCode (Imp.If cond tb fb) = do + cond' <- compileExp $ Imp.untyped cond + tb' <- collect $ compileCode tb + fb' <- collect $ compileCode fb + stm $ If cond' tb' fb' +compileCode (c1 Imp.:>>: c2) = do + compileCode c1 + compileCode c2 +compileCode (Imp.While cond body) = do + cond' <- compileExp $ Imp.untyped cond + body' <- collect $ compileCode body + stm $ While cond' body' +compileCode (Imp.For i bound body) = do + bound' <- compileExp bound + let i' = compileName i + body' <- collect $ compileCode body + counter <- pretty <$> newVName "counter" + one <- pretty <$> newVName "one" + stm $ Assign (Var i') $ simpleCall (compilePrimToNp (Imp.primExpType bound)) [Integer 0] + stm $ Assign (Var one) $ simpleCall (compilePrimToNp (Imp.primExpType bound)) [Integer 1] + stm $ + For counter (simpleCall "range" [bound']) $ + body' ++ [AssignOp "+" (Var i') (Var one)] +compileCode (Imp.SetScalar name exp1) = + stm =<< Assign <$> compileVar name <*> compileExp exp1 +compileCode Imp.DeclareMem {} = return () +compileCode (Imp.DeclareScalar v _ Cert) = do + v' <- compileVar v + stm $ Assign v' $ Var "True" +compileCode Imp.DeclareScalar {} = return () +compileCode (Imp.DeclareArray name (Space space) t vs) = + join $ + asks envStaticArray + <*> pure name + <*> pure space + <*> pure t + <*> pure vs +compileCode (Imp.DeclareArray name _ t vs) = do + let arr_name = compileName name <> "_arr" + -- It is important to store the Numpy array in a temporary variable + -- to prevent it from going "out-of-scope" before calling + -- unwrapArray (which internally uses the .ctype method); see + -- https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.ctypes.html + atInit $ + Assign (Field (Var "self") arr_name) $ case vs of + Imp.ArrayValues vs' -> + Call + (Var "np.array") + [ Arg $ List $ map compilePrimValue vs', + ArgKeyword "dtype" $ Var $ compilePrimToNp t + ] + Imp.ArrayZeros n -> + Call + (Var "np.zeros") + [ Arg $ Integer $ fromIntegral n, + ArgKeyword "dtype" $ Var $ compilePrimToNp t + ] + atInit $ + Assign (Field (Var "self") (compileName name)) $ + simpleCall "unwrapArray" [Field (Var "self") arr_name] + name' <- compileVar name + stm $ Assign name' $ Field (Var "self") (compileName name) +compileCode (Imp.Comment s code) = do + code' <- collect $ compileCode code + stm $ Comment s code' +compileCode (Imp.Assert e (Imp.ErrorMsg parts) (loc, locs)) = do + e' <- compileExp e + let onPart (Imp.ErrorString s) = return ("%s", String s) + onPart (Imp.ErrorInt32 x) = ("%d",) <$> compileExp x + onPart (Imp.ErrorInt64 x) = ("%d",) <$> compileExp x + (formatstrs, formatargs) <- unzip <$> mapM onPart parts + stm $ + Assert + e' + ( BinOp + "%" + (String $ "Error: " ++ concat formatstrs ++ "\n\nBacktrace:\n" ++ stacktrace) + (Tuple formatargs) + ) + where + stacktrace = prettyStacktrace 0 $ map locStr $ loc : locs +compileCode (Imp.Call dests fname args) = do + args' <- mapM compileArg args + dests' <- tupleOrSingle <$> mapM compileVar dests + let fname' + | isBuiltInFunction fname = futharkFun (pretty fname) + | otherwise = "self." ++ futharkFun (pretty fname) + call' = simpleCall fname' args' + -- If the function returns nothing (is called only for side + -- effects), take care not to assign to an empty tuple. + stm $ + if null dests + then Exp call' + else Assign dests' call' + where + compileArg (Imp.MemArg m) = compileVar m + compileArg (Imp.ExpArg e) = compileExp e +compileCode (Imp.SetMem dest src _) = + stm =<< Assign <$> compileVar dest <*> compileVar src +compileCode (Imp.Allocate name (Imp.Count (Imp.TPrimExp e)) (Imp.Space space)) = + join $ + asks envAllocate + <*> compileVar name + <*> compileExp e + <*> pure space +compileCode (Imp.Allocate name (Imp.Count (Imp.TPrimExp e)) _) = do + e' <- compileExp e + let allocate' = simpleCall "allocateMem" [e'] + stm =<< Assign <$> compileVar name <*> pure allocate' +compileCode (Imp.Free name _) = + stm =<< Assign <$> compileVar name <*> pure None +compileCode (Imp.Copy dest (Imp.Count destoffset) DefaultSpace src (Imp.Count srcoffset) DefaultSpace (Imp.Count size)) = do + destoffset' <- compileExp $ Imp.untyped destoffset + srcoffset' <- compileExp $ Imp.untyped srcoffset + dest' <- compileVar dest + src' <- compileVar src + size' <- compileExp $ Imp.untyped size + let offset_call1 = simpleCall "addressOffset" [dest', destoffset', Var "ct.c_byte"] + let offset_call2 = simpleCall "addressOffset" [src', srcoffset', Var "ct.c_byte"] + stm $ Exp $ simpleCall "ct.memmove" [offset_call1, offset_call2, size'] +compileCode (Imp.Copy dest (Imp.Count destoffset) destspace src (Imp.Count srcoffset) srcspace (Imp.Count size)) = do + copy <- asks envCopy + join $ + copy + <$> compileVar dest + <*> compileExp (Imp.untyped destoffset) + <*> pure destspace + <*> compileVar src + <*> compileExp (Imp.untyped srcoffset) + <*> pure srcspace + <*> compileExp (Imp.untyped size) + <*> pure (IntType Int32) -- FIXME +compileCode (Imp.Write dest (Imp.Count idx) elemtype (Imp.Space space) _ elemexp) = + join $ + asks envWriteScalar + <*> compileVar dest + <*> compileExp (Imp.untyped idx) + <*> pure elemtype + <*> pure space + <*> compileExp elemexp +compileCode (Imp.Write dest (Imp.Count idx) elemtype _ _ elemexp) = do + idx' <- compileExp $ Imp.untyped idx + elemexp' <- compileExp elemexp + dest' <- compileVar dest + let elemtype' = compilePrimType elemtype + ctype = simpleCall elemtype' [elemexp'] + stm $ Exp $ simpleCall "writeScalarArray" [dest', idx', ctype] compileCode Imp.Skip = return () diff --git a/src/Futhark/CodeGen/Backends/GenericPython/AST.hs b/src/Futhark/CodeGen/Backends/GenericPython/AST.hs --- a/src/Futhark/CodeGen/Backends/GenericPython/AST.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython/AST.hs @@ -1,88 +1,94 @@ module Futhark.CodeGen.Backends.GenericPython.AST - ( PyExp(..) - , PyIdx (..) - , PyArg (..) - , PyStmt(..) - , module Language.Futhark.Core - , PyProg(..) - , PyExcept(..) - , PyFunDef(..) - , PyClassDef(..) + ( PyExp (..), + PyIdx (..), + PyArg (..), + PyStmt (..), + module Language.Futhark.Core, + PyProg (..), + PyExcept (..), + PyFunDef (..), + PyClassDef (..), ) - where +where -import Language.Futhark.Core import Futhark.Util.Pretty - - -data UnOp = Not -- ^ Boolean negation. - | Complement -- ^ Bitwise complement. - | Negate -- ^ Numerical negation. - | Abs -- ^ Absolute/numerical value. - deriving (Eq, Show) - -data PyExp = Integer Integer - | Bool Bool - | Float Double - | String String - | RawStringLiteral String - | Var String - | BinOp String PyExp PyExp - | UnOp String PyExp - | Cond PyExp PyExp PyExp - | Index PyExp PyIdx - | Call PyExp [PyArg] - | Cast PyExp String - | Tuple [PyExp] - | List [PyExp] - | Field PyExp String - | Dict [(PyExp, PyExp)] - | Lambda String PyExp - | None - deriving (Eq, Show) +import Language.Futhark.Core -data PyIdx = IdxRange PyExp PyExp - | IdxExp PyExp - deriving (Eq, Show) +data UnOp + = -- | Boolean negation. + Not + | -- | Bitwise complement. + Complement + | -- | Numerical negation. + Negate + | -- | Absolute/numerical value. + Abs + deriving (Eq, Show) -data PyArg = ArgKeyword String PyExp - | Arg PyExp - deriving (Eq, Show) +data PyExp + = Integer Integer + | Bool Bool + | Float Double + | String String + | RawStringLiteral String + | Var String + | BinOp String PyExp PyExp + | UnOp String PyExp + | Cond PyExp PyExp PyExp + | Index PyExp PyIdx + | Call PyExp [PyArg] + | Cast PyExp String + | Tuple [PyExp] + | List [PyExp] + | Field PyExp String + | Dict [(PyExp, PyExp)] + | Lambda String PyExp + | None + deriving (Eq, Show) -data PyStmt = If PyExp [PyStmt] [PyStmt] - | Try [PyStmt] [PyExcept] - | While PyExp [PyStmt] - | For String PyExp [PyStmt] - | With PyExp [PyStmt] - | Assign PyExp PyExp - | AssignOp String PyExp PyExp - | Comment String [PyStmt] - | Assert PyExp PyExp - | Raise PyExp - | Exp PyExp - | Return PyExp - | Pass +data PyIdx + = IdxRange PyExp PyExp + | IdxExp PyExp + deriving (Eq, Show) - -- Definition-like statements. - | Import String (Maybe String) - | FunDef PyFunDef - | ClassDef PyClassDef +data PyArg + = ArgKeyword String PyExp + | Arg PyExp + deriving (Eq, Show) - -- Some arbitrary string of Python code. - | Escape String - deriving (Eq, Show) +data PyStmt + = If PyExp [PyStmt] [PyStmt] + | Try [PyStmt] [PyExcept] + | While PyExp [PyStmt] + | For String PyExp [PyStmt] + | With PyExp [PyStmt] + | Assign PyExp PyExp + | AssignOp String PyExp PyExp + | Comment String [PyStmt] + | Assert PyExp PyExp + | Raise PyExp + | Exp PyExp + | Return PyExp + | Pass + | -- Definition-like statements. + Import String (Maybe String) + | FunDef PyFunDef + | ClassDef PyClassDef + | -- Some arbitrary string of Python code. + Escape String + deriving (Eq, Show) data PyExcept = Catch PyExp [PyStmt] - deriving (Eq, Show) + deriving (Eq, Show) data PyFunDef = Def String [String] [PyStmt] - deriving (Eq, Show) + deriving (Eq, Show) data PyClassDef = Class String [PyStmt] - deriving (Eq, Show) + deriving (Eq, Show) newtype PyProg = PyProg [PyStmt] - deriving (Eq, Show) + deriving (Eq, Show) instance Pretty PyIdx where ppr (IdxExp e) = ppr e @@ -92,7 +98,6 @@ ppr (ArgKeyword k e) = text k <> equals <> ppr e ppr (Arg e) = ppr e - instance Pretty PyExp where ppr (Integer x) = ppr x ppr (Bool x) = ppr x @@ -103,102 +108,86 @@ ppr (RawStringLiteral s) = text "\"\"\"" <> text s <> text "\"\"\"" ppr (Var n) = text $ map (\x -> if x == '\'' then 'm' else x) n ppr (Field e s) = ppr e <> text "." <> text s - ppr (BinOp s e1 e2) = parens(ppr e1 <+> text s <+> ppr e2) + ppr (BinOp s e1 e2) = parens (ppr e1 <+> text s <+> ppr e2) ppr (UnOp s e) = text s <> parens (ppr e) ppr (Cond e1 e2 e3) = ppr e2 <+> text "if" <+> ppr e1 <+> text "else" <+> ppr e3 - ppr (Cast src bt) = text "ct.cast" <> - parens (ppr src <> text "," <+> - text "ct.POINTER" <> parens(text bt)) - ppr (Index src idx) = ppr src <> brackets(ppr idx) - ppr (Call fun exps) = ppr fun <> parens(commasep $ map ppr exps) - ppr (Tuple [dim]) = parens(ppr dim <> text ",") - ppr (Tuple dims) = parens(commasep $ map ppr dims) + ppr (Cast src bt) = + text "ct.cast" + <> parens + ( ppr src <> text "," + <+> text "ct.POINTER" <> parens (text bt) + ) + ppr (Index src idx) = ppr src <> brackets (ppr idx) + ppr (Call fun exps) = ppr fun <> parens (commasep $ map ppr exps) + ppr (Tuple [dim]) = parens (ppr dim <> text ",") + ppr (Tuple dims) = parens (commasep $ map ppr dims) ppr (List es) = brackets $ commasep $ map ppr es ppr (Dict kvs) = braces $ commasep $ map ppElem kvs - where ppElem (k, v) = ppr k <> colon <+> ppr v + where + ppElem (k, v) = ppr k <> colon <+> ppr v ppr (Lambda p e) = text "lambda" <+> text p <> text ":" <+> ppr e ppr None = text "None" instance Pretty PyStmt where ppr (If cond [] []) = - text "if" <+> ppr cond <> text ":" - indent 2 (text "pass") - + text "if" <+> ppr cond <> text ":" + indent 2 (text "pass") ppr (If cond [] fbranch) = - text "if" <+> ppr cond <> text ":" - indent 2 (text "pass") - text "else:" - indent 2 (stack $ map ppr fbranch) - + text "if" <+> ppr cond <> text ":" + indent 2 (text "pass") + text "else:" + indent 2 (stack $ map ppr fbranch) ppr (If cond tbranch []) = - text "if" <+> ppr cond <> text ":" - indent 2 (stack $ map ppr tbranch) - + text "if" <+> ppr cond <> text ":" + indent 2 (stack $ map ppr tbranch) ppr (If cond tbranch fbranch) = - text "if" <+> ppr cond <> text ":" - indent 2 (stack $ map ppr tbranch) - text "else:" - indent 2 (stack $ map ppr fbranch) - + text "if" <+> ppr cond <> text ":" + indent 2 (stack $ map ppr tbranch) + text "else:" + indent 2 (stack $ map ppr fbranch) ppr (Try pystms pyexcepts) = - text "try:" - indent 2 (stack $ map ppr pystms) - stack (map ppr pyexcepts) - + text "try:" + indent 2 (stack $ map ppr pystms) + stack (map ppr pyexcepts) ppr (While cond body) = - text "while" <+> ppr cond <> text ":" - indent 2 (stack $ map ppr body) - + text "while" <+> ppr cond <> text ":" + indent 2 (stack $ map ppr body) ppr (For i what body) = - text "for" <+> ppr i <+> text "in" <+> ppr what <> text ":" - indent 2 (stack $ map ppr body) - + text "for" <+> ppr i <+> text "in" <+> ppr what <> text ":" + indent 2 (stack $ map ppr body) ppr (With what body) = - text "with" <+> ppr what <> text ":" - indent 2 (stack $ map ppr body) - + text "with" <+> ppr what <> text ":" + indent 2 (stack $ map ppr body) ppr (Assign e1 e2) = ppr e1 <+> text "=" <+> ppr e2 - ppr (AssignOp op e1 e2) = ppr e1 <+> text (op ++ "=") <+> ppr e2 - ppr (Comment s body) = text "#" <> text s stack (map ppr body) - ppr (Assert e1 e2) = text "assert" <+> ppr e1 <> text "," <+> ppr e2 - ppr (Raise e) = text "raise" <+> ppr e - ppr (Exp c) = ppr c - ppr (Return e) = text "return" <+> ppr e - ppr Pass = text "pass" - ppr (Import from (Just as)) = text "import" <+> text from <+> text "as" <+> text as - ppr (Import from Nothing) = text "import" <+> text from - ppr (FunDef d) = ppr d - ppr (ClassDef d) = ppr d - ppr (Escape s) = stack $ map text $ lines s instance Pretty PyFunDef where ppr (Def fname params body) = - text "def" <+> text fname <> parens (commasep $ map ppr params) <> text ":" - indent 2 (stack (map ppr body)) + text "def" <+> text fname <> parens (commasep $ map ppr params) <> text ":" + indent 2 (stack (map ppr body)) instance Pretty PyClassDef where ppr (Class cname body) = - text "class" <+> text cname <> text ":" - indent 2 (stack (map ppr body)) + text "class" <+> text cname <> text ":" + indent 2 (stack (map ppr body)) instance Pretty PyExcept where ppr (Catch pyexp stms) = - text "except" <+> ppr pyexp <+> text "as e:" - indent 2 (stack $ map ppr stms) + text "except" <+> ppr pyexp <+> text "as e:" + indent 2 (stack $ map ppr stms) instance Pretty PyProg where ppr (PyProg stms) = stack (map ppr stms) diff --git a/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs b/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs --- a/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs @@ -1,11 +1,13 @@ {-# LANGUAGE TemplateHaskell #-} + module Futhark.CodeGen.Backends.GenericPython.Definitions - ( pyFunctions - , pyUtility - , pyValues - , pyPanic - , pyTuning - ) where + ( pyFunctions, + pyUtility, + pyValues, + pyPanic, + pyTuning, + ) +where import Data.FileEmbed diff --git a/src/Futhark/CodeGen/Backends/GenericPython/Options.hs b/src/Futhark/CodeGen/Backends/GenericPython/Options.hs --- a/src/Futhark/CodeGen/Backends/GenericPython/Options.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython/Options.hs @@ -3,11 +3,11 @@ -- Python code that will perform side effects, usually by setting some -- global variables. module Futhark.CodeGen.Backends.GenericPython.Options - ( Option (..) - , OptionArgument (..) - , generateOptionParser - ) - where + ( Option (..), + OptionArgument (..), + generateOptionParser, + ) +where import Futhark.CodeGen.Backends.GenericPython.AST @@ -16,16 +16,18 @@ -- -- When the statement is being executed, the argument (if any) will be -- stored in the variable @optarg@. -data Option = Option { optionLongName :: String - , optionShortName :: Maybe Char - , optionArgument :: OptionArgument - , optionAction :: [PyStmt] - } +data Option = Option + { optionLongName :: String, + optionShortName :: Maybe Char, + optionArgument :: OptionArgument, + optionAction :: [PyStmt] + } -- | Whether an option accepts an argument. -data OptionArgument = NoArgument - | RequiredArgument String - | OptionalArgument +data OptionArgument + = NoArgument + | RequiredArgument String + | OptionalArgument -- | Generate option parsing code that accepts the given command line options. Will read from @sys.argv@. -- @@ -33,41 +35,59 @@ -- terminate with error code 1. generateOptionParser :: [Option] -> [PyStmt] generateOptionParser options = - [Assign (Var "parser") - (Call (Var "argparse.ArgumentParser") - [ArgKeyword "description" $ - String "A compiled Futhark program."])] ++ - map parseOption options ++ - [Assign (Var "parser_result") $ - Call (Var "vars") [Arg $ Call (Var "parser.parse_args") [Arg $ Var "sys.argv[1:]"]]] ++ - map executeOption options - where parseOption option = - Exp $ Call (Var "parser.add_argument") $ - map (Arg . String) name_args ++ - argument_args - where name_args = maybe id ((:) . ('-':) . (:[])) (optionShortName option) - ["--" ++ optionLongName option] - argument_args = case optionArgument option of - RequiredArgument t -> - [ArgKeyword "action" (String "append"), - ArgKeyword "default" $ List [], - ArgKeyword "type" $ Var t] - - NoArgument -> - [ArgKeyword "action" (String "append_const"), - ArgKeyword "default" $ List [], - ArgKeyword "const" None] - - OptionalArgument -> - [ArgKeyword "action" (String "append"), - ArgKeyword "default" $ List [], - ArgKeyword "nargs" $ String "?"] + [ Assign + (Var "parser") + ( Call + (Var "argparse.ArgumentParser") + [ ArgKeyword "description" $ + String "A compiled Futhark program." + ] + ) + ] + ++ map parseOption options + ++ [ Assign (Var "parser_result") $ + Call (Var "vars") [Arg $ Call (Var "parser.parse_args") [Arg $ Var "sys.argv[1:]"]] + ] + ++ map executeOption options + where + parseOption option = + Exp $ + Call (Var "parser.add_argument") $ + map (Arg . String) name_args + ++ argument_args + where + name_args = + maybe + id + ((:) . ('-' :) . (: [])) + (optionShortName option) + ["--" ++ optionLongName option] + argument_args = case optionArgument option of + RequiredArgument t -> + [ ArgKeyword "action" (String "append"), + ArgKeyword "default" $ List [], + ArgKeyword "type" $ Var t + ] + NoArgument -> + [ ArgKeyword "action" (String "append_const"), + ArgKeyword "default" $ List [], + ArgKeyword "const" None + ] + OptionalArgument -> + [ ArgKeyword "action" (String "append"), + ArgKeyword "default" $ List [], + ArgKeyword "nargs" $ String "?" + ] - executeOption option = - For "optarg" (Index (Var "parser_result") $ - IdxExp $ String $ fieldName option) $ - optionAction option + executeOption option = + For + "optarg" + ( Index (Var "parser_result") $ + IdxExp $ String $ fieldName option + ) + $ optionAction option - fieldName = map escape . optionLongName - where escape '-' = '_' - escape c = c + fieldName = map escape . optionLongName + where + escape '-' = '_' + escape c = c 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 @@ -1,131 +1,179 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TupleSections #-} + module Futhark.CodeGen.Backends.PyOpenCL - ( compileProg - ) where + ( compileProg, + ) +where import Control.Monad import qualified Data.Map as M - -import Futhark.IR.KernelsMem (Prog, KernelsMem) -import Futhark.CodeGen.Backends.PyOpenCL.Boilerplate import qualified Futhark.CodeGen.Backends.GenericPython as Py -import qualified Futhark.CodeGen.ImpCode.OpenCL as Imp -import qualified Futhark.CodeGen.ImpGen.OpenCL as ImpGen import Futhark.CodeGen.Backends.GenericPython.AST -import Futhark.CodeGen.Backends.GenericPython.Options import Futhark.CodeGen.Backends.GenericPython.Definitions +import Futhark.CodeGen.Backends.GenericPython.Options +import Futhark.CodeGen.Backends.PyOpenCL.Boilerplate +import qualified Futhark.CodeGen.ImpCode.OpenCL as Imp +import qualified Futhark.CodeGen.ImpGen.OpenCL as ImpGen +import Futhark.IR.KernelsMem (KernelsMem, Prog) import Futhark.MonadFreshNames import Futhark.Util (zEncodeString) --maybe pass the config file rather than multiple arguments -compileProg :: MonadFreshNames m => - Maybe String -> Prog KernelsMem -> m (ImpGen.Warnings, String) +compileProg :: + MonadFreshNames m => + Maybe String -> + Prog KernelsMem -> + m (ImpGen.Warnings, String) compileProg module_name prog = do - (ws, Imp.Program opencl_code opencl_prelude - kernels types sizes failures prog') <- + ( ws, + Imp.Program + opencl_code + opencl_prelude + kernels + types + sizes + failures + prog' + ) <- ImpGen.compileProg prog --prepare the strings for assigning the kernels and set them as global - let assign = unlines $ - map (\x -> pretty $ - Assign (Var ("self."++zEncodeString (nameToString x)++"_var")) - (Var $ "program."++zEncodeString (nameToString x))) $ - M.keys kernels + let assign = + unlines $ + map + ( \x -> + pretty $ + Assign + (Var ("self." ++ zEncodeString (nameToString x) ++ "_var")) + (Var $ "program." ++ zEncodeString (nameToString x)) + ) + $ M.keys kernels let defines = - [Assign (Var "synchronous") $ Bool False, - Assign (Var "preferred_platform") None, - Assign (Var "preferred_device") None, - Assign (Var "default_threshold") None, - Assign (Var "default_group_size") None, - Assign (Var "default_num_groups") None, - Assign (Var "default_tile_size") None, - Assign (Var "fut_opencl_src") $ RawStringLiteral $ opencl_prelude ++ opencl_code, - Escape pyValues, - Escape pyFunctions, - Escape pyPanic, - Escape pyTuning] + [ Assign (Var "synchronous") $ Bool False, + Assign (Var "preferred_platform") None, + Assign (Var "preferred_device") None, + Assign (Var "default_threshold") None, + Assign (Var "default_group_size") None, + Assign (Var "default_num_groups") None, + Assign (Var "default_tile_size") None, + Assign (Var "fut_opencl_src") $ RawStringLiteral $ opencl_prelude ++ opencl_code, + Escape pyValues, + Escape pyFunctions, + Escape pyPanic, + Escape pyTuning + ] - let imports = [Import "sys" Nothing, - Import "numpy" $ Just "np", - Import "ctypes" $ Just "ct", - Escape openClPrelude, - Import "pyopencl.array" Nothing, - Import "time" Nothing] + let imports = + [ Import "sys" Nothing, + Import "numpy" $ Just "np", + Import "ctypes" $ Just "ct", + Escape openClPrelude, + Import "pyopencl.array" Nothing, + Import "time" Nothing + ] - let constructor = Py.Constructor [ "self" - , "command_queue=None" - , "interactive=False" - , "platform_pref=preferred_platform" - , "device_pref=preferred_device" - , "default_group_size=default_group_size" - , "default_num_groups=default_num_groups" - , "default_tile_size=default_tile_size" - , "default_threshold=default_threshold" - , "sizes=sizes"] - [Escape $ openClInit types assign sizes failures] - options = [ Option { optionLongName = "platform" - , optionShortName = Just 'p' - , optionArgument = RequiredArgument "str" - , optionAction = - [ Assign (Var "preferred_platform") $ Var "optarg" ] - } - , Option { optionLongName = "device" - , optionShortName = Just 'd' - , optionArgument = RequiredArgument "str" - , optionAction = - [ Assign (Var "preferred_device") $ Var "optarg" ] - } - , Option { optionLongName = "default-threshold" - , optionShortName = Nothing - , optionArgument = RequiredArgument "int" - , optionAction = - [ Assign (Var "default_threshold") $ Var "optarg" ] - } - , Option { optionLongName = "default-group-size" - , optionShortName = Nothing - , optionArgument = RequiredArgument "int" - , optionAction = - [ Assign (Var "default_group_size") $ Var "optarg" ] - } - , Option { optionLongName = "default-num-groups" - , optionShortName = Nothing - , optionArgument = RequiredArgument "int" - , optionAction = - [ Assign (Var "default_num_groups") $ Var "optarg" ] - } - , Option { optionLongName = "default-tile-size" - , optionShortName = Nothing - , optionArgument = RequiredArgument "int" - , optionAction = - [ Assign (Var "default_tile_size") $ Var "optarg" ] - } - , Option { optionLongName = "size" - , optionShortName = Nothing - , optionArgument = RequiredArgument "size_assignment" - , optionAction = - [Assign (Index (Var "sizes") - (IdxExp (Index (Var "optarg") - (IdxExp (Integer 0))))) - (Index (Var "optarg") (IdxExp (Integer 1))) - ] - } + let constructor = + Py.Constructor + [ "self", + "command_queue=None", + "interactive=False", + "platform_pref=preferred_platform", + "device_pref=preferred_device", + "default_group_size=default_group_size", + "default_num_groups=default_num_groups", + "default_tile_size=default_tile_size", + "default_threshold=default_threshold", + "sizes=sizes" + ] + [Escape $ openClInit types assign sizes failures] + options = + [ Option + { optionLongName = "platform", + optionShortName = Just 'p', + optionArgument = RequiredArgument "str", + optionAction = + [Assign (Var "preferred_platform") $ Var "optarg"] + }, + Option + { optionLongName = "device", + optionShortName = Just 'd', + optionArgument = RequiredArgument "str", + optionAction = + [Assign (Var "preferred_device") $ Var "optarg"] + }, + Option + { optionLongName = "default-threshold", + optionShortName = Nothing, + optionArgument = RequiredArgument "int", + optionAction = + [Assign (Var "default_threshold") $ Var "optarg"] + }, + Option + { optionLongName = "default-group-size", + optionShortName = Nothing, + optionArgument = RequiredArgument "int", + optionAction = + [Assign (Var "default_group_size") $ Var "optarg"] + }, + Option + { optionLongName = "default-num-groups", + optionShortName = Nothing, + optionArgument = RequiredArgument "int", + optionAction = + [Assign (Var "default_num_groups") $ Var "optarg"] + }, + Option + { optionLongName = "default-tile-size", + optionShortName = Nothing, + optionArgument = RequiredArgument "int", + optionAction = + [Assign (Var "default_tile_size") $ Var "optarg"] + }, + Option + { optionLongName = "size", + optionShortName = Nothing, + optionArgument = RequiredArgument "size_assignment", + optionAction = + [ Assign + ( Index + (Var "sizes") + ( IdxExp + ( Index + (Var "optarg") + (IdxExp (Integer 0)) + ) + ) + ) + (Index (Var "optarg") (IdxExp (Integer 1))) ] + } + ] - (ws,) <$> - Py.compileProg module_name constructor imports defines operations () - [Exp $ Py.simpleCall "sync" [Var "self"]] options prog' - where operations :: Py.Operations Imp.OpenCL () - operations = Py.Operations - { Py.opsCompiler = callKernel - , Py.opsWriteScalar = writeOpenCLScalar - , Py.opsReadScalar = readOpenCLScalar - , Py.opsAllocate = allocateOpenCLBuffer - , Py.opsCopy = copyOpenCLMemory - , Py.opsStaticArray = staticOpenCLArray - , Py.opsEntryOutput = packArrayOutput - , Py.opsEntryInput = unpackArrayInput - } + (ws,) + <$> Py.compileProg + module_name + constructor + imports + defines + operations + () + [Exp $ Py.simpleCall "sync" [Var "self"]] + options + prog' + where + operations :: Py.Operations Imp.OpenCL () + operations = + Py.Operations + { Py.opsCompiler = callKernel, + Py.opsWriteScalar = writeOpenCLScalar, + Py.opsReadScalar = readOpenCLScalar, + Py.opsAllocate = allocateOpenCLBuffer, + Py.opsCopy = copyOpenCLMemory, + Py.opsStaticArray = staticOpenCLArray, + Py.opsEntryOutput = packArrayOutput, + Py.opsEntryInput = unpackArrayInput + } -- We have many casts to 'long', because PyOpenCL may get confused at -- the 32-bit numbers that ImpCode uses for offsets and the like. @@ -135,18 +183,20 @@ callKernel :: Py.OpCompiler Imp.OpenCL () callKernel (Imp.GetSize v key) = do v' <- Py.compileVar v - Py.stm $ Assign v' $ - Index (Var "self.sizes") (IdxExp $ String $ pretty key) + Py.stm $ + Assign v' $ + Index (Var "self.sizes") (IdxExp $ String $ pretty key) callKernel (Imp.CmpSizeLe v key x) = do v' <- Py.compileVar v x' <- Py.compileExp x - Py.stm $ Assign v' $ - BinOp "<=" (Index (Var "self.sizes") (IdxExp $ String $ pretty key)) x' + Py.stm $ + Assign v' $ + BinOp "<=" (Index (Var "self.sizes") (IdxExp $ String $ pretty key)) x' callKernel (Imp.GetSizeMax v size_class) = do v' <- Py.compileVar v - Py.stm $ Assign v' $ - Var $ "self.max_" ++ pretty size_class - + Py.stm $ + Assign v' $ + Var $ "self.max_" ++ pretty size_class callKernel (Imp.LaunchKernel safety name args num_workgroups workgroup_size) = do num_workgroups' <- mapM (fmap asLong . Py.compileExp) num_workgroups workgroup_size' <- mapM (fmap asLong . Py.compileExp) workgroup_size @@ -158,45 +208,67 @@ Py.stm $ If cond body [] when (safety >= Imp.SafetyFull) $ - Py.stm $ Assign (Var "self.failure_is_an_option") $ - Py.compilePrimValue (Imp.IntValue (Imp.Int32Value 1)) - - where mult_exp = BinOp "*" + Py.stm $ + Assign (Var "self.failure_is_an_option") $ + Py.compilePrimValue (Imp.IntValue (Imp.Int32Value 1)) + where + mult_exp = BinOp "*" -launchKernel :: Imp.KernelName -> Imp.KernelSafety -> [PyExp] -> [PyExp] -> [Imp.KernelArg] - -> Py.CompilerM op s () +launchKernel :: + Imp.KernelName -> + Imp.KernelSafety -> + [PyExp] -> + [PyExp] -> + [Imp.KernelArg] -> + Py.CompilerM op s () launchKernel kernel_name safety kernel_dims workgroup_dims args = do let kernel_dims' = Tuple kernel_dims workgroup_dims' = Tuple workgroup_dims kernel_name' = "self." ++ zEncodeString (nameToString kernel_name) ++ "_var" args' <- mapM processKernelArg args - let failure_args = take (Imp.numFailureParams safety) - [Var "self.global_failure", - Var "self.failure_is_an_option", - Var "self.global_failure_args"] - Py.stm $ Exp $ Py.simpleCall (kernel_name' ++ ".set_args") $ - failure_args ++ args' - Py.stm $ Exp $ Py.simpleCall "cl.enqueue_nd_range_kernel" - [Var "self.queue", Var kernel_name', kernel_dims', workgroup_dims'] + let failure_args = + take + (Imp.numFailureParams safety) + [ Var "self.global_failure", + Var "self.failure_is_an_option", + Var "self.global_failure_args" + ] + Py.stm $ + Exp $ + Py.simpleCall (kernel_name' ++ ".set_args") $ + failure_args ++ args' + Py.stm $ + Exp $ + Py.simpleCall + "cl.enqueue_nd_range_kernel" + [Var "self.queue", Var kernel_name', kernel_dims', workgroup_dims'] finishIfSynchronous - where processKernelArg :: Imp.KernelArg -> Py.CompilerM op s PyExp - processKernelArg (Imp.ValueKArg e bt) = do - e' <- Py.compileExp e - return $ Py.simpleCall (Py.compilePrimToNp bt) [e'] - processKernelArg (Imp.MemKArg v) = Py.compileVar v - processKernelArg (Imp.SharedMemoryKArg (Imp.Count num_bytes)) = do - num_bytes' <- Py.compileExp num_bytes - return $ Py.simpleCall "cl.LocalMemory" [asLong num_bytes'] + where + processKernelArg :: Imp.KernelArg -> Py.CompilerM op s PyExp + processKernelArg (Imp.ValueKArg e bt) = do + e' <- Py.compileExp e + return $ Py.simpleCall (Py.compilePrimToNp bt) [e'] + processKernelArg (Imp.MemKArg v) = Py.compileVar v + processKernelArg (Imp.SharedMemoryKArg (Imp.Count num_bytes)) = do + num_bytes' <- Py.compileExp num_bytes + return $ Py.simpleCall "cl.LocalMemory" [asLong num_bytes'] writeOpenCLScalar :: Py.WriteScalar Imp.OpenCL () writeOpenCLScalar mem i bt "device" val = do - let nparr = Call (Var "np.array") - [Arg val, ArgKeyword "dtype" $ Var $ Py.compilePrimType bt] - Py.stm $ Exp $ Call (Var "cl.enqueue_copy") - [Arg $ Var "self.queue", Arg mem, Arg nparr, - ArgKeyword "device_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt), - ArgKeyword "is_blocking" $ Var "synchronous"] - + let nparr = + Call + (Var "np.array") + [Arg val, ArgKeyword "dtype" $ Var $ Py.compilePrimType bt] + Py.stm $ + Exp $ + Call + (Var "cl.enqueue_copy") + [ Arg $ Var "self.queue", + Arg mem, + Arg nparr, + ArgKeyword "device_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt), + ArgKeyword "is_blocking" $ Var "synchronous" + ] writeOpenCLScalar _ _ _ space _ = error $ "Cannot write to '" ++ space ++ "' memory space." @@ -204,25 +276,33 @@ readOpenCLScalar mem i bt "device" = do val <- newVName "read_res" let val' = Var $ pretty val - let nparr = Call (Var "np.empty") - [Arg $ Integer 1, - ArgKeyword "dtype" (Var $ Py.compilePrimType bt)] + let nparr = + Call + (Var "np.empty") + [ Arg $ Integer 1, + ArgKeyword "dtype" (Var $ Py.compilePrimType bt) + ] Py.stm $ Assign val' nparr - Py.stm $ Exp $ Call (Var "cl.enqueue_copy") - [Arg $ Var "self.queue", Arg val', Arg mem, - ArgKeyword "device_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt), - ArgKeyword "is_blocking" $ Var "synchronous"] + Py.stm $ + Exp $ + Call + (Var "cl.enqueue_copy") + [ Arg $ Var "self.queue", + Arg val', + Arg mem, + ArgKeyword "device_offset" $ BinOp "*" (asLong i) (Integer $ Imp.primByteSize bt), + ArgKeyword "is_blocking" $ Var "synchronous" + ] Py.stm $ Exp $ Py.simpleCall "sync" [Var "self"] return $ Index val' $ IdxExp $ Integer 0 - readOpenCLScalar _ _ _ space = error $ "Cannot read from '" ++ space ++ "' memory space." allocateOpenCLBuffer :: Py.Allocate Imp.OpenCL () allocateOpenCLBuffer mem size "device" = - Py.stm $ Assign mem $ - Py.simpleCall "opencl_alloc" [Var "self", size, String $ pretty mem] - + Py.stm $ + Assign mem $ + Py.simpleCall "opencl_alloc" [Var "self", size, String $ pretty mem] allocateOpenCLBuffer _ _ space = error $ "Cannot allocate in '" ++ space ++ "' space" @@ -231,53 +311,73 @@ let divide = BinOp "//" nbytes (Integer $ Imp.primByteSize bt) end = BinOp "+" destidx divide dest = Index destmem (IdxRange destidx end) - Py.stm $ ifNotZeroSize nbytes $ - Exp $ Call (Var "cl.enqueue_copy") - [Arg $ Var "self.queue", Arg dest, Arg srcmem, - ArgKeyword "device_offset" $ asLong srcidx, - ArgKeyword "is_blocking" $ Var "synchronous"] - + Py.stm $ + ifNotZeroSize nbytes $ + Exp $ + Call + (Var "cl.enqueue_copy") + [ Arg $ Var "self.queue", + Arg dest, + Arg srcmem, + ArgKeyword "device_offset" $ asLong srcidx, + ArgKeyword "is_blocking" $ Var "synchronous" + ] copyOpenCLMemory destmem destidx (Imp.Space "device") srcmem srcidx Imp.DefaultSpace nbytes bt = do let divide = BinOp "//" nbytes (Integer $ Imp.primByteSize bt) end = BinOp "+" srcidx divide src = Index srcmem (IdxRange srcidx end) - Py.stm $ ifNotZeroSize nbytes $ - Exp $ Call (Var "cl.enqueue_copy") - [Arg $ Var "self.queue", Arg destmem, Arg src, - ArgKeyword "device_offset" $ asLong destidx, - ArgKeyword "is_blocking" $ Var "synchronous"] - + Py.stm $ + ifNotZeroSize nbytes $ + Exp $ + Call + (Var "cl.enqueue_copy") + [ Arg $ Var "self.queue", + Arg destmem, + Arg src, + ArgKeyword "device_offset" $ asLong destidx, + ArgKeyword "is_blocking" $ Var "synchronous" + ] copyOpenCLMemory destmem destidx (Imp.Space "device") srcmem srcidx (Imp.Space "device") nbytes _ = do - Py.stm $ ifNotZeroSize nbytes $ - Exp $ Call (Var "cl.enqueue_copy") - [Arg $ Var "self.queue", Arg destmem, Arg srcmem, - ArgKeyword "dest_offset" $ asLong destidx, - ArgKeyword "src_offset" $ asLong srcidx, - ArgKeyword "byte_count" $ asLong nbytes] + Py.stm $ + ifNotZeroSize nbytes $ + Exp $ + Call + (Var "cl.enqueue_copy") + [ Arg $ Var "self.queue", + Arg destmem, + Arg srcmem, + ArgKeyword "dest_offset" $ asLong destidx, + ArgKeyword "src_offset" $ asLong srcidx, + ArgKeyword "byte_count" $ asLong nbytes + ] finishIfSynchronous - copyOpenCLMemory destmem destidx Imp.DefaultSpace srcmem srcidx Imp.DefaultSpace nbytes _ = Py.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes - -copyOpenCLMemory _ _ destspace _ _ srcspace _ _= +copyOpenCLMemory _ _ destspace _ _ srcspace _ _ = error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace staticOpenCLArray :: Py.StaticArray Imp.OpenCL () staticOpenCLArray name "device" t vs = do mapM_ Py.atInit <=< Py.collect $ do -- Create host-side Numpy array with intended values. - Py.stm $ Assign (Var name') $ case vs of - Imp.ArrayValues vs' -> - Call (Var "np.array") - [Arg $ List $ map Py.compilePrimValue vs', - ArgKeyword "dtype" $ Var $ Py.compilePrimToNp t] - Imp.ArrayZeros n -> - Call (Var "np.zeros") - [Arg $ Integer $ fromIntegral n, - ArgKeyword "dtype" $ Var $ Py.compilePrimToNp t] + Py.stm $ + Assign (Var name') $ case vs of + Imp.ArrayValues vs' -> + Call + (Var "np.array") + [ Arg $ List $ map Py.compilePrimValue vs', + ArgKeyword "dtype" $ Var $ Py.compilePrimToNp t + ] + Imp.ArrayZeros n -> + Call + (Var "np.zeros") + [ Arg $ Integer $ fromIntegral n, + ArgKeyword "dtype" $ Var $ Py.compilePrimToNp t + ] - let num_elems = case vs of Imp.ArrayValues vs' -> length vs' - Imp.ArrayZeros n -> n + let num_elems = case vs of + Imp.ArrayValues vs' -> length vs' + Imp.ArrayZeros n -> n -- Create memory block on the device. static_mem <- newVName "static_mem" @@ -285,58 +385,74 @@ allocateOpenCLBuffer (Var (Py.compileName static_mem)) size "device" -- Copy Numpy array to the device memory block. - Py.stm $ ifNotZeroSize size $ - Exp $ Call (Var "cl.enqueue_copy") - [Arg $ Var "self.queue", - Arg $ Var $ Py.compileName static_mem, - Arg $ Call (Var "normaliseArray") [Arg (Var name')], - ArgKeyword "is_blocking" $ Var "synchronous"] + Py.stm $ + ifNotZeroSize size $ + Exp $ + Call + (Var "cl.enqueue_copy") + [ Arg $ Var "self.queue", + Arg $ Var $ Py.compileName static_mem, + Arg $ Call (Var "normaliseArray") [Arg (Var name')], + ArgKeyword "is_blocking" $ Var "synchronous" + ] -- Store the memory block for later reference. - Py.stm $ Assign (Field (Var "self") name') $ - Var $ Py.compileName static_mem + Py.stm $ + Assign (Field (Var "self") name') $ + Var $ Py.compileName static_mem Py.stm $ Assign (Var name') (Field (Var "self") name') - where name' = Py.compileName name + where + name' = Py.compileName name staticOpenCLArray _ space _ _ = error $ "PyOpenCL backend cannot create static array in memory space '" ++ space ++ "'" packArrayOutput :: Py.EntryOutput Imp.OpenCL () packArrayOutput mem "device" bt ept dims = do mem' <- Py.compileVar mem - return $ Call (Var "cl.array.Array") - [Arg $ Var "self.queue", - Arg $ Tuple $ map Py.compileDim dims, - Arg $ Var $ Py.compilePrimTypeExt bt ept, - ArgKeyword "data" mem'] + return $ + Call + (Var "cl.array.Array") + [ Arg $ Var "self.queue", + Arg $ Tuple $ map Py.compileDim dims, + Arg $ Var $ Py.compilePrimTypeExt bt ept, + ArgKeyword "data" mem' + ] packArrayOutput _ sid _ _ _ = error $ "Cannot return array from " ++ sid ++ " space." unpackArrayInput :: Py.EntryInput Imp.OpenCL () unpackArrayInput mem "device" t s dims e = do let type_is_ok = - BinOp "and" - (BinOp "in" (Py.simpleCall "type" [e]) (List [Var "np.ndarray", Var "cl.array.Array"])) - (BinOp "==" (Field e "dtype") (Var (Py.compilePrimToExtNp t s))) + BinOp + "and" + (BinOp "in" (Py.simpleCall "type" [e]) (List [Var "np.ndarray", Var "cl.array.Array"])) + (BinOp "==" (Field e "dtype") (Var (Py.compilePrimToExtNp t s))) Py.stm $ Assert type_is_ok $ String "Parameter has unexpected type" - zipWithM_ (Py.unpackDim e) dims [0..] + zipWithM_ (Py.unpackDim e) dims [0 ..] let memsize' = Py.simpleCall "np.int64" [Field e "nbytes"] pyOpenCLArrayCase = [Assign mem $ Field e "data"] numpyArrayCase <- Py.collect $ do allocateOpenCLBuffer mem memsize' "device" - Py.stm $ ifNotZeroSize memsize' $ - Exp $ Call (Var "cl.enqueue_copy") - [Arg $ Var "self.queue", - Arg mem, - Arg $ Call (Var "normaliseArray") [Arg e], - ArgKeyword "is_blocking" $ Var "synchronous"] + Py.stm $ + ifNotZeroSize memsize' $ + Exp $ + Call + (Var "cl.enqueue_copy") + [ Arg $ Var "self.queue", + Arg mem, + Arg $ Call (Var "normaliseArray") [Arg e], + ArgKeyword "is_blocking" $ Var "synchronous" + ] - Py.stm $ If (BinOp "==" (Py.simpleCall "type" [e]) (Var "cl.array.Array")) - pyOpenCLArrayCase - numpyArrayCase + Py.stm $ + If + (BinOp "==" (Py.simpleCall "type" [e]) (Var "cl.array.Array")) + pyOpenCLArrayCase + numpyArrayCase unpackArrayInput _ sid _ _ _ _ = error $ "Cannot accept array from " ++ sid ++ " space." 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 @@ -1,25 +1,33 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE TemplateHaskell #-} + -- | Various boilerplate definitions for the PyOpenCL backend. module Futhark.CodeGen.Backends.PyOpenCL.Boilerplate - ( openClInit - , openClPrelude - ) where + ( openClInit, + openClPrelude, + ) +where import Control.Monad.Identity import Data.FileEmbed -import qualified Data.Text as T import qualified Data.Map as M -import NeatInterpolation (text) - +import qualified Data.Text as T +import qualified Futhark.CodeGen.Backends.GenericPython as Py +import Futhark.CodeGen.Backends.GenericPython.AST import Futhark.CodeGen.ImpCode.OpenCL - (PrimType(..), SizeClass(..), sizeDefault, - FailureMsg(..), ErrorMsg(..), ErrorMsgPart(..), errorMsgArgTypes) + ( ErrorMsg (..), + ErrorMsgPart (..), + FailureMsg (..), + PrimType (..), + SizeClass (..), + errorMsgArgTypes, + sizeDefault, + untyped, + ) import Futhark.CodeGen.OpenCL.Heuristics -import Futhark.CodeGen.Backends.GenericPython.AST -import qualified Futhark.CodeGen.Backends.GenericPython as Py import Futhark.Util.Pretty (prettyText) +import NeatInterpolation (text) errorMsgNumArgs :: ErrorMsg a -> Int errorMsgNumArgs = length . errorMsgArgTypes @@ -32,7 +40,9 @@ -- @initiatialize_opencl_object@ procedure. Should be put in the -- class constructor. openClInit :: [PrimType] -> String -> M.Map Name SizeClass -> [FailureMsg] -> String -openClInit types assign sizes failures = T.unpack [text| +openClInit types assign sizes failures = + T.unpack + [text| size_heuristics=$size_heuristics self.global_failure_args_max = $max_num_args self.failure_msgs=$failure_msgs @@ -52,51 +62,70 @@ all_sizes=$sizes') $assign' |] - where assign' = T.pack assign - size_heuristics = prettyText $ sizeHeuristicsToPython sizeHeuristicsTable - types' = prettyText $ map (show . pretty) types -- Looks enough like Python. - sizes' = prettyText $ sizeClassesToPython sizes - max_num_args = prettyText $ foldl max 0 $ map (errorMsgNumArgs . failureError) failures - failure_msgs = prettyText $ List $ map formatFailure failures + where + assign' = T.pack assign + size_heuristics = prettyText $ sizeHeuristicsToPython sizeHeuristicsTable + types' = prettyText $ map (show . pretty) types -- Looks enough like Python. + sizes' = prettyText $ sizeClassesToPython sizes + max_num_args = prettyText $ foldl max 0 $ map (errorMsgNumArgs . failureError) failures + failure_msgs = prettyText $ List $ map formatFailure failures formatFailure :: FailureMsg -> PyExp formatFailure (FailureMsg (ErrorMsg parts) backtrace) = String $ concatMap onPart parts ++ "\n" ++ formatEscape backtrace - where formatEscape = let escapeChar '{' = "{{" - escapeChar '}' = "}}" - escapeChar c = [c] - in concatMap escapeChar + where + formatEscape = + let escapeChar '{' = "{{" + escapeChar '}' = "}}" + escapeChar c = [c] + in concatMap escapeChar - onPart (ErrorString s) = formatEscape s - onPart ErrorInt32{} = "{}" + onPart (ErrorString s) = formatEscape s + onPart ErrorInt32 {} = "{}" + onPart ErrorInt64 {} = "{}" sizeClassesToPython :: M.Map Name SizeClass -> PyExp sizeClassesToPython = Dict . map f . M.toList - where f (size_name, size_class) = - (String $ pretty size_name, - Dict [(String "class", String $ pretty size_class), - (String "value", maybe None (Integer . fromIntegral) $ - sizeDefault size_class)]) + where + f (size_name, size_class) = + ( String $ pretty size_name, + Dict + [ (String "class", String $ pretty size_class), + ( String "value", + maybe None (Integer . fromIntegral) $ + sizeDefault size_class + ) + ] + ) sizeHeuristicsToPython :: [SizeHeuristic] -> PyExp sizeHeuristicsToPython = List . map f - where f (SizeHeuristic platform_name device_type which what) = - Tuple [String platform_name, - clDeviceType device_type, - which', - what'] - - where clDeviceType DeviceGPU = Var "cl.device_type.GPU" - clDeviceType DeviceCPU = Var "cl.device_type.CPU" + where + f (SizeHeuristic platform_name device_type which what) = + Tuple + [ String platform_name, + clDeviceType device_type, + which', + what' + ] + where + clDeviceType DeviceGPU = Var "cl.device_type.GPU" + clDeviceType DeviceCPU = Var "cl.device_type.CPU" - which' = case which of LockstepWidth -> String "lockstep_width" - NumGroups -> String "num_groups" - GroupSize -> String "group_size" - TileSize -> String "tile_size" - Threshold -> String "threshold" + which' = case which of + LockstepWidth -> String "lockstep_width" + NumGroups -> String "num_groups" + GroupSize -> String "group_size" + TileSize -> String "tile_size" + Threshold -> String "threshold" - what' = Lambda "device" $ runIdentity $ Py.compilePrimExp onLeaf what + what' = + Lambda "device" $ + runIdentity $ + Py.compilePrimExp onLeaf $ untyped what - onLeaf (DeviceInfo s) = - pure $ Py.simpleCall "device.get_info" - [Py.simpleCall "getattr" [Var "cl.device_info", String s]] + onLeaf (DeviceInfo s) = + pure $ + Py.simpleCall + "device.get_info" + [Py.simpleCall "getattr" [Var "cl.device_info", String s]] 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,76 +1,83 @@ {-# 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. module Futhark.CodeGen.Backends.SequentialC - ( compileProg - , GC.CParts(..) - , GC.asLibrary - , GC.asExecutable - ) where + ( compileProg, + GC.CParts (..), + GC.asLibrary, + GC.asExecutable, + ) +where import Control.Monad - -import qualified Language.C.Quote.OpenCL as C - -import Futhark.IR.SeqMem +import qualified Futhark.CodeGen.Backends.GenericC as GC import qualified Futhark.CodeGen.ImpCode.Sequential as Imp import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGen -import qualified Futhark.CodeGen.Backends.GenericC as GC +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] []) <=< - ImpGen.compileProg - where operations :: GC.Operations Imp.Sequential () - operations = GC.defaultOperations - { GC.opsCompiler = const $ return () - , GC.opsCopy = copySequentialMemory - } + (GC.compileProg "c" operations generateContext "" [DefaultSpace] []) + <=< ImpGen.compileProg + where + operations :: GC.Operations Imp.Sequential () + operations = + GC.defaultOperations + { GC.opsCompiler = const $ return (), + GC.opsCopy = copySequentialMemory + } - generateContext = do - cfg <- GC.publicDef "context_config" GC.InitDecl $ \s -> - ([C.cedecl|struct $id:s;|], - [C.cedecl|struct $id:s { int debugging; };|]) + 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) { + 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) { + 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) { + 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) { + 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 + (fields, init_fields) <- GC.contextContents - ctx <- GC.publicDef "context" GC.InitDecl $ \s -> - ([C.cedecl|struct $id:s;|], - [C.cedecl|struct $id:s { + ctx <- GC.publicDef "context" GC.InitDecl $ \s -> + ( [C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { int detail_memory; int debugging; int profiling; @@ -78,11 +85,12 @@ char *error; 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) { + 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; @@ -95,30 +103,33 @@ $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) { + 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) { + 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.publicDef_ "context_clear_caches" GC.MiscDecl $ \s -> - ([C.cedecl|int $id:s(struct $id:ctx* ctx);|], - [C.cedecl|int $id:s(struct $id:ctx* ctx) { + GC.publicDef_ "context_clear_caches" 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; - }|]) - + }|] + ) copySequentialMemory :: GC.Copy Imp.Sequential () copySequentialMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes = diff --git a/src/Futhark/CodeGen/Backends/SequentialPython.hs b/src/Futhark/CodeGen/Backends/SequentialPython.hs --- a/src/Futhark/CodeGen/Backends/SequentialPython.hs +++ b/src/Futhark/CodeGen/Backends/SequentialPython.hs @@ -1,37 +1,49 @@ module Futhark.CodeGen.Backends.SequentialPython - ( compileProg - ) where + ( compileProg, + ) +where import Control.Monad - -import Futhark.IR.SeqMem -import qualified Futhark.CodeGen.ImpCode.Sequential as Imp -import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGen import qualified Futhark.CodeGen.Backends.GenericPython as GenericPython -import Futhark.CodeGen.Backends.GenericPython.Definitions import Futhark.CodeGen.Backends.GenericPython.AST +import Futhark.CodeGen.Backends.GenericPython.Definitions +import qualified Futhark.CodeGen.ImpCode.Sequential as Imp +import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGen +import Futhark.IR.SeqMem import Futhark.MonadFreshNames -compileProg :: MonadFreshNames m => - Maybe String -> Prog SeqMem -> m (ImpGen.Warnings, String) +compileProg :: + MonadFreshNames m => + Maybe String -> + Prog SeqMem -> + m (ImpGen.Warnings, String) compileProg module_name = - ImpGen.compileProg >=> - traverse (GenericPython.compileProg - module_name - GenericPython.emptyConstructor - imports - defines - operations () [] []) - where imports = [Import "sys" Nothing, - Import "numpy" $ Just "np", - Import "ctypes" $ Just "ct", - Import "time" Nothing] - defines = [Escape pyValues, Escape pyFunctions, Escape pyPanic, Escape pyTuning] - operations :: GenericPython.Operations Imp.Sequential () - operations = GenericPython.defaultOperations - { GenericPython.opsCompiler = const $ return () - , GenericPython.opsCopy = copySequentialMemory - } + ImpGen.compileProg + >=> traverse + ( GenericPython.compileProg + module_name + GenericPython.emptyConstructor + imports + defines + operations + () + [] + [] + ) + where + imports = + [ Import "sys" Nothing, + Import "numpy" $ Just "np", + Import "ctypes" $ Just "ct", + Import "time" Nothing + ] + defines = [Escape pyValues, Escape pyFunctions, Escape pyPanic, Escape pyTuning] + operations :: GenericPython.Operations Imp.Sequential () + operations = + GenericPython.defaultOperations + { GenericPython.opsCompiler = const $ return (), + GenericPython.opsCopy = copySequentialMemory + } copySequentialMemory :: GenericPython.Copy Imp.Sequential () copySequentialMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes _bt = 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 @@ -1,27 +1,29 @@ {-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE Trustworthy #-} + -- | Simple C runtime representation. module Futhark.CodeGen.Backends.SimpleRep - ( tupleField - , funName - , defaultMemBlockType - , primTypeToCType - , signedPrimTypeToCType + ( tupleField, + funName, + defaultMemBlockType, + primTypeToCType, + signedPrimTypeToCType, -- * Primitive value operations - , cIntOps - , cFloat32Ops, cFloat32Funs - , cFloat64Ops, cFloat64Funs - , cFloatConvOps + cIntOps, + cFloat32Ops, + cFloat32Funs, + cFloat64Ops, + cFloat64Funs, + cFloatConvOps, ) - where - -import qualified Language.C.Syntax as C -import qualified Language.C.Quote.C as C +where import Futhark.CodeGen.ImpCode -import Futhark.Util.Pretty (prettyOneLine) import Futhark.Util (zEncodeString) +import Futhark.Util.Pretty (prettyOneLine) +import qualified Language.C.Quote.C as C +import qualified Language.C.Syntax as C -- | The C type corresponding to a signed integer type. intTypeToCType :: IntType -> C.Type @@ -64,7 +66,7 @@ -- | @funName f@ is the name of the C function corresponding to -- the Futhark function @f@. funName :: Name -> String -funName = ("futrts_"++) . zEncodeString . nameToString +funName = ("futrts_" ++) . zEncodeString . nameToString funName' :: String -> String funName' = funName . nameFromString @@ -74,92 +76,120 @@ defaultMemBlockType = [C.cty|char*|] cIntOps :: [C.Definition] -cIntOps = concatMap (`map` [minBound..maxBound]) ops - ++ cIntPrimFuns - where ops = [mkAdd, mkSub, mkMul, - mkUDiv, mkUDivUp, mkUMod, mkUDivSafe, mkUDivUpSafe, mkUModSafe, - mkSDiv, mkSDivUp, mkSMod, mkSDivSafe, mkSDivUpSafe, mkSModSafe, - mkSQuot, mkSRem, mkSQuotSafe, mkSRemSafe, - mkSMin, mkUMin, - mkSMax, mkUMax, - mkShl, mkLShr, mkAShr, - mkAnd, mkOr, mkXor, - mkUlt, mkUle, mkSlt, mkSle, - mkPow, - mkIToB, mkBToI - ] ++ - map mkSExt [minBound..maxBound] ++ - map mkZExt [minBound..maxBound] +cIntOps = + concatMap (`map` [minBound .. maxBound]) ops + ++ cIntPrimFuns + where + ops = + [ mkAdd, + mkSub, + mkMul, + mkUDiv, + mkUDivUp, + mkUMod, + mkUDivSafe, + mkUDivUpSafe, + mkUModSafe, + mkSDiv, + mkSDivUp, + mkSMod, + mkSDivSafe, + mkSDivUpSafe, + mkSModSafe, + mkSQuot, + mkSRem, + mkSQuotSafe, + mkSRemSafe, + mkSMin, + mkUMin, + mkSMax, + mkUMax, + mkShl, + mkLShr, + mkAShr, + mkAnd, + mkOr, + mkXor, + mkUlt, + mkUle, + mkSlt, + mkSle, + mkPow, + mkIToB, + mkBToI + ] + ++ map mkSExt [minBound .. maxBound] + ++ map mkZExt [minBound .. maxBound] - taggedI s Int8 = s ++ "8" - taggedI s Int16 = s ++ "16" - taggedI s Int32 = s ++ "32" - taggedI s Int64 = s ++ "64" + taggedI s Int8 = s ++ "8" + taggedI s Int16 = s ++ "16" + taggedI s Int32 = s ++ "32" + taggedI s Int64 = s ++ "64" - -- Use unsigned types for add/sub/mul so we can do - -- well-defined overflow. - mkAdd = simpleUintOp "add" [C.cexp|x + y|] - mkSub = simpleUintOp "sub" [C.cexp|x - y|] - mkMul = simpleUintOp "mul" [C.cexp|x * y|] - mkUDiv = simpleUintOp "udiv" [C.cexp|x / y|] - mkUDivUp = simpleUintOp "udiv_up" [C.cexp|(x+y-1) / y|] - mkUMod = simpleUintOp "umod" [C.cexp|x % y|] - mkUDivSafe = simpleUintOp "udiv_safe" [C.cexp|y == 0 ? 0 : x / y|] - mkUDivUpSafe = simpleUintOp "udiv_up_safe" [C.cexp|y == 0 ? 0 : (x+y-1) / y|] - mkUModSafe = simpleUintOp "umod_safe" [C.cexp|y == 0 ? 0 : x % y|] - mkUMax = simpleUintOp "umax" [C.cexp|x < y ? y : x|] - mkUMin = simpleUintOp "umin" [C.cexp|x < y ? x : y|] + -- Use unsigned types for add/sub/mul so we can do + -- well-defined overflow. + mkAdd = simpleUintOp "add" [C.cexp|x + y|] + mkSub = simpleUintOp "sub" [C.cexp|x - y|] + mkMul = simpleUintOp "mul" [C.cexp|x * y|] + mkUDiv = simpleUintOp "udiv" [C.cexp|x / y|] + mkUDivUp = simpleUintOp "udiv_up" [C.cexp|(x+y-1) / y|] + mkUMod = simpleUintOp "umod" [C.cexp|x % y|] + mkUDivSafe = simpleUintOp "udiv_safe" [C.cexp|y == 0 ? 0 : x / y|] + mkUDivUpSafe = simpleUintOp "udiv_up_safe" [C.cexp|y == 0 ? 0 : (x+y-1) / y|] + mkUModSafe = simpleUintOp "umod_safe" [C.cexp|y == 0 ? 0 : x % y|] + mkUMax = simpleUintOp "umax" [C.cexp|x < y ? y : x|] + mkUMin = simpleUintOp "umin" [C.cexp|x < y ? x : y|] - mkSDiv t = - let ct = intTypeToCType t - in [C.cedecl|static inline $ty:ct $id:(taggedI "sdiv" t)($ty:ct x, $ty:ct y) { + mkSDiv t = + let ct = intTypeToCType t + in [C.cedecl|static inline $ty:ct $id:(taggedI "sdiv" t)($ty:ct x, $ty:ct y) { $ty:ct q = x / y; $ty:ct r = x % y; return q - (((r != 0) && ((r < 0) != (y < 0))) ? 1 : 0); }|] - mkSDivUp t = - simpleIntOp "sdiv_up" [C.cexp|$id:(taggedI "sdiv" t)(x+y-1,y)|] t - mkSMod t = - let ct = intTypeToCType t - in [C.cedecl|static inline $ty:ct $id:(taggedI "smod" t)($ty:ct x, $ty:ct y) { + mkSDivUp t = + simpleIntOp "sdiv_up" [C.cexp|$id:(taggedI "sdiv" t)(x+y-1,y)|] t + mkSMod t = + let ct = intTypeToCType t + in [C.cedecl|static inline $ty:ct $id:(taggedI "smod" t)($ty:ct x, $ty:ct y) { $ty:ct r = x % y; return r + ((r == 0 || (x > 0 && y > 0) || (x < 0 && y < 0)) ? 0 : y); }|] - mkSDivSafe t = - simpleIntOp "sdiv_safe" [C.cexp|y == 0 ? 0 : $id:(taggedI "sdiv" t)(x,y)|] t - mkSDivUpSafe t = - simpleIntOp "sdiv_up_safe" [C.cexp|$id:(taggedI "sdiv_safe" t)(x+y-1,y)|] t - mkSModSafe t = - simpleIntOp "smod_safe" [C.cexp|y == 0 ? 0 : $id:(taggedI "smod" t)(x,y)|] t + mkSDivSafe t = + simpleIntOp "sdiv_safe" [C.cexp|y == 0 ? 0 : $id:(taggedI "sdiv" t)(x,y)|] t + mkSDivUpSafe t = + simpleIntOp "sdiv_up_safe" [C.cexp|$id:(taggedI "sdiv_safe" t)(x+y-1,y)|] t + mkSModSafe t = + simpleIntOp "smod_safe" [C.cexp|y == 0 ? 0 : $id:(taggedI "smod" t)(x,y)|] t - mkSQuot = simpleIntOp "squot" [C.cexp|x / y|] - mkSRem = simpleIntOp "srem" [C.cexp|x % y|] - mkSQuotSafe = simpleIntOp "squot_safe" [C.cexp|y == 0 ? 0 : x / y|] - mkSRemSafe = simpleIntOp "srem_safe" [C.cexp|y == 0 ? 0 : x % y|] - mkSMax = simpleIntOp "smax" [C.cexp|x < y ? y : x|] - mkSMin = simpleIntOp "smin" [C.cexp|x < y ? x : y|] - mkShl = simpleUintOp "shl" [C.cexp|x << y|] - mkLShr = simpleUintOp "lshr" [C.cexp|x >> y|] - mkAShr = simpleIntOp "ashr" [C.cexp|x >> y|] - mkAnd = simpleUintOp "and" [C.cexp|x & y|] - mkOr = simpleUintOp "or" [C.cexp|x | y|] - mkXor = simpleUintOp "xor" [C.cexp|x ^ y|] - mkUlt = uintCmpOp "ult" [C.cexp|x < y|] - mkUle = uintCmpOp "ule" [C.cexp|x <= y|] - mkSlt = intCmpOp "slt" [C.cexp|x < y|] - mkSle = intCmpOp "sle" [C.cexp|x <= y|] + mkSQuot = simpleIntOp "squot" [C.cexp|x / y|] + mkSRem = simpleIntOp "srem" [C.cexp|x % y|] + mkSQuotSafe = simpleIntOp "squot_safe" [C.cexp|y == 0 ? 0 : x / y|] + mkSRemSafe = simpleIntOp "srem_safe" [C.cexp|y == 0 ? 0 : x % y|] + mkSMax = simpleIntOp "smax" [C.cexp|x < y ? y : x|] + mkSMin = simpleIntOp "smin" [C.cexp|x < y ? x : y|] + mkShl = simpleUintOp "shl" [C.cexp|x << y|] + mkLShr = simpleUintOp "lshr" [C.cexp|x >> y|] + mkAShr = simpleIntOp "ashr" [C.cexp|x >> y|] + mkAnd = simpleUintOp "and" [C.cexp|x & y|] + mkOr = simpleUintOp "or" [C.cexp|x | y|] + mkXor = simpleUintOp "xor" [C.cexp|x ^ y|] + mkUlt = uintCmpOp "ult" [C.cexp|x < y|] + mkUle = uintCmpOp "ule" [C.cexp|x <= y|] + mkSlt = intCmpOp "slt" [C.cexp|x < y|] + mkSle = intCmpOp "sle" [C.cexp|x <= y|] - -- We define some operations as macros rather than functions, - -- because this allows us to use them as constant expressions - -- in things like array sizes and static initialisers. - macro name rhs = - [C.cedecl|$esc:("#define " ++ name ++ "(x) (" ++ prettyOneLine rhs ++ ")")|] + -- We define some operations as macros rather than functions, + -- because this allows us to use them as constant expressions + -- in things like array sizes and static initialisers. + macro name rhs = + [C.cedecl|$esc:("#define " ++ name ++ "(x) (" ++ prettyOneLine rhs ++ ")")|] - mkPow t = - let ct = intTypeToCType t - in [C.cedecl|static inline $ty:ct $id:(taggedI "pow" t)($ty:ct x, $ty:ct y) { + mkPow t = + let ct = intTypeToCType t + in [C.cedecl|static inline $ty:ct $id:(taggedI "pow" t)($ty:ct x, $ty:ct y) { $ty:ct res = 1, rem = y; while (rem != 0) { if (rem & 1) { @@ -171,42 +201,50 @@ return res; }|] - mkSExt from_t to_t = macro name [C.cexp|($ty:to_ct)(($ty:from_ct)x)|] - where name = "sext_"++pretty from_t++"_"++pretty to_t - from_ct = intTypeToCType from_t - to_ct = intTypeToCType to_t + mkSExt from_t to_t = macro name [C.cexp|($ty:to_ct)(($ty:from_ct)x)|] + where + name = "sext_" ++ pretty from_t ++ "_" ++ pretty to_t + from_ct = intTypeToCType from_t + to_ct = intTypeToCType to_t - mkZExt from_t to_t = macro name [C.cexp|($ty:to_ct)(($ty:from_ct)x)|] - where name = "zext_"++pretty from_t++"_"++pretty to_t - from_ct = uintTypeToCType from_t - to_ct = intTypeToCType to_t + mkZExt from_t to_t = macro name [C.cexp|($ty:to_ct)(($ty:from_ct)x)|] + where + name = "zext_" ++ pretty from_t ++ "_" ++ pretty to_t + from_ct = uintTypeToCType from_t + to_ct = intTypeToCType to_t - mkBToI to_t = - [C.cedecl|static inline $ty:to_ct + mkBToI to_t = + [C.cedecl|static inline $ty:to_ct $id:name($ty:from_ct x) { return x; } |] - where name = "btoi_bool_"++pretty to_t - from_ct = primTypeToCType Bool - to_ct = intTypeToCType to_t + where + name = "btoi_bool_" ++ pretty to_t + from_ct = primTypeToCType Bool + to_ct = intTypeToCType to_t - mkIToB from_t = - [C.cedecl|static inline $ty:to_ct + mkIToB from_t = + [C.cedecl|static inline $ty:to_ct $id:name($ty:from_ct x) { return x; } |] - where name = "itob_"++pretty from_t++"_bool" - to_ct = primTypeToCType Bool - from_ct = intTypeToCType from_t + where + name = "itob_" ++ pretty from_t ++ "_bool" + to_ct = primTypeToCType Bool + from_ct = intTypeToCType from_t - simpleUintOp s e t = - [C.cedecl|static inline $ty:ct $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] - where ct = uintTypeToCType t - simpleIntOp s e t = - [C.cedecl|static inline $ty:ct $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] - where ct = intTypeToCType t - intCmpOp s e t = - [C.cedecl|static inline typename bool $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] - where ct = intTypeToCType t - uintCmpOp s e t = - [C.cedecl|static inline typename bool $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] - where ct = uintTypeToCType t + simpleUintOp s e t = + [C.cedecl|static inline $ty:ct $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where + ct = uintTypeToCType t + simpleIntOp s e t = + [C.cedecl|static inline $ty:ct $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where + ct = intTypeToCType t + intCmpOp s e t = + [C.cedecl|static inline typename bool $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where + ct = intTypeToCType t + uintCmpOp s e t = + [C.cedecl|static inline typename bool $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where + ct = uintTypeToCType t cIntPrimFuns :: [C.Definition] cIntPrimFuns = @@ -479,56 +517,64 @@ cFloat64Ops :: [C.Definition] cFloatConvOps :: [C.Definition] (cFloat32Ops, cFloat64Ops, cFloatConvOps) = - ( map ($Float32) mkOps - , map ($Float64) mkOps - , [ mkFPConvFF "fpconv" from to | - from <- [minBound..maxBound], - to <- [minBound..maxBound] ]) - where taggedF s Float32 = s ++ "32" - taggedF s Float64 = s ++ "64" - convOp s from to = s ++ "_" ++ pretty from ++ "_" ++ pretty to + ( map ($ Float32) mkOps, + map ($ Float64) mkOps, + [ mkFPConvFF "fpconv" from to + | from <- [minBound .. maxBound], + to <- [minBound .. maxBound] + ] + ) + where + taggedF s Float32 = s ++ "32" + taggedF s Float64 = s ++ "64" + convOp s from to = s ++ "_" ++ pretty from ++ "_" ++ pretty to - mkOps = [mkFDiv, mkFAdd, mkFSub, mkFMul, mkFMin, mkFMax, mkPow, mkCmpLt, mkCmpLe] ++ - map (mkFPConvIF "sitofp") [minBound..maxBound] ++ - map (mkFPConvUF "uitofp") [minBound..maxBound] ++ - map (flip $ mkFPConvFI "fptosi") [minBound..maxBound] ++ - map (flip $ mkFPConvFU "fptoui") [minBound..maxBound] + mkOps = + [mkFDiv, mkFAdd, mkFSub, mkFMul, mkFMin, mkFMax, mkPow, mkCmpLt, mkCmpLe] + ++ map (mkFPConvIF "sitofp") [minBound .. maxBound] + ++ map (mkFPConvUF "uitofp") [minBound .. maxBound] + ++ map (flip $ mkFPConvFI "fptosi") [minBound .. maxBound] + ++ map (flip $ mkFPConvFU "fptoui") [minBound .. maxBound] - mkFDiv = simpleFloatOp "fdiv" [C.cexp|x / y|] - mkFAdd = simpleFloatOp "fadd" [C.cexp|x + y|] - mkFSub = simpleFloatOp "fsub" [C.cexp|x - y|] - mkFMul = simpleFloatOp "fmul" [C.cexp|x * y|] - mkFMin = simpleFloatOp "fmin" [C.cexp|fmin(x, y)|] - mkFMax = simpleFloatOp "fmax" [C.cexp|fmax(x, y)|] - mkCmpLt = floatCmpOp "cmplt" [C.cexp|x < y|] - mkCmpLe = floatCmpOp "cmple" [C.cexp|x <= y|] + mkFDiv = simpleFloatOp "fdiv" [C.cexp|x / y|] + mkFAdd = simpleFloatOp "fadd" [C.cexp|x + y|] + mkFSub = simpleFloatOp "fsub" [C.cexp|x - y|] + mkFMul = simpleFloatOp "fmul" [C.cexp|x * y|] + mkFMin = simpleFloatOp "fmin" [C.cexp|fmin(x, y)|] + mkFMax = simpleFloatOp "fmax" [C.cexp|fmax(x, y)|] + mkCmpLt = floatCmpOp "cmplt" [C.cexp|x < y|] + mkCmpLe = floatCmpOp "cmple" [C.cexp|x <= y|] - mkPow Float32 = - [C.cedecl|static inline float fpow32(float x, float y) { return pow(x, y); }|] - mkPow Float64 = - [C.cedecl|static inline double fpow64(double x, double y) { return pow(x, y); }|] + mkPow Float32 = + [C.cedecl|static inline float fpow32(float x, float y) { return pow(x, y); }|] + mkPow Float64 = + [C.cedecl|static inline double fpow64(double x, double y) { return pow(x, y); }|] - mkFPConv from_f to_f s from_t to_t = - [C.cedecl|static inline $ty:to_ct + mkFPConv from_f to_f s from_t to_t = + [C.cedecl|static inline $ty:to_ct $id:(convOp s from_t to_t)($ty:from_ct x) { return ($ty:to_ct)x;} |] - where from_ct = from_f from_t - to_ct = to_f to_t + where + from_ct = from_f from_t + to_ct = to_f to_t - mkFPConvFF = mkFPConv floatTypeToCType floatTypeToCType - mkFPConvFI = mkFPConv floatTypeToCType intTypeToCType - mkFPConvIF = mkFPConv intTypeToCType floatTypeToCType - mkFPConvFU = mkFPConv floatTypeToCType uintTypeToCType - mkFPConvUF = mkFPConv uintTypeToCType floatTypeToCType + mkFPConvFF = mkFPConv floatTypeToCType floatTypeToCType + mkFPConvFI = mkFPConv floatTypeToCType intTypeToCType + mkFPConvIF = mkFPConv intTypeToCType floatTypeToCType + mkFPConvFU = mkFPConv floatTypeToCType uintTypeToCType + mkFPConvUF = mkFPConv uintTypeToCType floatTypeToCType - simpleFloatOp s e t = - [C.cedecl|static inline $ty:ct $id:(taggedF s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] - where ct = floatTypeToCType t - floatCmpOp s e t = - [C.cedecl|static inline typename bool $id:(taggedF s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] - where ct = floatTypeToCType t + simpleFloatOp s e t = + [C.cedecl|static inline $ty:ct $id:(taggedF s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where + ct = floatTypeToCType t + floatCmpOp s e t = + [C.cedecl|static inline typename bool $id:(taggedF s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where + ct = floatTypeToCType t cFloat32Funs :: [C.Definition] -cFloat32Funs = [C.cunit| +cFloat32Funs = + [C.cunit| static inline float $id:(funName' "log32")(float x) { return log(x); } @@ -683,7 +729,8 @@ |] cFloat64Funs :: [C.Definition] -cFloat64Funs = [C.cunit| +cFloat64Funs = + [C.cunit| static inline double $id:(funName' "log64")(double x) { return log(x); } 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 @@ -1,7 +1,8 @@ {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TupleSections #-} {-# LANGUAGE Safe #-} {-# LANGUAGE Strict #-} +{-# LANGUAGE TupleSections #-} + -- | Imperative intermediate language used as a stepping stone in code generation. -- -- This is a generic representation parametrised on an extensible @@ -10,69 +11,74 @@ -- Originally inspired by the paper "Defunctionalizing Push Arrays" -- (FHPC '14). module Futhark.CodeGen.ImpCode - ( Definitions (..) - , Functions (..) - , Function - , FunctionT (..) - , Constants (..) - , ValueDesc (..) - , Signedness (..) - , ExternalValue (..) - , Param (..) - , paramName - , SubExp(..) - , MemSize - , DimSize - , Space (..) - , SpaceId - , Code (..) - , PrimValue (..) - , ExpLeaf (..) - , Exp - , Volatility (..) - , Arg (..) - , var - , vi32 - , index - , ErrorMsg(..) - , ErrorMsgPart(..) - , errorMsgArgTypes - , ArrayContents(..) - - , lexicalMemoryUsage - , calledFuncs + ( Definitions (..), + Functions (..), + Function, + FunctionT (..), + Constants (..), + ValueDesc (..), + Signedness (..), + ExternalValue (..), + Param (..), + paramName, + SubExp (..), + MemSize, + DimSize, + Space (..), + SpaceId, + Code (..), + PrimValue (..), + ExpLeaf (..), + Exp, + TExp, + Volatility (..), + Arg (..), + var, + vi32, + vi64, + index, + ErrorMsg (..), + ErrorMsgPart (..), + errorMsgArgTypes, + ArrayContents (..), + lexicalMemoryUsage, + calledFuncs, -- * Typed enumerations - , Bytes - , Elements - , elements - , bytes - , withElemType + Bytes, + Elements, + elements, + bytes, + withElemType, -- * Re-exports from other modules. - , module Language.Futhark.Core - , module Futhark.IR.Primitive - , module Futhark.Analysis.PrimExp - , module Futhark.IR.Kernels.Sizes - , module Futhark.IR.Prop.Names + module Language.Futhark.Core, + module Futhark.IR.Primitive, + module Futhark.Analysis.PrimExp, + module Futhark.IR.Kernels.Sizes, + module Futhark.IR.Prop.Names, ) - where +where import Data.List (intersperse) +import qualified Data.Map as M import qualified Data.Set as S import Data.Traversable -import qualified Data.Map as M - -import Language.Futhark.Core +import Futhark.Analysis.PrimExp +import Futhark.IR.Kernels.Sizes (Count (..)) +import Futhark.IR.Pretty () import Futhark.IR.Primitive -import Futhark.IR.Syntax - (SubExp(..), Space(..), SpaceId, - ErrorMsg(..), ErrorMsgPart(..), errorMsgArgTypes) import Futhark.IR.Prop.Names -import Futhark.IR.Pretty () -import Futhark.Analysis.PrimExp +import Futhark.IR.Syntax + ( ErrorMsg (..), + ErrorMsgPart (..), + Space (..), + SpaceId, + SubExp (..), + errorMsgArgTypes, + ) import Futhark.Util.Pretty hiding (space) -import Futhark.IR.Kernels.Sizes (Count(..)) +import Language.Futhark.Core -- | The size of a memory block. type MemSize = SubExp @@ -81,9 +87,10 @@ type DimSize = SubExp -- | An ImpCode function parameter. -data Param = MemParam VName Space - | ScalarParam VName PrimType - deriving (Show) +data Param + = MemParam VName Space + | ScalarParam VName PrimType + deriving (Show) -- | The name of a parameter. paramName :: Param -> VName @@ -91,9 +98,10 @@ paramName (ScalarParam name _) = name -- | A collection of imperative functions and constants. -data Definitions a = Definitions { defConsts :: Constants a - , defFuns :: Functions a - } +data Definitions a = Definitions + { defConsts :: Constants a, + defFuns :: Functions a + } -- | A collection of imperative functions. newtype Functions a = Functions [(Name, Function a)] @@ -106,51 +114,55 @@ -- | A collection of imperative constants. data Constants a = Constants - { constsDecl :: [Param] - -- ^ The constants that are made available to the functions. - , constsInit :: Code a - -- ^ Setting the value of the constants. Note that this must not + { -- | The constants that are made available to the functions. + constsDecl :: [Param], + -- | Setting the value of the constants. Note that this must not -- contain declarations of the names defined in 'constsDecl'. + constsInit :: Code a } -- | Since the core language does not care for signedness, but the -- source language does, entry point input/output information has -- metadata for integer types (and arrays containing these) that -- indicate whether they are really unsigned integers. -data Signedness = TypeUnsigned - | TypeDirect - deriving (Eq, Show) +data Signedness + = TypeUnsigned + | TypeDirect + deriving (Eq, Show) -- | A description of an externally meaningful value. -data ValueDesc = ArrayValue VName Space PrimType Signedness [DimSize] - -- ^ An array with memory block, memory block size, - -- memory space, element type, signedness of element - -- type (if applicable), and shape. - | ScalarValue PrimType Signedness VName - -- ^ A scalar value with signedness if applicable. - deriving (Eq, Show) +data ValueDesc + = -- | An array with memory block, memory block size, + -- memory space, element type, signedness of element + -- type (if applicable), and shape. + ArrayValue VName Space PrimType Signedness [DimSize] + | -- | A scalar value with signedness if applicable. + ScalarValue PrimType Signedness VName + deriving (Eq, Show) -- | ^ An externally visible value. This can be an opaque value -- (covering several physical internal values), or a single value that -- can be used externally. -data ExternalValue = OpaqueValue String [ValueDesc] - -- ^ The string is a human-readable description - -- with no other semantics. - | TransparentValue ValueDesc - deriving (Show) +data ExternalValue + = -- | The string is a human-readable description + -- with no other semantics. + OpaqueValue String [ValueDesc] + | TransparentValue ValueDesc + deriving (Show) -- | A imperative function, containing the body as well as its -- low-level inputs and outputs, as well as its high-level arguments -- and results. The latter are only used if the function is an entry -- point. -data FunctionT a = Function { functionEntry :: Bool - , functionOutput :: [Param] - , functionInput :: [Param] - , functionBody :: Code a - , functionResult :: [ExternalValue] - , functionArgs :: [ExternalValue] - } - deriving (Show) +data FunctionT a = Function + { functionEntry :: Bool, + functionOutput :: [Param], + functionInput :: [Param], + functionBody :: Code a, + functionResult :: [ExternalValue], + functionArgs :: [ExternalValue] + } + deriving (Show) -- | Type alias for namespace control. type Function = FunctionT @@ -158,100 +170,109 @@ -- | The contents of a statically declared constant array. Such -- arrays are always unidimensional, and reshaped if necessary in the -- code that uses them. -data ArrayContents = ArrayValues [PrimValue] - -- ^ Precisely these values. - | ArrayZeros Int - -- ^ This many zeroes. - deriving (Show) +data ArrayContents + = -- | Precisely these values. + ArrayValues [PrimValue] + | -- | This many zeroes. + ArrayZeros Int + deriving (Show) -- | A block of imperative code. Parameterised by an 'Op', which -- allows extensibility. Concrete uses of this type will instantiate -- the type parameter with e.g. a construct for launching GPU kernels. -data Code a = Skip - -- ^ No-op. Crucial for the 'Monoid' instance. - | Code a :>>: Code a - -- ^ Statement composition. Crucial for the 'Semigroup' instance. - | For VName IntType Exp (Code a) - -- ^ A for-loop iterating the given number of times. The - -- loop parameter starts counting from zero and will have - -- the given type. The bound is evaluated just once, - -- before the loop is entered. - | While Exp (Code a) - -- ^ While loop. The conditional is (of course) - -- re-evaluated before every iteration of the loop. - | DeclareMem VName Space - -- ^ Declare a memory block variable that will point to - -- memory in the given memory space. Note that this is - -- distinct from allocation. The memory block must be the - -- target of either an 'Allocate' or a 'SetMem' before it - -- can be used for reading or writing. - | DeclareScalar VName Volatility PrimType - -- ^ Declare a scalar variable with an initially undefined value. - | DeclareArray VName Space PrimType ArrayContents - -- ^ Create an array containing the given values. The - -- lifetime of the array will be the entire application. - -- This is mostly used for constant arrays, but also for - -- some bookkeeping data, like the synchronisation - -- counts used to implement reduction. - | Allocate VName (Count Bytes Exp) Space - -- ^ Memory space must match the corresponding - -- 'DeclareMem'. - | Free VName Space - -- ^ Indicate that some memory block will never again be - -- referenced via the indicated variable. However, it - -- may still be accessed through aliases. It is only - -- safe to actually deallocate the memory block if this - -- is the last reference. There is no guarantee that - -- all memory blocks will be freed with this statement. - -- Backends are free to ignore it entirely. - | Copy VName (Count Bytes Exp) Space VName (Count Bytes Exp) Space (Count Bytes Exp) - -- ^ Destination, offset in destination, destination - -- space, source, offset in source, offset space, number - -- of bytes. - | Write VName (Count Elements Exp) PrimType Space Volatility Exp - -- ^ @Write mem i t space vol v@ writes the value @v@ to - -- @mem@ offset by @i@ elements of type @t@. The - -- 'Space' argument is the memory space of @mem@ - -- (technically redundant, but convenient). Note that - -- /reading/ is done with an 'Exp' ('Index'). - | SetScalar VName Exp - -- ^ Set a scalar variable. - | SetMem VName VName Space - -- ^ Must be in same space. - | Call [VName] Name [Arg] - -- ^ Function call. The results are written to the - -- provided 'VName' variables. - | If Exp (Code a) (Code a) - -- ^ Conditional execution. - | Assert Exp (ErrorMsg Exp) (SrcLoc, [SrcLoc]) - -- ^ Assert that something must be true. Should it turn - -- out not to be true, then report a failure along with - -- the given error message. - | Comment String (Code a) - -- ^ Has the same semantics as the contained code, but - -- the comment should show up in generated code for ease - -- of inspection. - | DebugPrint String (Maybe Exp) - -- ^ Print the given value to the screen, somehow - -- annotated with the given string as a description. If - -- no type/value pair, just print the string. This has - -- no semantic meaning, but is used entirely for - -- debugging. Code generators are free to ignore this - -- statement. - | Op a - -- ^ Perform an extensible operation. - deriving (Show) +data Code a + = -- | No-op. Crucial for the 'Monoid' instance. + Skip + | -- | Statement composition. Crucial for the 'Semigroup' instance. + Code a :>>: Code a + | -- | A for-loop iterating the given number of times. + -- The loop parameter starts counting from zero and will + -- have the same (integer) type as the bound. The bound + -- is evaluated just once, before the loop is entered. + For VName Exp (Code a) + | -- | While loop. The conditional is (of course) + -- re-evaluated before every iteration of the loop. + While (TExp Bool) (Code a) + | -- | Declare a memory block variable that will point to + -- memory in the given memory space. Note that this is + -- distinct from allocation. The memory block must be the + -- target of either an 'Allocate' or a 'SetMem' before it + -- can be used for reading or writing. + DeclareMem VName Space + | -- | Declare a scalar variable with an initially undefined value. + DeclareScalar VName Volatility PrimType + | -- | Create an array containing the given values. The + -- lifetime of the array will be the entire application. + -- This is mostly used for constant arrays, but also for + -- some bookkeeping data, like the synchronisation + -- counts used to implement reduction. + DeclareArray VName Space PrimType ArrayContents + | -- | Memory space must match the corresponding + -- 'DeclareMem'. + Allocate VName (Count Bytes (TExp Int64)) Space + | -- | Indicate that some memory block will never again be + -- referenced via the indicated variable. However, it + -- may still be accessed through aliases. It is only + -- safe to actually deallocate the memory block if this + -- is the last reference. There is no guarantee that + -- all memory blocks will be freed with this statement. + -- Backends are free to ignore it entirely. + Free VName Space + | -- | Destination, offset in destination, destination + -- space, source, offset in source, offset space, number + -- of bytes. + Copy + VName + (Count Bytes (TExp Int64)) + Space + VName + (Count Bytes (TExp Int64)) + Space + (Count Bytes (TExp Int64)) + | -- | @Write mem i t space vol v@ writes the value @v@ to + -- @mem@ offset by @i@ elements of type @t@. The + -- 'Space' argument is the memory space of @mem@ + -- (technically redundant, but convenient). Note that + -- /reading/ is done with an 'Exp' ('Index'). + Write VName (Count Elements (TExp Int64)) PrimType Space Volatility Exp + | -- | Set a scalar variable. + SetScalar VName Exp + | -- | Must be in same space. + SetMem VName VName Space + | -- | Function call. The results are written to the + -- provided 'VName' variables. + Call [VName] Name [Arg] + | -- | Conditional execution. + If (TExp Bool) (Code a) (Code a) + | -- | Assert that something must be true. Should it turn + -- out not to be true, then report a failure along with + -- the given error message. + Assert Exp (ErrorMsg Exp) (SrcLoc, [SrcLoc]) + | -- | Has the same semantics as the contained code, but + -- the comment should show up in generated code for ease + -- of inspection. + Comment String (Code a) + | -- | Print the given value to the screen, somehow + -- annotated with the given string as a description. If + -- no type/value pair, just print the string. This has + -- no semantic meaning, but is used entirely for + -- debugging. Code generators are free to ignore this + -- statement. + DebugPrint String (Maybe Exp) + | -- | Perform an extensible operation. + Op a + deriving (Show) -- | The volatility of a memory access or variable. Feel free to -- ignore this for backends where it makes no sense (anything but C -- and similar low-level things) data Volatility = Volatile | Nonvolatile - deriving (Eq, Ord, Show) + deriving (Eq, Ord, Show) instance Semigroup (Code a) where - Skip <> y = y - x <> Skip = x - x <> y = x :>>: y + Skip <> y = y + x <> Skip = x + x <> y = x :>>: y instance Monoid (Code a) where mempty = Skip @@ -268,57 +289,64 @@ lexicalMemoryUsage :: Function a -> M.Map VName Space lexicalMemoryUsage func = M.filterWithKey (const . not . (`nameIn` nonlexical)) $ - declared $ functionBody func - where nonlexical = - set (functionBody func) <> - namesFromList (map paramName (functionOutput func)) + declared $ functionBody func + where + nonlexical = + set (functionBody func) + <> namesFromList (map paramName (functionOutput func)) - go f (x :>>: y) = f x <> f y - go f (If _ x y) = f x <> f y - go f (For _ _ _ x) = f x - go f (While _ x) = f x - go f (Comment _ x) = f x - go _ _ = mempty + go f (x :>>: y) = f x <> f y + go f (If _ x y) = f x <> f y + go f (For _ _ x) = f x + go f (While _ x) = f x + go f (Comment _ x) = f x + go _ _ = mempty - declared (DeclareMem mem space) = M.singleton mem space - declared x = go declared x + declared (DeclareMem mem space) = M.singleton mem space + declared x = go declared x - set (SetMem x y _) = namesFromList [x,y] - set (Call _ _ args) = foldMap onArg args - where onArg ExpArg{} = mempty - onArg (MemArg x) = oneName x - set x = go set x + set (SetMem x y _) = namesFromList [x, y] + set (Call _ _ args) = foldMap onArg args + where + onArg ExpArg {} = mempty + onArg (MemArg x) = oneName x + set x = go set x -- | The set of functions that are called by this code. Assumes there -- are no function calls in 'Op's. calledFuncs :: Code a -> S.Set Name calledFuncs (x :>>: y) = calledFuncs x <> calledFuncs y calledFuncs (If _ x y) = calledFuncs x <> calledFuncs y -calledFuncs (For _ _ _ x) = calledFuncs x +calledFuncs (For _ _ x) = calledFuncs x calledFuncs (While _ x) = calledFuncs x calledFuncs (Comment _ x) = calledFuncs x calledFuncs (Call _ f _) = S.singleton f calledFuncs _ = mempty -- | The leaves of an 'Exp'. -data ExpLeaf = ScalarVar VName - -- ^ A scalar variable. The type is stored in the - -- 'LeafExp' constructor itself. - | SizeOf PrimType - -- ^ The size of a primitive type. - | Index VName (Count Elements Exp) PrimType Space Volatility - -- ^ Reading a value from memory. The arguments have - -- the same meaning as with 'Write'. - deriving (Eq, Show) +data ExpLeaf + = -- | A scalar variable. The type is stored in the + -- 'LeafExp' constructor itself. + ScalarVar VName + | -- | The size of a primitive type. + SizeOf PrimType + | -- | Reading a value from memory. The arguments have + -- the same meaning as with 'Write'. + Index VName (Count Elements (TExp Int64)) PrimType Space Volatility + deriving (Eq, Show) -- | A side-effect free expression whose execution will produce a -- single primitive value. type Exp = PrimExp ExpLeaf +-- | Like 'Exp', but with a required/known type. +type TExp t = TPrimExp t ExpLeaf + -- | A function call argument. -data Arg = ExpArg Exp - | MemArg VName - deriving (Show) +data Arg + = ExpArg Exp + | MemArg VName + deriving (Show) -- | Phantom type for a count of elements. data Elements @@ -327,29 +355,33 @@ data Bytes -- | This expression counts elements. -elements :: Exp -> Count Elements Exp +elements :: a -> Count Elements a elements = Count -- | This expression counts bytes. -bytes :: Exp -> Count Bytes Exp +bytes :: a -> Count Bytes a bytes = Count -- | Convert a count of elements into a count of bytes, given the -- per-element size. -withElemType :: Count Elements Exp -> PrimType -> Count Bytes Exp +withElemType :: Count Elements (TExp Int64) -> PrimType -> Count Bytes (TExp Int64) withElemType (Count e) t = - bytes $ sExt Int64 e * LeafExp (SizeOf t) (IntType Int64) + bytes $ sExt64 e * isInt64 (LeafExp (SizeOf t) (IntType Int64)) -- | Turn a 'VName' into a 'Imp.ScalarVar'. var :: VName -> PrimType -> Exp var = LeafExp . ScalarVar -- | Turn a 'VName' into a v'Int32' 'Imp.ScalarVar'. -vi32 :: VName -> Exp -vi32 = flip var $ IntType Int32 +vi32 :: VName -> TExp Int32 +vi32 = TPrimExp . flip var (IntType Int32) +-- | Turn a 'VName' into a v'Int64' 'Imp.ScalarVar'. +vi64 :: VName -> TExp Int64 +vi64 = TPrimExp . flip var (IntType Int64) + -- | Concise wrapper for using 'Index'. -index :: VName -> Count Elements Exp -> PrimType -> Space -> Volatility -> Exp +index :: VName -> Count Elements (TExp Int64) -> PrimType -> Space -> Volatility -> Exp index arr i t s vol = LeafExp (Index arr i t s vol) t -- Prettyprinting definitions. @@ -360,25 +392,31 @@ instance Pretty op => Pretty (Functions op) where ppr (Functions funs) = stack $ intersperse mempty $ map ppFun funs - where ppFun (name, fun) = - text "Function " <> ppr name <> colon indent 2 (ppr fun) + where + ppFun (name, fun) = + text "Function " <> ppr name <> colon indent 2 (ppr fun) instance Pretty op => Pretty (Constants op) where ppr (Constants decls code) = - text "Constants:" indent 2 (stack $ map ppr decls) - mempty - text "Initialisation:" - indent 2 (ppr code) + text "Constants:" indent 2 (stack $ map ppr decls) + mempty + text "Initialisation:" + indent 2 (ppr code) instance Pretty op => Pretty (FunctionT op) where ppr (Function _ outs ins body results args) = - text "Inputs:" block ins - text "Outputs:" block outs - text "Arguments:" block args - text "Result:" block results - text "Body:" indent 2 (ppr body) - where block :: Pretty a => [a] -> Doc - block = indent 2 . stack . map ppr + text "Inputs:" block ins + text "Outputs:" + block outs + text "Arguments:" + block args + text "Result:" + block results + text "Body:" + indent 2 (ppr body) + where + block :: Pretty a => [a] -> Doc + block = indent 2 . stack . map ppr instance Pretty Param where ppr (ScalarParam name ptype) = ppr ptype <+> ppr name @@ -387,20 +425,23 @@ instance Pretty ValueDesc where ppr (ScalarValue t ept name) = ppr t <+> ppr name <> ept' - where ept' = case ept of TypeUnsigned -> text " (unsigned)" - TypeDirect -> mempty + where + ept' = case ept of + TypeUnsigned -> text " (unsigned)" + TypeDirect -> mempty ppr (ArrayValue mem space et ept shape) = foldr f (ppr et) shape <+> text "at" <+> ppr mem <> ppr space <+> ept' - where f e s = brackets $ s <> comma <> ppr e - ept' = case ept of TypeUnsigned -> text " (unsigned)" - TypeDirect -> mempty - + where + f e s = brackets $ s <> comma <> ppr e + ept' = case ept of + TypeUnsigned -> text " (unsigned)" + TypeDirect -> mempty instance Pretty ExternalValue where ppr (TransparentValue v) = ppr v ppr (OpaqueValue desc vs) = - text "opaque" <+> text desc <+> - nestedBlock "{" "}" (stack $ map ppr vs) + text "opaque" <+> text desc + <+> nestedBlock "{" "}" (stack $ map ppr vs) instance Pretty ArrayContents where ppr (ArrayValues vs) = braces (commasep $ map ppr vs) @@ -408,34 +449,40 @@ instance Pretty op => Pretty (Code op) where ppr (Op op) = ppr op - ppr Skip = text "skip" + ppr Skip = text "skip" ppr (c1 :>>: c2) = ppr c1 ppr c2 - ppr (For i it limit body) = - text "for" <+> ppr i <> text ":" <> ppr it <+> langle <+> ppr limit <+> text "{" - indent 2 (ppr body) - text "}" + ppr (For i limit body) = + text "for" <+> ppr i <+> langle <+> ppr limit <+> text "{" + indent 2 (ppr body) + text "}" ppr (While cond body) = - text "while" <+> ppr cond <+> text "{" - indent 2 (ppr body) - text "}" + text "while" <+> ppr cond <+> text "{" + indent 2 (ppr body) + text "}" ppr (DeclareMem name space) = text "var" <+> ppr name <> text ": mem" <> ppr space ppr (DeclareScalar name vol t) = text "var" <+> ppr name <> text ":" <+> vol' <> ppr t - where vol' = case vol of Volatile -> text "volatile " - Nonvolatile -> mempty + where + vol' = case vol of + Volatile -> text "volatile " + Nonvolatile -> mempty ppr (DeclareArray name space t vs) = - text "array" <+> ppr name <> text "@" <> ppr space <+> text ":" <+> ppr t <+> - equals <+> ppr vs + text "array" <+> ppr name <> text "@" <> ppr space <+> text ":" <+> ppr t + <+> equals + <+> ppr vs ppr (Allocate name e space) = ppr name <+> text "<-" <+> text "malloc" <> parens (ppr e) <> ppr space ppr (Free name space) = text "free" <> parens (ppr name) <> ppr space ppr (Write name i bt space vol val) = - ppr name <> langle <> vol' <> ppr bt <> ppr space <> rangle <> brackets (ppr i) <+> - text "<-" <+> ppr val - where vol' = case vol of Volatile -> text "volatile " - Nonvolatile -> mempty + ppr name <> langle <> vol' <> ppr bt <> ppr space <> rangle <> brackets (ppr i) + <+> text "<-" + <+> ppr val + where + vol' = case vol of + Volatile -> text "volatile " + Nonvolatile -> mempty ppr (SetScalar name val) = ppr name <+> text "<-" <+> ppr val ppr (SetMem dest from space) = @@ -443,21 +490,24 @@ ppr (Assert e msg _) = text "assert" <> parens (commasep [ppr msg, ppr e]) ppr (Copy dest destoffset destspace src srcoffset srcspace size) = - text "memcpy" <> - parens (ppMemLoc dest destoffset <> ppr destspace <> comma - ppMemLoc src srcoffset <> ppr srcspace <> comma - ppr size) - where ppMemLoc base offset = - ppr base <+> text "+" <+> ppr offset + text "memcpy" + <> parens + ( ppMemLoc dest destoffset <> ppr destspace <> comma + ppMemLoc src srcoffset <> ppr srcspace <> comma + ppr size + ) + where + ppMemLoc base offset = + ppr base <+> text "+" <+> ppr offset ppr (If cond tbranch fbranch) = - text "if" <+> ppr cond <+> text "then {" - indent 2 (ppr tbranch) - text "} else {" - indent 2 (ppr fbranch) - text "}" + text "if" <+> ppr cond <+> text "then {" + indent 2 (ppr tbranch) + text "} else {" + indent 2 (ppr fbranch) + text "}" ppr (Call dests fname args) = - commasep (map ppr dests) <+> text "<-" <+> - ppr fname <> parens (commasep $ map ppr args) + commasep (map ppr dests) <+> text "<-" + <+> ppr fname <> parens (commasep $ map ppr args) ppr (Comment s code) = text "--" <+> text s ppr code ppr (DebugPrint desc (Just e)) = @@ -474,9 +524,10 @@ ppr v ppr (Index v is bt space vol) = ppr v <> langle <> vol' <> ppr bt <> ppr space <> rangle <> brackets (ppr is) - where vol' = case vol of Volatile -> text "volatile " - Nonvolatile -> mempty - + where + vol' = case vol of + Volatile -> text "volatile " + Nonvolatile -> mempty ppr (SizeOf t) = text "sizeof" <> parens (ppr t) @@ -489,7 +540,8 @@ instance Traversable Functions where traverse f (Functions funs) = Functions <$> traverse f' funs - where f' (name, fun) = (name,) <$> traverse f fun + where + f' (name, fun) = (name,) <$> traverse f fun instance Functor FunctionT where fmap = fmapDefault @@ -510,8 +562,8 @@ instance Traversable Code where traverse f (x :>>: y) = (:>>:) <$> traverse f x <*> traverse f y - traverse f (For i it bound code) = - For i it bound <$> traverse f code + traverse f (For i bound code) = + For i bound <$> traverse f code traverse f (While cond code) = While cond <$> traverse f code traverse f (If cond x y) = @@ -553,7 +605,7 @@ declaredIn (DeclareArray name _ _ _) = oneName name declaredIn (If _ t f) = declaredIn t <> declaredIn f declaredIn (x :>>: y) = declaredIn x <> declaredIn y -declaredIn (For i _ _ body) = oneName i <> declaredIn body +declaredIn (For i _ body) = oneName i <> declaredIn body declaredIn (While _ body) = declaredIn body declaredIn (Comment _ body) = declaredIn body declaredIn _ = mempty @@ -567,15 +619,15 @@ fvBind (declaredIn x) $ freeIn' x <> freeIn' y freeIn' Skip = mempty - freeIn' (For i _ bound body) = + freeIn' (For i bound body) = fvBind (oneName i) $ freeIn' bound <> freeIn' body freeIn' (While cond body) = freeIn' cond <> freeIn' body freeIn' (DeclareMem _ space) = freeIn' space - freeIn' DeclareScalar{} = + freeIn' DeclareScalar {} = mempty - freeIn' DeclareArray{} = + freeIn' DeclareArray {} = mempty freeIn' (Allocate name size space) = freeIn' name <> freeIn' size <> freeIn' space diff --git a/src/Futhark/CodeGen/ImpCode/Kernels.hs b/src/Futhark/CodeGen/ImpCode/Kernels.hs --- a/src/Futhark/CodeGen/ImpCode/Kernels.hs +++ b/src/Futhark/CodeGen/ImpCode/Kernels.hs @@ -1,27 +1,28 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} + -- | Variation of "Futhark.CodeGen.ImpCode" that contains the notion -- of a kernel invocation. module Futhark.CodeGen.ImpCode.Kernels - ( Program - , Function - , FunctionT (Function) - , Code - , KernelCode - , KernelConst (..) - , KernelConstExp - , HostOp (..) - , KernelOp (..) - , Fence (..) - , AtomicOp (..) - , Kernel (..) - , KernelUse (..) - , module Futhark.CodeGen.ImpCode - , module Futhark.IR.Kernels.Sizes + ( Program, + Function, + FunctionT (Function), + Code, + KernelCode, + KernelConst (..), + KernelConstExp, + HostOp (..), + KernelOp (..), + Fence (..), + AtomicOp (..), + Kernel (..), + KernelUse (..), + module Futhark.CodeGen.ImpCode, + module Futhark.IR.Kernels.Sizes, ) - where +where -import Futhark.CodeGen.ImpCode hiding (Function, Code) +import Futhark.CodeGen.ImpCode hiding (Code, Function) import qualified Futhark.CodeGen.ImpCode as Imp import Futhark.IR.Kernels.Sizes import Futhark.IR.Pretty () @@ -41,48 +42,47 @@ -- | A run-time constant related to kernels. newtype KernelConst = SizeConst Name - deriving (Eq, Ord, Show) + deriving (Eq, Ord, Show) -- | An expression whose variables are kernel constants. type KernelConstExp = PrimExp KernelConst -- | An operation that runs on the host (CPU). -data HostOp = CallKernel Kernel - | GetSize VName Name SizeClass - | CmpSizeLe VName Name SizeClass Imp.Exp - | GetSizeMax VName SizeClass - deriving (Show) +data HostOp + = CallKernel Kernel + | GetSize VName Name SizeClass + | CmpSizeLe VName Name SizeClass Imp.Exp + | GetSizeMax VName SizeClass + deriving (Show) -- | A generic kernel containing arbitrary kernel code. data Kernel = Kernel - { kernelBody :: Imp.Code KernelOp - - , kernelUses :: [KernelUse] - -- ^ The host variables referenced by the kernel. - - , kernelNumGroups :: [Imp.Exp] - , kernelGroupSize :: [Imp.Exp] - , kernelName :: Name - -- ^ A short descriptive and _unique_ name - should be - -- alphanumeric and without spaces. - - , kernelFailureTolerant :: Bool - -- ^ If true, this kernel does not need to check - -- whether we are in a failing state, as it can cope. - -- Intuitively, it means that the kernel does not - -- depend on any non-scalar parameters to make control - -- flow decisions. Replication, transpose, and copy - -- kernels are examples of this. - } - deriving (Show) + { kernelBody :: Imp.Code KernelOp, + -- | The host variables referenced by the kernel. + kernelUses :: [KernelUse], + kernelNumGroups :: [Imp.Exp], + kernelGroupSize :: [Imp.Exp], + -- | A short descriptive and _unique_ name - should be + -- alphanumeric and without spaces. + kernelName :: Name, + -- | If true, this kernel does not need to check + -- whether we are in a failing state, as it can cope. + -- Intuitively, it means that the kernel does not + -- depend on any non-scalar parameters to make control + -- flow decisions. Replication, transpose, and copy + -- kernels are examples of this. + kernelFailureTolerant :: Bool + } + deriving (Show) -- | Information about a host-level variable that is used inside this -- kernel. When generating the actual kernel code, this is used to -- deduce which parameters are needed. -data KernelUse = ScalarUse VName PrimType - | MemoryUse VName - | ConstUse VName KernelConstExp - deriving (Eq, Show) +data KernelUse + = ScalarUse VName PrimType + | MemoryUse VName + | ConstUse VName KernelConstExp + deriving (Eq, Show) instance Pretty KernelConst where ppr (SizeConst key) = text "get_size" <> parens (ppr key) @@ -97,14 +97,15 @@ instance Pretty HostOp where ppr (GetSize dest key size_class) = - ppr dest <+> text "<-" <+> - text "get_size" <> parens (commasep [ppr key, ppr size_class]) + ppr dest <+> text "<-" + <+> text "get_size" <> parens (commasep [ppr key, ppr size_class]) ppr (GetSizeMax dest size_class) = ppr dest <+> text "<-" <+> text "get_size_max" <> parens (ppr size_class) ppr (CmpSizeLe dest name size_class x) = - ppr dest <+> text "<-" <+> - text "get_size" <> parens (commasep [ppr name, ppr size_class]) <+> - text "<" <+> ppr x + ppr dest <+> text "<-" + <+> text "get_size" <> parens (commasep [ppr name, ppr size_class]) + <+> text "<" + <+> ppr x ppr (CallKernel c) = ppr c @@ -119,58 +120,63 @@ freeIn' dest instance FreeIn Kernel where - freeIn' kernel = freeIn' (kernelBody kernel) <> - freeIn' [kernelNumGroups kernel, kernelGroupSize kernel] + freeIn' kernel = + freeIn' (kernelBody kernel) + <> freeIn' [kernelNumGroups kernel, kernelGroupSize kernel] instance Pretty Kernel where ppr kernel = - text "kernel" <+> brace - (text "groups" <+> brace (ppr $ kernelNumGroups kernel) - text "group_size" <+> brace (ppr $ kernelGroupSize kernel) - text "uses" <+> brace (commasep $ map ppr $ kernelUses kernel) - text "failure_tolerant" <+> brace (ppr $ kernelFailureTolerant kernel) - text "body" <+> brace (ppr $ kernelBody kernel)) + text "kernel" + <+> brace + ( text "groups" <+> brace (ppr $ kernelNumGroups kernel) + text "group_size" <+> brace (ppr $ kernelGroupSize kernel) + text "uses" <+> brace (commasep $ map ppr $ kernelUses kernel) + text "failure_tolerant" <+> brace (ppr $ kernelFailureTolerant kernel) + text "body" <+> brace (ppr $ kernelBody kernel) + ) -- | When we do a barrier or fence, is it at the local or global -- level? data Fence = FenceLocal | FenceGlobal - deriving (Show) + deriving (Show) -- | An operation that occurs within a kernel body. -data KernelOp = GetGroupId VName Int - | GetLocalId VName Int - | GetLocalSize VName Int - | GetGlobalSize VName Int - | GetGlobalId VName Int - | GetLockstepWidth VName - | Atomic Space AtomicOp - | Barrier Fence - | MemFence Fence - | LocalAlloc VName (Count Bytes Imp.Exp) - | ErrorSync Fence - -- ^ Perform a barrier and also check whether any - -- threads have failed an assertion. Make sure all - -- threads would reach all 'ErrorSync's if any of them - -- do. A failing assertion will jump to the next - -- following 'ErrorSync', so make sure it's not inside - -- control flow or similar. - deriving (Show) +data KernelOp + = GetGroupId VName Int + | GetLocalId VName Int + | GetLocalSize VName Int + | GetGlobalSize VName Int + | GetGlobalId VName Int + | GetLockstepWidth VName + | Atomic Space AtomicOp + | Barrier Fence + | MemFence Fence + | LocalAlloc VName (Count Bytes (Imp.TExp Int64)) + | -- | Perform a barrier and also check whether any + -- threads have failed an assertion. Make sure all + -- threads would reach all 'ErrorSync's if any of them + -- do. A failing assertion will jump to the next + -- following 'ErrorSync', so make sure it's not inside + -- control flow or similar. + ErrorSync Fence + deriving (Show) -- | Atomic operations return the value stored before the update. -- This old value is stored in the first 'VName'. The second 'VName' -- is the memory block to update. The 'Exp' is the new value. -data AtomicOp = AtomicAdd IntType VName VName (Count Elements Imp.Exp) Exp - | AtomicFAdd FloatType VName VName (Count Elements Imp.Exp) Exp - | AtomicSMax IntType VName VName (Count Elements Imp.Exp) Exp - | AtomicSMin IntType VName VName (Count Elements Imp.Exp) Exp - | AtomicUMax IntType VName VName (Count Elements Imp.Exp) Exp - | AtomicUMin IntType VName VName (Count Elements Imp.Exp) Exp - | AtomicAnd IntType VName VName (Count Elements Imp.Exp) Exp - | AtomicOr IntType VName VName (Count Elements Imp.Exp) Exp - | AtomicXor IntType VName VName (Count Elements Imp.Exp) Exp - | AtomicCmpXchg PrimType VName VName (Count Elements Imp.Exp) Exp Exp - | AtomicXchg PrimType VName VName (Count Elements Imp.Exp) Exp - deriving (Show) +data AtomicOp + = AtomicAdd IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicFAdd FloatType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicSMax IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicSMin IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicUMax IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicUMin IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicAnd IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicOr IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicXor IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicCmpXchg PrimType VName VName (Count Elements (Imp.TExp Int64)) Exp Exp + | AtomicXchg PrimType VName VName (Count Elements (Imp.TExp Int64)) Exp + deriving (Show) instance FreeIn AtomicOp where freeIn' (AtomicAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x @@ -187,74 +193,74 @@ instance Pretty KernelOp where ppr (GetGroupId dest i) = - ppr dest <+> "<-" <+> - "get_group_id" <> parens (ppr i) + ppr dest <+> "<-" + <+> "get_group_id" <> parens (ppr i) ppr (GetLocalId dest i) = - ppr dest <+> "<-" <+> - "get_local_id" <> parens (ppr i) + ppr dest <+> "<-" + <+> "get_local_id" <> parens (ppr i) ppr (GetLocalSize dest i) = - ppr dest <+> "<-" <+> - "get_local_size" <> parens (ppr i) + ppr dest <+> "<-" + <+> "get_local_size" <> parens (ppr i) ppr (GetGlobalSize dest i) = - ppr dest <+> "<-" <+> - "get_global_size" <> parens (ppr i) + ppr dest <+> "<-" + <+> "get_global_size" <> parens (ppr i) ppr (GetGlobalId dest i) = - ppr dest <+> "<-" <+> - "get_global_id" <> parens (ppr i) + ppr dest <+> "<-" + <+> "get_global_id" <> parens (ppr i) ppr (GetLockstepWidth dest) = - ppr dest <+> "<-" <+> - "get_lockstep_width()" + ppr dest <+> "<-" + <+> "get_lockstep_width()" ppr (Barrier FenceLocal) = - "local_barrier()" + "local_barrier()" ppr (Barrier FenceGlobal) = - "global_barrier()" + "global_barrier()" ppr (MemFence FenceLocal) = - "mem_fence_local()" + "mem_fence_local()" ppr (MemFence FenceGlobal) = - "mem_fence_global()" + "mem_fence_global()" ppr (LocalAlloc name size) = - ppr name <+> equals <+> "local_alloc" <> parens (ppr size) + ppr name <+> equals <+> "local_alloc" <> parens (ppr size) ppr (ErrorSync FenceLocal) = - "error_sync_local()" + "error_sync_local()" ppr (ErrorSync FenceGlobal) = - "error_sync_global()" + "error_sync_global()" ppr (Atomic _ (AtomicAdd t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_add_" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_add_" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) ppr (Atomic _ (AtomicFAdd t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_fadd_" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_fadd_" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) ppr (Atomic _ (AtomicSMax t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_smax" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_smax" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) ppr (Atomic _ (AtomicSMin t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_smin" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_smin" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) ppr (Atomic _ (AtomicUMax t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_umax" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_umax" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) ppr (Atomic _ (AtomicUMin t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_umin" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_umin" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) ppr (Atomic _ (AtomicAnd t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_and" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_and" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) ppr (Atomic _ (AtomicOr t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_or" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_or" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) ppr (Atomic _ (AtomicXor t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_xor" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_xor" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) ppr (Atomic _ (AtomicCmpXchg t old arr ind x y)) = - ppr old <+> "<-" <+> "atomic_cmp_xchg" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x, ppr y]) + ppr old <+> "<-" <+> "atomic_cmp_xchg" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x, ppr y]) ppr (Atomic _ (AtomicXchg t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_xchg" <> ppr t <> - parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr old <+> "<-" <+> "atomic_xchg" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) instance FreeIn KernelOp where freeIn' (Atomic _ op) = freeIn' op freeIn' _ = mempty brace :: Doc -> Doc -brace body = " {" indent 2 body "}" +brace body = " {" indent 2 body "}" diff --git a/src/Futhark/CodeGen/ImpCode/OpenCL.hs b/src/Futhark/CodeGen/ImpCode/OpenCL.hs --- a/src/Futhark/CodeGen/ImpCode/OpenCL.hs +++ b/src/Futhark/CodeGen/ImpCode/OpenCL.hs @@ -7,49 +7,49 @@ -- The imperative code has been augmented with a 'LaunchKernel' -- operation that allows one to execute an OpenCL kernel. module Futhark.CodeGen.ImpCode.OpenCL - ( Program (..) - , Function - , FunctionT (Function) - , Code - , KernelName - , KernelArg (..) - , OpenCL (..) - , KernelSafety(..) - , numFailureParams - , KernelTarget (..) - , FailureMsg(..) - , module Futhark.CodeGen.ImpCode - , module Futhark.IR.Kernels.Sizes - ) - where + ( Program (..), + Function, + FunctionT (Function), + Code, + KernelName, + KernelArg (..), + OpenCL (..), + KernelSafety (..), + numFailureParams, + KernelTarget (..), + FailureMsg (..), + module Futhark.CodeGen.ImpCode, + module Futhark.IR.Kernels.Sizes, + ) +where import qualified Data.Map as M - -import Futhark.CodeGen.ImpCode hiding (Function, Code) -import Futhark.IR.Kernels.Sizes +import Futhark.CodeGen.ImpCode hiding (Code, Function) import qualified Futhark.CodeGen.ImpCode as Imp - +import Futhark.IR.Kernels.Sizes import Futhark.Util.Pretty -- | An program calling OpenCL kernels. -data Program = Program { openClProgram :: String - , openClPrelude :: String - -- ^ Must be prepended to the program. - , openClKernelNames :: M.Map KernelName KernelSafety - , openClUsedTypes :: [PrimType] - -- ^ So we can detect whether the device is capable. - , openClSizes :: M.Map Name SizeClass - -- ^ Runtime-configurable constants. - , openClFailures :: [FailureMsg] - -- ^ Assertion failure error messages. - , hostDefinitions :: Definitions OpenCL - } +data Program = Program + { openClProgram :: String, + -- | Must be prepended to the program. + openClPrelude :: String, + openClKernelNames :: M.Map KernelName KernelSafety, + -- | So we can detect whether the device is capable. + openClUsedTypes :: [PrimType], + -- | Runtime-configurable constants. + openClSizes :: M.Map Name SizeClass, + -- | Assertion failure error messages. + openClFailures :: [FailureMsg], + hostDefinitions :: Definitions OpenCL + } -- | Something that can go wrong in a kernel. Part of the machinery -- for reporting error messages from within kernels. -data FailureMsg = FailureMsg { failureError :: ErrorMsg Exp - , failureBacktrace :: String - } +data FailureMsg = FailureMsg + { failureError :: ErrorMsg Exp, + failureBacktrace :: String + } -- | A function calling OpenCL kernels. type Function = Imp.Function OpenCL @@ -61,31 +61,32 @@ type KernelName = Name -- | An argument to be passed to a kernel. -data KernelArg = ValueKArg Exp PrimType - -- ^ Pass the value of this scalar expression as argument. - | MemKArg VName - -- ^ Pass this pointer as argument. - | SharedMemoryKArg (Count Bytes Exp) - -- ^ Create this much local memory per workgroup. - deriving (Show) +data KernelArg + = -- | Pass the value of this scalar expression as argument. + ValueKArg Exp PrimType + | -- | Pass this pointer as argument. + MemKArg VName + | -- | Create this much local memory per workgroup. + SharedMemoryKArg (Count Bytes Exp) + deriving (Show) -- | Whether a kernel can potentially fail (because it contains bounds -- checks and such). data MayFail = MayFail | CannotFail - deriving (Show) + deriving (Show) -- | Information about bounds checks and how sensitive it is to -- errors. Ordered by least demanding to most. data KernelSafety - = SafetyNone - -- ^ Does not need to know if we are in a failing state, and also + = -- | Does not need to know if we are in a failing state, and also -- cannot fail. - | SafetyCheap - -- ^ Needs to be told if there's a global failure, and that's it, + SafetyNone + | -- | Needs to be told if there's a global failure, and that's it, -- and cannot fail. - | SafetyFull - -- ^ Needs all parameters, may fail itself. - deriving (Eq, Ord, Show) + SafetyCheap + | -- | Needs all parameters, may fail itself. + SafetyFull + deriving (Eq, Ord, Show) -- | How many leading failure arguments we must pass when launching a -- kernel with these safety characteristics. @@ -95,16 +96,18 @@ numFailureParams SafetyFull = 3 -- | Host-level OpenCL operation. -data OpenCL = LaunchKernel KernelSafety KernelName [KernelArg] [Exp] [Exp] - | GetSize VName Name - | CmpSizeLe VName Name Exp - | GetSizeMax VName SizeClass - deriving (Show) +data OpenCL + = LaunchKernel KernelSafety KernelName [KernelArg] [Exp] [Exp] + | GetSize VName Name + | CmpSizeLe VName Name Exp + | GetSizeMax VName SizeClass + deriving (Show) -- | The target platform when compiling imperative code to a 'Program' -data KernelTarget = TargetOpenCL - | TargetCUDA - deriving (Eq) +data KernelTarget + = TargetOpenCL + | TargetCUDA + deriving (Eq) instance Pretty OpenCL where ppr = text . show diff --git a/src/Futhark/CodeGen/ImpCode/Sequential.hs b/src/Futhark/CodeGen/ImpCode/Sequential.hs --- a/src/Futhark/CodeGen/ImpCode/Sequential.hs +++ b/src/Futhark/CodeGen/ImpCode/Sequential.hs @@ -1,17 +1,16 @@ -- | Sequential imperative code. module Futhark.CodeGen.ImpCode.Sequential - ( Program - , Function - , FunctionT (Function) - , Code - , Sequential - , module Futhark.CodeGen.ImpCode - ) - where + ( Program, + Function, + FunctionT (Function), + Code, + Sequential, + module Futhark.CodeGen.ImpCode, + ) +where -import Futhark.CodeGen.ImpCode hiding (Function, Code) +import Futhark.CodeGen.ImpCode hiding (Code, Function) import qualified Futhark.CodeGen.ImpCode as Imp - import Futhark.Util.Pretty -- | An imperative program. diff --git a/src/Futhark/CodeGen/ImpGen.hs b/src/Futhark/CodeGen/ImpGen.hs --- a/src/Futhark/CodeGen/ImpGen.hs +++ b/src/Futhark/CodeGen/ImpGen.hs @@ -1,1418 +1,1736 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, LambdaCase, FlexibleInstances, MultiParamTypeClasses #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE Strict #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE Trustworthy #-} -module Futhark.CodeGen.ImpGen - ( -- * Entry Points - compileProg - - -- * Pluggable Compiler - , OpCompiler - , ExpCompiler - , CopyCompiler - , StmsCompiler - , AllocCompiler - , Operations (..) - , defaultOperations - , MemLocation (..) - , MemEntry (..) - , ScalarEntry (..) - - -- * Monadic Compiler Interface - , ImpM - , localDefaultSpace, askFunction - , newVNameForFun, nameForFun - , askEnv, localEnv - , localOps - , VTable - , getVTable - , localVTable - , subImpM - , subImpM_ - , emit - , emitFunction - , hasFunction - , collect - , collect' - , comment - , VarEntry (..) - , ArrayEntry (..) - - -- * Lookups - , lookupVar - , lookupArray - , lookupMemory - - -- * Building Blocks - , ToExp(..) - , compileAlloc - , everythingVolatile - , compileBody - , compileBody' - , compileLoopBody - , defCompileStms - , compileStms - , compileExp - , defCompileExp - , fullyIndexArray - , fullyIndexArray' - , copy - , copyDWIM - , copyDWIMFix - , copyElementWise - , typeSize - - -- * Constructing code. - , dLParams - , dFParams - , dScope - , dArray - , dPrim, dPrimVol_, dPrim_, dPrimV_, dPrimV, dPrimVE - - , sFor, sWhile - , sComment - , sIf, sWhen, sUnless - , sOp - , sDeclareMem, sAlloc, sAlloc_ - , sArray, sArrayInMem, sAllocArray, sAllocArrayPerm, sStaticArray - , sWrite, sUpdate - , sLoopNest - , (<--) - - , function - - , warn - , module Language.Futhark.Warnings - ) - where - -import Control.Monad.Reader -import Control.Monad.State -import Control.Monad.Writer -import Control.Parallel.Strategies -import Data.Bifunctor (first) -import qualified Data.DList as DL -import Data.Either -import qualified Data.Map.Strict as M -import qualified Data.Set as S -import Data.Maybe -import Data.List (find, sortOn, genericLength) - -import qualified Futhark.CodeGen.ImpCode as Imp -import Futhark.CodeGen.ImpCode - (Count, Bytes, Elements, - bytes, elements, withElemType) -import Futhark.IR.Mem -import Futhark.IR.SOACS (SOACS) -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Construct (fullSliceNum) -import Futhark.MonadFreshNames -import Futhark.Util -import Futhark.Util.Loc (noLoc) -import Language.Futhark.Warnings - --- | How to compile an t'Op'. -type OpCompiler lore r op = Pattern lore -> Op lore -> ImpM lore r op () - --- | How to compile some 'Stms'. -type StmsCompiler lore r op = Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op () - --- | How to compile an 'Exp'. -type ExpCompiler lore r op = Pattern lore -> Exp lore -> ImpM lore r op () - -type CopyCompiler lore r op = PrimType - -> MemLocation -> Slice Imp.Exp - -> MemLocation -> Slice Imp.Exp - -> ImpM lore r op () - --- | An alternate way of compiling an allocation. -type AllocCompiler lore r op = VName -> Count Bytes Imp.Exp -> ImpM lore r op () - -data Operations lore r op = Operations { opsExpCompiler :: ExpCompiler lore r op - , opsOpCompiler :: OpCompiler lore r op - , opsStmsCompiler :: StmsCompiler lore r op - , opsCopyCompiler :: CopyCompiler lore r op - , opsAllocCompilers :: M.Map Space (AllocCompiler lore r op) - } - --- | An operations set for which the expression compiler always --- returns 'defCompileExp'. -defaultOperations :: (Mem lore, FreeIn op) => - OpCompiler lore r op -> Operations lore r op -defaultOperations opc = Operations { opsExpCompiler = defCompileExp - , opsOpCompiler = opc - , opsStmsCompiler = defCompileStms - , opsCopyCompiler = defaultCopy - , opsAllocCompilers = mempty - } - --- | When an array is dared, this is where it is stored. -data MemLocation = MemLocation { memLocationName :: VName - , memLocationShape :: [Imp.DimSize] - , memLocationIxFun :: IxFun.IxFun Imp.Exp - } - deriving (Eq, Show) - -data ArrayEntry = ArrayEntry { - entryArrayLocation :: MemLocation - , entryArrayElemType :: PrimType - } - deriving (Show) - -entryArrayShape :: ArrayEntry -> [Imp.DimSize] -entryArrayShape = memLocationShape . entryArrayLocation - -newtype MemEntry = MemEntry { entryMemSpace :: Imp.Space } - deriving (Show) - -newtype ScalarEntry = ScalarEntry { - entryScalarType :: PrimType - } - deriving (Show) - --- | Every non-scalar variable must be associated with an entry. -data VarEntry lore = ArrayVar (Maybe (Exp lore)) ArrayEntry - | ScalarVar (Maybe (Exp lore)) ScalarEntry - | MemVar (Maybe (Exp lore)) MemEntry - deriving (Show) - --- | When compiling an expression, this is a description of where the --- result should end up. The integer is a reference to the construct --- that gave rise to this destination (for patterns, this will be the --- tag of the first name in the pattern). This can be used to make --- the generated code easier to relate to the original code. -data Destination = Destination { destinationTag :: Maybe Int - , valueDestinations :: [ValueDestination] } - deriving (Show) - -data ValueDestination = ScalarDestination VName - | MemoryDestination VName - | ArrayDestination (Maybe MemLocation) - -- ^ The 'MemLocation' is 'Just' if a copy if - -- required. If it is 'Nothing', then a - -- copy/assignment of a memory block somewhere - -- takes care of this array. - deriving (Show) - -data Env lore r op = Env { - envExpCompiler :: ExpCompiler lore r op - , envStmsCompiler :: StmsCompiler lore r op - , envOpCompiler :: OpCompiler lore r op - , envCopyCompiler :: CopyCompiler lore r op - , envAllocCompilers :: M.Map Space (AllocCompiler lore r op) - , envDefaultSpace :: Imp.Space - , envVolatility :: Imp.Volatility - , envEnv :: r - -- ^ User-extensible environment. - , envFunction :: Maybe Name - -- ^ Name of the function we are compiling, if any. - , envAttrs :: Attrs - -- ^ The set of attributes that are active on the enclosing - -- statements (including the one we are currently compiling). - } - -newEnv :: r -> Operations lore r op -> Imp.Space -> Env lore r op -newEnv r ops ds = - Env { envExpCompiler = opsExpCompiler ops - , envStmsCompiler = opsStmsCompiler ops - , envOpCompiler = opsOpCompiler ops - , envCopyCompiler = opsCopyCompiler ops - , envAllocCompilers = mempty - , envDefaultSpace = ds - , envVolatility = Imp.Nonvolatile - , envEnv = r - , envFunction = Nothing - , envAttrs = mempty - } - --- | The symbol table used during compilation. -type VTable lore = M.Map VName (VarEntry lore) - -data ImpState lore r op = - ImpState { stateVTable :: VTable lore - , stateFunctions :: Imp.Functions op - , stateCode :: Imp.Code op - , stateWarnings :: Warnings - , stateNameSource :: VNameSource - } - -newState :: VNameSource -> ImpState lore r op -newState = ImpState mempty mempty mempty mempty - -newtype ImpM lore r op a = - ImpM (ReaderT (Env lore r op) (State (ImpState lore r op)) a) - deriving (Functor, Applicative, Monad, - MonadState (ImpState lore r op), - MonadReader (Env lore r op)) - -instance MonadFreshNames (ImpM lore r op) where - getNameSource = gets stateNameSource - putNameSource src = modify $ \s -> s { stateNameSource = src } - --- Cannot be an KernelsMem scope because the index functions have --- the wrong leaves (VName instead of Imp.Exp). -instance HasScope SOACS (ImpM lore r op) where - askScope = gets $ M.map (LetName . entryType) . stateVTable - where entryType (MemVar _ memEntry) = - Mem (entryMemSpace memEntry) - entryType (ArrayVar _ arrayEntry) = - Array - (entryArrayElemType arrayEntry) - (Shape $ entryArrayShape arrayEntry) - NoUniqueness - entryType (ScalarVar _ scalarEntry) = - Prim $ entryScalarType scalarEntry - -runImpM :: ImpM lore r op a - -> r -> Operations lore r op -> Imp.Space -> ImpState lore r op - -> (a, ImpState lore r op) -runImpM (ImpM m) r ops space = runState (runReaderT m $ newEnv r ops space) - -subImpM_ :: r' -> Operations lore r' op' -> ImpM lore r' op' a - -> ImpM lore r op (Imp.Code op') -subImpM_ r ops m = snd <$> subImpM r ops m - -subImpM :: r' -> Operations lore r' op' -> ImpM lore r' op' a - -> ImpM lore r op (a, Imp.Code op') -subImpM r ops (ImpM m) = do - env <- ask - s <- get - - let env' = env { envExpCompiler = opsExpCompiler ops - , envStmsCompiler = opsStmsCompiler ops - , envCopyCompiler = opsCopyCompiler ops - , envOpCompiler = opsOpCompiler ops - , envAllocCompilers = opsAllocCompilers ops - , envEnv = r - } - s' = ImpState { stateVTable = stateVTable s - , stateFunctions = mempty - , stateCode = mempty - , stateNameSource = stateNameSource s - , stateWarnings = mempty - } - (x, s'') = runState (runReaderT m env') s' - - putNameSource $ stateNameSource s'' - warnings $ stateWarnings s'' - return (x, stateCode s'') - --- | Execute a code generation action, returning the code that was --- emitted. -collect :: ImpM lore r op () -> ImpM lore r op (Imp.Code op) -collect = fmap snd . collect' - -collect' :: ImpM lore r op a -> ImpM lore r op (a, Imp.Code op) -collect' m = do - prev_code <- gets stateCode - modify $ \s -> s { stateCode = mempty } - x <- m - new_code <- gets stateCode - modify $ \s -> s { stateCode = prev_code } - return (x, new_code) - --- | Execute a code generation action, wrapping the generated code --- within a 'Imp.Comment' with the given description. -comment :: String -> ImpM lore r op () -> ImpM lore r op () -comment desc m = do code <- collect m - emit $ Imp.Comment desc code - --- | Emit some generated imperative code. -emit :: Imp.Code op -> ImpM lore r op () -emit code = modify $ \s -> s { stateCode = stateCode s <> code } - -warnings :: Warnings -> ImpM lore r op () -warnings ws = modify $ \s -> s { stateWarnings = ws <> stateWarnings s} - --- | Emit a warning about something the user should be aware of. -warn :: Located loc => loc -> [loc] -> String -> ImpM lore r op () -warn loc locs problem = - warnings $ singleWarning' (srclocOf loc) (map srclocOf locs) problem - --- | Emit a function in the generated code. -emitFunction :: Name -> Imp.Function op -> ImpM lore r op () -emitFunction fname fun = do - Imp.Functions fs <- gets stateFunctions - modify $ \s -> s { stateFunctions = Imp.Functions $ (fname,fun) : fs } - --- | Check if a function of a given name exists. -hasFunction :: Name -> ImpM lore r op Bool -hasFunction fname = gets $ \s -> let Imp.Functions fs = stateFunctions s - in isJust $ lookup fname fs - -constsVTable :: Mem lore => Stms lore -> VTable lore -constsVTable = foldMap stmVtable - where stmVtable (Let pat _ e) = - foldMap (peVtable e) $ patternElements pat - peVtable e (PatElem name dec) = - M.singleton name $ memBoundToVarEntry (Just e) dec - -compileProg :: (Mem lore, FreeIn op, MonadFreshNames m) => - r -> Operations lore r op -> Imp.Space - -> Prog lore -> m (Warnings, Imp.Definitions op) -compileProg r ops space (Prog consts funs) = - modifyNameSource $ \src -> - let (_, ss) = - unzip $ parMap rpar (compileFunDef' src) funs - free_in_funs = - freeIn $ mconcat $ map stateFunctions ss - (consts', s') = - runImpM (compileConsts free_in_funs consts) r ops space $ - combineStates ss - in ((stateWarnings s', - Imp.Definitions consts' (stateFunctions s')), - stateNameSource s') - where compileFunDef' src fdef = - runImpM (compileFunDef fdef) r ops space - (newState src) { stateVTable = constsVTable consts } - - combineStates ss = - let Imp.Functions funs' = mconcat $ map stateFunctions ss - src = mconcat (map stateNameSource ss) - in (newState src) { stateFunctions = - Imp.Functions $ M.toList $ M.fromList funs' - , stateWarnings = - mconcat $ map stateWarnings ss - } - -compileConsts :: Names -> Stms lore -> ImpM lore r op (Imp.Constants op) -compileConsts used_consts stms = do - code <- collect $ compileStms used_consts stms $ pure () - pure $ uncurry Imp.Constants $ first DL.toList $ extract code - where -- Fish out those top-level declarations in the constant - -- initialisation code that are free in the functions. - extract (x Imp.:>>: y) = - extract x <> extract y - extract (Imp.DeclareMem name space) - | name `nameIn` used_consts = - (DL.singleton $ Imp.MemParam name space, - mempty) - extract (Imp.DeclareScalar name _ t) - | name `nameIn` used_consts = - (DL.singleton $ Imp.ScalarParam name t, - mempty) - extract s = - (mempty, s) - -compileInParam :: Mem lore => - FParam lore -> ImpM lore r op (Either Imp.Param ArrayDecl) -compileInParam fparam = case paramDec fparam of - MemPrim bt -> - return $ Left $ Imp.ScalarParam name bt - MemMem space -> - return $ Left $ Imp.MemParam name space - MemArray bt shape _ (ArrayIn mem ixfun) -> - return $ Right $ ArrayDecl name bt $ - MemLocation mem (shapeDims shape) $ fmap (toExp' int32) ixfun - where name = paramName fparam - -data ArrayDecl = ArrayDecl VName PrimType MemLocation - -fparamSizes :: Typed dec => Param dec -> S.Set VName -fparamSizes = S.fromList . subExpVars . arrayDims . paramType - -compileInParams :: Mem lore => - [FParam lore] -> [EntryPointType] - -> ImpM lore r op ([Imp.Param], [ArrayDecl], [Imp.ExternalValue]) -compileInParams params orig_epts = do - let (ctx_params, val_params) = - splitAt (length params - sum (map entryPointSize orig_epts)) params - (inparams, arrayds) <- partitionEithers <$> mapM compileInParam (ctx_params++val_params) - let findArray x = find (isArrayDecl x) arrayds - sizes = mconcat $ map fparamSizes $ ctx_params++val_params - - summaries = M.fromList $ mapMaybe memSummary params - where memSummary param - | MemMem space <- paramDec param = - Just (paramName param, space) - | otherwise = - Nothing - - findMemInfo :: VName -> Maybe Space - findMemInfo = flip M.lookup summaries - - mkValueDesc fparam signedness = - case (findArray $ paramName fparam, paramType fparam) of - (Just (ArrayDecl _ bt (MemLocation mem shape _)), _) -> do - memspace <- findMemInfo mem - Just $ Imp.ArrayValue mem memspace bt signedness shape - (_, Prim bt) - | paramName fparam `S.member` sizes -> - Nothing - | otherwise -> - Just $ Imp.ScalarValue bt signedness $ paramName fparam - _ -> - Nothing - - mkExts (TypeOpaque desc n:epts) fparams = - let (fparams',rest) = splitAt n fparams - in Imp.OpaqueValue desc - (mapMaybe (`mkValueDesc` Imp.TypeDirect) fparams') : - mkExts epts rest - mkExts (TypeUnsigned:epts) (fparam:fparams) = - maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeUnsigned) ++ - mkExts epts fparams - mkExts (TypeDirect:epts) (fparam:fparams) = - maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeDirect) ++ - mkExts epts fparams - mkExts _ _ = [] - - return (inparams, arrayds, mkExts orig_epts val_params) - where isArrayDecl x (ArrayDecl y _ _) = x == y - -compileOutParams :: Mem lore => - [RetType lore] -> [EntryPointType] - -> ImpM lore r op ([Imp.ExternalValue], [Imp.Param], Destination) -compileOutParams orig_rts orig_epts = do - ((extvs, dests), (outparams,ctx_dests)) <- - runWriterT $ evalStateT (mkExts orig_epts orig_rts) (M.empty, M.empty) - let ctx_dests' = map snd $ sortOn fst $ M.toList ctx_dests - return (extvs, outparams, Destination Nothing $ ctx_dests' <> dests) - where imp = lift . lift - - mkExts (TypeOpaque desc n:epts) rts = do - let (rts',rest) = splitAt n rts - (evs, dests) <- unzip <$> zipWithM mkParam rts' (repeat Imp.TypeDirect) - (more_values, more_dests) <- mkExts epts rest - return (Imp.OpaqueValue desc evs : more_values, - dests ++ more_dests) - mkExts (TypeUnsigned:epts) (rt:rts) = do - (ev,dest) <- mkParam rt Imp.TypeUnsigned - (more_values, more_dests) <- mkExts epts rts - return (Imp.TransparentValue ev : more_values, - dest : more_dests) - mkExts (TypeDirect:epts) (rt:rts) = do - (ev,dest) <- mkParam rt Imp.TypeDirect - (more_values, more_dests) <- mkExts epts rts - return (Imp.TransparentValue ev : more_values, - dest : more_dests) - mkExts _ _ = return ([], []) - - mkParam MemMem{} _ = - error "Functions may not explicitly return memory blocks." - mkParam (MemPrim t) ept = do - out <- imp $ newVName "scalar_out" - tell ([Imp.ScalarParam out t], mempty) - return (Imp.ScalarValue t ept out, ScalarDestination out) - mkParam (MemArray t shape _ dec) ept = do - space <- asks envDefaultSpace - memout <- case dec of - ReturnsNewBlock _ x _ixfun -> do - memout <- imp $ newVName "out_mem" - tell ([Imp.MemParam memout space], - M.singleton x $ MemoryDestination memout) - return memout - ReturnsInBlock memout _ -> - return memout - resultshape <- mapM inspectExtSize $ shapeDims shape - return (Imp.ArrayValue memout space t ept resultshape, - ArrayDestination Nothing) - - inspectExtSize (Ext x) = do - (memseen,arrseen) <- get - case M.lookup x arrseen of - Nothing -> do - out <- imp $ newVName "out_arrsize" - tell ([Imp.ScalarParam out int32], - M.singleton x $ ScalarDestination out) - put (memseen, M.insert x out arrseen) - return $ Var out - Just out -> - return $ Var out - inspectExtSize (Free se) = - return se - -compileFunDef :: Mem lore => - FunDef lore - -> ImpM lore r op () -compileFunDef (FunDef entry _ fname rettype params body) = - local (\env -> env { envFunction = Just fname }) $ do - ((outparams, inparams, results, args), body') <- collect' compile - emitFunction fname $ Imp.Function (isJust entry) outparams inparams body' results args - where params_entry = maybe (replicate (length params) TypeDirect) fst entry - ret_entry = maybe (replicate (length rettype) TypeDirect) snd entry - compile = do - (inparams, arrayds, args) <- compileInParams params params_entry - (results, outparams, Destination _ dests) <- compileOutParams rettype ret_entry - addFParams params - addArrays arrayds - - let Body _ stms ses = body - compileStms (freeIn ses) stms $ - forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se [] - - return (outparams, inparams, results, args) - -compileBody :: (Mem lore) => Pattern lore -> Body lore -> ImpM lore r op () -compileBody pat (Body _ bnds ses) = do - Destination _ dests <- destinationFromPattern pat - compileStms (freeIn ses) bnds $ - forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se [] - -compileBody' :: [Param dec] -> Body lore -> ImpM lore r op () -compileBody' params (Body _ bnds ses) = - compileStms (freeIn ses) bnds $ - forM_ (zip params ses) $ \(param, se) -> copyDWIM (paramName param) [] se [] - -compileLoopBody :: Typed dec => [Param dec] -> Body lore -> ImpM lore r op () -compileLoopBody mergeparams (Body _ bnds ses) = do - -- We cannot write the results to the merge parameters immediately, - -- as some of the results may actually *be* merge parameters, and - -- would thus be clobbered. Therefore, we first copy to new - -- variables mirroring the merge parameters, and then copy this - -- buffer to the merge parameters. This is efficient, because the - -- operations are all scalar operations. - tmpnames <- mapM (newVName . (++"_tmp") . baseString . paramName) mergeparams - compileStms (freeIn ses) bnds $ do - copy_to_merge_params <- forM (zip3 mergeparams tmpnames ses) $ \(p,tmp,se) -> - case typeOf p of - Prim pt -> do - emit $ Imp.DeclareScalar tmp Imp.Nonvolatile pt - emit $ Imp.SetScalar tmp $ toExp' pt se - return $ emit $ Imp.SetScalar (paramName p) $ Imp.var tmp pt - Mem space | Var v <- se -> do - emit $ Imp.DeclareMem tmp space - emit $ Imp.SetMem tmp v space - return $ emit $ Imp.SetMem (paramName p) tmp space - _ -> return $ return () - sequence_ copy_to_merge_params - -compileStms :: Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op () -compileStms alive_after_stms all_stms m = do - cb <- asks envStmsCompiler - cb alive_after_stms all_stms m - -defCompileStms :: (Mem lore, FreeIn op) => - Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op () -defCompileStms alive_after_stms all_stms m = - -- We keep track of any memory blocks produced by the statements, - -- and after the last time that memory block is used, we insert a - -- Free. This is very conservative, but can cut down on lifetimes - -- in some cases. - void $ compileStms' mempty $ stmsToList all_stms - where compileStms' allocs (Let pat aux e:bs) = do - dVars (Just e) (patternElements pat) - - e_code <- localAttrs (stmAuxAttrs aux) $ - collect $ compileExp pat e - (live_after, bs_code) <- collect' $ compileStms' (patternAllocs pat <> allocs) bs - let dies_here v = not (v `nameIn` live_after) && - v `nameIn` freeIn e_code - to_free = S.filter (dies_here . fst) allocs - - emit e_code - mapM_ (emit . uncurry Imp.Free) to_free - emit bs_code - - return $ freeIn e_code <> live_after - compileStms' _ [] = do - code <- collect m - emit code - return $ freeIn code <> alive_after_stms - - patternAllocs = S.fromList . mapMaybe isMemPatElem . patternElements - isMemPatElem pe = case patElemType pe of - Mem space -> Just (patElemName pe, space) - _ -> Nothing - -compileExp :: Pattern lore -> Exp lore -> ImpM lore r op () -compileExp pat e = do - ec <- asks envExpCompiler - ec pat e - -defCompileExp :: (Mem lore) => - Pattern lore -> Exp lore -> ImpM lore r op () - -defCompileExp pat (If cond tbranch fbranch _) = do - tcode <- collect $ compileBody pat tbranch - fcode <- collect $ compileBody pat fbranch - emit $ Imp.If (toExp' Bool cond) tcode fcode - -defCompileExp pat (Apply fname args _ _) = do - dest <- destinationFromPattern pat - targets <- funcallTargets dest - args' <- catMaybes <$> mapM compileArg args - emit $ Imp.Call targets fname args' - where compileArg (se, _) = do - t <- subExpType se - case (se, t) of - (_, Prim pt) -> return $ Just $ Imp.ExpArg $ toExp' pt se - (Var v, Mem{}) -> return $ Just $ Imp.MemArg v - _ -> return Nothing - -defCompileExp pat (BasicOp op) = defCompileBasicOp pat op - -defCompileExp pat (DoLoop ctx val form body) = do - attrs <- askAttrs - when ("unroll" `inAttrs` attrs) $ - warn (noLoc::SrcLoc) [] "#[unroll] on loop with unknown number of iterations." -- FIXME: no location. - - dFParams mergepat - forM_ merge $ \(p, se) -> - when ((==0) $ arrayRank $ paramType p) $ - copyDWIM (paramName p) [] se [] - - let doBody = compileLoopBody mergepat body - - case form of - ForLoop i it bound loopvars -> do - let setLoopParam (p,a) - | Prim _ <- paramType p = - copyDWIM (paramName p) [] (Var a) [DimFix $ Imp.vi32 i] - | otherwise = - return () - - dLParams $ map fst loopvars - sFor' i it (toExp' (IntType it) bound) $ - mapM_ setLoopParam loopvars >> doBody - WhileLoop cond -> - sWhile (Imp.var cond Bool) doBody - - Destination _ pat_dests <- destinationFromPattern pat - forM_ (zip pat_dests $ map (Var . paramName . fst) merge) $ \(d, r) -> - copyDWIMDest d [] r [] - - where merge = ctx ++ val - mergepat = map fst merge - -defCompileExp pat (Op op) = do - opc <- asks envOpCompiler - opc pat op - -defCompileBasicOp :: Mem lore => - Pattern lore -> BasicOp -> ImpM lore r op () - -defCompileBasicOp (Pattern _ [pe]) (SubExp se) = - copyDWIM (patElemName pe) [] se [] - -defCompileBasicOp (Pattern _ [pe]) (Opaque se) = - copyDWIM (patElemName pe) [] se [] - -defCompileBasicOp (Pattern _ [pe]) (UnOp op e) = do - e' <- toExp e - patElemName pe <-- Imp.UnOpExp op e' - -defCompileBasicOp (Pattern _ [pe]) (ConvOp conv e) = do - e' <- toExp e - patElemName pe <-- Imp.ConvOpExp conv e' - -defCompileBasicOp (Pattern _ [pe]) (BinOp bop x y) = do - x' <- toExp x - y' <- toExp y - patElemName pe <-- Imp.BinOpExp bop x' y' - -defCompileBasicOp (Pattern _ [pe]) (CmpOp bop x y) = do - x' <- toExp x - y' <- toExp y - patElemName pe <-- Imp.CmpOpExp bop x' y' - -defCompileBasicOp _ (Assert e msg loc) = do - e' <- toExp e - msg' <- traverse toExp msg - emit $ Imp.Assert e' msg' loc - - attrs <- askAttrs - when (AttrComp "warn" ["safety_checks"] `inAttrs` attrs) $ - uncurry warn loc "Safety check required at run-time." - -defCompileBasicOp (Pattern _ [pe]) (Index src slice) - | Just idxs <- sliceIndices slice = - copyDWIM (patElemName pe) [] (Var src) $ map (DimFix . toExp' int32) idxs - -defCompileBasicOp _ Index{} = - return () - -defCompileBasicOp (Pattern _ [pe]) (Update _ slice se) = - sUpdate (patElemName pe) (map (fmap (toExp' int32)) slice) se - -defCompileBasicOp (Pattern _ [pe]) (Replicate (Shape ds) se) = do - ds' <- mapM toExp ds - is <- replicateM (length ds) (newVName "i") - copy_elem <- collect $ copyDWIM (patElemName pe) (map (DimFix . Imp.vi32) is) se [] - emit $ foldl (.) id (zipWith (`Imp.For` Int32) is ds') copy_elem - -defCompileBasicOp _ Scratch{} = - return () - -defCompileBasicOp (Pattern [] [pe]) (Iota n e s it) = do - n' <- toExp n - e' <- toExp e - s' <- toExp s - sFor "i" n' $ \i -> do - let i' = sExt it i - x <- dPrimV "x" $ e' + i' * s' - copyDWIM (patElemName pe) [DimFix i] (Var x) [] - -defCompileBasicOp (Pattern _ [pe]) (Copy src) = - copyDWIM (patElemName pe) [] (Var src) [] - -defCompileBasicOp (Pattern _ [pe]) (Manifest _ src) = - copyDWIM (patElemName pe) [] (Var src) [] - -defCompileBasicOp (Pattern _ [pe]) (Concat i x ys _) = do - offs_glb <- dPrim "tmp_offs" int32 - emit $ Imp.SetScalar offs_glb 0 - - forM_ (x:ys) $ \y -> do - y_dims <- arrayDims <$> lookupType y - let rows = case drop i y_dims of - [] -> error $ "defCompileBasicOp Concat: empty array shape for " ++ pretty y - r:_ -> toExp' int32 r - skip_dims = take i y_dims - sliceAllDim d = DimSlice 0 d 1 - skip_slices = map (sliceAllDim . toExp' int32) skip_dims - destslice = skip_slices ++ [DimSlice (Imp.vi32 offs_glb) rows 1] - copyDWIM (patElemName pe) destslice (Var y) [] - emit $ Imp.SetScalar offs_glb $ Imp.var offs_glb int32 + rows - -defCompileBasicOp (Pattern [] [pe]) (ArrayLit es _) - | Just vs@(v:_) <- mapM isLiteral es = do - dest_mem <- entryArrayLocation <$> lookupArray (patElemName pe) - dest_space <- entryMemSpace <$> lookupMemory (memLocationName dest_mem) - let t = primValueType v - static_array <- newVNameForFun "static_array" - emit $ Imp.DeclareArray static_array dest_space t $ Imp.ArrayValues vs - let static_src = MemLocation static_array [intConst Int32 $ fromIntegral $ length es] $ - IxFun.iota [fromIntegral $ length es] - entry = MemVar Nothing $ MemEntry dest_space - addVar static_array entry - let slice = [DimSlice 0 (genericLength es) 1] - copy t dest_mem slice static_src slice - | otherwise = - forM_ (zip [0..] es) $ \(i,e) -> - copyDWIM (patElemName pe) [DimFix $ fromInteger i] e [] - - where isLiteral (Constant v) = Just v - isLiteral _ = Nothing - -defCompileBasicOp _ Rearrange{} = - return () - -defCompileBasicOp _ Rotate{} = - return () - -defCompileBasicOp _ Reshape{} = - return () - -defCompileBasicOp pat e = - error $ "ImpGen.defCompileBasicOp: Invalid pattern\n " ++ - pretty pat ++ "\nfor expression\n " ++ pretty e - --- | Note: a hack to be used only for functions. -addArrays :: [ArrayDecl] -> ImpM lore r op () -addArrays = mapM_ addArray - where addArray (ArrayDecl name bt location) = - addVar name $ - ArrayVar Nothing ArrayEntry - { entryArrayLocation = location - , entryArrayElemType = bt - } - --- | Like 'dFParams', but does not create new declarations. --- Note: a hack to be used only for functions. -addFParams :: Mem lore => [FParam lore] -> ImpM lore r op () -addFParams = mapM_ addFParam - where addFParam fparam = - addVar (paramName fparam) $ - memBoundToVarEntry Nothing $ noUniquenessReturns $ paramDec fparam - --- | Another hack. -addLoopVar :: VName -> IntType -> ImpM lore r op () -addLoopVar i it = addVar i $ ScalarVar Nothing $ ScalarEntry $ IntType it - -dVars :: Mem lore => - Maybe (Exp lore) -> [PatElem lore] -> ImpM lore r op () -dVars e = mapM_ dVar - where dVar = dScope e . scopeOfPatElem - -dFParams :: Mem lore => [FParam lore] -> ImpM lore r op () -dFParams = dScope Nothing . scopeOfFParams - -dLParams :: Mem lore => [LParam lore] -> ImpM lore r op () -dLParams = dScope Nothing . scopeOfLParams - -dPrimVol_ :: VName -> PrimType -> ImpM lore r op () -dPrimVol_ name t = do - emit $ Imp.DeclareScalar name Imp.Volatile t - addVar name $ ScalarVar Nothing $ ScalarEntry t - -dPrim_ :: VName -> PrimType -> ImpM lore r op () -dPrim_ name t = do - emit $ Imp.DeclareScalar name Imp.Nonvolatile t - addVar name $ ScalarVar Nothing $ ScalarEntry t - -dPrim :: String -> PrimType -> ImpM lore r op VName -dPrim name t = do name' <- newVName name - dPrim_ name' t - return name' - -dPrimV_ :: VName -> Imp.Exp -> ImpM lore r op () -dPrimV_ name e = do dPrim_ name $ primExpType e - name <-- e - -dPrimV :: String -> Imp.Exp -> ImpM lore r op VName -dPrimV name e = do name' <- dPrim name $ primExpType e - name' <-- e - return name' - -dPrimVE :: String -> Imp.Exp -> ImpM lore r op Imp.Exp -dPrimVE name e = do name' <- dPrim name $ primExpType e - name' <-- e - return $ Imp.var name' $ primExpType e - -memBoundToVarEntry :: Maybe (Exp lore) -> MemBound NoUniqueness - -> VarEntry lore -memBoundToVarEntry e (MemPrim bt) = - ScalarVar e ScalarEntry { entryScalarType = bt } -memBoundToVarEntry e (MemMem space) = - MemVar e $ MemEntry space -memBoundToVarEntry e (MemArray bt shape _ (ArrayIn mem ixfun)) = - let location = MemLocation mem (shapeDims shape) $ fmap (toExp' int32) ixfun - in ArrayVar e ArrayEntry { entryArrayLocation = location - , entryArrayElemType = bt - } - -infoDec :: Mem lore => - NameInfo lore - -> MemInfo SubExp NoUniqueness MemBind -infoDec (LetName dec) = dec -infoDec (FParamName dec) = noUniquenessReturns dec -infoDec (LParamName dec) = dec -infoDec (IndexName it) = MemPrim $ IntType it - -dInfo :: Mem lore => - Maybe (Exp lore) -> VName -> NameInfo lore - -> ImpM lore r op () -dInfo e name info = do - let entry = memBoundToVarEntry e $ infoDec info - case entry of - MemVar _ entry' -> - emit $ Imp.DeclareMem name $ entryMemSpace entry' - ScalarVar _ entry' -> - emit $ Imp.DeclareScalar name Imp.Nonvolatile $ entryScalarType entry' - ArrayVar _ _ -> - return () - addVar name entry - -dScope :: Mem lore => - Maybe (Exp lore) -> Scope lore -> ImpM lore r op () -dScope e = mapM_ (uncurry $ dInfo e) . M.toList - -dArray :: VName -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op () -dArray name bt shape membind = - addVar name $ - memBoundToVarEntry Nothing $ MemArray bt shape NoUniqueness membind - -everythingVolatile :: ImpM lore r op a -> ImpM lore r op a -everythingVolatile = local $ \env -> env { envVolatility = Imp.Volatile } - --- | Remove the array targets. -funcallTargets :: Destination -> ImpM lore r op [VName] -funcallTargets (Destination _ dests) = - concat <$> mapM funcallTarget dests - where funcallTarget (ScalarDestination name) = - return [name] - funcallTarget (ArrayDestination _) = - return [] - funcallTarget (MemoryDestination name) = - return [name] - --- | Compile things to 'Imp.Exp'. -class ToExp a where - -- | Compile to an 'Imp.Exp', where the type (must must still be a - -- primitive) is deduced monadically. - toExp :: a -> ImpM lore r op Imp.Exp - -- | Compile where we know the type in advance. - toExp' :: PrimType -> a -> Imp.Exp - -instance ToExp SubExp where - toExp (Constant v) = - return $ Imp.ValueExp v - toExp (Var v) = - lookupVar v >>= \case - ScalarVar _ (ScalarEntry pt) -> - return $ Imp.var v pt - _ -> error $ "toExp SubExp: SubExp is not a primitive type: " ++ pretty v - - toExp' _ (Constant v) = Imp.ValueExp v - toExp' t (Var v) = Imp.var v t - -instance ToExp (PrimExp VName) where - toExp = pure . fmap Imp.ScalarVar - toExp' _ = fmap Imp.ScalarVar - -addVar :: VName -> VarEntry lore -> ImpM lore r op () -addVar name entry = - modify $ \s -> s { stateVTable = M.insert name entry $ stateVTable s } - -localDefaultSpace :: Imp.Space -> ImpM lore r op a -> ImpM lore r op a -localDefaultSpace space = local (\env -> env { envDefaultSpace = space }) - -askFunction :: ImpM lore r op (Maybe Name) -askFunction = asks envFunction - --- | Generate a 'VName', prefixed with 'askFunction' if it exists. -newVNameForFun :: String -> ImpM lore r op VName -newVNameForFun s = do - fname <- fmap nameToString <$> askFunction - newVName $ maybe "" (++".") fname ++ s - --- | Generate a 'Name', prefixed with 'askFunction' if it exists. -nameForFun :: String -> ImpM lore r op Name -nameForFun s = do - fname <- askFunction - return $ maybe "" (<>".") fname <> nameFromString s - -askEnv :: ImpM lore r op r -askEnv = asks envEnv - -localEnv :: (r -> r) -> ImpM lore r op a -> ImpM lore r op a -localEnv f = local $ \env -> env { envEnv = f $ envEnv env } - --- | The active attributes, including those for the statement --- currently being compiled. -askAttrs :: ImpM lore r op Attrs -askAttrs = asks envAttrs - --- | Add more attributes to what is returning by 'askAttrs'. -localAttrs :: Attrs -> ImpM lore r op a -> ImpM lore r op a -localAttrs attrs = local $ \env -> env { envAttrs = attrs <> envAttrs env } - -localOps :: Operations lore r op -> ImpM lore r op a -> ImpM lore r op a -localOps ops = local $ \env -> - env { envExpCompiler = opsExpCompiler ops - , envStmsCompiler = opsStmsCompiler ops - , envCopyCompiler = opsCopyCompiler ops - , envOpCompiler = opsOpCompiler ops - , envAllocCompilers = opsAllocCompilers ops - } - --- | Get the current symbol table. -getVTable :: ImpM lore r op (VTable lore) -getVTable = gets stateVTable - -putVTable :: VTable lore -> ImpM lore r op () -putVTable vtable = modify $ \s -> s { stateVTable = vtable } - --- | Run an action with a modified symbol table. All changes to the --- symbol table will be reverted once the action is done! -localVTable :: (VTable lore -> VTable lore) -> ImpM lore r op a -> ImpM lore r op a -localVTable f m = do - old_vtable <- getVTable - putVTable $ f old_vtable - a <- m - putVTable old_vtable - return a - -lookupVar :: VName -> ImpM lore r op (VarEntry lore) -lookupVar name = do - res <- gets $ M.lookup name . stateVTable - case res of - Just entry -> return entry - _ -> error $ "Unknown variable: " ++ pretty name - -lookupArray :: VName -> ImpM lore r op ArrayEntry -lookupArray name = do - res <- lookupVar name - case res of - ArrayVar _ entry -> return entry - _ -> error $ "ImpGen.lookupArray: not an array: " ++ pretty name - -lookupMemory :: VName -> ImpM lore r op MemEntry -lookupMemory name = do - res <- lookupVar name - case res of - MemVar _ entry -> return entry - _ -> error $ "Unknown memory block: " ++ pretty name - -destinationFromPattern :: Mem lore => Pattern lore -> ImpM lore r op Destination -destinationFromPattern pat = - fmap (Destination (baseTag <$> maybeHead (patternNames pat))) . mapM inspect $ - patternElements pat - where inspect patElem = do - let name = patElemName patElem - entry <- lookupVar name - case entry of - ArrayVar _ (ArrayEntry MemLocation{} _) -> - return $ ArrayDestination Nothing - MemVar{} -> - return $ MemoryDestination name - - ScalarVar{} -> - return $ ScalarDestination name - -fullyIndexArray :: VName -> [Imp.Exp] - -> ImpM lore r op (VName, Imp.Space, Count Elements Imp.Exp) -fullyIndexArray name indices = do - arr <- lookupArray name - fullyIndexArray' (entryArrayLocation arr) indices - -fullyIndexArray' :: MemLocation -> [Imp.Exp] - -> ImpM lore r op (VName, Imp.Space, Count Elements Imp.Exp) -fullyIndexArray' (MemLocation mem _ ixfun) indices = do - space <- entryMemSpace <$> lookupMemory mem - let indices' = case space of - ScalarSpace ds _ -> - let (zero_is, is) = splitFromEnd (length ds) indices - in map (const 0) zero_is ++ is - _ -> indices - return (mem, space, - elements $ IxFun.index ixfun indices') - --- More complicated read/write operations that use index functions. - -copy :: CopyCompiler lore r op -copy bt dest destslice src srcslice = do - cc <- asks envCopyCompiler - cc bt dest destslice src srcslice - --- | Use an 'Imp.Copy' if possible, otherwise 'copyElementWise'. -defaultCopy :: CopyCompiler lore r op -defaultCopy bt dest destslice src srcslice - | Just destoffset <- - IxFun.linearWithOffset (IxFun.slice destIxFun destslice) bt_size, - Just srcoffset <- - IxFun.linearWithOffset (IxFun.slice srcIxFun srcslice) bt_size = do - srcspace <- entryMemSpace <$> lookupMemory srcmem - destspace <- entryMemSpace <$> lookupMemory destmem - if isScalarSpace srcspace || isScalarSpace destspace - then copyElementWise bt dest destslice src srcslice - else emit $ Imp.Copy - destmem (bytes destoffset) destspace - srcmem (bytes srcoffset) srcspace $ - num_elems `withElemType` bt - | otherwise = - copyElementWise bt dest destslice src srcslice - where bt_size = primByteSize bt - num_elems = Imp.elements $ product $ sliceDims srcslice - MemLocation destmem _ destIxFun = dest - MemLocation srcmem _ srcIxFun = src - isScalarSpace ScalarSpace{} = True - isScalarSpace _ = False - -copyElementWise :: CopyCompiler lore r op -copyElementWise bt dest destslice src srcslice = do - let bounds = sliceDims srcslice - is <- replicateM (length bounds) (newVName "i") - let ivars = map Imp.vi32 is - (destmem, destspace, destidx) <- - fullyIndexArray' dest $ fixSlice destslice ivars - (srcmem, srcspace, srcidx) <- - fullyIndexArray' src $ fixSlice srcslice ivars - vol <- asks envVolatility - emit $ foldl (.) id (zipWith (`Imp.For` Int32) is bounds) $ - Imp.Write destmem destidx bt destspace vol $ - Imp.index srcmem srcidx bt srcspace vol - --- | Copy from here to there; both destination and source may be --- indexeded. -copyArrayDWIM :: PrimType - -> MemLocation -> [DimIndex Imp.Exp] - -> MemLocation -> [DimIndex Imp.Exp] - -> ImpM lore r op (Imp.Code op) -copyArrayDWIM bt - destlocation@(MemLocation _ destshape _) destslice - srclocation@(MemLocation _ srcshape _) srcslice - - | Just destis <- mapM dimFix destslice, - Just srcis <- mapM dimFix srcslice, - length srcis == length srcshape, - length destis == length destshape = do - (targetmem, destspace, targetoffset) <- - fullyIndexArray' destlocation destis - (srcmem, srcspace, srcoffset) <- - fullyIndexArray' srclocation srcis - vol <- asks envVolatility - return $ Imp.Write targetmem targetoffset bt destspace vol $ - Imp.index srcmem srcoffset bt srcspace vol - - | otherwise = do - let destslice' = - fullSliceNum (map (toExp' int32) destshape) destslice - srcslice' = - fullSliceNum (map (toExp' int32) srcshape) srcslice - destrank = length $ sliceDims destslice' - srcrank = length $ sliceDims srcslice' - if destrank /= srcrank - then error $ "copyArrayDWIM: cannot copy to " ++ - pretty (memLocationName destlocation) ++ - " from " ++ pretty (memLocationName srclocation) ++ - " because ranks do not match (" ++ pretty destrank ++ - " vs " ++ pretty srcrank ++ ")" - else if destlocation == srclocation && destslice' == srcslice' - then return mempty -- Copy would be no-op. - else collect $ copy bt destlocation destslice' srclocation srcslice' - --- | Like 'copyDWIM', but the target is a 'ValueDestination' --- instead of a variable name. -copyDWIMDest :: ValueDestination -> [DimIndex Imp.Exp] -> SubExp -> [DimIndex Imp.Exp] - -> ImpM lore r op () - -copyDWIMDest _ _ (Constant v) (_:_) = - error $ - unwords ["copyDWIMDest: constant source", pretty v, "cannot be indexed."] -copyDWIMDest pat dest_slice (Constant v) [] = - case mapM dimFix dest_slice of - Nothing -> - error $ - unwords ["copyDWIMDest: constant source", pretty v, "with slice destination."] - Just dest_is -> - case pat of - ScalarDestination name -> - emit $ Imp.SetScalar name $ Imp.ValueExp v - MemoryDestination{} -> - error $ - unwords ["copyDWIMDest: constant source", pretty v, "cannot be written to memory destination."] - ArrayDestination (Just dest_loc) -> do - (dest_mem, dest_space, dest_i) <- - fullyIndexArray' dest_loc dest_is - vol <- asks envVolatility - emit $ Imp.Write dest_mem dest_i bt dest_space vol $ Imp.ValueExp v - ArrayDestination Nothing -> - error "copyDWIMDest: ArrayDestination Nothing" - where bt = primValueType v - -copyDWIMDest dest dest_slice (Var src) src_slice = do - src_entry <- lookupVar src - case (dest, src_entry) of - (MemoryDestination mem, MemVar _ (MemEntry space)) -> - emit $ Imp.SetMem mem src space - - (MemoryDestination{}, _) -> - error $ - unwords ["copyDWIMDest: cannot write", pretty src, "to memory destination."] - - (_, MemVar{}) -> - error $ - unwords ["copyDWIMDest: source", pretty src, "is a memory block."] - - (_, ScalarVar _ (ScalarEntry _)) | not $ null src_slice -> - error $ - unwords ["copyDWIMDest: prim-typed source", pretty src, "with slice", pretty src_slice] - - (ScalarDestination name, _) | not $ null dest_slice -> - error $ - unwords ["copyDWIMDest: prim-typed target", pretty name, "with slice", pretty dest_slice] - - (ScalarDestination name, ScalarVar _ (ScalarEntry pt)) -> - emit $ Imp.SetScalar name $ Imp.var src pt - - (ScalarDestination name, ArrayVar _ arr) - | Just src_is <- mapM dimFix src_slice, - length src_slice == length (entryArrayShape arr) -> do - let bt = entryArrayElemType arr - (mem, space, i) <- - fullyIndexArray' (entryArrayLocation arr) src_is - vol <- asks envVolatility - emit $ Imp.SetScalar name $ Imp.index mem i bt space vol - | otherwise -> - error $ - unwords ["copyDWIMDest: prim-typed target", pretty name, - "and array-typed source", pretty src, - "with slice", pretty src_slice] - - (ArrayDestination (Just dest_loc), ArrayVar _ src_arr) -> do - let src_loc = entryArrayLocation src_arr - bt = entryArrayElemType src_arr - emit =<< copyArrayDWIM bt dest_loc dest_slice src_loc src_slice - - (ArrayDestination (Just dest_loc), ScalarVar _ (ScalarEntry bt)) - | Just dest_is <- mapM dimFix dest_slice -> do - (dest_mem, dest_space, dest_i) <- fullyIndexArray' dest_loc dest_is - vol <- asks envVolatility - emit $ Imp.Write dest_mem dest_i bt dest_space vol (Imp.var src bt) - | otherwise -> - error $ - unwords ["copyDWIMDest: array-typed target and prim-typed source", pretty src, - "with slice", pretty dest_slice] - - (ArrayDestination Nothing, _) -> - return () -- Nothing to do; something else set some memory - -- somewhere. - --- | Copy from here to there; both destination and source be --- indexeded. If so, they better be arrays of enough dimensions. --- This function will generally just Do What I Mean, and Do The Right --- Thing. Both destination and source must be in scope. -copyDWIM :: VName -> [DimIndex Imp.Exp] -> SubExp -> [DimIndex Imp.Exp] - -> ImpM lore r op () -copyDWIM dest dest_slice src src_slice = do - dest_entry <- lookupVar dest - let dest_target = - case dest_entry of - ScalarVar _ _ -> - ScalarDestination dest - - ArrayVar _ (ArrayEntry (MemLocation mem shape ixfun) _) -> - ArrayDestination $ Just $ MemLocation mem shape ixfun - - MemVar _ _ -> - MemoryDestination dest - copyDWIMDest dest_target dest_slice src src_slice - --- | As 'copyDWIM', but implicitly 'DimFix'es the indexes. -copyDWIMFix :: VName -> [Imp.Exp] -> SubExp -> [Imp.Exp] -> ImpM lore r op () -copyDWIMFix dest dest_is src src_is = - copyDWIM dest (map DimFix dest_is) src (map DimFix src_is) - --- | @compileAlloc pat size space@ allocates @n@ bytes of memory in @space@, --- writing the result to @dest@, which must be a single --- 'MemoryDestination', -compileAlloc :: Mem lore => - Pattern lore -> SubExp -> Space - -> ImpM lore r op () -compileAlloc (Pattern [] [mem]) e space = do - e' <- Imp.bytes <$> toExp e - allocator <- asks $ M.lookup space . envAllocCompilers - case allocator of - Nothing -> emit $ Imp.Allocate (patElemName mem) e' space - Just allocator' -> allocator' (patElemName mem) e' -compileAlloc pat _ _ = - error $ "compileAlloc: Invalid pattern: " ++ pretty pat - --- | The number of bytes needed to represent the array in a --- straightforward contiguous format, as an 'Int64' expression. -typeSize :: Type -> Count Bytes Imp.Exp -typeSize t = - Imp.bytes $ sExt Int64 (Imp.LeafExp (Imp.SizeOf $ elemType t) int32) * - product (map (sExt Int64 . toExp' int32) (arrayDims t)) - ---- Building blocks for constructing code. - -sFor' :: VName -> IntType -> Imp.Exp -> ImpM lore r op () -> ImpM lore r op () -sFor' i it bound body = do - addLoopVar i it - body' <- collect body - emit $ Imp.For i it bound body' - -sFor :: String -> Imp.Exp -> (Imp.Exp -> ImpM lore r op ()) -> ImpM lore r op () -sFor i bound body = do - i' <- newVName i - it <- case primExpType bound of - IntType it -> return it - t -> error $ "sFor: bound " ++ pretty bound ++ " is of type " ++ pretty t - addLoopVar i' it - body' <- collect $ body $ Imp.var i' $ IntType it - emit $ Imp.For i' it bound body' - -sWhile :: Imp.Exp -> ImpM lore r op () -> ImpM lore r op () -sWhile cond body = do - body' <- collect body - emit $ Imp.While cond body' - -sComment :: String -> ImpM lore r op () -> ImpM lore r op () -sComment s code = do - code' <- collect code - emit $ Imp.Comment s code' - -sIf :: Imp.Exp -> ImpM lore r op () -> ImpM lore r op () -> ImpM lore r op () -sIf cond tbranch fbranch = do - tbranch' <- collect tbranch - fbranch' <- collect fbranch - emit $ Imp.If cond tbranch' fbranch' - -sWhen :: Imp.Exp -> ImpM lore r op () -> ImpM lore r op () -sWhen cond tbranch = sIf cond tbranch (return ()) - -sUnless :: Imp.Exp -> ImpM lore r op () -> ImpM lore r op () -sUnless cond = sIf cond (return ()) - -sOp :: op -> ImpM lore r op () -sOp = emit . Imp.Op - -sDeclareMem :: String -> Space -> ImpM lore r op VName -sDeclareMem name space = do - name' <- newVName name - emit $ Imp.DeclareMem name' space - addVar name' $ MemVar Nothing $ MemEntry space - return name' - -sAlloc_ :: VName -> Count Bytes Imp.Exp -> Space -> ImpM lore r op () -sAlloc_ name' size' space = do - allocator <- asks $ M.lookup space . envAllocCompilers - case allocator of - Nothing -> emit $ Imp.Allocate name' size' space - Just allocator' -> allocator' name' size' - -sAlloc :: String -> Count Bytes Imp.Exp -> Space -> ImpM lore r op VName -sAlloc name size space = do - name' <- sDeclareMem name space - sAlloc_ name' size space - return name' - -sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op VName -sArray name bt shape membind = do - name' <- newVName name - dArray name' bt shape membind - return name' - --- | Declare an array in row-major order in the given memory block. -sArrayInMem :: String -> PrimType -> ShapeBase SubExp -> VName -> ImpM lore r op VName -sArrayInMem name pt shape mem = - sArray name pt shape $ ArrayIn mem $ - IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape - --- | Like 'sAllocArray', but permute the in-memory representation of the indices as specified. -sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM lore r op VName -sAllocArrayPerm name pt shape space perm = do - let permuted_dims = rearrangeShape perm $ shapeDims shape - mem <- sAlloc (name ++ "_mem") (typeSize (Array pt shape NoUniqueness)) space - let iota_ixfun = IxFun.iota $ map (primExpFromSubExp int32) permuted_dims - sArray name pt shape $ - ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm - --- | Uses linear/iota index function. -sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM lore r op VName -sAllocArray name pt shape space = - sAllocArrayPerm name pt shape space [0..shapeRank shape-1] - --- | Uses linear/iota index function. -sStaticArray :: String -> Space -> PrimType -> Imp.ArrayContents -> ImpM lore r op VName -sStaticArray name space pt vs = do - let num_elems = case vs of Imp.ArrayValues vs' -> length vs' - Imp.ArrayZeros n -> fromIntegral n - shape = Shape [intConst Int32 $ toInteger num_elems] - mem <- newVNameForFun $ name ++ "_mem" - emit $ Imp.DeclareArray mem space pt vs - addVar mem $ MemVar Nothing $ MemEntry space - sArray name pt shape $ ArrayIn mem $ IxFun.iota [fromIntegral num_elems] - -sWrite :: VName -> [Imp.Exp] -> PrimExp Imp.ExpLeaf -> ImpM lore r op () -sWrite arr is v = do - (mem, space, offset) <- fullyIndexArray arr is - vol <- asks envVolatility - emit $ Imp.Write mem offset (primExpType v) space vol v - -sUpdate :: VName -> Slice Imp.Exp -> SubExp -> ImpM lore r op () -sUpdate arr slice v = copyDWIM arr slice v [] - -sLoopNest :: Shape - -> ([Imp.Exp] -> ImpM lore r op ()) - -> ImpM lore r op () -sLoopNest = sLoopNest' [] . shapeDims - where sLoopNest' is [] f = f $ reverse is - sLoopNest' is (d:ds) f = do - d' <- toExp d - sFor "nest_i" d' $ \i -> sLoopNest' (i:is) ds f - --- | ASsignment. -(<--) :: VName -> Imp.Exp -> ImpM lore r op () -x <-- e = emit $ Imp.SetScalar x e -infixl 3 <-- - --- | Constructing an ad-hoc function that does not --- correspond to any of the IR functions in the input program. -function :: Name -> [Imp.Param] -> [Imp.Param] -> ImpM lore r op () - -> ImpM lore r op () -function fname outputs inputs m = local newFunction $ do - body <- collect $ do - mapM_ addParam $ outputs ++ inputs - m - emitFunction fname $ Imp.Function False outputs inputs body [] [] - where addParam (Imp.MemParam name space) = - addVar name $ MemVar Nothing $ MemEntry space - addParam (Imp.ScalarParam name bt) = - addVar name $ ScalarVar Nothing $ ScalarEntry bt - newFunction env = env { envFunction = Just fname } +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE Strict #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TypeFamilies #-} + +module Futhark.CodeGen.ImpGen + ( -- * Entry Points + compileProg, + + -- * Pluggable Compiler + OpCompiler, + ExpCompiler, + CopyCompiler, + StmsCompiler, + AllocCompiler, + Operations (..), + defaultOperations, + MemLocation (..), + MemEntry (..), + ScalarEntry (..), + + -- * Monadic Compiler Interface + ImpM, + localDefaultSpace, + askFunction, + newVNameForFun, + nameForFun, + askEnv, + localEnv, + localOps, + VTable, + getVTable, + localVTable, + subImpM, + subImpM_, + emit, + emitFunction, + hasFunction, + collect, + collect', + comment, + VarEntry (..), + ArrayEntry (..), + + -- * Lookups + lookupVar, + lookupArray, + lookupMemory, + + -- * Building Blocks + TV, + mkTV, + tvSize, + tvExp, + tvVar, + ToExp (..), + compileAlloc, + everythingVolatile, + compileBody, + compileBody', + compileLoopBody, + defCompileStms, + compileStms, + compileExp, + defCompileExp, + fullyIndexArray, + fullyIndexArray', + copy, + copyDWIM, + copyDWIMFix, + copyElementWise, + typeSize, + isMapTransposeCopy, + + -- * Constructing code. + dLParams, + dFParams, + dScope, + dArray, + dPrim, + dPrimVol, + dPrim_, + dPrimV_, + dPrimV, + dPrimVE, + sFor, + sWhile, + sComment, + sIf, + sWhen, + sUnless, + sOp, + sDeclareMem, + sAlloc, + sAlloc_, + sArray, + sArrayInMem, + sAllocArray, + sAllocArrayPerm, + sStaticArray, + sWrite, + sUpdate, + sLoopNest, + (<--), + (<~~), + function, + warn, + module Language.Futhark.Warnings, + ) +where + +import Control.Monad.Reader +import Control.Monad.State +import Control.Monad.Writer +import Control.Parallel.Strategies +import Data.Bifunctor (first) +import qualified Data.DList as DL +import Data.Either +import Data.List (find, genericLength, sortOn) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import Futhark.CodeGen.ImpCode + ( Bytes, + Count, + Elements, + bytes, + elements, + withElemType, + ) +import qualified Futhark.CodeGen.ImpCode as Imp +import Futhark.CodeGen.ImpGen.Transpose +import Futhark.Construct hiding (ToExp (..)) +import Futhark.IR.Mem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.IR.SOACS (SOACS) +import Futhark.Util +import Futhark.Util.Loc (noLoc) +import Language.Futhark.Warnings + +-- | How to compile an t'Op'. +type OpCompiler lore r op = Pattern lore -> Op lore -> ImpM lore r op () + +-- | How to compile some 'Stms'. +type StmsCompiler lore r op = Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op () + +-- | How to compile an 'Exp'. +type ExpCompiler lore r op = Pattern lore -> Exp lore -> ImpM lore r op () + +type CopyCompiler lore r op = + PrimType -> + MemLocation -> + Slice (Imp.TExp Int64) -> + MemLocation -> + Slice (Imp.TExp Int64) -> + ImpM lore r op () + +-- | An alternate way of compiling an allocation. +type AllocCompiler lore r op = VName -> Count Bytes (Imp.TExp Int64) -> ImpM lore r op () + +data Operations lore r op = Operations + { opsExpCompiler :: ExpCompiler lore r op, + opsOpCompiler :: OpCompiler lore r op, + opsStmsCompiler :: StmsCompiler lore r op, + opsCopyCompiler :: CopyCompiler lore r op, + opsAllocCompilers :: M.Map Space (AllocCompiler lore r op) + } + +-- | An operations set for which the expression compiler always +-- returns 'defCompileExp'. +defaultOperations :: + (Mem lore, FreeIn op) => + OpCompiler lore r op -> + Operations lore r op +defaultOperations opc = + Operations + { opsExpCompiler = defCompileExp, + opsOpCompiler = opc, + opsStmsCompiler = defCompileStms, + opsCopyCompiler = defaultCopy, + opsAllocCompilers = mempty + } + +-- | When an array is dared, this is where it is stored. +data MemLocation = MemLocation + { memLocationName :: VName, + memLocationShape :: [Imp.DimSize], + memLocationIxFun :: IxFun.IxFun (Imp.TExp Int64) + } + deriving (Eq, Show) + +data ArrayEntry = ArrayEntry + { entryArrayLocation :: MemLocation, + entryArrayElemType :: PrimType + } + deriving (Show) + +entryArrayShape :: ArrayEntry -> [Imp.DimSize] +entryArrayShape = memLocationShape . entryArrayLocation + +newtype MemEntry = MemEntry {entryMemSpace :: Imp.Space} + deriving (Show) + +newtype ScalarEntry = ScalarEntry + { entryScalarType :: PrimType + } + deriving (Show) + +-- | Every non-scalar variable must be associated with an entry. +data VarEntry lore + = ArrayVar (Maybe (Exp lore)) ArrayEntry + | ScalarVar (Maybe (Exp lore)) ScalarEntry + | MemVar (Maybe (Exp lore)) MemEntry + deriving (Show) + +-- | When compiling an expression, this is a description of where the +-- result should end up. The integer is a reference to the construct +-- that gave rise to this destination (for patterns, this will be the +-- tag of the first name in the pattern). This can be used to make +-- the generated code easier to relate to the original code. +data Destination = Destination + { destinationTag :: Maybe Int, + valueDestinations :: [ValueDestination] + } + deriving (Show) + +data ValueDestination + = ScalarDestination VName + | MemoryDestination VName + | -- | The 'MemLocation' is 'Just' if a copy if + -- required. If it is 'Nothing', then a + -- copy/assignment of a memory block somewhere + -- takes care of this array. + ArrayDestination (Maybe MemLocation) + deriving (Show) + +data Env lore r op = Env + { envExpCompiler :: ExpCompiler lore r op, + envStmsCompiler :: StmsCompiler lore r op, + envOpCompiler :: OpCompiler lore r op, + envCopyCompiler :: CopyCompiler lore r op, + envAllocCompilers :: M.Map Space (AllocCompiler lore r op), + envDefaultSpace :: Imp.Space, + envVolatility :: Imp.Volatility, + -- | User-extensible environment. + envEnv :: r, + -- | Name of the function we are compiling, if any. + envFunction :: Maybe Name, + -- | The set of attributes that are active on the enclosing + -- statements (including the one we are currently compiling). + envAttrs :: Attrs + } + +newEnv :: r -> Operations lore r op -> Imp.Space -> Env lore r op +newEnv r ops ds = + Env + { envExpCompiler = opsExpCompiler ops, + envStmsCompiler = opsStmsCompiler ops, + envOpCompiler = opsOpCompiler ops, + envCopyCompiler = opsCopyCompiler ops, + envAllocCompilers = mempty, + envDefaultSpace = ds, + envVolatility = Imp.Nonvolatile, + envEnv = r, + envFunction = Nothing, + envAttrs = mempty + } + +-- | The symbol table used during compilation. +type VTable lore = M.Map VName (VarEntry lore) + +data ImpState lore r op = ImpState + { stateVTable :: VTable lore, + stateFunctions :: Imp.Functions op, + stateCode :: Imp.Code op, + stateWarnings :: Warnings, + stateNameSource :: VNameSource + } + +newState :: VNameSource -> ImpState lore r op +newState = ImpState mempty mempty mempty mempty + +newtype ImpM lore r op a + = ImpM (ReaderT (Env lore r op) (State (ImpState lore r op)) a) + deriving + ( Functor, + Applicative, + Monad, + MonadState (ImpState lore r op), + MonadReader (Env lore r op) + ) + +instance MonadFreshNames (ImpM lore r op) where + getNameSource = gets stateNameSource + putNameSource src = modify $ \s -> s {stateNameSource = src} + +-- Cannot be an KernelsMem scope because the index functions have +-- the wrong leaves (VName instead of Imp.Exp). +instance HasScope SOACS (ImpM lore r op) where + askScope = gets $ M.map (LetName . entryType) . stateVTable + where + entryType (MemVar _ memEntry) = + Mem (entryMemSpace memEntry) + entryType (ArrayVar _ arrayEntry) = + Array + (entryArrayElemType arrayEntry) + (Shape $ entryArrayShape arrayEntry) + NoUniqueness + entryType (ScalarVar _ scalarEntry) = + Prim $ entryScalarType scalarEntry + +runImpM :: + ImpM lore r op a -> + r -> + Operations lore r op -> + Imp.Space -> + ImpState lore r op -> + (a, ImpState lore r op) +runImpM (ImpM m) r ops space = runState (runReaderT m $ newEnv r ops space) + +subImpM_ :: + r' -> + Operations lore r' op' -> + ImpM lore r' op' a -> + ImpM lore r op (Imp.Code op') +subImpM_ r ops m = snd <$> subImpM r ops m + +subImpM :: + r' -> + Operations lore r' op' -> + ImpM lore r' op' a -> + ImpM lore r op (a, Imp.Code op') +subImpM r ops (ImpM m) = do + env <- ask + s <- get + + let env' = + env + { envExpCompiler = opsExpCompiler ops, + envStmsCompiler = opsStmsCompiler ops, + envCopyCompiler = opsCopyCompiler ops, + envOpCompiler = opsOpCompiler ops, + envAllocCompilers = opsAllocCompilers ops, + envEnv = r + } + s' = + ImpState + { stateVTable = stateVTable s, + stateFunctions = mempty, + stateCode = mempty, + stateNameSource = stateNameSource s, + stateWarnings = mempty + } + (x, s'') = runState (runReaderT m env') s' + + putNameSource $ stateNameSource s'' + warnings $ stateWarnings s'' + return (x, stateCode s'') + +-- | Execute a code generation action, returning the code that was +-- emitted. +collect :: ImpM lore r op () -> ImpM lore r op (Imp.Code op) +collect = fmap snd . collect' + +collect' :: ImpM lore r op a -> ImpM lore r op (a, Imp.Code op) +collect' m = do + prev_code <- gets stateCode + modify $ \s -> s {stateCode = mempty} + x <- m + new_code <- gets stateCode + modify $ \s -> s {stateCode = prev_code} + return (x, new_code) + +-- | Execute a code generation action, wrapping the generated code +-- within a 'Imp.Comment' with the given description. +comment :: String -> ImpM lore r op () -> ImpM lore r op () +comment desc m = do + code <- collect m + emit $ Imp.Comment desc code + +-- | Emit some generated imperative code. +emit :: Imp.Code op -> ImpM lore r op () +emit code = modify $ \s -> s {stateCode = stateCode s <> code} + +warnings :: Warnings -> ImpM lore r op () +warnings ws = modify $ \s -> s {stateWarnings = ws <> stateWarnings s} + +-- | Emit a warning about something the user should be aware of. +warn :: Located loc => loc -> [loc] -> String -> ImpM lore r op () +warn loc locs problem = + warnings $ singleWarning' (srclocOf loc) (map srclocOf locs) problem + +-- | Emit a function in the generated code. +emitFunction :: Name -> Imp.Function op -> ImpM lore r op () +emitFunction fname fun = do + Imp.Functions fs <- gets stateFunctions + modify $ \s -> s {stateFunctions = Imp.Functions $ (fname, fun) : fs} + +-- | Check if a function of a given name exists. +hasFunction :: Name -> ImpM lore r op Bool +hasFunction fname = gets $ \s -> + let Imp.Functions fs = stateFunctions s + in isJust $ lookup fname fs + +constsVTable :: Mem lore => Stms lore -> VTable lore +constsVTable = foldMap stmVtable + where + stmVtable (Let pat _ e) = + foldMap (peVtable e) $ patternElements pat + peVtable e (PatElem name dec) = + M.singleton name $ memBoundToVarEntry (Just e) dec + +compileProg :: + (Mem lore, FreeIn op, MonadFreshNames m) => + r -> + Operations lore r op -> + Imp.Space -> + Prog lore -> + m (Warnings, Imp.Definitions op) +compileProg r ops space (Prog consts funs) = + modifyNameSource $ \src -> + let (_, ss) = + unzip $ parMap rpar (compileFunDef' src) funs + free_in_funs = + freeIn $ mconcat $ map stateFunctions ss + (consts', s') = + runImpM (compileConsts free_in_funs consts) r ops space $ + combineStates ss + in ( ( stateWarnings s', + Imp.Definitions consts' (stateFunctions s') + ), + stateNameSource s' + ) + where + compileFunDef' src fdef = + runImpM + (compileFunDef fdef) + r + ops + space + (newState src) {stateVTable = constsVTable consts} + + combineStates ss = + let Imp.Functions funs' = mconcat $ map stateFunctions ss + src = mconcat (map stateNameSource ss) + in (newState src) + { stateFunctions = + Imp.Functions $ M.toList $ M.fromList funs', + stateWarnings = + mconcat $ map stateWarnings ss + } + +compileConsts :: Names -> Stms lore -> ImpM lore r op (Imp.Constants op) +compileConsts used_consts stms = do + code <- collect $ compileStms used_consts stms $ pure () + pure $ uncurry Imp.Constants $ first DL.toList $ extract code + where + -- Fish out those top-level declarations in the constant + -- initialisation code that are free in the functions. + extract (x Imp.:>>: y) = + extract x <> extract y + extract (Imp.DeclareMem name space) + | name `nameIn` used_consts = + ( DL.singleton $ Imp.MemParam name space, + mempty + ) + extract (Imp.DeclareScalar name _ t) + | name `nameIn` used_consts = + ( DL.singleton $ Imp.ScalarParam name t, + mempty + ) + extract s = + (mempty, s) + +compileInParam :: + Mem lore => + FParam lore -> + ImpM lore r op (Either Imp.Param ArrayDecl) +compileInParam fparam = case paramDec fparam of + MemPrim bt -> + return $ Left $ Imp.ScalarParam name bt + MemMem space -> + return $ Left $ Imp.MemParam name space + MemArray bt shape _ (ArrayIn mem ixfun) -> + return $ + Right $ + ArrayDecl name bt $ + MemLocation mem (shapeDims shape) $ fmap (fmap Imp.ScalarVar) ixfun + where + name = paramName fparam + +data ArrayDecl = ArrayDecl VName PrimType MemLocation + +fparamSizes :: Typed dec => Param dec -> S.Set VName +fparamSizes = S.fromList . subExpVars . arrayDims . paramType + +compileInParams :: + Mem lore => + [FParam lore] -> + [EntryPointType] -> + ImpM lore r op ([Imp.Param], [ArrayDecl], [Imp.ExternalValue]) +compileInParams params orig_epts = do + let (ctx_params, val_params) = + splitAt (length params - sum (map entryPointSize orig_epts)) params + (inparams, arrayds) <- partitionEithers <$> mapM compileInParam (ctx_params ++ val_params) + let findArray x = find (isArrayDecl x) arrayds + sizes = mconcat $ map fparamSizes $ ctx_params ++ val_params + + summaries = M.fromList $ mapMaybe memSummary params + where + memSummary param + | MemMem space <- paramDec param = + Just (paramName param, space) + | otherwise = + Nothing + + findMemInfo :: VName -> Maybe Space + findMemInfo = flip M.lookup summaries + + mkValueDesc fparam signedness = + case (findArray $ paramName fparam, paramType fparam) of + (Just (ArrayDecl _ bt (MemLocation mem shape _)), _) -> do + memspace <- findMemInfo mem + Just $ Imp.ArrayValue mem memspace bt signedness shape + (_, Prim bt) + | paramName fparam `S.member` sizes -> + Nothing + | otherwise -> + Just $ Imp.ScalarValue bt signedness $ paramName fparam + _ -> + Nothing + + mkExts (TypeOpaque desc n : epts) fparams = + let (fparams', rest) = splitAt n fparams + in Imp.OpaqueValue + desc + (mapMaybe (`mkValueDesc` Imp.TypeDirect) fparams') : + mkExts epts rest + mkExts (TypeUnsigned : epts) (fparam : fparams) = + maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeUnsigned) + ++ mkExts epts fparams + mkExts (TypeDirect : epts) (fparam : fparams) = + maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeDirect) + ++ mkExts epts fparams + mkExts _ _ = [] + + return (inparams, arrayds, mkExts orig_epts val_params) + where + isArrayDecl x (ArrayDecl y _ _) = x == y + +compileOutParams :: + Mem lore => + [RetType lore] -> + [EntryPointType] -> + ImpM lore r op ([Imp.ExternalValue], [Imp.Param], Destination) +compileOutParams orig_rts orig_epts = do + ((extvs, dests), (outparams, ctx_dests)) <- + runWriterT $ evalStateT (mkExts orig_epts orig_rts) (M.empty, M.empty) + let ctx_dests' = map snd $ sortOn fst $ M.toList ctx_dests + return (extvs, outparams, Destination Nothing $ ctx_dests' <> dests) + where + imp = lift . lift + + mkExts (TypeOpaque desc n : epts) rts = do + let (rts', rest) = splitAt n rts + (evs, dests) <- unzip <$> zipWithM mkParam rts' (repeat Imp.TypeDirect) + (more_values, more_dests) <- mkExts epts rest + return + ( Imp.OpaqueValue desc evs : more_values, + dests ++ more_dests + ) + mkExts (TypeUnsigned : epts) (rt : rts) = do + (ev, dest) <- mkParam rt Imp.TypeUnsigned + (more_values, more_dests) <- mkExts epts rts + return + ( Imp.TransparentValue ev : more_values, + dest : more_dests + ) + mkExts (TypeDirect : epts) (rt : rts) = do + (ev, dest) <- mkParam rt Imp.TypeDirect + (more_values, more_dests) <- mkExts epts rts + return + ( Imp.TransparentValue ev : more_values, + dest : more_dests + ) + mkExts _ _ = return ([], []) + + mkParam MemMem {} _ = + error "Functions may not explicitly return memory blocks." + mkParam (MemPrim t) ept = do + out <- imp $ newVName "scalar_out" + tell ([Imp.ScalarParam out t], mempty) + return (Imp.ScalarValue t ept out, ScalarDestination out) + mkParam (MemArray t shape _ dec) ept = do + space <- asks envDefaultSpace + memout <- case dec of + ReturnsNewBlock _ x _ixfun -> do + memout <- imp $ newVName "out_mem" + tell + ( [Imp.MemParam memout space], + M.singleton x $ MemoryDestination memout + ) + return memout + ReturnsInBlock memout _ -> + return memout + resultshape <- mapM inspectExtSize $ shapeDims shape + return + ( Imp.ArrayValue memout space t ept resultshape, + ArrayDestination Nothing + ) + + inspectExtSize (Ext x) = do + (memseen, arrseen) <- get + case M.lookup x arrseen of + Nothing -> do + out <- imp $ newVName "out_arrsize" + tell + ( [Imp.ScalarParam out int64], + M.singleton x $ ScalarDestination out + ) + put (memseen, M.insert x out arrseen) + return $ Var out + Just out -> + return $ Var out + inspectExtSize (Free se) = + return se + +compileFunDef :: + Mem lore => + FunDef lore -> + ImpM lore r op () +compileFunDef (FunDef entry _ fname rettype params body) = + local (\env -> env {envFunction = Just fname}) $ do + ((outparams, inparams, results, args), body') <- collect' compile + emitFunction fname $ Imp.Function (isJust entry) outparams inparams body' results args + where + params_entry = maybe (replicate (length params) TypeDirect) fst entry + ret_entry = maybe (replicate (length rettype) TypeDirect) snd entry + compile = do + (inparams, arrayds, args) <- compileInParams params params_entry + (results, outparams, Destination _ dests) <- compileOutParams rettype ret_entry + addFParams params + addArrays arrayds + + let Body _ stms ses = body + compileStms (freeIn ses) stms $ + forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se [] + + return (outparams, inparams, results, args) + +compileBody :: (Mem lore) => Pattern lore -> Body lore -> ImpM lore r op () +compileBody pat (Body _ bnds ses) = do + Destination _ dests <- destinationFromPattern pat + compileStms (freeIn ses) bnds $ + forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se [] + +compileBody' :: [Param dec] -> Body lore -> ImpM lore r op () +compileBody' params (Body _ bnds ses) = + compileStms (freeIn ses) bnds $ + forM_ (zip params ses) $ \(param, se) -> copyDWIM (paramName param) [] se [] + +compileLoopBody :: Typed dec => [Param dec] -> Body lore -> ImpM lore r op () +compileLoopBody mergeparams (Body _ bnds ses) = do + -- We cannot write the results to the merge parameters immediately, + -- as some of the results may actually *be* merge parameters, and + -- would thus be clobbered. Therefore, we first copy to new + -- variables mirroring the merge parameters, and then copy this + -- buffer to the merge parameters. This is efficient, because the + -- operations are all scalar operations. + tmpnames <- mapM (newVName . (++ "_tmp") . baseString . paramName) mergeparams + compileStms (freeIn ses) bnds $ do + copy_to_merge_params <- forM (zip3 mergeparams tmpnames ses) $ \(p, tmp, se) -> + case typeOf p of + Prim pt -> do + emit $ Imp.DeclareScalar tmp Imp.Nonvolatile pt + emit $ Imp.SetScalar tmp $ toExp' pt se + return $ emit $ Imp.SetScalar (paramName p) $ Imp.var tmp pt + Mem space | Var v <- se -> do + emit $ Imp.DeclareMem tmp space + emit $ Imp.SetMem tmp v space + return $ emit $ Imp.SetMem (paramName p) tmp space + _ -> return $ return () + sequence_ copy_to_merge_params + +compileStms :: Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op () +compileStms alive_after_stms all_stms m = do + cb <- asks envStmsCompiler + cb alive_after_stms all_stms m + +defCompileStms :: + (Mem lore, FreeIn op) => + Names -> + Stms lore -> + ImpM lore r op () -> + ImpM lore r op () +defCompileStms alive_after_stms all_stms m = + -- We keep track of any memory blocks produced by the statements, + -- and after the last time that memory block is used, we insert a + -- Free. This is very conservative, but can cut down on lifetimes + -- in some cases. + void $ compileStms' mempty $ stmsToList all_stms + where + compileStms' allocs (Let pat aux e : bs) = do + dVars (Just e) (patternElements pat) + + e_code <- + localAttrs (stmAuxAttrs aux) $ + collect $ compileExp pat e + (live_after, bs_code) <- collect' $ compileStms' (patternAllocs pat <> allocs) bs + let dies_here v = + not (v `nameIn` live_after) + && v `nameIn` freeIn e_code + to_free = S.filter (dies_here . fst) allocs + + emit e_code + mapM_ (emit . uncurry Imp.Free) to_free + emit bs_code + + return $ freeIn e_code <> live_after + compileStms' _ [] = do + code <- collect m + emit code + return $ freeIn code <> alive_after_stms + + patternAllocs = S.fromList . mapMaybe isMemPatElem . patternElements + isMemPatElem pe = case patElemType pe of + Mem space -> Just (patElemName pe, space) + _ -> Nothing + +compileExp :: Pattern lore -> Exp lore -> ImpM lore r op () +compileExp pat e = do + ec <- asks envExpCompiler + ec pat e + +defCompileExp :: + (Mem lore) => + Pattern lore -> + Exp lore -> + ImpM lore r op () +defCompileExp pat (If cond tbranch fbranch _) = + sIf (toBoolExp cond) (compileBody pat tbranch) (compileBody pat fbranch) +defCompileExp pat (Apply fname args _ _) = do + dest <- destinationFromPattern pat + targets <- funcallTargets dest + args' <- catMaybes <$> mapM compileArg args + emit $ Imp.Call targets fname args' + where + compileArg (se, _) = do + t <- subExpType se + case (se, t) of + (_, Prim pt) -> return $ Just $ Imp.ExpArg $ toExp' pt se + (Var v, Mem {}) -> return $ Just $ Imp.MemArg v + _ -> return Nothing +defCompileExp pat (BasicOp op) = defCompileBasicOp pat op +defCompileExp pat (DoLoop ctx val form body) = do + attrs <- askAttrs + when ("unroll" `inAttrs` attrs) $ + warn (noLoc :: SrcLoc) [] "#[unroll] on loop with unknown number of iterations." -- FIXME: no location. + dFParams mergepat + forM_ merge $ \(p, se) -> + when ((== 0) $ arrayRank $ paramType p) $ + copyDWIM (paramName p) [] se [] + + let doBody = compileLoopBody mergepat body + + case form of + ForLoop i _ bound loopvars -> do + let setLoopParam (p, a) + | Prim _ <- paramType p = + copyDWIM (paramName p) [] (Var a) [DimFix $ Imp.vi64 i] + | otherwise = + return () + + bound' <- toExp bound + + dLParams $ map fst loopvars + sFor' i bound' $ + mapM_ setLoopParam loopvars >> doBody + WhileLoop cond -> + sWhile (TPrimExp $ Imp.var cond Bool) doBody + + Destination _ pat_dests <- destinationFromPattern pat + forM_ (zip pat_dests $ map (Var . paramName . fst) merge) $ \(d, r) -> + copyDWIMDest d [] r [] + where + merge = ctx ++ val + mergepat = map fst merge +defCompileExp pat (Op op) = do + opc <- asks envOpCompiler + opc pat op + +defCompileBasicOp :: + Mem lore => + Pattern lore -> + BasicOp -> + ImpM lore r op () +defCompileBasicOp (Pattern _ [pe]) (SubExp se) = + copyDWIM (patElemName pe) [] se [] +defCompileBasicOp (Pattern _ [pe]) (Opaque se) = + copyDWIM (patElemName pe) [] se [] +defCompileBasicOp (Pattern _ [pe]) (UnOp op e) = do + e' <- toExp e + patElemName pe <~~ Imp.UnOpExp op e' +defCompileBasicOp (Pattern _ [pe]) (ConvOp conv e) = do + e' <- toExp e + patElemName pe <~~ Imp.ConvOpExp conv e' +defCompileBasicOp (Pattern _ [pe]) (BinOp bop x y) = do + x' <- toExp x + y' <- toExp y + patElemName pe <~~ Imp.BinOpExp bop x' y' +defCompileBasicOp (Pattern _ [pe]) (CmpOp bop x y) = do + x' <- toExp x + y' <- toExp y + patElemName pe <~~ Imp.CmpOpExp bop x' y' +defCompileBasicOp _ (Assert e msg loc) = do + e' <- toExp e + msg' <- traverse toExp msg + emit $ Imp.Assert e' msg' loc + + attrs <- askAttrs + when (AttrComp "warn" ["safety_checks"] `inAttrs` attrs) $ + uncurry warn loc "Safety check required at run-time." +defCompileBasicOp (Pattern _ [pe]) (Index src slice) + | Just idxs <- sliceIndices slice = + copyDWIM (patElemName pe) [] (Var src) $ map (DimFix . toInt64Exp) idxs +defCompileBasicOp _ Index {} = + return () +defCompileBasicOp (Pattern _ [pe]) (Update _ slice se) = + sUpdate (patElemName pe) (map (fmap toInt64Exp) slice) se +defCompileBasicOp (Pattern _ [pe]) (Replicate (Shape ds) se) = do + ds' <- mapM toExp ds + is <- replicateM (length ds) (newVName "i") + copy_elem <- collect $ copyDWIM (patElemName pe) (map (DimFix . Imp.vi64) is) se [] + emit $ foldl (.) id (zipWith Imp.For is ds') copy_elem +defCompileBasicOp _ Scratch {} = + return () +defCompileBasicOp (Pattern [] [pe]) (Iota n e s it) = do + e' <- toExp e + s' <- toExp s + sFor "i" (toInt64Exp n) $ \i -> do + let i' = sExt it $ untyped i + x <- + dPrimV "x" $ + TPrimExp $ + BinOpExp (Add it OverflowUndef) e' $ + BinOpExp (Mul it OverflowUndef) i' s' + copyDWIM (patElemName pe) [DimFix i] (Var (tvVar x)) [] +defCompileBasicOp (Pattern _ [pe]) (Copy src) = + copyDWIM (patElemName pe) [] (Var src) [] +defCompileBasicOp (Pattern _ [pe]) (Manifest _ src) = + copyDWIM (patElemName pe) [] (Var src) [] +defCompileBasicOp (Pattern _ [pe]) (Concat i x ys _) = do + offs_glb <- dPrimV "tmp_offs" 0 + + forM_ (x : ys) $ \y -> do + y_dims <- arrayDims <$> lookupType y + let rows = case drop i y_dims of + [] -> error $ "defCompileBasicOp Concat: empty array shape for " ++ pretty y + r : _ -> toInt64Exp r + skip_dims = take i y_dims + sliceAllDim d = DimSlice 0 d 1 + skip_slices = map (sliceAllDim . toInt64Exp) skip_dims + destslice = skip_slices ++ [DimSlice (tvExp offs_glb) rows 1] + copyDWIM (patElemName pe) destslice (Var y) [] + offs_glb <-- tvExp offs_glb + rows +defCompileBasicOp (Pattern [] [pe]) (ArrayLit es _) + | Just vs@(v : _) <- mapM isLiteral es = do + dest_mem <- entryArrayLocation <$> lookupArray (patElemName pe) + dest_space <- entryMemSpace <$> lookupMemory (memLocationName dest_mem) + let t = primValueType v + static_array <- newVNameForFun "static_array" + emit $ Imp.DeclareArray static_array dest_space t $ Imp.ArrayValues vs + let static_src = + MemLocation static_array [intConst Int64 $ fromIntegral $ length es] $ + IxFun.iota [fromIntegral $ length es] + entry = MemVar Nothing $ MemEntry dest_space + addVar static_array entry + let slice = [DimSlice 0 (genericLength es) 1] + copy t dest_mem slice static_src slice + | otherwise = + forM_ (zip [0 ..] es) $ \(i, e) -> + copyDWIM (patElemName pe) [DimFix $ fromInteger i] e [] + where + isLiteral (Constant v) = Just v + isLiteral _ = Nothing +defCompileBasicOp _ Rearrange {} = + return () +defCompileBasicOp _ Rotate {} = + return () +defCompileBasicOp _ Reshape {} = + return () +defCompileBasicOp pat e = + error $ + "ImpGen.defCompileBasicOp: Invalid pattern\n " + ++ pretty pat + ++ "\nfor expression\n " + ++ pretty e + +-- | Note: a hack to be used only for functions. +addArrays :: [ArrayDecl] -> ImpM lore r op () +addArrays = mapM_ addArray + where + addArray (ArrayDecl name bt location) = + addVar name $ + ArrayVar + Nothing + ArrayEntry + { entryArrayLocation = location, + entryArrayElemType = bt + } + +-- | Like 'dFParams', but does not create new declarations. +-- Note: a hack to be used only for functions. +addFParams :: Mem lore => [FParam lore] -> ImpM lore r op () +addFParams = mapM_ addFParam + where + addFParam fparam = + addVar (paramName fparam) $ + memBoundToVarEntry Nothing $ noUniquenessReturns $ paramDec fparam + +-- | Another hack. +addLoopVar :: VName -> IntType -> ImpM lore r op () +addLoopVar i it = addVar i $ ScalarVar Nothing $ ScalarEntry $ IntType it + +dVars :: + Mem lore => + Maybe (Exp lore) -> + [PatElem lore] -> + ImpM lore r op () +dVars e = mapM_ dVar + where + dVar = dScope e . scopeOfPatElem + +dFParams :: Mem lore => [FParam lore] -> ImpM lore r op () +dFParams = dScope Nothing . scopeOfFParams + +dLParams :: Mem lore => [LParam lore] -> ImpM lore r op () +dLParams = dScope Nothing . scopeOfLParams + +dPrimVol :: String -> PrimType -> Imp.TExp t -> ImpM lore r op (TV t) +dPrimVol name t e = do + name' <- newVName name + emit $ Imp.DeclareScalar name' Imp.Volatile t + addVar name' $ ScalarVar Nothing $ ScalarEntry t + name' <~~ untyped e + return $ TV name' t + +dPrim_ :: VName -> PrimType -> ImpM lore r op () +dPrim_ name t = do + emit $ Imp.DeclareScalar name Imp.Nonvolatile t + addVar name $ ScalarVar Nothing $ ScalarEntry t + +-- | The return type is polymorphic, so there is no guarantee it +-- actually matches the 'PrimType', but at least we have to use it +-- consistently. +dPrim :: String -> PrimType -> ImpM lore r op (TV t) +dPrim name t = do + name' <- newVName name + dPrim_ name' t + return $ TV name' t + +dPrimV_ :: VName -> Imp.TExp t -> ImpM lore r op () +dPrimV_ name e = do + dPrim_ name t + TV name t <-- e + where + t = primExpType $ untyped e + +dPrimV :: String -> Imp.TExp t -> ImpM lore r op (TV t) +dPrimV name e = do + name' <- dPrim name $ primExpType $ untyped e + name' <-- e + return name' + +dPrimVE :: String -> Imp.TExp t -> ImpM lore r op (Imp.TExp t) +dPrimVE name e = do + name' <- dPrim name $ primExpType $ untyped e + name' <-- e + return $ tvExp name' + +memBoundToVarEntry :: + Maybe (Exp lore) -> + MemBound NoUniqueness -> + VarEntry lore +memBoundToVarEntry e (MemPrim bt) = + ScalarVar e ScalarEntry {entryScalarType = bt} +memBoundToVarEntry e (MemMem space) = + MemVar e $ MemEntry space +memBoundToVarEntry e (MemArray bt shape _ (ArrayIn mem ixfun)) = + let location = MemLocation mem (shapeDims shape) $ fmap (fmap Imp.ScalarVar) ixfun + in ArrayVar + e + ArrayEntry + { entryArrayLocation = location, + entryArrayElemType = bt + } + +infoDec :: + Mem lore => + NameInfo lore -> + MemInfo SubExp NoUniqueness MemBind +infoDec (LetName dec) = dec +infoDec (FParamName dec) = noUniquenessReturns dec +infoDec (LParamName dec) = dec +infoDec (IndexName it) = MemPrim $ IntType it + +dInfo :: + Mem lore => + Maybe (Exp lore) -> + VName -> + NameInfo lore -> + ImpM lore r op () +dInfo e name info = do + let entry = memBoundToVarEntry e $ infoDec info + case entry of + MemVar _ entry' -> + emit $ Imp.DeclareMem name $ entryMemSpace entry' + ScalarVar _ entry' -> + emit $ Imp.DeclareScalar name Imp.Nonvolatile $ entryScalarType entry' + ArrayVar _ _ -> + return () + addVar name entry + +dScope :: + Mem lore => + Maybe (Exp lore) -> + Scope lore -> + ImpM lore r op () +dScope e = mapM_ (uncurry $ dInfo e) . M.toList + +dArray :: VName -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op () +dArray name bt shape membind = + addVar name $ + memBoundToVarEntry Nothing $ MemArray bt shape NoUniqueness membind + +everythingVolatile :: ImpM lore r op a -> ImpM lore r op a +everythingVolatile = local $ \env -> env {envVolatility = Imp.Volatile} + +-- | Remove the array targets. +funcallTargets :: Destination -> ImpM lore r op [VName] +funcallTargets (Destination _ dests) = + concat <$> mapM funcallTarget dests + where + funcallTarget (ScalarDestination name) = + return [name] + funcallTarget (ArrayDestination _) = + return [] + funcallTarget (MemoryDestination name) = + return [name] + +-- | A typed variable, which we can turn into a typed expression, or +-- use as the target for an assignment. This is used to aid in type +-- safety when doing code generation, by keeping the types straight. +-- It is still easy to cheat when you need to. +data TV t = TV VName PrimType + +-- | Create a typed variable from a name and a dynamic type. Note +-- that there is no guarantee that the dynamic type corresponds to the +-- inferred static type, but the latter will at least have to be used +-- consistently. +mkTV :: VName -> PrimType -> TV t +mkTV = TV + +-- | Convert a typed variable to a size (a SubExp). +tvSize :: TV t -> Imp.DimSize +tvSize = Var . tvVar + +-- | Convert a typed variable to a similarly typed expression. +tvExp :: TV t -> Imp.TExp t +tvExp (TV v t) = Imp.TPrimExp $ Imp.var v t + +-- | Extract the underlying variable name from a typed variable. +tvVar :: TV t -> VName +tvVar (TV v _) = v + +-- | Compile things to 'Imp.Exp'. +class ToExp a where + -- | Compile to an 'Imp.Exp', where the type (must must still be a + -- primitive) is deduced monadically. + toExp :: a -> ImpM lore r op Imp.Exp + + -- | Compile where we know the type in advance. + toExp' :: PrimType -> a -> Imp.Exp + + toInt32Exp :: a -> Imp.TExp Int32 + toInt32Exp = TPrimExp . toExp' int32 + + toInt64Exp :: a -> Imp.TExp Int64 + toInt64Exp = TPrimExp . toExp' int64 + + toBoolExp :: a -> Imp.TExp Bool + toBoolExp = TPrimExp . toExp' Bool + +instance ToExp SubExp where + toExp (Constant v) = + return $ Imp.ValueExp v + toExp (Var v) = + lookupVar v >>= \case + ScalarVar _ (ScalarEntry pt) -> + return $ Imp.var v pt + _ -> error $ "toExp SubExp: SubExp is not a primitive type: " ++ pretty v + + toExp' _ (Constant v) = Imp.ValueExp v + toExp' t (Var v) = Imp.var v t + +instance ToExp (PrimExp VName) where + toExp = pure . fmap Imp.ScalarVar + toExp' _ = fmap Imp.ScalarVar + +addVar :: VName -> VarEntry lore -> ImpM lore r op () +addVar name entry = + modify $ \s -> s {stateVTable = M.insert name entry $ stateVTable s} + +localDefaultSpace :: Imp.Space -> ImpM lore r op a -> ImpM lore r op a +localDefaultSpace space = local (\env -> env {envDefaultSpace = space}) + +askFunction :: ImpM lore r op (Maybe Name) +askFunction = asks envFunction + +-- | Generate a 'VName', prefixed with 'askFunction' if it exists. +newVNameForFun :: String -> ImpM lore r op VName +newVNameForFun s = do + fname <- fmap nameToString <$> askFunction + newVName $ maybe "" (++ ".") fname ++ s + +-- | Generate a 'Name', prefixed with 'askFunction' if it exists. +nameForFun :: String -> ImpM lore r op Name +nameForFun s = do + fname <- askFunction + return $ maybe "" (<> ".") fname <> nameFromString s + +askEnv :: ImpM lore r op r +askEnv = asks envEnv + +localEnv :: (r -> r) -> ImpM lore r op a -> ImpM lore r op a +localEnv f = local $ \env -> env {envEnv = f $ envEnv env} + +-- | The active attributes, including those for the statement +-- currently being compiled. +askAttrs :: ImpM lore r op Attrs +askAttrs = asks envAttrs + +-- | Add more attributes to what is returning by 'askAttrs'. +localAttrs :: Attrs -> ImpM lore r op a -> ImpM lore r op a +localAttrs attrs = local $ \env -> env {envAttrs = attrs <> envAttrs env} + +localOps :: Operations lore r op -> ImpM lore r op a -> ImpM lore r op a +localOps ops = local $ \env -> + env + { envExpCompiler = opsExpCompiler ops, + envStmsCompiler = opsStmsCompiler ops, + envCopyCompiler = opsCopyCompiler ops, + envOpCompiler = opsOpCompiler ops, + envAllocCompilers = opsAllocCompilers ops + } + +-- | Get the current symbol table. +getVTable :: ImpM lore r op (VTable lore) +getVTable = gets stateVTable + +putVTable :: VTable lore -> ImpM lore r op () +putVTable vtable = modify $ \s -> s {stateVTable = vtable} + +-- | Run an action with a modified symbol table. All changes to the +-- symbol table will be reverted once the action is done! +localVTable :: (VTable lore -> VTable lore) -> ImpM lore r op a -> ImpM lore r op a +localVTable f m = do + old_vtable <- getVTable + putVTable $ f old_vtable + a <- m + putVTable old_vtable + return a + +lookupVar :: VName -> ImpM lore r op (VarEntry lore) +lookupVar name = do + res <- gets $ M.lookup name . stateVTable + case res of + Just entry -> return entry + _ -> error $ "Unknown variable: " ++ pretty name + +lookupArray :: VName -> ImpM lore r op ArrayEntry +lookupArray name = do + res <- lookupVar name + case res of + ArrayVar _ entry -> return entry + _ -> error $ "ImpGen.lookupArray: not an array: " ++ pretty name + +lookupMemory :: VName -> ImpM lore r op MemEntry +lookupMemory name = do + res <- lookupVar name + case res of + MemVar _ entry -> return entry + _ -> error $ "Unknown memory block: " ++ pretty name + +destinationFromPattern :: Mem lore => Pattern lore -> ImpM lore r op Destination +destinationFromPattern pat = + fmap (Destination (baseTag <$> maybeHead (patternNames pat))) . mapM inspect $ + patternElements pat + where + inspect patElem = do + let name = patElemName patElem + entry <- lookupVar name + case entry of + ArrayVar _ (ArrayEntry MemLocation {} _) -> + return $ ArrayDestination Nothing + MemVar {} -> + return $ MemoryDestination name + ScalarVar {} -> + return $ ScalarDestination name + +fullyIndexArray :: + VName -> + [Imp.TExp Int64] -> + ImpM lore r op (VName, Imp.Space, Count Elements (Imp.TExp Int64)) +fullyIndexArray name indices = do + arr <- lookupArray name + fullyIndexArray' (entryArrayLocation arr) indices + +fullyIndexArray' :: + MemLocation -> + [Imp.TExp Int64] -> + ImpM lore r op (VName, Imp.Space, Count Elements (Imp.TExp Int64)) +fullyIndexArray' (MemLocation mem _ ixfun) indices = do + space <- entryMemSpace <$> lookupMemory mem + let indices' = case space of + ScalarSpace ds _ -> + let (zero_is, is) = splitFromEnd (length ds) indices + in map (const 0) zero_is ++ is + _ -> indices + return + ( mem, + space, + elements $ IxFun.index ixfun indices' + ) + +-- More complicated read/write operations that use index functions. + +copy :: CopyCompiler lore r op +copy bt dest destslice src srcslice = do + cc <- asks envCopyCompiler + cc bt dest destslice src srcslice + +-- | Is this copy really a mapping with transpose? +isMapTransposeCopy :: + PrimType -> + MemLocation -> + Slice (Imp.TExp Int64) -> + MemLocation -> + Slice (Imp.TExp Int64) -> + Maybe + ( Imp.TExp Int64, + Imp.TExp Int64, + Imp.TExp Int64, + Imp.TExp Int64, + Imp.TExp Int64 + ) +isMapTransposeCopy + bt + (MemLocation _ _ destIxFun) + destslice + (MemLocation _ _ srcIxFun) + srcslice + | Just (dest_offset, perm_and_destshape) <- IxFun.rearrangeWithOffset destIxFun' bt_size, + (perm, destshape) <- unzip perm_and_destshape, + Just src_offset <- IxFun.linearWithOffset srcIxFun' bt_size, + Just (r1, r2, _) <- isMapTranspose perm = + isOk destshape swap r1 r2 dest_offset src_offset + | Just dest_offset <- IxFun.linearWithOffset destIxFun' bt_size, + Just (src_offset, perm_and_srcshape) <- IxFun.rearrangeWithOffset srcIxFun' bt_size, + (perm, srcshape) <- unzip perm_and_srcshape, + Just (r1, r2, _) <- isMapTranspose perm = + isOk srcshape id r1 r2 dest_offset src_offset + | otherwise = + Nothing + where + bt_size = primByteSize bt + swap (x, y) = (y, x) + + destIxFun' = IxFun.slice destIxFun destslice + srcIxFun' = IxFun.slice srcIxFun srcslice + + isOk shape f r1 r2 dest_offset src_offset = do + let (num_arrays, size_x, size_y) = getSizes shape f r1 r2 + return + ( dest_offset, + src_offset, + num_arrays, + size_x, + size_y + ) + + getSizes shape f r1 r2 = + let (mapped, notmapped) = splitAt r1 shape + (pretrans, posttrans) = f $ splitAt r2 notmapped + in (product mapped, product pretrans, product posttrans) + +mapTransposeName :: PrimType -> String +mapTransposeName bt = "map_transpose_" ++ pretty bt + +mapTransposeForType :: PrimType -> ImpM lore r op Name +mapTransposeForType bt = do + let fname = nameFromString $ "builtin#" <> mapTransposeName bt + + exists <- hasFunction fname + unless exists $ emitFunction fname $ mapTransposeFunction fname bt + + return fname + +-- | Use an 'Imp.Copy' if possible, otherwise 'copyElementWise'. +defaultCopy :: CopyCompiler lore r op +defaultCopy pt dest destslice src srcslice + | Just + ( destoffset, + srcoffset, + num_arrays, + size_x, + size_y + ) <- + isMapTransposeCopy pt dest destslice src srcslice = do + fname <- mapTransposeForType pt + emit $ + Imp.Call + [] + fname + $ transposeArgs + pt + destmem + (bytes destoffset) + srcmem + (bytes srcoffset) + num_arrays + size_x + size_y + | Just destoffset <- + IxFun.linearWithOffset (IxFun.slice dest_ixfun destslice) pt_size, + Just srcoffset <- + IxFun.linearWithOffset (IxFun.slice src_ixfun srcslice) pt_size = do + srcspace <- entryMemSpace <$> lookupMemory srcmem + destspace <- entryMemSpace <$> lookupMemory destmem + if isScalarSpace srcspace || isScalarSpace destspace + then copyElementWise pt dest destslice src srcslice + else + emit $ + Imp.Copy + destmem + (bytes destoffset) + destspace + srcmem + (bytes srcoffset) + srcspace + $ num_elems `withElemType` pt + | otherwise = + copyElementWise pt dest destslice src srcslice + where + pt_size = primByteSize pt + num_elems = Imp.elements $ product $ sliceDims srcslice + + MemLocation destmem _ dest_ixfun = dest + MemLocation srcmem _ src_ixfun = src + + isScalarSpace ScalarSpace {} = True + isScalarSpace _ = False + +copyElementWise :: CopyCompiler lore r op +copyElementWise bt dest destslice src srcslice = do + let bounds = sliceDims srcslice + is <- replicateM (length bounds) (newVName "i") + let ivars = map Imp.vi64 is + (destmem, destspace, destidx) <- + fullyIndexArray' dest $ fixSlice destslice ivars + (srcmem, srcspace, srcidx) <- + fullyIndexArray' src $ fixSlice srcslice ivars + vol <- asks envVolatility + emit $ + foldl (.) id (zipWith Imp.For is $ map untyped bounds) $ + Imp.Write destmem destidx bt destspace vol $ + Imp.index srcmem srcidx bt srcspace vol + +-- | Copy from here to there; both destination and source may be +-- indexeded. +copyArrayDWIM :: + PrimType -> + MemLocation -> + [DimIndex (Imp.TExp Int64)] -> + MemLocation -> + [DimIndex (Imp.TExp Int64)] -> + ImpM lore r op (Imp.Code op) +copyArrayDWIM + bt + destlocation@(MemLocation _ destshape _) + destslice + srclocation@(MemLocation _ srcshape _) + srcslice + | Just destis <- mapM dimFix destslice, + Just srcis <- mapM dimFix srcslice, + length srcis == length srcshape, + length destis == length destshape = do + (targetmem, destspace, targetoffset) <- + fullyIndexArray' destlocation destis + (srcmem, srcspace, srcoffset) <- + fullyIndexArray' srclocation srcis + vol <- asks envVolatility + return $ + Imp.Write targetmem targetoffset bt destspace vol $ + Imp.index srcmem srcoffset bt srcspace vol + | otherwise = do + let destslice' = + fullSliceNum (map toInt64Exp destshape) destslice + srcslice' = + fullSliceNum (map toInt64Exp srcshape) srcslice + destrank = length $ sliceDims destslice' + srcrank = length $ sliceDims srcslice' + if destrank /= srcrank + then + error $ + "copyArrayDWIM: cannot copy to " + ++ pretty (memLocationName destlocation) + ++ " from " + ++ pretty (memLocationName srclocation) + ++ " because ranks do not match (" + ++ pretty destrank + ++ " vs " + ++ pretty srcrank + ++ ")" + else + if destlocation == srclocation && destslice' == srcslice' + then return mempty -- Copy would be no-op. + else collect $ copy bt destlocation destslice' srclocation srcslice' + +-- | Like 'copyDWIM', but the target is a 'ValueDestination' +-- instead of a variable name. +copyDWIMDest :: + ValueDestination -> + [DimIndex (Imp.TExp Int64)] -> + SubExp -> + [DimIndex (Imp.TExp Int64)] -> + ImpM lore r op () +copyDWIMDest _ _ (Constant v) (_ : _) = + error $ + unwords ["copyDWIMDest: constant source", pretty v, "cannot be indexed."] +copyDWIMDest pat dest_slice (Constant v) [] = + case mapM dimFix dest_slice of + Nothing -> + error $ + unwords ["copyDWIMDest: constant source", pretty v, "with slice destination."] + Just dest_is -> + case pat of + ScalarDestination name -> + emit $ Imp.SetScalar name $ Imp.ValueExp v + MemoryDestination {} -> + error $ + unwords ["copyDWIMDest: constant source", pretty v, "cannot be written to memory destination."] + ArrayDestination (Just dest_loc) -> do + (dest_mem, dest_space, dest_i) <- + fullyIndexArray' dest_loc dest_is + vol <- asks envVolatility + emit $ Imp.Write dest_mem dest_i bt dest_space vol $ Imp.ValueExp v + ArrayDestination Nothing -> + error "copyDWIMDest: ArrayDestination Nothing" + where + bt = primValueType v +copyDWIMDest dest dest_slice (Var src) src_slice = do + src_entry <- lookupVar src + case (dest, src_entry) of + (MemoryDestination mem, MemVar _ (MemEntry space)) -> + emit $ Imp.SetMem mem src space + (MemoryDestination {}, _) -> + error $ + unwords ["copyDWIMDest: cannot write", pretty src, "to memory destination."] + (_, MemVar {}) -> + error $ + unwords ["copyDWIMDest: source", pretty src, "is a memory block."] + (_, ScalarVar _ (ScalarEntry _)) + | not $ null src_slice -> + error $ + unwords ["copyDWIMDest: prim-typed source", pretty src, "with slice", pretty src_slice] + (ScalarDestination name, _) + | not $ null dest_slice -> + error $ + unwords ["copyDWIMDest: prim-typed target", pretty name, "with slice", pretty dest_slice] + (ScalarDestination name, ScalarVar _ (ScalarEntry pt)) -> + emit $ Imp.SetScalar name $ Imp.var src pt + (ScalarDestination name, ArrayVar _ arr) + | Just src_is <- mapM dimFix src_slice, + length src_slice == length (entryArrayShape arr) -> do + let bt = entryArrayElemType arr + (mem, space, i) <- + fullyIndexArray' (entryArrayLocation arr) src_is + vol <- asks envVolatility + emit $ Imp.SetScalar name $ Imp.index mem i bt space vol + | otherwise -> + error $ + unwords + [ "copyDWIMDest: prim-typed target", + pretty name, + "and array-typed source", + pretty src, + "with slice", + pretty src_slice + ] + (ArrayDestination (Just dest_loc), ArrayVar _ src_arr) -> do + let src_loc = entryArrayLocation src_arr + bt = entryArrayElemType src_arr + emit =<< copyArrayDWIM bt dest_loc dest_slice src_loc src_slice + (ArrayDestination (Just dest_loc), ScalarVar _ (ScalarEntry bt)) + | Just dest_is <- mapM dimFix dest_slice -> do + (dest_mem, dest_space, dest_i) <- fullyIndexArray' dest_loc dest_is + vol <- asks envVolatility + emit $ Imp.Write dest_mem dest_i bt dest_space vol (Imp.var src bt) + | otherwise -> + error $ + unwords + [ "copyDWIMDest: array-typed target and prim-typed source", + pretty src, + "with slice", + pretty dest_slice + ] + (ArrayDestination Nothing, _) -> + return () -- Nothing to do; something else set some memory + -- somewhere. + +-- | Copy from here to there; both destination and source be +-- indexeded. If so, they better be arrays of enough dimensions. +-- This function will generally just Do What I Mean, and Do The Right +-- Thing. Both destination and source must be in scope. +copyDWIM :: + VName -> + [DimIndex (Imp.TExp Int64)] -> + SubExp -> + [DimIndex (Imp.TExp Int64)] -> + ImpM lore r op () +copyDWIM dest dest_slice src src_slice = do + dest_entry <- lookupVar dest + let dest_target = + case dest_entry of + ScalarVar _ _ -> + ScalarDestination dest + ArrayVar _ (ArrayEntry (MemLocation mem shape ixfun) _) -> + ArrayDestination $ Just $ MemLocation mem shape ixfun + MemVar _ _ -> + MemoryDestination dest + copyDWIMDest dest_target dest_slice src src_slice + +-- | As 'copyDWIM', but implicitly 'DimFix'es the indexes. +copyDWIMFix :: + VName -> + [Imp.TExp Int64] -> + SubExp -> + [Imp.TExp Int64] -> + ImpM lore r op () +copyDWIMFix dest dest_is src src_is = + copyDWIM dest (map DimFix dest_is) src (map DimFix src_is) + +-- | @compileAlloc pat size space@ allocates @n@ bytes of memory in @space@, +-- writing the result to @dest@, which must be a single +-- 'MemoryDestination', +compileAlloc :: + Mem lore => + Pattern lore -> + SubExp -> + Space -> + ImpM lore r op () +compileAlloc (Pattern [] [mem]) e space = do + let e' = Imp.bytes $ toInt64Exp e + allocator <- asks $ M.lookup space . envAllocCompilers + case allocator of + Nothing -> emit $ Imp.Allocate (patElemName mem) e' space + Just allocator' -> allocator' (patElemName mem) e' +compileAlloc pat _ _ = + error $ "compileAlloc: Invalid pattern: " ++ pretty pat + +-- | The number of bytes needed to represent the array in a +-- straightforward contiguous format, as an 'Int64' expression. +typeSize :: Type -> Count Bytes (Imp.TExp Int64) +typeSize t = + Imp.bytes $ + isInt64 (Imp.LeafExp (Imp.SizeOf $ elemType t) int64) + * product (map (sExt64 . toInt64Exp) (arrayDims t)) + +--- Building blocks for constructing code. + +sFor' :: VName -> Imp.Exp -> ImpM lore r op () -> ImpM lore r op () +sFor' i bound body = do + let it = case primExpType bound of + IntType bound_t -> bound_t + t -> error $ "sFor': bound " ++ pretty bound ++ " is of type " ++ pretty t + addLoopVar i it + body' <- collect body + emit $ Imp.For i bound body' + +sFor :: String -> Imp.TExp t -> (Imp.TExp t -> ImpM lore r op ()) -> ImpM lore r op () +sFor i bound body = do + i' <- newVName i + sFor' i' (untyped bound) $ + body $ TPrimExp $ Imp.var i' $ primExpType $ untyped bound + +sWhile :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op () +sWhile cond body = do + body' <- collect body + emit $ Imp.While cond body' + +sComment :: String -> ImpM lore r op () -> ImpM lore r op () +sComment s code = do + code' <- collect code + emit $ Imp.Comment s code' + +sIf :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op () -> ImpM lore r op () +sIf cond tbranch fbranch = do + tbranch' <- collect tbranch + fbranch' <- collect fbranch + emit $ Imp.If cond tbranch' fbranch' + +sWhen :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op () +sWhen cond tbranch = sIf cond tbranch (return ()) + +sUnless :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op () +sUnless cond = sIf cond (return ()) + +sOp :: op -> ImpM lore r op () +sOp = emit . Imp.Op + +sDeclareMem :: String -> Space -> ImpM lore r op VName +sDeclareMem name space = do + name' <- newVName name + emit $ Imp.DeclareMem name' space + addVar name' $ MemVar Nothing $ MemEntry space + return name' + +sAlloc_ :: VName -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM lore r op () +sAlloc_ name' size' space = do + allocator <- asks $ M.lookup space . envAllocCompilers + case allocator of + Nothing -> emit $ Imp.Allocate name' size' space + Just allocator' -> allocator' name' size' + +sAlloc :: String -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM lore r op VName +sAlloc name size space = do + name' <- sDeclareMem name space + sAlloc_ name' size space + return name' + +sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op VName +sArray name bt shape membind = do + name' <- newVName name + dArray name' bt shape membind + return name' + +-- | Declare an array in row-major order in the given memory block. +sArrayInMem :: String -> PrimType -> ShapeBase SubExp -> VName -> ImpM lore r op VName +sArrayInMem name pt shape mem = + sArray name pt shape $ + ArrayIn mem $ + IxFun.iota $ map (isInt64 . primExpFromSubExp int64) $ shapeDims shape + +-- | Like 'sAllocArray', but permute the in-memory representation of the indices as specified. +sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM lore r op VName +sAllocArrayPerm name pt shape space perm = do + let permuted_dims = rearrangeShape perm $ shapeDims shape + mem <- sAlloc (name ++ "_mem") (typeSize (Array pt shape NoUniqueness)) space + let iota_ixfun = IxFun.iota $ map (isInt64 . primExpFromSubExp int64) permuted_dims + sArray name pt shape $ + ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm + +-- | Uses linear/iota index function. +sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM lore r op VName +sAllocArray name pt shape space = + sAllocArrayPerm name pt shape space [0 .. shapeRank shape -1] + +-- | Uses linear/iota index function. +sStaticArray :: String -> Space -> PrimType -> Imp.ArrayContents -> ImpM lore r op VName +sStaticArray name space pt vs = do + let num_elems = case vs of + Imp.ArrayValues vs' -> length vs' + Imp.ArrayZeros n -> fromIntegral n + shape = Shape [intConst Int64 $ toInteger num_elems] + mem <- newVNameForFun $ name ++ "_mem" + emit $ Imp.DeclareArray mem space pt vs + addVar mem $ MemVar Nothing $ MemEntry space + sArray name pt shape $ ArrayIn mem $ IxFun.iota [fromIntegral num_elems] + +sWrite :: VName -> [Imp.TExp Int64] -> Imp.Exp -> ImpM lore r op () +sWrite arr is v = do + (mem, space, offset) <- fullyIndexArray arr is + vol <- asks envVolatility + emit $ Imp.Write mem offset (primExpType v) space vol v + +sUpdate :: VName -> Slice (Imp.TExp Int64) -> SubExp -> ImpM lore r op () +sUpdate arr slice v = copyDWIM arr slice v [] + +sLoopNest :: + Shape -> + ([Imp.TExp Int64] -> ImpM lore r op ()) -> + ImpM lore r op () +sLoopNest = sLoopNest' [] . shapeDims + where + sLoopNest' is [] f = f $ reverse is + sLoopNest' is (d : ds) f = + sFor "nest_i" (toInt64Exp d) $ \i -> sLoopNest' (i : is) ds f + +-- | Untyped assignment. +(<~~) :: VName -> Imp.Exp -> ImpM lore r op () +x <~~ e = emit $ Imp.SetScalar x e + +infixl 3 <~~ + +-- | Typed assignment. +(<--) :: TV t -> Imp.TExp t -> ImpM lore r op () +TV x _ <-- e = emit $ Imp.SetScalar x $ untyped e + +infixl 3 <-- + +-- | Constructing an ad-hoc function that does not +-- correspond to any of the IR functions in the input program. +function :: + Name -> + [Imp.Param] -> + [Imp.Param] -> + ImpM lore r op () -> + ImpM lore r op () +function fname outputs inputs m = local newFunction $ do + body <- collect $ do + mapM_ addParam $ outputs ++ inputs + m + emitFunction fname $ Imp.Function False outputs inputs body [] [] + where + addParam (Imp.MemParam name space) = + addVar name $ MemVar Nothing $ MemEntry space + addParam (Imp.ScalarParam name bt) = + addVar name $ ScalarVar Nothing $ ScalarEntry bt + newFunction env = env {envFunction = Just fname} diff --git a/src/Futhark/CodeGen/ImpGen/CUDA.hs b/src/Futhark/CodeGen/ImpGen/CUDA.hs --- a/src/Futhark/CodeGen/ImpGen/CUDA.hs +++ b/src/Futhark/CodeGen/ImpGen/CUDA.hs @@ -1,14 +1,14 @@ module Futhark.CodeGen.ImpGen.CUDA - ( compileProg - , Warnings - ) where + ( compileProg, + Warnings, + ) +where import Data.Bifunctor (second) - -import Futhark.IR.KernelsMem import Futhark.CodeGen.ImpCode.OpenCL import Futhark.CodeGen.ImpGen.Kernels import Futhark.CodeGen.ImpGen.Kernels.ToOpenCL +import Futhark.IR.KernelsMem import Futhark.MonadFreshNames compileProg :: MonadFreshNames m => Prog KernelsMem -> m (Warnings, Program) 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 @@ -1,137 +1,148 @@ +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE TypeFamilies #-} + -- | Compile a 'KernelsMem' program to imperative code with kernels. -- This is mostly (but not entirely) the same process no matter if we -- are targeting OpenCL or CUDA. The important distinctions (the host -- level code) are introduced later. module Futhark.CodeGen.ImpGen.Kernels - ( compileProgOpenCL - , compileProgCUDA - , Warnings + ( compileProgOpenCL, + compileProgCUDA, + Warnings, ) - where +where import Control.Monad.Except import Data.Bifunctor (second) +import Data.List (foldl') import qualified Data.Map as M import Data.Maybe -import Data.List (foldl') - -import Prelude hiding (quot) - -import Futhark.Error -import Futhark.MonadFreshNames -import Futhark.IR.KernelsMem -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp import Futhark.CodeGen.ImpCode.Kernels (bytes) +import qualified Futhark.CodeGen.ImpCode.Kernels as Imp import Futhark.CodeGen.ImpGen hiding (compileProg) import qualified Futhark.CodeGen.ImpGen import Futhark.CodeGen.ImpGen.Kernels.Base +import Futhark.CodeGen.ImpGen.Kernels.SegHist import Futhark.CodeGen.ImpGen.Kernels.SegMap import Futhark.CodeGen.ImpGen.Kernels.SegRed import Futhark.CodeGen.ImpGen.Kernels.SegScan -import Futhark.CodeGen.ImpGen.Kernels.SegHist import Futhark.CodeGen.ImpGen.Kernels.Transpose -import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.CodeGen.SetDefaultSpace -import Futhark.Util.IntegralExp (quot, divUp, IntegralExp) +import Futhark.Error +import Futhark.IR.KernelsMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Util.IntegralExp (IntegralExp, divUp, quot) +import Prelude hiding (quot) callKernelOperations :: Operations KernelsMem HostEnv Imp.HostOp callKernelOperations = - Operations { opsExpCompiler = expCompiler - , opsCopyCompiler = callKernelCopy - , opsOpCompiler = opCompiler - , opsStmsCompiler = defCompileStms - , opsAllocCompilers = mempty - } + Operations + { opsExpCompiler = expCompiler, + opsCopyCompiler = callKernelCopy, + opsOpCompiler = opCompiler, + opsStmsCompiler = defCompileStms, + opsAllocCompilers = mempty + } openclAtomics, cudaAtomics :: AtomicBinOp (openclAtomics, cudaAtomics) = (flip lookup opencl, flip lookup cuda) - where opencl = [ (Add Int32 OverflowUndef, Imp.AtomicAdd Int32) - , (SMax Int32, Imp.AtomicSMax Int32) - , (SMin Int32, Imp.AtomicSMin Int32) - , (UMax Int32, Imp.AtomicUMax Int32) - , (UMin Int32, Imp.AtomicUMin Int32) - , (And Int32, Imp.AtomicAnd Int32) - , (Or Int32, Imp.AtomicOr Int32) - , (Xor Int32, Imp.AtomicXor Int32) - ] - cuda = opencl ++ [(FAdd Float32, Imp.AtomicFAdd Float32)] + where + opencl = + [ (Add Int32 OverflowUndef, Imp.AtomicAdd Int32), + (SMax Int32, Imp.AtomicSMax Int32), + (SMin Int32, Imp.AtomicSMin Int32), + (UMax Int32, Imp.AtomicUMax Int32), + (UMin Int32, Imp.AtomicUMin Int32), + (And Int32, Imp.AtomicAnd Int32), + (Or Int32, Imp.AtomicOr Int32), + (Xor Int32, Imp.AtomicXor Int32) + ] + cuda = opencl ++ [(FAdd Float32, Imp.AtomicFAdd Float32)] -compileProg :: MonadFreshNames m => HostEnv -> Prog KernelsMem - -> m (Warnings, Imp.Program) +compileProg :: + MonadFreshNames m => + HostEnv -> + Prog KernelsMem -> + m (Warnings, Imp.Program) compileProg env prog = - second (setDefaultSpace (Imp.Space "device")) <$> - Futhark.CodeGen.ImpGen.compileProg env callKernelOperations (Imp.Space "device") prog + second (setDefaultSpace (Imp.Space "device")) + <$> Futhark.CodeGen.ImpGen.compileProg env callKernelOperations (Imp.Space "device") prog -- | Compile a 'KernelsMem' program to low-level parallel code, with -- either CUDA or OpenCL characteristics. -compileProgOpenCL, compileProgCUDA - :: MonadFreshNames m => Prog KernelsMem -> m (Warnings, Imp.Program) +compileProgOpenCL, + compileProgCUDA :: + MonadFreshNames m => Prog KernelsMem -> m (Warnings, Imp.Program) compileProgOpenCL = compileProg $ HostEnv openclAtomics compileProgCUDA = compileProg $ HostEnv cudaAtomics -opCompiler :: Pattern KernelsMem -> Op KernelsMem - -> CallKernelGen () - +opCompiler :: + Pattern KernelsMem -> + Op KernelsMem -> + CallKernelGen () opCompiler dest (Alloc e space) = compileAlloc dest e space - opCompiler (Pattern _ [pe]) (Inner (SizeOp (GetSize key size_class))) = do fname <- askFunction - sOp $ Imp.GetSize (patElemName pe) (keyWithEntryPoint fname key) $ - sizeClassWithEntryPoint fname size_class - + sOp $ + Imp.GetSize (patElemName pe) (keyWithEntryPoint fname key) $ + sizeClassWithEntryPoint fname size_class opCompiler (Pattern _ [pe]) (Inner (SizeOp (CmpSizeLe key size_class x))) = do fname <- askFunction let size_class' = sizeClassWithEntryPoint fname size_class sOp . Imp.CmpSizeLe (patElemName pe) (keyWithEntryPoint fname key) size_class' =<< toExp x - opCompiler (Pattern _ [pe]) (Inner (SizeOp (GetSizeMax size_class))) = sOp $ Imp.GetSizeMax (patElemName pe) size_class - opCompiler (Pattern _ [pe]) (Inner (SizeOp (CalcNumGroups w64 max_num_groups_key group_size))) = do fname <- askFunction - max_num_groups <- dPrim "max_num_groups" int32 - sOp $ Imp.GetSize max_num_groups (keyWithEntryPoint fname max_num_groups_key) $ - sizeClassWithEntryPoint fname SizeNumGroups + max_num_groups :: TV Int32 <- dPrim "max_num_groups" int32 + sOp $ + Imp.GetSize (tvVar max_num_groups) (keyWithEntryPoint fname max_num_groups_key) $ + sizeClassWithEntryPoint fname SizeNumGroups -- If 'w' is small, we launch fewer groups than we normally would. -- We don't want any idle groups. -- -- The calculations are done with 64-bit integers to avoid overflow -- issues. - let num_groups_maybe_zero = BinOpExp (SMin Int64) - (toExp' int64 w64 `divUp` - sExt Int64 (toExp' int32 group_size)) $ - sExt Int64 (Imp.vi32 max_num_groups) + let num_groups_maybe_zero = + sMin64 + ( toInt64Exp w64 + `divUp` sExt64 (toInt32Exp group_size) + ) + $ sExt64 (tvExp max_num_groups) -- We also don't want zero groups. - let num_groups = BinOpExp (SMax Int64) 1 num_groups_maybe_zero - patElemName pe <-- sExt Int32 num_groups - + let num_groups = sMax64 1 num_groups_maybe_zero + mkTV (patElemName pe) int32 <-- sExt32 num_groups opCompiler dest (Inner (SegOp op)) = segOpCompiler dest op - opCompiler pat e = - compilerBugS $ "opCompiler: Invalid pattern\n " ++ - pretty pat ++ "\nfor expression\n " ++ pretty e + compilerBugS $ + "opCompiler: Invalid pattern\n " + ++ pretty pat + ++ "\nfor expression\n " + ++ pretty e sizeClassWithEntryPoint :: Maybe Name -> Imp.SizeClass -> Imp.SizeClass sizeClassWithEntryPoint fname (Imp.SizeThreshold path def) = Imp.SizeThreshold (map f path) def - where f (name, x) = (keyWithEntryPoint fname name, x) + where + f (name, x) = (keyWithEntryPoint fname name, x) sizeClassWithEntryPoint _ size_class = size_class -segOpCompiler :: Pattern KernelsMem -> SegOp SegLevel KernelsMem - -> CallKernelGen () +segOpCompiler :: + Pattern KernelsMem -> + SegOp SegLevel KernelsMem -> + CallKernelGen () segOpCompiler pat (SegMap lvl space _ kbody) = compileSegMap pat lvl space kbody -segOpCompiler pat (SegRed lvl@SegThread{} space reds _ kbody) = +segOpCompiler pat (SegRed lvl@SegThread {} space reds _ kbody) = compileSegRed pat lvl space reds kbody -segOpCompiler pat (SegScan lvl@SegThread{} space scans _ kbody) = +segOpCompiler pat (SegScan lvl@SegThread {} space scans _ kbody) = compileSegScan pat lvl space scans kbody segOpCompiler pat (SegHist (SegThread num_groups group_size _) space ops _ kbody) = compileSegHist pat num_groups group_size space ops kbody @@ -144,7 +155,7 @@ -- otherwise protected by their own multi-versioning branches deeper -- down. Currently the compiler will not generate multi-versioning -- that makes this a problem, but it might in the future. -checkLocalMemoryReqs :: Imp.Code -> CallKernelGen (Maybe Imp.Exp) +checkLocalMemoryReqs :: Imp.Code -> CallKernelGen (Maybe (Imp.TExp Bool)) checkLocalMemoryReqs code = do scope <- askScope let alloc_sizes = map (sum . localAllocSizes . Imp.kernelBody) $ getKernels code @@ -154,40 +165,35 @@ if any (`M.notMember` scope) (namesToList $ freeIn alloc_sizes) then return Nothing else do - local_memory_capacity <- dPrim "local_memory_capacity" int32 - sOp $ Imp.GetSizeMax local_memory_capacity SizeLocalMemory - - let local_memory_capacity_64 = - sExt Int64 $ Imp.vi32 local_memory_capacity - fits size = - unCount size .<=. local_memory_capacity_64 - return $ Just $ foldl' (.&&.) true (map fits alloc_sizes) + local_memory_capacity :: TV Int32 <- dPrim "local_memory_capacity" int32 + sOp $ Imp.GetSizeMax (tvVar local_memory_capacity) SizeLocalMemory - where getKernels = foldMap getKernel - getKernel (Imp.CallKernel k) = [k] - getKernel _ = [] + let local_memory_capacity_64 = + sExt64 $ tvExp local_memory_capacity + fits size = + unCount size .<=. local_memory_capacity_64 + return $ Just $ foldl' (.&&.) true (map fits alloc_sizes) + where + getKernels = foldMap getKernel + getKernel (Imp.CallKernel k) = [k] + getKernel _ = [] - localAllocSizes = foldMap localAllocSize - localAllocSize (Imp.LocalAlloc _ size) = [size] - localAllocSize _ = [] + localAllocSizes = foldMap localAllocSize + localAllocSize (Imp.LocalAlloc _ size) = [size] + localAllocSize _ = [] expCompiler :: ExpCompiler KernelsMem HostEnv Imp.HostOp - -- We generate a simple kernel for itoa and replicate. expCompiler (Pattern _ [pe]) (BasicOp (Iota n x s et)) = do - n' <- toExp n x' <- toExp x s' <- toExp s - sIota (patElemName pe) n' x' s' et - + sIota (patElemName pe) (toInt64Exp n) x' s' et expCompiler (Pattern _ [pe]) (BasicOp (Replicate _ se)) = sReplicate (patElemName pe) se - -- Allocation in the "local" space is just a placeholder. expCompiler _ (Op (Alloc _ (Space "local"))) = return () - -- This is a multi-versioning If created by incremental flattening. -- We need to augment the conditional with a check that any local -- memory requirements in tbranch are compatible with the hardware. @@ -199,40 +205,57 @@ fcode <- collect $ compileBody dest fbranch check <- checkLocalMemoryReqs tcode emit $ case check of - Nothing -> fcode - Just ok -> Imp.If (ok .&&. toExp' Bool cond) tcode fcode - + Nothing -> fcode + Just ok -> Imp.If (ok .&&. toBoolExp cond) tcode fcode expCompiler dest e = defCompileExp dest e callKernelCopy :: CopyCompiler KernelsMem HostEnv Imp.HostOp -callKernelCopy bt - destloc@(MemLocation destmem _ destIxFun) destslice - srcloc@(MemLocation srcmem srcshape srcIxFun) srcslice - | Just (destoffset, srcoffset, - num_arrays, size_x, size_y) <- - isMapTransposeKernel bt destloc destslice srcloc srcslice = do - +callKernelCopy + bt + destloc@(MemLocation destmem _ destIxFun) + destslice + srcloc@(MemLocation srcmem srcshape srcIxFun) + srcslice + | Just + ( destoffset, + srcoffset, + num_arrays, + size_x, + size_y + ) <- + isMapTransposeCopy bt destloc destslice srcloc srcslice = do fname <- mapTransposeForType bt - emit $ Imp.Call [] fname - [Imp.MemArg destmem, Imp.ExpArg destoffset, - Imp.MemArg srcmem, Imp.ExpArg srcoffset, - Imp.ExpArg num_arrays, Imp.ExpArg size_x, Imp.ExpArg size_y] - - | bt_size <- primByteSize bt, - Just destoffset <- - IxFun.linearWithOffset (IxFun.slice destIxFun destslice) bt_size, - Just srcoffset <- - IxFun.linearWithOffset (IxFun.slice srcIxFun srcslice) bt_size = do - let num_elems = Imp.elements $ product $ map (toExp' int32) srcshape - srcspace <- entryMemSpace <$> lookupMemory srcmem - destspace <- entryMemSpace <$> lookupMemory destmem - emit $ Imp.Copy - destmem (bytes destoffset) destspace - srcmem (bytes srcoffset) srcspace $ - num_elems `Imp.withElemType` bt - - | otherwise = sCopy bt destloc destslice srcloc srcslice + emit $ + Imp.Call + [] + fname + [ Imp.MemArg destmem, + Imp.ExpArg $ untyped destoffset, + Imp.MemArg srcmem, + Imp.ExpArg $ untyped srcoffset, + Imp.ExpArg $ untyped num_arrays, + Imp.ExpArg $ untyped size_x, + Imp.ExpArg $ untyped size_y + ] + | bt_size <- primByteSize bt, + Just destoffset <- + IxFun.linearWithOffset (IxFun.slice destIxFun destslice) bt_size, + Just srcoffset <- + IxFun.linearWithOffset (IxFun.slice srcIxFun srcslice) bt_size = do + let num_elems = Imp.elements $ product $ map toInt64Exp srcshape + srcspace <- entryMemSpace <$> lookupMemory srcmem + destspace <- entryMemSpace <$> lookupMemory destmem + emit $ + Imp.Copy + destmem + (bytes $ sExt64 destoffset) + destspace + srcmem + (bytes $ sExt64 srcoffset) + srcspace + $ num_elems `Imp.withElemType` bt + | otherwise = sCopy bt destloc destslice srcloc srcslice mapTransposeForType :: PrimType -> CallKernelGen Name mapTransposeForType bt = do @@ -244,129 +267,123 @@ return fname mapTransposeName :: PrimType -> String -mapTransposeName bt = "map_transpose_" ++ pretty bt +mapTransposeName bt = "gpu_map_transpose_" ++ pretty bt mapTransposeFunction :: PrimType -> Imp.Function mapTransposeFunction bt = Imp.Function False [] params transpose_code [] [] - - where params = [memparam destmem, intparam destoffset, - memparam srcmem, intparam srcoffset, - intparam num_arrays, intparam x, intparam y] - - space = Space "device" - memparam v = Imp.MemParam v space - intparam v = Imp.ScalarParam v $ IntType Int32 - - [destmem, destoffset, srcmem, srcoffset, - num_arrays, x, y, - mulx, muly, block] = - zipWith (VName . nameFromString) - ["destmem", - "destoffset", - "srcmem", - "srcoffset", - "num_arrays", - "x_elems", - "y_elems", - -- The following is only used for low width/height - -- transpose kernels - "mulx", - "muly", - "block" - ] - [0..] - - v32 v = Imp.var v int32 - - block_dim_int = 16 - - block_dim :: IntegralExp a => a - block_dim = 16 + where + params = + [ memparam destmem, + intparam destoffset, + memparam srcmem, + intparam srcoffset, + intparam num_arrays, + intparam x, + intparam y + ] - -- When an input array has either width==1 or height==1, performing a - -- transpose will be the same as performing a copy. - can_use_copy = - let onearr = CmpOpExp (CmpEq $ IntType Int32) (v32 num_arrays) 1 - height_is_one = CmpOpExp (CmpEq $ IntType Int32) (v32 y) 1 - width_is_one = CmpOpExp (CmpEq $ IntType Int32) (v32 x) 1 - in onearr .&&. (width_is_one .||. height_is_one) + space = Space "device" + memparam v = Imp.MemParam v space + intparam v = Imp.ScalarParam v $ IntType Int32 - transpose_code = - Imp.If input_is_empty mempty $ mconcat - [ Imp.DeclareScalar muly Imp.Nonvolatile (IntType Int32) - , Imp.SetScalar muly $ block_dim `quot` v32 x - , Imp.DeclareScalar mulx Imp.Nonvolatile (IntType Int32) - , Imp.SetScalar mulx $ block_dim `quot` v32 y - , Imp.If can_use_copy copy_code $ - Imp.If should_use_lowwidth (callTransposeKernel TransposeLowWidth) $ - Imp.If should_use_lowheight (callTransposeKernel TransposeLowHeight) $ - Imp.If should_use_small (callTransposeKernel TransposeSmall) $ - callTransposeKernel TransposeNormal] + [ destmem, + destoffset, + srcmem, + srcoffset, + num_arrays, + x, + y, + mulx, + muly, + block + ] = + zipWith + (VName . nameFromString) + [ "destmem", + "destoffset", + "srcmem", + "srcoffset", + "num_arrays", + "x_elems", + "y_elems", + -- The following is only used for low width/height + -- transpose kernels + "mulx", + "muly", + "block" + ] + [0 ..] - input_is_empty = - v32 num_arrays .==. 0 .||. v32 x .==. 0 .||. v32 y .==. 0 + block_dim_int = 16 - should_use_small = BinOpExp LogAnd - (CmpOpExp (CmpSle Int32) (v32 x) (block_dim `quot` 2)) - (CmpOpExp (CmpSle Int32) (v32 y) (block_dim `quot` 2)) + block_dim :: IntegralExp a => a + block_dim = 16 - should_use_lowwidth = BinOpExp LogAnd - (CmpOpExp (CmpSle Int32) (v32 x) (block_dim `quot` 2)) - (CmpOpExp (CmpSlt Int32) block_dim (v32 y)) + -- When an input array has either width==1 or height==1, performing a + -- transpose will be the same as performing a copy. + can_use_copy = + let onearr = Imp.vi32 num_arrays .==. 1 + height_is_one = Imp.vi32 y .==. 1 + width_is_one = Imp.vi32 x .==. 1 + in onearr .&&. (width_is_one .||. height_is_one) - should_use_lowheight = BinOpExp LogAnd - (CmpOpExp (CmpSle Int32) (v32 y) (block_dim `quot` 2)) - (CmpOpExp (CmpSlt Int32) block_dim (v32 x)) + transpose_code = + Imp.If input_is_empty mempty $ + mconcat + [ Imp.DeclareScalar muly Imp.Nonvolatile (IntType Int32), + Imp.SetScalar muly $ untyped $ block_dim `quot` Imp.vi32 x, + Imp.DeclareScalar mulx Imp.Nonvolatile (IntType Int32), + Imp.SetScalar mulx $ untyped $ block_dim `quot` Imp.vi32 y, + Imp.If can_use_copy copy_code $ + Imp.If should_use_lowwidth (callTransposeKernel TransposeLowWidth) $ + Imp.If should_use_lowheight (callTransposeKernel TransposeLowHeight) $ + Imp.If should_use_small (callTransposeKernel TransposeSmall) $ + callTransposeKernel TransposeNormal + ] - copy_code = - let num_bytes = - v32 x * v32 y * Imp.LeafExp (Imp.SizeOf bt) (IntType Int32) - in Imp.Copy - destmem (Imp.Count $ v32 destoffset) space - srcmem (Imp.Count $ v32 srcoffset) space - (Imp.Count num_bytes) + input_is_empty = + Imp.vi32 num_arrays .==. 0 .||. Imp.vi32 x .==. 0 .||. Imp.vi32 y .==. 0 - callTransposeKernel = - Imp.Op . Imp.CallKernel . - mapTransposeKernel (mapTransposeName bt) block_dim_int - (destmem, v32 destoffset, srcmem, v32 srcoffset, - v32 x, v32 y, - v32 mulx, v32 muly, v32 num_arrays, - block) bt + should_use_small = + Imp.vi32 x .<=. (block_dim `quot` 2) + .&&. Imp.vi32 y .<=. (block_dim `quot` 2) -isMapTransposeKernel :: PrimType - -> MemLocation -> Slice Imp.Exp - -> MemLocation -> Slice Imp.Exp - -> Maybe (Imp.Exp, Imp.Exp, - Imp.Exp, Imp.Exp, Imp.Exp) -isMapTransposeKernel bt - (MemLocation _ _ destIxFun) destslice - (MemLocation _ _ srcIxFun) srcslice - | Just (dest_offset, perm_and_destshape) <- IxFun.rearrangeWithOffset destIxFun' bt_size, - (perm, destshape) <- unzip perm_and_destshape, - Just src_offset <- IxFun.linearWithOffset srcIxFun' bt_size, - Just (r1, r2, _) <- isMapTranspose perm = - isOk destshape swap r1 r2 dest_offset src_offset - | Just dest_offset <- IxFun.linearWithOffset destIxFun' bt_size, - Just (src_offset, perm_and_srcshape) <- IxFun.rearrangeWithOffset srcIxFun' bt_size, - (perm, srcshape) <- unzip perm_and_srcshape, - Just (r1, r2, _) <- isMapTranspose perm = - isOk srcshape id r1 r2 dest_offset src_offset - | otherwise = - Nothing - where bt_size = primByteSize bt - swap (x,y) = (y,x) + should_use_lowwidth = + Imp.vi32 x .<=. (block_dim `quot` 2) + .&&. block_dim .<. Imp.vi32 y - destIxFun' = IxFun.slice destIxFun destslice - srcIxFun' = IxFun.slice srcIxFun srcslice + should_use_lowheight = + Imp.vi32 y .<=. (block_dim `quot` 2) + .&&. block_dim .<. Imp.vi32 x - isOk shape f r1 r2 dest_offset src_offset = do - let (num_arrays, size_x, size_y) = getSizes shape f r1 r2 - return (dest_offset, src_offset, - num_arrays, size_x, size_y) + copy_code = + let num_bytes = + sExt64 $ + Imp.vi32 x * Imp.vi32 y * isInt32 (Imp.LeafExp (Imp.SizeOf bt) (IntType Int32)) + in Imp.Copy + destmem + (Imp.Count $ sExt64 $ Imp.vi32 destoffset) + space + srcmem + (Imp.Count $ sExt64 $ Imp.vi32 srcoffset) + space + (Imp.Count num_bytes) - getSizes shape f r1 r2 = - let (mapped, notmapped) = splitAt r1 shape - (pretrans, posttrans) = f $ splitAt r2 notmapped - in (product mapped, product pretrans, product posttrans) + callTransposeKernel = + Imp.Op . Imp.CallKernel + . mapTransposeKernel + (mapTransposeName bt) + block_dim_int + ( destmem, + Imp.vi32 destoffset, + srcmem, + Imp.vi32 srcoffset, + Imp.vi32 x, + Imp.vi32 y, + Imp.vi32 mulx, + Imp.vi32 muly, + Imp.vi32 num_arrays, + block + ) + bt 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 @@ -1,1514 +1,1699 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE TypeFamilies #-} -module Futhark.CodeGen.ImpGen.Kernels.Base - ( KernelConstants (..) - , keyWithEntryPoint - , CallKernelGen - , InKernelGen - , HostEnv (..) - , KernelEnv (..) - , computeThreadChunkSize - , groupReduce - , groupScan - , isActive - , sKernelThread - , sKernelGroup - , sReplicate - , sIota - , sCopy - , compileThreadResult - , compileGroupResult - , virtualiseGroups - , groupLoop - , kernelLoop - , groupCoverSpace - , precomputeSegOpIDs - - , atomicUpdateLocking - , AtomicBinOp - , Locking(..) - , AtomicUpdate(..) - , DoAtomicUpdate - ) - where - -import Control.Monad.Except -import Data.Maybe -import qualified Data.Map.Strict as M -import qualified Data.Set as S -import Data.List (elemIndex, find, nub, zip4) - -import Prelude hiding (quot, rem) - -import Futhark.Error -import Futhark.MonadFreshNames -import Futhark.Transform.Rename -import Futhark.IR.KernelsMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen -import Futhark.Util.IntegralExp (divUp, quot, rem) -import Futhark.Util (chunks, maybeNth, mapAccumLM, takeLast, dropLast) - -newtype HostEnv = HostEnv - { hostAtomics :: AtomicBinOp } - -data KernelEnv = KernelEnv - { kernelAtomics :: AtomicBinOp - , kernelConstants :: KernelConstants - } - -type CallKernelGen = ImpM KernelsMem HostEnv Imp.HostOp -type InKernelGen = ImpM KernelsMem KernelEnv Imp.KernelOp - -data KernelConstants = - KernelConstants - { kernelGlobalThreadId :: Imp.Exp - , kernelLocalThreadId :: Imp.Exp - , kernelGroupId :: Imp.Exp - , kernelGlobalThreadIdVar :: VName - , kernelLocalThreadIdVar :: VName - , kernelGroupIdVar :: VName - , kernelNumGroups :: Imp.Exp - , kernelGroupSize :: Imp.Exp - , kernelNumThreads :: Imp.Exp - , kernelWaveSize :: Imp.Exp - , kernelThreadActive :: Imp.Exp - , kernelLocalIdMap :: M.Map [SubExp] [Imp.Exp] - -- ^ A mapping from dimensions of nested SegOps to already - -- computed local thread IDs. - } - -segOpSizes :: Stms KernelsMem -> S.Set [SubExp] -segOpSizes = onStms - where onStms = foldMap (onExp . stmExp) - onExp (Op (Inner (SegOp op))) = - S.singleton $ map snd $ unSegSpace $ segSpace op - onExp (If _ tbranch fbranch _) = - onStms (bodyStms tbranch) <> onStms (bodyStms fbranch) - onExp (DoLoop _ _ _ body) = - onStms (bodyStms body) - onExp _ = mempty - -precomputeSegOpIDs :: Stms KernelsMem -> InKernelGen a -> InKernelGen a -precomputeSegOpIDs stms m = do - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - new_ids <- M.fromList <$> mapM (mkMap ltid) (S.toList (segOpSizes stms)) - let f env = env { kernelConstants = - (kernelConstants env) { kernelLocalIdMap = new_ids } - } - localEnv f m - where mkMap ltid dims = do - dims' <- mapM toExp dims - ids' <- mapM (dPrimVE "ltid_pre") $ unflattenIndex dims' ltid - return (dims, ids') - -keyWithEntryPoint :: Maybe Name -> Name -> Name -keyWithEntryPoint fname key = - nameFromString $ maybe "" ((++".") . nameToString) fname ++ nameToString key - -allocLocal :: AllocCompiler KernelsMem r Imp.KernelOp -allocLocal mem size = - sOp $ Imp.LocalAlloc mem size - -kernelAlloc :: Pattern KernelsMem - -> SubExp -> Space - -> InKernelGen () -kernelAlloc (Pattern _ [_]) _ ScalarSpace{} = - -- Handled by the declaration of the memory block, which is then - -- translated to an actual scalar variable during C code generation. - return () -kernelAlloc (Pattern _ [mem]) size (Space "local") = do - size' <- toExp size - allocLocal (patElemName mem) $ Imp.bytes size' -kernelAlloc (Pattern _ [mem]) _ _ = - compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." -kernelAlloc dest _ _ = - error $ "Invalid target for in-kernel allocation: " ++ show dest - -splitSpace :: (ToExp w, ToExp i, ToExp elems_per_thread) => - Pattern KernelsMem -> SplitOrdering -> w -> i -> elems_per_thread - -> ImpM lore r op () -splitSpace (Pattern [] [size]) o w i elems_per_thread = do - num_elements <- Imp.elements <$> toExp w - i' <- toExp i - elems_per_thread' <- Imp.elements <$> toExp elems_per_thread - computeThreadChunkSize o i' elems_per_thread' num_elements (patElemName size) -splitSpace pat _ _ _ _ = - error $ "Invalid target for splitSpace: " ++ pretty pat - -compileThreadExp :: ExpCompiler KernelsMem KernelEnv Imp.KernelOp -compileThreadExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = - forM_ (zip [0..] es) $ \(i,e) -> - copyDWIMFix (patElemName dest) [fromIntegral (i::Int32)] e [] -compileThreadExp dest e = - defCompileExp dest e - - --- | Assign iterations of a for-loop to all threads in the kernel. --- The passed-in function is invoked with the (symbolic) iteration. --- 'threadOperations' will be in effect in the body. For --- multidimensional loops, use 'groupCoverSpace'. -kernelLoop :: Imp.Exp -> Imp.Exp -> Imp.Exp - -> (Imp.Exp -> InKernelGen ()) -> InKernelGen () -kernelLoop tid num_threads n f = - localOps threadOperations $ - if n == num_threads then - f tid - else do - -- Compute how many elements this thread is responsible for. - -- Formula: (n - tid) / num_threads (rounded up). - let elems_for_this = (n - tid) `divUp` num_threads - - sFor "i" elems_for_this $ \i -> f $ - i * num_threads + tid - --- | Assign iterations of a for-loop to threads in the workgroup. The --- passed-in function is invoked with the (symbolic) iteration. For --- multidimensional loops, use 'groupCoverSpace'. -groupLoop :: Imp.Exp - -> (Imp.Exp -> InKernelGen ()) -> InKernelGen () -groupLoop n f = do - constants <- kernelConstants <$> askEnv - kernelLoop (kernelLocalThreadId constants) (kernelGroupSize constants) n f - --- | Iterate collectively though a multidimensional space, such that --- all threads in the group participate. The passed-in function is --- invoked with a (symbolic) point in the index space. -groupCoverSpace :: [Imp.Exp] - -> ([Imp.Exp] -> InKernelGen ()) -> InKernelGen () -groupCoverSpace ds f = - groupLoop (product ds) $ f . unflattenIndex ds - -compileGroupExp :: ExpCompiler KernelsMem KernelEnv Imp.KernelOp --- The static arrays stuff does not work inside kernels. -compileGroupExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = - forM_ (zip [0..] es) $ \(i,e) -> - copyDWIMFix (patElemName dest) [fromIntegral (i::Int32)] e [] -compileGroupExp (Pattern _ [dest]) (BasicOp (Replicate ds se)) = do - ds' <- mapM toExp $ shapeDims ds - groupCoverSpace ds' $ \is -> - copyDWIMFix (patElemName dest) is se (drop (shapeRank ds) is) - sOp $ Imp.Barrier Imp.FenceLocal -compileGroupExp (Pattern _ [dest]) (BasicOp (Iota n e s _)) = do - n' <- toExp n - e' <- toExp e - s' <- toExp s - groupLoop n' $ \i' -> do - x <- dPrimV "x" $ e' + i' * s' - copyDWIMFix (patElemName dest) [i'] (Var x) [] - sOp $ Imp.Barrier Imp.FenceLocal -compileGroupExp dest e = - defCompileExp dest e - -sanityCheckLevel :: SegLevel -> InKernelGen () -sanityCheckLevel SegThread{} = return () -sanityCheckLevel SegGroup{} = - error "compileGroupOp: unexpected group-level SegOp." - -localThreadIDs :: [SubExp] -> InKernelGen [Imp.Exp] -localThreadIDs dims = do - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - dims' <- mapM toExp dims - fromMaybe (unflattenIndex dims' ltid) . - M.lookup dims . kernelLocalIdMap . kernelConstants <$> askEnv - -compileGroupSpace :: SegLevel -> SegSpace -> InKernelGen () -compileGroupSpace lvl space = do - sanityCheckLevel lvl - let (ltids, dims) = unzip $ unSegSpace space - zipWithM_ dPrimV_ ltids =<< localThreadIDs dims - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - dPrimV_ (segFlat space) ltid - --- Construct the necessary lock arrays for an intra-group histogram. -prepareIntraGroupSegHist :: Count GroupSize SubExp - -> [HistOp KernelsMem] - -> InKernelGen [[Imp.Exp] -> InKernelGen ()] -prepareIntraGroupSegHist group_size = - fmap snd . mapAccumLM onOp Nothing - where - onOp l op = do - - constants <- kernelConstants <$> askEnv - atomicBinOp <- kernelAtomics <$> askEnv - - let local_subhistos = histDest op - - case (l, atomicUpdateLocking atomicBinOp $ histOp op) of - (_, AtomicPrim f) -> return (l, f (Space "local") local_subhistos) - (_, AtomicCAS f) -> return (l, f (Space "local") local_subhistos) - (Just l', AtomicLocking f) -> return (l, f l' (Space "local") local_subhistos) - (Nothing, AtomicLocking f) -> do - locks <- newVName "locks" - num_locks <- toExp $ unCount group_size - - let dims = map (toExp' int32) $ - shapeDims (histShape op) ++ - [histWidth op] - l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims) - locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness - - locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local" - dArray locks int32 (arrayShape locks_t) $ - ArrayIn locks_mem $ IxFun.iota $ - map (primExpFromSubExp int32) $ arrayDims locks_t - - sComment "All locks start out unlocked" $ - groupCoverSpace [kernelGroupSize constants] $ \is -> - copyDWIMFix locks is (intConst Int32 0) [] - - return (Just l', f l' (Space "local") local_subhistos) - -whenActive :: SegLevel -> SegSpace -> InKernelGen () -> InKernelGen () -whenActive lvl space m - | SegNoVirtFull <- segVirt lvl = m - | otherwise = sWhen (isActive $ unSegSpace space) m - -compileGroupOp :: OpCompiler KernelsMem KernelEnv Imp.KernelOp - -compileGroupOp pat (Alloc size space) = - kernelAlloc pat size space - -compileGroupOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = - splitSpace pat o w i elems_per_thread - -compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do - void $ compileGroupSpace lvl space - - whenActive lvl space $ localOps threadOperations $ - compileStms mempty (kernelBodyStms body) $ - zipWithM_ (compileThreadResult space) (patternElements pat) $ - kernelBodyResult body - - sOp $ Imp.ErrorSync Imp.FenceLocal - -compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do - compileGroupSpace lvl space - let (ltids, dims) = unzip $ unSegSpace space - dims' <- mapM toExp dims - - whenActive lvl space $ - compileStms mempty (kernelBodyStms body) $ - forM_ (zip (patternNames pat) $ kernelBodyResult body) $ \(dest, res) -> - copyDWIMFix dest - (map (`Imp.var` int32) ltids) - (kernelResultSubExp res) [] - - sOp $ Imp.ErrorSync Imp.FenceLocal - - let segment_size = last dims' - crossesSegment from to = (to-from) .>. (to `rem` segment_size) - - -- groupScan needs to treat the scan output as a one-dimensional - -- array of scan elements, so we invent some new flattened arrays - -- here. XXX: this assumes that the original index function is just - -- row-major, but does not actually verify it. - dims_flat <- dPrimV "dims_flat" $ product dims' - let flattened pe = do - MemLocation mem _ _ <- - entryArrayLocation <$> lookupArray (patElemName pe) - let pe_t = typeOf pe - arr_dims = Var dims_flat : drop (length dims') (arrayDims pe_t) - sArray (baseString (patElemName pe) ++ "_flat") - (elemType pe_t) (Shape arr_dims) $ - ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) arr_dims - - num_scan_results = sum $ map (length . segBinOpNeutral) scans - - arrs_flat <- mapM flattened $ take num_scan_results $ patternElements pat - - forM_ scans $ \scan -> do - let scan_op = segBinOpLambda scan - groupScan (Just crossesSegment) (product dims') (product dims') scan_op arrs_flat - -compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do - compileGroupSpace lvl space - - let (ltids, dims) = unzip $ unSegSpace space - (red_pes, map_pes) = - splitAt (segBinOpResults ops) $ patternElements pat - - dims' <- mapM toExp dims - - let mkTempArr t = - sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local" - tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops - let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs - - whenActive lvl space $ - compileStms mempty (kernelBodyStms body) $ do - let (red_res, map_res) = - splitAt (segBinOpResults ops) $ kernelBodyResult body - forM_ (zip tmp_arrs red_res) $ \(dest, res) -> - copyDWIMFix dest (map (`Imp.var` int32) ltids) (kernelResultSubExp res) [] - zipWithM_ (compileThreadResult space) map_pes map_res - - sOp $ Imp.ErrorSync Imp.FenceLocal - - case dims' of - -- Nonsegmented case (or rather, a single segment) - this we can - -- handle directly with a group-level reduction. - [dim'] -> do - forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> - groupReduce dim' (segBinOpLambda op) tmps - - sOp $ Imp.ErrorSync Imp.FenceLocal - - forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> - copyDWIMFix (patElemName pe) [] (Var arr) [0] - - _ -> do - -- Segmented intra-group reductions are turned into (regular) - -- segmented scans. It is possible that this can be done - -- better, but at least this approach is simple. - - -- groupScan operates on flattened arrays. This does not - -- involve copying anything; merely playing with the index - -- function. - dims_flat <- dPrimV "dims_flat" $ product dims' - let flatten arr = do - ArrayEntry arr_loc pt <- lookupArray arr - let flat_shape = Shape $ Var dims_flat : - drop (length ltids) (memLocationShape arr_loc) - sArray "red_arr_flat" pt flat_shape $ - ArrayIn (memLocationName arr_loc) $ - IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims flat_shape - - let segment_size = last dims' - crossesSegment from to = (to-from) .>. (to `rem` segment_size) - - forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do - tmps_flat <- mapM flatten tmps - groupScan (Just crossesSegment) (product dims') (product dims') - (segBinOpLambda op) tmps_flat - - sOp $ Imp.ErrorSync Imp.FenceLocal - - forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> - copyDWIM (patElemName pe) [] (Var arr) - (map (unitSlice 0) (init dims') ++ [DimFix $ last dims'-1]) - - sOp $ Imp.Barrier Imp.FenceLocal - -compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do - compileGroupSpace lvl space - let ltids = map fst $ unSegSpace space - - -- We don't need the red_pes, because it is guaranteed by our type - -- rules that they occupy the same memory as the destinations for - -- the ops. - let num_red_res = length ops + sum (map (length . histNeutral) ops) - (_red_pes, map_pes) = - splitAt num_red_res $ patternElements pat - - ops' <- prepareIntraGroupSegHist (segGroupSize lvl) ops - - -- Ensure that all locks have been initialised. - sOp $ Imp.Barrier Imp.FenceLocal - - whenActive lvl space $ - compileStms mempty (kernelBodyStms kbody) $ do - let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody - (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res - zipWithM_ (compileThreadResult space) map_pes map_res - - let vs_per_op = chunks (map (length . histDest) ops) red_vs - - forM_ (zip4 red_is vs_per_op ops' ops) $ - \(bin, op_vs, do_op, HistOp dest_w _ _ _ shape lam) -> do - let bin' = toExp' int32 bin - dest_w' = toExp' int32 dest_w - bin_in_bounds = 0 .<=. bin' .&&. bin' .<. dest_w' - bin_is = map (`Imp.var` int32) (init ltids) ++ [bin'] - vs_params = takeLast (length op_vs) $ lambdaParams lam - - sComment "perform atomic updates" $ - sWhen bin_in_bounds $ do - dLParams $ lambdaParams lam - sLoopNest shape $ \is -> do - forM_ (zip vs_params op_vs) $ \(p, v) -> - copyDWIMFix (paramName p) [] v is - do_op (bin_is ++ is) - - sOp $ Imp.ErrorSync Imp.FenceLocal - -compileGroupOp pat _ = - compilerBugS $ "compileGroupOp: cannot compile rhs of binding " ++ pretty pat - -compileThreadOp :: OpCompiler KernelsMem KernelEnv Imp.KernelOp -compileThreadOp pat (Alloc size space) = - kernelAlloc pat size space -compileThreadOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = - splitSpace pat o w i elems_per_thread -compileThreadOp pat _ = - compilerBugS $ "compileThreadOp: cannot compile rhs of binding " ++ pretty pat - --- | Locking strategy used for an atomic update. -data Locking = - Locking { lockingArray :: VName - -- ^ Array containing the lock. - , lockingIsUnlocked :: Imp.Exp - -- ^ Value for us to consider the lock free. - , lockingToLock :: Imp.Exp - -- ^ What to write when we lock it. - , lockingToUnlock :: Imp.Exp - -- ^ What to write when we unlock it. - , lockingMapping :: [Imp.Exp] -> [Imp.Exp] - -- ^ A transformation from the logical lock index to the - -- physical position in the array. This can also be used - -- to make the lock array smaller. - } - --- | A function for generating code for an atomic update. Assumes --- that the bucket is in-bounds. -type DoAtomicUpdate lore r = - Space -> [VName] -> [Imp.Exp] -> ImpM lore r Imp.KernelOp () - --- | The mechanism that will be used for performing the atomic update. --- Approximates how efficient it will be. Ordered from most to least --- efficient. -data AtomicUpdate lore r - = AtomicPrim (DoAtomicUpdate lore r) - -- ^ Supported directly by primitive. - | AtomicCAS (DoAtomicUpdate lore r) - -- ^ Can be done by efficient swaps. - | AtomicLocking (Locking -> DoAtomicUpdate lore r) - -- ^ Requires explicit locking. - --- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? -type AtomicBinOp = - BinOp -> - Maybe (VName -> VName -> Count Imp.Elements Imp.Exp -> Imp.Exp -> Imp.AtomicOp) - --- | Do an atomic update corresponding to a binary operator lambda. -atomicUpdateLocking :: AtomicBinOp -> Lambda KernelsMem - -> AtomicUpdate KernelsMem KernelEnv - -atomicUpdateLocking atomicBinOp lam - | Just ops_and_ts <- splitOp lam, - all (\(_, t, _, _) -> primBitSize t == 32) 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. - forM_ (zip arrs ops_and_ts) $ \(a, (op, t, x, y)) -> do - - -- Common variables. - old <- dPrim "old" t - - (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket - - case opHasAtomicSupport space old arr' bucket_offset op of - Just f -> sOp $ f $ Imp.var y t - Nothing -> atomicUpdateCAS space t a old bucket x $ - x <-- Imp.BinOpExp op (Imp.var x t) (Imp.var y t) - - where opHasAtomicSupport space old arr' bucket' bop = do - let atomic f = Imp.Atomic space . f old arr' bucket' - atomic <$> atomicBinOp bop - - primOrCas ops - | all isPrim ops = AtomicPrim - | otherwise = AtomicCAS - - isPrim (op, _, _, _) = isJust $ atomicBinOp op - --- If the operator functions purely on single 32-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 - old <- dPrim "old" t - atomicUpdateCAS space t arr old bucket (paramName xp) $ - compileBody' [xp] $ lambdaBody op - -atomicUpdateLocking _ op = AtomicLocking $ \locking space arrs bucket -> do - old <- dPrim "old" int32 - continue <- newVName "continue" - dPrimVol_ continue Bool - continue <-- true - - -- Correctly index into locks. - (locks', _locks_space, locks_offset) <- - fullyIndexArray (lockingArray locking) $ lockingMapping locking bucket - - -- Critical section - let try_acquire_lock = - sOp $ Imp.Atomic space $ - Imp.AtomicCmpXchg int32 old locks' locks_offset - (lockingIsUnlocked locking) (lockingToLock locking) - lock_acquired = Imp.var old int32 .==. lockingIsUnlocked locking - -- Even the releasing is done with an atomic rather than a - -- simple write, for memory coherency reasons. - release_lock = - sOp $ Imp.Atomic space $ - Imp.AtomicCmpXchg int32 old locks' locks_offset - (lockingToLock locking) (lockingToUnlock locking) - break_loop = continue <-- false - - -- Preparing parameters. It is assumed that the caller has already - -- filled the arr_params. We copy the current value to the - -- accumulator parameters. - -- - -- Note the use of 'everythingVolatile' when reading and writing the - -- buckets. This was necessary to ensure correct execution on a - -- newer NVIDIA GPU (RTX 2080). The 'volatile' modifiers likely - -- make the writes pass through the (SM-local) L1 cache, which is - -- necessary here, because we are really doing device-wide - -- synchronisation without atomics (naughty!). - let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op - bind_acc_params = - everythingVolatile $ - sComment "bind lhs" $ - forM_ (zip acc_params arrs) $ \(acc_p, arr) -> - copyDWIMFix (paramName acc_p) [] (Var arr) bucket - - let op_body = sComment "execute operation" $ - compileBody' acc_params $ lambdaBody op - - do_hist = - everythingVolatile $ - sComment "update global result" $ - zipWithM_ (writeArray bucket) arrs $ map (Var . paramName) acc_params - - fence = case space of Space "local" -> sOp $ Imp.MemFence Imp.FenceLocal - _ -> sOp $ Imp.MemFence Imp.FenceGlobal - - - -- While-loop: Try to insert your value - sWhile (Imp.var continue Bool) $ do - try_acquire_lock - sWhen lock_acquired $ do - dLParams acc_params - bind_acc_params - op_body - do_hist - fence - release_lock - break_loop - fence - where writeArray bucket arr val = copyDWIMFix arr bucket val [] - -atomicUpdateCAS :: Space -> PrimType - -> VName -> VName - -> [Imp.Exp] -> VName - -> InKernelGen () - -> InKernelGen () -atomicUpdateCAS space t arr old bucket x do_op = do - -- Code generation target: - -- - -- old = d_his[idx]; - -- do { - -- assumed = old; - -- x = do_op(assumed, y); - -- old = atomicCAS(&d_his[idx], assumed, tmp); - -- } while(assumed != old); - assumed <- dPrim "assumed" t - run_loop <- dPrimV "run_loop" 1 - - -- XXX: CUDA may generate really bad code if this is not a volatile - -- read. Unclear why. The later reads are volatile, so maybe - -- that's it. - everythingVolatile $ copyDWIMFix old [] (Var arr) bucket - - (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket - - -- While-loop: Try to insert your value - let (toBits, fromBits) = - case t of FloatType Float32 -> (\v -> Imp.FunExp "to_bits32" [v] int32, - \v -> Imp.FunExp "from_bits32" [v] t) - _ -> (id, id) - sWhile (Imp.var run_loop int32) $ do - assumed <-- Imp.var old t - x <-- Imp.var assumed t - do_op - old_bits <- dPrim "old_bits" int32 - sOp $ Imp.Atomic space $ - Imp.AtomicCmpXchg int32 old_bits arr' bucket_offset - (toBits (Imp.var assumed t)) (toBits (Imp.var x t)) - old <-- fromBits (Imp.var old_bits int32) - sWhen (toBits (Imp.var assumed t) .==. Imp.var old_bits int32) - (run_loop <-- 0) - --- | Horizontally fission a lambda that models a binary operator. -splitOp :: ASTLore lore => Lambda lore -> Maybe [(BinOp, PrimType, VName, VName)] -splitOp lam = mapM splitStm $ bodyResult $ lambdaBody lam - where n = length $ lambdaReturnType lam - splitStm (Var res) = do - Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y))) <- - find (([res]==) . patternNames . stmPattern) $ - stmsToList $ bodyStms $ lambdaBody lam - i <- Var res `elemIndex` bodyResult (lambdaBody lam) - xp <- maybeNth i $ lambdaParams lam - yp <- maybeNth (n+i) $ lambdaParams lam - guard $ paramName xp == x - guard $ paramName yp == y - Prim t <- Just $ patElemType pe - return (op, t, paramName xp, paramName yp) - splitStm _ = Nothing - -computeKernelUses :: FreeIn a => - a -> [VName] - -> CallKernelGen [Imp.KernelUse] -computeKernelUses kernel_body bound_in_kernel = do - let actually_free = freeIn kernel_body `namesSubtract` namesFromList bound_in_kernel - -- Compute the variables that we need to pass to the kernel. - nub <$> readsFromSet actually_free - -readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] -readsFromSet free = - fmap catMaybes $ - forM (namesToList free) $ \var -> do - t <- lookupType var - vtable <- getVTable - case t of - Array {} -> return Nothing - Mem (Space "local") -> return Nothing - Mem {} -> return $ Just $ Imp.MemoryUse var - Prim bt -> - isConstExp vtable (Imp.var var bt) >>= \case - Just ce -> return $ Just $ Imp.ConstUse var ce - Nothing | bt == Cert -> return Nothing - | otherwise -> return $ Just $ Imp.ScalarUse var bt - -isConstExp :: VTable KernelsMem -> Imp.Exp - -> ImpM lore r op (Maybe Imp.KernelConstExp) -isConstExp vtable size = do - fname <- askFunction - let onLeaf (Imp.ScalarVar name) _ = lookupConstExp name - onLeaf (Imp.SizeOf pt) _ = Just $ primByteSize pt - onLeaf Imp.Index{} _ = Nothing - lookupConstExp name = - constExp =<< hasExp =<< M.lookup name vtable - constExp (Op (Inner (SizeOp (GetSize key _)))) = - Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32 - constExp e = primExpFromExp lookupConstExp e - return $ replaceInPrimExpM onLeaf size - where hasExp (ArrayVar e _) = e - hasExp (ScalarVar e _) = e - hasExp (MemVar e _) = e - -computeThreadChunkSize :: SplitOrdering - -> Imp.Exp - -> Imp.Count Imp.Elements Imp.Exp - -> Imp.Count Imp.Elements Imp.Exp - -> VName - -> ImpM lore r op () -computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = do - stride' <- toExp stride - chunk_var <-- - Imp.BinOpExp (SMin Int32) - (Imp.unCount elements_per_thread) - ((Imp.unCount num_elements - thread_index) `divUp` stride') - -computeThreadChunkSize SplitContiguous thread_index elements_per_thread num_elements chunk_var = do - starting_point <- dPrimV "starting_point" $ - thread_index * Imp.unCount elements_per_thread - remaining_elements <- dPrimV "remaining_elements" $ - Imp.unCount num_elements - Imp.var starting_point int32 - - let no_remaining_elements = Imp.var remaining_elements int32 .<=. 0 - beyond_bounds = Imp.unCount num_elements .<=. Imp.var starting_point int32 - - sIf (no_remaining_elements .||. beyond_bounds) - (chunk_var <-- 0) - (sIf is_last_thread - (chunk_var <-- Imp.unCount last_thread_elements) - (chunk_var <-- Imp.unCount elements_per_thread)) - where last_thread_elements = - num_elements - Imp.elements thread_index * elements_per_thread - is_last_thread = - Imp.unCount num_elements .<. - (thread_index + 1) * Imp.unCount elements_per_thread - -kernelInitialisationSimple :: Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp - -> CallKernelGen (KernelConstants, InKernelGen ()) -kernelInitialisationSimple (Count num_groups) (Count group_size) = do - global_tid <- newVName "global_tid" - local_tid <- newVName "local_tid" - group_id <- newVName "group_tid" - wave_size <- newVName "wave_size" - inner_group_size <- newVName "group_size" - let constants = - KernelConstants - (Imp.var global_tid int32) - (Imp.var local_tid int32) - (Imp.var group_id int32) - global_tid local_tid group_id - num_groups group_size (group_size*num_groups) - (Imp.var wave_size int32) - true - mempty - - let set_constants = do - dPrim_ global_tid int32 - dPrim_ local_tid int32 - dPrim_ inner_group_size int32 - dPrim_ wave_size int32 - dPrim_ group_id int32 - - sOp (Imp.GetGlobalId global_tid 0) - sOp (Imp.GetLocalId local_tid 0) - sOp (Imp.GetLocalSize inner_group_size 0) - sOp (Imp.GetLockstepWidth wave_size) - sOp (Imp.GetGroupId group_id 0) - - return (constants, set_constants) - -isActive :: [(VName, SubExp)] -> Imp.Exp -isActive limit = case actives of - [] -> Imp.ValueExp $ BoolValue True - x:xs -> foldl (.&&.) x xs - where (is, ws) = unzip limit - actives = zipWith active is $ map (toExp' Bool) ws - active i = (Imp.var i int32 .<.) - --- | Change every memory block to be in the global address space, --- except those who are in the local memory space. This only affects --- generated code - we still need to make sure that the memory is --- actually present on the device (and dared as variables in the --- kernel). -makeAllMemoryGlobal :: CallKernelGen a -> CallKernelGen a -makeAllMemoryGlobal = - localDefaultSpace (Imp.Space "global") . localVTable (M.map globalMemory) - where globalMemory (MemVar _ entry) - | entryMemSpace entry /= Space "local" = - MemVar Nothing entry { entryMemSpace = Imp.Space "global" } - globalMemory entry = - entry - -groupReduce :: Imp.Exp - -> Lambda KernelsMem - -> [VName] - -> InKernelGen () -groupReduce w lam arrs = do - offset <- dPrim "offset" int32 - groupReduceWithOffset offset w lam arrs - -groupReduceWithOffset :: VName - -> Imp.Exp - -> Lambda KernelsMem - -> [VName] - -> InKernelGen () -groupReduceWithOffset offset w lam arrs = do - constants <- kernelConstants <$> askEnv - - let local_tid = kernelLocalThreadId constants - global_tid = kernelGlobalThreadId constants - - barrier - | all primType $ lambdaReturnType lam = sOp $ Imp.Barrier Imp.FenceLocal - | otherwise = sOp $ Imp.Barrier Imp.FenceGlobal - - readReduceArgument param arr - | Prim _ <- paramType param = do - let i = local_tid + Imp.vi32 offset - copyDWIMFix (paramName param) [] (Var arr) [i] - | otherwise = do - let i = global_tid + Imp.vi32 offset - copyDWIMFix (paramName param) [] (Var arr) [i] - - writeReduceOpResult param arr - | Prim _ <- paramType param = - copyDWIMFix arr [local_tid] (Var $ paramName param) [] - | otherwise = - return () - - let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam - - skip_waves <- dPrim "skip_waves" int32 - dLParams $ lambdaParams lam - - offset <-- 0 - - comment "participating threads read initial accumulator" $ - sWhen (local_tid .<. w) $ - zipWithM_ readReduceArgument reduce_acc_params arrs - - let do_reduce = do comment "read array element" $ - zipWithM_ readReduceArgument reduce_arr_params arrs - comment "apply reduction operation" $ - compileBody' reduce_acc_params $ lambdaBody lam - comment "write result of operation" $ - zipWithM_ writeReduceOpResult reduce_acc_params arrs - in_wave_reduce = everythingVolatile do_reduce - - wave_size = kernelWaveSize constants - group_size = kernelGroupSize constants - wave_id = local_tid `quot` wave_size - in_wave_id = local_tid - wave_id * wave_size - num_waves = (group_size + wave_size - 1) `quot` wave_size - arg_in_bounds = local_tid + Imp.var offset int32 .<. w - - doing_in_wave_reductions = - Imp.var offset int32 .<. wave_size - apply_in_in_wave_iteration = - (in_wave_id .&. (2 * Imp.var offset int32 - 1)) .==. 0 - in_wave_reductions = do - offset <-- 1 - sWhile doing_in_wave_reductions $ do - sWhen (arg_in_bounds .&&. apply_in_in_wave_iteration) - in_wave_reduce - offset <-- Imp.var offset int32 * 2 - - doing_cross_wave_reductions = - Imp.var skip_waves int32 .<. num_waves - is_first_thread_in_wave = - in_wave_id .==. 0 - wave_not_skipped = - (wave_id .&. (2 * Imp.var skip_waves int32 - 1)) .==. 0 - apply_in_cross_wave_iteration = - arg_in_bounds .&&. is_first_thread_in_wave .&&. wave_not_skipped - cross_wave_reductions = do - skip_waves <-- 1 - sWhile doing_cross_wave_reductions $ do - barrier - offset <-- Imp.var skip_waves int32 * wave_size - sWhen apply_in_cross_wave_iteration - do_reduce - skip_waves <-- Imp.var skip_waves int32 * 2 - - in_wave_reductions - cross_wave_reductions - -groupScan :: Maybe (Imp.Exp -> Imp.Exp -> Imp.Exp) - -> Imp.Exp - -> Imp.Exp - -> Lambda KernelsMem - -> [VName] - -> InKernelGen () -groupScan seg_flag arrs_full_size w lam arrs = do - constants <- kernelConstants <$> askEnv - renamed_lam <- renameLambda lam - - let ltid = kernelLocalThreadId constants - (x_params, y_params) = splitAt (length arrs) $ lambdaParams lam - - dLParams (lambdaParams lam++lambdaParams renamed_lam) - - -- The scan works by splitting the group into blocks, which are - -- scanned separately. Typically, these blocks are smaller than - -- the lockstep width, which enables barrier-free execution inside - -- them. - -- - -- We hardcode the block size here. The only requirement is that - -- it should not be less than the square root of the group size. - -- With 32, we will work on groups of size 1024 or smaller, which - -- fits every device Troels has seen. Still, it would be nicer if - -- it were a runtime parameter. Some day. - let block_size = Imp.ValueExp $ IntValue $ Int32Value 32 - simd_width = kernelWaveSize constants - block_id = ltid `quot` block_size - in_block_id = ltid - block_id * block_size - doInBlockScan seg_flag' active = - inBlockScan constants seg_flag' arrs_full_size - simd_width block_size active arrs barrier - ltid_in_bounds = ltid .<. w - array_scan = not $ all primType $ lambdaReturnType lam - barrier | array_scan = - sOp $ Imp.Barrier Imp.FenceGlobal - | otherwise = - sOp $ Imp.Barrier Imp.FenceLocal - - group_offset = kernelGroupId constants * kernelGroupSize constants - - writeBlockResult p arr - | primType $ paramType p = - copyDWIM arr [DimFix block_id] (Var $ paramName p) [] - | otherwise = - copyDWIM arr [DimFix $ group_offset + block_id] (Var $ paramName p) [] - - readPrevBlockResult p arr - | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix $ block_id - 1] - | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + block_id - 1] - - doInBlockScan seg_flag ltid_in_bounds lam - barrier - - let is_first_block = block_id .==. 0 - when array_scan $ do - sComment "save correct values for first block" $ - sWhen is_first_block $ forM_ (zip x_params arrs) $ \(x, arr) -> - unless (primType $ paramType x) $ - copyDWIM arr [DimFix $ arrs_full_size + group_offset + block_size + ltid] (Var $ paramName x) [] - - barrier - - let last_in_block = in_block_id .==. block_size - 1 - sComment "last thread of block 'i' writes its result to offset 'i'" $ - sWhen (last_in_block .&&. ltid_in_bounds) $ everythingVolatile $ - zipWithM_ writeBlockResult x_params arrs - - barrier - - let first_block_seg_flag = do - flag_true <- seg_flag - Just $ \from to -> - flag_true (from*block_size+block_size-1) (to*block_size+block_size-1) - comment - "scan the first block, after which offset 'i' contains carry-in for block 'i+1'" $ - doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam - - barrier - - when array_scan $ do - sComment "move correct values for first block back a block" $ - sWhen is_first_block $ forM_ (zip x_params arrs) $ \(x, arr) -> - unless (primType $ paramType x) $ - copyDWIM - arr [DimFix $ arrs_full_size + group_offset + ltid] - (Var arr) [DimFix $ arrs_full_size + group_offset + block_size + ltid] - - barrier - - let read_carry_in = do - forM_ (zip x_params y_params) $ \(x,y) -> - copyDWIM (paramName y) [] (Var (paramName x)) [] - zipWithM_ readPrevBlockResult x_params arrs - - y_to_x = forM_ (zip x_params y_params) $ \(x,y) -> - when (primType (paramType x)) $ - copyDWIM (paramName x) [] (Var (paramName y)) [] - - op_to_x - | Nothing <- seg_flag = - compileBody' x_params $ lambdaBody lam - | Just flag_true <- seg_flag = do - inactive <- - dPrimVE "inactive" $ flag_true (block_id*block_size-1) ltid - sWhen inactive y_to_x - when array_scan barrier - sUnless inactive $ compileBody' x_params $ lambdaBody lam - - write_final_result = - forM_ (zip x_params arrs) $ \(p, arr) -> - when (primType $ paramType p) $ - copyDWIM arr [DimFix ltid] (Var $ paramName p) [] - - sComment "carry-in for every block except the first" $ - sUnless (is_first_block .||. Imp.UnOpExp Not ltid_in_bounds) $ do - sComment "read operands" read_carry_in - sComment "perform operation" op_to_x - sComment "write final result" write_final_result - - barrier - - sComment "restore correct values for first block" $ - sWhen is_first_block $ forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) -> - if primType (paramType y) - then copyDWIM arr [DimFix ltid] (Var $ paramName y) [] - else copyDWIM (paramName x) [] (Var arr) [DimFix $ arrs_full_size + group_offset + ltid] - - barrier - -inBlockScan :: KernelConstants - -> Maybe (Imp.Exp -> Imp.Exp -> Imp.Exp) - -> Imp.Exp - -> Imp.Exp - -> Imp.Exp - -> Imp.Exp - -> [VName] - -> InKernelGen () - -> Lambda KernelsMem - -> InKernelGen () -inBlockScan constants seg_flag arrs_full_size lockstep_width block_size active arrs barrier scan_lam = everythingVolatile $ do - skip_threads <- dPrim "skip_threads" int32 - let in_block_thread_active = - Imp.var skip_threads int32 .<=. in_block_id - actual_params = lambdaParams scan_lam - (x_params, y_params) = - splitAt (length actual_params `div` 2) actual_params - y_to_x = - forM_ (zip x_params y_params) $ \(x,y) -> - when (primType (paramType x)) $ - copyDWIM (paramName x) [] (Var (paramName y)) [] - - -- Set initial y values - sComment "read input for in-block scan" $ - sWhen active $ do - zipWithM_ readInitial y_params arrs - -- Since the final result is expected to be in x_params, we may - -- need to copy it there for the first thread in the block. - sWhen (in_block_id .==. 0) y_to_x - - when array_scan barrier - - let op_to_x - | Nothing <- seg_flag = - compileBody' x_params $ lambdaBody scan_lam - | Just flag_true <- seg_flag = do - inactive <- dPrimVE "inactive" $ - flag_true (ltid-Imp.var skip_threads int32) ltid - sWhen inactive y_to_x - when array_scan barrier - sUnless inactive $ compileBody' x_params $ lambdaBody scan_lam - - maybeBarrier = sWhen (lockstep_width .<=. Imp.var skip_threads int32) - barrier - - sComment "in-block scan (hopefully no barriers needed)" $ do - skip_threads <-- 1 - sWhile (Imp.var skip_threads int32 .<. block_size) $ do - sWhen (in_block_thread_active .&&. active) $ do - sComment "read operands" $ - zipWithM_ (readParam (Imp.vi32 skip_threads)) x_params arrs - sComment "perform operation" op_to_x - - maybeBarrier - - sWhen (in_block_thread_active .&&. active) $ - sComment "write result" $ - sequence_ $ zipWith3 writeResult x_params y_params arrs - - maybeBarrier - - skip_threads <-- Imp.var skip_threads int32 * 2 - - where block_id = ltid `quot` block_size - in_block_id = ltid - block_id * block_size - ltid = kernelLocalThreadId constants - gtid = kernelGlobalThreadId constants - array_scan = not $ all primType $ lambdaReturnType scan_lam - - readInitial p arr - | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix ltid] - | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix gtid] - - readParam behind p arr - | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix $ ltid - behind] - | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix $ gtid - behind + arrs_full_size] - - writeResult x y arr - | primType $ paramType x = do - copyDWIM arr [DimFix ltid] (Var $ paramName x) [] - copyDWIM (paramName y) [] (Var $ paramName x) [] - | otherwise = - copyDWIM (paramName y) [] (Var $ paramName x) [] - -computeMapKernelGroups :: Imp.Exp -> CallKernelGen (Imp.Exp, Imp.Exp) -computeMapKernelGroups kernel_size = do - group_size <- dPrim "group_size" int32 - fname <- askFunction - let group_size_var = Imp.var group_size int32 - group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty group_size - sOp $ Imp.GetSize group_size group_size_key Imp.SizeGroup - num_groups <- dPrimV "num_groups" $ kernel_size `divUp` group_size_var - return (Imp.var num_groups int32, Imp.var group_size int32) - -simpleKernelConstants :: Imp.Exp -> String - -> CallKernelGen (KernelConstants, InKernelGen ()) -simpleKernelConstants kernel_size desc = do - thread_gtid <- newVName $ desc ++ "_gtid" - thread_ltid <- newVName $ desc ++ "_ltid" - group_id <- newVName $ desc ++ "_gid" - (num_groups, group_size) <- computeMapKernelGroups kernel_size - let set_constants = do - dPrim_ thread_gtid int32 - dPrim_ thread_ltid int32 - dPrim_ group_id int32 - sOp (Imp.GetGlobalId thread_gtid 0) - sOp (Imp.GetLocalId thread_ltid 0) - sOp (Imp.GetGroupId group_id 0) - - - return (KernelConstants - (Imp.var thread_gtid int32) (Imp.var thread_ltid int32) (Imp.var group_id int32) - thread_gtid thread_ltid group_id - num_groups group_size (group_size*num_groups) 0 - (Imp.var thread_gtid int32 .<. kernel_size) - mempty, - - set_constants) - --- | For many kernels, we may not have enough physical groups to cover --- the logical iteration space. Some groups thus have to perform --- double duty; we put an outer loop to accomplish this. The --- advantage over just launching a bazillion threads is that the cost --- of memory expansion should be proportional to the number of --- *physical* threads (hardware parallelism), not the amount of --- application parallelism. -virtualiseGroups :: SegVirt - -> Imp.Exp - -> (VName -> InKernelGen ()) - -> InKernelGen () -virtualiseGroups SegVirt required_groups m = do - constants <- kernelConstants <$> askEnv - phys_group_id <- dPrim "phys_group_id" int32 - sOp $ Imp.GetGroupId phys_group_id 0 - let iterations = (required_groups - Imp.vi32 phys_group_id) `divUp` - kernelNumGroups constants - - sFor "i" iterations $ \i -> do - m =<< dPrimV "virt_group_id" (Imp.vi32 phys_group_id + i * kernelNumGroups constants) - -- Make sure the virtual group is actually done before we let - -- another virtual group have its way with it. - sOp $ Imp.Barrier Imp.FenceGlobal -virtualiseGroups _ _ m = do - gid <- kernelGroupIdVar . kernelConstants <$> askEnv - m gid - -sKernelThread :: String - -> Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp - -> VName - -> InKernelGen () - -> CallKernelGen () -sKernelThread = sKernel threadOperations kernelGlobalThreadId - -sKernelGroup :: String - -> Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp - -> VName - -> InKernelGen () - -> CallKernelGen () -sKernelGroup = sKernel groupOperations kernelGroupId - -sKernelFailureTolerant :: Bool - -> Operations KernelsMem KernelEnv Imp.KernelOp - -> KernelConstants - -> Name - -> InKernelGen () - -> CallKernelGen () -sKernelFailureTolerant tol ops constants name m = do - HostEnv atomics <- askEnv - body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants) ops m - uses <- computeKernelUses body mempty - emit $ Imp.Op $ Imp.CallKernel Imp.Kernel - { Imp.kernelBody = body - , Imp.kernelUses = uses - , Imp.kernelNumGroups = [kernelNumGroups constants] - , Imp.kernelGroupSize = [kernelGroupSize constants] - , Imp.kernelName = name - , Imp.kernelFailureTolerant = tol - } - -sKernel :: Operations KernelsMem KernelEnv Imp.KernelOp - -> (KernelConstants -> Imp.Exp) - -> String - -> Count NumGroups Imp.Exp - -> Count GroupSize Imp.Exp - -> VName - -> InKernelGen () - -> CallKernelGen () -sKernel ops flatf name num_groups group_size v f = do - (constants, set_constants) <- kernelInitialisationSimple num_groups group_size - name' <- nameForFun $ name ++ "_" ++ show (baseTag v) - sKernelFailureTolerant False ops constants name' $ do - set_constants - dPrimV_ v $ flatf constants - f - -copyInGroup :: CopyCompiler KernelsMem KernelEnv Imp.KernelOp -copyInGroup pt destloc destslice srcloc srcslice = do - dest_space <- entryMemSpace <$> lookupMemory (memLocationName destloc) - src_space <- entryMemSpace <$> lookupMemory (memLocationName srcloc) - - case (dest_space, src_space) of - (ScalarSpace destds _, ScalarSpace srcds _) -> do - let destslice' = - replicate (length destslice - length destds) (DimFix 0) ++ - takeLast (length destds) destslice - srcslice' = - replicate (length srcslice - length srcds) (DimFix 0) ++ - takeLast (length srcds) srcslice - copyElementWise pt destloc destslice' srcloc srcslice' - - _ -> do - groupCoverSpace (sliceDims destslice) $ \is -> - copyElementWise pt - destloc (map DimFix $ fixSlice destslice is) - srcloc (map DimFix $ fixSlice srcslice is) - sOp $ Imp.Barrier Imp.FenceLocal - -threadOperations, groupOperations :: Operations KernelsMem KernelEnv Imp.KernelOp -threadOperations = - (defaultOperations compileThreadOp) - { opsCopyCompiler = copyElementWise - , opsExpCompiler = compileThreadExp - , opsStmsCompiler = \_ -> defCompileStms mempty - , opsAllocCompilers = - M.fromList [ (Space "local", allocLocal) ] - } -groupOperations = - (defaultOperations compileGroupOp) - { opsCopyCompiler = copyInGroup - , opsExpCompiler = compileGroupExp - , opsStmsCompiler = \_ -> defCompileStms mempty - , opsAllocCompilers = - M.fromList [ (Space "local", allocLocal) ] - } - --- | Perform a Replicate with a kernel. -sReplicateKernel :: VName -> SubExp -> CallKernelGen () -sReplicateKernel arr se = do - t <- subExpType se - ds <- dropLast (arrayRank t) . arrayDims <$> lookupType arr - - dims <- mapM toExp $ ds ++ arrayDims t - (constants, set_constants) <- - simpleKernelConstants (product dims) "replicate" - - fname <- askFunction - let name = keyWithEntryPoint fname $ nameFromString $ - "replicate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - is' = unflattenIndex dims $ kernelGlobalThreadId constants - - sKernelFailureTolerant True threadOperations constants name $ do - set_constants - sWhen (kernelThreadActive constants) $ - copyDWIMFix arr is' se $ drop (length ds) is' - -replicateName :: PrimType -> String -replicateName bt = "replicate_" ++ pretty bt - -replicateForType :: PrimType -> CallKernelGen Name -replicateForType bt = do - let fname = nameFromString $ "builtin#" <> replicateName bt - - exists <- hasFunction fname - unless exists $ do - mem <- newVName "mem" - num_elems <- newVName "num_elems" - val <- newVName "val" - - let params = [Imp.MemParam mem (Space "device"), - Imp.ScalarParam num_elems int32, - Imp.ScalarParam val bt] - shape = Shape [Var num_elems] - function fname [] params $ do - arr <- sArray "arr" bt shape $ ArrayIn mem $ IxFun.iota $ - map (primExpFromSubExp int32) $ shapeDims shape - sReplicateKernel arr $ Var val - - return fname - -replicateIsFill :: VName -> SubExp -> CallKernelGen (Maybe (CallKernelGen ())) -replicateIsFill arr v = do - ArrayEntry (MemLocation arr_mem arr_shape arr_ixfun) _ <- lookupArray arr - v_t <- subExpType v - case v_t of - Prim v_t' - | IxFun.isLinear arr_ixfun -> return $ Just $ do - fname <- replicateForType v_t' - emit $ Imp.Call [] fname - [Imp.MemArg arr_mem, - Imp.ExpArg $ product $ map (toExp' int32) arr_shape, - Imp.ExpArg $ toExp' v_t' v] - _ -> return Nothing - --- | Perform a Replicate with a kernel. -sReplicate :: VName -> SubExp -> CallKernelGen () -sReplicate arr se = do - -- If the replicate is of a particularly common and simple form - -- (morally a memset()/fill), then we use a common function. - is_fill <- replicateIsFill arr se - - case is_fill of - Just m -> m - Nothing -> sReplicateKernel arr se - --- | Perform an Iota with a kernel. -sIotaKernel :: VName -> Imp.Exp -> Imp.Exp -> Imp.Exp -> IntType - -> CallKernelGen () -sIotaKernel arr n x s et = do - destloc <- entryArrayLocation <$> lookupArray arr - (constants, set_constants) <- simpleKernelConstants n "iota" - - fname <- askFunction - let name = keyWithEntryPoint fname $ nameFromString $ - "iota_" ++ pretty et ++ "_" ++ - show (baseTag $ kernelGlobalThreadIdVar constants) - - sKernelFailureTolerant True threadOperations constants name $ do - set_constants - let gtid = kernelGlobalThreadId constants - sWhen (kernelThreadActive constants) $ do - (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] - - emit $ - Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ - Imp.sExt et gtid * s + x - -iotaName :: IntType -> String -iotaName bt = "iota_" ++ pretty bt - -iotaForType :: IntType -> CallKernelGen Name -iotaForType bt = do - let fname = nameFromString $ "builtin#" <> iotaName bt - - exists <- hasFunction fname - unless exists $ do - mem <- newVName "mem" - n <- newVName "n" - x <- newVName "x" - s <- newVName "s" - - let params = [Imp.MemParam mem (Space "device"), - Imp.ScalarParam n int32, - Imp.ScalarParam x $ IntType bt, - Imp.ScalarParam s $ IntType bt] - shape = Shape [Var n] - n' = Imp.vi32 n - x' = Imp.var x $ IntType bt - s' = Imp.var s $ IntType bt - - function fname [] params $ do - arr <- sArray "arr" (IntType bt) shape $ ArrayIn mem $ IxFun.iota $ - map (primExpFromSubExp int32) $ shapeDims shape - sIotaKernel arr n' x' s' bt - - return fname - --- | Perform an Iota with a kernel. -sIota :: VName -> Imp.Exp -> Imp.Exp -> Imp.Exp -> IntType - -> CallKernelGen () -sIota arr n x s et = do - ArrayEntry (MemLocation arr_mem _ arr_ixfun) _ <- lookupArray arr - if IxFun.isLinear arr_ixfun then do - fname <- iotaForType et - emit $ Imp.Call [] fname - [Imp.MemArg arr_mem, Imp.ExpArg n, Imp.ExpArg x, Imp.ExpArg s] - else sIotaKernel arr n x s et - -sCopy :: CopyCompiler KernelsMem HostEnv Imp.HostOp -sCopy bt - destloc@(MemLocation destmem _ _) destslice - srcloc@(MemLocation srcmem _ _) srcslice - = do - -- Note that the shape of the destination and the source are - -- necessarily the same. - let shape = sliceDims srcslice - kernel_size = product shape - - (constants, set_constants) <- simpleKernelConstants kernel_size "copy" - - fname <- askFunction - let name = keyWithEntryPoint fname $ nameFromString $ - "copy_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - - sKernelFailureTolerant True threadOperations constants name $ do - set_constants - - let gtid = kernelGlobalThreadId constants - dest_is = unflattenIndex shape gtid - src_is = dest_is - - (_, destspace, destidx) <- - fullyIndexArray' destloc $ fixSlice destslice dest_is - (_, srcspace, srcidx) <- - fullyIndexArray' srcloc $ fixSlice srcslice src_is - - sWhen (gtid .<. kernel_size) $ emit $ - Imp.Write destmem destidx bt destspace Imp.Nonvolatile $ - Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile - -compileGroupResult :: SegSpace - -> PatElem KernelsMem -> KernelResult - -> InKernelGen () - -compileGroupResult _ pe (TileReturns [(w,per_group_elems)] what) = do - n <- toExp . arraySize 0 =<< lookupType what - - constants <- kernelConstants <$> askEnv - let ltid = kernelLocalThreadId constants - offset = toExp' int32 per_group_elems * kernelGroupId constants - - -- Avoid loop for the common case where each thread is statically - -- known to write at most one element. - localOps threadOperations $ - if toExp' int32 per_group_elems == kernelGroupSize constants - then sWhen (offset + ltid .<. toExp' int32 w) $ - copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid] - else - sFor "i" (n `divUp` kernelGroupSize constants) $ \i -> do - j <- fmap Imp.vi32 $ dPrimV "j" $ - kernelGroupSize constants * i + ltid - sWhen (j .<. n) $ copyDWIMFix (patElemName pe) [j + offset] (Var what) [j] - -compileGroupResult space pe (TileReturns dims what) = do - let gids = map fst $ unSegSpace space - out_tile_sizes = map (toExp' int32 . snd) dims - group_is = zipWith (*) (map Imp.vi32 gids) out_tile_sizes - local_is <- localThreadIDs $ map snd dims - is_for_thread <- mapM (dPrimV "thread_out_index") $ zipWith (+) group_is local_is - - localOps threadOperations $ - sWhen (isActive $ zip is_for_thread $ map fst dims) $ - copyDWIMFix (patElemName pe) (map Imp.vi32 is_for_thread) (Var what) local_is - -compileGroupResult space pe (Returns _ what) = do - constants <- kernelConstants <$> askEnv - in_local_memory <- arrayInLocalMemory what - let gids = map (Imp.vi32 . fst) $ unSegSpace space - - if not in_local_memory then - localOps threadOperations $ - sWhen (kernelLocalThreadId constants .==. 0) $ - copyDWIMFix (patElemName pe) gids what [] - else - -- If the result of the group is an array in local memory, we - -- store it by collective copying among all the threads of the - -- group. TODO: also do this if the array is in global memory - -- (but this is a bit more tricky, synchronisation-wise). - copyDWIMFix (patElemName pe) gids what [] - -compileGroupResult _ _ WriteReturns{} = - compilerLimitationS "compileGroupResult: WriteReturns not handled yet." - -compileGroupResult _ _ ConcatReturns{} = - compilerLimitationS "compileGroupResult: ConcatReturns not handled yet." - -compileThreadResult :: SegSpace - -> PatElem KernelsMem -> KernelResult - -> InKernelGen () - -compileThreadResult space pe (Returns _ what) = do - let is = map (Imp.vi32 . fst) $ unSegSpace space - copyDWIMFix (patElemName pe) is what [] - -compileThreadResult _ pe (ConcatReturns SplitContiguous _ per_thread_elems what) = do - constants <- kernelConstants <$> askEnv - let offset = toExp' int32 per_thread_elems * kernelGlobalThreadId constants - n <- toExp' int32 . arraySize 0 <$> lookupType what - copyDWIM (patElemName pe) [DimSlice offset n 1] (Var what) [] - -compileThreadResult _ pe (ConcatReturns (SplitStrided stride) _ _ what) = do - offset <- kernelGlobalThreadId . kernelConstants <$> askEnv - n <- toExp' int32 . arraySize 0 <$> lookupType what - copyDWIM (patElemName pe) [DimSlice offset n $ toExp' int32 stride] (Var what) [] - -compileThreadResult _ pe (WriteReturns rws _arr dests) = do - constants <- kernelConstants <$> askEnv - rws' <- mapM toExp rws - forM_ dests $ \(slice, e) -> do - slice' <- mapM (traverse toExp) slice - let condInBounds (DimFix i) rw = - 0 .<=. i .&&. i .<. rw - condInBounds (DimSlice i n s) rw = - 0 .<=. i .&&. i+n*s .<. rw - write = foldl (.&&.) (kernelThreadActive constants) $ - zipWith condInBounds slice' rws' - sWhen write $ copyDWIM (patElemName pe) slice' e [] - -compileThreadResult _ _ TileReturns{} = - compilerBugS "compileThreadResult: TileReturns unhandled." - -arrayInLocalMemory :: SubExp -> InKernelGen Bool -arrayInLocalMemory (Var name) = do - res <- lookupVar name - case res of - ArrayVar _ entry -> - (Space "local"==) . entryMemSpace <$> - lookupMemory (memLocationName (entryArrayLocation entry)) - _ -> return False -arrayInLocalMemory Constant{} = return False + +module Futhark.CodeGen.ImpGen.Kernels.Base + ( KernelConstants (..), + keyWithEntryPoint, + CallKernelGen, + InKernelGen, + HostEnv (..), + KernelEnv (..), + computeThreadChunkSize, + groupReduce, + groupScan, + isActive, + sKernelThread, + sKernelGroup, + sReplicate, + sIota, + sCopy, + compileThreadResult, + compileGroupResult, + virtualiseGroups, + groupLoop, + kernelLoop, + groupCoverSpace, + precomputeSegOpIDs, + atomicUpdateLocking, + AtomicBinOp, + Locking (..), + AtomicUpdate (..), + DoAtomicUpdate, + ) +where + +import Control.Monad.Except +import Data.List (elemIndex, find, nub, zip4) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import qualified Futhark.CodeGen.ImpCode.Kernels as Imp +import Futhark.CodeGen.ImpGen +import Futhark.Error +import Futhark.IR.KernelsMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Transform.Rename +import Futhark.Util (chunks, dropLast, mapAccumLM, maybeNth, takeLast) +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) + +newtype HostEnv = HostEnv + {hostAtomics :: AtomicBinOp} + +data KernelEnv = KernelEnv + { kernelAtomics :: AtomicBinOp, + kernelConstants :: KernelConstants + } + +type CallKernelGen = ImpM KernelsMem HostEnv Imp.HostOp + +type InKernelGen = ImpM KernelsMem KernelEnv Imp.KernelOp + +data KernelConstants = KernelConstants + { kernelGlobalThreadId :: Imp.TExp Int32, + kernelLocalThreadId :: Imp.TExp Int32, + kernelGroupId :: Imp.TExp Int32, + kernelGlobalThreadIdVar :: VName, + kernelLocalThreadIdVar :: VName, + kernelGroupIdVar :: VName, + kernelNumGroups :: Imp.TExp Int64, + kernelGroupSize :: Imp.TExp Int64, + kernelNumThreads :: Imp.TExp Int32, + kernelWaveSize :: Imp.TExp Int32, + kernelThreadActive :: Imp.TExp Bool, + -- | A mapping from dimensions of nested SegOps to already + -- computed local thread IDs. + kernelLocalIdMap :: M.Map [SubExp] [Imp.TExp Int32] + } + +segOpSizes :: Stms KernelsMem -> S.Set [SubExp] +segOpSizes = onStms + where + onStms = foldMap (onExp . stmExp) + onExp (Op (Inner (SegOp op))) = + S.singleton $ map snd $ unSegSpace $ segSpace op + onExp (If _ tbranch fbranch _) = + onStms (bodyStms tbranch) <> onStms (bodyStms fbranch) + onExp (DoLoop _ _ _ body) = + onStms (bodyStms body) + onExp _ = mempty + +precomputeSegOpIDs :: Stms KernelsMem -> InKernelGen a -> InKernelGen a +precomputeSegOpIDs stms m = do + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + new_ids <- M.fromList <$> mapM (mkMap ltid) (S.toList (segOpSizes stms)) + let f env = + env + { kernelConstants = + (kernelConstants env) {kernelLocalIdMap = new_ids} + } + localEnv f m + where + mkMap ltid dims = do + let dims' = map (sExt32 . toInt64Exp) dims + ids' <- mapM (dPrimVE "ltid_pre") $ unflattenIndex dims' ltid + return (dims, ids') + +keyWithEntryPoint :: Maybe Name -> Name -> Name +keyWithEntryPoint fname key = + nameFromString $ maybe "" ((++ ".") . nameToString) fname ++ nameToString key + +allocLocal :: AllocCompiler KernelsMem r Imp.KernelOp +allocLocal mem size = + sOp $ Imp.LocalAlloc mem size + +kernelAlloc :: + Pattern KernelsMem -> + SubExp -> + Space -> + InKernelGen () +kernelAlloc (Pattern _ [_]) _ ScalarSpace {} = + -- Handled by the declaration of the memory block, which is then + -- translated to an actual scalar variable during C code generation. + return () +kernelAlloc (Pattern _ [mem]) size (Space "local") = + allocLocal (patElemName mem) $ Imp.bytes $ toInt64Exp size +kernelAlloc (Pattern _ [mem]) _ _ = + compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." +kernelAlloc dest _ _ = + error $ "Invalid target for in-kernel allocation: " ++ show dest + +splitSpace :: + (ToExp w, ToExp i, ToExp elems_per_thread) => + Pattern KernelsMem -> + SplitOrdering -> + w -> + i -> + elems_per_thread -> + ImpM lore r op () +splitSpace (Pattern [] [size]) o w i elems_per_thread = do + num_elements <- Imp.elements . TPrimExp <$> toExp w + let i' = toInt64Exp i + elems_per_thread' <- Imp.elements . TPrimExp <$> toExp elems_per_thread + computeThreadChunkSize o i' elems_per_thread' num_elements (mkTV (patElemName size) int64) +splitSpace pat _ _ _ _ = + error $ "Invalid target for splitSpace: " ++ pretty pat + +compileThreadExp :: ExpCompiler KernelsMem KernelEnv Imp.KernelOp +compileThreadExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = + forM_ (zip [0 ..] es) $ \(i, e) -> + copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] +compileThreadExp dest e = + defCompileExp dest e + +-- | Assign iterations of a for-loop to all threads in the kernel. +-- The passed-in function is invoked with the (symbolic) iteration. +-- 'threadOperations' will be in effect in the body. For +-- multidimensional loops, use 'groupCoverSpace'. +kernelLoop :: + IntExp t => + Imp.TExp t -> + Imp.TExp t -> + Imp.TExp t -> + (Imp.TExp t -> InKernelGen ()) -> + InKernelGen () +kernelLoop tid num_threads n f = + localOps threadOperations $ + if n == num_threads + then f tid + else do + -- Compute how many elements this thread is responsible for. + -- Formula: (n - tid) / num_threads (rounded up). + let elems_for_this = (n - tid) `divUp` num_threads + + sFor "i" elems_for_this $ \i -> f $ i * num_threads + tid + +-- | Assign iterations of a for-loop to threads in the workgroup. The +-- passed-in function is invoked with the (symbolic) iteration. For +-- multidimensional loops, use 'groupCoverSpace'. +groupLoop :: + Imp.TExp Int64 -> + (Imp.TExp Int64 -> InKernelGen ()) -> + InKernelGen () +groupLoop n f = do + constants <- kernelConstants <$> askEnv + kernelLoop + (sExt64 $ kernelLocalThreadId constants) + (kernelGroupSize constants) + n + f + +-- | Iterate collectively though a multidimensional space, such that +-- all threads in the group participate. The passed-in function is +-- invoked with a (symbolic) point in the index space. +groupCoverSpace :: + [Imp.TExp Int64] -> + ([Imp.TExp Int64] -> InKernelGen ()) -> + InKernelGen () +groupCoverSpace ds f = + groupLoop (product ds) $ f . unflattenIndex ds + +compileGroupExp :: ExpCompiler KernelsMem KernelEnv Imp.KernelOp +-- The static arrays stuff does not work inside kernels. +compileGroupExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = + forM_ (zip [0 ..] es) $ \(i, e) -> + copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] +compileGroupExp (Pattern _ [dest]) (BasicOp (Replicate ds se)) = do + let ds' = map toInt64Exp $ shapeDims ds + groupCoverSpace ds' $ \is -> + copyDWIMFix (patElemName dest) is se (drop (shapeRank ds) is) + sOp $ Imp.Barrier Imp.FenceLocal +compileGroupExp (Pattern _ [dest]) (BasicOp (Iota n e s it)) = do + n' <- toExp n + e' <- toExp e + s' <- toExp s + groupLoop (TPrimExp n') $ \i' -> do + x <- + dPrimV "x" $ + TPrimExp $ + BinOpExp (Add it OverflowUndef) e' $ + BinOpExp (Mul it OverflowUndef) (untyped i') s' + copyDWIMFix (patElemName dest) [i'] (Var (tvVar x)) [] + sOp $ Imp.Barrier Imp.FenceLocal + +-- When generating code for a scalar in-place update, we must make +-- sure that only one thread performs the write. When writing an +-- array, the group-level copy code will take care of doing the right +-- thing. +compileGroupExp (Pattern _ [pe]) (BasicOp (Update _ slice se)) + | null $ sliceDims slice = do + sOp $ Imp.Barrier Imp.FenceLocal + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + sWhen (ltid .==. 0) $ + copyDWIM (patElemName pe) (map (fmap toInt64Exp) slice) se [] + sOp $ Imp.Barrier Imp.FenceLocal +compileGroupExp dest e = + defCompileExp dest e + +sanityCheckLevel :: SegLevel -> InKernelGen () +sanityCheckLevel SegThread {} = return () +sanityCheckLevel SegGroup {} = + error "compileGroupOp: unexpected group-level SegOp." + +localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64] +localThreadIDs dims = do + ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv + let dims' = map toInt64Exp dims + maybe (unflattenIndex dims' ltid) (map sExt64) + . M.lookup dims + . kernelLocalIdMap + . kernelConstants + <$> askEnv + +compileGroupSpace :: SegLevel -> SegSpace -> InKernelGen () +compileGroupSpace lvl space = do + sanityCheckLevel lvl + let (ltids, dims) = unzip $ unSegSpace space + zipWithM_ dPrimV_ ltids =<< localThreadIDs dims + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + dPrimV_ (segFlat space) ltid + +-- Construct the necessary lock arrays for an intra-group histogram. +prepareIntraGroupSegHist :: + Count GroupSize SubExp -> + [HistOp KernelsMem] -> + InKernelGen [[Imp.TExp Int64] -> InKernelGen ()] +prepareIntraGroupSegHist group_size = + fmap snd . mapAccumLM onOp Nothing + where + onOp l op = do + constants <- kernelConstants <$> askEnv + atomicBinOp <- kernelAtomics <$> askEnv + + let local_subhistos = histDest op + + case (l, atomicUpdateLocking atomicBinOp $ histOp op) of + (_, AtomicPrim f) -> return (l, f (Space "local") local_subhistos) + (_, AtomicCAS f) -> return (l, f (Space "local") local_subhistos) + (Just l', AtomicLocking f) -> return (l, f l' (Space "local") local_subhistos) + (Nothing, AtomicLocking f) -> do + locks <- newVName "locks" + + let num_locks = toInt64Exp $ unCount group_size + dims = map toInt64Exp $ shapeDims (histShape op) ++ [histWidth op] + l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims) + locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness + + locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local" + dArray locks int32 (arrayShape locks_t) $ + ArrayIn locks_mem $ + IxFun.iota $ + map pe64 $ arrayDims locks_t + + sComment "All locks start out unlocked" $ + groupCoverSpace [kernelGroupSize constants] $ \is -> + copyDWIMFix locks is (intConst Int32 0) [] + + return (Just l', f l' (Space "local") local_subhistos) + +whenActive :: SegLevel -> SegSpace -> InKernelGen () -> InKernelGen () +whenActive lvl space m + | SegNoVirtFull <- segVirt lvl = m + | otherwise = sWhen (isActive $ unSegSpace space) m + +compileGroupOp :: OpCompiler KernelsMem KernelEnv Imp.KernelOp +compileGroupOp pat (Alloc size space) = + kernelAlloc pat size space +compileGroupOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = + splitSpace pat o w i elems_per_thread +compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do + void $ compileGroupSpace lvl space + + whenActive lvl space $ + localOps threadOperations $ + compileStms mempty (kernelBodyStms body) $ + zipWithM_ (compileThreadResult space) (patternElements pat) $ + kernelBodyResult body + + sOp $ Imp.ErrorSync Imp.FenceLocal +compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do + compileGroupSpace lvl space + let (ltids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + + whenActive lvl space $ + compileStms mempty (kernelBodyStms body) $ + forM_ (zip (patternNames pat) $ kernelBodyResult body) $ \(dest, res) -> + copyDWIMFix + dest + (map Imp.vi64 ltids) + (kernelResultSubExp res) + [] + + sOp $ Imp.ErrorSync Imp.FenceLocal + + let segment_size = last dims' + crossesSegment from to = + (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size) + + -- groupScan needs to treat the scan output as a one-dimensional + -- array of scan elements, so we invent some new flattened arrays + -- here. XXX: this assumes that the original index function is just + -- row-major, but does not actually verify it. + dims_flat <- dPrimV "dims_flat" $ product dims' + let flattened pe = do + MemLocation mem _ _ <- + entryArrayLocation <$> lookupArray (patElemName pe) + let pe_t = typeOf pe + arr_dims = Var (tvVar dims_flat) : drop (length dims') (arrayDims pe_t) + sArray + (baseString (patElemName pe) ++ "_flat") + (elemType pe_t) + (Shape arr_dims) + $ ArrayIn mem $ IxFun.iota $ map pe64 arr_dims + + num_scan_results = sum $ map (length . segBinOpNeutral) scans + + arrs_flat <- mapM flattened $ take num_scan_results $ patternElements pat + + forM_ scans $ \scan -> do + let scan_op = segBinOpLambda scan + groupScan (Just crossesSegment) (product dims') (product dims') scan_op arrs_flat +compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do + compileGroupSpace lvl space + + let (ltids, dims) = unzip $ unSegSpace space + (red_pes, map_pes) = + splitAt (segBinOpResults ops) $ patternElements pat + + dims' = map toInt64Exp dims + + mkTempArr t = + sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local" + + tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops + let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs + + whenActive lvl space $ + compileStms mempty (kernelBodyStms body) $ do + let (red_res, map_res) = + splitAt (segBinOpResults ops) $ kernelBodyResult body + forM_ (zip tmp_arrs red_res) $ \(dest, res) -> + copyDWIMFix dest (map Imp.vi64 ltids) (kernelResultSubExp res) [] + zipWithM_ (compileThreadResult space) map_pes map_res + + sOp $ Imp.ErrorSync Imp.FenceLocal + + case dims' of + -- Nonsegmented case (or rather, a single segment) - this we can + -- handle directly with a group-level reduction. + [dim'] -> do + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> + groupReduce (sExt32 dim') (segBinOpLambda op) tmps + + sOp $ Imp.ErrorSync Imp.FenceLocal + + forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> + copyDWIMFix (patElemName pe) [] (Var arr) [0] + _ -> do + -- Segmented intra-group reductions are turned into (regular) + -- segmented scans. It is possible that this can be done + -- better, but at least this approach is simple. + + -- groupScan operates on flattened arrays. This does not + -- involve copying anything; merely playing with the index + -- function. + dims_flat <- dPrimV "dims_flat" $ product dims' + let flatten arr = do + ArrayEntry arr_loc pt <- lookupArray arr + let flat_shape = + Shape $ + Var (tvVar dims_flat) : + drop (length ltids) (memLocationShape arr_loc) + sArray "red_arr_flat" pt flat_shape $ + ArrayIn (memLocationName arr_loc) $ + IxFun.iota $ map pe64 $ shapeDims flat_shape + + let segment_size = last dims' + crossesSegment from to = + (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size) + + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do + tmps_flat <- mapM flatten tmps + groupScan + (Just crossesSegment) + (product dims') + (product dims') + (segBinOpLambda op) + tmps_flat + + sOp $ Imp.ErrorSync Imp.FenceLocal + + forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> + copyDWIM + (patElemName pe) + [] + (Var arr) + (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' -1]) + + sOp $ Imp.Barrier Imp.FenceLocal +compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do + compileGroupSpace lvl space + let ltids = map fst $ unSegSpace space + + -- We don't need the red_pes, because it is guaranteed by our type + -- rules that they occupy the same memory as the destinations for + -- the ops. + let num_red_res = length ops + sum (map (length . histNeutral) ops) + (_red_pes, map_pes) = + splitAt num_red_res $ patternElements pat + + ops' <- prepareIntraGroupSegHist (segGroupSize lvl) ops + + -- Ensure that all locks have been initialised. + sOp $ Imp.Barrier Imp.FenceLocal + + whenActive lvl space $ + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody + (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res + zipWithM_ (compileThreadResult space) map_pes map_res + + let vs_per_op = chunks (map (length . histDest) ops) red_vs + + forM_ (zip4 red_is vs_per_op ops' ops) $ + \(bin, op_vs, do_op, HistOp dest_w _ _ _ shape lam) -> do + let bin' = toInt64Exp bin + dest_w' = toInt64Exp dest_w + bin_in_bounds = 0 .<=. bin' .&&. bin' .<. dest_w' + bin_is = map Imp.vi64 (init ltids) ++ [bin'] + vs_params = takeLast (length op_vs) $ lambdaParams lam + + sComment "perform atomic updates" $ + sWhen bin_in_bounds $ do + dLParams $ lambdaParams lam + sLoopNest shape $ \is -> do + forM_ (zip vs_params op_vs) $ \(p, v) -> + copyDWIMFix (paramName p) [] v is + do_op (bin_is ++ is) + + sOp $ Imp.ErrorSync Imp.FenceLocal +compileGroupOp pat _ = + compilerBugS $ "compileGroupOp: cannot compile rhs of binding " ++ pretty pat + +compileThreadOp :: OpCompiler KernelsMem KernelEnv Imp.KernelOp +compileThreadOp pat (Alloc size space) = + kernelAlloc pat size space +compileThreadOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = + splitSpace pat o w i elems_per_thread +compileThreadOp pat _ = + compilerBugS $ "compileThreadOp: cannot compile rhs of binding " ++ pretty pat + +-- | Locking strategy used for an atomic update. +data Locking = Locking + { -- | Array containing the lock. + lockingArray :: VName, + -- | Value for us to consider the lock free. + lockingIsUnlocked :: Imp.TExp Int32, + -- | What to write when we lock it. + lockingToLock :: Imp.TExp Int32, + -- | What to write when we unlock it. + lockingToUnlock :: Imp.TExp Int32, + -- | A transformation from the logical lock index to the + -- physical position in the array. This can also be used + -- to make the lock array smaller. + lockingMapping :: [Imp.TExp Int64] -> [Imp.TExp Int64] + } + +-- | A function for generating code for an atomic update. Assumes +-- that the bucket is in-bounds. +type DoAtomicUpdate lore r = + Space -> [VName] -> [Imp.TExp Int64] -> ImpM lore r Imp.KernelOp () + +-- | The mechanism that will be used for performing the atomic update. +-- Approximates how efficient it will be. Ordered from most to least +-- efficient. +data AtomicUpdate lore r + = -- | Supported directly by primitive. + AtomicPrim (DoAtomicUpdate lore r) + | -- | Can be done by efficient swaps. + AtomicCAS (DoAtomicUpdate lore r) + | -- | Requires explicit locking. + AtomicLocking (Locking -> DoAtomicUpdate lore r) + +-- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? +type AtomicBinOp = + BinOp -> + Maybe (VName -> VName -> Count Imp.Elements (Imp.TExp Int64) -> Imp.Exp -> Imp.AtomicOp) + +-- | Do an atomic update corresponding to a binary operator lambda. +atomicUpdateLocking :: + AtomicBinOp -> + Lambda KernelsMem -> + AtomicUpdate KernelsMem KernelEnv +atomicUpdateLocking atomicBinOp lam + | Just ops_and_ts <- splitOp lam, + all (\(_, t, _, _) -> primBitSize t == 32) 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. + forM_ (zip arrs ops_and_ts) $ \(a, (op, t, x, y)) -> do + -- Common variables. + old <- dPrim "old" t + + (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket + + case opHasAtomicSupport space (tvVar old) arr' bucket_offset op of + Just f -> sOp $ f $ Imp.var y t + Nothing -> + atomicUpdateCAS space t a (tvVar old) bucket x $ + x <~~ Imp.BinOpExp op (Imp.var x t) (Imp.var y t) + where + opHasAtomicSupport space old arr' bucket' bop = do + let atomic f = Imp.Atomic space . f old arr' bucket' + atomic <$> atomicBinOp bop + + primOrCas ops + | all isPrim ops = AtomicPrim + | otherwise = AtomicCAS + + isPrim (op, _, _, _) = isJust $ atomicBinOp op + +-- If the operator functions purely on single 32-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 + old <- dPrim "old" t + atomicUpdateCAS space t arr (tvVar old) bucket (paramName xp) $ + compileBody' [xp] $ lambdaBody op +atomicUpdateLocking _ op = AtomicLocking $ \locking space arrs bucket -> do + old <- dPrim "old" int32 + continue <- dPrimVol "continue" Bool true + + -- Correctly index into locks. + (locks', _locks_space, locks_offset) <- + fullyIndexArray (lockingArray locking) $ lockingMapping locking bucket + + -- Critical section + let try_acquire_lock = + sOp $ + Imp.Atomic space $ + Imp.AtomicCmpXchg + int32 + (tvVar old) + locks' + locks_offset + (untyped $ lockingIsUnlocked locking) + (untyped $ lockingToLock locking) + lock_acquired = tvExp old .==. lockingIsUnlocked locking + -- Even the releasing is done with an atomic rather than a + -- simple write, for memory coherency reasons. + release_lock = + sOp $ + Imp.Atomic space $ + Imp.AtomicCmpXchg + int32 + (tvVar old) + locks' + locks_offset + (untyped $ lockingToLock locking) + (untyped $ lockingToUnlock locking) + break_loop = continue <-- false + + -- Preparing parameters. It is assumed that the caller has already + -- filled the arr_params. We copy the current value to the + -- accumulator parameters. + -- + -- Note the use of 'everythingVolatile' when reading and writing the + -- buckets. This was necessary to ensure correct execution on a + -- newer NVIDIA GPU (RTX 2080). The 'volatile' modifiers likely + -- make the writes pass through the (SM-local) L1 cache, which is + -- necessary here, because we are really doing device-wide + -- synchronisation without atomics (naughty!). + let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op + bind_acc_params = + everythingVolatile $ + sComment "bind lhs" $ + forM_ (zip acc_params arrs) $ \(acc_p, arr) -> + copyDWIMFix (paramName acc_p) [] (Var arr) bucket + + let op_body = + sComment "execute operation" $ + compileBody' acc_params $ lambdaBody op + + do_hist = + everythingVolatile $ + sComment "update global result" $ + zipWithM_ (writeArray bucket) arrs $ map (Var . paramName) acc_params + + fence = case space of + Space "local" -> sOp $ Imp.MemFence Imp.FenceLocal + _ -> sOp $ Imp.MemFence Imp.FenceGlobal + + -- While-loop: Try to insert your value + sWhile (tvExp continue) $ do + try_acquire_lock + sWhen lock_acquired $ do + dLParams acc_params + bind_acc_params + op_body + do_hist + fence + release_lock + break_loop + fence + where + writeArray bucket arr val = copyDWIMFix arr bucket val [] + +atomicUpdateCAS :: + Space -> + PrimType -> + VName -> + VName -> + [Imp.TExp Int64] -> + VName -> + InKernelGen () -> + InKernelGen () +atomicUpdateCAS space t arr old bucket x do_op = do + -- Code generation target: + -- + -- old = d_his[idx]; + -- do { + -- assumed = old; + -- x = do_op(assumed, y); + -- old = atomicCAS(&d_his[idx], assumed, tmp); + -- } while(assumed != old); + assumed <- tvVar <$> dPrim "assumed" t + run_loop <- dPrimV "run_loop" true + + -- XXX: CUDA may generate really bad code if this is not a volatile + -- read. Unclear why. The later reads are volatile, so maybe + -- that's it. + everythingVolatile $ copyDWIMFix old [] (Var arr) bucket + + (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket + + -- While-loop: Try to insert your value + let (toBits, fromBits) = + case t of + FloatType Float32 -> + ( \v -> Imp.FunExp "to_bits32" [v] int32, + \v -> Imp.FunExp "from_bits32" [v] t + ) + _ -> (id, id) + sWhile (tvExp run_loop) $ do + assumed <~~ Imp.var old t + x <~~ Imp.var assumed t + do_op + old_bits <- dPrim "old_bits" int32 + sOp $ + Imp.Atomic space $ + Imp.AtomicCmpXchg + int32 + (tvVar old_bits) + 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) + +-- | Horizontally fission a lambda that models a binary operator. +splitOp :: ASTLore lore => Lambda lore -> Maybe [(BinOp, PrimType, VName, VName)] +splitOp lam = mapM splitStm $ bodyResult $ lambdaBody lam + where + n = length $ lambdaReturnType lam + splitStm (Var res) = do + Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y))) <- + find (([res] ==) . patternNames . stmPattern) $ + stmsToList $ bodyStms $ lambdaBody lam + i <- Var res `elemIndex` bodyResult (lambdaBody lam) + xp <- maybeNth i $ lambdaParams lam + yp <- maybeNth (n + i) $ lambdaParams lam + guard $ paramName xp == x + guard $ paramName yp == y + Prim t <- Just $ patElemType pe + return (op, t, paramName xp, paramName yp) + splitStm _ = Nothing + +computeKernelUses :: + FreeIn a => + a -> + [VName] -> + CallKernelGen [Imp.KernelUse] +computeKernelUses kernel_body bound_in_kernel = do + let actually_free = freeIn kernel_body `namesSubtract` namesFromList bound_in_kernel + -- Compute the variables that we need to pass to the kernel. + nub <$> readsFromSet actually_free + +readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] +readsFromSet free = + fmap catMaybes $ + forM (namesToList free) $ \var -> do + t <- lookupType var + vtable <- getVTable + case t of + Array {} -> return Nothing + Mem (Space "local") -> return Nothing + Mem {} -> return $ Just $ Imp.MemoryUse var + Prim bt -> + isConstExp vtable (Imp.var var bt) >>= \case + Just ce -> return $ Just $ Imp.ConstUse var ce + Nothing + | bt == Cert -> return Nothing + | otherwise -> return $ Just $ Imp.ScalarUse var bt + +isConstExp :: + VTable KernelsMem -> + Imp.Exp -> + ImpM lore r op (Maybe Imp.KernelConstExp) +isConstExp vtable size = do + fname <- askFunction + let onLeaf (Imp.ScalarVar name) _ = lookupConstExp name + onLeaf (Imp.SizeOf pt) _ = Just $ ValueExp $ IntValue $ Int32Value $ primByteSize pt + onLeaf Imp.Index {} _ = Nothing + lookupConstExp name = + constExp =<< hasExp =<< M.lookup name vtable + constExp (Op (Inner (SizeOp (GetSize key _)))) = + Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32 + constExp e = primExpFromExp lookupConstExp e + return $ replaceInPrimExpM onLeaf size + where + hasExp (ArrayVar e _) = e + hasExp (ScalarVar e _) = e + hasExp (MemVar e _) = e + +computeThreadChunkSize :: + SplitOrdering -> + Imp.TExp Int64 -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + TV Int64 -> + ImpM lore r op () +computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = + chunk_var + <-- sMin64 + (Imp.unCount elements_per_thread) + ((Imp.unCount num_elements - thread_index) `divUp` toInt64Exp stride) +computeThreadChunkSize SplitContiguous thread_index elements_per_thread num_elements chunk_var = do + starting_point <- + dPrimV "starting_point" $ + thread_index * Imp.unCount elements_per_thread + remaining_elements <- + dPrimV "remaining_elements" $ + Imp.unCount num_elements - tvExp starting_point + + let no_remaining_elements = tvExp remaining_elements .<=. 0 + beyond_bounds = Imp.unCount num_elements .<=. tvExp starting_point + + sIf + (no_remaining_elements .||. beyond_bounds) + (chunk_var <-- 0) + ( sIf + is_last_thread + (chunk_var <-- Imp.unCount last_thread_elements) + (chunk_var <-- Imp.unCount elements_per_thread) + ) + where + last_thread_elements = + num_elements - Imp.elements thread_index * elements_per_thread + is_last_thread = + Imp.unCount num_elements + .<. (thread_index + 1) * Imp.unCount elements_per_thread + +kernelInitialisationSimple :: + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + CallKernelGen (KernelConstants, InKernelGen ()) +kernelInitialisationSimple (Count num_groups) (Count group_size) = do + global_tid <- newVName "global_tid" + local_tid <- newVName "local_tid" + group_id <- newVName "group_tid" + wave_size <- newVName "wave_size" + inner_group_size <- newVName "group_size" + let constants = + KernelConstants + (Imp.vi32 global_tid) + (Imp.vi32 local_tid) + (Imp.vi32 group_id) + global_tid + local_tid + group_id + num_groups + group_size + (sExt32 (group_size * num_groups)) + (Imp.vi32 wave_size) + true + mempty + + let set_constants = do + dPrim_ global_tid int32 + dPrim_ local_tid int32 + dPrim_ inner_group_size int64 + dPrim_ wave_size int32 + dPrim_ group_id int32 + + sOp (Imp.GetGlobalId global_tid 0) + sOp (Imp.GetLocalId local_tid 0) + sOp (Imp.GetLocalSize inner_group_size 0) + sOp (Imp.GetLockstepWidth wave_size) + sOp (Imp.GetGroupId group_id 0) + + return (constants, set_constants) + +isActive :: [(VName, SubExp)] -> Imp.TExp Bool +isActive limit = case actives of + [] -> true + x : xs -> foldl (.&&.) x xs + where + (is, ws) = unzip limit + actives = zipWith active is $ map toInt64Exp ws + active i = (Imp.vi64 i .<.) + +-- | Change every memory block to be in the global address space, +-- except those who are in the local memory space. This only affects +-- generated code - we still need to make sure that the memory is +-- actually present on the device (and dared as variables in the +-- kernel). +makeAllMemoryGlobal :: CallKernelGen a -> CallKernelGen a +makeAllMemoryGlobal = + localDefaultSpace (Imp.Space "global") . localVTable (M.map globalMemory) + where + globalMemory (MemVar _ entry) + | entryMemSpace entry /= Space "local" = + MemVar Nothing entry {entryMemSpace = Imp.Space "global"} + globalMemory entry = + entry + +groupReduce :: + Imp.TExp Int32 -> + Lambda KernelsMem -> + [VName] -> + InKernelGen () +groupReduce w lam arrs = do + offset <- dPrim "offset" int32 + groupReduceWithOffset offset w lam arrs + +groupReduceWithOffset :: + TV Int32 -> + Imp.TExp Int32 -> + Lambda KernelsMem -> + [VName] -> + InKernelGen () +groupReduceWithOffset offset w lam arrs = do + constants <- kernelConstants <$> askEnv + + let local_tid = kernelLocalThreadId constants + global_tid = kernelGlobalThreadId constants + + barrier + | all primType $ lambdaReturnType lam = sOp $ Imp.Barrier Imp.FenceLocal + | otherwise = sOp $ Imp.Barrier Imp.FenceGlobal + + readReduceArgument param arr + | Prim _ <- paramType param = do + let i = local_tid + tvExp offset + copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] + | otherwise = do + let i = global_tid + tvExp offset + copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] + + writeReduceOpResult param arr + | Prim _ <- paramType param = + copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) [] + | otherwise = + return () + + let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam + + skip_waves <- dPrimV "skip_waves" (1 :: Imp.TExp Int32) + dLParams $ lambdaParams lam + + offset <-- (0 :: Imp.TExp Int32) + + comment "participating threads read initial accumulator" $ + sWhen (local_tid .<. w) $ + zipWithM_ readReduceArgument reduce_acc_params arrs + + let do_reduce = do + comment "read array element" $ + zipWithM_ readReduceArgument reduce_arr_params arrs + comment "apply reduction operation" $ + compileBody' reduce_acc_params $ lambdaBody lam + comment "write result of operation" $ + zipWithM_ writeReduceOpResult reduce_acc_params arrs + in_wave_reduce = everythingVolatile do_reduce + + wave_size = kernelWaveSize constants + group_size = kernelGroupSize constants + wave_id = local_tid `quot` wave_size + in_wave_id = local_tid - wave_id * wave_size + num_waves = (sExt32 group_size + wave_size - 1) `quot` wave_size + arg_in_bounds = local_tid + tvExp offset .<. w + + doing_in_wave_reductions = + tvExp offset .<. wave_size + apply_in_in_wave_iteration = + (in_wave_id .&. (2 * tvExp offset - 1)) .==. 0 + in_wave_reductions = do + offset <-- (1 :: Imp.TExp Int32) + sWhile doing_in_wave_reductions $ do + sWhen + (arg_in_bounds .&&. apply_in_in_wave_iteration) + in_wave_reduce + offset <-- tvExp offset * 2 + + doing_cross_wave_reductions = + tvExp skip_waves .<. num_waves + is_first_thread_in_wave = + in_wave_id .==. 0 + wave_not_skipped = + (wave_id .&. (2 * tvExp skip_waves - 1)) .==. 0 + apply_in_cross_wave_iteration = + arg_in_bounds .&&. is_first_thread_in_wave .&&. wave_not_skipped + cross_wave_reductions = + sWhile doing_cross_wave_reductions $ do + barrier + offset <-- tvExp skip_waves * wave_size + sWhen + apply_in_cross_wave_iteration + do_reduce + skip_waves <-- tvExp skip_waves * 2 + + in_wave_reductions + cross_wave_reductions + +groupScan :: + Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Lambda KernelsMem -> + [VName] -> + InKernelGen () +groupScan seg_flag arrs_full_size w lam arrs = do + constants <- kernelConstants <$> askEnv + renamed_lam <- renameLambda lam + + let ltid32 = kernelLocalThreadId constants + ltid = sExt64 ltid32 + (x_params, y_params) = splitAt (length arrs) $ lambdaParams lam + + dLParams (lambdaParams lam ++ lambdaParams renamed_lam) + + ltid_in_bounds <- dPrimVE "ltid_in_bounds" $ ltid .<. w + + -- The scan works by splitting the group into blocks, which are + -- scanned separately. Typically, these blocks are smaller than + -- the lockstep width, which enables barrier-free execution inside + -- them. + -- + -- We hardcode the block size here. The only requirement is that + -- it should not be less than the square root of the group size. + -- With 32, we will work on groups of size 1024 or smaller, which + -- fits every device Troels has seen. Still, it would be nicer if + -- it were a runtime parameter. Some day. + let block_size = 32 + simd_width = kernelWaveSize constants + block_id = ltid32 `quot` block_size + in_block_id = ltid32 - block_id * block_size + doInBlockScan seg_flag' active = + inBlockScan + constants + seg_flag' + arrs_full_size + simd_width + block_size + active + arrs + barrier + array_scan = not $ all primType $ lambdaReturnType lam + barrier + | array_scan = + sOp $ Imp.Barrier Imp.FenceGlobal + | otherwise = + sOp $ Imp.Barrier Imp.FenceLocal + + group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants + + writeBlockResult p arr + | primType $ paramType p = + copyDWIM arr [DimFix $ sExt64 block_id] (Var $ paramName p) [] + | otherwise = + copyDWIM arr [DimFix $ group_offset + sExt64 block_id] (Var $ paramName p) [] + + readPrevBlockResult p arr + | primType $ paramType p = + copyDWIM (paramName p) [] (Var arr) [DimFix $ sExt64 block_id - 1] + | otherwise = + copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + sExt64 block_id - 1] + + doInBlockScan seg_flag ltid_in_bounds lam + barrier + + let is_first_block = block_id .==. 0 + when array_scan $ do + sComment "save correct values for first block" $ + sWhen is_first_block $ + forM_ (zip x_params arrs) $ \(x, arr) -> + unless (primType $ paramType x) $ + copyDWIM arr [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) [] + + barrier + + let last_in_block = in_block_id .==. block_size - 1 + sComment "last thread of block 'i' writes its result to offset 'i'" $ + sWhen (last_in_block .&&. ltid_in_bounds) $ + everythingVolatile $ + zipWithM_ writeBlockResult x_params arrs + + barrier + + let first_block_seg_flag = do + flag_true <- seg_flag + Just $ \from to -> + flag_true (from * block_size + block_size -1) (to * block_size + block_size -1) + comment + "scan the first block, after which offset 'i' contains carry-in for block 'i+1'" + $ doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam + + barrier + + when array_scan $ do + sComment "move correct values for first block back a block" $ + sWhen is_first_block $ + forM_ (zip x_params arrs) $ \(x, arr) -> + unless (primType $ paramType x) $ + copyDWIM + arr + [DimFix $ arrs_full_size + group_offset + ltid] + (Var arr) + [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] + + barrier + + let read_carry_in = do + forM_ (zip x_params y_params) $ \(x, y) -> + copyDWIM (paramName y) [] (Var (paramName x)) [] + zipWithM_ readPrevBlockResult x_params arrs + + y_to_x = forM_ (zip x_params y_params) $ \(x, y) -> + when (primType (paramType x)) $ + copyDWIM (paramName x) [] (Var (paramName y)) [] + + op_to_x + | Nothing <- seg_flag = + compileBody' x_params $ lambdaBody lam + | Just flag_true <- seg_flag = do + inactive <- + dPrimVE "inactive" $ flag_true (block_id * block_size -1) ltid32 + sWhen inactive y_to_x + when array_scan barrier + sUnless inactive $ compileBody' x_params $ lambdaBody lam + + write_final_result = + forM_ (zip x_params arrs) $ \(p, arr) -> + when (primType $ paramType p) $ + copyDWIM arr [DimFix ltid] (Var $ paramName p) [] + + sComment "carry-in for every block except the first" $ + sUnless (is_first_block .||. bNot ltid_in_bounds) $ do + sComment "read operands" read_carry_in + sComment "perform operation" op_to_x + sComment "write final result" write_final_result + + barrier + + sComment "restore correct values for first block" $ + sWhen is_first_block $ + forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) -> + if primType (paramType y) + then copyDWIM arr [DimFix ltid] (Var $ paramName y) [] + else copyDWIM (paramName x) [] (Var arr) [DimFix $ arrs_full_size + group_offset + ltid] + + barrier + +inBlockScan :: + KernelConstants -> + Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> + Imp.TExp Int64 -> + Imp.TExp Int32 -> + Imp.TExp Int32 -> + Imp.TExp Bool -> + [VName] -> + InKernelGen () -> + Lambda KernelsMem -> + InKernelGen () +inBlockScan constants seg_flag arrs_full_size lockstep_width block_size active arrs barrier scan_lam = everythingVolatile $ do + skip_threads <- dPrim "skip_threads" int32 + let in_block_thread_active = + tvExp skip_threads .<=. in_block_id + actual_params = lambdaParams scan_lam + (x_params, y_params) = + splitAt (length actual_params `div` 2) actual_params + y_to_x = + forM_ (zip x_params y_params) $ \(x, y) -> + when (primType (paramType x)) $ + copyDWIM (paramName x) [] (Var (paramName y)) [] + + -- Set initial y values + sComment "read input for in-block scan" $ + sWhen active $ do + zipWithM_ readInitial y_params arrs + -- Since the final result is expected to be in x_params, we may + -- need to copy it there for the first thread in the block. + sWhen (in_block_id .==. 0) y_to_x + + when array_scan barrier + + let op_to_x + | Nothing <- seg_flag = + compileBody' x_params $ lambdaBody scan_lam + | Just flag_true <- seg_flag = do + inactive <- + dPrimVE "inactive" $ + flag_true (ltid32 - tvExp skip_threads) ltid32 + sWhen inactive y_to_x + when array_scan barrier + sUnless inactive $ compileBody' x_params $ lambdaBody scan_lam + + maybeBarrier = + sWhen + (lockstep_width .<=. tvExp skip_threads) + barrier + + sComment "in-block scan (hopefully no barriers needed)" $ do + skip_threads <-- 1 + sWhile (tvExp skip_threads .<. block_size) $ do + sWhen (in_block_thread_active .&&. active) $ do + sComment "read operands" $ + zipWithM_ (readParam (sExt64 $ tvExp skip_threads)) x_params arrs + sComment "perform operation" op_to_x + + maybeBarrier + + sWhen (in_block_thread_active .&&. active) $ + sComment "write result" $ + sequence_ $ zipWith3 writeResult x_params y_params arrs + + maybeBarrier + + skip_threads <-- tvExp skip_threads * 2 + where + block_id = ltid32 `quot` block_size + in_block_id = ltid32 - block_id * block_size + ltid32 = kernelLocalThreadId constants + ltid = sExt64 ltid32 + gtid = sExt64 $ kernelGlobalThreadId constants + array_scan = not $ all primType $ lambdaReturnType scan_lam + + readInitial p arr + | primType $ paramType p = + copyDWIM (paramName p) [] (Var arr) [DimFix ltid] + | otherwise = + copyDWIM (paramName p) [] (Var arr) [DimFix gtid] + + readParam behind p arr + | primType $ paramType p = + copyDWIM (paramName p) [] (Var arr) [DimFix $ ltid - behind] + | otherwise = + copyDWIM (paramName p) [] (Var arr) [DimFix $ gtid - behind + arrs_full_size] + + writeResult x y arr + | primType $ paramType x = do + copyDWIM arr [DimFix ltid] (Var $ paramName x) [] + copyDWIM (paramName y) [] (Var $ paramName x) [] + | otherwise = + copyDWIM (paramName y) [] (Var $ paramName x) [] + +computeMapKernelGroups :: Imp.TExp Int64 -> CallKernelGen (Imp.TExp Int64, Imp.TExp Int64) +computeMapKernelGroups kernel_size = do + group_size <- dPrim "group_size" int64 + fname <- askFunction + let group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty $ tvVar group_size + sOp $ Imp.GetSize (tvVar group_size) group_size_key Imp.SizeGroup + num_groups <- dPrimV "num_groups" $ kernel_size `divUp` tvExp group_size + return (tvExp num_groups, tvExp group_size) + +simpleKernelConstants :: + Imp.TExp Int64 -> + String -> + CallKernelGen (KernelConstants, InKernelGen ()) +simpleKernelConstants kernel_size desc = do + thread_gtid <- newVName $ desc ++ "_gtid" + thread_ltid <- newVName $ desc ++ "_ltid" + group_id <- newVName $ desc ++ "_gid" + (num_groups, group_size) <- computeMapKernelGroups kernel_size + let set_constants = do + dPrim_ thread_gtid int32 + dPrim_ thread_ltid int32 + dPrim_ group_id int32 + sOp (Imp.GetGlobalId thread_gtid 0) + sOp (Imp.GetLocalId thread_ltid 0) + sOp (Imp.GetGroupId group_id 0) + + return + ( KernelConstants + (Imp.vi32 thread_gtid) + (Imp.vi32 thread_ltid) + (Imp.vi32 group_id) + thread_gtid + thread_ltid + group_id + num_groups + group_size + (sExt32 (group_size * num_groups)) + 0 + (Imp.vi64 thread_gtid .<. kernel_size) + mempty, + set_constants + ) + +-- | For many kernels, we may not have enough physical groups to cover +-- the logical iteration space. Some groups thus have to perform +-- double duty; we put an outer loop to accomplish this. The +-- advantage over just launching a bazillion threads is that the cost +-- of memory expansion should be proportional to the number of +-- *physical* threads (hardware parallelism), not the amount of +-- application parallelism. +virtualiseGroups :: + SegVirt -> + Imp.TExp Int32 -> + (Imp.TExp Int32 -> InKernelGen ()) -> + InKernelGen () +virtualiseGroups SegVirt required_groups m = do + constants <- kernelConstants <$> askEnv + phys_group_id <- dPrim "phys_group_id" int32 + sOp $ Imp.GetGroupId (tvVar phys_group_id) 0 + let iterations = + (required_groups - tvExp phys_group_id) + `divUp` sExt32 (kernelNumGroups constants) + + sFor "i" iterations $ \i -> do + m . tvExp + =<< dPrimV + "virt_group_id" + (tvExp phys_group_id + i * sExt32 (kernelNumGroups constants)) + -- Make sure the virtual group is actually done before we let + -- another virtual group have its way with it. + sOp $ Imp.Barrier Imp.FenceGlobal +virtualiseGroups _ _ m = do + gid <- kernelGroupIdVar . kernelConstants <$> askEnv + m $ Imp.vi32 gid + +sKernelThread :: + String -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + VName -> + InKernelGen () -> + CallKernelGen () +sKernelThread = sKernel threadOperations kernelGlobalThreadId + +sKernelGroup :: + String -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + VName -> + InKernelGen () -> + CallKernelGen () +sKernelGroup = sKernel groupOperations kernelGroupId + +sKernelFailureTolerant :: + Bool -> + Operations KernelsMem KernelEnv Imp.KernelOp -> + KernelConstants -> + Name -> + InKernelGen () -> + CallKernelGen () +sKernelFailureTolerant tol ops constants name m = do + HostEnv atomics <- askEnv + body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants) ops m + uses <- computeKernelUses body mempty + emit $ + Imp.Op $ + Imp.CallKernel + Imp.Kernel + { Imp.kernelBody = body, + Imp.kernelUses = uses, + Imp.kernelNumGroups = [untyped $ kernelNumGroups constants], + Imp.kernelGroupSize = [untyped $ kernelGroupSize constants], + Imp.kernelName = name, + Imp.kernelFailureTolerant = tol + } + +sKernel :: + Operations KernelsMem KernelEnv Imp.KernelOp -> + (KernelConstants -> Imp.TExp Int32) -> + String -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + VName -> + InKernelGen () -> + CallKernelGen () +sKernel ops flatf name num_groups group_size v f = do + (constants, set_constants) <- kernelInitialisationSimple num_groups group_size + name' <- nameForFun $ name ++ "_" ++ show (baseTag v) + sKernelFailureTolerant False ops constants name' $ do + set_constants + dPrimV_ v $ flatf constants + f + +copyInGroup :: CopyCompiler KernelsMem KernelEnv Imp.KernelOp +copyInGroup pt destloc destslice srcloc srcslice = do + dest_space <- entryMemSpace <$> lookupMemory (memLocationName destloc) + src_space <- entryMemSpace <$> lookupMemory (memLocationName srcloc) + + case (dest_space, src_space) of + (ScalarSpace destds _, ScalarSpace srcds _) -> do + let destslice' = + replicate (length destslice - length destds) (DimFix 0) + ++ takeLast (length destds) destslice + srcslice' = + replicate (length srcslice - length srcds) (DimFix 0) + ++ takeLast (length srcds) srcslice + copyElementWise pt destloc destslice' srcloc srcslice' + _ -> do + groupCoverSpace (sliceDims destslice) $ \is -> + copyElementWise + pt + destloc + (map DimFix $ fixSlice destslice is) + srcloc + (map DimFix $ fixSlice srcslice is) + sOp $ Imp.Barrier Imp.FenceLocal + +threadOperations, groupOperations :: Operations KernelsMem KernelEnv Imp.KernelOp +threadOperations = + (defaultOperations compileThreadOp) + { opsCopyCompiler = copyElementWise, + opsExpCompiler = compileThreadExp, + opsStmsCompiler = \_ -> defCompileStms mempty, + opsAllocCompilers = + M.fromList [(Space "local", allocLocal)] + } +groupOperations = + (defaultOperations compileGroupOp) + { opsCopyCompiler = copyInGroup, + opsExpCompiler = compileGroupExp, + opsStmsCompiler = \_ -> defCompileStms mempty, + opsAllocCompilers = + M.fromList [(Space "local", allocLocal)] + } + +-- | Perform a Replicate with a kernel. +sReplicateKernel :: VName -> SubExp -> CallKernelGen () +sReplicateKernel arr se = do + t <- subExpType se + ds <- dropLast (arrayRank t) . arrayDims <$> lookupType arr + + let dims = map toInt64Exp $ ds ++ arrayDims t + (constants, set_constants) <- + simpleKernelConstants (product $ map sExt64 dims) "replicate" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "replicate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) + is' = unflattenIndex dims $ sExt64 $ kernelGlobalThreadId constants + + sKernelFailureTolerant True threadOperations constants name $ do + set_constants + sWhen (kernelThreadActive constants) $ + copyDWIMFix arr is' se $ drop (length ds) is' + +replicateName :: PrimType -> String +replicateName bt = "replicate_" ++ pretty bt + +replicateForType :: PrimType -> CallKernelGen Name +replicateForType bt = do + let fname = nameFromString $ "builtin#" <> replicateName bt + + exists <- hasFunction fname + unless exists $ do + mem <- newVName "mem" + num_elems <- newVName "num_elems" + val <- newVName "val" + + let params = + [ Imp.MemParam mem (Space "device"), + Imp.ScalarParam num_elems int32, + Imp.ScalarParam val bt + ] + shape = Shape [Var num_elems] + function fname [] params $ do + arr <- + sArray "arr" bt shape $ + ArrayIn mem $ + IxFun.iota $ + map pe64 $ shapeDims shape + sReplicateKernel arr $ Var val + + return fname + +replicateIsFill :: VName -> SubExp -> CallKernelGen (Maybe (CallKernelGen ())) +replicateIsFill arr v = do + ArrayEntry (MemLocation arr_mem arr_shape arr_ixfun) _ <- lookupArray arr + v_t <- subExpType v + case v_t of + Prim v_t' + | IxFun.isLinear arr_ixfun -> return $ + Just $ do + fname <- replicateForType v_t' + emit $ + Imp.Call + [] + fname + [ Imp.MemArg arr_mem, + Imp.ExpArg $ untyped $ product $ map toInt64Exp arr_shape, + Imp.ExpArg $ toExp' v_t' v + ] + _ -> return Nothing + +-- | Perform a Replicate with a kernel. +sReplicate :: VName -> SubExp -> CallKernelGen () +sReplicate arr se = do + -- If the replicate is of a particularly common and simple form + -- (morally a memset()/fill), then we use a common function. + is_fill <- replicateIsFill arr se + + case is_fill of + Just m -> m + Nothing -> sReplicateKernel arr se + +-- | Perform an Iota with a kernel. +sIotaKernel :: + VName -> + Imp.TExp Int64 -> + Imp.Exp -> + Imp.Exp -> + IntType -> + CallKernelGen () +sIotaKernel arr n x s et = do + destloc <- entryArrayLocation <$> lookupArray arr + (constants, set_constants) <- simpleKernelConstants n "iota" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "iota_" ++ pretty et ++ "_" + ++ show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernelFailureTolerant True threadOperations constants name $ do + set_constants + let gtid = sExt64 $ kernelGlobalThreadId constants + sWhen (kernelThreadActive constants) $ do + (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] + + emit $ + Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ + BinOpExp + (Add et OverflowWrap) + (BinOpExp (Mul et OverflowWrap) (Imp.sExt et $ untyped gtid) s) + x + +iotaName :: IntType -> String +iotaName bt = "iota_" ++ pretty bt + +iotaForType :: IntType -> CallKernelGen Name +iotaForType bt = do + let fname = nameFromString $ "builtin#" <> iotaName bt + + exists <- hasFunction fname + unless exists $ do + mem <- newVName "mem" + n <- newVName "n" + x <- newVName "x" + s <- newVName "s" + + let params = + [ Imp.MemParam mem (Space "device"), + Imp.ScalarParam n int32, + Imp.ScalarParam x $ IntType bt, + Imp.ScalarParam s $ IntType bt + ] + shape = Shape [Var n] + n' = Imp.vi64 n + x' = Imp.var x $ IntType bt + s' = Imp.var s $ IntType bt + + function fname [] params $ do + arr <- + sArray "arr" (IntType bt) shape $ + ArrayIn mem $ + IxFun.iota $ + map pe64 $ shapeDims shape + sIotaKernel arr (sExt64 n') x' s' bt + + return fname + +-- | Perform an Iota with a kernel. +sIota :: + VName -> + Imp.TExp Int64 -> + Imp.Exp -> + Imp.Exp -> + IntType -> + CallKernelGen () +sIota arr n x s et = do + ArrayEntry (MemLocation arr_mem _ arr_ixfun) _ <- lookupArray arr + if IxFun.isLinear arr_ixfun + then do + fname <- iotaForType et + emit $ + Imp.Call + [] + fname + [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s] + else sIotaKernel arr n x s et + +sCopy :: CopyCompiler KernelsMem HostEnv Imp.HostOp +sCopy + bt + destloc@(MemLocation destmem _ _) + destslice + srcloc@(MemLocation srcmem _ _) + srcslice = + do + -- Note that the shape of the destination and the source are + -- necessarily the same. + let shape = sliceDims srcslice + kernel_size = product shape + + (constants, set_constants) <- simpleKernelConstants kernel_size "copy" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "copy_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernelFailureTolerant True threadOperations constants name $ do + set_constants + + let gtid = sExt64 $ kernelGlobalThreadId constants + dest_is = unflattenIndex shape gtid + src_is = dest_is + + (_, destspace, destidx) <- + fullyIndexArray' destloc $ fixSlice destslice dest_is + (_, srcspace, srcidx) <- + fullyIndexArray' srcloc $ fixSlice srcslice src_is + + sWhen (gtid .<. kernel_size) $ + emit $ + Imp.Write destmem destidx bt destspace Imp.Nonvolatile $ + Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile + +compileGroupResult :: + SegSpace -> + PatElem KernelsMem -> + KernelResult -> + InKernelGen () +compileGroupResult _ pe (TileReturns [(w, per_group_elems)] what) = do + n <- toInt64Exp . arraySize 0 <$> lookupType what + + constants <- kernelConstants <$> askEnv + let ltid = sExt64 $ kernelLocalThreadId constants + offset = + toInt64Exp per_group_elems + * sExt64 (kernelGroupId constants) + + -- Avoid loop for the common case where each thread is statically + -- known to write at most one element. + localOps threadOperations $ + if toInt64Exp per_group_elems == kernelGroupSize constants + then + sWhen (ltid + offset .<. toInt64Exp w) $ + copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid] + else sFor "i" (n `divUp` kernelGroupSize constants) $ \i -> do + j <- dPrimVE "j" $ kernelGroupSize constants * i + ltid + sWhen (j + offset .<. toInt64Exp w) $ + copyDWIMFix (patElemName pe) [j + offset] (Var what) [j] +compileGroupResult space pe (TileReturns dims what) = do + let gids = map fst $ unSegSpace space + out_tile_sizes = map (toInt64Exp . snd) dims + group_is = zipWith (*) (map Imp.vi64 gids) out_tile_sizes + local_is <- localThreadIDs $ map snd dims + is_for_thread <- + mapM (dPrimV "thread_out_index") $ + zipWith (+) group_is local_is + + 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 (Returns _ what) = do + constants <- kernelConstants <$> askEnv + in_local_memory <- arrayInLocalMemory what + let gids = map (Imp.vi64 . fst) $ unSegSpace space + + if not in_local_memory + then + localOps threadOperations $ + sWhen (kernelLocalThreadId constants .==. 0) $ + copyDWIMFix (patElemName pe) gids what [] + else -- If the result of the group is an array in local memory, we + -- store it by collective copying among all the threads of the + -- group. TODO: also do this if the array is in global memory + -- (but this is a bit more tricky, synchronisation-wise). + copyDWIMFix (patElemName pe) gids what [] +compileGroupResult _ _ WriteReturns {} = + compilerLimitationS "compileGroupResult: WriteReturns not handled yet." +compileGroupResult _ _ ConcatReturns {} = + compilerLimitationS "compileGroupResult: ConcatReturns not handled yet." + +compileThreadResult :: + SegSpace -> + PatElem KernelsMem -> + KernelResult -> + InKernelGen () +compileThreadResult space pe (Returns _ what) = do + let is = map (Imp.vi64 . fst) $ unSegSpace space + copyDWIMFix (patElemName pe) is what [] +compileThreadResult _ pe (ConcatReturns SplitContiguous _ per_thread_elems what) = do + constants <- kernelConstants <$> askEnv + let offset = + toInt64Exp per_thread_elems + * sExt64 (kernelGlobalThreadId constants) + n <- toInt64Exp . arraySize 0 <$> lookupType what + copyDWIM (patElemName pe) [DimSlice offset n 1] (Var what) [] +compileThreadResult _ pe (ConcatReturns (SplitStrided stride) _ _ what) = do + offset <- sExt64 . kernelGlobalThreadId . kernelConstants <$> askEnv + n <- toInt64Exp . arraySize 0 <$> lookupType what + copyDWIM (patElemName pe) [DimSlice offset n $ toInt64Exp stride] (Var what) [] +compileThreadResult _ pe (WriteReturns rws _arr dests) = do + constants <- kernelConstants <$> askEnv + let rws' = map toInt64Exp rws + forM_ dests $ \(slice, e) -> do + let slice' = map (fmap toInt64Exp) slice + condInBounds (DimFix i) rw = + 0 .<=. i .&&. i .<. rw + condInBounds (DimSlice i n s) rw = + 0 .<=. i .&&. i + n * s .<. rw + write = + foldl (.&&.) (kernelThreadActive constants) $ + zipWith condInBounds slice' rws' + sWhen write $ copyDWIM (patElemName pe) slice' e [] +compileThreadResult _ _ TileReturns {} = + compilerBugS "compileThreadResult: TileReturns unhandled." + +arrayInLocalMemory :: SubExp -> InKernelGen Bool +arrayInLocalMemory (Var name) = do + res <- lookupVar name + case res of + ArrayVar _ entry -> + (Space "local" ==) . entryMemSpace + <$> lookupMemory (memLocationName (entryArrayLocation entry)) + _ -> return False +arrayInLocalMemory Constant {} = return False 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 @@ -1,979 +1,1143 @@ -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE OverloadedStrings #-} --- | Our compilation strategy for 'SegHist' is based around avoiding --- bin conflicts. We do this by splitting the input into chunks, and --- for each chunk computing a single subhistogram. Then we combine --- the subhistograms using an ordinary segmented reduction ('SegRed'). --- --- There are some branches around to efficiently handle the case where --- we use only a single subhistogram (because it's large), so that we --- respect the asymptotics, and do not copy the destination array. --- --- We also use a heuristic strategy for computing subhistograms in --- local memory when possible. Given: --- --- H: total size of histograms in bytes, including any lock arrays. --- --- G: group size --- --- T: number of bytes of local memory each thread can be given without --- impacting occupancy (determined experimentally, e.g. 32). --- --- LMAX: maximum amount of local memory per workgroup (hard limit). --- --- We wish to compute: --- --- COOP: cooperation level (number of threads per subhistogram) --- --- LH: number of local memory subhistograms --- --- We do this as: --- --- COOP = ceil(H / T) --- LH = ceil((G*T)/H) --- if COOP <= G && H <= LMAX then --- use local memory --- else --- use global memory - -module Futhark.CodeGen.ImpGen.Kernels.SegHist - ( compileSegHist ) - where - -import Control.Monad.Except -import Data.Maybe -import Data.List (foldl', genericLength, zip4, zip6) - -import Prelude hiding (quot, rem) - -import Futhark.MonadFreshNames -import Futhark.IR.KernelsMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Pass.ExplicitAllocations() -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen -import Futhark.CodeGen.ImpGen.Kernels.SegRed (compileSegRed') -import Futhark.CodeGen.ImpGen.Kernels.Base -import Futhark.Util.IntegralExp (divUp, quot, rem) -import Futhark.Util (chunks, mapAccumLM, maxinum, splitFromEnd, takeLast) -import Futhark.Construct (fullSliceNum) - -data SubhistosInfo = SubhistosInfo { subhistosArray :: VName - , subhistosAlloc :: CallKernelGen () - } - -data SegHistSlug = SegHistSlug - { slugOp :: HistOp KernelsMem - , slugNumSubhistos :: VName - , slugSubhistos :: [SubhistosInfo] - , slugAtomicUpdate :: AtomicUpdate KernelsMem KernelEnv - } - -histoSpaceUsage :: HistOp KernelsMem - -> Imp.Count Imp.Bytes Imp.Exp -histoSpaceUsage op = - fmap (sExt Int32) $ sum $ - map (typeSize . - (`arrayOfRow` histWidth op) . - (`arrayOfShape` histShape op)) $ - lambdaReturnType $ histOp op - --- | Figure out how much memory is needed per histogram, both --- segmented and unsegmented,, and compute some other auxiliary --- information. -computeHistoUsage :: SegSpace - -> HistOp KernelsMem - -> CallKernelGen (Imp.Count Imp.Bytes Imp.Exp, - Imp.Count Imp.Bytes Imp.Exp, - SegHistSlug) -computeHistoUsage space op = do - let segment_dims = init $ unSegSpace space - num_segments = length segment_dims - - -- Create names for the intermediate array memory blocks, - -- memory block sizes, arrays, and number of subhistograms. - num_subhistos <- dPrim "num_subhistos" int32 - subhisto_infos <- forM (zip (histDest op) (histNeutral op)) $ \(dest, ne) -> do - dest_t <- lookupType dest - dest_mem <- entryArrayLocation <$> lookupArray dest - - subhistos_mem <- - sDeclareMem (baseString dest ++ "_subhistos_mem") (Space "device") - - let subhistos_shape = Shape (map snd segment_dims++[Var num_subhistos]) <> - stripDims num_segments (arrayShape dest_t) - subhistos_membind = ArrayIn subhistos_mem $ IxFun.iota $ - map (primExpFromSubExp int32) $ shapeDims subhistos_shape - subhistos <- sArray (baseString dest ++ "_subhistos") - (elemType dest_t) subhistos_shape subhistos_membind - - return $ SubhistosInfo subhistos $ do - let unitHistoCase = - emit $ - Imp.SetMem subhistos_mem (memLocationName dest_mem) $ - Space "device" - - multiHistoCase = do - let num_elems = foldl' (*) (Imp.var num_subhistos int32) $ - map (toExp' int32) $ arrayDims dest_t - - let subhistos_mem_size = - Imp.bytes $ - Imp.unCount (Imp.elements num_elems `Imp.withElemType` elemType dest_t) - - sAlloc_ subhistos_mem subhistos_mem_size $ Space "device" - sReplicate subhistos ne - subhistos_t <- lookupType subhistos - let slice = fullSliceNum (map (toExp' int32) $ arrayDims subhistos_t) $ - map (unitSlice 0 . toExp' int32 . snd) segment_dims ++ - [DimFix 0] - sUpdate subhistos slice $ Var dest - - sIf (Imp.var num_subhistos int32 .==. 1) unitHistoCase multiHistoCase - - let h = histoSpaceUsage op - segmented_h = h * product (map (Imp.bytes . toExp' int32) $ init $ segSpaceDims space) - - atomics <- hostAtomics <$> askEnv - - return (h, - segmented_h, - SegHistSlug op num_subhistos subhisto_infos $ - atomicUpdateLocking atomics $ histOp op) - -prepareAtomicUpdateGlobal :: Maybe Locking -> [VName] -> SegHistSlug - -> CallKernelGen (Maybe Locking, - [Imp.Exp] -> InKernelGen ()) -prepareAtomicUpdateGlobal l dests slug = - -- We need a separate lock array if the operators are not all of a - -- particularly simple form that permits pure atomic operations. - case (l, slugAtomicUpdate slug) of - (_, AtomicPrim f) -> return (l, f (Space "global") dests) - (_, AtomicCAS f) -> return (l, f (Space "global") dests) - (Just l', AtomicLocking f) -> return (l, f l' (Space "global") dests) - (Nothing, AtomicLocking f) -> do - -- The number of locks used here is too low, but since we are - -- currently forced to inline a huge list, I'm keeping it down - -- for now. Some quick experiments suggested that it has little - -- impact anyway (maybe the locking case is just too slow). - -- - -- A fun solution would also be to use a simple hashing - -- algorithm to ensure good distribution of locks. - let num_locks = 100151 - dims = map (toExp' int32) $ - shapeDims (histShape (slugOp slug)) ++ - [ Var (slugNumSubhistos slug) - , histWidth (slugOp slug)] - locks <- - sStaticArray "hist_locks" (Space "device") int32 $ - Imp.ArrayZeros num_locks - let l' = Locking locks 0 1 0 (pure . (`rem` fromIntegral num_locks) . flattenIndex dims) - return (Just l', f l' (Space "global") dests) - --- | Some kernel bodies are not safe (or efficient) to execute --- multiple times. -data Passage = MustBeSinglePass | MayBeMultiPass deriving (Eq, Ord) - -bodyPassage :: KernelBody KernelsMem -> Passage -bodyPassage kbody - | mempty == consumedInKernelBody (aliasAnalyseKernelBody kbody) = - MayBeMultiPass - | otherwise = - MustBeSinglePass - -prepareIntermediateArraysGlobal :: Passage -> Imp.Exp -> Imp.Exp -> [SegHistSlug] - -> CallKernelGen - (Imp.Exp, - [[Imp.Exp] -> InKernelGen ()]) -prepareIntermediateArraysGlobal passage hist_T hist_N slugs = do - -- The paper formulae assume there is only one histogram, but in our - -- implementation there can be multiple that have been horisontally - -- fused. We do a bit of trickery with summings and averages to - -- pretend there is really only one. For the case of a single - -- histogram, the actual calculations should be the same as in the - -- paper. - - -- The sum of all Hs. - hist_H <- dPrimVE "hist_H" . sum =<< mapM (toExp . histWidth . slugOp) slugs - - hist_RF <- dPrimVE "hist_RF" $ - sum (map (r64 . toExp' int32 . histRaceFactor . slugOp) slugs) - / r64 (genericLength slugs) - - hist_el_size <- dPrimVE "hist_el_size" $ sum $ map slugElAvgSize slugs - - hist_C_max <- dPrimVE "hist_C_max" $ - Imp.BinOpExp (FMin Float64) (r64 hist_T) $ r64 hist_H / hist_k_ct_min - - hist_M_min <- dPrimVE "hist_M_min" $ - Imp.BinOpExp (SMax Int32) 1 $ t64 $ r64 hist_T / hist_C_max - - -- Querying L2 cache size is not reliable. Instead we provide a - -- tunable knob with a hopefully sane default. - let hist_L2_def = 4 * 1024 * 1024 - hist_L2 <- dPrim "L2_size" int32 - entry <- askFunction - -- Equivalent to F_L2*L2 in paper. - sOp $ Imp.GetSize hist_L2 - (keyWithEntryPoint entry $ nameFromString (pretty hist_L2)) $ - Imp.SizeBespoke (nameFromString "L2_for_histogram") hist_L2_def - - let hist_L2_ln_sz = 16*4 -- L2 cache line size approximation - - hist_RACE_exp <- dPrimVE "hist_RACE_exp" $ - Imp.BinOpExp (FMax Float64) 1 $ - (hist_k_RF * hist_RF) / - (hist_L2_ln_sz / r64 hist_el_size) - - hist_S <- dPrim "hist_S" int32 - - -- For sparse histograms (H exceeds N) we only want a single chunk. - sIf (hist_N .<. hist_H) - (hist_S <-- 1) $ - hist_S <-- - case passage of - MayBeMultiPass -> - (hist_M_min * hist_H * hist_el_size) `divUp` - t64 (hist_F_L2 * r64 (Imp.vi32 hist_L2) * hist_RACE_exp) - MustBeSinglePass -> - 1 - - emit $ Imp.DebugPrint "Race expansion factor (RACE^exp)" $ Just hist_RACE_exp - emit $ Imp.DebugPrint "Number of chunks (S)" $ Just $ Imp.vi32 hist_S - - histograms <- snd <$> mapAccumLM (onOp (Imp.vi32 hist_L2) hist_M_min (Imp.vi32 hist_S) hist_RACE_exp) Nothing slugs - - return (Imp.vi32 hist_S, histograms) - where - hist_k_ct_min = 2 -- Chosen experimentally - hist_k_RF = 0.75 -- Chosen experimentally - hist_F_L2 = 0.4 -- Chosen experimentally - - r64 = ConvOpExp (SIToFP Int32 Float64) - t64 = ConvOpExp (FPToSI Float64 Int32) - - -- "Average element size" as computed by a formula that also takes - -- locking into account. - slugElAvgSize slug@(SegHistSlug op _ _ do_op) = - case do_op of - AtomicLocking{} -> - slugElSize slug `quot` (1+genericLength (lambdaReturnType (histOp op))) - _ -> - slugElSize slug `quot` genericLength (lambdaReturnType (histOp op)) - - -- "Average element size" as computed by a formula that also takes - -- locking into account. - slugElSize (SegHistSlug op _ _ do_op) = - case do_op of - AtomicLocking{} -> - sExt Int32 $ unCount $ - sum $ map (typeSize . (`arrayOfShape` histShape op)) $ - Prim int32 : lambdaReturnType (histOp op) - _ -> - sExt Int32 $ unCount $ sum $ - map (typeSize . (`arrayOfShape` histShape op)) $ - lambdaReturnType (histOp op) - - onOp hist_L2 hist_M_min hist_S hist_RACE_exp l slug = do - let SegHistSlug op num_subhistos subhisto_info do_op = slug - hist_H <- toExp $ histWidth op - - hist_H_chk <- dPrimVE "hist_H_chk" $ - hist_H `divUp` hist_S - - emit $ Imp.DebugPrint "Chunk size (H_chk)" $ Just hist_H_chk - - hist_k_max <- dPrimVE "hist_k_max" $ - Imp.BinOpExp (FMin Float64) - (hist_F_L2 * (r64 hist_L2 / r64 (slugElSize slug)) * hist_RACE_exp) - (r64 hist_N) - / r64 hist_T - - hist_u <- dPrimVE "hist_u" $ - case do_op of - AtomicPrim{} -> 2 - _ -> 1 - - hist_C <- dPrimVE "hist_C" $ - Imp.BinOpExp (FMin Float64) (r64 hist_T) $ - r64 (hist_u * hist_H_chk) / hist_k_max - - -- Number of subhistograms per result histogram. - hist_M <- dPrimVE "hist_M" $ - case slugAtomicUpdate slug of - AtomicPrim{} -> 1 - _ -> Imp.BinOpExp (SMax Int32) hist_M_min $ - t64 $ r64 hist_T / hist_C - - emit $ Imp.DebugPrint "Elements/thread in L2 cache (k_max)" $ Just hist_k_max - emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just hist_M - emit $ Imp.DebugPrint "Cooperation level (C)" $ Just hist_C - - -- num_subhistos is the variable we use to communicate back. - num_subhistos <-- hist_M - - -- Initialise sub-histograms. - -- - -- If hist_M is 1, then we just reuse the original - -- destination. The idea is to avoid a copy if we are writing a - -- small number of values into a very large prior histogram. - dests <- forM (zip (histDest op) subhisto_info) $ \(dest, info) -> do - dest_mem <- entryArrayLocation <$> lookupArray dest - - sub_mem <- fmap memLocationName $ - entryArrayLocation <$> - lookupArray (subhistosArray info) - - let unitHistoCase = - emit $ - Imp.SetMem sub_mem (memLocationName dest_mem) $ - Space "device" - - multiHistoCase = subhistosAlloc info - - sIf (hist_M .==. 1) unitHistoCase multiHistoCase - - return $ subhistosArray info - - (l', do_op') <- prepareAtomicUpdateGlobal l dests slug - - return (l', do_op') - -histKernelGlobalPass :: [PatElem KernelsMem] - -> Count NumGroups Imp.Exp - -> Count GroupSize Imp.Exp - -> SegSpace - -> [SegHistSlug] - -> KernelBody KernelsMem - -> [[Imp.Exp] -> InKernelGen ()] - -> Imp.Exp -> Imp.Exp - -> CallKernelGen () -histKernelGlobalPass map_pes num_groups group_size space slugs kbody histograms hist_S chk_i = do - - let (space_is, space_sizes) = unzip $ unSegSpace space - space_sizes_64 = map (sExt Int64 . toExp' int32) space_sizes - total_w_64 = product space_sizes_64 - - hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w -> do - w' <- toExp w - dPrimVE "hist_H_chk" $ w' `divUp` hist_S - - sKernelThread "seghist_global" num_groups group_size (segFlat space) $ do - constants <- kernelConstants <$> askEnv - - -- Compute subhistogram index for each thread, per histogram. - subhisto_inds <- forM slugs $ \slug -> - dPrimVE "subhisto_ind" $ - kernelGlobalThreadId constants `quot` - (kernelNumThreads constants `divUp` Imp.vi32 (slugNumSubhistos slug)) - - -- Loop over flat offsets into the input and output. The - -- calculation is done with 64-bit integers to avoid overflow, - -- but the final unflattened segment indexes are 32 bit. - let gtid = sExt Int64 $ kernelGlobalThreadId constants - num_threads = sExt Int64 $ kernelNumThreads constants - kernelLoop gtid num_threads total_w_64 $ \offset -> do - - -- Construct segment indices. - let setIndex v e = do dPrim_ v int32 - v <-- e - zipWithM_ setIndex space_is $ - map (sExt Int32) $ unflattenIndex space_sizes_64 offset - - -- We execute the bucket function once and update each histogram serially. - -- We apply the bucket function if j=offset+ltid is less than - -- num_elements. This also involves writing to the mapout - -- arrays. - let input_in_bounds = offset .<. total_w_64 - - sWhen input_in_bounds $ compileStms mempty (kernelBodyStms kbody) $ do - let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody - - sComment "save map-out results" $ - forM_ (zip map_pes map_res) $ \(pe, res) -> - copyDWIMFix (patElemName pe) - (map (Imp.vi32 . fst) $ unSegSpace space) - (kernelResultSubExp res) [] - - let (buckets, vs) = splitAt (length slugs) red_res - perOp = chunks $ map (length . histDest . slugOp) slugs - - sComment "perform atomic updates" $ - forM_ (zip6 (map slugOp slugs) histograms buckets (perOp vs) subhisto_inds hist_H_chks) $ - \(HistOp dest_w _ _ _ shape lam, - do_op, bucket, vs', subhisto_ind, hist_H_chk) -> do - - let chk_beg = chk_i * hist_H_chk - bucket' = toExp' int32 $ kernelResultSubExp bucket - dest_w' = toExp' int32 dest_w - bucket_in_bounds = chk_beg .<=. bucket' .&&. - bucket' .<. (chk_beg + hist_H_chk) .&&. - bucket' .<. dest_w' - vs_params = takeLast (length vs') $ lambdaParams lam - - sWhen bucket_in_bounds $ do - let bucket_is = map Imp.vi32 (init space_is) ++ - [subhisto_ind, bucket'] - dLParams $ lambdaParams lam - sLoopNest shape $ \is -> do - forM_ (zip vs_params vs') $ \(p, res) -> - copyDWIMFix (paramName p) [] (kernelResultSubExp res) is - do_op (bucket_is ++ is) - - -histKernelGlobal :: [PatElem KernelsMem] - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> SegSpace - -> [SegHistSlug] - -> KernelBody KernelsMem - -> CallKernelGen () -histKernelGlobal map_pes num_groups group_size space slugs kbody = do - num_groups' <- traverse toExp num_groups - group_size' <- traverse toExp group_size - let (_space_is, space_sizes) = unzip $ unSegSpace space - num_threads = unCount num_groups' * unCount group_size' - - emit $ Imp.DebugPrint "## Using global memory" Nothing - - (hist_S, histograms) <- - prepareIntermediateArraysGlobal (bodyPassage kbody) - num_threads (toExp' int32 $ last space_sizes) slugs - - sFor "chk_i" hist_S $ \chk_i -> - histKernelGlobalPass map_pes num_groups' group_size' space slugs kbody - histograms hist_S chk_i - -type InitLocalHistograms = [([VName], - SubExp -> - InKernelGen ([VName], - [Imp.Exp] -> InKernelGen ()))] - -prepareIntermediateArraysLocal :: VName - -> Count NumGroups Imp.Exp - -> SegSpace -> [SegHistSlug] - -> CallKernelGen InitLocalHistograms -prepareIntermediateArraysLocal num_subhistos_per_group groups_per_segment space slugs = do - num_segments <- dPrimVE "num_segments" $ - product $ map (toExp' int32 . snd) $ init $ unSegSpace space - mapM (onOp num_segments) slugs - where - onOp num_segments (SegHistSlug op num_subhistos subhisto_info do_op) = do - - num_subhistos <-- unCount groups_per_segment * num_segments - - emit $ Imp.DebugPrint "Number of subhistograms in global memory" $ - Just $ Imp.vi32 num_subhistos - - mk_op <- - case do_op of - AtomicPrim f -> return $ const $ return f - AtomicCAS f -> return $ const $ return f - AtomicLocking f -> return $ \hist_H_chk -> do - let lock_shape = - Shape $ Var num_subhistos_per_group : - shapeDims (histShape op) ++ - [hist_H_chk] - - dims <- mapM toExp $ shapeDims lock_shape - - locks <- sAllocArray "locks" int32 lock_shape $ Space "local" - - sComment "All locks start out unlocked" $ - groupCoverSpace dims $ \is -> - copyDWIMFix locks is (intConst Int32 0) [] - - return $ f $ Locking locks 0 1 0 id - - -- Initialise local-memory sub-histograms. These are - -- represented as two-dimensional arrays. - let init_local_subhistos hist_H_chk = do - local_subhistos <- - forM (histType op) $ \t -> do - let sub_local_shape = - Shape [Var num_subhistos_per_group] <> - (arrayShape t `setOuterDim` hist_H_chk) - sAllocArray "subhistogram_local" - (elemType t) sub_local_shape (Space "local") - - do_op' <- mk_op hist_H_chk - - return (local_subhistos, do_op' (Space "local") local_subhistos) - - -- Initialise global-memory sub-histograms. - glob_subhistos <- forM subhisto_info $ \info -> do - subhistosAlloc info - return $ subhistosArray info - - return (glob_subhistos, init_local_subhistos) - -histKernelLocalPass :: VName -> Count NumGroups Imp.Exp - -> [PatElem KernelsMem] - -> Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp - -> SegSpace - -> [SegHistSlug] - -> KernelBody KernelsMem - -> InitLocalHistograms -> Imp.Exp -> Imp.Exp - -> CallKernelGen () -histKernelLocalPass num_subhistos_per_group_var groups_per_segment map_pes num_groups group_size space slugs kbody - init_histograms hist_S chk_i = do - let (space_is, space_sizes) = unzip $ unSegSpace space - segment_is = init space_is - segment_dims = init space_sizes - (i_in_segment, segment_size) = last $ unSegSpace space - num_subhistos_per_group = Imp.var num_subhistos_per_group_var int32 - - segment_size' <- toExp segment_size - - num_segments <- dPrimVE "num_segments" $ - product $ map (toExp' int32) segment_dims - - hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w -> do - w' <- toExp w - dPrimV "hist_H_chk" $ w' `divUp` hist_S - - sKernelThread "seghist_local" num_groups group_size (segFlat space) $ - virtualiseGroups SegVirt (unCount groups_per_segment * num_segments) $ \group_id_var -> do - - constants <- kernelConstants <$> askEnv - - let group_id = Imp.vi32 group_id_var - - flat_segment_id <- dPrimVE "flat_segment_id" $ group_id `quot` unCount groups_per_segment - gid_in_segment <- dPrimVE "gid_in_segment" $ group_id `rem` unCount groups_per_segment - -- This pgtid is kind of a "virtualised physical" gtid - not the - -- same thing as the gtid used for the SegHist itself. - pgtid_in_segment <- dPrimVE "pgtid_in_segment" $ - gid_in_segment * kernelGroupSize constants + kernelLocalThreadId constants - threads_per_segment <- dPrimVE "threads_per_segment" $ - unCount groups_per_segment * kernelGroupSize constants - - -- Set segment indices. - zipWithM_ dPrimV_ segment_is $ - unflattenIndex (map (toExp' int32) segment_dims) flat_segment_id - - histograms <- forM (zip init_histograms hist_H_chks) $ - \((glob_subhistos, init_local_subhistos), hist_H_chk) -> do - (local_subhistos, do_op) <- init_local_subhistos $ Var hist_H_chk - return (zip glob_subhistos local_subhistos, hist_H_chk, do_op) - - -- Find index of local subhistograms updated by this thread. We - -- try to ensure, as much as possible, that threads in the same - -- warp use different subhistograms, to avoid conflicts. - thread_local_subhisto_i <- - dPrimVE "thread_local_subhisto_i" $ - kernelLocalThreadId constants `rem` num_subhistos_per_group - - let onSlugs f = forM_ (zip slugs histograms) $ \(slug, (dests, hist_H_chk, _)) -> do - let histo_dims = fmap (toExp' int32) $ Var hist_H_chk : - shapeDims (histShape (slugOp slug)) - histo_size <- dPrimVE "histo_size" $ product histo_dims - f slug dests (Imp.vi32 hist_H_chk) histo_dims histo_size - - let onAllHistograms f = - onSlugs $ \slug dests hist_H_chk histo_dims histo_size -> do - let group_hists_size = num_subhistos_per_group * histo_size - init_per_thread <- dPrimVE "init_per_thread" $ - group_hists_size `divUp` - kernelGroupSize constants - - forM_ (zip dests (histNeutral $ slugOp slug)) $ - \((dest_global, dest_local), ne) -> - sFor "local_i" init_per_thread $ \i -> do - j <- dPrimVE "j" $ - i * kernelGroupSize constants + - kernelLocalThreadId constants - j_offset <- dPrimVE "j_offset" $ - num_subhistos_per_group * histo_size * gid_in_segment + j - - local_subhisto_i <- dPrimVE "local_subhisto_i" $ j `quot` histo_size - let local_bucket_is = unflattenIndex histo_dims $ j `rem` histo_size - global_bucket_is = head local_bucket_is + chk_i * hist_H_chk : - tail local_bucket_is - global_subhisto_i <- dPrimVE "global_subhisto_i" $ j_offset `quot` histo_size - - sWhen (j .<. group_hists_size) $ - f dest_local dest_global (slugOp slug) ne - local_subhisto_i global_subhisto_i - local_bucket_is global_bucket_is - - sComment "initialize histograms in local memory" $ - onAllHistograms $ \dest_local dest_global op ne local_subhisto_i global_subhisto_i local_bucket_is global_bucket_is -> - sComment "First subhistogram is initialised from global memory; others with neutral element." $ do - let global_is = map Imp.vi32 segment_is ++ [0] ++ global_bucket_is - local_is = local_subhisto_i : local_bucket_is - sIf (global_subhisto_i .==. 0) - (copyDWIMFix dest_local local_is (Var dest_global) global_is) - (sLoopNest (histShape op) $ \is -> - copyDWIMFix dest_local (local_is++is) ne []) - - sOp $ Imp.Barrier Imp.FenceLocal - - kernelLoop pgtid_in_segment threads_per_segment segment_size' $ \ie -> do - dPrimV_ i_in_segment ie - - -- We execute the bucket function once and update each histogram - -- serially. This also involves writing to the mapout arrays if - -- this is the first chunk. - - compileStms mempty (kernelBodyStms kbody) $ do - - let (red_res, map_res) = splitFromEnd (length map_pes) $ - map kernelResultSubExp $ kernelBodyResult kbody - (buckets, vs) = splitAt (length slugs) red_res - perOp = chunks $ map (length . histDest . slugOp) slugs - - sWhen (chk_i .==. 0) $ - sComment "save map-out results" $ - forM_ (zip map_pes map_res) $ \(pe, se) -> - copyDWIMFix (patElemName pe) - (map Imp.vi32 space_is) se [] - - forM_ (zip4 (map slugOp slugs) histograms buckets (perOp vs)) $ - \(HistOp dest_w _ _ _ shape lam, - (_, hist_H_chk, do_op), bucket, vs') -> do - - let chk_beg = chk_i * Imp.vi32 hist_H_chk - bucket' = toExp' int32 bucket - dest_w' = toExp' int32 dest_w - bucket_in_bounds = bucket' .<. dest_w' .&&. - chk_beg .<=. bucket' .&&. - bucket' .<. (chk_beg + Imp.vi32 hist_H_chk) - bucket_is = [thread_local_subhisto_i, bucket' - chk_beg] - vs_params = takeLast (length vs') $ lambdaParams lam - - sComment "perform atomic updates" $ - sWhen bucket_in_bounds $ do - dLParams $ lambdaParams lam - sLoopNest shape $ \is -> do - forM_ (zip vs_params vs') $ \(p, v) -> - copyDWIMFix (paramName p) [] v is - do_op (bucket_is ++ is) - - sOp $ Imp.ErrorSync Imp.FenceGlobal - - sComment "Compact the multiple local memory subhistograms to result in global memory" $ - onSlugs $ \slug dests hist_H_chk histo_dims histo_size -> do - bins_per_thread <- dPrimVE "init_per_thread" $ - histo_size `divUp` kernelGroupSize constants - - trunc_H <- dPrimV "trunc_H" $ - Imp.BinOpExp (SMin Int32) hist_H_chk $ - toExp' int32 (histWidth (slugOp slug)) - - chk_i * head histo_dims - let trunc_histo_dims = map (toExp' int32) $ Var trunc_H : - shapeDims (histShape (slugOp slug)) - trunc_histo_size <- dPrimVE "histo_size" $ product trunc_histo_dims - - sFor "local_i" bins_per_thread $ \i -> do - j <- dPrimVE "j" $ - i * kernelGroupSize constants + kernelLocalThreadId constants - sWhen (j .<. trunc_histo_size) $ do - -- We are responsible for compacting the flat bin 'j', which - -- we immediately unflatten. - let local_bucket_is = unflattenIndex histo_dims j - global_bucket_is = head local_bucket_is + chk_i * hist_H_chk : - tail local_bucket_is - dLParams $ lambdaParams $ histOp $ slugOp slug - let (global_dests, local_dests) = unzip dests - (xparams, yparams) = splitAt (length local_dests) $ - lambdaParams $ histOp $ slugOp slug - - sComment "Read values from subhistogram 0." $ - forM_ (zip xparams local_dests) $ \(xp, subhisto) -> - copyDWIMFix - (paramName xp) [] - (Var subhisto) (0:local_bucket_is) - - sComment "Accumulate based on values in other subhistograms." $ - sFor "subhisto_id" (num_subhistos_per_group - 1) $ \subhisto_id -> do - forM_ (zip yparams local_dests) $ \(yp, subhisto) -> - copyDWIMFix - (paramName yp) [] - (Var subhisto) (subhisto_id + 1 : local_bucket_is) - compileBody' xparams $ lambdaBody $ histOp $ slugOp slug - - sComment "Put final bucket value in global memory." $ do - let global_is = - map Imp.vi32 segment_is ++ - [group_id `rem` unCount groups_per_segment] ++ - global_bucket_is - forM_ (zip xparams global_dests) $ \(xp, global_dest) -> - copyDWIMFix global_dest global_is (Var $ paramName xp) [] - -histKernelLocal :: VName -> Count NumGroups Imp.Exp - -> [PatElem KernelsMem] - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> SegSpace - -> Imp.Exp - -> [SegHistSlug] - -> KernelBody KernelsMem - -> CallKernelGen () -histKernelLocal num_subhistos_per_group_var groups_per_segment map_pes num_groups group_size space hist_S slugs kbody = do - num_groups' <- traverse toExp num_groups - group_size' <- traverse toExp group_size - let num_subhistos_per_group = Imp.var num_subhistos_per_group_var int32 - - emit $ Imp.DebugPrint "Number of local subhistograms per group" $ Just num_subhistos_per_group - - init_histograms <- - prepareIntermediateArraysLocal num_subhistos_per_group_var groups_per_segment space slugs - - sFor "chk_i" hist_S $ \chk_i -> - histKernelLocalPass - num_subhistos_per_group_var groups_per_segment map_pes num_groups' group_size' space slugs kbody - init_histograms hist_S chk_i - --- | The maximum number of passes we are willing to accept for this --- kind of atomic update. -slugMaxLocalMemPasses :: SegHistSlug -> Int -slugMaxLocalMemPasses slug = - case slugAtomicUpdate slug of - AtomicPrim _ -> 3 - AtomicCAS _ -> 4 - AtomicLocking _ -> 6 - -localMemoryCase :: [PatElem KernelsMem] - -> Imp.Exp - -> SegSpace - -> Imp.Exp -> Imp.Exp -> Imp.Exp -> Imp.Exp - -> [SegHistSlug] - -> KernelBody KernelsMem - -> CallKernelGen (Imp.Exp, CallKernelGen ()) -localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N _ slugs kbody = do - let space_sizes = segSpaceDims space - segment_dims = init space_sizes - segmented = not $ null segment_dims - - hist_L <- dPrim "hist_L" int32 - sOp $ Imp.GetSizeMax hist_L Imp.SizeLocalMemory - - max_group_size <- dPrim "max_group_size" int32 - sOp $ Imp.GetSizeMax max_group_size Imp.SizeGroup - let group_size = Imp.Count $ Var max_group_size - num_groups <- fmap (Imp.Count . Var) $ dPrimV "num_groups" $ - hist_T `divUp` toExp' int32 (unCount group_size) - let num_groups' = toExp' int32 <$> num_groups - group_size' = toExp' int32 <$> group_size - - let r64 = ConvOpExp (SIToFP Int32 Float64) - t64 = ConvOpExp (FPToSI Float64 Int32) - - -- M approximation. - hist_m' <- dPrimVE "hist_m_prime" $ - r64 (Imp.BinOpExp (SMin Int32) - (Imp.vi32 hist_L `quot` hist_el_size) - (hist_N `divUp` unCount num_groups')) - / r64 hist_H - - let hist_B = unCount group_size' - - -- M in the paper, but not adjusted for asymptotic efficiency. - hist_M0 <- dPrimVE "hist_M0" $ - Imp.BinOpExp (SMax Int32) 1 $ - Imp.BinOpExp (SMin Int32) (t64 hist_m') hist_B - - -- Minimal sequential chunking factor. - let q_small = 2 - - -- The number of segments/histograms produced.. - hist_Nout <- dPrimVE "hist_Nout" $ product $ map (toExp' int32) segment_dims - - hist_Nin <- dPrimVE "hist_Nin" $ toExp' int32 $ last space_sizes - - -- Maximum M for work efficiency. - work_asymp_M_max <- - if segmented then do - - hist_T_hist_min <- dPrimVE "hist_T_hist_min" $ - sExt Int32 $ - Imp.BinOpExp (SMin Int64) - (sExt Int64 hist_Nin * sExt Int64 hist_Nout) (sExt Int64 hist_T) - `divUp` - sExt Int64 hist_Nout - - -- Number of groups, rounded up. - let r = hist_T_hist_min `divUp` hist_B - - dPrimVE "work_asymp_M_max" $ hist_Nin `quot` (r * hist_H) - - else dPrimVE "work_asymp_M_max" $ - (hist_Nout * hist_N) `quot` - ((q_small * unCount num_groups' * hist_H) - `quot` genericLength slugs) - - -- Number of subhistograms per result histogram. - hist_M <- dPrimV "hist_M" $ - Imp.BinOpExp (SMin Int32) hist_M0 work_asymp_M_max - - -- hist_M may be zero (which we'll check for below), but we need it - -- for some divisions first, so crudely make a nonzero form. - let hist_M_nonzero = Imp.BinOpExp (SMax Int32) 1 $ Imp.vi32 hist_M - - -- "Cooperation factor" - the number of threads cooperatively - -- working on the same (sub)histogram. - hist_C <- dPrimVE "hist_C" $ - hist_B `divUp` hist_M_nonzero - - emit $ Imp.DebugPrint "local hist_M0" $ Just hist_M0 - emit $ Imp.DebugPrint "local work asymp M max" $ Just work_asymp_M_max - emit $ Imp.DebugPrint "local C" $ Just hist_C - emit $ Imp.DebugPrint "local B" $ Just hist_B - emit $ Imp.DebugPrint "local M" $ Just $ Imp.vi32 hist_M - emit $ Imp.DebugPrint "local memory needed" $ - Just $ hist_H * hist_el_size * Imp.vi32 hist_M - - -- local_mem_needed is what we need to keep a single bucket in local - -- memory - this is an absolute minimum. We can fit anything else - -- by doing multiple passes, although more than a few is - -- (heuristically) not efficient. - local_mem_needed <- dPrimVE "local_mem_needed" $ hist_el_size * Imp.vi32 hist_M - hist_S <- dPrimVE "hist_S" $ (hist_H * local_mem_needed) `divUp` Imp.vi32 hist_L - let max_S = case bodyPassage kbody of - MustBeSinglePass -> 1 - MayBeMultiPass -> fromIntegral $ maxinum $ map slugMaxLocalMemPasses slugs - - -- We only use local memory if the number of updates per histogram - -- at least matches the histogram size, as otherwise it is not - -- asymptotically efficient. This mostly matters for the segmented - -- case. - let pick_local = - hist_Nin .>=. hist_H - .&&. (local_mem_needed .<=. Imp.vi32 hist_L) - .&&. (hist_S .<=. max_S) - .&&. hist_C .<=. hist_B - .&&. Imp.vi32 hist_M .>. 0 - - groups_per_segment - | segmented = num_groups' `divUp` Imp.Count hist_Nout - | otherwise = num_groups' - - run = do - emit $ Imp.DebugPrint "## Using local memory" Nothing - emit $ Imp.DebugPrint "Histogram size (H)" $ Just hist_H - emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just $ Imp.vi32 hist_M - emit $ Imp.DebugPrint "Cooperation level (C)" $ Just hist_C - emit $ Imp.DebugPrint "Number of chunks (S)" $ Just hist_S - when segmented $ - emit $ Imp.DebugPrint "Groups per segment" $ Just $ unCount groups_per_segment - histKernelLocal hist_M groups_per_segment map_pes - num_groups group_size space hist_S slugs kbody - - return (pick_local, run) - --- | Generate code for a segmented histogram called from the host. -compileSegHist :: Pattern KernelsMem - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> SegSpace - -> [HistOp KernelsMem] - -> KernelBody KernelsMem - -> CallKernelGen () -compileSegHist (Pattern _ pes) num_groups group_size space ops kbody = do - -- Most of this function is not the histogram part itself, but - -- rather figuring out whether to use a local or global memory - -- strategy, as well as collapsing the subhistograms produced (which - -- are always in global memory, but their number may vary). - num_groups' <- traverse toExp num_groups - group_size' <- traverse toExp group_size - - dims <- mapM toExp $ segSpaceDims space - - let num_red_res = length ops + sum (map (length . histNeutral) ops) - (all_red_pes, map_pes) = splitAt num_red_res pes - segment_size = last dims - - (op_hs, op_seg_hs, slugs) <- unzip3 <$> mapM (computeHistoUsage space) ops - h <- dPrimVE "h" $ Imp.unCount $ sum op_hs - seg_h <- dPrimVE "seg_h" $ Imp.unCount $ sum op_seg_hs - - -- Check for emptyness to avoid division-by-zero. - sUnless (seg_h .==. 0) $ do - - -- Maximum group size (or actual, in this case). - let hist_B = unCount group_size' - - -- Size of a histogram. - hist_H <- dPrimVE "hist_H" $ sum $ map (toExp' int32 . histWidth) ops - - -- Size of a single histogram element. Actually the weighted - -- average of histogram elements in cases where we have more than - -- one histogram operation, plus any locks. - let lockSize slug = case slugAtomicUpdate slug of - AtomicLocking{} -> Just $ primByteSize int32 - _ -> Nothing - hist_el_size <- dPrimVE "hist_el_size" $ foldl' (+) (h `divUp` hist_H) $ - mapMaybe lockSize slugs - - -- Input elements contributing to each histogram. - hist_N <- dPrimVE "hist_N" segment_size - - -- Compute RF as the average RF over all the histograms. - hist_RF <- dPrimVE "hist_RF" $ - sum (map (toExp' int32. histRaceFactor . slugOp) slugs) - `quot` - genericLength slugs - - let hist_T = unCount num_groups' * unCount group_size' - emit $ Imp.DebugPrint "\n# SegHist" Nothing - emit $ Imp.DebugPrint "Number of threads (T)" $ Just hist_T - emit $ Imp.DebugPrint "Desired group size (B)" $ Just hist_B - emit $ Imp.DebugPrint "Histogram size (H)" $ Just hist_H - emit $ Imp.DebugPrint "Input elements per histogram (N)" $ Just hist_N - emit $ Imp.DebugPrint "Number of segments" $ - Just $ product $ map (toExp' int32 . snd) segment_dims - emit $ Imp.DebugPrint "Histogram element size (el_size)" $ Just hist_el_size - emit $ Imp.DebugPrint "Race factor (RF)" $ Just hist_RF - emit $ Imp.DebugPrint "Memory per set of subhistograms per segment" $ Just h - emit $ Imp.DebugPrint "Memory per set of subhistograms times segments" $ Just seg_h - - (use_local_memory, run_in_local_memory) <- - localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N hist_RF slugs kbody - - sIf use_local_memory run_in_local_memory $ - histKernelGlobal map_pes num_groups group_size space slugs kbody - - let pes_per_op = chunks (map (length . histDest) ops) all_red_pes - - forM_ (zip3 slugs pes_per_op ops) $ \(slug, red_pes, op) -> do - let num_histos = slugNumSubhistos slug - subhistos = map subhistosArray $ slugSubhistos slug - - let unitHistoCase = - -- This is OK because the memory blocks are at least as - -- large as the ones we are supposed to use for the result. - forM_ (zip red_pes subhistos) $ \(pe, subhisto) -> do - pe_mem <- memLocationName . entryArrayLocation <$> - lookupArray (patElemName pe) - subhisto_mem <- memLocationName . entryArrayLocation <$> - lookupArray subhisto - emit $ Imp.SetMem pe_mem subhisto_mem $ Space "device" - - sIf (Imp.var num_histos int32 .==. 1) unitHistoCase $ do - -- For the segmented reduction, we keep the segment dimensions - -- unchanged. To this, we add two dimensions: one over the number - -- of buckets, and one over the number of subhistograms. This - -- inner dimension is the one that is collapsed in the reduction. - let num_buckets = histWidth op - - bucket_id <- newVName "bucket_id" - subhistogram_id <- newVName "subhistogram_id" - vector_ids <- mapM (const $ newVName "vector_id") $ - shapeDims $ histShape op - - flat_gtid <- newVName "flat_gtid" - - let lvl = SegThread num_groups group_size SegVirt - segred_space = - SegSpace flat_gtid $ - segment_dims ++ - [(bucket_id, num_buckets)] ++ - zip vector_ids (shapeDims $ histShape op) ++ - [(subhistogram_id, Var num_histos)] - - let segred_op = SegBinOp Commutative (histOp op) (histNeutral op) mempty - compileSegRed' (Pattern [] red_pes) lvl segred_space [segred_op] $ \red_cont -> - red_cont $ flip map subhistos $ \subhisto -> - (Var subhisto, map Imp.vi32 $ - map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids) - - where segment_dims = init $ unSegSpace space +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Our compilation strategy for 'SegHist' is based around avoiding +-- bin conflicts. We do this by splitting the input into chunks, and +-- for each chunk computing a single subhistogram. Then we combine +-- the subhistograms using an ordinary segmented reduction ('SegRed'). +-- +-- There are some branches around to efficiently handle the case where +-- we use only a single subhistogram (because it's large), so that we +-- respect the asymptotics, and do not copy the destination array. +-- +-- We also use a heuristic strategy for computing subhistograms in +-- local memory when possible. Given: +-- +-- H: total size of histograms in bytes, including any lock arrays. +-- +-- G: group size +-- +-- T: number of bytes of local memory each thread can be given without +-- impacting occupancy (determined experimentally, e.g. 32). +-- +-- LMAX: maximum amount of local memory per workgroup (hard limit). +-- +-- We wish to compute: +-- +-- COOP: cooperation level (number of threads per subhistogram) +-- +-- LH: number of local memory subhistograms +-- +-- We do this as: +-- +-- COOP = ceil(H / T) +-- LH = ceil((G*T)/H) +-- if COOP <= G && H <= LMAX then +-- use local memory +-- else +-- use global memory +module Futhark.CodeGen.ImpGen.Kernels.SegHist (compileSegHist) where + +import Control.Monad.Except +import Data.List (foldl', genericLength, zip4, zip6) +import Data.Maybe +import qualified Futhark.CodeGen.ImpCode.Kernels as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.Kernels.Base +import Futhark.CodeGen.ImpGen.Kernels.SegRed (compileSegRed') +import Futhark.Construct (fullSliceNum) +import Futhark.IR.KernelsMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Pass.ExplicitAllocations () +import Futhark.Util (chunks, mapAccumLM, maxinum, splitFromEnd, takeLast) +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) + +data SubhistosInfo = SubhistosInfo + { subhistosArray :: VName, + subhistosAlloc :: CallKernelGen () + } + +data SegHistSlug = SegHistSlug + { slugOp :: HistOp KernelsMem, + slugNumSubhistos :: TV Int64, + slugSubhistos :: [SubhistosInfo], + slugAtomicUpdate :: AtomicUpdate KernelsMem KernelEnv + } + +histoSpaceUsage :: + HistOp KernelsMem -> + Imp.Count Imp.Bytes (Imp.TExp Int64) +histoSpaceUsage op = + sum $ + map + ( typeSize + . (`arrayOfRow` histWidth op) + . (`arrayOfShape` histShape op) + ) + $ lambdaReturnType $ histOp op + +-- | Figure out how much memory is needed per histogram, both +-- segmented and unsegmented,, and compute some other auxiliary +-- information. +computeHistoUsage :: + SegSpace -> + HistOp KernelsMem -> + CallKernelGen + ( Imp.Count Imp.Bytes (Imp.TExp Int64), + Imp.Count Imp.Bytes (Imp.TExp Int64), + SegHistSlug + ) +computeHistoUsage space op = do + let segment_dims = init $ unSegSpace space + num_segments = length segment_dims + + -- Create names for the intermediate array memory blocks, + -- memory block sizes, arrays, and number of subhistograms. + num_subhistos <- dPrim "num_subhistos" int32 + subhisto_infos <- forM (zip (histDest op) (histNeutral op)) $ \(dest, ne) -> do + dest_t <- lookupType dest + dest_mem <- entryArrayLocation <$> lookupArray dest + + subhistos_mem <- + sDeclareMem (baseString dest ++ "_subhistos_mem") (Space "device") + + let subhistos_shape = + Shape (map snd segment_dims ++ [tvSize num_subhistos]) + <> stripDims num_segments (arrayShape dest_t) + subhistos_membind = + ArrayIn subhistos_mem $ + IxFun.iota $ + map pe64 $ shapeDims subhistos_shape + subhistos <- + sArray + (baseString dest ++ "_subhistos") + (elemType dest_t) + subhistos_shape + subhistos_membind + + return $ + SubhistosInfo subhistos $ do + let unitHistoCase = + emit $ + Imp.SetMem subhistos_mem (memLocationName dest_mem) $ + Space "device" + + multiHistoCase = do + let num_elems = + foldl' (*) (sExt64 $ tvExp num_subhistos) $ + map toInt64Exp $ arrayDims dest_t + + let subhistos_mem_size = + Imp.bytes $ + Imp.unCount (Imp.elements num_elems `Imp.withElemType` elemType dest_t) + + sAlloc_ subhistos_mem subhistos_mem_size $ Space "device" + sReplicate subhistos ne + subhistos_t <- lookupType subhistos + let slice = + fullSliceNum (map toInt64Exp $ arrayDims subhistos_t) $ + map (unitSlice 0 . toInt64Exp . snd) segment_dims + ++ [DimFix 0] + sUpdate subhistos slice $ Var dest + + sIf (tvExp num_subhistos .==. 1) unitHistoCase multiHistoCase + + let h = histoSpaceUsage op + segmented_h = h * product (map (Imp.bytes . toInt64Exp) $ init $ segSpaceDims space) + + atomics <- hostAtomics <$> askEnv + + return + ( h, + segmented_h, + SegHistSlug op num_subhistos subhisto_infos $ + atomicUpdateLocking atomics $ histOp op + ) + +prepareAtomicUpdateGlobal :: + Maybe Locking -> + [VName] -> + SegHistSlug -> + CallKernelGen + ( Maybe Locking, + [Imp.TExp Int64] -> InKernelGen () + ) +prepareAtomicUpdateGlobal l dests slug = + -- We need a separate lock array if the operators are not all of a + -- particularly simple form that permits pure atomic operations. + case (l, slugAtomicUpdate slug) of + (_, AtomicPrim f) -> return (l, f (Space "global") dests) + (_, AtomicCAS f) -> return (l, f (Space "global") dests) + (Just l', AtomicLocking f) -> return (l, f l' (Space "global") dests) + (Nothing, AtomicLocking f) -> do + -- The number of locks used here is too low, but since we are + -- currently forced to inline a huge list, I'm keeping it down + -- for now. Some quick experiments suggested that it has little + -- impact anyway (maybe the locking case is just too slow). + -- + -- A fun solution would also be to use a simple hashing + -- algorithm to ensure good distribution of locks. + let num_locks = 100151 + dims = + map toInt64Exp $ + shapeDims (histShape (slugOp slug)) + ++ [ tvSize (slugNumSubhistos slug), + histWidth (slugOp slug) + ] + locks <- + sStaticArray "hist_locks" (Space "device") int32 $ + Imp.ArrayZeros num_locks + let l' = Locking locks 0 1 0 (pure . (`rem` fromIntegral num_locks) . flattenIndex dims) + return (Just l', f l' (Space "global") dests) + +-- | Some kernel bodies are not safe (or efficient) to execute +-- multiple times. +data Passage = MustBeSinglePass | MayBeMultiPass deriving (Eq, Ord) + +bodyPassage :: KernelBody KernelsMem -> Passage +bodyPassage kbody + | mempty == consumedInKernelBody (aliasAnalyseKernelBody kbody) = + MayBeMultiPass + | otherwise = + MustBeSinglePass + +prepareIntermediateArraysGlobal :: + Passage -> + Imp.TExp Int32 -> + Imp.TExp Int64 -> + [SegHistSlug] -> + CallKernelGen + ( Imp.TExp Int32, + [[Imp.TExp Int64] -> InKernelGen ()] + ) +prepareIntermediateArraysGlobal passage hist_T hist_N slugs = do + -- The paper formulae assume there is only one histogram, but in our + -- implementation there can be multiple that have been horisontally + -- fused. We do a bit of trickery with summings and averages to + -- pretend there is really only one. For the case of a single + -- histogram, the actual calculations should be the same as in the + -- paper. + + -- The sum of all Hs. + hist_H <- dPrimVE "hist_H" $ sum $ map (toInt64Exp . histWidth . slugOp) slugs + + hist_RF <- + dPrimVE "hist_RF" $ + sum (map (r64 . toInt64Exp . histRaceFactor . slugOp) slugs) + / genericLength slugs + + hist_el_size <- dPrimVE "hist_el_size" $ sum $ map slugElAvgSize slugs + + hist_C_max <- + dPrimVE "hist_C_max" $ + fMin64 (r64 hist_T) $ r64 hist_H / hist_k_ct_min + + hist_M_min <- + dPrimVE "hist_M_min" $ + sMax32 1 $ sExt32 $ t64 $ r64 hist_T / hist_C_max + + -- Querying L2 cache size is not reliable. Instead we provide a + -- tunable knob with a hopefully sane default. + let hist_L2_def = 4 * 1024 * 1024 + hist_L2 <- dPrim "L2_size" int32 + entry <- askFunction + -- Equivalent to F_L2*L2 in paper. + sOp $ + Imp.GetSize + (tvVar hist_L2) + (keyWithEntryPoint entry $ nameFromString (pretty (tvVar hist_L2))) + $ Imp.SizeBespoke (nameFromString "L2_for_histogram") hist_L2_def + + let hist_L2_ln_sz = 16 * 4 -- L2 cache line size approximation + hist_RACE_exp <- + dPrimVE "hist_RACE_exp" $ + fMax64 1 $ + (hist_k_RF * hist_RF) + / (hist_L2_ln_sz / r64 hist_el_size) + + hist_S <- dPrim "hist_S" int32 + + -- For sparse histograms (H exceeds N) we only want a single chunk. + sIf + (hist_N .<. hist_H) + (hist_S <-- (1 :: Imp.TExp Int32)) + $ hist_S + <-- case passage of + MayBeMultiPass -> + sExt32 $ + (sExt64 hist_M_min * hist_H * sExt64 hist_el_size) + `divUp` t64 (hist_F_L2 * r64 (tvExp hist_L2) * hist_RACE_exp) + MustBeSinglePass -> + 1 + + emit $ Imp.DebugPrint "Race expansion factor (RACE^exp)" $ Just $ untyped hist_RACE_exp + emit $ Imp.DebugPrint "Number of chunks (S)" $ Just $ untyped $ tvExp hist_S + + histograms <- + snd + <$> mapAccumLM + (onOp (tvExp hist_L2) hist_M_min (tvExp hist_S) hist_RACE_exp) + Nothing + slugs + + return (tvExp hist_S, histograms) + where + hist_k_ct_min = 2 -- Chosen experimentally + hist_k_RF = 0.75 -- Chosen experimentally + hist_F_L2 = 0.4 -- Chosen experimentally + r64 = isF64 . ConvOpExp (SIToFP Int32 Float64) . untyped + t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped + + -- "Average element size" as computed by a formula that also takes + -- locking into account. + slugElAvgSize slug@(SegHistSlug op _ _ do_op) = + case do_op of + AtomicLocking {} -> + slugElSize slug `quot` (1 + genericLength (lambdaReturnType (histOp op))) + _ -> + slugElSize slug `quot` genericLength (lambdaReturnType (histOp op)) + + -- "Average element size" as computed by a formula that also takes + -- locking into account. + slugElSize (SegHistSlug op _ _ do_op) = + case do_op of + AtomicLocking {} -> + sExt32 $ + unCount $ + sum $ + map (typeSize . (`arrayOfShape` histShape op)) $ + Prim int32 : lambdaReturnType (histOp op) + _ -> + sExt32 $ + unCount $ + sum $ + map (typeSize . (`arrayOfShape` histShape op)) $ + lambdaReturnType (histOp op) + + onOp hist_L2 hist_M_min hist_S hist_RACE_exp l slug = do + let SegHistSlug op num_subhistos subhisto_info do_op = slug + hist_H = toInt64Exp $ histWidth op + + hist_H_chk <- dPrimVE "hist_H_chk" $ hist_H `divUp` sExt64 hist_S + + emit $ Imp.DebugPrint "Chunk size (H_chk)" $ Just $ untyped hist_H_chk + + hist_k_max <- + dPrimVE "hist_k_max" $ + fMin64 + (hist_F_L2 * (r64 hist_L2 / r64 (slugElSize slug)) * hist_RACE_exp) + (r64 hist_N) + / r64 hist_T + + hist_u <- dPrimVE "hist_u" $ + case do_op of + AtomicPrim {} -> 2 + _ -> 1 + + hist_C <- + dPrimVE "hist_C" $ + fMin64 (r64 hist_T) $ r64 (hist_u * hist_H_chk) / hist_k_max + + -- Number of subhistograms per result histogram. + hist_M <- dPrimVE "hist_M" $ + case slugAtomicUpdate slug of + AtomicPrim {} -> 1 + _ -> sMax32 hist_M_min $ sExt32 $ t64 $ r64 hist_T / hist_C + + emit $ Imp.DebugPrint "Elements/thread in L2 cache (k_max)" $ Just $ untyped hist_k_max + emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just $ untyped hist_M + emit $ Imp.DebugPrint "Cooperation level (C)" $ Just $ untyped hist_C + + -- num_subhistos is the variable we use to communicate back. + num_subhistos <-- sExt64 hist_M + + -- Initialise sub-histograms. + -- + -- If hist_M is 1, then we just reuse the original + -- destination. The idea is to avoid a copy if we are writing a + -- small number of values into a very large prior histogram. + dests <- forM (zip (histDest op) subhisto_info) $ \(dest, info) -> do + dest_mem <- entryArrayLocation <$> lookupArray dest + + sub_mem <- + fmap memLocationName $ + entryArrayLocation + <$> lookupArray (subhistosArray info) + + let unitHistoCase = + emit $ + Imp.SetMem sub_mem (memLocationName dest_mem) $ + Space "device" + + multiHistoCase = subhistosAlloc info + + sIf (hist_M .==. 1) unitHistoCase multiHistoCase + + return $ subhistosArray info + + (l', do_op') <- prepareAtomicUpdateGlobal l dests slug + + return (l', do_op') + +histKernelGlobalPass :: + [PatElem KernelsMem] -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + SegSpace -> + [SegHistSlug] -> + KernelBody KernelsMem -> + [[Imp.TExp Int64] -> InKernelGen ()] -> + Imp.TExp Int32 -> + Imp.TExp Int32 -> + CallKernelGen () +histKernelGlobalPass map_pes num_groups group_size space slugs kbody histograms hist_S chk_i = do + let (space_is, space_sizes) = unzip $ unSegSpace space + space_sizes_64 = map (sExt64 . toInt64Exp) space_sizes + total_w_64 = product space_sizes_64 + + hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w -> + dPrimVE "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S + + sKernelThread "seghist_global" num_groups group_size (segFlat space) $ do + constants <- kernelConstants <$> askEnv + + -- Compute subhistogram index for each thread, per histogram. + subhisto_inds <- forM slugs $ \slug -> + dPrimVE "subhisto_ind" $ + kernelGlobalThreadId constants + `quot` ( kernelNumThreads constants + `divUp` sExt32 (tvExp (slugNumSubhistos slug)) + ) + + -- Loop over flat offsets into the input and output. The + -- calculation is done with 64-bit integers to avoid overflow, + -- but the final unflattened segment indexes are 32 bit. + let gtid = sExt64 $ kernelGlobalThreadId constants + num_threads = sExt64 $ kernelNumThreads constants + kernelLoop gtid num_threads total_w_64 $ \offset -> do + -- Construct segment indices. + zipWithM_ dPrimV_ space_is $ + map sExt32 $ unflattenIndex space_sizes_64 offset + + -- We execute the bucket function once and update each histogram serially. + -- We apply the bucket function if j=offset+ltid is less than + -- num_elements. This also involves writing to the mapout + -- arrays. + let input_in_bounds = offset .<. total_w_64 + + sWhen input_in_bounds $ + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody + + sComment "save map-out results" $ + forM_ (zip map_pes map_res) $ \(pe, res) -> + copyDWIMFix + (patElemName pe) + (map (Imp.vi64 . fst) $ unSegSpace space) + (kernelResultSubExp res) + [] + + let (buckets, vs) = splitAt (length slugs) red_res + perOp = chunks $ map (length . histDest . slugOp) slugs + + sComment "perform atomic updates" $ + forM_ (zip6 (map slugOp slugs) histograms buckets (perOp vs) subhisto_inds hist_H_chks) $ + \( HistOp dest_w _ _ _ shape lam, + do_op, + bucket, + vs', + subhisto_ind, + hist_H_chk + ) -> do + let chk_beg = sExt64 chk_i * hist_H_chk + bucket' = toInt64Exp $ kernelResultSubExp bucket + dest_w' = toInt64Exp dest_w + bucket_in_bounds = + chk_beg .<=. bucket' + .&&. bucket' .<. (chk_beg + hist_H_chk) + .&&. bucket' .<. dest_w' + vs_params = takeLast (length vs') $ lambdaParams lam + + sWhen bucket_in_bounds $ do + let bucket_is = + map Imp.vi64 (init space_is) + ++ [sExt64 subhisto_ind, bucket'] + dLParams $ lambdaParams lam + sLoopNest shape $ \is -> do + forM_ (zip vs_params vs') $ \(p, res) -> + copyDWIMFix (paramName p) [] (kernelResultSubExp res) is + do_op (bucket_is ++ is) + +histKernelGlobal :: + [PatElem KernelsMem] -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegHistSlug] -> + KernelBody KernelsMem -> + CallKernelGen () +histKernelGlobal map_pes num_groups group_size space slugs kbody = do + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + let (_space_is, space_sizes) = unzip $ unSegSpace space + num_threads = sExt32 $ unCount num_groups' * unCount group_size' + + emit $ Imp.DebugPrint "## Using global memory" Nothing + + (hist_S, histograms) <- + prepareIntermediateArraysGlobal + (bodyPassage kbody) + num_threads + (toInt64Exp $ last space_sizes) + slugs + + sFor "chk_i" hist_S $ \chk_i -> + histKernelGlobalPass + map_pes + num_groups' + group_size' + space + slugs + kbody + histograms + hist_S + chk_i + +type InitLocalHistograms = + [ ( [VName], + SubExp -> + InKernelGen + ( [VName], + [Imp.TExp Int64] -> InKernelGen () + ) + ) + ] + +prepareIntermediateArraysLocal :: + TV Int32 -> + Count NumGroups (Imp.TExp Int64) -> + SegSpace -> + [SegHistSlug] -> + CallKernelGen InitLocalHistograms +prepareIntermediateArraysLocal num_subhistos_per_group groups_per_segment space slugs = do + num_segments <- + dPrimVE "num_segments" $ + product $ map (toInt64Exp . snd) $ init $ unSegSpace space + mapM (onOp num_segments) slugs + where + onOp num_segments (SegHistSlug op num_subhistos subhisto_info do_op) = do + num_subhistos <-- sExt64 (unCount groups_per_segment) * num_segments + + emit $ + Imp.DebugPrint "Number of subhistograms in global memory" $ + Just $ untyped $ tvExp num_subhistos + + mk_op <- + case do_op of + AtomicPrim f -> return $ const $ return f + AtomicCAS f -> return $ const $ return f + AtomicLocking f -> return $ \hist_H_chk -> do + let lock_shape = + Shape $ + tvSize num_subhistos_per_group : + shapeDims (histShape op) + ++ [hist_H_chk] + + let dims = map toInt64Exp $ shapeDims lock_shape + + locks <- sAllocArray "locks" int32 lock_shape $ Space "local" + + sComment "All locks start out unlocked" $ + groupCoverSpace dims $ \is -> + copyDWIMFix locks is (intConst Int32 0) [] + + return $ f $ Locking locks 0 1 0 id + + -- Initialise local-memory sub-histograms. These are + -- represented as two-dimensional arrays. + let init_local_subhistos hist_H_chk = do + local_subhistos <- + forM (histType op) $ \t -> do + let sub_local_shape = + Shape [tvSize num_subhistos_per_group] + <> (arrayShape t `setOuterDim` hist_H_chk) + sAllocArray + "subhistogram_local" + (elemType t) + sub_local_shape + (Space "local") + + do_op' <- mk_op hist_H_chk + + return (local_subhistos, do_op' (Space "local") local_subhistos) + + -- Initialise global-memory sub-histograms. + glob_subhistos <- forM subhisto_info $ \info -> do + subhistosAlloc info + return $ subhistosArray info + + return (glob_subhistos, init_local_subhistos) + +histKernelLocalPass :: + TV Int32 -> + Count NumGroups (Imp.TExp Int64) -> + [PatElem KernelsMem] -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + SegSpace -> + [SegHistSlug] -> + KernelBody KernelsMem -> + InitLocalHistograms -> + Imp.TExp Int32 -> + Imp.TExp Int32 -> + CallKernelGen () +histKernelLocalPass + num_subhistos_per_group_var + groups_per_segment + map_pes + num_groups + group_size + space + slugs + kbody + init_histograms + hist_S + chk_i = do + let (space_is, space_sizes) = unzip $ unSegSpace space + segment_is = init space_is + segment_dims = init space_sizes + (i_in_segment, segment_size) = last $ unSegSpace space + num_subhistos_per_group = tvExp num_subhistos_per_group_var + segment_size' = toInt64Exp segment_size + + num_segments <- + dPrimVE "num_segments" $ + product $ map toInt64Exp segment_dims + + hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w -> + dPrimV "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S + + sKernelThread "seghist_local" num_groups group_size (segFlat space) $ + virtualiseGroups SegVirt (sExt32 $ unCount groups_per_segment * num_segments) $ \group_id -> do + constants <- kernelConstants <$> askEnv + + flat_segment_id <- dPrimVE "flat_segment_id" $ group_id `quot` sExt32 (unCount groups_per_segment) + gid_in_segment <- dPrimVE "gid_in_segment" $ group_id `rem` sExt32 (unCount groups_per_segment) + -- This pgtid is kind of a "virtualised physical" gtid - not the + -- same thing as the gtid used for the SegHist itself. + pgtid_in_segment <- + dPrimVE "pgtid_in_segment" $ + gid_in_segment * sExt32 (kernelGroupSize constants) + + kernelLocalThreadId constants + threads_per_segment <- + dPrimVE "threads_per_segment" $ + sExt32 $ unCount groups_per_segment * kernelGroupSize constants + + -- Set segment indices. + zipWithM_ dPrimV_ segment_is $ + unflattenIndex (map toInt64Exp segment_dims) $ sExt64 flat_segment_id + + histograms <- forM (zip init_histograms hist_H_chks) $ + \((glob_subhistos, init_local_subhistos), hist_H_chk) -> do + (local_subhistos, do_op) <- init_local_subhistos $ Var $ tvVar hist_H_chk + return (zip glob_subhistos local_subhistos, hist_H_chk, do_op) + + -- Find index of local subhistograms updated by this thread. We + -- try to ensure, as much as possible, that threads in the same + -- warp use different subhistograms, to avoid conflicts. + thread_local_subhisto_i <- + dPrimVE "thread_local_subhisto_i" $ + kernelLocalThreadId constants `rem` num_subhistos_per_group + + let onSlugs f = forM_ (zip slugs histograms) $ \(slug, (dests, hist_H_chk, _)) -> do + let histo_dims = + tvExp hist_H_chk : + map toInt64Exp (shapeDims (histShape (slugOp slug))) + histo_size <- dPrimVE "histo_size" $ product histo_dims + f slug dests (tvExp hist_H_chk) histo_dims histo_size + + let onAllHistograms f = + onSlugs $ \slug dests hist_H_chk histo_dims histo_size -> do + let group_hists_size = num_subhistos_per_group * sExt32 histo_size + init_per_thread <- + dPrimVE "init_per_thread" $ + group_hists_size + `divUp` sExt32 (kernelGroupSize constants) + + forM_ (zip dests (histNeutral $ slugOp slug)) $ + \((dest_global, dest_local), ne) -> + sFor "local_i" init_per_thread $ \i -> do + j <- + dPrimVE "j" $ + i * sExt32 (kernelGroupSize constants) + + kernelLocalThreadId constants + j_offset <- + dPrimVE "j_offset" $ + num_subhistos_per_group * sExt32 histo_size * gid_in_segment + j + + local_subhisto_i <- dPrimVE "local_subhisto_i" $ j `quot` sExt32 histo_size + let local_bucket_is = unflattenIndex histo_dims $ sExt64 $ j `rem` sExt32 histo_size + global_bucket_is = + head local_bucket_is + sExt64 chk_i * hist_H_chk : + tail local_bucket_is + global_subhisto_i <- dPrimVE "global_subhisto_i" $ j_offset `quot` sExt32 histo_size + + sWhen (j .<. group_hists_size) $ + f + dest_local + dest_global + (slugOp slug) + ne + local_subhisto_i + global_subhisto_i + local_bucket_is + global_bucket_is + + sComment "initialize histograms in local memory" $ + onAllHistograms $ \dest_local dest_global op ne local_subhisto_i global_subhisto_i local_bucket_is global_bucket_is -> + sComment "First subhistogram is initialised from global memory; others with neutral element." $ do + let global_is = map Imp.vi64 segment_is ++ [0] ++ global_bucket_is + local_is = sExt64 local_subhisto_i : local_bucket_is + sIf + (global_subhisto_i .==. 0) + (copyDWIMFix dest_local local_is (Var dest_global) global_is) + ( sLoopNest (histShape op) $ \is -> + copyDWIMFix dest_local (local_is ++ is) ne [] + ) + + sOp $ Imp.Barrier Imp.FenceLocal + + kernelLoop pgtid_in_segment threads_per_segment (sExt32 segment_size') $ \ie -> do + dPrimV_ i_in_segment ie + + -- We execute the bucket function once and update each histogram + -- serially. This also involves writing to the mapout arrays if + -- this is the first chunk. + + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = + splitFromEnd (length map_pes) $ + map kernelResultSubExp $ kernelBodyResult kbody + (buckets, vs) = splitAt (length slugs) red_res + perOp = chunks $ map (length . histDest . slugOp) slugs + + sWhen (chk_i .==. 0) $ + sComment "save map-out results" $ + forM_ (zip map_pes map_res) $ \(pe, se) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 space_is) + se + [] + + forM_ (zip4 (map slugOp slugs) histograms buckets (perOp vs)) $ + \( HistOp dest_w _ _ _ shape lam, + (_, hist_H_chk, do_op), + bucket, + vs' + ) -> do + let chk_beg = sExt64 chk_i * tvExp hist_H_chk + bucket' = toInt64Exp bucket + dest_w' = toInt64Exp dest_w + bucket_in_bounds = + bucket' .<. dest_w' + .&&. chk_beg .<=. bucket' + .&&. bucket' .<. (chk_beg + tvExp hist_H_chk) + bucket_is = [sExt64 thread_local_subhisto_i, bucket' - chk_beg] + vs_params = takeLast (length vs') $ lambdaParams lam + + sComment "perform atomic updates" $ + sWhen bucket_in_bounds $ do + dLParams $ lambdaParams lam + sLoopNest shape $ \is -> do + forM_ (zip vs_params vs') $ \(p, v) -> + copyDWIMFix (paramName p) [] v is + do_op (bucket_is ++ is) + + sOp $ Imp.ErrorSync Imp.FenceGlobal + + sComment "Compact the multiple local memory subhistograms to result in global memory" $ + onSlugs $ \slug dests hist_H_chk histo_dims histo_size -> do + bins_per_thread <- + dPrimVE "init_per_thread" $ + histo_size `divUp` sExt64 (kernelGroupSize constants) + + trunc_H <- + dPrimV "trunc_H" $ + sMin64 hist_H_chk $ + toInt64Exp (histWidth (slugOp slug)) + - sExt64 chk_i * head histo_dims + let trunc_histo_dims = + tvExp trunc_H : + map toInt64Exp (shapeDims (histShape (slugOp slug))) + trunc_histo_size <- dPrimVE "histo_size" $ product trunc_histo_dims + + sFor "local_i" bins_per_thread $ \i -> do + j <- + dPrimVE "j" $ + i * sExt64 (kernelGroupSize constants) + + sExt64 (kernelLocalThreadId constants) + sWhen (j .<. trunc_histo_size) $ do + -- We are responsible for compacting the flat bin 'j', which + -- we immediately unflatten. + let local_bucket_is = unflattenIndex histo_dims j + global_bucket_is = + head local_bucket_is + sExt64 chk_i * hist_H_chk : + tail local_bucket_is + dLParams $ lambdaParams $ histOp $ slugOp slug + let (global_dests, local_dests) = unzip dests + (xparams, yparams) = + splitAt (length local_dests) $ + lambdaParams $ histOp $ slugOp slug + + sComment "Read values from subhistogram 0." $ + forM_ (zip xparams local_dests) $ \(xp, subhisto) -> + copyDWIMFix + (paramName xp) + [] + (Var subhisto) + (0 : local_bucket_is) + + sComment "Accumulate based on values in other subhistograms." $ + sFor "subhisto_id" (num_subhistos_per_group - 1) $ \subhisto_id -> do + forM_ (zip yparams local_dests) $ \(yp, subhisto) -> + copyDWIMFix + (paramName yp) + [] + (Var subhisto) + (sExt64 subhisto_id + 1 : local_bucket_is) + compileBody' xparams $ lambdaBody $ histOp $ slugOp slug + + sComment "Put final bucket value in global memory." $ do + let global_is = + map Imp.vi64 segment_is + ++ [sExt64 group_id `rem` unCount groups_per_segment] + ++ global_bucket_is + forM_ (zip xparams global_dests) $ \(xp, global_dest) -> + copyDWIMFix global_dest global_is (Var $ paramName xp) [] + +histKernelLocal :: + TV Int32 -> + Count NumGroups (Imp.TExp Int64) -> + [PatElem KernelsMem] -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + Imp.TExp Int32 -> + [SegHistSlug] -> + KernelBody KernelsMem -> + CallKernelGen () +histKernelLocal num_subhistos_per_group_var groups_per_segment map_pes num_groups group_size space hist_S slugs kbody = do + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + num_subhistos_per_group = tvExp num_subhistos_per_group_var + + emit $ + Imp.DebugPrint "Number of local subhistograms per group" $ + Just $ untyped num_subhistos_per_group + + init_histograms <- + prepareIntermediateArraysLocal num_subhistos_per_group_var groups_per_segment space slugs + + sFor "chk_i" hist_S $ \chk_i -> + histKernelLocalPass + num_subhistos_per_group_var + groups_per_segment + map_pes + num_groups' + group_size' + space + slugs + kbody + init_histograms + hist_S + chk_i + +-- | The maximum number of passes we are willing to accept for this +-- kind of atomic update. +slugMaxLocalMemPasses :: SegHistSlug -> Int +slugMaxLocalMemPasses slug = + case slugAtomicUpdate slug of + AtomicPrim _ -> 3 + AtomicCAS _ -> 4 + AtomicLocking _ -> 6 + +localMemoryCase :: + [PatElem KernelsMem] -> + Imp.TExp Int32 -> + SegSpace -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Imp.TExp Int32 -> + [SegHistSlug] -> + KernelBody KernelsMem -> + CallKernelGen (Imp.TExp Bool, CallKernelGen ()) +localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N _ slugs kbody = do + let space_sizes = segSpaceDims space + segment_dims = init space_sizes + segmented = not $ null segment_dims + + hist_L <- dPrim "hist_L" int32 + sOp $ Imp.GetSizeMax (tvVar hist_L) Imp.SizeLocalMemory + + max_group_size <- dPrim "max_group_size" int32 + sOp $ Imp.GetSizeMax (tvVar max_group_size) Imp.SizeGroup + let group_size = Imp.Count $ Var $ tvVar max_group_size + num_groups <- + fmap (Imp.Count . tvSize) $ + dPrimV "num_groups" $ + hist_T `divUp` sExt32 (toInt64Exp (unCount group_size)) + let num_groups' = toInt64Exp <$> num_groups + group_size' = toInt64Exp <$> group_size + + let r64 = isF64 . ConvOpExp (SIToFP Int64 Float64) . untyped + t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped + + -- M approximation. + hist_m' <- + dPrimVE "hist_m_prime" $ + r64 + ( sMin64 + (sExt64 (tvExp hist_L `quot` hist_el_size)) + (hist_N `divUp` sExt64 (unCount num_groups')) + ) + / r64 hist_H + + let hist_B = unCount group_size' + + -- M in the paper, but not adjusted for asymptotic efficiency. + hist_M0 <- + dPrimVE "hist_M0" $ + sMax64 1 $ sMin64 (t64 hist_m') hist_B + + -- Minimal sequential chunking factor. + let q_small = 2 + + -- The number of segments/histograms produced.. + hist_Nout <- dPrimVE "hist_Nout" $ product $ map toInt64Exp segment_dims + + hist_Nin <- dPrimVE "hist_Nin" $ toInt64Exp $ last space_sizes + + -- Maximum M for work efficiency. + work_asymp_M_max <- + if segmented + then do + hist_T_hist_min <- + dPrimVE "hist_T_hist_min" $ + sExt32 $ + sMin64 (sExt64 hist_Nin * sExt64 hist_Nout) (sExt64 hist_T) + `divUp` sExt64 hist_Nout + + -- Number of groups, rounded up. + let r = hist_T_hist_min `divUp` sExt32 hist_B + + dPrimVE "work_asymp_M_max" $ hist_Nin `quot` (sExt64 r * hist_H) + else + dPrimVE "work_asymp_M_max" $ + (hist_Nout * hist_N) + `quot` ( (q_small * unCount num_groups' * hist_H) + `quot` genericLength slugs + ) + + -- Number of subhistograms per result histogram. + hist_M <- dPrimV "hist_M" $ sExt32 $ sMin64 hist_M0 work_asymp_M_max + + -- hist_M may be zero (which we'll check for below), but we need it + -- for some divisions first, so crudely make a nonzero form. + let hist_M_nonzero = sMax32 1 $ tvExp hist_M + + -- "Cooperation factor" - the number of threads cooperatively + -- working on the same (sub)histogram. + hist_C <- + dPrimVE "hist_C" $ + hist_B `divUp` sExt64 hist_M_nonzero + + emit $ Imp.DebugPrint "local hist_M0" $ Just $ untyped hist_M0 + emit $ Imp.DebugPrint "local work asymp M max" $ Just $ untyped work_asymp_M_max + emit $ Imp.DebugPrint "local C" $ Just $ untyped hist_C + emit $ Imp.DebugPrint "local B" $ Just $ untyped hist_B + emit $ Imp.DebugPrint "local M" $ Just $ untyped $ tvExp hist_M + emit $ + Imp.DebugPrint "local memory needed" $ + Just $ untyped $ hist_H * hist_el_size * sExt64 (tvExp hist_M) + + -- local_mem_needed is what we need to keep a single bucket in local + -- memory - this is an absolute minimum. We can fit anything else + -- by doing multiple passes, although more than a few is + -- (heuristically) not efficient. + local_mem_needed <- + dPrimVE "local_mem_needed" $ + hist_el_size * sExt64 (tvExp hist_M) + hist_S <- + dPrimVE "hist_S" $ + sExt32 $ + (hist_H * local_mem_needed) `divUp` tvExp hist_L + let max_S = case bodyPassage kbody of + MustBeSinglePass -> 1 + MayBeMultiPass -> fromIntegral $ maxinum $ map slugMaxLocalMemPasses slugs + + -- We only use local memory if the number of updates per histogram + -- at least matches the histogram size, as otherwise it is not + -- asymptotically efficient. This mostly matters for the segmented + -- case. + let pick_local = + hist_Nin .>=. hist_H + .&&. (local_mem_needed .<=. tvExp hist_L) + .&&. (hist_S .<=. max_S) + .&&. hist_C .<=. hist_B + .&&. tvExp hist_M .>. 0 + + groups_per_segment + | segmented = num_groups' `divUp` Imp.Count hist_Nout + | otherwise = num_groups' + + run = do + emit $ Imp.DebugPrint "## Using local memory" Nothing + emit $ Imp.DebugPrint "Histogram size (H)" $ Just $ untyped hist_H + emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just $ untyped $ tvExp hist_M + emit $ Imp.DebugPrint "Cooperation level (C)" $ Just $ untyped hist_C + emit $ Imp.DebugPrint "Number of chunks (S)" $ Just $ untyped hist_S + when segmented $ + emit $ Imp.DebugPrint "Groups per segment" $ Just $ untyped $ unCount groups_per_segment + histKernelLocal + hist_M + groups_per_segment + map_pes + num_groups + group_size + space + hist_S + slugs + kbody + + return (pick_local, run) + +-- | Generate code for a segmented histogram called from the host. +compileSegHist :: + Pattern KernelsMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [HistOp KernelsMem] -> + KernelBody KernelsMem -> + CallKernelGen () +compileSegHist (Pattern _ pes) num_groups group_size space ops kbody = do + -- Most of this function is not the histogram part itself, but + -- rather figuring out whether to use a local or global memory + -- strategy, as well as collapsing the subhistograms produced (which + -- are always in global memory, but their number may vary). + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + dims = map toInt64Exp $ segSpaceDims space + + num_red_res = length ops + sum (map (length . histNeutral) ops) + (all_red_pes, map_pes) = splitAt num_red_res pes + segment_size = last dims + + (op_hs, op_seg_hs, slugs) <- unzip3 <$> mapM (computeHistoUsage space) ops + h <- dPrimVE "h" $ Imp.unCount $ sum op_hs + seg_h <- dPrimVE "seg_h" $ Imp.unCount $ sum op_seg_hs + + -- Check for emptyness to avoid division-by-zero. + sUnless (seg_h .==. 0) $ do + -- Maximum group size (or actual, in this case). + let hist_B = unCount group_size' + + -- Size of a histogram. + hist_H <- dPrimVE "hist_H" $ sum $ map (toInt64Exp . histWidth) ops + + -- Size of a single histogram element. Actually the weighted + -- average of histogram elements in cases where we have more than + -- one histogram operation, plus any locks. + let lockSize slug = case slugAtomicUpdate slug of + AtomicLocking {} -> Just $ primByteSize int32 + _ -> Nothing + hist_el_size <- + dPrimVE "hist_el_size" $ + foldl' (+) (h `divUp` hist_H) $ + mapMaybe lockSize slugs + + -- Input elements contributing to each histogram. + hist_N <- dPrimVE "hist_N" segment_size + + -- Compute RF as the average RF over all the histograms. + hist_RF <- + dPrimVE "hist_RF" $ + sum (map (toInt32Exp . histRaceFactor . slugOp) slugs) + `quot` genericLength slugs + + let hist_T = sExt32 $ unCount num_groups' * unCount group_size' + emit $ Imp.DebugPrint "\n# SegHist" Nothing + emit $ Imp.DebugPrint "Number of threads (T)" $ Just $ untyped hist_T + emit $ Imp.DebugPrint "Desired group size (B)" $ Just $ untyped hist_B + emit $ Imp.DebugPrint "Histogram size (H)" $ Just $ untyped hist_H + emit $ Imp.DebugPrint "Input elements per histogram (N)" $ Just $ untyped hist_N + emit $ + Imp.DebugPrint "Number of segments" $ + Just $ untyped $ product $ map (toInt64Exp . snd) segment_dims + emit $ Imp.DebugPrint "Histogram element size (el_size)" $ Just $ untyped hist_el_size + emit $ Imp.DebugPrint "Race factor (RF)" $ Just $ untyped hist_RF + emit $ Imp.DebugPrint "Memory per set of subhistograms per segment" $ Just $ untyped h + emit $ Imp.DebugPrint "Memory per set of subhistograms times segments" $ Just $ untyped seg_h + + (use_local_memory, run_in_local_memory) <- + localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N hist_RF slugs kbody + + sIf use_local_memory run_in_local_memory $ + histKernelGlobal map_pes num_groups group_size space slugs kbody + + let pes_per_op = chunks (map (length . histDest) ops) all_red_pes + + forM_ (zip3 slugs pes_per_op ops) $ \(slug, red_pes, op) -> do + let num_histos = slugNumSubhistos slug + subhistos = map subhistosArray $ slugSubhistos slug + + let unitHistoCase = + -- This is OK because the memory blocks are at least as + -- large as the ones we are supposed to use for the result. + forM_ (zip red_pes subhistos) $ \(pe, subhisto) -> do + pe_mem <- + memLocationName . entryArrayLocation + <$> lookupArray (patElemName pe) + subhisto_mem <- + memLocationName . entryArrayLocation + <$> lookupArray subhisto + emit $ Imp.SetMem pe_mem subhisto_mem $ Space "device" + + sIf (tvExp num_histos .==. 1) unitHistoCase $ do + -- For the segmented reduction, we keep the segment dimensions + -- unchanged. To this, we add two dimensions: one over the number + -- of buckets, and one over the number of subhistograms. This + -- inner dimension is the one that is collapsed in the reduction. + let num_buckets = histWidth op + + bucket_id <- newVName "bucket_id" + subhistogram_id <- newVName "subhistogram_id" + vector_ids <- + mapM (const $ newVName "vector_id") $ + shapeDims $ histShape op + + flat_gtid <- newVName "flat_gtid" + + let lvl = SegThread num_groups group_size SegVirt + segred_space = + SegSpace flat_gtid $ + segment_dims + ++ [(bucket_id, num_buckets)] + ++ zip vector_ids (shapeDims $ histShape op) + ++ [(subhistogram_id, Var $ tvVar num_histos)] + + let segred_op = SegBinOp Commutative (histOp op) (histNeutral op) mempty + compileSegRed' (Pattern [] red_pes) lvl segred_space [segred_op] $ \red_cont -> + red_cont $ + flip map subhistos $ \subhisto -> + ( Var subhisto, + map Imp.vi64 $ + map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids + ) + where + segment_dims = init $ unSegSpace space diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs @@ -1,61 +1,59 @@ -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + -- | Code generation for 'SegMap' is quite straightforward. The only -- trick is virtualisation in case the physical number of threads is -- not sufficient to cover the logical thread space. This is handled -- by having actual workgroups run a loop to imitate multiple workgroups. -module Futhark.CodeGen.ImpGen.Kernels.SegMap - ( compileSegMap ) -where +module Futhark.CodeGen.ImpGen.Kernels.SegMap (compileSegMap) where import Control.Monad.Except - -import Prelude hiding (quot, rem) - -import Futhark.IR.KernelsMem -import Futhark.CodeGen.ImpGen.Kernels.Base -import Futhark.CodeGen.ImpGen import qualified Futhark.CodeGen.ImpCode.Kernels as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.Kernels.Base +import Futhark.IR.KernelsMem import Futhark.Util.IntegralExp (divUp) +import Prelude hiding (quot, rem) -- | Compile 'SegMap' instance code. -compileSegMap :: Pattern KernelsMem - -> SegLevel - -> SegSpace - -> KernelBody KernelsMem - -> CallKernelGen () - +compileSegMap :: + Pattern KernelsMem -> + SegLevel -> + SegSpace -> + KernelBody KernelsMem -> + CallKernelGen () compileSegMap pat lvl space kbody = do let (is, dims) = unzip $ unSegSpace space - dims' <- mapM toExp dims - - num_groups' <- traverse toExp $ segNumGroups lvl - group_size' <- traverse toExp $ segGroupSize lvl + dims' = map toInt64Exp dims + num_groups' = toInt64Exp <$> segNumGroups lvl + group_size' = toInt64Exp <$> segGroupSize lvl case lvl of - SegThread{} -> do + SegThread {} -> do emit $ Imp.DebugPrint "\n# SegMap" Nothing - let virt_num_groups = product dims' `divUp` unCount group_size' + let virt_num_groups = + sExt32 $ product dims' `divUp` unCount group_size' sKernelThread "segmap" num_groups' group_size' (segFlat space) $ virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do - local_tid <- kernelLocalThreadId . kernelConstants <$> askEnv - let global_tid = Imp.vi32 group_id * unCount group_size' + local_tid - - zipWithM_ dPrimV_ is $ unflattenIndex dims' global_tid + local_tid <- kernelLocalThreadId . kernelConstants <$> askEnv + let global_tid = + sExt64 group_id * sExt64 (unCount group_size') + + sExt64 local_tid - sWhen (isActive $ unSegSpace space) $ - compileStms mempty (kernelBodyStms kbody) $ - zipWithM_ (compileThreadResult space) (patternElements pat) $ - kernelBodyResult kbody + zipWithM_ dPrimV_ is $ + map sExt64 $ unflattenIndex (map sExt64 dims') global_tid - SegGroup{} -> + sWhen (isActive $ unSegSpace space) $ + compileStms mempty (kernelBodyStms kbody) $ + zipWithM_ (compileThreadResult space) (patternElements pat) $ + kernelBodyResult kbody + SegGroup {} -> sKernelGroup "segmap_intragroup" num_groups' group_size' (segFlat space) $ do - let virt_num_groups = product dims' - precomputeSegOpIDs (kernelBodyStms kbody) $ - virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do - - zipWithM_ dPrimV_ is $ unflattenIndex dims' $ Imp.vi32 group_id + let virt_num_groups = sExt32 $ product dims' + precomputeSegOpIDs (kernelBodyStms kbody) $ + virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do + zipWithM_ dPrimV_ is $ unflattenIndex dims' $ sExt64 group_id - compileStms mempty (kernelBodyStms kbody) $ - zipWithM_ (compileGroupResult space) (patternElements pat) $ - kernelBodyResult kbody + compileStms mempty (kernelBodyStms kbody) $ + zipWithM_ (compileGroupResult space) (patternElements pat) $ + kernelBodyResult kbody diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs @@ -1,5 +1,6 @@ -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + -- | We generate code for non-segmented/single-segment SegRed using -- the basic approach outlined in the paper "Design and GPGPU -- Performance of Futhark’s Redomap Construct" (ARRAY '16). The main @@ -42,27 +43,25 @@ -- large strategy, but at most 50% of the threads in the group would -- have any element to read, which becomes highly inefficient. module Futhark.CodeGen.ImpGen.Kernels.SegRed - ( compileSegRed - , compileSegRed' - , DoSegBody + ( compileSegRed, + compileSegRed', + DoSegBody, ) - where +where import Control.Monad.Except -import Data.Maybe import Data.List (genericLength, zip7) - -import Prelude hiding (quot, rem) - -import Futhark.Error -import Futhark.Transform.Rename -import Futhark.IR.KernelsMem +import Data.Maybe import qualified Futhark.CodeGen.ImpCode.Kernels as Imp import Futhark.CodeGen.ImpGen import Futhark.CodeGen.ImpGen.Kernels.Base +import Futhark.Error +import Futhark.IR.KernelsMem import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.Transform.Rename import Futhark.Util (chunks) import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) -- | The maximum number of operators we support in a single SegRed. -- This limit arises out of the static allocation of counters. @@ -73,47 +72,53 @@ -- for saving the results of the body. The results should be -- represented as a pairing of a t'SubExp' along with a list of -- indexes into that 'SubExp' for reading the result. -type DoSegBody = ([(SubExp, [Imp.Exp])] -> InKernelGen ()) -> InKernelGen () +type DoSegBody = ([(SubExp, [Imp.TExp Int64])] -> InKernelGen ()) -> InKernelGen () -- | Compile 'SegRed' instance to host-level code with calls to -- various kernels. -compileSegRed :: Pattern KernelsMem - -> SegLevel -> SegSpace - -> [SegBinOp KernelsMem] - -> KernelBody KernelsMem - -> CallKernelGen () +compileSegRed :: + Pattern KernelsMem -> + SegLevel -> + SegSpace -> + [SegBinOp KernelsMem] -> + KernelBody KernelsMem -> + CallKernelGen () compileSegRed pat lvl space reds body = compileSegRed' pat lvl space reds $ \red_cont -> - compileStms mempty (kernelBodyStms body) $ do - let (red_res, map_res) = splitAt (segBinOpResults reds) $ kernelBodyResult body + compileStms mempty (kernelBodyStms body) $ do + let (red_res, map_res) = splitAt (segBinOpResults reds) $ kernelBodyResult body - sComment "save map-out results" $ do - let map_arrs = drop (segBinOpResults reds) $ patternElements pat - zipWithM_ (compileThreadResult space) map_arrs map_res + sComment "save map-out results" $ do + let map_arrs = drop (segBinOpResults reds) $ patternElements pat + zipWithM_ (compileThreadResult space) map_arrs map_res - red_cont $ zip (map kernelResultSubExp red_res) $ repeat [] + red_cont $ zip (map kernelResultSubExp red_res) $ repeat [] -- | Like 'compileSegRed', but where the body is a monadic action. -compileSegRed' :: Pattern KernelsMem - -> SegLevel -> SegSpace - -> [SegBinOp KernelsMem] - -> DoSegBody - -> CallKernelGen () +compileSegRed' :: + Pattern KernelsMem -> + SegLevel -> + SegSpace -> + [SegBinOp KernelsMem] -> + DoSegBody -> + CallKernelGen () compileSegRed' pat lvl space reds body | genericLength reds > maxNumOps = - compilerLimitationS $ + compilerLimitationS $ "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported." - | [(_, Constant (IntValue (Int32Value 1))), _] <- unSegSpace space = - nonsegmentedReduction pat num_groups group_size space reds body + | [(_, Constant (IntValue (Int64Value 1))), _] <- unSegSpace space = + nonsegmentedReduction pat num_groups group_size space reds body | otherwise = do - group_size' <- toExp $ unCount group_size - segment_size <- toExp $ last $ segSpaceDims space - let use_small_segments = segment_size * 2 .<. group_size' - sIf use_small_segments - (smallSegmentsReduction pat num_groups group_size space reds body) - (largeSegmentsReduction pat num_groups group_size space reds body) - where num_groups = segNumGroups lvl - group_size = segGroupSize lvl + let group_size' = toInt32Exp $ unCount group_size + segment_size = toInt32Exp $ last $ segSpaceDims space + use_small_segments = segment_size * 2 .<. group_size' + sIf + use_small_segments + (smallSegmentsReduction pat num_groups group_size space reds body) + (largeSegmentsReduction pat num_groups group_size space reds body) + where + num_groups = segNumGroups lvl + group_size = segGroupSize lvl -- | Prepare intermediate arrays for the reduction. Prim-typed -- arguments go in local memory (so we need to do the allocation of @@ -121,9 +126,11 @@ -- global memory. Allocations for the former have already been -- performed. This policy is baked into how the allocations are done -- in ExplicitAllocations. -intermediateArrays :: Count GroupSize SubExp -> SubExp - -> SegBinOp KernelsMem - -> InKernelGen [VName] +intermediateArrays :: + Count GroupSize SubExp -> + SubExp -> + SegBinOp KernelsMem -> + InKernelGen [VName] intermediateArrays (Count group_size) num_threads (SegBinOp _ red_op nes _) = do let red_op_params = lambdaParams red_op (red_acc_params, _) = splitAt (length nes) red_op_params @@ -132,7 +139,7 @@ MemArray pt shape _ (ArrayIn mem _) -> do let shape' = Shape [num_threads] <> shape sArray "red_arr" pt shape' $ - ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape' + ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape' _ -> do let pt = elemType $ paramType p shape = Shape [group_size] @@ -144,121 +151,168 @@ -- because they are also used for keeping vectorised accumulators for -- first-stage reduction, if necessary. When actually storing group -- results, the first index is set to 0. -groupResultArrays :: Count NumGroups SubExp -> Count GroupSize SubExp - -> [SegBinOp KernelsMem] - -> CallKernelGen [[VName]] +groupResultArrays :: + Count NumGroups SubExp -> + Count GroupSize SubExp -> + [SegBinOp KernelsMem] -> + CallKernelGen [[VName]] groupResultArrays (Count virt_num_groups) (Count group_size) reds = forM reds $ \(SegBinOp _ lam _ shape) -> forM (lambdaReturnType lam) $ \t -> do - let pt = elemType t - full_shape = Shape [group_size, virt_num_groups] <> shape <> arrayShape t - -- Move the groupsize dimension last to ensure coalesced - -- memory access. - perm = [1..shapeRank full_shape-1] ++ [0] - sAllocArrayPerm "group_res_arr" pt full_shape (Space "device") perm + let pt = elemType t + full_shape = Shape [group_size, virt_num_groups] <> shape <> arrayShape t + -- Move the groupsize dimension last to ensure coalesced + -- memory access. + perm = [1 .. shapeRank full_shape -1] ++ [0] + sAllocArrayPerm "group_res_arr" pt full_shape (Space "device") perm -nonsegmentedReduction :: Pattern KernelsMem - -> Count NumGroups SubExp -> Count GroupSize SubExp -> SegSpace - -> [SegBinOp KernelsMem] - -> DoSegBody - -> CallKernelGen () +nonsegmentedReduction :: + Pattern KernelsMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegBinOp KernelsMem] -> + DoSegBody -> + CallKernelGen () nonsegmentedReduction segred_pat num_groups group_size space reds body = do let (gtids, dims) = unzip $ unSegSpace space - dims' <- mapM toExp dims - - num_groups' <- traverse toExp num_groups - group_size' <- traverse toExp group_size - - let global_tid = Imp.vi32 $ segFlat space + dims' = map toInt64Exp dims + num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + global_tid = Imp.vi32 $ segFlat space w = last dims' counter <- sStaticArray "counter" (Space "device") int32 $ - Imp.ArrayValues $ replicate (fromIntegral maxNumOps) $ IntValue $ Int32Value 0 + Imp.ArrayValues $ replicate (fromIntegral maxNumOps) $ IntValue $ Int32Value 0 reds_group_res_arrs <- groupResultArrays num_groups group_size reds - num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size' + num_threads <- + dPrimV "num_threads" $ + unCount num_groups' * unCount group_size' emit $ Imp.DebugPrint "\n# SegRed" Nothing sKernelThread "segred_nonseg" num_groups' group_size' (segFlat space) $ do constants <- kernelConstants <$> askEnv sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" - reds_arrs <- mapM (intermediateArrays group_size (Var num_threads)) reds + reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds -- Since this is the nonsegmented case, all outer segment IDs must -- necessarily be 0. - forM_ gtids $ \v -> dPrimV_ v 0 + forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int32) let num_elements = Imp.elements w - let elems_per_thread = num_elements `divUp` Imp.elements (kernelNumThreads constants) + elems_per_thread = + num_elements + `divUp` Imp.elements (sExt64 (kernelNumThreads constants)) - slugs <- mapM (segBinOpSlug - (kernelLocalThreadId constants) - (kernelGroupId constants)) $ - zip3 reds reds_arrs reds_group_res_arrs + slugs <- + mapM + ( segBinOpSlug + (kernelLocalThreadId constants) + (kernelGroupId constants) + ) + $ zip3 reds reds_arrs reds_group_res_arrs reds_op_renamed <- - reductionStageOne constants (zip gtids dims') num_elements - global_tid elems_per_thread num_threads - slugs body + reductionStageOne + constants + (zip gtids dims') + num_elements + global_tid + elems_per_thread + (tvVar num_threads) + slugs + body - let segred_pes = chunks (map (length . segBinOpNeutral) reds) $ - patternElements segred_pat - forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes - slugs reds_op_renamed [0..]) $ - \(SegBinOp _ red_op nes _, - red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do - let (red_x_params, red_y_params) = splitAt (length nes) $ lambdaParams red_op - reductionStageTwo constants pes (kernelGroupId constants) 0 [0] 0 - (kernelNumGroups constants) slug red_x_params red_y_params - red_op_renamed nes - 1 counter (ValueExp $ IntValue $ Int32Value i) - sync_arr group_res_arrs red_arrs + let segred_pes = + chunks (map (length . segBinOpNeutral) reds) $ + patternElements segred_pat + forM_ + ( zip7 + reds + reds_arrs + reds_group_res_arrs + segred_pes + slugs + reds_op_renamed + [0 ..] + ) + $ \( SegBinOp _ red_op nes _, + red_arrs, + group_res_arrs, + pes, + slug, + red_op_renamed, + i + ) -> do + let (red_x_params, red_y_params) = splitAt (length nes) $ lambdaParams red_op + reductionStageTwo + constants + pes + (kernelGroupId constants) + 0 + [0] + 0 + (sExt64 $ kernelNumGroups constants) + slug + red_x_params + red_y_params + red_op_renamed + nes + 1 + counter + (fromInteger i) + sync_arr + group_res_arrs + red_arrs -smallSegmentsReduction :: Pattern KernelsMem - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> SegSpace - -> [SegBinOp KernelsMem] - -> DoSegBody - -> CallKernelGen () +smallSegmentsReduction :: + Pattern KernelsMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegBinOp KernelsMem] -> + DoSegBody -> + CallKernelGen () smallSegmentsReduction (Pattern _ segred_pes) num_groups group_size space reds body = do let (gtids, dims) = unzip $ unSegSpace space - dims' <- mapM toExp dims + dims' = map toInt64Exp dims + segment_size = last dims' - let segment_size = last dims' -- Careful to avoid division by zero now. - segment_size_nonzero_v <- dPrimV "segment_size_nonzero" $ - BinOpExp (SMax Int32) 1 segment_size + segment_size_nonzero <- + dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size - num_groups' <- traverse toExp num_groups - group_size' <- traverse toExp group_size + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size' - let segment_size_nonzero = Imp.var segment_size_nonzero_v int32 - num_segments = product $ init dims' + let num_segments = product $ init dims' segments_per_group = unCount group_size' `quot` segment_size_nonzero - required_groups = num_segments `divUp` segments_per_group + required_groups = sExt32 $ num_segments `divUp` segments_per_group emit $ Imp.DebugPrint "\n# SegRed-small" Nothing - emit $ Imp.DebugPrint "num_segments" $ Just num_segments - emit $ Imp.DebugPrint "segment_size" $ Just segment_size - emit $ Imp.DebugPrint "segments_per_group" $ Just segments_per_group - emit $ Imp.DebugPrint "required_groups" $ Just required_groups + emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments + emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size + emit $ Imp.DebugPrint "segments_per_group" $ Just $ untyped segments_per_group + emit $ Imp.DebugPrint "required_groups" $ Just $ untyped required_groups sKernelThread "segred_small" num_groups' group_size' (segFlat space) $ do constants <- kernelConstants <$> askEnv - reds_arrs <- mapM (intermediateArrays group_size (Var num_threads)) reds + reds_arrs <- mapM (intermediateArrays group_size (Var $ tvVar num_threads)) reds -- We probably do not have enough actual workgroups to cover the -- entire iteration space. Some groups thus have to perform double -- duty; we put an outer loop to accomplish this. - virtualiseGroups SegVirt required_groups $ \group_id_var' -> do - let group_id' = Imp.vi32 group_id_var' + virtualiseGroups SegVirt required_groups $ \group_id' -> do -- Compute the 'n' input indices. The outer 'n-1' correspond to -- the segment ID, and are computed from the group id. The inner -- is computed from the local thread id, and may be out-of-bounds. - let ltid = kernelLocalThreadId constants - segment_index = (ltid `quot` segment_size_nonzero) + (group_id' * segments_per_group) + let ltid = sExt64 $ kernelLocalThreadId constants + segment_index = + (ltid `quot` segment_size_nonzero) + + (sExt64 group_id' * sExt64 segments_per_group) index_within_segment = ltid `rem` segment_size zipWithM_ dPrimV_ (init gtids) $ unflattenIndex (init dims') segment_index @@ -266,83 +320,105 @@ let out_of_bounds = forM_ (zip reds reds_arrs) $ \(SegBinOp _ _ nes _, red_arrs) -> - forM_ (zip red_arrs nes) $ \(arr, ne) -> - copyDWIMFix arr [ltid] ne [] + forM_ (zip red_arrs nes) $ \(arr, ne) -> + copyDWIMFix arr [ltid] ne [] in_bounds = body $ \red_res -> - sComment "save results to be reduced" $ do - let red_dests = zip (concat reds_arrs) $ repeat [ltid] - forM_ (zip red_dests red_res) $ \((d,d_is), (res, res_is)) -> - copyDWIMFix d d_is res res_is + sComment "save results to be reduced" $ do + let red_dests = zip (concat reds_arrs) $ repeat [ltid] + forM_ (zip red_dests red_res) $ \((d, d_is), (res, res_is)) -> + copyDWIMFix d d_is res res_is sComment "apply map function if in bounds" $ - sIf (segment_size .>. 0 .&&. - isActive (init $ zip gtids dims) .&&. - ltid .<. segment_size * segments_per_group) in_bounds out_of_bounds + sIf + ( segment_size .>. 0 + .&&. isActive (init $ zip gtids dims) + .&&. ltid .<. segment_size * segments_per_group + ) + in_bounds + out_of_bounds sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier. - - let crossesSegment from to = (to-from) .>. (to `rem` segment_size) + let crossesSegment from to = + (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size) sWhen (segment_size .>. 0) $ sComment "perform segmented scan to imitate reduction" $ - forM_ (zip reds reds_arrs) $ \(SegBinOp _ red_op _ _, red_arrs) -> - groupScan (Just crossesSegment) (Imp.vi32 num_threads) - (segment_size*segments_per_group) red_op red_arrs + forM_ (zip reds reds_arrs) $ \(SegBinOp _ red_op _ _, red_arrs) -> + groupScan + (Just crossesSegment) + (sExt64 $ tvExp num_threads) + (segment_size * segments_per_group) + red_op + red_arrs sOp $ Imp.Barrier Imp.FenceLocal sComment "save final values of segments" $ - sWhen (group_id' * segments_per_group + ltid .<. num_segments .&&. - ltid .<. segments_per_group) $ - forM_ (zip segred_pes (concat reds_arrs)) $ \(pe, arr) -> do - -- Figure out which segment result this thread should write... - let flat_segment_index = group_id' * segments_per_group + ltid - gtids' = unflattenIndex (init dims') flat_segment_index - copyDWIMFix (patElemName pe) gtids' - (Var arr) [(ltid+1) * segment_size_nonzero - 1] + sWhen + ( sExt64 group_id' * segments_per_group + sExt64 ltid .<. num_segments + .&&. ltid .<. segments_per_group + ) + $ forM_ (zip segred_pes (concat reds_arrs)) $ \(pe, arr) -> do + -- Figure out which segment result this thread should write... + let flat_segment_index = + sExt64 group_id' * segments_per_group + sExt64 ltid + gtids' = + unflattenIndex (init dims') flat_segment_index + copyDWIMFix + (patElemName pe) + gtids' + (Var arr) + [(ltid + 1) * segment_size_nonzero - 1] -- Finally another barrier, because we will be writing to the -- local memory array first thing in the next iteration. sOp $ Imp.Barrier Imp.FenceLocal -largeSegmentsReduction :: Pattern KernelsMem - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> SegSpace - -> [SegBinOp KernelsMem] - -> DoSegBody - -> CallKernelGen () +largeSegmentsReduction :: + Pattern KernelsMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegBinOp KernelsMem] -> + DoSegBody -> + CallKernelGen () largeSegmentsReduction segred_pat num_groups group_size space reds body = do let (gtids, dims) = unzip $ unSegSpace space - dims' <- mapM toExp dims - let segment_size = last dims' + dims' = map toInt64Exp dims num_segments = product $ init dims' - - num_groups' <- traverse toExp num_groups - group_size' <- traverse toExp group_size + segment_size = last dims' + num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size (groups_per_segment, elems_per_thread) <- - groupsPerSegmentAndElementsPerThread segment_size num_segments - num_groups' group_size' - virt_num_groups <- dPrimV "virt_num_groups" $ - groups_per_segment * num_segments + groupsPerSegmentAndElementsPerThread + segment_size + num_segments + num_groups' + group_size' + virt_num_groups <- + dPrimV "virt_num_groups" $ + groups_per_segment * num_segments - num_threads <- dPrimV "num_threads" $ - unCount num_groups' * unCount group_size' + num_threads <- + dPrimV "num_threads" $ + unCount num_groups' * unCount group_size' - threads_per_segment <- dPrimV "threads_per_segment" $ - groups_per_segment * unCount group_size' + threads_per_segment <- + dPrimV "threads_per_segment" $ + groups_per_segment * unCount group_size' emit $ Imp.DebugPrint "\n# SegRed-large" Nothing - emit $ Imp.DebugPrint "num_segments" $ Just num_segments - emit $ Imp.DebugPrint "segment_size" $ Just segment_size - emit $ Imp.DebugPrint "virt_num_groups" $ Just $ Imp.vi32 virt_num_groups - emit $ Imp.DebugPrint "num_groups" $ Just $ Imp.unCount num_groups' - emit $ Imp.DebugPrint "group_size" $ Just $ Imp.unCount group_size' - emit $ Imp.DebugPrint "elems_per_thread" $ Just $ Imp.unCount elems_per_thread - emit $ Imp.DebugPrint "groups_per_segment" $ Just groups_per_segment + emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments + emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size + emit $ Imp.DebugPrint "virt_num_groups" $ Just $ untyped $ tvExp virt_num_groups + emit $ Imp.DebugPrint "num_groups" $ Just $ untyped $ Imp.unCount num_groups' + emit $ Imp.DebugPrint "group_size" $ Just $ untyped $ Imp.unCount group_size' + emit $ Imp.DebugPrint "elems_per_thread" $ Just $ untyped $ Imp.unCount elems_per_thread + emit $ Imp.DebugPrint "groups_per_segment" $ Just $ untyped groups_per_segment - reds_group_res_arrs <- groupResultArrays (Count (Var virt_num_groups)) group_size reds + reds_group_res_arrs <- groupResultArrays (Count (tvSize virt_num_groups)) group_size reds -- In principle we should have a counter for every segment. Since -- the number of segments is a dynamic quantity, we would have to @@ -357,91 +433,133 @@ let num_counters = fromIntegral maxNumOps * 1024 counter <- sStaticArray "counter" (Space "device") int32 $ - Imp.ArrayZeros num_counters + Imp.ArrayZeros num_counters sKernelThread "segred_large" num_groups' group_size' (segFlat space) $ do constants <- kernelConstants <$> askEnv - reds_arrs <- mapM (intermediateArrays group_size (Var num_threads)) reds + reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" -- We probably do not have enough actual workgroups to cover the -- entire iteration space. Some groups thus have to perform double -- duty; we put an outer loop to accomplish this. - virtualiseGroups SegVirt (Imp.vi32 virt_num_groups) $ \group_id_var -> do + virtualiseGroups SegVirt (sExt32 (tvExp virt_num_groups)) $ \group_id -> do let segment_gtids = init gtids w = last dims - group_id = Imp.vi32 group_id_var local_tid = kernelLocalThreadId constants - flat_segment_id <- dPrimVE "flat_segment_id" $ - group_id `quot` groups_per_segment + flat_segment_id <- + dPrimVE "flat_segment_id" $ + group_id `quot` sExt32 groups_per_segment - global_tid <- dPrimVE "global_tid" $ - (group_id * unCount group_size' + local_tid) - `rem` (unCount group_size' * groups_per_segment) + global_tid <- + dPrimVE "global_tid" $ + (sExt64 group_id * sExt64 (unCount group_size') + sExt64 local_tid) + `rem` (sExt64 (unCount group_size') * groups_per_segment) - let first_group_for_segment = flat_segment_id * groups_per_segment + let first_group_for_segment = sExt64 flat_segment_id * groups_per_segment - zipWithM_ dPrimV_ segment_gtids $ unflattenIndex (init dims') flat_segment_id - dPrim_ (last gtids) int32 - num_elements <- Imp.elements <$> toExp w + zipWithM_ dPrimV_ segment_gtids $ + unflattenIndex (init dims') $ sExt64 flat_segment_id + dPrim_ (last gtids) int64 + let num_elements = Imp.elements $ toInt64Exp w - slugs <- mapM (segBinOpSlug local_tid group_id) $ - zip3 reds reds_arrs reds_group_res_arrs + slugs <- + mapM (segBinOpSlug local_tid group_id) $ + zip3 reds reds_arrs reds_group_res_arrs reds_op_renamed <- - reductionStageOne constants (zip gtids dims') num_elements - global_tid elems_per_thread threads_per_segment - slugs body + reductionStageOne + constants + (zip gtids dims') + num_elements + (sExt32 global_tid) + elems_per_thread + (tvVar threads_per_segment) + slugs + body - let segred_pes = chunks (map (length . segBinOpNeutral) reds) $ - patternElements segred_pat + let segred_pes = + chunks (map (length . segBinOpNeutral) reds) $ + patternElements segred_pat multiple_groups_per_segment = - forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes - slugs reds_op_renamed [0..]) $ - \(SegBinOp _ red_op nes _, red_arrs, group_res_arrs, pes, - slug, red_op_renamed, i) -> do - let (red_x_params, red_y_params) = - splitAt (length nes) $ lambdaParams red_op - reductionStageTwo constants pes - group_id flat_segment_id (map (`Imp.var` int32) segment_gtids) - first_group_for_segment groups_per_segment - slug red_x_params red_y_params red_op_renamed nes - (fromIntegral num_counters) counter (ValueExp $ IntValue $ Int32Value i) - sync_arr group_res_arrs red_arrs + forM_ + ( zip7 + reds + reds_arrs + reds_group_res_arrs + segred_pes + slugs + reds_op_renamed + [0 ..] + ) + $ \( SegBinOp _ red_op nes _, + red_arrs, + group_res_arrs, + pes, + slug, + red_op_renamed, + i + ) -> do + let (red_x_params, red_y_params) = + splitAt (length nes) $ lambdaParams red_op + reductionStageTwo + constants + pes + group_id + flat_segment_id + (map Imp.vi64 segment_gtids) + (sExt64 first_group_for_segment) + groups_per_segment + slug + red_x_params + red_y_params + red_op_renamed + nes + (fromIntegral num_counters) + counter + (fromInteger i) + sync_arr + group_res_arrs + red_arrs one_group_per_segment = comment "first thread in group saves final result to memory" $ - forM_ (zip slugs segred_pes) $ \(slug, pes) -> - sWhen (local_tid .==. 0) $ - forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) -> - copyDWIMFix (patElemName v) (map (`Imp.var` int32) segment_gtids) (Var acc) acc_is + forM_ (zip slugs segred_pes) $ \(slug, pes) -> + sWhen (local_tid .==. 0) $ + forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) -> + copyDWIMFix (patElemName v) (map Imp.vi64 segment_gtids) (Var acc) acc_is sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment -- Careful to avoid division by zero here. We have at least one group -- per segment. -groupsPerSegmentAndElementsPerThread :: Imp.Exp -> Imp.Exp - -> Count NumGroups Imp.Exp -> Count GroupSize Imp.Exp - -> CallKernelGen (Imp.Exp, Imp.Count Imp.Elements Imp.Exp) +groupsPerSegmentAndElementsPerThread :: + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + CallKernelGen + ( Imp.TExp Int64, + Imp.Count Imp.Elements (Imp.TExp Int64) + ) groupsPerSegmentAndElementsPerThread segment_size num_segments num_groups_hint group_size = do groups_per_segment <- dPrimVE "groups_per_segment" $ - unCount num_groups_hint `divUp` BinOpExp (SMax Int32) 1 num_segments + unCount num_groups_hint `divUp` sMax64 1 num_segments elements_per_thread <- dPrimVE "elements_per_thread" $ - segment_size `divUp` (unCount group_size * groups_per_segment) + segment_size `divUp` (unCount group_size * groups_per_segment) return (groups_per_segment, Imp.elements elements_per_thread) -- | A SegBinOp with auxiliary information. -data SegBinOpSlug = - SegBinOpSlug - { slugOp :: SegBinOp KernelsMem - , slugArrs :: [VName] - -- ^ The arrays used for computing the intra-group reduction +data SegBinOpSlug = SegBinOpSlug + { slugOp :: SegBinOp KernelsMem, + -- | The arrays used for computing the intra-group reduction -- (either local or global memory). - , slugAccs :: [(VName, [Imp.Exp])] - -- ^ Places to store accumulator in stage 1 reduction. + slugArrs :: [VName], + -- | Places to store accumulator in stage 1 reduction. + slugAccs :: [(VName, [Imp.TExp Int64])] } slugBody :: SegBinOpSlug -> Body KernelsMem @@ -463,70 +581,71 @@ accParams slug = take (length (slugNeutral slug)) $ slugParams slug nextParams slug = drop (length (slugNeutral slug)) $ slugParams slug -segBinOpSlug :: Imp.Exp -> Imp.Exp -> (SegBinOp KernelsMem, [VName], [VName]) -> InKernelGen SegBinOpSlug +segBinOpSlug :: Imp.TExp Int32 -> Imp.TExp Int32 -> (SegBinOp KernelsMem, [VName], [VName]) -> InKernelGen SegBinOpSlug segBinOpSlug local_tid group_id (op, group_res_arrs, param_arrs) = - SegBinOpSlug op group_res_arrs <$> - zipWithM mkAcc (lambdaParams (segBinOpLambda op)) param_arrs - where mkAcc p param_arr - | Prim t <- paramType p, - shapeRank (segBinOpShape op) == 0 = do - acc <- dPrim (baseString (paramName p) <> "_acc") t - return (acc, []) - | otherwise = - return (param_arr, [local_tid, group_id]) + SegBinOpSlug op group_res_arrs + <$> zipWithM mkAcc (lambdaParams (segBinOpLambda op)) param_arrs + where + mkAcc p param_arr + | Prim t <- paramType p, + shapeRank (segBinOpShape op) == 0 = do + acc <- dPrim (baseString (paramName p) <> "_acc") t + return (tvVar acc, []) + | otherwise = + return (param_arr, [sExt64 local_tid, sExt64 group_id]) -reductionStageZero :: KernelConstants - -> [(VName, Imp.Exp)] - -> Imp.Count Imp.Elements Imp.Exp - -> Imp.Exp - -> Imp.Count Imp.Elements Imp.Exp - -> VName - -> [SegBinOpSlug] - -> DoSegBody - -> InKernelGen ([Lambda KernelsMem], InKernelGen ()) +reductionStageZero :: + KernelConstants -> + [(VName, Imp.TExp Int64)] -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + Imp.TExp Int32 -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + VName -> + [SegBinOpSlug] -> + DoSegBody -> + InKernelGen ([Lambda KernelsMem], InKernelGen ()) reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do let (gtids, _dims) = unzip ispace - gtid = last gtids - local_tid = kernelLocalThreadId constants + gtid = mkTV (last gtids) int64 + local_tid = sExt64 $ kernelLocalThreadId constants -- Figure out how many elements this thread should process. - chunk_size <- dPrim "chunk_size" int32 + chunk_size <- dPrim "chunk_size" int64 let ordering = case slugsComm slugs of - Commutative -> SplitStrided $ Var threads_per_segment - Noncommutative -> SplitContiguous - computeThreadChunkSize ordering global_tid elems_per_thread num_elements chunk_size + Commutative -> SplitStrided $ Var threads_per_segment + Noncommutative -> SplitContiguous + computeThreadChunkSize ordering (sExt64 global_tid) elems_per_thread num_elements chunk_size dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs sComment "neutral-initialise the accumulators" $ forM_ slugs $ \slug -> - forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) -> - sLoopNest (slugShape slug) $ \vec_is -> - copyDWIMFix acc (acc_is++vec_is) ne [] + forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) -> + sLoopNest (slugShape slug) $ \vec_is -> + copyDWIMFix acc (acc_is ++ vec_is) ne [] slugs_op_renamed <- mapM (renameLambda . segBinOpLambda . slugOp) slugs let doTheReduction = forM_ (zip slugs_op_renamed slugs) $ \(slug_op_renamed, slug) -> - sLoopNest (slugShape slug) $ \vec_is -> do - comment "to reduce current chunk, first store our result in memory" $ do - forM_ (zip (slugParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> - copyDWIMFix (paramName p) [] (Var acc) (acc_is++vec_is) - - forM_ (zip (slugArrs slug) (slugParams slug)) $ \(arr, p) -> - when (primType $ paramType p) $ - copyDWIMFix arr [local_tid] (Var $ paramName p) [] + sLoopNest (slugShape slug) $ \vec_is -> do + comment "to reduce current chunk, first store our result in memory" $ do + forM_ (zip (slugParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> + copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is) - sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier. + forM_ (zip (slugArrs slug) (slugParams slug)) $ \(arr, p) -> + when (primType $ paramType p) $ + copyDWIMFix arr [local_tid] (Var $ paramName p) [] - groupReduce (kernelGroupSize constants) slug_op_renamed (slugArrs slug) + sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier. + groupReduce (sExt32 (kernelGroupSize constants)) slug_op_renamed (slugArrs slug) - sOp $ Imp.Barrier Imp.FenceLocal + sOp $ Imp.Barrier Imp.FenceLocal - sComment "first thread saves the result in accumulator" $ - sWhen (local_tid .==. 0) $ - forM_ (zip (slugAccs slug) (lambdaParams slug_op_renamed)) $ \((acc, acc_is), p) -> - copyDWIMFix acc (acc_is++vec_is) (Var $ paramName p) [] + sComment "first thread saves the result in accumulator" $ + sWhen (local_tid .==. 0) $ + forM_ (zip (slugAccs slug) (lambdaParams slug_op_renamed)) $ \((acc, acc_is), p) -> + copyDWIMFix acc (acc_is ++ vec_is) (Var $ paramName p) [] -- If this is a non-commutative reduction, each thread must run the -- loop the same number of iterations, because we will be performing @@ -534,42 +653,47 @@ let comm = slugsComm slugs (bound, check_bounds) = case comm of - Commutative -> (Imp.var chunk_size int32, id) - Noncommutative -> (Imp.unCount elems_per_thread, - sWhen (Imp.var gtid int32 .<. Imp.unCount num_elements)) + Commutative -> (tvExp chunk_size, id) + Noncommutative -> + ( Imp.unCount elems_per_thread, + sWhen (tvExp gtid .<. Imp.unCount num_elements) + ) sFor "i" bound $ \i -> do - gtid <-- - case comm of + gtid + <-- case comm of Commutative -> - global_tid + - Imp.var threads_per_segment int32 * i + sExt64 global_tid + + Imp.vi64 threads_per_segment * i Noncommutative -> - let index_in_segment = global_tid `quot` kernelGroupSize constants - in local_tid + - (index_in_segment * Imp.unCount elems_per_thread + i) * - kernelGroupSize constants - - check_bounds $ sComment "apply map function" $ - body $ \all_red_res -> do + let index_in_segment = global_tid `quot` sExt32 (kernelGroupSize constants) + in sExt64 local_tid + + (sExt64 index_in_segment * Imp.unCount elems_per_thread + i) + * sExt64 (kernelGroupSize constants) - let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res + check_bounds $ + sComment "apply map function" $ + body $ \all_red_res -> do + let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res - forM_ (zip slugs slugs_res) $ \(slug, red_res) -> - sLoopNest (slugShape slug) $ \vec_is -> do - sComment "load accumulator" $ - forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> - copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is) - sComment "load new values" $ - forM_ (zip (nextParams slug) red_res) $ \(p, (res, res_is)) -> - copyDWIMFix (paramName p) [] res (res_is ++ vec_is) - sComment "apply reduction operator" $ - compileStms mempty (bodyStms $ slugBody slug) $ - sComment "store in accumulator" $ - forM_ (zip - (slugAccs slug) - (bodyResult $ slugBody slug)) $ \((acc, acc_is), se) -> - copyDWIMFix acc (acc_is ++ vec_is) se [] + forM_ (zip slugs slugs_res) $ \(slug, red_res) -> + sLoopNest (slugShape slug) $ \vec_is -> do + sComment "load accumulator" $ + forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> + copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is) + sComment "load new values" $ + forM_ (zip (nextParams slug) red_res) $ \(p, (res, res_is)) -> + copyDWIMFix (paramName p) [] res (res_is ++ vec_is) + sComment "apply reduction operator" $ + compileStms mempty (bodyStms $ slugBody slug) $ + sComment "store in accumulator" $ + forM_ + ( zip + (slugAccs slug) + (bodyResult $ slugBody slug) + ) + $ \((acc, acc_is), se) -> + copyDWIMFix acc (acc_is ++ vec_is) se [] case comm of Noncommutative -> do @@ -577,23 +701,24 @@ sComment "first thread keeps accumulator; others reset to neutral element" $ do let reset_to_neutral = forM_ slugs $ \slug -> - forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) -> - sLoopNest (slugShape slug) $ \vec_is -> - copyDWIMFix acc (acc_is++vec_is) ne [] + forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) -> + sLoopNest (slugShape slug) $ \vec_is -> + copyDWIMFix acc (acc_is ++ vec_is) ne [] sUnless (local_tid .==. 0) reset_to_neutral _ -> return () return (slugs_op_renamed, doTheReduction) -reductionStageOne :: KernelConstants - -> [(VName, Imp.Exp)] - -> Imp.Count Imp.Elements Imp.Exp - -> Imp.Exp - -> Imp.Count Imp.Elements Imp.Exp - -> VName - -> [SegBinOpSlug] - -> DoSegBody - -> InKernelGen [Lambda KernelsMem] +reductionStageOne :: + KernelConstants -> + [(VName, Imp.TExp Int64)] -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + Imp.TExp Int32 -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + VName -> + [SegBinOpSlug] -> + DoSegBody -> + InKernelGen [Lambda KernelsMem] reductionStageOne constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do (slugs_op_renamed, doTheReduction) <- reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body @@ -601,105 +726,146 @@ case slugsComm slugs of Noncommutative -> forM_ slugs $ \slug -> - forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> - copyDWIMFix (paramName p) [] (Var acc) acc_is + forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> + copyDWIMFix (paramName p) [] (Var acc) acc_is _ -> doTheReduction return slugs_op_renamed -reductionStageTwo :: KernelConstants - -> [PatElem KernelsMem] - -> Imp.Exp - -> Imp.Exp - -> [Imp.Exp] - -> Imp.Exp - -> Imp.Exp - -> SegBinOpSlug - -> [LParam KernelsMem] -> [LParam KernelsMem] - -> Lambda KernelsMem -> [SubExp] - -> Imp.Exp -> VName -> Imp.Exp -> VName -> [VName] -> [VName] - -> InKernelGen () -reductionStageTwo constants segred_pes - group_id flat_segment_id segment_gtids first_group_for_segment groups_per_segment - slug red_x_params red_y_params - red_op_renamed nes - num_counters counter counter_i sync_arr group_res_arrs red_arrs = do - let local_tid = kernelLocalThreadId constants - group_size = kernelGroupSize constants - old_counter <- dPrim "old_counter" int32 - (counter_mem, _, counter_offset) <- fullyIndexArray counter [counter_i * num_counters + - flat_segment_id `rem` num_counters] - comment "first thread in group saves group result to global memory" $ - sWhen (local_tid .==. 0) $ do - forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) -> - copyDWIMFix v [0, group_id] (Var acc) acc_is - sOp $ Imp.MemFence Imp.FenceGlobal - -- Increment the counter, thus stating that our result is - -- available. - sOp $ Imp.Atomic DefaultSpace $ Imp.AtomicAdd Int32 old_counter counter_mem counter_offset 1 - -- Now check if we were the last group to write our result. If - -- so, it is our responsibility to produce the final result. - sWrite sync_arr [0] $ Imp.var old_counter int32 .==. groups_per_segment - 1 - - sOp $ Imp.Barrier Imp.FenceGlobal +reductionStageTwo :: + KernelConstants -> + [PatElem KernelsMem] -> + Imp.TExp Int32 -> + Imp.TExp Int32 -> + [Imp.TExp Int64] -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + SegBinOpSlug -> + [LParam KernelsMem] -> + [LParam KernelsMem] -> + Lambda KernelsMem -> + [SubExp] -> + Imp.TExp Int32 -> + VName -> + Imp.TExp Int32 -> + VName -> + [VName] -> + [VName] -> + InKernelGen () +reductionStageTwo + constants + segred_pes + group_id + flat_segment_id + segment_gtids + first_group_for_segment + groups_per_segment + slug + red_x_params + red_y_params + red_op_renamed + nes + num_counters + counter + counter_i + sync_arr + group_res_arrs + red_arrs = do + let local_tid = kernelLocalThreadId constants + group_size = kernelGroupSize constants + old_counter <- dPrim "old_counter" int32 + (counter_mem, _, counter_offset) <- + fullyIndexArray + counter + [ sExt64 $ + counter_i * num_counters + + flat_segment_id `rem` num_counters + ] + comment "first thread in group saves group result to global memory" $ + sWhen (local_tid .==. 0) $ do + forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) -> + copyDWIMFix v [0, sExt64 group_id] (Var acc) acc_is + sOp $ Imp.MemFence Imp.FenceGlobal + -- Increment the counter, thus stating that our result is + -- available. + sOp $ + Imp.Atomic DefaultSpace $ + Imp.AtomicAdd + Int32 + (tvVar old_counter) + counter_mem + counter_offset + $ untyped (1 :: Imp.TExp Int32) + -- Now check if we were the last group to write our result. If + -- so, it is our responsibility to produce the final result. + sWrite sync_arr [0] $ untyped $ tvExp old_counter .==. groups_per_segment - 1 - is_last_group <- dPrim "is_last_group" Bool - copyDWIMFix is_last_group [] (Var sync_arr) [0] - sWhen (Imp.var is_last_group Bool) $ do - -- The final group has written its result (and it was - -- us!), so read in all the group results and perform the - -- final stage of the reduction. But first, we reset the - -- counter so it is ready for next time. This is done - -- with an atomic to avoid warnings about write/write - -- races in oclgrind. - sWhen (local_tid .==. 0) $ - sOp $ Imp.Atomic DefaultSpace $ - Imp.AtomicAdd Int32 old_counter counter_mem counter_offset $ - negate groups_per_segment + sOp $ Imp.Barrier Imp.FenceGlobal - sLoopNest (slugShape slug) $ \vec_is -> do - -- There is no guarantee that the number of workgroups for the - -- segment is less than the workgroup size, so each thread may - -- have to read multiple elements. We do this in a sequential - -- way that may induce non-coalesced accesses, but the total - -- number of accesses should be tiny here. - comment "read in the per-group-results" $ do - read_per_thread <- dPrimVE "read_per_thread" $ - groups_per_segment `divUp` group_size + is_last_group <- dPrim "is_last_group" Bool + copyDWIMFix (tvVar is_last_group) [] (Var sync_arr) [0] + sWhen (tvExp is_last_group) $ do + -- The final group has written its result (and it was + -- us!), so read in all the group results and perform the + -- final stage of the reduction. But first, we reset the + -- counter so it is ready for next time. This is done + -- with an atomic to avoid warnings about write/write + -- races in oclgrind. + sWhen (local_tid .==. 0) $ + sOp $ + Imp.Atomic DefaultSpace $ + Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $ + untyped $ negate groups_per_segment - forM_ (zip red_x_params nes) $ \(p, ne) -> - copyDWIMFix (paramName p) [] ne [] + sLoopNest (slugShape slug) $ \vec_is -> do + -- There is no guarantee that the number of workgroups for the + -- segment is less than the workgroup size, so each thread may + -- have to read multiple elements. We do this in a sequential + -- way that may induce non-coalesced accesses, but the total + -- number of accesses should be tiny here. + comment "read in the per-group-results" $ do + read_per_thread <- + dPrimVE "read_per_thread" $ + groups_per_segment `divUp` sExt64 group_size - sFor "i" read_per_thread $ \i -> do + forM_ (zip red_x_params nes) $ \(p, ne) -> + copyDWIMFix (paramName p) [] ne [] - group_res_id <- dPrimVE "group_res_id" $ - local_tid * read_per_thread + i - index_of_group_res <- dPrimVE "index_of_group_res" $ - first_group_for_segment + group_res_id + sFor "i" read_per_thread $ \i -> do + group_res_id <- + dPrimVE "group_res_id" $ + sExt64 local_tid * read_per_thread + i + index_of_group_res <- + dPrimVE "index_of_group_res" $ + first_group_for_segment + group_res_id - sWhen (group_res_id .<. groups_per_segment) $ do - forM_ (zip red_y_params group_res_arrs) $ - \(p, group_res_arr) -> - copyDWIMFix (paramName p) [] - (Var group_res_arr) - ([0, index_of_group_res] ++ vec_is) + sWhen (group_res_id .<. groups_per_segment) $ do + forM_ (zip red_y_params group_res_arrs) $ + \(p, group_res_arr) -> + copyDWIMFix + (paramName p) + [] + (Var group_res_arr) + ([0, index_of_group_res] ++ vec_is) - compileStms mempty (bodyStms $ slugBody slug) $ - forM_ (zip red_x_params (bodyResult $ slugBody slug)) $ \(p, se) -> - copyDWIMFix (paramName p) [] se [] + compileStms mempty (bodyStms $ slugBody slug) $ + forM_ (zip red_x_params (bodyResult $ slugBody slug)) $ \(p, se) -> + copyDWIMFix (paramName p) [] se [] - forM_ (zip red_x_params red_arrs) $ \(p, arr) -> - when (primType $ paramType p) $ - copyDWIMFix arr [local_tid] (Var $ paramName p) [] + forM_ (zip red_x_params red_arrs) $ \(p, arr) -> + when (primType $ paramType p) $ + copyDWIMFix arr [sExt64 local_tid] (Var $ paramName p) [] - sOp $ Imp.Barrier Imp.FenceLocal + sOp $ Imp.Barrier Imp.FenceLocal - sComment "reduce the per-group results" $ do - groupReduce group_size red_op_renamed red_arrs + sComment "reduce the per-group results" $ do + groupReduce (sExt32 group_size) red_op_renamed red_arrs - sComment "and back to memory with the final result" $ - sWhen (local_tid .==. 0) $ - forM_ (zip segred_pes $ lambdaParams red_op_renamed) $ \(pe, p) -> - copyDWIMFix - (patElemName pe) (segment_gtids++vec_is) - (Var $ paramName p) [] + sComment "and back to memory with the final result" $ + sWhen (local_tid .==. 0) $ + forM_ (zip segred_pes $ lambdaParams red_op_renamed) $ \(pe, p) -> + copyDWIMFix + (patElemName pe) + (segment_gtids ++ vec_is) + (Var $ paramName p) + [] diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs @@ -1,86 +1,89 @@ -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + -- | Code generation for segmented and non-segmented scans. Uses a -- fairly inefficient two-pass algorithm. -module Futhark.CodeGen.ImpGen.Kernels.SegScan - ( compileSegScan ) - where +module Futhark.CodeGen.ImpGen.Kernels.SegScan (compileSegScan) where import Control.Monad.Except import Control.Monad.State -import Data.Maybe import Data.List (delete, find, foldl', zip4) - -import Prelude hiding (quot, rem) - -import Futhark.Transform.Rename -import Futhark.IR.KernelsMem +import Data.Maybe import qualified Futhark.CodeGen.ImpCode.Kernels as Imp import Futhark.CodeGen.ImpGen import Futhark.CodeGen.ImpGen.Kernels.Base +import Futhark.IR.KernelsMem import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Util.IntegralExp (divUp, quot, rem) +import Futhark.Transform.Rename import Futhark.Util (takeLast) +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) -- Aggressively try to reuse memory for different SegBinOps, because -- we will run them sequentially after another. -makeLocalArrays :: Count GroupSize SubExp -> SubExp -> [SegBinOp KernelsMem] - -> InKernelGen [[VName]] +makeLocalArrays :: + Count GroupSize SubExp -> + SubExp -> + [SegBinOp KernelsMem] -> + InKernelGen [[VName]] makeLocalArrays (Count group_size) num_threads scans = do (arrs, mems_and_sizes) <- runStateT (mapM onScan scans) mempty - let maxSize sizes = - Imp.bytes $ foldl' (Imp.BinOpExp (SMax Int32)) 1 $ - map Imp.unCount sizes + let maxSize sizes = Imp.bytes $ foldl' sMax64 1 $ map Imp.unCount sizes forM_ mems_and_sizes $ \(sizes, mem) -> sAlloc_ mem (maxSize sizes) (Space "local") return arrs - - where onScan (SegBinOp _ scan_op nes _) = do - let (scan_x_params, _scan_y_params) = - splitAt (length nes) $ lambdaParams scan_op - (arrs, used_mems) <- fmap unzip $ forM scan_x_params $ \p -> - case paramDec p of - MemArray pt shape _ (ArrayIn mem _) -> do - let shape' = Shape [num_threads] <> shape - arr <- lift $ sArray "scan_arr" pt shape' $ - ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape' - return (arr, []) - _ -> do - let pt = elemType $ paramType p - shape = Shape [group_size] - (sizes, mem') <- getMem pt shape - arr <- lift $ sArrayInMem "scan_arr" pt shape mem' - return (arr, [(sizes, mem')]) - modify (<>concat used_mems) - return arrs + where + onScan (SegBinOp _ scan_op nes _) = do + let (scan_x_params, _scan_y_params) = + splitAt (length nes) $ lambdaParams scan_op + (arrs, used_mems) <- fmap unzip $ + forM scan_x_params $ \p -> + case paramDec p of + MemArray pt shape _ (ArrayIn mem _) -> do + let shape' = Shape [num_threads] <> shape + arr <- + lift $ + sArray "scan_arr" pt shape' $ + ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape' + return (arr, []) + _ -> do + let pt = elemType $ paramType p + shape = Shape [group_size] + (sizes, mem') <- getMem pt shape + arr <- lift $ sArrayInMem "scan_arr" pt shape mem' + return (arr, [(sizes, mem')]) + modify (<> concat used_mems) + return arrs - getMem pt shape = do - let size = typeSize $ Array pt shape NoUniqueness - mems <- get - case (find ((size `elem`) . fst) mems, mems) of - (Just mem, _) -> do - modify $ delete mem - return mem - (Nothing, (size', mem) : mems') -> do - put mems' - return (size : size', mem) - (Nothing, []) -> do - mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "local" - return ([size], mem) + getMem pt shape = do + let size = typeSize $ Array pt shape NoUniqueness + mems <- get + case (find ((size `elem`) . fst) mems, mems) of + (Just mem, _) -> do + modify $ delete mem + return mem + (Nothing, (size', mem) : mems') -> do + put mems' + return (size : size', mem) + (Nothing, []) -> do + mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "local" + return ([size], mem) -type CrossesSegment = Maybe (Imp.Exp -> Imp.Exp -> Imp.Exp) +type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool) -localArrayIndex :: KernelConstants -> Type -> Imp.Exp +localArrayIndex :: KernelConstants -> Type -> Imp.TExp Int64 localArrayIndex constants t = if primType t - then kernelLocalThreadId constants - else kernelGlobalThreadId constants + then sExt64 (kernelLocalThreadId constants) + else sExt64 (kernelGlobalThreadId constants) barrierFor :: Lambda KernelsMem -> (Bool, Imp.Fence, InKernelGen ()) barrierFor scan_op = (array_scan, fence, sOp $ Imp.Barrier fence) - where array_scan = not $ all primType $ lambdaReturnType scan_op - fence | array_scan = Imp.FenceGlobal - | otherwise = Imp.FenceLocal + where + array_scan = not $ all primType $ lambdaReturnType scan_op + fence + | array_scan = Imp.FenceGlobal + | otherwise = Imp.FenceLocal xParams, yParams :: SegBinOp KernelsMem -> [LParam KernelsMem] xParams scan = @@ -88,72 +91,88 @@ yParams scan = drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) -writeToScanValues :: [VName] - -> ([PatElem KernelsMem], SegBinOp KernelsMem, [KernelResult]) - -> InKernelGen () +writeToScanValues :: + [VName] -> + ([PatElem KernelsMem], SegBinOp KernelsMem, [KernelResult]) -> + InKernelGen () writeToScanValues gtids (pes, scan, scan_res) | shapeRank (segBinOpShape scan) > 0 = - forM_ (zip pes scan_res) $ \(pe, res) -> - copyDWIMFix (patElemName pe) (map Imp.vi32 gtids) - (kernelResultSubExp res) [] + forM_ (zip pes scan_res) $ \(pe, res) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 gtids) + (kernelResultSubExp res) + [] | otherwise = - forM_ (zip (yParams scan) scan_res) $ \(p, res) -> + forM_ (zip (yParams scan) scan_res) $ \(p, res) -> copyDWIMFix (paramName p) [] (kernelResultSubExp res) [] -readToScanValues :: [Imp.Exp] -> [PatElem KernelsMem] -> SegBinOp KernelsMem - -> InKernelGen () +readToScanValues :: + [Imp.TExp Int64] -> + [PatElem KernelsMem] -> + SegBinOp KernelsMem -> + InKernelGen () readToScanValues is pes scan | shapeRank (segBinOpShape scan) > 0 = - forM_ (zip (yParams scan) pes) $ \(p, pe) -> + forM_ (zip (yParams scan) pes) $ \(p, pe) -> copyDWIMFix (paramName p) [] (Var (patElemName pe)) is | otherwise = - return () + return () -readCarries :: Imp.Exp -> [Imp.Exp] -> [Imp.Exp] - -> [PatElem KernelsMem] - -> SegBinOp KernelsMem - -> InKernelGen () +readCarries :: + Imp.TExp Int64 -> + [Imp.TExp Int64] -> + [Imp.TExp Int64] -> + [PatElem KernelsMem] -> + SegBinOp KernelsMem -> + InKernelGen () readCarries chunk_offset dims' vec_is pes scan | shapeRank (segBinOpShape scan) > 0 = do - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - -- We may have to reload the carries from the output of the - -- previous chunk. - sIf (chunk_offset .>. 0 .&&. ltid .==. 0) - (do let is = unflattenIndex dims' $ chunk_offset - 1 - forM_ (zip (xParams scan) pes) $ \(p, pe) -> - copyDWIMFix (paramName p) [] (Var (patElemName pe)) (is++vec_is)) - (forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> - copyDWIMFix (paramName p) [] ne []) + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + -- We may have to reload the carries from the output of the + -- previous chunk. + sIf + (chunk_offset .>. 0 .&&. ltid .==. 0) + ( do + let is = unflattenIndex dims' $ chunk_offset - 1 + forM_ (zip (xParams scan) pes) $ \(p, pe) -> + copyDWIMFix (paramName p) [] (Var (patElemName pe)) (is ++ vec_is) + ) + ( forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> + copyDWIMFix (paramName p) [] ne [] + ) | otherwise = - return () + return () -- | Produce partially scanned intervals; one per workgroup. -scanStage1 :: Pattern KernelsMem - -> Count NumGroups SubExp -> Count GroupSize SubExp -> SegSpace - -> [SegBinOp KernelsMem] - -> KernelBody KernelsMem - -> CallKernelGen (VName, Imp.Exp, CrossesSegment) +scanStage1 :: + Pattern KernelsMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegBinOp KernelsMem] -> + KernelBody KernelsMem -> + CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment) scanStage1 (Pattern _ all_pes) num_groups group_size space scans kbody = do - num_groups' <- traverse toExp num_groups - group_size' <- traverse toExp group_size - num_threads <- dPrimV "num_threads" $ - unCount num_groups' * unCount group_size' + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_groups' * unCount group_size' let (gtids, dims) = unzip $ unSegSpace space - dims' <- mapM toExp dims + dims' = map toInt64Exp dims let num_elements = product dims' - elems_per_thread = num_elements `divUp` Imp.vi32 num_threads + elems_per_thread = num_elements `divUp` sExt64 (tvExp num_threads) elems_per_group = unCount group_size' * elems_per_thread let crossesSegment = case reverse dims' of segment_size : _ : _ -> Just $ \from to -> - (to-from) .>. (to `rem` segment_size) + (to - from) .>. (to `rem` segment_size) _ -> Nothing sKernelThread "scan_stage1" num_groups' group_size' (segFlat space) $ do constants <- kernelConstants <$> askEnv - all_local_arrs <- makeLocalArrays group_size (Var num_threads) scans + all_local_arrs <- makeLocalArrays group_size (tvSize num_threads) scans -- The variables from scan_op will be used for the carry and such -- in the big chunking loop. @@ -163,18 +182,20 @@ copyDWIMFix (paramName p) [] ne [] sFor "j" elems_per_thread $ \j -> do - chunk_offset <- dPrimV "chunk_offset" $ - kernelGroupSize constants * j + - kernelGroupId constants * elems_per_group - flat_idx <- dPrimV "flat_idx" $ - Imp.var chunk_offset int32 + kernelLocalThreadId constants + chunk_offset <- + dPrimV "chunk_offset" $ + sExt64 (kernelGroupSize constants) * j + + sExt64 (kernelGroupId constants) * elems_per_group + flat_idx <- + dPrimV "flat_idx" $ + tvExp chunk_offset + sExt64 (kernelLocalThreadId constants) -- Construct segment indices. - zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ Imp.vi32 flat_idx + zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx let per_scan_pes = segBinOpChunks scans all_pes in_bounds = - foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi32 gtids) dims' + foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do let (all_scan_res, map_res) = @@ -184,241 +205,306 @@ sComment "write to-scan values to parameters" $ mapM_ (writeToScanValues gtids) $ - zip3 per_scan_pes scans per_scan_res + zip3 per_scan_pes scans per_scan_res sComment "write mapped values results to global memory" $ forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(pe, se) -> - copyDWIMFix (patElemName pe) (map Imp.vi32 gtids) - (kernelResultSubExp se) [] + copyDWIMFix + (patElemName pe) + (map Imp.vi64 gtids) + (kernelResultSubExp se) + [] sComment "threads in bounds read input" $ sWhen in_bounds when_in_bounds forM_ (zip3 per_scan_pes scans all_local_arrs) $ \(pes, scan@(SegBinOp _ scan_op nes vec_shape), local_arrs) -> - sComment "do one intra-group scan operation" $ do - let rets = lambdaReturnType scan_op - scan_x_params = xParams scan - (array_scan, fence, barrier) = barrierFor scan_op + sComment "do one intra-group scan operation" $ do + let rets = lambdaReturnType scan_op + scan_x_params = xParams scan + (array_scan, fence, barrier) = barrierFor scan_op - when array_scan barrier + when array_scan barrier - sLoopNest vec_shape $ \vec_is -> do - sComment "maybe restore some to-scan values to parameters, or read neutral" $ - sIf in_bounds - (do readToScanValues (map Imp.vi32 gtids++vec_is) pes scan - readCarries (Imp.vi32 chunk_offset) dims' vec_is pes scan) - (forM_ (zip (yParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> - copyDWIMFix (paramName p) [] ne []) + sLoopNest vec_shape $ \vec_is -> do + sComment "maybe restore some to-scan values to parameters, or read neutral" $ + sIf + in_bounds + ( do + readToScanValues (map Imp.vi64 gtids ++ vec_is) pes scan + readCarries (tvExp chunk_offset) dims' vec_is pes scan + ) + ( forM_ (zip (yParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> + copyDWIMFix (paramName p) [] ne [] + ) - sComment "combine with carry and write to local memory" $ - compileStms mempty (bodyStms $ lambdaBody scan_op) $ - forM_ (zip3 rets local_arrs (bodyResult $ lambdaBody scan_op)) $ - \(t, arr, se) -> copyDWIMFix arr [localArrayIndex constants t] se [] + sComment "combine with carry and write to local memory" $ + compileStms mempty (bodyStms $ lambdaBody scan_op) $ + forM_ (zip3 rets local_arrs (bodyResult $ lambdaBody scan_op)) $ + \(t, arr, se) -> + copyDWIMFix arr [localArrayIndex constants t] se [] - let crossesSegment' = do - f <- crossesSegment - Just $ \from to -> - let from' = from + Imp.var chunk_offset int32 - to' = to + Imp.var chunk_offset int32 - in f from' to' + let crossesSegment' = do + f <- crossesSegment + Just $ \from to -> + let from' = sExt64 from + tvExp chunk_offset + to' = sExt64 to + tvExp chunk_offset + in f from' to' - sOp $ Imp.ErrorSync fence + sOp $ Imp.ErrorSync fence - -- We need to avoid parameter name clashes. - scan_op_renamed <- renameLambda scan_op - groupScan crossesSegment' - (Imp.vi32 num_threads) - (kernelGroupSize constants) scan_op_renamed local_arrs + -- We need to avoid parameter name clashes. + scan_op_renamed <- renameLambda scan_op + groupScan + crossesSegment' + (sExt64 $ tvExp num_threads) + (sExt64 $ kernelGroupSize constants) + scan_op_renamed + local_arrs - sComment "threads in bounds write partial scan result" $ - sWhen in_bounds $ forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> - copyDWIMFix (patElemName pe) (map Imp.vi32 gtids++vec_is) - (Var arr) [localArrayIndex constants t] + sComment "threads in bounds write partial scan result" $ + sWhen in_bounds $ + forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 gtids ++ vec_is) + (Var arr) + [localArrayIndex constants t] - barrier + barrier - let load_carry = - forM_ (zip local_arrs scan_x_params) $ \(arr, p) -> - copyDWIMFix (paramName p) [] (Var arr) - [if primType $ paramType p - then kernelGroupSize constants - 1 - else (kernelGroupId constants+1) * kernelGroupSize constants - 1] - load_neutral = - forM_ (zip nes scan_x_params) $ \(ne, p) -> - copyDWIMFix (paramName p) [] ne [] + let load_carry = + forM_ (zip local_arrs scan_x_params) $ \(arr, p) -> + copyDWIMFix + (paramName p) + [] + (Var arr) + [ if primType $ paramType p + then sExt64 (kernelGroupSize constants) - 1 + else + (sExt64 (kernelGroupId constants) + 1) + * sExt64 (kernelGroupSize constants) - 1 + ] + load_neutral = + forM_ (zip nes scan_x_params) $ \(ne, p) -> + copyDWIMFix (paramName p) [] ne [] - sComment "first thread reads last element as carry-in for next iteration" $ do - crosses_segment <- dPrimVE "crosses_segment" $ - case crossesSegment of - Nothing -> false - Just f -> f (Imp.var chunk_offset int32 + - kernelGroupSize constants-1) - (Imp.var chunk_offset int32 + - kernelGroupSize constants) - should_load_carry <- dPrimVE "should_load_carry" $ - kernelLocalThreadId constants .==. 0 .&&. UnOpExp Not crosses_segment - sWhen should_load_carry load_carry - when array_scan barrier - sUnless should_load_carry load_neutral + sComment "first thread reads last element as carry-in for next iteration" $ do + crosses_segment <- dPrimVE "crosses_segment" $ + case crossesSegment of + Nothing -> false + Just f -> + f + ( tvExp chunk_offset + + sExt64 (kernelGroupSize constants) -1 + ) + ( tvExp chunk_offset + + sExt64 (kernelGroupSize constants) + ) + should_load_carry <- + dPrimVE "should_load_carry" $ + kernelLocalThreadId constants .==. 0 .&&. bNot crosses_segment + sWhen should_load_carry load_carry + when array_scan barrier + sUnless should_load_carry load_neutral - barrier + barrier return (num_threads, elems_per_group, crossesSegment) -scanStage2 :: Pattern KernelsMem - -> VName -> Imp.Exp -> Count NumGroups SubExp -> CrossesSegment -> SegSpace - -> [SegBinOp KernelsMem] - -> CallKernelGen () +scanStage2 :: + Pattern KernelsMem -> + TV Int32 -> + Imp.TExp Int64 -> + Count NumGroups SubExp -> + CrossesSegment -> + SegSpace -> + [SegBinOp KernelsMem] -> + CallKernelGen () scanStage2 (Pattern _ all_pes) stage1_num_threads elems_per_group num_groups crossesSegment space scans = do let (gtids, dims) = unzip $ unSegSpace space - dims' <- mapM toExp dims + dims' = map toInt64Exp dims -- Our group size is the number of groups for the stage 1 kernel. let group_size = Count $ unCount num_groups - group_size' <- traverse toExp group_size + group_size' = fmap toInt64Exp group_size let crossesSegment' = do f <- crossesSegment Just $ \from to -> - f ((from + 1) * elems_per_group - 1) ((to + 1) * elems_per_group - 1) + f + ((sExt64 from + 1) * elems_per_group - 1) + ((sExt64 to + 1) * elems_per_group - 1) - sKernelThread "scan_stage2" 1 group_size' (segFlat space) $ do + sKernelThread "scan_stage2" 1 group_size' (segFlat space) $ do constants <- kernelConstants <$> askEnv - per_scan_local_arrs <- makeLocalArrays group_size (Var stage1_num_threads) scans + per_scan_local_arrs <- makeLocalArrays group_size (tvSize stage1_num_threads) scans let per_scan_rets = map (lambdaReturnType . segBinOpLambda) scans per_scan_pes = segBinOpChunks scans all_pes - flat_idx <- dPrimV "flat_idx" $ - (kernelLocalThreadId constants + 1) * elems_per_group - 1 + flat_idx <- + dPrimV "flat_idx" $ + (sExt64 (kernelLocalThreadId constants) + 1) * elems_per_group - 1 -- Construct segment indices. - zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ Imp.var flat_idx int32 + zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx forM_ (zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $ \(SegBinOp _ scan_op nes vec_shape, local_arrs, rets, pes) -> sLoopNest vec_shape $ \vec_is -> do - let glob_is = map Imp.vi32 gtids ++ vec_is + let glob_is = map Imp.vi64 gtids ++ vec_is - in_bounds = - foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi32 gtids) dims' + in_bounds = + foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' - when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) -> - copyDWIMFix arr [localArrayIndex constants t] - (Var $ patElemName pe) glob_is + when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) -> + copyDWIMFix + arr + [localArrayIndex constants t] + (Var $ patElemName pe) + glob_is - when_out_of_bounds = forM_ (zip3 rets local_arrs nes) $ \(t, arr, ne) -> - copyDWIMFix arr [localArrayIndex constants t] ne [] - (_, _, barrier) = - barrierFor scan_op + when_out_of_bounds = forM_ (zip3 rets local_arrs nes) $ \(t, arr, ne) -> + copyDWIMFix arr [localArrayIndex constants t] ne [] + (_, _, barrier) = + barrierFor scan_op - sComment "threads in bound read carries; others get neutral element" $ - sIf in_bounds when_in_bounds when_out_of_bounds + sComment "threads in bound read carries; others get neutral element" $ + sIf in_bounds when_in_bounds when_out_of_bounds - barrier + barrier - groupScan crossesSegment' - (Imp.vi32 stage1_num_threads) (kernelGroupSize constants) scan_op local_arrs + groupScan + crossesSegment' + (sExt64 $ tvExp stage1_num_threads) + (sExt64 $ kernelGroupSize constants) + scan_op + local_arrs - sComment "threads in bounds write scanned carries" $ - sWhen in_bounds $ forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> - copyDWIMFix (patElemName pe) glob_is - (Var arr) [localArrayIndex constants t] + sComment "threads in bounds write scanned carries" $ + sWhen in_bounds $ + forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> + copyDWIMFix + (patElemName pe) + glob_is + (Var arr) + [localArrayIndex constants t] -scanStage3 :: Pattern KernelsMem - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> Imp.Exp -> CrossesSegment -> SegSpace - -> [SegBinOp KernelsMem] - -> CallKernelGen () +scanStage3 :: + Pattern KernelsMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + Imp.TExp Int64 -> + CrossesSegment -> + SegSpace -> + [SegBinOp KernelsMem] -> + CallKernelGen () scanStage3 (Pattern _ all_pes) num_groups group_size elems_per_group crossesSegment space scans = do - num_groups' <- traverse toExp num_groups - group_size' <- traverse toExp group_size - let (gtids, dims) = unzip $ unSegSpace space - dims' <- mapM toExp dims - required_groups <- dPrimVE "required_groups" $ - product dims' `divUp` unCount group_size' + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + (gtids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + required_groups <- + dPrimVE "required_groups" $ + sExt32 $ product dims' `divUp` sExt64 (unCount group_size') sKernelThread "scan_stage3" num_groups' group_size' (segFlat space) $ virtualiseGroups SegVirt required_groups $ \virt_group_id -> do - constants <- kernelConstants <$> askEnv + constants <- kernelConstants <$> askEnv - -- Compute our logical index. - flat_idx <- dPrimVE "flat_idx" $ - Imp.vi32 virt_group_id * unCount group_size' + - kernelLocalThreadId constants - zipWithM_ dPrimV_ gtids $ unflattenIndex dims' flat_idx + -- Compute our logical index. + flat_idx <- + dPrimVE "flat_idx" $ + sExt64 virt_group_id * sExt64 (unCount group_size') + + sExt64 (kernelLocalThreadId constants) + zipWithM_ dPrimV_ gtids $ unflattenIndex dims' flat_idx - -- Figure out which group this element was originally in. - orig_group <- dPrimV "orig_group" $ flat_idx `quot` elems_per_group - -- Then the index of the carry-in of the preceding group. - carry_in_flat_idx <- dPrimV "carry_in_flat_idx" $ - Imp.var orig_group int32 * elems_per_group - 1 - -- Figure out the logical index of the carry-in. - let carry_in_idx = unflattenIndex dims' $ Imp.var carry_in_flat_idx int32 + -- Figure out which group this element was originally in. + orig_group <- dPrimV "orig_group" $ flat_idx `quot` elems_per_group + -- Then the index of the carry-in of the preceding group. + carry_in_flat_idx <- + dPrimV "carry_in_flat_idx" $ + tvExp orig_group * elems_per_group - 1 + -- Figure out the logical index of the carry-in. + let carry_in_idx = unflattenIndex dims' $ tvExp carry_in_flat_idx - -- Apply the carry if we are not in the scan results for the first - -- group, and are not the last element in such a group (because - -- then the carry was updated in stage 2), and we are not crossing - -- a segment boundary. - let in_bounds = - foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi32 gtids) dims' - crosses_segment = fromMaybe false $ - crossesSegment <*> - pure (Imp.var carry_in_flat_idx int32) <*> - pure flat_idx - is_a_carry = flat_idx .==. - (Imp.var orig_group int32 + 1) * elems_per_group - 1 - no_carry_in = Imp.var orig_group int32 .==. 0 .||. is_a_carry .||. crosses_segment + -- Apply the carry if we are not in the scan results for the first + -- group, and are not the last element in such a group (because + -- then the carry was updated in stage 2), and we are not crossing + -- a segment boundary. + let in_bounds = + foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' + crosses_segment = + fromMaybe false $ + crossesSegment + <*> pure (tvExp carry_in_flat_idx) + <*> pure flat_idx + is_a_carry = flat_idx .==. (tvExp orig_group + 1) * elems_per_group - 1 + no_carry_in = tvExp orig_group .==. 0 .||. is_a_carry .||. crosses_segment - let per_scan_pes = segBinOpChunks scans all_pes - sWhen in_bounds $ sUnless no_carry_in $ - forM_ (zip per_scan_pes scans) $ - \(pes, SegBinOp _ scan_op nes vec_shape) -> do - dScope Nothing $ scopeOfLParams $ lambdaParams scan_op - let (scan_x_params, scan_y_params) = - splitAt (length nes) $ lambdaParams scan_op + let per_scan_pes = segBinOpChunks scans all_pes + sWhen in_bounds $ + sUnless no_carry_in $ + forM_ (zip per_scan_pes scans) $ + \(pes, SegBinOp _ scan_op nes vec_shape) -> do + dScope Nothing $ scopeOfLParams $ lambdaParams scan_op + let (scan_x_params, scan_y_params) = + splitAt (length nes) $ lambdaParams scan_op - sLoopNest vec_shape $ \vec_is -> do - forM_ (zip scan_x_params pes) $ \(p, pe) -> - copyDWIMFix (paramName p) [] - (Var $ patElemName pe) (carry_in_idx++vec_is) + sLoopNest vec_shape $ \vec_is -> do + forM_ (zip scan_x_params pes) $ \(p, pe) -> + copyDWIMFix + (paramName p) + [] + (Var $ patElemName pe) + (carry_in_idx ++ vec_is) - forM_ (zip scan_y_params pes) $ \(p, pe) -> - copyDWIMFix (paramName p) [] - (Var $ patElemName pe) (map Imp.vi32 gtids++vec_is) + forM_ (zip scan_y_params pes) $ \(p, pe) -> + copyDWIMFix + (paramName p) + [] + (Var $ patElemName pe) + (map Imp.vi64 gtids ++ vec_is) - compileBody' scan_x_params $ lambdaBody scan_op + compileBody' scan_x_params $ lambdaBody scan_op - forM_ (zip scan_x_params pes) $ \(p, pe) -> - copyDWIMFix (patElemName pe) (map Imp.vi32 gtids++vec_is) - (Var $ paramName p) [] + forM_ (zip scan_x_params pes) $ \(p, pe) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 gtids ++ vec_is) + (Var $ paramName p) + [] -- | Compile 'SegScan' instance to host-level code with calls to -- various kernels. -compileSegScan :: Pattern KernelsMem - -> SegLevel -> SegSpace - -> [SegBinOp KernelsMem] - -> KernelBody KernelsMem - -> CallKernelGen () +compileSegScan :: + Pattern KernelsMem -> + SegLevel -> + SegSpace -> + [SegBinOp KernelsMem] -> + KernelBody KernelsMem -> + CallKernelGen () compileSegScan pat lvl space scans kbody = sWhen (0 .<. n) $ do emit $ Imp.DebugPrint "\n# SegScan" Nothing -- Since stage 2 involves a group size equal to the number of groups -- used for stage 1, we have to cap this number to the maximum group -- size. - stage1_max_num_groups <- - dPrim "stage1_max_num_groups" int32 - sOp $ Imp.GetSizeMax stage1_max_num_groups SizeGroup + stage1_max_num_groups <- dPrim "stage1_max_num_groups" int32 + sOp $ Imp.GetSizeMax (tvVar stage1_max_num_groups) SizeGroup stage1_num_groups <- - fmap (Imp.Count . Var) $ dPrimV "stage1_num_groups" $ - Imp.BinOpExp (SMin Int32) (Imp.vi32 stage1_max_num_groups) $ - toExp' int32 $ Imp.unCount $ segNumGroups lvl + fmap (Imp.Count . tvSize) $ + dPrimV "stage1_num_groups" $ + sMin32 (tvExp stage1_max_num_groups) $ + toInt32Exp $ Imp.unCount $ segNumGroups lvl (stage1_num_threads, elems_per_group, crossesSegment) <- scanStage1 pat stage1_num_groups (segGroupSize lvl) space scans kbody - emit $ Imp.DebugPrint "elems_per_group" $ Just elems_per_group + emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group scanStage2 pat stage1_num_threads elems_per_group stage1_num_groups crossesSegment space scans scanStage3 pat (segNumGroups lvl) (segGroupSize lvl) elems_per_group crossesSegment space scans - where n = product $ map (toExp' int32) $ segSpaceDims space + where + n = product $ map toInt32Exp $ segSpaceDims space 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 @@ -1,26 +1,22 @@ {-# LANGUAGE QuasiQuotes #-} -{-# LANGUAGE TupleSections #-} {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TupleSections #-} + -- | This module defines a translation from imperative code with -- kernels to imperative code with OpenCL calls. module Futhark.CodeGen.ImpGen.Kernels.ToOpenCL - ( kernelsToOpenCL - , kernelsToCUDA + ( kernelsToOpenCL, + kernelsToCUDA, ) - where +where +import Control.Monad.Identity import Control.Monad.Reader import Control.Monad.State -import Control.Monad.Identity import Data.FileEmbed +import qualified Data.Map.Strict as M import Data.Maybe import qualified Data.Set as S -import qualified Data.Map.Strict as M - -import qualified Language.C.Syntax as C -import qualified Language.C.Quote.OpenCL as C -import qualified Language.C.Quote.CUDA as CUDAC - import qualified Futhark.CodeGen.Backends.GenericC as GC import Futhark.CodeGen.Backends.SimpleRep import Futhark.CodeGen.ImpCode.Kernels hiding (Program) @@ -32,91 +28,105 @@ import Futhark.MonadFreshNames import Futhark.Util (zEncodeString) import Futhark.Util.Pretty (prettyOneLine) +import qualified Language.C.Quote.CUDA as CUDAC +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C kernelsToCUDA, kernelsToOpenCL :: ImpKernels.Program -> ImpOpenCL.Program kernelsToCUDA = translateKernels TargetCUDA kernelsToOpenCL = translateKernels TargetOpenCL -- | Translate a kernels-program to an OpenCL-program. -translateKernels :: KernelTarget - -> ImpKernels.Program - -> ImpOpenCL.Program +translateKernels :: + KernelTarget -> + ImpKernels.Program -> + ImpOpenCL.Program translateKernels target prog = - let (prog', - ToOpenCL kernels device_funs used_types sizes failures) = - - (`runState` initialOpenCL) . (`runReaderT` defFuns prog) $ do - let ImpKernels.Definitions - (ImpKernels.Constants ps consts) - (ImpKernels.Functions funs) = prog - consts' <- traverse (onHostOp target) consts - funs' <- forM funs $ \(fname, fun) -> - (fname,) <$> traverse (onHostOp target) fun + let ( prog', + ToOpenCL kernels device_funs used_types sizes failures + ) = + (`runState` initialOpenCL) . (`runReaderT` defFuns prog) $ do + let ImpKernels.Definitions + (ImpKernels.Constants ps consts) + (ImpKernels.Functions funs) = prog + consts' <- traverse (onHostOp target) consts + funs' <- forM funs $ \(fname, fun) -> + (fname,) <$> traverse (onHostOp target) fun - return $ ImpOpenCL.Definitions - (ImpOpenCL.Constants ps consts') - (ImpOpenCL.Functions funs') + return $ + ImpOpenCL.Definitions + (ImpOpenCL.Constants ps consts') + (ImpOpenCL.Functions funs') (device_prototypes, device_defs) = unzip $ M.elems device_funs kernels' = M.map fst kernels opencl_code = openClCode $ map snd $ M.elems kernels opencl_prelude = - unlines [pretty $ genPrelude target used_types, - unlines $ map pretty device_prototypes, - unlines $ map pretty device_defs] - - in ImpOpenCL.Program opencl_code opencl_prelude kernels' - (S.toList used_types) (cleanSizes sizes) failures prog' - - where genPrelude TargetOpenCL = genOpenClPrelude - genPrelude TargetCUDA = const genCUDAPrelude + unlines + [ pretty $ genPrelude target used_types, + unlines $ map pretty device_prototypes, + unlines $ map pretty device_defs + ] + in ImpOpenCL.Program + opencl_code + opencl_prelude + kernels' + (S.toList used_types) + (cleanSizes sizes) + failures + prog' + where + genPrelude TargetOpenCL = genOpenClPrelude + genPrelude TargetCUDA = const genCUDAPrelude -- | Due to simplifications after kernel extraction, some threshold -- parameters may contain KernelPaths that reference threshold -- parameters that no longer exist. We remove these here. cleanSizes :: M.Map Name SizeClass -> M.Map Name SizeClass cleanSizes m = M.map clean m - where known = M.keys m - clean (SizeThreshold path def) = - SizeThreshold (filter ((`elem` known) . fst) path) def - clean s = s + where + known = M.keys m + clean (SizeThreshold path def) = + SizeThreshold (filter ((`elem` known) . fst) path) def + clean s = s -pointerQuals :: Monad m => String -> m [C.TypeQual] -pointerQuals "global" = return [C.ctyquals|__global|] -pointerQuals "local" = return [C.ctyquals|__local|] -pointerQuals "private" = return [C.ctyquals|__private|] -pointerQuals "constant" = return [C.ctyquals|__constant|] +pointerQuals :: Monad m => String -> m [C.TypeQual] +pointerQuals "global" = return [C.ctyquals|__global|] +pointerQuals "local" = return [C.ctyquals|__local|] +pointerQuals "private" = return [C.ctyquals|__private|] +pointerQuals "constant" = return [C.ctyquals|__constant|] pointerQuals "write_only" = return [C.ctyquals|__write_only|] -pointerQuals "read_only" = return [C.ctyquals|__read_only|] -pointerQuals "kernel" = return [C.ctyquals|__kernel|] -pointerQuals s = error $ "'" ++ s ++ "' is not an OpenCL kernel address space." +pointerQuals "read_only" = return [C.ctyquals|__read_only|] +pointerQuals "kernel" = return [C.ctyquals|__kernel|] +pointerQuals s = error $ "'" ++ s ++ "' is not an OpenCL kernel address space." -- In-kernel name and per-workgroup size in bytes. type LocalMemoryUse = (VName, Count Bytes Exp) -data KernelState = - KernelState { kernelLocalMemory :: [LocalMemoryUse] - , kernelFailures :: [FailureMsg] - , kernelNextSync :: Int - , kernelSyncPending :: Bool - -- ^ Has a potential failure occurred sine the last - -- ErrorSync? - , kernelHasBarriers :: Bool - } +data KernelState = KernelState + { kernelLocalMemory :: [LocalMemoryUse], + kernelFailures :: [FailureMsg], + kernelNextSync :: Int, + -- | Has a potential failure occurred sine the last + -- ErrorSync? + kernelSyncPending :: Bool, + kernelHasBarriers :: Bool + } newKernelState :: [FailureMsg] -> KernelState newKernelState failures = KernelState mempty failures 0 False False errorLabel :: KernelState -> String -errorLabel = ("error_"++) . show . kernelNextSync +errorLabel = ("error_" ++) . show . kernelNextSync -data ToOpenCL = ToOpenCL { clKernels :: M.Map KernelName (KernelSafety, C.Func) - , clDevFuns :: M.Map Name (C.Definition, C.Func) - , clUsedTypes :: S.Set PrimType - , clSizes :: M.Map Name SizeClass - , clFailures :: [FailureMsg] - } +data ToOpenCL = ToOpenCL + { clKernels :: M.Map KernelName (KernelSafety, C.Func), + clDevFuns :: M.Map Name (C.Definition, C.Func), + clUsedTypes :: S.Set PrimType, + clSizes :: M.Map Name SizeClass, + clFailures :: [FailureMsg] + } initialOpenCL :: ToOpenCL initialOpenCL = ToOpenCL mempty mempty mempty mempty mempty @@ -130,7 +140,7 @@ addSize :: Name -> SizeClass -> OnKernelM () addSize key sclass = - modify $ \s -> s { clSizes = M.insert key sclass $ clSizes s } + modify $ \s -> s {clSizes = M.insert key sclass $ clSizes s} onHostOp :: KernelTarget -> HostOp -> OnKernelM OpenCL onHostOp target (CallKernel k) = onKernel target k @@ -143,12 +153,17 @@ onHostOp _ (ImpKernels.GetSizeMax v size_class) = return $ ImpOpenCL.GetSizeMax v size_class -genGPUCode :: OpsMode -> KernelCode -> [FailureMsg] - -> GC.CompilerM KernelOp KernelState a - -> (a, GC.CompilerState KernelState) +genGPUCode :: + OpsMode -> + KernelCode -> + [FailureMsg] -> + GC.CompilerM KernelOp KernelState a -> + (a, GC.CompilerState KernelState) genGPUCode mode body failures = - GC.runCompilerM (inKernelOperations mode body) - blankNameSource (newKernelState failures) + GC.runCompilerM + (inKernelOperations mode body) + blankNameSource + (newKernelState failures) -- Compilation of a device function that is not not invoked from the -- host, but is invoked by (perhaps multiple) kernels. @@ -164,30 +179,32 @@ failures <- gets clFailures let params = - [[C.cparam|__global int *global_failure|], - [C.cparam|__global int *global_failure_args|]] + [ [C.cparam|__global int *global_failure|], + [C.cparam|__global typename int64_t *global_failure_args|] + ] (func, cstate) = genGPUCode FunMode (functionBody device_func) failures $ - GC.compileFun mempty params (fname, device_func) + GC.compileFun mempty params (fname, device_func) kstate = GC.compUserState cstate - modify $ \s -> s - { clUsedTypes = typesInCode (functionBody device_func) <> clUsedTypes s - , clDevFuns = M.insert fname func $ clDevFuns s - , clFailures = kernelFailures kstate - } + modify $ \s -> + s + { clUsedTypes = typesInCode (functionBody device_func) <> clUsedTypes s, + clDevFuns = M.insert fname func $ clDevFuns s, + clFailures = kernelFailures kstate + } -- Important to do this after the 'modify' call, so we propagate the -- right clFailures. void $ ensureDeviceFuns $ functionBody device_func - - where toDevice :: HostOp -> KernelOp - toDevice _ = bad + where + toDevice :: HostOp -> KernelOp + toDevice _ = bad - memParam MemParam{} = True - memParam ScalarParam{} = False + memParam MemParam {} = True + memParam ScalarParam {} = False - bad = compilerLimitationS "Cannot generate GPU functions that use arrays." + bad = compilerLimitationS "Cannot generate GPU functions that use arrays." -- Ensure that this device function is available, but don't regenerate -- it if it already exists. @@ -199,15 +216,16 @@ ensureDeviceFuns :: ImpKernels.KernelCode -> OnKernelM [Name] ensureDeviceFuns code = do let called = calledFuncs code - fmap catMaybes $ forM (S.toList called) $ \fname -> do - def <- asks $ lookupFunction fname - case def of - Just func -> do ensureDeviceFun fname func - return $ Just fname - Nothing -> return Nothing + fmap catMaybes $ + forM (S.toList called) $ \fname -> do + def <- asks $ lookupFunction fname + case def of + Just func -> do + ensureDeviceFun fname func + return $ Just fname + Nothing -> return Nothing onKernel :: KernelTarget -> Kernel -> OnKernelM OpenCL - onKernel target kernel = do called <- ensureDeviceFuns $ kernelBody kernel @@ -217,15 +235,15 @@ let (kernel_body, cstate) = genGPUCode KernelMode (kernelBody kernel) failures $ - GC.blockScope $ GC.compileCode $ kernelBody kernel + GC.blockScope $ GC.compileCode $ kernelBody kernel kstate = GC.compUserState cstate use_params = mapMaybe useAsParam $ kernelUses kernel (local_memory_args, local_memory_params, local_memory_init) = unzip3 $ - flip evalState (blankNameSource :: VNameSource) $ - mapM (prepareLocalMemory target) $ kernelLocalMemory kstate + flip evalState (blankNameSource :: VNameSource) $ + mapM (prepareLocalMemory target) $ kernelLocalMemory kstate -- CUDA has very strict restrictions on the number of blocks -- permitted along the 'y' and 'z' dimensions of the grid @@ -239,36 +257,45 @@ -- added automatically in CCUDA.hs. (perm_params, block_dim_init) = case (target, num_groups) of - (TargetCUDA, [_, _, _]) -> ([[C.cparam|const int block_dim0|], - [C.cparam|const int block_dim1|], - [C.cparam|const int block_dim2|]], - mempty) - _ -> (mempty, - [[C.citem|const int block_dim0 = 0;|], - [C.citem|const int block_dim1 = 1;|], - [C.citem|const int block_dim2 = 2;|]]) + (TargetCUDA, [_, _, _]) -> + ( [ [C.cparam|const int block_dim0|], + [C.cparam|const int block_dim1|], + [C.cparam|const int block_dim2|] + ], + mempty + ) + _ -> + ( mempty, + [ [C.citem|const int block_dim0 = 0;|], + [C.citem|const int block_dim1 = 1;|], + [C.citem|const int block_dim2 = 2;|] + ] + ) (const_defs, const_undefs) = unzip $ mapMaybe constDef $ kernelUses kernel let (safety, error_init) -- We conservatively assume that any called function can fail. | not $ null called = - (SafetyFull, []) - + (SafetyFull, []) | length (kernelFailures kstate) == length failures = - if kernelFailureTolerant kernel + if kernelFailureTolerant kernel then (SafetyNone, []) else -- No possible failures in this kernel, so if we make - -- it past an initial check, then we are good to go. - (SafetyCheap, - [C.citems|if (*global_failure >= 0) { return; }|]) + -- it past an initial check, then we are good to go. + ( SafetyCheap, + [C.citems|if (*global_failure >= 0) { return; }|] + ) | otherwise = - if not (kernelHasBarriers kstate) - then (SafetyFull, - [C.citems|if (*global_failure >= 0) { return; }|]) - else (SafetyFull, - [C.citems| + if not (kernelHasBarriers kstate) + then + ( SafetyFull, + [C.citems|if (*global_failure >= 0) { return; }|] + ) + else + ( SafetyFull, + [C.citems| volatile __local bool local_failure; if (failure_is_an_option) { int failed = *global_failure >= 0; @@ -279,17 +306,20 @@ // All threads write this value - it looks like CUDA has a compiler bug otherwise. local_failure = false; barrier(CLK_LOCAL_MEM_FENCE); - |]) + |] + ) failure_params = - [[C.cparam|__global int *global_failure|], - [C.cparam|int failure_is_an_option|], - [C.cparam|__global int *global_failure_args|]] + [ [C.cparam|__global int *global_failure|], + [C.cparam|int failure_is_an_option|], + [C.cparam|__global typename int64_t *global_failure_args|] + ] - params = perm_params ++ - take (numFailureParams safety) failure_params ++ - catMaybes local_memory_params ++ - use_params + params = + perm_params + ++ take (numFailureParams safety) failure_params + ++ catMaybes local_memory_params + ++ use_params kernel_fun = [C.cfun|__kernel void $id:name ($params:params) { @@ -303,62 +333,75 @@ $items:const_undefs }|] - modify $ \s -> s - { clKernels = M.insert name (safety, kernel_fun) $ clKernels s - , clUsedTypes = typesInKernel kernel <> clUsedTypes s - , clFailures = kernelFailures kstate - } + modify $ \s -> + s + { clKernels = M.insert name (safety, kernel_fun) $ clKernels s, + clUsedTypes = typesInKernel kernel <> clUsedTypes s, + clFailures = kernelFailures kstate + } -- The argument corresponding to the global_failure parameters is -- added automatically later. - let args = catMaybes local_memory_args ++ - kernelArgs kernel + let args = + catMaybes local_memory_args + ++ kernelArgs kernel return $ LaunchKernel safety name args num_groups group_size - where name = kernelName kernel - num_groups = kernelNumGroups kernel - group_size = kernelGroupSize kernel + where + name = kernelName kernel + num_groups = kernelNumGroups kernel + group_size = kernelGroupSize kernel - prepareLocalMemory TargetOpenCL (mem, size) = do - mem_aligned <- newVName $ baseString mem ++ "_aligned" - return (Just $ SharedMemoryKArg size, - Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|], - [C.citem|__local volatile char* restrict $id:mem = (__local volatile char*)$id:mem_aligned;|]) - prepareLocalMemory TargetCUDA (mem, size) = do - param <- newVName $ baseString mem ++ "_offset" - return (Just $ SharedMemoryKArg size, - Just [C.cparam|uint $id:param|], - [C.citem|volatile char *$id:mem = &shared_mem[$id:param];|]) + prepareLocalMemory TargetOpenCL (mem, size) = do + mem_aligned <- newVName $ baseString mem ++ "_aligned" + return + ( Just $ SharedMemoryKArg size, + Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|], + [C.citem|__local volatile char* restrict $id:mem = (__local volatile char*)$id:mem_aligned;|] + ) + prepareLocalMemory TargetCUDA (mem, size) = do + param <- newVName $ baseString mem ++ "_offset" + return + ( Just $ SharedMemoryKArg size, + Just [C.cparam|uint $id:param|], + [C.citem|volatile char *$id:mem = &shared_mem[$id:param];|] + ) useAsParam :: KernelUse -> Maybe C.Param useAsParam (ScalarUse name bt) = let ctp = case bt of -- OpenCL does not permit bool as a kernel parameter type. Bool -> [C.cty|unsigned char|] - _ -> GC.primTypeToCType bt - in Just [C.cparam|$ty:ctp $id:name|] + _ -> GC.primTypeToCType bt + in Just [C.cparam|$ty:ctp $id:name|] useAsParam (MemoryUse name) = Just [C.cparam|__global unsigned char *$id:name|] -useAsParam ConstUse{} = +useAsParam ConstUse {} = Nothing -- Constants are #defined as macros. Since a constant name in one -- kernel might potentially (although unlikely) also be used for -- something else in another kernel, we #undef them after the kernel. constDef :: KernelUse -> Maybe (C.BlockItem, C.BlockItem) -constDef (ConstUse v e) = Just ([C.citem|$escstm:def|], - [C.citem|$escstm:undef|]) - where e' = compilePrimExp e - def = "#define " ++ pretty (C.toIdent v mempty) ++ " (" ++ prettyOneLine e' ++ ")" - undef = "#undef " ++ pretty (C.toIdent v mempty) +constDef (ConstUse v e) = + Just + ( [C.citem|$escstm:def|], + [C.citem|$escstm:undef|] + ) + where + e' = compilePrimExp e + def = "#define " ++ pretty (C.toIdent v mempty) ++ " (" ++ prettyOneLine e' ++ ")" + undef = "#undef " ++ pretty (C.toIdent v mempty) constDef _ = Nothing openClCode :: [C.Func] -> String openClCode kernels = pretty [C.cunit|$edecls:funcs|] - where funcs = - [[C.cedecl|$func:kernel_func|] | - kernel_func <- kernels ] + where + funcs = + [ [C.cedecl|$func:kernel_func|] + | kernel_func <- kernels + ] atomicsDefs :: String atomicsDefs = $(embedStringFile "rts/c/atomics.h") @@ -366,13 +409,13 @@ genOpenClPrelude :: S.Set PrimType -> [C.Definition] genOpenClPrelude ts = -- Clang-based OpenCL implementations need this for 'static' to work. - [ [C.cedecl|$esc:("#ifdef cl_clang_storage_class_specifiers")|] - , [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_clang_storage_class_specifiers : enable")|] - , [C.cedecl|$esc:("#endif")|] - , [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_khr_byte_addressable_store : enable")|]] - ++ - [[C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_khr_fp64 : enable")|] | uses_float64] ++ - [C.cunit| + [ [C.cedecl|$esc:("#ifdef cl_clang_storage_class_specifiers")|], + [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_clang_storage_class_specifiers : enable")|], + [C.cedecl|$esc:("#endif")|], + [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_khr_byte_addressable_store : enable")|] + ] + ++ [[C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_khr_fp64 : enable")|] | uses_float64] + ++ [C.cunit| /* Some OpenCL programs dislike empty progams, or programs with no kernels. * Declare a dummy kernel to ensure they remain our friends. */ __kernel void dummy_kernel(__global unsigned char *dummy, int n) @@ -405,19 +448,26 @@ static inline void mem_fence_local() { mem_fence(CLK_LOCAL_MEM_FENCE); } -|] ++ - cIntOps ++ cFloat32Ops ++ cFloat32Funs ++ - (if uses_float64 then cFloat64Ops ++ cFloat64Funs ++ cFloatConvOps else []) - ++ [[C.cedecl|$esc:atomicsDefs|]] - where uses_float64 = FloatType Float64 `S.member` ts +|] + ++ cIntOps + ++ cFloat32Ops + ++ cFloat32Funs + ++ (if uses_float64 then cFloat64Ops ++ cFloat64Funs ++ cFloatConvOps else []) + ++ [[C.cedecl|$esc:atomicsDefs|]] + where + uses_float64 = FloatType Float64 `S.member` ts genCUDAPrelude :: [C.Definition] genCUDAPrelude = cudafy ++ ops - where ops = cIntOps ++ cFloat32Ops ++ cFloat32Funs ++ cFloat64Ops - ++ cFloat64Funs ++ cFloatConvOps - ++ [[C.cedecl|$esc:atomicsDefs|]] - cudafy = [CUDAC.cunit| + where + ops = + cIntOps ++ cFloat32Ops ++ cFloat32Funs ++ cFloat64Ops + ++ cFloat64Funs + ++ cFloatConvOps + ++ [[C.cedecl|$esc:atomicsDefs|]] + cudafy = + [CUDAC.cunit| $esc:("#define FUTHARK_CUDA") typedef char int8_t; @@ -523,14 +573,16 @@ compilePrimExp :: PrimExp KernelConst -> C.Exp compilePrimExp e = runIdentity $ GC.compilePrimExp compileKernelConst e - where compileKernelConst (SizeConst key) = - return [C.cexp|$id:(zEncodeString (pretty key))|] + where + compileKernelConst (SizeConst key) = + return [C.cexp|$id:(zEncodeString (pretty key))|] kernelArgs :: Kernel -> [KernelArg] kernelArgs = mapMaybe useToArg . kernelUses - where useToArg (MemoryUse mem) = Just $ MemKArg mem - useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt - useToArg ConstUse{} = Nothing + where + useToArg (MemoryUse mem) = Just $ MemKArg mem + useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt + useToArg ConstUse {} = Nothing nextErrorLabel :: GC.CompilerM KernelOp KernelState String nextErrorLabel = @@ -538,202 +590,209 @@ incErrorLabel :: GC.CompilerM KernelOp KernelState () incErrorLabel = - GC.modifyUserState $ \s -> s { kernelNextSync = kernelNextSync s + 1 } + GC.modifyUserState $ \s -> s {kernelNextSync = kernelNextSync s + 1} pendingError :: Bool -> GC.CompilerM KernelOp KernelState () pendingError b = - GC.modifyUserState $ \s -> s { kernelSyncPending = b } + GC.modifyUserState $ \s -> s {kernelSyncPending = b} hasCommunication :: ImpKernels.KernelCode -> Bool hasCommunication = any communicates - where communicates ErrorSync{} = True - communicates Barrier{} = True - communicates _ = False + where + communicates ErrorSync {} = True + communicates Barrier {} = True + communicates _ = False -- Whether we are generating code for a kernel or a device function. -- This has minor effects, such as exactly how failures are -- propagated. data OpsMode = KernelMode | FunMode deriving (Eq) -inKernelOperations :: OpsMode -> ImpKernels.KernelCode - -> GC.Operations KernelOp KernelState +inKernelOperations :: + OpsMode -> + ImpKernels.KernelCode -> + GC.Operations KernelOp KernelState inKernelOperations mode body = GC.Operations - { GC.opsCompiler = kernelOps - , GC.opsMemoryType = kernelMemoryType - , GC.opsWriteScalar = kernelWriteScalar - , GC.opsReadScalar = kernelReadScalar - , GC.opsAllocate = cannotAllocate - , GC.opsDeallocate = cannotDeallocate - , GC.opsCopy = copyInKernel - , GC.opsStaticArray = noStaticArrays - , GC.opsFatMemory = False - , GC.opsError = errorInKernel - , GC.opsCall = callInKernel - } - where has_communication = hasCommunication body - - fence FenceLocal = [C.cexp|CLK_LOCAL_MEM_FENCE|] - fence FenceGlobal = [C.cexp|CLK_GLOBAL_MEM_FENCE | CLK_LOCAL_MEM_FENCE|] - - kernelOps :: GC.OpCompiler KernelOp KernelState - kernelOps (GetGroupId v i) = - GC.stm [C.cstm|$id:v = get_group_id($int:i);|] - kernelOps (GetLocalId v i) = - GC.stm [C.cstm|$id:v = get_local_id($int:i);|] - kernelOps (GetLocalSize v i) = - GC.stm [C.cstm|$id:v = get_local_size($int:i);|] - kernelOps (GetGlobalId v i) = - GC.stm [C.cstm|$id:v = get_global_id($int:i);|] - kernelOps (GetGlobalSize v i) = - GC.stm [C.cstm|$id:v = get_global_size($int:i);|] - kernelOps (GetLockstepWidth v) = - GC.stm [C.cstm|$id:v = LOCKSTEP_WIDTH;|] - kernelOps (Barrier f) = do - GC.stm [C.cstm|barrier($exp:(fence f));|] - GC.modifyUserState $ \s -> s { kernelHasBarriers = True } - kernelOps (MemFence FenceLocal) = - GC.stm [C.cstm|mem_fence_local();|] - kernelOps (MemFence FenceGlobal) = - GC.stm [C.cstm|mem_fence_global();|] - kernelOps (LocalAlloc name size) = do - name' <- newVName $ pretty name ++ "_backing" - GC.modifyUserState $ \s -> - s { kernelLocalMemory = (name', size) : kernelLocalMemory s } - GC.stm [C.cstm|$id:name = (__local char*) $id:name';|] - kernelOps (ErrorSync f) = do - label <- nextErrorLabel - pending <- kernelSyncPending <$> GC.getUserState - when pending $ do - pendingError False - GC.stm [C.cstm|$id:label: barrier($exp:(fence f));|] - GC.stm [C.cstm|if (local_failure) { return; }|] - GC.stm [C.cstm|barrier(CLK_LOCAL_MEM_FENCE);|] -- intentional - GC.modifyUserState $ \s -> s { kernelHasBarriers = True } - incErrorLabel - kernelOps (Atomic space aop) = atomicOps space aop - - atomicCast s t = do - let volatile = [C.ctyquals|volatile|] - quals <- case s of Space sid -> pointerQuals sid - _ -> pointerQuals "global" - return [C.cty|$tyquals:(volatile++quals) $ty:t|] - - atomicSpace (Space sid) = sid - atomicSpace _ = "global" - - doAtomic s t old arr ind val op ty = do - ind' <- GC.compileExp $ unCount ind - val' <- GC.compileExp val - cast <- atomicCast s ty - GC.stm [C.cstm|$id:old = $id:op'(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|] - where op' = op ++ "_" ++ pretty t ++ "_" ++ atomicSpace s - - 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 (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) = - doAtomic s t old arr ind val "atomic_smin" [C.cty|int|] - - atomicOps s (AtomicUMax t old arr ind val) = - doAtomic s t old arr ind val "atomic_umax" [C.cty|unsigned int|] - - atomicOps s (AtomicUMin t old arr ind val) = - doAtomic s t old arr ind val "atomic_umin" [C.cty|unsigned int|] + { GC.opsCompiler = kernelOps, + GC.opsMemoryType = kernelMemoryType, + GC.opsWriteScalar = kernelWriteScalar, + GC.opsReadScalar = kernelReadScalar, + GC.opsAllocate = cannotAllocate, + GC.opsDeallocate = cannotDeallocate, + GC.opsCopy = copyInKernel, + GC.opsStaticArray = noStaticArrays, + GC.opsFatMemory = False, + GC.opsError = errorInKernel, + GC.opsCall = callInKernel, + GC.opsCritical = mempty + } + where + has_communication = hasCommunication body - atomicOps s (AtomicAnd t old arr ind val) = - doAtomic s t old arr ind val "atomic_and" [C.cty|int|] + fence FenceLocal = [C.cexp|CLK_LOCAL_MEM_FENCE|] + fence FenceGlobal = [C.cexp|CLK_GLOBAL_MEM_FENCE | CLK_LOCAL_MEM_FENCE|] - atomicOps s (AtomicOr t old arr ind val) = - doAtomic s t old arr ind val "atomic_or" [C.cty|int|] + kernelOps :: GC.OpCompiler KernelOp KernelState + kernelOps (GetGroupId v i) = + GC.stm [C.cstm|$id:v = get_group_id($int:i);|] + kernelOps (GetLocalId v i) = + GC.stm [C.cstm|$id:v = get_local_id($int:i);|] + kernelOps (GetLocalSize v i) = + GC.stm [C.cstm|$id:v = get_local_size($int:i);|] + kernelOps (GetGlobalId v i) = + GC.stm [C.cstm|$id:v = get_global_id($int:i);|] + kernelOps (GetGlobalSize v i) = + GC.stm [C.cstm|$id:v = get_global_size($int:i);|] + kernelOps (GetLockstepWidth v) = + GC.stm [C.cstm|$id:v = LOCKSTEP_WIDTH;|] + kernelOps (Barrier f) = do + GC.stm [C.cstm|barrier($exp:(fence f));|] + GC.modifyUserState $ \s -> s {kernelHasBarriers = True} + kernelOps (MemFence FenceLocal) = + GC.stm [C.cstm|mem_fence_local();|] + kernelOps (MemFence FenceGlobal) = + GC.stm [C.cstm|mem_fence_global();|] + kernelOps (LocalAlloc name size) = do + name' <- newVName $ pretty name ++ "_backing" + GC.modifyUserState $ \s -> + s {kernelLocalMemory = (name', fmap untyped size) : kernelLocalMemory s} + GC.stm [C.cstm|$id:name = (__local char*) $id:name';|] + kernelOps (ErrorSync f) = do + label <- nextErrorLabel + pending <- kernelSyncPending <$> GC.getUserState + when pending $ do + pendingError False + GC.stm [C.cstm|$id:label: barrier($exp:(fence f));|] + GC.stm [C.cstm|if (local_failure) { return; }|] + GC.stm [C.cstm|barrier(CLK_LOCAL_MEM_FENCE);|] -- intentional + GC.modifyUserState $ \s -> s {kernelHasBarriers = True} + incErrorLabel + kernelOps (Atomic space aop) = atomicOps space aop - atomicOps s (AtomicXor t old arr ind val) = - doAtomic s t old arr ind val "atomic_xor" [C.cty|int|] + atomicCast s t = do + let volatile = [C.ctyquals|volatile|] + quals <- case s of + Space sid -> pointerQuals sid + _ -> pointerQuals "global" + return [C.cty|$tyquals:(volatile++quals) $ty:t|] - atomicOps s (AtomicCmpXchg t old arr ind cmp val) = do - ind' <- GC.compileExp $ 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 + atomicSpace (Space sid) = sid + atomicSpace _ = "global" - atomicOps s (AtomicXchg t old arr ind val) = do - ind' <- GC.compileExp $ 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 + doAtomic s t old arr ind val op ty = do + ind' <- GC.compileExp $ untyped $ unCount ind + val' <- GC.compileExp val + cast <- atomicCast s ty + GC.stm [C.cstm|$id:old = $id:op'(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|] + where + op' = op ++ "_" ++ pretty t ++ "_" ++ atomicSpace s - cannotAllocate :: GC.Allocate KernelOp KernelState - cannotAllocate _ = - error "Cannot allocate memory in kernel" + 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 (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) = + doAtomic s t old arr ind val "atomic_smin" [C.cty|int|] + atomicOps s (AtomicUMax t old arr ind val) = + doAtomic s t old arr ind val "atomic_umax" [C.cty|unsigned int|] + atomicOps s (AtomicUMin t old arr ind val) = + doAtomic s t old arr ind val "atomic_umin" [C.cty|unsigned int|] + atomicOps s (AtomicAnd t old arr ind val) = + doAtomic s t old arr ind val "atomic_and" [C.cty|int|] + atomicOps s (AtomicOr t old arr ind val) = + 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 - cannotDeallocate :: GC.Deallocate KernelOp KernelState - cannotDeallocate _ _ = - error "Cannot deallocate memory in kernel" + cannotAllocate :: GC.Allocate KernelOp KernelState + cannotAllocate _ = + error "Cannot allocate memory in kernel" - copyInKernel :: GC.Copy KernelOp KernelState - copyInKernel _ _ _ _ _ _ _ = - error "Cannot bulk copy in kernel." + cannotDeallocate :: GC.Deallocate KernelOp KernelState + cannotDeallocate _ _ = + error "Cannot deallocate memory in kernel" - noStaticArrays :: GC.StaticArray KernelOp KernelState - noStaticArrays _ _ _ _ = - error "Cannot create static array in kernel." + copyInKernel :: GC.Copy KernelOp KernelState + copyInKernel _ _ _ _ _ _ _ = + error "Cannot bulk copy in kernel." - kernelMemoryType space = do - quals <- pointerQuals space - return [C.cty|$tyquals:quals $ty:defaultMemBlockType|] + noStaticArrays :: GC.StaticArray KernelOp KernelState + noStaticArrays _ _ _ _ = + error "Cannot create static array in kernel." - kernelWriteScalar = - GC.writeScalarPointerWithQuals pointerQuals + kernelMemoryType space = do + quals <- pointerQuals space + return [C.cty|$tyquals:quals $ty:defaultMemBlockType|] - kernelReadScalar = - GC.readScalarPointerWithQuals pointerQuals + kernelWriteScalar = + GC.writeScalarPointerWithQuals pointerQuals - whatNext = do - label <- nextErrorLabel - pendingError True - return $ if has_communication - then [C.citems|local_failure = true; goto $id:label;|] - else if mode == FunMode - then [C.citems|return 1;|] - else [C.citems|return;|] + kernelReadScalar = + GC.readScalarPointerWithQuals pointerQuals - callInKernel dests fname args - | isBuiltInFunction fname = - GC.opsCall GC.defaultOperations dests fname args + whatNext = do + label <- nextErrorLabel + pendingError True + return $ + if has_communication + then [C.citems|local_failure = true; goto $id:label;|] + else + if mode == FunMode + then [C.citems|return 1;|] + else [C.citems|return;|] - | otherwise = do - let out_args = [ [C.cexp|&$id:d|] | d <- dests ] - args' = [C.cexp|global_failure|] : [C.cexp|global_failure_args|] : - out_args ++ args + callInKernel dests fname args + | isBuiltInFunction fname = + GC.opsCall GC.defaultOperations dests fname args + | otherwise = do + let out_args = [[C.cexp|&$id:d|] | d <- dests] + args' = + [C.cexp|global_failure|] : + [C.cexp|global_failure_args|] : + out_args ++ args - what_next <- whatNext + what_next <- whatNext - GC.item [C.citem|if ($id:(funName fname)($args:args') != 0) { $items:what_next; }|] + GC.item [C.citem|if ($id:(funName fname)($args:args') != 0) { $items:what_next; }|] - errorInKernel msg@(ErrorMsg parts) backtrace = do - n <- length . kernelFailures <$> GC.getUserState - GC.modifyUserState $ \s -> - s { kernelFailures = kernelFailures s ++ [FailureMsg msg backtrace] } - let setArgs _ [] = return [] - setArgs i (ErrorString{} : parts') = setArgs i parts' - setArgs i (ErrorInt32 x : parts') = do - x' <- GC.compileExp x - stms <- setArgs (i+1) parts' - return $ [C.cstm|global_failure_args[$int:i] = $exp:x';|] : stms - argstms <- setArgs (0::Int) parts + errorInKernel msg@(ErrorMsg parts) backtrace = do + n <- length . kernelFailures <$> GC.getUserState + GC.modifyUserState $ \s -> + s {kernelFailures = kernelFailures s ++ [FailureMsg msg backtrace]} + let setArgs _ [] = return [] + setArgs i (ErrorString {} : parts') = setArgs i parts' + setArgs i (ErrorInt32 x : parts') = do + x' <- GC.compileExp x + stms <- setArgs (i + 1) parts' + return $ [C.cstm|global_failure_args[$int:i] = (typename int64_t)$exp:x';|] : stms + setArgs i (ErrorInt64 x : parts') = do + x' <- GC.compileExp x + stms <- setArgs (i + 1) parts' + return $ [C.cstm|global_failure_args[$int:i] = $exp:x';|] : stms + argstms <- setArgs (0 :: Int) parts - what_next <- whatNext + what_next <- whatNext - GC.stm [C.cstm|{ if (atomic_cmpxchg_i32_global(global_failure, -1, $int:n) == -1) + GC.stm + [C.cstm|{ if (atomic_cmpxchg_i32_global(global_failure, -1, $int:n) == -1) { $stms:argstms; } $items:what_next }|] @@ -746,37 +805,48 @@ typesInCode :: ImpKernels.KernelCode -> S.Set PrimType typesInCode Skip = mempty typesInCode (c1 :>>: c2) = typesInCode c1 <> typesInCode c2 -typesInCode (For _ it e c) = IntType it `S.insert` typesInExp e <> typesInCode c -typesInCode (While e c) = typesInExp e <> typesInCode c -typesInCode DeclareMem{} = mempty +typesInCode (For _ e c) = typesInExp e <> typesInCode c +typesInCode (While (TPrimExp e) c) = typesInExp e <> typesInCode c +typesInCode DeclareMem {} = mempty typesInCode (DeclareScalar _ _ t) = S.singleton t typesInCode (DeclareArray _ _ t _) = S.singleton t -typesInCode (Allocate _ (Count e) _) = typesInExp e -typesInCode Free{} = mempty -typesInCode (Copy _ (Count e1) _ _ (Count e2) _ (Count e3)) = - typesInExp e1 <> typesInExp e2 <> typesInExp e3 -typesInCode (Write _ (Count e1) t _ _ e2) = +typesInCode (Allocate _ (Count (TPrimExp e)) _) = typesInExp e +typesInCode Free {} = mempty +typesInCode + ( Copy + _ + (Count (TPrimExp e1)) + _ + _ + (Count (TPrimExp e2)) + _ + (Count (TPrimExp e3)) + ) = + typesInExp e1 <> typesInExp e2 <> typesInExp e3 +typesInCode (Write _ (Count (TPrimExp e1)) t _ _ e2) = typesInExp e1 <> S.singleton t <> typesInExp e2 typesInCode (SetScalar _ e) = typesInExp e -typesInCode SetMem{} = mempty +typesInCode SetMem {} = mempty typesInCode (Call _ _ es) = mconcat $ map typesInArg es - where typesInArg MemArg{} = mempty - typesInArg (ExpArg e) = typesInExp e -typesInCode (If e c1 c2) = + where + typesInArg MemArg {} = mempty + typesInArg (ExpArg e) = typesInExp e +typesInCode (If (TPrimExp e) c1 c2) = typesInExp e <> typesInCode c1 <> typesInCode c2 typesInCode (Assert e _ _) = typesInExp e typesInCode (Comment _ c) = typesInCode c typesInCode (DebugPrint _ v) = maybe mempty typesInExp v -typesInCode Op{} = mempty +typesInCode Op {} = mempty typesInExp :: Exp -> S.Set PrimType typesInExp (ValueExp v) = S.singleton $ primValueType v typesInExp (BinOpExp _ e1 e2) = typesInExp e1 <> typesInExp e2 typesInExp (CmpOpExp _ e1 e2) = typesInExp e1 <> typesInExp e2 typesInExp (ConvOpExp op e) = S.fromList [from, to] <> typesInExp e - where (from, to) = convOpType op + where + (from, to) = convOpType op typesInExp (UnOpExp _ e) = typesInExp e typesInExp (FunExp _ args t) = S.singleton t <> mconcat (map typesInExp args) -typesInExp (LeafExp (Index _ (Count e) t _ _) _) = S.singleton t <> typesInExp e -typesInExp (LeafExp ScalarVar{} _) = mempty +typesInExp (LeafExp (Index _ (Count (TPrimExp e)) t _ _) _) = S.singleton t <> typesInExp e +typesInExp (LeafExp ScalarVar {} _) = mempty typesInExp (LeafExp (SizeOf t) _) = S.singleton t diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs b/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs @@ -1,186 +1,263 @@ -- | Carefully optimised implementations of GPU transpositions. -- Written in ImpCode so we can compile it to both CUDA and OpenCL. module Futhark.CodeGen.ImpGen.Kernels.Transpose - ( TransposeType(..) - , TransposeArgs - , mapTransposeKernel + ( TransposeType (..), + TransposeArgs, + mapTransposeKernel, ) - where - -import Prelude hiding (quot, rem) +where import Futhark.CodeGen.ImpCode.Kernels import Futhark.IR.Prop.Types -import Futhark.Util.IntegralExp (IntegralExp, divUp, quot, rem) +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) -- | Which form of transposition to generate code for. -data TransposeType = TransposeNormal - | TransposeLowWidth - | TransposeLowHeight - | TransposeSmall -- ^ For small arrays that do not - -- benefit from coalescing. - deriving (Eq, Ord, Show) +data TransposeType + = TransposeNormal + | TransposeLowWidth + | TransposeLowHeight + | -- | For small arrays that do not + -- benefit from coalescing. + TransposeSmall + deriving (Eq, Ord, Show) -- | The types of the arguments accepted by a transposition function. -type TransposeArgs = (VName, Exp, - VName, Exp, - Exp, Exp, - Exp, Exp, Exp, - VName) +type TransposeArgs = + ( VName, + TExp Int32, + VName, + TExp Int32, + TExp Int32, + TExp Int32, + TExp Int32, + TExp Int32, + TExp Int32, + VName + ) -elemsPerThread :: IntegralExp a => a +elemsPerThread :: TExp Int32 elemsPerThread = 4 -mapTranspose :: Exp -> TransposeArgs -> PrimType -> TransposeType -> KernelCode +mapTranspose :: TExp Int32 -> TransposeArgs -> PrimType -> TransposeType -> KernelCode mapTranspose block_dim args t kind = case kind of TransposeSmall -> mconcat - [ get_ids - - , dec our_array_offset $ v32 get_global_id_0 `quot` (height*width) * (height*width) - - , dec x_index $ (v32 get_global_id_0 `rem` (height*width)) `quot` height - , dec y_index $ v32 get_global_id_0 `rem` height - - , dec odata_offset $ - (basic_odata_offset `quot` primByteSize t) + v32 our_array_offset - , dec idata_offset $ - (basic_idata_offset `quot` primByteSize t) + v32 our_array_offset - - , dec index_in $ v32 y_index * width + v32 x_index - , dec index_out $ v32 x_index * height + v32 y_index - - , If (v32 get_global_id_0 .<. width * height * num_arrays) - (Write odata (elements $ v32 odata_offset + v32 index_out) t (Space "global") Nonvolatile $ - index idata (elements $ v32 idata_offset + v32 index_in) t (Space "global") Nonvolatile) - mempty - ] - + [ get_ids, + dec our_array_offset $ vi32 get_global_id_0 `quot` (height * width) * (height * width), + dec x_index $ (vi32 get_global_id_0 `rem` (height * width)) `quot` height, + dec y_index $ vi32 get_global_id_0 `rem` height, + dec odata_offset $ + (basic_odata_offset `quot` primByteSize t) + vi32 our_array_offset, + dec idata_offset $ + (basic_idata_offset `quot` primByteSize t) + vi32 our_array_offset, + dec index_in $ vi32 y_index * width + vi32 x_index, + dec index_out $ vi32 x_index * height + vi32 y_index, + when + (vi32 get_global_id_0 .<. width * height * num_arrays) + ( Write odata (elements $ sExt64 $ vi32 odata_offset + vi32 index_out) t (Space "global") Nonvolatile $ + index idata (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) t (Space "global") Nonvolatile + ) + ] TransposeLowWidth -> - mkTranspose $ lowDimBody - (v32 get_group_id_0 * block_dim + (v32 get_local_id_0 `quot` muly)) - (v32 get_group_id_1 * block_dim * muly + v32 get_local_id_1 + - (v32 get_local_id_0 `rem` muly) * block_dim) - (v32 get_group_id_1* block_dim * muly + v32 get_local_id_0 + - (v32 get_local_id_1 `rem` muly) * block_dim) - (v32 get_group_id_0 * block_dim + (v32 get_local_id_1 `quot` muly)) - + mkTranspose $ + lowDimBody + (vi32 get_group_id_0 * block_dim + (vi32 get_local_id_0 `quot` muly)) + ( vi32 get_group_id_1 * block_dim * muly + vi32 get_local_id_1 + + (vi32 get_local_id_0 `rem` muly) * block_dim + ) + ( vi32 get_group_id_1 * block_dim * muly + vi32 get_local_id_0 + + (vi32 get_local_id_1 `rem` muly) * block_dim + ) + (vi32 get_group_id_0 * block_dim + (vi32 get_local_id_1 `quot` muly)) TransposeLowHeight -> - mkTranspose $ lowDimBody - (v32 get_group_id_0 * block_dim * mulx + v32 get_local_id_0 + - (v32 get_local_id_1 `rem` mulx) * block_dim) - (v32 get_group_id_1 * block_dim + (v32 get_local_id_1 `quot` mulx)) - (v32 get_group_id_1 * block_dim + (v32 get_local_id_0 `quot` mulx)) - (v32 get_group_id_0 * block_dim * mulx + v32 get_local_id_1 + - (v32 get_local_id_0 `rem` mulx) * block_dim) - + mkTranspose $ + lowDimBody + ( vi32 get_group_id_0 * block_dim * mulx + vi32 get_local_id_0 + + (vi32 get_local_id_1 `rem` mulx) * block_dim + ) + (vi32 get_group_id_1 * block_dim + (vi32 get_local_id_1 `quot` mulx)) + (vi32 get_group_id_1 * block_dim + (vi32 get_local_id_0 `quot` mulx)) + ( vi32 get_group_id_0 * block_dim * mulx + vi32 get_local_id_1 + + (vi32 get_local_id_0 `rem` mulx) * block_dim + ) TransposeNormal -> - mkTranspose $ mconcat - [ dec x_index $ v32 get_global_id_0 - , dec y_index $ v32 get_group_id_1 * tile_dim + v32 get_local_id_1 - , when (v32 x_index .<. width) $ - For j Int32 elemsPerThread $ - let i = v32 j * (tile_dim `quot` elemsPerThread) - in mconcat [ dec index_in $ (v32 y_index + i) * width + v32 x_index - , when (v32 y_index + i .<. height) $ - Write block (elements $ (v32 get_local_id_1 + i) * (tile_dim+1) - + v32 get_local_id_0) - t (Space "local") Nonvolatile $ - index idata (elements $ v32 idata_offset + v32 index_in) - t (Space "global") Nonvolatile] - , Op $ Barrier FenceLocal - , SetScalar x_index $ v32 get_group_id_1 * tile_dim + v32 get_local_id_0 - , SetScalar y_index $ v32 get_group_id_0 * tile_dim + v32 get_local_id_1 - , when (v32 x_index .<. height) $ - For j Int32 elemsPerThread $ - let i = v32 j * (tile_dim `quot` elemsPerThread) - in mconcat [ dec index_out $ (v32 y_index + i) * height + v32 x_index - , when (v32 y_index + i .<. width) $ - Write odata (elements $ v32 odata_offset + v32 index_out) - t (Space "global") Nonvolatile $ - index block (elements $ v32 get_local_id_0 * (tile_dim+1) - + v32 get_local_id_1+i) - t (Space "local") Nonvolatile - ] - ] - - where dec v e = DeclareScalar v Nonvolatile int32 <> SetScalar v e - v32 = flip var int32 - tile_dim = 2 * block_dim + mkTranspose $ + mconcat + [ dec x_index $ vi32 get_global_id_0, + dec y_index $ vi32 get_group_id_1 * tile_dim + vi32 get_local_id_1, + when (vi32 x_index .<. width) $ + For j (untyped elemsPerThread) $ + let i = vi32 j * (tile_dim `quot` elemsPerThread) + in mconcat + [ dec index_in $ (vi32 y_index + i) * width + vi32 x_index, + when (vi32 y_index + i .<. height) $ + Write + block + ( elements $ + sExt64 $ + (vi32 get_local_id_1 + i) * (tile_dim + 1) + + vi32 get_local_id_0 + ) + t + (Space "local") + Nonvolatile + $ index + idata + (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) + t + (Space "global") + Nonvolatile + ], + Op $ Barrier FenceLocal, + SetScalar x_index $ untyped $ vi32 get_group_id_1 * tile_dim + vi32 get_local_id_0, + SetScalar y_index $ untyped $ vi32 get_group_id_0 * tile_dim + vi32 get_local_id_1, + when (vi32 x_index .<. height) $ + For j (untyped elemsPerThread) $ + let i = vi32 j * (tile_dim `quot` elemsPerThread) + in mconcat + [ dec index_out $ (vi32 y_index + i) * height + vi32 x_index, + when (vi32 y_index + i .<. width) $ + Write + odata + (elements $ sExt64 $ vi32 odata_offset + vi32 index_out) + t + (Space "global") + Nonvolatile + $ index + block + ( elements $ + sExt64 $ + vi32 get_local_id_0 * (tile_dim + 1) + vi32 get_local_id_1 + i + ) + t + (Space "local") + Nonvolatile + ] + ] + where + dec v (TPrimExp e) = + DeclareScalar v Nonvolatile (primExpType e) <> SetScalar v e + tile_dim = 2 * block_dim - when a b = If a b mempty + when a b = If a b mempty - (odata, basic_odata_offset, idata, basic_idata_offset, - width, height, - mulx, muly, num_arrays, block) = args + ( odata, + basic_odata_offset, + idata, + basic_idata_offset, + width, + height, + mulx, + muly, + num_arrays, + block + ) = args - -- Be extremely careful when editing this list to ensure that - -- the names match up. Also, be careful that the tags on - -- these names do not conflict with the tags of the - -- surrounding code. We accomplish the latter by using very - -- low tags (normal variables start at least in the low - -- hundreds). - [ our_array_offset , x_index , y_index - , odata_offset, idata_offset, index_in, index_out - , get_global_id_0 - , get_local_id_0, get_local_id_1 - , get_group_id_0, get_group_id_1, get_group_id_2 - , j] = - zipWith (flip VName) [30..] $ map nameFromString - [ "our_array_offset" , "x_index" , "y_index" - , "odata_offset", "idata_offset", "index_in", "index_out" - , "get_global_id_0" - , "get_local_id_0", "get_local_id_1" - , "get_group_id_0", "get_group_id_1", "get_group_id_2" - , "j"] + -- Be extremely careful when editing this list to ensure that + -- the names match up. Also, be careful that the tags on + -- these names do not conflict with the tags of the + -- surrounding code. We accomplish the latter by using very + -- low tags (normal variables start at least in the low + -- hundreds). + [ our_array_offset, + x_index, + y_index, + odata_offset, + idata_offset, + index_in, + index_out, + get_global_id_0, + get_local_id_0, + get_local_id_1, + get_group_id_0, + get_group_id_1, + get_group_id_2, + j + ] = + zipWith (flip VName) [30 ..] $ + map + nameFromString + [ "our_array_offset", + "x_index", + "y_index", + "odata_offset", + "idata_offset", + "index_in", + "index_out", + "get_global_id_0", + "get_local_id_0", + "get_local_id_1", + "get_group_id_0", + "get_group_id_1", + "get_group_id_2", + "j" + ] - get_ids = - mconcat [ DeclareScalar get_global_id_0 Nonvolatile int32 - , Op $ GetGlobalId get_global_id_0 0 - , DeclareScalar get_local_id_0 Nonvolatile int32 - , Op $ GetLocalId get_local_id_0 0 - , DeclareScalar get_local_id_1 Nonvolatile int32 - , Op $ GetLocalId get_local_id_1 1 - , DeclareScalar get_group_id_0 Nonvolatile int32 - , Op $ GetGroupId get_group_id_0 0 - , DeclareScalar get_group_id_1 Nonvolatile int32 - , Op $ GetGroupId get_group_id_1 1 - , DeclareScalar get_group_id_2 Nonvolatile int32 - , Op $ GetGroupId get_group_id_2 2 - ] + get_ids = + mconcat + [ DeclareScalar get_global_id_0 Nonvolatile int32, + Op $ GetGlobalId get_global_id_0 0, + DeclareScalar get_local_id_0 Nonvolatile int32, + Op $ GetLocalId get_local_id_0 0, + DeclareScalar get_local_id_1 Nonvolatile int32, + Op $ GetLocalId get_local_id_1 1, + DeclareScalar get_group_id_0 Nonvolatile int32, + Op $ GetGroupId get_group_id_0 0, + DeclareScalar get_group_id_1 Nonvolatile int32, + Op $ GetGroupId get_group_id_1 1, + DeclareScalar get_group_id_2 Nonvolatile int32, + Op $ GetGroupId get_group_id_2 2 + ] - mkTranspose body = - mconcat - [ get_ids - , dec our_array_offset $ v32 get_group_id_2 * width * height - , dec odata_offset $ - (basic_odata_offset `quot` primByteSize t) + v32 our_array_offset - , dec idata_offset $ - (basic_idata_offset `quot` primByteSize t) + v32 our_array_offset - , body - ] + mkTranspose body = + mconcat + [ get_ids, + dec our_array_offset $ vi32 get_group_id_2 * width * height, + dec odata_offset $ + (basic_odata_offset `quot` primByteSize t) + vi32 our_array_offset, + dec idata_offset $ + (basic_idata_offset `quot` primByteSize t) + vi32 our_array_offset, + body + ] - lowDimBody x_in_index y_in_index x_out_index y_out_index = - mconcat - [ dec x_index x_in_index - , dec y_index y_in_index - , dec index_in $ v32 y_index * width + v32 x_index - , when (v32 x_index .<. width .&&. v32 y_index .<. height) $ - Write block (elements $ v32 get_local_id_1 * (block_dim+1) + v32 get_local_id_0) - t (Space "local") Nonvolatile $ - index idata (elements $ v32 idata_offset + v32 index_in) - t (Space "global") Nonvolatile - , Op $ Barrier FenceLocal - , SetScalar x_index x_out_index - , SetScalar y_index y_out_index - , dec index_out $ v32 y_index * height + v32 x_index - , when (v32 x_index .<. height .&&. v32 y_index .<. width) $ - Write odata (elements $ v32 odata_offset + v32 index_out) - t (Space "global") Nonvolatile $ - index block (elements $ v32 get_local_id_0 * (block_dim+1) + v32 get_local_id_1) - t (Space "local") Nonvolatile - ] + lowDimBody x_in_index y_in_index x_out_index y_out_index = + mconcat + [ dec x_index x_in_index, + dec y_index y_in_index, + dec index_in $ vi32 y_index * width + vi32 x_index, + when (vi32 x_index .<. width .&&. vi32 y_index .<. height) $ + Write + block + (elements $ sExt64 $ vi32 get_local_id_1 * (block_dim + 1) + vi32 get_local_id_0) + t + (Space "local") + Nonvolatile + $ index + idata + (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) + t + (Space "global") + Nonvolatile, + Op $ Barrier FenceLocal, + SetScalar x_index $ untyped x_out_index, + SetScalar y_index $ untyped y_out_index, + dec index_out $ vi32 y_index * height + vi32 x_index, + when (vi32 x_index .<. height .&&. vi32 y_index .<. width) $ + Write + odata + (elements $ sExt64 (vi32 odata_offset + vi32 index_out)) + t + (Space "global") + Nonvolatile + $ index + block + (elements $ sExt64 $ vi32 get_local_id_0 * (block_dim + 1) + vi32 get_local_id_1) + t + (Space "local") + Nonvolatile + ] -- | Generate a transpose kernel. There is special support to handle -- input arrays with low width, low height, or both. @@ -205,65 +282,105 @@ -- of block_dim*2 by block_dim*2+1 elements. Padding each row with -- an additional element prevents bank conflicts from occuring when -- the tile is accessed column-wise. -mapTransposeKernel :: String -> Integer -> TransposeArgs -> PrimType -> TransposeType - -> Kernel +mapTransposeKernel :: + String -> + Integer -> + TransposeArgs -> + PrimType -> + TransposeType -> + Kernel mapTransposeKernel desc block_dim_int args t kind = Kernel - { kernelBody = DeclareMem block (Space "local") <> - Op (LocalAlloc block block_size) <> - mapTranspose block_dim args t kind - , kernelUses = uses - , kernelNumGroups = num_groups - , kernelGroupSize = group_size - , kernelName = nameFromString name - , kernelFailureTolerant = True - } - where pad2DBytes k = k * (k + 1) * primByteSize t - block_size = - case kind of TransposeSmall -> 1 -- Not used, but AMD's - -- OpenCL does not like - -- zero-size local memory. - TransposeNormal -> fromInteger $ pad2DBytes $ 2*block_dim_int - TransposeLowWidth -> fromInteger $ pad2DBytes block_dim_int - TransposeLowHeight -> fromInteger $ pad2DBytes block_dim_int - block_dim = fromInteger block_dim_int + { kernelBody = + DeclareMem block (Space "local") + <> Op (LocalAlloc block block_size) + <> mapTranspose block_dim args t kind, + kernelUses = uses, + kernelNumGroups = map untyped num_groups, + kernelGroupSize = map untyped group_size, + kernelName = nameFromString name, + kernelFailureTolerant = True + } + where + pad2DBytes k = k * (k + 1) * primByteSize t + block_size = + bytes $ + case kind of + TransposeSmall -> 1 :: TExp Int64 + -- Not used, but AMD's OpenCL + -- does not like zero-size + -- local memory. + TransposeNormal -> fromInteger $ pad2DBytes $ 2 * block_dim_int + TransposeLowWidth -> fromInteger $ pad2DBytes block_dim_int + TransposeLowHeight -> fromInteger $ pad2DBytes block_dim_int + block_dim = fromInteger block_dim_int :: TExp Int32 - (odata, basic_odata_offset, idata, basic_idata_offset, - width, height, - mulx, muly, num_arrays, - block) = args + ( odata, + basic_odata_offset, + idata, + basic_idata_offset, + width, + height, + mulx, + muly, + num_arrays, + block + ) = args - (num_groups, group_size) = - case kind of - TransposeSmall -> - ([(num_arrays * width * height) `divUp` (block_dim * block_dim)], - [block_dim * block_dim]) - TransposeLowWidth -> - lowDimKernelAndGroupSize block_dim num_arrays width $ height `divUp` muly - TransposeLowHeight -> - lowDimKernelAndGroupSize block_dim num_arrays (width `divUp` mulx) height - TransposeNormal -> - let actual_dim = block_dim*2 - in ( [ width `divUp` actual_dim - , height `divUp` actual_dim - , num_arrays] - , [actual_dim, actual_dim `quot` elemsPerThread, 1]) + (num_groups, group_size) = + case kind of + TransposeSmall -> + ( [(num_arrays * width * height) `divUp` (block_dim * block_dim)], + [block_dim * block_dim] + ) + TransposeLowWidth -> + lowDimKernelAndGroupSize block_dim num_arrays width $ height `divUp` muly + TransposeLowHeight -> + lowDimKernelAndGroupSize block_dim num_arrays (width `divUp` mulx) height + TransposeNormal -> + let actual_dim = block_dim * 2 + in ( [ width `divUp` actual_dim, + height `divUp` actual_dim, + num_arrays + ], + [actual_dim, actual_dim `quot` elemsPerThread, 1] + ) - uses = map (`ScalarUse` int32) - (namesToList $ mconcat $ map freeIn - [basic_odata_offset, basic_idata_offset, num_arrays, - width, height, mulx, muly]) ++ - map MemoryUse [odata, idata] + uses = + map + (`ScalarUse` int32) + ( namesToList $ + mconcat $ + map + freeIn + [ basic_odata_offset, + basic_idata_offset, + num_arrays, + width, + height, + mulx, + muly + ] + ) + ++ map MemoryUse [odata, idata] - name = - case kind of TransposeSmall -> desc ++ "_small" - TransposeLowHeight -> desc ++ "_low_height" - TransposeLowWidth -> desc ++ "_low_width" - TransposeNormal -> desc + name = + case kind of + TransposeSmall -> desc ++ "_small" + TransposeLowHeight -> desc ++ "_low_height" + TransposeLowWidth -> desc ++ "_low_width" + TransposeNormal -> desc -lowDimKernelAndGroupSize :: Exp -> Exp -> Exp -> Exp -> ([Exp], [Exp]) +lowDimKernelAndGroupSize :: + TExp Int32 -> + TExp Int32 -> + TExp Int32 -> + TExp Int32 -> + ([TExp Int32], [TExp Int32]) lowDimKernelAndGroupSize block_dim num_arrays x_elems y_elems = - ([x_elems `divUp` block_dim, - y_elems `divUp` block_dim, - num_arrays], - [block_dim, block_dim, 1]) + ( [ x_elems `divUp` block_dim, + y_elems `divUp` block_dim, + num_arrays + ], + [block_dim, block_dim, 1] + ) diff --git a/src/Futhark/CodeGen/ImpGen/OpenCL.hs b/src/Futhark/CodeGen/ImpGen/OpenCL.hs --- a/src/Futhark/CodeGen/ImpGen/OpenCL.hs +++ b/src/Futhark/CodeGen/ImpGen/OpenCL.hs @@ -1,14 +1,14 @@ module Futhark.CodeGen.ImpGen.OpenCL - ( compileProg - , Warnings - ) where + ( compileProg, + Warnings, + ) +where import Data.Bifunctor (second) - -import Futhark.IR.KernelsMem import qualified Futhark.CodeGen.ImpCode.OpenCL as OpenCL import Futhark.CodeGen.ImpGen.Kernels import Futhark.CodeGen.ImpGen.Kernels.ToOpenCL +import Futhark.IR.KernelsMem import Futhark.MonadFreshNames compileProg :: MonadFreshNames m => Prog KernelsMem -> m (Warnings, OpenCL.Program) diff --git a/src/Futhark/CodeGen/ImpGen/Sequential.hs b/src/Futhark/CodeGen/ImpGen/Sequential.hs --- a/src/Futhark/CodeGen/ImpGen/Sequential.hs +++ b/src/Futhark/CodeGen/ImpGen/Sequential.hs @@ -1,11 +1,12 @@ -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + -- | Compile Futhark to sequential imperative code. module Futhark.CodeGen.ImpGen.Sequential - ( compileProg - , ImpGen.Warnings + ( compileProg, + ImpGen.Warnings, ) - where +where import qualified Futhark.CodeGen.ImpCode.Sequential as Imp import qualified Futhark.CodeGen.ImpGen as ImpGen @@ -15,7 +16,8 @@ -- | Compile a 'SeqMem' program to sequential imperative code. compileProg :: MonadFreshNames m => Prog SeqMem -> m (ImpGen.Warnings, Imp.Program) compileProg = ImpGen.compileProg () ops Imp.DefaultSpace - where ops = ImpGen.defaultOperations opCompiler - opCompiler dest (Alloc e space) = - ImpGen.compileAlloc dest e space - opCompiler _ (Inner ()) = pure () + where + ops = ImpGen.defaultOperations opCompiler + opCompiler dest (Alloc e space) = + ImpGen.compileAlloc dest e space + opCompiler _ (Inner ()) = pure () diff --git a/src/Futhark/CodeGen/ImpGen/Transpose.hs b/src/Futhark/CodeGen/ImpGen/Transpose.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Transpose.hs @@ -0,0 +1,214 @@ +-- | A cache-oblivious sequential transposition for CPU execution. +-- Generates a recursive function. +module Futhark.CodeGen.ImpGen.Transpose + ( mapTransposeFunction, + transposeArgs, + ) +where + +import Futhark.CodeGen.ImpCode +import Futhark.IR.Prop.Types +import Futhark.Util.IntegralExp +import Prelude hiding (quot) + +-- | Take well-typed arguments to the transpose function and produce +-- the actual argument list. +transposeArgs :: + PrimType -> + VName -> + Count Bytes (TExp Int64) -> + VName -> + Count Bytes (TExp Int64) -> + TExp Int64 -> + TExp Int64 -> + TExp Int64 -> + [Arg] +transposeArgs pt destmem destoffset srcmem srcoffset num_arrays m n = + [ MemArg destmem, + ExpArg $ untyped $ unCount destoffset `quot` primByteSize pt, + MemArg srcmem, + ExpArg $ untyped $ unCount srcoffset `quot` primByteSize pt, + ExpArg $ untyped num_arrays, + ExpArg $ untyped m, + ExpArg $ untyped n, + ExpArg $ untyped (0 :: TExp Int64), + ExpArg $ untyped m, + ExpArg $ untyped (0 :: TExp Int64), + ExpArg $ untyped n + ] + +-- | We need to know the name of the function we are generating, as +-- this function is recursive. +mapTransposeFunction :: Name -> PrimType -> Function op +mapTransposeFunction fname pt = + Function + False + [] + params + ( mconcat + [ dec r $ vi64 re - vi64 rb, + dec c $ vi64 ce - vi64 cb, + If (vi64 num_arrays .==. 1) doTranspose doMapTranspose + ] + ) + [] + [] + where + params = + [ memparam destmem, + intparam destoffset, + memparam srcmem, + intparam srcoffset, + intparam num_arrays, + intparam m, + intparam n, + intparam cb, + intparam ce, + intparam rb, + intparam re + ] + + memparam v = MemParam v DefaultSpace + intparam v = ScalarParam v int64 + + [ destmem, + destoffset, + srcmem, + srcoffset, + num_arrays, + n, + m, + rb, + re, + cb, + ce, + r, + c, + i, + j + ] = + zipWith + (VName . nameFromString) + [ "destmem", + "destoffset", + "srcmem", + "srcoffset", + "num_arrays", + "n", + "m", + "rb", + "re", + "cb", + "ce", + "r", + "c", + "i", + "j" -- local + ] + [0 ..] + + dec v e = DeclareScalar v Nonvolatile int32 <> SetScalar v (untyped e) + + naiveTranspose = + For j (untyped $ vi64 c) $ + For i (untyped $ vi64 r) $ + let i' = vi64 i + vi64 rb + j' = vi64 j + vi64 cb + in Write + destmem + (elements $ vi64 destoffset + j' * vi64 n + i') + pt + DefaultSpace + Nonvolatile + $ index + srcmem + (elements $ vi64 srcoffset + i' * vi64 m + j') + pt + DefaultSpace + Nonvolatile + + recArgs (cb', ce', rb', re') = + [ MemArg destmem, + ExpArg $ untyped $ vi64 destoffset, + MemArg srcmem, + ExpArg $ untyped $ vi64 srcoffset, + ExpArg $ untyped $ vi64 num_arrays, + ExpArg $ untyped $ vi64 m, + ExpArg $ untyped $ vi64 n, + ExpArg $ untyped cb', + ExpArg $ untyped ce', + ExpArg $ untyped rb', + ExpArg $ untyped re' + ] + + cutoff = 64 -- arbitrary + doTranspose = + mconcat + [ If + (vi64 r .<=. cutoff .&&. vi64 c .<=. cutoff) + naiveTranspose + $ If + (vi64 r .>=. vi64 c) + ( Call + [] + fname + ( recArgs + ( vi64 cb, + vi64 ce, + vi64 rb, + vi64 rb + (vi64 r `quot` 2) + ) + ) + <> Call + [] + fname + ( recArgs + ( vi64 cb, + vi64 ce, + vi64 rb + vi64 r `quot` 2, + vi64 re + ) + ) + ) + ( Call + [] + fname + ( recArgs + ( vi64 cb, + vi64 cb + (vi64 c `quot` 2), + vi64 rb, + vi64 re + ) + ) + <> Call + [] + fname + ( recArgs + ( vi64 cb + vi64 c `quot` 2, + vi64 ce, + vi64 rb, + vi64 re + ) + ) + ) + ] + + doMapTranspose = + -- In the map-transpose case, we assume that cb==rb==0, ce==m, + -- re==n. + For i (untyped $ vi64 num_arrays) $ + Call + [] + fname + [ MemArg destmem, + ExpArg $ untyped $ vi64 destoffset + vi64 i * vi64 m * vi64 n, + MemArg srcmem, + ExpArg $ untyped $ vi64 srcoffset + vi64 i * vi64 m * vi64 n, + ExpArg $ untyped (1 :: TExp Int64), + ExpArg $ untyped $ vi64 m, + ExpArg $ untyped $ vi64 n, + ExpArg $ untyped $ vi64 cb, + ExpArg $ untyped $ vi64 ce, + ExpArg $ untyped $ vi64 rb, + ExpArg $ untyped $ vi64 re + ] 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 @@ -10,13 +10,13 @@ -- We also use this to select reasonable default group sizes and group -- counts. module Futhark.CodeGen.OpenCL.Heuristics - ( SizeHeuristic (..) - , DeviceType (..) - , WhichSize (..) - , DeviceInfo (..) - , sizeHeuristicsTable - ) - where + ( SizeHeuristic (..), + DeviceType (..), + WhichSize (..), + DeviceInfo (..), + sizeHeuristicsTable, + ) +where import Futhark.Analysis.PrimExp import Futhark.Util.Pretty @@ -37,33 +37,32 @@ data WhichSize = LockstepWidth | NumGroups | GroupSize | TileSize | Threshold -- | A heuristic for setting the default value for something. -data SizeHeuristic = - SizeHeuristic { platformName :: String - , deviceType :: DeviceType - , heuristicSize :: WhichSize - , heuristicValue :: PrimExp DeviceInfo - } +data SizeHeuristic = SizeHeuristic + { platformName :: String, + deviceType :: DeviceType, + heuristicSize :: WhichSize, + heuristicValue :: TPrimExp Int32 DeviceInfo + } -- | All of our heuristics. sizeHeuristicsTable :: [SizeHeuristic] sizeHeuristicsTable = - [ SizeHeuristic "NVIDIA CUDA" DeviceGPU LockstepWidth $ constant 32 - , SizeHeuristic "AMD Accelerated Parallel Processing" DeviceGPU LockstepWidth $ constant 32 - , SizeHeuristic "" DeviceGPU LockstepWidth $ constant 1 - -- We calculate the number of groups to aim for 1024 threads per - -- compute unit if we also use the default group size. This seems - -- to perform well in practice. - , SizeHeuristic "" DeviceGPU NumGroups $ 4 * max_compute_units - , SizeHeuristic "" DeviceGPU GroupSize $ constant 256 - , SizeHeuristic "" DeviceGPU TileSize $ constant 32 - , SizeHeuristic "" DeviceGPU Threshold $ constant $ 32*1024 - - , SizeHeuristic "" DeviceCPU LockstepWidth $ constant 1 - , SizeHeuristic "" DeviceCPU NumGroups max_compute_units - , SizeHeuristic "" DeviceCPU GroupSize $ constant 32 - , SizeHeuristic "" DeviceCPU TileSize $ constant 4 - , SizeHeuristic "" DeviceCPU Threshold max_compute_units + [ SizeHeuristic "NVIDIA CUDA" DeviceGPU LockstepWidth 32, + SizeHeuristic "AMD Accelerated Parallel Processing" DeviceGPU LockstepWidth 32, + SizeHeuristic "" DeviceGPU LockstepWidth 1, + -- We calculate the number of groups to aim for 1024 threads per + -- compute unit if we also use the default group size. This seems + -- to perform well in practice. + SizeHeuristic "" DeviceGPU NumGroups $ 4 * max_compute_units, + SizeHeuristic "" DeviceGPU GroupSize 256, + SizeHeuristic "" DeviceGPU TileSize 32, + 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 Threshold max_compute_units ] - where constant = ValueExp . IntValue . Int32Value - max_compute_units = - LeafExp (DeviceInfo "MAX_COMPUTE_UNITS") $ IntType Int32 + where + max_compute_units = + TPrimExp $ LeafExp (DeviceInfo "MAX_COMPUTE_UNITS") $ IntType Int32 diff --git a/src/Futhark/CodeGen/SetDefaultSpace.hs b/src/Futhark/CodeGen/SetDefaultSpace.hs --- a/src/Futhark/CodeGen/SetDefaultSpace.hs +++ b/src/Futhark/CodeGen/SetDefaultSpace.hs @@ -3,9 +3,9 @@ -- to GPU memory for most of the pipeline, but final code generation -- assumes that 'DefaultSpace' is CPU memory. module Futhark.CodeGen.SetDefaultSpace - ( setDefaultSpace - ) - where + ( setDefaultSpace, + ) +where import Futhark.CodeGen.ImpCode @@ -14,18 +14,22 @@ setDefaultSpace :: Space -> Definitions op -> Definitions op setDefaultSpace space (Definitions (Constants ps consts) (Functions fundecs)) = Definitions - (Constants (map (setParamSpace space) ps) (setBodySpace space consts)) - (Functions [ (fname, setFunctionSpace space func) - | (fname, func) <- fundecs ]) + (Constants (map (setParamSpace space) ps) (setBodySpace space consts)) + ( Functions + [ (fname, setFunctionSpace space func) + | (fname, func) <- fundecs + ] + ) setFunctionSpace :: Space -> Function op -> Function op setFunctionSpace space (Function entry outputs inputs body results args) = - Function entry - (map (setParamSpace space) outputs) - (map (setParamSpace space) inputs) - (setBodySpace space body) - (map (setExtValueSpace space) results) - (map (setExtValueSpace space) args) + Function + entry + (map (setParamSpace space) outputs) + (map (setParamSpace space) inputs) + (setBodySpace space body) + (map (setExtValueSpace space) results) + (map (setExtValueSpace space) args) setParamSpace :: Space -> Param -> Param setParamSpace space (MemParam name DefaultSpace) = @@ -47,7 +51,7 @@ setBodySpace :: Space -> Code op -> Code op setBodySpace space (Allocate v e old_space) = - Allocate v (fmap (setExpSpace space) e) $ setSpace space old_space + Allocate v (fmap (setTExpSpace space) e) $ setSpace space old_space setBodySpace space (Free v old_space) = Free v $ setSpace space old_space setBodySpace space (DeclareMem name old_space) = @@ -56,22 +60,32 @@ DeclareArray name space t vs setBodySpace space (Copy dest dest_offset dest_space src src_offset src_space n) = Copy - dest (fmap (setExpSpace space) dest_offset) dest_space' - src (fmap (setExpSpace space) src_offset) src_space' $ - fmap (setExpSpace space) n - where dest_space' = setSpace space dest_space - src_space' = setSpace space src_space + dest + (fmap (setTExpSpace space) dest_offset) + dest_space' + src + (fmap (setTExpSpace space) src_offset) + src_space' + $ fmap (setTExpSpace space) n + where + dest_space' = setSpace space dest_space + src_space' = setSpace space src_space setBodySpace space (Write dest dest_offset bt dest_space vol e) = - Write dest (fmap (setExpSpace space) dest_offset) bt (setSpace space dest_space) - vol (setExpSpace space e) + Write + dest + (fmap (setTExpSpace space) dest_offset) + bt + (setSpace space dest_space) + vol + (setExpSpace space e) setBodySpace space (c1 :>>: c2) = setBodySpace space c1 :>>: setBodySpace space c2 -setBodySpace space (For i it e body) = - For i it (setExpSpace space e) $ setBodySpace space body +setBodySpace space (For i e body) = + For i (setExpSpace space e) $ setBodySpace space body setBodySpace space (While e body) = - While (setExpSpace space e) $ setBodySpace space body + While (setTExpSpace space e) $ setBodySpace space body setBodySpace space (If e c1 c2) = - If (setExpSpace space e) (setBodySpace space c1) (setBodySpace space c2) + If (setTExpSpace space e) (setBodySpace space c1) (setBodySpace space c2) setBodySpace space (Comment s c) = Comment s $ setBodySpace space c setBodySpace _ Skip = @@ -84,8 +98,9 @@ SetMem to from $ setSpace space old_space setBodySpace space (Call dests fname args) = Call dests fname $ map setArgSpace args - where setArgSpace (MemArg m) = MemArg m - setArgSpace (ExpArg e) = ExpArg $ setExpSpace space e + where + setArgSpace (MemArg m) = MemArg m + setArgSpace (ExpArg e) = ExpArg $ setExpSpace space e setBodySpace space (Assert e msg loc) = Assert (setExpSpace space e) msg loc setBodySpace space (DebugPrint s v) = @@ -95,10 +110,14 @@ setExpSpace :: Space -> Exp -> Exp setExpSpace space = fmap setLeafSpace - where setLeafSpace (Index mem i bt DefaultSpace vol) = - Index mem i bt space vol - setLeafSpace e = e + where + setLeafSpace (Index mem i bt DefaultSpace vol) = + Index mem i bt space vol + setLeafSpace e = e +setTExpSpace :: Space -> TExp t -> TExp t +setTExpSpace space = TPrimExp . setExpSpace space . untyped + setSpace :: Space -> Space -> Space setSpace space DefaultSpace = space -setSpace _ space = space +setSpace _ space = space diff --git a/src/Futhark/Compiler.hs b/src/Futhark/Compiler.hs --- a/src/Futhark/Compiler.hs +++ b/src/Futhark/Compiler.hs @@ -1,61 +1,62 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Safe #-} {-# LANGUAGE Strict #-} + -- | High-level API for invoking the Futhark compiler. module Futhark.Compiler - ( - runPipelineOnProgram - , runCompilerOnProgram - - , FutharkConfig (..) - , newFutharkConfig - , dumpError - , handleWarnings - - , module Futhark.Compiler.Program - , readProgram - , readLibrary - , readProgramOrDie - ) + ( runPipelineOnProgram, + runCompilerOnProgram, + FutharkConfig (..), + newFutharkConfig, + dumpError, + handleWarnings, + module Futhark.Compiler.Program, + readProgram, + readLibrary, + readProgramOrDie, + ) where import Control.Monad -import Control.Monad.Reader import Control.Monad.Except -import System.Exit (exitWith, ExitCode(..)) -import System.IO import qualified Data.Text.IO as T - import qualified Futhark.Analysis.Alias as Alias -import Futhark.Internalise -import Futhark.Pipeline -import Futhark.MonadFreshNames +import Futhark.Compiler.Program import Futhark.IR import qualified Futhark.IR.SOACS as I +import Futhark.Internalise +import Futhark.MonadFreshNames +import Futhark.Pipeline import qualified Futhark.TypeCheck as I -import Futhark.Compiler.Program import Futhark.Util.Log -import Futhark.Util.Pretty (prettyText, ppr) +import Futhark.Util.Pretty (ppr, prettyText) +import System.Exit (ExitCode (..), exitWith) +import System.IO -- | The compiler configuration. This only contains options related -- to core compiler functionality, such as reading the initial program -- and running passes. Options related to code generation are handled -- elsewhere. data FutharkConfig = FutharkConfig - { futharkVerbose :: (Verbosity, Maybe FilePath) - , futharkWarn :: Bool -- ^ Warn if True. - , futharkWerror :: Bool -- ^ If true, error on any warnings. - , futharkSafe :: Bool -- ^ If True, ignore @unsafe@. - } + { futharkVerbose :: (Verbosity, Maybe FilePath), + -- | Warn if True. + futharkWarn :: Bool, + -- | If true, error on any warnings. + futharkWerror :: Bool, + -- | If True, ignore @unsafe@. + futharkSafe :: Bool + } -- | The default compiler configuration. newFutharkConfig :: FutharkConfig -newFutharkConfig = FutharkConfig { futharkVerbose = (NotVerbose, Nothing) - , futharkWarn = True - , futharkWerror = False - , futharkSafe = False - } +newFutharkConfig = + FutharkConfig + { futharkVerbose = (NotVerbose, Nothing), + futharkWarn = True, + futharkWerror = False, + futharkSafe = False + } -- | Print a compiler error to stdout. The 'FutharkConfig' controls -- to which degree auxiliary information (e.g. the failing program) is @@ -75,19 +76,24 @@ T.hPutStrLn stderr "Known compiler limitation encountered. Sorry." T.hPutStrLn stderr "Revise your program or try a different Futhark compiler." report s info - where report s info = do - T.hPutStrLn stderr s - when (fst (futharkVerbose config) > NotVerbose) $ - maybe (T.hPutStr stderr) T.writeFile - (snd (futharkVerbose config)) $ info <> "\n" + where + report s info = do + T.hPutStrLn stderr s + when (fst (futharkVerbose config) > NotVerbose) $ + maybe + (T.hPutStr stderr) + T.writeFile + (snd (futharkVerbose config)) + $ info <> "\n" -- | Read a program from the given 'FilePath', run the given -- 'Pipeline', and finish up with the given 'Action'. -runCompilerOnProgram :: FutharkConfig - -> Pipeline I.SOACS lore - -> Action lore - -> FilePath - -> IO () +runCompilerOnProgram :: + FutharkConfig -> + Pipeline I.SOACS lore -> + Action lore -> + FilePath -> + IO () runCompilerOnProgram config pipeline action file = do res <- runFutharkM compile $ fst $ futharkVerbose config case res of @@ -96,25 +102,27 @@ exitWith $ ExitFailure 2 Right () -> return () - where compile = do - prog <- runPipelineOnProgram config pipeline file - when ((>NotVerbose) . fst $ futharkVerbose config) $ - logMsg $ "Running action " ++ actionName action - actionProcedure action prog - when ((>NotVerbose) . fst $ futharkVerbose config) $ - logMsg ("Done." :: String) + where + compile = do + prog <- runPipelineOnProgram config pipeline file + when ((> NotVerbose) . fst $ futharkVerbose config) $ + logMsg $ "Running action " ++ actionName action + actionProcedure action prog + when ((> NotVerbose) . fst $ futharkVerbose config) $ + logMsg ("Done." :: String) -- | Read a program from the given 'FilePath', run the given -- 'Pipeline', and return it. -runPipelineOnProgram :: FutharkConfig - -> Pipeline I.SOACS tolore - -> FilePath - -> FutharkM (Prog tolore) +runPipelineOnProgram :: + FutharkConfig -> + Pipeline I.SOACS tolore -> + FilePath -> + FutharkM (Prog tolore) runPipelineOnProgram config pipeline file = do when (pipelineVerbose pipeline_config) $ logMsg ("Reading and type-checking source program" :: String) (prog_imports, namesrc) <- - handleWarnings config $ (\(a,b,c) -> (a,(b,c))) <$> readProgram file + handleWarnings config $ (\(a, b, c) -> (a, (b, c))) <$> readProgram file putNameSource namesrc when (pipelineVerbose pipeline_config) $ @@ -124,27 +132,34 @@ logMsg ("Type-checking internalised program" :: String) typeCheckInternalProgram int_prog runPipeline pipeline pipeline_config int_prog - where pipeline_config = - PipelineConfig { pipelineVerbose = fst (futharkVerbose config) > NotVerbose - , pipelineValidate = True - } + where + pipeline_config = + PipelineConfig + { pipelineVerbose = fst (futharkVerbose config) > NotVerbose, + pipelineValidate = True + } typeCheckInternalProgram :: I.Prog I.SOACS -> FutharkM () typeCheckInternalProgram prog = case I.checkProg prog' of Left err -> internalErrorS ("After internalisation:\n" ++ show err) (ppr prog') Right () -> return () - where prog' = Alias.aliasAnalysis prog + where + prog' = Alias.aliasAnalysis prog -- | Read and type-check a Futhark program, including all imports. -readProgram :: (MonadError CompilerError m, MonadIO m) => - FilePath -> m (Warnings, Imports, VNameSource) +readProgram :: + (MonadError CompilerError m, MonadIO m) => + FilePath -> + m (Warnings, Imports, VNameSource) readProgram = readLibrary . pure -- | Read and type-check a collection of Futhark files, including all -- imports. -readLibrary :: (MonadError CompilerError m, MonadIO m) => - [FilePath] -> m (Warnings, Imports, VNameSource) +readLibrary :: + (MonadError CompilerError m, MonadIO m) => + [FilePath] -> + m (Warnings, Imports, VNameSource) readLibrary = readLibraryWithBasis emptyBasis -- | Not verbose, and terminates process on error. diff --git a/src/Futhark/Compiler/CLI.hs b/src/Futhark/Compiler/CLI.hs --- a/src/Futhark/Compiler/CLI.hs +++ b/src/Futhark/Compiler/CLI.hs @@ -3,131 +3,177 @@ -- A small amount of flexibility is provided for backend-specific -- options. module Futhark.Compiler.CLI - ( compilerMain - , CompilerOption - , CompilerMode(..) - , module Futhark.Pipeline - , module Futhark.Compiler - ) + ( compilerMain, + CompilerOption, + CompilerMode (..), + module Futhark.Pipeline, + module Futhark.Compiler, + ) where import Control.Monad import Data.Maybe -import System.FilePath -import System.Console.GetOpt -import System.IO - -import Futhark.Pipeline import Futhark.Compiler import Futhark.IR (Prog) import Futhark.IR.SOACS (SOACS) +import Futhark.Pipeline import Futhark.Util.Options +import System.Console.GetOpt +import System.FilePath +import System.IO -- | Run a parameterised Futhark compiler, where @cfg@ is a user-given -- configuration type. Call this from @main@. -compilerMain :: cfg -- ^ Initial configuration. - -> [CompilerOption cfg] -- ^ Options that affect the configuration. - -> String -- ^ The short action name (e.g. "compile to C"). - -> String -- ^ The longer action description. - -> Pipeline SOACS lore -- ^ The pipeline to use. - -> (FutharkConfig -> cfg -> CompilerMode -> FilePath -> Prog lore - -> FutharkM ()) - -- ^ The action to take on the result of the pipeline. - -> String -- ^ Program name - -> [String] -- ^ Command line arguments. - -> IO () +compilerMain :: + -- | Initial configuration. + cfg -> + -- | Options that affect the configuration. + [CompilerOption cfg] -> + -- | The short action name (e.g. "compile to C"). + String -> + -- | The longer action description. + String -> + -- | The pipeline to use. + Pipeline SOACS lore -> + -- | The action to take on the result of the pipeline. + ( FutharkConfig -> + cfg -> + CompilerMode -> + FilePath -> + Prog lore -> + FutharkM () + ) -> + -- | Program name + String -> + -- | Command line arguments. + [String] -> + IO () compilerMain cfg cfg_opts name desc pipeline doIt prog args = do hSetEncoding stdout utf8 hSetEncoding stderr utf8 - mainWithOptions (newCompilerConfig cfg) (commandLineOptions ++ map wrapOption cfg_opts) - "options... " inspectNonOptions prog args - where inspectNonOptions [file] config = Just $ compile config file - inspectNonOptions _ _ = Nothing + mainWithOptions + (newCompilerConfig cfg) + (commandLineOptions ++ map wrapOption cfg_opts) + "options... " + inspectNonOptions + prog + args + where + inspectNonOptions [file] config = Just $ compile config file + inspectNonOptions _ _ = Nothing - compile config filepath = - runCompilerOnProgram (futharkConfig config) - pipeline (action config filepath) filepath + compile config filepath = + runCompilerOnProgram + (futharkConfig config) + pipeline + (action config filepath) + filepath - action config filepath = - Action { actionName = name - , actionDescription = desc - , actionProcedure = - doIt (futharkConfig config) - (compilerConfig config) - (compilerMode config) - (outputFilePath filepath config) - } + action config filepath = + Action + { actionName = name, + actionDescription = desc, + actionProcedure = + doIt + (futharkConfig config) + (compilerConfig config) + (compilerMode config) + (outputFilePath filepath config) + } -- | An option that modifies the configuration of type @cfg@. type CompilerOption cfg = OptDescr (Either (IO ()) (cfg -> cfg)) -type CoreCompilerOption cfg = OptDescr (Either - (IO ()) - (CompilerConfig cfg -> CompilerConfig cfg)) +type CoreCompilerOption cfg = + OptDescr + ( Either + (IO ()) + (CompilerConfig cfg -> CompilerConfig cfg) + ) commandLineOptions :: [CoreCompilerOption cfg] commandLineOptions = - [ Option "o" [] - (ReqArg (\filename -> Right $ \config -> config { compilerOutput = Just filename }) - "FILE") - "Name of the compiled binary." - , Option "v" ["verbose"] - (OptArg (Right . incVerbosity) "FILE") - "Print verbose output on standard error; wrong program to FILE." - , Option [] ["library"] - (NoArg $ Right $ \config -> config { compilerMode = ToLibrary }) - "Generate a library instead of an executable." - , Option [] ["executable"] - (NoArg $ Right $ \config -> config { compilerMode = ToExecutable }) - "Generate an executable instead of a library (set by default)." - , Option "w" [] - (NoArg $ Right $ \config -> config { compilerWarn = False }) - "Disable all warnings." - , Option [] ["Werror"] - (NoArg $ Right $ \config -> config { compilerWerror = True }) - "Treat warnings as errors." - , Option [] ["safe"] - (NoArg $ Right $ \config -> config { compilerSafe = True }) - "Ignore 'unsafe' in code." + [ Option + "o" + [] + ( ReqArg + (\filename -> Right $ \config -> config {compilerOutput = Just filename}) + "FILE" + ) + "Name of the compiled binary.", + Option + "v" + ["verbose"] + (OptArg (Right . incVerbosity) "FILE") + "Print verbose output on standard error; wrong program to FILE.", + Option + [] + ["library"] + (NoArg $ Right $ \config -> config {compilerMode = ToLibrary}) + "Generate a library instead of an executable.", + Option + [] + ["executable"] + (NoArg $ Right $ \config -> config {compilerMode = ToExecutable}) + "Generate an executable instead of a library (set by default).", + Option + "w" + [] + (NoArg $ Right $ \config -> config {compilerWarn = False}) + "Disable all warnings.", + Option + [] + ["Werror"] + (NoArg $ Right $ \config -> config {compilerWerror = True}) + "Treat warnings as errors.", + Option + [] + ["safe"] + (NoArg $ Right $ \config -> config {compilerSafe = True}) + "Ignore 'unsafe' in code." ] wrapOption :: CompilerOption cfg -> CoreCompilerOption cfg wrapOption = fmap wrap - where wrap f = do - g <- f - return $ \cfg -> cfg { compilerConfig = g (compilerConfig cfg) } + where + wrap f = do + g <- f + return $ \cfg -> cfg {compilerConfig = g (compilerConfig cfg)} incVerbosity :: Maybe FilePath -> CompilerConfig cfg -> CompilerConfig cfg incVerbosity file cfg = - cfg { compilerVerbose = (v, file `mplus` snd (compilerVerbose cfg)) } - where v = case fst $ compilerVerbose cfg of - NotVerbose -> Verbose - Verbose -> VeryVerbose - VeryVerbose -> VeryVerbose + cfg {compilerVerbose = (v, file `mplus` snd (compilerVerbose cfg))} + where + v = case fst $ compilerVerbose cfg of + NotVerbose -> Verbose + Verbose -> VeryVerbose + VeryVerbose -> VeryVerbose -data CompilerConfig cfg = - CompilerConfig { compilerOutput :: Maybe FilePath - , compilerVerbose :: (Verbosity, Maybe FilePath) - , compilerMode :: CompilerMode - , compilerWerror :: Bool - , compilerSafe :: Bool - , compilerWarn :: Bool - , compilerConfig :: cfg - } +data CompilerConfig cfg = CompilerConfig + { compilerOutput :: Maybe FilePath, + compilerVerbose :: (Verbosity, Maybe FilePath), + compilerMode :: CompilerMode, + compilerWerror :: Bool, + compilerSafe :: Bool, + compilerWarn :: Bool, + compilerConfig :: cfg + } -- | Are we compiling a library or an executable? data CompilerMode = ToLibrary | ToExecutable deriving (Eq, Ord, Show) -- | The configuration of the compiler. newCompilerConfig :: cfg -> CompilerConfig cfg -newCompilerConfig x = CompilerConfig { compilerOutput = Nothing - , compilerVerbose = (NotVerbose, Nothing) - , compilerMode = ToExecutable - , compilerWerror = False - , compilerSafe = False - , compilerWarn = True - , compilerConfig = x - } +newCompilerConfig x = + CompilerConfig + { compilerOutput = Nothing, + compilerVerbose = (NotVerbose, Nothing), + compilerMode = ToExecutable, + compilerWerror = False, + compilerSafe = False, + compilerWarn = True, + compilerConfig = x + } outputFilePath :: FilePath -> CompilerConfig cfg -> FilePath outputFilePath srcfile = @@ -135,8 +181,9 @@ futharkConfig :: CompilerConfig cfg -> FutharkConfig futharkConfig config = - newFutharkConfig { futharkVerbose = compilerVerbose config - , futharkWerror = compilerWerror config - , futharkSafe = compilerSafe config - , futharkWarn = compilerWarn config - } + newFutharkConfig + { futharkVerbose = compilerVerbose config, + futharkWerror = compilerWerror config, + futharkSafe = compilerSafe config, + futharkWarn = compilerWarn config + } diff --git a/src/Futhark/Compiler/Program.hs b/src/Futhark/Compiler/Program.hs --- a/src/Futhark/Compiler/Program.hs +++ b/src/Futhark/Compiler/Program.hs @@ -1,84 +1,94 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TupleSections #-} + -- | Low-level compilation parts. Look at "Futhark.Compiler" for a -- more high-level API. module Futhark.Compiler.Program - ( readLibraryWithBasis - , readImports - , Imports - , FileModule(..) - , E.Warnings - - , Basis(..) - , emptyBasis - ) + ( readLibraryWithBasis, + readImports, + Imports, + FileModule (..), + E.Warnings, + Basis (..), + emptyBasis, + ) where import Control.Exception import Control.Monad +import Control.Monad.Except import Control.Monad.Reader import Control.Monad.State -import Control.Monad.Except -import Data.Maybe import Data.List (intercalate) -import qualified System.FilePath.Posix as Posix -import System.IO.Error +import Data.Maybe import qualified Data.Text as T import qualified Data.Text.IO as T - import Futhark.Error import Futhark.FreshNames -import Language.Futhark.Parser +import Futhark.Util.Pretty (ppr) import qualified Language.Futhark as E -import qualified Language.Futhark.TypeChecker as E -import Language.Futhark.Semantic +import Language.Futhark.Parser import Language.Futhark.Prelude -import Futhark.Util.Pretty (ppr) +import Language.Futhark.Semantic +import qualified Language.Futhark.TypeChecker as E +import qualified System.FilePath.Posix as Posix +import System.IO.Error -- | A little monad for reading and type-checking a Futhark program. type CompilerM m = ReaderT [FilePath] (StateT ReaderState m) -data ReaderState = ReaderState { alreadyImported :: Imports - , nameSource :: VNameSource - , warnings :: E.Warnings - } +data ReaderState = ReaderState + { alreadyImported :: Imports, + nameSource :: VNameSource, + warnings :: E.Warnings + } -- | Pre-typechecked imports, including a starting point for the name source. -data Basis = Basis { basisImports :: Imports - , basisNameSource :: VNameSource - , basisRoots :: [String] - -- ^ Files that should be implicitly opened. - } +data Basis = Basis + { basisImports :: Imports, + basisNameSource :: VNameSource, + -- | Files that should be implicitly opened. + basisRoots :: [String] + } -- | A basis that contains no imports, and has a properly initialised -- name source. emptyBasis :: Basis -emptyBasis = Basis { basisImports = mempty - , basisNameSource = src - , basisRoots = mempty - } - where src = newNameSource $ E.maxIntrinsicTag + 1 +emptyBasis = + Basis + { basisImports = mempty, + basisNameSource = src, + basisRoots = mempty + } + where + src = newNameSource $ E.maxIntrinsicTag + 1 -readImport :: (MonadError CompilerError m, MonadIO m) => - [ImportName] -> ImportName -> CompilerM m () +readImport :: + (MonadError CompilerError m, MonadIO m) => + [ImportName] -> + ImportName -> + CompilerM m () readImport steps include | include `elem` steps = - externalErrorS $ - "Import cycle: " ++ intercalate " -> " - (map includeToString $ reverse $ include:steps) + externalErrorS $ + "Import cycle: " + ++ intercalate + " -> " + (map includeToString $ reverse $ include : steps) | otherwise = do - already_done <- gets $ isJust . lookup (includeToString include) . alreadyImported + already_done <- gets $ isJust . lookup (includeToString include) . alreadyImported - unless already_done $ - uncurry (handleFile steps include) =<< readImportFile include + unless already_done $ + uncurry (handleFile steps include) =<< readImportFile include -handleFile :: (MonadIO m, MonadError CompilerError m) => - [ImportName] - -> ImportName - -> T.Text - -> FilePath - -> CompilerM m () +handleFile :: + (MonadIO m, MonadError CompilerError m) => + [ImportName] -> + ImportName -> + T.Text -> + FilePath -> + CompilerM m () handleFile steps include file_contents file_name = do prog <- case parseFuthark file_name file_contents of Left err -> externalErrorS $ show err @@ -98,94 +108,122 @@ externalError $ ppr err Right (m, ws, src') -> modify $ \s -> - s { alreadyImported = (includeToString include,m) : imports - , nameSource = src' - , warnings = warnings s <> ws + s + { alreadyImported = (includeToString include, m) : imports, + nameSource = src', + warnings = warnings s <> ws } - where steps' = include:steps + where + steps' = include : steps readFileSafely :: String -> IO (Maybe (Either String (String, T.Text))) readFileSafely filepath = (Just . Right . (filepath,) <$> T.readFile filepath) `catch` couldNotRead - where couldNotRead e - | isDoesNotExistError e = - return Nothing - | otherwise = - return $ Just $ Left $ show e + where + couldNotRead e + | isDoesNotExistError e = + return Nothing + | otherwise = + return $ Just $ Left $ show e -readImportFile :: (MonadError CompilerError m, MonadIO m) => - ImportName -> m (T.Text, FilePath) +readImportFile :: + (MonadError CompilerError m, MonadIO m) => + ImportName -> + m (T.Text, FilePath) readImportFile include = do -- First we try to find a file of the given name in the search path, -- then we look at the builtin library if we have to. For the -- builtins, we don't use the search path. r <- liftIO $ readFileSafely $ includeToFilePath include case (r, lookup prelude_str prelude) of - (Just (Right (filepath,s)), _) -> return (s, filepath) - (Just (Left e), _) -> externalErrorS e - (Nothing, Just t) -> return (t, prelude_str) - (Nothing, Nothing) -> externalErrorS not_found - where prelude_str = "/" Posix. includeToString include Posix.<.> "fut" + (Just (Right (filepath, s)), _) -> return (s, filepath) + (Just (Left e), _) -> externalErrorS e + (Nothing, Just t) -> return (t, prelude_str) + (Nothing, Nothing) -> externalErrorS not_found + where + prelude_str = "/" Posix. includeToString include Posix.<.> "fut" - not_found = - "Error at " ++ E.locStr (E.srclocOf include) ++ - ": could not find import '" ++ includeToString include ++ "'." + not_found = + "Error at " ++ E.locStr (E.srclocOf include) + ++ ": could not find import '" + ++ includeToString include + ++ "'." -- | Read Futhark files from some basis, and printing log messages if -- the first parameter is True. -readLibraryWithBasis :: (MonadError CompilerError m, MonadIO m) => - Basis -> [FilePath] - -> m (E.Warnings, - Imports, - VNameSource) +readLibraryWithBasis :: + (MonadError CompilerError m, MonadIO m) => + Basis -> + [FilePath] -> + m + ( E.Warnings, + Imports, + VNameSource + ) readLibraryWithBasis builtin fps = do - (_, imps, src) <- runCompilerM builtin $ - readImport [] $ mkInitialImport "/prelude/prelude" + (_, imps, src) <- + runCompilerM builtin $ + readImport [] $ mkInitialImport "/prelude/prelude" let basis = Basis imps src ["/prelude/prelude"] readLibrary' basis fps -- | Read and type-check a Futhark library (multiple files, relative -- to the same search path), including all imports. -readLibrary' :: (MonadError CompilerError m, MonadIO m) => - Basis -> [FilePath] - -> m (E.Warnings, - Imports, - VNameSource) +readLibrary' :: + (MonadError CompilerError m, MonadIO m) => + Basis -> + [FilePath] -> + m + ( E.Warnings, + Imports, + VNameSource + ) readLibrary' basis fps = runCompilerM basis $ mapM onFile fps - where onFile fp = do - r <- liftIO $ readFileSafely fp - case r of - Just (Right (_, fs)) -> - handleFile [] (mkInitialImport fp_name) fs fp - Just (Left e) -> externalErrorS e - Nothing -> externalErrorS $ fp ++ ": file not found." - where (fp_name, _) = Posix.splitExtension fp + where + onFile fp = do + r <- liftIO $ readFileSafely fp + case r of + Just (Right (_, fs)) -> + handleFile [] (mkInitialImport fp_name) fs fp + Just (Left e) -> externalErrorS e + Nothing -> externalErrorS $ fp ++ ": file not found." + where + (fp_name, _) = Posix.splitExtension fp -- | Read and type-check Futhark imports (no @.fut@ extension; may -- refer to baked-in prelude). This is an exotic operation that -- probably only makes sense in an interactive environment. -readImports :: (MonadError CompilerError m, MonadIO m) => - Basis -> [ImportName] - -> m (E.Warnings, - Imports, - VNameSource) +readImports :: + (MonadError CompilerError m, MonadIO m) => + Basis -> + [ImportName] -> + m + ( E.Warnings, + Imports, + VNameSource + ) readImports basis imps = runCompilerM basis $ mapM (readImport []) imps -runCompilerM :: Monad m => - Basis -> CompilerM m a - -> m (E.Warnings, [(String, FileModule)], VNameSource) +runCompilerM :: + Monad m => + Basis -> + CompilerM m a -> + m (E.Warnings, [(String, FileModule)], VNameSource) runCompilerM (Basis imports src roots) m = do let s = ReaderState (reverse imports) src mempty s' <- execStateT (runReaderT m roots) s - return (warnings s', - reverse $ alreadyImported s', - nameSource s') + return + ( warnings s', + reverse $ alreadyImported s', + nameSource s' + ) prependRoots :: [FilePath] -> E.UncheckedProg -> E.UncheckedProg prependRoots roots (E.Prog doc ds) = E.Prog doc $ map mkImport roots ++ ds - where mkImport fp = - -- We do not use ImportDec here, because we do not want the - -- type checker to issue a warning about a redundant import. - E.LocalDec (E.OpenDec (E.ModImport fp E.NoInfo mempty) mempty) mempty + where + mkImport fp = + -- We do not use ImportDec here, because we do not want the + -- type checker to issue a warning about a redundant import. + E.LocalDec (E.OpenDec (E.ModImport fp E.NoInfo mempty) mempty) mempty diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -1,7 +1,8 @@ {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE Safe #-} +{-# LANGUAGE TypeFamilies #-} + -- | = Constructing Futhark ASTs -- -- This module re-exports and defines a bunch of building blocks for @@ -52,82 +53,79 @@ -- (relatively) simple example of how to use these components. As are -- some of the high-level building blocks in this very module. module Futhark.Construct - ( letSubExp - , letSubExps - , letExp - , letTupExp - , letTupExp' - , letInPlace - - , eSubExp - , eIf - , eIf' - , eBinOp - , eCmpOp - , eConvOp - , eNot - , eSignum - , eCopy - , eAssert - , eBody - , eLambda - , eRoundToMultipleOf - , eSliceArray - , eBlank - , eAll - - , eOutOfBounds - , eWriteArray - - , asIntZ, asIntS - - , resultBody - , resultBodyM - , insertStmsM - , mapResult - - , foldBinOp - , binOpLambda - , cmpOpLambda - , sliceDim - , fullSlice - , fullSliceNum - , isFullSlice - , sliceAt - , ifCommon - - , module Futhark.Binder - - -- * Result types - , instantiateShapes - , instantiateShapes' - , removeExistentials + ( letSubExp, + letSubExps, + letExp, + letTupExp, + letTupExp', + letInPlace, + eSubExp, + eIf, + eIf', + eBinOp, + eCmpOp, + eConvOp, + eSignum, + eCopy, + eAssert, + eBody, + eLambda, + eRoundToMultipleOf, + eSliceArray, + eBlank, + eAll, + eOutOfBounds, + eWriteArray, + asIntZ, + asIntS, + resultBody, + resultBodyM, + insertStmsM, + mapResult, + foldBinOp, + binOpLambda, + cmpOpLambda, + sliceDim, + fullSlice, + fullSliceNum, + isFullSlice, + sliceAt, + ifCommon, + module Futhark.Binder, - -- * Convenience - , simpleMkLetNames + -- * Result types + instantiateShapes, + instantiateShapes', + removeExistentials, - , ToExp(..) - , toSubExp + -- * Convenience + simpleMkLetNames, + ToExp (..), + toSubExp, ) where -import qualified Data.Map.Strict as M -import Data.List (sortOn) import Control.Monad.Identity import Control.Monad.State import Control.Monad.Writer - -import Futhark.IR -import Futhark.MonadFreshNames +import Data.List (sortOn) +import qualified Data.Map.Strict as M import Futhark.Binder +import Futhark.IR -letSubExp :: MonadBinder m => - String -> Exp (Lore m) -> m SubExp +letSubExp :: + MonadBinder m => + String -> + Exp (Lore m) -> + m SubExp letSubExp _ (BasicOp (SubExp se)) = return se letSubExp desc e = Var <$> letExp desc e -letExp :: MonadBinder m => - String -> Exp (Lore m) -> m VName +letExp :: + MonadBinder m => + String -> + Exp (Lore m) -> + m VName letExp _ (BasicOp (SubExp (Var v))) = return v letExp desc e = do @@ -136,22 +134,31 @@ letBindNames vs e case vs of [v] -> return v - _ -> error $ "letExp: tuple-typed expression given:\n" ++ pretty e + _ -> error $ "letExp: tuple-typed expression given:\n" ++ pretty e -letInPlace :: MonadBinder m => - String -> VName -> Slice SubExp -> Exp (Lore m) - -> m VName +letInPlace :: + MonadBinder m => + String -> + VName -> + Slice SubExp -> + Exp (Lore m) -> + m VName letInPlace desc src slice e = do tmp <- letSubExp (desc ++ "_tmp") e letExp desc $ BasicOp $ Update src slice tmp -letSubExps :: MonadBinder m => - String -> [Exp (Lore m)] -> m [SubExp] +letSubExps :: + MonadBinder m => + String -> + [Exp (Lore m)] -> + m [SubExp] letSubExps desc = mapM $ letSubExp desc -letTupExp :: (MonadBinder m) => - String -> Exp (Lore m) - -> m [VName] +letTupExp :: + (MonadBinder m) => + String -> + Exp (Lore m) -> + m [VName] letTupExp _ (BasicOp (SubExp (Var v))) = return [v] letTupExp name e = do @@ -160,26 +167,36 @@ letBindNames names e return names -letTupExp' :: (MonadBinder m) => - String -> Exp (Lore m) - -> m [SubExp] +letTupExp' :: + (MonadBinder m) => + String -> + Exp (Lore m) -> + m [SubExp] letTupExp' _ (BasicOp (SubExp se)) = return [se] letTupExp' name ses = map Var <$> letTupExp name ses -eSubExp :: MonadBinder m => - SubExp -> m (Exp (Lore m)) +eSubExp :: + MonadBinder m => + SubExp -> + m (Exp (Lore m)) eSubExp = pure . BasicOp . SubExp -eIf :: (MonadBinder m, BranchType (Lore m) ~ ExtType) => - m (Exp (Lore m)) -> m (Body (Lore m)) -> m (Body (Lore m)) - -> m (Exp (Lore m)) +eIf :: + (MonadBinder m, BranchType (Lore m) ~ ExtType) => + m (Exp (Lore m)) -> + m (Body (Lore m)) -> + m (Body (Lore m)) -> + m (Exp (Lore m)) eIf ce te fe = eIf' ce te fe IfNormal -- | As 'eIf', but an 'IfSort' can be given. -eIf' :: (MonadBinder m, BranchType (Lore m) ~ ExtType) => - m (Exp (Lore m)) -> m (Body (Lore m)) -> m (Body (Lore m)) - -> IfSort - -> m (Exp (Lore m)) +eIf' :: + (MonadBinder m, BranchType (Lore m) ~ ExtType) => + m (Exp (Lore m)) -> + m (Body (Lore m)) -> + m (Body (Lore m)) -> + IfSort -> + m (Exp (Lore m)) eIf' ce te fe if_sort = do ce' <- letSubExp "cond" =<< ce te' <- insertStmsM te @@ -189,50 +206,61 @@ te'' <- addContextForBranch ts te' fe'' <- addContextForBranch ts fe' return $ If ce' te'' fe'' $ IfDec ts if_sort - where addContextForBranch ts (Body _ stms val_res) = do - body_ts <- extendedScope (traverse subExpType val_res) stmsscope - let ctx_res = map snd $ sortOn fst $ - M.toList $ shapeExtMapping ts body_ts - mkBodyM stms $ ctx_res++val_res - where stmsscope = scopeOf stms + where + addContextForBranch ts (Body _ stms val_res) = do + body_ts <- extendedScope (traverse subExpType val_res) stmsscope + let ctx_res = + map snd $ + sortOn fst $ + M.toList $ shapeExtMapping ts body_ts + mkBodyM stms $ ctx_res ++ val_res + where + stmsscope = scopeOf stms -- The type of a body. Watch out: this only works for the degenerate -- case where the body does not already return its context. bodyExtType :: (HasScope lore m, Monad m) => Body lore -> m [ExtType] bodyExtType (Body _ stms res) = - existentialiseExtTypes (M.keys stmsscope) . staticShapes <$> - extendedScope (traverse subExpType res) stmsscope - where stmsscope = scopeOf stms + existentialiseExtTypes (M.keys stmsscope) . staticShapes + <$> extendedScope (traverse subExpType res) stmsscope + where + stmsscope = scopeOf stms -eBinOp :: MonadBinder m => - BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) - -> m (Exp (Lore m)) +eBinOp :: + MonadBinder m => + BinOp -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) eBinOp op x y = do x' <- letSubExp "x" =<< x y' <- letSubExp "y" =<< y return $ BasicOp $ BinOp op x' y' -eCmpOp :: MonadBinder m => - CmpOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) - -> m (Exp (Lore m)) +eCmpOp :: + MonadBinder m => + CmpOp -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) eCmpOp op x y = do x' <- letSubExp "x" =<< x y' <- letSubExp "y" =<< y return $ BasicOp $ CmpOp op x' y' -eConvOp :: MonadBinder m => - ConvOp -> m (Exp (Lore m)) - -> m (Exp (Lore m)) +eConvOp :: + MonadBinder m => + ConvOp -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) eConvOp op x = do x' <- letSubExp "x" =<< x return $ BasicOp $ ConvOp op x' -eNot :: MonadBinder m => - m (Exp (Lore m)) -> m (Exp (Lore m)) -eNot e = BasicOp . UnOp Not <$> (letSubExp "not_arg" =<< e) - -eSignum :: MonadBinder m => - m (Exp (Lore m)) -> m (Exp (Lore m)) +eSignum :: + MonadBinder m => + m (Exp (Lore m)) -> + m (Exp (Lore m)) eSignum em = do e <- em e' <- letSubExp "signum_arg" e @@ -243,70 +271,101 @@ _ -> error $ "eSignum: operand " ++ pretty e ++ " has invalid type." -eCopy :: MonadBinder m => - m (Exp (Lore m)) -> m (Exp (Lore m)) +eCopy :: + MonadBinder m => + m (Exp (Lore m)) -> + m (Exp (Lore m)) eCopy e = BasicOp . Copy <$> (letExp "copy_arg" =<< e) -eAssert :: MonadBinder m => - m (Exp (Lore m)) -> ErrorMsg SubExp -> SrcLoc -> m (Exp (Lore m)) -eAssert e msg loc = do e' <- letSubExp "assert_arg" =<< e - return $ BasicOp $ Assert e' msg (loc, mempty) +eAssert :: + MonadBinder m => + m (Exp (Lore m)) -> + ErrorMsg SubExp -> + SrcLoc -> + m (Exp (Lore m)) +eAssert e msg loc = do + e' <- letSubExp "assert_arg" =<< e + return $ BasicOp $ Assert e' msg (loc, mempty) -eBody :: (MonadBinder m) => - [m (Exp (Lore m))] - -> m (Body (Lore m)) +eBody :: + (MonadBinder m) => + [m (Exp (Lore m))] -> + m (Body (Lore m)) eBody es = insertStmsM $ do - es' <- sequence es - xs <- mapM (letTupExp "x") es' - mkBodyM mempty $ map Var $ concat xs + es' <- sequence es + xs <- mapM (letTupExp "x") es' + mkBodyM mempty $ map Var $ concat xs -eLambda :: MonadBinder m => - Lambda (Lore m) -> [m (Exp (Lore m))] -> m [SubExp] -eLambda lam args = do zipWithM_ bindParam (lambdaParams lam) args - bodyBind $ lambdaBody lam - where bindParam param arg = letBindNames [paramName param] =<< arg +eLambda :: + MonadBinder m => + Lambda (Lore m) -> + [m (Exp (Lore m))] -> + m [SubExp] +eLambda lam args = do + zipWithM_ bindParam (lambdaParams lam) args + bodyBind $ lambdaBody lam + where + bindParam param arg = letBindNames [paramName param] =<< arg -eRoundToMultipleOf :: MonadBinder m => - IntType -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m)) +eRoundToMultipleOf :: + MonadBinder m => + IntType -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) eRoundToMultipleOf t x d = ePlus x (eMod (eMinus d (eMod x d)) d) - where eMod = eBinOp (SMod t Unsafe) - eMinus = eBinOp (Sub t OverflowWrap) - ePlus = eBinOp (Add t OverflowWrap) + where + eMod = eBinOp (SMod t Unsafe) + eMinus = eBinOp (Sub t OverflowWrap) + ePlus = eBinOp (Add t OverflowWrap) -- | Construct an 'Index' expressions that slices an array with unit stride. -eSliceArray :: MonadBinder m => - Int -> VName -> m (Exp (Lore m)) -> m (Exp (Lore m)) - -> m (Exp (Lore m)) +eSliceArray :: + MonadBinder m => + Int -> + VName -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) eSliceArray d arr i n = do arr_t <- lookupType arr - let skips = map (slice (constant (0::Int32))) $ take d $ arrayDims arr_t + let skips = map (slice (constant (0 :: Int64))) $ take d $ arrayDims arr_t i' <- letSubExp "slice_i" =<< i n' <- letSubExp "slice_n" =<< n return $ BasicOp $ Index arr $ fullSlice arr_t $ skips ++ [slice i' n'] - where slice j m = DimSlice j m (constant (1::Int32)) + where + slice j m = DimSlice j m (constant (1 :: Int64)) -- | Are these indexes out-of-bounds for the array? -eOutOfBounds :: MonadBinder m => - VName -> [m (Exp (Lore m))] -> m (Exp (Lore m)) +eOutOfBounds :: + MonadBinder m => + VName -> + [m (Exp (Lore m))] -> + m (Exp (Lore m)) eOutOfBounds arr is = do arr_t <- lookupType arr let ws = arrayDims arr_t is' <- mapM (letSubExp "write_i") =<< sequence is let checkDim w i = do - less_than_zero <- letSubExp "less_than_zero" $ - BasicOp $ CmpOp (CmpSlt Int32) i (constant (0::Int32)) - greater_than_size <- letSubExp "greater_than_size" $ - BasicOp $ CmpOp (CmpSle Int32) w i + less_than_zero <- + letSubExp "less_than_zero" $ + BasicOp $ CmpOp (CmpSlt Int64) i (constant (0 :: Int64)) + greater_than_size <- + letSubExp "greater_than_size" $ + BasicOp $ CmpOp (CmpSle Int64) w i letSubExp "outside_bounds_dim" $ BasicOp $ BinOp LogOr less_than_zero greater_than_size foldBinOp LogOr (constant False) =<< zipWithM checkDim ws is' -- | Write to an index of the array, if within bounds. Otherwise, -- nothing. Produces the updated array. -eWriteArray :: (MonadBinder m, BranchType (Lore m) ~ ExtType) => - VName -> [m (Exp (Lore m))] -> m (Exp (Lore m)) - -> m (Exp (Lore m)) +eWriteArray :: + (MonadBinder m, BranchType (Lore m) ~ ExtType) => + VName -> + [m (Exp (Lore m))] -> + m (Exp (Lore m)) -> + m (Exp (Lore m)) eWriteArray arr is v = do arr_t <- lookupType arr is' <- mapM (letSubExp "write_i") =<< sequence is @@ -317,19 +376,23 @@ outside_bounds_branch <- insertStmsM $ resultBodyM [Var arr] in_bounds_branch <- insertStmsM $ do - res <- letInPlace "write_out_inside_bounds" arr - (fullSlice arr_t (map DimFix is')) $ BasicOp $ SubExp v' + res <- + letInPlace + "write_out_inside_bounds" + arr + (fullSlice arr_t (map DimFix is')) + $ BasicOp $ SubExp v' resultBodyM [Var res] return $ If outside_bounds outside_bounds_branch in_bounds_branch $ - ifCommon [arr_t] + ifCommon [arr_t] -- | Construct an unspecified value of the given type. eBlank :: MonadBinder m => Type -> m (Exp (Lore m)) eBlank (Prim t) = return $ BasicOp $ SubExp $ Constant $ blankPrimValue t eBlank (Array t shape _) = return $ BasicOp $ Scratch t $ shapeDims shape -eBlank Mem{} = error "eBlank: cannot create blank memory" +eBlank Mem {} = error "eBlank: cannot create blank memory" -- | Sign-extend to the given integer type. asIntS :: MonadBinder m => IntType -> SubExp -> m SubExp @@ -339,8 +402,12 @@ asIntZ :: MonadBinder m => IntType -> SubExp -> m SubExp asIntZ = asInt ZExt -asInt :: MonadBinder m => - (IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp +asInt :: + MonadBinder m => + (IntType -> IntType -> ConvOp) -> + IntType -> + SubExp -> + m SubExp asInt ext to_it e = do e_t <- subExpType e case e_t of @@ -348,55 +415,71 @@ | to_it == from_it -> return e | otherwise -> letSubExp s $ BasicOp $ ConvOp (ext from_it to_it) e _ -> error "asInt: wrong type" - where s = case e of Var v -> baseString v - _ -> "to_" ++ pretty to_it - + where + s = case e of + Var v -> baseString v + _ -> "to_" ++ pretty to_it -- | Apply a binary operator to several subexpressions. A left-fold. -foldBinOp :: MonadBinder m => - BinOp -> SubExp -> [SubExp] -> m (Exp (Lore m)) +foldBinOp :: + MonadBinder m => + BinOp -> + SubExp -> + [SubExp] -> + m (Exp (Lore m)) foldBinOp _ ne [] = return $ BasicOp $ SubExp ne -foldBinOp bop ne (e:es) = +foldBinOp bop ne (e : es) = eBinOp bop (pure $ BasicOp $ SubExp e) (foldBinOp bop ne es) -- | True if all operands are true. eAll :: MonadBinder m => [SubExp] -> m (Exp (Lore m)) eAll [] = pure $ BasicOp $ SubExp $ constant True -eAll (x:xs) = foldBinOp LogAnd x xs +eAll (x : xs) = foldBinOp LogAnd x xs -- | Create a two-parameter lambda whose body applies the given binary -- operation to its arguments. It is assumed that both argument and -- result types are the same. (This assumption should be fixed at -- some point.) -binOpLambda :: (MonadBinder m, Bindable (Lore m)) => - BinOp -> PrimType -> m (Lambda (Lore m)) +binOpLambda :: + (MonadBinder m, Bindable (Lore m)) => + BinOp -> + PrimType -> + m (Lambda (Lore m)) binOpLambda bop t = binLambda (BinOp bop) t t -- | As 'binOpLambda', but for t'CmpOp's. -cmpOpLambda :: (MonadBinder m, Bindable (Lore m)) => - CmpOp -> m (Lambda (Lore m)) +cmpOpLambda :: + (MonadBinder m, Bindable (Lore m)) => + CmpOp -> + m (Lambda (Lore m)) cmpOpLambda cop = binLambda (CmpOp cop) (cmpOpType cop) Bool -binLambda :: (MonadBinder m, Bindable (Lore m)) => - (SubExp -> SubExp -> BasicOp) -> PrimType -> PrimType - -> m (Lambda (Lore m)) +binLambda :: + (MonadBinder m, Bindable (Lore m)) => + (SubExp -> SubExp -> BasicOp) -> + PrimType -> + PrimType -> + m (Lambda (Lore m)) binLambda bop arg_t ret_t = do - x <- newVName "x" - y <- newVName "y" + x <- newVName "x" + y <- newVName "y" body <- insertStmsM $ do res <- letSubExp "binlam_res" $ BasicOp $ bop (Var x) (Var y) return $ resultBody [res] - return Lambda { - lambdaParams = [Param x (Prim arg_t), - Param y (Prim arg_t)] - , lambdaReturnType = [Prim ret_t] - , lambdaBody = body - } + return + Lambda + { lambdaParams = + [ Param x (Prim arg_t), + Param y (Prim arg_t) + ], + lambdaReturnType = [Prim ret_t], + lambdaBody = body + } -- | Slice a full dimension of the given size. sliceDim :: SubExp -> DimIndex SubExp -sliceDim d = DimSlice (constant (0::Int32)) d (constant (1::Int32)) +sliceDim d = DimSlice (constant (0 :: Int64)) d (constant (1 :: Int64)) -- | @fullSlice t slice@ returns @slice@, but with 'DimSlice's of -- entire dimensions appended to the full dimensionality of @t@. This @@ -423,9 +506,10 @@ -- dimension, but also one that fixes all unit dimensions. isFullSlice :: Shape -> Slice SubExp -> Bool isFullSlice shape slice = and $ zipWith allOfIt (shapeDims shape) slice - where allOfIt (Constant v) DimFix{} = oneIsh v - allOfIt d (DimSlice _ n _) = d == n - allOfIt _ _ = False + where + allOfIt (Constant v) DimFix {} = oneIsh v + allOfIt d (DimSlice _ n _) = d == n + allOfIt _ _ = False ifCommon :: [Type] -> IfDec ExtType ifCommon ts = IfDec (staticShapes ts) IfNormal @@ -436,75 +520,96 @@ -- | Conveniently construct a body that contains no bindings - but -- this time, monadically! -resultBodyM :: MonadBinder m => - [SubExp] - -> m (Body (Lore m)) +resultBodyM :: + MonadBinder m => + [SubExp] -> + m (Body (Lore m)) resultBodyM = mkBodyM mempty -- | Evaluate the action, producing a body, then wrap it in all the -- bindings it created using 'addStm'. -insertStmsM :: (MonadBinder m) => - m (Body (Lore m)) -> m (Body (Lore m)) +insertStmsM :: + (MonadBinder m) => + m (Body (Lore m)) -> + m (Body (Lore m)) insertStmsM m = do (Body _ bnds res, otherbnds) <- collectStms m mkBodyM (otherbnds <> bnds) res -- | Change that result where evaluation of the body would stop. Also -- change type annotations at branches. -mapResult :: Bindable lore => - (Result -> Body lore) -> Body lore -> Body lore +mapResult :: + Bindable lore => + (Result -> Body lore) -> + Body lore -> + Body lore mapResult f (Body _ bnds res) = let Body _ bnds2 newres = f res - in mkBody (bnds<>bnds2) newres + in mkBody (bnds <> bnds2) newres -- | Instantiate all existential parts dimensions of the given -- type, using a monadic action to create the necessary t'SubExp's. -- You should call this function within some monad that allows you to -- collect the actions performed (say, 'Writer'). -instantiateShapes :: Monad m => - (Int -> m SubExp) - -> [TypeBase ExtShape u] - -> m [TypeBase Shape u] +instantiateShapes :: + Monad m => + (Int -> m SubExp) -> + [TypeBase ExtShape u] -> + m [TypeBase Shape u] instantiateShapes f ts = evalStateT (mapM instantiate ts) M.empty - where instantiate t = do - shape <- mapM instantiate' $ shapeDims $ arrayShape t - return $ t `setArrayShape` Shape shape - instantiate' (Ext x) = do - m <- get - case M.lookup x m of - Just se -> return se - Nothing -> do se <- lift $ f x - put $ M.insert x se m - return se - instantiate' (Free se) = return se + where + instantiate t = do + shape <- mapM instantiate' $ shapeDims $ arrayShape t + return $ t `setArrayShape` Shape shape + instantiate' (Ext x) = do + m <- get + case M.lookup x m of + Just se -> return se + Nothing -> do + se <- lift $ f x + put $ M.insert x se m + return se + instantiate' (Free se) = return se -instantiateShapes' :: MonadFreshNames m => - [TypeBase ExtShape u] - -> m ([TypeBase Shape u], [Ident]) +instantiateShapes' :: + MonadFreshNames m => + [TypeBase ExtShape u] -> + m ([TypeBase Shape u], [Ident]) instantiateShapes' ts = runWriterT $ instantiateShapes instantiate ts - where instantiate _ = do v <- lift $ newIdent "size" $ Prim int32 - tell [v] - return $ Var $ identName v + where + instantiate _ = do + v <- lift $ newIdent "size" $ Prim int64 + tell [v] + return $ Var $ identName v removeExistentials :: ExtType -> Type -> Type removeExistentials t1 t2 = - t1 `setArrayDims` - zipWith nonExistential - (shapeDims $ arrayShape t1) - (arrayDims t2) - where nonExistential (Ext _) dim = dim - nonExistential (Free dim) _ = dim + t1 + `setArrayDims` zipWith + nonExistential + (shapeDims $ arrayShape t1) + (arrayDims t2) + where + nonExistential (Ext _) dim = dim + nonExistential (Free dim) _ = dim -- | Can be used as the definition of 'mkLetNames' for a 'Bindable' -- instance for simple representations. -simpleMkLetNames :: (ExpDec lore ~ (), LetDec lore ~ Type, - MonadFreshNames m, TypedOp (Op lore), HasScope lore m) => - [VName] -> Exp lore -> m (Stm lore) +simpleMkLetNames :: + ( ExpDec lore ~ (), + LetDec lore ~ Type, + MonadFreshNames m, + TypedOp (Op lore), + HasScope lore m + ) => + [VName] -> + Exp lore -> + m (Stm lore) simpleMkLetNames names e = do et <- expExtType e (ts, shapes) <- instantiateShapes' et - let shapeElems = [ PatElem shape shapet | Ident shape shapet <- shapes ] + let shapeElems = [PatElem shape shapet | Ident shape shapet <- shapes] let valElems = zipWith PatElem names ts return $ Let (Pattern shapeElems valElems) (defAux ()) e diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs --- a/src/Futhark/Doc/Generator.hs +++ b/src/Futhark/Doc/Generator.hs @@ -1,33 +1,32 @@ {-# LANGUAGE OverloadedStrings #-} + -- | The core logic of @futhark doc@. module Futhark.Doc.Generator (renderFiles) where +import qualified CMarkGFM as GFM import Control.Arrow ((***)) import Control.Monad import Control.Monad.Reader import Control.Monad.Writer hiding (Sum) -import Data.List (sort, sortOn, intersperse, inits, tails, isPrefixOf, find, groupBy, partition) -import Data.Char (isSpace, isAlpha, toUpper) +import Data.Char (isAlpha, isSpace, toUpper) +import Data.List (find, groupBy, inits, intersperse, isPrefixOf, partition, sort, sortOn, tails) +import qualified Data.Map as M import Data.Maybe import Data.Ord -import qualified Data.Map as M import qualified Data.Set as S -import System.FilePath (splitPath, (), (-<.>), (<.>), makeRelative) -import Text.Blaze.Html5 (AttributeValue, Html, (!), toHtml) -import qualified Text.Blaze.Html5 as H -import qualified Text.Blaze.Html5.Attributes as A import Data.String (fromString) -import Data.Version import qualified Data.Text as T -import qualified CMarkGFM as GFM - -import Prelude hiding (abs) - -import Language.Futhark.Semantic -import Language.Futhark.TypeChecker.Monad hiding (warn) -import Language.Futhark +import Data.Version import Futhark.Util.Pretty (Doc, ppr) import Futhark.Version +import Language.Futhark +import Language.Futhark.Semantic +import Language.Futhark.TypeChecker.Monad hiding (warn) +import System.FilePath (makeRelative, splitPath, (-<.>), (<.>), ()) +import Text.Blaze.Html5 (AttributeValue, Html, toHtml, (!)) +import qualified Text.Blaze.Html5 as H +import qualified Text.Blaze.Html5.Attributes as A +import Prelude hiding (abs) docToHtml :: Doc -> Html docToHtml = toHtml . pretty @@ -44,7 +43,7 @@ joinBy :: Html -> [Html] -> Html joinBy _ [] = mempty joinBy _ [x] = x -joinBy sep (x:xs) = x <> foldMap (sep <>) xs +joinBy sep (x : xs) = x <> foldMap (sep <>) xs commas :: [Html] -> Html commas = joinBy ", " @@ -54,6 +53,7 @@ braces :: Html -> Html braces x = "{" <> x <> "}" + brackets :: Html -> Html brackets x = "[" <> x <> "]" @@ -64,17 +64,20 @@ -- are uninteresting. These are for example type parameters. type NoLink = S.Set VName -data Context = Context { ctxCurrent :: String - , ctxFileMod :: FileModule - , ctxImports :: Imports - , ctxNoLink :: NoLink - , ctxFileMap :: FileMap - , ctxVisibleMTys :: S.Set VName - -- ^ Local module types that show up in the - -- interface. These should be documented, - -- but clearly marked local. - } +data Context = Context + { ctxCurrent :: String, + ctxFileMod :: FileModule, + ctxImports :: Imports, + ctxNoLink :: NoLink, + ctxFileMap :: FileMap, + -- | Local module types that show up in the + -- interface. These should be documented, + -- but clearly marked local. + ctxVisibleMTys :: S.Set VName + } + type FileMap = M.Map VName (String, Namespace) + type DocM = ReaderT Context (WriterT Documented (Writer Warnings)) data IndexWhat = IndexValue | IxFun | IndexModule | IndexModuleType | IndexType @@ -91,7 +94,7 @@ noLink :: [VName] -> DocM a -> DocM a noLink names = local $ \ctx -> - ctx { ctxNoLink = S.fromList names <> ctxNoLink ctx } + ctx {ctxNoLink = S.fromList names <> ctxNoLink ctx} selfLink :: AttributeValue -> Html -> Html selfLink s = H.a ! A.id s ! A.href ("#" <> s) ! A.class_ "self_link" @@ -103,24 +106,28 @@ emptyRow = H.tr $ H.td mempty <> H.td mempty <> H.td mempty specRow :: Html -> Html -> Html -> Html -specRow a b c = H.tr $ (H.td ! A.class_ "spec_lhs") a <> - (H.td ! A.class_ "spec_eql") b <> - (H.td ! A.class_ "spec_rhs") c +specRow a b c = + H.tr $ + (H.td ! A.class_ "spec_lhs") a + <> (H.td ! A.class_ "spec_eql") b + <> (H.td ! A.class_ "spec_rhs") c vnameToFileMap :: Imports -> FileMap vnameToFileMap = mconcat . map forFile - where forFile (file, FileModule abs file_env _prog) = - mconcat (map (vname Type) (M.keys abs)) <> - forEnv file_env - where vname ns v = M.singleton (qualLeaf v) (file, ns) - vname' ((ns, _), v) = vname ns v + where + forFile (file, FileModule abs file_env _prog) = + mconcat (map (vname Type) (M.keys abs)) + <> forEnv file_env + where + vname ns v = M.singleton (qualLeaf v) (file, ns) + vname' ((ns, _), v) = vname ns v - forEnv env = - mconcat (map vname' $ M.toList $ envNameMap env) <> - mconcat (map forMty $ M.elems $ envSigTable env) - forMod (ModEnv env) = forEnv env - forMod ModFun{} = mempty - forMty = forMod . mtyMod + forEnv env = + mconcat (map vname' $ M.toList $ envNameMap env) + <> mconcat (map forMty $ M.elems $ envSigTable env) + forMod (ModEnv env) = forEnv env + forMod ModFun {} = mempty + forMty = forMod . mtyMod -- | @renderFiles important_imports imports@ produces HTML files -- documenting the type-checked program @imports@, with the files in @@ -129,32 +136,46 @@ -- or the relative links will be wrong. renderFiles :: [FilePath] -> Imports -> ([(FilePath, Html)], Warnings) renderFiles important_imports imports = runWriter $ do - (import_pages, documented) <- runWriterT $ forM imports $ \(current, fm) -> - let ctx = Context current fm imports mempty file_map - (progModuleTypes $ fileProg fm) in - flip runReaderT ctx $ do - - (first_paragraph, maybe_abstract, maybe_sections) <- headerDoc $ fileProg fm + (import_pages, documented) <- runWriterT $ + forM imports $ \(current, fm) -> + let ctx = + Context + current + fm + imports + mempty + file_map + (progModuleTypes $ fileProg fm) + in flip runReaderT ctx $ do + (first_paragraph, maybe_abstract, maybe_sections) <- headerDoc $ fileProg fm - synopsis <- (H.div ! A.id "module") <$> synopsisDecs (progDecs $ fileProg fm) + synopsis <- (H.div ! A.id "module") <$> synopsisDecs (progDecs $ fileProg fm) - description <- describeDecs $ progDecs $ fileProg fm + description <- describeDecs $ progDecs $ fileProg fm - return (current, - (H.docTypeHtml ! A.lang "en" $ - addBoilerplateWithNav important_imports imports ("doc" current) current $ - H.main $ maybe_abstract <> - selfLink "synopsis" (H.h2 "Synopsis") <> (H.div ! A.id "overview") synopsis <> - selfLink "description" (H.h2 "Description") <> description <> - maybe_sections, - first_paragraph)) + return + ( current, + ( H.docTypeHtml ! A.lang "en" $ + addBoilerplateWithNav important_imports imports ("doc" current) current $ + H.main $ + maybe_abstract + <> selfLink "synopsis" (H.h2 "Synopsis") + <> (H.div ! A.id "overview") synopsis + <> selfLink "description" (H.h2 "Description") + <> description + <> maybe_sections, + first_paragraph + ) + ) return $ - [("index.html", contentsPage important_imports $ map (fmap snd) import_pages), - ("doc-index.html", indexPage important_imports imports documented file_map)] - ++ map (importHtml *** fst) import_pages - where file_map = vnameToFileMap imports - importHtml import_name = "doc" import_name <.> "html" + [ ("index.html", contentsPage important_imports $ map (fmap snd) import_pages), + ("doc-index.html", indexPage important_imports imports documented file_map) + ] + ++ map (importHtml *** fst) import_pages + where + file_map = vnameToFileMap imports + importHtml import_name = "doc" import_name <.> "html" -- | The header documentation (which need not be present) can contain -- an abstract and further sections. @@ -166,135 +187,163 @@ first_paragraph <- docHtml $ Just $ DocComment (firstParagraph abstract) loc abstract' <- docHtml $ Just $ DocComment abstract loc more_sections' <- docHtml $ Just $ DocComment more_sections loc - return (first_paragraph, - selfLink "abstract" (H.h2 "Abstract") <> abstract', - more_sections') + return + ( first_paragraph, + selfLink "abstract" (H.h2 "Abstract") <> abstract', + more_sections' + ) _ -> return mempty - where splitHeaderDoc s = fromMaybe (s, mempty) $ - find (("\n##" `isPrefixOf`) . snd) $ - zip (inits s) (tails s) - firstParagraph = unlines . takeWhile (not . paragraphSeparator) . lines - paragraphSeparator = all isSpace - + where + splitHeaderDoc s = + fromMaybe (s, mempty) $ + find (("\n##" `isPrefixOf`) . snd) $ + zip (inits s) (tails s) + firstParagraph = unlines . takeWhile (not . paragraphSeparator) . lines + paragraphSeparator = all isSpace contentsPage :: [FilePath] -> [(String, Html)] -> Html contentsPage important_imports pages = - H.docTypeHtml $ addBoilerplate "index.html" "Futhark Library Documentation" $ - H.main $ H.h2 "Main libraries" <> - fileList important_pages <> - if null unimportant_pages then mempty else - H.h2 "Supporting libraries" <> - fileList unimportant_pages - where (important_pages, unimportant_pages) = - partition ((`elem` important_imports) . fst) pages + H.docTypeHtml $ + addBoilerplate "index.html" "Futhark Library Documentation" $ + H.main $ + H.h2 "Main libraries" + <> fileList important_pages + <> if null unimportant_pages + then mempty + else + H.h2 "Supporting libraries" + <> fileList unimportant_pages + where + (important_pages, unimportant_pages) = + partition ((`elem` important_imports) . fst) pages - fileList pages' = - H.dl ! A.class_ "file_list" $ - mconcat $ map linkTo $ sortOn fst pages' + fileList pages' = + H.dl ! A.class_ "file_list" $ + mconcat $ map linkTo $ sortOn fst pages' - linkTo (name, maybe_abstract) = - H.div ! A.class_ "file_desc" $ - (H.dt ! A.class_ "desc_header") (importLink "index.html" name) <> - (H.dd ! A.class_ "desc_doc") maybe_abstract + linkTo (name, maybe_abstract) = + H.div ! A.class_ "file_desc" $ + (H.dt ! A.class_ "desc_header") (importLink "index.html" name) + <> (H.dd ! A.class_ "desc_doc") maybe_abstract importLink :: FilePath -> String -> Html importLink current name = let file = relativise (makeRelative "/" $ "doc" name -<.> "html") current - in (H.a ! A.href (fromString file) $ fromString name) + in (H.a ! A.href (fromString file) $ fromString name) indexPage :: [FilePath] -> Imports -> Documented -> FileMap -> Html indexPage important_imports imports documented fm = - H.docTypeHtml $ addBoilerplateWithNav important_imports imports "doc-index.html" "Index" $ - H.main $ - (H.ul ! A.id "doc_index_list" $ - mconcat $ map initialListEntry $ - letter_group_links ++ [symbol_group_link]) <> - (H.table ! A.id "doc_index" $ - H.thead (H.tr $ H.td "Who" <> H.td "What" <> H.td "Where") <> - mconcat (letter_groups ++ [symbol_group])) - where (letter_names, sym_names) = - partition (isLetterName . baseString . fst) $ - sortOn (map toUpper . baseString . fst) $ + H.docTypeHtml $ + addBoilerplateWithNav important_imports imports "doc-index.html" "Index" $ + H.main $ + ( H.ul ! A.id "doc_index_list" $ + mconcat $ + map initialListEntry $ + letter_group_links ++ [symbol_group_link] + ) + <> ( H.table ! A.id "doc_index" $ + H.thead (H.tr $ H.td "Who" <> H.td "What" <> H.td "Where") + <> mconcat (letter_groups ++ [symbol_group]) + ) + where + (letter_names, sym_names) = + partition (isLetterName . baseString . fst) $ + sortOn (map toUpper . baseString . fst) $ mapMaybe isDocumented $ M.toList fm - isDocumented (k, (file, _)) = do - what <- M.lookup k documented - Just (k, (file, what)) + isDocumented (k, (file, _)) = do + what <- M.lookup k documented + Just (k, (file, what)) - (letter_groups, letter_group_links) = - unzip $ map tbodyForNames $ groupBy sameInitial letter_names - (symbol_group, symbol_group_link) = - tbodyForInitial "Symbols" sym_names + (letter_groups, letter_group_links) = + unzip $ map tbodyForNames $ groupBy sameInitial letter_names + (symbol_group, symbol_group_link) = + tbodyForInitial "Symbols" sym_names - isLetterName [] = False - isLetterName (c:_) = isAlpha c + isLetterName [] = False + isLetterName (c : _) = isAlpha c - sameInitial (x, _) (y, _) = - case (baseString x, baseString y) of - (x':_, y':_) -> toUpper x' == toUpper y' - _ -> False + sameInitial (x, _) (y, _) = + case (baseString x, baseString y) of + (x' : _, y' : _) -> toUpper x' == toUpper y' + _ -> False - tbodyForNames names@((s,_):_) = - tbodyForInitial (map toUpper $ take 1 $ baseString s) names - tbodyForNames _ = mempty + tbodyForNames names@((s, _) : _) = + tbodyForInitial (map toUpper $ take 1 $ baseString s) names + tbodyForNames _ = mempty - tbodyForInitial initial names = - (H.tbody $ mconcat $ initial' : map linkTo names, - initial) - where initial' = - H.tr $ H.td ! A.colspan "2" ! A.class_ "doc_index_initial" $ - H.a ! A.id (fromString initial) - ! A.href (fromString $ '#' : initial) - $ fromString initial + tbodyForInitial initial names = + ( H.tbody $ mconcat $ initial' : map linkTo names, + initial + ) + where + initial' = + H.tr $ + H.td ! A.colspan "2" ! A.class_ "doc_index_initial" $ + H.a ! A.id (fromString initial) + ! A.href (fromString $ '#' : initial) + $ fromString initial - initialListEntry initial = - H.li $ H.a ! A.href (fromString $ '#' : initial) $ fromString initial + initialListEntry initial = + H.li $ H.a ! A.href (fromString $ '#' : initial) $ fromString initial - linkTo (name, (file, what)) = - let link = (H.a ! A.href (fromString (makeRelative "/" $ "doc" vnameLink' name "" file))) $ - fromString $ baseString name - what' = case what of IndexValue -> "value" - IxFun -> "function" - IndexType -> "type" - IndexModuleType -> "module type" - IndexModule -> "module" - html_file = makeRelative "/" $ "doc" file -<.> "html" - in H.tr $ - (H.td ! A.class_ "doc_index_name" $ link) <> - (H.td ! A.class_ "doc_index_namespace" $ what') <> - (H.td ! A.class_ "doc_index_file" $ - (H.a ! A.href (fromString html_file) $ fromString file)) + linkTo (name, (file, what)) = + let link = + (H.a ! A.href (fromString (makeRelative "/" $ "doc" vnameLink' name "" file))) $ + fromString $ baseString name + what' = case what of + IndexValue -> "value" + IxFun -> "function" + IndexType -> "type" + IndexModuleType -> "module type" + IndexModule -> "module" + html_file = makeRelative "/" $ "doc" file -<.> "html" + in H.tr $ + (H.td ! A.class_ "doc_index_name" $ link) + <> (H.td ! A.class_ "doc_index_namespace" $ what') + <> ( H.td ! A.class_ "doc_index_file" $ + (H.a ! A.href (fromString html_file) $ fromString file) + ) addBoilerplate :: String -> String -> Html -> Html addBoilerplate current titleText content = - let headHtml = H.head $ - H.meta ! A.charset "utf-8" <> - H.title (fromString titleText) <> - H.link ! A.href (fromString $ relativise "style.css" current) - ! A.rel "stylesheet" - ! A.type_ "text/css" + let headHtml = + H.head $ + H.meta + ! A.charset "utf-8" + <> H.title (fromString titleText) + <> H.link + ! A.href (fromString $ relativise "style.css" current) + ! A.rel "stylesheet" + ! A.type_ "text/css" - navigation = H.ul ! A.id "navigation" $ - H.li (H.a ! A.href (fromString $ relativise "index.html" current) $ "Contents") <> - H.li (H.a ! A.href (fromString $ relativise "doc-index.html" current) $ "Index") + navigation = + H.ul ! A.id "navigation" $ + H.li (H.a ! A.href (fromString $ relativise "index.html" current) $ "Contents") + <> H.li (H.a ! A.href (fromString $ relativise "doc-index.html" current) $ "Index") madeByHtml = "Generated by " <> (H.a ! A.href futhark_doc_url) "futhark-doc" - <> " " <> fromString (showVersion version) - in headHtml <> - H.body ((H.div ! A.id "header") (H.h1 (toHtml titleText) <> navigation) <> - (H.div ! A.id "content") content <> - (H.div ! A.id "footer") madeByHtml) - where futhark_doc_url = - "https://futhark.readthedocs.io/en/latest/man/futhark-doc.html" + <> " " + <> fromString (showVersion version) + in headHtml + <> H.body + ( (H.div ! A.id "header") (H.h1 (toHtml titleText) <> navigation) + <> (H.div ! A.id "content") content + <> (H.div ! A.id "footer") madeByHtml + ) + where + futhark_doc_url = + "https://futhark.readthedocs.io/en/latest/man/futhark-doc.html" addBoilerplateWithNav :: [FilePath] -> Imports -> String -> String -> Html -> Html addBoilerplateWithNav important_imports imports current titleText content = addBoilerplate current titleText $ - (H.nav ! A.id "filenav" $ files) <> content - where files = H.ul $ mconcat $ map pp $ sort $ filter visible important_imports - pp name = H.li $ importLink current name - visible = (`elem` map fst imports) + (H.nav ! A.id "filenav" $ files) <> content + where + files = H.ul $ mconcat $ map pp $ sort $ filter visible important_imports + pp name = H.li $ importLink current name + visible = (`elem` map fst imports) synopsisDecs :: [Dec] -> DocM Html synopsisDecs decs = do @@ -302,8 +351,8 @@ fm <- asks ctxFileMod -- We add an empty row to avoid generating invalid HTML in cases -- where all rows are otherwise colspan=2. - (H.table ! A.class_ "specs") . (emptyRow<>) . mconcat <$> - sequence (mapMaybe (synopsisDec visible fm) decs) + (H.table ! A.class_ "specs") . (emptyRow <>) . mconcat + <$> sequence (mapMaybe (synopsisDec visible fm) decs) synopsisDec :: S.Set VName -> FileModule -> Dec -> Maybe (DocM Html) synopsisDec visible fm dec = case dec of @@ -313,21 +362,24 @@ TypeDec t -> synopsisType t OpenDec x _ | Just opened <- synopsisOpened x -> Just $ do - opened' <- opened - return $ fullRow $ keyword "open " <> opened' + opened' <- opened + return $ fullRow $ keyword "open " <> opened' | otherwise -> - Just $ return $ fullRow $ - keyword "open" <> fromString (" <" <> pretty x <> ">") + Just $ + return $ + fullRow $ + keyword "open" <> fromString (" <" <> pretty x <> ">") LocalDec (SigDec s) _ | sigName s `S.member` visible -> - synopsisModType (keyword "local" <> " ") s - LocalDec{} -> Nothing - ImportDec{} -> Nothing + synopsisModType (keyword "local" <> " ") s + LocalDec {} -> Nothing + ImportDec {} -> Nothing synopsisOpened :: ModExp -> Maybe (DocM Html) synopsisOpened (ModVar qn _) = Just $ qualNameHtml qn -synopsisOpened (ModParens me _) = do me' <- synopsisOpened me - Just $ parens <$> me' +synopsisOpened (ModParens me _) = do + me' <- synopsisOpened me + Just $ parens <$> me' synopsisOpened (ModImport _ (Info file) _) = Just $ do current <- asks ctxCurrent let dest = fromString $ relativise file current <> ".html" @@ -345,14 +397,18 @@ valBindHtml :: Html -> ValBind -> DocM (Html, Html, Html) valBindHtml name (ValBind _ _ retdecl (Info (rettype, _)) tparams params _ _ _ _) = do - let tparams' = mconcat $ map ((" "<>) . typeParamHtml) tparams - noLink' = noLink $ map typeParamName tparams ++ - map identName (S.toList $ mconcat $ map patternIdents params) + let tparams' = mconcat $ map ((" " <>) . typeParamHtml) tparams + noLink' = + noLink $ + map typeParamName tparams + ++ map identName (S.toList $ mconcat $ map patternIdents params) rettype' <- noLink' $ maybe (typeHtml rettype) typeExpHtml retdecl params' <- noLink' $ mapM patternHtml params - return (keyword "val " <> (H.span ! A.class_ "decl_name") name, - tparams', - mconcat (intersperse " -> " $ params' ++ [rettype'])) + return + ( keyword "val " <> (H.span ! A.class_ "decl_name") name, + tparams', + mconcat (intersperse " -> " $ params' ++ [rettype']) + ) synopsisModType :: Html -> SigBind -> Maybe (DocM Html) synopsisModType prefix sb = Just $ do @@ -363,16 +419,18 @@ synopsisMod :: FileModule -> ModBind -> Maybe (DocM Html) synopsisMod fm (ModBind name ps sig _ _ _) = - case sig of Nothing -> (proceed <=< envSig) <$> M.lookup name modtable - Just (s,_) -> Just $ proceed =<< synopsisSigExp s - where proceed sig' = do - let name' = vnameSynopsisDef name - ps' <- modParamHtml ps - return $ specRow (keyword "module " <> name') ": " (ps' <> sig') + case sig of + Nothing -> (proceed <=< envSig) <$> M.lookup name modtable + Just (s, _) -> Just $ proceed =<< synopsisSigExp s + where + proceed sig' = do + let name' = vnameSynopsisDef name + ps' <- modParamHtml ps + return $ specRow (keyword "module " <> name') ": " (ps' <> sig') - FileModule _abs Env { envModTable = modtable} _ = fm - envSig (ModEnv e) = renderEnv e - envSig (ModFun (FunSig _ _ (MTy _ m))) = envSig m + FileModule _abs Env {envModTable = modtable} _ = fm + envSig (ModEnv e) = renderEnv e + envSig (ModFun (FunSig _ _ (MTy _ m))) = envSig m synopsisType :: TypeBind -> Maybe (DocM Html) synopsisType tb = Just $ do @@ -413,8 +471,10 @@ let tps' = map typeParamHtml tps t' <- typeHtml t return $ - keyword "val " <> vnameHtml name <> - mconcat (map (" "<>) tps') <> ": " <> t' + keyword "val " <> vnameHtml name + <> mconcat (map (" " <>) tps') + <> ": " + <> t' typeHtml :: StructType -> DocM Html typeHtml t = case t of @@ -425,16 +485,17 @@ Scalar (Prim et) -> return $ primTypeHtml et Scalar (Record fs) | Just ts <- areTupleFields fs -> - parens . commas <$> mapM typeHtml ts + parens . commas <$> mapM typeHtml ts | otherwise -> - braces . commas <$> mapM ppField (M.toList fs) - where ppField (name, tp) = do - tp' <- typeHtml tp - return $ toHtml (nameToString name) <> ": " <> tp' + braces . commas <$> mapM ppField (M.toList fs) + where + ppField (name, tp) = do + tp' <- typeHtml tp + return $ toHtml (nameToString name) <> ": " <> tp' Scalar (TypeVar _ u et targs) -> do targs' <- mapM typeArgHtml targs et' <- typeNameHtml et - return $ prettyU u <> et' <> mconcat (map (" "<>) targs') + return $ prettyU u <> et' <> mconcat (map (" " <>) targs') Scalar (Arrow _ pname t1 t2) -> do t1' <- typeHtml t1 t2' <- typeHtml t2 @@ -444,8 +505,9 @@ Unnamed -> t1' <> " -> " <> t2' Scalar (Sum cs) -> pipes <$> mapM ppClause (sortConstrs cs) - where ppClause (n, ts) = joinBy " " . (ppConstr n :) <$> mapM typeHtml ts - ppConstr name = "#" <> toHtml (nameToString name) + where + ppClause (n, ts) = joinBy " " . (ppConstr n :) <$> mapM typeHtml ts + ppConstr name = "#" <> toHtml (nameToString name) prettyShapeDecl :: ShapeDecl (DimDecl VName) -> DocM Html prettyShapeDecl (ShapeDecl ds) = @@ -459,8 +521,13 @@ modParamHtml [] = return mempty modParamHtml (ModParam pname psig _ _ : mps) = liftM2 f (synopsisSigExp psig) (modParamHtml mps) - where f se params = "(" <> vnameHtml pname <> - ": " <> se <> ") -> " <> params + where + f se params = + "(" <> vnameHtml pname + <> ": " + <> se + <> ") -> " + <> params synopsisSigExp :: SigExpBase Info VName -> DocM Html synopsisSigExp e = case e of @@ -471,16 +538,18 @@ s' <- synopsisSigExp s t' <- typeDeclHtml t v' <- qualNameHtml v - let ps' = mconcat $ map ((" "<>) . typeParamHtml) ps + let ps' = mconcat $ map ((" " <>) . typeParamHtml) ps return $ s' <> keyword " with " <> v' <> ps' <> " = " <> t' SigArrow Nothing e1 e2 _ -> liftM2 f (synopsisSigExp e1) (synopsisSigExp e2) - where f e1' e2' = e1' <> " -> " <> e2' + where + f e1' e2' = e1' <> " -> " <> e2' SigArrow (Just v) e1 e2 _ -> - do let name = vnameHtml v - e1' <- synopsisSigExp e1 - e2' <- noLink [v] $ synopsisSigExp e2 - return $ "(" <> name <> ": " <> e1' <> ") -> " <> e2' + do + let name = vnameHtml v + e1' <- synopsisSigExp e1 + e2' <- noLink [v] $ synopsisSigExp e2 + return $ "(" <> name <> ": " <> e1' <> ") -> " <> e2' keyword :: String -> Html keyword = (H.span ! A.class_ "keyword") . fromString @@ -497,29 +566,35 @@ vnameSynopsisDef :: VName -> Html vnameSynopsisDef (VName name tag) = H.span ! A.id (fromString (show tag ++ "s")) $ - H.a ! A.href (fromString ("#" ++ show tag)) $ renderName name + H.a ! A.href (fromString ("#" ++ show tag)) $ renderName name vnameSynopsisRef :: VName -> Html -vnameSynopsisRef v = H.a ! A.class_ "synopsis_link" - ! A.href (fromString ("#" ++ show (baseTag v) ++ "s")) $ - "↑" +vnameSynopsisRef v = + H.a ! A.class_ "synopsis_link" + ! A.href (fromString ("#" ++ show (baseTag v) ++ "s")) + $ "↑" synopsisSpec :: SpecBase Info VName -> DocM Html synopsisSpec spec = case spec of TypeAbbrSpec tpsig -> fullRow <$> typeBindHtml (vnameSynopsisDef $ typeAlias tpsig) tpsig TypeSpec l name ps _ _ -> - return $ fullRow $ keyword l' <> vnameSynopsisDef name <> mconcat (map ((" "<>) . typeParamHtml) ps) - where l' = case l of Unlifted -> "type " - SizeLifted -> "type~ " - Lifted -> "type^ " + return $ fullRow $ keyword l' <> vnameSynopsisDef name <> mconcat (map ((" " <>) . typeParamHtml) ps) + where + l' = case l of + Unlifted -> "type " + SizeLifted -> "type~ " + Lifted -> "type^ " ValSpec name tparams rettype _ _ -> do let tparams' = map typeParamHtml tparams - rettype' <- noLink (map typeParamName tparams) $ - typeDeclHtml rettype - return $ specRow - (keyword "val " <> vnameSynopsisDef name) - (mconcat (map (" "<>) tparams') <> ": ") rettype' + rettype' <- + noLink (map typeParamName tparams) $ + typeDeclHtml rettype + return $ + specRow + (keyword "val " <> vnameSynopsisDef name) + (mconcat (map (" " <>) tparams') <> ": ") + rettype' ModSpec name sig _ _ -> specRow (keyword "module " <> vnameSynopsisDef name) ": " <$> synopsisSigExp sig IncludeSpec e _ -> fullRow . (keyword "include " <>) <$> synopsisSigExp e @@ -529,24 +604,26 @@ typeExpHtml :: TypeExp VName -> DocM Html typeExpHtml e = case e of - TEUnique t _ -> ("*"<>) <$> typeExpHtml t + TEUnique t _ -> ("*" <>) <$> typeExpHtml t TEArray at d _ -> do at' <- typeExpHtml at d' <- dimExpHtml d return $ brackets d' <> at' TETuple ts _ -> parens . commas <$> mapM typeExpHtml ts TERecord fs _ -> braces . commas <$> mapM ppField fs - where ppField (name, t) = do - t' <- typeExpHtml t - return $ toHtml (nameToString name) <> ": " <> t' - TEVar name _ -> qualNameHtml name + where + ppField (name, t) = do + t' <- typeExpHtml t + return $ toHtml (nameToString name) <> ": " <> t' + TEVar name _ -> qualNameHtml name TEApply t arg _ -> do t' <- typeExpHtml t arg' <- typeArgExpHtml arg return $ t' <> " " <> arg' TEArrow pname t1 t2 _ -> do - t1' <- case t1 of TEArrow{} -> parens <$> typeExpHtml t1 - _ -> typeExpHtml t1 + t1' <- case t1 of + TEArrow {} -> parens <$> typeExpHtml t1 + _ -> typeExpHtml t1 t2' <- typeExpHtml t2 return $ case pname of Just v -> @@ -554,23 +631,26 @@ Nothing -> t1' <> " -> " <> t2' TESum cs _ -> pipes <$> mapM ppClause cs - where ppClause (n, ts) = joinBy " " . (ppConstr n :) <$> mapM typeExpHtml ts - ppConstr name = "#" <> toHtml (nameToString name) + where + ppClause (n, ts) = joinBy " " . (ppConstr n :) <$> mapM typeExpHtml ts + ppConstr name = "#" <> toHtml (nameToString name) qualNameHtml :: QualName VName -> DocM Html qualNameHtml (QualName names vname@(VName name tag)) = if tag <= maxIntrinsicTag - then return $ renderName name - else f <$> ref - where prefix :: Html - prefix = mapM_ ((<> ".") . renderName . baseName) names - f (Just s) = H.a ! A.href (fromString s) $ prefix <> renderName name - f Nothing = prefix <> renderName name + then return $ renderName name + else f <$> ref + where + prefix :: Html + prefix = mapM_ ((<> ".") . renderName . baseName) names + f (Just s) = H.a ! A.href (fromString s) $ prefix <> renderName name + f Nothing = prefix <> renderName name - ref = do boring <- asks $ S.member vname . ctxNoLink - if boring - then return Nothing - else Just <$> vnameLink vname + ref = do + boring <- asks $ S.member vname . ctxNoLink + if boring + then return Nothing + else Just <$> vnameLink vname vnameLink' :: VName -> String -> String -> String vnameLink :: VName -> DocM String @@ -578,7 +658,6 @@ current <- asks ctxCurrent file <- asks $ maybe current fst . M.lookup vname . ctxFileMap return $ vnameLink' vname current file - vnameLink' (VName _ tag) current file = if file == current then "#" ++ show tag @@ -592,8 +671,8 @@ let (pat_param, t) = patternParam pat t' <- typeHtml t return $ case pat_param of - Named v -> parens (vnameHtml v <> ": " <> t') - Unnamed -> t' + Named v -> parens (vnameHtml v <> ": " <> t') + Unnamed -> t' relativise :: FilePath -> FilePath -> FilePath relativise dest src = @@ -621,42 +700,48 @@ typeAbbrevHtml :: Liftedness -> Html -> [TypeParam] -> Html typeAbbrevHtml l name params = - what <> name <> mconcat (map ((" "<>) . typeParamHtml) params) - where what = keyword $ "type" ++ pretty l ++ " " + what <> name <> mconcat (map ((" " <>) . typeParamHtml) params) + where + what = keyword $ "type" ++ pretty l ++ " " docHtml :: Maybe DocComment -> DocM Html docHtml (Just (DocComment doc loc)) = - H.preEscapedText . - GFM.commonmarkToHtml [] [GFM.extAutolink] . - T.pack <$> identifierLinks loc doc + H.preEscapedText + . GFM.commonmarkToHtml [] [GFM.extAutolink] + . T.pack + <$> identifierLinks loc doc docHtml Nothing = return mempty identifierLinks :: SrcLoc -> String -> DocM String identifierLinks _ [] = return [] identifierLinks loc s | Just ((name, namespace, file), s') <- identifierReference s = do - let proceed x = (x<>) <$> identifierLinks loc s' - unknown = proceed $ "`" <> name <> "`" - case knownNamespace namespace of - Just namespace' -> do - maybe_v <- lookupName (namespace', name, file) - case maybe_v of - Nothing -> do - warn loc $ - "Identifier '" <> name <> "' not found in namespace '" <> - namespace <> "'" <> maybe "" (" in file "<>) file <> "." - unknown - Just v' -> do - link <- vnameLink v' - proceed $ "[`" <> name <> "`](" <> link <> ")" - _ -> do - warn loc $ "Unknown namespace '" <> namespace <> "'." - unknown - where knownNamespace "term" = Just Term - knownNamespace "mtype" = Just Signature - knownNamespace "type" = Just Type - knownNamespace _ = Nothing -identifierLinks loc (c:s') = (c:) <$> identifierLinks loc s' + let proceed x = (x <>) <$> identifierLinks loc s' + unknown = proceed $ "`" <> name <> "`" + case knownNamespace namespace of + Just namespace' -> do + maybe_v <- lookupName (namespace', name, file) + case maybe_v of + Nothing -> do + warn loc $ + "Identifier '" <> name <> "' not found in namespace '" + <> namespace + <> "'" + <> maybe "" (" in file " <>) file + <> "." + unknown + Just v' -> do + link <- vnameLink v' + proceed $ "[`" <> name <> "`](" <> link <> ")" + _ -> do + warn loc $ "Unknown namespace '" <> namespace <> "'." + unknown + where + knownNamespace "term" = Just Term + knownNamespace "mtype" = Just Signature + knownNamespace "type" = Just Type + knownNamespace _ = Nothing +identifierLinks loc (c : s') = (c :) <$> identifierLinks loc s' lookupName :: (Namespace, String, Maybe FilePath) -> DocM (Maybe VName) lookupName (namespace, name, file) = do @@ -668,83 +753,84 @@ Just qn -> return $ Just $ qualLeaf qn lookupEnvForFile :: Maybe FilePath -> DocM (Maybe Env) -lookupEnvForFile Nothing = asks $ Just . fileEnv . ctxFileMod +lookupEnvForFile Nothing = asks $ Just . fileEnv . ctxFileMod lookupEnvForFile (Just file) = asks $ fmap fileEnv . lookup file . ctxImports -describeGeneric :: VName - -> IndexWhat - -> Maybe DocComment - -> (Html -> DocM Html) - -> DocM Html +describeGeneric :: + VName -> + IndexWhat -> + Maybe DocComment -> + (Html -> DocM Html) -> + DocM Html describeGeneric name what doc f = do name' <- H.span ! A.class_ "decl_name" <$> vnameDescDef name what decl_type <- f name' doc' <- docHtml doc let decl_doc = H.dd ! A.class_ "desc_doc" $ doc' - decl_header = (H.dt ! A.class_ "desc_header") $ - vnameSynopsisRef name <> decl_type + decl_header = + (H.dt ! A.class_ "desc_header") $ + vnameSynopsisRef name <> decl_type return $ decl_header <> decl_doc -describeGenericMod :: VName - -> IndexWhat - -> SigExp - -> Maybe DocComment - -> (Html -> DocM Html) - -> DocM Html +describeGenericMod :: + VName -> + IndexWhat -> + SigExp -> + Maybe DocComment -> + (Html -> DocM Html) -> + DocM Html describeGenericMod name what se doc f = do name' <- H.span ! A.class_ "decl_name" <$> vnameDescDef name what decl_type <- f name' doc' <- case se of - SigSpecs specs _ -> (<>) <$> docHtml doc <*> describeSpecs specs - _ -> docHtml doc + SigSpecs specs _ -> (<>) <$> docHtml doc <*> describeSpecs specs + _ -> docHtml doc let decl_doc = H.dd ! A.class_ "desc_doc" $ doc' - decl_header = (H.dt ! A.class_ "desc_header") $ - vnameSynopsisRef name <> decl_type + decl_header = + (H.dt ! A.class_ "desc_header") $ + vnameSynopsisRef name <> decl_type return $ decl_header <> decl_doc describeDecs :: [Dec] -> DocM Html describeDecs decs = do visible <- asks ctxVisibleMTys - H.dl . mconcat <$> - mapM (fmap $ H.div ! A.class_ "decl_description") - (mapMaybe (describeDec visible) decs) + H.dl . mconcat + <$> mapM + (fmap $ H.div ! A.class_ "decl_description") + (mapMaybe (describeDec visible) decs) describeDec :: S.Set VName -> Dec -> Maybe (DocM Html) describeDec _ (ValDec vb) = Just $ describeGeneric (valBindName vb) (valBindWhat vb) (valBindDoc vb) $ \name -> do - (lhs, mhs, rhs) <- valBindHtml name vb - return $ lhs <> mhs <> ": " <> rhs - -describeDec _ (TypeDec vb) = Just $ - describeGeneric (typeAlias vb) IndexType (typeDoc vb) (`typeBindHtml` vb) - + (lhs, mhs, rhs) <- valBindHtml name vb + return $ lhs <> mhs <> ": " <> rhs +describeDec _ (TypeDec vb) = + Just $ + describeGeneric (typeAlias vb) IndexType (typeDoc vb) (`typeBindHtml` vb) describeDec _ (SigDec (SigBind name se doc _)) = Just $ describeGenericMod name IndexModuleType se doc $ \name' -> - return $ keyword "module type " <> name' - + return $ keyword "module type " <> name' describeDec _ (ModDec mb) = Just $ describeGeneric (modName mb) IndexModule (modDoc mb) $ \name' -> - return $ keyword "module " <> name' - -describeDec _ OpenDec{} = Nothing - + return $ keyword "module " <> name' +describeDec _ OpenDec {} = Nothing describeDec visible (LocalDec (SigDec (SigBind name se doc _)) _) | name `S.member` visible = Just $ - describeGenericMod name IndexModuleType se doc $ \name' -> - return $ keyword "local module type " <> name' - -describeDec _ LocalDec{} = Nothing -describeDec _ ImportDec{} = Nothing + describeGenericMod name IndexModuleType se doc $ \name' -> + return $ keyword "local module type " <> name' +describeDec _ LocalDec {} = Nothing +describeDec _ ImportDec {} = Nothing valBindWhat :: ValBind -> IndexWhat -valBindWhat vb | null (valBindParams vb), - orderZero (fst $ unInfo $ valBindRetType vb) = - IndexValue - | otherwise = - IxFun +valBindWhat vb + | null (valBindParams vb), + orderZero (fst $ unInfo $ valBindRetType vb) = + IndexValue + | otherwise = + IxFun describeSpecs :: [Spec] -> DocM Html describeSpecs specs = @@ -753,29 +839,35 @@ describeSpec :: Spec -> DocM Html describeSpec (ValSpec name tparams t doc _) = describeGeneric name what doc $ \name' -> do - let tparams' = mconcat $ map ((" "<>) . typeParamHtml) tparams - t' <- noLink (map typeParamName tparams) $ - typeExpHtml $ declaredType t - return $ keyword "val " <> name' <> tparams' <> ": " <> t' - where what = if orderZero (unInfo $ expandedType t) - then IndexValue else IxFun + let tparams' = mconcat $ map ((" " <>) . typeParamHtml) tparams + t' <- + noLink (map typeParamName tparams) $ + typeExpHtml $ declaredType t + return $ keyword "val " <> name' <> tparams' <> ": " <> t' + where + what = + if orderZero (unInfo $ expandedType t) + then IndexValue + else IxFun describeSpec (TypeAbbrSpec vb) = describeGeneric (typeAlias vb) IndexType (typeDoc vb) (`typeBindHtml` vb) describeSpec (TypeSpec l name tparams doc _) = describeGeneric name IndexType doc $ - return . (\name' -> typeAbbrevHtml l name' tparams) + return . (\name' -> typeAbbrevHtml l name' tparams) describeSpec (ModSpec name se doc _) = describeGenericMod name IndexModule se doc $ \name' -> - case se of - SigSpecs{} -> return $ keyword "module " <> name' - _ -> do se' <- synopsisSigExp se - return $ keyword "module " <> name' <> ": " <> se' + case se of + SigSpecs {} -> return $ keyword "module " <> name' + _ -> do + se' <- synopsisSigExp se + return $ keyword "module " <> name' <> ": " <> se' describeSpec (IncludeSpec sig _) = do sig' <- synopsisSigExp sig doc' <- docHtml Nothing - let decl_header = (H.dt ! A.class_ "desc_header") $ - (H.span ! A.class_ "synopsis_link") mempty <> - keyword "include " <> - sig' + let decl_header = + (H.dt ! A.class_ "desc_header") $ + (H.span ! A.class_ "synopsis_link") mempty + <> keyword "include " + <> sig' decl_doc = H.dd ! A.class_ "desc_doc" $ doc' return $ decl_header <> decl_doc diff --git a/src/Futhark/Error.hs b/src/Futhark/Error.hs --- a/src/Futhark/Error.hs +++ b/src/Futhark/Error.hs @@ -1,19 +1,18 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE Safe #-} + -- | Futhark error definitions. module Futhark.Error - ( CompilerError(..) - , ErrorClass(..) - - , externalError - , externalErrorS - - , InternalError(..) - , compilerBug - , compilerBugS - , compilerLimitation - , compilerLimitationS - , internalErrorS + ( CompilerError (..), + ErrorClass (..), + externalError, + externalErrorS, + InternalError (..), + compilerBug, + compilerBugS, + compilerLimitation, + compilerLimitationS, + internalErrorS, ) where @@ -25,18 +24,19 @@ -- | There are two classes of internal errors: actual bugs, and -- implementation limitations. The latter are already known and need -- not be reported. -data ErrorClass = CompilerBug - | CompilerLimitation - deriving (Eq, Ord, Show) +data ErrorClass + = CompilerBug + | CompilerLimitation + deriving (Eq, Ord, Show) -- | A compiler error. -data CompilerError = - ExternalError Doc - -- ^ An error that happened due to something the user did, such as +data CompilerError + = -- | An error that happened due to something the user did, such as -- provide incorrect code or options. - | InternalError T.Text T.Text ErrorClass - -- ^ An internal compiler error. The second text is extra data + ExternalError Doc + | -- | An internal compiler error. The second text is extra data -- for debugging, which can be written to a file. + InternalError T.Text T.Text ErrorClass instance Show CompilerError where show (ExternalError s) = pretty s diff --git a/src/Futhark/FreshNames.hs b/src/Futhark/FreshNames.hs --- a/src/Futhark/FreshNames.hs +++ b/src/Futhark/FreshNames.hs @@ -1,15 +1,16 @@ {-# LANGUAGE DeriveLift #-} + -- | This module provides facilities for generating unique names. module Futhark.FreshNames - ( VNameSource - , blankNameSource - , newNameSource - , newName - ) where - -import Language.Haskell.TH.Syntax (Lift) + ( VNameSource, + blankNameSource, + newNameSource, + newName, + ) +where import Language.Futhark.Core +import Language.Haskell.TH.Syntax (Lift) -- | A name source is conceptually an infinite sequence of names with -- no repeating entries. In practice, when asked for a name, the name @@ -30,7 +31,8 @@ -- | Produce a fresh name, using the given name as a template. newName :: VNameSource -> VName -> (VName, VNameSource) newName (VNameSource i) k = i' `seq` (VName (baseName k) i, VNameSource i') - where i' = i+1 + where + i' = i + 1 -- | A blank name source. blankNameSource :: VNameSource diff --git a/src/Futhark/IR.hs b/src/Futhark/IR.hs --- a/src/Futhark/IR.hs +++ b/src/Futhark/IR.hs @@ -1,17 +1,18 @@ {-# LANGUAGE Safe #-} + -- | A convenient re-export of basic AST modules. Note that -- "Futhark.IR.Lore" is not exported, as this would -- cause name clashes. You are advised to use a qualified import of -- the lore module, if you need it. module Futhark.IR - ( module Futhark.IR.Prop - , module Futhark.IR.Traversals - , module Futhark.IR.Pretty - , module Futhark.IR.Syntax - ) + ( module Futhark.IR.Prop, + module Futhark.IR.Traversals, + module Futhark.IR.Pretty, + module Futhark.IR.Syntax, + ) where -import Futhark.IR.Syntax +import Futhark.IR.Pretty import Futhark.IR.Prop +import Futhark.IR.Syntax import Futhark.IR.Traversals -import Futhark.IR.Pretty 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 @@ -1,67 +1,79 @@ -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + -- | A representation where all bindings are annotated with aliasing -- information. module Futhark.IR.Aliases - ( -- * The Lore definition - Aliases - , AliasDec (..) - , VarAliases - , ConsumedInExp - , BodyAliasing - , module Futhark.IR.Prop.Aliases - -- * Module re-exports - , module Futhark.IR.Prop - , module Futhark.IR.Traversals - , module Futhark.IR.Pretty - , module Futhark.IR.Syntax - -- * Adding aliases - , addAliasesToPattern - , mkAliasedLetStm - , mkAliasedBody - , mkPatternAliases - , mkBodyAliases - -- * Removing aliases - , removeProgAliases - , removeFunDefAliases - , removeExpAliases - , removeStmAliases - , removeLambdaAliases - , removePatternAliases - , removeScopeAliases - -- * Tracking aliases - , AliasesAndConsumed - , trackAliases - , consumedInStms - ) + ( -- * The Lore definition + Aliases, + AliasDec (..), + VarAliases, + ConsumedInExp, + BodyAliasing, + module Futhark.IR.Prop.Aliases, + + -- * Module re-exports + module Futhark.IR.Prop, + module Futhark.IR.Traversals, + module Futhark.IR.Pretty, + module Futhark.IR.Syntax, + + -- * Adding aliases + addAliasesToPattern, + mkAliasedLetStm, + mkAliasedBody, + mkPatternAliases, + mkBodyAliases, + + -- * Removing aliases + removeProgAliases, + removeFunDefAliases, + removeExpAliases, + removeStmAliases, + removeLambdaAliases, + removePatternAliases, + removeScopeAliases, + + -- * Tracking aliases + AliasesAndConsumed, + trackAliases, + consumedInStms, + ) where import Control.Monad.Identity import Control.Monad.Reader -import Data.Maybe import qualified Data.Map.Strict as M - -import Futhark.IR.Syntax +import Data.Maybe +import Futhark.Analysis.Rephrase +import Futhark.Binder +import Futhark.IR.Pretty import Futhark.IR.Prop import Futhark.IR.Prop.Aliases +import Futhark.IR.Syntax import Futhark.IR.Traversals -import Futhark.IR.Pretty import Futhark.Transform.Rename -import Futhark.Binder import Futhark.Transform.Substitute -import Futhark.Analysis.Rephrase import qualified Futhark.Util.Pretty as PP +import GHC.Generics +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic -- | The lore for the basic representation. data Aliases lore -- | A wrapper around 'AliasDec to get around the fact that we need an -- 'Ord' instance, which 'AliasDec does not have. -newtype AliasDec = AliasDec { unAliases :: Names } - deriving (Show) +newtype AliasDec = AliasDec {unAliases :: Names} + deriving (Show, Generic) +instance SexpIso AliasDec where + sexpIso = with $ \vname -> sexpIso >>> vname + instance Semigroup AliasDec where x <> y = AliasDec $ unAliases x <> unAliases y @@ -96,8 +108,10 @@ -- consumed inside of it. type BodyAliasing = ([VarAliases], ConsumedInExp) -instance (Decorations lore, CanBeAliased (Op lore)) => - Decorations (Aliases lore) where +instance + (Decorations lore, CanBeAliased (Op lore)) => + Decorations (Aliases lore) + where type LetDec (Aliases lore) = (VarAliases, LetDec lore) type ExpDec (Aliases lore) = (ConsumedInExp, ExpDec lore) type BodyDec (Aliases lore) = (BodyAliasing, BodyDec lore) @@ -110,10 +124,12 @@ instance AliasesOf (VarAliases, dec) where aliasesOf = unAliases . fst -instance FreeDec AliasDec where +instance FreeDec AliasDec -withoutAliases :: (HasScope (Aliases lore) m, Monad m) => - ReaderT (Scope lore) m a -> m a +withoutAliases :: + (HasScope (Aliases lore) m, Monad m) => + ReaderT (Scope lore) m a -> + m a withoutAliases m = do scope <- asksScope removeScopeAliases runReaderT m scope @@ -126,42 +142,49 @@ bodyAliases = map unAliases . fst . fst . bodyDec consumedInBody = unAliases . snd . fst . bodyDec -instance PrettyAnnot (PatElemT dec) => - PrettyAnnot (PatElemT (VarAliases, dec)) where - +instance + PrettyAnnot (PatElemT dec) => + PrettyAnnot (PatElemT (VarAliases, dec)) + where ppAnnot (PatElem name (AliasDec als, dec)) = let alias_comment = PP.oneLine <$> aliasComment name als - in case (alias_comment, ppAnnot (PatElem name dec)) of - (_, Nothing) -> - alias_comment - (Just alias_comment', Just inner_comment) -> - Just $ alias_comment' PP. inner_comment - (Nothing, Just inner_comment) -> - Just inner_comment - + in case (alias_comment, ppAnnot (PatElem name dec)) of + (_, Nothing) -> + alias_comment + (Just alias_comment', Just inner_comment) -> + Just $ alias_comment' PP. inner_comment + (Nothing, Just inner_comment) -> + Just inner_comment instance (ASTLore lore, CanBeAliased (Op lore)) => PrettyLore (Aliases lore) where ppExpLore (consumed, inner) e = - maybeComment $ catMaybes [exp_dec, - merge_dec, - ppExpLore inner $ removeExpAliases e] - where merge_dec = - case e of - DoLoop _ merge _ body -> - let mergeParamAliases fparam als - | primType (paramType fparam) = - Nothing - | otherwise = - resultAliasComment (paramName fparam) als - in maybeComment $ catMaybes $ - zipWith mergeParamAliases (map fst merge) $ - bodyAliases body - _ -> Nothing + maybeComment $ + catMaybes + [ exp_dec, + merge_dec, + ppExpLore inner $ removeExpAliases e + ] + where + merge_dec = + case e of + DoLoop _ merge _ body -> + let mergeParamAliases fparam als + | primType (paramType fparam) = + Nothing + | otherwise = + resultAliasComment (paramName fparam) als + in maybeComment $ + catMaybes $ + zipWith mergeParamAliases (map fst merge) $ + bodyAliases body + _ -> Nothing - exp_dec = case namesToList $ unAliases consumed of - [] -> Nothing - als -> Just $ PP.oneLine $ - PP.text "-- Consumes " <> PP.commasep (map PP.ppr als) + exp_dec = case namesToList $ unAliases consumed of + [] -> Nothing + als -> + Just $ + PP.oneLine $ + PP.text "-- Consumes " <> PP.commasep (map PP.ppr als) maybeComment :: [PP.Doc] -> Maybe PP.Doc maybeComment [] = Nothing @@ -171,117 +194,153 @@ aliasComment name als = case namesToList als of [] -> Nothing - als' -> Just $ PP.oneLine $ - PP.text "-- " <> PP.ppr name <> PP.text " aliases " <> - PP.commasep (map PP.ppr als') + als' -> + Just $ + PP.oneLine $ + PP.text "-- " <> PP.ppr name <> PP.text " aliases " + <> PP.commasep (map PP.ppr als') resultAliasComment :: PP.Pretty a => a -> Names -> Maybe PP.Doc resultAliasComment name als = case namesToList als of [] -> Nothing - als' -> Just $ PP.oneLine $ - PP.text "-- Result of " <> PP.ppr name <> PP.text " aliases " <> - PP.commasep (map PP.ppr als') + als' -> + Just $ + PP.oneLine $ + PP.text "-- Result of " <> PP.ppr name <> PP.text " aliases " + <> PP.commasep (map PP.ppr als') removeAliases :: CanBeAliased (Op lore) => Rephraser Identity (Aliases lore) lore -removeAliases = Rephraser { rephraseExpLore = return . snd - , rephraseLetBoundLore = return . snd - , rephraseBodyLore = return . snd - , rephraseFParamLore = return - , rephraseLParamLore = return - , rephraseRetType = return - , rephraseBranchType = return - , rephraseOp = return . removeOpAliases - } +removeAliases = + Rephraser + { rephraseExpLore = return . snd, + rephraseLetBoundLore = return . snd, + rephraseBodyLore = return . snd, + rephraseFParamLore = return, + rephraseLParamLore = return, + rephraseRetType = return, + rephraseBranchType = return, + rephraseOp = return . removeOpAliases + } removeScopeAliases :: Scope (Aliases lore) -> Scope lore removeScopeAliases = M.map unAlias - where unAlias (LetName (_, dec)) = LetName dec - unAlias (FParamName dec) = FParamName dec - unAlias (LParamName dec) = LParamName dec - unAlias (IndexName it) = IndexName it + where + unAlias (LetName (_, dec)) = LetName dec + unAlias (FParamName dec) = FParamName dec + unAlias (LParamName dec) = LParamName dec + unAlias (IndexName it) = IndexName it -removeProgAliases :: CanBeAliased (Op lore) => - Prog (Aliases lore) -> Prog lore +removeProgAliases :: + CanBeAliased (Op lore) => + Prog (Aliases lore) -> + Prog lore removeProgAliases = runIdentity . rephraseProg removeAliases -removeFunDefAliases :: CanBeAliased (Op lore) => - FunDef (Aliases lore) -> FunDef lore +removeFunDefAliases :: + CanBeAliased (Op lore) => + FunDef (Aliases lore) -> + FunDef lore removeFunDefAliases = runIdentity . rephraseFunDef removeAliases -removeExpAliases :: CanBeAliased (Op lore) => - Exp (Aliases lore) -> Exp lore +removeExpAliases :: + CanBeAliased (Op lore) => + Exp (Aliases lore) -> + Exp lore removeExpAliases = runIdentity . rephraseExp removeAliases -removeStmAliases :: CanBeAliased (Op lore) => - Stm (Aliases lore) -> Stm lore +removeStmAliases :: + CanBeAliased (Op lore) => + Stm (Aliases lore) -> + Stm lore removeStmAliases = runIdentity . rephraseStm removeAliases -removeLambdaAliases :: CanBeAliased (Op lore) => - Lambda (Aliases lore) -> Lambda lore +removeLambdaAliases :: + CanBeAliased (Op lore) => + Lambda (Aliases lore) -> + Lambda lore removeLambdaAliases = runIdentity . rephraseLambda removeAliases -removePatternAliases :: PatternT (AliasDec, a) - -> PatternT a +removePatternAliases :: + PatternT (AliasDec, a) -> + PatternT a removePatternAliases = runIdentity . rephrasePattern (return . snd) -addAliasesToPattern :: (ASTLore lore, CanBeAliased (Op lore), Typed dec) => - PatternT dec -> Exp (Aliases lore) - -> PatternT (VarAliases, dec) +addAliasesToPattern :: + (ASTLore lore, CanBeAliased (Op lore), Typed dec) => + PatternT dec -> + Exp (Aliases lore) -> + PatternT (VarAliases, dec) addAliasesToPattern pat e = uncurry Pattern $ mkPatternAliases pat e -mkAliasedBody :: (ASTLore lore, CanBeAliased (Op lore)) => - BodyDec lore -> Stms (Aliases lore) -> Result -> Body (Aliases lore) +mkAliasedBody :: + (ASTLore lore, CanBeAliased (Op lore)) => + BodyDec lore -> + Stms (Aliases lore) -> + Result -> + Body (Aliases lore) mkAliasedBody innerlore bnds res = Body (mkBodyAliases bnds res, innerlore) bnds res -mkPatternAliases :: (Aliased lore, Typed dec) => - PatternT dec -> Exp lore - -> ([PatElemT (VarAliases, dec)], - [PatElemT (VarAliases, dec)]) +mkPatternAliases :: + (Aliased lore, Typed dec) => + PatternT dec -> + Exp lore -> + ( [PatElemT (VarAliases, dec)], + [PatElemT (VarAliases, dec)] + ) mkPatternAliases pat e = -- Some part of the pattern may be the context. This does not have -- aliases from expAliases, so we use a hack to compute aliases of -- the context. let als = expAliases e ++ repeat mempty -- In case the pattern has - -- more elements (this - -- implies a type error). + -- more elements (this + -- implies a type error). context_als = mkContextAliases pat e - in (zipWith annotateBindee (patternContextElements pat) context_als, - zipWith annotateBindee (patternValueElements pat) als) - where annotateBindee bindee names = - bindee `setPatElemLore` (AliasDec names', patElemDec bindee) - where names' = - case patElemType bindee of - Array {} -> names - Mem _ -> names - _ -> mempty + in ( zipWith annotateBindee (patternContextElements pat) context_als, + zipWith annotateBindee (patternValueElements pat) als + ) + where + annotateBindee bindee names = + bindee `setPatElemLore` (AliasDec names', patElemDec bindee) + where + names' = + case patElemType bindee of + Array {} -> names + Mem _ -> names + _ -> mempty -mkContextAliases :: Aliased lore => - PatternT dec -> Exp lore -> [Names] +mkContextAliases :: + Aliased lore => + PatternT dec -> + Exp lore -> + [Names] mkContextAliases pat (DoLoop ctxmerge valmerge _ body) = let ctx = map fst ctxmerge init_als = zip mergenames $ map (subExpAliases . snd) $ ctxmerge ++ valmerge expand als = als <> mconcat (mapMaybe (`lookup` init_als) (namesToList als)) - merge_als = zip mergenames $ - map ((`namesSubtract` mergenames_set) . expand) $ - bodyAliases body - in if length ctx == length (patternContextElements pat) - then map (fromMaybe mempty . flip lookup merge_als . paramName) ctx - else map (const mempty) $ patternContextElements pat - where mergenames = map (paramName . fst) $ ctxmerge ++ valmerge - mergenames_set = namesFromList mergenames + merge_als = + zip mergenames $ + map ((`namesSubtract` mergenames_set) . expand) $ + bodyAliases body + in if length ctx == length (patternContextElements pat) + then map (fromMaybe mempty . flip lookup merge_als . paramName) ctx + else map (const mempty) $ patternContextElements pat + where + mergenames = map (paramName . fst) $ ctxmerge ++ valmerge + mergenames_set = namesFromList mergenames mkContextAliases pat (If _ tbranch fbranch _) = take (length $ patternContextNames pat) $ - zipWith (<>) (bodyAliases tbranch) (bodyAliases fbranch) + zipWith (<>) (bodyAliases tbranch) (bodyAliases fbranch) mkContextAliases pat _ = replicate (length $ patternContextElements pat) mempty -mkBodyAliases :: Aliased lore => - Stms lore - -> Result - -> BodyAliasing +mkBodyAliases :: + Aliased lore => + Stms lore -> + Result -> + BodyAliasing mkBodyAliases bnds res = -- We need to remove the names that are bound in bnds from the alias -- and consumption sets. We do this by computing the transitive @@ -292,56 +351,70 @@ foldMap (namesFromList . patternNames . stmPattern) bnds aliases' = map (`namesSubtract` boundNames) aliases consumed' = consumed `namesSubtract` boundNames - in (map AliasDec aliases', AliasDec consumed') + in (map AliasDec aliases', AliasDec consumed') -mkStmsAliases :: Aliased lore => - Stms lore -> [SubExp] - -> ([Names], Names) +mkStmsAliases :: + Aliased lore => + Stms lore -> + [SubExp] -> + ([Names], Names) mkStmsAliases bnds res = delve mempty $ stmsToList bnds - where delve (aliasmap, consumed) [] = - (map (aliasClosure aliasmap . subExpAliases) res, - consumed) - delve (aliasmap, consumed) (bnd:bnds') = - delve (trackAliases (aliasmap, consumed) bnd) bnds' - aliasClosure aliasmap names = - names <> mconcat (map look $ namesToList names) - where look k = M.findWithDefault mempty k aliasmap + where + delve (aliasmap, consumed) [] = + ( map (aliasClosure aliasmap . subExpAliases) res, + consumed + ) + delve (aliasmap, consumed) (bnd : bnds') = + delve (trackAliases (aliasmap, consumed) bnd) bnds' + aliasClosure aliasmap names = + names <> mconcat (map look $ namesToList names) + 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) +type AliasesAndConsumed = + ( M.Map VName Names, + Names + ) -trackAliases :: Aliased lore => - AliasesAndConsumed -> Stm lore - -> AliasesAndConsumed +trackAliases :: + Aliased lore => + AliasesAndConsumed -> + Stm lore -> + AliasesAndConsumed trackAliases (aliasmap, consumed) bnd = let pat = stmPattern bnd - als = M.fromList $ - zip (patternNames pat) (map addAliasesOfAliases $ patternAliases pat) + als = + M.fromList $ + zip (patternNames pat) (map addAliasesOfAliases $ patternAliases pat) aliasmap' = als <> aliasmap consumed' = consumed <> addAliasesOfAliases (consumedInStm bnd) - in (aliasmap', consumed') - where addAliasesOfAliases names = names <> aliasesOfAliases names - aliasesOfAliases = mconcat . map look . namesToList - look k = M.findWithDefault mempty k aliasmap + in (aliasmap', consumed') + where + addAliasesOfAliases names = names <> aliasesOfAliases names + aliasesOfAliases = mconcat . map look . namesToList + look k = M.findWithDefault mempty k aliasmap -mkAliasedLetStm :: (ASTLore lore, CanBeAliased (Op lore)) => - Pattern lore - -> StmAux (ExpDec lore) -> Exp (Aliases lore) - -> Stm (Aliases lore) +mkAliasedLetStm :: + (ASTLore lore, CanBeAliased (Op lore)) => + Pattern lore -> + StmAux (ExpDec lore) -> + Exp (Aliases lore) -> + Stm (Aliases lore) mkAliasedLetStm pat (StmAux cs attrs dec) e = - Let (addAliasesToPattern pat e) - (StmAux cs attrs (AliasDec $ consumedInExp e, dec)) - e + Let + (addAliasesToPattern pat e) + (StmAux cs attrs (AliasDec $ consumedInExp e, dec)) + e instance (Bindable lore, CanBeAliased (Op lore)) => Bindable (Aliases lore) where mkExpDec pat e = let dec = mkExpDec (removePatternAliases pat) $ removeExpAliases e - in (AliasDec $ consumedInExp e, dec) + in (AliasDec $ consumedInExp e, dec) mkExpPat ctx val e = addAliasesToPattern (mkExpPat ctx val $ removeExpAliases e) e @@ -354,6 +427,6 @@ mkBody bnds res = let Body bodylore _ _ = mkBody (fmap removeStmAliases bnds) res - in mkAliasedBody bodylore bnds res + in mkAliasedBody bodylore bnds res -instance (ASTLore (Aliases lore), Bindable (Aliases lore)) => BinderOps (Aliases lore) where +instance (ASTLore (Aliases lore), Bindable (Aliases lore)) => BinderOps (Aliases lore) diff --git a/src/Futhark/IR/Decorations.hs b/src/Futhark/IR/Decorations.hs --- a/src/Futhark/IR/Decorations.hs +++ b/src/Futhark/IR/Decorations.hs @@ -1,52 +1,101 @@ -{-# LANGUAGE TypeFamilies, FlexibleContexts #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + -- | The core Futhark AST is parameterised by a @lore@ type parameter, -- which is then used to invoke the type families defined here. module Futhark.IR.Decorations - ( Decorations (..) - , module Futhark.IR.RetType - ) - where + ( Decorations (..), + module Futhark.IR.RetType, + ) +where import qualified Data.Kind - -import Futhark.IR.Syntax.Core -import Futhark.IR.RetType import Futhark.IR.Prop.Types +import Futhark.IR.RetType +import Futhark.IR.Syntax.Core +import Language.SexpGrammar (SexpIso) -- | A collection of type families, along with constraints specifying -- that the types they map to should satisfy some minimal -- requirements. -class (Show (LetDec l), Show (ExpDec l), Show (BodyDec l), Show (FParamInfo l), Show (LParamInfo l), Show (RetType l), Show (BranchType l), Show (Op l), - Eq (LetDec l), Eq (ExpDec l), Eq (BodyDec l), Eq (FParamInfo l), Eq (LParamInfo l), Eq (RetType l), Eq (BranchType l), Eq (Op l), - Ord (LetDec l), Ord (ExpDec l), Ord (BodyDec l), Ord (FParamInfo l), Ord (LParamInfo l), Ord (RetType l), Ord (BranchType l), Ord (Op l), - IsRetType (RetType l), IsBodyType (BranchType l), - Typed (FParamInfo l), Typed (LParamInfo l), Typed (LetDec l), - DeclTyped (FParamInfo l)) - => Decorations l where +class + ( Show (LetDec l), + Show (ExpDec l), + Show (BodyDec l), + Show (FParamInfo l), + Show (LParamInfo l), + Show (RetType l), + Show (BranchType l), + Show (Op l), + Eq (LetDec l), + Eq (ExpDec l), + Eq (BodyDec l), + Eq (FParamInfo l), + Eq (LParamInfo l), + Eq (RetType l), + Eq (BranchType l), + Eq (Op l), + Ord (LetDec l), + Ord (ExpDec l), + Ord (BodyDec l), + Ord (FParamInfo l), + Ord (LParamInfo l), + Ord (RetType l), + Ord (BranchType l), + Ord (Op l), + IsRetType (RetType l), + IsBodyType (BranchType l), + Typed (FParamInfo l), + Typed (LParamInfo l), + Typed (LetDec l), + DeclTyped (FParamInfo l), + SexpIso (LetDec l), + SexpIso (ExpDec l), + SexpIso (BodyDec l), + SexpIso (FParamInfo l), + SexpIso (LParamInfo l), + SexpIso (RetType l), + SexpIso (BranchType l), + SexpIso (Op l) + ) => + Decorations l + where -- | Decoration for every let-pattern element. type LetDec l :: Data.Kind.Type + type LetDec l = Type + -- | Decoration for every expression. type ExpDec l :: Data.Kind.Type + type ExpDec l = () + -- | Decoration for every body. type BodyDec l :: Data.Kind.Type + type BodyDec l = () + -- | Decoration for every (non-lambda) function parameter. type FParamInfo l :: Data.Kind.Type + type FParamInfo l = DeclType + -- | Decoration for every lambda function parameter. type LParamInfo l :: Data.Kind.Type + type LParamInfo l = Type -- | The return type decoration of function calls. type RetType l :: Data.Kind.Type + type RetType l = DeclExtType -- | The return type decoration of branches. type BranchType l :: Data.Kind.Type + type BranchType l = ExtType -- | Extensible operation. type Op l :: Data.Kind.Type + type Op l = () diff --git a/src/Futhark/IR/Kernels.hs b/src/Futhark/IR/Kernels.hs --- a/src/Futhark/IR/Kernels.hs +++ b/src/Futhark/IR/Kernels.hs @@ -1,29 +1,31 @@ -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} + -- | A representation with flat parallelism via GPU-oriented kernels. module Futhark.IR.Kernels - ( -- * The Lore definition - Kernels - -- * Module re-exports - , module Futhark.IR.Prop - , module Futhark.IR.Traversals - , module Futhark.IR.Pretty - , module Futhark.IR.Syntax - , module Futhark.IR.Kernels.Kernel - , module Futhark.IR.Kernels.Sizes - , module Futhark.IR.SOACS.SOAC - ) + ( -- * The Lore definition + Kernels, + + -- * Module re-exports + module Futhark.IR.Prop, + module Futhark.IR.Traversals, + module Futhark.IR.Pretty, + module Futhark.IR.Syntax, + module Futhark.IR.Kernels.Kernel, + module Futhark.IR.Kernels.Sizes, + module Futhark.IR.SOACS.SOAC, + ) where -import Futhark.IR.Syntax +import Futhark.Binder +import Futhark.Construct import Futhark.IR.Kernels.Kernel import Futhark.IR.Kernels.Sizes +import Futhark.IR.Pretty import Futhark.IR.Prop +import Futhark.IR.SOACS.SOAC hiding (HistOp (..)) +import Futhark.IR.Syntax import Futhark.IR.Traversals -import Futhark.IR.Pretty -import Futhark.IR.SOACS.SOAC hiding (HistOp(..)) -import Futhark.Binder -import Futhark.Construct import qualified Futhark.TypeCheck as TypeCheck -- | The phantom data type for the kernels representation. @@ -31,15 +33,17 @@ instance Decorations Kernels where type Op Kernels = HostOp Kernels (SOAC Kernels) + instance ASTLore Kernels where expTypesFromPattern = return . expExtTypesFromPattern instance TypeCheck.CheckableOp Kernels where checkOp = typeCheckKernelsOp Nothing - where typeCheckKernelsOp lvl = - typeCheckHostOp (typeCheckKernelsOp . Just) lvl typeCheckSOAC + where + typeCheckKernelsOp lvl = + typeCheckHostOp (typeCheckKernelsOp . Just) lvl typeCheckSOAC -instance TypeCheck.Checkable Kernels where +instance TypeCheck.Checkable Kernels instance Bindable Kernels where mkBody = Body () @@ -47,9 +51,9 @@ mkExpDec _ _ = () mkLetNames = simpleMkLetNames -instance BinderOps Kernels where +instance BinderOps Kernels -instance PrettyLore Kernels where +instance PrettyLore Kernels instance HasSegOp Kernels where type SegOpLevel Kernels = SegLevel 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 @@ -1,81 +1,117 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + module Futhark.IR.Kernels.Kernel ( -- * Size operations - SizeOp(..) + SizeOp (..), -- * Host operations - , HostOp(..) - , typeCheckHostOp + HostOp (..), + typeCheckHostOp, -- * SegOp refinements - , SegLevel(..) + SegLevel (..), -- * Reexports - , module Futhark.IR.Kernels.Sizes - , module Futhark.IR.SegOp + module Futhark.IR.Kernels.Sizes, + module Futhark.IR.SegOp, ) where -import Futhark.IR +import Control.Category +import Futhark.Analysis.Metrics import qualified Futhark.Analysis.SymbolTable as ST -import qualified Futhark.Util.Pretty as PP -import Futhark.Util.Pretty - ((), (<+>), ppr, commasep, parens, text) -import Futhark.Transform.Substitute -import Futhark.Transform.Rename -import Futhark.Optimise.Simplify.Lore -import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.IR.Prop.Aliases +import Futhark.IR import Futhark.IR.Aliases (Aliases) -import Futhark.IR.SegOp import Futhark.IR.Kernels.Sizes +import Futhark.IR.Prop.Aliases +import Futhark.IR.SegOp +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Optimise.Simplify.Lore +import Futhark.Transform.Rename +import Futhark.Transform.Substitute import qualified Futhark.TypeCheck as TC -import Futhark.Analysis.Metrics +import Futhark.Util.Pretty + ( commasep, + parens, + ppr, + text, + (<+>), + (), + ) +import qualified Futhark.Util.Pretty as PP +import GHC.Generics (Generic) +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) -- | At which level the *body* of a t'SegOp' executes. -data SegLevel = SegThread { segNumGroups :: Count NumGroups SubExp - , segGroupSize :: Count GroupSize SubExp - , segVirt :: SegVirt } - | SegGroup { segNumGroups :: Count NumGroups SubExp - , segGroupSize :: Count GroupSize SubExp - , segVirt :: SegVirt } - deriving (Eq, Ord, Show) +data SegLevel + = SegThread + { segNumGroups :: Count NumGroups SubExp, + segGroupSize :: Count GroupSize SubExp, + segVirt :: SegVirt + } + | SegGroup + { segNumGroups :: Count NumGroups SubExp, + segGroupSize :: Count GroupSize SubExp, + segVirt :: SegVirt + } + deriving (Eq, Ord, Show, Generic) +instance SexpIso SegLevel where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "thread") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "group") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End + instance PP.Pretty SegLevel where ppr lvl = - lvl' - PP.parens (text "#groups=" <> ppr (segNumGroups lvl) <> PP.semi <+> - text "groupsize=" <> ppr (segGroupSize lvl) <> - case segVirt lvl of - SegNoVirt -> mempty - SegNoVirtFull -> PP.semi <+> text "full" - SegVirt -> PP.semi <+> text "virtualise") - - where lvl' = case lvl of SegThread{} -> "_thread" - SegGroup{} -> "_group" + lvl' + PP.parens + ( text "#groups=" <> ppr (segNumGroups lvl) <> PP.semi + <+> text "groupsize=" + <> ppr (segGroupSize lvl) + <> case segVirt lvl of + SegNoVirt -> mempty + SegNoVirtFull -> PP.semi <+> text "full" + SegVirt -> PP.semi <+> text "virtualise" + ) + where + lvl' = case lvl of + SegThread {} -> "_thread" + SegGroup {} -> "_group" instance Engine.Simplifiable SegLevel where simplify (SegThread num_groups group_size virt) = - SegThread <$> traverse Engine.simplify num_groups <*> - traverse Engine.simplify group_size <*> pure virt + SegThread <$> traverse Engine.simplify num_groups + <*> traverse Engine.simplify group_size + <*> pure virt simplify (SegGroup num_groups group_size virt) = - SegGroup <$> traverse Engine.simplify num_groups <*> - traverse Engine.simplify group_size <*> pure virt + SegGroup <$> traverse Engine.simplify num_groups + <*> traverse Engine.simplify group_size + <*> pure virt instance Substitute SegLevel where substituteNames substs (SegThread num_groups group_size virt) = SegThread - (substituteNames substs num_groups) (substituteNames substs group_size) virt + (substituteNames substs num_groups) + (substituteNames substs group_size) + virt substituteNames substs (SegGroup num_groups group_size virt) = SegGroup - (substituteNames substs num_groups) (substituteNames substs group_size) virt + (substituteNames substs num_groups) + (substituteNames substs group_size) + virt instance Rename SegLevel where rename = substituteRename @@ -88,8 +124,7 @@ -- | A simple size-level query or computation. data SizeOp - = SplitSpace SplitOrdering SubExp SubExp SubExp - -- ^ @SplitSpace o w i elems_per_thread@. + = -- | @SplitSpace o w i elems_per_thread@. -- -- Computes how to divide array elements to -- threads in a kernel. Returns the number of @@ -110,42 +145,54 @@ -- the thread will receive elements @i, -- i+stride, i+2*stride, ..., -- i+(elems_per_thread-1)*stride@. - | GetSize Name SizeClass - -- ^ Produce some runtime-configurable size. - | GetSizeMax SizeClass - -- ^ The maximum size of some class. - | CmpSizeLe Name SizeClass SubExp - -- ^ Compare size (likely a threshold) with some integer value. - | CalcNumGroups SubExp Name SubExp - -- ^ @CalcNumGroups w max_num_groups group_size@ calculates the + SplitSpace SplitOrdering SubExp SubExp SubExp + | -- | Produce some runtime-configurable size. + GetSize Name SizeClass + | -- | The maximum size of some class. + GetSizeMax SizeClass + | -- | Compare size (likely a threshold) with some integer value. + CmpSizeLe Name SizeClass SubExp + | -- | @CalcNumGroups w max_num_groups group_size@ calculates the -- number of GPU workgroups to use for an input of the given size. -- The @Name@ is a size name. Note that @w@ is an i64 to avoid -- overflow issues. - deriving (Eq, Ord, Show) + CalcNumGroups SubExp Name SubExp + deriving (Eq, Ord, Show, Generic) +instance SexpIso SizeOp where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "split-space") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "get-size") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "get-size-max") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "cmp-size-le") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "calc-num-groups") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End + instance Substitute SizeOp where substituteNames subst (SplitSpace o w i elems_per_thread) = SplitSpace - (substituteNames subst o) - (substituteNames subst w) - (substituteNames subst i) - (substituteNames subst elems_per_thread) + (substituteNames subst o) + (substituteNames subst w) + (substituteNames subst i) + (substituteNames subst elems_per_thread) substituteNames substs (CmpSizeLe name sclass x) = CmpSizeLe name sclass (substituteNames substs x) substituteNames substs (CalcNumGroups w max_num_groups group_size) = CalcNumGroups - (substituteNames substs w) - max_num_groups - (substituteNames substs group_size) + (substituteNames substs w) + max_num_groups + (substituteNames substs group_size) substituteNames _ op = op instance Rename SizeOp where rename (SplitSpace o w i elems_per_thread) = SplitSpace - <$> rename o - <*> rename w - <*> rename i - <*> rename elems_per_thread + <$> rename o + <*> rename w + <*> rename i + <*> rename elems_per_thread rename (CmpSizeLe name sclass x) = CmpSizeLe name sclass <$> rename x rename (CalcNumGroups w max_num_groups group_size) = @@ -157,11 +204,11 @@ cheapOp _ = True instance TypedOp SizeOp where - opType SplitSpace{} = pure [Prim int32] - opType (GetSize _ _) = pure [Prim int32] - opType (GetSizeMax _) = pure [Prim int32] - opType CmpSizeLe{} = pure [Prim Bool] - opType CalcNumGroups{} = pure [Prim int32] + opType SplitSpace {} = pure [Prim int64] + opType (GetSize _ _) = pure [Prim int64] + opType (GetSizeMax _) = pure [Prim int64] + opType CmpSizeLe {} = pure [Prim Bool] + opType CalcNumGroups {} = pure [Prim int64] instance AliasedOp SizeOp where opAliases _ = [mempty] @@ -176,51 +223,60 @@ instance PP.Pretty SizeOp where ppr (SplitSpace o w i elems_per_thread) = - text "splitSpace" <> suff <> - parens (commasep [ppr w, ppr i, ppr elems_per_thread]) - where suff = case o of SplitContiguous -> mempty - SplitStrided stride -> text "Strided" <> parens (ppr stride) - + text "splitSpace" <> suff + <> parens (commasep [ppr w, ppr i, ppr elems_per_thread]) + where + suff = case o of + SplitContiguous -> mempty + SplitStrided stride -> text "Strided" <> parens (ppr stride) ppr (GetSize name size_class) = text "get_size" <> parens (commasep [ppr name, ppr size_class]) - ppr (GetSizeMax size_class) = text "get_size_max" <> parens (commasep [ppr size_class]) - ppr (CmpSizeLe name size_class x) = - text "get_size" <> parens (commasep [ppr name, ppr size_class]) <+> - text "<=" <+> ppr x - + text "get_size" <> parens (commasep [ppr name, ppr size_class]) + <+> text "<=" + <+> ppr x ppr (CalcNumGroups w max_num_groups group_size) = text "calc_num_groups" <> parens (commasep [ppr w, ppr max_num_groups, ppr group_size]) instance OpMetrics SizeOp where - opMetrics SplitSpace{} = seen "SplitSpace" - opMetrics GetSize{} = seen "GetSize" - opMetrics GetSizeMax{} = seen "GetSizeMax" - opMetrics CmpSizeLe{} = seen "CmpSizeLe" - opMetrics CalcNumGroups{} = seen "CalcNumGroups" + opMetrics SplitSpace {} = seen "SplitSpace" + opMetrics GetSize {} = seen "GetSize" + opMetrics GetSizeMax {} = seen "GetSizeMax" + opMetrics CmpSizeLe {} = seen "CmpSizeLe" + opMetrics CalcNumGroups {} = seen "CalcNumGroups" typeCheckSizeOp :: TC.Checkable lore => SizeOp -> TC.TypeM lore () typeCheckSizeOp (SplitSpace o w i elems_per_thread) = do case o of - SplitContiguous -> return () - SplitStrided stride -> TC.require [Prim int32] stride - mapM_ (TC.require [Prim int32]) [w, i, elems_per_thread] -typeCheckSizeOp GetSize{} = return () -typeCheckSizeOp GetSizeMax{} = return () -typeCheckSizeOp (CmpSizeLe _ _ x) = TC.require [Prim int32] x -typeCheckSizeOp (CalcNumGroups w _ group_size) = do TC.require [Prim int64] w - TC.require [Prim int32] group_size + SplitContiguous -> return () + SplitStrided stride -> TC.require [Prim int64] stride + mapM_ (TC.require [Prim int64]) [w, i, elems_per_thread] +typeCheckSizeOp GetSize {} = return () +typeCheckSizeOp GetSizeMax {} = return () +typeCheckSizeOp (CmpSizeLe _ _ x) = TC.require [Prim int64] x +typeCheckSizeOp (CalcNumGroups w _ group_size) = do + TC.require [Prim int64] w + TC.require [Prim int64] group_size -- | A host-level operation; parameterised by what else it can do. data HostOp lore op - = SegOp (SegOp SegLevel lore) - -- ^ A segmented operation. + = -- | A segmented operation. + SegOp (SegOp SegLevel lore) | SizeOp SizeOp | OtherOp op - deriving (Eq, Ord, Show) + deriving (Eq, Ord, Show, Generic) +instance (SexpIso op, Decorations lore) => SexpIso (HostOp lore op) where + sexpIso = + match $ + With (. sexpIso) $ + With (. sexpIso) $ + With + (. sexpIso) + End + instance (ASTLore lore, Substitute op) => Substitute (HostOp lore op) where substituteNames substs (SegOp op) = SegOp $ substituteNames substs op @@ -295,28 +351,32 @@ opMetrics (OtherOp op) = opMetrics op opMetrics (SizeOp op) = opMetrics op -checkSegLevel :: TC.Checkable lore => - Maybe SegLevel -> SegLevel -> TC.TypeM lore () +checkSegLevel :: + TC.Checkable lore => + Maybe SegLevel -> + SegLevel -> + TC.TypeM lore () checkSegLevel Nothing lvl = do - TC.require [Prim int32] $ unCount $ segNumGroups lvl - TC.require [Prim int32] $ unCount $ segGroupSize lvl -checkSegLevel (Just SegThread{}) _ = + TC.require [Prim int64] $ unCount $ segNumGroups lvl + TC.require [Prim int64] $ unCount $ segGroupSize lvl +checkSegLevel (Just SegThread {}) _ = TC.bad $ TC.TypeError "SegOps cannot occur when already at thread level." checkSegLevel (Just x) y | x == y = TC.bad $ TC.TypeError $ "Already at at level " ++ pretty x | segNumGroups x /= segNumGroups y || segGroupSize x /= segGroupSize y = - TC.bad $ TC.TypeError "Physical layout for SegLevel does not match parent SegLevel." + TC.bad $ TC.TypeError "Physical layout for SegLevel does not match parent SegLevel." | otherwise = - return () + return () -typeCheckHostOp :: TC.Checkable lore => - (SegLevel -> OpWithAliases (Op lore) -> TC.TypeM lore ()) - -> Maybe SegLevel - -> (op -> TC.TypeM lore ()) - -> HostOp (Aliases lore) op - -> TC.TypeM lore () +typeCheckHostOp :: + TC.Checkable lore => + (SegLevel -> OpWithAliases (Op lore) -> TC.TypeM lore ()) -> + Maybe SegLevel -> + (op -> TC.TypeM lore ()) -> + HostOp (Aliases lore) op -> + TC.TypeM lore () typeCheckHostOp checker lvl _ (SegOp op) = TC.checkOpWith (checker $ segLevel op) $ - typeCheckSegOp (checkSegLevel lvl) op + typeCheckSegOp (checkSegLevel lvl) op typeCheckHostOp _ _ f (OtherOp op) = f op typeCheckHostOp _ _ _ (SizeOp op) = typeCheckSizeOp 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 @@ -1,31 +1,31 @@ -{-# OPTIONS_GHC -fno-warn-orphans #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ConstraintKinds #-} -module Futhark.IR.Kernels.Simplify - ( simplifyKernels - , simplifyLambda +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} - , Kernels +module Futhark.IR.Kernels.Simplify + ( simplifyKernels, + simplifyLambda, + Kernels, - -- * Building blocks - , simplifyKernelOp - ) + -- * Building blocks + simplifyKernelOp, + ) where +import qualified Futhark.Analysis.SymbolTable as ST import Futhark.IR.Kernels -import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Rules -import Futhark.Optimise.Simplify.Lore -import Futhark.MonadFreshNames -import Futhark.Tools -import Futhark.Pass import qualified Futhark.IR.SOACS.Simplify as SOAC +import Futhark.MonadFreshNames import qualified Futhark.Optimise.Simplify as Simplify +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Optimise.Simplify.Lore import Futhark.Optimise.Simplify.Rule -import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Optimise.Simplify.Rules +import Futhark.Pass +import Futhark.Tools import qualified Futhark.Transform.FirstOrderTransform as FOT simpleKernels :: Simplify.SimpleOps Kernels @@ -35,30 +35,35 @@ simplifyKernels = Simplify.simplifyProg simpleKernels kernelRules Simplify.noExtraHoistBlockers -simplifyLambda :: (HasScope Kernels m, MonadFreshNames m) => - Lambda Kernels -> m (Lambda Kernels) +simplifyLambda :: + (HasScope Kernels m, MonadFreshNames m) => + Lambda Kernels -> + m (Lambda Kernels) simplifyLambda = Simplify.simplifyLambda simpleKernels kernelRules Engine.noExtraHoistBlockers -simplifyKernelOp :: (Engine.SimplifiableLore lore, - BodyDec lore ~ ()) => - Simplify.SimplifyOp lore op - -> HostOp lore op - -> Engine.SimpleM lore (HostOp (Wise lore) (OpWithWisdom op), Stms (Wise lore)) - +simplifyKernelOp :: + ( Engine.SimplifiableLore lore, + BodyDec lore ~ () + ) => + Simplify.SimplifyOp lore op -> + HostOp lore op -> + Engine.SimpleM lore (HostOp (Wise lore) (OpWithWisdom op), Stms (Wise lore)) simplifyKernelOp f (OtherOp op) = do (op', stms) <- f op return (OtherOp op', stms) - simplifyKernelOp _ (SegOp op) = do (op', hoisted) <- simplifySegOp op return (SegOp op', hoisted) - simplifyKernelOp _ (SizeOp (SplitSpace o w i elems_per_thread)) = - (,) <$> (SizeOp <$> - (SplitSpace <$> Engine.simplify o <*> Engine.simplify w - <*> Engine.simplify i <*> Engine.simplify elems_per_thread)) - <*> pure mempty + (,) + <$> ( SizeOp + <$> ( SplitSpace <$> Engine.simplify o <*> Engine.simplify w + <*> Engine.simplify i + <*> Engine.simplify elems_per_thread + ) + ) + <*> pure mempty simplifyKernelOp _ (SizeOp (GetSize key size_class)) = return (SizeOp $ GetSize key size_class, mempty) simplifyKernelOp _ (SizeOp (GetSizeMax size_class)) = @@ -70,7 +75,7 @@ w' <- Engine.simplify w return (SizeOp $ CalcNumGroups w' max_num_groups group_size, mempty) -instance BinderOps (Wise Kernels) where +instance BinderOps (Wise Kernels) instance HasSegOp (Wise Kernels) where type SegOpLevel (Wise Kernels) = SegLevel @@ -84,13 +89,16 @@ soacOp = OtherOp kernelRules :: RuleBook (Wise Kernels) -kernelRules = standardRules <> segOpRules <> - ruleBook - [ RuleOp redomapIotaToLoop - , RuleOp SOAC.simplifyKnownIterationSOAC - , RuleOp SOAC.removeReplicateMapping ] - [ RuleBasicOp removeUnnecessaryCopy - , RuleOp SOAC.liftIdentityMapping ] +kernelRules = + standardRules <> segOpRules + <> ruleBook + [ RuleOp redomapIotaToLoop, + RuleOp SOAC.simplifyKnownIterationSOAC, + RuleOp SOAC.removeReplicateMapping + ] + [ RuleBasicOp removeUnnecessaryCopy, + RuleOp SOAC.liftIdentityMapping + ] -- We turn reductions over (solely) iotas into do-loops, because there -- is no useful structure here anyway. This is mostly a hack to work @@ -99,7 +107,7 @@ redomapIotaToLoop :: TopDownRuleOp (Wise Kernels) redomapIotaToLoop vtable pat aux (OtherOp soac@(Screma _ form [arr])) | Just _ <- isRedomapSOAC form, - Just (Iota{}, _) <- ST.lookupBasicOp arr vtable = - Simplify $ certifying (stmAuxCerts aux) $ FOT.transformSOAC pat soac + Just (Iota {}, _) <- ST.lookupBasicOp arr vtable = + Simplify $ certifying (stmAuxCerts aux) $ FOT.transformSOAC pat soac redomapIotaToLoop _ _ _ _ = Skip 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 @@ -1,24 +1,33 @@ +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Trustworthy #-} + -- | In the context of this module, a "size" is any kind of tunable -- (run-time) constant. module Futhark.IR.Kernels.Sizes - ( SizeClass (..) - , sizeDefault - , KernelPath - , Count(..) - , NumGroups, GroupSize, NumThreads + ( SizeClass (..), + sizeDefault, + KernelPath, + Count (..), + NumGroups, + GroupSize, + NumThreads, ) - where +where -import Data.Int (Int32) +import Control.Category +import Data.Int (Int64) import Data.Traversable - -import Futhark.Util.Pretty +import Futhark.IR.Prop.Names (FreeIn) import Futhark.Transform.Substitute -import Language.Futhark.Core (Name) import Futhark.Util.IntegralExp (IntegralExp) -import Futhark.IR.Prop.Names (FreeIn) +import Futhark.Util.Pretty +import GHC.Generics (Generic) +import Language.Futhark.Core (Name) +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) -- | An indication of which comparisons have been performed to get to -- this point, as well as the result of each comparison. @@ -26,22 +35,36 @@ -- | The class of some kind of configurable size. Each class may -- impose constraints on the valid values. -data SizeClass = SizeThreshold KernelPath (Maybe Int32) - -- ^ A threshold with an optional default. - | SizeGroup - | SizeNumGroups - | SizeTile - | SizeLocalMemory - -- ^ Likely not useful on its own, but querying the - -- maximum can be handy. - | SizeBespoke Name Int32 - -- ^ A bespoke size with a default. - deriving (Eq, Ord, Show) +data SizeClass + = -- | A threshold with an optional default. + SizeThreshold KernelPath (Maybe Int64) + | SizeGroup + | SizeNumGroups + | SizeTile + | -- | Likely not useful on its own, but querying the + -- maximum can be handy. + SizeLocalMemory + | -- | A bespoke size with a default. + SizeBespoke Name Int64 + deriving (Eq, Ord, Show, Generic) +instance SexpIso SizeClass where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "threshold") >>> Sexp.el sexpIso >>> props (Sexp.optKey "default" (iso fromIntegral fromIntegral . Sexp.int)))) $ + 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 + instance Pretty SizeClass where ppr (SizeThreshold path _) = text $ "threshold (" ++ unwords (map pStep path) ++ ")" - where pStep (v, True) = pretty v - pStep (v, False) = '!' : pretty v + where + pStep (v, True) = pretty v + pStep (v, False) = '!' : pretty v ppr SizeGroup = text "group_size" ppr SizeNumGroups = text "num_groups" ppr SizeTile = text "tile_size" @@ -49,14 +72,17 @@ ppr (SizeBespoke k _) = ppr k -- | The default value for the size. If 'Nothing', that means the backend gets to decide. -sizeDefault :: SizeClass -> Maybe Int32 +sizeDefault :: SizeClass -> Maybe Int64 sizeDefault (SizeThreshold _ x) = x sizeDefault (SizeBespoke _ x) = Just x sizeDefault _ = Nothing -- | A wrapper supporting a phantom type for indicating what we are counting. -newtype Count u e = Count { unCount :: e } - deriving (Eq, Ord, Show, Num, IntegralExp, FreeIn, Pretty, Substitute) +newtype Count u e = Count {unCount :: e} + deriving (Eq, Ord, Show, Num, IntegralExp, FreeIn, Pretty, Substitute, Generic) + +instance SexpIso e => SexpIso (Count u e) where + sexpIso = with $ \count -> sexpIso >>> count instance Functor (Count u) where fmap = fmapDefault diff --git a/src/Futhark/IR/KernelsMem.hs b/src/Futhark/IR/KernelsMem.hs --- a/src/Futhark/IR/KernelsMem.hs +++ b/src/Futhark/IR/KernelsMem.hs @@ -1,47 +1,44 @@ -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE TypeFamilies #-} + module Futhark.IR.KernelsMem - ( KernelsMem + ( KernelsMem, - -- * Simplification - , simplifyProg - , simplifyStms - , simpleKernelsMem + -- * Simplification + simplifyProg, + simplifyStms, + simpleKernelsMem, -- * Module re-exports - , module Futhark.IR.Mem - , module Futhark.IR.Kernels.Kernel + module Futhark.IR.Mem, + module Futhark.IR.Kernels.Kernel, ) - where +where import Futhark.Analysis.PrimExp.Convert import qualified Futhark.Analysis.UsageTable as UT -import Futhark.MonadFreshNames -import Futhark.Pass -import Futhark.IR.Syntax -import Futhark.IR.Prop -import Futhark.IR.Traversals -import Futhark.IR.Pretty import Futhark.IR.Kernels.Kernel import Futhark.IR.Kernels.Simplify (simplifyKernelOp) -import qualified Futhark.TypeCheck as TC import Futhark.IR.Mem import Futhark.IR.Mem.Simplify -import Futhark.Pass.ExplicitAllocations (BinderOps(..), mkLetNamesB', mkLetNamesB'') +import Futhark.MonadFreshNames import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Pass +import Futhark.Pass.ExplicitAllocations (BinderOps (..), mkLetNamesB', mkLetNamesB'') +import qualified Futhark.TypeCheck as TC data KernelsMem instance Decorations KernelsMem where - type LetDec KernelsMem = LetDecMem + type LetDec KernelsMem = LetDecMem type FParamInfo KernelsMem = FParamMem type LParamInfo KernelsMem = LParamMem - type RetType KernelsMem = RetTypeMem + type RetType KernelsMem = RetTypeMem type BranchType KernelsMem = BranchTypeMem - type Op KernelsMem = MemOp (HostOp KernelsMem ()) + type Op KernelsMem = MemOp (HostOp KernelsMem ()) instance ASTLore KernelsMem where expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern @@ -52,14 +49,15 @@ opReturns (Inner (SegOp op)) = segOpReturns op opReturns k = extReturns <$> opType k -instance PrettyLore KernelsMem where +instance PrettyLore KernelsMem instance TC.CheckableOp KernelsMem where checkOp = typeCheckMemoryOp Nothing - where typeCheckMemoryOp _ (Alloc size _) = - TC.require [Prim int64] size - typeCheckMemoryOp lvl (Inner op) = - typeCheckHostOp (typeCheckMemoryOp . Just) lvl (const $ return ()) op + where + typeCheckMemoryOp _ (Alloc size _) = + TC.require [Prim int64] size + typeCheckMemoryOp lvl (Inner op) = + typeCheckHostOp (typeCheckMemoryOp . Just) lvl (const $ return ()) op instance TC.Checkable KernelsMem where checkFParamLore = checkMemInfo @@ -85,10 +83,13 @@ simplifyProg :: Prog KernelsMem -> PassM (Prog KernelsMem) simplifyProg = simplifyProgGeneric simpleKernelsMem -simplifyStms :: (HasScope KernelsMem m, MonadFreshNames m) => - Stms KernelsMem - -> m (Engine.SymbolTable (Engine.Wise KernelsMem), - Stms KernelsMem) +simplifyStms :: + (HasScope KernelsMem m, MonadFreshNames m) => + Stms KernelsMem -> + m + ( Engine.SymbolTable (Engine.Wise KernelsMem), + Stms KernelsMem + ) simplifyStms = simplifyStmsGeneric simpleKernelsMem simpleKernelsMem :: Engine.SimpleOps KernelsMem @@ -100,7 +101,7 @@ -- usages for those sizes. This is necessary so the simplifier -- will hoist those sizes out as far as possible (most -- importantly, past the versioning If). - usage (SegOp (SegMap SegGroup{} _ _ kbody)) = localAllocs kbody + usage (SegOp (SegMap SegGroup {} _ _ kbody)) = localAllocs kbody usage _ = mempty localAllocs = foldMap stmLocalAlloc . kernelBodyStms stmLocalAlloc = expLocalAlloc . stmExp 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 @@ -1,1003 +1,1230 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ConstraintKinds #-} --- | Building blocks for defining representations where every array --- is given information about which memory block is it based in, and --- how array elements map to memory block offsets. --- --- There are two primary concepts you will need to understand: --- --- 1. Memory blocks, which are Futhark values of type v'Mem' --- (parametrized with their size). These correspond to arbitrary --- blocks of memory, and are created using the 'Alloc' operation. --- --- 2. Index functions, which describe a mapping from the index space --- of an array (eg. a two-dimensional space for an array of type --- @[[int]]@) to a one-dimensional offset into a memory block. --- Thus, index functions describe how arbitrary-dimensional arrays --- are mapped to the single-dimensional world of memory. --- --- At a conceptual level, imagine that we have a two-dimensional array --- @a@ of 32-bit integers, consisting of @n@ rows of @m@ elements --- each. This array could be represented in classic row-major format --- with an index function like the following: --- --- @ --- f(i,j) = i * m + j --- @ --- --- When we want to know the location of element @a[2,3]@, we simply --- call the index function as @f(2,3)@ and obtain @2*m+3@. We could --- also have chosen another index function, one that represents the --- array in column-major (or "transposed") format: --- --- @ --- f(i,j) = j * n + i --- @ --- --- Index functions are not Futhark-level functions, but a special --- construct that the final code generator will eventually use to --- generate concrete access code. By modifying the index functions we --- can change how an array is represented in memory, which can permit --- memory access pattern optimisations. --- --- Every time we bind an array, whether in a @let@-binding, @loop@ --- merge parameter, or @lambda@ parameter, we have an annotation --- specifying a memory block and an index function. In some cases, --- such as @let@-bindings for many expressions, we are free to specify --- an arbitrary index function and memory block - for example, we get --- to decide where 'Copy' stores its result - but in other cases the --- type rules of the expression chooses for us. For example, 'Index' --- always produces an array in the same memory block as its input, and --- with the same index function, except with some indices fixed. -module Futhark.IR.Mem - ( LetDecMem - , FParamMem - , LParamMem - , RetTypeMem - , BranchTypeMem - - , MemOp (..) - , MemInfo (..) - , MemBound - , MemBind (..) - , MemReturn (..) - , IxFun - , ExtIxFun - , isStaticIxFun - , ExpReturns - , BodyReturns - , FunReturns - , noUniquenessReturns - , bodyReturnsToExpReturns - , Mem - , AllocOp(..) - , OpReturns(..) - , varReturns - , expReturns - , extReturns - , lookupMemInfo - , subExpMemInfo - , lookupArraySummary - , existentialiseIxFun - - -- * Type checking parts - , matchBranchReturnType - , matchPatternToExp - , matchFunctionReturnType - , matchLoopResultMem - , bodyReturnsFromPattern - , checkMemInfo - - -- * Module re-exports - , module Futhark.IR.Prop - , module Futhark.IR.Traversals - , module Futhark.IR.Pretty - , module Futhark.IR.Syntax - , module Futhark.Analysis.PrimExp.Convert - ) -where - -import Data.Maybe -import Control.Monad.State -import Control.Monad.Reader -import Control.Monad.Except -import qualified Data.Map.Strict as M -import Data.Foldable (traverse_, toList) -import Data.List (elemIndex, find) -import qualified Data.Set as S - -import Futhark.Analysis.Metrics -import Futhark.IR.Syntax -import Futhark.IR.Prop -import Futhark.IR.Prop.Aliases -import Futhark.IR.Traversals -import Futhark.IR.Pretty -import Futhark.Transform.Rename -import Futhark.Transform.Substitute -import qualified Futhark.TypeCheck as TC -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Analysis.PrimExp.Convert -import Futhark.Analysis.PrimExp.Simplify -import Futhark.Util -import qualified Futhark.Util.Pretty as PP -import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore -import Futhark.IR.Aliases - (Aliases, removeScopeAliases, removeExpAliases, removePatternAliases) -import qualified Futhark.Analysis.SymbolTable as ST - -type LetDecMem = MemInfo SubExp NoUniqueness MemBind -type FParamMem = MemInfo SubExp Uniqueness MemBind -type LParamMem = MemInfo SubExp NoUniqueness MemBind -type RetTypeMem = FunReturns -type BranchTypeMem = BodyReturns - --- | The class of ops that have memory allocation. -class AllocOp op where - allocOp :: SubExp -> Space -> op - -type Mem lore = (AllocOp (Op lore), - FParamInfo lore ~ FParamMem, - LParamInfo lore ~ LParamMem, - LetDec lore ~ LetDecMem, - RetType lore ~ RetTypeMem, - BranchType lore ~ BranchTypeMem, - CanBeAliased (Op lore), - ASTLore lore, Decorations lore, - TC.Checkable lore, - OpReturns lore) - -instance IsRetType FunReturns where - primRetType = MemPrim - applyRetType = applyFunReturns - -instance IsBodyType BodyReturns where - primBodyType = MemPrim - -data MemOp inner = Alloc SubExp Space - -- ^ Allocate a memory block. This really should not be an - -- expression, but what are you gonna do... - | Inner inner - deriving (Eq, Ord, Show) - -instance AllocOp (MemOp inner) where - allocOp = Alloc - -instance FreeIn inner => FreeIn (MemOp inner) where - freeIn' (Alloc size _) = freeIn' size - freeIn' (Inner k) = freeIn' k - -instance TypedOp inner => TypedOp (MemOp inner) where - opType (Alloc _ space) = pure [Mem space] - opType (Inner k) = opType k - -instance AliasedOp inner => AliasedOp (MemOp inner) where - opAliases Alloc{} = [mempty] - opAliases (Inner k) = opAliases k - - consumedInOp Alloc{} = mempty - consumedInOp (Inner k) = consumedInOp k - -instance CanBeAliased inner => CanBeAliased (MemOp inner) where - type OpWithAliases (MemOp inner) = MemOp (OpWithAliases inner) - 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 - -instance Rename inner => Rename (MemOp inner) where - rename (Alloc size space) = Alloc <$> rename size <*> pure space - rename (Inner k) = Inner <$> rename k - -instance Substitute inner => Substitute (MemOp inner) where - substituteNames subst (Alloc size space) = Alloc (substituteNames subst size) space - substituteNames subst (Inner k) = Inner $ substituteNames subst k - -instance PP.Pretty inner => PP.Pretty (MemOp inner) where - ppr (Alloc e DefaultSpace) = PP.text "alloc" <> PP.apply [PP.ppr e] - ppr (Alloc e s) = PP.text "alloc" <> PP.apply [PP.ppr e, PP.ppr s] - ppr (Inner k) = PP.ppr k - -instance OpMetrics inner => OpMetrics (MemOp inner) where - opMetrics Alloc{} = seen "Alloc" - opMetrics (Inner k) = opMetrics k - -instance IsOp inner => IsOp (MemOp inner) where - safeOp Alloc{} = False - safeOp (Inner k) = safeOp k - cheapOp (Inner k) = cheapOp k - cheapOp Alloc{} = True - -instance CanBeWise inner => CanBeWise (MemOp inner) where - type OpWithWisdom (MemOp inner) = MemOp (OpWithWisdom inner) - removeOpWisdom (Alloc size space) = Alloc size space - removeOpWisdom (Inner k) = Inner $ removeOpWisdom k - -instance ST.IndexOp inner => ST.IndexOp (MemOp inner) where - indexOp vtable k (Inner op) is = ST.indexOp vtable k op is - indexOp _ _ _ _ = Nothing - --- | The index function representation used for memory annotations. -type IxFun = IxFun.IxFun (PrimExp VName) - --- | An index function that may contain existential variables. -type ExtIxFun = IxFun.IxFun (PrimExp (Ext VName)) - --- | A summary of the memory information for every let-bound --- identifier, function parameter, and return value. Parameterisered --- over uniqueness, dimension, and auxiliary array information. -data MemInfo d u ret = MemPrim PrimType - -- ^ A primitive value. - | MemMem Space - -- ^ A memory block. - | MemArray PrimType (ShapeBase d) u ret - -- ^ The array is stored in the named memory block, - -- and with the given index function. The index - -- function maps indices in the array to /element/ - -- offset, /not/ byte offsets! To translate to byte - -- offsets, multiply the offset with the size of the - -- array element type. - deriving (Eq, Show, Ord) --- XXX Ord? - -type MemBound u = MemInfo SubExp u MemBind - -instance FixExt ret => DeclExtTyped (MemInfo ExtSize Uniqueness ret) where - declExtTypeOf (MemPrim pt) = Prim pt - declExtTypeOf (MemMem space) = Mem space - declExtTypeOf (MemArray pt shape u _) = Array pt shape u - -instance FixExt ret => ExtTyped (MemInfo ExtSize NoUniqueness ret) where - extTypeOf (MemPrim pt) = Prim pt - extTypeOf (MemMem space) = Mem space - extTypeOf (MemArray pt shape u _) = Array pt shape u - -instance FixExt ret => FixExt (MemInfo ExtSize u ret) where - fixExt _ _ (MemPrim pt) = MemPrim pt - fixExt _ _ (MemMem space) = MemMem space - fixExt i se (MemArray pt shape u ret) = - MemArray pt (fixExt i se shape) u (fixExt i se ret) - -instance Typed (MemInfo SubExp Uniqueness ret) where - typeOf = fromDecl . declTypeOf - -instance Typed (MemInfo SubExp NoUniqueness ret) where - typeOf (MemPrim pt) = Prim pt - typeOf (MemMem space) = Mem space - typeOf (MemArray bt shape u _) = Array bt shape u - -instance DeclTyped (MemInfo SubExp Uniqueness ret) where - declTypeOf (MemPrim bt) = Prim bt - declTypeOf (MemMem space) = Mem space - declTypeOf (MemArray bt shape u _) = Array bt shape u - -instance (FreeIn d, FreeIn ret) => FreeIn (MemInfo d u ret) where - freeIn' (MemArray _ shape _ ret) = freeIn' shape <> freeIn' ret - freeIn' (MemMem s) = freeIn' s - freeIn' MemPrim{} = mempty - -instance (Substitute d, Substitute ret) => Substitute (MemInfo d u ret) where - substituteNames subst (MemArray bt shape u ret) = - MemArray bt - (substituteNames subst shape) u - (substituteNames subst ret) - substituteNames _ (MemMem space) = - MemMem space - substituteNames _ (MemPrim bt) = - MemPrim bt - -instance (Substitute d, Substitute ret) => Rename (MemInfo d u ret) where - rename = substituteRename - -simplifyIxFun :: Engine.SimplifiableLore lore => - IxFun -> Engine.SimpleM lore IxFun -simplifyIxFun = traverse simplifyPrimExp - -simplifyExtIxFun :: Engine.SimplifiableLore lore => - ExtIxFun -> Engine.SimpleM lore ExtIxFun -simplifyExtIxFun = traverse simplifyExtPrimExp - -isStaticIxFun :: ExtIxFun -> Maybe IxFun -isStaticIxFun = traverse $ traverse inst - where inst Ext{} = Nothing - inst (Free x) = Just x - -instance (Engine.Simplifiable d, Engine.Simplifiable ret) => - Engine.Simplifiable (MemInfo d u ret) where - simplify (MemPrim bt) = - return $ MemPrim bt - simplify (MemMem space) = - pure $ MemMem space - simplify (MemArray bt shape u ret) = - MemArray bt <$> Engine.simplify shape <*> pure u <*> Engine.simplify ret - -instance (PP.Pretty (TypeBase (ShapeBase d) u), - PP.Pretty d, PP.Pretty u, PP.Pretty ret) => PP.Pretty (MemInfo d u ret) where - ppr (MemPrim bt) = PP.ppr bt - ppr (MemMem DefaultSpace) = PP.text "mem" - ppr (MemMem s) = PP.text "mem" <> PP.ppr s - ppr (MemArray bt shape u ret) = - PP.ppr (Array bt shape u) <> PP.text "@" <> PP.ppr ret - -instance PP.Pretty (Param (MemInfo SubExp Uniqueness ret)) where - ppr = PP.ppr . fmap declTypeOf - -instance PP.Pretty (Param (MemInfo SubExp NoUniqueness ret)) where - ppr = PP.ppr . fmap typeOf - -instance PP.Pretty (PatElemT (MemInfo SubExp NoUniqueness ret)) where - ppr = PP.ppr . fmap typeOf - --- | Memory information for an array bound somewhere in the program. -data MemBind = ArrayIn VName IxFun - -- ^ Located in this memory block with this index - -- function. - deriving (Show) - -instance Eq MemBind where - _ == _ = True - -instance Ord MemBind where - _ `compare` _ = EQ - -instance Rename MemBind where - rename = substituteRename - -instance Substitute MemBind where - substituteNames substs (ArrayIn ident ixfun) = - ArrayIn (substituteNames substs ident) (substituteNames substs ixfun) - -instance PP.Pretty MemBind where - ppr (ArrayIn mem ixfun) = - PP.text "@" <> PP.ppr mem <> PP.text "->" PP. PP.ppr ixfun - -instance FreeIn MemBind where - freeIn' (ArrayIn mem ixfun) = freeIn' mem <> freeIn' ixfun - --- | A description of the memory properties of an array being returned --- by an operation. -data MemReturn = ReturnsInBlock VName ExtIxFun - -- ^ The array is located in a memory block that is - -- already in scope. - | ReturnsNewBlock Space Int ExtIxFun - -- ^ The operation returns a new (existential) memory - -- block. - deriving (Show) - -instance Eq MemReturn where - _ == _ = True - -instance Ord MemReturn where - _ `compare` _ = EQ - -instance Rename MemReturn where - rename = substituteRename - -instance Substitute MemReturn where - substituteNames substs (ReturnsInBlock ident ixfun) = - ReturnsInBlock (substituteNames substs ident) (substituteNames substs ixfun) - substituteNames substs (ReturnsNewBlock space i ixfun) = - ReturnsNewBlock space i (substituteNames substs ixfun) - -instance FixExt MemReturn where - fixExt i (Var v) (ReturnsNewBlock _ j ixfun) - | j == i = ReturnsInBlock v $ fixExtIxFun i - (primExpFromSubExp int32 (Var v)) ixfun - fixExt i se (ReturnsNewBlock space j ixfun) = - ReturnsNewBlock space j' - (fixExtIxFun i (primExpFromSubExp int32 se) ixfun) - where j' | i < j = j-1 - | otherwise = j - fixExt i se (ReturnsInBlock mem ixfun) = - ReturnsInBlock mem (fixExtIxFun i (primExpFromSubExp int32 se) ixfun) - -fixExtIxFun :: Int -> PrimExp VName -> ExtIxFun -> ExtIxFun -fixExtIxFun i e = fmap $ replaceInPrimExp update - where update (Ext j) t | j > i = LeafExp (Ext $ j - 1) t - | j == i = fmap Free e - | otherwise = LeafExp (Ext j) t - update (Free x) t = LeafExp (Free x) t - -leafExp :: Int -> PrimExp (Ext a) -leafExp i = LeafExp (Ext i) int32 - -existentialiseIxFun :: [VName] -> IxFun -> ExtIxFun -existentialiseIxFun ctx = IxFun.substituteInIxFun ctx' . fmap (fmap Free) - where ctx' = M.map leafExp $ M.fromList $ zip (map Free ctx) [0..] - -instance PP.Pretty MemReturn where - ppr (ReturnsInBlock v ixfun) = - PP.parens $ PP.text (pretty v) <> PP.text "->" PP. PP.ppr ixfun - ppr (ReturnsNewBlock space i ixfun) = - PP.text ("?" ++ show i) <> PP.ppr space <> PP.text "->" PP. PP.ppr ixfun - -instance FreeIn MemReturn where - freeIn' (ReturnsInBlock v ixfun) = freeIn' v <> freeIn' ixfun - freeIn' (ReturnsNewBlock space _ ixfun) = freeIn' space <> freeIn' ixfun - -instance Engine.Simplifiable MemReturn where - simplify (ReturnsNewBlock space i ixfun) = - ReturnsNewBlock space i <$> simplifyExtIxFun ixfun - simplify (ReturnsInBlock v ixfun) = - ReturnsInBlock <$> Engine.simplify v <*> simplifyExtIxFun ixfun - - -instance Engine.Simplifiable MemBind where - simplify (ArrayIn mem ixfun) = - ArrayIn <$> Engine.simplify mem <*> simplifyIxFun ixfun - -instance Engine.Simplifiable [FunReturns] where - simplify = mapM Engine.simplify - --- | The memory return of an expression. An array is annotated with --- @Maybe MemReturn@, which can be interpreted as the expression --- either dictating exactly where the array is located when it is --- returned (if 'Just'), or able to put it whereever the binding --- prefers (if 'Nothing'). --- --- This is necessary to capture the difference between an expression --- that is just an array-typed variable, in which the array being --- "returned" is located where it already is, and a @copy@ expression, --- whose entire purpose is to store an existing array in some --- arbitrary location. This is a consequence of the design decision --- never to have implicit memory copies. -type ExpReturns = MemInfo ExtSize NoUniqueness (Maybe MemReturn) - --- | The return of a body, which must always indicate where --- returned arrays are located. -type BodyReturns = MemInfo ExtSize NoUniqueness MemReturn - --- | The memory return of a function, which must always indicate where --- returned arrays are located. -type FunReturns = MemInfo ExtSize Uniqueness MemReturn - -maybeReturns :: MemInfo d u r -> MemInfo d u (Maybe r) -maybeReturns (MemArray bt shape u ret) = - MemArray bt shape u $ Just ret -maybeReturns (MemPrim bt) = - MemPrim bt -maybeReturns (MemMem space) = - MemMem space - -noUniquenessReturns :: MemInfo d u r -> MemInfo d NoUniqueness r -noUniquenessReturns (MemArray bt shape _ r) = - MemArray bt shape NoUniqueness r -noUniquenessReturns (MemPrim bt) = - MemPrim bt -noUniquenessReturns (MemMem space) = - MemMem space - -funReturnsToExpReturns :: FunReturns -> ExpReturns -funReturnsToExpReturns = noUniquenessReturns . maybeReturns - -bodyReturnsToExpReturns :: BodyReturns -> ExpReturns -bodyReturnsToExpReturns = noUniquenessReturns . maybeReturns - -matchRetTypeToResult :: Mem lore => - [FunReturns] -> Result -> TC.TypeM lore () -matchRetTypeToResult rettype result = do - scope <- askScope - result_ts <- runReaderT (mapM subExpMemInfo result) $ removeScopeAliases scope - matchReturnType rettype result result_ts - -matchFunctionReturnType :: Mem lore => - [FunReturns] -> Result -> TC.TypeM lore () -matchFunctionReturnType rettype result = do - matchRetTypeToResult rettype result - mapM_ checkResultSubExp result - where checkResultSubExp Constant{} = - return () - checkResultSubExp (Var v) = do - dec <- varMemInfo v - case dec of - MemPrim _ -> return () - MemMem{} -> return () - MemArray _ _ _ (ArrayIn _ ixfun) - | IxFun.isLinear ixfun -> - return () - | otherwise -> - TC.bad $ TC.TypeError $ - "Array " ++ pretty v ++ - " returned by function, but has nontrivial index function " ++ - pretty ixfun - -matchLoopResultMem :: Mem lore => - [FParam (Aliases lore)] -> [FParam (Aliases lore)] - -> [SubExp] -> TC.TypeM lore () -matchLoopResultMem ctx val = matchRetTypeToResult rettype - where ctx_names = map paramName ctx - - -- Invent a ReturnType so we can pretend that the loop body is - -- actually returning from a function. - rettype = map (toRet . paramDec) val - - toExtV v - | Just i <- v `elemIndex` ctx_names = Ext i - | otherwise = Free v - - toExtSE (Var v) = Var <$> toExtV v - toExtSE (Constant v) = Free $ Constant v - - toRet (MemPrim t) = - MemPrim t - toRet (MemMem space) = - MemMem space - toRet (MemArray pt shape u (ArrayIn mem ixfun)) - | Just i <- mem `elemIndex` ctx_names, - Param _ (MemMem space) : _ <- drop i ctx = - MemArray pt shape' u $ ReturnsNewBlock space i ixfun' - | otherwise = - MemArray pt shape' u $ ReturnsInBlock mem ixfun' - where shape' = fmap toExtSE shape - ixfun' = existentialiseIxFun ctx_names ixfun - -matchBranchReturnType :: Mem lore => - [BodyReturns] - -> Body (Aliases lore) - -> TC.TypeM lore () -matchBranchReturnType rettype (Body _ stms res) = do - scope <- askScope - ts <- runReaderT (mapM subExpMemInfo res) $ removeScopeAliases (scope <> scopeOf stms) - matchReturnType rettype res ts - --- | Helper function for index function unification. --- --- The first return value maps a VName (wrapped in 'Free') to its Int --- (wrapped in 'Ext'). In case of duplicates, it is mapped to the --- *first* Int that occurs. --- --- The second return value maps each Int (wrapped in an 'Ext') to a --- 'LeafExp' 'Ext' with the Int at which its associated VName first --- occurs. -getExtMaps :: [(VName,Int)] -> (M.Map (Ext VName) (PrimExp (Ext VName)), - M.Map (Ext VName) (PrimExp (Ext VName))) -getExtMaps ctx_lst_ids = - (M.map leafExp $ M.mapKeys Free $ M.fromListWith (flip const) ctx_lst_ids, - M.fromList $ - mapMaybe (traverse (fmap (\i -> LeafExp (Ext i) int32) . - (`lookup` ctx_lst_ids)) . - uncurry (flip (,)) . fmap Ext) ctx_lst_ids) - -matchReturnType :: PP.Pretty u => - [MemInfo ExtSize u MemReturn] - -> [SubExp] - -> [MemInfo SubExp NoUniqueness MemBind] - -> TC.TypeM lore () -matchReturnType rettype res ts = do - let (ctx_ts, val_ts) = splitFromEnd (length rettype) ts - (ctx_res, _val_res) = splitFromEnd (length rettype) res - - existentialiseIxFun0 :: IxFun -> ExtIxFun - existentialiseIxFun0 = fmap $ fmap Free - - fetchCtx i = case maybeNth i $ zip ctx_res ctx_ts of - Nothing -> throwError $ "Cannot find context variable " ++ - show i ++ " in context results: " ++ pretty ctx_res - Just (se, t) -> return (se, t) - - checkReturn (MemPrim x) (MemPrim y) - | x == y = return () - checkReturn (MemMem x) (MemMem y) - | x == y = return () - checkReturn (MemArray x_pt x_shape _ x_ret) - (MemArray y_pt y_shape _ y_ret) - | x_pt == y_pt, shapeRank x_shape == shapeRank y_shape = do - zipWithM_ checkDim (shapeDims x_shape) (shapeDims y_shape) - checkMemReturn x_ret y_ret - checkReturn x y = - throwError $ unwords ["Expected ", pretty x, " but got ", pretty y] - - checkDim (Free x) y - | x == y = return () - | otherwise = throwError $ unwords ["Expected dim", pretty x, - "but got", pretty y] - checkDim (Ext i) y = do - (x, _) <- fetchCtx i - unless (x == y) $ - throwError $ unwords ["Expected ext dim", pretty i, "=>", pretty x, - "but got", pretty y] - - extsInMemInfo :: MemInfo ExtSize u MemReturn -> S.Set Int - extsInMemInfo (MemArray _ shp _ ret) = - extInShape shp <> extInMemReturn ret - extsInMemInfo _ = S.empty - - checkMemReturn (ReturnsInBlock x_mem x_ixfun) (ArrayIn y_mem y_ixfun) - | x_mem == y_mem = - unless (IxFun.closeEnough x_ixfun $ existentialiseIxFun0 y_ixfun) $ - throwError $ unwords ["Index function unification failed (ReturnsInBlock)", - "\nixfun of body result: ", pretty y_ixfun, - "\nixfun of return type: ", pretty x_ixfun, - "\nand context elements: ", pretty ctx_res] - checkMemReturn (ReturnsNewBlock x_space x_ext x_ixfun) - (ArrayIn y_mem y_ixfun) = do - (x_mem, x_mem_type) <- fetchCtx x_ext - unless (IxFun.closeEnough x_ixfun $ existentialiseIxFun0 y_ixfun) $ - throwError $ unwords ["Index function unification failed (ReturnsNewBlock)", - "\nixfun of body result: ", pretty y_ixfun, - "\nixfun of return type: ", pretty x_ixfun, - "\nand context elements: ", pretty ctx_res] - case x_mem_type of - MemMem y_space -> - unless (x_space == y_space) $ - throwError $ unwords ["Expected memory", pretty y_mem, "in space", pretty x_space, - "but actually in space", pretty y_space] - t -> - throwError $ unwords ["Expected memory", pretty x_ext, "=>", pretty x_mem, - "but but has type", pretty t] - checkMemReturn x y = - throwError $ unwords ["Expected array in", pretty x, - "but array returned in", pretty y] - - bad :: String -> TC.TypeM lore a - bad s = TC.bad $ TC.TypeError $ PP.pretty $ - "Return type" PP. - PP.indent 2 (ppTuple' rettype) PP. - "cannot match returns of results" PP. - PP.indent 2 (ppTuple' ts) PP. - PP.text s - - unless (length (S.unions $ map extsInMemInfo rettype) == length ctx_res) $ - TC.bad $ TC.TypeError $ "Too many context parameters for the number of " ++ - "existentials in the return type! type:\n " ++ - prettyTuple rettype ++ - "\ncannot match context parameters:\n " ++ prettyTuple ctx_res - - - either bad return =<< runExceptT (zipWithM_ checkReturn rettype val_ts) - -matchPatternToExp :: (Mem lore) => - Pattern (Aliases lore) - -> Exp (Aliases lore) - -> TC.TypeM lore () -matchPatternToExp pat e = do - scope <- asksScope removeScopeAliases - rt <- runReaderT (expReturns $ removeExpAliases e) scope - - let (ctxs, vals) = bodyReturnsFromPattern $ removePatternAliases pat - (ctx_ids, _ctx_ts) = unzip ctxs - (_val_ids, val_ts) = unzip vals - (ctx_map_ids, ctx_map_exts) = - getExtMaps $ zip ctx_ids [0..length ctx_ids - 1] - - let rt_exts = foldMap extInExpReturns rt - - unless (length val_ts == length rt && - and (zipWith (matches ctx_map_ids ctx_map_exts) val_ts rt) && - M.keysSet ctx_map_exts `S.isSubsetOf` S.map Ext rt_exts) $ - TC.bad $ TC.TypeError $ "Expression type:\n " ++ prettyTuple rt ++ - "\ncannot match pattern type:\n " ++ prettyTuple val_ts ++ - "\nwith context elements: " ++ pretty ctx_ids - where matches _ _ (MemPrim x) (MemPrim y) = x == y - matches _ _ (MemMem x_space) (MemMem y_space) = - x_space == y_space - matches ctxids ctxexts (MemArray x_pt x_shape _ x_ret) (MemArray y_pt y_shape _ y_ret) = - x_pt == y_pt && x_shape == y_shape && - case (x_ret, y_ret) of - (ReturnsInBlock _ x_ixfun, Just (ReturnsInBlock _ y_ixfun)) -> - let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun - y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun - in IxFun.closeEnough x_ixfun' y_ixfun' - (ReturnsInBlock _ x_ixfun, - Just (ReturnsNewBlock _ _ y_ixfun)) -> - let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun - y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun - in IxFun.closeEnough x_ixfun' y_ixfun' - (ReturnsNewBlock _ x_i x_ixfun, - Just (ReturnsNewBlock _ y_i y_ixfun)) -> - let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun - y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun - in x_i == y_i && IxFun.closeEnough x_ixfun' y_ixfun' - (_, Nothing) -> True - _ -> False - matches _ _ _ _ = False - - extInExpReturns :: ExpReturns -> S.Set Int - extInExpReturns (MemArray _ shape _ mem_return) = - extInShape shape <> maybe S.empty extInMemReturn mem_return - extInExpReturns _ = mempty - - -extInShape :: ShapeBase (Ext SubExp) -> S.Set Int -extInShape shape = S.fromList $ mapMaybe isExt $ shapeDims shape - -extInMemReturn :: MemReturn -> S.Set Int -extInMemReturn (ReturnsInBlock _ extixfn) = extInIxFn extixfn -extInMemReturn (ReturnsNewBlock _ i extixfn) = - S.singleton i <> extInIxFn extixfn - -extInIxFn :: ExtIxFun -> S.Set Int -extInIxFn ixfun = S.fromList $ concatMap (mapMaybe isExt . toList) ixfun - -varMemInfo :: Mem lore => - VName -> TC.TypeM lore (MemInfo SubExp NoUniqueness MemBind) -varMemInfo name = do - dec <- TC.lookupVar name - - case dec of - LetName (_, summary) -> return summary - FParamName summary -> return $ noUniquenessReturns summary - LParamName summary -> return summary - IndexName it -> return $ MemPrim $ IntType it - -nameInfoToMemInfo :: Mem lore => NameInfo lore -> MemBound NoUniqueness -nameInfoToMemInfo info = - case info of - FParamName summary -> noUniquenessReturns summary - LParamName summary -> summary - LetName summary -> summary - IndexName it -> MemPrim $ IntType it - -lookupMemInfo :: (HasScope lore m, Mem lore) => - VName -> m (MemInfo SubExp NoUniqueness MemBind) -lookupMemInfo = fmap nameInfoToMemInfo . lookupInfo - -subExpMemInfo :: (HasScope lore m, Monad m, Mem lore) => - SubExp -> m (MemInfo SubExp NoUniqueness MemBind) -subExpMemInfo (Var v) = lookupMemInfo v -subExpMemInfo (Constant v) = return $ MemPrim $ primValueType v - -lookupArraySummary :: (Mem lore, HasScope lore m, Monad m) => - VName -> m (VName, IxFun.IxFun (PrimExp VName)) -lookupArraySummary name = do - summary <- lookupMemInfo name - case summary of - MemArray _ _ _ (ArrayIn mem ixfun) -> - return (mem, ixfun) - _ -> - error $ "Variable " ++ pretty name ++ " does not look like an array." - -checkMemInfo :: TC.Checkable lore => - VName -> MemInfo SubExp u MemBind - -> TC.TypeM lore () -checkMemInfo _ (MemPrim _) = return () -checkMemInfo _ (MemMem (ScalarSpace d _)) = mapM_ (TC.require [Prim int32]) d -checkMemInfo _ (MemMem _) = return () -checkMemInfo name (MemArray _ shape _ (ArrayIn v ixfun)) = do - t <- lookupType v - case t of - Mem{} -> - return () - _ -> - TC.bad $ TC.TypeError $ - "Variable " ++ pretty v ++ - " used as memory block, but is of type " ++ - pretty t ++ "." - - TC.context ("in index function " ++ pretty ixfun) $ do - traverse_ (TC.requirePrimExp int32) ixfun - let ixfun_rank = IxFun.rank ixfun - ident_rank = shapeRank shape - unless (ixfun_rank == ident_rank) $ - TC.bad $ TC.TypeError $ - "Arity of index function (" ++ pretty ixfun_rank ++ - ") does not match rank of array " ++ pretty name ++ - " (" ++ show ident_rank ++ ")" - -bodyReturnsFromPattern :: PatternT (MemBound NoUniqueness) - -> ([(VName,BodyReturns)], [(VName,BodyReturns)]) -bodyReturnsFromPattern pat = - (map asReturns $ patternContextElements pat, - map asReturns $ patternValueElements pat) - where ctx = patternContextElements pat - - ext (Var v) - | Just (i, _) <- find ((==v) . patElemName . snd) $ zip [0..] ctx = - Ext i - ext se = Free se - - asReturns pe = - (patElemName pe, - case patElemDec pe of - MemPrim pt -> MemPrim pt - MemMem space -> MemMem space - MemArray pt shape u (ArrayIn mem ixfun) -> - MemArray pt (Shape $ map ext $ shapeDims shape) u $ - case find ((==mem) . patElemName . snd) $ zip [0..] ctx of - Just (i, PatElem _ (MemMem space)) -> - ReturnsNewBlock space i $ - existentialiseIxFun (map patElemName ctx) ixfun - _ -> ReturnsInBlock mem $ existentialiseIxFun [] ixfun - ) - -instance (PP.Pretty u, PP.Pretty r) => PrettyAnnot (PatElemT (MemInfo SubExp u r)) where - ppAnnot = bindeeAnnot patElemName patElemDec - -instance (PP.Pretty u, PP.Pretty r) => PrettyAnnot (Param (MemInfo SubExp u r)) where - ppAnnot = bindeeAnnot paramName paramDec - -bindeeAnnot :: (PP.Pretty u, PP.Pretty r) => - (a -> VName) -> (a -> MemInfo SubExp u r) - -> a -> Maybe PP.Doc -bindeeAnnot bindeeName bindeeLore bindee = - case bindeeLore bindee of - dec@MemArray{} -> - Just $ - PP.stack $ map (("-- "<>) . PP.text) $ lines $ - pretty (PP.ppr (bindeeName bindee) PP.<+> ":" PP.<+> PP.ppr dec) - MemMem {} -> - Nothing - MemPrim _ -> - Nothing - -extReturns :: [ExtType] -> [ExpReturns] -extReturns ts = - evalState (mapM addDec ts) 0 - where addDec (Prim bt) = - return $ MemPrim bt - addDec (Mem space) = - return $ MemMem space - addDec t@(Array bt shape u) - | existential t = do - i <- get <* modify (+1) - return $ MemArray bt shape u $ Just $ - ReturnsNewBlock DefaultSpace i $ - IxFun.iota $ map convert $ shapeDims shape - | otherwise = - return $ MemArray bt shape u Nothing - convert (Ext i) = LeafExp (Ext i) int32 - convert (Free v) = Free <$> primExpFromSubExp int32 v - -arrayVarReturns :: (HasScope lore m, Monad m, Mem lore) => - VName - -> m (PrimType, Shape, VName, IxFun.IxFun (PrimExp VName)) -arrayVarReturns v = do - summary <- lookupMemInfo v - case summary of - MemArray et shape _ (ArrayIn mem ixfun) -> - return (et, Shape $ shapeDims shape, mem, ixfun) - _ -> - error $ "arrayVarReturns: " ++ pretty v ++ " is not an array." - -varReturns :: (HasScope lore m, Monad m, Mem lore) => - VName -> m ExpReturns -varReturns v = do - summary <- lookupMemInfo v - case summary of - MemPrim bt -> - return $ MemPrim bt - MemArray et shape _ (ArrayIn mem ixfun) -> - return $ MemArray et (fmap Free shape) NoUniqueness $ - Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun - MemMem space -> - return $ MemMem space - --- | The return information of an expression. This can be seen as the --- "return type with memory annotations" of the expression. -expReturns :: (Monad m, HasScope lore m, - Mem lore) => - Exp lore -> m [ExpReturns] - -expReturns (BasicOp (SubExp (Var v))) = - pure <$> varReturns v - -expReturns (BasicOp (Opaque (Var v))) = - pure <$> varReturns v - -expReturns (BasicOp (Reshape newshape v)) = do - (et, _, mem, ixfun) <- arrayVarReturns v - return [MemArray et (Shape $ map (Free . newDim) newshape) NoUniqueness $ - Just $ ReturnsInBlock mem $ existentialiseIxFun [] $ - IxFun.reshape ixfun $ map (fmap $ primExpFromSubExp int32) newshape] - -expReturns (BasicOp (Rearrange perm v)) = do - (et, Shape dims, mem, ixfun) <- arrayVarReturns v - let ixfun' = IxFun.permute ixfun perm - dims' = rearrangeShape perm dims - return [MemArray et (Shape $ map Free dims') NoUniqueness $ - Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun'] - -expReturns (BasicOp (Rotate offsets v)) = do - (et, Shape dims, mem, ixfun) <- arrayVarReturns v - let offsets' = map (primExpFromSubExp int32) offsets - ixfun' = IxFun.rotate ixfun offsets' - return [MemArray et (Shape $ map Free dims) NoUniqueness $ - Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun'] - -expReturns (BasicOp (Index v slice)) = do - info <- sliceInfo v slice - case info of - MemArray et shape u (ArrayIn mem ixfun) -> - return [MemArray et (fmap Free shape) u $ - Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun] - MemPrim pt -> return [MemPrim pt] - MemMem space -> return [MemMem space] - -expReturns (BasicOp (Update v _ _)) = - pure <$> varReturns v - -expReturns (BasicOp op) = - extReturns . staticShapes <$> primOpType op - -expReturns e@(DoLoop ctx val _ _) = do - t <- expExtType e - zipWithM typeWithDec t $ map fst val - where typeWithDec t p = - case (t, paramDec p) of - (Array bt shape u, MemArray _ _ _ (ArrayIn mem ixfun)) - | Just (i, mem_p) <- isMergeVar mem, - Mem space <- paramType mem_p -> - return $ MemArray bt shape u $ Just $ ReturnsNewBlock space i ixfun' - | otherwise -> - return (MemArray bt shape u $ - Just $ ReturnsInBlock mem ixfun') - where ixfun' = existentialiseIxFun (map paramName mergevars) ixfun - (Array{}, _) -> - error "expReturns: Array return type but not array merge variable." - (Prim bt, _) -> - return $ MemPrim bt - (Mem{}, _) -> - error "expReturns: loop returns memory block explicitly." - isMergeVar v = find ((==v) . paramName . snd) $ zip [0..] mergevars - mergevars = map fst $ ctx ++ val - -expReturns (Apply _ _ ret _) = - return $ map funReturnsToExpReturns ret - -expReturns (If _ _ _ (IfDec ret _)) = - return $ map bodyReturnsToExpReturns ret - -expReturns (Op op) = - opReturns op - -sliceInfo :: (Monad m, HasScope lore m, Mem lore) => - VName - -> Slice SubExp -> m (MemInfo SubExp NoUniqueness MemBind) -sliceInfo v slice = do - (et, _, mem, ixfun) <- arrayVarReturns v - case sliceDims slice of - [] -> return $ MemPrim et - dims -> - return $ MemArray et (Shape dims) NoUniqueness $ - ArrayIn mem $ IxFun.slice ixfun - (map (fmap (primExpFromSubExp int32)) slice) - -class TypedOp (Op lore) => OpReturns lore where - opReturns :: (Monad m, HasScope lore m) => - Op lore -> m [ExpReturns] - opReturns op = extReturns <$> opType op - -applyFunReturns :: Typed dec => - [FunReturns] - -> [Param dec] - -> [(SubExp,Type)] - -> Maybe [FunReturns] -applyFunReturns rets params args - | Just _ <- applyRetType rettype params args = - Just $ map correctDims rets - | otherwise = - Nothing - where rettype = map declExtTypeOf rets - parammap :: M.Map VName (SubExp, Type) - parammap = M.fromList $ - zip (map paramName params) args - - substSubExp (Var v) - | Just (se,_) <- M.lookup v parammap = se - substSubExp se = se - - correctDims (MemPrim t) = - MemPrim t - correctDims (MemMem space) = - MemMem space - correctDims (MemArray et shape u memsummary) = - MemArray et (correctShape shape) u $ - correctSummary memsummary - - correctShape = Shape . map correctDim . shapeDims - correctDim (Ext i) = Ext i - correctDim (Free se) = Free $ substSubExp se - - correctSummary (ReturnsNewBlock space i ixfun) = - ReturnsNewBlock space i ixfun - correctSummary (ReturnsInBlock mem ixfun) = - -- FIXME: we should also do a replacement in ixfun here. - ReturnsInBlock mem' ixfun - where mem' = case M.lookup mem parammap of - Just (Var v, _) -> v - _ -> mem +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | Building blocks for defining representations where every array +-- is given information about which memory block is it based in, and +-- how array elements map to memory block offsets. +-- +-- There are two primary concepts you will need to understand: +-- +-- 1. Memory blocks, which are Futhark values of type v'Mem' +-- (parametrized with their size). These correspond to arbitrary +-- blocks of memory, and are created using the 'Alloc' operation. +-- +-- 2. Index functions, which describe a mapping from the index space +-- of an array (eg. a two-dimensional space for an array of type +-- @[[int]]@) to a one-dimensional offset into a memory block. +-- Thus, index functions describe how arbitrary-dimensional arrays +-- are mapped to the single-dimensional world of memory. +-- +-- At a conceptual level, imagine that we have a two-dimensional array +-- @a@ of 32-bit integers, consisting of @n@ rows of @m@ elements +-- each. This array could be represented in classic row-major format +-- with an index function like the following: +-- +-- @ +-- f(i,j) = i * m + j +-- @ +-- +-- When we want to know the location of element @a[2,3]@, we simply +-- call the index function as @f(2,3)@ and obtain @2*m+3@. We could +-- also have chosen another index function, one that represents the +-- array in column-major (or "transposed") format: +-- +-- @ +-- f(i,j) = j * n + i +-- @ +-- +-- Index functions are not Futhark-level functions, but a special +-- construct that the final code generator will eventually use to +-- generate concrete access code. By modifying the index functions we +-- can change how an array is represented in memory, which can permit +-- memory access pattern optimisations. +-- +-- Every time we bind an array, whether in a @let@-binding, @loop@ +-- merge parameter, or @lambda@ parameter, we have an annotation +-- specifying a memory block and an index function. In some cases, +-- such as @let@-bindings for many expressions, we are free to specify +-- an arbitrary index function and memory block - for example, we get +-- to decide where 'Copy' stores its result - but in other cases the +-- type rules of the expression chooses for us. For example, 'Index' +-- always produces an array in the same memory block as its input, and +-- with the same index function, except with some indices fixed. +module Futhark.IR.Mem + ( LetDecMem, + FParamMem, + LParamMem, + RetTypeMem, + BranchTypeMem, + MemOp (..), + MemInfo (..), + MemBound, + MemBind (..), + MemReturn (..), + IxFun, + ExtIxFun, + isStaticIxFun, + ExpReturns, + BodyReturns, + FunReturns, + noUniquenessReturns, + bodyReturnsToExpReturns, + Mem, + AllocOp (..), + OpReturns (..), + varReturns, + expReturns, + extReturns, + lookupMemInfo, + subExpMemInfo, + lookupArraySummary, + existentialiseIxFun, + + -- * Type checking parts + matchBranchReturnType, + matchPatternToExp, + matchFunctionReturnType, + matchLoopResultMem, + bodyReturnsFromPattern, + checkMemInfo, + + -- * Module re-exports + module Futhark.IR.Prop, + module Futhark.IR.Traversals, + module Futhark.IR.Pretty, + module Futhark.IR.Syntax, + module Futhark.Analysis.PrimExp.Convert, + ) +where + +import Control.Category +import Control.Monad.Except +import Control.Monad.Reader +import Control.Monad.State +import Data.Foldable (toList, traverse_) +import Data.List (elemIndex, find) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import Futhark.Analysis.Metrics +import Futhark.Analysis.PrimExp.Convert +import Futhark.Analysis.PrimExp.Simplify +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.IR.Aliases + ( Aliases, + removeExpAliases, + removePatternAliases, + removeScopeAliases, + ) +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.IR.Pretty +import Futhark.IR.Prop +import Futhark.IR.Prop.Aliases +import Futhark.IR.Syntax +import Futhark.IR.Traversals +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Optimise.Simplify.Lore +import Futhark.Transform.Rename +import Futhark.Transform.Substitute +import qualified Futhark.TypeCheck as TC +import Futhark.Util +import qualified Futhark.Util.Pretty as PP +import GHC.Generics (Generic) +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) + +type LetDecMem = MemInfo SubExp NoUniqueness MemBind + +type FParamMem = MemInfo SubExp Uniqueness MemBind + +type LParamMem = MemInfo SubExp NoUniqueness MemBind + +type RetTypeMem = FunReturns + +type BranchTypeMem = BodyReturns + +-- | The class of ops that have memory allocation. +class AllocOp op where + allocOp :: SubExp -> Space -> op + +type Mem lore = + ( AllocOp (Op lore), + FParamInfo lore ~ FParamMem, + LParamInfo lore ~ LParamMem, + LetDec lore ~ LetDecMem, + RetType lore ~ RetTypeMem, + BranchType lore ~ BranchTypeMem, + ASTLore lore, + Decorations lore, + OpReturns lore + ) + +instance IsRetType FunReturns where + primRetType = MemPrim + applyRetType = applyFunReturns + +instance IsBodyType BodyReturns where + primBodyType = MemPrim + +data MemOp inner + = -- | Allocate a memory block. This really should not be an + -- expression, but what are you gonna do... + Alloc SubExp Space + | Inner inner + deriving (Eq, Ord, Show, Generic) + +instance SexpIso inner => SexpIso (MemOp inner) where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "alloc") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "inner") >>> Sexp.el sexpIso)) + End + +instance AllocOp (MemOp inner) where + allocOp = Alloc + +instance FreeIn inner => FreeIn (MemOp inner) where + freeIn' (Alloc size _) = freeIn' size + freeIn' (Inner k) = freeIn' k + +instance TypedOp inner => TypedOp (MemOp inner) where + opType (Alloc _ space) = pure [Mem space] + opType (Inner k) = opType k + +instance AliasedOp inner => AliasedOp (MemOp inner) where + opAliases Alloc {} = [mempty] + opAliases (Inner k) = opAliases k + + consumedInOp Alloc {} = mempty + consumedInOp (Inner k) = consumedInOp k + +instance CanBeAliased inner => CanBeAliased (MemOp inner) where + type OpWithAliases (MemOp inner) = MemOp (OpWithAliases inner) + 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 + +instance Rename inner => Rename (MemOp inner) where + rename (Alloc size space) = Alloc <$> rename size <*> pure space + rename (Inner k) = Inner <$> rename k + +instance Substitute inner => Substitute (MemOp inner) where + substituteNames subst (Alloc size space) = Alloc (substituteNames subst size) space + substituteNames subst (Inner k) = Inner $ substituteNames subst k + +instance PP.Pretty inner => PP.Pretty (MemOp inner) where + ppr (Alloc e DefaultSpace) = PP.text "alloc" <> PP.apply [PP.ppr e] + ppr (Alloc e s) = PP.text "alloc" <> PP.apply [PP.ppr e, PP.ppr s] + ppr (Inner k) = PP.ppr k + +instance OpMetrics inner => OpMetrics (MemOp inner) where + opMetrics Alloc {} = seen "Alloc" + opMetrics (Inner k) = opMetrics k + +instance IsOp inner => IsOp (MemOp inner) where + safeOp (Alloc (Constant (IntValue (Int64Value k))) _) = k >= 0 + safeOp Alloc {} = False + safeOp (Inner k) = safeOp k + cheapOp (Inner k) = cheapOp k + cheapOp Alloc {} = True + +instance CanBeWise inner => CanBeWise (MemOp inner) where + type OpWithWisdom (MemOp inner) = MemOp (OpWithWisdom inner) + removeOpWisdom (Alloc size space) = Alloc size space + removeOpWisdom (Inner k) = Inner $ removeOpWisdom k + +instance ST.IndexOp inner => ST.IndexOp (MemOp inner) where + indexOp vtable k (Inner op) is = ST.indexOp vtable k op is + indexOp _ _ _ _ = Nothing + +-- | The index function representation used for memory annotations. +type IxFun = IxFun.IxFun (TPrimExp Int64 VName) + +-- | An index function that may contain existential variables. +type ExtIxFun = IxFun.IxFun (TPrimExp Int64 (Ext VName)) + +-- | A summary of the memory information for every let-bound +-- identifier, function parameter, and return value. Parameterisered +-- over uniqueness, dimension, and auxiliary array information. +data MemInfo d u ret + = -- | A primitive value. + MemPrim PrimType + | -- | A memory block. + MemMem Space + | -- | The array is stored in the named memory block, and with the + -- given index function. The index function maps indices in the + -- array to /element/ offset, /not/ byte offsets! To translate to + -- byte offsets, multiply the offset with the size of the array + -- element type. + MemArray PrimType (ShapeBase d) u ret + deriving (Eq, Show, Ord, Generic) --- XXX Ord? + +instance (SexpIso d, SexpIso u, SexpIso ret) => SexpIso (MemInfo d u ret) where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "prim") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "mem") >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "array") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End + +type MemBound u = MemInfo SubExp u MemBind + +instance FixExt ret => DeclExtTyped (MemInfo ExtSize Uniqueness ret) where + declExtTypeOf (MemPrim pt) = Prim pt + declExtTypeOf (MemMem space) = Mem space + declExtTypeOf (MemArray pt shape u _) = Array pt shape u + +instance FixExt ret => ExtTyped (MemInfo ExtSize NoUniqueness ret) where + extTypeOf (MemPrim pt) = Prim pt + extTypeOf (MemMem space) = Mem space + extTypeOf (MemArray pt shape u _) = Array pt shape u + +instance FixExt ret => FixExt (MemInfo ExtSize u ret) where + fixExt _ _ (MemPrim pt) = MemPrim pt + fixExt _ _ (MemMem space) = MemMem space + fixExt i se (MemArray pt shape u ret) = + MemArray pt (fixExt i se shape) u (fixExt i se ret) + +instance Typed (MemInfo SubExp Uniqueness ret) where + typeOf = fromDecl . declTypeOf + +instance Typed (MemInfo SubExp NoUniqueness ret) where + typeOf (MemPrim pt) = Prim pt + typeOf (MemMem space) = Mem space + typeOf (MemArray bt shape u _) = Array bt shape u + +instance DeclTyped (MemInfo SubExp Uniqueness ret) where + declTypeOf (MemPrim bt) = Prim bt + declTypeOf (MemMem space) = Mem space + declTypeOf (MemArray bt shape u _) = Array bt shape u + +instance (FreeIn d, FreeIn ret) => FreeIn (MemInfo d u ret) where + freeIn' (MemArray _ shape _ ret) = freeIn' shape <> freeIn' ret + freeIn' (MemMem s) = freeIn' s + freeIn' MemPrim {} = mempty + +instance (Substitute d, Substitute ret) => Substitute (MemInfo d u ret) where + substituteNames subst (MemArray bt shape u ret) = + MemArray + bt + (substituteNames subst shape) + u + (substituteNames subst ret) + substituteNames _ (MemMem space) = + MemMem space + substituteNames _ (MemPrim bt) = + MemPrim bt + +instance (Substitute d, Substitute ret) => Rename (MemInfo d u ret) where + rename = substituteRename + +simplifyIxFun :: + Engine.SimplifiableLore lore => + IxFun -> + Engine.SimpleM lore IxFun +simplifyIxFun = traverse $ fmap isInt64 . simplifyPrimExp . untyped + +simplifyExtIxFun :: + Engine.SimplifiableLore lore => + ExtIxFun -> + Engine.SimpleM lore ExtIxFun +simplifyExtIxFun = traverse $ fmap isInt64 . simplifyExtPrimExp . untyped + +isStaticIxFun :: ExtIxFun -> Maybe IxFun +isStaticIxFun = traverse $ traverse inst + where + inst Ext {} = Nothing + inst (Free x) = Just x + +instance + (Engine.Simplifiable d, Engine.Simplifiable ret) => + Engine.Simplifiable (MemInfo d u ret) + where + simplify (MemPrim bt) = + return $ MemPrim bt + simplify (MemMem space) = + pure $ MemMem space + simplify (MemArray bt shape u ret) = + MemArray bt <$> Engine.simplify shape <*> pure u <*> Engine.simplify ret + +instance + ( PP.Pretty (TypeBase (ShapeBase d) u), + PP.Pretty d, + PP.Pretty u, + PP.Pretty ret + ) => + PP.Pretty (MemInfo d u ret) + where + ppr (MemPrim bt) = PP.ppr bt + ppr (MemMem DefaultSpace) = PP.text "mem" + ppr (MemMem s) = PP.text "mem" <> PP.ppr s + ppr (MemArray bt shape u ret) = + PP.ppr (Array bt shape u) <> PP.text "@" <> PP.ppr ret + +instance PP.Pretty (Param (MemInfo SubExp Uniqueness ret)) where + ppr = PP.ppr . fmap declTypeOf + +instance PP.Pretty (Param (MemInfo SubExp NoUniqueness ret)) where + ppr = PP.ppr . fmap typeOf + +instance PP.Pretty (PatElemT (MemInfo SubExp NoUniqueness ret)) where + ppr = PP.ppr . fmap typeOf + +-- | Memory information for an array bound somewhere in the program. +data MemBind + = -- | Located in this memory block with this index + -- function. + ArrayIn VName IxFun + deriving (Show, Generic) + +instance SexpIso MemBind where + sexpIso = with $ \membind -> + Sexp.list (Sexp.el sexpIso >>> Sexp.el sexpIso) + >>> membind + +instance Eq MemBind where + _ == _ = True + +instance Ord MemBind where + _ `compare` _ = EQ + +instance Rename MemBind where + rename = substituteRename + +instance Substitute MemBind where + substituteNames substs (ArrayIn ident ixfun) = + ArrayIn (substituteNames substs ident) (substituteNames substs ixfun) + +instance PP.Pretty MemBind where + ppr (ArrayIn mem ixfun) = + PP.text "@" <> PP.ppr mem <> PP.text "->" PP. PP.ppr ixfun + +instance FreeIn MemBind where + freeIn' (ArrayIn mem ixfun) = freeIn' mem <> freeIn' ixfun + +-- | A description of the memory properties of an array being returned +-- by an operation. +data MemReturn + = -- | The array is located in a memory block that is + -- already in scope. + ReturnsInBlock VName ExtIxFun + | -- | The operation returns a new (existential) memory + -- block. + ReturnsNewBlock Space Int ExtIxFun + deriving (Show, Generic) + +instance SexpIso MemReturn where + sexpIso = + match $ + With + ( . + Sexp.list + ( Sexp.el (Sexp.sym "returns-in-block") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + ) + $ With + ( . + Sexp.list + ( Sexp.el (Sexp.sym "returns-new-block") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + ) + End + +instance Eq MemReturn where + _ == _ = True + +instance Ord MemReturn where + _ `compare` _ = EQ + +instance Rename MemReturn where + rename = substituteRename + +instance Substitute MemReturn where + substituteNames substs (ReturnsInBlock ident ixfun) = + ReturnsInBlock (substituteNames substs ident) (substituteNames substs ixfun) + substituteNames substs (ReturnsNewBlock space i ixfun) = + ReturnsNewBlock space i (substituteNames substs ixfun) + +instance FixExt MemReturn where + fixExt i (Var v) (ReturnsNewBlock _ j ixfun) + | j == i = + ReturnsInBlock v $ + fixExtIxFun + i + (primExpFromSubExp int64 (Var v)) + ixfun + fixExt i se (ReturnsNewBlock space j ixfun) = + ReturnsNewBlock + space + j' + (fixExtIxFun i (primExpFromSubExp int64 se) ixfun) + where + j' + | i < j = j -1 + | otherwise = j + fixExt i se (ReturnsInBlock mem ixfun) = + ReturnsInBlock mem (fixExtIxFun i (primExpFromSubExp int64 se) ixfun) + +fixExtIxFun :: Int -> PrimExp VName -> ExtIxFun -> ExtIxFun +fixExtIxFun i e = fmap $ isInt64 . replaceInPrimExp update . untyped + where + update (Ext j) t + | j > i = LeafExp (Ext $ j - 1) t + | j == i = fmap Free e + | otherwise = LeafExp (Ext j) t + update (Free x) t = LeafExp (Free x) t + +leafExp :: Int -> TPrimExp Int64 (Ext a) +leafExp i = isInt64 $ LeafExp (Ext i) int64 + +existentialiseIxFun :: [VName] -> IxFun -> ExtIxFun +existentialiseIxFun ctx = IxFun.substituteInIxFun ctx' . fmap (fmap Free) + where + ctx' = M.map leafExp $ M.fromList $ zip (map Free ctx) [0 ..] + +instance PP.Pretty MemReturn where + ppr (ReturnsInBlock v ixfun) = + PP.parens $ PP.text (pretty v) <> PP.text "->" PP. PP.ppr ixfun + ppr (ReturnsNewBlock space i ixfun) = + PP.text ("?" ++ show i) <> PP.ppr space <> PP.text "->" PP. PP.ppr ixfun + +instance FreeIn MemReturn where + freeIn' (ReturnsInBlock v ixfun) = freeIn' v <> freeIn' ixfun + freeIn' (ReturnsNewBlock space _ ixfun) = freeIn' space <> freeIn' ixfun + +instance Engine.Simplifiable MemReturn where + simplify (ReturnsNewBlock space i ixfun) = + ReturnsNewBlock space i <$> simplifyExtIxFun ixfun + simplify (ReturnsInBlock v ixfun) = + ReturnsInBlock <$> Engine.simplify v <*> simplifyExtIxFun ixfun + +instance Engine.Simplifiable MemBind where + simplify (ArrayIn mem ixfun) = + ArrayIn <$> Engine.simplify mem <*> simplifyIxFun ixfun + +instance Engine.Simplifiable [FunReturns] where + simplify = mapM Engine.simplify + +-- | The memory return of an expression. An array is annotated with +-- @Maybe MemReturn@, which can be interpreted as the expression +-- either dictating exactly where the array is located when it is +-- returned (if 'Just'), or able to put it whereever the binding +-- prefers (if 'Nothing'). +-- +-- This is necessary to capture the difference between an expression +-- that is just an array-typed variable, in which the array being +-- "returned" is located where it already is, and a @copy@ expression, +-- whose entire purpose is to store an existing array in some +-- arbitrary location. This is a consequence of the design decision +-- never to have implicit memory copies. +type ExpReturns = MemInfo ExtSize NoUniqueness (Maybe MemReturn) + +-- | The return of a body, which must always indicate where +-- returned arrays are located. +type BodyReturns = MemInfo ExtSize NoUniqueness MemReturn + +-- | The memory return of a function, which must always indicate where +-- returned arrays are located. +type FunReturns = MemInfo ExtSize Uniqueness MemReturn + +maybeReturns :: MemInfo d u r -> MemInfo d u (Maybe r) +maybeReturns (MemArray bt shape u ret) = + MemArray bt shape u $ Just ret +maybeReturns (MemPrim bt) = + MemPrim bt +maybeReturns (MemMem space) = + MemMem space + +noUniquenessReturns :: MemInfo d u r -> MemInfo d NoUniqueness r +noUniquenessReturns (MemArray bt shape _ r) = + MemArray bt shape NoUniqueness r +noUniquenessReturns (MemPrim bt) = + MemPrim bt +noUniquenessReturns (MemMem space) = + MemMem space + +funReturnsToExpReturns :: FunReturns -> ExpReturns +funReturnsToExpReturns = noUniquenessReturns . maybeReturns + +bodyReturnsToExpReturns :: BodyReturns -> ExpReturns +bodyReturnsToExpReturns = noUniquenessReturns . maybeReturns + +matchRetTypeToResult :: + (Mem lore, TC.Checkable lore) => + [FunReturns] -> + Result -> + TC.TypeM lore () +matchRetTypeToResult rettype result = do + scope <- askScope + result_ts <- runReaderT (mapM subExpMemInfo result) $ removeScopeAliases scope + matchReturnType rettype result result_ts + +matchFunctionReturnType :: + (Mem lore, TC.Checkable lore) => + [FunReturns] -> + Result -> + TC.TypeM lore () +matchFunctionReturnType rettype result = do + matchRetTypeToResult rettype result + mapM_ checkResultSubExp result + where + checkResultSubExp Constant {} = + return () + checkResultSubExp (Var v) = do + dec <- varMemInfo v + case dec of + MemPrim _ -> return () + MemMem {} -> return () + MemArray _ _ _ (ArrayIn _ ixfun) + | IxFun.isLinear ixfun -> + return () + | otherwise -> + TC.bad $ + TC.TypeError $ + "Array " ++ pretty v + ++ " returned by function, but has nontrivial index function " + ++ pretty ixfun + +matchLoopResultMem :: + (Mem lore, TC.Checkable lore) => + [FParam (Aliases lore)] -> + [FParam (Aliases lore)] -> + [SubExp] -> + TC.TypeM lore () +matchLoopResultMem ctx val = matchRetTypeToResult rettype + where + ctx_names = map paramName ctx + + -- Invent a ReturnType so we can pretend that the loop body is + -- actually returning from a function. + rettype = map (toRet . paramDec) val + + toExtV v + | Just i <- v `elemIndex` ctx_names = Ext i + | otherwise = Free v + + toExtSE (Var v) = Var <$> toExtV v + toExtSE (Constant v) = Free $ Constant v + + toRet (MemPrim t) = + MemPrim t + toRet (MemMem space) = + MemMem space + toRet (MemArray pt shape u (ArrayIn mem ixfun)) + | Just i <- mem `elemIndex` ctx_names, + Param _ (MemMem space) : _ <- drop i ctx = + MemArray pt shape' u $ ReturnsNewBlock space i ixfun' + | otherwise = + MemArray pt shape' u $ ReturnsInBlock mem ixfun' + where + shape' = fmap toExtSE shape + ixfun' = existentialiseIxFun ctx_names ixfun + +matchBranchReturnType :: + (Mem lore, TC.Checkable lore) => + [BodyReturns] -> + Body (Aliases lore) -> + TC.TypeM lore () +matchBranchReturnType rettype (Body _ stms res) = do + scope <- askScope + ts <- runReaderT (mapM subExpMemInfo res) $ removeScopeAliases (scope <> scopeOf stms) + matchReturnType rettype res ts + +-- | Helper function for index function unification. +-- +-- The first return value maps a VName (wrapped in 'Free') to its Int +-- (wrapped in 'Ext'). In case of duplicates, it is mapped to the +-- *first* Int that occurs. +-- +-- The second return value maps each Int (wrapped in an 'Ext') to a +-- 'LeafExp' 'Ext' with the Int at which its associated VName first +-- occurs. +getExtMaps :: + [(VName, Int)] -> + ( M.Map (Ext VName) (TPrimExp Int64 (Ext VName)), + M.Map (Ext VName) (TPrimExp Int64 (Ext VName)) + ) +getExtMaps ctx_lst_ids = + ( M.map leafExp $ M.mapKeys Free $ M.fromListWith (flip const) ctx_lst_ids, + M.fromList $ + mapMaybe + ( traverse + ( fmap (\i -> isInt64 $ LeafExp (Ext i) int64) + . (`lookup` ctx_lst_ids) + ) + . uncurry (flip (,)) + . fmap Ext + ) + ctx_lst_ids + ) + +matchReturnType :: + PP.Pretty u => + [MemInfo ExtSize u MemReturn] -> + [SubExp] -> + [MemInfo SubExp NoUniqueness MemBind] -> + TC.TypeM lore () +matchReturnType rettype res ts = do + let (ctx_ts, val_ts) = splitFromEnd (length rettype) ts + (ctx_res, _val_res) = splitFromEnd (length rettype) res + + existentialiseIxFun0 :: IxFun -> ExtIxFun + existentialiseIxFun0 = fmap $ fmap Free + + fetchCtx i = case maybeNth i $ zip ctx_res ctx_ts of + Nothing -> + throwError $ + "Cannot find context variable " + ++ show i + ++ " in context results: " + ++ pretty ctx_res + Just (se, t) -> return (se, t) + + checkReturn (MemPrim x) (MemPrim y) + | x == y = return () + checkReturn (MemMem x) (MemMem y) + | x == y = return () + checkReturn + (MemArray x_pt x_shape _ x_ret) + (MemArray y_pt y_shape _ y_ret) + | x_pt == y_pt, + shapeRank x_shape == shapeRank y_shape = do + zipWithM_ checkDim (shapeDims x_shape) (shapeDims y_shape) + checkMemReturn x_ret y_ret + checkReturn x y = + throwError $ unwords ["Expected", pretty x, "but got", pretty y] + + checkDim (Free x) y + | x == y = return () + | otherwise = + throwError $ + unwords + [ "Expected dim", + pretty x, + "but got", + pretty y + ] + checkDim (Ext i) y = do + (x, _) <- fetchCtx i + unless (x == y) $ + throwError $ + unwords + [ "Expected ext dim", + pretty i, + "=>", + pretty x, + "but got", + pretty y + ] + + extsInMemInfo :: MemInfo ExtSize u MemReturn -> S.Set Int + extsInMemInfo (MemArray _ shp _ ret) = + extInShape shp <> extInMemReturn ret + extsInMemInfo _ = S.empty + + checkMemReturn (ReturnsInBlock x_mem x_ixfun) (ArrayIn y_mem y_ixfun) + | x_mem == y_mem = + unless (IxFun.closeEnough x_ixfun $ existentialiseIxFun0 y_ixfun) $ + throwError $ + unwords + [ "Index function unification failed (ReturnsInBlock)", + "\nixfun of body result: ", + pretty y_ixfun, + "\nixfun of return type: ", + pretty x_ixfun, + "\nand context elements: ", + pretty ctx_res + ] + checkMemReturn + (ReturnsNewBlock x_space x_ext x_ixfun) + (ArrayIn y_mem y_ixfun) = do + (x_mem, x_mem_type) <- fetchCtx x_ext + unless (IxFun.closeEnough x_ixfun $ existentialiseIxFun0 y_ixfun) $ + throwError $ + unwords + [ "Index function unification failed (ReturnsNewBlock)", + "\nixfun of body result: ", + pretty y_ixfun, + "\nixfun of return type: ", + pretty x_ixfun, + "\nand context elements: ", + pretty ctx_res + ] + case x_mem_type of + MemMem y_space -> + unless (x_space == y_space) $ + throwError $ + unwords + [ "Expected memory", + pretty y_mem, + "in space", + pretty x_space, + "but actually in space", + pretty y_space + ] + t -> + throwError $ + unwords + [ "Expected memory", + pretty x_ext, + "=>", + pretty x_mem, + "but but has type", + pretty t + ] + checkMemReturn x y = + throwError $ + unwords + [ "Expected array in", + pretty x, + "but array returned in", + pretty y + ] + + bad :: String -> TC.TypeM lore a + bad s = + TC.bad $ + TC.TypeError $ + PP.pretty $ + "Return type" + PP. PP.indent 2 (ppTuple' rettype) + PP. "cannot match returns of results" + PP. PP.indent 2 (ppTuple' ts) + PP. PP.text s + + unless (length (S.unions $ map extsInMemInfo rettype) == length ctx_res) $ + TC.bad $ + TC.TypeError $ + "Too many context parameters for the number of " + ++ "existentials in the return type! type:\n " + ++ prettyTuple rettype + ++ "\ncannot match context parameters:\n " + ++ prettyTuple ctx_res + + either bad return =<< runExceptT (zipWithM_ checkReturn rettype val_ts) + +matchPatternToExp :: + (Mem lore, TC.Checkable lore) => + Pattern (Aliases lore) -> + Exp (Aliases lore) -> + TC.TypeM lore () +matchPatternToExp pat e = do + scope <- asksScope removeScopeAliases + rt <- runReaderT (expReturns $ removeExpAliases e) scope + + let (ctxs, vals) = bodyReturnsFromPattern $ removePatternAliases pat + (ctx_ids, _ctx_ts) = unzip ctxs + (_val_ids, val_ts) = unzip vals + (ctx_map_ids, ctx_map_exts) = + getExtMaps $ zip ctx_ids [0 .. length ctx_ids - 1] + + let rt_exts = foldMap extInExpReturns rt + + unless + ( length val_ts == length rt + && and (zipWith (matches ctx_map_ids ctx_map_exts) val_ts rt) + && M.keysSet ctx_map_exts `S.isSubsetOf` S.map Ext rt_exts + ) + $ TC.bad $ + TC.TypeError $ + "Expression type:\n " ++ prettyTuple rt + ++ "\ncannot match pattern type:\n " + ++ prettyTuple val_ts + ++ "\nwith context elements: " + ++ pretty ctx_ids + where + matches _ _ (MemPrim x) (MemPrim y) = x == y + matches _ _ (MemMem x_space) (MemMem y_space) = + x_space == y_space + matches ctxids ctxexts (MemArray x_pt x_shape _ x_ret) (MemArray y_pt y_shape _ y_ret) = + x_pt == y_pt && x_shape == y_shape + && case (x_ret, y_ret) of + (ReturnsInBlock _ x_ixfun, Just (ReturnsInBlock _ y_ixfun)) -> + let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun + y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun + in IxFun.closeEnough x_ixfun' y_ixfun' + ( ReturnsInBlock _ x_ixfun, + Just (ReturnsNewBlock _ _ y_ixfun) + ) -> + let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun + y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun + in IxFun.closeEnough x_ixfun' y_ixfun' + ( ReturnsNewBlock _ x_i x_ixfun, + Just (ReturnsNewBlock _ y_i y_ixfun) + ) -> + let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun + y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun + in x_i == y_i && IxFun.closeEnough x_ixfun' y_ixfun' + (_, Nothing) -> True + _ -> False + matches _ _ _ _ = False + + extInExpReturns :: ExpReturns -> S.Set Int + extInExpReturns (MemArray _ shape _ mem_return) = + extInShape shape <> maybe S.empty extInMemReturn mem_return + extInExpReturns _ = mempty + +extInShape :: ShapeBase (Ext SubExp) -> S.Set Int +extInShape shape = S.fromList $ mapMaybe isExt $ shapeDims shape + +extInMemReturn :: MemReturn -> S.Set Int +extInMemReturn (ReturnsInBlock _ extixfn) = extInIxFn extixfn +extInMemReturn (ReturnsNewBlock _ i extixfn) = + S.singleton i <> extInIxFn extixfn + +extInIxFn :: ExtIxFun -> S.Set Int +extInIxFn ixfun = S.fromList $ concatMap (mapMaybe isExt . toList) ixfun + +varMemInfo :: + Mem lore => + VName -> + TC.TypeM lore (MemInfo SubExp NoUniqueness MemBind) +varMemInfo name = do + dec <- TC.lookupVar name + + case dec of + LetName (_, summary) -> return summary + FParamName summary -> return $ noUniquenessReturns summary + LParamName summary -> return summary + IndexName it -> return $ MemPrim $ IntType it + +nameInfoToMemInfo :: Mem lore => NameInfo lore -> MemBound NoUniqueness +nameInfoToMemInfo info = + case info of + FParamName summary -> noUniquenessReturns summary + LParamName summary -> summary + LetName summary -> summary + IndexName it -> MemPrim $ IntType it + +lookupMemInfo :: + (HasScope lore m, Mem lore) => + VName -> + m (MemInfo SubExp NoUniqueness MemBind) +lookupMemInfo = fmap nameInfoToMemInfo . lookupInfo + +subExpMemInfo :: + (HasScope lore m, Monad m, Mem lore) => + SubExp -> + m (MemInfo SubExp NoUniqueness MemBind) +subExpMemInfo (Var v) = lookupMemInfo v +subExpMemInfo (Constant v) = return $ MemPrim $ primValueType v + +lookupArraySummary :: + (Mem lore, HasScope lore m, Monad m) => + VName -> + m (VName, IxFun.IxFun (TPrimExp Int64 VName)) +lookupArraySummary name = do + summary <- lookupMemInfo name + case summary of + MemArray _ _ _ (ArrayIn mem ixfun) -> + return (mem, ixfun) + _ -> + error $ "Variable " ++ pretty name ++ " does not look like an array." + +checkMemInfo :: + TC.Checkable lore => + VName -> + MemInfo SubExp u MemBind -> + TC.TypeM lore () +checkMemInfo _ (MemPrim _) = return () +checkMemInfo _ (MemMem (ScalarSpace d _)) = mapM_ (TC.require [Prim int64]) d +checkMemInfo _ (MemMem _) = return () +checkMemInfo name (MemArray _ shape _ (ArrayIn v ixfun)) = do + t <- lookupType v + case t of + Mem {} -> + return () + _ -> + TC.bad $ + TC.TypeError $ + "Variable " ++ pretty v + ++ " used as memory block, but is of type " + ++ pretty t + ++ "." + + TC.context ("in index function " ++ pretty ixfun) $ do + traverse_ (TC.requirePrimExp int64 . untyped) ixfun + let ixfun_rank = IxFun.rank ixfun + ident_rank = shapeRank shape + unless (ixfun_rank == ident_rank) $ + TC.bad $ + TC.TypeError $ + "Arity of index function (" ++ pretty ixfun_rank + ++ ") does not match rank of array " + ++ pretty name + ++ " (" + ++ show ident_rank + ++ ")" + +bodyReturnsFromPattern :: + PatternT (MemBound NoUniqueness) -> + ([(VName, BodyReturns)], [(VName, BodyReturns)]) +bodyReturnsFromPattern pat = + ( map asReturns $ patternContextElements pat, + map asReturns $ patternValueElements pat + ) + where + ctx = patternContextElements pat + + ext (Var v) + | Just (i, _) <- find ((== v) . patElemName . snd) $ zip [0 ..] ctx = + Ext i + ext se = Free se + + asReturns pe = + ( patElemName pe, + case patElemDec pe of + MemPrim pt -> MemPrim pt + MemMem space -> MemMem space + MemArray pt shape u (ArrayIn mem ixfun) -> + MemArray pt (Shape $ map ext $ shapeDims shape) u $ + case find ((== mem) . patElemName . snd) $ zip [0 ..] ctx of + Just (i, PatElem _ (MemMem space)) -> + ReturnsNewBlock space i $ + existentialiseIxFun (map patElemName ctx) ixfun + _ -> ReturnsInBlock mem $ existentialiseIxFun [] ixfun + ) + +instance (PP.Pretty u, PP.Pretty r) => PrettyAnnot (PatElemT (MemInfo SubExp u r)) where + ppAnnot = bindeeAnnot patElemName patElemDec + +instance (PP.Pretty u, PP.Pretty r) => PrettyAnnot (Param (MemInfo SubExp u r)) where + ppAnnot = bindeeAnnot paramName paramDec + +bindeeAnnot :: + (PP.Pretty u, PP.Pretty r) => + (a -> VName) -> + (a -> MemInfo SubExp u r) -> + a -> + Maybe PP.Doc +bindeeAnnot bindeeName bindeeLore bindee = + case bindeeLore bindee of + dec@MemArray {} -> + Just $ + PP.stack $ + map (("-- " <>) . PP.text) $ + lines $ + pretty (PP.ppr (bindeeName bindee) PP.<+> ":" PP.<+> PP.ppr dec) + MemMem {} -> + Nothing + MemPrim _ -> + Nothing + +extReturns :: [ExtType] -> [ExpReturns] +extReturns ts = + evalState (mapM addDec ts) 0 + where + addDec (Prim bt) = + return $ MemPrim bt + addDec (Mem space) = + return $ MemMem space + addDec t@(Array bt shape u) + | existential t = do + i <- get <* modify (+ 1) + return $ + MemArray bt shape u $ + Just $ + ReturnsNewBlock DefaultSpace i $ + IxFun.iota $ map convert $ shapeDims shape + | otherwise = + return $ MemArray bt shape u Nothing + convert (Ext i) = le64 (Ext i) + convert (Free v) = Free <$> pe64 v + +arrayVarReturns :: + (HasScope lore m, Monad m, Mem lore) => + VName -> + m (PrimType, Shape, VName, IxFun) +arrayVarReturns v = do + summary <- lookupMemInfo v + case summary of + MemArray et shape _ (ArrayIn mem ixfun) -> + return (et, Shape $ shapeDims shape, mem, ixfun) + _ -> + error $ "arrayVarReturns: " ++ pretty v ++ " is not an array." + +varReturns :: + (HasScope lore m, Monad m, Mem lore) => + VName -> + m ExpReturns +varReturns v = do + summary <- lookupMemInfo v + case summary of + MemPrim bt -> + return $ MemPrim bt + MemArray et shape _ (ArrayIn mem ixfun) -> + return $ + MemArray et (fmap Free shape) NoUniqueness $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun + MemMem space -> + return $ MemMem space + +-- | The return information of an expression. This can be seen as the +-- "return type with memory annotations" of the expression. +expReturns :: + ( Monad m, + HasScope lore m, + Mem lore + ) => + Exp lore -> + m [ExpReturns] +expReturns (BasicOp (SubExp (Var v))) = + pure <$> varReturns v +expReturns (BasicOp (Opaque (Var v))) = + pure <$> varReturns v +expReturns (BasicOp (Reshape newshape v)) = do + (et, _, mem, ixfun) <- arrayVarReturns v + return + [ MemArray et (Shape $ map (Free . newDim) newshape) NoUniqueness $ + Just $ + ReturnsInBlock mem $ + existentialiseIxFun [] $ + IxFun.reshape ixfun $ map (fmap pe64) newshape + ] +expReturns (BasicOp (Rearrange perm v)) = do + (et, Shape dims, mem, ixfun) <- arrayVarReturns v + let ixfun' = IxFun.permute ixfun perm + dims' = rearrangeShape perm dims + return + [ MemArray et (Shape $ map Free dims') NoUniqueness $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun' + ] +expReturns (BasicOp (Rotate offsets v)) = do + (et, Shape dims, mem, ixfun) <- arrayVarReturns v + let offsets' = map pe64 offsets + ixfun' = IxFun.rotate ixfun offsets' + return + [ MemArray et (Shape $ map Free dims) NoUniqueness $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun' + ] +expReturns (BasicOp (Index v slice)) = do + info <- sliceInfo v slice + case info of + MemArray et shape u (ArrayIn mem ixfun) -> + return + [ MemArray et (fmap Free shape) u $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun + ] + MemPrim pt -> return [MemPrim pt] + MemMem space -> return [MemMem space] +expReturns (BasicOp (Update v _ _)) = + pure <$> varReturns v +expReturns (BasicOp op) = + extReturns . staticShapes <$> primOpType op +expReturns e@(DoLoop ctx val _ _) = do + t <- expExtType e + zipWithM typeWithDec t $ map fst val + where + typeWithDec t p = + case (t, paramDec p) of + ( Array bt shape u, + MemArray _ _ _ (ArrayIn mem ixfun) + ) + | Just (i, mem_p) <- isMergeVar mem, + Mem space <- paramType mem_p -> + return $ MemArray bt shape u $ Just $ ReturnsNewBlock space i ixfun' + | otherwise -> + return + ( MemArray bt shape u $ + Just $ ReturnsInBlock mem ixfun' + ) + where + ixfun' = existentialiseIxFun (map paramName mergevars) ixfun + (Array {}, _) -> + error "expReturns: Array return type but not array merge variable." + (Prim bt, _) -> + return $ MemPrim bt + (Mem {}, _) -> + error "expReturns: loop returns memory block explicitly." + isMergeVar v = find ((== v) . paramName . snd) $ zip [0 ..] mergevars + mergevars = map fst $ ctx ++ val +expReturns (Apply _ _ ret _) = + return $ map funReturnsToExpReturns ret +expReturns (If _ _ _ (IfDec ret _)) = + return $ map bodyReturnsToExpReturns ret +expReturns (Op op) = + opReturns op + +sliceInfo :: + (Monad m, HasScope lore m, Mem lore) => + VName -> + Slice SubExp -> + m (MemInfo SubExp NoUniqueness MemBind) +sliceInfo v slice = do + (et, _, mem, ixfun) <- arrayVarReturns v + case sliceDims slice of + [] -> return $ MemPrim et + dims -> + return $ + MemArray et (Shape dims) NoUniqueness $ + ArrayIn mem $ + IxFun.slice + ixfun + (map (fmap (isInt64 . primExpFromSubExp int64)) slice) + +class TypedOp (Op lore) => OpReturns lore where + opReturns :: + (Monad m, HasScope lore m) => + Op lore -> + m [ExpReturns] + opReturns op = extReturns <$> opType op + +applyFunReturns :: + Typed dec => + [FunReturns] -> + [Param dec] -> + [(SubExp, Type)] -> + Maybe [FunReturns] +applyFunReturns rets params args + | Just _ <- applyRetType rettype params args = + Just $ map correctDims rets + | otherwise = + Nothing + where + rettype = map declExtTypeOf rets + parammap :: M.Map VName (SubExp, Type) + parammap = + M.fromList $ + zip (map paramName params) args + + substSubExp (Var v) + | Just (se, _) <- M.lookup v parammap = se + substSubExp se = se + + correctDims (MemPrim t) = + MemPrim t + correctDims (MemMem space) = + MemMem space + correctDims (MemArray et shape u memsummary) = + MemArray et (correctShape shape) u $ + correctSummary memsummary + + correctShape = Shape . map correctDim . shapeDims + correctDim (Ext i) = Ext i + correctDim (Free se) = Free $ substSubExp se + + correctSummary (ReturnsNewBlock space i ixfun) = + ReturnsNewBlock space i ixfun + correctSummary (ReturnsInBlock mem ixfun) = + -- FIXME: we should also do a replacement in ixfun here. + ReturnsInBlock mem' ixfun + where + mem' = case M.lookup mem parammap of + Just (Var v, _) -> v + _ -> mem diff --git a/src/Futhark/IR/Mem/IxFun.hs b/src/Futhark/IR/Mem/IxFun.hs --- a/src/Futhark/IR/Mem/IxFun.hs +++ b/src/Futhark/IR/Mem/IxFun.hs @@ -1,828 +1,1066 @@ -{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} --- | This module contains a representation for the index function based on --- linear-memory accessor descriptors; see Zhu, Hoeflinger and David work. -module Futhark.IR.Mem.IxFun - ( IxFun(..) - , index - , iota - , permute - , rotate - , reshape - , slice - , rebase - , shape - , rank - , linearWithOffset - , rearrangeWithOffset - , isDirect - , isLinear - , substituteInIxFun - , leastGeneralGeneralization - , existentialize - , closeEnough - ) - where - -import Prelude hiding (mod) -import Data.List (sort, sortBy, zip4, zip5, zipWith5) -import qualified Data.List.NonEmpty as NE -import Data.List.NonEmpty (NonEmpty(..)) -import Data.Function (on) -import Data.Maybe (isJust) -import Control.Monad.Identity -import Control.Monad.State -import Control.Monad.Writer -import qualified Data.Map.Strict as M - -import Futhark.Analysis.PrimExp (PrimExp(..), primExpType) -import Futhark.IR.Syntax.Core (Ext(..)) -import Futhark.Transform.Substitute -import Futhark.Transform.Rename -import Futhark.IR.Syntax - (ShapeChange, DimChange(..), DimIndex(..), Slice, unitSlice, dimFix) -import Futhark.IR.Prop -import Futhark.Util.IntegralExp -import Futhark.Util.Pretty -import Futhark.Analysis.PrimExp.Convert (substituteInPrimExp) -import qualified Futhark.Analysis.PrimExp.Generalize as PEG - -type Shape num = [num] -type Indices num = [num] -type Permutation = [Int] - -data Monotonicity = Inc | Dec | Unknown - -- ^ monotonously increasing, decreasing or unknown - deriving (Show, Eq) - -data LMADDim num = LMADDim { ldStride :: num - , ldRotate :: num - , ldShape :: num - , ldPerm :: Int - , ldMon :: Monotonicity - } - deriving (Show, Eq) - --- | LMAD's representation consists of a general offset and for each dimension a --- stride, rotate factor, number of elements (or shape), permutation, and --- monotonicity. Note that the permutation is not strictly necessary in that the --- permutation can be performed directly on LMAD dimensions, but then it is --- difficult to extract the permutation back from an LMAD. --- --- LMAD algebra is closed under composition w.r.t. operators such as --- permute, index and slice. However, other operations, such as --- reshape, cannot always be represented inside the LMAD algebra. --- --- It follows that the general representation of an index function is a list of --- LMADS, in which each following LMAD in the list implicitly corresponds to an --- irregular reshaping operation. --- --- However, we expect that the common case is when the index function is one --- LMAD -- we call this the "nice" representation. --- --- Finally, the list of LMADs is kept in an @IxFun@ together with the shape of --- the original array, and a bit to indicate whether the index function is --- contiguous, i.e., if we instantiate all the points of the current index --- function, do we get a contiguous memory interval? --- --- By definition, the LMAD denotes the set of points (simplified): --- --- \{ o + \Sigma_{j=0}^{k} ((i_j+r_j) `mod` n_j)*s_j, --- \forall i_j such that 0<=i_j Pretty (LMAD num) where - ppr (LMAD offset dims) = - braces $ semisep [ text "offset: " <> oneLine (ppr offset) - , text "strides: " <> p ldStride - , text "rotates: " <> p ldRotate - , text "shape: " <> p ldShape - , text "permutation: " <> p ldPerm - , text "monotonicity: " <> p ldMon - ] - where p f = oneLine $ brackets $ commasep $ map (ppr . f) dims - -instance Pretty num => Pretty (IxFun num) where - ppr (IxFun lmads oshp cg) = - braces $ semisep [ text "base: " <> brackets (commasep $ map ppr oshp) - , text "contiguous: " <> text (show cg) - , text "LMADs: " <> brackets (commasep $ NE.toList $ NE.map ppr lmads) - ] - - -instance Substitute num => Substitute (LMAD num) where - substituteNames substs = fmap $ substituteNames substs - -instance Substitute num => Substitute (IxFun num) where - substituteNames substs = fmap $ substituteNames substs - -instance Substitute num => Rename (LMAD num) where - rename = substituteRename - -instance Substitute num => Rename (IxFun num) where - rename = substituteRename - - -instance FreeIn num => FreeIn (LMAD num) where - freeIn' = foldMap freeIn' - -instance FreeIn num => FreeIn (IxFun num) where - freeIn' = foldMap freeIn' - -instance Functor LMAD where - fmap f = runIdentity . traverse (return . f) - -instance Functor IxFun where - fmap f = runIdentity . traverse (return . f) - - -instance Foldable LMAD where - foldMap f = execWriter . traverse (tell . f) - -instance Foldable IxFun where - foldMap f = execWriter . traverse (tell . f) - - -instance Traversable LMAD where - traverse f (LMAD offset dims) = - LMAD <$> f offset <*> traverse f' dims - where f' (LMADDim s r n p m) = - LMADDim <$> f s <*> f r <*> f n <*> pure p <*> pure m - -instance Traversable IxFun where - traverse f (IxFun lmads oshp cg) = - IxFun <$> traverse (traverse f) lmads <*> traverse f oshp <*> pure cg - -(++@) :: [a] -> NonEmpty a -> NonEmpty a -es ++@ (ne :| nes) = case es of - e : es' -> e :| es' ++ [ne] ++ nes - [] -> ne :| nes - -(@++@) :: NonEmpty a -> NonEmpty a -> NonEmpty a -(x :| xs) @++@ (y :| ys) = x :| xs ++ [y] ++ ys - -invertMonotonicity :: Monotonicity -> Monotonicity -invertMonotonicity Inc = Dec -invertMonotonicity Dec = Inc -invertMonotonicity Unknown = Unknown - -lmadPermutation :: LMAD num -> Permutation -lmadPermutation = map ldPerm . lmadDims - -setLMADPermutation :: Permutation -> LMAD num -> LMAD num -setLMADPermutation perm lmad = - lmad { lmadDims = zipWith (\dim p -> dim { ldPerm = p }) (lmadDims lmad) perm } - -setLMADShape :: Shape num -> LMAD num -> LMAD num -setLMADShape shp lmad = lmad { lmadDims = zipWith (\dim s -> dim { ldShape = s }) (lmadDims lmad) shp } - --- | Substitute a name with a PrimExp in an LMAD. -substituteInLMAD :: Ord a => M.Map a (PrimExp a) -> LMAD (PrimExp a) - -> LMAD (PrimExp a) -substituteInLMAD tab (LMAD offset dims) = - let offset' = substituteInPrimExp tab offset - dims' = map (\(LMADDim s r n p m) -> - LMADDim - (substituteInPrimExp tab s) - (substituteInPrimExp tab r) - (substituteInPrimExp tab n) - p m) - dims - in LMAD offset' dims' - --- | Substitute a name with a PrimExp in an index function. -substituteInIxFun :: (Ord a) => M.Map a (PrimExp a) -> IxFun (PrimExp a) - -> IxFun (PrimExp a) -substituteInIxFun tab (IxFun lmads oshp cg) = - IxFun (NE.map (substituteInLMAD tab) lmads) - (map (substituteInPrimExp tab) oshp) - cg - --- | Is this is a row-major array? -isDirect :: (Eq num, IntegralExp num) => IxFun num -> Bool -isDirect ixfun@(IxFun (LMAD offset dims :| []) oshp True) = - let strides_expected = reverse $ scanl (*) 1 (reverse (tail oshp)) - in hasContiguousPerm ixfun && - length oshp == length dims && - offset == 0 && - all (\(LMADDim s r n p _, m, d, se) -> - s == se && r == 0 && n == d && p == m) - (zip4 dims [0..length dims - 1] oshp strides_expected) -isDirect _ = False - --- | Does the index function have an ascending permutation? -hasContiguousPerm :: IxFun num -> Bool -hasContiguousPerm (IxFun (lmad :| []) _ _) = - let perm = lmadPermutation lmad - in perm == sort perm -hasContiguousPerm _ = False - --- | Shape of an index function. -shape :: (Eq num, IntegralExp num) => IxFun num -> Shape num -shape (IxFun (lmad :| _) _ _) = lmadShape lmad - --- | Shape of an LMAD. -lmadShape :: (Eq num, IntegralExp num) => LMAD num -> Shape num -lmadShape lmad = permuteInv (lmadPermutation lmad) $ lmadShapeBase lmad - --- | Shape of an LMAD, ignoring permutations. -lmadShapeBase :: (Eq num, IntegralExp num) => LMAD num -> Shape num -lmadShapeBase = map ldShape . lmadDims - --- | Compute the flat memory index for a complete set @inds@ of array indices --- and a certain element size @elem_size@. -index :: (IntegralExp num, Eq num) => - IxFun num -> Indices num -> num -index = indexFromLMADs . ixfunLMADs - where indexFromLMADs :: (IntegralExp num, Eq num) => - NonEmpty (LMAD num) -> Indices num -> num - indexFromLMADs (lmad :| []) inds = indexLMAD lmad inds - indexFromLMADs (lmad1 :| lmad2 : lmads) inds = - let i_flat = indexLMAD lmad1 inds - new_inds = unflattenIndex (permuteFwd (lmadPermutation lmad2) $ lmadShapeBase lmad2) i_flat - in indexFromLMADs (lmad2 :| lmads) new_inds - - -- | Compute the flat index of an LMAD. - indexLMAD :: (IntegralExp num, Eq num) => - LMAD num -> Indices num -> num - indexLMAD lmad@(LMAD off dims) inds = - let prod = sum $ zipWith flatOneDim - (map (\(LMADDim s r n _ _) -> (s, r, n)) dims) - (permuteInv (lmadPermutation lmad) inds) - in off + prod - --- | iota. -iota :: IntegralExp num => Shape num -> IxFun num -iota ns = - let rs = replicate (length ns) 0 - in IxFun (makeRotIota Inc 0 (zip rs ns) :| []) ns True - --- | Permute dimensions. -permute :: IntegralExp num => - IxFun num -> Permutation -> IxFun num -permute (IxFun (lmad :| lmads) oshp cg) perm_new = - let perm_cur = lmadPermutation lmad - perm = map (perm_cur !!) perm_new - in IxFun (setLMADPermutation perm lmad :| lmads) oshp cg - --- | Rotate an index function. -rotate :: (Eq num, IntegralExp num) => - IxFun num -> Indices num -> IxFun num -rotate (IxFun (lmad@(LMAD off dims) :| lmads) oshp cg) offs = - let dims' = zipWith (\(LMADDim s r n p f) o -> - if s == 0 then LMADDim 0 0 n p Unknown - else LMADDim s (r + o) n p f - ) dims (permuteInv (lmadPermutation lmad) offs) - in IxFun (LMAD off dims' :| lmads) oshp cg - --- | Handle the case where a slice can stay within a single LMAD. -sliceOneLMAD :: (Eq num, IntegralExp num) => - IxFun num -> Slice num -> Maybe (IxFun num) -sliceOneLMAD (IxFun (lmad@(LMAD _ ldims) :| lmads) oshp cg) is = do - let perm = lmadPermutation lmad - is' = permuteInv perm is - cg' = cg && slicePreservesContiguous lmad is' - guard $ harmlessRotation lmad is' - let lmad' = foldl sliceOne (LMAD (lmadOffset lmad) []) $ zip is' ldims - -- need to remove the fixed dims from the permutation - perm' = updatePerm perm $ map fst $ filter (isJust . dimFix . snd) $ - zip [0..length is' - 1] is' - - return $ IxFun (setLMADPermutation perm' lmad' :| lmads) oshp cg' - where updatePerm ps inds = foldl (\acc p -> acc ++ decrease p) [] ps - where decrease p = - let d = foldl (\n i -> if i == p then -1 - else if i > p - then n - else if n /= -1 then n + 1 - else n - ) 0 inds - in [p - d | d /= -1] - - harmlessRotation' :: (Eq num, IntegralExp num) => - LMADDim num -> DimIndex num -> Bool - harmlessRotation' _ (DimFix _) = True - harmlessRotation' (LMADDim 0 _ _ _ _) _ = True - harmlessRotation' (LMADDim _ 0 _ _ _) _ = True - harmlessRotation' (LMADDim _ _ n _ _) dslc - | dslc == DimSlice (n - 1) n (-1) || - dslc == unitSlice 0 n = True - harmlessRotation' _ _ = False - - harmlessRotation :: (Eq num, IntegralExp num) => - LMAD num -> Slice num -> Bool - harmlessRotation (LMAD _ dims) iss = - and $ zipWith harmlessRotation' dims iss - - -- XXX: TODO: what happens to r on a negative-stride slice; is there - -- such a case? - sliceOne :: (Eq num, IntegralExp num) => - LMAD num -> (DimIndex num, LMADDim num) -> LMAD num - sliceOne (LMAD off dims) (DimFix i, LMADDim s r n _ _) = - LMAD (off + flatOneDim (s, r, n) i) dims - sliceOne (LMAD off dims) (DimSlice _ ne _, LMADDim 0 _ _ p _) = - LMAD off (dims ++ [LMADDim 0 0 ne p Unknown]) - sliceOne (LMAD off dims) (dmind, dim@(LMADDim _ _ n _ _)) - | dmind == unitSlice 0 n = LMAD off (dims ++ [dim]) - sliceOne (LMAD off dims) (dmind, LMADDim s r n p m) - | dmind == DimSlice (n - 1) n (-1) = - let r' = if r == 0 then 0 else n - r - off' = off + flatOneDim (s, 0, n) (n - 1) - in LMAD off' (dims ++ [LMADDim (s * (-1)) r' n p (invertMonotonicity m)]) - sliceOne (LMAD off dims) (DimSlice b ne 0, LMADDim s r n p _) = - LMAD (off + flatOneDim (s, r, n) b) (dims ++ [LMADDim 0 0 ne p Unknown]) - sliceOne (LMAD off dims) (DimSlice bs ns ss, LMADDim s 0 _ p m) = - let m' = case sgn ss of - Just 1 -> m - Just (-1) -> invertMonotonicity m - _ -> Unknown - in LMAD (off + s * bs) (dims ++ [LMADDim (ss * s) 0 ns p m']) - sliceOne _ _ = error "slice: reached impossible case" - - slicePreservesContiguous :: (Eq num, IntegralExp num) => - LMAD num -> Slice num -> Bool - slicePreservesContiguous (LMAD _ dims) slc = - -- remove from the slice the LMAD dimensions that have stride 0. - -- If the LMAD was contiguous in mem, then these dims will not - -- influence the contiguousness of the result. - -- Also normalize the input slice, i.e., 0-stride and size-1 - -- slices are rewritten as DimFixed. - let (dims', slc') = unzip $ - filter ((/= 0) . ldStride . fst) $ - zip dims $ map normIndex slc - -- Check that: - -- 1. a clean split point exists between Fixed and Sliced dims - -- 2. the outermost sliced dim has +/- 1 stride AND is unrotated or full. - -- 3. the rest of inner sliced dims are full. - (_, success) = - foldl (\(found, res) (slcdim, LMADDim _ r n _ _) -> - case (slcdim, found) of - (DimFix{}, True ) -> (found, False) - (DimFix{}, False) -> (found, res) - (DimSlice _ ne ds, False) -> -- outermost sliced dim: +/-1 stride - let res' = (r == 0 || n == ne) && (ds == 1 || ds == -1) - in (True, res && res') - (DimSlice _ ne ds, True) -> -- inner sliced dim: needs to be full - let res' = (n == ne) && (ds == 1 || ds == -1) - in (found, res && res') - ) (False, True) $ zip slc' dims' - in success - - normIndex :: (Eq num, IntegralExp num) => - DimIndex num -> DimIndex num - normIndex (DimSlice b 1 _) = DimFix b - normIndex (DimSlice b _ 0) = DimFix b - normIndex d = d - --- | Slice an index function. -slice :: (Eq num, IntegralExp num) => - IxFun num -> Slice num -> IxFun num -slice _ [] = error "slice: empty slice" -slice ixfun@(IxFun (lmad@(LMAD _ _) :| lmads) oshp cg) dim_slices - -- Avoid identity slicing. - | dim_slices == map (unitSlice 0) (shape ixfun) = ixfun - | Just ixfun' <- sliceOneLMAD ixfun dim_slices = ixfun' - | otherwise = - case sliceOneLMAD (iota (lmadShape lmad)) dim_slices of - Just (IxFun (lmad' :| []) _ cg') -> - IxFun (lmad' :| lmad : lmads) oshp (cg && cg') - _ -> error "slice: reached impossible case" - --- | Handle the simple case where all reshape dimensions are coercions. -reshapeCoercion :: (Eq num, IntegralExp num) => - IxFun num -> ShapeChange num -> Maybe (IxFun num) -reshapeCoercion (IxFun (lmad@(LMAD off dims) :| lmads) _ cg) newshape = do - let perm = lmadPermutation lmad - (head_coercions, reshapes, tail_coercions) <- splitCoercions newshape - let hd_len = length head_coercions - num_coercions = hd_len + length tail_coercions - dims' = permuteFwd perm dims - mid_dims = take (length dims - num_coercions) $ drop hd_len dims' - num_rshps = length reshapes - guard (num_rshps == 0 || (num_rshps == 1 && length mid_dims == 1)) - let dims'' = permuteInv perm $ - zipWith (\ld n -> ld { ldShape = n }) - dims' (newDims newshape) - lmad' = LMAD off dims'' - return $ IxFun (lmad' :| lmads) (newDims newshape) cg - --- | Handle the case where a reshape operation can stay inside a single LMAD. --- --- There are four conditions that all must hold for the result of a reshape --- operation to remain in the one-LMAD domain: --- --- (1) the permutation of the underlying LMAD must leave unchanged --- the LMAD dimensions that were *not* reshape coercions. --- (2) the repetition of dimensions of the underlying LMAD must --- refer only to the coerced-dimensions of the reshape operation. --- (3) similarly, the rotated dimensions must refer only to --- dimensions that are coerced by the reshape operation. --- (4) finally, the underlying memory is contiguous (and monotonous). --- --- If any of these conditions do not hold, then the reshape operation will --- conservatively add a new LMAD to the list, leading to a representation that --- provides less opportunities for further analysis. -reshapeOneLMAD :: (Eq num, IntegralExp num) => - IxFun num -> ShapeChange num -> Maybe (IxFun num) -reshapeOneLMAD ixfun@(IxFun (lmad@(LMAD off dims) :| lmads) _ cg) newshape = do - let perm = lmadPermutation lmad - (head_coercions, reshapes, tail_coercions) <- splitCoercions newshape - let hd_len = length head_coercions - num_coercions = hd_len + length tail_coercions - dims_perm = permuteFwd perm dims - mid_dims = take (length dims - num_coercions) $ drop hd_len dims_perm - -- Ignore rotates, as we only care about not having rotates in the - -- dimensions that aren't coercions (@mid_dims@), which we check - -- separately. - mon = ixfunMonotonicityRots True ixfun - - guard $ - -- checking conditions (2) and (3) - all (\ (LMADDim s r _ _ _) -> s /= 0 && r == 0) mid_dims && - -- checking condition (1) - consecutive hd_len (map ldPerm mid_dims) && - -- checking condition (4) - hasContiguousPerm ixfun && cg && (mon == Inc || mon == Dec) - - -- make new permutation - let rsh_len = length reshapes - diff = length newshape - length dims - iota_shape = [0..length newshape-1] - perm' = map (\i -> let ind = if i < hd_len - then i else i - diff - in if (i >= hd_len) && (i < hd_len + rsh_len) - then i -- already checked mid_dims not affected - else let p = ldPerm (dims !! ind) - in if p < hd_len - then p - else p + diff - ) iota_shape - -- split the dimensions - (support_inds, repeat_inds) = - foldl (\(sup, rpt) (i, shpdim, ip) -> - case (i < hd_len, i >= hd_len + rsh_len, shpdim) of - (True, _, DimCoercion n) -> - case dims_perm !! i of - (LMADDim 0 _ _ _ _) -> ( sup, (ip, n) : rpt ) - (LMADDim _ r _ _ _) -> ( (ip, (r, n)) : sup, rpt ) - (_, True, DimCoercion n) -> - case dims_perm !! (i-diff) of - (LMADDim 0 _ _ _ _) -> ( sup, (ip, n) : rpt ) - (LMADDim _ r _ _ _) -> ( (ip, (r, n)) : sup, rpt ) - (False, False, _) -> - ( (ip, (0, newDim shpdim)) : sup, rpt ) - -- already checked that the reshaped - -- dims cannot be rotates - _ -> error "reshape: reached impossible case" - ) ([], []) $ reverse $ zip3 iota_shape newshape perm' - - (sup_inds, support) = unzip $ sortBy (compare `on` fst) support_inds - (rpt_inds, repeats) = unzip repeat_inds - LMAD off' dims_sup = makeRotIota mon off support - repeats' = map (\n -> LMADDim 0 0 n 0 Unknown) repeats - dims' = map snd $ sortBy (compare `on` fst) - $ zip sup_inds dims_sup ++ zip rpt_inds repeats' - lmad' = LMAD off' dims' - return $ IxFun (setLMADPermutation perm' lmad' :| lmads) (newDims newshape) cg - where consecutive _ [] = True - consecutive i [p]= i == p - consecutive i ps = and $ zipWith (==) ps [i, i+1..] - -splitCoercions :: (Eq num, IntegralExp num) => - ShapeChange num -> Maybe (ShapeChange num, ShapeChange num, ShapeChange num) -splitCoercions newshape' = do - let (head_coercions, newshape'') = span isCoercion newshape' - (reshapes, tail_coercions) = break isCoercion newshape'' - guard (all isCoercion tail_coercions) - return (head_coercions, reshapes, tail_coercions) - where isCoercion DimCoercion{} = True - isCoercion _ = False - --- | Reshape an index function. -reshape :: (Eq num, IntegralExp num) => - IxFun num -> ShapeChange num -> IxFun num -reshape ixfun new_shape - | Just ixfun' <- reshapeCoercion ixfun new_shape = ixfun' - | Just ixfun' <- reshapeOneLMAD ixfun new_shape = ixfun' -reshape (IxFun (lmad0 :| lmad0s) oshp cg) new_shape = - case iota (newDims new_shape) of - IxFun (lmad :| []) _ _ -> IxFun (lmad :| lmad0 : lmad0s) oshp cg - _ -> error "reshape: reached impossible case" - --- | The number of dimensions in the domain of the input function. -rank :: IntegralExp num => - IxFun num -> Int -rank (IxFun (LMAD _ sss :| _) _ _) = length sss - --- | Handle the case where a rebase operation can stay within m + n - 1 LMADs, --- where m is the number of LMADs in the index function, and n is the number of --- LMADs in the new base. If both index function have only on LMAD, this means --- that we stay within the single-LMAD domain. --- --- We can often stay in that domain if the original ixfun is essentially a --- slice, e.g. `x[i, (k1,m,s1), (k2,n,s2)] = orig`. --- --- XXX: TODO: handle repetitions in both lmads. --- --- How to handle repeated dimensions in the original? --- --- (a) Shave them off of the last lmad of original --- (b) Compose the result from (a) with the first --- lmad of the new base --- (c) apply a repeat operation on the result of (b). --- --- However, I strongly suspect that for in-place update what we need is actually --- the INVERSE of the rebase function, i.e., given an index function new-base --- and another one orig, compute the index function ixfun0 such that: --- --- new-base == rebase ixfun0 ixfun, or equivalently: --- new-base == ixfun o ixfun0 --- --- because then I can go bottom up and compose with ixfun0 all the index --- functions corresponding to the memory block associated with ixfun. -rebaseNice :: (Eq num, IntegralExp num) => - IxFun num -> IxFun num -> Maybe (IxFun num) -rebaseNice - new_base@(IxFun (lmad_base :| lmads_base) _ cg_base) - ixfun@(IxFun lmads shp cg) = do - let (lmad :| lmads') = NE.reverse lmads - dims = lmadDims lmad - perm = lmadPermutation lmad - perm_base = lmadPermutation lmad_base - - guard $ - -- Core rebase condition. - base ixfun == shape new_base - -- Conservative safety conditions: ixfun is contiguous and has known - -- monotonicity for all dimensions. - && cg && all ((/= Unknown) . ldMon) dims - -- XXX: We should be able to handle some basic cases where both index - -- functions have non-trivial permutations. - && (hasContiguousPerm ixfun || hasContiguousPerm new_base) - -- We need the permutations to be of the same size if we want to compose - -- them. They don't have to be of the same size if the ixfun has a trivial - -- permutation. Supporting this latter case allows us to rebase when ixfun - -- has been created by slicing with fixed dimensions. - && (length perm == length perm_base || hasContiguousPerm ixfun) - -- To not have to worry about ixfun having non-1 strides, we also check that - -- it is a row-major array (modulo permutation, which is handled - -- separately). Accept a non-full innermost dimension. XXX: Maybe this can - -- be less conservative? - && and (zipWith3 (\sn ld inner -> sn == ldShape ld || (inner && ldStride ld == 1)) - shp dims (replicate (length dims - 1) False ++ [True])) - - -- Compose permutations, reverse strides and adjust offset if necessary. - let perm_base' = if hasContiguousPerm ixfun - then perm_base - else map (perm !!) perm_base - lmad_base' = setLMADPermutation perm_base' lmad_base - dims_base = lmadDims lmad_base' - n_fewer_dims = length dims_base - length dims - (dims_base', offs_contrib) = unzip $ - zipWith (\(LMADDim s1 r1 n1 p1 _) (LMADDim _ r2 _ _ m2) -> - let (s', off') | m2 == Inc = (s1, 0) - | otherwise = (s1 * (-1), s1 * (n1 - 1)) - r' | m2 == Inc = if r2 == 0 then r1 else r1 + r2 - | r1 == 0 = r2 - | r2 == 0 = n1 - r1 - | otherwise = n1 - r1 + r2 - in (LMADDim s' r' n1 (p1 - n_fewer_dims) Inc, off')) - -- If @dims@ is morally a slice, it might have fewer dimensions than - -- @dims_base@. Drop extraneous outer dimensions. - (drop n_fewer_dims dims_base) dims - off_base = lmadOffset lmad_base' + sum offs_contrib - lmad_base'' - | lmadOffset lmad == 0 = LMAD off_base dims_base' - | otherwise = - -- If the innermost dimension of the ixfun was not full (but still - -- had a stride of 1), add its offset relative to the new base. - setLMADShape (lmadShape lmad) - (LMAD (off_base + ldStride (last dims_base) * lmadOffset lmad) - dims_base') - new_base' = IxFun (lmad_base'' :| lmads_base) shp cg_base - IxFun lmads_base' _ _ = new_base' - lmads'' = lmads' ++@ lmads_base' - return $ IxFun lmads'' shp (cg && cg_base) - --- | Rebase an index function on top of a new base. -rebase :: (Eq num, IntegralExp num) => - IxFun num -> IxFun num -> IxFun num -rebase new_base@(IxFun lmads_base shp_base cg_base) ixfun@(IxFun lmads shp cg) - | Just ixfun' <- rebaseNice new_base ixfun = ixfun' - -- In the general case just concatenate LMADs since this refers to index - -- function composition, which is always safe. - | otherwise = - let (lmads_base', shp_base') = - if base ixfun == shape new_base - then (lmads_base, shp_base) - else let IxFun lmads' shp_base'' _ = reshape new_base $ map DimCoercion shp - in (lmads', shp_base'') - in IxFun (lmads @++@ lmads_base') shp_base' (cg && cg_base) - -ixfunMonotonicity :: (Eq num, IntegralExp num) => IxFun num -> Monotonicity -ixfunMonotonicity = ixfunMonotonicityRots False - --- | If the memory support of the index function is contiguous and row-major --- (i.e., no transpositions, repetitions, rotates, etc.), then this should --- return the offset from which the memory-support of this index function --- starts. -linearWithOffset :: (Eq num, IntegralExp num) => - IxFun num -> num -> Maybe num -linearWithOffset ixfun@(IxFun (lmad :| []) _ cg) elem_size - | hasContiguousPerm ixfun && cg && ixfunMonotonicity ixfun == Inc = - Just $ lmadOffset lmad * elem_size -linearWithOffset _ _ = Nothing - --- | Similar restrictions to @linearWithOffset@ except for transpositions, which --- are returned together with the offset. -rearrangeWithOffset :: (Eq num, IntegralExp num) => - IxFun num -> num -> Maybe (num, [(Int,num)]) -rearrangeWithOffset (IxFun (lmad :| []) oshp cg) elem_size = do - -- Note that @cg@ describes whether the index function is - -- contiguous, *ignoring permutations*. This function requires that - -- functionality. - let perm = lmadPermutation lmad - perm_contig = [0..length perm-1] - offset <- linearWithOffset - (IxFun (setLMADPermutation perm_contig lmad :| []) oshp cg) elem_size - return (offset, zip perm (permuteFwd perm (lmadShapeBase lmad))) -rearrangeWithOffset _ _ = Nothing - --- | Is this a row-major array starting at offset zero? -isLinear :: (Eq num, IntegralExp num) => IxFun num -> Bool -isLinear = (== Just 0) . flip linearWithOffset 1 - -permuteFwd :: Permutation -> [a] -> [a] -permuteFwd ps elems = map (elems !!) ps - -permuteInv :: Permutation -> [a] -> [a] -permuteInv ps elems = map snd $ sortBy (compare `on` fst) $ zip ps elems - -flatOneDim :: (Eq num, IntegralExp num) => - (num, num, num) -> num -> num -flatOneDim (s, r, n) i - | s == 0 = 0 - | r == 0 = i * s - | otherwise = ((i + r) `mod` n) * s - --- | Generalised iota with user-specified offset and strides. -makeRotIota :: IntegralExp num => - Monotonicity -> num -> [(num, num)] -> LMAD num -makeRotIota mon off support - | mon == Inc || mon == Dec = - let rk = length support - (rs, ns) = unzip support - ss0 = reverse $ take rk $ scanl (*) 1 $ reverse ns - ss = if mon == Inc - then ss0 - else map (* (-1)) ss0 - ps = map fromIntegral [0..rk-1] - fi = replicate rk mon - in LMAD off $ zipWith5 LMADDim ss rs ns ps fi - | otherwise = error "makeRotIota: requires Inc or Dec" - --- | Check monotonicity of an index function. -ixfunMonotonicityRots :: (Eq num, IntegralExp num) => - Bool -> IxFun num -> Monotonicity -ixfunMonotonicityRots ignore_rots (IxFun (lmad :| lmads) _ _) = - let mon0 = lmadMonotonicityRots lmad - in if all ((== mon0) . lmadMonotonicityRots) lmads - then mon0 - else Unknown - where lmadMonotonicityRots :: (Eq num, IntegralExp num) => - LMAD num -> Monotonicity - lmadMonotonicityRots (LMAD _ dims) - | all (isMonDim Inc) dims = Inc - | all (isMonDim Dec) dims = Dec - | otherwise = Unknown - - isMonDim :: (Eq num, IntegralExp num) => - Monotonicity -> LMADDim num -> Bool - isMonDim mon (LMADDim s r _ _ ldmon) = - s == 0 || ((ignore_rots || r == 0) && mon == ldmon) - --- | Generalization (anti-unification) --- --- Anti-unification of two index functions is supported under the following conditions: --- 0. Both index functions are represented by ONE lmad (assumed common case!) --- 1. The support array of the two indexfuns have the same dimensionality --- (we can relax this condition if we use a 1D support, as we probably should!) --- 2. The contiguous property and the per-dimension monotonicity are the same --- (otherwise we might loose important information; this can be relaxed!) --- 3. Most importantly, both index functions correspond to the same permutation --- (since the permutation is represented by INTs, this restriction cannot --- be relaxed, unless we move to a gated-LMAD representation!) -leastGeneralGeneralization :: Eq v => IxFun (PrimExp v) -> IxFun (PrimExp v) -> - Maybe (IxFun (PrimExp (Ext v)), [(PrimExp v, PrimExp v)]) -leastGeneralGeneralization (IxFun (lmad1 :| []) oshp1 ctg1) (IxFun (lmad2 :| []) oshp2 ctg2) = do - guard $ - length oshp1 == length oshp2 && - ctg1 == ctg2 && - map ldPerm (lmadDims lmad1) == map ldPerm (lmadDims lmad2) && - lmadDMon lmad1 == lmadDMon lmad2 - let (ctg, dperm, dmon) = (ctg1, lmadPermutation lmad1, lmadDMon lmad1) - (dshp, m1) <- generalize [] (lmadDShp lmad1) (lmadDShp lmad2) - (oshp, m2) <- generalize m1 oshp1 oshp2 - (dstd, m3) <- generalize m2 (lmadDSrd lmad1) (lmadDSrd lmad2) - (drot, m4) <- generalize m3 (lmadDRot lmad1) (lmadDRot lmad2) - let (offt, m5) = PEG.leastGeneralGeneralization m4 (lmadOffset lmad1) (lmadOffset lmad2) - let lmad_dims = map (\(a,b,c,d,e) -> LMADDim a b c d e) $ - zip5 dstd drot dshp dperm dmon - lmad = LMAD offt lmad_dims - return (IxFun (lmad :| []) oshp ctg, m5) - where lmadDMon = map ldMon . lmadDims - lmadDSrd = map ldStride . lmadDims - lmadDShp = map ldShape . lmadDims - lmadDRot = map ldRotate . lmadDims - generalize m l1 l2 = - foldM (\(l_acc, m') (pe1,pe2) -> do - let (e, m'') = PEG.leastGeneralGeneralization m' pe1 pe2 - return (l_acc++[e], m'') - ) ([], m) (zip l1 l2) -leastGeneralGeneralization _ _ = Nothing - -isSequential :: [Int] -> Bool -isSequential xs = - all (uncurry (==)) $ zip xs [0..] - -existentializeExp :: PrimExp v -> State [PrimExp v] (PrimExp (Ext v)) -existentializeExp e = do - i <- gets length - modify (++ [e]) - let t = primExpType e - return $ LeafExp (Ext i) t - --- We require that there's only one lmad, and that the index function is contiguous, and the base shape has only one dimension -existentialize :: (Eq v, Pretty v) => - IxFun (PrimExp v) -> State [PrimExp v] (Maybe (IxFun (PrimExp (Ext v)))) -existentialize (IxFun (lmad :| []) oshp True) - | all ((== 0) . ldRotate) (lmadDims lmad), - length (lmadShape lmad) == length oshp, - isSequential (map ldPerm $ lmadDims lmad) = do - oshp' <- mapM existentializeExp oshp - lmadOffset' <- existentializeExp $ lmadOffset lmad - lmadDims' <- mapM existentializeLMADDim $ lmadDims lmad - let lmad' = LMAD lmadOffset' lmadDims' - return $ Just $ IxFun (lmad' :| []) oshp' True - where - existentializeLMADDim :: LMADDim (PrimExp v) -> State [PrimExp v] (LMADDim (PrimExp (Ext v))) - existentializeLMADDim (LMADDim str rot shp perm mon) = do - stride' <- existentializeExp str - shape' <- existentializeExp shp - return $ LMADDim stride' (fmap Free rot) shape' perm mon - - -- oshp' = LeafExp (Ext 0) - -- lmad' = LMAD lmadOffset' lmadDims' - -- lmadOffset' = LeafExp (Ext 1) - -- (_, lmadDims', lmadDimSubsts) = foldr generalizeDim (2, [], []) $ lmadDims lmad - -- substs = oshp : lmadOffset lmad' : lmadDimSubsts - - -- generalizeDim :: (Int, [LMADDim num]) -> LMADDim num -> (Int, [LMADDim num]) - -- generalizeDim (i, acc) (LMADDim stride rotate shape perm mon) = - -- (i + 3, - -- LMADDim (LeafExp $ Ext i) (LeafExp $ Ext $ i + 1) (LeafExp $ Ext $ i + 2) perm mon, - -- [stride, rotate, shape]) -existentialize _ = return Nothing - --- | When comparing index functions as part of the type check in KernelsMem, --- we may run into problems caused by the simplifier. As index functions can be --- generalized over if-then-else expressions, the simplifier might hoist some of --- the code from inside the if-then-else (computing the offset of an array, for --- instance), but now the type checker cannot verify that the generalized index --- function is valid, because some of the existentials are computed somewhere --- else. To Work around this, we've had to relax the KernelsMem type-checker --- a bit, specifically, we've introduced this function to verify whether two --- index functions are "close enough" that we can assume that they match. We use --- this instead of `ixfun1 == ixfun2` and hope that it's good enough. -closeEnough :: IxFun num -> IxFun num -> Bool -closeEnough ixf1 ixf2 = - (length (base ixf1) == length (base ixf2)) && - (NE.length (ixfunLMADs ixf1) == NE.length (ixfunLMADs ixf2)) && - all closeEnoughLMADs (NE.zip (ixfunLMADs ixf1) (ixfunLMADs ixf2)) - where - closeEnoughLMADs :: (LMAD num, LMAD num) -> Bool - closeEnoughLMADs (lmad1, lmad2) = - length (lmadDims lmad1) == length (lmadDims lmad2) && - map ldPerm (lmadDims lmad1) == - map ldPerm (lmadDims lmad2) +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE OverloadedStrings #-} +{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} + +-- | This module contains a representation for the index function based on +-- linear-memory accessor descriptors; see Zhu, Hoeflinger and David work. +module Futhark.IR.Mem.IxFun + ( IxFun (..), + index, + iota, + permute, + rotate, + reshape, + slice, + rebase, + shape, + rank, + linearWithOffset, + rearrangeWithOffset, + isDirect, + isLinear, + substituteInIxFun, + leastGeneralGeneralization, + existentialize, + closeEnough, + ) +where + +import Control.Category +import Control.Monad.Identity +import Control.Monad.State +import Control.Monad.Writer +import Data.Function (on) +import Data.List (sort, sortBy, zip4, zip5, zipWith5) +import Data.List.NonEmpty (NonEmpty (..)) +import qualified Data.List.NonEmpty as NE +import qualified Data.Map.Strict as M +import Data.Maybe (isJust) +import Futhark.Analysis.PrimExp + ( IntExp, + PrimExp (..), + TPrimExp (..), + primExpType, + ) +import Futhark.Analysis.PrimExp.Convert (substituteInPrimExp) +import qualified Futhark.Analysis.PrimExp.Generalize as PEG +import Futhark.IR.Prop +import Futhark.IR.Syntax + ( DimChange (..), + DimIndex (..), + ShapeChange, + Slice, + dimFix, + unitSlice, + ) +import Futhark.IR.Syntax.Core (Ext (..)) +import Futhark.Transform.Rename +import Futhark.Transform.Substitute +import Futhark.Util.IntegralExp +import Futhark.Util.Pretty +import GHC.Generics (Generic) +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, mod, (.)) + +type Shape num = [num] + +type Indices num = [num] + +type Permutation = [Int] + +data Monotonicity + = Inc + | Dec + | -- | monotonously increasing, decreasing or unknown + Unknown + deriving (Show, Eq, Generic) + +instance SexpIso Monotonicity where + sexpIso = + match $ + With (. Sexp.sym "inc") $ + With (. Sexp.sym "dec") $ + With + (. Sexp.sym "unknown") + End + +data LMADDim num = LMADDim + { ldStride :: num, + ldRotate :: num, + ldShape :: num, + ldPerm :: Int, + ldMon :: Monotonicity + } + deriving (Show, Eq, Generic) + +instance SexpIso num => SexpIso (LMADDim num) where + sexpIso = with $ \lmaddim -> + Sexp.list + ( Sexp.el (Sexp.sym "dim") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> lmaddim + +-- | LMAD's representation consists of a general offset and for each dimension a +-- stride, rotate factor, number of elements (or shape), permutation, and +-- monotonicity. Note that the permutation is not strictly necessary in that the +-- permutation can be performed directly on LMAD dimensions, but then it is +-- difficult to extract the permutation back from an LMAD. +-- +-- LMAD algebra is closed under composition w.r.t. operators such as +-- permute, index and slice. However, other operations, such as +-- reshape, cannot always be represented inside the LMAD algebra. +-- +-- It follows that the general representation of an index function is a list of +-- LMADS, in which each following LMAD in the list implicitly corresponds to an +-- irregular reshaping operation. +-- +-- However, we expect that the common case is when the index function is one +-- LMAD -- we call this the "nice" representation. +-- +-- Finally, the list of LMADs is kept in an @IxFun@ together with the shape of +-- the original array, and a bit to indicate whether the index function is +-- contiguous, i.e., if we instantiate all the points of the current index +-- function, do we get a contiguous memory interval? +-- +-- By definition, the LMAD denotes the set of points (simplified): +-- +-- \{ o + \Sigma_{j=0}^{k} ((i_j+r_j) `mod` n_j)*s_j, +-- \forall i_j such that 0<=i_j SexpIso (LMAD num) where + sexpIso = with $ \lmad -> + Sexp.list + ( Sexp.el (Sexp.sym "lmad") + >>> Sexp.el sexpIso + >>> Sexp.rest sexpIso + ) + >>> lmad + +-- | An index function is a mapping from a multidimensional array +-- index space (the domain) to a one-dimensional memory index space. +-- Essentially, it explains where the element at position @[i,j,p]@ of +-- some array is stored inside the flat one-dimensional array that +-- constitutes its memory. For example, we can use this to +-- distinguish row-major and column-major representations. +-- +-- An index function is represented as a sequence of 'LMAD's. +data IxFun num = IxFun + { ixfunLMADs :: NonEmpty (LMAD num), + base :: Shape num, + -- | ignoring permutations, is the index function contiguous? + ixfunContig :: Bool + } + deriving (Show, Eq, Generic) + +instance SexpIso num => SexpIso (IxFun num) where + sexpIso = with $ \ixfun -> + Sexp.list + ( Sexp.el (Sexp.sym "ixfun") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> ixfun + +instance Pretty Monotonicity where + ppr = text . show + +instance Pretty num => Pretty (LMAD num) where + ppr (LMAD offset dims) = + braces $ + semisep + [ text "offset: " <> oneLine (ppr offset), + text "strides: " <> p ldStride, + text "rotates: " <> p ldRotate, + text "shape: " <> p ldShape, + text "permutation: " <> p ldPerm, + text "monotonicity: " <> p ldMon + ] + where + p f = oneLine $ brackets $ commasep $ map (ppr . f) dims + +instance Pretty num => Pretty (IxFun num) where + ppr (IxFun lmads oshp cg) = + braces $ + semisep + [ text "base: " <> brackets (commasep $ map ppr oshp), + text "contiguous: " <> text (show cg), + text "LMADs: " <> brackets (commasep $ NE.toList $ NE.map ppr lmads) + ] + +instance Substitute num => Substitute (LMAD num) where + substituteNames substs = fmap $ substituteNames substs + +instance Substitute num => Substitute (IxFun num) where + substituteNames substs = fmap $ substituteNames substs + +instance Substitute num => Rename (LMAD num) where + rename = substituteRename + +instance Substitute num => Rename (IxFun num) where + rename = substituteRename + +instance FreeIn num => FreeIn (LMAD num) where + freeIn' = foldMap freeIn' + +instance FreeIn num => FreeIn (IxFun num) where + freeIn' = foldMap freeIn' + +instance Functor LMAD where + fmap f = runIdentity . traverse (return . f) + +instance Functor IxFun where + fmap f = runIdentity . traverse (return . f) + +instance Foldable LMAD where + foldMap f = execWriter . traverse (tell . f) + +instance Foldable IxFun where + foldMap f = execWriter . traverse (tell . f) + +instance Traversable LMAD where + traverse f (LMAD offset dims) = + LMAD <$> f offset <*> traverse f' dims + where + f' (LMADDim s r n p m) = + LMADDim <$> f s <*> f r <*> f n <*> pure p <*> pure m + +instance Traversable IxFun where + traverse f (IxFun lmads oshp cg) = + IxFun <$> traverse (traverse f) lmads <*> traverse f oshp <*> pure cg + +(++@) :: [a] -> NonEmpty a -> NonEmpty a +es ++@ (ne :| nes) = case es of + e : es' -> e :| es' ++ [ne] ++ nes + [] -> ne :| nes + +(@++@) :: NonEmpty a -> NonEmpty a -> NonEmpty a +(x :| xs) @++@ (y :| ys) = x :| xs ++ [y] ++ ys + +invertMonotonicity :: Monotonicity -> Monotonicity +invertMonotonicity Inc = Dec +invertMonotonicity Dec = Inc +invertMonotonicity Unknown = Unknown + +lmadPermutation :: LMAD num -> Permutation +lmadPermutation = map ldPerm . lmadDims + +setLMADPermutation :: Permutation -> LMAD num -> LMAD num +setLMADPermutation perm lmad = + lmad {lmadDims = zipWith (\dim p -> dim {ldPerm = p}) (lmadDims lmad) perm} + +setLMADShape :: Shape num -> LMAD num -> LMAD num +setLMADShape shp lmad = lmad {lmadDims = zipWith (\dim s -> dim {ldShape = s}) (lmadDims lmad) shp} + +-- | Substitute a name with a PrimExp in an LMAD. +substituteInLMAD :: + Ord a => + M.Map a (PrimExp a) -> + LMAD (PrimExp a) -> + LMAD (PrimExp a) +substituteInLMAD tab (LMAD offset dims) = + let offset' = substituteInPrimExp tab offset + dims' = + map + ( \(LMADDim s r n p m) -> + LMADDim + (substituteInPrimExp tab s) + (substituteInPrimExp tab r) + (substituteInPrimExp tab n) + p + m + ) + dims + in LMAD offset' dims' + +-- | Substitute a name with a PrimExp in an index function. +substituteInIxFun :: + Ord a => + M.Map a (TPrimExp t a) -> + IxFun (TPrimExp t a) -> + IxFun (TPrimExp t a) +substituteInIxFun tab (IxFun lmads oshp cg) = + IxFun + (NE.map (fmap TPrimExp . substituteInLMAD tab' . fmap untyped) lmads) + (map (TPrimExp . substituteInPrimExp tab' . untyped) oshp) + cg + where + tab' = fmap untyped tab + +-- | Is this is a row-major array? +isDirect :: (Eq num, IntegralExp num) => IxFun num -> Bool +isDirect ixfun@(IxFun (LMAD offset dims :| []) oshp True) = + let strides_expected = reverse $ scanl (*) 1 (reverse (tail oshp)) + in hasContiguousPerm ixfun + && length oshp == length dims + && offset == 0 + && all + ( \(LMADDim s r n p _, m, d, se) -> + s == se && r == 0 && n == d && p == m + ) + (zip4 dims [0 .. length dims - 1] oshp strides_expected) +isDirect _ = False + +-- | Does the index function have an ascending permutation? +hasContiguousPerm :: IxFun num -> Bool +hasContiguousPerm (IxFun (lmad :| []) _ _) = + let perm = lmadPermutation lmad + in perm == sort perm +hasContiguousPerm _ = False + +-- | Shape of an index function. +shape :: (Eq num, IntegralExp num) => IxFun num -> Shape num +shape (IxFun (lmad :| _) _ _) = lmadShape lmad + +-- | Shape of an LMAD. +lmadShape :: (Eq num, IntegralExp num) => LMAD num -> Shape num +lmadShape lmad = permuteInv (lmadPermutation lmad) $ lmadShapeBase lmad + +-- | Shape of an LMAD, ignoring permutations. +lmadShapeBase :: (Eq num, IntegralExp num) => LMAD num -> Shape num +lmadShapeBase = map ldShape . lmadDims + +-- | Compute the flat memory index for a complete set @inds@ of array indices +-- and a certain element size @elem_size@. +index :: + (IntegralExp num, Eq num) => + IxFun num -> + Indices num -> + num +index = indexFromLMADs . ixfunLMADs + where + indexFromLMADs :: + (IntegralExp num, Eq num) => + NonEmpty (LMAD num) -> + Indices num -> + num + indexFromLMADs (lmad :| []) inds = indexLMAD lmad inds + indexFromLMADs (lmad1 :| lmad2 : lmads) inds = + let i_flat = indexLMAD lmad1 inds + new_inds = unflattenIndex (permuteFwd (lmadPermutation lmad2) $ lmadShapeBase lmad2) i_flat + in indexFromLMADs (lmad2 :| lmads) new_inds + indexLMAD :: + (IntegralExp num, Eq num) => + LMAD num -> + Indices num -> + num + indexLMAD lmad@(LMAD off dims) inds = + let prod = + sum $ + zipWith + flatOneDim + (map (\(LMADDim s r n _ _) -> (s, r, n)) dims) + (permuteInv (lmadPermutation lmad) inds) + in off + prod + +-- | iota. +iota :: IntegralExp num => Shape num -> IxFun num +iota ns = + let rs = replicate (length ns) 0 + in IxFun (makeRotIota Inc 0 (zip rs ns) :| []) ns True + +-- | Permute dimensions. +permute :: + IntegralExp num => + IxFun num -> + Permutation -> + IxFun num +permute (IxFun (lmad :| lmads) oshp cg) perm_new = + let perm_cur = lmadPermutation lmad + perm = map (perm_cur !!) perm_new + in IxFun (setLMADPermutation perm lmad :| lmads) oshp cg + +-- | Rotate an index function. +rotate :: + (Eq num, IntegralExp num) => + IxFun num -> + Indices num -> + IxFun num +rotate (IxFun (lmad@(LMAD off dims) :| lmads) oshp cg) offs = + let dims' = + zipWith + ( \(LMADDim s r n p f) o -> + if s == 0 + then LMADDim 0 0 n p Unknown + else LMADDim s (r + o) n p f + ) + dims + (permuteInv (lmadPermutation lmad) offs) + in IxFun (LMAD off dims' :| lmads) oshp cg + +-- | Handle the case where a slice can stay within a single LMAD. +sliceOneLMAD :: + (Eq num, IntegralExp num) => + IxFun num -> + Slice num -> + Maybe (IxFun num) +sliceOneLMAD (IxFun (lmad@(LMAD _ ldims) :| lmads) oshp cg) is = do + let perm = lmadPermutation lmad + is' = permuteInv perm is + cg' = cg && slicePreservesContiguous lmad is' + guard $ harmlessRotation lmad is' + let lmad' = foldl sliceOne (LMAD (lmadOffset lmad) []) $ zip is' ldims + -- need to remove the fixed dims from the permutation + perm' = + updatePerm perm $ + map fst $ + filter (isJust . dimFix . snd) $ + zip [0 .. length is' - 1] is' + + return $ IxFun (setLMADPermutation perm' lmad' :| lmads) oshp cg' + where + updatePerm ps inds = foldl (\acc p -> acc ++ decrease p) [] ps + where + decrease p = + let d = + foldl + ( \n i -> + if i == p + then -1 + else + if i > p + then n + else + if n /= -1 + then n + 1 + else n + ) + 0 + inds + in [p - d | d /= -1] + + harmlessRotation' :: + (Eq num, IntegralExp num) => + LMADDim num -> + DimIndex num -> + Bool + harmlessRotation' _ (DimFix _) = True + harmlessRotation' (LMADDim 0 _ _ _ _) _ = True + harmlessRotation' (LMADDim _ 0 _ _ _) _ = True + harmlessRotation' (LMADDim _ _ n _ _) dslc + | dslc == DimSlice (n - 1) n (-1) + || dslc == unitSlice 0 n = + True + harmlessRotation' _ _ = False + + harmlessRotation :: + (Eq num, IntegralExp num) => + LMAD num -> + Slice num -> + Bool + harmlessRotation (LMAD _ dims) iss = + and $ zipWith harmlessRotation' dims iss + + -- XXX: TODO: what happens to r on a negative-stride slice; is there + -- such a case? + sliceOne :: + (Eq num, IntegralExp num) => + LMAD num -> + (DimIndex num, LMADDim num) -> + LMAD num + sliceOne (LMAD off dims) (DimFix i, LMADDim s r n _ _) = + LMAD (off + flatOneDim (s, r, n) i) dims + sliceOne (LMAD off dims) (DimSlice _ ne _, LMADDim 0 _ _ p _) = + LMAD off (dims ++ [LMADDim 0 0 ne p Unknown]) + sliceOne (LMAD off dims) (dmind, dim@(LMADDim _ _ n _ _)) + | dmind == unitSlice 0 n = LMAD off (dims ++ [dim]) + sliceOne (LMAD off dims) (dmind, LMADDim s r n p m) + | dmind == DimSlice (n - 1) n (-1) = + let r' = if r == 0 then 0 else n - r + off' = off + flatOneDim (s, 0, n) (n - 1) + in LMAD off' (dims ++ [LMADDim (s * (-1)) r' n p (invertMonotonicity m)]) + sliceOne (LMAD off dims) (DimSlice b ne 0, LMADDim s r n p _) = + LMAD (off + flatOneDim (s, r, n) b) (dims ++ [LMADDim 0 0 ne p Unknown]) + sliceOne (LMAD off dims) (DimSlice bs ns ss, LMADDim s 0 _ p m) = + let m' = case sgn ss of + Just 1 -> m + Just (-1) -> invertMonotonicity m + _ -> Unknown + in LMAD (off + s * bs) (dims ++ [LMADDim (ss * s) 0 ns p m']) + sliceOne _ _ = error "slice: reached impossible case" + + slicePreservesContiguous :: + (Eq num, IntegralExp num) => + LMAD num -> + Slice num -> + Bool + slicePreservesContiguous (LMAD _ dims) slc = + -- remove from the slice the LMAD dimensions that have stride 0. + -- If the LMAD was contiguous in mem, then these dims will not + -- influence the contiguousness of the result. + -- Also normalize the input slice, i.e., 0-stride and size-1 + -- slices are rewritten as DimFixed. + let (dims', slc') = + unzip $ + filter ((/= 0) . ldStride . fst) $ + zip dims $ map normIndex slc + -- Check that: + -- 1. a clean split point exists between Fixed and Sliced dims + -- 2. the outermost sliced dim has +/- 1 stride AND is unrotated or full. + -- 3. the rest of inner sliced dims are full. + (_, success) = + foldl + ( \(found, res) (slcdim, LMADDim _ r n _ _) -> + case (slcdim, found) of + (DimFix {}, True) -> (found, False) + (DimFix {}, False) -> (found, res) + (DimSlice _ ne ds, False) -> + -- outermost sliced dim: +/-1 stride + let res' = (r == 0 || n == ne) && (ds == 1 || ds == -1) + in (True, res && res') + (DimSlice _ ne ds, True) -> + -- inner sliced dim: needs to be full + let res' = (n == ne) && (ds == 1 || ds == -1) + in (found, res && res') + ) + (False, True) + $ zip slc' dims' + in success + + normIndex :: + (Eq num, IntegralExp num) => + DimIndex num -> + DimIndex num + normIndex (DimSlice b 1 _) = DimFix b + normIndex (DimSlice b _ 0) = DimFix b + normIndex d = d + +-- | Slice an index function. +slice :: + (Eq num, IntegralExp num) => + IxFun num -> + Slice num -> + IxFun num +slice _ [] = error "slice: empty slice" +slice ixfun@(IxFun (lmad@(LMAD _ _) :| lmads) oshp cg) dim_slices + -- Avoid identity slicing. + | dim_slices == map (unitSlice 0) (shape ixfun) = ixfun + | Just ixfun' <- sliceOneLMAD ixfun dim_slices = ixfun' + | otherwise = + case sliceOneLMAD (iota (lmadShape lmad)) dim_slices of + Just (IxFun (lmad' :| []) _ cg') -> + IxFun (lmad' :| lmad : lmads) oshp (cg && cg') + _ -> error "slice: reached impossible case" + +-- | Handle the simple case where all reshape dimensions are coercions. +reshapeCoercion :: + (Eq num, IntegralExp num) => + IxFun num -> + ShapeChange num -> + Maybe (IxFun num) +reshapeCoercion (IxFun (lmad@(LMAD off dims) :| lmads) _ cg) newshape = do + let perm = lmadPermutation lmad + (head_coercions, reshapes, tail_coercions) <- splitCoercions newshape + let hd_len = length head_coercions + num_coercions = hd_len + length tail_coercions + dims' = permuteFwd perm dims + mid_dims = take (length dims - num_coercions) $ drop hd_len dims' + num_rshps = length reshapes + guard (num_rshps == 0 || (num_rshps == 1 && length mid_dims == 1)) + let dims'' = + permuteInv perm $ + zipWith + (\ld n -> ld {ldShape = n}) + dims' + (newDims newshape) + lmad' = LMAD off dims'' + return $ IxFun (lmad' :| lmads) (newDims newshape) cg + +-- | Handle the case where a reshape operation can stay inside a single LMAD. +-- +-- There are four conditions that all must hold for the result of a reshape +-- operation to remain in the one-LMAD domain: +-- +-- (1) the permutation of the underlying LMAD must leave unchanged +-- the LMAD dimensions that were *not* reshape coercions. +-- (2) the repetition of dimensions of the underlying LMAD must +-- refer only to the coerced-dimensions of the reshape operation. +-- (3) similarly, the rotated dimensions must refer only to +-- dimensions that are coerced by the reshape operation. +-- (4) finally, the underlying memory is contiguous (and monotonous). +-- +-- If any of these conditions do not hold, then the reshape operation will +-- conservatively add a new LMAD to the list, leading to a representation that +-- provides less opportunities for further analysis. +reshapeOneLMAD :: + (Eq num, IntegralExp num) => + IxFun num -> + ShapeChange num -> + Maybe (IxFun num) +reshapeOneLMAD ixfun@(IxFun (lmad@(LMAD off dims) :| lmads) _ cg) newshape = do + let perm = lmadPermutation lmad + (head_coercions, reshapes, tail_coercions) <- splitCoercions newshape + let hd_len = length head_coercions + num_coercions = hd_len + length tail_coercions + dims_perm = permuteFwd perm dims + mid_dims = take (length dims - num_coercions) $ drop hd_len dims_perm + -- Ignore rotates, as we only care about not having rotates in the + -- dimensions that aren't coercions (@mid_dims@), which we check + -- separately. + mon = ixfunMonotonicityRots True ixfun + + guard $ + -- checking conditions (2) and (3) + all (\(LMADDim s r _ _ _) -> s /= 0 && r == 0) mid_dims + && + -- checking condition (1) + consecutive hd_len (map ldPerm mid_dims) + && + -- checking condition (4) + hasContiguousPerm ixfun + && cg + && (mon == Inc || mon == Dec) + + -- make new permutation + let rsh_len = length reshapes + diff = length newshape - length dims + iota_shape = [0 .. length newshape -1] + perm' = + map + ( \i -> + let ind = + if i < hd_len + then i + else i - diff + in if (i >= hd_len) && (i < hd_len + rsh_len) + then i -- already checked mid_dims not affected + else + let p = ldPerm (dims !! ind) + in if p < hd_len + then p + else p + diff + ) + iota_shape + -- split the dimensions + (support_inds, repeat_inds) = + foldl + ( \(sup, rpt) (i, shpdim, ip) -> + case (i < hd_len, i >= hd_len + rsh_len, shpdim) of + (True, _, DimCoercion n) -> + case dims_perm !! i of + (LMADDim 0 _ _ _ _) -> (sup, (ip, n) : rpt) + (LMADDim _ r _ _ _) -> ((ip, (r, n)) : sup, rpt) + (_, True, DimCoercion n) -> + case dims_perm !! (i - diff) of + (LMADDim 0 _ _ _ _) -> (sup, (ip, n) : rpt) + (LMADDim _ r _ _ _) -> ((ip, (r, n)) : sup, rpt) + (False, False, _) -> + ((ip, (0, newDim shpdim)) : sup, rpt) + -- already checked that the reshaped + -- dims cannot be rotates + _ -> error "reshape: reached impossible case" + ) + ([], []) + $ reverse $ zip3 iota_shape newshape perm' + + (sup_inds, support) = unzip $ sortBy (compare `on` fst) support_inds + (rpt_inds, repeats) = unzip repeat_inds + LMAD off' dims_sup = makeRotIota mon off support + repeats' = map (\n -> LMADDim 0 0 n 0 Unknown) repeats + dims' = + map snd $ + sortBy (compare `on` fst) $ + zip sup_inds dims_sup ++ zip rpt_inds repeats' + lmad' = LMAD off' dims' + return $ IxFun (setLMADPermutation perm' lmad' :| lmads) (newDims newshape) cg + where + consecutive _ [] = True + consecutive i [p] = i == p + consecutive i ps = and $ zipWith (==) ps [i, i + 1 ..] + +splitCoercions :: + (Eq num, IntegralExp num) => + ShapeChange num -> + Maybe (ShapeChange num, ShapeChange num, ShapeChange num) +splitCoercions newshape' = do + let (head_coercions, newshape'') = span isCoercion newshape' + (reshapes, tail_coercions) = break isCoercion newshape'' + guard (all isCoercion tail_coercions) + return (head_coercions, reshapes, tail_coercions) + where + isCoercion DimCoercion {} = True + isCoercion _ = False + +-- | Reshape an index function. +reshape :: + (Eq num, IntegralExp num) => + IxFun num -> + ShapeChange num -> + IxFun num +reshape ixfun new_shape + | Just ixfun' <- reshapeCoercion ixfun new_shape = ixfun' + | Just ixfun' <- reshapeOneLMAD ixfun new_shape = ixfun' +reshape (IxFun (lmad0 :| lmad0s) oshp cg) new_shape = + case iota (newDims new_shape) of + IxFun (lmad :| []) _ _ -> IxFun (lmad :| lmad0 : lmad0s) oshp cg + _ -> error "reshape: reached impossible case" + +-- | The number of dimensions in the domain of the input function. +rank :: + IntegralExp num => + IxFun num -> + Int +rank (IxFun (LMAD _ sss :| _) _ _) = length sss + +-- | Handle the case where a rebase operation can stay within m + n - 1 LMADs, +-- where m is the number of LMADs in the index function, and n is the number of +-- LMADs in the new base. If both index function have only on LMAD, this means +-- that we stay within the single-LMAD domain. +-- +-- We can often stay in that domain if the original ixfun is essentially a +-- slice, e.g. `x[i, (k1,m,s1), (k2,n,s2)] = orig`. +-- +-- XXX: TODO: handle repetitions in both lmads. +-- +-- How to handle repeated dimensions in the original? +-- +-- (a) Shave them off of the last lmad of original +-- (b) Compose the result from (a) with the first +-- lmad of the new base +-- (c) apply a repeat operation on the result of (b). +-- +-- However, I strongly suspect that for in-place update what we need is actually +-- the INVERSE of the rebase function, i.e., given an index function new-base +-- and another one orig, compute the index function ixfun0 such that: +-- +-- new-base == rebase ixfun0 ixfun, or equivalently: +-- new-base == ixfun o ixfun0 +-- +-- because then I can go bottom up and compose with ixfun0 all the index +-- functions corresponding to the memory block associated with ixfun. +rebaseNice :: + (Eq num, IntegralExp num) => + IxFun num -> + IxFun num -> + Maybe (IxFun num) +rebaseNice + new_base@(IxFun (lmad_base :| lmads_base) _ cg_base) + ixfun@(IxFun lmads shp cg) = do + let (lmad :| lmads') = NE.reverse lmads + dims = lmadDims lmad + perm = lmadPermutation lmad + perm_base = lmadPermutation lmad_base + + guard $ + -- Core rebase condition. + base ixfun == shape new_base + -- Conservative safety conditions: ixfun is contiguous and has known + -- monotonicity for all dimensions. + && cg + && all ((/= Unknown) . ldMon) dims + -- XXX: We should be able to handle some basic cases where both index + -- functions have non-trivial permutations. + && (hasContiguousPerm ixfun || hasContiguousPerm new_base) + -- We need the permutations to be of the same size if we want to compose + -- them. They don't have to be of the same size if the ixfun has a trivial + -- permutation. Supporting this latter case allows us to rebase when ixfun + -- has been created by slicing with fixed dimensions. + && (length perm == length perm_base || hasContiguousPerm ixfun) + -- To not have to worry about ixfun having non-1 strides, we also check that + -- it is a row-major array (modulo permutation, which is handled + -- separately). Accept a non-full innermost dimension. XXX: Maybe this can + -- be less conservative? + && and + ( zipWith3 + (\sn ld inner -> sn == ldShape ld || (inner && ldStride ld == 1)) + shp + dims + (replicate (length dims - 1) False ++ [True]) + ) + + -- Compose permutations, reverse strides and adjust offset if necessary. + let perm_base' = + if hasContiguousPerm ixfun + then perm_base + else map (perm !!) perm_base + lmad_base' = setLMADPermutation perm_base' lmad_base + dims_base = lmadDims lmad_base' + n_fewer_dims = length dims_base - length dims + (dims_base', offs_contrib) = + unzip $ + zipWith + ( \(LMADDim s1 r1 n1 p1 _) (LMADDim _ r2 _ _ m2) -> + let (s', off') + | m2 == Inc = (s1, 0) + | otherwise = (s1 * (-1), s1 * (n1 - 1)) + r' + | m2 == Inc = if r2 == 0 then r1 else r1 + r2 + | r1 == 0 = r2 + | r2 == 0 = n1 - r1 + | otherwise = n1 - r1 + r2 + in (LMADDim s' r' n1 (p1 - n_fewer_dims) Inc, off') + ) + -- If @dims@ is morally a slice, it might have fewer dimensions than + -- @dims_base@. Drop extraneous outer dimensions. + (drop n_fewer_dims dims_base) + dims + off_base = lmadOffset lmad_base' + sum offs_contrib + lmad_base'' + | lmadOffset lmad == 0 = LMAD off_base dims_base' + | otherwise = + -- If the innermost dimension of the ixfun was not full (but still + -- had a stride of 1), add its offset relative to the new base. + setLMADShape + (lmadShape lmad) + ( LMAD + (off_base + ldStride (last dims_base) * lmadOffset lmad) + dims_base' + ) + new_base' = IxFun (lmad_base'' :| lmads_base) shp cg_base + IxFun lmads_base' _ _ = new_base' + lmads'' = lmads' ++@ lmads_base' + return $ IxFun lmads'' shp (cg && cg_base) + +-- | Rebase an index function on top of a new base. +rebase :: + (Eq num, IntegralExp num) => + IxFun num -> + IxFun num -> + IxFun num +rebase new_base@(IxFun lmads_base shp_base cg_base) ixfun@(IxFun lmads shp cg) + | Just ixfun' <- rebaseNice new_base ixfun = ixfun' + -- In the general case just concatenate LMADs since this refers to index + -- function composition, which is always safe. + | otherwise = + let (lmads_base', shp_base') = + if base ixfun == shape new_base + then (lmads_base, shp_base) + else + let IxFun lmads' shp_base'' _ = reshape new_base $ map DimCoercion shp + in (lmads', shp_base'') + in IxFun (lmads @++@ lmads_base') shp_base' (cg && cg_base) + +ixfunMonotonicity :: (Eq num, IntegralExp num) => IxFun num -> Monotonicity +ixfunMonotonicity = ixfunMonotonicityRots False + +-- | If the memory support of the index function is contiguous and row-major +-- (i.e., no transpositions, repetitions, rotates, etc.), then this should +-- return the offset from which the memory-support of this index function +-- starts. +linearWithOffset :: + (Eq num, IntegralExp num) => + IxFun num -> + num -> + Maybe num +linearWithOffset ixfun@(IxFun (lmad :| []) _ cg) elem_size + | hasContiguousPerm ixfun && cg && ixfunMonotonicity ixfun == Inc = + Just $ lmadOffset lmad * elem_size +linearWithOffset _ _ = Nothing + +-- | Similar restrictions to @linearWithOffset@ except for transpositions, which +-- are returned together with the offset. +rearrangeWithOffset :: + (Eq num, IntegralExp num) => + IxFun num -> + num -> + Maybe (num, [(Int, num)]) +rearrangeWithOffset (IxFun (lmad :| []) oshp cg) elem_size = do + -- Note that @cg@ describes whether the index function is + -- contiguous, *ignoring permutations*. This function requires that + -- functionality. + let perm = lmadPermutation lmad + perm_contig = [0 .. length perm -1] + offset <- + linearWithOffset + (IxFun (setLMADPermutation perm_contig lmad :| []) oshp cg) + elem_size + return (offset, zip perm (permuteFwd perm (lmadShapeBase lmad))) +rearrangeWithOffset _ _ = Nothing + +-- | Is this a row-major array starting at offset zero? +isLinear :: (Eq num, IntegralExp num) => IxFun num -> Bool +isLinear = (== Just 0) . flip linearWithOffset 1 + +permuteFwd :: Permutation -> [a] -> [a] +permuteFwd ps elems = map (elems !!) ps + +permuteInv :: Permutation -> [a] -> [a] +permuteInv ps elems = map snd $ sortBy (compare `on` fst) $ zip ps elems + +flatOneDim :: + (Eq num, IntegralExp num) => + (num, num, num) -> + num -> + num +flatOneDim (s, r, n) i + | s == 0 = 0 + | r == 0 = i * s + | otherwise = ((i + r) `mod` n) * s + +-- | Generalised iota with user-specified offset and strides. +makeRotIota :: + IntegralExp num => + Monotonicity -> + num -> + [(num, num)] -> + LMAD num +makeRotIota mon off support + | mon == Inc || mon == Dec = + let rk = length support + (rs, ns) = unzip support + ss0 = reverse $ take rk $ scanl (*) 1 $ reverse ns + ss = + if mon == Inc + then ss0 + else map (* (-1)) ss0 + ps = map fromIntegral [0 .. rk -1] + fi = replicate rk mon + in LMAD off $ zipWith5 LMADDim ss rs ns ps fi + | otherwise = error "makeRotIota: requires Inc or Dec" + +-- | Check monotonicity of an index function. +ixfunMonotonicityRots :: + (Eq num, IntegralExp num) => + Bool -> + IxFun num -> + Monotonicity +ixfunMonotonicityRots ignore_rots (IxFun (lmad :| lmads) _ _) = + let mon0 = lmadMonotonicityRots lmad + in if all ((== mon0) . lmadMonotonicityRots) lmads + then mon0 + else Unknown + where + lmadMonotonicityRots :: + (Eq num, IntegralExp num) => + LMAD num -> + Monotonicity + lmadMonotonicityRots (LMAD _ dims) + | all (isMonDim Inc) dims = Inc + | all (isMonDim Dec) dims = Dec + | otherwise = Unknown + + isMonDim :: + (Eq num, IntegralExp num) => + Monotonicity -> + LMADDim num -> + Bool + isMonDim mon (LMADDim s r _ _ ldmon) = + s == 0 || ((ignore_rots || r == 0) && mon == ldmon) + +-- | Generalization (anti-unification) +-- +-- Anti-unification of two index functions is supported under the following conditions: +-- 0. Both index functions are represented by ONE lmad (assumed common case!) +-- 1. The support array of the two indexfuns have the same dimensionality +-- (we can relax this condition if we use a 1D support, as we probably should!) +-- 2. The contiguous property and the per-dimension monotonicity are the same +-- (otherwise we might loose important information; this can be relaxed!) +-- 3. Most importantly, both index functions correspond to the same permutation +-- (since the permutation is represented by INTs, this restriction cannot +-- be relaxed, unless we move to a gated-LMAD representation!) +leastGeneralGeneralization :: + Eq v => + IxFun (PrimExp v) -> + IxFun (PrimExp v) -> + Maybe (IxFun (PrimExp (Ext v)), [(PrimExp v, PrimExp v)]) +leastGeneralGeneralization (IxFun (lmad1 :| []) oshp1 ctg1) (IxFun (lmad2 :| []) oshp2 ctg2) = do + guard $ + length oshp1 == length oshp2 + && ctg1 == ctg2 + && map ldPerm (lmadDims lmad1) == map ldPerm (lmadDims lmad2) + && lmadDMon lmad1 == lmadDMon lmad2 + let (ctg, dperm, dmon) = (ctg1, lmadPermutation lmad1, lmadDMon lmad1) + (dshp, m1) <- generalize [] (lmadDShp lmad1) (lmadDShp lmad2) + (oshp, m2) <- generalize m1 oshp1 oshp2 + (dstd, m3) <- generalize m2 (lmadDSrd lmad1) (lmadDSrd lmad2) + (drot, m4) <- generalize m3 (lmadDRot lmad1) (lmadDRot lmad2) + let (offt, m5) = PEG.leastGeneralGeneralization m4 (lmadOffset lmad1) (lmadOffset lmad2) + let lmad_dims = + map (\(a, b, c, d, e) -> LMADDim a b c d e) $ + zip5 dstd drot dshp dperm dmon + lmad = LMAD offt lmad_dims + return (IxFun (lmad :| []) oshp ctg, m5) + where + lmadDMon = map ldMon . lmadDims + lmadDSrd = map ldStride . lmadDims + lmadDShp = map ldShape . lmadDims + lmadDRot = map ldRotate . lmadDims + generalize m l1 l2 = + foldM + ( \(l_acc, m') (pe1, pe2) -> do + let (e, m'') = PEG.leastGeneralGeneralization m' pe1 pe2 + return (l_acc ++ [e], m'') + ) + ([], m) + (zip l1 l2) +leastGeneralGeneralization _ _ = Nothing + +isSequential :: [Int] -> Bool +isSequential xs = + all (uncurry (==)) $ zip xs [0 ..] + +existentializeExp :: TPrimExp t v -> State [TPrimExp t v] (TPrimExp t (Ext v)) +existentializeExp e = do + i <- gets length + modify (++ [e]) + let t = primExpType $ untyped e + return $ TPrimExp $ LeafExp (Ext i) t + +-- We require that there's only one lmad, and that the index function is contiguous, and the base shape has only one dimension +existentialize :: + (IntExp t, Eq v, Pretty v) => + IxFun (TPrimExp t v) -> + State [TPrimExp t v] (Maybe (IxFun (TPrimExp t (Ext v)))) +existentialize (IxFun (lmad :| []) oshp True) + | all ((== 0) . ldRotate) (lmadDims lmad), + length (lmadShape lmad) == length oshp, + isSequential (map ldPerm $ lmadDims lmad) = do + oshp' <- mapM existentializeExp oshp + lmadOffset' <- existentializeExp $ lmadOffset lmad + lmadDims' <- mapM existentializeLMADDim $ lmadDims lmad + let lmad' = LMAD lmadOffset' lmadDims' + return $ Just $ IxFun (lmad' :| []) oshp' True + where + existentializeLMADDim :: + LMADDim (TPrimExp t v) -> + State [TPrimExp t v] (LMADDim (TPrimExp t (Ext v))) + existentializeLMADDim (LMADDim str rot shp perm mon) = do + stride' <- existentializeExp str + shape' <- existentializeExp shp + return $ LMADDim stride' (fmap Free rot) shape' perm mon + +-- oshp' = LeafExp (Ext 0) +-- lmad' = LMAD lmadOffset' lmadDims' +-- lmadOffset' = LeafExp (Ext 1) +-- (_, lmadDims', lmadDimSubsts) = foldr generalizeDim (2, [], []) $ lmadDims lmad +-- substs = oshp : lmadOffset lmad' : lmadDimSubsts + +-- generalizeDim :: (Int, [LMADDim num]) -> LMADDim num -> (Int, [LMADDim num]) +-- generalizeDim (i, acc) (LMADDim stride rotate shape perm mon) = +-- (i + 3, +-- LMADDim (LeafExp $ Ext i) (LeafExp $ Ext $ i + 1) (LeafExp $ Ext $ i + 2) perm mon, +-- [stride, rotate, shape]) +existentialize _ = return Nothing + +-- | When comparing index functions as part of the type check in KernelsMem, +-- we may run into problems caused by the simplifier. As index functions can be +-- generalized over if-then-else expressions, the simplifier might hoist some of +-- the code from inside the if-then-else (computing the offset of an array, for +-- instance), but now the type checker cannot verify that the generalized index +-- function is valid, because some of the existentials are computed somewhere +-- else. To Work around this, we've had to relax the KernelsMem type-checker +-- a bit, specifically, we've introduced this function to verify whether two +-- index functions are "close enough" that we can assume that they match. We use +-- this instead of `ixfun1 == ixfun2` and hope that it's good enough. +closeEnough :: IxFun num -> IxFun num -> Bool +closeEnough ixf1 ixf2 = + (length (base ixf1) == length (base ixf2)) + && (NE.length (ixfunLMADs ixf1) == NE.length (ixfunLMADs ixf2)) + && all closeEnoughLMADs (NE.zip (ixfunLMADs ixf1) (ixfunLMADs ixf2)) + where + closeEnoughLMADs :: (LMAD num, LMAD num) -> Bool + closeEnoughLMADs (lmad1, lmad2) = + length (lmadDims lmad1) == length (lmadDims lmad2) + && map ldPerm (lmadDims lmad1) + == map ldPerm (lmadDims lmad2) diff --git a/src/Futhark/IR/Mem/Simplify.hs b/src/Futhark/IR/Mem/Simplify.hs --- a/src/Futhark/IR/Mem/Simplify.hs +++ b/src/Futhark/IR/Mem/Simplify.hs @@ -1,101 +1,121 @@ +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE TypeFamilies #-} + module Futhark.IR.Mem.Simplify - ( simplifyProgGeneric - , simplifyStmsGeneric - , simpleGeneric - , SimplifyMemory - ) + ( simplifyProgGeneric, + simplifyStmsGeneric, + simpleGeneric, + SimplifyMemory, + ) where import Control.Monad import Data.List (find) - -import qualified Futhark.IR.Syntax as AST -import Futhark.IR.Syntax - hiding (Prog, BasicOp, Exp, Body, Stm, - Pattern, PatElem, Lambda, FunDef, FParam, LParam, RetType) -import Futhark.IR.Mem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Pass.ExplicitAllocations (simplifiable) import qualified Futhark.Analysis.SymbolTable as ST import qualified Futhark.Analysis.UsageTable as UT -import qualified Futhark.Optimise.Simplify.Engine as Engine -import qualified Futhark.Optimise.Simplify as Simplify import Futhark.Construct -import Futhark.Pass -import Futhark.Optimise.Simplify.Rules -import Futhark.Optimise.Simplify.Rule +import Futhark.IR.Mem +import qualified Futhark.IR.Mem.IxFun as IxFun +import qualified Futhark.IR.Syntax as AST +import qualified Futhark.Optimise.Simplify as Simplify +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.Pass +import Futhark.Pass.ExplicitAllocations (simplifiable) import Futhark.Util -simpleGeneric :: (SimplifyMemory lore, Op lore ~ MemOp inner) => - (OpWithWisdom inner -> UT.UsageTable) - -> Simplify.SimplifyOp lore inner - -> Simplify.SimpleOps lore +simpleGeneric :: + (SimplifyMemory lore, Op lore ~ MemOp inner) => + (OpWithWisdom inner -> UT.UsageTable) -> + Simplify.SimplifyOp lore inner -> + Simplify.SimpleOps lore simpleGeneric = simplifiable -simplifyProgGeneric :: (SimplifyMemory lore, Op lore ~ MemOp inner) => - Simplify.SimpleOps lore - -> Prog lore -> PassM (Prog lore) +simplifyProgGeneric :: + (SimplifyMemory lore, Op lore ~ MemOp inner) => + Simplify.SimpleOps lore -> + Prog lore -> + PassM (Prog lore) simplifyProgGeneric ops = - Simplify.simplifyProg ops callKernelRules - blockers { Engine.blockHoistBranch = blockAllocs } - where blockAllocs vtable _ (Let _ _ (Op Alloc{})) = - not $ ST.simplifyMemory vtable - -- Do not hoist statements that produce arrays. This is - -- because in the KernelsMem representation, multiple - -- arrays can be located in the same memory block, and moving - -- their creation out of a branch can thus cause memory - -- corruption. At this point in the compiler we have probably - -- already moved all the array creations that matter. - blockAllocs _ _ (Let pat _ _) = - not $ all primType $ patternTypes pat + Simplify.simplifyProg + ops + callKernelRules + blockers {Engine.blockHoistBranch = blockAllocs} + where + blockAllocs vtable _ (Let _ _ (Op Alloc {})) = + not $ ST.simplifyMemory vtable + -- Do not hoist statements that produce arrays. This is + -- because in the KernelsMem representation, multiple + -- arrays can be located in the same memory block, and moving + -- their creation out of a branch can thus cause memory + -- corruption. At this point in the compiler we have probably + -- already moved all the array creations that matter. + blockAllocs _ _ (Let pat _ _) = + not $ all primType $ patternTypes pat -simplifyStmsGeneric :: (HasScope lore m, MonadFreshNames m, - SimplifyMemory lore, Op lore ~ MemOp inner) => - Simplify.SimpleOps lore -> Stms lore - -> m (ST.SymbolTable (Wise lore), Stms lore) +simplifyStmsGeneric :: + ( HasScope lore m, + MonadFreshNames m, + SimplifyMemory lore, + Op lore ~ MemOp inner + ) => + Simplify.SimpleOps lore -> + Stms lore -> + m (ST.SymbolTable (Wise lore), Stms lore) simplifyStmsGeneric ops stms = do scope <- askScope - Simplify.simplifyStms ops callKernelRules blockers - scope stms + Simplify.simplifyStms + ops + callKernelRules + blockers + scope + stms isResultAlloc :: Op lore ~ MemOp op => Engine.BlockPred lore -isResultAlloc _ usage (Let (AST.Pattern [] [bindee]) _ (Op Alloc{})) = +isResultAlloc _ usage (Let (AST.Pattern [] [bindee]) _ (Op Alloc {})) = UT.isInResult (patElemName bindee) usage isResultAlloc _ _ _ = False isAlloc :: Op lore ~ MemOp op => Engine.BlockPred lore -isAlloc _ _ (Let _ _ (Op Alloc{})) = True -isAlloc _ _ _ = False +isAlloc _ _ (Let _ _ (Op Alloc {})) = True +isAlloc _ _ _ = False -blockers :: (Op lore ~ MemOp inner) => - Simplify.HoistBlockers lore -blockers = Engine.noExtraHoistBlockers { - Engine.blockHoistPar = isAlloc - , Engine.blockHoistSeq = isResultAlloc - , Engine.isAllocation = isAlloc mempty mempty - } +blockers :: + (Op lore ~ MemOp inner) => + Simplify.HoistBlockers lore +blockers = + Engine.noExtraHoistBlockers + { Engine.blockHoistPar = isAlloc, + Engine.blockHoistSeq = isResultAlloc, + Engine.isAllocation = isAlloc mempty mempty + } -- | Some constraints that must hold for the simplification rules to work. type SimplifyMemory lore = - (Simplify.SimplifiableLore lore, - ExpDec lore ~ (), - BodyDec lore ~ (), - AllocOp (Op (Wise lore)), - CanBeWise (Op lore), - BinderOps (Wise lore), - Mem lore) + ( Simplify.SimplifiableLore lore, + ExpDec lore ~ (), + BodyDec lore ~ (), + AllocOp (Op (Wise lore)), + CanBeWise (Op lore), + BinderOps (Wise lore), + Mem lore + ) callKernelRules :: SimplifyMemory lore => RuleBook (Wise lore) -callKernelRules = standardRules <> - ruleBook [RuleBasicOp copyCopyToCopy, - RuleBasicOp removeIdentityCopy, - RuleIf unExistentialiseMemory] [] +callKernelRules = + standardRules + <> ruleBook + [ RuleBasicOp copyCopyToCopy, + RuleBasicOp removeIdentityCopy, + RuleIf unExistentialiseMemory, + RuleOp decertifySafeAlloc + ] + [] -- | If a branch is returning some existential memory, but the size of -- the array is not existential, and the index function of the array @@ -106,101 +126,119 @@ | ST.simplifyMemory vtable, fixable <- foldl hasConcretisableMemory mempty $ patternElements pat, not $ null fixable = Simplify $ do - - -- Create non-existential memory blocks big enough to hold the - -- arrays. - (arr_to_mem, oldmem_to_mem) <- - fmap unzip $ forM fixable $ \(arr_pe, mem_size, oldmem, space) -> do + -- Create non-existential memory blocks big enough to hold the + -- arrays. + (arr_to_mem, oldmem_to_mem) <- + fmap unzip $ + forM fixable $ \(arr_pe, mem_size, oldmem, space) -> do size <- toSubExp "size" mem_size mem <- letExp "mem" $ Op $ allocOp size space return ((patElemName arr_pe, mem), (oldmem, mem)) - -- Update the branches to contain Copy expressions putting the - -- arrays where they are expected. - let updateBody body = insertStmsM $ do - res <- bodyBind body - resultBodyM =<< - zipWithM updateResult (patternElements pat) res - updateResult pat_elem (Var v) - | Just mem <- lookup (patElemName pat_elem) arr_to_mem, - (_, MemArray pt shape u (ArrayIn _ ixfun)) <- patElemDec pat_elem = do - v_copy <- newVName $ baseString v <> "_nonext_copy" - let v_pat = Pattern [] [PatElem v_copy $ - MemArray pt shape u $ ArrayIn mem ixfun] - addStm $ mkWiseLetStm v_pat (defAux ()) $ BasicOp (Copy v) - return $ Var v_copy - | Just mem <- lookup (patElemName pat_elem) oldmem_to_mem = - return $ Var mem - updateResult _ se = - return se - tbranch' <- updateBody tbranch - fbranch' <- updateBody fbranch - letBind pat $ If cond tbranch' fbranch' ifdec - where onlyUsedIn name here = not $ any ((name `nameIn`) . freeIn) $ - filter ((/=here) . patElemName) $ - patternValueElements pat - knownSize Constant{} = True - knownSize (Var v) = not $ inContext v - inContext = (`elem` patternContextNames pat) + -- Update the branches to contain Copy expressions putting the + -- arrays where they are expected. + let updateBody body = insertStmsM $ do + res <- bodyBind body + resultBodyM + =<< zipWithM updateResult (patternElements pat) res + updateResult pat_elem (Var v) + | Just mem <- lookup (patElemName pat_elem) arr_to_mem, + (_, MemArray pt shape u (ArrayIn _ ixfun)) <- patElemDec pat_elem = do + v_copy <- newVName $ baseString v <> "_nonext_copy" + let v_pat = + Pattern + [] + [ PatElem v_copy $ + MemArray pt shape u $ ArrayIn mem ixfun + ] + addStm $ mkWiseLetStm v_pat (defAux ()) $ BasicOp (Copy v) + return $ Var v_copy + | Just mem <- lookup (patElemName pat_elem) oldmem_to_mem = + return $ Var mem + updateResult _ se = + return se + tbranch' <- updateBody tbranch + fbranch' <- updateBody fbranch + letBind pat $ If cond tbranch' fbranch' ifdec + where + onlyUsedIn name here = + not $ + any ((name `nameIn`) . freeIn) $ + filter ((/= here) . patElemName) $ + patternValueElements pat + knownSize Constant {} = True + knownSize (Var v) = not $ inContext v + inContext = (`elem` patternContextNames pat) - hasConcretisableMemory fixable pat_elem - | (_, MemArray pt shape _ (ArrayIn mem ixfun)) <- patElemDec pat_elem, - Just (j, Mem space) <- - fmap patElemType <$> find ((mem==) . patElemName . snd) - (zip [(0::Int)..] $ patternElements pat), - Just tse <- maybeNth j $ bodyResult tbranch, - Just fse <- maybeNth j $ bodyResult fbranch, - mem `onlyUsedIn` patElemName pat_elem, - all knownSize (shapeDims shape), - not $ freeIn ixfun `namesIntersect` namesFromList (patternNames pat), - fse /= tse = - let mem_size = - sExt Int64 $ product $ primByteSize pt : IxFun.base ixfun - in (pat_elem, mem_size, mem, space) : fixable - | otherwise = - fixable + hasConcretisableMemory fixable pat_elem + | (_, MemArray pt shape _ (ArrayIn mem ixfun)) <- patElemDec pat_elem, + Just (j, Mem space) <- + fmap patElemType + <$> find + ((mem ==) . patElemName . snd) + (zip [(0 :: Int) ..] $ patternElements pat), + Just tse <- maybeNth j $ bodyResult tbranch, + Just fse <- maybeNth j $ bodyResult fbranch, + mem `onlyUsedIn` patElemName pat_elem, + all knownSize (shapeDims shape), + not $ freeIn ixfun `namesIntersect` namesFromList (patternNames pat), + fse /= tse = + let mem_size = + untyped $ product $ primByteSize pt : map sExt64 (IxFun.base ixfun) + in (pat_elem, mem_size, mem, space) : fixable + | otherwise = + fixable unExistentialiseMemory _ _ _ _ = Skip -- | If we are copying something that is itself a copy, just copy the -- original one instead. -copyCopyToCopy :: (BinderOps lore, - LetDec lore ~ (VarWisdom, MemBound u)) => - TopDownRuleBasicOp lore +copyCopyToCopy :: + ( BinderOps lore, + LetDec lore ~ (VarWisdom, MemBound u) + ) => + TopDownRuleBasicOp lore copyCopyToCopy vtable pat@(Pattern [] [pat_elem]) _ (Copy v1) | Just (BasicOp (Copy v2), v1_cs) <- ST.lookupExp v1 vtable, - Just (_, MemArray _ _ _ (ArrayIn srcmem src_ixfun)) <- ST.entryLetBoundDec =<< ST.lookup v1 vtable, - Just (Mem src_space) <- ST.lookupType srcmem vtable, - (_, MemArray _ _ _ (ArrayIn destmem dest_ixfun)) <- patElemDec pat_elem, - Just (Mem dest_space) <- ST.lookupType destmem vtable, - - src_space == dest_space, dest_ixfun == src_ixfun = - - Simplify $ certifying v1_cs $ letBind pat $ BasicOp $ Copy v2 - + src_space == dest_space, + dest_ixfun == src_ixfun = + Simplify $ certifying v1_cs $ letBind pat $ BasicOp $ Copy v2 copyCopyToCopy vtable pat _ (Copy v0) | Just (BasicOp (Rearrange perm v1), v0_cs) <- ST.lookupExp v0 vtable, Just (BasicOp (Copy v2), v1_cs) <- ST.lookupExp v1 vtable = Simplify $ do - v0' <- certifying (v0_cs<>v1_cs) $ - letExp "rearrange_v0" $ BasicOp $ Rearrange perm v2 - letBind pat $ BasicOp $ Copy v0' - + v0' <- + certifying (v0_cs <> v1_cs) $ + letExp "rearrange_v0" $ BasicOp $ Rearrange perm v2 + letBind pat $ BasicOp $ Copy v0' copyCopyToCopy _ _ _ _ = Skip -- | If the destination of a copy is the same as the source, just -- remove it. -removeIdentityCopy :: (BinderOps lore, - LetDec lore ~ (VarWisdom, MemBound u)) => - TopDownRuleBasicOp lore +removeIdentityCopy :: + ( BinderOps lore, + LetDec lore ~ (VarWisdom, MemBound u) + ) => + TopDownRuleBasicOp lore removeIdentityCopy vtable pat@(Pattern [] [pe]) _ (Copy v) | (_, MemArray _ _ _ (ArrayIn dest_mem dest_ixfun)) <- patElemDec pe, Just (_, MemArray _ _ _ (ArrayIn src_mem src_ixfun)) <- ST.entryLetBoundDec =<< ST.lookup v vtable, - dest_mem == src_mem, dest_ixfun == src_ixfun = - Simplify $ letBind pat $ BasicOp $ SubExp $ Var v - + dest_mem == src_mem, + dest_ixfun == src_ixfun = + Simplify $ letBind pat $ BasicOp $ SubExp $ Var v removeIdentityCopy _ _ _ _ = Skip + +-- If an allocation is statically known to be safe, then we can remove +-- the certificates on it. This can help hoist things that would +-- otherwise be stuck inside loops or branches. +decertifySafeAlloc :: SimplifyMemory lore => TopDownRuleOp (Wise lore) +decertifySafeAlloc _ pat (StmAux cs attrs _) op + | cs /= mempty, + [Mem _] <- patternTypes pat, + safeOp op = + Simplify $ attributing attrs $ letBind pat $ Op op +decertifySafeAlloc _ _ _ _ = Skip 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 @@ -1,26 +1,26 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} + -- | Futhark prettyprinter. This module defines 'Pretty' instances -- for the AST defined in "Futhark.IR.Syntax", -- but also a number of convenience functions if you don't want to use -- the interface from 'Pretty'. module Futhark.IR.Pretty - ( prettyTuple - , pretty - , PrettyAnnot (..) - , PrettyLore (..) - , ppTuple' + ( prettyTuple, + pretty, + PrettyAnnot (..), + PrettyLore (..), + ppTuple', ) - where - -import Data.Foldable (toList) -import Data.Maybe +where -import Futhark.Util.Pretty -import Futhark.IR.Prop.Patterns -import Futhark.IR.Syntax +import Data.Foldable (toList) +import Data.Maybe +import Futhark.IR.Prop.Patterns +import Futhark.IR.Syntax +import Futhark.Util.Pretty -- | Class for values that may have some prettyprinted annotation. class PrettyAnnot a where @@ -36,20 +36,28 @@ ppAnnot = const Nothing -- | The class of lores whose annotations can be prettyprinted. -class (Decorations lore, - Pretty (RetType lore), - Pretty (BranchType lore), - Pretty (Param (FParamInfo lore)), - Pretty (Param (LParamInfo lore)), - Pretty (PatElemT (LetDec lore)), - PrettyAnnot (PatElem lore), - PrettyAnnot (FParam lore), - PrettyAnnot (LParam lore), - Pretty (Op lore)) => PrettyLore lore where +class + ( Decorations lore, + Pretty (RetType lore), + Pretty (BranchType lore), + Pretty (Param (FParamInfo lore)), + Pretty (Param (LParamInfo lore)), + Pretty (PatElemT (LetDec lore)), + PrettyAnnot (PatElem lore), + PrettyAnnot (FParam lore), + PrettyAnnot (LParam lore), + Pretty (Op lore) + ) => + 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 + Just $ + stack $ + map (text . ("-- " ++)) $ + lines $ + pretty $ + text "Branch returns:" <+> ppTuple' ts ppExpLore _ _ = Nothing commastack :: [Doc] -> Doc @@ -62,7 +70,7 @@ ppr _ = mempty instance Pretty Commutativity where - ppr Commutative = text "commutative" + ppr Commutative = text "commutative" ppr Noncommutative = text "noncommutative" instance Pretty Shape where @@ -70,7 +78,7 @@ instance Pretty a => Pretty (Ext a) where ppr (Free e) = ppr e - ppr (Ext x) = text "?" <> text (show x) + ppr (Ext x) = text "?" <> text (show x) instance Pretty ExtShape where ppr = brackets . commasep . map ppr . shapeDims @@ -102,7 +110,7 @@ ppr ident = ppr (identType ident) <+> ppr (identName ident) instance Pretty SubExp where - ppr (Var v) = ppr v + ppr (Var v) = ppr v ppr (Constant v) = ppr v instance Pretty Certificates where @@ -115,8 +123,9 @@ instance PrettyLore lore => Pretty (Body lore) where ppr (Body _ stms res) | null stms = braces (commasep $ map ppr res) - | otherwise = stack (map ppr $ stmsToList stms) - text "in" <+> braces (commasep $ map ppr res) + | otherwise = + stack (map ppr $ stmsToList stms) + text "in" <+> braces (commasep $ map ppr res) instance Pretty Attr where ppr (AttrAtom v) = ppr v @@ -124,7 +133,8 @@ attrAnnots :: Attrs -> [Doc] attrAnnots = map f . toList . unAttrs - where f v = text "#[" <> ppr v <> text "]" + where + f v = text "#[" <> ppr v <> text "]" stmAttrAnnots :: Stm lore -> [Doc] stmAttrAnnots = attrAnnots . stmAuxAttrs . stmAux @@ -132,49 +142,53 @@ instance Pretty (PatElemT dec) => Pretty (PatternT dec) where ppr pat = ppPattern (patternContextElements pat) (patternValueElements pat) -instance Pretty (PatElemT b) => Pretty (PatElemT (a,b)) where +instance Pretty (PatElemT b) => Pretty (PatElemT (a, b)) where ppr = ppr . fmap snd instance Pretty (PatElemT Type) where ppr (PatElem name t) = ppr t <+> ppr name -instance Pretty (Param b) => Pretty (Param (a,b)) where +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 t + <+> ppr name instance Pretty (Param Type) where ppr (Param name t) = - ppr t <+> - ppr name + ppr t + <+> ppr name instance PrettyLore lore => Pretty (Stm lore) where ppr bnd@(Let pat (StmAux cs _ dec) 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' - where e' | linebreak = ppr cs ppr e - | otherwise = ppr cs <> ppr e - linebreak = case e of - DoLoop{} -> True - Op{} -> True - If{} -> True - Apply{} -> True - BasicOp ArrayLit{} -> False - BasicOp Assert{} -> True - _ -> cs /= mempty + 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' + where + e' + | linebreak = ppr cs ppr e + | otherwise = ppr cs <> ppr e + linebreak = case e of + DoLoop {} -> True + Op {} -> True + If {} -> True + Apply {} -> True + BasicOp ArrayLit {} -> False + BasicOp Assert {} -> True + _ -> cs /= mempty - stmannot = - case stmAttrAnnots bnd <> - mapMaybe ppAnnot (patternElements $ stmPattern bnd) of - [] -> id - annots -> (align (stack annots) ) + stmannot = + case stmAttrAnnots bnd + <> mapMaybe ppAnnot (patternElements $ stmPattern bnd) of + [] -> id + annots -> (align (stack annots) ) instance Pretty BasicOp where ppr (SubExp se) = ppr se @@ -184,20 +198,23 @@ ppr (ArrayLit es rt) = case rt of Array {} -> brackets $ commastack $ map ppr es - _ -> brackets $ commasep $ map ppr es + _ -> brackets $ commasep $ map ppr es 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) = text (convOpFun conv) <+> ppr fromtype <+> ppr x <+> text "to" <+> ppr totype - where (fromtype, totype) = convOpType conv + where + (fromtype, totype) = convOpType conv ppr (UnOp op e) = ppr op <+> pprPrec 9 e ppr (Index v idxs) = ppr v <> brackets (commasep (map ppr idxs)) ppr (Update src idxs se) = - ppr src <+> text "with" <+> brackets (commasep (map ppr idxs)) <+> - text "<-" <+> ppr se + ppr src <+> text "with" <+> brackets (commasep (map ppr idxs)) + <+> text "<-" + <+> ppr se ppr (Iota e x s et) = text "iota" <> et' <> apply [ppr e, ppr x, ppr s] - where et' = text $ show $ primBitSize $ IntType et + where + et' = text $ show $ primBitSize $ IntType et ppr (Replicate ne ve) = text "replicate" <> apply [ppr ne, align (ppr ve)] ppr (Scratch t shape) = @@ -217,64 +234,88 @@ instance Pretty a => Pretty (ErrorMsg a) where ppr (ErrorMsg parts) = commasep $ map p parts - where p (ErrorString s) = text $ show s - p (ErrorInt32 x) = ppr x + where + p (ErrorString s) = text $ show s + p (ErrorInt32 x) = ppr x + p (ErrorInt64 x) = ppr x instance PrettyLore lore => Pretty (Exp lore) where ppr (If c t f (IfDec _ ifsort)) = - text "if" <+> info' <+> ppr c - text "then" <+> maybeNest t <+> - text "else" <+> maybeNest f - where info' = case ifsort of IfNormal -> mempty - IfFallback -> text "" - IfEquiv -> text "" - maybeNest b | null $ bodyStms b = ppr b - | otherwise = nestedBlock "{" "}" $ ppr b + text "if" <+> info' <+> ppr c + text "then" + <+> maybeNest t + <+> text "else" + <+> maybeNest f + where + info' = case ifsort of + IfNormal -> mempty + IfFallback -> text "" + IfEquiv -> text "" + maybeNest b + | null $ bodyStms b = ppr b + | otherwise = nestedBlock "{" "}" $ ppr b ppr (BasicOp op) = ppr op ppr (Apply fname args _ (safety, _, _)) = text (nameToString fname) <> safety' <> apply (map (align . pprArg) args) - where pprArg (arg, Consume) = text "*" <> ppr arg - pprArg (arg, _) = ppr arg - safety' = case safety of Unsafe -> text "" - Safe -> mempty + where + pprArg (arg, Consume) = text "*" <> ppr arg + pprArg (arg, _) = ppr arg + safety' = case safety of + Unsafe -> text "" + Safe -> mempty ppr (Op op) = ppr op ppr (DoLoop ctx val form loopbody) = - annot (mapMaybe ppAnnot (ctxparams++valparams)) $ - text "loop" <+> ppPattern ctxparams valparams <+> - equals <+> ppTuple' (ctxinit++valinit) - (case form of - ForLoop i it bound [] -> - text "for" <+> align (ppr i <> text ":" <> ppr it <+> - text "<" <+> align (ppr bound)) - ForLoop i it bound loop_vars -> - annot (mapMaybe (ppAnnot . fst) loop_vars) $ - text "for" <+> align (ppr i <> text ":" <> ppr it <+> - text "<" <+> align (ppr bound) - stack (map pprLoopVar loop_vars)) - WhileLoop cond -> - text "while" <+> ppr cond - ) <+> text "do" <+> nestedBlock "{" "}" (ppr loopbody) - where (ctxparams, ctxinit) = unzip ctx - (valparams, valinit) = unzip val - pprLoopVar (p,a) = ppr p <+> text "in" <+> ppr a + annot (mapMaybe ppAnnot (ctxparams ++ valparams)) $ + text "loop" <+> ppPattern ctxparams valparams + <+> equals + <+> ppTuple' (ctxinit ++ valinit) + ( case form of + ForLoop i it bound [] -> + text "for" + <+> align + ( ppr i <> text ":" <> ppr it + <+> text "<" + <+> align (ppr bound) + ) + ForLoop i it bound loop_vars -> + annot (mapMaybe (ppAnnot . fst) loop_vars) $ + text "for" + <+> align + ( ppr i <> text ":" <> ppr it + <+> text "<" + <+> align (ppr bound) + stack (map pprLoopVar loop_vars) + ) + WhileLoop cond -> + text "while" <+> ppr cond + ) + <+> text "do" + <+> nestedBlock "{" "}" (ppr loopbody) + where + (ctxparams, ctxinit) = unzip ctx + (valparams, valinit) = unzip val + pprLoopVar (p, a) = ppr p <+> text "in" <+> ppr a instance PrettyLore lore => Pretty (Lambda lore) where ppr (Lambda [] _ []) = text "nilFn" ppr (Lambda params body rettype) = annot (mapMaybe ppAnnot params) $ - text "fn" <+> ppTuple' rettype <+/> - align (parens (commasep (map ppr params))) <+> - text "=>" indent 2 (ppr body) + text "fn" <+> ppTuple' rettype + <+/> align (parens (commasep (map ppr params))) + <+> text "=>" indent 2 (ppr body) instance PrettyLore lore => Pretty (FunDef lore) where ppr (FunDef entry attrs name rettype fparams body) = annot (mapMaybe ppAnnot fparams <> attrAnnots attrs) $ - text fun <+> ppTuple' rettype <+/> - text (nameToString name) <+> - apply (map ppr fparams) <+> - equals <+> nestedBlock "{" "}" (ppr body) - where fun | isJust entry = "entry" - | otherwise = "fun" + text fun <+> ppTuple' rettype + <+/> text (nameToString name) + <+> apply (map ppr fparams) + <+> equals + <+> nestedBlock "{" "}" (ppr body) + where + fun + | isJust entry = "entry" + | otherwise = "fun" instance PrettyLore lore => Pretty (Prog lore) where ppr (Prog consts funs) = @@ -282,10 +323,10 @@ instance Pretty d => Pretty (DimChange d) where ppr (DimCoercion se) = text "~" <> ppr se - ppr (DimNew se) = ppr se + ppr (DimNew se) = ppr se instance Pretty d => Pretty (DimIndex d) where - ppr (DimFix i) = ppr i + ppr (DimFix i) = ppr i ppr (DimSlice i n s) = ppr i <> text ":+" <> ppr n <> text "*" <> ppr s ppPattern :: (Pretty a, Pretty b) => [a] -> [b] -> Doc 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 @@ -1,1369 +1,1717 @@ -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE Safe #-} --- | Definitions of primitive types, the values that inhabit these --- types, and operations on these values. A primitive value can also --- be called a scalar. --- --- Essentially, this module describes the subset of the (internal) --- Futhark language that operates on primitive types. -module Futhark.IR.Primitive - ( -- * Types - IntType (..), allIntTypes - , FloatType (..), allFloatTypes - , PrimType (..), allPrimTypes - - -- * Values - , IntValue(..) - , intValue, intValueType, valueIntegral - , FloatValue(..) - , floatValue, floatValueType - , PrimValue(..) - , primValueType - , blankPrimValue - - -- * Operations - , Overflow (..) - , Safety(..) - , UnOp (..), allUnOps - , BinOp (..), allBinOps - , ConvOp (..), allConvOps - , CmpOp (..), allCmpOps - - -- ** Unary Operations - , doUnOp - , doComplement - , doAbs, doFAbs - , doSSignum, doUSignum - - -- ** Binary Operations - , doBinOp - , doAdd, doMul, doSDiv, doSMod - , doPow - - -- ** Conversion Operations - , doConvOp - , doZExt, doSExt - , doFPConv - , doFPToUI, doFPToSI - , doUIToFP, doSIToFP - , intToInt64, intToWord64 - - -- * Comparison Operations - , doCmpOp - , doCmpEq - , doCmpUlt, doCmpUle - , doCmpSlt, doCmpSle - , doFCmpLt, doFCmpLe - - -- * Type Of - , binOpType - , unOpType - , cmpOpType - , convOpType - - -- * Primitive functions - , primFuns - - -- * Utility - , zeroIsh - , zeroIshInt - , oneIsh - , oneIshInt - , negativeIsh - , primBitSize - , primByteSize - , intByteSize - , floatByteSize - , commutativeBinOp - - -- * Prettyprinting - , convOpFun - , prettySigned - ) - where - -import Control.Applicative -import qualified Data.Binary.Get as G -import qualified Data.Binary.Put as P -import Data.Bits -import Data.Fixed (mod') -- Weird location. -import Data.Int (Int16, Int32, Int64, Int8) -import qualified Data.Map as M -import Data.Word - -import Prelude - -import Futhark.Util.Pretty -import Futhark.Util (roundFloat, ceilFloat, floorFloat, - roundDouble, ceilDouble, floorDouble, - lgamma, lgammaf, tgamma, tgammaf) - --- | An integer type, ordered by size. Note that signedness is not a --- property of the type, but a property of the operations performed on --- values of these types. -data IntType = Int8 - | Int16 - | Int32 - | Int64 - deriving (Eq, Ord, Show, Enum, Bounded) - -instance Pretty IntType where - ppr Int8 = text "i8" - ppr Int16 = text "i16" - ppr Int32 = text "i32" - ppr Int64 = text "i64" - --- | A list of all integer types. -allIntTypes :: [IntType] -allIntTypes = [minBound..maxBound] - --- | A floating point type. -data FloatType = Float32 - | Float64 - deriving (Eq, Ord, Show, Enum, Bounded) - -instance Pretty FloatType where - ppr Float32 = text "f32" - ppr Float64 = text "f64" - --- | A list of all floating-point types. -allFloatTypes :: [FloatType] -allFloatTypes = [minBound..maxBound] - --- | Low-level primitive types. -data PrimType = IntType IntType - | FloatType FloatType - | Bool - | Cert - deriving (Eq, Ord, Show) - -instance Enum PrimType where - toEnum 0 = IntType Int8 - toEnum 1 = IntType Int16 - toEnum 2 = IntType Int32 - toEnum 3 = IntType Int64 - toEnum 4 = FloatType Float32 - toEnum 5 = FloatType Float64 - toEnum 6 = Bool - toEnum _ = Cert - - fromEnum (IntType Int8) = 0 - fromEnum (IntType Int16) = 1 - fromEnum (IntType Int32) = 2 - fromEnum (IntType Int64) = 3 - fromEnum (FloatType Float32) = 4 - fromEnum (FloatType Float64) = 5 - fromEnum Bool = 6 - fromEnum Cert = 7 - -instance Bounded PrimType where - minBound = IntType Int8 - maxBound = Cert - -instance Pretty PrimType where - ppr (IntType t) = ppr t - ppr (FloatType t) = ppr t - ppr Bool = text "bool" - ppr Cert = text "cert" - --- | A list of all primitive types. -allPrimTypes :: [PrimType] -allPrimTypes = map IntType allIntTypes ++ - map FloatType allFloatTypes ++ - [Bool, Cert] - --- | An integer value. -data IntValue = Int8Value !Int8 - | Int16Value !Int16 - | Int32Value !Int32 - | Int64Value !Int64 - deriving (Eq, Ord, Show) - -instance Pretty IntValue where - ppr (Int8Value v) = text $ show v ++ "i8" - ppr (Int16Value v) = text $ show v ++ "i16" - ppr (Int32Value v) = text $ show v ++ "i32" - ppr (Int64Value v) = text $ show v ++ "i64" - --- | Create an t'IntValue' from a type and an 'Integer'. -intValue :: Integral int => IntType -> int -> IntValue -intValue Int8 = Int8Value . fromIntegral -intValue Int16 = Int16Value . fromIntegral -intValue Int32 = Int32Value . fromIntegral -intValue Int64 = Int64Value . fromIntegral - --- | The type of an integer value. -intValueType :: IntValue -> IntType -intValueType Int8Value{} = Int8 -intValueType Int16Value{} = Int16 -intValueType Int32Value{} = Int32 -intValueType Int64Value{} = Int64 - --- | Convert an t'IntValue' to any 'Integral' type. -valueIntegral :: Integral int => IntValue -> int -valueIntegral (Int8Value v) = fromIntegral v -valueIntegral (Int16Value v) = fromIntegral v -valueIntegral (Int32Value v) = fromIntegral v -valueIntegral (Int64Value v) = fromIntegral v - --- | A floating-point value. -data FloatValue = Float32Value !Float - | Float64Value !Double - deriving (Show) - -instance Eq FloatValue where - Float32Value x == Float32Value y = isNaN x && isNaN y || x == y - Float64Value x == Float64Value y = isNaN x && isNaN y || x == y - Float32Value _ == Float64Value _ = False - Float64Value _ == Float32Value _ = False - --- The derived Ord instance does not handle NaNs correctly. -instance Ord FloatValue where - Float32Value x <= Float32Value y = x <= y - Float64Value x <= Float64Value y = x <= y - Float32Value _ <= Float64Value _ = True - Float64Value _ <= Float32Value _ = False - - Float32Value x < Float32Value y = x < y - Float64Value x < Float64Value y = x < y - Float32Value _ < Float64Value _ = True - Float64Value _ < Float32Value _ = False - - (>) = flip (<) - (>=) = flip (<=) - -instance Pretty FloatValue where - ppr (Float32Value v) - | isInfinite v, v >= 0 = text "f32.inf" - | isInfinite v, v < 0 = text "-f32.inf" - | isNaN v = text "f32.nan" - | otherwise = text $ show v ++ "f32" - ppr (Float64Value v) - | isInfinite v, v >= 0 = text "f64.inf" - | isInfinite v, v < 0 = text "-f64.inf" - | isNaN v = text "f64.nan" - | otherwise = text $ show v ++ "f64" - --- | Create a t'FloatValue' from a type and a 'Rational'. -floatValue :: Real num => FloatType -> num -> FloatValue -floatValue Float32 = Float32Value . fromRational . toRational -floatValue Float64 = Float64Value . fromRational . toRational - --- | The type of a floating-point value. -floatValueType :: FloatValue -> FloatType -floatValueType Float32Value{} = Float32 -floatValueType Float64Value{} = Float64 - --- | Non-array values. -data PrimValue = IntValue !IntValue - | FloatValue !FloatValue - | BoolValue !Bool - | Checked -- ^ The only value of type @cert@. - deriving (Eq, Ord, Show) - -instance Pretty PrimValue where - ppr (IntValue v) = ppr v - ppr (BoolValue True) = text "true" - ppr (BoolValue False) = text "false" - ppr (FloatValue v) = ppr v - ppr Checked = text "checked" - --- | The type of a basic value. -primValueType :: PrimValue -> PrimType -primValueType (IntValue v) = IntType $ intValueType v -primValueType (FloatValue v) = FloatType $ floatValueType v -primValueType BoolValue{} = Bool -primValueType Checked = Cert - --- | A "blank" value of the given primitive type - this is zero, or --- whatever is close to it. Don't depend on this value, but use it --- for e.g. creating arrays to be populated by do-loops. -blankPrimValue :: PrimType -> PrimValue -blankPrimValue (IntType Int8) = IntValue $ Int8Value 0 -blankPrimValue (IntType Int16) = IntValue $ Int16Value 0 -blankPrimValue (IntType Int32) = IntValue $ Int32Value 0 -blankPrimValue (IntType Int64) = IntValue $ Int64Value 0 -blankPrimValue (FloatType Float32) = FloatValue $ Float32Value 0.0 -blankPrimValue (FloatType Float64) = FloatValue $ Float64Value 0.0 -blankPrimValue Bool = BoolValue False -blankPrimValue Cert = Checked - --- | Various unary operators. It is a bit ad-hoc what is a unary --- operator and what is a built-in function. Perhaps these should all --- go away eventually. -data UnOp = Not -- ^ E.g., @! True == False@. - | Complement IntType -- ^ E.g., @~(~1) = 1@. - | Abs IntType -- ^ @abs(-2) = 2@. - | FAbs FloatType -- ^ @fabs(-2.0) = 2.0@. - | SSignum IntType -- ^ Signed sign function: @ssignum(-2)@ = -1. - | USignum IntType -- ^ Unsigned sign function: @usignum(2)@ = 1. - deriving (Eq, Ord, Show) - --- | What to do in case of arithmetic overflow. Futhark's semantics --- are that overflow does wraparound, but for generated code (like --- address arithmetic), it can be beneficial for overflow to be --- undefined behaviour, as it allows better optimisation of things --- such as GPU kernels. --- --- Note that all values of this type are considered equal for 'Eq' and --- 'Ord'. -data Overflow = OverflowWrap | OverflowUndef - deriving (Show) - -instance Eq Overflow where - _ == _ = True - -instance Ord Overflow where - _ `compare` _ = EQ - --- | Whether something is safe or unsafe (mostly function calls, and --- in the context of whether operations are dynamically checked). --- When we inline an 'Unsafe' function, we remove all safety checks in --- its body. The 'Ord' instance picks 'Unsafe' as being less than --- 'Safe'. --- --- For operations like integer division, a safe division will not --- explode the computer in case of division by zero, but instead --- return some unspecified value. This always involves a run-time --- check, so generally the unsafe variant is what the compiler will --- insert, but guarded by an explicit assertion elsewhere. Safe --- operations are useful when the optimiser wants to move e.g. a --- division to a location where the divisor may be zero, but where the --- result will only be used when it is non-zero (so it doesn't matter --- what result is provided with a zero divisor, as long as the program --- keeps running). -data Safety = Unsafe | Safe deriving (Eq, Ord, Show) - --- | Binary operators. These correspond closely to the binary operators in --- LLVM. Most are parametrised by their expected input and output --- types. -data BinOp = Add IntType Overflow -- ^ Integer addition. - | FAdd FloatType -- ^ Floating-point addition. - - | Sub IntType Overflow -- ^ Integer subtraction. - | FSub FloatType -- ^ Floating-point subtraction. - - | Mul IntType Overflow -- ^ Integer multiplication. - | FMul FloatType -- ^ Floating-point multiplication. - - | UDiv IntType Safety - -- ^ Unsigned integer division. Rounds towards - -- negativity infinity. Note: this is different - -- from LLVM. - | UDivUp IntType Safety - -- ^ Unsigned integer division. Rounds towards positive - -- infinity. - - | SDiv IntType Safety - -- ^ Signed integer division. Rounds towards - -- negativity infinity. Note: this is different - -- from LLVM. - | SDivUp IntType Safety - -- ^ Signed integer division. Rounds towards positive - -- infinity. - - | FDiv FloatType -- ^ Floating-point division. - | FMod FloatType -- ^ Floating-point modulus. - - | UMod IntType Safety - -- ^ Unsigned integer modulus; the countepart to 'UDiv'. - | SMod IntType Safety - -- ^ Signed integer modulus; the countepart to 'SDiv'. - - | SQuot IntType Safety - -- ^ Signed integer division. Rounds towards zero. This - -- corresponds to the @sdiv@ instruction in LLVM and - -- integer division in C. - | SRem IntType Safety - -- ^ Signed integer division. Rounds towards zero. This - -- corresponds to the @srem@ instruction in LLVM and - -- integer modulo in C. - - | SMin IntType - -- ^ Returns the smallest of two signed integers. - | UMin IntType - -- ^ Returns the smallest of two unsigned integers. - | FMin FloatType - -- ^ Returns the smallest of two floating-point numbers. - | SMax IntType - -- ^ Returns the greatest of two signed integers. - | UMax IntType - -- ^ Returns the greatest of two unsigned integers. - | FMax FloatType - -- ^ Returns the greatest of two floating-point numbers. - - | Shl IntType -- ^ Left-shift. - | LShr IntType -- ^ Logical right-shift, zero-extended. - | AShr IntType -- ^ Arithmetic right-shift, sign-extended. - - | And IntType -- ^ Bitwise and. - | Or IntType -- ^ Bitwise or. - | Xor IntType -- ^ Bitwise exclusive-or. - - | Pow IntType -- ^ Integer exponentiation. - | FPow FloatType -- ^ Floating-point exponentiation. - - | LogAnd -- ^ Boolean and - not short-circuiting. - | LogOr -- ^ Boolean or - not short-circuiting. - deriving (Eq, Ord, Show) - --- | Comparison operators are like 'BinOp's, but they always return a --- boolean value. The somewhat ugly constructor names are straight --- out of LLVM. -data CmpOp = CmpEq PrimType -- ^ All types equality. - | CmpUlt IntType -- ^ Unsigned less than. - | CmpUle IntType -- ^ Unsigned less than or equal. - | CmpSlt IntType -- ^ Signed less than. - | CmpSle IntType -- ^ Signed less than or equal. - - -- Comparison operators for floating-point values. TODO: extend - -- this to handle NaNs and such, like the LLVM fcmp instruction. - | FCmpLt FloatType -- ^ Floating-point less than. - | FCmpLe FloatType -- ^ Floating-point less than or equal. - - -- Boolean comparison. - | CmpLlt -- ^ Boolean less than. - | CmpLle -- ^ Boolean less than or equal. - deriving (Eq, Ord, Show) - --- | Conversion operators try to generalise the @from t0 x to t1@ --- instructions from LLVM. -data ConvOp = ZExt IntType IntType - -- ^ Zero-extend the former integer type to the latter. - -- If the new type is smaller, the result is a - -- truncation. - | SExt IntType IntType - -- ^ Sign-extend the former integer type to the latter. - -- If the new type is smaller, the result is a - -- truncation. - | FPConv FloatType FloatType - -- ^ Convert value of the former floating-point type to - -- the latter. If the new type is smaller, the result - -- is a truncation. - | FPToUI FloatType IntType - -- ^ Convert a floating-point value to the nearest - -- unsigned integer (rounding towards zero). - | FPToSI FloatType IntType - -- ^ Convert a floating-point value to the nearest - -- signed integer (rounding towards zero). - | UIToFP IntType FloatType - -- ^ Convert an unsigned integer to a floating-point value. - | SIToFP IntType FloatType - -- ^ Convert a signed integer to a floating-point value. - | IToB IntType - -- ^ Convert an integer to a boolean value. Zero - -- becomes false; anything else is true. - | BToI IntType - -- ^ Convert a boolean to an integer. True is converted - -- to 1 and False to 0. - deriving (Eq, Ord, Show) - --- | A list of all unary operators for all types. -allUnOps :: [UnOp] -allUnOps = Not : - map Complement [minBound..maxBound] ++ - map Abs [minBound..maxBound] ++ - map FAbs [minBound..maxBound] ++ - map SSignum [minBound..maxBound] ++ - map USignum [minBound..maxBound] - --- | A list of all binary operators for all types. -allBinOps :: [BinOp] -allBinOps = concat [ map (`Add` OverflowWrap) allIntTypes - , map FAdd allFloatTypes - , map (`Sub` OverflowWrap) allIntTypes - , map FSub allFloatTypes - , map (`Mul` OverflowWrap) allIntTypes - , map FMul allFloatTypes - , map (`UDiv` Unsafe) allIntTypes - , map (`UDivUp` Unsafe) allIntTypes - , map (`SDiv` Unsafe) allIntTypes - , map (`SDivUp` Unsafe) allIntTypes - , map FDiv allFloatTypes - , map FMod allFloatTypes - , map (`UMod` Unsafe) allIntTypes - , map (`SMod` Unsafe) allIntTypes - , map (`SQuot` Unsafe) allIntTypes - , map (`SRem` Unsafe) allIntTypes - , map SMin allIntTypes - , map UMin allIntTypes - , map FMin allFloatTypes - , map SMax allIntTypes - , map UMax allIntTypes - , map FMax allFloatTypes - , map Shl allIntTypes - , map LShr allIntTypes - , map AShr allIntTypes - , map And allIntTypes - , map Or allIntTypes - , map Xor allIntTypes - , map Pow allIntTypes - , map FPow allFloatTypes - , [LogAnd, LogOr] - ] - --- | A list of all comparison operators for all types. -allCmpOps :: [CmpOp] -allCmpOps = concat [ map CmpEq allPrimTypes - , map CmpUlt allIntTypes - , map CmpUle allIntTypes - , map CmpSlt allIntTypes - , map CmpSle allIntTypes - , map FCmpLt allFloatTypes - , map FCmpLe allFloatTypes - ] - --- | A list of all conversion operators for all types. -allConvOps :: [ConvOp] -allConvOps = concat [ ZExt <$> allIntTypes <*> allIntTypes - , SExt <$> allIntTypes <*> allIntTypes - , FPConv <$> allFloatTypes <*> allFloatTypes - , FPToUI <$> allFloatTypes <*> allIntTypes - , FPToSI <$> allFloatTypes <*> allIntTypes - , UIToFP <$> allIntTypes <*> allFloatTypes - , SIToFP <$> allIntTypes <*> allFloatTypes - , IToB <$> allIntTypes - , BToI <$> allIntTypes - ] - --- | Apply an 'UnOp' to an operand. Returns 'Nothing' if the --- application is mistyped. -doUnOp :: UnOp -> PrimValue -> Maybe PrimValue -doUnOp Not (BoolValue b) = Just $ BoolValue $ not b -doUnOp Complement{} (IntValue v) = Just $ IntValue $ doComplement v -doUnOp Abs{} (IntValue v) = Just $ IntValue $ doAbs v -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 _ _ = Nothing - --- | E.g., @~(~1) = 1@. -doComplement :: IntValue -> IntValue -doComplement v = intValue (intValueType v) $ complement $ intToInt64 v - --- | @abs(-2) = 2@. -doAbs :: IntValue -> IntValue -doAbs v = intValue (intValueType v) $ abs $ intToInt64 v - --- | @abs(-2.0) = 2.0@. -doFAbs :: FloatValue -> FloatValue -doFAbs v = floatValue (floatValueType v) $ abs $ floatToDouble v - --- | @ssignum(-2)@ = -1. -doSSignum :: IntValue -> IntValue -doSSignum v = intValue (intValueType v) $ signum $ intToInt64 v - --- | @usignum(-2)@ = -1. -doUSignum :: IntValue -> IntValue -doUSignum v = intValue (intValueType v) $ signum $ intToWord64 v - --- | Apply a 'BinOp' to an operand. Returns 'Nothing' if the --- application is mistyped, or outside the domain (e.g. division by --- zero). -doBinOp :: BinOp -> PrimValue -> PrimValue -> Maybe PrimValue -doBinOp Add{} = doIntBinOp doAdd -doBinOp FAdd{} = doFloatBinOp (+) (+) -doBinOp Sub{} = doIntBinOp doSub -doBinOp FSub{} = doFloatBinOp (-) (-) -doBinOp Mul{} = doIntBinOp doMul -doBinOp FMul{} = doFloatBinOp (*) (*) -doBinOp UDiv{} = doRiskyIntBinOp doUDiv -doBinOp UDivUp{} = doRiskyIntBinOp doUDivUp -doBinOp SDiv{} = doRiskyIntBinOp doSDiv -doBinOp SDivUp{} = doRiskyIntBinOp doSDivUp -doBinOp FDiv{} = doFloatBinOp (/) (/) -doBinOp FMod{} = doFloatBinOp mod' mod' -doBinOp UMod{} = doRiskyIntBinOp doUMod -doBinOp SMod{} = doRiskyIntBinOp doSMod -doBinOp SQuot{} = doRiskyIntBinOp doSQuot -doBinOp SRem{} = doRiskyIntBinOp doSRem -doBinOp SMin{} = doIntBinOp doSMin -doBinOp UMin{} = doIntBinOp doUMin -doBinOp FMin{} = doFloatBinOp min min -doBinOp SMax{} = doIntBinOp doSMax -doBinOp UMax{} = doIntBinOp doUMax -doBinOp FMax{} = doFloatBinOp max max -doBinOp Shl{} = doIntBinOp doShl -doBinOp LShr{} = doIntBinOp doLShr -doBinOp AShr{} = doIntBinOp doAShr -doBinOp And{} = doIntBinOp doAnd -doBinOp Or{} = doIntBinOp doOr -doBinOp Xor{} = doIntBinOp doXor -doBinOp Pow{} = doRiskyIntBinOp doPow -doBinOp FPow{} = doFloatBinOp (**) (**) -doBinOp LogAnd{} = doBoolBinOp (&&) -doBinOp LogOr{} = doBoolBinOp (||) - -doIntBinOp :: (IntValue -> IntValue -> IntValue) -> PrimValue -> PrimValue - -> Maybe PrimValue -doIntBinOp f (IntValue v1) (IntValue v2) = - Just $ IntValue $ f v1 v2 -doIntBinOp _ _ _ = Nothing - -doRiskyIntBinOp :: (IntValue -> IntValue -> Maybe IntValue) -> PrimValue -> PrimValue - -> Maybe PrimValue -doRiskyIntBinOp f (IntValue v1) (IntValue v2) = - IntValue <$> f v1 v2 -doRiskyIntBinOp _ _ _ = Nothing - -doFloatBinOp :: (Float -> Float -> Float) - -> (Double -> Double -> Double) - -> PrimValue -> PrimValue - -> Maybe PrimValue -doFloatBinOp f32 _ (FloatValue (Float32Value v1)) (FloatValue (Float32Value v2)) = - Just $ FloatValue $ Float32Value $ f32 v1 v2 -doFloatBinOp _ f64 (FloatValue (Float64Value v1)) (FloatValue (Float64Value v2)) = - Just $ FloatValue $ Float64Value $ f64 v1 v2 -doFloatBinOp _ _ _ _ = Nothing - -doBoolBinOp :: (Bool -> Bool -> Bool) -> PrimValue -> PrimValue - -> Maybe PrimValue -doBoolBinOp f (BoolValue v1) (BoolValue v2) = - Just $ BoolValue $ f v1 v2 -doBoolBinOp _ _ _ = Nothing - --- | Integer addition. -doAdd :: IntValue -> IntValue -> IntValue -doAdd v1 v2 = intValue (intValueType v1) $ intToInt64 v1 + intToInt64 v2 - --- | Integer subtraction. -doSub :: IntValue -> IntValue -> IntValue -doSub v1 v2 = intValue (intValueType v1) $ intToInt64 v1 - intToInt64 v2 - --- | Integer multiplication. -doMul :: IntValue -> IntValue -> IntValue -doMul v1 v2 = intValue (intValueType v1) $ intToInt64 v1 * intToInt64 v2 - --- | Unsigned integer division. Rounds towards negativity infinity. --- Note: this is different from LLVM. -doUDiv :: IntValue -> IntValue -> Maybe IntValue -doUDiv v1 v2 - | zeroIshInt v2 = Nothing - | otherwise = Just $ intValue (intValueType v1) $ - intToWord64 v1 `div` intToWord64 v2 - --- | Unsigned integer division. Rounds towards positive infinity. -doUDivUp :: IntValue -> IntValue -> Maybe IntValue -doUDivUp v1 v2 - | zeroIshInt v2 = Nothing - | otherwise = Just $ intValue (intValueType v1) $ - (intToWord64 v1 + intToWord64 v2 - 1) `div` intToWord64 v2 - --- | Signed integer division. Rounds towards negativity infinity. --- Note: this is different from LLVM. -doSDiv :: IntValue -> IntValue -> Maybe IntValue -doSDiv v1 v2 - | zeroIshInt v2 = Nothing - | otherwise = Just $ intValue (intValueType v1) $ - intToInt64 v1 `div` intToInt64 v2 - --- | Signed integer division. Rounds towards positive infinity. -doSDivUp :: IntValue -> IntValue -> Maybe IntValue -doSDivUp v1 v2 - | zeroIshInt v2 = Nothing - | otherwise = Just $ intValue (intValueType v1) $ - (intToInt64 v1 + intToInt64 v2 - 1) `div` intToInt64 v2 - --- | Unsigned integer modulus; the countepart to 'UDiv'. -doUMod :: IntValue -> IntValue -> Maybe IntValue -doUMod v1 v2 - | zeroIshInt v2 = Nothing - | otherwise = Just $ intValue (intValueType v1) $ intToWord64 v1 `mod` intToWord64 v2 - --- | Signed integer modulus; the countepart to 'SDiv'. -doSMod :: IntValue -> IntValue -> Maybe IntValue -doSMod v1 v2 - | zeroIshInt v2 = Nothing - | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `mod` intToInt64 v2 - --- | Signed integer division. Rounds towards zero. --- This corresponds to the @sdiv@ instruction in LLVM. -doSQuot :: IntValue -> IntValue -> Maybe IntValue -doSQuot v1 v2 - | zeroIshInt v2 = Nothing - | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `quot` intToInt64 v2 - --- | Signed integer division. Rounds towards zero. --- This corresponds to the @srem@ instruction in LLVM. -doSRem :: IntValue -> IntValue -> Maybe IntValue -doSRem v1 v2 - | zeroIshInt v2 = Nothing - | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `rem` intToInt64 v2 - --- | Minimum of two signed integers. -doSMin :: IntValue -> IntValue -> IntValue -doSMin v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `min` intToInt64 v2 - --- | Minimum of two unsigned integers. -doUMin :: IntValue -> IntValue -> IntValue -doUMin v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `min` intToWord64 v2 - --- | Maximum of two signed integers. -doSMax :: IntValue -> IntValue -> IntValue -doSMax v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `max` intToInt64 v2 - --- | Maximum of two unsigned integers. -doUMax :: IntValue -> IntValue -> IntValue -doUMax v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `max` intToWord64 v2 - --- | Left-shift. -doShl :: IntValue -> IntValue -> IntValue -doShl v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `shift` intToInt v2 - --- | Logical right-shift, zero-extended. -doLShr :: IntValue -> IntValue -> IntValue -doLShr v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `shift` negate (intToInt v2) - --- | Arithmetic right-shift, sign-extended. -doAShr :: IntValue -> IntValue -> IntValue -doAShr v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `shift` negate (intToInt v2) - --- | Bitwise and. -doAnd :: IntValue -> IntValue -> IntValue -doAnd v1 v2 = intValue (intValueType v1) $ intToWord64 v1 .&. intToWord64 v2 - --- | Bitwise or. -doOr :: IntValue -> IntValue -> IntValue -doOr v1 v2 = intValue (intValueType v1) $ intToWord64 v1 .|. intToWord64 v2 - --- | Bitwise exclusive-or. -doXor :: IntValue -> IntValue -> IntValue -doXor v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `xor` intToWord64 v2 - --- | Signed integer exponentatation. -doPow :: IntValue -> IntValue -> Maybe IntValue -doPow v1 v2 - | negativeIshInt v2 = Nothing - | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 ^ intToInt64 v2 - --- | Apply a 'ConvOp' to an operand. Returns 'Nothing' if the --- application is mistyped. -doConvOp :: ConvOp -> PrimValue -> Maybe PrimValue -doConvOp (ZExt _ to) (IntValue v) = Just $ IntValue $ doZExt v to -doConvOp (SExt _ to) (IntValue v) = Just $ IntValue $ doSExt v to -doConvOp (FPConv _ to) (FloatValue v) = Just $ FloatValue $ doFPConv v to -doConvOp (FPToUI _ to) (FloatValue v) = Just $ IntValue $ doFPToUI v to -doConvOp (FPToSI _ to) (FloatValue v) = Just $ IntValue $ doFPToSI v to -doConvOp (UIToFP _ to) (IntValue v) = Just $ FloatValue $ doUIToFP v to -doConvOp (SIToFP _ to) (IntValue v) = Just $ FloatValue $ doSIToFP v to -doConvOp (IToB _) (IntValue v) = Just $ BoolValue $ intToInt64 v /= 0 -doConvOp (BToI to) (BoolValue v) = Just $ IntValue $ intValue to $ if v then 1 else 0::Int -doConvOp _ _ = Nothing - --- | Zero-extend the given integer value to the size of the given --- type. If the type is smaller than the given value, the result is a --- truncation. -doZExt :: IntValue -> IntType -> IntValue -doZExt (Int8Value x) t = intValue t $ toInteger (fromIntegral x :: Word8) -doZExt (Int16Value x) t = intValue t $ toInteger (fromIntegral x :: Word16) -doZExt (Int32Value x) t = intValue t $ toInteger (fromIntegral x :: Word32) -doZExt (Int64Value x) t = intValue t $ toInteger (fromIntegral x :: Word64) - --- | Sign-extend the given integer value to the size of the given --- type. If the type is smaller than the given value, the result is a --- truncation. -doSExt :: IntValue -> IntType -> IntValue -doSExt (Int8Value x) t = intValue t $ toInteger x -doSExt (Int16Value x) t = intValue t $ toInteger x -doSExt (Int32Value x) t = intValue t $ toInteger x -doSExt (Int64Value x) t = intValue t $ toInteger x - --- | 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 - --- | Convert a floating-point value to the nearest --- unsigned integer (rounding towards zero). -doFPToUI :: FloatValue -> IntType -> IntValue -doFPToUI v t = intValue t (truncate $ floatToDouble v :: Word64) - --- | Convert a floating-point value to the nearest --- signed integer (rounding towards zero). -doFPToSI :: FloatValue -> IntType -> IntValue -doFPToSI v t = intValue t (truncate $ floatToDouble v :: Word64) - --- | Convert an unsigned integer to a floating-point value. -doUIToFP :: IntValue -> FloatType -> FloatValue -doUIToFP v t = floatValue t $ intToWord64 v - --- | Convert a signed integer to a floating-point value. -doSIToFP :: IntValue -> FloatType -> FloatValue -doSIToFP v t = floatValue t $ intToInt64 v - --- | Apply a 'CmpOp' to an operand. Returns 'Nothing' if the --- application is mistyped. -doCmpOp :: CmpOp -> PrimValue -> PrimValue -> Maybe Bool -doCmpOp CmpEq{} v1 v2 = Just $ doCmpEq v1 v2 -doCmpOp CmpUlt{} (IntValue v1) (IntValue v2) = Just $ doCmpUlt v1 v2 -doCmpOp CmpUle{} (IntValue v1) (IntValue v2) = Just $ doCmpUle v1 v2 -doCmpOp CmpSlt{} (IntValue v1) (IntValue v2) = Just $ doCmpSlt v1 v2 -doCmpOp CmpSle{} (IntValue v1) (IntValue v2) = Just $ doCmpSle v1 v2 -doCmpOp FCmpLt{} (FloatValue v1) (FloatValue v2) = Just $ doFCmpLt v1 v2 -doCmpOp FCmpLe{} (FloatValue v1) (FloatValue v2) = Just $ doFCmpLe v1 v2 -doCmpOp CmpLlt{} (BoolValue v1) (BoolValue v2) = Just $ not v1 && v2 -doCmpOp CmpLle{} (BoolValue v1) (BoolValue v2) = Just $ not (v1 && not v2) -doCmpOp _ _ _ = Nothing - --- | Compare any two primtive values for exact equality. -doCmpEq :: PrimValue -> PrimValue -> Bool -doCmpEq (FloatValue (Float32Value v1)) (FloatValue (Float32Value v2)) = v1 == v2 -doCmpEq (FloatValue (Float64Value v1)) (FloatValue (Float64Value v2)) = v1 == v2 -doCmpEq v1 v2 = v1 == v2 - --- | Unsigned less than. -doCmpUlt :: IntValue -> IntValue -> Bool -doCmpUlt v1 v2 = intToWord64 v1 < intToWord64 v2 - --- | Unsigned less than or equal. -doCmpUle :: IntValue -> IntValue -> Bool -doCmpUle v1 v2 = intToWord64 v1 <= intToWord64 v2 - --- | Signed less than. -doCmpSlt :: IntValue -> IntValue -> Bool -doCmpSlt = (<) - --- | Signed less than or equal. -doCmpSle :: IntValue -> IntValue -> Bool -doCmpSle = (<=) - --- | Floating-point less than. -doFCmpLt :: FloatValue -> FloatValue -> Bool -doFCmpLt = (<) - --- | Floating-point less than or equal. -doFCmpLe :: FloatValue -> FloatValue -> Bool -doFCmpLe = (<=) - --- | Translate an t'IntValue' to 'Word64'. This is guaranteed to fit. -intToWord64 :: IntValue -> Word64 -intToWord64 (Int8Value v) = fromIntegral (fromIntegral v :: Word8) -intToWord64 (Int16Value v) = fromIntegral (fromIntegral v :: Word16) -intToWord64 (Int32Value v) = fromIntegral (fromIntegral v :: Word32) -intToWord64 (Int64Value v) = fromIntegral (fromIntegral v :: Word64) - --- | Translate an t'IntValue' to t'Int64'. This is guaranteed to fit. -intToInt64 :: IntValue -> Int64 -intToInt64 (Int8Value v) = fromIntegral v -intToInt64 (Int16Value v) = fromIntegral v -intToInt64 (Int32Value v) = fromIntegral v -intToInt64 (Int64Value v) = fromIntegral v - --- | Careful - there is no guarantee this will fit. -intToInt :: IntValue -> Int -intToInt = fromIntegral . intToInt64 - -floatToDouble :: FloatValue -> Double -floatToDouble (Float32Value v) = fromRational $ toRational v -floatToDouble (Float64Value v) = v - --- | The result type of a binary operator. -binOpType :: BinOp -> PrimType -binOpType (Add t _) = IntType t -binOpType (Sub t _) = IntType t -binOpType (Mul t _) = IntType t -binOpType (SDiv t _) = IntType t -binOpType (SDivUp t _) = IntType t -binOpType (SMod t _) = IntType t -binOpType (SQuot t _) = IntType t -binOpType (SRem t _) = IntType t -binOpType (UDiv t _) = IntType t -binOpType (UDivUp t _) = IntType t -binOpType (UMod t _) = IntType t -binOpType (SMin t) = IntType t -binOpType (UMin t) = IntType t -binOpType (FMin t) = FloatType t -binOpType (SMax t) = IntType t -binOpType (UMax t) = IntType t -binOpType (FMax t) = FloatType t -binOpType (Shl t) = IntType t -binOpType (LShr t) = IntType t -binOpType (AShr t) = IntType t -binOpType (And t) = IntType t -binOpType (Or t) = IntType t -binOpType (Xor t) = IntType t -binOpType (Pow t) = IntType t -binOpType (FPow t) = FloatType t -binOpType LogAnd = Bool -binOpType LogOr = Bool -binOpType (FAdd t) = FloatType t -binOpType (FSub t) = FloatType t -binOpType (FMul t) = FloatType t -binOpType (FDiv t) = FloatType t -binOpType (FMod t) = FloatType t - --- | The operand types of a comparison operator. -cmpOpType :: CmpOp -> PrimType -cmpOpType (CmpEq t) = t -cmpOpType (CmpSlt t) = IntType t -cmpOpType (CmpSle t) = IntType t -cmpOpType (CmpUlt t) = IntType t -cmpOpType (CmpUle t) = IntType t -cmpOpType (FCmpLt t) = FloatType t -cmpOpType (FCmpLe t) = FloatType t -cmpOpType CmpLlt = Bool -cmpOpType CmpLle = Bool - --- | The operand and result type of a unary operator. -unOpType :: UnOp -> PrimType -unOpType (SSignum t) = IntType t -unOpType (USignum t) = IntType t -unOpType Not = Bool -unOpType (Complement t) = IntType t -unOpType (Abs t) = IntType t -unOpType (FAbs t) = FloatType t - --- | The input and output types of a conversion operator. -convOpType :: ConvOp -> (PrimType, PrimType) -convOpType (ZExt from to) = (IntType from, IntType to) -convOpType (SExt from to) = (IntType from, IntType to) -convOpType (FPConv from to) = (FloatType from, FloatType to) -convOpType (FPToUI from to) = (FloatType from, IntType to) -convOpType (FPToSI from to) = (FloatType from, IntType to) -convOpType (UIToFP from to) = (IntType from, FloatType to) -convOpType (SIToFP from to) = (IntType from, FloatType to) -convOpType (IToB from) = (IntType from, Bool) -convOpType (BToI to) = (Bool, IntType to) - -floatToWord :: Float -> Word32 -floatToWord = G.runGet G.getWord32le . P.runPut . P.putFloatle - -wordToFloat :: Word32 -> Float -wordToFloat = G.runGet G.getFloatle . P.runPut . P.putWord32le - -doubleToWord :: Double -> Word64 -doubleToWord = G.runGet G.getWord64le . P.runPut . P.putDoublele - -wordToDouble :: Word64 -> Double -wordToDouble = G.runGet G.getDoublele . P.runPut . P.putWord64le - --- | A mapping from names of primitive functions to their parameter --- types, their result type, and a function for evaluating them. -primFuns :: M.Map String ([PrimType], PrimType, - [PrimValue] -> Maybe PrimValue) -primFuns = M.fromList - [ f32 "sqrt32" sqrt, f64 "sqrt64" sqrt - , f32 "log32" log, f64 "log64" log - , f32 "log10_32" (logBase 10), f64 "log10_64" (logBase 10) - , f32 "log2_32" (logBase 2), f64 "log2_64" (logBase 2) - , f32 "exp32" exp, f64 "exp64" exp - - , f32 "sin32" sin, f64 "sin64" sin - , f32 "sinh32" sinh, f64 "sinh64" sinh - , f32 "cos32" cos, f64 "cos64" cos - , f32 "cosh32" cosh, f64 "cosh64" cosh - , f32 "tan32" tan, f64 "tan64" tan - , f32 "tanh32" tanh, f64 "tanh64" tanh - , f32 "asin32" asin, f64 "asin64" asin - , f32 "asinh32" asinh, f64 "asinh64" asinh - , f32 "acos32" acos, f64 "acos64" acos - , f32 "acosh32" acosh, f64 "acosh64" acosh - , f32 "atan32" atan, f64 "atan64" atan - , f32 "atanh32" atanh, f64 "atanh64" atanh - - , f32 "round32" roundFloat, f64 "round64" roundDouble - , f32 "ceil32" ceilFloat, f64 "ceil64" ceilDouble - , f32 "floor32" floorFloat, f64 "floor64" floorDouble - , f32 "gamma32" tgammaf, f64 "gamma64" tgamma - , f32 "lgamma32" lgammaf, f64 "lgamma64" lgamma - - , i8 "clz8" $ IntValue . Int32Value . fromIntegral . countLeadingZeros - , i16 "clz16" $ IntValue . Int32Value . fromIntegral . countLeadingZeros - , i32 "clz32" $ IntValue . Int32Value . fromIntegral . countLeadingZeros - , i64 "clz64" $ IntValue . Int32Value . fromIntegral . countLeadingZeros - - , i8 "ctz8" $ IntValue . Int32Value . fromIntegral . countTrailingZeros - , i16 "ctz16" $ IntValue . Int32Value . fromIntegral . countTrailingZeros - , i32 "ctz32" $ IntValue . Int32Value . fromIntegral . countTrailingZeros - , i64 "ctz64" $ IntValue . Int32Value . fromIntegral . countTrailingZeros - - , i8 "popc8" $ IntValue . Int32Value . fromIntegral . popCount - , i16 "popc16" $ IntValue . Int32Value . fromIntegral . popCount - , i32 "popc32" $ IntValue . Int32Value . fromIntegral . popCount - , i64 "popc64" $ IntValue . Int32Value . fromIntegral . popCount - - , ("mad_hi8", ([IntType Int8, IntType Int8, IntType Int8], IntType Int8, - \case - [IntValue (Int8Value a), IntValue (Int8Value b), IntValue (Int8Value c)] -> - Just $ IntValue . Int8Value $ mad_hi8 (Int8Value a) (Int8Value b) c - _ -> Nothing - )) - , ("mad_hi16", ([IntType Int16, IntType Int16, IntType Int16], IntType Int16, - \case - [IntValue (Int16Value a), IntValue (Int16Value b), IntValue (Int16Value c)] -> - Just $ IntValue . Int16Value $ mad_hi16 (Int16Value a) (Int16Value b) c - _ -> Nothing - )) - , ("mad_hi32", ([IntType Int32, IntType Int32, IntType Int32], IntType Int32, - \case - [IntValue (Int32Value a), IntValue (Int32Value b), IntValue (Int32Value c)] -> - Just $ IntValue . Int32Value $ mad_hi32 (Int32Value a) (Int32Value b) c - _ -> Nothing - )) - , ("mad_hi64", ([IntType Int64, IntType Int64, IntType Int64], IntType Int64, - \case - [IntValue (Int64Value a), IntValue (Int64Value b), IntValue (Int64Value c)] -> - Just $ IntValue . Int64Value $ mad_hi64 (Int64Value a) (Int64Value b) c - _ -> Nothing - )) - - , ("mul_hi8", ([IntType Int8, IntType Int8], IntType Int8, - \case - [IntValue (Int8Value a), IntValue (Int8Value b)] -> - Just $ IntValue . Int8Value $ mul_hi8 (Int8Value a) (Int8Value b) - _ -> Nothing - )) - , ("mul_hi16", ([IntType Int16, IntType Int16], IntType Int16, - \case - [IntValue (Int16Value a), IntValue (Int16Value b)] -> - Just $ IntValue . Int16Value $ mul_hi16 (Int16Value a) (Int16Value b) - _ -> Nothing - )) - , ("mul_hi32", ([IntType Int32, IntType Int32], IntType Int32, - \case - [IntValue (Int32Value a), IntValue (Int32Value b)] -> - Just $ IntValue . Int32Value $ mul_hi32 (Int32Value a) (Int32Value b) - _ -> Nothing - )) - , ("mul_hi64", ([IntType Int64, IntType Int64], IntType Int64, - \case - [IntValue (Int64Value a), IntValue (Int64Value b)] -> - Just $ IntValue . Int64Value $ mul_hi64 (Int64Value a) (Int64Value b) - _ -> Nothing - )) - - , ("atan2_32", - ([FloatType Float32, FloatType Float32], FloatType Float32, - \case - [FloatValue (Float32Value x), FloatValue (Float32Value y)] -> - Just $ FloatValue $ Float32Value $ atan2 x y - _ -> Nothing)) - , ("atan2_64", - ([FloatType Float64, FloatType Float64], FloatType Float64, - \case - [FloatValue (Float64Value x), FloatValue (Float64Value y)] -> - Just $ FloatValue $ Float64Value $ atan2 x y - _ -> Nothing)) - - , ("isinf32", - ([FloatType Float32], Bool, - \case - [FloatValue (Float32Value x)] -> Just $ BoolValue $ isInfinite x - _ -> Nothing)) - , ("isinf64", - ([FloatType Float64], Bool, - \case - [FloatValue (Float64Value x)] -> Just $ BoolValue $ isInfinite x - _ -> Nothing)) - - , ("isnan32", - ([FloatType Float32], Bool, - \case - [FloatValue (Float32Value x)] -> Just $ BoolValue $ isNaN x - _ -> Nothing)) - , ("isnan64", - ([FloatType Float64], Bool, - \case - [FloatValue (Float64Value x)] -> Just $ BoolValue $ isNaN x - _ -> Nothing)) - - , ("to_bits32", - ([FloatType Float32], IntType Int32, - \case - [FloatValue (Float32Value x)] -> - Just $ IntValue $ Int32Value $ fromIntegral $ floatToWord x - _ -> Nothing)) - , ("to_bits64", - ([FloatType Float64], IntType Int64, - \case - [FloatValue (Float64Value x)] -> - Just $ IntValue $ Int64Value $ fromIntegral $ doubleToWord x - _ -> Nothing)) - - , ("from_bits32", - ([IntType Int32], FloatType Float32, - \case - [IntValue (Int32Value x)] -> - Just $ FloatValue $ Float32Value $ wordToFloat $ fromIntegral x - _ -> Nothing)) - , ("from_bits64", - ([IntType Int64], FloatType Float64, - \case - [IntValue (Int64Value x)] -> - Just $ FloatValue $ Float64Value $ wordToDouble $ fromIntegral x - _ -> Nothing)) - - , f32_3 "lerp32" (\v0 v1 t -> v0 + (v1-v0)*max 0 (min 1 t)) - , f64_3 "lerp64" (\v0 v1 t -> v0 + (v1-v0)*max 0 (min 1 t)) - - , f32_3 "mad32" (\a b c -> a*b+c) - , f64_3 "mad64" (\a b c -> a*b+c) - - , f32_3 "fma32" (\a b c -> a*b+c) - , f64_3 "fma64" (\a b c -> a*b+c) - - ] - where i8 s f = (s, ([IntType Int8], IntType Int32, i8PrimFun f)) - i16 s f = (s, ([IntType Int16], IntType Int32, i16PrimFun f)) - i32 s f = (s, ([IntType Int32], IntType Int32, i32PrimFun f)) - i64 s f = (s, ([IntType Int64], IntType Int32, i64PrimFun f)) - f32 s f = (s, ([FloatType Float32], FloatType Float32, f32PrimFun f)) - f64 s f = (s, ([FloatType Float64], FloatType Float64, f64PrimFun f)) - f32_3 s f = (s, ([FloatType Float32,FloatType Float32,FloatType Float32], - FloatType Float32, f32PrimFun3 f)) - f64_3 s f = (s, ([FloatType Float64,FloatType Float64,FloatType Float64], - FloatType Float64, f64PrimFun3 f)) - - i8PrimFun f [IntValue (Int8Value x)] = - Just $ f x - i8PrimFun _ _ = Nothing - - i16PrimFun f [IntValue (Int16Value x)] = - Just $ f x - i16PrimFun _ _ = Nothing - - i32PrimFun f [IntValue (Int32Value x)] = - Just $ f x - i32PrimFun _ _ = Nothing - - i64PrimFun f [IntValue (Int64Value x)] = - Just $ f x - i64PrimFun _ _ = Nothing - - f32PrimFun f [FloatValue (Float32Value x)] = - Just $ FloatValue $ Float32Value $ f x - f32PrimFun _ _ = Nothing - - f64PrimFun f [FloatValue (Float64Value x)] = - Just $ FloatValue $ Float64Value $ f x - f64PrimFun _ _ = Nothing - - f32PrimFun3 f [FloatValue (Float32Value a), - FloatValue (Float32Value b), - FloatValue (Float32Value c)] = - Just $ FloatValue $ Float32Value $ f a b c - f32PrimFun3 _ _ = Nothing - - f64PrimFun3 f [FloatValue (Float64Value a), - FloatValue (Float64Value b), - FloatValue (Float64Value c)] = - Just $ FloatValue $ Float64Value $ f a b c - f64PrimFun3 _ _ = Nothing - --- | Is the given value kind of zero? -zeroIsh :: PrimValue -> Bool -zeroIsh (IntValue k) = zeroIshInt k -zeroIsh (FloatValue (Float32Value k)) = k == 0 -zeroIsh (FloatValue (Float64Value k)) = k == 0 -zeroIsh (BoolValue False) = True -zeroIsh _ = False - --- | Is the given value kind of one? -oneIsh :: PrimValue -> Bool -oneIsh (IntValue k) = oneIshInt k -oneIsh (FloatValue (Float32Value k)) = k == 1 -oneIsh (FloatValue (Float64Value k)) = k == 1 -oneIsh (BoolValue True) = True -oneIsh _ = False - --- | Is the given value kind of negative? -negativeIsh :: PrimValue -> Bool -negativeIsh (IntValue k) = negativeIshInt k -negativeIsh (FloatValue (Float32Value k)) = k < 0 -negativeIsh (FloatValue (Float64Value k)) = k < 0 -negativeIsh (BoolValue _) = False -negativeIsh Checked = False - --- | Is the given integer value kind of zero? -zeroIshInt :: IntValue -> Bool -zeroIshInt (Int8Value k) = k == 0 -zeroIshInt (Int16Value k) = k == 0 -zeroIshInt (Int32Value k) = k == 0 -zeroIshInt (Int64Value k) = k == 0 - --- | Is the given integer value kind of one? -oneIshInt :: IntValue -> Bool -oneIshInt (Int8Value k) = k == 1 -oneIshInt (Int16Value k) = k == 1 -oneIshInt (Int32Value k) = k == 1 -oneIshInt (Int64Value k) = k == 1 - --- | Is the given integer value kind of negative? -negativeIshInt :: IntValue -> Bool -negativeIshInt (Int8Value k) = k < 0 -negativeIshInt (Int16Value k) = k < 0 -negativeIshInt (Int32Value k) = k < 0 -negativeIshInt (Int64Value k) = k < 0 - --- | The size of a value of a given primitive type in bites. -primBitSize :: PrimType -> Int -primBitSize = (*8) . primByteSize - --- | The size of a value of a given primitive type in eight-bit bytes. -primByteSize :: Num a => PrimType -> a -primByteSize (IntType t) = intByteSize t -primByteSize (FloatType t) = floatByteSize t -primByteSize Bool = 1 -primByteSize Cert = 1 - --- | The size of a value of a given integer type in eight-bit bytes. -intByteSize :: Num a => IntType -> a -intByteSize Int8 = 1 -intByteSize Int16 = 2 -intByteSize Int32 = 4 -intByteSize Int64 = 8 - --- | The size of a value of a given floating-point type in eight-bit bytes. -floatByteSize :: Num a => FloatType -> a -floatByteSize Float32 = 4 -floatByteSize Float64 = 8 - --- | True if the given binary operator is commutative. -commutativeBinOp :: BinOp -> Bool -commutativeBinOp Add{} = True -commutativeBinOp FAdd{} = True -commutativeBinOp Mul{} = True -commutativeBinOp FMul{} = True -commutativeBinOp And{} = True -commutativeBinOp Or{} = True -commutativeBinOp Xor{} = True -commutativeBinOp LogOr{} = True -commutativeBinOp LogAnd{} = True -commutativeBinOp SMax{} = True -commutativeBinOp SMin{} = True -commutativeBinOp UMax{} = True -commutativeBinOp UMin{} = True -commutativeBinOp FMax{} = True -commutativeBinOp FMin{} = True -commutativeBinOp _ = False - --- Prettyprinting instances - -instance Pretty BinOp where - ppr (Add t OverflowWrap) = taggedI "add" t - ppr (Add t OverflowUndef) = taggedI "add_nw" t - ppr (Sub t OverflowWrap) = taggedI "sub" t - ppr (Sub t OverflowUndef) = taggedI "sub_nw" t - ppr (Mul t OverflowWrap) = taggedI "mul" t - ppr (Mul t OverflowUndef) = taggedI "mul_nw" t - ppr (FAdd t) = taggedF "fadd" t - ppr (FSub t) = taggedF "fsub" t - ppr (FMul t) = taggedF "fmul" t - ppr (UDiv t Safe) = taggedI "udiv_safe" t - ppr (UDiv t Unsafe) = taggedI "udiv" t - ppr (UDivUp t Safe) = taggedI "udiv_up_safe" t - ppr (UDivUp t Unsafe) = taggedI "udiv_up" t - ppr (UMod t Safe) = taggedI "umod_safe" t - ppr (UMod t Unsafe) = taggedI "umod" t - ppr (SDiv t Safe) = taggedI "sdiv_safe" t - ppr (SDiv t Unsafe) = taggedI "sdiv" t - ppr (SDivUp t Safe) = taggedI "sdiv_up_safe" t - ppr (SDivUp t Unsafe) = taggedI "sdiv_up" t - ppr (SMod t Safe) = taggedI "smod_safe" t - ppr (SMod t Unsafe) = taggedI "smod" t - ppr (SQuot t Safe) = taggedI "squot_safe" t - ppr (SQuot t Unsafe) = taggedI "squot" t - ppr (SRem t Safe) = taggedI "srem_safe" t - ppr (SRem t Unsafe) = taggedI "srem" t - ppr (FDiv t) = taggedF "fdiv" t - ppr (FMod t) = taggedF "fmod" t - ppr (SMin t) = taggedI "smin" t - ppr (UMin t) = taggedI "umin" t - ppr (FMin t) = taggedF "fmin" t - ppr (SMax t) = taggedI "smax" t - ppr (UMax t) = taggedI "umax" t - ppr (FMax t) = taggedF "fmax" t - ppr (Shl t) = taggedI "shl" t - ppr (LShr t) = taggedI "lshr" t - ppr (AShr t) = taggedI "ashr" t - ppr (And t) = taggedI "and" t - ppr (Or t) = taggedI "or" t - ppr (Xor t) = taggedI "xor" t - ppr (Pow t) = taggedI "pow" t - ppr (FPow t) = taggedF "fpow" t - ppr LogAnd = text "logand" - ppr LogOr = text "logor" - -instance Pretty CmpOp where - ppr (CmpEq t) = text "eq_" <> ppr t - ppr (CmpUlt t) = taggedI "ult" t - ppr (CmpUle t) = taggedI "ule" t - ppr (CmpSlt t) = taggedI "slt" t - ppr (CmpSle t) = taggedI "sle" t - ppr (FCmpLt t) = taggedF "lt" t - ppr (FCmpLe t) = taggedF "le" t - ppr CmpLlt = text "llt" - ppr CmpLle = text "lle" - -instance Pretty ConvOp where - ppr op = convOp (convOpFun op) from to - where (from, to) = convOpType op - -instance Pretty UnOp where - ppr Not = text "not" - ppr (Abs t) = taggedI "abs" t - ppr (FAbs t) = taggedF "fabs" t - ppr (SSignum t) = taggedI "ssignum" t - ppr (USignum t) = taggedI "usignum" t - ppr (Complement t) = taggedI "complement" t - --- | The human-readable name for a 'ConvOp'. This is used to expose --- the 'ConvOp' in the @intrinsics@ module of a Futhark program. -convOpFun :: ConvOp -> String -convOpFun ZExt{} = "zext" -convOpFun SExt{} = "sext" -convOpFun FPConv{} = "fpconv" -convOpFun FPToUI{} = "fptoui" -convOpFun FPToSI{} = "fptosi" -convOpFun UIToFP{} = "uitofp" -convOpFun SIToFP{} = "sitofp" -convOpFun IToB{} = "itob" -convOpFun BToI{} = "btoi" - -taggedI :: String -> IntType -> Doc -taggedI s Int8 = text $ s ++ "8" -taggedI s Int16 = text $ s ++ "16" -taggedI s Int32 = text $ s ++ "32" -taggedI s Int64 = text $ s ++ "64" - -taggedF :: String -> FloatType -> Doc -taggedF s Float32 = text $ s ++ "32" -taggedF s Float64 = text $ s ++ "64" - -convOp :: (Pretty from, Pretty to) => String -> from -> to -> Doc -convOp s from to = text s <> text "_" <> ppr from <> text "_" <> ppr to - --- | True if signed. Only makes a difference for integer types. -prettySigned :: Bool -> PrimType -> String -prettySigned True (IntType it) = 'u' : drop 1 (pretty it) -prettySigned _ t = pretty t - -mul_hi8 :: IntValue -> IntValue -> Int8 -mul_hi8 a b = - let a' = intToWord64 a - b' = intToWord64 b - in fromIntegral (shiftR (a' * b') 8) - -mul_hi16 :: IntValue -> IntValue -> Int16 -mul_hi16 a b = - let a' = intToWord64 a - b' = intToWord64 b - in fromIntegral (shiftR (a' * b') 16) - -mul_hi32 :: IntValue -> IntValue -> Int32 -mul_hi32 a b = - let a' = intToWord64 a - b' = intToWord64 b - in fromIntegral (shiftR (a' * b') 32) - -mul_hi64 :: IntValue -> IntValue -> Int64 -mul_hi64 a b = - let a' = (toInteger . intToWord64) a - b' = (toInteger . intToWord64) b - in fromIntegral (shiftR (a' * b') 64) +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE Safe #-} +{-# LANGUAGE TypeOperators #-} + +-- | Definitions of primitive types, the values that inhabit these +-- types, and operations on these values. A primitive value can also +-- be called a scalar. +-- +-- Essentially, this module describes the subset of the (internal) +-- Futhark language that operates on primitive types. +module Futhark.IR.Primitive + ( -- * Types + IntType (..), + allIntTypes, + FloatType (..), + allFloatTypes, + PrimType (..), + allPrimTypes, + module Data.Int, + + -- * Values + IntValue (..), + intValue, + intValueType, + valueIntegral, + FloatValue (..), + floatValue, + floatValueType, + PrimValue (..), + primValueType, + blankPrimValue, + + -- * Operations + Overflow (..), + Safety (..), + UnOp (..), + allUnOps, + BinOp (..), + allBinOps, + ConvOp (..), + allConvOps, + CmpOp (..), + allCmpOps, + + -- ** Unary Operations + doUnOp, + doComplement, + doAbs, + doFAbs, + doSSignum, + doUSignum, + + -- ** Binary Operations + doBinOp, + doAdd, + doMul, + doSDiv, + doSMod, + doPow, + + -- ** Conversion Operations + doConvOp, + doZExt, + doSExt, + doFPConv, + doFPToUI, + doFPToSI, + doUIToFP, + doSIToFP, + intToInt64, + intToWord64, + + -- * Comparison Operations + doCmpOp, + doCmpEq, + doCmpUlt, + doCmpUle, + doCmpSlt, + doCmpSle, + doFCmpLt, + doFCmpLe, + + -- * Type Of + binOpType, + unOpType, + cmpOpType, + convOpType, + + -- * Primitive functions + primFuns, + + -- * Utility + zeroIsh, + zeroIshInt, + oneIsh, + oneIshInt, + negativeIsh, + primBitSize, + primByteSize, + intByteSize, + floatByteSize, + commutativeBinOp, + + -- * Prettyprinting + convOpFun, + prettySigned, + ) +where + +import Control.Category +import qualified Data.Binary.Get as G +import qualified Data.Binary.Put as P +import Data.Bits +import Data.Fixed (mod') -- Weird location. +import Data.Int (Int16, Int32, Int64, Int8) +import qualified Data.Map as M +import qualified Data.Text as T +import Data.Word +import Futhark.Util + ( ceilDouble, + ceilFloat, + floorDouble, + floorFloat, + lgamma, + lgammaf, + roundDouble, + roundFloat, + tgamma, + tgammaf, + ) +import Futhark.Util.Pretty +import GHC.Generics (Generic) +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Text.Read (readMaybe) +import Prelude hiding (id, (.)) + +-- | An integer type, ordered by size. Note that signedness is not a +-- property of the type, but a property of the operations performed on +-- values of these types. +data IntType + = Int8 + | Int16 + | Int32 + | Int64 + deriving (Eq, Ord, Show, Enum, Bounded, Generic) + +instance SexpIso IntType where + sexpIso = + match $ + With (sym "i8" >>>) $ + With (sym "i16" >>>) $ + With (sym "i32" >>>) $ + With + (sym "i64" >>>) + End + +instance Pretty IntType where + ppr Int8 = text "i8" + ppr Int16 = text "i16" + ppr Int32 = text "i32" + ppr Int64 = text "i64" + +-- | A list of all integer types. +allIntTypes :: [IntType] +allIntTypes = [minBound .. maxBound] + +-- | A floating point type. +data FloatType + = Float32 + | Float64 + deriving (Eq, Ord, Show, Enum, Bounded, Generic) + +instance SexpIso FloatType where + sexpIso = + match $ + With (sym "f32" >>>) $ + With + (sym "f64" >>>) + End + +instance Pretty FloatType where + ppr Float32 = text "f32" + ppr Float64 = text "f64" + +-- | A list of all floating-point types. +allFloatTypes :: [FloatType] +allFloatTypes = [minBound .. maxBound] + +-- | Low-level primitive types. +data PrimType + = IntType IntType + | FloatType FloatType + | Bool + | Cert + deriving (Eq, Ord, Show, Generic) + +instance SexpIso PrimType where + sexpIso = + match $ + With (. sexpIso) $ + With (. sexpIso) $ + With (sym "bool" >>>) $ + With + (sym "cert" >>>) + End + +instance Enum PrimType where + toEnum 0 = IntType Int8 + toEnum 1 = IntType Int16 + toEnum 2 = IntType Int32 + toEnum 3 = IntType Int64 + toEnum 4 = FloatType Float32 + toEnum 5 = FloatType Float64 + toEnum 6 = Bool + toEnum _ = Cert + + fromEnum (IntType Int8) = 0 + fromEnum (IntType Int16) = 1 + fromEnum (IntType Int32) = 2 + fromEnum (IntType Int64) = 3 + fromEnum (FloatType Float32) = 4 + fromEnum (FloatType Float64) = 5 + fromEnum Bool = 6 + fromEnum Cert = 7 + +instance Bounded PrimType where + minBound = IntType Int8 + maxBound = Cert + +instance Pretty PrimType where + ppr (IntType t) = ppr t + ppr (FloatType t) = ppr t + ppr Bool = text "bool" + ppr Cert = text "cert" + +-- | A list of all primitive types. +allPrimTypes :: [PrimType] +allPrimTypes = + map IntType allIntTypes + ++ map FloatType allFloatTypes + ++ [Bool, Cert] + +numberIso :: (Read a, Pretty a) => T.Text -> Grammar p (T.Text :- t) (a :- t) +numberIso postfix = partialOsi (fromS postfix) (toS postfix) + where + toS t s = prettyText s <> t + + fromS t s + | t `T.isSuffixOf` s, + Just v <- readMaybe $ T.unpack $ T.dropEnd (T.length t) s = + Right v + | otherwise = + Left $ expected $ "Couldn't parse " <> t + +-- | An integer value. +data IntValue + = Int8Value !Int8 + | Int16Value !Int16 + | Int32Value !Int32 + | Int64Value !Int64 + deriving (Eq, Ord, Show, Generic) + +instance SexpIso IntValue where + sexpIso = + match $ + With (. numberIso "i8" . Sexp.symbol) $ + With (. numberIso "i16" . Sexp.symbol) $ + With (. numberIso "i32" . Sexp.symbol) $ + With + (. numberIso "i64" . Sexp.symbol) + End + +instance Pretty IntValue where + ppr (Int8Value v) = text $ show v ++ "i8" + ppr (Int16Value v) = text $ show v ++ "i16" + ppr (Int32Value v) = text $ show v ++ "i32" + ppr (Int64Value v) = text $ show v ++ "i64" + +-- | Create an t'IntValue' from a type and an 'Integer'. +intValue :: Integral int => IntType -> int -> IntValue +intValue Int8 = Int8Value . fromIntegral +intValue Int16 = Int16Value . fromIntegral +intValue Int32 = Int32Value . fromIntegral +intValue Int64 = Int64Value . fromIntegral + +-- | The type of an integer value. +intValueType :: IntValue -> IntType +intValueType Int8Value {} = Int8 +intValueType Int16Value {} = Int16 +intValueType Int32Value {} = Int32 +intValueType Int64Value {} = Int64 + +-- | Convert an t'IntValue' to any 'Integral' type. +valueIntegral :: Integral int => IntValue -> int +valueIntegral (Int8Value v) = fromIntegral v +valueIntegral (Int16Value v) = fromIntegral v +valueIntegral (Int32Value v) = fromIntegral v +valueIntegral (Int64Value v) = fromIntegral v + +-- | A floating-point value. +data FloatValue + = Float32Value !Float + | Float64Value !Double + deriving (Show, Generic) + +instance Eq FloatValue where + Float32Value x == Float32Value y = isNaN x && isNaN y || x == y + Float64Value x == Float64Value y = isNaN x && isNaN y || x == y + Float32Value _ == Float64Value _ = False + Float64Value _ == Float32Value _ = False + +-- The derived Ord instance does not handle NaNs correctly. +instance Ord FloatValue where + Float32Value x <= Float32Value y = x <= y + Float64Value x <= Float64Value y = x <= y + Float32Value _ <= Float64Value _ = True + Float64Value _ <= Float32Value _ = False + + Float32Value x < Float32Value y = x < y + Float64Value x < Float64Value y = x < y + Float32Value _ < Float64Value _ = True + Float64Value _ < Float32Value _ = False + + (>) = flip (<) + (>=) = flip (<=) + +instance SexpIso FloatValue where + sexpIso = + match $ + With (. numberIso "f32" . Sexp.symbol) $ + With + (. numberIso "f64" . Sexp.symbol) + End + +instance Pretty FloatValue where + ppr (Float32Value v) + | isInfinite v, v >= 0 = text "f32.inf" + | isInfinite v, v < 0 = text "-f32.inf" + | isNaN v = text "f32.nan" + | otherwise = text $ show v ++ "f32" + ppr (Float64Value v) + | isInfinite v, v >= 0 = text "f64.inf" + | isInfinite v, v < 0 = text "-f64.inf" + | isNaN v = text "f64.nan" + | otherwise = text $ show v ++ "f64" + +-- | Create a t'FloatValue' from a type and a 'Rational'. +floatValue :: Real num => FloatType -> num -> FloatValue +floatValue Float32 = Float32Value . fromRational . toRational +floatValue Float64 = Float64Value . fromRational . toRational + +-- | The type of a floating-point value. +floatValueType :: FloatValue -> FloatType +floatValueType Float32Value {} = Float32 +floatValueType Float64Value {} = Float64 + +-- | Non-array values. +data PrimValue + = IntValue !IntValue + | FloatValue !FloatValue + | BoolValue !Bool + | -- | The only value of type @cert@. + Checked + deriving (Eq, Ord, Show, Generic) + +instance SexpIso PrimValue where + sexpIso = + match $ + With (. sexpIso) $ + With (. sexpIso) $ + With (. sexpIso) $ + With + (sym "checked" >>>) + End + +instance Pretty PrimValue where + ppr (IntValue v) = ppr v + ppr (BoolValue True) = text "true" + ppr (BoolValue False) = text "false" + ppr (FloatValue v) = ppr v + ppr Checked = text "checked" + +-- | The type of a basic value. +primValueType :: PrimValue -> PrimType +primValueType (IntValue v) = IntType $ intValueType v +primValueType (FloatValue v) = FloatType $ floatValueType v +primValueType BoolValue {} = Bool +primValueType Checked = Cert + +-- | A "blank" value of the given primitive type - this is zero, or +-- whatever is close to it. Don't depend on this value, but use it +-- for e.g. creating arrays to be populated by do-loops. +blankPrimValue :: PrimType -> PrimValue +blankPrimValue (IntType Int8) = IntValue $ Int8Value 0 +blankPrimValue (IntType Int16) = IntValue $ Int16Value 0 +blankPrimValue (IntType Int32) = IntValue $ Int32Value 0 +blankPrimValue (IntType Int64) = IntValue $ Int64Value 0 +blankPrimValue (FloatType Float32) = FloatValue $ Float32Value 0.0 +blankPrimValue (FloatType Float64) = FloatValue $ Float64Value 0.0 +blankPrimValue Bool = BoolValue False +blankPrimValue Cert = Checked + +-- | Various unary operators. It is a bit ad-hoc what is a unary +-- operator and what is a built-in function. Perhaps these should all +-- go away eventually. +data UnOp + = -- | E.g., @! True == False@. + Not + | -- | E.g., @~(~1) = 1@. + Complement IntType + | -- | @abs(-2) = 2@. + Abs IntType + | -- | @fabs(-2.0) = 2.0@. + FAbs FloatType + | -- | Signed sign function: @ssignum(-2)@ = -1. + SSignum IntType + | -- | Unsigned sign function: @usignum(2)@ = 1. + USignum IntType + deriving (Eq, Ord, Show, Generic) + +instance SexpIso UnOp where + sexpIso = + match $ + With (. Sexp.sym "not") $ + With (. Sexp.list (Sexp.el (Sexp.sym "complement") >>> Sexp.el sexpIso)) $ + 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 + +-- | What to do in case of arithmetic overflow. Futhark's semantics +-- are that overflow does wraparound, but for generated code (like +-- address arithmetic), it can be beneficial for overflow to be +-- undefined behaviour, as it allows better optimisation of things +-- such as GPU kernels. +-- +-- Note that all values of this type are considered equal for 'Eq' and +-- 'Ord'. +data Overflow = OverflowWrap | OverflowUndef + deriving (Show, Generic) + +instance SexpIso Overflow where + sexpIso = + match $ + With (. Sexp.sym "wrap") $ + With + (. Sexp.sym "undef") + End + +instance Eq Overflow where + _ == _ = True + +instance Ord Overflow where + _ `compare` _ = EQ + +-- | Whether something is safe or unsafe (mostly function calls, and +-- in the context of whether operations are dynamically checked). +-- When we inline an 'Unsafe' function, we remove all safety checks in +-- its body. The 'Ord' instance picks 'Unsafe' as being less than +-- 'Safe'. +-- +-- For operations like integer division, a safe division will not +-- explode the computer in case of division by zero, but instead +-- return some unspecified value. This always involves a run-time +-- check, so generally the unsafe variant is what the compiler will +-- insert, but guarded by an explicit assertion elsewhere. Safe +-- operations are useful when the optimiser wants to move e.g. a +-- division to a location where the divisor may be zero, but where the +-- result will only be used when it is non-zero (so it doesn't matter +-- what result is provided with a zero divisor, as long as the program +-- keeps running). +data Safety = Unsafe | Safe deriving (Eq, Ord, Show, Generic) + +instance SexpIso Safety where + sexpIso = + match $ + With (. Sexp.sym "unsafe") $ + With + (. Sexp.sym "safe") + End + +-- | Binary operators. These correspond closely to the binary operators in +-- LLVM. Most are parametrised by their expected input and output +-- types. +data BinOp + = -- | Integer addition. + Add IntType Overflow + | -- | Floating-point addition. + FAdd FloatType + | -- | Integer subtraction. + Sub IntType Overflow + | -- | Floating-point subtraction. + FSub FloatType + | -- | Integer multiplication. + Mul IntType Overflow + | -- | Floating-point multiplication. + FMul FloatType + | -- | Unsigned integer division. Rounds towards + -- negativity infinity. Note: this is different + -- from LLVM. + UDiv IntType Safety + | -- | Unsigned integer division. Rounds towards positive + -- infinity. + UDivUp IntType Safety + | -- | Signed integer division. Rounds towards + -- negativity infinity. Note: this is different + -- from LLVM. + SDiv IntType Safety + | -- | Signed integer division. Rounds towards positive + -- infinity. + SDivUp IntType Safety + | -- | Floating-point division. + FDiv FloatType + | -- | Floating-point modulus. + FMod FloatType + | -- | Unsigned integer modulus; the countepart to 'UDiv'. + UMod IntType Safety + | -- | Signed integer modulus; the countepart to 'SDiv'. + SMod IntType Safety + | -- | Signed integer division. Rounds towards zero. This + -- corresponds to the @sdiv@ instruction in LLVM and + -- integer division in C. + SQuot IntType Safety + | -- | Signed integer division. Rounds towards zero. This + -- corresponds to the @srem@ instruction in LLVM and + -- integer modulo in C. + SRem IntType Safety + | -- | Returns the smallest of two signed integers. + SMin IntType + | -- | Returns the smallest of two unsigned integers. + UMin IntType + | -- | Returns the smallest of two floating-point numbers. + FMin FloatType + | -- | Returns the greatest of two signed integers. + SMax IntType + | -- | Returns the greatest of two unsigned integers. + UMax IntType + | -- | Returns the greatest of two floating-point numbers. + FMax FloatType + | -- | Left-shift. + Shl IntType + | -- | Logical right-shift, zero-extended. + LShr IntType + | -- | Arithmetic right-shift, sign-extended. + AShr IntType + | -- | Bitwise and. + And IntType + | -- | Bitwise or. + Or IntType + | -- | Bitwise exclusive-or. + Xor IntType + | -- | Integer exponentiation. + Pow IntType + | -- | Floating-point exponentiation. + FPow FloatType + | -- | Boolean and - not short-circuiting. + LogAnd + | -- | Boolean or - not short-circuiting. + LogOr + deriving (Eq, Ord, Show, Generic) + +-- | Comparison operators are like 'BinOp's, but they always return a +-- boolean value. The somewhat ugly constructor names are straight +-- out of LLVM. +data CmpOp + = -- | All types equality. + CmpEq PrimType + | -- | Unsigned less than. + CmpUlt IntType + | -- | Unsigned less than or equal. + CmpUle IntType + | -- | Signed less than. + CmpSlt IntType + | -- | Signed less than or equal. + CmpSle IntType + | -- Comparison operators for floating-point values. TODO: extend + -- this to handle NaNs and such, like the LLVM fcmp instruction. + + -- | Floating-point less than. + FCmpLt FloatType + | -- | Floating-point less than or equal. + FCmpLe FloatType + | -- Boolean comparison. + + -- | Boolean less than. + CmpLlt + | -- | Boolean less than or equal. + CmpLle + deriving (Eq, Ord, Show, Generic) + +-- | Conversion operators try to generalise the @from t0 x to t1@ +-- instructions from LLVM. +data ConvOp + = -- | Zero-extend the former integer type to the latter. + -- If the new type is smaller, the result is a + -- truncation. + ZExt IntType IntType + | -- | Sign-extend the former integer type to the latter. + -- If the new type is smaller, the result is a + -- truncation. + SExt IntType IntType + | -- | Convert value of the former floating-point type to + -- the latter. If the new type is smaller, the result + -- is a truncation. + FPConv FloatType FloatType + | -- | Convert a floating-point value to the nearest + -- unsigned integer (rounding towards zero). + FPToUI FloatType IntType + | -- | Convert a floating-point value to the nearest + -- signed integer (rounding towards zero). + FPToSI FloatType IntType + | -- | Convert an unsigned integer to a floating-point value. + UIToFP IntType FloatType + | -- | Convert a signed integer to a floating-point value. + SIToFP IntType FloatType + | -- | Convert an integer to a boolean value. Zero + -- becomes false; anything else is true. + IToB IntType + | -- | Convert a boolean to an integer. True is converted + -- to 1 and False to 0. + BToI IntType + deriving (Eq, Ord, Show, Generic) + +-- | A list of all unary operators for all types. +allUnOps :: [UnOp] +allUnOps = + Not : + map Complement [minBound .. maxBound] + ++ map Abs [minBound .. maxBound] + ++ map FAbs [minBound .. maxBound] + ++ map SSignum [minBound .. maxBound] + ++ map USignum [minBound .. maxBound] + +-- | A list of all binary operators for all types. +allBinOps :: [BinOp] +allBinOps = + concat + [ map (`Add` OverflowWrap) allIntTypes, + map FAdd allFloatTypes, + map (`Sub` OverflowWrap) allIntTypes, + map FSub allFloatTypes, + map (`Mul` OverflowWrap) allIntTypes, + map FMul allFloatTypes, + map (`UDiv` Unsafe) allIntTypes, + map (`UDivUp` Unsafe) allIntTypes, + map (`SDiv` Unsafe) allIntTypes, + map (`SDivUp` Unsafe) allIntTypes, + map FDiv allFloatTypes, + map FMod allFloatTypes, + map (`UMod` Unsafe) allIntTypes, + map (`SMod` Unsafe) allIntTypes, + map (`SQuot` Unsafe) allIntTypes, + map (`SRem` Unsafe) allIntTypes, + map SMin allIntTypes, + map UMin allIntTypes, + map FMin allFloatTypes, + map SMax allIntTypes, + map UMax allIntTypes, + map FMax allFloatTypes, + map Shl allIntTypes, + map LShr allIntTypes, + map AShr allIntTypes, + map And allIntTypes, + map Or allIntTypes, + map Xor allIntTypes, + map Pow allIntTypes, + map FPow allFloatTypes, + [LogAnd, LogOr] + ] + +-- | A list of all comparison operators for all types. +allCmpOps :: [CmpOp] +allCmpOps = + concat + [ map CmpEq allPrimTypes, + map CmpUlt allIntTypes, + map CmpUle allIntTypes, + map CmpSlt allIntTypes, + map CmpSle allIntTypes, + map FCmpLt allFloatTypes, + map FCmpLe allFloatTypes + ] + +-- | A list of all conversion operators for all types. +allConvOps :: [ConvOp] +allConvOps = + concat + [ ZExt <$> allIntTypes <*> allIntTypes, + SExt <$> allIntTypes <*> allIntTypes, + FPConv <$> allFloatTypes <*> allFloatTypes, + FPToUI <$> allFloatTypes <*> allIntTypes, + FPToSI <$> allFloatTypes <*> allIntTypes, + UIToFP <$> allIntTypes <*> allFloatTypes, + SIToFP <$> allIntTypes <*> allFloatTypes, + IToB <$> allIntTypes, + BToI <$> allIntTypes + ] + +-- | Apply an 'UnOp' to an operand. Returns 'Nothing' if the +-- application is mistyped. +doUnOp :: UnOp -> PrimValue -> Maybe PrimValue +doUnOp Not (BoolValue b) = Just $ BoolValue $ not b +doUnOp Complement {} (IntValue v) = Just $ IntValue $ doComplement v +doUnOp Abs {} (IntValue v) = Just $ IntValue $ doAbs v +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 _ _ = Nothing + +-- | E.g., @~(~1) = 1@. +doComplement :: IntValue -> IntValue +doComplement v = intValue (intValueType v) $ complement $ intToInt64 v + +-- | @abs(-2) = 2@. +doAbs :: IntValue -> IntValue +doAbs v = intValue (intValueType v) $ abs $ intToInt64 v + +-- | @abs(-2.0) = 2.0@. +doFAbs :: FloatValue -> FloatValue +doFAbs v = floatValue (floatValueType v) $ abs $ floatToDouble v + +-- | @ssignum(-2)@ = -1. +doSSignum :: IntValue -> IntValue +doSSignum v = intValue (intValueType v) $ signum $ intToInt64 v + +-- | @usignum(-2)@ = -1. +doUSignum :: IntValue -> IntValue +doUSignum v = intValue (intValueType v) $ signum $ intToWord64 v + +-- | Apply a 'BinOp' to an operand. Returns 'Nothing' if the +-- application is mistyped, or outside the domain (e.g. division by +-- zero). +doBinOp :: BinOp -> PrimValue -> PrimValue -> Maybe PrimValue +doBinOp Add {} = doIntBinOp doAdd +doBinOp FAdd {} = doFloatBinOp (+) (+) +doBinOp Sub {} = doIntBinOp doSub +doBinOp FSub {} = doFloatBinOp (-) (-) +doBinOp Mul {} = doIntBinOp doMul +doBinOp FMul {} = doFloatBinOp (*) (*) +doBinOp UDiv {} = doRiskyIntBinOp doUDiv +doBinOp UDivUp {} = doRiskyIntBinOp doUDivUp +doBinOp SDiv {} = doRiskyIntBinOp doSDiv +doBinOp SDivUp {} = doRiskyIntBinOp doSDivUp +doBinOp FDiv {} = doFloatBinOp (/) (/) +doBinOp FMod {} = doFloatBinOp mod' mod' +doBinOp UMod {} = doRiskyIntBinOp doUMod +doBinOp SMod {} = doRiskyIntBinOp doSMod +doBinOp SQuot {} = doRiskyIntBinOp doSQuot +doBinOp SRem {} = doRiskyIntBinOp doSRem +doBinOp SMin {} = doIntBinOp doSMin +doBinOp UMin {} = doIntBinOp doUMin +doBinOp FMin {} = doFloatBinOp min min +doBinOp SMax {} = doIntBinOp doSMax +doBinOp UMax {} = doIntBinOp doUMax +doBinOp FMax {} = doFloatBinOp max max +doBinOp Shl {} = doIntBinOp doShl +doBinOp LShr {} = doIntBinOp doLShr +doBinOp AShr {} = doIntBinOp doAShr +doBinOp And {} = doIntBinOp doAnd +doBinOp Or {} = doIntBinOp doOr +doBinOp Xor {} = doIntBinOp doXor +doBinOp Pow {} = doRiskyIntBinOp doPow +doBinOp FPow {} = doFloatBinOp (**) (**) +doBinOp LogAnd {} = doBoolBinOp (&&) +doBinOp LogOr {} = doBoolBinOp (||) + +doIntBinOp :: + (IntValue -> IntValue -> IntValue) -> + PrimValue -> + PrimValue -> + Maybe PrimValue +doIntBinOp f (IntValue v1) (IntValue v2) = + Just $ IntValue $ f v1 v2 +doIntBinOp _ _ _ = Nothing + +doRiskyIntBinOp :: + (IntValue -> IntValue -> Maybe IntValue) -> + PrimValue -> + PrimValue -> + Maybe PrimValue +doRiskyIntBinOp f (IntValue v1) (IntValue v2) = + IntValue <$> f v1 v2 +doRiskyIntBinOp _ _ _ = Nothing + +doFloatBinOp :: + (Float -> Float -> Float) -> + (Double -> Double -> Double) -> + PrimValue -> + PrimValue -> + Maybe PrimValue +doFloatBinOp f32 _ (FloatValue (Float32Value v1)) (FloatValue (Float32Value v2)) = + Just $ FloatValue $ Float32Value $ f32 v1 v2 +doFloatBinOp _ f64 (FloatValue (Float64Value v1)) (FloatValue (Float64Value v2)) = + Just $ FloatValue $ Float64Value $ f64 v1 v2 +doFloatBinOp _ _ _ _ = Nothing + +doBoolBinOp :: + (Bool -> Bool -> Bool) -> + PrimValue -> + PrimValue -> + Maybe PrimValue +doBoolBinOp f (BoolValue v1) (BoolValue v2) = + Just $ BoolValue $ f v1 v2 +doBoolBinOp _ _ _ = Nothing + +-- | Integer addition. +doAdd :: IntValue -> IntValue -> IntValue +doAdd v1 v2 = intValue (intValueType v1) $ intToInt64 v1 + intToInt64 v2 + +-- | Integer subtraction. +doSub :: IntValue -> IntValue -> IntValue +doSub v1 v2 = intValue (intValueType v1) $ intToInt64 v1 - intToInt64 v2 + +-- | Integer multiplication. +doMul :: IntValue -> IntValue -> IntValue +doMul v1 v2 = intValue (intValueType v1) $ intToInt64 v1 * intToInt64 v2 + +-- | Unsigned integer division. Rounds towards negativity infinity. +-- Note: this is different from LLVM. +doUDiv :: IntValue -> IntValue -> Maybe IntValue +doUDiv v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = + Just $ + intValue (intValueType v1) $ + intToWord64 v1 `div` intToWord64 v2 + +-- | Unsigned integer division. Rounds towards positive infinity. +doUDivUp :: IntValue -> IntValue -> Maybe IntValue +doUDivUp v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = + Just $ + intValue (intValueType v1) $ + (intToWord64 v1 + intToWord64 v2 - 1) `div` intToWord64 v2 + +-- | Signed integer division. Rounds towards negativity infinity. +-- Note: this is different from LLVM. +doSDiv :: IntValue -> IntValue -> Maybe IntValue +doSDiv v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = + Just $ + intValue (intValueType v1) $ + intToInt64 v1 `div` intToInt64 v2 + +-- | Signed integer division. Rounds towards positive infinity. +doSDivUp :: IntValue -> IntValue -> Maybe IntValue +doSDivUp v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = + Just $ + intValue (intValueType v1) $ + (intToInt64 v1 + intToInt64 v2 - 1) `div` intToInt64 v2 + +-- | Unsigned integer modulus; the countepart to 'UDiv'. +doUMod :: IntValue -> IntValue -> Maybe IntValue +doUMod v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToWord64 v1 `mod` intToWord64 v2 + +-- | Signed integer modulus; the countepart to 'SDiv'. +doSMod :: IntValue -> IntValue -> Maybe IntValue +doSMod v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `mod` intToInt64 v2 + +-- | Signed integer division. Rounds towards zero. +-- This corresponds to the @sdiv@ instruction in LLVM. +doSQuot :: IntValue -> IntValue -> Maybe IntValue +doSQuot v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `quot` intToInt64 v2 + +-- | Signed integer division. Rounds towards zero. +-- This corresponds to the @srem@ instruction in LLVM. +doSRem :: IntValue -> IntValue -> Maybe IntValue +doSRem v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `rem` intToInt64 v2 + +-- | Minimum of two signed integers. +doSMin :: IntValue -> IntValue -> IntValue +doSMin v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `min` intToInt64 v2 + +-- | Minimum of two unsigned integers. +doUMin :: IntValue -> IntValue -> IntValue +doUMin v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `min` intToWord64 v2 + +-- | Maximum of two signed integers. +doSMax :: IntValue -> IntValue -> IntValue +doSMax v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `max` intToInt64 v2 + +-- | Maximum of two unsigned integers. +doUMax :: IntValue -> IntValue -> IntValue +doUMax v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `max` intToWord64 v2 + +-- | Left-shift. +doShl :: IntValue -> IntValue -> IntValue +doShl v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `shift` intToInt v2 + +-- | Logical right-shift, zero-extended. +doLShr :: IntValue -> IntValue -> IntValue +doLShr v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `shift` negate (intToInt v2) + +-- | Arithmetic right-shift, sign-extended. +doAShr :: IntValue -> IntValue -> IntValue +doAShr v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `shift` negate (intToInt v2) + +-- | Bitwise and. +doAnd :: IntValue -> IntValue -> IntValue +doAnd v1 v2 = intValue (intValueType v1) $ intToWord64 v1 .&. intToWord64 v2 + +-- | Bitwise or. +doOr :: IntValue -> IntValue -> IntValue +doOr v1 v2 = intValue (intValueType v1) $ intToWord64 v1 .|. intToWord64 v2 + +-- | Bitwise exclusive-or. +doXor :: IntValue -> IntValue -> IntValue +doXor v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `xor` intToWord64 v2 + +-- | Signed integer exponentatation. +doPow :: IntValue -> IntValue -> Maybe IntValue +doPow v1 v2 + | negativeIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 ^ intToInt64 v2 + +-- | Apply a 'ConvOp' to an operand. Returns 'Nothing' if the +-- application is mistyped. +doConvOp :: ConvOp -> PrimValue -> Maybe PrimValue +doConvOp (ZExt _ to) (IntValue v) = Just $ IntValue $ doZExt v to +doConvOp (SExt _ to) (IntValue v) = Just $ IntValue $ doSExt v to +doConvOp (FPConv _ to) (FloatValue v) = Just $ FloatValue $ doFPConv v to +doConvOp (FPToUI _ to) (FloatValue v) = Just $ IntValue $ doFPToUI v to +doConvOp (FPToSI _ to) (FloatValue v) = Just $ IntValue $ doFPToSI v to +doConvOp (UIToFP _ to) (IntValue v) = Just $ FloatValue $ doUIToFP v to +doConvOp (SIToFP _ to) (IntValue v) = Just $ FloatValue $ doSIToFP v to +doConvOp (IToB _) (IntValue v) = Just $ BoolValue $ intToInt64 v /= 0 +doConvOp (BToI to) (BoolValue v) = Just $ IntValue $ intValue to $ if v then 1 else 0 :: Int +doConvOp _ _ = Nothing + +-- | Zero-extend the given integer value to the size of the given +-- type. If the type is smaller than the given value, the result is a +-- truncation. +doZExt :: IntValue -> IntType -> IntValue +doZExt (Int8Value x) t = intValue t $ toInteger (fromIntegral x :: Word8) +doZExt (Int16Value x) t = intValue t $ toInteger (fromIntegral x :: Word16) +doZExt (Int32Value x) t = intValue t $ toInteger (fromIntegral x :: Word32) +doZExt (Int64Value x) t = intValue t $ toInteger (fromIntegral x :: Word64) + +-- | Sign-extend the given integer value to the size of the given +-- type. If the type is smaller than the given value, the result is a +-- truncation. +doSExt :: IntValue -> IntType -> IntValue +doSExt (Int8Value x) t = intValue t $ toInteger x +doSExt (Int16Value x) t = intValue t $ toInteger x +doSExt (Int32Value x) t = intValue t $ toInteger x +doSExt (Int64Value x) t = intValue t $ toInteger x + +-- | 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 + +-- | Convert a floating-point value to the nearest +-- unsigned integer (rounding towards zero). +doFPToUI :: FloatValue -> IntType -> IntValue +doFPToUI v t = intValue t (truncate $ floatToDouble v :: Word64) + +-- | Convert a floating-point value to the nearest +-- signed integer (rounding towards zero). +doFPToSI :: FloatValue -> IntType -> IntValue +doFPToSI v t = intValue t (truncate $ floatToDouble v :: Word64) + +-- | Convert an unsigned integer to a floating-point value. +doUIToFP :: IntValue -> FloatType -> FloatValue +doUIToFP v t = floatValue t $ intToWord64 v + +-- | Convert a signed integer to a floating-point value. +doSIToFP :: IntValue -> FloatType -> FloatValue +doSIToFP v t = floatValue t $ intToInt64 v + +-- | Apply a 'CmpOp' to an operand. Returns 'Nothing' if the +-- application is mistyped. +doCmpOp :: CmpOp -> PrimValue -> PrimValue -> Maybe Bool +doCmpOp CmpEq {} v1 v2 = Just $ doCmpEq v1 v2 +doCmpOp CmpUlt {} (IntValue v1) (IntValue v2) = Just $ doCmpUlt v1 v2 +doCmpOp CmpUle {} (IntValue v1) (IntValue v2) = Just $ doCmpUle v1 v2 +doCmpOp CmpSlt {} (IntValue v1) (IntValue v2) = Just $ doCmpSlt v1 v2 +doCmpOp CmpSle {} (IntValue v1) (IntValue v2) = Just $ doCmpSle v1 v2 +doCmpOp FCmpLt {} (FloatValue v1) (FloatValue v2) = Just $ doFCmpLt v1 v2 +doCmpOp FCmpLe {} (FloatValue v1) (FloatValue v2) = Just $ doFCmpLe v1 v2 +doCmpOp CmpLlt {} (BoolValue v1) (BoolValue v2) = Just $ not v1 && v2 +doCmpOp CmpLle {} (BoolValue v1) (BoolValue v2) = Just $ not (v1 && not v2) +doCmpOp _ _ _ = Nothing + +-- | Compare any two primtive values for exact equality. +doCmpEq :: PrimValue -> PrimValue -> Bool +doCmpEq (FloatValue (Float32Value v1)) (FloatValue (Float32Value v2)) = v1 == v2 +doCmpEq (FloatValue (Float64Value v1)) (FloatValue (Float64Value v2)) = v1 == v2 +doCmpEq v1 v2 = v1 == v2 + +-- | Unsigned less than. +doCmpUlt :: IntValue -> IntValue -> Bool +doCmpUlt v1 v2 = intToWord64 v1 < intToWord64 v2 + +-- | Unsigned less than or equal. +doCmpUle :: IntValue -> IntValue -> Bool +doCmpUle v1 v2 = intToWord64 v1 <= intToWord64 v2 + +-- | Signed less than. +doCmpSlt :: IntValue -> IntValue -> Bool +doCmpSlt = (<) + +-- | Signed less than or equal. +doCmpSle :: IntValue -> IntValue -> Bool +doCmpSle = (<=) + +-- | Floating-point less than. +doFCmpLt :: FloatValue -> FloatValue -> Bool +doFCmpLt = (<) + +-- | Floating-point less than or equal. +doFCmpLe :: FloatValue -> FloatValue -> Bool +doFCmpLe = (<=) + +-- | Translate an t'IntValue' to 'Word64'. This is guaranteed to fit. +intToWord64 :: IntValue -> Word64 +intToWord64 (Int8Value v) = fromIntegral (fromIntegral v :: Word8) +intToWord64 (Int16Value v) = fromIntegral (fromIntegral v :: Word16) +intToWord64 (Int32Value v) = fromIntegral (fromIntegral v :: Word32) +intToWord64 (Int64Value v) = fromIntegral (fromIntegral v :: Word64) + +-- | Translate an t'IntValue' to t'Int64'. This is guaranteed to fit. +intToInt64 :: IntValue -> Int64 +intToInt64 (Int8Value v) = fromIntegral v +intToInt64 (Int16Value v) = fromIntegral v +intToInt64 (Int32Value v) = fromIntegral v +intToInt64 (Int64Value v) = fromIntegral v + +-- | Careful - there is no guarantee this will fit. +intToInt :: IntValue -> Int +intToInt = fromIntegral . intToInt64 + +floatToDouble :: FloatValue -> Double +floatToDouble (Float32Value v) = fromRational $ toRational v +floatToDouble (Float64Value v) = v + +-- | The result type of a binary operator. +binOpType :: BinOp -> PrimType +binOpType (Add t _) = IntType t +binOpType (Sub t _) = IntType t +binOpType (Mul t _) = IntType t +binOpType (SDiv t _) = IntType t +binOpType (SDivUp t _) = IntType t +binOpType (SMod t _) = IntType t +binOpType (SQuot t _) = IntType t +binOpType (SRem t _) = IntType t +binOpType (UDiv t _) = IntType t +binOpType (UDivUp t _) = IntType t +binOpType (UMod t _) = IntType t +binOpType (SMin t) = IntType t +binOpType (UMin t) = IntType t +binOpType (FMin t) = FloatType t +binOpType (SMax t) = IntType t +binOpType (UMax t) = IntType t +binOpType (FMax t) = FloatType t +binOpType (Shl t) = IntType t +binOpType (LShr t) = IntType t +binOpType (AShr t) = IntType t +binOpType (And t) = IntType t +binOpType (Or t) = IntType t +binOpType (Xor t) = IntType t +binOpType (Pow t) = IntType t +binOpType (FPow t) = FloatType t +binOpType LogAnd = Bool +binOpType LogOr = Bool +binOpType (FAdd t) = FloatType t +binOpType (FSub t) = FloatType t +binOpType (FMul t) = FloatType t +binOpType (FDiv t) = FloatType t +binOpType (FMod t) = FloatType t + +-- | The operand types of a comparison operator. +cmpOpType :: CmpOp -> PrimType +cmpOpType (CmpEq t) = t +cmpOpType (CmpSlt t) = IntType t +cmpOpType (CmpSle t) = IntType t +cmpOpType (CmpUlt t) = IntType t +cmpOpType (CmpUle t) = IntType t +cmpOpType (FCmpLt t) = FloatType t +cmpOpType (FCmpLe t) = FloatType t +cmpOpType CmpLlt = Bool +cmpOpType CmpLle = Bool + +-- | The operand and result type of a unary operator. +unOpType :: UnOp -> PrimType +unOpType (SSignum t) = IntType t +unOpType (USignum t) = IntType t +unOpType Not = Bool +unOpType (Complement t) = IntType t +unOpType (Abs t) = IntType t +unOpType (FAbs t) = FloatType t + +-- | The input and output types of a conversion operator. +convOpType :: ConvOp -> (PrimType, PrimType) +convOpType (ZExt from to) = (IntType from, IntType to) +convOpType (SExt from to) = (IntType from, IntType to) +convOpType (FPConv from to) = (FloatType from, FloatType to) +convOpType (FPToUI from to) = (FloatType from, IntType to) +convOpType (FPToSI from to) = (FloatType from, IntType to) +convOpType (UIToFP from to) = (IntType from, FloatType to) +convOpType (SIToFP from to) = (IntType from, FloatType to) +convOpType (IToB from) = (IntType from, Bool) +convOpType (BToI to) = (Bool, IntType to) + +floatToWord :: Float -> Word32 +floatToWord = G.runGet G.getWord32le . P.runPut . P.putFloatle + +wordToFloat :: Word32 -> Float +wordToFloat = G.runGet G.getFloatle . P.runPut . P.putWord32le + +doubleToWord :: Double -> Word64 +doubleToWord = G.runGet G.getWord64le . P.runPut . P.putDoublele + +wordToDouble :: Word64 -> Double +wordToDouble = G.runGet G.getDoublele . P.runPut . P.putWord64le + +-- | A mapping from names of primitive functions to their parameter +-- types, their result type, and a function for evaluating them. +primFuns :: + M.Map + String + ( [PrimType], + PrimType, + [PrimValue] -> Maybe PrimValue + ) +primFuns = + M.fromList + [ f32 "sqrt32" sqrt, + f64 "sqrt64" sqrt, + f32 "log32" log, + f64 "log64" log, + f32 "log10_32" (logBase 10), + f64 "log10_64" (logBase 10), + f32 "log2_32" (logBase 2), + f64 "log2_64" (logBase 2), + f32 "exp32" exp, + f64 "exp64" exp, + f32 "sin32" sin, + f64 "sin64" sin, + f32 "sinh32" sinh, + f64 "sinh64" sinh, + f32 "cos32" cos, + f64 "cos64" cos, + f32 "cosh32" cosh, + f64 "cosh64" cosh, + f32 "tan32" tan, + f64 "tan64" tan, + f32 "tanh32" tanh, + f64 "tanh64" tanh, + f32 "asin32" asin, + f64 "asin64" asin, + f32 "asinh32" asinh, + f64 "asinh64" asinh, + f32 "acos32" acos, + f64 "acos64" acos, + f32 "acosh32" acosh, + f64 "acosh64" acosh, + f32 "atan32" atan, + f64 "atan64" atan, + f32 "atanh32" atanh, + f64 "atanh64" atanh, + f32 "round32" roundFloat, + f64 "round64" roundDouble, + f32 "ceil32" ceilFloat, + f64 "ceil64" ceilDouble, + f32 "floor32" floorFloat, + f64 "floor64" floorDouble, + f32 "gamma32" tgammaf, + f64 "gamma64" tgamma, + f32 "lgamma32" lgammaf, + f64 "lgamma64" lgamma, + i8 "clz8" $ IntValue . Int32Value . fromIntegral . countLeadingZeros, + i16 "clz16" $ IntValue . Int32Value . fromIntegral . countLeadingZeros, + i32 "clz32" $ IntValue . Int32Value . fromIntegral . countLeadingZeros, + i64 "clz64" $ IntValue . Int32Value . fromIntegral . countLeadingZeros, + i8 "ctz8" $ IntValue . Int32Value . fromIntegral . countTrailingZeros, + i16 "ctz16" $ IntValue . Int32Value . fromIntegral . countTrailingZeros, + i32 "ctz32" $ IntValue . Int32Value . fromIntegral . countTrailingZeros, + i64 "ctz64" $ IntValue . Int32Value . fromIntegral . countTrailingZeros, + i8 "popc8" $ IntValue . Int32Value . fromIntegral . popCount, + i16 "popc16" $ IntValue . Int32Value . fromIntegral . popCount, + i32 "popc32" $ IntValue . Int32Value . fromIntegral . popCount, + i64 "popc64" $ IntValue . Int32Value . fromIntegral . popCount, + ( "mad_hi8", + ( [IntType Int8, IntType Int8, IntType Int8], + IntType Int8, + \case + [IntValue (Int8Value a), IntValue (Int8Value b), IntValue (Int8Value c)] -> + Just $ IntValue . Int8Value $ mad_hi8 (Int8Value a) (Int8Value b) c + _ -> Nothing + ) + ), + ( "mad_hi16", + ( [IntType Int16, IntType Int16, IntType Int16], + IntType Int16, + \case + [IntValue (Int16Value a), IntValue (Int16Value b), IntValue (Int16Value c)] -> + Just $ IntValue . Int16Value $ mad_hi16 (Int16Value a) (Int16Value b) c + _ -> Nothing + ) + ), + ( "mad_hi32", + ( [IntType Int32, IntType Int32, IntType Int32], + IntType Int32, + \case + [IntValue (Int32Value a), IntValue (Int32Value b), IntValue (Int32Value c)] -> + Just $ IntValue . Int32Value $ mad_hi32 (Int32Value a) (Int32Value b) c + _ -> Nothing + ) + ), + ( "mad_hi64", + ( [IntType Int64, IntType Int64, IntType Int64], + IntType Int64, + \case + [IntValue (Int64Value a), IntValue (Int64Value b), IntValue (Int64Value c)] -> + Just $ IntValue . Int64Value $ mad_hi64 (Int64Value a) (Int64Value b) c + _ -> Nothing + ) + ), + ( "mul_hi8", + ( [IntType Int8, IntType Int8], + IntType Int8, + \case + [IntValue (Int8Value a), IntValue (Int8Value b)] -> + Just $ IntValue . Int8Value $ mul_hi8 (Int8Value a) (Int8Value b) + _ -> Nothing + ) + ), + ( "mul_hi16", + ( [IntType Int16, IntType Int16], + IntType Int16, + \case + [IntValue (Int16Value a), IntValue (Int16Value b)] -> + Just $ IntValue . Int16Value $ mul_hi16 (Int16Value a) (Int16Value b) + _ -> Nothing + ) + ), + ( "mul_hi32", + ( [IntType Int32, IntType Int32], + IntType Int32, + \case + [IntValue (Int32Value a), IntValue (Int32Value b)] -> + Just $ IntValue . Int32Value $ mul_hi32 (Int32Value a) (Int32Value b) + _ -> Nothing + ) + ), + ( "mul_hi64", + ( [IntType Int64, IntType Int64], + IntType Int64, + \case + [IntValue (Int64Value a), IntValue (Int64Value b)] -> + Just $ IntValue . Int64Value $ mul_hi64 (Int64Value a) (Int64Value b) + _ -> Nothing + ) + ), + ( "atan2_32", + ( [FloatType Float32, FloatType Float32], + FloatType Float32, + \case + [FloatValue (Float32Value x), FloatValue (Float32Value y)] -> + Just $ FloatValue $ Float32Value $ atan2 x y + _ -> Nothing + ) + ), + ( "atan2_64", + ( [FloatType Float64, FloatType Float64], + FloatType Float64, + \case + [FloatValue (Float64Value x), FloatValue (Float64Value y)] -> + Just $ FloatValue $ Float64Value $ atan2 x y + _ -> Nothing + ) + ), + ( "isinf32", + ( [FloatType Float32], + Bool, + \case + [FloatValue (Float32Value x)] -> Just $ BoolValue $ isInfinite x + _ -> Nothing + ) + ), + ( "isinf64", + ( [FloatType Float64], + Bool, + \case + [FloatValue (Float64Value x)] -> Just $ BoolValue $ isInfinite x + _ -> Nothing + ) + ), + ( "isnan32", + ( [FloatType Float32], + Bool, + \case + [FloatValue (Float32Value x)] -> Just $ BoolValue $ isNaN x + _ -> Nothing + ) + ), + ( "isnan64", + ( [FloatType Float64], + Bool, + \case + [FloatValue (Float64Value x)] -> Just $ BoolValue $ isNaN x + _ -> Nothing + ) + ), + ( "to_bits32", + ( [FloatType Float32], + IntType Int32, + \case + [FloatValue (Float32Value x)] -> + Just $ IntValue $ Int32Value $ fromIntegral $ floatToWord x + _ -> Nothing + ) + ), + ( "to_bits64", + ( [FloatType Float64], + IntType Int64, + \case + [FloatValue (Float64Value x)] -> + Just $ IntValue $ Int64Value $ fromIntegral $ doubleToWord x + _ -> Nothing + ) + ), + ( "from_bits32", + ( [IntType Int32], + FloatType Float32, + \case + [IntValue (Int32Value x)] -> + Just $ FloatValue $ Float32Value $ wordToFloat $ fromIntegral x + _ -> Nothing + ) + ), + ( "from_bits64", + ( [IntType Int64], + FloatType Float64, + \case + [IntValue (Int64Value x)] -> + Just $ FloatValue $ Float64Value $ wordToDouble $ fromIntegral x + _ -> Nothing + ) + ), + f32_3 "lerp32" (\v0 v1 t -> v0 + (v1 - v0) * max 0 (min 1 t)), + f64_3 "lerp64" (\v0 v1 t -> v0 + (v1 - v0) * max 0 (min 1 t)), + f32_3 "mad32" (\a b c -> a * b + c), + f64_3 "mad64" (\a b c -> a * b + c), + f32_3 "fma32" (\a b c -> a * b + c), + f64_3 "fma64" (\a b c -> a * b + c) + ] + where + i8 s f = (s, ([IntType Int8], IntType Int32, i8PrimFun f)) + i16 s f = (s, ([IntType Int16], IntType Int32, i16PrimFun f)) + i32 s f = (s, ([IntType Int32], IntType Int32, i32PrimFun f)) + i64 s f = (s, ([IntType Int64], IntType Int32, i64PrimFun f)) + f32 s f = (s, ([FloatType Float32], FloatType Float32, f32PrimFun f)) + f64 s f = (s, ([FloatType Float64], FloatType Float64, f64PrimFun f)) + f32_3 s f = + ( s, + ( [FloatType Float32, FloatType Float32, FloatType Float32], + FloatType Float32, + f32PrimFun3 f + ) + ) + f64_3 s f = + ( s, + ( [FloatType Float64, FloatType Float64, FloatType Float64], + FloatType Float64, + f64PrimFun3 f + ) + ) + + i8PrimFun f [IntValue (Int8Value x)] = + Just $ f x + i8PrimFun _ _ = Nothing + + i16PrimFun f [IntValue (Int16Value x)] = + Just $ f x + i16PrimFun _ _ = Nothing + + i32PrimFun f [IntValue (Int32Value x)] = + Just $ f x + i32PrimFun _ _ = Nothing + + i64PrimFun f [IntValue (Int64Value x)] = + Just $ f x + i64PrimFun _ _ = Nothing + + f32PrimFun f [FloatValue (Float32Value x)] = + Just $ FloatValue $ Float32Value $ f x + f32PrimFun _ _ = Nothing + + f64PrimFun f [FloatValue (Float64Value x)] = + Just $ FloatValue $ Float64Value $ f x + f64PrimFun _ _ = Nothing + + f32PrimFun3 + f + [ FloatValue (Float32Value a), + FloatValue (Float32Value b), + FloatValue (Float32Value c) + ] = + Just $ FloatValue $ Float32Value $ f a b c + f32PrimFun3 _ _ = Nothing + + f64PrimFun3 + f + [ FloatValue (Float64Value a), + FloatValue (Float64Value b), + FloatValue (Float64Value c) + ] = + Just $ FloatValue $ Float64Value $ f a b c + f64PrimFun3 _ _ = Nothing + +-- | Is the given value kind of zero? +zeroIsh :: PrimValue -> Bool +zeroIsh (IntValue k) = zeroIshInt k +zeroIsh (FloatValue (Float32Value k)) = k == 0 +zeroIsh (FloatValue (Float64Value k)) = k == 0 +zeroIsh (BoolValue False) = True +zeroIsh _ = False + +-- | Is the given value kind of one? +oneIsh :: PrimValue -> Bool +oneIsh (IntValue k) = oneIshInt k +oneIsh (FloatValue (Float32Value k)) = k == 1 +oneIsh (FloatValue (Float64Value k)) = k == 1 +oneIsh (BoolValue True) = True +oneIsh _ = False + +-- | Is the given value kind of negative? +negativeIsh :: PrimValue -> Bool +negativeIsh (IntValue k) = negativeIshInt k +negativeIsh (FloatValue (Float32Value k)) = k < 0 +negativeIsh (FloatValue (Float64Value k)) = k < 0 +negativeIsh (BoolValue _) = False +negativeIsh Checked = False + +-- | Is the given integer value kind of zero? +zeroIshInt :: IntValue -> Bool +zeroIshInt (Int8Value k) = k == 0 +zeroIshInt (Int16Value k) = k == 0 +zeroIshInt (Int32Value k) = k == 0 +zeroIshInt (Int64Value k) = k == 0 + +-- | Is the given integer value kind of one? +oneIshInt :: IntValue -> Bool +oneIshInt (Int8Value k) = k == 1 +oneIshInt (Int16Value k) = k == 1 +oneIshInt (Int32Value k) = k == 1 +oneIshInt (Int64Value k) = k == 1 + +-- | Is the given integer value kind of negative? +negativeIshInt :: IntValue -> Bool +negativeIshInt (Int8Value k) = k < 0 +negativeIshInt (Int16Value k) = k < 0 +negativeIshInt (Int32Value k) = k < 0 +negativeIshInt (Int64Value k) = k < 0 + +-- | The size of a value of a given primitive type in bites. +primBitSize :: PrimType -> Int +primBitSize = (* 8) . primByteSize + +-- | The size of a value of a given primitive type in eight-bit bytes. +primByteSize :: Num a => PrimType -> a +primByteSize (IntType t) = intByteSize t +primByteSize (FloatType t) = floatByteSize t +primByteSize Bool = 1 +primByteSize Cert = 1 + +-- | The size of a value of a given integer type in eight-bit bytes. +intByteSize :: Num a => IntType -> a +intByteSize Int8 = 1 +intByteSize Int16 = 2 +intByteSize Int32 = 4 +intByteSize Int64 = 8 + +-- | The size of a value of a given floating-point type in eight-bit bytes. +floatByteSize :: Num a => FloatType -> a +floatByteSize Float32 = 4 +floatByteSize Float64 = 8 + +-- | True if the given binary operator is commutative. +commutativeBinOp :: BinOp -> Bool +commutativeBinOp Add {} = True +commutativeBinOp FAdd {} = True +commutativeBinOp Mul {} = True +commutativeBinOp FMul {} = True +commutativeBinOp And {} = True +commutativeBinOp Or {} = True +commutativeBinOp Xor {} = True +commutativeBinOp LogOr {} = True +commutativeBinOp LogAnd {} = True +commutativeBinOp SMax {} = True +commutativeBinOp SMin {} = True +commutativeBinOp UMax {} = True +commutativeBinOp UMin {} = True +commutativeBinOp FMax {} = True +commutativeBinOp FMin {} = True +commutativeBinOp _ = False + +-- SexpIso instances + +instance SexpIso BinOp where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "add") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fadd") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "sub") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fsub") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "mul") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fmul") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "udiv") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "udivup") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "sdiv") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "sdivup") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fdiv") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fmod") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "umod") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "smod") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "squot") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "srem") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "smin") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "umin") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fmin") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "smax") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "umax") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fmap") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "shl") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "lshr") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "ashr") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "and") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "or") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "xor") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "pow") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fpow") >>> Sexp.el sexpIso)) $ + With (. Sexp.sym "logand") $ + With + (. Sexp.sym "logor") + End + +instance SexpIso CmpOp where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "eq") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "ult") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "ule") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "slt") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "sle") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "lt") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "le") >>> Sexp.el sexpIso)) $ + With (. Sexp.sym "llt") $ + With + (. Sexp.sym "lle") + End + +instance SexpIso ConvOp where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "zext") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "sext") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fpconv") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fptoui") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "fptosi") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "uitofp") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "sitofp") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "itob") >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "btoi") >>> Sexp.el sexpIso)) + End + +-- Prettyprinting instances + +instance Pretty BinOp where + ppr (Add t OverflowWrap) = taggedI "add" t + ppr (Add t OverflowUndef) = taggedI "add_nw" t + ppr (Sub t OverflowWrap) = taggedI "sub" t + ppr (Sub t OverflowUndef) = taggedI "sub_nw" t + ppr (Mul t OverflowWrap) = taggedI "mul" t + ppr (Mul t OverflowUndef) = taggedI "mul_nw" t + ppr (FAdd t) = taggedF "fadd" t + ppr (FSub t) = taggedF "fsub" t + ppr (FMul t) = taggedF "fmul" t + ppr (UDiv t Safe) = taggedI "udiv_safe" t + ppr (UDiv t Unsafe) = taggedI "udiv" t + ppr (UDivUp t Safe) = taggedI "udiv_up_safe" t + ppr (UDivUp t Unsafe) = taggedI "udiv_up" t + ppr (UMod t Safe) = taggedI "umod_safe" t + ppr (UMod t Unsafe) = taggedI "umod" t + ppr (SDiv t Safe) = taggedI "sdiv_safe" t + ppr (SDiv t Unsafe) = taggedI "sdiv" t + ppr (SDivUp t Safe) = taggedI "sdiv_up_safe" t + ppr (SDivUp t Unsafe) = taggedI "sdiv_up" t + ppr (SMod t Safe) = taggedI "smod_safe" t + ppr (SMod t Unsafe) = taggedI "smod" t + ppr (SQuot t Safe) = taggedI "squot_safe" t + ppr (SQuot t Unsafe) = taggedI "squot" t + ppr (SRem t Safe) = taggedI "srem_safe" t + ppr (SRem t Unsafe) = taggedI "srem" t + ppr (FDiv t) = taggedF "fdiv" t + ppr (FMod t) = taggedF "fmod" t + ppr (SMin t) = taggedI "smin" t + ppr (UMin t) = taggedI "umin" t + ppr (FMin t) = taggedF "fmin" t + ppr (SMax t) = taggedI "smax" t + ppr (UMax t) = taggedI "umax" t + ppr (FMax t) = taggedF "fmax" t + ppr (Shl t) = taggedI "shl" t + ppr (LShr t) = taggedI "lshr" t + ppr (AShr t) = taggedI "ashr" t + ppr (And t) = taggedI "and" t + ppr (Or t) = taggedI "or" t + ppr (Xor t) = taggedI "xor" t + ppr (Pow t) = taggedI "pow" t + ppr (FPow t) = taggedF "fpow" t + ppr LogAnd = text "logand" + ppr LogOr = text "logor" + +instance Pretty CmpOp where + ppr (CmpEq t) = text "eq_" <> ppr t + ppr (CmpUlt t) = taggedI "ult" t + ppr (CmpUle t) = taggedI "ule" t + ppr (CmpSlt t) = taggedI "slt" t + ppr (CmpSle t) = taggedI "sle" t + ppr (FCmpLt t) = taggedF "lt" t + ppr (FCmpLe t) = taggedF "le" t + ppr CmpLlt = text "llt" + ppr CmpLle = text "lle" + +instance Pretty ConvOp where + ppr op = convOp (convOpFun op) from to + where + (from, to) = convOpType op + +instance Pretty UnOp where + ppr Not = text "not" + ppr (Abs t) = taggedI "abs" t + ppr (FAbs t) = taggedF "fabs" t + ppr (SSignum t) = taggedI "ssignum" t + ppr (USignum t) = taggedI "usignum" t + ppr (Complement t) = taggedI "complement" t + +-- | The human-readable name for a 'ConvOp'. This is used to expose +-- the 'ConvOp' in the @intrinsics@ module of a Futhark program. +convOpFun :: ConvOp -> String +convOpFun ZExt {} = "zext" +convOpFun SExt {} = "sext" +convOpFun FPConv {} = "fpconv" +convOpFun FPToUI {} = "fptoui" +convOpFun FPToSI {} = "fptosi" +convOpFun UIToFP {} = "uitofp" +convOpFun SIToFP {} = "sitofp" +convOpFun IToB {} = "itob" +convOpFun BToI {} = "btoi" + +taggedI :: String -> IntType -> Doc +taggedI s Int8 = text $ s ++ "8" +taggedI s Int16 = text $ s ++ "16" +taggedI s Int32 = text $ s ++ "32" +taggedI s Int64 = text $ s ++ "64" + +taggedF :: String -> FloatType -> Doc +taggedF s Float32 = text $ s ++ "32" +taggedF s Float64 = text $ s ++ "64" + +convOp :: (Pretty from, Pretty to) => String -> from -> to -> Doc +convOp s from to = text s <> text "_" <> ppr from <> text "_" <> ppr to + +-- | True if signed. Only makes a difference for integer types. +prettySigned :: Bool -> PrimType -> String +prettySigned True (IntType it) = 'u' : drop 1 (pretty it) +prettySigned _ t = pretty t + +mul_hi8 :: IntValue -> IntValue -> Int8 +mul_hi8 a b = + let a' = intToWord64 a + b' = intToWord64 b + in fromIntegral (shiftR (a' * b') 8) + +mul_hi16 :: IntValue -> IntValue -> Int16 +mul_hi16 a b = + let a' = intToWord64 a + b' = intToWord64 b + in fromIntegral (shiftR (a' * b') 16) + +mul_hi32 :: IntValue -> IntValue -> Int32 +mul_hi32 a b = + let a' = intToWord64 a + b' = intToWord64 b + in fromIntegral (shiftR (a' * b') 32) + +mul_hi64 :: IntValue -> IntValue -> Int64 +mul_hi64 a b = + let a' = (toInteger . intToWord64) a + b' = (toInteger . intToWord64) b + in fromIntegral (shiftR (a' * b') 64) mad_hi8 :: IntValue -> IntValue -> Int8 -> Int8 mad_hi8 a b c = mul_hi8 a b + c diff --git a/src/Futhark/IR/Prop.hs b/src/Futhark/IR/Prop.hs --- a/src/Futhark/IR/Prop.hs +++ b/src/Futhark/IR/Prop.hs @@ -1,80 +1,81 @@ -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Safe #-} +{-# LANGUAGE TypeFamilies #-} + -- | This module provides various simple ways to query and manipulate -- fundamental Futhark terms, such as types and values. The intent is to -- keep "Futhark.IRrsentation.AST.Syntax" simple, and put whatever -- embellishments we need here. This is an internal, desugared -- representation. module Futhark.IR.Prop - ( module Futhark.IR.Prop.Reshape - , module Futhark.IR.Prop.Rearrange - , module Futhark.IR.Prop.Types - , module Futhark.IR.Prop.Constants - , module Futhark.IR.Prop.TypeOf - , module Futhark.IR.Prop.Patterns - , module Futhark.IR.Prop.Names - , module Futhark.IR.RetType - - -- * Built-in functions - , isBuiltInFunction - , builtInFunctions + ( module Futhark.IR.Prop.Reshape, + module Futhark.IR.Prop.Rearrange, + module Futhark.IR.Prop.Types, + module Futhark.IR.Prop.Constants, + module Futhark.IR.Prop.TypeOf, + module Futhark.IR.Prop.Patterns, + module Futhark.IR.Prop.Names, + module Futhark.IR.RetType, - -- * Extra tools - , asBasicOp - , safeExp - , subExpVars - , subExpVar - , shapeVars - , commutativeLambda - , entryPointSize - , defAux - , stmCerts - , certify - , expExtTypesFromPattern - , attrsForAssert + -- * Built-in functions + isBuiltInFunction, + builtInFunctions, - , ASTConstraints - , IsOp (..) - , ASTLore (..) + -- * Extra tools + asBasicOp, + safeExp, + subExpVars, + subExpVar, + shapeVars, + commutativeLambda, + entryPointSize, + defAux, + stmCerts, + certify, + expExtTypesFromPattern, + attrsForAssert, + ASTConstraints, + IsOp (..), + ASTLore (..), ) - where +where import Data.List (find) -import Data.Maybe (mapMaybe, isJust) import qualified Data.Map.Strict as M +import Data.Maybe (isJust, mapMaybe) import qualified Data.Set as S - -import Futhark.IR.Prop.Reshape -import Futhark.IR.Prop.Rearrange -import Futhark.IR.Prop.Types +import Futhark.IR.Pretty import Futhark.IR.Prop.Constants -import Futhark.IR.Prop.Patterns import Futhark.IR.Prop.Names +import Futhark.IR.Prop.Patterns +import Futhark.IR.Prop.Rearrange +import Futhark.IR.Prop.Reshape import Futhark.IR.Prop.TypeOf +import Futhark.IR.Prop.Types import Futhark.IR.RetType import Futhark.IR.Syntax -import Futhark.IR.Pretty import Futhark.Transform.Rename (Rename, Renameable) -import Futhark.Transform.Substitute (Substitute, Substitutable) +import Futhark.Transform.Substitute (Substitutable, Substitute) import Futhark.Util.Pretty +import Language.SexpGrammar as Sexp -- | @isBuiltInFunction k@ is 'True' if @k@ is an element of 'builtInFunctions'. isBuiltInFunction :: Name -> Bool isBuiltInFunction fnm = fnm `M.member` builtInFunctions -- | A map of all built-in functions and their types. -builtInFunctions :: M.Map Name (PrimType,[PrimType]) +builtInFunctions :: M.Map Name (PrimType, [PrimType]) builtInFunctions = M.fromList $ map namify $ M.toList primFuns - where namify (k,(paramts,ret,_)) = (nameFromString k, (ret, paramts)) + where + namify (k, (paramts, ret, _)) = (nameFromString k, (ret, paramts)) -- | If the expression is a t'BasicOp', return it, otherwise 'Nothing'. asBasicOp :: Exp lore -> Maybe BasicOp asBasicOp (BasicOp op) = Just op -asBasicOp _ = Nothing +asBasicOp _ = Nothing -- | An expression is safe if it is always well-defined (assuming that -- any required certificates have been checked) in any context. For @@ -82,57 +83,54 @@ -- bounds. On the other hand, adding two numbers cannot fail. safeExp :: IsOp (Op lore) => Exp lore -> Bool safeExp (BasicOp op) = safeBasicOp op - where safeBasicOp (BinOp (SDiv _ Safe) _ _) = True - safeBasicOp (BinOp (SDivUp _ Safe) _ _) = True - safeBasicOp (BinOp (SQuot _ Safe) _ _) = True - safeBasicOp (BinOp (UDiv _ Safe) _ _) = True - safeBasicOp (BinOp (UDivUp _ Safe) _ _) = True - safeBasicOp (BinOp (SMod _ Safe) _ _) = True - safeBasicOp (BinOp (SRem _ Safe) _ _) = True - safeBasicOp (BinOp (UMod _ Safe) _ _) = True - - safeBasicOp (BinOp SDiv{} _ (Constant y)) = not $ zeroIsh y - safeBasicOp (BinOp SDiv{} _ _) = False - safeBasicOp (BinOp SDivUp{} _ (Constant y)) = not $ zeroIsh y - safeBasicOp (BinOp SDivUp{} _ _) = False - safeBasicOp (BinOp UDiv{} _ (Constant y)) = not $ zeroIsh y - safeBasicOp (BinOp UDiv{} _ _) = False - safeBasicOp (BinOp UDivUp{} _ (Constant y)) = not $ zeroIsh y - safeBasicOp (BinOp UDivUp{} _ _) = False - safeBasicOp (BinOp SMod{} _ (Constant y)) = not $ zeroIsh y - safeBasicOp (BinOp SMod{} _ _) = False - safeBasicOp (BinOp UMod{} _ (Constant y)) = not $ zeroIsh y - safeBasicOp (BinOp UMod{} _ _) = False - - safeBasicOp (BinOp SQuot{} _ (Constant y)) = not $ zeroIsh y - safeBasicOp (BinOp SQuot{} _ _) = False - safeBasicOp (BinOp SRem{} _ (Constant y)) = not $ zeroIsh y - safeBasicOp (BinOp SRem{} _ _) = False - - safeBasicOp (BinOp Pow{} _ (Constant y)) = not $ negativeIsh y - safeBasicOp (BinOp Pow{} _ _) = False - safeBasicOp ArrayLit{} = True - safeBasicOp BinOp{} = True - safeBasicOp SubExp{} = True - safeBasicOp UnOp{} = True - safeBasicOp CmpOp{} = True - safeBasicOp ConvOp{} = True - safeBasicOp Scratch{} = True - safeBasicOp Concat{} = True - safeBasicOp Reshape{} = True - safeBasicOp Rearrange{} = True - safeBasicOp Manifest{} = True - safeBasicOp Iota{} = True - safeBasicOp Replicate{} = True - safeBasicOp Copy{} = True - safeBasicOp _ = False - + where + safeBasicOp (BinOp (SDiv _ Safe) _ _) = True + safeBasicOp (BinOp (SDivUp _ Safe) _ _) = True + safeBasicOp (BinOp (SQuot _ Safe) _ _) = True + safeBasicOp (BinOp (UDiv _ Safe) _ _) = True + safeBasicOp (BinOp (UDivUp _ Safe) _ _) = True + safeBasicOp (BinOp (SMod _ Safe) _ _) = True + safeBasicOp (BinOp (SRem _ Safe) _ _) = True + safeBasicOp (BinOp (UMod _ Safe) _ _) = True + safeBasicOp (BinOp SDiv {} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp SDiv {} _ _) = False + safeBasicOp (BinOp SDivUp {} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp SDivUp {} _ _) = False + safeBasicOp (BinOp UDiv {} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp UDiv {} _ _) = False + safeBasicOp (BinOp UDivUp {} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp UDivUp {} _ _) = False + safeBasicOp (BinOp SMod {} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp SMod {} _ _) = False + safeBasicOp (BinOp UMod {} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp UMod {} _ _) = False + safeBasicOp (BinOp SQuot {} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp SQuot {} _ _) = False + safeBasicOp (BinOp SRem {} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp SRem {} _ _) = False + safeBasicOp (BinOp Pow {} _ (Constant y)) = not $ negativeIsh y + safeBasicOp (BinOp Pow {} _ _) = False + safeBasicOp ArrayLit {} = True + safeBasicOp BinOp {} = True + safeBasicOp SubExp {} = True + safeBasicOp UnOp {} = True + safeBasicOp CmpOp {} = True + safeBasicOp ConvOp {} = True + safeBasicOp Scratch {} = True + safeBasicOp Concat {} = True + safeBasicOp Reshape {} = True + safeBasicOp Rearrange {} = True + safeBasicOp Manifest {} = True + safeBasicOp Iota {} = True + safeBasicOp Replicate {} = True + safeBasicOp Copy {} = True + safeBasicOp _ = False safeExp (DoLoop _ _ _ body) = safeBody body safeExp (Apply fname _ _ _) = isBuiltInFunction fname safeExp (If _ tbranch fbranch _) = - all (safeExp . stmExp) (bodyStms tbranch) && - all (safeExp . stmExp) (bodyStms fbranch) + all (safeExp . stmExp) (bodyStms tbranch) + && all (safeExp . stmExp) (bodyStms fbranch) safeExp (Op op) = safeOp op safeBody :: IsOp (Op lore) => Body lore -> Bool @@ -145,8 +143,8 @@ -- | If the t'SubExp' is a 'Var' return the variable name. subExpVar :: SubExp -> Maybe VName -subExpVar (Var v) = Just v -subExpVar Constant{} = Nothing +subExpVar (Var v) = Just v +subExpVar Constant {} = Nothing -- | Return the variable dimension sizes. May contain -- duplicates. @@ -161,19 +159,19 @@ commutativeLambda lam = let body = lambdaBody lam n2 = length (lambdaParams lam) `div` 2 - (xps,yps) = splitAt n2 (lambdaParams lam) + (xps, yps) = splitAt n2 (lambdaParams lam) okComponent c = isJust $ find (okBinOp c) $ bodyStms body - okBinOp (xp,yp,Var r) (Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y)))) = - patElemName pe == r && - commutativeBinOp op && - ((x == paramName xp && y == paramName yp) || - (y == paramName xp && x == paramName yp)) + okBinOp (xp, yp, Var r) (Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y)))) = + patElemName pe == r + && commutativeBinOp op + && ( (x == paramName xp && y == paramName yp) + || (y == paramName xp && x == paramName yp) + ) okBinOp _ _ = False - - in n2 * 2 == length (lambdaParams lam) && - n2 == length (bodyResult body) && - all okComponent (zip3 xps yps $ bodyResult body) + in n2 * 2 == length (lambdaParams lam) + && n2 == length (bodyResult body) + && all okComponent (zip3 xps yps $ bodyResult body) -- | How many value parameters are accepted by this entry point? This -- is used to determine which of the function parameters correspond to @@ -195,17 +193,18 @@ -- | Add certificates to a statement. certify :: Certificates -> Stm lore -> Stm lore certify cs1 (Let pat (StmAux cs2 attrs dec) e) = - Let pat (StmAux (cs2<>cs1) attrs dec) e + Let pat (StmAux (cs2 <> cs1) attrs dec) e -- | A handy shorthand for properties that we usually want to things -- we stuff into ASTs. type ASTConstraints a = - (Eq a, Ord a, Show a, Rename a, Substitute a, FreeIn a, Pretty a) + (Eq a, Ord a, Show a, Rename a, Substitute a, FreeIn a, Pretty a, SexpIso a) -- | A type class for operations. class (ASTConstraints op, TypedOp op) => IsOp op where -- | Like 'safeExp', but for arbitrary ops. safeOp :: op -> Bool + -- | Should we try to hoist this out of branches? cheapOp :: op -> Bool @@ -215,35 +214,40 @@ -- | Lore-specific attributes; also means the lore supports some basic -- facilities. -class (Decorations lore, - - PrettyLore lore, - - Renameable lore, Substitutable lore, - FreeDec (ExpDec lore), - FreeIn (LetDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (RetType lore), - FreeIn (BranchType lore), - - IsOp (Op lore)) => ASTLore lore where +class + ( Decorations lore, + PrettyLore lore, + Renameable lore, + Substitutable lore, + FreeDec (ExpDec lore), + FreeIn (LetDec lore), + FreeDec (BodyDec lore), + FreeIn (FParamInfo lore), + FreeIn (LParamInfo lore), + FreeIn (RetType lore), + FreeIn (BranchType lore), + IsOp (Op lore) + ) => + ASTLore lore + where -- | Given a pattern, construct the type of a body that would match -- it. An implementation for many lores would be -- 'expExtTypesFromPattern'. - expTypesFromPattern :: (HasScope lore m, Monad m) => - Pattern lore -> m [BranchType lore] + expTypesFromPattern :: + (HasScope lore m, Monad m) => + Pattern lore -> + m [BranchType lore] -- | Construct the type of an expression that would match the pattern. expExtTypesFromPattern :: Typed dec => PatternT dec -> [ExtType] expExtTypesFromPattern pat = existentialiseExtTypes (patternContextNames pat) $ - staticShapes $ map patElemType $ patternValueElements pat + staticShapes $ map patElemType $ patternValueElements pat -- | Keep only those attributes that are relevant for 'Assert' -- expressions. attrsForAssert :: Attrs -> Attrs attrsForAssert (Attrs attrs) = Attrs $ S.filter attrForAssert attrs - where attrForAssert = (==AttrComp "warn" ["safety_checks"]) + where + attrForAssert = (== AttrComp "warn" ["safety_checks"]) 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 @@ -1,5 +1,7 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-} -{-# Language FlexibleInstances, FlexibleContexts #-} + -- | The IR tracks aliases, mostly to ensure the soundness of in-place -- updates, but it can also be used for other things (such as memory -- optimisations). This module contains the raw building blocks for @@ -10,35 +12,42 @@ -- Thus, they do not take aliases-of-aliases into account. See -- "Futhark.Analysis.Alias" if this is not what you want. module Futhark.IR.Prop.Aliases - ( subExpAliases - , expAliases - , patternAliases - , Aliased (..) - , AliasesOf (..) - -- * Consumption - , consumedInStm - , consumedInExp - , consumedByLambda - -- * Extensibility - , AliasedOp (..) - , CanBeAliased (..) - ) - where + ( subExpAliases, + expAliases, + patternAliases, + Aliased (..), + AliasesOf (..), + -- * Consumption + consumedInStm, + consumedInExp, + consumedByLambda, + + -- * Extensibility + AliasedOp (..), + CanBeAliased (..), + ) +where + import Control.Arrow (first) import qualified Data.Kind - import Futhark.IR.Prop (IsOp) -import Futhark.IR.Syntax +import Futhark.IR.Prop.Names import Futhark.IR.Prop.Patterns import Futhark.IR.Prop.Types -import Futhark.IR.Prop.Names +import Futhark.IR.Syntax -- | The class of lores that contain aliasing information. -class (Decorations lore, AliasedOp (Op lore), - AliasesOf (LetDec lore)) => Aliased lore where +class + ( Decorations lore, + AliasedOp (Op lore), + AliasesOf (LetDec lore) + ) => + Aliased lore + where -- | The aliases of the body results. bodyAliases :: Body lore -> [Names] + -- | The variables consumed in the body. consumedInBody :: Body lore -> Names @@ -47,71 +56,76 @@ -- | The alises of a subexpression. subExpAliases :: SubExp -> Names -subExpAliases Constant{} = mempty -subExpAliases (Var v) = vnameAliases v +subExpAliases Constant {} = mempty +subExpAliases (Var v) = vnameAliases v basicOpAliases :: BasicOp -> [Names] basicOpAliases (SubExp se) = [subExpAliases se] basicOpAliases (Opaque se) = [subExpAliases se] basicOpAliases (ArrayLit _ _) = [mempty] -basicOpAliases BinOp{} = [mempty] -basicOpAliases ConvOp{} = [mempty] -basicOpAliases CmpOp{} = [mempty] -basicOpAliases UnOp{} = [mempty] +basicOpAliases BinOp {} = [mempty] +basicOpAliases ConvOp {} = [mempty] +basicOpAliases CmpOp {} = [mempty] +basicOpAliases UnOp {} = [mempty] basicOpAliases (Index ident _) = [vnameAliases ident] -basicOpAliases Update{} = [mempty] -basicOpAliases Iota{} = [mempty] -basicOpAliases Replicate{} = [mempty] -basicOpAliases Scratch{} = [mempty] +basicOpAliases Update {} = [mempty] +basicOpAliases Iota {} = [mempty] +basicOpAliases Replicate {} = [mempty] +basicOpAliases Scratch {} = [mempty] basicOpAliases (Reshape _ e) = [vnameAliases e] basicOpAliases (Rearrange _ e) = [vnameAliases e] basicOpAliases (Rotate _ e) = [vnameAliases e] -basicOpAliases Concat{} = [mempty] -basicOpAliases Copy{} = [mempty] -basicOpAliases Manifest{} = [mempty] -basicOpAliases Assert{} = [mempty] +basicOpAliases Concat {} = [mempty] +basicOpAliases Copy {} = [mempty] +basicOpAliases Manifest {} = [mempty] +basicOpAliases Assert {} = [mempty] ifAliases :: ([Names], Names) -> ([Names], Names) -> [Names] -ifAliases (als1,cons1) (als2,cons2) = +ifAliases (als1, cons1) (als2, cons2) = map (`namesSubtract` cons) $ zipWith mappend als1 als2 - where cons = cons1 <> cons2 + where + cons = cons1 <> cons2 funcallAliases :: [(SubExp, Diet)] -> [TypeBase shape Uniqueness] -> [Names] funcallAliases args t = - returnAliases t [(subExpAliases se, d) | (se,d) <- args ] + returnAliases t [(subExpAliases se, d) | (se, d) <- args] -- | The aliases of an expression, one per non-context value returned. expAliases :: (Aliased lore) => Exp lore -> [Names] expAliases (If _ tb fb dec) = drop (length all_aliases - length ts) all_aliases - where ts = ifReturns dec - all_aliases = ifAliases - (bodyAliases tb, consumedInBody tb) - (bodyAliases fb, consumedInBody fb) + where + ts = ifReturns dec + all_aliases = + ifAliases + (bodyAliases tb, consumedInBody tb) + (bodyAliases fb, consumedInBody fb) expAliases (BasicOp op) = basicOpAliases op expAliases (DoLoop ctxmerge valmerge _ loopbody) = map (`namesSubtract` merge_names) val_aliases - where (_ctx_aliases, val_aliases) = - splitAt (length ctxmerge) $ bodyAliases loopbody - merge_names = namesFromList $ map (paramName . fst) $ ctxmerge ++ valmerge + where + (_ctx_aliases, val_aliases) = + splitAt (length ctxmerge) $ bodyAliases loopbody + merge_names = namesFromList $ map (paramName . fst) $ ctxmerge ++ valmerge expAliases (Apply _ args t _) = funcallAliases args $ map declExtTypeOf t expAliases (Op op) = opAliases op returnAliases :: [TypeBase shape Uniqueness] -> [(Names, Diet)] -> [Names] returnAliases rts args = map returnType' rts - where returnType' (Array _ _ Nonunique) = - mconcat $ map (uncurry maskAliases) args - returnType' (Array _ _ Unique) = - mempty - returnType' (Prim _) = - mempty - returnType' Mem{} = - error "returnAliases Mem" + where + returnType' (Array _ _ Nonunique) = + mconcat $ map (uncurry maskAliases) args + returnType' (Array _ _ Unique) = + mempty + returnType' (Prim _) = + mempty + returnType' Mem {} = + error "returnAliases Mem" maskAliases :: Names -> Diet -> Names -maskAliases _ Consume = mempty -maskAliases _ ObservePrim = mempty +maskAliases _ Consume = mempty +maskAliases _ ObservePrim = mempty maskAliases als Observe = als -- | The variables consumed in this statement. @@ -122,13 +136,16 @@ consumedInExp :: (Aliased lore) => Exp lore -> Names consumedInExp (Apply _ args _ _) = mconcat (map (consumeArg . first subExpAliases) args) - where consumeArg (als, Consume) = als - consumeArg _ = mempty + where + consumeArg (als, Consume) = als + consumeArg _ = mempty consumedInExp (If _ tb fb _) = consumedInBody tb <> consumedInBody fb consumedInExp (DoLoop _ merge _ _) = - mconcat (map (subExpAliases . snd) $ - filter (unique . paramDeclType . fst) merge) + mconcat + ( map (subExpAliases . snd) $ + filter (unique . paramDeclType . fst) merge + ) consumedInExp (BasicOp (Update src _ _)) = oneName src consumedInExp (Op op) = consumedInOp op consumedInExp _ = mempty diff --git a/src/Futhark/IR/Prop/Constants.hs b/src/Futhark/IR/Prop/Constants.hs --- a/src/Futhark/IR/Prop/Constants.hs +++ b/src/Futhark/IR/Prop/Constants.hs @@ -1,12 +1,11 @@ -- | Possibly convenient facilities for constructing constants. module Futhark.IR.Prop.Constants - ( - IsValue (..) - , constant - , intConst - , floatConst - ) - where + ( IsValue (..), + constant, + intConst, + floatConst, + ) +where import Futhark.IR.Syntax.Core @@ -17,9 +16,6 @@ -- (LogVal True) loc@. class IsValue a where value :: a -> PrimValue - -instance IsValue Int where - value = IntValue . Int32Value . fromIntegral instance IsValue Int8 where value = IntValue . Int8Value diff --git a/src/Futhark/IR/Prop/Names.hs b/src/Futhark/IR/Prop/Names.hs --- a/src/Futhark/IR/Prop/Names.hs +++ b/src/Futhark/IR/Prop/Names.hs @@ -1,55 +1,72 @@ -{-# LANGUAGE FlexibleInstances, FlexibleContexts, UndecidableInstances #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE Safe #-} +{-# LANGUAGE UndecidableInstances #-} + -- | Facilities for determining which names are used in some syntactic -- construct. The most important interface is the 'FreeIn' class and -- its instances, but for reasons related to the Haskell type system, -- some constructs have specialised functions. module Futhark.IR.Prop.Names - ( -- * Free names - Names - , namesIntMap - , nameIn - , oneName - , namesFromList - , namesToList - , namesIntersection - , namesIntersect - , namesSubtract - , mapNames - -- * Class - , FreeIn (..) - , freeIn - -- * Specialised Functions - , freeInStmsAndRes - -- * Bound Names - , boundInBody - , boundByStm - , boundByStms - , boundByLambda - -- * Efficient computation - , FreeDec(..) - , FV - , fvBind - , fvName - , fvNames - ) - where + ( -- * Free names + Names, + namesIntMap, + nameIn, + oneName, + namesFromList, + namesToList, + namesIntersection, + namesIntersect, + namesSubtract, + mapNames, + -- * Class + FreeIn (..), + freeIn, + + -- * Specialised Functions + freeInStmsAndRes, + + -- * Bound Names + boundInBody, + boundByStm, + boundByStms, + boundByLambda, + + -- * Efficient computation + FreeDec (..), + FV, + fvBind, + fvName, + fvNames, + ) +where + +import Control.Category import Control.Monad.State.Strict +import Data.Foldable import qualified Data.IntMap.Strict as IM import qualified Data.Map.Strict as M -import Data.Foldable - -import Futhark.IR.Syntax -import Futhark.IR.Traversals import Futhark.IR.Prop.Patterns import Futhark.IR.Prop.Scope +import Futhark.IR.Syntax +import Futhark.IR.Traversals import Futhark.Util.Pretty +import GHC.Generics +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) -- | A set of names. Note that the 'Ord' instance is a dummy that -- treats everything as 'EQ' if '==', and otherwise 'LT'. newtype Names = Names (IM.IntMap VName) - deriving (Eq, Show) + deriving (Eq, Show, Generic) +instance SexpIso Names where + sexpIso = with $ \names -> + (iso IM.fromList IM.toList . sexpIso) >>> names + -- | Retrieve the data structure underlying the names representation. namesIntMap :: Names -> IM.IntMap VName namesIntMap (Names m) = m @@ -99,7 +116,8 @@ mapNames f vs = namesFromList $ map f $ namesToList vs -- | A computation to build a free variable set. -newtype FV = FV { unFV :: Names } +newtype FV = FV {unFV :: Names} + -- Right now the variable set is just stored explicitly, without the -- fancy functional representation that GHC uses. Turns out it's -- faster this way. @@ -122,32 +140,39 @@ fvNames :: Names -> FV fvNames = FV -freeWalker :: (FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (Op lore)) => - Walker lore (State FV) -freeWalker = identityWalker { - walkOnSubExp = modify . (<>) . freeIn' - , walkOnBody = \scope body -> do - modify $ (<>) $ freeIn' body - modify $ fvBind (namesFromList (M.keys scope)) - , walkOnVName = modify . (<>) . fvName - , walkOnOp = modify . (<>) . freeIn' - } +freeWalker :: + ( FreeDec (ExpDec lore), + FreeDec (BodyDec lore), + FreeIn (FParamInfo lore), + FreeIn (LParamInfo lore), + FreeIn (LetDec lore), + FreeIn (Op lore) + ) => + Walker lore (State FV) +freeWalker = + identityWalker + { walkOnSubExp = modify . (<>) . freeIn', + walkOnBody = \scope body -> do + modify $ (<>) $ freeIn' body + modify $ fvBind (namesFromList (M.keys scope)), + walkOnVName = modify . (<>) . fvName, + walkOnOp = modify . (<>) . freeIn' + } -- | Return the set of variable names that are free in the given -- statements and result. Filters away the names that are bound by -- the statements. -freeInStmsAndRes :: (FreeIn (Op lore), - FreeIn (LetDec lore), - FreeIn (LParamInfo lore), - FreeIn (FParamInfo lore), - FreeDec (BodyDec lore), - FreeDec (ExpDec lore)) => - Stms lore -> Result -> FV +freeInStmsAndRes :: + ( FreeIn (Op lore), + FreeIn (LetDec lore), + FreeIn (LParamInfo lore), + FreeIn (FParamInfo lore), + FreeDec (BodyDec lore), + FreeDec (ExpDec lore) + ) => + Stms lore -> + Result -> + FV freeInStmsAndRes stms res = fvBind (boundByStms stms) $ foldMap freeIn' stms <> freeIn' res @@ -170,71 +195,94 @@ instance FreeIn Int where freeIn' = const mempty -instance (FreeIn a, FreeIn b) => FreeIn (a,b) where - freeIn' (a,b) = freeIn' a <> freeIn' b +instance (FreeIn a, FreeIn b) => FreeIn (a, b) where + freeIn' (a, b) = freeIn' a <> freeIn' b -instance (FreeIn a, FreeIn b, FreeIn c) => FreeIn (a,b,c) where - freeIn' (a,b,c) = freeIn' a <> freeIn' b <> freeIn' c +instance (FreeIn a, FreeIn b, FreeIn c) => FreeIn (a, b, c) where + freeIn' (a, b, c) = freeIn' a <> freeIn' b <> freeIn' c instance FreeIn a => FreeIn [a] where freeIn' = foldMap freeIn' -instance (FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (RetType lore), - FreeIn (Op lore)) => FreeIn (FunDef lore) where +instance + ( FreeDec (ExpDec lore), + FreeDec (BodyDec lore), + FreeIn (FParamInfo lore), + FreeIn (LParamInfo lore), + FreeIn (LetDec lore), + FreeIn (RetType lore), + FreeIn (Op lore) + ) => + FreeIn (FunDef lore) + where freeIn' (FunDef _ _ _ rettype params body) = fvBind (namesFromList $ map paramName params) $ - freeIn' rettype <> freeIn' params <> freeIn' body + freeIn' rettype <> freeIn' params <> freeIn' body -instance (FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (Op lore)) => FreeIn (Lambda lore) where +instance + ( FreeDec (ExpDec lore), + FreeDec (BodyDec lore), + FreeIn (FParamInfo lore), + FreeIn (LParamInfo lore), + FreeIn (LetDec lore), + FreeIn (Op lore) + ) => + FreeIn (Lambda lore) + where freeIn' (Lambda params body rettype) = fvBind (namesFromList $ map paramName params) $ - freeIn' rettype <> freeIn' params <> freeIn' body + freeIn' rettype <> freeIn' params <> freeIn' body -instance (FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (Op lore)) => FreeIn (Body lore) where +instance + ( FreeDec (ExpDec lore), + FreeDec (BodyDec lore), + FreeIn (FParamInfo lore), + FreeIn (LParamInfo lore), + FreeIn (LetDec lore), + FreeIn (Op lore) + ) => + FreeIn (Body lore) + where freeIn' (Body dec stms res) = precomputed dec $ freeIn' dec <> freeInStmsAndRes stms res -instance (FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (Op lore)) => FreeIn (Exp lore) where +instance + ( FreeDec (ExpDec lore), + FreeDec (BodyDec lore), + FreeIn (FParamInfo lore), + FreeIn (LParamInfo lore), + FreeIn (LetDec lore), + FreeIn (Op lore) + ) => + FreeIn (Exp lore) + where freeIn' (DoLoop ctxmerge valmerge form loopbody) = let (ctxparams, ctxinits) = unzip ctxmerge (valparams, valinits) = unzip valmerge - bound_here = namesFromList $ M.keys $ - scopeOf form <> - scopeOfFParams (ctxparams ++ valparams) - in fvBind bound_here $ - freeIn' (ctxinits ++ valinits) <> freeIn' form <> - freeIn' (ctxparams ++ valparams) <> freeIn' loopbody + bound_here = + namesFromList $ + M.keys $ + scopeOf form + <> scopeOfFParams (ctxparams ++ valparams) + in fvBind bound_here $ + freeIn' (ctxinits ++ valinits) <> freeIn' form + <> freeIn' (ctxparams ++ valparams) + <> freeIn' loopbody freeIn' e = execState (walkExpM freeWalker e) mempty -instance (FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (Op lore)) => FreeIn (Stm lore) where +instance + ( FreeDec (ExpDec lore), + FreeDec (BodyDec lore), + FreeIn (FParamInfo lore), + FreeIn (LParamInfo lore), + FreeIn (LetDec lore), + FreeIn (Op lore) + ) => + FreeIn (Stm lore) + where freeIn' (Let pat (StmAux cs attrs dec) e) = - freeIn' cs <> freeIn' attrs <> - precomputed dec (freeIn' dec <> freeIn' e <> freeIn' pat) + freeIn' cs <> freeIn' attrs + <> precomputed dec (freeIn' dec <> freeIn' e <> freeIn' pat) instance FreeIn (Stm lore) => FreeIn (Stms lore) where freeIn' = foldMap freeIn' @@ -256,7 +304,7 @@ instance FreeIn SubExp where freeIn' (Var v) = freeIn' v - freeIn' Constant{} = mempty + freeIn' Constant {} = mempty instance FreeIn Space where freeIn' (ScalarSpace d _) = freeIn' d @@ -268,12 +316,12 @@ instance FreeIn d => FreeIn (Ext d) where freeIn' (Free x) = freeIn' x - freeIn' (Ext _) = mempty + freeIn' (Ext _) = mempty instance FreeIn shape => FreeIn (TypeBase shape u) where freeIn' (Array _ shape _) = freeIn' shape - freeIn' (Mem s) = freeIn' s - freeIn' (Prim _) = mempty + freeIn' (Mem s) = freeIn' s + freeIn' (Prim _) = mempty instance FreeIn dec => FreeIn (Param dec) where freeIn' (Param _ dec) = freeIn' dec @@ -294,7 +342,8 @@ instance FreeIn dec => FreeIn (PatternT dec) where freeIn' (Pattern context values) = fvBind bound_here $ freeIn' $ context ++ values - where bound_here = namesFromList $ map patElemName $ context ++ values + where + bound_here = namesFromList $ map patElemName $ context ++ values instance FreeIn Certificates where freeIn' (Certificates cs) = freeIn' cs @@ -315,14 +364,14 @@ precomputed :: dec -> FV -> FV precomputed _ = id -instance FreeDec () where +instance FreeDec () -instance (FreeDec a, FreeIn b) => FreeDec (a,b) where - precomputed (a,_) = precomputed a +instance (FreeDec a, FreeIn b) => FreeDec (a, b) where + precomputed (a, _) = precomputed a instance FreeDec a => FreeDec [a] where precomputed [] = id - precomputed (a:_) = precomputed a + precomputed (a : _) = precomputed a instance FreeDec a => FreeDec (Maybe a) where precomputed Nothing = id diff --git a/src/Futhark/IR/Prop/Patterns.hs b/src/Futhark/IR/Prop/Patterns.hs --- a/src/Futhark/IR/Prop/Patterns.hs +++ b/src/Futhark/IR/Prop/Patterns.hs @@ -1,34 +1,36 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + -- | Inspecing and modifying t'Pattern's, function parameters and -- pattern elements. module Futhark.IR.Prop.Patterns - ( - -- * Function parameters - paramIdent - , paramType - , paramDeclType - -- * Pattern elements - , patElemIdent - , patElemType - , setPatElemLore - , patternElements - , patternIdents - , patternContextIdents - , patternValueIdents - , patternNames - , patternValueNames - , patternContextNames - , patternTypes - , patternValueTypes - , patternSize - -- * Pattern construction - , basicPattern - ) - where + ( -- * Function parameters + paramIdent, + paramType, + paramDeclType, + -- * Pattern elements + patElemIdent, + patElemType, + setPatElemLore, + patternElements, + patternIdents, + patternContextIdents, + patternValueIdents, + patternNames, + patternValueNames, + patternContextNames, + patternTypes, + patternValueTypes, + patternSize, + + -- * Pattern construction + basicPattern, + ) +where + +import Futhark.IR.Prop.Types (DeclTyped (..), Typed (..)) import Futhark.IR.Syntax -import Futhark.IR.Prop.Types (Typed(..), DeclTyped(..)) -- | The 'Type' of a parameter. paramType :: Typed dec => Param dec -> Type @@ -98,4 +100,5 @@ basicPattern :: [Ident] -> [Ident] -> PatternT Type basicPattern context values = Pattern (map patElem context) (map patElem values) - where patElem (Ident name t) = PatElem name t + where + patElem (Ident name t) = PatElem name t diff --git a/src/Futhark/IR/Prop/Rearrange.hs b/src/Futhark/IR/Prop/Rearrange.hs --- a/src/Futhark/IR/Prop/Rearrange.hs +++ b/src/Futhark/IR/Prop/Rearrange.hs @@ -1,40 +1,42 @@ -- | A rearrangement is a generalisation of transposition, where the -- dimensions are arbitrarily permuted. module Futhark.IR.Prop.Rearrange - ( rearrangeShape - , rearrangeInverse - , rearrangeReach - , rearrangeCompose - , isPermutationOf - , transposeIndex - , isMapTranspose - ) where + ( rearrangeShape, + rearrangeInverse, + rearrangeReach, + rearrangeCompose, + isPermutationOf, + transposeIndex, + isMapTranspose, + ) +where import Data.List (sortOn, tails) - import Futhark.Util -- | Calculate the given permutation of the list. It is an error if -- the permutation goes out of bounds. rearrangeShape :: [Int] -> [a] -> [a] rearrangeShape perm l = map pick perm - where pick i - | 0 <= i, i < n = l!!i - | otherwise = - error $ show perm ++ " is not a valid permutation for input." - n = length l + where + pick i + | 0 <= i, i < n = l !! i + | otherwise = + error $ show perm ++ " is not a valid permutation for input." + n = length l -- | Produce the inverse permutation. rearrangeInverse :: [Int] -> [Int] -rearrangeInverse perm = map snd $ sortOn fst $ zip perm [0..] +rearrangeInverse perm = map snd $ sortOn fst $ zip perm [0 ..] -- | Return the first dimension not affected by the permutation. For -- example, the permutation @[1,0,2]@ would return @2@. rearrangeReach :: [Int] -> Int -rearrangeReach perm = case dropWhile (uncurry (/=)) $ zip (tails perm) (tails [0..n-1]) of - [] -> n + 1 - (perm',_):_ -> n - length perm' - where n = length perm +rearrangeReach perm = case dropWhile (uncurry (/=)) $ zip (tails perm) (tails [0 .. n -1]) of + [] -> n + 1 + (perm', _) : _ -> n - length perm' + where + n = length perm -- | Compose two permutations, with the second given permutation being -- applied first. @@ -49,15 +51,16 @@ isPermutationOf l1 l2 = case mapAccumLM (pick 0) (map Just l2) l1 of Just (l2', perm) - | all (==Nothing) l2' -> Just perm - _ -> Nothing - where pick :: Eq a => Int -> [Maybe a] -> a -> Maybe ([Maybe a], Int) - pick _ [] _ = Nothing - pick i (x:xs) y - | Just y == x = Just (Nothing : xs, i) - | otherwise = do - (xs', v) <- pick (i+1) xs y - return (x : xs', v) + | all (== Nothing) l2' -> Just perm + _ -> Nothing + where + pick :: Eq a => Int -> [Maybe a] -> a -> Maybe ([Maybe a], Int) + pick _ [] _ = Nothing + pick i (x : xs) y + | Just y == x = Just (Nothing : xs, i) + | otherwise = do + (xs', v) <- pick (i + 1) xs y + return (x : xs', v) -- | If @l@ is an index into the array @a@, then @transposeIndex k n -- l@ is an index to the same element in the array @transposeArray k n @@ -65,14 +68,14 @@ transposeIndex :: Int -> Int -> [a] -> [a] transposeIndex k n l | k + n >= length l = - let n' = ((k + n) `mod` length l)-k - in transposeIndex k n' l + let n' = ((k + n) `mod` length l) - k + in transposeIndex k n' l | n < 0, - (pre,needle:end) <- splitAt k l, - (beg,mid) <- splitAt (length pre+n) pre = + (pre, needle : end) <- splitAt k l, + (beg, mid) <- splitAt (length pre + n) pre = beg ++ [needle] ++ mid ++ end - | (beg,needle:post) <- splitAt k l, - (mid,end) <- splitAt n post = + | (beg, needle : post) <- splitAt k l, + (mid, end) <- splitAt n post = beg ++ mid ++ [needle] ++ end | otherwise = l @@ -87,22 +90,24 @@ -- undefined. isMapTranspose :: [Int] -> Maybe (Int, Int, Int) isMapTranspose perm - | posttrans == [length mapped..length mapped+length posttrans-1], - not $ null pretrans, not $ null posttrans = - Just (length mapped, length pretrans, length posttrans) + | posttrans == [length mapped .. length mapped + length posttrans -1], + not $ null pretrans, + not $ null posttrans = + Just (length mapped, length pretrans, length posttrans) | otherwise = - Nothing - where (mapped, notmapped) = findIncreasingFrom 0 perm - (pretrans, posttrans) = findTransposed notmapped + Nothing + where + (mapped, notmapped) = findIncreasingFrom 0 perm + (pretrans, posttrans) = findTransposed notmapped - findIncreasingFrom x (i:is) - | i == x = - let (js, ps) = findIncreasingFrom (x+1) is - in (i : js, ps) - findIncreasingFrom _ is = - ([], is) + findIncreasingFrom x (i : is) + | i == x = + let (js, ps) = findIncreasingFrom (x + 1) is + in (i : js, ps) + findIncreasingFrom _ is = + ([], is) - findTransposed [] = - ([], []) - findTransposed (i:is) = - findIncreasingFrom i (i:is) + findTransposed [] = + ([], []) + findTransposed (i : is) = + findIncreasingFrom i (i : is) diff --git a/src/Futhark/IR/Prop/Reshape.hs b/src/Futhark/IR/Prop/Reshape.hs --- a/src/Futhark/IR/Prop/Reshape.hs +++ b/src/Futhark/IR/Prop/Reshape.hs @@ -1,45 +1,42 @@ -- | Facilities for creating, inspecting, and simplifying reshape and -- coercion operations. module Futhark.IR.Prop.Reshape - ( - -- * Basic tools - newDim - , newDims - , newShape + ( -- * Basic tools + newDim, + newDims, + newShape, - -- * Construction - , shapeCoerce + -- * Construction + shapeCoerce, - -- * Execution - , reshapeOuter - , reshapeInner + -- * Execution + reshapeOuter, + reshapeInner, - -- * Inspection - , shapeCoercion + -- * Inspection + shapeCoercion, - -- * Simplification - , fuseReshape - , informReshape + -- * Simplification + fuseReshape, + informReshape, - -- * Shape calculations - , reshapeIndex - , flattenIndex - , unflattenIndex - , sliceSizes - ) - where + -- * Shape calculations + reshapeIndex, + flattenIndex, + unflattenIndex, + sliceSizes, + ) +where import Data.Foldable - -import Prelude hiding (sum, product, quot) - import Futhark.IR.Syntax import Futhark.Util.IntegralExp +import Prelude hiding (product, quot, sum) -- | The new dimension. newDim :: DimChange d -> d newDim (DimCoercion se) = se -newDim (DimNew se) = se +newDim (DimNew se) = se -- | The new dimensions resulting from a reshape operation. newDims :: ShapeChange d -> [d] @@ -60,9 +57,10 @@ reshapeOuter :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp reshapeOuter newshape n oldshape = newshape ++ map coercion_or_new (drop n (shapeDims oldshape)) - where coercion_or_new - | length newshape == n = DimCoercion - | otherwise = DimNew + where + coercion_or_new + | length newshape == n = DimCoercion + | otherwise = DimNew -- | @reshapeInner newshape n oldshape@ returns a 'Reshape' expression -- that replaces the inner @m-n@ dimensions (where @m@ is the rank of @@ -70,17 +68,19 @@ reshapeInner :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp reshapeInner newshape n oldshape = map coercion_or_new (take n (shapeDims oldshape)) ++ newshape - where coercion_or_new - | length newshape == m-n = DimCoercion - | otherwise = DimNew - m = shapeRank oldshape + where + coercion_or_new + | length newshape == m - n = DimCoercion + | otherwise = DimNew + m = shapeRank oldshape -- | If the shape change is nothing but shape coercions, return the new dimensions. Otherwise, return -- 'Nothing'. shapeCoercion :: ShapeChange d -> Maybe [d] shapeCoercion = mapM dimCoercion - where dimCoercion (DimCoercion d) = Just d - dimCoercion (DimNew _) = Nothing + where + dimCoercion (DimCoercion d) = Just d + dimCoercion (DimNew _) = Nothing -- | @fuseReshape s1 s2@ creates a new 'ShapeChange' that is -- semantically the same as first applying @s1@ and then @s2@. This @@ -89,14 +89,15 @@ fuseReshape :: Eq d => ShapeChange d -> ShapeChange d -> ShapeChange d fuseReshape s1 s2 | length s1 == length s2 = - zipWith comb s1 s2 - where comb (DimNew _) (DimCoercion d2) = - DimNew d2 - comb (DimCoercion d1) (DimNew d2) - | d1 == d2 = DimCoercion d2 - | otherwise = DimNew d2 - comb _ d2 = - d2 + zipWith comb s1 s2 + where + comb (DimNew _) (DimCoercion d2) = + DimNew d2 + comb (DimCoercion d1) (DimNew d2) + | d1 == d2 = DimCoercion d2 + | otherwise = DimNew d2 + comb _ d2 = + d2 -- TODO: intelligently handle case where s1 is a prefix of s2. fuseReshape _ s2 = s2 @@ -106,10 +107,11 @@ informReshape shape sc | length shape == length sc = zipWith inform shape sc - where inform d1 (DimNew d2) - | d1 == d2 = DimCoercion d2 - inform _ dc = - dc + where + inform d1 (DimNew d2) + | d1 == d2 = DimCoercion d2 + inform _ dc = + dc informReshape _ sc = sc -- | @reshapeIndex to_dims from_dims is@ transforms the index list @@ -117,20 +119,30 @@ -- list @is'@, which is into an array of shape @to_dims@. @is@ must -- have the same length as @from_dims@, and @is'@ will have the same -- length as @to_dims@. -reshapeIndex :: IntegralExp num => - [num] -> [num] -> [num] -> [num] +reshapeIndex :: + IntegralExp num => + [num] -> + [num] -> + [num] -> + [num] reshapeIndex to_dims from_dims is = unflattenIndex to_dims $ flattenIndex from_dims is -- | @unflattenIndex dims i@ computes a list of indices into an array -- with dimension @dims@ given the flat index @i@. The resulting list -- will have the same size as @dims@. -unflattenIndex :: IntegralExp num => - [num] -> num -> [num] +unflattenIndex :: + IntegralExp num => + [num] -> + num -> + [num] unflattenIndex = unflattenIndexFromSlices . drop 1 . sliceSizes -unflattenIndexFromSlices :: IntegralExp num => - [num] -> num -> [num] +unflattenIndexFromSlices :: + IntegralExp num => + [num] -> + num -> + [num] unflattenIndexFromSlices [] _ = [] unflattenIndexFromSlices (size : slices) i = (i `quot` size) : unflattenIndexFromSlices slices (i - (i `quot` size) * size) @@ -138,20 +150,26 @@ -- | @flattenIndex dims is@ computes the flat index of @is@ into an -- array with dimensions @dims@. The length of @dims@ and @is@ must -- be the same. -flattenIndex :: IntegralExp num => - [num] -> [num] -> num +flattenIndex :: + IntegralExp num => + [num] -> + [num] -> + num flattenIndex dims is = sum $ zipWith (*) is slicesizes - where slicesizes = drop 1 $ sliceSizes dims + where + slicesizes = drop 1 $ sliceSizes dims -- | Given a length @n@ list of dimensions @dims@, @sizeSizes dims@ -- will compute a length @n+1@ list of the size of each possible array -- slice. The first element of this list will be the product of -- @dims@, and the last element will be 1. -sliceSizes :: IntegralExp num => - [num] -> [num] +sliceSizes :: + IntegralExp num => + [num] -> + [num] sliceSizes [] = [1] -sliceSizes (n:ns) = +sliceSizes (n : ns) = product (n : ns) : sliceSizes ns {- HLINT ignore sliceSizes -} diff --git a/src/Futhark/IR/Prop/Scope.hs b/src/Futhark/IR/Prop/Scope.hs --- a/src/Futhark/IR/Prop/Scope.hs +++ b/src/Futhark/IR/Prop/Scope.hs @@ -1,12 +1,13 @@ -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + -- | The core Futhark AST does not contain type information when we -- use a variable. Therefore, most transformations expect to be able -- to access some kind of symbol table that maps names to their types. @@ -16,42 +17,42 @@ -- also provided to communicate that some monad or applicative functor -- maintains type information. module Futhark.IR.Prop.Scope - ( HasScope (..) - , NameInfo (..) - , LocalScope (..) - , Scope - , Scoped(..) - , inScopeOf - , scopeOfLParams - , scopeOfFParams - , scopeOfPattern - , scopeOfPatElem - - , SameScope - , castScope + ( HasScope (..), + NameInfo (..), + LocalScope (..), + Scope, + Scoped (..), + inScopeOf, + scopeOfLParams, + scopeOfFParams, + scopeOfPattern, + scopeOfPatElem, + SameScope, + castScope, - -- * Extended type environment - , ExtendedScope - , extendedScope - ) where + -- * Extended type environment + ExtendedScope, + extendedScope, + ) +where import Control.Monad.Except -import Control.Monad.Reader -import qualified Control.Monad.RWS.Strict import qualified Control.Monad.RWS.Lazy +import qualified Control.Monad.RWS.Strict +import Control.Monad.Reader import qualified Data.Map.Strict as M - import Futhark.IR.Decorations -import Futhark.IR.Syntax -import Futhark.IR.Prop.Types -import Futhark.IR.Prop.Patterns import Futhark.IR.Pretty () +import Futhark.IR.Prop.Patterns +import Futhark.IR.Prop.Types +import Futhark.IR.Syntax -- | How some name in scope was bound. -data NameInfo lore = LetName (LetDec lore) - | FParamName (FParamInfo lore) - | LParamName (LParamInfo lore) - | IndexName IntType +data NameInfo lore + = LetName (LetDec lore) + | FParamName (FParamInfo lore) + | LParamName (LParamInfo lore) + | IndexName IntType deriving instance Decorations lore => Show (NameInfo lore) @@ -81,9 +82,11 @@ lookupInfo :: VName -> m (NameInfo lore) lookupInfo name = asksScope (M.findWithDefault notFound name) - where notFound = - error $ "Scope.lookupInfo: Name " ++ pretty name ++ - " not found in type environment." + where + notFound = + error $ + "Scope.lookupInfo: Name " ++ pretty name + ++ " not found in type environment." -- | Return the type environment contained in the applicative -- functor. @@ -94,19 +97,25 @@ asksScope :: (Scope lore -> a) -> m a asksScope f = f <$> askScope -instance (Applicative m, Monad m, Decorations lore) => - HasScope lore (ReaderT (Scope lore) m) where +instance + (Applicative m, Monad m, Decorations lore) => + HasScope lore (ReaderT (Scope lore) m) + where askScope = ask instance (Monad m, HasScope lore m) => HasScope lore (ExceptT e m) where askScope = lift askScope -instance (Applicative m, Monad m, Monoid w, Decorations lore) => - HasScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m) where +instance + (Applicative m, Monad m, Monoid w, Decorations lore) => + HasScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m) + where askScope = ask -instance (Applicative m, Monad m, Monoid w, Decorations lore) => - HasScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) w s m) where +instance + (Applicative m, Monad m, Monoid w, Decorations lore) => + HasScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) w s m) + where askScope = ask -- | The class of monads that not only provide a 'Scope', but also @@ -121,16 +130,22 @@ instance (Monad m, LocalScope lore m) => LocalScope lore (ExceptT e m) where localScope = mapExceptT . localScope -instance (Applicative m, Monad m, Decorations lore) => - LocalScope lore (ReaderT (Scope lore) m) where +instance + (Applicative m, Monad m, Decorations lore) => + LocalScope lore (ReaderT (Scope lore) m) + where localScope = local . M.union -instance (Applicative m, Monad m, Monoid w, Decorations lore) => - LocalScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m) where +instance + (Applicative m, Monad m, Monoid w, Decorations lore) => + LocalScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m) + where localScope = local . M.union -instance (Applicative m, Monad m, Monoid w, Decorations lore) => - LocalScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) w s m) where +instance + (Applicative m, Monad m, Monoid w, Decorations lore) => + LocalScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) w s m) + where localScope = local . M.union -- | The class of things that can provide a scope. There is no @@ -173,34 +188,46 @@ scopeOfPatElem (PatElem name dec) = M.singleton name $ LetName dec -- | The scope of some lambda parameters. -scopeOfLParams :: LParamInfo lore ~ dec => - [Param dec] -> Scope lore +scopeOfLParams :: + LParamInfo lore ~ dec => + [Param dec] -> + Scope lore scopeOfLParams = M.fromList . map f - where f param = (paramName param, LParamName $ paramDec param) + where + f param = (paramName param, LParamName $ paramDec param) -- | The scope of some function or loop parameters. -scopeOfFParams :: FParamInfo lore ~ dec => - [Param dec] -> Scope lore +scopeOfFParams :: + FParamInfo lore ~ dec => + [Param dec] -> + Scope lore scopeOfFParams = M.fromList . map f - where f param = (paramName param, FParamName $ paramDec param) + where + f param = (paramName param, FParamName $ paramDec param) instance Scoped lore (Lambda lore) where scopeOf lam = scopeOfLParams $ lambdaParams lam -- | A constraint that indicates two lores have the same 'NameInfo' -- representation. -type SameScope lore1 lore2 = (LetDec lore1 ~ LetDec lore2, - FParamInfo lore1 ~ FParamInfo lore2, - LParamInfo lore1 ~ LParamInfo lore2) +type SameScope lore1 lore2 = + ( LetDec lore1 ~ LetDec lore2, + FParamInfo lore1 ~ FParamInfo lore2, + LParamInfo lore1 ~ LParamInfo lore2 + ) -- | If two scopes are really the same, then you can convert one to -- the other. -castScope :: SameScope fromlore tolore => - Scope fromlore -> Scope tolore +castScope :: + SameScope fromlore tolore => + Scope fromlore -> + Scope tolore castScope = M.map castNameInfo -castNameInfo :: SameScope fromlore tolore => - NameInfo fromlore -> NameInfo tolore +castNameInfo :: + SameScope fromlore tolore => + NameInfo fromlore -> + NameInfo tolore castNameInfo (LetName dec) = LetName dec castNameInfo (FParamName dec) = FParamName dec castNameInfo (LParamName dec) = LParamName dec @@ -210,18 +237,25 @@ -- Its 'lookupType' method will first look in the extended 'Scope', -- and then use the 'lookupType' method of the underlying monad. newtype ExtendedScope lore m a = ExtendedScope (ReaderT (Scope lore) m a) - deriving (Functor, Applicative, Monad, - MonadReader (Scope lore)) + deriving + ( Functor, + Applicative, + Monad, + MonadReader (Scope lore) + ) -instance (HasScope lore m, Monad m) => - HasScope lore (ExtendedScope lore m) where +instance + (HasScope lore m, Monad m) => + HasScope lore (ExtendedScope lore m) + where lookupType name = do res <- asks $ fmap typeOf . M.lookup name maybe (ExtendedScope $ lift $ lookupType name) return res askScope = asks M.union <*> ExtendedScope (lift askScope) -- | Run a computation in the extended type environment. -extendedScope :: ExtendedScope lore m a - -> Scope lore - -> m a +extendedScope :: + ExtendedScope lore m a -> + Scope lore -> + m a extendedScope (ExtendedScope m) = runReaderT m diff --git a/src/Futhark/IR/Prop/TypeOf.hs b/src/Futhark/IR/Prop/TypeOf.hs --- a/src/Futhark/IR/Prop/TypeOf.hs +++ b/src/Futhark/IR/Prop/TypeOf.hs @@ -2,6 +2,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeFamilies #-} + -- | This module provides facilities for obtaining the types of -- various Futhark constructs. Typically, you will need to execute -- these in a context where type information is available as a @@ -16,44 +17,45 @@ -- "Futhark.IR.Mem" exposes functionality for -- also obtaining information about the storage location of results. module Futhark.IR.Prop.TypeOf - ( - expExtType - , expExtTypeSize - , subExpType - , primOpType - , mapType + ( expExtType, + expExtTypeSize, + subExpType, + primOpType, + mapType, - -- * Return type - , module Futhark.IR.RetType - -- * Type environment - , module Futhark.IR.Prop.Scope + -- * Return type + module Futhark.IR.RetType, - -- * Extensibility - , TypedOp(..) - ) - where + -- * Type environment + module Futhark.IR.Prop.Scope, -import Data.Maybe + -- * Extensibility + TypedOp (..), + ) +where -import Futhark.IR.Syntax +import Data.Maybe +import Futhark.IR.Prop.Constants +import Futhark.IR.Prop.Patterns import Futhark.IR.Prop.Reshape +import Futhark.IR.Prop.Scope import Futhark.IR.Prop.Types -import Futhark.IR.Prop.Patterns -import Futhark.IR.Prop.Constants import Futhark.IR.RetType -import Futhark.IR.Prop.Scope +import Futhark.IR.Syntax -- | The type of a subexpression. subExpType :: HasScope t m => SubExp -> m Type subExpType (Constant val) = pure $ Prim $ primValueType val -subExpType (Var name) = lookupType name +subExpType (Var name) = lookupType name -- | @mapType f arrts@ wraps each element in the return type of @f@ in -- an array with size equal to the outermost dimension of the first -- element of @arrts@. mapType :: SubExp -> Lambda lore -> [Type] -mapType outersize f = [ arrayOf t (Shape [outersize]) NoUniqueness - | t <- lambdaReturnType f ] +mapType outersize f = + [ arrayOf t (Shape [outersize]) NoUniqueness + | t <- lambdaReturnType f + ] -- | The type of a primitive operation. primOpType :: HasScope lore m => BasicOp -> m [Type] @@ -63,21 +65,23 @@ pure <$> subExpType se primOpType (ArrayLit es rt) = pure [arrayOf rt (Shape [n]) NoUniqueness] - where n = Constant (value (length es)) + where + n = intConst Int64 $ toInteger $ length es primOpType (BinOp bop _ _) = pure [Prim $ binOpType bop] primOpType (UnOp _ x) = pure <$> subExpType x -primOpType CmpOp{} = +primOpType CmpOp {} = pure [Prim Bool] primOpType (ConvOp conv _) = pure [Prim $ snd $ convOpType conv] primOpType (Index ident slice) = result <$> lookupType ident - where result t = [Prim (elemType t) `arrayOfShape` shape] - shape = Shape $ mapMaybe dimSize slice - dimSize (DimSlice _ d _) = Just d - dimSize DimFix{} = Nothing + where + result t = [Prim (elemType t) `arrayOfShape` shape] + shape = Shape $ mapMaybe dimSize slice + dimSize (DimSlice _ d _) = Just d + dimSize DimFix {} = Nothing primOpType (Update src _ _) = pure <$> lookupType src primOpType (Iota n _ _ et) = @@ -90,42 +94,50 @@ pure [arrayOf (Prim t) (Shape shape) NoUniqueness] primOpType (Reshape [] e) = result <$> lookupType e - where result t = [Prim $ elemType t] + where + result t = [Prim $ elemType t] primOpType (Reshape shape e) = result <$> lookupType e - where result t = [t `setArrayShape` newShape shape] + where + result t = [t `setArrayShape` newShape shape] primOpType (Rearrange perm e) = result <$> lookupType e - where result t = [rearrangeType perm t] + where + result t = [rearrangeType perm t] primOpType (Rotate _ e) = pure <$> lookupType e primOpType (Concat i x _ ressize) = result <$> lookupType x - where result xt = [setDimSize i xt ressize] + where + result xt = [setDimSize i xt ressize] primOpType (Copy v) = pure <$> lookupType v primOpType (Manifest _ v) = pure <$> lookupType v -primOpType Assert{} = +primOpType Assert {} = pure [Prim Cert] -- | The type of an expression. -expExtType :: (HasScope lore m, TypedOp (Op lore)) => - Exp lore -> m [ExtType] +expExtType :: + (HasScope lore m, TypedOp (Op lore)) => + Exp lore -> + m [ExtType] expExtType (Apply _ _ rt _) = pure $ map (fromDecl . declExtTypeOf) rt -expExtType (If _ _ _ rt) = pure $ map extTypeOf $ ifReturns rt +expExtType (If _ _ _ rt) = pure $ map extTypeOf $ ifReturns rt expExtType (DoLoop ctxmerge valmerge _ _) = pure $ loopExtType (map (paramIdent . fst) ctxmerge) (map (paramIdent . fst) valmerge) -expExtType (BasicOp op) = staticShapes <$> primOpType op -expExtType (Op op) = opType op +expExtType (BasicOp op) = staticShapes <$> primOpType op +expExtType (Op op) = opType op -- | The number of values returned by an expression. -expExtTypeSize :: (Decorations lore, TypedOp (Op lore)) => - Exp lore -> Int +expExtTypeSize :: + (Decorations lore, TypedOp (Op lore)) => + Exp lore -> + Int expExtTypeSize = length . feelBad . expExtType -- FIXME, this is a horrible quick hack. -newtype FeelBad lore a = FeelBad { feelBad :: a } +newtype FeelBad lore a = FeelBad {feelBad :: a} instance Functor (FeelBad lore) where fmap f = FeelBad . f . feelBad @@ -135,7 +147,7 @@ f <*> x = FeelBad $ feelBad f $ feelBad x instance Decorations lore => HasScope lore (FeelBad lore) where - lookupType = const $ pure $ Prim $ IntType Int32 + lookupType = const $ pure $ Prim $ IntType Int64 askScope = pure mempty -- | Given the context and value merge parameters of a Futhark @loop@, @@ -143,7 +155,8 @@ loopExtType :: [Ident] -> [Ident] -> [ExtType] loopExtType ctx val = existentialiseExtTypes inaccessible $ staticShapes $ map identType val - where inaccessible = map identName ctx + where + inaccessible = map identName ctx -- | Any operation must define an instance of this class, which -- describes the type of the operation (at the value level). diff --git a/src/Futhark/IR/Prop/Types.hs b/src/Futhark/IR/Prop/Types.hs --- a/src/Futhark/IR/Prop/Types.hs +++ b/src/Futhark/IR/Prop/Types.hs @@ -1,82 +1,79 @@ -{-# LANGUAGE FlexibleContexts, FlexibleInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} + -- | Functions for inspecting and constructing various types. module Futhark.IR.Prop.Types - ( - rankShaped - , arrayRank - , arrayShape - , setArrayShape - , existential - , uniqueness - , unique - , staticShapes - , staticShapes1 - , primType - - , arrayOf - , arrayOfRow - , arrayOfShape - , setOuterSize - , setDimSize - , setOuterDim - , setDim - , setArrayDims - , peelArray - , stripArray - , arrayDims - , arrayExtDims - , shapeSize - , arraySize - , arraysSize - , rowType - , elemType - - , transposeType - , rearrangeType - - , mapOnExtType - , mapOnType - - , diet - - , subtypeOf - , subtypesOf - - , toDecl - , fromDecl - - , isExt - , isFree - , extractShapeContext - , shapeContext - , hasStaticShape - , generaliseExtTypes - , existentialiseExtTypes - , shapeExtMapping + ( rankShaped, + arrayRank, + arrayShape, + setArrayShape, + existential, + uniqueness, + unique, + staticShapes, + staticShapes1, + primType, + arrayOf, + arrayOfRow, + arrayOfShape, + setOuterSize, + setDimSize, + setOuterDim, + setDim, + setArrayDims, + peelArray, + stripArray, + arrayDims, + arrayExtDims, + shapeSize, + arraySize, + arraysSize, + rowType, + elemType, + transposeType, + rearrangeType, + mapOnExtType, + mapOnType, + diet, + subtypeOf, + subtypesOf, + toDecl, + fromDecl, + isExt, + isFree, + extractShapeContext, + shapeContext, + hasStaticShape, + generaliseExtTypes, + existentialiseExtTypes, + shapeExtMapping, - -- * Abbreviations - , int8, int16, int32, int64 - , float32, float64 + -- * Abbreviations + int8, + int16, + int32, + int64, + float32, + float64, - -- * The Typed typeclass - , Typed (..) - , DeclTyped (..) - , ExtTyped (..) - , DeclExtTyped (..) - , SetType (..) - , FixExt (..) - ) - where + -- * The Typed typeclass + Typed (..), + DeclTyped (..), + ExtTyped (..), + DeclExtTyped (..), + SetType (..), + FixExt (..), + ) +where import Control.Monad.State -import Data.Maybe import Data.List (elemIndex, foldl') -import qualified Data.Set as S import qualified Data.Map.Strict as M - -import Futhark.IR.Syntax.Core +import Data.Maybe +import qualified Data.Set as S import Futhark.IR.Prop.Constants import Futhark.IR.Prop.Rearrange +import Futhark.IR.Syntax.Core -- | Remove shape information from a type. rankShaped :: ArrayShape shape => TypeBase shape u -> TypeBase Rank u @@ -94,33 +91,37 @@ -- 'mempty'. arrayShape :: ArrayShape shape => TypeBase shape u -> shape arrayShape (Array _ ds _) = ds -arrayShape _ = mempty +arrayShape _ = mempty -- | Modify the shape of an array - for non-arrays, this does nothing. -modifyArrayShape :: ArrayShape newshape => - (oldshape -> newshape) - -> TypeBase oldshape u - -> TypeBase newshape u +modifyArrayShape :: + ArrayShape newshape => + (oldshape -> newshape) -> + TypeBase oldshape u -> + TypeBase newshape u modifyArrayShape f (Array t ds u) | shapeRank ds' == 0 = Prim t - | otherwise = Array t (f ds) u - where ds' = f ds -modifyArrayShape _ (Prim t) = Prim t + | otherwise = Array t (f ds) u + where + ds' = f ds +modifyArrayShape _ (Prim t) = Prim t modifyArrayShape _ (Mem space) = Mem space -- | Set the shape of an array. If the given type is not an -- array, return the type unchanged. -setArrayShape :: ArrayShape newshape => - TypeBase oldshape u - -> newshape - -> TypeBase newshape u +setArrayShape :: + ArrayShape newshape => + TypeBase oldshape u -> + newshape -> + TypeBase newshape u setArrayShape t ds = modifyArrayShape (const ds) t -- | True if the given type has a dimension that is existentially sized. existential :: ExtType -> Bool existential = any ext . shapeDims . arrayShape - where ext (Ext _) = True - ext (Free _) = False + where + ext (Ext _) = True + ext (Free _) = False -- | Return the uniqueness of a type. uniqueness :: TypeBase shape Uniqueness -> Uniqueness @@ -129,7 +130,7 @@ -- | @unique t@ is 'True' if the type of the argument is unique. unique :: TypeBase shape Uniqueness -> Bool -unique = (==Unique) . uniqueness +unique = (== Unique) . uniqueness -- | Convert types with non-existential shapes to types with -- non-existential shapes. Only the representation is changed, so all @@ -154,24 +155,29 @@ -- uniqueness of @t@. If the shape @s@ has rank 0, then the @t@ will -- be returned, although if it is an array, with the uniqueness -- changed to @u@. -arrayOf :: ArrayShape shape => - TypeBase shape u_unused -> shape -> u -> TypeBase shape u +arrayOf :: + ArrayShape shape => + TypeBase shape u_unused -> + shape -> + u -> + TypeBase shape u arrayOf (Array et size1 _) size2 u = Array et (size2 <> size1) u arrayOf (Prim et) s _ | 0 <- shapeRank s = Prim et arrayOf (Prim et) size u = Array et size u -arrayOf Mem{} _ _ = +arrayOf Mem {} _ _ = error "arrayOf Mem" -- | Construct an array whose rows are the given type, and the outer -- size is the given dimension. This is just a convenient wrapper -- around 'arrayOf'. -arrayOfRow :: ArrayShape (ShapeBase d) => - TypeBase (ShapeBase d) NoUniqueness - -> d - -> TypeBase (ShapeBase d) NoUniqueness +arrayOfRow :: + ArrayShape (ShapeBase d) => + TypeBase (ShapeBase d) NoUniqueness -> + d -> + TypeBase (ShapeBase d) NoUniqueness arrayOfRow t size = arrayOf t (Shape [size]) NoUniqueness -- | Construct an array whose rows are the given type, and the outer @@ -187,14 +193,21 @@ -- | Replace the size of the outermost dimension of an array. If the -- given type is not an array, it is returned unchanged. -setOuterSize :: ArrayShape (ShapeBase d) => - TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u +setOuterSize :: + ArrayShape (ShapeBase d) => + TypeBase (ShapeBase d) u -> + d -> + TypeBase (ShapeBase d) u setOuterSize = setDimSize 0 -- | Replace the size of the given dimension of an array. If the -- given type is not an array, it is returned unchanged. -setDimSize :: ArrayShape (ShapeBase d) => - Int -> TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u +setDimSize :: + ArrayShape (ShapeBase d) => + Int -> + TypeBase (ShapeBase d) u -> + d -> + TypeBase (ShapeBase d) u setDimSize i t e = t `setArrayShape` setDim i (arrayShape t) e -- | Replace the outermost dimension of an array shape. @@ -203,17 +216,20 @@ -- | Replace the specified dimension of an array shape. setDim :: Int -> ShapeBase d -> d -> ShapeBase d -setDim i (Shape ds) e = Shape $ take i ds ++ e : drop (i+1) ds +setDim i (Shape ds) e = Shape $ take i ds ++ e : drop (i + 1) ds -- | @peelArray n t@ returns the type resulting from peeling the first -- @n@ array dimensions from @t@. Returns @Nothing@ if @t@ has less -- than @n@ dimensions. -peelArray :: ArrayShape shape => - Int -> TypeBase shape u -> Maybe (TypeBase shape u) +peelArray :: + ArrayShape shape => + Int -> + TypeBase shape u -> + Maybe (TypeBase shape u) peelArray 0 t = Just t peelArray n (Array et shape u) | shapeRank shape == n = Just $ Prim et - | shapeRank shape > n = Just $ Array et (stripDims n shape) u + | shapeRank shape > n = Just $ Array et (stripDims n shape) u peelArray _ _ = Nothing -- | @stripArray n t@ removes the @n@ outermost layers of the array. @@ -222,7 +238,7 @@ stripArray :: ArrayShape shape => Int -> TypeBase shape u -> TypeBase shape u stripArray n (Array et shape u) | n < shapeRank shape = Array et (stripDims n shape) u - | otherwise = Prim et + | otherwise = Prim et stripArray _ t = t -- | Return the size of the given dimension. If the dimension does @@ -230,7 +246,7 @@ shapeSize :: Int -> Shape -> SubExp shapeSize i shape = case drop i $ shapeDims shape of e : _ -> e - [] -> constant (0 :: Int32) + [] -> constant (0 :: Int64) -- | Return the dimensions of a type - for non-arrays, this is the -- empty list. @@ -251,8 +267,8 @@ -- the given type list. If the dimension does not exist, or no types -- are given, the zero constant is returned. arraysSize :: Int -> [TypeBase Shape u] -> SubExp -arraysSize _ [] = constant (0 :: Int32) -arraysSize i (t:_) = arraySize i t +arraysSize _ [] = constant (0 :: Int64) +arraysSize i (t : _) = arraySize i t -- | Return the immediate row-type of an array. For @[[int]]@, this -- would be @[int]@. @@ -261,20 +277,20 @@ -- | A type is a primitive type if it is not an array or memory block. primType :: TypeBase shape u -> Bool -primType Array{} = False -primType Mem{} = False +primType Array {} = False +primType Mem {} = False primType _ = True -- | Returns the bottommost type of an array. For @[[int]]@, this -- would be @int@. If the given type is not an array, it is returned. elemType :: TypeBase shape u -> PrimType elemType (Array t _ _) = t -elemType (Prim t) = t -elemType Mem{} = error "elemType Mem" +elemType (Prim t) = t +elemType Mem {} = error "elemType Mem" -- | Swap the two outer dimensions of the type. transposeType :: Type -> Type -transposeType = rearrangeType [1,0] +transposeType = rearrangeType [1, 0] -- | Rearrange the dimensions of the type. If the length of the -- permutation does not match the rank of the type, the permutation @@ -282,13 +298,15 @@ rearrangeType :: [Int] -> Type -> Type rearrangeType perm t = t `setArrayShape` Shape (rearrangeShape perm' $ arrayDims t) - where perm' = perm ++ [length perm .. arrayRank t - 1] + where + perm' = perm ++ [length perm .. arrayRank t - 1] -- | Transform any t'SubExp's in the type. -mapOnExtType :: Monad m => - (SubExp -> m SubExp) - -> TypeBase ExtShape u - -> m (TypeBase ExtShape u) +mapOnExtType :: + Monad m => + (SubExp -> m SubExp) -> + TypeBase ExtShape u -> + m (TypeBase ExtShape u) mapOnExtType _ (Prim bt) = return $ Prim bt mapOnExtType _ (Mem space) = @@ -297,10 +315,11 @@ Array t <$> (Shape <$> mapM (traverse f) (shapeDims shape)) <*> pure u -- | Transform any t'SubExp's in the type. -mapOnType :: Monad m => - (SubExp -> m SubExp) - -> TypeBase Shape u - -> m (TypeBase Shape u) +mapOnType :: + Monad m => + (SubExp -> m SubExp) -> + TypeBase Shape u -> + m (TypeBase Shape u) mapOnType _ (Prim bt) = return $ Prim bt mapOnType _ (Mem space) = pure $ Mem space mapOnType f (Array t shape u) = @@ -312,18 +331,19 @@ diet (Prim _) = ObservePrim diet (Array _ _ Unique) = Consume diet (Array _ _ Nonunique) = Observe -diet Mem{} = Observe +diet Mem {} = Observe -- | @x \`subtypeOf\` y@ is true if @x@ is a subtype of @y@ (or equal to -- @y@), meaning @x@ is valid whenever @y@ is. -subtypeOf :: (Ord u, ArrayShape shape) => - TypeBase shape u - -> TypeBase shape u - -> Bool +subtypeOf :: + (Ord u, ArrayShape shape) => + TypeBase shape u -> + TypeBase shape u -> + Bool subtypeOf (Array t1 shape1 u1) (Array t2 shape2 u2) = - u2 <= u1 && - t1 == t2 && - shape1 `subShapeOf` shape2 + u2 <= u1 + && t1 == t2 + && shape1 `subShapeOf` shape2 subtypeOf (Prim t1) (Prim t2) = t1 == t2 subtypeOf (Mem space1) (Mem space2) = space1 == space2 subtypeOf _ _ = False @@ -331,24 +351,28 @@ -- | @xs \`subtypesOf\` ys@ is true if @xs@ is the same size as @ys@, -- and each element in @xs@ is a subtype of the corresponding element -- in @ys@.. -subtypesOf :: (Ord u, ArrayShape shape) => - [TypeBase shape u] - -> [TypeBase shape u] - -> Bool -subtypesOf xs ys = length xs == length ys && - and (zipWith subtypeOf xs ys) +subtypesOf :: + (Ord u, ArrayShape shape) => + [TypeBase shape u] -> + [TypeBase shape u] -> + Bool +subtypesOf xs ys = + length xs == length ys + && and (zipWith subtypeOf xs ys) -- | Add the given uniqueness information to the types. -toDecl :: TypeBase shape NoUniqueness - -> Uniqueness - -> TypeBase shape Uniqueness +toDecl :: + TypeBase shape NoUniqueness -> + Uniqueness -> + TypeBase shape Uniqueness toDecl (Prim bt) _ = Prim bt toDecl (Array et shape _) u = Array et shape u toDecl (Mem space) _ = Mem space -- | Remove uniqueness information from the type. -fromDecl :: TypeBase shape Uniqueness - -> TypeBase shape NoUniqueness +fromDecl :: + TypeBase shape Uniqueness -> + TypeBase shape NoUniqueness fromDecl (Prim bt) = Prim bt fromDecl (Array et shape _) = Array et shape NoUniqueness fromDecl (Mem space) = Mem space @@ -370,22 +394,27 @@ extractShapeContext :: [TypeBase ExtShape u] -> [[a]] -> [a] extractShapeContext ts shapes = evalState (concat <$> zipWithM extract ts shapes) S.empty - where extract t shape = - catMaybes <$> zipWithM extract' (shapeDims $ arrayShape t) shape - extract' (Ext x) v = do - seen <- gets $ S.member x - if seen then return Nothing - else do modify $ S.insert x - return $ Just v - extract' (Free _) _ = return Nothing + where + extract t shape = + catMaybes <$> zipWithM extract' (shapeDims $ arrayShape t) shape + extract' (Ext x) v = do + seen <- gets $ S.member x + if seen + then return Nothing + else do + modify $ S.insert x + return $ Just v + extract' (Free _) _ = return Nothing -- | The set of identifiers used for the shape context in the given -- 'ExtType's. shapeContext :: [TypeBase ExtShape u] -> S.Set Int -shapeContext = S.fromList - . concatMap (mapMaybe ext . shapeDims . arrayShape) - where ext (Ext x) = Just x - ext (Free _) = Nothing +shapeContext = + S.fromList + . concatMap (mapMaybe ext . shapeDims . arrayShape) + where + ext (Ext x) = Just x + ext (Free _) = Nothing -- | If all dimensions of the given 'ExtShape' are statically known, -- change to the corresponding 'Shape'. @@ -397,29 +426,37 @@ -- | Given two lists of 'ExtType's of the same length, return a list -- of 'ExtType's that is a subtype of the two operands. -generaliseExtTypes :: [TypeBase ExtShape u] - -> [TypeBase ExtShape u] - -> [TypeBase ExtShape u] +generaliseExtTypes :: + [TypeBase ExtShape u] -> + [TypeBase ExtShape u] -> + [TypeBase ExtShape u] generaliseExtTypes rt1 rt2 = evalState (zipWithM unifyExtShapes rt1 rt2) (0, M.empty) - where unifyExtShapes t1 t2 = - setArrayShape t1 . Shape <$> - zipWithM unifyExtDims + where + unifyExtShapes t1 t2 = + setArrayShape t1 . Shape + <$> zipWithM + unifyExtDims (shapeDims $ arrayShape t1) (shapeDims $ arrayShape t2) - unifyExtDims (Free se1) (Free se2) - | se1 == se2 = return $ Free se1 -- Arbitrary - | otherwise = do (n,m) <- get - put (n + 1, m) - return $ Ext n - unifyExtDims (Ext x) (Ext y) - | x == y = Ext <$> (maybe (new x) return =<< - gets (M.lookup x . snd)) - unifyExtDims (Ext x) _ = Ext <$> new x - unifyExtDims _ (Ext x) = Ext <$> new x - new x = do (n,m) <- get - put (n + 1, M.insert x n m) - return n + unifyExtDims (Free se1) (Free se2) + | se1 == se2 = return $ Free se1 -- Arbitrary + | otherwise = do + (n, m) <- get + put (n + 1, m) + return $ Ext n + unifyExtDims (Ext x) (Ext y) + | x == y = + Ext + <$> ( maybe (new x) return + =<< gets (M.lookup x . snd) + ) + unifyExtDims (Ext x) _ = Ext <$> new x + unifyExtDims _ (Ext x) = Ext <$> new x + new x = do + (n, m) <- get + put (n + 1, M.insert x n m) + return n -- | Given a list of 'ExtType's and a list of "forbidden" names, -- modify the dimensions of the 'ExtType's such that they are 'Ext' @@ -427,24 +464,30 @@ -- forbidden names. existentialiseExtTypes :: [VName] -> [ExtType] -> [ExtType] existentialiseExtTypes inaccessible = map makeBoundShapesFree - where makeBoundShapesFree = - modifyArrayShape $ fmap checkDim - checkDim (Free (Var v)) - | Just i <- v `elemIndex` inaccessible = - Ext i - checkDim d = d + where + makeBoundShapesFree = + modifyArrayShape $ fmap checkDim + checkDim (Free (Var v)) + | Just i <- v `elemIndex` inaccessible = + Ext i + checkDim d = d -- | Produce a mapping for the dimensions context. shapeExtMapping :: [TypeBase ExtShape u] -> [TypeBase Shape u1] -> M.Map Int SubExp shapeExtMapping = dimMapping arrayExtDims arrayDims match mappend - where match Free{} _ = mempty - match (Ext i) dim = M.singleton i dim + where + match Free {} _ = mempty + match (Ext i) dim = M.singleton i dim -dimMapping :: Monoid res => - (t1 -> [dim1]) -> (t2 -> [dim2]) -> (dim1 -> dim2 -> res) - -> (res -> res -> res) - -> [t1] -> [t2] - -> res +dimMapping :: + Monoid res => + (t1 -> [dim1]) -> + (t2 -> [dim2]) -> + (dim1 -> dim2 -> res) -> + (res -> res -> res) -> + [t1] -> + [t2] -> + res dimMapping getDims1 getDims2 f comb ts1 ts2 = foldl' comb mempty $ concat $ zipWith (zipWith f) (map getDims1 ts1) (map getDims2 ts2) @@ -491,7 +534,7 @@ instance Typed dec => Typed (PatElemT dec) where typeOf = typeOf . patElemDec -instance Typed b => Typed (a,b) where +instance Typed b => Typed (a, b) where typeOf = typeOf . snd -- | Typeclass for things that contain 'DeclType's. @@ -549,9 +592,10 @@ fixExt i se = fmap $ fixExt i se instance FixExt ExtSize where - fixExt i se (Ext j) | j > i = Ext $ j - 1 - | j == i = Free se - | otherwise = Ext j + fixExt i se (Ext j) + | j > i = Ext $ j - 1 + | j == i = Free se + | otherwise = Ext j fixExt _ _ (Free x) = Free x instance FixExt () where diff --git a/src/Futhark/IR/RetType.hs b/src/Futhark/IR/RetType.hs --- a/src/Futhark/IR/RetType.hs +++ b/src/Futhark/IR/RetType.hs @@ -1,18 +1,19 @@ -{-# LANGUAGE FlexibleInstances, TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} + -- | This module exports a type class covering representations of -- function return types. module Futhark.IR.RetType - ( - IsBodyType (..) - , IsRetType (..) - , expectedTypes - ) - where + ( IsBodyType (..), + IsRetType (..), + expectedTypes, + ) +where +import Control.Monad.Identity import qualified Data.Map.Strict as M - -import Futhark.IR.Syntax.Core import Futhark.IR.Prop.Types +import Futhark.IR.Syntax.Core -- | A type representing the return type of a body. It should contain -- at least the information contained in a list of 'ExtType's, but may @@ -35,50 +36,46 @@ -- | Given a function return type, the parameters of the function, -- and the arguments for a concrete call, return the instantiated -- return type for the concrete call, if valid. - applyRetType :: Typed dec => - [rt] - -> [Param dec] - -> [(SubExp, Type)] - -> Maybe [rt] + applyRetType :: + Typed dec => + [rt] -> + [Param dec] -> + [(SubExp, Type)] -> + Maybe [rt] -- | Given shape parameter names and value parameter types, produce the -- types of arguments accepted. expectedTypes :: Typed t => [VName] -> [t] -> [SubExp] -> [Type] expectedTypes shapes value_ts args = map (correctDims . typeOf) value_ts - where parammap :: M.Map VName SubExp - parammap = M.fromList $ zip shapes args - - correctDims t = - t `setArrayShape` - Shape (map correctDim $ shapeDims $ arrayShape t) + where + parammap :: M.Map VName SubExp + parammap = M.fromList $ zip shapes args - correctDim (Constant v) = Constant v - correctDim (Var v) - | Just se <- M.lookup v parammap = se - | otherwise = Var v + correctDims = runIdentity . mapOnType (pure . f) + where + f (Var v) + | Just se <- M.lookup v parammap = se + f se = se instance IsRetType DeclExtType where primRetType = Prim applyRetType extret params args = - if length args == length params && - and (zipWith subtypeOf argtypes $ - expectedTypes (map paramName params) params $ map fst args) - then Just $ map correctExtDims extret - else Nothing - where argtypes = map snd args - - parammap :: M.Map VName SubExp - parammap = M.fromList $ zip (map paramName params) (map fst args) - - correctExtDims t = - t `setArrayShape` - Shape (map correctExtDim $ shapeDims $ arrayShape t) + if length args == length params + && and + ( zipWith subtypeOf argtypes $ + expectedTypes (map paramName params) params $ map fst args + ) + then Just $ map correctExtDims extret + else Nothing + where + argtypes = map snd args - correctExtDim (Ext i) = Ext i - correctExtDim (Free d) = Free $ correctDim d + parammap :: M.Map VName SubExp + parammap = M.fromList $ zip (map paramName params) (map fst args) - correctDim (Constant v) = Constant v - correctDim (Var v) + correctExtDims = runIdentity . mapOnExtType (pure . f) + where + f (Var v) | Just se <- M.lookup v parammap = se - | otherwise = Var v + f se = se diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs --- a/src/Futhark/IR/SOACS.hs +++ b/src/Futhark/IR/SOACS.hs @@ -1,43 +1,54 @@ -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE Safe #-} +{-# LANGUAGE TypeFamilies #-} + -- | A simple representation with SOACs and nested parallelism. module Futhark.IR.SOACS - ( -- * The Lore definition - SOACS - -- * Syntax types - , Body - , Stm - , Pattern - , Exp - , Lambda - , FParam - , LParam - , RetType - , PatElem - -- * Module re-exports - , module Futhark.IR.Prop - , module Futhark.IR.Traversals - , module Futhark.IR.Pretty - , module Futhark.IR.Syntax - , module Futhark.IR.SOACS.SOAC - , AST.LambdaT(Lambda) - , AST.BodyT(Body) - , AST.PatternT(Pattern) - , AST.PatElemT(PatElem) - ) + ( -- * The Lore definition + SOACS, + + -- * Syntax types + Body, + Stm, + Pattern, + Exp, + Lambda, + FParam, + LParam, + RetType, + PatElem, + + -- * Module re-exports + module Futhark.IR.Prop, + module Futhark.IR.Traversals, + module Futhark.IR.Pretty, + module Futhark.IR.Syntax, + module Futhark.IR.SOACS.SOAC, + AST.LambdaT (Lambda), + AST.BodyT (Body), + AST.PatternT (Pattern), + AST.PatElemT (PatElem), + ) where -import qualified Futhark.IR.Syntax as AST -import Futhark.IR.Syntax - hiding (Exp, Body, Stm, - Pattern, Lambda, FParam, LParam, RetType, PatElem) -import Futhark.IR.SOACS.SOAC -import Futhark.IR.Prop -import Futhark.IR.Traversals -import Futhark.IR.Pretty import Futhark.Binder import Futhark.Construct +import Futhark.IR.Pretty +import Futhark.IR.Prop +import Futhark.IR.SOACS.SOAC +import Futhark.IR.Syntax hiding + ( Body, + Exp, + FParam, + LParam, + Lambda, + PatElem, + Pattern, + RetType, + Stm, + ) +import qualified Futhark.IR.Syntax as AST +import Futhark.IR.Traversals import qualified Futhark.TypeCheck as TypeCheck -- This module could be written much nicer if Haskell had functors @@ -54,19 +65,27 @@ expTypesFromPattern = return . expExtTypesFromPattern type Exp = AST.Exp SOACS + type Body = AST.Body SOACS + type Stm = AST.Stm SOACS + type Pattern = AST.Pattern SOACS + type Lambda = AST.Lambda SOACS + type FParam = AST.FParam SOACS + type LParam = AST.LParam SOACS + type RetType = AST.RetType SOACS + type PatElem = AST.PatElem SOACS instance TypeCheck.CheckableOp SOACS where checkOp = typeCheckSOAC -instance TypeCheck.Checkable SOACS where +instance TypeCheck.Checkable SOACS instance Bindable SOACS where mkBody = AST.Body () @@ -74,6 +93,6 @@ mkExpDec _ _ = () mkLetNames = simpleMkLetNames -instance BinderOps SOACS where +instance BinderOps SOACS -instance PrettyLore SOACS where +instance PrettyLore SOACS 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 @@ -1,83 +1,85 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + -- | Definition of /Second-Order Array Combinators/ (SOACs), which are -- the main form of parallelism in the early stages of the compiler. module Futhark.IR.SOACS.SOAC - ( SOAC(..) - , StreamOrd(..) - , StreamForm(..) - , ScremaForm(..) - , HistOp(..) - , Scan(..) - , scanResults - , singleScan - , Reduce(..) - , redResults - , singleReduce - - -- * Utility - , getStreamAccums - , scremaType - , soacType - - , typeCheckSOAC - - , mkIdentityLambda - , isIdentityLambda - , nilFn - , scanomapSOAC - , redomapSOAC - , scanSOAC - , reduceSOAC - , mapSOAC - , isScanomapSOAC - , isRedomapSOAC - , isScanSOAC - , isReduceSOAC - , isMapSOAC + ( SOAC (..), + StreamOrd (..), + StreamForm (..), + ScremaForm (..), + HistOp (..), + Scan (..), + scanResults, + singleScan, + Reduce (..), + redResults, + singleReduce, - , ppScrema - , ppHist + -- * Utility + getStreamAccums, + scremaType, + soacType, + typeCheckSOAC, + mkIdentityLambda, + isIdentityLambda, + nilFn, + scanomapSOAC, + redomapSOAC, + scanSOAC, + reduceSOAC, + mapSOAC, + isScanomapSOAC, + isRedomapSOAC, + isScanSOAC, + isReduceSOAC, + isMapSOAC, + ppScrema, + ppHist, - -- * Generic traversal - , SOACMapper(..) - , identitySOACMapper - , mapSOACM - ) - where + -- * Generic traversal + SOACMapper (..), + identitySOACMapper, + mapSOACM, + ) +where +import Control.Category +import Control.Monad.Identity import Control.Monad.State.Strict import Control.Monad.Writer -import Control.Monad.Identity +import Data.List (intersperse) import qualified Data.Map.Strict as M import Data.Maybe -import Data.List (intersperse) - -import Futhark.IR import qualified Futhark.Analysis.Alias as Alias -import qualified Futhark.Util.Pretty as PP -import Futhark.Util.Pretty (ppr, Doc, Pretty, parens, comma, (), (<+>), commasep, text) +import Futhark.Analysis.Metrics +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Construct +import Futhark.IR +import Futhark.IR.Aliases (Aliases, removeLambdaAliases) import Futhark.IR.Prop.Aliases -import Futhark.Transform.Substitute -import Futhark.Transform.Rename import Futhark.Optimise.Simplify.Lore -import Futhark.IR.Aliases (Aliases, removeLambdaAliases) -import qualified Futhark.Analysis.SymbolTable as ST -import Futhark.Analysis.PrimExp.Convert +import Futhark.Transform.Rename +import Futhark.Transform.Substitute import qualified Futhark.TypeCheck as TC -import Futhark.Analysis.Metrics -import Futhark.Construct -import Futhark.Util (maybeNth, chunks) +import Futhark.Util (chunks, maybeNth) +import Futhark.Util.Pretty (Doc, Pretty, comma, commasep, parens, ppr, text, (<+>), ()) +import qualified Futhark.Util.Pretty as PP +import GHC.Generics (Generic) +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) -- | A second-order array combinator (SOAC). -data SOAC lore = - Stream SubExp (StreamForm lore) (Lambda lore) [VName] - | Scatter SubExp (Lambda lore) [VName] [(SubExp, Int, VName)] - -- ^ @Scatter @ +data SOAC lore + = Stream SubExp (StreamForm lore) (Lambda lore) [VName] + | -- | @Scatter @ -- -- @@ -92,65 +94,142 @@ -- [index_0, index_1, ..., index_n, value_0, value_1, ..., value_n] -- -- This must be consistent along all Scatter-related optimisations. - | Hist SubExp [HistOp lore] (Lambda lore) [VName] - -- ^ @Hist @ + Scatter SubExp (Lambda lore) [VName] [(SubExp, Int, VName)] + | -- | @Hist @ -- -- The first SubExp is the length of the input arrays. The first -- list describes the operations to perform. The t'Lambda' is the -- bucket function. Finally comes the input images. - | Screma SubExp (ScremaForm lore) [VName] - -- ^ A combination of scan, reduction, and map. The first + Hist SubExp [HistOp lore] (Lambda lore) [VName] + | -- | A combination of scan, reduction, and map. The first -- t'SubExp' is the size of the input arrays. - deriving (Eq, Ord, Show) + Screma SubExp (ScremaForm lore) [VName] + deriving (Eq, Ord, Show, Generic) +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 "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 + (. Sexp.list (Sexp.el (Sexp.sym "screma") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End + -- | Information about computing a single histogram. -data HistOp lore = HistOp { histWidth :: SubExp - , histRaceFactor :: SubExp - -- ^ Race factor @RF@ means that only @1/RF@ - -- bins are used. - , histDest :: [VName] - , histNeutral :: [SubExp] - , histOp :: Lambda lore - } - deriving (Eq, Ord, Show) +data HistOp lore = HistOp + { histWidth :: SubExp, + -- | Race factor @RF@ means that only @1/RF@ + -- bins are used. + histRaceFactor :: SubExp, + histDest :: [VName], + histNeutral :: [SubExp], + histOp :: Lambda lore + } + deriving (Eq, Ord, Show, Generic) +instance Decorations lore => SexpIso (HistOp lore) where + sexpIso = with $ \histop -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> histop + -- | Is the stream chunk required to correspond to a contiguous -- subsequence of the original input ('InOrder') or not? 'Disorder' -- streams can be more efficient, but not all algorithms work with -- this. -data StreamOrd = InOrder | Disorder - deriving (Eq, Ord, Show) +data StreamOrd = InOrder | Disorder + deriving (Eq, Ord, Show, Generic) +instance SexpIso StreamOrd where + sexpIso = + match $ + With (. Sexp.sym "in-order") $ + With + (. Sexp.sym "disorder") + End + -- | What kind of stream is this? -data StreamForm lore = - Parallel StreamOrd Commutativity (Lambda lore) [SubExp] +data StreamForm lore + = Parallel StreamOrd Commutativity (Lambda lore) [SubExp] | Sequential [SubExp] - deriving (Eq, Ord, Show) + deriving (Eq, Ord, Show, Generic) +instance Decorations lore => SexpIso (StreamForm lore) where + sexpIso = + match $ + With + ( . + Sexp.list + ( Sexp.el (Sexp.sym "parallel") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.rest sexpIso + ) + ) + $ With + ( . + Sexp.list + ( Sexp.el (Sexp.sym "sequential") + >>> Sexp.rest sexpIso + ) + ) + End + -- | The essential parts of a 'Screma' factored out (everything -- except the input arrays). -data ScremaForm lore = ScremaForm - [Scan lore] - [Reduce lore] - (Lambda lore) - deriving (Eq, Ord, Show) +data ScremaForm lore + = ScremaForm + [Scan lore] + [Reduce lore] + (Lambda lore) + deriving (Eq, Ord, Show, Generic) +instance Decorations lore => SexpIso (ScremaForm lore) where + sexpIso = with $ \scremaform -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> scremaform + singleBinOp :: Bindable lore => [Lambda lore] -> Lambda lore singleBinOp lams = - Lambda { lambdaParams = concatMap xParams lams ++ concatMap yParams lams - , lambdaReturnType = concatMap lambdaReturnType lams - , lambdaBody = mkBody (mconcat (map (bodyStms . lambdaBody) lams)) - (concatMap (bodyResult . lambdaBody) lams) - } - where xParams lam = take (length (lambdaReturnType lam)) (lambdaParams lam) - yParams lam = drop (length (lambdaReturnType lam)) (lambdaParams lam) + Lambda + { lambdaParams = concatMap xParams lams ++ concatMap yParams lams, + lambdaReturnType = concatMap lambdaReturnType lams, + lambdaBody = + mkBody + (mconcat (map (bodyStms . lambdaBody) lams)) + (concatMap (bodyResult . lambdaBody) lams) + } + where + xParams lam = take (length (lambdaReturnType lam)) (lambdaParams lam) + yParams lam = drop (length (lambdaReturnType lam)) (lambdaParams lam) -- | How to compute a single scan result. -data Scan lore = Scan { scanLambda :: Lambda lore - , scanNeutral :: [SubExp] - } - deriving (Eq, Ord, Show) +data Scan lore = Scan + { scanLambda :: Lambda lore, + scanNeutral :: [SubExp] + } + deriving (Eq, Ord, Show, Generic) +instance Decorations lore => SexpIso (Scan lore) where + sexpIso = with $ \scan -> + Sexp.list + ( Sexp.el (Sexp.sym "scan") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> scan + -- | How many reduction results are produced by these 'Scan's? scanResults :: [Scan lore] -> Int scanResults = sum . map (length . scanNeutral) @@ -160,15 +239,26 @@ singleScan scans = let scan_nes = concatMap scanNeutral scans scan_lam = singleBinOp $ map scanLambda scans - in Scan scan_lam scan_nes + in Scan scan_lam scan_nes -- | How to compute a single reduction result. -data Reduce lore = Reduce { redComm :: Commutativity - , redLambda :: Lambda lore - , redNeutral :: [SubExp] - } - deriving (Eq, Ord, Show) +data Reduce lore = Reduce + { redComm :: Commutativity, + redLambda :: Lambda lore, + redNeutral :: [SubExp] + } + deriving (Eq, Ord, Show, Generic) +instance Decorations lore => SexpIso (Reduce lore) where + sexpIso = with $ \red -> + Sexp.list + ( Sexp.el (Sexp.sym "reduce") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> red + -- | How many reduction results are produced by these 'Reduce's? redResults :: [Reduce lore] -> Int redResults = sum . map (length . redNeutral) @@ -178,32 +268,40 @@ singleReduce reds = let red_nes = concatMap redNeutral reds red_lam = singleBinOp $ map redLambda reds - in Reduce (mconcat (map redComm reds)) red_lam red_nes + in Reduce (mconcat (map redComm reds)) red_lam red_nes -- | The types produced by a single 'Screma', given the size of the -- input array. scremaType :: SubExp -> ScremaForm lore -> [Type] scremaType w (ScremaForm scans reds map_lam) = scan_tps ++ red_tps ++ map (`arrayOfRow` w) map_tps - where scan_tps = map (`arrayOfRow` w) $ - concatMap (lambdaReturnType . scanLambda) scans - red_tps = concatMap (lambdaReturnType . redLambda) reds - map_tps = drop (length scan_tps + length red_tps) $ lambdaReturnType map_lam + where + scan_tps = + map (`arrayOfRow` w) $ + concatMap (lambdaReturnType . scanLambda) scans + red_tps = concatMap (lambdaReturnType . redLambda) reds + map_tps = drop (length scan_tps + length red_tps) $ lambdaReturnType map_lam -- | Construct a lambda that takes parameters of the given types and -- simply returns them unchanged. -mkIdentityLambda :: (Bindable lore, MonadFreshNames m) => - [Type] -> m (Lambda lore) +mkIdentityLambda :: + (Bindable lore, MonadFreshNames m) => + [Type] -> + m (Lambda lore) mkIdentityLambda ts = do params <- mapM (newParam "x") ts - return Lambda { lambdaParams = params - , lambdaBody = mkBody mempty $ map (Var . paramName) params - , lambdaReturnType = ts } + return + Lambda + { lambdaParams = params, + lambdaBody = mkBody mempty $ map (Var . paramName) params, + lambdaReturnType = ts + } -- | Is the given lambda an identity lambda? isIdentityLambda :: Lambda lore -> Bool -isIdentityLambda lam = bodyResult (lambdaBody lam) == - map (Var . paramName) (lambdaParams lam) +isIdentityLambda lam = + bodyResult (lambdaBody lam) + == map (Var . paramName) (lambdaParams lam) -- | A lambda with no parameters that returns no values. nilFn :: Bindable lore => Lambda lore @@ -221,17 +319,23 @@ -- | Construct a Screma with possibly multiple scans, and identity map -- function. -scanSOAC :: (Bindable lore, MonadFreshNames m) => - [Scan lore] -> m (ScremaForm lore) +scanSOAC :: + (Bindable lore, MonadFreshNames m) => + [Scan lore] -> + m (ScremaForm lore) scanSOAC scans = scanomapSOAC scans <$> mkIdentityLambda ts - where ts = concatMap (lambdaReturnType . scanLambda) scans + where + ts = concatMap (lambdaReturnType . scanLambda) scans -- | Construct a Screma with possibly multiple reductions, and -- identity map function. -reduceSOAC :: (Bindable lore, MonadFreshNames m) => - [Reduce lore] -> m (ScremaForm lore) +reduceSOAC :: + (Bindable lore, MonadFreshNames m) => + [Reduce lore] -> + m (ScremaForm lore) reduceSOAC reds = redomapSOAC reds <$> mkIdentityLambda ts - where ts = concatMap (lambdaReturnType . redLambda) reds + where + ts = concatMap (lambdaReturnType . redLambda) reds -- | Construct a Screma corresponding to a map. mapSOAC :: Lambda lore -> ScremaForm lore @@ -246,9 +350,10 @@ -- | Does this Screma correspond to pure scan? isScanSOAC :: ScremaForm lore -> Maybe [Scan lore] -isScanSOAC form = do (scans, map_lam) <- isScanomapSOAC form - guard $ isIdentityLambda map_lam - return scans +isScanSOAC form = do + (scans, map_lam) <- isScanomapSOAC form + guard $ isIdentityLambda map_lam + return scans -- | Does this Screma correspond to a reduce-map composition? isRedomapSOAC :: ScremaForm lore -> Maybe ([Reduce lore], Lambda lore) @@ -259,9 +364,10 @@ -- | Does this Screma correspond to a pure reduce? isReduceSOAC :: ScremaForm lore -> Maybe [Reduce lore] -isReduceSOAC form = do (reds, map_lam) <- isRedomapSOAC form - guard $ isIdentityLambda map_lam - return reds +isReduceSOAC form = do + (reds, map_lam) <- isRedomapSOAC form + guard $ isIdentityLambda map_lam + return reds -- | Does this Screma correspond to a simple map, without any -- reduction or scan results? @@ -272,100 +378,132 @@ return map_lam -- | Like 'Mapper', but just for 'SOAC's. -data SOACMapper flore tlore m = SOACMapper { - mapOnSOACSubExp :: SubExp -> m SubExp - , mapOnSOACLambda :: Lambda flore -> m (Lambda tlore) - , mapOnSOACVName :: VName -> m VName +data SOACMapper flore tlore m = SOACMapper + { mapOnSOACSubExp :: SubExp -> m SubExp, + mapOnSOACLambda :: Lambda flore -> m (Lambda tlore), + mapOnSOACVName :: VName -> m VName } -- | A mapper that simply returns the SOAC verbatim. identitySOACMapper :: Monad m => SOACMapper lore lore m -identitySOACMapper = SOACMapper { mapOnSOACSubExp = return - , mapOnSOACLambda = return - , mapOnSOACVName = return - } +identitySOACMapper = + SOACMapper + { mapOnSOACSubExp = return, + mapOnSOACLambda = return, + mapOnSOACVName = return + } -- | Map a monadic action across the immediate children of a -- SOAC. The mapping does not descend recursively into subexpressions -- and is done left-to-right. -mapSOACM :: (Applicative m, Monad m) => - SOACMapper flore tlore m -> SOAC flore -> m (SOAC tlore) +mapSOACM :: + (Applicative m, Monad m) => + SOACMapper flore tlore m -> + SOAC flore -> + m (SOAC tlore) mapSOACM tv (Stream size form lam arrs) = - Stream <$> mapOnSOACSubExp tv size <*> - mapOnStreamForm form <*> mapOnSOACLambda tv lam <*> - 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 + Stream <$> mapOnSOACSubExp tv size + <*> mapOnStreamForm form + <*> mapOnSOACLambda tv lam + <*> 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 mapSOACM tv (Scatter len lam ivs as) = Scatter - <$> mapOnSOACSubExp tv len - <*> mapOnSOACLambda tv lam - <*> mapM (mapOnSOACVName tv) ivs - <*> mapM (\(aw,an,a) -> (,,) <$> mapOnSOACSubExp tv aw <*> - pure an <*> mapOnSOACVName tv a) as + <$> mapOnSOACSubExp tv len + <*> mapOnSOACLambda tv lam + <*> mapM (mapOnSOACVName tv) ivs + <*> mapM + ( \(aw, an, a) -> + (,,) <$> mapOnSOACSubExp tv aw + <*> pure an + <*> mapOnSOACVName tv a + ) + as mapSOACM tv (Hist len ops bucket_fun imgs) = Hist - <$> mapOnSOACSubExp tv len - <*> mapM (\(HistOp e rf arrs nes op) -> - HistOp <$> mapOnSOACSubExp tv e - <*> mapOnSOACSubExp tv rf - <*> mapM (mapOnSOACVName tv) arrs - <*> mapM (mapOnSOACSubExp tv) nes - <*> mapOnSOACLambda tv op) ops - <*> mapOnSOACLambda tv bucket_fun - <*> mapM (mapOnSOACVName tv) imgs + <$> mapOnSOACSubExp tv len + <*> mapM + ( \(HistOp e rf arrs nes op) -> + HistOp <$> mapOnSOACSubExp tv e + <*> mapOnSOACSubExp tv rf + <*> mapM (mapOnSOACVName tv) arrs + <*> mapM (mapOnSOACSubExp tv) nes + <*> mapOnSOACLambda tv op + ) + ops + <*> mapOnSOACLambda tv bucket_fun + <*> mapM (mapOnSOACVName tv) imgs mapSOACM tv (Screma w (ScremaForm scans reds map_lam) arrs) = - Screma <$> mapOnSOACSubExp tv w <*> - (ScremaForm <$> - forM scans (\(Scan red_lam red_nes) -> - Scan <$> mapOnSOACLambda tv red_lam <*> - mapM (mapOnSOACSubExp tv) red_nes) <*> - forM reds (\(Reduce comm red_lam red_nes) -> - Reduce comm <$> mapOnSOACLambda tv red_lam <*> - mapM (mapOnSOACSubExp tv) red_nes) <*> - mapOnSOACLambda tv map_lam) - <*> mapM (mapOnSOACVName tv) arrs + Screma <$> mapOnSOACSubExp tv w + <*> ( ScremaForm + <$> forM + scans + ( \(Scan red_lam red_nes) -> + Scan <$> mapOnSOACLambda tv red_lam + <*> mapM (mapOnSOACSubExp tv) red_nes + ) + <*> forM + reds + ( \(Reduce comm red_lam red_nes) -> + Reduce comm <$> mapOnSOACLambda tv red_lam + <*> mapM (mapOnSOACSubExp tv) red_nes + ) + <*> mapOnSOACLambda tv map_lam + ) + <*> mapM (mapOnSOACVName tv) arrs instance ASTLore lore => FreeIn (SOAC lore) where freeIn' = flip execState mempty . mapSOACM free - where walk f x = modify (<>f x) >> return x - free = SOACMapper { mapOnSOACSubExp = walk freeIn' - , mapOnSOACLambda = walk freeIn' - , mapOnSOACVName = walk freeIn' - } + where + walk f x = modify (<> f x) >> return x + free = + SOACMapper + { mapOnSOACSubExp = walk freeIn', + mapOnSOACLambda = walk freeIn', + mapOnSOACVName = walk freeIn' + } instance ASTLore lore => Substitute (SOAC lore) where substituteNames subst = runIdentity . mapSOACM substitute - where substitute = - SOACMapper { mapOnSOACSubExp = return . substituteNames subst - , mapOnSOACLambda = return . substituteNames subst - , mapOnSOACVName = return . substituteNames subst - } + where + substitute = + SOACMapper + { mapOnSOACSubExp = return . substituteNames subst, + mapOnSOACLambda = return . substituteNames subst, + mapOnSOACVName = return . substituteNames subst + } instance ASTLore lore => Rename (SOAC lore) where rename = mapSOACM renamer - where renamer = SOACMapper rename rename rename + where + renamer = SOACMapper rename rename rename -- | The type of a SOAC. soacType :: SOAC lore -> [Type] soacType (Stream outersize form lam _) = 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 + 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 - where val_ts = concatMap (take 1) $ chunks ns $ - drop (sum ns) $ lambdaReturnType lam - (ws, ns, _) = unzip3 as + where + val_ts = + concatMap (take 1) $ + chunks ns $ + drop (sum ns) $ lambdaReturnType lam + (ws, ns, _) = unzip3 as soacType (Hist _len ops _bucket_fun _imgs) = do op <- ops map (`arrayOfRow` histWidth op) (lambdaReturnType $ histOp op) @@ -382,56 +520,79 @@ -- and reduce functions are always considered "fresh". consumedInOp (Screma _ (ScremaForm _ _ map_lam) arrs) = mapNames consumedArray $ consumedByLambda map_lam - where consumedArray v = fromMaybe v $ lookup v params_to_arrs - params_to_arrs = zip (map paramName $ lambdaParams map_lam) arrs + 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) = - namesFromList $ subExpVars $ - case form of Sequential accs -> - map (consumedArray accs) $ namesToList $ consumedByLambda lam - Parallel _ _ _ accs -> - map (consumedArray accs) $ namesToList $ consumedByLambda lam - where consumedArray accs v = fromMaybe (Var v) $ lookup v $ paramsToInput accs - -- Drop the chunk parameter, which cannot alias anything. - paramsToInput accs = zip - (map paramName $ drop 1 $ lambdaParams lam) - (accs++map Var arrs) + namesFromList $ + subExpVars $ + case form of + Sequential accs -> + map (consumedArray accs) $ namesToList $ consumedByLambda lam + Parallel _ _ _ accs -> + map (consumedArray accs) $ namesToList $ consumedByLambda lam + where + consumedArray accs v = fromMaybe (Var v) $ lookup v $ paramsToInput accs + -- Drop the chunk parameter, which cannot alias anything. + paramsToInput accs = + zip + (map paramName $ drop 1 $ lambdaParams lam) + (accs ++ map Var arrs) consumedInOp (Scatter _ _ _ as) = namesFromList $ map (\(_, _, a) -> a) as consumedInOp (Hist _ ops _ _) = namesFromList $ concatMap histDest ops -mapHistOp :: (Lambda flore -> Lambda tlore) - -> HistOp flore -> HistOp tlore +mapHistOp :: + (Lambda flore -> Lambda tlore) -> + HistOp flore -> + HistOp tlore mapHistOp f (HistOp w rf dests nes lam) = HistOp w rf dests nes $ f lam -instance (ASTLore lore, - ASTLore (Aliases lore), - CanBeAliased (Op lore)) => CanBeAliased (SOAC lore) where +instance + ( ASTLore lore, + ASTLore (Aliases lore), + CanBeAliased (Op lore) + ) => + CanBeAliased (SOAC lore) + where type OpWithAliases (SOAC lore) = SOAC (Aliases lore) addOpAliases (Stream size form lam arr) = - Stream size (analyseStreamForm form) - (Alias.analyseLambda lam) arr - where analyseStreamForm (Parallel o comm lam0 acc) = - Parallel o comm (Alias.analyseLambda lam0) acc - analyseStreamForm (Sequential acc) = Sequential acc + Stream + size + (analyseStreamForm form) + (Alias.analyseLambda lam) + 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) = - Hist len (map (mapHistOp Alias.analyseLambda) ops) - (Alias.analyseLambda bucket_fun) imgs + Hist + len + (map (mapHistOp Alias.analyseLambda) ops) + (Alias.analyseLambda bucket_fun) + imgs addOpAliases (Screma w (ScremaForm scans reds map_lam) arrs) = - Screma w (ScremaForm - (map onScan scans) - (map onRed reds) - (Alias.analyseLambda map_lam)) - arrs - where onRed red = red { redLambda = Alias.analyseLambda $ redLambda red } - onScan scan = scan { scanLambda = Alias.analyseLambda $ scanLambda scan } + Screma + w + ( ScremaForm + (map onScan scans) + (map onRed reds) + (Alias.analyseLambda map_lam) + ) + arrs + where + onRed red = red {redLambda = Alias.analyseLambda $ redLambda red} + onScan scan = scan {scanLambda = Alias.analyseLambda $ scanLambda scan} removeOpAliases = runIdentity . mapSOACM remove - where remove = SOACMapper return (return . removeLambdaAliases) return + where + remove = SOACMapper return (return . removeLambdaAliases) return instance ASTLore lore => IsOp (SOAC lore) where safeOp _ = False @@ -442,7 +603,7 @@ substNamesInType _ (Mem space) = Mem space substNamesInType subs (Array btp shp u) = let shp' = Shape $ map (substNamesInSubExp subs) (shapeDims shp) - in Array btp shp' u + in Array btp shp' u substNamesInSubExp :: M.Map VName SubExp -> SubExp -> SubExp substNamesInSubExp _ e@(Constant _) = e @@ -453,77 +614,81 @@ type OpWithWisdom (SOAC lore) = SOAC (Wise lore) removeOpWisdom = runIdentity . mapSOACM remove - where remove = SOACMapper return (return . removeLambdaWisdom) return + where + remove = SOACMapper return (return . removeLambdaWisdom) return instance Decorations lore => ST.IndexOp (SOAC lore) where indexOp vtable k soac [i] = do - (lam,se,arr_params,arrs) <- lambdaAndSubExp soac + (lam, se, arr_params, arrs) <- lambdaAndSubExp soac let arr_indexes = M.fromList $ catMaybes $ zipWith arrIndex arr_params arrs arr_indexes' = foldl expandPrimExpTable arr_indexes $ bodyStms $ lambdaBody lam case se of Var v -> uncurry (flip ST.Indexed) <$> M.lookup v arr_indexes' _ -> Nothing - where lambdaAndSubExp (Screma _ (ScremaForm scans reds map_lam) arrs) = - nthMapOut (scanResults scans + redResults reds) map_lam arrs - lambdaAndSubExp _ = - Nothing + where + lambdaAndSubExp (Screma _ (ScremaForm scans reds map_lam) arrs) = + nthMapOut (scanResults scans + redResults reds) map_lam arrs + lambdaAndSubExp _ = + Nothing - nthMapOut num_accs lam arrs = do - se <- maybeNth (num_accs+k) $ bodyResult $ lambdaBody lam - return (lam, se, drop num_accs $ lambdaParams lam, arrs) + nthMapOut num_accs lam arrs = do + se <- maybeNth (num_accs + k) $ bodyResult $ lambdaBody lam + return (lam, se, drop num_accs $ lambdaParams lam, arrs) - arrIndex p arr = do - ST.Indexed cs pe <- ST.index' arr [i] vtable - return (paramName p, (pe,cs)) + arrIndex p arr = do + ST.Indexed cs pe <- ST.index' arr [i] vtable + return (paramName p, (pe, cs)) - expandPrimExpTable table stm - | [v] <- patternNames $ stmPattern stm, - Just (pe,cs) <- - runWriterT $ primExpFromExp (asPrimExp table) $ stmExp stm, - all (`ST.elem` vtable) (unCertificates $ stmCerts stm) = - M.insert v (pe, stmCerts stm <> cs) table - | otherwise = - table + expandPrimExpTable table stm + | [v] <- patternNames $ stmPattern stm, + Just (pe, cs) <- + runWriterT $ primExpFromExp (asPrimExp table) $ stmExp stm, + all (`ST.elem` vtable) (unCertificates $ stmCerts stm) = + M.insert v (pe, stmCerts stm <> cs) table + | otherwise = + table - asPrimExp table v - | Just (e,cs) <- M.lookup v table = tell cs >> return e - | Just (Prim pt) <- ST.lookupType v vtable = - return $ LeafExp v pt - | otherwise = lift Nothing + asPrimExp table v + | Just (e, cs) <- M.lookup v table = tell cs >> return e + | Just (Prim pt) <- ST.lookupType v vtable = + return $ LeafExp v pt + | otherwise = lift Nothing indexOp _ _ _ _ = Nothing -- | 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 - TC.require [Prim int32] size + TC.require [Prim int64] size accargs <- mapM TC.checkArg accexps arrargs <- mapM lookupType arrexps _ <- TC.checkSOACArrayArgs size arrexps let chunk = head $ lambdaParams lam let asArg t = (t, mempty) - inttp = Prim int32 - lamarrs'= map (`setOuterSize` Var (paramName chunk)) arrargs - let acc_len= length accexps + inttp = Prim int64 + lamarrs' = map (`setOuterSize` Var (paramName chunk)) arrargs + let acc_len = length accexps let lamrtp = take acc_len $ lambdaReturnType lam unless (map TC.argType accargs == lamrtp) $ TC.bad $ TC.TypeError "Stream with inconsistent accumulator type in lambda." -- check reduce's lambda, if any _ <- case form of - Parallel _ _ lam0 _ -> do - let acct = map TC.argType accargs - outerRetType = lambdaReturnType lam0 - TC.checkLambda lam0 $ map TC.noArgAliases $ accargs ++ accargs - unless (acct == outerRetType) $ - TC.bad $ TC.TypeError $ - "Initial value is of type " ++ prettyTuple acct ++ - ", but stream's reduce lambda returns type " ++ prettyTuple outerRetType ++ "." - _ -> return () + Parallel _ _ lam0 _ -> do + let acct = map TC.argType accargs + outerRetType = lambdaReturnType lam0 + TC.checkLambda lam0 $ map TC.noArgAliases $ accargs ++ accargs + unless (acct == outerRetType) $ + TC.bad $ + TC.TypeError $ + "Initial value is of type " ++ prettyTuple acct + ++ ", but stream's reduce lambda returns type " + ++ prettyTuple outerRetType + ++ "." + _ -> 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' TC.checkLambda lam $ asArg inttp : accargs ++ fake_lamarrs' - typeCheckSOAC (Scatter w lam ivs as) = do -- Requirements: -- @@ -533,7 +698,7 @@ -- 1. The number of index types must be equal to the number of value types -- and the number of writes to arrays in @as@. -- - -- 2. Each index type must have the type i32. + -- 2. Each index type must have the type i64. -- -- 3. Each array in @as@ and the value types must have the same type -- @@ -547,7 +712,7 @@ -- Code: -- First check the input size. - TC.require [Prim int32] w + TC.require [Prim int64] w -- 0. let (_as_ws, as_ns, _as_vs) = unzip3 as @@ -557,16 +722,17 @@ rtsV = drop rtsLen rts -- 1. - unless (rtsLen == sum as_ns) - $ TC.bad $ TC.TypeError "Scatter: Uneven number of index types, value types, and arrays outputs." + unless (rtsLen == sum as_ns) $ + TC.bad $ TC.TypeError "Scatter: Uneven number of index types, value types, and arrays outputs." -- 2. - forM_ rtsI $ \rtI -> unless (Prim int32 == rtI) $ - TC.bad $ TC.TypeError "Scatter: Index return type must be i32." + forM_ rtsI $ \rtI -> + unless (Prim int64 == rtI) $ + TC.bad $ TC.TypeError "Scatter: Index return type must be i64." forM_ (zip (chunks as_ns rtsV) as) $ \(rtVs, (aw, _, a)) -> do - -- All lengths must have type i32. - TC.require [Prim int32] aw + -- All lengths must have type i64. + TC.require [Prim int64] aw -- 3. forM_ rtVs $ \rtV -> TC.requireI [rtV `arrayOfRow` aw] a @@ -577,23 +743,25 @@ -- 5. arrargs <- TC.checkSOACArrayArgs w ivs TC.checkLambda lam arrargs - typeCheckSOAC (Hist len ops bucket_fun imgs) = do - TC.require [Prim int32] len + TC.require [Prim int64] len -- Check the operators. forM_ ops $ \(HistOp dest_w rf dests nes op) -> do nes' <- mapM TC.checkArg nes - TC.require [Prim int32] dest_w - TC.require [Prim int32] rf + TC.require [Prim int64] dest_w + TC.require [Prim int64] rf -- Operator type must match the type of neutral elements. TC.checkLambda op $ map TC.noArgAliases $ nes' ++ nes' let nes_t = map TC.argType nes' unless (nes_t == lambdaReturnType op) $ - TC.bad $ TC.TypeError $ "Operator has return type " ++ - prettyTuple (lambdaReturnType op) ++ " but neutral element has type " ++ - prettyTuple nes_t + TC.bad $ + TC.TypeError $ + "Operator has return type " + ++ prettyTuple (lambdaReturnType op) + ++ " but neutral element has type " + ++ prettyTuple nes_t -- Arrays must have proper type. forM_ (zip nes_t dests) $ \(t, dest) -> do @@ -607,48 +775,62 @@ -- Return type of bucket function must be an index for each -- operation followed by the values to write. nes_ts <- concat <$> mapM (mapM subExpType . histNeutral) ops - let bucket_ret_t = replicate (length ops) (Prim int32) ++ nes_ts + let bucket_ret_t = replicate (length ops) (Prim int64) ++ nes_ts unless (bucket_ret_t == lambdaReturnType bucket_fun) $ - TC.bad $ TC.TypeError $ "Bucket function has return type " ++ - prettyTuple (lambdaReturnType bucket_fun) ++ " but should have type " ++ - prettyTuple bucket_ret_t - + TC.bad $ + TC.TypeError $ + "Bucket function has return type " + ++ prettyTuple (lambdaReturnType bucket_fun) + ++ " but should have type " + ++ prettyTuple bucket_ret_t typeCheckSOAC (Screma w (ScremaForm scans reds map_lam) arrs) = do - TC.require [Prim int32] w + TC.require [Prim int64] w arrs' <- TC.checkSOACArrayArgs w arrs TC.checkLambda map_lam $ map TC.noArgAliases arrs' - scan_nes' <- fmap concat $ forM scans $ \(Scan scan_lam scan_nes) -> do - scan_nes' <- mapM TC.checkArg scan_nes - let scan_t = map TC.argType scan_nes' - TC.checkLambda scan_lam $ map TC.noArgAliases $ scan_nes' ++ scan_nes' - unless (scan_t == lambdaReturnType scan_lam) $ - TC.bad $ TC.TypeError $ "Scan function returns type " ++ - prettyTuple (lambdaReturnType scan_lam) ++ " but neutral element has type " ++ - prettyTuple scan_t - return scan_nes' + scan_nes' <- fmap concat $ + forM scans $ \(Scan scan_lam scan_nes) -> do + scan_nes' <- mapM TC.checkArg scan_nes + let scan_t = map TC.argType scan_nes' + TC.checkLambda scan_lam $ map TC.noArgAliases $ scan_nes' ++ scan_nes' + unless (scan_t == lambdaReturnType scan_lam) $ + TC.bad $ + TC.TypeError $ + "Scan function returns type " + ++ prettyTuple (lambdaReturnType scan_lam) + ++ " but neutral element has type " + ++ prettyTuple scan_t + return scan_nes' - red_nes' <- fmap concat $ forM reds $ \(Reduce _ red_lam red_nes) -> do - red_nes' <- mapM TC.checkArg red_nes - let red_t = map TC.argType red_nes' - TC.checkLambda red_lam $ map TC.noArgAliases $ red_nes' ++ red_nes' - unless (red_t == lambdaReturnType red_lam) $ - TC.bad $ TC.TypeError $ "Reduce function returns type " ++ - prettyTuple (lambdaReturnType red_lam) ++ " but neutral element has type " ++ - prettyTuple red_t - return red_nes' + red_nes' <- fmap concat $ + forM reds $ \(Reduce _ red_lam red_nes) -> do + red_nes' <- mapM TC.checkArg red_nes + let red_t = map TC.argType red_nes' + TC.checkLambda red_lam $ map TC.noArgAliases $ red_nes' ++ red_nes' + unless (red_t == lambdaReturnType red_lam) $ + TC.bad $ + TC.TypeError $ + "Reduce function returns type " + ++ prettyTuple (lambdaReturnType red_lam) + ++ " but neutral element has type " + ++ prettyTuple red_t + return red_nes' let map_lam_ts = lambdaReturnType map_lam - unless (take (length scan_nes' + length red_nes') map_lam_ts == - map TC.argType (scan_nes'++ red_nes')) $ - TC.bad $ TC.TypeError $ "Map function return type " ++ prettyTuple map_lam_ts ++ - " wrong for given scan and reduction functions." + unless + ( take (length scan_nes' + length red_nes') map_lam_ts + == map TC.argType (scan_nes' ++ red_nes') + ) + $ TC.bad $ + TC.TypeError $ + "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 +getStreamAccums (Sequential accs) = accs instance OpMetrics (Op lore) => OpMetrics (SOAC lore) where opMetrics (Stream _ _ lam _) = @@ -658,61 +840,77 @@ opMetrics (Hist _len ops bucket_fun _imgs) = inside "Hist" $ mapM_ (lambdaMetrics . histOp) ops >> lambdaMetrics bucket_fun opMetrics (Screma _ (ScremaForm scans reds map_lam) _) = - inside "Screma" $ do mapM_ (lambdaMetrics . scanLambda) scans - mapM_ (lambdaMetrics . redLambda) reds - lambdaMetrics map_lam + inside "Screma" $ do + mapM_ (lambdaMetrics . scanLambda) scans + mapM_ (lambdaMetrics . redLambda) reds + lambdaMetrics map_lam instance PrettyLore lore => PP.Pretty (SOAC lore) where ppr (Stream size form lam arrs) = case form of - Parallel o comm lam0 acc -> - 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 - commasep ( PP.braces (commasep $ map ppr acc) : map ppr arrs )) - Sequential acc -> - text "streamSeq" <> - parens (ppr size <> comma ppr lam <> comma - commasep ( PP.braces (commasep $ map ppr acc) : map ppr arrs )) + Parallel o comm lam0 acc -> + 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 + commasep (PP.braces (commasep $ map ppr acc) : map ppr arrs) + ) + Sequential acc -> + 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)) (map (\(_, n, a) -> (n, a)) as) ppr (Hist len ops bucket_fun imgs) = ppHist len ops bucket_fun imgs ppr (Screma w (ScremaForm scans reds map_lam) arrs) - | null scans, null reds = - text "map" <> - parens (ppr w <> comma - ppr map_lam <> comma - commasep (map ppr arrs)) - + | null scans, + null reds = + text "map" + <> parens + ( ppr w <> comma + ppr map_lam <> comma + commasep (map ppr arrs) + ) | null scans = - text "redomap" <> - parens (ppr w <> comma - PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr reds) <> comma - ppr map_lam <> comma - commasep (map ppr arrs)) - + text "redomap" + <> parens + ( ppr w <> comma + PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr reds) <> comma + ppr map_lam <> comma + commasep (map ppr arrs) + ) | null reds = - text "scanomap" <> - parens (ppr w <> comma - PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr scans) <> comma - ppr map_lam <> comma - commasep (map ppr arrs)) - + text "scanomap" + <> parens + ( ppr w <> comma + PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr scans) <> comma + ppr map_lam <> comma + commasep (map ppr arrs) + ) ppr (Screma w form arrs) = ppScrema w form arrs -- | Prettyprint the given Screma. -ppScrema :: (PrettyLore lore, Pretty inp) => - SubExp -> ScremaForm lore -> [inp] -> Doc +ppScrema :: + (PrettyLore lore, Pretty inp) => + SubExp -> + ScremaForm lore -> + [inp] -> + Doc ppScrema w (ScremaForm scans reds map_lam) arrs = - text "screma" <> - parens (ppr w <> comma - PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr scans) <> comma - PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr reds) <> comma - ppr map_lam <> comma - commasep (map ppr arrs)) + text "screma" + <> parens + ( ppr w <> comma + PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr scans) <> comma + PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppr reds) <> comma + ppr map_lam <> comma + commasep (map ppr arrs) + ) instance PrettyLore lore => Pretty (Scan lore) where ppr (Scan scan_lam scan_nes) = @@ -724,31 +922,50 @@ instance PrettyLore lore => Pretty (Reduce lore) where ppr (Reduce comm red_lam red_nes) = - ppComm comm <> ppr red_lam <> comma - PP.braces (commasep $ map ppr red_nes) + ppComm comm <> ppr red_lam <> comma + PP.braces (commasep $ map ppr red_nes) -- | Prettyprint the given histogram operation. -ppHist :: (PrettyLore lore, Pretty inp) => - SubExp -> [HistOp lore] -> Lambda lore -> [inp] -> Doc +ppHist :: + (PrettyLore lore, Pretty inp) => + SubExp -> + [HistOp lore] -> + Lambda lore -> + [inp] -> + Doc ppHist len ops bucket_fun imgs = - text "hist" <> - parens (ppr len <> comma - PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppOp ops) <> comma - ppr bucket_fun <> comma - commasep (map ppr imgs)) - where ppOp (HistOp w rf dests nes op) = - ppr w <> comma <+> ppr rf <> comma <+> PP.braces (commasep $ map ppr dests) <> comma - PP.braces (commasep $ map ppr nes) <> comma ppr op + text "hist" + <> parens + ( ppr len <> comma + PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppOp ops) <> comma + ppr bucket_fun <> comma + commasep (map ppr imgs) + ) + where + ppOp (HistOp w rf dests nes op) = + ppr w <> comma <+> ppr rf <> comma <+> PP.braces (commasep $ map ppr dests) <> comma + PP.braces (commasep $ map ppr nes) <> comma + ppr op -ppSOAC :: (Pretty fn, Pretty v) => - String -> SubExp -> [fn] -> Maybe [SubExp] -> [v] -> Doc +ppSOAC :: + (Pretty fn, Pretty v) => + String -> + SubExp -> + [fn] -> + Maybe [SubExp] -> + [v] -> + Doc ppSOAC name size funs es as = - text name <> parens (ppr size <> comma - ppList funs - commasep (es' ++ map ppr as)) - where es' = maybe [] ((:[]) . ppTuple') es + text name + <> parens + ( ppr size <> comma + ppList funs + commasep (es' ++ map ppr as) + ) + where + es' = maybe [] ((: []) . ppTuple') es ppList :: Pretty a => [a] -> Doc ppList as = case map ppr as of - [] -> mempty - a':as' -> foldl () (a' <> comma) $ map (<> comma) as' + [] -> mempty + a' : as' -> foldl () (a' <> comma) $ map (<> comma) as' 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 @@ -1,810 +1,945 @@ -{-# OPTIONS_GHC -fno-warn-orphans #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -module Futhark.IR.SOACS.Simplify - ( simplifySOACS - , simplifyLambda - , simplifyFun - , simplifyStms - , simplifyConsts - - , simpleSOACS - , simplifySOAC - - , soacRules - - , HasSOAC(..) - , simplifyKnownIterationSOAC - , removeReplicateMapping - , liftIdentityMapping - - , SOACS - ) -where - -import Control.Monad -import Control.Monad.Identity -import Control.Monad.State -import Control.Monad.Writer -import Data.Foldable -import Data.Either -import Data.List (partition, transpose, unzip6, zip6) -import Data.Maybe -import qualified Data.Map.Strict as M -import qualified Data.Set as S - -import Futhark.IR.SOACS -import qualified Futhark.IR as AST -import Futhark.IR.Prop.Aliases -import qualified Futhark.Optimise.Simplify.Engine as Engine -import qualified Futhark.Optimise.Simplify as Simplify -import Futhark.Optimise.Simplify.Rules -import Futhark.MonadFreshNames -import Futhark.Optimise.Simplify.Rule -import Futhark.Optimise.Simplify.ClosedForm -import Futhark.Optimise.Simplify.Lore -import Futhark.Tools -import Futhark.Pass -import qualified Futhark.Analysis.SymbolTable as ST -import qualified Futhark.Analysis.UsageTable as UT -import Futhark.Analysis.DataDependencies -import Futhark.Transform.Rename -import Futhark.Util - -simpleSOACS :: Simplify.SimpleOps SOACS -simpleSOACS = Simplify.bindableSimpleOps simplifySOAC - -simplifySOACS :: Prog SOACS -> PassM (Prog SOACS) -simplifySOACS = - Simplify.simplifyProg simpleSOACS soacRules Engine.noExtraHoistBlockers - -simplifyFun :: MonadFreshNames m => - ST.SymbolTable (Wise SOACS) -> FunDef SOACS -> m (FunDef SOACS) -simplifyFun = - Simplify.simplifyFun simpleSOACS soacRules Engine.noExtraHoistBlockers - -simplifyLambda :: (HasScope SOACS m, MonadFreshNames m) => - Lambda -> m Lambda -simplifyLambda = - Simplify.simplifyLambda simpleSOACS soacRules Engine.noExtraHoistBlockers - -simplifyStms :: (HasScope SOACS m, MonadFreshNames m) => - Stms SOACS -> m (ST.SymbolTable (Wise SOACS), Stms SOACS) -simplifyStms stms = do - scope <- askScope - Simplify.simplifyStms simpleSOACS soacRules Engine.noExtraHoistBlockers - scope stms - -simplifyConsts :: MonadFreshNames m => - Stms SOACS -> m (ST.SymbolTable (Wise SOACS), Stms SOACS) -simplifyConsts = - Simplify.simplifyStms simpleSOACS soacRules Engine.noExtraHoistBlockers mempty - -simplifySOAC :: Simplify.SimplifiableLore lore => - Simplify.SimplifyOp lore (SOAC lore) -simplifySOAC (Stream outerdim form lam arr) = do - outerdim' <- Engine.simplify outerdim - (form', form_hoisted) <- simplifyStreamForm form - arr' <- mapM Engine.simplify arr - (lam', lam_hoisted) <- Engine.simplifyLambda lam - return (Stream outerdim' form' lam' arr', form_hoisted <> lam_hoisted) - where simplifyStreamForm (Parallel o comm lam0 acc) = do - acc' <- mapM Engine.simplify acc - (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) - -simplifySOAC (Scatter len lam ivs as) = do - len' <- Engine.simplify len - (lam', hoisted) <- Engine.simplifyLambda lam - ivs' <- mapM Engine.simplify ivs - as' <- mapM Engine.simplify as - return (Scatter len' lam' ivs' as', hoisted) - -simplifySOAC (Hist w ops bfun imgs) = do - w' <- Engine.simplify w - (ops', hoisted) <- fmap unzip $ forM ops $ \(HistOp dests_w rf dests nes op) -> do - dests_w' <- Engine.simplify dests_w - rf' <- Engine.simplify rf - dests' <- Engine.simplify dests - nes' <- mapM Engine.simplify nes - (op', hoisted) <- Engine.simplifyLambda op - return (HistOp dests_w' rf' dests' nes' op', hoisted) - imgs' <- mapM Engine.simplify imgs - (bfun', bfun_hoisted) <- Engine.simplifyLambda bfun - return (Hist w' ops' bfun' imgs', mconcat hoisted <> bfun_hoisted) - -simplifySOAC (Screma w (ScremaForm scans reds map_lam) arrs) = do - (scans', scans_hoisted) <- fmap unzip $ forM scans $ \(Scan lam nes) -> do - (lam', hoisted) <- Engine.simplifyLambda lam - nes' <- Engine.simplify nes - return (Scan lam' nes', hoisted) - - (reds', reds_hoisted) <- fmap unzip $ forM reds $ \(Reduce comm lam nes) -> do - (lam', hoisted) <- Engine.simplifyLambda lam - nes' <- Engine.simplify nes - return (Reduce comm lam' nes', hoisted) - - (map_lam', map_lam_hoisted) <- Engine.simplifyLambda map_lam - - (,) <$> (Screma <$> Engine.simplify w <*> - pure (ScremaForm scans' reds' map_lam') <*> - Engine.simplify arrs) <*> - pure (mconcat scans_hoisted <> mconcat reds_hoisted <> map_lam_hoisted) - -instance BinderOps (Wise SOACS) where - -fixLambdaParams :: (MonadBinder m, Bindable (Lore m), BinderOps (Lore m)) => - AST.Lambda (Lore m) -> [Maybe SubExp] -> m (AST.Lambda (Lore m)) -fixLambdaParams lam fixes = do - body <- runBodyBinder $ localScope (scopeOfLParams $ lambdaParams lam) $ do - zipWithM_ maybeFix (lambdaParams lam) fixes' - return $ lambdaBody lam - return lam { lambdaBody = body - , lambdaParams = map fst $ filter (isNothing . snd) $ - zip (lambdaParams lam) fixes' } - where fixes' = fixes ++ repeat Nothing - maybeFix p (Just x) = letBindNames [paramName p] $ BasicOp $ SubExp x - maybeFix _ Nothing = return () - -removeLambdaResults :: [Bool] -> AST.Lambda lore -> AST.Lambda lore -removeLambdaResults keep lam = lam { lambdaBody = lam_body' - , lambdaReturnType = ret } - where keep' :: [a] -> [a] - keep' = map snd . filter fst . zip (keep ++ repeat True) - lam_body = lambdaBody lam - lam_body' = lam_body { bodyResult = keep' $ bodyResult lam_body } - ret = keep' $ lambdaReturnType lam - -soacRules :: RuleBook (Wise SOACS) -soacRules = standardRules <> ruleBook topDownRules bottomUpRules - --- | Does this lore contain 'SOAC's in its t'Op's? A lore must be an --- instance of this class for the simplification rules to work. -class HasSOAC lore where - asSOAC :: Op lore -> Maybe (SOAC lore) - soacOp :: SOAC lore -> Op lore - -instance HasSOAC (Wise SOACS) where - asSOAC = Just - soacOp = id - -topDownRules :: [TopDownRule (Wise SOACS)] -topDownRules = [RuleOp hoistCertificates, - RuleOp removeReplicateMapping, - RuleOp removeReplicateWrite, - RuleOp removeUnusedSOACInput, - RuleOp simplifyClosedFormReduce, - RuleOp simplifyKnownIterationSOAC, - RuleOp fuseConcatScatter, - RuleOp simplifyMapIota, - RuleOp moveTransformToInput - ] - -bottomUpRules :: [BottomUpRule (Wise SOACS)] -bottomUpRules = [RuleOp removeDeadMapping, - RuleOp removeDeadReduction, - RuleOp removeDeadWrite, - RuleBasicOp removeUnnecessaryCopy, - RuleOp liftIdentityMapping, - RuleOp liftIdentityStreaming, - RuleOp removeDuplicateMapOutput, - RuleOp mapOpToOp - ] - --- Any certificates attached to a trivial Stm in the body might as --- well be applied to the SOAC itself. -hoistCertificates :: TopDownRuleOp (Wise SOACS) -hoistCertificates vtable pat aux soac - | (soac', hoisted) <- runState (mapSOACM mapper soac) mempty, - hoisted /= mempty = - Simplify $ auxing aux $ certifying hoisted $ letBind pat $ Op soac' - where mapper = identitySOACMapper { mapOnSOACLambda = onLambda } - onLambda lam = do - stms' <- mapM onStm $ bodyStms $ lambdaBody lam - return lam { lambdaBody = - mkBody stms' $ bodyResult $ lambdaBody lam } - onStm (Let se_pat se_aux (BasicOp (SubExp se))) = do - let (invariant, variant) = - partition (`ST.elem` vtable) $ - unCertificates $ stmAuxCerts se_aux - se_aux' = se_aux { stmAuxCerts = Certificates variant } - modify (Certificates invariant<>) - return $ Let se_pat se_aux' $ BasicOp $ SubExp se - onStm stm = return stm -hoistCertificates _ _ _ _ = - Skip - -liftIdentityMapping :: forall lore. - (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => - BottomUpRuleOp (Wise lore) -liftIdentityMapping (_, usages) 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 - free = freeIn $ lambdaBody fun - rettype = lambdaReturnType fun - ses = bodyResult $ lambdaBody fun - - freeOrConst (Var v) = v `nameIn` free - freeOrConst Constant{} = True - - 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') - checkInvariance (outId, e, t) (invariant, mapresult, rettype') - | freeOrConst e = ((Pattern [] [outId], BasicOp $ Replicate (Shape [w]) e) : invariant, - mapresult, - rettype') - | otherwise = (invariant, - (outId, e) : mapresult, - t : rettype') - - case foldr checkInvariance ([], [], []) $ - zip3 (patternElements pat) ses rettype of - ([], _, _) -> Skip - (invariant, mapresult, rettype') -> Simplify $ do - let (pat', ses') = unzip mapresult - fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = ses' } - , lambdaReturnType = rettype' - } - mapM_ (uncurry letBind) invariant - auxing aux $ - letBindNames (map patElemName pat') $ Op $ soacOp $ Screma w (mapSOAC fun') arrs -liftIdentityMapping _ _ _ _ = Skip - -liftIdentityStreaming :: BottomUpRuleOp (Wise SOACS) -liftIdentityStreaming _ (Pattern [] pes) aux (Stream w form lam arrs) - | (variant_map, invariant_map) <- - partitionEithers $ map isInvariantRes $ zip3 map_ts map_pes map_res, - not $ null invariant_map = Simplify $ do - - forM_ invariant_map $ \(pe, arr) -> - letBind (Pattern [] [pe]) $ BasicOp $ Copy arr - - let (variant_map_ts, variant_map_pes, variant_map_res) = unzip3 variant_map - lam' = lam { lambdaBody = (lambdaBody lam) { bodyResult = fold_res ++ variant_map_res } - , lambdaReturnType = fold_ts ++ variant_map_ts } - - auxing aux $ letBind (Pattern [] $ fold_pes ++ variant_map_pes) $ - Op $ Stream w form lam' arrs - where num_folds = length $ getStreamAccums form - (fold_pes, map_pes) = splitAt num_folds pes - (fold_ts, map_ts) = splitAt num_folds $ lambdaReturnType lam - lam_res = bodyResult $ lambdaBody lam - (fold_res, map_res) = splitAt num_folds lam_res - params_to_arrs = zip (map paramName $ drop (1 + num_folds) $ lambdaParams lam) arrs - - isInvariantRes (_, pe, Var v) - | Just arr <- lookup v params_to_arrs = - Right (pe, arr) - isInvariantRes x = - Left x -liftIdentityStreaming _ _ _ _ = Skip - --- | Remove all arguments to the map that are simply replicates. --- These can be turned into free variables instead. -removeReplicateMapping :: (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => - TopDownRuleOp (Wise lore) -removeReplicateMapping vtable pat aux op - | Just (Screma w form arrs) <- asSOAC op, - Just fun <- isMapSOAC form, - Just (bnds, fun', arrs') <- removeReplicateInput vtable fun arrs = Simplify $ do - forM_ bnds $ \(vs,cs,e) -> certifying cs $ letBindNames vs e - auxing aux $ letBind pat $ Op $ soacOp $ Screma w (mapSOAC fun') arrs' -removeReplicateMapping _ _ _ _ = Skip - --- | Like 'removeReplicateMapping', but for 'Scatter'. -removeReplicateWrite :: TopDownRuleOp (Wise SOACS) -removeReplicateWrite vtable pat aux (Scatter len lam ivs as) - | Just (bnds, lam', ivs') <- removeReplicateInput vtable lam ivs = Simplify $ do - forM_ bnds $ \(vs,cs,e) -> certifying cs $ letBindNames vs e - auxing aux $ letBind pat $ Op $ Scatter len lam' ivs' as -removeReplicateWrite _ _ _ _ = Skip - -removeReplicateInput :: Aliased lore => - ST.SymbolTable lore - -> AST.Lambda lore -> [VName] - -> Maybe ([([VName], Certificates, AST.Exp lore)], - AST.Lambda lore, [VName]) -removeReplicateInput vtable fun arrs - | not $ null parameterBnds = do - let (arr_params', arrs') = unzip params_and_arrs - fun' = fun { lambdaParams = acc_params <> arr_params' } - return (parameterBnds, fun', arrs') - | otherwise = Nothing - - where params = lambdaParams fun - (acc_params, arr_params) = - splitAt (length params - length arrs) params - (params_and_arrs, parameterBnds) = - partitionEithers $ zipWith isReplicateAndNotConsumed arr_params arrs - - isReplicateAndNotConsumed p v - | Just (BasicOp (Replicate (Shape (_:ds)) e), v_cs) <- - ST.lookupExp v vtable, - not $ paramName p `nameIn` consumedByLambda fun = - Right ([paramName p], - v_cs, - case ds of - [] -> BasicOp $ SubExp e - _ -> BasicOp $ Replicate (Shape ds) e) - | otherwise = - Left (p, v) - --- | Remove inputs that are not used inside the SOAC. -removeUnusedSOACInput :: TopDownRuleOp (Wise SOACS) -removeUnusedSOACInput _ pat aux (Screma w (ScremaForm scan reduce map_lam) arrs) - | (used,unused) <- partition usedInput params_and_arrs, - not (null unused) = Simplify $ do - let (used_params, used_arrs) = unzip used - map_lam' = map_lam { lambdaParams = used_params } - auxing aux $ letBind pat $ Op $ Screma w (ScremaForm scan reduce map_lam') used_arrs - where params_and_arrs = zip (lambdaParams map_lam) arrs - used_in_body = freeIn $ lambdaBody map_lam - usedInput (param, _) = paramName param `nameIn` used_in_body -removeUnusedSOACInput _ _ _ _ = Skip - -removeDeadMapping :: BottomUpRuleOp (Wise SOACS) -removeDeadMapping (_, used) pat aux (Screma w form arrs) - | Just fun <- isMapSOAC form = - let ses = bodyResult $ lambdaBody fun - isUsed (bindee, _, _) = (`UT.used` used) $ patElemName bindee - (pat',ses', ts') = unzip3 $ filter isUsed $ - zip3 (patternElements pat) ses $ lambdaReturnType fun - fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = ses' } - , lambdaReturnType = ts' - } - in if pat /= Pattern [] pat' - then Simplify $ auxing aux $ - letBind (Pattern [] pat') $ Op $ Screma w (mapSOAC fun') arrs - else Skip -removeDeadMapping _ _ _ _ = Skip - -removeDuplicateMapOutput :: BottomUpRuleOp (Wise SOACS) -removeDuplicateMapOutput (_, used) pat aux (Screma w form arrs) - | Just fun <- isMapSOAC form = - let ses = bodyResult $ lambdaBody fun - ts = lambdaReturnType fun - pes = patternValueElements pat - ses_ts_pes = zip3 ses ts pes - (ses_ts_pes', copies) = - foldl checkForDuplicates (mempty,mempty) ses_ts_pes - in if null copies then Skip - else Simplify $ do - let (ses', ts', pes') = unzip3 ses_ts_pes' - pat' = Pattern [] pes' - fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = ses' } - , lambdaReturnType = ts' } - 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 - where checkForDuplicates (ses_ts_pes',copies) (se,t,pe) - | Just (_,_,pe') <- find (\(x,_,_) -> x == se) ses_ts_pes' = - -- This subexp has been returned before, producing the - -- array pe'. - (ses_ts_pes', (pe', pe) : copies) - | otherwise = (ses_ts_pes' ++ [(se,t,pe)], copies) -removeDuplicateMapOutput _ _ _ _ = Skip - --- Mapping some operations becomes an extension of that operation. -mapOpToOp :: BottomUpRuleOp (Wise SOACS) - -mapOpToOp (_, used) pat aux1 e - | Just (map_pe, cs, w, BasicOp (Reshape newshape reshape_arr), [p], [arr]) <- - isMapWithOp pat e, - paramName p == reshape_arr, - not $ UT.isConsumed (patElemName map_pe) used = Simplify $ do - let redim | isJust $ shapeCoercion newshape = DimCoercion w - | otherwise = DimNew w - certifying (stmAuxCerts aux1 <> cs) $ letBind pat $ - BasicOp $ Reshape (redim : newshape) arr - - | Just (_, cs, _, - BasicOp (Concat d arr arrs dw), ps, outer_arr : outer_arrs) <- - isMapWithOp pat e, - (arr:arrs) == map paramName ps = - Simplify $ certifying (stmAuxCerts aux1 <> cs) $ letBind pat $ - BasicOp $ Concat (d+1) outer_arr outer_arrs dw - - | Just (map_pe, cs, _, - BasicOp (Rearrange perm rearrange_arr), [p], [arr]) <- - isMapWithOp pat e, - paramName p == rearrange_arr, - not $ UT.isConsumed (patElemName map_pe) used = - Simplify $ certifying (stmAuxCerts aux1 <> cs) $ letBind pat $ - BasicOp $ Rearrange (0 : map (1+) perm) arr - - | Just (map_pe, cs, _, BasicOp (Rotate rots rotate_arr), [p], [arr]) <- - isMapWithOp pat e, - paramName p == rotate_arr, - not $ UT.isConsumed (patElemName map_pe) used = - Simplify $ certifying (stmAuxCerts aux1 <> cs) $ letBind pat $ - BasicOp $ Rotate (intConst Int32 0 : rots) arr - -mapOpToOp _ _ _ _ = Skip - -isMapWithOp :: PatternT dec - -> SOAC (Wise SOACS) - -> Maybe (PatElemT dec, Certificates, SubExp, - AST.Exp (Wise SOACS), [Param Type], [VName]) -isMapWithOp pat e - | Pattern [] [map_pe] <- pat, - Screma w form arrs <- e, - Just map_lam <- isMapSOAC form, - [Let (Pattern [] [pe]) aux2 e'] <- - stmsToList $ bodyStms $ lambdaBody map_lam, - [Var r] <- bodyResult $ lambdaBody map_lam, - r == patElemName pe = - Just (map_pe, stmAuxCerts aux2, w, e', lambdaParams map_lam, arrs) - | otherwise = Nothing - --- | Some of the results of a reduction (or really: Redomap) may be --- dead. We remove them here. The trick is that we need to look at --- the data dependencies to see that the "dead" result is not --- actually used for computing one of the live ones. -removeDeadReduction :: BottomUpRuleOp (Wise SOACS) -removeDeadReduction (_, used) pat aux (Screma w form arrs) - | Just ([Reduce comm redlam nes], maplam) <- isRedomapSOAC form, - not $ all (`UT.used` used) $ patternNames pat, -- Quick/cheap check - - let (red_pes, map_pes) = splitAt (length nes) $ patternElements pat, - let redlam_deps = dataDependencies $ lambdaBody redlam, - let redlam_res = bodyResult $ lambdaBody redlam, - let redlam_params = lambdaParams redlam, - let used_after = map snd $ filter ((`UT.used` used) . patElemName . fst) $ - zip red_pes redlam_params, - let necessary = findNecessaryForReturned (`elem` used_after) - (zip redlam_params $ redlam_res <> redlam_res) redlam_deps, - let alive_mask = map ((`nameIn` necessary) . paramName) redlam_params, - - not $ all (==True) alive_mask = Simplify $ do - - let fixDeadToNeutral lives ne = if lives then Nothing else Just ne - dead_fix = zipWith fixDeadToNeutral alive_mask nes - (used_red_pes, _, used_nes) = - unzip3 $ filter (\(_,x,_) -> paramName x `nameIn` necessary) $ - zip3 red_pes redlam_params nes - - let maplam' = removeLambdaResults (take (length nes) alive_mask) maplam - redlam' <- removeLambdaResults (take (length nes) alive_mask) <$> fixLambdaParams redlam (dead_fix++dead_fix) - - auxing aux $ letBind (Pattern [] $ used_red_pes ++ map_pes) $ - Op $ Screma w (redomapSOAC [Reduce comm redlam' used_nes] maplam') arrs - -removeDeadReduction _ _ _ _ = Skip - --- | If we are writing to an array that is never used, get rid of it. -removeDeadWrite :: BottomUpRuleOp (Wise SOACS) -removeDeadWrite (_, used) pat aux (Scatter w fun arrs dests) = - let (i_ses, v_ses) = splitAt (length dests) $ bodyResult $ lambdaBody fun - (i_ts, v_ts) = splitAt (length dests) $ lambdaReturnType fun - isUsed (bindee, _, _, _, _, _) = (`UT.used` used) $ patElemName bindee - (pat', i_ses', v_ses', i_ts', v_ts', dests') = - unzip6 $ filter isUsed $ - zip6 (patternElements pat) i_ses v_ses i_ts v_ts dests - fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = i_ses' ++ v_ses' } - , lambdaReturnType = i_ts' ++ v_ts' - } - in if pat /= Pattern [] pat' - then Simplify $ auxing aux $ - letBind (Pattern [] pat') $ Op $ Scatter w fun' arrs dests' - else Skip -removeDeadWrite _ _ _ _ = Skip - --- handles now concatenation of more than two arrays -fuseConcatScatter :: TopDownRuleOp (Wise SOACS) -fuseConcatScatter vtable pat _ (Scatter _ fun arrs dests) - | Just (ws@(w':_), xss, css) <- unzip3 <$> mapM isConcat arrs, - xivs <- transpose xss, - all (w'==) ws = Simplify $ do - let r = length xivs - fun2s <- mapM (\_ -> renameLambda fun) [1 .. r-1] - let fun_n = length $ lambdaReturnType fun - (fun_is, fun_vs) = unzip $ map (splitAt (fun_n `div` 2) . - bodyResult . lambdaBody ) (fun:fun2s) - (its, vts) = unzip $ replicate r $ - splitAt (fun_n `div` 2) $ lambdaReturnType fun - new_stmts = mconcat $ map (bodyStms . lambdaBody) (fun:fun2s) - let fun' = Lambda - { lambdaParams = mconcat $ map lambdaParams (fun:fun2s) - , lambdaBody = mkBody new_stmts $ - mix fun_is <> mix fun_vs - , lambdaReturnType = mix its <> mix vts - } - certifying (mconcat css) $ - letBind pat $ Op $ Scatter w' fun' (concat xivs) $ map (incWrites r) dests - where sizeOf :: VName -> Maybe SubExp - sizeOf x = arraySize 0 . typeOf <$> ST.lookup x vtable - mix = concat . transpose - incWrites r (w, n, a) = (w, n*r, a) -- ToDO: is it (n*r) or (n+r-1)?? - isConcat v = case ST.lookupExp v vtable of - Just (BasicOp (Concat 0 x ys _), cs) -> do - x_w <- sizeOf x - y_ws<- mapM sizeOf ys - guard $ all (x_w==) y_ws - return (x_w, x:ys, cs) - Just (BasicOp (Reshape reshape arr), cs) -> do - guard $ isJust $ shapeCoercion reshape - (a, b, cs') <- isConcat arr - return (a, b, cs <> cs') - _ -> Nothing - -fuseConcatScatter _ _ _ _ = Skip - -simplifyClosedFormReduce :: TopDownRuleOp (Wise SOACS) -simplifyClosedFormReduce _ pat _ (Screma (Constant w) form _) - | Just nes <- concatMap redNeutral . fst <$> isRedomapSOAC form, - zeroIsh w = - Simplify $ forM_ (zip (patternNames pat) nes) $ \(v, ne) -> - letBindNames [v] $ BasicOp $ SubExp ne -simplifyClosedFormReduce vtable pat _ (Screma _ form arrs) - | Just [Reduce _ red_fun nes] <- isReduceSOAC form = - Simplify $ foldClosedForm (`ST.lookupExp` vtable) pat red_fun nes arrs -simplifyClosedFormReduce _ _ _ _ = Skip - --- For now we just remove singleton SOACs. -simplifyKnownIterationSOAC :: (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => - TopDownRuleOp (Wise lore) -simplifyKnownIterationSOAC _ pat _ op - | Just (Screma (Constant k) (ScremaForm scans reds map_lam) arrs) <- asSOAC op, - oneIsh k = Simplify $ do - - let (Reduce _ red_lam red_nes) = singleReduce reds - (Scan scan_lam scan_nes) = singleScan scans - (scan_pes, red_pes, map_pes) = splitAt3 (length scan_nes) (length red_nes) $ - patternElements pat - bindMapParam p a = do - a_t <- lookupType a - letBindNames [paramName p] $ - BasicOp $ Index a $ fullSlice a_t [DimFix $ constant (0::Int32)] - bindArrayResult pe se = - letBindNames [patElemName pe] $ - BasicOp $ ArrayLit [se] $ rowType $ patElemType pe - bindResult pe se = - letBindNames [patElemName pe] $ BasicOp $ SubExp se - - zipWithM_ bindMapParam (lambdaParams map_lam) arrs - (to_scan, to_red, map_res) <- splitAt3 (length scan_nes) (length red_nes) <$> - bodyBind (lambdaBody map_lam) - scan_res <- eLambda scan_lam $ map eSubExp $ scan_nes ++ to_scan - red_res <- eLambda red_lam $ map eSubExp $ red_nes ++ to_red - - zipWithM_ bindArrayResult scan_pes scan_res - 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, - oneIsh k = Simplify $ do - let nes = getStreamAccums form - (chunk_param, acc_params, slice_params) = - partitionChunkedFoldParameters (length nes) (lambdaParams fold_lam) - - letBindNames [paramName chunk_param] $ - BasicOp $ SubExp $ intConst Int32 1 - - forM_ (zip acc_params nes) $ \(p, ne) -> - letBindNames [paramName p] $ BasicOp $ SubExp ne - - forM_ (zip slice_params arrs) $ \(p, arr) -> - letBindNames [paramName p] $ BasicOp $ SubExp $ Var arr - - res <- bodyBind $ lambdaBody fold_lam - - forM_ (zip (patternNames pat) res) $ \(v, se) -> - letBindNames [v] $ BasicOp $ SubExp se - -simplifyKnownIterationSOAC _ _ _ _ = Skip - -data ArrayOp = ArrayIndexing Certificates VName (Slice SubExp) - | ArrayRearrange Certificates VName [Int] - | ArrayRotate Certificates VName [SubExp] - | ArrayVar Certificates VName -- ^ Never constructed. - deriving (Eq, Ord, Show) - -arrayOpArr :: ArrayOp -> VName -arrayOpArr (ArrayIndexing _ arr _) = arr -arrayOpArr (ArrayRearrange _ arr _) = arr -arrayOpArr (ArrayRotate _ arr _) = arr -arrayOpArr (ArrayVar _ arr) = arr - -arrayOpCerts :: ArrayOp -> Certificates -arrayOpCerts (ArrayIndexing cs _ _) = cs -arrayOpCerts (ArrayRearrange cs _ _) = cs -arrayOpCerts (ArrayRotate cs _ _) = cs -arrayOpCerts (ArrayVar cs _) = cs - -isArrayOp :: Certificates -> AST.Exp (Wise SOACS) -> Maybe ArrayOp -isArrayOp cs (BasicOp (Index arr slice)) = - Just $ ArrayIndexing cs arr slice -isArrayOp cs (BasicOp (Rearrange perm arr)) = - Just $ ArrayRearrange cs arr perm -isArrayOp cs (BasicOp (Rotate rots arr)) = - Just $ ArrayRotate cs arr rots -isArrayOp _ _ = - Nothing - -fromArrayOp :: ArrayOp -> (Certificates, AST.Exp (Wise SOACS)) -fromArrayOp (ArrayIndexing cs arr slice) = (cs, BasicOp $ Index arr slice) -fromArrayOp (ArrayRearrange cs arr perm) = (cs, BasicOp $ Rearrange perm arr) -fromArrayOp (ArrayRotate cs arr rots) = (cs, BasicOp $ Rotate rots arr) -fromArrayOp (ArrayVar cs arr) = (cs, BasicOp $ SubExp $ Var arr) - -arrayOps :: AST.Body (Wise SOACS) -> S.Set (AST.Pattern (Wise SOACS), ArrayOp) -arrayOps = mconcat . map onStm . stmsToList . bodyStms - where onStm (Let pat aux e) = - case isArrayOp (stmAuxCerts aux) e of - Just op -> S.singleton (pat, op) - Nothing -> execState (walkExpM walker e) mempty - onOp = execWriter . mapSOACM identitySOACMapper { mapOnSOACLambda = onLambda } - onLambda lam = do tell $ arrayOps $ lambdaBody lam - return lam - walker = identityWalker { walkOnBody = const $ modify . (<>) . arrayOps - , walkOnOp = modify . (<>) . onOp } - -replaceArrayOps :: M.Map ArrayOp ArrayOp - -> AST.Body (Wise SOACS) -> AST.Body (Wise SOACS) -replaceArrayOps substs (Body _ stms res) = - mkBody (fmap onStm stms) res - where onStm (Let pat aux e) = - let (cs', e') = onExp (stmAuxCerts aux) e - in certify cs' $ - mkLet' (patternContextIdents pat) (patternValueIdents pat) aux e' - onExp cs e - | Just op <- isArrayOp cs e, - Just op' <- M.lookup op substs = - fromArrayOp op' - onExp cs e = (cs, mapExp mapper e) - mapper = identityMapper { mapOnBody = const $ return . replaceArrayOps substs - , mapOnOp = return . onOp } - onOp = runIdentity . mapSOACM identitySOACMapper { mapOnSOACLambda = return . onLambda } - onLambda lam = lam { lambdaBody = replaceArrayOps substs $ lambdaBody lam } - --- Turn --- --- map (\i -> ... xs[i] ...) (iota n) --- --- into --- --- map (\i x -> ... x ...) (iota n) xs --- --- This is not because we want to encourage the map-iota pattern, but --- it may be present in generated code. This is an unfortunately --- expensive simplification rule, since it requires multiple passes --- over the entire lambda body. It only handles the very simplest --- case - if you find yourself planning to extend it to handle more --- complex situations (rotate or whatnot), consider turning it into a --- separate compiler pass instead. -simplifyMapIota :: TopDownRuleOp (Wise SOACS) -simplifyMapIota vtable pat aux (Screma w (ScremaForm scan reduce map_lam) arrs) - | Just (p, _) <- find isIota (zip (lambdaParams map_lam) arrs), - indexings <- filter (indexesWith (paramName p)) $ map snd $ S.toList $ - arrayOps $ lambdaBody map_lam, - not $ null indexings = Simplify $ do - -- For each indexing with iota, add the corresponding array to - -- the Screma, and construct a new lambda parameter. - (more_arrs, more_params, replacements) <- - unzip3 . catMaybes <$> mapM mapOverArr indexings - let substs = M.fromList $ zip indexings replacements - map_lam' = map_lam { lambdaParams = lambdaParams map_lam <> more_params - , lambdaBody = replaceArrayOps substs $ - lambdaBody map_lam - } - - auxing aux $ - letBind pat $ Op $ Screma w (ScremaForm scan reduce map_lam') (arrs <> more_arrs) - where isIota (_, arr) = case ST.lookupBasicOp arr vtable of - Just (Iota _ (Constant o) (Constant s) _, _) -> - zeroIsh o && oneIsh s - _ -> False - - indexesWith v (ArrayIndexing cs arr (DimFix (Var i) : _)) - | arr `ST.elem` vtable, - all (`ST.elem` vtable) $ unCertificates cs = - i == v - indexesWith _ _ = False - - mapOverArr (ArrayIndexing cs arr slice) = do - arr_elem <- newVName $ baseString arr ++ "_elem" - arr_t <- lookupType arr - arr' <- if arraySize 0 arr_t == w - then return arr - else certifying cs $ letExp (baseString arr ++ "_prefix") $ - BasicOp $ Index arr $ - fullSlice arr_t [DimSlice (intConst Int32 0) w (intConst Int32 1)] - return $ Just (arr', - Param arr_elem (rowType arr_t), - ArrayIndexing cs arr_elem (drop 1 slice)) - - mapOverArr _ = return Nothing - -simplifyMapIota _ _ _ _ = Skip - --- If a Screma's map function contains a transformation --- (e.g. transpose) on a parameter, create a new parameter --- corresponding to that transformation performed on the rows of the --- full array. -moveTransformToInput :: TopDownRuleOp (Wise SOACS) -moveTransformToInput vtable pat aux (Screma w (ScremaForm scan reduce map_lam) arrs) - | ops <- map snd $ filter arrayIsMapParam $ S.toList $ arrayOps $ lambdaBody map_lam, - not $ null ops = Simplify $ do - (more_arrs, more_params, replacements) <- - unzip3 . catMaybes <$> mapM mapOverArr ops - - when (null more_arrs) cannotSimplify - - let substs = M.fromList $ zip ops replacements - map_lam' = map_lam { lambdaParams = lambdaParams map_lam <> more_params - , lambdaBody = replaceArrayOps substs $ - lambdaBody map_lam - } - - auxing aux $ - letBind pat $ Op $ Screma w (ScremaForm scan reduce map_lam') (arrs <> more_arrs) - - where map_param_names = map paramName (lambdaParams map_lam) - topLevelPattern = (`elem` fmap stmPattern (bodyStms (lambdaBody map_lam))) - onlyUsedOnce arr = - case filter ((arr `nameIn`) . freeIn) $ stmsToList $ bodyStms $ lambdaBody map_lam of - _ : _ : _ -> False - _ -> True - - -- It's not just about whether the array is a parameter; - -- everything else must be map-invariant. - arrayIsMapParam (pat', ArrayIndexing cs arr slice) = - arr `elem` map_param_names && - all (`ST.elem` vtable) (namesToList $ freeIn cs <> freeIn slice) && - not (null slice) && - (not (null $ sliceDims slice) || (topLevelPattern pat' && onlyUsedOnce arr)) - arrayIsMapParam (_, ArrayRearrange cs arr perm) = - arr `elem` map_param_names && - all (`ST.elem` vtable) (namesToList $ freeIn cs) && - not (null perm) - arrayIsMapParam (_, ArrayRotate cs arr rots) = - arr `elem` map_param_names && - all (`ST.elem` vtable) (namesToList $ freeIn cs <> freeIn rots) - arrayIsMapParam (_, ArrayVar{}) = - False - - mapOverArr op - | Just (_, arr) <- find ((==arrayOpArr op) . fst) (zip map_param_names arrs) = do - arr_t <- lookupType arr - let whole_dim = DimSlice (intConst Int32 0) (arraySize 0 arr_t) (intConst Int32 1) - arr_transformed <- certifying (arrayOpCerts op) $ - letExp (baseString arr ++ "_transformed") $ - case op of - ArrayIndexing _ _ slice -> - BasicOp $ Index arr $ whole_dim : slice - ArrayRearrange _ _ perm -> - BasicOp $ Rearrange (0 : map (+1) perm) arr - ArrayRotate _ _ rots -> - BasicOp $ Rotate (intConst Int32 0 : rots) arr - ArrayVar{} -> - BasicOp $ SubExp $ Var arr - arr_transformed_t <- lookupType arr_transformed - arr_transformed_row <- newVName $ baseString arr ++ "_transformed_row" - return $ Just (arr_transformed, - Param arr_transformed_row (rowType arr_transformed_t), - ArrayVar mempty arr_transformed_row) - - mapOverArr _ = return Nothing - +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Futhark.IR.SOACS.Simplify + ( simplifySOACS, + simplifyLambda, + simplifyFun, + simplifyStms, + simplifyConsts, + simpleSOACS, + simplifySOAC, + soacRules, + HasSOAC (..), + simplifyKnownIterationSOAC, + removeReplicateMapping, + liftIdentityMapping, + SOACS, + ) +where + +import Control.Monad +import Control.Monad.Identity +import Control.Monad.State +import Control.Monad.Writer +import Data.Either +import Data.Foldable +import Data.List (partition, transpose, unzip6, zip6) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import Futhark.Analysis.DataDependencies +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import qualified Futhark.IR as AST +import Futhark.IR.Prop.Aliases +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.Pass +import Futhark.Tools +import Futhark.Transform.Rename +import Futhark.Util + +simpleSOACS :: Simplify.SimpleOps SOACS +simpleSOACS = Simplify.bindableSimpleOps simplifySOAC + +simplifySOACS :: Prog SOACS -> PassM (Prog SOACS) +simplifySOACS = + Simplify.simplifyProg simpleSOACS soacRules Engine.noExtraHoistBlockers + +simplifyFun :: + MonadFreshNames m => + ST.SymbolTable (Wise SOACS) -> + FunDef SOACS -> + m (FunDef SOACS) +simplifyFun = + Simplify.simplifyFun simpleSOACS soacRules Engine.noExtraHoistBlockers + +simplifyLambda :: + (HasScope SOACS m, MonadFreshNames m) => + Lambda -> + m Lambda +simplifyLambda = + Simplify.simplifyLambda simpleSOACS soacRules Engine.noExtraHoistBlockers + +simplifyStms :: + (HasScope SOACS m, MonadFreshNames m) => + Stms SOACS -> + m (ST.SymbolTable (Wise SOACS), Stms SOACS) +simplifyStms stms = do + scope <- askScope + Simplify.simplifyStms + simpleSOACS + soacRules + Engine.noExtraHoistBlockers + scope + stms + +simplifyConsts :: + MonadFreshNames m => + Stms SOACS -> + m (ST.SymbolTable (Wise SOACS), Stms SOACS) +simplifyConsts = + Simplify.simplifyStms simpleSOACS soacRules Engine.noExtraHoistBlockers mempty + +simplifySOAC :: + Simplify.SimplifiableLore lore => + Simplify.SimplifyOp lore (SOAC lore) +simplifySOAC (Stream outerdim form lam arr) = do + outerdim' <- Engine.simplify outerdim + (form', form_hoisted) <- simplifyStreamForm form + arr' <- mapM Engine.simplify arr + (lam', lam_hoisted) <- Engine.simplifyLambda lam + return (Stream outerdim' form' lam' arr', form_hoisted <> lam_hoisted) + where + simplifyStreamForm (Parallel o comm lam0 acc) = do + acc' <- mapM Engine.simplify acc + (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) +simplifySOAC (Scatter len lam ivs as) = do + len' <- Engine.simplify len + (lam', hoisted) <- Engine.simplifyLambda lam + ivs' <- mapM Engine.simplify ivs + as' <- mapM Engine.simplify as + return (Scatter len' lam' ivs' as', hoisted) +simplifySOAC (Hist w ops bfun imgs) = do + w' <- Engine.simplify w + (ops', hoisted) <- fmap unzip $ + forM ops $ \(HistOp dests_w rf dests nes op) -> do + dests_w' <- Engine.simplify dests_w + rf' <- Engine.simplify rf + dests' <- Engine.simplify dests + nes' <- mapM Engine.simplify nes + (op', hoisted) <- Engine.simplifyLambda op + return (HistOp dests_w' rf' dests' nes' op', hoisted) + imgs' <- mapM Engine.simplify imgs + (bfun', bfun_hoisted) <- Engine.simplifyLambda bfun + return (Hist w' ops' bfun' imgs', mconcat hoisted <> bfun_hoisted) +simplifySOAC (Screma w (ScremaForm scans reds map_lam) arrs) = do + (scans', scans_hoisted) <- fmap unzip $ + forM scans $ \(Scan lam nes) -> do + (lam', hoisted) <- Engine.simplifyLambda lam + nes' <- Engine.simplify nes + return (Scan lam' nes', hoisted) + + (reds', reds_hoisted) <- fmap unzip $ + forM reds $ \(Reduce comm lam nes) -> do + (lam', hoisted) <- Engine.simplifyLambda lam + nes' <- Engine.simplify nes + return (Reduce comm lam' nes', hoisted) + + (map_lam', map_lam_hoisted) <- Engine.simplifyLambda map_lam + + (,) + <$> ( Screma <$> Engine.simplify w + <*> pure (ScremaForm scans' reds' map_lam') + <*> Engine.simplify arrs + ) + <*> pure (mconcat scans_hoisted <> mconcat reds_hoisted <> map_lam_hoisted) + +instance BinderOps (Wise SOACS) + +fixLambdaParams :: + (MonadBinder m, Bindable (Lore m), BinderOps (Lore m)) => + AST.Lambda (Lore m) -> + [Maybe SubExp] -> + m (AST.Lambda (Lore m)) +fixLambdaParams lam fixes = do + body <- runBodyBinder $ + localScope (scopeOfLParams $ lambdaParams lam) $ do + zipWithM_ maybeFix (lambdaParams lam) fixes' + return $ lambdaBody lam + return + lam + { lambdaBody = body, + lambdaParams = + map fst $ + filter (isNothing . snd) $ + zip (lambdaParams lam) fixes' + } + where + fixes' = fixes ++ repeat Nothing + maybeFix p (Just x) = letBindNames [paramName p] $ BasicOp $ SubExp x + maybeFix _ Nothing = return () + +removeLambdaResults :: [Bool] -> AST.Lambda lore -> AST.Lambda lore +removeLambdaResults keep lam = + lam + { lambdaBody = lam_body', + lambdaReturnType = ret + } + where + keep' :: [a] -> [a] + keep' = map snd . filter fst . zip (keep ++ repeat True) + lam_body = lambdaBody lam + lam_body' = lam_body {bodyResult = keep' $ bodyResult lam_body} + ret = keep' $ lambdaReturnType lam + +soacRules :: RuleBook (Wise SOACS) +soacRules = standardRules <> ruleBook topDownRules bottomUpRules + +-- | Does this lore contain 'SOAC's in its t'Op's? A lore must be an +-- instance of this class for the simplification rules to work. +class HasSOAC lore where + asSOAC :: Op lore -> Maybe (SOAC lore) + soacOp :: SOAC lore -> Op lore + +instance HasSOAC (Wise SOACS) where + asSOAC = Just + soacOp = id + +topDownRules :: [TopDownRule (Wise SOACS)] +topDownRules = + [ RuleOp hoistCertificates, + RuleOp removeReplicateMapping, + RuleOp removeReplicateWrite, + RuleOp removeUnusedSOACInput, + RuleOp simplifyClosedFormReduce, + RuleOp simplifyKnownIterationSOAC, + RuleOp fuseConcatScatter, + RuleOp simplifyMapIota, + RuleOp moveTransformToInput + ] + +bottomUpRules :: [BottomUpRule (Wise SOACS)] +bottomUpRules = + [ RuleOp removeDeadMapping, + RuleOp removeDeadReduction, + RuleOp removeDeadWrite, + RuleBasicOp removeUnnecessaryCopy, + RuleOp liftIdentityMapping, + RuleOp liftIdentityStreaming, + RuleOp removeDuplicateMapOutput, + RuleOp mapOpToOp + ] + +-- Any certificates attached to a trivial Stm in the body might as +-- well be applied to the SOAC itself. +hoistCertificates :: TopDownRuleOp (Wise SOACS) +hoistCertificates vtable pat aux soac + | (soac', hoisted) <- runState (mapSOACM mapper soac) mempty, + hoisted /= mempty = + Simplify $ auxing aux $ certifying hoisted $ letBind pat $ Op soac' + where + mapper = identitySOACMapper {mapOnSOACLambda = onLambda} + onLambda lam = do + stms' <- mapM onStm $ bodyStms $ lambdaBody lam + return + lam + { lambdaBody = + mkBody stms' $ bodyResult $ lambdaBody lam + } + onStm (Let se_pat se_aux (BasicOp (SubExp se))) = do + let (invariant, variant) = + partition (`ST.elem` vtable) $ + unCertificates $ stmAuxCerts se_aux + se_aux' = se_aux {stmAuxCerts = Certificates variant} + modify (Certificates invariant <>) + return $ Let se_pat se_aux' $ BasicOp $ SubExp se + onStm stm = return stm +hoistCertificates _ _ _ _ = + Skip + +liftIdentityMapping :: + forall lore. + (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => + BottomUpRuleOp (Wise lore) +liftIdentityMapping (_, usages) 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 + free = freeIn $ lambdaBody fun + rettype = lambdaReturnType fun + ses = bodyResult $ lambdaBody fun + + freeOrConst (Var v) = v `nameIn` free + freeOrConst Constant {} = True + + 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' + ) + checkInvariance (outId, e, t) (invariant, mapresult, rettype') + | freeOrConst e = + ( (Pattern [] [outId], BasicOp $ Replicate (Shape [w]) e) : invariant, + mapresult, + rettype' + ) + | otherwise = + ( invariant, + (outId, e) : mapresult, + t : rettype' + ) + + case foldr checkInvariance ([], [], []) $ + zip3 (patternElements pat) ses rettype of + ([], _, _) -> Skip + (invariant, mapresult, rettype') -> Simplify $ do + let (pat', ses') = unzip mapresult + fun' = + fun + { lambdaBody = (lambdaBody fun) {bodyResult = ses'}, + lambdaReturnType = rettype' + } + mapM_ (uncurry letBind) invariant + auxing aux $ + letBindNames (map patElemName pat') $ Op $ soacOp $ Screma w (mapSOAC fun') arrs +liftIdentityMapping _ _ _ _ = Skip + +liftIdentityStreaming :: BottomUpRuleOp (Wise SOACS) +liftIdentityStreaming _ (Pattern [] pes) aux (Stream w form lam arrs) + | (variant_map, invariant_map) <- + partitionEithers $ map isInvariantRes $ zip3 map_ts map_pes map_res, + not $ null invariant_map = Simplify $ do + forM_ invariant_map $ \(pe, arr) -> + letBind (Pattern [] [pe]) $ BasicOp $ Copy arr + + let (variant_map_ts, variant_map_pes, variant_map_res) = unzip3 variant_map + lam' = + lam + { lambdaBody = (lambdaBody lam) {bodyResult = fold_res ++ variant_map_res}, + lambdaReturnType = fold_ts ++ variant_map_ts + } + + auxing aux $ + letBind (Pattern [] $ fold_pes ++ variant_map_pes) $ + Op $ Stream w form lam' arrs + where + num_folds = length $ getStreamAccums form + (fold_pes, map_pes) = splitAt num_folds pes + (fold_ts, map_ts) = splitAt num_folds $ lambdaReturnType lam + lam_res = bodyResult $ lambdaBody lam + (fold_res, map_res) = splitAt num_folds lam_res + params_to_arrs = zip (map paramName $ drop (1 + num_folds) $ lambdaParams lam) arrs + + isInvariantRes (_, pe, Var v) + | Just arr <- lookup v params_to_arrs = + Right (pe, arr) + isInvariantRes x = + Left x +liftIdentityStreaming _ _ _ _ = Skip + +-- | Remove all arguments to the map that are simply replicates. +-- These can be turned into free variables instead. +removeReplicateMapping :: + (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => + TopDownRuleOp (Wise lore) +removeReplicateMapping vtable pat aux op + | Just (Screma w form arrs) <- asSOAC op, + Just fun <- isMapSOAC form, + Just (bnds, fun', arrs') <- removeReplicateInput vtable fun arrs = Simplify $ do + forM_ bnds $ \(vs, cs, e) -> certifying cs $ letBindNames vs e + auxing aux $ letBind pat $ Op $ soacOp $ Screma w (mapSOAC fun') arrs' +removeReplicateMapping _ _ _ _ = Skip + +-- | Like 'removeReplicateMapping', but for 'Scatter'. +removeReplicateWrite :: TopDownRuleOp (Wise SOACS) +removeReplicateWrite vtable pat aux (Scatter len lam ivs as) + | Just (bnds, lam', ivs') <- removeReplicateInput vtable lam ivs = Simplify $ do + forM_ bnds $ \(vs, cs, e) -> certifying cs $ letBindNames vs e + auxing aux $ letBind pat $ Op $ Scatter len lam' ivs' as +removeReplicateWrite _ _ _ _ = Skip + +removeReplicateInput :: + Aliased lore => + ST.SymbolTable lore -> + AST.Lambda lore -> + [VName] -> + Maybe + ( [([VName], Certificates, AST.Exp lore)], + AST.Lambda lore, + [VName] + ) +removeReplicateInput vtable fun arrs + | not $ null parameterBnds = do + let (arr_params', arrs') = unzip params_and_arrs + fun' = fun {lambdaParams = acc_params <> arr_params'} + return (parameterBnds, fun', arrs') + | otherwise = Nothing + where + params = lambdaParams fun + (acc_params, arr_params) = + splitAt (length params - length arrs) params + (params_and_arrs, parameterBnds) = + partitionEithers $ zipWith isReplicateAndNotConsumed arr_params arrs + + isReplicateAndNotConsumed p v + | Just (BasicOp (Replicate (Shape (_ : ds)) e), v_cs) <- + ST.lookupExp v vtable, + not $ paramName p `nameIn` consumedByLambda fun = + Right + ( [paramName p], + v_cs, + case ds of + [] -> BasicOp $ SubExp e + _ -> BasicOp $ Replicate (Shape ds) e + ) + | otherwise = + Left (p, v) + +-- | Remove inputs that are not used inside the SOAC. +removeUnusedSOACInput :: TopDownRuleOp (Wise SOACS) +removeUnusedSOACInput _ pat aux (Screma w (ScremaForm scan reduce map_lam) arrs) + | (used, unused) <- partition usedInput params_and_arrs, + not (null unused) = Simplify $ do + let (used_params, used_arrs) = unzip used + map_lam' = map_lam {lambdaParams = used_params} + auxing aux $ letBind pat $ Op $ Screma w (ScremaForm scan reduce map_lam') used_arrs + where + params_and_arrs = zip (lambdaParams map_lam) arrs + used_in_body = freeIn $ lambdaBody map_lam + usedInput (param, _) = paramName param `nameIn` used_in_body +removeUnusedSOACInput _ _ _ _ = Skip + +removeDeadMapping :: BottomUpRuleOp (Wise SOACS) +removeDeadMapping (_, used) pat aux (Screma w form arrs) + | Just fun <- isMapSOAC form = + let ses = bodyResult $ lambdaBody fun + isUsed (bindee, _, _) = (`UT.used` used) $ patElemName bindee + (pat', ses', ts') = + unzip3 $ + filter isUsed $ + zip3 (patternElements pat) ses $ lambdaReturnType fun + fun' = + fun + { lambdaBody = (lambdaBody fun) {bodyResult = ses'}, + lambdaReturnType = ts' + } + in if pat /= Pattern [] pat' + then + Simplify $ + auxing aux $ + letBind (Pattern [] pat') $ Op $ Screma w (mapSOAC fun') arrs + else Skip +removeDeadMapping _ _ _ _ = Skip + +removeDuplicateMapOutput :: BottomUpRuleOp (Wise SOACS) +removeDuplicateMapOutput (_, used) pat aux (Screma w form arrs) + | Just fun <- isMapSOAC form = + let ses = bodyResult $ lambdaBody fun + ts = lambdaReturnType fun + pes = patternValueElements pat + ses_ts_pes = zip3 ses ts pes + (ses_ts_pes', copies) = + foldl checkForDuplicates (mempty, mempty) ses_ts_pes + in if null copies + then Skip + else Simplify $ do + let (ses', ts', pes') = unzip3 ses_ts_pes' + pat' = Pattern [] pes' + fun' = + fun + { lambdaBody = (lambdaBody fun) {bodyResult = ses'}, + lambdaReturnType = ts' + } + 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 + where + checkForDuplicates (ses_ts_pes', copies) (se, t, pe) + | Just (_, _, pe') <- find (\(x, _, _) -> x == se) ses_ts_pes' = + -- This subexp has been returned before, producing the + -- array pe'. + (ses_ts_pes', (pe', pe) : copies) + | otherwise = (ses_ts_pes' ++ [(se, t, pe)], copies) +removeDuplicateMapOutput _ _ _ _ = Skip + +-- Mapping some operations becomes an extension of that operation. +mapOpToOp :: BottomUpRuleOp (Wise SOACS) +mapOpToOp (_, used) pat aux1 e + | Just (map_pe, cs, w, BasicOp (Reshape newshape reshape_arr), [p], [arr]) <- + isMapWithOp pat e, + paramName p == reshape_arr, + not $ UT.isConsumed (patElemName map_pe) used = Simplify $ do + let redim + | isJust $ shapeCoercion newshape = DimCoercion w + | otherwise = DimNew w + certifying (stmAuxCerts aux1 <> cs) $ + letBind pat $ + BasicOp $ Reshape (redim : newshape) arr + | Just + ( _, + cs, + _, + BasicOp (Concat d arr arrs dw), + ps, + outer_arr : outer_arrs + ) <- + isMapWithOp pat e, + (arr : arrs) == map paramName ps = + Simplify $ + certifying (stmAuxCerts aux1 <> cs) $ + letBind pat $ + BasicOp $ Concat (d + 1) outer_arr outer_arrs dw + | Just + ( map_pe, + cs, + _, + BasicOp (Rearrange perm rearrange_arr), + [p], + [arr] + ) <- + isMapWithOp pat e, + paramName p == rearrange_arr, + not $ UT.isConsumed (patElemName map_pe) used = + Simplify $ + certifying (stmAuxCerts aux1 <> cs) $ + letBind pat $ + BasicOp $ Rearrange (0 : map (1 +) perm) arr + | Just (map_pe, cs, _, BasicOp (Rotate rots rotate_arr), [p], [arr]) <- + isMapWithOp pat e, + paramName p == rotate_arr, + not $ UT.isConsumed (patElemName map_pe) used = + Simplify $ + certifying (stmAuxCerts aux1 <> cs) $ + letBind pat $ + BasicOp $ Rotate (intConst Int64 0 : rots) arr +mapOpToOp _ _ _ _ = Skip + +isMapWithOp :: + PatternT dec -> + SOAC (Wise SOACS) -> + Maybe + ( PatElemT dec, + Certificates, + SubExp, + AST.Exp (Wise SOACS), + [Param Type], + [VName] + ) +isMapWithOp pat e + | Pattern [] [map_pe] <- pat, + Screma w form arrs <- e, + Just map_lam <- isMapSOAC form, + [Let (Pattern [] [pe]) aux2 e'] <- + stmsToList $ bodyStms $ lambdaBody map_lam, + [Var r] <- bodyResult $ lambdaBody map_lam, + r == patElemName pe = + Just (map_pe, stmAuxCerts aux2, w, e', lambdaParams map_lam, arrs) + | otherwise = Nothing + +-- | Some of the results of a reduction (or really: Redomap) may be +-- dead. We remove them here. The trick is that we need to look at +-- the data dependencies to see that the "dead" result is not +-- actually used for computing one of the live ones. +removeDeadReduction :: BottomUpRuleOp (Wise SOACS) +removeDeadReduction (_, used) pat aux (Screma w form arrs) + | Just ([Reduce comm redlam nes], maplam) <- isRedomapSOAC form, + not $ all (`UT.used` used) $ patternNames pat, -- Quick/cheap check + let (red_pes, map_pes) = splitAt (length nes) $ patternElements pat, + let redlam_deps = dataDependencies $ lambdaBody redlam, + let redlam_res = bodyResult $ lambdaBody redlam, + let redlam_params = lambdaParams redlam, + let used_after = + map snd $ + filter ((`UT.used` used) . patElemName . fst) $ + zip red_pes redlam_params, + let necessary = + findNecessaryForReturned + (`elem` used_after) + (zip redlam_params $ redlam_res <> redlam_res) + redlam_deps, + let alive_mask = map ((`nameIn` necessary) . paramName) redlam_params, + not $ all (== True) alive_mask = Simplify $ do + let fixDeadToNeutral lives ne = if lives then Nothing else Just ne + dead_fix = zipWith fixDeadToNeutral alive_mask nes + (used_red_pes, _, used_nes) = + unzip3 $ + filter (\(_, x, _) -> paramName x `nameIn` necessary) $ + zip3 red_pes redlam_params nes + + let maplam' = removeLambdaResults (take (length nes) alive_mask) maplam + redlam' <- removeLambdaResults (take (length nes) alive_mask) <$> fixLambdaParams redlam (dead_fix ++ dead_fix) + + auxing aux $ + letBind (Pattern [] $ used_red_pes ++ map_pes) $ + Op $ Screma w (redomapSOAC [Reduce comm redlam' used_nes] maplam') arrs +removeDeadReduction _ _ _ _ = Skip + +-- | If we are writing to an array that is never used, get rid of it. +removeDeadWrite :: BottomUpRuleOp (Wise SOACS) +removeDeadWrite (_, used) pat aux (Scatter w fun arrs dests) = + let (i_ses, v_ses) = splitAt (length dests) $ bodyResult $ lambdaBody fun + (i_ts, v_ts) = splitAt (length dests) $ lambdaReturnType fun + isUsed (bindee, _, _, _, _, _) = (`UT.used` used) $ patElemName bindee + (pat', i_ses', v_ses', i_ts', v_ts', dests') = + unzip6 $ + filter isUsed $ + zip6 (patternElements pat) i_ses v_ses i_ts v_ts dests + fun' = + fun + { lambdaBody = (lambdaBody fun) {bodyResult = i_ses' ++ v_ses'}, + lambdaReturnType = i_ts' ++ v_ts' + } + in if pat /= Pattern [] pat' + then + Simplify $ + auxing aux $ + letBind (Pattern [] pat') $ Op $ Scatter w fun' arrs dests' + else Skip +removeDeadWrite _ _ _ _ = Skip + +-- handles now concatenation of more than two arrays +fuseConcatScatter :: TopDownRuleOp (Wise SOACS) +fuseConcatScatter vtable pat _ (Scatter _ fun arrs dests) + | Just (ws@(w' : _), xss, css) <- unzip3 <$> mapM isConcat arrs, + xivs <- transpose xss, + all (w' ==) ws = Simplify $ do + let r = length xivs + fun2s <- mapM (\_ -> renameLambda fun) [1 .. r -1] + let fun_n = length $ lambdaReturnType fun + (fun_is, fun_vs) = + unzip $ + map + ( splitAt (fun_n `div` 2) + . bodyResult + . lambdaBody + ) + (fun : fun2s) + (its, vts) = + unzip $ + replicate r $ + splitAt (fun_n `div` 2) $ lambdaReturnType fun + new_stmts = mconcat $ map (bodyStms . lambdaBody) (fun : fun2s) + let fun' = + Lambda + { lambdaParams = mconcat $ map lambdaParams (fun : fun2s), + lambdaBody = + mkBody new_stmts $ + mix fun_is <> mix fun_vs, + lambdaReturnType = mix its <> mix vts + } + certifying (mconcat css) $ + letBind pat $ Op $ Scatter w' fun' (concat xivs) $ map (incWrites r) dests + where + sizeOf :: VName -> Maybe SubExp + sizeOf x = arraySize 0 . typeOf <$> ST.lookup x vtable + mix = concat . transpose + incWrites r (w, n, a) = (w, n * r, a) -- ToDO: is it (n*r) or (n+r-1)?? + isConcat v = case ST.lookupExp v vtable of + Just (BasicOp (Concat 0 x ys _), cs) -> do + x_w <- sizeOf x + y_ws <- mapM sizeOf ys + guard $ all (x_w ==) y_ws + return (x_w, x : ys, cs) + Just (BasicOp (Reshape reshape arr), cs) -> do + guard $ isJust $ shapeCoercion reshape + (a, b, cs') <- isConcat arr + return (a, b, cs <> cs') + _ -> Nothing +fuseConcatScatter _ _ _ _ = Skip + +simplifyClosedFormReduce :: TopDownRuleOp (Wise SOACS) +simplifyClosedFormReduce _ pat _ (Screma (Constant w) form _) + | Just nes <- concatMap redNeutral . fst <$> isRedomapSOAC form, + zeroIsh w = + Simplify $ + forM_ (zip (patternNames pat) nes) $ \(v, ne) -> + letBindNames [v] $ BasicOp $ SubExp ne +simplifyClosedFormReduce vtable pat _ (Screma _ form arrs) + | Just [Reduce _ red_fun nes] <- isReduceSOAC form = + Simplify $ foldClosedForm (`ST.lookupExp` vtable) pat red_fun nes arrs +simplifyClosedFormReduce _ _ _ _ = Skip + +-- For now we just remove singleton SOACs. +simplifyKnownIterationSOAC :: + (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => + TopDownRuleOp (Wise lore) +simplifyKnownIterationSOAC _ pat _ op + | Just (Screma (Constant k) (ScremaForm scans reds map_lam) arrs) <- asSOAC op, + oneIsh k = Simplify $ do + let (Reduce _ red_lam red_nes) = singleReduce reds + (Scan scan_lam scan_nes) = singleScan scans + (scan_pes, red_pes, map_pes) = + splitAt3 (length scan_nes) (length red_nes) $ + patternElements pat + bindMapParam p a = do + a_t <- lookupType a + letBindNames [paramName p] $ + BasicOp $ Index a $ fullSlice a_t [DimFix $ constant (0 :: Int64)] + bindArrayResult pe se = + letBindNames [patElemName pe] $ + BasicOp $ ArrayLit [se] $ rowType $ patElemType pe + bindResult pe se = + letBindNames [patElemName pe] $ BasicOp $ SubExp se + + zipWithM_ bindMapParam (lambdaParams map_lam) arrs + (to_scan, to_red, map_res) <- + splitAt3 (length scan_nes) (length red_nes) + <$> bodyBind (lambdaBody map_lam) + scan_res <- eLambda scan_lam $ map eSubExp $ scan_nes ++ to_scan + red_res <- eLambda red_lam $ map eSubExp $ red_nes ++ to_red + + zipWithM_ bindArrayResult scan_pes scan_res + 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, + oneIsh k = Simplify $ do + let nes = getStreamAccums form + (chunk_param, acc_params, slice_params) = + partitionChunkedFoldParameters (length nes) (lambdaParams fold_lam) + + letBindNames [paramName chunk_param] $ + BasicOp $ SubExp $ intConst Int64 1 + + forM_ (zip acc_params nes) $ \(p, ne) -> + letBindNames [paramName p] $ BasicOp $ SubExp ne + + forM_ (zip slice_params arrs) $ \(p, arr) -> + letBindNames [paramName p] $ BasicOp $ SubExp $ Var arr + + res <- bodyBind $ lambdaBody fold_lam + + forM_ (zip (patternNames pat) res) $ \(v, se) -> + letBindNames [v] $ BasicOp $ SubExp se +simplifyKnownIterationSOAC _ _ _ _ = Skip + +data ArrayOp + = ArrayIndexing Certificates VName (Slice SubExp) + | ArrayRearrange Certificates VName [Int] + | ArrayRotate Certificates VName [SubExp] + | -- | Never constructed. + ArrayVar Certificates VName + deriving (Eq, Ord, Show) + +arrayOpArr :: ArrayOp -> VName +arrayOpArr (ArrayIndexing _ arr _) = arr +arrayOpArr (ArrayRearrange _ arr _) = arr +arrayOpArr (ArrayRotate _ arr _) = arr +arrayOpArr (ArrayVar _ arr) = arr + +arrayOpCerts :: ArrayOp -> Certificates +arrayOpCerts (ArrayIndexing cs _ _) = cs +arrayOpCerts (ArrayRearrange cs _ _) = cs +arrayOpCerts (ArrayRotate cs _ _) = cs +arrayOpCerts (ArrayVar cs _) = cs + +isArrayOp :: Certificates -> AST.Exp (Wise SOACS) -> Maybe ArrayOp +isArrayOp cs (BasicOp (Index arr slice)) = + Just $ ArrayIndexing cs arr slice +isArrayOp cs (BasicOp (Rearrange perm arr)) = + Just $ ArrayRearrange cs arr perm +isArrayOp cs (BasicOp (Rotate rots arr)) = + Just $ ArrayRotate cs arr rots +isArrayOp _ _ = + Nothing + +fromArrayOp :: ArrayOp -> (Certificates, AST.Exp (Wise SOACS)) +fromArrayOp (ArrayIndexing cs arr slice) = (cs, BasicOp $ Index arr slice) +fromArrayOp (ArrayRearrange cs arr perm) = (cs, BasicOp $ Rearrange perm arr) +fromArrayOp (ArrayRotate cs arr rots) = (cs, BasicOp $ Rotate rots arr) +fromArrayOp (ArrayVar cs arr) = (cs, BasicOp $ SubExp $ Var arr) + +arrayOps :: AST.Body (Wise SOACS) -> S.Set (AST.Pattern (Wise SOACS), ArrayOp) +arrayOps = mconcat . map onStm . stmsToList . bodyStms + where + onStm (Let pat aux e) = + case isArrayOp (stmAuxCerts aux) e of + Just op -> S.singleton (pat, op) + Nothing -> execState (walkExpM walker e) mempty + onOp = execWriter . mapSOACM identitySOACMapper {mapOnSOACLambda = onLambda} + onLambda lam = do + tell $ arrayOps $ lambdaBody lam + return lam + walker = + identityWalker + { walkOnBody = const $ modify . (<>) . arrayOps, + walkOnOp = modify . (<>) . onOp + } + +replaceArrayOps :: + M.Map ArrayOp ArrayOp -> + AST.Body (Wise SOACS) -> + AST.Body (Wise SOACS) +replaceArrayOps substs (Body _ stms res) = + mkBody (fmap onStm stms) res + where + onStm (Let pat aux e) = + let (cs', e') = onExp (stmAuxCerts aux) e + in certify cs' $ + mkLet' (patternContextIdents pat) (patternValueIdents pat) aux e' + onExp cs e + | Just op <- isArrayOp cs e, + Just op' <- M.lookup op substs = + fromArrayOp op' + onExp cs e = (cs, mapExp mapper e) + mapper = + identityMapper + { mapOnBody = const $ return . replaceArrayOps substs, + mapOnOp = return . onOp + } + onOp = runIdentity . mapSOACM identitySOACMapper {mapOnSOACLambda = return . onLambda} + onLambda lam = lam {lambdaBody = replaceArrayOps substs $ lambdaBody lam} + +-- Turn +-- +-- map (\i -> ... xs[i] ...) (iota n) +-- +-- into +-- +-- map (\i x -> ... x ...) (iota n) xs +-- +-- This is not because we want to encourage the map-iota pattern, but +-- it may be present in generated code. This is an unfortunately +-- expensive simplification rule, since it requires multiple passes +-- over the entire lambda body. It only handles the very simplest +-- case - if you find yourself planning to extend it to handle more +-- complex situations (rotate or whatnot), consider turning it into a +-- separate compiler pass instead. +simplifyMapIota :: TopDownRuleOp (Wise SOACS) +simplifyMapIota vtable pat aux (Screma w (ScremaForm scan reduce map_lam) arrs) + | Just (p, _) <- find isIota (zip (lambdaParams map_lam) arrs), + indexings <- + filter (indexesWith (paramName p)) $ + map snd $ + S.toList $ + arrayOps $ lambdaBody map_lam, + not $ null indexings = Simplify $ do + -- For each indexing with iota, add the corresponding array to + -- the Screma, and construct a new lambda parameter. + (more_arrs, more_params, replacements) <- + unzip3 . catMaybes <$> mapM mapOverArr indexings + let substs = M.fromList $ zip indexings replacements + map_lam' = + map_lam + { lambdaParams = lambdaParams map_lam <> more_params, + lambdaBody = + replaceArrayOps substs $ + lambdaBody map_lam + } + + auxing aux $ + letBind pat $ Op $ Screma w (ScremaForm scan reduce map_lam') (arrs <> more_arrs) + where + isIota (_, arr) = case ST.lookupBasicOp arr vtable of + Just (Iota _ (Constant o) (Constant s) _, _) -> + zeroIsh o && oneIsh s + _ -> False + + indexesWith v (ArrayIndexing cs arr (DimFix (Var i) : _)) + | arr `ST.elem` vtable, + all (`ST.elem` vtable) $ unCertificates cs = + i == v + indexesWith _ _ = False + + mapOverArr (ArrayIndexing cs arr slice) = do + arr_elem <- newVName $ baseString arr ++ "_elem" + arr_t <- lookupType arr + arr' <- + if arraySize 0 arr_t == w + then return arr + else + certifying cs $ + letExp (baseString arr ++ "_prefix") $ + BasicOp $ + Index arr $ + fullSlice arr_t [DimSlice (intConst Int64 0) w (intConst Int64 1)] + return $ + Just + ( arr', + Param arr_elem (rowType arr_t), + ArrayIndexing cs arr_elem (drop 1 slice) + ) + mapOverArr _ = return Nothing +simplifyMapIota _ _ _ _ = Skip + +-- If a Screma's map function contains a transformation +-- (e.g. transpose) on a parameter, create a new parameter +-- corresponding to that transformation performed on the rows of the +-- full array. +moveTransformToInput :: TopDownRuleOp (Wise SOACS) +moveTransformToInput vtable pat aux (Screma w (ScremaForm scan reduce map_lam) arrs) + | ops <- map snd $ filter arrayIsMapParam $ S.toList $ arrayOps $ lambdaBody map_lam, + not $ null ops = Simplify $ do + (more_arrs, more_params, replacements) <- + unzip3 . catMaybes <$> mapM mapOverArr ops + + when (null more_arrs) cannotSimplify + + let substs = M.fromList $ zip ops replacements + map_lam' = + map_lam + { lambdaParams = lambdaParams map_lam <> more_params, + lambdaBody = + replaceArrayOps substs $ + lambdaBody map_lam + } + + auxing aux $ + letBind pat $ Op $ Screma w (ScremaForm scan reduce map_lam') (arrs <> more_arrs) + where + map_param_names = map paramName (lambdaParams map_lam) + topLevelPattern = (`elem` fmap stmPattern (bodyStms (lambdaBody map_lam))) + onlyUsedOnce arr = + case filter ((arr `nameIn`) . freeIn) $ stmsToList $ bodyStms $ lambdaBody map_lam of + _ : _ : _ -> False + _ -> True + + -- It's not just about whether the array is a parameter; + -- everything else must be map-invariant. + arrayIsMapParam (pat', ArrayIndexing cs arr slice) = + arr `elem` map_param_names + && all (`ST.elem` vtable) (namesToList $ freeIn cs <> freeIn slice) + && not (null slice) + && (not (null $ sliceDims slice) || (topLevelPattern pat' && onlyUsedOnce arr)) + arrayIsMapParam (_, ArrayRearrange cs arr perm) = + arr `elem` map_param_names + && all (`ST.elem` vtable) (namesToList $ freeIn cs) + && not (null perm) + arrayIsMapParam (_, ArrayRotate cs arr rots) = + arr `elem` map_param_names + && all (`ST.elem` vtable) (namesToList $ freeIn cs <> freeIn rots) + arrayIsMapParam (_, ArrayVar {}) = + False + + mapOverArr op + | Just (_, arr) <- find ((== arrayOpArr op) . fst) (zip map_param_names arrs) = do + arr_t <- lookupType arr + let whole_dim = DimSlice (intConst Int64 0) (arraySize 0 arr_t) (intConst Int64 1) + arr_transformed <- certifying (arrayOpCerts op) $ + letExp (baseString arr ++ "_transformed") $ + case op of + ArrayIndexing _ _ slice -> + BasicOp $ Index arr $ whole_dim : slice + ArrayRearrange _ _ perm -> + BasicOp $ Rearrange (0 : map (+ 1) perm) arr + ArrayRotate _ _ rots -> + BasicOp $ Rotate (intConst Int64 0 : rots) arr + ArrayVar {} -> + BasicOp $ SubExp $ Var arr + arr_transformed_t <- lookupType arr_transformed + arr_transformed_row <- newVName $ baseString arr ++ "_transformed_row" + return $ + Just + ( arr_transformed, + Param arr_transformed_row (rowType arr_transformed_t), + ArrayVar mempty arr_transformed_row + ) + mapOverArr _ = return Nothing moveTransformToInput _ _ _ _ = Skip 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 @@ -1,1145 +1,1445 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE ScopedTypeVariables #-} --- | Segmented operations. These correspond to perfect @map@ nests on --- top of /something/, except that the @map@s are conceptually only --- over @iota@s (so there will be explicit indexing inside them). -module Futhark.IR.SegOp - ( SegOp(..) - , SegVirt(..) - , segLevel - , segSpace - , typeCheckSegOp - , SegSpace(..) - , scopeOfSegSpace - , segSpaceDims - - -- * Details - , HistOp(..) - , histType - , SegBinOp(..) - , segBinOpResults - , segBinOpChunks - , KernelBody(..) - , aliasAnalyseKernelBody - , consumedInKernelBody - , ResultManifest(..) - , KernelResult(..) - , kernelResultSubExp - , SplitOrdering(..) - - -- ** Generic traversal - , SegOpMapper(..) - , identitySegOpMapper - , mapSegOpM - - -- * Simplification - , simplifySegOp - , HasSegOp(..) - , segOpRules - - -- * Memory - , segOpReturns - ) -where - -import Control.Monad.State.Strict -import Control.Monad.Writer hiding (mapM_) -import Control.Monad.Identity hiding (mapM_) -import Data.Bifunctor (first) -import qualified Data.Map.Strict as M -import Data.Maybe -import Data.List - (intersperse, foldl', partition, isPrefixOf, groupBy) - -import Futhark.IR -import qualified Futhark.Analysis.Alias as Alias -import qualified Futhark.Analysis.SymbolTable as ST -import qualified Futhark.Analysis.UsageTable as UT -import Futhark.Analysis.PrimExp.Convert -import qualified Futhark.Util.Pretty as PP -import Futhark.Util.Pretty - ((), (<+>), ppr, commasep, Pretty, parens, text) -import Futhark.Transform.Substitute -import Futhark.Transform.Rename -import Futhark.Optimise.Simplify.Lore -import Futhark.IR.Prop.Aliases -import Futhark.IR.Aliases - (Aliases, removeLambdaAliases, removeStmAliases) -import Futhark.IR.Mem -import qualified Futhark.TypeCheck as TC -import Futhark.Analysis.Metrics -import Futhark.Util (maybeNth, chunks) -import Futhark.Optimise.Simplify.Rule -import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Tools - --- | How an array is split into chunks. -data SplitOrdering = SplitContiguous - | SplitStrided SubExp - deriving (Eq, Ord, Show) - -instance FreeIn SplitOrdering where - freeIn' SplitContiguous = mempty - freeIn' (SplitStrided stride) = freeIn' stride - -instance Substitute SplitOrdering where - substituteNames _ SplitContiguous = - SplitContiguous - substituteNames subst (SplitStrided stride) = - SplitStrided $ substituteNames subst stride - -instance Rename SplitOrdering where - rename SplitContiguous = - pure SplitContiguous - rename (SplitStrided stride) = - SplitStrided <$> rename stride - --- | An operator for 'SegHist'. -data HistOp lore = - HistOp { histWidth :: SubExp - , histRaceFactor :: SubExp - , histDest :: [VName] - , histNeutral :: [SubExp] - , histShape :: Shape - -- ^ In case this operator is semantically a vectorised - -- operator (corresponding to a perfect map nest in the - -- SOACS representation), these are the logical - -- "dimensions". This is used to generate more efficient - -- code. - , histOp :: Lambda lore - } - deriving (Eq, Ord, Show) - --- | The type of a histogram produced by a 'HistOp'. This can be --- different from the type of the 'histDest's in case we are --- dealing with a segmented histogram. -histType :: HistOp lore -> [Type] -histType op = map ((`arrayOfRow` histWidth op) . - (`arrayOfShape` histShape op)) $ - lambdaReturnType $ histOp op - --- | An operator for 'SegScan' and 'SegRed'. -data SegBinOp lore = - SegBinOp { segBinOpComm :: Commutativity - , segBinOpLambda :: Lambda lore - , segBinOpNeutral :: [SubExp] - , segBinOpShape :: Shape - -- ^ In case this operator is semantically a vectorised - -- operator (corresponding to a perfect map nest in the - -- SOACS representation), these are the logical - -- "dimensions". This is used to generate more efficient - -- code. - } - deriving (Eq, Ord, Show) - --- | How many reduction results are produced by these 'SegBinOp's? -segBinOpResults :: [SegBinOp lore] -> Int -segBinOpResults = sum . map (length . segBinOpNeutral) - --- | Split some list into chunks equal to the number of values --- returned by each 'SegBinOp' -segBinOpChunks :: [SegBinOp lore] -> [a] -> [[a]] -segBinOpChunks = chunks . map (length . segBinOpNeutral) - --- | The body of a 'SegOp'. -data KernelBody lore = KernelBody { kernelBodyLore :: BodyDec lore - , kernelBodyStms :: Stms lore - , kernelBodyResult :: [KernelResult] - } - -deriving instance Decorations lore => Ord (KernelBody lore) -deriving instance Decorations lore => Show (KernelBody lore) -deriving instance Decorations lore => Eq (KernelBody lore) - --- | Metadata about whether there is a subtle point to this --- 'KernelResult'. This is used to protect things like tiling, which --- might otherwise be removed by the simplifier because they're --- semantically redundant. This has no semantic effect and can be --- ignored at code generation. -data ResultManifest - = ResultNoSimplify - -- ^ Don't simplify this one! - | ResultMaySimplify - -- ^ Go nuts. - | ResultPrivate - -- ^ The results produced are only used within the - -- same physical thread later on, and can thus be - -- kept in registers. - deriving (Eq, Show, Ord) - --- | A 'KernelBody' does not return an ordinary 'Result'. Instead, it --- returns a list of these. -data KernelResult = Returns ResultManifest SubExp - -- ^ Each "worker" in the kernel returns this. - -- Whether this is a result-per-thread or a - -- result-per-group depends on where the 'SegOp' occurs. - | WriteReturns - [SubExp] -- Size of array. Must match number of dims. - VName -- Which array - [(Slice SubExp, SubExp)] - -- Arbitrary number of index/value pairs. - | ConcatReturns - SplitOrdering -- Permuted? - SubExp -- The final size. - SubExp -- Per-thread/group (max) chunk size. - VName -- Chunk by this worker. - | TileReturns - [(SubExp, SubExp)] -- Total/tile for each dimension - VName -- Tile written by this worker. - -- The TileReturns must not expect more than one - -- result to be written per physical thread. - deriving (Eq, Show, Ord) - --- | Get the root t'SubExp' corresponding values for a 'KernelResult'. -kernelResultSubExp :: KernelResult -> SubExp -kernelResultSubExp (Returns _ se) = se -kernelResultSubExp (WriteReturns _ arr _) = Var arr -kernelResultSubExp (ConcatReturns _ _ _ v) = Var v -kernelResultSubExp (TileReturns _ v) = Var v - -instance FreeIn KernelResult where - freeIn' (Returns _ what) = freeIn' what - freeIn' (WriteReturns rws arr res) = freeIn' rws <> freeIn' arr <> freeIn' res - freeIn' (ConcatReturns o w per_thread_elems v) = - freeIn' o <> freeIn' w <> freeIn' per_thread_elems <> freeIn' v - freeIn' (TileReturns dims v) = - freeIn' dims <> freeIn' v - -instance ASTLore lore => FreeIn (KernelBody lore) where - freeIn' (KernelBody dec stms res) = - fvBind bound_in_stms $ freeIn' dec <> freeIn' stms <> freeIn' res - where bound_in_stms = foldMap boundByStm stms - -instance ASTLore lore => Substitute (KernelBody lore) where - substituteNames subst (KernelBody dec stms res) = - KernelBody - (substituteNames subst dec) - (substituteNames subst stms) - (substituteNames subst res) - -instance Substitute KernelResult where - substituteNames subst (Returns manifest se) = - Returns manifest (substituteNames subst se) - substituteNames subst (WriteReturns rws arr res) = - WriteReturns - (substituteNames subst rws) (substituteNames subst arr) - (substituteNames subst res) - substituteNames subst (ConcatReturns o w per_thread_elems v) = - ConcatReturns - (substituteNames subst o) - (substituteNames subst w) - (substituteNames subst per_thread_elems) - (substituteNames subst v) - substituteNames subst (TileReturns dims v) = - TileReturns (substituteNames subst dims) (substituteNames subst v) - -instance ASTLore lore => Rename (KernelBody lore) where - rename (KernelBody dec stms res) = do - dec' <- rename dec - renamingStms stms $ \stms' -> - KernelBody dec' stms' <$> rename res - -instance Rename KernelResult where - rename = substituteRename - --- | Perform alias analysis on a 'KernelBody'. -aliasAnalyseKernelBody :: (ASTLore lore, - CanBeAliased (Op lore)) => - KernelBody lore - -> KernelBody (Aliases lore) -aliasAnalyseKernelBody (KernelBody dec stms res) = - let Body dec' stms' _ = Alias.analyseBody mempty $ Body dec stms [] - in KernelBody dec' stms' res - -removeKernelBodyAliases :: CanBeAliased (Op lore) => - KernelBody (Aliases lore) -> KernelBody lore -removeKernelBodyAliases (KernelBody (_, dec) stms res) = - KernelBody dec (fmap removeStmAliases stms) res - -removeKernelBodyWisdom :: CanBeWise (Op lore) => - KernelBody (Wise lore) -> KernelBody lore -removeKernelBodyWisdom (KernelBody dec stms res) = - let Body dec' stms' _ = removeBodyWisdom $ Body dec stms [] - in KernelBody dec' stms' res - --- | The variables consumed in the kernel body. -consumedInKernelBody :: Aliased lore => - KernelBody lore -> Names -consumedInKernelBody (KernelBody dec stms res) = - consumedInBody (Body dec stms []) <> mconcat (map consumedByReturn res) - where consumedByReturn (WriteReturns _ a _) = oneName a - consumedByReturn _ = mempty - -checkKernelBody :: TC.Checkable lore => - [Type] -> KernelBody (Aliases lore) -> TC.TypeM lore () -checkKernelBody ts (KernelBody (_, dec) stms kres) = do - TC.checkBodyLore dec - TC.checkStms stms $ do - unless (length ts == length kres) $ - TC.bad $ TC.TypeError $ "Kernel return type is " ++ prettyTuple ts ++ - ", but body returns " ++ show (length kres) ++ " values." - zipWithM_ checkKernelResult kres ts - - where checkKernelResult (Returns _ what) t = - TC.require [t] what - checkKernelResult (WriteReturns rws arr res) t = do - mapM_ (TC.require [Prim int32]) rws - arr_t <- lookupType arr - forM_ res $ \(slice, e) -> do - mapM_ (traverse $ TC.require [Prim int32]) slice - TC.require [t] e - unless (arr_t == t `arrayOfShape` Shape rws) $ - TC.bad $ TC.TypeError $ "WriteReturns returning " ++ - pretty e ++ " of type " ++ pretty t ++ ", shape=" ++ pretty rws ++ - ", but destination array has type " ++ pretty arr_t - TC.consume =<< TC.lookupAliases arr - checkKernelResult (ConcatReturns o w per_thread_elems v) t = do - case o of - SplitContiguous -> return () - SplitStrided stride -> TC.require [Prim int32] stride - TC.require [Prim int32] w - TC.require [Prim int32] per_thread_elems - vt <- lookupType v - unless (vt == t `arrayOfRow` arraySize 0 vt) $ - TC.bad $ TC.TypeError $ "Invalid type for ConcatReturns " ++ pretty v - checkKernelResult (TileReturns dims v) t = do - forM_ dims $ \(dim, tile) -> do - TC.require [Prim int32] dim - TC.require [Prim int32] tile - vt <- lookupType v - unless (vt == t `arrayOfShape` Shape (map snd dims)) $ - TC.bad $ TC.TypeError $ "Invalid type for TileReturns " ++ pretty v - -kernelBodyMetrics :: OpMetrics (Op lore) => KernelBody lore -> MetricsM () -kernelBodyMetrics = mapM_ stmMetrics . kernelBodyStms - -instance PrettyLore lore => Pretty (KernelBody lore) where - ppr (KernelBody _ stms res) = - PP.stack (map ppr (stmsToList stms)) - text "return" <+> PP.braces (PP.commasep $ map ppr res) - -instance Pretty KernelResult where - ppr (Returns ResultNoSimplify what) = - text "returns (manifest)" <+> ppr what - ppr (Returns ResultPrivate what) = - text "returns (private)" <+> ppr what - ppr (Returns ResultMaySimplify what) = - text "returns" <+> ppr what - ppr (WriteReturns rws arr res) = - ppr arr <+> text "with" <+> PP.apply (map ppRes res) - where ppRes (is, e) = - PP.brackets (PP.commasep $ zipWith f is rws) <+> text "<-" <+> ppr e - f i rw = ppr i <+> text "<" <+> ppr rw - ppr (ConcatReturns o w per_thread_elems v) = - text "concat" <> suff <> - parens (commasep [ppr w, ppr per_thread_elems]) <+> ppr v - where suff = case o of SplitContiguous -> mempty - SplitStrided stride -> text "Strided" <> parens (ppr stride) - ppr (TileReturns dims v) = - text "tile" <> - parens (commasep $ map onDim dims) <+> ppr v - where onDim (dim, tile) = ppr dim <+> text "/" <+> ppr tile - - --- | Do we need group-virtualisation when generating code for the --- segmented operation? In most cases, we do, but for some simple --- kernels, we compute the full number of groups in advance, and then --- virtualisation is an unnecessary (but generally very small) --- overhead. This only really matters for fairly trivial but very --- wide @map@ kernels where each thread performs constant-time work on --- scalars. -data SegVirt - = SegVirt - | SegNoVirt - | SegNoVirtFull - -- ^ Not only do we not need virtualisation, but we _guarantee_ - -- that all physical threads participate in the work. This can - -- save some checks in code generation. - deriving (Eq, Ord, Show) - --- | Index space of a 'SegOp'. -data SegSpace = SegSpace { segFlat :: VName - -- ^ Flat physical index corresponding to the - -- dimensions (at code generation used for a - -- thread ID or similar). - , unSegSpace :: [(VName, SubExp)] - } - deriving (Eq, Ord, Show) - - --- | The sizes spanned by the indexes of the 'SegSpace'. -segSpaceDims :: SegSpace -> [SubExp] -segSpaceDims (SegSpace _ space) = map snd space - --- | A 'Scope' containing all the identifiers brought into scope by --- this 'SegSpace'. -scopeOfSegSpace :: SegSpace -> Scope lore -scopeOfSegSpace (SegSpace phys space) = - M.fromList $ zip (phys : map fst space) $ repeat $ IndexName Int32 - -checkSegSpace :: TC.Checkable lore => SegSpace -> TC.TypeM lore () -checkSegSpace (SegSpace _ dims) = - mapM_ (TC.require [Prim int32] . snd) dims - --- | A 'SegOp' is semantically a perfectly nested stack of maps, on --- top of some bottommost computation (scalar computation, reduction, --- scan, or histogram). The 'SegSpace' encodes the original map --- structure. --- --- All 'SegOp's are parameterised by the representation of their body, --- as well as a *level*. The *level* is a representation-specific bit --- of information. For example, in GPU backends, it is used to --- indicate whether the 'SegOp' is expected to run at the thread-level --- or the group-level. -data SegOp lvl lore - = SegMap lvl SegSpace [Type] (KernelBody lore) - | SegRed lvl SegSpace [SegBinOp lore] [Type] (KernelBody lore) - -- ^ The KernelSpace must always have at least two dimensions, - -- implying that the result of a SegRed is always an array. - | SegScan lvl SegSpace [SegBinOp lore] [Type] (KernelBody lore) - | SegHist lvl SegSpace [HistOp lore] [Type] (KernelBody lore) - deriving (Eq, Ord, Show) - --- | The level of a 'SegOp'. -segLevel :: SegOp lvl lore -> lvl -segLevel (SegMap lvl _ _ _) = lvl -segLevel (SegRed lvl _ _ _ _) = lvl -segLevel (SegScan lvl _ _ _ _) = lvl -segLevel (SegHist lvl _ _ _ _) = lvl - --- | The space of a 'SegOp'. -segSpace :: SegOp lvl lore -> SegSpace -segSpace (SegMap _ lvl _ _) = lvl -segSpace (SegRed _ lvl _ _ _) = lvl -segSpace (SegScan _ lvl _ _ _) = lvl -segSpace (SegHist _ lvl _ _ _) = lvl - -segResultShape :: SegSpace -> Type -> KernelResult -> Type -segResultShape _ t (WriteReturns rws _ _) = - t `arrayOfShape` Shape rws -segResultShape space t (Returns _ _) = - foldr (flip arrayOfRow) t $ segSpaceDims space -segResultShape _ t (ConcatReturns _ w _ _) = - t `arrayOfRow` w -segResultShape _ t (TileReturns dims _) = - t `arrayOfShape` Shape (map fst dims) - --- | The return type of a 'SegOp'. -segOpType :: SegOp lvl lore -> [Type] -segOpType (SegMap _ space ts kbody) = - zipWith (segResultShape space) ts $ kernelBodyResult kbody -segOpType (SegRed _ space reds ts kbody) = - red_ts ++ - zipWith (segResultShape space) map_ts - (drop (length red_ts) $ kernelBodyResult kbody) - where map_ts = drop (length red_ts) ts - segment_dims = init $ segSpaceDims space - red_ts = do - op <- reds - let shape = Shape segment_dims <> segBinOpShape op - map (`arrayOfShape` shape) (lambdaReturnType $ segBinOpLambda op) -segOpType (SegScan _ space scans ts kbody) = - scan_ts ++ - zipWith (segResultShape space) map_ts - (drop (length scan_ts) $ kernelBodyResult kbody) - where map_ts = drop (length scan_ts) ts - scan_ts = do - op <- scans - let shape = Shape (segSpaceDims space) <> segBinOpShape op - map (`arrayOfShape` shape) (lambdaReturnType $ segBinOpLambda op) -segOpType (SegHist _ space ops _ _) = do - op <- ops - let shape = Shape (segment_dims <> [histWidth op]) <> histShape op - map (`arrayOfShape` shape) (lambdaReturnType $ histOp op) - where dims = segSpaceDims space - segment_dims = init dims - -instance TypedOp (SegOp lvl lore) where - opType = pure . staticShapes . segOpType - -instance (ASTLore lore, Aliased lore, ASTConstraints lvl) => - AliasedOp (SegOp lvl lore) where - opAliases = map (const mempty) . segOpType - - consumedInOp (SegMap _ _ _ kbody) = - consumedInKernelBody kbody - consumedInOp (SegRed _ _ _ _ kbody) = - consumedInKernelBody kbody - consumedInOp (SegScan _ _ _ _ kbody) = - consumedInKernelBody kbody - consumedInOp (SegHist _ _ ops _ kbody) = - namesFromList (concatMap histDest ops) <> consumedInKernelBody kbody - --- | Type check a 'SegOp', given a checker for its level. -typeCheckSegOp :: TC.Checkable lore => - (lvl -> TC.TypeM lore ()) - -> SegOp lvl (Aliases lore) -> TC.TypeM lore () -typeCheckSegOp checkLvl (SegMap lvl space ts kbody) = do - checkLvl lvl - checkScanRed space [] ts kbody - -typeCheckSegOp checkLvl (SegRed lvl space reds ts body) = do - checkLvl lvl - checkScanRed space reds' ts body - where reds' = zip3 - (map segBinOpLambda reds) - (map segBinOpNeutral reds) - (map segBinOpShape reds) - -typeCheckSegOp checkLvl (SegScan lvl space scans ts body) = do - checkLvl lvl - checkScanRed space scans' ts body - where scans' = zip3 - (map segBinOpLambda scans) - (map segBinOpNeutral scans) - (map segBinOpShape scans) - -typeCheckSegOp checkLvl (SegHist lvl space ops ts kbody) = do - checkLvl lvl - checkSegSpace space - mapM_ TC.checkType ts - - TC.binding (scopeOfSegSpace space) $ do - nes_ts <- forM ops $ \(HistOp dest_w rf dests nes shape op) -> do - TC.require [Prim int32] dest_w - TC.require [Prim int32] rf - nes' <- mapM TC.checkArg nes - mapM_ (TC.require [Prim int32]) $ shapeDims shape - - -- Operator type must match the type of neutral elements. - let stripVecDims = stripArray $ shapeRank shape - TC.checkLambda op $ map (TC.noArgAliases . first stripVecDims) $ nes' ++ nes' - let nes_t = map TC.argType nes' - unless (nes_t == lambdaReturnType op) $ - TC.bad $ TC.TypeError $ "SegHist operator has return type " ++ - prettyTuple (lambdaReturnType op) ++ " but neutral element has type " ++ - prettyTuple nes_t - - -- Arrays must have proper type. - let dest_shape = Shape (segment_dims <> [dest_w]) <> shape - forM_ (zip nes_t dests) $ \(t, dest) -> do - TC.requireI [t `arrayOfShape` dest_shape] dest - TC.consume =<< TC.lookupAliases dest - - return $ map (`arrayOfShape` shape) nes_t - - checkKernelBody ts kbody - - -- Return type of bucket function must be an index for each - -- operation followed by the values to write. - let bucket_ret_t = replicate (length ops) (Prim int32) ++ concat nes_ts - unless (bucket_ret_t == ts) $ - TC.bad $ TC.TypeError $ "SegHist body has return type " ++ - prettyTuple ts ++ " but should have type " ++ - prettyTuple bucket_ret_t - - where segment_dims = init $ segSpaceDims space - -checkScanRed :: TC.Checkable lore => - SegSpace - -> [(Lambda (Aliases lore), [SubExp], Shape)] - -> [Type] - -> KernelBody (Aliases lore) - -> TC.TypeM lore () -checkScanRed space ops ts kbody = do - checkSegSpace space - mapM_ TC.checkType ts - - TC.binding (scopeOfSegSpace space) $ do - ne_ts <- forM ops $ \(lam, nes, shape) -> do - mapM_ (TC.require [Prim int32]) $ shapeDims shape - nes' <- mapM TC.checkArg nes - - -- Operator type must match the type of neutral elements. - TC.checkLambda lam $ map TC.noArgAliases $ nes' ++ nes' - let nes_t = map TC.argType nes' - - unless (lambdaReturnType lam == nes_t) $ - TC.bad $ TC.TypeError "wrong type for operator or neutral elements." - - return $ map (`arrayOfShape` shape) nes_t - - let expecting = concat ne_ts - got = take (length expecting) ts - unless (expecting == got) $ - TC.bad $ TC.TypeError $ - "Wrong return for body (does not match neutral elements; expected " ++ - pretty expecting ++ "; found " ++ - pretty got ++ ")" - - checkKernelBody ts kbody - --- | Like 'Mapper', but just for 'SegOp's. -data SegOpMapper lvl flore tlore m = SegOpMapper { - mapOnSegOpSubExp :: SubExp -> m SubExp - , mapOnSegOpLambda :: Lambda flore -> m (Lambda tlore) - , mapOnSegOpBody :: KernelBody flore -> m (KernelBody tlore) - , mapOnSegOpVName :: VName -> m VName - , mapOnSegOpLevel :: lvl -> m lvl - } - --- | A mapper that simply returns the 'SegOp' verbatim. -identitySegOpMapper :: Monad m => SegOpMapper lvl lore lore m -identitySegOpMapper = SegOpMapper { mapOnSegOpSubExp = return - , mapOnSegOpLambda = return - , mapOnSegOpBody = return - , mapOnSegOpVName = return - , mapOnSegOpLevel = return - } - -mapOnSegSpace :: Monad f => - SegOpMapper lvl flore tlore f -> SegSpace -> f SegSpace -mapOnSegSpace tv (SegSpace phys dims) = - SegSpace phys <$> traverse (traverse $ mapOnSegOpSubExp tv) dims - -mapSegBinOp :: Monad m => - SegOpMapper lvl flore tlore m - -> SegBinOp flore -> m (SegBinOp tlore) -mapSegBinOp tv (SegBinOp comm red_op nes shape) = - SegBinOp comm - <$> mapOnSegOpLambda tv red_op - <*> mapM (mapOnSegOpSubExp tv) nes - <*> (Shape <$> mapM (mapOnSegOpSubExp tv) (shapeDims shape)) - --- | Apply a 'SegOpMapper' to the given 'SegOp'. -mapSegOpM :: (Applicative m, Monad m) => - SegOpMapper lvl flore tlore m - -> SegOp lvl flore -> m (SegOp lvl tlore) -mapSegOpM tv (SegMap lvl space ts body) = - SegMap - <$> mapOnSegOpLevel tv lvl - <*> mapOnSegSpace tv space - <*> mapM (mapOnSegOpType tv) ts - <*> mapOnSegOpBody tv body -mapSegOpM tv (SegRed lvl space reds ts lam) = - SegRed - <$> mapOnSegOpLevel tv lvl - <*> mapOnSegSpace tv space - <*> mapM (mapSegBinOp tv) reds - <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts - <*> mapOnSegOpBody tv lam -mapSegOpM tv (SegScan lvl space scans ts body) = - SegScan - <$> mapOnSegOpLevel tv lvl - <*> mapOnSegSpace tv space - <*> mapM (mapSegBinOp tv) scans - <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts - <*> mapOnSegOpBody tv body -mapSegOpM tv (SegHist lvl space ops ts body) = - SegHist - <$> mapOnSegOpLevel tv lvl - <*> mapOnSegSpace tv space - <*> mapM onHistOp ops - <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts - <*> mapOnSegOpBody tv body - where onHistOp (HistOp w rf arrs nes shape op) = - HistOp <$> mapOnSegOpSubExp tv w - <*> mapOnSegOpSubExp tv rf - <*> mapM (mapOnSegOpVName tv) arrs - <*> mapM (mapOnSegOpSubExp tv) nes - <*> (Shape <$> mapM (mapOnSegOpSubExp tv) (shapeDims shape)) - <*> mapOnSegOpLambda tv op - -mapOnSegOpType :: Monad m => - SegOpMapper lvl flore tlore m -> Type -> m Type -mapOnSegOpType _tv t@Prim{} = pure t -mapOnSegOpType tv (Array pt shape u) = Array pt <$> f shape <*> pure u - where f (Shape dims) = Shape <$> mapM (mapOnSegOpSubExp tv) dims -mapOnSegOpType _tv (Mem s) = pure $ Mem s - -instance (ASTLore lore, Substitute lvl) => - Substitute (SegOp lvl lore) where - substituteNames subst = runIdentity . mapSegOpM substitute - where substitute = - SegOpMapper { mapOnSegOpSubExp = return . substituteNames subst - , mapOnSegOpLambda = return . substituteNames subst - , mapOnSegOpBody = return . substituteNames subst - , mapOnSegOpVName = return . substituteNames subst - , mapOnSegOpLevel = return . substituteNames subst - } - -instance (ASTLore lore, ASTConstraints lvl) => - Rename (SegOp lvl lore) where - rename = mapSegOpM renamer - where renamer = SegOpMapper rename rename rename rename rename - -instance (ASTLore lore, FreeIn (LParamInfo lore), FreeIn lvl) => - FreeIn (SegOp lvl lore) where - freeIn' e = flip execState mempty $ mapSegOpM free e - where walk f x = modify (<>f x) >> return x - free = SegOpMapper { mapOnSegOpSubExp = walk freeIn' - , mapOnSegOpLambda = walk freeIn' - , mapOnSegOpBody = walk freeIn' - , mapOnSegOpVName = walk freeIn' - , mapOnSegOpLevel = walk freeIn' - } - -instance OpMetrics (Op lore) => OpMetrics (SegOp lvl lore) where - opMetrics (SegMap _ _ _ body) = - inside "SegMap" $ kernelBodyMetrics body - opMetrics (SegRed _ _ reds _ body) = - inside "SegRed" $ do mapM_ (lambdaMetrics . segBinOpLambda) reds - kernelBodyMetrics body - opMetrics (SegScan _ _ scans _ body) = - inside "SegScan" $ do mapM_ (lambdaMetrics . segBinOpLambda) scans - kernelBodyMetrics body - opMetrics (SegHist _ _ ops _ body) = - inside "SegHist" $ do mapM_ (lambdaMetrics . histOp) ops - kernelBodyMetrics body - -instance Pretty SegSpace where - ppr (SegSpace phys dims) = parens (commasep $ do (i,d) <- dims - return $ ppr i <+> "<" <+> ppr d) <+> - parens (text "~" <> ppr phys) - -instance PrettyLore lore => Pretty (SegBinOp lore) where - ppr (SegBinOp comm lam nes shape) = - PP.braces (PP.commasep $ map ppr nes) <> PP.comma - ppr shape <> PP.comma - comm' <> ppr lam - where comm' = case comm of Commutative -> text "commutative " - Noncommutative -> mempty - -instance (PrettyLore lore, PP.Pretty lvl) => PP.Pretty (SegOp lvl lore) where - ppr (SegMap lvl space ts body) = - text "segmap" <> ppr lvl - PP.align (ppr space) <+> - PP.colon <+> ppTuple' ts <+> PP.nestedBlock "{" "}" (ppr body) - - ppr (SegRed lvl space reds ts body) = - text "segred" <> ppr lvl - PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppr reds)) - PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> - PP.nestedBlock "{" "}" (ppr body) - - ppr (SegScan lvl space scans ts body) = - text "segscan" <> ppr lvl - PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppr scans)) - PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> - PP.nestedBlock "{" "}" (ppr body) - - ppr (SegHist lvl space ops ts body) = - text "seghist" <> ppr lvl - ppr lvl - PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp ops)) - PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> - PP.nestedBlock "{" "}" (ppr body) - where ppOp (HistOp w rf dests nes shape op) = - ppr w <> PP.comma <+> ppr rf <> PP.comma - PP.braces (PP.commasep $ map ppr dests) <> PP.comma - PP.braces (PP.commasep $ map ppr nes) <> PP.comma - ppr shape <> PP.comma - ppr op - -instance (ASTLore lore, ASTLore (Aliases lore), - CanBeAliased (Op lore), ASTConstraints lvl) => - CanBeAliased (SegOp lvl lore) where - type OpWithAliases (SegOp lvl lore) = SegOp lvl (Aliases lore) - - addOpAliases = runIdentity . mapSegOpM alias - where alias = SegOpMapper return (return . Alias.analyseLambda) - (return . aliasAnalyseKernelBody) return return - - removeOpAliases = runIdentity . mapSegOpM remove - where remove = SegOpMapper return (return . removeLambdaAliases) - (return . removeKernelBodyAliases) return return - -instance (CanBeWise (Op lore), ASTLore lore, ASTConstraints lvl) => - CanBeWise (SegOp lvl lore) where - type OpWithWisdom (SegOp lvl lore) = SegOp lvl (Wise lore) - - removeOpWisdom = runIdentity . mapSegOpM remove - where remove = SegOpMapper return - (return . removeLambdaWisdom) - (return . removeKernelBodyWisdom) - return return - -instance ASTLore lore => ST.IndexOp (SegOp lvl lore) where - indexOp vtable k (SegMap _ space _ kbody) is = do - Returns ResultMaySimplify se <- maybeNth k $ kernelBodyResult kbody - guard $ length gtids <= length is - let idx_table = M.fromList $ zip gtids $ map (ST.Indexed mempty) is - idx_table' = foldl expandIndexedTable idx_table $ kernelBodyStms kbody - case se of - Var v -> M.lookup v idx_table' - _ -> Nothing - - where (gtids, _) = unzip $ unSegSpace space - -- Indexes in excess of what is used to index through the - -- segment dimensions. - excess_is = drop (length gtids) is - - expandIndexedTable table stm - | [v] <- patternNames $ stmPattern stm, - Just (pe,cs) <- - runWriterT $ primExpFromExp (asPrimExp table) $ stmExp stm = - M.insert v (ST.Indexed (stmCerts stm <> cs) pe) table - - | [v] <- patternNames $ stmPattern stm, - BasicOp (Index arr slice) <- stmExp stm, - length (sliceDims slice) == length excess_is, - arr `ST.elem` vtable, - Just (slice', cs) <- asPrimExpSlice table slice = - let idx = ST.IndexedArray (stmCerts stm <> cs) - arr (fixSlice slice' excess_is) - in M.insert v idx table - - | otherwise = - table - - asPrimExpSlice table = - runWriterT . mapM (traverse (primExpFromSubExpM (asPrimExp table))) - - asPrimExp table v - | Just (ST.Indexed cs e) <- M.lookup v table = tell cs >> return e - | Just (Prim pt) <- ST.lookupType v vtable = - return $ LeafExp v pt - | otherwise = lift Nothing - - indexOp _ _ _ _ = Nothing - -instance (ASTLore lore, ASTConstraints lvl) => - IsOp (SegOp lvl lore) where - cheapOp _ = False - safeOp _ = True - ---- Simplification - -instance Engine.Simplifiable SplitOrdering where - simplify SplitContiguous = - return SplitContiguous - simplify (SplitStrided stride) = - SplitStrided <$> Engine.simplify stride - -instance Engine.Simplifiable SegSpace where - simplify (SegSpace phys dims) = - SegSpace phys <$> mapM (traverse Engine.simplify) dims - -instance Engine.Simplifiable KernelResult where - simplify (Returns manifest what) = - Returns manifest <$> Engine.simplify what - simplify (WriteReturns ws a res) = - WriteReturns <$> Engine.simplify ws <*> Engine.simplify a <*> Engine.simplify res - simplify (ConcatReturns o w pte what) = - ConcatReturns - <$> Engine.simplify o - <*> Engine.simplify w - <*> Engine.simplify pte - <*> Engine.simplify what - simplify (TileReturns dims what) = - TileReturns <$> Engine.simplify dims <*> Engine.simplify what - -mkWiseKernelBody :: (ASTLore lore, CanBeWise (Op lore)) => - BodyDec lore -> Stms (Wise lore) -> [KernelResult] -> KernelBody (Wise lore) -mkWiseKernelBody dec bnds res = - let Body dec' _ _ = mkWiseBody dec bnds res_vs - in KernelBody dec' bnds res - where res_vs = map kernelResultSubExp res - -mkKernelBodyM :: MonadBinder m => - Stms (Lore m) -> [KernelResult] - -> m (KernelBody (Lore m)) -mkKernelBodyM stms kres = do - Body dec' _ _ <- mkBodyM stms res_ses - return $ KernelBody dec' stms kres - where res_ses = map kernelResultSubExp kres - -simplifyKernelBody :: (Engine.SimplifiableLore lore, BodyDec lore ~ ()) => - SegSpace -> KernelBody lore - -> Engine.SimpleM lore (KernelBody (Wise lore), Stms (Wise lore)) -simplifyKernelBody space (KernelBody _ stms res) = do - par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers - - ((body_stms, body_res), hoisted) <- - Engine.localVtable (<>scope_vtable) $ - Engine.localVtable (\vtable -> vtable { ST.simplifyMemory = True }) $ - Engine.blockIf (Engine.hasFree bound_here - `Engine.orIf` Engine.isOp - `Engine.orIf` par_blocker - `Engine.orIf` Engine.isConsumed) $ - Engine.simplifyStms stms $ do - res' <- Engine.localVtable (ST.hideCertified $ namesFromList $ M.keys $ scopeOf stms) $ - mapM Engine.simplify res - return ((res', UT.usages $ freeIn res'), mempty) - - return (mkWiseKernelBody () body_stms body_res, hoisted) - - where scope_vtable = segSpaceSymbolTable space - bound_here = namesFromList $ M.keys $ scopeOfSegSpace space - -segSpaceSymbolTable :: ASTLore lore => SegSpace -> ST.SymbolTable lore -segSpaceSymbolTable (SegSpace flat gtids_and_dims) = - foldl' f (ST.fromScope $ M.singleton flat $ IndexName Int32) gtids_and_dims - where f vtable (gtid, dim) = ST.insertLoopVar gtid Int32 dim vtable - -simplifySegBinOp :: Engine.SimplifiableLore lore => - SegBinOp lore - -> Engine.SimpleM lore (SegBinOp (Wise lore), Stms (Wise lore)) -simplifySegBinOp (SegBinOp comm lam nes shape) = do - (lam', hoisted) <- - Engine.localVtable (\vtable -> vtable { ST.simplifyMemory = True }) $ - Engine.simplifyLambda lam - shape' <- Engine.simplify shape - nes' <- mapM Engine.simplify nes - return (SegBinOp comm lam' nes' shape', hoisted) - --- | Simplify the given 'SegOp'. -simplifySegOp :: (Engine.SimplifiableLore lore, - BodyDec lore ~ (), - Engine.Simplifiable lvl) => - SegOp lvl lore - -> Engine.SimpleM lore (SegOp lvl (Wise lore), Stms (Wise lore)) -simplifySegOp (SegMap lvl space ts kbody) = do - (lvl', space', ts') <- Engine.simplify (lvl, space, ts) - (kbody', body_hoisted) <- simplifyKernelBody space kbody - return (SegMap lvl' space' ts' kbody', - body_hoisted) - -simplifySegOp (SegRed lvl space reds ts kbody) = do - (lvl', space', ts') <- Engine.simplify (lvl, space, ts) - (reds', reds_hoisted) <- Engine.localVtable (<>scope_vtable) $ - unzip <$> mapM simplifySegBinOp reds - (kbody', body_hoisted) <- simplifyKernelBody space kbody - - return (SegRed lvl' space' reds' ts' kbody', - mconcat reds_hoisted <> body_hoisted) - where scope = scopeOfSegSpace space - scope_vtable = ST.fromScope scope - -simplifySegOp (SegScan lvl space scans ts kbody) = do - (lvl', space', ts') <- Engine.simplify (lvl, space, ts) - (scans', scans_hoisted) <- Engine.localVtable (<>scope_vtable) $ - unzip <$> mapM simplifySegBinOp scans - (kbody', body_hoisted) <- simplifyKernelBody space kbody - - return (SegScan lvl' space' scans' ts' kbody', - mconcat scans_hoisted <> body_hoisted) - where scope = scopeOfSegSpace space - scope_vtable = ST.fromScope scope - -simplifySegOp (SegHist lvl space ops ts kbody) = do - (lvl', space', ts') <- Engine.simplify (lvl, space, ts) - - (ops', ops_hoisted) <- fmap unzip $ forM ops $ - \(HistOp w rf arrs nes dims lam) -> do - w' <- Engine.simplify w - rf' <- Engine.simplify rf - arrs' <- Engine.simplify arrs - nes' <- Engine.simplify nes - dims' <- Engine.simplify dims - (lam', op_hoisted) <- - Engine.localVtable (<>scope_vtable) $ - Engine.localVtable (\vtable -> vtable { ST.simplifyMemory = True }) $ - Engine.simplifyLambda lam - return (HistOp w' rf' arrs' nes' dims' lam', - op_hoisted) - - (kbody', body_hoisted) <- simplifyKernelBody space kbody - - return (SegHist lvl' space' ops' ts' kbody', - mconcat ops_hoisted <> body_hoisted) - - where scope = scopeOfSegSpace space - scope_vtable = ST.fromScope scope - --- | Does this lore contain 'SegOp's in its t'Op's? A lore must be an --- instance of this class for the simplification rules to work. -class HasSegOp lore where - type SegOpLevel lore - asSegOp :: Op lore -> Maybe (SegOp (SegOpLevel lore) lore) - segOp :: SegOp (SegOpLevel lore) lore -> Op lore - --- | Simplification rules for simplifying 'SegOp's. -segOpRules :: (HasSegOp lore, BinderOps lore, Bindable lore) => - RuleBook lore -segOpRules = - ruleBook [ RuleOp segOpRuleTopDown ] [ RuleOp segOpRuleBottomUp ] - -segOpRuleTopDown :: (HasSegOp lore, BinderOps lore, Bindable lore) => - TopDownRuleOp lore -segOpRuleTopDown vtable pat dec op - | Just op' <- asSegOp op = - topDownSegOp vtable pat dec op' - | otherwise = - Skip - -segOpRuleBottomUp :: (HasSegOp lore, BinderOps lore) => - BottomUpRuleOp lore -segOpRuleBottomUp vtable pat dec op - | Just op' <- asSegOp op = - bottomUpSegOp vtable pat dec op' - | otherwise = - Skip - -topDownSegOp :: (HasSegOp lore, BinderOps lore, Bindable lore) => - ST.SymbolTable lore - -> Pattern lore - -> StmAux (ExpDec lore) - -> SegOp (SegOpLevel lore) lore - -> Rule lore - --- If a SegOp produces something invariant to the SegOp, turn it --- into a replicate. -topDownSegOp vtable (Pattern [] kpes) dec (SegMap lvl space ts (KernelBody _ kstms kres)) = Simplify $ do - (ts', kpes', kres') <- - unzip3 <$> filterM checkForInvarianceResult (zip3 ts kpes kres) - - -- Check if we did anything at all. - when (kres == kres') - cannotSimplify - - kbody <- mkKernelBodyM kstms kres' - addStm $ Let (Pattern [] kpes') dec $ Op $ segOp $ - SegMap lvl space ts' kbody - - where isInvariant Constant{} = True - isInvariant (Var v) = isJust $ ST.lookup v vtable - - checkForInvarianceResult (_, pe, Returns rm se) - | rm == ResultMaySimplify, - isInvariant se = do - letBindNames [patElemName pe] $ - BasicOp $ Replicate (Shape $ segSpaceDims space) se - return False - checkForInvarianceResult _ = - return True - --- If a SegRed contains two reduction operations that have the same --- vector shape, merge them together. This saves on communication --- overhead, but can in principle lead to more local memory usage. -topDownSegOp _ (Pattern [] pes) _ (SegRed lvl space ops ts kbody) - | length ops > 1, - op_groupings <- groupBy sameShape $ zip ops $ chunks (map (length . segBinOpNeutral) ops) $ - zip3 red_pes red_ts red_res, - any ((>1) . length) op_groupings = Simplify $ do - let (ops', aux) = unzip $ mapMaybe combineOps op_groupings - (red_pes', red_ts', red_res') = unzip3 $ concat aux - pes' = red_pes' ++ map_pes - ts' = red_ts' ++ map_ts - kbody' = kbody { kernelBodyResult = red_res' ++ map_res } - letBind (Pattern [] pes') $ Op $ segOp $ SegRed lvl space ops' ts' kbody' - where (red_pes, map_pes) = splitAt (segBinOpResults ops) pes - (red_ts, map_ts) = splitAt (segBinOpResults ops) ts - (red_res, map_res) = splitAt (segBinOpResults ops) $ kernelBodyResult kbody - - sameShape (op1, _) (op2, _) = segBinOpShape op1 == segBinOpShape op2 - - combineOps [] = Nothing - combineOps (x:xs) = Just $ foldl' combine x xs - - combine (op1, op1_aux) (op2, op2_aux) = - let lam1 = segBinOpLambda op1 - lam2 = segBinOpLambda op2 - (op1_xparams, op1_yparams) = - splitAt (length (segBinOpNeutral op1)) $ lambdaParams lam1 - (op2_xparams, op2_yparams) = - splitAt (length (segBinOpNeutral op2)) $ lambdaParams lam2 - lam = Lambda { lambdaParams = op1_xparams ++ op2_xparams ++ - op1_yparams ++ op2_yparams - , lambdaReturnType = lambdaReturnType lam1 ++ lambdaReturnType lam2 - , lambdaBody = - mkBody (bodyStms (lambdaBody lam1) <> bodyStms (lambdaBody lam2)) $ - bodyResult (lambdaBody lam1) <> bodyResult (lambdaBody lam2) - } - in (SegBinOp { segBinOpComm = segBinOpComm op1 <> segBinOpComm op2 - , segBinOpLambda = lam - , segBinOpNeutral = segBinOpNeutral op1 ++ segBinOpNeutral op2 - , segBinOpShape = segBinOpShape op1 -- Same as shape of op2 due to the grouping. - }, - op1_aux ++ op2_aux) -topDownSegOp _ _ _ _ = Skip - -bottomUpSegOp :: (HasSegOp lore, BinderOps lore) => - (ST.SymbolTable lore, UT.UsageTable) - -> Pattern lore - -> StmAux (ExpDec lore) - -> SegOp (SegOpLevel lore) lore - -> Rule lore - --- Some SegOp results can be moved outside the SegOp, which can --- simplify further analysis. -bottomUpSegOp (vtable, used) (Pattern [] kpes) dec (SegMap lvl space kts (KernelBody _ kstms kres)) = Simplify $ do - - -- Iterate through the bindings. For each, we check whether it is - -- in kres and can be moved outside. If so, we remove it from kres - -- and kpes and make it a binding outside. - (kpes', kts', kres', kstms') <- localScope (scopeOfSegSpace space) $ - foldM distribute (kpes, kts, kres, mempty) kstms - - when (kpes' == kpes) - cannotSimplify - - kbody <- localScope (scopeOfSegSpace space) $ - mkKernelBodyM kstms' kres' - - addStm $ Let (Pattern [] kpes') dec $ Op $ segOp $ - SegMap lvl space kts' kbody - where - free_in_kstms = foldMap freeIn kstms - - sliceWithGtidsFixed stm - | Let _ _ (BasicOp (Index arr slice)) <- stm, - space_slice <- map (DimFix . Var . fst) $ unSegSpace space, - space_slice `isPrefixOf` slice, - remaining_slice <- drop (length space_slice) slice, - all (isJust . flip ST.lookup vtable) $ namesToList $ - freeIn arr <> freeIn remaining_slice = - Just (remaining_slice, arr) - - | otherwise = - Nothing - - distribute (kpes', kts', kres', kstms') stm - | Let (Pattern [] [pe]) _ _ <- stm, - Just (remaining_slice, arr) <- sliceWithGtidsFixed stm, - Just (kpe, kpes'', kts'', kres'') <- isResult kpes' kts' kres' pe = do - let outer_slice = map (\d -> DimSlice - (constant (0::Int32)) - d - (constant (1::Int32))) $ - segSpaceDims space - index kpe' = letBind (Pattern [] [kpe']) $ BasicOp $ Index arr $ - outer_slice <> remaining_slice - if patElemName kpe `UT.isConsumed` used - then do precopy <- newVName $ baseString (patElemName kpe) <> "_precopy" - index kpe { patElemName = precopy } - letBind (Pattern [] [kpe]) $ BasicOp $ Copy precopy - else index kpe - return (kpes'', kts'', kres'', - if patElemName pe `nameIn` free_in_kstms - then kstms' <> oneStm stm - else kstms') - - distribute (kpes', kts', kres', kstms') stm = - return (kpes', kts', kres', kstms' <> oneStm stm) - - isResult kpes' kts' kres' pe = - case partition matches $ zip3 kpes' kts' kres' of - ([(kpe,_,_)], kpes_and_kres) - | (kpes'', kts'', kres'') <- unzip3 kpes_and_kres -> - Just (kpe, kpes'', kts'', kres'') - _ -> Nothing - where matches (_, _, Returns _ (Var v)) = v == patElemName pe - matches _ = False -bottomUpSegOp _ _ _ _ = Skip - ---- Memory - -kernelBodyReturns :: (Mem lore, HasScope lore m, Monad m) => - KernelBody lore -> [ExpReturns] -> m [ExpReturns] -kernelBodyReturns = zipWithM correct . kernelBodyResult - where correct (WriteReturns _ arr _) _ = varReturns arr - correct _ ret = return ret - --- | Like 'segOpType', but for memory representations. -segOpReturns :: (Mem lore, Monad m, HasScope lore m) => - SegOp lvl lore -> m [ExpReturns] +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | Segmented operations. These correspond to perfect @map@ nests on +-- top of /something/, except that the @map@s are conceptually only +-- over @iota@s (so there will be explicit indexing inside them). +module Futhark.IR.SegOp + ( SegOp (..), + SegVirt (..), + segLevel, + segSpace, + typeCheckSegOp, + SegSpace (..), + scopeOfSegSpace, + segSpaceDims, + + -- * Details + HistOp (..), + histType, + SegBinOp (..), + segBinOpResults, + segBinOpChunks, + KernelBody (..), + aliasAnalyseKernelBody, + consumedInKernelBody, + ResultManifest (..), + KernelResult (..), + kernelResultSubExp, + SplitOrdering (..), + + -- ** Generic traversal + SegOpMapper (..), + identitySegOpMapper, + mapSegOpM, + + -- * Simplification + simplifySegOp, + HasSegOp (..), + segOpRules, + + -- * Memory + segOpReturns, + ) +where + +import Control.Category +import Control.Monad.Identity hiding (mapM_) +import Control.Monad.State.Strict +import Control.Monad.Writer hiding (mapM_) +import Data.Bifunctor (first) +import Data.List + ( foldl', + groupBy, + intersperse, + isPrefixOf, + partition, + ) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Futhark.Analysis.Alias as Alias +import Futhark.Analysis.Metrics +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.IR +import Futhark.IR.Aliases + ( Aliases, + removeLambdaAliases, + removeStmAliases, + ) +import Futhark.IR.Mem +import Futhark.IR.Prop.Aliases +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rule +import Futhark.Tools +import Futhark.Transform.Rename +import Futhark.Transform.Substitute +import qualified Futhark.TypeCheck as TC +import Futhark.Util (chunks, maybeNth) +import Futhark.Util.Pretty + ( Pretty, + commasep, + parens, + ppr, + text, + (<+>), + (), + ) +import qualified Futhark.Util.Pretty as PP +import GHC.Generics (Generic) +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) + +-- | How an array is split into chunks. +data SplitOrdering + = SplitContiguous + | SplitStrided SubExp + deriving (Eq, Ord, Show, Generic) + +instance SexpIso SplitOrdering where + sexpIso = + match $ + With (. Sexp.sym "contiguous") $ + With + (. Sexp.list (Sexp.el (Sexp.sym "strided") >>> Sexp.el sexpIso)) + End + +instance FreeIn SplitOrdering where + freeIn' SplitContiguous = mempty + freeIn' (SplitStrided stride) = freeIn' stride + +instance Substitute SplitOrdering where + substituteNames _ SplitContiguous = + SplitContiguous + substituteNames subst (SplitStrided stride) = + SplitStrided $ substituteNames subst stride + +instance Rename SplitOrdering where + rename SplitContiguous = + pure SplitContiguous + rename (SplitStrided stride) = + SplitStrided <$> rename stride + +-- | An operator for 'SegHist'. +data HistOp lore = HistOp + { histWidth :: SubExp, + histRaceFactor :: SubExp, + histDest :: [VName], + histNeutral :: [SubExp], + -- | In case this operator is semantically a vectorised + -- operator (corresponding to a perfect map nest in the + -- SOACS representation), these are the logical + -- "dimensions". This is used to generate more efficient + -- code. + histShape :: Shape, + histOp :: Lambda lore + } + deriving (Eq, Ord, Show, Generic) + +instance Decorations lore => SexpIso (HistOp lore) where + sexpIso = with $ \histop -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> histop + +-- | The type of a histogram produced by a 'HistOp'. This can be +-- different from the type of the 'histDest's in case we are +-- dealing with a segmented histogram. +histType :: HistOp lore -> [Type] +histType op = + map + ( (`arrayOfRow` histWidth op) + . (`arrayOfShape` histShape op) + ) + $ lambdaReturnType $ histOp op + +-- | An operator for 'SegScan' and 'SegRed'. +data SegBinOp lore = SegBinOp + { segBinOpComm :: Commutativity, + segBinOpLambda :: Lambda lore, + segBinOpNeutral :: [SubExp], + -- | In case this operator is semantically a vectorised + -- operator (corresponding to a perfect map nest in the + -- SOACS representation), these are the logical + -- "dimensions". This is used to generate more efficient + -- code. + segBinOpShape :: Shape + } + deriving (Eq, Ord, Show, Generic) + +instance Decorations lore => SexpIso (SegBinOp lore) where + sexpIso = with $ \segbinop -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> segbinop + +-- | How many reduction results are produced by these 'SegBinOp's? +segBinOpResults :: [SegBinOp lore] -> Int +segBinOpResults = sum . map (length . segBinOpNeutral) + +-- | Split some list into chunks equal to the number of values +-- returned by each 'SegBinOp' +segBinOpChunks :: [SegBinOp lore] -> [a] -> [[a]] +segBinOpChunks = chunks . map (length . segBinOpNeutral) + +-- | The body of a 'SegOp'. +data KernelBody lore = KernelBody + { kernelBodyLore :: BodyDec lore, + kernelBodyStms :: Stms lore, + kernelBodyResult :: [KernelResult] + } + deriving (Generic) + +instance Decorations lore => SexpIso (KernelBody lore) where + sexpIso = with $ \kernelbody -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> kernelbody + +deriving instance Decorations lore => Ord (KernelBody lore) + +deriving instance Decorations lore => Show (KernelBody lore) + +deriving instance Decorations lore => Eq (KernelBody lore) + +-- | Metadata about whether there is a subtle point to this +-- 'KernelResult'. This is used to protect things like tiling, which +-- might otherwise be removed by the simplifier because they're +-- semantically redundant. This has no semantic effect and can be +-- ignored at code generation. +data ResultManifest + = -- | Don't simplify this one! + ResultNoSimplify + | -- | Go nuts. + ResultMaySimplify + | -- | The results produced are only used within the + -- same physical thread later on, and can thus be + -- kept in registers. + ResultPrivate + deriving (Eq, Show, Ord, Generic) + +instance SexpIso ResultManifest where + sexpIso = + match $ + With (. Sexp.sym "no-simplify") $ + With (. Sexp.sym "may-simplify") $ + With + (. Sexp.sym "private") + End + +-- | A 'KernelBody' does not return an ordinary 'Result'. Instead, it +-- returns a list of these. +data KernelResult + = -- | Each "worker" in the kernel returns this. + -- Whether this is a result-per-thread or a + -- result-per-group depends on where the 'SegOp' occurs. + Returns ResultManifest SubExp + | WriteReturns + [SubExp] -- Size of array. Must match number of dims. + VName -- Which array + [(Slice SubExp, SubExp)] + | -- Arbitrary number of index/value pairs. + ConcatReturns + SplitOrdering -- Permuted? + SubExp -- The final size. + SubExp -- Per-thread/group (max) chunk size. + VName -- Chunk by this worker. + | TileReturns + [(SubExp, SubExp)] -- Total/tile for each dimension + VName -- Tile written by this worker. + -- The TileReturns must not expect more than one + -- result to be written per physical thread. + deriving (Eq, Show, Ord, Generic) + +instance SexpIso KernelResult where + sexpIso = + match $ + 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 + +-- | Get the root t'SubExp' corresponding values for a 'KernelResult'. +kernelResultSubExp :: KernelResult -> SubExp +kernelResultSubExp (Returns _ se) = se +kernelResultSubExp (WriteReturns _ arr _) = Var arr +kernelResultSubExp (ConcatReturns _ _ _ v) = Var v +kernelResultSubExp (TileReturns _ v) = Var v + +instance FreeIn KernelResult where + freeIn' (Returns _ what) = freeIn' what + freeIn' (WriteReturns rws arr res) = freeIn' rws <> freeIn' arr <> freeIn' res + freeIn' (ConcatReturns o w per_thread_elems v) = + freeIn' o <> freeIn' w <> freeIn' per_thread_elems <> freeIn' v + freeIn' (TileReturns dims v) = + freeIn' dims <> freeIn' v + +instance ASTLore lore => FreeIn (KernelBody lore) where + freeIn' (KernelBody dec stms res) = + fvBind bound_in_stms $ freeIn' dec <> freeIn' stms <> freeIn' res + where + bound_in_stms = foldMap boundByStm stms + +instance ASTLore lore => Substitute (KernelBody lore) where + substituteNames subst (KernelBody dec stms res) = + KernelBody + (substituteNames subst dec) + (substituteNames subst stms) + (substituteNames subst res) + +instance Substitute KernelResult where + substituteNames subst (Returns manifest se) = + Returns manifest (substituteNames subst se) + substituteNames subst (WriteReturns rws arr res) = + WriteReturns + (substituteNames subst rws) + (substituteNames subst arr) + (substituteNames subst res) + substituteNames subst (ConcatReturns o w per_thread_elems v) = + ConcatReturns + (substituteNames subst o) + (substituteNames subst w) + (substituteNames subst per_thread_elems) + (substituteNames subst v) + substituteNames subst (TileReturns dims v) = + TileReturns (substituteNames subst dims) (substituteNames subst v) + +instance ASTLore lore => Rename (KernelBody lore) where + rename (KernelBody dec stms res) = do + dec' <- rename dec + renamingStms stms $ \stms' -> + KernelBody dec' stms' <$> rename res + +instance Rename KernelResult where + rename = substituteRename + +-- | Perform alias analysis on a 'KernelBody'. +aliasAnalyseKernelBody :: + ( ASTLore lore, + CanBeAliased (Op lore) + ) => + KernelBody lore -> + KernelBody (Aliases lore) +aliasAnalyseKernelBody (KernelBody dec stms res) = + let Body dec' stms' _ = Alias.analyseBody mempty $ Body dec stms [] + in KernelBody dec' stms' res + +removeKernelBodyAliases :: + CanBeAliased (Op lore) => + KernelBody (Aliases lore) -> + KernelBody lore +removeKernelBodyAliases (KernelBody (_, dec) stms res) = + KernelBody dec (fmap removeStmAliases stms) res + +removeKernelBodyWisdom :: + CanBeWise (Op lore) => + KernelBody (Wise lore) -> + KernelBody lore +removeKernelBodyWisdom (KernelBody dec stms res) = + let Body dec' stms' _ = removeBodyWisdom $ Body dec stms [] + in KernelBody dec' stms' res + +-- | The variables consumed in the kernel body. +consumedInKernelBody :: + Aliased lore => + KernelBody lore -> + Names +consumedInKernelBody (KernelBody dec stms res) = + consumedInBody (Body dec stms []) <> mconcat (map consumedByReturn res) + where + consumedByReturn (WriteReturns _ a _) = oneName a + consumedByReturn _ = mempty + +checkKernelBody :: + TC.Checkable lore => + [Type] -> + KernelBody (Aliases lore) -> + TC.TypeM lore () +checkKernelBody ts (KernelBody (_, dec) stms kres) = do + TC.checkBodyLore dec + TC.checkStms stms $ do + unless (length ts == length kres) $ + TC.bad $ + TC.TypeError $ + "Kernel return type is " ++ prettyTuple ts + ++ ", but body returns " + ++ show (length kres) + ++ " values." + zipWithM_ checkKernelResult kres ts + where + checkKernelResult (Returns _ what) t = + TC.require [t] what + checkKernelResult (WriteReturns rws arr res) t = do + mapM_ (TC.require [Prim int64]) rws + arr_t <- lookupType arr + forM_ res $ \(slice, e) -> do + mapM_ (traverse $ TC.require [Prim int64]) slice + TC.require [t] e + unless (arr_t == t `arrayOfShape` Shape rws) $ + TC.bad $ + TC.TypeError $ + "WriteReturns returning " + ++ pretty e + ++ " of type " + ++ pretty t + ++ ", shape=" + ++ pretty rws + ++ ", but destination array has type " + ++ pretty arr_t + TC.consume =<< TC.lookupAliases arr + checkKernelResult (ConcatReturns o w per_thread_elems v) t = do + case o of + SplitContiguous -> return () + SplitStrided stride -> TC.require [Prim int64] stride + TC.require [Prim int64] w + TC.require [Prim int64] per_thread_elems + vt <- lookupType v + unless (vt == t `arrayOfRow` arraySize 0 vt) $ + TC.bad $ TC.TypeError $ "Invalid type for ConcatReturns " ++ pretty v + checkKernelResult (TileReturns dims v) t = do + forM_ dims $ \(dim, tile) -> do + TC.require [Prim int64] dim + TC.require [Prim int64] tile + vt <- lookupType v + unless (vt == t `arrayOfShape` Shape (map snd dims)) $ + TC.bad $ TC.TypeError $ "Invalid type for TileReturns " ++ pretty v + +kernelBodyMetrics :: OpMetrics (Op lore) => KernelBody lore -> MetricsM () +kernelBodyMetrics = mapM_ stmMetrics . kernelBodyStms + +instance PrettyLore lore => Pretty (KernelBody lore) where + ppr (KernelBody _ stms res) = + PP.stack (map ppr (stmsToList stms)) + text "return" <+> PP.braces (PP.commasep $ map ppr res) + +instance Pretty KernelResult where + ppr (Returns ResultNoSimplify what) = + text "returns (manifest)" <+> ppr what + ppr (Returns ResultPrivate what) = + text "returns (private)" <+> ppr what + ppr (Returns ResultMaySimplify what) = + text "returns" <+> ppr what + ppr (WriteReturns rws arr res) = + ppr arr <+> text "with" <+> PP.apply (map ppRes res) + where + ppRes (is, e) = + PP.brackets (PP.commasep $ zipWith f is rws) <+> text "<-" <+> ppr e + f i rw = ppr i <+> text "<" <+> ppr rw + ppr (ConcatReturns o w per_thread_elems v) = + text "concat" <> suff + <> parens (commasep [ppr w, ppr per_thread_elems]) <+> ppr v + where + suff = case o of + SplitContiguous -> mempty + SplitStrided stride -> text "Strided" <> parens (ppr stride) + ppr (TileReturns dims v) = + text "tile" + <> parens (commasep $ map onDim dims) <+> ppr v + where + onDim (dim, tile) = ppr dim <+> text "/" <+> ppr tile + +-- | Do we need group-virtualisation when generating code for the +-- segmented operation? In most cases, we do, but for some simple +-- kernels, we compute the full number of groups in advance, and then +-- virtualisation is an unnecessary (but generally very small) +-- overhead. This only really matters for fairly trivial but very +-- wide @map@ kernels where each thread performs constant-time work on +-- scalars. +data SegVirt + = SegVirt + | SegNoVirt + | -- | Not only do we not need virtualisation, but we _guarantee_ + -- that all physical threads participate in the work. This can + -- save some checks in code generation. + SegNoVirtFull + deriving (Eq, Ord, Show, Generic) + +instance SexpIso SegVirt where + sexpIso = + match $ + With (. Sexp.sym "virt") $ + With (. Sexp.sym "no-virt") $ + With + (. Sexp.sym "no-virt-ful") + End + +-- | Index space of a 'SegOp'. +data SegSpace = SegSpace + { -- | Flat physical index corresponding to the + -- dimensions (at code generation used for a + -- thread ID or similar). + segFlat :: VName, + unSegSpace :: [(VName, SubExp)] + } + deriving (Eq, Ord, Show, Generic) + +instance SexpIso SegSpace where + sexpIso = with $ \segspace -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> segspace + +-- | The sizes spanned by the indexes of the 'SegSpace'. +segSpaceDims :: SegSpace -> [SubExp] +segSpaceDims (SegSpace _ space) = map snd space + +-- | A 'Scope' containing all the identifiers brought into scope by +-- this 'SegSpace'. +scopeOfSegSpace :: SegSpace -> Scope lore +scopeOfSegSpace (SegSpace phys space) = + M.fromList $ zip (phys : map fst space) $ repeat $ IndexName Int64 + +checkSegSpace :: TC.Checkable lore => SegSpace -> TC.TypeM lore () +checkSegSpace (SegSpace _ dims) = + mapM_ (TC.require [Prim int64] . snd) dims + +-- | A 'SegOp' is semantically a perfectly nested stack of maps, on +-- top of some bottommost computation (scalar computation, reduction, +-- scan, or histogram). The 'SegSpace' encodes the original map +-- structure. +-- +-- All 'SegOp's are parameterised by the representation of their body, +-- as well as a *level*. The *level* is a representation-specific bit +-- of information. For example, in GPU backends, it is used to +-- indicate whether the 'SegOp' is expected to run at the thread-level +-- or the group-level. +data SegOp lvl lore + = SegMap lvl SegSpace [Type] (KernelBody lore) + | -- | The KernelSpace must always have at least two dimensions, + -- implying that the result of a SegRed is always an array. + SegRed lvl SegSpace [SegBinOp lore] [Type] (KernelBody lore) + | SegScan lvl SegSpace [SegBinOp lore] [Type] (KernelBody lore) + | SegHist lvl SegSpace [HistOp lore] [Type] (KernelBody lore) + deriving (Eq, Ord, Show, Generic) + +instance (SexpIso lvl, Decorations lore) => SexpIso (SegOp lvl lore) where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "segmap") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "segred") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "segscan") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "seghist") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End + +-- | The level of a 'SegOp'. +segLevel :: SegOp lvl lore -> lvl +segLevel (SegMap lvl _ _ _) = lvl +segLevel (SegRed lvl _ _ _ _) = lvl +segLevel (SegScan lvl _ _ _ _) = lvl +segLevel (SegHist lvl _ _ _ _) = lvl + +-- | The space of a 'SegOp'. +segSpace :: SegOp lvl lore -> SegSpace +segSpace (SegMap _ lvl _ _) = lvl +segSpace (SegRed _ lvl _ _ _) = lvl +segSpace (SegScan _ lvl _ _ _) = lvl +segSpace (SegHist _ lvl _ _ _) = lvl + +segResultShape :: SegSpace -> Type -> KernelResult -> Type +segResultShape _ t (WriteReturns rws _ _) = + t `arrayOfShape` Shape rws +segResultShape space t (Returns _ _) = + foldr (flip arrayOfRow) t $ segSpaceDims space +segResultShape _ t (ConcatReturns _ w _ _) = + t `arrayOfRow` w +segResultShape _ t (TileReturns dims _) = + t `arrayOfShape` Shape (map fst dims) + +-- | The return type of a 'SegOp'. +segOpType :: SegOp lvl lore -> [Type] +segOpType (SegMap _ space ts kbody) = + zipWith (segResultShape space) ts $ kernelBodyResult kbody +segOpType (SegRed _ space reds ts kbody) = + red_ts + ++ zipWith + (segResultShape space) + map_ts + (drop (length red_ts) $ kernelBodyResult kbody) + where + map_ts = drop (length red_ts) ts + segment_dims = init $ segSpaceDims space + red_ts = do + op <- reds + let shape = Shape segment_dims <> segBinOpShape op + map (`arrayOfShape` shape) (lambdaReturnType $ segBinOpLambda op) +segOpType (SegScan _ space scans ts kbody) = + scan_ts + ++ zipWith + (segResultShape space) + map_ts + (drop (length scan_ts) $ kernelBodyResult kbody) + where + map_ts = drop (length scan_ts) ts + scan_ts = do + op <- scans + let shape = Shape (segSpaceDims space) <> segBinOpShape op + map (`arrayOfShape` shape) (lambdaReturnType $ segBinOpLambda op) +segOpType (SegHist _ space ops _ _) = do + op <- ops + let shape = Shape (segment_dims <> [histWidth op]) <> histShape op + map (`arrayOfShape` shape) (lambdaReturnType $ histOp op) + where + dims = segSpaceDims space + segment_dims = init dims + +instance TypedOp (SegOp lvl lore) where + opType = pure . staticShapes . segOpType + +instance + (ASTLore lore, Aliased lore, ASTConstraints lvl) => + AliasedOp (SegOp lvl lore) + where + opAliases = map (const mempty) . segOpType + + consumedInOp (SegMap _ _ _ kbody) = + consumedInKernelBody kbody + consumedInOp (SegRed _ _ _ _ kbody) = + consumedInKernelBody kbody + consumedInOp (SegScan _ _ _ _ kbody) = + consumedInKernelBody kbody + consumedInOp (SegHist _ _ ops _ kbody) = + namesFromList (concatMap histDest ops) <> consumedInKernelBody kbody + +-- | Type check a 'SegOp', given a checker for its level. +typeCheckSegOp :: + TC.Checkable lore => + (lvl -> TC.TypeM lore ()) -> + SegOp lvl (Aliases lore) -> + TC.TypeM lore () +typeCheckSegOp checkLvl (SegMap lvl space ts kbody) = do + checkLvl lvl + checkScanRed space [] ts kbody +typeCheckSegOp checkLvl (SegRed lvl space reds ts body) = do + checkLvl lvl + checkScanRed space reds' ts body + where + reds' = + zip3 + (map segBinOpLambda reds) + (map segBinOpNeutral reds) + (map segBinOpShape reds) +typeCheckSegOp checkLvl (SegScan lvl space scans ts body) = do + checkLvl lvl + checkScanRed space scans' ts body + where + scans' = + zip3 + (map segBinOpLambda scans) + (map segBinOpNeutral scans) + (map segBinOpShape scans) +typeCheckSegOp checkLvl (SegHist lvl space ops ts kbody) = do + checkLvl lvl + checkSegSpace space + mapM_ TC.checkType ts + + TC.binding (scopeOfSegSpace space) $ do + nes_ts <- forM ops $ \(HistOp dest_w rf dests nes shape op) -> do + TC.require [Prim int64] dest_w + TC.require [Prim int64] rf + nes' <- mapM TC.checkArg nes + mapM_ (TC.require [Prim int64]) $ shapeDims shape + + -- Operator type must match the type of neutral elements. + let stripVecDims = stripArray $ shapeRank shape + TC.checkLambda op $ map (TC.noArgAliases . first stripVecDims) $ nes' ++ nes' + let nes_t = map TC.argType nes' + unless (nes_t == lambdaReturnType op) $ + TC.bad $ + TC.TypeError $ + "SegHist operator has return type " + ++ prettyTuple (lambdaReturnType op) + ++ " but neutral element has type " + ++ prettyTuple nes_t + + -- Arrays must have proper type. + let dest_shape = Shape (segment_dims <> [dest_w]) <> shape + forM_ (zip nes_t dests) $ \(t, dest) -> do + TC.requireI [t `arrayOfShape` dest_shape] dest + TC.consume =<< TC.lookupAliases dest + + return $ map (`arrayOfShape` shape) nes_t + + checkKernelBody ts kbody + + -- Return type of bucket function must be an index for each + -- operation followed by the values to write. + let bucket_ret_t = replicate (length ops) (Prim int64) ++ concat nes_ts + unless (bucket_ret_t == ts) $ + TC.bad $ + TC.TypeError $ + "SegHist body has return type " + ++ prettyTuple ts + ++ " but should have type " + ++ prettyTuple bucket_ret_t + where + segment_dims = init $ segSpaceDims space + +checkScanRed :: + TC.Checkable lore => + SegSpace -> + [(Lambda (Aliases lore), [SubExp], Shape)] -> + [Type] -> + KernelBody (Aliases lore) -> + TC.TypeM lore () +checkScanRed space ops ts kbody = do + checkSegSpace space + mapM_ TC.checkType ts + + TC.binding (scopeOfSegSpace space) $ do + ne_ts <- forM ops $ \(lam, nes, shape) -> do + mapM_ (TC.require [Prim int64]) $ shapeDims shape + nes' <- mapM TC.checkArg nes + + -- Operator type must match the type of neutral elements. + TC.checkLambda lam $ map TC.noArgAliases $ nes' ++ nes' + let nes_t = map TC.argType nes' + + unless (lambdaReturnType lam == nes_t) $ + TC.bad $ TC.TypeError "wrong type for operator or neutral elements." + + return $ map (`arrayOfShape` shape) nes_t + + let expecting = concat ne_ts + got = take (length expecting) ts + unless (expecting == got) $ + TC.bad $ + TC.TypeError $ + "Wrong return for body (does not match neutral elements; expected " + ++ pretty expecting + ++ "; found " + ++ pretty got + ++ ")" + + checkKernelBody ts kbody + +-- | Like 'Mapper', but just for 'SegOp's. +data SegOpMapper lvl flore tlore m = SegOpMapper + { mapOnSegOpSubExp :: SubExp -> m SubExp, + mapOnSegOpLambda :: Lambda flore -> m (Lambda tlore), + mapOnSegOpBody :: KernelBody flore -> m (KernelBody tlore), + mapOnSegOpVName :: VName -> m VName, + mapOnSegOpLevel :: lvl -> m lvl + } + +-- | A mapper that simply returns the 'SegOp' verbatim. +identitySegOpMapper :: Monad m => SegOpMapper lvl lore lore m +identitySegOpMapper = + SegOpMapper + { mapOnSegOpSubExp = return, + mapOnSegOpLambda = return, + mapOnSegOpBody = return, + mapOnSegOpVName = return, + mapOnSegOpLevel = return + } + +mapOnSegSpace :: + Monad f => + SegOpMapper lvl flore tlore f -> + SegSpace -> + f SegSpace +mapOnSegSpace tv (SegSpace phys dims) = + SegSpace phys <$> traverse (traverse $ mapOnSegOpSubExp tv) dims + +mapSegBinOp :: + Monad m => + SegOpMapper lvl flore tlore m -> + SegBinOp flore -> + m (SegBinOp tlore) +mapSegBinOp tv (SegBinOp comm red_op nes shape) = + SegBinOp comm + <$> mapOnSegOpLambda tv red_op + <*> mapM (mapOnSegOpSubExp tv) nes + <*> (Shape <$> mapM (mapOnSegOpSubExp tv) (shapeDims shape)) + +-- | Apply a 'SegOpMapper' to the given 'SegOp'. +mapSegOpM :: + (Applicative m, Monad m) => + SegOpMapper lvl flore tlore m -> + SegOp lvl flore -> + m (SegOp lvl tlore) +mapSegOpM tv (SegMap lvl space ts body) = + SegMap + <$> mapOnSegOpLevel tv lvl + <*> mapOnSegSpace tv space + <*> mapM (mapOnSegOpType tv) ts + <*> mapOnSegOpBody tv body +mapSegOpM tv (SegRed lvl space reds ts lam) = + SegRed + <$> mapOnSegOpLevel tv lvl + <*> mapOnSegSpace tv space + <*> mapM (mapSegBinOp tv) reds + <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts + <*> mapOnSegOpBody tv lam +mapSegOpM tv (SegScan lvl space scans ts body) = + SegScan + <$> mapOnSegOpLevel tv lvl + <*> mapOnSegSpace tv space + <*> mapM (mapSegBinOp tv) scans + <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts + <*> mapOnSegOpBody tv body +mapSegOpM tv (SegHist lvl space ops ts body) = + SegHist + <$> mapOnSegOpLevel tv lvl + <*> mapOnSegSpace tv space + <*> mapM onHistOp ops + <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts + <*> mapOnSegOpBody tv body + where + onHistOp (HistOp w rf arrs nes shape op) = + HistOp <$> mapOnSegOpSubExp tv w + <*> mapOnSegOpSubExp tv rf + <*> mapM (mapOnSegOpVName tv) arrs + <*> mapM (mapOnSegOpSubExp tv) nes + <*> (Shape <$> mapM (mapOnSegOpSubExp tv) (shapeDims shape)) + <*> mapOnSegOpLambda tv op + +mapOnSegOpType :: + Monad m => + SegOpMapper lvl flore tlore m -> + Type -> + m Type +mapOnSegOpType _tv t@Prim {} = pure t +mapOnSegOpType tv (Array pt shape u) = Array pt <$> f shape <*> pure u + where + f (Shape dims) = Shape <$> mapM (mapOnSegOpSubExp tv) dims +mapOnSegOpType _tv (Mem s) = pure $ Mem s + +instance + (ASTLore lore, Substitute lvl) => + Substitute (SegOp lvl lore) + where + substituteNames subst = runIdentity . mapSegOpM substitute + where + substitute = + SegOpMapper + { mapOnSegOpSubExp = return . substituteNames subst, + mapOnSegOpLambda = return . substituteNames subst, + mapOnSegOpBody = return . substituteNames subst, + mapOnSegOpVName = return . substituteNames subst, + mapOnSegOpLevel = return . substituteNames subst + } + +instance + (ASTLore lore, ASTConstraints lvl) => + Rename (SegOp lvl lore) + where + rename = mapSegOpM renamer + where + renamer = SegOpMapper rename rename rename rename rename + +instance + (ASTLore lore, FreeIn (LParamInfo lore), FreeIn lvl) => + FreeIn (SegOp lvl lore) + where + freeIn' e = flip execState mempty $ mapSegOpM free e + where + walk f x = modify (<> f x) >> return x + free = + SegOpMapper + { mapOnSegOpSubExp = walk freeIn', + mapOnSegOpLambda = walk freeIn', + mapOnSegOpBody = walk freeIn', + mapOnSegOpVName = walk freeIn', + mapOnSegOpLevel = walk freeIn' + } + +instance OpMetrics (Op lore) => OpMetrics (SegOp lvl lore) where + opMetrics (SegMap _ _ _ body) = + inside "SegMap" $ kernelBodyMetrics body + opMetrics (SegRed _ _ reds _ body) = + inside "SegRed" $ do + mapM_ (lambdaMetrics . segBinOpLambda) reds + kernelBodyMetrics body + opMetrics (SegScan _ _ scans _ body) = + inside "SegScan" $ do + mapM_ (lambdaMetrics . segBinOpLambda) scans + kernelBodyMetrics body + opMetrics (SegHist _ _ ops _ body) = + inside "SegHist" $ do + mapM_ (lambdaMetrics . histOp) ops + kernelBodyMetrics body + +instance Pretty SegSpace where + ppr (SegSpace phys dims) = + parens + ( commasep $ do + (i, d) <- dims + return $ ppr i <+> "<" <+> ppr d + ) + <+> parens (text "~" <> ppr phys) + +instance PrettyLore lore => Pretty (SegBinOp lore) where + ppr (SegBinOp comm lam nes shape) = + PP.braces (PP.commasep $ map ppr nes) <> PP.comma + ppr shape <> PP.comma + comm' <> ppr lam + where + comm' = case comm of + Commutative -> text "commutative " + Noncommutative -> mempty + +instance (PrettyLore lore, PP.Pretty lvl) => PP.Pretty (SegOp lvl lore) where + ppr (SegMap lvl space ts body) = + text "segmap" <> ppr lvl + PP.align (ppr space) + <+> PP.colon + <+> ppTuple' ts + <+> PP.nestedBlock "{" "}" (ppr body) + ppr (SegRed lvl space reds ts body) = + text "segred" <> ppr lvl + PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppr reds)) + PP.align (ppr space) + <+> PP.colon + <+> ppTuple' ts + <+> PP.nestedBlock "{" "}" (ppr body) + ppr (SegScan lvl space scans ts body) = + text "segscan" <> ppr lvl + PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppr scans)) + PP.align (ppr space) + <+> PP.colon + <+> ppTuple' ts + <+> PP.nestedBlock "{" "}" (ppr body) + ppr (SegHist lvl space ops ts body) = + text "seghist" <> ppr lvl + ppr lvl + PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp ops)) + PP.align (ppr space) + <+> PP.colon + <+> ppTuple' ts + <+> PP.nestedBlock "{" "}" (ppr body) + where + ppOp (HistOp w rf dests nes shape op) = + ppr w <> PP.comma <+> ppr rf <> PP.comma + PP.braces (PP.commasep $ map ppr dests) <> PP.comma + PP.braces (PP.commasep $ map ppr nes) <> PP.comma + ppr shape <> PP.comma + ppr op + +instance + ( ASTLore lore, + ASTLore (Aliases lore), + CanBeAliased (Op lore), + ASTConstraints lvl + ) => + CanBeAliased (SegOp lvl lore) + where + type OpWithAliases (SegOp lvl lore) = SegOp lvl (Aliases lore) + + addOpAliases = runIdentity . mapSegOpM alias + where + alias = + SegOpMapper + return + (return . Alias.analyseLambda) + (return . aliasAnalyseKernelBody) + return + return + + removeOpAliases = runIdentity . mapSegOpM remove + where + remove = + SegOpMapper + return + (return . removeLambdaAliases) + (return . removeKernelBodyAliases) + return + return + +instance + (CanBeWise (Op lore), ASTLore lore, ASTConstraints lvl) => + CanBeWise (SegOp lvl lore) + where + type OpWithWisdom (SegOp lvl lore) = SegOp lvl (Wise lore) + + removeOpWisdom = runIdentity . mapSegOpM remove + where + remove = + SegOpMapper + return + (return . removeLambdaWisdom) + (return . removeKernelBodyWisdom) + return + return + +instance ASTLore lore => ST.IndexOp (SegOp lvl lore) where + indexOp vtable k (SegMap _ space _ kbody) is = do + Returns ResultMaySimplify se <- maybeNth k $ kernelBodyResult kbody + guard $ length gtids <= length is + let idx_table = M.fromList $ zip gtids $ map (ST.Indexed mempty . untyped) is + idx_table' = foldl' expandIndexedTable idx_table $ kernelBodyStms kbody + case se of + Var v -> M.lookup v idx_table' + _ -> Nothing + where + (gtids, _) = unzip $ unSegSpace space + -- Indexes in excess of what is used to index through the + -- segment dimensions. + excess_is = drop (length gtids) is + + expandIndexedTable table stm + | [v] <- patternNames $ stmPattern stm, + Just (pe, cs) <- + runWriterT $ primExpFromExp (asPrimExp table) $ stmExp stm = + M.insert v (ST.Indexed (stmCerts stm <> cs) pe) table + | [v] <- patternNames $ stmPattern stm, + BasicOp (Index arr slice) <- stmExp stm, + length (sliceDims slice) == length excess_is, + arr `ST.elem` vtable, + Just (slice', cs) <- asPrimExpSlice table slice = + let idx = + ST.IndexedArray + (stmCerts stm <> cs) + arr + (fixSlice (map (fmap isInt64) slice') excess_is) + in M.insert v idx table + | otherwise = + table + + asPrimExpSlice table = + runWriterT . mapM (traverse (primExpFromSubExpM (asPrimExp table))) + + asPrimExp table v + | Just (ST.Indexed cs e) <- M.lookup v table = tell cs >> return e + | Just (Prim pt) <- ST.lookupType v vtable = + return $ LeafExp v pt + | otherwise = lift Nothing + indexOp _ _ _ _ = Nothing + +instance + (ASTLore lore, ASTConstraints lvl) => + IsOp (SegOp lvl lore) + where + cheapOp _ = False + safeOp _ = True + +--- Simplification + +instance Engine.Simplifiable SplitOrdering where + simplify SplitContiguous = + return SplitContiguous + simplify (SplitStrided stride) = + SplitStrided <$> Engine.simplify stride + +instance Engine.Simplifiable SegSpace where + simplify (SegSpace phys dims) = + SegSpace phys <$> mapM (traverse Engine.simplify) dims + +instance Engine.Simplifiable KernelResult where + simplify (Returns manifest what) = + Returns manifest <$> Engine.simplify what + simplify (WriteReturns ws a res) = + WriteReturns <$> Engine.simplify ws <*> Engine.simplify a <*> Engine.simplify res + simplify (ConcatReturns o w pte what) = + ConcatReturns + <$> Engine.simplify o + <*> Engine.simplify w + <*> Engine.simplify pte + <*> Engine.simplify what + simplify (TileReturns dims what) = + TileReturns <$> Engine.simplify dims <*> Engine.simplify what + +mkWiseKernelBody :: + (ASTLore lore, CanBeWise (Op lore)) => + BodyDec lore -> + Stms (Wise lore) -> + [KernelResult] -> + KernelBody (Wise lore) +mkWiseKernelBody dec bnds res = + let Body dec' _ _ = mkWiseBody dec bnds res_vs + in KernelBody dec' bnds res + where + res_vs = map kernelResultSubExp res + +mkKernelBodyM :: + MonadBinder m => + Stms (Lore m) -> + [KernelResult] -> + m (KernelBody (Lore m)) +mkKernelBodyM stms kres = do + Body dec' _ _ <- mkBodyM stms res_ses + return $ KernelBody dec' stms kres + where + res_ses = map kernelResultSubExp kres + +simplifyKernelBody :: + (Engine.SimplifiableLore lore, BodyDec lore ~ ()) => + SegSpace -> + KernelBody lore -> + Engine.SimpleM lore (KernelBody (Wise lore), Stms (Wise lore)) +simplifyKernelBody space (KernelBody _ stms res) = do + par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers + + ((body_stms, body_res), hoisted) <- + Engine.localVtable (<> scope_vtable) $ + Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $ + Engine.blockIf + ( Engine.hasFree bound_here + `Engine.orIf` Engine.isOp + `Engine.orIf` par_blocker + `Engine.orIf` Engine.isConsumed + ) + $ Engine.simplifyStms stms $ do + res' <- + Engine.localVtable (ST.hideCertified $ namesFromList $ M.keys $ scopeOf stms) $ + mapM Engine.simplify res + return ((res', UT.usages $ freeIn res'), mempty) + + return (mkWiseKernelBody () body_stms body_res, hoisted) + where + scope_vtable = segSpaceSymbolTable space + bound_here = namesFromList $ M.keys $ scopeOfSegSpace space + +segSpaceSymbolTable :: ASTLore lore => SegSpace -> ST.SymbolTable lore +segSpaceSymbolTable (SegSpace flat gtids_and_dims) = + foldl' f (ST.fromScope $ M.singleton flat $ IndexName Int64) gtids_and_dims + where + f vtable (gtid, dim) = ST.insertLoopVar gtid Int64 dim vtable + +simplifySegBinOp :: + Engine.SimplifiableLore lore => + SegBinOp lore -> + Engine.SimpleM lore (SegBinOp (Wise lore), Stms (Wise lore)) +simplifySegBinOp (SegBinOp comm lam nes shape) = do + (lam', hoisted) <- + Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $ + Engine.simplifyLambda lam + shape' <- Engine.simplify shape + nes' <- mapM Engine.simplify nes + return (SegBinOp comm lam' nes' shape', hoisted) + +-- | Simplify the given 'SegOp'. +simplifySegOp :: + ( Engine.SimplifiableLore lore, + BodyDec lore ~ (), + Engine.Simplifiable lvl + ) => + SegOp lvl lore -> + Engine.SimpleM lore (SegOp lvl (Wise lore), Stms (Wise lore)) +simplifySegOp (SegMap lvl space ts kbody) = do + (lvl', space', ts') <- Engine.simplify (lvl, space, ts) + (kbody', body_hoisted) <- simplifyKernelBody space kbody + return + ( SegMap lvl' space' ts' kbody', + body_hoisted + ) +simplifySegOp (SegRed lvl space reds ts kbody) = do + (lvl', space', ts') <- Engine.simplify (lvl, space, ts) + (reds', reds_hoisted) <- + Engine.localVtable (<> scope_vtable) $ + unzip <$> mapM simplifySegBinOp reds + (kbody', body_hoisted) <- simplifyKernelBody space kbody + + return + ( SegRed lvl' space' reds' ts' kbody', + mconcat reds_hoisted <> body_hoisted + ) + where + scope = scopeOfSegSpace space + scope_vtable = ST.fromScope scope +simplifySegOp (SegScan lvl space scans ts kbody) = do + (lvl', space', ts') <- Engine.simplify (lvl, space, ts) + (scans', scans_hoisted) <- + Engine.localVtable (<> scope_vtable) $ + unzip <$> mapM simplifySegBinOp scans + (kbody', body_hoisted) <- simplifyKernelBody space kbody + + return + ( SegScan lvl' space' scans' ts' kbody', + mconcat scans_hoisted <> body_hoisted + ) + where + scope = scopeOfSegSpace space + scope_vtable = ST.fromScope scope +simplifySegOp (SegHist lvl space ops ts kbody) = do + (lvl', space', ts') <- Engine.simplify (lvl, space, ts) + + (ops', ops_hoisted) <- fmap unzip $ + forM ops $ + \(HistOp w rf arrs nes dims lam) -> do + w' <- Engine.simplify w + rf' <- Engine.simplify rf + arrs' <- Engine.simplify arrs + nes' <- Engine.simplify nes + dims' <- Engine.simplify dims + (lam', op_hoisted) <- + Engine.localVtable (<> scope_vtable) $ + Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $ + Engine.simplifyLambda lam + return + ( HistOp w' rf' arrs' nes' dims' lam', + op_hoisted + ) + + (kbody', body_hoisted) <- simplifyKernelBody space kbody + + return + ( SegHist lvl' space' ops' ts' kbody', + mconcat ops_hoisted <> body_hoisted + ) + where + scope = scopeOfSegSpace space + scope_vtable = ST.fromScope scope + +-- | Does this lore contain 'SegOp's in its t'Op's? A lore must be an +-- instance of this class for the simplification rules to work. +class HasSegOp lore where + type SegOpLevel lore + asSegOp :: Op lore -> Maybe (SegOp (SegOpLevel lore) lore) + segOp :: SegOp (SegOpLevel lore) lore -> Op lore + +-- | Simplification rules for simplifying 'SegOp's. +segOpRules :: + (HasSegOp lore, BinderOps lore, Bindable lore) => + RuleBook lore +segOpRules = + ruleBook [RuleOp segOpRuleTopDown] [RuleOp segOpRuleBottomUp] + +segOpRuleTopDown :: + (HasSegOp lore, BinderOps lore, Bindable lore) => + TopDownRuleOp lore +segOpRuleTopDown vtable pat dec op + | Just op' <- asSegOp op = + topDownSegOp vtable pat dec op' + | otherwise = + Skip + +segOpRuleBottomUp :: + (HasSegOp lore, BinderOps lore) => + BottomUpRuleOp lore +segOpRuleBottomUp vtable pat dec op + | Just op' <- asSegOp op = + bottomUpSegOp vtable pat dec op' + | otherwise = + Skip + +topDownSegOp :: + (HasSegOp lore, BinderOps lore, Bindable lore) => + ST.SymbolTable lore -> + Pattern lore -> + StmAux (ExpDec lore) -> + SegOp (SegOpLevel lore) lore -> + Rule lore +-- If a SegOp produces something invariant to the SegOp, turn it +-- into a replicate. +topDownSegOp vtable (Pattern [] kpes) dec (SegMap lvl space ts (KernelBody _ kstms kres)) = Simplify $ do + (ts', kpes', kres') <- + unzip3 <$> filterM checkForInvarianceResult (zip3 ts kpes kres) + + -- Check if we did anything at all. + when + (kres == kres') + cannotSimplify + + kbody <- mkKernelBodyM kstms kres' + addStm $ + Let (Pattern [] kpes') dec $ + Op $ + segOp $ + SegMap lvl space ts' kbody + where + isInvariant Constant {} = True + isInvariant (Var v) = isJust $ ST.lookup v vtable + + checkForInvarianceResult (_, pe, Returns rm se) + | rm == ResultMaySimplify, + isInvariant se = do + letBindNames [patElemName pe] $ + BasicOp $ Replicate (Shape $ segSpaceDims space) se + return False + checkForInvarianceResult _ = + return True + +-- If a SegRed contains two reduction operations that have the same +-- vector shape, merge them together. This saves on communication +-- overhead, but can in principle lead to more local memory usage. +topDownSegOp _ (Pattern [] pes) _ (SegRed lvl space ops ts kbody) + | length ops > 1, + op_groupings <- + groupBy sameShape $ + zip ops $ + chunks (map (length . segBinOpNeutral) ops) $ + zip3 red_pes red_ts red_res, + any ((> 1) . length) op_groupings = Simplify $ do + let (ops', aux) = unzip $ mapMaybe combineOps op_groupings + (red_pes', red_ts', red_res') = unzip3 $ concat aux + pes' = red_pes' ++ map_pes + ts' = red_ts' ++ map_ts + kbody' = kbody {kernelBodyResult = red_res' ++ map_res} + letBind (Pattern [] pes') $ Op $ segOp $ SegRed lvl space ops' ts' kbody' + where + (red_pes, map_pes) = splitAt (segBinOpResults ops) pes + (red_ts, map_ts) = splitAt (segBinOpResults ops) ts + (red_res, map_res) = splitAt (segBinOpResults ops) $ kernelBodyResult kbody + + sameShape (op1, _) (op2, _) = segBinOpShape op1 == segBinOpShape op2 + + combineOps [] = Nothing + combineOps (x : xs) = Just $ foldl' combine x xs + + combine (op1, op1_aux) (op2, op2_aux) = + let lam1 = segBinOpLambda op1 + lam2 = segBinOpLambda op2 + (op1_xparams, op1_yparams) = + splitAt (length (segBinOpNeutral op1)) $ lambdaParams lam1 + (op2_xparams, op2_yparams) = + splitAt (length (segBinOpNeutral op2)) $ lambdaParams lam2 + lam = + Lambda + { lambdaParams = + op1_xparams ++ op2_xparams + ++ op1_yparams + ++ op2_yparams, + lambdaReturnType = lambdaReturnType lam1 ++ lambdaReturnType lam2, + lambdaBody = + mkBody (bodyStms (lambdaBody lam1) <> bodyStms (lambdaBody lam2)) $ + bodyResult (lambdaBody lam1) <> bodyResult (lambdaBody lam2) + } + in ( SegBinOp + { segBinOpComm = segBinOpComm op1 <> segBinOpComm op2, + segBinOpLambda = lam, + segBinOpNeutral = segBinOpNeutral op1 ++ segBinOpNeutral op2, + segBinOpShape = segBinOpShape op1 -- Same as shape of op2 due to the grouping. + }, + op1_aux ++ op2_aux + ) +topDownSegOp _ _ _ _ = Skip + +bottomUpSegOp :: + (HasSegOp lore, BinderOps lore) => + (ST.SymbolTable lore, UT.UsageTable) -> + Pattern lore -> + StmAux (ExpDec lore) -> + SegOp (SegOpLevel lore) lore -> + Rule lore +-- Some SegOp results can be moved outside the SegOp, which can +-- simplify further analysis. +bottomUpSegOp (vtable, used) (Pattern [] kpes) dec (SegMap lvl space kts (KernelBody _ kstms kres)) = Simplify $ do + -- Iterate through the bindings. For each, we check whether it is + -- in kres and can be moved outside. If so, we remove it from kres + -- and kpes and make it a binding outside. + (kpes', kts', kres', kstms') <- + localScope (scopeOfSegSpace space) $ + foldM distribute (kpes, kts, kres, mempty) kstms + + when + (kpes' == kpes) + cannotSimplify + + kbody <- + localScope (scopeOfSegSpace space) $ + mkKernelBodyM kstms' kres' + + addStm $ + Let (Pattern [] kpes') dec $ + Op $ + segOp $ + SegMap lvl space kts' kbody + where + free_in_kstms = foldMap freeIn kstms + + sliceWithGtidsFixed stm + | Let _ _ (BasicOp (Index arr slice)) <- stm, + space_slice <- map (DimFix . Var . fst) $ unSegSpace space, + space_slice `isPrefixOf` slice, + remaining_slice <- drop (length space_slice) slice, + all (isJust . flip ST.lookup vtable) $ + namesToList $ + freeIn arr <> freeIn remaining_slice = + Just (remaining_slice, arr) + | otherwise = + Nothing + + distribute (kpes', kts', kres', kstms') stm + | Let (Pattern [] [pe]) _ _ <- stm, + Just (remaining_slice, arr) <- sliceWithGtidsFixed stm, + Just (kpe, kpes'', kts'', kres'') <- isResult kpes' kts' kres' pe = do + let outer_slice = + map + ( \d -> + DimSlice + (constant (0 :: Int64)) + d + (constant (1 :: Int64)) + ) + $ segSpaceDims space + index kpe' = + letBind (Pattern [] [kpe']) $ + BasicOp $ + Index arr $ + outer_slice <> remaining_slice + if patElemName kpe `UT.isConsumed` used + then do + precopy <- newVName $ baseString (patElemName kpe) <> "_precopy" + index kpe {patElemName = precopy} + letBind (Pattern [] [kpe]) $ BasicOp $ Copy precopy + else index kpe + return + ( kpes'', + kts'', + kres'', + if patElemName pe `nameIn` free_in_kstms + then kstms' <> oneStm stm + else kstms' + ) + distribute (kpes', kts', kres', kstms') stm = + return (kpes', kts', kres', kstms' <> oneStm stm) + + isResult kpes' kts' kres' pe = + case partition matches $ zip3 kpes' kts' kres' of + ([(kpe, _, _)], kpes_and_kres) + | (kpes'', kts'', kres'') <- unzip3 kpes_and_kres -> + Just (kpe, kpes'', kts'', kres'') + _ -> Nothing + where + matches (_, _, Returns _ (Var v)) = v == patElemName pe + matches _ = False +bottomUpSegOp _ _ _ _ = Skip + +--- Memory + +kernelBodyReturns :: + (Mem lore, HasScope lore m, Monad m) => + KernelBody lore -> + [ExpReturns] -> + m [ExpReturns] +kernelBodyReturns = zipWithM correct . kernelBodyResult + where + correct (WriteReturns _ arr _) _ = varReturns arr + correct _ ret = return ret + +-- | Like 'segOpType', but for memory representations. +segOpReturns :: + (Mem lore, Monad m, HasScope lore m) => + SegOp lvl lore -> + m [ExpReturns] segOpReturns k@(SegMap _ _ _ kbody) = kernelBodyReturns kbody =<< (extReturns <$> opType k) segOpReturns k@(SegRed _ _ _ _ kbody) = diff --git a/src/Futhark/IR/Seq.hs b/src/Futhark/IR/Seq.hs --- a/src/Futhark/IR/Seq.hs +++ b/src/Futhark/IR/Seq.hs @@ -1,32 +1,33 @@ -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} + -- | A sequential representation. module Futhark.IR.Seq - ( -- * The Lore definition - Seq + ( -- * The Lore definition + Seq, - -- * Simplification - , simplifyProg + -- * Simplification + simplifyProg, - -- * Module re-exports - , module Futhark.IR.Prop - , module Futhark.IR.Traversals - , module Futhark.IR.Pretty - , module Futhark.IR.Syntax - ) + -- * Module re-exports + module Futhark.IR.Prop, + module Futhark.IR.Traversals, + module Futhark.IR.Pretty, + module Futhark.IR.Syntax, + ) where -import Futhark.Pass -import Futhark.IR.Syntax -import Futhark.IR.Prop -import Futhark.IR.Traversals -import Futhark.IR.Pretty import Futhark.Binder import Futhark.Construct -import qualified Futhark.TypeCheck as TypeCheck -import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.IR.Pretty +import Futhark.IR.Prop +import Futhark.IR.Syntax +import Futhark.IR.Traversals import qualified Futhark.Optimise.Simplify as Simplify +import qualified Futhark.Optimise.Simplify.Engine as Engine import Futhark.Optimise.Simplify.Rules +import Futhark.Pass +import qualified Futhark.TypeCheck as TypeCheck -- | The phantom type for the Seq representation. data Seq @@ -40,7 +41,7 @@ instance TypeCheck.CheckableOp Seq where checkOp = pure -instance TypeCheck.Checkable Seq where +instance TypeCheck.Checkable Seq instance Bindable Seq where mkBody = Body () @@ -48,11 +49,11 @@ mkExpDec _ _ = () mkLetNames = simpleMkLetNames -instance BinderOps Seq where +instance BinderOps Seq -instance PrettyLore Seq where +instance PrettyLore Seq -instance BinderOps (Engine.Wise Seq) where +instance BinderOps (Engine.Wise Seq) simpleSeq :: Simplify.SimpleOps Seq simpleSeq = Simplify.bindableSimpleOps (const $ pure ((), mempty)) @@ -60,4 +61,5 @@ -- | Simplify a sequential program. simplifyProg :: Prog Seq -> PassM (Prog Seq) simplifyProg = Simplify.simplifyProg simpleSeq standardRules blockers - where blockers = Engine.noExtraHoistBlockers + where + blockers = Engine.noExtraHoistBlockers diff --git a/src/Futhark/IR/SeqMem.hs b/src/Futhark/IR/SeqMem.hs --- a/src/Futhark/IR/SeqMem.hs +++ b/src/Futhark/IR/SeqMem.hs @@ -1,43 +1,40 @@ -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE TypeFamilies #-} + module Futhark.IR.SeqMem - ( SeqMem + ( SeqMem, - -- * Simplification - , simplifyProg - , simpleSeqMem + -- * Simplification + simplifyProg, + simpleSeqMem, -- * Module re-exports - , module Futhark.IR.Mem - , module Futhark.IR.Kernels.Kernel + module Futhark.IR.Mem, + module Futhark.IR.Kernels.Kernel, ) - where +where import Futhark.Analysis.PrimExp.Convert -import Futhark.Pass -import Futhark.IR.Syntax -import Futhark.IR.Prop -import Futhark.IR.Traversals -import Futhark.IR.Pretty import Futhark.IR.Kernels.Kernel -import qualified Futhark.TypeCheck as TC import Futhark.IR.Mem import Futhark.IR.Mem.Simplify -import Futhark.Pass.ExplicitAllocations (BinderOps(..), mkLetNamesB', mkLetNamesB'') import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Pass +import Futhark.Pass.ExplicitAllocations (BinderOps (..), mkLetNamesB', mkLetNamesB'') +import qualified Futhark.TypeCheck as TC data SeqMem instance Decorations SeqMem where - type LetDec SeqMem = LetDecMem + type LetDec SeqMem = LetDecMem type FParamInfo SeqMem = FParamMem type LParamInfo SeqMem = LParamMem - type RetType SeqMem = RetTypeMem + type RetType SeqMem = RetTypeMem type BranchType SeqMem = BranchTypeMem - type Op SeqMem = MemOp () + type Op SeqMem = MemOp () instance ASTLore SeqMem where expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern @@ -46,7 +43,7 @@ opReturns (Alloc _ space) = return [MemMem space] opReturns (Inner ()) = pure [] -instance PrettyLore SeqMem where +instance PrettyLore SeqMem instance TC.CheckableOp SeqMem where checkOp (Alloc size _) = 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 @@ -1,7 +1,13 @@ -{-# LANGUAGE TypeFamilies, FlexibleContexts, FlexibleInstances, StandaloneDeriving #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE Strict #-} {-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TypeFamilies #-} + -- | = Definition of the Futhark core language IR -- -- For actually /constructing/ ASTs, see "Futhark.Construct". @@ -96,97 +102,112 @@ -- instances for it. See the source of "Futhark.IR.Seq" -- for what is likely the simplest example. module Futhark.IR.Syntax - ( - module Language.Futhark.Core - , module Futhark.IR.Decorations - , module Futhark.IR.Syntax.Core + ( module Language.Futhark.Core, + module Futhark.IR.Decorations, + module Futhark.IR.Syntax.Core, - -- * Types - , Uniqueness(..) - , NoUniqueness(..) - , Rank(..) - , ArrayShape(..) - , Space (..) - , TypeBase(..) - , Diet(..) + -- * Types + Uniqueness (..), + NoUniqueness (..), + Rank (..), + ArrayShape (..), + Space (..), + TypeBase (..), + Diet (..), - -- * Attributes - , Attr(..) - , Attrs(..) - , oneAttr - , inAttrs - , withoutAttrs + -- * Attributes + Attr (..), + Attrs (..), + oneAttr, + inAttrs, + withoutAttrs, - -- * Abstract syntax tree - , Ident (..) - , SubExp(..) - , PatElem - , PatElemT (..) - , PatternT (..) - , Pattern - , StmAux(..) - , Stm(..) - , Stms - , Result - , BodyT(..) - , Body - , BasicOp (..) - , UnOp (..) - , BinOp (..) - , CmpOp (..) - , ConvOp (..) - , DimChange (..) - , ShapeChange - , ExpT(..) - , Exp - , LoopForm (..) - , IfDec (..) - , IfSort (..) - , Safety (..) - , LambdaT(..) - , Lambda + -- * Abstract syntax tree + Ident (..), + SubExp (..), + PatElem, + PatElemT (..), + PatternT (..), + Pattern, + StmAux (..), + Stm (..), + Stms, + Result, + BodyT (..), + Body, + BasicOp (..), + UnOp (..), + BinOp (..), + CmpOp (..), + ConvOp (..), + DimChange (..), + ShapeChange, + ExpT (..), + Exp, + LoopForm (..), + IfDec (..), + IfSort (..), + Safety (..), + LambdaT (..), + Lambda, - -- * Definitions - , Param (..) - , FParam - , LParam - , FunDef (..) - , EntryPoint - , EntryPointType(..) - , Prog(..) + -- * Definitions + Param (..), + FParam, + LParam, + FunDef (..), + EntryPoint, + EntryPointType (..), + Prog (..), - -- * Utils - , oneStm - , stmsFromList - , stmsToList - , stmsHead + -- * Utils + oneStm, + stmsFromList, + stmsToList, + stmsHead, ) - where +where -import qualified Data.Set as S +import Control.Category import Data.Foldable import qualified Data.Sequence as Seq +import qualified Data.Set as S import Data.String +import qualified Data.Text as T import Data.Traversable (fmapDefault, foldMapDefault) - -import Language.Futhark.Core import Futhark.IR.Decorations import Futhark.IR.Syntax.Core +import GHC.Generics (Generic) +import Language.Futhark.Core +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) -- | A single attribute. data Attr = AttrAtom Name | AttrComp Name [Attr] - deriving (Ord, Show, Eq) + deriving (Ord, Show, Eq, Generic) +instance SexpIso Attr where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "atom") >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "comp") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End + instance IsString Attr where fromString = AttrAtom . fromString -- | Every statement is associated with a set of attributes, which can -- have various effects throughout the compiler. -newtype Attrs = Attrs { unAttrs :: S.Set Attr } - deriving (Ord, Show, Eq, Monoid, Semigroup) +newtype Attrs = Attrs {unAttrs :: S.Set Attr} + deriving (Ord, Show, Eq, Monoid, Semigroup, Generic) +instance SexpIso Attrs where + sexpIso = with $ \attrs -> sexpIso >>> attrs + -- | Construct 'Attrs' from a single 'Attr'. oneAttr :: Attr -> Attrs oneAttr = Attrs . S.singleton @@ -203,16 +224,24 @@ type PatElem lore = PatElemT (LetDec lore) -- | A pattern is conceptually just a list of names and their types. -data PatternT dec = - Pattern { patternContextElements :: [PatElemT dec] - -- ^ existential context (sizes and memory blocks) - , patternValueElements :: [PatElemT dec] - -- ^ "real" values - } - deriving (Ord, Show, Eq) +data PatternT dec = Pattern + { -- | existential context (sizes and memory blocks) + patternContextElements :: [PatElemT dec], + -- | "real" values + patternValueElements :: [PatElemT dec] + } + deriving (Ord, Show, Eq, Generic) +instance SexpIso dec => SexpIso (PatternT dec) where + sexpIso = with $ \vname -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> vname + instance Semigroup (PatternT dec) where - Pattern cs1 vs1 <> Pattern cs2 vs2 = Pattern (cs1++cs2) (vs1++vs2) + Pattern cs1 vs1 <> Pattern cs2 vs2 = Pattern (cs1 ++ cs2) (vs1 ++ vs2) instance Monoid (PatternT dec) where mempty = Pattern [] [] @@ -231,32 +260,60 @@ type Pattern lore = PatternT (LetDec lore) -- | Auxilliary Information associated with a statement. -data StmAux dec = StmAux { stmAuxCerts :: !Certificates - , stmAuxAttrs :: Attrs - , stmAuxDec :: dec - } - deriving (Ord, Show, Eq) +data StmAux dec = StmAux + { stmAuxCerts :: !Certificates, + stmAuxAttrs :: Attrs, + stmAuxDec :: dec + } + deriving (Ord, Show, Eq, Generic) +instance SexpIso dec => SexpIso (StmAux dec) where + sexpIso = with $ \vname -> + Sexp.list + ( Sexp.el (Sexp.sym "aux") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> vname + instance Semigroup dec => Semigroup (StmAux dec) where StmAux cs1 attrs1 dec1 <> StmAux cs2 attrs2 dec2 = - StmAux (cs1<>cs2) (attrs1<>attrs2) (dec1<>dec2) + StmAux (cs1 <> cs2) (attrs1 <> attrs2) (dec1 <> dec2) -- | A local variable binding. -data Stm lore = Let { stmPattern :: Pattern lore - -- ^ Pattern. - , stmAux :: StmAux (ExpDec lore) - -- ^ Auxiliary information statement. - , stmExp :: Exp lore - -- ^ Expression. - } +data Stm lore = Let + { -- | Pattern. + stmPattern :: Pattern lore, + -- | Auxiliary information statement. + stmAux :: StmAux (ExpDec lore), + -- | Expression. + stmExp :: Exp lore + } + deriving (Generic) +instance Decorations lore => SexpIso (Stm lore) where + sexpIso = with $ \stm -> + Sexp.list + ( Sexp.el (Sexp.sym "let") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> stm + deriving instance Decorations lore => Ord (Stm lore) + deriving instance Decorations lore => Show (Stm lore) + deriving instance Decorations lore => Eq (Stm lore) -- | A sequence of statements. type Stms lore = Seq.Seq (Stm lore) +instance Decorations lore => SexpIso (Stms lore) where + sexpIso = iso stmsFromList stmsToList . sexpIso + -- | A single statement. oneStm :: Stm lore -> Stms lore oneStm = Seq.singleton @@ -271,21 +328,35 @@ -- | The first statement in the sequence, if any. stmsHead :: Stms lore -> Maybe (Stm lore, Stms lore) -stmsHead stms = case Seq.viewl stms of stm Seq.:< stms' -> Just (stm, stms') - Seq.EmptyL -> Nothing +stmsHead stms = case Seq.viewl stms of + stm Seq.:< stms' -> Just (stm, stms') + Seq.EmptyL -> Nothing -- | The result of a body is a sequence of subexpressions. type Result = [SubExp] -- | A body consists of a number of bindings, terminating in a result -- (essentially a tuple literal). -data BodyT lore = Body { bodyDec :: BodyDec lore - , bodyStms :: Stms lore - , bodyResult :: Result - } +data BodyT lore = Body + { bodyDec :: BodyDec lore, + bodyStms :: Stms lore, + bodyResult :: Result + } + deriving (Generic) +instance Decorations lore => SexpIso (BodyT lore) where + sexpIso = with $ \stm -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> stm + deriving instance Decorations lore => Ord (BodyT lore) + deriving instance Decorations lore => Show (BodyT lore) + deriving instance Decorations lore => Eq (BodyT lore) -- | Type alias for namespace reasons. @@ -295,14 +366,23 @@ -- disambiguate "real" reshapes, that change the actual shape of the -- array, from type coercions that are just present to make the types -- work out. The two constructors are considered equal for purposes of 'Eq'. -data DimChange d = DimCoercion d - -- ^ The new dimension is guaranteed to be numerically - -- equal to the old one. - | DimNew d - -- ^ The new dimension is not necessarily numerically - -- equal to the old one. - deriving (Ord, Show) +data DimChange d + = -- | The new dimension is guaranteed to be numerically + -- equal to the old one. + DimCoercion d + | -- | The new dimension is not necessarily numerically + -- equal to the old one. + DimNew d + deriving (Ord, Show, Generic) +instance SexpIso d => SexpIso (DimChange d) where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "coercion") >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "new") >>> Sexp.el sexpIso)) + End + instance Eq d => Eq (DimChange d) where DimCoercion x == DimNew y = x == y DimCoercion x == DimCoercion y = x == y @@ -311,15 +391,15 @@ instance Functor DimChange where fmap f (DimCoercion d) = DimCoercion $ f d - fmap f (DimNew d) = DimNew $ f d + fmap f (DimNew d) = DimNew $ f d instance Foldable DimChange where foldMap f (DimCoercion d) = f d - foldMap f (DimNew d) = f d + foldMap f (DimNew d) = f d instance Traversable DimChange where traverse f (DimCoercion d) = DimCoercion <$> f d - traverse f (DimNew d) = DimNew <$> f d + traverse f (DimNew d) = DimNew <$> f d -- | A list of 'DimChange's, indicating the new dimensions of an array. type ShapeChange d = [DimChange d] @@ -327,148 +407,229 @@ -- | A primitive operation that returns something of known size and -- does not itself contain any bindings. data BasicOp - = SubExp SubExp - -- ^ A variable or constant. - - | Opaque SubExp - -- ^ Semantically and operationally just identity, but is + = -- | A variable or constant. + SubExp SubExp + | -- | Semantically and operationally just identity, but is -- invisible/impenetrable to optimisations (hopefully). This is -- just a hack to avoid optimisation (so, to work around compiler -- limitations). - - | ArrayLit [SubExp] Type - -- ^ Array literals, e.g., @[ [1+x, 3], [2, 1+4] ]@. + 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 - - | UnOp UnOp SubExp - -- ^ Unary operation. - - | BinOp BinOp SubExp SubExp - -- ^ Binary operation. - - | CmpOp CmpOp SubExp SubExp - -- ^ Comparison - result type is always boolean. - - | ConvOp ConvOp SubExp - -- ^ Conversion "casting". - - | Assert SubExp (ErrorMsg SubExp) (SrcLoc, [SrcLoc]) - -- ^ Turn a boolean into a certificate, halting the program with the - -- given error message if the boolean is false. - - -- Primitive array operations - - | Index VName (Slice SubExp) - -- ^ The certificates for bounds-checking are part of the 'Stm'. - - | Update VName (Slice SubExp) SubExp - -- ^ An in-place update of the given array at the given position. - -- Consumes the array. - - | Concat Int VName [VName] SubExp - -- ^ @concat@0([1],[2, 3, 4]) = [1, 2, 3, 4]@. - - | Copy VName - -- ^ Copy the given array. The result will not alias anything. - - | Manifest [Int] VName - -- ^ Manifest an array with dimensions represented in the given - -- order. The result will not alias anything. - - -- Array construction. - | Iota SubExp SubExp SubExp IntType - -- ^ @iota(n, x, s) = [x,x+s,..,x+(n-1)*s]@. - -- - -- The t'IntType' indicates the type of the array returned and the - -- offset/stride arguments, but not the length argument. - - | Replicate Shape SubExp - -- ^ @replicate([3][2],1) = [[1,1], [1,1], [1,1]]@ + ArrayLit [SubExp] Type + | -- | Unary operation. + UnOp UnOp SubExp + | -- | Binary operation. + BinOp BinOp SubExp SubExp + | -- | Comparison - result type is always boolean. + CmpOp CmpOp SubExp SubExp + | -- | Conversion "casting". + ConvOp ConvOp SubExp + | -- | Turn a boolean into a certificate, halting the program with the + -- given error message if the boolean is false. + Assert SubExp (ErrorMsg SubExp) (SrcLoc, [SrcLoc]) + | -- Primitive array operations - | Scratch PrimType [SubExp] - -- ^ Create array of given type and shape, with undefined elements. + -- | The certificates for bounds-checking are part of the 'Stm'. + Index VName (Slice SubExp) + | -- | An in-place update of the given array at the given position. + -- Consumes the array. + Update VName (Slice SubExp) SubExp + | -- | @concat@0([1],[2, 3, 4]) = [1, 2, 3, 4]@. + Concat Int VName [VName] SubExp + | -- | Copy the given array. The result will not alias anything. + Copy VName + | -- | Manifest an array with dimensions represented in the given + -- order. The result will not alias anything. + Manifest [Int] VName + | -- Array construction. - -- Array index space transformation. - | Reshape (ShapeChange SubExp) VName - -- ^ 1st arg is the new shape, 2nd arg is the input array *) + -- | @iota(n, x, s) = [x,x+s,..,x+(n-1)*s]@. + -- + -- The t'IntType' indicates the type of the array returned and the + -- offset/stride arguments, but not the length argument. + Iota SubExp SubExp SubExp IntType + | -- | @replicate([3][2],1) = [[1,1], [1,1], [1,1]]@ + Replicate Shape SubExp + | -- | Create array of given type and shape, with undefined elements. + Scratch PrimType [SubExp] + | -- Array index space transformation. - | Rearrange [Int] VName - -- ^ Permute the dimensions of the input array. The list - -- of integers is a list of dimensions (0-indexed), which - -- must be a permutation of @[0,n-1]@, where @n@ is the - -- number of dimensions in the input array. + -- | 1st arg is the new shape, 2nd arg is the input array *) + Reshape (ShapeChange SubExp) VName + | -- | Permute the dimensions of the input array. The list + -- of integers is a list of dimensions (0-indexed), which + -- must be a permutation of @[0,n-1]@, where @n@ is the + -- number of dimensions in the input array. + Rearrange [Int] VName + | -- | Rotate the dimensions of the input array. The list of + -- subexpressions specify how much each dimension is rotated. The + -- length of this list must be equal to the rank of the array. + Rotate [SubExp] VName + deriving (Eq, Ord, Show, Generic) - | Rotate [SubExp] VName - -- ^ Rotate the dimensions of the input array. The list of - -- subexpressions specify how much each dimension is rotated. The - -- length of this list must be equal to the rank of the array. - deriving (Eq, Ord, Show) +instance SexpIso BasicOp where + sexpIso = + match $ + With (. Sexp.list (Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "opaque") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "array") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "assert") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el assertHelper)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "index") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "update") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "concat") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "copy") >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "manifest") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "iota") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "replicate") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "scratch") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "reshape") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "rearrange") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "rotate") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End + where + assertHelper = + with $ \tuple -> + Sexp.list + ( Sexp.el (iso (const mempty) (T.pack . show) . Sexp.symbol) + >>> Sexp.rest (iso (const mempty) (T.pack . show) . Sexp.symbol) + ) + >>> tuple -- | The root Futhark expression type. The v'Op' constructor contains -- a lore-specific operation. Do-loops, branches and function calls -- are special. Everything else is a simple t'BasicOp'. data ExpT lore - = BasicOp BasicOp - -- ^ A simple (non-recursive) operation. - - | Apply Name [(SubExp, Diet)] [RetType lore] (Safety, SrcLoc, [SrcLoc]) - - | If SubExp (BodyT lore) (BodyT lore) (IfDec (BranchType lore)) - - | DoLoop [(FParam lore, SubExp)] [(FParam lore, SubExp)] (LoopForm lore) (BodyT lore) - -- ^ @loop {a} = {v} (for i < n|while b) do b@. The merge + = -- | A simple (non-recursive) operation. + BasicOp BasicOp + | Apply Name [(SubExp, Diet)] [RetType lore] (Safety, SrcLoc, [SrcLoc]) + | If SubExp (BodyT lore) (BodyT lore) (IfDec (BranchType lore)) + | -- | @loop {a} = {v} (for i < n|while b) do b@. The merge -- parameters are divided into context and value part. - + DoLoop [(FParam lore, SubExp)] [(FParam lore, SubExp)] (LoopForm lore) (BodyT lore) | Op (Op lore) + deriving (Generic) +instance Decorations lore => SexpIso (ExpT lore) where + sexpIso = + match $ + With (. sexpIso) $ + With (. Sexp.list (Sexp.el (Sexp.sym "apply") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el applyHelper)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "if") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "loop") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. sexpIso) + End + where + applyHelper = + with $ \triple -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el (iso (const mempty) (T.pack . show) . Sexp.symbol) + >>> Sexp.rest (iso (const mempty) (T.pack . show) . Sexp.symbol) + ) + >>> triple + deriving instance Decorations lore => Eq (ExpT lore) + deriving instance Decorations lore => Show (ExpT lore) + deriving instance Decorations lore => Ord (ExpT lore) -- | For-loop or while-loop? -data LoopForm lore = ForLoop VName IntType SubExp [(LParam lore,VName)] - | WhileLoop VName +data LoopForm lore + = ForLoop VName IntType SubExp [(LParam lore, VName)] + | WhileLoop VName + deriving (Generic) +instance Decorations lore => SexpIso (LoopForm lore) where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "for") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "while") >>> Sexp.el sexpIso)) + End + deriving instance Decorations lore => Eq (LoopForm lore) + deriving instance Decorations lore => Show (LoopForm lore) + deriving instance Decorations lore => Ord (LoopForm lore) -- | Data associated with a branch. -data IfDec rt = IfDec { ifReturns :: [rt] - , ifSort :: IfSort - } - deriving (Eq, Show, Ord) +data IfDec rt = IfDec + { ifReturns :: [rt], + ifSort :: IfSort + } + deriving (Eq, Show, Ord, Generic) +instance SexpIso rt => SexpIso (IfDec rt) where + sexpIso = with $ \stm -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> stm + -- | What kind of branch is this? This has no semantic meaning, but -- provides hints to simplifications. -data IfSort = IfNormal - -- ^ An ordinary branch. - | IfFallback - -- ^ A branch where the "true" case is what we are - -- actually interested in, and the "false" case is only - -- present as a fallback for when the true case cannot - -- be safely evaluated. The compiler is permitted to - -- optimise away the branch if the true case contains - -- only safe statements. - | IfEquiv - -- ^ Both of these branches are semantically equivalent, - -- and it is fine to eliminate one if it turns out to - -- have problems (e.g. contain things we cannot generate - -- code for). - deriving (Eq, Show, Ord) +data IfSort + = -- | An ordinary branch. + IfNormal + | -- | A branch where the "true" case is what we are + -- actually interested in, and the "false" case is only + -- present as a fallback for when the true case cannot + -- be safely evaluated. The compiler is permitted to + -- optimise away the branch if the true case contains + -- only safe statements. + IfFallback + | -- | Both of these branches are semantically equivalent, + -- and it is fine to eliminate one if it turns out to + -- have problems (e.g. contain things we cannot generate + -- code for). + IfEquiv + deriving (Eq, Show, Ord, Generic) +instance SexpIso IfSort where + sexpIso = + match $ + With (sym "normal" >>>) $ + With (sym "fallback" >>>) $ + With + (sym "equiv" >>>) + End + -- | A type alias for namespace control. type Exp = ExpT -- | Anonymous function for use in a SOAC. -data LambdaT lore = Lambda { lambdaParams :: [LParam lore] - , lambdaBody :: BodyT lore - , lambdaReturnType :: [Type] - } +data LambdaT lore = Lambda + { lambdaParams :: [LParam lore], + lambdaBody :: BodyT lore, + lambdaReturnType :: [Type] + } + deriving (Generic) +instance Decorations lore => SexpIso (LambdaT lore) where + sexpIso = with $ \lambdat -> + Sexp.list + ( Sexp.el (Sexp.sym "lambda") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> lambdat + deriving instance Decorations lore => Eq (LambdaT lore) + deriving instance Decorations lore => Show (LambdaT lore) + deriving instance Decorations lore => Ord (LambdaT lore) -- | Type alias for namespacing reasons. @@ -481,18 +642,35 @@ type LParam lore = Param (LParamInfo lore) -- | Function Declarations -data FunDef lore = FunDef { funDefEntryPoint :: Maybe EntryPoint - -- ^ Contains a value if this function is - -- an entry point. - , funDefAttrs :: Attrs - , funDefName :: Name - , funDefRetType :: [RetType lore] - , funDefParams :: [FParam lore] - , funDefBody :: BodyT lore - } +data FunDef lore = FunDef + { -- | Contains a value if this function is + -- an entry point. + funDefEntryPoint :: Maybe EntryPoint, + funDefAttrs :: Attrs, + funDefName :: Name, + funDefRetType :: [RetType lore], + funDefParams :: [FParam lore], + funDefBody :: BodyT lore + } + deriving (Generic) +instance Decorations lore => SexpIso (FunDef lore) where + sexpIso = with $ \fundef -> + Sexp.list + ( Sexp.el (Sexp.sym "fundef") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> fundef + deriving instance Decorations lore => Eq (FunDef lore) + deriving instance Decorations lore => Show (FunDef lore) + deriving instance Decorations lore => Ord (FunDef lore) -- | Information about the parameters and return value of an entry @@ -502,26 +680,45 @@ -- | Every entry point argument and return value has an annotation -- indicating how it maps to the original source program type. -data EntryPointType = TypeUnsigned - -- ^ Is an unsigned integer or array of unsigned - -- integers. - | TypeOpaque String Int - -- ^ A black box type comprising this many core - -- values. The string is a human-readable - -- description with no other semantics. - | TypeDirect - -- ^ Maps directly. - deriving (Eq, Show, Ord) +data EntryPointType + = -- | Is an unsigned integer or array of unsigned + -- integers. + TypeUnsigned + | -- | A black box type comprising this many core + -- values. The string is a human-readable + -- description with no other semantics. + TypeOpaque String Int + | -- | Maps directly. + TypeDirect + deriving (Eq, Show, Ord, Generic) +instance SexpIso EntryPointType where + sexpIso = + match $ + With (. Sexp.sym "unsigned") $ + With (. Sexp.list (Sexp.el (Sexp.sym "opaque") >>> Sexp.el (iso T.unpack T.pack . sexpIso) >>> Sexp.el sexpIso)) $ + With + (. Sexp.sym "direct") + End + -- | An entire Futhark program. data Prog lore = Prog - { progConsts :: Stms lore - -- ^ Top-level constants that are computed at program startup, and + { -- | Top-level constants that are computed at program startup, and -- which are in scope inside all functions. - - , progFuns :: [FunDef lore] - -- ^ The functions comprising the program. All funtions are also + progConsts :: Stms lore, + -- | The functions comprising the program. All funtions are also -- available in scope in the definitions of the constants, so be -- careful not to introduce circular dependencies (not currently -- checked). - } deriving (Eq, Ord, Show) + progFuns :: [FunDef lore] + } + deriving (Eq, Ord, Show, Generic) + +instance Decorations lore => SexpIso (Prog lore) where + sexpIso = with $ \prog -> + Sexp.list + ( Sexp.el (Sexp.sym "prog") + >>> Sexp.el sexpIso + >>> Sexp.rest sexpIso + ) + >>> prog 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 @@ -1,72 +1,83 @@ +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Safe #-} {-# LANGUAGE Strict #-} + -- | The most primitive ("core") aspects of the AST. Split out of -- "Futhark.IR.Syntax" in order for -- "Futhark.IR.Decorations" to use these definitions. This -- module is re-exported from "Futhark.IR.Syntax" and -- there should be no reason to include it explicitly. module Futhark.IR.Syntax.Core - ( - module Language.Futhark.Core - , module Futhark.IR.Primitive + ( module Language.Futhark.Core, + module Futhark.IR.Primitive, - -- * Types - , Uniqueness(..) - , NoUniqueness(..) - , ShapeBase(..) - , Shape - , Ext(..) - , ExtSize - , ExtShape - , Rank(..) - , ArrayShape(..) - , Space (..) - , SpaceId - , TypeBase(..) - , Type - , ExtType - , DeclType - , DeclExtType - , Diet(..) - , ErrorMsg (..) - , ErrorMsgPart (..) - , errorMsgArgTypes + -- * Types + Uniqueness (..), + NoUniqueness (..), + ShapeBase (..), + Shape, + Ext (..), + ExtSize, + ExtShape, + Rank (..), + ArrayShape (..), + Space (..), + SpaceId, + TypeBase (..), + Type, + ExtType, + DeclType, + DeclExtType, + Diet (..), + ErrorMsg (..), + ErrorMsgPart (..), + errorMsgArgTypes, - -- * Values - , PrimValue(..) + -- * Values + PrimValue (..), - -- * Abstract syntax tree - , Ident (..) - , Certificates(..) - , SubExp(..) - , Param (..) - , DimIndex (..) - , Slice - , dimFix - , sliceIndices - , sliceDims - , unitSlice - , fixSlice - , sliceSlice - , PatElemT (..) - ) where + -- * Abstract syntax tree + Ident (..), + Certificates (..), + SubExp (..), + Param (..), + DimIndex (..), + Slice, + dimFix, + sliceIndices, + sliceDims, + unitSlice, + fixSlice, + sliceSlice, + PatElemT (..), + ) +where +import Control.Category import Control.Monad.State +import qualified Data.Map.Strict as M import Data.Maybe import Data.String -import qualified Data.Map.Strict as M +import qualified Data.Text as T import Data.Traversable (fmapDefault, foldMapDefault) - -import Language.Futhark.Core import Futhark.IR.Primitive +import GHC.Generics +import Language.Futhark.Core +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) -- | The size of an array type as a list of its dimension sizes, with -- the type of sizes being parametric. -newtype ShapeBase d = Shape { shapeDims :: [d] } - deriving (Eq, Ord, Show) +newtype ShapeBase d = Shape {shapeDims :: [d]} + deriving (Eq, Ord, Show, Generic) +instance SexpIso d => SexpIso (ShapeBase d) where + sexpIso = with $ \vname -> sexpIso >>> vname + instance Functor ShapeBase where fmap = fmapDefault @@ -87,10 +98,19 @@ type Shape = ShapeBase SubExp -- | Something that may be existential. -data Ext a = Ext Int - | Free a - deriving (Eq, Ord, Show) +data Ext a + = Ext Int + | Free a + deriving (Eq, Ord, Show, Generic) +instance SexpIso a => SexpIso (Ext a) where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "ext") >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "free") >>> Sexp.el sexpIso)) + End + instance Functor Ext where fmap = fmapDefault @@ -111,15 +131,25 @@ -- | The size of an array type as merely the number of dimensions, -- with no further information. newtype Rank = Rank Int - deriving (Show, Eq, Ord) + deriving (Show, Eq, Ord, Generic) +instance SexpIso Rank where + sexpIso = with $ \rank -> + Sexp.list + ( Sexp.el (Sexp.sym "rank") + >>> Sexp.el sexpIso + ) + >>> rank + -- | A class encompassing types containing array shape information. class (Monoid a, Eq a, Ord a) => ArrayShape a where -- | Return the rank of an array with the given size. shapeRank :: a -> Int + -- | @stripDims n shape@ strips the outer @n@ dimensions from -- @shape@. stripDims :: Int -> a -> a + -- | Check whether one shape if a subset of another shape. subShapeOf :: a -> a -> Bool @@ -134,18 +164,21 @@ subShapeOf (Shape ds1) (Shape ds2) = -- Must agree on Free dimensions, and ds1 may not be existential -- where ds2 is Free. Existentials must also be congruent. - length ds1 == length ds2 && - evalState (and <$> zipWithM subDimOf ds1 ds2) M.empty - where subDimOf (Free se1) (Free se2) = return $ se1 == se2 - subDimOf (Ext _) (Free _) = return False - subDimOf (Free _) (Ext _) = return True - subDimOf (Ext x) (Ext y) = do - extmap <- get - case M.lookup y extmap of - Just ywas | ywas == x -> return True - | otherwise -> return False - Nothing -> do put $ M.insert y x extmap - return True + length ds1 == length ds2 + && evalState (and <$> zipWithM subDimOf ds1 ds2) M.empty + where + subDimOf (Free se1) (Free se2) = return $ se1 == se2 + subDimOf (Ext _) (Free _) = return False + subDimOf (Free _) (Ext _) = return True + subDimOf (Ext x) (Ext y) = do + extmap <- get + case M.lookup y extmap of + Just ywas + | ywas == x -> return True + | otherwise -> return False + Nothing -> do + put $ M.insert y x extmap + return True instance Semigroup Rank where Rank x <> Rank y = Rank $ x + y @@ -164,28 +197,51 @@ -- used to distinguish between constant, global and shared memory -- spaces. In GPU-enabled host code, it is used to distinguish -- between host memory ('DefaultSpace') and GPU space. -data Space = DefaultSpace - | Space SpaceId - | ScalarSpace [SubExp] PrimType - -- ^ A special kind of memory that is a statically sized - -- array of some primitive type. Used for private memory - -- on GPUs. - deriving (Show, Eq, Ord) +data Space + = DefaultSpace + | Space SpaceId + | -- | A special kind of memory that is a statically sized + -- array of some primitive type. Used for private memory + -- on GPUs. + ScalarSpace [SubExp] PrimType + deriving (Show, Eq, Ord, Generic) +instance SexpIso Space where + sexpIso = + match $ + With (Sexp.sym "default" >>>) $ + With (. Sexp.list (Sexp.el (sym "space") >>> Sexp.el (iso T.unpack T.pack . sexpIso))) $ + With + (. Sexp.list (Sexp.el (sym "scalar-space") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End + -- | A string representing a specific non-default memory space. type SpaceId = String -- | A fancier name for @()@ - encodes no uniqueness information. data NoUniqueness = NoUniqueness - deriving (Eq, Ord, Show) + deriving (Eq, Ord, Show, Generic) +instance SexpIso NoUniqueness where + sexpIso = with (. sym "no-uniqueness") + -- | A Futhark type is either an array or an element type. When -- comparing types for equality with '==', shapes must match. -data TypeBase shape u = Prim PrimType - | Array PrimType shape u - | Mem Space - deriving (Show, Eq, Ord) +data TypeBase shape u + = Prim PrimType + | Array PrimType shape u + | Mem Space + deriving (Show, Eq, Ord, Generic) +instance (SexpIso shape, SexpIso u) => SexpIso (TypeBase shape u) where + sexpIso = + match $ + With (. sexpIso) $ + With (. Sexp.list (Sexp.el (sym "array") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (sym "mem") >>> Sexp.el sexpIso)) + End + -- | A type with shape information, used for describing the type of -- variables. type Type = TypeBase Shape NoUniqueness @@ -207,21 +263,44 @@ -- example, we might say that a function taking three arguments of -- types @([int], *[int], [int])@ has diet @[Observe, Consume, -- Observe]@. -data Diet = Consume -- ^ Consumes this value. - | Observe -- ^ Only observes value in this position, does - -- not consume. A result may alias this. - | ObservePrim -- ^ As 'Observe', but the result will not - -- alias, because the parameter does not carry - -- aliases. - deriving (Eq, Ord, Show) +data Diet + = -- | Consumes this value. + Consume + | -- | Only observes value in this position, does + -- not consume. A result may alias this. + Observe + | -- | As 'Observe', but the result will not + -- alias, because the parameter does not carry + -- aliases. + ObservePrim + deriving (Eq, Ord, Show, Generic) +instance SexpIso Diet where + sexpIso = + match $ + With (Sexp.sym "consume" >>>) $ + With (Sexp.sym "observe" >>>) $ + With + (Sexp.sym "observe-prim" >>>) + End + -- | An identifier consists of its name and the type of the value -- bound to the identifier. -data Ident = Ident { identName :: VName - , identType :: Type - } - deriving (Show) +data Ident = Ident + { identName :: VName, + identType :: Type + } + deriving (Show, Generic) +instance SexpIso Ident where + sexpIso = with $ \vname -> + Sexp.list + ( Sexp.el (Sexp.sym "ident") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> vname + instance Eq Ident where x == y = identName x == identName y @@ -229,9 +308,12 @@ x `compare` y = identName x `compare` identName y -- | A list of names used for certificates in some expressions. -newtype Certificates = Certificates { unCertificates :: [VName] } - deriving (Eq, Ord, Show) +newtype Certificates = Certificates {unCertificates :: [VName]} + deriving (Eq, Ord, Show, Generic) +instance SexpIso Certificates where + sexpIso = with $ \certificates -> sexpIso >>> certificates + instance Semigroup Certificates where Certificates x <> Certificates y = Certificates (x <> y) @@ -241,18 +323,37 @@ -- | A subexpression is either a scalar constant or a variable. One -- important property is that evaluation of a subexpression is -- guaranteed to complete in constant time. -data SubExp = Constant PrimValue - | Var VName - deriving (Show, Eq, Ord) +data SubExp + = Constant PrimValue + | Var VName + deriving (Show, Eq, Ord, Generic) +instance SexpIso SubExp where + sexpIso = + match $ + With (. sexpIso) $ + With + (. sexpIso) + End + -- | A function or lambda parameter. data Param dec = Param - { paramName :: VName - -- ^ Name of the parameter. - , paramDec :: dec - -- ^ Function parameter decoration. - } deriving (Ord, Show, Eq) + { -- | Name of the parameter. + paramName :: VName, + -- | Function parameter decoration. + paramDec :: dec + } + deriving (Ord, Show, Eq, Generic) +instance SexpIso dec => SexpIso (Param dec) where + sexpIso = with $ \vname -> + Sexp.list + ( Sexp.el (Sexp.sym "param") + >>> Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> vname + instance Foldable Param where foldMap = foldMapDefault @@ -263,12 +364,22 @@ traverse f (Param name dec) = Param name <$> f dec -- | How to index a single dimension of an array. -data DimIndex d = DimFix - d -- ^ Fix index in this dimension. - | DimSlice d d d - -- ^ @DimSlice start_offset num_elems stride@. - deriving (Eq, Ord, Show) +data DimIndex d + = -- | Fix index in this dimension. + DimFix + d + | -- | @DimSlice start_offset num_elems stride@. + DimSlice d d d + deriving (Eq, Ord, Show, Generic) +instance SexpIso d => SexpIso (DimIndex d) where + sexpIso = + match $ + With (. sexpIso) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "slice") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End + instance Functor DimIndex where fmap f (DimFix i) = DimFix $ f i fmap f (DimSlice i j s) = DimSlice (f i) (f j) (f s) @@ -299,8 +410,9 @@ -- | The dimensions of the array produced by this slice. sliceDims :: Slice d -> [d] sliceDims = mapMaybe dimSlice - where dimSlice (DimSlice _ d _) = Just d - dimSlice DimFix{} = Nothing + where + dimSlice (DimSlice _ d _) = Just d + dimSlice DimFix {} = Nothing -- | A slice with a stride of one. unitSlice :: Num d => d -> d -> DimIndex d @@ -309,33 +421,42 @@ -- | Fix the 'DimSlice's of a slice. The number of indexes must equal -- the length of 'sliceDims' for the slice. fixSlice :: Num d => Slice d -> [d] -> [d] -fixSlice (DimFix j:mis') is' = +fixSlice (DimFix j : mis') is' = j : fixSlice mis' is' -fixSlice (DimSlice orig_k _ orig_s:mis') (i:is') = - (orig_k+i*orig_s) : fixSlice mis' is' +fixSlice (DimSlice orig_k _ orig_s : mis') (i : is') = + (orig_k + i * orig_s) : fixSlice mis' is' fixSlice _ _ = [] -- | Further slice the 'DimSlice's of a slice. The number of slices -- must equal the length of 'sliceDims' for the slice. sliceSlice :: Num d => Slice d -> Slice d -> Slice d -sliceSlice (DimFix j:js') is' = +sliceSlice (DimFix j : js') is' = DimFix j : sliceSlice js' is' -sliceSlice (DimSlice j _ s:js') (DimFix i:is') = - DimFix (j + (i*s)) : sliceSlice js' is' -sliceSlice (DimSlice j _ s0:js') (DimSlice i n s1:is') = - DimSlice (j+(s0*i)) n (s0*s1) : sliceSlice js' is' +sliceSlice (DimSlice j _ s : js') (DimFix i : is') = + DimFix (j + (i * s)) : sliceSlice js' is' +sliceSlice (DimSlice j _ s0 : js') (DimSlice i n s1 : is') = + DimSlice (j + (s0 * i)) n (s0 * s1) : sliceSlice js' is' sliceSlice _ _ = [] -- | An element of a pattern - consisting of a name and an addditional -- parametric decoration. This decoration is what is expected to -- contain the type of the resulting variable. -data PatElemT dec = PatElem { patElemName :: VName - -- ^ The name being bound. - , patElemDec :: dec - -- ^ Pattern element decoration. - } - deriving (Ord, Show, Eq) +data PatElemT dec = PatElem + { -- | The name being bound. + patElemName :: VName, + -- | Pattern element decoration. + patElemDec :: dec + } + deriving (Ord, Show, Eq, Generic) +instance SexpIso dec => SexpIso (PatElemT dec) where + sexpIso = with $ \pe -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> pe + instance Functor PatElemT where fmap = fmapDefault @@ -349,16 +470,33 @@ -- | An error message is a list of error parts, which are concatenated -- to form the final message. newtype ErrorMsg a = ErrorMsg [ErrorMsgPart a] - deriving (Eq, Ord, Show) + deriving (Eq, Ord, Show, Generic) +instance SexpIso a => SexpIso (ErrorMsg a) where + sexpIso = with $ \errormsg -> sexpIso >>> errormsg + instance IsString (ErrorMsg a) where fromString = ErrorMsg . pure . fromString -- | A part of an error message. -data ErrorMsgPart a = ErrorString String -- ^ A literal string. - | ErrorInt32 a -- ^ A run-time integer value. - deriving (Eq, Ord, Show) +data ErrorMsgPart a + = -- | A literal string. + ErrorString String + | -- | A run-time integer value. + ErrorInt32 a + | -- | A bigger run-time integer value. + ErrorInt64 a + deriving (Eq, Ord, Show, Generic) +instance SexpIso a => SexpIso (ErrorMsgPart a) where + sexpIso = + match $ + With (. Sexp.list (Sexp.el (Sexp.sym "error-string") . Sexp.el (iso T.unpack T.pack . sexpIso))) $ + With (. Sexp.list (Sexp.el (Sexp.sym "error-int32") . Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "error-int64") . Sexp.el sexpIso)) + End + instance IsString (ErrorMsgPart a) where fromString = ErrorString @@ -374,18 +512,23 @@ instance Functor ErrorMsgPart where fmap _ (ErrorString s) = ErrorString s fmap f (ErrorInt32 a) = ErrorInt32 $ f a + fmap f (ErrorInt64 a) = ErrorInt64 $ f a instance Foldable ErrorMsgPart where - foldMap _ ErrorString{} = mempty + foldMap _ ErrorString {} = mempty foldMap f (ErrorInt32 a) = f a + foldMap f (ErrorInt64 a) = f a instance Traversable ErrorMsgPart where traverse _ (ErrorString s) = pure $ ErrorString s traverse f (ErrorInt32 a) = ErrorInt32 <$> f a + traverse f (ErrorInt64 a) = ErrorInt64 <$> f a -- | How many non-constant parts does the error message have, and what -- is their type? errorMsgArgTypes :: ErrorMsg a -> [PrimType] errorMsgArgTypes (ErrorMsg parts) = mapMaybe onPart parts - where onPart ErrorString{} = Nothing - onPart ErrorInt32{} = Just $ IntType Int32 + where + onPart ErrorString {} = Nothing + onPart ErrorInt32 {} = Just $ IntType Int32 + onPart ErrorInt64 {} = Just $ IntType Int64 diff --git a/src/Futhark/IR/Traversals.hs b/src/Futhark/IR/Traversals.hs --- a/src/Futhark/IR/Traversals.hs +++ b/src/Futhark/IR/Traversals.hs @@ -1,4 +1,5 @@ {-# LANGUAGE Safe #-} + -- | -- -- Functions for generic traversals across Futhark syntax trees. The @@ -23,68 +24,72 @@ -- The "Futhark.Transform.Rename" module is a simple example of how to -- use this facility. module Futhark.IR.Traversals - ( - -- * Mapping - Mapper(..) - , identityMapper - , mapExpM - , mapExp + ( -- * Mapping + Mapper (..), + identityMapper, + mapExpM, + mapExp, - -- * Walking - , Walker(..) - , identityWalker - , walkExpM + -- * Walking + Walker (..), + identityWalker, + walkExpM, ) - where +where import Control.Monad import Control.Monad.Identity -import qualified Data.Traversable import Data.Foldable (traverse_) - -import Futhark.IR.Syntax +import qualified Data.Traversable import Futhark.IR.Prop.Scope import Futhark.IR.Prop.Types (mapOnType) +import Futhark.IR.Syntax -- | Express a monad mapping operation on a syntax node. Each element -- of this structure expresses the operation to be performed on a -- given child. -data Mapper flore tlore m = Mapper { - mapOnSubExp :: SubExp -> m SubExp - , mapOnBody :: Scope tlore -> Body flore -> m (Body tlore) - -- ^ Most bodies are enclosed in a scope, which is passed along +data Mapper flore tlore m = Mapper + { mapOnSubExp :: SubExp -> m SubExp, + -- | Most bodies are enclosed in a scope, which is passed along -- for convenience. - , mapOnVName :: VName -> m VName - , mapOnRetType :: RetType flore -> m (RetType tlore) - , mapOnBranchType :: BranchType flore -> m (BranchType tlore) - , mapOnFParam :: FParam flore -> m (FParam tlore) - , mapOnLParam :: LParam flore -> m (LParam tlore) - , mapOnOp :: Op flore -> m (Op tlore) + mapOnBody :: Scope tlore -> Body flore -> m (Body tlore), + mapOnVName :: VName -> m VName, + mapOnRetType :: RetType flore -> m (RetType tlore), + mapOnBranchType :: BranchType flore -> m (BranchType tlore), + mapOnFParam :: FParam flore -> m (FParam tlore), + mapOnLParam :: LParam flore -> m (LParam tlore), + mapOnOp :: Op flore -> m (Op tlore) } -- | A mapper that simply returns the tree verbatim. identityMapper :: Monad m => Mapper lore lore m -identityMapper = Mapper { - mapOnSubExp = return - , mapOnBody = const return - , mapOnVName = return - , mapOnRetType = return - , mapOnBranchType = return - , mapOnFParam = return - , mapOnLParam = return - , mapOnOp = return - } +identityMapper = + Mapper + { mapOnSubExp = return, + mapOnBody = const return, + mapOnVName = return, + mapOnRetType = return, + mapOnBranchType = return, + mapOnFParam = return, + mapOnLParam = return, + mapOnOp = return + } -- | Map a monadic action across the immediate children of an -- expression. Importantly, the mapping does not descend recursively -- into subexpressions. The mapping is done left-to-right. -mapExpM :: (Applicative m, Monad m) => - Mapper flore tlore m -> Exp flore -> m (Exp tlore) +mapExpM :: + (Applicative m, Monad m) => + Mapper flore tlore m -> + Exp flore -> + m (Exp tlore) mapExpM tv (BasicOp (SubExp se)) = BasicOp <$> (SubExp <$> mapOnSubExp tv se) mapExpM tv (BasicOp (ArrayLit els rowt)) = - BasicOp <$> (ArrayLit <$> mapM (mapOnSubExp tv) els <*> - mapOnType (mapOnSubExp tv) rowt) + BasicOp + <$> ( ArrayLit <$> mapM (mapOnSubExp tv) els + <*> mapOnType (mapOnSubExp tv) rowt + ) mapExpM tv (BasicOp (BinOp bop x y)) = BasicOp <$> (BinOp bop <$> mapOnSubExp tv x <*> mapOnSubExp tv y) mapExpM tv (BasicOp (CmpOp op x y)) = @@ -94,18 +99,23 @@ mapExpM tv (BasicOp (UnOp op x)) = BasicOp <$> (UnOp op <$> mapOnSubExp tv x) mapExpM tv (If c texp fexp (IfDec ts s)) = - If <$> mapOnSubExp tv c <*> mapOnBody tv mempty texp <*> mapOnBody tv mempty fexp <*> - (IfDec <$> mapM (mapOnBranchType tv) ts <*> pure s) + If <$> mapOnSubExp tv c <*> mapOnBody tv mempty texp <*> mapOnBody tv mempty fexp + <*> (IfDec <$> mapM (mapOnBranchType tv) ts <*> pure s) mapExpM tv (Apply fname args ret loc) = do args' <- forM args $ \(arg, d) -> - (,) <$> mapOnSubExp tv arg <*> pure d + (,) <$> mapOnSubExp tv arg <*> pure d Apply fname args' <$> mapM (mapOnRetType tv) ret <*> pure loc mapExpM tv (BasicOp (Index arr slice)) = - BasicOp <$> (Index <$> mapOnVName tv arr <*> - mapM (traverse (mapOnSubExp tv)) slice) + BasicOp + <$> ( Index <$> mapOnVName tv arr + <*> mapM (traverse (mapOnSubExp tv)) slice + ) mapExpM tv (BasicOp (Update arr slice se)) = - BasicOp <$> (Update <$> mapOnVName tv arr <*> - mapM (traverse (mapOnSubExp tv)) slice <*> mapOnSubExp tv se) + BasicOp + <$> ( Update <$> mapOnVName tv arr + <*> mapM (traverse (mapOnSubExp tv)) slice + <*> mapOnSubExp tv se + ) mapExpM tv (BasicOp (Iota n x s et)) = BasicOp <$> (Iota <$> mapOnSubExp tv n <*> mapOnSubExp tv x <*> mapOnSubExp tv s <*> pure et) mapExpM tv (BasicOp (Replicate shape vexp)) = @@ -113,17 +123,22 @@ mapExpM tv (BasicOp (Scratch t shape)) = BasicOp <$> (Scratch t <$> mapM (mapOnSubExp tv) shape) mapExpM tv (BasicOp (Reshape shape arrexp)) = - BasicOp <$> (Reshape <$> - mapM (Data.Traversable.traverse (mapOnSubExp tv)) shape <*> - mapOnVName tv arrexp) + BasicOp + <$> ( Reshape + <$> mapM (Data.Traversable.traverse (mapOnSubExp tv)) shape + <*> mapOnVName tv arrexp + ) mapExpM tv (BasicOp (Rearrange perm e)) = BasicOp <$> (Rearrange perm <$> mapOnVName tv e) mapExpM tv (BasicOp (Rotate es e)) = BasicOp <$> (Rotate <$> mapM (mapOnSubExp tv) es <*> mapOnVName tv e) mapExpM tv (BasicOp (Concat i x ys size)) = - BasicOp <$> (Concat i <$> - mapOnVName tv x <*> mapM (mapOnVName tv) ys <*> - mapOnSubExp tv size) + BasicOp + <$> ( Concat i + <$> mapOnVName tv x + <*> mapM (mapOnVName tv) ys + <*> mapOnSubExp tv size + ) mapExpM tv (BasicOp (Copy e)) = BasicOp <$> (Copy <$> mapOnVName tv e) mapExpM tv (BasicOp (Manifest perm e)) = @@ -136,25 +151,31 @@ ctxparams' <- mapM (mapOnFParam tv) ctxparams valparams' <- mapM (mapOnFParam tv) valparams form' <- mapOnLoopForm tv form - let scope = scopeOf form' <> scopeOfFParams (ctxparams'++valparams') - DoLoop <$> - (zip ctxparams' <$> mapM (mapOnSubExp tv) ctxinits) <*> - (zip valparams' <$> mapM (mapOnSubExp tv) valinits) <*> - pure form' <*> mapOnBody tv scope loopbody - where (ctxparams,ctxinits) = unzip ctxmerge - (valparams,valinits) = unzip valmerge + let scope = scopeOf form' <> scopeOfFParams (ctxparams' ++ valparams') + DoLoop + <$> (zip ctxparams' <$> mapM (mapOnSubExp tv) ctxinits) + <*> (zip valparams' <$> mapM (mapOnSubExp tv) valinits) + <*> pure form' + <*> mapOnBody tv scope loopbody + where + (ctxparams, ctxinits) = unzip ctxmerge + (valparams, valinits) = unzip valmerge mapExpM tv (Op op) = Op <$> mapOnOp tv op mapOnShape :: Monad m => Mapper flore tlore m -> Shape -> m Shape mapOnShape tv (Shape ds) = Shape <$> mapM (mapOnSubExp tv) ds -mapOnLoopForm :: Monad m => - Mapper flore tlore m -> LoopForm flore -> m (LoopForm tlore) +mapOnLoopForm :: + Monad m => + Mapper flore tlore m -> + LoopForm flore -> + m (LoopForm tlore) mapOnLoopForm tv (ForLoop i it bound loop_vars) = - ForLoop <$> mapOnVName tv i <*> pure it <*> mapOnSubExp tv bound <*> - (zip <$> mapM (mapOnLParam tv) loop_lparams <*> mapM (mapOnVName tv) loop_arrs) - where (loop_lparams,loop_arrs) = unzip loop_vars + ForLoop <$> mapOnVName tv i <*> pure it <*> mapOnSubExp tv bound + <*> (zip <$> mapM (mapOnLParam tv) loop_lparams <*> mapM (mapOnVName tv) loop_arrs) + where + (loop_lparams, loop_arrs) = unzip loop_vars mapOnLoopForm tv (WhileLoop cond) = WhileLoop <$> mapOnVName tv cond @@ -165,43 +186,46 @@ -- | Express a monad expression on a syntax node. Each element of -- this structure expresses the action to be performed on a given -- child. -data Walker lore m = Walker { - walkOnSubExp :: SubExp -> m () - , walkOnBody :: Scope lore -> Body lore -> m () - , walkOnVName :: VName -> m () - , walkOnRetType :: RetType lore -> m () - , walkOnBranchType :: BranchType lore -> m () - , walkOnFParam :: FParam lore -> m () - , walkOnLParam :: LParam lore -> m () - , walkOnOp :: Op lore -> m () +data Walker lore m = Walker + { walkOnSubExp :: SubExp -> m (), + walkOnBody :: Scope lore -> Body lore -> m (), + walkOnVName :: VName -> m (), + walkOnRetType :: RetType lore -> m (), + walkOnBranchType :: BranchType lore -> m (), + walkOnFParam :: FParam lore -> m (), + walkOnLParam :: LParam lore -> m (), + walkOnOp :: Op lore -> m () } -- | A no-op traversal. identityWalker :: Monad m => Walker lore m -identityWalker = Walker { - walkOnSubExp = const $ return () - , walkOnBody = const $ const $ return () - , walkOnVName = const $ return () - , walkOnRetType = const $ return () - , walkOnBranchType = const $ return () - , walkOnFParam = const $ return () - , walkOnLParam = const $ return () - , walkOnOp = const $ return () - } +identityWalker = + Walker + { walkOnSubExp = const $ return (), + walkOnBody = const $ const $ return (), + walkOnVName = const $ return (), + walkOnRetType = const $ return (), + walkOnBranchType = const $ return (), + walkOnFParam = const $ return (), + walkOnLParam = const $ return (), + walkOnOp = const $ return () + } walkOnShape :: Monad m => Walker lore m -> Shape -> m () walkOnShape tv (Shape ds) = mapM_ (walkOnSubExp tv) ds walkOnType :: Monad m => Walker lore m -> Type -> m () -walkOnType _ Prim{} = return () -walkOnType _ Mem{} = return () +walkOnType _ Prim {} = return () +walkOnType _ Mem {} = return () walkOnType tv (Array _ shape _) = walkOnShape tv shape walkOnLoopForm :: Monad m => Walker lore m -> LoopForm lore -> m () walkOnLoopForm tv (ForLoop i _ bound loop_vars) = - walkOnVName tv i >> walkOnSubExp tv bound >> - mapM_ (walkOnLParam tv) loop_lparams >> mapM_ (walkOnVName tv) loop_arrs - where (loop_lparams,loop_arrs) = unzip loop_vars + walkOnVName tv i >> walkOnSubExp tv bound + >> mapM_ (walkOnLParam tv) loop_lparams + >> mapM_ (walkOnVName tv) loop_arrs + where + (loop_lparams, loop_arrs) = unzip loop_vars walkOnLoopForm tv (WhileLoop cond) = walkOnVName tv cond @@ -229,9 +253,9 @@ walkExpM tv (BasicOp (Index arr slice)) = walkOnVName tv arr >> mapM_ (traverse_ (walkOnSubExp tv)) slice walkExpM tv (BasicOp (Update arr slice se)) = - walkOnVName tv arr >> - mapM_ (traverse_ (walkOnSubExp tv)) slice >> - walkOnSubExp tv se + walkOnVName tv arr + >> mapM_ (traverse_ (walkOnSubExp tv)) slice + >> walkOnSubExp tv se walkExpM tv (BasicOp (Iota n x s _)) = walkOnSubExp tv n >> walkOnSubExp tv x >> walkOnSubExp tv s walkExpM tv (BasicOp (Replicate shape vexp)) = @@ -260,9 +284,10 @@ walkOnLoopForm tv form mapM_ (walkOnSubExp tv) ctxinits mapM_ (walkOnSubExp tv) valinits - let scope = scopeOfFParams (ctxparams++valparams) <> scopeOf form + let scope = scopeOfFParams (ctxparams ++ valparams) <> scopeOf form walkOnBody tv scope loopbody - where (ctxparams,ctxinits) = unzip ctxmerge - (valparams,valinits) = unzip valmerge + where + (ctxparams, ctxinits) = unzip ctxmerge + (valparams, valinits) = unzip valmerge walkExpM tv (Op op) = walkOnOp tv op diff --git a/src/Futhark/Internalise.hs b/src/Futhark/Internalise.hs --- a/src/Futhark/Internalise.hs +++ b/src/Futhark/Internalise.hs @@ -1,1900 +1,2181 @@ -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE Strict #-} -{-# LANGUAGE Safe #-} --- | --- --- This module implements a transformation from source to core --- Futhark. --- -module Futhark.Internalise (internaliseProg) where - -import Control.Monad.State -import Control.Monad.Reader -import Data.Bitraversable -import qualified Data.Map.Strict as M -import qualified Data.Set as S -import Data.List (find, intercalate, intersperse, nub, transpose) -import qualified Data.List.NonEmpty as NE - -import Language.Futhark as E hiding (TypeArg) -import Language.Futhark.Semantic (Imports) -import Futhark.IR.SOACS as I hiding (stmPattern) -import Futhark.Transform.Rename as I -import Futhark.MonadFreshNames -import Futhark.Tools -import Futhark.Util (splitAt3) - -import Futhark.Internalise.Monad as I -import Futhark.Internalise.AccurateSizes -import Futhark.Internalise.TypesValues -import Futhark.Internalise.Bindings -import Futhark.Internalise.Lambdas -import Futhark.Internalise.Defunctorise as Defunctorise -import Futhark.Internalise.Defunctionalise as Defunctionalise -import Futhark.Internalise.Monomorphise as Monomorphise - --- | Convert a program in source Futhark to a program in the Futhark --- core language. -internaliseProg :: MonadFreshNames m => - Bool -> Imports -> m (I.Prog SOACS) -internaliseProg always_safe prog = do - prog_decs <- Defunctorise.transformProg prog - prog_decs' <- Monomorphise.transformProg prog_decs - prog_decs'' <- Defunctionalise.transformProg prog_decs' - (consts, funs) <- - runInternaliseM always_safe (internaliseValBinds prog_decs'') - I.renameProg $ I.Prog consts funs - -internaliseAttr :: E.AttrInfo -> Attr -internaliseAttr (E.AttrAtom v) = I.AttrAtom v -internaliseAttr (E.AttrComp f attrs) = I.AttrComp f $ map internaliseAttr attrs - -internaliseAttrs :: [E.AttrInfo] -> Attrs -internaliseAttrs = mconcat . map (oneAttr . internaliseAttr) - -internaliseValBinds :: [E.ValBind] -> InternaliseM () -internaliseValBinds = mapM_ internaliseValBind - -internaliseFunName :: VName -> [E.Pattern] -> InternaliseM Name -internaliseFunName ofname [] = return $ nameFromString $ pretty ofname ++ "f" -internaliseFunName ofname _ = do - info <- lookupFunction' ofname - -- In some rare cases involving local functions, the same function - -- name may be re-used in multiple places. We check whether the - -- function name has already been used, and generate a new one if - -- so. - case info of - Just _ -> nameFromString . pretty <$> newNameFromString (baseString ofname) - Nothing -> return $ nameFromString $ pretty ofname - -internaliseValBind :: E.ValBind -> InternaliseM () -internaliseValBind fb@(E.ValBind entry fname retdecl (Info (rettype, _)) tparams params body _ attrs loc) = do - localConstsScope $ bindingParams tparams params $ \shapeparams params' -> do - let shapenames = map I.paramName shapeparams - normal_params = shapenames ++ map I.paramName (concat params') - normal_param_names = namesFromList normal_params - - fname' <- internaliseFunName fname params - - msg <- case retdecl of - Just dt -> errorMsg . - ("Function return value does not match shape of type ":) <$> - typeExpForError dt - Nothing -> return $ errorMsg ["Function return value does not match shape of declared return type."] - - ((rettype', body_res), body_stms) <- collectStms $ do - body_res <- internaliseExp "res" body - rettype_bad <- internaliseReturnType rettype - let rettype' = zeroExts rettype_bad - return (rettype', body_res) - body' <- ensureResultExtShape msg loc (map I.fromDecl rettype') $ - mkBody body_stms body_res - - constants <- allConsts - let free_in_fun = freeIn body' - `namesSubtract` normal_param_names - `namesSubtract` constants - - used_free_params <- forM (namesToList free_in_fun) $ \v -> do - v_t <- lookupType v - return $ Param v $ toDecl v_t Nonunique - - let free_shape_params = map (`Param` I.Prim int32) $ - concatMap (I.shapeVars . I.arrayShape . I.paramType) used_free_params - free_params = nub $ free_shape_params ++ used_free_params - all_params = free_params ++ shapeparams ++ concat params' - - let fd = I.FunDef Nothing (internaliseAttrs attrs) fname' - rettype' all_params body' - - if null params' - then bindConstant fname fd - else bindFunction fname fd - (fname', - map I.paramName free_params, - shapenames, - map declTypeOf $ concat params', - all_params, - applyRetType rettype' all_params) - - case entry of Just (Info entry') -> generateEntryPoint entry' fb - Nothing -> return () - - where - -- | Recompute existential sizes to start from zero. - -- Necessary because some convoluted constructions will start - -- them from somewhere else. - zeroExts ts = generaliseExtTypes ts ts - -allDimsFreshInType :: MonadFreshNames m => E.PatternType -> m E.PatternType -allDimsFreshInType = bitraverse onDim pure - where onDim (E.NamedDim v) = - E.NamedDim . E.qualName <$> newVName (baseString $ E.qualLeaf v) - onDim _ = - E.NamedDim . E.qualName <$> newVName "size" - --- | Replace all named dimensions with a fresh name, and remove all --- constant dimensions. The point is to remove the constraints, but --- keep the names around. We use this for constructing the entry --- point parameters. -allDimsFreshInPat :: MonadFreshNames m => E.Pattern -> m E.Pattern -allDimsFreshInPat (PatternAscription p _ _) = - allDimsFreshInPat p -allDimsFreshInPat (PatternParens p _) = - allDimsFreshInPat p -allDimsFreshInPat (Id v (Info t) loc) = - Id v <$> (Info <$> allDimsFreshInType t) <*> pure loc -allDimsFreshInPat (TuplePattern ps loc) = - TuplePattern <$> mapM allDimsFreshInPat ps <*> pure loc -allDimsFreshInPat (RecordPattern ps loc) = - RecordPattern <$> mapM (traverse allDimsFreshInPat) ps <*> pure loc -allDimsFreshInPat (Wildcard (Info t) loc) = - Wildcard <$> (Info <$> allDimsFreshInType t) <*> pure loc -allDimsFreshInPat (PatternLit e (Info t) loc) = - PatternLit e <$> (Info <$> allDimsFreshInType t) <*> pure loc -allDimsFreshInPat (PatternConstr c (Info t) pats loc) = - PatternConstr c <$> (Info <$> allDimsFreshInType t) <*> - mapM allDimsFreshInPat pats <*> pure loc - -generateEntryPoint :: E.EntryPoint -> E.ValBind -> InternaliseM () -generateEntryPoint (E.EntryPoint e_paramts e_rettype) vb = localConstsScope $ do - let (E.ValBind _ ofname _ (Info (rettype, _)) _ params _ _ attrs loc) = vb - -- We replace all shape annotations, so there should be no constant - -- parameters here. - params_fresh <- mapM allDimsFreshInPat params - let tparams = map (`E.TypeParamDim` mempty) $ S.toList $ - mconcat $ map E.patternDimNames params_fresh - bindingParams tparams params_fresh $ \shapeparams params' -> do - entry_rettype <- internaliseEntryReturnType $ anySizes rettype - let entry' = entryPoint (zip e_paramts params') (e_rettype, entry_rettype) - args = map (I.Var . I.paramName) $ concat params' - - entry_body <- insertStmsM $ do - -- Special case the (rare) situation where the entry point is - -- not a function. - maybe_const <- lookupConst ofname - vals <- case maybe_const of - Just ses -> - return ses - Nothing -> - fst <$> funcall "entry_result" (E.qualName ofname) args loc - ctx <- extractShapeContext (concat entry_rettype) <$> - mapM (fmap I.arrayDims . subExpType) vals - resultBodyM (ctx ++ vals) - - addFunDef $ - I.FunDef (Just entry') (internaliseAttrs attrs) - (baseName ofname) - (concat entry_rettype) - (shapeparams ++ concat params') entry_body - -entryPoint :: [(E.EntryType, [I.FParam])] - -> (E.EntryType, - [[I.TypeBase ExtShape Uniqueness]]) - -> I.EntryPoint -entryPoint params (eret, crets) = - (concatMap (entryPointType . preParam) params, - case (isTupleRecord $ entryType eret, - entryAscribed eret) of - (Just ts, Just (E.TETuple e_ts _)) -> - concatMap entryPointType $ - zip (zipWith E.EntryType ts (map Just e_ts)) crets - (Just ts, Nothing) -> - concatMap entryPointType $ - zip (map (`E.EntryType` Nothing) ts) crets - _ -> - entryPointType (eret, concat crets)) - where preParam (e_t, ps) = (e_t, staticShapes $ map I.paramDeclType ps) - - entryPointType (t, ts) - | E.Scalar (E.Prim E.Unsigned{}) <- E.entryType t = - [I.TypeUnsigned] - | E.Array _ _ (E.Prim E.Unsigned{}) _ <- E.entryType t = - [I.TypeUnsigned] - | E.Scalar E.Prim{} <- E.entryType t = - [I.TypeDirect] - | E.Array _ _ E.Prim{} _ <- E.entryType t = - [I.TypeDirect] - | otherwise = - [I.TypeOpaque desc $ length ts] - where desc = maybe (pretty t') typeExpOpaqueName $ E.entryAscribed t - t' = noSizes (E.entryType t) `E.setUniqueness` Nonunique - - -- | We remove dimension arguments such that we hopefully end - -- up with a simpler type name for the entry point. The - -- intent is that if an entry point uses a type 'nasty [w] [h]', - -- then we should turn that into an opaque type just called - -- 'nasty'. Also, we try to give arrays of opaques a nicer name. - typeExpOpaqueName (TEApply te TypeArgExpDim{} _) = - typeExpOpaqueName te - typeExpOpaqueName (TEArray te _ _) = - let (d, te') = withoutDims te - in "arr_" ++ typeExpOpaqueName te' ++ - "_" ++ show (1 + d) ++ "d" - typeExpOpaqueName te = pretty te - - withoutDims (TEArray te _ _) = - let (d, te') = withoutDims te - in (d+1, te') - withoutDims te = (0::Int, te) - -internaliseIdent :: E.Ident -> InternaliseM I.VName -internaliseIdent (E.Ident name (Info tp) loc) = - case tp of - E.Scalar E.Prim{} -> return name - _ -> error $ "Futhark.Internalise.internaliseIdent: asked to internalise non-prim-typed ident '" - ++ pretty name ++ " of type " ++ pretty tp ++ - " at " ++ locStr loc ++ "." - -internaliseBody :: E.Exp -> InternaliseM Body -internaliseBody e = insertStmsM $ resultBody <$> internaliseExp "res" e - -bodyFromStms :: InternaliseM (Result, a) - -> InternaliseM (Body, a) -bodyFromStms m = do - ((res, a), stms) <- collectStms m - (,a) <$> mkBodyM stms res - -internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp] - -internaliseExp desc (E.Parens e _) = - internaliseExp desc e - -internaliseExp desc (E.QualParens _ e _) = - internaliseExp desc e - -internaliseExp desc (E.StringLit vs _) = - fmap pure $ letSubExp desc $ - I.BasicOp $ I.ArrayLit (map constant vs) $ I.Prim int8 - -internaliseExp _ (E.Var (E.QualName _ name) (Info t) loc) = do - subst <- lookupSubst name - case subst of - Just substs -> return substs - Nothing -> do - -- If this identifier is the name of a constant, we have to turn it - -- into a call to the corresponding function. - is_const <- lookupConst name - case is_const of - Just ses -> return ses - Nothing -> (:[]) . I.Var <$> internaliseIdent (E.Ident name (Info t) loc) - -internaliseExp desc (E.Index e idxs (Info ret, Info retext) loc) = do - vs <- internaliseExpToVars "indexed" e - dims <- case vs of [] -> return [] -- Will this happen? - v:_ -> I.arrayDims <$> lookupType v - (idxs', cs) <- internaliseSlice loc dims idxs - let index v = do v_t <- lookupType v - return $ I.BasicOp $ I.Index v $ fullSlice v_t idxs' - ses <- certifying cs $ letSubExps desc =<< mapM index vs - bindExtSizes (E.toStruct ret) retext ses - return ses - --- XXX: we map empty records and tuples to bools, because otherwise --- arrays of unit will lose their sizes. -internaliseExp _ (E.TupLit [] _) = - return [constant True] -internaliseExp _ (E.RecordLit [] _) = - return [constant True] - -internaliseExp desc (E.TupLit es _) = concat <$> mapM (internaliseExp desc) es - -internaliseExp desc (E.RecordLit orig_fields _) = - concatMap snd . sortFields . M.unions <$> mapM internaliseField orig_fields - where internaliseField (E.RecordFieldExplicit name e _) = - M.singleton name <$> internaliseExp desc e - internaliseField (E.RecordFieldImplicit name t loc) = - internaliseField $ E.RecordFieldExplicit (baseName name) - (E.Var (E.qualName name) t loc) loc - -internaliseExp desc (E.ArrayLit es (Info arr_t) loc) - -- If this is a multidimensional array literal of primitives, we - -- treat it specially by flattening it out followed by a reshape. - -- This cuts down on the amount of statements that are produced, and - -- thus allows us to efficiently handle huge array literals - a - -- corner case, but an important one. - | Just ((eshape,e'):es') <- mapM isArrayLiteral es, - not $ null eshape, - all ((eshape==) . fst) es', - Just basetype <- E.peelArray (length eshape) arr_t = do - let flat_lit = E.ArrayLit (e' ++ concatMap snd es') (Info basetype) loc - new_shape = length es:eshape - flat_arrs <- internaliseExpToVars "flat_literal" flat_lit - forM flat_arrs $ \flat_arr -> do - flat_arr_t <- lookupType flat_arr - let new_shape' = reshapeOuter (map (DimNew . constant) new_shape) - 1 $ I.arrayShape flat_arr_t - letSubExp desc $ I.BasicOp $ I.Reshape new_shape' flat_arr - - | otherwise = do - es' <- mapM (internaliseExp "arr_elem") es - arr_t_ext <- internaliseReturnType (E.toStruct arr_t) - - rowtypes <- - case mapM (fmap rowType . hasStaticShape . I.fromDecl) arr_t_ext of - Just ts -> pure ts - Nothing -> - -- XXX: the monomorphiser may create single-element array - -- literals with an unknown row type. In those cases we - -- need to look at the types of the actual elements. - -- Fixing this in the monomorphiser is a lot more tricky - -- than just working around it here. - case es' of - [] -> error $ "internaliseExp ArrayLit: existential type: " ++ pretty arr_t - e':_ -> mapM subExpType e' - - let arraylit ks rt = do - ks' <- mapM (ensureShape - "shape of element differs from shape of first element" - loc rt "elem_reshaped") ks - return $ I.BasicOp $ I.ArrayLit ks' rt - - letSubExps desc =<< - if null es' - then mapM (arraylit []) rowtypes - else zipWithM arraylit (transpose es') rowtypes - - where isArrayLiteral :: E.Exp -> Maybe ([Int],[E.Exp]) - isArrayLiteral (E.ArrayLit inner_es _ _) = do - (eshape,e):inner_es' <- mapM isArrayLiteral inner_es - guard $ all ((eshape==) . fst) inner_es' - return (length inner_es:eshape, e ++ concatMap snd inner_es') - isArrayLiteral e = - Just ([], [e]) - -internaliseExp desc (E.Range start maybe_second end (Info ret, Info retext) loc) = do - start' <- internaliseExp1 "range_start" start - end' <- internaliseExp1 "range_end" $ case end of - DownToExclusive e -> e - ToInclusive e -> e - UpToExclusive e -> e - maybe_second' <- - traverse (internaliseExp1 "range_second") maybe_second - - -- Construct an error message in case the range is invalid. - let conv = case E.typeOf start of - E.Scalar (E.Prim (E.Unsigned _)) -> asIntZ Int32 - _ -> asIntS Int32 - start'_i32 <- conv start' - end'_i32 <- conv end' - maybe_second'_i32 <- traverse conv maybe_second' - let errmsg = - errorMsg $ - ["Range "] ++ - [ErrorInt32 start'_i32] ++ - (case maybe_second'_i32 of - Nothing -> [] - Just second_i32 -> ["..", ErrorInt32 second_i32]) ++ - (case end of - DownToExclusive{} -> ["..>"] - ToInclusive{} -> ["..."] - UpToExclusive{} -> ["..<"]) ++ - [ErrorInt32 end'_i32, " is invalid."] - - (it, le_op, lt_op) <- - case E.typeOf start of - E.Scalar (E.Prim (E.Signed it)) -> return (it, CmpSle it, CmpSlt it) - E.Scalar (E.Prim (E.Unsigned it)) -> return (it, CmpUle it, CmpUlt it) - start_t -> error $ "Start value in range has type " ++ pretty start_t - - let one = intConst it 1 - negone = intConst it (-1) - default_step = case end of DownToExclusive{} -> negone - ToInclusive{} -> one - UpToExclusive{} -> one - - (step, step_zero) <- case maybe_second' of - Just second' -> do - subtracted_step <- letSubExp "subtracted_step" $ - I.BasicOp $ I.BinOp (I.Sub it I.OverflowWrap) second' start' - step_zero <- letSubExp "step_zero" $ I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) start' second' - return (subtracted_step, step_zero) - Nothing -> - return (default_step, constant False) - - step_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum it) step - step_sign_i32 <- asIntS Int32 step_sign - - bounds_invalid_downwards <- letSubExp "bounds_invalid_downwards" $ - I.BasicOp $ I.CmpOp le_op start' end' - bounds_invalid_upwards <- letSubExp "bounds_invalid_upwards" $ - I.BasicOp $ I.CmpOp lt_op end' start' - - (distance, step_wrong_dir, bounds_invalid) <- case end of - DownToExclusive{} -> do - step_wrong_dir <- letSubExp "step_wrong_dir" $ - I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) step_sign one - distance <- letSubExp "distance" $ - I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) start' end' - distance_i32 <- asIntZ Int32 distance - return (distance_i32, step_wrong_dir, bounds_invalid_downwards) - UpToExclusive{} -> do - step_wrong_dir <- letSubExp "step_wrong_dir" $ - I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) step_sign negone - distance <- letSubExp "distance" $ I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) end' start' - distance_i32 <- asIntZ Int32 distance - return (distance_i32, step_wrong_dir, bounds_invalid_upwards) - ToInclusive{} -> do - downwards <- letSubExp "downwards" $ - I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) step_sign negone - distance_downwards_exclusive <- - letSubExp "distance_downwards_exclusive" $ - I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) start' end' - distance_upwards_exclusive <- - letSubExp "distance_upwards_exclusive" $ - I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) end' start' - - bounds_invalid <- letSubExp "bounds_invalid" $ - I.If downwards - (resultBody [bounds_invalid_downwards]) - (resultBody [bounds_invalid_upwards]) $ - ifCommon [I.Prim I.Bool] - distance_exclusive <- letSubExp "distance_exclusive" $ - I.If downwards - (resultBody [distance_downwards_exclusive]) - (resultBody [distance_upwards_exclusive]) $ - ifCommon [I.Prim $ IntType it] - distance_exclusive_i32 <- asIntZ Int32 distance_exclusive - distance <- letSubExp "distance" $ - I.BasicOp $ I.BinOp (Add Int32 I.OverflowWrap) - distance_exclusive_i32 (intConst Int32 1) - return (distance, constant False, bounds_invalid) - - step_invalid <- letSubExp "step_invalid" $ - I.BasicOp $ I.BinOp I.LogOr step_wrong_dir step_zero - - invalid <- letSubExp "range_invalid" $ - I.BasicOp $ I.BinOp I.LogOr step_invalid bounds_invalid - valid <- letSubExp "valid" $ I.BasicOp $ I.UnOp I.Not invalid - cs <- assert "range_valid_c" valid errmsg loc - - step_i32 <- asIntS Int32 step - pos_step <- letSubExp "pos_step" $ - I.BasicOp $ I.BinOp (Mul Int32 I.OverflowWrap) step_i32 step_sign_i32 - - num_elems <- certifying cs $ - letSubExp "num_elems" $ - I.BasicOp $ I.BinOp (SDivUp Int32 I.Unsafe) distance pos_step - - se <- letSubExp desc (I.BasicOp $ I.Iota num_elems start' step it) - bindExtSizes (E.toStruct ret) retext [se] - return [se] - -internaliseExp desc (E.Ascript e _ _) = - internaliseExp desc e - -internaliseExp desc (E.Coerce e (TypeDecl dt (Info et)) (Info ret, Info retext) loc) = do - ses <- internaliseExp desc e - ts <- internaliseReturnType et - dt' <- typeExpForError dt - bindExtSizes (E.toStruct ret) retext ses - forM (zip ses ts) $ \(e',t') -> do - dims <- arrayDims <$> subExpType e' - let parts = ["Value of (core language) shape ("] ++ - intersperse ", " (map ErrorInt32 dims) ++ - [") cannot match shape of type `"] ++ dt' ++ ["`."] - ensureExtShape (errorMsg parts) loc (I.fromDecl t') desc e' - -internaliseExp desc (E.Negate e _) = do - e' <- internaliseExp1 "negate_arg" e - et <- subExpType e' - case et of I.Prim (I.IntType t) -> - letTupExp' desc $ I.BasicOp $ I.BinOp (I.Sub t I.OverflowWrap) (I.intConst t 0) e' - I.Prim (I.FloatType t) -> - letTupExp' desc $ I.BasicOp $ I.BinOp (I.FSub t) (I.floatConst t 0) e' - _ -> error "Futhark.Internalise.internaliseExp: non-numeric type in Negate" - -internaliseExp desc e@E.Apply{} = do - (qfname, args, ret, retext) <- findFuncall e - -- Argument evaluation is outermost-in so that any existential sizes - -- created by function applications can be brought into scope. - let fname = nameFromString $ pretty $ baseName $ qualLeaf qfname - loc = srclocOf e - arg_desc = nameToString fname ++ "_arg" - - -- Some functions are magical (overloaded) and we handle that here. - ses <- - case () of - -- Overloaded functions never take array arguments (except - -- equality, but those cannot be existential), so we can safely - -- ignore the existential dimensions. - () | Just internalise <- isOverloadedFunction qfname (map fst args) loc -> - internalise desc - - | Just (rettype, _) <- M.lookup fname I.builtInFunctions -> do - let tag ses = [ (se, I.Observe) | se <- ses ] - args' <- reverse <$> mapM (internaliseArg arg_desc) (reverse args) - let args'' = concatMap tag args' - letTupExp' desc $ I.Apply fname args'' [I.Prim rettype] - (Safe, loc, []) - - | otherwise -> do - args' <- concat . reverse <$> mapM (internaliseArg arg_desc) (reverse args) - fst <$> funcall desc qfname args' loc - - bindExtSizes ret retext ses - return ses - -internaliseExp desc (E.LetPat pat e body (Info ret, Info retext) _) = do - ses <- internalisePat desc pat e body (internaliseExp desc) - bindExtSizes (E.toStruct ret) retext ses - return ses - -internaliseExp desc (E.LetFun ofname (tparams, params, retdecl, Info rettype, body) letbody _ loc) = do - internaliseValBind $ - E.ValBind Nothing ofname retdecl (Info (rettype, [])) tparams params body Nothing mempty loc - internaliseExp desc letbody - -internaliseExp desc (E.DoLoop sparams mergepat mergeexp form loopbody (Info (ret, retext)) loc) = do - ses <- internaliseExp "loop_init" mergeexp - ((loopbody', (form', shapepat, mergepat', mergeinit')), initstms) <- - collectStms $ handleForm ses form - - addStms initstms - mergeinit_ts' <- mapM subExpType mergeinit' - - ctxinit <- argShapes (map I.paramName shapepat) mergepat' mergeinit_ts' - - let ctxmerge = zip shapepat ctxinit - valmerge = zip mergepat' mergeinit' - dropCond = case form of E.While{} -> drop 1 - _ -> id - - -- Ensure that the result of the loop matches the shapes of the - -- merge parameters. XXX: Ideally they should already match (by - -- the source language type rules), but some of our - -- transformations (esp. defunctionalisation) strips out some size - -- information. For a type-correct source program, these reshapes - -- should simplify away. - let merge = ctxmerge ++ valmerge - merge_ts = map (I.paramType . fst) merge - loopbody'' <- localScope (scopeOfFParams $ map fst merge) $ - inScopeOf form' $ insertStmsM $ - resultBodyM - =<< ensureArgShapes - "shape of loop result does not match shapes in loop parameter" - loc (map (I.paramName . fst) ctxmerge) merge_ts - =<< bodyBind loopbody' - - attrs <- asks envAttrs - loop_res <- map I.Var . dropCond <$> attributing attrs - (letTupExp desc (I.DoLoop ctxmerge valmerge form' loopbody'')) - bindExtSizes (E.toStruct ret) retext loop_res - return loop_res - - where - sparams' = map (`TypeParamDim` mempty) sparams - - forLoop mergepat' shapepat mergeinit form' = - bodyFromStms $ inScopeOf form' $ do - ses <- internaliseExp "loopres" loopbody - sets <- mapM subExpType ses - shapeargs <- argShapes (map I.paramName shapepat) mergepat' sets - return (shapeargs ++ ses, - (form', - shapepat, - mergepat', - mergeinit)) - - handleForm mergeinit (E.ForIn x arr) = do - arr' <- internaliseExpToVars "for_in_arr" arr - arr_ts <- mapM lookupType arr' - let w = arraysSize 0 arr_ts - - i <- newVName "i" - - bindingLoopParams sparams' mergepat $ - \shapepat mergepat' -> - bindingLambdaParams [x] (map rowType arr_ts) $ \x_params -> do - let loopvars = zip x_params arr' - forLoop mergepat' shapepat mergeinit $ - I.ForLoop i Int32 w loopvars - - handleForm mergeinit (E.For i num_iterations) = do - num_iterations' <- internaliseExp1 "upper_bound" num_iterations - i' <- internaliseIdent i - num_iterations_t <- I.subExpType num_iterations' - it <- case num_iterations_t of - I.Prim (IntType it) -> return it - _ -> error "internaliseExp DoLoop: invalid type" - - bindingLoopParams sparams' mergepat $ - \shapepat mergepat' -> - forLoop mergepat' shapepat mergeinit $ - I.ForLoop i' it num_iterations' [] - - handleForm mergeinit (E.While cond) = - bindingLoopParams sparams' mergepat $ \shapepat mergepat' -> do - mergeinit_ts <- mapM subExpType mergeinit - -- We need to insert 'cond' twice - once for the initial - -- condition (do we enter the loop at all?), and once with the - -- result values of the loop (do we continue into the next - -- iteration?). This is safe, as the type rules for the - -- external language guarantees that 'cond' does not consume - -- anything. - shapeinit <- argShapes (map I.paramName shapepat) mergepat' mergeinit_ts - - (loop_initial_cond, init_loop_cond_bnds) <- collectStms $ do - forM_ (zip shapepat shapeinit) $ \(p, se) -> - letBindNames [paramName p] $ BasicOp $ SubExp se - forM_ (zip mergepat' mergeinit) $ \(p, se) -> - unless (se == I.Var (paramName p)) $ - letBindNames [paramName p] $ BasicOp $ - case se of I.Var v | not $ primType $ paramType p -> - Reshape (map DimCoercion $ arrayDims $ paramType p) v - _ -> SubExp se - internaliseExp1 "loop_cond" cond - - addStms init_loop_cond_bnds - - bodyFromStms $ do - ses <- internaliseExp "loopres" loopbody - sets <- mapM subExpType ses - loop_while <- newParam "loop_while" $ I.Prim I.Bool - shapeargs <- argShapes (map I.paramName shapepat) mergepat' sets - - -- Careful not to clobber anything. - loop_end_cond_body <- renameBody <=< insertStmsM $ do - forM_ (zip shapepat shapeargs) $ \(p, se) -> - unless (se == I.Var (paramName p)) $ - letBindNames [paramName p] $ BasicOp $ SubExp se - forM_ (zip mergepat' ses) $ \(p, se) -> - unless (se == I.Var (paramName p)) $ - letBindNames [paramName p] $ BasicOp $ - case se of I.Var v | not $ primType $ paramType p -> - Reshape (map DimCoercion $ arrayDims $ paramType p) v - _ -> SubExp se - resultBody <$> internaliseExp "loop_cond" cond - loop_end_cond <- bodyBind loop_end_cond_body - - return (shapeargs++loop_end_cond++ses, - (I.WhileLoop $ I.paramName loop_while, - shapepat, - loop_while : mergepat', - loop_initial_cond : mergeinit)) - -internaliseExp desc (E.LetWith name src idxs ve body t loc) = do - let pat = E.Id (E.identName name) (E.identType name) loc - src_t = E.fromStruct <$> E.identType src - e = E.Update (E.Var (E.qualName $ E.identName src) src_t loc) idxs ve loc - internaliseExp desc $ E.LetPat pat e body (t, Info []) loc - -internaliseExp desc (E.Update src slice ve loc) = do - ves <- internaliseExp "lw_val" ve - srcs <- internaliseExpToVars "src" src - dims <- case srcs of - [] -> return [] -- Will this happen? - v:_ -> I.arrayDims <$> lookupType v - (idxs', cs) <- internaliseSlice loc dims slice - - let comb sname ve' = do - sname_t <- lookupType sname - let full_slice = fullSlice sname_t idxs' - rowtype = sname_t `setArrayDims` sliceDims full_slice - ve'' <- ensureShape "shape of value does not match shape of source array" - loc rowtype "lw_val_correct_shape" ve' - letInPlace desc sname full_slice $ BasicOp $ SubExp ve'' - certifying cs $ map I.Var <$> zipWithM comb srcs ves - -internaliseExp desc (E.RecordUpdate src fields ve _ _) = do - src' <- internaliseExp desc src - ve' <- internaliseExp desc ve - replace (E.typeOf src `setAliases` ()) fields ve' src' - where replace (E.Scalar (E.Record m)) (f:fs) ve' src' - | Just t <- M.lookup f m = do - i <- fmap sum $ mapM (internalisedTypeSize . snd) $ - takeWhile ((/=f) . fst) $ sortFields m - k <- internalisedTypeSize t - let (bef, to_update, aft) = splitAt3 i k src' - src'' <- replace t fs ve' to_update - return $ bef ++ src'' ++ aft - replace _ _ ve' _ = return ve' - -internaliseExp desc (E.Attr attr e _) = - local f $ internaliseExp desc e - where attrs = oneAttr $ internaliseAttr attr - f env | "unsafe" `inAttrs` attrs, - not $ envSafe env = - env { envDoBoundsChecks = False } - | otherwise = - env { envAttrs = envAttrs env <> attrs } - -internaliseExp desc (E.Assert e1 e2 (Info check) loc) = do - e1' <- internaliseExp1 "assert_cond" e1 - c <- assert "assert_c" e1' (errorMsg [ErrorString $ "Assertion is false: " <> check]) loc - -- Make sure there are some bindings to certify. - certifying c $ mapM rebind =<< internaliseExp desc e2 - where rebind v = do - v' <- newVName "assert_res" - letBindNames [v'] $ I.BasicOp $ I.SubExp v - return $ I.Var v' - -internaliseExp _ (E.Constr c es (Info (E.Scalar (E.Sum fs))) _) = do - (ts, constr_map) <- internaliseSumType $ M.map (map E.toStruct) fs - es' <- concat <$> mapM (internaliseExp "payload") es - - let noExt _ = return $ intConst Int32 0 - ts' <- instantiateShapes noExt $ map fromDecl ts - - case M.lookup c constr_map of - Just (i, js) -> - (intConst Int8 (toInteger i):) <$> clauses 0 ts' (zip js es') - Nothing -> - error "internaliseExp Constr: missing constructor" - - where clauses j (t:ts) js_to_es - | Just e <- j `lookup` js_to_es = - (e:) <$> clauses (j+1) ts js_to_es - | otherwise = do - blank <- letSubExp "zero" =<< eBlank t - (blank:) <$> clauses (j+1) ts js_to_es - clauses _ [] _ = - return [] - -internaliseExp _ (E.Constr _ _ (Info t) loc) = - error $ "internaliseExp: constructor with type " ++ pretty t ++ " at " ++ locStr loc - -internaliseExp desc (E.Match e cs (Info ret, Info retext) _) = do - ses <- internaliseExp (desc ++ "_scrutinee") e - res <- - case NE.uncons cs of - (CasePat pCase eCase _, Nothing) -> do - (_, pertinent) <- generateCond pCase ses - internalisePat' pCase pertinent eCase (internaliseExp desc) - (c, Just cs') -> do - let CasePat pLast eLast _ = NE.last cs' - bFalse <- do - (_, pertinent) <- generateCond pLast ses - eLast' <- internalisePat' pLast pertinent eLast internaliseBody - foldM (\bf c' -> eBody $ return $ generateCaseIf ses c' bf) eLast' $ - reverse $ NE.init cs' - letTupExp' desc =<< generateCaseIf ses c bFalse - bindExtSizes (E.toStruct ret) retext res - return res - --- The "interesting" cases are over, now it's mostly boilerplate. - -internaliseExp _ (E.Literal v _) = - return [I.Constant $ internalisePrimValue v] - -internaliseExp _ (E.IntLit v (Info t) _) = - case t of - E.Scalar (E.Prim (E.Signed it)) -> - return [I.Constant $ I.IntValue $ intValue it v] - E.Scalar (E.Prim (E.Unsigned it)) -> - return [I.Constant $ I.IntValue $ intValue it v] - E.Scalar (E.Prim (E.FloatType ft)) -> - return [I.Constant $ I.FloatValue $ floatValue ft v] - _ -> error $ "internaliseExp: nonsensical type for integer literal: " ++ pretty t - -internaliseExp _ (E.FloatLit v (Info t) _) = - case t of - E.Scalar (E.Prim (E.FloatType ft)) -> - return [I.Constant $ I.FloatValue $ floatValue ft v] - _ -> error $ "internaliseExp: nonsensical type for float literal: " ++ pretty t - -internaliseExp desc (E.If ce te fe (Info ret, Info retext) _) = do - ses <- letTupExp' desc =<< - eIf (BasicOp . SubExp <$> internaliseExp1 "cond" ce) - (internaliseBody te) (internaliseBody fe) - bindExtSizes (E.toStruct ret) retext ses - return ses - --- Builtin operators are handled specially because they are --- overloaded. -internaliseExp desc (E.BinOp (op, _) _ (xe,_) (ye,_) _ _ loc) - | Just internalise <- isOverloadedFunction op [xe, ye] loc = - internalise desc - --- User-defined operators are just the same as a function call. -internaliseExp desc (E.BinOp (op, oploc) (Info t) - (xarg, Info (xt,xext)) (yarg, Info (yt,yext)) - _ (Info retext) loc) = - internaliseExp desc $ - E.Apply (E.Apply (E.Var op (Info t) oploc) xarg (Info (E.diet xt, xext)) - (Info $ foldFunType [E.fromStruct yt] t, Info []) loc) - yarg (Info (E.diet yt, yext)) (Info t, Info retext) loc - -internaliseExp desc (E.Project k e (Info rt) _) = do - n <- internalisedTypeSize $ rt `setAliases` () - i' <- fmap sum $ mapM internalisedTypeSize $ - case E.typeOf e `setAliases` () of - E.Scalar (Record fs) -> - map snd $ takeWhile ((/=k) . fst) $ sortFields fs - t -> [t] - take n . drop i' <$> internaliseExp desc e - -internaliseExp _ e@E.Lambda{} = - error $ "internaliseExp: Unexpected lambda at " ++ locStr (srclocOf e) - -internaliseExp _ e@E.OpSection{} = - error $ "internaliseExp: Unexpected operator section at " ++ locStr (srclocOf e) - -internaliseExp _ e@E.OpSectionLeft{} = - error $ "internaliseExp: Unexpected left operator section at " ++ locStr (srclocOf e) - -internaliseExp _ e@E.OpSectionRight{} = - error $ "internaliseExp: Unexpected right operator section at " ++ locStr (srclocOf e) - -internaliseExp _ e@E.ProjectSection{} = - error $ "internaliseExp: Unexpected projection section at " ++ locStr (srclocOf e) - -internaliseExp _ e@E.IndexSection{} = - error $ "internaliseExp: Unexpected index section at " ++ locStr (srclocOf e) - -internaliseArg :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp] -internaliseArg desc (arg, argdim) = do - arg' <- internaliseExp desc arg - case (arg', argdim) of - ([se], Just d) -> letBindNames [d] $ BasicOp $ SubExp se - _ -> return () - return arg' - -generateCond :: E.Pattern -> [I.SubExp] -> InternaliseM (I.SubExp, [I.SubExp]) -generateCond orig_p orig_ses = do - (cmps, pertinent, _) <- compares orig_p orig_ses - cmp <- letSubExp "matches" =<< eAll cmps - return (cmp, pertinent) - where - -- Literals are always primitive values. - compares (E.PatternLit e _ _) (se:ses) = do - e' <- internaliseExp1 "constant" e - t' <- elemType <$> subExpType se - cmp <- letSubExp "match_lit" $ I.BasicOp $ I.CmpOp (I.CmpEq t') e' se - return ([cmp], [se], ses) - - compares (E.PatternConstr c (Info (E.Scalar (E.Sum fs))) pats _) (se:ses) = do - (payload_ts, m) <- internaliseSumType $ M.map (map toStruct) fs - case M.lookup c m of - Just (i, payload_is) -> do - let i' = intConst Int8 $ toInteger i - let (payload_ses, ses') = splitAt (length payload_ts) ses - cmp <- letSubExp "match_constr" $ I.BasicOp $ I.CmpOp (I.CmpEq int8) i' se - (cmps, pertinent, _) <- comparesMany pats $ map (payload_ses!!) payload_is - return (cmp : cmps, pertinent, ses') - Nothing -> - error "generateCond: missing constructor" - - compares (E.PatternConstr _ (Info t) _ _) _ = - error $ "generateCond: PatternConstr has nonsensical type: " ++ pretty t - - compares (E.Id _ t loc) ses = - compares (E.Wildcard t loc) ses - - compares (E.Wildcard (Info t) _) ses = do - n <- internalisedTypeSize $ E.toStruct t - let (id_ses, rest_ses) = splitAt n ses - return ([], id_ses, rest_ses) - - compares (E.PatternParens pat _) ses = - compares pat ses - - compares (E.TuplePattern pats _) ses = - comparesMany pats ses - - compares (E.RecordPattern fs _) ses = - comparesMany (map snd $ E.sortFields $ M.fromList fs) ses - - compares (E.PatternAscription pat _ _) ses = - compares pat ses - - compares pat [] = - error $ "generateCond: No values left for pattern " ++ pretty pat - - comparesMany [] ses = return ([], [], ses) - comparesMany (pat:pats) ses = do - (cmps1, pertinent1, ses') <- compares pat ses - (cmps2, pertinent2, ses'') <- comparesMany pats ses' - return (cmps1 <> cmps2, - pertinent1 <> pertinent2, - ses'') - -generateCaseIf :: [I.SubExp] -> Case -> I.Body -> InternaliseM I.Exp -generateCaseIf ses (CasePat p eCase _) bFail = do - (cond, pertinent) <- generateCond p ses - eCase' <- internalisePat' p pertinent eCase internaliseBody - eIf (eSubExp cond) (return eCase') (return bFail) - -internalisePat :: String -> E.Pattern -> E.Exp - -> E.Exp -> (E.Exp -> InternaliseM a) - -> InternaliseM a -internalisePat desc p e body m = do - ses <- internaliseExp desc' e - internalisePat' p ses body m - where desc' = case S.toList $ E.patternIdents p of - [v] -> baseString $ E.identName v - _ -> desc - -internalisePat' :: E.Pattern -> [I.SubExp] - -> E.Exp -> (E.Exp -> InternaliseM a) - -> InternaliseM a -internalisePat' p ses body m = do - ses_ts <- mapM subExpType ses - stmPattern p ses_ts $ \pat_names -> do - forM_ (zip pat_names ses) $ \(v,se) -> - letBindNames [v] $ I.BasicOp $ I.SubExp se - m body - -internaliseSlice :: SrcLoc - -> [SubExp] - -> [E.DimIndex] - -> InternaliseM ([I.DimIndex SubExp], Certificates) -internaliseSlice loc dims idxs = do - (idxs', oks, parts) <- unzip3 <$> zipWithM internaliseDimIndex dims idxs - ok <- letSubExp "index_ok" =<< eAll oks - let msg = errorMsg $ ["Index ["] ++ intercalate [", "] parts ++ - ["] out of bounds for array of shape ["] ++ - intersperse "][" (map ErrorInt32 $ take (length idxs) dims) ++ ["]."] - c <- assert "index_certs" ok msg loc - return (idxs', c) - -internaliseDimIndex :: SubExp -> E.DimIndex - -> InternaliseM (I.DimIndex SubExp, SubExp, [ErrorMsgPart SubExp]) -internaliseDimIndex w (E.DimFix i) = do - (i', _) <- internaliseDimExp "i" i - let lowerBound = I.BasicOp $ - I.CmpOp (I.CmpSle I.Int32) (I.constant (0 :: I.Int32)) i' - upperBound = I.BasicOp $ - I.CmpOp (I.CmpSlt I.Int32) i' w - ok <- letSubExp "bounds_check" =<< eBinOp I.LogAnd (pure lowerBound) (pure upperBound) - return (I.DimFix i', ok, [ErrorInt32 i']) - --- Special-case an important common case that otherwise leads to horrible code. -internaliseDimIndex w (E.DimSlice Nothing Nothing - (Just (E.Negate (E.IntLit 1 _ _) _))) = do - w_minus_1 <- letSubExp "w_minus_1" $ - BasicOp $ I.BinOp (Sub Int32 I.OverflowWrap) w one - return (I.DimSlice w_minus_1 w $ intConst Int32 (-1), - constant True, - mempty) - where one = constant (1::Int32) - -internaliseDimIndex w (E.DimSlice i j s) = do - s' <- maybe (return one) (fmap fst . internaliseDimExp "s") s - s_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum Int32) s' - backwards <- letSubExp "backwards" $ I.BasicOp $ I.CmpOp (I.CmpEq int32) s_sign negone - w_minus_1 <- letSubExp "w_minus_1" $ BasicOp $ I.BinOp (Sub Int32 I.OverflowWrap) w one - let i_def = letSubExp "i_def" $ I.If backwards - (resultBody [w_minus_1]) - (resultBody [zero]) $ ifCommon [I.Prim int32] - j_def = letSubExp "j_def" $ I.If backwards - (resultBody [negone]) - (resultBody [w]) $ ifCommon [I.Prim int32] - i' <- maybe i_def (fmap fst . internaliseDimExp "i") i - j' <- maybe j_def (fmap fst . internaliseDimExp "j") j - j_m_i <- letSubExp "j_m_i" $ BasicOp $ I.BinOp (Sub Int32 I.OverflowWrap) j' i' - -- Something like a division-rounding-up, but accomodating negative - -- operands. - let divRounding x y = - eBinOp (SQuot Int32 Unsafe) - (eBinOp (Add Int32 I.OverflowWrap) x - (eBinOp (Sub Int32 I.OverflowWrap) y (eSignum $ toExp s'))) y - n <- letSubExp "n" =<< divRounding (toExp j_m_i) (toExp s') - - -- Bounds checks depend on whether we are slicing forwards or - -- backwards. If forwards, we must check '0 <= i && i <= j'. If - -- backwards, '-1 <= j && j <= i'. In both cases, we check '0 <= - -- i+n*s && i+(n-1)*s < w'. We only check if the slice is nonempty. - empty_slice <- letSubExp "empty_slice" $ I.BasicOp $ I.CmpOp (CmpEq int32) n zero - - m <- letSubExp "m" $ I.BasicOp $ I.BinOp (Sub Int32 I.OverflowWrap) n one - m_t_s <- letSubExp "m_t_s" $ I.BasicOp $ I.BinOp (Mul Int32 I.OverflowWrap) m s' - i_p_m_t_s <- letSubExp "i_p_m_t_s" $ I.BasicOp $ I.BinOp (Add Int32 I.OverflowWrap) i' m_t_s - zero_leq_i_p_m_t_s <- letSubExp "zero_leq_i_p_m_t_s" $ - I.BasicOp $ I.CmpOp (I.CmpSle Int32) zero i_p_m_t_s - i_p_m_t_s_leq_w <- letSubExp "i_p_m_t_s_leq_w" $ - I.BasicOp $ I.CmpOp (I.CmpSle Int32) i_p_m_t_s w - i_p_m_t_s_lth_w <- letSubExp "i_p_m_t_s_leq_w" $ - I.BasicOp $ I.CmpOp (I.CmpSlt Int32) i_p_m_t_s w - - zero_lte_i <- letSubExp "zero_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) zero i' - i_lte_j <- letSubExp "i_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) i' j' - forwards_ok <- letSubExp "forwards_ok" =<< - eAll [zero_lte_i, zero_lte_i, i_lte_j, zero_leq_i_p_m_t_s, i_p_m_t_s_lth_w] - - negone_lte_j <- letSubExp "negone_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) negone j' - j_lte_i <- letSubExp "j_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) j' i' - backwards_ok <- letSubExp "backwards_ok" =<< - eAll - [negone_lte_j, negone_lte_j, j_lte_i, zero_leq_i_p_m_t_s, i_p_m_t_s_leq_w] - - slice_ok <- letSubExp "slice_ok" $ I.If backwards - (resultBody [backwards_ok]) - (resultBody [forwards_ok]) $ - ifCommon [I.Prim I.Bool] - ok_or_empty <- letSubExp "ok_or_empty" $ - I.BasicOp $ I.BinOp I.LogOr empty_slice slice_ok - - let parts = case (i, j, s) of - (_, _, Just{}) -> - [maybe "" (const $ ErrorInt32 i') i, ":", - maybe "" (const $ ErrorInt32 j') j, ":", - ErrorInt32 s'] - (_, Just{}, _) -> - [maybe "" (const $ ErrorInt32 i') i, ":", - ErrorInt32 j'] ++ - maybe mempty (const [":", ErrorInt32 s']) s - (_, Nothing, Nothing) -> - [ErrorInt32 i', ":"] - return (I.DimSlice i' n s', ok_or_empty, parts) - where zero = constant (0::Int32) - negone = constant (-1::Int32) - one = constant (1::Int32) - -internaliseScanOrReduce :: String -> String - -> (SubExp -> I.Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)) - -> (E.Exp, E.Exp, E.Exp, SrcLoc) - -> InternaliseM [SubExp] -internaliseScanOrReduce desc what f (lam, ne, arr, loc) = do - arrs <- internaliseExpToVars (what++"_arr") arr - nes <- internaliseExp (what++"_ne") ne - nes' <- forM (zip nes arrs) $ \(ne', arr') -> do - rowtype <- I.stripArray 1 <$> lookupType arr' - ensureShape - "Row shape of input array does not match shape of neutral element" - loc rowtype (what++"_ne_right_shape") ne' - nests <- mapM I.subExpType nes' - arrts <- mapM lookupType arrs - lam' <- internaliseFoldLambda internaliseLambda lam nests arrts - w <- arraysSize 0 <$> mapM lookupType arrs - letTupExp' desc . I.Op =<< f w lam' nes' arrs - -internaliseHist :: String - -> E.Exp -> E.Exp -> E.Exp -> E.Exp -> E.Exp -> E.Exp -> SrcLoc - -> InternaliseM [SubExp] -internaliseHist desc rf hist op ne buckets img loc = do - rf' <- internaliseExp1 "hist_rf" rf - ne' <- internaliseExp "hist_ne" ne - hist' <- internaliseExpToVars "hist_hist" hist - buckets' <- letExp "hist_buckets" . BasicOp . SubExp =<< - internaliseExp1 "hist_buckets" buckets - img' <- internaliseExpToVars "hist_img" img - - -- reshape neutral element to have same size as the destination array - ne_shp <- forM (zip ne' hist') $ \(n, h) -> do - rowtype <- I.stripArray 1 <$> lookupType h - ensureShape - "Row shape of destination array does not match shape of neutral element" - loc rowtype "hist_ne_right_shape" n - ne_ts <- mapM I.subExpType ne_shp - his_ts <- mapM lookupType hist' - op' <- internaliseFoldLambda internaliseLambda op ne_ts his_ts - - -- reshape return type of bucket function to have same size as neutral element - -- (modulo the index) - bucket_param <- newParam "bucket_p" $ I.Prim int32 - img_params <- mapM (newParam "img_p" . rowType) =<< mapM lookupType img' - let params = bucket_param : img_params - rettype = I.Prim int32 : ne_ts - body = mkBody mempty $ map (I.Var . paramName) params - body' <- localScope (scopeOfLParams params) $ - ensureResultShape - "Row shape of value array does not match row shape of hist target" - (srclocOf img) rettype body - - -- get sizes of histogram and image arrays - w_hist <- arraysSize 0 <$> mapM lookupType hist' - w_img <- arraysSize 0 <$> mapM lookupType img' - - -- Generate an assertion and reshapes to ensure that buckets' and - -- img' are the same size. - b_shape <- I.arrayShape <$> lookupType buckets' - let b_w = shapeSize 0 b_shape - cmp <- letSubExp "bucket_cmp" $ I.BasicOp $ I.CmpOp (I.CmpEq I.int32) b_w w_img - c <- assert "bucket_cert" cmp - "length of index and value array does not match" loc - buckets'' <- certifying c $ letExp (baseString buckets') $ - I.BasicOp $ I.Reshape (reshapeOuter [DimCoercion w_img] 1 b_shape) buckets' - - letTupExp' desc $ I.Op $ - I.Hist w_img [HistOp w_hist rf' hist' ne_shp op'] (I.Lambda params body' rettype) $ buckets'' : img' - -internaliseStreamMap :: String -> StreamOrd -> E.Exp -> E.Exp - -> InternaliseM [SubExp] -internaliseStreamMap desc o lam arr = do - 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 - -internaliseStreamRed :: String -> StreamOrd -> Commutativity - -> E.Exp -> E.Exp -> E.Exp - -> InternaliseM [SubExp] -internaliseStreamRed desc o comm lam0 lam arr = do - arrs <- internaliseExpToVars "stream_input" arr - rowts <- mapM (fmap I.rowType . lookupType) arrs - (lam_params, lam_body) <- - internaliseStreamLambda internaliseLambda lam rowts - let (chunk_param, _, lam_val_params) = - partitionChunkedFoldParameters 0 lam_params - - -- Synthesize neutral elements by applying the fold function - -- to an empty chunk. - letBindNames [I.paramName chunk_param] $ - I.BasicOp $ I.SubExp $ constant (0::Int32) - forM_ lam_val_params $ \p -> - letBindNames [I.paramName p] $ - I.BasicOp $ I.Scratch (I.elemType $ I.paramType p) $ - I.arrayDims $ I.paramType p - nes <- bodyBind =<< renameBody lam_body - - nes_ts <- mapM I.subExpType nes - outsz <- arraysSize 0 <$> mapM lookupType arrs - let acc_arr_tps = [ I.arrayOf t (I.Shape [outsz]) NoUniqueness | t <- nes_ts ] - lam0' <- internaliseFoldLambda internaliseLambda lam0 nes_ts acc_arr_tps - - let lam0_acc_params = take (length nes) $ I.lambdaParams lam0' - lam_acc_params <- forM lam0_acc_params $ \p -> do - name <- newVName $ baseString $ I.paramName p - return p { I.paramName = name } - - -- Make sure the chunk size parameter comes first. - let lam_params' = chunk_param : lam_acc_params ++ lam_val_params - - body_with_lam0 <- - ensureResultShape "shape of result does not match shape of initial value" - (srclocOf lam0) nes_ts <=< insertStmsM $ localScope (scopeOfLParams lam_params') $ do - lam_res <- bodyBind lam_body - lam_res' <- ensureArgShapes - "shape of chunk function result does not match shape of initial value" - (srclocOf lam) [] (map I.typeOf $ I.lambdaParams lam0') lam_res - new_lam_res <- eLambda lam0' $ map eSubExp $ - map (I.Var . paramName) lam_acc_params ++ lam_res' - return $ resultBody new_lam_res - - let form = I.Parallel o comm lam0' nes - lam' = I.Lambda { lambdaParams = lam_params' - , lambdaBody = body_with_lam0 - , lambdaReturnType = nes_ts } - w <- arraysSize 0 <$> mapM lookupType arrs - letTupExp' desc $ I.Op $ I.Stream w form lam' arrs - -internaliseExp1 :: String -> E.Exp -> InternaliseM I.SubExp -internaliseExp1 desc e = do - vs <- internaliseExp desc e - case vs of [se] -> return se - _ -> error "Internalise.internaliseExp1: was passed not just a single subexpression" - --- | Promote to dimension type as appropriate for the original type. --- Also return original type. -internaliseDimExp :: String -> E.Exp -> InternaliseM (I.SubExp, IntType) -internaliseDimExp s e = do - e' <- internaliseExp1 s e - case E.typeOf e of - E.Scalar (E.Prim (Signed it)) -> (,it) <$> asIntS Int32 e' - E.Scalar (E.Prim (Unsigned it)) -> (,it) <$> asIntZ Int32 e' - _ -> error "internaliseDimExp: bad type" - -internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName] -internaliseExpToVars desc e = - mapM asIdent =<< internaliseExp desc e - where asIdent (I.Var v) = return v - asIdent se = letExp desc $ I.BasicOp $ I.SubExp se - -internaliseOperation :: String - -> E.Exp - -> (I.VName -> InternaliseM I.BasicOp) - -> InternaliseM [I.SubExp] -internaliseOperation s e op = do - vs <- internaliseExpToVars s e - letSubExps s =<< mapM (fmap I.BasicOp . op) vs - -certifyingNonzero :: SrcLoc -> IntType -> SubExp - -> InternaliseM a - -> InternaliseM a -certifyingNonzero loc t x m = do - zero <- letSubExp "zero" $ I.BasicOp $ - CmpOp (CmpEq (IntType t)) x (intConst t 0) - nonzero <- letSubExp "nonzero" $ I.BasicOp $ UnOp Not zero - c <- assert "nonzero_cert" nonzero "division by zero" loc - certifying c m - -certifyingNonnegative :: SrcLoc -> IntType -> SubExp - -> InternaliseM a - -> InternaliseM a -certifyingNonnegative loc t x m = do - nonnegative <- letSubExp "nonnegative" $ I.BasicOp $ - CmpOp (CmpSle t) (intConst t 0) x - c <- assert "nonzero_cert" nonnegative "negative exponent" loc - certifying c m - -internaliseBinOp :: SrcLoc -> String - -> E.BinOp - -> I.SubExp -> I.SubExp - -> E.PrimType - -> E.PrimType - -> InternaliseM [I.SubExp] -internaliseBinOp _ desc E.Plus x y (E.Signed t) _ = - simpleBinOp desc (I.Add t I.OverflowWrap) x y -internaliseBinOp _ desc E.Plus x y (E.Unsigned t) _ = - simpleBinOp desc (I.Add t I.OverflowWrap) x y -internaliseBinOp _ desc E.Plus x y (E.FloatType t) _ = - simpleBinOp desc (I.FAdd t) x y -internaliseBinOp _ desc E.Minus x y (E.Signed t) _ = - simpleBinOp desc (I.Sub t I.OverflowWrap) x y -internaliseBinOp _ desc E.Minus x y (E.Unsigned t) _ = - simpleBinOp desc (I.Sub t I.OverflowWrap) x y -internaliseBinOp _ desc E.Minus x y (E.FloatType t) _ = - simpleBinOp desc (I.FSub t) x y -internaliseBinOp _ desc E.Times x y (E.Signed t) _ = - simpleBinOp desc (I.Mul t I.OverflowWrap) x y -internaliseBinOp _ desc E.Times x y (E.Unsigned t) _ = - simpleBinOp desc (I.Mul t I.OverflowWrap) x y -internaliseBinOp _ desc E.Times x y (E.FloatType t) _ = - simpleBinOp desc (I.FMul t) x y -internaliseBinOp loc desc E.Divide x y (E.Signed t) _ = - certifyingNonzero loc t y $ - simpleBinOp desc (I.SDiv t I.Unsafe) x y -internaliseBinOp loc desc E.Divide x y (E.Unsigned t) _ = - certifyingNonzero loc t y $ - simpleBinOp desc (I.UDiv t I.Unsafe) x y -internaliseBinOp _ desc E.Divide x y (E.FloatType t) _ = - simpleBinOp desc (I.FDiv t) x y -internaliseBinOp _ desc E.Pow x y (E.FloatType t) _ = - simpleBinOp desc (I.FPow t) x y -internaliseBinOp loc desc E.Pow x y (E.Signed t) _ = - certifyingNonnegative loc t y $ - simpleBinOp desc (I.Pow t) x y -internaliseBinOp _ desc E.Pow x y (E.Unsigned t) _ = - simpleBinOp desc (I.Pow t) x y -internaliseBinOp loc desc E.Mod x y (E.Signed t) _ = - certifyingNonzero loc t y $ - simpleBinOp desc (I.SMod t I.Unsafe) x y -internaliseBinOp loc desc E.Mod x y (E.Unsigned t) _ = - certifyingNonzero loc t y $ - simpleBinOp desc (I.UMod t I.Unsafe) x y -internaliseBinOp _ desc E.Mod x y (E.FloatType t) _ = - simpleBinOp desc (I.FMod t) x y -internaliseBinOp loc desc E.Quot x y (E.Signed t) _ = - certifyingNonzero loc t y $ - simpleBinOp desc (I.SQuot t I.Unsafe) x y -internaliseBinOp loc desc E.Quot x y (E.Unsigned t) _ = - certifyingNonzero loc t y $ - simpleBinOp desc (I.UDiv t I.Unsafe) x y -internaliseBinOp loc desc E.Rem x y (E.Signed t) _ = - certifyingNonzero loc t y $ - simpleBinOp desc (I.SRem t I.Unsafe) x y -internaliseBinOp loc desc E.Rem x y (E.Unsigned t) _ = - certifyingNonzero loc t y $ - simpleBinOp desc (I.UMod t I.Unsafe) x y -internaliseBinOp _ desc E.ShiftR x y (E.Signed t) _ = - simpleBinOp desc (I.AShr t) x y -internaliseBinOp _ desc E.ShiftR x y (E.Unsigned t) _ = - simpleBinOp desc (I.LShr t) x y -internaliseBinOp _ desc E.ShiftL x y (E.Signed t) _ = - simpleBinOp desc (I.Shl t) x y -internaliseBinOp _ desc E.ShiftL x y (E.Unsigned t) _ = - simpleBinOp desc (I.Shl t) x y -internaliseBinOp _ desc E.Band x y (E.Signed t) _ = - simpleBinOp desc (I.And t) x y -internaliseBinOp _ desc E.Band x y (E.Unsigned t) _ = - simpleBinOp desc (I.And t) x y -internaliseBinOp _ desc E.Xor x y (E.Signed t) _ = - simpleBinOp desc (I.Xor t) x y -internaliseBinOp _ desc E.Xor x y (E.Unsigned t) _ = - simpleBinOp desc (I.Xor t) x y -internaliseBinOp _ desc E.Bor x y (E.Signed t) _ = - simpleBinOp desc (I.Or t) x y -internaliseBinOp _ desc E.Bor x y (E.Unsigned t) _ = - simpleBinOp desc (I.Or t) x y - -internaliseBinOp _ desc E.Equal x y t _ = - simpleCmpOp desc (I.CmpEq $ internalisePrimType t) x y -internaliseBinOp _ desc E.NotEqual x y t _ = do - eq <- letSubExp (desc++"true") $ I.BasicOp $ I.CmpOp (I.CmpEq $ internalisePrimType t) x y - fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp I.Not eq -internaliseBinOp _ desc E.Less x y (E.Signed t) _ = - simpleCmpOp desc (I.CmpSlt t) x y -internaliseBinOp _ desc E.Less x y (E.Unsigned t) _ = - simpleCmpOp desc (I.CmpUlt t) x y -internaliseBinOp _ desc E.Leq x y (E.Signed t) _ = - simpleCmpOp desc (I.CmpSle t) x y -internaliseBinOp _ desc E.Leq x y (E.Unsigned t) _ = - simpleCmpOp desc (I.CmpUle t) x y -internaliseBinOp _ desc E.Greater x y (E.Signed t) _ = - simpleCmpOp desc (I.CmpSlt t) y x -- Note the swapped x and y -internaliseBinOp _ desc E.Greater x y (E.Unsigned t) _ = - simpleCmpOp desc (I.CmpUlt t) y x -- Note the swapped x and y -internaliseBinOp _ desc E.Geq x y (E.Signed t) _ = - simpleCmpOp desc (I.CmpSle t) y x -- Note the swapped x and y -internaliseBinOp _ desc E.Geq x y (E.Unsigned t) _ = - simpleCmpOp desc (I.CmpUle t) y x -- Note the swapped x and y -internaliseBinOp _ desc E.Less x y (E.FloatType t) _ = - simpleCmpOp desc (I.FCmpLt t) x y -internaliseBinOp _ desc E.Leq x y (E.FloatType t) _ = - simpleCmpOp desc (I.FCmpLe t) x y -internaliseBinOp _ desc E.Greater x y (E.FloatType t) _ = - simpleCmpOp desc (I.FCmpLt t) y x -- Note the swapped x and y -internaliseBinOp _ desc E.Geq x y (E.FloatType t) _ = - simpleCmpOp desc (I.FCmpLe t) y x -- Note the swapped x and y - --- Relational operators for booleans. -internaliseBinOp _ desc E.Less x y E.Bool _ = - simpleCmpOp desc I.CmpLlt x y -internaliseBinOp _ desc E.Leq x y E.Bool _ = - simpleCmpOp desc I.CmpLle x y -internaliseBinOp _ desc E.Greater x y E.Bool _ = - simpleCmpOp desc I.CmpLlt y x -- Note the swapped x and y -internaliseBinOp _ desc E.Geq x y E.Bool _ = - simpleCmpOp desc I.CmpLle y x -- Note the swapped x and y - -internaliseBinOp _ _ op _ _ t1 t2 = - error $ "Invalid binary operator " ++ pretty op ++ - " with operand types " ++ pretty t1 ++ ", " ++ pretty t2 - -simpleBinOp :: String - -> I.BinOp - -> I.SubExp -> I.SubExp - -> InternaliseM [I.SubExp] -simpleBinOp desc bop x y = - letTupExp' desc $ I.BasicOp $ I.BinOp bop x y - -simpleCmpOp :: String - -> I.CmpOp - -> I.SubExp -> I.SubExp - -> InternaliseM [I.SubExp] -simpleCmpOp desc op x y = - letTupExp' desc $ I.BasicOp $ I.CmpOp op x y - -findFuncall :: E.Exp -> InternaliseM (E.QualName VName, - [(E.Exp, Maybe VName)], - E.StructType, - [VName]) -findFuncall (E.Var fname (Info t) _) = - return (fname, [], E.toStruct t, []) -findFuncall (E.Apply f arg (Info (_, argext)) (Info ret, Info retext) _) = do - (fname, args, _, _) <- findFuncall f - return (fname, args ++ [(arg, argext)], E.toStruct ret, retext) -findFuncall e = - error $ "Invalid function expression in application: " ++ pretty e - -internaliseLambda :: InternaliseLambda - -internaliseLambda (E.Parens e _) rowtypes = - internaliseLambda e rowtypes - -internaliseLambda (E.Lambda params body _ (Info (_, rettype)) _) rowtypes = - bindingLambdaParams params rowtypes $ \params' -> do - body' <- internaliseBody body - rettype' <- internaliseLambdaReturnType rettype - return (params', body', rettype') - -internaliseLambda e _ = error $ "internaliseLambda: unexpected expression:\n" ++ pretty e - --- | Some operators and functions are overloaded or otherwise special --- - we detect and treat them here. -isOverloadedFunction :: E.QualName VName -> [E.Exp] -> SrcLoc - -> Maybe (String -> InternaliseM [SubExp]) -isOverloadedFunction qname args loc = do - guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag - let handlers = [handleSign, - handleIntrinsicOps, - handleOps, - handleSOACs, - handleRest] - msum [h args $ baseString $ qualLeaf qname | h <- handlers] - where - handleSign [x] "sign_i8" = Just $ toSigned I.Int8 x - handleSign [x] "sign_i16" = Just $ toSigned I.Int16 x - handleSign [x] "sign_i32" = Just $ toSigned I.Int32 x - handleSign [x] "sign_i64" = Just $ toSigned I.Int64 x - - handleSign [x] "unsign_i8" = Just $ toUnsigned I.Int8 x - handleSign [x] "unsign_i16" = Just $ toUnsigned I.Int16 x - handleSign [x] "unsign_i32" = Just $ toUnsigned I.Int32 x - handleSign [x] "unsign_i64" = Just $ toUnsigned I.Int64 x - - handleSign _ _ = Nothing - - handleIntrinsicOps [x] s - | Just unop <- find ((==s) . pretty) allUnOps = Just $ \desc -> do - x' <- internaliseExp1 "x" x - fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop x' - handleIntrinsicOps [TupLit [x,y] _] s - | Just bop <- find ((==s) . pretty) allBinOps = Just $ \desc -> do - x' <- internaliseExp1 "x" x - y' <- internaliseExp1 "y" y - fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y' - | Just cmp <- find ((==s) . pretty) allCmpOps = Just $ \desc -> do - x' <- internaliseExp1 "x" x - y' <- internaliseExp1 "y" y - fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' y' - handleIntrinsicOps [x] s - | Just conv <- find ((==s) . pretty) allConvOps = Just $ \desc -> do - x' <- internaliseExp1 "x" x - fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x' - handleIntrinsicOps _ _ = Nothing - - -- Short-circuiting operators are magical. - handleOps [x,y] "&&" = Just $ \desc -> - internaliseExp desc $ - E.If x y (E.Literal (E.BoolValue False) mempty) (Info $ E.Scalar $ E.Prim E.Bool, Info []) mempty - handleOps [x,y] "||" = Just $ \desc -> - internaliseExp desc $ - E.If x (E.Literal (E.BoolValue True) mempty) y (Info $ E.Scalar $ E.Prim E.Bool, Info []) mempty - - -- Handle equality and inequality specially, to treat the case of - -- arrays. - handleOps [xe,ye] op - | Just cmp_f <- isEqlOp op = Just $ \desc -> do - xe' <- internaliseExp "x" xe - ye' <- internaliseExp "y" ye - rs <- zipWithM (doComparison desc) xe' ye' - cmp_f desc =<< letSubExp "eq" =<< eAll rs - where isEqlOp "!=" = Just $ \desc eq -> - letTupExp' desc $ I.BasicOp $ I.UnOp I.Not eq - isEqlOp "==" = Just $ \_ eq -> - return [eq] - isEqlOp _ = Nothing - - doComparison desc x y = do - x_t <- I.subExpType x - y_t <- I.subExpType y - case x_t of - I.Prim t -> letSubExp desc $ I.BasicOp $ I.CmpOp (I.CmpEq t) x y - _ -> do - let x_dims = I.arrayDims x_t - y_dims = I.arrayDims y_t - dims_match <- forM (zip x_dims y_dims) $ \(x_dim, y_dim) -> - letSubExp "dim_eq" $ I.BasicOp $ I.CmpOp (I.CmpEq int32) x_dim y_dim - shapes_match <- letSubExp "shapes_match" =<< eAll dims_match - compare_elems_body <- runBodyBinder $ do - -- Flatten both x and y. - x_num_elems <- letSubExp "x_num_elems" =<< - foldBinOp (I.Mul Int32 I.OverflowUndef) (constant (1::Int32)) x_dims - x' <- letExp "x" $ I.BasicOp $ I.SubExp x - y' <- letExp "x" $ I.BasicOp $ I.SubExp y - x_flat <- letExp "x_flat" $ I.BasicOp $ I.Reshape [I.DimNew x_num_elems] x' - y_flat <- letExp "y_flat" $ I.BasicOp $ I.Reshape [I.DimNew x_num_elems] y' - - -- Compare the elements. - cmp_lam <- cmpOpLambda $ I.CmpEq (elemType x_t) - cmps <- letExp "cmps" $ I.Op $ - I.Screma x_num_elems (I.mapSOAC cmp_lam) [x_flat, y_flat] - - -- Check that all were equal. - and_lam <- binOpLambda I.LogAnd I.Bool - reduce <- I.reduceSOAC [Reduce Commutative and_lam [constant True]] - all_equal <- letSubExp "all_equal" $ I.Op $ I.Screma x_num_elems reduce [cmps] - return $ resultBody [all_equal] - - letSubExp "arrays_equal" $ - I.If shapes_match compare_elems_body (resultBody [constant False]) $ - ifCommon [I.Prim I.Bool] - - handleOps [x,y] name - | Just bop <- find ((name==) . pretty) [minBound..maxBound::E.BinOp] = - Just $ \desc -> do - x' <- internaliseExp1 "x" x - y' <- internaliseExp1 "y" y - case (E.typeOf x, E.typeOf y) of - (E.Scalar (E.Prim t1), E.Scalar (E.Prim t2)) -> - internaliseBinOp loc desc bop x' y' t1 t2 - _ -> error "Futhark.Internalise.internaliseExp: non-primitive type in BinOp." - - handleOps _ _ = Nothing - - handleSOACs [TupLit [lam, arr] _] "map" = Just $ \desc -> do - arr' <- internaliseExpToVars "map_arr" arr - lam' <- internaliseMapLambda internaliseLambda lam $ map I.Var arr' - w <- arraysSize 0 <$> mapM lookupType arr' - letTupExp' desc $ I.Op $ - I.Screma w (I.mapSOAC lam') arr' - - handleSOACs [TupLit [k, lam, arr] _] "partition" = do - k' <- fromIntegral <$> isInt32 k - Just $ \_desc -> do - arrs <- internaliseExpToVars "partition_input" arr - lam' <- internalisePartitionLambda internaliseLambda k' lam $ map I.Var arrs - uncurry (++) <$> partitionWithSOACS k' lam' arrs - where isInt32 (Literal (SignedValue (Int32Value k')) _) = Just k' - isInt32 (IntLit k' (Info (E.Scalar (E.Prim (Signed Int32)))) _) = Just $ fromInteger k' - isInt32 _ = Nothing - - handleSOACs [TupLit [lam, ne, arr] _] "reduce" = Just $ \desc -> - internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr, loc) - where reduce w red_lam nes arrs = - I.Screma w <$> - I.reduceSOAC [Reduce Noncommutative red_lam nes] <*> pure arrs - - handleSOACs [TupLit [lam, ne, arr] _] "reduce_comm" = Just $ \desc -> - internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr, loc) - where reduce w red_lam nes arrs = - I.Screma w <$> - I.reduceSOAC [Reduce Commutative red_lam nes] <*> pure arrs - - handleSOACs [TupLit [lam, ne, arr] _] "scan" = Just $ \desc -> - internaliseScanOrReduce desc "scan" reduce (lam, ne, arr, loc) - where reduce w scan_lam nes arrs = - I.Screma w <$> I.scanSOAC [Scan scan_lam nes] <*> pure arrs - - handleSOACs [TupLit [op, f, arr] _] "reduce_stream" = Just $ \desc -> - internaliseStreamRed desc InOrder Noncommutative op f arr - - handleSOACs [TupLit [op, f, arr] _] "reduce_stream_per" = Just $ \desc -> - internaliseStreamRed desc Disorder Commutative op f arr - - handleSOACs [TupLit [f, arr] _] "map_stream" = Just $ \desc -> - internaliseStreamMap desc InOrder f arr - - handleSOACs [TupLit [f, arr] _] "map_stream_per" = Just $ \desc -> - internaliseStreamMap desc Disorder f arr - - handleSOACs [TupLit [rf, dest, op, ne, buckets, img] _] "hist" = Just $ \desc -> - internaliseHist desc rf dest op ne buckets img loc - - handleSOACs _ _ = Nothing - - 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 [n, m, arr] _] "unflatten" = Just $ \desc -> do - arrs <- internaliseExpToVars "unflatten_arr" arr - n' <- internaliseExp1 "n" n - m' <- internaliseExp1 "m" m - -- The unflattened dimension needs to have the same number of elements - -- as the original dimension. - old_dim <- I.arraysSize 0 <$> mapM lookupType arrs - dim_ok <- letSubExp "dim_ok" =<< - eCmpOp (I.CmpEq I.int32) - (eBinOp (I.Mul Int32 I.OverflowUndef) (eSubExp n') (eSubExp m')) - (eSubExp old_dim) - dim_ok_cert <- assert "dim_ok_cert" dim_ok - "new shape has different number of elements than old shape" loc - certifying dim_ok_cert $ forM arrs $ \arr' -> do - arr_t <- lookupType arr' - letSubExp desc $ I.BasicOp $ - I.Reshape (reshapeOuter [DimNew n', DimNew m'] 1 $ I.arrayShape arr_t) arr' - - handleRest [arr] "flatten" = Just $ \desc -> do - arrs <- internaliseExpToVars "flatten_arr" arr - forM arrs $ \arr' -> do - arr_t <- lookupType arr' - let n = arraySize 0 arr_t - m = arraySize 1 arr_t - k <- letSubExp "flat_dim" $ I.BasicOp $ I.BinOp (Mul Int32 I.OverflowUndef) n m - letSubExp desc $ I.BasicOp $ - I.Reshape (reshapeOuter [DimNew k] 2 $ I.arrayShape arr_t) arr' - - handleRest [TupLit [x, y] _] "concat" = Just $ \desc -> do - xs <- internaliseExpToVars "concat_x" x - ys <- internaliseExpToVars "concat_y" y - outer_size <- arraysSize 0 <$> mapM lookupType xs - let sumdims xsize ysize = letSubExp "conc_tmp" $ I.BasicOp $ - I.BinOp (I.Add I.Int32 I.OverflowUndef) xsize ysize - ressize <- foldM sumdims outer_size =<< - mapM (fmap (arraysSize 0) . mapM lookupType) [ys] - - let conc xarr yarr = - I.BasicOp $ I.Concat 0 xarr [yarr] ressize - letSubExps desc $ zipWith conc xs ys - - handleRest [TupLit [offset, e] _] "rotate" = Just $ \desc -> do - offset' <- internaliseExp1 "rotation_offset" offset - internaliseOperation desc e $ \v -> do - r <- I.arrayRank <$> lookupType v - let zero = intConst Int32 0 - offsets = offset' : replicate (r-1) zero - return $ I.Rotate offsets v - - handleRest [e] "transpose" = Just $ \desc -> - internaliseOperation desc e $ \v -> do - r <- I.arrayRank <$> lookupType v - return $ I.Rearrange ([1,0] ++ [2..r-1]) v - - handleRest [TupLit [x, y] _] "zip" = Just $ \desc -> - (++) <$> internaliseExp (desc ++ "_zip_x") x - <*> internaliseExp (desc ++ "_zip_y") y - - handleRest [x] "unzip" = Just $ flip internaliseExp x - - handleRest [x] "trace" = Just $ flip internaliseExp x - handleRest [x] "break" = Just $ flip internaliseExp x - - handleRest _ _ = Nothing - - toSigned int_to e desc = do - e' <- internaliseExp1 "trunc_arg" e - case E.typeOf e of - E.Scalar (E.Prim E.Bool) -> - letTupExp' desc $ I.If e' (resultBody [intConst int_to 1]) - (resultBody [intConst int_to 0]) $ - ifCommon [I.Prim $ I.IntType int_to] - E.Scalar (E.Prim (E.Signed int_from)) -> - letTupExp' desc $ I.BasicOp $ I.ConvOp (I.SExt int_from int_to) e' - E.Scalar (E.Prim (E.Unsigned int_from)) -> - letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' - E.Scalar (E.Prim (E.FloatType float_from)) -> - letTupExp' desc $ I.BasicOp $ I.ConvOp (I.FPToSI float_from int_to) e' - _ -> error "Futhark.Internalise: non-numeric type in ToSigned" - - toUnsigned int_to e desc = do - e' <- internaliseExp1 "trunc_arg" e - case E.typeOf e of - E.Scalar (E.Prim E.Bool) -> - letTupExp' desc $ I.If e' (resultBody [intConst int_to 1]) - (resultBody [intConst int_to 0]) $ - ifCommon [I.Prim $ I.IntType int_to] - E.Scalar (E.Prim (E.Signed int_from)) -> - letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' - E.Scalar (E.Prim (E.Unsigned int_from)) -> - letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' - E.Scalar (E.Prim (E.FloatType float_from)) -> - letTupExp' desc $ I.BasicOp $ I.ConvOp (I.FPToUI float_from int_to) e' - _ -> error "Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned" - - complementF e desc = do - e' <- internaliseExp1 "complement_arg" e - et <- subExpType e' - case et of I.Prim (I.IntType t) -> - letTupExp' desc $ I.BasicOp $ I.UnOp (I.Complement t) e' - I.Prim I.Bool -> - letTupExp' desc $ I.BasicOp $ I.UnOp I.Not e' - _ -> - 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 - 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 - sv_ts <- mapM lookupType svs - - svs' <- forM (zip svs sv_ts) $ \(sv,sv_t) -> do - let sv_shape = I.arrayShape sv_t - sv_w = arraySize 0 sv_t - - -- Generate an assertion and reshapes to ensure that sv and si' are the same - -- size. - cmp <- letSubExp "write_cmp" $ I.BasicOp $ - I.CmpOp (I.CmpEq I.int32) si_w sv_w - c <- assert "write_cert" cmp - "length of index and value array does not match" loc - certifying c $ 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" - 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 - 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 - results <- forM outs $ \name -> - letSubExp "write_res" $ I.BasicOp $ I.SubExp $ I.Var name - ensureResultShape "scatter value has wrong size" loc - bodyTypes $ resultBody results - - let lam = I.Lambda { I.lambdaParams = bodyParams - , I.lambdaReturnType = bodyTypes - , I.lambdaBody = body - } - sivs = si' : svs' - - let sa_ws = map (arraySize 0) sa_ts - letTupExp' desc $ I.Op $ I.Scatter si_w lam sivs $ zip3 sa_ws (repeat 1) sas - -funcall :: String -> QualName VName -> [SubExp] -> SrcLoc - -> InternaliseM ([SubExp], [I.ExtType]) -funcall desc (QualName _ fname) args loc = do - (fname', closure, shapes, value_paramts, fun_params, rettype_fun) <- - lookupFunction fname - argts <- mapM subExpType args - - shapeargs <- argShapes shapes fun_params argts - let diets = replicate (length closure + length shapeargs) I.ObservePrim ++ - map I.diet value_paramts - args' <- ensureArgShapes "function arguments of wrong shape" - loc (map I.paramName fun_params) - (map I.paramType fun_params) (map I.Var closure ++ shapeargs ++ args) - argts' <- mapM subExpType args' - case rettype_fun $ zip args' argts' of - Nothing -> error $ "Cannot apply " ++ pretty fname ++ " to arguments\n " ++ - pretty args' ++ "\nof types\n " ++ - pretty argts' ++ - "\nFunction has parameters\n " ++ pretty fun_params - Just ts -> do - safety <- askSafety - attrs <- asks envAttrs - ses <- attributing attrs $ letTupExp' desc $ - I.Apply fname' (zip args' diets) ts (safety, loc, mempty) - return (ses, map I.fromDecl ts) - --- Bind existential names defined by an expression, based on the --- concrete values that expression evaluated to. This most --- importantly should be done after function calls, but also --- everything else that can produce existentials in the source --- language. -bindExtSizes :: E.StructType -> [VName] -> [SubExp] -> InternaliseM () -bindExtSizes ret retext ses = do - ts <- internaliseType ret - ses_ts <- mapM subExpType ses - - let combine t1 t2 = - mconcat $ zipWith combine' (arrayExtDims t1) (arrayDims t2) - combine' (I.Free (I.Var v)) se - | v `elem` retext = M.singleton v se - combine' _ _ = mempty - - forM_ (M.toList $ mconcat $ zipWith combine ts ses_ts) $ \(v, se) -> - letBindNames [v] $ BasicOp $ SubExp se - -askSafety :: InternaliseM Safety -askSafety = do check <- asks envDoBoundsChecks - return $ if check then I.Safe else I.Unsafe - --- Implement partitioning using maps, scans and writes. -partitionWithSOACS :: Int -> I.Lambda -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp]) -partitionWithSOACS k lam arrs = do - arr_ts <- mapM lookupType arrs - let w = arraysSize 0 arr_ts - classes_and_increments <- letTupExp "increments" $ I.Op $ I.Screma w (mapSOAC lam) arrs - (classes, increments) <- case classes_and_increments of - classes : increments -> return (classes, take k increments) - _ -> error "partitionWithSOACS" - - add_lam_x_params <- - replicateM k $ I.Param <$> newVName "x" <*> pure (I.Prim int32) - add_lam_y_params <- - replicateM k $ I.Param <$> newVName "y" <*> pure (I.Prim int32) - add_lam_body <- runBodyBinder $ - localScope (scopeOfLParams $ add_lam_x_params++add_lam_y_params) $ - fmap resultBody $ forM (zip add_lam_x_params add_lam_y_params) $ \(x,y) -> - letSubExp "z" $ I.BasicOp $ I.BinOp (I.Add Int32 I.OverflowUndef) - (I.Var $ I.paramName x) (I.Var $ I.paramName y) - let add_lam = I.Lambda { I.lambdaBody = add_lam_body - , I.lambdaParams = add_lam_x_params ++ add_lam_y_params - , I.lambdaReturnType = replicate k $ I.Prim int32 - } - nes = replicate (length increments) $ constant (0::Int32) - - scan <- I.scanSOAC [I.Scan add_lam nes] - all_offsets <- letTupExp "offsets" $ I.Op $ I.Screma w scan increments - - -- We have the offsets for each of the partitions, but we also need - -- the total sizes, which are the last elements in the offests. We - -- just have to be careful in case the array is empty. - last_index <- letSubExp "last_index" $ I.BasicOp $ I.BinOp (I.Sub Int32 OverflowUndef) w $ constant (1::Int32) - nonempty_body <- runBodyBinder $ fmap resultBody $ forM all_offsets $ \offset_array -> - letSubExp "last_offset" $ I.BasicOp $ I.Index offset_array [I.DimFix last_index] - let empty_body = resultBody $ replicate k $ constant (0::Int32) - is_empty <- letSubExp "is_empty" $ I.BasicOp $ I.CmpOp (CmpEq int32) w $ constant (0::Int32) - sizes <- letTupExp "partition_size" $ - I.If is_empty empty_body nonempty_body $ - ifCommon $ replicate k $ I.Prim int32 - - -- The total size of all partitions must necessarily be equal to the - -- size of the input array. - - -- Create scratch arrays for the result. - blanks <- forM arr_ts $ \arr_t -> - letExp "partition_dest" $ I.BasicOp $ - Scratch (elemType arr_t) (w : drop 1 (I.arrayDims arr_t)) - - -- Now write into the result. - write_lam <- do - c_param <- I.Param <$> newVName "c" <*> pure (I.Prim int32) - offset_params <- replicateM k $ I.Param <$> newVName "offset" <*> pure (I.Prim int32) - value_params <- forM arr_ts $ \arr_t -> - I.Param <$> newVName "v" <*> pure (I.rowType arr_t) - (offset, offset_stms) <- collectStms $ mkOffsetLambdaBody (map I.Var sizes) - (I.Var $ I.paramName c_param) 0 offset_params - return I.Lambda { I.lambdaParams = c_param : offset_params ++ value_params - , I.lambdaReturnType = replicate (length arr_ts) (I.Prim int32) ++ - map I.rowType arr_ts - , I.lambdaBody = mkBody offset_stms $ - replicate (length arr_ts) offset ++ - map (I.Var . I.paramName) value_params - } - results <- letTupExp "partition_res" $ I.Op $ I.Scatter w - write_lam (classes : all_offsets ++ arrs) $ - zip3 (repeat w) (repeat 1) blanks - sizes' <- letSubExp "partition_sizes" $ I.BasicOp $ - I.ArrayLit (map I.Var sizes) $ I.Prim int32 - return (map I.Var results, [sizes']) - where - mkOffsetLambdaBody :: [SubExp] - -> SubExp - -> Int - -> [I.LParam] - -> InternaliseM SubExp - mkOffsetLambdaBody _ _ _ [] = - return $ constant (-1::Int32) - mkOffsetLambdaBody sizes c i (p:ps) = do - is_this_one <- letSubExp "is_this_one" $ I.BasicOp $ I.CmpOp (CmpEq int32) c (constant i) - next_one <- mkOffsetLambdaBody sizes c (i+1) ps - this_one <- letSubExp "this_offset" =<< - foldBinOp (Add Int32 OverflowUndef) (constant (-1::Int32)) - (I.Var (I.paramName p) : take i sizes) - letSubExp "total_res" $ I.If is_this_one - (resultBody [this_one]) (resultBody [next_one]) $ ifCommon [I.Prim int32] - -typeExpForError :: E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp] -typeExpForError (E.TEVar qn _) = - return [ErrorString $ pretty qn] -typeExpForError (E.TEUnique te _) = - ("*":) <$> typeExpForError te -typeExpForError (E.TEArray te d _) = do - d' <- dimExpForError d - te' <- typeExpForError te - return $ ["[", d', "]"] ++ te' -typeExpForError (E.TETuple tes _) = do - tes' <- mapM typeExpForError tes - return $ ["("] ++ intercalate [", "] tes' ++ [")"] -typeExpForError (E.TERecord fields _) = do - fields' <- mapM onField fields - return $ ["{"] ++ intercalate [", "] fields' ++ ["}"] - where onField (k, te) = - (ErrorString (pretty k ++ ": "):) <$> typeExpForError te -typeExpForError (E.TEArrow _ t1 t2 _) = do - t1' <- typeExpForError t1 - t2' <- typeExpForError t2 - return $ t1' ++ [" -> "] ++ t2' -typeExpForError (E.TEApply t arg _) = do - t' <- typeExpForError t - arg' <- case arg of TypeArgExpType argt -> typeExpForError argt - TypeArgExpDim d _ -> pure <$> dimExpForError d - return $ t' ++ [" "] ++ arg' -typeExpForError (E.TESum cs _) = do - cs' <- mapM (onClause . snd) cs - return $ intercalate [" | "] cs' - where onClause c = do - c' <- mapM typeExpForError c - return $ intercalate [" "] c' - -dimExpForError :: E.DimExp VName -> InternaliseM (ErrorMsgPart SubExp) -dimExpForError (DimExpNamed d _) = do - substs <- lookupSubst $ E.qualLeaf d - d' <- case substs of - Just [v] -> return v - _ -> return $ I.Var $ E.qualLeaf d - return $ ErrorInt32 d' -dimExpForError (DimExpConst d _) = - return $ ErrorString $ pretty d -dimExpForError DimExpAny = return "" - --- A smart constructor that compacts neighbouring literals for easier --- reading in the IR. -errorMsg :: [ErrorMsgPart a] -> ErrorMsg a -errorMsg = ErrorMsg . compact - where compact [] = [] - compact (ErrorString x : ErrorString y : parts) = - compact (ErrorString (x++y) : parts) - compact (x:y) = x : compact y +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE Safe #-} +{-# LANGUAGE Strict #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} + +-- | +-- +-- This module implements a transformation from source to core +-- Futhark. +module Futhark.Internalise (internaliseProg) where + +import Control.Monad.Reader +import Data.Bitraversable +import Data.List (find, intercalate, intersperse, nub, transpose) +import qualified Data.List.NonEmpty as NE +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Futhark.IR.SOACS as I hiding (stmPattern) +import Futhark.Internalise.AccurateSizes +import Futhark.Internalise.Bindings +import Futhark.Internalise.Defunctionalise as Defunctionalise +import Futhark.Internalise.Defunctorise as Defunctorise +import Futhark.Internalise.Lambdas +import Futhark.Internalise.Monad as I +import Futhark.Internalise.Monomorphise as Monomorphise +import Futhark.Internalise.TypesValues +import Futhark.Transform.Rename as I +import Futhark.Util (splitAt3) +import Language.Futhark as E hiding (TypeArg) +import Language.Futhark.Semantic (Imports) + +-- | Convert a program in source Futhark to a program in the Futhark +-- core language. +internaliseProg :: + MonadFreshNames m => + Bool -> + Imports -> + m (I.Prog SOACS) +internaliseProg always_safe prog = do + prog_decs <- Defunctorise.transformProg prog + prog_decs' <- Monomorphise.transformProg prog_decs + prog_decs'' <- Defunctionalise.transformProg prog_decs' + (consts, funs) <- + runInternaliseM always_safe (internaliseValBinds prog_decs'') + I.renameProg $ I.Prog consts funs + +internaliseAttr :: E.AttrInfo -> Attr +internaliseAttr (E.AttrAtom v) = I.AttrAtom v +internaliseAttr (E.AttrComp f attrs) = I.AttrComp f $ map internaliseAttr attrs + +internaliseAttrs :: [E.AttrInfo] -> Attrs +internaliseAttrs = mconcat . map (oneAttr . internaliseAttr) + +internaliseValBinds :: [E.ValBind] -> InternaliseM () +internaliseValBinds = mapM_ internaliseValBind + +internaliseFunName :: VName -> [E.Pattern] -> InternaliseM Name +internaliseFunName ofname [] = return $ nameFromString $ pretty ofname ++ "f" +internaliseFunName ofname _ = do + info <- lookupFunction' ofname + -- In some rare cases involving local functions, the same function + -- name may be re-used in multiple places. We check whether the + -- function name has already been used, and generate a new one if + -- so. + case info of + Just _ -> nameFromString . pretty <$> newNameFromString (baseString ofname) + Nothing -> return $ nameFromString $ pretty ofname + +internaliseValBind :: E.ValBind -> InternaliseM () +internaliseValBind fb@(E.ValBind entry fname retdecl (Info (rettype, _)) tparams params body _ attrs loc) = do + localConstsScope $ + bindingParams tparams params $ \shapeparams params' -> do + let shapenames = map I.paramName shapeparams + normal_params = shapenames ++ map I.paramName (concat params') + normal_param_names = namesFromList normal_params + + fname' <- internaliseFunName fname params + + msg <- case retdecl of + Just dt -> + errorMsg + . ("Function return value does not match shape of type " :) + <$> typeExpForError dt + Nothing -> return $ errorMsg ["Function return value does not match shape of declared return type."] + + ((rettype', body_res), body_stms) <- collectStms $ do + body_res <- internaliseExp "res" body + rettype_bad <- internaliseReturnType rettype + let rettype' = zeroExts rettype_bad + return (rettype', body_res) + body' <- + ensureResultExtShape msg loc (map I.fromDecl rettype') $ + mkBody body_stms body_res + + constants <- allConsts + let free_in_fun = + freeIn body' + `namesSubtract` normal_param_names + `namesSubtract` constants + + used_free_params <- forM (namesToList free_in_fun) $ \v -> do + v_t <- lookupType v + return $ Param v $ toDecl v_t Nonunique + + let free_shape_params = + map (`Param` I.Prim int64) $ + concatMap (I.shapeVars . I.arrayShape . I.paramType) used_free_params + free_params = nub $ free_shape_params ++ used_free_params + all_params = free_params ++ shapeparams ++ concat params' + + let fd = + I.FunDef + Nothing + (internaliseAttrs attrs) + fname' + rettype' + all_params + body' + + if null params' + then bindConstant fname fd + else + bindFunction + fname + fd + ( fname', + map I.paramName free_params, + shapenames, + map declTypeOf $ concat params', + all_params, + applyRetType rettype' all_params + ) + + case entry of + Just (Info entry') -> generateEntryPoint entry' fb + Nothing -> return () + where + zeroExts ts = generaliseExtTypes ts ts + +allDimsFreshInType :: MonadFreshNames m => E.PatternType -> m E.PatternType +allDimsFreshInType = bitraverse onDim pure + where + onDim (E.NamedDim v) = + E.NamedDim . E.qualName <$> newVName (baseString $ E.qualLeaf v) + onDim _ = + E.NamedDim . E.qualName <$> newVName "size" + +-- | Replace all named dimensions with a fresh name, and remove all +-- constant dimensions. The point is to remove the constraints, but +-- keep the names around. We use this for constructing the entry +-- point parameters. +allDimsFreshInPat :: MonadFreshNames m => E.Pattern -> m E.Pattern +allDimsFreshInPat (PatternAscription p _ _) = + allDimsFreshInPat p +allDimsFreshInPat (PatternParens p _) = + allDimsFreshInPat p +allDimsFreshInPat (Id v (Info t) loc) = + Id v <$> (Info <$> allDimsFreshInType t) <*> pure loc +allDimsFreshInPat (TuplePattern ps loc) = + TuplePattern <$> mapM allDimsFreshInPat ps <*> pure loc +allDimsFreshInPat (RecordPattern ps loc) = + RecordPattern <$> mapM (traverse allDimsFreshInPat) ps <*> pure loc +allDimsFreshInPat (Wildcard (Info t) loc) = + Wildcard <$> (Info <$> allDimsFreshInType t) <*> pure loc +allDimsFreshInPat (PatternLit e (Info t) loc) = + PatternLit e <$> (Info <$> allDimsFreshInType t) <*> pure loc +allDimsFreshInPat (PatternConstr c (Info t) pats loc) = + PatternConstr c <$> (Info <$> allDimsFreshInType t) + <*> mapM allDimsFreshInPat pats + <*> pure loc + +generateEntryPoint :: E.EntryPoint -> E.ValBind -> InternaliseM () +generateEntryPoint (E.EntryPoint e_paramts e_rettype) vb = localConstsScope $ do + let (E.ValBind _ ofname _ (Info (rettype, _)) _ params _ _ attrs loc) = vb + -- We replace all shape annotations, so there should be no constant + -- parameters here. + params_fresh <- mapM allDimsFreshInPat params + let tparams = + map (`E.TypeParamDim` mempty) $ + S.toList $ + mconcat $ map E.patternDimNames params_fresh + bindingParams tparams params_fresh $ \shapeparams params' -> do + entry_rettype <- internaliseEntryReturnType $ anySizes rettype + let entry' = entryPoint (zip e_paramts params') (e_rettype, entry_rettype) + args = map (I.Var . I.paramName) $ concat params' + + entry_body <- insertStmsM $ do + -- Special case the (rare) situation where the entry point is + -- not a function. + maybe_const <- lookupConst ofname + vals <- case maybe_const of + Just ses -> + return ses + Nothing -> + fst <$> funcall "entry_result" (E.qualName ofname) args loc + ctx <- + extractShapeContext (concat entry_rettype) + <$> mapM (fmap I.arrayDims . subExpType) vals + resultBodyM (ctx ++ vals) + + addFunDef $ + I.FunDef + (Just entry') + (internaliseAttrs attrs) + (baseName ofname) + (concat entry_rettype) + (shapeparams ++ concat params') + entry_body + +entryPoint :: + [(E.EntryType, [I.FParam])] -> + ( E.EntryType, + [[I.TypeBase ExtShape Uniqueness]] + ) -> + I.EntryPoint +entryPoint params (eret, crets) = + ( concatMap (entryPointType . preParam) params, + case ( isTupleRecord $ entryType eret, + entryAscribed eret + ) of + (Just ts, Just (E.TETuple e_ts _)) -> + concatMap entryPointType $ + zip (zipWith E.EntryType ts (map Just e_ts)) crets + (Just ts, Nothing) -> + concatMap entryPointType $ + zip (map (`E.EntryType` Nothing) ts) crets + _ -> + entryPointType (eret, concat crets) + ) + where + preParam (e_t, ps) = (e_t, staticShapes $ map I.paramDeclType ps) + + entryPointType (t, ts) + | E.Scalar (E.Prim E.Unsigned {}) <- E.entryType t = + [I.TypeUnsigned] + | E.Array _ _ (E.Prim E.Unsigned {}) _ <- E.entryType t = + [I.TypeUnsigned] + | E.Scalar E.Prim {} <- E.entryType t = + [I.TypeDirect] + | E.Array _ _ E.Prim {} _ <- E.entryType t = + [I.TypeDirect] + | otherwise = + [I.TypeOpaque desc $ length ts] + where + desc = maybe (pretty t') typeExpOpaqueName $ E.entryAscribed t + t' = noSizes (E.entryType t) `E.setUniqueness` Nonunique + typeExpOpaqueName (TEApply te TypeArgExpDim {} _) = + typeExpOpaqueName te + typeExpOpaqueName (TEArray te _ _) = + let (d, te') = withoutDims te + in "arr_" ++ typeExpOpaqueName te' + ++ "_" + ++ show (1 + d) + ++ "d" + typeExpOpaqueName te = pretty te + + withoutDims (TEArray te _ _) = + let (d, te') = withoutDims te + in (d + 1, te') + withoutDims te = (0 :: Int, te) + +internaliseIdent :: E.Ident -> InternaliseM I.VName +internaliseIdent (E.Ident name (Info tp) loc) = + case tp of + E.Scalar E.Prim {} -> return name + _ -> + error $ + "Futhark.Internalise.internaliseIdent: asked to internalise non-prim-typed ident '" + ++ pretty name + ++ " of type " + ++ pretty tp + ++ " at " + ++ locStr loc + ++ "." + +internaliseBody :: E.Exp -> InternaliseM Body +internaliseBody e = insertStmsM $ resultBody <$> internaliseExp "res" e + +bodyFromStms :: + InternaliseM (Result, a) -> + InternaliseM (Body, a) +bodyFromStms m = do + ((res, a), stms) <- collectStms m + (,a) <$> mkBodyM stms res + +internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp] +internaliseExp desc (E.Parens e _) = + internaliseExp desc e +internaliseExp desc (E.QualParens _ e _) = + internaliseExp desc e +internaliseExp desc (E.StringLit vs _) = + fmap pure $ + letSubExp desc $ + I.BasicOp $ I.ArrayLit (map constant vs) $ I.Prim int8 +internaliseExp _ (E.Var (E.QualName _ name) (Info t) loc) = do + subst <- lookupSubst name + case subst of + Just substs -> return substs + Nothing -> do + -- If this identifier is the name of a constant, we have to turn it + -- into a call to the corresponding function. + is_const <- lookupConst name + case is_const of + Just ses -> return ses + Nothing -> (: []) . I.Var <$> internaliseIdent (E.Ident name (Info t) loc) +internaliseExp desc (E.Index e idxs (Info ret, Info retext) loc) = do + vs <- internaliseExpToVars "indexed" e + dims <- case vs of + [] -> return [] -- Will this happen? + v : _ -> I.arrayDims <$> lookupType v + (idxs', cs) <- internaliseSlice loc dims idxs + let index v = do + v_t <- lookupType v + return $ I.BasicOp $ I.Index v $ fullSlice v_t idxs' + ses <- certifying cs $ letSubExps desc =<< mapM index vs + bindExtSizes (E.toStruct ret) retext ses + return ses + +-- XXX: we map empty records and tuples to bools, because otherwise +-- arrays of unit will lose their sizes. +internaliseExp _ (E.TupLit [] _) = + return [constant True] +internaliseExp _ (E.RecordLit [] _) = + return [constant True] +internaliseExp desc (E.TupLit es _) = concat <$> mapM (internaliseExp desc) es +internaliseExp desc (E.RecordLit orig_fields _) = + concatMap snd . sortFields . M.unions <$> mapM internaliseField orig_fields + where + internaliseField (E.RecordFieldExplicit name e _) = + M.singleton name <$> internaliseExp desc e + internaliseField (E.RecordFieldImplicit name t loc) = + internaliseField $ + E.RecordFieldExplicit + (baseName name) + (E.Var (E.qualName name) t loc) + loc +internaliseExp desc (E.ArrayLit es (Info arr_t) loc) + -- If this is a multidimensional array literal of primitives, we + -- treat it specially by flattening it out followed by a reshape. + -- This cuts down on the amount of statements that are produced, and + -- thus allows us to efficiently handle huge array literals - a + -- corner case, but an important one. + | Just ((eshape, e') : es') <- mapM isArrayLiteral es, + not $ null eshape, + all ((eshape ==) . fst) es', + Just basetype <- E.peelArray (length eshape) arr_t = do + let flat_lit = E.ArrayLit (e' ++ concatMap snd es') (Info basetype) loc + new_shape = length es : eshape + flat_arrs <- internaliseExpToVars "flat_literal" flat_lit + forM flat_arrs $ \flat_arr -> do + flat_arr_t <- lookupType flat_arr + let new_shape' = + reshapeOuter + (map (DimNew . intConst Int64 . toInteger) new_shape) + 1 + $ I.arrayShape flat_arr_t + letSubExp desc $ I.BasicOp $ I.Reshape new_shape' flat_arr + | otherwise = do + es' <- mapM (internaliseExp "arr_elem") es + arr_t_ext <- internaliseReturnType (E.toStruct arr_t) + + rowtypes <- + case mapM (fmap rowType . hasStaticShape . I.fromDecl) arr_t_ext of + Just ts -> pure ts + Nothing -> + -- XXX: the monomorphiser may create single-element array + -- literals with an unknown row type. In those cases we + -- need to look at the types of the actual elements. + -- Fixing this in the monomorphiser is a lot more tricky + -- than just working around it here. + case es' of + [] -> error $ "internaliseExp ArrayLit: existential type: " ++ pretty arr_t + e' : _ -> mapM subExpType e' + + let arraylit ks rt = do + ks' <- + mapM + ( ensureShape + "shape of element differs from shape of first element" + loc + rt + "elem_reshaped" + ) + ks + return $ I.BasicOp $ I.ArrayLit ks' rt + + letSubExps desc + =<< if null es' + then mapM (arraylit []) rowtypes + else zipWithM arraylit (transpose es') rowtypes + where + isArrayLiteral :: E.Exp -> Maybe ([Int], [E.Exp]) + isArrayLiteral (E.ArrayLit inner_es _ _) = do + (eshape, e) : inner_es' <- mapM isArrayLiteral inner_es + guard $ all ((eshape ==) . fst) inner_es' + return (length inner_es : eshape, e ++ concatMap snd inner_es') + isArrayLiteral e = + Just ([], [e]) +internaliseExp desc (E.Range start maybe_second end (Info ret, Info retext) loc) = do + start' <- internaliseExp1 "range_start" start + end' <- internaliseExp1 "range_end" $ case end of + DownToExclusive e -> e + ToInclusive e -> e + UpToExclusive e -> e + maybe_second' <- + traverse (internaliseExp1 "range_second") maybe_second + + -- Construct an error message in case the range is invalid. + let conv = case E.typeOf start of + E.Scalar (E.Prim (E.Unsigned _)) -> asIntZ Int64 + _ -> asIntS Int64 + start'_i64 <- conv start' + end'_i64 <- conv end' + maybe_second'_i64 <- traverse conv maybe_second' + let errmsg = + errorMsg $ + ["Range "] + ++ [ErrorInt64 start'_i64] + ++ ( case maybe_second'_i64 of + Nothing -> [] + Just second_i64 -> ["..", ErrorInt64 second_i64] + ) + ++ ( case end of + DownToExclusive {} -> ["..>"] + ToInclusive {} -> ["..."] + UpToExclusive {} -> ["..<"] + ) + ++ [ErrorInt64 end'_i64, " is invalid."] + + (it, le_op, lt_op) <- + case E.typeOf start of + E.Scalar (E.Prim (E.Signed it)) -> return (it, CmpSle it, CmpSlt it) + E.Scalar (E.Prim (E.Unsigned it)) -> return (it, CmpUle it, CmpUlt it) + start_t -> error $ "Start value in range has type " ++ pretty start_t + + let one = intConst it 1 + negone = intConst it (-1) + default_step = case end of + DownToExclusive {} -> negone + ToInclusive {} -> one + UpToExclusive {} -> one + + (step, step_zero) <- case maybe_second' of + Just second' -> do + subtracted_step <- + letSubExp "subtracted_step" $ + I.BasicOp $ I.BinOp (I.Sub it I.OverflowWrap) second' start' + step_zero <- letSubExp "step_zero" $ I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) start' second' + return (subtracted_step, step_zero) + Nothing -> + return (default_step, constant False) + + step_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum it) step + step_sign_i64 <- asIntS Int64 step_sign + + bounds_invalid_downwards <- + letSubExp "bounds_invalid_downwards" $ + I.BasicOp $ I.CmpOp le_op start' end' + bounds_invalid_upwards <- + letSubExp "bounds_invalid_upwards" $ + I.BasicOp $ I.CmpOp lt_op end' start' + + (distance, step_wrong_dir, bounds_invalid) <- case end of + DownToExclusive {} -> do + step_wrong_dir <- + letSubExp "step_wrong_dir" $ + I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) step_sign one + distance <- + letSubExp "distance" $ + I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) start' end' + distance_i64 <- asIntS Int64 distance + return (distance_i64, step_wrong_dir, bounds_invalid_downwards) + UpToExclusive {} -> do + step_wrong_dir <- + letSubExp "step_wrong_dir" $ + I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) step_sign negone + distance <- letSubExp "distance" $ I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) end' start' + distance_i64 <- asIntS Int64 distance + return (distance_i64, step_wrong_dir, bounds_invalid_upwards) + ToInclusive {} -> do + downwards <- + letSubExp "downwards" $ + I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) step_sign negone + distance_downwards_exclusive <- + letSubExp "distance_downwards_exclusive" $ + I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) start' end' + distance_upwards_exclusive <- + letSubExp "distance_upwards_exclusive" $ + I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) end' start' + + bounds_invalid <- + letSubExp "bounds_invalid" $ + I.If + downwards + (resultBody [bounds_invalid_downwards]) + (resultBody [bounds_invalid_upwards]) + $ ifCommon [I.Prim I.Bool] + distance_exclusive <- + letSubExp "distance_exclusive" $ + I.If + downwards + (resultBody [distance_downwards_exclusive]) + (resultBody [distance_upwards_exclusive]) + $ ifCommon [I.Prim $ IntType it] + distance_exclusive_i64 <- asIntS Int64 distance_exclusive + distance <- + letSubExp "distance" $ + I.BasicOp $ + I.BinOp + (Add Int64 I.OverflowWrap) + distance_exclusive_i64 + (intConst Int64 1) + return (distance, constant False, bounds_invalid) + + step_invalid <- + letSubExp "step_invalid" $ + I.BasicOp $ I.BinOp I.LogOr step_wrong_dir step_zero + + invalid <- + letSubExp "range_invalid" $ + I.BasicOp $ I.BinOp I.LogOr step_invalid bounds_invalid + valid <- letSubExp "valid" $ I.BasicOp $ I.UnOp I.Not invalid + cs <- assert "range_valid_c" valid errmsg loc + + step_i64 <- asIntS Int64 step + pos_step <- + letSubExp "pos_step" $ + I.BasicOp $ I.BinOp (Mul Int64 I.OverflowWrap) step_i64 step_sign_i64 + + num_elems <- + certifying cs $ + letSubExp "num_elems" $ + I.BasicOp $ I.BinOp (SDivUp Int64 I.Unsafe) distance pos_step + + se <- letSubExp desc (I.BasicOp $ I.Iota num_elems start' step it) + bindExtSizes (E.toStruct ret) retext [se] + return [se] +internaliseExp desc (E.Ascript e _ _) = + internaliseExp desc e +internaliseExp desc (E.Coerce e (TypeDecl dt (Info et)) (Info ret, Info retext) loc) = do + ses <- internaliseExp desc e + ts <- internaliseReturnType et + dt' <- typeExpForError dt + bindExtSizes (E.toStruct ret) retext ses + forM (zip ses ts) $ \(e', t') -> do + dims <- arrayDims <$> subExpType e' + let parts = + ["Value of (core language) shape ("] + ++ intersperse ", " (map ErrorInt64 dims) + ++ [") cannot match shape of type `"] + ++ dt' + ++ ["`."] + ensureExtShape (errorMsg parts) loc (I.fromDecl t') desc e' +internaliseExp desc (E.Negate e _) = do + e' <- internaliseExp1 "negate_arg" e + et <- subExpType e' + case et of + I.Prim (I.IntType t) -> + letTupExp' desc $ I.BasicOp $ I.BinOp (I.Sub t I.OverflowWrap) (I.intConst t 0) e' + I.Prim (I.FloatType t) -> + letTupExp' desc $ I.BasicOp $ I.BinOp (I.FSub t) (I.floatConst t 0) e' + _ -> error "Futhark.Internalise.internaliseExp: non-numeric type in Negate" +internaliseExp desc e@E.Apply {} = do + (qfname, args, ret, retext) <- findFuncall e + -- Argument evaluation is outermost-in so that any existential sizes + -- created by function applications can be brought into scope. + let fname = nameFromString $ pretty $ baseName $ qualLeaf qfname + loc = srclocOf e + arg_desc = nameToString fname ++ "_arg" + + -- Some functions are magical (overloaded) and we handle that here. + ses <- + case () of + -- Overloaded functions never take array arguments (except + -- equality, but those cannot be existential), so we can safely + -- ignore the existential dimensions. + () + | Just internalise <- isOverloadedFunction qfname (map fst args) loc -> + internalise desc + | Just (rettype, _) <- M.lookup fname I.builtInFunctions -> do + let tag ses = [(se, I.Observe) | se <- ses] + args' <- reverse <$> mapM (internaliseArg arg_desc) (reverse args) + let args'' = concatMap tag args' + letTupExp' desc $ + I.Apply + fname + args'' + [I.Prim rettype] + (Safe, loc, []) + | otherwise -> do + args' <- concat . reverse <$> mapM (internaliseArg arg_desc) (reverse args) + fst <$> funcall desc qfname args' loc + + bindExtSizes ret retext ses + return ses +internaliseExp desc (E.LetPat pat e body (Info ret, Info retext) _) = do + ses <- internalisePat desc pat e body (internaliseExp desc) + bindExtSizes (E.toStruct ret) retext ses + return ses +internaliseExp desc (E.LetFun ofname (tparams, params, retdecl, Info rettype, body) letbody _ loc) = do + internaliseValBind $ + E.ValBind Nothing ofname retdecl (Info (rettype, [])) tparams params body Nothing mempty loc + internaliseExp desc letbody +internaliseExp desc (E.DoLoop sparams mergepat mergeexp form loopbody (Info (ret, retext)) loc) = do + ses <- internaliseExp "loop_init" mergeexp + ((loopbody', (form', shapepat, mergepat', mergeinit')), initstms) <- + collectStms $ handleForm ses form + + addStms initstms + mergeinit_ts' <- mapM subExpType mergeinit' + + ctxinit <- argShapes (map I.paramName shapepat) mergepat' mergeinit_ts' + + let ctxmerge = zip shapepat ctxinit + valmerge = zip mergepat' mergeinit' + dropCond = case form of + E.While {} -> drop 1 + _ -> id + + -- Ensure that the result of the loop matches the shapes of the + -- merge parameters. XXX: Ideally they should already match (by + -- the source language type rules), but some of our + -- transformations (esp. defunctionalisation) strips out some size + -- information. For a type-correct source program, these reshapes + -- should simplify away. + let merge = ctxmerge ++ valmerge + merge_ts = map (I.paramType . fst) merge + loopbody'' <- + localScope (scopeOfFParams $ map fst merge) $ + inScopeOf form' $ + insertStmsM $ + resultBodyM + =<< ensureArgShapes + "shape of loop result does not match shapes in loop parameter" + loc + (map (I.paramName . fst) ctxmerge) + merge_ts + =<< bodyBind loopbody' + + attrs <- asks envAttrs + loop_res <- + map I.Var . dropCond + <$> attributing + attrs + (letTupExp desc (I.DoLoop ctxmerge valmerge form' loopbody'')) + bindExtSizes (E.toStruct ret) retext loop_res + return loop_res + where + sparams' = map (`TypeParamDim` mempty) sparams + + forLoop mergepat' shapepat mergeinit form' = + bodyFromStms $ + inScopeOf form' $ do + ses <- internaliseExp "loopres" loopbody + sets <- mapM subExpType ses + shapeargs <- argShapes (map I.paramName shapepat) mergepat' sets + return + ( shapeargs ++ ses, + ( form', + shapepat, + mergepat', + mergeinit + ) + ) + + handleForm mergeinit (E.ForIn x arr) = do + arr' <- internaliseExpToVars "for_in_arr" arr + arr_ts <- mapM lookupType arr' + let w = arraysSize 0 arr_ts + + i <- newVName "i" + + bindingLoopParams sparams' mergepat $ + \shapepat mergepat' -> + bindingLambdaParams [x] (map rowType arr_ts) $ \x_params -> do + let loopvars = zip x_params arr' + forLoop mergepat' shapepat mergeinit $ + I.ForLoop i Int64 w loopvars + handleForm mergeinit (E.For i num_iterations) = do + num_iterations' <- internaliseExp1 "upper_bound" num_iterations + i' <- internaliseIdent i + num_iterations_t <- I.subExpType num_iterations' + it <- case num_iterations_t of + I.Prim (IntType it) -> return it + _ -> error "internaliseExp DoLoop: invalid type" + + bindingLoopParams sparams' mergepat $ + \shapepat mergepat' -> + forLoop mergepat' shapepat mergeinit $ + I.ForLoop i' it num_iterations' [] + handleForm mergeinit (E.While cond) = + bindingLoopParams sparams' mergepat $ \shapepat mergepat' -> do + mergeinit_ts <- mapM subExpType mergeinit + -- We need to insert 'cond' twice - once for the initial + -- condition (do we enter the loop at all?), and once with the + -- result values of the loop (do we continue into the next + -- iteration?). This is safe, as the type rules for the + -- external language guarantees that 'cond' does not consume + -- anything. + shapeinit <- argShapes (map I.paramName shapepat) mergepat' mergeinit_ts + + (loop_initial_cond, init_loop_cond_bnds) <- collectStms $ do + forM_ (zip shapepat shapeinit) $ \(p, se) -> + letBindNames [paramName p] $ BasicOp $ SubExp se + forM_ (zip mergepat' mergeinit) $ \(p, se) -> + unless (se == I.Var (paramName p)) $ + letBindNames [paramName p] $ + BasicOp $ + case se of + I.Var v + | not $ primType $ paramType p -> + Reshape (map DimCoercion $ arrayDims $ paramType p) v + _ -> SubExp se + internaliseExp1 "loop_cond" cond + + addStms init_loop_cond_bnds + + bodyFromStms $ do + ses <- internaliseExp "loopres" loopbody + sets <- mapM subExpType ses + loop_while <- newParam "loop_while" $ I.Prim I.Bool + shapeargs <- argShapes (map I.paramName shapepat) mergepat' sets + + -- Careful not to clobber anything. + loop_end_cond_body <- renameBody <=< insertStmsM $ do + forM_ (zip shapepat shapeargs) $ \(p, se) -> + unless (se == I.Var (paramName p)) $ + letBindNames [paramName p] $ BasicOp $ SubExp se + forM_ (zip mergepat' ses) $ \(p, se) -> + unless (se == I.Var (paramName p)) $ + letBindNames [paramName p] $ + BasicOp $ + case se of + I.Var v + | not $ primType $ paramType p -> + Reshape (map DimCoercion $ arrayDims $ paramType p) v + _ -> SubExp se + resultBody <$> internaliseExp "loop_cond" cond + loop_end_cond <- bodyBind loop_end_cond_body + + return + ( shapeargs ++ loop_end_cond ++ ses, + ( I.WhileLoop $ I.paramName loop_while, + shapepat, + loop_while : mergepat', + loop_initial_cond : mergeinit + ) + ) +internaliseExp desc (E.LetWith name src idxs ve body t loc) = do + let pat = E.Id (E.identName name) (E.identType name) loc + src_t = E.fromStruct <$> E.identType src + e = E.Update (E.Var (E.qualName $ E.identName src) src_t loc) idxs ve loc + internaliseExp desc $ E.LetPat pat e body (t, Info []) loc +internaliseExp desc (E.Update src slice ve loc) = do + ves <- internaliseExp "lw_val" ve + srcs <- internaliseExpToVars "src" src + dims <- case srcs of + [] -> return [] -- Will this happen? + v : _ -> I.arrayDims <$> lookupType v + (idxs', cs) <- internaliseSlice loc dims slice + + let comb sname ve' = do + sname_t <- lookupType sname + let full_slice = fullSlice sname_t idxs' + rowtype = sname_t `setArrayDims` sliceDims full_slice + ve'' <- + ensureShape + "shape of value does not match shape of source array" + loc + rowtype + "lw_val_correct_shape" + ve' + letInPlace desc sname full_slice $ BasicOp $ SubExp ve'' + certifying cs $ map I.Var <$> zipWithM comb srcs ves +internaliseExp desc (E.RecordUpdate src fields ve _ _) = do + src' <- internaliseExp desc src + ve' <- internaliseExp desc ve + replace (E.typeOf src `setAliases` ()) fields ve' src' + where + replace (E.Scalar (E.Record m)) (f : fs) ve' src' + | Just t <- M.lookup f m = do + i <- + fmap sum $ + mapM (internalisedTypeSize . snd) $ + takeWhile ((/= f) . fst) $ sortFields m + k <- internalisedTypeSize t + let (bef, to_update, aft) = splitAt3 i k src' + src'' <- replace t fs ve' to_update + return $ bef ++ src'' ++ aft + replace _ _ ve' _ = return ve' +internaliseExp desc (E.Attr attr e _) = + local f $ internaliseExp desc e + where + attrs = oneAttr $ internaliseAttr attr + f env + | "unsafe" `inAttrs` attrs, + not $ envSafe env = + env {envDoBoundsChecks = False} + | otherwise = + env {envAttrs = envAttrs env <> attrs} +internaliseExp desc (E.Assert e1 e2 (Info check) loc) = do + e1' <- internaliseExp1 "assert_cond" e1 + c <- assert "assert_c" e1' (errorMsg [ErrorString $ "Assertion is false: " <> check]) loc + -- Make sure there are some bindings to certify. + certifying c $ mapM rebind =<< internaliseExp desc e2 + where + rebind v = do + v' <- newVName "assert_res" + letBindNames [v'] $ I.BasicOp $ I.SubExp v + return $ I.Var v' +internaliseExp _ (E.Constr c es (Info (E.Scalar (E.Sum fs))) _) = do + (ts, constr_map) <- internaliseSumType $ M.map (map E.toStruct) fs + es' <- concat <$> mapM (internaliseExp "payload") es + + let noExt _ = return $ intConst Int64 0 + ts' <- instantiateShapes noExt $ map fromDecl ts + + case M.lookup c constr_map of + Just (i, js) -> + (intConst Int8 (toInteger i) :) <$> clauses 0 ts' (zip js es') + Nothing -> + error "internaliseExp Constr: missing constructor" + where + clauses j (t : ts) js_to_es + | Just e <- j `lookup` js_to_es = + (e :) <$> clauses (j + 1) ts js_to_es + | otherwise = do + blank <- letSubExp "zero" =<< eBlank t + (blank :) <$> clauses (j + 1) ts js_to_es + clauses _ [] _ = + return [] +internaliseExp _ (E.Constr _ _ (Info t) loc) = + error $ "internaliseExp: constructor with type " ++ pretty t ++ " at " ++ locStr loc +internaliseExp desc (E.Match e cs (Info ret, Info retext) _) = do + ses <- internaliseExp (desc ++ "_scrutinee") e + res <- + case NE.uncons cs of + (CasePat pCase eCase _, Nothing) -> do + (_, pertinent) <- generateCond pCase ses + internalisePat' pCase pertinent eCase (internaliseExp desc) + (c, Just cs') -> do + let CasePat pLast eLast _ = NE.last cs' + bFalse <- do + (_, pertinent) <- generateCond pLast ses + eLast' <- internalisePat' pLast pertinent eLast internaliseBody + foldM (\bf c' -> eBody $ return $ generateCaseIf ses c' bf) eLast' $ + reverse $ NE.init cs' + letTupExp' desc =<< generateCaseIf ses c bFalse + bindExtSizes (E.toStruct ret) retext res + return res + +-- The "interesting" cases are over, now it's mostly boilerplate. + +internaliseExp _ (E.Literal v _) = + return [I.Constant $ internalisePrimValue v] +internaliseExp _ (E.IntLit v (Info t) _) = + case t of + E.Scalar (E.Prim (E.Signed it)) -> + return [I.Constant $ I.IntValue $ intValue it v] + E.Scalar (E.Prim (E.Unsigned it)) -> + return [I.Constant $ I.IntValue $ intValue it v] + E.Scalar (E.Prim (E.FloatType ft)) -> + return [I.Constant $ I.FloatValue $ floatValue ft v] + _ -> error $ "internaliseExp: nonsensical type for integer literal: " ++ pretty t +internaliseExp _ (E.FloatLit v (Info t) _) = + case t of + E.Scalar (E.Prim (E.FloatType ft)) -> + return [I.Constant $ I.FloatValue $ floatValue ft v] + _ -> error $ "internaliseExp: nonsensical type for float literal: " ++ pretty t +internaliseExp desc (E.If ce te fe (Info ret, Info retext) _) = do + ses <- + letTupExp' desc + =<< eIf + (BasicOp . SubExp <$> internaliseExp1 "cond" ce) + (internaliseBody te) + (internaliseBody fe) + bindExtSizes (E.toStruct ret) retext ses + return ses + +-- Builtin operators are handled specially because they are +-- overloaded. +internaliseExp desc (E.BinOp (op, _) _ (xe, _) (ye, _) _ _ loc) + | Just internalise <- isOverloadedFunction op [xe, ye] loc = + internalise desc +-- User-defined operators are just the same as a function call. +internaliseExp + desc + ( E.BinOp + (op, oploc) + (Info t) + (xarg, Info (xt, xext)) + (yarg, Info (yt, yext)) + _ + (Info retext) + loc + ) = + internaliseExp desc $ + E.Apply + ( E.Apply + (E.Var op (Info t) oploc) + xarg + (Info (E.diet xt, xext)) + (Info $ foldFunType [E.fromStruct yt] t, Info []) + loc + ) + yarg + (Info (E.diet yt, yext)) + (Info t, Info retext) + loc +internaliseExp desc (E.Project k e (Info rt) _) = do + n <- internalisedTypeSize $ rt `setAliases` () + i' <- fmap sum $ + mapM internalisedTypeSize $ + case E.typeOf e `setAliases` () of + E.Scalar (Record fs) -> + map snd $ takeWhile ((/= k) . fst) $ sortFields fs + t -> [t] + take n . drop i' <$> internaliseExp desc e +internaliseExp _ e@E.Lambda {} = + error $ "internaliseExp: Unexpected lambda at " ++ locStr (srclocOf e) +internaliseExp _ e@E.OpSection {} = + error $ "internaliseExp: Unexpected operator section at " ++ locStr (srclocOf e) +internaliseExp _ e@E.OpSectionLeft {} = + error $ "internaliseExp: Unexpected left operator section at " ++ locStr (srclocOf e) +internaliseExp _ e@E.OpSectionRight {} = + error $ "internaliseExp: Unexpected right operator section at " ++ locStr (srclocOf e) +internaliseExp _ e@E.ProjectSection {} = + error $ "internaliseExp: Unexpected projection section at " ++ locStr (srclocOf e) +internaliseExp _ e@E.IndexSection {} = + error $ "internaliseExp: Unexpected index section at " ++ locStr (srclocOf e) + +internaliseArg :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp] +internaliseArg desc (arg, argdim) = do + arg' <- internaliseExp desc arg + case (arg', argdim) of + ([se], Just d) -> letBindNames [d] $ BasicOp $ SubExp se + _ -> return () + return arg' + +generateCond :: E.Pattern -> [I.SubExp] -> InternaliseM (I.SubExp, [I.SubExp]) +generateCond orig_p orig_ses = do + (cmps, pertinent, _) <- compares orig_p orig_ses + cmp <- letSubExp "matches" =<< eAll cmps + return (cmp, pertinent) + where + -- Literals are always primitive values. + compares (E.PatternLit e _ _) (se : ses) = do + e' <- internaliseExp1 "constant" e + t' <- elemType <$> subExpType se + cmp <- letSubExp "match_lit" $ I.BasicOp $ I.CmpOp (I.CmpEq t') e' se + return ([cmp], [se], ses) + compares (E.PatternConstr c (Info (E.Scalar (E.Sum fs))) pats _) (se : ses) = do + (payload_ts, m) <- internaliseSumType $ M.map (map toStruct) fs + case M.lookup c m of + Just (i, payload_is) -> do + let i' = intConst Int8 $ toInteger i + let (payload_ses, ses') = splitAt (length payload_ts) ses + cmp <- letSubExp "match_constr" $ I.BasicOp $ I.CmpOp (I.CmpEq int8) i' se + (cmps, pertinent, _) <- comparesMany pats $ map (payload_ses !!) payload_is + return (cmp : cmps, pertinent, ses') + Nothing -> + error "generateCond: missing constructor" + compares (E.PatternConstr _ (Info t) _ _) _ = + error $ "generateCond: PatternConstr has nonsensical type: " ++ pretty t + compares (E.Id _ t loc) ses = + compares (E.Wildcard t loc) ses + compares (E.Wildcard (Info t) _) ses = do + n <- internalisedTypeSize $ E.toStruct t + let (id_ses, rest_ses) = splitAt n ses + return ([], id_ses, rest_ses) + compares (E.PatternParens pat _) ses = + compares pat ses + compares (E.TuplePattern pats _) ses = + comparesMany pats ses + compares (E.RecordPattern fs _) ses = + comparesMany (map snd $ E.sortFields $ M.fromList fs) ses + compares (E.PatternAscription pat _ _) ses = + compares pat ses + compares pat [] = + error $ "generateCond: No values left for pattern " ++ pretty pat + + comparesMany [] ses = return ([], [], ses) + comparesMany (pat : pats) ses = do + (cmps1, pertinent1, ses') <- compares pat ses + (cmps2, pertinent2, ses'') <- comparesMany pats ses' + return + ( cmps1 <> cmps2, + pertinent1 <> pertinent2, + ses'' + ) + +generateCaseIf :: [I.SubExp] -> Case -> I.Body -> InternaliseM I.Exp +generateCaseIf ses (CasePat p eCase _) bFail = do + (cond, pertinent) <- generateCond p ses + eCase' <- internalisePat' p pertinent eCase internaliseBody + eIf (eSubExp cond) (return eCase') (return bFail) + +internalisePat :: + String -> + E.Pattern -> + E.Exp -> + E.Exp -> + (E.Exp -> InternaliseM a) -> + InternaliseM a +internalisePat desc p e body m = do + ses <- internaliseExp desc' e + internalisePat' p ses body m + where + desc' = case S.toList $ E.patternIdents p of + [v] -> baseString $ E.identName v + _ -> desc + +internalisePat' :: + E.Pattern -> + [I.SubExp] -> + E.Exp -> + (E.Exp -> InternaliseM a) -> + InternaliseM a +internalisePat' p ses body m = do + ses_ts <- mapM subExpType ses + stmPattern p ses_ts $ \pat_names -> do + forM_ (zip pat_names ses) $ \(v, se) -> + letBindNames [v] $ I.BasicOp $ I.SubExp se + m body + +internaliseSlice :: + SrcLoc -> + [SubExp] -> + [E.DimIndex] -> + InternaliseM ([I.DimIndex SubExp], Certificates) +internaliseSlice loc dims idxs = do + (idxs', oks, parts) <- unzip3 <$> zipWithM internaliseDimIndex dims idxs + ok <- letSubExp "index_ok" =<< eAll oks + let msg = + errorMsg $ + ["Index ["] ++ intercalate [", "] parts + ++ ["] out of bounds for array of shape ["] + ++ intersperse "][" (map ErrorInt64 $ take (length idxs) dims) + ++ ["]."] + c <- assert "index_certs" ok msg loc + return (idxs', c) + +internaliseDimIndex :: + SubExp -> + E.DimIndex -> + InternaliseM (I.DimIndex SubExp, SubExp, [ErrorMsgPart SubExp]) +internaliseDimIndex w (E.DimFix i) = do + (i', _) <- internaliseDimExp "i" i + let lowerBound = + I.BasicOp $ + I.CmpOp (I.CmpSle I.Int64) (I.constant (0 :: I.Int64)) i' + upperBound = + I.BasicOp $ + I.CmpOp (I.CmpSlt I.Int64) i' w + ok <- letSubExp "bounds_check" =<< eBinOp I.LogAnd (pure lowerBound) (pure upperBound) + return (I.DimFix i', ok, [ErrorInt64 i']) + +-- Special-case an important common case that otherwise leads to horrible code. +internaliseDimIndex + w + ( E.DimSlice + Nothing + Nothing + (Just (E.Negate (E.IntLit 1 _ _) _)) + ) = do + w_minus_1 <- + letSubExp "w_minus_1" $ + BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) w one + return + ( I.DimSlice w_minus_1 w $ intConst Int64 (-1), + constant True, + mempty + ) + where + one = constant (1 :: Int64) +internaliseDimIndex w (E.DimSlice i j s) = do + s' <- maybe (return one) (fmap fst . internaliseDimExp "s") s + s_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum Int64) s' + backwards <- letSubExp "backwards" $ I.BasicOp $ I.CmpOp (I.CmpEq int64) s_sign negone + w_minus_1 <- letSubExp "w_minus_1" $ BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) w one + let i_def = + letSubExp "i_def" $ + I.If + backwards + (resultBody [w_minus_1]) + (resultBody [zero]) + $ ifCommon [I.Prim int64] + j_def = + letSubExp "j_def" $ + I.If + backwards + (resultBody [negone]) + (resultBody [w]) + $ ifCommon [I.Prim int64] + i' <- maybe i_def (fmap fst . internaliseDimExp "i") i + j' <- maybe j_def (fmap fst . internaliseDimExp "j") j + j_m_i <- letSubExp "j_m_i" $ BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) j' i' + -- Something like a division-rounding-up, but accomodating negative + -- operands. + let divRounding x y = + eBinOp + (SQuot Int64 Unsafe) + ( eBinOp + (Add Int64 I.OverflowWrap) + x + (eBinOp (Sub Int64 I.OverflowWrap) y (eSignum $ toExp s')) + ) + y + n <- letSubExp "n" =<< divRounding (toExp j_m_i) (toExp s') + + -- Bounds checks depend on whether we are slicing forwards or + -- backwards. If forwards, we must check '0 <= i && i <= j'. If + -- backwards, '-1 <= j && j <= i'. In both cases, we check '0 <= + -- i+n*s && i+(n-1)*s < w'. We only check if the slice is nonempty. + empty_slice <- letSubExp "empty_slice" $ I.BasicOp $ I.CmpOp (CmpEq int64) n zero + + m <- letSubExp "m" $ I.BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) n one + m_t_s <- letSubExp "m_t_s" $ I.BasicOp $ I.BinOp (Mul Int64 I.OverflowWrap) m s' + i_p_m_t_s <- letSubExp "i_p_m_t_s" $ I.BasicOp $ I.BinOp (Add Int64 I.OverflowWrap) i' m_t_s + zero_leq_i_p_m_t_s <- + letSubExp "zero_leq_i_p_m_t_s" $ + I.BasicOp $ I.CmpOp (I.CmpSle Int64) zero i_p_m_t_s + i_p_m_t_s_leq_w <- + letSubExp "i_p_m_t_s_leq_w" $ + I.BasicOp $ I.CmpOp (I.CmpSle Int64) i_p_m_t_s w + i_p_m_t_s_lth_w <- + letSubExp "i_p_m_t_s_leq_w" $ + I.BasicOp $ I.CmpOp (I.CmpSlt Int64) i_p_m_t_s w + + zero_lte_i <- letSubExp "zero_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) zero i' + i_lte_j <- letSubExp "i_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) i' j' + forwards_ok <- + letSubExp "forwards_ok" + =<< eAll [zero_lte_i, zero_lte_i, i_lte_j, zero_leq_i_p_m_t_s, i_p_m_t_s_lth_w] + + negone_lte_j <- letSubExp "negone_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) negone j' + j_lte_i <- letSubExp "j_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int64) j' i' + backwards_ok <- + letSubExp "backwards_ok" + =<< eAll + [negone_lte_j, negone_lte_j, j_lte_i, zero_leq_i_p_m_t_s, i_p_m_t_s_leq_w] + + slice_ok <- + letSubExp "slice_ok" $ + I.If + backwards + (resultBody [backwards_ok]) + (resultBody [forwards_ok]) + $ ifCommon [I.Prim I.Bool] + ok_or_empty <- + letSubExp "ok_or_empty" $ + I.BasicOp $ I.BinOp I.LogOr empty_slice slice_ok + + let parts = case (i, j, s) of + (_, _, Just {}) -> + [ maybe "" (const $ ErrorInt64 i') i, + ":", + maybe "" (const $ ErrorInt64 j') j, + ":", + ErrorInt64 s' + ] + (_, Just {}, _) -> + [ maybe "" (const $ ErrorInt64 i') i, + ":", + ErrorInt64 j' + ] + ++ maybe mempty (const [":", ErrorInt64 s']) s + (_, Nothing, Nothing) -> + [ErrorInt64 i', ":"] + return (I.DimSlice i' n s', ok_or_empty, parts) + where + zero = constant (0 :: Int64) + negone = constant (-1 :: Int64) + one = constant (1 :: Int64) + +internaliseScanOrReduce :: + String -> + String -> + (SubExp -> I.Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)) -> + (E.Exp, E.Exp, E.Exp, SrcLoc) -> + InternaliseM [SubExp] +internaliseScanOrReduce desc what f (lam, ne, arr, loc) = do + arrs <- internaliseExpToVars (what ++ "_arr") arr + nes <- internaliseExp (what ++ "_ne") ne + nes' <- forM (zip nes arrs) $ \(ne', arr') -> do + rowtype <- I.stripArray 1 <$> lookupType arr' + ensureShape + "Row shape of input array does not match shape of neutral element" + loc + rowtype + (what ++ "_ne_right_shape") + ne' + nests <- mapM I.subExpType nes' + arrts <- mapM lookupType arrs + lam' <- internaliseFoldLambda internaliseLambda lam nests arrts + w <- arraysSize 0 <$> mapM lookupType arrs + letTupExp' desc . I.Op =<< f w lam' nes' arrs + +internaliseHist :: + String -> + E.Exp -> + E.Exp -> + E.Exp -> + E.Exp -> + E.Exp -> + E.Exp -> + SrcLoc -> + InternaliseM [SubExp] +internaliseHist desc rf hist op ne buckets img loc = do + rf' <- internaliseExp1 "hist_rf" rf + ne' <- internaliseExp "hist_ne" ne + hist' <- internaliseExpToVars "hist_hist" hist + buckets' <- + letExp "hist_buckets" . BasicOp . SubExp + =<< internaliseExp1 "hist_buckets" buckets + img' <- internaliseExpToVars "hist_img" img + + -- reshape neutral element to have same size as the destination array + ne_shp <- forM (zip ne' hist') $ \(n, h) -> do + rowtype <- I.stripArray 1 <$> lookupType h + ensureShape + "Row shape of destination array does not match shape of neutral element" + loc + rowtype + "hist_ne_right_shape" + n + ne_ts <- mapM I.subExpType ne_shp + his_ts <- mapM lookupType hist' + op' <- internaliseFoldLambda internaliseLambda op ne_ts his_ts + + -- reshape return type of bucket function to have same size as neutral element + -- (modulo the index) + bucket_param <- newParam "bucket_p" $ I.Prim int64 + img_params <- mapM (newParam "img_p" . rowType) =<< mapM lookupType img' + let params = bucket_param : img_params + rettype = I.Prim int64 : ne_ts + body = mkBody mempty $ map (I.Var . paramName) params + body' <- + localScope (scopeOfLParams params) $ + ensureResultShape + "Row shape of value array does not match row shape of hist target" + (srclocOf img) + rettype + body + + -- get sizes of histogram and image arrays + w_hist <- arraysSize 0 <$> mapM lookupType hist' + w_img <- arraysSize 0 <$> mapM lookupType img' + + -- Generate an assertion and reshapes to ensure that buckets' and + -- img' are the same size. + b_shape <- I.arrayShape <$> lookupType buckets' + let b_w = shapeSize 0 b_shape + cmp <- letSubExp "bucket_cmp" $ I.BasicOp $ I.CmpOp (I.CmpEq I.int64) b_w w_img + c <- + assert + "bucket_cert" + cmp + "length of index and value array does not match" + loc + buckets'' <- + certifying c $ + letExp (baseString buckets') $ + I.BasicOp $ I.Reshape (reshapeOuter [DimCoercion w_img] 1 b_shape) buckets' + + letTupExp' desc $ + I.Op $ + I.Hist w_img [HistOp w_hist rf' hist' ne_shp op'] (I.Lambda params body' rettype) $ buckets'' : img' + +internaliseStreamMap :: + String -> + StreamOrd -> + E.Exp -> + E.Exp -> + InternaliseM [SubExp] +internaliseStreamMap desc o lam arr = do + 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 + +internaliseStreamRed :: + String -> + StreamOrd -> + Commutativity -> + E.Exp -> + E.Exp -> + E.Exp -> + InternaliseM [SubExp] +internaliseStreamRed desc o comm lam0 lam arr = do + arrs <- internaliseExpToVars "stream_input" arr + rowts <- mapM (fmap I.rowType . lookupType) arrs + (lam_params, lam_body) <- + internaliseStreamLambda internaliseLambda lam rowts + let (chunk_param, _, lam_val_params) = + partitionChunkedFoldParameters 0 lam_params + + -- Synthesize neutral elements by applying the fold function + -- to an empty chunk. + letBindNames [I.paramName chunk_param] $ + I.BasicOp $ I.SubExp $ constant (0 :: Int64) + forM_ lam_val_params $ \p -> + letBindNames [I.paramName p] $ + I.BasicOp $ + I.Scratch (I.elemType $ I.paramType p) $ + I.arrayDims $ I.paramType p + nes <- bodyBind =<< renameBody lam_body + + nes_ts <- mapM I.subExpType nes + outsz <- arraysSize 0 <$> mapM lookupType arrs + let acc_arr_tps = [I.arrayOf t (I.Shape [outsz]) NoUniqueness | t <- nes_ts] + lam0' <- internaliseFoldLambda internaliseLambda lam0 nes_ts acc_arr_tps + + let lam0_acc_params = take (length nes) $ I.lambdaParams lam0' + lam_acc_params <- forM lam0_acc_params $ \p -> do + name <- newVName $ baseString $ I.paramName p + return p {I.paramName = name} + + -- Make sure the chunk size parameter comes first. + let lam_params' = chunk_param : lam_acc_params ++ lam_val_params + + body_with_lam0 <- + ensureResultShape + "shape of result does not match shape of initial value" + (srclocOf lam0) + nes_ts + <=< insertStmsM + $ localScope (scopeOfLParams lam_params') $ do + lam_res <- bodyBind lam_body + lam_res' <- + ensureArgShapes + "shape of chunk function result does not match shape of initial value" + (srclocOf lam) + [] + (map I.typeOf $ I.lambdaParams lam0') + lam_res + new_lam_res <- + eLambda lam0' $ + map eSubExp $ + map (I.Var . paramName) lam_acc_params ++ lam_res' + return $ resultBody new_lam_res + + let form = I.Parallel o comm lam0' nes + lam' = + I.Lambda + { lambdaParams = lam_params', + lambdaBody = body_with_lam0, + lambdaReturnType = nes_ts + } + w <- arraysSize 0 <$> mapM lookupType arrs + letTupExp' desc $ I.Op $ I.Stream w form lam' arrs + +internaliseExp1 :: String -> E.Exp -> InternaliseM I.SubExp +internaliseExp1 desc e = do + vs <- internaliseExp desc e + case vs of + [se] -> return se + _ -> error "Internalise.internaliseExp1: was passed not just a single subexpression" + +-- | Promote to dimension type as appropriate for the original type. +-- Also return original type. +internaliseDimExp :: String -> E.Exp -> InternaliseM (I.SubExp, IntType) +internaliseDimExp s e = do + e' <- internaliseExp1 s e + case E.typeOf e of + E.Scalar (E.Prim (Signed it)) -> (,it) <$> asIntS Int64 e' + _ -> error "internaliseDimExp: bad type" + +internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName] +internaliseExpToVars desc e = + mapM asIdent =<< internaliseExp desc e + where + asIdent (I.Var v) = return v + asIdent se = letExp desc $ I.BasicOp $ I.SubExp se + +internaliseOperation :: + String -> + E.Exp -> + (I.VName -> InternaliseM I.BasicOp) -> + InternaliseM [I.SubExp] +internaliseOperation s e op = do + vs <- internaliseExpToVars s e + letSubExps s =<< mapM (fmap I.BasicOp . op) vs + +certifyingNonzero :: + SrcLoc -> + IntType -> + SubExp -> + InternaliseM a -> + InternaliseM a +certifyingNonzero loc t x m = do + zero <- + letSubExp "zero" $ + I.BasicOp $ + CmpOp (CmpEq (IntType t)) x (intConst t 0) + nonzero <- letSubExp "nonzero" $ I.BasicOp $ UnOp Not zero + c <- assert "nonzero_cert" nonzero "division by zero" loc + certifying c m + +certifyingNonnegative :: + SrcLoc -> + IntType -> + SubExp -> + InternaliseM a -> + InternaliseM a +certifyingNonnegative loc t x m = do + nonnegative <- + letSubExp "nonnegative" $ + I.BasicOp $ + CmpOp (CmpSle t) (intConst t 0) x + c <- assert "nonzero_cert" nonnegative "negative exponent" loc + certifying c m + +internaliseBinOp :: + SrcLoc -> + String -> + E.BinOp -> + I.SubExp -> + I.SubExp -> + E.PrimType -> + E.PrimType -> + InternaliseM [I.SubExp] +internaliseBinOp _ desc E.Plus x y (E.Signed t) _ = + simpleBinOp desc (I.Add t I.OverflowWrap) x y +internaliseBinOp _ desc E.Plus x y (E.Unsigned t) _ = + simpleBinOp desc (I.Add t I.OverflowWrap) x y +internaliseBinOp _ desc E.Plus x y (E.FloatType t) _ = + simpleBinOp desc (I.FAdd t) x y +internaliseBinOp _ desc E.Minus x y (E.Signed t) _ = + simpleBinOp desc (I.Sub t I.OverflowWrap) x y +internaliseBinOp _ desc E.Minus x y (E.Unsigned t) _ = + simpleBinOp desc (I.Sub t I.OverflowWrap) x y +internaliseBinOp _ desc E.Minus x y (E.FloatType t) _ = + simpleBinOp desc (I.FSub t) x y +internaliseBinOp _ desc E.Times x y (E.Signed t) _ = + simpleBinOp desc (I.Mul t I.OverflowWrap) x y +internaliseBinOp _ desc E.Times x y (E.Unsigned t) _ = + simpleBinOp desc (I.Mul t I.OverflowWrap) x y +internaliseBinOp _ desc E.Times x y (E.FloatType t) _ = + simpleBinOp desc (I.FMul t) x y +internaliseBinOp loc desc E.Divide x y (E.Signed t) _ = + certifyingNonzero loc t y $ + simpleBinOp desc (I.SDiv t I.Unsafe) x y +internaliseBinOp loc desc E.Divide x y (E.Unsigned t) _ = + certifyingNonzero loc t y $ + simpleBinOp desc (I.UDiv t I.Unsafe) x y +internaliseBinOp _ desc E.Divide x y (E.FloatType t) _ = + simpleBinOp desc (I.FDiv t) x y +internaliseBinOp _ desc E.Pow x y (E.FloatType t) _ = + simpleBinOp desc (I.FPow t) x y +internaliseBinOp loc desc E.Pow x y (E.Signed t) _ = + certifyingNonnegative loc t y $ + simpleBinOp desc (I.Pow t) x y +internaliseBinOp _ desc E.Pow x y (E.Unsigned t) _ = + simpleBinOp desc (I.Pow t) x y +internaliseBinOp loc desc E.Mod x y (E.Signed t) _ = + certifyingNonzero loc t y $ + simpleBinOp desc (I.SMod t I.Unsafe) x y +internaliseBinOp loc desc E.Mod x y (E.Unsigned t) _ = + certifyingNonzero loc t y $ + simpleBinOp desc (I.UMod t I.Unsafe) x y +internaliseBinOp _ desc E.Mod x y (E.FloatType t) _ = + simpleBinOp desc (I.FMod t) x y +internaliseBinOp loc desc E.Quot x y (E.Signed t) _ = + certifyingNonzero loc t y $ + simpleBinOp desc (I.SQuot t I.Unsafe) x y +internaliseBinOp loc desc E.Quot x y (E.Unsigned t) _ = + certifyingNonzero loc t y $ + simpleBinOp desc (I.UDiv t I.Unsafe) x y +internaliseBinOp loc desc E.Rem x y (E.Signed t) _ = + certifyingNonzero loc t y $ + simpleBinOp desc (I.SRem t I.Unsafe) x y +internaliseBinOp loc desc E.Rem x y (E.Unsigned t) _ = + certifyingNonzero loc t y $ + simpleBinOp desc (I.UMod t I.Unsafe) x y +internaliseBinOp _ desc E.ShiftR x y (E.Signed t) _ = + simpleBinOp desc (I.AShr t) x y +internaliseBinOp _ desc E.ShiftR x y (E.Unsigned t) _ = + simpleBinOp desc (I.LShr t) x y +internaliseBinOp _ desc E.ShiftL x y (E.Signed t) _ = + simpleBinOp desc (I.Shl t) x y +internaliseBinOp _ desc E.ShiftL x y (E.Unsigned t) _ = + simpleBinOp desc (I.Shl t) x y +internaliseBinOp _ desc E.Band x y (E.Signed t) _ = + simpleBinOp desc (I.And t) x y +internaliseBinOp _ desc E.Band x y (E.Unsigned t) _ = + simpleBinOp desc (I.And t) x y +internaliseBinOp _ desc E.Xor x y (E.Signed t) _ = + simpleBinOp desc (I.Xor t) x y +internaliseBinOp _ desc E.Xor x y (E.Unsigned t) _ = + simpleBinOp desc (I.Xor t) x y +internaliseBinOp _ desc E.Bor x y (E.Signed t) _ = + simpleBinOp desc (I.Or t) x y +internaliseBinOp _ desc E.Bor x y (E.Unsigned t) _ = + simpleBinOp desc (I.Or t) x y +internaliseBinOp _ desc E.Equal x y t _ = + simpleCmpOp desc (I.CmpEq $ internalisePrimType t) x y +internaliseBinOp _ desc E.NotEqual x y t _ = do + eq <- letSubExp (desc ++ "true") $ I.BasicOp $ I.CmpOp (I.CmpEq $ internalisePrimType t) x y + fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp I.Not eq +internaliseBinOp _ desc E.Less x y (E.Signed t) _ = + simpleCmpOp desc (I.CmpSlt t) x y +internaliseBinOp _ desc E.Less x y (E.Unsigned t) _ = + simpleCmpOp desc (I.CmpUlt t) x y +internaliseBinOp _ desc E.Leq x y (E.Signed t) _ = + simpleCmpOp desc (I.CmpSle t) x y +internaliseBinOp _ desc E.Leq x y (E.Unsigned t) _ = + simpleCmpOp desc (I.CmpUle t) x y +internaliseBinOp _ desc E.Greater x y (E.Signed t) _ = + simpleCmpOp desc (I.CmpSlt t) y x -- Note the swapped x and y +internaliseBinOp _ desc E.Greater x y (E.Unsigned t) _ = + simpleCmpOp desc (I.CmpUlt t) y x -- Note the swapped x and y +internaliseBinOp _ desc E.Geq x y (E.Signed t) _ = + simpleCmpOp desc (I.CmpSle t) y x -- Note the swapped x and y +internaliseBinOp _ desc E.Geq x y (E.Unsigned t) _ = + simpleCmpOp desc (I.CmpUle t) y x -- Note the swapped x and y +internaliseBinOp _ desc E.Less x y (E.FloatType t) _ = + simpleCmpOp desc (I.FCmpLt t) x y +internaliseBinOp _ desc E.Leq x y (E.FloatType t) _ = + simpleCmpOp desc (I.FCmpLe t) x y +internaliseBinOp _ desc E.Greater x y (E.FloatType t) _ = + simpleCmpOp desc (I.FCmpLt t) y x -- Note the swapped x and y +internaliseBinOp _ desc E.Geq x y (E.FloatType t) _ = + simpleCmpOp desc (I.FCmpLe t) y x -- Note the swapped x and y + +-- Relational operators for booleans. +internaliseBinOp _ desc E.Less x y E.Bool _ = + simpleCmpOp desc I.CmpLlt x y +internaliseBinOp _ desc E.Leq x y E.Bool _ = + simpleCmpOp desc I.CmpLle x y +internaliseBinOp _ desc E.Greater x y E.Bool _ = + simpleCmpOp desc I.CmpLlt y x -- Note the swapped x and y +internaliseBinOp _ desc E.Geq x y E.Bool _ = + simpleCmpOp desc I.CmpLle y x -- Note the swapped x and y +internaliseBinOp _ _ op _ _ t1 t2 = + error $ + "Invalid binary operator " ++ pretty op + ++ " with operand types " + ++ pretty t1 + ++ ", " + ++ pretty t2 + +simpleBinOp :: + String -> + I.BinOp -> + I.SubExp -> + I.SubExp -> + InternaliseM [I.SubExp] +simpleBinOp desc bop x y = + letTupExp' desc $ I.BasicOp $ I.BinOp bop x y + +simpleCmpOp :: + String -> + I.CmpOp -> + I.SubExp -> + I.SubExp -> + InternaliseM [I.SubExp] +simpleCmpOp desc op x y = + letTupExp' desc $ I.BasicOp $ I.CmpOp op x y + +findFuncall :: + E.Exp -> + InternaliseM + ( E.QualName VName, + [(E.Exp, Maybe VName)], + E.StructType, + [VName] + ) +findFuncall (E.Var fname (Info t) _) = + return (fname, [], E.toStruct t, []) +findFuncall (E.Apply f arg (Info (_, argext)) (Info ret, Info retext) _) = do + (fname, args, _, _) <- findFuncall f + return (fname, args ++ [(arg, argext)], E.toStruct ret, retext) +findFuncall e = + error $ "Invalid function expression in application: " ++ pretty e + +internaliseLambda :: InternaliseLambda +internaliseLambda (E.Parens e _) rowtypes = + internaliseLambda e rowtypes +internaliseLambda (E.Lambda params body _ (Info (_, rettype)) _) rowtypes = + bindingLambdaParams params rowtypes $ \params' -> do + body' <- internaliseBody body + rettype' <- internaliseLambdaReturnType rettype + return (params', body', rettype') +internaliseLambda e _ = error $ "internaliseLambda: unexpected expression:\n" ++ pretty e + +-- | Some operators and functions are overloaded or otherwise special +-- - we detect and treat them here. +isOverloadedFunction :: + E.QualName VName -> + [E.Exp] -> + SrcLoc -> + Maybe (String -> InternaliseM [SubExp]) +isOverloadedFunction qname args loc = do + guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag + let handlers = + [ handleSign, + handleIntrinsicOps, + handleOps, + handleSOACs, + handleRest + ] + msum [h args $ baseString $ qualLeaf qname | h <- handlers] + where + handleSign [x] "sign_i8" = Just $ toSigned I.Int8 x + handleSign [x] "sign_i16" = Just $ toSigned I.Int16 x + handleSign [x] "sign_i32" = Just $ toSigned I.Int32 x + handleSign [x] "sign_i64" = Just $ toSigned I.Int64 x + handleSign [x] "unsign_i8" = Just $ toUnsigned I.Int8 x + handleSign [x] "unsign_i16" = Just $ toUnsigned I.Int16 x + handleSign [x] "unsign_i32" = Just $ toUnsigned I.Int32 x + handleSign [x] "unsign_i64" = Just $ toUnsigned I.Int64 x + handleSign _ _ = Nothing + + handleIntrinsicOps [x] s + | Just unop <- find ((== s) . pretty) allUnOps = Just $ \desc -> do + x' <- internaliseExp1 "x" x + fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop x' + handleIntrinsicOps [TupLit [x, y] _] s + | Just bop <- find ((== s) . pretty) allBinOps = Just $ \desc -> do + x' <- internaliseExp1 "x" x + y' <- internaliseExp1 "y" y + fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y' + | Just cmp <- find ((== s) . pretty) allCmpOps = Just $ \desc -> do + x' <- internaliseExp1 "x" x + y' <- internaliseExp1 "y" y + fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' y' + handleIntrinsicOps [x] s + | Just conv <- find ((== s) . pretty) allConvOps = Just $ \desc -> do + x' <- internaliseExp1 "x" x + fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x' + handleIntrinsicOps _ _ = Nothing + + -- Short-circuiting operators are magical. + handleOps [x, y] "&&" = Just $ \desc -> + internaliseExp desc $ + E.If x y (E.Literal (E.BoolValue False) mempty) (Info $ E.Scalar $ E.Prim E.Bool, Info []) mempty + handleOps [x, y] "||" = Just $ \desc -> + internaliseExp desc $ + E.If x (E.Literal (E.BoolValue True) mempty) y (Info $ E.Scalar $ E.Prim E.Bool, Info []) mempty + -- Handle equality and inequality specially, to treat the case of + -- arrays. + handleOps [xe, ye] op + | Just cmp_f <- isEqlOp op = Just $ \desc -> do + xe' <- internaliseExp "x" xe + ye' <- internaliseExp "y" ye + rs <- zipWithM (doComparison desc) xe' ye' + cmp_f desc =<< letSubExp "eq" =<< eAll rs + where + isEqlOp "!=" = Just $ \desc eq -> + letTupExp' desc $ I.BasicOp $ I.UnOp I.Not eq + isEqlOp "==" = Just $ \_ eq -> + return [eq] + isEqlOp _ = Nothing + + doComparison desc x y = do + x_t <- I.subExpType x + y_t <- I.subExpType y + case x_t of + I.Prim t -> letSubExp desc $ I.BasicOp $ I.CmpOp (I.CmpEq t) x y + _ -> do + let x_dims = I.arrayDims x_t + y_dims = I.arrayDims y_t + dims_match <- forM (zip x_dims y_dims) $ \(x_dim, y_dim) -> + letSubExp "dim_eq" $ I.BasicOp $ I.CmpOp (I.CmpEq int64) x_dim y_dim + shapes_match <- letSubExp "shapes_match" =<< eAll dims_match + compare_elems_body <- runBodyBinder $ do + -- Flatten both x and y. + x_num_elems <- + letSubExp "x_num_elems" + =<< foldBinOp (I.Mul Int64 I.OverflowUndef) (constant (1 :: Int64)) x_dims + x' <- letExp "x" $ I.BasicOp $ I.SubExp x + y' <- letExp "x" $ I.BasicOp $ I.SubExp y + x_flat <- letExp "x_flat" $ I.BasicOp $ I.Reshape [I.DimNew x_num_elems] x' + y_flat <- letExp "y_flat" $ I.BasicOp $ I.Reshape [I.DimNew x_num_elems] y' + + -- Compare the elements. + cmp_lam <- cmpOpLambda $ I.CmpEq (elemType x_t) + cmps <- + letExp "cmps" $ + I.Op $ + I.Screma x_num_elems (I.mapSOAC cmp_lam) [x_flat, y_flat] + + -- Check that all were equal. + and_lam <- binOpLambda I.LogAnd I.Bool + reduce <- I.reduceSOAC [Reduce Commutative and_lam [constant True]] + all_equal <- letSubExp "all_equal" $ I.Op $ I.Screma x_num_elems reduce [cmps] + return $ resultBody [all_equal] + + letSubExp "arrays_equal" $ + I.If shapes_match compare_elems_body (resultBody [constant False]) $ + ifCommon [I.Prim I.Bool] + handleOps [x, y] name + | Just bop <- find ((name ==) . pretty) [minBound .. maxBound :: E.BinOp] = + Just $ \desc -> do + x' <- internaliseExp1 "x" x + y' <- internaliseExp1 "y" y + case (E.typeOf x, E.typeOf y) of + (E.Scalar (E.Prim t1), E.Scalar (E.Prim t2)) -> + internaliseBinOp loc desc bop x' y' t1 t2 + _ -> error "Futhark.Internalise.internaliseExp: non-primitive type in BinOp." + handleOps _ _ = Nothing + + handleSOACs [TupLit [lam, arr] _] "map" = Just $ \desc -> do + arr' <- internaliseExpToVars "map_arr" arr + lam' <- internaliseMapLambda internaliseLambda lam $ map I.Var arr' + w <- arraysSize 0 <$> mapM lookupType arr' + letTupExp' desc $ + I.Op $ + I.Screma w (I.mapSOAC lam') arr' + handleSOACs [TupLit [k, lam, arr] _] "partition" = do + k' <- fromIntegral <$> fromInt32 k + Just $ \_desc -> do + arrs <- internaliseExpToVars "partition_input" arr + lam' <- internalisePartitionLambda internaliseLambda k' lam $ map I.Var arrs + uncurry (++) <$> partitionWithSOACS (fromIntegral k') lam' arrs + where + fromInt32 (Literal (SignedValue (Int32Value k')) _) = Just k' + fromInt32 (IntLit k' (Info (E.Scalar (E.Prim (Signed Int32)))) _) = Just $ fromInteger k' + fromInt32 _ = Nothing + handleSOACs [TupLit [lam, ne, arr] _] "reduce" = Just $ \desc -> + internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr, loc) + where + reduce w red_lam nes arrs = + I.Screma w + <$> I.reduceSOAC [Reduce Noncommutative red_lam nes] <*> pure arrs + handleSOACs [TupLit [lam, ne, arr] _] "reduce_comm" = Just $ \desc -> + internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr, loc) + where + reduce w red_lam nes arrs = + I.Screma w + <$> I.reduceSOAC [Reduce Commutative red_lam nes] <*> pure arrs + handleSOACs [TupLit [lam, ne, arr] _] "scan" = Just $ \desc -> + internaliseScanOrReduce desc "scan" reduce (lam, ne, arr, loc) + where + reduce w scan_lam nes arrs = + I.Screma w <$> I.scanSOAC [Scan scan_lam nes] <*> pure arrs + handleSOACs [TupLit [op, f, arr] _] "reduce_stream" = Just $ \desc -> + internaliseStreamRed desc InOrder Noncommutative op f arr + handleSOACs [TupLit [op, f, arr] _] "reduce_stream_per" = Just $ \desc -> + internaliseStreamRed desc Disorder Commutative op f arr + handleSOACs [TupLit [f, arr] _] "map_stream" = Just $ \desc -> + internaliseStreamMap desc InOrder f arr + handleSOACs [TupLit [f, arr] _] "map_stream_per" = Just $ \desc -> + internaliseStreamMap desc Disorder f arr + handleSOACs [TupLit [rf, dest, op, ne, buckets, img] _] "hist" = Just $ \desc -> + internaliseHist desc rf dest op ne buckets img loc + handleSOACs _ _ = Nothing + + 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 [n, m, arr] _] "unflatten" = Just $ \desc -> do + arrs <- internaliseExpToVars "unflatten_arr" arr + n' <- internaliseExp1 "n" n + m' <- internaliseExp1 "m" m + -- The unflattened dimension needs to have the same number of elements + -- as the original dimension. + old_dim <- I.arraysSize 0 <$> mapM lookupType arrs + dim_ok <- + letSubExp "dim_ok" + =<< eCmpOp + (I.CmpEq I.int64) + (eBinOp (I.Mul Int64 I.OverflowUndef) (eSubExp n') (eSubExp m')) + (eSubExp old_dim) + dim_ok_cert <- + assert + "dim_ok_cert" + dim_ok + "new shape has different number of elements than old shape" + loc + certifying dim_ok_cert $ + forM arrs $ \arr' -> do + arr_t <- lookupType arr' + letSubExp desc $ + I.BasicOp $ + I.Reshape (reshapeOuter [DimNew n', DimNew m'] 1 $ I.arrayShape arr_t) arr' + handleRest [arr] "flatten" = Just $ \desc -> do + arrs <- internaliseExpToVars "flatten_arr" arr + forM arrs $ \arr' -> do + arr_t <- lookupType arr' + let n = arraySize 0 arr_t + m = arraySize 1 arr_t + k <- letSubExp "flat_dim" $ I.BasicOp $ I.BinOp (Mul Int64 I.OverflowUndef) n m + letSubExp desc $ + I.BasicOp $ + I.Reshape (reshapeOuter [DimNew k] 2 $ I.arrayShape arr_t) arr' + handleRest [TupLit [x, y] _] "concat" = Just $ \desc -> do + xs <- internaliseExpToVars "concat_x" x + ys <- internaliseExpToVars "concat_y" y + outer_size <- arraysSize 0 <$> mapM lookupType xs + let sumdims xsize ysize = + letSubExp "conc_tmp" $ + I.BasicOp $ + I.BinOp (I.Add I.Int64 I.OverflowUndef) xsize ysize + ressize <- + foldM sumdims outer_size + =<< mapM (fmap (arraysSize 0) . mapM lookupType) [ys] + + let conc xarr yarr = + I.BasicOp $ I.Concat 0 xarr [yarr] ressize + letSubExps desc $ zipWith conc xs ys + handleRest [TupLit [offset, e] _] "rotate" = Just $ \desc -> do + offset' <- internaliseExp1 "rotation_offset" offset + internaliseOperation desc e $ \v -> do + r <- I.arrayRank <$> lookupType v + let zero = intConst Int64 0 + offsets = offset' : replicate (r -1) zero + return $ I.Rotate offsets v + handleRest [e] "transpose" = Just $ \desc -> + internaliseOperation desc e $ \v -> do + r <- I.arrayRank <$> lookupType v + return $ I.Rearrange ([1, 0] ++ [2 .. r -1]) v + handleRest [TupLit [x, y] _] "zip" = Just $ \desc -> + (++) <$> internaliseExp (desc ++ "_zip_x") x + <*> internaliseExp (desc ++ "_zip_y") y + handleRest [x] "unzip" = Just $ flip internaliseExp x + handleRest [x] "trace" = Just $ flip internaliseExp x + handleRest [x] "break" = Just $ flip internaliseExp x + handleRest _ _ = Nothing + + toSigned int_to e desc = do + e' <- internaliseExp1 "trunc_arg" e + case E.typeOf e of + E.Scalar (E.Prim E.Bool) -> + letTupExp' desc $ + I.If + e' + (resultBody [intConst int_to 1]) + (resultBody [intConst int_to 0]) + $ ifCommon [I.Prim $ I.IntType int_to] + E.Scalar (E.Prim (E.Signed int_from)) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.SExt int_from int_to) e' + E.Scalar (E.Prim (E.Unsigned int_from)) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' + E.Scalar (E.Prim (E.FloatType float_from)) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.FPToSI float_from int_to) e' + _ -> error "Futhark.Internalise: non-numeric type in ToSigned" + + toUnsigned int_to e desc = do + e' <- internaliseExp1 "trunc_arg" e + case E.typeOf e of + E.Scalar (E.Prim E.Bool) -> + letTupExp' desc $ + I.If + e' + (resultBody [intConst int_to 1]) + (resultBody [intConst int_to 0]) + $ ifCommon [I.Prim $ I.IntType int_to] + E.Scalar (E.Prim (E.Signed int_from)) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' + E.Scalar (E.Prim (E.Unsigned int_from)) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' + E.Scalar (E.Prim (E.FloatType float_from)) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.FPToUI float_from int_to) e' + _ -> error "Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned" + + complementF e desc = do + e' <- internaliseExp1 "complement_arg" e + et <- subExpType e' + case et of + I.Prim (I.IntType t) -> + letTupExp' desc $ I.BasicOp $ I.UnOp (I.Complement t) e' + I.Prim I.Bool -> + letTupExp' desc $ I.BasicOp $ I.UnOp I.Not e' + _ -> + 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 + 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 + sv_ts <- mapM lookupType svs + + svs' <- forM (zip svs sv_ts) $ \(sv, sv_t) -> do + let sv_shape = I.arrayShape sv_t + sv_w = arraySize 0 sv_t + + -- Generate an assertion and reshapes to ensure that sv and si' are the same + -- size. + cmp <- + letSubExp "write_cmp" $ + I.BasicOp $ + I.CmpOp (I.CmpEq I.int64) si_w sv_w + c <- + assert + "write_cert" + cmp + "length of index and value array does not match" + loc + certifying c $ + 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" + 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 + 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 + results <- forM outs $ \name -> + letSubExp "write_res" $ I.BasicOp $ I.SubExp $ I.Var name + ensureResultShape + "scatter value has wrong size" + loc + bodyTypes + $ resultBody results + + let lam = + I.Lambda + { I.lambdaParams = bodyParams, + I.lambdaReturnType = bodyTypes, + I.lambdaBody = body + } + sivs = si' : svs' + + let sa_ws = map (arraySize 0) sa_ts + letTupExp' desc $ I.Op $ I.Scatter si_w lam sivs $ zip3 sa_ws (repeat 1) sas + +funcall :: + String -> + QualName VName -> + [SubExp] -> + SrcLoc -> + InternaliseM ([SubExp], [I.ExtType]) +funcall desc (QualName _ fname) args loc = do + (fname', closure, shapes, value_paramts, fun_params, rettype_fun) <- + lookupFunction fname + argts <- mapM subExpType args + + shapeargs <- argShapes shapes fun_params argts + let diets = + replicate (length closure + length shapeargs) I.ObservePrim + ++ map I.diet value_paramts + args' <- + ensureArgShapes + "function arguments of wrong shape" + loc + (map I.paramName fun_params) + (map I.paramType fun_params) + (map I.Var closure ++ shapeargs ++ args) + argts' <- mapM subExpType args' + case rettype_fun $ zip args' argts' of + Nothing -> + error $ + "Cannot apply " ++ pretty fname ++ " to arguments\n " + ++ pretty args' + ++ "\nof types\n " + ++ pretty argts' + ++ "\nFunction has parameters\n " + ++ pretty fun_params + Just ts -> do + safety <- askSafety + attrs <- asks envAttrs + ses <- + attributing attrs $ + letTupExp' desc $ + I.Apply fname' (zip args' diets) ts (safety, loc, mempty) + return (ses, map I.fromDecl ts) + +-- Bind existential names defined by an expression, based on the +-- concrete values that expression evaluated to. This most +-- importantly should be done after function calls, but also +-- everything else that can produce existentials in the source +-- language. +bindExtSizes :: E.StructType -> [VName] -> [SubExp] -> InternaliseM () +bindExtSizes ret retext ses = do + ts <- internaliseType ret + ses_ts <- mapM subExpType ses + + let combine t1 t2 = + mconcat $ zipWith combine' (arrayExtDims t1) (arrayDims t2) + combine' (I.Free (I.Var v)) se + | v `elem` retext = M.singleton v se + combine' _ _ = mempty + + forM_ (M.toList $ mconcat $ zipWith combine ts ses_ts) $ \(v, se) -> + letBindNames [v] $ BasicOp $ SubExp se + +askSafety :: InternaliseM Safety +askSafety = do + check <- asks envDoBoundsChecks + return $ if check then I.Safe else I.Unsafe + +-- Implement partitioning using maps, scans and writes. +partitionWithSOACS :: Int -> I.Lambda -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp]) +partitionWithSOACS k lam arrs = do + arr_ts <- mapM lookupType arrs + let w = arraysSize 0 arr_ts + classes_and_increments <- letTupExp "increments" $ I.Op $ I.Screma w (mapSOAC lam) arrs + (classes, increments) <- case classes_and_increments of + classes : increments -> return (classes, take k increments) + _ -> error "partitionWithSOACS" + + add_lam_x_params <- + replicateM k $ I.Param <$> newVName "x" <*> pure (I.Prim int64) + add_lam_y_params <- + replicateM k $ I.Param <$> newVName "y" <*> pure (I.Prim int64) + add_lam_body <- runBodyBinder $ + localScope (scopeOfLParams $ add_lam_x_params ++ add_lam_y_params) $ + fmap resultBody $ + forM (zip add_lam_x_params add_lam_y_params) $ \(x, y) -> + letSubExp "z" $ + I.BasicOp $ + I.BinOp + (I.Add Int64 I.OverflowUndef) + (I.Var $ I.paramName x) + (I.Var $ I.paramName y) + let add_lam = + I.Lambda + { I.lambdaBody = add_lam_body, + I.lambdaParams = add_lam_x_params ++ add_lam_y_params, + I.lambdaReturnType = replicate k $ I.Prim int64 + } + nes = replicate (length increments) $ intConst Int64 0 + + scan <- I.scanSOAC [I.Scan add_lam nes] + all_offsets <- letTupExp "offsets" $ I.Op $ I.Screma w scan increments + + -- We have the offsets for each of the partitions, but we also need + -- the total sizes, which are the last elements in the offests. We + -- just have to be careful in case the array is empty. + last_index <- letSubExp "last_index" $ I.BasicOp $ I.BinOp (I.Sub Int64 OverflowUndef) w $ constant (1 :: Int64) + nonempty_body <- runBodyBinder $ + fmap resultBody $ + forM all_offsets $ \offset_array -> + letSubExp "last_offset" $ I.BasicOp $ I.Index offset_array [I.DimFix last_index] + let empty_body = resultBody $ replicate k $ constant (0 :: Int64) + is_empty <- letSubExp "is_empty" $ I.BasicOp $ I.CmpOp (CmpEq int64) w $ constant (0 :: Int64) + sizes <- + letTupExp "partition_size" $ + I.If is_empty empty_body nonempty_body $ + ifCommon $ replicate k $ I.Prim int64 + + -- The total size of all partitions must necessarily be equal to the + -- size of the input array. + + -- Create scratch arrays for the result. + blanks <- forM arr_ts $ \arr_t -> + letExp "partition_dest" $ + I.BasicOp $ + Scratch (elemType arr_t) (w : drop 1 (I.arrayDims arr_t)) + + -- Now write into the result. + write_lam <- do + c_param <- I.Param <$> newVName "c" <*> pure (I.Prim int64) + offset_params <- replicateM k $ I.Param <$> newVName "offset" <*> pure (I.Prim int64) + value_params <- forM arr_ts $ \arr_t -> + I.Param <$> newVName "v" <*> pure (I.rowType arr_t) + (offset, offset_stms) <- + collectStms $ + mkOffsetLambdaBody + (map I.Var sizes) + (I.Var $ I.paramName c_param) + 0 + offset_params + return + I.Lambda + { I.lambdaParams = c_param : offset_params ++ value_params, + I.lambdaReturnType = + replicate (length arr_ts) (I.Prim int64) + ++ map I.rowType arr_ts, + I.lambdaBody = + mkBody offset_stms $ + replicate (length arr_ts) offset + ++ map (I.Var . I.paramName) value_params + } + results <- + letTupExp "partition_res" $ + I.Op $ + I.Scatter + w + write_lam + (classes : all_offsets ++ arrs) + $ zip3 (repeat w) (repeat 1) blanks + sizes' <- + letSubExp "partition_sizes" $ + I.BasicOp $ + I.ArrayLit (map I.Var sizes) $ I.Prim int64 + return (map I.Var results, [sizes']) + where + mkOffsetLambdaBody :: + [SubExp] -> + SubExp -> + Int -> + [I.LParam] -> + InternaliseM SubExp + mkOffsetLambdaBody _ _ _ [] = + return $ constant (-1 :: Int64) + mkOffsetLambdaBody sizes c i (p : ps) = do + is_this_one <- + letSubExp "is_this_one" $ + I.BasicOp $ + I.CmpOp (CmpEq int64) c $ + intConst Int64 $ toInteger i + next_one <- mkOffsetLambdaBody sizes c (i + 1) ps + this_one <- + letSubExp "this_offset" + =<< foldBinOp + (Add Int64 OverflowUndef) + (constant (-1 :: Int64)) + (I.Var (I.paramName p) : take i sizes) + letSubExp "total_res" $ + I.If + is_this_one + (resultBody [this_one]) + (resultBody [next_one]) + $ ifCommon [I.Prim int64] + +typeExpForError :: E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp] +typeExpForError (E.TEVar qn _) = + return [ErrorString $ pretty qn] +typeExpForError (E.TEUnique te _) = + ("*" :) <$> typeExpForError te +typeExpForError (E.TEArray te d _) = do + d' <- dimExpForError d + te' <- typeExpForError te + return $ ["[", d', "]"] ++ te' +typeExpForError (E.TETuple tes _) = do + tes' <- mapM typeExpForError tes + return $ ["("] ++ intercalate [", "] tes' ++ [")"] +typeExpForError (E.TERecord fields _) = do + fields' <- mapM onField fields + return $ ["{"] ++ intercalate [", "] fields' ++ ["}"] + where + onField (k, te) = + (ErrorString (pretty k ++ ": ") :) <$> typeExpForError te +typeExpForError (E.TEArrow _ t1 t2 _) = do + t1' <- typeExpForError t1 + t2' <- typeExpForError t2 + return $ t1' ++ [" -> "] ++ t2' +typeExpForError (E.TEApply t arg _) = do + t' <- typeExpForError t + arg' <- case arg of + TypeArgExpType argt -> typeExpForError argt + TypeArgExpDim d _ -> pure <$> dimExpForError d + return $ t' ++ [" "] ++ arg' +typeExpForError (E.TESum cs _) = do + cs' <- mapM (onClause . snd) cs + return $ intercalate [" | "] cs' + where + onClause c = do + c' <- mapM typeExpForError c + return $ intercalate [" "] c' + +dimExpForError :: E.DimExp VName -> InternaliseM (ErrorMsgPart SubExp) +dimExpForError (DimExpNamed d _) = do + substs <- lookupSubst $ E.qualLeaf d + d' <- case substs of + Just [v] -> return v + _ -> return $ I.Var $ E.qualLeaf d + return $ ErrorInt64 d' +dimExpForError (DimExpConst d _) = + return $ ErrorString $ pretty d +dimExpForError DimExpAny = return "" + +-- A smart constructor that compacts neighbouring literals for easier +-- reading in the IR. +errorMsg :: [ErrorMsgPart a] -> ErrorMsg a +errorMsg = ErrorMsg . compact + where + compact [] = [] + compact (ErrorString x : ErrorString y : parts) = + compact (ErrorString (x ++ y) : parts) + compact (x : y) = x : compact y diff --git a/src/Futhark/Internalise/AccurateSizes.hs b/src/Futhark/Internalise/AccurateSizes.hs --- a/src/Futhark/Internalise/AccurateSizes.hs +++ b/src/Futhark/Internalise/AccurateSizes.hs @@ -1,67 +1,87 @@ {-# LANGUAGE FlexibleContexts #-} + module Futhark.Internalise.AccurateSizes - ( argShapes - , ensureResultShape - , ensureResultExtShape - , ensureExtShape - , ensureShape - , ensureArgShapes + ( argShapes, + ensureResultShape, + ensureResultExtShape, + ensureExtShape, + ensureShape, + ensureArgShapes, ) - where +where import Control.Monad -import Data.Maybe import qualified Data.Map.Strict as M - +import Data.Maybe import Futhark.Construct -import Futhark.Internalise.Monad import Futhark.IR.SOACS +import Futhark.Internalise.Monad import Futhark.Util (takeLast) -shapeMapping :: HasScope SOACS m => - [FParam] -> [Type] - -> m (M.Map VName SubExp) +shapeMapping :: + HasScope SOACS m => + [FParam] -> + [Type] -> + m (M.Map VName SubExp) shapeMapping all_params value_arg_types = mconcat <$> zipWithM f value_params value_arg_types - where value_params = takeLast (length value_arg_types) all_params + where + value_params = takeLast (length value_arg_types) all_params - f (Param _ t1@Array{}) t2@Array{} = - pure $ M.fromList $ mapMaybe match $ zip (arrayDims t1) (arrayDims t2) - f _ _ = - pure mempty + f (Param _ t1@Array {}) t2@Array {} = + pure $ M.fromList $ mapMaybe match $ zip (arrayDims t1) (arrayDims t2) + f _ _ = + pure mempty - match (Var v, se) = Just (v, se) - match _ = Nothing + match (Var v, se) = Just (v, se) + match _ = Nothing -argShapes :: (HasScope SOACS m, Monad m) => - [VName] -> [FParam] -> [Type] -> m [SubExp] +argShapes :: + (HasScope SOACS m, Monad m) => + [VName] -> + [FParam] -> + [Type] -> + m [SubExp] argShapes shapes all_params valargts = do mapping <- shapeMapping all_params valargts let addShape name = case M.lookup name mapping of Just se -> se - _ -> intConst Int32 0 -- FIXME: we only need this because - -- the defunctionaliser throws away - -- sizes. + _ -> intConst Int64 0 -- FIXME: we only need this because + -- the defunctionaliser throws away + -- sizes. return $ map addShape shapes -ensureResultShape :: ErrorMsg SubExp -> SrcLoc -> [Type] -> Body - -> InternaliseM Body +ensureResultShape :: + ErrorMsg SubExp -> + SrcLoc -> + [Type] -> + Body -> + InternaliseM Body ensureResultShape msg loc = ensureResultExtShape msg loc . staticShapes -ensureResultExtShape :: ErrorMsg SubExp -> SrcLoc -> [ExtType] -> Body - -> InternaliseM Body +ensureResultExtShape :: + ErrorMsg SubExp -> + SrcLoc -> + [ExtType] -> + Body -> + InternaliseM Body ensureResultExtShape msg loc rettype body = insertStmsM $ do - reses <- bodyBind =<< - ensureResultExtShapeNoCtx msg loc rettype body + reses <- + bodyBind + =<< ensureResultExtShapeNoCtx msg loc rettype body ts <- mapM subExpType reses let ctx = extractShapeContext rettype $ map arrayDims ts mkBodyM mempty $ ctx ++ reses -ensureResultExtShapeNoCtx :: ErrorMsg SubExp -> SrcLoc -> [ExtType] -> Body - -> InternaliseM Body +ensureResultExtShapeNoCtx :: + ErrorMsg SubExp -> + SrcLoc -> + [ExtType] -> + Body -> + InternaliseM Body ensureResultExtShapeNoCtx msg loc rettype body = insertStmsM $ do es <- bodyBind body @@ -70,38 +90,60 @@ rettype' = foldr (uncurry fixExt) rettype $ M.toList ext_mapping assertProperShape t se = let name = "result_proper_shape" - in ensureExtShape msg loc t name se + in ensureExtShape msg loc t name se resultBodyM =<< zipWithM assertProperShape rettype' es -ensureExtShape :: ErrorMsg SubExp -> SrcLoc -> ExtType -> String -> SubExp - -> InternaliseM SubExp +ensureExtShape :: + ErrorMsg SubExp -> + SrcLoc -> + ExtType -> + String -> + SubExp -> + InternaliseM SubExp ensureExtShape msg loc t name orig - | Array{} <- t, Var v <- orig = + | Array {} <- t, + Var v <- orig = Var <$> ensureShapeVar msg loc t name v | otherwise = return orig -ensureShape :: ErrorMsg SubExp -> SrcLoc -> Type -> String -> SubExp - -> InternaliseM SubExp +ensureShape :: + ErrorMsg SubExp -> + SrcLoc -> + Type -> + String -> + SubExp -> + InternaliseM SubExp ensureShape msg loc = ensureExtShape msg loc . staticShapes1 -- | Reshape the arguments to a function so that they fit the expected -- shape declarations. Not used to change rank of arguments. Assumes -- everything is otherwise type-correct. -ensureArgShapes :: (Typed (TypeBase Shape u)) => - ErrorMsg SubExp -> SrcLoc -> [VName] -> [TypeBase Shape u] -> [SubExp] - -> InternaliseM [SubExp] +ensureArgShapes :: + (Typed (TypeBase Shape u)) => + ErrorMsg SubExp -> + SrcLoc -> + [VName] -> + [TypeBase Shape u] -> + [SubExp] -> + InternaliseM [SubExp] ensureArgShapes msg loc shapes paramts args = zipWithM ensureArgShape (expectedTypes shapes paramts args) args - where ensureArgShape _ (Constant v) = return $ Constant v - ensureArgShape t (Var v) - | arrayRank t < 1 = return $ Var v - | otherwise = - ensureShape msg loc t (baseString v) $ Var v + where + ensureArgShape _ (Constant v) = return $ Constant v + ensureArgShape t (Var v) + | arrayRank t < 1 = return $ Var v + | otherwise = + ensureShape msg loc t (baseString v) $ Var v -ensureShapeVar :: ErrorMsg SubExp -> SrcLoc -> ExtType -> String -> VName - -> InternaliseM VName +ensureShapeVar :: + ErrorMsg SubExp -> + SrcLoc -> + ExtType -> + String -> + VName -> + InternaliseM VName ensureShapeVar msg loc t name v - | Array{} <- t = do + | Array {} <- t = do newdims <- arrayDims . removeExistentials t <$> lookupType v olddims <- arrayDims <$> lookupType v if newdims == olddims @@ -112,5 +154,6 @@ cs <- assert "empty_or_match_cert" all_match msg loc certifying cs $ letExp name $ shapeCoerce newdims v | otherwise = return v - where checkDim desired has = - letSubExp "dim_match" $ BasicOp $ CmpOp (CmpEq int32) desired has + where + checkDim desired has = + letSubExp "dim_match" $ BasicOp $ CmpOp (CmpEq int64) desired has diff --git a/src/Futhark/Internalise/Bindings.hs b/src/Futhark/Internalise/Bindings.hs --- a/src/Futhark/Internalise/Bindings.hs +++ b/src/Futhark/Internalise/Bindings.hs @@ -1,93 +1,98 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE Safe #-} + -- | Internalising bindings. module Futhark.Internalise.Bindings - ( - bindingParams - , bindingLoopParams - , bindingLambdaParams - , stmPattern + ( bindingParams, + bindingLoopParams, + bindingLambdaParams, + stmPattern, ) - where +where -import Control.Monad.State hiding (mapM) import Control.Monad.Reader hiding (mapM) - import qualified Data.Map.Strict as M - -import Language.Futhark as E hiding (matchDims) import qualified Futhark.IR.SOACS as I -import Futhark.MonadFreshNames import Futhark.Internalise.Monad import Futhark.Internalise.TypesValues import Futhark.Util +import Language.Futhark as E hiding (matchDims) -bindingParams :: [E.TypeParam] -> [E.Pattern] - -> ([I.FParam] -> [[I.FParam]] -> InternaliseM a) - -> InternaliseM a +bindingParams :: + [E.TypeParam] -> + [E.Pattern] -> + ([I.FParam] -> [[I.FParam]] -> InternaliseM a) -> + InternaliseM a bindingParams tparams params m = do flattened_params <- mapM flattenPattern params let params_idents = concat flattened_params params_ts <- internaliseParamTypes $ - map (flip E.setAliases () . E.unInfo . E.identType) params_idents + map (flip E.setAliases () . E.unInfo . E.identType) params_idents let num_param_idents = map length flattened_params num_param_ts = map (sum . map length) $ chunks num_param_idents params_ts - let shape_params = [ I.Param v $ I.Prim I.int32 | E.TypeParamDim v _ <- tparams ] - shape_subst = M.fromList [ (I.paramName p, [I.Var $ I.paramName p]) | p <- shape_params ] + let shape_params = [I.Param v $ I.Prim I.int64 | E.TypeParamDim v _ <- tparams] + shape_subst = M.fromList [(I.paramName p, [I.Var $ I.paramName p]) | p <- shape_params] bindingFlatPattern params_idents (concat params_ts) $ \valueparams -> - I.localScope (I.scopeOfFParams $ shape_params++concat valueparams) $ - substitutingVars shape_subst $ m shape_params $ - chunks num_param_ts (concat valueparams) + I.localScope (I.scopeOfFParams $ shape_params ++ concat valueparams) $ + substitutingVars shape_subst $ + m shape_params $ + chunks num_param_ts (concat valueparams) -bindingLoopParams :: [E.TypeParam] -> E.Pattern - -> ([I.FParam] -> [I.FParam] -> InternaliseM a) - -> InternaliseM a +bindingLoopParams :: + [E.TypeParam] -> + E.Pattern -> + ([I.FParam] -> [I.FParam] -> InternaliseM a) -> + InternaliseM a bindingLoopParams tparams pat m = do pat_idents <- flattenPattern pat pat_ts <- internaliseLoopParamType (E.patternStructType pat) - let shape_params = [ I.Param v $ I.Prim I.int32 | E.TypeParamDim v _ <- tparams ] - shape_subst = M.fromList [ (I.paramName p, [I.Var $ I.paramName p]) | p <- shape_params ] + let shape_params = [I.Param v $ I.Prim I.int64 | E.TypeParamDim v _ <- tparams] + shape_subst = M.fromList [(I.paramName p, [I.Var $ I.paramName p]) | p <- shape_params] bindingFlatPattern pat_idents pat_ts $ \valueparams -> - I.localScope (I.scopeOfFParams $ shape_params++concat valueparams) $ - substitutingVars shape_subst $ m shape_params $ concat valueparams + I.localScope (I.scopeOfFParams $ shape_params ++ concat valueparams) $ + substitutingVars shape_subst $ m shape_params $ concat valueparams -bindingLambdaParams :: [E.Pattern] -> [I.Type] - -> ([I.LParam] -> InternaliseM a) - -> InternaliseM a +bindingLambdaParams :: + [E.Pattern] -> + [I.Type] -> + ([I.LParam] -> InternaliseM a) -> + InternaliseM a bindingLambdaParams params ts m = do params_idents <- concat <$> mapM flattenPattern params bindingFlatPattern params_idents ts $ \params' -> I.localScope (I.scopeOfLParams $ concat params') $ m $ concat params' -processFlatPattern :: Show t => [E.Ident] -> [t] - -> InternaliseM ([[I.Param t]], VarSubstitutions) +processFlatPattern :: + Show t => + [E.Ident] -> + [t] -> + InternaliseM ([[I.Param t]], VarSubstitutions) processFlatPattern x y = processFlatPattern' [] x y where - processFlatPattern' pat [] _ = do + processFlatPattern' pat [] _ = do let (vs, substs) = unzip pat substs' = M.fromList substs idents = reverse vs return (idents, substs') - - processFlatPattern' pat (p:rest) ts = do + processFlatPattern' pat (p : rest) ts = do (ps, subst, rest_ts) <- handleMapping ts <$> internaliseBindee p processFlatPattern' ((ps, (E.identName p, map (I.Var . I.paramName) subst)) : pat) rest rest_ts handleMapping ts [] = ([], [], ts) - handleMapping ts (r:rs) = - let (ps, reps, ts') = handleMapping' ts r - (pss, repss, ts'') = handleMapping ts' rs - in (ps++pss, reps:repss, ts'') + handleMapping ts (r : rs) = + let (ps, reps, ts') = handleMapping' ts r + (pss, repss, ts'') = handleMapping ts' rs + in (ps ++ pss, reps : repss, ts'') - handleMapping' (t:ts) vname = + handleMapping' (t : ts) vname = let v' = I.Param vname t - in ([v'], v', ts) + in ([v'], v', ts) handleMapping' [] _ = error $ "processFlatPattern: insufficient identifiers in pattern." ++ show (x, y) @@ -99,44 +104,50 @@ 1 -> return [name] _ -> replicateM n $ newVName $ baseString name -bindingFlatPattern :: Show t => [E.Ident] -> [t] - -> ([[I.Param t]] -> InternaliseM a) - -> InternaliseM a +bindingFlatPattern :: + Show t => + [E.Ident] -> + [t] -> + ([[I.Param t]] -> InternaliseM a) -> + InternaliseM a bindingFlatPattern idents ts m = do (ps, substs) <- processFlatPattern idents ts - local (\env -> env { envSubsts = substs `M.union` envSubsts env}) $ + local (\env -> env {envSubsts = substs `M.union` envSubsts env}) $ m ps -- | Flatten a pattern. Returns a list of identifiers. The -- structural type of each identifier is returned separately. flattenPattern :: MonadFreshNames m => E.Pattern -> m [E.Ident] flattenPattern = flattenPattern' - where flattenPattern' (E.PatternParens p _) = - flattenPattern' p - flattenPattern' (E.Wildcard t loc) = do - name <- newVName "nameless" - flattenPattern' $ E.Id name t loc - flattenPattern' (E.Id v (Info t) loc) = - return [E.Ident v (Info t) loc] - -- XXX: treat empty tuples and records as bool. - flattenPattern' (E.TuplePattern [] loc) = - flattenPattern' (E.Wildcard (Info $ E.Scalar $ E.Prim E.Bool) loc) - flattenPattern' (E.RecordPattern [] loc) = - flattenPattern' (E.Wildcard (Info $ E.Scalar $ E.Prim E.Bool) loc) - flattenPattern' (E.TuplePattern pats _) = - concat <$> mapM flattenPattern' pats - flattenPattern' (E.RecordPattern fs loc) = - flattenPattern' $ E.TuplePattern (map snd $ sortFields $ M.fromList fs) loc - flattenPattern' (E.PatternAscription p _ _) = - flattenPattern' p - flattenPattern' (E.PatternLit _ t loc) = - flattenPattern' $ E.Wildcard t loc - flattenPattern' (E.PatternConstr _ _ ps _) = - concat <$> mapM flattenPattern' ps + where + flattenPattern' (E.PatternParens p _) = + flattenPattern' p + flattenPattern' (E.Wildcard t loc) = do + name <- newVName "nameless" + flattenPattern' $ E.Id name t loc + flattenPattern' (E.Id v (Info t) loc) = + return [E.Ident v (Info t) loc] + -- XXX: treat empty tuples and records as bool. + flattenPattern' (E.TuplePattern [] loc) = + flattenPattern' (E.Wildcard (Info $ E.Scalar $ E.Prim E.Bool) loc) + flattenPattern' (E.RecordPattern [] loc) = + flattenPattern' (E.Wildcard (Info $ E.Scalar $ E.Prim E.Bool) loc) + flattenPattern' (E.TuplePattern pats _) = + concat <$> mapM flattenPattern' pats + flattenPattern' (E.RecordPattern fs loc) = + flattenPattern' $ E.TuplePattern (map snd $ sortFields $ M.fromList fs) loc + flattenPattern' (E.PatternAscription p _ _) = + flattenPattern' p + flattenPattern' (E.PatternLit _ t loc) = + flattenPattern' $ E.Wildcard t loc + flattenPattern' (E.PatternConstr _ _ ps _) = + concat <$> mapM flattenPattern' ps -stmPattern :: E.Pattern -> [I.Type] - -> ([VName] -> InternaliseM a) - -> InternaliseM a +stmPattern :: + E.Pattern -> + [I.Type] -> + ([VName] -> InternaliseM a) -> + InternaliseM a stmPattern pat ts m = do pat' <- flattenPattern pat let addShapeStms l = 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,1094 +1,1247 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE Trustworthy #-} --- | Defunctionalization of typed, monomorphic Futhark programs without modules. -module Futhark.Internalise.Defunctionalise - ( transformProg ) where - -import qualified Control.Arrow as Arrow -import Control.Monad.Identity -import Control.Monad.State -import Control.Monad.RWS hiding (Sum) -import Data.Bifunctor -import Data.Bitraversable -import Data.Foldable -import Data.List (sortOn, nub, partition, tails) -import qualified Data.List.NonEmpty as NE -import Data.Maybe -import qualified Data.Map.Strict as M -import qualified Data.Set as S -import qualified Data.Sequence as Seq - -import Futhark.MonadFreshNames -import Language.Futhark -import Language.Futhark.Traversals -import Futhark.IR.Pretty () - --- | An expression or an extended 'Lambda' (with size parameters, --- which AST lambdas do not support). -data ExtExp = ExtLambda [TypeParam] [Pattern] Exp (Aliasing, StructType) SrcLoc - | ExtExp Exp - deriving (Show) - --- | A static value stores additional information about the result of --- defunctionalization of an expression, aside from the residual expression. -data StaticVal = Dynamic PatternType - | LambdaSV [VName] Pattern StructType ExtExp Env - -- ^ The 'VName's are shape parameters that are bound - -- by the 'Pattern'. - | RecordSV [(Name, StaticVal)] - | SumSV Name [StaticVal] [(Name, [PatternType])] - -- ^ The constructor that is actually present, plus - -- the others that are not. - | DynamicFun (Exp, StaticVal) StaticVal - | IntrinsicSV - deriving (Show) - --- | Environment mapping variable names to their associated static value. -type Env = M.Map VName StaticVal - -localEnv :: Env -> DefM a -> DefM a -localEnv env = local $ Arrow.second (env<>) - --- Even when using a "new" environment (for evaluating closures) we --- still ram the global environment of DynamicFuns in there. -localNewEnv :: Env -> DefM a -> DefM a -localNewEnv env = local $ \(globals, old_env) -> - (globals, M.filterWithKey (\k _ -> k `S.member` globals) old_env <> env) - -extendEnv :: VName -> StaticVal -> DefM a -> DefM a -extendEnv vn sv = localEnv (M.singleton vn sv) - -askEnv :: DefM Env -askEnv = asks snd - -isGlobal :: VName -> DefM a -> DefM a -isGlobal v = local $ Arrow.first (S.insert v) - --- | Returns the defunctionalization environment restricted --- to the given set of variable names and types. -restrictEnvTo :: NameSet -> DefM Env -restrictEnvTo (NameSet m) = restrict <$> ask - where restrict (globals, env) = M.mapMaybeWithKey keep env - where keep k sv = do guard $ not $ k `S.member` globals - u <- M.lookup k m - Just $ restrict' u sv - restrict' Nonunique (Dynamic t) = - Dynamic $ t `setUniqueness` Nonunique - restrict' _ (Dynamic t) = - Dynamic t - restrict' u (LambdaSV dims pat t e env) = - LambdaSV dims pat t e $ M.map (restrict' u) env - restrict' u (RecordSV fields) = - RecordSV $ map (fmap $ restrict' u) fields - restrict' u (SumSV c svs fields) = - SumSV c (map (restrict' u) svs) fields - restrict' u (DynamicFun (e, sv1) sv2) = - DynamicFun (e, restrict' u sv1) $ restrict' u sv2 - restrict' _ IntrinsicSV = IntrinsicSV - --- | Defunctionalization monad. The Reader environment tracks both --- the current Env as well as the set of globally defined dynamic --- functions. This is used to avoid unnecessarily large closure --- environments. -newtype DefM a = DefM (RWS (S.Set VName, Env) (Seq.Seq ValBind) VNameSource a) - deriving (Functor, Applicative, Monad, - MonadReader (S.Set VName, Env), - MonadWriter (Seq.Seq ValBind), - MonadFreshNames) - --- | Run a computation in the defunctionalization monad. Returns the result of --- the computation, a new name source, and a list of lifted function declations. -runDefM :: VNameSource -> DefM a -> (a, VNameSource, Seq.Seq ValBind) -runDefM src (DefM m) = runRWS m mempty src - -collectFuns :: DefM a -> DefM (a, Seq.Seq ValBind) -collectFuns m = pass $ do - (x, decs) <- listen m - return ((x, decs), const mempty) - --- | Looks up the associated static value for a given name in the environment. -lookupVar :: SrcLoc -> VName -> DefM StaticVal -lookupVar loc x = do - env <- askEnv - case M.lookup x env of - Just sv -> return sv - Nothing -- If the variable is unknown, it may refer to the 'intrinsics' - -- module, which we will have to treat specially. - | baseTag x <= maxIntrinsicTag -> return IntrinsicSV - | otherwise -> -- Anything not in scope is going to be an - -- existential size. - return $ Dynamic $ Scalar $ Prim $ Signed Int32 - | otherwise -> error $ "Variable " ++ pretty x ++ " at " - ++ locStr loc ++ " is out of scope." - --- Like patternDimNames, but ignores sizes that are only found in --- funtion types. -arraySizes :: StructType -> S.Set VName -arraySizes (Scalar Arrow{}) = mempty -arraySizes (Scalar (Record fields)) = foldMap arraySizes fields -arraySizes (Scalar (Sum cs)) = foldMap (foldMap arraySizes) cs -arraySizes (Scalar (TypeVar _ _ _ targs)) = - mconcat $ map f targs - where f (TypeArgDim (NamedDim d) _) = S.singleton $ qualLeaf d - f TypeArgDim{} = mempty - f (TypeArgType t _) = arraySizes t -arraySizes (Scalar Prim{}) = mempty -arraySizes (Array _ _ t shape) = - arraySizes (Scalar t) <> foldMap dimName (shapeDims shape) - where dimName :: DimDecl VName -> S.Set VName - dimName (NamedDim qn) = S.singleton $ qualLeaf qn - dimName _ = mempty - -patternArraySizes :: Pattern -> S.Set VName -patternArraySizes = arraySizes . patternStructType - -dimMapping :: Monoid a => - TypeBase (DimDecl VName) a - -> TypeBase (DimDecl VName) a - -> M.Map VName VName -dimMapping t1 t2 = execState (matchDims f t1 t2) mempty - where f (NamedDim d1) (NamedDim d2) = do - modify $ M.insert (qualLeaf d1) (qualLeaf d2) - return $ NamedDim d1 - f d _ = return d - -defuncFun :: [TypeParam] -> [Pattern] -> Exp -> (Aliasing, StructType) -> SrcLoc - -> DefM (Exp, StaticVal) -defuncFun tparams pats e0 (closure, ret) loc = do - when (any isTypeParam tparams) $ - error $ "Received a lambda with type parameters at " ++ locStr loc - ++ ", but the defunctionalizer expects a monomorphic input program." - -- Extract the first parameter of the lambda and "push" the - -- remaining ones (if there are any) into the body of the lambda. - let (dims, pat, ret', e0') = case pats of - [] -> error "Received a lambda with no parameters." - [pat'] -> (map typeParamName tparams, pat', ret, ExtExp e0) - (pat' : pats') -> - -- Split shape parameters into those that are determined by - -- the first pattern, and those that are determined by later - -- patterns. - let bound_by_pat = (`S.member` patternArraySizes pat') . typeParamName - (pat_dims, rest_dims) = partition bound_by_pat tparams - in (map typeParamName pat_dims, pat', - foldFunType (map (toStruct . patternType) pats') ret, - ExtLambda rest_dims pats' e0 (closure, ret) loc) - - -- Construct a record literal that closes over the environment of - -- the lambda. Closed-over 'DynamicFun's are converted to their - -- closure representation. - let used = freeVars (Lambda pats e0 Nothing (Info (closure, ret)) loc) - `without` mconcat (map oneName dims) - used_env <- restrictEnvTo used - - -- The closure parts that are sizes are proactively turned into size - -- parameters. - let sizes_of_arrays = foldMap (arraySizes . toStruct . typeFromSV') used_env <> - patternArraySizes pat - notSize = not . (`S.member` sizes_of_arrays) - (fields, env) = unzip $ map closureFromDynamicFun $ - filter (notSize . fst) $ M.toList used_env - env' = M.fromList env - closure_dims = S.toList sizes_of_arrays - - global <- asks fst - - return (RecordLit fields loc, - LambdaSV (nub $ filter (`S.notMember` global) $ - dims<>closure_dims) pat ret' e0' env') - - where closureFromDynamicFun (vn, DynamicFun (clsr_env, sv) _) = - let name = nameFromString $ pretty vn - in (RecordFieldExplicit name clsr_env mempty, (vn, sv)) - - closureFromDynamicFun (vn, sv) = - let name = nameFromString $ pretty vn - tp' = typeFromSV' sv - in (RecordFieldExplicit name - (Var (qualName vn) (Info tp') mempty) mempty, (vn, sv)) - --- | Defunctionalization of an expression. Returns the residual expression and --- the associated static value in the defunctionalization monad. -defuncExp :: Exp -> DefM (Exp, StaticVal) - -defuncExp e@Literal{} = - return (e, Dynamic $ typeOf e) - -defuncExp e@IntLit{} = - return (e, Dynamic $ typeOf e) - -defuncExp e@FloatLit{} = - return (e, Dynamic $ typeOf e) - -defuncExp e@StringLit{} = - return (e, Dynamic $ typeOf e) - -defuncExp (Parens e loc) = do - (e', sv) <- defuncExp e - return (Parens e' loc, sv) - -defuncExp (QualParens qn e loc) = do - (e', sv) <- defuncExp e - return (QualParens qn e' loc, sv) - -defuncExp (TupLit es loc) = do - (es', svs) <- unzip <$> mapM defuncExp es - return (TupLit es' loc, RecordSV $ zip tupleFieldNames svs) - -defuncExp (RecordLit fs loc) = do - (fs', names_svs) <- unzip <$> mapM defuncField fs - return (RecordLit fs' loc, RecordSV names_svs) - - where defuncField (RecordFieldExplicit vn e loc') = do - (e', sv) <- defuncExp e - return (RecordFieldExplicit vn e' loc', (vn, sv)) - defuncField (RecordFieldImplicit vn _ loc') = do - sv <- lookupVar loc' vn - case sv of - -- If the implicit field refers to a dynamic function, we - -- convert it to an explicit field with a record closing over - -- the environment and bind the corresponding static value. - DynamicFun (e, sv') _ -> let vn' = baseName vn - in return (RecordFieldExplicit vn' e loc', - (vn', sv')) - -- The field may refer to a functional expression, so we get the - -- type from the static value and not the one from the AST. - _ -> let tp = Info $ typeFromSV' sv - in return (RecordFieldImplicit vn tp loc', (baseName vn, sv)) - -defuncExp (ArrayLit es t@(Info t') loc) = do - es' <- mapM defuncExp' es - return (ArrayLit es' t loc, Dynamic t') - -defuncExp (Range e1 me incl t@(Info t', _) loc) = do - e1' <- defuncExp' e1 - me' <- mapM defuncExp' me - incl' <- mapM defuncExp' incl - return (Range e1' me' incl' t loc, Dynamic t') - -defuncExp e@(Var qn _ loc) = do - sv <- lookupVar loc (qualLeaf qn) - case sv of - -- If the variable refers to a dynamic function, we return its closure - -- representation (i.e., a record expression capturing the free variables - -- and a 'LambdaSV' static value) instead of the variable itself. - DynamicFun closure _ -> return closure - -- Intrinsic functions used as variables are eta-expanded, so we - -- can get rid of them. - IntrinsicSV -> do - (pats, body, tp) <- etaExpand (typeOf e) e - defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty - _ -> let tp = typeFromSV' sv - in return (Var qn (Info tp) loc, sv) - -defuncExp (Ascript e0 tydecl loc) - | orderZero (typeOf e0) = do (e0', sv) <- defuncExp e0 - return (Ascript e0' tydecl loc, sv) - | otherwise = defuncExp e0 - -defuncExp (Coerce e0 tydecl t loc) - | orderZero (typeOf e0) = do (e0', sv) <- defuncExp e0 - return (Coerce e0' tydecl t loc, sv) - | otherwise = defuncExp e0 - -defuncExp (LetPat pat e1 e2 (Info t, retext) loc) = do - (e1', sv1) <- defuncExp e1 - let env = matchPatternSV pat sv1 - pat' = updatePattern' pat sv1 - (e2', sv2) <- localEnv env $ defuncExp e2 - -- To maintain any sizes going out of scope, we need to compute the - -- old size substitution induced by retext and also apply it to the - -- newly computed body type. - let mapping = dimMapping (typeOf e2) t - subst v = fromMaybe v $ M.lookup v mapping - t' = first (fmap subst) $ typeOf e2' - return (LetPat pat' e1' e2' (Info t', retext) loc, sv2) - --- Local functions are handled by rewriting them to lambdas, so that --- the same machinery can be re-used. But we may have to eta-expand --- first. -defuncExp (LetFun vn (dims, pats, _, Info ret, e1) e2 let_t loc) - | Scalar Arrow{} <- ret = do - (body_pats, e1', ret') <- etaExpand (fromStruct ret) e1 - let f = (dims, pats <> body_pats, Nothing, Info ret', e1') - defuncExp $ LetFun vn f e2 let_t loc - - | otherwise = do - (e1', sv1) <- defuncFun dims pats e1 (mempty, ret) loc - (e2', sv2) <- localEnv (M.singleton vn sv1) $ defuncExp e2 - return (LetPat (Id vn (Info (typeOf e1')) loc) e1' e2' (Info $ typeOf e2', Info []) loc, - sv2) - -defuncExp (If e1 e2 e3 tp loc) = do - (e1', _ ) <- defuncExp e1 - (e2', sv) <- defuncExp e2 - (e3', _ ) <- defuncExp e3 - return (If e1' e2' e3' tp loc, sv) - -defuncExp e@(Apply f@(Var f' _ _) arg d (t, ext) loc) - | baseTag (qualLeaf f') <= maxIntrinsicTag, - TupLit es tuploc <- arg = do - -- defuncSoacExp also works fine for non-SOACs. - es' <- mapM defuncSoacExp es - return (Apply f (TupLit es' tuploc) d (t, ext) loc, - Dynamic $ typeOf e) - -defuncExp e@Apply{} = defuncApply 0 e - -defuncExp (Negate e0 loc) = do - (e0', sv) <- defuncExp e0 - return (Negate e0' loc, sv) - -defuncExp (Lambda pats e0 _ (Info (closure, ret)) loc) = - defuncFun [] pats e0 (closure, ret) loc - --- Operator sections are expected to be converted to lambda-expressions --- by the monomorphizer, so they should no longer occur at this point. -defuncExp OpSection{} = error "defuncExp: unexpected operator section." -defuncExp OpSectionLeft{} = error "defuncExp: unexpected operator section." -defuncExp OpSectionRight{} = error "defuncExp: unexpected operator section." -defuncExp ProjectSection{} = error "defuncExp: unexpected projection section." -defuncExp IndexSection{} = error "defuncExp: unexpected projection section." - -defuncExp (DoLoop sparams pat e1 form e3 ret loc) = do - (e1', sv1) <- defuncExp e1 - let env1 = matchPatternSV pat sv1 - (form', env2) <- case form of - For v e2 -> do e2' <- defuncExp' e2 - return (For v e2', envFromIdent v) - ForIn pat2 e2 -> do e2' <- defuncExp' e2 - return (ForIn pat2 e2', envFromPattern pat2) - While e2 -> do e2' <- localEnv env1 $ defuncExp' e2 - return (While e2', mempty) - (e3', sv) <- localEnv (env1 <> env2) $ defuncExp e3 - return (DoLoop sparams pat e1' form' e3' ret loc, sv) - where envFromIdent (Ident vn (Info tp) _) = - M.singleton vn $ Dynamic tp - --- We handle BinOps by turning them into ordinary function applications. -defuncExp (BinOp (qn, qnloc) (Info t) - (e1, Info (pt1, ext1)) (e2, Info (pt2, ext2)) - (Info ret) (Info retext) loc) = - defuncExp $ Apply (Apply (Var qn (Info t) qnloc) - e1 (Info (diet pt1, ext1)) - (Info (Scalar $ Arrow mempty Unnamed (fromStruct pt2) ret), Info []) loc) - e2 (Info (diet pt2, ext2)) (Info ret, Info retext) loc - -defuncExp (Project vn e0 tp@(Info tp') loc) = do - (e0', sv0) <- defuncExp e0 - case sv0 of - RecordSV svs -> case lookup vn svs of - Just sv -> return (Project vn e0' (Info $ typeFromSV' sv) loc, sv) - Nothing -> error "Invalid record projection." - Dynamic _ -> return (Project vn e0' tp loc, Dynamic tp') - _ -> error $ "Projection of an expression with static value " ++ show sv0 - -defuncExp (LetWith id1 id2 idxs e1 body t loc) = do - e1' <- defuncExp' e1 - sv1 <- lookupVar (identSrcLoc id2) $ identName id2 - idxs' <- mapM defuncDimIndex idxs - (body', sv) <- extendEnv (identName id1) sv1 $ defuncExp body - return (LetWith id1 id2 idxs' e1' body' t loc, sv) - -defuncExp expr@(Index e0 idxs info loc) = do - e0' <- defuncExp' e0 - idxs' <- mapM defuncDimIndex idxs - return (Index e0' idxs' info loc, Dynamic $ typeOf expr) - -defuncExp (Update e1 idxs e2 loc) = do - (e1', sv) <- defuncExp e1 - idxs' <- mapM defuncDimIndex idxs - e2' <- defuncExp' e2 - return (Update e1' idxs' e2' loc, sv) - --- Note that we might change the type of the record field here. This --- is not permitted in the type checker due to problems with type --- inference, but it actually works fine. -defuncExp (RecordUpdate e1 fs e2 _ loc) = do - (e1', sv1) <- defuncExp e1 - (e2', sv2) <- defuncExp e2 - let sv = staticField sv1 sv2 fs - return (RecordUpdate e1' fs e2' (Info $ typeFromSV' sv1) loc, - sv) - where staticField (RecordSV svs) sv2 (f:fs') = - case lookup f svs of - Just sv -> RecordSV $ - (f, staticField sv sv2 fs') : filter ((/=f) . fst) svs - Nothing -> error "Invalid record projection." - staticField (Dynamic t@(Scalar Record{})) sv2 fs'@(_:_) = - staticField (svFromType t) sv2 fs' - staticField _ sv2 _ = sv2 - -defuncExp (Assert e1 e2 desc loc) = do - (e1', _) <- defuncExp e1 - (e2', sv) <- defuncExp e2 - return (Assert e1' e2' desc loc, sv) - -defuncExp (Constr name es (Info (Scalar (Sum all_fs))) loc) = do - (es', svs) <- unzip <$> mapM defuncExp es - let sv = SumSV name svs $ M.toList $ - name `M.delete` M.map (map defuncType) all_fs - return (Constr name es' (Info (typeFromSV' sv)) loc, sv) - where defuncType :: Monoid als => - TypeBase (DimDecl VName) als - -> TypeBase (DimDecl VName) als - defuncType (Array as u t shape) = Array as u (defuncScalar t) shape - defuncType (Scalar t) = Scalar $ defuncScalar t - - defuncScalar :: Monoid als => - ScalarTypeBase (DimDecl VName) als - -> ScalarTypeBase (DimDecl VName) als - defuncScalar (Record fs) = Record $ M.map defuncType fs - defuncScalar Arrow{} = Record mempty - defuncScalar (Sum fs) = Sum $ M.map (map defuncType) fs - defuncScalar (Prim t) = Prim t - defuncScalar (TypeVar as u tn targs) = TypeVar as u tn targs - -defuncExp (Constr name _ (Info t) loc) = - error $ "Constructor " ++ pretty name ++ " given type " ++ - pretty t ++ " at " ++ locStr loc - -defuncExp (Match e cs t loc) = do - (e', sv) <- defuncExp e - csPairs <- mapM (defuncCase sv) cs - let cs' = fmap fst csPairs - sv' = snd $ NE.head csPairs - return (Match e' cs' t loc, sv') - -defuncExp (Attr info e loc) = do - (e', sv) <- defuncExp e - return (Attr info e' loc, sv) - --- | Same as 'defuncExp', except it ignores the static value. -defuncExp' :: Exp -> DefM Exp -defuncExp' = fmap fst . defuncExp - -defuncExtExp :: ExtExp -> DefM (Exp, StaticVal) -defuncExtExp (ExtExp e) = defuncExp e -defuncExtExp (ExtLambda tparams pats e0 (closure, ret) loc) = - defuncFun tparams pats e0 (closure, ret) loc - -defuncCase :: StaticVal -> Case -> DefM (Case, StaticVal) -defuncCase sv (CasePat p e loc) = do - let p' = updatePattern' p sv - env = matchPatternSV p sv - (e', sv') <- localEnv env $ defuncExp e - return (CasePat p' e' loc, sv') - --- | Defunctionalize the function argument to a SOAC by eta-expanding if --- necessary and then defunctionalizing the body of the introduced lambda. -defuncSoacExp :: Exp -> DefM Exp -defuncSoacExp e@OpSection{} = return e -defuncSoacExp e@OpSectionLeft{} = return e -defuncSoacExp e@OpSectionRight{} = return e -defuncSoacExp e@ProjectSection{} = return e - -defuncSoacExp (Parens e loc) = - Parens <$> defuncSoacExp e <*> pure loc - -defuncSoacExp (Lambda params e0 decl tp loc) = do - let env = foldMap envFromPattern params - e0' <- localEnv env $ defuncSoacExp e0 - return $ Lambda params e0' decl tp loc - -defuncSoacExp e - | Scalar Arrow{} <- typeOf e = do - (pats, body, tp) <- etaExpand (typeOf e) e - let env = foldMap envFromPattern pats - body' <- localEnv env $ defuncExp' body - return $ Lambda pats body' Nothing (Info (mempty, tp)) mempty - | otherwise = defuncExp' e - -etaExpand :: PatternType -> Exp -> DefM ([Pattern], Exp, StructType) -etaExpand e_t e = do - let (ps, ret) = getType e_t - (pats, vars) <- fmap unzip . forM ps $ \(p, t) -> do - x <- case p of Named x -> pure x - Unnamed -> newNameFromString "x" - return (Id x (Info t) mempty, - Var (qualName x) (Info t) mempty) - let e' = foldl' (\e1 (e2, t2, argtypes) -> - Apply e1 e2 (Info (diet t2, Nothing)) - (Info (foldFunType argtypes ret), Info []) mempty) - e $ zip3 vars (map snd ps) (drop 1 $ tails $ map snd ps) - return (pats, e', toStruct ret) - - where getType (Scalar (Arrow _ p t1 t2)) = - let (ps, r) = getType t2 in ((p,t1) : ps, r) - getType t = ([], t) - --- | Defunctionalize an indexing of a single array dimension. -defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName) -defuncDimIndex (DimFix e1) = DimFix . fst <$> defuncExp e1 -defuncDimIndex (DimSlice me1 me2 me3) = - DimSlice <$> defunc' me1 <*> defunc' me2 <*> defunc' me3 - where defunc' = mapM defuncExp' - --- | Defunctionalize a let-bound function, while preserving parameters --- that have order 0 types (i.e., non-functional). -defuncLet :: [TypeParam] -> [Pattern] -> Exp -> StructType - -> DefM ([TypeParam], [Pattern], Exp, StaticVal) -defuncLet dims ps@(pat:pats) body rettype - | patternOrderZero pat = do - - let bound_by_pat = (`S.member` patternDimNames pat) . typeParamName - -- Take care to not include more size parameters than necessary. - (pat_dims, rest_dims) = partition bound_by_pat dims - env = envFromPattern pat <> envFromShapeParams pat_dims - (rest_dims', pats', body', sv) <- localEnv env $ defuncLet rest_dims pats body rettype - closure <- defuncFun dims ps body (mempty, rettype) mempty - return (pat_dims ++ rest_dims', pat : pats', body', DynamicFun closure sv) - | otherwise = do - (e, sv) <- defuncFun dims ps body (mempty, rettype) mempty - return ([], [], e, sv) - -defuncLet _ [] body rettype = do - (body', sv) <- defuncExp body - return ([], [], body', imposeType sv rettype) - where imposeType Dynamic{} t = - Dynamic $ fromStruct t - imposeType (RecordSV fs1) (Scalar (Record fs2)) = - RecordSV $ M.toList $ M.intersectionWith imposeType (M.fromList fs1) fs2 - imposeType sv _ = sv - -sizesForAll :: MonadFreshNames m => [Pattern] -> m ([VName], [Pattern]) -sizesForAll params = do - (params', sizes) <- runStateT (mapM (astMap tv) params) [] - return (sizes, params') - where tv = identityMapper { mapOnPatternType = bitraverse onDim pure } - onDim AnyDim = do v <- lift $ newVName "size" - modify (v:) - pure $ NamedDim $ qualName v - onDim d = pure d - --- | Defunctionalize an application expression at a given depth of application. --- Calls to dynamic (first-order) functions are preserved at much as possible, --- but a new lifted function is created if a dynamic function is only partially --- applied. -defuncApply :: Int -> Exp -> DefM (Exp, StaticVal) -defuncApply depth e@(Apply e1 e2 d t@(Info ret, Info ext) loc) = do - let (argtypes, _) = unfoldFunType ret - (e1', sv1) <- defuncApply (depth+1) e1 - (e2', sv2) <- defuncExp e2 - let e' = Apply e1' e2' d t loc - case sv1 of - LambdaSV dims pat e0_t e0 closure_env -> do - let env' = matchPatternSV pat sv2 - env_dim = envFromDimNames dims - (e0', sv) <- localNewEnv (env' <> closure_env <> env_dim) $ defuncExtExp e0 - - let closure_pat = buildEnvPattern closure_env - pat' = updatePattern pat sv2 - - globals <- asks fst - - -- Lift lambda to top-level function definition. We put in - -- a lot of effort to try to infer the uniqueness attributes - -- of the lifted function, but this is ultimately all a sham - -- and a hack. There is some piece we're missing. - let params = [closure_pat, pat'] - params_for_rettype = params ++ svParams sv1 ++ svParams sv2 - svParams (LambdaSV _ sv_pat _ _ _) = [sv_pat] - svParams _ = [] - rettype = buildRetType closure_env params_for_rettype e0_t $ typeOf e0' - - already_bound = globals <> S.fromList dims <> - S.map identName (foldMap patternIdents params) - more_dims = S.toList $ - S.filter (`S.notMember` already_bound) $ - foldMap patternArraySizes params - - -- Embed some information about the original function - -- into the name of the lifted function, to make the - -- result slightly more human-readable. - liftedName i (Var f _ _) = - "lifted_" ++ show i ++ "_" ++ baseString (qualLeaf f) - liftedName i (Apply f _ _ _ _) = - liftedName (i+1) f - liftedName _ _ = "lifted" - - -- Ensure that no parameter sizes are AnyDim. The internaliser - -- expects this. This is easy, because they are all - -- first-order. - (missing_dims, params') <- sizesForAll params - - fname <- newNameFromString $ liftedName (0::Int) e1 - liftValDec fname rettype (dims ++ more_dims ++ missing_dims) - params' e0' - - let t1 = toStruct $ typeOf e1' - t2 = toStruct $ typeOf e2' - fname' = qualName fname - fname'' = Var fname' (Info (Scalar $ Arrow mempty Unnamed (fromStruct t1) $ - Scalar $ Arrow mempty Unnamed (fromStruct t2) rettype)) - loc - - -- FIXME: what if this application returns both a function - -- and a value? - callret | orderZero ret = (Info ret, Info ext) - | otherwise = (Info rettype, Info ext) - - return (Parens (Apply (Apply fname'' e1' - (Info (Observe, Nothing)) - (Info $ Scalar $ Arrow mempty Unnamed (fromStruct t2) rettype, - Info []) - loc) - e2' d callret loc) mempty, sv) - - -- If e1 is a dynamic function, we just leave the application in place, - -- but we update the types since it may be partially applied or return - -- a higher-order term. - DynamicFun _ sv -> - let (argtypes', rettype) = dynamicFunType sv argtypes - restype = foldFunType argtypes' rettype `setAliases` aliases ret - -- FIXME: what if this application returns both a function - -- and a value? - callret | orderZero ret = (Info ret, Info ext) - | otherwise = (Info restype, Info ext) - apply_e = Apply e1' e2' d callret loc - in return (apply_e, sv) - - -- Propagate the 'IntrinsicsSV' until we reach the outermost application, - -- where we construct a dynamic static value with the appropriate type. - IntrinsicSV - | depth == 0 -> - -- If the intrinsic is fully applied, then we are done. - -- Otherwise we need to eta-expand it and recursively - -- defunctionalise. XXX: might it be better to simply - -- eta-expand immediately any time we encounter a - -- non-fully-applied intrinsic? - if null argtypes - then return (e', Dynamic $ typeOf e) - else do (pats, body, tp) <- etaExpand (typeOf e') e' - defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty - | otherwise -> return (e', IntrinsicSV) - - _ -> error $ "Application of an expression that is neither a static lambda " - ++ "nor a dynamic function, but has static value: " ++ show sv1 - -defuncApply depth e@(Var qn (Info t) loc) = do - let (argtypes, _) = unfoldFunType t - sv <- lookupVar loc (qualLeaf qn) - case sv of - DynamicFun _ _ - | fullyApplied sv depth -> - -- We still need to update the types in case the dynamic - -- function returns a higher-order term. - let (argtypes', rettype) = dynamicFunType sv argtypes - in return (Var qn (Info (foldFunType argtypes' rettype)) loc, sv) - - | otherwise -> do - fname <- newName $ qualLeaf qn - let (dims, pats, e0, sv') = liftDynFun sv depth - pats_names = S.map identName $ mconcat $ map patternIdents pats - notInPats = (`S.notMember` pats_names) - dims' = filter notInPats dims - (argtypes', rettype) = dynamicFunType sv' argtypes - liftValDec fname (fromStruct rettype) dims' pats e0 - return (Var (qualName fname) - (Info (foldFunType argtypes' $ fromStruct rettype)) loc, sv') - - IntrinsicSV -> return (e, IntrinsicSV) - - _ -> return (Var qn (Info (typeFromSV' sv)) loc, sv) - -defuncApply depth (Parens e _) = defuncApply depth e - -defuncApply _ expr = defuncExp expr - --- | Check if a 'StaticVal' and a given application depth corresponds --- to a fully applied dynamic function. -fullyApplied :: StaticVal -> Int -> Bool -fullyApplied (DynamicFun _ sv) depth - | depth == 0 = False - | depth > 0 = fullyApplied sv (depth-1) -fullyApplied _ _ = True - --- | Converts a dynamic function 'StaticVal' into a list of --- dimensions, a list of parameters, a function body, and the --- appropriate static value for applying the function at the given --- depth of partial application. -liftDynFun :: StaticVal -> Int -> ([VName], [Pattern], Exp, StaticVal) -liftDynFun (DynamicFun (e, sv) _) 0 = ([], [], e, sv) -liftDynFun (DynamicFun clsr@(_, LambdaSV dims pat _ _ _) sv) d - | d > 0 = let (dims', pats, e', sv') = liftDynFun sv (d-1) - in (nub $ dims ++ dims', pat : pats, e', DynamicFun clsr sv') -liftDynFun sv _ = error $ "Tried to lift a StaticVal " ++ show sv - ++ ", but expected a dynamic function." - --- | Converts a pattern to an environment that binds the individual names of the --- pattern to their corresponding types wrapped in a 'Dynamic' static value. -envFromPattern :: Pattern -> Env -envFromPattern pat = case pat of - TuplePattern ps _ -> foldMap envFromPattern ps - RecordPattern fs _ -> foldMap (envFromPattern . snd) fs - PatternParens p _ -> envFromPattern p - Id vn (Info t) _ -> M.singleton vn $ Dynamic t - Wildcard _ _ -> mempty - PatternAscription p _ _ -> envFromPattern p - PatternLit{} -> mempty - PatternConstr _ _ ps _ -> foldMap envFromPattern ps - --- | Create an environment that binds the shape parameters. -envFromShapeParams :: [TypeParamBase VName] -> Env -envFromShapeParams = envFromDimNames . map dim - where dim (TypeParamDim vn _) = vn - dim tparam = error $ - "The defunctionalizer expects a monomorphic input program,\n" ++ - "but it received a type parameter " ++ pretty tparam ++ - " at " ++ locStr (srclocOf tparam) ++ "." - -envFromDimNames :: [VName] -> Env -envFromDimNames = M.fromList . flip zip (repeat $ Dynamic $ Scalar $ Prim $ Signed Int32) - --- | Create a new top-level value declaration with the given function name, --- return type, list of parameters, and body expression. -liftValDec :: VName -> PatternType -> [VName] -> [Pattern] -> Exp -> DefM () -liftValDec fname rettype dims pats body = tell $ Seq.singleton dec - where dims' = map (`TypeParamDim` mempty) dims - -- FIXME: this pass is still not correctly size-preserving, so - -- forget those return sizes that we forgot to propagate along - -- the way. Hopefully the internaliser is conservative and - -- will insert reshapes... - bound_here = S.fromList dims <> S.map identName (foldMap patternIdents pats) - anyDimIfNotBound (NamedDim v) - | qualLeaf v `S.member` bound_here = NamedDim v - | otherwise = AnyDim - anyDimIfNotBound d = d - rettype_st = first anyDimIfNotBound $ toStruct rettype - - dec = ValBind - { valBindEntryPoint = Nothing - , valBindName = fname - , valBindRetDecl = Nothing - , valBindRetType = Info (rettype_st, []) - , valBindTypeParams = dims' - , valBindParams = pats - , valBindBody = body - , valBindDoc = Nothing - , valBindAttrs = mempty - , valBindLocation = mempty - } - --- | Given a closure environment, construct a record pattern that --- binds the closed over variables. -buildEnvPattern :: Env -> Pattern -buildEnvPattern env = RecordPattern (map buildField $ M.toList env) mempty - where buildField (vn, sv) = - (nameFromString (pretty vn), - Id vn (Info $ snd $ typeFromSV sv) mempty) - --- | Given a closure environment pattern and the type of a term, --- construct the type of that term, where uniqueness is set to --- `Nonunique` for those arrays that are bound in the environment or --- pattern (except if they are unique there). This ensures that a --- lifted function can create unique arrays as long as they do not --- alias any of its parameters. XXX: it is not clear that this is a --- sufficient property, unfortunately. -buildRetType :: Env -> [Pattern] -> StructType -> PatternType -> PatternType -buildRetType env pats = comb - where bound = foldMap oneName (M.keys env) <> foldMap patternVars pats - boundAsUnique v = - maybe False (unique . unInfo . identType) $ - find ((==v) . identName) $ S.toList $ foldMap patternIdents pats - problematic v = (v `member` bound) && not (boundAsUnique v) - comb (Scalar (Record fs_annot)) (Scalar (Record fs_got)) = - Scalar $ Record $ M.intersectionWith comb fs_annot fs_got - comb (Scalar (Sum cs_annot)) (Scalar (Sum cs_got)) = - Scalar $ Sum $ M.intersectionWith (zipWith comb) cs_annot cs_got - comb (Scalar Arrow{}) t = - descend t - comb got et = - descend $ fromStruct got `setAliases` aliases et - - descend t@Array{} - | any (problematic . aliasVar) (aliases t) = t `setUniqueness` Nonunique - descend (Scalar (Record t)) = Scalar $ Record $ fmap descend t - descend t = t - --- | Compute the corresponding type for a given static value. -typeFromSV :: StaticVal -> ([VName], PatternType) -typeFromSV (Dynamic tp) = - (mempty, tp) -typeFromSV (LambdaSV sizes _ _ _ env) = - (sizes <> env_sizes, - Scalar $ Record $ M.fromList $ map (fmap snd) env') - where env' = map (bimap (nameFromString . pretty) typeFromSV) $ M.toList env - env_sizes = concatMap (fst . snd) env' -typeFromSV (RecordSV ls) = - let ts = map (fmap typeFromSV) ls - in (concatMap (fst . snd) ts, - Scalar $ Record $ M.fromList $ map (fmap snd) ts) -typeFromSV (DynamicFun (_, sv) _) = - typeFromSV sv -typeFromSV (SumSV name svs fields) = - let (sizes, svs') = unzip $ map typeFromSV svs - in (concat sizes, - Scalar $ Sum $ M.insert name svs' $ M.fromList fields) -typeFromSV IntrinsicSV = - error "Tried to get the type from the static value of an intrinsic." - -typeFromSV' :: StaticVal -> PatternType -typeFromSV' sv = - let (sizes, t) = typeFromSV sv - in unscopeType (S.fromList sizes) t - --- | Construct the type for a fully-applied dynamic function from its --- static value and the original types of its arguments. -dynamicFunType :: StaticVal -> [PatternType] -> ([PatternType], PatternType) -dynamicFunType (DynamicFun _ sv) (p:ps) = - let (ps', ret) = dynamicFunType sv ps in (p : ps', ret) -dynamicFunType sv _ = ([], typeFromSV' sv) - --- | Match a pattern with its static value. Returns an environment with --- the identifier components of the pattern mapped to the corresponding --- subcomponents of the static value. -matchPatternSV :: PatternBase Info VName -> StaticVal -> Env -matchPatternSV (TuplePattern ps _) (RecordSV ls) = - mconcat $ zipWith (\p (_, sv) -> matchPatternSV p sv) ps ls -matchPatternSV (RecordPattern ps _) (RecordSV ls) - | ps' <- sortOn fst ps, ls' <- sortOn fst ls, - map fst ps' == map fst ls' = - mconcat $ zipWith (\(_, p) (_, sv) -> matchPatternSV p sv) ps' ls' -matchPatternSV (PatternParens pat _) sv = matchPatternSV pat sv -matchPatternSV (Id vn (Info t) _) sv = - -- When matching a pattern with a zero-order STaticVal, the type of - -- the pattern wins out. This is important when matching a - -- nonunique pattern with a unique value. - if orderZeroSV sv - then M.singleton vn $ Dynamic t - else M.singleton vn sv -matchPatternSV (Wildcard _ _) _ = mempty -matchPatternSV (PatternAscription pat _ _) sv = matchPatternSV pat sv -matchPatternSV PatternLit{} _ = mempty -matchPatternSV (PatternConstr c1 _ ps _) (SumSV c2 ls fs) - | c1 == c2 = - mconcat $ zipWith matchPatternSV ps ls - | Just ts <- lookup c1 fs = - mconcat $ zipWith matchPatternSV ps $ map svFromType ts - | otherwise = - error $ "matchPatternSV: missing constructor in type: " ++ pretty c1 -matchPatternSV (PatternConstr c1 _ ps _) (Dynamic (Scalar (Sum fs))) - | Just ts <- M.lookup c1 fs = - mconcat $ zipWith matchPatternSV ps $ map svFromType ts - | otherwise = - error $ "matchPatternSV: missing constructor in type: " ++ pretty c1 -matchPatternSV pat (Dynamic t) = matchPatternSV pat $ svFromType t -matchPatternSV pat sv = error $ "Tried to match pattern " ++ pretty pat - ++ " with static value " ++ show sv ++ "." - -orderZeroSV :: StaticVal -> Bool -orderZeroSV Dynamic{} = True -orderZeroSV (RecordSV fields) = all (orderZeroSV . snd) fields -orderZeroSV _ = False - --- | Given a pattern and the static value for the defunctionalized argument, --- update the pattern to reflect the changes in the types. -updatePattern :: Pattern -> StaticVal -> Pattern -updatePattern (TuplePattern ps loc) (RecordSV svs) = - TuplePattern (zipWith updatePattern ps $ map snd svs) loc -updatePattern (RecordPattern ps loc) (RecordSV svs) - | ps' <- sortOn fst ps, svs' <- sortOn fst svs = - RecordPattern (zipWith (\(n, p) (_, sv) -> - (n, updatePattern p sv)) ps' svs') loc -updatePattern (PatternParens pat loc) sv = - PatternParens (updatePattern pat sv) loc -updatePattern (Id vn (Info tp) loc) sv = - Id vn (Info $ comb tp (snd (typeFromSV sv) `setUniqueness` Nonunique)) loc - -- Preserve any original zeroth-order types. - where comb (Scalar Arrow{}) t2 = t2 - comb (Scalar (Record m1)) (Scalar (Record m2)) = - Scalar $ Record $ M.intersectionWith comb m1 m2 - comb (Scalar (Sum m1)) (Scalar (Sum m2)) = - Scalar $ Sum $ M.intersectionWith (zipWith comb) m1 m2 - comb t1 _ = t1 -- t1 must be array or prim. -updatePattern pat@(Wildcard (Info tp) loc) sv - | orderZero tp = pat - | otherwise = Wildcard (Info $ snd $ typeFromSV sv) loc -updatePattern (PatternAscription pat tydecl loc) sv - | orderZero . unInfo $ expandedType tydecl = - PatternAscription (updatePattern pat sv) tydecl loc - | otherwise = updatePattern pat sv -updatePattern p@PatternLit{} _ = p -updatePattern pat@(PatternConstr c1 (Info t) ps loc) sv@(SumSV _ svs _) - | orderZero t = pat - | otherwise = PatternConstr c1 (Info t') ps' loc - where t' = snd (typeFromSV sv) `setUniqueness` Nonunique - ps' = zipWith updatePattern ps svs -updatePattern (PatternConstr c1 _ ps loc) (Dynamic t) = - PatternConstr c1 (Info t) ps loc -updatePattern pat (Dynamic t) = updatePattern pat (svFromType t) -updatePattern pat sv = - error $ "Tried to update pattern " ++ pretty pat - ++ "to reflect the static value " ++ show sv - --- Like updatePattern, but discard sizes. This is used for --- let-bindings, where we might otherwise introduce sizes that are --- free. -updatePattern' :: Pattern -> StaticVal -> Pattern -updatePattern' pat sv = - let pat' = updatePattern pat sv - (sizes, _) = typeFromSV sv - tr = identityMapper { mapOnPatternType = - pure . unscopeType (S.fromList sizes) - } - in runIdentity $ astMap tr pat' - --- | Convert a record (or tuple) type to a record static value. This is used for --- "unwrapping" tuples and records that are nested in 'Dynamic' static values. -svFromType :: PatternType -> StaticVal -svFromType (Scalar (Record fs)) = RecordSV . M.toList $ M.map svFromType fs -svFromType t = Dynamic t - --- A set of names where we also track uniqueness. -newtype NameSet = NameSet (M.Map VName Uniqueness) deriving (Show) - -instance Semigroup NameSet where - NameSet x <> NameSet y = NameSet $ M.unionWith max x y - -instance Monoid NameSet where - mempty = NameSet mempty - -without :: NameSet -> NameSet -> NameSet -without (NameSet x) (NameSet y) = NameSet $ x `M.difference` y - -member :: VName -> NameSet -> Bool -member v (NameSet m) = v `M.member` m - -ident :: Ident -> NameSet -ident v = NameSet $ M.singleton (identName v) (uniqueness $ unInfo $ identType v) - -oneName :: VName -> NameSet -oneName v = NameSet $ M.singleton v Nonunique - -names :: S.Set VName -> NameSet -names = foldMap oneName - --- | Compute the set of free variables of an expression. -freeVars :: Exp -> NameSet -freeVars expr = case expr of - Literal{} -> mempty - IntLit{} -> mempty - FloatLit{} -> mempty - StringLit{} -> mempty - Parens e _ -> freeVars e - QualParens _ e _ -> freeVars e - TupLit es _ -> foldMap freeVars es - - RecordLit fs _ -> foldMap freeVarsField fs - where freeVarsField (RecordFieldExplicit _ e _) = freeVars e - freeVarsField (RecordFieldImplicit vn t _) = ident $ Ident vn t mempty - - ArrayLit es t _ -> foldMap freeVars es <> - names (typeDimNames $ unInfo t) - Range e me incl _ _ -> freeVars e <> foldMap freeVars me <> - foldMap freeVars incl - Var qn (Info t) _ -> NameSet $ M.singleton (qualLeaf qn) $ uniqueness t - Ascript e t _ -> freeVars e <> names (typeDimNames $ unInfo $ expandedType t) - Coerce e t _ _ -> freeVars e <> names (typeDimNames $ unInfo $ expandedType t) - LetPat pat e1 e2 _ _ -> freeVars e1 <> ((names (patternDimNames pat) <> freeVars e2) - `without` patternVars pat) - - LetFun vn (_, pats, _, _, e1) e2 _ _ -> - ((freeVars e1 <> names (foldMap patternDimNames pats)) - `without` foldMap patternVars pats) <> - (freeVars e2 `without` oneName vn) - - If e1 e2 e3 _ _ -> freeVars e1 <> freeVars e2 <> freeVars e3 - Apply e1 e2 _ _ _ -> freeVars e1 <> freeVars e2 - Negate e _ -> freeVars e - Lambda pats e0 _ _ _ -> (names (foldMap patternDimNames pats) <> freeVars e0) - `without` foldMap patternVars pats - OpSection{} -> mempty - OpSectionLeft _ _ e _ _ _ -> freeVars e - OpSectionRight _ _ e _ _ _ -> freeVars e - ProjectSection{} -> mempty - IndexSection idxs _ _ -> foldMap freeDimIndex idxs - - DoLoop sparams pat e1 form e3 _ _ -> - let (e2fv, e2ident) = formVars form - in freeVars e1 <> e2fv <> - (freeVars e3 `without` - (names (S.fromList sparams) <> patternVars pat <> e2ident)) - where formVars (For v e2) = (freeVars e2, ident v) - formVars (ForIn p e2) = (freeVars e2, patternVars p) - formVars (While e2) = (freeVars e2, mempty) - - BinOp (qn, _) _ (e1, _) (e2, _) _ _ _ -> oneName (qualLeaf qn) <> - freeVars e1 <> freeVars e2 - Project _ e _ _ -> freeVars e - - LetWith id1 id2 idxs e1 e2 _ _ -> - ident id2 <> foldMap freeDimIndex idxs <> freeVars e1 <> - (freeVars e2 `without` ident id1) - - Index e idxs _ _ -> freeVars e <> foldMap freeDimIndex idxs - Update e1 idxs e2 _ -> freeVars e1 <> foldMap freeDimIndex idxs <> freeVars e2 - RecordUpdate e1 _ e2 _ _ -> freeVars e1 <> freeVars e2 - - Assert e1 e2 _ _ -> freeVars e1 <> freeVars e2 - Constr _ es _ _ -> foldMap freeVars es - Attr _ e _ -> freeVars e - Match e cs _ _ -> freeVars e <> foldMap caseFV cs - where caseFV (CasePat p eCase _) = (names (patternDimNames p) <> freeVars eCase) - `without` patternVars p - -freeDimIndex :: DimIndexBase Info VName -> NameSet -freeDimIndex (DimFix e) = freeVars e -freeDimIndex (DimSlice me1 me2 me3) = - foldMap (foldMap freeVars) [me1, me2, me3] - --- | Extract all the variable names bound in a pattern. -patternVars :: Pattern -> NameSet -patternVars = mconcat . map ident . S.toList . patternIdents - --- | Defunctionalize a top-level value binding. Returns the --- transformed result as well as an environment that binds the name of --- the value binding to the static value of the transformed body. The --- boolean is true if the function is a 'DynamicFun'. -defuncValBind :: ValBind -> DefM (ValBind, Env, Bool) - --- Eta-expand entry points with a functional return type. -defuncValBind (ValBind entry name _ (Info (rettype, retext)) tparams params body _ attrs loc) - | Scalar Arrow{} <- rettype = do - (body_pats, body', rettype') <- etaExpand (fromStruct rettype) body - defuncValBind $ ValBind entry name Nothing - (Info (rettype', retext)) - tparams (params <> body_pats) body' Nothing attrs loc - -defuncValBind valbind@(ValBind _ name retdecl (Info (rettype, retext)) tparams params body _ _ _) = do - (tparams', params', body', sv) <- defuncLet tparams params body rettype - let rettype' = combineTypeShapes rettype $ anySizes $ toStruct $ typeOf body' - (missing_dims, params'') <- sizesForAll params' - return ( valbind { valBindRetDecl = retdecl - , valBindRetType = Info (if null params' - then rettype' `setUniqueness` Nonunique - else rettype', - retext) - , valBindTypeParams = tparams' ++ - map (`TypeParamDim` mempty) missing_dims - , valBindParams = params'' - , valBindBody = body' - } - , M.singleton name sv - , case sv of DynamicFun{} -> True - Dynamic{} -> True - _ -> False) - --- | Defunctionalize a list of top-level declarations. -defuncVals :: [ValBind] -> DefM (Seq.Seq ValBind) -defuncVals [] = return mempty -defuncVals (valbind : ds) = do - ((valbind', env, dyn), defs) <- collectFuns $ defuncValBind valbind - ds' <- localEnv env $ if dyn - then isGlobal (valBindName valbind') $ defuncVals ds - else defuncVals ds - return $ defs <> Seq.singleton valbind' <> ds' - --- | Transform a list of top-level value bindings. May produce new --- lifted function definitions, which are placed in front of the --- resulting list of declarations. -transformProg :: MonadFreshNames m => [ValBind] -> m [ValBind] -transformProg decs = modifyNameSource $ \namesrc -> - let (decs', namesrc', liftedDecs) = runDefM namesrc $ defuncVals decs - in (toList $ liftedDecs <> decs', namesrc') +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE Trustworthy #-} + +-- | Defunctionalization of typed, monomorphic Futhark programs without modules. +module Futhark.Internalise.Defunctionalise (transformProg) where + +import qualified Control.Arrow as Arrow +import Control.Monad.Identity +import Control.Monad.RWS hiding (Sum) +import Control.Monad.State +import Data.Bifunctor +import Data.Bitraversable +import Data.Foldable +import Data.List (nub, partition, sortOn, tails) +import qualified Data.List.NonEmpty as NE +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Sequence as Seq +import qualified Data.Set as S +import Futhark.IR.Pretty () +import Futhark.MonadFreshNames +import Language.Futhark +import Language.Futhark.Traversals + +-- | An expression or an extended 'Lambda' (with size parameters, +-- which AST lambdas do not support). +data ExtExp + = ExtLambda [TypeParam] [Pattern] Exp (Aliasing, StructType) SrcLoc + | ExtExp Exp + deriving (Show) + +-- | A static value stores additional information about the result of +-- defunctionalization of an expression, aside from the residual expression. +data StaticVal + = Dynamic PatternType + | -- | The 'VName's are shape parameters that are bound + -- by the 'Pattern'. + LambdaSV [VName] Pattern StructType ExtExp Env + | RecordSV [(Name, StaticVal)] + | -- | The constructor that is actually present, plus + -- the others that are not. + SumSV Name [StaticVal] [(Name, [PatternType])] + | DynamicFun (Exp, StaticVal) StaticVal + | IntrinsicSV + deriving (Show) + +-- | Environment mapping variable names to their associated static value. +type Env = M.Map VName StaticVal + +localEnv :: Env -> DefM a -> DefM a +localEnv env = local $ Arrow.second (env <>) + +-- Even when using a "new" environment (for evaluating closures) we +-- still ram the global environment of DynamicFuns in there. +localNewEnv :: Env -> DefM a -> DefM a +localNewEnv env = local $ \(globals, old_env) -> + (globals, M.filterWithKey (\k _ -> k `S.member` globals) old_env <> env) + +extendEnv :: VName -> StaticVal -> DefM a -> DefM a +extendEnv vn sv = localEnv (M.singleton vn sv) + +askEnv :: DefM Env +askEnv = asks snd + +isGlobal :: VName -> DefM a -> DefM a +isGlobal v = local $ Arrow.first (S.insert v) + +-- | Returns the defunctionalization environment restricted +-- to the given set of variable names and types. +restrictEnvTo :: NameSet -> DefM Env +restrictEnvTo (NameSet m) = restrict <$> ask + where + restrict (globals, env) = M.mapMaybeWithKey keep env + where + keep k sv = do + guard $ not $ k `S.member` globals + u <- M.lookup k m + Just $ restrict' u sv + restrict' Nonunique (Dynamic t) = + Dynamic $ t `setUniqueness` Nonunique + restrict' _ (Dynamic t) = + Dynamic t + restrict' u (LambdaSV dims pat t e env) = + LambdaSV dims pat t e $ M.map (restrict' u) env + restrict' u (RecordSV fields) = + RecordSV $ map (fmap $ restrict' u) fields + restrict' u (SumSV c svs fields) = + SumSV c (map (restrict' u) svs) fields + restrict' u (DynamicFun (e, sv1) sv2) = + DynamicFun (e, restrict' u sv1) $ restrict' u sv2 + restrict' _ IntrinsicSV = IntrinsicSV + +-- | Defunctionalization monad. The Reader environment tracks both +-- the current Env as well as the set of globally defined dynamic +-- functions. This is used to avoid unnecessarily large closure +-- environments. +newtype DefM a = DefM (RWS (S.Set VName, Env) (Seq.Seq ValBind) VNameSource a) + deriving + ( Functor, + Applicative, + Monad, + MonadReader (S.Set VName, Env), + MonadWriter (Seq.Seq ValBind), + MonadFreshNames + ) + +-- | Run a computation in the defunctionalization monad. Returns the result of +-- the computation, a new name source, and a list of lifted function declations. +runDefM :: VNameSource -> DefM a -> (a, VNameSource, Seq.Seq ValBind) +runDefM src (DefM m) = runRWS m mempty src + +collectFuns :: DefM a -> DefM (a, Seq.Seq ValBind) +collectFuns m = pass $ do + (x, decs) <- listen m + return ((x, decs), const mempty) + +-- | Looks up the associated static value for a given name in the environment. +lookupVar :: SrcLoc -> VName -> DefM StaticVal +lookupVar loc x = do + env <- askEnv + case M.lookup x env of + Just sv -> return sv + Nothing -- If the variable is unknown, it may refer to the 'intrinsics' + -- module, which we will have to treat specially. + | baseTag x <= maxIntrinsicTag -> return IntrinsicSV + | otherwise -> -- Anything not in scope is going to be an + -- existential size. + return $ Dynamic $ Scalar $ Prim $ Signed Int64 + | otherwise -> + error $ + "Variable " ++ pretty x ++ " at " + ++ locStr loc + ++ " is out of scope." + +-- Like patternDimNames, but ignores sizes that are only found in +-- funtion types. +arraySizes :: StructType -> S.Set VName +arraySizes (Scalar Arrow {}) = mempty +arraySizes (Scalar (Record fields)) = foldMap arraySizes fields +arraySizes (Scalar (Sum cs)) = foldMap (foldMap arraySizes) cs +arraySizes (Scalar (TypeVar _ _ _ targs)) = + mconcat $ map f targs + where + f (TypeArgDim (NamedDim d) _) = S.singleton $ qualLeaf d + f TypeArgDim {} = mempty + f (TypeArgType t _) = arraySizes t +arraySizes (Scalar Prim {}) = mempty +arraySizes (Array _ _ t shape) = + arraySizes (Scalar t) <> foldMap dimName (shapeDims shape) + where + dimName :: DimDecl VName -> S.Set VName + dimName (NamedDim qn) = S.singleton $ qualLeaf qn + dimName _ = mempty + +patternArraySizes :: Pattern -> S.Set VName +patternArraySizes = arraySizes . patternStructType + +dimMapping :: + Monoid a => + TypeBase (DimDecl VName) a -> + TypeBase (DimDecl VName) a -> + M.Map VName VName +dimMapping t1 t2 = execState (matchDims f t1 t2) mempty + where + f (NamedDim d1) (NamedDim d2) = do + modify $ M.insert (qualLeaf d1) (qualLeaf d2) + return $ NamedDim d1 + f d _ = return d + +defuncFun :: + [TypeParam] -> + [Pattern] -> + Exp -> + (Aliasing, StructType) -> + SrcLoc -> + DefM (Exp, StaticVal) +defuncFun tparams pats e0 (closure, ret) loc = do + when (any isTypeParam tparams) $ + error $ + "Received a lambda with type parameters at " ++ locStr loc + ++ ", but the defunctionalizer expects a monomorphic input program." + -- Extract the first parameter of the lambda and "push" the + -- remaining ones (if there are any) into the body of the lambda. + let (dims, pat, ret', e0') = case pats of + [] -> error "Received a lambda with no parameters." + [pat'] -> (map typeParamName tparams, pat', ret, ExtExp e0) + (pat' : pats') -> + -- Split shape parameters into those that are determined by + -- the first pattern, and those that are determined by later + -- patterns. + let bound_by_pat = (`S.member` patternArraySizes pat') . typeParamName + (pat_dims, rest_dims) = partition bound_by_pat tparams + in ( map typeParamName pat_dims, + pat', + foldFunType (map (toStruct . patternType) pats') ret, + ExtLambda rest_dims pats' e0 (closure, ret) loc + ) + + -- Construct a record literal that closes over the environment of + -- the lambda. Closed-over 'DynamicFun's are converted to their + -- closure representation. + let used = + freeVars (Lambda pats e0 Nothing (Info (closure, ret)) loc) + `without` mconcat (map oneName dims) + used_env <- restrictEnvTo used + + -- The closure parts that are sizes are proactively turned into size + -- parameters. + let sizes_of_arrays = + foldMap (arraySizes . toStruct . typeFromSV') used_env + <> patternArraySizes pat + notSize = not . (`S.member` sizes_of_arrays) + (fields, env) = + unzip $ + map closureFromDynamicFun $ + filter (notSize . fst) $ M.toList used_env + env' = M.fromList env + closure_dims = S.toList sizes_of_arrays + + global <- asks fst + + return + ( RecordLit fields loc, + LambdaSV + ( nub $ + filter (`S.notMember` global) $ + dims <> closure_dims + ) + pat + ret' + e0' + env' + ) + where + closureFromDynamicFun (vn, DynamicFun (clsr_env, sv) _) = + let name = nameFromString $ pretty vn + in (RecordFieldExplicit name clsr_env mempty, (vn, sv)) + closureFromDynamicFun (vn, sv) = + let name = nameFromString $ pretty vn + tp' = typeFromSV' sv + in ( RecordFieldExplicit + name + (Var (qualName vn) (Info tp') mempty) + mempty, + (vn, sv) + ) + +-- | Defunctionalization of an expression. Returns the residual expression and +-- the associated static value in the defunctionalization monad. +defuncExp :: Exp -> DefM (Exp, StaticVal) +defuncExp e@Literal {} = + return (e, Dynamic $ typeOf e) +defuncExp e@IntLit {} = + return (e, Dynamic $ typeOf e) +defuncExp e@FloatLit {} = + return (e, Dynamic $ typeOf e) +defuncExp e@StringLit {} = + return (e, Dynamic $ typeOf e) +defuncExp (Parens e loc) = do + (e', sv) <- defuncExp e + return (Parens e' loc, sv) +defuncExp (QualParens qn e loc) = do + (e', sv) <- defuncExp e + return (QualParens qn e' loc, sv) +defuncExp (TupLit es loc) = do + (es', svs) <- unzip <$> mapM defuncExp es + return (TupLit es' loc, RecordSV $ zip tupleFieldNames svs) +defuncExp (RecordLit fs loc) = do + (fs', names_svs) <- unzip <$> mapM defuncField fs + return (RecordLit fs' loc, RecordSV names_svs) + where + defuncField (RecordFieldExplicit vn e loc') = do + (e', sv) <- defuncExp e + return (RecordFieldExplicit vn e' loc', (vn, sv)) + defuncField (RecordFieldImplicit vn _ loc') = do + sv <- lookupVar loc' vn + case sv of + -- If the implicit field refers to a dynamic function, we + -- convert it to an explicit field with a record closing over + -- the environment and bind the corresponding static value. + DynamicFun (e, sv') _ -> + let vn' = baseName vn + in return + ( RecordFieldExplicit vn' e loc', + (vn', sv') + ) + -- The field may refer to a functional expression, so we get the + -- type from the static value and not the one from the AST. + _ -> + let tp = Info $ typeFromSV' sv + in return (RecordFieldImplicit vn tp loc', (baseName vn, sv)) +defuncExp (ArrayLit es t@(Info t') loc) = do + es' <- mapM defuncExp' es + return (ArrayLit es' t loc, Dynamic t') +defuncExp (Range e1 me incl t@(Info t', _) loc) = do + e1' <- defuncExp' e1 + me' <- mapM defuncExp' me + incl' <- mapM defuncExp' incl + return (Range e1' me' incl' t loc, Dynamic t') +defuncExp e@(Var qn _ loc) = do + sv <- lookupVar loc (qualLeaf qn) + case sv of + -- If the variable refers to a dynamic function, we return its closure + -- representation (i.e., a record expression capturing the free variables + -- and a 'LambdaSV' static value) instead of the variable itself. + DynamicFun closure _ -> return closure + -- Intrinsic functions used as variables are eta-expanded, so we + -- can get rid of them. + IntrinsicSV -> do + (pats, body, tp) <- etaExpand (typeOf e) e + defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty + _ -> + let tp = typeFromSV' sv + in return (Var qn (Info tp) loc, sv) +defuncExp (Ascript e0 tydecl loc) + | orderZero (typeOf e0) = do + (e0', sv) <- defuncExp e0 + return (Ascript e0' tydecl loc, sv) + | otherwise = defuncExp e0 +defuncExp (Coerce e0 tydecl t loc) + | orderZero (typeOf e0) = do + (e0', sv) <- defuncExp e0 + return (Coerce e0' tydecl t loc, sv) + | otherwise = defuncExp e0 +defuncExp (LetPat pat e1 e2 (Info t, retext) loc) = do + (e1', sv1) <- defuncExp e1 + let env = matchPatternSV pat sv1 + pat' = updatePattern' pat sv1 + (e2', sv2) <- localEnv env $ defuncExp e2 + -- To maintain any sizes going out of scope, we need to compute the + -- old size substitution induced by retext and also apply it to the + -- newly computed body type. + let mapping = dimMapping (typeOf e2) t + subst v = fromMaybe v $ M.lookup v mapping + t' = first (fmap subst) $ typeOf e2' + return (LetPat pat' e1' e2' (Info t', retext) loc, sv2) + +-- Local functions are handled by rewriting them to lambdas, so that +-- the same machinery can be re-used. But we may have to eta-expand +-- first. +defuncExp (LetFun vn (dims, pats, _, Info ret, e1) e2 let_t loc) + | Scalar Arrow {} <- ret = do + (body_pats, e1', ret') <- etaExpand (fromStruct ret) e1 + let f = (dims, pats <> body_pats, Nothing, Info ret', e1') + defuncExp $ LetFun vn f e2 let_t loc + | otherwise = do + (e1', sv1) <- defuncFun dims pats e1 (mempty, ret) loc + (e2', sv2) <- localEnv (M.singleton vn sv1) $ defuncExp e2 + return + ( LetPat (Id vn (Info (typeOf e1')) loc) e1' e2' (Info $ typeOf e2', Info []) loc, + sv2 + ) +defuncExp (If e1 e2 e3 tp loc) = do + (e1', _) <- defuncExp e1 + (e2', sv) <- defuncExp e2 + (e3', _) <- defuncExp e3 + return (If e1' e2' e3' tp loc, sv) +defuncExp e@(Apply f@(Var f' _ _) arg d (t, ext) loc) + | baseTag (qualLeaf f') <= maxIntrinsicTag, + TupLit es tuploc <- arg = do + -- defuncSoacExp also works fine for non-SOACs. + es' <- mapM defuncSoacExp es + return + ( Apply f (TupLit es' tuploc) d (t, ext) loc, + Dynamic $ typeOf e + ) +defuncExp e@Apply {} = defuncApply 0 e +defuncExp (Negate e0 loc) = do + (e0', sv) <- defuncExp e0 + return (Negate e0' loc, sv) +defuncExp (Lambda pats e0 _ (Info (closure, ret)) loc) = + defuncFun [] pats e0 (closure, ret) loc +-- Operator sections are expected to be converted to lambda-expressions +-- by the monomorphizer, so they should no longer occur at this point. +defuncExp OpSection {} = error "defuncExp: unexpected operator section." +defuncExp OpSectionLeft {} = error "defuncExp: unexpected operator section." +defuncExp OpSectionRight {} = error "defuncExp: unexpected operator section." +defuncExp ProjectSection {} = error "defuncExp: unexpected projection section." +defuncExp IndexSection {} = error "defuncExp: unexpected projection section." +defuncExp (DoLoop sparams pat e1 form e3 ret loc) = do + (e1', sv1) <- defuncExp e1 + let env1 = matchPatternSV pat sv1 + (form', env2) <- case form of + For v e2 -> do + e2' <- defuncExp' e2 + return (For v e2', envFromIdent v) + ForIn pat2 e2 -> do + e2' <- defuncExp' e2 + return (ForIn pat2 e2', envFromPattern pat2) + While e2 -> do + e2' <- localEnv env1 $ defuncExp' e2 + return (While e2', mempty) + (e3', sv) <- localEnv (env1 <> env2) $ defuncExp e3 + return (DoLoop sparams pat e1' form' e3' ret loc, sv) + where + envFromIdent (Ident vn (Info tp) _) = + M.singleton vn $ Dynamic tp + +-- We handle BinOps by turning them into ordinary function applications. +defuncExp + ( BinOp + (qn, qnloc) + (Info t) + (e1, Info (pt1, ext1)) + (e2, Info (pt2, ext2)) + (Info ret) + (Info retext) + loc + ) = + defuncExp $ + Apply + ( Apply + (Var qn (Info t) qnloc) + e1 + (Info (diet pt1, ext1)) + (Info (Scalar $ Arrow mempty Unnamed (fromStruct pt2) ret), Info []) + loc + ) + e2 + (Info (diet pt2, ext2)) + (Info ret, Info retext) + loc +defuncExp (Project vn e0 tp@(Info tp') loc) = do + (e0', sv0) <- defuncExp e0 + case sv0 of + RecordSV svs -> case lookup vn svs of + Just sv -> return (Project vn e0' (Info $ typeFromSV' sv) loc, sv) + Nothing -> error "Invalid record projection." + Dynamic _ -> return (Project vn e0' tp loc, Dynamic tp') + _ -> error $ "Projection of an expression with static value " ++ show sv0 +defuncExp (LetWith id1 id2 idxs e1 body t loc) = do + e1' <- defuncExp' e1 + sv1 <- lookupVar (identSrcLoc id2) $ identName id2 + idxs' <- mapM defuncDimIndex idxs + (body', sv) <- extendEnv (identName id1) sv1 $ defuncExp body + return (LetWith id1 id2 idxs' e1' body' t loc, sv) +defuncExp expr@(Index e0 idxs info loc) = do + e0' <- defuncExp' e0 + idxs' <- mapM defuncDimIndex idxs + return (Index e0' idxs' info loc, Dynamic $ typeOf expr) +defuncExp (Update e1 idxs e2 loc) = do + (e1', sv) <- defuncExp e1 + idxs' <- mapM defuncDimIndex idxs + e2' <- defuncExp' e2 + return (Update e1' idxs' e2' loc, sv) + +-- Note that we might change the type of the record field here. This +-- is not permitted in the type checker due to problems with type +-- inference, but it actually works fine. +defuncExp (RecordUpdate e1 fs e2 _ loc) = do + (e1', sv1) <- defuncExp e1 + (e2', sv2) <- defuncExp e2 + let sv = staticField sv1 sv2 fs + return + ( RecordUpdate e1' fs e2' (Info $ typeFromSV' sv1) loc, + sv + ) + where + staticField (RecordSV svs) sv2 (f : fs') = + case lookup f svs of + Just sv -> + RecordSV $ + (f, staticField sv sv2 fs') : filter ((/= f) . fst) svs + Nothing -> error "Invalid record projection." + staticField (Dynamic t@(Scalar Record {})) sv2 fs'@(_ : _) = + staticField (svFromType t) sv2 fs' + staticField _ sv2 _ = sv2 +defuncExp (Assert e1 e2 desc loc) = do + (e1', _) <- defuncExp e1 + (e2', sv) <- defuncExp e2 + return (Assert e1' e2' desc loc, sv) +defuncExp (Constr name es (Info (Scalar (Sum all_fs))) loc) = do + (es', svs) <- unzip <$> mapM defuncExp es + let sv = + SumSV name svs $ + M.toList $ + name `M.delete` M.map (map defuncType) all_fs + return (Constr name es' (Info (typeFromSV' sv)) loc, sv) + where + defuncType :: + Monoid als => + TypeBase (DimDecl VName) als -> + TypeBase (DimDecl VName) als + defuncType (Array as u t shape) = Array as u (defuncScalar t) shape + defuncType (Scalar t) = Scalar $ defuncScalar t + + defuncScalar :: + Monoid als => + ScalarTypeBase (DimDecl VName) als -> + ScalarTypeBase (DimDecl VName) als + defuncScalar (Record fs) = Record $ M.map defuncType fs + defuncScalar Arrow {} = Record mempty + defuncScalar (Sum fs) = Sum $ M.map (map defuncType) fs + defuncScalar (Prim t) = Prim t + defuncScalar (TypeVar as u tn targs) = TypeVar as u tn targs +defuncExp (Constr name _ (Info t) loc) = + error $ + "Constructor " ++ pretty name ++ " given type " + ++ pretty t + ++ " at " + ++ locStr loc +defuncExp (Match e cs t loc) = do + (e', sv) <- defuncExp e + csPairs <- mapM (defuncCase sv) cs + let cs' = fmap fst csPairs + sv' = snd $ NE.head csPairs + return (Match e' cs' t loc, sv') +defuncExp (Attr info e loc) = do + (e', sv) <- defuncExp e + return (Attr info e' loc, sv) + +-- | Same as 'defuncExp', except it ignores the static value. +defuncExp' :: Exp -> DefM Exp +defuncExp' = fmap fst . defuncExp + +defuncExtExp :: ExtExp -> DefM (Exp, StaticVal) +defuncExtExp (ExtExp e) = defuncExp e +defuncExtExp (ExtLambda tparams pats e0 (closure, ret) loc) = + defuncFun tparams pats e0 (closure, ret) loc + +defuncCase :: StaticVal -> Case -> DefM (Case, StaticVal) +defuncCase sv (CasePat p e loc) = do + let p' = updatePattern' p sv + env = matchPatternSV p sv + (e', sv') <- localEnv env $ defuncExp e + return (CasePat p' e' loc, sv') + +-- | Defunctionalize the function argument to a SOAC by eta-expanding if +-- necessary and then defunctionalizing the body of the introduced lambda. +defuncSoacExp :: Exp -> DefM Exp +defuncSoacExp e@OpSection {} = return e +defuncSoacExp e@OpSectionLeft {} = return e +defuncSoacExp e@OpSectionRight {} = return e +defuncSoacExp e@ProjectSection {} = return e +defuncSoacExp (Parens e loc) = + Parens <$> defuncSoacExp e <*> pure loc +defuncSoacExp (Lambda params e0 decl tp loc) = do + let env = foldMap envFromPattern params + e0' <- localEnv env $ defuncSoacExp e0 + return $ Lambda params e0' decl tp loc +defuncSoacExp e + | Scalar Arrow {} <- typeOf e = do + (pats, body, tp) <- etaExpand (typeOf e) e + let env = foldMap envFromPattern pats + body' <- localEnv env $ defuncExp' body + return $ Lambda pats body' Nothing (Info (mempty, tp)) mempty + | otherwise = defuncExp' e + +etaExpand :: PatternType -> Exp -> DefM ([Pattern], Exp, StructType) +etaExpand e_t e = do + let (ps, ret) = getType e_t + (pats, vars) <- fmap unzip . forM ps $ \(p, t) -> do + x <- case p of + Named x -> pure x + Unnamed -> newNameFromString "x" + return + ( Id x (Info t) mempty, + Var (qualName x) (Info t) mempty + ) + let e' = + foldl' + ( \e1 (e2, t2, argtypes) -> + Apply + e1 + e2 + (Info (diet t2, Nothing)) + (Info (foldFunType argtypes ret), Info []) + mempty + ) + e + $ zip3 vars (map snd ps) (drop 1 $ tails $ map snd ps) + return (pats, e', toStruct ret) + where + getType (Scalar (Arrow _ p t1 t2)) = + let (ps, r) = getType t2 in ((p, t1) : ps, r) + getType t = ([], t) + +-- | Defunctionalize an indexing of a single array dimension. +defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName) +defuncDimIndex (DimFix e1) = DimFix . fst <$> defuncExp e1 +defuncDimIndex (DimSlice me1 me2 me3) = + DimSlice <$> defunc' me1 <*> defunc' me2 <*> defunc' me3 + where + defunc' = mapM defuncExp' + +-- | Defunctionalize a let-bound function, while preserving parameters +-- that have order 0 types (i.e., non-functional). +defuncLet :: + [TypeParam] -> + [Pattern] -> + Exp -> + StructType -> + DefM ([TypeParam], [Pattern], Exp, StaticVal) +defuncLet dims ps@(pat : pats) body rettype + | patternOrderZero pat = do + let bound_by_pat = (`S.member` patternDimNames pat) . typeParamName + -- Take care to not include more size parameters than necessary. + (pat_dims, rest_dims) = partition bound_by_pat dims + env = envFromPattern pat <> envFromShapeParams pat_dims + (rest_dims', pats', body', sv) <- localEnv env $ defuncLet rest_dims pats body rettype + closure <- defuncFun dims ps body (mempty, rettype) mempty + return (pat_dims ++ rest_dims', pat : pats', body', DynamicFun closure sv) + | otherwise = do + (e, sv) <- defuncFun dims ps body (mempty, rettype) mempty + return ([], [], e, sv) +defuncLet _ [] body rettype = do + (body', sv) <- defuncExp body + return ([], [], body', imposeType sv rettype) + where + imposeType Dynamic {} t = + Dynamic $ fromStruct t + imposeType (RecordSV fs1) (Scalar (Record fs2)) = + RecordSV $ M.toList $ M.intersectionWith imposeType (M.fromList fs1) fs2 + imposeType sv _ = sv + +sizesForAll :: MonadFreshNames m => [Pattern] -> m ([VName], [Pattern]) +sizesForAll params = do + (params', sizes) <- runStateT (mapM (astMap tv) params) [] + return (sizes, params') + where + tv = identityMapper {mapOnPatternType = bitraverse onDim pure} + onDim AnyDim = do + v <- lift $ newVName "size" + modify (v :) + pure $ NamedDim $ qualName v + onDim d = pure d + +-- | Defunctionalize an application expression at a given depth of application. +-- Calls to dynamic (first-order) functions are preserved at much as possible, +-- but a new lifted function is created if a dynamic function is only partially +-- applied. +defuncApply :: Int -> Exp -> DefM (Exp, StaticVal) +defuncApply depth e@(Apply e1 e2 d t@(Info ret, Info ext) loc) = do + let (argtypes, _) = unfoldFunType ret + (e1', sv1) <- defuncApply (depth + 1) e1 + (e2', sv2) <- defuncExp e2 + let e' = Apply e1' e2' d t loc + case sv1 of + LambdaSV dims pat e0_t e0 closure_env -> do + let env' = matchPatternSV pat sv2 + env_dim = envFromDimNames dims + (e0', sv) <- localNewEnv (env' <> closure_env <> env_dim) $ defuncExtExp e0 + + let closure_pat = buildEnvPattern closure_env + pat' = updatePattern pat sv2 + + globals <- asks fst + + -- Lift lambda to top-level function definition. We put in + -- a lot of effort to try to infer the uniqueness attributes + -- of the lifted function, but this is ultimately all a sham + -- and a hack. There is some piece we're missing. + let params = [closure_pat, pat'] + params_for_rettype = params ++ svParams sv1 ++ svParams sv2 + svParams (LambdaSV _ sv_pat _ _ _) = [sv_pat] + svParams _ = [] + rettype = buildRetType closure_env params_for_rettype e0_t $ typeOf e0' + + already_bound = + globals <> S.fromList dims + <> S.map identName (foldMap patternIdents params) + more_dims = + S.toList $ + S.filter (`S.notMember` already_bound) $ + foldMap patternArraySizes params + + -- Embed some information about the original function + -- into the name of the lifted function, to make the + -- result slightly more human-readable. + liftedName i (Var f _ _) = + "lifted_" ++ show i ++ "_" ++ baseString (qualLeaf f) + liftedName i (Apply f _ _ _ _) = + liftedName (i + 1) f + liftedName _ _ = "lifted" + + -- Ensure that no parameter sizes are AnyDim. The internaliser + -- expects this. This is easy, because they are all + -- first-order. + (missing_dims, params') <- sizesForAll params + + fname <- newNameFromString $ liftedName (0 :: Int) e1 + liftValDec + fname + rettype + (dims ++ more_dims ++ missing_dims) + params' + e0' + + let t1 = toStruct $ typeOf e1' + t2 = toStruct $ typeOf e2' + fname' = qualName fname + fname'' = + Var + fname' + ( Info + ( Scalar $ + Arrow mempty Unnamed (fromStruct t1) $ + Scalar $ Arrow mempty Unnamed (fromStruct t2) rettype + ) + ) + loc + + -- FIXME: what if this application returns both a function + -- and a value? + callret + | orderZero ret = (Info ret, Info ext) + | otherwise = (Info rettype, Info ext) + + return + ( Parens + ( Apply + ( Apply + fname'' + e1' + (Info (Observe, Nothing)) + ( Info $ Scalar $ Arrow mempty Unnamed (fromStruct t2) rettype, + Info [] + ) + loc + ) + e2' + d + callret + loc + ) + mempty, + sv + ) + + -- If e1 is a dynamic function, we just leave the application in place, + -- but we update the types since it may be partially applied or return + -- a higher-order term. + DynamicFun _ sv -> + let (argtypes', rettype) = dynamicFunType sv argtypes + restype = foldFunType argtypes' rettype `setAliases` aliases ret + -- FIXME: what if this application returns both a function + -- and a value? + callret + | orderZero ret = (Info ret, Info ext) + | otherwise = (Info restype, Info ext) + apply_e = Apply e1' e2' d callret loc + in return (apply_e, sv) + -- Propagate the 'IntrinsicsSV' until we reach the outermost application, + -- where we construct a dynamic static value with the appropriate type. + IntrinsicSV + | depth == 0 -> + -- If the intrinsic is fully applied, then we are done. + -- Otherwise we need to eta-expand it and recursively + -- defunctionalise. XXX: might it be better to simply + -- eta-expand immediately any time we encounter a + -- non-fully-applied intrinsic? + if null argtypes + then return (e', Dynamic $ typeOf e) + else do + (pats, body, tp) <- etaExpand (typeOf e') e' + defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty + | otherwise -> return (e', IntrinsicSV) + _ -> + error $ + "Application of an expression that is neither a static lambda " + ++ "nor a dynamic function, but has static value: " + ++ show sv1 +defuncApply depth e@(Var qn (Info t) loc) = do + let (argtypes, _) = unfoldFunType t + sv <- lookupVar loc (qualLeaf qn) + case sv of + DynamicFun _ _ + | fullyApplied sv depth -> + -- We still need to update the types in case the dynamic + -- function returns a higher-order term. + let (argtypes', rettype) = dynamicFunType sv argtypes + in return (Var qn (Info (foldFunType argtypes' rettype)) loc, sv) + | otherwise -> do + fname <- newName $ qualLeaf qn + let (dims, pats, e0, sv') = liftDynFun sv depth + pats_names = S.map identName $ mconcat $ map patternIdents pats + notInPats = (`S.notMember` pats_names) + dims' = filter notInPats dims + (argtypes', rettype) = dynamicFunType sv' argtypes + liftValDec fname (fromStruct rettype) dims' pats e0 + return + ( Var + (qualName fname) + (Info (foldFunType argtypes' $ fromStruct rettype)) + loc, + sv' + ) + IntrinsicSV -> return (e, IntrinsicSV) + _ -> return (Var qn (Info (typeFromSV' sv)) loc, sv) +defuncApply depth (Parens e _) = defuncApply depth e +defuncApply _ expr = defuncExp expr + +-- | Check if a 'StaticVal' and a given application depth corresponds +-- to a fully applied dynamic function. +fullyApplied :: StaticVal -> Int -> Bool +fullyApplied (DynamicFun _ sv) depth + | depth == 0 = False + | depth > 0 = fullyApplied sv (depth -1) +fullyApplied _ _ = True + +-- | Converts a dynamic function 'StaticVal' into a list of +-- dimensions, a list of parameters, a function body, and the +-- appropriate static value for applying the function at the given +-- depth of partial application. +liftDynFun :: StaticVal -> Int -> ([VName], [Pattern], Exp, StaticVal) +liftDynFun (DynamicFun (e, sv) _) 0 = ([], [], e, sv) +liftDynFun (DynamicFun clsr@(_, LambdaSV dims pat _ _ _) sv) d + | d > 0 = + let (dims', pats, e', sv') = liftDynFun sv (d -1) + in (nub $ dims ++ dims', pat : pats, e', DynamicFun clsr sv') +liftDynFun sv _ = + error $ + "Tried to lift a StaticVal " ++ show sv + ++ ", but expected a dynamic function." + +-- | Converts a pattern to an environment that binds the individual names of the +-- pattern to their corresponding types wrapped in a 'Dynamic' static value. +envFromPattern :: Pattern -> Env +envFromPattern pat = case pat of + TuplePattern ps _ -> foldMap envFromPattern ps + RecordPattern fs _ -> foldMap (envFromPattern . snd) fs + PatternParens p _ -> envFromPattern p + Id vn (Info t) _ -> M.singleton vn $ Dynamic t + Wildcard _ _ -> mempty + PatternAscription p _ _ -> envFromPattern p + PatternLit {} -> mempty + PatternConstr _ _ ps _ -> foldMap envFromPattern ps + +-- | Create an environment that binds the shape parameters. +envFromShapeParams :: [TypeParamBase VName] -> Env +envFromShapeParams = envFromDimNames . map dim + where + dim (TypeParamDim vn _) = vn + dim tparam = + error $ + "The defunctionalizer expects a monomorphic input program,\n" + ++ "but it received a type parameter " + ++ pretty tparam + ++ " at " + ++ locStr (srclocOf tparam) + ++ "." + +envFromDimNames :: [VName] -> Env +envFromDimNames = M.fromList . flip zip (repeat $ Dynamic $ Scalar $ Prim $ Signed Int64) + +-- | Create a new top-level value declaration with the given function name, +-- return type, list of parameters, and body expression. +liftValDec :: VName -> PatternType -> [VName] -> [Pattern] -> Exp -> DefM () +liftValDec fname rettype dims pats body = tell $ Seq.singleton dec + where + dims' = map (`TypeParamDim` mempty) dims + -- FIXME: this pass is still not correctly size-preserving, so + -- forget those return sizes that we forgot to propagate along + -- the way. Hopefully the internaliser is conservative and + -- will insert reshapes... + bound_here = S.fromList dims <> S.map identName (foldMap patternIdents pats) + anyDimIfNotBound (NamedDim v) + | qualLeaf v `S.member` bound_here = NamedDim v + | otherwise = AnyDim + anyDimIfNotBound d = d + rettype_st = first anyDimIfNotBound $ toStruct rettype + + dec = + ValBind + { valBindEntryPoint = Nothing, + valBindName = fname, + valBindRetDecl = Nothing, + valBindRetType = Info (rettype_st, []), + valBindTypeParams = dims', + valBindParams = pats, + valBindBody = body, + valBindDoc = Nothing, + valBindAttrs = mempty, + valBindLocation = mempty + } + +-- | Given a closure environment, construct a record pattern that +-- binds the closed over variables. +buildEnvPattern :: Env -> Pattern +buildEnvPattern env = RecordPattern (map buildField $ M.toList env) mempty + where + buildField (vn, sv) = + ( nameFromString (pretty vn), + Id vn (Info $ snd $ typeFromSV sv) mempty + ) + +-- | Given a closure environment pattern and the type of a term, +-- construct the type of that term, where uniqueness is set to +-- `Nonunique` for those arrays that are bound in the environment or +-- pattern (except if they are unique there). This ensures that a +-- lifted function can create unique arrays as long as they do not +-- alias any of its parameters. XXX: it is not clear that this is a +-- sufficient property, unfortunately. +buildRetType :: Env -> [Pattern] -> StructType -> PatternType -> PatternType +buildRetType env pats = comb + where + bound = foldMap oneName (M.keys env) <> foldMap patternVars pats + boundAsUnique v = + maybe False (unique . unInfo . identType) $ + find ((== v) . identName) $ S.toList $ foldMap patternIdents pats + problematic v = (v `member` bound) && not (boundAsUnique v) + comb (Scalar (Record fs_annot)) (Scalar (Record fs_got)) = + Scalar $ Record $ M.intersectionWith comb fs_annot fs_got + comb (Scalar (Sum cs_annot)) (Scalar (Sum cs_got)) = + Scalar $ Sum $ M.intersectionWith (zipWith comb) cs_annot cs_got + comb (Scalar Arrow {}) t = + descend t + comb got et = + descend $ fromStruct got `setAliases` aliases et + + descend t@Array {} + | any (problematic . aliasVar) (aliases t) = t `setUniqueness` Nonunique + descend (Scalar (Record t)) = Scalar $ Record $ fmap descend t + descend t = t + +-- | Compute the corresponding type for a given static value. +typeFromSV :: StaticVal -> ([VName], PatternType) +typeFromSV (Dynamic tp) = + (mempty, tp) +typeFromSV (LambdaSV sizes _ _ _ env) = + ( sizes <> env_sizes, + Scalar $ Record $ M.fromList $ map (fmap snd) env' + ) + where + env' = map (bimap (nameFromString . pretty) typeFromSV) $ M.toList env + env_sizes = concatMap (fst . snd) env' +typeFromSV (RecordSV ls) = + let ts = map (fmap typeFromSV) ls + in ( concatMap (fst . snd) ts, + Scalar $ Record $ M.fromList $ map (fmap snd) ts + ) +typeFromSV (DynamicFun (_, sv) _) = + typeFromSV sv +typeFromSV (SumSV name svs fields) = + let (sizes, svs') = unzip $ map typeFromSV svs + in ( concat sizes, + Scalar $ Sum $ M.insert name svs' $ M.fromList fields + ) +typeFromSV IntrinsicSV = + error "Tried to get the type from the static value of an intrinsic." + +typeFromSV' :: StaticVal -> PatternType +typeFromSV' sv = + let (sizes, t) = typeFromSV sv + in unscopeType (S.fromList sizes) t + +-- | Construct the type for a fully-applied dynamic function from its +-- static value and the original types of its arguments. +dynamicFunType :: StaticVal -> [PatternType] -> ([PatternType], PatternType) +dynamicFunType (DynamicFun _ sv) (p : ps) = + let (ps', ret) = dynamicFunType sv ps in (p : ps', ret) +dynamicFunType sv _ = ([], typeFromSV' sv) + +-- | Match a pattern with its static value. Returns an environment with +-- the identifier components of the pattern mapped to the corresponding +-- subcomponents of the static value. +matchPatternSV :: PatternBase Info VName -> StaticVal -> Env +matchPatternSV (TuplePattern ps _) (RecordSV ls) = + mconcat $ zipWith (\p (_, sv) -> matchPatternSV p sv) ps ls +matchPatternSV (RecordPattern ps _) (RecordSV ls) + | ps' <- sortOn fst ps, + ls' <- sortOn fst ls, + map fst ps' == map fst ls' = + mconcat $ zipWith (\(_, p) (_, sv) -> matchPatternSV p sv) ps' ls' +matchPatternSV (PatternParens pat _) sv = matchPatternSV pat sv +matchPatternSV (Id vn (Info t) _) sv = + -- When matching a pattern with a zero-order STaticVal, the type of + -- the pattern wins out. This is important when matching a + -- nonunique pattern with a unique value. + if orderZeroSV sv + then M.singleton vn $ Dynamic t + else M.singleton vn sv +matchPatternSV (Wildcard _ _) _ = mempty +matchPatternSV (PatternAscription pat _ _) sv = matchPatternSV pat sv +matchPatternSV PatternLit {} _ = mempty +matchPatternSV (PatternConstr c1 _ ps _) (SumSV c2 ls fs) + | c1 == c2 = + mconcat $ zipWith matchPatternSV ps ls + | Just ts <- lookup c1 fs = + mconcat $ zipWith matchPatternSV ps $ map svFromType ts + | otherwise = + error $ "matchPatternSV: missing constructor in type: " ++ pretty c1 +matchPatternSV (PatternConstr c1 _ ps _) (Dynamic (Scalar (Sum fs))) + | Just ts <- M.lookup c1 fs = + mconcat $ zipWith matchPatternSV ps $ map svFromType ts + | otherwise = + error $ "matchPatternSV: missing constructor in type: " ++ pretty c1 +matchPatternSV pat (Dynamic t) = matchPatternSV pat $ svFromType t +matchPatternSV pat sv = + error $ + "Tried to match pattern " ++ pretty pat + ++ " with static value " + ++ show sv + ++ "." + +orderZeroSV :: StaticVal -> Bool +orderZeroSV Dynamic {} = True +orderZeroSV (RecordSV fields) = all (orderZeroSV . snd) fields +orderZeroSV _ = False + +-- | Given a pattern and the static value for the defunctionalized argument, +-- update the pattern to reflect the changes in the types. +updatePattern :: Pattern -> StaticVal -> Pattern +updatePattern (TuplePattern ps loc) (RecordSV svs) = + TuplePattern (zipWith updatePattern ps $ map snd svs) loc +updatePattern (RecordPattern ps loc) (RecordSV svs) + | ps' <- sortOn fst ps, + svs' <- sortOn fst svs = + RecordPattern + ( zipWith + ( \(n, p) (_, sv) -> + (n, updatePattern p sv) + ) + ps' + svs' + ) + loc +updatePattern (PatternParens pat loc) sv = + PatternParens (updatePattern pat sv) loc +updatePattern (Id vn (Info tp) loc) sv = + Id vn (Info $ comb tp (snd (typeFromSV sv) `setUniqueness` Nonunique)) loc + where + -- Preserve any original zeroth-order types. + comb (Scalar Arrow {}) t2 = t2 + comb (Scalar (Record m1)) (Scalar (Record m2)) = + Scalar $ Record $ M.intersectionWith comb m1 m2 + comb (Scalar (Sum m1)) (Scalar (Sum m2)) = + Scalar $ Sum $ M.intersectionWith (zipWith comb) m1 m2 + comb t1 _ = t1 -- t1 must be array or prim. +updatePattern pat@(Wildcard (Info tp) loc) sv + | orderZero tp = pat + | otherwise = Wildcard (Info $ snd $ typeFromSV sv) loc +updatePattern (PatternAscription pat tydecl loc) sv + | orderZero . unInfo $ expandedType tydecl = + PatternAscription (updatePattern pat sv) tydecl loc + | otherwise = updatePattern pat sv +updatePattern p@PatternLit {} _ = p +updatePattern pat@(PatternConstr c1 (Info t) ps loc) sv@(SumSV _ svs _) + | orderZero t = pat + | otherwise = PatternConstr c1 (Info t') ps' loc + where + t' = snd (typeFromSV sv) `setUniqueness` Nonunique + ps' = zipWith updatePattern ps svs +updatePattern (PatternConstr c1 _ ps loc) (Dynamic t) = + PatternConstr c1 (Info t) ps loc +updatePattern pat (Dynamic t) = updatePattern pat (svFromType t) +updatePattern pat sv = + error $ + "Tried to update pattern " ++ pretty pat + ++ "to reflect the static value " + ++ show sv + +-- Like updatePattern, but discard sizes. This is used for +-- let-bindings, where we might otherwise introduce sizes that are +-- free. +updatePattern' :: Pattern -> StaticVal -> Pattern +updatePattern' pat sv = + let pat' = updatePattern pat sv + (sizes, _) = typeFromSV sv + tr = + identityMapper + { mapOnPatternType = + pure . unscopeType (S.fromList sizes) + } + in runIdentity $ astMap tr pat' + +-- | Convert a record (or tuple) type to a record static value. This is used for +-- "unwrapping" tuples and records that are nested in 'Dynamic' static values. +svFromType :: PatternType -> StaticVal +svFromType (Scalar (Record fs)) = RecordSV . M.toList $ M.map svFromType fs +svFromType t = Dynamic t + +-- A set of names where we also track uniqueness. +newtype NameSet = NameSet (M.Map VName Uniqueness) deriving (Show) + +instance Semigroup NameSet where + NameSet x <> NameSet y = NameSet $ M.unionWith max x y + +instance Monoid NameSet where + mempty = NameSet mempty + +without :: NameSet -> NameSet -> NameSet +without (NameSet x) (NameSet y) = NameSet $ x `M.difference` y + +member :: VName -> NameSet -> Bool +member v (NameSet m) = v `M.member` m + +ident :: Ident -> NameSet +ident v = NameSet $ M.singleton (identName v) (uniqueness $ unInfo $ identType v) + +oneName :: VName -> NameSet +oneName v = NameSet $ M.singleton v Nonunique + +names :: S.Set VName -> NameSet +names = foldMap oneName + +-- | Compute the set of free variables of an expression. +freeVars :: Exp -> NameSet +freeVars expr = case expr of + Literal {} -> mempty + IntLit {} -> mempty + FloatLit {} -> mempty + StringLit {} -> mempty + Parens e _ -> freeVars e + QualParens _ e _ -> freeVars e + TupLit es _ -> foldMap freeVars es + RecordLit fs _ -> foldMap freeVarsField fs + where + freeVarsField (RecordFieldExplicit _ e _) = freeVars e + freeVarsField (RecordFieldImplicit vn t _) = ident $ Ident vn t mempty + ArrayLit es t _ -> + foldMap freeVars es + <> names (typeDimNames $ unInfo t) + Range e me incl _ _ -> + freeVars e <> foldMap freeVars me + <> foldMap freeVars incl + Var qn (Info t) _ -> NameSet $ M.singleton (qualLeaf qn) $ uniqueness t + Ascript e t _ -> freeVars e <> names (typeDimNames $ unInfo $ expandedType t) + Coerce e t _ _ -> freeVars e <> names (typeDimNames $ unInfo $ expandedType t) + LetPat pat e1 e2 _ _ -> + freeVars e1 + <> ( (names (patternDimNames pat) <> freeVars e2) + `without` patternVars pat + ) + LetFun vn (_, pats, _, _, e1) e2 _ _ -> + ( (freeVars e1 <> names (foldMap patternDimNames pats)) + `without` foldMap patternVars pats + ) + <> (freeVars e2 `without` oneName vn) + If e1 e2 e3 _ _ -> freeVars e1 <> freeVars e2 <> freeVars e3 + Apply e1 e2 _ _ _ -> freeVars e1 <> freeVars e2 + Negate e _ -> freeVars e + Lambda pats e0 _ _ _ -> + (names (foldMap patternDimNames pats) <> freeVars e0) + `without` foldMap patternVars pats + OpSection {} -> mempty + OpSectionLeft _ _ e _ _ _ -> freeVars e + OpSectionRight _ _ e _ _ _ -> freeVars e + ProjectSection {} -> mempty + IndexSection idxs _ _ -> foldMap freeDimIndex idxs + DoLoop sparams pat e1 form e3 _ _ -> + let (e2fv, e2ident) = formVars form + in freeVars e1 <> e2fv + <> ( freeVars e3 + `without` (names (S.fromList sparams) <> patternVars pat <> e2ident) + ) + where + formVars (For v e2) = (freeVars e2, ident v) + formVars (ForIn p e2) = (freeVars e2, patternVars p) + formVars (While e2) = (freeVars e2, mempty) + BinOp (qn, _) _ (e1, _) (e2, _) _ _ _ -> + oneName (qualLeaf qn) + <> freeVars e1 + <> freeVars e2 + Project _ e _ _ -> freeVars e + LetWith id1 id2 idxs e1 e2 _ _ -> + ident id2 <> foldMap freeDimIndex idxs <> freeVars e1 + <> (freeVars e2 `without` ident id1) + Index e idxs _ _ -> freeVars e <> foldMap freeDimIndex idxs + Update e1 idxs e2 _ -> freeVars e1 <> foldMap freeDimIndex idxs <> freeVars e2 + RecordUpdate e1 _ e2 _ _ -> freeVars e1 <> freeVars e2 + Assert e1 e2 _ _ -> freeVars e1 <> freeVars e2 + Constr _ es _ _ -> foldMap freeVars es + Attr _ e _ -> freeVars e + Match e cs _ _ -> freeVars e <> foldMap caseFV cs + where + caseFV (CasePat p eCase _) = + (names (patternDimNames p) <> freeVars eCase) + `without` patternVars p + +freeDimIndex :: DimIndexBase Info VName -> NameSet +freeDimIndex (DimFix e) = freeVars e +freeDimIndex (DimSlice me1 me2 me3) = + foldMap (foldMap freeVars) [me1, me2, me3] + +-- | Extract all the variable names bound in a pattern. +patternVars :: Pattern -> NameSet +patternVars = mconcat . map ident . S.toList . patternIdents + +-- | Defunctionalize a top-level value binding. Returns the +-- transformed result as well as an environment that binds the name of +-- the value binding to the static value of the transformed body. The +-- boolean is true if the function is a 'DynamicFun'. +defuncValBind :: ValBind -> DefM (ValBind, Env, Bool) +-- Eta-expand entry points with a functional return type. +defuncValBind (ValBind entry name _ (Info (rettype, retext)) tparams params body _ attrs loc) + | Scalar Arrow {} <- rettype = do + (body_pats, body', rettype') <- etaExpand (fromStruct rettype) body + defuncValBind $ + ValBind + entry + name + Nothing + (Info (rettype', retext)) + tparams + (params <> body_pats) + body' + Nothing + attrs + loc +defuncValBind valbind@(ValBind _ name retdecl (Info (rettype, retext)) tparams params body _ _ _) = do + (tparams', params', body', sv) <- defuncLet tparams params body rettype + let rettype' = combineTypeShapes rettype $ anySizes $ toStruct $ typeOf body' + (missing_dims, params'') <- sizesForAll params' + return + ( valbind + { valBindRetDecl = retdecl, + valBindRetType = + Info + ( if null params' + then rettype' `setUniqueness` Nonunique + else rettype', + retext + ), + valBindTypeParams = + tparams' + ++ map (`TypeParamDim` mempty) missing_dims, + valBindParams = params'', + valBindBody = body' + }, + M.singleton name sv, + case sv of + DynamicFun {} -> True + Dynamic {} -> True + _ -> False + ) + +-- | Defunctionalize a list of top-level declarations. +defuncVals :: [ValBind] -> DefM (Seq.Seq ValBind) +defuncVals [] = return mempty +defuncVals (valbind : ds) = do + ((valbind', env, dyn), defs) <- collectFuns $ defuncValBind valbind + ds' <- + localEnv env $ + if dyn + then isGlobal (valBindName valbind') $ defuncVals ds + else defuncVals ds + return $ defs <> Seq.singleton valbind' <> ds' + +-- | Transform a list of top-level value bindings. May produce new +-- lifted function definitions, which are placed in front of the +-- resulting list of declarations. +transformProg :: MonadFreshNames m => [ValBind] -> m [ValBind] +transformProg decs = modifyNameSource $ \namesrc -> + let (decs', namesrc', liftedDecs) = runDefM namesrc $ defuncVals decs + in (toList $ liftedDecs <> decs', namesrc') 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 @@ -1,22 +1,21 @@ --- | Partially evaluate all modules away from a source Futhark --- program. This is implemented as a source-to-source transformation. {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE Trustworthy #-} + +-- | Partially evaluate all modules away from a source Futhark +-- program. This is implemented as a source-to-source transformation. module Futhark.Internalise.Defunctorise (transformProg) where -import Control.Monad.RWS.Strict import Control.Monad.Identity +import Control.Monad.RWS.Strict import qualified Data.DList as DL import qualified Data.Map as M -import qualified Data.Set as S import Data.Maybe - -import Prelude hiding (mod, abs) - +import qualified Data.Set as S import Futhark.MonadFreshNames import Language.Futhark +import Language.Futhark.Semantic (FileModule (..), Imports) import Language.Futhark.Traversals -import Language.Futhark.Semantic (Imports, FileModule(..)) +import Prelude hiding (abs, mod) -- | A substitution from names in the original program to names in the -- generated/residual program. @@ -24,53 +23,60 @@ lookupSubst :: VName -> Substitutions -> VName lookupSubst v substs = case M.lookup v substs of - Just v' | v' /= v -> lookupSubst v' substs - _ -> v + Just v' | v' /= v -> lookupSubst v' substs + _ -> v -data Mod = ModFun TySet Scope ModParam ModExp - -- ^ A pairing of a lexical closure and a module function. - | ModMod Scope - -- ^ A non-parametric module. - deriving (Show) +data Mod + = -- | A pairing of a lexical closure and a module function. + ModFun TySet Scope ModParam ModExp + | -- | A non-parametric module. + ModMod Scope + deriving (Show) modScope :: Mod -> Scope modScope (ModMod scope) = scope -modScope ModFun{} = mempty +modScope ModFun {} = mempty -data Scope = Scope { scopeSubsts :: Substitutions - , scopeMods :: M.Map VName Mod - } - deriving (Show) +data Scope = Scope + { scopeSubsts :: Substitutions, + scopeMods :: M.Map VName Mod + } + deriving (Show) lookupSubstInScope :: QualName VName -> Scope -> (QualName VName, Scope) lookupSubstInScope qn@(QualName quals name) scope@(Scope substs mods) = case quals of [] -> (qualName $ lookupSubst name substs, scope) - q:qs -> + q : qs -> let q' = lookupSubst q substs - in case M.lookup q' mods of - Just (ModMod mod_scope) -> lookupSubstInScope (QualName qs name) mod_scope - _ -> (qn, scope) + in case M.lookup q' mods of + Just (ModMod mod_scope) -> lookupSubstInScope (QualName qs name) mod_scope + _ -> (qn, scope) instance Semigroup Scope where - Scope ss1 mt1 <> Scope ss2 mt2 = Scope (ss1<>ss2) (mt1<>mt2) + Scope ss1 mt1 <> Scope ss2 mt2 = Scope (ss1 <> ss2) (mt1 <> mt2) instance Monoid Scope where mempty = Scope mempty mempty type TySet = S.Set VName -data Env = Env { envScope :: Scope - , envGenerating :: Bool - , envImports :: M.Map String Scope - , envAbs :: TySet - } +data Env = Env + { envScope :: Scope, + envGenerating :: Bool, + envImports :: M.Map String Scope, + envAbs :: TySet + } newtype TransformM a = TransformM (RWS Env (DL.DList Dec) VNameSource a) - deriving (Applicative, Functor, Monad, - MonadFreshNames, - MonadReader Env, - MonadWriter (DL.DList Dec)) + deriving + ( Applicative, + Functor, + Monad, + MonadFreshNames, + MonadReader Env, + MonadWriter (DL.DList Dec) + ) emit :: Dec -> TransformM () emit = tell . DL.singleton @@ -79,20 +85,24 @@ askScope = asks envScope localScope :: (Scope -> Scope) -> TransformM a -> TransformM a -localScope f = local $ \env -> env { envScope = f $ envScope env } +localScope f = local $ \env -> env {envScope = f $ envScope env} extendScope :: Scope -> TransformM a -> TransformM a extendScope (Scope substs mods) = localScope $ \scope -> - scope { scopeSubsts = M.map (forward (scopeSubsts scope)) substs <> scopeSubsts scope - , scopeMods = mods <> scopeMods scope } - where forward old_substs v = fromMaybe v $ M.lookup v old_substs + scope + { scopeSubsts = M.map (forward (scopeSubsts scope)) substs <> scopeSubsts scope, + scopeMods = mods <> scopeMods scope + } + where + forward old_substs v = fromMaybe v $ M.lookup v old_substs substituting :: Substitutions -> TransformM a -> TransformM a -substituting substs = extendScope mempty { scopeSubsts = substs } +substituting substs = extendScope mempty {scopeSubsts = substs} boundName :: VName -> TransformM VName -boundName v = do g <- asks envGenerating - if g then newName v else return v +boundName v = do + g <- asks envGenerating + if g then newName v else return v bindingNames :: [VName] -> TransformM Scope -> TransformM Scope bindingNames names m = do @@ -101,35 +111,41 @@ substituting substs $ mappend <$> m <*> pure (Scope substs mempty) generating :: TransformM a -> TransformM a -generating = local $ \env -> env { envGenerating = True } +generating = local $ \env -> env {envGenerating = True} bindingImport :: String -> Scope -> TransformM a -> TransformM a bindingImport name scope = local $ \env -> - env { envImports = M.insert name scope $ envImports env } + env {envImports = M.insert name scope $ envImports env} bindingAbs :: TySet -> TransformM a -> TransformM a bindingAbs abs = local $ \env -> - env { envAbs = abs <> envAbs env } + env {envAbs = abs <> envAbs env} lookupImport :: String -> TransformM Scope lookupImport name = maybe bad return =<< asks (M.lookup name . envImports) - where bad = error $ "Unknown import: " ++ name + where + bad = error $ "Unknown import: " ++ name lookupMod' :: QualName VName -> Scope -> Either String Mod lookupMod' mname scope = let (mname', scope') = lookupSubstInScope mname scope - in maybe (Left $ bad mname') Right $ M.lookup (qualLeaf mname') $ scopeMods scope' - where bad mname' = "Unknown module: " ++ pretty mname ++ " (" ++ pretty mname' ++ ")" + in maybe (Left $ bad mname') Right $ M.lookup (qualLeaf mname') $ scopeMods scope' + where + bad mname' = "Unknown module: " ++ pretty mname ++ " (" ++ pretty mname' ++ ")" lookupMod :: QualName VName -> TransformM Mod lookupMod mname = either error return . lookupMod' mname =<< askScope runTransformM :: VNameSource -> TransformM a -> (a, VNameSource, DL.DList Dec) runTransformM src (TransformM m) = runRWS m env src - where env = Env mempty False mempty mempty + where + env = Env mempty False mempty mempty -maybeAscript :: SrcLoc -> Maybe (SigExp, Info (M.Map VName VName)) -> ModExp - -> ModExp +maybeAscript :: + SrcLoc -> + Maybe (SigExp, Info (M.Map VName VName)) -> + ModExp -> + ModExp maybeAscript loc (Just (mtye, substs)) me = ModAscript me mtye substs loc maybeAscript _ Nothing me = me @@ -137,12 +153,14 @@ substituteInMod substs (ModMod (Scope mod_substs mod_mods)) = -- Forward all substitutions. ModMod $ Scope substs' $ M.map (substituteInMod substs) mod_mods - where forward v = lookupSubst v $ mod_substs <> substs - substs' = M.map forward substs + where + forward v = lookupSubst v $ mod_substs <> substs + substs' = M.map forward substs substituteInMod substs (ModFun abs (Scope mod_substs mod_mods) mparam mbody) = - ModFun abs (Scope (substs'<>mod_substs) mod_mods) mparam mbody - where forward v = lookupSubst v mod_substs - substs' = M.map forward substs + ModFun abs (Scope (substs' <> mod_substs) mod_mods) mparam mbody + where + forward v = lookupSubst v mod_substs + substs' = M.map forward substs extendAbsTypes :: Substitutions -> TransformM a -> TransformM a extendAbsTypes ascript_substs m = do @@ -150,8 +168,11 @@ -- Some abstract types may have a different name on the inside, and -- we need to make them visible, because substitutions involving -- abstract types must be lifted out in transformModBind. - let subst_abs = S.fromList $ map snd $ filter ((`S.member` abs) . fst) $ - M.toList ascript_substs + let subst_abs = + S.fromList $ + map snd $ + filter ((`S.member` abs) . fst) $ + M.toList ascript_substs bindingAbs subst_abs m evalModExp :: ModExp -> TransformM Mod @@ -161,7 +182,7 @@ evalModExp (ModImport _ (Info fpath) _) = ModMod <$> lookupImport fpath evalModExp (ModAscript me _ (Info ascript_substs) _) = extendAbsTypes ascript_substs $ - substituteInMod ascript_substs <$> evalModExp me + substituteInMod ascript_substs <$> evalModExp me evalModExp (ModApply f arg (Info p_substs) (Info b_substs) loc) = do f_mod <- evalModExp f arg_mod <- evalModExp arg @@ -170,33 +191,43 @@ 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)) $ 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 - addSubsts _ substs (ModMod (Scope msubsts mods)) = - ModMod $ Scope (substs<>msubsts) mods - addSubstsModMod substs (ModMod (Scope msubsts mods)) = - ModMod $ Scope (substs<>msubsts) mods - addSubstsModMod _ m = m + 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 + ) + ) + $ 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 + addSubsts _ substs (ModMod (Scope msubsts mods)) = + ModMod $ Scope (substs <> msubsts) mods + addSubstsModMod substs (ModMod (Scope msubsts mods)) = + ModMod $ Scope (substs <> msubsts) mods + addSubstsModMod _ m = m evalModExp (ModLambda p ascript e loc) = do scope <- askScope abs <- asks envAbs @@ -210,24 +241,26 @@ transformNames x = do scope <- askScope return $ runIdentity $ astMap (substituter scope) x - where substituter scope = - ASTMapper { mapOnExp = onExp scope - , mapOnName = \v -> - return $ qualLeaf $ fst $ lookupSubstInScope (qualName v) scope - , mapOnQualName = \v -> - return $ fst $ lookupSubstInScope v scope - , mapOnStructType = astMap (substituter scope) - , mapOnPatternType = astMap (substituter scope) - } - onExp scope e = - -- One expression is tricky, because it interacts with scoping rules. - case e of - QualParens (mn, _) e' _ -> - case lookupMod' mn scope of - Left err -> error err - Right mod -> - astMap (substituter $ modScope mod<>scope) e' - _ -> astMap (substituter scope) e + where + substituter scope = + ASTMapper + { mapOnExp = onExp scope, + mapOnName = \v -> + return $ qualLeaf $ fst $ lookupSubstInScope (qualName v) scope, + mapOnQualName = \v -> + return $ fst $ lookupSubstInScope v scope, + mapOnStructType = astMap (substituter scope), + mapOnPatternType = astMap (substituter scope) + } + onExp scope e = + -- One expression is tricky, because it interacts with scoping rules. + case e of + QualParens (mn, _) e' _ -> + case lookupMod' mn scope of + Left err -> error err + Right mod -> + astMap (substituter $ modScope mod <> scope) e' + _ -> astMap (substituter scope) e transformTypeExp :: TypeExp VName -> TransformM (TypeExp VName) transformTypeExp = transformNames @@ -255,21 +288,29 @@ transformTypeBind :: TypeBind -> TransformM () transformTypeBind (TypeBind name l tparams te doc loc) = do name' <- transformName name - emit . TypeDec =<< (TypeBind name' l <$> traverse transformNames tparams - <*> transformTypeDecl te <*> pure doc <*> pure loc) + emit . TypeDec + =<< ( TypeBind name' l <$> traverse transformNames tparams + <*> transformTypeDecl te + <*> pure doc + <*> pure loc + ) transformModBind :: ModBind -> TransformM Scope transformModBind mb = do let addParam p me = ModLambda p Nothing me $ srclocOf me - mod <- evalModExp $ foldr addParam - (maybeAscript (srclocOf mb) (modSignature mb) $ modExp mb) $ - modParams mb + mod <- + evalModExp $ + foldr + addParam + (maybeAscript (srclocOf mb) (modSignature mb) $ modExp mb) + $ modParams mb mname <- transformName $ modName mb abs <- asks envAbs -- Copy substitutions involving abstract types out, because they are -- always resolved at the outermost level. - let abs_substs = M.filterWithKey (const . flip S.member abs) $ - scopeSubsts $ modScope mod + let abs_substs = + M.filterWithKey (const . flip S.member abs) $ + scopeSubsts $ modScope mod return $ Scope abs_substs $ M.singleton mname mod transformDecs :: [Dec] -> TransformM Scope @@ -298,28 +339,32 @@ extendScope scope $ mappend <$> transformDecs ds' <*> pure scope ImportDec name name' loc : ds' -> let d = LocalDec (OpenDec (ModImport name name' loc) loc) loc - in transformDecs $ d : ds' + in transformDecs $ d : ds' transformImports :: Imports -> TransformM () transformImports [] = return () -transformImports ((name,imp):imps) = do +transformImports ((name, imp) : imps) = do let abs = S.fromList $ map qualLeaf $ M.keys $ fileAbs imp - scope <- censor (fmap maybeHideEntryPoint) $ - bindingAbs abs $ transformDecs $ progDecs $ fileProg imp + scope <- + censor (fmap maybeHideEntryPoint) $ + bindingAbs abs $ transformDecs $ progDecs $ fileProg imp bindingAbs abs $ bindingImport name scope $ transformImports imps where -- Only the "main" file (last import) is allowed to have entry points. permit_entry_points = null imps maybeHideEntryPoint (ValDec vdec) = - ValDec vdec { valBindEntryPoint = - if permit_entry_points - then valBindEntryPoint vdec - else Nothing } + ValDec + vdec + { valBindEntryPoint = + if permit_entry_points + then valBindEntryPoint vdec + else Nothing + } maybeHideEntryPoint d = d -- | Perform defunctorisation. transformProg :: MonadFreshNames m => Imports -> m [Dec] transformProg prog = modifyNameSource $ \namesrc -> let ((), namesrc', prog') = runTransformM namesrc $ transformImports prog - in (DL.toList prog', namesrc') + in (DL.toList prog', namesrc') diff --git a/src/Futhark/Internalise/Lambdas.hs b/src/Futhark/Internalise/Lambdas.hs --- a/src/Futhark/Internalise/Lambdas.hs +++ b/src/Futhark/Internalise/Lambdas.hs @@ -1,124 +1,151 @@ {-# LANGUAGE FlexibleContexts #-} + module Futhark.Internalise.Lambdas - ( InternaliseLambda - , internaliseMapLambda - , internaliseStreamMapLambda - , internaliseFoldLambda - , internaliseStreamLambda - , internalisePartitionLambda + ( InternaliseLambda, + internaliseMapLambda, + internaliseStreamMapLambda, + internaliseFoldLambda, + internaliseStreamLambda, + internalisePartitionLambda, ) - where +where -import Language.Futhark as E import Futhark.IR.SOACS as I -import Futhark.MonadFreshNames -import Futhark.Internalise.Monad import Futhark.Internalise.AccurateSizes +import Futhark.Internalise.Monad +import Language.Futhark as E -- | A function for internalising lambdas. type InternaliseLambda = E.Exp -> [I.Type] -> InternaliseM ([I.LParam], I.Body, [I.Type]) -internaliseMapLambda :: InternaliseLambda - -> E.Exp - -> [I.SubExp] - -> InternaliseM I.Lambda +internaliseMapLambda :: + InternaliseLambda -> + E.Exp -> + [I.SubExp] -> + InternaliseM I.Lambda internaliseMapLambda internaliseLambda lam args = do argtypes <- mapM I.subExpType args let rowtypes = map I.rowType argtypes (params, body, rettype) <- internaliseLambda lam rowtypes - body' <- localScope (scopeOfLParams params) $ - ensureResultShape - (ErrorMsg [ErrorString "not all iterations produce same shape"]) - (srclocOf lam) rettype body + body' <- + localScope (scopeOfLParams params) $ + ensureResultShape + (ErrorMsg [ErrorString "not all iterations produce same shape"]) + (srclocOf lam) + rettype + body return $ I.Lambda params body' rettype -internaliseStreamMapLambda :: InternaliseLambda - -> E.Exp - -> [I.SubExp] - -> InternaliseM I.Lambda +internaliseStreamMapLambda :: + InternaliseLambda -> + E.Exp -> + [I.SubExp] -> + InternaliseM I.Lambda internaliseStreamMapLambda internaliseLambda lam args = do chunk_size <- newVName "chunk_size" - let chunk_param = I.Param chunk_size (I.Prim int32) + let chunk_param = I.Param chunk_size (I.Prim int64) outer = (`setOuterSize` I.Var chunk_size) localScope (scopeOfLParams [chunk_param]) $ do argtypes <- mapM I.subExpType args (lam_params, orig_body, rettype) <- - internaliseLambda lam $ I.Prim int32 : map outer argtypes + internaliseLambda lam $ I.Prim int64 : map outer argtypes let orig_chunk_param : params = lam_params body <- runBodyBinder $ do letBindNames [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size return orig_body - body' <- localScope (scopeOfLParams params) $ insertStmsM $ do - letBindNames [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size - ensureResultShape - (ErrorMsg [ErrorString "not all iterations produce same shape"]) - (srclocOf lam) (map outer rettype) body - return $ I.Lambda (chunk_param:params) body' (map outer rettype) + body' <- localScope (scopeOfLParams params) $ + insertStmsM $ do + letBindNames [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size + ensureResultShape + (ErrorMsg [ErrorString "not all iterations produce same shape"]) + (srclocOf lam) + (map outer rettype) + body + return $ I.Lambda (chunk_param : params) body' (map outer rettype) -internaliseFoldLambda :: InternaliseLambda - -> E.Exp - -> [I.Type] -> [I.Type] - -> InternaliseM I.Lambda +internaliseFoldLambda :: + InternaliseLambda -> + E.Exp -> + [I.Type] -> + [I.Type] -> + InternaliseM I.Lambda internaliseFoldLambda internaliseLambda lam acctypes arrtypes = do let rowtypes = map I.rowType arrtypes (params, body, rettype) <- internaliseLambda lam $ acctypes ++ rowtypes - let rettype' = [ t `I.setArrayShape` I.arrayShape shape - | (t,shape) <- zip rettype acctypes ] + let rettype' = + [ t `I.setArrayShape` I.arrayShape shape + | (t, shape) <- zip rettype acctypes + ] -- The result of the body must have the exact same shape as the -- initial accumulator. We accomplish this with an assertion and -- reshape(). - body' <- localScope (scopeOfLParams params) $ - ensureResultShape - (ErrorMsg [ErrorString "shape of result does not match shape of initial value"]) - (srclocOf lam) rettype' body + body' <- + localScope (scopeOfLParams params) $ + ensureResultShape + (ErrorMsg [ErrorString "shape of result does not match shape of initial value"]) + (srclocOf lam) + rettype' + body return $ I.Lambda params body' rettype' -internaliseStreamLambda :: InternaliseLambda - -> E.Exp - -> [I.Type] - -> InternaliseM ([LParam], Body) +internaliseStreamLambda :: + InternaliseLambda -> + E.Exp -> + [I.Type] -> + InternaliseM ([LParam], Body) internaliseStreamLambda internaliseLambda lam rowts = do chunk_size <- newVName "chunk_size" - let chunk_param = I.Param chunk_size $ I.Prim int32 + let chunk_param = I.Param chunk_size $ I.Prim int64 chunktypes = map (`arrayOfRow` I.Var chunk_size) rowts localScope (scopeOfLParams [chunk_param]) $ do (lam_params, orig_body, _) <- - internaliseLambda lam $ I.Prim int32 : chunktypes + internaliseLambda lam $ I.Prim int64 : chunktypes let orig_chunk_param : params = lam_params body <- runBodyBinder $ do letBindNames [paramName orig_chunk_param] $ I.BasicOp $ I.SubExp $ I.Var chunk_size return orig_body - return (chunk_param:params, body) + return (chunk_param : params, body) -- Given @k@ lambdas, this will return a lambda that returns an -- (k+2)-element tuple of integers. The first element is the -- equivalence class ID in the range [0,k]. The remaining are all zero -- except for possibly one element. -internalisePartitionLambda :: InternaliseLambda - -> Int - -> E.Exp - -> [I.SubExp] - -> InternaliseM I.Lambda +internalisePartitionLambda :: + InternaliseLambda -> + Int -> + E.Exp -> + [I.SubExp] -> + InternaliseM I.Lambda internalisePartitionLambda internaliseLambda k lam args = do argtypes <- mapM I.subExpType args let rowtypes = map I.rowType argtypes (params, body, _) <- internaliseLambda lam rowtypes - body' <- localScope (scopeOfLParams params) $ - lambdaWithIncrement body + body' <- + localScope (scopeOfLParams params) $ + lambdaWithIncrement body return $ I.Lambda params body' rettype - where rettype = replicate (k+2) $ I.Prim int32 - result i = map constant $ (fromIntegral i :: Int32) : - (replicate i 0 ++ [1::Int32] ++ replicate (k-i) 0) + where + rettype = replicate (k + 2) $ I.Prim int64 + result i = + map constant $ + fromIntegral i : + (replicate i 0 ++ [1 :: Int64] ++ replicate (k - i) 0) - mkResult _ i | i >= k = return $ result i - mkResult eq_class i = do - is_i <- letSubExp "is_i" $ BasicOp $ CmpOp (CmpEq int32) eq_class (constant i) - fmap (map I.Var) . letTupExp "part_res" =<< - eIf (eSubExp is_i) (pure $ resultBody $ result i) - (resultBody <$> mkResult eq_class (i+1)) + mkResult _ i | i >= k = return $ result i + mkResult eq_class i = do + is_i <- + letSubExp "is_i" $ + BasicOp $ + CmpOp (CmpEq int64) eq_class $ + intConst Int64 $ toInteger i + fmap (map I.Var) . letTupExp "part_res" + =<< eIf + (eSubExp is_i) + (pure $ resultBody $ result i) + (resultBody <$> mkResult eq_class (i + 1)) - lambdaWithIncrement :: I.Body -> InternaliseM I.Body - lambdaWithIncrement lam_body = runBodyBinder $ do - eq_class <- head <$> bodyBind lam_body - resultBody <$> mkResult eq_class 0 + lambdaWithIncrement :: I.Body -> InternaliseM I.Body + lambdaWithIncrement lam_body = runBodyBinder $ do + eq_class <- head <$> bodyBind lam_body + resultBody <$> mkResult eq_class 0 diff --git a/src/Futhark/Internalise/Monad.hs b/src/Futhark/Internalise/Monad.hs --- a/src/Futhark/Internalise/Monad.hs +++ b/src/Futhark/Internalise/Monad.hs @@ -1,46 +1,38 @@ {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Trustworthy #-} -module Futhark.Internalise.Monad - ( InternaliseM - , runInternaliseM - , throwError - , VarSubstitutions - , InternaliseEnv (..) - , Closure - , FunInfo - - , substitutingVars - , lookupSubst - , addFunDef - - , lookupFunction - , lookupFunction' - , lookupConst - , allConsts - - , bindFunction - , bindConstant - - , localConstsScope +{-# LANGUAGE TypeFamilies #-} - , assert +module Futhark.Internalise.Monad + ( InternaliseM, + runInternaliseM, + throwError, + VarSubstitutions, + InternaliseEnv (..), + Closure, + FunInfo, + substitutingVars, + lookupSubst, + addFunDef, + lookupFunction, + lookupFunction', + lookupConst, + allConsts, + bindFunction, + bindConstant, + localConstsScope, + assert, -- * Convenient reexports - , module Futhark.Tools + module Futhark.Tools, ) - where +where import Control.Monad.Except -import Control.Monad.State -import Control.Monad.Reader -import Control.Monad.Writer import Control.Monad.RWS import qualified Data.Map.Strict as M - import Futhark.IR.SOACS import Futhark.MonadFreshNames import Futhark.Tools @@ -50,10 +42,14 @@ -- corresponds to the closure of a locally defined function. type Closure = [VName] -type FunInfo = (Name, Closure, - [VName], [DeclType], - [FParam], - [(SubExp,Type)] -> Maybe [DeclExtType]) +type FunInfo = + ( Name, + Closure, + [VName], + [DeclType], + [FParam], + [(SubExp, Type)] -> Maybe [DeclExtType] + ) type FunTable = M.Map VName FunInfo @@ -61,46 +57,55 @@ -- internalised subexpressions. type VarSubstitutions = M.Map VName [SubExp] -data InternaliseEnv = InternaliseEnv { - envSubsts :: VarSubstitutions - , envDoBoundsChecks :: Bool - , envSafe :: Bool - , envAttrs :: Attrs +data InternaliseEnv = InternaliseEnv + { envSubsts :: VarSubstitutions, + envDoBoundsChecks :: Bool, + envSafe :: Bool, + envAttrs :: Attrs } -data InternaliseState = InternaliseState { - stateNameSource :: VNameSource - , stateFunTable :: FunTable - , stateConstSubsts :: VarSubstitutions - , stateConstScope :: Scope SOACS - , stateConsts :: Names +data InternaliseState = InternaliseState + { stateNameSource :: VNameSource, + stateFunTable :: FunTable, + stateConstSubsts :: VarSubstitutions, + stateConstScope :: Scope SOACS, + stateConsts :: Names } data InternaliseResult = InternaliseResult (Stms SOACS) [FunDef SOACS] instance Semigroup InternaliseResult where InternaliseResult xs1 ys1 <> InternaliseResult xs2 ys2 = - InternaliseResult (xs1<>xs2) (ys1<>ys2) + InternaliseResult (xs1 <> xs2) (ys1 <> ys2) instance Monoid InternaliseResult where mempty = InternaliseResult mempty mempty -newtype InternaliseM a = InternaliseM (BinderT SOACS - (RWS - InternaliseEnv - InternaliseResult - InternaliseState) - a) - deriving (Functor, Applicative, Monad, - MonadReader InternaliseEnv, - MonadState InternaliseState, - MonadFreshNames, - HasScope SOACS, - LocalScope SOACS) +newtype InternaliseM a + = InternaliseM + ( BinderT + SOACS + ( RWS + InternaliseEnv + InternaliseResult + InternaliseState + ) + a + ) + deriving + ( Functor, + Applicative, + Monad, + MonadReader InternaliseEnv, + MonadState InternaliseState, + MonadFreshNames, + HasScope SOACS, + LocalScope SOACS + ) instance (Monoid w, Monad m) => MonadFreshNames (RWST r w InternaliseState m) where getNameSource = gets stateNameSource - putNameSource src = modify $ \s -> s { stateNameSource = src } + putNameSource src = modify $ \s -> s {stateNameSource = src} instance MonadBinder InternaliseM where type Lore InternaliseM = SOACS @@ -111,30 +116,35 @@ addStms = InternaliseM . addStms collectStms (InternaliseM m) = InternaliseM $ collectStms m -runInternaliseM :: MonadFreshNames m => - Bool -> InternaliseM () - -> m (Stms SOACS, [FunDef SOACS]) +runInternaliseM :: + MonadFreshNames m => + Bool -> + InternaliseM () -> + m (Stms SOACS, [FunDef SOACS]) runInternaliseM safe (InternaliseM m) = modifyNameSource $ \src -> - let ((_, consts), s, InternaliseResult _ funs) = - runRWS (runBinderT m mempty) newEnv (newState src) - in ((consts, funs), stateNameSource s) - where newEnv = InternaliseEnv { - envSubsts = mempty - , envDoBoundsChecks = True - , envSafe = safe - , envAttrs = mempty - } - newState src = - InternaliseState { stateNameSource = src - , stateFunTable = mempty - , stateConstSubsts = mempty - , stateConsts = mempty - , stateConstScope = mempty - } + let ((_, consts), s, InternaliseResult _ funs) = + runRWS (runBinderT m mempty) newEnv (newState src) + in ((consts, funs), stateNameSource s) + where + newEnv = + InternaliseEnv + { envSubsts = mempty, + envDoBoundsChecks = True, + envSafe = safe, + envAttrs = mempty + } + newState src = + InternaliseState + { stateNameSource = src, + stateFunTable = mempty, + stateConstSubsts = mempty, + stateConsts = mempty, + stateConstScope = mempty + } substitutingVars :: VarSubstitutions -> InternaliseM a -> InternaliseM a -substitutingVars substs = local $ \env -> env { envSubsts = substs <> envSubsts env } +substitutingVars substs = local $ \env -> env {envSubsts = substs <> envSubsts env} lookupSubst :: VName -> InternaliseM (Maybe [SubExp]) lookupSubst v = do @@ -152,7 +162,8 @@ lookupFunction :: VName -> InternaliseM FunInfo lookupFunction fname = maybe bad return =<< lookupFunction' fname - where bad = error $ "Internalise.lookupFunction: Function '" ++ pretty fname ++ "' not found." + where + bad = error $ "Internalise.lookupFunction: Function '" ++ pretty fname ++ "' not found." lookupConst :: VName -> InternaliseM (Maybe [SubExp]) lookupConst fname = gets $ M.lookup fname . stateConstSubsts @@ -163,19 +174,21 @@ bindFunction :: VName -> FunDef SOACS -> FunInfo -> InternaliseM () bindFunction fname fd info = do addFunDef fd - modify $ \s -> s { stateFunTable = M.insert fname info $ stateFunTable s } + modify $ \s -> s {stateFunTable = M.insert fname info $ stateFunTable s} bindConstant :: VName -> FunDef SOACS -> InternaliseM () bindConstant cname fd = do let stms = bodyStms $ funDefBody fd - substs = takeLast (length (funDefRetType fd)) $ - bodyResult $ funDefBody fd + substs = + takeLast (length (funDefRetType fd)) $ + bodyResult $ funDefBody fd const_names = namesFromList $ M.keys $ scopeOf stms addStms stms modify $ \s -> - s { stateConstSubsts = M.insert cname substs $ stateConstSubsts s - , stateConstScope = scopeOf stms <> stateConstScope s - , stateConsts = const_names <> stateConsts s + s + { stateConstSubsts = M.insert cname substs $ stateConstSubsts s, + stateConstScope = scopeOf stms <> stateConstScope s, + stateConsts = const_names <> stateConsts s } localConstsScope :: InternaliseM a -> InternaliseM a @@ -186,25 +199,32 @@ -- | Construct an 'Assert' statement, but taking attributes into -- account. Always use this function, and never construct 'Assert' -- directly in the internaliser! -assert :: String -> SubExp -> ErrorMsg SubExp -> SrcLoc - -> InternaliseM Certificates +assert :: + String -> + SubExp -> + ErrorMsg SubExp -> + SrcLoc -> + InternaliseM Certificates assert desc se msg loc = assertingOne $ do attrs <- asks $ attrsForAssert . envAttrs - attributing attrs $ letExp desc $ - BasicOp $ Assert se msg (loc, mempty) + attributing attrs $ + letExp desc $ + BasicOp $ Assert se msg (loc, mempty) -- | Execute the given action if 'envDoBoundsChecks' is true, otherwise -- just return an empty list. -asserting :: InternaliseM Certificates - -> InternaliseM Certificates +asserting :: + InternaliseM Certificates -> + InternaliseM Certificates asserting m = do doBoundsChecks <- asks envDoBoundsChecks if doBoundsChecks - then m - else return mempty + then m + else return mempty -- | Execute the given action if 'envDoBoundsChecks' is true, otherwise -- just return an empty list. -assertingOne :: InternaliseM VName - -> InternaliseM Certificates +assertingOne :: + InternaliseM VName -> + InternaliseM Certificates assertingOne m = asserting $ Certificates . pure <$> m 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 @@ -1,6 +1,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE Trustworthy #-} + -- | This monomorphization module converts a well-typed, polymorphic, -- module-free Futhark program into an equivalent monomorphic program. -- @@ -23,30 +24,28 @@ -- -- Note that these changes are unfortunately not visible in the AST -- representation. -module Futhark.Internalise.Monomorphise - ( transformProg ) where +module Futhark.Internalise.Monomorphise (transformProg) where -import Control.Monad.Identity -import Control.Monad.RWS hiding (Sum) -import Control.Monad.State -import Control.Monad.Writer hiding (Sum) -import Data.Bitraversable -import Data.Bifunctor -import Data.List (partition) +import Control.Monad.Identity +import Control.Monad.RWS hiding (Sum) +import Control.Monad.State +import Control.Monad.Writer hiding (Sum) +import Data.Bifunctor +import Data.Bitraversable +import Data.Foldable +import Data.List (partition) import qualified Data.Map.Strict as M -import Data.Maybe -import qualified Data.Set as S +import Data.Maybe import qualified Data.Sequence as Seq -import Data.Foldable - -import Futhark.MonadFreshNames -import Language.Futhark -import Language.Futhark.Traversals -import Language.Futhark.Semantic (TypeBinding(..)) -import Language.Futhark.TypeChecker.Types +import qualified Data.Set as S +import Futhark.MonadFreshNames +import Language.Futhark +import Language.Futhark.Semantic (TypeBinding (..)) +import Language.Futhark.Traversals +import Language.Futhark.TypeChecker.Types -i32 :: TypeBase dim als -i32 = Scalar $ Prim $ Signed Int32 +i64 :: TypeBase dim als +i64 = Scalar $ Prim $ Signed Int64 -- The monomorphization monad reads 'PolyBinding's and writes -- 'ValBind's. The 'TypeParam's in the 'ValBind's can only be size @@ -55,10 +54,19 @@ -- Each 'Polybinding' is also connected with the 'RecordReplacements' -- that were active when the binding was defined. This is used only -- in local functions. -data PolyBinding = PolyBinding RecordReplacements - (VName, [TypeParam], [Pattern], - Maybe (TypeExp VName), StructType, [VName], Exp, - [AttrInfo], SrcLoc) +data PolyBinding + = PolyBinding + RecordReplacements + ( VName, + [TypeParam], + [Pattern], + Maybe (TypeExp VName), + StructType, + [VName], + Exp, + [AttrInfo], + SrcLoc + ) -- Mapping from record names to the variable names that contain the -- fields. This is used because the monomorphiser also expands all @@ -69,41 +77,55 @@ -- Monomorphization environment mapping names of polymorphic functions -- to a representation of their corresponding function bindings. -data Env = Env { envPolyBindings :: M.Map VName PolyBinding - , envTypeBindings :: M.Map VName TypeBinding - , envRecordReplacements :: RecordReplacements - } +data Env = Env + { envPolyBindings :: M.Map VName PolyBinding, + envTypeBindings :: M.Map VName TypeBinding, + envRecordReplacements :: RecordReplacements + } instance Semigroup Env where Env tb1 pb1 rr1 <> Env tb2 pb2 rr2 = Env (tb1 <> tb2) (pb1 <> pb2) (rr1 <> rr2) instance Monoid Env where - mempty = Env mempty mempty mempty + mempty = Env mempty mempty mempty localEnv :: Env -> MonoM a -> MonoM a localEnv env = local (env <>) extendEnv :: VName -> PolyBinding -> MonoM a -> MonoM a -extendEnv vn binding = localEnv - mempty { envPolyBindings = M.singleton vn binding } +extendEnv vn binding = + localEnv + mempty {envPolyBindings = M.singleton vn binding} withRecordReplacements :: RecordReplacements -> MonoM a -> MonoM a -withRecordReplacements rr = localEnv mempty { envRecordReplacements = rr } +withRecordReplacements rr = localEnv mempty {envRecordReplacements = rr} replaceRecordReplacements :: RecordReplacements -> MonoM a -> MonoM a -replaceRecordReplacements rr = local $ \env -> env { envRecordReplacements = rr } +replaceRecordReplacements rr = local $ \env -> env {envRecordReplacements = rr} -- The monomorphization monad. -newtype MonoM a = MonoM (RWST Env (Seq.Seq (VName, ValBind)) VNameSource - (State Lifts) a) - deriving (Functor, Applicative, Monad, - MonadReader Env, - MonadWriter (Seq.Seq (VName, ValBind)), - MonadFreshNames) +newtype MonoM a + = MonoM + ( RWST + Env + (Seq.Seq (VName, ValBind)) + VNameSource + (State Lifts) + a + ) + deriving + ( Functor, + Applicative, + Monad, + MonadReader Env, + MonadWriter (Seq.Seq (VName, ValBind)), + MonadFreshNames + ) runMonoM :: VNameSource -> MonoM a -> ((a, Seq.Seq (VName, ValBind)), VNameSource) runMonoM src (MonoM m) = ((a, defs), src') - where (a, src', defs) = evalState (runRWST m mempty src) mempty + where + (a, src', defs) = evalState (runRWST m mempty src) mempty lookupFun :: VName -> MonoM (Maybe PolyBinding) lookupFun vn = do @@ -125,11 +147,12 @@ monoType :: TypeBase (DimDecl VName) als -> MonoType monoType = runIdentity . traverseDims onDim . toStruct - where onDim bound _ (NamedDim d) - -- A locally bound size. - | qualLeaf d `S.member` bound = pure False - onDim _ _ AnyDim = pure False - onDim _ _ _ = pure True + where + onDim bound _ (NamedDim d) + -- A locally bound size. + | qualLeaf d `S.member` bound = pure False + onDim _ _ AnyDim = pure False + onDim _ _ _ = pure True -- Mapping from function name and instance list to a new function name in case -- the function has already been instantiated with those concrete types. @@ -152,95 +175,109 @@ transformFName loc fname t | baseTag (qualLeaf fname) <= maxIntrinsicTag = return $ var fname | otherwise = do - maybe_fname <- lookupLifted (qualLeaf fname) (monoType t) - maybe_funbind <- lookupFun $ qualLeaf fname - t' <- removeTypeVariablesInType t - case (maybe_fname, maybe_funbind) of - -- The function has already been monomorphised. - (Just (fname', infer), _) -> - return $ applySizeArgs fname' t' $ infer t' - -- An intrinsic function. - (Nothing, Nothing) -> return $ var fname - -- A polymorphic function. - (Nothing, Just funbind) -> do - (fname', infer, funbind') <- monomorphiseBinding False funbind (monoType t') - tell $ Seq.singleton (qualLeaf fname, funbind') - addLifted (qualLeaf fname) (monoType t) (fname', infer) - return $ applySizeArgs fname' t' $ infer t' - - where var fname' = Var fname' (Info (fromStruct t)) loc + maybe_fname <- lookupLifted (qualLeaf fname) (monoType t) + maybe_funbind <- lookupFun $ qualLeaf fname + t' <- removeTypeVariablesInType t + case (maybe_fname, maybe_funbind) of + -- The function has already been monomorphised. + (Just (fname', infer), _) -> + return $ applySizeArgs fname' t' $ infer t' + -- An intrinsic function. + (Nothing, Nothing) -> return $ var fname + -- A polymorphic function. + (Nothing, Just funbind) -> do + (fname', infer, funbind') <- monomorphiseBinding False funbind (monoType t') + tell $ Seq.singleton (qualLeaf fname, funbind') + addLifted (qualLeaf fname) (monoType t) (fname', infer) + return $ applySizeArgs fname' t' $ infer t' + where + var fname' = Var fname' (Info (fromStruct t)) loc - applySizeArg (i, f) size_arg = - (i-1, - Apply f size_arg (Info (Observe, Nothing)) - (Info (foldFunType (replicate i i32) (fromStruct t)), Info []) - loc) + applySizeArg (i, f) size_arg = + ( i -1, + Apply + f + size_arg + (Info (Observe, Nothing)) + (Info (foldFunType (replicate i i64) (fromStruct t)), Info []) + loc + ) - applySizeArgs fname' t' size_args = - snd $ foldl' applySizeArg (length size_args - 1, - Var (qualName fname') - (Info (foldFunType (map (const i32) size_args) - (fromStruct t'))) - loc) + applySizeArgs fname' t' size_args = + snd $ + foldl' + applySizeArg + ( length size_args - 1, + Var + (qualName fname') + ( Info + ( foldFunType + (map (const i64) size_args) + (fromStruct t') + ) + ) + loc + ) size_args -- This carries out record replacements in the alias information of a type. transformType :: TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing) transformType t = do rrs <- asks envRecordReplacements - let replace (AliasBound v) | Just d <- M.lookup v rrs = - S.fromList $ map (AliasBound . fst) $ M.elems d + let replace (AliasBound v) + | Just d <- M.lookup v rrs = + S.fromList $ map (AliasBound . fst) $ M.elems d replace x = S.singleton x -- As an attempt at an optimisation, only transform the aliases if -- they refer to a variable we have record-replaced. - return $ if any ((`M.member` rrs) . aliasVar) $ aliases t - then second (mconcat . map replace . S.toList) t - else t + return $ + if any ((`M.member` rrs) . aliasVar) $ aliases t + then second (mconcat . map replace . S.toList) t + else t sizesForPat :: MonadFreshNames m => Pattern -> m ([VName], Pattern) sizesForPat pat = do (params', sizes) <- runStateT (astMap tv pat) [] return (sizes, params') - where tv = identityMapper { mapOnPatternType = bitraverse onDim pure } - onDim AnyDim = do v <- lift $ newVName "size" - modify (v:) - pure $ NamedDim $ qualName v - onDim d = pure d + where + tv = identityMapper {mapOnPatternType = bitraverse onDim pure} + onDim AnyDim = do + v <- lift $ newVName "size" + modify (v :) + pure $ NamedDim $ qualName v + onDim d = pure d -- Monomorphization of expressions. transformExp :: Exp -> MonoM Exp -transformExp e@Literal{} = return e -transformExp e@IntLit{} = return e -transformExp e@FloatLit{} = return e -transformExp e@StringLit{} = return e - +transformExp e@Literal {} = return e +transformExp e@IntLit {} = return e +transformExp e@FloatLit {} = return e +transformExp e@StringLit {} = return e transformExp (Parens e loc) = Parens <$> transformExp e <*> pure loc - transformExp (QualParens qn e loc) = QualParens qn <$> transformExp e <*> pure loc - transformExp (TupLit es loc) = TupLit <$> mapM transformExp es <*> pure loc - transformExp (RecordLit fs loc) = RecordLit <$> mapM transformField fs <*> pure loc - where transformField (RecordFieldExplicit name e loc') = - RecordFieldExplicit name <$> transformExp e <*> pure loc' - transformField (RecordFieldImplicit v t _) = do - t' <- traverse transformType t - transformField $ RecordFieldExplicit (baseName v) - (Var (qualName v) t' loc) loc - + where + transformField (RecordFieldExplicit name e loc') = + RecordFieldExplicit name <$> transformExp e <*> pure loc' + transformField (RecordFieldImplicit v t _) = do + t' <- traverse transformType t + transformField $ + RecordFieldExplicit + (baseName v) + (Var (qualName v) t' loc) + loc transformExp (ArrayLit es t loc) = ArrayLit <$> mapM transformExp es <*> traverse transformType t <*> pure loc - transformExp (Range e1 me incl tp loc) = do e1' <- transformExp e1 me' <- mapM transformExp me incl' <- mapM transformExp incl return $ Range e1' me' incl' tp loc - transformExp (Var fname (Info t) loc) = do maybe_fs <- lookupRecordReplacement $ qualLeaf fname case maybe_fs of @@ -253,105 +290,121 @@ Nothing -> do t' <- transformType t transformFName loc fname (toStruct t') - transformExp (Ascript e tp loc) = Ascript <$> transformExp e <*> pure tp <*> pure loc - transformExp (Coerce e tp (Info t, ext) loc) = do noticeDims t - Coerce <$> transformExp e <*> pure tp <*> - ((,) <$> (Info <$> transformType t) <*> pure ext) <*> pure loc - + Coerce <$> transformExp e <*> pure tp + <*> ((,) <$> (Info <$> transformType t) <*> pure ext) + <*> pure loc transformExp (LetPat pat e1 e2 (Info t, retext) loc) = do (pat', rr) <- transformPattern pat t' <- transformType t - LetPat pat' <$> transformExp e1 <*> - withRecordReplacements rr (transformExp e2) <*> - pure (Info t', retext) <*> pure loc - + LetPat pat' <$> transformExp e1 + <*> withRecordReplacements rr (transformExp e2) + <*> pure (Info t', retext) + <*> pure loc transformExp (LetFun fname (tparams, params, retdecl, Info ret, body) e e_t loc) | any isTypeParam tparams = do - -- Retrieve the lifted monomorphic function bindings that are produced, - -- filter those that are monomorphic versions of the current let-bound - -- function and insert them at this point, and propagate the rest. - rr <- asks envRecordReplacements - let funbind = PolyBinding rr (fname, tparams, params, retdecl, ret, [], body, mempty, loc) - pass $ do - (e', bs) <- listen $ extendEnv fname funbind $ transformExp e - -- Do not remember this one for next time we monomorphise this - -- function. - modifyLifts $ filter ((/=fname) . fst . fst) - let (bs_local, bs_prop) = Seq.partition ((== fname) . fst) bs - return (unfoldLetFuns (map snd $ toList bs_local) e', const bs_prop) - + -- Retrieve the lifted monomorphic function bindings that are produced, + -- filter those that are monomorphic versions of the current let-bound + -- function and insert them at this point, and propagate the rest. + rr <- asks envRecordReplacements + let funbind = PolyBinding rr (fname, tparams, params, retdecl, ret, [], body, mempty, loc) + pass $ do + (e', bs) <- listen $ extendEnv fname funbind $ transformExp e + -- Do not remember this one for next time we monomorphise this + -- function. + modifyLifts $ filter ((/= fname) . fst . fst) + let (bs_local, bs_prop) = Seq.partition ((== fname) . fst) bs + return (unfoldLetFuns (map snd $ toList bs_local) e', const bs_prop) | otherwise = do - body' <- transformExp body - LetFun fname (tparams, params, retdecl, Info ret, body') <$> - transformExp e <*> traverse transformType e_t <*> pure loc - + body' <- transformExp body + LetFun fname (tparams, params, retdecl, Info ret, body') + <$> transformExp e <*> traverse transformType e_t <*> pure loc transformExp (If e1 e2 e3 (tp, retext) loc) = do e1' <- transformExp e1 e2' <- transformExp e2 e3' <- transformExp e3 tp' <- traverse transformType tp return $ If e1' e2' e3' (tp', retext) loc - transformExp (Apply e1 e2 d (ret, ext) loc) = do e1' <- transformExp e1 e2' <- transformExp e2 ret' <- traverse transformType ret return $ Apply e1' e2' d (ret', ext) loc - transformExp (Negate e loc) = Negate <$> transformExp e <*> pure loc - transformExp (Lambda params e0 decl tp loc) = do e0' <- transformExp e0 return $ Lambda params e0' decl tp loc - transformExp (OpSection qn t loc) = transformExp $ Var qn t loc - -transformExp (OpSectionLeft fname (Info t) e - (Info (xtype, xargext), Info ytype) (Info rettype, Info retext) loc) = do - fname' <- transformFName loc fname $ toStruct t - e' <- transformExp e - desugarBinOpSection fname' (Just e') Nothing - t (xtype, xargext) (ytype, Nothing) (rettype, retext) loc - -transformExp (OpSectionRight fname (Info t) e - (Info xtype, Info (ytype, yargext)) (Info rettype) loc) = do - fname' <- transformFName loc fname $ toStruct t - e' <- transformExp e - desugarBinOpSection fname' Nothing (Just e') - t (xtype, Nothing) (ytype, yargext) (rettype, []) loc - +transformExp + ( OpSectionLeft + fname + (Info t) + e + (Info (xtype, xargext), Info ytype) + (Info rettype, Info retext) + loc + ) = do + fname' <- transformFName loc fname $ toStruct t + e' <- transformExp e + desugarBinOpSection + fname' + (Just e') + Nothing + t + (xtype, xargext) + (ytype, Nothing) + (rettype, retext) + loc +transformExp + ( OpSectionRight + fname + (Info t) + e + (Info xtype, Info (ytype, yargext)) + (Info rettype) + loc + ) = do + fname' <- transformFName loc fname $ toStruct t + e' <- transformExp e + desugarBinOpSection + fname' + Nothing + (Just e') + t + (xtype, Nothing) + (ytype, yargext) + (rettype, []) + loc transformExp (ProjectSection fields (Info t) loc) = desugarProjectSection fields t loc - transformExp (IndexSection idxs (Info t) loc) = desugarIndexSection idxs t loc - transformExp (DoLoop sparams pat e1 form e3 ret loc) = do e1' <- transformExp e1 form' <- case form of - For ident e2 -> For ident <$> transformExp e2 + For ident e2 -> For ident <$> transformExp e2 ForIn pat2 e2 -> ForIn pat2 <$> transformExp e2 - While e2 -> While <$> transformExp e2 + While e2 -> While <$> transformExp e2 e3' <- transformExp e3 -- Maybe monomorphisation introduced new arrays to the loop, and -- maybe they have AnyDim sizes. This is not allowed. Invent some -- sizes for them. (pat_sizes, pat') <- sizesForPat pat - return $ DoLoop (sparams++pat_sizes) pat' e1' form' e3' ret loc - + return $ DoLoop (sparams ++ pat_sizes) pat' e1' form' e3' ret loc transformExp (BinOp (fname, oploc) (Info t) (e1, d1) (e2, d2) tp ext loc) = do fname' <- transformFName loc fname $ toStruct t e1' <- transformExp e1 e2' <- transformExp e2 case fname' of - Var fname'' _ _ | orderZero (typeOf e1'), orderZero (typeOf e2') -> - return $ BinOp (fname'', oploc) (Info t) (e1', d1) (e2', d2) tp ext loc + Var fname'' _ _ + | orderZero (typeOf e1'), + orderZero (typeOf e2') -> + return $ BinOp (fname'', oploc) (Info t) (e1', d1) (e2', d2) tp ext loc _ -> do -- We have to flip the arguments to the function, because -- operator application is left-to-right, while function @@ -361,62 +414,79 @@ -- involves existential sizes will necessarily go through here. (x_param_e, x_param) <- makeVarParam e1' (y_param_e, y_param) <- makeVarParam e2' - return $ LetPat x_param e1' - (LetPat y_param e2' - (applyOp fname' x_param_e y_param_e) (tp, Info mempty) mempty) - (tp, Info mempty) mempty - where applyOp fname' x y = - Apply (Apply fname' x (Info (Observe, snd (unInfo d1))) - (Info (foldFunType [fromStruct $ fst (unInfo d2)] (unInfo tp)), - Info mempty) loc) - y (Info (Observe, snd (unInfo d2))) (tp, ext) loc - - makeVarParam arg = do - let argtype = typeOf arg - x <- newNameFromString "binop_p" - return (Var (qualName x) (Info argtype) mempty, - Id x (Info $ fromStruct argtype) mempty) - + return $ + LetPat + x_param + e1' + ( LetPat + y_param + e2' + (applyOp fname' x_param_e y_param_e) + (tp, Info mempty) + mempty + ) + (tp, Info mempty) + mempty + where + applyOp fname' x y = + Apply + ( Apply + fname' + x + (Info (Observe, snd (unInfo d1))) + ( Info (foldFunType [fromStruct $ fst (unInfo d2)] (unInfo tp)), + Info mempty + ) + loc + ) + y + (Info (Observe, snd (unInfo d2))) + (tp, ext) + loc + makeVarParam arg = do + let argtype = typeOf arg + x <- newNameFromString "binop_p" + return + ( Var (qualName x) (Info argtype) mempty, + Id x (Info $ fromStruct argtype) mempty + ) transformExp (Project n e tp loc) = do maybe_fs <- case e of Var qn _ _ -> lookupRecordReplacement (qualLeaf qn) - _ -> return Nothing + _ -> return Nothing case maybe_fs of - Just m | Just (v, _) <- M.lookup n m -> - return $ Var (qualName v) tp loc + Just m + | Just (v, _) <- M.lookup n m -> + return $ Var (qualName v) tp loc _ -> do e' <- transformExp e return $ Project n e' tp loc - transformExp (LetWith id1 id2 idxs e1 body (Info t) loc) = do idxs' <- mapM transformDimIndex idxs e1' <- transformExp e1 body' <- transformExp body t' <- transformType t return $ LetWith id1 id2 idxs' e1' body' (Info t') loc - transformExp (Index e0 idxs info loc) = Index <$> transformExp e0 <*> mapM transformDimIndex idxs <*> pure info <*> pure loc - transformExp (Update e1 idxs e2 loc) = Update <$> transformExp e1 <*> mapM transformDimIndex idxs - <*> transformExp e2 <*> pure loc - + <*> transformExp e2 + <*> pure loc transformExp (RecordUpdate e1 fs e2 t loc) = RecordUpdate <$> transformExp e1 <*> pure fs - <*> transformExp e2 <*> pure t <*> pure loc - + <*> transformExp e2 + <*> pure t + <*> pure loc transformExp (Assert e1 e2 desc loc) = Assert <$> transformExp e1 <*> transformExp e2 <*> pure desc <*> pure loc - transformExp (Constr name all_es t loc) = Constr name <$> mapM transformExp all_es <*> pure t <*> pure loc - transformExp (Match e cs (t, retext) loc) = - Match <$> transformExp e <*> mapM transformCase cs <*> - ((,) <$> traverse transformType t <*> pure retext) <*> pure loc - + Match <$> transformExp e <*> mapM transformCase cs + <*> ((,) <$> traverse transformType t <*> pure retext) + <*> pure loc transformExp (Attr info e loc) = Attr info <$> transformExp e <*> pure loc @@ -429,54 +499,78 @@ transformDimIndex (DimFix e) = DimFix <$> transformExp e transformDimIndex (DimSlice me1 me2 me3) = DimSlice <$> trans me1 <*> trans me2 <*> trans me3 - where trans = mapM transformExp + where + trans = mapM transformExp -- Transform an operator section into a lambda. -desugarBinOpSection :: Exp -> Maybe Exp -> Maybe Exp - -> PatternType - -> (StructType, Maybe VName) -> (StructType, Maybe VName) - -> (PatternType, [VName]) -> SrcLoc -> MonoM Exp +desugarBinOpSection :: + Exp -> + Maybe Exp -> + Maybe Exp -> + PatternType -> + (StructType, Maybe VName) -> + (StructType, Maybe VName) -> + (PatternType, [VName]) -> + SrcLoc -> + MonoM Exp desugarBinOpSection op e_left e_right t (xtype, xext) (ytype, yext) (rettype, retext) loc = do (e1, p1) <- makeVarParam e_left $ fromStruct xtype (e2, p2) <- makeVarParam e_right $ fromStruct ytype - let apply_left = Apply op e1 (Info (Observe, xext)) - (Info $ foldFunType [fromStruct ytype] t, Info []) loc - body = Apply apply_left e2 (Info (Observe, yext)) - (Info rettype, Info retext) loc + let apply_left = + Apply + op + e1 + (Info (Observe, xext)) + (Info $ foldFunType [fromStruct ytype] t, Info []) + loc + body = + Apply + apply_left + e2 + (Info (Observe, yext)) + (Info rettype, Info retext) + loc rettype' = toStruct rettype return $ Lambda (p1 ++ p2) body Nothing (Info (mempty, rettype')) loc - - where makeVarParam (Just e) _ = return (e, []) - makeVarParam Nothing argtype = do - x <- newNameFromString "x" - return (Var (qualName x) (Info argtype) mempty, - [Id x (Info $ fromStruct argtype) mempty]) + where + makeVarParam (Just e) _ = return (e, []) + makeVarParam Nothing argtype = do + x <- newNameFromString "x" + return + ( Var (qualName x) (Info argtype) mempty, + [Id x (Info $ fromStruct argtype) mempty] + ) desugarProjectSection :: [Name] -> PatternType -> SrcLoc -> MonoM Exp desugarProjectSection fields (Scalar (Arrow _ _ t1 t2)) loc = do p <- newVName "project_p" let body = foldl project (Var (qualName p) (Info t1) mempty) fields return $ Lambda [Id p (Info t1) mempty] body Nothing (Info (mempty, toStruct t2)) loc - where project e field = - case typeOf e of - Scalar (Record fs) - | Just t <- M.lookup field fs -> - Project field e (Info t) mempty - t -> error $ "desugarOpSection: type " ++ pretty t ++ - " does not have field " ++ pretty field -desugarProjectSection _ t _ = error $ "desugarOpSection: not a function type: " ++ pretty t + where + project e field = + case typeOf e of + Scalar (Record fs) + | Just t <- M.lookup field fs -> + Project field e (Info t) mempty + t -> + error $ + "desugarOpSection: type " ++ pretty t + ++ " does not have field " + ++ pretty field +desugarProjectSection _ t _ = error $ "desugarOpSection: not a function type: " ++ pretty t desugarIndexSection :: [DimIndex] -> PatternType -> SrcLoc -> MonoM Exp desugarIndexSection idxs (Scalar (Arrow _ _ t1 t2)) loc = do p <- newVName "index_i" let body = Index (Var (qualName p) (Info t1) loc) idxs (Info t2, Info []) loc return $ Lambda [Id p (Info t1) mempty] body Nothing (Info (mempty, toStruct t2)) loc -desugarIndexSection _ t _ = error $ "desugarIndexSection: not a function type: " ++ pretty t +desugarIndexSection _ t _ = error $ "desugarIndexSection: not a function type: " ++ pretty t noticeDims :: TypeBase (DimDecl VName) as -> MonoM () noticeDims = mapM_ notice . nestedDims - where notice (NamedDim v) = void $ transformFName mempty v i32 - notice _ = return () + where + notice (NamedDim v) = void $ transformFName mempty v i64 + notice _ = return () -- Convert a collection of 'ValBind's to a nested sequence of let-bound, -- monomorphic functions with the given expression at the bottom. @@ -484,18 +578,25 @@ unfoldLetFuns [] e = e unfoldLetFuns (ValBind _ fname _ (Info (rettype, _)) dim_params params body _ _ loc : rest) e = LetFun fname (dim_params, params, Nothing, Info rettype, body) e' (Info e_t) loc - where e' = unfoldLetFuns rest e - e_t = typeOf e' + where + e' = unfoldLetFuns rest e + e_t = typeOf e' transformPattern :: Pattern -> MonoM (Pattern, RecordReplacements) transformPattern (Id v (Info (Scalar (Record fs))) loc) = do let fs' = M.toList fs - (fs_ks, fs_ts) <- fmap unzip $ forM fs' $ \(f, ft) -> - (,) <$> newVName (nameToString f) <*> transformType ft - return (RecordPattern (zip (map fst fs') - (zipWith3 Id fs_ks (map Info fs_ts) $ repeat loc)) - loc, - M.singleton v $ M.fromList $ zip (map fst fs') $ zip fs_ks fs_ts) + (fs_ks, fs_ts) <- fmap unzip $ + forM fs' $ \(f, ft) -> + (,) <$> newVName (nameToString f) <*> transformType ft + return + ( RecordPattern + ( zip + (map fst fs') + (zipWith3 Id fs_ks (map Info fs_ts) $ repeat loc) + ) + loc, + M.singleton v $ M.fromList $ zip (map fst fs') $ zip fs_ks fs_ts + ) transformPattern (Id v t loc) = return (Id v t loc, mempty) transformPattern (TuplePattern pats loc) = do (pats', rrs) <- unzip <$> mapM transformPattern pats @@ -526,36 +627,40 @@ type DimInst = M.Map VName (DimDecl VName) -dimMapping :: Monoid a => - TypeBase (DimDecl VName) a - -> TypeBase (DimDecl VName) a - -> DimInst +dimMapping :: + Monoid a => + TypeBase (DimDecl VName) a -> + TypeBase (DimDecl VName) a -> + DimInst dimMapping t1 t2 = execState (matchDims f t1 t2) mempty - where f (NamedDim d1) d2 = do - modify $ M.insert (qualLeaf d1) d2 - return $ NamedDim d1 - f d _ = return d + where + f (NamedDim d1) d2 = do + modify $ M.insert (qualLeaf d1) d2 + return $ NamedDim d1 + f d _ = return d inferSizeArgs :: [TypeParam] -> StructType -> StructType -> [Exp] inferSizeArgs tparams bind_t t = mapMaybe (tparamArg (dimMapping bind_t t)) tparams - where tparamArg dinst tp = - case M.lookup (typeParamName tp) dinst of - Just (NamedDim d) -> - Just $ Var d (Info i32) mempty - Just (ConstDim x) -> - Just $ Literal (SignedValue $ Int32Value $ fromIntegral x) mempty - _ -> - Nothing + where + tparamArg dinst tp = + case M.lookup (typeParamName tp) dinst of + Just (NamedDim d) -> + Just $ Var d (Info i64) mempty + 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, True) -> modify $ S.insert $ qualLeaf v - _ -> return () - return d1 + where + onDims d1 d2 = do + case (d1, d2) of + (NamedDim v, True) -> modify $ S.insert $ qualLeaf v + _ -> return () + return d1 -- Monomorphising higher-order functions can result in function types -- where the same named parameter occurs in multiple spots. When @@ -566,128 +671,153 @@ -- anyway. noNamedParams :: MonoType -> MonoType noNamedParams = f - where f (Array () u t shape) = Array () u (f' t) shape - f (Scalar t) = Scalar $ f' t - f' (Arrow () _ t1 t2) = - Arrow () Unnamed (f t1) (f t2) - f' (Record fs) = - Record $ fmap f fs - f' (Sum cs) = - Sum $ fmap (map f) cs - f' t = t + where + f (Array () u t shape) = Array () u (f' t) shape + f (Scalar t) = Scalar $ f' t + f' (Arrow () _ t1 t2) = + Arrow () Unnamed (f t1) (f t2) + f' (Record fs) = + Record $ fmap f fs + f' (Sum cs) = + Sum $ fmap (map f) cs + f' t = t -- Monomorphise a polymorphic function at the types given in the instance -- list. Monomorphises the body of the function as well. Returns the fresh name -- of the generated monomorphic function and its 'ValBind' representation. -monomorphiseBinding :: Bool -> PolyBinding -> MonoType - -> MonoM (VName, InferSizeArgs, ValBind) +monomorphiseBinding :: + Bool -> + PolyBinding -> + MonoType -> + MonoM (VName, InferSizeArgs, ValBind) monomorphiseBinding entry (PolyBinding rr (name, tparams, params, retdecl, rettype, retext, body, attrs, loc)) t = replaceRecordReplacements rr $ do - let bind_t = foldFunType (map patternStructType params) rettype - (substs, t_shape_params) <- typeSubstsM loc (noSizes bind_t) $ noNamedParams t - let substs' = M.map Subst substs - rettype' = substTypesAny (`M.lookup` substs') rettype - substPatternType = - substTypesAny (fmap (fmap fromStruct) . (`M.lookup` substs')) - 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) $ - shape_params ++ t_shape_params - - (params'', rrs) <- unzip <$> mapM transformPattern params' + let bind_t = foldFunType (map patternStructType params) rettype + (substs, t_shape_params) <- typeSubstsM loc (noSizes bind_t) $ noNamedParams t + let substs' = M.map Subst substs + rettype' = substTypesAny (`M.lookup` substs') rettype + substPatternType = + substTypesAny (fmap (fmap fromStruct) . (`M.lookup` substs')) + 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) $ + shape_params ++ t_shape_params - mapM_ noticeDims $ rettype : map patternStructType params'' + (params'', rrs) <- unzip <$> mapM transformPattern params' - body' <- updateExpTypes (`M.lookup` substs') body - body'' <- withRecordReplacements (mconcat rrs) $ transformExp body' - name' <- if null tparams && not entry then return name else newName name + mapM_ noticeDims $ rettype : map patternStructType params'' - return (name', - inferSizeArgs shape_params_explicit bind_t', - if entry - then toValBinding name' - (shape_params_explicit++shape_params_implicit) params'' - (rettype', retext) body'' - else toValBinding name' shape_params_implicit - (map shapeParam shape_params_explicit ++ params'') - (rettype', retext) body'') + body' <- updateExpTypes (`M.lookup` substs') body + body'' <- withRecordReplacements (mconcat rrs) $ transformExp body' + name' <- if null tparams && not entry then return name else newName name - where shape_params = filter (not . isTypeParam) tparams + return + ( name', + inferSizeArgs shape_params_explicit bind_t', + if entry + then + toValBinding + name' + (shape_params_explicit ++ shape_params_implicit) + params'' + (rettype', retext) + body'' + else + toValBinding + name' + shape_params_implicit + (map shapeParam shape_params_explicit ++ params'') + (rettype', retext) + body'' + ) + where + shape_params = filter (not . isTypeParam) tparams - updateExpTypes substs = astMap $ mapper substs - mapper substs = ASTMapper { mapOnExp = astMap $ mapper substs - , mapOnName = pure - , mapOnQualName = pure - , mapOnStructType = pure . applySubst substs - , mapOnPatternType = pure . applySubst substs - } + updateExpTypes substs = astMap $ mapper substs + mapper substs = + ASTMapper + { mapOnExp = astMap $ mapper substs, + mapOnName = pure, + mapOnQualName = pure, + mapOnStructType = pure . applySubst substs, + mapOnPatternType = pure . applySubst substs + } - shapeParam tp = Id (typeParamName tp) (Info i32) $ srclocOf tp + shapeParam tp = Id (typeParamName tp) (Info i64) $ srclocOf tp - toValBinding name' tparams' params'' rettype' body'' = - ValBind { valBindEntryPoint = Nothing - , valBindName = name' - , valBindRetDecl = retdecl - , valBindRetType = Info rettype' - , valBindTypeParams = tparams' - , valBindParams = params'' - , valBindBody = body'' - , valBindDoc = Nothing - , valBindAttrs = attrs - , valBindLocation = loc - } + toValBinding name' tparams' params'' rettype' body'' = + ValBind + { valBindEntryPoint = Nothing, + valBindName = name', + valBindRetDecl = retdecl, + valBindRetType = Info rettype', + valBindTypeParams = tparams', + valBindParams = params'', + valBindBody = body'', + valBindDoc = Nothing, + valBindAttrs = attrs, + valBindLocation = loc + } -typeSubstsM :: MonadFreshNames m => - SrcLoc -> TypeBase () () -> MonoType - -> m (M.Map VName StructType, [TypeParam]) +typeSubstsM :: + MonadFreshNames m => + SrcLoc -> + TypeBase () () -> + MonoType -> + m (M.Map VName StructType, [TypeParam]) typeSubstsM loc orig_t1 orig_t2 = let m = sub orig_t1 orig_t2 - in runWriterT $ execStateT m mempty - - where sub t1@Array{} t2@Array{} - | Just t1' <- peelArray (arrayRank t1) t1, - Just t2' <- peelArray (arrayRank t1) t2 = - sub t1' t2' - sub (Scalar (TypeVar _ _ v _)) t = addSubst v t - sub (Scalar (Record fields1)) (Scalar (Record fields2)) = - zipWithM_ sub - (map snd $ sortFields fields1) (map snd $ sortFields fields2) - sub (Scalar Prim{}) (Scalar Prim{}) = return () - sub (Scalar (Arrow _ _ t1a t1b)) (Scalar (Arrow _ _ t2a t2b)) = do - sub t1a t2a - sub t1b t2b - sub (Scalar (Sum cs1)) (Scalar (Sum cs2)) = - zipWithM_ typeSubstClause (sortConstrs cs1) (sortConstrs cs2) - where typeSubstClause (_, ts1) (_, ts2) = zipWithM sub ts1 ts2 - sub t1@(Scalar Sum{}) t2 = sub t1 t2 - sub t1 t2@(Scalar Sum{}) = sub t1 t2 - - sub t1 t2 = error $ unlines ["typeSubstsM: mismatched types:", pretty t1, pretty t2] + in runWriterT $ execStateT m mempty + where + sub t1@Array {} t2@Array {} + | Just t1' <- peelArray (arrayRank t1) t1, + Just t2' <- peelArray (arrayRank t1) t2 = + sub t1' t2' + sub (Scalar (TypeVar _ _ v _)) t = addSubst v t + sub (Scalar (Record fields1)) (Scalar (Record fields2)) = + zipWithM_ + sub + (map snd $ sortFields fields1) + (map snd $ sortFields fields2) + sub (Scalar Prim {}) (Scalar Prim {}) = return () + sub (Scalar (Arrow _ _ t1a t1b)) (Scalar (Arrow _ _ t2a t2b)) = do + sub t1a t2a + sub t1b t2b + sub (Scalar (Sum cs1)) (Scalar (Sum cs2)) = + zipWithM_ typeSubstClause (sortConstrs cs1) (sortConstrs cs2) + where + typeSubstClause (_, ts1) (_, ts2) = zipWithM sub ts1 ts2 + sub t1@(Scalar Sum {}) t2 = sub t1 t2 + sub t1 t2@(Scalar Sum {}) = sub t1 t2 + sub t1 t2 = error $ unlines ["typeSubstsM: mismatched types:", pretty t1, pretty t2] - addSubst (TypeName _ v) t = do - exists <- gets $ M.member v - unless exists $ do - t' <- bitraverse onDim pure t - modify $ M.insert v t' + addSubst (TypeName _ v) t = do + exists <- gets $ M.member v + unless exists $ do + t' <- bitraverse onDim pure t + modify $ M.insert v t' - onDim True = do d <- lift $ lift $ newVName "d" - tell [TypeParamDim d loc] - return $ NamedDim $ qualName d - onDim False = return AnyDim + onDim True = do + d <- lift $ lift $ newVName "d" + tell [TypeParamDim d loc] + return $ NamedDim $ qualName d + onDim False = return AnyDim -- Perform a given substitution on the types in a pattern. substPattern :: Bool -> (PatternType -> PatternType) -> Pattern -> Pattern substPattern entry f pat = case pat of - TuplePattern pats loc -> TuplePattern (map (substPattern entry f) pats) loc - RecordPattern fs loc -> RecordPattern (map substField fs) loc - where substField (n, p) = (n, substPattern entry f p) - PatternParens p loc -> PatternParens (substPattern entry f p) loc - Id vn (Info tp) loc -> Id vn (Info $ f tp) loc - Wildcard (Info tp) loc -> Wildcard (Info $ f tp) loc - PatternAscription p td loc | entry -> PatternAscription (substPattern False f p) td loc - | otherwise -> substPattern False f p - PatternLit e (Info tp) loc -> PatternLit e (Info $ f tp) loc + TuplePattern pats loc -> TuplePattern (map (substPattern entry f) pats) loc + RecordPattern fs loc -> RecordPattern (map substField fs) loc + where + substField (n, p) = (n, substPattern entry f p) + PatternParens p loc -> PatternParens (substPattern entry f p) loc + Id vn (Info tp) loc -> Id vn (Info $ f tp) loc + Wildcard (Info tp) loc -> Wildcard (Info $ f tp) loc + PatternAscription p td loc + | entry -> PatternAscription (substPattern False f p) td loc + | otherwise -> substPattern False f p + PatternLit e (Info tp) loc -> PatternLit e (Info $ f tp) loc PatternConstr n (Info tp) ps loc -> PatternConstr n (Info $ f tp) ps loc toPolyBinding :: ValBind -> PolyBinding @@ -698,20 +828,23 @@ removeTypeVariables :: Bool -> ValBind -> MonoM ValBind removeTypeVariables entry valbind@(ValBind _ _ _ (Info (rettype, retext)) _ pats body _ _ _) = do subs <- asks $ M.map TypeSub . envTypeBindings - let mapper = ASTMapper { - mapOnExp = astMap mapper - , mapOnName = pure - , mapOnQualName = pure - , mapOnStructType = pure . substituteTypes subs - , mapOnPatternType = pure . substituteTypes subs - } + let mapper = + ASTMapper + { mapOnExp = astMap mapper, + mapOnName = pure, + mapOnQualName = pure, + mapOnStructType = pure . substituteTypes subs, + mapOnPatternType = pure . substituteTypes subs + } body' <- astMap mapper body - return valbind { valBindRetType = Info (substituteTypes subs rettype, retext) - , valBindParams = map (substPattern entry $ substituteTypes subs) pats - , valBindBody = body' - } + return + valbind + { valBindRetType = Info (substituteTypes subs rettype, retext), + valBindParams = map (substPattern entry $ substituteTypes subs) pats, + valBindBody = body' + } removeTypeVariablesInType :: StructType -> MonoM StructType removeTypeVariablesInType t = do @@ -720,17 +853,20 @@ transformValBind :: ValBind -> MonoM Env transformValBind valbind = do - valbind' <- toPolyBinding <$> - removeTypeVariables (isJust (valBindEntryPoint valbind)) valbind + valbind' <- + toPolyBinding + <$> removeTypeVariables (isJust (valBindEntryPoint valbind)) valbind when (isJust $ valBindEntryPoint valbind) $ do - t <- removeTypeVariablesInType $ foldFunType - (map patternStructType (valBindParams valbind)) $ - fst $ unInfo $ valBindRetType valbind + t <- + removeTypeVariablesInType $ + foldFunType + (map patternStructType (valBindParams valbind)) + $ fst $ unInfo $ valBindRetType valbind (name, _, valbind'') <- monomorphiseBinding True valbind' $ monoType t - tell $ Seq.singleton (name, valbind'' { valBindEntryPoint = valBindEntryPoint valbind}) + tell $ Seq.singleton (name, valbind'' {valBindEntryPoint = valBindEntryPoint valbind}) - return mempty { envPolyBindings = M.singleton (valBindName valbind) valbind' } + return mempty {envPolyBindings = M.singleton (valBindName valbind) valbind'} transformTypeBind :: TypeBind -> MonoM Env transformTypeBind (TypeBind name l tparams tydecl _ _) = do @@ -738,25 +874,26 @@ noticeDims $ unInfo $ expandedType tydecl let tp = substituteTypes subs . unInfo $ expandedType tydecl tbinding = TypeAbbr l tparams tp - return mempty { envTypeBindings = M.singleton name tbinding } + return mempty {envTypeBindings = M.singleton name tbinding} transformDecs :: [Dec] -> MonoM () transformDecs [] = return () transformDecs (ValDec valbind : ds) = do env <- transformValBind valbind localEnv env $ transformDecs ds - transformDecs (TypeDec typebind : ds) = do env <- transformTypeBind typebind localEnv env $ transformDecs ds - transformDecs (dec : _) = - error $ "The monomorphization module expects a module-free " ++ - "input program, but received: " ++ pretty dec + error $ + "The monomorphization module expects a module-free " + ++ "input program, but received: " + ++ pretty dec -- | Monomorphise a list of top-level declarations. A module-free input program -- is expected, so only value declarations and type declaration are accepted. transformProg :: MonadFreshNames m => [Dec] -> m [ValBind] transformProg decs = - fmap (toList . fmap snd . snd) $ modifyNameSource $ \namesrc -> - runMonoM namesrc $ transformDecs decs + fmap (toList . fmap snd . snd) $ + modifyNameSource $ \namesrc -> + runMonoM namesrc $ transformDecs decs diff --git a/src/Futhark/Internalise/TypesValues.hs b/src/Futhark/Internalise/TypesValues.hs --- a/src/Futhark/Internalise/TypesValues.hs +++ b/src/Futhark/Internalise/TypesValues.hs @@ -1,32 +1,32 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE Trustworthy #-} + module Futhark.Internalise.TypesValues - ( - -- * Internalising types - internaliseReturnType - , internaliseLambdaReturnType - , internaliseEntryReturnType - , internaliseType - , internaliseParamTypes - , internaliseLoopParamType - , internalisePrimType - , internalisedTypeSize - , internaliseSumType + ( -- * Internalising types + internaliseReturnType, + internaliseLambdaReturnType, + internaliseEntryReturnType, + internaliseType, + internaliseParamTypes, + internaliseLoopParamType, + internalisePrimType, + internalisedTypeSize, + internaliseSumType, - -- * Internalising values - , internalisePrimValue + -- * Internalising values + internalisePrimValue, ) - where +where + import Control.Monad.Reader import Control.Monad.State import Data.List (delete, find, foldl') -import Data.Maybe import qualified Data.Map.Strict as M - -import qualified Language.Futhark as E +import Data.Maybe import Futhark.IR.SOACS as I import Futhark.Internalise.Monad +import qualified Language.Futhark as E internaliseUniqueness :: E.Uniqueness -> I.Uniqueness internaliseUniqueness E.Nonunique = I.Nonunique @@ -34,79 +34,92 @@ type TypeState = Int -newtype InternaliseTypeM a = - InternaliseTypeM (StateT TypeState InternaliseM a) +newtype InternaliseTypeM a + = InternaliseTypeM (StateT TypeState InternaliseM a) deriving (Functor, Applicative, Monad, MonadState TypeState) liftInternaliseM :: InternaliseM a -> InternaliseTypeM a liftInternaliseM = InternaliseTypeM . lift -runInternaliseTypeM :: InternaliseTypeM a - -> InternaliseM a +runInternaliseTypeM :: + InternaliseTypeM a -> + InternaliseM a runInternaliseTypeM (InternaliseTypeM m) = evalStateT m 0 -internaliseParamTypes :: [E.TypeBase (E.DimDecl VName) ()] - -> InternaliseM [[I.TypeBase Shape Uniqueness]] +internaliseParamTypes :: + [E.TypeBase (E.DimDecl VName) ()] -> + InternaliseM [[I.TypeBase Shape Uniqueness]] internaliseParamTypes ts = runInternaliseTypeM $ mapM (fmap (map onType) . internaliseTypeM) ts - where onType = fromMaybe bad . hasStaticShape - bad = error $ "internaliseParamTypes: " ++ pretty ts + where + onType = fromMaybe bad . hasStaticShape + bad = error $ "internaliseParamTypes: " ++ pretty ts -internaliseLoopParamType :: E.TypeBase (E.DimDecl VName) () - -> InternaliseM [I.TypeBase Shape Uniqueness] +internaliseLoopParamType :: + E.TypeBase (E.DimDecl VName) () -> + InternaliseM [I.TypeBase Shape Uniqueness] internaliseLoopParamType et = concat <$> internaliseParamTypes [et] -internaliseReturnType :: E.TypeBase (E.DimDecl VName) () - -> InternaliseM [I.TypeBase ExtShape Uniqueness] +internaliseReturnType :: + E.TypeBase (E.DimDecl VName) () -> + InternaliseM [I.TypeBase ExtShape Uniqueness] internaliseReturnType et = runInternaliseTypeM (internaliseTypeM et) -internaliseLambdaReturnType :: E.TypeBase (E.DimDecl VName) () - -> InternaliseM [I.TypeBase Shape NoUniqueness] +internaliseLambdaReturnType :: + E.TypeBase (E.DimDecl VName) () -> + InternaliseM [I.TypeBase Shape NoUniqueness] internaliseLambdaReturnType = fmap (map fromDecl) . internaliseLoopParamType -- | As 'internaliseReturnType', but returns components of a top-level -- tuple type piecemeal. -internaliseEntryReturnType :: E.TypeBase (E.DimDecl VName) () - -> InternaliseM [[I.TypeBase ExtShape Uniqueness]] +internaliseEntryReturnType :: + E.TypeBase (E.DimDecl VName) () -> + InternaliseM [[I.TypeBase ExtShape Uniqueness]] internaliseEntryReturnType et = - runInternaliseTypeM $ mapM internaliseTypeM $ - case E.isTupleRecord et of - Just ets | not $ null ets -> ets - _ -> [et] + runInternaliseTypeM $ + mapM internaliseTypeM $ + case E.isTupleRecord et of + Just ets | not $ null ets -> ets + _ -> [et] -internaliseType :: E.TypeBase (E.DimDecl VName) () - -> InternaliseM [I.TypeBase I.ExtShape Uniqueness] +internaliseType :: + E.TypeBase (E.DimDecl VName) () -> + InternaliseM [I.TypeBase I.ExtShape Uniqueness] internaliseType = runInternaliseTypeM . internaliseTypeM newId :: InternaliseTypeM Int -newId = do i <- get - put $ i + 1 - return i +newId = do + i <- get + put $ i + 1 + return i -internaliseDim :: E.DimDecl VName - -> InternaliseTypeM ExtSize +internaliseDim :: + E.DimDecl VName -> + InternaliseTypeM ExtSize internaliseDim d = case d of E.AnyDim -> Ext <$> newId - E.ConstDim n -> return $ Free $ intConst I.Int32 $ toInteger n + E.ConstDim n -> return $ Free $ intConst I.Int64 $ toInteger n E.NamedDim name -> namedDim name - where namedDim (E.QualName _ name) = do - subst <- liftInternaliseM $ lookupSubst name - case subst of - Just [v] -> return $ I.Free v - _ -> return $ I.Free $ I.Var name + where + namedDim (E.QualName _ name) = do + subst <- liftInternaliseM $ lookupSubst name + case subst of + Just [v] -> return $ I.Free v + _ -> return $ I.Free $ I.Var name -internaliseTypeM :: E.StructType - -> InternaliseTypeM [I.TypeBase ExtShape Uniqueness] +internaliseTypeM :: + E.StructType -> + InternaliseTypeM [I.TypeBase ExtShape Uniqueness] internaliseTypeM orig_t = case orig_t of E.Array _ u et shape -> do dims <- internaliseShape shape ets <- internaliseTypeM $ E.Scalar et - return [I.arrayOf et' (Shape dims) $ internaliseUniqueness u | et' <- ets ] + return [I.arrayOf et' (Shape dims) $ internaliseUniqueness u | et' <- ets] E.Scalar (E.Prim bt) -> return [I.Prim $ internalisePrimType bt] E.Scalar (E.Record ets) @@ -114,43 +127,54 @@ -- arrays of unit will lose their sizes. | null ets -> return [I.Prim I.Bool] | otherwise -> - concat <$> mapM (internaliseTypeM . snd) (E.sortFields ets) - E.Scalar E.TypeVar{} -> + concat <$> mapM (internaliseTypeM . snd) (E.sortFields ets) + E.Scalar E.TypeVar {} -> error "internaliseTypeM: cannot handle type variable." - E.Scalar E.Arrow{} -> + E.Scalar E.Arrow {} -> error $ "internaliseTypeM: cannot handle function type: " ++ pretty orig_t E.Scalar (E.Sum cs) -> do - (ts, _) <- internaliseConstructors <$> - traverse (fmap concat . mapM internaliseTypeM) cs + (ts, _) <- + internaliseConstructors + <$> traverse (fmap concat . mapM internaliseTypeM) cs return $ I.Prim (I.IntType I.Int8) : ts - - where internaliseShape = mapM internaliseDim . E.shapeDims + where + internaliseShape = mapM internaliseDim . E.shapeDims -internaliseConstructors :: M.Map Name [I.TypeBase ExtShape Uniqueness] - -> ([I.TypeBase ExtShape Uniqueness], - M.Map Name (Int, [Int])) +internaliseConstructors :: + M.Map Name [I.TypeBase ExtShape Uniqueness] -> + ( [I.TypeBase ExtShape Uniqueness], + M.Map Name (Int, [Int]) + ) internaliseConstructors cs = - foldl' onConstructor mempty $ zip (E.sortConstrs cs) [0..] - where onConstructor (ts, mapping) ((c, c_ts), i) = - let (_, js, new_ts) = - foldl' f (zip ts [0..], mempty, mempty) c_ts - in (ts ++ new_ts, M.insert c (i, js) mapping) - where f (ts', js, new_ts) t - | Just (_, j) <- find ((==t) . fst) ts' = - (delete (t, j) ts', - js ++ [j], - new_ts) - | otherwise = - (ts', - js ++ [length ts + length new_ts], - new_ts ++ [t]) + foldl' onConstructor mempty $ zip (E.sortConstrs cs) [0 ..] + where + onConstructor (ts, mapping) ((c, c_ts), i) = + let (_, js, new_ts) = + foldl' f (zip ts [0 ..], mempty, mempty) c_ts + in (ts ++ new_ts, M.insert c (i, js) mapping) + where + f (ts', js, new_ts) t + | Just (_, j) <- find ((== t) . fst) ts' = + ( delete (t, j) ts', + js ++ [j], + new_ts + ) + | otherwise = + ( ts', + js ++ [length ts + length new_ts], + new_ts ++ [t] + ) -internaliseSumType :: M.Map Name [E.StructType] - -> InternaliseM ([I.TypeBase ExtShape Uniqueness], - M.Map Name (Int, [Int])) +internaliseSumType :: + M.Map Name [E.StructType] -> + InternaliseM + ( [I.TypeBase ExtShape Uniqueness], + M.Map Name (Int, [Int]) + ) internaliseSumType cs = - runInternaliseTypeM $ internaliseConstructors <$> - traverse (fmap concat . mapM internaliseTypeM) cs + runInternaliseTypeM $ + internaliseConstructors + <$> traverse (fmap concat . mapM internaliseTypeM) cs -- | How many core language values are needed to represent one source -- language value of the given type? diff --git a/src/Futhark/MonadFreshNames.hs b/src/Futhark/MonadFreshNames.hs --- a/src/Futhark/MonadFreshNames.hs +++ b/src/Futhark/MonadFreshNames.hs @@ -1,4 +1,6 @@ -{-# LANGUAGE FlexibleInstances, UndecidableInstances #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE UndecidableInstances #-} + -- | This module provides a monadic facility similar (and built on top -- of) "Futhark.FreshNames". The removes the need for a (small) amount of -- boilerplate, at the cost of using some GHC extensions. The idea is @@ -6,30 +8,30 @@ -- 'MonadFreshNames', you can automatically use the name generation -- functions exported by this module. module Futhark.MonadFreshNames - ( MonadFreshNames (..) - , modifyNameSource - , newName - , newNameFromString - , newVName - , newIdent - , newIdent' - , newParam - , module Futhark.FreshNames - ) where + ( MonadFreshNames (..), + modifyNameSource, + newName, + newNameFromString, + newVName, + newIdent, + newIdent', + newParam, + module Futhark.FreshNames, + ) +where import Control.Monad.Except +import qualified Control.Monad.RWS.Lazy +import qualified Control.Monad.RWS.Strict +import Control.Monad.Reader import qualified Control.Monad.State.Lazy import qualified Control.Monad.State.Strict +import qualified Control.Monad.Trans.Maybe import qualified Control.Monad.Writer.Lazy import qualified Control.Monad.Writer.Strict -import qualified Control.Monad.RWS.Lazy -import qualified Control.Monad.RWS.Strict -import qualified Control.Monad.Trans.Maybe -import Control.Monad.Reader - -import Futhark.IR.Syntax -import qualified Futhark.FreshNames as FreshNames import Futhark.FreshNames hiding (newName) +import qualified Futhark.FreshNames as FreshNames +import Futhark.IR.Syntax -- | A monad that stores a name source. The following is a good -- instance for a monad in which the only state is a @NameSource vn@: @@ -51,13 +53,17 @@ getNameSource = Control.Monad.State.Strict.get putNameSource = Control.Monad.State.Strict.put -instance (Applicative im, Monad im, Monoid w) => - MonadFreshNames (Control.Monad.RWS.Lazy.RWST r w VNameSource im) where +instance + (Applicative im, Monad im, Monoid w) => + MonadFreshNames (Control.Monad.RWS.Lazy.RWST r w VNameSource im) + where getNameSource = Control.Monad.RWS.Lazy.get putNameSource = Control.Monad.RWS.Lazy.put -instance (Applicative im, Monad im, Monoid w) => - MonadFreshNames (Control.Monad.RWS.Strict.RWST r w VNameSource im) where +instance + (Applicative im, Monad im, Monoid w) => + MonadFreshNames (Control.Monad.RWS.Strict.RWST r w VNameSource im) + where getNameSource = Control.Monad.RWS.Strict.get putNameSource = Control.Monad.RWS.Strict.put @@ -65,10 +71,11 @@ -- one, using 'getNameSource' and 'putNameSource' before and after the -- computation. modifyNameSource :: MonadFreshNames m => (VNameSource -> (a, VNameSource)) -> m a -modifyNameSource m = do src <- getNameSource - let (x,src') = m src - putNameSource src' - return x +modifyNameSource m = do + src <- getNameSource + let (x, src') = m src + putNameSource src' + return x -- | Produce a fresh name, using the given name as a template. newName :: MonadFreshNames m => VName -> m VName @@ -87,24 +94,33 @@ newVName = newID . nameFromString -- | Produce a fresh 'Ident', using the given name as a template. -newIdent :: MonadFreshNames m => - String -> Type -> m Ident +newIdent :: + MonadFreshNames m => + String -> + Type -> + m Ident newIdent s t = do s' <- newID $ nameFromString s return $ Ident s' t -- | Produce a fresh 'Ident', using the given 'Ident' as a template, -- but possibly modifying the name. -newIdent' :: MonadFreshNames m => - (String -> String) - -> Ident -> m Ident +newIdent' :: + MonadFreshNames m => + (String -> String) -> + Ident -> + m Ident newIdent' f ident = - newIdent (f $ nameToString $ baseName $ identName ident) - (identType ident) + newIdent + (f $ nameToString $ baseName $ identName ident) + (identType ident) -- | Produce a fresh 'Param', using the given name as a template. -newParam :: MonadFreshNames m => - String -> dec -> m (Param dec) +newParam :: + MonadFreshNames m => + String -> + dec -> + m (Param dec) newParam s t = do s' <- newID $ nameFromString s return $ Param s' t @@ -116,22 +132,30 @@ getNameSource = lift getNameSource putNameSource = lift . putNameSource -instance (MonadFreshNames m, Monoid s) => - MonadFreshNames (Control.Monad.Writer.Lazy.WriterT s m) where +instance + (MonadFreshNames m, Monoid s) => + MonadFreshNames (Control.Monad.Writer.Lazy.WriterT s m) + where getNameSource = lift getNameSource putNameSource = lift . putNameSource -instance (MonadFreshNames m, Monoid s) => - MonadFreshNames (Control.Monad.Writer.Strict.WriterT s m) where +instance + (MonadFreshNames m, Monoid s) => + MonadFreshNames (Control.Monad.Writer.Strict.WriterT s m) + where getNameSource = lift getNameSource putNameSource = lift . putNameSource -instance MonadFreshNames m => - MonadFreshNames (Control.Monad.Trans.Maybe.MaybeT m) where +instance + MonadFreshNames m => + MonadFreshNames (Control.Monad.Trans.Maybe.MaybeT m) + where getNameSource = lift getNameSource putNameSource = lift . putNameSource -instance MonadFreshNames m => - MonadFreshNames (ExceptT e m) where +instance + MonadFreshNames m => + MonadFreshNames (ExceptT e m) + where getNameSource = lift getNameSource putNameSource = lift . putNameSource 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 @@ -1,7 +1,8 @@ -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE UndecidableInstances #-} + -- | This module implements common-subexpression elimination. This -- module does not actually remove the duplicate, but only replaces -- one with a diference to the other. E.g: @@ -27,27 +28,30 @@ -- equal to any other, since the variable names will be distinct. -- This affects SOACs in particular. module Futhark.Optimise.CSE - ( performCSE - , performCSEOnFunDef - , performCSEOnStms - , CSEInOp - ) - where + ( performCSE, + performCSEOnFunDef, + performCSEOnStms, + CSEInOp, + ) +where import Control.Monad.Reader import qualified Data.Map.Strict as M - import Futhark.Analysis.Alias import Futhark.IR -import Futhark.IR.Prop.Aliases import Futhark.IR.Aliases - (removeProgAliases, removeFunDefAliases, removeStmAliases, - Aliases, consumedInStms) + ( Aliases, + consumedInStms, + removeFunDefAliases, + removeProgAliases, + removeStmAliases, + ) import qualified Futhark.IR.Kernels.Kernel as Kernel -import qualified Futhark.IR.SOACS.SOAC as SOAC import qualified Futhark.IR.Mem as Memory -import Futhark.Transform.Substitute +import Futhark.IR.Prop.Aliases +import qualified Futhark.IR.SOACS.SOAC as SOAC import Futhark.Pass +import Futhark.Transform.Substitute -- | Perform CSE on every function in a program. -- @@ -55,19 +59,26 @@ -- expressions producing arrays. This should be disabled when the lore has -- memory information, since at that point arrays have identity beyond their -- value. -performCSE :: (ASTLore lore, CanBeAliased (Op lore), - CSEInOp (OpWithAliases (Op lore))) => - Bool -> Pass lore lore +performCSE :: + ( ASTLore lore, + CanBeAliased (Op lore), + CSEInOp (OpWithAliases (Op lore)) + ) => + Bool -> + Pass lore lore performCSE cse_arrays = Pass "CSE" "Combine common subexpressions." $ - fmap removeProgAliases . - intraproceduralTransformationWithConsts onConsts onFun . - aliasAnalysis - where onConsts stms = - pure $ fst $ - runReader (cseInStms (consumedInStms stms) (stmsToList stms) (return ())) - (newCSEState cse_arrays) - onFun _ = pure . cseInFunDef cse_arrays + fmap removeProgAliases + . intraproceduralTransformationWithConsts onConsts onFun + . aliasAnalysis + where + onConsts stms = + pure $ + fst $ + runReader + (cseInStms (consumedInStms stms) (stmsToList stms) (return ())) + (newCSEState cse_arrays) + onFun _ = pure . cseInFunDef cse_arrays -- | Perform CSE on a single function. -- @@ -75,9 +86,14 @@ -- expressions producing arrays. This should be disabled when the lore has -- memory information, since at that point arrays have identity beyond their -- value. -performCSEOnFunDef :: (ASTLore lore, CanBeAliased (Op lore), - CSEInOp (OpWithAliases (Op lore))) => - Bool -> FunDef lore -> FunDef lore +performCSEOnFunDef :: + ( ASTLore lore, + CanBeAliased (Op lore), + CSEInOp (OpWithAliases (Op lore)) + ) => + Bool -> + FunDef lore -> + FunDef lore performCSEOnFunDef cse_arrays = removeFunDefAliases . cseInFunDef cse_arrays . analyseFun @@ -87,106 +103,139 @@ -- expressions producing arrays. This should be disabled when the lore has -- memory information, since at that point arrays have identity beyond their -- value. -performCSEOnStms :: (ASTLore lore, CanBeAliased (Op lore), - CSEInOp (OpWithAliases (Op lore))) => - Bool -> Stms lore -> Stms lore +performCSEOnStms :: + ( ASTLore lore, + CanBeAliased (Op lore), + CSEInOp (OpWithAliases (Op lore)) + ) => + Bool -> + Stms lore -> + Stms lore performCSEOnStms cse_arrays = fmap removeStmAliases . f . fst . analyseStms mempty - where f stms = - fst $ runReader (cseInStms (consumedInStms stms) - (stmsToList stms) (return ())) + where + f stms = + fst $ + runReader + ( cseInStms + (consumedInStms stms) + (stmsToList stms) + (return ()) + ) (newCSEState cse_arrays) -cseInFunDef :: (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => - Bool -> FunDef lore -> FunDef lore +cseInFunDef :: + (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => + Bool -> + FunDef lore -> + FunDef lore cseInFunDef cse_arrays fundec = - fundec { funDefBody = - runReader (cseInBody ds $ funDefBody fundec) $ newCSEState cse_arrays - } - where ds = map (diet . declExtTypeOf) $ funDefRetType fundec + fundec + { funDefBody = + runReader (cseInBody ds $ funDefBody fundec) $ newCSEState cse_arrays + } + where + ds = map (diet . declExtTypeOf) $ funDefRetType fundec type CSEM lore = Reader (CSEState lore) -cseInBody :: (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => - [Diet] -> Body lore -> CSEM lore (Body lore) +cseInBody :: + (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => + [Diet] -> + Body lore -> + CSEM lore (Body lore) cseInBody ds (Body bodydec bnds res) = do (bnds', res') <- cseInStms (res_cons <> consumedInStms bnds) (stmsToList bnds) $ do - CSEState (_, nsubsts) _ <- ask - return $ substituteNames nsubsts res + CSEState (_, nsubsts) _ <- ask + return $ substituteNames nsubsts res return $ Body bodydec bnds' res' - where res_cons = mconcat $ zipWith consumeResult ds res - consumeResult Consume se = freeIn se - consumeResult _ _ = mempty + where + res_cons = mconcat $ zipWith consumeResult ds res + consumeResult Consume se = freeIn se + consumeResult _ _ = mempty -cseInLambda :: (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => - Lambda lore -> CSEM lore (Lambda lore) +cseInLambda :: + (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => + Lambda lore -> + CSEM lore (Lambda lore) cseInLambda lam = do body' <- cseInBody (map (const Observe) $ lambdaReturnType lam) $ lambdaBody lam - return lam { lambdaBody = body' } + return lam {lambdaBody = body'} -cseInStms :: (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => - Names -> [Stm lore] - -> CSEM lore a - -> CSEM lore (Stms lore, a) -cseInStms _ [] m = do a <- m - return (mempty, a) -cseInStms consumed (bnd:bnds) m = +cseInStms :: + (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => + Names -> + [Stm lore] -> + CSEM lore a -> + CSEM lore (Stms lore, a) +cseInStms _ [] m = do + a <- m + return (mempty, a) +cseInStms consumed (bnd : bnds) m = cseInStm consumed bnd $ \bnd' -> do (bnds', a) <- cseInStms consumed bnds m bnd'' <- mapM nestedCSE bnd' - return (stmsFromList bnd''<>bnds', a) - where nestedCSE bnd' = do - let ds = map patElemDiet $ patternValueElements $ stmPattern bnd' - e <- mapExpM (cse ds) $ stmExp bnd' - return bnd' { stmExp = e } + return (stmsFromList bnd'' <> bnds', a) + where + nestedCSE bnd' = do + let ds = map patElemDiet $ patternValueElements $ stmPattern bnd' + e <- mapExpM (cse ds) $ stmExp bnd' + return bnd' {stmExp = e} - cse ds = identityMapper { mapOnBody = const $ cseInBody ds - , mapOnOp = cseInOp - } + cse ds = + identityMapper + { mapOnBody = const $ cseInBody ds, + mapOnOp = cseInOp + } - patElemDiet pe | patElemName pe `nameIn` consumed = Consume - | otherwise = Observe + patElemDiet pe + | patElemName pe `nameIn` consumed = Consume + | otherwise = Observe -cseInStm :: ASTLore lore => - Names -> Stm lore - -> ([Stm lore] -> CSEM lore a) - -> CSEM lore a +cseInStm :: + ASTLore lore => + Names -> + Stm lore -> + ([Stm lore] -> CSEM lore a) -> + CSEM lore a cseInStm consumed (Let pat (StmAux cs attrs edec) e) m = do CSEState (esubsts, nsubsts) cse_arrays <- ask let e' = substituteNames nsubsts e pat' = substituteNames nsubsts pat - if any (bad cse_arrays) $ patternValueElements pat then - m [Let pat' (StmAux cs attrs edec) e'] - else - case M.lookup (edec, e') esubsts of + if any (bad cse_arrays) $ patternValueElements pat + then m [Let pat' (StmAux cs attrs edec) e'] + else case M.lookup (edec, e') esubsts of Just subpat -> local (addNameSubst pat' subpat) $ do let lets = [ Let (Pattern [] [patElem']) (StmAux cs attrs edec) $ BasicOp $ SubExp $ Var $ patElemName patElem - | (name,patElem) <- zip (patternNames pat') $ patternElements subpat , - let patElem' = patElem { patElemName = name } + | (name, patElem) <- zip (patternNames pat') $ patternElements subpat, + let patElem' = patElem {patElemName = name} ] m lets - _ -> local (addExpSubst pat' edec e') $ - m [Let pat' (StmAux cs attrs edec) e'] + _ -> + local (addExpSubst pat' edec e') $ + m [Let pat' (StmAux cs attrs edec) e'] + where + bad cse_arrays pe + | Mem {} <- patElemType pe = True + | Array {} <- patElemType pe, not cse_arrays = True + | patElemName pe `nameIn` consumed = True + | otherwise = False - where bad cse_arrays pe - | Mem{} <- patElemType pe = True - | Array{} <- patElemType pe, not cse_arrays = True - | patElemName pe `nameIn` consumed = True - | otherwise = False +type ExpressionSubstitutions lore = + M.Map + (ExpDec lore, Exp lore) + (Pattern lore) -type ExpressionSubstitutions lore = M.Map - (ExpDec lore, Exp lore) - (Pattern lore) type NameSubstitutions = M.Map VName VName data CSEState lore = CSEState - { _cseSubstitutions :: (ExpressionSubstitutions lore, NameSubstitutions) - , _cseArrays :: Bool - } + { _cseSubstitutions :: (ExpressionSubstitutions lore, NameSubstitutions), + _cseArrays :: Bool + } newCSEState :: Bool -> CSEState lore newCSEState = CSEState (M.empty, M.empty) @@ -198,12 +247,15 @@ addNameSubst pat subpat (CSEState (esubsts, nsubsts) cse_arrays) = CSEState (esubsts, mkSubsts pat subpat `M.union` nsubsts) cse_arrays -addExpSubst :: ASTLore lore => - Pattern lore -> ExpDec lore -> Exp lore - -> CSEState lore - -> CSEState lore +addExpSubst :: + ASTLore lore => + Pattern lore -> + ExpDec lore -> + Exp lore -> + CSEState lore -> + CSEState lore addExpSubst pat edec e (CSEState (esubsts, nsubsts) cse_arrays) = - CSEState (M.insert (edec,e) pat esubsts, nsubsts) cse_arrays + CSEState (M.insert (edec, e) pat esubsts, nsubsts) cse_arrays -- | The operations that permit CSE. class CSEInOp op where @@ -218,30 +270,44 @@ CSEState _ cse_arrays <- ask return $ runReader m $ newCSEState cse_arrays -instance (ASTLore lore, Aliased lore, - CSEInOp (Op lore), CSEInOp op) => CSEInOp (Kernel.HostOp lore op) where +instance + ( ASTLore lore, + Aliased lore, + CSEInOp (Op lore), + CSEInOp op + ) => + CSEInOp (Kernel.HostOp lore op) + where cseInOp (Kernel.SegOp op) = Kernel.SegOp <$> cseInOp op cseInOp (Kernel.OtherOp op) = Kernel.OtherOp <$> cseInOp op cseInOp x = return x -instance (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => - CSEInOp (Kernel.SegOp lvl lore) where - cseInOp = subCSE . - Kernel.mapSegOpM - (Kernel.SegOpMapper return cseInLambda cseInKernelBody return return) +instance + (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => + CSEInOp (Kernel.SegOp lvl lore) + where + cseInOp = + subCSE + . Kernel.mapSegOpM + (Kernel.SegOpMapper return cseInLambda cseInKernelBody return return) -cseInKernelBody :: (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => - Kernel.KernelBody lore -> CSEM lore (Kernel.KernelBody lore) +cseInKernelBody :: + (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => + Kernel.KernelBody lore -> + CSEM lore (Kernel.KernelBody lore) cseInKernelBody (Kernel.KernelBody bodydec bnds res) = do Body _ bnds' _ <- cseInBody (map (const Observe) res) $ Body bodydec bnds [] return $ Kernel.KernelBody bodydec bnds' res instance CSEInOp op => CSEInOp (Memory.MemOp op) where - cseInOp o@Memory.Alloc{} = return o + cseInOp o@Memory.Alloc {} = return o cseInOp (Memory.Inner k) = Memory.Inner <$> subCSE (cseInOp k) -instance (ASTLore lore, - CanBeAliased (Op lore), - CSEInOp (OpWithAliases (Op lore))) => - CSEInOp (SOAC.SOAC (Aliases lore)) where +instance + ( ASTLore lore, + CanBeAliased (Op lore), + CSEInOp (OpWithAliases (Op lore)) + ) => + CSEInOp (SOAC.SOAC (Aliases lore)) + where cseInOp = subCSE . SOAC.mapSOACM (SOAC.SOACMapper return cseInLambda return) diff --git a/src/Futhark/Optimise/DoubleBuffer.hs b/src/Futhark/Optimise/DoubleBuffer.hs --- a/src/Futhark/Optimise/DoubleBuffer.hs +++ b/src/Futhark/Optimise/DoubleBuffer.hs @@ -1,10 +1,11 @@ {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + -- | The simplification engine is only willing to hoist allocations -- out of loops if the memory block resulting from the allocation is -- dead at the end of the loop. If it is not, we may cause data @@ -21,117 +22,132 @@ -- value. This has the effect of making the memory block returned by -- the array non-existential, which is important for later memory -- expansion to work. -module Futhark.Optimise.DoubleBuffer - ( doubleBuffer ) - where +module Futhark.Optimise.DoubleBuffer (doubleBuffer) where -import Control.Monad.State -import Control.Monad.Writer -import Control.Monad.Reader +import Control.Monad.Reader +import Control.Monad.State +import Control.Monad.Writer +import Data.List (find) import qualified Data.Map.Strict as M -import Data.Maybe -import Data.List (find) - -import Futhark.Construct -import Futhark.IR -import Futhark.Pass.ExplicitAllocations (arraySizeInBytesExp) -import Futhark.Pass.ExplicitAllocations.Kernels () +import Data.Maybe +import Futhark.Construct +import Futhark.IR +import Futhark.IR.KernelsMem import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.IR.KernelsMem -import Futhark.Pass -import Futhark.Util (maybeHead) +import Futhark.Pass +import Futhark.Pass.ExplicitAllocations (arraySizeInBytesExp) +import Futhark.Pass.ExplicitAllocations.Kernels () +import Futhark.Util (maybeHead) -- | The double buffering pass definition. doubleBuffer :: Pass KernelsMem KernelsMem doubleBuffer = - Pass { passName = "Double buffer" - , passDescription = "Perform double buffering for merge parameters of sequential loops." - , passFunction = intraproceduralTransformation optimise - } - where optimise scope stms = modifyNameSource $ \src -> - let m = runDoubleBufferM $ localScope scope $ - fmap stmsFromList $ optimiseStms $ stmsToList stms - in runState (runReaderT m env) src + Pass + { passName = "Double buffer", + passDescription = "Perform double buffering for merge parameters of sequential loops.", + passFunction = intraproceduralTransformation optimise + } + where + optimise scope stms = modifyNameSource $ \src -> + let m = + runDoubleBufferM $ + localScope scope $ + fmap stmsFromList $ optimiseStms $ stmsToList stms + in runState (runReaderT m env) src - env = Env mempty doNotTouchLoop - doNotTouchLoop ctx val body = return (mempty, ctx, val, body) + env = Env mempty doNotTouchLoop + doNotTouchLoop ctx val body = return (mempty, ctx, val, body) -data Env = Env { envScope :: Scope KernelsMem - , envOptimiseLoop :: OptimiseLoop - } +data Env = Env + { envScope :: Scope KernelsMem, + envOptimiseLoop :: OptimiseLoop + } -newtype DoubleBufferM a = - DoubleBufferM { runDoubleBufferM :: ReaderT Env (State VNameSource) a } +newtype DoubleBufferM a = DoubleBufferM {runDoubleBufferM :: ReaderT Env (State VNameSource) a} deriving (Functor, Applicative, Monad, MonadReader Env, MonadFreshNames) instance HasScope KernelsMem DoubleBufferM where askScope = asks envScope instance LocalScope KernelsMem DoubleBufferM where - localScope scope = local $ \env -> env { envScope = envScope env <> scope } + localScope scope = local $ \env -> env {envScope = envScope env <> scope} optimiseBody :: Body KernelsMem -> DoubleBufferM (Body KernelsMem) optimiseBody body = do bnds' <- optimiseStms $ stmsToList $ bodyStms body - return $ body { bodyStms = stmsFromList bnds' } + return $ body {bodyStms = stmsFromList bnds'} optimiseStms :: [Stm KernelsMem] -> DoubleBufferM [Stm KernelsMem] optimiseStms [] = return [] -optimiseStms (e:es) = do +optimiseStms (e : es) = do e_es <- optimiseStm e es' <- localScope (castScope $ scopeOf e_es) $ optimiseStms es return $ e_es ++ es' optimiseStm :: Stm KernelsMem -> DoubleBufferM [Stm KernelsMem] optimiseStm (Let pat aux (DoLoop ctx val form body)) = do - body' <- localScope (scopeOf form <> scopeOfFParams (map fst $ ctx++val)) $ - optimiseBody body + body' <- + localScope (scopeOf form <> scopeOfFParams (map fst $ ctx ++ val)) $ + optimiseBody body opt_loop <- asks envOptimiseLoop (bnds, ctx', val', body'') <- opt_loop ctx val body' return $ bnds ++ [Let pat aux $ DoLoop ctx' val' form body''] optimiseStm (Let pat aux e) = pure . Let pat aux <$> mapExpM optimise e - where optimise = identityMapper { mapOnBody = \_ x -> - optimiseBody x :: DoubleBufferM (Body KernelsMem) - , mapOnOp = optimiseOp - } + where + optimise = + identityMapper + { mapOnBody = \_ x -> + optimiseBody x :: DoubleBufferM (Body KernelsMem), + mapOnOp = optimiseOp + } -optimiseOp :: Op KernelsMem - -> DoubleBufferM (Op KernelsMem) +optimiseOp :: + Op KernelsMem -> + DoubleBufferM (Op KernelsMem) optimiseOp (Inner (SegOp op)) = local inSegOp $ Inner . SegOp <$> mapSegOpM mapper op - where mapper = identitySegOpMapper - { mapOnSegOpLambda = optimiseLambda - , mapOnSegOpBody = optimiseKernelBody - } - inSegOp env = env { envOptimiseLoop = optimiseLoop } + where + mapper = + identitySegOpMapper + { mapOnSegOpLambda = optimiseLambda, + mapOnSegOpBody = optimiseKernelBody + } + inSegOp env = env {envOptimiseLoop = optimiseLoop} optimiseOp op = return op -optimiseKernelBody :: KernelBody KernelsMem - -> DoubleBufferM (KernelBody KernelsMem) +optimiseKernelBody :: + KernelBody KernelsMem -> + DoubleBufferM (KernelBody KernelsMem) optimiseKernelBody kbody = do stms' <- optimiseStms $ stmsToList $ kernelBodyStms kbody - return $ kbody { kernelBodyStms = stmsFromList stms' } + return $ kbody {kernelBodyStms = stmsFromList stms'} optimiseLambda :: Lambda KernelsMem -> DoubleBufferM (Lambda KernelsMem) optimiseLambda lam = do body <- localScope (castScope $ scopeOf lam) $ optimiseBody $ lambdaBody lam - return lam { lambdaBody = body } + return lam {lambdaBody = body} type OptimiseLoop = - [(FParam KernelsMem, SubExp)] -> [(FParam KernelsMem, SubExp)] -> Body KernelsMem - -> DoubleBufferM ([Stm KernelsMem], - [(FParam KernelsMem, SubExp)], - [(FParam KernelsMem, SubExp)], - Body KernelsMem) + [(FParam KernelsMem, SubExp)] -> + [(FParam KernelsMem, SubExp)] -> + Body KernelsMem -> + DoubleBufferM + ( [Stm KernelsMem], + [(FParam KernelsMem, SubExp)], + [(FParam KernelsMem, SubExp)], + Body KernelsMem + ) optimiseLoop :: OptimiseLoop optimiseLoop ctx val body = do -- We start out by figuring out which of the merge variables should -- be double-buffered. - buffered <- doubleBufferMergeParams - (zip (map fst ctx) (bodyResult body)) (map fst merge) - (boundInBody body) + buffered <- + doubleBufferMergeParams + (zip (map fst ctx) (bodyResult body)) + (map fst merge) + (boundInBody body) -- Then create the allocations of the buffers and copies of the -- initial values. (merge', allocs) <- allocStms merge buffered @@ -140,124 +156,144 @@ (ctx', val') = splitAt (length ctx) merge' -- Modify the initial merge p return (allocs, ctx', val', body') - where merge = ctx ++ val + where + merge = ctx ++ val -- | The booleans indicate whether we should also play with the -- initial merge values. -data DoubleBuffer = BufferAlloc VName (PrimExp VName) Space Bool - | BufferCopy VName IxFun VName Bool - -- ^ First name is the memory block to copy to, - -- second is the name of the array copy. - | NoBuffer - deriving (Show) +data DoubleBuffer + = BufferAlloc VName (PrimExp VName) Space Bool + | -- | First name is the memory block to copy to, + -- second is the name of the array copy. + BufferCopy VName IxFun VName Bool + | NoBuffer + deriving (Show) -doubleBufferMergeParams :: MonadFreshNames m => - [(FParam KernelsMem, SubExp)] - -> [FParam KernelsMem] -> Names - -> m [DoubleBuffer] +doubleBufferMergeParams :: + MonadFreshNames m => + [(FParam KernelsMem, SubExp)] -> + [FParam KernelsMem] -> + Names -> + m [DoubleBuffer] doubleBufferMergeParams ctx_and_res val_params bound_in_loop = evalStateT (mapM buffer val_params) M.empty - where loopVariant v = v `nameIn` bound_in_loop || - v `elem` map (paramName . fst) ctx_and_res + where + loopVariant v = + v `nameIn` bound_in_loop + || v `elem` map (paramName . fst) ctx_and_res - loopInvariantSize (Constant v) = - Just (Constant v, True) - loopInvariantSize (Var v) = - case find ((==v) . paramName . fst) ctx_and_res of - Just (_, Constant val) -> - Just (Constant val, False) - Just (_, Var v') | not $ loopVariant v' -> - Just (Var v', False) - Just _ -> - Nothing - Nothing -> - Just (Var v, True) + loopInvariantSize (Constant v) = + Just (Constant v, True) + loopInvariantSize (Var v) = + case find ((== v) . paramName . fst) ctx_and_res of + Just (_, Constant val) -> + Just (Constant val, False) + Just (_, Var v') + | not $ loopVariant v' -> + Just (Var v', False) + Just _ -> + Nothing + Nothing -> + Just (Var v, True) - sizeForMem mem = maybeHead $ mapMaybe (arrayInMem . paramDec) val_params - where arrayInMem (MemArray pt shape _ (ArrayIn arraymem ixfun)) - | IxFun.isDirect ixfun, - Just (dims, b) <- - mapAndUnzipM loopInvariantSize $ shapeDims shape, - mem == arraymem = - Just (arraySizeInBytesExp $ - Array pt (Shape dims) NoUniqueness, - or b) - arrayInMem _ = Nothing + sizeForMem mem = maybeHead $ mapMaybe (arrayInMem . paramDec) val_params + where + arrayInMem (MemArray pt shape _ (ArrayIn arraymem ixfun)) + | IxFun.isDirect ixfun, + Just (dims, b) <- + mapAndUnzipM loopInvariantSize $ shapeDims shape, + mem == arraymem = + Just + ( arraySizeInBytesExp $ + Array pt (Shape dims) NoUniqueness, + or b + ) + arrayInMem _ = Nothing - buffer fparam = case paramType fparam of - Mem space - | Just (size, b) <- sizeForMem $ paramName fparam -> do - -- Let us double buffer this! - bufname <- lift $ newVName "double_buffer_mem" - modify $ M.insert (paramName fparam) (bufname, b) - return $ BufferAlloc bufname size space b - Array {} - | MemArray _ _ _ (ArrayIn mem ixfun) <- paramDec fparam -> do - buffered <- gets $ M.lookup mem - case buffered of - Just (bufname, b) -> do - copyname <- lift $ newVName "double_buffer_array" - return $ BufferCopy bufname ixfun copyname b - Nothing -> - return NoBuffer - _ -> return NoBuffer + buffer fparam = case paramType fparam of + Mem space + | Just (size, b) <- sizeForMem $ paramName fparam -> do + -- Let us double buffer this! + bufname <- lift $ newVName "double_buffer_mem" + modify $ M.insert (paramName fparam) (bufname, b) + return $ BufferAlloc bufname size space b + Array {} + | MemArray _ _ _ (ArrayIn mem ixfun) <- paramDec fparam -> do + buffered <- gets $ M.lookup mem + case buffered of + Just (bufname, b) -> do + copyname <- lift $ newVName "double_buffer_array" + return $ BufferCopy bufname ixfun copyname b + Nothing -> + return NoBuffer + _ -> return NoBuffer -allocStms :: [(FParam KernelsMem, SubExp)] -> [DoubleBuffer] - -> DoubleBufferM ([(FParam KernelsMem, SubExp)], [Stm KernelsMem]) +allocStms :: + [(FParam KernelsMem, SubExp)] -> + [DoubleBuffer] -> + DoubleBufferM ([(FParam KernelsMem, SubExp)], [Stm KernelsMem]) allocStms merge = runWriterT . zipWithM allocation merge - where allocation m@(Param pname _, _) (BufferAlloc name size space b) = do - stms <- lift $ runBinder_ $ do - size' <- toSubExp "double_buffer_size" size - letBindNames [name] $ Op $ Alloc size' space - tell $ stmsToList stms - if b then return (Param pname $ MemMem space, Var name) - else return m - allocation (f, Var v) (BufferCopy mem _ _ b) | b = do - v_copy <- lift $ newVName $ baseString v ++ "_double_buffer_copy" - (_v_mem, v_ixfun) <- lift $ lookupArraySummary v - let bt = elemType $ paramType f - shape = arrayShape $ paramType f - bound = MemArray bt shape NoUniqueness $ ArrayIn mem v_ixfun - tell [Let (Pattern [] [PatElem v_copy bound]) (defAux ()) $ - BasicOp $ Copy v] - -- It is important that we treat this as a consumption, to - -- avoid the Copy from being hoisted out of any enclosing - -- loops. Since we re-use (=overwrite) memory in the loop, - -- the copy is critical for initialisation. See issue #816. - let uniqueMemInfo (MemArray pt pshape _ ret) = - MemArray pt pshape Unique ret - uniqueMemInfo info = info - return (uniqueMemInfo <$> f, Var v_copy) - allocation (f, se) _ = - return (f, se) + where + allocation m@(Param pname _, _) (BufferAlloc name size space b) = do + stms <- lift $ + runBinder_ $ do + size' <- toSubExp "double_buffer_size" size + letBindNames [name] $ Op $ Alloc size' space + tell $ stmsToList stms + if b + then return (Param pname $ MemMem space, Var name) + else return m + allocation (f, Var v) (BufferCopy mem _ _ b) | b = do + v_copy <- lift $ newVName $ baseString v ++ "_double_buffer_copy" + (_v_mem, v_ixfun) <- lift $ lookupArraySummary v + let bt = elemType $ paramType f + shape = arrayShape $ paramType f + bound = MemArray bt shape NoUniqueness $ ArrayIn mem v_ixfun + tell + [ Let (Pattern [] [PatElem v_copy bound]) (defAux ()) $ + BasicOp $ Copy v + ] + -- It is important that we treat this as a consumption, to + -- avoid the Copy from being hoisted out of any enclosing + -- loops. Since we re-use (=overwrite) memory in the loop, + -- the copy is critical for initialisation. See issue #816. + let uniqueMemInfo (MemArray pt pshape _ ret) = + MemArray pt pshape Unique ret + uniqueMemInfo info = info + return (uniqueMemInfo <$> f, Var v_copy) + allocation (f, se) _ = + return (f, se) -doubleBufferResult :: [FParam KernelsMem] -> [DoubleBuffer] - -> Body KernelsMem -> Body KernelsMem +doubleBufferResult :: + [FParam KernelsMem] -> + [DoubleBuffer] -> + Body KernelsMem -> + Body KernelsMem doubleBufferResult valparams buffered (Body () bnds res) = let (ctx_res, val_res) = splitAt (length res - length valparams) res - (copybnds,val_res') = + (copybnds, val_res') = unzip $ zipWith3 buffer valparams buffered val_res - in Body () (bnds<>stmsFromList (catMaybes copybnds)) $ ctx_res ++ val_res' - where buffer _ (BufferAlloc bufname _ _ _) _ = - (Nothing, Var bufname) - - buffer fparam (BufferCopy bufname ixfun copyname _) (Var v) = - -- To construct the copy we will need to figure out its type - -- based on the type of the function parameter. - let t = resultType $ paramType fparam - summary = MemArray (elemType t) (arrayShape t) NoUniqueness $ ArrayIn bufname ixfun - copybnd = Let (Pattern [] [PatElem copyname summary]) (defAux ()) $ - BasicOp $ Copy v - in (Just copybnd, Var copyname) - - buffer _ _ se = - (Nothing, se) + in Body () (bnds <> stmsFromList (catMaybes copybnds)) $ ctx_res ++ val_res' + where + buffer _ (BufferAlloc bufname _ _ _) _ = + (Nothing, Var bufname) + buffer fparam (BufferCopy bufname ixfun copyname _) (Var v) = + -- To construct the copy we will need to figure out its type + -- based on the type of the function parameter. + let t = resultType $ paramType fparam + summary = MemArray (elemType t) (arrayShape t) NoUniqueness $ ArrayIn bufname ixfun + copybnd = + Let (Pattern [] [PatElem copyname summary]) (defAux ()) $ + BasicOp $ Copy v + in (Just copybnd, Var copyname) + buffer _ _ se = + (Nothing, se) - parammap = M.fromList $ zip (map paramName valparams) res + parammap = M.fromList $ zip (map paramName valparams) res - resultType t = t `setArrayDims` map substitute (arrayDims t) + resultType t = t `setArrayDims` map substitute (arrayDims t) - substitute (Var v) - | Just replacement <- M.lookup v parammap = replacement - substitute se = - se + substitute (Var v) + | Just replacement <- M.lookup v parammap = replacement + substitute se = + se 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 @@ -1,980 +1,1125 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} --- | Perform horizontal and vertical fusion of SOACs. See the paper --- /A T2 Graph-Reduction Approach To Fusion/ for the basic idea (some --- extensions discussed in /Design and GPGPU Performance of Futhark’s --- Redomap Construct/). -module Futhark.Optimise.Fusion ( fuseSOACs ) - where - -import Control.Monad.State -import Control.Monad.Reader -import Control.Monad.Except -import Data.Maybe -import qualified Data.Map.Strict as M -import qualified Data.Set as S -import qualified Data.List as L - -import Futhark.IR.Prop.Aliases -import Futhark.IR.SOACS hiding (SOAC(..)) -import qualified Futhark.IR.Aliases as Aliases -import qualified Futhark.IR.SOACS as Futhark -import Futhark.MonadFreshNames -import Futhark.IR.SOACS.Simplify -import Futhark.Optimise.Fusion.LoopKernel -import Futhark.Construct -import qualified Futhark.Analysis.HORep.SOAC as SOAC -import qualified Futhark.Analysis.Alias as Alias -import Futhark.Transform.Rename -import Futhark.Transform.Substitute -import Futhark.Pass -import Futhark.Util (maxinum) - -data VarEntry = IsArray VName (NameInfo SOACS) Names SOAC.Input - | IsNotArray (NameInfo SOACS) - -varEntryType :: VarEntry -> NameInfo SOACS -varEntryType (IsArray _ dec _ _) = - dec -varEntryType (IsNotArray dec) = - dec - -varEntryAliases :: VarEntry -> Names -varEntryAliases (IsArray _ _ x _) = x -varEntryAliases _ = mempty - -data FusionGEnv = FusionGEnv { - soacs :: M.Map VName [VName] - -- ^ Mapping from variable name to its entire family. - , varsInScope:: M.Map VName VarEntry - , fusedRes :: FusedRes - } - -lookupArr :: VName -> FusionGEnv -> Maybe SOAC.Input -lookupArr v env = asArray =<< M.lookup v (varsInScope env) - where asArray (IsArray _ _ _ input) = Just input - asArray IsNotArray{} = Nothing - -newtype Error = Error String - -instance Show Error where - show (Error msg) = "Fusion error:\n" ++ msg - -newtype FusionGM a = FusionGM (ExceptT Error (StateT VNameSource (Reader FusionGEnv)) a) - deriving (Monad, Applicative, Functor, - MonadError Error, - MonadState VNameSource, - MonadReader FusionGEnv) - -instance MonadFreshNames FusionGM where - getNameSource = get - putNameSource = put - -instance HasScope SOACS FusionGM where - askScope = asks $ toScope . varsInScope - where toScope = M.map varEntryType - ------------------------------------------------------------------------- ---- Monadic Helpers: bind/new/runFusionGatherM, etc ------------------------------------------------------------------------- - --- | Binds an array name to the set of used-array vars -bindVar :: FusionGEnv -> (Ident, Names) -> FusionGEnv -bindVar env (Ident name t, aliases) = - env { varsInScope = M.insert name entry $ varsInScope env } - where entry = case t of - Array {} -> IsArray name (LetName t) aliases' $ SOAC.identInput $ Ident name t - _ -> IsNotArray $ LetName t - expand = maybe mempty varEntryAliases . flip M.lookup (varsInScope env) - aliases' = aliases <> mconcat (map expand $ namesToList aliases) - -bindVars :: FusionGEnv -> [(Ident, Names)] -> FusionGEnv -bindVars = foldl bindVar - -binding :: [(Ident, Names)] -> FusionGM a -> FusionGM a -binding vs = local (`bindVars` vs) - -gatherStmPattern :: Pattern -> Exp -> FusionGM FusedRes -> FusionGM FusedRes -gatherStmPattern pat e = binding $ zip idents aliases - where idents = patternIdents pat - aliases = replicate (length (patternContextNames pat)) mempty ++ - expAliases (Alias.analyseExp mempty e) - -bindingPat :: Pattern -> FusionGM a -> FusionGM a -bindingPat = binding . (`zip` repeat mempty) . patternIdents - -bindingParams :: Typed t => [Param t] -> FusionGM a -> FusionGM a -bindingParams = binding . (`zip` repeat mempty) . map paramIdent - --- | Binds an array name to the set of soac-produced vars -bindingFamilyVar :: [VName] -> FusionGEnv -> Ident -> FusionGEnv -bindingFamilyVar faml env (Ident nm t) = - env { soacs = M.insert nm faml $ soacs env - , varsInScope = M.insert nm (IsArray nm (LetName t) mempty $ - SOAC.identInput $ Ident nm t) $ - varsInScope env - } - -varAliases :: VName -> FusionGM Names -varAliases v = asks $ (oneName v<>) . maybe mempty varEntryAliases . - M.lookup v . varsInScope - -varsAliases :: Names -> FusionGM Names -varsAliases = fmap mconcat . mapM varAliases . namesToList - -checkForUpdates :: FusedRes -> Exp -> FusionGM FusedRes -checkForUpdates res (BasicOp (Update src is _)) = do - res' <- foldM addVarToInfusible res $ - src : namesToList (mconcat $ map freeIn is) - aliases <- varAliases src - let inspectKer k = k { inplace = aliases <> inplace k } - return res' { kernels = M.map inspectKer $ kernels res' } -checkForUpdates res _ = return res - --- | Updates the environment: (i) the @soacs@ (map) by binding each pattern --- element identifier to all pattern elements (identifiers) and (ii) the --- variables in scope (map) by inserting each (pattern-array) name. --- Finally, if the binding is an in-place update, then the @inplace@ field --- of each (result) kernel is updated with the new in-place updates. -bindingFamily :: Pattern -> FusionGM FusedRes -> FusionGM FusedRes -bindingFamily pat = local bind - where idents = patternIdents pat - family = patternNames pat - bind env = foldl (bindingFamilyVar family) env idents - -bindingTransform :: PatElem -> VName -> SOAC.ArrayTransform -> FusionGM a -> FusionGM a -bindingTransform pe srcname trns = local $ \env -> - case M.lookup srcname $ varsInScope env of - Just (IsArray src' _ aliases input) -> - env { varsInScope = - M.insert vname - (IsArray src' (LetName dec) (oneName srcname <> aliases) $ - trns `SOAC.addTransform` input) $ - varsInScope env - } - _ -> bindVar env (patElemIdent pe, oneName vname) - where vname = patElemName pe - dec = patElemDec pe - --- | Binds the fusion result to the environment. -bindRes :: FusedRes -> FusionGM a -> FusionGM a -bindRes rrr = local (\x -> x { fusedRes = rrr }) - --- | The fusion transformation runs in this monad. The mutable --- state refers to the fresh-names engine. --- The reader hides the vtable that associates ... to ... (fill in, please). --- The 'Either' monad is used for error handling. -runFusionGatherM :: MonadFreshNames m => - FusionGM a -> FusionGEnv -> m (Either Error a) -runFusionGatherM (FusionGM a) env = - modifyNameSource $ \src -> runReader (runStateT (runExceptT a) src) env - ------------------------------------------------------------------------- ---- Fusion Entry Points: gather the to-be-fused kernels@pgm level --- ---- and fuse them in a second pass! --- ------------------------------------------------------------------------- - --- | The pass definition. -fuseSOACs :: Pass SOACS SOACS -fuseSOACs = - Pass { passName = "Fuse SOACs" - , passDescription = "Perform higher-order optimisation, i.e., fusion." - , passFunction = \prog -> - simplifySOACS =<< renameProg =<< - intraproceduralTransformationWithConsts - (fuseConsts (freeIn (progFuns prog))) fuseFun prog - } - -fuseConsts :: Names -> Stms SOACS -> PassM (Stms SOACS) -fuseConsts used_consts consts = - fuseStms mempty consts $ map Var $ namesToList used_consts - -fuseFun :: Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS) -fuseFun consts fun = do - stms <- fuseStms (scopeOf consts <> scopeOfFParams (funDefParams fun)) - (bodyStms $ funDefBody fun) - (bodyResult $ funDefBody fun) - let body = (funDefBody fun) { bodyStms = stms } - return fun { funDefBody = body } - -fuseStms :: Scope SOACS -> Stms SOACS -> Result -> PassM (Stms SOACS) -fuseStms scope stms res = do - let env = FusionGEnv { soacs = M.empty - , varsInScope = mempty - , fusedRes = mempty - } - k <- cleanFusionResult <$> - liftEitherM (runFusionGatherM - (binding scope' $ fusionGatherStms mempty (stmsToList stms) res) - env) - if not $ rsucc k - then return stms - else liftEitherM $ runFusionGatherM (binding scope' $ bindRes k $ fuseInStms stms) env - where scope' = map toBind $ M.toList scope - toBind (k, t) = (Ident k $ typeOf t, mempty) - ---------------------------------------------------- ---------------------------------------------------- ----- RESULT's Data Structure ---------------------------------------------------- ---------------------------------------------------- - --- | A type used for (hopefully) uniquely referring a producer SOAC. --- The uniquely identifying value is the name of the first array --- returned from the SOAC. -newtype KernName = KernName { unKernName :: VName } - deriving (Eq, Ord, Show) - -data FusedRes = FusedRes { - rsucc :: Bool - -- ^ Whether we have fused something anywhere. - - , outArr :: M.Map VName KernName - -- ^ Associates an array to the name of the - -- SOAC kernel that has produced it. - - , inpArr :: M.Map VName (S.Set KernName) - -- ^ Associates an array to the names of the - -- SOAC kernels that uses it. These sets include - -- only the SOAC input arrays used as full variables, i.e., no `a[i]'. - - , infusible :: Names - -- ^ the (names of) arrays that are not fusible, i.e., - -- - -- 1. they are either used other than input to SOAC kernels, or - -- - -- 2. are used as input to at least two different kernels that - -- are not located on disjoint control-flow branches, or - -- - -- 3. are used in the lambda expression of SOACs - - , kernels :: M.Map KernName FusedKer - -- ^ The map recording the uses - } - -instance Semigroup FusedRes where - res1 <> res2 = - FusedRes (rsucc res1 || rsucc res2) - (outArr res1 `M.union` outArr res2) - (M.unionWith S.union (inpArr res1) (inpArr res2) ) - (infusible res1 <> infusible res2) - (kernels res1 `M.union` kernels res2) - -instance Monoid FusedRes where - mempty = FusedRes { rsucc = False, outArr = M.empty, inpArr = M.empty, - infusible = mempty, kernels = M.empty } - -isInpArrInResModKers :: FusedRes -> S.Set KernName -> VName -> Bool -isInpArrInResModKers ress kers nm = - case M.lookup nm (inpArr ress) of - Nothing -> False - Just s -> not $ S.null $ s `S.difference` kers - -getKersWithInpArrs :: FusedRes -> [VName] -> S.Set KernName -getKersWithInpArrs ress = - S.unions . mapMaybe (`M.lookup` inpArr ress) - --- | extend the set of names to include all the names --- produced via SOACs (by querring the vtable's soac) -expandSoacInpArr :: [VName] -> FusionGM [VName] -expandSoacInpArr = - foldM (\y nm -> do bnd <- asks $ M.lookup nm . soacs - case bnd of - Nothing -> return (y++[nm]) - Just nns -> return (y++nns ) - ) [] - ----------------------------------------------------------------------- ----------------------------------------------------------------------- - -soacInputs :: SOAC -> FusionGM ([VName], [VName]) -soacInputs soac = do - let (inp_idds, other_idds) = getIdentArr $ SOAC.inputs soac - (inp_nms0, other_nms0) = (inp_idds, other_idds) - inp_nms <- expandSoacInpArr inp_nms0 - other_nms <- expandSoacInpArr other_nms0 - return (inp_nms, other_nms) - -addNewKerWithInfusible :: FusedRes -> ([Ident], StmAux (), SOAC, Names) -> Names -> FusionGM FusedRes -addNewKerWithInfusible res (idd, aux, soac, consumed) ufs = do - nm_ker <- KernName <$> newVName "ker" - scope <- askScope - let out_nms = map identName idd - new_ker = newKernel aux soac consumed out_nms scope - comb = M.unionWith S.union - os' = M.fromList [(arr,nm_ker) | arr <- out_nms] - `M.union` outArr res - is' = M.fromList [(arr,S.singleton nm_ker) - | arr <- map SOAC.inputArray $ SOAC.inputs soac] - `comb` inpArr res - return $ FusedRes (rsucc res) os' is' ufs - (M.insert nm_ker new_ker (kernels res)) - -lookupInput :: VName -> FusionGM (Maybe SOAC.Input) -lookupInput name = asks $ lookupArr name - -inlineSOACInput :: SOAC.Input -> FusionGM SOAC.Input -inlineSOACInput (SOAC.Input ts v t) = do - maybe_inp <- lookupInput v - case maybe_inp of - Nothing -> - return $ SOAC.Input ts v t - Just (SOAC.Input ts2 v2 t2) -> - return $ SOAC.Input (ts2<>ts) v2 t2 - -inlineSOACInputs :: SOAC -> FusionGM SOAC -inlineSOACInputs soac = do - inputs' <- mapM inlineSOACInput $ SOAC.inputs soac - return $ inputs' `SOAC.setInputs` soac - - --- | Attempts to fuse between SOACs. Input: --- @rem_bnds@ are the bindings remaining in the current body after @orig_soac@. --- @lam_used_nms@ the infusible names --- @res@ the fusion result (before processing the current soac) --- @orig_soac@ and @out_idds@ the current SOAC and its binding pattern --- @consumed@ is the set of names consumed by the SOAC. --- Output: a new Fusion Result (after processing the current SOAC binding) -greedyFuse :: [Stm] -> Names -> FusedRes -> (Pattern, StmAux (), SOAC, Names) - -> FusionGM FusedRes -greedyFuse rem_bnds lam_used_nms res (out_idds, aux, orig_soac, consumed) = do - soac <- inlineSOACInputs orig_soac - (inp_nms, other_nms) <- soacInputs soac - -- Assumption: the free vars in lambda are already in @infusible res@. - let out_nms = patternNames out_idds - isInfusible = (`nameIn` infusible res) - is_screma = case orig_soac of - SOAC.Screma _ form _ -> - (isJust (isRedomapSOAC form) || isJust (isScanomapSOAC form)) && - not (isJust (isReduceSOAC form) || isJust (isScanSOAC form)) - _ -> False - -- - -- Conditions for fusion: - -- If current soac is a replicate OR (current soac a redomap/scanomap AND - -- (i) none of @out_idds@ belongs to the infusible set) - -- THEN try applying producer-consumer fusion - -- ELSE try applying horizontal fusion - -- (without duplicating computation in both cases) - - (ok_kers_compat, fused_kers, fused_nms, old_kers, oldker_nms) <- - if is_screma || any isInfusible out_nms - then horizontGreedyFuse rem_bnds res (out_idds, aux, soac, consumed) - else prodconsGreedyFuse res (out_idds, aux, soac, consumed) - -- - -- (ii) check whether fusing @soac@ will violate any in-place update - -- restriction, e.g., would move an input array past its in-place update. - let all_used_names = namesToList $ mconcat [lam_used_nms, namesFromList inp_nms, namesFromList other_nms] - has_inplace ker = any (`nameIn` inplace ker) all_used_names - ok_inplace = not $ any has_inplace old_kers - -- - -- (iii) there are some kernels that use some of `out_idds' as inputs - -- (iv) and producer-consumer or horizontal fusion succeeds with those. - let fusible_ker = not (null old_kers) && ok_inplace && ok_kers_compat - -- - -- Start constructing the fusion's result: - -- (i) inparr ids other than vars will be added to infusible list, - -- (ii) will also become part of the infusible set the inparr vars - -- that also appear as inparr of another kernel, - -- BUT which said kernel is not the one we are fusing with (now)! - let mod_kerS = if fusible_ker then S.fromList oldker_nms else mempty - let used_inps = filter (isInpArrInResModKers res mod_kerS) inp_nms - let ufs = mconcat [infusible res, namesFromList used_inps, - namesFromList other_nms `namesSubtract` - namesFromList (map SOAC.inputArray $ SOAC.inputs soac)] - let comb = M.unionWith S.union - - if not fusible_ker then - addNewKerWithInfusible res (patternIdents out_idds, aux, soac, consumed) ufs - else do - -- Need to suitably update `inpArr': - -- (i) first remove the inpArr bindings of the old kernel - let inpArr' = - foldl (\inpa (kold, knm) -> - S.foldl' - (\inpp nm -> - case M.lookup nm inpp of - Nothing -> inpp - Just s -> let new_set = S.delete knm s - in if S.null new_set - then M.delete nm inpp - else M.insert nm new_set inpp - ) - inpa $ arrInputs kold - ) - (inpArr res) (zip old_kers oldker_nms) - -- (ii) then add the inpArr bindings of the new kernel - let fused_ker_nms = zip fused_nms fused_kers - inpArr''= foldl (\inpa' (knm, knew) -> - M.fromList [ (k, S.singleton knm) - | k <- S.toList $ arrInputs knew ] - `comb` inpa' - ) - inpArr' fused_ker_nms - -- Update the kernels map (why not delete the ones that have been fused?) - let kernels' = M.fromList fused_ker_nms `M.union` kernels res - -- nothing to do for `outArr' (since we have not added a new kernel) - -- DO IMPROVEMENT: attempt to fuse the resulting kernel AGAIN until it fails, - -- but make sure NOT to add a new kernel! - return $ FusedRes True (outArr res) inpArr'' ufs kernels' - -prodconsGreedyFuse :: FusedRes -> (Pattern, StmAux (), SOAC, Names) - -> FusionGM (Bool, [FusedKer], [KernName], [FusedKer], [KernName]) -prodconsGreedyFuse res (out_idds, aux, soac, consumed) = do - let out_nms = patternNames out_idds -- Extract VNames from output patterns - to_fuse_knmSet = getKersWithInpArrs res out_nms -- Find kernels which consume outputs - to_fuse_knms = S.toList to_fuse_knmSet - lookup_kern k = case M.lookup k (kernels res) of - Nothing -> throwError $ Error - ("In Fusion.hs, greedyFuse, comp of to_fuse_kers: " - ++ "kernel name not found in kernels field!") - Just ker -> return ker - to_fuse_kers <- mapM lookup_kern to_fuse_knms -- Get all consumer kernels - -- try producer-consumer fusion - (ok_kers_compat, fused_kers) <- do - kers <- forM to_fuse_kers $ - attemptFusion mempty (patternNames out_idds) soac consumed - case sequence kers of - Nothing -> return (False, []) - Just kers' -> return (True, map certifyKer kers') - return (ok_kers_compat, fused_kers, to_fuse_knms, to_fuse_kers, to_fuse_knms) - where certifyKer k = k { kerAux = kerAux k <> aux } - -horizontGreedyFuse :: [Stm] -> FusedRes -> (Pattern, StmAux (), SOAC, Names) - -> FusionGM (Bool, [FusedKer], [KernName], [FusedKer], [KernName]) -horizontGreedyFuse rem_bnds res (out_idds, aux, soac, consumed) = do - (inp_nms, _) <- soacInputs soac - let out_nms = patternNames out_idds - infusible_nms = namesFromList $ filter (`nameIn` infusible res) out_nms - out_arr_nms = case soac of - -- 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 - _ -> out_nms - to_fuse_knms1 = S.toList $ getKersWithInpArrs res (out_arr_nms++inp_nms) - to_fuse_knms2 = getKersWithSameInpSize (SOAC.width soac) res - to_fuse_knms = S.toList $ S.fromList $ to_fuse_knms1 ++ to_fuse_knms2 - lookupKernel k = case M.lookup k (kernels res) of - Nothing -> throwError $ Error - ("In Fusion.hs, greedyFuse, comp of to_fuse_kers: " - ++ "kernel name not found in kernels field!") - Just ker -> return ker - - -- For each kernel get the index in the bindings where the kernel is - -- located and sort based on the index so that partial fusion may - -- succeed. We use the last position where one of the kernel - -- outputs occur. - let bnd_nms = map (patternNames . stmPattern) rem_bnds - kernminds <- forM to_fuse_knms $ \ker_nm -> do - ker <- lookupKernel ker_nm - case mapMaybe (\out_nm -> L.findIndex (elem out_nm) bnd_nms) (outNames ker) of - [] -> return Nothing - is -> return $ Just (ker,ker_nm,maxinum is) - - scope <- askScope - let kernminds' = L.sortBy (\(_,_,i1) (_,_,i2)->compare i1 i2) $ catMaybes kernminds - soac_kernel = newKernel aux soac consumed out_nms scope - - -- now try to fuse kernels one by one (in a fold); @ok_ind@ is the index of the - -- kernel until which fusion succeded, and @fused_ker@ is the resulting kernel. - use_scope <- (<>scopeOf rem_bnds) <$> askScope - (_,ok_ind,_,fused_ker,_) <- - foldM (\(cur_ok,n,prev_ind,cur_ker,ufus_nms) (ker, _ker_nm, bnd_ind) -> do - -- check that we still try fusion and that the intermediate - -- bindings do not use the results of cur_ker - let curker_outnms = outNames cur_ker - curker_outset = namesFromList curker_outnms - new_ufus_nms = namesFromList $ outNames ker ++ namesToList ufus_nms - -- disable horizontal fusion in the case when an output array of - -- producer SOAC is a non-trivially transformed input of the consumer - out_transf_ok = let ker_inp = SOAC.inputs $ fsoac ker - unfuse1 = namesFromList (map SOAC.inputArray ker_inp) `namesSubtract` - namesFromList (mapMaybe SOAC.isVarInput ker_inp) - unfuse2 = namesIntersection curker_outset ufus_nms - in not $ unfuse1 `namesIntersect` unfuse2 - -- Disable horizontal fusion if consumer has any - -- output transforms. - cons_no_out_transf = SOAC.nullTransforms $ outputTransform ker - - consumer_ok <- do let consumer_bnd = rem_bnds !! bnd_ind - maybesoac <- runReaderT (SOAC.fromExp $ stmExp consumer_bnd) use_scope - case maybesoac of - -- check that consumer's lambda body does not use - -- directly the produced arrays (e.g., see noFusion3.fut). - Right conssoac -> return $ not $ - curker_outset - `namesIntersect` - freeIn (lambdaBody $ SOAC.lambda conssoac) - Left _ -> return True - - let interm_bnds_ok = cur_ok && consumer_ok && out_transf_ok && cons_no_out_transf && - foldl (\ok bnd-> ok && -- hardwired to False after first fail - -- (i) check that the in-between bindings do - -- not use the result of current kernel OR - not (curker_outset `namesIntersect` freeIn (stmExp bnd)) || - --(ii) that the pattern-binding corresponds to - -- the result of the consumer kernel; in the - -- latter case it means it corresponds to a - -- kernel that has been fused in the consumer, - -- hence it should be ignored - not ( null $ curker_outnms `L.intersect` - patternNames (stmPattern bnd)) - ) True (drop (prev_ind+1) $ take bnd_ind rem_bnds) - if not interm_bnds_ok then return (False,n,bnd_ind,cur_ker,mempty) - else do new_ker <- attemptFusion ufus_nms (outNames cur_ker) - (fsoac cur_ker) (fusedConsumed cur_ker) ker - case new_ker of - Nothing -> return (False, n,bnd_ind,cur_ker,mempty) - Just krn-> - let krn' = krn { kerAux = aux <> kerAux krn } - in return (True,n+1,bnd_ind,krn',new_ufus_nms) - ) (True,0,0,soac_kernel,infusible_nms) kernminds' - - -- Find the kernels we have fused into and the name of the last such - -- kernel (if any). - let (to_fuse_kers', to_fuse_knms',_) = unzip3 $ take ok_ind kernminds' - new_kernms = drop (ok_ind-1) to_fuse_knms' - - return (ok_ind>0, [fused_ker], new_kernms, to_fuse_kers', to_fuse_knms') - - where getKersWithSameInpSize :: SubExp -> FusedRes -> [KernName] - getKersWithSameInpSize sz ress = - map fst $ filter (\ (_,ker) -> sz == SOAC.width (fsoac ker)) $ M.toList $ kernels ress - ------------------------------------------------------------------------- ------------------------------------------------------------------------- ------------------------------------------------------------------------- ---- Fusion Gather for EXPRESSIONS and BODIES, --- ---- i.e., where work is being done: --- ---- i) bottom-up AbSyn traversal (backward analysis) --- ---- ii) soacs are fused greedily iff does not duplicate computation--- ---- E.g., (y1, y2, y3) = mapT(f, x1, x2[i]) --- ---- (z1, z2) = mapT(g1, y1, y2) --- ---- (q1, q2) = mapT(g2, y3, z1, a, y3) --- ---- res = reduce(op, ne, q1, q2, z2, y1, y3) --- ---- can be fused if y1,y2,y3, z1,z2, q1,q2 are not used elsewhere: --- ---- res = redomap(op, \(x1,x2i,a)-> --- ---- let (y1,y2,y3) = f (x1, x2i) in--- ---- let (z1,z2) = g1(y1, y2) in--- ---- let (q1,q2) = g2(y3, z1, a, y3) in--- ---- (q1, q2, z2, y1, y3) --- ---- x1, x2[i], a) --- ------------------------------------------------------------------------- ------------------------------------------------------------------------- ------------------------------------------------------------------------- - -fusionGatherBody :: FusedRes -> Body -> FusionGM FusedRes -fusionGatherBody fres (Body _ stms res) = - fusionGatherStms fres (stmsToList stms) res - -fusionGatherStms :: FusedRes -> [Stm] -> Result -> FusionGM FusedRes - --- Some forms of do-loops can profitably be considered streamSeqs. We --- are careful to ensure that the generated nested loop cannot itself --- be considered a stream, to avoid infinite recursion. -fusionGatherStms fres (Let (Pattern [] pes) bndtp - (DoLoop [] merge (ForLoop i it w loop_vars) body) : bnds) res - | not $ null loop_vars = do - let (merge_params,merge_init) = unzip merge - (loop_params,loop_arrs) = unzip loop_vars - chunk_size <- newVName "chunk_size" - offset <- newVName "offset" - let chunk_param = Param chunk_size $ Prim int32 - offset_param = Param offset $ Prim $ IntType it - - acc_params <- forM merge_params $ \p -> - Param <$> newVName (baseString (paramName p) ++ "_outer") <*> - pure (paramType p) - - chunked_params <- forM loop_vars $ \(p,arr) -> - Param <$> newVName (baseString arr ++ "_chunk") <*> - pure (paramType p `arrayOfRow` Futhark.Var chunk_size) - - let lam_params = chunk_param : acc_params ++ [offset_param] ++ chunked_params - - lam_body <- runBodyBinder $ localScope (scopeOfLParams lam_params) $ do - let merge' = zip merge_params $ map (Futhark.Var . paramName) acc_params - j <- newVName "j" - loop_body <- runBodyBinder $ do - forM_ (zip loop_params chunked_params) $ \(p,a_p) -> - letBindNames [paramName p] $ BasicOp $ Index (paramName a_p) $ - fullSlice (paramType a_p) [DimFix $ Futhark.Var j] - letBindNames [i] $ BasicOp $ BinOp (Add it OverflowUndef) (Futhark.Var offset) (Futhark.Var j) - return body - eBody [pure $ - DoLoop [] merge' (ForLoop j it (Futhark.Var chunk_size) []) loop_body, - pure $ - BasicOp $ BinOp (Add Int32 OverflowUndef) (Futhark.Var offset) (Futhark.Var chunk_size)] - let lam = Lambda { lambdaParams = lam_params - , lambdaBody = lam_body - , lambdaReturnType = map paramType $ acc_params ++ [offset_param] - } - stream = Futhark.Stream w (Sequential $ merge_init ++ [intConst it 0]) lam 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 - -- identify the SOAC in the second phase of fusion. - discard <- newVName "discard" - let discard_pe = PatElem discard $ Prim int32 - - fusionGatherStms fres - (Let (Pattern [] (pes<>[discard_pe])) bndtp (Op stream) : bnds) res - -fusionGatherStms fres (bnd@(Let pat _ e):bnds) res = do - maybesoac <- SOAC.fromExp e - case maybesoac of - Right soac@(SOAC.Scatter _len lam _ivs _as) -> do - -- We put the variables produced by Scatter into the infusible - -- set to force horizontal fusion. It is not possible to - -- producer/consumer-fuse Scatter anyway. - fres' <- addNamesToInfusible fres $ namesFromList $ patternNames pat - mapLike fres' soac lam - - Right soac@(SOAC.Hist _ _ lam _) -> do - -- We put the variables produced by Hist into the infusible - -- set to force horizontal fusion. It is not possible to - -- producer/consumer-fuse Hist anyway. - fres' <- addNamesToInfusible fres $ namesFromList $ patternNames pat - mapLike fres' soac lam - - 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 - -- 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 - - _ | [pe] <- patternValueElements pat, - Just (src,trns) <- SOAC.transformFromExp (stmCerts bnd) e -> - bindingTransform pe src trns $ fusionGatherStms fres bnds res - | otherwise -> do - let pat_vars = map (BasicOp . SubExp . Var) $ patternNames pat - bres <- gatherStmPattern pat e $ fusionGatherStms fres bnds res - bres' <- checkForUpdates bres e - foldM fusionGatherExp bres' (e:pat_vars) - - where aux = stmAux bnd - rem_bnds = bnd : bnds - consumed = consumedInExp $ Alias.analyseExp mempty e - - reduceLike soac lambdas nes = do - (used_lam, lres) <- foldM fusionGatherLam (mempty, fres) lambdas - bres <- bindingFamily pat $ fusionGatherStms lres bnds res - bres' <- foldM fusionGatherSubExp bres nes - consumed' <- varsAliases consumed - greedyFuse rem_bnds used_lam bres' (pat, aux, soac, consumed') - - mapLike fres' soac lambda = do - bres <- bindingFamily pat $ fusionGatherStms fres' bnds res - (used_lam, blres) <- fusionGatherLam (mempty, bres) lambda - consumed' <- varsAliases consumed - greedyFuse rem_bnds used_lam blres (pat, aux, soac, consumed') - -fusionGatherStms fres [] res = - foldM fusionGatherExp fres $ map (BasicOp . SubExp) res - -fusionGatherExp :: FusedRes -> Exp -> FusionGM FusedRes - ------------------------------------------ ----- Index/If ---- ------------------------------------------ - -fusionGatherExp fres (DoLoop ctx val form loop_body) = do - fres' <- addNamesToInfusible fres $ freeIn form <> freeIn ctx <> freeIn val - let form_idents = - case form of - ForLoop i _ _ loopvars -> - Ident i (Prim int32) : map (paramIdent . fst) loopvars - WhileLoop{} -> [] - - new_res <- binding (zip (form_idents ++ map (paramIdent . fst) (ctx<>val)) $ - repeat mempty) $ - fusionGatherBody mempty loop_body - -- make the inpArr infusible, so that they - -- cannot be fused from outside the loop: - let (inp_arrs, _) = unzip $ M.toList $ inpArr new_res - let new_res' = new_res { infusible = infusible new_res <> mconcat (map oneName inp_arrs) } - -- merge new_res with fres' - return $ new_res' <> fres' - -fusionGatherExp fres (If cond e_then e_else _) = do - then_res <- fusionGatherBody mempty e_then - else_res <- fusionGatherBody mempty e_else - let both_res = then_res <> else_res - fres' <- fusionGatherSubExp fres cond - mergeFusionRes fres' both_res - ------------------------------------------------------------------------------------ ---- Errors: all SOACs, (because normalization ensures they appear ---- directly in let exp, i.e., let x = e) ------------------------------------------------------------------------------------ - -fusionGatherExp _ (Op Futhark.Screma{}) = errorIllegal "screma" -fusionGatherExp _ (Op Futhark.Scatter{}) = errorIllegal "write" - ------------------------------------ ----- Generic Traversal ---- ------------------------------------ - -fusionGatherExp fres e = addNamesToInfusible fres $ freeIn e - -fusionGatherSubExp :: FusedRes -> SubExp -> FusionGM FusedRes -fusionGatherSubExp fres (Var idd) = addVarToInfusible fres idd -fusionGatherSubExp fres _ = return fres - -addNamesToInfusible :: FusedRes -> Names -> FusionGM FusedRes -addNamesToInfusible fres = foldM addVarToInfusible fres . namesToList - -addVarToInfusible :: FusedRes -> VName -> FusionGM FusedRes -addVarToInfusible fres name = do - trns <- asks $ lookupArr name - let name' = case trns of - Nothing -> name - Just (SOAC.Input _ orig _) -> orig - return fres { infusible = oneName name' <> infusible fres } - --- Lambdas create a new scope. Disallow fusing from outside lambda by --- adding inp_arrs to the infusible set. -fusionGatherLam :: (Names, FusedRes) -> Lambda -> FusionGM (Names, FusedRes) -fusionGatherLam (u_set,fres) (Lambda idds body _) = do - new_res <- bindingParams idds $ fusionGatherBody mempty body - -- make the inpArr infusible, so that they - -- cannot be fused from outside the lambda: - let inp_arrs = namesFromList $ M.keys $ inpArr new_res - let unfus = infusible new_res <> inp_arrs - bnds <- asks $ M.keys . varsInScope - let unfus' = unfus `namesIntersection` namesFromList bnds - -- merge fres with new_res' - let new_res' = new_res { infusible = unfus' } - -- merge new_res with fres' - return (u_set <> unfus', new_res' <> fres) - -------------------------------------------------------------- -------------------------------------------------------------- ---- FINALLY, Substitute the kernels in function -------------------------------------------------------------- -------------------------------------------------------------- - -fuseInStms :: Stms SOACS -> FusionGM (Stms SOACS) -fuseInStms stms - | Just (Let pat aux e, stms') <- stmsHead stms = do - stms'' <- bindingPat pat $ fuseInStms stms' - soac_bnds <- replaceSOAC pat aux e - pure $ soac_bnds <> stms'' - | otherwise = - pure mempty - -fuseInBody :: Body -> FusionGM Body -fuseInBody (Body _ stms res) = - Body () <$> fuseInStms stms <*> pure res - -fuseInExp :: Exp -> FusionGM Exp - --- Handle loop specially because we need to bind the types of the --- merge variables. -fuseInExp (DoLoop ctx val form loopbody) = - binding (zip form_idents $ repeat mempty) $ - bindingParams (map fst $ ctx ++ val) $ - DoLoop ctx val form <$> fuseInBody loopbody - where form_idents = case form of - WhileLoop{} -> [] - ForLoop i it _ loopvars -> - Ident i (Prim $ IntType it) : - map (paramIdent . fst) loopvars - -fuseInExp e = mapExpM fuseIn e - -fuseIn :: Mapper SOACS SOACS FusionGM -fuseIn = identityMapper { - mapOnBody = const fuseInBody - , mapOnOp = mapSOACM identitySOACMapper { mapOnSOACLambda = fuseInLambda } - } - -fuseInLambda :: Lambda -> FusionGM Lambda -fuseInLambda (Lambda params body rtp) = do - body' <- bindingParams params $ fuseInBody body - return $ Lambda params body' rtp - -replaceSOAC :: Pattern -> StmAux () -> Exp -> FusionGM (Stms SOACS) -replaceSOAC (Pattern _ []) _ _ = return mempty -replaceSOAC pat@(Pattern _ (patElem : _)) aux e = do - fres <- asks fusedRes - let pat_nm = patElemName patElem - names = patternIdents pat - case M.lookup pat_nm (outArr fres) of - Nothing -> - oneStm . Let pat aux <$> fuseInExp e - Just knm -> - case M.lookup knm (kernels fres) of - Nothing -> throwError $ Error - ("In Fusion.hs, replaceSOAC, outArr in ker_name " - ++"which is not in Res: "++pretty (unKernName knm)) - Just ker -> do - when (null $ fusedVars ker) $ - throwError $ Error - ("In Fusion.hs, replaceSOAC, unfused kernel " - ++"still in result: "++pretty names) - insertKerSOAC aux (outNames ker) ker - -insertKerSOAC :: StmAux () -> [VName] -> FusedKer -> FusionGM (Stms SOACS) -insertKerSOAC aux names ker = do - new_soac' <- finaliseSOAC $ fsoac ker - runBinder_ $ do - 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 - 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 - --- | Perform simplification and fusion inside the lambda(s) of a SOAC. -finaliseSOAC :: SOAC.SOAC SOACS -> FusionGM (SOAC.SOAC SOACS) -finaliseSOAC new_soac = - case new_soac of - SOAC.Screma w (ScremaForm scans reds map_lam) arrs -> do - scans' <- forM scans $ \(Scan scan_lam scan_nes) -> do - scan_lam' <- simplifyAndFuseInLambda scan_lam - return $ Scan scan_lam' scan_nes - - reds' <- forM reds $ \(Reduce comm red_lam red_nes) -> do - red_lam' <- simplifyAndFuseInLambda red_lam - return $ Reduce comm red_lam' red_nes - - map_lam' <- simplifyAndFuseInLambda map_lam - - return $ SOAC.Screma w (ScremaForm scans' reds' map_lam') arrs - SOAC.Scatter w lam inps dests -> do - lam' <- simplifyAndFuseInLambda lam - return $ SOAC.Scatter w lam' inps dests - SOAC.Hist w ops lam arrs -> do - lam' <- simplifyAndFuseInLambda lam - return $ SOAC.Hist w ops lam' arrs - SOAC.Stream w form lam inps -> do - lam' <- simplifyAndFuseInLambda lam - return $ SOAC.Stream w form lam' inps - -simplifyAndFuseInLambda :: Lambda -> FusionGM Lambda -simplifyAndFuseInLambda lam = do - lam' <- simplifyLambda lam - (_, nfres) <- fusionGatherLam (mempty, mkFreshFusionRes) lam' - let nfres' = cleanFusionResult nfres - bindRes nfres' $ fuseInLambda lam' - -copyNewlyConsumed :: Names - -> Futhark.SOAC (Aliases.Aliases SOACS) - -> Binder SOACS (Futhark.SOAC SOACS) -copyNewlyConsumed was_consumed soac = - case soac of - Futhark.Screma w (Futhark.ScremaForm scans reds map_lam) arrs -> do - -- Copy any arrays that are consumed now, but were not in the - -- constituents. - arrs' <- mapM copyConsumedArr arrs - -- Any consumed free variables will have to be copied inside the - -- lambda, and we have to substitute the name of the copy for - -- the original. - map_lam' <- copyFreeInLambda map_lam - - let scans' = map (\scan -> scan { scanLambda = - Aliases.removeLambdaAliases - (scanLambda scan)}) - scans - - let reds' = map (\red -> red { redLambda = - Aliases.removeLambdaAliases - (redLambda red)}) - reds - - return $ Futhark.Screma w (Futhark.ScremaForm scans' reds' map_lam') arrs' - - _ -> return $ removeOpAliases soac - where consumed = consumedInOp soac - newly_consumed = consumed `namesSubtract` was_consumed - - copyConsumedArr a - | a `nameIn` newly_consumed = - letExp (baseString a <> "_copy") $ BasicOp $ Copy a - | otherwise = return a - - copyFreeInLambda lam = do - let free_consumed = consumedByLambda lam `namesSubtract` - namesFromList (map paramName $ lambdaParams lam) - (bnds, subst) <- - foldM copyFree (mempty, mempty) $ namesToList free_consumed - let lam' = Aliases.removeLambdaAliases lam - return $ if null bnds - then lam' - else lam' { lambdaBody = - insertStms bnds $ - substituteNames subst $ lambdaBody lam' - } - - copyFree (bnds, subst) v = do - v_copy <- newVName $ baseString v <> "_copy" - copy <- mkLetNamesM [v_copy] $ BasicOp $ Copy v - return (oneStm copy<>bnds, M.insert v v_copy subst) - ---------------------------------------------------- ---------------------------------------------------- ----- HELPERS ---------------------------------------------------- ---------------------------------------------------- - --- | Get a new fusion result, i.e., for when entering a new scope, --- e.g., a new lambda or a new loop. -mkFreshFusionRes :: FusedRes -mkFreshFusionRes = - FusedRes { rsucc = False, outArr = M.empty, inpArr = M.empty, - infusible = mempty, kernels = M.empty } - -mergeFusionRes :: FusedRes -> FusedRes -> FusionGM FusedRes -mergeFusionRes res1 res2 = do - let ufus_mres = infusible res1 <> infusible res2 - inp_both <- expandSoacInpArr $ M.keys $ inpArr res1 `M.intersection` inpArr res2 - let m_unfus = ufus_mres <> mconcat (map oneName inp_both) - return $ FusedRes (rsucc res1 || rsucc res2) - (outArr res1 `M.union` outArr res2) - (M.unionWith S.union (inpArr res1) (inpArr res2) ) - m_unfus - (kernels res1 `M.union` kernels res2) - - --- | The expression arguments are supposed to be array-type exps. --- Returns a tuple, in which the arrays that are vars are in the --- first element of the tuple, and the one which are indexed or --- transposes (or otherwise transformed) should be in the second. --- --- E.g., for expression `mapT(f, a, b[i])', the result should be --- `([a],[b])' -getIdentArr :: [SOAC.Input] -> ([VName], [VName]) -getIdentArr = foldl comb ([],[]) - where comb (vs,os) (SOAC.Input ts idd _) - | SOAC.nullTransforms ts = (idd:vs, os) - comb (vs, os) inp = - (vs, SOAC.inputArray inp : os) - -cleanFusionResult :: FusedRes -> FusedRes -cleanFusionResult fres = - let newks = M.filter (not . null . fusedVars) (kernels fres) - newoa = M.filter (`M.member` newks) (outArr fres) - newia = M.map (S.filter (`M.member` newks)) (inpArr fres) - in fres { outArr = newoa, inpArr = newia, kernels = newks } - --------------- ---- Errors --- --------------- - -errorIllegal :: String -> FusionGM FusedRes -errorIllegal soac_name = - throwError $ Error - ("In Fusion.hs, soac "++soac_name++" appears illegally in pgm!") +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Perform horizontal and vertical fusion of SOACs. See the paper +-- /A T2 Graph-Reduction Approach To Fusion/ for the basic idea (some +-- extensions discussed in /Design and GPGPU Performance of Futhark’s +-- Redomap Construct/). +module Futhark.Optimise.Fusion (fuseSOACs) where + +import Control.Monad.Except +import Control.Monad.Reader +import Control.Monad.State +import qualified Data.List as L +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import qualified Futhark.Analysis.Alias as Alias +import qualified Futhark.Analysis.HORep.SOAC as SOAC +import Futhark.Construct +import qualified Futhark.IR.Aliases as Aliases +import Futhark.IR.Prop.Aliases +import Futhark.IR.SOACS hiding (SOAC (..)) +import qualified Futhark.IR.SOACS as Futhark +import Futhark.IR.SOACS.Simplify +import Futhark.Optimise.Fusion.LoopKernel +import Futhark.Pass +import Futhark.Transform.Rename +import Futhark.Transform.Substitute +import Futhark.Util (maxinum) + +data VarEntry + = IsArray VName (NameInfo SOACS) Names SOAC.Input + | IsNotArray (NameInfo SOACS) + +varEntryType :: VarEntry -> NameInfo SOACS +varEntryType (IsArray _ dec _ _) = + dec +varEntryType (IsNotArray dec) = + dec + +varEntryAliases :: VarEntry -> Names +varEntryAliases (IsArray _ _ x _) = x +varEntryAliases _ = mempty + +data FusionGEnv = FusionGEnv + { -- | Mapping from variable name to its entire family. + soacs :: M.Map VName [VName], + varsInScope :: M.Map VName VarEntry, + fusedRes :: FusedRes + } + +lookupArr :: VName -> FusionGEnv -> Maybe SOAC.Input +lookupArr v env = asArray =<< M.lookup v (varsInScope env) + where + asArray (IsArray _ _ _ input) = Just input + asArray IsNotArray {} = Nothing + +newtype Error = Error String + +instance Show Error where + show (Error msg) = "Fusion error:\n" ++ msg + +newtype FusionGM a = FusionGM (ExceptT Error (StateT VNameSource (Reader FusionGEnv)) a) + deriving + ( Monad, + Applicative, + Functor, + MonadError Error, + MonadState VNameSource, + MonadReader FusionGEnv + ) + +instance MonadFreshNames FusionGM where + getNameSource = get + putNameSource = put + +instance HasScope SOACS FusionGM where + askScope = asks $ toScope . varsInScope + where + toScope = M.map varEntryType + +------------------------------------------------------------------------ +--- Monadic Helpers: bind/new/runFusionGatherM, etc +------------------------------------------------------------------------ + +-- | Binds an array name to the set of used-array vars +bindVar :: FusionGEnv -> (Ident, Names) -> FusionGEnv +bindVar env (Ident name t, aliases) = + env {varsInScope = M.insert name entry $ varsInScope env} + where + entry = case t of + Array {} -> IsArray name (LetName t) aliases' $ SOAC.identInput $ Ident name t + _ -> IsNotArray $ LetName t + expand = maybe mempty varEntryAliases . flip M.lookup (varsInScope env) + aliases' = aliases <> mconcat (map expand $ namesToList aliases) + +bindVars :: FusionGEnv -> [(Ident, Names)] -> FusionGEnv +bindVars = foldl bindVar + +binding :: [(Ident, Names)] -> FusionGM a -> FusionGM a +binding vs = local (`bindVars` vs) + +gatherStmPattern :: Pattern -> Exp -> FusionGM FusedRes -> FusionGM FusedRes +gatherStmPattern pat e = binding $ zip idents aliases + where + idents = patternIdents pat + aliases = + replicate (length (patternContextNames pat)) mempty + ++ expAliases (Alias.analyseExp mempty e) + +bindingPat :: Pattern -> FusionGM a -> FusionGM a +bindingPat = binding . (`zip` repeat mempty) . patternIdents + +bindingParams :: Typed t => [Param t] -> FusionGM a -> FusionGM a +bindingParams = binding . (`zip` repeat mempty) . map paramIdent + +-- | Binds an array name to the set of soac-produced vars +bindingFamilyVar :: [VName] -> FusionGEnv -> Ident -> FusionGEnv +bindingFamilyVar faml env (Ident nm t) = + env + { soacs = M.insert nm faml $ soacs env, + varsInScope = + M.insert + nm + ( IsArray nm (LetName t) mempty $ + SOAC.identInput $ Ident nm t + ) + $ varsInScope env + } + +varAliases :: VName -> FusionGM Names +varAliases v = + asks $ + (oneName v <>) . maybe mempty varEntryAliases + . M.lookup v + . varsInScope + +varsAliases :: Names -> FusionGM Names +varsAliases = fmap mconcat . mapM varAliases . namesToList + +checkForUpdates :: FusedRes -> Exp -> FusionGM FusedRes +checkForUpdates res (BasicOp (Update src is _)) = do + res' <- + foldM addVarToInfusible res $ + src : namesToList (mconcat $ map freeIn is) + aliases <- varAliases src + let inspectKer k = k {inplace = aliases <> inplace k} + return res' {kernels = M.map inspectKer $ kernels res'} +checkForUpdates res _ = return res + +-- | Updates the environment: (i) the @soacs@ (map) by binding each pattern +-- element identifier to all pattern elements (identifiers) and (ii) the +-- variables in scope (map) by inserting each (pattern-array) name. +-- Finally, if the binding is an in-place update, then the @inplace@ field +-- of each (result) kernel is updated with the new in-place updates. +bindingFamily :: Pattern -> FusionGM FusedRes -> FusionGM FusedRes +bindingFamily pat = local bind + where + idents = patternIdents pat + family = patternNames pat + bind env = foldl (bindingFamilyVar family) env idents + +bindingTransform :: PatElem -> VName -> SOAC.ArrayTransform -> FusionGM a -> FusionGM a +bindingTransform pe srcname trns = local $ \env -> + case M.lookup srcname $ varsInScope env of + Just (IsArray src' _ aliases input) -> + env + { varsInScope = + M.insert + vname + ( IsArray src' (LetName dec) (oneName srcname <> aliases) $ + trns `SOAC.addTransform` input + ) + $ varsInScope env + } + _ -> bindVar env (patElemIdent pe, oneName vname) + where + vname = patElemName pe + dec = patElemDec pe + +-- | Binds the fusion result to the environment. +bindRes :: FusedRes -> FusionGM a -> FusionGM a +bindRes rrr = local (\x -> x {fusedRes = rrr}) + +-- | The fusion transformation runs in this monad. The mutable +-- state refers to the fresh-names engine. +-- The reader hides the vtable that associates ... to ... (fill in, please). +-- The 'Either' monad is used for error handling. +runFusionGatherM :: + MonadFreshNames m => + FusionGM a -> + FusionGEnv -> + m (Either Error a) +runFusionGatherM (FusionGM a) env = + modifyNameSource $ \src -> runReader (runStateT (runExceptT a) src) env + +------------------------------------------------------------------------ +--- Fusion Entry Points: gather the to-be-fused kernels@pgm level --- +--- and fuse them in a second pass! --- +------------------------------------------------------------------------ + +-- | The pass definition. +fuseSOACs :: Pass SOACS SOACS +fuseSOACs = + Pass + { passName = "Fuse SOACs", + passDescription = "Perform higher-order optimisation, i.e., fusion.", + passFunction = \prog -> + simplifySOACS =<< renameProg + =<< intraproceduralTransformationWithConsts + (fuseConsts (freeIn (progFuns prog))) + fuseFun + prog + } + +fuseConsts :: Names -> Stms SOACS -> PassM (Stms SOACS) +fuseConsts used_consts consts = + fuseStms mempty consts $ map Var $ namesToList used_consts + +fuseFun :: Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS) +fuseFun consts fun = do + stms <- + fuseStms + (scopeOf consts <> scopeOfFParams (funDefParams fun)) + (bodyStms $ funDefBody fun) + (bodyResult $ funDefBody fun) + let body = (funDefBody fun) {bodyStms = stms} + return fun {funDefBody = body} + +fuseStms :: Scope SOACS -> Stms SOACS -> Result -> PassM (Stms SOACS) +fuseStms scope stms res = do + let env = + FusionGEnv + { soacs = M.empty, + varsInScope = mempty, + fusedRes = mempty + } + k <- + cleanFusionResult + <$> liftEitherM + ( runFusionGatherM + (binding scope' $ fusionGatherStms mempty (stmsToList stms) res) + env + ) + if not $ rsucc k + then return stms + else liftEitherM $ runFusionGatherM (binding scope' $ bindRes k $ fuseInStms stms) env + where + scope' = map toBind $ M.toList scope + toBind (k, t) = (Ident k $ typeOf t, mempty) + +--------------------------------------------------- +--------------------------------------------------- +---- RESULT's Data Structure +--------------------------------------------------- +--------------------------------------------------- + +-- | A type used for (hopefully) uniquely referring a producer SOAC. +-- The uniquely identifying value is the name of the first array +-- returned from the SOAC. +newtype KernName = KernName {unKernName :: VName} + deriving (Eq, Ord, Show) + +data FusedRes = FusedRes + { -- | Whether we have fused something anywhere. + rsucc :: Bool, + -- | Associates an array to the name of the + -- SOAC kernel that has produced it. + outArr :: M.Map VName KernName, + -- | Associates an array to the names of the + -- SOAC kernels that uses it. These sets include + -- only the SOAC input arrays used as full variables, i.e., no `a[i]'. + inpArr :: M.Map VName (S.Set KernName), + -- | the (names of) arrays that are not fusible, i.e., + -- + -- 1. they are either used other than input to SOAC kernels, or + -- + -- 2. are used as input to at least two different kernels that + -- are not located on disjoint control-flow branches, or + -- + -- 3. are used in the lambda expression of SOACs + infusible :: Names, + -- | The map recording the uses + kernels :: M.Map KernName FusedKer + } + +instance Semigroup FusedRes where + res1 <> res2 = + FusedRes + (rsucc res1 || rsucc res2) + (outArr res1 `M.union` outArr res2) + (M.unionWith S.union (inpArr res1) (inpArr res2)) + (infusible res1 <> infusible res2) + (kernels res1 `M.union` kernels res2) + +instance Monoid FusedRes where + mempty = + FusedRes + { rsucc = False, + outArr = M.empty, + inpArr = M.empty, + infusible = mempty, + kernels = M.empty + } + +isInpArrInResModKers :: FusedRes -> S.Set KernName -> VName -> Bool +isInpArrInResModKers ress kers nm = + case M.lookup nm (inpArr ress) of + Nothing -> False + Just s -> not $ S.null $ s `S.difference` kers + +getKersWithInpArrs :: FusedRes -> [VName] -> S.Set KernName +getKersWithInpArrs ress = + S.unions . mapMaybe (`M.lookup` inpArr ress) + +-- | extend the set of names to include all the names +-- produced via SOACs (by querring the vtable's soac) +expandSoacInpArr :: [VName] -> FusionGM [VName] +expandSoacInpArr = + foldM + ( \y nm -> do + bnd <- asks $ M.lookup nm . soacs + case bnd of + Nothing -> return (y ++ [nm]) + Just nns -> return (y ++ nns) + ) + [] + +---------------------------------------------------------------------- +---------------------------------------------------------------------- + +soacInputs :: SOAC -> FusionGM ([VName], [VName]) +soacInputs soac = do + let (inp_idds, other_idds) = getIdentArr $ SOAC.inputs soac + (inp_nms0, other_nms0) = (inp_idds, other_idds) + inp_nms <- expandSoacInpArr inp_nms0 + other_nms <- expandSoacInpArr other_nms0 + return (inp_nms, other_nms) + +addNewKerWithInfusible :: FusedRes -> ([Ident], StmAux (), SOAC, Names) -> Names -> FusionGM FusedRes +addNewKerWithInfusible res (idd, aux, soac, consumed) ufs = do + nm_ker <- KernName <$> newVName "ker" + scope <- askScope + let out_nms = map identName idd + new_ker = newKernel aux soac consumed out_nms scope + comb = M.unionWith S.union + os' = + M.fromList [(arr, nm_ker) | arr <- out_nms] + `M.union` outArr res + is' = + M.fromList + [ (arr, S.singleton nm_ker) + | arr <- map SOAC.inputArray $ SOAC.inputs soac + ] + `comb` inpArr res + return $ + FusedRes + (rsucc res) + os' + is' + ufs + (M.insert nm_ker new_ker (kernels res)) + +lookupInput :: VName -> FusionGM (Maybe SOAC.Input) +lookupInput name = asks $ lookupArr name + +inlineSOACInput :: SOAC.Input -> FusionGM SOAC.Input +inlineSOACInput (SOAC.Input ts v t) = do + maybe_inp <- lookupInput v + case maybe_inp of + Nothing -> + return $ SOAC.Input ts v t + Just (SOAC.Input ts2 v2 t2) -> + return $ SOAC.Input (ts2 <> ts) v2 t2 + +inlineSOACInputs :: SOAC -> FusionGM SOAC +inlineSOACInputs soac = do + inputs' <- mapM inlineSOACInput $ SOAC.inputs soac + return $ inputs' `SOAC.setInputs` soac + +-- | Attempts to fuse between SOACs. Input: +-- @rem_bnds@ are the bindings remaining in the current body after @orig_soac@. +-- @lam_used_nms@ the infusible names +-- @res@ the fusion result (before processing the current soac) +-- @orig_soac@ and @out_idds@ the current SOAC and its binding pattern +-- @consumed@ is the set of names consumed by the SOAC. +-- Output: a new Fusion Result (after processing the current SOAC binding) +greedyFuse :: + [Stm] -> + Names -> + FusedRes -> + (Pattern, StmAux (), SOAC, Names) -> + FusionGM FusedRes +greedyFuse rem_bnds lam_used_nms res (out_idds, aux, orig_soac, consumed) = do + soac <- inlineSOACInputs orig_soac + (inp_nms, other_nms) <- soacInputs soac + -- Assumption: the free vars in lambda are already in @infusible res@. + let out_nms = patternNames out_idds + isInfusible = (`nameIn` infusible res) + is_screma = case orig_soac of + SOAC.Screma _ form _ -> + (isJust (isRedomapSOAC form) || isJust (isScanomapSOAC form)) + && not (isJust (isReduceSOAC form) || isJust (isScanSOAC form)) + _ -> False + -- + -- Conditions for fusion: + -- If current soac is a replicate OR (current soac a redomap/scanomap AND + -- (i) none of @out_idds@ belongs to the infusible set) + -- THEN try applying producer-consumer fusion + -- ELSE try applying horizontal fusion + -- (without duplicating computation in both cases) + + (ok_kers_compat, fused_kers, fused_nms, old_kers, oldker_nms) <- + if is_screma || any isInfusible out_nms + then horizontGreedyFuse rem_bnds res (out_idds, aux, soac, consumed) + else prodconsGreedyFuse res (out_idds, aux, soac, consumed) + -- + -- (ii) check whether fusing @soac@ will violate any in-place update + -- restriction, e.g., would move an input array past its in-place update. + let all_used_names = namesToList $ mconcat [lam_used_nms, namesFromList inp_nms, namesFromList other_nms] + has_inplace ker = any (`nameIn` inplace ker) all_used_names + ok_inplace = not $ any has_inplace old_kers + -- + -- (iii) there are some kernels that use some of `out_idds' as inputs + -- (iv) and producer-consumer or horizontal fusion succeeds with those. + let fusible_ker = not (null old_kers) && ok_inplace && ok_kers_compat + -- + -- Start constructing the fusion's result: + -- (i) inparr ids other than vars will be added to infusible list, + -- (ii) will also become part of the infusible set the inparr vars + -- that also appear as inparr of another kernel, + -- BUT which said kernel is not the one we are fusing with (now)! + let mod_kerS = if fusible_ker then S.fromList oldker_nms else mempty + let used_inps = filter (isInpArrInResModKers res mod_kerS) inp_nms + let ufs = + mconcat + [ infusible res, + namesFromList used_inps, + namesFromList other_nms + `namesSubtract` namesFromList (map SOAC.inputArray $ SOAC.inputs soac) + ] + let comb = M.unionWith S.union + + if not fusible_ker + then addNewKerWithInfusible res (patternIdents out_idds, aux, soac, consumed) ufs + else do + -- Need to suitably update `inpArr': + -- (i) first remove the inpArr bindings of the old kernel + let inpArr' = + foldl + ( \inpa (kold, knm) -> + S.foldl' + ( \inpp nm -> + case M.lookup nm inpp of + Nothing -> inpp + Just s -> + let new_set = S.delete knm s + in if S.null new_set + then M.delete nm inpp + else M.insert nm new_set inpp + ) + inpa + $ arrInputs kold + ) + (inpArr res) + (zip old_kers oldker_nms) + -- (ii) then add the inpArr bindings of the new kernel + let fused_ker_nms = zip fused_nms fused_kers + inpArr'' = + foldl + ( \inpa' (knm, knew) -> + M.fromList + [ (k, S.singleton knm) + | k <- S.toList $ arrInputs knew + ] + `comb` inpa' + ) + inpArr' + fused_ker_nms + -- Update the kernels map (why not delete the ones that have been fused?) + let kernels' = M.fromList fused_ker_nms `M.union` kernels res + -- nothing to do for `outArr' (since we have not added a new kernel) + -- DO IMPROVEMENT: attempt to fuse the resulting kernel AGAIN until it fails, + -- but make sure NOT to add a new kernel! + return $ FusedRes True (outArr res) inpArr'' ufs kernels' + +prodconsGreedyFuse :: + FusedRes -> + (Pattern, StmAux (), SOAC, Names) -> + FusionGM (Bool, [FusedKer], [KernName], [FusedKer], [KernName]) +prodconsGreedyFuse res (out_idds, aux, soac, consumed) = do + let out_nms = patternNames out_idds -- Extract VNames from output patterns + to_fuse_knmSet = getKersWithInpArrs res out_nms -- Find kernels which consume outputs + to_fuse_knms = S.toList to_fuse_knmSet + lookup_kern k = case M.lookup k (kernels res) of + Nothing -> + throwError $ + Error + ( "In Fusion.hs, greedyFuse, comp of to_fuse_kers: " + ++ "kernel name not found in kernels field!" + ) + Just ker -> return ker + to_fuse_kers <- mapM lookup_kern to_fuse_knms -- Get all consumer kernels + -- try producer-consumer fusion + (ok_kers_compat, fused_kers) <- do + kers <- + forM to_fuse_kers $ + attemptFusion mempty (patternNames out_idds) soac consumed + case sequence kers of + Nothing -> return (False, []) + Just kers' -> return (True, map certifyKer kers') + return (ok_kers_compat, fused_kers, to_fuse_knms, to_fuse_kers, to_fuse_knms) + where + certifyKer k = k {kerAux = kerAux k <> aux} + +horizontGreedyFuse :: + [Stm] -> + FusedRes -> + (Pattern, StmAux (), SOAC, Names) -> + FusionGM (Bool, [FusedKer], [KernName], [FusedKer], [KernName]) +horizontGreedyFuse rem_bnds res (out_idds, aux, soac, consumed) = do + (inp_nms, _) <- soacInputs soac + let out_nms = patternNames out_idds + infusible_nms = namesFromList $ filter (`nameIn` infusible res) out_nms + out_arr_nms = case soac of + -- 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 + _ -> out_nms + to_fuse_knms1 = S.toList $ getKersWithInpArrs res (out_arr_nms ++ inp_nms) + to_fuse_knms2 = getKersWithSameInpSize (SOAC.width soac) res + to_fuse_knms = S.toList $ S.fromList $ to_fuse_knms1 ++ to_fuse_knms2 + lookupKernel k = case M.lookup k (kernels res) of + Nothing -> + throwError $ + Error + ( "In Fusion.hs, greedyFuse, comp of to_fuse_kers: " + ++ "kernel name not found in kernels field!" + ) + Just ker -> return ker + + -- For each kernel get the index in the bindings where the kernel is + -- located and sort based on the index so that partial fusion may + -- succeed. We use the last position where one of the kernel + -- outputs occur. + let bnd_nms = map (patternNames . stmPattern) rem_bnds + kernminds <- forM to_fuse_knms $ \ker_nm -> do + ker <- lookupKernel ker_nm + case mapMaybe (\out_nm -> L.findIndex (elem out_nm) bnd_nms) (outNames ker) of + [] -> return Nothing + is -> return $ Just (ker, ker_nm, maxinum is) + + scope <- askScope + let kernminds' = L.sortBy (\(_, _, i1) (_, _, i2) -> compare i1 i2) $ catMaybes kernminds + soac_kernel = newKernel aux soac consumed out_nms scope + + -- now try to fuse kernels one by one (in a fold); @ok_ind@ is the index of the + -- kernel until which fusion succeded, and @fused_ker@ is the resulting kernel. + use_scope <- (<> scopeOf rem_bnds) <$> askScope + (_, ok_ind, _, fused_ker, _) <- + foldM + ( \(cur_ok, n, prev_ind, cur_ker, ufus_nms) (ker, _ker_nm, bnd_ind) -> do + -- check that we still try fusion and that the intermediate + -- bindings do not use the results of cur_ker + let curker_outnms = outNames cur_ker + curker_outset = namesFromList curker_outnms + new_ufus_nms = namesFromList $ outNames ker ++ namesToList ufus_nms + -- disable horizontal fusion in the case when an output array of + -- producer SOAC is a non-trivially transformed input of the consumer + out_transf_ok = + let ker_inp = SOAC.inputs $ fsoac ker + unfuse1 = + namesFromList (map SOAC.inputArray ker_inp) + `namesSubtract` namesFromList (mapMaybe SOAC.isVarInput ker_inp) + unfuse2 = namesIntersection curker_outset ufus_nms + in not $ unfuse1 `namesIntersect` unfuse2 + -- Disable horizontal fusion if consumer has any + -- output transforms. + cons_no_out_transf = SOAC.nullTransforms $ outputTransform ker + + consumer_ok <- do + let consumer_bnd = rem_bnds !! bnd_ind + maybesoac <- runReaderT (SOAC.fromExp $ stmExp consumer_bnd) use_scope + case maybesoac of + -- check that consumer's lambda body does not use + -- directly the produced arrays (e.g., see noFusion3.fut). + Right conssoac -> + return $ + not $ + curker_outset + `namesIntersect` freeIn (lambdaBody $ SOAC.lambda conssoac) + Left _ -> return True + + let interm_bnds_ok = + cur_ok && consumer_ok && out_transf_ok && cons_no_out_transf + && foldl + ( \ok bnd -> + ok + && not (curker_outset `namesIntersect` freeIn (stmExp bnd)) -- hardwired to False after first fail + -- (i) check that the in-between bindings do + -- not use the result of current kernel OR + || + --(ii) that the pattern-binding corresponds to + -- the result of the consumer kernel; in the + -- latter case it means it corresponds to a + -- kernel that has been fused in the consumer, + -- hence it should be ignored + not + ( null $ + curker_outnms + `L.intersect` patternNames (stmPattern bnd) + ) + ) + True + (drop (prev_ind + 1) $ take bnd_ind rem_bnds) + if not interm_bnds_ok + then return (False, n, bnd_ind, cur_ker, mempty) + else do + new_ker <- + attemptFusion + ufus_nms + (outNames cur_ker) + (fsoac cur_ker) + (fusedConsumed cur_ker) + ker + case new_ker of + Nothing -> return (False, n, bnd_ind, cur_ker, mempty) + Just krn -> + let krn' = krn {kerAux = aux <> kerAux krn} + in return (True, n + 1, bnd_ind, krn', new_ufus_nms) + ) + (True, 0, 0, soac_kernel, infusible_nms) + kernminds' + + -- Find the kernels we have fused into and the name of the last such + -- kernel (if any). + let (to_fuse_kers', to_fuse_knms', _) = unzip3 $ take ok_ind kernminds' + new_kernms = drop (ok_ind -1) to_fuse_knms' + + return (ok_ind > 0, [fused_ker], new_kernms, to_fuse_kers', to_fuse_knms') + where + getKersWithSameInpSize :: SubExp -> FusedRes -> [KernName] + getKersWithSameInpSize sz ress = + map fst $ filter (\(_, ker) -> sz == SOAC.width (fsoac ker)) $ M.toList $ kernels ress + +------------------------------------------------------------------------ +------------------------------------------------------------------------ +------------------------------------------------------------------------ +--- Fusion Gather for EXPRESSIONS and BODIES, --- +--- i.e., where work is being done: --- +--- i) bottom-up AbSyn traversal (backward analysis) --- +--- ii) soacs are fused greedily iff does not duplicate computation--- +--- E.g., (y1, y2, y3) = mapT(f, x1, x2[i]) --- +--- (z1, z2) = mapT(g1, y1, y2) --- +--- (q1, q2) = mapT(g2, y3, z1, a, y3) --- +--- res = reduce(op, ne, q1, q2, z2, y1, y3) --- +--- can be fused if y1,y2,y3, z1,z2, q1,q2 are not used elsewhere: --- +--- res = redomap(op, \(x1,x2i,a)-> --- +--- let (y1,y2,y3) = f (x1, x2i) in--- +--- let (z1,z2) = g1(y1, y2) in--- +--- let (q1,q2) = g2(y3, z1, a, y3) in--- +--- (q1, q2, z2, y1, y3) --- +--- x1, x2[i], a) --- +------------------------------------------------------------------------ +------------------------------------------------------------------------ +------------------------------------------------------------------------ + +fusionGatherBody :: FusedRes -> Body -> FusionGM FusedRes +fusionGatherBody fres (Body _ stms res) = + fusionGatherStms fres (stmsToList stms) res + +fusionGatherStms :: FusedRes -> [Stm] -> Result -> FusionGM FusedRes +-- Some forms of do-loops can profitably be considered streamSeqs. We +-- are careful to ensure that the generated nested loop cannot itself +-- be considered a stream, to avoid infinite recursion. +fusionGatherStms + fres + ( Let + (Pattern [] pes) + bndtp + (DoLoop [] merge (ForLoop i it w loop_vars) body) + : bnds + ) + res + | not $ null loop_vars = do + let (merge_params, merge_init) = unzip merge + (loop_params, loop_arrs) = unzip loop_vars + chunk_size <- newVName "chunk_size" + offset <- newVName "offset" + let chunk_param = Param chunk_size $ Prim int64 + offset_param = Param offset $ Prim $ IntType it + + acc_params <- forM merge_params $ \p -> + Param <$> newVName (baseString (paramName p) ++ "_outer") + <*> pure (paramType p) + + chunked_params <- forM loop_vars $ \(p, arr) -> + Param <$> newVName (baseString arr ++ "_chunk") + <*> pure (paramType p `arrayOfRow` Futhark.Var chunk_size) + + let lam_params = chunk_param : acc_params ++ [offset_param] ++ chunked_params + + lam_body <- runBodyBinder $ + localScope (scopeOfLParams lam_params) $ do + let merge' = zip merge_params $ map (Futhark.Var . paramName) acc_params + j <- newVName "j" + loop_body <- runBodyBinder $ do + forM_ (zip loop_params chunked_params) $ \(p, a_p) -> + letBindNames [paramName p] $ + BasicOp $ + Index (paramName a_p) $ + fullSlice (paramType a_p) [DimFix $ Futhark.Var j] + letBindNames [i] $ BasicOp $ BinOp (Add it OverflowUndef) (Futhark.Var offset) (Futhark.Var j) + return body + eBody + [ pure $ + DoLoop [] merge' (ForLoop j it (Futhark.Var chunk_size) []) loop_body, + pure $ + BasicOp $ BinOp (Add Int64 OverflowUndef) (Futhark.Var offset) (Futhark.Var chunk_size) + ] + let lam = + Lambda + { lambdaParams = lam_params, + lambdaBody = lam_body, + lambdaReturnType = map paramType $ acc_params ++ [offset_param] + } + stream = Futhark.Stream w (Sequential $ merge_init ++ [intConst it 0]) lam 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 + -- identify the SOAC in the second phase of fusion. + discard <- newVName "discard" + let discard_pe = PatElem discard $ Prim int64 + + fusionGatherStms + fres + (Let (Pattern [] (pes <> [discard_pe])) bndtp (Op stream) : bnds) + res +fusionGatherStms fres (bnd@(Let pat _ e) : bnds) res = do + maybesoac <- SOAC.fromExp e + case maybesoac of + Right soac@(SOAC.Scatter _len lam _ivs _as) -> do + -- We put the variables produced by Scatter into the infusible + -- set to force horizontal fusion. It is not possible to + -- producer/consumer-fuse Scatter anyway. + fres' <- addNamesToInfusible fres $ namesFromList $ patternNames pat + mapLike fres' soac lam + Right soac@(SOAC.Hist _ _ lam _) -> do + -- We put the variables produced by Hist into the infusible + -- set to force horizontal fusion. It is not possible to + -- producer/consumer-fuse Hist anyway. + fres' <- addNamesToInfusible fres $ namesFromList $ patternNames pat + mapLike fres' soac lam + 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 + -- 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 + _ + | [pe] <- patternValueElements pat, + Just (src, trns) <- SOAC.transformFromExp (stmCerts bnd) e -> + bindingTransform pe src trns $ fusionGatherStms fres bnds res + | otherwise -> do + let pat_vars = map (BasicOp . SubExp . Var) $ patternNames pat + bres <- gatherStmPattern pat e $ fusionGatherStms fres bnds res + bres' <- checkForUpdates bres e + foldM fusionGatherExp bres' (e : pat_vars) + where + aux = stmAux bnd + rem_bnds = bnd : bnds + consumed = consumedInExp $ Alias.analyseExp mempty e + + reduceLike soac lambdas nes = do + (used_lam, lres) <- foldM fusionGatherLam (mempty, fres) lambdas + bres <- bindingFamily pat $ fusionGatherStms lres bnds res + bres' <- foldM fusionGatherSubExp bres nes + consumed' <- varsAliases consumed + greedyFuse rem_bnds used_lam bres' (pat, aux, soac, consumed') + + mapLike fres' soac lambda = do + bres <- bindingFamily pat $ fusionGatherStms fres' bnds res + (used_lam, blres) <- fusionGatherLam (mempty, bres) lambda + consumed' <- varsAliases consumed + greedyFuse rem_bnds used_lam blres (pat, aux, soac, consumed') +fusionGatherStms fres [] res = + foldM fusionGatherExp fres $ map (BasicOp . SubExp) res + +fusionGatherExp :: FusedRes -> Exp -> FusionGM FusedRes +----------------------------------------- +---- Index/If ---- +----------------------------------------- + +fusionGatherExp fres (DoLoop ctx val form loop_body) = do + fres' <- addNamesToInfusible fres $ freeIn form <> freeIn ctx <> freeIn val + let form_idents = + case form of + ForLoop i it _ loopvars -> + Ident i (Prim (IntType it)) : map (paramIdent . fst) loopvars + WhileLoop {} -> [] + + new_res <- + binding + ( zip (form_idents ++ map (paramIdent . fst) (ctx <> val)) $ + repeat mempty + ) + $ fusionGatherBody mempty loop_body + -- make the inpArr infusible, so that they + -- cannot be fused from outside the loop: + let (inp_arrs, _) = unzip $ M.toList $ inpArr new_res + let new_res' = new_res {infusible = infusible new_res <> mconcat (map oneName inp_arrs)} + -- merge new_res with fres' + return $ new_res' <> fres' +fusionGatherExp fres (If cond e_then e_else _) = do + then_res <- fusionGatherBody mempty e_then + else_res <- fusionGatherBody mempty e_else + let both_res = then_res <> else_res + fres' <- fusionGatherSubExp fres cond + mergeFusionRes fres' both_res + +----------------------------------------------------------------------------------- +--- Errors: all SOACs, (because normalization ensures they appear +--- directly in let exp, i.e., let x = e) +----------------------------------------------------------------------------------- + +fusionGatherExp _ (Op Futhark.Screma {}) = errorIllegal "screma" +fusionGatherExp _ (Op Futhark.Scatter {}) = errorIllegal "write" +----------------------------------- +---- Generic Traversal ---- +----------------------------------- + +fusionGatherExp fres e = addNamesToInfusible fres $ freeIn e + +fusionGatherSubExp :: FusedRes -> SubExp -> FusionGM FusedRes +fusionGatherSubExp fres (Var idd) = addVarToInfusible fres idd +fusionGatherSubExp fres _ = return fres + +addNamesToInfusible :: FusedRes -> Names -> FusionGM FusedRes +addNamesToInfusible fres = foldM addVarToInfusible fres . namesToList + +addVarToInfusible :: FusedRes -> VName -> FusionGM FusedRes +addVarToInfusible fres name = do + trns <- asks $ lookupArr name + let name' = case trns of + Nothing -> name + Just (SOAC.Input _ orig _) -> orig + return fres {infusible = oneName name' <> infusible fres} + +-- Lambdas create a new scope. Disallow fusing from outside lambda by +-- adding inp_arrs to the infusible set. +fusionGatherLam :: (Names, FusedRes) -> Lambda -> FusionGM (Names, FusedRes) +fusionGatherLam (u_set, fres) (Lambda idds body _) = do + new_res <- bindingParams idds $ fusionGatherBody mempty body + -- make the inpArr infusible, so that they + -- cannot be fused from outside the lambda: + let inp_arrs = namesFromList $ M.keys $ inpArr new_res + let unfus = infusible new_res <> inp_arrs + bnds <- asks $ M.keys . varsInScope + let unfus' = unfus `namesIntersection` namesFromList bnds + -- merge fres with new_res' + let new_res' = new_res {infusible = unfus'} + -- merge new_res with fres' + return (u_set <> unfus', new_res' <> fres) + +------------------------------------------------------------- +------------------------------------------------------------- +--- FINALLY, Substitute the kernels in function +------------------------------------------------------------- +------------------------------------------------------------- + +fuseInStms :: Stms SOACS -> FusionGM (Stms SOACS) +fuseInStms stms + | Just (Let pat aux e, stms') <- stmsHead stms = do + stms'' <- bindingPat pat $ fuseInStms stms' + soac_bnds <- replaceSOAC pat aux e + pure $ soac_bnds <> stms'' + | otherwise = + pure mempty + +fuseInBody :: Body -> FusionGM Body +fuseInBody (Body _ stms res) = + Body () <$> fuseInStms stms <*> pure res + +fuseInExp :: Exp -> FusionGM Exp +-- Handle loop specially because we need to bind the types of the +-- merge variables. +fuseInExp (DoLoop ctx val form loopbody) = + binding (zip form_idents $ repeat mempty) $ + bindingParams (map fst $ ctx ++ val) $ + DoLoop ctx val form <$> fuseInBody loopbody + where + form_idents = case form of + WhileLoop {} -> [] + ForLoop i it _ loopvars -> + Ident i (Prim $ IntType it) : + map (paramIdent . fst) loopvars +fuseInExp e = mapExpM fuseIn e + +fuseIn :: Mapper SOACS SOACS FusionGM +fuseIn = + identityMapper + { mapOnBody = const fuseInBody, + mapOnOp = mapSOACM identitySOACMapper {mapOnSOACLambda = fuseInLambda} + } + +fuseInLambda :: Lambda -> FusionGM Lambda +fuseInLambda (Lambda params body rtp) = do + body' <- bindingParams params $ fuseInBody body + return $ Lambda params body' rtp + +replaceSOAC :: Pattern -> StmAux () -> Exp -> FusionGM (Stms SOACS) +replaceSOAC (Pattern _ []) _ _ = return mempty +replaceSOAC pat@(Pattern _ (patElem : _)) aux e = do + fres <- asks fusedRes + let pat_nm = patElemName patElem + names = patternIdents pat + case M.lookup pat_nm (outArr fres) of + Nothing -> + oneStm . Let pat aux <$> fuseInExp e + Just knm -> + case M.lookup knm (kernels fres) of + Nothing -> + throwError $ + Error + ( "In Fusion.hs, replaceSOAC, outArr in ker_name " + ++ "which is not in Res: " + ++ pretty (unKernName knm) + ) + Just ker -> do + when (null $ fusedVars ker) $ + throwError $ + Error + ( "In Fusion.hs, replaceSOAC, unfused kernel " + ++ "still in result: " + ++ pretty names + ) + insertKerSOAC aux (outNames ker) ker + +insertKerSOAC :: StmAux () -> [VName] -> FusedKer -> FusionGM (Stms SOACS) +insertKerSOAC aux names ker = do + new_soac' <- finaliseSOAC $ fsoac ker + runBinder_ $ do + 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 + 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 + +-- | Perform simplification and fusion inside the lambda(s) of a SOAC. +finaliseSOAC :: SOAC.SOAC SOACS -> FusionGM (SOAC.SOAC SOACS) +finaliseSOAC new_soac = + case new_soac of + SOAC.Screma w (ScremaForm scans reds map_lam) arrs -> do + scans' <- forM scans $ \(Scan scan_lam scan_nes) -> do + scan_lam' <- simplifyAndFuseInLambda scan_lam + return $ Scan scan_lam' scan_nes + + reds' <- forM reds $ \(Reduce comm red_lam red_nes) -> do + red_lam' <- simplifyAndFuseInLambda red_lam + return $ Reduce comm red_lam' red_nes + + map_lam' <- simplifyAndFuseInLambda map_lam + + return $ SOAC.Screma w (ScremaForm scans' reds' map_lam') arrs + SOAC.Scatter w lam inps dests -> do + lam' <- simplifyAndFuseInLambda lam + return $ SOAC.Scatter w lam' inps dests + SOAC.Hist w ops lam arrs -> do + lam' <- simplifyAndFuseInLambda lam + return $ SOAC.Hist w ops lam' arrs + SOAC.Stream w form lam inps -> do + lam' <- simplifyAndFuseInLambda lam + return $ SOAC.Stream w form lam' inps + +simplifyAndFuseInLambda :: Lambda -> FusionGM Lambda +simplifyAndFuseInLambda lam = do + lam' <- simplifyLambda lam + (_, nfres) <- fusionGatherLam (mempty, mkFreshFusionRes) lam' + let nfres' = cleanFusionResult nfres + bindRes nfres' $ fuseInLambda lam' + +copyNewlyConsumed :: + Names -> + Futhark.SOAC (Aliases.Aliases SOACS) -> + Binder SOACS (Futhark.SOAC SOACS) +copyNewlyConsumed was_consumed soac = + case soac of + Futhark.Screma w (Futhark.ScremaForm scans reds map_lam) arrs -> do + -- Copy any arrays that are consumed now, but were not in the + -- constituents. + arrs' <- mapM copyConsumedArr arrs + -- Any consumed free variables will have to be copied inside the + -- lambda, and we have to substitute the name of the copy for + -- the original. + map_lam' <- copyFreeInLambda map_lam + + let scans' = + map + ( \scan -> + scan + { scanLambda = + Aliases.removeLambdaAliases + (scanLambda scan) + } + ) + scans + + let reds' = + map + ( \red -> + red + { redLambda = + Aliases.removeLambdaAliases + (redLambda red) + } + ) + reds + + return $ Futhark.Screma w (Futhark.ScremaForm scans' reds' map_lam') arrs' + _ -> return $ removeOpAliases soac + where + consumed = consumedInOp soac + newly_consumed = consumed `namesSubtract` was_consumed + + copyConsumedArr a + | a `nameIn` newly_consumed = + letExp (baseString a <> "_copy") $ BasicOp $ Copy a + | otherwise = return a + + copyFreeInLambda lam = do + let free_consumed = + consumedByLambda lam + `namesSubtract` namesFromList (map paramName $ lambdaParams lam) + (bnds, subst) <- + foldM copyFree (mempty, mempty) $ namesToList free_consumed + let lam' = Aliases.removeLambdaAliases lam + return $ + if null bnds + then lam' + else + lam' + { lambdaBody = + insertStms bnds $ + substituteNames subst $ lambdaBody lam' + } + + copyFree (bnds, subst) v = do + v_copy <- newVName $ baseString v <> "_copy" + copy <- mkLetNamesM [v_copy] $ BasicOp $ Copy v + return (oneStm copy <> bnds, M.insert v v_copy subst) + +--------------------------------------------------- +--------------------------------------------------- +---- HELPERS +--------------------------------------------------- +--------------------------------------------------- + +-- | Get a new fusion result, i.e., for when entering a new scope, +-- e.g., a new lambda or a new loop. +mkFreshFusionRes :: FusedRes +mkFreshFusionRes = + FusedRes + { rsucc = False, + outArr = M.empty, + inpArr = M.empty, + infusible = mempty, + kernels = M.empty + } + +mergeFusionRes :: FusedRes -> FusedRes -> FusionGM FusedRes +mergeFusionRes res1 res2 = do + let ufus_mres = infusible res1 <> infusible res2 + inp_both <- expandSoacInpArr $ M.keys $ inpArr res1 `M.intersection` inpArr res2 + let m_unfus = ufus_mres <> mconcat (map oneName inp_both) + return $ + FusedRes + (rsucc res1 || rsucc res2) + (outArr res1 `M.union` outArr res2) + (M.unionWith S.union (inpArr res1) (inpArr res2)) + m_unfus + (kernels res1 `M.union` kernels res2) + +-- | The expression arguments are supposed to be array-type exps. +-- Returns a tuple, in which the arrays that are vars are in the +-- first element of the tuple, and the one which are indexed or +-- transposes (or otherwise transformed) should be in the second. +-- +-- E.g., for expression `mapT(f, a, b[i])', the result should be +-- `([a],[b])' +getIdentArr :: [SOAC.Input] -> ([VName], [VName]) +getIdentArr = foldl comb ([], []) + where + comb (vs, os) (SOAC.Input ts idd _) + | SOAC.nullTransforms ts = (idd : vs, os) + comb (vs, os) inp = + (vs, SOAC.inputArray inp : os) + +cleanFusionResult :: FusedRes -> FusedRes +cleanFusionResult fres = + let newks = M.filter (not . null . fusedVars) (kernels fres) + newoa = M.filter (`M.member` newks) (outArr fres) + newia = M.map (S.filter (`M.member` newks)) (inpArr fres) + in fres {outArr = newoa, inpArr = newia, kernels = newks} + +-------------- +--- Errors --- +-------------- + +errorIllegal :: String -> FusionGM FusedRes +errorIllegal soac_name = + throwError $ + Error + ("In Fusion.hs, soac " ++ soac_name ++ " appears illegally in pgm!") diff --git a/src/Futhark/Optimise/Fusion/Composing.hs b/src/Futhark/Optimise/Fusion/Composing.hs --- a/src/Futhark/Optimise/Fusion/Composing.hs +++ b/src/Futhark/Optimise/Fusion/Composing.hs @@ -11,22 +11,20 @@ -- -- The module will, however, remove duplicate inputs after fusion. module Futhark.Optimise.Fusion.Composing - ( fuseMaps - , fuseRedomap - , mergeReduceOps + ( fuseMaps, + fuseRedomap, + mergeReduceOps, ) - where +where import Data.List (mapAccumL) import qualified Data.Map.Strict as M import Data.Maybe - import qualified Futhark.Analysis.HORep.SOAC as SOAC - -import Futhark.IR -import Futhark.Binder (Bindable(..), insertStm, insertStms, mkLet) +import Futhark.Binder (Bindable (..), insertStm, insertStms, mkLet) import Futhark.Construct (mapResult) -import Futhark.Util (splitAt3, takeLast, dropLast) +import Futhark.IR +import Futhark.Util (dropLast, splitAt3, takeLast) -- | @fuseMaps lam1 inp1 out1 lam2 inp2@ fuses the function @lam1@ into -- @lam2@. Both functions must be mapping functions, although @lam2@ @@ -45,167 +43,252 @@ -- -- The result is the fused function, and a list of the array inputs -- expected by the SOAC containing the fused function. -fuseMaps :: Bindable lore => - Names -- ^ The producer var names that still need to be returned - -> Lambda lore -- ^ Function of SOAC to be fused. - -> [SOAC.Input] -- ^ Input of SOAC to be fused. - -> [(VName,Ident)] -- ^ Output of SOAC to be fused. The - -- first identifier is the name of the - -- actual output, where the second output - -- is an identifier that can be used to - -- bind a single element of that output. - -> Lambda lore -- ^ Function to be fused with. - -> [SOAC.Input] -- ^ Input of SOAC to be fused with. - -> (Lambda lore, [SOAC.Input]) -- ^ The fused lambda and the inputs of - -- the resulting SOAC. +fuseMaps :: + Bindable lore => + -- | The producer var names that still need to be returned + Names -> + -- | Function of SOAC to be fused. + Lambda lore -> + -- | Input of SOAC to be fused. + [SOAC.Input] -> + -- | Output of SOAC to be fused. The + -- first identifier is the name of the + -- actual output, where the second output + -- is an identifier that can be used to + -- bind a single element of that output. + [(VName, Ident)] -> + -- | Function to be fused with. + Lambda lore -> + -- | Input of SOAC to be fused with. + [SOAC.Input] -> + -- | The fused lambda and the inputs of + -- the resulting SOAC. + (Lambda lore, [SOAC.Input]) fuseMaps unfus_nms lam1 inp1 out1 lam2 inp2 = (lam2', M.elems inputmap) - where lam2' = - lam2 { lambdaParams = [ Param name t - | Ident name t <- lam2redparams ++ M.keys inputmap ] - , lambdaBody = new_body2' - } - new_body2 = let bnds res = [ mkLet [] [p] $ BasicOp $ SubExp e - | (p,e) <- zip pat res] - bindLambda res = - stmsFromList (bnds res) `insertStms` makeCopiesInner (lambdaBody lam2) - in makeCopies $ mapResult bindLambda (lambdaBody lam1) - new_body2_rses = bodyResult new_body2 - new_body2'= new_body2 { bodyResult = new_body2_rses ++ - map (Var . identName) unfus_pat } - -- infusible variables are added at the end of the result/pattern/type - (lam2redparams, unfus_pat, pat, inputmap, makeCopies, makeCopiesInner) = - fuseInputs unfus_nms lam1 inp1 out1 lam2 inp2 - --(unfus_accpat, unfus_arrpat) = splitAt (length unfus_accs) unfus_pat + where + lam2' = + lam2 + { lambdaParams = + [ Param name t + | Ident name t <- lam2redparams ++ M.keys inputmap + ], + lambdaBody = new_body2' + } + new_body2 = + let bnds res = + [ mkLet [] [p] $ BasicOp $ SubExp e + | (p, e) <- zip pat res + ] + bindLambda res = + stmsFromList (bnds res) `insertStms` makeCopiesInner (lambdaBody lam2) + in makeCopies $ mapResult bindLambda (lambdaBody lam1) + new_body2_rses = bodyResult new_body2 + new_body2' = + new_body2 + { bodyResult = + new_body2_rses + ++ map (Var . identName) unfus_pat + } + -- infusible variables are added at the end of the result/pattern/type + (lam2redparams, unfus_pat, pat, inputmap, makeCopies, makeCopiesInner) = + fuseInputs unfus_nms lam1 inp1 out1 lam2 inp2 -fuseInputs :: Bindable lore => - Names - -> Lambda lore -> [SOAC.Input] -> [(VName,Ident)] - -> Lambda lore -> [SOAC.Input] - -> ([Ident], [Ident], [Ident], - M.Map Ident SOAC.Input, - Body lore -> Body lore, Body lore -> Body lore) +--(unfus_accpat, unfus_arrpat) = splitAt (length unfus_accs) unfus_pat + +fuseInputs :: + Bindable lore => + Names -> + Lambda lore -> + [SOAC.Input] -> + [(VName, Ident)] -> + Lambda lore -> + [SOAC.Input] -> + ( [Ident], + [Ident], + [Ident], + M.Map Ident SOAC.Input, + Body lore -> Body lore, + Body lore -> Body lore + ) fuseInputs unfus_nms lam1 inp1 out1 lam2 inp2 = (lam2redparams, unfus_vars, outbnds, inputmap, makeCopies, makeCopiesInner) - where (lam2redparams, lam2arrparams) = - splitAt (length lam2params - length inp2) lam2params - lam1params = map paramIdent $ lambdaParams lam1 - lam2params = map paramIdent $ lambdaParams lam2 - lam1inputmap = M.fromList $ zip lam1params inp1 - lam2inputmap = M.fromList $ zip lam2arrparams inp2 - (lam2inputmap', makeCopiesInner) = removeDuplicateInputs lam2inputmap - originputmap = lam1inputmap `M.union` lam2inputmap' - outins = uncurry (outParams $ map fst out1) $ - unzip $ M.toList lam2inputmap' - outbnds= filterOutParams out1 outins - (inputmap, makeCopies) = - removeDuplicateInputs $ originputmap `M.difference` outins - -- Cosmin: @unfus_vars@ is supposed to be the lam2 vars corresponding to unfus_nms (?) - getVarParPair x = case SOAC.isVarInput (snd x) of - Just nm -> Just (nm, fst x) - Nothing -> Nothing --should not be reached! - outinsrev = M.fromList $ mapMaybe getVarParPair $ M.toList outins - unfusible outname - | outname `nameIn` unfus_nms = - outname `M.lookup` M.union outinsrev (M.fromList out1) - unfusible _ = Nothing - unfus_vars= mapMaybe (unfusible . fst) out1 + where + (lam2redparams, lam2arrparams) = + splitAt (length lam2params - length inp2) lam2params + lam1params = map paramIdent $ lambdaParams lam1 + lam2params = map paramIdent $ lambdaParams lam2 + lam1inputmap = M.fromList $ zip lam1params inp1 + lam2inputmap = M.fromList $ zip lam2arrparams inp2 + (lam2inputmap', makeCopiesInner) = removeDuplicateInputs lam2inputmap + originputmap = lam1inputmap `M.union` lam2inputmap' + outins = + uncurry (outParams $ map fst out1) $ + unzip $ M.toList lam2inputmap' + outbnds = filterOutParams out1 outins + (inputmap, makeCopies) = + removeDuplicateInputs $ originputmap `M.difference` outins + -- Cosmin: @unfus_vars@ is supposed to be the lam2 vars corresponding to unfus_nms (?) + getVarParPair x = case SOAC.isVarInput (snd x) of + Just nm -> Just (nm, fst x) + Nothing -> Nothing --should not be reached! + outinsrev = M.fromList $ mapMaybe getVarParPair $ M.toList outins + unfusible outname + | outname `nameIn` unfus_nms = + outname `M.lookup` M.union outinsrev (M.fromList out1) + unfusible _ = Nothing + unfus_vars = mapMaybe (unfusible . fst) out1 -outParams :: [VName] -> [Ident] -> [SOAC.Input] - -> M.Map Ident SOAC.Input +outParams :: + [VName] -> + [Ident] -> + [SOAC.Input] -> + M.Map Ident SOAC.Input outParams out1 lam2arrparams inp2 = M.fromList $ mapMaybe isOutParam $ zip lam2arrparams inp2 - where isOutParam (p, inp) - | Just a <- SOAC.isVarInput inp, - a `elem` out1 = Just (p, inp) - isOutParam _ = Nothing + where + isOutParam (p, inp) + | Just a <- SOAC.isVarInput inp, + a `elem` out1 = + Just (p, inp) + isOutParam _ = Nothing -filterOutParams :: [(VName,Ident)] - -> M.Map Ident SOAC.Input - -> [Ident] +filterOutParams :: + [(VName, Ident)] -> + M.Map Ident SOAC.Input -> + [Ident] filterOutParams out1 outins = snd $ mapAccumL checkUsed outUsage out1 - where outUsage = M.foldlWithKey' add M.empty outins - where add m p inp = - case SOAC.isVarInput inp of - Just v -> M.insertWith (++) v [p] m - Nothing -> m + where + outUsage = M.foldlWithKey' add M.empty outins + where + add m p inp = + case SOAC.isVarInput inp of + Just v -> M.insertWith (++) v [p] m + Nothing -> m - checkUsed m (a,ra) = - case M.lookup a m of - Just (p:ps) -> (M.insert a ps m, p) - _ -> (m, ra) + checkUsed m (a, ra) = + case M.lookup a m of + Just (p : ps) -> (M.insert a ps m, p) + _ -> (m, ra) -removeDuplicateInputs :: Bindable lore => - M.Map Ident SOAC.Input - -> (M.Map Ident SOAC.Input, Body lore -> Body lore) +removeDuplicateInputs :: + Bindable lore => + M.Map Ident SOAC.Input -> + (M.Map Ident SOAC.Input, Body lore -> Body lore) removeDuplicateInputs = fst . M.foldlWithKey' comb ((M.empty, id), M.empty) - where comb ((parmap, inner), arrmap) par arr = - case M.lookup arr arrmap of - Nothing -> ((M.insert par arr parmap, inner), - M.insert arr (identName par) arrmap) - Just par' -> ((parmap, inner . forward par par'), - arrmap) - forward to from b = - mkLet [] [to] (BasicOp $ SubExp $ Var from) - `insertStm` b - -fuseRedomap :: Bindable lore => - Names -> [VName] - -> Lambda lore -> [SubExp] -> [SubExp] -> [SOAC.Input] - -> [(VName,Ident)] - -> Lambda lore -> [SubExp] -> [SubExp] -> [SOAC.Input] - -> (Lambda lore, [SOAC.Input]) -fuseRedomap unfus_nms outVars p_lam p_scan_nes p_red_nes p_inparr outPairs - c_lam c_scan_nes c_red_nes c_inparr = - -- We hack the implementation of map o redomap to handle this case: - -- (i) we remove the accumulator formal paramter and corresponding - -- (body) result from from redomap's fold-lambda body - let p_num_nes = length p_scan_nes + length p_red_nes - unfus_arrs = filter (`nameIn` unfus_nms) outVars - p_lam_body = lambdaBody p_lam - (p_lam_scan_ts, p_lam_red_ts, p_lam_map_ts) = - splitAt3 (length p_scan_nes) (length p_red_nes) $ lambdaReturnType p_lam - (p_lam_scan_res, p_lam_red_res, p_lam_map_res) = - splitAt3 (length p_scan_nes) (length p_red_nes) $ bodyResult p_lam_body - p_lam_hacked = p_lam { lambdaParams = takeLast (length p_inparr) $ lambdaParams p_lam - , lambdaBody = p_lam_body { bodyResult = p_lam_map_res } - , lambdaReturnType = p_lam_map_ts } + where + comb ((parmap, inner), arrmap) par arr = + case M.lookup arr arrmap of + Nothing -> + ( (M.insert par arr parmap, inner), + M.insert arr (identName par) arrmap + ) + Just par' -> + ( (parmap, inner . forward par par'), + arrmap + ) + forward to from b = + mkLet [] [to] (BasicOp $ SubExp $ Var from) + `insertStm` b - -- (ii) we remove the accumulator's (global) output result from - -- @outPairs@, then ``map o redomap'' fuse the two lambdas - -- (in the usual way), and construct the extra return types - -- for the arrays that fall through. - (res_lam, new_inp) = fuseMaps (namesFromList unfus_arrs) p_lam_hacked p_inparr - (drop p_num_nes outPairs) c_lam c_inparr - (res_lam_scan_ts, res_lam_red_ts, res_lam_map_ts) = - splitAt3 (length c_scan_nes) (length c_red_nes) $ lambdaReturnType res_lam - (_,extra_map_ts) = unzip $ filter (\(nm,_)-> nm `elem` unfus_arrs) $ - zip (drop p_num_nes outVars) $ drop p_num_nes $ - lambdaReturnType p_lam +fuseRedomap :: + Bindable lore => + Names -> + [VName] -> + Lambda lore -> + [SubExp] -> + [SubExp] -> + [SOAC.Input] -> + [(VName, Ident)] -> + Lambda lore -> + [SubExp] -> + [SubExp] -> + [SOAC.Input] -> + (Lambda lore, [SOAC.Input]) +fuseRedomap + unfus_nms + outVars + p_lam + p_scan_nes + p_red_nes + p_inparr + outPairs + c_lam + c_scan_nes + c_red_nes + c_inparr = + -- We hack the implementation of map o redomap to handle this case: + -- (i) we remove the accumulator formal paramter and corresponding + -- (body) result from from redomap's fold-lambda body + let p_num_nes = length p_scan_nes + length p_red_nes + unfus_arrs = filter (`nameIn` unfus_nms) outVars + p_lam_body = lambdaBody p_lam + (p_lam_scan_ts, p_lam_red_ts, p_lam_map_ts) = + splitAt3 (length p_scan_nes) (length p_red_nes) $ lambdaReturnType p_lam + (p_lam_scan_res, p_lam_red_res, p_lam_map_res) = + splitAt3 (length p_scan_nes) (length p_red_nes) $ bodyResult p_lam_body + p_lam_hacked = + p_lam + { lambdaParams = takeLast (length p_inparr) $ lambdaParams p_lam, + lambdaBody = p_lam_body {bodyResult = p_lam_map_res}, + lambdaReturnType = p_lam_map_ts + } - -- (iii) Finally, we put back the accumulator's formal parameter and - -- (body) result in the first position of the obtained lambda. - accpars = dropLast (length p_inparr) $ lambdaParams p_lam - res_body = lambdaBody res_lam - (res_lam_scan_res, res_lam_red_res, res_lam_map_res) = - splitAt3 (length c_scan_nes) (length c_red_nes) $ bodyResult res_body - res_body'= res_body { bodyResult = p_lam_scan_res ++ res_lam_scan_res ++ - p_lam_red_res ++ res_lam_red_res ++ - res_lam_map_res } - res_lam' = res_lam { lambdaParams = accpars ++ lambdaParams res_lam - , lambdaBody = res_body' - , lambdaReturnType = p_lam_scan_ts ++ res_lam_scan_ts ++ - p_lam_red_ts ++ res_lam_red_ts ++ - res_lam_map_ts ++ extra_map_ts - } - in (res_lam', new_inp) + -- (ii) we remove the accumulator's (global) output result from + -- @outPairs@, then ``map o redomap'' fuse the two lambdas + -- (in the usual way), and construct the extra return types + -- for the arrays that fall through. + (res_lam, new_inp) = + fuseMaps + (namesFromList unfus_arrs) + p_lam_hacked + p_inparr + (drop p_num_nes outPairs) + c_lam + c_inparr + (res_lam_scan_ts, res_lam_red_ts, res_lam_map_ts) = + splitAt3 (length c_scan_nes) (length c_red_nes) $ lambdaReturnType res_lam + (_, extra_map_ts) = + unzip $ + filter (\(nm, _) -> nm `elem` unfus_arrs) $ + zip (drop p_num_nes outVars) $ + drop p_num_nes $ + lambdaReturnType p_lam + -- (iii) Finally, we put back the accumulator's formal parameter and + -- (body) result in the first position of the obtained lambda. + accpars = dropLast (length p_inparr) $ lambdaParams p_lam + res_body = lambdaBody res_lam + (res_lam_scan_res, res_lam_red_res, res_lam_map_res) = + splitAt3 (length c_scan_nes) (length c_red_nes) $ bodyResult res_body + res_body' = + res_body + { bodyResult = + p_lam_scan_res ++ res_lam_scan_res + ++ p_lam_red_res + ++ res_lam_red_res + ++ res_lam_map_res + } + res_lam' = + res_lam + { lambdaParams = accpars ++ lambdaParams res_lam, + lambdaBody = res_body', + lambdaReturnType = + p_lam_scan_ts ++ res_lam_scan_ts + ++ p_lam_red_ts + ++ res_lam_red_ts + ++ res_lam_map_ts + ++ extra_map_ts + } + in (res_lam', new_inp) mergeReduceOps :: Lambda lore -> Lambda lore -> Lambda lore mergeReduceOps (Lambda par1 bdy1 rtp1) (Lambda par2 bdy2 rtp2) = - let body' = Body (bodyDec bdy1) - (bodyStms bdy1 <> bodyStms bdy2) - (bodyResult bdy1 ++ bodyResult bdy2) + let body' = + Body + (bodyDec bdy1) + (bodyStms bdy1 <> bodyStms bdy2) + (bodyResult bdy1 ++ bodyResult bdy2) (len1, len2) = (length rtp1, length rtp2) - par' = take len1 par1 ++ take len2 par2 ++ drop len1 par1 ++ drop len2 par2 - in Lambda par' body' (rtp1++rtp2) + par' = take len1 par1 ++ take len2 par2 ++ drop len1 par1 ++ drop len2 par2 + in Lambda par' body' (rtp1 ++ rtp2) 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 @@ -1,785 +1,942 @@ {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -module Futhark.Optimise.Fusion.LoopKernel - ( FusedKer(..) - , newKernel - , inputs - , setInputs - , arrInputs - , transformOutput - , attemptFusion - , SOAC - , MapNest - ) - where - -import Control.Applicative -import Control.Arrow (first) -import Control.Monad -import Control.Monad.Reader -import Control.Monad.State -import qualified Data.Set as S -import qualified Data.Map.Strict as M -import Data.Maybe -import Data.List (find, (\\), tails) - -import Futhark.IR.SOACS hiding (SOAC(..)) -import qualified Futhark.IR.SOACS as Futhark -import Futhark.Transform.Rename (renameLambda) -import Futhark.Transform.Substitute -import Futhark.MonadFreshNames -import qualified Futhark.Analysis.HORep.SOAC as SOAC -import qualified Futhark.Analysis.HORep.MapNest as MapNest -import Futhark.Pass.ExtractKernels.ISRWIM (rwimPossible) -import Futhark.Optimise.Fusion.Composing -import Futhark.Construct -import Futhark.Util (splitAt3) - -newtype TryFusion a = TryFusion (ReaderT (Scope SOACS) - (StateT VNameSource Maybe) - a) - deriving (Functor, Applicative, Alternative, Monad, MonadFail, - MonadFreshNames, - HasScope SOACS, - LocalScope SOACS) - -tryFusion :: MonadFreshNames m => - TryFusion a -> Scope SOACS -> m (Maybe a) -tryFusion (TryFusion m) types = modifyNameSource $ \src -> - case runStateT (runReaderT m types) src of - Just (x, src') -> (Just x, src') - Nothing -> (Nothing, src) - -liftMaybe :: Maybe a -> TryFusion a -liftMaybe Nothing = fail "Nothing" -liftMaybe (Just x) = return x - -type SOAC = SOAC.SOAC SOACS -type MapNest = MapNest.MapNest SOACS - --- XXX: This function is very gross. -transformOutput :: SOAC.ArrayTransforms -> [VName] -> [Ident] - -> Binder SOACS () -transformOutput ts names = descend ts - where descend ts' validents = - case SOAC.viewf ts' of - SOAC.EmptyF -> - forM_ (zip names validents) $ \(k, valident) -> - letBindNames [k] $ BasicOp $ SubExp $ Var $ identName valident - t SOAC.:< ts'' -> do - let (es,css) = unzip $ map (applyTransform t) validents - mkPat (Ident nm tp) = Pattern [] [PatElem nm tp] - opts <- concat <$> mapM primOpType es - newIds <- forM (zip names opts) $ \(k, opt) -> - newIdent (baseString k) opt - forM_ (zip3 css newIds es) $ \(cs,ids,e) -> - certifying cs $ letBind (mkPat ids) (BasicOp e) - descend ts'' newIds - -applyTransform :: SOAC.ArrayTransform -> Ident -> (BasicOp, Certificates) -applyTransform (SOAC.Rearrange cs perm) v = - (Rearrange perm' $ identName v, cs) - where perm' = perm ++ drop (length perm) [0..arrayRank (identType v)-1] -applyTransform (SOAC.Reshape cs shape) v = - (Reshape shape $ identName v, cs) -applyTransform (SOAC.ReshapeOuter cs shape) v = - let shapes = reshapeOuter shape 1 $ arrayShape $ identType v - in (Reshape shapes $ identName v, cs) -applyTransform (SOAC.ReshapeInner cs shape) v = - let shapes = reshapeInner shape 1 $ arrayShape $ identType v - in (Reshape shapes $ identName v, cs) -applyTransform (SOAC.Replicate cs n) v = - (Replicate n $ Var $ identName v, cs) - -inputToOutput :: SOAC.Input -> Maybe (SOAC.ArrayTransform, SOAC.Input) -inputToOutput (SOAC.Input ts ia iat) = - case SOAC.viewf ts of - t SOAC.:< ts' -> Just (t, SOAC.Input ts' ia iat) - SOAC.EmptyF -> Nothing - -data FusedKer = FusedKer { - fsoac :: SOAC - -- ^ the SOAC expression, e.g., mapT( f(a,b), x, y ) - - , inplace :: Names - -- ^ Variables used in in-place updates in the kernel itself, as - -- well as on the path to the kernel from the current position. - -- This is used to avoid fusion that would violate in-place - -- restrictions. - - , fusedVars :: [VName] - -- ^ whether at least a fusion has been performed. - - , fusedConsumed :: Names - -- ^ The set of variables that were consumed by the SOACs - -- contributing to this kernel. Note that, by the type rules, the - -- final SOAC may actually consume _more_ than its original - -- contributors, which implies the need for 'Copy' expressions. - - , kernelScope :: Scope SOACS - -- ^ The names in scope at the kernel. - - , outputTransform :: SOAC.ArrayTransforms - , outNames :: [VName] - , kerAux :: StmAux () - } deriving (Show) - -newKernel :: StmAux () -> SOAC -> Names -> [VName] -> Scope SOACS -> FusedKer -newKernel aux soac consumed out_nms scope = - FusedKer { fsoac = soac - , inplace = consumed - , fusedVars = [] - , fusedConsumed = consumed - , outputTransform = SOAC.noTransforms - , outNames = out_nms - , kernelScope = scope - , kerAux = aux - } - -arrInputs :: FusedKer -> S.Set VName -arrInputs = S.fromList . map SOAC.inputArray . inputs - -inputs :: FusedKer -> [SOAC.Input] -inputs = SOAC.inputs . fsoac - -setInputs :: [SOAC.Input] -> FusedKer -> FusedKer -setInputs inps ker = ker { fsoac = inps `SOAC.setInputs` fsoac ker } - -tryOptimizeSOAC :: Names -> [VName] -> SOAC -> Names -> FusedKer - -> TryFusion FusedKer -tryOptimizeSOAC unfus_nms outVars soac consumed ker = do - (soac', ots) <- optimizeSOAC Nothing soac mempty - let ker' = map (addInitialTransformIfRelevant ots) (inputs ker) `setInputs` ker - outIdents = zipWith Ident outVars $ SOAC.typeOf soac' - ker'' = fixInputTypes outIdents ker' - applyFusionRules unfus_nms outVars soac' consumed ker'' - where addInitialTransformIfRelevant ots inp - | SOAC.inputArray inp `elem` outVars = - SOAC.addInitialTransforms ots inp - | otherwise = - inp - -tryOptimizeKernel :: Names -> [VName] -> SOAC -> Names -> FusedKer - -> TryFusion FusedKer -tryOptimizeKernel unfus_nms outVars soac consumed ker = do - ker' <- optimizeKernel (Just outVars) ker - applyFusionRules unfus_nms outVars soac consumed ker' - -tryExposeInputs :: Names -> [VName] -> SOAC -> Names -> FusedKer - -> TryFusion FusedKer -tryExposeInputs unfus_nms outVars soac consumed ker = do - (ker', ots) <- exposeInputs outVars ker - if SOAC.nullTransforms ots - then fuseSOACwithKer unfus_nms outVars soac consumed ker' - else do - (soac', ots') <- pullOutputTransforms soac ots - let outIdents = zipWith Ident outVars $ SOAC.typeOf soac' - ker'' = fixInputTypes outIdents ker' - if SOAC.nullTransforms ots' - then applyFusionRules unfus_nms outVars soac' consumed ker'' - else fail "tryExposeInputs could not pull SOAC transforms" - -fixInputTypes :: [Ident] -> FusedKer -> FusedKer -fixInputTypes outIdents ker = - ker { fsoac = fixInputTypes' $ fsoac ker } - where fixInputTypes' soac = - map fixInputType (SOAC.inputs soac) `SOAC.setInputs` soac - fixInputType (SOAC.Input ts v _) - | Just v' <- find ((==v) . identName) outIdents = - SOAC.Input ts v $ identType v' - fixInputType inp = inp - -applyFusionRules :: Names -> [VName] -> SOAC -> Names -> FusedKer - -> TryFusion FusedKer -applyFusionRules unfus_nms outVars soac consumed ker = - tryOptimizeSOAC unfus_nms outVars soac consumed ker <|> - tryOptimizeKernel unfus_nms outVars soac consumed ker <|> - fuseSOACwithKer unfus_nms outVars soac consumed ker <|> - tryExposeInputs unfus_nms outVars soac consumed ker - -attemptFusion :: MonadFreshNames m => - Names -> [VName] -> SOAC -> Names -> FusedKer - -> m (Maybe FusedKer) -attemptFusion unfus_nms outVars soac consumed ker = - fmap removeUnusedParamsFromKer <$> - tryFusion (applyFusionRules unfus_nms outVars soac consumed ker) - (kernelScope ker) - -removeUnusedParamsFromKer :: FusedKer -> FusedKer -removeUnusedParamsFromKer ker = - case soac of SOAC.Screma {} -> ker { fsoac = soac' } - _ -> ker - where soac = fsoac ker - l = SOAC.lambda soac - inps = SOAC.inputs soac - (l', inps') = removeUnusedParams l inps - soac' = l' `SOAC.setLambda` - (inps' `SOAC.setInputs` soac) - -removeUnusedParams :: Lambda -> [SOAC.Input] -> (Lambda, [SOAC.Input]) -removeUnusedParams l inps = - (l { lambdaParams = ps' }, inps') - where pInps = zip (lambdaParams l) inps - (ps', inps') = case (unzip $ filter (used . fst) pInps, pInps) of - (([], []), (p,inp):_) -> ([p], [inp]) - ((ps_, inps_), _) -> (ps_, inps_) - used p = paramName p `nameIn` freeVars - freeVars = freeIn $ lambdaBody l - --- | Check that the consumer uses at least one output of the producer --- unmodified. -mapFusionOK :: [VName] -> FusedKer -> Bool -mapFusionOK outVars ker = any (`elem` inpIds) outVars - where inpIds = mapMaybe SOAC.isVarishInput (inputs ker) - --- | Check that the consumer uses all the outputs of the producer unmodified. -mapWriteFusionOK :: [VName] -> FusedKer -> Bool -mapWriteFusionOK outVars ker = all (`elem` inpIds) outVars - where inpIds = mapMaybe SOAC.isVarishInput (inputs ker) - --- | The brain of this module: Fusing a SOAC with a Kernel. -fuseSOACwithKer :: Names -> [VName] -> SOAC -> Names -> FusedKer - -> TryFusion FusedKer -fuseSOACwithKer unfus_set outVars soac_p soac_p_consumed ker = do - -- We are fusing soac_p into soac_c, i.e, the output of soac_p is going - -- into soac_c. - let soac_c = fsoac ker - inp_p_arr = SOAC.inputs soac_p - horizFuse= unfus_set /= mempty && - SOAC.width soac_p == SOAC.width soac_c - inp_c_arr = SOAC.inputs soac_c - lam_p = SOAC.lambda soac_p - lam_c = SOAC.lambda soac_c - w = SOAC.width soac_p - returned_outvars = filter (`nameIn` unfus_set) outVars - success res_outnms res_soac = do - let fusedVars_new = fusedVars ker++outVars - -- Avoid name duplication, because the producer lambda is not - -- removed from the program until much later. - uniq_lam <- renameLambda $ SOAC.lambda res_soac - return $ ker { fsoac = uniq_lam `SOAC.setLambda` res_soac - , fusedVars = fusedVars_new - , inplace = inplace ker <> soac_p_consumed - , fusedConsumed = fusedConsumed ker <> soac_p_consumed - , outNames = res_outnms - } - - outPairs <- forM (zip outVars $ map rowType $ SOAC.typeOf soac_p) $ \(outVar, t) -> do - outVar' <- newVName $ baseString outVar ++ "_elem" - return (outVar, Ident outVar' t) - - let mapLikeFusionCheck = - let (res_lam, new_inp) = fuseMaps unfus_set lam_p inp_p_arr outPairs lam_c inp_c_arr - (extra_nms,extra_rtps) = unzip $ filter ((`nameIn` unfus_set) . fst) $ - zip outVars $ map (stripArray 1) $ SOAC.typeOf soac_p - res_lam' = res_lam { lambdaReturnType = lambdaReturnType res_lam ++ extra_rtps } - in (extra_nms, res_lam', new_inp) - - when (horizFuse && not (SOAC.nullTransforms $ outputTransform ker)) $ - fail "Horizontal fusion is invalid in the presence of output transforms." - - case (soac_c, soac_p) of - _ | SOAC.width soac_p /= SOAC.width soac_c -> fail "SOAC widths must match." - - (SOAC.Screma _ (ScremaForm scans_c reds_c _) _, - SOAC.Screma _ (ScremaForm scans_p reds_p _) _) - | mapFusionOK (drop (Futhark.scanResults scans_p+Futhark.redResults reds_p) outVars) ker - || horizFuse -> do - let red_nes_p = concatMap redNeutral reds_p - red_nes_c = concatMap redNeutral reds_c - scan_nes_p = concatMap scanNeutral scans_p - scan_nes_c = concatMap scanNeutral scans_c - (res_lam', new_inp) = fuseRedomap unfus_set outVars - lam_p scan_nes_p red_nes_p inp_p_arr - outPairs - lam_c scan_nes_c red_nes_c inp_c_arr - (soac_p_scanout, soac_p_redout, _soac_p_mapout) = - splitAt3 (length scan_nes_p) (length red_nes_p) outVars - (soac_c_scanout, soac_c_redout, soac_c_mapout) = - splitAt3 (length scan_nes_c) (length red_nes_c) $ outNames ker - unfus_arrs = returned_outvars \\ (soac_p_scanout++soac_p_redout) - success (soac_p_scanout ++ soac_c_scanout ++ - soac_p_redout ++ soac_c_redout ++ - soac_c_mapout ++ unfus_arrs) $ - SOAC.Screma w (ScremaForm (scans_p ++ scans_c) (reds_p ++ reds_c) res_lam') - new_inp - - ------------------ - -- Scatter fusion -- - ------------------ - - -- Map-Scatter fusion. - -- - -- The 'inplace' mechanism for kernels already takes care of - -- checking that the Scatter is not writing to any array used in - -- the Map. - (SOAC.Scatter _len _lam _ivs dests, - SOAC.Screma _ form _) - | isJust $ isMapSOAC form, - -- 1. all arrays produced by the map are ONLY used (consumed) - -- by the scatter, i.e., not used elsewhere. - not (any (`nameIn` unfus_set) outVars), - -- 2. all arrays produced by the map are input to the scatter. - mapWriteFusionOK outVars ker -> do - let (extra_nms, res_lam', new_inp) = mapLikeFusionCheck - success (outNames ker ++ extra_nms) $ - SOAC.Scatter w res_lam' new_inp dests - - -- Map-Hist fusion. - -- - -- The 'inplace' mechanism for kernels already takes care of - -- checking that the Hist is not writing to any array used in - -- the Map. - (SOAC.Hist _ ops _ _, - SOAC.Screma _ form _) - | isJust $ isMapSOAC form, - -- 1. all arrays produced by the map are ONLY used (consumed) - -- by the hist, i.e., not used elsewhere. - not (any (`nameIn` unfus_set) outVars), - -- 2. all arrays produced by the map are input to the scatter. - mapWriteFusionOK outVars ker -> do - let (extra_nms, res_lam', new_inp) = mapLikeFusionCheck - success (outNames ker ++ extra_nms) $ - SOAC.Hist w ops res_lam' new_inp - - -- Hist-Hist fusion - (SOAC.Hist _ ops_c _ _, - SOAC.Hist _ ops_p _ _) - | horizFuse -> do - let p_num_buckets = length ops_p - c_num_buckets = length ops_c - (body_p, body_c) = (lambdaBody lam_p, lambdaBody lam_c) - body' = - Body { bodyDec = bodyDec body_p -- body_p and body_c have the same lores - , bodyStms = bodyStms body_p <> bodyStms body_c - , bodyResult = take c_num_buckets (bodyResult body_c) ++ - take p_num_buckets (bodyResult body_p) ++ - drop c_num_buckets (bodyResult body_c) ++ - drop p_num_buckets (bodyResult body_p) - } - lam' = - Lambda { lambdaParams = lambdaParams lam_c ++ lambdaParams lam_p - , lambdaBody = body' - , lambdaReturnType = replicate (c_num_buckets+p_num_buckets) (Prim int32) ++ - drop c_num_buckets (lambdaReturnType lam_c) ++ - drop p_num_buckets (lambdaReturnType lam_p) - } - success (outNames ker ++ returned_outvars) $ - SOAC.Hist w (ops_c <> ops_p) lam' (inp_c_arr <> inp_p_arr) - - -- Scatter-write fusion. - (SOAC.Scatter _len2 _lam_c ivs2 as2, - SOAC.Scatter _len_p _lam_p ivs_p as_p) - | horizFuse -> do - let zipW xs ys = ys_p ++ xs_p ++ ys2 ++ xs2 - where lenx = length xs `div` 2 - xs_p = take lenx xs - xs2 = drop lenx xs - leny = length ys `div` 2 - ys_p = take leny ys - ys2 = drop leny ys - let (body_p, body2) = (lambdaBody lam_p, lambdaBody lam_c) - let body' = Body { bodyDec = bodyDec body_p -- body_p and body2 have the same lores - , bodyStms = bodyStms body_p <> bodyStms body2 - , bodyResult = zipW (bodyResult body_p) (bodyResult body2) - } - let lam' = Lambda { lambdaParams = lambdaParams lam_p ++ lambdaParams lam_c - , lambdaBody = body' - , lambdaReturnType = zipW (lambdaReturnType lam_p) (lambdaReturnType lam_c) - } - success (outNames ker ++ returned_outvars) $ - SOAC.Scatter w lam' (ivs_p ++ ivs2) (as2 ++ as_p) - - (SOAC.Scatter {}, _) -> - fail "Cannot fuse a write with anything else than a write or a map" - (_, SOAC.Scatter {}) -> - fail "Cannot fuse a write with anything else than a write or a map" - - ---------------------------- - -- Stream-Stream Fusions: -- - ---------------------------- - (SOAC.Stream _ Sequential{} _ _, SOAC.Stream _ form_p@Sequential{} _ _) - | mapFusionOK (drop (length $ getStreamAccums form_p) 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{} _ _) -> - fail "Fusion conditions not met for two SEQ streams!" - - (SOAC.Stream _ Sequential{} _ _, SOAC.Stream{}) -> - fail "Cannot fuse a parallel with a sequential Stream!" - - (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 - -- fuse two PARALLEL streams - (res_nms, res_stream) <- fuseStreamHelper (outNames ker) unfus_set outVars outPairs soac_c soac_p - success res_nms res_stream - - (SOAC.Stream{}, SOAC.Stream {}) -> - fail "Fusion conditions not met for two PAR streams!" - - ------------------------------------------------------------------- - --- If one is a stream, translate the other to a stream as well.--- - --- This does not get in trouble (infinite computation) because --- - --- scan's translation to Stream introduces a hindrance to --- - --- (horizontal fusion), hence repeated application is for the--- - --- moment impossible. However, if with a dependence-graph rep--- - --- we could run in an infinite recursion, i.e., repeatedly --- - --- fusing map o scan into an infinity of Stream levels! --- - ------------------------------------------------------------------- - (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, - -- and recursively perform stream-stream fusion. - (soac_p', newacc_ids) <- SOAC.soacToStream soac_p - soac_p'' <- case form2 of - Sequential{} -> toSeqStream soac_p' - _ -> return soac_p' - if soac_p' == soac_p - then fail "SOAC could not be turned into stream." - else fuseSOACwithKer unfus_set (map identName newacc_ids++outVars) soac_p'' soac_p_consumed ker - - (_, SOAC.Screma _ form _) | Just _ <- Futhark.isScanSOAC form -> do - -- A Scan soac can be currently only fused as a (sequential) stream, - -- hence it is first translated to a (sequential) Stream and then - -- fusion with a kernel is attempted. - (soac_p', newacc_ids) <- SOAC.soacToStream soac_p - 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 - -- 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. - -- The newly created stream corresponding to soac_c borrows the - -- sequential/parallel property of the soac_p stream. - (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' - _ -> return soac_c' - - fuseSOACwithKer unfus_set outVars soac_p soac_p_consumed $ - ker { fsoac = soac_c'', outNames = map identName newacc_ids ++ outNames ker } - - --------------------------------- - --- DEFAULT, CANNOT FUSE CASE --- - --------------------------------- - _ -> fail "Cannot fuse" - -getStreamOrder :: StreamForm lore -> StreamOrd -getStreamOrder (Parallel o _ _ _) = o -getStreamOrder (Sequential _) = InOrder - -fuseStreamHelper :: [VName] -> Names -> [VName] -> [(VName,Ident)] - -> SOAC -> SOAC -> TryFusion ([VName], SOAC) -fuseStreamHelper out_kernms unfus_set outVars outPairs - (SOAC.Stream w2 form2 lam2 inp2_arr) - (SOAC.Stream _ form1 lam1 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 - chunk2 = head $ lambdaParams lam2 - hmnms = M.fromList [(paramName chunk2, paramName chunk1)] - lam20 = substituteNames hmnms lam2 - lam1' = lam1 { lambdaParams = tail $ lambdaParams lam1 } - lam2' = lam20 { lambdaParams = tail $ lambdaParams lam20 } - (res_lam', new_inp) = fuseRedomap unfus_set outVars - lam1' [] nes1 - inp1_arr outPairs - lam2' [] (getStreamAccums form2) - inp2_arr - res_lam'' = res_lam' { lambdaParams = chunk1 : lambdaParams res_lam' } - unfus_accs = take (length nes1) outVars - unfus_arrs = filter (`nameIn` unfus_set) outVars - res_form <- mergeForms form2 form1 - return (unfus_accs ++ out_kernms ++ unfus_arrs, - SOAC.Stream w2 res_form res_lam'' 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 _ _ = 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 _ = fail "toSeqStream expects a stream, but given a SOAC." - --- Here follows optimizations and transforms to expose fusability. - -optimizeKernel :: Maybe [VName] -> FusedKer -> TryFusion FusedKer -optimizeKernel inp ker = do - (soac, resTrans) <- optimizeSOAC inp (fsoac ker) startTrans - return $ ker { fsoac = soac - , outputTransform = resTrans - } - where startTrans = outputTransform ker - -optimizeSOAC :: Maybe [VName] -> SOAC -> SOAC.ArrayTransforms - -> TryFusion (SOAC, SOAC.ArrayTransforms) -optimizeSOAC inp soac os = do - res <- foldM comb (False, soac, os) optimizations - case res of - (False, _, _) -> fail "No optimisation applied" - (True, soac', os') -> return (soac', os') - where comb (changed, soac', os') f = do - (soac'', os'') <- f inp soac' os - return (True, soac'', os'') - <|> return (changed, soac', os') - -type Optimization = Maybe [VName] - -> SOAC - -> SOAC.ArrayTransforms - -> TryFusion (SOAC, SOAC.ArrayTransforms) - -optimizations :: [Optimization] -optimizations = [iswim] - -iswim :: Maybe [VName] -> SOAC -> SOAC.ArrayTransforms - -> TryFusion (SOAC, SOAC.ArrayTransforms) -iswim _ (SOAC.Screma w form arrs) ots - | Just [Futhark.Scan scan_fun nes] <- Futhark.isScanSOAC form, - Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible scan_fun, - Just nes_names <- mapM subExpVar nes = do - - let nes_idents = zipWith Ident nes_names $ lambdaReturnType scan_fun - map_nes = map SOAC.identInput nes_idents - map_arrs' = map_nes ++ map (SOAC.transposeInput 0 1) arrs - (scan_acc_params, scan_elem_params) = - splitAt (length arrs) $ lambdaParams scan_fun - map_params = map removeParamOuterDim scan_acc_params ++ - map (setParamOuterDimTo w) scan_elem_params - map_rettype = map (`setOuterSize` w) $ lambdaReturnType scan_fun - - scan_params = lambdaParams map_fun - scan_body = lambdaBody map_fun - scan_rettype = lambdaReturnType map_fun - scan_fun' = Lambda scan_params scan_body scan_rettype - nes' = map Var $ take (length map_nes) $ map paramName map_params - arrs' = drop (length map_nes) $ map paramName map_params - - scan_form <- scanSOAC [Futhark.Scan scan_fun' nes'] - - let map_body = mkBody (oneStm $ - Let (setPatternOuterDimTo w map_pat) (defAux ()) $ - Op $ Futhark.Screma w scan_form arrs') $ - map Var $ patternNames map_pat - map_fun' = Lambda map_params map_body map_rettype - perm = case lambdaReturnType map_fun of - [] -> [] - t:_ -> 1 : 0 : [2..arrayRank t] - - return (SOAC.Screma map_w (ScremaForm [] [] map_fun') map_arrs', - ots SOAC.|> SOAC.Rearrange map_cs perm) - -iswim _ _ _ = - fail "ISWIM does not apply." - -removeParamOuterDim :: LParam -> LParam -removeParamOuterDim param = - let t = rowType $ paramType param - in param { paramDec = t } - -setParamOuterDimTo :: SubExp -> LParam -> LParam -setParamOuterDimTo w param = - let t = paramType param `setOuterSize` w - in param { paramDec = t } - -setPatternOuterDimTo :: SubExp -> Pattern -> Pattern -setPatternOuterDimTo w = fmap (`setOuterSize` w) - --- Now for fiddling with transpositions... - -commonTransforms :: [VName] -> [SOAC.Input] - -> (SOAC.ArrayTransforms, [SOAC.Input]) -commonTransforms interesting inps = commonTransforms' inps' - where inps' = [ (SOAC.inputArray inp `elem` interesting, inp) - | inp <- inps ] - -commonTransforms' :: [(Bool, SOAC.Input)] -> (SOAC.ArrayTransforms, [SOAC.Input]) -commonTransforms' inps = - case foldM inspect (Nothing, []) inps of - Just (Just mot, inps') -> first (mot SOAC.<|) $ commonTransforms' $ reverse inps' - _ -> (SOAC.noTransforms, map snd inps) - where inspect (mot, prev) (True, inp) = - case (mot, inputToOutput inp) of - (Nothing, Just (ot, inp')) -> Just (Just ot, (True, inp') : prev) - (Just ot1, Just (ot2, inp')) - | ot1 == ot2 -> Just (Just ot2, (True, inp') : prev) - _ -> Nothing - inspect (mot, prev) inp = Just (mot,inp:prev) - -mapDepth :: MapNest -> Int -mapDepth (MapNest.MapNest _ lam levels _) = - min resDims (length levels) + 1 - where resDims = minDim $ case levels of - [] -> lambdaReturnType lam - nest:_ -> MapNest.nestingReturnType nest - minDim [] = 0 - minDim (t:ts) = foldl min (arrayRank t) $ map arrayRank ts - -pullRearrange :: SOAC -> SOAC.ArrayTransforms - -> TryFusion (SOAC, SOAC.ArrayTransforms) -pullRearrange soac ots = do - nest <- liftMaybe =<< MapNest.fromSOAC soac - SOAC.Rearrange cs perm SOAC.:< ots' <- return $ SOAC.viewf ots - if rearrangeReach perm <= mapDepth nest then do - let -- Expand perm to cover the full extent of the input dimensionality - perm' inp = take r perm ++ [length perm..r-1] - where r = SOAC.inputRank inp - addPerm inp = SOAC.addTransform (SOAC.Rearrange cs $ perm' inp) inp - inputs' = map addPerm $ MapNest.inputs nest - soac' <- MapNest.toSOAC $ - inputs' `MapNest.setInputs` rearrangeReturnTypes nest perm - return (soac', ots') - else fail "Cannot pull transpose" - -pushRearrange :: [VName] -> SOAC -> SOAC.ArrayTransforms - -> TryFusion (SOAC, SOAC.ArrayTransforms) -pushRearrange inpIds soac ots = do - nest <- liftMaybe =<< MapNest.fromSOAC soac - (perm, inputs') <- liftMaybe $ fixupInputs inpIds $ MapNest.inputs nest - if rearrangeReach perm <= mapDepth nest then do - let invertRearrange = SOAC.Rearrange mempty $ rearrangeInverse perm - soac' <- MapNest.toSOAC $ - inputs' `MapNest.setInputs` - rearrangeReturnTypes nest perm - return (soac', invertRearrange SOAC.<| ots) - else fail "Cannot push transpose" - --- | Actually also rearranges indices. -rearrangeReturnTypes :: MapNest -> [Int] -> MapNest -rearrangeReturnTypes nest@(MapNest.MapNest w body nestings inps) perm = - MapNest.MapNest w - body - (zipWith setReturnType - nestings $ - drop 1 $ iterate (map rowType) ts) - inps - where origts = MapNest.typeOf nest - -- The permutation may be deeper than the rank of the type, - -- but it is required that it is an identity permutation - -- beyond that. This is supposed to be checked as an - -- invariant by whoever calls rearrangeReturnTypes. - rearrangeType' t = rearrangeType (take (arrayRank t) perm) t - ts = map rearrangeType' origts - - setReturnType nesting t' = - nesting { MapNest.nestingReturnType = t' } - -fixupInputs :: [VName] -> [SOAC.Input] -> Maybe ([Int], [SOAC.Input]) -fixupInputs inpIds inps = - case mapMaybe inputRearrange $ filter exposable inps of - perm:_ -> do inps' <- mapM (fixupInput (rearrangeReach perm) perm) inps - return (perm, inps') - _ -> Nothing - where exposable = (`elem` inpIds) . SOAC.inputArray - - inputRearrange (SOAC.Input ts _ _) - | _ SOAC.:> SOAC.Rearrange _ perm <- SOAC.viewl ts = Just perm - inputRearrange _ = Nothing - - fixupInput d perm inp - | r <- SOAC.inputRank inp, - r >= d = - Just $ SOAC.addTransform (SOAC.Rearrange mempty $ take r perm) inp - | otherwise = Nothing - -pullReshape :: SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms) -pullReshape (SOAC.Screma _ form inps) ots - | Just maplam <- Futhark.isMapSOAC form, - SOAC.Reshape cs shape SOAC.:< ots' <- SOAC.viewf ots, - all primType $ lambdaReturnType maplam = do - let mapw' = case reverse $ newDims shape of - [] -> intConst Int32 0 - d:_ -> d - inputs' = map (SOAC.addTransform $ SOAC.ReshapeOuter cs shape) inps - inputTypes = map SOAC.inputType inputs' - - let outersoac :: ([SOAC.Input] -> SOAC) -> (SubExp, [SubExp]) - -> TryFusion ([SOAC.Input] -> SOAC) - outersoac inner (w, outershape) = do - let addDims t = arrayOf t (Shape outershape) NoUniqueness - retTypes = map addDims $ lambdaReturnType maplam - - ps <- forM inputTypes $ \inpt -> - newParam "pullReshape_param" $ - stripArray (length shape-length outershape) inpt - - inner_body <- runBodyBinder $ - eBody [SOAC.toExp $ inner $ map (SOAC.identInput . paramIdent) ps] - let inner_fun = Lambda { lambdaParams = ps - , lambdaReturnType = retTypes - , lambdaBody = inner_body - } - return $ SOAC.Screma w $ Futhark.mapSOAC inner_fun - - op' <- foldM outersoac (SOAC.Screma mapw' $ Futhark.mapSOAC maplam) $ - zip (drop 1 $ reverse $ newDims shape) $ - drop 1 $ reverse $ drop 1 $ tails $ newDims shape - return (op' inputs', ots') -pullReshape _ _ = fail "Cannot pull reshape" - --- Tie it all together in exposeInputs (for making inputs to a --- consumer available) and pullOutputTransforms (for moving --- output-transforms of a producer to its inputs instead). - -exposeInputs :: [VName] -> FusedKer - -> TryFusion (FusedKer, SOAC.ArrayTransforms) -exposeInputs inpIds ker = - (exposeInputs' =<< pushRearrange') <|> - (exposeInputs' =<< pullRearrange') <|> - exposeInputs' ker - where ot = outputTransform ker - - pushRearrange' = do - (soac', ot') <- pushRearrange inpIds (fsoac ker) ot - return ker { fsoac = soac' - , outputTransform = ot' - } - - pullRearrange' = do - (soac',ot') <- pullRearrange (fsoac ker) ot - unless (SOAC.nullTransforms ot') $ - fail "pullRearrange was not enough" - return ker { fsoac = soac' - , outputTransform = SOAC.noTransforms - } - - exposeInputs' ker' = - case commonTransforms inpIds $ inputs ker' of - (ot', inps') | all exposed inps' -> - return (ker' { fsoac = inps' `SOAC.setInputs` fsoac ker'}, ot') - _ -> fail "Cannot expose" - - exposed (SOAC.Input ts _ _) - | SOAC.nullTransforms ts = True - exposed inp = SOAC.inputArray inp `notElem` inpIds - -outputTransformPullers :: [SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms)] -outputTransformPullers = [pullRearrange, pullReshape] - -pullOutputTransforms :: SOAC -> SOAC.ArrayTransforms - -> TryFusion (SOAC, SOAC.ArrayTransforms) -pullOutputTransforms = attempt outputTransformPullers - where attempt [] _ _ = fail "Cannot pull anything" - attempt (p:ps) soac ots = do - (soac',ots') <- p soac ots - if SOAC.nullTransforms ots' then return (soac', SOAC.noTransforms) - else pullOutputTransforms soac' ots' <|> return (soac', ots') - <|> attempt ps soac ots +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} + +module Futhark.Optimise.Fusion.LoopKernel + ( FusedKer (..), + newKernel, + inputs, + setInputs, + arrInputs, + transformOutput, + attemptFusion, + SOAC, + MapNest, + ) +where + +import Control.Applicative +import Control.Arrow (first) +import Control.Monad +import Control.Monad.Reader +import Control.Monad.State +import Data.List (find, tails, (\\)) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import qualified Futhark.Analysis.HORep.MapNest as MapNest +import qualified Futhark.Analysis.HORep.SOAC as SOAC +import Futhark.Construct +import Futhark.IR.SOACS hiding (SOAC (..)) +import qualified Futhark.IR.SOACS as Futhark +import Futhark.Optimise.Fusion.Composing +import Futhark.Pass.ExtractKernels.ISRWIM (rwimPossible) +import Futhark.Transform.Rename (renameLambda) +import Futhark.Transform.Substitute +import Futhark.Util (splitAt3) + +newtype TryFusion a + = TryFusion + ( ReaderT + (Scope SOACS) + (StateT VNameSource Maybe) + a + ) + deriving + ( Functor, + Applicative, + Alternative, + Monad, + MonadFail, + MonadFreshNames, + HasScope SOACS, + LocalScope SOACS + ) + +tryFusion :: + MonadFreshNames m => + TryFusion a -> + Scope SOACS -> + m (Maybe a) +tryFusion (TryFusion m) types = modifyNameSource $ \src -> + case runStateT (runReaderT m types) src of + Just (x, src') -> (Just x, src') + Nothing -> (Nothing, src) + +liftMaybe :: Maybe a -> TryFusion a +liftMaybe Nothing = fail "Nothing" +liftMaybe (Just x) = return x + +type SOAC = SOAC.SOAC SOACS + +type MapNest = MapNest.MapNest SOACS + +-- XXX: This function is very gross. +transformOutput :: + SOAC.ArrayTransforms -> + [VName] -> + [Ident] -> + Binder SOACS () +transformOutput ts names = descend ts + where + descend ts' validents = + case SOAC.viewf ts' of + SOAC.EmptyF -> + forM_ (zip names validents) $ \(k, valident) -> + letBindNames [k] $ BasicOp $ SubExp $ Var $ identName valident + t SOAC.:< ts'' -> do + let (es, css) = unzip $ map (applyTransform t) validents + mkPat (Ident nm tp) = Pattern [] [PatElem nm tp] + opts <- concat <$> mapM primOpType es + newIds <- forM (zip names opts) $ \(k, opt) -> + newIdent (baseString k) opt + forM_ (zip3 css newIds es) $ \(cs, ids, e) -> + certifying cs $ letBind (mkPat ids) (BasicOp e) + descend ts'' newIds + +applyTransform :: SOAC.ArrayTransform -> Ident -> (BasicOp, Certificates) +applyTransform (SOAC.Rearrange cs perm) v = + (Rearrange perm' $ identName v, cs) + where + perm' = perm ++ drop (length perm) [0 .. arrayRank (identType v) -1] +applyTransform (SOAC.Reshape cs shape) v = + (Reshape shape $ identName v, cs) +applyTransform (SOAC.ReshapeOuter cs shape) v = + let shapes = reshapeOuter shape 1 $ arrayShape $ identType v + in (Reshape shapes $ identName v, cs) +applyTransform (SOAC.ReshapeInner cs shape) v = + let shapes = reshapeInner shape 1 $ arrayShape $ identType v + in (Reshape shapes $ identName v, cs) +applyTransform (SOAC.Replicate cs n) v = + (Replicate n $ Var $ identName v, cs) + +inputToOutput :: SOAC.Input -> Maybe (SOAC.ArrayTransform, SOAC.Input) +inputToOutput (SOAC.Input ts ia iat) = + case SOAC.viewf ts of + t SOAC.:< ts' -> Just (t, SOAC.Input ts' ia iat) + SOAC.EmptyF -> Nothing + +data FusedKer = FusedKer + { -- | the SOAC expression, e.g., mapT( f(a,b), x, y ) + fsoac :: SOAC, + -- | Variables used in in-place updates in the kernel itself, as + -- well as on the path to the kernel from the current position. + -- This is used to avoid fusion that would violate in-place + -- restrictions. + inplace :: Names, + -- | whether at least a fusion has been performed. + fusedVars :: [VName], + -- | The set of variables that were consumed by the SOACs + -- contributing to this kernel. Note that, by the type rules, the + -- final SOAC may actually consume _more_ than its original + -- contributors, which implies the need for 'Copy' expressions. + fusedConsumed :: Names, + -- | The names in scope at the kernel. + kernelScope :: Scope SOACS, + outputTransform :: SOAC.ArrayTransforms, + outNames :: [VName], + kerAux :: StmAux () + } + deriving (Show) + +newKernel :: StmAux () -> SOAC -> Names -> [VName] -> Scope SOACS -> FusedKer +newKernel aux soac consumed out_nms scope = + FusedKer + { fsoac = soac, + inplace = consumed, + fusedVars = [], + fusedConsumed = consumed, + outputTransform = SOAC.noTransforms, + outNames = out_nms, + kernelScope = scope, + kerAux = aux + } + +arrInputs :: FusedKer -> S.Set VName +arrInputs = S.fromList . map SOAC.inputArray . inputs + +inputs :: FusedKer -> [SOAC.Input] +inputs = SOAC.inputs . fsoac + +setInputs :: [SOAC.Input] -> FusedKer -> FusedKer +setInputs inps ker = ker {fsoac = inps `SOAC.setInputs` fsoac ker} + +tryOptimizeSOAC :: + Names -> + [VName] -> + SOAC -> + Names -> + FusedKer -> + TryFusion FusedKer +tryOptimizeSOAC unfus_nms outVars soac consumed ker = do + (soac', ots) <- optimizeSOAC Nothing soac mempty + let ker' = map (addInitialTransformIfRelevant ots) (inputs ker) `setInputs` ker + outIdents = zipWith Ident outVars $ SOAC.typeOf soac' + ker'' = fixInputTypes outIdents ker' + applyFusionRules unfus_nms outVars soac' consumed ker'' + where + addInitialTransformIfRelevant ots inp + | SOAC.inputArray inp `elem` outVars = + SOAC.addInitialTransforms ots inp + | otherwise = + inp + +tryOptimizeKernel :: + Names -> + [VName] -> + SOAC -> + Names -> + FusedKer -> + TryFusion FusedKer +tryOptimizeKernel unfus_nms outVars soac consumed ker = do + ker' <- optimizeKernel (Just outVars) ker + applyFusionRules unfus_nms outVars soac consumed ker' + +tryExposeInputs :: + Names -> + [VName] -> + SOAC -> + Names -> + FusedKer -> + TryFusion FusedKer +tryExposeInputs unfus_nms outVars soac consumed ker = do + (ker', ots) <- exposeInputs outVars ker + if SOAC.nullTransforms ots + then fuseSOACwithKer unfus_nms outVars soac consumed ker' + else do + (soac', ots') <- pullOutputTransforms soac ots + let outIdents = zipWith Ident outVars $ SOAC.typeOf soac' + ker'' = fixInputTypes outIdents ker' + if SOAC.nullTransforms ots' + then applyFusionRules unfus_nms outVars soac' consumed ker'' + else fail "tryExposeInputs could not pull SOAC transforms" + +fixInputTypes :: [Ident] -> FusedKer -> FusedKer +fixInputTypes outIdents ker = + ker {fsoac = fixInputTypes' $ fsoac ker} + where + fixInputTypes' soac = + map fixInputType (SOAC.inputs soac) `SOAC.setInputs` soac + fixInputType (SOAC.Input ts v _) + | Just v' <- find ((== v) . identName) outIdents = + SOAC.Input ts v $ identType v' + fixInputType inp = inp + +applyFusionRules :: + Names -> + [VName] -> + SOAC -> + Names -> + FusedKer -> + TryFusion FusedKer +applyFusionRules unfus_nms outVars soac consumed ker = + tryOptimizeSOAC unfus_nms outVars soac consumed ker + <|> tryOptimizeKernel unfus_nms outVars soac consumed ker + <|> fuseSOACwithKer unfus_nms outVars soac consumed ker + <|> tryExposeInputs unfus_nms outVars soac consumed ker + +attemptFusion :: + MonadFreshNames m => + Names -> + [VName] -> + SOAC -> + Names -> + FusedKer -> + m (Maybe FusedKer) +attemptFusion unfus_nms outVars soac consumed ker = + fmap removeUnusedParamsFromKer + <$> tryFusion + (applyFusionRules unfus_nms outVars soac consumed ker) + (kernelScope ker) + +removeUnusedParamsFromKer :: FusedKer -> FusedKer +removeUnusedParamsFromKer ker = + case soac of + SOAC.Screma {} -> ker {fsoac = soac'} + _ -> ker + where + soac = fsoac ker + l = SOAC.lambda soac + inps = SOAC.inputs soac + (l', inps') = removeUnusedParams l inps + soac' = + l' + `SOAC.setLambda` (inps' `SOAC.setInputs` soac) + +removeUnusedParams :: Lambda -> [SOAC.Input] -> (Lambda, [SOAC.Input]) +removeUnusedParams l inps = + (l {lambdaParams = ps'}, inps') + where + pInps = zip (lambdaParams l) inps + (ps', inps') = case (unzip $ filter (used . fst) pInps, pInps) of + (([], []), (p, inp) : _) -> ([p], [inp]) + ((ps_, inps_), _) -> (ps_, inps_) + used p = paramName p `nameIn` freeVars + freeVars = freeIn $ lambdaBody l + +-- | Check that the consumer uses at least one output of the producer +-- unmodified. +mapFusionOK :: [VName] -> FusedKer -> Bool +mapFusionOK outVars ker = any (`elem` inpIds) outVars + where + inpIds = mapMaybe SOAC.isVarishInput (inputs ker) + +-- | Check that the consumer uses all the outputs of the producer unmodified. +mapWriteFusionOK :: [VName] -> FusedKer -> Bool +mapWriteFusionOK outVars ker = all (`elem` inpIds) outVars + where + inpIds = mapMaybe SOAC.isVarishInput (inputs ker) + +-- | The brain of this module: Fusing a SOAC with a Kernel. +fuseSOACwithKer :: + Names -> + [VName] -> + SOAC -> + Names -> + FusedKer -> + TryFusion FusedKer +fuseSOACwithKer unfus_set outVars soac_p soac_p_consumed ker = do + -- We are fusing soac_p into soac_c, i.e, the output of soac_p is going + -- into soac_c. + let soac_c = fsoac ker + inp_p_arr = SOAC.inputs soac_p + horizFuse = + unfus_set /= mempty + && SOAC.width soac_p == SOAC.width soac_c + inp_c_arr = SOAC.inputs soac_c + lam_p = SOAC.lambda soac_p + lam_c = SOAC.lambda soac_c + w = SOAC.width soac_p + returned_outvars = filter (`nameIn` unfus_set) outVars + success res_outnms res_soac = do + let fusedVars_new = fusedVars ker ++ outVars + -- Avoid name duplication, because the producer lambda is not + -- removed from the program until much later. + uniq_lam <- renameLambda $ SOAC.lambda res_soac + return $ + ker + { fsoac = uniq_lam `SOAC.setLambda` res_soac, + fusedVars = fusedVars_new, + inplace = inplace ker <> soac_p_consumed, + fusedConsumed = fusedConsumed ker <> soac_p_consumed, + outNames = res_outnms + } + + outPairs <- forM (zip outVars $ map rowType $ SOAC.typeOf soac_p) $ \(outVar, t) -> do + outVar' <- newVName $ baseString outVar ++ "_elem" + return (outVar, Ident outVar' t) + + let mapLikeFusionCheck = + let (res_lam, new_inp) = fuseMaps unfus_set lam_p inp_p_arr outPairs lam_c inp_c_arr + (extra_nms, extra_rtps) = + unzip $ + filter ((`nameIn` unfus_set) . fst) $ + zip outVars $ map (stripArray 1) $ SOAC.typeOf soac_p + res_lam' = res_lam {lambdaReturnType = lambdaReturnType res_lam ++ extra_rtps} + in (extra_nms, res_lam', new_inp) + + when (horizFuse && not (SOAC.nullTransforms $ outputTransform ker)) $ + fail "Horizontal fusion is invalid in the presence of output transforms." + + case (soac_c, soac_p) of + _ | SOAC.width soac_p /= SOAC.width soac_c -> fail "SOAC widths must match." + ( SOAC.Screma _ (ScremaForm scans_c reds_c _) _, + SOAC.Screma _ (ScremaForm scans_p reds_p _) _ + ) + | mapFusionOK (drop (Futhark.scanResults scans_p + Futhark.redResults reds_p) outVars) ker + || horizFuse -> do + let red_nes_p = concatMap redNeutral reds_p + red_nes_c = concatMap redNeutral reds_c + scan_nes_p = concatMap scanNeutral scans_p + scan_nes_c = concatMap scanNeutral scans_c + (res_lam', new_inp) = + fuseRedomap + unfus_set + outVars + lam_p + scan_nes_p + red_nes_p + inp_p_arr + outPairs + lam_c + scan_nes_c + red_nes_c + inp_c_arr + (soac_p_scanout, soac_p_redout, _soac_p_mapout) = + splitAt3 (length scan_nes_p) (length red_nes_p) outVars + (soac_c_scanout, soac_c_redout, soac_c_mapout) = + splitAt3 (length scan_nes_c) (length red_nes_c) $ outNames ker + unfus_arrs = returned_outvars \\ (soac_p_scanout ++ soac_p_redout) + success + ( soac_p_scanout ++ soac_c_scanout + ++ soac_p_redout + ++ soac_c_redout + ++ soac_c_mapout + ++ unfus_arrs + ) + $ SOAC.Screma + w + (ScremaForm (scans_p ++ scans_c) (reds_p ++ reds_c) res_lam') + new_inp + + ------------------ + -- Scatter fusion -- + ------------------ + + -- Map-Scatter fusion. + -- + -- The 'inplace' mechanism for kernels already takes care of + -- checking that the Scatter is not writing to any array used in + -- the Map. + ( SOAC.Scatter _len _lam _ivs dests, + SOAC.Screma _ form _ + ) + | isJust $ isMapSOAC form, + -- 1. all arrays produced by the map are ONLY used (consumed) + -- by the scatter, i.e., not used elsewhere. + not (any (`nameIn` unfus_set) outVars), + -- 2. all arrays produced by the map are input to the scatter. + mapWriteFusionOK outVars ker -> do + let (extra_nms, res_lam', new_inp) = mapLikeFusionCheck + success (outNames ker ++ extra_nms) $ + SOAC.Scatter w res_lam' new_inp dests + + -- Map-Hist fusion. + -- + -- The 'inplace' mechanism for kernels already takes care of + -- checking that the Hist is not writing to any array used in + -- the Map. + ( SOAC.Hist _ ops _ _, + SOAC.Screma _ form _ + ) + | isJust $ isMapSOAC form, + -- 1. all arrays produced by the map are ONLY used (consumed) + -- by the hist, i.e., not used elsewhere. + not (any (`nameIn` unfus_set) outVars), + -- 2. all arrays produced by the map are input to the scatter. + mapWriteFusionOK outVars ker -> do + let (extra_nms, res_lam', new_inp) = mapLikeFusionCheck + success (outNames ker ++ extra_nms) $ + SOAC.Hist w ops res_lam' new_inp + + -- Hist-Hist fusion + ( SOAC.Hist _ ops_c _ _, + SOAC.Hist _ ops_p _ _ + ) + | horizFuse -> do + let p_num_buckets = length ops_p + c_num_buckets = length ops_c + (body_p, body_c) = (lambdaBody lam_p, lambdaBody lam_c) + body' = + Body + { bodyDec = bodyDec body_p, -- body_p and body_c have the same lores + bodyStms = bodyStms body_p <> bodyStms body_c, + bodyResult = + take c_num_buckets (bodyResult body_c) + ++ take p_num_buckets (bodyResult body_p) + ++ drop c_num_buckets (bodyResult body_c) + ++ drop p_num_buckets (bodyResult body_p) + } + lam' = + Lambda + { lambdaParams = lambdaParams lam_c ++ lambdaParams lam_p, + lambdaBody = body', + lambdaReturnType = + replicate (c_num_buckets + p_num_buckets) (Prim int64) + ++ drop c_num_buckets (lambdaReturnType lam_c) + ++ drop p_num_buckets (lambdaReturnType lam_p) + } + success (outNames ker ++ returned_outvars) $ + SOAC.Hist w (ops_c <> ops_p) lam' (inp_c_arr <> inp_p_arr) + + -- Scatter-write fusion. + ( SOAC.Scatter _len2 _lam_c ivs2 as2, + SOAC.Scatter _len_p _lam_p ivs_p as_p + ) + | horizFuse -> do + let zipW xs ys = ys_p ++ xs_p ++ ys2 ++ xs2 + where + lenx = length xs `div` 2 + xs_p = take lenx xs + xs2 = drop lenx xs + leny = length ys `div` 2 + ys_p = take leny ys + ys2 = drop leny ys + let (body_p, body2) = (lambdaBody lam_p, lambdaBody lam_c) + let body' = + Body + { bodyDec = bodyDec body_p, -- body_p and body2 have the same lores + bodyStms = bodyStms body_p <> bodyStms body2, + bodyResult = zipW (bodyResult body_p) (bodyResult body2) + } + let lam' = + Lambda + { lambdaParams = lambdaParams lam_p ++ lambdaParams lam_c, + lambdaBody = body', + lambdaReturnType = zipW (lambdaReturnType lam_p) (lambdaReturnType lam_c) + } + success (outNames ker ++ returned_outvars) $ + SOAC.Scatter w lam' (ivs_p ++ ivs2) (as2 ++ as_p) + (SOAC.Scatter {}, _) -> + fail "Cannot fuse a write with anything else than a write or a map" + (_, SOAC.Scatter {}) -> + fail "Cannot fuse a write with anything else than a write or a map" + ---------------------------- + -- Stream-Stream Fusions: -- + ---------------------------- + (SOAC.Stream _ Sequential {} _ _, SOAC.Stream _ form_p@Sequential {} _ _) + | mapFusionOK (drop (length $ getStreamAccums form_p) 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 {} _ _) -> + fail "Fusion conditions not met for two SEQ streams!" + (SOAC.Stream _ Sequential {} _ _, SOAC.Stream {}) -> + fail "Cannot fuse a parallel with a sequential Stream!" + (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 + -- fuse two PARALLEL streams + (res_nms, res_stream) <- fuseStreamHelper (outNames ker) unfus_set outVars outPairs soac_c soac_p + success res_nms res_stream + (SOAC.Stream {}, SOAC.Stream {}) -> + fail "Fusion conditions not met for two PAR streams!" + ------------------------------------------------------------------- + --- If one is a stream, translate the other to a stream as well.--- + --- This does not get in trouble (infinite computation) because --- + --- scan's translation to Stream introduces a hindrance to --- + --- (horizontal fusion), hence repeated application is for the--- + --- moment impossible. However, if with a dependence-graph rep--- + --- we could run in an infinite recursion, i.e., repeatedly --- + --- fusing map o scan into an infinity of Stream levels! --- + ------------------------------------------------------------------- + (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, + -- and recursively perform stream-stream fusion. + (soac_p', newacc_ids) <- SOAC.soacToStream soac_p + soac_p'' <- case form2 of + Sequential {} -> toSeqStream soac_p' + _ -> return soac_p' + if soac_p' == soac_p + then fail "SOAC could not be turned into stream." + else fuseSOACwithKer unfus_set (map identName newacc_ids ++ outVars) soac_p'' soac_p_consumed ker + (_, SOAC.Screma _ form _) | Just _ <- Futhark.isScanSOAC form -> do + -- A Scan soac can be currently only fused as a (sequential) stream, + -- hence it is first translated to a (sequential) Stream and then + -- fusion with a kernel is attempted. + (soac_p', newacc_ids) <- SOAC.soacToStream soac_p + 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 + -- 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. + -- The newly created stream corresponding to soac_c borrows the + -- sequential/parallel property of the soac_p stream. + (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' + _ -> return soac_c' + + fuseSOACwithKer unfus_set outVars soac_p soac_p_consumed $ + ker {fsoac = soac_c'', outNames = map identName newacc_ids ++ outNames ker} + + --------------------------------- + --- DEFAULT, CANNOT FUSE CASE --- + --------------------------------- + _ -> fail "Cannot fuse" + +getStreamOrder :: StreamForm lore -> StreamOrd +getStreamOrder (Parallel o _ _ _) = o +getStreamOrder (Sequential _) = InOrder + +fuseStreamHelper :: + [VName] -> + Names -> + [VName] -> + [(VName, Ident)] -> + SOAC -> + SOAC -> + TryFusion ([VName], SOAC) +fuseStreamHelper + out_kernms + unfus_set + outVars + outPairs + (SOAC.Stream w2 form2 lam2 inp2_arr) + (SOAC.Stream _ form1 lam1 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 + chunk2 = head $ lambdaParams lam2 + hmnms = M.fromList [(paramName chunk2, paramName chunk1)] + lam20 = substituteNames hmnms lam2 + lam1' = lam1 {lambdaParams = tail $ lambdaParams lam1} + lam2' = lam20 {lambdaParams = tail $ lambdaParams lam20} + (res_lam', new_inp) = + fuseRedomap + unfus_set + outVars + lam1' + [] + nes1 + inp1_arr + outPairs + lam2' + [] + (getStreamAccums form2) + inp2_arr + res_lam'' = res_lam' {lambdaParams = chunk1 : lambdaParams res_lam'} + unfus_accs = take (length nes1) outVars + unfus_arrs = filter (`nameIn` unfus_set) outVars + res_form <- mergeForms form2 form1 + return + ( unfus_accs ++ out_kernms ++ unfus_arrs, + SOAC.Stream w2 res_form res_lam'' 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 _ _ = 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 _ = fail "toSeqStream expects a stream, but given a SOAC." + +-- Here follows optimizations and transforms to expose fusability. + +optimizeKernel :: Maybe [VName] -> FusedKer -> TryFusion FusedKer +optimizeKernel inp ker = do + (soac, resTrans) <- optimizeSOAC inp (fsoac ker) startTrans + return $ + ker + { fsoac = soac, + outputTransform = resTrans + } + where + startTrans = outputTransform ker + +optimizeSOAC :: + Maybe [VName] -> + SOAC -> + SOAC.ArrayTransforms -> + TryFusion (SOAC, SOAC.ArrayTransforms) +optimizeSOAC inp soac os = do + res <- foldM comb (False, soac, os) optimizations + case res of + (False, _, _) -> fail "No optimisation applied" + (True, soac', os') -> return (soac', os') + where + comb (changed, soac', os') f = + do + (soac'', os'') <- f inp soac' os + return (True, soac'', os'') + <|> return (changed, soac', os') + +type Optimization = + Maybe [VName] -> + SOAC -> + SOAC.ArrayTransforms -> + TryFusion (SOAC, SOAC.ArrayTransforms) + +optimizations :: [Optimization] +optimizations = [iswim] + +iswim :: + Maybe [VName] -> + SOAC -> + SOAC.ArrayTransforms -> + TryFusion (SOAC, SOAC.ArrayTransforms) +iswim _ (SOAC.Screma w form arrs) ots + | Just [Futhark.Scan scan_fun nes] <- Futhark.isScanSOAC form, + Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible scan_fun, + Just nes_names <- mapM subExpVar nes = do + let nes_idents = zipWith Ident nes_names $ lambdaReturnType scan_fun + map_nes = map SOAC.identInput nes_idents + map_arrs' = map_nes ++ map (SOAC.transposeInput 0 1) arrs + (scan_acc_params, scan_elem_params) = + splitAt (length arrs) $ lambdaParams scan_fun + map_params = + map removeParamOuterDim scan_acc_params + ++ map (setParamOuterDimTo w) scan_elem_params + map_rettype = map (`setOuterSize` w) $ lambdaReturnType scan_fun + + scan_params = lambdaParams map_fun + scan_body = lambdaBody map_fun + scan_rettype = lambdaReturnType map_fun + scan_fun' = Lambda scan_params scan_body scan_rettype + nes' = map Var $ take (length map_nes) $ map paramName map_params + arrs' = drop (length map_nes) $ map paramName map_params + + scan_form <- scanSOAC [Futhark.Scan scan_fun' nes'] + + let map_body = + mkBody + ( oneStm $ + Let (setPatternOuterDimTo w map_pat) (defAux ()) $ + Op $ Futhark.Screma w scan_form arrs' + ) + $ map Var $ patternNames map_pat + map_fun' = Lambda map_params map_body map_rettype + perm = case lambdaReturnType map_fun of + [] -> [] + t : _ -> 1 : 0 : [2 .. arrayRank t] + + return + ( SOAC.Screma map_w (ScremaForm [] [] map_fun') map_arrs', + ots SOAC.|> SOAC.Rearrange map_cs perm + ) +iswim _ _ _ = + fail "ISWIM does not apply." + +removeParamOuterDim :: LParam -> LParam +removeParamOuterDim param = + let t = rowType $ paramType param + in param {paramDec = t} + +setParamOuterDimTo :: SubExp -> LParam -> LParam +setParamOuterDimTo w param = + let t = paramType param `setOuterSize` w + in param {paramDec = t} + +setPatternOuterDimTo :: SubExp -> Pattern -> Pattern +setPatternOuterDimTo w = fmap (`setOuterSize` w) + +-- Now for fiddling with transpositions... + +commonTransforms :: + [VName] -> + [SOAC.Input] -> + (SOAC.ArrayTransforms, [SOAC.Input]) +commonTransforms interesting inps = commonTransforms' inps' + where + inps' = + [ (SOAC.inputArray inp `elem` interesting, inp) + | inp <- inps + ] + +commonTransforms' :: [(Bool, SOAC.Input)] -> (SOAC.ArrayTransforms, [SOAC.Input]) +commonTransforms' inps = + case foldM inspect (Nothing, []) inps of + Just (Just mot, inps') -> first (mot SOAC.<|) $ commonTransforms' $ reverse inps' + _ -> (SOAC.noTransforms, map snd inps) + where + inspect (mot, prev) (True, inp) = + case (mot, inputToOutput inp) of + (Nothing, Just (ot, inp')) -> Just (Just ot, (True, inp') : prev) + (Just ot1, Just (ot2, inp')) + | ot1 == ot2 -> Just (Just ot2, (True, inp') : prev) + _ -> Nothing + inspect (mot, prev) inp = Just (mot, inp : prev) + +mapDepth :: MapNest -> Int +mapDepth (MapNest.MapNest _ lam levels _) = + min resDims (length levels) + 1 + where + resDims = minDim $ case levels of + [] -> lambdaReturnType lam + nest : _ -> MapNest.nestingReturnType nest + minDim [] = 0 + minDim (t : ts) = foldl min (arrayRank t) $ map arrayRank ts + +pullRearrange :: + SOAC -> + SOAC.ArrayTransforms -> + TryFusion (SOAC, SOAC.ArrayTransforms) +pullRearrange soac ots = do + nest <- liftMaybe =<< MapNest.fromSOAC soac + SOAC.Rearrange cs perm SOAC.:< ots' <- return $ SOAC.viewf ots + if rearrangeReach perm <= mapDepth nest + then do + let -- Expand perm to cover the full extent of the input dimensionality + perm' inp = take r perm ++ [length perm .. r -1] + where + r = SOAC.inputRank inp + addPerm inp = SOAC.addTransform (SOAC.Rearrange cs $ perm' inp) inp + inputs' = map addPerm $ MapNest.inputs nest + soac' <- + MapNest.toSOAC $ + inputs' `MapNest.setInputs` rearrangeReturnTypes nest perm + return (soac', ots') + else fail "Cannot pull transpose" + +pushRearrange :: + [VName] -> + SOAC -> + SOAC.ArrayTransforms -> + TryFusion (SOAC, SOAC.ArrayTransforms) +pushRearrange inpIds soac ots = do + nest <- liftMaybe =<< MapNest.fromSOAC soac + (perm, inputs') <- liftMaybe $ fixupInputs inpIds $ MapNest.inputs nest + if rearrangeReach perm <= mapDepth nest + then do + let invertRearrange = SOAC.Rearrange mempty $ rearrangeInverse perm + soac' <- + MapNest.toSOAC $ + inputs' + `MapNest.setInputs` rearrangeReturnTypes nest perm + return (soac', invertRearrange SOAC.<| ots) + else fail "Cannot push transpose" + +-- | Actually also rearranges indices. +rearrangeReturnTypes :: MapNest -> [Int] -> MapNest +rearrangeReturnTypes nest@(MapNest.MapNest w body nestings inps) perm = + MapNest.MapNest + w + body + ( zipWith + setReturnType + nestings + $ drop 1 $ iterate (map rowType) ts + ) + inps + where + origts = MapNest.typeOf nest + -- The permutation may be deeper than the rank of the type, + -- but it is required that it is an identity permutation + -- beyond that. This is supposed to be checked as an + -- invariant by whoever calls rearrangeReturnTypes. + rearrangeType' t = rearrangeType (take (arrayRank t) perm) t + ts = map rearrangeType' origts + + setReturnType nesting t' = + nesting {MapNest.nestingReturnType = t'} + +fixupInputs :: [VName] -> [SOAC.Input] -> Maybe ([Int], [SOAC.Input]) +fixupInputs inpIds inps = + case mapMaybe inputRearrange $ filter exposable inps of + perm : _ -> do + inps' <- mapM (fixupInput (rearrangeReach perm) perm) inps + return (perm, inps') + _ -> Nothing + where + exposable = (`elem` inpIds) . SOAC.inputArray + + inputRearrange (SOAC.Input ts _ _) + | _ SOAC.:> SOAC.Rearrange _ perm <- SOAC.viewl ts = Just perm + inputRearrange _ = Nothing + + fixupInput d perm inp + | r <- SOAC.inputRank inp, + r >= d = + Just $ SOAC.addTransform (SOAC.Rearrange mempty $ take r perm) inp + | otherwise = Nothing + +pullReshape :: SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms) +pullReshape (SOAC.Screma _ form inps) ots + | Just maplam <- Futhark.isMapSOAC form, + SOAC.Reshape cs shape SOAC.:< ots' <- SOAC.viewf ots, + all primType $ lambdaReturnType maplam = do + let mapw' = case reverse $ newDims shape of + [] -> intConst Int64 0 + d : _ -> d + inputs' = map (SOAC.addTransform $ SOAC.ReshapeOuter cs shape) inps + inputTypes = map SOAC.inputType inputs' + + let outersoac :: + ([SOAC.Input] -> SOAC) -> + (SubExp, [SubExp]) -> + TryFusion ([SOAC.Input] -> SOAC) + outersoac inner (w, outershape) = do + let addDims t = arrayOf t (Shape outershape) NoUniqueness + retTypes = map addDims $ lambdaReturnType maplam + + ps <- forM inputTypes $ \inpt -> + newParam "pullReshape_param" $ + stripArray (length shape - length outershape) inpt + + inner_body <- + runBodyBinder $ + eBody [SOAC.toExp $ inner $ map (SOAC.identInput . paramIdent) ps] + let inner_fun = + Lambda + { lambdaParams = ps, + lambdaReturnType = retTypes, + lambdaBody = inner_body + } + return $ SOAC.Screma w $ Futhark.mapSOAC inner_fun + + op' <- + foldM outersoac (SOAC.Screma mapw' $ Futhark.mapSOAC maplam) $ + zip (drop 1 $ reverse $ newDims shape) $ + drop 1 $ reverse $ drop 1 $ tails $ newDims shape + return (op' inputs', ots') +pullReshape _ _ = fail "Cannot pull reshape" + +-- Tie it all together in exposeInputs (for making inputs to a +-- consumer available) and pullOutputTransforms (for moving +-- output-transforms of a producer to its inputs instead). + +exposeInputs :: + [VName] -> + FusedKer -> + TryFusion (FusedKer, SOAC.ArrayTransforms) +exposeInputs inpIds ker = + (exposeInputs' =<< pushRearrange') + <|> (exposeInputs' =<< pullRearrange') + <|> exposeInputs' ker + where + ot = outputTransform ker + + pushRearrange' = do + (soac', ot') <- pushRearrange inpIds (fsoac ker) ot + return + ker + { fsoac = soac', + outputTransform = ot' + } + + pullRearrange' = do + (soac', ot') <- pullRearrange (fsoac ker) ot + unless (SOAC.nullTransforms ot') $ + fail "pullRearrange was not enough" + return + ker + { fsoac = soac', + outputTransform = SOAC.noTransforms + } + + exposeInputs' ker' = + case commonTransforms inpIds $ inputs ker' of + (ot', inps') + | all exposed inps' -> + return (ker' {fsoac = inps' `SOAC.setInputs` fsoac ker'}, ot') + _ -> fail "Cannot expose" + + exposed (SOAC.Input ts _ _) + | SOAC.nullTransforms ts = True + exposed inp = SOAC.inputArray inp `notElem` inpIds + +outputTransformPullers :: [SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms)] +outputTransformPullers = [pullRearrange, pullReshape] + +pullOutputTransforms :: + SOAC -> + SOAC.ArrayTransforms -> + TryFusion (SOAC, SOAC.ArrayTransforms) +pullOutputTransforms = attempt outputTransformPullers + where + attempt [] _ _ = fail "Cannot pull anything" + attempt (p : ps) soac ots = + do + (soac', ots') <- p soac ots + if SOAC.nullTransforms ots' + then return (soac', SOAC.noTransforms) + else pullOutputTransforms soac' ots' <|> return (soac', ots') + <|> attempt ps soac ots diff --git a/src/Futhark/Optimise/InPlaceLowering.hs b/src/Futhark/Optimise/InPlaceLowering.hs --- a/src/Futhark/Optimise/InPlaceLowering.hs +++ b/src/Futhark/Optimise/InPlaceLowering.hs @@ -1,10 +1,11 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE UndecidableInstances #-} + -- | This module implements an optimisation that moves in-place -- updates into/before loops where possible, with the end goal of -- minimising memory copies. As an example, consider this program: @@ -63,21 +64,19 @@ -- FIXME: the implementation is not finished yet. Specifically, not -- all of the above conditions are checked. module Futhark.Optimise.InPlaceLowering - ( inPlaceLoweringKernels - , inPlaceLoweringSeq - ) + ( inPlaceLoweringKernels, + inPlaceLoweringSeq, + ) where import Control.Monad.RWS import qualified Data.Map.Strict as M - import Futhark.Analysis.Alias +import Futhark.Binder import Futhark.IR.Aliases import Futhark.IR.Kernels import Futhark.IR.Seq (Seq) import Futhark.Optimise.InPlaceLowering.LowerIntoStm -import Futhark.MonadFreshNames -import Futhark.Binder import Futhark.Pass -- | Apply the in-place lowering optimisation to the given program. @@ -89,49 +88,62 @@ inPlaceLoweringSeq = inPlaceLowering pure lowerUpdate -- | Apply the in-place lowering optimisation to the given program. -inPlaceLowering :: Constraints lore => - OnOp lore -> LowerUpdate lore (ForwardingM lore) - -> Pass lore lore +inPlaceLowering :: + Constraints lore => + OnOp lore -> + LowerUpdate lore (ForwardingM lore) -> + Pass lore lore inPlaceLowering onOp lower = Pass "In-place lowering" "Lower in-place updates into loops" $ - fmap removeProgAliases . - intraproceduralTransformationWithConsts optimiseConsts optimiseFunDef . - aliasAnalysis - where optimiseConsts stms = - modifyNameSource $ runForwardingM lower onOp $ + fmap removeProgAliases + . intraproceduralTransformationWithConsts optimiseConsts optimiseFunDef + . aliasAnalysis + where + optimiseConsts stms = + modifyNameSource $ + runForwardingM lower onOp $ stmsFromList <$> optimiseStms (stmsToList stms) (pure ()) - optimiseFunDef consts fundec = - modifyNameSource $ runForwardingM lower onOp $ - descend (stmsToList consts) $ bindingFParams (funDefParams fundec) $ do - body <- optimiseBody $ funDefBody fundec - return $ fundec { funDefBody = body } + optimiseFunDef consts fundec = + modifyNameSource $ + runForwardingM lower onOp $ + descend (stmsToList consts) $ + bindingFParams (funDefParams fundec) $ do + body <- optimiseBody $ funDefBody fundec + return $ fundec {funDefBody = body} - descend [] m = m - descend (stm:stms) m = bindingStm stm $ descend stms m + descend [] m = m + descend (stm : stms) m = bindingStm stm $ descend stms m type Constraints lore = (Bindable lore, CanBeAliased (Op lore)) -optimiseBody :: Constraints lore => - Body (Aliases lore) -> ForwardingM lore (Body (Aliases lore)) +optimiseBody :: + Constraints lore => + Body (Aliases lore) -> + ForwardingM lore (Body (Aliases lore)) optimiseBody (Body als bnds res) = do - bnds' <- deepen $ optimiseStms (stmsToList bnds) $ - mapM_ seen res + bnds' <- + deepen $ + optimiseStms (stmsToList bnds) $ + mapM_ seen res return $ Body als (stmsFromList bnds') res - where seen Constant{} = return () - seen (Var v) = seenVar v + where + seen Constant {} = return () + seen (Var v) = seenVar v -optimiseStms :: Constraints lore => - [Stm (Aliases lore)] -> ForwardingM lore () - -> ForwardingM lore [Stm (Aliases lore)] +optimiseStms :: + Constraints lore => + [Stm (Aliases lore)] -> + ForwardingM lore () -> + ForwardingM lore [Stm (Aliases lore)] optimiseStms [] m = m >> return [] - -optimiseStms (bnd:bnds) m = do +optimiseStms (bnd : bnds) m = do (bnds', bup) <- tapBottomUp $ bindingStm bnd $ optimiseStms bnds m bnd' <- optimiseInStm bnd case filter ((`elem` boundHere) . updateValue) $ forwardThese bup of - [] -> do checkIfForwardableUpdate bnd' - return $ bnd':bnds' + [] -> do + checkIfForwardableUpdate bnd' + return $ bnd' : bnds' updates -> do lower <- asks topLowerUpdate scope <- askScope @@ -145,20 +157,23 @@ -- Condition (5) and (7) are assumed to be checked by -- lowerUpdate. case lower scope bnd' updates of - Just lowering -> do new_bnds <- lowering - new_bnds' <- optimiseStms new_bnds $ - tell bup { forwardThese = [] } - return $ new_bnds' ++ filter notUpdated bnds' - Nothing -> do checkIfForwardableUpdate bnd' - return $ bnd':bnds' - - where boundHere = patternNames $ stmPattern bnd + Just lowering -> do + new_bnds <- lowering + new_bnds' <- + optimiseStms new_bnds $ + tell bup {forwardThese = []} + return $ new_bnds' ++ filter notUpdated bnds' + Nothing -> do + checkIfForwardableUpdate bnd' + return $ bnd' : bnds' + where + boundHere = patternNames $ stmPattern bnd - checkIfForwardableUpdate (Let pat (StmAux cs _ _) e) - | Pattern [] [PatElem v dec] <- pat, - BasicOp (Update src slice (Var ve)) <- e = - maybeForward ve v dec cs src slice - checkIfForwardableUpdate _ = return () + checkIfForwardableUpdate (Let pat (StmAux cs _ _) e) + | Pattern [] [PatElem v dec] <- pat, + BasicOp (Update src slice (Var ve)) <- e = + maybeForward ve v dec cs src slice + checkIfForwardableUpdate _ = return () optimiseInStm :: Constraints lore => Stm (Aliases lore) -> ForwardingM lore (Stm (Aliases lore)) optimiseInStm (Let pat dec e) = @@ -167,46 +182,55 @@ optimiseExp :: Constraints lore => Exp (Aliases lore) -> ForwardingM lore (Exp (Aliases lore)) optimiseExp (DoLoop ctx val form body) = bindingScope (scopeOf form) $ - bindingFParams (map fst $ ctx ++ val) $ - DoLoop ctx val form <$> optimiseBody body + bindingFParams (map fst $ ctx ++ val) $ + DoLoop ctx val form <$> optimiseBody body optimiseExp (Op op) = do f <- asks topOnOp Op <$> f op optimiseExp e = mapExpM optimise e - where optimise = identityMapper { mapOnBody = const optimiseBody - } + where + optimise = + identityMapper + { mapOnBody = const optimiseBody + } + onKernelOp :: OnOp Kernels onKernelOp (SegOp op) = bindingScope (scopeOfSegSpace (segSpace op)) $ do - let mapper = identitySegOpMapper { mapOnSegOpBody = onKernelBody } + let mapper = identitySegOpMapper {mapOnSegOpBody = onKernelBody} onKernelBody kbody = do - stms <- deepen $ optimiseStms (stmsToList (kernelBodyStms kbody)) $ - mapM_ seenVar $ namesToList $ freeIn $ kernelBodyResult kbody - return kbody { kernelBodyStms = stmsFromList stms } + stms <- + deepen $ + optimiseStms (stmsToList (kernelBodyStms kbody)) $ + mapM_ seenVar $ namesToList $ freeIn $ kernelBodyResult kbody + return kbody {kernelBodyStms = stmsFromList stms} SegOp <$> mapSegOpM mapper op onKernelOp op = return op -data Entry lore = Entry { entryNumber :: Int - , entryAliases :: Names - , entryDepth :: Int - , entryOptimisable :: Bool - , entryType :: NameInfo (Aliases lore) - } +data Entry lore = Entry + { entryNumber :: Int, + entryAliases :: Names, + entryDepth :: Int, + entryOptimisable :: Bool, + entryType :: NameInfo (Aliases lore) + } type VTable lore = M.Map VName (Entry lore) type OnOp lore = Op (Aliases lore) -> ForwardingM lore (Op (Aliases lore)) -data TopDown lore = TopDown { topDownCounter :: Int - , topDownTable :: VTable lore - , topDownDepth :: Int - , topLowerUpdate :: LowerUpdate lore (ForwardingM lore) - , topOnOp :: OnOp lore - } +data TopDown lore = TopDown + { topDownCounter :: Int, + topDownTable :: VTable lore, + topDownDepth :: Int, + topLowerUpdate :: LowerUpdate lore (ForwardingM lore), + topOnOp :: OnOp lore + } -data BottomUp lore = BottomUp { bottomUpSeen :: Names - , forwardThese :: [DesiredUpdate (LetDec (Aliases lore))] - } +data BottomUp lore = BottomUp + { bottomUpSeen :: Names, + forwardThese :: [DesiredUpdate (LetDec (Aliases lore))] + } instance Semigroup (BottomUp lore) where BottomUp seen1 forward1 <> BottomUp seen2 forward2 = @@ -216,10 +240,14 @@ mempty = BottomUp mempty mempty newtype ForwardingM lore a = ForwardingM (RWS (TopDown lore) (BottomUp lore) VNameSource a) - deriving (Monad, Applicative, Functor, - MonadReader (TopDown lore), - MonadWriter (BottomUp lore), - MonadState VNameSource) + deriving + ( Monad, + Applicative, + Functor, + MonadReader (TopDown lore), + MonadWriter (BottomUp lore), + MonadState VNameSource + ) instance MonadFreshNames (ForwardingM lore) where getNameSource = get @@ -228,63 +256,81 @@ instance Constraints lore => HasScope (Aliases lore) (ForwardingM lore) where askScope = M.map entryType <$> asks topDownTable -runForwardingM :: LowerUpdate lore (ForwardingM lore) -> OnOp lore -> ForwardingM lore a - -> VNameSource -> (a, VNameSource) -runForwardingM f g (ForwardingM m) src = let (x, src', _) = runRWS m emptyTopDown src - in (x, src') - where emptyTopDown = TopDown { topDownCounter = 0 - , topDownTable = M.empty - , topDownDepth = 0 - , topLowerUpdate = f - , topOnOp = g - } +runForwardingM :: + LowerUpdate lore (ForwardingM lore) -> + OnOp lore -> + ForwardingM lore a -> + VNameSource -> + (a, VNameSource) +runForwardingM f g (ForwardingM m) src = + let (x, src', _) = runRWS m emptyTopDown src + in (x, src') + where + emptyTopDown = + TopDown + { topDownCounter = 0, + topDownTable = M.empty, + topDownDepth = 0, + topLowerUpdate = f, + topOnOp = g + } -bindingParams :: (dec -> NameInfo (Aliases lore)) - -> [Param dec] - -> ForwardingM lore a - -> ForwardingM lore a +bindingParams :: + (dec -> NameInfo (Aliases lore)) -> + [Param dec] -> + ForwardingM lore a -> + ForwardingM lore a bindingParams f params = local $ \(TopDown n vtable d x y) -> let entry fparam = - (paramName fparam, - Entry n mempty d False $ f $ paramDec fparam) + ( paramName fparam, + Entry n mempty d False $ f $ paramDec fparam + ) entries = M.fromList $ map entry params - in TopDown (n+1) (M.union entries vtable) d x y + in TopDown (n + 1) (M.union entries vtable) d x y -bindingFParams :: [FParam (Aliases lore)] - -> ForwardingM lore a - -> ForwardingM lore a +bindingFParams :: + [FParam (Aliases lore)] -> + ForwardingM lore a -> + ForwardingM lore a bindingFParams = bindingParams FParamName -bindingScope :: Scope (Aliases lore) - -> ForwardingM lore a - -> ForwardingM lore a +bindingScope :: + Scope (Aliases lore) -> + ForwardingM lore a -> + ForwardingM lore a bindingScope scope = local $ \(TopDown n vtable d x y) -> let entries = M.map entry scope infoAliases (LetName (aliases, _)) = unAliases aliases infoAliases _ = mempty entry info = Entry n (infoAliases info) d False info - in TopDown (n+1) (entries<>vtable) d x y + in TopDown (n + 1) (entries <> vtable) d x y -bindingStm :: Stm (Aliases lore) - -> ForwardingM lore a - -> ForwardingM lore a +bindingStm :: + Stm (Aliases lore) -> + ForwardingM lore a -> + ForwardingM lore a bindingStm (Let pat _ _) = local $ \(TopDown n vtable d x y) -> let entries = M.fromList $ map entry $ patternElements pat entry patElem = let (aliases, _) = patElemDec patElem - in (patElemName patElem, - Entry n (unAliases aliases) d True $ LetName $ patElemDec patElem) - in TopDown (n+1) (M.union entries vtable) d x y + in ( patElemName patElem, + Entry n (unAliases aliases) d True $ LetName $ patElemDec patElem + ) + in TopDown (n + 1) (M.union entries vtable) d x y bindingNumber :: VName -> ForwardingM lore Int bindingNumber name = do res <- asks $ fmap entryNumber . M.lookup name . topDownTable - case res of Just n -> return n - Nothing -> error $ "bindingNumber: variable " ++ - pretty name ++ " not found." + case res of + Just n -> return n + Nothing -> + error $ + "bindingNumber: variable " + ++ pretty name + ++ " not found." deepen :: ForwardingM lore a -> ForwardingM lore a -deepen = local $ \env -> env { topDownDepth = topDownDepth env + 1 } +deepen = local $ \env -> env {topDownDepth = topDownDepth env + 1} areAvailableBefore :: Names -> VName -> ForwardingM lore Bool areAvailableBefore names point = do @@ -296,37 +342,53 @@ isInCurrentBody name = do current <- asks topDownDepth res <- asks $ fmap entryDepth . M.lookup name . topDownTable - case res of Just d -> return $ d == current - Nothing -> error $ "isInCurrentBody: variable " ++ - pretty name ++ " not found." + case res of + Just d -> return $ d == current + Nothing -> + error $ + "isInCurrentBody: variable " + ++ pretty name + ++ " not found." isOptimisable :: VName -> ForwardingM lore Bool isOptimisable name = do res <- asks $ fmap entryOptimisable . M.lookup name . topDownTable - case res of Just b -> return b - Nothing -> error $ "isOptimisable: variable " ++ - pretty name ++ " not found." + case res of + Just b -> return b + Nothing -> + error $ + "isOptimisable: variable " + ++ pretty name + ++ " not found." seenVar :: VName -> ForwardingM lore () seenVar name = do - aliases <- asks $ - maybe mempty entryAliases . - M.lookup name . topDownTable - tell $ mempty { bottomUpSeen = oneName name <> aliases } + aliases <- + asks $ + maybe mempty entryAliases + . M.lookup name + . topDownTable + tell $ mempty {bottomUpSeen = oneName name <> aliases} tapBottomUp :: ForwardingM lore a -> ForwardingM lore (a, BottomUp lore) -tapBottomUp m = do (x,bup) <- listen m - return (x, bup) +tapBottomUp m = do + (x, bup) <- listen m + return (x, bup) -maybeForward :: Constraints lore => - VName - -> VName -> LetDec (Aliases lore) - -> Certificates -> VName -> Slice SubExp - -> ForwardingM lore () +maybeForward :: + Constraints lore => + VName -> + VName -> + LetDec (Aliases lore) -> + Certificates -> + VName -> + Slice SubExp -> + ForwardingM lore () maybeForward v dest_nm dest_dec cs src slice = do -- Checks condition (2) - available <- (freeIn src <> freeIn slice <> freeIn cs) - `areAvailableBefore` v + available <- + (freeIn src <> freeIn slice <> freeIn cs) + `areAvailableBefore` v -- Check condition (3) samebody <- isInCurrentBody v -- Check condition (6) @@ -334,4 +396,4 @@ not_prim <- not . primType <$> lookupType v when (available && samebody && optimisable && not_prim) $ do let fwd = DesiredUpdate dest_nm dest_dec cs src slice v - tell mempty { forwardThese = [fwd] } + tell mempty {forwardThese = [fwd]} diff --git a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs --- a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs +++ b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs @@ -1,72 +1,88 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + module Futhark.Optimise.InPlaceLowering.LowerIntoStm - ( lowerUpdateKernels - , lowerUpdate - , LowerUpdate - , DesiredUpdate (..) - ) where + ( lowerUpdateKernels, + lowerUpdate, + LowerUpdate, + DesiredUpdate (..), + ) +where import Control.Monad import Control.Monad.Writer -import Data.List (find, unzip4) -import Data.Maybe (mapMaybe) import Data.Either +import Data.List (find, unzip4) import qualified Data.Map as M - +import Data.Maybe (mapMaybe) import Futhark.Analysis.PrimExp.Convert -import Futhark.IR.Prop.Aliases +import Futhark.Construct import Futhark.IR.Aliases import Futhark.IR.Kernels -import Futhark.Construct import Futhark.Optimise.InPlaceLowering.SubstituteIndices -data DesiredUpdate dec = - DesiredUpdate { updateName :: VName -- ^ Name of result. - , updateType :: dec -- ^ Type of result. - , updateCertificates :: Certificates - , updateSource :: VName - , updateIndices :: Slice SubExp - , updateValue :: VName - } +data DesiredUpdate dec = DesiredUpdate + { -- | Name of result. + updateName :: VName, + -- | Type of result. + updateType :: dec, + updateCertificates :: Certificates, + updateSource :: VName, + updateIndices :: Slice SubExp, + updateValue :: VName + } deriving (Show) instance Functor DesiredUpdate where - f `fmap` u = u { updateType = f $ updateType u } + f `fmap` u = u {updateType = f $ updateType u} updateHasValue :: VName -> DesiredUpdate dec -> Bool -updateHasValue name = (name==) . updateValue +updateHasValue name = (name ==) . updateValue -type LowerUpdate lore m = Scope (Aliases lore) - -> Stm (Aliases lore) - -> [DesiredUpdate (LetDec (Aliases lore))] - -> Maybe (m [Stm (Aliases lore)]) +type LowerUpdate lore m = + Scope (Aliases lore) -> + Stm (Aliases lore) -> + [DesiredUpdate (LetDec (Aliases lore))] -> + Maybe (m [Stm (Aliases lore)]) -lowerUpdate :: (MonadFreshNames m, Bindable lore, - LetDec lore ~ Type, CanBeAliased (Op lore)) => LowerUpdate lore m +lowerUpdate :: + ( MonadFreshNames m, + Bindable lore, + LetDec lore ~ Type, + CanBeAliased (Op lore) + ) => + LowerUpdate lore m lowerUpdate scope (Let pat aux (DoLoop ctx val form body)) updates = do canDo <- lowerUpdateIntoLoop scope updates pat ctx val form body Just $ do (prebnds, postbnds, ctxpat, valpat, ctx', val', body') <- canDo return $ - prebnds ++ [certify (stmAuxCerts aux) $ - mkLet ctxpat valpat $ DoLoop ctx' val' form body'] ++ postbnds -lowerUpdate _ + prebnds + ++ [ certify (stmAuxCerts aux) $ + mkLet ctxpat valpat $ DoLoop ctx' val' form body' + ] + ++ postbnds +lowerUpdate + _ (Let pat aux (BasicOp (SubExp (Var v)))) [DesiredUpdate bindee_nm bindee_dec cs src is val] - | patternNames pat == [src] = - let is' = fullSlice (typeOf bindee_dec) is - in Just $ - return [certify (stmAuxCerts aux <> cs) $ - mkLet [] [Ident bindee_nm $ typeOf bindee_dec] $ - BasicOp $ Update v is' $ Var val] + | patternNames pat == [src] = + let is' = fullSlice (typeOf bindee_dec) is + in Just $ + return + [ certify (stmAuxCerts aux <> cs) $ + mkLet [] [Ident bindee_nm $ typeOf bindee_dec] $ + BasicOp $ Update v is' $ Var val + ] lowerUpdate _ _ _ = Nothing lowerUpdateKernels :: MonadFreshNames m => LowerUpdate Kernels m -lowerUpdateKernels scope - (Let pat aux (Op (SegOp (SegMap lvl space ts kbody)))) updates - | all ((`elem` patternNames pat) . updateValue) updates = do +lowerUpdateKernels + scope + (Let pat aux (Op (SegOp (SegMap lvl space ts kbody)))) + updates + | all ((`elem` patternNames pat) . updateValue) updates = do mk <- lowerUpdatesIntoSegMap scope pat updates space kbody Just $ do (pat', kbody', poststms) <- mk @@ -76,71 +92,88 @@ stmsToList poststms lowerUpdateKernels scope stm updates = lowerUpdate scope stm updates -lowerUpdatesIntoSegMap :: MonadFreshNames m => - Scope (Aliases Kernels) - -> Pattern (Aliases Kernels) - -> [DesiredUpdate (LetDec (Aliases Kernels))] - -> SegSpace - -> KernelBody (Aliases Kernels) - -> Maybe (m (Pattern (Aliases Kernels), - KernelBody (Aliases Kernels), - Stms (Aliases Kernels))) +lowerUpdatesIntoSegMap :: + MonadFreshNames m => + Scope (Aliases Kernels) -> + Pattern (Aliases Kernels) -> + [DesiredUpdate (LetDec (Aliases Kernels))] -> + SegSpace -> + KernelBody (Aliases Kernels) -> + Maybe + ( m + ( Pattern (Aliases Kernels), + KernelBody (Aliases Kernels), + Stms (Aliases Kernels) + ) + ) lowerUpdatesIntoSegMap scope pat updates kspace kbody = do -- The updates are all-or-nothing. Being more liberal would require -- changes to the in-place-lowering pass itself. mk <- zipWithM onRet (patternElements pat) (kernelBodyResult kbody) return $ do (pes, bodystms, krets, poststms) <- unzip4 <$> sequence mk - return (Pattern [] pes, - kbody { kernelBodyStms = kernelBodyStms kbody <> mconcat bodystms - , kernelBodyResult = krets - }, - mconcat poststms) - - where (gtids, _dims) = unzip $ unSegSpace kspace - - onRet (PatElem v v_dec) ret - | Just (DesiredUpdate bindee_nm bindee_dec _cs src slice _val) <- - find ((==v) . updateValue) updates = do - - Returns _ se <- Just ret + return + ( Pattern [] pes, + kbody + { kernelBodyStms = kernelBodyStms kbody <> mconcat bodystms, + kernelBodyResult = krets + }, + mconcat poststms + ) + where + (gtids, _dims) = unzip $ unSegSpace kspace - Just $ do - let pexp = primExpFromSubExp int32 - (slice', bodystms) <- flip runBinderT scope $ - traverse (toSubExp "index") $ - fixSlice (map (fmap pexp) slice) $ - map (pexp . Var) gtids + onRet (PatElem v v_dec) ret + | Just (DesiredUpdate bindee_nm bindee_dec _cs src slice _val) <- + find ((== v) . updateValue) updates = do + Returns _ se <- Just ret - let res_dims = arrayDims $ snd bindee_dec - ret' = WriteReturns res_dims src [(map DimFix slice', se)] + Just $ do + (slice', bodystms) <- + flip runBinderT scope $ + traverse (toSubExp "index") $ + fixSlice (map (fmap pe32) slice) $ + map (pe32 . Var) gtids - return (PatElem bindee_nm bindee_dec, - bodystms, - ret', - oneStm $ mkLet [] [Ident v $ typeOf v_dec] $ - BasicOp $ Index bindee_nm slice) + let res_dims = arrayDims $ snd bindee_dec + ret' = WriteReturns res_dims src [(map DimFix slice', se)] - onRet pe ret = - Just $ return (pe, mempty, ret, mempty) + return + ( PatElem bindee_nm bindee_dec, + bodystms, + ret', + oneStm $ + mkLet [] [Ident v $ typeOf v_dec] $ + BasicOp $ Index bindee_nm slice + ) + onRet pe ret = + Just $ return (pe, mempty, ret, mempty) -lowerUpdateIntoLoop :: (Bindable lore, BinderOps lore, - Aliased lore, LetDec lore ~ (als, Type), - MonadFreshNames m) => - Scope lore - -> [DesiredUpdate (LetDec lore)] - -> Pattern lore - -> [(FParam lore, SubExp)] - -> [(FParam lore, SubExp)] - -> LoopForm lore - -> Body lore - -> Maybe (m ([Stm lore], - [Stm lore], - [Ident], - [Ident], - [(FParam lore, SubExp)], - [(FParam lore, SubExp)], - Body lore)) +lowerUpdateIntoLoop :: + ( Bindable lore, + BinderOps lore, + Aliased lore, + LetDec lore ~ (als, Type), + MonadFreshNames m + ) => + Scope lore -> + [DesiredUpdate (LetDec lore)] -> + Pattern lore -> + [(FParam lore, SubExp)] -> + [(FParam lore, SubExp)] -> + LoopForm lore -> + Body lore -> + Maybe + ( m + ( [Stm lore], + [Stm lore], + [Ident], + [Ident], + [(FParam lore, SubExp)], + [(FParam lore, SubExp)], + Body lore + ) + ) lowerUpdateIntoLoop scope updates pat ctx val form body = do -- Algorithm: -- @@ -176,124 +209,154 @@ Just $ do in_place_map <- mk_in_place_map - (val',prebnds,postbnds) <- mkMerges in_place_map - let (ctxpat,valpat) = mkResAndPat in_place_map + (val', prebnds, postbnds) <- mkMerges in_place_map + let (ctxpat, valpat) = mkResAndPat in_place_map idxsubsts = indexSubstitutions in_place_map (idxsubsts', newbnds) <- substituteIndices idxsubsts $ bodyStms body (body_res, res_bnds) <- manipulateResult in_place_map idxsubsts' - let body' = mkBody (newbnds<>res_bnds) body_res + let body' = mkBody (newbnds <> res_bnds) body_res return (prebnds, postbnds, ctxpat, valpat, ctx, val', body') - where usedInBody = mconcat $ map expandAliases $ namesToList $ freeIn body <> freeIn form - expandAliases v = case M.lookup v scope of - Just (LetName dec) -> oneName v <> aliasesOf dec - _ -> oneName v - resmap = zip (bodyResult body) $ patternValueIdents pat + where + usedInBody = mconcat $ map expandAliases $ namesToList $ freeIn body <> freeIn form + expandAliases v = case M.lookup v scope of + Just (LetName dec) -> oneName v <> aliasesOf dec + _ -> oneName v + resmap = zip (bodyResult body) $ patternValueIdents pat - mkMerges :: (MonadFreshNames m, Bindable lore) => - [LoopResultSummary (als, Type)] - -> m ([(Param DeclType, SubExp)], [Stm lore], [Stm lore]) - mkMerges summaries = do - ((origmerge, extramerge), (prebnds, postbnds)) <- - runWriterT $ partitionEithers <$> mapM mkMerge summaries - return (origmerge ++ extramerge, prebnds, postbnds) + mkMerges :: + (MonadFreshNames m, Bindable lore) => + [LoopResultSummary (als, Type)] -> + m ([(Param DeclType, SubExp)], [Stm lore], [Stm lore]) + mkMerges summaries = do + ((origmerge, extramerge), (prebnds, postbnds)) <- + runWriterT $ partitionEithers <$> mapM mkMerge summaries + return (origmerge ++ extramerge, prebnds, postbnds) - mkMerge summary - | Just (update, mergename, mergedec) <- relatedUpdate summary = do - source <- newVName "modified_source" - let source_t = snd $ updateType update - elmident = Ident - (updateValue update) - (source_t `setArrayDims` sliceDims (updateIndices update)) - tell ([mkLet [] [Ident source source_t] $ BasicOp $ Update - (updateSource update) - (fullSlice source_t $ updateIndices update) $ - snd $ mergeParam summary], - [mkLet [] [elmident] $ BasicOp $ Index - (updateName update) - (fullSlice source_t $ updateIndices update)]) - return $ Right (Param - mergename - (toDecl (typeOf mergedec) Unique), - Var source) - | otherwise = return $ Left $ mergeParam summary + mkMerge summary + | Just (update, mergename, mergedec) <- relatedUpdate summary = do + source <- newVName "modified_source" + let source_t = snd $ updateType update + elmident = + Ident + (updateValue update) + (source_t `setArrayDims` sliceDims (updateIndices update)) + tell + ( [ mkLet [] [Ident source source_t] $ + BasicOp $ + Update + (updateSource update) + (fullSlice source_t $ updateIndices update) + $ snd $ mergeParam summary + ], + [ mkLet [] [elmident] $ + BasicOp $ + Index + (updateName update) + (fullSlice source_t $ updateIndices update) + ] + ) + return $ + Right + ( Param + mergename + (toDecl (typeOf mergedec) Unique), + Var source + ) + | otherwise = return $ Left $ mergeParam summary - mkResAndPat summaries = - let (origpat,extrapat) = partitionEithers $ map mkResAndPat' summaries - in (patternContextIdents pat, - origpat ++ extrapat) + mkResAndPat summaries = + let (origpat, extrapat) = partitionEithers $ map mkResAndPat' summaries + in ( patternContextIdents pat, + origpat ++ extrapat + ) - mkResAndPat' summary - | Just (update, _, _) <- relatedUpdate summary = - Right (Ident (updateName update) (snd $ updateType update)) - | otherwise = - Left (inPatternAs summary) + mkResAndPat' summary + | Just (update, _, _) <- relatedUpdate summary = + Right (Ident (updateName update) (snd $ updateType update)) + | otherwise = + Left (inPatternAs summary) -summariseLoop :: MonadFreshNames m => - [DesiredUpdate (als, Type)] - -> Names - -> [(SubExp, Ident)] - -> [(Param DeclType, SubExp)] - -> Maybe (m [LoopResultSummary (als, Type)]) +summariseLoop :: + MonadFreshNames m => + [DesiredUpdate (als, Type)] -> + Names -> + [(SubExp, Ident)] -> + [(Param DeclType, SubExp)] -> + Maybe (m [LoopResultSummary (als, Type)]) summariseLoop updates usedInBody resmap merge = sequence <$> zipWithM summariseLoopResult resmap merge - where summariseLoopResult (se, v) (fparam, mergeinit) - | Just update <- find (updateHasValue $ identName v) updates = - if updateSource update `nameIn` usedInBody - then Nothing - else if hasLoopInvariantShape fparam then Just $ do - lowered_array <- newVName "lowered_array" - return LoopResultSummary { resultSubExp = se - , inPatternAs = v - , mergeParam = (fparam, mergeinit) - , relatedUpdate = Just (update, - lowered_array, - updateType update) - } - else Nothing - summariseLoopResult _ _ = - Nothing -- XXX: conservative; but this entire pass is going away. - - hasLoopInvariantShape = all loopInvariant . arrayDims . paramType + where + summariseLoopResult (se, v) (fparam, mergeinit) + | Just update <- find (updateHasValue $ identName v) updates = + if updateSource update `nameIn` usedInBody + then Nothing + else + if hasLoopInvariantShape fparam + then Just $ do + lowered_array <- newVName "lowered_array" + return + LoopResultSummary + { resultSubExp = se, + inPatternAs = v, + mergeParam = (fparam, mergeinit), + relatedUpdate = + Just + ( update, + lowered_array, + updateType update + ) + } + else Nothing + summariseLoopResult _ _ = + Nothing -- XXX: conservative; but this entire pass is going away. + hasLoopInvariantShape = all loopInvariant . arrayDims . paramType - merge_param_names = map (paramName . fst) merge + merge_param_names = map (paramName . fst) merge - loopInvariant (Var v) = v `notElem` merge_param_names - loopInvariant Constant{} = True + loopInvariant (Var v) = v `notElem` merge_param_names + loopInvariant Constant {} = True -data LoopResultSummary dec = - LoopResultSummary { resultSubExp :: SubExp - , inPatternAs :: Ident - , mergeParam :: (Param DeclType, SubExp) - , relatedUpdate :: Maybe (DesiredUpdate dec, VName, dec) - } +data LoopResultSummary dec = LoopResultSummary + { resultSubExp :: SubExp, + inPatternAs :: Ident, + mergeParam :: (Param DeclType, SubExp), + relatedUpdate :: Maybe (DesiredUpdate dec, VName, dec) + } deriving (Show) -indexSubstitutions :: [LoopResultSummary dec] - -> IndexSubstitutions dec +indexSubstitutions :: + [LoopResultSummary dec] -> + IndexSubstitutions dec indexSubstitutions = mapMaybe getSubstitution - where getSubstitution res = do - (DesiredUpdate _ _ cs _ is _, nm, dec) <- relatedUpdate res - let name = paramName $ fst $ mergeParam res - return (name, (cs, nm, dec, is)) + where + getSubstitution res = do + (DesiredUpdate _ _ cs _ is _, nm, dec) <- relatedUpdate res + let name = paramName $ fst $ mergeParam res + return (name, (cs, nm, dec, is)) -manipulateResult :: (Bindable lore, MonadFreshNames m) => - [LoopResultSummary (LetDec lore)] - -> IndexSubstitutions (LetDec lore) - -> m (Result, Stms lore) +manipulateResult :: + (Bindable lore, MonadFreshNames m) => + [LoopResultSummary (LetDec lore)] -> + IndexSubstitutions (LetDec lore) -> + m (Result, Stms lore) manipulateResult summaries substs = do - let (orig_ses,updated_ses) = partitionEithers $ map unchangedRes summaries + let (orig_ses, updated_ses) = partitionEithers $ map unchangedRes summaries (subst_ses, res_bnds) <- runWriterT $ zipWithM substRes updated_ses substs return (orig_ses ++ subst_ses, stmsFromList res_bnds) where unchangedRes summary = case relatedUpdate summary of Nothing -> Left $ resultSubExp summary - Just _ -> Right $ resultSubExp summary + Just _ -> Right $ resultSubExp summary substRes (Var res_v) (subst_v, (_, nm, _, _)) | res_v == subst_v = return $ Var nm substRes res_se (_, (cs, nm, dec, is)) = do - v' <- newIdent' (++"_updated") $ Ident nm $ typeOf dec - tell [certify cs $ mkLet [] [v'] $ BasicOp $ - Update nm (fullSlice (typeOf dec) is) res_se] + v' <- newIdent' (++ "_updated") $ Ident nm $ typeOf dec + tell + [ certify cs $ + mkLet [] [v'] $ + BasicOp $ + Update nm (fullSlice (typeOf dec) is) res_se + ] return $ Var $ identName v' diff --git a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs --- a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs +++ b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs @@ -1,136 +1,176 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + -- | This module exports facilities for transforming array accesses in -- a list of 'Stm's (intended to be the bindings in a body). The -- idea is that you can state that some variable @x@ is in fact an -- array indexing @v[i0,i1,...]@. module Futhark.Optimise.InPlaceLowering.SubstituteIndices - ( - substituteIndices - , IndexSubstitution - , IndexSubstitutions - ) where + ( substituteIndices, + IndexSubstitution, + IndexSubstitutions, + ) +where import Control.Monad import qualified Data.Map.Strict as M - -import Futhark.IR.Prop.Aliases -import Futhark.IR import Futhark.Construct +import Futhark.IR +import Futhark.IR.Prop.Aliases import Futhark.Util type IndexSubstitution dec = (Certificates, VName, dec, Slice SubExp) + type IndexSubstitutions dec = [(VName, IndexSubstitution dec)] -typeEnvFromSubstitutions :: LetDec lore ~ dec => - IndexSubstitutions dec -> Scope lore +typeEnvFromSubstitutions :: + LetDec lore ~ dec => + IndexSubstitutions dec -> + Scope lore typeEnvFromSubstitutions = M.fromList . map (fromSubstitution . snd) - where fromSubstitution (_, name, t, _) = - (name, LetName t) + where + fromSubstitution (_, name, t, _) = + (name, LetName t) -- | Perform the substitution. -substituteIndices :: (MonadFreshNames m, BinderOps lore, Bindable lore, - Aliased lore, LetDec lore ~ dec) => - IndexSubstitutions dec -> Stms lore - -> m (IndexSubstitutions dec, Stms lore) +substituteIndices :: + ( MonadFreshNames m, + BinderOps lore, + Bindable lore, + Aliased lore, + LetDec lore ~ dec + ) => + IndexSubstitutions dec -> + Stms lore -> + m (IndexSubstitutions dec, Stms lore) substituteIndices substs bnds = runBinderT (substituteIndicesInStms substs bnds) types - where types = typeEnvFromSubstitutions substs + where + types = typeEnvFromSubstitutions substs -substituteIndicesInStms :: (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => - IndexSubstitutions (LetDec (Lore m)) - -> Stms (Lore m) - -> m (IndexSubstitutions (LetDec (Lore m))) +substituteIndicesInStms :: + (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => + IndexSubstitutions (LetDec (Lore m)) -> + Stms (Lore m) -> + m (IndexSubstitutions (LetDec (Lore m))) substituteIndicesInStms = foldM substituteIndicesInStm -substituteIndicesInStm :: (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => - IndexSubstitutions (LetDec (Lore m)) - -> Stm (Lore m) - -> m (IndexSubstitutions (LetDec (Lore m))) +substituteIndicesInStm :: + (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => + IndexSubstitutions (LetDec (Lore m)) -> + Stm (Lore m) -> + m (IndexSubstitutions (LetDec (Lore m))) substituteIndicesInStm substs (Let pat lore e) = do e' <- substituteIndicesInExp substs e (substs', pat') <- substituteIndicesInPattern substs pat addStm $ Let pat' lore e' return substs' -substituteIndicesInPattern :: (MonadBinder m, LetDec (Lore m) ~ dec) => - IndexSubstitutions (LetDec (Lore m)) - -> PatternT dec - -> m (IndexSubstitutions (LetDec (Lore m)), PatternT dec) +substituteIndicesInPattern :: + (MonadBinder m, LetDec (Lore m) ~ dec) => + IndexSubstitutions (LetDec (Lore m)) -> + PatternT dec -> + m (IndexSubstitutions (LetDec (Lore m)), PatternT dec) substituteIndicesInPattern substs pat = do (substs', context) <- mapAccumLM sub substs $ patternContextElements pat (substs'', values) <- mapAccumLM sub substs' $ patternValueElements pat return (substs'', Pattern context values) - where sub substs' patElem = return (substs', patElem) + where + sub substs' patElem = return (substs', patElem) -substituteIndicesInExp :: (MonadBinder m, Bindable (Lore m), Aliased (Lore m), - LetDec (Lore m) ~ dec) => - IndexSubstitutions (LetDec (Lore m)) - -> Exp (Lore m) - -> m (Exp (Lore m)) +substituteIndicesInExp :: + ( MonadBinder m, + Bindable (Lore m), + Aliased (Lore m), + LetDec (Lore m) ~ dec + ) => + IndexSubstitutions (LetDec (Lore m)) -> + Exp (Lore m) -> + m (Exp (Lore m)) substituteIndicesInExp substs e = do substs' <- copyAnyConsumed e - let substitute = identityMapper { mapOnSubExp = substituteIndicesInSubExp substs' - , mapOnVName = substituteIndicesInVar substs' - , mapOnBody = const $ substituteIndicesInBody substs' - } + let substitute = + identityMapper + { mapOnSubExp = substituteIndicesInSubExp substs', + mapOnVName = substituteIndicesInVar substs', + mapOnBody = const $ substituteIndicesInBody substs' + } mapExpM substitute e - where copyAnyConsumed = - let consumingSubst substs' v - | Just (cs2, src2, src2dec, is2) <- lookup v substs = do - row <- certifying cs2 $ - letExp (baseString v ++ "_row") $ - BasicOp $ Index src2 $ fullSlice (typeOf src2dec) is2 - row_copy <- letExp (baseString v ++ "_row_copy") $ - BasicOp $ Copy row - return $ update v v (mempty, - row_copy, - src2dec `setType` - (typeOf src2dec `setArrayDims` - sliceDims is2), - []) substs' - consumingSubst substs' _ = - return substs' - in foldM consumingSubst substs . namesToList . consumedInExp + where + copyAnyConsumed = + let consumingSubst substs' v + | Just (cs2, src2, src2dec, is2) <- lookup v substs = do + row <- + certifying cs2 $ + letExp (baseString v ++ "_row") $ + BasicOp $ Index src2 $ fullSlice (typeOf src2dec) is2 + row_copy <- + letExp (baseString v ++ "_row_copy") $ + BasicOp $ Copy row + return $ + update + v + v + ( mempty, + row_copy, + src2dec + `setType` ( typeOf src2dec + `setArrayDims` sliceDims is2 + ), + [] + ) + substs' + consumingSubst substs' _ = + return substs' + in foldM consumingSubst substs . namesToList . consumedInExp -substituteIndicesInSubExp :: MonadBinder m => - IndexSubstitutions (LetDec (Lore m)) - -> SubExp - -> m SubExp +substituteIndicesInSubExp :: + MonadBinder m => + IndexSubstitutions (LetDec (Lore m)) -> + SubExp -> + m SubExp substituteIndicesInSubExp substs (Var v) = Var <$> substituteIndicesInVar substs v substituteIndicesInSubExp _ se = return se -substituteIndicesInVar :: MonadBinder m => - IndexSubstitutions (LetDec (Lore m)) - -> VName - -> m VName +substituteIndicesInVar :: + MonadBinder m => + IndexSubstitutions (LetDec (Lore m)) -> + VName -> + m VName substituteIndicesInVar substs v | Just (cs2, src2, _, []) <- lookup v substs = certifying cs2 $ - letExp (baseString src2) $ BasicOp $ SubExp $ Var src2 + letExp (baseString src2) $ BasicOp $ SubExp $ Var src2 | Just (cs2, src2, src2_dec, is2) <- lookup v substs = certifying cs2 $ - letExp "idx" $ BasicOp $ Index src2 $ fullSlice (typeOf src2_dec) is2 + letExp "idx" $ BasicOp $ Index src2 $ fullSlice (typeOf src2_dec) is2 | otherwise = return v -substituteIndicesInBody :: (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => - IndexSubstitutions (LetDec (Lore m)) - -> Body (Lore m) - -> m (Body (Lore m)) +substituteIndicesInBody :: + (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => + IndexSubstitutions (LetDec (Lore m)) -> + Body (Lore m) -> + m (Body (Lore m)) substituteIndicesInBody substs (Body _ stms res) = do - (substs', stms') <- inScopeOf stms $ - collectStms $ substituteIndicesInStms substs stms - (res', res_stms) <- inScopeOf stms' $ - collectStms $ mapM (substituteIndicesInSubExp substs') res - mkBodyM (stms'<>res_stms) res' + (substs', stms') <- + inScopeOf stms $ + collectStms $ substituteIndicesInStms substs stms + (res', res_stms) <- + inScopeOf stms' $ + collectStms $ mapM (substituteIndicesInSubExp substs') res + mkBodyM (stms' <> res_stms) res' -update :: VName -> VName -> IndexSubstitution dec -> IndexSubstitutions dec - -> IndexSubstitutions dec +update :: + VName -> + VName -> + IndexSubstitution dec -> + IndexSubstitutions dec -> + IndexSubstitutions dec update needle name subst ((othername, othersubst) : substs) - | needle == othername = (name, subst) : substs - | otherwise = (othername, othersubst) : update needle name subst substs -update needle _ _ [] = error $ "Cannot find substitution for " ++ pretty needle + | needle == othername = (name, subst) : substs + | otherwise = (othername, othersubst) : update needle name subst substs +update needle _ _ [] = error $ "Cannot find substitution for " ++ pretty needle diff --git a/src/Futhark/Optimise/InliningDeadFun.hs b/src/Futhark/Optimise/InliningDeadFun.hs --- a/src/Futhark/Optimise/InliningDeadFun.hs +++ b/src/Futhark/Optimise/InliningDeadFun.hs @@ -1,33 +1,38 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} + -- | This module implements a compiler pass for inlining functions, -- then removing those that have become dead. module Futhark.Optimise.InliningDeadFun - ( inlineFunctions - , removeDeadFunctions + ( inlineFunctions, + removeDeadFunctions, ) - where +where import Control.Monad.Identity import Control.Monad.State +import Control.Parallel.Strategies import Data.List (partition) -import Data.Maybe import qualified Data.Map.Strict as M +import Data.Maybe import qualified Data.Set as S -import Control.Parallel.Strategies - +import Futhark.Analysis.CallGraph +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Binder import Futhark.IR.SOACS import Futhark.IR.SOACS.Simplify - (simpleSOACS, simplifyFun, simplifyConsts) + ( simpleSOACS, + simplifyConsts, + simplifyFun, + ) import Futhark.Optimise.CSE import Futhark.Optimise.Simplify.Lore (addScopeWisdom) +import Futhark.Pass import Futhark.Transform.CopyPropagate - (copyPropagateInProg, copyPropagateInFun) -import qualified Futhark.Analysis.SymbolTable as ST + ( copyPropagateInFun, + copyPropagateInProg, + ) import Futhark.Transform.Rename -import Futhark.Analysis.CallGraph -import Futhark.Binder -import Futhark.Pass parMapM :: MonadFreshNames m => (a -> State VNameSource b) -> [a] -> m [b] -- The special-casing of [] is quite important here! If 'as' is @@ -36,154 +41,175 @@ parMapM _ [] = pure [] parMapM f as = modifyNameSource $ \src -> - let f' a = runState (f a) src - (bs, srcs) = unzip $ parMap rpar f' as - in (bs, mconcat srcs) + let f' a = runState (f a) src + (bs, srcs) = unzip $ parMap rpar f' as + in (bs, mconcat srcs) aggInlineFunctions :: MonadFreshNames m => Prog SOACS -> m (Prog SOACS) aggInlineFunctions prog = let Prog consts funs = prog - in uncurry Prog . fmap (filter keep) <$> - recurse 0 (ST.fromScope (addScopeWisdom (scopeOf consts)), consts, funs) - where fdmap fds = - M.fromList $ zip (map funDefName fds) fds - - cg = buildCallGraph prog - noninlined = findNoninlined prog + in uncurry Prog . fmap (filter keep) + <$> recurse 0 (ST.fromScope (addScopeWisdom (scopeOf consts)), consts, funs) + where + fdmap fds = + M.fromList $ zip (map funDefName fds) fds - noCallsTo which fundec = - not $ any (`S.member` which) $ allCalledBy (funDefName fundec) cg + cg = buildCallGraph prog + noninlined = findNoninlined prog - -- The inverse rate at which we perform full simplification - -- after inlining. For the other steps we just do copy - -- propagation. The rate here has been determined - -- heuristically and is probably not optimal for any given - -- program. - simplifyRate :: Int - simplifyRate = 4 + noCallsTo which fundec = + not $ any (`S.member` which) $ allCalledBy (funDefName fundec) cg - -- We apply simplification after every round of inlining, - -- because it is more efficient to shrink the program as soon - -- as possible, rather than wait until it has balooned after - -- full inlining. - recurse i (vtable, consts, funs) = do - let remaining = S.fromList $ map funDefName funs - (to_be_inlined, maybe_inline_in) = - partition (noCallsTo remaining) funs - (not_to_inline_in, to_inline_in) = - partition (noCallsTo - (S.fromList $ map funDefName to_be_inlined)) - maybe_inline_in - keep_although_inlined = filter keep to_be_inlined - if null to_be_inlined - then return (consts, funs) - else do + -- The inverse rate at which we perform full simplification + -- after inlining. For the other steps we just do copy + -- propagation. The rate here has been determined + -- heuristically and is probably not optimal for any given + -- program. + simplifyRate :: Int + simplifyRate = 4 - (vtable', consts') <- - if any ((`calledByConsts` cg) . funDefName) to_be_inlined - then simplifyConsts . performCSEOnStms True =<< - inlineInStms (fdmap to_be_inlined) consts + -- We apply simplification after every round of inlining, + -- because it is more efficient to shrink the program as soon + -- as possible, rather than wait until it has balooned after + -- full inlining. + recurse i (vtable, consts, funs) = do + let remaining = S.fromList $ map funDefName funs + (to_be_inlined, maybe_inline_in) = + partition (noCallsTo remaining) funs + (not_to_inline_in, to_inline_in) = + partition + ( noCallsTo + (S.fromList $ map funDefName to_be_inlined) + ) + maybe_inline_in + keep_although_inlined = filter keep to_be_inlined + if null to_be_inlined + then return (consts, funs) + else do + (vtable', consts') <- + if any ((`calledByConsts` cg) . funDefName) to_be_inlined + then + simplifyConsts . performCSEOnStms True + =<< inlineInStms (fdmap to_be_inlined) consts else pure (vtable, consts) - let simplifyFun' fd - | i `rem` simplifyRate == 0 = - copyPropagateInFun simpleSOACS vtable' . - performCSEOnFunDef True =<< - simplifyFun vtable' fd - | otherwise = - copyPropagateInFun simpleSOACS vtable' fd + let simplifyFun' fd + | i `rem` simplifyRate == 0 = + copyPropagateInFun simpleSOACS vtable' + . performCSEOnFunDef True + =<< simplifyFun vtable' fd + | otherwise = + copyPropagateInFun simpleSOACS vtable' fd - let onFun = simplifyFun' <=< - inlineInFunDef (fdmap to_be_inlined) - to_inline_in' <- parMapM onFun to_inline_in - fmap (keep_although_inlined<>) <$> - recurse (i+1) + let onFun = + simplifyFun' + <=< inlineInFunDef (fdmap to_be_inlined) + to_inline_in' <- parMapM onFun to_inline_in + fmap (keep_although_inlined <>) + <$> recurse + (i + 1) (vtable', consts', not_to_inline_in <> to_inline_in') - keep fd = - isJust (funDefEntryPoint fd) || - funDefName fd `S.member` noninlined + keep fd = + isJust (funDefEntryPoint fd) + || funDefName fd `S.member` noninlined -- | @inlineInFunDef constf fdmap caller@ inlines in @calleer@ the -- functions in @fdmap@ that are called as @constf@. At this point the -- preconditions are that if @fdmap@ is not empty, and, more -- importantly, the functions in @fdmap@ do not call any other -- functions. -inlineInFunDef :: MonadFreshNames m => - M.Map Name (FunDef SOACS) -> FunDef SOACS - -> m (FunDef SOACS) +inlineInFunDef :: + MonadFreshNames m => + M.Map Name (FunDef SOACS) -> + FunDef SOACS -> + m (FunDef SOACS) inlineInFunDef fdmap (FunDef entry attrs name rtp args body) = FunDef entry attrs name rtp args <$> inlineInBody fdmap body -inlineFunction :: MonadFreshNames m => - Pattern - -> StmAux dec - -> [(SubExp, Diet)] - -> (Safety, SrcLoc, [SrcLoc]) - -> FunDef SOACS - -> m [Stm] -inlineFunction pat aux args (safety,loc,locs) fun = do +inlineFunction :: + MonadFreshNames m => + Pattern -> + StmAux dec -> + [(SubExp, Diet)] -> + (Safety, SrcLoc, [SrcLoc]) -> + FunDef SOACS -> + m [Stm] +inlineFunction pat aux args (safety, loc, locs) fun = do Body _ stms res <- - renameBody $ mkBody - (stmsFromList param_stms <> stmsFromList body_stms) - (bodyResult (funDefBody fun)) + renameBody $ + mkBody + (stmsFromList param_stms <> stmsFromList body_stms) + (bodyResult (funDefBody fun)) let res_stms = - certify (stmAuxCerts aux) <$> - zipWith bindSubExp (patternIdents pat) res + certify (stmAuxCerts aux) + <$> zipWith bindSubExp (patternIdents pat) res pure $ stmsToList stms <> res_stms - where param_stms = - zipWith bindSubExp - (map paramIdent $ funDefParams fun) (map fst args) + where + param_stms = + zipWith + bindSubExp + (map paramIdent $ funDefParams fun) + (map fst args) - body_stms = - stmsToList $ - addLocations (stmAuxAttrs aux) safety (filter notmempty (loc:locs)) $ + body_stms = + stmsToList $ + addLocations (stmAuxAttrs aux) safety (filter notmempty (loc : locs)) $ bodyStms $ funDefBody fun - -- Note that the sizes of arrays may not be correct at this - -- point - it is crucial that we run copy propagation before - -- the type checker sees this! - bindSubExp ident se = - mkLet [] [ident] $ BasicOp $ SubExp se + -- Note that the sizes of arrays may not be correct at this + -- point - it is crucial that we run copy propagation before + -- the type checker sees this! + bindSubExp ident se = + mkLet [] [ident] $ BasicOp $ SubExp se - notmempty = (/=mempty) . locOf + notmempty = (/= mempty) . locOf -inlineInStms :: MonadFreshNames m => - M.Map Name (FunDef SOACS) -> Stms SOACS -> m (Stms SOACS) +inlineInStms :: + MonadFreshNames m => + M.Map Name (FunDef SOACS) -> + Stms SOACS -> + m (Stms SOACS) inlineInStms fdmap stms = bodyStms <$> inlineInBody fdmap (mkBody stms []) -inlineInBody :: MonadFreshNames m => - M.Map Name (FunDef SOACS) -> Body -> m Body +inlineInBody :: + MonadFreshNames m => + M.Map Name (FunDef SOACS) -> + Body -> + m Body inlineInBody fdmap = onBody - where inline (Let pat aux (Apply fname args _ what) : rest) - | Just fd <- M.lookup fname fdmap, - not $ "noinline" `inAttrs` funDefAttrs fd, - not $ "noinline" `inAttrs` stmAuxAttrs aux = - (<>) <$> inlineFunction pat aux args what fd <*> inline rest - - inline (stm : rest) = - (:) <$> onStm stm <*> inline rest - inline [] = - pure mempty + where + inline (Let pat aux (Apply fname args _ what) : rest) + | Just fd <- M.lookup fname fdmap, + not $ "noinline" `inAttrs` funDefAttrs fd, + not $ "noinline" `inAttrs` stmAuxAttrs aux = + (<>) <$> inlineFunction pat aux args what fd <*> inline rest + inline (stm : rest) = + (:) <$> onStm stm <*> inline rest + inline [] = + pure mempty - onBody (Body dec stms res) = - Body dec . stmsFromList <$> inline (stmsToList stms) <*> pure res + onBody (Body dec stms res) = + Body dec . stmsFromList <$> inline (stmsToList stms) <*> pure res - onStm (Let pat aux e) = - Let pat aux <$> mapExpM inliner e + onStm (Let pat aux e) = + Let pat aux <$> mapExpM inliner e - inliner = - identityMapper { mapOnBody = const onBody - , mapOnOp = onSOAC - } + inliner = + identityMapper + { mapOnBody = const onBody, + mapOnOp = onSOAC + } - onSOAC = - mapSOACM identitySOACMapper - { mapOnSOACLambda = onLambda } + onSOAC = + mapSOACM + identitySOACMapper + { mapOnSOACLambda = onLambda + } - onLambda (Lambda params body ret) = - Lambda params <$> onBody body <*> pure ret + onLambda (Lambda params body ret) = + Lambda params <$> onBody body <*> pure ret -- Propagate source locations and attributes to the inlined -- statements. Attributes are propagated only when applicable (this @@ -191,53 +217,71 @@ -- specially here). addLocations :: Attrs -> Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS addLocations attrs caller_safety more_locs = fmap onStm - where onStm (Let pat aux (Apply fname args t (safety, loc,locs))) = - Let pat aux' $ - Apply fname args t (min caller_safety safety, loc,locs++more_locs) - where aux' = aux { stmAuxAttrs = attrs <> stmAuxAttrs aux } - onStm (Let pat aux (BasicOp (Assert cond desc (loc,locs)))) = - Let pat (withAttrs (attrsForAssert attrs) aux) $ - case caller_safety of - Safe -> BasicOp $ Assert cond desc (loc,locs++more_locs) - Unsafe -> BasicOp $ SubExp $ Constant Checked - onStm (Let pat aux (Op soac)) = - Let pat (withAttrs attrs' aux) $ Op $ runIdentity $ mapSOACM - identitySOACMapper { mapOnSOACLambda = return . onLambda - } soac - where attrs' = attrs `withoutAttrs` for_assert - for_assert = attrsForAssert attrs - onLambda lam = - lam { lambdaBody = onBody for_assert $ lambdaBody lam } - onStm (Let pat aux e) = - Let pat aux $ onExp e + where + onStm (Let pat aux (Apply fname args t (safety, loc, locs))) = + Let pat aux' $ + Apply fname args t (min caller_safety safety, loc, locs ++ more_locs) + where + aux' = aux {stmAuxAttrs = attrs <> stmAuxAttrs aux} + onStm (Let pat aux (BasicOp (Assert cond desc (loc, locs)))) = + Let pat (withAttrs (attrsForAssert attrs) aux) $ + case caller_safety of + Safe -> BasicOp $ Assert cond desc (loc, locs ++ more_locs) + Unsafe -> BasicOp $ SubExp $ Constant Checked + onStm (Let pat aux (Op soac)) = + Let pat (withAttrs attrs' aux) $ + Op $ + runIdentity $ + mapSOACM + identitySOACMapper + { mapOnSOACLambda = return . onLambda + } + soac + where + attrs' = attrs `withoutAttrs` for_assert + for_assert = attrsForAssert attrs + onLambda lam = + lam {lambdaBody = onBody for_assert $ lambdaBody lam} + onStm (Let pat aux e) = + Let pat aux $ onExp e - onExp = mapExp identityMapper - { mapOnBody = const $ return . onBody attrs } + onExp = + mapExp + identityMapper + { mapOnBody = const $ return . onBody attrs + } - withAttrs attrs' aux = aux { stmAuxAttrs = attrs' <> stmAuxAttrs aux } + withAttrs attrs' aux = aux {stmAuxAttrs = attrs' <> stmAuxAttrs aux} - onBody attrs' body = - body { bodyStms = addLocations attrs' caller_safety more_locs $ - bodyStms body } + onBody attrs' body = + body + { bodyStms = + addLocations attrs' caller_safety more_locs $ + bodyStms body + } -- | Inline all functions and remove the resulting dead functions. inlineFunctions :: Pass SOACS SOACS inlineFunctions = - Pass { passName = "Inline functions" - , passDescription = "Inline and remove resulting dead functions." - , passFunction = copyPropagateInProg simpleSOACS <=< aggInlineFunctions - } + Pass + { passName = "Inline functions", + passDescription = "Inline and remove resulting dead functions.", + passFunction = copyPropagateInProg simpleSOACS <=< aggInlineFunctions + } -- | @removeDeadFunctions prog@ removes the functions that are unreachable from -- the main function from the program. removeDeadFunctions :: Pass SOACS SOACS removeDeadFunctions = - Pass { passName = "Remove dead functions" - , passDescription = "Remove the functions that are unreachable from entry points" - , passFunction = return . pass - } - where pass prog = - let cg = buildCallGraph prog - live_funs = filter ((`isFunInCallGraph` cg) . funDefName) $ - progFuns prog - in prog { progFuns = live_funs } + Pass + { passName = "Remove dead functions", + passDescription = "Remove the functions that are unreachable from entry points", + passFunction = return . pass + } + where + pass prog = + let cg = buildCallGraph prog + live_funs = + filter ((`isFunInCallGraph` cg) . funDefName) $ + progFuns prog + in prog {progFuns = live_funs} diff --git a/src/Futhark/Optimise/Simplify.hs b/src/Futhark/Optimise/Simplify.hs --- a/src/Futhark/Optimise/Simplify.hs +++ b/src/Futhark/Optimise/Simplify.hs @@ -1,49 +1,51 @@ -{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE TupleSections #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE Strict #-} -module Futhark.Optimise.Simplify - ( simplifyProg - , simplifySomething - , simplifyFun - , simplifyLambda - , simplifyStms +{-# LANGUAGE TupleSections #-} - , Engine.SimpleOps (..) - , Engine.SimpleM - , Engine.SimplifyOp - , Engine.bindableSimpleOps - , Engine.noExtraHoistBlockers - , Engine.neverHoist - , Engine.SimplifiableLore - , Engine.HoistBlockers - , RuleBook +module Futhark.Optimise.Simplify + ( simplifyProg, + simplifySomething, + simplifyFun, + simplifyLambda, + simplifyStms, + Engine.SimpleOps (..), + Engine.SimpleM, + Engine.SimplifyOp, + Engine.bindableSimpleOps, + Engine.noExtraHoistBlockers, + Engine.neverHoist, + Engine.SimplifiableLore, + Engine.HoistBlockers, + RuleBook, ) - where +where import Data.Bifunctor (second) +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT import Futhark.IR import Futhark.MonadFreshNames import qualified Futhark.Optimise.Simplify.Engine as Engine -import qualified Futhark.Analysis.SymbolTable as ST -import qualified Futhark.Analysis.UsageTable as UT -import Futhark.Optimise.Simplify.Rule import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rule import Futhark.Pass -- | Simplify the given program. Even if the output differs from the -- output, meaningful simplification may not have taken place - the -- order of bindings may simply have been rearranged. -simplifyProg :: Engine.SimplifiableLore lore => - Engine.SimpleOps lore - -> RuleBook (Engine.Wise lore) - -> Engine.HoistBlockers lore - -> Prog lore - -> PassM (Prog lore) +simplifyProg :: + Engine.SimplifiableLore lore => + Engine.SimpleOps lore -> + RuleBook (Engine.Wise lore) -> + Engine.HoistBlockers lore -> + Prog lore -> + PassM (Prog lore) simplifyProg simpl rules blockers (Prog consts funs) = do (consts_vtable, consts') <- - simplifyConsts (UT.usages $ foldMap (freeIn . funDefBody) funs) - (mempty, consts) + simplifyConsts + (UT.usages $ foldMap (freeIn . funDefBody) funs) + (mempty, consts) funs' <- parPass (simplifyFun' consts_vtable) funs let funs_uses = UT.usages $ foldMap (freeIn . funDefBody) funs' @@ -51,88 +53,110 @@ (_, consts'') <- simplifyConsts funs_uses (mempty, consts') return $ Prog consts'' funs' - - where simplifyFun' consts_vtable = - simplifySomething - (Engine.localVtable (consts_vtable<>) . Engine.simplifyFun) - removeFunDefWisdom - simpl rules blockers mempty + where + simplifyFun' consts_vtable = + simplifySomething + (Engine.localVtable (consts_vtable <>) . Engine.simplifyFun) + removeFunDefWisdom + simpl + rules + blockers + mempty - simplifyConsts uses = - simplifySomething (onConsts uses . snd) - (second (removeStmWisdom<$>)) - simpl rules blockers mempty + simplifyConsts uses = + simplifySomething + (onConsts uses . snd) + (second (removeStmWisdom <$>)) + simpl + rules + blockers + mempty - onConsts uses consts' = do - (_, consts'') <- - Engine.simplifyStms consts' (pure ((), mempty)) - (consts''', _) <- - Engine.hoistStms rules (Engine.isFalse False) mempty uses consts'' - return (ST.insertStms consts''' mempty, consts''') + onConsts uses consts' = do + (_, consts'') <- + Engine.simplifyStms consts' (pure ((), mempty)) + (consts''', _) <- + Engine.hoistStms rules (Engine.isFalse False) mempty uses consts'' + return (ST.insertStms consts''' mempty, consts''') -- | Run a simplification operation to convergence. -simplifySomething :: (MonadFreshNames m, Engine.SimplifiableLore lore) => - (a -> Engine.SimpleM lore b) - -> (b -> a) - -> Engine.SimpleOps lore - -> RuleBook (Wise lore) - -> Engine.HoistBlockers lore - -> ST.SymbolTable (Wise lore) - -> a - -> m a +simplifySomething :: + (MonadFreshNames m, Engine.SimplifiableLore lore) => + (a -> Engine.SimpleM lore b) -> + (b -> a) -> + Engine.SimpleOps lore -> + RuleBook (Wise lore) -> + Engine.HoistBlockers lore -> + ST.SymbolTable (Wise lore) -> + a -> + m a simplifySomething f g simpl rules blockers vtable x = do - let f' x' = Engine.localVtable (vtable<>) $ f x' + let f' x' = Engine.localVtable (vtable <>) $ f x' loopUntilConvergence env simpl f' g x - where env = Engine.emptyEnv rules blockers + where + env = Engine.emptyEnv rules blockers -- | Simplify the given function. Even if the output differs from the -- output, meaningful simplification may not have taken place - the -- order of bindings may simply have been rearranged. Runs in a loop -- until convergence. -simplifyFun :: (MonadFreshNames m, Engine.SimplifiableLore lore) => - Engine.SimpleOps lore - -> RuleBook (Engine.Wise lore) - -> Engine.HoistBlockers lore - -> ST.SymbolTable (Wise lore) - -> FunDef lore - -> m (FunDef lore) +simplifyFun :: + (MonadFreshNames m, Engine.SimplifiableLore lore) => + Engine.SimpleOps lore -> + RuleBook (Engine.Wise lore) -> + Engine.HoistBlockers lore -> + ST.SymbolTable (Wise lore) -> + FunDef lore -> + m (FunDef lore) simplifyFun = simplifySomething Engine.simplifyFun removeFunDefWisdom -- | Simplify just a single t'Lambda'. -simplifyLambda :: (MonadFreshNames m, HasScope lore m, - Engine.SimplifiableLore lore) => - Engine.SimpleOps lore - -> RuleBook (Engine.Wise lore) - -> Engine.HoistBlockers lore - -> Lambda lore - -> m (Lambda lore) +simplifyLambda :: + ( MonadFreshNames m, + HasScope lore m, + Engine.SimplifiableLore lore + ) => + Engine.SimpleOps lore -> + RuleBook (Engine.Wise lore) -> + Engine.HoistBlockers lore -> + Lambda lore -> + m (Lambda lore) simplifyLambda simpl rules blockers orig_lam = do vtable <- ST.fromScope . addScopeWisdom <$> askScope - simplifySomething Engine.simplifyLambdaNoHoisting - removeLambdaWisdom simpl rules blockers vtable orig_lam + simplifySomething + Engine.simplifyLambdaNoHoisting + removeLambdaWisdom + simpl + rules + blockers + vtable + orig_lam -- | Simplify a sequence of 'Stm's. -simplifyStms :: (MonadFreshNames m, Engine.SimplifiableLore lore) => - Engine.SimpleOps lore - -> RuleBook (Engine.Wise lore) - -> Engine.HoistBlockers lore - -> Scope lore - -> Stms lore - -> m (ST.SymbolTable (Wise lore), Stms lore) +simplifyStms :: + (MonadFreshNames m, Engine.SimplifiableLore lore) => + Engine.SimpleOps lore -> + RuleBook (Engine.Wise lore) -> + Engine.HoistBlockers lore -> + Scope lore -> + Stms lore -> + m (ST.SymbolTable (Wise lore), Stms lore) simplifyStms simpl rules blockers scope = simplifySomething f g simpl rules blockers vtable . (mempty,) - where vtable = ST.fromScope $ addScopeWisdom scope - f (_, stms) = - Engine.simplifyStms stms ((,mempty) <$> Engine.askVtable) - g = second $ fmap removeStmWisdom + where + vtable = ST.fromScope $ addScopeWisdom scope + f (_, stms) = + Engine.simplifyStms stms ((,mempty) <$> Engine.askVtable) + g = second $ fmap removeStmWisdom -loopUntilConvergence :: (MonadFreshNames m, Engine.SimplifiableLore lore) => - Engine.Env lore - -> Engine.SimpleOps lore - -> (a -> Engine.SimpleM lore b) - -> (b -> a) - -> a - -> m a +loopUntilConvergence :: + (MonadFreshNames m, Engine.SimplifiableLore lore) => + Engine.Env lore -> + Engine.SimpleOps lore -> + (a -> Engine.SimpleM lore b) -> + (b -> a) -> + a -> + m a loopUntilConvergence env simpl f g x = do (x', changed) <- modifyNameSource $ Engine.runSimpleM (f x) simpl env if changed then loopUntilConvergence env simpl f g (g x') else return $ g x' diff --git a/src/Futhark/Optimise/Simplify/ClosedForm.hs b/src/Futhark/Optimise/Simplify/ClosedForm.hs --- a/src/Futhark/Optimise/Simplify/ClosedForm.hs +++ b/src/Futhark/Optimise/Simplify/ClosedForm.hs @@ -1,5 +1,6 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE Safe #-} + -- | This module implements facilities for determining whether a -- reduction or fold can be expressed in a closed form (i.e. not as a -- SOAC). @@ -8,20 +9,19 @@ -- 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 + ( foldClosedForm, + loopClosedForm, + VarLookup, ) where import Control.Monad -import Data.Maybe import qualified Data.Map.Strict as M - +import Data.Maybe import Futhark.Construct import Futhark.IR -import Futhark.Transform.Rename 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) @@ -42,137 +42,186 @@ -- | @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 :: + (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 + t <- case patternTypes pat of + [Prim t] -> return t + _ -> cannotSimplify - closedBody <- checkResults (patternNames pat) inputsize mempty knownBnds - (map paramName (lambdaParams lam)) - (lambdaBody lam) accs + 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 int32) inputsize (intConst Int32 0) - letBind pat =<< (If (Var isEmpty) - <$> resultBodyM accs - <*> renameBody closedBody - <*> pure (IfDec [primBodyType t] IfNormal)) - where knownBnds = determineKnownBindings look lam accs arrs + 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 -> SubExp -> Body lore - -> RuleM lore () -loopClosedForm pat merge i bound body = do - t <- case patternTypes pat of [Prim t] -> return t - _ -> cannotSimplify +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 knownBnds - (map identName mergeidents) body mergeexp + closedBody <- + checkResults + mergenames + bound + i + it + knownBnds + (map identName mergeidents) + body + mergeexp isEmpty <- newVName "bound_is_zero" letBindNames [isEmpty] $ - BasicOp $ CmpOp (CmpSlt Int32) bound (intConst Int32 0) + 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 + 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 - -> M.Map VName SubExp - -> [VName] -- ^ Lambda-bound - -> Body lore - -> [SubExp] - -> RuleM lore (Body lore) -checkResults pat size untouchable knownBnds params body accs = do - ((), bnds) <- collectStms $ - zipWithM_ checkResult (zip pat res) (zip accparams accs) +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 + where + bndMap = makeBindMap body + (accparams, _) = splitAt (length accs) params + res = bodyResult body - 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 + nonFree = boundInBody body <> namesFromList params <> untouchable - 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 (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 - checkResult _ _ = cannotSimplify + 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 + asFreeSubExp :: SubExp -> Maybe SubExp + asFreeSubExp (Var v) + | v `nameIn` nonFree = M.lookup v knownBnds + asFreeSubExp se = Just se - properIntSize Int32 = Just $ return size - properIntSize t = Just $ letSubExp "converted_size" $ - BasicOp $ ConvOp (SExt Int32 t) size + 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 Int32 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 :: + 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 + 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 + 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 + where + isSingletonStm (Let pat _ e) = case patternNames pat of + [v] -> Just (v, e) + _ -> Nothing diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs --- a/src/Futhark/Optimise/Simplify/Engine.hs +++ b/src/Futhark/Optimise/Simplify/Engine.hs @@ -1,885 +1,1035 @@ -{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, TypeFamilies, FlexibleContexts #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE Strict #-} --- | --- --- Perform general rule-based simplification based on data dependency --- information. This module will: --- --- * Perform common-subexpression elimination (CSE). --- --- * Hoist expressions out of loops (including lambdas) and --- branches. This is done as aggressively as possible. --- --- * Apply simplification rules (see --- "Futhark.Optimise.Simplification.Rules"). --- --- If you just want to run the simplifier as simply as possible, you --- may prefer to use the "Futhark.Optimise.Simplify" module. --- -module Futhark.Optimise.Simplify.Engine - ( -- * Monadic interface - SimpleM - , runSimpleM - , SimpleOps (..) - , SimplifyOp - , bindableSimpleOps - - , Env (envHoistBlockers, envRules) - , emptyEnv - , HoistBlockers(..) - , neverBlocks - , noExtraHoistBlockers - , neverHoist - , BlockPred - , orIf - , hasFree - , isConsumed - , isFalse - , isOp - , isNotSafe - , asksEngineEnv - , askVtable - , localVtable - - -- * Building blocks - , SimplifiableLore - , Simplifiable (..) - , simplifyStms - , simplifyFun - , simplifyLambda - , simplifyLambdaNoHoisting - , bindLParams - , simplifyBody - , SimplifiedBody - , ST.SymbolTable - - , hoistStms - , blockIf - - , module Futhark.Optimise.Simplify.Lore - ) where - -import Control.Monad.Writer -import Control.Monad.Reader -import Control.Monad.State.Strict -import Data.Either -import Data.List (find, foldl', nub, mapAccumL) -import Data.Maybe - -import Futhark.IR -import Futhark.IR.Prop.Aliases -import Futhark.Optimise.Simplify.Rule -import qualified Futhark.Analysis.SymbolTable as ST -import qualified Futhark.Analysis.UsageTable as UT -import Futhark.Construct -import Futhark.Optimise.Simplify.Lore -import Futhark.Util (splitFromEnd) - -data HoistBlockers lore = HoistBlockers - { blockHoistPar :: BlockPred (Wise lore) - -- ^ Blocker for hoisting out of parallel loops. - , blockHoistSeq :: BlockPred (Wise lore) - -- ^ Blocker for hoisting out of sequential loops. - , blockHoistBranch :: BlockPred (Wise lore) - -- ^ Blocker for hoisting out of branches. - , isAllocation :: Stm (Wise lore) -> Bool - } - -noExtraHoistBlockers :: HoistBlockers lore -noExtraHoistBlockers = - HoistBlockers neverBlocks neverBlocks neverBlocks (const False) - -neverHoist :: HoistBlockers lore -neverHoist = - HoistBlockers alwaysBlocks alwaysBlocks alwaysBlocks (const False) - -data Env lore = Env { envRules :: RuleBook (Wise lore) - , envHoistBlockers :: HoistBlockers lore - , envVtable :: ST.SymbolTable (Wise lore) - } - -emptyEnv :: RuleBook (Wise lore) -> HoistBlockers lore -> Env lore -emptyEnv rules blockers = - Env { envRules = rules - , envHoistBlockers = blockers - , envVtable = mempty - } - -type Protect m = SubExp -> Pattern (Lore m) -> Op (Lore m) -> Maybe (m ()) - -data SimpleOps lore = - SimpleOps { mkExpDecS :: ST.SymbolTable (Wise lore) - -> Pattern (Wise lore) -> Exp (Wise lore) - -> SimpleM lore (ExpDec (Wise lore)) - , mkBodyS :: ST.SymbolTable (Wise lore) - -> Stms (Wise lore) -> Result - -> SimpleM lore (Body (Wise lore)) - , protectHoistedOpS :: Protect (Binder (Wise lore)) - -- ^ Make a hoisted Op safe. The SubExp is a boolean - -- that is true when the value of the statement will - -- actually be used. - , opUsageS :: Op (Wise lore) -> UT.UsageTable - , simplifyOpS :: SimplifyOp lore (Op lore) - } - -type SimplifyOp lore op = op -> SimpleM lore (OpWithWisdom op, Stms (Wise lore)) - -bindableSimpleOps :: (SimplifiableLore lore, Bindable lore) => - SimplifyOp lore (Op lore) -> SimpleOps lore -bindableSimpleOps = - SimpleOps mkExpDecS' mkBodyS' protectHoistedOpS' (const mempty) - where mkExpDecS' _ pat e = return $ mkExpDec pat e - mkBodyS' _ bnds res = return $ mkBody bnds res - protectHoistedOpS' _ _ _ = Nothing - -newtype SimpleM lore a = - SimpleM (ReaderT (SimpleOps lore, Env lore) - (State (VNameSource, Bool, Certificates)) a) - deriving (Applicative, Functor, Monad, - MonadReader (SimpleOps lore, Env lore), - MonadState (VNameSource, Bool, Certificates)) - -instance MonadFreshNames (SimpleM lore) where - putNameSource src = modify $ \(_, b, c) -> (src, b, c) - getNameSource = gets $ \(a, _, _) -> a - -instance SimplifiableLore lore => HasScope (Wise lore) (SimpleM lore) where - askScope = ST.toScope <$> askVtable - lookupType name = do - vtable <- askVtable - case ST.lookupType name vtable of - Just t -> return t - Nothing -> error $ - "SimpleM.lookupType: cannot find variable " ++ - pretty name ++ " in symbol table." - -instance SimplifiableLore lore => - LocalScope (Wise lore) (SimpleM lore) where - localScope types = localVtable (<>ST.fromScope types) - -runSimpleM :: SimpleM lore a - -> SimpleOps lore - -> Env lore - -> VNameSource - -> ((a, Bool), VNameSource) -runSimpleM (SimpleM m) simpl env src = - let (x, (src', b, _)) = runState (runReaderT m (simpl, env)) (src, False, mempty) - in ((x, b), src') - -askEngineEnv :: SimpleM lore (Env lore) -askEngineEnv = asks snd - -asksEngineEnv :: (Env lore -> a) -> SimpleM lore a -asksEngineEnv f = f <$> askEngineEnv - -askVtable :: SimpleM lore (ST.SymbolTable (Wise lore)) -askVtable = asksEngineEnv envVtable - -localVtable :: (ST.SymbolTable (Wise lore) -> ST.SymbolTable (Wise lore)) - -> SimpleM lore a -> SimpleM lore a -localVtable f = local $ \(ops, env) -> (ops, env { envVtable = f $ envVtable env }) - -collectCerts :: SimpleM lore a -> SimpleM lore (a, Certificates) -collectCerts m = do x <- m - (a, b, cs) <- get - put (a, b, mempty) - return (x, cs) - --- | Mark that we have changed something and it would be a good idea --- to re-run the simplifier. -changed :: SimpleM lore () -changed = modify $ \(src, _, cs) -> (src, True, cs) - -usedCerts :: Certificates -> SimpleM lore () -usedCerts cs = modify $ \(a, b, c) -> (a, b, cs <> c) - -enterLoop :: SimpleM lore a -> SimpleM lore a -enterLoop = localVtable ST.deepen - -bindFParams :: SimplifiableLore lore => [FParam (Wise lore)] -> SimpleM lore a -> SimpleM lore a -bindFParams params = - localVtable $ ST.insertFParams params - -bindLParams :: SimplifiableLore lore => [LParam (Wise lore)] -> SimpleM lore a -> SimpleM lore a -bindLParams params = - localVtable $ \vtable -> foldr ST.insertLParam vtable params - -bindArrayLParams :: SimplifiableLore lore => [LParam (Wise lore)] -> SimpleM lore a - -> SimpleM lore a -bindArrayLParams params = - localVtable $ \vtable -> foldl' (flip ST.insertLParam) vtable params - -bindLoopVar :: SimplifiableLore lore => VName -> IntType -> SubExp -> SimpleM lore a -> SimpleM lore a -bindLoopVar var it bound = - localVtable $ ST.insertLoopVar var it bound - --- | We are willing to hoist potentially unsafe statements out of --- branches, but they most be protected by adding a branch on top of --- them. (This means such hoisting is not worth it unless they are in --- turn hoisted out of a loop somewhere.) -protectIfHoisted :: SimplifiableLore lore => - SubExp -- ^ Branch condition. - -> Bool -- ^ Which side of the branch are we - -- protecting here? - -> SimpleM lore (a, Stms (Wise lore)) - -> SimpleM lore (a, Stms (Wise lore)) -protectIfHoisted cond side m = do - (x, stms) <- m - ops <- asks $ protectHoistedOpS . fst - runBinder $ do - if not $ all (safeExp . stmExp) stms - then do cond' <- if side then return cond - else letSubExp "cond_neg" $ BasicOp $ UnOp Not cond - mapM_ (protectIf ops unsafeOrCostly cond') stms - else addStms stms - return x - where unsafeOrCostly e = not (safeExp e) || not (cheapExp e) - --- | We are willing to hoist potentially unsafe statements out of --- loops, but they most be protected by adding a branch on top of --- them. -protectLoopHoisted :: SimplifiableLore lore => - [(FParam (Wise lore),SubExp)] - -> [(FParam (Wise lore),SubExp)] - -> LoopForm (Wise lore) - -> SimpleM lore (a, Stms (Wise lore)) - -> SimpleM lore (a, Stms (Wise lore)) -protectLoopHoisted ctx val form m = do - (x, stms) <- m - ops <- asks $ protectHoistedOpS . fst - runBinder $ do - if not $ all (safeExp . stmExp) stms - then do is_nonempty <- checkIfNonEmpty - mapM_ (protectIf ops (not . safeExp) is_nonempty) stms - else addStms stms - return x - where checkIfNonEmpty = - case form of - WhileLoop cond - | Just (_, cond_init) <- - find ((==cond) . paramName . fst) $ ctx ++ val -> - return cond_init - | otherwise -> return $ constant True -- infinite loop - ForLoop _ it bound _ -> - letSubExp "loop_nonempty" $ - BasicOp $ CmpOp (CmpSlt it) (intConst it 0) bound - -protectIf :: MonadBinder m => - Protect m - -> (Exp (Lore m) -> Bool) - -> SubExp -> Stm (Lore m) -> m () -protectIf _ _ taken (Let pat aux - (If cond taken_body untaken_body (IfDec if_ts IfFallback))) = do - cond' <- letSubExp "protect_cond_conj" $ BasicOp $ BinOp LogAnd taken cond - auxing aux $ - letBind pat $ If cond' taken_body untaken_body $ - IfDec if_ts IfFallback -protectIf _ _ taken (Let pat aux (BasicOp (Assert cond msg loc))) = do - not_taken <- letSubExp "loop_not_taken" $ BasicOp $ UnOp Not taken - cond' <- letSubExp "protect_assert_disj" $ BasicOp $ BinOp LogOr not_taken cond - auxing aux $ letBind pat $ BasicOp $ Assert cond' msg loc -protectIf protect _ taken (Let pat aux (Op op)) - | Just m <- protect taken pat op = - auxing aux m -protectIf _ f taken (Let pat aux e) - | f e = - case makeSafe e of - Just e' -> - auxing aux $ letBind pat e' - Nothing -> do - taken_body <- eBody [pure e] - untaken_body <- eBody $ map (emptyOfType $ patternContextNames pat) - (patternValueTypes pat) - if_ts <- expTypesFromPattern pat - auxing aux $ - letBind pat $ If taken taken_body untaken_body $ - IfDec if_ts IfFallback -protectIf _ _ _ stm = - addStm stm - -makeSafe :: Exp lore -> Maybe (Exp lore) -makeSafe (BasicOp (BinOp (SDiv t _) x y)) = - Just $ BasicOp (BinOp (SDiv t Safe) x y) -makeSafe (BasicOp (BinOp (SDivUp t _) x y)) = - Just $ BasicOp (BinOp (SDivUp t Safe) x y) -makeSafe (BasicOp (BinOp (SQuot t _) x y)) = - Just $ BasicOp (BinOp (SQuot t Safe) x y) -makeSafe (BasicOp (BinOp (UDiv t _) x y)) = - Just $ BasicOp (BinOp (UDiv t Safe) x y) -makeSafe (BasicOp (BinOp (UDivUp t _) x y)) = - Just $ BasicOp (BinOp (UDivUp t Safe) x y) -makeSafe (BasicOp (BinOp (SMod t _) x y)) = - Just $ BasicOp (BinOp (SMod t Safe) x y) -makeSafe (BasicOp (BinOp (SRem t _) x y)) = - Just $ BasicOp (BinOp (SRem t Safe) x y) -makeSafe (BasicOp (BinOp (UMod t _) x y)) = - Just $ BasicOp (BinOp (UMod t Safe) x y) -makeSafe _ = - Nothing - -emptyOfType :: MonadBinder m => [VName] -> Type -> m (Exp (Lore m)) -emptyOfType _ Mem{} = - error "emptyOfType: Cannot hoist non-existential memory." -emptyOfType _ (Prim pt) = - return $ BasicOp $ SubExp $ Constant $ blankPrimValue pt -emptyOfType ctx_names (Array pt shape _) = do - let dims = map zeroIfContext $ shapeDims shape - return $ BasicOp $ Scratch pt dims - where zeroIfContext (Var v) | v `elem` ctx_names = intConst Int32 0 - zeroIfContext se = se - --- | Statements that are not worth hoisting out of loops, because they --- are unsafe, and added safety (by 'protectLoopHoisted') may inhibit --- further optimisation.. -notWorthHoisting :: ASTLore lore => BlockPred lore -notWorthHoisting _ _ (Let pat _ e) = - not (safeExp e) && any ((>0) . arrayRank) (patternTypes pat) - -hoistStms :: SimplifiableLore lore => - RuleBook (Wise lore) -> BlockPred (Wise lore) - -> ST.SymbolTable (Wise lore) -> UT.UsageTable - -> Stms (Wise lore) - -> SimpleM lore (Stms (Wise lore), - Stms (Wise lore)) -hoistStms rules block vtable uses orig_stms = do - (blocked, hoisted) <- simplifyStmsBottomUp vtable uses orig_stms - unless (null hoisted) changed - return (stmsFromList blocked, stmsFromList hoisted) - where simplifyStmsBottomUp vtable' uses' stms = do - (_, stms') <- simplifyStmsBottomUp' vtable' uses' stms - -- We need to do a final pass to ensure that nothing is - -- hoisted past something that it depends on. - let (blocked, hoisted) = partitionEithers $ blockUnhoistedDeps stms' - return (blocked, hoisted) - - simplifyStmsBottomUp' vtable' uses' stms = do - opUsage <- asks $ opUsageS . fst - let usageInStm stm = - UT.usageInStm stm <> - case stmExp stm of - Op op -> opUsage op - _ -> mempty - foldM (hoistable usageInStm) (uses',[]) $ reverse $ zip (stmsToList stms) vtables - where vtables = scanl (flip ST.insertStm) vtable' $ stmsToList stms - - hoistable usageInStm (uses',stms) (stm, vtable') - | not $ any (`UT.isUsedDirectly` uses') $ provides stm = -- Dead statement. - return (uses', stms) - | otherwise = do - res <- localVtable (const vtable') $ - bottomUpSimplifyStm rules (vtable', uses') stm - case res of - Nothing -- Nothing to optimise - see if hoistable. - | block vtable' uses' stm -> - return (expandUsage usageInStm vtable' uses' stm - `UT.without` provides stm, - Left stm : stms) - | otherwise -> - return (expandUsage usageInStm vtable' uses' stm, - Right stm : stms) - Just optimstms -> do - changed - (uses'',stms') <- simplifyStmsBottomUp' vtable' uses' optimstms - return (uses'', stms'++stms) - -blockUnhoistedDeps :: ASTLore lore => - [Either (Stm lore) (Stm lore)] - -> [Either (Stm lore) (Stm lore)] -blockUnhoistedDeps = snd . mapAccumL block mempty - where block blocked (Left need) = - (blocked <> namesFromList (provides need), Left need) - block blocked (Right need) - | blocked `namesIntersect` freeIn need = - (blocked <> namesFromList (provides need), Left need) - | otherwise = - (blocked, Right need) - -provides :: Stm lore -> [VName] -provides = patternNames . stmPattern - -expandUsage :: (ASTLore lore, Aliased lore) => - (Stm lore -> UT.UsageTable) -> ST.SymbolTable lore -> UT.UsageTable - -> Stm lore -> UT.UsageTable -expandUsage usageInStm vtable utable stm@(Let pat _ e) = - UT.expand (`ST.lookupAliases` vtable) (usageInStm stm <> usageThroughAliases) <> - (if any (`UT.isSize` utable) (patternNames pat) - then UT.sizeUsages (freeIn e) - else mempty) <> - utable - where usageThroughAliases = - mconcat $ mapMaybe usageThroughBindeeAliases $ - zip (patternNames pat) (patternAliases pat) - usageThroughBindeeAliases (name, aliases) = do - uses <- UT.lookup name utable - return $ mconcat $ map (`UT.usage` uses) $ namesToList aliases - -type BlockPred lore = ST.SymbolTable lore -> UT.UsageTable -> Stm lore -> Bool - -neverBlocks :: BlockPred lore -neverBlocks _ _ _ = False - -alwaysBlocks :: BlockPred lore -alwaysBlocks _ _ _ = True - -isFalse :: Bool -> BlockPred lore -isFalse b _ _ _ = not b - -orIf :: BlockPred lore -> BlockPred lore -> BlockPred lore -orIf p1 p2 body vtable need = p1 body vtable need || p2 body vtable need - -andAlso :: BlockPred lore -> BlockPred lore -> BlockPred lore -andAlso p1 p2 body vtable need = p1 body vtable need && p2 body vtable need - -isConsumed :: BlockPred lore -isConsumed _ utable = any (`UT.isConsumed` utable) . patternNames . stmPattern - -isOp :: BlockPred lore -isOp _ _ (Let _ _ Op{}) = True -isOp _ _ _ = False - -constructBody :: SimplifiableLore lore => Stms (Wise lore) -> Result - -> SimpleM lore (Body (Wise lore)) -constructBody stms res = - fmap fst $ runBinder $ insertStmsM $ do addStms stms - resultBodyM res - -type SimplifiedBody lore a = ((a, UT.UsageTable), Stms (Wise lore)) - -blockIf :: SimplifiableLore lore => - BlockPred (Wise lore) - -> SimpleM lore (SimplifiedBody lore a) - -> SimpleM lore ((Stms (Wise lore), a), Stms (Wise lore)) -blockIf block m = do - ((x, usages), stms) <- m - vtable <- askVtable - rules <- asksEngineEnv envRules - (blocked, hoisted) <- hoistStms rules block vtable usages stms - return ((blocked, x), hoisted) - -hasFree :: ASTLore lore => Names -> BlockPred lore -hasFree ks _ _ need = ks `namesIntersect` freeIn need - -isNotSafe :: ASTLore lore => BlockPred lore -isNotSafe _ _ = not . safeExp . stmExp - -isInPlaceBound :: BlockPred m -isInPlaceBound _ _ = isUpdate . stmExp - where isUpdate (BasicOp Update{}) = True - isUpdate _ = False - -isNotCheap :: ASTLore lore => BlockPred lore -isNotCheap _ _ = not . cheapStm - -cheapStm :: ASTLore lore => Stm lore -> Bool -cheapStm = cheapExp . stmExp - -cheapExp :: ASTLore lore => Exp lore -> Bool -cheapExp (BasicOp BinOp{}) = True -cheapExp (BasicOp SubExp{}) = True -cheapExp (BasicOp UnOp{}) = True -cheapExp (BasicOp CmpOp{}) = True -cheapExp (BasicOp ConvOp{}) = True -cheapExp (BasicOp Copy{}) = False -cheapExp (BasicOp Manifest{}) = False -cheapExp DoLoop{} = False -cheapExp (If _ tbranch fbranch _) = all cheapStm (bodyStms tbranch) && - all cheapStm (bodyStms fbranch) -cheapExp (Op op) = cheapOp op -cheapExp _ = True -- Used to be False, but - -- let's try it out. - -stmIs :: (Stm lore -> Bool) -> BlockPred lore -stmIs f _ _ = f - -loopInvariantStm :: ASTLore lore => ST.SymbolTable lore -> Stm lore -> Bool -loopInvariantStm vtable = - all (`nameIn` ST.availableAtClosestLoop vtable) . namesToList . freeIn - -hoistCommon :: SimplifiableLore lore => - SubExp -> IfSort - -> SimplifiedBody lore Result - -> SimplifiedBody lore Result - -> SimpleM lore (Body (Wise lore), - Body (Wise lore), - Stms (Wise lore)) -hoistCommon cond ifsort ((res1, usages1), stms1) ((res2, usages2), stms2) = do - is_alloc_fun <- asksEngineEnv $ isAllocation . envHoistBlockers - branch_blocker <- asksEngineEnv $ blockHoistBranch . envHoistBlockers - vtable <- askVtable - let -- We are unwilling to hoist things that are unsafe or costly, - -- *except* if they are invariant to the most enclosing loop, - -- because in that case they will also be hoisted past that - -- loop. - -- - -- We also try very hard to hoist allocations or anything that - -- contributes to memory or array size, because that will allow - -- allocations to be hoisted. - cond_loop_invariant = - all (`nameIn` ST.availableAtClosestLoop vtable) $ namesToList $ freeIn cond - - desirableToHoist stm = - is_alloc_fun stm || - (ST.loopDepth vtable > 0 && - cond_loop_invariant && - ifsort /= IfFallback && - loopInvariantStm vtable stm) - - -- No matter what, we always want to hoist constants as much as - -- possible. - isNotHoistableBnd _ _ (Let _ _ (BasicOp ArrayLit{})) = False - isNotHoistableBnd _ _ (Let _ _ (BasicOp SubExp{})) = False - isNotHoistableBnd _ usages (Let pat _ _) - | any (`UT.isSize` usages) $ patternNames pat = - False - isNotHoistableBnd _ _ _ = - -- Hoist aggressively out of versioning branches. - ifsort /= IfEquiv - - block = branch_blocker `orIf` - ((isNotSafe `orIf` isNotCheap) `andAlso` stmIs (not . desirableToHoist)) - `orIf` isInPlaceBound `orIf` isNotHoistableBnd - - rules <- asksEngineEnv envRules - (body1_bnds', safe1) <- protectIfHoisted cond True $ - hoistStms rules block vtable usages1 stms1 - (body2_bnds', safe2) <- protectIfHoisted cond False $ - hoistStms rules block vtable usages2 stms2 - let hoistable = safe1 <> safe2 - body1' <- constructBody body1_bnds' res1 - body2' <- constructBody body2_bnds' res2 - return (body1', body2', hoistable) - --- | Simplify a single body. The @[Diet]@ only covers the value --- elements, because the context cannot be consumed. -simplifyBody :: SimplifiableLore lore => - [Diet] -> Body lore -> SimpleM lore (SimplifiedBody lore Result) -simplifyBody ds (Body _ bnds res) = - simplifyStms bnds $ do res' <- simplifyResult ds res - return (res', mempty) - --- | Simplify a single 'Result'. The @[Diet]@ only covers the value --- elements, because the context cannot be consumed. -simplifyResult :: SimplifiableLore lore => - [Diet] -> Result -> SimpleM lore (Result, UT.UsageTable) -simplifyResult ds res = do - let (ctx_res, val_res) = splitFromEnd (length ds) res - -- Copy propagation is a little trickier here, because there is no - -- place to put the certificates when copy-propagating a certified - -- statement. However, for results in the *context*, it is OK to - -- just throw away the certificates, because for the program to be - -- type-correct, those statements must anyway be used (or - -- copy-propagated into) the statements producing the value result. - (ctx_res', _ctx_res_cs) <- collectCerts $ mapM simplify ctx_res - val_res' <- mapM simplify' val_res - - let consumption = consumeResult $ zip ds val_res' - res' = ctx_res' <> val_res' - return (res', UT.usages (freeIn res') <> consumption) - - where simplify' (Var name) = do - bnd <- ST.lookupSubExp name <$> askVtable - case bnd of - Just (Constant v, cs) - | cs == mempty -> return $ Constant v - Just (Var id', cs) - | cs == mempty -> return $ Var id' - _ -> return $ Var name - simplify' (Constant v) = - return $ Constant v - -isDoLoopResult :: Result -> UT.UsageTable -isDoLoopResult = mconcat . map checkForVar - where checkForVar (Var ident) = UT.inResultUsage ident - checkForVar _ = mempty - -simplifyStms :: SimplifiableLore lore => - Stms lore -> SimpleM lore (a, Stms (Wise lore)) - -> SimpleM lore (a, Stms (Wise lore)) -simplifyStms stms m = - case stmsHead stms of - Nothing -> inspectStms mempty m - Just (Let pat (StmAux stm_cs attrs dec) e, stms') -> do - stm_cs' <- simplify stm_cs - ((e', e_stms), e_cs) <- collectCerts $ simplifyExp e - (pat', pat_cs) <- collectCerts $ simplifyPattern pat - let cs = stm_cs'<>e_cs<>pat_cs - inspectStms e_stms $ - inspectStm (mkWiseLetStm pat' (StmAux cs attrs dec) e') $ - simplifyStms stms' m - -inspectStm :: SimplifiableLore lore => - Stm (Wise lore) -> SimpleM lore (a, Stms (Wise lore)) - -> SimpleM lore (a, Stms (Wise lore)) -inspectStm = inspectStms . oneStm - -inspectStms :: SimplifiableLore lore => - Stms (Wise lore) - -> SimpleM lore (a, Stms (Wise lore)) - -> SimpleM lore (a, Stms (Wise lore)) -inspectStms stms m = - case stmsHead stms of - Nothing -> m - Just (stm, stms') -> do - vtable <- askVtable - rules <- asksEngineEnv envRules - simplified <- topDownSimplifyStm rules vtable stm - case simplified of - Just newbnds -> changed >> inspectStms (newbnds <> stms') m - Nothing -> do (x, stms'') <- localVtable (ST.insertStm stm) $ inspectStms stms' m - return (x, oneStm stm <> stms'') - -simplifyOp :: Op lore -> SimpleM lore (Op (Wise lore), Stms (Wise lore)) -simplifyOp op = do f <- asks $ simplifyOpS . fst - f op - -simplifyExp :: SimplifiableLore lore => - Exp lore -> SimpleM lore (Exp (Wise lore), Stms (Wise lore)) - -simplifyExp (If cond tbranch fbranch (IfDec ts ifsort)) = do - -- Here, we have to check whether 'cond' puts a bound on some free - -- variable, and if so, chomp it. We should also try to do CSE - -- across branches. - cond' <- simplify cond - ts' <- mapM simplify ts - -- FIXME: we have to be conservative about the diet here, because we - -- lack proper ifnormation. Something is wrong with the order in - -- which the simplifier does things - it should be purely bottom-up - -- (or else, If expressions should indicate explicitly the diet of - -- their return types). - let ds = map (const Consume) ts - tbranch' <- simplifyBody ds tbranch - fbranch' <- simplifyBody ds fbranch - (tbranch'',fbranch'', hoisted) <- hoistCommon cond' ifsort tbranch' fbranch' - return (If cond' tbranch'' fbranch'' $ IfDec ts' ifsort, hoisted) - -simplifyExp (DoLoop ctx val form loopbody) = do - let (ctxparams, ctxinit) = unzip ctx - (valparams, valinit) = unzip val - ctxparams' <- mapM (traverse simplify) ctxparams - ctxinit' <- mapM simplify ctxinit - valparams' <- mapM (traverse simplify) valparams - valinit' <- mapM simplify valinit - let ctx' = zip ctxparams' ctxinit' - val' = zip valparams' valinit' - diets = map (diet . paramDeclType) valparams' - (form', boundnames, wrapbody) <- case form of - ForLoop loopvar it boundexp loopvars -> do - boundexp' <- simplify boundexp - let (loop_params, loop_arrs) = unzip loopvars - loop_params' <- mapM (traverse simplify) loop_params - loop_arrs' <- mapM simplify loop_arrs - let form' = ForLoop loopvar it boundexp' (zip loop_params' loop_arrs') - return (form', - namesFromList (loopvar : map paramName loop_params') <> fparamnames, - bindLoopVar loopvar it boundexp' . - protectLoopHoisted ctx' val' form' . - bindArrayLParams loop_params') - WhileLoop cond -> do - cond' <- simplify cond - return (WhileLoop cond', - fparamnames, - protectLoopHoisted ctx' val' (WhileLoop cond')) - seq_blocker <- asksEngineEnv $ blockHoistSeq . envHoistBlockers - ((loopstms, loopres), hoisted) <- - enterLoop $ consumeMerge $ - bindFParams (ctxparams'++valparams') $ wrapbody $ - blockIf - (hasFree boundnames `orIf` isConsumed - `orIf` seq_blocker `orIf` notWorthHoisting) $ do - ((res, uses), stms) <- simplifyBody diets loopbody - return ((res, uses <> isDoLoopResult res), stms) - loopbody' <- constructBody loopstms loopres - return (DoLoop ctx' val' form' loopbody', hoisted) - where fparamnames = - namesFromList (map (paramName . fst) $ ctx++val) - consumeMerge = - localVtable $ flip (foldl' (flip ST.consume)) $ namesToList consumed_by_merge - consumed_by_merge = - freeIn $ map snd $ filter (unique . paramDeclType . fst) val - -simplifyExp (Op op) = do (op', stms) <- simplifyOp op - return (Op op', stms) - --- Special case for simplification of commutative BinOps where we --- arrange the operands in sorted order. This can make expressions --- more identical, which helps CSE. -simplifyExp (BasicOp (BinOp op x y)) - | commutativeBinOp op = do - x' <- simplify x - y' <- simplify y - return (BasicOp $ BinOp op (min x' y') (max x' y'), mempty) - -simplifyExp e = do e' <- simplifyExpBase e - return (e', mempty) - -simplifyExpBase :: SimplifiableLore lore => - Exp lore -> SimpleM lore (Exp (Wise lore)) -simplifyExpBase = mapExpM hoist - where hoist = Mapper { - -- Bodies are handled explicitly because we need to - -- provide their result diet. - mapOnBody = - error "Unhandled body in simplification engine." - , mapOnSubExp = simplify - -- Lambdas are handled explicitly because we need to - -- bind their parameters. - , mapOnVName = simplify - , mapOnRetType = simplify - , mapOnBranchType = simplify - , mapOnFParam = - error "Unhandled FParam in simplification engine." - , mapOnLParam = - error "Unhandled LParam in simplification engine." - , mapOnOp = - error "Unhandled Op in simplification engine." - } - -type SimplifiableLore lore = (ASTLore lore, - Simplifiable (LetDec lore), - Simplifiable (FParamInfo lore), - Simplifiable (LParamInfo lore), - Simplifiable (RetType lore), - Simplifiable (BranchType lore), - CanBeWise (Op lore), - ST.IndexOp (OpWithWisdom (Op lore)), - BinderOps (Wise lore), - IsOp (Op lore)) - -class Simplifiable e where - simplify :: SimplifiableLore lore => e -> SimpleM lore e - -instance (Simplifiable a, Simplifiable b) => Simplifiable (a, b) where - simplify (x,y) = (,) <$> simplify x <*> simplify y - -instance (Simplifiable a, Simplifiable b, Simplifiable c) => - Simplifiable (a, b, c) where - simplify (x,y,z) = (,,) <$> simplify x <*> simplify y <*> simplify z - --- Convenient for Scatter. -instance Simplifiable Int where - simplify = pure - -instance Simplifiable a => Simplifiable (Maybe a) where - simplify Nothing = return Nothing - simplify (Just x) = Just <$> simplify x - -instance Simplifiable a => Simplifiable [a] where - simplify = mapM simplify - -instance Simplifiable SubExp where - simplify (Var name) = do - bnd <- ST.lookupSubExp name <$> askVtable - case bnd of - Just (Constant v, cs) -> do changed - usedCerts cs - return $ Constant v - Just (Var id', cs) -> do changed - usedCerts cs - return $ Var id' - _ -> return $ Var name - simplify (Constant v) = - return $ Constant v - -simplifyPattern :: (SimplifiableLore lore, Simplifiable dec) => - PatternT dec - -> SimpleM lore (PatternT dec) -simplifyPattern pat = - Pattern <$> - mapM inspect (patternContextElements pat) <*> - mapM inspect (patternValueElements pat) - where inspect (PatElem name lore) = PatElem name <$> simplify lore - -instance Simplifiable () where - simplify = pure - -instance Simplifiable VName where - simplify v = do - se <- ST.lookupSubExp v <$> askVtable - case se of - Just (Var v', cs) -> do changed - usedCerts cs - return v' - _ -> return v - -instance Simplifiable d => Simplifiable (ShapeBase d) where - simplify = fmap Shape . simplify . shapeDims - -instance Simplifiable ExtSize where - simplify (Free se) = Free <$> simplify se - simplify (Ext x) = return $ Ext x - -instance Simplifiable Space where - simplify (ScalarSpace ds t) = ScalarSpace <$> simplify ds <*> pure t - simplify s = pure s - -instance Simplifiable shape => Simplifiable (TypeBase shape u) where - simplify (Array et shape u) = do - shape' <- simplify shape - return $ Array et shape' u - simplify (Mem space) = - Mem <$> simplify space - simplify (Prim bt) = - return $ Prim bt - -instance Simplifiable d => Simplifiable (DimIndex d) where - simplify (DimFix i) = DimFix <$> simplify i - simplify (DimSlice i n s) = DimSlice <$> simplify i <*> simplify n <*> simplify s - -simplifyLambda :: SimplifiableLore lore => - Lambda lore - -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)) -simplifyLambda lam = do - par_blocker <- asksEngineEnv $ blockHoistPar . envHoistBlockers - simplifyLambdaMaybeHoist par_blocker lam - -simplifyLambdaNoHoisting :: SimplifiableLore lore => - Lambda lore - -> SimpleM lore (Lambda (Wise lore)) -simplifyLambdaNoHoisting lam = - fst <$> simplifyLambdaMaybeHoist (isFalse False) lam - -simplifyLambdaMaybeHoist :: SimplifiableLore lore => - BlockPred (Wise lore) -> Lambda lore - -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)) -simplifyLambdaMaybeHoist blocked lam@(Lambda params body rettype) = do - params' <- mapM (traverse simplify) params - let paramnames = namesFromList $ boundByLambda lam - ((lamstms, lamres), hoisted) <- - enterLoop $ - bindLParams params' $ - blockIf (blocked `orIf` hasFree paramnames `orIf` isConsumed) $ - simplifyBody (map (const Observe) rettype) body - body' <- constructBody lamstms lamres - rettype' <- simplify rettype - return (Lambda params' body' rettype', hoisted) - -consumeResult :: [(Diet, SubExp)] -> UT.UsageTable -consumeResult = mconcat . map inspect - where inspect (Consume, se) = - mconcat $ map UT.consumedUsage $ namesToList $ subExpAliases se - inspect _ = mempty - -instance Simplifiable Certificates where - simplify (Certificates ocs) = Certificates . nub . concat <$> mapM check ocs - where check idd = do - vv <- ST.lookupSubExp idd <$> askVtable - case vv of - Just (Constant Checked, Certificates cs) -> return cs - Just (Var idd', _) -> return [idd'] - _ -> return [idd] - - -insertAllStms :: SimplifiableLore lore => - SimpleM lore (SimplifiedBody lore Result) - -> SimpleM lore (Body (Wise lore)) -insertAllStms = uncurry constructBody . fst <=< blockIf (isFalse False) - - -simplifyFun :: SimplifiableLore lore => - FunDef lore -> SimpleM lore (FunDef (Wise lore)) +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE Strict #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- +-- Perform general rule-based simplification based on data dependency +-- information. This module will: +-- +-- * Perform common-subexpression elimination (CSE). +-- +-- * Hoist expressions out of loops (including lambdas) and +-- branches. This is done as aggressively as possible. +-- +-- * Apply simplification rules (see +-- "Futhark.Optimise.Simplification.Rules"). +-- +-- If you just want to run the simplifier as simply as possible, you +-- may prefer to use the "Futhark.Optimise.Simplify" module. +module Futhark.Optimise.Simplify.Engine + ( -- * Monadic interface + SimpleM, + runSimpleM, + SimpleOps (..), + SimplifyOp, + bindableSimpleOps, + Env (envHoistBlockers, envRules), + emptyEnv, + HoistBlockers (..), + neverBlocks, + noExtraHoistBlockers, + neverHoist, + BlockPred, + orIf, + hasFree, + isConsumed, + isFalse, + isOp, + isNotSafe, + asksEngineEnv, + askVtable, + localVtable, + + -- * Building blocks + SimplifiableLore, + Simplifiable (..), + simplifyStms, + simplifyFun, + simplifyLambda, + simplifyLambdaNoHoisting, + bindLParams, + simplifyBody, + SimplifiedBody, + ST.SymbolTable, + hoistStms, + blockIf, + module Futhark.Optimise.Simplify.Lore, + ) +where + +import Control.Monad.Reader +import Control.Monad.State.Strict +import Data.Either +import Data.List (find, foldl', mapAccumL, nub) +import Data.Maybe +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.Lore +import Futhark.Optimise.Simplify.Rule +import Futhark.Util (splitFromEnd) + +data HoistBlockers lore = HoistBlockers + { -- | Blocker for hoisting out of parallel loops. + blockHoistPar :: BlockPred (Wise lore), + -- | Blocker for hoisting out of sequential loops. + blockHoistSeq :: BlockPred (Wise lore), + -- | Blocker for hoisting out of branches. + blockHoistBranch :: BlockPred (Wise lore), + isAllocation :: Stm (Wise lore) -> Bool + } + +noExtraHoistBlockers :: HoistBlockers lore +noExtraHoistBlockers = + HoistBlockers neverBlocks neverBlocks neverBlocks (const False) + +neverHoist :: HoistBlockers lore +neverHoist = + HoistBlockers alwaysBlocks alwaysBlocks alwaysBlocks (const False) + +data Env lore = Env + { envRules :: RuleBook (Wise lore), + envHoistBlockers :: HoistBlockers lore, + envVtable :: ST.SymbolTable (Wise lore) + } + +emptyEnv :: RuleBook (Wise lore) -> HoistBlockers lore -> Env lore +emptyEnv rules blockers = + Env + { envRules = rules, + envHoistBlockers = blockers, + envVtable = mempty + } + +type Protect m = SubExp -> Pattern (Lore m) -> Op (Lore m) -> Maybe (m ()) + +data SimpleOps lore = SimpleOps + { mkExpDecS :: + ST.SymbolTable (Wise lore) -> + Pattern (Wise lore) -> + Exp (Wise lore) -> + SimpleM lore (ExpDec (Wise lore)), + mkBodyS :: + ST.SymbolTable (Wise lore) -> + Stms (Wise lore) -> + Result -> + SimpleM lore (Body (Wise lore)), + -- | Make a hoisted Op safe. The SubExp is a boolean + -- that is true when the value of the statement will + -- actually be used. + protectHoistedOpS :: Protect (Binder (Wise lore)), + opUsageS :: Op (Wise lore) -> UT.UsageTable, + simplifyOpS :: SimplifyOp lore (Op lore) + } + +type SimplifyOp lore op = op -> SimpleM lore (OpWithWisdom op, Stms (Wise lore)) + +bindableSimpleOps :: + (SimplifiableLore lore, Bindable lore) => + SimplifyOp lore (Op lore) -> + SimpleOps lore +bindableSimpleOps = + SimpleOps mkExpDecS' mkBodyS' protectHoistedOpS' (const mempty) + where + mkExpDecS' _ pat e = return $ mkExpDec pat e + mkBodyS' _ bnds res = return $ mkBody bnds res + protectHoistedOpS' _ _ _ = Nothing + +newtype SimpleM lore a + = SimpleM + ( ReaderT + (SimpleOps lore, Env lore) + (State (VNameSource, Bool, Certificates)) + a + ) + deriving + ( Applicative, + Functor, + Monad, + MonadReader (SimpleOps lore, Env lore), + MonadState (VNameSource, Bool, Certificates) + ) + +instance MonadFreshNames (SimpleM lore) where + putNameSource src = modify $ \(_, b, c) -> (src, b, c) + getNameSource = gets $ \(a, _, _) -> a + +instance SimplifiableLore lore => HasScope (Wise lore) (SimpleM lore) where + askScope = ST.toScope <$> askVtable + lookupType name = do + vtable <- askVtable + case ST.lookupType name vtable of + Just t -> return t + Nothing -> + error $ + "SimpleM.lookupType: cannot find variable " + ++ pretty name + ++ " in symbol table." + +instance + SimplifiableLore lore => + LocalScope (Wise lore) (SimpleM lore) + where + localScope types = localVtable (<> ST.fromScope types) + +runSimpleM :: + SimpleM lore a -> + SimpleOps lore -> + Env lore -> + VNameSource -> + ((a, Bool), VNameSource) +runSimpleM (SimpleM m) simpl env src = + let (x, (src', b, _)) = runState (runReaderT m (simpl, env)) (src, False, mempty) + in ((x, b), src') + +askEngineEnv :: SimpleM lore (Env lore) +askEngineEnv = asks snd + +asksEngineEnv :: (Env lore -> a) -> SimpleM lore a +asksEngineEnv f = f <$> askEngineEnv + +askVtable :: SimpleM lore (ST.SymbolTable (Wise lore)) +askVtable = asksEngineEnv envVtable + +localVtable :: + (ST.SymbolTable (Wise lore) -> ST.SymbolTable (Wise lore)) -> + SimpleM lore a -> + SimpleM lore a +localVtable f = local $ \(ops, env) -> (ops, env {envVtable = f $ envVtable env}) + +collectCerts :: SimpleM lore a -> SimpleM lore (a, Certificates) +collectCerts m = do + x <- m + (a, b, cs) <- get + put (a, b, mempty) + return (x, cs) + +-- | Mark that we have changed something and it would be a good idea +-- to re-run the simplifier. +changed :: SimpleM lore () +changed = modify $ \(src, _, cs) -> (src, True, cs) + +usedCerts :: Certificates -> SimpleM lore () +usedCerts cs = modify $ \(a, b, c) -> (a, b, cs <> c) + +enterLoop :: SimpleM lore a -> SimpleM lore a +enterLoop = localVtable ST.deepen + +bindFParams :: SimplifiableLore lore => [FParam (Wise lore)] -> SimpleM lore a -> SimpleM lore a +bindFParams params = + localVtable $ ST.insertFParams params + +bindLParams :: SimplifiableLore lore => [LParam (Wise lore)] -> SimpleM lore a -> SimpleM lore a +bindLParams params = + localVtable $ \vtable -> foldr ST.insertLParam vtable params + +bindArrayLParams :: + SimplifiableLore lore => + [LParam (Wise lore)] -> + SimpleM lore a -> + SimpleM lore a +bindArrayLParams params = + localVtable $ \vtable -> foldl' (flip ST.insertLParam) vtable params + +bindMerge :: + SimplifiableLore lore => + [(FParam (Wise lore), SubExp, SubExp)] -> + SimpleM lore a -> + SimpleM lore a +bindMerge = localVtable . ST.insertLoopMerge + +bindLoopVar :: SimplifiableLore lore => VName -> IntType -> SubExp -> SimpleM lore a -> SimpleM lore a +bindLoopVar var it bound = + localVtable $ ST.insertLoopVar var it bound + +-- | We are willing to hoist potentially unsafe statements out of +-- branches, but they most be protected by adding a branch on top of +-- them. (This means such hoisting is not worth it unless they are in +-- turn hoisted out of a loop somewhere.) +protectIfHoisted :: + SimplifiableLore lore => + -- | Branch condition. + SubExp -> + -- | Which side of the branch are we + -- protecting here? + Bool -> + SimpleM lore (a, Stms (Wise lore)) -> + SimpleM lore (a, Stms (Wise lore)) +protectIfHoisted cond side m = do + (x, stms) <- m + ops <- asks $ protectHoistedOpS . fst + runBinder $ do + if not $ all (safeExp . stmExp) stms + then do + cond' <- + if side + then return cond + else letSubExp "cond_neg" $ BasicOp $ UnOp Not cond + mapM_ (protectIf ops unsafeOrCostly cond') stms + else addStms stms + return x + where + unsafeOrCostly e = not (safeExp e) || not (cheapExp e) + +-- | We are willing to hoist potentially unsafe statements out of +-- loops, but they most be protected by adding a branch on top of +-- them. +protectLoopHoisted :: + SimplifiableLore lore => + [(FParam (Wise lore), SubExp)] -> + [(FParam (Wise lore), SubExp)] -> + LoopForm (Wise lore) -> + SimpleM lore (a, Stms (Wise lore)) -> + SimpleM lore (a, Stms (Wise lore)) +protectLoopHoisted ctx val form m = do + (x, stms) <- m + ops <- asks $ protectHoistedOpS . fst + runBinder $ do + if not $ all (safeExp . stmExp) stms + then do + is_nonempty <- checkIfNonEmpty + mapM_ (protectIf ops (not . safeExp) is_nonempty) stms + else addStms stms + return x + where + checkIfNonEmpty = + case form of + WhileLoop cond + | Just (_, cond_init) <- + find ((== cond) . paramName . fst) $ ctx ++ val -> + return cond_init + | otherwise -> return $ constant True -- infinite loop + ForLoop _ it bound _ -> + letSubExp "loop_nonempty" $ + BasicOp $ CmpOp (CmpSlt it) (intConst it 0) bound + +protectIf :: + MonadBinder m => + Protect m -> + (Exp (Lore m) -> Bool) -> + SubExp -> + Stm (Lore m) -> + m () +protectIf + _ + _ + taken + ( Let + pat + aux + (If cond taken_body untaken_body (IfDec if_ts IfFallback)) + ) = do + cond' <- letSubExp "protect_cond_conj" $ BasicOp $ BinOp LogAnd taken cond + auxing aux $ + letBind pat $ + If cond' taken_body untaken_body $ + IfDec if_ts IfFallback +protectIf _ _ taken (Let pat aux (BasicOp (Assert cond msg loc))) = do + not_taken <- letSubExp "loop_not_taken" $ BasicOp $ UnOp Not taken + cond' <- letSubExp "protect_assert_disj" $ BasicOp $ BinOp LogOr not_taken cond + auxing aux $ letBind pat $ BasicOp $ Assert cond' msg loc +protectIf protect _ taken (Let pat aux (Op op)) + | Just m <- protect taken pat op = + auxing aux m +protectIf _ f taken (Let pat aux e) + | f e = + case makeSafe e of + Just e' -> + auxing aux $ letBind pat e' + Nothing -> do + taken_body <- eBody [pure e] + untaken_body <- + eBody $ + map + (emptyOfType $ patternContextNames pat) + (patternValueTypes pat) + if_ts <- expTypesFromPattern pat + auxing aux $ + letBind pat $ + If taken taken_body untaken_body $ + IfDec if_ts IfFallback +protectIf _ _ _ stm = + addStm stm + +makeSafe :: Exp lore -> Maybe (Exp lore) +makeSafe (BasicOp (BinOp (SDiv t _) x y)) = + Just $ BasicOp (BinOp (SDiv t Safe) x y) +makeSafe (BasicOp (BinOp (SDivUp t _) x y)) = + Just $ BasicOp (BinOp (SDivUp t Safe) x y) +makeSafe (BasicOp (BinOp (SQuot t _) x y)) = + Just $ BasicOp (BinOp (SQuot t Safe) x y) +makeSafe (BasicOp (BinOp (UDiv t _) x y)) = + Just $ BasicOp (BinOp (UDiv t Safe) x y) +makeSafe (BasicOp (BinOp (UDivUp t _) x y)) = + Just $ BasicOp (BinOp (UDivUp t Safe) x y) +makeSafe (BasicOp (BinOp (SMod t _) x y)) = + Just $ BasicOp (BinOp (SMod t Safe) x y) +makeSafe (BasicOp (BinOp (SRem t _) x y)) = + Just $ BasicOp (BinOp (SRem t Safe) x y) +makeSafe (BasicOp (BinOp (UMod t _) x y)) = + Just $ BasicOp (BinOp (UMod t Safe) x y) +makeSafe _ = + Nothing + +emptyOfType :: MonadBinder m => [VName] -> Type -> m (Exp (Lore m)) +emptyOfType _ Mem {} = + error "emptyOfType: Cannot hoist non-existential memory." +emptyOfType _ (Prim pt) = + return $ BasicOp $ SubExp $ Constant $ blankPrimValue pt +emptyOfType ctx_names (Array pt shape _) = do + let dims = map zeroIfContext $ shapeDims shape + return $ BasicOp $ Scratch pt dims + where + zeroIfContext (Var v) | v `elem` ctx_names = intConst Int32 0 + zeroIfContext se = se + +-- | Statements that are not worth hoisting out of loops, because they +-- are unsafe, and added safety (by 'protectLoopHoisted') may inhibit +-- further optimisation.. +notWorthHoisting :: ASTLore lore => BlockPred lore +notWorthHoisting _ _ (Let pat _ e) = + not (safeExp e) && any ((> 0) . arrayRank) (patternTypes pat) + +hoistStms :: + SimplifiableLore lore => + RuleBook (Wise lore) -> + BlockPred (Wise lore) -> + ST.SymbolTable (Wise lore) -> + UT.UsageTable -> + Stms (Wise lore) -> + SimpleM + lore + ( Stms (Wise lore), + Stms (Wise lore) + ) +hoistStms rules block vtable uses orig_stms = do + (blocked, hoisted) <- simplifyStmsBottomUp vtable uses orig_stms + unless (null hoisted) changed + return (stmsFromList blocked, stmsFromList hoisted) + where + simplifyStmsBottomUp vtable' uses' stms = do + (_, stms') <- simplifyStmsBottomUp' vtable' uses' stms + -- We need to do a final pass to ensure that nothing is + -- hoisted past something that it depends on. + let (blocked, hoisted) = partitionEithers $ blockUnhoistedDeps stms' + return (blocked, hoisted) + + simplifyStmsBottomUp' vtable' uses' stms = do + opUsage <- asks $ opUsageS . fst + let usageInStm stm = + UT.usageInStm stm + <> case stmExp stm of + Op op -> opUsage op + _ -> mempty + foldM (hoistable usageInStm) (uses', []) $ reverse $ zip (stmsToList stms) vtables + where + vtables = scanl (flip ST.insertStm) vtable' $ stmsToList stms + + hoistable usageInStm (uses', stms) (stm, vtable') + | not $ any (`UT.isUsedDirectly` uses') $ provides stm -- Dead statement. + = + return (uses', stms) + | otherwise = do + res <- + localVtable (const vtable') $ + bottomUpSimplifyStm rules (vtable', uses') stm + case res of + Nothing -- Nothing to optimise - see if hoistable. + | block vtable' uses' stm -> + return + ( expandUsage usageInStm vtable' uses' stm + `UT.without` provides stm, + Left stm : stms + ) + | otherwise -> + return + ( expandUsage usageInStm vtable' uses' stm, + Right stm : stms + ) + Just optimstms -> do + changed + (uses'', stms') <- simplifyStmsBottomUp' vtable' uses' optimstms + return (uses'', stms' ++ stms) + +blockUnhoistedDeps :: + ASTLore lore => + [Either (Stm lore) (Stm lore)] -> + [Either (Stm lore) (Stm lore)] +blockUnhoistedDeps = snd . mapAccumL block mempty + where + block blocked (Left need) = + (blocked <> namesFromList (provides need), Left need) + block blocked (Right need) + | blocked `namesIntersect` freeIn need = + (blocked <> namesFromList (provides need), Left need) + | otherwise = + (blocked, Right need) + +provides :: Stm lore -> [VName] +provides = patternNames . stmPattern + +expandUsage :: + (ASTLore lore, Aliased lore) => + (Stm lore -> UT.UsageTable) -> + ST.SymbolTable lore -> + UT.UsageTable -> + Stm lore -> + UT.UsageTable +expandUsage usageInStm vtable utable stm@(Let pat _ e) = + UT.expand (`ST.lookupAliases` vtable) (usageInStm stm <> usageThroughAliases) + <> ( if any (`UT.isSize` utable) (patternNames pat) + then UT.sizeUsages (freeIn e) + else mempty + ) + <> utable + where + usageThroughAliases = + mconcat $ + mapMaybe usageThroughBindeeAliases $ + zip (patternNames pat) (patternAliases pat) + usageThroughBindeeAliases (name, aliases) = do + uses <- UT.lookup name utable + return $ mconcat $ map (`UT.usage` uses) $ namesToList aliases + +type BlockPred lore = ST.SymbolTable lore -> UT.UsageTable -> Stm lore -> Bool + +neverBlocks :: BlockPred lore +neverBlocks _ _ _ = False + +alwaysBlocks :: BlockPred lore +alwaysBlocks _ _ _ = True + +isFalse :: Bool -> BlockPred lore +isFalse b _ _ _ = not b + +orIf :: BlockPred lore -> BlockPred lore -> BlockPred lore +orIf p1 p2 body vtable need = p1 body vtable need || p2 body vtable need + +andAlso :: BlockPred lore -> BlockPred lore -> BlockPred lore +andAlso p1 p2 body vtable need = p1 body vtable need && p2 body vtable need + +isConsumed :: BlockPred lore +isConsumed _ utable = any (`UT.isConsumed` utable) . patternNames . stmPattern + +isOp :: BlockPred lore +isOp _ _ (Let _ _ Op {}) = True +isOp _ _ _ = False + +constructBody :: + SimplifiableLore lore => + Stms (Wise lore) -> + Result -> + SimpleM lore (Body (Wise lore)) +constructBody stms res = + fmap fst $ + runBinder $ + insertStmsM $ do + addStms stms + resultBodyM res + +type SimplifiedBody lore a = ((a, UT.UsageTable), Stms (Wise lore)) + +blockIf :: + SimplifiableLore lore => + BlockPred (Wise lore) -> + SimpleM lore (SimplifiedBody lore a) -> + SimpleM lore ((Stms (Wise lore), a), Stms (Wise lore)) +blockIf block m = do + ((x, usages), stms) <- m + vtable <- askVtable + rules <- asksEngineEnv envRules + (blocked, hoisted) <- hoistStms rules block vtable usages stms + return ((blocked, x), hoisted) + +hasFree :: ASTLore lore => Names -> BlockPred lore +hasFree ks _ _ need = ks `namesIntersect` freeIn need + +isNotSafe :: ASTLore lore => BlockPred lore +isNotSafe _ _ = not . safeExp . stmExp + +isInPlaceBound :: BlockPred m +isInPlaceBound _ _ = isUpdate . stmExp + where + isUpdate (BasicOp Update {}) = True + isUpdate _ = False + +isNotCheap :: ASTLore lore => BlockPred lore +isNotCheap _ _ = not . cheapStm + +cheapStm :: ASTLore lore => Stm lore -> Bool +cheapStm = cheapExp . stmExp + +cheapExp :: ASTLore lore => Exp lore -> Bool +cheapExp (BasicOp BinOp {}) = True +cheapExp (BasicOp SubExp {}) = True +cheapExp (BasicOp UnOp {}) = True +cheapExp (BasicOp CmpOp {}) = True +cheapExp (BasicOp ConvOp {}) = True +cheapExp (BasicOp Copy {}) = False +cheapExp (BasicOp Manifest {}) = False +cheapExp DoLoop {} = False +cheapExp (If _ tbranch fbranch _) = + all cheapStm (bodyStms tbranch) + && all cheapStm (bodyStms fbranch) +cheapExp (Op op) = cheapOp op +cheapExp _ = True -- Used to be False, but +-- let's try it out. + +stmIs :: (Stm lore -> Bool) -> BlockPred lore +stmIs f _ _ = f + +loopInvariantStm :: ASTLore lore => ST.SymbolTable lore -> Stm lore -> Bool +loopInvariantStm vtable = + all (`nameIn` ST.availableAtClosestLoop vtable) . namesToList . freeIn + +hoistCommon :: + SimplifiableLore lore => + SubExp -> + IfSort -> + SimplifiedBody lore Result -> + SimplifiedBody lore Result -> + SimpleM + lore + ( Body (Wise lore), + Body (Wise lore), + Stms (Wise lore) + ) +hoistCommon cond ifsort ((res1, usages1), stms1) ((res2, usages2), stms2) = do + is_alloc_fun <- asksEngineEnv $ isAllocation . envHoistBlockers + branch_blocker <- asksEngineEnv $ blockHoistBranch . envHoistBlockers + vtable <- askVtable + let -- We are unwilling to hoist things that are unsafe or costly, + + -- because in that case they will also be hoisted past that + -- loop. + -- + -- We also try very hard to hoist allocations or anything that + -- contributes to memory or array size, because that will allow + -- allocations to be hoisted. + cond_loop_invariant = + all (`nameIn` ST.availableAtClosestLoop vtable) $ namesToList $ freeIn cond + + desirableToHoist stm = + is_alloc_fun stm + || ( ST.loopDepth vtable > 0 + && cond_loop_invariant + && ifsort /= IfFallback + && loopInvariantStm vtable stm + ) + + -- No matter what, we always want to hoist constants as much as + -- possible. + isNotHoistableBnd _ _ (Let _ _ (BasicOp ArrayLit {})) = False + isNotHoistableBnd _ _ (Let _ _ (BasicOp SubExp {})) = False + isNotHoistableBnd _ usages (Let pat _ _) + | any (`UT.isSize` usages) $ patternNames pat = + False + isNotHoistableBnd _ _ stm + | is_alloc_fun stm = False + isNotHoistableBnd _ _ _ = + -- Hoist aggressively out of versioning branches. + ifsort /= IfEquiv + + block = + branch_blocker + `orIf` ((isNotSafe `orIf` isNotCheap) `andAlso` stmIs (not . desirableToHoist)) + `orIf` isInPlaceBound + `orIf` isNotHoistableBnd + + rules <- asksEngineEnv envRules + (body1_bnds', safe1) <- + protectIfHoisted cond True $ + hoistStms rules block vtable usages1 stms1 + (body2_bnds', safe2) <- + protectIfHoisted cond False $ + hoistStms rules block vtable usages2 stms2 + let hoistable = safe1 <> safe2 + body1' <- constructBody body1_bnds' res1 + body2' <- constructBody body2_bnds' res2 + return (body1', body2', hoistable) + +-- | Simplify a single body. The @[Diet]@ only covers the value +-- elements, because the context cannot be consumed. +simplifyBody :: + SimplifiableLore lore => + [Diet] -> + Body lore -> + SimpleM lore (SimplifiedBody lore Result) +simplifyBody ds (Body _ bnds res) = + simplifyStms bnds $ do + res' <- simplifyResult ds res + return (res', mempty) + +-- | Simplify a single 'Result'. The @[Diet]@ only covers the value +-- elements, because the context cannot be consumed. +simplifyResult :: + SimplifiableLore lore => + [Diet] -> + Result -> + SimpleM lore (Result, UT.UsageTable) +simplifyResult ds res = do + let (ctx_res, val_res) = splitFromEnd (length ds) res + -- Copy propagation is a little trickier here, because there is no + -- place to put the certificates when copy-propagating a certified + -- statement. However, for results in the *context*, it is OK to + -- just throw away the certificates, because for the program to be + -- type-correct, those statements must anyway be used (or + -- copy-propagated into) the statements producing the value result. + (ctx_res', _ctx_res_cs) <- collectCerts $ mapM simplify ctx_res + val_res' <- mapM simplify' val_res + + let consumption = consumeResult $ zip ds val_res' + res' = ctx_res' <> val_res' + return (res', UT.usages (freeIn res') <> consumption) + where + simplify' (Var name) = do + bnd <- ST.lookupSubExp name <$> askVtable + case bnd of + Just (Constant v, cs) + | cs == mempty -> return $ Constant v + Just (Var id', cs) + | cs == mempty -> return $ Var id' + _ -> return $ Var name + simplify' (Constant v) = + return $ Constant v + +isDoLoopResult :: Result -> UT.UsageTable +isDoLoopResult = mconcat . map checkForVar + where + checkForVar (Var ident) = UT.inResultUsage ident + checkForVar _ = mempty + +simplifyStms :: + SimplifiableLore lore => + Stms lore -> + SimpleM lore (a, Stms (Wise lore)) -> + SimpleM lore (a, Stms (Wise lore)) +simplifyStms stms m = + case stmsHead stms of + Nothing -> inspectStms mempty m + Just (Let pat (StmAux stm_cs attrs dec) e, stms') -> do + stm_cs' <- simplify stm_cs + ((e', e_stms), e_cs) <- collectCerts $ simplifyExp e + (pat', pat_cs) <- collectCerts $ simplifyPattern pat + let cs = stm_cs' <> e_cs <> pat_cs + inspectStms e_stms $ + inspectStm (mkWiseLetStm pat' (StmAux cs attrs dec) e') $ + simplifyStms stms' m + +inspectStm :: + SimplifiableLore lore => + Stm (Wise lore) -> + SimpleM lore (a, Stms (Wise lore)) -> + SimpleM lore (a, Stms (Wise lore)) +inspectStm = inspectStms . oneStm + +inspectStms :: + SimplifiableLore lore => + Stms (Wise lore) -> + SimpleM lore (a, Stms (Wise lore)) -> + SimpleM lore (a, Stms (Wise lore)) +inspectStms stms m = + case stmsHead stms of + Nothing -> m + Just (stm, stms') -> do + vtable <- askVtable + rules <- asksEngineEnv envRules + simplified <- topDownSimplifyStm rules vtable stm + case simplified of + Just newbnds -> changed >> inspectStms (newbnds <> stms') m + Nothing -> do + (x, stms'') <- localVtable (ST.insertStm stm) $ inspectStms stms' m + return (x, oneStm stm <> stms'') + +simplifyOp :: Op lore -> SimpleM lore (Op (Wise lore), Stms (Wise lore)) +simplifyOp op = do + f <- asks $ simplifyOpS . fst + f op + +simplifyExp :: + SimplifiableLore lore => + Exp lore -> + SimpleM lore (Exp (Wise lore), Stms (Wise lore)) +simplifyExp (If cond tbranch fbranch (IfDec ts ifsort)) = do + -- Here, we have to check whether 'cond' puts a bound on some free + -- variable, and if so, chomp it. We should also try to do CSE + -- across branches. + cond' <- simplify cond + ts' <- mapM simplify ts + -- FIXME: we have to be conservative about the diet here, because we + -- lack proper ifnormation. Something is wrong with the order in + -- which the simplifier does things - it should be purely bottom-up + -- (or else, If expressions should indicate explicitly the diet of + -- their return types). + let ds = map (const Consume) ts + tbranch' <- simplifyBody ds tbranch + fbranch' <- simplifyBody ds fbranch + (tbranch'', fbranch'', hoisted) <- hoistCommon cond' ifsort tbranch' fbranch' + return (If cond' tbranch'' fbranch'' $ IfDec ts' ifsort, hoisted) +simplifyExp (DoLoop ctx val form loopbody) = do + let (ctxparams, ctxinit) = unzip ctx + (valparams, valinit) = unzip val + ctxparams' <- mapM (traverse simplify) ctxparams + ctxinit' <- mapM simplify ctxinit + valparams' <- mapM (traverse simplify) valparams + valinit' <- mapM simplify valinit + let ctx' = zip ctxparams' ctxinit' + val' = zip valparams' valinit' + diets = map (diet . paramDeclType) valparams' + (form', boundnames, wrapbody) <- case form of + ForLoop loopvar it boundexp loopvars -> do + boundexp' <- simplify boundexp + let (loop_params, loop_arrs) = unzip loopvars + loop_params' <- mapM (traverse simplify) loop_params + loop_arrs' <- mapM simplify loop_arrs + let form' = ForLoop loopvar it boundexp' (zip loop_params' loop_arrs') + return + ( form', + namesFromList (loopvar : map paramName loop_params') <> fparamnames, + bindLoopVar loopvar it boundexp' + . protectLoopHoisted ctx' val' form' + . bindArrayLParams loop_params' + ) + WhileLoop cond -> do + cond' <- simplify cond + return + ( WhileLoop cond', + fparamnames, + protectLoopHoisted ctx' val' (WhileLoop cond') + ) + seq_blocker <- asksEngineEnv $ blockHoistSeq . envHoistBlockers + ((loopstms, loopres), hoisted) <- + enterLoop $ + consumeMerge $ + bindMerge (zipWith withRes (ctx' ++ val') (bodyResult loopbody)) $ + wrapbody $ + blockIf + ( hasFree boundnames `orIf` isConsumed + `orIf` seq_blocker + `orIf` notWorthHoisting + ) + $ do + ((res, uses), stms) <- simplifyBody diets loopbody + return ((res, uses <> isDoLoopResult res), stms) + loopbody' <- constructBody loopstms loopres + return (DoLoop ctx' val' form' loopbody', hoisted) + where + fparamnames = + namesFromList (map (paramName . fst) $ ctx ++ val) + consumeMerge = + localVtable $ flip (foldl' (flip ST.consume)) $ namesToList consumed_by_merge + consumed_by_merge = + freeIn $ map snd $ filter (unique . paramDeclType . fst) val + withRes (p, x) y = (p, x, y) +simplifyExp (Op op) = do + (op', stms) <- simplifyOp op + return (Op op', stms) + +-- Special case for simplification of commutative BinOps where we +-- arrange the operands in sorted order. This can make expressions +-- more identical, which helps CSE. +simplifyExp (BasicOp (BinOp op x y)) + | commutativeBinOp op = do + x' <- simplify x + y' <- simplify y + return (BasicOp $ BinOp op (min x' y') (max x' y'), mempty) +simplifyExp e = do + e' <- simplifyExpBase e + return (e', mempty) + +simplifyExpBase :: + SimplifiableLore lore => + Exp lore -> + SimpleM lore (Exp (Wise lore)) +simplifyExpBase = mapExpM hoist + where + hoist = + Mapper + { -- Bodies are handled explicitly because we need to + -- provide their result diet. + mapOnBody = + error "Unhandled body in simplification engine.", + mapOnSubExp = simplify, + -- Lambdas are handled explicitly because we need to + -- bind their parameters. + mapOnVName = simplify, + mapOnRetType = simplify, + mapOnBranchType = simplify, + mapOnFParam = + error "Unhandled FParam in simplification engine.", + mapOnLParam = + error "Unhandled LParam in simplification engine.", + mapOnOp = + error "Unhandled Op in simplification engine." + } + +type SimplifiableLore lore = + ( ASTLore lore, + Simplifiable (LetDec lore), + Simplifiable (FParamInfo lore), + Simplifiable (LParamInfo lore), + Simplifiable (RetType lore), + Simplifiable (BranchType lore), + CanBeWise (Op lore), + ST.IndexOp (OpWithWisdom (Op lore)), + BinderOps (Wise lore), + IsOp (Op lore) + ) + +class Simplifiable e where + simplify :: SimplifiableLore lore => e -> SimpleM lore e + +instance (Simplifiable a, Simplifiable b) => Simplifiable (a, b) where + simplify (x, y) = (,) <$> simplify x <*> simplify y + +instance + (Simplifiable a, Simplifiable b, Simplifiable c) => + Simplifiable (a, b, c) + where + simplify (x, y, z) = (,,) <$> simplify x <*> simplify y <*> simplify z + +-- Convenient for Scatter. +instance Simplifiable Int where + simplify = pure + +instance Simplifiable a => Simplifiable (Maybe a) where + simplify Nothing = return Nothing + simplify (Just x) = Just <$> simplify x + +instance Simplifiable a => Simplifiable [a] where + simplify = mapM simplify + +instance Simplifiable SubExp where + simplify (Var name) = do + bnd <- ST.lookupSubExp name <$> askVtable + case bnd of + Just (Constant v, cs) -> do + changed + usedCerts cs + return $ Constant v + Just (Var id', cs) -> do + changed + usedCerts cs + return $ Var id' + _ -> return $ Var name + simplify (Constant v) = + return $ Constant v + +simplifyPattern :: + (SimplifiableLore lore, Simplifiable dec) => + PatternT dec -> + SimpleM lore (PatternT dec) +simplifyPattern pat = + Pattern + <$> mapM inspect (patternContextElements pat) + <*> mapM inspect (patternValueElements pat) + where + inspect (PatElem name lore) = PatElem name <$> simplify lore + +instance Simplifiable () where + simplify = pure + +instance Simplifiable VName where + simplify v = do + se <- ST.lookupSubExp v <$> askVtable + case se of + Just (Var v', cs) -> do + changed + usedCerts cs + return v' + _ -> return v + +instance Simplifiable d => Simplifiable (ShapeBase d) where + simplify = fmap Shape . simplify . shapeDims + +instance Simplifiable ExtSize where + simplify (Free se) = Free <$> simplify se + simplify (Ext x) = return $ Ext x + +instance Simplifiable Space where + simplify (ScalarSpace ds t) = ScalarSpace <$> simplify ds <*> pure t + simplify s = pure s + +instance Simplifiable shape => Simplifiable (TypeBase shape u) where + simplify (Array et shape u) = do + shape' <- simplify shape + return $ Array et shape' u + simplify (Mem space) = + Mem <$> simplify space + simplify (Prim bt) = + return $ Prim bt + +instance Simplifiable d => Simplifiable (DimIndex d) where + simplify (DimFix i) = DimFix <$> simplify i + simplify (DimSlice i n s) = DimSlice <$> simplify i <*> simplify n <*> simplify s + +simplifyLambda :: + SimplifiableLore lore => + Lambda lore -> + SimpleM lore (Lambda (Wise lore), Stms (Wise lore)) +simplifyLambda lam = do + par_blocker <- asksEngineEnv $ blockHoistPar . envHoistBlockers + simplifyLambdaMaybeHoist par_blocker lam + +simplifyLambdaNoHoisting :: + SimplifiableLore lore => + Lambda lore -> + SimpleM lore (Lambda (Wise lore)) +simplifyLambdaNoHoisting lam = + fst <$> simplifyLambdaMaybeHoist (isFalse False) lam + +simplifyLambdaMaybeHoist :: + SimplifiableLore lore => + BlockPred (Wise lore) -> + Lambda lore -> + SimpleM lore (Lambda (Wise lore), Stms (Wise lore)) +simplifyLambdaMaybeHoist blocked lam@(Lambda params body rettype) = do + params' <- mapM (traverse simplify) params + let paramnames = namesFromList $ boundByLambda lam + ((lamstms, lamres), hoisted) <- + enterLoop $ + bindLParams params' $ + blockIf (blocked `orIf` hasFree paramnames `orIf` isConsumed) $ + simplifyBody (map (const Observe) rettype) body + body' <- constructBody lamstms lamres + rettype' <- simplify rettype + return (Lambda params' body' rettype', hoisted) + +consumeResult :: [(Diet, SubExp)] -> UT.UsageTable +consumeResult = mconcat . map inspect + where + inspect (Consume, se) = + mconcat $ map UT.consumedUsage $ namesToList $ subExpAliases se + inspect _ = mempty + +instance Simplifiable Certificates where + simplify (Certificates ocs) = Certificates . nub . concat <$> mapM check ocs + where + check idd = do + vv <- ST.lookupSubExp idd <$> askVtable + case vv of + Just (Constant Checked, Certificates cs) -> return cs + Just (Var idd', _) -> return [idd'] + _ -> return [idd] + +insertAllStms :: + SimplifiableLore lore => + SimpleM lore (SimplifiedBody lore Result) -> + SimpleM lore (Body (Wise lore)) +insertAllStms = uncurry constructBody . fst <=< blockIf (isFalse False) + +simplifyFun :: + SimplifiableLore lore => + FunDef lore -> + SimpleM lore (FunDef (Wise lore)) simplifyFun (FunDef entry attrs fname rettype params body) = do rettype' <- simplify rettype params' <- mapM (traverse simplify) params diff --git a/src/Futhark/Optimise/Simplify/Lore.hs b/src/Futhark/Optimise/Simplify/Lore.hs --- a/src/Futhark/Optimise/Simplify/Lore.hs +++ b/src/Futhark/Optimise/Simplify/Lore.hs @@ -1,51 +1,63 @@ -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + -- | Definition of the lore used by the simplification engine. module Futhark.Optimise.Simplify.Lore - ( - Wise - , VarWisdom (..) - , ExpWisdom - , removeStmWisdom - , removeLambdaWisdom - , removeFunDefWisdom - , removeExpWisdom - , removePatternWisdom - , removeBodyWisdom - , removeScopeWisdom - , addScopeWisdom - , addWisdomToPattern - , mkWiseBody - , mkWiseLetStm - , mkWiseExpDec - - , CanBeWise (..) - ) - where + ( Wise, + VarWisdom (..), + ExpWisdom, + removeStmWisdom, + removeLambdaWisdom, + removeFunDefWisdom, + removeExpWisdom, + removePatternWisdom, + removeBodyWisdom, + removeScopeWisdom, + addScopeWisdom, + addWisdomToPattern, + mkWiseBody, + mkWiseLetStm, + mkWiseExpDec, + CanBeWise (..), + ) +where +import Control.Category import Control.Monad.Identity import Control.Monad.Reader import qualified Data.Kind import qualified Data.Map.Strict as M - +import Futhark.Analysis.Rephrase +import Futhark.Binder import Futhark.IR -import Futhark.IR.Prop.Aliases import Futhark.IR.Aliases - (unAliases, AliasDec (..), VarAliases, ConsumedInExp) + ( AliasDec (..), + ConsumedInExp, + VarAliases, + unAliases, + ) import qualified Futhark.IR.Aliases as Aliases -import Futhark.Binder +import Futhark.IR.Prop.Aliases import Futhark.Transform.Rename import Futhark.Transform.Substitute -import Futhark.Analysis.Rephrase +import GHC.Generics (Generic) +import Language.SexpGrammar as Sexp hiding (cons) +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) data Wise lore -- | The wisdom of the let-bound variable. -newtype VarWisdom = VarWisdom { varWisdomAliases :: VarAliases } - deriving (Eq, Ord, Show) +newtype VarWisdom = VarWisdom {varWisdomAliases :: VarAliases} + deriving (Eq, Ord, Show, Generic) +instance SexpIso VarWisdom where + sexpIso = with $ \varwisdom -> sexpIso >>> varwisdom + instance Rename VarWisdom where rename = substituteRename @@ -57,11 +69,20 @@ freeIn' (VarWisdom als) = freeIn' als -- | Wisdom about an expression. -data ExpWisdom = ExpWisdom { _expWisdomConsumed :: ConsumedInExp - , expWisdomFree :: AliasDec - } - deriving (Eq, Ord, Show) +data ExpWisdom = ExpWisdom + { _expWisdomConsumed :: ConsumedInExp, + expWisdomFree :: AliasDec + } + deriving (Eq, Ord, Show, Generic) +instance SexpIso ExpWisdom where + sexpIso = with $ \expwisdom -> + Sexp.list + ( Sexp.el sexpIso + >>> Sexp.el sexpIso + ) + >>> expwisdom + instance FreeIn ExpWisdom where freeIn' = mempty @@ -71,28 +92,36 @@ instance Substitute ExpWisdom where substituteNames substs (ExpWisdom cons free) = ExpWisdom - (substituteNames substs cons) - (substituteNames substs free) + (substituteNames substs cons) + (substituteNames substs free) instance Rename ExpWisdom where rename = substituteRename -- | Wisdom about a body. -data BodyWisdom = BodyWisdom { bodyWisdomAliases :: [VarAliases] - , bodyWisdomConsumed :: ConsumedInExp - , bodyWisdomFree :: AliasDec - } - deriving (Eq, Ord, Show) +data BodyWisdom = BodyWisdom + { bodyWisdomAliases :: [VarAliases], + bodyWisdomConsumed :: ConsumedInExp, + bodyWisdomFree :: AliasDec + } + deriving (Eq, Ord, Show, Generic) +instance SexpIso BodyWisdom where + sexpIso = with $ \bodywisdom -> + Sexp.list + ( Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso + ) + >>> bodywisdom + instance Rename BodyWisdom where rename = substituteRename instance Substitute BodyWisdom where substituteNames substs (BodyWisdom als cons free) = BodyWisdom - (substituteNames substs als) - (substituteNames substs cons) - (substituteNames substs free) + (substituteNames substs als) + (substituteNames substs cons) + (substituteNames substs free) instance FreeIn BodyWisdom where freeIn' (BodyWisdom als cons free) = @@ -101,8 +130,12 @@ instance FreeDec BodyWisdom where precomputed = const . fvNames . unAliases . bodyWisdomFree -instance (Decorations lore, - CanBeWise (Op lore)) => Decorations (Wise lore) where +instance + ( Decorations lore, + CanBeWise (Op lore) + ) => + Decorations (Wise lore) + where type LetDec (Wise lore) = (VarWisdom, LetDec lore) type ExpDec (Wise lore) = (ExpWisdom, ExpDec lore) type BodyDec (Wise lore) = (BodyWisdom, BodyDec lore) @@ -112,9 +145,10 @@ type BranchType (Wise lore) = BranchType lore type Op (Wise lore) = OpWithWisdom (Op lore) -withoutWisdom :: (HasScope (Wise lore) m, Monad m) => - ReaderT (Scope lore) m a -> - m a +withoutWisdom :: + (HasScope (Wise lore) m, Monad m) => + ReaderT (Scope lore) m a -> + m a withoutWisdom m = do scope <- asksScope removeScopeWisdom runReaderT m scope @@ -137,29 +171,33 @@ consumedInBody = unAliases . bodyWisdomConsumed . fst . bodyDec removeWisdom :: CanBeWise (Op lore) => Rephraser Identity (Wise lore) lore -removeWisdom = Rephraser { rephraseExpLore = return . snd - , rephraseLetBoundLore = return . snd - , rephraseBodyLore = return . snd - , rephraseFParamLore = return - , rephraseLParamLore = return - , rephraseRetType = return - , rephraseBranchType = return - , rephraseOp = return . removeOpWisdom - } +removeWisdom = + Rephraser + { rephraseExpLore = return . snd, + rephraseLetBoundLore = return . snd, + rephraseBodyLore = return . snd, + rephraseFParamLore = return, + rephraseLParamLore = return, + rephraseRetType = return, + rephraseBranchType = return, + rephraseOp = return . removeOpWisdom + } removeScopeWisdom :: Scope (Wise lore) -> Scope lore removeScopeWisdom = M.map unAlias - where unAlias (LetName (_, dec)) = LetName dec - unAlias (FParamName dec) = FParamName dec - unAlias (LParamName dec) = LParamName dec - unAlias (IndexName it) = IndexName it + where + unAlias (LetName (_, dec)) = LetName dec + unAlias (FParamName dec) = FParamName dec + unAlias (LParamName dec) = LParamName dec + unAlias (IndexName it) = IndexName it addScopeWisdom :: Scope lore -> Scope (Wise lore) addScopeWisdom = M.map alias - where alias (LetName dec) = LetName (VarWisdom mempty, dec) - alias (FParamName dec) = FParamName dec - alias (LParamName dec) = LParamName dec - alias (IndexName it) = IndexName it + where + alias (LetName dec) = LetName (VarWisdom mempty, dec) + alias (FParamName dec) = FParamName dec + alias (LParamName dec) = LParamName dec + alias (IndexName it) = IndexName it removeFunDefWisdom :: CanBeWise (Op lore) => FunDef (Wise lore) -> FunDef lore removeFunDefWisdom = runIdentity . rephraseFunDef removeWisdom @@ -179,42 +217,64 @@ removePatternWisdom :: PatternT (VarWisdom, a) -> PatternT a removePatternWisdom = runIdentity . rephrasePattern (return . snd) -addWisdomToPattern :: (ASTLore lore, CanBeWise (Op lore)) => - Pattern lore - -> Exp (Wise lore) - -> Pattern (Wise lore) +addWisdomToPattern :: + (ASTLore lore, CanBeWise (Op lore)) => + Pattern lore -> + Exp (Wise lore) -> + Pattern (Wise lore) addWisdomToPattern pat e = Pattern (map f ctx) (map f val) - where (ctx, val) = Aliases.mkPatternAliases pat e - f pe = let (als, dec) = patElemDec pe - in pe `setPatElemLore` (VarWisdom als, dec) + where + (ctx, val) = Aliases.mkPatternAliases pat e + f pe = + let (als, dec) = patElemDec pe + in pe `setPatElemLore` (VarWisdom als, dec) -mkWiseBody :: (ASTLore lore, CanBeWise (Op lore)) => - BodyDec lore -> Stms (Wise lore) -> Result -> Body (Wise lore) +mkWiseBody :: + (ASTLore lore, CanBeWise (Op lore)) => + BodyDec lore -> + Stms (Wise lore) -> + Result -> + Body (Wise lore) mkWiseBody innerlore bnds res = - Body (BodyWisdom aliases consumed (AliasDec $ freeIn $ freeInStmsAndRes bnds res), - innerlore) bnds res - where (aliases, consumed) = Aliases.mkBodyAliases bnds res + Body + ( BodyWisdom aliases consumed (AliasDec $ freeIn $ freeInStmsAndRes bnds res), + innerlore + ) + bnds + res + where + (aliases, consumed) = Aliases.mkBodyAliases bnds res -mkWiseLetStm :: (ASTLore lore, CanBeWise (Op lore)) => - Pattern lore - -> StmAux (ExpDec lore) -> Exp (Wise lore) - -> Stm (Wise lore) +mkWiseLetStm :: + (ASTLore lore, CanBeWise (Op lore)) => + Pattern lore -> + StmAux (ExpDec lore) -> + Exp (Wise lore) -> + Stm (Wise lore) mkWiseLetStm pat (StmAux cs attrs dec) e = let pat' = addWisdomToPattern pat e - in Let pat' (StmAux cs attrs $ mkWiseExpDec pat' dec e) e + in Let pat' (StmAux cs attrs $ mkWiseExpDec pat' dec e) e -mkWiseExpDec :: (ASTLore lore, CanBeWise (Op lore)) => - Pattern (Wise lore) -> ExpDec lore -> Exp (Wise lore) - -> ExpDec (Wise lore) +mkWiseExpDec :: + (ASTLore lore, CanBeWise (Op lore)) => + Pattern (Wise lore) -> + ExpDec lore -> + Exp (Wise lore) -> + ExpDec (Wise lore) mkWiseExpDec pat explore e = - (ExpWisdom - (AliasDec $ consumedInExp e) - (AliasDec $ freeIn pat <> freeIn explore <> freeIn e), - explore) + ( ExpWisdom + (AliasDec $ consumedInExp e) + (AliasDec $ freeIn pat <> freeIn explore <> freeIn e), + explore + ) -instance (Bindable lore, - CanBeWise (Op lore)) => Bindable (Wise lore) where +instance + ( Bindable lore, + CanBeWise (Op lore) + ) => + Bindable (Wise lore) + where mkExpPat ctx val e = addWisdomToPattern (mkExpPat ctx val $ removeExpWisdom e) e @@ -229,10 +289,14 @@ mkBody bnds res = let Body bodylore _ _ = mkBody (fmap removeStmWisdom bnds) res - in mkWiseBody bodylore bnds res + in mkWiseBody bodylore bnds res -class (AliasedOp (OpWithWisdom op), - IsOp (OpWithWisdom op)) => CanBeWise op where +class + ( AliasedOp (OpWithWisdom op), + IsOp (OpWithWisdom op) + ) => + CanBeWise op + where type OpWithWisdom op :: Data.Kind.Type removeOpWisdom :: OpWithWisdom op -> op diff --git a/src/Futhark/Optimise/Simplify/Rule.hs b/src/Futhark/Optimise/Simplify/Rule.hs --- a/src/Futhark/Optimise/Simplify/Rule.hs +++ b/src/Futhark/Optimise/Simplify/Rule.hs @@ -1,8 +1,9 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + -- | This module defines the concept of a simplification rule for -- bindings. The intent is that you pass some context (such as symbol -- table) and a binding, and is given back a sequence of bindings that @@ -13,57 +14,63 @@ -- fusion algorithm in @Futhark.Optimise.Fusion.Fusion@, which must be implemented -- as its own pass. module Futhark.Optimise.Simplify.Rule - ( -- * The rule monad - RuleM - , cannotSimplify - , liftMaybe + ( -- * The rule monad + RuleM, + cannotSimplify, + liftMaybe, - -- * Rule definition - , Rule(..) - , SimplificationRule(..) - , RuleGeneric - , RuleBasicOp - , RuleIf - , RuleDoLoop + -- * Rule definition + Rule (..), + SimplificationRule (..), + RuleGeneric, + RuleBasicOp, + RuleIf, + RuleDoLoop, - -- * Top-down rules - , TopDown - , TopDownRule - , TopDownRuleGeneric - , TopDownRuleBasicOp - , TopDownRuleIf - , TopDownRuleDoLoop - , TopDownRuleOp + -- * Top-down rules + TopDown, + TopDownRule, + TopDownRuleGeneric, + TopDownRuleBasicOp, + TopDownRuleIf, + TopDownRuleDoLoop, + TopDownRuleOp, - -- * Bottom-up rules - , BottomUp - , BottomUpRule - , BottomUpRuleGeneric - , BottomUpRuleBasicOp - , BottomUpRuleIf - , BottomUpRuleDoLoop - , BottomUpRuleOp + -- * Bottom-up rules + BottomUp, + BottomUpRule, + BottomUpRuleGeneric, + BottomUpRuleBasicOp, + BottomUpRuleIf, + BottomUpRuleDoLoop, + BottomUpRuleOp, - -- * Assembling rules - , RuleBook - , ruleBook + -- * Assembling rules + RuleBook, + ruleBook, - -- * Applying rules - , topDownSimplifyStm - , bottomUpSimplifyStm - ) where + -- * Applying rules + topDownSimplifyStm, + bottomUpSimplifyStm, + ) +where import Control.Monad.State - import qualified Futhark.Analysis.SymbolTable as ST import qualified Futhark.Analysis.UsageTable as UT -import Futhark.IR import Futhark.Binder +import Futhark.IR -- | The monad in which simplification rules are evaluated. newtype RuleM lore a = RuleM (BinderT lore (StateT VNameSource Maybe) a) - deriving (Functor, Applicative, Monad, - MonadFreshNames, HasScope lore, LocalScope lore) + deriving + ( Functor, + Applicative, + Monad, + MonadFreshNames, + HasScope lore, + LocalScope lore + ) instance (ASTLore lore, BinderOps lore) => MonadBinder (RuleM lore) where type Lore (RuleM lore) = lore @@ -76,8 +83,11 @@ -- | Execute a 'RuleM' action. If succesful, returns the result and a -- list of new bindings. -simplify :: Scope lore -> VNameSource -> Rule lore - -> Maybe (Stms lore, VNameSource) +simplify :: + Scope lore -> + VNameSource -> + Rule lore -> + Maybe (Stms lore, VNameSource) simplify _ _ Skip = Nothing simplify scope src (Simplify (RuleM m)) = runStateT (runBinderT_ m scope) src @@ -90,43 +100,73 @@ liftMaybe (Just x) = return x -- | An efficient way of encoding whether a simplification rule should even be attempted. -data Rule lore = Simplify (RuleM lore ()) -- ^ Give it a shot. - | Skip -- ^ Don't bother. +data Rule lore + = -- | Give it a shot. + Simplify (RuleM lore ()) + | -- | Don't bother. + Skip type RuleGeneric lore a = a -> Stm lore -> Rule lore -type RuleBasicOp lore a = (a -> Pattern lore -> StmAux (ExpDec lore) -> - BasicOp -> Rule lore) -type RuleIf lore a = a -> Pattern lore -> StmAux (ExpDec lore) -> - (SubExp, BodyT lore, BodyT lore, - IfDec (BranchType lore)) -> - Rule lore -type RuleDoLoop lore a = a -> Pattern lore -> StmAux (ExpDec lore) -> - ([(FParam lore, SubExp)], [(FParam lore, SubExp)], - LoopForm lore, BodyT lore) -> - Rule lore -type RuleOp lore a = a -> Pattern lore -> StmAux (ExpDec lore) -> - Op lore -> Rule lore +type RuleBasicOp lore a = + ( a -> + Pattern lore -> + StmAux (ExpDec lore) -> + BasicOp -> + Rule lore + ) + +type RuleIf lore a = + a -> + Pattern lore -> + StmAux (ExpDec lore) -> + ( SubExp, + BodyT lore, + BodyT lore, + IfDec (BranchType lore) + ) -> + Rule lore + +type RuleDoLoop lore a = + a -> + Pattern lore -> + StmAux (ExpDec lore) -> + ( [(FParam lore, SubExp)], + [(FParam lore, SubExp)], + LoopForm lore, + BodyT lore + ) -> + Rule lore + +type RuleOp lore a = + a -> + Pattern lore -> + StmAux (ExpDec lore) -> + Op lore -> + Rule lore + -- | A simplification rule takes some argument and a statement, and -- tries to simplify the statement. -data SimplificationRule lore a = RuleGeneric (RuleGeneric lore a) - | RuleBasicOp (RuleBasicOp lore a) - | RuleIf (RuleIf lore a) - | RuleDoLoop (RuleDoLoop lore a) - | RuleOp (RuleOp lore a) +data SimplificationRule lore a + = RuleGeneric (RuleGeneric lore a) + | RuleBasicOp (RuleBasicOp lore a) + | RuleIf (RuleIf lore a) + | RuleDoLoop (RuleDoLoop lore a) + | RuleOp (RuleOp lore a) -- | A collection of rules grouped by which forms of statements they -- may apply to. -data Rules lore a = Rules { rulesAny :: [SimplificationRule lore a] - , rulesBasicOp :: [SimplificationRule lore a] - , rulesIf :: [SimplificationRule lore a] - , rulesDoLoop :: [SimplificationRule lore a] - , rulesOp :: [SimplificationRule lore a] - } +data Rules lore a = Rules + { rulesAny :: [SimplificationRule lore a], + rulesBasicOp :: [SimplificationRule lore a], + rulesIf :: [SimplificationRule lore a], + rulesDoLoop :: [SimplificationRule lore a], + rulesOp :: [SimplificationRule lore a] + } instance Semigroup (Rules lore a) where Rules as1 bs1 cs1 ds1 es1 <> Rules as2 bs2 cs2 ds2 es2 = - Rules (as1<>as2) (bs1<>bs2) (cs1<>cs2) (ds1<>ds2) (es1<>es2) + Rules (as1 <> as2) (bs1 <> bs2) (cs1 <> cs2) (ds1 <> ds2) (es1 <> es2) instance Monoid (Rules lore a) where mempty = Rules mempty mempty mempty mempty mempty @@ -136,10 +176,15 @@ type TopDown lore = ST.SymbolTable lore type TopDownRuleGeneric lore = RuleGeneric lore (TopDown lore) + type TopDownRuleBasicOp lore = RuleBasicOp lore (TopDown lore) + type TopDownRuleIf lore = RuleIf lore (TopDown lore) + type TopDownRuleDoLoop lore = RuleDoLoop lore (TopDown lore) + type TopDownRuleOp lore = RuleOp lore (TopDown lore) + type TopDownRule lore = SimplificationRule lore (TopDown lore) -- | Context for a rule applied during bottom-up traversal of the @@ -147,10 +192,15 @@ type BottomUp lore = (ST.SymbolTable lore, UT.UsageTable) type BottomUpRuleGeneric lore = RuleGeneric lore (BottomUp lore) + type BottomUpRuleBasicOp lore = RuleBasicOp lore (BottomUp lore) + type BottomUpRuleIf lore = RuleIf lore (BottomUp lore) + type BottomUpRuleDoLoop lore = RuleDoLoop lore (BottomUp lore) + type BottomUpRuleOp lore = RuleOp lore (BottomUp lore) + type BottomUpRule lore = SimplificationRule lore (BottomUp lore) -- | A collection of top-down rules. @@ -160,54 +210,60 @@ type BottomUpRules lore = Rules lore (BottomUp lore) -- | A collection of both top-down and bottom-up rules. -data RuleBook lore = RuleBook { bookTopDownRules :: TopDownRules lore - , bookBottomUpRules :: BottomUpRules lore - } +data RuleBook lore = RuleBook + { bookTopDownRules :: TopDownRules lore, + bookBottomUpRules :: BottomUpRules lore + } instance Semigroup (RuleBook lore) where - RuleBook ts1 bs1 <> RuleBook ts2 bs2 = RuleBook (ts1<>ts2) (bs1<>bs2) + RuleBook ts1 bs1 <> RuleBook ts2 bs2 = RuleBook (ts1 <> ts2) (bs1 <> bs2) instance Monoid (RuleBook lore) where mempty = RuleBook mempty mempty -- | Construct a rule book from a collection of rules. -ruleBook :: [TopDownRule m] - -> [BottomUpRule m] - -> RuleBook m +ruleBook :: + [TopDownRule m] -> + [BottomUpRule m] -> + RuleBook m ruleBook topdowns bottomups = RuleBook (groupRules topdowns) (groupRules bottomups) - where groupRules :: [SimplificationRule m a] -> Rules m a - groupRules rs = Rules rs - (filter forBasicOp rs) - (filter forIf rs) - (filter forDoLoop rs) - (filter forOp rs) + where + groupRules :: [SimplificationRule m a] -> Rules m a + groupRules rs = + Rules + rs + (filter forBasicOp rs) + (filter forIf rs) + (filter forDoLoop rs) + (filter forOp rs) - forBasicOp RuleBasicOp{} = True - forBasicOp RuleGeneric{} = True - forBasicOp _ = False + forBasicOp RuleBasicOp {} = True + forBasicOp RuleGeneric {} = True + forBasicOp _ = False - forIf RuleIf{} = True - forIf RuleGeneric{} = True - forIf _ = False + forIf RuleIf {} = True + forIf RuleGeneric {} = True + forIf _ = False - forDoLoop RuleDoLoop{} = True - forDoLoop RuleGeneric{} = True - forDoLoop _ = False + forDoLoop RuleDoLoop {} = True + forDoLoop RuleGeneric {} = True + forDoLoop _ = False - forOp RuleOp{} = True - forOp RuleGeneric{} = True - forOp _ = False + forOp RuleOp {} = True + forOp RuleGeneric {} = True + forOp _ = False -- | @simplifyStm lookup bnd@ performs simplification of the -- binding @bnd@. If simplification is possible, a replacement list -- of bindings is returned, that bind at least the same names as the -- original binding (and possibly more, for intermediate results). -topDownSimplifyStm :: (MonadFreshNames m, HasScope lore m) => - RuleBook lore - -> ST.SymbolTable lore - -> Stm lore - -> m (Maybe (Stms lore)) +topDownSimplifyStm :: + (MonadFreshNames m, HasScope lore m) => + RuleBook lore -> + ST.SymbolTable lore -> + Stm lore -> + m (Maybe (Stms lore)) topDownSimplifyStm = applyRules . bookTopDownRules -- | @simplifyStm uses bnd@ performs simplification of the binding @@ -215,19 +271,21 @@ -- bindings is returned, that bind at least the same names as the -- original binding (and possibly more, for intermediate results). -- The first argument is the set of names used after this binding. -bottomUpSimplifyStm :: (MonadFreshNames m, HasScope lore m) => - RuleBook lore - -> (ST.SymbolTable lore, UT.UsageTable) - -> Stm lore - -> m (Maybe (Stms lore)) +bottomUpSimplifyStm :: + (MonadFreshNames m, HasScope lore m) => + RuleBook lore -> + (ST.SymbolTable lore, UT.UsageTable) -> + Stm lore -> + m (Maybe (Stms lore)) bottomUpSimplifyStm = applyRules . bookBottomUpRules rulesForStm :: Stm lore -> Rules lore a -> [SimplificationRule lore a] -rulesForStm stm = case stmExp stm of BasicOp{} -> rulesBasicOp - DoLoop{} -> rulesDoLoop - Op{} -> rulesOp - If{} -> rulesIf - _ -> rulesAny +rulesForStm stm = case stmExp stm of + BasicOp {} -> rulesBasicOp + DoLoop {} -> rulesDoLoop + Op {} -> rulesOp + If {} -> rulesIf + _ -> rulesAny applyRule :: SimplificationRule lore a -> a -> Stm lore -> Rule lore applyRule (RuleGeneric f) a stm = f a stm @@ -241,19 +299,21 @@ applyRule _ _ _ = Skip -applyRules :: (MonadFreshNames m, HasScope lore m) => - Rules lore a -> a -> Stm lore - -> m (Maybe (Stms lore)) +applyRules :: + (MonadFreshNames m, HasScope lore m) => + Rules lore a -> + a -> + Stm lore -> + m (Maybe (Stms lore)) applyRules all_rules context stm = do scope <- askScope modifyNameSource $ \src -> - let applyRules' [] = Nothing - applyRules' (rule:rules) = + let applyRules' [] = Nothing + applyRules' (rule : rules) = case simplify scope src (applyRule rule context stm) of Just x -> Just x Nothing -> applyRules' rules - - in case applyRules' $ rulesForStm stm all_rules of - Just (stms, src') -> (Just stms, src') - Nothing -> (Nothing, src) + in case applyRules' $ rulesForStm stm all_rules of + Just (stms, src') -> (Just stms, src') + Nothing -> (Nothing, src) 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,1326 +1,1486 @@ {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE Safe #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE OverloadedStrings #-} --- | 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, isSuffixOf, partition, sort) -import Data.Maybe -import qualified Data.Map.Strict as M - -import qualified Futhark.Analysis.SymbolTable as ST -import qualified Futhark.Analysis.UsageTable as UT -import Futhark.Analysis.DataDependencies -import Futhark.Optimise.Simplify.ClosedForm -import Futhark.Optimise.Simplify.Rule -import Futhark.Analysis.PrimExp.Convert -import Futhark.IR -import Futhark.IR.Prop.Aliases -import Futhark.Transform.Rename -import Futhark.Construct -import Futhark.Util - -topDownRules :: (BinderOps lore, Aliased lore) => [TopDownRule lore] -topDownRules = [ RuleDoLoop hoistLoopInvariantMergeVariables - , RuleDoLoop simplifyClosedFormLoop - , RuleDoLoop simplifyKnownIterationLoop - , RuleDoLoop simplifyLoopVariables - , RuleGeneric constantFoldPrimFun - , RuleIf ruleIf - , RuleIf hoistBranchInvariant - , RuleBasicOp ruleBasicOp - ] - -bottomUpRules :: BinderOps lore => [BottomUpRule lore] -bottomUpRules = [ RuleDoLoop removeRedundantMergeVariables - , RuleIf removeDeadBranchResult - , RuleBasicOp simplifyIndex - , RuleBasicOp simplifyConcat - ] - -asInt32PrimExp :: PrimExp v -> PrimExp v -asInt32PrimExp pe - | IntType it <- primExpType pe, it /= Int32 = - sExt Int32 pe - | otherwise = - pe - --- | 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 _ 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, [], []) $ - zip 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 - ((mergeParam,mergeInit), resExp) - (invariant, explpat', merge', resExps) - | not (unique (paramDeclType mergeParam)) || - arrayRank (paramDeclType mergeParam) == 1, - isInvariant resExp, - -- 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 the corresponding - -- subexp in the result is EITHER: - -- - -- (0) a variable of the same name as the parameter, where - -- all existential parameters are already known to be - -- invariant - isInvariant (Var v2) - | paramName mergeParam == v2 = - allExistentialInvariant - (namesFromList $ map (identName . fst) invariant) mergeParam - -- (1) or identical to the initial value of the parameter. - isInvariant _ = mergeInit == resExp - - checkInvariance ((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 _ bound [], body) = - Simplify $ loopClosedForm pat val (oneName i) 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 Int32 0) w (intConst Int32 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 - -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 Int32 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{} e1 e2) - | isCt0 e2 = subExpRes e1 - -- Cases for simplifying (a+b)-b and permutations. - | Var v1 <- e1, - Just (BasicOp (BinOp Add{} e1_a e1_b), cs) <- look v1, - e1_a == e2 = Just (SubExp e1_b, cs) - | Var v1 <- e1, - Just (BasicOp (BinOp Add{} e1_a e1_b), cs) <- look v1, - e1_b == e2 = Just (SubExp e1_a, cs) - | Var v2 <- e2, - Just (BasicOp (BinOp Add{} e2_a e2_b), cs) <- look v2, - e2_a == e1 = Just (SubExp e2_b, cs) - | Var v2 <- e1, - Just (BasicOp (BinOp Add{} e2_a e2_b), cs) <- look v2, - e2_b == e1 = Just (SubExp 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 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 $ - fmap (SubExpResult cs) $ toSubExp "index_iota" $ - sExt to_it (primExpFromSubExp (IntType from_it) ii) - * primExpFromSubExp (IntType to_it) s - + 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 - i_offset'' <- toSubExp "iota_offset" $ - primExpFromSubExp (IntType to_it) x + - primExpFromSubExp (IntType to_it) s * - primExpFromSubExp (IntType to_it) i_offset' - i_stride'' <- letSubExp "iota_offset" $ - BasicOp $ BinOp (Mul Int32 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 Int32 OverflowWrap) i o - letSubExp "rot_i" (BasicOp $ BinOp (SMod Int32 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::Int32)) 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, - 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) - | 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 Int32 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 Int32) start i - (thisres, thisbnds) <- collectStms $ do - i' <- letSubExp "index_concat_i" $ BasicOp $ BinOp (Sub Int32 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 (Int32Value 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::Int32)) : inds' - - _ -> Nothing - - where defOf v = do (BasicOp op, def_cs) <- ST.lookupExp v vtable - return (op, def_cs) - - -- | A crude heuristic for determining when a PrimExp is - -- worth inlining over keeping it in an array and reading it - -- from memory. - 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 - -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 - --- 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) - --- If concatenating a bunch of array literals (or equivalent --- replicate), just construct the array literal instead. -simplifyConcat (vtable, _) pat aux (Concat 0 x xs _) - | Just (vs, vcs) <- unzip <$> mapM isArrayLit (x:xs) = Simplify $ do - rt <- rowType <$> lookupType x - certifying (mconcat vcs) $ auxing aux $ - letBind pat $ BasicOp $ ArrayLit (concat vs) rt - where isArrayLit v - | Just (Replicate shape se, vcs) <- ST.lookupBasicOp v vtable, - unitShape shape = Just ([se], vcs) - | Just (ArrayLit ses _, vcs) <- ST.lookupBasicOp v vtable = - Just (ses, vcs) - | otherwise = - Nothing - - unitShape = (==Shape [Constant $ IntValue $ Int32Value 1]) - -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 Int32 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 - --- | Turn in-place updates that replace an entire array into just --- array literals. -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 - --- | Simplify a chain of in-place updates and copies. This chain is --- often produced by in-place lowering. -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 - --- | If we are comparing X against the result of a branch of the form --- @if P then Y else Z@ then replace comparison with '(P && X == Y) || --- (!P && X == Z'). This may allow us to get rid of a branch, and the --- extra comparisons may be constant-folded out. Question: maybe we --- should have some more checks to ensure that we only do this if that --- is actually the case, such as if we will obtain at least one --- constant-to-constant comparison? -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 - --- | Turn array literals with identical elements into replicates. -ruleBasicOp _ pat _ (ArrayLit (se:ses) _) - | all (==se) ses = - Simplify $ let n = constant (fromIntegral (length ses) + 1 :: Int32) - 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 (primExpFromSubExp int32) oldshape) - (map (primExpFromSubExp int32) $ newDims newshape) - (map (primExpFromSubExp int32) inds) - new_inds' <- - mapM (toSubExp "new_index" . asInt32PrimExp) 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 Int32 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 Int32 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 Int32 1) (intConst Int32 1)] - slice_y' = slice_y_bef ++ [DimSlice j (intConst Int32 1) (intConst Int32 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 (Int32Value 0)) <- x, - Var v <- y, - Just _ <- ST.lookupLoopVar v vtable = - Simplify $ auxing aux $ letBind pat $ BasicOp $ SubExp $ constant True - --- Simplify away i BottomUpRuleIf lore -removeDeadBranchResult (_, used) pat _ (e1, tb, fb, IfDec rettype ifsort) - | -- Only if there is no existential context... - patternSize pat == length rettype, - -- Figure out which of the names in 'pat' are used... - patused <- map (`UT.isUsedDirectly` used) $ patternNames pat, - -- If they are not all used, then this rule applies. - not (and patused) = - -- Remove the parts of the branch-results that correspond to dead - -- return value bindings. Note that this leaves dead code in the - -- branch bodies, but that will be removed later. - let tses = bodyResult tb - fses = bodyResult fb - pick :: [a] -> [a] - pick = map snd . filter fst . zip patused - tb' = tb { bodyResult = pick tses } - fb' = fb { bodyResult = pick fses } - pat' = pick $ patternElements pat - rettype' = pick rettype - in Simplify $ letBind (Pattern [] pat') $ If e1 tb' fb' $ IfDec rettype' ifsort - - | otherwise = Skip - +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE Safe #-} +{-# 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 {} e1 e2) + | isCt0 e2 = subExpRes e1 + -- Cases for simplifying (a+b)-b and permutations. + | Var v1 <- e1, + Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1, + e1_a == e2 = + Just (SubExp e1_b, cs) + | Var v1 <- e1, + Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1, + e1_b == e2 = + Just (SubExp e1_a, cs) + | Var v2 <- e2, + Just (BasicOp (BinOp Add {} e2_a e2_b), cs) <- look v2, + e2_a == e1 = + Just (SubExp e2_b, cs) + | Var v2 <- e1, + Just (BasicOp (BinOp Add {} e2_a e2_b), cs) <- look v2, + e2_b == e1 = + Just (SubExp 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, + 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 BottomUpRuleIf lore +removeDeadBranchResult (_, used) pat _ (e1, tb, fb, IfDec rettype ifsort) + | -- Only if there is no existential context... + patternSize pat == length rettype, + -- Figure out which of the names in 'pat' are used... + patused <- map (`UT.isUsedDirectly` used) $ patternNames pat, + -- If they are not all used, then this rule applies. + not (and patused) = + -- Remove the parts of the branch-results that correspond to dead + -- return value bindings. Note that this leaves dead code in the + -- branch bodies, but that will be removed later. + let tses = bodyResult tb + fses = bodyResult fb + pick :: [a] -> [a] + pick = map snd . filter fst . zip patused + tb' = tb {bodyResult = pick tses} + fb' = fb {bodyResult = pick fses} + pat' = pick $ patternElements pat + rettype' = pick rettype + in Simplify $ letBind (Pattern [] pat') $ If e1 tb' fb' $ IfDec rettype' ifsort + | otherwise = Skip -- Some helper functions diff --git a/src/Futhark/Optimise/Sink.hs b/src/Futhark/Optimise/Sink.hs --- a/src/Futhark/Optimise/Sink.hs +++ b/src/Futhark/Optimise/Sink.hs @@ -1,5 +1,6 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + -- | "Sinking" is conceptually the opposite of hoisting. The idea is -- to take code that looks like this: -- @@ -41,14 +42,12 @@ -- This pass is defined on the Kernels representation. This is not -- because we do anything kernel-specific here, but simply because -- more explicit indexing is going on after SOACs are gone. - module Futhark.Optimise.Sink (sink) where import Control.Monad.State import Data.List (foldl') import qualified Data.Map as M import qualified Data.Set as S - import qualified Futhark.Analysis.Alias as Alias import qualified Futhark.Analysis.SymbolTable as ST import Futhark.IR.Aliases @@ -56,8 +55,11 @@ import Futhark.Pass type SinkLore = Aliases Kernels + type SymbolTable = ST.SymbolTable SinkLore + type Sinking = M.Map VName (Stm SinkLore) + type Sunk = S.Set VName -- | Given a statement, compute how often each of its free variables @@ -68,120 +70,141 @@ case stmExp stm of If cond tbranch fbranch _ -> free cond 1 `comb` free tbranch 1 `comb` free fbranch 1 - Op{} -> free stm 2 - DoLoop{} -> free stm 2 + Op {} -> free stm 2 + DoLoop {} -> free stm 2 _ -> free stm 1 - where free x k = M.fromList $ zip (namesToList $ freeIn x) $ repeat k - comb = M.unionWith (+) + where + free x k = M.fromList $ zip (namesToList $ freeIn x) $ repeat k + comb = M.unionWith (+) -optimiseBranch :: SymbolTable -> Sinking -> Body SinkLore - -> (Body SinkLore, Sunk) +optimiseBranch :: + SymbolTable -> + Sinking -> + Body SinkLore -> + (Body SinkLore, Sunk) optimiseBranch vtable sinking (Body dec stms res) = let (stms', stms_sunk) = optimiseStms vtable sinking' stms $ freeIn res - in (Body dec (sunk_stms <> stms') res, - sunk <> stms_sunk) - where free_in_stms = freeIn stms <> freeIn res - (sinking_here, sinking') = M.partitionWithKey sunkHere sinking - sunk_stms = stmsFromList $ M.elems sinking_here - sunkHere v stm = - v `nameIn` free_in_stms && - all (`ST.available` vtable) (namesToList (freeIn stm)) - sunk = S.fromList $ concatMap (patternNames . stmPattern) sunk_stms + in ( Body dec (sunk_stms <> stms') res, + sunk <> stms_sunk + ) + where + free_in_stms = freeIn stms <> freeIn res + (sinking_here, sinking') = M.partitionWithKey sunkHere sinking + sunk_stms = stmsFromList $ M.elems sinking_here + sunkHere v stm = + v `nameIn` free_in_stms + && all (`ST.available` vtable) (namesToList (freeIn stm)) + sunk = S.fromList $ concatMap (patternNames . stmPattern) sunk_stms -optimiseStms :: SymbolTable -> Sinking -> Stms SinkLore -> Names - -> (Stms SinkLore, Sunk) +optimiseStms :: + SymbolTable -> + Sinking -> + Stms SinkLore -> + Names -> + (Stms SinkLore, Sunk) optimiseStms init_vtable init_sinking all_stms free_in_res = let (all_stms', sunk) = optimiseStms' init_vtable init_sinking $ stmsToList all_stms - in (stmsFromList all_stms', sunk) + in (stmsFromList all_stms', sunk) where - multiplicities = foldl' (M.unionWith (+)) - (M.fromList (zip (namesToList free_in_res) (repeat 1))) - (map multiplicity $ stmsToList all_stms) + multiplicities = + foldl' + (M.unionWith (+)) + (M.fromList (zip (namesToList free_in_res) (repeat 1))) + (map multiplicity $ stmsToList all_stms) optimiseStms' _ _ [] = ([], mempty) - optimiseStms' vtable sinking (stm : stms) - | BasicOp Index{} <- stmExp stm, + | BasicOp Index {} <- stmExp stm, [pe] <- patternElements (stmPattern stm), primType $ patElemType pe, - maybe True (==1) $ M.lookup (patElemName pe) multiplicities = - let (stms', sunk) = - optimiseStms' vtable' (M.insert (patElemName pe) stm sinking) stms - in if patElemName pe `S.member` sunk - then (stms', sunk) - else (stm : stms', sunk) - + maybe True (== 1) $ M.lookup (patElemName pe) multiplicities = + let (stms', sunk) = + optimiseStms' vtable' (M.insert (patElemName pe) stm sinking) stms + in if patElemName pe `S.member` sunk + then (stms', sunk) + else (stm : stms', sunk) | If cond tbranch fbranch ret <- stmExp stm = - let (tbranch', tsunk) = optimiseBranch vtable sinking tbranch - (fbranch', fsunk) = optimiseBranch vtable sinking fbranch - (stms', sunk) = optimiseStms' vtable' sinking stms - in (stm { stmExp = If cond tbranch' fbranch' ret } : stms', - tsunk <> fsunk <> sunk) - + let (tbranch', tsunk) = optimiseBranch vtable sinking tbranch + (fbranch', fsunk) = optimiseBranch vtable sinking fbranch + (stms', sunk) = optimiseStms' vtable' sinking stms + in ( stm {stmExp = If cond tbranch' fbranch' ret} : stms', + tsunk <> fsunk <> sunk + ) | Op (SegOp op) <- stmExp stm = - let scope = scopeOfSegSpace $ segSpace op - (stms', stms_sunk) = optimiseStms' vtable' sinking stms - (op', op_sunk) = runState (mapSegOpM (opMapper scope) op) mempty - in (stm { stmExp = Op (SegOp op') } : stms', - stms_sunk <> op_sunk) - + let scope = scopeOfSegSpace $ segSpace op + (stms', stms_sunk) = optimiseStms' vtable' sinking stms + (op', op_sunk) = runState (mapSegOpM (opMapper scope) op) mempty + in ( stm {stmExp = Op (SegOp op')} : stms', + stms_sunk <> op_sunk + ) | otherwise = - let (stms', stms_sunk) = optimiseStms' vtable' sinking stms - (e', stm_sunk) = runState (mapExpM mapper (stmExp stm)) mempty - in (stm { stmExp = e' } : stms', - stm_sunk <> stms_sunk) - - where vtable' = ST.insertStm stm vtable - mapper = - identityMapper - { mapOnBody = \scope body -> do - let (body', sunk) = - optimiseBody (ST.fromScope scope <> vtable) sinking body - modify (<>sunk) - return body' - } + let (stms', stms_sunk) = optimiseStms' vtable' sinking stms + (e', stm_sunk) = runState (mapExpM mapper (stmExp stm)) mempty + in ( stm {stmExp = e'} : stms', + stm_sunk <> stms_sunk + ) + where + vtable' = ST.insertStm stm vtable + mapper = + identityMapper + { mapOnBody = \scope body -> do + let (body', sunk) = + optimiseBody (ST.fromScope scope <> vtable) sinking body + modify (<> sunk) + return body' + } - opMapper scope = - identitySegOpMapper - { mapOnSegOpLambda = \lam -> do - let (body, sunk) = - optimiseBody op_vtable sinking $ + opMapper scope = + identitySegOpMapper + { mapOnSegOpLambda = \lam -> do + let (body, sunk) = + optimiseBody op_vtable sinking $ lambdaBody lam - modify (<>sunk) - return lam { lambdaBody = body } - - , mapOnSegOpBody = \body -> do - let (body', sunk) = - optimiseKernelBody op_vtable sinking body - modify (<>sunk) - return body' - } - where op_vtable = ST.fromScope scope <> vtable + modify (<> sunk) + return lam {lambdaBody = body}, + mapOnSegOpBody = \body -> do + let (body', sunk) = + optimiseKernelBody op_vtable sinking body + modify (<> sunk) + return body' + } + where + op_vtable = ST.fromScope scope <> vtable -optimiseBody :: SymbolTable -> Sinking -> Body SinkLore - -> (Body SinkLore, Sunk) +optimiseBody :: + SymbolTable -> + Sinking -> + Body SinkLore -> + (Body SinkLore, Sunk) optimiseBody vtable sinking (Body dec stms res) = let (stms', sunk) = optimiseStms vtable sinking stms $ freeIn res - in (Body dec stms' res, sunk) + in (Body dec stms' res, sunk) -optimiseKernelBody :: SymbolTable -> Sinking -> KernelBody SinkLore - -> (KernelBody SinkLore, Sunk) +optimiseKernelBody :: + SymbolTable -> + Sinking -> + KernelBody SinkLore -> + (KernelBody SinkLore, Sunk) optimiseKernelBody vtable sinking (KernelBody dec stms res) = let (stms', sunk) = optimiseStms vtable sinking stms $ freeIn res - in (KernelBody dec stms' res, sunk) + in (KernelBody dec stms' res, sunk) -- | The pass definition. sink :: Pass Kernels Kernels -sink = Pass "sink" "move memory loads closer to their uses" $ - fmap removeProgAliases . - intraproceduralTransformationWithConsts onConsts onFun . - Alias.aliasAnalysis - where onFun _ fd = do - let vtable = ST.insertFParams (funDefParams fd) mempty - (body, _) = optimiseBody vtable mempty $ funDefBody fd - return fd { funDefBody = body } +sink = + Pass "sink" "move memory loads closer to their uses" $ + fmap removeProgAliases + . intraproceduralTransformationWithConsts onConsts onFun + . Alias.aliasAnalysis + where + onFun _ fd = do + let vtable = ST.insertFParams (funDefParams fd) mempty + (body, _) = optimiseBody vtable mempty $ funDefBody fd + return fd {funDefBody = body} - onConsts consts = - pure $ fst $ optimiseStms mempty mempty consts $ - namesFromList $ M.keys $ scopeOf consts + onConsts consts = + pure $ + fst $ + optimiseStms mempty mempty consts $ + namesFromList $ M.keys $ scopeOf consts 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 @@ -1,930 +1,1271 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} --- | Perform a restricted form of loop tiling within SegMaps. We only --- tile primitive types, to avoid excessive local memory use. -module Futhark.Optimise.TileLoops - ( tileLoops ) - where - -import Control.Monad.State -import Control.Monad.Reader -import qualified Data.Sequence as Seq -import qualified Data.Map.Strict as M -import Data.List (foldl') - -import Prelude hiding (quot) - -import Futhark.MonadFreshNames -import Futhark.IR.Kernels -import Futhark.Transform.Rename -import Futhark.Pass -import Futhark.Tools - --- | The pass definition. -tileLoops :: Pass Kernels Kernels -tileLoops = Pass "tile loops" "Tile stream loops inside kernels" $ - \(Prog consts funs) -> - Prog consts <$> mapM optimiseFunDef funs - -optimiseFunDef :: MonadFreshNames m => FunDef Kernels -> m (FunDef Kernels) -optimiseFunDef fundec = do - body' <- modifyNameSource $ runState $ - runReaderT m (scopeOfFParams (funDefParams fundec)) - return fundec { funDefBody = body' } - where m = optimiseBody $ funDefBody fundec - -type TileM = ReaderT (Scope Kernels) (State VNameSource) - -optimiseBody :: Body Kernels -> TileM (Body Kernels) -optimiseBody (Body () bnds res) = localScope (scopeOf bnds) $ - Body () <$> (mconcat <$> mapM optimiseStm (stmsToList bnds)) <*> pure res - -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') - where initial_variance = M.map mempty $ scopeOfSegSpace space - -optimiseStm (Let pat aux e) = - pure <$> (Let pat aux <$> mapExpM optimise e) - where optimise = identityMapper { mapOnBody = \scope -> localScope scope . optimiseBody } - -tileInKernelBody :: Names -> VarianceTable - -> SegLevel -> SegSpace -> [Type] -> KernelBody Kernels - -> TileM (Stms Kernels, (SegLevel, SegSpace, KernelBody Kernels)) -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 $ - Body () (kernelBodyStms kbody) kbody_res - case maybe_tiled of - Just (host_stms, tiling, tiledBody) -> do - (res', stms') <- - runBinder $ mapM (tilingTileReturns tiling) =<< tiledBody mempty - return (host_stms, (tilingLevel tiling, - tilingSpace tiling, - KernelBody () stms' res')) - Nothing -> - return (mempty, (lvl, initial_kspace, kbody)) - | otherwise = - return (mempty, (lvl, initial_kspace, kbody)) - where isSimpleResult (Returns _ se) = Just se - isSimpleResult _ = Nothing - -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) = - descend mempty $ stmsToList initial_kstms - where - variance = varianceInStms initial_variance initial_kstms - - descend _ [] = - return Nothing - - descend prestms (stm_to_tile:poststms) - - -- 1D tiling of redomap. - | (gtid, kdim) : top_space_rev <- reverse $ unSegSpace initial_space, - Just (w, arrs, form) <- tileable stm_to_tile, - not $ any (nameIn gtid . - flip (M.findWithDefault mempty) variance) arrs, - not $ gtid `nameIn` branch_variant, - (prestms', poststms') <- - preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms), - used <- freeIn stm_to_tile <> freeIn stms_res = - - Just . injectPrelude initial_space private variance prestms' used <$> - tileGeneric (tiling1d $ reverse top_space_rev) - initial_lvl res_ts (stmPattern stm_to_tile) - gtid kdim - w form (zip arrs $ repeat [0]) poststms' stms_res - - -- 2D tiling of redomap. - | (gtids, kdims) <- unzip $ unSegSpace initial_space, - Just (w, arrs, form) <- tileable stm_to_tile, - Just inner_perm <- mapM (invariantToOneOfTwoInnerDims branch_variant variance gtids) arrs, - gtid_y : gtid_x : top_gtids_rev <- reverse gtids, - kdim_y : kdim_x : top_kdims_rev <- reverse kdims, - (prestms', poststms') <- - preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms), - used <- freeIn stm_to_tile <> freeIn stms_res = - - Just . injectPrelude initial_space private variance prestms' used <$> - tileGeneric (tiling2d $ reverse $ zip top_gtids_rev top_kdims_rev) - initial_lvl res_ts (stmPattern stm_to_tile) - (gtid_x, gtid_y) (kdim_x, kdim_y) - w form (zip arrs inner_perm) poststms' stms_res - - -- Tiling inside for-loop. - | DoLoop [] merge (ForLoop i it bound []) loopbody <- stmExp stm_to_tile, - (prestms', poststms') <- - preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms) = do - - let branch_variant' = - branch_variant <> - mconcat (map (flip (M.findWithDefault mempty) variance) - (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 - (map paramType merge_params) $ mkBody (bodyStms loopbody) (bodyResult loopbody) - - case maybe_tiled of - Nothing -> next - Just tiled -> - Just <$> tileDoLoop initial_space variance prestms' - (freeIn loopbody <> freeIn merge) tiled - res_ts (stmPattern stm_to_tile) (stmAux stm_to_tile) - merge i it bound poststms' stms_res - - | otherwise = next - - where next = localScope (scopeOf stm_to_tile) $ - descend (prestms <> oneStm stm_to_tile) poststms - --- | Move statements from prelude to postlude if they are not used in --- the tiled statement anyway. -preludeToPostlude :: VarianceTable - -> Stms Kernels -> Stm Kernels -> Stms Kernels - -> (Stms Kernels, Stms Kernels) -preludeToPostlude variance prelude stm_to_tile postlude = - (prelude_used, prelude_not_used <> postlude) - where used_in_tiled = freeIn stm_to_tile - - used_in_stm_variant = - (used_in_tiled<>) $ mconcat $ - map (flip (M.findWithDefault mempty) variance) $ - namesToList used_in_tiled - - used stm = any (`nameIn` used_in_stm_variant) $ - patternNames $ stmPattern stm - - (prelude_used, prelude_not_used) = - Seq.partition used prelude - --- | Partition prelude statements preceding a tiled loop (or something --- containing a tiled loop) into three categories: --- --- 1) Group-level statements that are invariant to the threads in the group. --- --- 2) Thread-variant statements that should be computed once with a segmap_thread_scalar. --- --- 3) Thread-variant statements that should be recomputed whenever --- they are needed. --- --- The third category duplicates computation, so we only want to do it --- when absolutely necessary. Currently, this is necessary for --- results that are views of an array (slicing, rotate, etc), because --- these cannot be efficiently represented by a scalar segmap (they'll --- be manifested in memory). -partitionPrelude :: VarianceTable -> Stms Kernels -> Names - -> (Stms Kernels, Stms Kernels, Stms Kernels) -partitionPrelude variance prestms private = - (invariant_prestms, precomputed_variant_prestms, recomputed_variant_prestms) - where - invariantTo names stm = - case patternNames (stmPattern stm) of - [] -> True -- Does not matter. - v:_ -> not $ any (`nameIn` names) $ namesToList $ - M.findWithDefault mempty v variance - (invariant_prestms, variant_prestms) = - Seq.partition (invariantTo private) prestms - - mustBeInlinedExp (BasicOp (Index _ slice)) = not $ null $ sliceDims slice - mustBeInlinedExp (BasicOp Rotate{}) = True - mustBeInlinedExp (BasicOp Rearrange{}) = True - mustBeInlinedExp (BasicOp Reshape{}) = True - mustBeInlinedExp _ = False - mustBeInlined = mustBeInlinedExp . stmExp - - must_be_inlined = namesFromList $ concatMap (patternNames . stmPattern) $ - stmsToList $ Seq.filter mustBeInlined variant_prestms - recompute stm = - any (`nameIn` must_be_inlined) (patternNames (stmPattern stm)) || - not (invariantTo must_be_inlined stm) - (recomputed_variant_prestms, precomputed_variant_prestms) = - Seq.partition recompute variant_prestms - --- Anything that is variant to the "private" names should be --- 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) = - (host_stms, tiling, tiledBody') - where private' = private <> namesFromList - (map fst $ - filter (`notElem` unSegSpace (tilingSpace tiling)) $ - unSegSpace initial_space) - - tiledBody' privstms = do - let (invariant_prestms, - precomputed_variant_prestms, - recomputed_variant_prestms) = - partitionPrelude variance prestms private' - - addStms invariant_prestms - - let live_set = namesToList $ liveSet precomputed_variant_prestms $ - used <> freeIn recomputed_variant_prestms - prelude_arrs <- inScopeOf precomputed_variant_prestms $ - doPrelude tiling precomputed_variant_prestms live_set - - let prelude_privstms = - PrivStms recomputed_variant_prestms $ - mkReadPreludeValues prelude_arrs live_set - - tiledBody (prelude_privstms <> privstms) - -tileDoLoop :: SegSpace -> VarianceTable - -> Stms Kernels -> Names - -> (Stms Kernels, Tiling, TiledBody) - -> [Type] -> Pattern Kernels -> StmAux (ExpDec Kernels) - -> [(FParam Kernels, SubExp)] -> VName -> IntType -> SubExp - -> Stms Kernels -> 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 (invariant_prestms, - precomputed_variant_prestms, - recomputed_variant_prestms) = - partitionPrelude variance prestms tiled_kdims - - let (mergeparams, mergeinits) = unzip merge - - -- Expand the loop merge parameters to be arrays. - tileDim t = arrayOf t (tilingTileShape tiling) $ uniqueness t - - tiledBody' privstms = inScopeOf host_stms $ do - addStms invariant_prestms - - let live_set = namesToList $ liveSet precomputed_variant_prestms used_in_body - prelude_arrs <- inScopeOf precomputed_variant_prestms $ - doPrelude tiling precomputed_variant_prestms live_set - - mergeparams' <- forM mergeparams $ \(Param pname pt) -> - Param <$> newVName (baseString pname ++ "_group") <*> pure (tileDim pt) - - let merge_ts = map paramType mergeparams - - let inloop_privstms = - PrivStms recomputed_variant_prestms $ - mkReadPreludeValues prelude_arrs live_set - - mergeinit' <- - fmap (map Var) $ certifying (stmAuxCerts aux) $ - tilingSegMap tiling "tiled_loopinit" (scalarLevel tiling) ResultPrivate $ - \in_bounds slice -> - fmap (map Var) $ protectOutOfBounds "loopinit" in_bounds merge_ts $ do - addPrivStms slice inloop_privstms - addPrivStms slice privstms - return mergeinits - - let merge' = zip mergeparams' mergeinit' - - let indexMergeParams slice = - localScope (scopeOfFParams mergeparams') $ - forM_ (zip mergeparams mergeparams') $ \(to, from) -> - letBindNames [paramName to] $ BasicOp $ Index (paramName from) $ - fullSlice (paramType from) slice - - loopbody' <- runBodyBinder $ resultBody . map Var <$> - tiledBody (PrivStms mempty indexMergeParams <> privstms <> inloop_privstms) - accs' <- letTupExp "tiled_inside_loop" $ - DoLoop [] merge' (ForLoop i it bound []) loopbody' - - postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts - - return (host_stms, tiling, tiledBody') - - where tiled_kdims = namesFromList $ map fst $ - filter (`notElem` unSegSpace (tilingSpace tiling)) $ - unSegSpace initial_space - -doPrelude :: Tiling -> Stms Kernels -> [VName] -> Binder Kernels [VName] -doPrelude tiling prestms prestms_live = - -- Create a SegMap that takes care of the prelude for every thread. - tilingSegMap tiling "prelude" (scalarLevel tiling) ResultPrivate $ - \in_bounds _slice -> do - ts <- mapM lookupType prestms_live - fmap (map Var) $ letTupExp "pre" =<< - eIf (toExp in_bounds) - (do addStms prestms - resultBodyM $ map Var prestms_live) - (eBody $ map eBlank ts) - -liveSet :: FreeIn a => Stms Kernels -> a -> Names -liveSet stms after = - namesFromList (concatMap (patternNames . stmPattern) stms) `namesIntersection` - freeIn after - -tileable :: Stm Kernels - -> Maybe (SubExp, [VName], - (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels)) -tileable stm - | Op (OtherOp (Screma w form arrs)) <- stmExp stm, - Just (reds, map_lam) <- isRedomapSOAC form, - Reduce red_comm red_lam red_nes <- singleReduce reds, - 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 - --- | Statements that we insert directly into every thread-private --- SegMaps. This is for things that cannot efficiently be computed --- once in advance in the prelude SegMap, primarily (exclusively?) --- array slicing operations. -data PrivStms = PrivStms (Stms Kernels) ReadPrelude - -privStms :: Stms Kernels -> PrivStms -privStms stms = PrivStms stms $ const $ return () - -addPrivStms :: Slice SubExp -> PrivStms -> Binder Kernels () -addPrivStms local_slice (PrivStms stms readPrelude) = do - readPrelude local_slice - addStms stms - -instance Semigroup PrivStms where - PrivStms stms_x readPrelude_x <> PrivStms stms_y readPrelude_y = - PrivStms stms_z readPrelude_z - where stms_z = stms_x <> stms_y - readPrelude_z slice = readPrelude_x slice >> readPrelude_y slice - -instance Monoid PrivStms where - mempty = privStms mempty - -type ReadPrelude = Slice SubExp -> Binder Kernels () - --- | Information about a loop that has been tiled inside a kernel, as --- well as the kinds of changes that we would then like to perform on --- the kernel. -data Tiling = - Tiling - { tilingSegMap :: String -> SegLevel -> ResultManifest - -> (PrimExp VName -> Slice SubExp -> Binder Kernels [SubExp]) - -> Binder Kernels [VName] - -- The boolean PrimExp indicates whether they are in-bounds. - - , tilingReadTile :: TileKind -> PrivStms - -> SubExp -> [(VName, [Int])] - -> Binder Kernels [VName] - - , tilingProcessTile :: PrivStms - -> Commutativity -> Lambda Kernels -> Lambda Kernels - -> [(VName, [Int])] -> [VName] - -> Binder Kernels [VName] - - , tilingProcessResidualTile :: PrivStms - -> Commutativity -> Lambda Kernels -> Lambda Kernels - -> SubExp -> [VName] -> SubExp - -> [(VName, [Int])] - -> Binder Kernels [VName] - - , tilingTileReturns :: VName -> Binder Kernels KernelResult - - , tilingSpace :: SegSpace - - , tilingTileShape :: Shape - - , tilingLevel :: SegLevel - - , tilingNumWholeTiles :: Binder Kernels SubExp - } - -type DoTiling gtids kdims = - SegLevel -> gtids -> kdims -> SubExp -> Binder Kernels Tiling - -scalarLevel :: Tiling -> SegLevel -scalarLevel tiling = - SegThread (segNumGroups lvl) (segGroupSize lvl) SegNoVirt - where lvl = tilingLevel tiling - -protectOutOfBounds :: String -> PrimExp VName -> [Type] -> Binder Kernels [SubExp] - -> Binder Kernels [VName] -protectOutOfBounds desc in_bounds ts m = - letTupExp desc =<< eIf (toExp in_bounds) (resultBody <$> m) (eBody $ map eBlank ts) - -postludeGeneric :: Tiling -> PrivStms - -> Pattern Kernels -> [VName] - -> Stms Kernels -> Result -> [Type] - -> Binder Kernels [VName] -postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts = - tilingSegMap tiling "thread_res" (scalarLevel tiling) ResultPrivate $ \in_bounds slice -> do - -- Read our per-thread result from the tiled loop. - forM_ (zip (patternNames pat) accs') $ \(us, everyone) -> do - everyone_t <- lookupType everyone - letBindNames [us] $ BasicOp $ Index everyone $ fullSlice everyone_t slice - - if poststms == mempty - then do -- The privstms may still be necessary for the result. - addPrivStms slice privstms - return poststms_res - - else - fmap (map Var) $ protectOutOfBounds "postlude" in_bounds res_ts $ do - addPrivStms slice privstms - addStms poststms - return poststms_res - -type TiledBody = PrivStms -> Binder Kernels [VName] - -tileGeneric :: DoTiling gtids kdims - -> SegLevel - -> [Type] - -> Pattern Kernels - -> gtids - -> kdims - -> SubExp - -> (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels) - -> [(VName, [Int])] - -> Stms Kernels -> Result - -> TileM (Stms Kernels, Tiling, TiledBody) -tileGeneric doTiling initial_lvl res_ts pat gtids kdims w form arrs_and_perms poststms poststms_res = do - - (tiling, tiling_stms) <- runBinder $ doTiling initial_lvl gtids kdims w - - return (tiling_stms, tiling, tiledBody tiling) - - where - (red_comm, red_lam, red_nes, map_lam) = form - - tiledBody :: Tiling -> PrivStms -> Binder Kernels [VName] - tiledBody tiling privstms = do - let tile_shape = tilingTileShape tiling - - num_whole_tiles <- tilingNumWholeTiles tiling - - -- We don't use a Replicate here, because we want to enforce a - -- scalar memory space. - mergeinits <- tilingSegMap tiling "mergeinit" (scalarLevel tiling) ResultPrivate $ \in_bounds slice -> - -- Constant neutral elements (a common case) do not need protection from OOB. - if freeIn red_nes == mempty - then return red_nes - else fmap (map Var) $ protectOutOfBounds "neutral" in_bounds (lambdaReturnType red_lam) $ do - addPrivStms slice privstms - return red_nes - - merge <- forM (zip (lambdaParams red_lam) mergeinits) $ \(p, mergeinit) -> - (,) <$> - newParam (baseString (paramName p) ++ "_merge") - (paramType p `arrayOfShape` tile_shape `toDecl` Unique) <*> - pure (Var mergeinit) - - tile_id <- newVName "tile_id" - let loopform = ForLoop tile_id Int32 num_whole_tiles [] - loopbody <- renameBody <=< runBodyBinder $ inScopeOf loopform $ - localScope (scopeOfFParams $ map fst merge) $ do - - -- Collectively read a tile. - tile <- tilingReadTile tiling TilePartial privstms (Var tile_id) arrs_and_perms - - -- Now each thread performs a traversal of the tile and - -- updates its accumulator. - resultBody . map Var <$> - tilingProcessTile tiling privstms - red_comm red_lam map_lam - (zip tile (map snd arrs_and_perms)) (map (paramName . fst) merge) - - accs <- letTupExp "accs" $ DoLoop [] merge loopform loopbody - - -- We possibly have to traverse a residual tile. - red_lam' <- renameLambda red_lam - map_lam' <- renameLambda map_lam - accs' <- tilingProcessResidualTile tiling privstms - red_comm red_lam' map_lam' - num_whole_tiles accs w arrs_and_perms - - -- Create a SegMap that takes care of the postlude for every thread. - postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts - -data TileKind = TilePartial | TileFull - -mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude -mkReadPreludeValues prestms_live_arrs prestms_live slice = - fmap mconcat $ forM (zip prestms_live_arrs prestms_live) $ \(arr, v) -> do - arr_t <- lookupType arr - letBindNames [v] $ BasicOp $ Index arr $ fullSlice arr_t slice - -tileReturns :: [(VName, SubExp)] -> [(SubExp, SubExp)] -> VName -> Binder Kernels KernelResult -tileReturns dims_on_top dims arr = do - let unit_dims = replicate (length dims_on_top) (intConst Int32 1) - arr' <- if null dims_on_top then return arr - else do arr_t <- lookupType arr - let new_shape = unit_dims ++ arrayDims arr_t - letExp (baseString arr) $ BasicOp $ Reshape (map DimNew new_shape) arr - let tile_dims = zip (map snd dims_on_top) unit_dims ++ dims - return $ TileReturns tile_dims arr' - -segMap1D :: String - -> SegLevel -> ResultManifest - -> (VName -> Binder Kernels [SubExp]) - -> Binder Kernels [VName] -segMap1D desc lvl manifest f = do - ltid <- newVName "ltid" - ltid_flat <- newVName "ltid_flat" - let space = SegSpace ltid_flat [(ltid, unCount $ segGroupSize lvl)] - - ((ts, res), stms) <- runBinder $ do - res <- f ltid - 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' - -reconstructGtids1D :: Count GroupSize SubExp -> VName -> VName -> VName - -> Binder Kernels () -reconstructGtids1D group_size gtid gid ltid = - letBindNames [gtid] =<< - toExp (LeafExp gid int32 * - primExpFromSubExp int32 (unCount group_size) + - LeafExp ltid int32) - -readTile1D :: SubExp -> VName -> VName - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> TileKind -> PrivStms - -> SubExp - -> [(VName, [Int])] - -> Binder Kernels [VName] -readTile1D - tile_size gid gtid num_groups group_size - kind privstms tile_id arrs_and_perms = - - segMap1D "full_tile" (SegThread num_groups group_size SegNoVirt) ResultNoSimplify $ \ltid -> do - j <- letSubExp "j" =<< - toExp (primExpFromSubExp int32 tile_id * - primExpFromSubExp int32 tile_size + - LeafExp ltid int32) - - reconstructGtids1D group_size gtid gid ltid - addPrivStms [DimFix $ Var ltid] privstms - - let arrs = map fst arrs_and_perms - arr_ts <- mapM lookupType arrs - let tile_ts = map rowType arr_ts - w = arraysSize 0 arr_ts - - let readTileElem arr = - -- No need for fullSlice because we are tiling only prims. - letExp "tile_elem" $ BasicOp $ Index arr [DimFix j] - fmap (map Var) $ - case kind of - TilePartial -> - letTupExp "pre" =<< eIf (toExp $ primExpFromSubExp int32 j .<. - primExpFromSubExp int32 w) - (resultBody <$> mapM (fmap Var . readTileElem) arrs) - (eBody $ map eBlank tile_ts) - TileFull -> - mapM readTileElem arrs - -processTile1D :: VName -> VName -> SubExp -> SubExp - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> PrivStms - -> Commutativity -> Lambda Kernels -> Lambda Kernels - -> [(VName, [Int])] -> [VName] - -> Binder Kernels [VName] -processTile1D - gid gtid kdim tile_size num_groups group_size - privstms - red_comm red_lam map_lam tiles_and_perm accs = do - - let tile = map fst tiles_and_perm - - segMap1D "acc" (SegThread num_groups group_size SegNoVirt) ResultPrivate $ \ltid -> do - - reconstructGtids1D group_size gtid gid ltid - addPrivStms [DimFix $ Var ltid] privstms - - -- We replace the neutral elements with the accumulators (this is - -- OK because the parallel semantics are not used after this - -- point). - thread_accs <- forM accs $ \acc -> - letSubExp "acc" $ BasicOp $ Index acc [DimFix $ Var ltid] - let form' = redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam - - fmap (map Var) $ - letTupExp "acc" =<< eIf (toExp $ LeafExp gtid int32 .<. primExpFromSubExp int32 kdim) - (eBody [pure $ Op $ OtherOp $ Screma tile_size form' tile]) - (resultBodyM thread_accs) - -processResidualTile1D :: VName -> VName -> SubExp -> SubExp - -> Count NumGroups SubExp -> Count GroupSize SubExp -> PrivStms - -> Commutativity -> Lambda Kernels -> Lambda Kernels - -> SubExp -> [VName] -> SubExp -> [(VName, [Int])] - -> Binder Kernels [VName] -processResidualTile1D - gid gtid kdim tile_size num_groups group_size privstms red_comm red_lam map_lam - num_whole_tiles accs w arrs_and_perms = do - -- The number of residual elements that are not covered by - -- the whole tiles. - residual_input <- letSubExp "residual_input" $ - BasicOp $ BinOp (SRem Int32 Unsafe) w tile_size - - letTupExp "acc_after_residual" =<< - eIf (toExp $ primExpFromSubExp int32 residual_input .==. 0) - (resultBodyM $ map Var accs) - (nonemptyTile residual_input) - - where - nonemptyTile residual_input = runBodyBinder $ do - -- Collectively construct a tile. Threads that are out-of-bounds - -- provide a blank dummy value. - full_tile <- readTile1D tile_size gid gtid num_groups group_size - TilePartial privstms num_whole_tiles arrs_and_perms - tile <- forM full_tile $ \tile -> - letExp "partial_tile" $ BasicOp $ Index tile - [DimSlice (intConst Int32 0) residual_input (intConst Int32 1)] - - -- Now each thread performs a traversal of the tile and - -- updates its accumulator. - resultBody . map Var <$> processTile1D - gid gtid kdim residual_input num_groups group_size privstms - red_comm red_lam map_lam (zip tile $ repeat [0]) accs - -tiling1d :: [(VName,SubExp)] -> DoTiling VName SubExp -tiling1d dims_on_top initial_lvl gtid kdim w = do - gid <- newVName "gid" - gid_flat <- newVName "gid_flat" - - (lvl, space) <- - if null dims_on_top - then return (SegGroup (segNumGroups initial_lvl) (segGroupSize initial_lvl) $ segVirt initial_lvl, - SegSpace gid_flat [(gid, unCount $ segNumGroups initial_lvl)]) - else do - group_size <- letSubExp "computed_group_size" $ - BasicOp $ BinOp (SMin Int32) (unCount (segGroupSize initial_lvl)) kdim - - -- How many groups we need to exhaust the innermost dimension. - ldim <- letSubExp "ldim" $ - BasicOp $ BinOp (SDivUp Int32 Unsafe) kdim group_size - - num_groups <- letSubExp "computed_num_groups" =<< - foldBinOp (Mul Int32 OverflowUndef) ldim (map snd dims_on_top) - - return (SegGroup (Count num_groups) (Count group_size) SegNoVirt, - SegSpace gid_flat $ dims_on_top ++ [(gid, ldim)]) - let tile_size = unCount $ segGroupSize lvl - - - return Tiling - { tilingSegMap = \desc lvl' manifest f -> segMap1D desc lvl' manifest $ \ltid -> do - letBindNames [gtid] =<< - toExp (LeafExp gid int32 * primExpFromSubExp int32 tile_size + - LeafExp ltid int32) - f (LeafExp gtid int32 .<. primExpFromSubExp int32 kdim) - [DimFix $ Var ltid] - - , tilingReadTile = - readTile1D tile_size gid gtid (segNumGroups lvl) (segGroupSize lvl) - - , tilingProcessTile = - processTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl) - - , tilingProcessResidualTile = - processResidualTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl) - - , tilingTileReturns = tileReturns dims_on_top [(kdim, tile_size)] - - , tilingTileShape = Shape [tile_size] - , tilingNumWholeTiles = letSubExp "num_whole_tiles" $ - BasicOp $ BinOp (SQuot Int32 Unsafe) w tile_size - , tilingLevel = lvl - , tilingSpace = space - } - -invariantToOneOfTwoInnerDims :: Names -> M.Map VName Names -> [VName] -> VName - -> Maybe [Int] -invariantToOneOfTwoInnerDims branch_variant variance dims arr = do - j : i : _ <- Just $ reverse dims - let variant_to = M.findWithDefault mempty arr variance - branch_invariant = not $ nameIn j branch_variant || nameIn i branch_variant - if branch_invariant && i `nameIn` variant_to && not (j `nameIn` variant_to) then - Just [0,1] - else if branch_invariant && j `nameIn` variant_to && not (i `nameIn` variant_to) then - Just [1,0] - else - Nothing - -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 -> (VName, VName) -> (VName, VName) -> (VName, VName) - -> Binder Kernels () -reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) = do - -- Reconstruct the original gtids from gid_x/gid_y and ltid_x/ltid_y. - letBindNames [gtid_x] =<< - toExp (LeafExp gid_x int32 * primExpFromSubExp int32 tile_size + - LeafExp ltid_x int32) - letBindNames [gtid_y] =<< - toExp (LeafExp gid_y int32 * primExpFromSubExp int32 tile_size + - LeafExp ltid_y int32) - -readTile2D :: (SubExp, SubExp) -> (VName, VName) -> (VName, VName) -> SubExp - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> TileKind -> PrivStms -> SubExp - -> [(VName, [Int])] - -> Binder Kernels [VName] -readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size num_groups group_size kind privstms tile_id arrs_and_perms = - segMap2D "full_tile" (SegThread num_groups group_size SegNoVirtFull) - ResultNoSimplify (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do - i <- letSubExp "i" =<< - toExp (primExpFromSubExp int32 tile_id * - primExpFromSubExp int32 tile_size + - LeafExp ltid_x int32) - j <- letSubExp "j" =<< - toExp (primExpFromSubExp int32 tile_id * - primExpFromSubExp int32 tile_size + - LeafExp ltid_y int32) - - reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) - addPrivStms [DimFix $ Var ltid_x, DimFix $ Var ltid_y] privstms - - let (arrs, perms) = unzip arrs_and_perms - arr_ts <- mapM lookupType arrs - let tile_ts = map rowType arr_ts - w = arraysSize 0 arr_ts - - let readTileElem arr perm = - -- No need for fullSlice because we are tiling only prims. - letExp "tile_elem" $ BasicOp $ Index arr - [DimFix $ last $ rearrangeShape perm [i,j]] - readTileElemIfInBounds (tile_t, arr, perm) = do - let idx = last $ rearrangeShape perm [i,j] - othercheck = last $ rearrangeShape perm - [ LeafExp gtid_y int32 .<. primExpFromSubExp int32 kdim_y - , LeafExp gtid_x int32 .<. primExpFromSubExp int32 kdim_x - ] - eIf (toExp $ - primExpFromSubExp int32 idx .<. primExpFromSubExp int32 w .&&. othercheck) - (eBody [return $ BasicOp $ Index arr [DimFix idx]]) - (eBody [eBlank tile_t]) - - fmap (map Var) $ - case kind of - TilePartial -> - mapM (letExp "pre" <=< readTileElemIfInBounds) (zip3 tile_ts arrs perms) - TileFull -> - zipWithM readTileElem arrs perms - -processTile2D :: (VName, VName) -> (VName, VName) -> (SubExp, SubExp) -> SubExp - -> Count NumGroups SubExp -> Count GroupSize SubExp - -> PrivStms - -> Commutativity -> Lambda Kernels -> Lambda Kernels - -> [(VName,[Int])] -> [VName] - -> Binder Kernels [VName] -processTile2D - (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size num_groups group_size - privstms red_comm red_lam map_lam tiles_and_perms accs = do - - -- Might be truncated in case of a partial tile. - actual_tile_size <- arraysSize 0 <$> mapM (lookupType . fst) tiles_and_perms - - segMap2D "acc" (SegThread num_groups group_size SegNoVirtFull) - ResultPrivate (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do - reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) - - addPrivStms [DimFix $ Var ltid_x, DimFix $ Var ltid_y] privstms - - -- We replace the neutral elements with the accumulators (this is - -- OK because the parallel semantics are not used after this - -- point). - thread_accs <- forM accs $ \acc -> - letSubExp "acc" $ BasicOp $ Index acc [DimFix $ Var ltid_x, DimFix $ Var ltid_y] - let form' = redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam - - tiles' <- forM tiles_and_perms $ \(tile, perm) -> do - tile_t <- lookupType tile - letExp "tile" $ BasicOp $ Index tile $ sliceAt tile_t (head perm) - [DimFix $ Var $ head $ rearrangeShape perm [ltid_x, ltid_y]] - - fmap (map Var) $ - letTupExp "acc" =<< eIf (toExp $ - LeafExp gtid_x int32 .<. primExpFromSubExp int32 kdim_x .&&. - LeafExp gtid_y int32 .<. primExpFromSubExp int32 kdim_y) - (eBody [pure $ Op $ OtherOp $ Screma actual_tile_size form' tiles']) - (resultBodyM thread_accs) - -processResidualTile2D :: (VName, VName) -> (VName, VName) -> (SubExp, SubExp) -> SubExp - -> Count NumGroups SubExp -> Count GroupSize SubExp -> PrivStms - -> Commutativity -> Lambda Kernels -> Lambda Kernels - -> SubExp -> [VName] -> SubExp -> [(VName, [Int])] - -> Binder Kernels [VName] -processResidualTile2D - gids gtids kdims tile_size num_groups group_size privstms red_comm red_lam map_lam - num_whole_tiles accs w arrs_and_perms = do - -- The number of residual elements that are not covered by - -- the whole tiles. - residual_input <- letSubExp "residual_input" $ - BasicOp $ BinOp (SRem Int32 Unsafe) w tile_size - - letTupExp "acc_after_residual" =<< - eIf (toExp $ primExpFromSubExp int32 residual_input .==. 0) - (resultBodyM $ map Var accs) - (nonemptyTile residual_input) - - where - nonemptyTile residual_input = renameBody <=< runBodyBinder $ do - -- Collectively construct a tile. Threads that are out-of-bounds - -- provide a blank dummy value. - full_tile <- readTile2D kdims gtids gids tile_size num_groups group_size - TilePartial privstms num_whole_tiles arrs_and_perms - - tile <- forM full_tile $ \tile -> - letExp "partial_tile" $ BasicOp $ Index tile - [DimSlice (intConst Int32 0) residual_input (intConst Int32 1), - DimSlice (intConst Int32 0) residual_input (intConst Int32 1)] - - -- Now each thread performs a traversal of the tile and - -- updates its accumulator. - resultBody . map Var <$> - processTile2D gids gtids kdims tile_size num_groups group_size - privstms red_comm red_lam map_lam - (zip tile (map snd arrs_and_perms)) accs - -tiling2d :: [(VName,SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp) -tiling2d dims_on_top _initial_lvl (gtid_x, gtid_y) (kdim_x, kdim_y) w = do - gid_x <- newVName "gid_x" - gid_y <- newVName "gid_y" - - tile_size_key <- nameFromString . pretty <$> newVName "tile_size" - tile_size <- letSubExp "tile_size" $ Op $ SizeOp $ GetSize tile_size_key SizeTile - group_size <- letSubExp "group_size" $ BasicOp $ BinOp (Mul Int32 OverflowUndef) tile_size tile_size - - num_groups_x <- letSubExp "num_groups_x" $ - BasicOp $ BinOp (SDivUp Int32 Unsafe) kdim_x tile_size - num_groups_y <- letSubExp "num_groups_y" $ - BasicOp $ BinOp (SDivUp Int32 Unsafe) kdim_y tile_size - - num_groups <- letSubExp "num_groups_top" =<< - foldBinOp (Mul Int32 OverflowUndef) num_groups_x - (num_groups_y : map snd dims_on_top) - - gid_flat <- newVName "gid_flat" - let lvl = SegGroup (Count num_groups) (Count group_size) SegNoVirtFull - space = SegSpace gid_flat $ - dims_on_top ++ [(gid_x, num_groups_x), (gid_y, num_groups_y)] - - return Tiling - { tilingSegMap = \desc lvl' manifest f -> - segMap2D desc lvl' manifest (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do - reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) - f (LeafExp gtid_x int32 .<. primExpFromSubExp int32 kdim_x .&&. - LeafExp gtid_y int32 .<. primExpFromSubExp int32 kdim_y) - [DimFix $ Var ltid_x, DimFix $ Var ltid_y] - - , tilingReadTile = readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size (segNumGroups lvl) (segGroupSize lvl) - , tilingProcessTile = processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl) - , tilingProcessResidualTile = processResidualTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl) - - , tilingTileReturns = tileReturns dims_on_top [(kdim_x, tile_size), (kdim_y, tile_size)] - - , tilingTileShape = Shape [tile_size, tile_size] - , tilingNumWholeTiles = letSubExp "num_whole_tiles" $ - BasicOp $ BinOp (SQuot Int32 Unsafe) w tile_size - , 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) + +-- | Perform a restricted form of loop tiling within SegMaps. We only +-- tile primitive types, to avoid excessive local memory use. +module Futhark.Optimise.TileLoops (tileLoops) where + +import Control.Monad.Reader +import Control.Monad.State +import Data.List (foldl') +import qualified Data.Map.Strict as M +import qualified Data.Sequence as Seq +import Futhark.IR.Kernels +import Futhark.MonadFreshNames +import Futhark.Pass +import Futhark.Tools +import Futhark.Transform.Rename +import Prelude hiding (quot) + +-- | The pass definition. +tileLoops :: Pass Kernels Kernels +tileLoops = + Pass "tile loops" "Tile stream loops inside kernels" $ + intraproceduralTransformation onStms + where + onStms scope stms = + modifyNameSource $ + 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 + +optimiseStms :: Stms Kernels -> TileM (Stms Kernels) +optimiseStms stms = + localScope (scopeOf stms) $ + 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') + where + initial_variance = M.map mempty $ scopeOfSegSpace space +optimiseStm (Let pat aux e) = + pure <$> (Let pat aux <$> mapExpM optimise e) + where + optimise = identityMapper {mapOnBody = \scope -> localScope scope . optimiseBody} + +tileInKernelBody :: + Names -> + VarianceTable -> + SegLevel -> + SegSpace -> + [Type] -> + KernelBody Kernels -> + TileM (Stms Kernels, (SegLevel, SegSpace, KernelBody Kernels)) +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 $ + Body () (kernelBodyStms kbody) kbody_res + case maybe_tiled of + Just (host_stms, tiling, tiledBody) -> do + (res', stms') <- + runBinder $ mapM (tilingTileReturns tiling) =<< tiledBody mempty + return + ( host_stms, + ( tilingLevel tiling, + tilingSpace tiling, + KernelBody () stms' res' + ) + ) + Nothing -> + return (mempty, (lvl, initial_kspace, kbody)) + | otherwise = + return (mempty, (lvl, initial_kspace, kbody)) + where + isSimpleResult (Returns _ se) = Just se + isSimpleResult _ = Nothing + +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) = + descend mempty $ stmsToList initial_kstms + where + variance = varianceInStms initial_variance initial_kstms + + descend _ [] = + return Nothing + descend prestms (stm_to_tile : poststms) + -- 1D tiling of redomap. + | (gtid, kdim) : top_space_rev <- reverse $ unSegSpace initial_space, + Just (w, arrs, form) <- tileable stm_to_tile, + not $ + any + ( nameIn gtid + . flip (M.findWithDefault mempty) variance + ) + arrs, + not $ gtid `nameIn` branch_variant, + (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 + <$> tileGeneric + (tiling1d $ reverse top_space_rev) + initial_lvl + res_ts + (stmPattern stm_to_tile) + gtid + kdim + w + form + (zip arrs $ repeat [0]) + poststms' + stms_res + -- 2D tiling of redomap. + | (gtids, kdims) <- unzip $ unSegSpace initial_space, + Just (w, arrs, form) <- tileable stm_to_tile, + Just inner_perm <- mapM (invariantToOneOfTwoInnerDims branch_variant variance gtids) arrs, + gtid_y : gtid_x : top_gtids_rev <- reverse gtids, + kdim_y : kdim_x : top_kdims_rev <- reverse kdims, + (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 + <$> tileGeneric + (tiling2d $ reverse $ zip top_gtids_rev top_kdims_rev) + initial_lvl + res_ts + (stmPattern stm_to_tile) + (gtid_x, gtid_y) + (kdim_x, kdim_y) + w + form + (zip arrs inner_perm) + poststms' + stms_res + -- Tiling inside for-loop. + | DoLoop [] merge (ForLoop i it bound []) loopbody <- stmExp stm_to_tile, + (prestms', poststms') <- + preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms) = do + let branch_variant' = + branch_variant + <> mconcat + ( map + (flip (M.findWithDefault mempty) variance) + (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 + (map paramType merge_params) + $ mkBody (bodyStms loopbody) (bodyResult loopbody) + + case maybe_tiled of + Nothing -> next + Just tiled -> + Just + <$> tileDoLoop + initial_space + variance + prestms' + (freeIn loopbody <> freeIn merge) + tiled + res_ts + (stmPattern stm_to_tile) + (stmAux stm_to_tile) + merge + i + it + bound + poststms' + stms_res + | otherwise = next + where + next = + localScope (scopeOf stm_to_tile) $ + descend (prestms <> oneStm stm_to_tile) poststms + +-- | Move statements from prelude to postlude if they are not used in +-- the tiled statement anyway. +preludeToPostlude :: + VarianceTable -> + Stms Kernels -> + Stm Kernels -> + Stms Kernels -> + (Stms Kernels, Stms Kernels) +preludeToPostlude variance prelude stm_to_tile postlude = + (prelude_used, prelude_not_used <> postlude) + where + used_in_tiled = freeIn stm_to_tile + + used_in_stm_variant = + (used_in_tiled <>) $ + mconcat $ + map (flip (M.findWithDefault mempty) variance) $ + namesToList used_in_tiled + + used stm = + any (`nameIn` used_in_stm_variant) $ + patternNames $ stmPattern stm + + (prelude_used, prelude_not_used) = + Seq.partition used prelude + +-- | Partition prelude statements preceding a tiled loop (or something +-- containing a tiled loop) into three categories: +-- +-- 1) Group-level statements that are invariant to the threads in the group. +-- +-- 2) Thread-variant statements that should be computed once with a segmap_thread_scalar. +-- +-- 3) Thread-variant statements that should be recomputed whenever +-- they are needed. +-- +-- The third category duplicates computation, so we only want to do it +-- when absolutely necessary. Currently, this is necessary for +-- results that are views of an array (slicing, rotate, etc), because +-- these cannot be efficiently represented by a scalar segmap (they'll +-- be manifested in memory). +partitionPrelude :: + VarianceTable -> + Stms Kernels -> + Names -> + (Stms Kernels, Stms Kernels, Stms Kernels) +partitionPrelude variance prestms private = + (invariant_prestms, precomputed_variant_prestms, recomputed_variant_prestms) + where + invariantTo names stm = + case patternNames (stmPattern stm) of + [] -> True -- Does not matter. + v : _ -> + not $ + any (`nameIn` names) $ + namesToList $ + M.findWithDefault mempty v variance + (invariant_prestms, variant_prestms) = + Seq.partition (invariantTo private) prestms + + mustBeInlinedExp (BasicOp (Index _ slice)) = not $ null $ sliceDims slice + mustBeInlinedExp (BasicOp Rotate {}) = True + mustBeInlinedExp (BasicOp Rearrange {}) = True + mustBeInlinedExp (BasicOp Reshape {}) = True + mustBeInlinedExp _ = False + mustBeInlined = mustBeInlinedExp . stmExp + + must_be_inlined = + namesFromList $ + concatMap (patternNames . stmPattern) $ + stmsToList $ Seq.filter mustBeInlined variant_prestms + recompute stm = + any (`nameIn` must_be_inlined) (patternNames (stmPattern stm)) + || not (invariantTo must_be_inlined stm) + (recomputed_variant_prestms, precomputed_variant_prestms) = + Seq.partition recompute variant_prestms + +-- Anything that is variant to the "private" names should be +-- 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) = + (host_stms, tiling, tiledBody') + where + private' = + private + <> namesFromList + ( map fst $ + filter (`notElem` unSegSpace (tilingSpace tiling)) $ + unSegSpace initial_space + ) + + tiledBody' privstms = do + let ( invariant_prestms, + precomputed_variant_prestms, + recomputed_variant_prestms + ) = + partitionPrelude variance prestms private' + + addStms invariant_prestms + + let live_set = + namesToList $ + liveSet precomputed_variant_prestms $ + used <> freeIn recomputed_variant_prestms + prelude_arrs <- + inScopeOf precomputed_variant_prestms $ + doPrelude tiling precomputed_variant_prestms live_set + + let prelude_privstms = + PrivStms recomputed_variant_prestms $ + mkReadPreludeValues prelude_arrs live_set + + tiledBody (prelude_privstms <> privstms) + +tileDoLoop :: + SegSpace -> + VarianceTable -> + Stms Kernels -> + Names -> + (Stms Kernels, Tiling, TiledBody) -> + [Type] -> + Pattern Kernels -> + StmAux (ExpDec Kernels) -> + [(FParam Kernels, SubExp)] -> + VName -> + IntType -> + SubExp -> + Stms Kernels -> + 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 ( invariant_prestms, + precomputed_variant_prestms, + recomputed_variant_prestms + ) = + partitionPrelude variance prestms tiled_kdims + + let (mergeparams, mergeinits) = unzip merge + + -- Expand the loop merge parameters to be arrays. + tileDim t = arrayOf t (tilingTileShape tiling) $ uniqueness t + + tiledBody' privstms = inScopeOf host_stms $ do + addStms invariant_prestms + + let live_set = + namesToList $ + liveSet precomputed_variant_prestms $ + freeIn recomputed_variant_prestms + <> used_in_body + <> freeIn poststms + + prelude_arrs <- + inScopeOf precomputed_variant_prestms $ + doPrelude tiling precomputed_variant_prestms live_set + + mergeparams' <- forM mergeparams $ \(Param pname pt) -> + Param <$> newVName (baseString pname ++ "_group") <*> pure (tileDim pt) + + let merge_ts = map paramType mergeparams + + let inloop_privstms = + PrivStms recomputed_variant_prestms $ + mkReadPreludeValues prelude_arrs live_set + + mergeinit' <- + fmap (map Var) $ + certifying (stmAuxCerts aux) $ + tilingSegMap tiling "tiled_loopinit" (scalarLevel tiling) ResultPrivate $ + \in_bounds slice -> + fmap (map Var) $ + protectOutOfBounds "loopinit" in_bounds merge_ts $ do + addPrivStms slice inloop_privstms + addPrivStms slice privstms + return mergeinits + + let merge' = zip mergeparams' mergeinit' + + let indexMergeParams slice = + localScope (scopeOfFParams mergeparams') $ + forM_ (zip mergeparams mergeparams') $ \(to, from) -> + letBindNames [paramName to] $ + BasicOp $ + Index (paramName from) $ + fullSlice (paramType from) slice + + loopbody' <- + runBodyBinder $ + resultBody . map Var + <$> tiledBody (PrivStms mempty indexMergeParams <> privstms <> inloop_privstms) + accs' <- + letTupExp "tiled_inside_loop" $ + DoLoop [] merge' (ForLoop i it bound []) loopbody' + + postludeGeneric tiling inloop_privstms pat accs' poststms poststms_res res_ts + + return (host_stms, tiling, tiledBody') + where + tiled_kdims = + namesFromList $ + map fst $ + filter (`notElem` unSegSpace (tilingSpace tiling)) $ + unSegSpace initial_space + +doPrelude :: Tiling -> Stms Kernels -> [VName] -> Binder Kernels [VName] +doPrelude tiling prestms prestms_live = + -- Create a SegMap that takes care of the prelude for every thread. + tilingSegMap tiling "prelude" (scalarLevel tiling) ResultPrivate $ + \in_bounds _slice -> do + ts <- mapM lookupType prestms_live + fmap (map Var) $ + letTupExp "pre" + =<< eIf + (toExp in_bounds) + ( do + addStms prestms + resultBodyM $ map Var prestms_live + ) + (eBody $ map eBlank ts) + +liveSet :: FreeIn a => Stms Kernels -> a -> Names +liveSet stms after = + namesFromList (concatMap (patternNames . stmPattern) stms) + `namesIntersection` freeIn after + +tileable :: + Stm Kernels -> + Maybe + ( SubExp, + [VName], + (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels) + ) +tileable stm + | Op (OtherOp (Screma w form arrs)) <- stmExp stm, + Just (reds, map_lam) <- isRedomapSOAC form, + Reduce red_comm red_lam red_nes <- singleReduce reds, + 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 + +-- | Statements that we insert directly into every thread-private +-- SegMaps. This is for things that cannot efficiently be computed +-- once in advance in the prelude SegMap, primarily (exclusively?) +-- array slicing operations. +data PrivStms = PrivStms (Stms Kernels) ReadPrelude + +privStms :: Stms Kernels -> PrivStms +privStms stms = PrivStms stms $ const $ return () + +addPrivStms :: Slice SubExp -> PrivStms -> Binder Kernels () +addPrivStms local_slice (PrivStms stms readPrelude) = do + readPrelude local_slice + addStms stms + +instance Semigroup PrivStms where + PrivStms stms_x readPrelude_x <> PrivStms stms_y readPrelude_y = + PrivStms stms_z readPrelude_z + where + stms_z = stms_x <> stms_y + readPrelude_z slice = readPrelude_x slice >> readPrelude_y slice + +instance Monoid PrivStms where + mempty = privStms mempty + +type ReadPrelude = Slice SubExp -> Binder Kernels () + +-- | Information about a loop that has been tiled inside a kernel, as +-- well as the kinds of changes that we would then like to perform on +-- the kernel. +data Tiling = Tiling + { tilingSegMap :: + String -> + SegLevel -> + ResultManifest -> + (PrimExp VName -> Slice SubExp -> Binder Kernels [SubExp]) -> + Binder Kernels [VName], + -- The boolean PrimExp indicates whether they are in-bounds. + + tilingReadTile :: + TileKind -> + PrivStms -> + SubExp -> + [(VName, [Int])] -> + Binder Kernels [VName], + tilingProcessTile :: + PrivStms -> + Commutativity -> + Lambda Kernels -> + Lambda Kernels -> + [(VName, [Int])] -> + [VName] -> + Binder Kernels [VName], + tilingProcessResidualTile :: + PrivStms -> + Commutativity -> + Lambda Kernels -> + Lambda Kernels -> + SubExp -> + [VName] -> + SubExp -> + [(VName, [Int])] -> + Binder Kernels [VName], + tilingTileReturns :: VName -> Binder Kernels KernelResult, + tilingSpace :: SegSpace, + tilingTileShape :: Shape, + tilingLevel :: SegLevel, + tilingNumWholeTiles :: Binder Kernels SubExp + } + +type DoTiling gtids kdims = + SegLevel -> gtids -> kdims -> SubExp -> Binder Kernels Tiling + +scalarLevel :: Tiling -> SegLevel +scalarLevel tiling = + SegThread (segNumGroups lvl) (segGroupSize lvl) SegNoVirt + where + lvl = tilingLevel tiling + +protectOutOfBounds :: + String -> + PrimExp VName -> + [Type] -> + Binder Kernels [SubExp] -> + Binder Kernels [VName] +protectOutOfBounds desc in_bounds ts m = + letTupExp desc =<< eIf (toExp in_bounds) (resultBody <$> m) (eBody $ map eBlank ts) + +postludeGeneric :: + Tiling -> + PrivStms -> + Pattern Kernels -> + [VName] -> + Stms Kernels -> + Result -> + [Type] -> + Binder Kernels [VName] +postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts = + tilingSegMap tiling "thread_res" (scalarLevel tiling) ResultPrivate $ \in_bounds slice -> do + -- Read our per-thread result from the tiled loop. + forM_ (zip (patternNames pat) accs') $ \(us, everyone) -> do + everyone_t <- lookupType everyone + letBindNames [us] $ BasicOp $ Index everyone $ fullSlice everyone_t slice + + if poststms == mempty + then do + -- The privstms may still be necessary for the result. + addPrivStms slice privstms + return poststms_res + else fmap (map Var) $ + protectOutOfBounds "postlude" in_bounds res_ts $ do + addPrivStms slice privstms + addStms poststms + return poststms_res + +type TiledBody = PrivStms -> Binder Kernels [VName] + +tileGeneric :: + DoTiling gtids kdims -> + SegLevel -> + [Type] -> + Pattern Kernels -> + gtids -> + kdims -> + SubExp -> + (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels) -> + [(VName, [Int])] -> + Stms Kernels -> + Result -> + TileM (Stms Kernels, Tiling, TiledBody) +tileGeneric doTiling initial_lvl res_ts pat gtids kdims w form arrs_and_perms poststms poststms_res = do + (tiling, tiling_stms) <- runBinder $ doTiling initial_lvl gtids kdims w + + return (tiling_stms, tiling, tiledBody tiling) + where + (red_comm, red_lam, red_nes, map_lam) = form + + tiledBody :: Tiling -> PrivStms -> Binder Kernels [VName] + tiledBody tiling privstms = do + let tile_shape = tilingTileShape tiling + + num_whole_tiles <- tilingNumWholeTiles tiling + + -- We don't use a Replicate here, because we want to enforce a + -- scalar memory space. + mergeinits <- tilingSegMap tiling "mergeinit" (scalarLevel tiling) ResultPrivate $ \in_bounds slice -> + -- Constant neutral elements (a common case) do not need protection from OOB. + if freeIn red_nes == mempty + then return red_nes + else fmap (map Var) $ + protectOutOfBounds "neutral" in_bounds (lambdaReturnType red_lam) $ do + addPrivStms slice privstms + return red_nes + + merge <- forM (zip (lambdaParams red_lam) mergeinits) $ \(p, mergeinit) -> + (,) + <$> newParam + (baseString (paramName p) ++ "_merge") + (paramType p `arrayOfShape` tile_shape `toDecl` Unique) + <*> pure (Var mergeinit) + + tile_id <- newVName "tile_id" + let loopform = ForLoop tile_id Int64 num_whole_tiles [] + loopbody <- renameBody <=< runBodyBinder $ + inScopeOf loopform $ + localScope (scopeOfFParams $ map fst merge) $ do + -- Collectively read a tile. + tile <- tilingReadTile tiling TilePartial privstms (Var tile_id) arrs_and_perms + + -- Now each thread performs a traversal of the tile and + -- updates its accumulator. + resultBody . map Var + <$> tilingProcessTile + tiling + privstms + red_comm + red_lam + map_lam + (zip tile (map snd arrs_and_perms)) + (map (paramName . fst) merge) + + accs <- letTupExp "accs" $ DoLoop [] merge loopform loopbody + + -- We possibly have to traverse a residual tile. + red_lam' <- renameLambda red_lam + map_lam' <- renameLambda map_lam + accs' <- + tilingProcessResidualTile + tiling + privstms + red_comm + red_lam' + map_lam' + num_whole_tiles + accs + w + arrs_and_perms + + -- Create a SegMap that takes care of the postlude for every thread. + postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts + +data TileKind = TilePartial | TileFull + +mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude +mkReadPreludeValues prestms_live_arrs prestms_live slice = + fmap mconcat $ + forM (zip prestms_live_arrs prestms_live) $ \(arr, v) -> do + arr_t <- lookupType arr + letBindNames [v] $ BasicOp $ Index arr $ fullSlice arr_t slice + +tileReturns :: [(VName, SubExp)] -> [(SubExp, SubExp)] -> VName -> Binder Kernels KernelResult +tileReturns dims_on_top dims arr = do + let unit_dims = replicate (length dims_on_top) (intConst Int64 1) + arr' <- + if null dims_on_top + then return arr + else do + arr_t <- lookupType arr + let new_shape = unit_dims ++ arrayDims arr_t + letExp (baseString arr) $ BasicOp $ Reshape (map DimNew new_shape) arr + let tile_dims = zip (map snd dims_on_top) unit_dims ++ dims + return $ TileReturns tile_dims arr' + +segMap1D :: + String -> + SegLevel -> + ResultManifest -> + (VName -> Binder Kernels [SubExp]) -> + Binder Kernels [VName] +segMap1D desc lvl manifest f = do + ltid <- newVName "ltid" + ltid_flat <- newVName "ltid_flat" + let space = SegSpace ltid_flat [(ltid, unCount $ segGroupSize lvl)] + + ((ts, res), stms) <- runBinder $ do + res <- f ltid + 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' + +reconstructGtids1D :: + Count GroupSize SubExp -> + VName -> + VName -> + VName -> + Binder Kernels () +reconstructGtids1D group_size gtid gid ltid = + letBindNames [gtid] + =<< toExp (le64 gid * pe64 (unCount group_size) + le64 ltid) + +readTile1D :: + SubExp -> + VName -> + VName -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + TileKind -> + PrivStms -> + SubExp -> + [(VName, [Int])] -> + Binder Kernels [VName] +readTile1D + tile_size + gid + gtid + num_groups + group_size + kind + privstms + tile_id + arrs_and_perms = + segMap1D "full_tile" (SegThread num_groups group_size SegNoVirt) ResultNoSimplify $ \ltid -> do + j <- + letSubExp "j" + =<< toExp (pe64 tile_id * pe64 tile_size + le64 ltid) + + reconstructGtids1D group_size gtid gid ltid + addPrivStms [DimFix $ Var ltid] privstms + + let arrs = map fst arrs_and_perms + arr_ts <- mapM lookupType arrs + let tile_ts = map rowType arr_ts + w = arraysSize 0 arr_ts + + let readTileElem arr = + -- No need for fullSlice because we are tiling only prims. + letExp "tile_elem" $ BasicOp $ Index arr [DimFix j] + fmap (map Var) $ + case kind of + TilePartial -> + letTupExp "pre" + =<< eIf + (toExp $ pe64 j .<. pe64 w) + (resultBody <$> mapM (fmap Var . readTileElem) arrs) + (eBody $ map eBlank tile_ts) + TileFull -> + mapM readTileElem arrs + +processTile1D :: + VName -> + VName -> + SubExp -> + SubExp -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + PrivStms -> + Commutativity -> + Lambda Kernels -> + Lambda Kernels -> + [(VName, [Int])] -> + [VName] -> + Binder Kernels [VName] +processTile1D + gid + gtid + kdim + tile_size + num_groups + group_size + privstms + red_comm + red_lam + map_lam + tiles_and_perm + accs = do + let tile = map fst tiles_and_perm + + segMap1D "acc" (SegThread num_groups group_size SegNoVirt) ResultPrivate $ \ltid -> do + reconstructGtids1D group_size gtid gid ltid + addPrivStms [DimFix $ Var ltid] privstms + + -- We replace the neutral elements with the accumulators (this is + -- OK because the parallel semantics are not used after this + -- point). + thread_accs <- forM accs $ \acc -> + letSubExp "acc" $ BasicOp $ Index acc [DimFix $ Var ltid] + let form' = redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam + + fmap (map Var) $ + letTupExp "acc" + =<< eIf + (toExp $ le64 gtid .<. pe64 kdim) + (eBody [pure $ Op $ OtherOp $ Screma tile_size form' tile]) + (resultBodyM thread_accs) + +processResidualTile1D :: + VName -> + VName -> + SubExp -> + SubExp -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + PrivStms -> + Commutativity -> + Lambda Kernels -> + Lambda Kernels -> + SubExp -> + [VName] -> + SubExp -> + [(VName, [Int])] -> + Binder Kernels [VName] +processResidualTile1D + gid + gtid + kdim + tile_size + num_groups + group_size + privstms + red_comm + red_lam + map_lam + num_whole_tiles + accs + w + arrs_and_perms = do + -- The number of residual elements that are not covered by + -- the whole tiles. + residual_input <- + letSubExp "residual_input" $ + BasicOp $ BinOp (SRem Int64 Unsafe) w tile_size + + letTupExp "acc_after_residual" + =<< eIf + (toExp $ pe64 residual_input .==. 0) + (resultBodyM $ map Var accs) + (nonemptyTile residual_input) + where + nonemptyTile residual_input = runBodyBinder $ do + -- Collectively construct a tile. Threads that are out-of-bounds + -- provide a blank dummy value. + full_tile <- + readTile1D + tile_size + gid + gtid + num_groups + group_size + TilePartial + privstms + num_whole_tiles + arrs_and_perms + tile <- forM full_tile $ \tile -> + letExp "partial_tile" $ + BasicOp $ + Index + tile + [DimSlice (intConst Int64 0) residual_input (intConst Int64 1)] + + -- Now each thread performs a traversal of the tile and + -- updates its accumulator. + resultBody . map Var + <$> processTile1D + gid + gtid + kdim + residual_input + num_groups + group_size + privstms + red_comm + red_lam + map_lam + (zip tile $ repeat [0]) + accs + +tiling1d :: [(VName, SubExp)] -> DoTiling VName SubExp +tiling1d dims_on_top initial_lvl gtid kdim w = do + gid <- newVName "gid" + gid_flat <- newVName "gid_flat" + + (lvl, space) <- + if null dims_on_top + then + return + ( SegGroup (segNumGroups initial_lvl) (segGroupSize initial_lvl) $ segVirt initial_lvl, + SegSpace gid_flat [(gid, unCount $ segNumGroups initial_lvl)] + ) + else do + group_size <- + letSubExp "computed_group_size" $ + BasicOp $ BinOp (SMin Int64) (unCount (segGroupSize initial_lvl)) kdim + + -- How many groups we need to exhaust the innermost dimension. + ldim <- + letSubExp "ldim" $ + BasicOp $ BinOp (SDivUp Int64 Unsafe) kdim group_size + + num_groups <- + letSubExp "computed_num_groups" + =<< foldBinOp (Mul Int64 OverflowUndef) ldim (map snd dims_on_top) + + return + ( SegGroup (Count num_groups) (Count group_size) SegNoVirt, + SegSpace gid_flat $ dims_on_top ++ [(gid, ldim)] + ) + let tile_size = unCount $ segGroupSize lvl + + return + Tiling + { tilingSegMap = \desc lvl' manifest f -> segMap1D desc lvl' manifest $ \ltid -> do + letBindNames [gtid] + =<< toExp (le64 gid * pe64 tile_size + le64 ltid) + f (untyped $ le64 gtid .<. pe64 kdim) [DimFix $ Var ltid], + tilingReadTile = + readTile1D tile_size gid gtid (segNumGroups lvl) (segGroupSize lvl), + tilingProcessTile = + processTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl), + tilingProcessResidualTile = + processResidualTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl), + tilingTileReturns = tileReturns dims_on_top [(kdim, tile_size)], + tilingTileShape = Shape [tile_size], + tilingNumWholeTiles = + letSubExp "num_whole_tiles" $ + BasicOp $ BinOp (SQuot Int64 Unsafe) w tile_size, + tilingLevel = lvl, + tilingSpace = space + } + +invariantToOneOfTwoInnerDims :: + Names -> + M.Map VName Names -> + [VName] -> + VName -> + Maybe [Int] +invariantToOneOfTwoInnerDims branch_variant variance dims arr = do + j : i : _ <- Just $ reverse dims + let variant_to = M.findWithDefault mempty arr variance + branch_invariant = not $ nameIn j branch_variant || nameIn i branch_variant + if branch_invariant && i `nameIn` variant_to && not (j `nameIn` variant_to) + then Just [0, 1] + else + if branch_invariant && j `nameIn` variant_to && not (i `nameIn` variant_to) + then Just [1, 0] + else Nothing + +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 -> + (VName, VName) -> + (VName, VName) -> + (VName, VName) -> + Binder Kernels () +reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) = do + -- Reconstruct the original gtids from gid_x/gid_y and ltid_x/ltid_y. + letBindNames [gtid_x] + =<< toExp (le64 gid_x * pe64 tile_size + le64 ltid_x) + letBindNames [gtid_y] + =<< toExp (le64 gid_y * pe64 tile_size + le64 ltid_y) + +readTile2D :: + (SubExp, SubExp) -> + (VName, VName) -> + (VName, VName) -> + SubExp -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + TileKind -> + PrivStms -> + SubExp -> + [(VName, [Int])] -> + Binder Kernels [VName] +readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size num_groups group_size kind privstms tile_id arrs_and_perms = + segMap2D + "full_tile" + (SegThread num_groups group_size SegNoVirtFull) + ResultNoSimplify + (tile_size, tile_size) + $ \(ltid_x, ltid_y) -> do + i <- + letSubExp "i" + =<< toExp (pe64 tile_id * pe64 tile_size + le64 ltid_x) + j <- + letSubExp "j" + =<< toExp (pe64 tile_id * pe64 tile_size + le64 ltid_y) + + reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) + addPrivStms [DimFix $ Var ltid_x, DimFix $ Var ltid_y] privstms + + let (arrs, perms) = unzip arrs_and_perms + arr_ts <- mapM lookupType arrs + let tile_ts = map rowType arr_ts + w = arraysSize 0 arr_ts + + let readTileElem arr perm = + -- No need for fullSlice because we are tiling only prims. + letExp "tile_elem" $ + BasicOp $ + Index + arr + [DimFix $ last $ rearrangeShape perm [i, j]] + readTileElemIfInBounds (tile_t, arr, perm) = do + let idx = last $ rearrangeShape perm [i, j] + othercheck = + last $ + rearrangeShape + perm + [ le64 gtid_y .<. pe64 kdim_y, + le64 gtid_x .<. pe64 kdim_x + ] + eIf + (toExp $ pe64 idx .<. pe64 w .&&. othercheck) + (eBody [return $ BasicOp $ Index arr [DimFix idx]]) + (eBody [eBlank tile_t]) + + fmap (map Var) $ + case kind of + TilePartial -> + mapM (letExp "pre" <=< readTileElemIfInBounds) (zip3 tile_ts arrs perms) + TileFull -> + zipWithM readTileElem arrs perms + +processTile2D :: + (VName, VName) -> + (VName, VName) -> + (SubExp, SubExp) -> + SubExp -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + PrivStms -> + Commutativity -> + Lambda Kernels -> + Lambda Kernels -> + [(VName, [Int])] -> + [VName] -> + Binder Kernels [VName] +processTile2D + (gid_x, gid_y) + (gtid_x, gtid_y) + (kdim_x, kdim_y) + tile_size + num_groups + group_size + privstms + red_comm + red_lam + map_lam + tiles_and_perms + accs = do + -- Might be truncated in case of a partial tile. + actual_tile_size <- arraysSize 0 <$> mapM (lookupType . fst) tiles_and_perms + + segMap2D + "acc" + (SegThread num_groups group_size SegNoVirtFull) + ResultPrivate + (tile_size, tile_size) + $ \(ltid_x, ltid_y) -> do + reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) + + addPrivStms [DimFix $ Var ltid_x, DimFix $ Var ltid_y] privstms + + -- We replace the neutral elements with the accumulators (this is + -- OK because the parallel semantics are not used after this + -- point). + thread_accs <- forM accs $ \acc -> + letSubExp "acc" $ BasicOp $ Index acc [DimFix $ Var ltid_x, DimFix $ Var ltid_y] + let form' = redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam + + tiles' <- forM tiles_and_perms $ \(tile, perm) -> do + tile_t <- lookupType tile + letExp "tile" $ + BasicOp $ + Index tile $ + sliceAt + tile_t + (head perm) + [DimFix $ Var $ head $ rearrangeShape perm [ltid_x, ltid_y]] + + fmap (map Var) $ + letTupExp "acc" + =<< eIf + ( toExp $ le64 gtid_x .<. pe64 kdim_x .&&. le64 gtid_y .<. pe64 kdim_y + ) + (eBody [pure $ Op $ OtherOp $ Screma actual_tile_size form' tiles']) + (resultBodyM thread_accs) + +processResidualTile2D :: + (VName, VName) -> + (VName, VName) -> + (SubExp, SubExp) -> + SubExp -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + PrivStms -> + Commutativity -> + Lambda Kernels -> + Lambda Kernels -> + SubExp -> + [VName] -> + SubExp -> + [(VName, [Int])] -> + Binder Kernels [VName] +processResidualTile2D + gids + gtids + kdims + tile_size + num_groups + group_size + privstms + red_comm + red_lam + map_lam + num_whole_tiles + accs + w + arrs_and_perms = do + -- The number of residual elements that are not covered by + -- the whole tiles. + residual_input <- + letSubExp "residual_input" $ + BasicOp $ BinOp (SRem Int64 Unsafe) w tile_size + + letTupExp "acc_after_residual" + =<< eIf + (toExp $ pe64 residual_input .==. 0) + (resultBodyM $ map Var accs) + (nonemptyTile residual_input) + where + nonemptyTile residual_input = renameBody <=< runBodyBinder $ do + -- Collectively construct a tile. Threads that are out-of-bounds + -- provide a blank dummy value. + full_tile <- + readTile2D + kdims + gtids + gids + tile_size + num_groups + group_size + TilePartial + privstms + num_whole_tiles + arrs_and_perms + + tile <- forM full_tile $ \tile -> + letExp "partial_tile" $ + BasicOp $ + Index + tile + [ DimSlice (intConst Int64 0) residual_input (intConst Int64 1), + DimSlice (intConst Int64 0) residual_input (intConst Int64 1) + ] + + -- Now each thread performs a traversal of the tile and + -- updates its accumulator. + resultBody . map Var + <$> processTile2D + gids + gtids + kdims + tile_size + num_groups + group_size + privstms + red_comm + red_lam + map_lam + (zip tile (map snd arrs_and_perms)) + accs + +tiling2d :: [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp) +tiling2d dims_on_top _initial_lvl (gtid_x, gtid_y) (kdim_x, kdim_y) w = do + gid_x <- newVName "gid_x" + gid_y <- newVName "gid_y" + + tile_size_key <- nameFromString . pretty <$> newVName "tile_size" + tile_size <- letSubExp "tile_size" $ Op $ SizeOp $ GetSize tile_size_key SizeTile + group_size <- letSubExp "group_size" $ BasicOp $ BinOp (Mul Int64 OverflowUndef) tile_size tile_size + + num_groups_x <- + letSubExp "num_groups_x" $ + BasicOp $ BinOp (SDivUp Int64 Unsafe) kdim_x tile_size + num_groups_y <- + letSubExp "num_groups_y" $ + BasicOp $ BinOp (SDivUp Int64 Unsafe) kdim_y tile_size + + num_groups <- + letSubExp "num_groups_top" + =<< foldBinOp + (Mul Int64 OverflowUndef) + num_groups_x + (num_groups_y : map snd dims_on_top) + + gid_flat <- newVName "gid_flat" + let lvl = SegGroup (Count num_groups) (Count group_size) SegNoVirtFull + space = + SegSpace gid_flat $ + dims_on_top ++ [(gid_x, num_groups_x), (gid_y, num_groups_y)] + + return + Tiling + { tilingSegMap = \desc lvl' manifest f -> + segMap2D desc lvl' manifest (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do + reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) + f + ( untyped $ + le64 gtid_x .<. pe64 kdim_x + .&&. le64 gtid_y .<. pe64 kdim_y + ) + [DimFix $ Var ltid_x, DimFix $ Var ltid_y], + tilingReadTile = readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size (segNumGroups lvl) (segGroupSize lvl), + tilingProcessTile = processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl), + tilingProcessResidualTile = processResidualTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl), + tilingTileReturns = tileReturns dims_on_top [(kdim_x, tile_size), (kdim_y, tile_size)], + tilingTileShape = Shape [tile_size, tile_size], + tilingNumWholeTiles = + letSubExp "num_whole_tiles" $ + BasicOp $ BinOp (SQuot Int64 Unsafe) w tile_size, + 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/Unstream.hs b/src/Futhark/Optimise/Unstream.hs --- a/src/Futhark/Optimise/Unstream.hs +++ b/src/Futhark/Optimise/Unstream.hs @@ -1,5 +1,6 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + -- | Sequentialise any remaining SOACs. It is very important that -- this is run *after* any access-pattern-related optimisation, -- because this pass will destroy information. @@ -19,12 +20,11 @@ -- kept together. module Futhark.Optimise.Unstream (unstream) where -import Control.Monad.State import Control.Monad.Reader - -import Futhark.MonadFreshNames +import Control.Monad.State import Futhark.IR.Kernels import Futhark.IR.Kernels.Simplify (simplifyKernels) +import Futhark.MonadFreshNames import Futhark.Pass import Futhark.Tools import qualified Futhark.Transform.FirstOrderTransform as FOT @@ -33,19 +33,21 @@ -- | The pass definition. unstream :: Pass Kernels Kernels -unstream = Pass "unstream" "sequentialise remaining SOACs" $ - intraproceduralTransformation (optimise SeqStreams) - >=> simplifyKernels - >=> intraproceduralTransformation (optimise SeqAll) - where optimise stage scope stms = - modifyNameSource $ runState $ runReaderT (optimiseStms stage stms) scope +unstream = + Pass "unstream" "sequentialise remaining SOACs" $ + intraproceduralTransformation (optimise SeqStreams) + >=> simplifyKernels + >=> intraproceduralTransformation (optimise SeqAll) + where + optimise stage scope stms = + modifyNameSource $ runState $ runReaderT (optimiseStms stage stms) scope type UnstreamM = ReaderT (Scope Kernels) (State VNameSource) optimiseStms :: Stage -> Stms Kernels -> UnstreamM (Stms Kernels) optimiseStms stage stms = localScope (scopeOf stms) $ - stmsFromList . concat <$> mapM (optimiseStm stage) (stmsToList stms) + stmsFromList . concat <$> mapM (optimiseStm stage) (stmsToList stms) optimiseBody :: Stage -> Body Kernels -> UnstreamM (Body Kernels) optimiseBody stage (Body () stms res) = @@ -54,38 +56,40 @@ optimiseKernelBody :: Stage -> KernelBody Kernels -> UnstreamM (KernelBody Kernels) optimiseKernelBody stage (KernelBody () stms res) = localScope (scopeOf stms) $ - KernelBody () - <$> (stmsFromList . concat <$> mapM (optimiseStm stage) (stmsToList stms)) - <*> pure res + KernelBody () + <$> (stmsFromList . concat <$> mapM (optimiseStm stage) (stmsToList stms)) + <*> pure res optimiseLambda :: Stage -> Lambda Kernels -> UnstreamM (Lambda Kernels) optimiseLambda stage lam = localScope (scopeOfLParams $ lambdaParams lam) $ do body <- optimiseBody stage $ lambdaBody lam - return lam { lambdaBody = body } + return lam {lambdaBody = body} sequentialise :: Stage -> SOAC Kernels -> Bool -sequentialise SeqStreams Stream{} = True +sequentialise SeqStreams Stream {} = True sequentialise SeqStreams _ = False sequentialise SeqAll _ = True optimiseStm :: Stage -> Stm Kernels -> UnstreamM [Stm Kernels] - optimiseStm stage (Let pat aux (Op (OtherOp soac))) | sequentialise stage soac = do - stms <- runBinder_ $ FOT.transformSOAC pat soac - fmap concat $ localScope (scopeOf stms) $ mapM (optimiseStm stage) $ stmsToList stms - | otherwise = do - -- Still sequentialise whatever's inside. - pure <$> (Let pat aux . Op . OtherOp <$> mapSOACM optimise soac) - where optimise = identitySOACMapper { mapOnSOACLambda = optimiseLambda stage } - + stms <- runBinder_ $ FOT.transformSOAC pat soac + fmap concat $ localScope (scopeOf stms) $ mapM (optimiseStm stage) $ stmsToList stms + | otherwise = + -- Still sequentialise whatever's inside. + pure <$> (Let pat aux . Op . OtherOp <$> mapSOACM optimise soac) + where + optimise = identitySOACMapper {mapOnSOACLambda = optimiseLambda stage} optimiseStm stage (Let pat aux (Op (SegOp op))) = localScope (scopeOfSegSpace $ segSpace op) $ - pure <$> (Let pat aux . Op . SegOp <$> mapSegOpM optimise op) - where optimise = identitySegOpMapper { mapOnSegOpBody = optimiseKernelBody stage - , mapOnSegOpLambda = optimiseLambda stage - } - + pure <$> (Let pat aux . Op . SegOp <$> mapSegOpM optimise op) + where + optimise = + identitySegOpMapper + { mapOnSegOpBody = optimiseKernelBody stage, + mapOnSegOpLambda = optimiseLambda stage + } optimiseStm stage (Let pat aux e) = pure <$> (Let pat aux <$> mapExpM optimise e) - where optimise = identityMapper { mapOnBody = \scope -> localScope scope . optimiseBody stage } + where + optimise = identityMapper {mapOnBody = \scope -> localScope scope . optimiseBody stage} diff --git a/src/Futhark/Pass.hs b/src/Futhark/Pass.hs --- a/src/Futhark/Pass.hs +++ b/src/Futhark/Pass.hs @@ -1,36 +1,36 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE Strict #-} + -- | Definition of a polymorphic (generic) pass that can work with -- programs of any lore. module Futhark.Pass - ( PassM - , runPassM - , liftEither - , liftEitherM - , Pass (..) - , passLongOption - , parPass - , intraproceduralTransformation - , intraproceduralTransformationWithConsts - ) where + ( PassM, + runPassM, + liftEither, + liftEitherM, + Pass (..), + passLongOption, + parPass, + intraproceduralTransformation, + intraproceduralTransformationWithConsts, + ) +where -import Control.Monad.Writer.Strict import Control.Monad.State.Strict +import Control.Monad.Writer.Strict import Control.Parallel.Strategies import Data.Char import Data.Either - -import Prelude hiding (log) - import Futhark.Error import Futhark.IR -import Futhark.Util.Log import Futhark.MonadFreshNames +import Futhark.Util.Log +import Prelude hiding (log) -- | The monad in which passes execute. newtype PassM a = PassM (WriterT Log (State VNameSource) a) - deriving (Functor, Applicative, Monad) + deriving (Functor, Applicative, Monad) instance MonadLogger PassM where addLog = PassM . tell @@ -41,14 +41,16 @@ -- | Execute a 'PassM' action, yielding logging information and either -- an error text or a result. -runPassM :: MonadFreshNames m => - PassM a -> m (a, Log) +runPassM :: + MonadFreshNames m => + PassM a -> + m (a, Log) runPassM (PassM m) = modifyNameSource $ runState (runWriterT m) -- | Turn an 'Either' computation into a 'PassM'. If the 'Either' is -- 'Left', the result is a 'CompilerBug'. liftEither :: Show err => Either err a -> PassM a -liftEither (Left e) = compilerBugS $ show e +liftEither (Left e) = compilerBugS $ show e liftEither (Right v) = return v -- | Turn an 'Either' monadic computation into a 'PassM'. If the 'Either' is @@ -58,23 +60,24 @@ -- | A compiler pass transforming a 'Prog' of a given lore to a 'Prog' -- of another lore. -data Pass fromlore tolore = - Pass { passName :: String - -- ^ Name of the pass. Keep this short and simple. It will - -- be used to automatically generate a command-line option - -- name via 'passLongOption'. - , passDescription :: String - -- ^ A slightly longer description, which will show up in the - -- command-line help text. - , passFunction :: Prog fromlore -> PassM (Prog tolore) - } +data Pass fromlore tolore = Pass + { -- | Name of the pass. Keep this short and simple. It will + -- be used to automatically generate a command-line option + -- name via 'passLongOption'. + passName :: String, + -- | A slightly longer description, which will show up in the + -- command-line help text. + passDescription :: String, + passFunction :: Prog fromlore -> PassM (Prog tolore) + } -- | Take the name of the pass, turn spaces into dashes, and make all -- characters lowercase. passLongOption :: Pass fromlore tolore -> String passLongOption = map (spaceToDash . toLower) . passName - where spaceToDash ' ' = '-' - spaceToDash c = c + where + spaceToDash ' ' = '-' + spaceToDash c = c -- | Apply a 'PassM' operation in parallel to multiple elements, -- joining together the name sources and logs, and propagating any @@ -83,23 +86,25 @@ parPass f as = do (x, log) <- modifyNameSource $ \src -> let (bs, logs, srcs) = unzip3 $ parMap rpar (f' src) as - in ((bs, mconcat logs), mconcat srcs) + in ((bs, mconcat logs), mconcat srcs) addLog log return x - - where f' src a = - let ((x', log), src') = runState (runPassM (f a)) src - in (x', log, src') + where + f' src a = + let ((x', log), src') = runState (runPassM (f a)) src + in (x', log, src') -- | Apply some operation to the top-level constants. Then applies an -- operation to all the function function definitions, which are also -- given the transformed constants so they can be brought into scope. -- The function definition transformations are run in parallel (with -- 'parPass'), since they cannot affect each other. -intraproceduralTransformationWithConsts :: (Stms fromlore -> PassM (Stms tolore)) - -> (Stms tolore -> FunDef fromlore -> PassM (FunDef tolore)) - -> Prog fromlore -> PassM (Prog tolore) +intraproceduralTransformationWithConsts :: + (Stms fromlore -> PassM (Stms tolore)) -> + (Stms tolore -> FunDef fromlore -> PassM (FunDef tolore)) -> + Prog fromlore -> + PassM (Prog tolore) intraproceduralTransformationWithConsts ct ft (Prog consts funs) = do consts' <- ct consts funs' <- parPass (ft consts') funs @@ -107,13 +112,16 @@ -- | Like 'intraproceduralTransformationWithConsts', but do not change -- the top-level constants, and simply pass along their 'Scope'. -intraproceduralTransformation :: (Scope lore -> Stms lore -> PassM (Stms lore)) - -> Prog lore - -> PassM (Prog lore) +intraproceduralTransformation :: + (Scope lore -> Stms lore -> PassM (Stms lore)) -> + Prog lore -> + PassM (Prog lore) intraproceduralTransformation f = intraproceduralTransformationWithConsts (f mempty) f' - where f' consts fd = do - stms <- f - (scopeOf consts<>scopeOfFParams (funDefParams fd)) - (bodyStms $ funDefBody fd) - return fd { funDefBody = (funDefBody fd) { bodyStms = stms } } + where + f' consts fd = do + stms <- + f + (scopeOf consts <> scopeOfFParams (funDefParams fd)) + (bodyStms $ funDefBody fd) + return fd {funDefBody = (funDefBody fd) {bodyStms = stms}} 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 @@ -1,64 +1,70 @@ -{-# LANGUAGE TypeFamilies, FlexibleContexts, GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} + -- | Expand allocations inside of maps when possible. -module Futhark.Pass.ExpandAllocations - ( expandAllocations ) -where +module Futhark.Pass.ExpandAllocations (expandAllocations) where -import Control.Monad.Identity import Control.Monad.Except -import Control.Monad.State import Control.Monad.Reader +import Control.Monad.State import Control.Monad.Writer +import Data.List (foldl') import qualified Data.Map.Strict as M import Data.Maybe -import Data.List (foldl') - -import Prelude hiding (quot) - import Futhark.Analysis.Rephrase +import qualified Futhark.Analysis.SymbolTable as ST import Futhark.Error -import Futhark.MonadFreshNames -import Futhark.Tools -import Futhark.Pass import Futhark.IR -import Futhark.IR.KernelsMem import qualified Futhark.IR.Kernels.Simplify as Kernels +import Futhark.IR.KernelsMem import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Optimise.Simplify.Lore (addScopeWisdom) +import Futhark.Pass +import Futhark.Pass.ExplicitAllocations.Kernels (explicitAllocationsInStms) import Futhark.Pass.ExtractKernels.BlockedKernel (nonSegRed) import Futhark.Pass.ExtractKernels.ToKernels (segThread) -import Futhark.Pass.ExplicitAllocations.Kernels (explicitAllocationsInStms) -import Futhark.Transform.Rename (renameStm) +import Futhark.Tools import Futhark.Transform.CopyPropagate (copyPropagateInFun) -import Futhark.Optimise.Simplify.Lore (addScopeWisdom) -import qualified Futhark.Analysis.SymbolTable as ST -import Futhark.Util.IntegralExp +import Futhark.Transform.Rename (renameStm) import Futhark.Util (mapAccumLM) +import Futhark.Util.IntegralExp +import Prelude hiding (quot) -- | The memory expansion pass definition. expandAllocations :: Pass KernelsMem KernelsMem expandAllocations = Pass "expand allocations" "Expand allocations" $ - \(Prog consts funs) -> do - consts' <- - modifyNameSource $ limitationOnLeft . runStateT (runReaderT (transformStms consts) mempty) - Prog consts' <$> mapM (transformFunDef $ scopeOf consts') funs - -- Cannot use intraproceduralTransformation because it might create - -- duplicate size keys (which are not fixed by renamer, and size - -- keys must currently be globally unique). + \(Prog consts funs) -> do + consts' <- + modifyNameSource $ limitationOnLeft . runStateT (runReaderT (transformStms consts) mempty) + Prog consts' <$> mapM (transformFunDef $ scopeOf consts') funs +-- Cannot use intraproceduralTransformation because it might create +-- duplicate size keys (which are not fixed by renamer, and size +-- keys must currently be globally unique). + type ExpandM = ReaderT (Scope KernelsMem) (StateT VNameSource (Either String)) limitationOnLeft :: Either String a -> a limitationOnLeft = either compilerLimitationS id -transformFunDef :: Scope KernelsMem -> FunDef KernelsMem - -> PassM (FunDef KernelsMem) +transformFunDef :: + Scope KernelsMem -> + FunDef KernelsMem -> + PassM (FunDef KernelsMem) transformFunDef scope fundec = do body' <- modifyNameSource $ limitationOnLeft . runStateT (runReaderT m mempty) - copyPropagateInFun simpleKernelsMem - (ST.fromScope (addScopeWisdom scope)) fundec { funDefBody = body' } - where m = localScope scope $ inScopeOf fundec $ - transformBody $ funDefBody fundec + copyPropagateInFun + simpleKernelsMem + (ST.fromScope (addScopeWisdom scope)) + fundec {funDefBody = body'} + where + m = + localScope scope $ + inScopeOf fundec $ + transformBody $ funDefBody fundec transformBody :: Body KernelsMem -> ExpandM (Body KernelsMem) transformBody (Body () stms res) = Body () <$> transformStms stms <*> pure res @@ -68,7 +74,6 @@ inScopeOf stms $ mconcat <$> mapM transformStm (stmsToList stms) transformStm :: Stm KernelsMem -> ExpandM (Stms KernelsMem) - -- It is possible that we are unable to expand allocations in some -- code versions. If so, we can remove the offending branch. Only if -- both versions fail do we propagate the error. @@ -84,18 +89,20 @@ return $ oneStm $ Let pat aux $ If cond tbranch'' fbranch'' (IfDec ts IfEquiv) (Left e, _) -> throwError e - - where bindRes pe se = Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se - - useBranch b = - bodyStms b <> - stmsFromList (zipWith bindRes (patternElements pat) (bodyResult b)) + where + bindRes pe se = Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se + useBranch b = + bodyStms b + <> stmsFromList (zipWith bindRes (patternElements pat) (bodyResult b)) transformStm (Let pat aux e) = do (bnds, e') <- transformExp =<< mapExpM transform e return $ bnds <> oneStm (Let pat aux e') - where transform = identityMapper { mapOnBody = \scope -> localScope scope . transformBody - } + where + transform = + identityMapper + { mapOnBody = \scope -> localScope scope . transformBody + } nameInfoConv :: NameInfo KernelsMem -> NameInfo KernelsMem nameInfoConv (LetName mem_info) = LetName mem_info @@ -104,41 +111,47 @@ nameInfoConv (IndexName it) = IndexName it transformExp :: Exp KernelsMem -> ExpandM (Stms KernelsMem, Exp KernelsMem) - transformExp (Op (Inner (SegOp (SegMap lvl space ts kbody)))) = do (alloc_stms, (_, kbody')) <- transformScanRed lvl space [] kbody - return (alloc_stms, - Op $ Inner $ SegOp $ SegMap lvl space ts kbody') - + return + ( alloc_stms, + Op $ Inner $ SegOp $ SegMap lvl space ts kbody' + ) transformExp (Op (Inner (SegOp (SegRed lvl space reds ts kbody)))) = do (alloc_stms, (lams, kbody')) <- transformScanRed lvl space (map segBinOpLambda reds) kbody - let reds' = zipWith (\red lam -> red { segBinOpLambda = lam }) reds lams - return (alloc_stms, - Op $ Inner $ SegOp $ SegRed lvl space reds' ts kbody') - + let reds' = zipWith (\red lam -> red {segBinOpLambda = lam}) reds lams + return + ( alloc_stms, + Op $ Inner $ SegOp $ SegRed lvl space reds' ts kbody' + ) transformExp (Op (Inner (SegOp (SegScan lvl space scans ts kbody)))) = do (alloc_stms, (lams, kbody')) <- transformScanRed lvl space (map segBinOpLambda scans) kbody - let scans' = zipWith (\red lam -> red { segBinOpLambda = lam }) scans lams - return (alloc_stms, - Op $ Inner $ SegOp $ SegScan lvl space scans' ts kbody') - + let scans' = zipWith (\red lam -> red {segBinOpLambda = lam}) scans lams + return + ( alloc_stms, + Op $ Inner $ SegOp $ SegScan lvl space scans' ts kbody' + ) transformExp (Op (Inner (SegOp (SegHist lvl space ops ts kbody)))) = do (alloc_stms, (lams', kbody')) <- transformScanRed lvl space lams kbody let ops' = zipWith onOp ops lams' - return (alloc_stms, - Op $ Inner $ SegOp $ SegHist lvl space ops' ts kbody') - where lams = map histOp ops - onOp op lam = op { histOp = lam } - + return + ( alloc_stms, + Op $ Inner $ SegOp $ SegHist lvl space ops' ts kbody' + ) + where + lams = map histOp ops + onOp op lam = op {histOp = lam} transformExp e = return (mempty, e) -transformScanRed :: SegLevel -> SegSpace - -> [Lambda KernelsMem] - -> KernelBody KernelsMem - -> ExpandM (Stms KernelsMem, ([Lambda KernelsMem], KernelBody KernelsMem)) +transformScanRed :: + SegLevel -> + SegSpace -> + [Lambda KernelsMem] -> + KernelBody KernelsMem -> + ExpandM (Stms KernelsMem, ([Lambda KernelsMem], KernelBody KernelsMem)) transformScanRed lvl space ops kbody = do bound_outside <- asks $ namesFromList . M.keys let (kbody', kbody_allocs) = @@ -150,9 +163,9 @@ (variant_allocs, invariant_allocs) = M.partition variantAlloc allocs case lvl of - SegGroup{} + SegGroup {} | not $ null variant_allocs -> - throwError "Cannot handle invariant allocations in SegGroup." + throwError "Cannot handle invariant allocations in SegGroup." _ -> return () @@ -160,185 +173,251 @@ ops'' <- forM ops' $ \op' -> localScope (scopeOf op') $ offsetMemoryInLambda op' return (alloc_stms, (ops'', kbody'')) - - where bound_in_kernel = namesFromList $ M.keys $ scopeOfSegSpace space <> - scopeOf (kernelBodyStms kbody) + where + bound_in_kernel = + namesFromList $ + M.keys $ + scopeOfSegSpace space + <> scopeOf (kernelBodyStms kbody) -allocsForBody :: Extraction - -> Extraction - -> SegLevel -> SegSpace - -> KernelBody KernelsMem - -> (Stms KernelsMem -> KernelBody KernelsMem -> OffsetM b) - -> ExpandM b +allocsForBody :: + Extraction -> + Extraction -> + SegLevel -> + SegSpace -> + KernelBody KernelsMem -> + (Stms KernelsMem -> KernelBody KernelsMem -> OffsetM b) -> + ExpandM b allocsForBody variant_allocs invariant_allocs lvl space kbody' m = do (alloc_offsets, alloc_stms) <- - memoryRequirements lvl space - (kernelBodyStms kbody') variant_allocs invariant_allocs + memoryRequirements + lvl + space + (kernelBodyStms kbody') + variant_allocs + invariant_allocs scope <- askScope let scope' = scopeOfSegSpace space <> M.map nameInfoConv scope - either throwError pure $ runOffsetM scope' alloc_offsets $ do - kbody'' <- offsetMemoryInKernelBody kbody' - m alloc_stms kbody'' + either throwError pure $ + runOffsetM scope' alloc_offsets $ do + kbody'' <- offsetMemoryInKernelBody kbody' + m alloc_stms kbody'' -memoryRequirements :: SegLevel -> SegSpace - -> Stms KernelsMem - -> Extraction -> Extraction - -> ExpandM (RebaseMap, Stms KernelsMem) +memoryRequirements :: + SegLevel -> + SegSpace -> + Stms KernelsMem -> + Extraction -> + Extraction -> + ExpandM (RebaseMap, Stms KernelsMem) memoryRequirements lvl space kstms variant_allocs invariant_allocs = do - ((num_threads, num_groups64, num_threads64), num_threads_stms) <- runBinder $ do - num_threads <- letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int32 OverflowUndef) - (unCount $ segNumGroups lvl) (unCount $ segGroupSize lvl) - num_groups64 <- letSubExp "num_groups64" $ - BasicOp $ ConvOp (SExt Int32 Int64) (unCount $ segNumGroups lvl) - num_threads64 <- letSubExp "num_threads64" $ BasicOp $ ConvOp (SExt Int32 Int64) num_threads - return (num_threads, num_groups64, num_threads64) + (num_threads, num_threads_stms) <- + runBinder $ + letSubExp "num_threads" $ + BasicOp $ + BinOp + (Mul Int64 OverflowUndef) + (unCount $ segNumGroups lvl) + (unCount $ segGroupSize lvl) (invariant_alloc_stms, invariant_alloc_offsets) <- - inScopeOf num_threads_stms $ expandedInvariantAllocations - (num_threads64, num_groups64, segNumGroups lvl, segGroupSize lvl) - space invariant_allocs + inScopeOf num_threads_stms $ + expandedInvariantAllocations + (num_threads, segNumGroups lvl, segGroupSize lvl) + space + invariant_allocs (variant_alloc_stms, variant_alloc_offsets) <- - inScopeOf num_threads_stms $ expandedVariantAllocations - num_threads space kstms variant_allocs + inScopeOf num_threads_stms $ + expandedVariantAllocations + num_threads + space + kstms + variant_allocs - return (invariant_alloc_offsets <> variant_alloc_offsets, - num_threads_stms <> invariant_alloc_stms <> variant_alloc_stms) + return + ( invariant_alloc_offsets <> variant_alloc_offsets, + num_threads_stms <> invariant_alloc_stms <> variant_alloc_stms + ) -- | A description of allocations that have been extracted, and how -- much memory (and which space) is needed. type Extraction = M.Map VName (SegLevel, SubExp, Space) -extractKernelBodyAllocations :: SegLevel -> Names -> Names -> KernelBody KernelsMem - -> (KernelBody KernelsMem, - Extraction) +extractKernelBodyAllocations :: + SegLevel -> + Names -> + Names -> + KernelBody KernelsMem -> + ( KernelBody KernelsMem, + Extraction + ) extractKernelBodyAllocations lvl bound_outside bound_kernel = extractGenericBodyAllocations lvl bound_outside bound_kernel kernelBodyStms $ - \stms kbody -> kbody { kernelBodyStms = stms } + \stms kbody -> kbody {kernelBodyStms = stms} -extractBodyAllocations :: SegLevel -> Names -> Names -> Body KernelsMem - -> (Body KernelsMem, Extraction) +extractBodyAllocations :: + SegLevel -> + Names -> + Names -> + Body KernelsMem -> + (Body KernelsMem, Extraction) extractBodyAllocations lvl bound_outside bound_kernel = extractGenericBodyAllocations lvl bound_outside bound_kernel bodyStms $ - \stms body -> body { bodyStms = stms } + \stms body -> body {bodyStms = stms} -extractLambdaAllocations :: SegLevel -> Names -> Names -> Lambda KernelsMem - -> (Lambda KernelsMem, Extraction) -extractLambdaAllocations lvl bound_outside bound_kernel lam = (lam { lambdaBody = body' }, allocs) - where (body', allocs) = extractBodyAllocations lvl bound_outside bound_kernel $ lambdaBody lam +extractLambdaAllocations :: + SegLevel -> + Names -> + Names -> + Lambda KernelsMem -> + (Lambda KernelsMem, Extraction) +extractLambdaAllocations lvl bound_outside bound_kernel lam = (lam {lambdaBody = body'}, allocs) + where + (body', allocs) = extractBodyAllocations lvl bound_outside bound_kernel $ lambdaBody lam -extractGenericBodyAllocations :: SegLevel -> Names -> Names - -> (body -> Stms KernelsMem) - -> (Stms KernelsMem -> body -> body) - -> body - -> (body, - Extraction) +extractGenericBodyAllocations :: + SegLevel -> + Names -> + Names -> + (body -> Stms KernelsMem) -> + (Stms KernelsMem -> body -> body) -> + body -> + ( body, + Extraction + ) extractGenericBodyAllocations lvl bound_outside bound_kernel get_stms set_stms body = - let (stms, allocs) = runWriter $ fmap catMaybes $ - mapM (extractStmAllocations lvl bound_outside bound_kernel) $ - stmsToList $ get_stms body - in (set_stms (stmsFromList stms) body, allocs) + let (stms, allocs) = + runWriter $ + fmap catMaybes $ + mapM (extractStmAllocations lvl bound_outside bound_kernel) $ + stmsToList $ get_stms body + in (set_stms (stmsFromList stms) body, allocs) expandable :: Space -> Bool expandable (Space "local") = False -expandable ScalarSpace{} = False +expandable ScalarSpace {} = False expandable _ = True -extractStmAllocations :: SegLevel -> Names -> Names -> Stm KernelsMem - -> Writer Extraction (Maybe (Stm KernelsMem)) +extractStmAllocations :: + SegLevel -> + Names -> + Names -> + 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 - tell $ M.singleton (patElemName patElem) (lvl, size, space) - return Nothing - - where expandableSize (Var v) = v `nameIn` bound_outside || v `nameIn` bound_kernel - expandableSize Constant{} = True - boundInKernel (Var v) = v `nameIn` bound_kernel - boundInKernel Constant{} = False - + tell $ M.singleton (patElemName patElem) (lvl, size, space) + return Nothing + where + expandableSize (Var v) = v `nameIn` bound_outside || v `nameIn` bound_kernel + expandableSize Constant {} = True + boundInKernel (Var v) = v `nameIn` bound_kernel + boundInKernel Constant {} = False extractStmAllocations lvl bound_outside bound_kernel stm = do e <- mapExpM (expMapper lvl) $ stmExp stm - return $ Just $ stm { stmExp = e } - where expMapper lvl' = identityMapper { mapOnBody = const $ onBody lvl' - , mapOnOp = onOp - } - - onBody lvl' body = do - let (body', allocs) = extractBodyAllocations lvl' bound_outside bound_kernel body - tell allocs - return body' - - onOp (Inner (SegOp op)) = - Inner . SegOp <$> mapSegOpM (opMapper (segLevel op)) op - onOp op = return op + return $ Just $ stm {stmExp = e} + where + expMapper lvl' = + identityMapper + { mapOnBody = const $ onBody lvl', + mapOnOp = onOp + } - opMapper lvl' = identitySegOpMapper { mapOnSegOpLambda = onLambda lvl' - , mapOnSegOpBody = onKernelBody lvl' - } + onBody lvl' body = do + let (body', allocs) = extractBodyAllocations lvl' bound_outside bound_kernel body + tell allocs + return body' - onKernelBody lvl' body = do - let (body', allocs) = extractKernelBodyAllocations lvl' bound_outside bound_kernel body - tell allocs - return body' + onOp (Inner (SegOp op)) = + Inner . SegOp <$> mapSegOpM (opMapper (segLevel op)) op + onOp op = return op - onLambda lvl' lam = do - body <- onBody lvl' $ lambdaBody lam - return lam { lambdaBody = body } + opMapper lvl' = + identitySegOpMapper + { mapOnSegOpLambda = onLambda lvl', + mapOnSegOpBody = onKernelBody lvl' + } -expandedInvariantAllocations :: (SubExp, SubExp, - Count NumGroups SubExp, Count GroupSize SubExp) - -> SegSpace - -> Extraction - -> ExpandM (Stms KernelsMem, RebaseMap) -expandedInvariantAllocations (num_threads64, num_groups64, - Count num_groups, Count group_size) - segspace - invariant_allocs = do - -- We expand the invariant allocations by adding an inner dimension - -- equal to the number of kernel threads. - (alloc_bnds, rebases) <- unzip <$> mapM expand (M.toList invariant_allocs) + onKernelBody lvl' body = do + let (body', allocs) = extractKernelBodyAllocations lvl' bound_outside bound_kernel body + tell allocs + return body' - return (mconcat alloc_bnds, mconcat rebases) - where expand (mem, (lvl, per_thread_size, space)) = do - total_size <- newVName "total_size" - let sizepat = Pattern [] [PatElem total_size $ MemPrim int64] - allocpat = Pattern [] [PatElem mem $ MemMem space] - num_users = case lvl of SegThread{} -> num_threads64 - SegGroup{} -> num_groups64 - return (stmsFromList - [Let sizepat (defAux ()) $ - BasicOp $ BinOp (Mul Int64 OverflowUndef) num_users per_thread_size, - Let allocpat (defAux ()) $ - Op $ Alloc (Var total_size) space], - M.singleton mem $ newBase lvl) + onLambda lvl' lam = do + body <- onBody lvl' $ lambdaBody lam + return lam {lambdaBody = body} - untouched d = DimSlice (fromInt32 0) d (fromInt32 1) +expandedInvariantAllocations :: + ( SubExp, + Count NumGroups SubExp, + Count GroupSize SubExp + ) -> + SegSpace -> + Extraction -> + ExpandM (Stms KernelsMem, RebaseMap) +expandedInvariantAllocations + ( num_threads, + Count num_groups, + Count group_size + ) + segspace + invariant_allocs = do + -- We expand the invariant allocations by adding an inner dimension + -- equal to the number of kernel threads. + (alloc_bnds, rebases) <- unzip <$> mapM expand (M.toList invariant_allocs) - newBase SegThread{} (old_shape, _) = - let num_dims = length old_shape - perm = num_dims : [0..num_dims-1] - root_ixfun = IxFun.iota (old_shape - ++ [primExpFromSubExp int32 num_groups * - primExpFromSubExp int32 group_size]) - permuted_ixfun = IxFun.permute root_ixfun perm - offset_ixfun = IxFun.slice permuted_ixfun $ - DimFix (LeafExp (segFlat segspace) int32) : - map untouched old_shape - in offset_ixfun + return (mconcat alloc_bnds, mconcat rebases) + where + expand (mem, (lvl, per_thread_size, space)) = do + total_size <- newVName "total_size" + let sizepat = Pattern [] [PatElem total_size $ MemPrim int64] + allocpat = Pattern [] [PatElem mem $ MemMem space] + num_users = case lvl of + SegThread {} -> num_threads + SegGroup {} -> num_groups + return + ( stmsFromList + [ Let sizepat (defAux ()) $ + BasicOp $ BinOp (Mul Int64 OverflowUndef) num_users per_thread_size, + Let allocpat (defAux ()) $ + Op $ Alloc (Var total_size) space + ], + M.singleton mem $ newBase lvl + ) - newBase SegGroup{} (old_shape, _) = - let root_ixfun = IxFun.iota (primExpFromSubExp int32 num_groups : old_shape) - offset_ixfun = IxFun.slice root_ixfun $ - DimFix (LeafExp (segFlat segspace) int32) : - map untouched old_shape - in offset_ixfun + untouched d = DimSlice 0 d 1 + newBase SegThread {} (old_shape, _) = + let num_dims = length old_shape + perm = num_dims : [0 .. num_dims -1] + root_ixfun = + IxFun.iota + ( old_shape + ++ [ pe64 num_groups * pe64 group_size + ] + ) + permuted_ixfun = IxFun.permute root_ixfun perm + offset_ixfun = + IxFun.slice permuted_ixfun $ + DimFix (le64 (segFlat segspace)) : + map untouched old_shape + in offset_ixfun + newBase SegGroup {} (old_shape, _) = + let root_ixfun = IxFun.iota (pe64 num_groups : old_shape) + offset_ixfun = + IxFun.slice root_ixfun $ + DimFix (le64 (segFlat segspace)) : + map untouched old_shape + in offset_ixfun -expandedVariantAllocations :: SubExp - -> SegSpace -> Stms KernelsMem - -> Extraction - -> ExpandM (Stms KernelsMem, RebaseMap) +expandedVariantAllocations :: + SubExp -> + SegSpace -> + Stms KernelsMem -> + Extraction -> + ExpandM (Stms KernelsMem, RebaseMap) expandedVariantAllocations _ _ _ variant_allocs | null variant_allocs = return (mempty, mempty) expandedVariantAllocations num_threads kspace kstms variant_allocs = do @@ -354,49 +433,67 @@ slice_stms' <- transformStms slice_stms_tmp let variant_allocs' :: [(VName, (SubExp, SubExp, Space))] - variant_allocs' = concat $ zipWith memInfo (map snd sizes_to_blocks) - (zip offsets size_sums) + variant_allocs' = + concat $ + zipWith + memInfo + (map snd sizes_to_blocks) + (zip offsets size_sums) memInfo blocks (offset, total_size) = - [ (mem, (Var offset, Var total_size, space)) | (mem, space) <- blocks ] + [(mem, (Var offset, Var total_size, space)) | (mem, space) <- blocks] -- We expand the invariant allocations by adding an inner dimension -- equal to the sum of the sizes required by different threads. (alloc_bnds, rebases) <- unzip <$> mapM expand variant_allocs' return (slice_stms' <> stmsFromList alloc_bnds, mconcat rebases) - where expand (mem, (offset, total_size, space)) = do - let allocpat = Pattern [] [PatElem mem $ MemMem space] - return (Let allocpat (defAux ()) $ Op $ Alloc total_size space, - M.singleton mem $ newBase offset) + where + expand (mem, (offset, total_size, space)) = do + let allocpat = Pattern [] [PatElem mem $ MemMem space] + return + ( Let allocpat (defAux ()) $ Op $ Alloc total_size space, + M.singleton mem $ newBase offset + ) - num_threads' = primExpFromSubExp int32 num_threads - gtid = LeafExp (segFlat kspace) int32 + num_threads' = pe64 num_threads + gtid = le64 $ segFlat kspace - -- For the variant allocations, we add an inner dimension, - -- which is then offset by a thread-specific amount. - newBase size_per_thread (old_shape, pt) = - let pt_size = fromInt32 $ primByteSize pt - elems_per_thread = sExt Int32 - (primExpFromSubExp int64 size_per_thread) - `quot` pt_size - root_ixfun = IxFun.iota [elems_per_thread, num_threads'] - offset_ixfun = IxFun.slice root_ixfun - [DimSlice (fromInt32 0) num_threads' (fromInt32 1), - DimFix gtid] - shapechange = if length old_shape == 1 - then map DimCoercion old_shape - else map DimNew old_shape - in IxFun.reshape offset_ixfun shapechange + -- For the variant allocations, we add an inner dimension, + -- which is then offset by a thread-specific amount. + newBase size_per_thread (old_shape, pt) = + let elems_per_thread = + pe64 size_per_thread `quot` primByteSize pt + root_ixfun = IxFun.iota [elems_per_thread, num_threads'] + offset_ixfun = + IxFun.slice + root_ixfun + [ DimSlice 0 num_threads' 1, + DimFix gtid + ] + shapechange = + if length old_shape == 1 + then map DimCoercion old_shape + else map DimNew old_shape + in IxFun.reshape offset_ixfun shapechange -- | A map from memory block names to new index function bases. - -type RebaseMap = M.Map VName (([PrimExp VName], PrimType) -> IxFun) +type RebaseMap = M.Map VName (([TPrimExp Int64 VName], PrimType) -> IxFun) -newtype OffsetM a = OffsetM (ReaderT (Scope KernelsMem) - (ReaderT RebaseMap (Either String)) a) - deriving (Applicative, Functor, Monad, - HasScope KernelsMem, LocalScope KernelsMem, - MonadError String) +newtype OffsetM a + = OffsetM + ( ReaderT + (Scope KernelsMem) + (ReaderT RebaseMap (Either String)) + a + ) + deriving + ( Applicative, + Functor, + Monad, + HasScope KernelsMem, + LocalScope KernelsMem, + MonadError String + ) runOffsetM :: Scope KernelsMem -> RebaseMap -> OffsetM a -> Either String a runOffsetM scope offsets (OffsetM m) = @@ -405,7 +502,7 @@ askRebaseMap :: OffsetM RebaseMap askRebaseMap = OffsetM $ lift ask -lookupNewBase :: VName -> ([PrimExp VName], PrimType) -> OffsetM (Maybe IxFun) +lookupNewBase :: VName -> ([TPrimExp Int64 VName], PrimType) -> OffsetM (Maybe IxFun) lookupNewBase name x = do offsets <- askRebaseMap return $ ($ x) <$> M.lookup name offsets @@ -413,17 +510,23 @@ offsetMemoryInKernelBody :: KernelBody KernelsMem -> OffsetM (KernelBody KernelsMem) offsetMemoryInKernelBody kbody = do scope <- askScope - stms' <- stmsFromList . snd <$> - mapAccumLM (\scope' -> localScope scope' . offsetMemoryInStm) scope - (stmsToList $ kernelBodyStms kbody) - return kbody { kernelBodyStms = stms' } + stms' <- + stmsFromList . snd + <$> mapAccumLM + (\scope' -> localScope scope' . offsetMemoryInStm) + scope + (stmsToList $ kernelBodyStms kbody) + return kbody {kernelBodyStms = stms'} offsetMemoryInBody :: Body KernelsMem -> OffsetM (Body KernelsMem) offsetMemoryInBody (Body dec stms res) = do scope <- askScope - stms' <- stmsFromList . snd <$> - mapAccumLM (\scope' -> localScope scope' . offsetMemoryInStm) scope - (stmsToList stms) + stms' <- + stmsFromList . snd + <$> mapAccumLM + (\scope' -> localScope scope' . offsetMemoryInStm) + scope + (stmsToList stms) return $ Body dec stms' res offsetMemoryInStm :: Stm KernelsMem -> OffsetM (Scope KernelsMem, Stm KernelsMem) @@ -434,67 +537,81 @@ -- Try to recompute the index function. Fall back to creating rebase -- operations with the RebaseMap. rts <- runReaderT (expReturns e') scope - let pat'' = Pattern (patternContextElements pat') - (zipWith pick (patternValueElements pat') rts) + let pat'' = + Pattern + (patternContextElements pat') + (zipWith pick (patternValueElements pat') rts) stm = Let pat'' dec e' let scope' = scopeOf stm <> scope return (scope', stm) - where pick :: PatElemT (MemInfo SubExp NoUniqueness MemBind) -> - ExpReturns -> PatElemT (MemInfo SubExp NoUniqueness MemBind) - pick (PatElem name (MemArray pt s u _ret)) - (MemArray _ _ _ (Just (ReturnsInBlock m extixfun))) - | Just ixfun <- instantiateIxFun extixfun = - PatElem name (MemArray pt s u (ArrayIn m ixfun)) - pick p _ = p + where + pick :: + PatElemT (MemInfo SubExp NoUniqueness MemBind) -> + ExpReturns -> + PatElemT (MemInfo SubExp NoUniqueness MemBind) + pick + (PatElem name (MemArray pt s u _ret)) + (MemArray _ _ _ (Just (ReturnsInBlock m extixfun))) + | Just ixfun <- instantiateIxFun extixfun = + PatElem name (MemArray pt s u (ArrayIn m ixfun)) + pick p _ = p - instantiateIxFun :: ExtIxFun -> Maybe IxFun - instantiateIxFun = traverse (traverse inst) - where inst Ext{} = Nothing - inst (Free x) = return x + instantiateIxFun :: ExtIxFun -> Maybe IxFun + instantiateIxFun = traverse (traverse inst) + where + inst Ext {} = Nothing + inst (Free x) = return x offsetMemoryInPattern :: Pattern KernelsMem -> OffsetM (Pattern KernelsMem) offsetMemoryInPattern (Pattern ctx vals) = do mapM_ inspectCtx ctx Pattern ctx <$> mapM inspectVal vals - where inspectVal patElem = do - new_dec <- offsetMemoryInMemBound $ patElemDec patElem - return patElem { patElemDec = new_dec } - inspectCtx patElem - | Mem space <- patElemType patElem, - expandable space = - throwError $ unwords ["Cannot deal with existential memory block", - pretty (patElemName patElem), - "when expanding inside kernels."] - | otherwise = return () + where + inspectVal patElem = do + new_dec <- offsetMemoryInMemBound $ patElemDec patElem + return patElem {patElemDec = new_dec} + inspectCtx patElem + | Mem space <- patElemType patElem, + expandable space = + throwError $ + unwords + [ "Cannot deal with existential memory block", + pretty (patElemName patElem), + "when expanding inside kernels." + ] + | otherwise = return () offsetMemoryInParam :: Param (MemBound u) -> OffsetM (Param (MemBound u)) offsetMemoryInParam fparam = do fparam' <- offsetMemoryInMemBound $ paramDec fparam - return fparam { paramDec = fparam' } + return fparam {paramDec = fparam'} offsetMemoryInMemBound :: MemBound u -> OffsetM (MemBound u) offsetMemoryInMemBound summary@(MemArray pt shape u (ArrayIn mem ixfun)) = do new_base <- lookupNewBase mem (IxFun.base ixfun, pt) - return $ fromMaybe summary $ do - new_base' <- new_base - return $ MemArray pt shape u $ ArrayIn mem $ IxFun.rebase new_base' ixfun + return $ + fromMaybe summary $ do + new_base' <- new_base + return $ MemArray pt shape u $ ArrayIn mem $ IxFun.rebase new_base' ixfun offsetMemoryInMemBound summary = return summary offsetMemoryInBodyReturns :: BodyReturns -> OffsetM BodyReturns offsetMemoryInBodyReturns br@(MemArray pt shape u (ReturnsInBlock mem ixfun)) | Just ixfun' <- isStaticIxFun ixfun = do - new_base <- lookupNewBase mem (IxFun.base ixfun', pt) - return $ fromMaybe br $ do + new_base <- lookupNewBase mem (IxFun.base ixfun', pt) + return $ + fromMaybe br $ do new_base' <- new_base return $ - MemArray pt shape u $ ReturnsInBlock mem $ - IxFun.rebase (fmap (fmap Free) new_base') ixfun + MemArray pt shape u $ + ReturnsInBlock mem $ + IxFun.rebase (fmap (fmap Free) new_base') ixfun offsetMemoryInBodyReturns br = return br offsetMemoryInLambda :: Lambda KernelsMem -> OffsetM (Lambda KernelsMem) offsetMemoryInLambda lam = inScopeOf lam $ do body <- offsetMemoryInBody $ lambdaBody lam - return $ lam { lambdaBody = body } + return $ lam {lambdaBody = body} offsetMemoryInExp :: Exp KernelsMem -> OffsetM (Exp KernelsMem) offsetMemoryInExp (DoLoop ctx val form body) = do @@ -505,20 +622,23 @@ body' <- localScope (scopeOfFParams ctxparams' <> scopeOfFParams valparams' <> scopeOf form) (offsetMemoryInBody body) return $ DoLoop (zip ctxparams' ctxinit) (zip valparams' valinit) form body' offsetMemoryInExp e = mapExpM recurse e - where recurse = identityMapper - { mapOnBody = \bscope -> localScope bscope . offsetMemoryInBody - , mapOnBranchType = offsetMemoryInBodyReturns - , mapOnOp = onOp - } - onOp (Inner (SegOp op)) = - Inner . SegOp <$> - localScope (scopeOfSegSpace (segSpace op)) (mapSegOpM segOpMapper op) - where segOpMapper = - identitySegOpMapper { mapOnSegOpBody = offsetMemoryInKernelBody - , mapOnSegOpLambda = offsetMemoryInLambda - } - onOp op = return op - + where + recurse = + identityMapper + { mapOnBody = \bscope -> localScope bscope . offsetMemoryInBody, + mapOnBranchType = offsetMemoryInBodyReturns, + mapOnOp = onOp + } + onOp (Inner (SegOp op)) = + Inner . SegOp + <$> localScope (scopeOfSegSpace (segSpace op)) (mapSegOpM segOpMapper op) + where + segOpMapper = + identitySegOpMapper + { mapOnSegOpBody = offsetMemoryInKernelBody, + mapOnSegOpLambda = offsetMemoryInLambda + } + onOp op = return op ---- Slicing allocation sizes out of a kernel. @@ -534,7 +654,7 @@ unAllocStms nested = fmap (stmsFromList . catMaybes) . mapM (unAllocStm nested) . stmsToList - unAllocStm nested stm@(Let _ _ (Op Alloc{})) + unAllocStm nested stm@(Let _ _ (Op Alloc {})) | nested = throwError $ "Cannot handle nested allocation: " ++ pretty stm | otherwise = return Nothing unAllocStm _ (Let pat dec e) = @@ -547,53 +667,68 @@ unAllocPattern pat@(Pattern ctx val) = Pattern <$> maybe bad return (mapM (rephrasePatElem unMem) ctx) - <*> maybe bad return (mapM (rephrasePatElem unMem) val) - where bad = Left $ "Cannot handle memory in pattern " ++ pretty pat + <*> maybe bad return (mapM (rephrasePatElem unMem) val) + where + bad = Left $ "Cannot handle memory in pattern " ++ pretty pat - unAllocOp Alloc{} = Left "unAllocOp: unhandled Alloc" - unAllocOp (Inner OtherOp{}) = Left "unAllocOp: unhandled OtherOp" + unAllocOp Alloc {} = Left "unAllocOp: unhandled Alloc" + unAllocOp (Inner OtherOp {}) = Left "unAllocOp: unhandled OtherOp" unAllocOp (Inner (SizeOp op)) = return $ SizeOp op unAllocOp (Inner (SegOp op)) = SegOp <$> mapSegOpM mapper op - where mapper = identitySegOpMapper { mapOnSegOpLambda = unAllocLambda - , mapOnSegOpBody = unAllocKernelBody - } + where + mapper = + identitySegOpMapper + { mapOnSegOpLambda = unAllocLambda, + mapOnSegOpBody = unAllocKernelBody + } unParam p = maybe bad return $ traverse unMem p - where bad = Left $ "Cannot handle memory-typed parameter '" ++ pretty p ++ "'" + where + bad = Left $ "Cannot handle memory-typed parameter '" ++ pretty p ++ "'" unT t = maybe bad return $ unMem t - where bad = Left $ "Cannot handle memory type '" ++ pretty t ++ "'" + where + bad = Left $ "Cannot handle memory type '" ++ pretty t ++ "'" - unAlloc' = Mapper { mapOnBody = const unAllocBody - , mapOnRetType = unT - , mapOnBranchType = unT - , mapOnFParam = unParam - , mapOnLParam = unParam - , mapOnOp = unAllocOp - , mapOnSubExp = Right - , mapOnVName = Right - } + unAlloc' = + Mapper + { mapOnBody = const unAllocBody, + mapOnRetType = unT, + mapOnBranchType = unT, + mapOnFParam = unParam, + mapOnLParam = unParam, + mapOnOp = unAllocOp, + mapOnSubExp = Right, + mapOnVName = Right + } unMem :: MemInfo d u ret -> Maybe (TypeBase (ShapeBase d) u) unMem (MemPrim pt) = Just $ Prim pt unMem (MemArray pt shape u _) = Just $ Array pt shape u -unMem MemMem{} = Nothing +unMem MemMem {} = Nothing unAllocScope :: Scope KernelsMem -> Scope Kernels.Kernels unAllocScope = M.mapMaybe unInfo - where unInfo (LetName dec) = LetName <$> unMem dec - unInfo (FParamName dec) = FParamName <$> unMem dec - unInfo (LParamName dec) = LParamName <$> unMem dec - unInfo (IndexName it) = Just $ IndexName it + where + unInfo (LetName dec) = LetName <$> unMem dec + unInfo (FParamName dec) = FParamName <$> unMem dec + unInfo (LParamName dec) = LParamName <$> unMem dec + unInfo (IndexName it) = Just $ IndexName it -removeCommonSizes :: Extraction - -> [(SubExp, [(VName, Space)])] +removeCommonSizes :: + Extraction -> + [(SubExp, [(VName, Space)])] removeCommonSizes = M.toList . foldl' comb mempty . M.toList - where comb m (mem, (_, size, space)) = M.insertWith (++) size [(mem, space)] m + where + comb m (mem, (_, size, space)) = M.insertWith (++) size [(mem, space)] m -sliceKernelSizes :: SubExp -> [SubExp] -> SegSpace -> Stms KernelsMem - -> ExpandM (Stms Kernels.Kernels, [VName], [VName]) +sliceKernelSizes :: + SubExp -> + [SubExp] -> + SegSpace -> + Stms KernelsMem -> + ExpandM (Stms Kernels.Kernels, [VName], [VName]) sliceKernelSizes num_threads sizes space kstms = do kstms' <- either throwError return $ unAllocKernelsStms kstms let num_sizes = length sizes @@ -604,55 +739,66 @@ (max_lam, _) <- flip runBinderT kernels_scope $ do xs <- replicateM num_sizes $ newParam "x" (Prim int64) ys <- replicateM num_sizes $ newParam "y" (Prim int64) - (zs, stms) <- localScope (scopeOfLParams $ xs ++ ys) $ collectStms $ - forM (zip xs ys) $ \(x,y) -> - letSubExp "z" $ BasicOp $ BinOp (SMax Int64) (Var $ paramName x) (Var $ paramName y) + (zs, stms) <- localScope (scopeOfLParams $ xs ++ ys) $ + collectStms $ + forM (zip xs ys) $ \(x, y) -> + letSubExp "z" $ BasicOp $ BinOp (SMax Int64) (Var $ paramName x) (Var $ paramName y) return $ Lambda (xs ++ ys) (mkBody stms zs) i64s - flat_gtid_lparam <- Param <$> newVName "flat_gtid" <*> pure (Prim (IntType Int32)) + flat_gtid_lparam <- Param <$> newVName "flat_gtid" <*> pure (Prim (IntType Int64)) (size_lam', _) <- flip runBinderT kernels_scope $ do params <- replicateM num_sizes $ newParam "x" (Prim int64) - (zs, stms) <- localScope (scopeOfLParams params <> - scopeOfLParams [flat_gtid_lparam]) $ collectStms $ do - - -- Even though this SegRed is one-dimensional, we need to - -- provide indexes corresponding to the original potentially - -- multi-dimensional construct. - let (kspace_gtids, kspace_dims) = unzip $ unSegSpace space - new_inds = unflattenIndex - (map (primExpFromSubExp int32) kspace_dims) - (primExpFromSubExp int32 $ Var $ paramName flat_gtid_lparam) - zipWithM_ letBindNames (map pure kspace_gtids) =<< mapM toExp new_inds + (zs, stms) <- localScope + ( scopeOfLParams params + <> scopeOfLParams [flat_gtid_lparam] + ) + $ collectStms $ do + -- Even though this SegRed is one-dimensional, we need to + -- provide indexes corresponding to the original potentially + -- multi-dimensional construct. + let (kspace_gtids, kspace_dims) = unzip $ unSegSpace space + new_inds = + unflattenIndex + (map pe64 kspace_dims) + (pe64 $ Var $ paramName flat_gtid_lparam) + zipWithM_ letBindNames (map pure kspace_gtids) =<< mapM toExp new_inds - mapM_ addStm kstms' - return sizes + mapM_ addStm kstms' + return sizes localScope (scopeOfSegSpace space) $ Kernels.simplifyLambda (Lambda [flat_gtid_lparam] (Body () stms zs) i64s) ((maxes_per_thread, size_sums), slice_stms) <- flip runBinderT kernels_scope $ do - num_threads_64 <- letSubExp "num_threads" $ - BasicOp $ ConvOp (SExt Int32 Int64) num_threads - - pat <- basicPattern [] <$> replicateM num_sizes - (newIdent "max_per_thread" $ Prim int64) + pat <- + basicPattern [] + <$> replicateM + num_sizes + (newIdent "max_per_thread" $ Prim int64) - w <- letSubExp "size_slice_w" =<< - foldBinOp (Mul Int32 OverflowUndef) (intConst Int32 1) (segSpaceDims space) + w <- + letSubExp "size_slice_w" + =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) (segSpaceDims space) - thread_space_iota <- letExp "thread_space_iota" $ BasicOp $ - Iota w (intConst Int32 0) (intConst Int32 1) Int32 - let red_op = SegBinOp Commutative max_lam - (replicate num_sizes $ intConst Int64 0) mempty + thread_space_iota <- + letExp "thread_space_iota" $ + BasicOp $ + Iota w (intConst Int64 0) (intConst Int64 1) Int64 + let red_op = + SegBinOp + Commutative + max_lam + (replicate num_sizes $ intConst Int64 0) + mempty lvl <- segThread "segred" - addStms =<< mapM renameStm =<< - nonSegRed lvl pat w [red_op] size_lam' [thread_space_iota] + addStms =<< mapM renameStm + =<< nonSegRed lvl pat w [red_op] size_lam' [thread_space_iota] size_sums <- forM (patternNames pat) $ \threads_max -> letExp "size_sum" $ - BasicOp $ BinOp (Mul Int64 OverflowUndef) (Var threads_max) num_threads_64 + BasicOp $ BinOp (Mul Int64 OverflowUndef) (Var threads_max) num_threads return (patternNames pat, size_sums) 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 @@ -1,976 +1,1181 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE DefaultSignatures #-} --- | A generic transformation for adding memory allocations to a --- Futhark program. Specialised by specific representations in --- submodules. -module Futhark.Pass.ExplicitAllocations - ( explicitAllocationsGeneric - , explicitAllocationsInStmsGeneric - , ExpHint(..) - , defaultExpHints - - , Allocable - , Allocator(..) - , AllocM - , AllocEnv(..) - , SizeSubst(..) - , allocInStms - , allocForArray - - , simplifiable - , arraySizeInBytesExp - - , mkLetNamesB' - , mkLetNamesB'' - - -- * Module re-exports - -- - -- These are highly likely to be needed by any downstream - -- users. - , module Control.Monad.Reader - , module Futhark.MonadFreshNames - , module Futhark.Pass - , module Futhark.Tools - ) -where - -import Control.Monad.State -import Control.Monad.Writer -import Control.Monad.Reader -import Control.Monad.RWS.Strict -import qualified Data.Map.Strict as M -import qualified Data.Set as S -import Data.Maybe -import Data.List (foldl', zip4, partition, sort) - -import qualified Futhark.Analysis.UsageTable as UT -import Futhark.Optimise.Simplify.Lore (mkWiseBody) -import Futhark.MonadFreshNames -import Futhark.IR.Mem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Tools -import Futhark.Optimise.Simplify.Engine (SimpleOps (..)) -import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Pass -import Futhark.Util (splitFromEnd, takeLast) - -data AllocStm = SizeComputation VName (PrimExp VName) - | Allocation VName SubExp Space - | ArrayCopy VName VName - deriving (Eq, Ord, Show) - -bindAllocStm :: (MonadBinder m, Op (Lore m) ~ MemOp inner) => - AllocStm -> m () -bindAllocStm (SizeComputation name pe) = - letBindNames [name] =<< toExp (coerceIntPrimExp Int64 pe) -bindAllocStm (Allocation name size space) = - letBindNames [name] $ Op $ Alloc size space -bindAllocStm (ArrayCopy name src) = - letBindNames [name] $ BasicOp $ Copy src - -class (MonadFreshNames m, HasScope lore m, Mem lore) => - Allocator lore m where - addAllocStm :: AllocStm -> m () - askDefaultSpace :: m Space - - default addAllocStm :: (Allocable fromlore lore, - m ~ AllocM fromlore lore) - => AllocStm -> m () - addAllocStm (SizeComputation name se) = - letBindNames [name] =<< toExp (coerceIntPrimExp Int64 se) - addAllocStm (Allocation name size space) = - letBindNames [name] $ Op $ allocOp size space - addAllocStm (ArrayCopy name src) = - letBindNames [name] $ BasicOp $ Copy src - - -- | The subexpression giving the number of elements we should - -- allocate space for. See 'ChunkMap' comment. - dimAllocationSize :: SubExp -> m SubExp - - default dimAllocationSize :: m ~ AllocM fromlore lore - => SubExp -> m SubExp - dimAllocationSize (Var v) = - -- It is important to recurse here, as the substitution may itself - -- be a chunk size. - maybe (return $ Var v) dimAllocationSize =<< asks (M.lookup v . chunkMap) - dimAllocationSize size = - return size - - -- | Get those names that are known to be constants at run-time. - askConsts :: m (S.Set VName) - - expHints :: Exp lore -> m [ExpHint] - expHints = defaultExpHints - -allocateMemory :: Allocator lore m => - String -> SubExp -> Space -> m VName -allocateMemory desc size space = do - v <- newVName desc - addAllocStm $ Allocation v size space - return v - -computeSize :: Allocator lore m => - String -> PrimExp VName -> m SubExp -computeSize desc se = do - v <- newVName desc - addAllocStm $ SizeComputation v se - return $ Var v - -type Allocable fromlore tolore = - (PrettyLore fromlore, PrettyLore tolore, - Mem tolore, - FParamInfo fromlore ~ DeclType, - LParamInfo fromlore ~ Type, - BranchType fromlore ~ ExtType, - RetType fromlore ~ DeclExtType, - BodyDec fromlore ~ (), - BodyDec tolore ~ (), - ExpDec tolore ~ (), - SizeSubst (Op tolore), - BinderOps tolore) - --- | A mapping from chunk names to their maximum size. XXX FIXME --- HACK: This is part of a hack to add loop-invariant allocations to --- reduce kernels, because memory expansion does not use range --- analysis yet (it should). -type ChunkMap = M.Map VName SubExp - -data AllocEnv fromlore tolore = - AllocEnv { chunkMap :: ChunkMap - , aggressiveReuse :: Bool - -- ^ Aggressively try to reuse memory in do-loops - - -- should be True inside kernels, False outside. - , allocSpace :: Space - -- ^ When allocating memory, put it in this memory space. - -- This is primarily used to ensure that group-wide - -- statements store their results in local memory. - , envConsts :: S.Set VName - -- ^ The set of names that are known to be constants at - -- kernel compile time. - , allocInOp :: Op fromlore -> AllocM fromlore tolore (Op tolore) - , envExpHints :: Exp tolore -> AllocM fromlore tolore [ExpHint] - } - --- | Monad for adding allocations to an entire program. -newtype AllocM fromlore tolore a = - AllocM (BinderT tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a) - deriving (Applicative, Functor, Monad, - MonadFreshNames, - HasScope tolore, - LocalScope tolore, - MonadReader (AllocEnv fromlore tolore)) - -instance (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - MonadBinder (AllocM fromlore tolore) where - type Lore (AllocM fromlore tolore) = tolore - - mkExpDecM _ _ = return () - - mkLetNamesM names e = do - pat <- patternWithAllocations names e - return $ Let pat (defAux ()) e - - mkBodyM bnds res = return $ Body () bnds res - - addStms = AllocM . addStms - collectStms (AllocM m) = AllocM $ collectStms m - -instance (Allocable fromlore tolore) => - Allocator tolore (AllocM fromlore tolore) where - expHints e = do - f <- asks envExpHints - f e - askDefaultSpace = asks allocSpace - - askConsts = asks envConsts - -runAllocM :: MonadFreshNames m => - (Op fromlore -> AllocM fromlore tolore (Op tolore)) - -> (Exp tolore -> AllocM fromlore tolore [ExpHint]) - -> AllocM fromlore tolore a -> m a -runAllocM handleOp hints (AllocM m) = - fmap fst $ modifyNameSource $ runState $ runReaderT (runBinderT m mempty) env - where env = AllocEnv { chunkMap = mempty - , aggressiveReuse = False - , allocSpace = DefaultSpace - , envConsts = mempty - , allocInOp = handleOp - , envExpHints = hints - } - --- | Monad for adding allocations to a single pattern. -newtype PatAllocM lore a = PatAllocM (RWS - (Scope lore) - [AllocStm] - VNameSource - a) - deriving (Applicative, Functor, Monad, - HasScope lore, - MonadWriter [AllocStm], - MonadFreshNames) - -instance Mem lore => Allocator lore (PatAllocM lore) where - addAllocStm = tell . pure - dimAllocationSize = return - askDefaultSpace = return DefaultSpace - askConsts = pure mempty - -runPatAllocM :: MonadFreshNames m => - PatAllocM lore a -> Scope lore - -> m (a, [AllocStm]) -runPatAllocM (PatAllocM m) mems = - modifyNameSource $ frob . runRWS m mems - where frob (a,s,w) = ((a,w),s) - -arraySizeInBytesExp :: Type -> PrimExp VName -arraySizeInBytesExp t = - foldl' (*) - (ValueExp $ IntValue $ Int64Value $ primByteSize $ elemType t) $ - map (sExt Int64 .primExpFromSubExp int32) (arrayDims t) - -arraySizeInBytesExpM :: Allocator lore m => Type -> m (PrimExp VName) -arraySizeInBytesExpM t = do - dims <- mapM dimAllocationSize (arrayDims t) - let dim_prod_i32 = product $ map (sExt Int64 . primExpFromSubExp int32) dims - let elm_size_i64 = ValueExp $ IntValue $ Int64Value $ primByteSize $ elemType t - return $ product [ dim_prod_i32, elm_size_i64 ] - -arraySizeInBytes :: Allocator lore m => Type -> m SubExp -arraySizeInBytes = computeSize "bytes" <=< arraySizeInBytesExpM - --- | Allocate memory for a value of the given type. -allocForArray :: Allocator lore m => - Type -> Space -> m VName -allocForArray t space = do - size <- arraySizeInBytes t - allocateMemory "mem" size space - -allocsForStm :: (Allocator lore m, ExpDec lore ~ ()) => - [Ident] -> [Ident] -> Exp lore - -> m (Stm lore) -allocsForStm sizeidents validents e = do - rts <- expReturns e - hints <- expHints e - (ctxElems, valElems) <- allocsForPattern sizeidents validents rts hints - return $ Let (Pattern ctxElems valElems) (defAux ()) e - -patternWithAllocations :: (Allocator lore m, ExpDec lore ~ ()) => - [VName] - -> Exp lore - -> m (Pattern lore) -patternWithAllocations names e = do - (ts',sizes) <- instantiateShapes' =<< expExtType e - let identForBindage name t = - pure $ Ident name t - vals <- sequence [ identForBindage name t | (name, t) <- zip names ts' ] - stmPattern <$> allocsForStm sizes vals e - -allocsForPattern :: Allocator lore m => - [Ident] -> [Ident] -> [ExpReturns] -> [ExpHint] - -> m ([PatElem lore], - [PatElem lore]) -allocsForPattern sizeidents validents rts hints = do - let sizes' = [ PatElem size $ MemPrim int32 | size <- map identName sizeidents ] - (vals, (exts, mems)) <- - runWriterT $ forM (zip3 validents rts hints) $ \(ident, rt, hint) -> do - let shape = arrayShape $ identType ident - case rt of - MemPrim _ -> do - summary <- lift $ summaryForBindage (identType ident) hint - return $ PatElem (identName ident) summary - - MemMem space -> - return $ PatElem (identName ident) $ - MemMem space - - MemArray bt _ u (Just (ReturnsInBlock mem extixfun)) -> do - (patels, ixfn) <- instantiateExtIxFun ident extixfun - tell (patels, []) - - return $ PatElem (identName ident) $ - MemArray bt shape u $ - ArrayIn mem ixfn - - MemArray _ extshape _ Nothing - | Just _ <- knownShape extshape -> do - summary <- lift $ summaryForBindage (identType ident) hint - return $ PatElem (identName ident) summary - - MemArray bt _ u (Just (ReturnsNewBlock space _ extixfn)) -> do - -- treat existential index function first - (patels, ixfn) <- instantiateExtIxFun ident extixfn - tell (patels, []) - - memid <- lift $ mkMemIdent ident space - tell ([], [PatElem (identName memid) $ MemMem space]) - return $ PatElem (identName ident) $ MemArray bt shape u $ - ArrayIn (identName memid) ixfn - - _ -> error "Impossible case reached in allocsForPattern!" - - return (sizes' <> exts <> mems, - vals) - where knownShape = mapM known . shapeDims - known (Free v) = Just v - known Ext{} = Nothing - - mkMemIdent :: (MonadFreshNames m) => Ident -> Space -> m Ident - mkMemIdent ident space = do - let memname = baseString (identName ident) <> "_mem" - newIdent memname $ Mem space - - instantiateExtIxFun :: MonadFreshNames m => - Ident -> ExtIxFun -> - m ([PatElemT (MemInfo d u ret)], IxFun) - instantiateExtIxFun idd ext_ixfn = do - let isAndPtps = S.toList $ - foldMap onlyExts $ - foldMap leafExpTypes ext_ixfn - - -- Find the existentials that reuse the sizeidents, and - -- those that need new pattern elements. Assumes that the - -- Exts form a contiguous interval of integers. - let (size_exts, new_exts) = - span (( do - v <- newVName $ baseString (identName idd) <> "_ixfn" - return ((Ext i, LeafExp (Free v) t), - PatElem v $ MemPrim t) - let size_substs = zipWith (\(i, t) ident -> - (Ext i, LeafExp (Free (identName ident)) t)) - size_exts sizeidents - substs = M.fromList $ new_substs <> size_substs - ixfn <- instantiateIxFun $ IxFun.substituteInIxFun substs ext_ixfn - - return (patels, ixfn) - -onlyExts :: (Ext a, PrimType) -> S.Set (Int, PrimType) -onlyExts (Free _, _) = S.empty -onlyExts (Ext i, t) = S.singleton (i, t) - - -instantiateIxFun :: Monad m => ExtIxFun -> m IxFun -instantiateIxFun = traverse $ traverse inst - where inst Ext{} = error "instantiateIxFun: not yet" - inst (Free x) = return x - -summaryForBindage :: Allocator lore m => - Type -> ExpHint - -> m (MemBound NoUniqueness) -summaryForBindage (Prim bt) _ = - return $ MemPrim bt -summaryForBindage (Mem space) _ = - return $ MemMem space -summaryForBindage t@(Array bt shape u) NoHint = do - m <- allocForArray t =<< askDefaultSpace - return $ directIxFun bt shape u m t -summaryForBindage t (Hint ixfun space) = do - let bt = elemType t - bytes <- computeSize "bytes" $ - product [product $ map (sExt Int64) $ IxFun.base ixfun, - fromIntegral (primByteSize (elemType t)::Int64)] - m <- allocateMemory "mem" bytes space - return $ MemArray bt (arrayShape t) NoUniqueness $ ArrayIn m ixfun - -lookupMemSpace :: (HasScope lore m, Monad m) => VName -> m Space -lookupMemSpace v = do - t <- lookupType v - case t of - Mem space -> return space - _ -> error $ "lookupMemSpace: " ++ pretty v ++ " is not a memory block." - -directIxFun :: PrimType -> Shape -> u -> VName -> Type -> MemBound u -directIxFun bt shape u mem t = - let ixf = IxFun.iota $ map (primExpFromSubExp int32) $ arrayDims t - in MemArray bt shape u $ ArrayIn mem ixf - - -allocInFParams :: (Allocable fromlore tolore) => - [(FParam fromlore, Space)] -> - ([FParam tolore] -> AllocM fromlore tolore a) - -> AllocM fromlore tolore a -allocInFParams params m = do - (valparams, (ctxparams, memparams)) <- - runWriterT $ mapM (uncurry allocInFParam) params - let params' = ctxparams <> memparams <> valparams - summary = scopeOfFParams params' - localScope summary $ m params' - -allocInFParam :: (Allocable fromlore tolore) => - FParam fromlore - -> Space - -> WriterT ([FParam tolore], [FParam tolore]) - (AllocM fromlore tolore) (FParam tolore) -allocInFParam param pspace = - case paramDeclType param of - Array bt shape u -> do - let memname = baseString (paramName param) <> "_mem" - ixfun = IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape - mem <- lift $ newVName memname - tell ([], [Param mem $ MemMem pspace]) - return param { paramDec = MemArray bt shape u $ ArrayIn mem ixfun } - Prim bt -> - return param { paramDec = MemPrim bt } - Mem space -> - return param { paramDec = MemMem space } - -allocInMergeParams :: (Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore)) => - [(FParam fromlore,SubExp)] - -> ([FParam tolore] - -> [FParam tolore] - -> ([SubExp] -> AllocM fromlore tolore ([SubExp], [SubExp])) - -> AllocM fromlore tolore a) - -> AllocM fromlore tolore a -allocInMergeParams merge m = do - ((valparams, handle_loop_subexps), (ctx_params, mem_params)) <- - runWriterT $ unzip <$> mapM allocInMergeParam merge - let mergeparams' = ctx_params <> mem_params <> valparams - summary = scopeOfFParams mergeparams' - - mk_loop_res ses = do - (valargs, (ctxargs, memargs)) <- - runWriterT $ zipWithM ($) handle_loop_subexps ses - return (ctxargs <> memargs, valargs) - - localScope summary $ m (ctx_params <> mem_params) valparams mk_loop_res - where - allocInMergeParam :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - (Param DeclType, SubExp) -> - WriterT - ([FParam tolore], [FParam tolore]) - (AllocM fromlore tolore) - (FParam tolore, SubExp -> WriterT ([SubExp], [SubExp]) (AllocM fromlore tolore) SubExp) - allocInMergeParam (mergeparam, Var v) - | Array bt shape u <- paramDeclType mergeparam = do - (mem', _) <- lift $ lookupArraySummary v - mem_space <- lift $ lookupMemSpace mem' - - (_, ext_ixfun, substs, _) <- lift $ existentializeArray mem_space v - - (ctx_params, param_ixfun_substs) <- - unzip <$> - mapM (\primExp -> do - let pt = primExpType primExp - vname <- lift $ newVName "ctx_param_ext" - return (Param vname $ MemPrim pt, - fmap Free $ primExpFromSubExp int32 $ Var vname)) - substs - - tell (ctx_params, []) - - param_ixfun <- instantiateIxFun $ - IxFun.substituteInIxFun (M.fromList $ zip (fmap Ext [0..]) param_ixfun_substs) - ext_ixfun - - mem_name <- newVName "mem_param" - tell ([], [Param mem_name $ MemMem mem_space]) - - return (mergeparam { paramDec = MemArray bt shape u $ ArrayIn mem_name param_ixfun }, - ensureArrayIn mem_space) - - allocInMergeParam (mergeparam, _) = doDefault mergeparam =<< lift askDefaultSpace - - doDefault mergeparam space = do - mergeparam' <- allocInFParam mergeparam space - return (mergeparam', linearFuncallArg (paramType mergeparam) space) - - --- Returns the existentialized index function, the list of substituted values and the memory location. -existentializeArray :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - Space -> VName -> AllocM fromlore tolore (SubExp, ExtIxFun, [PrimExp VName], VName) -existentializeArray space v = do - (mem', ixfun) <- lookupArraySummary v - sp <- lookupMemSpace mem' - - let (ext_ixfun', substs') = runState (IxFun.existentialize ixfun) [] - - case (ext_ixfun', sp == space) of - (Just x, True) -> return (Var v, x, substs', mem') - _ -> do - (mem, subexp) <- allocLinearArray space (baseString v) v - ixfun' <- fromJust <$> subExpIxFun subexp - let (ext_ixfun, substs) = runState (IxFun.existentialize ixfun') [] - return (subexp, fromJust ext_ixfun, substs, mem) - - - -ensureArrayIn :: (Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore)) => - Space -> SubExp - -> WriterT ([SubExp], [SubExp]) (AllocM fromlore tolore) SubExp -ensureArrayIn _ (Constant v) = - error $ "ensureArrayIn: " ++ pretty v ++ " cannot be an array." -ensureArrayIn space (Var v) = do - (sub_exp, _, substs, mem) <- lift $ existentializeArray space v - (ctx_vals, _) <- - unzip <$> - mapM (\s -> do - vname <- lift $ letExp "ctx_val" =<< toExp s - return (Var vname, fmap Free $ primExpFromSubExp int32 $ Var vname)) - substs - - tell (ctx_vals, [Var mem]) - - return sub_exp - -ensureDirectArray :: (Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore)) => - Maybe Space -> VName -> AllocM fromlore tolore (VName, SubExp) -ensureDirectArray space_ok v = do - (mem, ixfun) <- lookupArraySummary v - mem_space <- lookupMemSpace mem - default_space <- askDefaultSpace - if IxFun.isDirect ixfun && maybe True (==mem_space) space_ok - then return (mem, Var v) - else needCopy (fromMaybe default_space space_ok) - where needCopy space = - -- We need to do a new allocation, copy 'v', and make a new - -- binding for the size of the memory block. - allocLinearArray space (baseString v) v - -allocLinearArray :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - Space -> String -> VName - -> AllocM fromlore tolore (VName, SubExp) -allocLinearArray space s v = do - t <- lookupType v - mem <- allocForArray t space - v' <- newIdent (s ++ "_linear") t - let ixfun = directIxFun (elemType t) (arrayShape t) NoUniqueness mem t - let pat = Pattern [] [PatElem (identName v') ixfun] - addStm $ Let pat (defAux ()) $ BasicOp $ Copy v - return (mem, Var $ identName v') - -funcallArgs :: (Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore)) => - [(SubExp,Diet)] -> AllocM fromlore tolore [(SubExp,Diet)] -funcallArgs args = do - (valargs, (ctx_args, mem_and_size_args)) <- runWriterT $ forM args $ \(arg,d) -> do - t <- lift $ subExpType arg - space <- lift askDefaultSpace - arg' <- linearFuncallArg t space arg - return (arg', d) - return $ map (,Observe) (ctx_args <> mem_and_size_args) <> valargs - -linearFuncallArg :: (Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore)) => - Type -> Space -> SubExp - -> WriterT ([SubExp], [SubExp]) (AllocM fromlore tolore) SubExp -linearFuncallArg Array{} space (Var v) = do - (mem, arg') <- lift $ ensureDirectArray (Just space) v - tell ([], [Var mem]) - return arg' -linearFuncallArg _ _ arg = - return arg - -explicitAllocationsGeneric :: (Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore)) => - (Op fromlore -> AllocM fromlore tolore (Op tolore)) - -> (Exp tolore -> AllocM fromlore tolore [ExpHint]) - -> Pass fromlore tolore -explicitAllocationsGeneric handleOp hints = - Pass "explicit allocations" "Transform program to explicit memory representation" $ - intraproceduralTransformationWithConsts onStms allocInFun - where onStms stms = runAllocM handleOp hints $ allocInStms stms pure - - allocInFun consts (FunDef entry attrs fname rettype params fbody) = - runAllocM handleOp hints $ inScopeOf consts $ - allocInFParams (zip params $ repeat DefaultSpace) $ \params' -> do - fbody' <- insertStmsM $ allocInFunBody - (map (const $ Just DefaultSpace) rettype) fbody - return $ FunDef entry attrs fname (memoryInDeclExtType rettype) params' fbody' - -explicitAllocationsInStmsGeneric :: (MonadFreshNames m, HasScope tolore m, - Allocable fromlore tolore) => - (Op fromlore -> AllocM fromlore tolore (Op tolore)) - -> (Exp tolore -> AllocM fromlore tolore [ExpHint]) - -> Stms fromlore -> m (Stms tolore) -explicitAllocationsInStmsGeneric handleOp hints stms = do - scope <- askScope - runAllocM handleOp hints $ localScope scope $ allocInStms stms return - -memoryInDeclExtType :: [DeclExtType] -> [FunReturns] -memoryInDeclExtType ts = evalState (mapM addMem ts) $ startOfFreeIDRange ts - where addMem (Prim t) = return $ MemPrim t - addMem Mem{} = error "memoryInDeclExtType: too much memory" - addMem (Array bt shape u) = do - i <- get <* modify (+1) - return $ MemArray bt shape u $ ReturnsNewBlock DefaultSpace i $ - IxFun.iota $ map convert $ shapeDims shape - - convert (Ext i) = LeafExp (Ext i) int32 - convert (Free v) = Free <$> primExpFromSubExp int32 v - -startOfFreeIDRange :: [TypeBase ExtShape u] -> Int -startOfFreeIDRange = S.size . shapeContext - -bodyReturnMemCtx :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - SubExp -> AllocM fromlore tolore [SubExp] -bodyReturnMemCtx Constant{} = - return [] -bodyReturnMemCtx (Var v) = do - info <- lookupMemInfo v - case info of - MemPrim{} -> return [] - MemMem{} -> return [] -- should not happen - MemArray _ _ _ (ArrayIn mem _) -> return [Var mem] - -allocInFunBody :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - [Maybe Space] -> Body fromlore -> AllocM fromlore tolore (Body tolore) -allocInFunBody space_oks (Body _ bnds res) = - allocInStms bnds $ \bnds' -> do - (res'', allocs) <- collectStms $ do - res' <- zipWithM ensureDirect space_oks' res - let (ctx_res, val_res) = splitFromEnd num_vals res' - mem_ctx_res <- concat <$> mapM bodyReturnMemCtx val_res - return $ ctx_res <> mem_ctx_res <> val_res - return $ Body () (bnds'<>allocs) res'' - where num_vals = length space_oks - space_oks' = replicate (length res - num_vals) Nothing ++ space_oks - -ensureDirect :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - Maybe Space -> SubExp -> AllocM fromlore tolore SubExp -ensureDirect _ se@Constant{} = return se -ensureDirect space_ok (Var v) = do - bt <- primType <$> lookupType v - if bt - then return $ Var v - else do (_, v') <- ensureDirectArray space_ok v - return v' - -allocInStms :: (Allocable fromlore tolore) => - Stms fromlore -> (Stms tolore -> AllocM fromlore tolore a) - -> AllocM fromlore tolore a -allocInStms origstms m = allocInStms' (stmsToList origstms) mempty - where allocInStms' [] stms' = - m stms' - allocInStms' (x:xs) stms' = do - allocstms <- allocInStm' x - localScope (scopeOf allocstms) $ do - let stms_substs = foldMap sizeSubst allocstms - stms_consts = foldMap stmConsts allocstms - f env = env { chunkMap = stms_substs <> chunkMap env - , envConsts = stms_consts <> envConsts env - } - local f $ allocInStms' xs (stms'<>allocstms) - allocInStm' stm = - collectStms_ $ auxing (stmAux stm) $ allocInStm stm - -allocInStm :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - Stm fromlore -> AllocM fromlore tolore () -allocInStm (Let (Pattern sizeElems valElems) _ e) = do - e' <- allocInExp e - let sizeidents = map patElemIdent sizeElems - validents = map patElemIdent valElems - bnd <- allocsForStm sizeidents validents e' - addStm bnd - -allocInExp :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - Exp fromlore -> AllocM fromlore tolore (Exp tolore) -allocInExp (DoLoop ctx val form (Body () bodybnds bodyres)) = - allocInMergeParams ctx $ \_ ctxparams' _ -> - allocInMergeParams val $ - \new_ctx_params valparams' mk_loop_val -> do - form' <- allocInLoopForm form - localScope (scopeOf form') $ do - (valinit_ctx, valinit') <- mk_loop_val valinit - body' <- insertStmsM $ allocInStms bodybnds $ \bodybnds' -> do - ((val_ses,valres'),val_retbnds) <- collectStms $ mk_loop_val valres - return $ Body () (bodybnds'<>val_retbnds) (ctxres++val_ses++valres') - return $ - DoLoop - (zip (ctxparams'++new_ctx_params) (ctxinit++valinit_ctx)) - (zip valparams' valinit') - form' body' - where (_ctxparams, ctxinit) = unzip ctx - (_valparams, valinit) = unzip val - (ctxres, valres) = splitAt (length ctx) bodyres -allocInExp (Apply fname args rettype loc) = do - args' <- funcallArgs args - return $ Apply fname args' (memoryInDeclExtType rettype) loc -allocInExp (If cond tbranch0 fbranch0 (IfDec rets ifsort)) = do - let num_rets = length rets - -- switch to the explicit-mem rep, but do nothing about results - (tbranch, tm_ixfs) <- allocInIfBody num_rets tbranch0 - (fbranch, fm_ixfs) <- allocInIfBody num_rets fbranch0 - tspaces <- mkSpaceOks num_rets tbranch - fspaces <- mkSpaceOks num_rets fbranch - -- try to generalize (antiunify) the index functions of the then and else bodies - let sp_substs = zipWith generalize (zip tspaces tm_ixfs) (zip fspaces fm_ixfs) - (spaces, subs) = unzip sp_substs - tsubs = map (selectSub fst) subs - fsubs = map (selectSub snd) subs - (tbranch', trets) <- addResCtxInIfBody rets tbranch spaces tsubs - (fbranch', frets) <- addResCtxInIfBody rets fbranch spaces fsubs - if frets /= trets then error "In allocInExp, IF case: antiunification of then/else produce different ExtInFn!" - else do -- above is a sanity check; implementation continues on else branch - let res_then = bodyResult tbranch' - res_else = bodyResult fbranch' - size_ext = length res_then - length trets - (ind_ses0, r_then_else) = - partition (\(r_then, r_else, _) -> r_then == r_else) $ - zip3 res_then res_else [0 .. size_ext - 1] - (r_then_ext, r_else_ext, _) = unzip3 r_then_else - ind_ses = zipWith (\(se, _, i) k -> (i-k, se)) ind_ses0 - [0 .. length ind_ses0 - 1] - rets'' = foldl (\acc (i, se) -> fixExt i se acc) trets ind_ses - tbranch'' = tbranch' { bodyResult = r_then_ext ++ drop size_ext res_then } - fbranch'' = fbranch' { bodyResult = r_else_ext ++ drop size_ext res_else } - res_if_expr = If cond tbranch'' fbranch'' $ IfDec rets'' ifsort - return res_if_expr - where generalize :: (Maybe Space, Maybe IxFun) -> (Maybe Space, Maybe IxFun) - -> (Maybe Space, Maybe (ExtIxFun, [(PrimExp VName, PrimExp VName)])) - generalize (Just sp1, Just ixf1) (Just sp2, Just ixf2) = - if sp1 /= sp2 then (Just sp1, Nothing) - else case IxFun.leastGeneralGeneralization ixf1 ixf2 of - Just (ixf, m) -> (Just sp1, Just (ixf, m)) - Nothing -> (Just sp1, Nothing) - generalize (mbsp1, _) _ = (mbsp1, Nothing) - - selectSub :: ((a, a) -> a) -> Maybe (ExtIxFun, [(a, a)]) -> - Maybe (ExtIxFun, [a]) - selectSub f (Just (ixfn, m)) = Just (ixfn, map f m) - selectSub _ Nothing = Nothing - - -- | Just introduces the new representation (index functions); but - -- does not unify (e.g., does not ensures direct); implementation - -- extends `allocInBodyNoDirect`, but also return `MemBind` - allocInIfBody :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - Int -> Body fromlore -> AllocM fromlore tolore (Body tolore, [Maybe IxFun]) - allocInIfBody num_vals (Body _ bnds res) = - allocInStms bnds $ \bnds' -> do - let (_, val_res) = splitFromEnd num_vals res - mem_ixfs <- mapM subExpIxFun val_res - return (Body () bnds' res, mem_ixfs) -allocInExp e = mapExpM alloc e - where alloc = - identityMapper { mapOnBody = error "Unhandled Body in ExplicitAllocations" - , mapOnRetType = error "Unhandled RetType in ExplicitAllocations" - , mapOnBranchType = error "Unhandled BranchType in ExplicitAllocations" - , mapOnFParam = error "Unhandled FParam in ExplicitAllocations" - , mapOnLParam = error "Unhandled LParam in ExplicitAllocations" - , mapOnOp = \op -> do handle <- asks allocInOp - handle op - } - - - -subExpIxFun :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - SubExp -> AllocM fromlore tolore (Maybe IxFun) -subExpIxFun Constant{} = return Nothing -subExpIxFun (Var v) = do - info <- lookupMemInfo v - case info of - MemArray _ptp _shp _u (ArrayIn _ ixf) -> return $ Just ixf - _ -> return Nothing - - -addResCtxInIfBody :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - [ExtType] -> Body tolore -> [Maybe Space] -> - [Maybe (ExtIxFun, [PrimExp VName])] -> - AllocM fromlore tolore (Body tolore, [BodyReturns]) -addResCtxInIfBody ifrets (Body _ bnds res) spaces substs = do - let num_vals = length ifrets - (ctx_res, val_res) = splitFromEnd num_vals res - ((res', bodyrets'), all_body_stms) <- collectStms $ do - mapM_ addStm bnds - (val_res', ext_ses_res, mem_ctx_res, bodyrets, total_existentials) <- - foldM helper ([], [], [], [], length ctx_res) (zip4 ifrets val_res substs spaces) - return (ctx_res <> ext_ses_res <> mem_ctx_res <> val_res', - -- We need to adjust the ReturnsNewBlock existentials, because they - -- should always be numbered _after_ all other existentials in the - -- return values. - reverse $ fst $ foldl adjustNewBlockExistential ([], total_existentials) bodyrets) - body' <- mkBodyM all_body_stms res' - return (body', bodyrets') - where - helper (res_acc, ext_acc, ctx_acc, br_acc, k) (ifr, r, mbixfsub, sp) = - case mbixfsub of - Nothing -> do - -- does NOT generalize/antiunify; ensure direct - r' <- ensureDirect sp r - mem_ctx_r <- bodyReturnMemCtx r' - let body_ret = inspect ifr sp - return (res_acc ++ [r'], - ext_acc, - ctx_acc ++ mem_ctx_r, - br_acc ++ [body_ret], - k) - Just (ixfn, m) -> do -- generalizes - let i = length m - ext_ses <- mapM (toSubExp "ixfn_exist") m - mem_ctx_r <- bodyReturnMemCtx r - let sp' = fromMaybe DefaultSpace sp - ixfn' = fmap (adjustExtPE k) ixfn - exttp = case ifr of - Array pt shp' u -> - MemArray pt shp' u $ - ReturnsNewBlock sp' 0 ixfn' - _ -> error "Impossible case reached in addResCtxInIfBody" - return (res_acc ++ [r], - ext_acc ++ ext_ses, - ctx_acc ++ mem_ctx_r, - br_acc ++ [exttp], - k + i) - - adjustNewBlockExistential :: ([BodyReturns], Int) -> BodyReturns -> ([BodyReturns], Int) - adjustNewBlockExistential (acc, k) (MemArray pt shp u (ReturnsNewBlock space _ ixfun)) = - (MemArray pt shp u (ReturnsNewBlock space k ixfun) : acc, k + 1) - adjustNewBlockExistential (acc, k) x = (x : acc, k) - - inspect (Array pt shape u) space = - let space' = fromMaybe DefaultSpace space - bodyret = MemArray pt shape u $ ReturnsNewBlock space' 0 $ - IxFun.iota $ map convert $ shapeDims shape - in bodyret - inspect (Prim pt) _ = MemPrim pt - inspect (Mem space) _ = MemMem space - - convert (Ext i) = LeafExp (Ext i) int32 - convert (Free v) = Free <$> primExpFromSubExp int32 v - - adjustExtV :: Int -> Ext VName -> Ext VName - adjustExtV _ (Free v) = Free v - adjustExtV k (Ext i) = Ext (k + i) - - adjustExtPE :: Int -> PrimExp (Ext VName) -> PrimExp (Ext VName) - adjustExtPE k = fmap (adjustExtV k) - -mkSpaceOks :: (Mem tolore, LocalScope tolore m) => - Int -> Body tolore -> m [Maybe Space] -mkSpaceOks num_vals (Body _ stms res) = - inScopeOf stms $ - mapM mkSpaceOK $ takeLast num_vals res - where mkSpaceOK (Var v) = do - v_info <- lookupMemInfo v - case v_info of MemArray _ _ _ (ArrayIn mem _) -> do - mem_info <- lookupMemInfo mem - case mem_info of MemMem space -> return $ Just space - _ -> return Nothing - _ -> return Nothing - mkSpaceOK _ = return Nothing - -allocInLoopForm :: (Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore)) => - LoopForm fromlore -> AllocM fromlore tolore (LoopForm tolore) -allocInLoopForm (WhileLoop v) = return $ WhileLoop v -allocInLoopForm (ForLoop i it n loopvars) = - ForLoop i it n <$> mapM allocInLoopVar loopvars - where allocInLoopVar (p,a) = do - (mem, ixfun) <- lookupArraySummary a - case paramType p of - Array bt shape u -> do - dims <- map (primExpFromSubExp int32) . arrayDims <$> lookupType a - let ixfun' = IxFun.slice ixfun $ - fullSliceNum dims [DimFix $ LeafExp i int32] - return (p { paramDec = MemArray bt shape u $ ArrayIn mem ixfun' }, a) - Prim bt -> - return (p { paramDec = MemPrim bt }, a) - Mem space -> - return (p { paramDec = MemMem space }, a) - -class SizeSubst op where - opSizeSubst :: PatternT dec -> op -> ChunkMap - opIsConst :: op -> Bool - opIsConst = const False - -instance SizeSubst () where - opSizeSubst _ _ = mempty - -instance SizeSubst op => SizeSubst (MemOp op) where - opSizeSubst pat (Inner op) = opSizeSubst pat op - opSizeSubst _ _ = mempty - - opIsConst (Inner op) = opIsConst op - opIsConst _ = False - -sizeSubst :: SizeSubst (Op lore) => Stm lore -> ChunkMap -sizeSubst (Let pat _ (Op op)) = opSizeSubst pat op -sizeSubst _ = mempty - -stmConsts :: SizeSubst (Op lore) => Stm lore -> S.Set VName -stmConsts (Let pat _ (Op op)) - | opIsConst op = S.fromList $ patternNames pat -stmConsts _ = mempty - -mkLetNamesB' :: (Op (Lore m) ~ MemOp inner, - MonadBinder m, ExpDec (Lore m) ~ (), - Allocator (Lore m) (PatAllocM (Lore m))) => - ExpDec (Lore m) -> [VName] -> Exp (Lore m) -> m (Stm (Lore m)) -mkLetNamesB' dec names e = do - scope <- askScope - pat <- bindPatternWithAllocations scope names e - return $ Let pat (defAux dec) e - -mkLetNamesB'' :: (Op (Lore m) ~ MemOp inner, ExpDec lore ~ (), - HasScope (Engine.Wise lore) m, Allocator lore (PatAllocM lore), - MonadBinder m, Engine.CanBeWise (Op lore)) => - [VName] -> Exp (Engine.Wise lore) - -> m (Stm (Engine.Wise lore)) -mkLetNamesB'' names e = do - scope <- Engine.removeScopeWisdom <$> askScope - (pat, prestms) <- runPatAllocM (patternWithAllocations names $ Engine.removeExpWisdom e) scope - mapM_ bindAllocStm prestms - let pat' = Engine.addWisdomToPattern pat e - dec = Engine.mkWiseExpDec pat' () e - return $ Let pat' (defAux dec) e - -simplifiable :: (Engine.SimplifiableLore lore, - ExpDec lore ~ (), - BodyDec lore ~ (), - Op lore ~ MemOp inner, - Allocator lore (PatAllocM lore)) => - (Engine.OpWithWisdom inner -> UT.UsageTable) - -> (inner -> Engine.SimpleM lore (Engine.OpWithWisdom inner, Stms (Engine.Wise lore))) - -> SimpleOps lore -simplifiable innerUsage simplifyInnerOp = - SimpleOps mkExpDecS' mkBodyS' protectOp opUsage simplifyOp - where mkExpDecS' _ pat e = - return $ Engine.mkWiseExpDec pat () e - - mkBodyS' _ bnds res = return $ mkWiseBody () bnds res - - protectOp taken pat (Alloc size space) = Just $ do - tbody <- resultBodyM [size] - fbody <- resultBodyM [intConst Int64 0] - size' <- letSubExp "hoisted_alloc_size" $ - If taken tbody fbody $ IfDec [MemPrim int64] IfFallback - letBind pat $ Op $ Alloc size' space - protectOp _ _ _ = Nothing - - opUsage (Alloc (Var size) _) = - UT.sizeUsage size - opUsage (Alloc _ _) = - mempty - opUsage (Inner inner) = - innerUsage inner - - simplifyOp (Alloc size space) = - (,) <$> (Alloc <$> Engine.simplify size <*> pure space) <*> pure mempty - simplifyOp (Inner k) = do (k', hoisted) <- simplifyInnerOp k - return (Inner k', hoisted) - -bindPatternWithAllocations :: (MonadBinder m, - ExpDec lore ~ (), - Op (Lore m) ~ MemOp inner, - Allocator lore (PatAllocM lore)) => - Scope lore -> [VName] -> Exp lore - -> m (Pattern lore) -bindPatternWithAllocations types names e = do - (pat,prebnds) <- runPatAllocM (patternWithAllocations names e) types - mapM_ bindAllocStm prebnds - return pat - -data ExpHint = NoHint - | Hint IxFun Space +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | A generic transformation for adding memory allocations to a +-- Futhark program. Specialised by specific representations in +-- submodules. +module Futhark.Pass.ExplicitAllocations + ( explicitAllocationsGeneric, + explicitAllocationsInStmsGeneric, + ExpHint (..), + defaultExpHints, + Allocable, + Allocator (..), + AllocM, + AllocEnv (..), + SizeSubst (..), + allocInStms, + allocForArray, + simplifiable, + arraySizeInBytesExp, + mkLetNamesB', + mkLetNamesB'', + + -- * Module re-exports + + -- + -- These are highly likely to be needed by any downstream + -- users. + module Control.Monad.Reader, + module Futhark.MonadFreshNames, + module Futhark.Pass, + module Futhark.Tools, + ) +where + +import Control.Monad.RWS.Strict +import Control.Monad.Reader +import Control.Monad.State +import Control.Monad.Writer +import Data.List (foldl', partition, sort, zip4) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.IR.Mem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Optimise.Simplify.Engine (SimpleOps (..)) +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Optimise.Simplify.Lore (mkWiseBody) +import Futhark.Pass +import Futhark.Tools +import Futhark.Util (splitFromEnd, takeLast) + +data AllocStm + = SizeComputation VName (PrimExp VName) + | Allocation VName SubExp Space + | ArrayCopy VName VName + deriving (Eq, Ord, Show) + +bindAllocStm :: + (MonadBinder m, Op (Lore m) ~ MemOp inner) => + AllocStm -> + m () +bindAllocStm (SizeComputation name pe) = + letBindNames [name] =<< toExp (coerceIntPrimExp Int64 pe) +bindAllocStm (Allocation name size space) = + letBindNames [name] $ Op $ Alloc size space +bindAllocStm (ArrayCopy name src) = + letBindNames [name] $ BasicOp $ Copy src + +class + (MonadFreshNames m, HasScope lore m, Mem lore) => + Allocator lore m + where + addAllocStm :: AllocStm -> m () + askDefaultSpace :: m Space + + default addAllocStm :: + ( Allocable fromlore lore, + m ~ AllocM fromlore lore + ) => + AllocStm -> + m () + addAllocStm (SizeComputation name se) = + letBindNames [name] =<< toExp (coerceIntPrimExp Int64 se) + addAllocStm (Allocation name size space) = + letBindNames [name] $ Op $ allocOp size space + addAllocStm (ArrayCopy name src) = + letBindNames [name] $ BasicOp $ Copy src + + -- | The subexpression giving the number of elements we should + -- allocate space for. See 'ChunkMap' comment. + dimAllocationSize :: SubExp -> m SubExp + default dimAllocationSize :: + m ~ AllocM fromlore lore => + SubExp -> + m SubExp + dimAllocationSize (Var v) = + -- It is important to recurse here, as the substitution may itself + -- be a chunk size. + maybe (return $ Var v) dimAllocationSize =<< asks (M.lookup v . chunkMap) + dimAllocationSize size = + return size + + -- | Get those names that are known to be constants at run-time. + askConsts :: m (S.Set VName) + + expHints :: Exp lore -> m [ExpHint] + expHints = defaultExpHints + +allocateMemory :: + Allocator lore m => + String -> + SubExp -> + Space -> + m VName +allocateMemory desc size space = do + v <- newVName desc + addAllocStm $ Allocation v size space + return v + +computeSize :: + Allocator lore m => + String -> + PrimExp VName -> + m SubExp +computeSize desc se = do + v <- newVName desc + addAllocStm $ SizeComputation v se + return $ Var v + +type Allocable fromlore tolore = + ( PrettyLore fromlore, + PrettyLore tolore, + Mem tolore, + FParamInfo fromlore ~ DeclType, + LParamInfo fromlore ~ Type, + BranchType fromlore ~ ExtType, + RetType fromlore ~ DeclExtType, + BodyDec fromlore ~ (), + BodyDec tolore ~ (), + ExpDec tolore ~ (), + SizeSubst (Op tolore), + BinderOps tolore + ) + +-- | A mapping from chunk names to their maximum size. XXX FIXME +-- HACK: This is part of a hack to add loop-invariant allocations to +-- reduce kernels, because memory expansion does not use range +-- analysis yet (it should). +type ChunkMap = M.Map VName SubExp + +data AllocEnv fromlore tolore = AllocEnv + { chunkMap :: ChunkMap, + -- | Aggressively try to reuse memory in do-loops - + -- should be True inside kernels, False outside. + aggressiveReuse :: Bool, + -- | When allocating memory, put it in this memory space. + -- This is primarily used to ensure that group-wide + -- statements store their results in local memory. + allocSpace :: Space, + -- | The set of names that are known to be constants at + -- kernel compile time. + envConsts :: S.Set VName, + allocInOp :: Op fromlore -> AllocM fromlore tolore (Op tolore), + envExpHints :: Exp tolore -> AllocM fromlore tolore [ExpHint] + } + +-- | Monad for adding allocations to an entire program. +newtype AllocM fromlore tolore a + = AllocM (BinderT tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a) + deriving + ( Applicative, + Functor, + Monad, + MonadFreshNames, + HasScope tolore, + LocalScope tolore, + MonadReader (AllocEnv fromlore tolore) + ) + +instance + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + MonadBinder (AllocM fromlore tolore) + where + type Lore (AllocM fromlore tolore) = tolore + + mkExpDecM _ _ = return () + + mkLetNamesM names e = do + pat <- patternWithAllocations names e + return $ Let pat (defAux ()) e + + mkBodyM bnds res = return $ Body () bnds res + + addStms = AllocM . addStms + collectStms (AllocM m) = AllocM $ collectStms m + +instance + (Allocable fromlore tolore) => + Allocator tolore (AllocM fromlore tolore) + where + expHints e = do + f <- asks envExpHints + f e + askDefaultSpace = asks allocSpace + + askConsts = asks envConsts + +runAllocM :: + MonadFreshNames m => + (Op fromlore -> AllocM fromlore tolore (Op tolore)) -> + (Exp tolore -> AllocM fromlore tolore [ExpHint]) -> + AllocM fromlore tolore a -> + m a +runAllocM handleOp hints (AllocM m) = + fmap fst $ modifyNameSource $ runState $ runReaderT (runBinderT m mempty) env + where + env = + AllocEnv + { chunkMap = mempty, + aggressiveReuse = False, + allocSpace = DefaultSpace, + envConsts = mempty, + allocInOp = handleOp, + envExpHints = hints + } + +-- | Monad for adding allocations to a single pattern. +newtype PatAllocM lore a + = PatAllocM + ( RWS + (Scope lore) + [AllocStm] + VNameSource + a + ) + deriving + ( Applicative, + Functor, + Monad, + HasScope lore, + MonadWriter [AllocStm], + MonadFreshNames + ) + +instance Mem lore => Allocator lore (PatAllocM lore) where + addAllocStm = tell . pure + dimAllocationSize = return + askDefaultSpace = return DefaultSpace + askConsts = pure mempty + +runPatAllocM :: + MonadFreshNames m => + PatAllocM lore a -> + Scope lore -> + m (a, [AllocStm]) +runPatAllocM (PatAllocM m) mems = + modifyNameSource $ frob . runRWS m mems + where + frob (a, s, w) = ((a, w), s) + +elemSize :: Num a => Type -> a +elemSize = primByteSize . elemType + +arraySizeInBytesExp :: Type -> PrimExp VName +arraySizeInBytesExp t = + untyped $ foldl' (*) (elemSize t) $ map pe64 (arrayDims t) + +arraySizeInBytesExpM :: Allocator lore m => Type -> m (PrimExp VName) +arraySizeInBytesExpM t = do + dims <- mapM dimAllocationSize (arrayDims t) + let dim_prod_i64 = product $ map pe64 dims + elm_size_i64 = primByteSize $ elemType t + return $ untyped $ dim_prod_i64 * elm_size_i64 + +arraySizeInBytes :: Allocator lore m => Type -> m SubExp +arraySizeInBytes = computeSize "bytes" <=< arraySizeInBytesExpM + +-- | Allocate memory for a value of the given type. +allocForArray :: + Allocator lore m => + Type -> + Space -> + m VName +allocForArray t space = do + size <- arraySizeInBytes t + allocateMemory "mem" size space + +allocsForStm :: + (Allocator lore m, ExpDec lore ~ ()) => + [Ident] -> + [Ident] -> + Exp lore -> + m (Stm lore) +allocsForStm sizeidents validents e = do + rts <- expReturns e + hints <- expHints e + (ctxElems, valElems) <- allocsForPattern sizeidents validents rts hints + return $ Let (Pattern ctxElems valElems) (defAux ()) e + +patternWithAllocations :: + (Allocator lore m, ExpDec lore ~ ()) => + [VName] -> + Exp lore -> + m (Pattern lore) +patternWithAllocations names e = do + (ts', sizes) <- instantiateShapes' =<< expExtType e + let identForBindage name t = + pure $ Ident name t + vals <- sequence [identForBindage name t | (name, t) <- zip names ts'] + stmPattern <$> allocsForStm sizes vals e + +allocsForPattern :: + Allocator lore m => + [Ident] -> + [Ident] -> + [ExpReturns] -> + [ExpHint] -> + m + ( [PatElem lore], + [PatElem lore] + ) +allocsForPattern sizeidents validents rts hints = do + let sizes' = [PatElem size $ MemPrim int64 | size <- map identName sizeidents] + (vals, (exts, mems)) <- + runWriterT $ + forM (zip3 validents rts hints) $ \(ident, rt, hint) -> do + let shape = arrayShape $ identType ident + case rt of + MemPrim _ -> do + summary <- lift $ summaryForBindage (identType ident) hint + return $ PatElem (identName ident) summary + MemMem space -> + return $ + PatElem (identName ident) $ + MemMem space + MemArray bt _ u (Just (ReturnsInBlock mem extixfun)) -> do + (patels, ixfn) <- instantiateExtIxFun ident extixfun + tell (patels, []) + + return $ + PatElem (identName ident) $ + MemArray bt shape u $ + ArrayIn mem ixfn + MemArray _ extshape _ Nothing + | Just _ <- knownShape extshape -> do + summary <- lift $ summaryForBindage (identType ident) hint + return $ PatElem (identName ident) summary + MemArray bt _ u (Just (ReturnsNewBlock space _ extixfn)) -> do + -- treat existential index function first + (patels, ixfn) <- instantiateExtIxFun ident extixfn + tell (patels, []) + + memid <- lift $ mkMemIdent ident space + tell ([], [PatElem (identName memid) $ MemMem space]) + return $ + PatElem (identName ident) $ + MemArray bt shape u $ + ArrayIn (identName memid) ixfn + _ -> error "Impossible case reached in allocsForPattern!" + + return + ( sizes' <> exts <> mems, + vals + ) + where + knownShape = mapM known . shapeDims + known (Free v) = Just v + known Ext {} = Nothing + + mkMemIdent :: (MonadFreshNames m) => Ident -> Space -> m Ident + mkMemIdent ident space = do + let memname = baseString (identName ident) <> "_mem" + newIdent memname $ Mem space + + instantiateExtIxFun :: + MonadFreshNames m => + Ident -> + ExtIxFun -> + m ([PatElemT (MemInfo d u ret)], IxFun) + instantiateExtIxFun idd ext_ixfn = do + let isAndPtps = + S.toList $ + foldMap onlyExts $ + foldMap (leafExpTypes . untyped) ext_ixfn + + -- Find the existentials that reuse the sizeidents, and + -- those that need new pattern elements. Assumes that the + -- Exts form a contiguous interval of integers. + let (size_exts, new_exts) = + span ((< length sizeidents) . fst) $ sort isAndPtps + (new_substs, patels) <- + fmap unzip $ + forM new_exts $ \(i, t) -> do + v <- newVName $ baseString (identName idd) <> "_ixfn" + return + ( (Ext i, LeafExp (Free v) t), + PatElem v $ MemPrim t + ) + let size_substs = + zipWith + ( \(i, t) ident -> + (Ext i, LeafExp (Free (identName ident)) t) + ) + size_exts + sizeidents + substs = M.fromList $ new_substs <> size_substs + ixfn <- instantiateIxFun $ IxFun.substituteInIxFun (fmap isInt64 substs) ext_ixfn + + return (patels, ixfn) + +onlyExts :: (Ext a, PrimType) -> S.Set (Int, PrimType) +onlyExts (Free _, _) = S.empty +onlyExts (Ext i, t) = S.singleton (i, t) + +instantiateIxFun :: Monad m => ExtIxFun -> m IxFun +instantiateIxFun = traverse $ traverse inst + where + inst Ext {} = error "instantiateIxFun: not yet" + inst (Free x) = return x + +summaryForBindage :: + Allocator lore m => + Type -> + ExpHint -> + m (MemBound NoUniqueness) +summaryForBindage (Prim bt) _ = + return $ MemPrim bt +summaryForBindage (Mem space) _ = + return $ MemMem space +summaryForBindage t@(Array bt shape u) NoHint = do + m <- allocForArray t =<< askDefaultSpace + return $ directIxFun bt shape u m t +summaryForBindage t (Hint ixfun space) = do + let bt = elemType t + bytes <- + computeSize "bytes" $ + untyped $ + product + [ product $ IxFun.base ixfun, + primByteSize (elemType t) + ] + m <- allocateMemory "mem" bytes space + return $ MemArray bt (arrayShape t) NoUniqueness $ ArrayIn m ixfun + +lookupMemSpace :: (HasScope lore m, Monad m) => VName -> m Space +lookupMemSpace v = do + t <- lookupType v + case t of + Mem space -> return space + _ -> error $ "lookupMemSpace: " ++ pretty v ++ " is not a memory block." + +directIxFun :: PrimType -> Shape -> u -> VName -> Type -> MemBound u +directIxFun bt shape u mem t = + let ixf = IxFun.iota $ map pe64 $ arrayDims t + in MemArray bt shape u $ ArrayIn mem ixf + +allocInFParams :: + (Allocable fromlore tolore) => + [(FParam fromlore, Space)] -> + ([FParam tolore] -> AllocM fromlore tolore a) -> + AllocM fromlore tolore a +allocInFParams params m = do + (valparams, (ctxparams, memparams)) <- + runWriterT $ mapM (uncurry allocInFParam) params + let params' = ctxparams <> memparams <> valparams + summary = scopeOfFParams params' + localScope summary $ m params' + +allocInFParam :: + (Allocable fromlore tolore) => + FParam fromlore -> + Space -> + WriterT + ([FParam tolore], [FParam tolore]) + (AllocM fromlore tolore) + (FParam tolore) +allocInFParam param pspace = + case paramDeclType param of + Array bt shape u -> do + let memname = baseString (paramName param) <> "_mem" + ixfun = IxFun.iota $ map pe64 $ shapeDims shape + mem <- lift $ newVName memname + tell ([], [Param mem $ MemMem pspace]) + return param {paramDec = MemArray bt shape u $ ArrayIn mem ixfun} + Prim bt -> + return param {paramDec = MemPrim bt} + Mem space -> + return param {paramDec = MemMem space} + +allocInMergeParams :: + ( Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore) + ) => + [(FParam fromlore, SubExp)] -> + ( [FParam tolore] -> + [FParam tolore] -> + ([SubExp] -> AllocM fromlore tolore ([SubExp], [SubExp])) -> + AllocM fromlore tolore a + ) -> + AllocM fromlore tolore a +allocInMergeParams merge m = do + ((valparams, handle_loop_subexps), (ctx_params, mem_params)) <- + runWriterT $ unzip <$> mapM allocInMergeParam merge + let mergeparams' = ctx_params <> mem_params <> valparams + summary = scopeOfFParams mergeparams' + + mk_loop_res ses = do + (valargs, (ctxargs, memargs)) <- + runWriterT $ zipWithM ($) handle_loop_subexps ses + return (ctxargs <> memargs, valargs) + + localScope summary $ m (ctx_params <> mem_params) valparams mk_loop_res + where + allocInMergeParam :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Param DeclType, SubExp) -> + WriterT + ([FParam tolore], [FParam tolore]) + (AllocM fromlore tolore) + (FParam tolore, SubExp -> WriterT ([SubExp], [SubExp]) (AllocM fromlore tolore) SubExp) + allocInMergeParam (mergeparam, Var v) + | Array bt shape u <- paramDeclType mergeparam = do + (mem', _) <- lift $ lookupArraySummary v + mem_space <- lift $ lookupMemSpace mem' + + (_, ext_ixfun, substs, _) <- lift $ existentializeArray mem_space v + + (ctx_params, param_ixfun_substs) <- + unzip + <$> mapM + ( \_ -> do + vname <- lift $ newVName "ctx_param_ext" + return + ( Param vname $ MemPrim int64, + fmap Free $ pe64 $ Var vname + ) + ) + substs + + tell (ctx_params, []) + + param_ixfun <- + instantiateIxFun $ + IxFun.substituteInIxFun + (M.fromList $ zip (fmap Ext [0 ..]) param_ixfun_substs) + ext_ixfun + + mem_name <- newVName "mem_param" + tell ([], [Param mem_name $ MemMem mem_space]) + + return + ( mergeparam {paramDec = MemArray bt shape u $ ArrayIn mem_name param_ixfun}, + ensureArrayIn mem_space + ) + allocInMergeParam (mergeparam, _) = doDefault mergeparam =<< lift askDefaultSpace + + doDefault mergeparam space = do + mergeparam' <- allocInFParam mergeparam space + return (mergeparam', linearFuncallArg (paramType mergeparam) space) + +-- Returns the existentialized index function, the list of substituted values and the memory location. +existentializeArray :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Space -> + VName -> + AllocM fromlore tolore (SubExp, ExtIxFun, [TPrimExp Int64 VName], VName) +existentializeArray space v = do + (mem', ixfun) <- lookupArraySummary v + sp <- lookupMemSpace mem' + + let (ext_ixfun', substs') = runState (IxFun.existentialize ixfun) [] + + case (ext_ixfun', sp == space) of + (Just x, True) -> return (Var v, x, substs', mem') + _ -> do + (mem, subexp) <- allocLinearArray space (baseString v) v + ixfun' <- fromJust <$> subExpIxFun subexp + let (ext_ixfun, substs) = runState (IxFun.existentialize ixfun') [] + return (subexp, fromJust ext_ixfun, substs, mem) + +ensureArrayIn :: + ( Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore) + ) => + Space -> + SubExp -> + WriterT ([SubExp], [SubExp]) (AllocM fromlore tolore) SubExp +ensureArrayIn _ (Constant v) = + error $ "ensureArrayIn: " ++ pretty v ++ " cannot be an array." +ensureArrayIn space (Var v) = do + (sub_exp, _, substs, mem) <- lift $ existentializeArray space v + (ctx_vals, _) <- + unzip + <$> mapM + ( \s -> do + vname <- lift $ letExp "ctx_val" =<< toExp s + return (Var vname, fmap Free $ primExpFromSubExp int64 $ Var vname) + ) + substs + + tell (ctx_vals, [Var mem]) + + return sub_exp + +ensureDirectArray :: + ( Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore) + ) => + Maybe Space -> + VName -> + AllocM fromlore tolore (VName, SubExp) +ensureDirectArray space_ok v = do + (mem, ixfun) <- lookupArraySummary v + mem_space <- lookupMemSpace mem + default_space <- askDefaultSpace + if IxFun.isDirect ixfun && maybe True (== mem_space) space_ok + then return (mem, Var v) + else needCopy (fromMaybe default_space space_ok) + where + needCopy space = + -- We need to do a new allocation, copy 'v', and make a new + -- binding for the size of the memory block. + allocLinearArray space (baseString v) v + +allocLinearArray :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Space -> + String -> + VName -> + AllocM fromlore tolore (VName, SubExp) +allocLinearArray space s v = do + t <- lookupType v + mem <- allocForArray t space + v' <- newIdent (s ++ "_linear") t + let ixfun = directIxFun (elemType t) (arrayShape t) NoUniqueness mem t + let pat = Pattern [] [PatElem (identName v') ixfun] + addStm $ Let pat (defAux ()) $ BasicOp $ Copy v + return (mem, Var $ identName v') + +funcallArgs :: + ( Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore) + ) => + [(SubExp, Diet)] -> + AllocM fromlore tolore [(SubExp, Diet)] +funcallArgs args = do + (valargs, (ctx_args, mem_and_size_args)) <- runWriterT $ + forM args $ \(arg, d) -> do + t <- lift $ subExpType arg + space <- lift askDefaultSpace + arg' <- linearFuncallArg t space arg + return (arg', d) + return $ map (,Observe) (ctx_args <> mem_and_size_args) <> valargs + +linearFuncallArg :: + ( Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore) + ) => + Type -> + Space -> + SubExp -> + WriterT ([SubExp], [SubExp]) (AllocM fromlore tolore) SubExp +linearFuncallArg Array {} space (Var v) = do + (mem, arg') <- lift $ ensureDirectArray (Just space) v + tell ([], [Var mem]) + return arg' +linearFuncallArg _ _ arg = + return arg + +explicitAllocationsGeneric :: + ( Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore) + ) => + (Op fromlore -> AllocM fromlore tolore (Op tolore)) -> + (Exp tolore -> AllocM fromlore tolore [ExpHint]) -> + Pass fromlore tolore +explicitAllocationsGeneric handleOp hints = + Pass "explicit allocations" "Transform program to explicit memory representation" $ + intraproceduralTransformationWithConsts onStms allocInFun + where + onStms stms = runAllocM handleOp hints $ allocInStms stms pure + + allocInFun consts (FunDef entry attrs fname rettype params fbody) = + runAllocM handleOp hints $ + inScopeOf consts $ + allocInFParams (zip params $ repeat DefaultSpace) $ \params' -> do + fbody' <- + insertStmsM $ + allocInFunBody + (map (const $ Just DefaultSpace) rettype) + fbody + return $ FunDef entry attrs fname (memoryInDeclExtType rettype) params' fbody' + +explicitAllocationsInStmsGeneric :: + ( MonadFreshNames m, + HasScope tolore m, + Allocable fromlore tolore + ) => + (Op fromlore -> AllocM fromlore tolore (Op tolore)) -> + (Exp tolore -> AllocM fromlore tolore [ExpHint]) -> + Stms fromlore -> + m (Stms tolore) +explicitAllocationsInStmsGeneric handleOp hints stms = do + scope <- askScope + runAllocM handleOp hints $ localScope scope $ allocInStms stms return + +memoryInDeclExtType :: [DeclExtType] -> [FunReturns] +memoryInDeclExtType ts = evalState (mapM addMem ts) $ startOfFreeIDRange ts + where + addMem (Prim t) = return $ MemPrim t + addMem Mem {} = error "memoryInDeclExtType: too much memory" + addMem (Array bt shape u) = do + i <- get <* modify (+ 1) + return $ + MemArray bt shape u $ + ReturnsNewBlock DefaultSpace i $ + IxFun.iota $ map convert $ shapeDims shape + + convert (Ext i) = le64 $ Ext i + convert (Free v) = Free <$> pe64 v + +startOfFreeIDRange :: [TypeBase ExtShape u] -> Int +startOfFreeIDRange = S.size . shapeContext + +bodyReturnMemCtx :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + SubExp -> + AllocM fromlore tolore [SubExp] +bodyReturnMemCtx Constant {} = + return [] +bodyReturnMemCtx (Var v) = do + info <- lookupMemInfo v + case info of + MemPrim {} -> return [] + MemMem {} -> return [] -- should not happen + MemArray _ _ _ (ArrayIn mem _) -> return [Var mem] + +allocInFunBody :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + [Maybe Space] -> + Body fromlore -> + AllocM fromlore tolore (Body tolore) +allocInFunBody space_oks (Body _ bnds res) = + allocInStms bnds $ \bnds' -> do + (res'', allocs) <- collectStms $ do + res' <- zipWithM ensureDirect space_oks' res + let (ctx_res, val_res) = splitFromEnd num_vals res' + mem_ctx_res <- concat <$> mapM bodyReturnMemCtx val_res + return $ ctx_res <> mem_ctx_res <> val_res + return $ Body () (bnds' <> allocs) res'' + where + num_vals = length space_oks + space_oks' = replicate (length res - num_vals) Nothing ++ space_oks + +ensureDirect :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Maybe Space -> + SubExp -> + AllocM fromlore tolore SubExp +ensureDirect _ se@Constant {} = return se +ensureDirect space_ok (Var v) = do + bt <- primType <$> lookupType v + if bt + then return $ Var v + else do + (_, v') <- ensureDirectArray space_ok v + return v' + +allocInStms :: + (Allocable fromlore tolore) => + Stms fromlore -> + (Stms tolore -> AllocM fromlore tolore a) -> + AllocM fromlore tolore a +allocInStms origstms m = allocInStms' (stmsToList origstms) mempty + where + allocInStms' [] stms' = + m stms' + allocInStms' (x : xs) stms' = do + allocstms <- allocInStm' x + localScope (scopeOf allocstms) $ do + let stms_substs = foldMap sizeSubst allocstms + stms_consts = foldMap stmConsts allocstms + f env = + env + { chunkMap = stms_substs <> chunkMap env, + envConsts = stms_consts <> envConsts env + } + local f $ allocInStms' xs (stms' <> allocstms) + allocInStm' stm = + collectStms_ $ auxing (stmAux stm) $ allocInStm stm + +allocInStm :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Stm fromlore -> + AllocM fromlore tolore () +allocInStm (Let (Pattern sizeElems valElems) _ e) = do + e' <- allocInExp e + let sizeidents = map patElemIdent sizeElems + validents = map patElemIdent valElems + bnd <- allocsForStm sizeidents validents e' + addStm bnd + +allocInExp :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Exp fromlore -> + AllocM fromlore tolore (Exp tolore) +allocInExp (DoLoop ctx val form (Body () bodybnds bodyres)) = + allocInMergeParams ctx $ \_ ctxparams' _ -> + allocInMergeParams val $ + \new_ctx_params valparams' mk_loop_val -> do + form' <- allocInLoopForm form + localScope (scopeOf form') $ do + (valinit_ctx, valinit') <- mk_loop_val valinit + body' <- insertStmsM $ + allocInStms bodybnds $ \bodybnds' -> do + ((val_ses, valres'), val_retbnds) <- collectStms $ mk_loop_val valres + return $ Body () (bodybnds' <> val_retbnds) (ctxres ++ val_ses ++ valres') + return $ + DoLoop + (zip (ctxparams' ++ new_ctx_params) (ctxinit ++ valinit_ctx)) + (zip valparams' valinit') + form' + body' + where + (_ctxparams, ctxinit) = unzip ctx + (_valparams, valinit) = unzip val + (ctxres, valres) = splitAt (length ctx) bodyres +allocInExp (Apply fname args rettype loc) = do + args' <- funcallArgs args + return $ Apply fname args' (memoryInDeclExtType rettype) loc +allocInExp (If cond tbranch0 fbranch0 (IfDec rets ifsort)) = do + let num_rets = length rets + -- switch to the explicit-mem rep, but do nothing about results + (tbranch, tm_ixfs) <- allocInIfBody num_rets tbranch0 + (fbranch, fm_ixfs) <- allocInIfBody num_rets fbranch0 + tspaces <- mkSpaceOks num_rets tbranch + fspaces <- mkSpaceOks num_rets fbranch + -- try to generalize (antiunify) the index functions of the then and else bodies + let sp_substs = zipWith generalize (zip tspaces tm_ixfs) (zip fspaces fm_ixfs) + (spaces, subs) = unzip sp_substs + tsubs = map (selectSub fst) subs + fsubs = map (selectSub snd) subs + (tbranch', trets) <- addResCtxInIfBody rets tbranch spaces tsubs + (fbranch', frets) <- addResCtxInIfBody rets fbranch spaces fsubs + if frets /= trets + then error "In allocInExp, IF case: antiunification of then/else produce different ExtInFn!" + else do + -- above is a sanity check; implementation continues on else branch + let res_then = bodyResult tbranch' + res_else = bodyResult fbranch' + size_ext = length res_then - length trets + (ind_ses0, r_then_else) = + partition (\(r_then, r_else, _) -> r_then == r_else) $ + zip3 res_then res_else [0 .. size_ext - 1] + (r_then_ext, r_else_ext, _) = unzip3 r_then_else + ind_ses = + zipWith + (\(se, _, i) k -> (i - k, se)) + ind_ses0 + [0 .. length ind_ses0 - 1] + rets'' = foldl (\acc (i, se) -> fixExt i se acc) trets ind_ses + tbranch'' = tbranch' {bodyResult = r_then_ext ++ drop size_ext res_then} + fbranch'' = fbranch' {bodyResult = r_else_ext ++ drop size_ext res_else} + res_if_expr = If cond tbranch'' fbranch'' $ IfDec rets'' ifsort + return res_if_expr + where + generalize :: + (Maybe Space, Maybe IxFun) -> + (Maybe Space, Maybe IxFun) -> + (Maybe Space, Maybe (ExtIxFun, [(TPrimExp Int64 VName, TPrimExp Int64 VName)])) + generalize (Just sp1, Just ixf1) (Just sp2, Just ixf2) = + if sp1 /= sp2 + then (Just sp1, Nothing) + else case IxFun.leastGeneralGeneralization (fmap untyped ixf1) (fmap untyped ixf2) of + Just (ixf, m) -> + ( Just sp1, + Just + ( fmap TPrimExp ixf, + zip (map (TPrimExp . fst) m) (map (TPrimExp . snd) m) + ) + ) + Nothing -> (Just sp1, Nothing) + generalize (mbsp1, _) _ = (mbsp1, Nothing) + + selectSub :: + ((a, a) -> a) -> + Maybe (ExtIxFun, [(a, a)]) -> + Maybe (ExtIxFun, [a]) + selectSub f (Just (ixfn, m)) = Just (ixfn, map f m) + selectSub _ Nothing = Nothing + allocInIfBody :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Int -> + Body fromlore -> + AllocM fromlore tolore (Body tolore, [Maybe IxFun]) + allocInIfBody num_vals (Body _ bnds res) = + allocInStms bnds $ \bnds' -> do + let (_, val_res) = splitFromEnd num_vals res + mem_ixfs <- mapM subExpIxFun val_res + return (Body () bnds' res, mem_ixfs) +allocInExp e = mapExpM alloc e + where + alloc = + identityMapper + { mapOnBody = error "Unhandled Body in ExplicitAllocations", + mapOnRetType = error "Unhandled RetType in ExplicitAllocations", + mapOnBranchType = error "Unhandled BranchType in ExplicitAllocations", + mapOnFParam = error "Unhandled FParam in ExplicitAllocations", + mapOnLParam = error "Unhandled LParam in ExplicitAllocations", + mapOnOp = \op -> do + handle <- asks allocInOp + handle op + } + +subExpIxFun :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + SubExp -> + AllocM fromlore tolore (Maybe IxFun) +subExpIxFun Constant {} = return Nothing +subExpIxFun (Var v) = do + info <- lookupMemInfo v + case info of + MemArray _ptp _shp _u (ArrayIn _ ixf) -> return $ Just ixf + _ -> return Nothing + +addResCtxInIfBody :: + (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + [ExtType] -> + Body tolore -> + [Maybe Space] -> + [Maybe (ExtIxFun, [TPrimExp Int64 VName])] -> + AllocM fromlore tolore (Body tolore, [BodyReturns]) +addResCtxInIfBody ifrets (Body _ bnds res) spaces substs = do + let num_vals = length ifrets + (ctx_res, val_res) = splitFromEnd num_vals res + ((res', bodyrets'), all_body_stms) <- collectStms $ do + mapM_ addStm bnds + (val_res', ext_ses_res, mem_ctx_res, bodyrets, total_existentials) <- + foldM helper ([], [], [], [], length ctx_res) (zip4 ifrets val_res substs spaces) + return + ( ctx_res <> ext_ses_res <> mem_ctx_res <> val_res', + -- We need to adjust the ReturnsNewBlock existentials, because they + -- should always be numbered _after_ all other existentials in the + -- return values. + reverse $ fst $ foldl adjustNewBlockExistential ([], total_existentials) bodyrets + ) + body' <- mkBodyM all_body_stms res' + return (body', bodyrets') + where + helper (res_acc, ext_acc, ctx_acc, br_acc, k) (ifr, r, mbixfsub, sp) = + case mbixfsub of + Nothing -> do + -- does NOT generalize/antiunify; ensure direct + r' <- ensureDirect sp r + mem_ctx_r <- bodyReturnMemCtx r' + let body_ret = inspect ifr sp + return + ( res_acc ++ [r'], + ext_acc, + ctx_acc ++ mem_ctx_r, + br_acc ++ [body_ret], + k + ) + Just (ixfn, m) -> do + -- generalizes + let i = length m + ext_ses <- mapM (toSubExp "ixfn_exist") m + mem_ctx_r <- bodyReturnMemCtx r + let sp' = fromMaybe DefaultSpace sp + ixfn' = fmap (adjustExtPE k) ixfn + exttp = case ifr of + Array pt shp' u -> + MemArray pt shp' u $ + ReturnsNewBlock sp' 0 ixfn' + _ -> error "Impossible case reached in addResCtxInIfBody" + return + ( res_acc ++ [r], + ext_acc ++ ext_ses, + ctx_acc ++ mem_ctx_r, + br_acc ++ [exttp], + k + i + ) + + adjustNewBlockExistential :: ([BodyReturns], Int) -> BodyReturns -> ([BodyReturns], Int) + adjustNewBlockExistential (acc, k) (MemArray pt shp u (ReturnsNewBlock space _ ixfun)) = + (MemArray pt shp u (ReturnsNewBlock space k ixfun) : acc, k + 1) + adjustNewBlockExistential (acc, k) x = (x : acc, k) + + inspect (Array pt shape u) space = + let space' = fromMaybe DefaultSpace space + bodyret = + MemArray pt shape u $ + ReturnsNewBlock space' 0 $ + IxFun.iota $ map convert $ shapeDims shape + in bodyret + inspect (Prim pt) _ = MemPrim pt + inspect (Mem space) _ = MemMem space + + convert (Ext i) = le64 (Ext i) + convert (Free v) = Free <$> pe64 v + + adjustExtV :: Int -> Ext VName -> Ext VName + adjustExtV _ (Free v) = Free v + adjustExtV k (Ext i) = Ext (k + i) + + adjustExtPE :: Int -> TPrimExp t (Ext VName) -> TPrimExp t (Ext VName) + adjustExtPE k = fmap (adjustExtV k) + +mkSpaceOks :: + (Mem tolore, LocalScope tolore m) => + Int -> + Body tolore -> + m [Maybe Space] +mkSpaceOks num_vals (Body _ stms res) = + inScopeOf stms $ + mapM mkSpaceOK $ takeLast num_vals res + where + mkSpaceOK (Var v) = do + v_info <- lookupMemInfo v + case v_info of + MemArray _ _ _ (ArrayIn mem _) -> do + mem_info <- lookupMemInfo mem + case mem_info of + MemMem space -> return $ Just space + _ -> return Nothing + _ -> return Nothing + mkSpaceOK _ = return Nothing + +allocInLoopForm :: + ( Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore) + ) => + LoopForm fromlore -> + AllocM fromlore tolore (LoopForm tolore) +allocInLoopForm (WhileLoop v) = return $ WhileLoop v +allocInLoopForm (ForLoop i it n loopvars) = + ForLoop i it n <$> mapM allocInLoopVar loopvars + where + allocInLoopVar (p, a) = do + (mem, ixfun) <- lookupArraySummary a + case paramType p of + Array bt shape u -> do + dims <- map pe64 . arrayDims <$> lookupType a + let ixfun' = + IxFun.slice ixfun $ + fullSliceNum dims [DimFix $ le64 i] + return (p {paramDec = MemArray bt shape u $ ArrayIn mem ixfun'}, a) + Prim bt -> + return (p {paramDec = MemPrim bt}, a) + Mem space -> + return (p {paramDec = MemMem space}, a) + +class SizeSubst op where + opSizeSubst :: PatternT dec -> op -> ChunkMap + opIsConst :: op -> Bool + opIsConst = const False + +instance SizeSubst () where + opSizeSubst _ _ = mempty + +instance SizeSubst op => SizeSubst (MemOp op) where + opSizeSubst pat (Inner op) = opSizeSubst pat op + opSizeSubst _ _ = mempty + + opIsConst (Inner op) = opIsConst op + opIsConst _ = False + +sizeSubst :: SizeSubst (Op lore) => Stm lore -> ChunkMap +sizeSubst (Let pat _ (Op op)) = opSizeSubst pat op +sizeSubst _ = mempty + +stmConsts :: SizeSubst (Op lore) => Stm lore -> S.Set VName +stmConsts (Let pat _ (Op op)) + | opIsConst op = S.fromList $ patternNames pat +stmConsts _ = mempty + +mkLetNamesB' :: + ( Op (Lore m) ~ MemOp inner, + MonadBinder m, + ExpDec (Lore m) ~ (), + Allocator (Lore m) (PatAllocM (Lore m)) + ) => + ExpDec (Lore m) -> + [VName] -> + Exp (Lore m) -> + m (Stm (Lore m)) +mkLetNamesB' dec names e = do + scope <- askScope + pat <- bindPatternWithAllocations scope names e + return $ Let pat (defAux dec) e + +mkLetNamesB'' :: + ( Op (Lore m) ~ MemOp inner, + ExpDec lore ~ (), + HasScope (Engine.Wise lore) m, + Allocator lore (PatAllocM lore), + MonadBinder m, + Engine.CanBeWise (Op lore) + ) => + [VName] -> + Exp (Engine.Wise lore) -> + m (Stm (Engine.Wise lore)) +mkLetNamesB'' names e = do + scope <- Engine.removeScopeWisdom <$> askScope + (pat, prestms) <- runPatAllocM (patternWithAllocations names $ Engine.removeExpWisdom e) scope + mapM_ bindAllocStm prestms + let pat' = Engine.addWisdomToPattern pat e + dec = Engine.mkWiseExpDec pat' () e + return $ Let pat' (defAux dec) e + +simplifiable :: + ( Engine.SimplifiableLore lore, + ExpDec lore ~ (), + BodyDec lore ~ (), + Op lore ~ MemOp inner, + Allocator lore (PatAllocM lore) + ) => + (Engine.OpWithWisdom inner -> UT.UsageTable) -> + (inner -> Engine.SimpleM lore (Engine.OpWithWisdom inner, Stms (Engine.Wise lore))) -> + SimpleOps lore +simplifiable innerUsage simplifyInnerOp = + SimpleOps mkExpDecS' mkBodyS' protectOp opUsage simplifyOp + where + mkExpDecS' _ pat e = + return $ Engine.mkWiseExpDec pat () e + + mkBodyS' _ bnds res = return $ mkWiseBody () bnds res + + protectOp taken pat (Alloc size space) = Just $ do + tbody <- resultBodyM [size] + fbody <- resultBodyM [intConst Int64 0] + size' <- + letSubExp "hoisted_alloc_size" $ + If taken tbody fbody $ IfDec [MemPrim int64] IfFallback + letBind pat $ Op $ Alloc size' space + protectOp _ _ _ = Nothing + + opUsage (Alloc (Var size) _) = + UT.sizeUsage size + opUsage (Alloc _ _) = + mempty + opUsage (Inner inner) = + innerUsage inner + + simplifyOp (Alloc size space) = + (,) <$> (Alloc <$> Engine.simplify size <*> pure space) <*> pure mempty + simplifyOp (Inner k) = do + (k', hoisted) <- simplifyInnerOp k + return (Inner k', hoisted) + +bindPatternWithAllocations :: + ( MonadBinder m, + ExpDec lore ~ (), + Op (Lore m) ~ MemOp inner, + Allocator lore (PatAllocM lore) + ) => + Scope lore -> + [VName] -> + Exp lore -> + m (Pattern lore) +bindPatternWithAllocations types names e = do + (pat, prebnds) <- runPatAllocM (patternWithAllocations names e) types + mapM_ bindAllocStm prebnds + return pat + +data ExpHint + = NoHint + | Hint IxFun Space defaultExpHints :: (Monad m, ASTLore lore) => Exp lore -> m [ExpHint] defaultExpHints e = return $ replicate (expExtTypeSize e) NoHint diff --git a/src/Futhark/Pass/ExplicitAllocations/Kernels.hs b/src/Futhark/Pass/ExplicitAllocations/Kernels.hs --- a/src/Futhark/Pass/ExplicitAllocations/Kernels.hs +++ b/src/Futhark/Pass/ExplicitAllocations/Kernels.hs @@ -2,19 +2,19 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-} {-# OPTIONS_GHC -fno-warn-orphans #-} + -- | Facilities for converting a 'Kernels' program to 'KernelsMem'. module Futhark.Pass.ExplicitAllocations.Kernels - ( explicitAllocations - , explicitAllocationsInStms - ) + ( explicitAllocations, + explicitAllocationsInStms, + ) where import qualified Data.Map as M import qualified Data.Set as S - -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.IR.KernelsMem import Futhark.IR.Kernels +import Futhark.IR.KernelsMem +import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.Pass.ExplicitAllocations import Futhark.Pass.ExplicitAllocations.SegOp @@ -23,43 +23,56 @@ M.singleton (patElemName size) elems_per_thread opSizeSubst _ _ = mempty - opIsConst (SizeOp GetSize{}) = True - opIsConst (SizeOp GetSizeMax{}) = True + opIsConst (SizeOp GetSize {}) = True + opIsConst (SizeOp GetSizeMax {}) = True opIsConst _ = False instance SizeSubst (SegOp lvl lore) where opSizeSubst _ _ = mempty allocAtLevel :: SegLevel -> AllocM fromlore tlore a -> AllocM fromlore tlore a -allocAtLevel lvl = local $ \env -> env { allocSpace = space - , aggressiveReuse = True - } - where space = case lvl of SegThread{} -> DefaultSpace - SegGroup{} -> Space "local" +allocAtLevel lvl = local $ \env -> + env + { allocSpace = space, + aggressiveReuse = True + } + where + space = case lvl of + SegThread {} -> DefaultSpace + SegGroup {} -> Space "local" -handleSegOp :: SegOp SegLevel Kernels - -> AllocM Kernels KernelsMem (SegOp SegLevel KernelsMem) +handleSegOp :: + SegOp SegLevel Kernels -> + AllocM Kernels KernelsMem (SegOp SegLevel KernelsMem) handleSegOp op = do - num_threads <- letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int32 OverflowUndef) - (unCount (segNumGroups lvl)) (unCount (segGroupSize lvl)) + num_threads <- + letSubExp "num_threads" $ + BasicOp $ + BinOp + (Mul Int64 OverflowUndef) + (unCount (segNumGroups lvl)) + (unCount (segGroupSize lvl)) allocAtLevel lvl $ mapSegOpM (mapper num_threads) op - where scope = scopeOfSegSpace $ segSpace op - lvl = segLevel op - mapper num_threads = - identitySegOpMapper - { mapOnSegOpBody = - localScope scope . local f . allocInKernelBody - , mapOnSegOpLambda = - local inThread . - allocInBinOpLambda num_threads (segSpace op) - } - f = case segLevel op of SegThread{} -> inThread - SegGroup{} -> inGroup - inThread env = env { envExpHints = inThreadExpHints } - inGroup env = env { envExpHints = inGroupExpHints } + where + scope = scopeOfSegSpace $ segSpace op + lvl = segLevel op + mapper num_threads = + identitySegOpMapper + { mapOnSegOpBody = + localScope scope . local f . allocInKernelBody, + mapOnSegOpLambda = + local inThread + . allocInBinOpLambda num_threads (segSpace op) + } + f = case segLevel op of + SegThread {} -> inThread + SegGroup {} -> inGroup + inThread env = env {envExpHints = inThreadExpHints} + inGroup env = env {envExpHints = inGroupExpHints} -handleHostOp :: HostOp Kernels (SOAC Kernels) - -> AllocM Kernels KernelsMem (MemOp (HostOp KernelsMem ())) +handleHostOp :: + HostOp Kernels (SOAC Kernels) -> + AllocM Kernels KernelsMem (MemOp (HostOp KernelsMem ())) handleHostOp (SizeOp op) = return $ Inner $ SizeOp op handleHostOp (OtherOp op) = @@ -72,104 +85,111 @@ dims <- arrayDims <$> lookupType v let perm_inv = rearrangeInverse perm dims' = rearrangeShape perm dims - ixfun = IxFun.permute (IxFun.iota $ map (primExpFromSubExp int32) dims') - perm_inv + ixfun = IxFun.permute (IxFun.iota $ map pe64 dims') perm_inv return [Hint ixfun DefaultSpace] - -kernelExpHints (Op (Inner (SegOp (SegMap lvl@SegThread{} space ts body)))) = +kernelExpHints (Op (Inner (SegOp (SegMap lvl@SegThread {} space ts body)))) = zipWithM (mapResultHint lvl space) ts $ kernelBodyResult body - -kernelExpHints (Op (Inner (SegOp (SegRed lvl@SegThread{} space reds ts body)))) = +kernelExpHints (Op (Inner (SegOp (SegRed lvl@SegThread {} space reds ts body)))) = (map (const NoHint) red_res <>) <$> zipWithM (mapResultHint lvl space) (drop num_reds ts) map_res - where num_reds = segBinOpResults reds - (red_res, map_res) = splitAt num_reds $ kernelBodyResult body - + where + num_reds = segBinOpResults reds + (red_res, map_res) = splitAt num_reds $ kernelBodyResult body kernelExpHints e = return $ replicate (expExtTypeSize e) NoHint -mapResultHint :: Allocator lore m => - SegLevel -> SegSpace -> Type -> KernelResult -> m ExpHint +mapResultHint :: + Allocator lore m => + SegLevel -> + SegSpace -> + Type -> + KernelResult -> + m ExpHint mapResultHint lvl space = hint - where num_threads = primExpFromSubExp int32 (unCount $ segNumGroups lvl) * - primExpFromSubExp int32 (unCount $ segGroupSize lvl) - - -- Heuristic: do not rearrange for returned arrays that are - -- sufficiently small. - coalesceReturnOfShape _ [] = False - coalesceReturnOfShape bs [Constant (IntValue (Int32Value d))] = bs * d > 4 - coalesceReturnOfShape _ _ = True - - hint t Returns{} - | coalesceReturnOfShape (primByteSize (elemType t)) $ arrayDims t = do - let space_dims = segSpaceDims space - t_dims <- mapM dimAllocationSize $ arrayDims t - return $ Hint (innermost space_dims t_dims) DefaultSpace - - hint t (ConcatReturns SplitStrided{} w _ _) = do - t_dims <- mapM dimAllocationSize $ arrayDims t - return $ Hint (innermost [w] t_dims) DefaultSpace + where + num_threads = + pe64 (unCount $ segNumGroups lvl) * pe64 (unCount $ segGroupSize lvl) - hint Prim{} (ConcatReturns SplitContiguous w elems_per_thread _) = do - let ixfun_base = IxFun.iota [num_threads, primExpFromSubExp int32 elems_per_thread] - ixfun_tr = IxFun.permute ixfun_base [1,0] - ixfun = IxFun.reshape ixfun_tr $ map (DimNew . primExpFromSubExp int32) [w] - return $ Hint ixfun DefaultSpace + -- Heuristic: do not rearrange for returned arrays that are + -- sufficiently small. + coalesceReturnOfShape _ [] = False + coalesceReturnOfShape bs [Constant (IntValue (Int64Value d))] = bs * d > 4 + coalesceReturnOfShape _ _ = True - hint _ _ = return NoHint + hint t Returns {} + | coalesceReturnOfShape (primByteSize (elemType t)) $ arrayDims t = do + let space_dims = segSpaceDims space + t_dims <- mapM dimAllocationSize $ arrayDims t + return $ Hint (innermost space_dims t_dims) DefaultSpace + hint t (ConcatReturns SplitStrided {} w _ _) = do + t_dims <- mapM dimAllocationSize $ arrayDims t + return $ Hint (innermost [w] t_dims) DefaultSpace + hint Prim {} (ConcatReturns SplitContiguous w elems_per_thread _) = do + let ixfun_base = IxFun.iota [sExt64 num_threads, pe64 elems_per_thread] + ixfun_tr = IxFun.permute ixfun_base [1, 0] + ixfun = IxFun.reshape ixfun_tr $ map (DimNew . pe64) [w] + return $ Hint ixfun DefaultSpace + hint _ _ = return NoHint innermost :: [SubExp] -> [SubExp] -> IxFun innermost space_dims t_dims = let r = length t_dims dims = space_dims ++ t_dims - perm = [length space_dims..length space_dims+r-1] ++ - [0..length space_dims-1] + perm = + [length space_dims .. length space_dims + r -1] + ++ [0 .. length space_dims -1] perm_inv = rearrangeInverse perm dims_perm = rearrangeShape perm dims - ixfun_base = IxFun.iota $ map (primExpFromSubExp int32) dims_perm + ixfun_base = IxFun.iota $ map pe64 dims_perm ixfun_rearranged = IxFun.permute ixfun_base perm_inv - in ixfun_rearranged + in ixfun_rearranged semiStatic :: S.Set VName -> SubExp -> Bool -semiStatic _ Constant{} = True +semiStatic _ Constant {} = True semiStatic consts (Var v) = v `S.member` consts inGroupExpHints :: Allocator KernelsMem m => Exp KernelsMem -> m [ExpHint] inGroupExpHints (Op (Inner (SegOp (SegMap _ space ts body)))) | any private $ kernelBodyResult body = do - consts <- askConsts - return $ do - (t, r) <- zip ts $ kernelBodyResult body - return $ - if private r && all (semiStatic consts) (arrayDims t) - then let seg_dims = map (primExpFromSubExp int32) $ segSpaceDims space - dims = seg_dims ++ map (primExpFromSubExp int32) (arrayDims t) - nilSlice d = DimSlice 0 d 0 - in Hint (IxFun.slice (IxFun.iota dims) $ - fullSliceNum dims $ map nilSlice seg_dims) $ - ScalarSpace (arrayDims t) $ elemType t + consts <- askConsts + return $ do + (t, r) <- zip ts $ kernelBodyResult body + return $ + if private r && all (semiStatic consts) (arrayDims t) + then + let seg_dims = map pe64 $ segSpaceDims space + dims = seg_dims ++ map pe64 (arrayDims t) + nilSlice d = DimSlice 0 d 0 + in Hint + ( IxFun.slice (IxFun.iota dims) $ + fullSliceNum dims $ map nilSlice seg_dims + ) + $ ScalarSpace (arrayDims t) $ elemType t else NoHint - where private (Returns ResultPrivate _) = True - private _ = False + where + private (Returns ResultPrivate _) = True + private _ = False inGroupExpHints e = return $ replicate (expExtTypeSize e) NoHint inThreadExpHints :: Allocator KernelsMem m => Exp KernelsMem -> m [ExpHint] inThreadExpHints e = do consts <- askConsts mapM (maybePrivate consts) =<< expExtType e - where maybePrivate consts t - | Just (Array pt shape _) <- hasStaticShape t, - all (semiStatic consts) $ shapeDims shape = do - let ixfun = IxFun.iota $ map (primExpFromSubExp int32) $ - shapeDims shape - return $ Hint ixfun $ ScalarSpace (shapeDims shape) pt - | otherwise = - return NoHint + where + maybePrivate consts t + | Just (Array pt shape _) <- hasStaticShape t, + all (semiStatic consts) $ shapeDims shape = do + let ixfun = IxFun.iota $ map pe64 $ shapeDims shape + return $ Hint ixfun $ ScalarSpace (shapeDims shape) pt + | otherwise = + return NoHint -- | The pass from 'Kernels' to 'KernelsMem'. explicitAllocations :: Pass Kernels KernelsMem explicitAllocations = explicitAllocationsGeneric handleHostOp kernelExpHints -- | Convert some 'Kernels' stms to 'KernelsMem'. -explicitAllocationsInStms :: (MonadFreshNames m, HasScope KernelsMem m) => - Stms Kernels -> m (Stms KernelsMem) +explicitAllocationsInStms :: + (MonadFreshNames m, HasScope KernelsMem m) => + Stms Kernels -> + m (Stms KernelsMem) explicitAllocationsInStms = explicitAllocationsInStmsGeneric handleHostOp kernelExpHints diff --git a/src/Futhark/Pass/ExplicitAllocations/SegOp.hs b/src/Futhark/Pass/ExplicitAllocations/SegOp.hs --- a/src/Futhark/Pass/ExplicitAllocations/SegOp.hs +++ b/src/Futhark/Pass/ExplicitAllocations/SegOp.hs @@ -1,73 +1,94 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-} + module Futhark.Pass.ExplicitAllocations.SegOp - ( allocInKernelBody - , allocInBinOpLambda - ) + ( allocInKernelBody, + allocInBinOpLambda, + ) where -import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.IR.KernelsMem +import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.Pass.ExplicitAllocations -allocInKernelBody :: Allocable fromlore tolore => - KernelBody fromlore - -> AllocM fromlore tolore (KernelBody tolore) +allocInKernelBody :: + Allocable fromlore tolore => + KernelBody fromlore -> + AllocM fromlore tolore (KernelBody tolore) allocInKernelBody (KernelBody () stms res) = allocInStms stms $ \stms' -> return $ KernelBody () stms' res -allocInLambda :: Allocable fromlore tolore => - [LParam tolore] -> Body fromlore -> [Type] - -> AllocM fromlore tolore (Lambda tolore) +allocInLambda :: + Allocable fromlore tolore => + [LParam tolore] -> + Body fromlore -> + [Type] -> + AllocM fromlore tolore (Lambda tolore) allocInLambda params body rettype = do body' <- localScope (scopeOfLParams params) $ - allocInStms (bodyStms body) $ \bnds' -> - return $ Body () bnds' $ bodyResult body + allocInStms (bodyStms body) $ \bnds' -> + return $ Body () bnds' $ bodyResult body return $ Lambda params body' rettype -allocInBinOpParams :: Allocable fromlore tolore => - SubExp - -> PrimExp VName -> PrimExp VName - -> [LParam fromlore] - -> [LParam fromlore] - -> AllocM fromlore tolore ([LParam tolore], [LParam tolore]) +allocInBinOpParams :: + Allocable fromlore tolore => + SubExp -> + TPrimExp Int64 VName -> + TPrimExp Int64 VName -> + [LParam fromlore] -> + [LParam fromlore] -> + AllocM fromlore tolore ([LParam tolore], [LParam tolore]) allocInBinOpParams num_threads my_id other_id xs ys = unzip <$> zipWithM alloc xs ys - where alloc x y = - case paramType x of - Array bt shape u -> do - twice_num_threads <- - letSubExp "twice_num_threads" $ - BasicOp $ BinOp (Mul Int32 OverflowUndef) num_threads $ intConst Int32 2 - let t = paramType x `arrayOfRow` twice_num_threads - mem <- allocForArray t DefaultSpace - -- XXX: this iota ixfun is a bit inefficient; leading to - -- uncoalesced access. - let base_dims = map (primExpFromSubExp int32) (arrayDims t) - ixfun_base = IxFun.iota base_dims - ixfun_x = IxFun.slice ixfun_base $ - fullSliceNum base_dims [DimFix my_id] - ixfun_y = IxFun.slice ixfun_base $ - fullSliceNum base_dims [DimFix other_id] - return (x { paramDec = MemArray bt shape u $ ArrayIn mem ixfun_x }, - y { paramDec = MemArray bt shape u $ ArrayIn mem ixfun_y }) - Prim bt -> - return (x { paramDec = MemPrim bt }, - y { paramDec = MemPrim bt }) - Mem space -> - return (x { paramDec = MemMem space }, - y { paramDec = MemMem space }) + where + alloc x y = + case paramType x of + Array bt shape u -> do + twice_num_threads <- + letSubExp "twice_num_threads" $ + BasicOp $ BinOp (Mul Int64 OverflowUndef) num_threads $ intConst Int64 2 + let t = paramType x `arrayOfRow` twice_num_threads + mem <- allocForArray t DefaultSpace + -- XXX: this iota ixfun is a bit inefficient; leading to + -- uncoalesced access. + let base_dims = map pe64 $ arrayDims t + ixfun_base = IxFun.iota base_dims + ixfun_x = + IxFun.slice ixfun_base $ + fullSliceNum base_dims [DimFix my_id] + ixfun_y = + IxFun.slice ixfun_base $ + fullSliceNum base_dims [DimFix other_id] + return + ( x {paramDec = MemArray bt shape u $ ArrayIn mem ixfun_x}, + y {paramDec = MemArray bt shape u $ ArrayIn mem ixfun_y} + ) + Prim bt -> + return + ( x {paramDec = MemPrim bt}, + y {paramDec = MemPrim bt} + ) + Mem space -> + return + ( x {paramDec = MemMem space}, + y {paramDec = MemMem space} + ) -allocInBinOpLambda :: Allocable fromlore tolore => - SubExp -> SegSpace -> Lambda fromlore - -> AllocM fromlore tolore (Lambda tolore) +allocInBinOpLambda :: + Allocable fromlore tolore => + SubExp -> + SegSpace -> + Lambda fromlore -> + AllocM fromlore tolore (Lambda tolore) allocInBinOpLambda num_threads (SegSpace flat _) lam = do let (acc_params, arr_params) = splitAt (length (lambdaParams lam) `div` 2) $ lambdaParams lam - index_x = LeafExp flat int32 - index_y = index_x + primExpFromSubExp int32 num_threads + index_x = TPrimExp $ LeafExp flat int64 + index_y = index_x + pe64 num_threads (acc_params', arr_params') <- allocInBinOpParams num_threads index_x index_y acc_params arr_params - allocInLambda (acc_params' ++ arr_params') - (lambdaBody lam) (lambdaReturnType lam) + allocInLambda + (acc_params' ++ arr_params') + (lambdaBody lam) + (lambdaReturnType lam) diff --git a/src/Futhark/Pass/ExplicitAllocations/Seq.hs b/src/Futhark/Pass/ExplicitAllocations/Seq.hs --- a/src/Futhark/Pass/ExplicitAllocations/Seq.hs +++ b/src/Futhark/Pass/ExplicitAllocations/Seq.hs @@ -1,15 +1,16 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# OPTIONS_GHC -fno-warn-orphans #-} + module Futhark.Pass.ExplicitAllocations.Seq - ( explicitAllocations - , simplifiable - ) + ( explicitAllocations, + simplifiable, + ) where -import Futhark.Pass -import Futhark.IR.SeqMem import Futhark.IR.Seq +import Futhark.IR.SeqMem +import Futhark.Pass import Futhark.Pass.ExplicitAllocations explicitAllocations :: Pass Seq SeqMem 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 @@ -1,12 +1,13 @@ {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE LambdaCase #-} -{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE RankNTypes #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + -- | Kernel extraction. -- -- In the following, I will use the term "width" to denote the amount @@ -157,43 +158,41 @@ -- intermediate array (@b@ above) must be written to main memory. An -- often better approach is to just turn the entire @redomap@ into a -- single kernel. --- module Futhark.Pass.ExtractKernels (extractKernels) where import Control.Monad.Identity import Control.Monad.RWS.Strict import Control.Monad.Reader import Data.Maybe - -import Prelude hiding (log) - -import Futhark.IR.SOACS -import Futhark.IR.SOACS.Simplify (simplifyStms) import qualified Futhark.IR.Kernels as Out import Futhark.IR.Kernels.Kernel +import Futhark.IR.SOACS +import Futhark.IR.SOACS.Simplify (simplifyStms) import Futhark.MonadFreshNames -import Futhark.Tools -import qualified Futhark.Transform.FirstOrderTransform as FOT -import Futhark.Transform.Rename import Futhark.Pass -import Futhark.Pass.ExtractKernels.Distribution +import Futhark.Pass.ExtractKernels.BlockedKernel import Futhark.Pass.ExtractKernels.DistributeNests +import Futhark.Pass.ExtractKernels.Distribution import Futhark.Pass.ExtractKernels.ISRWIM -import Futhark.Pass.ExtractKernels.BlockedKernel import Futhark.Pass.ExtractKernels.Intragroup import Futhark.Pass.ExtractKernels.StreamKernel import Futhark.Pass.ExtractKernels.ToKernels +import Futhark.Tools +import qualified Futhark.Transform.FirstOrderTransform as FOT +import Futhark.Transform.Rename import Futhark.Util import Futhark.Util.Log +import Prelude hiding (log) -- | Transform a program using SOACs to a program using explicit -- kernels, using the kernel extraction transformation. extractKernels :: Pass SOACS Out.Kernels extractKernels = - Pass { passName = "extract kernels" - , passDescription = "Perform kernel extraction" - , passFunction = transformProg - } + Pass + { passName = "extract kernels", + passDescription = "Perform kernel extraction", + passFunction = transformProg + } transformProg :: Prog SOACS -> PassM (Prog Out.Kernels) transformProg (Prog consts funs) = do @@ -204,175 +203,190 @@ -- In order to generate more stable threshold names, we keep track of -- the numbers used for thresholds separately from the ordinary name -- source, -data State = State { stateNameSource :: VNameSource - , stateThresholdCounter :: Int - } +data State = State + { stateNameSource :: VNameSource, + stateThresholdCounter :: Int + } newtype DistribM a = DistribM (RWS (Scope Out.Kernels) Log State a) - deriving (Functor, Applicative, Monad, - HasScope Out.Kernels, LocalScope Out.Kernels, - MonadState State, - MonadLogger) + deriving + ( Functor, + Applicative, + Monad, + HasScope Out.Kernels, + LocalScope Out.Kernels, + MonadState State, + MonadLogger + ) instance MonadFreshNames DistribM where getNameSource = gets stateNameSource - putNameSource src = modify $ \s -> s { stateNameSource = src } + putNameSource src = modify $ \s -> s {stateNameSource = src} -runDistribM :: (MonadLogger m, MonadFreshNames m) => - DistribM a -> m a +runDistribM :: + (MonadLogger m, MonadFreshNames m) => + DistribM a -> + m a runDistribM (DistribM m) = do (x, msgs) <- modifyNameSource $ \src -> let (x, s, msgs) = runRWS m mempty (State src 0) - in ((x, msgs), stateNameSource s) + in ((x, msgs), stateNameSource s) addLog msgs return x -transformFunDef :: (MonadFreshNames m, MonadLogger m) => - Scope Out.Kernels -> FunDef SOACS - -> m (Out.FunDef Out.Kernels) +transformFunDef :: + (MonadFreshNames m, MonadLogger m) => + Scope Out.Kernels -> + FunDef SOACS -> + m (Out.FunDef Out.Kernels) transformFunDef scope (FunDef entry attrs name rettype params body) = runDistribM $ do - body' <- localScope (scope <> scopeOfFParams params) $ - transformBody mempty body + body' <- + localScope (scope <> scopeOfFParams params) $ + transformBody mempty body return $ FunDef entry attrs name rettype params body' type KernelsStms = Stms Out.Kernels transformBody :: KernelPath -> Body -> DistribM (Out.Body Out.Kernels) -transformBody path body = do bnds <- transformStms path $ stmsToList $ bodyStms body - return $ mkBody bnds $ bodyResult body +transformBody path body = do + bnds <- transformStms path $ stmsToList $ bodyStms body + return $ mkBody bnds $ bodyResult body transformStms :: KernelPath -> [Stm] -> DistribM KernelsStms transformStms _ [] = return mempty -transformStms path (bnd:bnds) = +transformStms path (bnd : bnds) = sequentialisedUnbalancedStm bnd >>= \case Nothing -> do bnd' <- transformStm path bnd inScopeOf bnd' $ - (bnd'<>) <$> transformStms path bnds + (bnd' <>) <$> transformStms path bnds Just bnds' -> transformStms path $ stmsToList bnds' <> bnds unbalancedLambda :: Lambda -> Bool unbalancedLambda lam = unbalancedBody - (namesFromList $ map paramName $ lambdaParams lam) $ - lambdaBody lam - - where subExpBound (Var i) bound = i `nameIn` bound - subExpBound (Constant _) _ = False - - unbalancedBody bound body = - any (unbalancedStm (bound <> boundInBody body) . stmExp) $ - bodyStms body - - -- XXX - our notion of balancing is probably still too naive. - unbalancedStm bound (Op (Stream w _ _ _)) = - w `subExpBound` bound - unbalancedStm bound (Op (Screma w _ _)) = - w `subExpBound` bound - unbalancedStm _ Op{} = - False - unbalancedStm _ DoLoop{} = False + (namesFromList $ map paramName $ lambdaParams lam) + $ lambdaBody lam + where + subExpBound (Var i) bound = i `nameIn` bound + subExpBound (Constant _) _ = False - unbalancedStm bound (If cond tbranch fbranch _) = - cond `subExpBound` bound && - (unbalancedBody bound tbranch || unbalancedBody bound fbranch) + unbalancedBody bound body = + any (unbalancedStm (bound <> boundInBody body) . stmExp) $ + bodyStms body - unbalancedStm _ (BasicOp _) = - False - unbalancedStm _ (Apply fname _ _ _) = - not $ isBuiltInFunction fname + -- XXX - our notion of balancing is probably still too naive. + unbalancedStm bound (Op (Stream w _ _ _)) = + w `subExpBound` bound + unbalancedStm bound (Op (Screma w _ _)) = + w `subExpBound` bound + unbalancedStm _ Op {} = + False + unbalancedStm _ DoLoop {} = False + unbalancedStm bound (If cond tbranch fbranch _) = + cond `subExpBound` bound + && (unbalancedBody bound tbranch || unbalancedBody bound fbranch) + unbalancedStm _ (BasicOp _) = + False + unbalancedStm _ (Apply fname _ _ _) = + not $ isBuiltInFunction fname sequentialisedUnbalancedStm :: Stm -> DistribM (Maybe (Stms SOACS)) sequentialisedUnbalancedStm (Let pat _ (Op soac@(Screma _ form _))) | Just (_, lam2) <- isRedomapSOAC form, - unbalancedLambda lam2, lambdaContainsParallelism lam2 = do - types <- asksScope scopeForSOACs - Just . snd <$> runBinderT (FOT.transformSOAC pat soac) types + unbalancedLambda lam2, + lambdaContainsParallelism lam2 = do + types <- asksScope scopeForSOACs + Just . snd <$> runBinderT (FOT.transformSOAC pat soac) types sequentialisedUnbalancedStm _ = return Nothing -cmpSizeLe :: String -> Out.SizeClass -> [SubExp] - -> DistribM ((SubExp, Name), Out.Stms Out.Kernels) +cmpSizeLe :: + String -> + Out.SizeClass -> + [SubExp] -> + DistribM ((SubExp, Name), Out.Stms Out.Kernels) cmpSizeLe desc size_class to_what = do x <- gets stateThresholdCounter - modify $ \s -> s { stateThresholdCounter = x + 1} + modify $ \s -> s {stateThresholdCounter = x + 1} let size_key = nameFromString $ desc ++ "_" ++ show x runBinder $ do - to_what' <- letSubExp "comparatee" =<< - foldBinOp (Mul Int32 OverflowUndef) (intConst Int32 1) to_what + to_what' <- + letSubExp "comparatee" + =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) to_what cmp_res <- letSubExp desc $ Op $ SizeOp $ CmpSizeLe size_key size_class to_what' return (cmp_res, size_key) -kernelAlternatives :: (MonadFreshNames m, HasScope Out.Kernels m) => - Out.Pattern Out.Kernels - -> Out.Body Out.Kernels - -> [(SubExp, Out.Body Out.Kernels)] - -> m (Out.Stms Out.Kernels) +kernelAlternatives :: + (MonadFreshNames m, HasScope Out.Kernels m) => + Out.Pattern Out.Kernels -> + Out.Body Out.Kernels -> + [(SubExp, Out.Body Out.Kernels)] -> + m (Out.Stms Out.Kernels) kernelAlternatives pat default_body [] = runBinder_ $ do ses <- bodyBind default_body forM_ (zip (patternNames pat) ses) $ \(name, se) -> letBindNames [name] $ BasicOp $ SubExp se -kernelAlternatives pat default_body ((cond,alt):alts) = runBinder_ $ do - alts_pat <- fmap (Pattern []) $ forM (patternElements pat) $ \pe -> do - name <- newVName $ baseString $ patElemName pe - return pe { patElemName = name } +kernelAlternatives pat default_body ((cond, alt) : alts) = runBinder_ $ do + alts_pat <- fmap (Pattern []) $ + forM (patternElements pat) $ \pe -> do + name <- newVName $ baseString $ patElemName pe + return pe {patElemName = name} alt_stms <- kernelAlternatives alts_pat default_body alts let alt_body = mkBody alt_stms $ map Var $ patternValueNames alts_pat - letBind pat $ If cond alt alt_body $ - IfDec (staticShapes (patternTypes pat)) IfEquiv + letBind pat $ + If cond alt alt_body $ + IfDec (staticShapes (patternTypes pat)) IfEquiv transformStm :: KernelPath -> Stm -> DistribM KernelsStms - transformStm _ stm | "sequential" `inAttrs` stmAuxAttrs (stmAux stm) = runBinder_ $ FOT.transformStmRecursively stm - transformStm path (Let pat aux (Op soac)) | "sequential_outer" `inAttrs` stmAuxAttrs aux = - transformStms path . stmsToList . fmap (certify (stmAuxCerts aux)) =<< - runBinder_ (FOT.transformSOAC pat soac) - + transformStms path . stmsToList . fmap (certify (stmAuxCerts aux)) + =<< runBinder_ (FOT.transformSOAC pat soac) transformStm path (Let pat aux (If c tb fb rt)) = do tb' <- transformBody path tb fb' <- transformBody path fb return $ oneStm $ Let pat aux $ If c tb' fb' rt - transformStm path (Let pat aux (DoLoop ctx val form body)) = - localScope (castScope (scopeOf form) <> - scopeOfFParams mergeparams) $ - oneStm . Let pat aux . DoLoop ctx val form' <$> transformBody path body - where mergeparams = map fst $ ctx ++ val - form' = case form of - WhileLoop cond -> - WhileLoop cond - ForLoop i it bound ps -> - ForLoop i it bound ps - + localScope + ( castScope (scopeOf form) + <> scopeOfFParams mergeparams + ) + $ oneStm . Let pat aux . DoLoop ctx val form' <$> transformBody path body + where + mergeparams = map fst $ ctx ++ val + form' = case form of + WhileLoop cond -> + WhileLoop cond + ForLoop i it bound ps -> + ForLoop i it bound ps transformStm path (Let pat aux (Op (Screma w form arrs))) | Just lam <- isMapSOAC form = - onMap path $ MapLoop pat aux w lam arrs - + onMap path $ MapLoop pat aux w lam arrs transformStm path (Let res_pat (StmAux cs _ _) (Op (Screma w form arrs))) | Just scans <- isScanSOAC form, 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) - + types <- asksScope scopeForSOACs + 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 - let scan_lam'' = soacsLambdaToKernels scan_lam' - 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 [] []) + scan_ops <- forM scans $ \(Scan scan_lam nes) -> do + (scan_lam', nes', shape) <- determineReduceOp scan_lam nes + let scan_lam'' = soacsLambdaToKernels scan_lam' + 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 [] [] + ) -- We are only willing to generate code for scanomaps that do not -- involve array accumulators, and do not have parallelism in their @@ -382,170 +396,186 @@ -- we will still crash in code generation). However, if the map -- lambda is already identity, let's just go ahead here. | Just (scans, map_lam) <- isScanomapSOAC form, - (all primType (concatMap (lambdaReturnType . scanLambda) scans) && - not (lambdaContainsParallelism map_lam)) || isIdentityLambda map_lam = runBinder_ $ do - - scan_ops <- forM scans $ \(Scan scan_lam nes) -> do - let scan_lam' = soacsLambdaToKernels scan_lam - return $ SegBinOp Noncommutative scan_lam' nes mempty - - let map_lam' = soacsLambdaToKernels map_lam - lvl <- segThreadCapped [w] "segscan" $ NoRecommendation SegNoVirt - addStms =<< segScan lvl res_pat w scan_ops map_lam' arrs [] [] + ( all primType (concatMap (lambdaReturnType . scanLambda) scans) + && not (lambdaContainsParallelism map_lam) + ) + || isIdentityLambda map_lam = runBinder_ $ do + scan_ops <- forM scans $ \(Scan scan_lam nes) -> do + let scan_lam' = soacsLambdaToKernels scan_lam + return $ SegBinOp Noncommutative scan_lam' nes mempty + let map_lam' = soacsLambdaToKernels map_lam + lvl <- segThreadCapped [w] "segscan" $ NoRecommendation SegNoVirt + addStms =<< segScan lvl res_pat w scan_ops map_lam' arrs [] [] transformStm path (Let res_pat aux (Op (Screma w form arrs))) | Just [Reduce comm red_fun nes] <- isReduceSOAC form, - let comm' | commutativeLambda red_fun = Commutative - | otherwise = comm, + let comm' + | commutativeLambda red_fun = Commutative + | otherwise = comm, Just do_irwim <- irwim res_pat w comm' red_fun $ zip nes arrs = do - types <- asksScope scopeForSOACs - (_, bnds) <- fst <$> runBinderT (simplifyStms =<< collectStms_ (auxing aux do_irwim)) types - transformStms path $ stmsToList bnds - + types <- asksScope scopeForSOACs + (_, bnds) <- fst <$> runBinderT (simplifyStms =<< collectStms_ (auxing aux do_irwim)) types + transformStms path $ stmsToList bnds transformStm path (Let pat aux@(StmAux cs _ _) (Op (Screma w form arrs))) | Just (reds, map_lam) <- isRedomapSOAC form = do - - let paralleliseOuter = runBinder_ $ do - red_ops <- forM reds $ \(Reduce comm red_lam nes) -> do - (red_lam', nes', shape) <- determineReduceOp red_lam nes - let comm' | commutativeLambda red_lam' = Commutative - | otherwise = comm - red_lam'' = soacsLambdaToKernels red_lam' - 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) + let paralleliseOuter = runBinder_ $ do + red_ops <- forM reds $ \(Reduce comm red_lam nes) -> do + (red_lam', nes', shape) <- determineReduceOp red_lam nes + let comm' + | commutativeLambda red_lam' = Commutative + | otherwise = comm + red_lam'' = soacsLambdaToKernels red_lam' + 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 + ) - outerParallelBody = - renameBody =<< - (mkBody <$> paralleliseOuter <*> pure (map Var (patternNames pat))) + 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] - where comm' | commutativeLambda red_lam = Commutative - | otherwise = comm - (Reduce comm red_lam nes) = singleReduce reds + paralleliseInner path' = do + (mapbnd, redbnd) <- redomapToMapAndReduce pat (w, comm', red_lam, map_lam, nes, arrs) + transformStms path' [certify cs mapbnd, certify cs redbnd] + where + comm' + | commutativeLambda red_lam = Commutative + | otherwise = comm + (Reduce comm red_lam nes) = singleReduce reds - innerParallelBody path' = - renameBody =<< - (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) + innerParallelBody path' = + renameBody + =<< (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) - if not (lambdaContainsParallelism map_lam) || - "sequential_inner" `inAttrs` stmAuxAttrs aux - then paralleliseOuter - else do - ((outer_suff, outer_suff_key), suff_stms) <- - sufficientParallelism "suff_outer_redomap" [w] path Nothing + if not (lambdaContainsParallelism map_lam) + || "sequential_inner" `inAttrs` stmAuxAttrs aux + then paralleliseOuter + else do + ((outer_suff, outer_suff_key), suff_stms) <- + sufficientParallelism "suff_outer_redomap" [w] path Nothing - outer_stms <- outerParallelBody - inner_stms <- innerParallelBody ((outer_suff_key, False):path) + outer_stms <- outerParallelBody + inner_stms <- innerParallelBody ((outer_suff_key, False) : path) - (suff_stms<>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] + (suff_stms <>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] -- 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))) | 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) - + -- 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))) | "sequential_inner" `inAttrs` stmAuxAttrs aux = - paralleliseOuter path + paralleliseOuter path | otherwise = do - ((outer_suff, outer_suff_key), suff_stms) <- - sufficientParallelism "suff_outer_stream" [w] path Nothing - - outer_stms <- outerParallelBody ((outer_suff_key, True) : path) - inner_stms <- innerParallelBody ((outer_suff_key, False) : path) + ((outer_suff, outer_suff_key), suff_stms) <- + sufficientParallelism "suff_outer_stream" [w] path Nothing - (suff_stms<>) <$> - kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] + outer_stms <- outerParallelBody ((outer_suff_key, True) : path) + inner_stms <- innerParallelBody ((outer_suff_key, False) : path) + (suff_stms <>) + <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] where paralleliseOuter path' | not $ all primType $ lambdaReturnType red_fun = do - -- Split into a chunked map and a reduction, with the latter - -- further transformed. - let fold_fun' = soacsLambdaToKernels fold_fun - - let (red_pat_elems, concat_pat_elems) = - splitAt (length nes) $ patternValueElements pat - red_pat = Pattern [] red_pat_elems + -- Split into a chunked map and a reduction, with the latter + -- further transformed. + let fold_fun' = soacsLambdaToKernels fold_fun - ((num_threads, red_results), stms) <- - streamMap segThreadCapped - (map (baseString . patElemName) red_pat_elems) concat_pat_elems w - Noncommutative fold_fun' nes arrs + let (red_pat_elems, concat_pat_elems) = + splitAt (length nes) $ patternValueElements pat + red_pat = Pattern [] red_pat_elems - reduce_soac <- reduceSOAC [Reduce comm' red_fun nes] + ((num_threads, red_results), stms) <- + streamMap + segThreadCapped + (map (baseString . patElemName) red_pat_elems) + concat_pat_elems + w + Noncommutative + fold_fun' + nes + arrs - (stms<>) <$> - inScopeOf stms - (transformStm path' $ Let red_pat aux { stmAuxAttrs = mempty } $ - Op (Screma num_threads reduce_soac red_results)) + reduce_soac <- reduceSOAC [Reduce comm' red_fun nes] + (stms <>) + <$> inScopeOf + stms + ( transformStm path' $ + Let red_pat aux {stmAuxAttrs = mempty} $ + Op (Screma num_threads reduce_soac red_results) + ) | otherwise = do - let red_fun_sequential = soacsLambdaToKernels red_fun - fold_fun_sequential = soacsLambdaToKernels fold_fun - fmap (certify cs) <$> - streamRed segThreadCapped - pat w comm' red_fun_sequential fold_fun_sequential nes arrs + let red_fun_sequential = soacsLambdaToKernels red_fun + fold_fun_sequential = soacsLambdaToKernels fold_fun + fmap (certify cs) + <$> streamRed + segThreadCapped + pat + w + comm' + red_fun_sequential + fold_fun_sequential + nes + arrs outerParallelBody path' = - renameBody =<< - (mkBody <$> paralleliseOuter path' <*> pure (map Var (patternNames pat))) + renameBody + =<< (mkBody <$> paralleliseOuter path' <*> pure (map Var (patternNames pat))) 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 =<< - (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) - - comm' | commutativeLambda red_fun, o /= InOrder = Commutative - | otherwise = comm + renameBody + =<< (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) + comm' + | commutativeLambda red_fun, o /= InOrder = Commutative + | otherwise = comm transformStm path (Let pat (StmAux cs _ _) (Op (Screma w form arrs))) = do -- This screma is too complicated for us to immediately do -- anything, so split it up and try again. scope <- asksScope scopeForSOACs - transformStms path . map (certify cs) . stmsToList . snd =<< - runBinderT (dissectScrema pat w form arrs) scope - + 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 -- 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' 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 ] + 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] body = KernelBody () kstms krets - inputs = do (p, p_a) <- zip (lambdaParams lam') ivs - return $ KernelInput (paramName p) (paramType p) p_a [Var write_i] + 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 (map rowType $ patternTypes pat) body certifying cs $ do addStms stms letBind pat $ Op $ SegOp kernel - transformStm _ (Let orig_pat (StmAux cs _ _) (Op (Hist w ops bucket_fun imgs))) = do let bfun' = soacsLambdaToKernels bucket_fun @@ -555,13 +585,17 @@ runBinder_ $ do lvl <- segThreadCapped [w] "seghist" $ NoRecommendation SegNoVirt addStms =<< histKernel onLambda lvl orig_pat [] [] cs w ops bfun' imgs - where onLambda = pure . soacsLambdaToKernels - + where + onLambda = pure . soacsLambdaToKernels transformStm _ bnd = runBinder_ $ FOT.transformStmRecursively bnd -sufficientParallelism :: String -> [SubExp] -> KernelPath -> Maybe Int32 - -> DistribM ((SubExp, Name), Out.Stms Out.Kernels) +sufficientParallelism :: + String -> + [SubExp] -> + KernelPath -> + Maybe Int64 -> + DistribM ((SubExp, Name), Out.Stms Out.Kernels) sufficientParallelism desc ws path def = cmpSizeLe desc (Out.SizeThreshold path def) ws @@ -570,73 +604,78 @@ -- parallelism inside a loop. worthIntraGroup :: Lambda -> Bool worthIntraGroup lam = bodyInterest (lambdaBody lam) > 1 - where bodyInterest body = - sum $ interest <$> bodyStms body - interest stm - | "sequential" `inAttrs` attrs = - 0::Int - | Op (Screma w form _) <- stmExp stm, - Just lam' <- isMapSOAC form = - mapLike w lam' - | Op (Scatter w lam' _ _) <- stmExp stm = - mapLike w lam' - | DoLoop _ _ _ body <- stmExp stm = - bodyInterest body * 10 - | If _ tbody fbody _ <- stmExp stm = - max (bodyInterest tbody) (bodyInterest fbody) - | Op (Screma w (ScremaForm _ _ lam') _) <- stmExp stm = - zeroIfTooSmall w + bodyInterest (lambdaBody lam') - | Op (Stream _ (Sequential _) lam' _) <- stmExp stm = - bodyInterest $ lambdaBody lam' - | otherwise = - 0 - - where attrs = stmAuxAttrs $ stmAux stm - sequential_inner = "sequential_inner" `inAttrs` attrs + where + bodyInterest body = + sum $ interest <$> bodyStms body + interest stm + | "sequential" `inAttrs` attrs = + 0 :: Int + | Op (Screma w form _) <- stmExp stm, + Just lam' <- isMapSOAC form = + mapLike w lam' + | Op (Scatter w lam' _ _) <- stmExp stm = + mapLike w lam' + | DoLoop _ _ _ body <- stmExp stm = + bodyInterest body * 10 + | If _ tbody fbody _ <- stmExp stm = + max (bodyInterest tbody) (bodyInterest fbody) + | Op (Screma w (ScremaForm _ _ lam') _) <- stmExp stm = + zeroIfTooSmall w + bodyInterest (lambdaBody lam') + | Op (Stream _ (Sequential _) lam' _) <- stmExp stm = + bodyInterest $ lambdaBody lam' + | otherwise = + 0 + where + attrs = stmAuxAttrs $ stmAux stm + sequential_inner = "sequential_inner" `inAttrs` attrs - zeroIfTooSmall (Constant (IntValue x)) - | intToInt64 x < 32 = 0 - zeroIfTooSmall _ = 1 + zeroIfTooSmall (Constant (IntValue x)) + | intToInt64 x < 32 = 0 + zeroIfTooSmall _ = 1 - mapLike w lam' = - if sequential_inner - then 0 - else max (zeroIfTooSmall w) (bodyInterest (lambdaBody lam')) + mapLike w lam' = + if sequential_inner + then 0 + else max (zeroIfTooSmall w) (bodyInterest (lambdaBody lam')) -- | A lambda is worth sequentialising if it contains enough nested -- parallelism of an interesting kind. worthSequentialising :: Lambda -> Bool worthSequentialising lam = bodyInterest (lambdaBody lam) > 1 - where bodyInterest body = - sum $ interest <$> bodyStms body - interest stm - | "sequential" `inAttrs` attrs = - 0::Int - | Op (Screma _ form@(ScremaForm _ _ lam') _) <- stmExp stm, - isJust $ isMapSOAC form = - if sequential_inner - then 0 - else bodyInterest (lambdaBody lam') - | Op Scatter{} <- stmExp stm = - 0 -- Basically a map. - | DoLoop _ _ _ body <- stmExp stm = - bodyInterest body * 10 - | Op (Screma _ form@(ScremaForm _ _ lam') _) <- stmExp stm = - 1 + bodyInterest (lambdaBody lam') + - -- Give this a bigger score if it's a redomap, as these - -- are often tileable and thus benefit more from - -- sequentialisation. - case isRedomapSOAC form of - Just _ -> 1 - Nothing -> 0 - | otherwise = - 0 - - where attrs = stmAuxAttrs $ stmAux stm - sequential_inner = "sequential_inner" `inAttrs` attrs + where + bodyInterest body = + sum $ interest <$> bodyStms body + interest stm + | "sequential" `inAttrs` attrs = + 0 :: Int + | Op (Screma _ form@(ScremaForm _ _ lam') _) <- stmExp stm, + isJust $ isMapSOAC form = + if sequential_inner + then 0 + else bodyInterest (lambdaBody lam') + | Op Scatter {} <- stmExp stm = + 0 -- Basically a map. + | DoLoop _ _ _ body <- stmExp stm = + bodyInterest body * 10 + | Op (Screma _ form@(ScremaForm _ _ lam') _) <- stmExp stm = + 1 + bodyInterest (lambdaBody lam') + + + -- Give this a bigger score if it's a redomap, as these + -- are often tileable and thus benefit more from + -- sequentialisation. + case isRedomapSOAC form of + Just _ -> 1 + Nothing -> 0 + | otherwise = + 0 + where + attrs = stmAuxAttrs $ stmAux stm + sequential_inner = "sequential_inner" `inAttrs` attrs -onTopLevelStms :: KernelPath -> Stms SOACS - -> DistNestT Out.Kernels DistribM KernelsStms +onTopLevelStms :: + KernelPath -> + Stms SOACS -> + DistNestT Out.Kernels DistribM KernelsStms onTopLevelStms path stms = liftInner $ transformStms path $ stmsToList stms @@ -644,30 +683,36 @@ onMap path (MapLoop pat aux w lam arrs) = do types <- askScope let loopnest = MapNesting pat aux w $ zip (lambdaParams lam) arrs - env path' = DistEnv - { distNest = singleNesting (Nesting mempty loopnest) - , distScope = scopeOfPattern pat <> - scopeForKernels (scopeOf lam) <> - types - , distOnInnerMap = onInnerMap path' - , distOnTopLevelStms = onTopLevelStms path' - , distSegLevel = segThreadCapped - , distOnSOACSStms = pure . oneStm . soacsStmToKernels - , distOnSOACSLambda = pure . soacsLambdaToKernels - } + env path' = + DistEnv + { distNest = singleNesting (Nesting mempty loopnest), + distScope = + scopeOfPattern pat + <> scopeForKernels (scopeOf lam) + <> types, + distOnInnerMap = onInnerMap path', + distOnTopLevelStms = onTopLevelStms path', + distSegLevel = segThreadCapped, + distOnSOACSStms = pure . oneStm . soacsStmToKernels, + distOnSOACSLambda = pure . soacsLambdaToKernels + } exploitInnerParallelism path' = runDistNestT (env path') $ - distributeMapBodyStms acc (bodyStms $ lambdaBody lam) + distributeMapBodyStms acc (bodyStms $ lambdaBody lam) let exploitOuterParallelism path' = do let lam' = soacsLambdaToKernels lam - runDistNestT (env path') $ distribute $ - addStmsToAcc (bodyStms $ lambdaBody lam') acc + runDistNestT (env path') $ + distribute $ + addStmsToAcc (bodyStms $ lambdaBody lam') acc onMap' (newKernel loopnest) path exploitOuterParallelism exploitInnerParallelism pat lam - where acc = DistAcc { distTargets = singleTarget (pat, bodyResult $ lambdaBody lam) - , distStms = mempty - } + where + acc = + DistAcc + { distTargets = singleTarget (pat, bodyResult $ lambdaBody lam), + distStms = mempty + } onlyExploitIntra :: Attrs -> Bool onlyExploitIntra attrs = @@ -676,27 +721,29 @@ mayExploitOuter :: Attrs -> Bool mayExploitOuter attrs = not $ - AttrComp "incremental_flattening" ["no_outer"] `inAttrs` attrs || - AttrComp "incremental_flattening" ["only_inner"] `inAttrs` attrs + AttrComp "incremental_flattening" ["no_outer"] `inAttrs` attrs + || AttrComp "incremental_flattening" ["only_inner"] `inAttrs` attrs mayExploitIntra :: Attrs -> Bool mayExploitIntra attrs = not $ - AttrComp "incremental_flattening" ["no_intra"] `inAttrs` attrs || - AttrComp "incremental_flattening" ["only_inner"] `inAttrs` attrs + AttrComp "incremental_flattening" ["no_intra"] `inAttrs` attrs + || AttrComp "incremental_flattening" ["only_inner"] `inAttrs` attrs -- The minimum amount of inner parallelism we require (by default) in -- intra-group versions. Less than this is usually pointless on a GPU -- (but we allow tuning to change it). -intraMinInnerPar :: Int32 +intraMinInnerPar :: Int64 intraMinInnerPar = 32 -- One NVIDIA warp -onMap' :: KernelNest -> KernelPath - -> (KernelPath -> DistribM (Out.Stms Out.Kernels)) - -> (KernelPath -> DistribM (Out.Stms Out.Kernels)) - -> Pattern - -> Lambda - -> DistribM (Out.Stms Out.Kernels) +onMap' :: + KernelNest -> + KernelPath -> + (KernelPath -> DistribM (Out.Stms Out.Kernels)) -> + (KernelPath -> DistribM (Out.Stms Out.Kernels)) -> + Pattern -> + Lambda -> + DistribM (Out.Stms Out.Kernels) onMap' loopnest path mk_seq_stms mk_par_stms pat lam = do let nest_ws = kernelNestWidths loopnest res = map Var $ patternNames pat @@ -707,41 +754,49 @@ ((outer_suff, outer_suff_key), outer_suff_stms) <- sufficientParallelism "suff_outer_par" nest_ws path Nothing - intra <- if onlyExploitIntra (stmAuxAttrs aux) || - (worthIntraGroup lam && mayExploitIntra attrs) then - flip runReaderT types $ intraGroupParallelise loopnest lam - else return Nothing - seq_body <- renameBody =<< mkBody <$> - mk_seq_stms ((outer_suff_key, True) : path) <*> pure res - let seq_alts = [(outer_suff, seq_body) - | worthSequentialising lam, mayExploitOuter attrs] + intra <- + if onlyExploitIntra (stmAuxAttrs aux) + || (worthIntraGroup lam && mayExploitIntra attrs) + then flip runReaderT types $ intraGroupParallelise loopnest lam + else return Nothing + seq_body <- + renameBody =<< mkBody + <$> mk_seq_stms ((outer_suff_key, True) : path) <*> pure res + let seq_alts = + [ (outer_suff, seq_body) + | worthSequentialising lam, + mayExploitOuter attrs + ] case intra of Nothing -> do - par_body <- renameBody =<< mkBody <$> - mk_par_stms ((outer_suff_key, False) : path) <*> pure res + par_body <- + renameBody =<< mkBody + <$> mk_par_stms ((outer_suff_key, False) : path) <*> pure res if "sequential_inner" `inAttrs` attrs then kernelAlternatives pat seq_body [] - else (outer_suff_stms<>) <$> kernelAlternatives pat par_body seq_alts - + else (outer_suff_stms <>) <$> kernelAlternatives pat par_body seq_alts Just ((_intra_min_par, intra_avail_par), group_size, log, intra_prelude, intra_stms) -> do addLog log -- We must check that all intra-group parallelism fits in a group. ((intra_ok, intra_suff_key), intra_suff_stms) <- do - ((intra_suff, suff_key), check_suff_stms) <- - sufficientParallelism "suff_intra_par" [intra_avail_par] - ((outer_suff_key, False) : path) (Just intraMinInnerPar) + sufficientParallelism + "suff_intra_par" + [intra_avail_par] + ((outer_suff_key, False) : path) + (Just intraMinInnerPar) runBinder $ do - addStms intra_prelude max_group_size <- letSubExp "max_group_size" $ Op $ SizeOp $ Out.GetSizeMax Out.SizeGroup - fits <- letSubExp "fits" $ BasicOp $ - CmpOp (CmpSle Int32) group_size max_group_size + fits <- + letSubExp "fits" $ + BasicOp $ + CmpOp (CmpSle Int64) group_size max_group_size addStms check_suff_stms @@ -750,36 +805,45 @@ group_par_body <- renameBody $ mkBody intra_stms res - par_body <- renameBody =<< mkBody <$> - mk_par_stms ([(outer_suff_key, False), - (intra_suff_key, False)] - ++ path) <*> pure res + par_body <- + renameBody =<< mkBody + <$> mk_par_stms + ( [ (outer_suff_key, False), + (intra_suff_key, False) + ] + ++ path + ) + <*> pure res if "sequential_inner" `inAttrs` attrs then kernelAlternatives pat seq_body [] else - if onlyExploitIntra attrs - then (intra_suff_stms<>) <$> kernelAlternatives pat group_par_body [] - else ((outer_suff_stms<>intra_suff_stms)<>) <$> - kernelAlternatives pat par_body (seq_alts ++ [(intra_ok, group_par_body)]) + if onlyExploitIntra attrs + then (intra_suff_stms <>) <$> kernelAlternatives pat group_par_body [] + else + ((outer_suff_stms <> intra_suff_stms) <>) + <$> kernelAlternatives pat par_body (seq_alts ++ [(intra_ok, group_par_body)]) -onInnerMap :: KernelPath -> MapLoop -> DistAcc Out.Kernels - -> DistNestT Out.Kernels DistribM (DistAcc Out.Kernels) +onInnerMap :: + KernelPath -> + MapLoop -> + DistAcc Out.Kernels -> + DistNestT Out.Kernels DistribM (DistAcc Out.Kernels) onInnerMap path maploop@(MapLoop pat aux w lam arrs) acc - | unbalancedLambda lam, lambdaContainsParallelism lam = - addStmToAcc (mapLoopStm maploop) acc + | unbalancedLambda lam, + lambdaContainsParallelism lam = + addStmToAcc (mapLoopStm maploop) acc | otherwise = - distributeSingleStm acc (mapLoopStm maploop) >>= \case + distributeSingleStm acc (mapLoopStm maploop) >>= \case Just (post_kernels, res, nest, acc') | Just (perm, _pat_unused) <- permutationAndMissing pat res -> do - addPostStms post_kernels - multiVersion perm nest acc' + addPostStms post_kernels + multiVersion perm nest acc' _ -> distributeMap maploop acc - where discardTargets acc' = -- FIXME: work around bogus targets. - acc' { distTargets = singleTarget (mempty, mempty) } + acc' {distTargets = singleTarget (mempty, mempty)} multiVersion perm nest acc' = do -- The kernel can be distributed by itself, so now we can @@ -788,40 +852,46 @@ dist_env <- ask let extra_scope = targetsScope $ distTargets acc' - stms <- liftInner $ localScope extra_scope $ do - let maploop' = MapLoop pat aux w lam arrs + stms <- liftInner $ + localScope extra_scope $ do + let maploop' = MapLoop pat aux w lam arrs - exploitInnerParallelism path' = do - let dist_env' = - dist_env { distOnTopLevelStms = onTopLevelStms path' - , distOnInnerMap = onInnerMap path' - } - runDistNestT dist_env' $ - inNesting nest $ localScope extra_scope $ - discardTargets <$> distributeMap maploop' acc { distStms = mempty } + exploitInnerParallelism path' = do + let dist_env' = + dist_env + { distOnTopLevelStms = onTopLevelStms path', + distOnInnerMap = onInnerMap path' + } + runDistNestT dist_env' $ + inNesting nest $ + localScope extra_scope $ + discardTargets <$> distributeMap maploop' acc {distStms = mempty} - -- Normally the permutation is for the output pattern, but - -- we can't really change that, so we change the result - -- order instead. - let lam_res' = rearrangeShape perm $ bodyResult $ lambdaBody lam - lam' = lam { lambdaBody = (lambdaBody lam) { bodyResult = lam_res' } } - map_nesting = MapNesting pat aux w $ zip (lambdaParams lam) arrs - nest' = pushInnerKernelNesting (pat, lam_res') map_nesting nest + -- Normally the permutation is for the output pattern, but + -- we can't really change that, so we change the result + -- order instead. + let lam_res' = rearrangeShape perm $ bodyResult $ lambdaBody lam + lam' = lam {lambdaBody = (lambdaBody lam) {bodyResult = lam_res'}} + map_nesting = MapNesting pat aux w $ zip (lambdaParams lam) arrs + nest' = pushInnerKernelNesting (pat, lam_res') map_nesting nest - -- XXX: we do not construct a new KernelPath when - -- sequentialising. This is only OK as long as further - -- versioning does not take place down that branch (it currently - -- does not). - (sequentialised_kernel, nestw_bnds) <- localScope extra_scope $ do - let sequentialised_lam = soacsLambdaToKernels lam' - constructKernel segThreadCapped nest' $ lambdaBody sequentialised_lam + -- XXX: we do not construct a new KernelPath when + -- sequentialising. This is only OK as long as further + -- versioning does not take place down that branch (it currently + -- does not). + (sequentialised_kernel, nestw_bnds) <- localScope extra_scope $ do + let sequentialised_lam = soacsLambdaToKernels lam' + constructKernel segThreadCapped nest' $ lambdaBody sequentialised_lam - let outer_pat = loopNestingPattern $ fst nest - (nestw_bnds<>) <$> - onMap' nest' path - (const $ return $ oneStm sequentialised_kernel) - exploitInnerParallelism - outer_pat lam' + let outer_pat = loopNestingPattern $ fst nest + (nestw_bnds <>) + <$> onMap' + nest' + path + (const $ return $ oneStm sequentialised_kernel) + exploitInnerParallelism + outer_pat + lam' postStm stms return acc' diff --git a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs --- a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +++ b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs @@ -1,45 +1,44 @@ +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE ConstraintKinds #-} -module Futhark.Pass.ExtractKernels.BlockedKernel - ( DistLore - , MkSegLevel - , ThreadRecommendation(..) - , segRed - , nonSegRed - , segScan - , segHist - , segMap - , mapKernel - , KernelInput(..) - , readKernelInput - - , mkSegSpace - , dummyDim - ) - where +module Futhark.Pass.ExtractKernels.BlockedKernel + ( DistLore, + MkSegLevel, + ThreadRecommendation (..), + segRed, + nonSegRed, + segScan, + segHist, + segMap, + mapKernel, + KernelInput (..), + readKernelInput, + mkSegSpace, + dummyDim, + ) +where import Control.Monad import Control.Monad.Writer import Data.List () - -import Prelude hiding (quot) - import Futhark.Analysis.PrimExp import Futhark.IR import Futhark.IR.SegOp import Futhark.MonadFreshNames import Futhark.Tools import Futhark.Transform.Rename +import Prelude hiding (quot) -- | Constraints pertinent to performing distribution/flattening. -type DistLore lore = (Bindable lore, - HasSegOp lore, - BinderOps lore, - LetDec lore ~ Type, - ExpDec lore ~ (), - BodyDec lore ~ ()) +type DistLore lore = + ( Bindable lore, + HasSegOp lore, + BinderOps lore, + LetDec lore ~ Type, + ExpDec lore ~ (), + BodyDec lore ~ () + ) data ThreadRecommendation = ManyThreads | NoRecommendation SegVirt @@ -49,71 +48,86 @@ mkSegSpace :: MonadFreshNames m => [(VName, SubExp)] -> m SegSpace mkSegSpace dims = SegSpace <$> newVName "phys_tid" <*> pure dims -prepareRedOrScan :: (MonadBinder m, DistLore (Lore m)) => - SubExp - -> Lambda (Lore m) - -> [VName] -> [(VName, SubExp)] -> [KernelInput] - -> m (SegSpace, KernelBody (Lore m)) +prepareRedOrScan :: + (MonadBinder m, DistLore (Lore m)) => + SubExp -> + Lambda (Lore m) -> + [VName] -> + [(VName, SubExp)] -> + [KernelInput] -> + m (SegSpace, KernelBody (Lore m)) prepareRedOrScan w map_lam arrs ispace inps = do gtid <- newVName "gtid" space <- mkSegSpace $ ispace ++ [(gtid, w)] - kbody <- fmap (uncurry (flip (KernelBody ()))) $ runBinder $ - localScope (scopeOfSegSpace space) $ do - mapM_ readKernelInput inps - forM_ (zip (lambdaParams map_lam) arrs) $ \(p, arr) -> do - arr_t <- lookupType arr - letBindNames [paramName p] $ - BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid] - map (Returns ResultMaySimplify) <$> bodyBind (lambdaBody map_lam) + kbody <- fmap (uncurry (flip (KernelBody ()))) $ + runBinder $ + localScope (scopeOfSegSpace space) $ do + mapM_ readKernelInput inps + forM_ (zip (lambdaParams map_lam) arrs) $ \(p, arr) -> do + arr_t <- lookupType arr + letBindNames [paramName p] $ + BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid] + map (Returns ResultMaySimplify) <$> bodyBind (lambdaBody map_lam) return (space, kbody) -segRed :: (MonadFreshNames m, DistLore lore, HasScope lore m) => - SegOpLevel lore - -> Pattern lore - -> SubExp -- segment size - -> [SegBinOp lore] - -> Lambda lore - -> [VName] - -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this reduction - -> [KernelInput] -- inps = inputs that can be looked up by using the gtids from ispace - -> m (Stms lore) +segRed :: + (MonadFreshNames m, DistLore lore, HasScope lore m) => + SegOpLevel lore -> + Pattern lore -> + SubExp -> -- segment size + [SegBinOp lore] -> + Lambda lore -> + [VName] -> + [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this reduction + [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace + m (Stms lore) segRed lvl pat w ops map_lam arrs ispace inps = runBinder_ $ do (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps - letBind pat $ Op $ segOp $ - SegRed lvl kspace ops (lambdaReturnType map_lam) kbody + letBind pat $ + Op $ + segOp $ + SegRed lvl kspace ops (lambdaReturnType map_lam) kbody -segScan :: (MonadFreshNames m, DistLore lore, HasScope lore m) => - SegOpLevel lore - -> Pattern lore - -> SubExp -- segment size - -> [SegBinOp lore] -> Lambda lore - -> [VName] - -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this scan - -> [KernelInput] -- inps = inputs that can be looked up by using the gtids from ispace - -> m (Stms lore) +segScan :: + (MonadFreshNames m, DistLore lore, HasScope lore m) => + SegOpLevel lore -> + Pattern lore -> + SubExp -> -- segment size + [SegBinOp lore] -> + Lambda lore -> + [VName] -> + [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this scan + [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace + m (Stms lore) segScan lvl pat w ops map_lam arrs ispace inps = runBinder_ $ do (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps - letBind pat $ Op $ segOp $ - SegScan lvl kspace ops (lambdaReturnType map_lam) kbody + letBind pat $ + Op $ + segOp $ + SegScan lvl kspace ops (lambdaReturnType map_lam) kbody -segMap :: (MonadFreshNames m, DistLore lore, HasScope lore m) => - SegOpLevel lore - -> Pattern lore - -> SubExp -- segment size - -> Lambda lore - -> [VName] - -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this map - -> [KernelInput] -- inps = inputs that can be looked up by using the gtids from ispace - -> m (Stms lore) +segMap :: + (MonadFreshNames m, DistLore lore, HasScope lore m) => + SegOpLevel lore -> + Pattern lore -> + SubExp -> -- segment size + Lambda lore -> + [VName] -> + [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this map + [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace + m (Stms lore) segMap lvl pat w map_lam arrs ispace inps = runBinder_ $ do (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps - letBind pat $ Op $ segOp $ - SegMap lvl kspace (lambdaReturnType map_lam) kbody + letBind pat $ + Op $ + segOp $ + SegMap lvl kspace (lambdaReturnType map_lam) kbody -dummyDim :: (MonadFreshNames m, MonadBinder m, DistLore (Lore m)) => - Pattern (Lore m) - -> m (Pattern (Lore m), [(VName, SubExp)], m ()) +dummyDim :: + (MonadFreshNames m, MonadBinder m, DistLore (Lore m)) => + Pattern (Lore m) -> + m (Pattern (Lore m), [(VName, SubExp)], m ()) dummyDim pat = do -- We add a unit-size segment on top to ensure that the result -- of the SegRed is an array, which we then immediately index. @@ -121,68 +135,83 @@ -- device afterwards, as this may save an expensive -- host-device copy (scalars are kept on the host, but arrays -- may be on the device). - let addDummyDim t = t `arrayOfRow` intConst Int32 1 + let addDummyDim t = t `arrayOfRow` intConst Int64 1 pat' <- fmap addDummyDim <$> renamePattern pat dummy <- newVName "dummy" - let ispace = [(dummy, intConst Int32 1)] + let ispace = [(dummy, intConst Int64 1)] - return (pat', ispace, - forM_ (zip (patternNames pat') (patternNames pat)) $ \(from, to) -> do - from_t <- lookupType from - letBindNames [to] $ BasicOp $ Index from $ - fullSlice from_t [DimFix $ intConst Int32 0]) + return + ( pat', + ispace, + forM_ (zip (patternNames pat') (patternNames pat)) $ \(from, to) -> do + from_t <- lookupType from + letBindNames [to] $ + BasicOp $ + Index from $ + fullSlice from_t [DimFix $ intConst Int64 0] + ) -nonSegRed :: (MonadFreshNames m, DistLore lore, HasScope lore m) => - SegOpLevel lore - -> Pattern lore - -> SubExp - -> [SegBinOp lore] - -> Lambda lore - -> [VName] - -> m (Stms lore) +nonSegRed :: + (MonadFreshNames m, DistLore lore, HasScope lore m) => + SegOpLevel lore -> + Pattern lore -> + SubExp -> + [SegBinOp lore] -> + Lambda lore -> + [VName] -> + m (Stms lore) nonSegRed lvl pat w ops map_lam arrs = runBinder_ $ do (pat', ispace, read_dummy) <- dummyDim pat addStms =<< segRed lvl pat' w ops map_lam arrs ispace [] read_dummy -segHist :: (DistLore lore, MonadFreshNames m, HasScope lore m) => - SegOpLevel lore - -> Pattern lore - -> SubExp - -> [(VName,SubExp)] -- ^ Segment indexes and sizes. - -> [KernelInput] - -> [HistOp lore] - -> Lambda lore -> [VName] - -> m (Stms lore) +segHist :: + (DistLore lore, MonadFreshNames m, HasScope lore m) => + SegOpLevel lore -> + Pattern lore -> + SubExp -> + -- | Segment indexes and sizes. + [(VName, SubExp)] -> + [KernelInput] -> + [HistOp lore] -> + Lambda lore -> + [VName] -> + m (Stms lore) segHist lvl pat arr_w ispace inps ops lam arrs = runBinder_ $ do gtid <- newVName "gtid" space <- mkSegSpace $ ispace ++ [(gtid, arr_w)] - kbody <- fmap (uncurry (flip $ KernelBody ())) $ runBinder $ - localScope (scopeOfSegSpace space) $ do - mapM_ readKernelInput inps - forM_ (zip (lambdaParams lam) arrs) $ \(p, arr) -> do - arr_t <- lookupType arr - letBindNames [paramName p] $ - BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid] - map (Returns ResultMaySimplify) <$> bodyBind (lambdaBody lam) + kbody <- fmap (uncurry (flip $ KernelBody ())) $ + runBinder $ + localScope (scopeOfSegSpace space) $ do + mapM_ readKernelInput inps + forM_ (zip (lambdaParams lam) arrs) $ \(p, arr) -> do + arr_t <- lookupType arr + letBindNames [paramName p] $ + BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid] + map (Returns ResultMaySimplify) <$> bodyBind (lambdaBody lam) letBind pat $ Op $ segOp $ SegHist lvl space ops (lambdaReturnType lam) kbody -mapKernelSkeleton :: (DistLore lore, HasScope lore m, MonadFreshNames m) => - [(VName, SubExp)] -> [KernelInput] - -> m (SegSpace, Stms lore) +mapKernelSkeleton :: + (DistLore lore, HasScope lore m, MonadFreshNames m) => + [(VName, SubExp)] -> + [KernelInput] -> + m (SegSpace, Stms lore) mapKernelSkeleton ispace inputs = do read_input_bnds <- runBinder_ $ mapM readKernelInput inputs space <- mkSegSpace ispace return (space, read_input_bnds) -mapKernel :: (DistLore lore, HasScope lore m, MonadFreshNames m) => - MkSegLevel lore m - -> [(VName, SubExp)] -> [KernelInput] - -> [Type] -> KernelBody lore - -> m (SegOp (SegOpLevel lore) lore, Stms lore) +mapKernel :: + (DistLore lore, HasScope lore m, MonadFreshNames m) => + MkSegLevel lore m -> + [(VName, SubExp)] -> + [KernelInput] -> + [Type] -> + KernelBody lore -> + m (SegOp (SegOpLevel lore) lore, Stms lore) mapKernel mk_lvl ispace inputs rts (KernelBody () kstms krets) = runBinderT' $ do (space, read_input_stms) <- mapKernelSkeleton ispace inputs @@ -199,18 +228,22 @@ return $ SegMap lvl space rts kbody' -data KernelInput = KernelInput { kernelInputName :: VName - , kernelInputType :: Type - , kernelInputArray :: VName - , kernelInputIndices :: [SubExp] - } - deriving (Show) +data KernelInput = KernelInput + { kernelInputName :: VName, + kernelInputType :: Type, + kernelInputArray :: VName, + kernelInputIndices :: [SubExp] + } + deriving (Show) -readKernelInput :: (DistLore (Lore m), MonadBinder m) => - KernelInput -> m () +readKernelInput :: + (DistLore (Lore m), MonadBinder m) => + KernelInput -> + m () readKernelInput inp = do let pe = PatElem (kernelInputName inp) $ kernelInputType inp arr_t <- lookupType $ kernelInputArray inp letBind (Pattern [] [pe]) $ - BasicOp $ Index (kernelInputArray inp) $ - fullSlice arr_t $ map DimFix $ kernelInputIndices inp + BasicOp $ + Index (kernelInputArray inp) $ + fullSlice arr_t $ map DimFix $ kernelInputIndices inp 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 @@ -1,1065 +1,1202 @@ {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE OverloadedStrings #-} -module Futhark.Pass.ExtractKernels.DistributeNests - ( MapLoop(..) - , mapLoopStm - - , bodyContainsParallelism - , lambdaContainsParallelism - , determineReduceOp - , histKernel - - , DistEnv (..) - , DistAcc (..) - , runDistNestT - , DistNestT - , liftInner - - , distributeMap - - , distribute - , distributeSingleStm - , distributeMapBodyStms - , addStmsToAcc - , addStmToAcc - , permutationAndMissing - , addPostStms - , postStm - , inNesting - ) -where - -import Control.Arrow (first) -import Control.Monad.Identity -import Control.Monad.RWS.Strict -import Control.Monad.Reader -import Control.Monad.Writer.Strict -import Control.Monad.Trans.Maybe -import Data.Maybe -import Data.List (find, partition, tails) -import qualified Data.Map as M - -import Futhark.IR -import qualified Futhark.IR.SOACS as SOACS -import Futhark.IR.SOACS.SOAC hiding (HistOp, histDest) -import Futhark.IR.SOACS (SOACS) -import Futhark.IR.SOACS.Simplify (simpleSOACS, simplifyStms) -import Futhark.IR.SegOp -import Futhark.MonadFreshNames -import Futhark.Tools -import Futhark.Transform.Rename -import Futhark.Transform.CopyPropagate -import qualified Futhark.Transform.FirstOrderTransform as FOT -import Futhark.Pass.ExtractKernels.Distribution -import Futhark.Pass.ExtractKernels.ISRWIM -import Futhark.Pass.ExtractKernels.BlockedKernel -import Futhark.Pass.ExtractKernels.Interchange -import Futhark.Util -import Futhark.Util.Log - -scopeForSOACs :: SameScope lore SOACS => Scope lore -> Scope SOACS -scopeForSOACs = castScope - -data MapLoop = MapLoop SOACS.Pattern (StmAux ()) SubExp SOACS.Lambda [VName] - -mapLoopStm :: MapLoop -> Stm SOACS -mapLoopStm (MapLoop pat aux w lam arrs) = - Let pat aux $ Op $ Screma w (mapSOAC lam) arrs - -data DistEnv lore m = - DistEnv { distNest :: Nestings - , distScope :: Scope lore - , distOnTopLevelStms :: Stms SOACS -> DistNestT lore m (Stms lore) - , distOnInnerMap :: MapLoop -> DistAcc lore - -> DistNestT lore m (DistAcc lore) - , distOnSOACSStms :: Stm SOACS -> Binder lore (Stms lore) - , distOnSOACSLambda :: Lambda SOACS -> Binder lore (Lambda lore) - , distSegLevel :: MkSegLevel lore m - } - -data DistAcc lore = - DistAcc { distTargets :: Targets - , distStms :: Stms lore - } - -data DistRes lore = - DistRes { accPostStms :: PostStms lore - , accLog :: Log - } - -instance Semigroup (DistRes lore) where - DistRes ks1 log1 <> DistRes ks2 log2 = - DistRes (ks1 <> ks2) (log1 <> log2) - -instance Monoid (DistRes lore) where - mempty = DistRes mempty mempty - -newtype PostStms lore = PostStms { unPostStms :: Stms lore } - -instance Semigroup (PostStms lore) where - PostStms xs <> PostStms ys = PostStms $ ys <> xs - -instance Monoid (PostStms lore) where - mempty = PostStms mempty - -typeEnvFromDistAcc :: DistLore lore => DistAcc lore -> Scope lore -typeEnvFromDistAcc = scopeOfPattern . fst . outerTarget . distTargets - -addStmsToAcc :: Stms lore -> DistAcc lore -> DistAcc lore -addStmsToAcc stms acc = - acc { distStms = stms <> distStms acc } - -addStmToAcc :: (MonadFreshNames m, DistLore lore) => - Stm SOACS -> DistAcc lore - -> DistNestT lore m (DistAcc lore) -addStmToAcc stm acc = do - onSoacs <- asks distOnSOACSStms - (stm', _) <- runBinder $ onSoacs stm - return acc { distStms = stm' <> distStms acc } - -soacsLambda :: (MonadFreshNames m, DistLore lore) => - Lambda SOACS -> DistNestT lore m (Lambda lore) -soacsLambda lam = do - onLambda <- asks distOnSOACSLambda - fst <$> runBinder (onLambda lam) - -newtype DistNestT lore m a = - DistNestT (ReaderT (DistEnv lore m) (WriterT (DistRes lore) m) a) - deriving (Functor, Applicative, Monad, - MonadReader (DistEnv lore m), - MonadWriter (DistRes lore)) - -liftInner :: (LocalScope lore m, DistLore lore) => m a -> DistNestT lore m a -liftInner m = do - outer_scope <- askScope - DistNestT $ lift $ lift $ do - inner_scope <- askScope - localScope (outer_scope `M.difference` inner_scope) m - -instance MonadFreshNames m => MonadFreshNames (DistNestT lore m) where - getNameSource = DistNestT $ lift getNameSource - putNameSource = DistNestT . lift . putNameSource - -instance (Monad m, ASTLore lore) => HasScope lore (DistNestT lore m) where - askScope = asks distScope - -instance (Monad m, ASTLore lore) => LocalScope lore (DistNestT lore m) where - localScope types = local $ \env -> - env { distScope = types <> distScope env } - -instance Monad m => MonadLogger (DistNestT lore m) where - addLog msgs = tell mempty { accLog = msgs } - -runDistNestT :: (MonadLogger m, DistLore lore) => - DistEnv lore m -> DistNestT lore m (DistAcc lore) -> m (Stms lore) -runDistNestT env (DistNestT m) = do - (acc, res) <- runWriterT $ runReaderT m env - addLog $ accLog res - -- There may be a few final targets remaining - these correspond to - -- arrays that are identity mapped, and must have statements - -- inserted here. - return $ - unPostStms (accPostStms res) <> - identityStms (outerTarget $ distTargets acc) - where outermost = nestingLoop $ - case distNest env of (nest, []) -> nest - (_, nest : _) -> nest - params_to_arrs = map (first paramName) $ - loopNestingParamsAndArrs outermost - - identityStms (rem_pat, res) = - stmsFromList $ zipWith identityStm (patternValueElements rem_pat) res - identityStm pe (Var v) - | Just arr <- lookup v params_to_arrs = - Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ Copy arr - identityStm pe se = - Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ - Replicate (Shape [loopNestingWidth outermost]) se - -addPostStms :: Monad m => PostStms lore -> DistNestT lore m () -addPostStms ks = tell $ mempty { accPostStms = ks } - -postStm :: Monad m => Stms lore -> DistNestT lore m () -postStm stms = addPostStms $ PostStms stms - -withStm :: (Monad m, DistLore lore) => - Stm SOACS -> DistNestT lore m a -> DistNestT lore m a -withStm stm = local $ \env -> - env { distScope = - castScope (scopeOf stm) <> distScope env - , distNest = - letBindInInnerNesting provided $ - distNest env - } - where provided = namesFromList $ patternNames $ stmPattern stm - -leavingNesting :: (MonadFreshNames m, DistLore lore) => - DistAcc lore -> DistNestT lore m (DistAcc lore) - -leavingNesting acc = - case popInnerTarget $ distTargets acc of - Nothing -> - error "The kernel targets list is unexpectedly small" - - Just ((pat, res), newtargets) - | not $ null $ distStms acc -> do - -- Any statements left over correspond to something that - -- could not be distributed because it would cause irregular - -- arrays. These must be reconstructed into a a Map SOAC - -- that will be sequentialised. XXX: life would be better if - -- we were able to distribute irregular parallelism. - (Nesting _ inner, _) <- asks distNest - let MapNesting _ aux w params_and_arrs = inner - body = Body () (distStms acc) res - used_in_body = freeIn body - (used_params, used_arrs) = - unzip $ - filter ((`nameIn` used_in_body) . paramName . fst) params_and_arrs - lam' = Lambda { lambdaParams = used_params - , lambdaBody = body - , lambdaReturnType = map rowType $ patternTypes pat - } - stms <- runBinder_ $ auxing aux $ FOT.transformSOAC pat $ - Screma w (mapSOAC lam') used_arrs - - return $ acc { distTargets = newtargets, distStms = stms } - - | otherwise -> do - -- Any results left over correspond to a Replicate or a Copy in - -- the parent nesting, depending on whether the argument is a - -- parameter of the innermost nesting. - (Nesting _ inner_nesting, _) <- asks distNest - let w = loopNestingWidth inner_nesting - aux = loopNestingAux inner_nesting - inps = loopNestingParamsAndArrs inner_nesting - - remnantStm pe (Var v) - | Just (_, arr) <- find ((==v) . paramName . fst) inps = - Let (Pattern [] [pe]) aux $ - BasicOp $ Copy arr - remnantStm pe se = - Let (Pattern [] [pe]) aux $ - BasicOp $ Replicate (Shape [w]) se - - stms = - stmsFromList $ zipWith remnantStm (patternElements pat) res - - return $ acc { distTargets = newtargets, distStms = stms } - -mapNesting :: (MonadFreshNames m, DistLore lore) => - PatternT Type -> StmAux () -> SubExp -> Lambda SOACS -> [VName] - -> DistNestT lore m (DistAcc lore) - -> DistNestT lore m (DistAcc lore) -mapNesting pat aux w lam arrs m = - local extend $ leavingNesting =<< m - where nest = Nesting mempty $ - MapNesting pat aux w $ - zip (lambdaParams lam) arrs - extend env = env { distNest = pushInnerNesting nest $ distNest env - , distScope = castScope (scopeOf lam) <> distScope env - } - -inNesting :: (Monad m, DistLore lore) => - KernelNest -> DistNestT lore m a -> DistNestT lore m a -inNesting (outer, nests) = local $ \env -> - env { distNest = (inner, nests') - , distScope = foldMap scopeOfLoopNesting (outer : nests) <> distScope env - } - where (inner, nests') = - case reverse nests of - [] -> (asNesting outer, []) - (inner' : ns) -> (asNesting inner', map asNesting $ outer : reverse ns) - asNesting = Nesting mempty - -bodyContainsParallelism :: Body SOACS -> Bool -bodyContainsParallelism = any isParallelStm . bodyStms - where isParallelStm stm = - isMap (stmExp stm) && - not ("sequential" `inAttrs` stmAuxAttrs (stmAux stm)) - isMap Op{} = True - isMap _ = False - -lambdaContainsParallelism :: Lambda SOACS -> Bool -lambdaContainsParallelism = bodyContainsParallelism . lambdaBody - -distributeMapBodyStms :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> Stms SOACS -> DistNestT lore m (DistAcc lore) -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 - types <- asksScope scopeForSOACs - stream_stms <- - snd <$> runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types - (_, stream_stms') <- - runReaderT (copyPropagateInStms simpleSOACS types stream_stms) types - onStms acc $ stmsToList (fmap (certify cs) stream_stms') ++ stms - - onStms acc (stm:stms) = - -- It is important that stm is in scope if 'maybeDistributeStm' - -- wants to distribute, even if this causes the slightly silly - -- situation that stm is in scope of itself. - withStm stm $ maybeDistributeStm stm =<< onStms acc stms - -onInnerMap :: Monad m => MapLoop -> DistAcc lore -> DistNestT lore m (DistAcc lore) -onInnerMap loop acc = do - f <- asks distOnInnerMap - f loop acc - -onTopLevelStms :: Monad m => Stms SOACS -> DistNestT lore m () -onTopLevelStms stms = do - f <- asks distOnTopLevelStms - postStm =<< f stms - -maybeDistributeStm :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - Stm SOACS -> DistAcc lore - -> DistNestT lore m (DistAcc lore) - -maybeDistributeStm stm acc - | "sequential" `inAttrs` stmAuxAttrs (stmAux stm) = - addStmToAcc stm acc - -maybeDistributeStm (Let pat aux (Op soac)) acc - | "sequential_outer" `inAttrs` stmAuxAttrs aux = - distributeMapBodyStms acc . fmap (certify (stmAuxCerts aux)) =<< - runBinder_ (FOT.transformSOAC pat soac) - -maybeDistributeStm stm@(Let pat _ (Op (Screma w form arrs))) acc - | Just lam <- isMapSOAC form = - -- Only distribute inside the map if we can distribute everything - -- following the map. - distributeIfPossible acc >>= \case - Nothing -> addStmToAcc stm acc - Just acc' -> distribute =<< onInnerMap (MapLoop pat (stmAux stm) w lam arrs) acc' - -maybeDistributeStm bnd@(Let pat _ (DoLoop [] val form@ForLoop{} body)) acc - | null (patternContextElements pat), bodyContainsParallelism body = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | not $ freeIn form `namesIntersect` boundInKernelNest nest, - Just (perm, pat_unused) <- permutationAndMissing pat res -> - -- We need to pretend pat_unused was used anyway, by adding - -- it to the kernel nest. - localScope (typeEnvFromDistAcc acc') $ do - addPostStms kernels - nest' <- expandKernelNest pat_unused nest - types <- asksScope scopeForSOACs - - -- Simplification is key to hoisting out statements that - -- were variant to the loop, but invariant to the outer maps - -- (which are now innermost). - stms <- - (`runReaderT` types) $ - fmap snd . simplifyStms =<< - interchangeLoops nest' (SeqLoop perm pat val form body) - onTopLevelStms stms - return acc' - _ -> - addStmToAcc bnd acc - -maybeDistributeStm stm@(Let pat _ (If cond tbranch fbranch ret)) acc - | null (patternContextElements pat), - bodyContainsParallelism tbranch || bodyContainsParallelism fbranch || - not (all primType (ifReturns ret)) = - distributeSingleStm acc stm >>= \case - Just (kernels, res, nest, acc') - | not $ - (freeIn cond <> freeIn ret) `namesIntersect` boundInKernelNest nest, - Just (perm, pat_unused) <- permutationAndMissing pat res -> - -- We need to pretend pat_unused was used anyway, by adding - -- it to the kernel nest. - localScope (typeEnvFromDistAcc acc') $ do - nest' <- expandKernelNest pat_unused nest - addPostStms kernels - types <- asksScope scopeForSOACs - let branch = Branch perm pat cond tbranch fbranch ret - stms <- - (`runReaderT` types) $ - fmap snd . simplifyStms =<< - interchangeBranch nest' branch - onTopLevelStms stms - return acc' - _ -> - addStmToAcc stm acc - -maybeDistributeStm (Let pat aux (Op (Screma w form arrs))) acc - | Just [Reduce comm lam nes] <- isReduceSOAC form, - Just m <- irwim pat w comm lam $ zip nes arrs = do - types <- asksScope scopeForSOACs - (_, bnds) <- runBinderT (auxing aux m) types - distributeMapBodyStms acc bnds - --- Parallelise segmented scatters. -maybeDistributeStm bnd@(Let pat (StmAux cs _ _) (Op (Scatter w lam ivs as))) acc = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | Just (perm, pat_unused) <- permutationAndMissing pat res -> - localScope (typeEnvFromDistAcc acc') $ do - nest' <- expandKernelNest pat_unused nest - lam' <- soacsLambda lam - addPostStms kernels - postStm =<< segmentedScatterKernel nest' perm pat cs w lam' ivs as - return acc' - _ -> - addStmToAcc bnd acc - --- Parallelise segmented Hist. -maybeDistributeStm bnd@(Let pat (StmAux cs _ _) (Op (Hist w ops lam as))) acc = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | Just (perm, pat_unused) <- permutationAndMissing pat res -> - localScope (typeEnvFromDistAcc acc') $ do - lam' <- soacsLambda lam - nest' <- expandKernelNest pat_unused nest - addPostStms kernels - postStm =<< segmentedHistKernel nest' perm cs w ops lam' as - return acc' - _ -> - addStmToAcc bnd acc - --- Parallelise Index slices if the result is going to be returned --- directly from the kernel. This is because we would otherwise have --- to sequentialise writing the result, which may be costly. -maybeDistributeStm stm@(Let (Pattern [] [pe]) - aux - (BasicOp (Index arr slice))) acc - | not $ null $ sliceDims slice, - Var (patElemName pe) `elem` snd (innerTarget (distTargets acc)) = - distributeSingleStm acc stm >>= \case - Just (kernels, _res, nest, acc') -> - localScope (typeEnvFromDistAcc acc') $ do - addPostStms kernels - postStm =<< segmentedGatherKernel nest (stmAuxCerts aux) arr slice - return acc' - _ -> - addStmToAcc stm acc - --- If the scan can be distributed by itself, we will turn it into a --- segmented scan. --- --- If the scan cannot be distributed by itself, it will be --- sequentialised in the default case for this function. -maybeDistributeStm bnd@(Let pat (StmAux cs _ _) (Op (Screma w form arrs))) acc - | Just (scans, map_lam) <- isScanomapSOAC form, - Scan lam nes <- singleScan scans = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | Just (perm, pat_unused) <- permutationAndMissing pat res -> - -- We need to pretend pat_unused was used anyway, by adding - -- it to the kernel nest. - localScope (typeEnvFromDistAcc acc') $ do - nest' <- expandKernelNest pat_unused nest - map_lam' <- soacsLambda map_lam - lam' <- soacsLambda lam - localScope (typeEnvFromDistAcc acc') $ - segmentedScanomapKernel nest' perm w lam' map_lam' nes arrs >>= - kernelOrNot cs bnd acc kernels acc' - _ -> - addStmToAcc bnd acc - --- if the reduction can be distributed by itself, we will turn it into a --- segmented reduce. --- --- If the reduction cannot be distributed by itself, it will be --- sequentialised in the default case for this function. -maybeDistributeStm bnd@(Let pat (StmAux cs _ _) (Op (Screma w form arrs))) acc - | Just (reds, map_lam) <- isRedomapSOAC form, - Reduce comm lam nes <- singleReduce reds = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | Just (perm, pat_unused) <- permutationAndMissing pat res -> - -- We need to pretend pat_unused was used anyway, by adding - -- it to the kernel nest. - localScope (typeEnvFromDistAcc acc') $ do - nest' <- expandKernelNest pat_unused nest - - lam' <- soacsLambda lam - map_lam' <- soacsLambda map_lam - - let comm' | commutativeLambda lam = Commutative - | otherwise = comm - - regularSegmentedRedomapKernel nest' perm w comm' lam' map_lam' nes arrs >>= - kernelOrNot cs bnd acc kernels acc' - _ -> - addStmToAcc bnd acc - -maybeDistributeStm (Let pat (StmAux cs _ _) (Op (Screma w form arrs))) acc = do - -- This Screma is too complicated for us to immediately do - -- anything, so split it up and try again. - scope <- asksScope scopeForSOACs - distributeMapBodyStms acc . fmap (certify cs) . snd =<< - runBinderT (dissectScrema pat w form arrs) scope - -maybeDistributeStm (Let pat aux (BasicOp (Replicate (Shape (d:ds)) v))) acc - | [t] <- patternTypes pat = do - tmp <- newVName "tmp" - let rowt = rowType t - newbnd = Let pat aux $ Op $ Screma d (mapSOAC lam) [] - tmpbnd = Let (Pattern [] [PatElem tmp rowt]) aux $ - BasicOp $ Replicate (Shape ds) v - lam = Lambda { lambdaReturnType = [rowt] - , lambdaParams = [] - , lambdaBody = mkBody (oneStm tmpbnd) [Var tmp] - } - maybeDistributeStm newbnd acc - -maybeDistributeStm bnd@(Let _ aux (BasicOp Copy{})) acc = - distributeSingleUnaryStm acc bnd $ \_ outerpat arr -> - return $ oneStm $ Let outerpat aux $ BasicOp $ Copy arr - --- Opaques are applied to the full array, because otherwise they can --- drastically inhibit parallelisation in some cases. -maybeDistributeStm bnd@(Let (Pattern [] [pe]) aux (BasicOp Opaque{})) acc - | not $ primType $ typeOf pe = - distributeSingleUnaryStm acc bnd $ \_ outerpat arr -> - return $ oneStm $ Let outerpat aux $ BasicOp $ Copy arr - -maybeDistributeStm bnd@(Let _ aux (BasicOp (Rearrange perm _))) acc = - distributeSingleUnaryStm acc bnd $ \nest outerpat arr -> do - let r = length (snd nest) + 1 - perm' = [0..r-1] ++ map (+r) perm - -- We need to add a copy, because the original map nest - -- will have produced an array without aliases, and so must we. - arr' <- newVName $ baseString arr - arr_t <- lookupType arr - return $ stmsFromList - [Let (Pattern [] [PatElem arr' arr_t]) aux $ BasicOp $ Copy arr, - Let outerpat aux $ BasicOp $ Rearrange perm' arr'] - -maybeDistributeStm bnd@(Let _ aux (BasicOp (Reshape reshape _))) acc = - distributeSingleUnaryStm acc bnd $ \nest outerpat arr -> do - let reshape' = map DimNew (kernelNestWidths nest) ++ - map DimNew (newDims reshape) - return $ oneStm $ Let outerpat aux $ BasicOp $ Reshape reshape' arr - -maybeDistributeStm stm@(Let _ aux (BasicOp (Rotate rots _))) acc = - distributeSingleUnaryStm acc stm $ \nest outerpat arr -> do - let rots' = map (const $ intConst Int32 0) (kernelNestWidths nest) ++ rots - return $ oneStm $ Let outerpat aux $ BasicOp $ Rotate rots' arr - -maybeDistributeStm stm@(Let pat aux (BasicOp (Update arr slice (Var v)))) acc - | not $ null $ sliceDims slice = - distributeSingleStm acc stm >>= \case - Just (kernels, res, nest, acc') - | res == map Var (patternNames $ stmPattern stm), - Just (perm, pat_unused) <- permutationAndMissing pat res -> do - addPostStms kernels - localScope (typeEnvFromDistAcc acc') $ do - nest' <- expandKernelNest pat_unused nest - postStm =<< - segmentedUpdateKernel nest' perm (stmAuxCerts aux) arr slice v - return acc' - - _ -> addStmToAcc stm acc - --- XXX? This rule is present to avoid the case where an in-place --- update is distributed as its own kernel, as this would mean thread --- then writes the entire array that it updated. This is problematic --- because the in-place updates is O(1), but writing the array is --- O(n). It is OK if the in-place update is preceded, followed, or --- nested inside a sequential loop or similar, because that will --- probably be O(n) by itself. As a hack, we only distribute if there --- does not appear to be a loop following. The better solution is to --- depend on memory block merging for this optimisation, but it is not --- ready yet. -maybeDistributeStm (Let pat aux (BasicOp (Update arr [DimFix i] v))) acc - | [t] <- patternTypes pat, - arrayRank t == 1, - not $ any (amortises . stmExp) $ distStms acc = do - let w = arraySize 0 t - et = stripArray 1 t - lam = Lambda { lambdaParams = [] - , lambdaReturnType = [Prim int32, et] - , lambdaBody = mkBody mempty [i, v] } - maybeDistributeStm (Let pat aux $ Op $ Scatter (intConst Int32 1) lam [] [(w, 1, arr)]) acc - where amortises DoLoop{} = True - amortises Op{} = True - amortises _ = False - -maybeDistributeStm stm@(Let _ aux (BasicOp (Concat d x xs w))) acc = - distributeSingleStm acc stm >>= \case - Just (kernels, _, nest, acc') -> - localScope (typeEnvFromDistAcc acc') $ - segmentedConcat nest >>= - kernelOrNot (stmAuxCerts aux) stm acc kernels acc' - _ -> - addStmToAcc stm acc - - where segmentedConcat nest = - isSegmentedOp nest [0] mempty mempty [] (x:xs) $ - \pat _ _ _ (x':xs') -> - let d' = d + length (snd nest) + 1 - in addStm $ Let pat aux $ BasicOp $ Concat d' x' xs' w - -maybeDistributeStm bnd acc = - addStmToAcc bnd acc - -distributeSingleUnaryStm :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> Stm SOACS - -> (KernelNest -> PatternT Type -> VName -> DistNestT lore m (Stms lore)) - -> DistNestT lore m (DistAcc lore) -distributeSingleUnaryStm acc bnd f = - distributeSingleStm acc bnd >>= \case - Just (kernels, res, nest, acc') - | res == map Var (patternNames $ stmPattern bnd), - (outer, _) <- nest, - [(arr_p, arr)] <- loopNestingParamsAndArrs outer, - boundInKernelNest nest `namesIntersection` freeIn bnd - == oneName (paramName arr_p), - perfectlyMapped arr nest -> do - addPostStms kernels - let outerpat = loopNestingPattern $ fst nest - localScope (typeEnvFromDistAcc acc') $ do - postStm =<< f nest outerpat arr - return acc' - _ -> addStmToAcc bnd acc - - where perfectlyMapped arr (outer, nest) - | [(p, arr')] <- loopNestingParamsAndArrs outer, - arr == arr' = - case nest of [] -> True - x:xs -> perfectlyMapped (paramName p) (x, xs) - | otherwise = - False - -distribute :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> DistNestT lore m (DistAcc lore) -distribute acc = - fromMaybe acc <$> distributeIfPossible acc - -mkSegLevel :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistNestT lore m (MkSegLevel lore (DistNestT lore m)) -mkSegLevel = do - mk_lvl <- asks distSegLevel - return $ \w desc r -> do - (lvl, stms) <- lift $ liftInner $ runBinderT' $ mk_lvl w desc r - addStms stms - return lvl - -distributeIfPossible :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> DistNestT lore m (Maybe (DistAcc lore)) -distributeIfPossible acc = do - nest <- asks distNest - mk_lvl <- mkSegLevel - tryDistribute mk_lvl nest (distTargets acc) (distStms acc) >>= \case - Nothing -> return Nothing - Just (targets, kernel) -> do - postStm kernel - return $ Just DistAcc { distTargets = targets - , distStms = mempty - } - -distributeSingleStm :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> Stm SOACS - -> DistNestT lore m (Maybe (PostStms lore, - Result, - KernelNest, - DistAcc lore)) -distributeSingleStm acc bnd = do - nest <- asks distNest - mk_lvl <- mkSegLevel - tryDistribute mk_lvl nest (distTargets acc) (distStms acc) >>= \case - Nothing -> return Nothing - Just (targets, distributed_bnds) -> - tryDistributeStm nest targets bnd >>= \case - Nothing -> return Nothing - Just (res, targets', new_kernel_nest) -> - return $ Just (PostStms distributed_bnds, - res, - new_kernel_nest, - DistAcc { distTargets = targets' - , distStms = mempty - }) - -segmentedScatterKernel :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - KernelNest - -> [Int] - -> PatternT Type - -> Certificates - -> SubExp - -> Lambda lore - -> [VName] -> [(SubExp,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 - -- things are different because a scatter is not a reduction or - -- scan. - -- - -- First, pretend that the scatter is also part of the nesting. The - -- KernelNest we produce here is technically not sensible, but it's - -- good enough for flatKernel to work. - let nesting = - MapNesting scatter_pat (StmAux cs mempty ()) scatter_w $ zip (lambdaParams lam) ivs - nest' = - pushInnerKernelNesting (scatter_pat, bodyResult $ lambdaBody lam) nesting nest - (ispace, kernel_inps) <- flatKernel nest' - - let (as_ws, as_ns, as) = unzip3 dests - - -- The input/output arrays ('as') _must_ correspond to some kernel - -- input, or else the original nested scatter would have been - -- ill-typed. Find them. - as_inps <- mapM (findInput kernel_inps) as - - mk_lvl <- mkSegLevel - - let rts = concatMap (take 1) $ chunks as_ns $ - drop (sum as_ns) $ lambdaReturnType lam - (is,vs) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam - - -- Maybe add certificates to the indices. - (is', k_body_stms) <- runBinder $ do - addStms $ bodyStms $ lambdaBody lam - forM is $ \i -> - if cs == mempty - then return i - 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 - - (k, k_bnds) <- mapKernel mk_lvl ispace kernel_inps rts k_body - - traverse renameStm <=< runBinder_ $ do - addStms k_bnds - - let pat = Pattern [] $ rearrangeShape perm $ - patternValueElements $ loopNestingPattern $ fst nest - - letBind pat $ Op $ segOp k - where findInput kernel_inps a = - maybe bad return $ find ((==a) . kernelInputName) kernel_inps - bad = error "Ill-typed nested scatter encountered." - - inPlaceReturn ispace (aw, inp, is_vs) = - WriteReturns (init ws++[aw]) (kernelInputArray inp) - [ (map DimFix $ map Var (init gtids)++[i], v) | (i,v) <- is_vs ] - where (gtids,ws) = unzip ispace - -segmentedUpdateKernel :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - KernelNest - -> [Int] - -> Certificates - -> VName - -> Slice SubExp - -> VName - -> DistNestT lore m (Stms lore) -segmentedUpdateKernel nest perm cs arr slice v = do - (base_ispace, kernel_inps) <- flatKernel nest - let slice_dims = sliceDims slice - slice_gtids <- replicateM (length slice_dims) (newVName "gtid_slice") - - let ispace = base_ispace ++ zip slice_gtids slice_dims - - ((res_t, res), kstms) <- runBinder $ do - - -- Compute indexes into full array. - v' <- certifying cs $ - letSubExp "v" $ BasicOp $ Index v $ map (DimFix . Var) slice_gtids - let pexp = primExpFromSubExp int32 - slice_is <- traverse (toSubExp "index") $ - fixSlice (map (fmap pexp) slice) $ map (pexp . Var) slice_gtids - - let write_is = map (Var . fst) base_ispace ++ slice_is - arr' = maybe (error "incorrectly typed Update") kernelInputArray $ - find ((==arr) . kernelInputName) kernel_inps - arr_t <- lookupType arr' - v_t <- subExpType v' - return (v_t, - WriteReturns (arrayDims arr_t) arr' [(map DimFix write_is, v')]) - - mk_lvl <- mkSegLevel - (k, prestms) <- mapKernel mk_lvl ispace kernel_inps [res_t] $ - KernelBody () kstms [res] - - traverse renameStm <=< runBinder_ $ do - addStms prestms - - let pat = Pattern [] $ rearrangeShape perm $ - patternValueElements $ loopNestingPattern $ fst nest - - letBind pat $ Op $ segOp k - -segmentedGatherKernel :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - KernelNest - -> Certificates - -> VName - -> Slice SubExp - -> DistNestT lore m (Stms lore) -segmentedGatherKernel nest cs arr slice = do - let slice_dims = sliceDims slice - slice_gtids <- replicateM (length slice_dims) (newVName "gtid_slice") - - (base_ispace, kernel_inps) <- flatKernel nest - let ispace = base_ispace ++ zip slice_gtids slice_dims - - ((res_t, res), kstms) <- runBinder $ do - -- Compute indexes into full array. - slice'' <- subExpSlice $ - sliceSlice (primExpSlice slice) $ - primExpSlice $ map (DimFix . Var) slice_gtids - v' <- certifying cs $ letSubExp "v" $ BasicOp $ Index arr slice'' - v_t <- subExpType v' - return (v_t, Returns ResultMaySimplify v') - - mk_lvl <- mkSegLevel - (k, prestms) <- mapKernel mk_lvl ispace kernel_inps [res_t] $ - KernelBody () kstms [res] - - traverse renameStm <=< runBinder_ $ do - addStms prestms - - let pat = Pattern [] $ patternValueElements $ loopNestingPattern $ fst nest - - letBind pat $ Op $ segOp k - -segmentedHistKernel :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - KernelNest - -> [Int] - -> Certificates - -> SubExp - -> [SOACS.HistOp SOACS] - -> Lambda lore - -> [VName] - -> DistNestT lore m (Stms lore) -segmentedHistKernel nest perm cs hist_w ops lam arrs = do - -- We replicate some of the checking done by 'isSegmentedOp', but - -- things are different because a Hist is not a reduction or - -- scan. - (ispace, inputs) <- flatKernel nest - let orig_pat = Pattern [] $ rearrangeShape perm $ - patternValueElements $ loopNestingPattern $ fst nest - - -- The input/output arrays _must_ correspond to some kernel input, - -- or else the original nested Hist would have been ill-typed. - -- Find them. - ops' <- forM ops $ \(SOACS.HistOp num_bins rf dests nes op) -> - SOACS.HistOp num_bins rf - <$> mapM (fmap kernelInputArray . findInput inputs) dests - <*> pure nes - <*> pure op - - mk_lvl <- asks distSegLevel - onLambda <- asks distOnSOACSLambda - let onLambda' = fmap fst . runBinder . onLambda - liftInner $ runBinderT'_ $ do - -- It is important not to launch unnecessarily many threads for - -- histograms, because it may mean we unnecessarily need to reduce - -- subhistograms as well. - lvl <- mk_lvl (hist_w : map snd ispace) "seghist" $ NoRecommendation SegNoVirt - addStms =<< - histKernel onLambda' lvl orig_pat ispace inputs cs hist_w ops' lam arrs - where findInput kernel_inps a = - maybe bad return $ find ((==a) . kernelInputName) kernel_inps - bad = error "Ill-typed nested Hist encountered." - -histKernel :: (MonadBinder m, DistLore (Lore m)) => - (Lambda SOACS -> m (Lambda (Lore m))) - -> SegOpLevel (Lore m) - -> PatternT Type -> [(VName, SubExp)] -> [KernelInput] - -> Certificates -> SubExp -> [SOACS.HistOp SOACS] - -> Lambda (Lore m) -> [VName] - -> m (Stms (Lore m)) -histKernel onLambda lvl orig_pat ispace inputs cs hist_w ops lam arrs = runBinderT'_ $ do - ops' <- forM ops $ \(SOACS.HistOp num_bins rf dests nes op) -> do - (op', nes', shape) <- determineReduceOp op nes - op'' <- lift $ onLambda op' - return $ HistOp num_bins rf dests nes' shape op'' - - let isDest = flip elem $ concatMap histDest ops' - inputs' = filter (not . isDest . kernelInputArray) inputs - - certifying cs $ - addStms =<< traverse renameStm =<< - segHist lvl orig_pat hist_w ispace inputs' ops' lam arrs - -determineReduceOp :: (MonadBinder m, Lore m ~ lore) => - Lambda SOACS -> [SubExp] - -> m (Lambda SOACS, [SubExp], Shape) -determineReduceOp lam nes = - -- FIXME? We are assuming that the accumulator is a replicate, and - -- we fish out its value in a gross way. - case mapM subExpVar nes of - Just ne_vs' -> do - let (shape, lam') = isVectorMap lam - nes' <- forM ne_vs' $ \ne_v -> do - ne_v_t <- lookupType ne_v - letSubExp "hist_ne" $ - BasicOp $ Index ne_v $ fullSlice ne_v_t $ - replicate (shapeRank shape) $ DimFix $ intConst Int32 0 - return (lam', nes', shape) - - Nothing -> - return (lam, nes, mempty) - -isVectorMap :: Lambda SOACS -> (Shape, Lambda SOACS) -isVectorMap lam - | [Let (Pattern [] pes) _ (Op (Screma w form arrs))] <- - stmsToList $ bodyStms $ lambdaBody lam, - bodyResult (lambdaBody lam) == map (Var . patElemName) pes, - Just map_lam <- isMapSOAC form, - arrs == map paramName (lambdaParams lam) = - let (shape, lam') = isVectorMap map_lam - in (Shape [w] <> shape, lam') - | otherwise = (mempty, lam) - -segmentedScanomapKernel :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - KernelNest - -> [Int] - -> SubExp - -> Lambda lore -> Lambda lore - -> [SubExp] -> [VName] - -> DistNestT lore m (Maybe (Stms lore)) -segmentedScanomapKernel nest perm segment_size lam map_lam nes arrs = do - mk_lvl <- asks distSegLevel - isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $ - \pat ispace inps nes' _ -> do - let scan_op = SegBinOp Noncommutative lam nes' mempty - lvl <- mk_lvl (segment_size : map snd ispace) "segscan" $ NoRecommendation SegNoVirt - addStms =<< traverse renameStm =<< - segScan lvl pat segment_size [scan_op] map_lam arrs ispace inps - -regularSegmentedRedomapKernel :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - KernelNest - -> [Int] - -> SubExp -> Commutativity - -> Lambda lore -> Lambda lore - -> [SubExp] -> [VName] - -> DistNestT lore m (Maybe (Stms lore)) -regularSegmentedRedomapKernel nest perm segment_size comm lam map_lam nes arrs = do - mk_lvl <- asks distSegLevel - isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $ - \pat ispace inps nes' _ -> do - let red_op = SegBinOp comm lam nes' mempty - lvl <- mk_lvl (segment_size : map snd ispace) "segred" $ NoRecommendation SegNoVirt - addStms =<< traverse renameStm =<< - segRed lvl pat segment_size [red_op] map_lam arrs ispace inps - -isSegmentedOp :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - KernelNest - -> [Int] - -> Names -> Names - -> [SubExp] -> [VName] - -> (PatternT Type - -> [(VName, SubExp)] - -> [KernelInput] - -> [SubExp] -> [VName] - -> BinderT lore m ()) - -> DistNestT lore m (Maybe (Stms lore)) -isSegmentedOp nest perm free_in_op _free_in_fold_op nes arrs m = runMaybeT $ do - -- We must verify that array inputs to the operation are inputs to - -- the outermost loop nesting or free in the loop nest. Nothing - -- free in the op may be bound by the nest. Furthermore, the - -- neutral elements must be free in the loop nest. - -- - -- We must summarise any names from free_in_op that are bound in the - -- nest, and describe how to obtain them given segment indices. - - let bound_by_nest = boundInKernelNest nest - - (ispace, kernel_inps) <- flatKernel nest - - when (free_in_op `namesIntersect` bound_by_nest) $ - fail "Non-fold lambda uses nest-bound parameters." - - let indices = map fst ispace - - prepareNe (Var v) | v `nameIn` bound_by_nest = - fail "Neutral element bound in nest" - prepareNe ne = return ne - - prepareArr arr = - case find ((==arr) . kernelInputName) kernel_inps of - Just inp - | kernelInputIndices inp == map Var indices -> - return $ return $ kernelInputArray inp - Nothing | not (arr `nameIn` bound_by_nest) -> - -- This input is something that is free inside - -- the loop nesting. We will have to replicate - -- it. - return $ - letExp (baseString arr ++ "_repd") - (BasicOp $ Replicate (Shape $ map snd ispace) $ Var arr) - _ -> - fail "Input not free, perfectly mapped, or outermost." - - nes' <- mapM prepareNe nes - - mk_arrs <- mapM prepareArr arrs - - lift $ liftInner $ runBinderT'_ $ do - nested_arrs <- sequence mk_arrs - - let pat = Pattern [] $ rearrangeShape perm $ - patternValueElements $ loopNestingPattern $ fst nest - - m pat ispace kernel_inps nes' nested_arrs - -permutationAndMissing :: PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type]) -permutationAndMissing pat res = do - let pes = patternValueElements pat - (_used,unused) = - partition ((`nameIn` freeIn res) . patElemName) pes - res_expanded = res ++ map (Var . patElemName) unused - perm <- map (Var . patElemName) pes `isPermutationOf` res_expanded - return (perm, unused) - --- Add extra pattern elements to every kernel nesting level. -expandKernelNest :: MonadFreshNames m => - [PatElemT Type] -> KernelNest -> m KernelNest -expandKernelNest pes (outer_nest, inner_nests) = do - let outer_size = loopNestingWidth outer_nest : - map loopNestingWidth inner_nests - inner_sizes = tails $ map loopNestingWidth inner_nests - outer_nest' <- expandWith outer_nest outer_size - inner_nests' <- zipWithM expandWith inner_nests inner_sizes - return (outer_nest', inner_nests') - where expandWith nest dims = do - pes' <- mapM (expandPatElemWith dims) pes - return nest { loopNestingPattern = - Pattern [] $ - patternElements (loopNestingPattern nest) <> pes' - } - - expandPatElemWith dims pe = do - name <- newVName $ baseString $ patElemName pe - return pe { patElemName = name - , patElemDec = patElemType pe `arrayOfShape` Shape dims - } - -kernelOrNot :: (MonadFreshNames m, DistLore lore) => - Certificates -> Stm SOACS -> DistAcc lore - -> PostStms lore -> DistAcc lore -> Maybe (Stms lore) - -> DistNestT lore m (DistAcc lore) -kernelOrNot cs bnd acc _ _ Nothing = - addStmToAcc (certify cs bnd) acc -kernelOrNot cs _ _ kernels acc' (Just bnds) = do - addPostStms kernels - postStm $ fmap (certify cs) bnds - return acc' - -distributeMap :: (MonadFreshNames m, LocalScope lore m, DistLore lore) => - MapLoop -> DistAcc lore - -> DistNestT lore m (DistAcc lore) -distributeMap (MapLoop pat aux w lam arrs) acc = - distribute =<< - mapNesting pat aux w lam arrs - (distribute =<< distributeMapBodyStms acc' lam_bnds) - - where acc' = DistAcc { distTargets = pushInnerTarget - (pat, bodyResult $ lambdaBody lam) $ - distTargets acc - , distStms = mempty - } - - lam_bnds = bodyStms $ lambdaBody lam +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + +module Futhark.Pass.ExtractKernels.DistributeNests + ( MapLoop (..), + mapLoopStm, + bodyContainsParallelism, + lambdaContainsParallelism, + determineReduceOp, + histKernel, + DistEnv (..), + DistAcc (..), + runDistNestT, + DistNestT, + liftInner, + distributeMap, + distribute, + distributeSingleStm, + distributeMapBodyStms, + addStmsToAcc, + addStmToAcc, + permutationAndMissing, + addPostStms, + postStm, + inNesting, + ) +where + +import Control.Arrow (first) +import Control.Monad.Identity +import Control.Monad.RWS.Strict +import Control.Monad.Reader +import Control.Monad.Trans.Maybe +import Control.Monad.Writer.Strict +import Data.List (find, partition, tails) +import qualified Data.Map as M +import Data.Maybe +import Futhark.IR +import Futhark.IR.SOACS (SOACS) +import qualified Futhark.IR.SOACS as SOACS +import Futhark.IR.SOACS.SOAC hiding (HistOp, histDest) +import Futhark.IR.SOACS.Simplify (simpleSOACS, simplifyStms) +import Futhark.IR.SegOp +import Futhark.MonadFreshNames +import Futhark.Pass.ExtractKernels.BlockedKernel +import Futhark.Pass.ExtractKernels.Distribution +import Futhark.Pass.ExtractKernels.ISRWIM +import Futhark.Pass.ExtractKernels.Interchange +import Futhark.Tools +import Futhark.Transform.CopyPropagate +import qualified Futhark.Transform.FirstOrderTransform as FOT +import Futhark.Transform.Rename +import Futhark.Util +import Futhark.Util.Log + +scopeForSOACs :: SameScope lore SOACS => Scope lore -> Scope SOACS +scopeForSOACs = castScope + +data MapLoop = MapLoop SOACS.Pattern (StmAux ()) SubExp SOACS.Lambda [VName] + +mapLoopStm :: MapLoop -> Stm SOACS +mapLoopStm (MapLoop pat aux w lam arrs) = + Let pat aux $ Op $ Screma w (mapSOAC lam) arrs + +data DistEnv lore m = DistEnv + { distNest :: Nestings, + distScope :: Scope lore, + distOnTopLevelStms :: Stms SOACS -> DistNestT lore m (Stms lore), + distOnInnerMap :: + MapLoop -> + DistAcc lore -> + DistNestT lore m (DistAcc lore), + distOnSOACSStms :: Stm SOACS -> Binder lore (Stms lore), + distOnSOACSLambda :: Lambda SOACS -> Binder lore (Lambda lore), + distSegLevel :: MkSegLevel lore m + } + +data DistAcc lore = DistAcc + { distTargets :: Targets, + distStms :: Stms lore + } + +data DistRes lore = DistRes + { accPostStms :: PostStms lore, + accLog :: Log + } + +instance Semigroup (DistRes lore) where + DistRes ks1 log1 <> DistRes ks2 log2 = + DistRes (ks1 <> ks2) (log1 <> log2) + +instance Monoid (DistRes lore) where + mempty = DistRes mempty mempty + +newtype PostStms lore = PostStms {unPostStms :: Stms lore} + +instance Semigroup (PostStms lore) where + PostStms xs <> PostStms ys = PostStms $ ys <> xs + +instance Monoid (PostStms lore) where + mempty = PostStms mempty + +typeEnvFromDistAcc :: DistLore lore => DistAcc lore -> Scope lore +typeEnvFromDistAcc = scopeOfPattern . fst . outerTarget . distTargets + +addStmsToAcc :: Stms lore -> DistAcc lore -> DistAcc lore +addStmsToAcc stms acc = + acc {distStms = stms <> distStms acc} + +addStmToAcc :: + (MonadFreshNames m, DistLore lore) => + Stm SOACS -> + DistAcc lore -> + DistNestT lore m (DistAcc lore) +addStmToAcc stm acc = do + onSoacs <- asks distOnSOACSStms + (stm', _) <- runBinder $ onSoacs stm + return acc {distStms = stm' <> distStms acc} + +soacsLambda :: + (MonadFreshNames m, DistLore lore) => + Lambda SOACS -> + DistNestT lore m (Lambda lore) +soacsLambda lam = do + onLambda <- asks distOnSOACSLambda + fst <$> runBinder (onLambda lam) + +newtype DistNestT lore m a + = DistNestT (ReaderT (DistEnv lore m) (WriterT (DistRes lore) m) a) + deriving + ( Functor, + Applicative, + Monad, + MonadReader (DistEnv lore m), + MonadWriter (DistRes lore) + ) + +liftInner :: (LocalScope lore m, DistLore lore) => m a -> DistNestT lore m a +liftInner m = do + outer_scope <- askScope + DistNestT $ + lift $ + lift $ do + inner_scope <- askScope + localScope (outer_scope `M.difference` inner_scope) m + +instance MonadFreshNames m => MonadFreshNames (DistNestT lore m) where + getNameSource = DistNestT $ lift getNameSource + putNameSource = DistNestT . lift . putNameSource + +instance (Monad m, ASTLore lore) => HasScope lore (DistNestT lore m) where + askScope = asks distScope + +instance (Monad m, ASTLore lore) => LocalScope lore (DistNestT lore m) where + localScope types = local $ \env -> + env {distScope = types <> distScope env} + +instance Monad m => MonadLogger (DistNestT lore m) where + addLog msgs = tell mempty {accLog = msgs} + +runDistNestT :: + (MonadLogger m, DistLore lore) => + DistEnv lore m -> + DistNestT lore m (DistAcc lore) -> + m (Stms lore) +runDistNestT env (DistNestT m) = do + (acc, res) <- runWriterT $ runReaderT m env + addLog $ accLog res + -- There may be a few final targets remaining - these correspond to + -- arrays that are identity mapped, and must have statements + -- inserted here. + return $ + unPostStms (accPostStms res) + <> identityStms (outerTarget $ distTargets acc) + where + outermost = nestingLoop $ + case distNest env of + (nest, []) -> nest + (_, nest : _) -> nest + params_to_arrs = + map (first paramName) $ + loopNestingParamsAndArrs outermost + + identityStms (rem_pat, res) = + stmsFromList $ zipWith identityStm (patternValueElements rem_pat) res + identityStm pe (Var v) + | Just arr <- lookup v params_to_arrs = + Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ Copy arr + identityStm pe se = + Let (Pattern [] [pe]) (defAux ()) $ + BasicOp $ + Replicate (Shape [loopNestingWidth outermost]) se + +addPostStms :: Monad m => PostStms lore -> DistNestT lore m () +addPostStms ks = tell $ mempty {accPostStms = ks} + +postStm :: Monad m => Stms lore -> DistNestT lore m () +postStm stms = addPostStms $ PostStms stms + +withStm :: + (Monad m, DistLore lore) => + Stm SOACS -> + DistNestT lore m a -> + DistNestT lore m a +withStm stm = local $ \env -> + env + { distScope = + castScope (scopeOf stm) <> distScope env, + distNest = + letBindInInnerNesting provided $ + distNest env + } + where + provided = namesFromList $ patternNames $ stmPattern stm + +leavingNesting :: + (MonadFreshNames m, DistLore lore) => + DistAcc lore -> + DistNestT lore m (DistAcc lore) +leavingNesting acc = + case popInnerTarget $ distTargets acc of + Nothing -> + error "The kernel targets list is unexpectedly small" + Just ((pat, res), newtargets) + | not $ null $ distStms acc -> do + -- Any statements left over correspond to something that + -- could not be distributed because it would cause irregular + -- arrays. These must be reconstructed into a a Map SOAC + -- that will be sequentialised. XXX: life would be better if + -- we were able to distribute irregular parallelism. + (Nesting _ inner, _) <- asks distNest + let MapNesting _ aux w params_and_arrs = inner + body = Body () (distStms acc) res + used_in_body = freeIn body + (used_params, used_arrs) = + unzip $ + filter ((`nameIn` used_in_body) . paramName . fst) params_and_arrs + lam' = + Lambda + { lambdaParams = used_params, + lambdaBody = body, + lambdaReturnType = map rowType $ patternTypes pat + } + stms <- + runBinder_ $ + auxing aux $ + FOT.transformSOAC pat $ + Screma w (mapSOAC lam') used_arrs + + return $ acc {distTargets = newtargets, distStms = stms} + | otherwise -> do + -- Any results left over correspond to a Replicate or a Copy in + -- the parent nesting, depending on whether the argument is a + -- parameter of the innermost nesting. + (Nesting _ inner_nesting, _) <- asks distNest + let w = loopNestingWidth inner_nesting + aux = loopNestingAux inner_nesting + inps = loopNestingParamsAndArrs inner_nesting + + remnantStm pe (Var v) + | Just (_, arr) <- find ((== v) . paramName . fst) inps = + Let (Pattern [] [pe]) aux $ + BasicOp $ Copy arr + remnantStm pe se = + Let (Pattern [] [pe]) aux $ + BasicOp $ Replicate (Shape [w]) se + + stms = + stmsFromList $ zipWith remnantStm (patternElements pat) res + + return $ acc {distTargets = newtargets, distStms = stms} + +mapNesting :: + (MonadFreshNames m, DistLore lore) => + PatternT Type -> + StmAux () -> + SubExp -> + Lambda SOACS -> + [VName] -> + DistNestT lore m (DistAcc lore) -> + DistNestT lore m (DistAcc lore) +mapNesting pat aux w lam arrs m = + local extend $ leavingNesting =<< m + where + nest = + Nesting mempty $ + MapNesting pat aux w $ + zip (lambdaParams lam) arrs + extend env = + env + { distNest = pushInnerNesting nest $ distNest env, + distScope = castScope (scopeOf lam) <> distScope env + } + +inNesting :: + (Monad m, DistLore lore) => + KernelNest -> + DistNestT lore m a -> + DistNestT lore m a +inNesting (outer, nests) = local $ \env -> + env + { distNest = (inner, nests'), + distScope = foldMap scopeOfLoopNesting (outer : nests) <> distScope env + } + where + (inner, nests') = + case reverse nests of + [] -> (asNesting outer, []) + (inner' : ns) -> (asNesting inner', map asNesting $ outer : reverse ns) + asNesting = Nesting mempty + +bodyContainsParallelism :: Body SOACS -> Bool +bodyContainsParallelism = any isParallelStm . bodyStms + where + isParallelStm stm = + isMap (stmExp stm) + && not ("sequential" `inAttrs` stmAuxAttrs (stmAux stm)) + isMap Op {} = True + isMap _ = False + +lambdaContainsParallelism :: Lambda SOACS -> Bool +lambdaContainsParallelism = bodyContainsParallelism . lambdaBody + +distributeMapBodyStms :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + DistAcc lore -> + Stms SOACS -> + DistNestT lore m (DistAcc lore) +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 + types <- asksScope scopeForSOACs + stream_stms <- + snd <$> runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types + (_, stream_stms') <- + runReaderT (copyPropagateInStms simpleSOACS types stream_stms) types + onStms acc $ stmsToList (fmap (certify cs) stream_stms') ++ stms + onStms acc (stm : stms) = + -- It is important that stm is in scope if 'maybeDistributeStm' + -- wants to distribute, even if this causes the slightly silly + -- situation that stm is in scope of itself. + withStm stm $ maybeDistributeStm stm =<< onStms acc stms + +onInnerMap :: Monad m => MapLoop -> DistAcc lore -> DistNestT lore m (DistAcc lore) +onInnerMap loop acc = do + f <- asks distOnInnerMap + f loop acc + +onTopLevelStms :: Monad m => Stms SOACS -> DistNestT lore m () +onTopLevelStms stms = do + f <- asks distOnTopLevelStms + postStm =<< f stms + +maybeDistributeStm :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + Stm SOACS -> + DistAcc lore -> + DistNestT lore m (DistAcc lore) +maybeDistributeStm stm acc + | "sequential" `inAttrs` stmAuxAttrs (stmAux stm) = + addStmToAcc stm acc +maybeDistributeStm (Let pat aux (Op soac)) acc + | "sequential_outer" `inAttrs` stmAuxAttrs aux = + distributeMapBodyStms acc . fmap (certify (stmAuxCerts aux)) + =<< runBinder_ (FOT.transformSOAC pat soac) +maybeDistributeStm stm@(Let pat _ (Op (Screma w form arrs))) acc + | Just lam <- isMapSOAC form = + -- Only distribute inside the map if we can distribute everything + -- following the map. + distributeIfPossible acc >>= \case + Nothing -> addStmToAcc stm acc + Just acc' -> distribute =<< onInnerMap (MapLoop pat (stmAux stm) w lam arrs) acc' +maybeDistributeStm bnd@(Let pat _ (DoLoop [] val form@ForLoop {} body)) acc + | null (patternContextElements pat), + bodyContainsParallelism body = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | not $ freeIn form `namesIntersect` boundInKernelNest nest, + Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromDistAcc acc') $ do + addPostStms kernels + nest' <- expandKernelNest pat_unused nest + types <- asksScope scopeForSOACs + + -- Simplification is key to hoisting out statements that + -- were variant to the loop, but invariant to the outer maps + -- (which are now innermost). + stms <- + (`runReaderT` types) $ + fmap snd . simplifyStms + =<< interchangeLoops nest' (SeqLoop perm pat val form body) + onTopLevelStms stms + return acc' + _ -> + addStmToAcc bnd acc +maybeDistributeStm stm@(Let pat _ (If cond tbranch fbranch ret)) acc + | null (patternContextElements pat), + bodyContainsParallelism tbranch || bodyContainsParallelism fbranch + || not (all primType (ifReturns ret)) = + distributeSingleStm acc stm >>= \case + Just (kernels, res, nest, acc') + | not $ + (freeIn cond <> freeIn ret) `namesIntersect` boundInKernelNest nest, + Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromDistAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + addPostStms kernels + types <- asksScope scopeForSOACs + let branch = Branch perm pat cond tbranch fbranch ret + stms <- + (`runReaderT` types) $ + fmap snd . simplifyStms + =<< interchangeBranch nest' branch + onTopLevelStms stms + return acc' + _ -> + addStmToAcc stm acc +maybeDistributeStm (Let pat aux (Op (Screma w form arrs))) acc + | Just [Reduce comm lam nes] <- isReduceSOAC form, + Just m <- irwim pat w comm lam $ zip nes arrs = do + types <- asksScope scopeForSOACs + (_, bnds) <- runBinderT (auxing aux m) types + distributeMapBodyStms acc bnds + +-- Parallelise segmented scatters. +maybeDistributeStm bnd@(Let pat (StmAux cs _ _) (Op (Scatter w lam ivs as))) acc = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + localScope (typeEnvFromDistAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + lam' <- soacsLambda lam + addPostStms kernels + postStm =<< segmentedScatterKernel nest' perm pat cs w lam' ivs as + return acc' + _ -> + addStmToAcc bnd acc +-- Parallelise segmented Hist. +maybeDistributeStm bnd@(Let pat (StmAux cs _ _) (Op (Hist w ops lam as))) acc = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + localScope (typeEnvFromDistAcc acc') $ do + lam' <- soacsLambda lam + nest' <- expandKernelNest pat_unused nest + addPostStms kernels + postStm =<< segmentedHistKernel nest' perm cs w ops lam' as + return acc' + _ -> + addStmToAcc bnd acc +-- Parallelise Index slices if the result is going to be returned +-- directly from the kernel. This is because we would otherwise have +-- to sequentialise writing the result, which may be costly. +maybeDistributeStm + stm@( Let + (Pattern [] [pe]) + aux + (BasicOp (Index arr slice)) + ) + acc + | not $ null $ sliceDims slice, + Var (patElemName pe) `elem` snd (innerTarget (distTargets acc)) = + distributeSingleStm acc stm >>= \case + Just (kernels, _res, nest, acc') -> + localScope (typeEnvFromDistAcc acc') $ do + addPostStms kernels + postStm =<< segmentedGatherKernel nest (stmAuxCerts aux) arr slice + return acc' + _ -> + addStmToAcc stm acc +-- If the scan can be distributed by itself, we will turn it into a +-- segmented scan. +-- +-- If the scan cannot be distributed by itself, it will be +-- sequentialised in the default case for this function. +maybeDistributeStm bnd@(Let pat (StmAux cs _ _) (Op (Screma w form arrs))) acc + | Just (scans, map_lam) <- isScanomapSOAC form, + Scan lam nes <- singleScan scans = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromDistAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + map_lam' <- soacsLambda map_lam + lam' <- soacsLambda lam + localScope (typeEnvFromDistAcc acc') $ + segmentedScanomapKernel nest' perm w lam' map_lam' nes arrs + >>= kernelOrNot cs bnd acc kernels acc' + _ -> + addStmToAcc bnd acc +-- if the reduction can be distributed by itself, we will turn it into a +-- segmented reduce. +-- +-- If the reduction cannot be distributed by itself, it will be +-- sequentialised in the default case for this function. +maybeDistributeStm bnd@(Let pat (StmAux cs _ _) (Op (Screma w form arrs))) acc + | Just (reds, map_lam) <- isRedomapSOAC form, + Reduce comm lam nes <- singleReduce reds = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromDistAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + + lam' <- soacsLambda lam + map_lam' <- soacsLambda map_lam + + let comm' + | commutativeLambda lam = Commutative + | otherwise = comm + + regularSegmentedRedomapKernel nest' perm w comm' lam' map_lam' nes arrs + >>= kernelOrNot cs bnd acc kernels acc' + _ -> + addStmToAcc bnd acc +maybeDistributeStm (Let pat (StmAux cs _ _) (Op (Screma w form arrs))) acc = do + -- This Screma is too complicated for us to immediately do + -- anything, so split it up and try again. + scope <- asksScope scopeForSOACs + distributeMapBodyStms acc . fmap (certify cs) . snd + =<< runBinderT (dissectScrema pat w form arrs) scope +maybeDistributeStm (Let pat aux (BasicOp (Replicate (Shape (d : ds)) v))) acc + | [t] <- patternTypes pat = do + tmp <- newVName "tmp" + let rowt = rowType t + newbnd = Let pat aux $ Op $ Screma d (mapSOAC lam) [] + tmpbnd = + Let (Pattern [] [PatElem tmp rowt]) aux $ + BasicOp $ Replicate (Shape ds) v + lam = + Lambda + { lambdaReturnType = [rowt], + lambdaParams = [], + lambdaBody = mkBody (oneStm tmpbnd) [Var tmp] + } + maybeDistributeStm newbnd acc +maybeDistributeStm bnd@(Let _ aux (BasicOp Copy {})) acc = + distributeSingleUnaryStm acc bnd $ \_ outerpat arr -> + return $ oneStm $ Let outerpat aux $ BasicOp $ Copy arr +-- Opaques are applied to the full array, because otherwise they can +-- drastically inhibit parallelisation in some cases. +maybeDistributeStm bnd@(Let (Pattern [] [pe]) aux (BasicOp Opaque {})) acc + | not $ primType $ typeOf pe = + distributeSingleUnaryStm acc bnd $ \_ outerpat arr -> + return $ oneStm $ Let outerpat aux $ BasicOp $ Copy arr +maybeDistributeStm bnd@(Let _ aux (BasicOp (Rearrange perm _))) acc = + distributeSingleUnaryStm acc bnd $ \nest outerpat arr -> do + let r = length (snd nest) + 1 + perm' = [0 .. r -1] ++ map (+ r) perm + -- We need to add a copy, because the original map nest + -- will have produced an array without aliases, and so must we. + arr' <- newVName $ baseString arr + arr_t <- lookupType arr + return $ + stmsFromList + [ Let (Pattern [] [PatElem arr' arr_t]) aux $ BasicOp $ Copy arr, + Let outerpat aux $ BasicOp $ Rearrange perm' arr' + ] +maybeDistributeStm bnd@(Let _ aux (BasicOp (Reshape reshape _))) acc = + distributeSingleUnaryStm acc bnd $ \nest outerpat arr -> do + let reshape' = + map DimNew (kernelNestWidths nest) + ++ map DimNew (newDims reshape) + return $ oneStm $ Let outerpat aux $ BasicOp $ Reshape reshape' arr +maybeDistributeStm stm@(Let _ aux (BasicOp (Rotate rots _))) acc = + distributeSingleUnaryStm acc stm $ \nest outerpat arr -> do + let rots' = map (const $ intConst Int64 0) (kernelNestWidths nest) ++ rots + return $ oneStm $ Let outerpat aux $ BasicOp $ Rotate rots' arr +maybeDistributeStm stm@(Let pat aux (BasicOp (Update arr slice (Var v)))) acc + | not $ null $ sliceDims slice = + distributeSingleStm acc stm >>= \case + Just (kernels, res, nest, acc') + | res == map Var (patternNames $ stmPattern stm), + Just (perm, pat_unused) <- permutationAndMissing pat res -> do + addPostStms kernels + localScope (typeEnvFromDistAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + postStm + =<< segmentedUpdateKernel nest' perm (stmAuxCerts aux) arr slice v + return acc' + _ -> addStmToAcc stm acc +-- XXX? This rule is present to avoid the case where an in-place +-- update is distributed as its own kernel, as this would mean thread +-- then writes the entire array that it updated. This is problematic +-- because the in-place updates is O(1), but writing the array is +-- O(n). It is OK if the in-place update is preceded, followed, or +-- nested inside a sequential loop or similar, because that will +-- probably be O(n) by itself. As a hack, we only distribute if there +-- does not appear to be a loop following. The better solution is to +-- depend on memory block merging for this optimisation, but it is not +-- ready yet. +maybeDistributeStm (Let pat aux (BasicOp (Update arr [DimFix i] v))) acc + | [t] <- patternTypes pat, + arrayRank t == 1, + not $ any (amortises . stmExp) $ distStms acc = do + let w = arraySize 0 t + et = stripArray 1 t + lam = + Lambda + { lambdaParams = [], + lambdaReturnType = [Prim int64, et], + lambdaBody = mkBody mempty [i, v] + } + maybeDistributeStm (Let pat aux $ Op $ Scatter (intConst Int64 1) lam [] [(w, 1, arr)]) acc + where + amortises DoLoop {} = True + amortises Op {} = True + amortises _ = False +maybeDistributeStm stm@(Let _ aux (BasicOp (Concat d x xs w))) acc = + distributeSingleStm acc stm >>= \case + Just (kernels, _, nest, acc') -> + localScope (typeEnvFromDistAcc acc') $ + segmentedConcat nest + >>= kernelOrNot (stmAuxCerts aux) stm acc kernels acc' + _ -> + addStmToAcc stm acc + where + segmentedConcat nest = + isSegmentedOp nest [0] mempty mempty [] (x : xs) $ + \pat _ _ _ (x' : xs') -> + let d' = d + length (snd nest) + 1 + in addStm $ Let pat aux $ BasicOp $ Concat d' x' xs' w +maybeDistributeStm bnd acc = + addStmToAcc bnd acc + +distributeSingleUnaryStm :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + DistAcc lore -> + Stm SOACS -> + (KernelNest -> PatternT Type -> VName -> DistNestT lore m (Stms lore)) -> + DistNestT lore m (DistAcc lore) +distributeSingleUnaryStm acc bnd f = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | res == map Var (patternNames $ stmPattern bnd), + (outer, _) <- nest, + [(arr_p, arr)] <- loopNestingParamsAndArrs outer, + boundInKernelNest nest `namesIntersection` freeIn bnd + == oneName (paramName arr_p), + perfectlyMapped arr nest -> do + addPostStms kernels + let outerpat = loopNestingPattern $ fst nest + localScope (typeEnvFromDistAcc acc') $ do + postStm =<< f nest outerpat arr + return acc' + _ -> addStmToAcc bnd acc + where + perfectlyMapped arr (outer, nest) + | [(p, arr')] <- loopNestingParamsAndArrs outer, + arr == arr' = + case nest of + [] -> True + x : xs -> perfectlyMapped (paramName p) (x, xs) + | otherwise = + False + +distribute :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + DistAcc lore -> + DistNestT lore m (DistAcc lore) +distribute acc = + fromMaybe acc <$> distributeIfPossible acc + +mkSegLevel :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + DistNestT lore m (MkSegLevel lore (DistNestT lore m)) +mkSegLevel = do + mk_lvl <- asks distSegLevel + return $ \w desc r -> do + (lvl, stms) <- lift $ liftInner $ runBinderT' $ mk_lvl w desc r + addStms stms + return lvl + +distributeIfPossible :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + DistAcc lore -> + DistNestT lore m (Maybe (DistAcc lore)) +distributeIfPossible acc = do + nest <- asks distNest + mk_lvl <- mkSegLevel + tryDistribute mk_lvl nest (distTargets acc) (distStms acc) >>= \case + Nothing -> return Nothing + Just (targets, kernel) -> do + postStm kernel + return $ + Just + DistAcc + { distTargets = targets, + distStms = mempty + } + +distributeSingleStm :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + DistAcc lore -> + Stm SOACS -> + DistNestT + lore + m + ( Maybe + ( PostStms lore, + Result, + KernelNest, + DistAcc lore + ) + ) +distributeSingleStm acc bnd = do + nest <- asks distNest + mk_lvl <- mkSegLevel + tryDistribute mk_lvl nest (distTargets acc) (distStms acc) >>= \case + Nothing -> return Nothing + Just (targets, distributed_bnds) -> + tryDistributeStm nest targets bnd >>= \case + Nothing -> return Nothing + Just (res, targets', new_kernel_nest) -> + return $ + Just + ( PostStms distributed_bnds, + res, + new_kernel_nest, + DistAcc + { distTargets = targets', + distStms = mempty + } + ) + +segmentedScatterKernel :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + KernelNest -> + [Int] -> + PatternT Type -> + Certificates -> + SubExp -> + Lambda lore -> + [VName] -> + [(SubExp, 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 + -- things are different because a scatter is not a reduction or + -- scan. + -- + -- First, pretend that the scatter is also part of the nesting. The + -- KernelNest we produce here is technically not sensible, but it's + -- good enough for flatKernel to work. + let nesting = + MapNesting scatter_pat (StmAux cs mempty ()) scatter_w $ zip (lambdaParams lam) ivs + nest' = + pushInnerKernelNesting (scatter_pat, bodyResult $ lambdaBody lam) nesting nest + (ispace, kernel_inps) <- flatKernel nest' + + let (as_ws, as_ns, as) = unzip3 dests + + -- The input/output arrays ('as') _must_ correspond to some kernel + -- input, or else the original nested scatter would have been + -- ill-typed. Find them. + as_inps <- mapM (findInput kernel_inps) as + + mk_lvl <- mkSegLevel + + let rts = + concatMap (take 1) $ + chunks as_ns $ + drop (sum as_ns) $ lambdaReturnType lam + (is, vs) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam + + -- Maybe add certificates to the indices. + (is', k_body_stms) <- runBinder $ do + addStms $ bodyStms $ lambdaBody lam + forM is $ \i -> + if cs == mempty + then return i + 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 + + (k, k_bnds) <- mapKernel mk_lvl ispace kernel_inps rts k_body + + traverse renameStm <=< runBinder_ $ do + addStms k_bnds + + let pat = + Pattern [] $ + rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + + letBind pat $ Op $ segOp k + where + findInput kernel_inps a = + maybe bad return $ find ((== a) . kernelInputName) kernel_inps + bad = error "Ill-typed nested scatter encountered." + + inPlaceReturn ispace (aw, inp, is_vs) = + WriteReturns + (init ws ++ [aw]) + (kernelInputArray inp) + [(map DimFix $ map Var (init gtids) ++ [i], v) | (i, v) <- is_vs] + where + (gtids, ws) = unzip ispace + +segmentedUpdateKernel :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + KernelNest -> + [Int] -> + Certificates -> + VName -> + Slice SubExp -> + VName -> + DistNestT lore m (Stms lore) +segmentedUpdateKernel nest perm cs arr slice v = do + (base_ispace, kernel_inps) <- flatKernel nest + let slice_dims = sliceDims slice + slice_gtids <- replicateM (length slice_dims) (newVName "gtid_slice") + + let ispace = base_ispace ++ zip slice_gtids slice_dims + + ((res_t, res), kstms) <- runBinder $ do + -- Compute indexes into full array. + v' <- + certifying cs $ + letSubExp "v" $ BasicOp $ Index v $ map (DimFix . Var) slice_gtids + slice_is <- + traverse (toSubExp "index") $ + fixSlice (map (fmap pe64) slice) $ map (pe64 . Var) slice_gtids + + let write_is = map (Var . fst) base_ispace ++ slice_is + arr' = + maybe (error "incorrectly typed Update") kernelInputArray $ + find ((== arr) . kernelInputName) kernel_inps + arr_t <- lookupType arr' + v_t <- subExpType v' + return + ( v_t, + WriteReturns (arrayDims arr_t) arr' [(map DimFix write_is, v')] + ) + + mk_lvl <- mkSegLevel + (k, prestms) <- + mapKernel mk_lvl ispace kernel_inps [res_t] $ + KernelBody () kstms [res] + + traverse renameStm <=< runBinder_ $ do + addStms prestms + + let pat = + Pattern [] $ + rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + + letBind pat $ Op $ segOp k + +segmentedGatherKernel :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + KernelNest -> + Certificates -> + VName -> + Slice SubExp -> + DistNestT lore m (Stms lore) +segmentedGatherKernel nest cs arr slice = do + let slice_dims = sliceDims slice + slice_gtids <- replicateM (length slice_dims) (newVName "gtid_slice") + + (base_ispace, kernel_inps) <- flatKernel nest + let ispace = base_ispace ++ zip slice_gtids slice_dims + + ((res_t, res), kstms) <- runBinder $ do + -- Compute indexes into full array. + slice'' <- + subExpSlice $ + sliceSlice (primExpSlice slice) $ + primExpSlice $ map (DimFix . Var) slice_gtids + v' <- certifying cs $ letSubExp "v" $ BasicOp $ Index arr slice'' + v_t <- subExpType v' + return (v_t, Returns ResultMaySimplify v') + + mk_lvl <- mkSegLevel + (k, prestms) <- + mapKernel mk_lvl ispace kernel_inps [res_t] $ + KernelBody () kstms [res] + + traverse renameStm <=< runBinder_ $ do + addStms prestms + + let pat = Pattern [] $ patternValueElements $ loopNestingPattern $ fst nest + + letBind pat $ Op $ segOp k + +segmentedHistKernel :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + KernelNest -> + [Int] -> + Certificates -> + SubExp -> + [SOACS.HistOp SOACS] -> + Lambda lore -> + [VName] -> + DistNestT lore m (Stms lore) +segmentedHistKernel nest perm cs hist_w ops lam arrs = do + -- We replicate some of the checking done by 'isSegmentedOp', but + -- things are different because a Hist is not a reduction or + -- scan. + (ispace, inputs) <- flatKernel nest + let orig_pat = + Pattern [] $ + rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + + -- The input/output arrays _must_ correspond to some kernel input, + -- or else the original nested Hist would have been ill-typed. + -- Find them. + ops' <- forM ops $ \(SOACS.HistOp num_bins rf dests nes op) -> + SOACS.HistOp num_bins rf + <$> mapM (fmap kernelInputArray . findInput inputs) dests + <*> pure nes + <*> pure op + + mk_lvl <- asks distSegLevel + onLambda <- asks distOnSOACSLambda + let onLambda' = fmap fst . runBinder . onLambda + liftInner $ + runBinderT'_ $ do + -- It is important not to launch unnecessarily many threads for + -- histograms, because it may mean we unnecessarily need to reduce + -- subhistograms as well. + lvl <- mk_lvl (hist_w : map snd ispace) "seghist" $ NoRecommendation SegNoVirt + addStms + =<< histKernel onLambda' lvl orig_pat ispace inputs cs hist_w ops' lam arrs + where + findInput kernel_inps a = + maybe bad return $ find ((== a) . kernelInputName) kernel_inps + bad = error "Ill-typed nested Hist encountered." + +histKernel :: + (MonadBinder m, DistLore (Lore m)) => + (Lambda SOACS -> m (Lambda (Lore m))) -> + SegOpLevel (Lore m) -> + PatternT Type -> + [(VName, SubExp)] -> + [KernelInput] -> + Certificates -> + SubExp -> + [SOACS.HistOp SOACS] -> + Lambda (Lore m) -> + [VName] -> + m (Stms (Lore m)) +histKernel onLambda lvl orig_pat ispace inputs cs hist_w ops lam arrs = runBinderT'_ $ do + ops' <- forM ops $ \(SOACS.HistOp num_bins rf dests nes op) -> do + (op', nes', shape) <- determineReduceOp op nes + op'' <- lift $ onLambda op' + return $ HistOp num_bins rf dests nes' shape op'' + + let isDest = flip elem $ concatMap histDest ops' + inputs' = filter (not . isDest . kernelInputArray) inputs + + certifying cs $ + addStms =<< traverse renameStm + =<< segHist lvl orig_pat hist_w ispace inputs' ops' lam arrs + +determineReduceOp :: + (MonadBinder m, Lore m ~ lore) => + Lambda SOACS -> + [SubExp] -> + m (Lambda SOACS, [SubExp], Shape) +determineReduceOp lam nes = + -- FIXME? We are assuming that the accumulator is a replicate, and + -- we fish out its value in a gross way. + case mapM subExpVar nes of + Just ne_vs' -> do + let (shape, lam') = isVectorMap lam + nes' <- forM ne_vs' $ \ne_v -> do + ne_v_t <- lookupType ne_v + letSubExp "hist_ne" $ + BasicOp $ + Index ne_v $ + fullSlice ne_v_t $ + replicate (shapeRank shape) $ DimFix $ intConst Int64 0 + return (lam', nes', shape) + Nothing -> + return (lam, nes, mempty) + +isVectorMap :: Lambda SOACS -> (Shape, Lambda SOACS) +isVectorMap lam + | [Let (Pattern [] pes) _ (Op (Screma w form arrs))] <- + stmsToList $ bodyStms $ lambdaBody lam, + bodyResult (lambdaBody lam) == map (Var . patElemName) pes, + Just map_lam <- isMapSOAC form, + arrs == map paramName (lambdaParams lam) = + let (shape, lam') = isVectorMap map_lam + in (Shape [w] <> shape, lam') + | otherwise = (mempty, lam) + +segmentedScanomapKernel :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + KernelNest -> + [Int] -> + SubExp -> + Lambda lore -> + Lambda lore -> + [SubExp] -> + [VName] -> + DistNestT lore m (Maybe (Stms lore)) +segmentedScanomapKernel nest perm segment_size lam map_lam nes arrs = do + mk_lvl <- asks distSegLevel + isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $ + \pat ispace inps nes' _ -> do + let scan_op = SegBinOp Noncommutative lam nes' mempty + lvl <- mk_lvl (segment_size : map snd ispace) "segscan" $ NoRecommendation SegNoVirt + addStms =<< traverse renameStm + =<< segScan lvl pat segment_size [scan_op] map_lam arrs ispace inps + +regularSegmentedRedomapKernel :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + KernelNest -> + [Int] -> + SubExp -> + Commutativity -> + Lambda lore -> + Lambda lore -> + [SubExp] -> + [VName] -> + DistNestT lore m (Maybe (Stms lore)) +regularSegmentedRedomapKernel nest perm segment_size comm lam map_lam nes arrs = do + mk_lvl <- asks distSegLevel + isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $ + \pat ispace inps nes' _ -> do + let red_op = SegBinOp comm lam nes' mempty + lvl <- mk_lvl (segment_size : map snd ispace) "segred" $ NoRecommendation SegNoVirt + addStms =<< traverse renameStm + =<< segRed lvl pat segment_size [red_op] map_lam arrs ispace inps + +isSegmentedOp :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + KernelNest -> + [Int] -> + Names -> + Names -> + [SubExp] -> + [VName] -> + ( PatternT Type -> + [(VName, SubExp)] -> + [KernelInput] -> + [SubExp] -> + [VName] -> + BinderT lore m () + ) -> + DistNestT lore m (Maybe (Stms lore)) +isSegmentedOp nest perm free_in_op _free_in_fold_op nes arrs m = runMaybeT $ do + -- We must verify that array inputs to the operation are inputs to + -- the outermost loop nesting or free in the loop nest. Nothing + -- free in the op may be bound by the nest. Furthermore, the + -- neutral elements must be free in the loop nest. + -- + -- We must summarise any names from free_in_op that are bound in the + -- nest, and describe how to obtain them given segment indices. + + let bound_by_nest = boundInKernelNest nest + + (ispace, kernel_inps) <- flatKernel nest + + when (free_in_op `namesIntersect` bound_by_nest) $ + fail "Non-fold lambda uses nest-bound parameters." + + let indices = map fst ispace + + prepareNe (Var v) + | v `nameIn` bound_by_nest = + fail "Neutral element bound in nest" + prepareNe ne = return ne + + prepareArr arr = + case find ((== arr) . kernelInputName) kernel_inps of + Just inp + | kernelInputIndices inp == map Var indices -> + return $ return $ kernelInputArray inp + Nothing + | not (arr `nameIn` bound_by_nest) -> + -- This input is something that is free inside + -- the loop nesting. We will have to replicate + -- it. + return $ + letExp + (baseString arr ++ "_repd") + (BasicOp $ Replicate (Shape $ map snd ispace) $ Var arr) + _ -> + fail "Input not free, perfectly mapped, or outermost." + + nes' <- mapM prepareNe nes + + mk_arrs <- mapM prepareArr arrs + + lift $ + liftInner $ + runBinderT'_ $ do + nested_arrs <- sequence mk_arrs + + let pat = + Pattern [] $ + rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + + m pat ispace kernel_inps nes' nested_arrs + +permutationAndMissing :: PatternT Type -> [SubExp] -> Maybe ([Int], [PatElemT Type]) +permutationAndMissing pat res = do + let pes = patternValueElements pat + (_used, unused) = + partition ((`nameIn` freeIn res) . patElemName) pes + res_expanded = res ++ map (Var . patElemName) unused + perm <- map (Var . patElemName) pes `isPermutationOf` res_expanded + return (perm, unused) + +-- Add extra pattern elements to every kernel nesting level. +expandKernelNest :: + MonadFreshNames m => + [PatElemT Type] -> + KernelNest -> + m KernelNest +expandKernelNest pes (outer_nest, inner_nests) = do + let outer_size = + loopNestingWidth outer_nest : + map loopNestingWidth inner_nests + inner_sizes = tails $ map loopNestingWidth inner_nests + outer_nest' <- expandWith outer_nest outer_size + inner_nests' <- zipWithM expandWith inner_nests inner_sizes + return (outer_nest', inner_nests') + where + expandWith nest dims = do + pes' <- mapM (expandPatElemWith dims) pes + return + nest + { loopNestingPattern = + Pattern [] $ + patternElements (loopNestingPattern nest) <> pes' + } + + expandPatElemWith dims pe = do + name <- newVName $ baseString $ patElemName pe + return + pe + { patElemName = name, + patElemDec = patElemType pe `arrayOfShape` Shape dims + } + +kernelOrNot :: + (MonadFreshNames m, DistLore lore) => + Certificates -> + Stm SOACS -> + DistAcc lore -> + PostStms lore -> + DistAcc lore -> + Maybe (Stms lore) -> + DistNestT lore m (DistAcc lore) +kernelOrNot cs bnd acc _ _ Nothing = + addStmToAcc (certify cs bnd) acc +kernelOrNot cs _ _ kernels acc' (Just bnds) = do + addPostStms kernels + postStm $ fmap (certify cs) bnds + return acc' + +distributeMap :: + (MonadFreshNames m, LocalScope lore m, DistLore lore) => + MapLoop -> + DistAcc lore -> + DistNestT lore m (DistAcc lore) +distributeMap (MapLoop pat aux w lam arrs) acc = + distribute + =<< mapNesting + pat + aux + w + lam + arrs + (distribute =<< distributeMapBodyStms acc' lam_bnds) + where + acc' = + DistAcc + { distTargets = + pushInnerTarget + (pat, bodyResult $ lambdaBody lam) + $ distTargets acc, + distStms = mempty + } + + lam_bnds = bodyStms $ lambdaBody lam 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 @@ -1,67 +1,67 @@ +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TupleSections #-} {-# LANGUAGE LambdaCase #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ConstraintKinds #-} -module Futhark.Pass.ExtractKernels.Distribution - ( - Target - , Targets - , ppTargets - , singleTarget - , outerTarget - , innerTarget - , pushInnerTarget - , popInnerTarget - , targetsScope - - , LoopNesting (..) - , ppLoopNesting - , scopeOfLoopNesting - - , Nesting (..) - , Nestings - , ppNestings - , letBindInInnerNesting - , singleNesting - , pushInnerNesting - - , KernelNest - , ppKernelNest - , newKernel - , innermostKernelNesting - , pushKernelNesting - , pushInnerKernelNesting - , kernelNestLoops - , kernelNestWidths - , boundInKernelNest - , boundInKernelNests - , flatKernel - , constructKernel +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} - , tryDistribute - , tryDistributeStm - ) - where +module Futhark.Pass.ExtractKernels.Distribution + ( Target, + Targets, + ppTargets, + singleTarget, + outerTarget, + innerTarget, + pushInnerTarget, + popInnerTarget, + targetsScope, + LoopNesting (..), + ppLoopNesting, + scopeOfLoopNesting, + Nesting (..), + Nestings, + ppNestings, + letBindInInnerNesting, + singleNesting, + pushInnerNesting, + KernelNest, + ppKernelNest, + newKernel, + innermostKernelNesting, + pushKernelNesting, + pushInnerKernelNesting, + kernelNestLoops, + kernelNestWidths, + boundInKernelNest, + boundInKernelNests, + flatKernel, + constructKernel, + tryDistribute, + tryDistributeStm, + ) +where import Control.Monad.RWS.Strict import Control.Monad.Trans.Maybe -import qualified Data.Map.Strict as M import Data.Foldable -import Data.Maybe import Data.List (elemIndex, sortOn) - +import qualified Data.Map.Strict as M +import Data.Maybe import Futhark.IR import Futhark.IR.SegOp import Futhark.MonadFreshNames +import Futhark.Pass.ExtractKernels.BlockedKernel + ( DistLore, + KernelInput (..), + MkSegLevel, + mapKernel, + readKernelInput, + ) import Futhark.Tools -import Futhark.Util import Futhark.Transform.Rename +import Futhark.Util import Futhark.Util.Log -import Futhark.Pass.ExtractKernels.BlockedKernel - (DistLore, mapKernel, KernelInput(..), readKernelInput, MkSegLevel) type Target = (PatternT Type, Result) @@ -70,15 +70,17 @@ -- element of a pattern must be present as the result of the -- immediately enclosing target. This is ensured by 'pushInnerTarget' -- by removing unused pattern elements. -data Targets = Targets { _innerTarget :: Target - , _outerTargets :: [Target] - } +data Targets = Targets + { _innerTarget :: Target, + _outerTargets :: [Target] + } ppTargets :: Targets -> String ppTargets (Targets target targets) = unlines $ map ppTarget $ targets ++ [target] - where ppTarget (pat, res) = - pretty pat ++ " <- " ++ pretty res + where + ppTarget (pat, res) = + pretty pat ++ " <- " ++ pretty res singleTarget :: Target -> Targets singleTarget = flip Targets [] @@ -97,16 +99,17 @@ pushInnerTarget :: Target -> Targets -> Targets pushInnerTarget (pat, res) (Targets inner_target targets) = Targets (pat', res') (targets ++ [inner_target]) - where (pes', res') = unzip $ filter (used . fst) $ zip (patternElements pat) res - pat' = Pattern [] pes' - inner_used = freeIn $ snd inner_target - used pe = patElemName pe `nameIn` inner_used + where + (pes', res') = unzip $ filter (used . fst) $ zip (patternElements pat) res + pat' = Pattern [] pes' + inner_used = freeIn $ snd inner_target + used pe = patElemName pe `nameIn` inner_used popInnerTarget :: Targets -> Maybe (Target, Targets) popInnerTarget (Targets t ts) = case reverse ts of - x:xs -> Just (t, Targets x $ reverse xs) - [] -> Nothing + x : xs -> Just (t, Targets x $ reverse xs) + [] -> Nothing targetScope :: DistLore lore => Target -> Scope lore targetScope = scopeOfPattern . fst @@ -114,50 +117,53 @@ targetsScope :: DistLore lore => Targets -> Scope lore targetsScope (Targets t ts) = mconcat $ map targetScope $ t : ts -data LoopNesting = MapNesting { loopNestingPattern :: PatternT Type - , loopNestingAux :: StmAux () - , loopNestingWidth :: SubExp - , loopNestingParamsAndArrs :: [(Param Type, VName)] - } - deriving (Show) +data LoopNesting = MapNesting + { loopNestingPattern :: PatternT Type, + loopNestingAux :: StmAux (), + loopNestingWidth :: SubExp, + loopNestingParamsAndArrs :: [(Param Type, VName)] + } + deriving (Show) scopeOfLoopNesting :: DistLore lore => LoopNesting -> Scope lore scopeOfLoopNesting = scopeOfLParams . map fst . loopNestingParamsAndArrs ppLoopNesting :: LoopNesting -> String ppLoopNesting (MapNesting _ _ _ params_and_arrs) = - pretty (map fst params_and_arrs) ++ - " <- " ++ - pretty (map snd params_and_arrs) + pretty (map fst params_and_arrs) + ++ " <- " + ++ pretty (map snd params_and_arrs) loopNestingParams :: LoopNesting -> [Param Type] -loopNestingParams = map fst . loopNestingParamsAndArrs +loopNestingParams = map fst . loopNestingParamsAndArrs instance FreeIn LoopNesting where freeIn' (MapNesting pat aux w params_and_arrs) = - freeIn' pat <> - freeIn' aux <> - freeIn' w <> - freeIn' params_and_arrs + freeIn' pat + <> freeIn' aux + <> freeIn' w + <> freeIn' params_and_arrs -data Nesting = Nesting { nestingLetBound :: Names - , nestingLoop :: LoopNesting - } - deriving (Show) +data Nesting = Nesting + { nestingLetBound :: Names, + nestingLoop :: LoopNesting + } + deriving (Show) letBindInNesting :: Names -> Nesting -> Nesting letBindInNesting newnames (Nesting oldnames loop) = Nesting (oldnames <> newnames) loop - -- ^ First pair element is the very innermost ("current") nest. In -- the list, the outermost nest comes first. + type Nestings = (Nesting, [Nesting]) ppNestings :: Nestings -> String ppNestings (nesting, nestings) = unlines $ map ppNesting $ nestings ++ [nesting] - where ppNesting (Nesting _ loop) = - ppLoopNesting loop + where + ppNesting (Nesting _ loop) = + ppLoopNesting loop singleNesting :: Nesting -> Nestings singleNesting = (,[]) @@ -169,15 +175,15 @@ -- | Both parameters and let-bound. boundInNesting :: Nesting -> Names boundInNesting nesting = - namesFromList (map paramName (loopNestingParams loop)) <> - nestingLetBound nesting - where loop = nestingLoop nesting + namesFromList (map paramName (loopNestingParams loop)) + <> nestingLetBound nesting + where + loop = nestingLoop nesting letBindInInnerNesting :: Names -> Nestings -> Nestings letBindInInnerNesting names (nest, nestings) = (letBindInNesting names nest, nestings) - -- | Note: first element is *outermost* nesting. This is different -- from the similar types elsewhere! type KernelNest = (LoopNesting, [LoopNesting]) @@ -195,8 +201,9 @@ -- list, also taking care to swap patterns if necessary. pushKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest pushKernelNesting target newnest (nest, nests) = - (fixNestingPatternOrder newnest target (loopNestingPattern nest), - nest : nests) + ( fixNestingPatternOrder newnest target (loopNestingPattern nest), + nest : nests + ) -- | Add new innermost nesting, pushing the current outermost to the -- list. It is important that the 'Target' has the right order @@ -204,18 +211,20 @@ pushInnerKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest pushInnerKernelNesting target newnest (nest, nests) = (nest, nests ++ [fixNestingPatternOrder newnest target (loopNestingPattern innermost)]) - where innermost = case reverse nests of - [] -> nest - n:_ -> n + where + innermost = case reverse nests of + [] -> nest + n : _ -> n fixNestingPatternOrder :: LoopNesting -> Target -> PatternT Type -> LoopNesting -fixNestingPatternOrder nest (_,res) inner_pat = - nest { loopNestingPattern = basicPattern [] pat' } - where pat = loopNestingPattern nest - pat' = map fst fixed_target - fixed_target = sortOn posInInnerPat $ zip (patternValueIdents pat) res - posInInnerPat (_, Var v) = fromMaybe 0 $ elemIndex v $ patternNames inner_pat - posInInnerPat _ = 0 +fixNestingPatternOrder nest (_, res) inner_pat = + nest {loopNestingPattern = basicPattern [] pat'} + where + pat = loopNestingPattern nest + pat' = map fst fixed_target + fixed_target = sortOn posInInnerPat $ zip (patternValueIdents pat) res + posInInnerPat (_, Var v) = fromMaybe 0 $ elemIndex v $ patternNames inner_pat + posInInnerPat _ = 0 newKernel :: LoopNesting -> KernelNest newKernel nest = (nest, []) @@ -227,17 +236,23 @@ boundInKernelNest = mconcat . boundInKernelNests boundInKernelNests :: KernelNest -> [Names] -boundInKernelNests = map (namesFromList . - map (paramName . fst) . - loopNestingParamsAndArrs) . - kernelNestLoops +boundInKernelNests = + map + ( namesFromList + . map (paramName . fst) + . loopNestingParamsAndArrs + ) + . kernelNestLoops kernelNestWidths :: KernelNest -> [SubExp] kernelNestWidths = map loopNestingWidth . kernelNestLoops -constructKernel :: (DistLore lore, MonadFreshNames m, LocalScope lore m) => - MkSegLevel lore m -> KernelNest -> Body lore - -> m (Stm lore, Stms lore) +constructKernel :: + (DistLore lore, MonadFreshNames m, LocalScope lore m) => + MkSegLevel lore m -> + KernelNest -> + Body lore -> + m (Stm lore, Stms lore) constructKernel mk_lvl kernel_nest inner_body = runBinderT' $ do (ispace, inps) <- flatKernel kernel_nest let aux = loopNestingAux first_nest @@ -245,10 +260,11 @@ pat = loopNestingPattern first_nest rts = map (stripArray (length ispace)) $ patternTypes pat - inner_body' <- fmap (uncurry (flip (KernelBody ()))) $ runBinder $ - localScope ispace_scope $ do - mapM_ readKernelInput $ filter inputIsUsed inps - map (Returns ResultMaySimplify) <$> bodyBind inner_body + inner_body' <- fmap (uncurry (flip (KernelBody ()))) $ + runBinder $ + localScope ispace_scope $ do + mapM_ readKernelInput $ filter inputIsUsed inps + map (Returns ResultMaySimplify) <$> bodyBind inner_body (segop, aux_stms) <- lift $ mapKernel mk_lvl ispace [] rts inner_body' @@ -264,260 +280,323 @@ -- (1) The index space. -- -- (2) The kernel inputs - note that some of these may be unused. -flatKernel :: MonadFreshNames m => - KernelNest - -> m ([(VName, SubExp)], - [KernelInput]) +flatKernel :: + MonadFreshNames m => + KernelNest -> + m + ( [(VName, SubExp)], + [KernelInput] + ) flatKernel (MapNesting _ _ nesting_w params_and_arrs, []) = do i <- newVName "gtid" - let inps = [ KernelInput pname ptype arr [Var i] | - (Param pname ptype, arr) <- params_and_arrs ] - return ([(i,nesting_w)], inps) - + let inps = + [ KernelInput pname ptype arr [Var i] + | (Param pname ptype, arr) <- params_and_arrs + ] + return ([(i, nesting_w)], inps) flatKernel (MapNesting _ _ nesting_w params_and_arrs, nest : nests) = do i <- newVName "gtid" (ispace, inps) <- flatKernel (nest, nests) let inps' = map fixupInput inps isParam inp = - snd <$> find ((==kernelInputArray inp) . paramName . fst) params_and_arrs + snd <$> find ((== kernelInputArray inp) . paramName . fst) params_and_arrs fixupInput inp | Just arr <- isParam inp = - inp { kernelInputArray = arr - , kernelInputIndices = Var i : kernelInputIndices inp } + inp + { kernelInputArray = arr, + kernelInputIndices = Var i : kernelInputIndices inp + } | otherwise = - inp + inp return ((i, nesting_w) : ispace, extra_inps i <> inps') - where extra_inps i = - [ KernelInput pname ptype arr [Var i] | - (Param pname ptype, arr) <- params_and_arrs ] + where + extra_inps i = + [ KernelInput pname ptype arr [Var i] + | (Param pname ptype, arr) <- params_and_arrs + ] -- | Description of distribution to do. -data DistributionBody = DistributionBody { - distributionTarget :: Targets - , distributionFreeInBody :: Names - , distributionIdentityMap :: M.Map VName Ident - , distributionExpandTarget :: Target -> Target - -- ^ Also related to avoiding identity mapping. +data DistributionBody = DistributionBody + { distributionTarget :: Targets, + distributionFreeInBody :: Names, + distributionIdentityMap :: M.Map VName Ident, + -- | Also related to avoiding identity mapping. + distributionExpandTarget :: Target -> Target } distributionInnerPattern :: DistributionBody -> PatternT Type distributionInnerPattern = fst . innerTarget . distributionTarget -distributionBodyFromStms :: ASTLore lore => - Targets -> Stms lore -> (DistributionBody, Result) +distributionBodyFromStms :: + ASTLore lore => + Targets -> + Stms lore -> + (DistributionBody, Result) distributionBodyFromStms (Targets (inner_pat, inner_res) targets) stms = let bound_by_stms = namesFromList $ M.keys $ scopeOf stms (inner_pat', inner_res', inner_identity_map, inner_expand_target) = removeIdentityMappingGeneral bound_by_stms inner_pat inner_res - in (DistributionBody - { distributionTarget = Targets (inner_pat', inner_res') targets - , distributionFreeInBody = foldMap freeIn stms `namesSubtract` bound_by_stms - , distributionIdentityMap = inner_identity_map - , distributionExpandTarget = inner_expand_target - }, - inner_res') + in ( DistributionBody + { distributionTarget = Targets (inner_pat', inner_res') targets, + distributionFreeInBody = foldMap freeIn stms `namesSubtract` bound_by_stms, + distributionIdentityMap = inner_identity_map, + distributionExpandTarget = inner_expand_target + }, + inner_res' + ) -distributionBodyFromStm :: ASTLore lore => - Targets -> Stm lore -> (DistributionBody, Result) +distributionBodyFromStm :: + ASTLore lore => + Targets -> + Stm lore -> + (DistributionBody, Result) distributionBodyFromStm targets bnd = distributionBodyFromStms targets $ oneStm bnd -createKernelNest :: (MonadFreshNames m, HasScope t m) => - Nestings - -> DistributionBody - -> m (Maybe (Targets, KernelNest)) +createKernelNest :: + (MonadFreshNames m, HasScope t m) => + Nestings -> + DistributionBody -> + m (Maybe (Targets, KernelNest)) createKernelNest (inner_nest, nests) distrib_body = do let Targets target targets = distributionTarget distrib_body unless (length nests == length targets) $ - error $ "Nests and targets do not match!\n" ++ - "nests: " ++ ppNestings (inner_nest, nests) ++ - "\ntargets:" ++ ppTargets (Targets target targets) + error $ + "Nests and targets do not match!\n" + ++ "nests: " + ++ ppNestings (inner_nest, nests) + ++ "\ntargets:" + ++ ppTargets (Targets target targets) runMaybeT $ fmap prepare $ recurse $ zip nests targets + where + prepare (x, _, z) = (z, x) + bound_in_nest = + mconcat $ map boundInNesting $ inner_nest : nests - where prepare (x, _, z) = (z, x) - bound_in_nest = - mconcat $ map boundInNesting $ inner_nest : nests - -- | Can something of this type be taken outside the nest? - -- I.e. are none of its dimensions bound inside the nest. - distributableType = - (==mempty) . namesIntersection bound_in_nest . freeIn . arrayDims + distributableType = + (== mempty) . namesIntersection bound_in_nest . freeIn . arrayDims - distributeAtNesting :: (HasScope t m, MonadFreshNames m) => - Nesting - -> PatternT Type - -> (LoopNesting -> KernelNest, Names) - -> M.Map VName Ident - -> [Ident] - -> (Target -> Targets) - -> MaybeT m (KernelNest, Names, Targets) - distributeAtNesting - (Nesting nest_let_bound nest) - pat - (add_to_kernel, free_in_kernel) - identity_map - inner_returned_arrs - addTarget = do - let nest'@(MapNesting _ aux w params_and_arrs) = - removeUnusedNestingParts free_in_kernel nest - (params,arrs) = unzip params_and_arrs - param_names = namesFromList $ map paramName params - free_in_kernel' = - (freeIn nest' <> free_in_kernel) `namesSubtract` param_names - required_from_nest = - free_in_kernel' `namesIntersection` nest_let_bound + distributeAtNesting :: + (HasScope t m, MonadFreshNames m) => + Nesting -> + PatternT Type -> + (LoopNesting -> KernelNest, Names) -> + M.Map VName Ident -> + [Ident] -> + (Target -> Targets) -> + MaybeT m (KernelNest, Names, Targets) + distributeAtNesting + (Nesting nest_let_bound nest) + pat + (add_to_kernel, free_in_kernel) + identity_map + inner_returned_arrs + addTarget = do + let nest'@(MapNesting _ aux w params_and_arrs) = + removeUnusedNestingParts free_in_kernel nest + (params, arrs) = unzip params_and_arrs + param_names = namesFromList $ map paramName params + free_in_kernel' = + (freeIn nest' <> free_in_kernel) `namesSubtract` param_names + required_from_nest = + free_in_kernel' `namesIntersection` nest_let_bound - required_from_nest_idents <- - forM (namesToList required_from_nest) $ \name -> do - t <- lift $ lookupType name - return $ Ident name t + required_from_nest_idents <- + forM (namesToList required_from_nest) $ \name -> do + t <- lift $ lookupType name + return $ Ident name t - (free_params, free_arrs, bind_in_target) <- - fmap unzip3 $ - forM (inner_returned_arrs++required_from_nest_idents) $ - \(Ident pname ptype) -> - case M.lookup pname identity_map of - Nothing -> do - arr <- newIdent (baseString pname ++ "_r") $ - arrayOfRow ptype w - return (Param pname ptype, - arr, - True) - Just arr -> - return (Param pname ptype, - arr, - False) + (free_params, free_arrs, bind_in_target) <- + fmap unzip3 $ + forM (inner_returned_arrs ++ required_from_nest_idents) $ + \(Ident pname ptype) -> + case M.lookup pname identity_map of + Nothing -> do + arr <- + newIdent (baseString pname ++ "_r") $ + arrayOfRow ptype w + return + ( Param pname ptype, + arr, + True + ) + Just arr -> + return + ( Param pname ptype, + arr, + False + ) - let free_arrs_pat = - basicPattern [] $ map snd $ - filter fst $ zip bind_in_target free_arrs - free_params_pat = - map snd $ filter fst $ zip bind_in_target free_params + let free_arrs_pat = + basicPattern [] $ + map snd $ + filter fst $ zip bind_in_target free_arrs + free_params_pat = + map snd $ filter fst $ zip bind_in_target free_params - (actual_params, actual_arrs) = - (params++free_params, - arrs++map identName free_arrs) - actual_param_names = - namesFromList $ map paramName actual_params + (actual_params, actual_arrs) = + ( params ++ free_params, + arrs ++ map identName free_arrs + ) + actual_param_names = + namesFromList $ map paramName actual_params - nest'' = - removeUnusedNestingParts free_in_kernel $ + nest'' = + removeUnusedNestingParts free_in_kernel $ MapNesting pat aux w $ zip actual_params actual_arrs - free_in_kernel'' = - (freeIn nest'' <> free_in_kernel) `namesSubtract` actual_param_names - - unless (all (distributableType . paramType) $ - loopNestingParams nest'') $ - fail "Would induce irregular array" - return (add_to_kernel nest'', - - free_in_kernel'', - - addTarget (free_arrs_pat, map (Var . paramName) free_params_pat)) + free_in_kernel'' = + (freeIn nest'' <> free_in_kernel) `namesSubtract` actual_param_names - recurse :: (HasScope t m, MonadFreshNames m) => - [(Nesting,Target)] - -> MaybeT m (KernelNest, Names, Targets) - recurse [] = - distributeAtNesting - inner_nest - (distributionInnerPattern distrib_body) - (newKernel, - distributionFreeInBody distrib_body `namesIntersection` bound_in_nest) - (distributionIdentityMap distrib_body) - [] $ - singleTarget . distributionExpandTarget distrib_body + unless + ( all (distributableType . paramType) $ + loopNestingParams nest'' + ) + $ fail "Would induce irregular array" + return + ( add_to_kernel nest'', + free_in_kernel'', + addTarget (free_arrs_pat, map (Var . paramName) free_params_pat) + ) - recurse ((nest, (pat,res)) : nests') = do - (kernel@(outer, _), kernel_free, kernel_targets) <- recurse nests' + recurse :: + (HasScope t m, MonadFreshNames m) => + [(Nesting, Target)] -> + MaybeT m (KernelNest, Names, Targets) + recurse [] = + distributeAtNesting + inner_nest + (distributionInnerPattern distrib_body) + ( newKernel, + distributionFreeInBody distrib_body `namesIntersection` bound_in_nest + ) + (distributionIdentityMap distrib_body) + [] + $ singleTarget . distributionExpandTarget distrib_body + recurse ((nest, (pat, res)) : nests') = do + (kernel@(outer, _), kernel_free, kernel_targets) <- recurse nests' - let (pat', res', identity_map, expand_target) = - removeIdentityMappingFromNesting - (namesFromList $ patternNames $ loopNestingPattern outer) pat res + let (pat', res', identity_map, expand_target) = + removeIdentityMappingFromNesting + (namesFromList $ patternNames $ loopNestingPattern outer) + pat + res - distributeAtNesting - nest - pat' - (\k -> pushKernelNesting (pat',res') k kernel, - kernel_free) - identity_map - (patternIdents $ fst $ outerTarget kernel_targets) - ((`pushOuterTarget` kernel_targets) . expand_target) + distributeAtNesting + nest + pat' + ( \k -> pushKernelNesting (pat', res') k kernel, + kernel_free + ) + identity_map + (patternIdents $ fst $ outerTarget kernel_targets) + ((`pushOuterTarget` kernel_targets) . expand_target) removeUnusedNestingParts :: Names -> LoopNesting -> LoopNesting removeUnusedNestingParts used (MapNesting pat aux w params_and_arrs) = MapNesting pat aux w $ zip used_params used_arrs - where (params,arrs) = unzip params_and_arrs - (used_params, used_arrs) = - unzip $ - filter ((`nameIn` used) . paramName . fst) $ + where + (params, arrs) = unzip params_and_arrs + (used_params, used_arrs) = + unzip $ + filter ((`nameIn` used) . paramName . fst) $ zip params arrs -removeIdentityMappingGeneral :: Names -> PatternT Type -> Result - -> (PatternT Type, - Result, - M.Map VName Ident, - Target -> Target) +removeIdentityMappingGeneral :: + Names -> + PatternT Type -> + Result -> + ( PatternT Type, + Result, + M.Map VName Ident, + Target -> Target + ) removeIdentityMappingGeneral bound pat res = let (identities, not_identities) = mapEither isIdentity $ zip (patternElements pat) res (not_identity_patElems, not_identity_res) = unzip not_identities (identity_patElems, identity_res) = unzip identities expandTarget (tpat, tres) = - (Pattern [] $ patternElements tpat ++ identity_patElems, - tres ++ map Var identity_res) - identity_map = M.fromList $ zip identity_res $ - map patElemIdent identity_patElems - in (Pattern [] not_identity_patElems, - not_identity_res, - identity_map, - expandTarget) - where isIdentity (patElem, Var v) - | not (v `nameIn` bound) = Left (patElem, v) - isIdentity x = Right x + ( Pattern [] $ patternElements tpat ++ identity_patElems, + tres ++ map Var identity_res + ) + identity_map = + M.fromList $ + zip identity_res $ + map patElemIdent identity_patElems + in ( Pattern [] not_identity_patElems, + not_identity_res, + identity_map, + expandTarget + ) + where + isIdentity (patElem, Var v) + | not (v `nameIn` bound) = Left (patElem, v) + isIdentity x = Right x -removeIdentityMappingFromNesting :: Names -> PatternT Type -> Result - -> (PatternT Type, - Result, - M.Map VName Ident, - Target -> Target) +removeIdentityMappingFromNesting :: + Names -> + PatternT Type -> + Result -> + ( PatternT Type, + Result, + M.Map VName Ident, + Target -> Target + ) removeIdentityMappingFromNesting bound_in_nesting pat res = let (pat', res', identity_map, expand_target) = removeIdentityMappingGeneral bound_in_nesting pat res - in (pat', res', identity_map, expand_target) + in (pat', res', identity_map, expand_target) -tryDistribute :: (DistLore lore, MonadFreshNames m, - LocalScope lore m, MonadLogger m) => - MkSegLevel lore m -> Nestings -> Targets -> Stms lore - -> m (Maybe (Targets, Stms lore)) -tryDistribute _ _ targets stms | null stms = - -- No point in distributing an empty kernel. - return $ Just (targets, mempty) +tryDistribute :: + ( DistLore lore, + MonadFreshNames m, + LocalScope lore m, + MonadLogger m + ) => + MkSegLevel lore m -> + Nestings -> + Targets -> + Stms lore -> + m (Maybe (Targets, Stms lore)) +tryDistribute _ _ targets stms + | null stms = + -- No point in distributing an empty kernel. + return $ Just (targets, mempty) tryDistribute mk_lvl nest targets stms = - createKernelNest nest dist_body >>= - \case - Just (targets', distributed) -> do - (kernel_bnd, w_bnds) <- - localScope (targetsScope targets') $ - constructKernel mk_lvl distributed $ mkBody stms inner_body_res - distributed' <- renameStm kernel_bnd - logMsg $ "distributing\n" ++ - unlines (map pretty $ stmsToList stms) ++ - pretty (snd $ innerTarget targets) ++ - "\nas\n" ++ pretty distributed' ++ - "\ndue to targets\n" ++ ppTargets targets ++ - "\nand with new targets\n" ++ ppTargets targets' - return $ Just (targets', w_bnds <> oneStm distributed') - Nothing -> - return Nothing - where (dist_body, inner_body_res) = distributionBodyFromStms targets stms + createKernelNest nest dist_body + >>= \case + Just (targets', distributed) -> do + (kernel_bnd, w_bnds) <- + localScope (targetsScope targets') $ + constructKernel mk_lvl distributed $ mkBody stms inner_body_res + distributed' <- renameStm kernel_bnd + logMsg $ + "distributing\n" + ++ unlines (map pretty $ stmsToList stms) + ++ pretty (snd $ innerTarget targets) + ++ "\nas\n" + ++ pretty distributed' + ++ "\ndue to targets\n" + ++ ppTargets targets + ++ "\nand with new targets\n" + ++ ppTargets targets' + return $ Just (targets', w_bnds <> oneStm distributed') + Nothing -> + return Nothing + where + (dist_body, inner_body_res) = distributionBodyFromStms targets stms -tryDistributeStm :: (MonadFreshNames m, HasScope t m, ASTLore lore) => - Nestings -> Targets -> Stm lore - -> m (Maybe (Result, Targets, KernelNest)) +tryDistributeStm :: + (MonadFreshNames m, HasScope t m, ASTLore lore) => + Nestings -> + Targets -> + Stm lore -> + m (Maybe (Result, Targets, KernelNest)) tryDistributeStm nest targets bnd = fmap addRes <$> createKernelNest nest dist_body - where (dist_body, res) = distributionBodyFromStm targets bnd - addRes (targets', kernel_nest) = (res, targets', kernel_nest) + where + (dist_body, res) = distributionBodyFromStm targets bnd + addRes (targets', kernel_nest) = (res, targets', kernel_nest) diff --git a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs --- a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs +++ b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs @@ -1,124 +1,152 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + -- | Interchanging scans with inner maps. module Futhark.Pass.ExtractKernels.ISRWIM - ( iswim - , irwim - , rwimPossible - ) - where + ( iswim, + irwim, + rwimPossible, + ) +where import Control.Arrow (first) import Control.Monad.State - -import Futhark.MonadFreshNames import Futhark.IR.SOACS +import Futhark.MonadFreshNames import Futhark.Tools -- | Interchange Scan With Inner Map. Tries to turn a @scan(map)@ into a -- @map(scan) -iswim :: (MonadBinder m, Lore m ~ SOACS) => - Pattern - -> SubExp - -> Lambda - -> [(SubExp, VName)] - -> Maybe (m ()) +iswim :: + (MonadBinder m, Lore m ~ SOACS) => + Pattern -> + SubExp -> + Lambda -> + [(SubExp, VName)] -> + Maybe (m ()) iswim res_pat w scan_fun scan_input | Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible scan_fun = Just $ do - let (accs, arrs) = unzip scan_input - arrs' <- transposedArrays arrs - accs' <- mapM (letExp "acc" . BasicOp . SubExp) accs + let (accs, arrs) = unzip scan_input + arrs' <- transposedArrays arrs + accs' <- mapM (letExp "acc" . BasicOp . SubExp) accs - let map_arrs' = accs' ++ arrs' - (scan_acc_params, scan_elem_params) = - splitAt (length arrs) $ lambdaParams scan_fun - map_params = map removeParamOuterDim scan_acc_params ++ - map (setParamOuterDimTo w) scan_elem_params - map_rettype = map (setOuterDimTo w) $ lambdaReturnType scan_fun + let map_arrs' = accs' ++ arrs' + (scan_acc_params, scan_elem_params) = + splitAt (length arrs) $ lambdaParams scan_fun + map_params = + map removeParamOuterDim scan_acc_params + ++ map (setParamOuterDimTo w) scan_elem_params + map_rettype = map (setOuterDimTo w) $ lambdaReturnType scan_fun - scan_params = lambdaParams map_fun - scan_body = lambdaBody map_fun - scan_rettype = lambdaReturnType map_fun - scan_fun' = Lambda scan_params scan_body scan_rettype - scan_input' = map (first Var) $ - uncurry zip $ splitAt (length arrs') $ map paramName map_params - (nes', scan_arrs) = unzip scan_input' + scan_params = lambdaParams map_fun + scan_body = lambdaBody map_fun + scan_rettype = lambdaReturnType map_fun + scan_fun' = Lambda scan_params scan_body scan_rettype + scan_input' = + map (first Var) $ + uncurry zip $ splitAt (length arrs') $ map paramName map_params + (nes', scan_arrs) = unzip scan_input' - scan_soac <- scanSOAC [Scan scan_fun' nes'] - let map_body = mkBody (oneStm $ Let (setPatternOuterDimTo w map_pat) (defAux ()) $ - Op $ Screma w scan_soac scan_arrs) $ - map Var $ patternNames map_pat - map_fun' = Lambda map_params map_body map_rettype + scan_soac <- scanSOAC [Scan scan_fun' nes'] + let map_body = + mkBody + ( oneStm $ + Let (setPatternOuterDimTo w map_pat) (defAux ()) $ + Op $ Screma w scan_soac scan_arrs + ) + $ map Var $ patternNames map_pat + map_fun' = Lambda map_params map_body map_rettype - res_pat' <- fmap (basicPattern []) $ - mapM (newIdent' (<>"_transposed") . transposeIdentType) $ - patternValueIdents res_pat + res_pat' <- + fmap (basicPattern []) $ + mapM (newIdent' (<> "_transposed") . transposeIdentType) $ + patternValueIdents res_pat - addStm $ Let res_pat' (StmAux map_cs mempty ()) $ Op $ Screma map_w - (mapSOAC map_fun') map_arrs' + addStm $ + Let res_pat' (StmAux map_cs mempty ()) $ + Op $ + Screma + map_w + (mapSOAC map_fun') + map_arrs' - forM_ (zip (patternValueIdents res_pat) - (patternValueIdents res_pat')) $ \(to, from) -> do - let perm = [1,0] ++ [2..arrayRank (identType from)-1] - addStm $ Let (basicPattern [] [to]) (defAux ()) $ - BasicOp $ Rearrange perm $ identName from + forM_ + ( zip + (patternValueIdents res_pat) + (patternValueIdents res_pat') + ) + $ \(to, from) -> do + let perm = [1, 0] ++ [2 .. arrayRank (identType from) -1] + addStm $ + Let (basicPattern [] [to]) (defAux ()) $ + BasicOp $ Rearrange perm $ identName from | otherwise = Nothing -- | Interchange Reduce With Inner Map. Tries to turn a @reduce(map)@ into a -- @map(reduce) -irwim :: (MonadBinder m, Lore m ~ SOACS) => - Pattern - -> SubExp - -> Commutativity -> Lambda - -> [(SubExp, VName)] - -> Maybe (m ()) +irwim :: + (MonadBinder m, Lore m ~ SOACS) => + Pattern -> + SubExp -> + Commutativity -> + Lambda -> + [(SubExp, VName)] -> + Maybe (m ()) irwim res_pat w comm red_fun red_input | Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible red_fun = Just $ do - let (accs, arrs) = unzip red_input - arrs' <- transposedArrays arrs - -- FIXME? Can we reasonably assume that the accumulator is a - -- replicate? We also assume that it is non-empty. - let indexAcc (Var v) = do - v_t <- lookupType v - letSubExp "acc" $ BasicOp $ Index v $ - fullSlice v_t [DimFix $ intConst Int32 0] - indexAcc Constant{} = - error "irwim: array accumulator is a constant." - accs' <- mapM indexAcc accs + let (accs, arrs) = unzip red_input + arrs' <- transposedArrays arrs + -- FIXME? Can we reasonably assume that the accumulator is a + -- replicate? We also assume that it is non-empty. + let indexAcc (Var v) = do + v_t <- lookupType v + letSubExp "acc" $ + BasicOp $ + Index v $ + fullSlice v_t [DimFix $ intConst Int64 0] + indexAcc Constant {} = + error "irwim: array accumulator is a constant." + accs' <- mapM indexAcc accs - let (_red_acc_params, red_elem_params) = - splitAt (length arrs) $ lambdaParams red_fun - map_rettype = map rowType $ lambdaReturnType red_fun - map_params = map (setParamOuterDimTo w) red_elem_params + let (_red_acc_params, red_elem_params) = + splitAt (length arrs) $ lambdaParams red_fun + map_rettype = map rowType $ lambdaReturnType red_fun + map_params = map (setParamOuterDimTo w) red_elem_params - red_params = lambdaParams map_fun - red_body = lambdaBody map_fun - red_rettype = lambdaReturnType map_fun - red_fun' = Lambda red_params red_body red_rettype - red_input' = zip accs' $ map paramName map_params - red_pat = stripPatternOuterDim map_pat + red_params = lambdaParams map_fun + red_body = lambdaBody map_fun + red_rettype = lambdaReturnType map_fun + red_fun' = Lambda red_params red_body red_rettype + red_input' = zip accs' $ map paramName map_params + red_pat = stripPatternOuterDim map_pat - map_body <- - case irwim red_pat w comm red_fun' red_input' of - Nothing -> do - reduce_soac <- reduceSOAC [Reduce comm red_fun' $ map fst red_input'] - return $ mkBody (oneStm $ Let red_pat (defAux ()) $ - Op $ Screma w reduce_soac $ map snd red_input') $ - map Var $ patternNames map_pat - Just m -> localScope (scopeOfLParams map_params) $ do - map_body_bnds <- collectStms_ m - return $ mkBody map_body_bnds $ map Var $ patternNames map_pat + map_body <- + case irwim red_pat w comm red_fun' red_input' of + Nothing -> do + reduce_soac <- reduceSOAC [Reduce comm red_fun' $ map fst red_input'] + return $ + mkBody + ( oneStm $ + Let red_pat (defAux ()) $ + Op $ Screma w reduce_soac $ map snd red_input' + ) + $ map Var $ patternNames map_pat + Just m -> localScope (scopeOfLParams map_params) $ do + map_body_bnds <- collectStms_ m + return $ mkBody map_body_bnds $ map Var $ patternNames map_pat - let map_fun' = Lambda map_params map_body map_rettype + let map_fun' = Lambda map_params map_body map_rettype - addStm $ Let res_pat (StmAux map_cs mempty ()) $ + addStm $ + Let res_pat (StmAux map_cs mempty ()) $ Op $ Screma map_w (mapSOAC map_fun') arrs' | otherwise = Nothing -- | Does this reduce operator contain an inner map, and if so, what -- does that map look like? -rwimPossible :: Lambda - -> Maybe (Pattern, Certificates, SubExp, Lambda) +rwimPossible :: + Lambda -> + Maybe (Pattern, Certificates, SubExp, Lambda) rwimPossible fun | Body _ stms res <- lambdaBody fun, [bnd] <- stmsToList stms, -- Body has a single binding @@ -127,30 +155,30 @@ Op (Screma map_w form map_arrs) <- stmExp bnd, Just map_fun <- isMapSOAC form, map paramName (lambdaParams fun) == map_arrs = - Just (map_pat, stmCerts bnd, map_w, map_fun) + Just (map_pat, stmCerts bnd, map_w, map_fun) | otherwise = - Nothing + Nothing transposedArrays :: MonadBinder m => [VName] -> m [VName] transposedArrays arrs = forM arrs $ \arr -> do t <- lookupType arr - let perm = [1,0] ++ [2..arrayRank t-1] + let perm = [1, 0] ++ [2 .. arrayRank t -1] letExp (baseString arr) $ BasicOp $ Rearrange perm arr removeParamOuterDim :: LParam -> LParam removeParamOuterDim param = let t = rowType $ paramType param - in param { paramDec = t } + in param {paramDec = t} setParamOuterDimTo :: SubExp -> LParam -> LParam setParamOuterDimTo w param = let t = setOuterDimTo w $ paramType param - in param { paramDec = t } + in param {paramDec = t} setIdentOuterDimTo :: SubExp -> Ident -> Ident setIdentOuterDimTo w ident = let t = setOuterDimTo w $ identType ident - in ident { identType = t } + in ident {identType = t} setOuterDimTo :: SubExp -> Type -> Type setOuterDimTo w t = @@ -162,11 +190,11 @@ transposeIdentType :: Ident -> Ident transposeIdentType ident = - ident { identType = transposeType $ identType ident } + ident {identType = transposeType $ identType ident} stripIdentOuterDim :: Ident -> Ident stripIdentOuterDim ident = - ident { identType = rowType $ identType ident } + ident {identType = rowType $ identType ident} stripPatternOuterDim :: Pattern -> Pattern stripPatternOuterDim pat = diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs --- a/src/Futhark/Pass/ExtractKernels/Interchange.hs +++ b/src/Futhark/Pass/ExtractKernels/Interchange.hs @@ -1,5 +1,6 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + -- | It is well known that fully parallel loops can always be -- interchanged inwards with a sequential loop. This module -- implements that transformation. @@ -7,23 +8,25 @@ -- This is also where we implement loop-switching (for branches), -- which is semantically similar to interchange. module Futhark.Pass.ExtractKernels.Interchange - ( - SeqLoop (..) - , interchangeLoops - , Branch (..) - , interchangeBranch - ) where + ( SeqLoop (..), + interchangeLoops, + Branch (..), + interchangeBranch, + ) +where import Control.Monad.RWS.Strict -import Data.Maybe import Data.List (find) - -import Futhark.Pass.ExtractKernels.Distribution - (LoopNesting(..), KernelNest, kernelNestLoops) +import Data.Maybe import Futhark.IR.SOACS import Futhark.MonadFreshNames -import Futhark.Transform.Rename +import Futhark.Pass.ExtractKernels.Distribution + ( KernelNest, + LoopNesting (..), + kernelNestLoops, + ) import Futhark.Tools +import Futhark.Transform.Rename -- | An encoding of a sequential do-loop with no existential context, -- alongside its result pattern. @@ -33,21 +36,26 @@ seqLoopStm (SeqLoop _ pat merge form body) = Let pat (defAux ()) $ DoLoop [] merge form body -interchangeLoop :: (MonadBinder m, LocalScope SOACS m) => - (VName -> Maybe VName) -> SeqLoop -> LoopNesting - -> m SeqLoop +interchangeLoop :: + (MonadBinder m, LocalScope SOACS m) => + (VName -> Maybe VName) -> + SeqLoop -> + LoopNesting -> + m SeqLoop interchangeLoop isMapParameter (SeqLoop perm loop_pat merge form body) (MapNesting pat aux w params_and_arrs) = do merge_expanded <- localScope (scopeOfLParams $ map fst params_and_arrs) $ - mapM expand merge + mapM expand merge let loop_pat_expanded = Pattern [] $ map expandPatElem $ patternElements loop_pat - new_params = [ Param pname $ fromDecl ptype - | (Param pname ptype, _) <- merge ] + new_params = + [ Param pname $ fromDecl ptype + | (Param pname ptype, _) <- merge + ] new_arrs = map (paramName . fst) merge_expanded rettype = map rowType $ patternTypes loop_pat_expanded @@ -57,59 +65,68 @@ -- it is not used anymore. This might happen if the parameter was -- used just as the inital value of a merge parameter. ((params', arrs'), pre_copy_bnds) <- - runBinder $ localScope (scopeOfLParams new_params) $ - unzip . catMaybes <$> mapM copyOrRemoveParam params_and_arrs + runBinder $ + localScope (scopeOfLParams new_params) $ + unzip . catMaybes <$> mapM copyOrRemoveParam params_and_arrs - let lam = Lambda (params'<>new_params) body rettype - map_bnd = Let loop_pat_expanded aux $ - Op $ Screma w (mapSOAC lam) $ arrs' <> new_arrs + let lam = Lambda (params' <> new_params) body rettype + map_bnd = + Let loop_pat_expanded aux $ + Op $ Screma w (mapSOAC lam) $ arrs' <> new_arrs res = map Var $ patternNames loop_pat_expanded pat' = Pattern [] $ rearrangeShape perm $ patternValueElements pat return $ SeqLoop perm pat' merge_expanded form $ - mkBody (pre_copy_bnds<>oneStm map_bnd) res - where free_in_body = freeIn body + mkBody (pre_copy_bnds <> oneStm map_bnd) res + where + free_in_body = freeIn body - copyOrRemoveParam (param, arr) - | not (paramName param `nameIn` free_in_body) = - return Nothing - | otherwise = - return $ Just (param, arr) + copyOrRemoveParam (param, arr) + | not (paramName param `nameIn` free_in_body) = + return Nothing + | otherwise = + return $ Just (param, arr) - expandedInit _ (Var v) - | Just arr <- isMapParameter v = - return $ Var arr - expandedInit param_name se = - letSubExp (param_name <> "_expanded_init") $ - BasicOp $ Replicate (Shape [w]) se + expandedInit _ (Var v) + | Just arr <- isMapParameter v = + return $ Var arr + expandedInit param_name se = + letSubExp (param_name <> "_expanded_init") $ + BasicOp $ Replicate (Shape [w]) se - expand (merge_param, merge_init) = do - expanded_param <- - newParam (param_name <> "_expanded") $ + expand (merge_param, merge_init) = do + expanded_param <- + newParam (param_name <> "_expanded") $ arrayOf (paramDeclType merge_param) (Shape [w]) $ - uniqueness $ declTypeOf merge_param - expanded_init <- expandedInit param_name merge_init - return (expanded_param, expanded_init) - where param_name = baseString $ paramName merge_param + uniqueness $ declTypeOf merge_param + expanded_init <- expandedInit param_name merge_init + return (expanded_param, expanded_init) + where + param_name = baseString $ paramName merge_param - expandPatElem (PatElem name t) = - PatElem name $ arrayOfRow t w + expandPatElem (PatElem name t) = + PatElem name $ arrayOfRow t w -- | Given a (parallel) map nesting and an inner sequential loop, move -- the maps inside the sequential loop. The result is several -- statements - one of these will be the loop, which will then contain -- statements with @map@ expressions. -interchangeLoops :: (MonadFreshNames m, HasScope SOACS m) => - KernelNest -> SeqLoop - -> m (Stms SOACS) +interchangeLoops :: + (MonadFreshNames m, HasScope SOACS m) => + KernelNest -> + SeqLoop -> + m (Stms SOACS) interchangeLoops nest loop = do (loop', bnds) <- - runBinder $ foldM (interchangeLoop isMapParameter) loop $ - reverse $ kernelNestLoops nest + runBinder $ + foldM (interchangeLoop isMapParameter) loop $ + reverse $ kernelNestLoops nest return $ bnds <> oneStm (seqLoopStm loop') - where isMapParameter v = - fmap snd $ find ((==v) . paramName . fst) $ + where + isMapParameter v = + fmap snd $ + find ((== v) . paramName . fst) $ concatMap loopNestingParamsAndArrs $ kernelNestLoops nest data Branch = Branch [Int] Pattern SubExp Body Body (IfDec (BranchType SOACS)) @@ -118,8 +135,11 @@ branchStm (Branch _ pat cond tbranch fbranch ret) = Let pat (defAux ()) $ If cond tbranch fbranch ret -interchangeBranch1 :: (MonadBinder m) => - Branch -> LoopNesting -> m Branch +interchangeBranch1 :: + (MonadBinder m) => + Branch -> + LoopNesting -> + m Branch interchangeBranch1 (Branch perm branch_pat cond tbranch fbranch (IfDec ret if_sort)) (MapNesting pat aux w params_and_arrs) = do @@ -132,7 +152,7 @@ branch_pat' = Pattern [] $ map (fmap (`arrayOfRow` w)) $ patternElements branch_pat - mkBranch branch = (renameBody=<<) $ do + mkBranch branch = (renameBody =<<) $ do let lam = Lambda params branch lam_ret res = map Var $ patternNames branch_pat' map_bnd = Let branch_pat' aux $ Op $ Screma w (mapSOAC lam) arrs @@ -140,11 +160,15 @@ tbranch' <- mkBranch tbranch fbranch' <- mkBranch fbranch - return $ Branch [0..patternSize pat-1] pat' cond tbranch' fbranch' $ - IfDec ret' if_sort + return $ + Branch [0 .. patternSize pat -1] pat' cond tbranch' fbranch' $ + IfDec ret' if_sort -interchangeBranch :: (MonadFreshNames m, HasScope SOACS m) => - KernelNest -> Branch -> m (Stms SOACS) +interchangeBranch :: + (MonadFreshNames m, HasScope SOACS m) => + KernelNest -> + Branch -> + m (Stms SOACS) interchangeBranch nest loop = do (loop', bnds) <- runBinder $ foldM interchangeBranch1 loop $ reverse $ kernelNestLoops nest 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 @@ -1,34 +1,31 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} + -- | Extract limited nested parallelism for execution inside -- individual kernel workgroups. -module Futhark.Pass.ExtractKernels.Intragroup - (intraGroupParallelise) -where +module Futhark.Pass.ExtractKernels.Intragroup (intraGroupParallelise) where import Control.Monad.Identity import Control.Monad.RWS import Control.Monad.Trans.Maybe import qualified Data.Map.Strict as M import qualified Data.Set as S - -import Prelude hiding (log) - import Futhark.Analysis.PrimExp.Convert -import Futhark.IR.SOACS import qualified Futhark.IR.Kernels as Out import Futhark.IR.Kernels.Kernel hiding (HistOp) +import Futhark.IR.SOACS import Futhark.MonadFreshNames -import Futhark.Tools +import Futhark.Pass.ExtractKernels.BlockedKernel import Futhark.Pass.ExtractKernels.DistributeNests import Futhark.Pass.ExtractKernels.Distribution -import Futhark.Pass.ExtractKernels.BlockedKernel import Futhark.Pass.ExtractKernels.ToKernels +import Futhark.Tools import qualified Futhark.Transform.FirstOrderTransform as FOT import Futhark.Util (chunks) import Futhark.Util.Log +import Prelude hiding (log) -- | Convert the statements inside a map nest to kernel statements, -- attempting to parallelise any remaining (top-level) parallel @@ -42,16 +39,27 @@ -- -- We distinguish between "minimum group size" and "maximum -- exploitable parallelism". -intraGroupParallelise :: (MonadFreshNames m, LocalScope Out.Kernels m) => - KernelNest -> Lambda - -> m (Maybe ((SubExp, SubExp), SubExp, Log, - Out.Stms Out.Kernels, Out.Stms Out.Kernels)) +intraGroupParallelise :: + (MonadFreshNames m, LocalScope Out.Kernels m) => + KernelNest -> + Lambda -> + m + ( Maybe + ( (SubExp, SubExp), + SubExp, + Log, + Out.Stms Out.Kernels, + Out.Stms Out.Kernels + ) + ) intraGroupParallelise knest lam = runMaybeT $ do (ispace, inps) <- lift $ flatKernel knest - (num_groups, w_stms) <- lift $ runBinder $ - letSubExp "intra_num_groups" =<< - foldBinOp (Mul Int32 OverflowUndef) (intConst Int32 1) (map snd ispace) + (num_groups, w_stms) <- + lift $ + runBinder $ + letSubExp "intra_num_groups" + =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) (map snd ispace) let body = lambdaBody lam @@ -59,66 +67,80 @@ let intra_lvl = SegThread (Count num_groups) (Count $ Var group_size) SegNoVirt (wss_min, wss_avail, log, kbody) <- - lift $ localScope (scopeOfLParams $ lambdaParams lam) $ - intraGroupParalleliseBody intra_lvl body + lift $ + localScope (scopeOfLParams $ lambdaParams lam) $ + intraGroupParalleliseBody intra_lvl body outside_scope <- lift askScope -- outside_scope may also contain the inputs, even though those are -- not actually available outside the kernel. let available v = - v `M.member` outside_scope && - v `notElem` map kernelInputName inps + v `M.member` outside_scope + && v `notElem` map kernelInputName inps unless (all available $ namesToList $ freeIn (wss_min ++ wss_avail)) $ fail "Irregular parallelism" - ((intra_avail_par, kspace, read_input_stms), prelude_stms) <- lift $ runBinder $ do - let foldBinOp' _ [] = eSubExp $ intConst Int32 0 - foldBinOp' bop (x:xs) = foldBinOp bop x xs - ws_min <- mapM (letSubExp "one_intra_par_min" <=< foldBinOp' (Mul Int32 OverflowUndef)) $ - filter (not . null) wss_min - ws_avail <- mapM (letSubExp "one_intra_par_avail" <=< foldBinOp' (Mul Int32 OverflowUndef)) $ - filter (not . null) wss_avail + ((intra_avail_par, kspace, read_input_stms), prelude_stms) <- lift $ + runBinder $ do + let foldBinOp' _ [] = eSubExp $ intConst Int64 0 + foldBinOp' bop (x : xs) = foldBinOp bop x xs + ws_min <- + mapM (letSubExp "one_intra_par_min" <=< foldBinOp' (Mul Int64 OverflowUndef)) $ + filter (not . null) wss_min + ws_avail <- + mapM (letSubExp "one_intra_par_avail" <=< foldBinOp' (Mul Int64 OverflowUndef)) $ + filter (not . null) wss_avail - -- The amount of parallelism available *in the worst case* is - -- equal to the smallest parallel loop. - intra_avail_par <- letSubExp "intra_avail_par" =<< foldBinOp' (SMin Int32) ws_avail + -- The amount of parallelism available *in the worst case* is + -- equal to the smallest parallel loop. + intra_avail_par <- letSubExp "intra_avail_par" =<< foldBinOp' (SMin Int64) ws_avail - -- The group size is either the maximum of the minimum parallelism - -- exploited, or the desired parallelism (bounded by the max group - -- size) in case there is no minimum. - letBindNames [group_size] =<< - if null ws_min - then eBinOp (SMin Int32) - (eSubExp =<< letSubExp "max_group_size" (Op $ SizeOp $ Out.GetSizeMax Out.SizeGroup)) - (eSubExp intra_avail_par) - else foldBinOp' (SMax Int32) ws_min + -- The group size is either the maximum of the minimum parallelism + -- exploited, or the desired parallelism (bounded by the max group + -- size) in case there is no minimum. + letBindNames [group_size] + =<< if null ws_min + then + eBinOp + (SMin Int64) + (eSubExp =<< letSubExp "max_group_size" (Op $ SizeOp $ Out.GetSizeMax Out.SizeGroup)) + (eSubExp intra_avail_par) + else foldBinOp' (SMax Int64) ws_min - let inputIsUsed input = kernelInputName input `nameIn` freeIn body - used_inps = filter inputIsUsed inps + let inputIsUsed input = kernelInputName input `nameIn` freeIn body + used_inps = filter inputIsUsed inps - addStms w_stms - read_input_stms <- runBinder_ $ mapM readKernelInput used_inps - space <- mkSegSpace ispace - return (intra_avail_par, space, read_input_stms) + addStms w_stms + read_input_stms <- runBinder_ $ mapM readKernelInput used_inps + space <- mkSegSpace ispace + return (intra_avail_par, space, read_input_stms) - let kbody' = kbody { kernelBodyStms = read_input_stms <> kernelBodyStms kbody } + let kbody' = kbody {kernelBodyStms = read_input_stms <> kernelBodyStms kbody} let nested_pat = loopNestingPattern first_nest rts = map (length ispace `stripArray`) $ patternTypes nested_pat lvl = SegGroup (Count num_groups) (Count $ Var group_size) SegNoVirt - kstm = Let nested_pat aux $ - Op $ SegOp $ SegMap lvl kspace rts kbody' + kstm = + Let nested_pat aux $ + Op $ SegOp $ SegMap lvl kspace rts kbody' let intra_min_par = intra_avail_par - return ((intra_min_par, intra_avail_par), Var group_size, log, - prelude_stms, oneStm kstm) - where first_nest = fst knest - aux = loopNestingAux first_nest + return + ( (intra_min_par, intra_avail_par), + Var group_size, + log, + prelude_stms, + oneStm kstm + ) + where + first_nest = fst knest + aux = loopNestingAux first_nest -data Acc = Acc { accMinPar :: S.Set [SubExp] - , accAvailPar :: S.Set [SubExp] - , accLog :: Log - } +data Acc = Acc + { accMinPar :: S.Set [SubExp], + accAvailPar :: S.Set [SubExp], + accLog :: Log + } instance Semigroup Acc where Acc min_x avail_x log_x <> Acc min_y avail_y log_y = @@ -131,20 +153,25 @@ BinderT Out.Kernels (RWS () Acc VNameSource) instance MonadLogger IntraGroupM where - addLog log = tell mempty { accLog = log } + addLog log = tell mempty {accLog = log} -runIntraGroupM :: (MonadFreshNames m, HasScope Out.Kernels m) => - IntraGroupM () -> m (Acc, Out.Stms Out.Kernels) +runIntraGroupM :: + (MonadFreshNames m, HasScope Out.Kernels m) => + IntraGroupM () -> + m (Acc, Out.Stms Out.Kernels) runIntraGroupM m = do scope <- castScope <$> askScope modifyNameSource $ \src -> let (((), kstms), src', acc) = runRWS (runBinderT m scope) () src - in ((acc, kstms), src') + in ((acc, kstms), src') parallelMin :: [SubExp] -> IntraGroupM () -parallelMin ws = tell mempty { accMinPar = S.singleton ws - , accAvailPar = S.singleton ws - } +parallelMin ws = + tell + mempty + { accMinPar = S.singleton ws, + accAvailPar = S.singleton ws + } intraGroupBody :: SegLevel -> Body -> IntraGroupM (Out.Body Out.Kernels) intraGroupBody lvl body = do @@ -159,71 +186,70 @@ case e of DoLoop ctx val form loopbody -> localScope (scopeOf form') $ - localScope (scopeOfFParams $ map fst $ ctx ++ val) $ do - loopbody' <- intraGroupBody lvl loopbody - certifying (stmAuxCerts aux) $ - letBind pat $ DoLoop ctx val form' loopbody' - where form' = case form of - ForLoop i it bound inps -> ForLoop i it bound inps - WhileLoop cond -> WhileLoop cond - + localScope (scopeOfFParams $ map fst $ ctx ++ val) $ do + loopbody' <- intraGroupBody lvl loopbody + certifying (stmAuxCerts aux) $ + letBind pat $ DoLoop ctx val form' loopbody' + where + form' = case form of + ForLoop i it bound inps -> ForLoop i it bound inps + WhileLoop cond -> WhileLoop cond If cond tbody fbody ifdec -> do tbody' <- intraGroupBody lvl tbody fbody' <- intraGroupBody lvl fbody certifying (stmAuxCerts aux) $ letBind pat $ If cond tbody' fbody' ifdec - Op soac | "sequential_outer" `inAttrs` stmAuxAttrs aux -> - intraGroupStms lvl . fmap (certify (stmAuxCerts aux)) =<< - runBinder_ (FOT.transformSOAC pat soac) - + intraGroupStms lvl . fmap (certify (stmAuxCerts aux)) + =<< runBinder_ (FOT.transformSOAC pat soac) Op (Screma w form arrs) | Just lam <- isMapSOAC form -> do - let loopnest = MapNesting pat aux w $ zip (lambdaParams lam) arrs - env = DistEnv { distNest = - singleNesting $ Nesting mempty loopnest - , distScope = - scopeOfPattern pat <> - scopeForKernels (scopeOf lam) <> scope - , distOnInnerMap = - distributeMap - , distOnTopLevelStms = - liftInner . collectStms_ . intraGroupStms lvl - , distSegLevel = \minw _ _ -> do - lift $ parallelMin minw - return lvl - , distOnSOACSStms = - pure . oneStm . soacsStmToKernels - , distOnSOACSLambda = - pure . soacsLambdaToKernels - } - acc = DistAcc { distTargets = singleTarget (pat, bodyResult $ lambdaBody lam) - , distStms = mempty - } - - addStms =<< - runDistNestT env (distributeMapBodyStms acc (bodyStms $ lambdaBody lam)) + let loopnest = MapNesting pat aux w $ zip (lambdaParams lam) arrs + env = + DistEnv + { distNest = + singleNesting $ Nesting mempty loopnest, + distScope = + scopeOfPattern pat + <> scopeForKernels (scopeOf lam) + <> scope, + distOnInnerMap = + distributeMap, + distOnTopLevelStms = + liftInner . collectStms_ . intraGroupStms lvl, + distSegLevel = \minw _ _ -> do + lift $ parallelMin minw + return lvl, + distOnSOACSStms = + pure . oneStm . soacsStmToKernels, + distOnSOACSLambda = + pure . soacsLambdaToKernels + } + acc = + DistAcc + { distTargets = singleTarget (pat, bodyResult $ lambdaBody lam), + distStms = mempty + } + addStms + =<< runDistNestT env (distributeMapBodyStms acc (bodyStms $ lambdaBody lam)) Op (Screma w form arrs) | Just (scans, mapfun) <- isScanomapSOAC form, Scan scanfun nes <- singleScan scans -> do - let scanfun' = soacsLambdaToKernels scanfun - mapfun' = soacsLambdaToKernels mapfun - certifying (stmAuxCerts aux) $ - addStms =<< segScan lvl' pat w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' arrs [] [] - parallelMin [w] - + let scanfun' = soacsLambdaToKernels scanfun + mapfun' = soacsLambdaToKernels mapfun + certifying (stmAuxCerts aux) $ + addStms =<< segScan lvl' pat w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' arrs [] [] + parallelMin [w] Op (Screma w form arrs) | Just (reds, map_lam) <- isRedomapSOAC form, Reduce comm red_lam nes <- singleReduce reds -> do - let red_lam' = soacsLambdaToKernels red_lam - map_lam' = soacsLambdaToKernels map_lam - certifying (stmAuxCerts aux) $ - addStms =<< segRed lvl' pat w [SegBinOp comm red_lam' nes mempty] map_lam' arrs [] [] - parallelMin [w] - - + let red_lam' = soacsLambdaToKernels red_lam + map_lam' = soacsLambdaToKernels map_lam + certifying (stmAuxCerts aux) $ + addStms =<< segRed lvl' pat w [SegBinOp comm red_lam' nes mempty] map_lam' arrs [] [] + parallelMin [w] Op (Hist w ops bucket_fun arrs) -> do ops' <- forM ops $ \(HistOp num_bins rf dests nes op) -> do (op', nes', shape) <- determineReduceOp op nes @@ -234,18 +260,16 @@ certifying (stmAuxCerts aux) $ addStms =<< segHist lvl' pat w [] [] ops' bucket_fun' arrs parallelMin [w] - Op (Stream w (Sequential accs) lam arrs) | chunk_size_param : _ <- lambdaParams lam -> do - types <- asksScope castScope - ((), stream_bnds) <- - runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types - let replace (Var v) | v == paramName chunk_size_param = w - replace se = se - replaceSets (Acc x y log) = - Acc (S.map (map replace) x) (S.map (map replace) y) log - censor replaceSets $ intraGroupStms lvl stream_bnds - + types <- asksScope castScope + ((), stream_bnds) <- + runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types + let replace (Var v) | v == paramName chunk_size_param = w + replace se = se + replaceSets (Acc x y log) = + Acc (S.map (map replace) x) (S.map (map replace) y) log + censor replaceSets $ intraGroupStms lvl stream_bnds Op (Scatter w lam ivs dests) -> do write_i <- newVName "write_i" space <- mkSegSpace [(write_i, w)] @@ -253,14 +277,17 @@ let lam' = soacsLambdaToKernels lam (dests_ws, dests_ns, dests_vs) = unzip3 dests (i_res, v_res) = splitAt (sum dests_ns) $ 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 ] - inputs = do (p, p_a) <- zip (lambdaParams lam') ivs - return $ KernelInput (paramName p) (paramType p) p_a [Var write_i] + 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] + inputs = do + (p, p_a) <- zip (lambdaParams lam') ivs + return $ KernelInput (paramName p) (paramType p) p_a [Var write_i] - kstms <- runBinder_ $ localScope (scopeOfSegSpace space) $ do - mapM_ readKernelInput inputs - addStms $ bodyStms $ lambdaBody lam' + kstms <- runBinder_ $ + localScope (scopeOfSegSpace space) $ do + mapM_ readKernelInput inputs + addStms $ bodyStms $ lambdaBody lam' certifying (stmAuxCerts aux) $ do let ts = map rowType $ patternTypes pat @@ -268,18 +295,23 @@ letBind pat $ Op $ SegOp $ SegMap lvl' space ts body parallelMin [w] - _ -> addStm $ soacsStmToKernels stm intraGroupStms :: SegLevel -> Stms SOACS -> IntraGroupM () intraGroupStms lvl = mapM_ (intraGroupStm lvl) -intraGroupParalleliseBody :: (MonadFreshNames m, HasScope Out.Kernels m) => - SegLevel -> Body - -> m ([[SubExp]], [[SubExp]], Log, Out.KernelBody Out.Kernels) +intraGroupParalleliseBody :: + (MonadFreshNames m, HasScope Out.Kernels m) => + SegLevel -> + Body -> + m ([[SubExp]], [[SubExp]], Log, Out.KernelBody Out.Kernels) intraGroupParalleliseBody lvl body = do (Acc min_ws avail_ws log, kstms) <- runIntraGroupM $ intraGroupStms lvl $ bodyStms body - return (S.toList min_ws, S.toList avail_ws, log, - KernelBody () kstms $ map (Returns ResultMaySimplify) $ bodyResult body) + return + ( S.toList min_ws, + S.toList avail_ws, + log, + KernelBody () kstms $ map (Returns ResultMaySimplify) $ bodyResult body + ) diff --git a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs b/src/Futhark/Pass/ExtractKernels/StreamKernel.hs --- a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs +++ b/src/Futhark/Pass/ExtractKernels/StreamKernel.hs @@ -1,105 +1,147 @@ +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE ConstraintKinds #-} + module Futhark.Pass.ExtractKernels.StreamKernel - ( segThreadCapped - , streamRed - , streamMap + ( segThreadCapped, + streamRed, + streamMap, ) - where +where import Control.Monad import Control.Monad.Writer import Data.List () - -import Prelude hiding (quot) - import Futhark.Analysis.PrimExp import Futhark.IR -import Futhark.IR.Kernels - hiding (Prog, Body, Stm, Pattern, PatElem, - BasicOp, Exp, Lambda, FunDef, FParam, LParam, RetType) +import Futhark.IR.Kernels hiding + ( BasicOp, + Body, + Exp, + FParam, + FunDef, + LParam, + Lambda, + PatElem, + Pattern, + Prog, + RetType, + Stm, + ) +import Futhark.MonadFreshNames import Futhark.Pass.ExtractKernels.BlockedKernel import Futhark.Pass.ExtractKernels.ToKernels -import Futhark.MonadFreshNames import Futhark.Tools +import Prelude hiding (quot) -data KernelSize = KernelSize { kernelElementsPerThread :: SubExp - -- ^ Int64 - , kernelNumThreads :: SubExp - -- ^ Int32 - } - deriving (Eq, Ord, Show) +data KernelSize = KernelSize + { -- | Int64 + kernelElementsPerThread :: SubExp, + -- | Int32 + kernelNumThreads :: SubExp + } + deriving (Eq, Ord, Show) -numberOfGroups :: (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => - String -> SubExp -> SubExp -> m (SubExp, SubExp) -numberOfGroups desc w64 group_size = do +numberOfGroups :: + (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => + String -> + SubExp -> + SubExp -> + m (SubExp, SubExp) +numberOfGroups desc w group_size = do max_num_groups_key <- nameFromString . pretty <$> newVName (desc ++ "_num_groups") - num_groups <- letSubExp "num_groups" $ - Op $ SizeOp $ CalcNumGroups w64 max_num_groups_key group_size - num_threads <- letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int32 OverflowUndef) num_groups group_size + num_groups <- + letSubExp "num_groups" $ + Op $ SizeOp $ CalcNumGroups w max_num_groups_key group_size + num_threads <- + letSubExp "num_threads" $ + BasicOp $ BinOp (Mul Int64 OverflowUndef) num_groups group_size return (num_groups, num_threads) -blockedKernelSize :: (MonadBinder m, Lore m ~ Kernels) => - String -> SubExp -> m KernelSize +blockedKernelSize :: + (MonadBinder m, Lore m ~ Kernels) => + String -> + SubExp -> + m KernelSize blockedKernelSize desc w = do group_size <- getSize (desc ++ "_group_size") SizeGroup - w64 <- letSubExp "w64" $ BasicOp $ ConvOp (SExt Int32 Int64) w - (_, num_threads) <- numberOfGroups desc w64 group_size + (_, num_threads) <- numberOfGroups desc w group_size per_thread_elements <- - letSubExp "per_thread_elements" =<< - eBinOp (SDivUp Int64 Unsafe) (eSubExp w64) (toExp =<< asIntS Int64 num_threads) + letSubExp "per_thread_elements" + =<< eBinOp (SDivUp Int64 Unsafe) (eSubExp w) (eSubExp num_threads) return $ KernelSize per_thread_elements num_threads -splitArrays :: (MonadBinder m, Lore m ~ Kernels) => - VName -> [VName] - -> SplitOrdering -> SubExp -> SubExp -> SubExp -> [VName] - -> m () +splitArrays :: + (MonadBinder m, Lore m ~ Kernels) => + VName -> + [VName] -> + SplitOrdering -> + SubExp -> + SubExp -> + SubExp -> + [VName] -> + m () splitArrays chunk_size split_bound ordering w i elems_per_i arrs = do letBindNames [chunk_size] $ Op $ SizeOp $ SplitSpace ordering w i elems_per_i case ordering of - SplitContiguous -> do - offset <- letSubExp "slice_offset" $ BasicOp $ BinOp (Mul Int32 OverflowUndef) i elems_per_i + SplitContiguous -> do + offset <- letSubExp "slice_offset" $ BasicOp $ BinOp (Mul Int64 OverflowUndef) i elems_per_i zipWithM_ (contiguousSlice offset) split_bound arrs SplitStrided stride -> zipWithM_ (stridedSlice stride) split_bound arrs - where contiguousSlice offset slice_name arr = do - arr_t <- lookupType arr - let slice = fullSlice arr_t [DimSlice offset (Var chunk_size) (constant (1::Int32))] - letBindNames [slice_name] $ BasicOp $ Index arr slice + where + contiguousSlice offset slice_name arr = do + arr_t <- lookupType arr + let slice = fullSlice arr_t [DimSlice offset (Var chunk_size) (constant (1 :: Int64))] + letBindNames [slice_name] $ BasicOp $ Index arr slice - stridedSlice stride slice_name arr = do - arr_t <- lookupType arr - let slice = fullSlice arr_t [DimSlice i (Var chunk_size) stride] - letBindNames [slice_name] $ BasicOp $ Index arr slice + stridedSlice stride slice_name arr = do + arr_t <- lookupType arr + let slice = fullSlice arr_t [DimSlice i (Var chunk_size) stride] + letBindNames [slice_name] $ BasicOp $ Index arr slice -partitionChunkedKernelFoldParameters :: Int -> [Param dec] - -> (VName, Param dec, [Param dec], [Param dec]) +partitionChunkedKernelFoldParameters :: + Int -> + [Param dec] -> + (VName, Param dec, [Param dec], [Param dec]) partitionChunkedKernelFoldParameters num_accs (i_param : chunk_param : params) = let (acc_params, arr_params) = splitAt num_accs params - in (paramName i_param, chunk_param, acc_params, arr_params) + in (paramName i_param, chunk_param, acc_params, arr_params) partitionChunkedKernelFoldParameters _ _ = error "partitionChunkedKernelFoldParameters: lambda takes too few parameters" -blockedPerThread :: (MonadBinder m, Lore m ~ Kernels) => - VName -> SubExp -> KernelSize -> StreamOrd -> Lambda (Lore m) - -> Int -> [VName] - -> m ([PatElemT Type], [PatElemT Type]) +blockedPerThread :: + (MonadBinder m, Lore m ~ Kernels) => + VName -> + SubExp -> + KernelSize -> + StreamOrd -> + Lambda (Lore m) -> + Int -> + [VName] -> + m ([PatElemT Type], [PatElemT Type]) blockedPerThread thread_gtid w kernel_size ordering lam num_nonconcat arrs = do let (_, chunk_size, [], arr_params) = partitionChunkedKernelFoldParameters 0 $ lambdaParams lam ordering' = - case ordering of InOrder -> SplitContiguous - Disorder -> SplitStrided $ kernelNumThreads kernel_size + case ordering of + InOrder -> SplitContiguous + Disorder -> SplitStrided $ kernelNumThreads kernel_size red_ts = take num_nonconcat $ lambdaReturnType lam map_ts = map rowType $ drop num_nonconcat $ lambdaReturnType lam - per_thread <- asIntS Int32 $ kernelElementsPerThread kernel_size - splitArrays (paramName chunk_size) (map paramName arr_params) ordering' w - (Var thread_gtid) per_thread arrs + per_thread <- asIntS Int64 $ kernelElementsPerThread kernel_size + splitArrays + (paramName chunk_size) + (map paramName arr_params) + ordering' + w + (Var thread_gtid) + per_thread + arrs chunk_red_pes <- forM red_ts $ \red_t -> do pe_name <- newVName "chunk_fold_red" @@ -112,21 +154,26 @@ splitAt num_nonconcat $ bodyResult $ lambdaBody lam addStms $ - bodyStms (lambdaBody lam) <> - stmsFromList - [ Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se - | (pe,se) <- zip chunk_red_pes chunk_red_ses ] <> - stmsFromList - [ Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se - | (pe,se) <- zip chunk_map_pes chunk_map_ses ] + bodyStms (lambdaBody lam) + <> stmsFromList + [ Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se + | (pe, se) <- zip chunk_red_pes chunk_red_ses + ] + <> stmsFromList + [ Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se + | (pe, se) <- zip chunk_map_pes chunk_map_ses + ] return (chunk_red_pes, chunk_map_pes) -- | Given a chunked fold lambda that takes its initial accumulator -- value as parameters, bind those parameters to the neutral element -- instead. -kerneliseLambda :: MonadFreshNames m => - [SubExp] -> Lambda Kernels -> m (Lambda Kernels) +kerneliseLambda :: + MonadFreshNames m => + [SubExp] -> + Lambda Kernels -> + m (Lambda Kernels) kerneliseLambda nes lam = do thread_index <- newVName "thread_index" let thread_index_param = Param thread_index $ Prim int32 @@ -135,59 +182,69 @@ mkAccInit p (Var v) | not $ primType $ paramType p = - mkLet [] [paramIdent p] $ BasicOp $ Copy v + mkLet [] [paramIdent p] $ BasicOp $ Copy v mkAccInit p x = mkLet [] [paramIdent p] $ BasicOp $ SubExp x acc_init_bnds = stmsFromList $ zipWith mkAccInit fold_acc_params nes - return lam { lambdaBody = insertStms acc_init_bnds $ - lambdaBody lam - , lambdaParams = thread_index_param : - fold_chunk_param : - fold_inp_params - } + return + lam + { lambdaBody = + insertStms acc_init_bnds $ + lambdaBody lam, + lambdaParams = + thread_index_param : + fold_chunk_param : + fold_inp_params + } -prepareStream :: (MonadBinder m, Lore m ~ Kernels) => - KernelSize - -> [(VName, SubExp)] - -> SubExp - -> Commutativity - -> Lambda Kernels - -> [SubExp] - -> [VName] - -> m (SubExp, SegSpace, [Type], KernelBody Kernels) +prepareStream :: + (MonadBinder m, Lore m ~ Kernels) => + KernelSize -> + [(VName, SubExp)] -> + SubExp -> + Commutativity -> + Lambda Kernels -> + [SubExp] -> + [VName] -> + m (SubExp, SegSpace, [Type], KernelBody Kernels) prepareStream size ispace w comm fold_lam nes arrs = do let (KernelSize elems_per_thread num_threads) = size let (ordering, split_ordering) = - case comm of Commutative -> (Disorder, SplitStrided num_threads) - Noncommutative -> (InOrder, SplitContiguous) + case comm of + Commutative -> (Disorder, SplitStrided num_threads) + Noncommutative -> (InOrder, SplitContiguous) fold_lam' <- kerneliseLambda nes fold_lam - elems_per_thread_32 <- asIntS Int32 elems_per_thread - gtid <- newVName "gtid" space <- mkSegSpace $ ispace ++ [(gtid, num_threads)] - kbody <- fmap (uncurry (flip (KernelBody ()))) $ runBinder $ - localScope (scopeOfSegSpace space) $ do - (chunk_red_pes, chunk_map_pes) <- - blockedPerThread gtid w size ordering fold_lam' (length nes) arrs - let concatReturns pe = - ConcatReturns split_ordering w elems_per_thread_32 $ patElemName pe - return (map (Returns ResultMaySimplify . Var . patElemName) chunk_red_pes ++ - map concatReturns chunk_map_pes) + kbody <- fmap (uncurry (flip (KernelBody ()))) $ + runBinder $ + localScope (scopeOfSegSpace space) $ do + (chunk_red_pes, chunk_map_pes) <- + blockedPerThread gtid w size ordering fold_lam' (length nes) arrs + let concatReturns pe = + ConcatReturns split_ordering w elems_per_thread $ patElemName pe + return + ( map (Returns ResultMaySimplify . Var . patElemName) chunk_red_pes + ++ map concatReturns chunk_map_pes + ) let (redout_ts, mapout_ts) = splitAt (length nes) $ lambdaReturnType fold_lam ts = redout_ts ++ map rowType mapout_ts return (num_threads, space, ts, kbody) -streamRed :: (MonadFreshNames m, HasScope Kernels m) => - MkSegLevel Kernels m - -> Pattern Kernels - -> SubExp - -> Commutativity - -> Lambda Kernels -> Lambda Kernels - -> [SubExp] -> [VName] - -> m (Stms Kernels) +streamRed :: + (MonadFreshNames m, HasScope Kernels m) => + MkSegLevel Kernels m -> + Pattern Kernels -> + SubExp -> + Commutativity -> + Lambda Kernels -> + Lambda Kernels -> + [SubExp] -> + [VName] -> + m (Stms Kernels) streamRed mk_lvl pat w comm red_lam fold_lam nes arrs = runBinderT'_ $ do -- The strategy here is to rephrase the stream reduction as a -- non-segmented SegRed that does explicit chunking within its body. @@ -201,17 +258,30 @@ (_, kspace, ts, kbody) <- prepareStream size ispace w comm fold_lam nes arrs lvl <- mk_lvl [w] "stream_red" $ NoRecommendation SegNoVirt - letBind pat' $ Op $ SegOp $ SegRed lvl kspace - [SegBinOp comm red_lam nes mempty] ts kbody + letBind pat' $ + Op $ + SegOp $ + SegRed + lvl + kspace + [SegBinOp comm red_lam nes mempty] + ts + kbody read_dummy -- Similar to streamRed, but without the last reduction. -streamMap :: (MonadFreshNames m, HasScope Kernels m) => - MkSegLevel Kernels m - -> [String] -> [PatElem Kernels] -> SubExp - -> Commutativity -> Lambda Kernels -> [SubExp] -> [VName] - -> m ((SubExp, [VName]), Stms Kernels) +streamMap :: + (MonadFreshNames m, HasScope Kernels m) => + MkSegLevel Kernels m -> + [String] -> + [PatElem Kernels] -> + SubExp -> + Commutativity -> + Lambda Kernels -> + [SubExp] -> + [VName] -> + m ((SubExp, [VName]), Stms Kernels) streamMap mk_lvl out_desc mapout_pes w comm fold_lam nes arrs = runBinderT' $ do size <- blockedKernelSize "stream_map" w @@ -233,19 +303,20 @@ -- array. segThreadCapped :: MonadFreshNames m => MkSegLevel Kernels m segThreadCapped ws desc r = do - w64 <- letSubExp "nest_size" =<< - foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) =<< - mapM (asIntS Int64) ws + w <- + letSubExp "nest_size" + =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) ws group_size <- getSize (desc ++ "_group_size") SizeGroup case r of ManyThreads -> do - usable_groups <- letSubExp "segmap_usable_groups" . - BasicOp . ConvOp (SExt Int64 Int32) =<< - letSubExp "segmap_usable_groups_64" =<< - eBinOp (SDivUp Int64 Unsafe) (eSubExp w64) - (eSubExp =<< asIntS Int64 group_size) + usable_groups <- + letSubExp "segmap_usable_groups" + =<< eBinOp + (SDivUp Int64 Unsafe) + (eSubExp w) + (eSubExp =<< asIntS Int64 group_size) return $ SegThread (Count usable_groups) (Count group_size) SegNoVirt NoRecommendation v -> do - (num_groups, _) <- numberOfGroups desc w64 group_size + (num_groups, _) <- numberOfGroups desc w group_size return $ SegThread (Count num_groups) (Count group_size) v diff --git a/src/Futhark/Pass/ExtractKernels/ToKernels.hs b/src/Futhark/Pass/ExtractKernels/ToKernels.hs --- a/src/Futhark/Pass/ExtractKernels/ToKernels.hs +++ b/src/Futhark/Pass/ExtractKernels/ToKernels.hs @@ -1,63 +1,75 @@ +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE ConstraintKinds #-} -module Futhark.Pass.ExtractKernels.ToKernels - ( getSize - , segThread - , soacsLambdaToKernels - , soacsStmToKernels - , scopeForKernels - , scopeForSOACs - ) - where +module Futhark.Pass.ExtractKernels.ToKernels + ( getSize, + segThread, + soacsLambdaToKernels, + soacsStmToKernels, + scopeForKernels, + scopeForSOACs, + ) +where import Control.Monad.Identity import Data.List () - import Futhark.Analysis.Rephrase import Futhark.IR +import Futhark.IR.Kernels import Futhark.IR.SOACS (SOACS) import qualified Futhark.IR.SOACS.SOAC as SOAC -import Futhark.IR.Kernels import Futhark.Tools -getSize :: (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => - String -> SizeClass -> m SubExp +getSize :: + (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => + String -> + SizeClass -> + m SubExp getSize desc size_class = do size_key <- nameFromString . pretty <$> newVName desc letSubExp desc $ Op $ SizeOp $ GetSize size_key size_class -segThread :: (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => - String -> m SegLevel +segThread :: + (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => + String -> + m SegLevel segThread desc = SegThread <$> (Count <$> getSize (desc ++ "_num_groups") SizeNumGroups) <*> (Count <$> getSize (desc ++ "_group_size") SizeGroup) <*> pure SegVirt -injectSOACS :: (Monad m, - SameScope from to, - ExpDec from ~ ExpDec to, - BodyDec from ~ BodyDec to, - RetType from ~ RetType to, - BranchType from ~ BranchType to, - Op from ~ SOAC from) => - (SOAC to -> Op to) -> Rephraser m from to -injectSOACS f = Rephraser { rephraseExpLore = return - , rephraseBodyLore = return - , rephraseLetBoundLore = return - , rephraseFParamLore = return - , rephraseLParamLore = return - , rephraseOp = fmap f . onSOAC - , rephraseRetType = return - , rephraseBranchType = return - } - where onSOAC = SOAC.mapSOACM mapper - mapper = SOAC.SOACMapper { SOAC.mapOnSOACSubExp = return - , SOAC.mapOnSOACVName = return - , SOAC.mapOnSOACLambda = rephraseLambda $ injectSOACS f - } +injectSOACS :: + ( Monad m, + SameScope from to, + ExpDec from ~ ExpDec to, + BodyDec from ~ BodyDec to, + RetType from ~ RetType to, + BranchType from ~ BranchType to, + Op from ~ SOAC from + ) => + (SOAC to -> Op to) -> + Rephraser m from to +injectSOACS f = + Rephraser + { rephraseExpLore = return, + rephraseBodyLore = return, + rephraseLetBoundLore = return, + rephraseFParamLore = return, + rephraseLParamLore = return, + rephraseOp = fmap f . onSOAC, + rephraseRetType = return, + rephraseBranchType = return + } + where + onSOAC = SOAC.mapSOACM mapper + mapper = + SOAC.SOACMapper + { SOAC.mapOnSOACSubExp = return, + SOAC.mapOnSOACVName = return, + SOAC.mapOnSOACLambda = rephraseLambda $ injectSOACS f + } soacsStmToKernels :: Stm SOACS -> Stm Kernels soacsStmToKernels = runIdentity . rephraseStm (injectSOACS OtherOp) diff --git a/src/Futhark/Pass/FirstOrderTransform.hs b/src/Futhark/Pass/FirstOrderTransform.hs --- a/src/Futhark/Pass/FirstOrderTransform.hs +++ b/src/Futhark/Pass/FirstOrderTransform.hs @@ -1,5 +1,6 @@ -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + -- | Transform any SOACs to @for@-loops. -- -- Example: @@ -19,19 +20,18 @@ -- let ys'[i] = y -- in ys' -- @ -module Futhark.Pass.FirstOrderTransform - ( firstOrderTransform ) - where +module Futhark.Pass.FirstOrderTransform (firstOrderTransform) where -import Futhark.Transform.FirstOrderTransform (FirstOrderLore, transformFunDef, transformConsts) import Futhark.IR.SOACS (SOACS, scopeOf) import Futhark.Pass +import Futhark.Transform.FirstOrderTransform (FirstOrderLore, transformConsts, transformFunDef) -- | The first-order transformation pass. firstOrderTransform :: FirstOrderLore lore => Pass SOACS lore firstOrderTransform = Pass - "first order transform" - "Transform all SOACs to for-loops." $ - intraproceduralTransformationWithConsts - transformConsts (transformFunDef . scopeOf) + "first order transform" + "Transform all SOACs to for-loops." + $ intraproceduralTransformationWithConsts + transformConsts + (transformFunDef . scopeOf) diff --git a/src/Futhark/Pass/KernelBabysitting.hs b/src/Futhark/Pass/KernelBabysitting.hs --- a/src/Futhark/Pass/KernelBabysitting.hs +++ b/src/Futhark/Pass/KernelBabysitting.hs @@ -1,33 +1,46 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + -- | Do various kernel optimisations - mostly related to coalescing. -module Futhark.Pass.KernelBabysitting ( babysitKernels ) - where +module Futhark.Pass.KernelBabysitting (babysitKernels) where import Control.Arrow (first) import Control.Monad.State.Strict -import qualified Data.Map.Strict as M import Data.Foldable import Data.List (elemIndex, isPrefixOf, sort) +import qualified Data.Map.Strict as M import Data.Maybe - -import Futhark.MonadFreshNames import Futhark.IR -import Futhark.IR.Kernels - hiding (Prog, Body, Stm, Pattern, PatElem, - BasicOp, Exp, Lambda, FunDef, FParam, LParam, RetType) -import Futhark.Tools +import Futhark.IR.Kernels hiding + ( BasicOp, + Body, + Exp, + FParam, + FunDef, + LParam, + Lambda, + PatElem, + Pattern, + Prog, + RetType, + Stm, + ) +import Futhark.MonadFreshNames import Futhark.Pass +import Futhark.Tools import Futhark.Util -- | The pass definition. babysitKernels :: Pass Kernels Kernels -babysitKernels = Pass "babysit kernels" - "Transpose kernel input arrays for better performance." $ - intraproceduralTransformation onStms - where onStms scope stms = do - let m = localScope scope $ transformStms mempty stms - fmap fst $ modifyNameSource $ runState (runBinderT m M.empty) +babysitKernels = + Pass + "babysit kernels" + "Transpose kernel input arrays for better performance." + $ intraproceduralTransformation onStms + where + onStms scope stms = do + let m = localScope scope $ transformStms mempty stms + fmap fst $ modifyNameSource $ runState (runBinderT m M.empty) type BabysitM = Binder Kernels @@ -54,37 +67,46 @@ Just (Let _ _ (BasicOp (Reshape _ arr))) -> nonlinearInMemory arr m Just (Let _ _ (BasicOp (Manifest perm _))) -> Just $ Just perm Just (Let pat _ (Op (SegOp (SegMap _ _ ts _)))) -> - nonlinear =<< find ((==name) . patElemName . fst) - (zip (patternElements pat) ts) + nonlinear + =<< find + ((== name) . patElemName . fst) + (zip (patternElements pat) ts) _ -> Nothing - where nonlinear (pe, t) - | inner_r <- arrayRank t, inner_r > 0 = do - let outer_r = arrayRank (patElemType pe) - inner_r - return $ Just $ rearrangeInverse $ [inner_r..inner_r+outer_r-1] ++ [0..inner_r-1] - | otherwise = Nothing + where + nonlinear (pe, t) + | inner_r <- arrayRank t, + inner_r > 0 = do + let outer_r = arrayRank (patElemType pe) - inner_r + return $ Just $ rearrangeInverse $ [inner_r .. inner_r + outer_r -1] ++ [0 .. inner_r -1] + | otherwise = Nothing transformStm :: ExpMap -> Stm Kernels -> BabysitM ExpMap transformStm expmap (Let pat aux (Op (SegOp op))) = do - let mapper = identitySegOpMapper - { mapOnSegOpBody = - transformKernelBody expmap (segLevel op) (segSpace op) - } + let mapper = + identitySegOpMapper + { mapOnSegOpBody = + transformKernelBody expmap (segLevel op) (segSpace op) + } op' <- mapSegOpM mapper op let stm' = Let pat aux $ Op $ SegOp op' addStm stm' - return $ M.fromList [ (name, stm') | name <- patternNames pat ] <> expmap + return $ M.fromList [(name, stm') | name <- patternNames pat] <> expmap transformStm expmap (Let pat aux e) = do e' <- mapExpM (transform expmap) e let bnd' = Let pat aux e' addStm bnd' - return $ M.fromList [ (name, bnd') | name <- patternNames pat ] <> expmap + return $ M.fromList [(name, bnd') | name <- patternNames pat] <> expmap transform :: ExpMap -> Mapper Kernels Kernels BabysitM transform expmap = - identityMapper { mapOnBody = \scope -> localScope scope . transformBody expmap } + identityMapper {mapOnBody = \scope -> localScope scope . transformBody expmap} -transformKernelBody :: ExpMap -> SegLevel -> SegSpace -> KernelBody Kernels - -> BabysitM (KernelBody Kernels) +transformKernelBody :: + ExpMap -> + SegLevel -> + SegSpace -> + KernelBody Kernels -> + BabysitM (KernelBody Kernels) transformKernelBody expmap lvl space kbody = do -- Go spelunking for accesses to arrays that are defined outside the -- kernel body and where the indices are kernel thread indices. @@ -92,212 +114,245 @@ let thread_gids = map fst $ unSegSpace space thread_local = namesFromList $ segFlat space : thread_gids free_ker_vars = freeIn kbody `namesSubtract` getKerVariantIds space - num_threads <- letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int32 OverflowUndef) - (unCount $ segNumGroups lvl) (unCount $ segGroupSize lvl) - evalStateT (traverseKernelBodyArrayIndexes - free_ker_vars - thread_local - (scope <> scopeOfSegSpace space) - (ensureCoalescedAccess expmap (unSegSpace space) num_threads) - kbody) + num_threads <- + letSubExp "num_threads" $ + BasicOp $ + BinOp + (Mul Int64 OverflowUndef) + (unCount $ segNumGroups lvl) + (unCount $ segGroupSize lvl) + evalStateT + ( traverseKernelBodyArrayIndexes + free_ker_vars + thread_local + (scope <> scopeOfSegSpace space) + (ensureCoalescedAccess expmap (unSegSpace space) num_threads) + kbody + ) mempty - where getKerVariantIds = namesFromList . M.keys . scopeOfSegSpace + where + getKerVariantIds = namesFromList . M.keys . scopeOfSegSpace type ArrayIndexTransform m = Names -> - (VName -> Bool) -> -- thread local? - (VName -> SubExp -> Bool)-> -- variant to a certain gid (given as first param)? - (SubExp -> Maybe SubExp) -> -- split substitution? - Scope Kernels -> -- type environment - VName -> Slice SubExp -> m (Maybe (VName, Slice SubExp)) + (VName -> Bool) -> -- thread local? + (VName -> SubExp -> Bool) -> -- variant to a certain gid (given as first param)? + (SubExp -> Maybe SubExp) -> -- split substitution? + Scope Kernels -> -- type environment + VName -> + Slice SubExp -> + m (Maybe (VName, Slice SubExp)) -traverseKernelBodyArrayIndexes :: (Applicative f, Monad f) => - Names - -> Names - -> Scope Kernels - -> ArrayIndexTransform f - -> KernelBody Kernels - -> f (KernelBody Kernels) +traverseKernelBodyArrayIndexes :: + (Applicative f, Monad f) => + Names -> + Names -> + Scope Kernels -> + ArrayIndexTransform f -> + KernelBody Kernels -> + f (KernelBody Kernels) traverseKernelBodyArrayIndexes free_ker_vars thread_variant outer_scope f (KernelBody () kstms kres) = - KernelBody () . stmsFromList <$> - mapM (onStm (varianceInStms mempty kstms, - mkSizeSubsts kstms, - outer_scope)) (stmsToList kstms) <*> - pure kres - where onLambda (variance, szsubst, scope) lam = - (\body' -> lam { lambdaBody = body' }) <$> - onBody (variance, szsubst, scope') (lambdaBody lam) - where scope' = scope <> scopeOfLParams (lambdaParams lam) - - onBody (variance, szsubst, scope) (Body bdec stms bres) = do - stms' <- stmsFromList <$> mapM (onStm (variance', szsubst', scope')) (stmsToList stms) - pure $ Body bdec stms' bres - where variance' = varianceInStms variance stms - szsubst' = mkSizeSubsts stms <> szsubst - scope' = scope <> scopeOf stms + KernelBody () . stmsFromList + <$> mapM + ( onStm + ( varianceInStms mempty kstms, + mkSizeSubsts kstms, + outer_scope + ) + ) + (stmsToList kstms) + <*> pure kres + where + onLambda (variance, szsubst, scope) lam = + (\body' -> lam {lambdaBody = body'}) + <$> onBody (variance, szsubst, scope') (lambdaBody lam) + where + scope' = scope <> scopeOfLParams (lambdaParams lam) - onStm (variance, szsubst, _) (Let pat dec (BasicOp (Index arr is))) = - Let pat dec . oldOrNew <$> f free_ker_vars isThreadLocal isGidVariant sizeSubst outer_scope arr is - where oldOrNew Nothing = - BasicOp $ Index arr is - oldOrNew (Just (arr', is')) = - BasicOp $ Index arr' is' + onBody (variance, szsubst, scope) (Body bdec stms bres) = do + stms' <- stmsFromList <$> mapM (onStm (variance', szsubst', scope')) (stmsToList stms) + pure $ Body bdec stms' bres + where + variance' = varianceInStms variance stms + szsubst' = mkSizeSubsts stms <> szsubst + scope' = scope <> scopeOf stms - isGidVariant gid (Var v) = - gid == v || nameIn gid (M.findWithDefault (oneName v) v variance) - isGidVariant _ _ = False + onStm (variance, szsubst, _) (Let pat dec (BasicOp (Index arr is))) = + Let pat dec . oldOrNew <$> f free_ker_vars isThreadLocal isGidVariant sizeSubst outer_scope arr is + where + oldOrNew Nothing = + BasicOp $ Index arr is + oldOrNew (Just (arr', is')) = + BasicOp $ Index arr' is' - isThreadLocal v = - thread_variant `namesIntersect` - M.findWithDefault (oneName v) v variance + isGidVariant gid (Var v) = + gid == v || nameIn gid (M.findWithDefault (oneName v) v variance) + isGidVariant _ _ = False - sizeSubst (Constant v) = Just $ Constant v - sizeSubst (Var v) - | v `M.member` outer_scope = Just $ Var v - | Just v' <- M.lookup v szsubst = sizeSubst v' - | otherwise = Nothing + isThreadLocal v = + thread_variant + `namesIntersect` M.findWithDefault (oneName v) v variance - onStm (variance, szsubst, scope) (Let pat dec e) = - Let pat dec <$> mapExpM (mapper (variance, szsubst, scope)) e + sizeSubst (Constant v) = Just $ Constant v + sizeSubst (Var v) + | v `M.member` outer_scope = Just $ Var v + | Just v' <- M.lookup v szsubst = sizeSubst v' + | otherwise = Nothing + onStm (variance, szsubst, scope) (Let pat dec e) = + Let pat dec <$> mapExpM (mapper (variance, szsubst, scope)) e - onOp ctx (OtherOp soac) = - OtherOp <$> mapSOACM identitySOACMapper{ mapOnSOACLambda = onLambda ctx } soac - onOp _ op = return op + onOp ctx (OtherOp soac) = + OtherOp <$> mapSOACM identitySOACMapper {mapOnSOACLambda = onLambda ctx} soac + onOp _ op = return op - mapper ctx = identityMapper { mapOnBody = const (onBody ctx) - , mapOnOp = onOp ctx } + mapper ctx = + identityMapper + { mapOnBody = const (onBody ctx), + mapOnOp = onOp ctx + } - mkSizeSubsts = foldMap mkStmSizeSubst - where mkStmSizeSubst (Let (Pattern [] [pe]) _ (Op (SizeOp (SplitSpace _ _ _ elems_per_i)))) = - M.singleton (patElemName pe) elems_per_i - mkStmSizeSubst _ = mempty + mkSizeSubsts = foldMap mkStmSizeSubst + where + mkStmSizeSubst (Let (Pattern [] [pe]) _ (Op (SizeOp (SplitSpace _ _ _ elems_per_i)))) = + M.singleton (patElemName pe) elems_per_i + mkStmSizeSubst _ = mempty type Replacements = M.Map (VName, Slice SubExp) VName -ensureCoalescedAccess :: MonadBinder m => - ExpMap - -> [(VName,SubExp)] - -> SubExp - -> ArrayIndexTransform (StateT Replacements m) -ensureCoalescedAccess expmap thread_space num_threads free_ker_vars isThreadLocal - isGidVariant sizeSubst outer_scope arr slice = do - seen <- gets $ M.lookup (arr, slice) - - case (seen, isThreadLocal arr, typeOf <$> M.lookup arr outer_scope) of - -- Already took care of this case elsewhere. - (Just arr', _, _) -> - pure $ Just (arr', slice) +ensureCoalescedAccess :: + MonadBinder m => + ExpMap -> + [(VName, SubExp)] -> + SubExp -> + ArrayIndexTransform (StateT Replacements m) +ensureCoalescedAccess + expmap + thread_space + num_threads + free_ker_vars + isThreadLocal + isGidVariant + sizeSubst + outer_scope + arr + slice = do + seen <- gets $ M.lookup (arr, slice) - (Nothing, False, Just t) - -- We are fully indexing the array with thread IDs, but the - -- indices are in a permuted order. - | Just is <- sliceIndices slice, - length is == arrayRank t, - Just is' <- coalescedIndexes free_ker_vars isGidVariant (map Var thread_gids) is, - Just perm <- is' `isPermutationOf` is -> + case (seen, isThreadLocal arr, typeOf <$> M.lookup arr outer_scope) of + -- Already took care of this case elsewhere. + (Just arr', _, _) -> + pure $ Just (arr', slice) + (Nothing, False, Just t) + -- We are fully indexing the array with thread IDs, but the + -- indices are in a permuted order. + | Just is <- sliceIndices slice, + length is == arrayRank t, + Just is' <- coalescedIndexes free_ker_vars isGidVariant (map Var thread_gids) is, + Just perm <- is' `isPermutationOf` is -> replace =<< lift (rearrangeInput (nonlinearInMemory arr expmap) perm arr) - - -- Check whether the access is already coalesced because of a - -- previous rearrange being applied to the current array: - -- 1. get the permutation of the source-array rearrange - -- 2. apply it to the slice - -- 3. check that the innermost index is actually the gid - -- of the innermost kernel dimension. - -- If so, the access is already coalesced, nothing to do! - -- (Cosmin's Heuristic.) - | Just (Let _ _ (BasicOp (Rearrange perm _))) <- M.lookup arr expmap, - not $ null perm, - not $ null thread_gids, - inner_gid <- last thread_gids, - length slice >= length perm, - slice' <- map (slice !!) perm, - DimFix inner_ind <- last slice', - not $ null thread_gids, - isGidVariant inner_gid inner_ind -> + -- Check whether the access is already coalesced because of a + -- previous rearrange being applied to the current array: + -- 1. get the permutation of the source-array rearrange + -- 2. apply it to the slice + -- 3. check that the innermost index is actually the gid + -- of the innermost kernel dimension. + -- If so, the access is already coalesced, nothing to do! + -- (Cosmin's Heuristic.) + | Just (Let _ _ (BasicOp (Rearrange perm _))) <- M.lookup arr expmap, + not $ null perm, + not $ null thread_gids, + inner_gid <- last thread_gids, + length slice >= length perm, + slice' <- map (slice !!) perm, + DimFix inner_ind <- last slice', + not $ null thread_gids, + isGidVariant inner_gid inner_ind -> return Nothing - - -- We are not fully indexing an array, but the remaining slice - -- is invariant to the innermost-kernel dimension. We assume - -- the remaining slice will be sequentially streamed, hence - -- tiling will be applied later and will solve coalescing. - -- Hence nothing to do at this point. (Cosmin's Heuristic.) - | (is, rem_slice) <- splitSlice slice, - not $ null rem_slice, - allDimAreSlice rem_slice, - Nothing <- M.lookup arr expmap, - not $ tooSmallSlice (primByteSize (elemType t)) rem_slice, - is /= map Var (take (length is) thread_gids) || length is == length thread_gids, - not (null thread_gids || null is), - not (last thread_gids `nameIn` (freeIn is <> freeIn rem_slice)) -> + -- We are not fully indexing an array, but the remaining slice + -- is invariant to the innermost-kernel dimension. We assume + -- the remaining slice will be sequentially streamed, hence + -- tiling will be applied later and will solve coalescing. + -- Hence nothing to do at this point. (Cosmin's Heuristic.) + | (is, rem_slice) <- splitSlice slice, + not $ null rem_slice, + allDimAreSlice rem_slice, + Nothing <- M.lookup arr expmap, + not $ tooSmallSlice (primByteSize (elemType t)) rem_slice, + is /= map Var (take (length is) thread_gids) || length is == length thread_gids, + not (null thread_gids || null is), + not (last thread_gids `nameIn` (freeIn is <> freeIn rem_slice)) -> return Nothing - - -- We are not fully indexing the array, and the indices are not - -- a proper prefix of the thread indices, and some indices are - -- thread local, so we assume (HEURISTIC!) that the remaining - -- dimensions will be traversed sequentially. - | (is, rem_slice) <- splitSlice slice, - not $ null rem_slice, - not $ tooSmallSlice (primByteSize (elemType t)) rem_slice, - is /= map Var (take (length is) thread_gids) || length is == length thread_gids, - any isThreadLocal (namesToList $ freeIn is) -> do + -- We are not fully indexing the array, and the indices are not + -- a proper prefix of the thread indices, and some indices are + -- thread local, so we assume (HEURISTIC!) that the remaining + -- dimensions will be traversed sequentially. + | (is, rem_slice) <- splitSlice slice, + not $ null rem_slice, + not $ tooSmallSlice (primByteSize (elemType t)) rem_slice, + is /= map Var (take (length is) thread_gids) || length is == length thread_gids, + any isThreadLocal (namesToList $ freeIn is) -> do let perm = coalescingPermutation (length is) $ arrayRank t replace =<< lift (rearrangeInput (nonlinearInMemory arr expmap) perm arr) - -- We are taking a slice of the array with a unit stride. We - -- assume that the slice will be traversed sequentially. - -- - -- We will really want to treat the sliced dimension like two - -- dimensions so we can transpose them. This may require - -- padding. - | (is, rem_slice) <- splitSlice slice, - and $ zipWith (==) is $ map Var thread_gids, - DimSlice offset len (Constant stride):_ <- rem_slice, - isThreadLocalSubExp offset, - Just {} <- sizeSubst len, - oneIsh stride -> do - let num_chunks = if null is - then primExpFromSubExp int32 num_threads - else coerceIntPrimExp Int32 $ - product $ map (primExpFromSubExp int32) $ - drop (length is) thread_gdims + -- We are taking a slice of the array with a unit stride. We + -- assume that the slice will be traversed sequentially. + -- + -- We will really want to treat the sliced dimension like two + -- dimensions so we can transpose them. This may require + -- padding. + | (is, rem_slice) <- splitSlice slice, + and $ zipWith (==) is $ map Var thread_gids, + DimSlice offset len (Constant stride) : _ <- rem_slice, + isThreadLocalSubExp offset, + Just {} <- sizeSubst len, + oneIsh stride -> do + let num_chunks = + if null is + then untyped $ pe32 num_threads + else + untyped $ + product $ + map pe64 $ + drop (length is) thread_gdims replace =<< lift (rearrangeSlice (length is) (arraySize (length is) t) num_chunks arr) - -- Everything is fine... assuming that the array is in row-major - -- order! Make sure that is the case. - | Just{} <- nonlinearInMemory arr expmap -> + -- Everything is fine... assuming that the array is in row-major + -- order! Make sure that is the case. + | Just {} <- nonlinearInMemory arr expmap -> case sliceIndices slice of - Just is | Just _ <- coalescedIndexes free_ker_vars isGidVariant (map Var thread_gids) is -> - replace =<< lift (rowMajorArray arr) - | otherwise -> - return Nothing + Just is + | Just _ <- coalescedIndexes free_ker_vars isGidVariant (map Var thread_gids) is -> + replace =<< lift (rowMajorArray arr) + | otherwise -> + return Nothing _ -> replace =<< lift (rowMajorArray arr) - - _ -> return Nothing - - where (thread_gids, thread_gdims) = unzip thread_space + _ -> return Nothing + where + (thread_gids, thread_gdims) = unzip thread_space - replace arr' = do - modify $ M.insert (arr, slice) arr' - return $ Just (arr', slice) + replace arr' = do + modify $ M.insert (arr, slice) arr' + return $ Just (arr', slice) - isThreadLocalSubExp (Var v) = isThreadLocal v - isThreadLocalSubExp Constant{} = False + isThreadLocalSubExp (Var v) = isThreadLocal v + isThreadLocalSubExp Constant {} = False -- Heuristic for avoiding rearranging too small arrays. tooSmallSlice :: Int32 -> Slice SubExp -> Bool -tooSmallSlice bs = fst . foldl comb (True,bs) . sliceDims - where comb (True, x) (Constant (IntValue (Int32Value d))) = (d*x < 4, d*x) - comb (_, x) _ = (False, x) +tooSmallSlice bs = fst . foldl comb (True, bs) . sliceDims + where + comb (True, x) (Constant (IntValue (Int32Value d))) = (d * x < 4, d * x) + comb (_, x) _ = (False, x) splitSlice :: Slice SubExp -> ([SubExp], Slice SubExp) splitSlice [] = ([], []) -splitSlice (DimFix i:is) = first (i:) $ splitSlice is +splitSlice (DimFix i : is) = first (i :) $ splitSlice is splitSlice is = ([], is) allDimAreSlice :: Slice SubExp -> Bool allDimAreSlice [] = True -allDimAreSlice (DimFix _:_) = False -allDimAreSlice (_:is) = allDimAreSlice is +allDimAreSlice (DimFix _ : _) = False +allDimAreSlice (_ : is) = allDimAreSlice is -- Try to move thread indexes into their proper position. coalescedIndexes :: Names -> (VName -> SubExp -> Bool) -> [SubExp] -> [SubExp] -> Maybe [SubExp] @@ -310,59 +365,65 @@ -- 3. the indexes are a prefix of the thread indexes, because that -- means multiple threads will be accessing the same element. | any isCt is = - Nothing + Nothing | any (`nameIn` free_ker_vars) (subExpVars is) = - Nothing + Nothing | is `isPrefixOf` tgids = - Nothing + Nothing | not (null tgids), not (null is), Var innergid <- last tgids, num_is > 0 && isGidVariant innergid (last is) = - Just is + Just is -- 3. Otherwise try fix coalescing | otherwise = - Just $ reverse $ foldl move (reverse is) $ zip [0..] (reverse tgids) - where num_is = length is + Just $ reverse $ foldl move (reverse is) $ zip [0 ..] (reverse tgids) + where + num_is = length is - move is_rev (i, tgid) - -- If tgid is in is_rev anywhere but at position i, and - -- position i exists, we move it to position i instead. - | Just j <- elemIndex tgid is_rev, i /= j, i < num_is = - swap i j is_rev - | otherwise = - is_rev + move is_rev (i, tgid) + -- If tgid is in is_rev anywhere but at position i, and + -- position i exists, we move it to position i instead. + | Just j <- elemIndex tgid is_rev, + i /= j, + i < num_is = + swap i j is_rev + | otherwise = + is_rev - swap i j l - | Just ix <- maybeNth i l, - Just jx <- maybeNth j l = - update i jx $ update j ix l - | otherwise = - error $ "coalescedIndexes swap: invalid indices" ++ show (i, j, l) + swap i j l + | Just ix <- maybeNth i l, + Just jx <- maybeNth j l = + update i jx $ update j ix l + | otherwise = + error $ "coalescedIndexes swap: invalid indices" ++ show (i, j, l) - update 0 x (_:ys) = x : ys - update i x (y:ys) = y : update (i-1) x ys - update _ _ [] = error "coalescedIndexes: update" + update 0 x (_ : ys) = x : ys + update i x (y : ys) = y : update (i -1) x ys + update _ _ [] = error "coalescedIndexes: update" - isCt :: SubExp -> Bool - isCt (Constant _) = True - isCt (Var _) = False + isCt :: SubExp -> Bool + isCt (Constant _) = True + isCt (Var _) = False coalescingPermutation :: Int -> Int -> [Int] coalescingPermutation num_is rank = - [num_is..rank-1] ++ [0..num_is-1] + [num_is .. rank -1] ++ [0 .. num_is -1] -rearrangeInput :: MonadBinder m => - Maybe (Maybe [Int]) -> [Int] -> VName -> m VName +rearrangeInput :: + MonadBinder m => + Maybe (Maybe [Int]) -> + [Int] -> + VName -> + m VName rearrangeInput (Just (Just current_perm)) perm arr | current_perm == perm = return arr -- Already has desired representation. - rearrangeInput Nothing perm arr | sort perm == perm = return arr -- We don't know the current - -- representation, but the indexing - -- is linear, so let's hope the - -- array is too. -rearrangeInput (Just Just{}) perm arr + -- representation, but the indexing + -- is linear, so let's hope the + -- array is too. +rearrangeInput (Just Just {}) perm arr | sort perm == perm = rowMajorArray arr -- We just want a row-major array, no tricks. rearrangeInput manifest perm arr = do -- We may first manifest the array to ensure that it is flat in @@ -372,60 +433,73 @@ letExp (baseString arr ++ "_coalesced") $ BasicOp $ Manifest perm manifested -rowMajorArray :: MonadBinder m => - VName -> m VName +rowMajorArray :: + MonadBinder m => + VName -> + m VName rowMajorArray arr = do rank <- arrayRank <$> lookupType arr - letExp (baseString arr ++ "_rowmajor") $ BasicOp $ Manifest [0..rank-1] arr + letExp (baseString arr ++ "_rowmajor") $ BasicOp $ Manifest [0 .. rank -1] arr -rearrangeSlice :: MonadBinder m => - Int -> SubExp -> PrimExp VName -> VName - -> m VName +rearrangeSlice :: + MonadBinder m => + Int -> + SubExp -> + PrimExp VName -> + VName -> + m VName rearrangeSlice d w num_chunks arr = do num_chunks' <- toSubExp "num_chunks" num_chunks (w_padded, padding) <- paddedScanReduceInput w num_chunks' - per_chunk <- letSubExp "per_chunk" $ - BasicOp $ BinOp (SQuot Int32 Unsafe) w_padded num_chunks' + per_chunk <- + letSubExp "per_chunk" $ + BasicOp $ BinOp (SQuot Int64 Unsafe) w_padded num_chunks' arr_t <- lookupType arr arr_padded <- padArray w_padded padding arr_t rearrange num_chunks' w_padded per_chunk (baseString arr) arr_padded arr_t - - where padArray w_padded padding arr_t = do - let arr_shape = arrayShape arr_t - padding_shape = setDim d arr_shape padding - arr_padding <- - letExp (baseString arr <> "_padding") $ - BasicOp $ Scratch (elemType arr_t) (shapeDims padding_shape) - letExp (baseString arr <> "_padded") $ - BasicOp $ Concat d arr [arr_padding] w_padded + where + padArray w_padded padding arr_t = do + let arr_shape = arrayShape arr_t + padding_shape = setDim d arr_shape padding + arr_padding <- + letExp (baseString arr <> "_padding") $ + BasicOp $ Scratch (elemType arr_t) (shapeDims padding_shape) + letExp (baseString arr <> "_padded") $ + BasicOp $ Concat d arr [arr_padding] w_padded - rearrange num_chunks' w_padded per_chunk arr_name arr_padded arr_t = do - let arr_dims = arrayDims arr_t - pre_dims = take d arr_dims - post_dims = drop (d+1) arr_dims - extradim_shape = Shape $ pre_dims ++ [num_chunks', per_chunk] ++ post_dims - tr_perm = [0..d-1] ++ map (+d) ([1] ++ [2..shapeRank extradim_shape-1-d] ++ [0]) - arr_extradim <- - letExp (arr_name <> "_extradim") $ - BasicOp $ Reshape (map DimNew $ shapeDims extradim_shape) arr_padded - arr_extradim_tr <- - letExp (arr_name <> "_extradim_tr") $ - BasicOp $ Manifest tr_perm arr_extradim - arr_inv_tr <- letExp (arr_name <> "_inv_tr") $ - BasicOp $ Reshape (map DimCoercion pre_dims ++ map DimNew (w_padded : post_dims)) - arr_extradim_tr - letExp (arr_name <> "_inv_tr_init") =<< - eSliceArray d arr_inv_tr (eSubExp $ constant (0::Int32)) (eSubExp w) + rearrange num_chunks' w_padded per_chunk arr_name arr_padded arr_t = do + let arr_dims = arrayDims arr_t + pre_dims = take d arr_dims + post_dims = drop (d + 1) arr_dims + extradim_shape = Shape $ pre_dims ++ [num_chunks', per_chunk] ++ post_dims + tr_perm = [0 .. d -1] ++ map (+ d) ([1] ++ [2 .. shapeRank extradim_shape -1 - d] ++ [0]) + arr_extradim <- + letExp (arr_name <> "_extradim") $ + BasicOp $ Reshape (map DimNew $ shapeDims extradim_shape) arr_padded + arr_extradim_tr <- + letExp (arr_name <> "_extradim_tr") $ + BasicOp $ Manifest tr_perm arr_extradim + arr_inv_tr <- + letExp (arr_name <> "_inv_tr") $ + BasicOp $ + Reshape + (map DimCoercion pre_dims ++ map DimNew (w_padded : post_dims)) + arr_extradim_tr + letExp (arr_name <> "_inv_tr_init") + =<< eSliceArray d arr_inv_tr (eSubExp $ constant (0 :: Int64)) (eSubExp w) -paddedScanReduceInput :: MonadBinder m => - SubExp -> SubExp - -> m (SubExp, SubExp) +paddedScanReduceInput :: + MonadBinder m => + SubExp -> + SubExp -> + m (SubExp, SubExp) paddedScanReduceInput w stride = do - w_padded <- letSubExp "padded_size" =<< - eRoundToMultipleOf Int32 (eSubExp w) (eSubExp stride) - padding <- letSubExp "padding" $ BasicOp $ BinOp (Sub Int32 OverflowUndef) w_padded w + w_padded <- + letSubExp "padded_size" + =<< eRoundToMultipleOf Int64 (eSubExp w) (eSubExp stride) + padding <- letSubExp "padding" $ BasicOp $ BinOp (Sub Int64 OverflowUndef) w_padded w return (w_padded, padding) --- Computing variance. @@ -438,6 +512,7 @@ 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) + 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/Pass/Simplify.hs b/src/Futhark/Pass/Simplify.hs --- a/src/Futhark/Pass/Simplify.hs +++ b/src/Futhark/Pass/Simplify.hs @@ -1,25 +1,26 @@ {-# LANGUAGE FlexibleContexts #-} + module Futhark.Pass.Simplify - ( simplify - , simplifySOACS - , simplifySeq - , simplifyKernels - , simplifyKernelsMem - , simplifySeqMem + ( simplify, + simplifySOACS, + simplifySeq, + simplifyKernels, + simplifyKernelsMem, + simplifySeqMem, ) - where +where -import qualified Futhark.IR.SOACS.Simplify as SOACS import qualified Futhark.IR.Kernels.Simplify as Kernels -import qualified Futhark.IR.Seq as Seq import qualified Futhark.IR.KernelsMem as KernelsMem +import qualified Futhark.IR.SOACS.Simplify as SOACS +import qualified Futhark.IR.Seq as Seq import qualified Futhark.IR.SeqMem as SeqMem - -import Futhark.Pass import Futhark.IR.Syntax +import Futhark.Pass -simplify :: (Prog lore -> PassM (Prog lore)) - -> Pass lore lore +simplify :: + (Prog lore -> PassM (Prog lore)) -> + Pass lore lore simplify = Pass "simplify" "Perform simple enabling optimisations." simplifySOACS :: Pass SOACS.SOACS SOACS.SOACS diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs --- a/src/Futhark/Passes.hs +++ b/src/Futhark/Passes.hs @@ -1,23 +1,28 @@ {-# LANGUAGE FlexibleContexts #-} + -- | Optimisation pipelines. module Futhark.Passes - ( standardPipeline - , sequentialPipeline - , kernelsPipeline - , sequentialCpuPipeline - , gpuPipeline + ( standardPipeline, + sequentialPipeline, + kernelsPipeline, + sequentialCpuPipeline, + gpuPipeline, ) where import Control.Category ((>>>)) - +import Futhark.IR.Kernels (Kernels) +import Futhark.IR.KernelsMem (KernelsMem) +import Futhark.IR.SOACS (SOACS) +import Futhark.IR.Seq (Seq) +import Futhark.IR.SeqMem (SeqMem) import Futhark.Optimise.CSE +import Futhark.Optimise.DoubleBuffer import Futhark.Optimise.Fusion import Futhark.Optimise.InPlaceLowering import Futhark.Optimise.InliningDeadFun import Futhark.Optimise.Sink import Futhark.Optimise.TileLoops -import Futhark.Optimise.DoubleBuffer import Futhark.Optimise.Unstream import Futhark.Pass.ExpandAllocations import qualified Futhark.Pass.ExplicitAllocations.Kernels as Kernels @@ -27,69 +32,69 @@ import Futhark.Pass.KernelBabysitting import Futhark.Pass.Simplify import Futhark.Pipeline -import Futhark.IR.KernelsMem (KernelsMem) -import Futhark.IR.SeqMem (SeqMem) -import Futhark.IR.Kernels (Kernels) -import Futhark.IR.Seq (Seq) -import Futhark.IR.SOACS (SOACS) standardPipeline :: Pipeline SOACS SOACS standardPipeline = - passes [ simplifySOACS - , inlineFunctions - , simplifySOACS - , performCSE True - , simplifySOACS - -- We run fusion twice - , fuseSOACs - , performCSE True - , simplifySOACS - , fuseSOACs - , performCSE True - , simplifySOACS - , removeDeadFunctions - ] + passes + [ simplifySOACS, + inlineFunctions, + simplifySOACS, + performCSE True, + simplifySOACS, + -- We run fusion twice + fuseSOACs, + performCSE True, + simplifySOACS, + fuseSOACs, + performCSE True, + simplifySOACS, + removeDeadFunctions + ] kernelsPipeline :: Pipeline SOACS Kernels kernelsPipeline = - standardPipeline >>> - onePass extractKernels >>> - passes [ simplifyKernels - , babysitKernels - , tileLoops - , unstream - , performCSE True - , simplifyKernels - , sink - , inPlaceLoweringKernels - ] + standardPipeline + >>> onePass extractKernels + >>> passes + [ simplifyKernels, + babysitKernels, + tileLoops, + unstream, + performCSE True, + simplifyKernels, + sink, + inPlaceLoweringKernels + ] sequentialPipeline :: Pipeline SOACS Seq sequentialPipeline = - standardPipeline >>> - onePass firstOrderTransform >>> - passes [ simplifySeq - , inPlaceLoweringSeq - ] + standardPipeline + >>> onePass firstOrderTransform + >>> passes + [ simplifySeq, + inPlaceLoweringSeq + ] sequentialCpuPipeline :: Pipeline SOACS SeqMem sequentialCpuPipeline = - sequentialPipeline >>> - onePass Seq.explicitAllocations >>> - passes [ performCSE False - , simplifySeqMem - , simplifySeqMem - ] + sequentialPipeline + >>> onePass Seq.explicitAllocations + >>> passes + [ performCSE False, + simplifySeqMem, + simplifySeqMem + ] gpuPipeline :: Pipeline SOACS KernelsMem gpuPipeline = - kernelsPipeline >>> - onePass Kernels.explicitAllocations >>> - passes [ simplifyKernelsMem - , performCSE False - , simplifyKernelsMem - , doubleBuffer - , simplifyKernelsMem - , expandAllocations - , simplifyKernelsMem - ] + kernelsPipeline + >>> onePass Kernels.explicitAllocations + >>> passes + [ simplifyKernelsMem, + performCSE False, + simplifyKernelsMem, + doubleBuffer, + simplifyKernelsMem, + expandAllocations, + simplifyKernelsMem + ] diff --git a/src/Futhark/Pipeline.hs b/src/Futhark/Pipeline.hs --- a/src/Futhark/Pipeline.hs +++ b/src/Futhark/Pipeline.hs @@ -1,5 +1,8 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, OverloadedStrings #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Trustworthy #-} + -- | Definition of the core compiler driver building blocks. The -- spine of the compiler is the 'FutharkM' monad, although note that -- individual passes are pure functions, and do not use the 'FutharkM' @@ -10,152 +13,171 @@ -- final program (still in IR), then running an 'Action', which is -- usually a code generator. module Futhark.Pipeline - ( Pipeline - , PipelineConfig (..) - , Action (..) - - , FutharkM - , runFutharkM - , Verbosity(..) - - , module Futhark.Error - - , onePass - , passes - , runPipeline - ) - where + ( Pipeline, + PipelineConfig (..), + Action (..), + FutharkM, + runFutharkM, + Verbosity (..), + module Futhark.Error, + onePass, + passes, + runPipeline, + ) +where import Control.Category import Control.Monad -import Control.Monad.Writer.Strict hiding (pass) import Control.Monad.Except -import Control.Monad.State import Control.Monad.Reader +import Control.Monad.State +import Control.Monad.Writer.Strict hiding (pass) import qualified Data.Text as T import qualified Data.Text.IO as T import Data.Time.Clock -import System.IO -import Text.Printf - -import Prelude hiding (id, (.)) - import qualified Futhark.Analysis.Alias as Alias import Futhark.Error -import Futhark.IR (Prog, PrettyLore) -import Futhark.TypeCheck +import Futhark.IR (PrettyLore, Prog) +import Futhark.MonadFreshNames import Futhark.Pass +import Futhark.TypeCheck import Futhark.Util.Log import Futhark.Util.Pretty (prettyText) -import Futhark.MonadFreshNames +import System.IO +import Text.Printf +import Prelude hiding (id, (.)) -- | How much information to print to stderr while the compiler is running. data Verbosity - = NotVerbose -- ^ Silence is golden. - | Verbose -- ^ Print messages about which pass is running. - | VeryVerbose -- ^ Also print logs from individual passes. + = -- | Silence is golden. + NotVerbose + | -- | Print messages about which pass is running. + Verbose + | -- | Also print logs from individual passes. + VeryVerbose deriving (Eq, Ord) -newtype FutharkEnv = FutharkEnv { futharkVerbose :: Verbosity } +newtype FutharkEnv = FutharkEnv {futharkVerbose :: Verbosity} -data FutharkState = FutharkState { futharkPrevLog :: UTCTime - , futharkNameSource :: VNameSource } +data FutharkState = FutharkState + { futharkPrevLog :: UTCTime, + futharkNameSource :: VNameSource + } -- | The main Futhark compiler driver monad - basically some state -- tracking on top if 'IO'. newtype FutharkM a = FutharkM (ExceptT CompilerError (StateT FutharkState (ReaderT FutharkEnv IO)) a) - deriving (Applicative, Functor, Monad, - MonadError CompilerError, - MonadState FutharkState, - MonadReader FutharkEnv, - MonadIO) + deriving + ( Applicative, + Functor, + Monad, + MonadError CompilerError, + MonadState FutharkState, + MonadReader FutharkEnv, + MonadIO + ) instance MonadFreshNames FutharkM where getNameSource = gets futharkNameSource - putNameSource src = modify $ \s -> s { futharkNameSource = src } + putNameSource src = modify $ \s -> s {futharkNameSource = src} instance MonadLogger FutharkM where addLog = mapM_ perLine . T.lines . toText - where perLine msg = do - verb <- asks $ (>=Verbose) . futharkVerbose - prev <- gets futharkPrevLog - now <- liftIO getCurrentTime - let delta :: Double - delta = fromRational $ toRational (now `diffUTCTime` prev) - prefix = printf "[ +%.6f] " delta - modify $ \s -> s { futharkPrevLog = now } - when verb $ liftIO $ T.hPutStrLn stderr $ T.pack prefix <> msg + where + perLine msg = do + verb <- asks $ (>= Verbose) . futharkVerbose + prev <- gets futharkPrevLog + now <- liftIO getCurrentTime + let delta :: Double + delta = fromRational $ toRational (now `diffUTCTime` prev) + prefix = printf "[ +%.6f] " delta + modify $ \s -> s {futharkPrevLog = now} + when verb $ liftIO $ T.hPutStrLn stderr $ T.pack prefix <> msg -- | Run a 'FutharkM' action. runFutharkM :: FutharkM a -> Verbosity -> IO (Either CompilerError a) runFutharkM (FutharkM m) verbose = do s <- FutharkState <$> getCurrentTime <*> pure blankNameSource runReaderT (evalStateT (runExceptT m) s) newEnv - where newEnv = FutharkEnv verbose + where + newEnv = FutharkEnv verbose -- | A compilation always ends with some kind of action. -data Action lore = - Action { actionName :: String - , actionDescription :: String - , actionProcedure :: Prog lore -> FutharkM () - } +data Action lore = Action + { actionName :: String, + actionDescription :: String, + actionProcedure :: Prog lore -> FutharkM () + } -- | Configuration object for running a compiler pipeline. -data PipelineConfig = - PipelineConfig { pipelineVerbose :: Bool - , pipelineValidate :: Bool - } +data PipelineConfig = PipelineConfig + { pipelineVerbose :: Bool, + pipelineValidate :: Bool + } -- | A compiler pipeline is conceptually a function from programs to -- programs, where the actual representation may change. Pipelines -- can be composed using their 'Category' instance. -newtype Pipeline fromlore tolore = - Pipeline { unPipeline :: PipelineConfig -> Prog fromlore -> FutharkM (Prog tolore) } +newtype Pipeline fromlore tolore = Pipeline {unPipeline :: PipelineConfig -> Prog fromlore -> FutharkM (Prog tolore)} instance Category Pipeline where id = Pipeline $ const return p2 . p1 = Pipeline perform - where perform cfg prog = - runPipeline p2 cfg =<< runPipeline p1 cfg prog + where + perform cfg prog = + runPipeline p2 cfg =<< runPipeline p1 cfg prog -- | Run the pipeline on the given program. -runPipeline :: Pipeline fromlore tolore - -> PipelineConfig - -> Prog fromlore - -> FutharkM (Prog tolore) +runPipeline :: + Pipeline fromlore tolore -> + PipelineConfig -> + Prog fromlore -> + FutharkM (Prog tolore) runPipeline = unPipeline -- | Construct a pipeline from a single compiler pass. -onePass :: Checkable tolore => - Pass fromlore tolore -> Pipeline fromlore tolore +onePass :: + Checkable tolore => + Pass fromlore tolore -> + Pipeline fromlore tolore onePass pass = Pipeline perform - where perform cfg prog = do - when (pipelineVerbose cfg) $ logMsg $ - "Running pass " <> T.pack (passName pass) - prog' <- runPass pass prog - let prog'' = Alias.aliasAnalysis prog' - when (pipelineValidate cfg) $ - case checkProg prog'' of - Left err -> validationError pass prog'' $ show err - Right () -> return () - return prog' + where + perform cfg prog = do + when (pipelineVerbose cfg) $ + logMsg $ + "Running pass " <> T.pack (passName pass) + prog' <- runPass pass prog + let prog'' = Alias.aliasAnalysis prog' + when (pipelineValidate cfg) $ + case checkProg prog'' of + Left err -> validationError pass prog'' $ show err + Right () -> return () + return prog' -- | Create a pipeline from a list of passes. -passes :: Checkable lore => - [Pass lore lore] -> Pipeline lore lore +passes :: + Checkable lore => + [Pass lore lore] -> + Pipeline lore lore passes = foldl (>>>) id . map onePass -validationError :: PrettyLore lore => - Pass fromlore tolore -> Prog lore -> String -> FutharkM a +validationError :: + PrettyLore lore => + Pass fromlore tolore -> + Prog lore -> + String -> + FutharkM a validationError pass prog err = throwError $ InternalError msg (prettyText prog) CompilerBug - where msg = "Type error after pass '" <> T.pack (passName pass) <> "':\n" <> T.pack err + where + msg = "Type error after pass '" <> T.pack (passName pass) <> "':\n" <> T.pack err -runPass :: Pass fromlore tolore - -> Prog fromlore - -> FutharkM (Prog tolore) +runPass :: + Pass fromlore tolore -> + Prog fromlore -> + FutharkM (Prog tolore) runPass pass prog = do (prog', logged) <- runPassM (passFunction pass prog) - verb <- asks $ (>=VeryVerbose) . futharkVerbose + verb <- asks $ (>= VeryVerbose) . futharkVerbose when verb $ addLog logged return prog' diff --git a/src/Futhark/Pkg/Info.hs b/src/Futhark/Pkg/Info.hs --- a/src/Futhark/Pkg/Info.hs +++ b/src/Futhark/Pkg/Info.hs @@ -1,44 +1,43 @@ {-# LANGUAGE OverloadedStrings #-} + -- | Obtaining information about packages over THE INTERNET! module Futhark.Pkg.Info ( -- * Package info - PkgInfo(..) - , lookupPkgRev - , pkgInfo - , PkgRevInfo (..) - , GetManifest (getManifest) - , downloadZipball + PkgInfo (..), + lookupPkgRev, + pkgInfo, + PkgRevInfo (..), + GetManifest (getManifest), + downloadZipball, -- * Package registry - , PkgRegistry - , MonadPkgRegistry(..) - , lookupPackage - , lookupPackageRev - , lookupNewestRev + PkgRegistry, + MonadPkgRegistry (..), + lookupPackage, + lookupPackageRev, + lookupNewestRev, ) - where +where +import qualified Codec.Archive.Zip as Zip import Control.Monad.IO.Class -import Data.Maybe +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as LBS import Data.IORef +import Data.List (foldl', intersperse) import qualified Data.Map as M +import Data.Maybe import qualified Data.Text as T -import qualified Data.ByteString as BS -import qualified Data.ByteString.Lazy as LBS import qualified Data.Text.Encoding as T -import Data.List (foldl', intersperse) -import qualified System.FilePath.Posix as Posix +import Data.Time (UTCTime, defaultTimeLocale, formatTime, getCurrentTime) +import Futhark.Pkg.Types +import Futhark.Util (maybeHead) +import Futhark.Util.Log import System.Exit +import qualified System.FilePath.Posix as Posix import System.IO - -import qualified Codec.Archive.Zip as Zip -import Data.Time (UTCTime, UTCTime, defaultTimeLocale, formatTime, getCurrentTime) import System.Process.ByteString (readProcessWithExitCode) -import Futhark.Pkg.Types -import Futhark.Util.Log -import Futhark.Util (maybeHead) - -- | Download URL via shelling out to @curl@. curl :: String -> IO (Either String BS.ByteString) curl url = do @@ -47,8 +46,9 @@ liftIO $ readProcessWithExitCode "curl" ["-L", url] mempty case code of ExitFailure 127 -> - return $ Left $ - "'" <> unwords ["curl", "-L", url] <> "' failed (program not found?)." + return $ + Left $ + "'" <> unwords ["curl", "-L", url] <> "' failed (program not found?)." ExitFailure _ -> do liftIO $ BS.hPutStr stderr err return $ Left $ "'" <> unwords ["curl", "-L", url] <> "' failed." @@ -59,7 +59,7 @@ -- fetch them on-demand. It would be a waste to fetch it information -- for every version of every package if we only actually need a small -- subset of them. -newtype GetManifest m = GetManifest { getManifest :: m PkgManifest } +newtype GetManifest m = GetManifest {getManifest :: m PkgManifest} instance Show (GetManifest m) where show _ = "#" @@ -69,24 +69,25 @@ -- | Information about a version of a single package. The version -- number is stored separately. -data PkgRevInfo m = PkgRevInfo { pkgRevZipballUrl :: T.Text - , pkgRevZipballDir :: FilePath - -- ^ The directory inside the zipball - -- containing the @lib@ directory, in - -- which the package files themselves - -- are stored (Based on the package - -- path). - , pkgRevCommit :: T.Text - -- ^ The commit ID can be used for - -- verification ("freezing"), by - -- storing what it was at the time this - -- version was last selected. - , pkgRevGetManifest :: GetManifest m - , pkgRevTime :: UTCTime - -- ^ Timestamp for when the revision - -- was made (rarely used). - } - deriving (Eq, Show) +data PkgRevInfo m = PkgRevInfo + { pkgRevZipballUrl :: T.Text, + -- | The directory inside the zipball + -- containing the @lib@ directory, in + -- which the package files themselves + -- are stored (Based on the package + -- path). + pkgRevZipballDir :: FilePath, + -- | The commit ID can be used for + -- verification ("freezing"), by + -- storing what it was at the time this + -- version was last selected. + pkgRevCommit :: T.Text, + pkgRevGetManifest :: GetManifest m, + -- | Timestamp for when the revision + -- was made (rarely used). + pkgRevTime :: UTCTime + } + deriving (Eq, Show) -- | Create memoisation around a 'GetManifest' action to ensure that -- multiple inspections of the same revisions will not result in @@ -94,23 +95,27 @@ memoiseGetManifest :: MonadIO m => GetManifest m -> m (GetManifest m) memoiseGetManifest (GetManifest m) = do ref <- liftIO $ newIORef Nothing - return $ GetManifest $ do - v <- liftIO $ readIORef ref - case v of Just v' -> return v' - Nothing -> do - v' <- m - liftIO $ writeIORef ref $ Just v' - return v' + return $ + GetManifest $ do + v <- liftIO $ readIORef ref + case v of + Just v' -> return v' + Nothing -> do + v' <- m + liftIO $ writeIORef ref $ Just v' + return v' -- | Download the zip archive corresponding to a specific package -- version. -downloadZipball :: (MonadLogger m, MonadIO m, MonadFail m) => - PkgRevInfo m -> m Zip.Archive +downloadZipball :: + (MonadLogger m, MonadIO m, MonadFail m) => + PkgRevInfo m -> + m Zip.Archive downloadZipball info = do let url = pkgRevZipballUrl info logMsg $ "Downloading " <> T.unpack url - let bad = fail . (("When downloading " <> T.unpack url <> ": ")<>) + let bad = fail . (("When downloading " <> T.unpack url <> ": ") <>) http <- liftIO $ curl $ T.unpack url case http of Left e -> bad e @@ -121,11 +126,12 @@ -- | Information about a package. The name of the package is stored -- separately. -data PkgInfo m = PkgInfo { pkgVersions :: M.Map SemVer (PkgRevInfo m) - , pkgLookupCommit :: Maybe T.Text -> m (PkgRevInfo m) - -- ^ Look up information about a specific - -- commit, or HEAD in case of Nothing. - } +data PkgInfo m = PkgInfo + { pkgVersions :: M.Map SemVer (PkgRevInfo m), + -- | Look up information about a specific + -- commit, or HEAD in case of Nothing. + pkgLookupCommit :: Maybe T.Text -> m (PkgRevInfo m) + } -- | Lookup information about a given version of a package. lookupPkgRev :: SemVer -> PkgInfo m -> Maybe (PkgRevInfo m) @@ -135,31 +141,40 @@ majorRevOfPkg p = case T.splitOn "@" p of [p', v] | [(v', "")] <- reads $ T.unpack v -> (p', [v']) - _ -> (p, [0, 1]) + _ -> (p, [0, 1]) -- | Retrieve information about a package based on its package path. -- This uses Semantic Import Versioning when interacting with -- repositories. For example, a package @github.com/user/repo@ will -- match version 0.* or 1.* tags only, a package -- @github.com/user/repo/v2@ will match 2.* tags, and so forth.. -pkgInfo :: (MonadIO m, MonadLogger m, MonadFail m) => - PkgPath -> m (Either T.Text (PkgInfo m)) +pkgInfo :: + (MonadIO m, MonadLogger m, MonadFail m) => + PkgPath -> + m (Either T.Text (PkgInfo m)) pkgInfo path | ["github.com", owner, repo] <- T.splitOn "/" path = - let (repo', vs) = majorRevOfPkg repo - in ghPkgInfo owner repo' vs - | "github.com": owner : repo : _ <- T.splitOn "/" path = - return $ Left $ T.intercalate "\n" - [nope, "Do you perhaps mean 'github.com/" <> owner <> "/" <> repo <> "'?"] + let (repo', vs) = majorRevOfPkg repo + in ghPkgInfo owner repo' vs + | "github.com" : owner : repo : _ <- T.splitOn "/" path = + return $ + Left $ + T.intercalate + "\n" + [nope, "Do you perhaps mean 'github.com/" <> owner <> "/" <> repo <> "'?"] | ["gitlab.com", owner, repo] <- T.splitOn "/" path = - let (repo', vs) = majorRevOfPkg repo - in glPkgInfo owner repo' vs - | "gitlab.com": owner : repo : _ <- T.splitOn "/" path = - return $ Left $ T.intercalate "\n" - [nope, "Do you perhaps mean 'gitlab.com/" <> owner <> "/" <> repo <> "'?"] + let (repo', vs) = majorRevOfPkg repo + in glPkgInfo owner repo' vs + | "gitlab.com" : owner : repo : _ <- T.splitOn "/" path = + return $ + Left $ + T.intercalate + "\n" + [nope, "Do you perhaps mean 'gitlab.com/" <> owner <> "/" <> repo <> "'?"] | otherwise = - return $ Left nope - where nope = "Unable to handle package paths of the form '" <> path <> "'" + return $ Left nope + where + nope = "Unable to handle package paths of the form '" <> path <> "'" -- For GitHub, we unfortunately cannot use the (otherwise very nice) -- GitHub web API, because it is rate-limited to 60 requests per hour @@ -182,14 +197,23 @@ -- couple of generic functions that are used to implement support for -- both. -ghglRevGetManifest :: (MonadIO m, MonadLogger m, MonadFail m) => - T.Text -> T.Text -> T.Text -> T.Text -> GetManifest m +ghglRevGetManifest :: + (MonadIO m, MonadLogger m, MonadFail m) => + T.Text -> + T.Text -> + T.Text -> + T.Text -> + GetManifest m ghglRevGetManifest url owner repo tag = GetManifest $ do logMsg $ "Downloading package manifest from " <> url - let path = T.unpack $ owner <> "/" <> repo <> "@" <> - tag <> "/" <> T.pack futharkPkg - msg = (("When reading " <> path <> ": ")<>) + let path = + T.unpack $ + owner <> "/" <> repo <> "@" + <> tag + <> "/" + <> T.pack futharkPkg + msg = (("When reading " <> path <> ": ") <>) http <- liftIO $ curl $ T.unpack url case http of Left e -> fail e @@ -201,75 +225,142 @@ Left e -> fail $ msg $ errorBundlePretty e Right pm -> return pm -ghglLookupCommit :: (MonadIO m, MonadLogger m, MonadFail m) => - T.Text -> T.Text -> (T.Text -> T.Text) - -> T.Text -> T.Text -> T.Text -> T.Text -> T.Text -> m (PkgRevInfo m) +ghglLookupCommit :: + (MonadIO m, MonadLogger m, MonadFail m) => + T.Text -> + T.Text -> + (T.Text -> T.Text) -> + T.Text -> + T.Text -> + T.Text -> + T.Text -> + T.Text -> + m (PkgRevInfo m) ghglLookupCommit archive_url manifest_url mk_zip_dir owner repo d ref hash = do gd <- memoiseGetManifest $ ghglRevGetManifest manifest_url owner repo ref let dir = Posix.addTrailingPathSeparator $ T.unpack $ mk_zip_dir d time <- liftIO getCurrentTime -- FIXME return $ PkgRevInfo archive_url dir hash gd time -ghglPkgInfo :: (MonadIO m, MonadLogger m, MonadFail m) => - T.Text - -> (T.Text -> T.Text) -> (T.Text -> T.Text) -> (T.Text -> T.Text) - -> T.Text -> T.Text -> [Word] -> m (Either T.Text (PkgInfo m)) +ghglPkgInfo :: + (MonadIO m, MonadLogger m, MonadFail m) => + T.Text -> + (T.Text -> T.Text) -> + (T.Text -> T.Text) -> + (T.Text -> T.Text) -> + T.Text -> + T.Text -> + [Word] -> + m (Either T.Text (PkgInfo m)) ghglPkgInfo repo_url mk_archive_url mk_manifest_url mk_zip_dir owner repo versions = do logMsg $ "Retrieving list of tags from " <> repo_url remote_lines <- T.lines . T.decodeUtf8 <$> gitCmd ["ls-remote", T.unpack repo_url] - head_ref <- maybe (fail $ "Cannot find HEAD ref for " <> T.unpack repo_url) return $ - maybeHead $ mapMaybe isHeadRef remote_lines + head_ref <- + maybe (fail $ "Cannot find HEAD ref for " <> T.unpack repo_url) return $ + maybeHead $ mapMaybe isHeadRef remote_lines let def = fromMaybe head_ref rev_info <- M.fromList . catMaybes <$> mapM revInfo remote_lines - return $ Right $ PkgInfo rev_info $ \r -> - ghglLookupCommit - (mk_archive_url (def r)) (mk_manifest_url (def r)) mk_zip_dir - owner repo (def r) (def r) (def r) - where isHeadRef l - | [hash, "HEAD"] <- T.words l = Just hash - | otherwise = Nothing + return $ + Right $ + PkgInfo rev_info $ \r -> + ghglLookupCommit + (mk_archive_url (def r)) + (mk_manifest_url (def r)) + mk_zip_dir + owner + repo + (def r) + (def r) + (def r) + where + isHeadRef l + | [hash, "HEAD"] <- T.words l = Just hash + | otherwise = Nothing - revInfo l - | [hash, ref] <- T.words l, - ["refs", "tags", t] <- T.splitOn "/" ref, - "v" `T.isPrefixOf` t, - Right v <- semver $ T.drop 1 t, - _svMajor v `elem` versions = do - pinfo <- ghglLookupCommit - (mk_archive_url t) (mk_manifest_url t) mk_zip_dir - owner repo (prettySemVer v) t hash - return $ Just (v, pinfo) - | otherwise = return Nothing + revInfo l + | [hash, ref] <- T.words l, + ["refs", "tags", t] <- T.splitOn "/" ref, + "v" `T.isPrefixOf` t, + Right v <- semver $ T.drop 1 t, + _svMajor v `elem` versions = do + pinfo <- + ghglLookupCommit + (mk_archive_url t) + (mk_manifest_url t) + mk_zip_dir + owner + repo + (prettySemVer v) + t + hash + return $ Just (v, pinfo) + | otherwise = return Nothing -ghPkgInfo :: (MonadIO m, MonadLogger m, MonadFail m) => - T.Text -> T.Text -> [Word] -> m (Either T.Text (PkgInfo m)) +ghPkgInfo :: + (MonadIO m, MonadLogger m, MonadFail m) => + T.Text -> + T.Text -> + [Word] -> + m (Either T.Text (PkgInfo m)) ghPkgInfo owner repo versions = - ghglPkgInfo repo_url mk_archive_url mk_manifest_url mk_zip_dir - owner repo versions - where repo_url = "https://github.com/" <> owner <> "/" <> repo - mk_archive_url r = repo_url <> "/archive/" <> r <> ".zip" - mk_manifest_url r = "https://raw.githubusercontent.com/" <> - owner <> "/" <> repo <> "/" <> - r <> "/" <> T.pack futharkPkg - mk_zip_dir r = repo <> "-" <> r + ghglPkgInfo + repo_url + mk_archive_url + mk_manifest_url + mk_zip_dir + owner + repo + versions + where + repo_url = "https://github.com/" <> owner <> "/" <> repo + mk_archive_url r = repo_url <> "/archive/" <> r <> ".zip" + mk_manifest_url r = + "https://raw.githubusercontent.com/" + <> owner + <> "/" + <> repo + <> "/" + <> r + <> "/" + <> T.pack futharkPkg + mk_zip_dir r = repo <> "-" <> r -glPkgInfo :: (MonadIO m, MonadLogger m, MonadFail m) => - T.Text -> T.Text -> [Word] -> m (Either T.Text (PkgInfo m)) +glPkgInfo :: + (MonadIO m, MonadLogger m, MonadFail m) => + T.Text -> + T.Text -> + [Word] -> + m (Either T.Text (PkgInfo m)) glPkgInfo owner repo versions = - ghglPkgInfo repo_url mk_archive_url mk_manifest_url mk_zip_dir - owner repo versions - where base_url = "https://gitlab.com/" <> owner <> "/" <> repo - repo_url = base_url <> ".git" - mk_archive_url r = base_url <> "/-/archive/" <> r <> - "/" <> repo <> "-" <> r <> ".zip" - mk_manifest_url r = base_url <> "/raw/" <> - r <> "/" <> T.pack futharkPkg - mk_zip_dir r - | Right _ <- semver r = repo <> "-v" <> r - | otherwise = repo <> "-" <> r + ghglPkgInfo + repo_url + mk_archive_url + mk_manifest_url + mk_zip_dir + owner + repo + versions + where + base_url = "https://gitlab.com/" <> owner <> "/" <> repo + repo_url = base_url <> ".git" + mk_archive_url r = + base_url <> "/-/archive/" <> r + <> "/" + <> repo + <> "-" + <> r + <> ".zip" + mk_manifest_url r = + base_url <> "/raw/" + <> r + <> "/" + <> T.pack futharkPkg + mk_zip_dir r + | Right _ <- semver r = repo <> "-v" <> r + | otherwise = repo <> "-" <> r -- | A package registry is a mapping from package paths to information -- about the package. It is unlikely that any given registry is @@ -297,8 +388,10 @@ modifyPkgRegistry f = putPkgRegistry . f =<< getPkgRegistry -- | Given a package path, look up information about that package. -lookupPackage :: MonadPkgRegistry m => - PkgPath -> m (PkgInfo m) +lookupPackage :: + MonadPkgRegistry m => + PkgPath -> + m (PkgInfo m) lookupPackage p = do r@(PkgRegistry m) <- getPkgRegistry case lookupKnownPackage p r of @@ -312,50 +405,67 @@ putPkgRegistry $ PkgRegistry $ M.insert p pinfo m return pinfo -lookupPackageCommit :: MonadPkgRegistry m => - PkgPath -> Maybe T.Text -> m (SemVer, PkgRevInfo m) +lookupPackageCommit :: + MonadPkgRegistry m => + PkgPath -> + Maybe T.Text -> + m (SemVer, PkgRevInfo m) lookupPackageCommit p ref = do pinfo <- lookupPackage p rev_info <- pkgLookupCommit pinfo ref - let timestamp = T.pack $ formatTime defaultTimeLocale "%Y%m%d%H%M%S" $ - pkgRevTime rev_info + let timestamp = + T.pack $ + formatTime defaultTimeLocale "%Y%m%d%H%M%S" $ + pkgRevTime rev_info v = commitVersion timestamp $ pkgRevCommit rev_info - pinfo' = pinfo { pkgVersions = M.insert v rev_info $ pkgVersions pinfo } + pinfo' = pinfo {pkgVersions = M.insert v rev_info $ pkgVersions pinfo} modifyPkgRegistry $ \(PkgRegistry m) -> PkgRegistry $ M.insert p pinfo' m return (v, rev_info) -- | Look up information about a specific version of a package. -lookupPackageRev :: MonadPkgRegistry m => - PkgPath -> SemVer -> m (PkgRevInfo m) +lookupPackageRev :: + MonadPkgRegistry m => + PkgPath -> + SemVer -> + m (PkgRevInfo m) lookupPackageRev p v | Just commit <- isCommitVersion v = - snd <$> lookupPackageCommit p (Just commit) + snd <$> lookupPackageCommit p (Just commit) | otherwise = do - pinfo <- lookupPackage p - case lookupPkgRev v pinfo of - Nothing -> - let versions = case M.keys $ pkgVersions pinfo of - [] -> "Package " <> p <> " has no versions. Invalid package path?" - ks -> "Known versions: " <> - T.concat (intersperse ", " $ map prettySemVer ks) - major | (_, vs) <- majorRevOfPkg p, - _svMajor v `notElem` vs = - "\nFor major version " <> T.pack (show (_svMajor v)) <> - ", use package path " <> p <> "@" <> T.pack (show (_svMajor v)) - | otherwise = mempty - in fail $ T.unpack $ - "package " <> p <> " does not have a version " <> prettySemVer v <> ".\n" <> - versions <> major - Just v' -> return v' + pinfo <- lookupPackage p + case lookupPkgRev v pinfo of + Nothing -> + let versions = case M.keys $ pkgVersions pinfo of + [] -> "Package " <> p <> " has no versions. Invalid package path?" + ks -> + "Known versions: " + <> T.concat (intersperse ", " $ map prettySemVer ks) + major + | (_, vs) <- majorRevOfPkg p, + _svMajor v `notElem` vs = + "\nFor major version " <> T.pack (show (_svMajor v)) + <> ", use package path " + <> p + <> "@" + <> T.pack (show (_svMajor v)) + | otherwise = mempty + in fail $ + T.unpack $ + "package " <> p <> " does not have a version " <> prettySemVer v <> ".\n" + <> versions + <> major + Just v' -> return v' -- | Find the newest version of a package. -lookupNewestRev :: MonadPkgRegistry m => - PkgPath -> m SemVer +lookupNewestRev :: + MonadPkgRegistry m => + PkgPath -> + m SemVer lookupNewestRev p = do pinfo <- lookupPackage p case M.keys $ pkgVersions pinfo of [] -> do logMsg $ "Package " <> p <> " has no released versions. Using HEAD." fst <$> lookupPackageCommit p Nothing - v:vs -> return $ foldl' max v vs + v : vs -> return $ foldl' max v vs diff --git a/src/Futhark/Pkg/Solve.hs b/src/Futhark/Pkg/Solve.hs --- a/src/Futhark/Pkg/Solve.hs +++ b/src/Futhark/Pkg/Solve.hs @@ -1,26 +1,25 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} + -- | Dependency solver -- -- This is a relatively simple problem due to the choice of the -- Minimum Package Version algorithm. In fact, the only failure mode -- is referencing an unknown package or revision. module Futhark.Pkg.Solve - ( solveDeps - , solveDepsPure - , PkgRevDepInfo - ) where + ( solveDeps, + solveDepsPure, + PkgRevDepInfo, + ) +where +import Control.Monad.Free.Church import Control.Monad.State -import qualified Data.Set as S import qualified Data.Map as M +import qualified Data.Set as S import qualified Data.Text as T - -import Control.Monad.Free.Church - import Futhark.Pkg.Info import Futhark.Pkg.Types - import Prelude data PkgOp a = OpGetDeps PkgPath SemVer (Maybe T.Text) (PkgRevDeps -> a) @@ -32,7 +31,7 @@ -- that are not reachable from the root. Also contains the -- dependencies of each package. newtype RoughBuildList = RoughBuildList (M.Map PkgPath (SemVer, [PkgPath])) - deriving (Show) + deriving (Show) emptyRoughBuildList :: RoughBuildList emptyRoughBuildList = RoughBuildList mempty @@ -45,12 +44,13 @@ buildList :: S.Set PkgPath -> RoughBuildList -> BuildList buildList roots (RoughBuildList pkgs) = BuildList $ execState (mapM_ addPkg roots) mempty - where addPkg p = case M.lookup p pkgs of - Nothing -> return () - Just (v, deps) -> do - listed <- gets $ M.member p - modify $ M.insert p v - unless listed $ mapM_ addPkg deps + where + addPkg p = case M.lookup p pkgs of + Nothing -> return () + Just (v, deps) -> do + listed <- gets $ M.member p + modify $ M.insert p v + unless listed $ mapM_ addPkg deps type SolveM = StateT RoughBuildList (F PkgOp) @@ -62,38 +62,51 @@ -- dependencies. doSolveDeps :: PkgRevDeps -> SolveM () doSolveDeps (PkgRevDeps deps) = mapM_ add $ M.toList deps - where add (p, (v, maybe_h)) = do - RoughBuildList l <- get - case M.lookup p l of - -- Already satisfied? - Just (cur_v, _) | v <= cur_v -> return () - -- No; add 'p' and its dependencies. - _ -> do - PkgRevDeps p_deps <- getDeps p v maybe_h - put $ RoughBuildList $ M.insert p (v, M.keys p_deps) l - mapM_ add $ M.toList p_deps + where + add (p, (v, maybe_h)) = do + RoughBuildList l <- get + case M.lookup p l of + -- Already satisfied? + Just (cur_v, _) | v <= cur_v -> return () + -- No; add 'p' and its dependencies. + _ -> do + PkgRevDeps p_deps <- getDeps p v maybe_h + put $ RoughBuildList $ M.insert p (v, M.keys p_deps) l + mapM_ add $ M.toList p_deps -- | Run the solver, producing both a package registry containing -- a cache of the lookups performed, as well as a build list. -solveDeps :: MonadPkgRegistry m => - PkgRevDeps -> m BuildList -solveDeps deps = buildList (depRoots deps) <$> runF - (execStateT (doSolveDeps deps) emptyRoughBuildList) - return step - where step (OpGetDeps p v h c) = do - pinfo <- lookupPackageRev p v +solveDeps :: + MonadPkgRegistry m => + PkgRevDeps -> + m BuildList +solveDeps deps = + buildList (depRoots deps) + <$> runF + (execStateT (doSolveDeps deps) emptyRoughBuildList) + return + step + where + step (OpGetDeps p v h c) = do + pinfo <- lookupPackageRev p v - checkHash p v pinfo h + checkHash p v pinfo h - d <- fmap pkgRevDeps . getManifest $ pkgRevGetManifest pinfo - c d + d <- fmap pkgRevDeps . getManifest $ pkgRevGetManifest pinfo + c d - checkHash _ _ _ Nothing = return () - checkHash p v pinfo (Just h) - | h == pkgRevCommit pinfo = return () - | otherwise = fail $ T.unpack $ "Package " <> p <> " " <> prettySemVer v <> - " has commit hash " <> pkgRevCommit pinfo <> - ", but expected " <> h <> " from package manifest." + checkHash _ _ _ Nothing = return () + checkHash p v pinfo (Just h) + | h == pkgRevCommit pinfo = return () + | otherwise = + fail $ + T.unpack $ + "Package " <> p <> " " <> prettySemVer v + <> " has commit hash " + <> pkgRevCommit pinfo + <> ", but expected " + <> h + <> " from package manifest." -- | A mapping of package revisions to the dependencies of that -- package. Can be considered a 'PkgRegistry' without the option of @@ -104,10 +117,14 @@ -- | Perform package resolution with only pre-known information. This -- is useful for testing. solveDepsPure :: PkgRevDepInfo -> PkgRevDeps -> Either T.Text BuildList -solveDepsPure r deps = buildList (depRoots deps) <$> runF - (execStateT (doSolveDeps deps) emptyRoughBuildList) - Right step - where step (OpGetDeps p v _ c) = do - let errmsg = "Unknown package/version: " <> p <> "-" <> prettySemVer v - d <- maybe (Left errmsg) Right $ M.lookup (p,v) r - c d +solveDepsPure r deps = + buildList (depRoots deps) + <$> runF + (execStateT (doSolveDeps deps) emptyRoughBuildList) + Right + step + where + step (OpGetDeps p v _ c) = do + let errmsg = "Unknown package/version: " <> p <> "-" <> prettySemVer v + d <- maybe (Left errmsg) Right $ M.lookup (p, v) r + c d diff --git a/src/Futhark/Pkg/Types.hs b/src/Futhark/Pkg/Types.hs --- a/src/Futhark/Pkg/Types.hs +++ b/src/Futhark/Pkg/Types.hs @@ -1,57 +1,57 @@ {-# LANGUAGE OverloadedStrings #-} + -- | Types (and a few other simple definitions) for futhark-pkg. module Futhark.Pkg.Types - ( PkgPath - , pkgPathFilePath - , PkgRevDeps(..) - , module Data.Versions + ( PkgPath, + pkgPathFilePath, + PkgRevDeps (..), + module Data.Versions, - -- * Versions - , commitVersion - , isCommitVersion - , parseVersion + -- * Versions + commitVersion, + isCommitVersion, + parseVersion, - -- * Package manifests - , PkgManifest(..) - , newPkgManifest - , pkgRevDeps - , pkgDir - , addRequiredToManifest - , removeRequiredFromManifest - , prettyPkgManifest - , Comment - , Commented(..) - , Required(..) - , futharkPkg + -- * Package manifests + PkgManifest (..), + newPkgManifest, + pkgRevDeps, + pkgDir, + addRequiredToManifest, + removeRequiredFromManifest, + prettyPkgManifest, + Comment, + Commented (..), + Required (..), + futharkPkg, - -- * Parsing package manifests - , parsePkgManifest - , parsePkgManifestFromFile - , errorBundlePretty + -- * Parsing package manifests + parsePkgManifest, + parsePkgManifestFromFile, + errorBundlePretty, - -- * Build list - , BuildList(..) - , prettyBuildList - ) where + -- * Build list + BuildList (..), + prettyBuildList, + ) +where import Control.Applicative import Control.Monad import Data.Either import Data.Foldable import Data.List (sortOn) +import qualified Data.Map as M import Data.Maybe -import Data.Traversable -import Data.Void import qualified Data.Text as T import qualified Data.Text.IO as T -import qualified Data.Map as M +import Data.Traversable +import Data.Versions (SemVer (..), VUnit (..), prettySemVer, semver) +import Data.Void import System.FilePath import qualified System.FilePath.Posix as Posix - -import Data.Versions (semver, SemVer(..), VUnit(..), prettySemVer) import Text.Megaparsec hiding (many, some) import Text.Megaparsec.Char - import Prelude -- | A package path is a unique identifier for a package, for example @@ -86,14 +86,15 @@ semver' :: Parsec Void T.Text SemVer semver' = SemVer <$> majorP <*> minorP <*> patchP <*> preRel <*> metaData - where majorP = digitsP <* char '.' - minorP = majorP - patchP = digitsP - digitsP = read <$> ((T.unpack <$> string "0") <|> some digitChar) - preRel = maybe [] pure <$> optional preRel' - preRel' = char '-' *> (pure . Str . T.pack <$> some digitChar) - metaData = maybe [] pure <$> optional metaData' - metaData' = char '+' *> (pure . Str . T.pack <$> some alphaNumChar) + where + majorP = digitsP <* char '.' + minorP = majorP + patchP = digitsP + digitsP = read <$> ((T.unpack <$> string "0") <|> some digitChar) + preRel = maybe [] pure <$> optional preRel' + preRel' = char '-' *> (pure . Str . T.pack <$> some digitChar) + metaData = maybe [] pure <$> optional metaData' + metaData' = char '+' *> (pure . Str . T.pack <$> some alphaNumChar) -- | The dependencies of a (revision of a) package is a mapping from -- package paths to minimum versions (and an optional hash pinning). @@ -114,10 +115,11 @@ -- | Wraps a value with an annotation of preceding line comments. -- This is important to our goal of being able to programmatically -- modify the @futhark.pkg@ file while keeping comments intact. -data Commented a = Commented { comments :: [Comment] - , commented :: a - } - deriving (Show, Eq) +data Commented a = Commented + { comments :: [Comment], + commented :: a + } + deriving (Show, Eq) instance Functor Commented where fmap = fmapDefault @@ -130,17 +132,17 @@ -- | An entry in the @required@ section of a @futhark.pkg@ file. data Required = Required - { requiredPkg :: PkgPath - -- ^ Name of the required package. - , requiredPkgRev :: SemVer - -- ^ The minimum revision. - , requiredHash :: Maybe T.Text - -- ^ An optional hash indicating what - -- this revision looked like the last - -- time we saw it. Used for integrity - -- checking. - } - deriving (Show, Eq) + { -- | Name of the required package. + requiredPkg :: PkgPath, + -- | The minimum revision. + requiredPkgRev :: SemVer, + -- | An optional hash indicating what + -- this revision looked like the last + -- time we saw it. Used for integrity + -- checking. + requiredHash :: Maybe T.Text + } + deriving (Show, Eq) -- | The name of the file containing the futhark-pkg manifest. futharkPkg :: FilePath @@ -149,12 +151,13 @@ -- | A structure corresponding to a @futhark.pkg@ file, including -- comments. It is an invariant that duplicate required packages do -- not occcur (the parser will verify this). -data PkgManifest = PkgManifest { manifestPkgPath :: Commented (Maybe PkgPath) - -- ^ The name of the package. - , manifestRequire :: Commented [Either Comment Required] - , manifestEndComments :: [Comment] - } - deriving (Show, Eq) +data PkgManifest = PkgManifest + { -- | The name of the package. + manifestPkgPath :: Commented (Maybe PkgPath), + manifestRequire :: Commented [Either Comment Required], + manifestEndComments :: [Comment] + } + deriving (Show, Eq) -- | Possibly given a package path, construct an otherwise-empty manifest file. newPkgManifest :: Maybe PkgPath -> PkgManifest @@ -165,54 +168,70 @@ -- @futhark.pkg@ file. prettyPkgManifest :: PkgManifest -> T.Text prettyPkgManifest (PkgManifest name required endcs) = - T.unlines $ concat [ prettyComments name - , maybe [] (pure . ("package "<>) . (<>"\n")) $ commented name - , prettyComments required - , ["require {"] - , map ((" "<>) . prettyRequired) $ commented required - , ["}"] - , map prettyComment endcs - ] - where prettyComments = map prettyComment . comments - prettyComment = ("--"<>) - prettyRequired (Left c) = prettyComment c - prettyRequired (Right (Required p r h)) = - T.unwords $ catMaybes [Just p, - Just $ prettySemVer r, - ("#"<>) <$> h] + T.unlines $ + concat + [ prettyComments name, + maybe [] (pure . ("package " <>) . (<> "\n")) $ commented name, + prettyComments required, + ["require {"], + map ((" " <>) . prettyRequired) $ commented required, + ["}"], + map prettyComment endcs + ] + where + prettyComments = map prettyComment . comments + prettyComment = ("--" <>) + prettyRequired (Left c) = prettyComment c + prettyRequired (Right (Required p r h)) = + T.unwords $ + catMaybes + [ Just p, + Just $ prettySemVer r, + ("#" <>) <$> h + ] -- | The required packages listed in a package manifest. pkgRevDeps :: PkgManifest -> PkgRevDeps -pkgRevDeps = PkgRevDeps . M.fromList . mapMaybe onR . - commented . manifestRequire - where onR (Right r) = Just (requiredPkg r, (requiredPkgRev r, requiredHash r)) - onR (Left _) = Nothing +pkgRevDeps = + PkgRevDeps . M.fromList . mapMaybe onR + . commented + . manifestRequire + where + onR (Right r) = Just (requiredPkg r, (requiredPkgRev r, requiredHash r)) + onR (Left _) = Nothing -- | Where in the corresponding repository archive we can expect to -- find the package files. pkgDir :: PkgManifest -> Maybe Posix.FilePath -pkgDir = fmap (Posix.addTrailingPathSeparator . ("lib" Posix.) . - T.unpack) . commented . manifestPkgPath +pkgDir = + fmap + ( Posix.addTrailingPathSeparator . ("lib" Posix.) + . T.unpack + ) + . commented + . manifestPkgPath -- | Add new required package to the package manifest. If the package -- was already present, return the old version. addRequiredToManifest :: Required -> PkgManifest -> (PkgManifest, Maybe Required) addRequiredToManifest new_r pm = let (old, requires') = mapAccumL add Nothing $ commented $ manifestRequire pm - in (if isJust old - then pm { manifestRequire = requires' <$ manifestRequire pm } - else pm { manifestRequire = (++[Right new_r]) <$> manifestRequire pm }, - old) - where add acc (Left c) = (acc, Left c) - add acc (Right r) - | requiredPkg r == requiredPkg new_r = (Just r, Right new_r) - | otherwise = (acc, Right r) + in ( if isJust old + then pm {manifestRequire = requires' <$ manifestRequire pm} + else pm {manifestRequire = (++ [Right new_r]) <$> manifestRequire pm}, + old + ) + where + add acc (Left c) = (acc, Left c) + add acc (Right r) + | requiredPkg r == requiredPkg new_r = (Just r, Right new_r) + | otherwise = (acc, Right r) -- | Check if the manifest specifies a required package with the given -- package path. requiredInManifest :: PkgPath -> PkgManifest -> Maybe Required requiredInManifest p = - find ((==p) . requiredPkg) . rights . commented . manifestRequire + find ((== p) . requiredPkg) . rights . commented . manifestRequire -- | Remove a required package from the manifest. Returns 'Nothing' -- if the package was not found in the manifest, and otherwise the new @@ -220,9 +239,12 @@ removeRequiredFromManifest :: PkgPath -> PkgManifest -> Maybe (PkgManifest, Required) removeRequiredFromManifest p pm = do r <- requiredInManifest p pm - return (pm { manifestRequire = filter (not . matches) <$> manifestRequire pm }, - r) - where matches = either (const False) ((==p) . requiredPkg) + return + ( pm {manifestRequire = filter (not . matches) <$> manifestRequire pm}, + r + ) + where + matches = either (const False) ((== p) . requiredPkg) --- Parsing futhark.pkg. @@ -234,42 +256,50 @@ p <- optional $ lexstr "package" *> pPkgPath space c2 <- pComments - required <- (lexstr "require" *> - braces (many $ (Left <$> pComment) <|> (Right <$> pRequired))) - <|> pure [] + required <- + ( lexstr "require" + *> braces (many $ (Left <$> pComment) <|> (Right <$> pRequired)) + ) + <|> pure [] c3 <- pComments eof return $ PkgManifest (Commented c1 p) (Commented c2 required) c3 - where lexeme :: Parser a -> Parser a - lexeme p = p <* space - - lexeme' p = p <* spaceNoEol + where + lexeme :: Parser a -> Parser a + lexeme p = p <* space - lexstr :: T.Text -> Parser () - lexstr = void . try . lexeme . string + lexeme' p = p <* spaceNoEol - braces :: Parser a -> Parser a - braces p = lexstr "{" *> p <* lexstr "}" + lexstr :: T.Text -> Parser () + lexstr = void . try . lexeme . string - spaceNoEol = many $ oneOf (" \t" :: String) + braces :: Parser a -> Parser a + braces p = lexstr "{" *> p <* lexstr "}" - pPkgPath = T.pack <$> some (alphaNumChar <|> oneOf ("@-/.:" :: String)) - "package path" + spaceNoEol = many $ oneOf (" \t" :: String) - pRequired = space *> (Required <$> lexeme' pPkgPath - <*> lexeme' semver' - <*> optional (lexeme' pHash)) <* space - "package requirement" + pPkgPath = + T.pack <$> some (alphaNumChar <|> oneOf ("@-/.:" :: String)) + "package path" - pHash = char '#' *> (T.pack <$> some alphaNumChar) + pRequired = + space + *> ( Required <$> lexeme' pPkgPath + <*> lexeme' semver' + <*> optional (lexeme' pHash) + ) + <* space + "package requirement" - pComment = lexeme $ T.pack <$> (string "--" >> anySingle `manyTill` (void eol <|> eof)) + pHash = char '#' *> (T.pack <$> some alphaNumChar) - pComments :: Parser [Comment] - pComments = catMaybes <$> many (comment <|> blankLine) - where comment = Just <$> pComment - blankLine = some spaceChar >> pure Nothing + pComment = lexeme $ T.pack <$> (string "--" >> anySingle `manyTill` (void eol <|> eof)) + pComments :: Parser [Comment] + pComments = catMaybes <$> many (comment <|> blankLine) + where + comment = Just <$> pComment + blankLine = some spaceChar >> pure Nothing -- | Parse a text as a 'PkgManifest'. The 'FilePath' is used for any error messages. parsePkgManifest :: FilePath -> T.Text -> Either (ParseErrorBundle T.Text Void) PkgManifest @@ -285,11 +315,12 @@ -- | A mapping from package paths to their chosen revisions. This is -- the result of the version solver. -newtype BuildList = BuildList { unBuildList :: M.Map PkgPath SemVer } - deriving (Eq, Show) +newtype BuildList = BuildList {unBuildList :: M.Map PkgPath SemVer} + deriving (Eq, Show) -- | Prettyprint a build list; one package per line and -- newline-terminated. prettyBuildList :: BuildList -> T.Text prettyBuildList (BuildList m) = T.unlines $ map f $ sortOn fst $ M.toList m - where f (p, v) = T.unwords [p, "=>", prettySemVer v] + where + f (p, v) = T.unwords [p, "=>", prettySemVer v] diff --git a/src/Futhark/Test.hs b/src/Futhark/Test.hs --- a/src/Futhark/Test.hs +++ b/src/Futhark/Test.hs @@ -1,95 +1,97 @@ -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TupleSections #-} + -- | Facilities for reading Futhark test programs. A Futhark test -- program is an ordinary Futhark program where an initial comment -- block specifies input- and output-sets. module Futhark.Test - ( testSpecFromFile - , testSpecFromFileOrDie - , testSpecsFromPaths - , testSpecsFromPathsOrDie - , valuesFromByteString - , getValues - , getValuesBS - , compareValues - , compareValues1 - , testRunReferenceOutput - , getExpectedResult - , compileProgram - , runProgram - , ensureReferenceOutput - , determineTuning - , binaryName - , Mismatch - - , ProgramTest (..) - , StructureTest (..) - , StructurePipeline (..) - , WarningTest (..) - , TestAction (..) - , ExpectedError (..) - , InputOutputs (..) - , TestRun (..) - , ExpectedResult (..) - , Success(..) - , Values (..) - , Value - ) - where + ( testSpecFromFile, + testSpecFromFileOrDie, + testSpecsFromPaths, + testSpecsFromPathsOrDie, + valuesFromByteString, + getValues, + getValuesBS, + compareValues, + compareValues1, + testRunReferenceOutput, + getExpectedResult, + compileProgram, + runProgram, + ensureReferenceOutput, + determineTuning, + binaryName, + Mismatch, + ProgramTest (..), + StructureTest (..), + StructurePipeline (..), + WarningTest (..), + TestAction (..), + ExpectedError (..), + InputOutputs (..), + TestRun (..), + ExpectedResult (..), + Success (..), + Values (..), + Value, + ) +where +import Codec.Compression.GZip +import Codec.Compression.Zlib.Internal (DecompressError) import Control.Applicative -import qualified Data.ByteString.Lazy as BS -import qualified Data.ByteString as SBS import Control.Exception (catch) +import qualified Control.Exception.Base as E import Control.Monad import Control.Monad.Except -import qualified Data.Map.Strict as M +import qualified Data.ByteString as SBS +import qualified Data.ByteString.Lazy as BS import Data.Char import Data.Functor -import Data.Maybe import Data.List (foldl') +import qualified Data.Map.Strict as M +import Data.Maybe import qualified Data.Text as T -import qualified Data.Text.IO as T import qualified Data.Text.Encoding as T +import qualified Data.Text.IO as T import Data.Void -import System.FilePath -import Codec.Compression.GZip -import Codec.Compression.Zlib.Internal (DecompressError) -import qualified Control.Exception.Base as E - -import Text.Megaparsec hiding (many, some) -import Text.Megaparsec.Char -import Text.Regex.TDFA -import System.Directory -import System.Exit -import System.Process.ByteString (readProcessWithExitCode) -import System.IO (withFile, IOMode(..), hFileSize, hClose) -import System.IO.Error -import System.IO.Temp - -import Prelude - import Futhark.Analysis.Metrics import Futhark.IR.Primitive - (IntType(..), intValue, intByteSize, - FloatType(..), floatValue, floatByteSize) + ( FloatType (..), + IntType (..), + floatByteSize, + floatValue, + intByteSize, + intValue, + ) import Futhark.Test.Values import Futhark.Util (directoryContents, pmapIO) import Futhark.Util.Pretty (pretty, prettyText) -import Language.Futhark.Syntax (PrimType(..), PrimValue(..)) -import Language.Futhark.Prop (primValueType, primByteSize) +import Language.Futhark.Prop (primByteSize, primValueType) +import Language.Futhark.Syntax (PrimType (..), PrimValue (..)) +import System.Directory +import System.Exit +import System.FilePath +import System.IO (IOMode (..), hClose, hFileSize, withFile) +import System.IO.Error +import System.IO.Temp +import System.Process.ByteString (readProcessWithExitCode) +import Text.Megaparsec hiding (many, some) +import Text.Megaparsec.Char +import Text.Regex.TDFA +import Prelude -- | Description of a test to be carried out on a Futhark program. -- The Futhark program is stored separately. -data ProgramTest = - ProgramTest { testDescription :: - T.Text - , testTags :: - [T.Text] - , testAction :: - TestAction - } +data ProgramTest = ProgramTest + { testDescription :: + T.Text, + testTags :: + [T.Text], + testAction :: + TestAction + } deriving (Show) -- | How to test a program. @@ -99,29 +101,33 @@ deriving (Show) -- | Input and output pairs for some entry point(s). -data InputOutputs = InputOutputs { iosEntryPoint :: T.Text - , iosTestRuns :: [TestRun] } +data InputOutputs = InputOutputs + { iosEntryPoint :: T.Text, + iosTestRuns :: [TestRun] + } deriving (Show) -- | The error expected for a negative test. -data ExpectedError = AnyError - | ThisError T.Text Regex +data ExpectedError + = AnyError + | ThisError T.Text Regex instance Show ExpectedError where show AnyError = "AnyError" show (ThisError r _) = "ThisError " ++ show r -- | How a program can be transformed. -data StructurePipeline = KernelsPipeline - | SOACSPipeline - | SequentialCpuPipeline - | GpuPipeline - deriving (Show) +data StructurePipeline + = KernelsPipeline + | SOACSPipeline + | SequentialCpuPipeline + | GpuPipeline + deriving (Show) -- | A structure test specifies a compilation pipeline, as well as -- metrics for the program coming out the other end. data StructureTest = StructureTest StructurePipeline AstMetrics - deriving (Show) + deriving (Show) -- | A warning test requires that a warning matching the regular -- expression is produced. The program must also compile succesfully. @@ -132,27 +138,29 @@ -- | A condition for execution, input, and expected result. data TestRun = TestRun - { runTags :: [String] - , runInput :: Values - , runExpectedResult :: ExpectedResult Success - , runIndex :: Int - , runDescription :: String - } - deriving (Show) + { runTags :: [String], + runInput :: Values, + runExpectedResult :: ExpectedResult Success, + runIndex :: Int, + runDescription :: String + } + deriving (Show) -- | Several Values - either literally, or by reference to a file, or -- to be generated on demand. -data Values = Values [Value] - | InFile FilePath - | GenValues [GenValue] - deriving (Show) +data Values + = Values [Value] + | InFile FilePath + | GenValues [GenValue] + deriving (Show) -data GenValue = GenValue [Int] PrimType - -- ^ Generate a value of the given rank and primitive - -- type. Scalars are considered 0-ary arrays. - | GenPrim PrimValue - -- ^ A fixed non-randomised primitive value. - deriving (Show) +data GenValue + = -- | Generate a value of the given rank and primitive + -- type. Scalars are considered 0-ary arrays. + GenValue [Int] PrimType + | -- | A fixed non-randomised primitive value. + GenPrim PrimValue + deriving (Show) -- | A prettyprinted representation of type of value produced by a -- 'GenValue'. @@ -164,18 +172,21 @@ -- | How a test case is expected to terminate. data ExpectedResult values - = Succeeds (Maybe values) -- ^ Execution suceeds, with or without - -- expected result values. - | RunTimeFailure ExpectedError -- ^ Execution fails with this error. + = -- | Execution suceeds, with or without + -- expected result values. + Succeeds (Maybe values) + | -- | Execution fails with this error. + RunTimeFailure ExpectedError deriving (Show) -- | The result expected from a succesful execution. -data Success = SuccessValues Values - -- ^ These values are expected. - | SuccessGenerateValues - -- ^ Compute expected values from executing a known-good - -- reference implementation. - deriving (Show) +data Success + = -- | These values are expected. + SuccessValues Values + | -- | Compute expected values from executing a known-good + -- reference implementation. + SuccessGenerateValues + deriving (Show) type Parser = Parsec Void T.Text @@ -197,31 +208,43 @@ braces p = lexstr "{" *> p <* lexstr "}" parseNatural :: Parser Int -parseNatural = lexeme $ foldl' (\acc x -> acc * 10 + x) 0 . - map num <$> some digitChar - where num c = ord c - ord '0' +parseNatural = + lexeme $ + foldl' (\acc x -> acc * 10 + x) 0 + . map num + <$> some digitChar + where + num c = ord c - ord '0' parseDescription :: Parser T.Text parseDescription = lexeme $ T.pack <$> (anySingle `manyTill` parseDescriptionSeparator) parseDescriptionSeparator :: Parser () -parseDescriptionSeparator = try (string descriptionSeparator >> - void (satisfy isSpace `manyTill` newline)) <|> eof +parseDescriptionSeparator = + try + ( string descriptionSeparator + >> void (satisfy isSpace `manyTill` newline) + ) + <|> eof descriptionSeparator :: T.Text descriptionSeparator = "==" parseTags :: Parser [T.Text] parseTags = lexstr "tags" *> braces (many parseTag) <|> pure [] - where parseTag = T.pack <$> lexeme (some $ satisfy tagConstituent) + where + parseTag = T.pack <$> lexeme (some $ satisfy tagConstituent) tagConstituent :: Char -> Bool tagConstituent c = isAlphaNum c || c == '_' || c == '-' parseAction :: Parser TestAction -parseAction = CompileTimeFailure <$> (lexstr' "error:" *> parseExpectedError) <|> - (RunCases <$> parseInputOutputs <*> - many parseExpectedStructure <*> many parseWarning) +parseAction = + CompileTimeFailure <$> (lexstr' "error:" *> parseExpectedError) + <|> ( RunCases <$> parseInputOutputs + <*> many parseExpectedStructure + <*> many parseWarning + ) parseInputOutputs :: Parser [InputOutputs] parseInputOutputs = do @@ -231,125 +254,156 @@ parseEntryPoints :: Parser [T.Text] parseEntryPoints = (lexstr "entry:" *> many entry <* space) <|> pure ["main"] - where constituent c = not (isSpace c) && c /= '}' - entry = lexeme' $ T.pack <$> some (satisfy constituent) + where + constituent c = not (isSpace c) && c /= '}' + entry = lexeme' $ T.pack <$> some (satisfy constituent) parseRunTags :: Parser [String] parseRunTags = many parseTag - where parseTag = try $ lexeme $ do s <- some $ satisfy tagConstituent - guard $ s `notElem` ["input", "structure", "warning"] - return s + where + parseTag = try $ + lexeme $ do + s <- some $ satisfy tagConstituent + guard $ s `notElem` ["input", "structure", "warning"] + return s parseRunCases :: Parser [TestRun] -parseRunCases = parseRunCases' (0::Int) - where parseRunCases' i = (:) <$> parseRunCase i <*> parseRunCases' (i+1) - <|> pure [] - parseRunCase i = do - tags <- parseRunTags - lexstr "input" - input <- if "random" `elem` tags - then parseRandomValues - else parseValues - expr <- parseExpectedResult - return $ TestRun tags input expr i $ desc i input +parseRunCases = parseRunCases' (0 :: Int) + where + parseRunCases' i = + (:) <$> parseRunCase i <*> parseRunCases' (i + 1) + <|> pure [] + parseRunCase i = do + tags <- parseRunTags + lexstr "input" + input <- + if "random" `elem` tags + then parseRandomValues + else parseValues + expr <- parseExpectedResult + return $ TestRun tags input expr i $ desc i input - -- If the file is gzipped, we strip the 'gz' extension from - -- the dataset name. This makes it more convenient to rename - -- from 'foo.in' to 'foo.in.gz', as the reported dataset name - -- does not change (which would make comparisons to historical - -- data harder). - desc _ (InFile path) - | takeExtension path == ".gz" = dropExtension path - | otherwise = path - desc i (Values vs) = - -- Turn linebreaks into space. - "#" ++ show i ++ " (\"" ++ unwords (lines vs') ++ "\")" - where vs' = case unwords (map pretty vs) of - s | length s > 50 -> take 50 s ++ "..." - | otherwise -> s - desc _ (GenValues gens) = - unwords $ map genValueType gens + -- If the file is gzipped, we strip the 'gz' extension from + -- the dataset name. This makes it more convenient to rename + -- from 'foo.in' to 'foo.in.gz', as the reported dataset name + -- does not change (which would make comparisons to historical + -- data harder). + desc _ (InFile path) + | takeExtension path == ".gz" = dropExtension path + | otherwise = path + desc i (Values vs) = + -- Turn linebreaks into space. + "#" ++ show i ++ " (\"" ++ unwords (lines vs') ++ "\")" + where + vs' = case unwords (map pretty vs) of + s + | length s > 50 -> take 50 s ++ "..." + | otherwise -> s + desc _ (GenValues gens) = + unwords $ map genValueType gens parseExpectedResult :: Parser (ExpectedResult Success) parseExpectedResult = - (lexstr "auto" *> lexstr "output" $> Succeeds (Just SuccessGenerateValues)) <|> - (Succeeds . Just . SuccessValues <$> (lexstr "output" *> parseValues)) <|> - (RunTimeFailure <$> (lexstr "error:" *> parseExpectedError)) <|> - pure (Succeeds Nothing) + (lexstr "auto" *> lexstr "output" $> Succeeds (Just SuccessGenerateValues)) + <|> (Succeeds . Just . SuccessValues <$> (lexstr "output" *> parseValues)) + <|> (RunTimeFailure <$> (lexstr "error:" *> parseExpectedError)) + <|> pure (Succeeds Nothing) parseExpectedError :: Parser ExpectedError parseExpectedError = lexeme $ do s <- T.strip <$> restOfLine if T.null s then return AnyError - -- blankCompOpt creates a regular expression that treats - -- newlines like ordinary characters, which is what we want. - else ThisError s <$> makeRegexOptsM blankCompOpt defaultExecOpt (T.unpack s) + else -- blankCompOpt creates a regular expression that treats + -- newlines like ordinary characters, which is what we want. + ThisError s <$> makeRegexOptsM blankCompOpt defaultExecOpt (T.unpack s) parseRandomValues :: Parser Values parseRandomValues = GenValues <$> between (lexstr "{") (lexstr "}") (many parseGenValue) parseGenValue :: Parser GenValue -parseGenValue = choice [ GenValue <$> many dim <*> parsePrimType - , lexeme $ GenPrim <$> choice [i8, i16, i32, i64, - u8, u16, u32, u64, - f32, f64, - int SignedValue Int32 ""] - ] - where digits = some (satisfy isDigit) - dim = between (lexstr "[") (lexstr "]") $ - lexeme $ read <$> digits - - readint :: String -> Integer - readint = read -- To avoid warnings. - - readfloat :: String -> Double - readfloat = read -- To avoid warnings. +parseGenValue = + choice + [ GenValue <$> many dim <*> parsePrimType, + lexeme $ + GenPrim + <$> choice + [ i8, + i16, + i32, + i64, + u8, + u16, + u32, + u64, + f32, + f64, + int SignedValue Int32 "" + ] + ] + where + digits = some (satisfy isDigit) + dim = + between (lexstr "[") (lexstr "]") $ + lexeme $ read <$> digits - int f t s = try $ lexeme $ f . intValue t . readint <$> digits <* - string s <* - notFollowedBy (satisfy isAlphaNum) - i8 = int SignedValue Int8 "i8" - i16 = int SignedValue Int16 "i16" - i32 = int SignedValue Int32 "i32" - i64 = int SignedValue Int64 "i64" - u8 = int UnsignedValue Int8 "u8" - u16 = int UnsignedValue Int16 "u16" - u32 = int UnsignedValue Int32 "u32" - u64 = int UnsignedValue Int64 "u64" + readint :: String -> Integer + readint = read -- To avoid warnings. + readfloat :: String -> Double + readfloat = read -- To avoid warnings. + int f t s = + try $ + lexeme $ + f . intValue t . readint <$> digits + <* string s + <* notFollowedBy (satisfy isAlphaNum) + i8 = int SignedValue Int8 "i8" + i16 = int SignedValue Int16 "i16" + i32 = int SignedValue Int32 "i32" + i64 = int SignedValue Int64 "i64" + u8 = int UnsignedValue Int8 "u8" + u16 = int UnsignedValue Int16 "u16" + u32 = int UnsignedValue Int32 "u32" + u64 = int UnsignedValue Int64 "u64" - optSuffix s suff = do - s' <- s - ((s'<>) <$> suff) <|> pure s' + optSuffix s suff = do + s' <- s + ((s' <>) <$> suff) <|> pure s' - float f t s = try $ lexeme $ f . floatValue t . readfloat <$> - (digits `optSuffix` (char '.' *> (("."<>) <$> digits))) <* - string s <* - notFollowedBy (satisfy isAlphaNum) - f32 = float FloatValue Float32 "f32" - f64 = float FloatValue Float64 "f64" + float f t s = + try $ + lexeme $ + f . floatValue t . readfloat + <$> (digits `optSuffix` (char '.' *> (("." <>) <$> digits))) + <* string s + <* notFollowedBy (satisfy isAlphaNum) + f32 = float FloatValue Float32 "f32" + f64 = float FloatValue Float64 "f64" parsePrimType :: Parser PrimType parsePrimType = - choice [ lexstr "i8" $> Signed Int8 - , lexstr "i16" $> Signed Int16 - , lexstr "i32" $> Signed Int32 - , lexstr "i64" $> Signed Int64 - , lexstr "u8" $> Unsigned Int8 - , lexstr "u16" $> Unsigned Int16 - , lexstr "u32" $> Unsigned Int32 - , lexstr "u64" $> Unsigned Int64 - , lexstr "f32" $> FloatType Float32 - , lexstr "f64" $> FloatType Float64 - , lexstr "bool" $> Bool - ] + choice + [ lexstr "i8" $> Signed Int8, + lexstr "i16" $> Signed Int16, + lexstr "i32" $> Signed Int32, + lexstr "i64" $> Signed Int64, + lexstr "u8" $> Unsigned Int8, + lexstr "u16" $> Unsigned Int16, + lexstr "u32" $> Unsigned Int32, + lexstr "u64" $> Unsigned Int64, + lexstr "f32" $> FloatType Float32, + lexstr "f64" $> FloatType Float64, + lexstr "bool" $> Bool + ] parseValues :: Parser Values -parseValues = do s <- parseBlock - case valuesFromByteString "input block contents" $ BS.fromStrict $ T.encodeUtf8 s of - Left err -> fail err - Right vs -> return $ Values vs - <|> lexstr "@" *> lexeme (InFile . T.unpack <$> nextWord) +parseValues = + do + s <- parseBlock + case valuesFromByteString "input block contents" $ BS.fromStrict $ T.encodeUtf8 s of + Left err -> fail err + Right vs -> return $ Values vs + <|> lexstr "@" *> lexeme (InFile . T.unpack <$> nextWord) parseBlock :: Parser T.Text parseBlock = lexeme $ braces (T.pack <$> parseBlockBody 0) @@ -357,11 +411,11 @@ parseBlockBody :: Int -> Parser String parseBlockBody n = do c <- lookAhead anySingle - case (c,n) of + case (c, n) of ('}', 0) -> return mempty - ('}', _) -> (:) <$> anySingle <*> parseBlockBody (n-1) - ('{', _) -> (:) <$> anySingle <*> parseBlockBody (n+1) - _ -> (:) <$> anySingle <*> parseBlockBody n + ('}', _) -> (:) <$> anySingle <*> parseBlockBody (n -1) + ('{', _) -> (:) <$> anySingle <*> parseBlockBody (n + 1) + _ -> (:) <$> anySingle <*> parseBlockBody n restOfLine :: Parser T.Text restOfLine = T.pack <$> (anySingle `manyTill` (void newline <|> eof)) @@ -371,25 +425,31 @@ parseWarning :: Parser WarningTest parseWarning = lexstr "warning:" >> parseExpectedWarning - where parseExpectedWarning = lexeme $ do - s <- T.strip <$> restOfLine - ExpectedWarning s <$> makeRegexOptsM blankCompOpt defaultExecOpt (T.unpack s) + where + parseExpectedWarning = lexeme $ do + s <- T.strip <$> restOfLine + ExpectedWarning s <$> makeRegexOptsM blankCompOpt defaultExecOpt (T.unpack s) parseExpectedStructure :: Parser StructureTest parseExpectedStructure = - lexstr "structure" *> - (StructureTest <$> optimisePipeline <*> parseMetrics) + lexstr "structure" + *> (StructureTest <$> optimisePipeline <*> parseMetrics) optimisePipeline :: Parser StructurePipeline -optimisePipeline = lexstr "distributed" $> KernelsPipeline <|> - lexstr "gpu" $> GpuPipeline <|> - lexstr "cpu" $> SequentialCpuPipeline <|> - pure SOACSPipeline +optimisePipeline = + lexstr "distributed" $> KernelsPipeline + <|> lexstr "gpu" $> GpuPipeline + <|> lexstr "cpu" $> SequentialCpuPipeline + <|> pure SOACSPipeline parseMetrics :: Parser AstMetrics -parseMetrics = braces $ fmap (AstMetrics . M.fromList) $ many $ - (,) <$> (T.pack <$> lexeme (some (satisfy constituent))) <*> parseNatural - where constituent c = isAlpha c || c == '/' +parseMetrics = + braces $ + fmap (AstMetrics . M.fromList) $ + many $ + (,) <$> (T.pack <$> lexeme (some (satisfy constituent))) <*> parseNatural + where + constituent c = isAlpha c || c == '/' testSpec :: Parser ProgramTest testSpec = @@ -397,20 +457,24 @@ parserState :: Int -> FilePath -> s -> State s e parserState line name t = - State { stateInput = t - , stateOffset = 0 - , statePosState = PosState - { pstateInput = t - , pstateOffset = 0 - , pstateSourcePos = SourcePos - { sourceName = name - , sourceLine = mkPos line - , sourceColumn = mkPos 3 } - , pstateTabWidth = defaultTabWidth - , pstateLinePrefix = "-- "} - , stateParseErrors = [] - } - + State + { stateInput = t, + stateOffset = 0, + statePosState = + PosState + { pstateInput = t, + pstateOffset = 0, + pstateSourcePos = + SourcePos + { sourceName = name, + sourceLine = mkPos line, + sourceColumn = mkPos 3 + }, + pstateTabWidth = defaultTabWidth, + pstateLinePrefix = "-- " + }, + stateParseErrors = [] + } readTestSpec :: Int -> String -> T.Text -> Either (ParseErrorBundle T.Text Void) ProgramTest readTestSpec line name t = @@ -418,8 +482,9 @@ readInputOutputs :: Int -> String -> T.Text -> Either (ParseErrorBundle T.Text Void) [InputOutputs] readInputOutputs line name t = - snd $ runParser' (parseDescription *> space *> parseInputOutputs <* eof) $ - parserState line name t + snd $ + runParser' (parseDescription *> space *> parseInputOutputs <* eof) $ + parserState line name t commentPrefix :: T.Text commentPrefix = T.pack "--" @@ -432,54 +497,58 @@ testSpecFromFile path = do blocks_or_err <- (Right . testBlocks <$> T.readFile path) - `catch` couldNotRead + `catch` couldNotRead case blocks_or_err of Left err -> return $ Left err Right blocks -> do let (first_spec_line, first_spec, rest_specs) = - case blocks of [] -> (0, mempty, []) - (n,s):ss -> (n, s, ss) - case readTestSpec (1+first_spec_line) path first_spec of + case blocks of + [] -> (0, mempty, []) + (n, s) : ss -> (n, s, ss) + case readTestSpec (1 + first_spec_line) path first_spec of Left err -> return $ Left $ errorBundlePretty err - Right v -> return $ foldM moreCases v rest_specs - - where moreCases test (lineno, cases) = - case readInputOutputs lineno path cases of - Left err -> Left $ errorBundlePretty err - Right cases' -> - case testAction test of - RunCases old_cases structures warnings -> - Right test { testAction = RunCases (old_cases ++ cases') structures warnings } - _ -> Left "Secondary test block provided, but primary test block specifies compilation error." + Right v -> return $ foldM moreCases v rest_specs + where + moreCases test (lineno, cases) = + case readInputOutputs lineno path cases of + Left err -> Left $ errorBundlePretty err + Right cases' -> + case testAction test of + RunCases old_cases structures warnings -> + Right test {testAction = RunCases (old_cases ++ cases') structures warnings} + _ -> Left "Secondary test block provided, but primary test block specifies compilation error." -- | Like 'testSpecFromFile', but kills the process on error. testSpecFromFileOrDie :: FilePath -> IO ProgramTest testSpecFromFileOrDie prog = do spec_or_err <- testSpecFromFile prog case spec_or_err of - Left err -> do putStrLn err - exitFailure + Left err -> do + putStrLn err + exitFailure Right spec -> return spec testBlocks :: T.Text -> [(Int, T.Text)] testBlocks = mapMaybe isTestBlock . commentBlocks - where isTestBlock (n,block) - | any ((" " <> descriptionSeparator) `T.isPrefixOf`) block = - Just (n, T.unlines block) - | otherwise = - Nothing + where + isTestBlock (n, block) + | any ((" " <> descriptionSeparator) `T.isPrefixOf`) block = + Just (n, T.unlines block) + | otherwise = + Nothing commentBlocks :: T.Text -> [(Int, [T.Text])] -commentBlocks = commentBlocks' . zip [0..] . T.lines - where isComment = (commentPrefix `T.isPrefixOf`) - commentBlocks' ls = - let ls' = dropWhile (not . isComment . snd) ls - in case ls' of +commentBlocks = commentBlocks' . zip [0 ..] . T.lines + where + isComment = (commentPrefix `T.isPrefixOf`) + commentBlocks' ls = + let ls' = dropWhile (not . isComment . snd) ls + in case ls' of [] -> [] - (n,_) : _ -> + (n, _) : _ -> let (block, ls'') = span (isComment . snd) ls' block' = map (T.drop 2 . snd) block - in (n, block') : commentBlocks' ls'' + in (n, block') : commentBlocks' ls'' -- | Read test specifications from the given path, which can be a file -- or directory containing @.fut@ files and further directories. @@ -495,23 +564,27 @@ -- | Read test specifications from the given paths, which can be a -- files or directories containing @.fut@ files and further -- directories. -testSpecsFromPaths :: [FilePath] - -> IO (Either String [(FilePath, ProgramTest)]) +testSpecsFromPaths :: + [FilePath] -> + IO (Either String [(FilePath, ProgramTest)]) testSpecsFromPaths = fmap (fmap concat . sequence) . mapM testSpecsFromPath -- | Like 'testSpecsFromPaths', but kills the process on errors. -testSpecsFromPathsOrDie :: [FilePath] - -> IO [(FilePath, ProgramTest)] +testSpecsFromPathsOrDie :: + [FilePath] -> + IO [(FilePath, ProgramTest)] testSpecsFromPathsOrDie dirs = do specs_or_err <- testSpecsFromPaths dirs case specs_or_err of - Left err -> do putStrLn err - exitFailure + Left err -> do + putStrLn err + exitFailure Right specs -> return specs testPrograms :: FilePath -> IO [FilePath] testPrograms dir = filter isFut <$> directoryContents dir - where isFut = (==".fut") . takeExtension + where + isFut = (== ".fut") . takeExtension -- | Try to parse a several values from a byte string. The 'String' -- parameter is used for error messages. @@ -528,11 +601,13 @@ getValues dir v = do s <- getValuesBS dir v case valuesFromByteString file s of - Left e -> fail e + Left e -> fail e Right vs -> return vs - where file = case v of Values{} -> "" - InFile f -> f - GenValues{} -> "" + where + file = case v of + Values {} -> "" + InFile f -> f + GenValues {} -> "" -- | Extract a pretty representation of some 'Values'. In the IO -- monad because this might involve reading from a file. There is no @@ -542,16 +617,17 @@ return $ BS.fromStrict $ T.encodeUtf8 $ T.unlines $ map prettyText vs getValuesBS dir (InFile file) = case takeExtension file of - ".gz" -> liftIO $ do - s <- E.try readAndDecompress - case s of - Left e -> fail $ show file ++ ": " ++ show (e :: DecompressError) - Right s' -> return s' - - _ -> liftIO $ BS.readFile file' - where file' = dir file - readAndDecompress = do s <- BS.readFile file' - E.evaluate $ decompress s + ".gz" -> liftIO $ do + s <- E.try readAndDecompress + case s of + Left e -> fail $ show file ++ ": " ++ show (e :: DecompressError) + Right s' -> return s' + _ -> liftIO $ BS.readFile file' + where + file' = dir file + readAndDecompress = do + s <- BS.readFile file' + E.evaluate $ decompress s getValuesBS dir (GenValues gens) = mconcat <$> mapM (getGenBS dir) gens @@ -568,30 +644,37 @@ getGenBS :: MonadIO m => FilePath -> GenValue -> m BS.ByteString getGenBS dir gen = do liftIO $ createDirectoryIfMissing True $ dir "data" - exists_and_proper_size <- liftIO $ - withFile (dir file) ReadMode (fmap (== genFileSize gen) . hFileSize) - `catch` \ex -> if isDoesNotExistError ex then return False - else E.throw ex - unless exists_and_proper_size $ liftIO $ do - s <- genValues [gen] - withTempFile (dir "data") (genFileName gen) $ \tmpfile h -> do - hClose h -- We will be writing and reading this ourselves. - SBS.writeFile tmpfile s - renameFile tmpfile $ dir file + exists_and_proper_size <- + liftIO $ + withFile (dir file) ReadMode (fmap (== genFileSize gen) . hFileSize) + `catch` \ex -> + if isDoesNotExistError ex + then return False + else E.throw ex + unless exists_and_proper_size $ + liftIO $ do + s <- genValues [gen] + withTempFile (dir "data") (genFileName gen) $ \tmpfile h -> do + hClose h -- We will be writing and reading this ourselves. + SBS.writeFile tmpfile s + renameFile tmpfile $ dir file getValuesBS dir $ InFile file - where file = "data" genFileName gen + where + file = "data" genFileName gen genValues :: [GenValue] -> IO SBS.ByteString genValues gens = do - (code, stdout, stderr) <- readProcessWithExitCode "futhark" ("dataset":args) mempty + (code, stdout, stderr) <- readProcessWithExitCode "futhark" ("dataset" : args) mempty case code of ExitSuccess -> return stdout ExitFailure e -> - fail $ "'futhark dataset' failed with exit code " ++ show e ++ " and stderr:\n" ++ - map (chr . fromIntegral) (SBS.unpack stderr) - where args = "-b" : concatMap argForGen gens - argForGen g = ["-g", genValueType g] + fail $ + "'futhark dataset' failed with exit code " ++ show e ++ " and stderr:\n" + ++ map (chr . fromIntegral) (SBS.unpack stderr) + where + args = "-b" : concatMap argForGen gens + argForGen g = ["-g", genValueType g] genFileName :: GenValue -> FilePath genFileName gen = genValueType gen ++ ".in" @@ -600,49 +683,63 @@ -- whether an existing file is broken/truncated. genFileSize :: GenValue -> Integer genFileSize = genSize - where header_size = 1 + 1 + 1 + 4 -- 'b' - - genSize (GenValue ds t) = header_size + toInteger (length ds) * 8 + - product (map toInteger ds) * primSize t - genSize (GenPrim v) = - header_size + primByteSize (primValueType v) + where + header_size = 1 + 1 + 1 + 4 -- 'b' + genSize (GenValue ds t) = + header_size + toInteger (length ds) * 8 + + product (map toInteger ds) * primSize t + genSize (GenPrim v) = + header_size + primByteSize (primValueType v) - primSize (Signed it) = intByteSize it - primSize (Unsigned it) = intByteSize it - primSize (FloatType ft) = floatByteSize ft - primSize Bool = 1 + primSize (Signed it) = intByteSize it + primSize (Unsigned it) = intByteSize it + primSize (FloatType ft) = floatByteSize ft + primSize Bool = 1 -- | When/if generating a reference output file for this run, what -- should it be called? Includes the "data/" folder. testRunReferenceOutput :: FilePath -> T.Text -> TestRun -> FilePath testRunReferenceOutput prog entry tr = "data" - takeBaseName prog - <> ":" <> T.unpack entry - <> "-" <> map clean (runDescription tr) - <.> "out" - where clean '/' = '_' -- Would this ever happen? - clean ' ' = '_' - clean c = c + takeBaseName prog + <> ":" + <> T.unpack entry + <> "-" + <> map clean (runDescription tr) + <.> "out" + where + clean '/' = '_' -- Would this ever happen? + clean ' ' = '_' + clean c = c -- | Get the values corresponding to an expected result, if any. -getExpectedResult :: (MonadFail m, MonadIO m) => - FilePath -> T.Text -> TestRun - -> m (ExpectedResult [Value]) +getExpectedResult :: + (MonadFail m, MonadIO m) => + FilePath -> + T.Text -> + TestRun -> + m (ExpectedResult [Value]) getExpectedResult prog entry tr = case runExpectedResult tr of (Succeeds (Just (SuccessValues vals))) -> Succeeds . Just <$> getValues (takeDirectory prog) vals Succeeds (Just SuccessGenerateValues) -> - getExpectedResult prog entry - tr { runExpectedResult = Succeeds $ Just $ SuccessValues $ InFile $ - testRunReferenceOutput prog entry tr } + getExpectedResult + prog + entry + tr + { runExpectedResult = + Succeeds $ + Just $ + SuccessValues $ + InFile $ + testRunReferenceOutput prog entry tr + } Succeeds Nothing -> return $ Succeeds Nothing RunTimeFailure err -> return $ RunTimeFailure err - -- | The name we use for compiled programs. binaryName :: FilePath -> FilePath binaryName = dropExtension @@ -651,19 +748,24 @@ -- @program@ with the command @futhark backend extra-options...@, and -- returns stdout and stderr of the compiler. Throws an IO exception -- containing stderr if compilation fails. -compileProgram :: (MonadIO m, MonadError [T.Text] m) => - [String] -> FilePath -> String -> FilePath - -> m (SBS.ByteString, SBS.ByteString) +compileProgram :: + (MonadIO m, MonadError [T.Text] m) => + [String] -> + FilePath -> + String -> + FilePath -> + m (SBS.ByteString, SBS.ByteString) compileProgram extra_options futhark backend program = do - (futcode, stdout, stderr) <- liftIO $ readProcessWithExitCode futhark (backend:options) "" + (futcode, stdout, stderr) <- liftIO $ readProcessWithExitCode futhark (backend : options) "" case futcode of ExitFailure 127 -> throwError [progNotFound $ T.pack futhark] - ExitFailure _ -> throwError [T.decodeUtf8 stderr] - ExitSuccess -> return () + ExitFailure _ -> throwError [T.decodeUtf8 stderr] + ExitSuccess -> return () return (stdout, stderr) - where binOutputf = binaryName program - options = [program, "-o", binOutputf] ++ extra_options - progNotFound s = s <> ": command not found" + where + binOutputf = binaryName program + options = [program, "-o", binOutputf] ++ extra_options + progNotFound s = s <> ": command not found" -- | @runProgram runner extra_options prog entry input@ runs the -- Futhark program @prog@ (which must have the @.fut@ suffix), @@ -673,10 +775,14 @@ -- "interpreter" for the compiled program (e.g. @python@ when using -- the Python backends). The @extra_options@ are passed to the -- program. -runProgram :: MonadIO m => - String -> [String] - -> String -> T.Text -> Values - -> m (ExitCode, SBS.ByteString, SBS.ByteString) +runProgram :: + MonadIO m => + String -> + [String] -> + String -> + T.Text -> + Values -> + m (ExitCode, SBS.ByteString, SBS.ByteString) runProgram runner extra_options prog entry input = do let progbin = binaryName prog dir = takeDirectory prog @@ -693,43 +799,53 @@ -- | Ensure that any reference output files exist, or create them (by -- compiling the program with the reference compiler and running it on -- the input) if necessary. -ensureReferenceOutput :: (MonadIO m, MonadError [T.Text] m) => - Maybe Int -> FilePath -> String -> FilePath -> [InputOutputs] - -> m () +ensureReferenceOutput :: + (MonadIO m, MonadError [T.Text] m) => + Maybe Int -> + FilePath -> + String -> + FilePath -> + [InputOutputs] -> + m () ensureReferenceOutput concurrency futhark compiler prog ios = do missing <- filterM isReferenceMissing $ concatMap entryAndRuns ios unless (null missing) $ do void $ compileProgram [] futhark compiler prog - res <- liftIO $ flip (pmapIO concurrency) missing $ \(entry, tr) -> do - (code, stdout, stderr) <- runProgram "" ["-b"] prog entry $ runInput tr - case code of - ExitFailure e -> - return $ Left - [T.pack $ "Reference dataset generation failed with exit code " ++ - show e ++ " and stderr:\n" ++ - map (chr . fromIntegral) (SBS.unpack stderr)] - ExitSuccess -> do - let f = file (entry, tr) - liftIO $ createDirectoryIfMissing True $ takeDirectory f - SBS.writeFile f stdout - return $ Right () + res <- liftIO $ + flip (pmapIO concurrency) missing $ \(entry, tr) -> do + (code, stdout, stderr) <- runProgram "" ["-b"] prog entry $ runInput tr + case code of + ExitFailure e -> + return $ + Left + [ T.pack $ + "Reference dataset generation failed with exit code " + ++ show e + ++ " and stderr:\n" + ++ map (chr . fromIntegral) (SBS.unpack stderr) + ] + ExitSuccess -> do + let f = file (entry, tr) + liftIO $ createDirectoryIfMissing True $ takeDirectory f + SBS.writeFile f stdout + return $ Right () case sequence_ res of Left err -> throwError err Right () -> return () - - where file (entry, tr) = - takeDirectory prog testRunReferenceOutput prog entry tr + where + file (entry, tr) = + takeDirectory prog testRunReferenceOutput prog entry tr - entryAndRuns (InputOutputs entry rts) = map (entry,) rts + entryAndRuns (InputOutputs entry rts) = map (entry,) rts - isReferenceMissing (entry, tr) - | Succeeds (Just SuccessGenerateValues) <- runExpectedResult tr = - liftIO . fmap not . doesFileExist . file $ (entry, tr) - | otherwise = - return False + isReferenceMissing (entry, tr) + | Succeeds (Just SuccessGenerateValues) <- runExpectedResult tr = + liftIO . fmap not . doesFileExist . file $ (entry, tr) + | otherwise = + return False -- | Determine the --tuning options to pass to the program. The first -- argument is the extension of the tuning file, or 'Nothing' if none @@ -739,6 +855,9 @@ determineTuning (Just ext) program = do exists <- liftIO $ doesFileExist (program <.> ext) if exists - then return (["--tuning", program <.> ext], - " (using " <> takeFileName (program <.> ext) <> ")") + then + return + ( ["--tuning", program <.> ext], + " (using " <> takeFileName (program <.> ext) <> ")" + ) else return ([], mempty) 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 @@ -1,6 +1,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Strict #-} {-# LANGUAGE Trustworthy #-} + -- | This module defines an efficient value representation as well as -- parsing and comparison functions. This is because the standard -- Futhark parser is not able to cope with large values (like arrays @@ -8,46 +9,45 @@ -- here does not support tuples, so don't use those as input/output -- for your test programs. module Futhark.Test.Values - ( Value(..) - , Vector + ( Value (..), + Vector, - -- * Reading Values - , readValues + -- * Reading Values + readValues, - -- * Types of values - , ValueType(..) - , valueType + -- * Types of values + ValueType (..), + valueType, - -- * Comparing Values - , compareValues - , compareValues1 - , Mismatch - ) - where + -- * Comparing Values + compareValues, + compareValues1, + Mismatch, + ) +where import Control.Monad import Control.Monad.ST import Data.Binary -import Data.Binary.Put import Data.Binary.Get +import Data.Binary.Put import qualified Data.ByteString.Lazy.Char8 as BS -import Data.Int (Int8, Int16, Int32, Int64) -import Data.Char (isSpace, ord, chr) +import Data.Char (chr, isSpace, ord) +import Data.Int (Int16, Int32, Int64, Int8) import Data.Vector.Binary -import qualified Data.Vector.Unboxed.Mutable as UMVec -import qualified Data.Vector.Unboxed as UVec import Data.Vector.Generic (freeze) - -import qualified Language.Futhark.Syntax as F -import Language.Futhark.Pretty() -import Futhark.IR.Primitive (PrimValue) -import Language.Futhark.Parser.Lexer -import qualified Futhark.Util.Pretty as PP -import Futhark.IR.Prop.Constants (IsValue(..)) +import qualified Data.Vector.Unboxed as UVec +import qualified Data.Vector.Unboxed.Mutable as UMVec import Futhark.IR.Pretty () -import Futhark.Util.Pretty +import Futhark.IR.Primitive (PrimValue) +import Futhark.IR.Prop.Constants (IsValue (..)) import Futhark.Util (maybeHead) -import Futhark.Util.Loc (Pos(..)) +import Futhark.Util.Loc (Pos (..)) +import Futhark.Util.Pretty +import qualified Futhark.Util.Pretty as PP +import Language.Futhark.Parser.Lexer +import Language.Futhark.Pretty () +import qualified Language.Futhark.Syntax as F type STVector s = UMVec.STVector s @@ -57,21 +57,19 @@ -- | An efficiently represented Futhark value. Use 'pretty' to get a -- human-readable representation, and v'put' to obtain binary a -- representation. -data Value = Int8Value (Vector Int) (Vector Int8) - | Int16Value (Vector Int) (Vector Int16) - | Int32Value (Vector Int) (Vector Int32) - | Int64Value (Vector Int) (Vector Int64) - - | Word8Value (Vector Int) (Vector Word8) - | Word16Value (Vector Int) (Vector Word16) - | Word32Value (Vector Int) (Vector Word32) - | Word64Value (Vector Int) (Vector Word64) - - | Float32Value (Vector Int) (Vector Float) - | Float64Value (Vector Int) (Vector Double) - - | BoolValue (Vector Int) (Vector Bool) - deriving Show +data Value + = Int8Value (Vector Int) (Vector Int8) + | Int16Value (Vector Int) (Vector Int16) + | Int32Value (Vector Int) (Vector Int32) + | Int64Value (Vector Int) (Vector Int64) + | Word8Value (Vector Int) (Vector Word8) + | Word16Value (Vector Int) (Vector Word16) + | Word32Value (Vector Int) (Vector Word32) + | Word64Value (Vector Int) (Vector Word64) + | Float32Value (Vector Int) (Vector Float) + | Float64Value (Vector Int) (Vector Double) + | BoolValue (Vector Int) (Vector Bool) + deriving (Show) binaryFormatVersion :: Word8 binaryFormatVersion = 2 @@ -88,8 +86,9 @@ put (Float32Value shape vs) = putBinaryValue " f32" shape vs putFloatle put (Float64Value shape vs) = putBinaryValue " f64" shape vs putDoublele put (BoolValue shape vs) = putBinaryValue "bool" shape vs $ putInt8 . boolToInt - where boolToInt True = 1 - boolToInt False = 0 + where + boolToInt True = 1 + boolToInt False = 0 get = do first <- getInt8 @@ -121,14 +120,20 @@ " f32" -> get' (Float32Value shape') getFloatle num_elems " f64" -> get' (Float64Value shape') getDoublele num_elems "bool" -> get' (BoolValue shape') getBool num_elems - s -> fail $ "Cannot parse binary values of type " ++ show s - where getBool = (/=0) <$> getWord8 + s -> fail $ "Cannot parse binary values of type " ++ show s + where + getBool = (/= 0) <$> getWord8 - get' mk get_elem num_elems = - mk <$> genericGetVectorWith (pure num_elems) get_elem + get' mk get_elem num_elems = + mk <$> genericGetVectorWith (pure num_elems) get_elem -putBinaryValue :: UVec.Unbox a => - String -> Vector Int -> Vector a -> (a -> Put) -> Put +putBinaryValue :: + UVec.Unbox a => + String -> + Vector Int -> + Vector a -> + (a -> Put) -> + Put putBinaryValue tstr shape vs putv = do putInt8 $ fromIntegral $ ord 'b' putWord8 binaryFormatVersion @@ -138,10 +143,12 @@ mapM_ putv $ UVec.toList vs instance PP.Pretty Value where - ppr v | product (valueShape v) == 0 = - text "empty" <> - parens (dims <> ppr (valueElemType v)) - where dims = mconcat $ map (brackets . ppr) $ valueShape v + ppr v + | product (valueShape v) == 0 = + text "empty" + <> parens (dims <> ppr (valueElemType v)) + where + dims = mconcat $ map (brackets . ppr) $ valueShape v ppr (Int8Value shape vs) = pprArray (UVec.toList shape) vs ppr (Int16Value shape vs) = pprArray (UVec.toList shape) vs ppr (Int32Value shape vs) = pprArray (UVec.toList shape) vs @@ -157,20 +164,23 @@ pprArray :: (UVec.Unbox a, F.IsPrimValue a) => [Int] -> UVec.Vector a -> Doc pprArray [] vs = ppr $ F.primValue $ UVec.head vs -pprArray (d:ds) vs = - brackets $ cat $ punctuate separator $ map (pprArray ds . slice) [0..d-1] - where slice_size = product ds - slice i = UVec.slice (i*slice_size) slice_size vs - separator | null ds = comma <> space - | otherwise = comma <> line +pprArray (d : ds) vs = + brackets $ cat $ punctuate separator $ map (pprArray ds . slice) [0 .. d -1] + where + slice_size = product ds + slice i = UVec.slice (i * slice_size) slice_size vs + separator + | null ds = comma <> space + | otherwise = comma <> line -- | A representation of the simple values we represent in this module. data ValueType = ValueType [Int] F.PrimType - deriving (Show) + deriving (Show) instance PP.Pretty ValueType where ppr (ValueType ds t) = mconcat (map pprDim ds) <> ppr t - where pprDim d = brackets $ ppr d + where + pprDim d = brackets $ ppr d -- | A textual description of the type of a value. Follows Futhark -- type notation, and contains the exact dimension sizes if an array. @@ -178,17 +188,17 @@ valueType v = ValueType (valueShape v) $ valueElemType v valueElemType :: Value -> F.PrimType -valueElemType Int8Value{} = F.Signed F.Int8 -valueElemType Int16Value{} = F.Signed F.Int16 -valueElemType Int32Value{} = F.Signed F.Int32 -valueElemType Int64Value{} = F.Signed F.Int64 -valueElemType Word8Value{} = F.Unsigned F.Int8 -valueElemType Word16Value{} = F.Unsigned F.Int16 -valueElemType Word32Value{} = F.Unsigned F.Int32 -valueElemType Word64Value{} = F.Unsigned F.Int64 -valueElemType Float32Value{} = F.FloatType F.Float32 -valueElemType Float64Value{} = F.FloatType F.Float64 -valueElemType BoolValue{} = F.Bool +valueElemType Int8Value {} = F.Signed F.Int8 +valueElemType Int16Value {} = F.Signed F.Int16 +valueElemType Int32Value {} = F.Signed F.Int32 +valueElemType Int64Value {} = F.Signed F.Int64 +valueElemType Word8Value {} = F.Unsigned F.Int8 +valueElemType Word16Value {} = F.Unsigned F.Int16 +valueElemType Word32Value {} = F.Unsigned F.Int32 +valueElemType Word64Value {} = F.Unsigned F.Int64 +valueElemType Float32Value {} = F.FloatType F.Float32 +valueElemType Float64Value {} = F.FloatType F.Float64 +valueElemType BoolValue {} = F.Bool valueShape :: Value -> [Int] valueShape (Int8Value shape _) = UVec.toList shape @@ -206,10 +216,11 @@ -- The parser dropRestOfLine, dropSpaces :: BS.ByteString -> BS.ByteString -dropRestOfLine = BS.drop 1 . BS.dropWhile (/='\n') +dropRestOfLine = BS.drop 1 . BS.dropWhile (/= '\n') dropSpaces t = case BS.dropWhile isSpace t of - t' | "--" `BS.isPrefixOf` t' -> dropSpaces $ dropRestOfLine t' - | otherwise -> t' + t' + | "--" `BS.isPrefixOf` t' -> dropSpaces $ dropRestOfLine t' + | otherwise -> t' type ReadValue v = BS.ByteString -> Maybe (v, BS.ByteString) @@ -226,48 +237,57 @@ -- (Used elements, shape, elements, remaining input) type State s v = (Int, Vector Int, STVector s v, BS.ByteString) -readArrayElemsST :: UMVec.Unbox v => - Int -> Int -> ReadValue v -> State s v - -> ST s (Maybe (Int, State s v)) +readArrayElemsST :: + UMVec.Unbox v => + Int -> + Int -> + ReadValue v -> + State s v -> + ST s (Maybe (Int, State s v)) readArrayElemsST j r rv s = do ms <- readRankedArrayOfST r rv s case ms of Just (i, shape, arr, t) | Just t' <- symbol ',' t -> - readArrayElemsST (j+1) r rv (i, shape, arr, t') + readArrayElemsST (j + 1) r rv (i, shape, arr, t') | otherwise -> return $ Just (j, (i, shape, arr, t)) _ -> return $ Just (0, s) updateShape :: Int -> Int -> Vector Int -> Maybe (Vector Int) updateShape d n shape - | old_n < 0 = Just $ shape UVec.// [(r-d, n)] + | old_n < 0 = Just $ shape UVec.// [(r - d, n)] | old_n == n = Just shape - | otherwise = Nothing - where r = UVec.length shape - old_n = shape UVec.! (r-d) + | otherwise = Nothing + where + r = UVec.length shape + old_n = shape UVec.! (r - d) growIfFilled :: UVec.Unbox v => Int -> STVector s v -> ST s (STVector s v) growIfFilled i arr = if i >= capacity - then UMVec.grow arr capacity - else return arr - where capacity = UMVec.length arr + then UMVec.grow arr capacity + else return arr + where + capacity = UMVec.length arr -readRankedArrayOfST :: UMVec.Unbox v => - Int -> ReadValue v -> State s v - -> ST s (Maybe (State s v)) +readRankedArrayOfST :: + UMVec.Unbox v => + Int -> + ReadValue v -> + State s v -> + ST s (Maybe (State s v)) readRankedArrayOfST 0 rv (i, shape, arr, t) | Just (v, t') <- rv t = do - arr' <- growIfFilled i arr - UMVec.write arr' i v - return $ Just (i+1, shape, arr', t') + arr' <- growIfFilled i arr + UMVec.write arr' i v + return $ Just (i + 1, shape, arr', t') readRankedArrayOfST r rv (i, shape, arr, t) | Just t' <- symbol '[' t = do - ms <- readArrayElemsST 1 (r-1) rv (i, shape, arr, t') - return $ do - (j, s) <- ms - closeArray r j s + ms <- readArrayElemsST 1 (r -1) rv (i, shape, arr, t') + return $ do + (j, s) <- ms + closeArray r j s readRankedArrayOfST _ _ _ = return Nothing @@ -277,8 +297,12 @@ shape' <- updateShape r j shape return (i, shape', arr, t') -readRankedArrayOf :: UMVec.Unbox v => - Int -> ReadValue v -> BS.ByteString -> Maybe (Vector Int, Vector v, BS.ByteString) +readRankedArrayOf :: + UMVec.Unbox v => + Int -> + ReadValue v -> + BS.ByteString -> + Maybe (Vector Int, Vector v, BS.ByteString) readRankedArrayOf r rv t = runST $ do arr <- UMVec.new 1024 ms <- readRankedArrayOfST r rv (0, UVec.replicate r (-1), arr, t) @@ -300,120 +324,135 @@ readIntegral :: Integral int => (Token -> Maybe int) -> ReadValue int readIntegral f t = do v <- case fst <$> scanTokens (Pos "" 1 1 0) a of - Right [L _ NEGATE, L _ (INTLIT x)] -> Just $ negate $ fromIntegral x - Right [L _ (INTLIT x)] -> Just $ fromIntegral x - Right [L _ tok] -> f tok - Right [L _ NEGATE, L _ tok] -> negate <$> f tok - _ -> Nothing + Right [L _ NEGATE, L _ (INTLIT x)] -> Just $ negate $ fromIntegral x + Right [L _ (INTLIT x)] -> Just $ fromIntegral x + Right [L _ tok] -> f tok + Right [L _ NEGATE, L _ tok] -> negate <$> f tok + _ -> Nothing return (v, dropSpaces b) - where (a,b) = BS.span constituent t + where + (a, b) = BS.span constituent t readInt8 :: ReadValue Int8 readInt8 = readIntegral f - where f (I8LIT x) = Just x - f _ = Nothing + where + f (I8LIT x) = Just x + f _ = Nothing readInt16 :: ReadValue Int16 readInt16 = readIntegral f - where f (I16LIT x) = Just x - f _ = Nothing + where + f (I16LIT x) = Just x + f _ = Nothing readInt32 :: ReadValue Int32 readInt32 = readIntegral f - where f (I32LIT x) = Just x - f _ = Nothing + where + f (I32LIT x) = Just x + f _ = Nothing readInt64 :: ReadValue Int64 readInt64 = readIntegral f - where f (I64LIT x) = Just x - f _ = Nothing + where + f (I64LIT x) = Just x + f _ = Nothing readWord8 :: ReadValue Word8 readWord8 = readIntegral f - where f (U8LIT x) = Just x - f _ = Nothing + where + f (U8LIT x) = Just x + f _ = Nothing readWord16 :: ReadValue Word16 readWord16 = readIntegral f - where f (U16LIT x) = Just x - f _ = Nothing + where + f (U16LIT x) = Just x + f _ = Nothing readWord32 :: ReadValue Word32 readWord32 = readIntegral f - where f (U32LIT x) = Just x - f _ = Nothing + where + f (U32LIT x) = Just x + f _ = Nothing readWord64 :: ReadValue Word64 readWord64 = readIntegral f - where f (U64LIT x) = Just x - f _ = Nothing + where + f (U64LIT x) = Just x + f _ = Nothing readFloat :: RealFloat float => ([Token] -> Maybe float) -> ReadValue float readFloat f t = do v <- case map unLoc . fst <$> scanTokens (Pos "" 1 1 0) a of - Right [NEGATE, FLOATLIT x] -> Just $ negate $ fromDouble x - Right [FLOATLIT x] -> Just $ fromDouble x - Right (NEGATE : toks) -> negate <$> f toks - Right toks -> f toks - _ -> Nothing + Right [NEGATE, FLOATLIT x] -> Just $ negate $ fromDouble x + Right [FLOATLIT x] -> Just $ fromDouble x + Right (NEGATE : toks) -> negate <$> f toks + Right toks -> f toks + _ -> Nothing return (v, dropSpaces b) - where (a,b) = BS.span constituent t - fromDouble = uncurry encodeFloat . decodeFloat - unLoc (L _ x) = x + where + (a, b) = BS.span constituent t + fromDouble = uncurry encodeFloat . decodeFloat + unLoc (L _ x) = x readFloat32 :: ReadValue Float readFloat32 = readFloat lexFloat32 - where lexFloat32 [F32LIT x] = Just x - lexFloat32 [ID "f32", PROJ_FIELD "inf"] = Just $ 1/0 - lexFloat32 [ID "f32", PROJ_FIELD "nan"] = Just $ 0/0 - lexFloat32 _ = Nothing + where + lexFloat32 [F32LIT x] = Just x + lexFloat32 [ID "f32", PROJ_FIELD "inf"] = Just $ 1 / 0 + lexFloat32 [ID "f32", PROJ_FIELD "nan"] = Just $ 0 / 0 + lexFloat32 _ = Nothing readFloat64 :: ReadValue Double readFloat64 = readFloat lexFloat64 - where lexFloat64 [F64LIT x] = Just x - lexFloat64 [ID "f64", PROJ_FIELD "inf"] = Just $ 1/0 - lexFloat64 [ID "f64", PROJ_FIELD "nan"] = Just $ 0/0 - lexFloat64 _ = Nothing + where + lexFloat64 [F64LIT x] = Just x + lexFloat64 [ID "f64", PROJ_FIELD "inf"] = Just $ 1 / 0 + lexFloat64 [ID "f64", PROJ_FIELD "nan"] = Just $ 0 / 0 + lexFloat64 _ = Nothing readBool :: ReadValue Bool -readBool t = do v <- case fst <$> scanTokens (Pos "" 1 1 0) a of - Right [L _ TRUE] -> Just True - Right [L _ FALSE] -> Just False - _ -> Nothing - return (v, dropSpaces b) - where (a,b) = BS.span constituent t +readBool t = do + v <- case fst <$> scanTokens (Pos "" 1 1 0) a of + Right [L _ TRUE] -> Just True + Right [L _ FALSE] -> Just False + _ -> Nothing + return (v, dropSpaces b) + where + (a, b) = BS.span constituent t readPrimType :: ReadValue String readPrimType t = do pt <- case fst <$> scanTokens (Pos "" 1 1 0) a of - Right [L _ (ID s)] -> Just $ F.nameToString s - _ -> Nothing + Right [L _ (ID s)] -> Just $ F.nameToString s + _ -> Nothing return (pt, dropSpaces b) - where (a,b) = BS.span constituent t + where + (a, b) = BS.span constituent t readEmptyArrayOfShape :: [Int] -> BS.ByteString -> Maybe (Value, BS.ByteString) readEmptyArrayOfShape shape t | Just t' <- symbol '[' t, Just (d, t'') <- readIntegral (const Nothing) t', - Just t''' <- symbol ']' t'' = readEmptyArrayOfShape (shape++[d]) t''' - + Just t''' <- symbol ']' t'' = + readEmptyArrayOfShape (shape ++ [d]) t''' | otherwise = do - (pt, t') <- readPrimType t - guard $ elem 0 shape - v <- case pt of - "i8" -> Just $ Int8Value (UVec.fromList shape) UVec.empty - "i16" -> Just $ Int16Value (UVec.fromList shape) UVec.empty - "i32" -> Just $ Int32Value (UVec.fromList shape) UVec.empty - "i64" -> Just $ Int64Value (UVec.fromList shape) UVec.empty - "u8" -> Just $ Word8Value (UVec.fromList shape) UVec.empty - "u16" -> Just $ Word16Value (UVec.fromList shape) UVec.empty - "u32" -> Just $ Word32Value (UVec.fromList shape) UVec.empty - "u64" -> Just $ Word64Value (UVec.fromList shape) UVec.empty - "f32" -> Just $ Float32Value (UVec.fromList shape) UVec.empty - "f64" -> Just $ Float64Value (UVec.fromList shape) UVec.empty - "bool" -> Just $ BoolValue (UVec.fromList shape) UVec.empty - _ -> Nothing - return (v, t') + (pt, t') <- readPrimType t + guard $ elem 0 shape + v <- case pt of + "i8" -> Just $ Int8Value (UVec.fromList shape) UVec.empty + "i16" -> Just $ Int16Value (UVec.fromList shape) UVec.empty + "i32" -> Just $ Int32Value (UVec.fromList shape) UVec.empty + "i64" -> Just $ Int64Value (UVec.fromList shape) UVec.empty + "u8" -> Just $ Word8Value (UVec.fromList shape) UVec.empty + "u16" -> Just $ Word16Value (UVec.fromList shape) UVec.empty + "u32" -> Just $ Word32Value (UVec.fromList shape) UVec.empty + "u64" -> Just $ Word64Value (UVec.fromList shape) UVec.empty + "f32" -> Just $ Float32Value (UVec.fromList shape) UVec.empty + "f64" -> Just $ Float64Value (UVec.fromList shape) UVec.empty + "bool" -> Just $ BoolValue (UVec.fromList shape) UVec.empty + _ -> Nothing + return (v, t') readEmptyArray :: BS.ByteString -> Maybe (Value, BS.ByteString) readEmptyArray t = do @@ -425,54 +464,55 @@ readValue :: BS.ByteString -> Maybe (Value, BS.ByteString) readValue full_t | Right (t', _, v) <- decodeOrFail full_t = - Just (v, dropSpaces t') + Just (v, dropSpaces t') | otherwise = readEmptyArray full_t `mplus` insideBrackets 0 full_t - where insideBrackets r t = maybe (tryValueAndReadValue r t) (insideBrackets (r+1)) $ symbol '[' t - tryWith f mk r t - | Just _ <- f t = do - (shape, arr, rest_t) <- readRankedArrayOf r f full_t - return (mk shape arr, rest_t) - | otherwise = Nothing - tryValueAndReadValue r t = - -- 32-bit signed integers come first such that we parse - -- unsuffixed integer constants as of that type. - tryWith readInt32 Int32Value r t `mplus` - tryWith readInt8 Int8Value r t `mplus` - tryWith readInt16 Int16Value r t `mplus` - tryWith readInt64 Int64Value r t `mplus` - - tryWith readWord8 Word8Value r t `mplus` - tryWith readWord16 Word16Value r t `mplus` - tryWith readWord32 Word32Value r t `mplus` - tryWith readWord64 Word64Value r t `mplus` - - tryWith readFloat64 Float64Value r t `mplus` - tryWith readFloat32 Float32Value r t `mplus` - - tryWith readBool BoolValue r t + where + insideBrackets r t = maybe (tryValueAndReadValue r t) (insideBrackets (r + 1)) $ symbol '[' t + tryWith f mk r t + | Just _ <- f t = do + (shape, arr, rest_t) <- readRankedArrayOf r f full_t + return (mk shape arr, rest_t) + | otherwise = Nothing + tryValueAndReadValue r t = + -- 32-bit signed integers come first such that we parse + -- unsuffixed integer constants as of that type. + tryWith readInt32 Int32Value r t + `mplus` tryWith readInt8 Int8Value r t + `mplus` tryWith readInt16 Int16Value r t + `mplus` tryWith readInt64 Int64Value r t + `mplus` tryWith readWord8 Word8Value r t + `mplus` tryWith readWord16 Word16Value r t + `mplus` tryWith readWord32 Word32Value r t + `mplus` tryWith readWord64 Word64Value r t + `mplus` tryWith readFloat64 Float64Value r t + `mplus` tryWith readFloat32 Float32Value r t + `mplus` tryWith readBool BoolValue r t -- | Parse Futhark values from the given bytestring. readValues :: BS.ByteString -> Maybe [Value] readValues = readValues' . dropSpaces - where readValues' t - | BS.null t = Just [] - | otherwise = do (a, t') <- readValue t - (a:) <$> readValues' t' + where + readValues' t + | BS.null t = Just [] + | otherwise = do + (a, t') <- readValue t + (a :) <$> readValues' t' -- Comparisons -- | Two values differ in some way. The 'Show' instance produces a -- human-readable explanation. -data Mismatch = PrimValueMismatch (Int,Int) PrimValue PrimValue - -- ^ The position the value number and a flat index - -- into the array. - | ArrayShapeMismatch Int [Int] [Int] - | TypeMismatch Int String String - | ValueCountMismatch Int Int +data Mismatch + = -- | The position the value number and a flat index + -- into the array. + PrimValueMismatch (Int, Int) PrimValue PrimValue + | ArrayShapeMismatch Int [Int] [Int] + | TypeMismatch Int String String + | ValueCountMismatch Int Int instance Show Mismatch where - show (PrimValueMismatch (i,j) got expected) = - explainMismatch (i,j) "" got expected + show (PrimValueMismatch (i, j) got expected) = + explainMismatch (i, j) "" got expected show (ArrayShapeMismatch i got expected) = explainMismatch i "array of shape " got expected show (TypeMismatch i got expected) = @@ -490,9 +530,10 @@ compareValues :: [Value] -> [Value] -> [Mismatch] compareValues got expected | n /= m = [ValueCountMismatch n m] - | otherwise = concat $ zipWith3 compareValue [0..] got expected - where n = length got - m = length expected + | otherwise = concat $ zipWith3 compareValue [0 ..] got expected + where + n = length got + m = length expected -- | As 'compareValues', but only reports one mismatch. compareValues1 :: [Value] -> [Value] -> Maybe Mismatch @@ -527,38 +568,47 @@ _ -> [TypeMismatch i (pretty $ valueElemType got_v) (pretty $ valueElemType expected_v)] | otherwise = - [ArrayShapeMismatch i (valueShape got_v) (valueShape expected_v)] - where compareNum tol = compareGen $ compareElement tol - compareFloat tol = compareGen $ compareFloatElement tol + [ArrayShapeMismatch i (valueShape got_v) (valueShape expected_v)] + where + compareNum tol = compareGen $ compareElement tol + compareFloat tol = compareGen $ compareFloatElement tol - compareGen cmp got expected = - concat $ - zipWith cmp (UVec.toList $ UVec.indexed got) (UVec.toList expected) + compareGen cmp got expected = + concat $ + zipWith cmp (UVec.toList $ UVec.indexed got) (UVec.toList expected) - compareElement tol (j, got) expected - | comparePrimValue tol got expected = [] - | otherwise = [PrimValueMismatch (i,j) (value got) (value expected)] + compareElement tol (j, got) expected + | comparePrimValue tol got expected = [] + | otherwise = [PrimValueMismatch (i, j) (value got) (value expected)] - compareFloatElement tol (j, got) expected - | isNaN got, isNaN expected = [] - | isInfinite got, isInfinite expected, - signum got == signum expected = [] - | otherwise = compareElement tol (j, got) expected + compareFloatElement tol (j, got) expected + | isNaN got, isNaN expected = [] + | isInfinite got, + isInfinite expected, + signum got == signum expected = + [] + | otherwise = compareElement tol (j, got) expected - compareBool (j, got) expected - | got == expected = [] - | otherwise = [PrimValueMismatch (i,j) (value got) (value expected)] + compareBool (j, got) expected + | got == expected = [] + | otherwise = [PrimValueMismatch (i, j) (value got) (value expected)] -comparePrimValue :: (Ord num, Num num) => - num -> num -> num -> Bool +comparePrimValue :: + (Ord num, Num num) => + num -> + num -> + num -> + Bool comparePrimValue tol x y = diff < tol - where diff = abs $ x - y + where + diff = abs $ x - y minTolerance :: Fractional a => a minTolerance = 0.002 -- 0.2% tolerance :: (RealFloat a, UVec.Unbox a) => Vector a -> a tolerance = UVec.foldl tolerance' minTolerance . UVec.filter (not . nanOrInf) - where tolerance' t v = max t $ minTolerance * v - nanOrInf x = isInfinite x || isNaN x + where + tolerance' t v = max t $ minTolerance * v + nanOrInf x = isInfinite x || isNaN x diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs --- a/src/Futhark/Tools.hs +++ b/src/Futhark/Tools.hs @@ -1,29 +1,25 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + -- | An unstructured grab-bag of various tools and inspection -- functions that didn't really fit anywhere else. module Futhark.Tools - ( - module Futhark.Construct - - , redomapToMapAndReduce - , dissectScrema - , sequentialStreamWholeArray - - , partitionChunkedFoldParameters + ( module Futhark.Construct, + redomapToMapAndReduce, + dissectScrema, + sequentialStreamWholeArray, + partitionChunkedFoldParameters, - -- * Primitive expressions - , module Futhark.Analysis.PrimExp.Convert + -- * Primitive expressions + module Futhark.Analysis.PrimExp.Convert, ) where import Control.Monad.Identity - +import Futhark.Analysis.PrimExp.Convert +import Futhark.Construct import Futhark.IR import Futhark.IR.SOACS.SOAC -import Futhark.MonadFreshNames -import Futhark.Construct -import Futhark.Analysis.PrimExp.Convert import Futhark.Util -- | Turns a binding of a @redomap@ into two seperate bindings, a @@ -33,30 +29,42 @@ -- pattern with new 'Ident's for the result of the @map@. -- -- Only handles a pattern with an empty 'patternContextElements'. -redomapToMapAndReduce :: (MonadFreshNames m, Bindable lore, - ExpDec lore ~ (), Op lore ~ SOAC lore) => - Pattern lore - -> ( SubExp - , Commutativity - , LambdaT lore, LambdaT lore, [SubExp] - , [VName]) - -> m (Stm lore, Stm lore) -redomapToMapAndReduce (Pattern [] patelems) - (w, comm, redlam, map_lam, accs, arrs) = do - (map_pat, red_pat, red_args) <- - splitScanOrRedomap patelems w map_lam accs - let map_bnd = mkLet [] map_pat $ Op $ Screma w (mapSOAC map_lam) arrs - (nes, red_arrs) = unzip red_args - red_bnd <- Let red_pat (defAux ()) . Op <$> - (Screma w <$> reduceSOAC [Reduce comm redlam nes] <*> pure red_arrs) - return (map_bnd, red_bnd) +redomapToMapAndReduce :: + ( MonadFreshNames m, + Bindable lore, + ExpDec lore ~ (), + Op lore ~ SOAC lore + ) => + Pattern lore -> + ( SubExp, + Commutativity, + LambdaT lore, + LambdaT lore, + [SubExp], + [VName] + ) -> + m (Stm lore, Stm lore) +redomapToMapAndReduce + (Pattern [] patelems) + (w, comm, redlam, map_lam, accs, arrs) = do + (map_pat, red_pat, red_args) <- + splitScanOrRedomap patelems w map_lam accs + let map_bnd = mkLet [] map_pat $ Op $ Screma w (mapSOAC map_lam) arrs + (nes, red_arrs) = unzip red_args + red_bnd <- + Let red_pat (defAux ()) . Op + <$> (Screma w <$> reduceSOAC [Reduce comm redlam nes] <*> pure red_arrs) + return (map_bnd, red_bnd) redomapToMapAndReduce _ _ = error "redomapToMapAndReduce does not handle a non-empty 'patternContextElements'" -splitScanOrRedomap :: (Typed dec, MonadFreshNames m) => - [PatElemT dec] - -> SubExp -> LambdaT lore -> [SubExp] - -> m ([Ident], PatternT dec, [(SubExp, VName)]) +splitScanOrRedomap :: + (Typed dec, MonadFreshNames m) => + [PatElemT dec] -> + SubExp -> + LambdaT lore -> + [SubExp] -> + m ([Ident], PatternT dec, [(SubExp, VName)]) splitScanOrRedomap patelems w map_lam accs = do let (acc_patelems, arr_patelems) = splitAt (length accs) patelems (acc_ts, _arr_ts) = splitAt (length accs) $ lambdaReturnType map_lam @@ -74,10 +82,16 @@ -- Redomap. This is used to handle Scremas that are so complicated -- that we cannot directly generate efficient parallel code for them. -- In essense, what happens is the opposite of horisontal fusion. -dissectScrema :: (MonadBinder m, Op (Lore m) ~ SOAC (Lore m), - Bindable (Lore m)) => - Pattern (Lore m) -> SubExp -> ScremaForm (Lore m) -> [VName] - -> m () +dissectScrema :: + ( MonadBinder m, + Op (Lore m) ~ SOAC (Lore m), + Bindable (Lore m) + ) => + Pattern (Lore m) -> + SubExp -> + ScremaForm (Lore m) -> + [VName] -> + m () dissectScrema pat w (ScremaForm scans reds map_lam) arrs = do let num_reds = redResults reds num_scans = scanResults scans @@ -95,11 +109,14 @@ -- | Turn a stream SOAC into statements that apply the stream lambda -- to the entire input. -sequentialStreamWholeArray :: (MonadBinder m, Bindable (Lore m)) => - Pattern (Lore m) - -> SubExp -> [SubExp] - -> LambdaT (Lore m) -> [VName] - -> m () +sequentialStreamWholeArray :: + (MonadBinder m, Bindable (Lore m)) => + Pattern (Lore m) -> + SubExp -> + [SubExp] -> + LambdaT (Lore m) -> + [VName] -> + m () sequentialStreamWholeArray pat w nes lam arrs = do -- We just set the chunksize to w and inline the lambda body. There -- is no difference between parallel and sequential streams here. @@ -116,8 +133,9 @@ -- Finally, the array parameters are set to the arrays (but reshaped -- to make the types work out; this will be simplified rapidly). forM_ (zip arr_params arrs) $ \(p, arr) -> - letBindNames [paramName p] $ BasicOp $ - Reshape (map DimCoercion $ arrayDims $ paramType p) arr + letBindNames [paramName p] $ + BasicOp $ + Reshape (map DimCoercion $ arrayDims $ paramType p) arr -- Then we just inline the lambda body. mapM_ addStm $ bodyStms $ lambdaBody lam @@ -129,17 +147,19 @@ case (arrayDims $ patElemType pe, se) of (dims, Var v) | not $ null dims -> - letBindNames [patElemName pe] $ BasicOp $ Reshape (map DimCoercion dims) v + letBindNames [patElemName pe] $ BasicOp $ Reshape (map DimCoercion dims) v _ -> letBindNames [patElemName pe] $ BasicOp $ SubExp se -- | Split the parameters of a stream reduction lambda into the chunk -- size parameter, the accumulator parameters, and the input chunk -- parameters. The integer argument is how many accumulators are -- used. -partitionChunkedFoldParameters :: Int -> [Param dec] - -> (Param dec, [Param dec], [Param dec]) +partitionChunkedFoldParameters :: + Int -> + [Param dec] -> + (Param dec, [Param dec], [Param dec]) partitionChunkedFoldParameters _ [] = error "partitionChunkedFoldParameters: lambda takes no parameters" partitionChunkedFoldParameters num_accs (chunk_param : params) = let (acc_params, arr_params) = splitAt num_accs params - in (chunk_param, acc_params, arr_params) + in (chunk_param, acc_params, arr_params) diff --git a/src/Futhark/Transform/CopyPropagate.hs b/src/Futhark/Transform/CopyPropagate.hs --- a/src/Futhark/Transform/CopyPropagate.hs +++ b/src/Futhark/Transform/CopyPropagate.hs @@ -1,41 +1,45 @@ -{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} + -- | Perform copy propagation. This is done by invoking the -- simplifier with no rules, so hoisting and dead-code elimination may -- also take place. module Futhark.Transform.CopyPropagate - ( copyPropagateInProg - , copyPropagateInStms - , copyPropagateInFun - ) + ( copyPropagateInProg, + copyPropagateInStms, + copyPropagateInFun, + ) where -import Futhark.Pass -import Futhark.MonadFreshNames +import qualified Futhark.Analysis.SymbolTable as ST import Futhark.IR +import Futhark.MonadFreshNames import Futhark.Optimise.Simplify -import qualified Futhark.Analysis.SymbolTable as ST import Futhark.Optimise.Simplify.Lore (Wise) +import Futhark.Pass -- | Run copy propagation on an entire program. -copyPropagateInProg :: SimplifiableLore lore => - SimpleOps lore - -> Prog lore - -> PassM (Prog lore) +copyPropagateInProg :: + SimplifiableLore lore => + SimpleOps lore -> + Prog lore -> + PassM (Prog lore) copyPropagateInProg simpl = simplifyProg simpl mempty neverHoist -- | Run copy propagation on some statements. -copyPropagateInStms :: (MonadFreshNames m, SimplifiableLore lore) => - SimpleOps lore - -> Scope lore - -> Stms lore - -> m (ST.SymbolTable (Wise lore), Stms lore) +copyPropagateInStms :: + (MonadFreshNames m, SimplifiableLore lore) => + SimpleOps lore -> + Scope lore -> + Stms lore -> + m (ST.SymbolTable (Wise lore), Stms lore) copyPropagateInStms simpl = simplifyStms simpl mempty neverHoist -- | Run copy propagation on a function. -copyPropagateInFun :: (MonadFreshNames m, SimplifiableLore lore) => - SimpleOps lore - -> ST.SymbolTable (Wise lore) - -> FunDef lore - -> m (FunDef lore) +copyPropagateInFun :: + (MonadFreshNames m, SimplifiableLore lore) => + SimpleOps lore -> + ST.SymbolTable (Wise lore) -> + FunDef lore -> + m (FunDef lore) copyPropagateInFun simpl = simplifyFun simpl mempty neverHoist 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 @@ -1,28 +1,27 @@ {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} + -- | The code generator cannot handle the array combinators (@map@ and -- friends), so this module was written to transform them into the -- equivalent do-loops. The transformation is currently rather naive, -- and - it's certainly worth considering when we can express such -- transformations in-place. module Futhark.Transform.FirstOrderTransform - ( transformFunDef - , transformConsts - - , FirstOrderLore - , Transformer - , transformStmRecursively - , transformLambda - , transformSOAC + ( transformFunDef, + transformConsts, + FirstOrderLore, + Transformer, + transformStmRecursively, + transformLambda, + transformSOAC, ) - where +where import Control.Monad.Except import Control.Monad.State -import qualified Data.Map.Strict as M import Data.List (zip4) - +import qualified Data.Map.Strict as M import qualified Futhark.IR as AST import Futhark.IR.SOACS import Futhark.MonadFreshNames @@ -32,65 +31,84 @@ -- | The constraints that must hold for a lore in order to be the -- target of first-order transformation. type FirstOrderLore lore = - (Bindable lore, BinderOps lore, - LetDec SOACS ~ LetDec lore, - LParamInfo SOACS ~ LParamInfo lore) + ( Bindable lore, + BinderOps lore, + LetDec SOACS ~ LetDec lore, + LParamInfo SOACS ~ LParamInfo lore + ) -- | First-order-transform a single function, with the given scope -- provided by top-level constants. -transformFunDef :: (MonadFreshNames m, FirstOrderLore tolore) => - Scope tolore -> FunDef SOACS -> m (AST.FunDef tolore) +transformFunDef :: + (MonadFreshNames m, FirstOrderLore tolore) => + Scope tolore -> + FunDef SOACS -> + m (AST.FunDef tolore) transformFunDef consts_scope (FunDef entry attrs fname rettype params body) = do - (body',_) <- modifyNameSource $ runState $ runBinderT m consts_scope + (body', _) <- modifyNameSource $ runState $ runBinderT m consts_scope return $ FunDef entry attrs fname rettype params body' - where m = localScope (scopeOfFParams params) $ insertStmsM $ transformBody body + where + m = localScope (scopeOfFParams params) $ insertStmsM $ transformBody body -- | First-order-transform these top-level constants. -transformConsts :: (MonadFreshNames m, FirstOrderLore tolore) => - Stms SOACS -> m (AST.Stms tolore) +transformConsts :: + (MonadFreshNames m, FirstOrderLore tolore) => + Stms SOACS -> + m (AST.Stms tolore) transformConsts stms = fmap snd $ modifyNameSource $ runState $ runBinderT m mempty - where m = mapM_ transformStmRecursively stms + where + m = mapM_ transformStmRecursively stms -- | The constraints that a monad must uphold in order to be used for -- first-order transformation. -type Transformer m = (MonadBinder m, LocalScope (Lore m) m, - Bindable (Lore m), BinderOps (Lore m), - LParamInfo SOACS ~ LParamInfo (Lore m)) +type Transformer m = + ( MonadBinder m, + LocalScope (Lore m) m, + Bindable (Lore m), + BinderOps (Lore m), + LParamInfo SOACS ~ LParamInfo (Lore m) + ) -transformBody :: (Transformer m, LetDec (Lore m) ~ LetDec SOACS) => - Body -> m (AST.Body (Lore m)) +transformBody :: + (Transformer m, LetDec (Lore m) ~ LetDec SOACS) => + Body -> + m (AST.Body (Lore m)) transformBody (Body () bnds res) = insertStmsM $ do mapM_ transformStmRecursively bnds return $ resultBody res -- | First transform any nested t'Body' or t'Lambda' elements, then -- apply 'transformSOAC' if the expression is a SOAC. -transformStmRecursively :: (Transformer m, LetDec (Lore m) ~ LetDec SOACS) => - Stm -> m () - +transformStmRecursively :: + (Transformer m, LetDec (Lore m) ~ LetDec SOACS) => + Stm -> + m () transformStmRecursively (Let pat aux (Op soac)) = auxing aux $ transformSOAC pat =<< mapSOACM soacTransform soac - where soacTransform = identitySOACMapper { mapOnSOACLambda = transformLambda } - + where + soacTransform = identitySOACMapper {mapOnSOACLambda = transformLambda} transformStmRecursively (Let pat aux e) = auxing aux $ letBind pat =<< mapExpM transform e - where transform = identityMapper { mapOnBody = \scope -> localScope scope . transformBody - , mapOnRetType = return - , mapOnBranchType = return - , mapOnFParam = return - , mapOnLParam = return - , mapOnOp = error "Unhandled Op in first order transform" - } + where + transform = + identityMapper + { mapOnBody = \scope -> localScope scope . transformBody, + mapOnRetType = return, + mapOnBranchType = return, + mapOnFParam = return, + mapOnLParam = return, + mapOnOp = error "Unhandled Op in first order transform" + } -- | Transform a single 'SOAC' into a do-loop. The body of the lambda -- is untouched, and may or may not contain further 'SOAC's depending -- on the given lore. -transformSOAC :: Transformer m => - AST.Pattern (Lore m) - -> SOAC (Lore m) - -> m () - +transformSOAC :: + Transformer m => + AST.Pattern (Lore m) -> + SOAC (Lore m) -> + m () transformSOAC pat (Screma w form@(ScremaForm scans reds map_lam) arrs) = do -- Start by combining all the reduction parts into a single operator let Reduce _ red_lam red_nes = singleReduce reds @@ -116,54 +134,68 @@ redout_params <- mapM (newParam "redout" . flip toDecl Nonunique) $ lambdaReturnType red_lam mapout_params <- mapM (newParam "mapout" . flip toDecl Unique) map_arr_ts - let merge = concat [zip scanacc_params scan_nes, - zip scanout_params $ map Var scan_arrs, - zip redout_params red_nes, - zip mapout_params $ map Var map_arrs] + let merge = + concat + [ zip scanacc_params scan_nes, + zip scanout_params $ map Var scan_arrs, + zip redout_params red_nes, + zip mapout_params $ map Var map_arrs + ] i <- newVName "i" - let loopform = ForLoop i Int32 w [] + let loopform = ForLoop i Int64 w [] loop_body <- runBodyBinder $ - localScope (scopeOfFParams $ map fst merge) $ - inScopeOf loopform $ do - - forM_ (zip (lambdaParams map_lam) arrs) $ \(p, arr) -> do - arr_t <- lookupType arr - letBindNames [paramName p] $ BasicOp $ Index arr $ - fullSlice arr_t [DimFix $ Var i] + localScope (scopeOfFParams $ map fst merge) $ + inScopeOf loopform $ do + forM_ (zip (lambdaParams map_lam) arrs) $ \(p, arr) -> do + arr_t <- lookupType arr + letBindNames [paramName p] $ + BasicOp $ + Index arr $ + fullSlice arr_t [DimFix $ Var i] - -- Insert the statements of the lambda. We have taken care to - -- ensure that the parameters are bound at this point. - mapM_ addStm $ bodyStms $ lambdaBody map_lam - -- Split into scan results, reduce results, and map results. - let (scan_res, red_res, map_res) = - splitAt3 (length scan_nes) (length red_nes) $ - bodyResult $ lambdaBody map_lam + -- Insert the statements of the lambda. We have taken care to + -- ensure that the parameters are bound at this point. + mapM_ addStm $ bodyStms $ lambdaBody map_lam + -- Split into scan results, reduce results, and map results. + let (scan_res, red_res, map_res) = + splitAt3 (length scan_nes) (length red_nes) $ + bodyResult $ lambdaBody map_lam - scan_res' <- eLambda scan_lam $ map (pure . BasicOp . SubExp) $ - map (Var . paramName) scanacc_params ++ scan_res - red_res' <- eLambda red_lam $ map (pure . BasicOp . SubExp) $ - map (Var . paramName) redout_params ++ red_res + scan_res' <- + eLambda scan_lam $ + map (pure . BasicOp . SubExp) $ + map (Var . paramName) scanacc_params ++ scan_res + red_res' <- + eLambda red_lam $ + map (pure . BasicOp . SubExp) $ + map (Var . paramName) redout_params ++ red_res - -- Write the scan accumulator to the scan result arrays. - scan_outarrs <- letwith (map paramName scanout_params) (pexp (Var i)) $ - map (BasicOp . SubExp) scan_res' + -- Write the scan accumulator to the scan result arrays. + scan_outarrs <- + letwith (map paramName scanout_params) (pexp (Var i)) $ + map (BasicOp . SubExp) scan_res' - -- Write the map results to the map result arrays. - map_outarrs <- letwith (map paramName mapout_params) (pexp (Var i)) $ - map (BasicOp . SubExp) map_res + -- Write the map results to the map result arrays. + map_outarrs <- + letwith (map paramName mapout_params) (pexp (Var i)) $ + map (BasicOp . SubExp) map_res - return $ resultBody $ concat [scan_res', - map Var scan_outarrs, - red_res', - map Var map_outarrs] + return $ + resultBody $ + concat + [ scan_res', + map Var scan_outarrs, + red_res', + map Var map_outarrs + ] -- We need to discard the final scan accumulators, as they are not -- bound in the original pattern. - names <- (++patternNames pat) - <$> replicateM (length scanacc_params) (newVName "discard") + names <- + (++ patternNames pat) + <$> replicateM (length scanacc_params) (newVName "discard") letBindNames names $ DoLoop [] merge loopform loop_body - transformSOAC pat (Stream w stream_form lam arrs) = do -- Create a loop that repeatedly applies the lambda body to a -- chunksize of 1. Hopefully this will lead to this outer loop @@ -176,40 +208,50 @@ mapout_merge <- forM (drop (length nes) $ lambdaReturnType lam) $ \t -> let t' = t `setOuterSize` w scratch = BasicOp $ Scratch (elemType t') (arrayDims t') - in (,) - <$> newParam "stream_mapout" (toDecl t' Unique) - <*> letSubExp "stream_mapout_scratch" scratch + in (,) + <$> newParam "stream_mapout" (toDecl t' Unique) + <*> letSubExp "stream_mapout_scratch" scratch - let merge = zip (map (fmap (`toDecl` Nonunique)) fold_params) nes ++ - mapout_merge + let merge = + zip (map (fmap (`toDecl` Nonunique)) fold_params) nes + ++ mapout_merge merge_params = map fst merge mapout_params = map fst mapout_merge i <- newVName "i" - let loop_form = ForLoop i Int32 w [] + let loop_form = ForLoop i Int64 w [] letBindNames [paramName chunk_size_param] $ - BasicOp $ SubExp $ intConst Int32 1 + BasicOp $ SubExp $ intConst Int64 1 - loop_body <- runBodyBinder $ localScope (scopeOf loop_form <> - scopeOfFParams merge_params) $ do - let slice = - [DimSlice (Var i) (Var (paramName chunk_size_param)) (intConst Int32 1)] - forM_ (zip chunk_params arrs) $ \(p, arr) -> - letBindNames [paramName p] $ BasicOp $ Index arr $ - fullSlice (paramType p) slice + loop_body <- runBodyBinder $ + localScope + ( scopeOf loop_form + <> scopeOfFParams merge_params + ) + $ do + let slice = + [DimSlice (Var i) (Var (paramName chunk_size_param)) (intConst Int64 1)] + forM_ (zip chunk_params arrs) $ \(p, arr) -> + letBindNames [paramName p] $ + BasicOp $ + Index arr $ + fullSlice (paramType p) slice - (res, mapout_res) <- splitAt (length nes) <$> bodyBind (lambdaBody lam) + (res, mapout_res) <- splitAt (length nes) <$> bodyBind (lambdaBody lam) - mapout_res' <- forM (zip mapout_params mapout_res) $ \(p, se) -> - letSubExp "mapout_res" $ BasicOp $ Update (paramName p) - (fullSlice (paramType p) slice) se + mapout_res' <- forM (zip mapout_params mapout_res) $ \(p, se) -> + letSubExp "mapout_res" $ + BasicOp $ + Update + (paramName p) + (fullSlice (paramType p) slice) + se - resultBodyM $ res ++ mapout_res' + resultBodyM $ res ++ mapout_res' letBind pat $ DoLoop [] merge loop_form loop_body - transformSOAC pat (Scatter len lam ivs as) = do iter <- newVName "write_iter" @@ -222,101 +264,118 @@ -- Scatter is in-place, so we use the input array as the output array. let merge = loopMerge asOuts $ map Var as_vs loopBody <- runBodyBinder $ - localScope (M.insert iter (IndexName Int32) $ - scopeOfFParams $ map fst merge) $ do - ivs' <- forM ivs $ \iv -> do - iv_t <- lookupType iv - 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'' + localScope + ( M.insert iter (IndexName Int64) $ + scopeOfFParams $ map fst merge + ) + $ do + ivs' <- forM ivs $ \iv -> do + iv_t <- lookupType iv + letSubExp "write_iv" $ BasicOp $ Index iv $ fullSlice iv_t [DimFix $ Var iter] + ivs'' <- bindLambda lam (map (BasicOp . SubExp) 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) + let indexes = chunks as_ns $ take ivsLen ivs'' + values = chunks as_ns $ drop ivsLen ivs'' - foldM saveInArray arr $ zip indexes' values' - return $ resultBody (map Var ress) - letBind pat $ DoLoop [] merge (ForLoop iter Int32 len []) loopBody + 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) + foldM saveInArray arr $ zip indexes' values' + return $ resultBody (map Var ress) + letBind pat $ DoLoop [] merge (ForLoop iter Int64 len []) loopBody transformSOAC pat (Hist len ops bucket_fun imgs) = do iter <- newVName "iter" -- Bind arguments to parameters for the merge-variables. - hists_ts <- mapM lookupType $ concatMap histDest ops + hists_ts <- mapM lookupType $ concatMap histDest ops hists_out <- mapM (newIdent "dests") hists_ts let merge = loopMerge hists_out $ concatMap (map Var . histDest) ops -- Bind lambda-bodies for operators. loopBody <- runBodyBinder $ - localScope (M.insert iter (IndexName Int32) $ - scopeOfFParams $ map fst merge) $ do - - -- Bind images to parameters of bucket function. - imgs' <- forM imgs $ \img -> do - img_t <- lookupType img - letSubExp "pixel" $ BasicOp $ Index img $ fullSlice img_t [DimFix $ Var iter] - imgs'' <- bindLambda bucket_fun $ map (BasicOp . SubExp) imgs' + localScope + ( M.insert iter (IndexName Int64) $ + scopeOfFParams $ map fst merge + ) + $ do + -- Bind images to parameters of bucket function. + imgs' <- forM imgs $ \img -> do + img_t <- lookupType img + letSubExp "pixel" $ BasicOp $ Index img $ fullSlice img_t [DimFix $ Var iter] + imgs'' <- bindLambda bucket_fun $ map (BasicOp . SubExp) imgs' - -- Split out values from bucket function. - let lens = length ops - inds = take lens imgs'' - vals = chunks (map (length . lambdaReturnType . histOp) ops) $ drop lens imgs'' - hists_out' = chunks (map (length . lambdaReturnType . histOp) ops) $ - map identName hists_out + -- Split out values from bucket function. + let lens = length ops + inds = take lens imgs'' + vals = chunks (map (length . lambdaReturnType . histOp) ops) $ drop lens imgs'' + hists_out' = + chunks (map (length . lambdaReturnType . histOp) ops) $ + map identName hists_out - hists_out'' <- forM (zip4 hists_out' ops inds vals) $ \(hist, op, idx, val) -> do - -- Check whether the indexes are in-bound. If they are not, we - -- return the histograms unchanged. - let outside_bounds_branch = insertStmsM $ resultBodyM $ map Var hist - oob = case hist of [] -> eSubExp $ constant True - arr:_ -> eOutOfBounds arr [eSubExp idx] + hists_out'' <- forM (zip4 hists_out' ops inds vals) $ \(hist, op, idx, val) -> do + -- Check whether the indexes are in-bound. If they are not, we + -- return the histograms unchanged. + let outside_bounds_branch = insertStmsM $ resultBodyM $ map Var hist + oob = case hist of + [] -> eSubExp $ constant True + arr : _ -> eOutOfBounds arr [eSubExp idx] - letTupExp "new_histo" <=< - eIf oob outside_bounds_branch $ do - -- Read values from histogram. - h_val <- forM hist $ \arr -> do - arr_t <- lookupType arr - letSubExp "read_hist" $ BasicOp $ Index arr $ fullSlice arr_t [DimFix idx] + letTupExp "new_histo" + <=< eIf oob outside_bounds_branch + $ do + -- Read values from histogram. + h_val <- forM hist $ \arr -> do + arr_t <- lookupType arr + letSubExp "read_hist" $ BasicOp $ Index arr $ fullSlice arr_t [DimFix idx] - -- Apply operator. - h_val' <- bindLambda (histOp op) $ + -- Apply operator. + h_val' <- + bindLambda (histOp op) $ map (BasicOp . SubExp) $ h_val ++ val - -- Write values back to histograms. - hist' <- forM (zip hist h_val') $ \(arr, v) -> do - arr_t <- lookupType arr - letInPlace "hist_out" arr (fullSlice arr_t [DimFix idx]) $ - BasicOp $ SubExp v + -- Write values back to histograms. + hist' <- forM (zip hist h_val') $ \(arr, v) -> do + arr_t <- lookupType arr + letInPlace "hist_out" arr (fullSlice arr_t [DimFix idx]) $ + BasicOp $ SubExp v - return $ resultBody $ map Var hist' + return $ resultBody $ map Var hist' - return $ resultBody $ map Var $ concat hists_out'' + return $ resultBody $ map Var $ concat hists_out'' -- Wrap up the above into a for-loop. - letBind pat $ DoLoop [] merge (ForLoop iter Int32 len []) loopBody + letBind pat $ DoLoop [] merge (ForLoop iter Int64 len []) loopBody -- | Recursively first-order-transform a lambda. -transformLambda :: (MonadFreshNames m, - Bindable lore, BinderOps lore, - LocalScope somelore m, - SameScope somelore lore, - LetDec lore ~ LetDec SOACS) => - Lambda -> m (AST.Lambda lore) +transformLambda :: + ( MonadFreshNames m, + Bindable lore, + BinderOps lore, + LocalScope somelore m, + SameScope somelore lore, + LetDec lore ~ LetDec SOACS + ) => + Lambda -> + m (AST.Lambda lore) transformLambda (Lambda params body rettype) = do - body' <- runBodyBinder $ - localScope (scopeOfLParams params) $ - transformBody body + body' <- + runBodyBinder $ + localScope (scopeOfLParams params) $ + transformBody body return $ Lambda params body' rettype resultArray :: Transformer m => [Type] -> m [VName] resultArray = mapM oneArray - where oneArray t = letExp "result" $ BasicOp $ Scratch (elemType t) (arrayDims t) + where + oneArray t = letExp "result" $ BasicOp $ Scratch (elemType t) (arrayDims t) -letwith :: Transformer m => - [VName] -> m (AST.Exp (Lore m)) -> [AST.Exp (Lore m)] - -> m [VName] +letwith :: + Transformer m => + [VName] -> + m (AST.Exp (Lore m)) -> + [AST.Exp (Lore m)] -> + m [VName] letwith ks i vs = do vs' <- letSubExps "values" vs i' <- letSubExp "i" =<< i @@ -328,19 +387,23 @@ pexp :: Applicative f => SubExp -> f (AST.Exp lore) pexp = pure . BasicOp . SubExp -bindLambda :: Transformer m => - AST.Lambda (Lore m) -> [AST.Exp (Lore m)] - -> m [SubExp] +bindLambda :: + Transformer m => + AST.Lambda (Lore m) -> + [AST.Exp (Lore m)] -> + m [SubExp] bindLambda (Lambda params body _) args = do forM_ (zip params args) $ \(param, arg) -> if primType $ paramType param - then letBindNames [paramName param] arg - else letBindNames [paramName param] =<< eCopy (pure arg) + then letBindNames [paramName param] arg + else letBindNames [paramName param] =<< eCopy (pure arg) bodyBind body loopMerge :: [Ident] -> [SubExp] -> [(Param DeclType, SubExp)] loopMerge vars = loopMerge' $ zip vars $ repeat Unique -loopMerge' :: [(Ident,Uniqueness)] -> [SubExp] -> [(Param DeclType, SubExp)] -loopMerge' vars vals = [ (Param pname $ toDecl ptype u, val) - | ((Ident pname ptype, u),val) <- zip vars vals ] +loopMerge' :: [(Ident, Uniqueness)] -> [SubExp] -> [(Param DeclType, SubExp)] +loopMerge' vars vals = + [ (Param pname $ toDecl ptype u, val) + | ((Ident pname ptype, u), val) <- zip vars vals + ] diff --git a/src/Futhark/Transform/Rename.hs b/src/Futhark/Transform/Rename.hs --- a/src/Futhark/Transform/Rename.hs +++ b/src/Futhark/Transform/Rename.hs @@ -1,106 +1,128 @@ -{-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE Trustworthy #-} +{-# LANGUAGE UndecidableInstances #-} + -- | This module provides facilities for transforming Futhark programs -- such that names are unique, via the 'renameProg' function. module Futhark.Transform.Rename - ( - -- * Renaming programs - renameProg - -- * Renaming parts of a program. - -- - -- These all require execution in a 'MonadFreshNames' environment. - , renameExp - , renameStm - , renameBody - , renameLambda - , renamePattern - -- * Renaming annotations - , RenameM - , substituteRename - , renamingStms - , Rename (..) - , Renameable + ( -- * Renaming programs + renameProg, + + -- * Renaming parts of a program. + + -- + -- These all require execution in a 'MonadFreshNames' environment. + renameExp, + renameStm, + renameBody, + renameLambda, + renamePattern, + + -- * Renaming annotations + RenameM, + substituteRename, + renamingStms, + Rename (..), + Renameable, ) - where +where -import Control.Monad.State import Control.Monad.Reader +import Control.Monad.State import qualified Data.Map.Strict as M import Data.Maybe - -import Futhark.IR.Syntax -import Futhark.IR.Traversals +import Futhark.FreshNames hiding (newName) import Futhark.IR.Prop.Names import Futhark.IR.Prop.Patterns -import Futhark.FreshNames hiding (newName) -import Futhark.MonadFreshNames (MonadFreshNames(..), modifyNameSource, newName) +import Futhark.IR.Syntax +import Futhark.IR.Traversals +import Futhark.MonadFreshNames (MonadFreshNames (..), modifyNameSource, newName) import Futhark.Transform.Substitute runRenamer :: RenameM a -> VNameSource -> (a, VNameSource) runRenamer (RenameM m) src = runReader (runStateT m src) env - where env = RenameEnv M.empty + where + env = RenameEnv M.empty -- | Rename variables such that each is unique. The semantics of the -- program are unaffected, under the assumption that the program was -- correct to begin with. In particular, the renaming may make an -- invalid program valid. -renameProg :: (Renameable lore, MonadFreshNames m) => - Prog lore -> m (Prog lore) -renameProg prog = modifyNameSource $ runRenamer $ - renamingStms (progConsts prog) $ \consts -> do - funs <- mapM rename (progFuns prog) - return prog { progConsts = consts, progFuns = funs } +renameProg :: + (Renameable lore, MonadFreshNames m) => + Prog lore -> + m (Prog lore) +renameProg prog = modifyNameSource $ + runRenamer $ + renamingStms (progConsts prog) $ \consts -> do + funs <- mapM rename (progFuns prog) + return prog {progConsts = consts, progFuns = funs} -- | Rename bound variables such that each is unique. The semantics -- of the expression is unaffected, under the assumption that the -- expression was correct to begin with. Any free variables are left -- untouched. -renameExp :: (Renameable lore, MonadFreshNames m) => - Exp lore -> m (Exp lore) +renameExp :: + (Renameable lore, MonadFreshNames m) => + Exp lore -> + m (Exp lore) renameExp = modifyNameSource . runRenamer . rename -- | Rename bound variables such that each is unique. The semantics -- of the binding is unaffected, under the assumption that the -- binding was correct to begin with. Any free variables are left -- untouched, as are the names in the pattern of the binding. -renameStm :: (Renameable lore, MonadFreshNames m) => - Stm lore -> m (Stm lore) +renameStm :: + (Renameable lore, MonadFreshNames m) => + Stm lore -> + m (Stm lore) renameStm binding = do e <- renameExp $ stmExp binding - return binding { stmExp = e } + return binding {stmExp = e} -- | Rename bound variables such that each is unique. The semantics -- of the body is unaffected, under the assumption that the body was -- correct to begin with. Any free variables are left untouched. -renameBody :: (Renameable lore, MonadFreshNames m) => - Body lore -> m (Body lore) +renameBody :: + (Renameable lore, MonadFreshNames m) => + Body lore -> + m (Body lore) renameBody = modifyNameSource . runRenamer . rename -- | Rename bound variables such that each is unique. The semantics -- of the lambda is unaffected, under the assumption that the body was -- correct to begin with. Any free variables are left untouched. -- Note in particular that the parameters of the lambda are renamed. -renameLambda :: (Renameable lore, MonadFreshNames m) => - Lambda lore -> m (Lambda lore) +renameLambda :: + (Renameable lore, MonadFreshNames m) => + Lambda lore -> + m (Lambda lore) renameLambda = modifyNameSource . runRenamer . rename -- | Produce an equivalent pattern but with each pattern element given -- a new name. -renamePattern :: (Rename dec, MonadFreshNames m) => - PatternT dec -> m (PatternT dec) +renamePattern :: + (Rename dec, MonadFreshNames m) => + PatternT dec -> + m (PatternT dec) renamePattern = modifyNameSource . runRenamer . rename' - where rename' pat = bind (patternNames pat) $ rename pat + where + rename' pat = bind (patternNames pat) $ rename pat -newtype RenameEnv = RenameEnv { envNameMap :: M.Map VName VName } +newtype RenameEnv = RenameEnv {envNameMap :: M.Map VName VName} -- | The monad in which renaming is performed. newtype RenameM a = RenameM (StateT VNameSource (Reader RenameEnv) a) - deriving (Functor, Applicative, Monad, - MonadFreshNames, MonadReader RenameEnv) + deriving + ( Functor, + Applicative, + Monad, + MonadFreshNames, + MonadReader RenameEnv + ) -- | Produce a map of the substitutions that should be performed by -- the renamer. @@ -128,15 +150,15 @@ instance Rename a => Rename [a] where rename = mapM rename -instance (Rename a, Rename b) => Rename (a,b) where - rename (a,b) = (,) <$> rename a <*> rename b +instance (Rename a, Rename b) => Rename (a, b) where + rename (a, b) = (,) <$> rename a <*> rename b -instance (Rename a, Rename b, Rename c) => Rename (a,b,c) where - rename (a,b,c) = do +instance (Rename a, Rename b, Rename c) => Rename (a, b, c) where + rename (a, b, c) = do a' <- rename a b' <- rename b c' <- rename c - return (a',b',c') + return (a', b', c') instance Rename a => Rename (Maybe a) where rename = maybe (return Nothing) (fmap Just . rename) @@ -156,18 +178,24 @@ -- This works because map union prefers elements from left -- operand. local (bind' vars') body - where bind' vars' env = env { envNameMap = M.fromList (zip vars vars') - `M.union` envNameMap env } + where + bind' vars' env = + env + { envNameMap = + M.fromList (zip vars vars') + `M.union` envNameMap env + } -- | Rename some statements, then execute an action with the name -- substitutions induced by the statements active. renamingStms :: Renameable lore => Stms lore -> (Stms lore -> RenameM a) -> RenameM a renamingStms stms m = descend mempty stms - where descend stms' rem_stms = case stmsHead rem_stms of - Nothing -> m stms' - Just (stm, rem_stms') -> bind (patternNames $ stmPattern stm) $ do - stm' <- rename stm - descend (stms' <> oneStm stm') rem_stms' + where + descend stms' rem_stms = case stmsHead rem_stms of + Nothing -> m stms' + Just (stm, rem_stms') -> bind (patternNames $ stmPattern stm) $ do + stm' <- rename stm + descend (stms' <> oneStm stm') rem_stms' instance Renameable lore => Rename (FunDef lore) where rename (FunDef entry attrs fname ret params body) = @@ -178,7 +206,7 @@ return $ FunDef entry attrs fname ret' params' body' instance Rename SubExp where - rename (Var v) = Var <$> rename v + rename (Var v) = Var <$> rename v rename (Constant v) = return $ Constant v instance Rename dec => Rename (Param dec) where @@ -223,40 +251,54 @@ let (loop_params, loop_arrs) = unzip loop_vars boundexp' <- rename boundexp loop_arrs' <- rename loop_arrs - bind (map paramName (ctxparams++valparams) ++ - map paramName loop_params) $ do - ctxparams' <- mapM rename ctxparams - valparams' <- mapM rename valparams - loop_params' <- mapM rename loop_params - i' <- rename i - loopbody' <- rename loopbody - return $ DoLoop - (zip ctxparams' ctxinit') (zip valparams' valinit') - (ForLoop i' it boundexp' $ - zip loop_params' loop_arrs') loopbody' + bind + ( map paramName (ctxparams ++ valparams) + ++ map paramName loop_params + ) + $ do + ctxparams' <- mapM rename ctxparams + valparams' <- mapM rename valparams + loop_params' <- mapM rename loop_params + i' <- rename i + loopbody' <- rename loopbody + return $ + DoLoop + (zip ctxparams' ctxinit') + (zip valparams' valinit') + ( ForLoop i' it boundexp' $ + zip loop_params' loop_arrs' + ) + loopbody' WhileLoop cond -> - bind (map paramName $ ctxparams++valparams) $ do + bind (map paramName $ ctxparams ++ valparams) $ do ctxparams' <- mapM rename ctxparams valparams' <- mapM rename valparams loopbody' <- rename loopbody - cond' <- rename cond - return $ DoLoop - (zip ctxparams' ctxinit') (zip valparams' valinit') - (WhileLoop cond') loopbody' + cond' <- rename cond + return $ + DoLoop + (zip ctxparams' ctxinit') + (zip valparams' valinit') + (WhileLoop cond') + loopbody' rename e = mapExpM mapper e - where mapper = Mapper { - mapOnBody = const rename - , mapOnSubExp = rename - , mapOnVName = rename - , mapOnRetType = rename - , mapOnBranchType = rename - , mapOnFParam = rename - , mapOnLParam = rename - , mapOnOp = rename - } + where + mapper = + Mapper + { mapOnBody = const rename, + mapOnSubExp = rename, + mapOnVName = rename, + mapOnRetType = rename, + mapOnBranchType = rename, + mapOnFParam = rename, + mapOnLParam = rename, + mapOnOp = rename + } -instance Rename shape => - Rename (TypeBase shape u) where +instance + Rename shape => + Rename (TypeBase shape u) + where rename (Array et size u) = do size' <- rename size return $ Array et size' u @@ -282,21 +324,23 @@ instance Rename ExtSize where rename (Free se) = Free <$> rename se - rename (Ext x) = return $ Ext x + rename (Ext x) = return $ Ext x instance Rename () where rename = return instance Rename d => Rename (DimIndex d) where - rename (DimFix i) = DimFix <$> rename i + rename (DimFix i) = DimFix <$> rename i rename (DimSlice i n s) = DimSlice <$> rename i <*> rename n <*> rename s -- | Lores in which all annotations are renameable. -type Renameable lore = (Rename (LetDec lore), - Rename (ExpDec lore), - Rename (BodyDec lore), - Rename (FParamInfo lore), - Rename (LParamInfo lore), - Rename (RetType lore), - Rename (BranchType lore), - Rename (Op lore)) +type Renameable lore = + ( Rename (LetDec lore), + Rename (ExpDec lore), + Rename (BodyDec lore), + Rename (FParamInfo lore), + Rename (LParamInfo lore), + Rename (RetType lore), + Rename (BranchType lore), + Rename (Op lore) + ) diff --git a/src/Futhark/Transform/Substitute.hs b/src/Futhark/Transform/Substitute.hs --- a/src/Futhark/Transform/Substitute.hs +++ b/src/Futhark/Transform/Substitute.hs @@ -1,26 +1,27 @@ -{-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE Safe #-} +{-# LANGUAGE UndecidableInstances #-} + -- | -- -- This module contains facilities for replacing variable names in -- syntactic constructs. module Futhark.Transform.Substitute - (Substitutions, - Substitute(..), - Substitutable) - where + ( Substitutions, + Substitute (..), + Substitutable, + ) +where import Control.Monad.Identity import qualified Data.Map.Strict as M - -import Futhark.IR.Syntax -import Futhark.IR.Traversals -import Futhark.IR.Prop.Scope import Futhark.Analysis.PrimExp import Futhark.IR.Prop.Names +import Futhark.IR.Prop.Scope +import Futhark.IR.Syntax +import Futhark.IR.Traversals -- | The substitutions to be made are given by a mapping from names to -- names. @@ -40,22 +41,24 @@ instance Substitute (Stm lore) => Substitute (Stms lore) where substituteNames substs = fmap $ substituteNames substs -instance (Substitute a, Substitute b) => Substitute (a,b) where - substituteNames substs (x,y) = +instance (Substitute a, Substitute b) => Substitute (a, b) where + substituteNames substs (x, y) = (substituteNames substs x, substituteNames substs y) -instance (Substitute a, Substitute b, Substitute c) => Substitute (a,b,c) where - substituteNames substs (x,y,z) = - (substituteNames substs x, - substituteNames substs y, - substituteNames substs z) +instance (Substitute a, Substitute b, Substitute c) => Substitute (a, b, c) where + substituteNames substs (x, y, z) = + ( substituteNames substs x, + substituteNames substs y, + substituteNames substs z + ) -instance (Substitute a, Substitute b, Substitute c, Substitute d) => Substitute (a,b,c,d) where - substituteNames substs (x,y,z,u) = - (substituteNames substs x, - substituteNames substs y, - substituteNames substs z, - substituteNames substs u) +instance (Substitute a, Substitute b, Substitute c, Substitute d) => Substitute (a, b, c, d) where + substituteNames substs (x, y, z, u) = + ( substituteNames substs x, + substituteNames substs y, + substituteNames substs z, + substituteNames substs u + ) instance Substitute a => Substitute (Maybe a) where substituteNames substs = fmap $ substituteNames substs @@ -83,15 +86,15 @@ instance Substitute dec => Substitute (StmAux dec) where substituteNames substs (StmAux cs attrs dec) = StmAux - (substituteNames substs cs) - (substituteNames substs attrs) - (substituteNames substs dec) + (substituteNames substs cs) + (substituteNames substs attrs) + (substituteNames substs dec) instance Substitute dec => Substitute (Param dec) where substituteNames substs (Param name dec) = Param - (substituteNames substs name) - (substituteNames substs dec) + (substituteNames substs name) + (substituteNames substs dec) instance Substitute dec => Substitute (PatternT dec) where substituteNames substs (Pattern context values) = @@ -104,28 +107,29 @@ instance Substitutable lore => Substitute (Stm lore) where substituteNames substs (Let pat annot e) = Let - (substituteNames substs pat) - (substituteNames substs annot) - (substituteNames substs e) + (substituteNames substs pat) + (substituteNames substs annot) + (substituteNames substs e) instance Substitutable lore => Substitute (Body lore) where substituteNames substs (Body dec stms res) = Body - (substituteNames substs dec) - (substituteNames substs stms) - (substituteNames substs res) + (substituteNames substs dec) + (substituteNames substs stms) + (substituteNames substs res) replace :: Substitutable lore => M.Map VName VName -> Mapper lore lore Identity -replace substs = Mapper { - mapOnVName = return . substituteNames substs - , mapOnSubExp = return . substituteNames substs - , mapOnBody = const $ return . substituteNames substs - , mapOnRetType = return . substituteNames substs - , mapOnBranchType = return . substituteNames substs - , mapOnFParam = return . substituteNames substs - , mapOnLParam = return . substituteNames substs - , mapOnOp = return . substituteNames substs - } +replace substs = + Mapper + { mapOnVName = return . substituteNames substs, + mapOnSubExp = return . substituteNames substs, + mapOnBody = const $ return . substituteNames substs, + mapOnRetType = return . substituteNames substs, + mapOnBranchType = return . substituteNames substs, + mapOnFParam = return . substituteNames substs, + mapOnLParam = return . substituteNames substs, + mapOnOp = return . substituteNames substs + } instance Substitute Rank where substituteNames _ = id @@ -139,7 +143,7 @@ instance Substitute d => Substitute (Ext d) where substituteNames substs (Free x) = Free $ substituteNames substs x - substituteNames _ (Ext x) = Ext x + substituteNames _ (Ext x) = Ext x instance Substitute Names where substituteNames = mapNames . substituteNames @@ -154,14 +158,15 @@ instance Substitutable lore => Substitute (Lambda lore) where substituteNames substs (Lambda params body rettype) = Lambda - (substituteNames substs params) - (substituteNames substs body) - (map (substituteNames substs) rettype) + (substituteNames substs params) + (substituteNames substs body) + (map (substituteNames substs) rettype) instance Substitute Ident where substituteNames substs v = - v { identName = substituteNames substs $ identName v - , identType = substituteNames substs $ identType v + v + { identName = substituteNames substs $ identName v, + identType = substituteNames substs $ identType v } instance Substitute d => Substitute (DimChange d) where @@ -173,6 +178,10 @@ instance Substitute v => Substitute (PrimExp v) where substituteNames substs = fmap $ substituteNames substs +instance Substitute v => Substitute (TPrimExp t v) where + substituteNames substs = + TPrimExp . fmap (substituteNames substs) . untyped + instance Substitutable lore => Substitute (NameInfo lore) where substituteNames subst (LetName dec) = LetName $ substituteNames subst dec @@ -188,12 +197,14 @@ -- | Lores in which all annotations support name -- substitution. -type Substitutable lore = (Decorations lore, - Substitute (ExpDec lore), - Substitute (BodyDec lore), - Substitute (LetDec lore), - Substitute (FParamInfo lore), - Substitute (LParamInfo lore), - Substitute (RetType lore), - Substitute (BranchType lore), - Substitute (Op lore)) +type Substitutable lore = + ( Decorations lore, + Substitute (ExpDec lore), + Substitute (BodyDec lore), + Substitute (LetDec lore), + Substitute (FParamInfo lore), + Substitute (LParamInfo lore), + Substitute (RetType lore), + Substitute (BranchType lore), + Substitute (Op lore) + ) diff --git a/src/Futhark/TypeCheck.hs b/src/Futhark/TypeCheck.hs --- a/src/Futhark/TypeCheck.hs +++ b/src/Futhark/TypeCheck.hs @@ -1,1142 +1,1389 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, TypeFamilies, ScopedTypeVariables #-} -{-# LANGUAGE DefaultSignatures #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE Strict #-} -{-# LANGUAGE Trustworthy #-} --- | The type checker checks whether the program is type-consistent. -module Futhark.TypeCheck - ( -- * Interface - checkProg - , TypeError (..) - , ErrorCase (..) - - -- * Extensionality - , TypeM - , bad - , context - , message - , Checkable (..) - , CheckableOp (..) - , lookupVar - , lookupAliases - , checkOpWith - - -- * Checkers - , require - , requireI - , requirePrimExp - , checkSubExp - , checkExp - , checkStms - , checkStm - , checkType - , checkExtType - , matchExtPattern - , matchExtBranchType - , argType - , argAliases - , noArgAliases - , checkArg - , checkSOACArrayArgs - , checkLambda - , checkBody - , consume - , consumeOnlyParams - , binding - ) - where - -import Control.Parallel.Strategies -import Control.Monad.Reader -import Control.Monad.Writer -import Control.Monad.State -import Control.Monad.RWS.Strict -import Data.List (find, intercalate, sort) -import qualified Data.Map.Strict as M -import qualified Data.Set as S -import Data.Maybe - -import Futhark.Analysis.PrimExp -import Futhark.Construct (instantiateShapes) -import Futhark.IR.Aliases -import Futhark.Util -import Futhark.Util.Pretty (Pretty, prettyDoc, indent, ppr, text, (<+>), align) - --- | Information about an error during type checking. The 'Show' --- instance for this type produces a human-readable description. -data ErrorCase lore = - TypeError String - | UnexpectedType (Exp lore) Type [Type] - | ReturnTypeError Name [ExtType] [ExtType] - | DupDefinitionError Name - | DupParamError Name VName - | DupPatternError VName - | InvalidPatternError (Pattern (Aliases lore)) [ExtType] (Maybe String) - | UnknownVariableError VName - | UnknownFunctionError Name - | ParameterMismatch (Maybe Name) [Type] [Type] - | SlicingError Int Int - | BadAnnotation String Type Type - | ReturnAliased Name VName - | UniqueReturnAliased Name - | NotAnArray VName Type - | PermutationError [Int] Int (Maybe VName) - -instance Checkable lore => Show (ErrorCase lore) where - show (TypeError msg) = - "Type error:\n" ++ msg - show (UnexpectedType e _ []) = - "Type of expression\n" ++ - prettyDoc 160 (indent 2 $ ppr e) ++ - "\ncannot have any type - possibly a bug in the type checker." - show (UnexpectedType e t ts) = - "Type of expression\n" ++ - prettyDoc 160 (indent 2 $ ppr e) ++ - "\nmust be one of " ++ intercalate ", " (map pretty ts) ++ ", but is " ++ - pretty t ++ "." - show (ReturnTypeError fname rettype bodytype) = - "Declaration of function " ++ nameToString fname ++ - " declares return type\n " ++ prettyTuple rettype ++ - "\nBut body has type\n " ++ prettyTuple bodytype - show (DupDefinitionError name) = - "Duplicate definition of function " ++ nameToString name ++ "" - show (DupParamError funname paramname) = - "Parameter " ++ pretty paramname ++ - " mentioned multiple times in argument list of function " ++ - nameToString funname ++ "." - show (DupPatternError name) = - "Variable " ++ pretty name ++ " bound twice in pattern." - show (InvalidPatternError pat t desc) = - "Pattern " ++ pretty pat ++ - " cannot match value of type " ++ prettyTuple t ++ end - where end = case desc of Nothing -> "." - Just desc' -> ":\n" ++ desc' - show (UnknownVariableError name) = - "Use of unknown variable " ++ pretty name ++ "." - show (UnknownFunctionError fname) = - "Call of unknown function " ++ nameToString fname ++ "." - show (ParameterMismatch fname expected got) = - "In call of " ++ fname' ++ ":\n" ++ - "expecting " ++ show nexpected ++ " arguments of type(s)\n" ++ - intercalate ", " (map pretty expected) ++ - "\nGot " ++ show ngot ++ - " arguments of types\n" ++ - intercalate ", " (map pretty got) - where nexpected = length expected - ngot = length got - fname' = maybe "anonymous function" (("function "++) . nameToString) fname - show (SlicingError dims got) = - show got ++ " indices given, but type of indexee has " ++ show dims ++ " dimension(s)." - show (BadAnnotation desc expected got) = - "Annotation of \"" ++ desc ++ "\" type of expression is " ++ pretty expected ++ - ", but derived to be " ++ pretty got ++ "." - show (ReturnAliased fname name) = - "Unique return value of function " ++ nameToString fname ++ - " is aliased to " ++ pretty name ++ ", which is not consumed." - show (UniqueReturnAliased fname) = - "A unique tuple element of return value of function " ++ - nameToString fname ++ " is aliased to some other tuple component." - show (NotAnArray e t) = - "The expression " ++ pretty e ++ - " is expected to be an array, but is " ++ pretty t ++ "." - show (PermutationError perm rank name) = - "The permutation (" ++ intercalate ", " (map show perm) ++ - ") is not valid for array " ++ name' ++ "of rank " ++ show rank ++ "." - where name' = maybe "" ((++" ") . pretty) name - --- | A type error. -data TypeError lore = Error [String] (ErrorCase lore) - -instance Checkable lore => Show (TypeError lore) where - show (Error [] err) = - show err - show (Error msgs err) = - intercalate "\n" msgs ++ "\n" ++ show err - --- | A tuple of a return type and a list of parameters, possibly --- named. -type FunBinding lore = ([RetType (Aliases lore)], [FParam (Aliases lore)]) - -type VarBinding lore = NameInfo (Aliases lore) - -data Usage = Consumed - | Observed - deriving (Eq, Ord, Show) - -data Occurence = Occurence { observed :: Names - , consumed :: Names - } - deriving (Eq, Show) - -observation :: Names -> Occurence -observation = flip Occurence mempty - -consumption :: Names -> Occurence -consumption = Occurence mempty - -nullOccurence :: Occurence -> Bool -nullOccurence occ = observed occ == mempty && consumed occ == mempty - -type Occurences = [Occurence] - -allConsumed :: Occurences -> Names -allConsumed = mconcat . map consumed - -seqOccurences :: Occurences -> Occurences -> Occurences -seqOccurences occurs1 occurs2 = - filter (not . nullOccurence) (map filt occurs1) ++ occurs2 - where filt occ = - occ { observed = observed occ `namesSubtract` postcons } - postcons = allConsumed occurs2 - -altOccurences :: Occurences -> Occurences -> Occurences -altOccurences occurs1 occurs2 = - filter (not . nullOccurence) (map filt occurs1) ++ occurs2 - where filt occ = - occ { consumed = consumed occ `namesSubtract` postcons - , observed = observed occ `namesSubtract` postcons } - postcons = allConsumed occurs2 - -unOccur :: Names -> Occurences -> Occurences -unOccur to_be_removed = filter (not . nullOccurence) . map unOccur' - where unOccur' occ = - occ { observed = observed occ `namesSubtract` to_be_removed - , consumed = consumed occ `namesSubtract` to_be_removed - } - --- | The 'Consumption' data structure is used to keep track of which --- variables have been consumed, as well as whether a violation has been detected. -data Consumption = ConsumptionError String - | Consumption Occurences - deriving (Show) - -instance Semigroup Consumption where - ConsumptionError e <> _ = ConsumptionError e - _ <> ConsumptionError e = ConsumptionError e - Consumption o1 <> Consumption o2 - | v:_ <- namesToList $ consumed_in_o1 `namesIntersection` used_in_o2 = - ConsumptionError $ "Variable " <> pretty v <> " referenced after being consumed." - | otherwise = - Consumption $ o1 `seqOccurences` o2 - where consumed_in_o1 = mconcat $ map consumed o1 - used_in_o2 = mconcat $ map consumed o2 <> map observed o2 - -instance Monoid Consumption where - mempty = Consumption mempty - --- | The environment contains a variable table and a function table. --- Type checking happens with access to this environment. The --- function table is only initialised at the very beginning, but the --- variable table will be extended during type-checking when --- let-expressions are encountered. -data Env lore = - Env { envVtable :: M.Map VName (VarBinding lore) - , envFtable :: M.Map Name (FunBinding lore) - , envCheckOp :: OpWithAliases (Op lore) -> TypeM lore () - , envContext :: [String] - } - --- | The type checker runs in this monad. -newtype TypeM lore a = TypeM (RWST - (Env lore) -- Reader - Consumption -- Writer - Names -- State - (Either (TypeError lore)) -- Inner monad - a) - deriving (Monad, Functor, Applicative, - MonadReader (Env lore), - MonadWriter Consumption, - MonadState Names) - -instance Checkable lore => - HasScope (Aliases lore) (TypeM lore) where - lookupType = fmap typeOf . lookupVar - askScope = asks $ M.fromList . mapMaybe varType . M.toList . envVtable - where varType (name, dec) = Just (name, dec) - -runTypeM :: Env lore -> TypeM lore a - -> Either (TypeError lore) (a, Consumption) -runTypeM env (TypeM m) = evalRWST m env mempty - -bad :: ErrorCase lore -> TypeM lore a -bad e = do - messages <- asks envContext - TypeM $ lift $ Left $ Error (reverse messages) e - --- | Add information about what is being type-checked to the current --- context. Liberal use of this combinator makes it easier to track --- type errors, as the strings are added to type errors signalled via --- 'bad'. -context :: String - -> TypeM lore a - -> TypeM lore a -context s = local $ \env -> env { envContext = s : envContext env} - -message :: Pretty a => - String -> a -> String -message s x = prettyDoc 80 $ - text s <+> align (ppr x) - --- | Mark a name as bound. If the name has been bound previously in --- the program, report a type error. -bound :: VName -> TypeM lore () -bound name = do already_seen <- gets $ nameIn name - when already_seen $ - bad $ TypeError $ "Name " ++ pretty name ++ " bound twice" - modify (<>oneName name) - -occur :: Occurences -> TypeM lore () -occur = tell . Consumption . filter (not . nullOccurence) - --- | Proclaim that we have made read-only use of the given variable. --- No-op unless the variable is array-typed. -observe :: Checkable lore => - VName -> TypeM lore () -observe name = do - dec <- lookupVar name - unless (primType $ typeOf dec) $ - occur [observation $ oneName name <> aliases dec] - --- | Proclaim that we have written to the given variables. -consume :: Checkable lore => Names -> TypeM lore () -consume als = do - scope <- askScope - let isArray = maybe False ((>0) . arrayRank . typeOf) . (`M.lookup` scope) - occur [consumption $ namesFromList $ filter isArray $ namesToList als] - -collectOccurences :: TypeM lore a -> TypeM lore (a, Occurences) -collectOccurences m = pass $ do - (x, c) <- listen m - o <- checkConsumption c - return ((x, o), const mempty) - -checkOpWith :: (OpWithAliases (Op lore) -> TypeM lore ()) - -> TypeM lore a -> TypeM lore a -checkOpWith checker = local $ \env -> env { envCheckOp = checker } - -checkConsumption :: Consumption -> TypeM lore Occurences -checkConsumption (ConsumptionError e) = bad $ TypeError e -checkConsumption (Consumption os) = return os - -alternative :: TypeM lore a -> TypeM lore b -> TypeM lore (a,b) -alternative m1 m2 = pass $ do - (x, c1) <- listen m1 - (y, c2) <- listen m2 - os1 <- checkConsumption c1 - os2 <- checkConsumption c2 - let usage = Consumption $ os1 `altOccurences` os2 - return ((x, y), const usage) - --- | Permit consumption of only the specified names. If one of these --- names is consumed, the consumption will be rewritten to be a --- consumption of the corresponding alias set. Consumption of --- anything else will result in a type error. -consumeOnlyParams :: [(VName, Names)] -> TypeM lore a -> TypeM lore a -consumeOnlyParams consumable m = do - (x, os) <- collectOccurences m - tell . Consumption =<< mapM inspect os - return x - where inspect o = do - new_consumed <- mconcat <$> mapM wasConsumed (namesToList $ consumed o) - return o { consumed = new_consumed } - wasConsumed v - | Just als <- lookup v consumable = return als - | otherwise = - bad $ TypeError $ - unlines [pretty v ++ " was invalidly consumed.", - what ++ " can be consumed here."] - what | null consumable = "Nothing" - | otherwise = "Only " ++ intercalate ", " (map (pretty . fst) consumable) - --- | Given the immediate aliases, compute the full transitive alias --- set (including the immediate aliases). -expandAliases :: Names -> Env lore -> Names -expandAliases names env = names <> aliasesOfAliases - where aliasesOfAliases = mconcat . map look . namesToList $ names - look k = case M.lookup k $ envVtable env of - Just (LetName (als, _)) -> unAliases als - _ -> mempty - -binding :: Checkable lore => - Scope (Aliases lore) - -> TypeM lore a - -> TypeM lore a -binding bnds = check . local (`bindVars` bnds) - where bindVars = M.foldlWithKey' bindVar - boundnames = M.keys bnds - - bindVar env name (LetName (AliasDec als, dec)) = - let als' | primType (typeOf dec) = mempty - | otherwise = expandAliases als env - in env { envVtable = - M.insert name (LetName (AliasDec als', dec)) $ envVtable env - } - bindVar env name dec = - env { envVtable = M.insert name dec $ envVtable env } - - -- Check whether the bound variables have been used correctly - -- within their scope. - check m = do - mapM_ bound $ M.keys bnds - (a, os) <- collectOccurences m - tell $ Consumption $ unOccur (namesFromList boundnames) os - return a - -lookupVar :: VName -> TypeM lore (NameInfo (Aliases lore)) -lookupVar name = do - bnd <- asks $ M.lookup name . envVtable - case bnd of - Nothing -> bad $ UnknownVariableError name - Just dec -> return dec - -lookupAliases :: Checkable lore => VName -> TypeM lore Names -lookupAliases name = do - info <- lookupVar name - return $ if primType $ typeOf info - then mempty - else oneName name <> aliases info - -aliases :: NameInfo (Aliases lore) -> Names -aliases (LetName (als, _)) = unAliases als -aliases _ = mempty - -subExpAliasesM :: Checkable lore => SubExp -> TypeM lore Names -subExpAliasesM Constant{} = return mempty -subExpAliasesM (Var v) = lookupAliases v - -lookupFun :: Checkable lore => - Name - -> [SubExp] - -> TypeM lore ([RetType lore], [DeclType]) -lookupFun fname args = do - bnd <- asks $ M.lookup fname . envFtable - case bnd of - Nothing -> bad $ UnknownFunctionError fname - Just (ftype, params) -> do - argts <- mapM subExpType args - case applyRetType ftype params $ zip args argts of - Nothing -> - bad $ ParameterMismatch (Just fname) (map paramType params) argts - Just rt -> - return (rt, map paramDeclType params) - --- | @checkAnnotation loc s t1 t2@ checks if @t2@ is equal to --- @t1@. If not, a 'BadAnnotation' is raised. -checkAnnotation :: String -> Type -> Type - -> TypeM lore () -checkAnnotation desc t1 t2 - | t2 == t1 = return () - | otherwise = bad $ BadAnnotation desc t1 t2 - --- | @require ts se@ causes a '(TypeError vn)' if the type of @se@ is --- not a subtype of one of the types in @ts@. -require :: Checkable lore => [Type] -> SubExp -> TypeM lore () -require ts se = do - t <- checkSubExp se - unless (t `elem` ts) $ - bad $ UnexpectedType (BasicOp $ SubExp se) t ts - --- | Variant of 'require' working on variable names. -requireI :: Checkable lore => [Type] -> VName -> TypeM lore () -requireI ts ident = require ts $ Var ident - -checkArrIdent :: Checkable lore => - VName -> TypeM lore Type -checkArrIdent v = do - t <- lookupType v - case t of - Array{} -> return t - _ -> bad $ NotAnArray v t - --- | Type check a program containing arbitrary type information, --- yielding either a type error or a program with complete type --- information. -checkProg :: Checkable lore => - Prog (Aliases lore) -> Either (TypeError lore) () -checkProg (Prog consts funs) = do - let typeenv = Env { envVtable = M.empty - , envFtable = mempty - , envContext = [] - , envCheckOp = checkOp - } - let onFunction ftable vtable fun = - fmap fst $ runTypeM typeenv $ - local (\env -> env { envFtable = ftable, envVtable = vtable }) $ - checkFun fun - (ftable, _) <- runTypeM typeenv buildFtable - (vtable, _) <- runTypeM typeenv { envFtable = ftable } $ - checkStms consts $ asks envVtable - sequence_ $ parMap rpar (onFunction ftable vtable) funs - where - buildFtable = do table <- initialFtable - foldM expand table funs - expand ftable (FunDef _ _ name ret params _) - | M.member name ftable = - bad $ DupDefinitionError name - | otherwise = - return $ M.insert name (ret,params) ftable - -initialFtable :: Checkable lore => - TypeM lore (M.Map Name (FunBinding lore)) -initialFtable = fmap M.fromList $ mapM addBuiltin $ M.toList builtInFunctions - where addBuiltin (fname, (t, ts)) = do - ps <- mapM (primFParam name) ts - return (fname, ([primRetType t], ps)) - name = VName (nameFromString "x") 0 - -checkFun :: Checkable lore => - FunDef (Aliases lore) -> TypeM lore () -checkFun (FunDef _ _ fname rettype params body) = - context ("In function " ++ nameToString fname) $ - checkFun' (fname, - map declExtTypeOf rettype, - funParamsToNameInfos params) consumable $ do - checkFunParams params - checkRetType rettype - context "When checking function body" $ checkFunBody rettype body - where consumable = [ (paramName param, mempty) - | param <- params - , unique $ paramDeclType param - ] - -funParamsToNameInfos :: [FParam lore] - -> [(VName, NameInfo (Aliases lore))] -funParamsToNameInfos = map nameTypeAndLore - where nameTypeAndLore fparam = (paramName fparam, - FParamName $ paramDec fparam) - -checkFunParams :: Checkable lore => - [FParam lore] -> TypeM lore () -checkFunParams = mapM_ $ \param -> - context ("In function parameter " ++ pretty param) $ - checkFParamLore (paramName param) (paramDec param) - -checkLambdaParams :: Checkable lore => - [LParam lore] -> TypeM lore () -checkLambdaParams = mapM_ $ \param -> - context ("In lambda parameter " ++ pretty param) $ - checkLParamLore (paramName param) (paramDec param) - -checkFun' :: Checkable lore => - (Name, - [DeclExtType], - [(VName, NameInfo (Aliases lore))]) - -> [(VName, Names)] - -> TypeM lore [Names] - -> TypeM lore () -checkFun' (fname, rettype, params) consumable check = do - checkNoDuplicateParams - binding (M.fromList params) $ - consumeOnlyParams consumable $ do - body_aliases <- check - scope <- askScope - let isArray = maybe False ((>0) . arrayRank . typeOf) . (`M.lookup` scope) - context ("When checking the body aliases: " ++ - pretty (map namesToList body_aliases)) $ - checkReturnAlias $ map (namesFromList . filter isArray . namesToList) body_aliases - where param_names = map fst params - - checkNoDuplicateParams = foldM_ expand [] param_names - - expand seen pname - | Just _ <- find (==pname) seen = - bad $ DupParamError fname pname - | otherwise = - return $ pname : seen - - -- | Check that unique return values do not alias a - -- non-consumed parameter. - checkReturnAlias = - foldM_ checkReturnAlias' mempty . returnAliasing rettype - - checkReturnAlias' seen (Unique, names) - | any (`S.member` S.map fst seen) $ namesToList names = - bad $ UniqueReturnAliased fname - | otherwise = do - consume names - return $ seen <> tag Unique names - checkReturnAlias' seen (Nonunique, names) - | any (`S.member` seen) $ tag Unique names = - bad $ UniqueReturnAliased fname - | otherwise = return $ seen <> tag Nonunique names - - tag u = S.fromList . map (,u) . namesToList - - returnAliasing expected got = - reverse $ - zip (reverse (map uniqueness expected) ++ repeat Nonunique) $ - reverse got - -checkSubExp :: Checkable lore => SubExp -> TypeM lore Type -checkSubExp (Constant val) = - return $ Prim $ primValueType val -checkSubExp (Var ident) = context ("In subexp " ++ pretty ident) $ do - observe ident - lookupType ident - -checkStms :: Checkable lore => - Stms (Aliases lore) -> TypeM lore a - -> TypeM lore a -checkStms origbnds m = delve $ stmsToList origbnds - where delve (stm@(Let pat _ e):bnds) = do - context ("In expression of statement " ++ pretty pat) $ - checkExp e - checkStm stm $ - delve bnds - delve [] = - m - -checkResult :: Checkable lore => - Result -> TypeM lore () -checkResult = mapM_ checkSubExp - -checkFunBody :: Checkable lore => - [RetType lore] - -> Body (Aliases lore) - -> TypeM lore [Names] -checkFunBody rt (Body (_,lore) bnds res) = do - checkBodyLore lore - checkStms bnds $ do - context "When checking body result" $ checkResult res - context "When matching declared return type to result of body" $ - matchReturnType rt res - map (`namesSubtract` bound_here) <$> mapM subExpAliasesM res - where bound_here = namesFromList $ M.keys $ scopeOf bnds - -checkLambdaBody :: Checkable lore => - [Type] -> Body (Aliases lore) -> TypeM lore [Names] -checkLambdaBody ret (Body (_,lore) bnds res) = do - checkBodyLore lore - checkStms bnds $ do - checkLambdaResult ret res - map (`namesSubtract` bound_here) <$> mapM subExpAliasesM res - where bound_here = namesFromList $ M.keys $ scopeOf bnds - -checkLambdaResult :: Checkable lore => - [Type] -> Result -> TypeM lore () -checkLambdaResult ts es - | length ts /= length es = - bad $ TypeError $ - "Lambda has return type " ++ prettyTuple ts ++ - " describing " ++ show (length ts) ++ " values, but body returns " ++ - show (length es) ++ " values: " ++ prettyTuple es - | otherwise = forM_ (zip ts es) $ \(t, e) -> do - et <- checkSubExp e - unless (et == t) $ - bad $ TypeError $ - "Subexpression " ++ pretty e ++ " has type " ++ pretty et ++ - " but expected " ++ pretty t - -checkBody :: Checkable lore => - Body (Aliases lore) -> TypeM lore [Names] -checkBody (Body (_,lore) bnds res) = do - checkBodyLore lore - checkStms bnds $ do - checkResult res - map (`namesSubtract` bound_here) <$> mapM subExpAliasesM res - where bound_here = namesFromList $ M.keys $ scopeOf bnds - -checkBasicOp :: Checkable lore => BasicOp -> TypeM lore () - -checkBasicOp (SubExp es) = - void $ checkSubExp es - -checkBasicOp (Opaque es) = - void $ checkSubExp es - -checkBasicOp (ArrayLit [] _) = - return () - -checkBasicOp (ArrayLit (e:es') t) = do - let check elemt eleme = do - elemet <- checkSubExp eleme - unless (elemet == elemt) $ - bad $ TypeError $ pretty elemet ++ - " is not of expected type " ++ pretty elemt ++ "." - et <- checkSubExp e - - -- Compare that type with the one given for the array literal. - checkAnnotation "array-element" t et - - mapM_ (check et) es' - -checkBasicOp (UnOp op e) = require [Prim $ unOpType op] e - -checkBasicOp (BinOp op e1 e2) = checkBinOpArgs (binOpType op) e1 e2 - -checkBasicOp (CmpOp op e1 e2) = checkCmpOp op e1 e2 - -checkBasicOp (ConvOp op e) = require [Prim $ fst $ convOpType op] e - -checkBasicOp (Index ident idxes) = do - vt <- lookupType ident - observe ident - when (arrayRank vt /= length idxes) $ - bad $ SlicingError (arrayRank vt) (length idxes) - mapM_ checkDimIndex idxes - -checkBasicOp (Update src idxes se) = do - src_t <- checkArrIdent src - when (arrayRank src_t /= length idxes) $ - bad $ SlicingError (arrayRank src_t) (length idxes) - - se_aliases <- subExpAliasesM se - when (src `nameIn` se_aliases) $ - bad $ TypeError "The target of an Update must not alias the value to be written." - - mapM_ checkDimIndex idxes - require [Prim (elemType src_t) `arrayOfShape` Shape (sliceDims idxes)] se - consume =<< lookupAliases src - -checkBasicOp (Iota e x s et) = do - require [Prim int32] e - require [Prim $ IntType et] x - require [Prim $ IntType et] s - -checkBasicOp (Replicate (Shape dims) valexp) = do - mapM_ (require [Prim int32]) dims - void $ checkSubExp valexp - -checkBasicOp (Scratch _ shape) = - mapM_ checkSubExp shape - -checkBasicOp (Reshape newshape arrexp) = do - rank <- arrayRank <$> checkArrIdent arrexp - mapM_ (require [Prim int32] . newDim) newshape - zipWithM_ (checkDimChange rank) newshape [0..] - where checkDimChange _ (DimNew _) _ = - return () - checkDimChange rank (DimCoercion se) i - | i >= rank = - bad $ TypeError $ - "Asked to coerce dimension " ++ show i ++ " to " ++ pretty se ++ - ", but array " ++ pretty arrexp ++ " has only " ++ pretty rank ++ " dimensions" - | otherwise = - return () - -checkBasicOp (Rearrange perm arr) = do - arrt <- lookupType arr - let rank = arrayRank arrt - when (length perm /= rank || sort perm /= [0..rank-1]) $ - bad $ PermutationError perm rank $ Just arr - -checkBasicOp (Rotate rots arr) = do - arrt <- lookupType arr - let rank = arrayRank arrt - mapM_ (require [Prim int32]) rots - when (length rots /= rank) $ - bad $ TypeError $ "Cannot rotate " ++ show (length rots) ++ - " dimensions of " ++ show rank ++ "-dimensional array." - -checkBasicOp (Concat i arr1exp arr2exps ressize) = do - arr1t <- checkArrIdent arr1exp - arr2ts <- mapM checkArrIdent arr2exps - let success = all ((== dropAt i 1 (arrayDims arr1t)). - dropAt i 1 . arrayDims) arr2ts - unless success $ - bad $ TypeError $ - "Types of arguments to concat do not match. Got " ++ - pretty arr1t ++ " and " ++ intercalate ", " (map pretty arr2ts) - require [Prim int32] ressize - -checkBasicOp (Copy e) = - void $ checkArrIdent e - -checkBasicOp (Manifest perm arr) = - checkBasicOp $ Rearrange perm arr -- Basically same thing! - -checkBasicOp (Assert e _ _) = - require [Prim Bool] e - -matchLoopResultExt :: Checkable lore => - [Param DeclType] -> [Param DeclType] - -> [SubExp] -> TypeM lore () -matchLoopResultExt ctx val loopres = do - let rettype_ext = - existentialiseExtTypes (map paramName ctx) $ - staticShapes $ map typeOf $ ctx ++ val - - bodyt <- mapM subExpType loopres - - case instantiateShapes (`maybeNth` loopres) rettype_ext of - Nothing -> bad $ ReturnTypeError (nameFromString "") - rettype_ext (staticShapes bodyt) - Just rettype' -> - unless (bodyt `subtypesOf` rettype') $ - bad $ ReturnTypeError (nameFromString "") - (staticShapes rettype') (staticShapes bodyt) - -checkExp :: Checkable lore => - Exp (Aliases lore) -> TypeM lore () - -checkExp (BasicOp op) = checkBasicOp op - -checkExp (If e1 e2 e3 info) = do - require [Prim Bool] e1 - _ <- checkBody e2 `alternative` checkBody e3 - context "in true branch" $ matchBranchType (ifReturns info) e2 - context "in false branch" $ matchBranchType (ifReturns info) e3 - -checkExp (Apply fname args rettype_annot _) = do - (rettype_derived, paramtypes) <- lookupFun fname $ map fst args - argflows <- mapM (checkArg . fst) args - when (rettype_derived /= rettype_annot) $ - bad $ TypeError $ "Expected apply result type " ++ pretty rettype_derived - ++ " but annotation is " ++ pretty rettype_annot - checkFuncall (Just fname) paramtypes argflows - -checkExp (DoLoop ctxmerge valmerge form loopbody) = do - let merge = ctxmerge ++ valmerge - (mergepat, mergeexps) = unzip merge - mergeargs <- mapM checkArg mergeexps - - let val_free = freeIn $ map fst valmerge - usedInVal p = paramName p `nameIn` val_free - case find (not . usedInVal . fst) ctxmerge of - Just p -> - bad $ TypeError $ "Loop context parameter " ++ pretty p ++ " unused." - Nothing -> - return () - - binding (scopeOf form) $ do - case form of - ForLoop loopvar it boundexp loopvars -> do - iparam <- primFParam loopvar $ IntType it - let funparams = iparam : mergepat - paramts = map paramDeclType funparams - - forM_ loopvars $ \(p,a) -> do - a_t <- lookupType a - observe a - case peelArray 1 a_t of - Just a_t_r -> do - checkLParamLore (paramName p) $ paramDec p - unless (a_t_r `subtypeOf` typeOf (paramDec p)) $ - bad $ TypeError $ "Loop parameter " ++ pretty p ++ - " not valid for element of " ++ pretty a ++ ", which has row type " ++ pretty a_t_r - _ -> bad $ TypeError $ "Cannot loop over " ++ pretty a ++ - " of type " ++ pretty a_t - - boundarg <- checkArg boundexp - checkFuncall Nothing paramts $ boundarg : mergeargs - - WhileLoop cond -> do - case find ((==cond) . paramName . fst) merge of - Just (condparam,_) -> - unless (paramType condparam == Prim Bool) $ - bad $ TypeError $ - "Conditional '" ++ pretty cond ++ "' of while-loop is not boolean, but " ++ - pretty (paramType condparam) ++ "." - Nothing -> - bad $ TypeError $ - "Conditional '" ++ pretty cond ++ "' of while-loop is not a merge variable." - let funparams = mergepat - paramts = map paramDeclType funparams - checkFuncall Nothing paramts mergeargs - - let rettype = map paramDeclType mergepat - consumable = [ (paramName param, mempty) - | param <- mergepat, - unique $ paramDeclType param - ] - - context "Inside the loop body" $ - checkFun' (nameFromString "", - staticShapes rettype, - funParamsToNameInfos mergepat) consumable $ do - checkFunParams mergepat - checkBodyLore $ snd $ bodyDec loopbody - - checkStms (bodyStms loopbody) $ do - checkResult $ bodyResult loopbody - - context "When matching result of body with loop parameters" $ - matchLoopResult (map fst ctxmerge) (map fst valmerge) $ - bodyResult loopbody - - let bound_here = namesFromList $ M.keys $ - scopeOf $ bodyStms loopbody - map (`namesSubtract` bound_here) <$> - mapM subExpAliasesM (bodyResult loopbody) - -checkExp (Op op) = do checker <- asks envCheckOp - checker op - -checkSOACArrayArgs :: Checkable lore => - SubExp -> [VName] -> TypeM lore [Arg] -checkSOACArrayArgs width vs = - forM vs $ \v -> do - (vt, v') <- checkSOACArrayArg v - let argSize = arraySize 0 vt - unless (argSize == width) $ - bad $ TypeError $ - "SOAC argument " ++ pretty v ++ " has outer size " ++ - pretty argSize ++ ", but width of SOAC is " ++ - pretty width - return v' - where checkSOACArrayArg ident = do - (t, als) <- checkArg $ Var ident - case peelArray 1 t of - Nothing -> bad $ TypeError $ - "SOAC argument " ++ pretty ident ++ " is not an array" - Just rt -> return (t, (rt, als)) - -checkType :: Checkable lore => - TypeBase Shape u -> TypeM lore () -checkType (Mem (ScalarSpace d _)) = mapM_ (require [Prim int32]) d -checkType t = mapM_ checkSubExp $ arrayDims t - -checkExtType :: Checkable lore => - TypeBase ExtShape u - -> TypeM lore () -checkExtType = mapM_ checkExtDim . shapeDims . arrayShape - where checkExtDim (Free se) = void $ checkSubExp se - checkExtDim (Ext _) = return () - -checkCmpOp :: Checkable lore => - CmpOp -> SubExp -> SubExp - -> TypeM lore () -checkCmpOp (CmpEq t) x y = do - require [Prim t] x - require [Prim t] y -checkCmpOp (CmpUlt t) x y = checkBinOpArgs (IntType t) x y -checkCmpOp (CmpUle t) x y = checkBinOpArgs (IntType t) x y -checkCmpOp (CmpSlt t) x y = checkBinOpArgs (IntType t) x y -checkCmpOp (CmpSle t) x y = checkBinOpArgs (IntType t) x y -checkCmpOp (FCmpLt t) x y = checkBinOpArgs (FloatType t) x y -checkCmpOp (FCmpLe t) x y = checkBinOpArgs (FloatType t) x y -checkCmpOp CmpLlt x y = checkBinOpArgs Bool x y -checkCmpOp CmpLle x y = checkBinOpArgs Bool x y - -checkBinOpArgs :: Checkable lore => - PrimType -> SubExp -> SubExp -> TypeM lore () -checkBinOpArgs t e1 e2 = do - require [Prim t] e1 - require [Prim t] e2 - -checkPatElem :: Checkable lore => - PatElemT (LetDec lore) -> TypeM lore () -checkPatElem (PatElem name dec) = context ("When checking pattern element " ++ pretty name) $ - checkLetBoundLore name dec - -checkDimIndex :: Checkable lore => - DimIndex SubExp -> TypeM lore () -checkDimIndex (DimFix i) = require [Prim int32] i -checkDimIndex (DimSlice i n s) = mapM_ (require [Prim int32]) [i,n,s] - -checkStm :: Checkable lore => - Stm (Aliases lore) - -> TypeM lore a - -> TypeM lore a -checkStm stm@(Let pat (StmAux (Certificates cs) _ (_,dec)) e) m = do - context "When checking certificates" $ mapM_ (requireI [Prim Cert]) cs - context "When checking expression annotation" $ checkExpLore dec - context ("When matching\n" ++ message " " pat ++ "\nwith\n" ++ message " " e) $ - matchPattern pat e - binding (maybeWithoutAliases $ scopeOf stm) $ do - mapM_ checkPatElem (patternElements $ removePatternAliases pat) - m - where - -- FIXME: this is wrong. However, the core language type system - -- is not strong enough to fully capture the aliases we want (see - -- issue #803). Since we eventually inline everything anyway, and - -- our intra-procedural alias analysis is much simpler and - -- correct, I could not justify spending time on improving the - -- inter-procedural alias analysis. If we ever stop inlining - -- everything, probably we need to go back and refine this. - maybeWithoutAliases = - case stmExp stm of - Apply{} -> M.map withoutAliases - _ -> id - withoutAliases (LetName (_, ldec)) = LetName (mempty, ldec) - withoutAliases info = info - -matchExtPattern :: Checkable lore => - Pattern (Aliases lore) -> [ExtType] -> TypeM lore () -matchExtPattern pat ts = - unless (expExtTypesFromPattern pat == ts) $ - bad $ InvalidPatternError pat ts Nothing - -matchExtReturnType :: Checkable lore => - [ExtType] -> Result -> TypeM lore () -matchExtReturnType rettype res = do - ts <- mapM subExpType res - matchExtReturns rettype res ts - -matchExtBranchType :: Checkable lore => - [ExtType] -> Body (Aliases lore) -> TypeM lore () -matchExtBranchType rettype (Body _ stms res) = do - ts <- extendedScope (traverse subExpType res) stmscope - matchExtReturns rettype res ts - where stmscope = scopeOf stms - -matchExtReturns :: [ExtType] -> Result -> [Type] -> TypeM lore () -matchExtReturns rettype res ts = do - let problem :: TypeM lore a - problem = bad $ TypeError $ unlines [ "Type annotation is" - , " " ++ prettyTuple rettype - , "But result returns type" - , " " ++ prettyTuple ts ] - - let (ctx_res, val_res) = splitFromEnd (length rettype) res - (ctx_ts, val_ts) = splitFromEnd (length rettype) ts - - unless (length val_res == length rettype) problem - - let num_exts = length $ S.fromList $ - concatMap (mapMaybe isExt . arrayExtDims) rettype - unless (num_exts == length ctx_res) $ - bad $ TypeError $ - "Number of context results does not match number of existentials in the return type.\n" ++ - "Type:\n " ++ - prettyTuple rettype ++ - "\ncannot match context parameters:\n " ++ prettyTuple ctx_res - - let ctx_vals = zip ctx_res ctx_ts - instantiateExt i = case maybeNth i ctx_vals of - Just (se, Prim (IntType Int32)) -> return se - _ -> problem - - rettype' <- instantiateShapes instantiateExt rettype - - unless (rettype' == val_ts) problem - -validApply :: ArrayShape shape => - [TypeBase shape Uniqueness] - -> [TypeBase shape NoUniqueness] - -> Bool -validApply expected got = - length got == length expected && - and (zipWith subtypeOf - (map rankShaped got) - (map (fromDecl . rankShaped) expected)) - -type Arg = (Type, Names) - -argType :: Arg -> Type -argType (t, _) = t - --- | Remove all aliases from the 'Arg'. -argAliases :: Arg -> Names -argAliases (_, als) = als - -noArgAliases :: Arg -> Arg -noArgAliases (t, _) = (t, mempty) - -checkArg :: Checkable lore => - SubExp -> TypeM lore Arg -checkArg arg = do argt <- checkSubExp arg - als <- subExpAliasesM arg - return (argt, als) - -checkFuncall :: Maybe Name - -> [DeclType] -> [Arg] - -> TypeM lore () -checkFuncall fname paramts args = do - let argts = map argType args - unless (validApply paramts argts) $ - bad $ ParameterMismatch fname - (map fromDecl paramts) $ - map argType args - forM_ (zip (map diet paramts) args) $ \(d, (_, als)) -> - occur [consumption (consumeArg als d)] - where consumeArg als Consume = als - consumeArg _ _ = mempty - -checkLambda :: Checkable lore => - Lambda (Aliases lore) -> [Arg] -> TypeM lore () -checkLambda (Lambda params body rettype) args = do - let fname = nameFromString "" - if length params == length args then do - checkFuncall Nothing - (map ((`toDecl` Nonunique) . paramType) params) args - let consumable = zip (map paramName params) (map argAliases args) - checkFun' (fname, - staticShapes $ map (`toDecl` Nonunique) rettype, - [ (paramName param, - LParamName $ paramDec param) - | param <- params ]) consumable $ do - checkLambdaParams params - mapM_ checkType rettype - checkLambdaBody rettype body - else bad $ TypeError $ "Anonymous function defined with " ++ show (length params) ++ " parameters, but expected to take " ++ show (length args) ++ " arguments." - -checkPrimExp :: Checkable lore => PrimExp VName -> TypeM lore () -checkPrimExp ValueExp{} = return () -checkPrimExp (LeafExp v pt) = requireI [Prim pt] v -checkPrimExp (BinOpExp op x y) = do requirePrimExp (binOpType op) x - requirePrimExp (binOpType op) y -checkPrimExp (CmpOpExp op x y) = do requirePrimExp (cmpOpType op) x - requirePrimExp (cmpOpType op) y -checkPrimExp (UnOpExp op x) = requirePrimExp (unOpType op) x -checkPrimExp (ConvOpExp op x) = requirePrimExp (fst $ convOpType op) x -checkPrimExp (FunExp h args t) = do - (h_ts, h_ret, _) <- maybe (bad $ TypeError $ "Unknown function: " ++ h) - return $ M.lookup h primFuns - when (length h_ts /= length args) $ - bad $ TypeError $ "Function expects " ++ show (length h_ts) ++ - " parameters, but given " ++ show (length args) ++ " arguments." - when (h_ret /= t) $ - bad $ TypeError $ "Function return annotation is " ++ pretty t ++ - ", but expected " ++ pretty h_ret - zipWithM_ requirePrimExp h_ts args - -requirePrimExp :: Checkable lore => PrimType -> PrimExp VName -> TypeM lore () -requirePrimExp t e = context ("in PrimExp " ++ pretty e) $ do - checkPrimExp e - unless (primExpType e == t) $ bad $ TypeError $ - pretty e ++ " must have type " ++ pretty t - -class ASTLore lore => CheckableOp lore where - checkOp :: OpWithAliases (Op lore) -> TypeM lore () - -- ^ Used at top level; can be locally changed with 'checkOpWith'. - --- | The class of lores that can be type-checked. -class (ASTLore lore, CanBeAliased (Op lore), CheckableOp lore) => Checkable lore where - checkExpLore :: ExpDec lore -> TypeM lore () - checkBodyLore :: BodyDec lore -> TypeM lore () - checkFParamLore :: VName -> FParamInfo lore -> TypeM lore () - checkLParamLore :: VName -> LParamInfo lore -> TypeM lore () - checkLetBoundLore :: VName -> LetDec lore -> TypeM lore () - checkRetType :: [RetType lore] -> TypeM lore () - matchPattern :: Pattern (Aliases lore) -> Exp (Aliases lore) -> TypeM lore () - primFParam :: VName -> PrimType -> TypeM lore (FParam (Aliases lore)) - matchReturnType :: [RetType lore] -> Result -> TypeM lore () - matchBranchType :: [BranchType lore] -> Body (Aliases lore) -> TypeM lore () - matchLoopResult :: [FParam (Aliases lore)] -> [FParam (Aliases lore)] - -> [SubExp] -> TypeM lore () - - default checkExpLore :: ExpDec lore ~ () => ExpDec lore -> TypeM lore () - checkExpLore = return - - default checkBodyLore :: BodyDec lore ~ () => BodyDec lore -> TypeM lore () - checkBodyLore = return - - default checkFParamLore :: FParamInfo lore ~ DeclType => VName -> FParamInfo lore -> TypeM lore () - checkFParamLore _ = checkType - - default checkLParamLore :: LParamInfo lore ~ Type => VName -> LParamInfo lore -> TypeM lore () - checkLParamLore _ = checkType - - default checkLetBoundLore :: LetDec lore ~ Type => VName -> LetDec lore -> TypeM lore () - checkLetBoundLore _ = checkType - - default checkRetType :: RetType lore ~ DeclExtType => [RetType lore] -> TypeM lore () - checkRetType = mapM_ $ checkExtType . declExtTypeOf - - default matchPattern :: Pattern (Aliases lore) -> Exp (Aliases lore) -> TypeM lore () - matchPattern pat = matchExtPattern pat <=< expExtType - - default primFParam :: FParamInfo lore ~ DeclType => VName -> PrimType -> TypeM lore (FParam (Aliases lore)) - primFParam name t = return $ Param name (Prim t) - - default matchReturnType :: RetType lore ~ DeclExtType => [RetType lore] -> Result -> TypeM lore () - matchReturnType = matchExtReturnType . map fromDecl - - default matchBranchType :: BranchType lore ~ ExtType => [BranchType lore] -> Body (Aliases lore) -> TypeM lore () - matchBranchType = matchExtBranchType - - default matchLoopResult :: FParamInfo lore ~ DeclType => - [FParam (Aliases lore)] -> [FParam (Aliases lore)] - -> [SubExp] -> TypeM lore () +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE Strict #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} + +-- | The type checker checks whether the program is type-consistent. +module Futhark.TypeCheck + ( -- * Interface + checkProg, + TypeError (..), + ErrorCase (..), + + -- * Extensionality + TypeM, + bad, + context, + message, + Checkable (..), + CheckableOp (..), + lookupVar, + lookupAliases, + checkOpWith, + + -- * Checkers + require, + requireI, + requirePrimExp, + checkSubExp, + checkExp, + checkStms, + checkStm, + checkType, + checkExtType, + matchExtPattern, + matchExtBranchType, + argType, + argAliases, + noArgAliases, + checkArg, + checkSOACArrayArgs, + checkLambda, + checkBody, + consume, + consumeOnlyParams, + binding, + ) +where + +import Control.Monad.RWS.Strict +import Control.Parallel.Strategies +import Data.List (find, intercalate, sort) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import Futhark.Analysis.PrimExp +import Futhark.Construct (instantiateShapes) +import Futhark.IR.Aliases +import Futhark.Util +import Futhark.Util.Pretty (Pretty, align, indent, ppr, prettyDoc, text, (<+>)) + +-- | Information about an error during type checking. The 'Show' +-- instance for this type produces a human-readable description. +data ErrorCase lore + = TypeError String + | UnexpectedType (Exp lore) Type [Type] + | ReturnTypeError Name [ExtType] [ExtType] + | DupDefinitionError Name + | DupParamError Name VName + | DupPatternError VName + | InvalidPatternError (Pattern (Aliases lore)) [ExtType] (Maybe String) + | UnknownVariableError VName + | UnknownFunctionError Name + | ParameterMismatch (Maybe Name) [Type] [Type] + | SlicingError Int Int + | BadAnnotation String Type Type + | ReturnAliased Name VName + | UniqueReturnAliased Name + | NotAnArray VName Type + | PermutationError [Int] Int (Maybe VName) + +instance Checkable lore => Show (ErrorCase lore) where + show (TypeError msg) = + "Type error:\n" ++ msg + show (UnexpectedType e _ []) = + "Type of expression\n" + ++ prettyDoc 160 (indent 2 $ ppr e) + ++ "\ncannot have any type - possibly a bug in the type checker." + show (UnexpectedType e t ts) = + "Type of expression\n" + ++ prettyDoc 160 (indent 2 $ ppr e) + ++ "\nmust be one of " + ++ intercalate ", " (map pretty ts) + ++ ", but is " + ++ pretty t + ++ "." + show (ReturnTypeError fname rettype bodytype) = + "Declaration of function " ++ nameToString fname + ++ " declares return type\n " + ++ prettyTuple rettype + ++ "\nBut body has type\n " + ++ prettyTuple bodytype + show (DupDefinitionError name) = + "Duplicate definition of function " ++ nameToString name ++ "" + show (DupParamError funname paramname) = + "Parameter " ++ pretty paramname + ++ " mentioned multiple times in argument list of function " + ++ nameToString funname + ++ "." + show (DupPatternError name) = + "Variable " ++ pretty name ++ " bound twice in pattern." + show (InvalidPatternError pat t desc) = + "Pattern " ++ pretty pat + ++ " cannot match value of type " + ++ prettyTuple t + ++ end + where + end = case desc of + Nothing -> "." + Just desc' -> ":\n" ++ desc' + show (UnknownVariableError name) = + "Use of unknown variable " ++ pretty name ++ "." + show (UnknownFunctionError fname) = + "Call of unknown function " ++ nameToString fname ++ "." + show (ParameterMismatch fname expected got) = + "In call of " ++ fname' ++ ":\n" + ++ "expecting " + ++ show nexpected + ++ " arguments of type(s)\n" + ++ intercalate ", " (map pretty expected) + ++ "\nGot " + ++ show ngot + ++ " arguments of types\n" + ++ intercalate ", " (map pretty got) + where + nexpected = length expected + ngot = length got + fname' = maybe "anonymous function" (("function " ++) . nameToString) fname + show (SlicingError dims got) = + show got ++ " indices given, but type of indexee has " ++ show dims ++ " dimension(s)." + show (BadAnnotation desc expected got) = + "Annotation of \"" ++ desc ++ "\" type of expression is " ++ pretty expected + ++ ", but derived to be " + ++ pretty got + ++ "." + show (ReturnAliased fname name) = + "Unique return value of function " ++ nameToString fname + ++ " is aliased to " + ++ pretty name + ++ ", which is not consumed." + show (UniqueReturnAliased fname) = + "A unique tuple element of return value of function " + ++ nameToString fname + ++ " is aliased to some other tuple component." + show (NotAnArray e t) = + "The expression " ++ pretty e + ++ " is expected to be an array, but is " + ++ pretty t + ++ "." + show (PermutationError perm rank name) = + "The permutation (" ++ intercalate ", " (map show perm) + ++ ") is not valid for array " + ++ name' + ++ "of rank " + ++ show rank + ++ "." + where + name' = maybe "" ((++ " ") . pretty) name + +-- | A type error. +data TypeError lore = Error [String] (ErrorCase lore) + +instance Checkable lore => Show (TypeError lore) where + show (Error [] err) = + show err + show (Error msgs err) = + intercalate "\n" msgs ++ "\n" ++ show err + +-- | A tuple of a return type and a list of parameters, possibly +-- named. +type FunBinding lore = ([RetType (Aliases lore)], [FParam (Aliases lore)]) + +type VarBinding lore = NameInfo (Aliases lore) + +data Usage + = Consumed + | Observed + deriving (Eq, Ord, Show) + +data Occurence = Occurence + { observed :: Names, + consumed :: Names + } + deriving (Eq, Show) + +observation :: Names -> Occurence +observation = flip Occurence mempty + +consumption :: Names -> Occurence +consumption = Occurence mempty + +nullOccurence :: Occurence -> Bool +nullOccurence occ = observed occ == mempty && consumed occ == mempty + +type Occurences = [Occurence] + +allConsumed :: Occurences -> Names +allConsumed = mconcat . map consumed + +seqOccurences :: Occurences -> Occurences -> Occurences +seqOccurences occurs1 occurs2 = + filter (not . nullOccurence) (map filt occurs1) ++ occurs2 + where + filt occ = + occ {observed = observed occ `namesSubtract` postcons} + postcons = allConsumed occurs2 + +altOccurences :: Occurences -> Occurences -> Occurences +altOccurences occurs1 occurs2 = + filter (not . nullOccurence) (map filt occurs1) ++ occurs2 + where + filt occ = + occ + { consumed = consumed occ `namesSubtract` postcons, + observed = observed occ `namesSubtract` postcons + } + postcons = allConsumed occurs2 + +unOccur :: Names -> Occurences -> Occurences +unOccur to_be_removed = filter (not . nullOccurence) . map unOccur' + where + unOccur' occ = + occ + { observed = observed occ `namesSubtract` to_be_removed, + consumed = consumed occ `namesSubtract` to_be_removed + } + +-- | The 'Consumption' data structure is used to keep track of which +-- variables have been consumed, as well as whether a violation has been detected. +data Consumption + = ConsumptionError String + | Consumption Occurences + deriving (Show) + +instance Semigroup Consumption where + ConsumptionError e <> _ = ConsumptionError e + _ <> ConsumptionError e = ConsumptionError e + Consumption o1 <> Consumption o2 + | v : _ <- namesToList $ consumed_in_o1 `namesIntersection` used_in_o2 = + ConsumptionError $ "Variable " <> pretty v <> " referenced after being consumed." + | otherwise = + Consumption $ o1 `seqOccurences` o2 + where + consumed_in_o1 = mconcat $ map consumed o1 + used_in_o2 = mconcat $ map consumed o2 <> map observed o2 + +instance Monoid Consumption where + mempty = Consumption mempty + +-- | The environment contains a variable table and a function table. +-- Type checking happens with access to this environment. The +-- function table is only initialised at the very beginning, but the +-- variable table will be extended during type-checking when +-- let-expressions are encountered. +data Env lore = Env + { envVtable :: M.Map VName (VarBinding lore), + envFtable :: M.Map Name (FunBinding lore), + envCheckOp :: OpWithAliases (Op lore) -> TypeM lore (), + envContext :: [String] + } + +-- | The type checker runs in this monad. +newtype TypeM lore a + = TypeM + ( RWST + (Env lore) -- Reader + Consumption -- Writer + Names -- State + (Either (TypeError lore)) -- Inner monad + a + ) + deriving + ( Monad, + Functor, + Applicative, + MonadReader (Env lore), + MonadWriter Consumption, + MonadState Names + ) + +instance + Checkable lore => + HasScope (Aliases lore) (TypeM lore) + where + lookupType = fmap typeOf . lookupVar + askScope = asks $ M.fromList . mapMaybe varType . M.toList . envVtable + where + varType (name, dec) = Just (name, dec) + +runTypeM :: + Env lore -> + TypeM lore a -> + Either (TypeError lore) (a, Consumption) +runTypeM env (TypeM m) = evalRWST m env mempty + +bad :: ErrorCase lore -> TypeM lore a +bad e = do + messages <- asks envContext + TypeM $ lift $ Left $ Error (reverse messages) e + +-- | Add information about what is being type-checked to the current +-- context. Liberal use of this combinator makes it easier to track +-- type errors, as the strings are added to type errors signalled via +-- 'bad'. +context :: + String -> + TypeM lore a -> + TypeM lore a +context s = local $ \env -> env {envContext = s : envContext env} + +message :: + Pretty a => + String -> + a -> + String +message s x = + prettyDoc 80 $ + text s <+> align (ppr x) + +-- | Mark a name as bound. If the name has been bound previously in +-- the program, report a type error. +bound :: VName -> TypeM lore () +bound name = do + already_seen <- gets $ nameIn name + when already_seen $ + bad $ TypeError $ "Name " ++ pretty name ++ " bound twice" + modify (<> oneName name) + +occur :: Occurences -> TypeM lore () +occur = tell . Consumption . filter (not . nullOccurence) + +-- | Proclaim that we have made read-only use of the given variable. +-- No-op unless the variable is array-typed. +observe :: + Checkable lore => + VName -> + TypeM lore () +observe name = do + dec <- lookupVar name + unless (primType $ typeOf dec) $ + occur [observation $ oneName name <> aliases dec] + +-- | Proclaim that we have written to the given variables. +consume :: Checkable lore => Names -> TypeM lore () +consume als = do + scope <- askScope + let isArray = maybe False ((> 0) . arrayRank . typeOf) . (`M.lookup` scope) + occur [consumption $ namesFromList $ filter isArray $ namesToList als] + +collectOccurences :: TypeM lore a -> TypeM lore (a, Occurences) +collectOccurences m = pass $ do + (x, c) <- listen m + o <- checkConsumption c + return ((x, o), const mempty) + +checkOpWith :: + (OpWithAliases (Op lore) -> TypeM lore ()) -> + TypeM lore a -> + TypeM lore a +checkOpWith checker = local $ \env -> env {envCheckOp = checker} + +checkConsumption :: Consumption -> TypeM lore Occurences +checkConsumption (ConsumptionError e) = bad $ TypeError e +checkConsumption (Consumption os) = return os + +alternative :: TypeM lore a -> TypeM lore b -> TypeM lore (a, b) +alternative m1 m2 = pass $ do + (x, c1) <- listen m1 + (y, c2) <- listen m2 + os1 <- checkConsumption c1 + os2 <- checkConsumption c2 + let usage = Consumption $ os1 `altOccurences` os2 + return ((x, y), const usage) + +-- | Permit consumption of only the specified names. If one of these +-- names is consumed, the consumption will be rewritten to be a +-- consumption of the corresponding alias set. Consumption of +-- anything else will result in a type error. +consumeOnlyParams :: [(VName, Names)] -> TypeM lore a -> TypeM lore a +consumeOnlyParams consumable m = do + (x, os) <- collectOccurences m + tell . Consumption =<< mapM inspect os + return x + where + inspect o = do + new_consumed <- mconcat <$> mapM wasConsumed (namesToList $ consumed o) + return o {consumed = new_consumed} + wasConsumed v + | Just als <- lookup v consumable = return als + | otherwise = + bad $ + TypeError $ + unlines + [ pretty v ++ " was invalidly consumed.", + what ++ " can be consumed here." + ] + what + | null consumable = "Nothing" + | otherwise = "Only " ++ intercalate ", " (map (pretty . fst) consumable) + +-- | Given the immediate aliases, compute the full transitive alias +-- set (including the immediate aliases). +expandAliases :: Names -> Env lore -> Names +expandAliases names env = names <> aliasesOfAliases + where + aliasesOfAliases = mconcat . map look . namesToList $ names + look k = case M.lookup k $ envVtable env of + Just (LetName (als, _)) -> unAliases als + _ -> mempty + +binding :: + Checkable lore => + Scope (Aliases lore) -> + TypeM lore a -> + TypeM lore a +binding bnds = check . local (`bindVars` bnds) + where + bindVars = M.foldlWithKey' bindVar + boundnames = M.keys bnds + + bindVar env name (LetName (AliasDec als, dec)) = + let als' + | primType (typeOf dec) = mempty + | otherwise = expandAliases als env + in env + { envVtable = + M.insert name (LetName (AliasDec als', dec)) $ envVtable env + } + bindVar env name dec = + env {envVtable = M.insert name dec $ envVtable env} + + -- Check whether the bound variables have been used correctly + -- within their scope. + check m = do + mapM_ bound $ M.keys bnds + (a, os) <- collectOccurences m + tell $ Consumption $ unOccur (namesFromList boundnames) os + return a + +lookupVar :: VName -> TypeM lore (NameInfo (Aliases lore)) +lookupVar name = do + bnd <- asks $ M.lookup name . envVtable + case bnd of + Nothing -> bad $ UnknownVariableError name + Just dec -> return dec + +lookupAliases :: Checkable lore => VName -> TypeM lore Names +lookupAliases name = do + info <- lookupVar name + return $ + if primType $ typeOf info + then mempty + else oneName name <> aliases info + +aliases :: NameInfo (Aliases lore) -> Names +aliases (LetName (als, _)) = unAliases als +aliases _ = mempty + +subExpAliasesM :: Checkable lore => SubExp -> TypeM lore Names +subExpAliasesM Constant {} = return mempty +subExpAliasesM (Var v) = lookupAliases v + +lookupFun :: + Checkable lore => + Name -> + [SubExp] -> + TypeM lore ([RetType lore], [DeclType]) +lookupFun fname args = do + bnd <- asks $ M.lookup fname . envFtable + case bnd of + Nothing -> bad $ UnknownFunctionError fname + Just (ftype, params) -> do + argts <- mapM subExpType args + case applyRetType ftype params $ zip args argts of + Nothing -> + bad $ ParameterMismatch (Just fname) (map paramType params) argts + Just rt -> + return (rt, map paramDeclType params) + +-- | @checkAnnotation loc s t1 t2@ checks if @t2@ is equal to +-- @t1@. If not, a 'BadAnnotation' is raised. +checkAnnotation :: + String -> + Type -> + Type -> + TypeM lore () +checkAnnotation desc t1 t2 + | t2 == t1 = return () + | otherwise = bad $ BadAnnotation desc t1 t2 + +-- | @require ts se@ causes a '(TypeError vn)' if the type of @se@ is +-- not a subtype of one of the types in @ts@. +require :: Checkable lore => [Type] -> SubExp -> TypeM lore () +require ts se = do + t <- checkSubExp se + unless (t `elem` ts) $ + bad $ UnexpectedType (BasicOp $ SubExp se) t ts + +-- | Variant of 'require' working on variable names. +requireI :: Checkable lore => [Type] -> VName -> TypeM lore () +requireI ts ident = require ts $ Var ident + +checkArrIdent :: + Checkable lore => + VName -> + TypeM lore Type +checkArrIdent v = do + t <- lookupType v + case t of + Array {} -> return t + _ -> bad $ NotAnArray v t + +-- | Type check a program containing arbitrary type information, +-- yielding either a type error or a program with complete type +-- information. +checkProg :: + Checkable lore => + Prog (Aliases lore) -> + Either (TypeError lore) () +checkProg (Prog consts funs) = do + let typeenv = + Env + { envVtable = M.empty, + envFtable = mempty, + envContext = [], + envCheckOp = checkOp + } + let onFunction ftable vtable fun = + fmap fst $ + runTypeM typeenv $ + local (\env -> env {envFtable = ftable, envVtable = vtable}) $ + checkFun fun + (ftable, _) <- runTypeM typeenv buildFtable + (vtable, _) <- + runTypeM typeenv {envFtable = ftable} $ + checkStms consts $ asks envVtable + sequence_ $ parMap rpar (onFunction ftable vtable) funs + where + buildFtable = do + table <- initialFtable + foldM expand table funs + expand ftable (FunDef _ _ name ret params _) + | M.member name ftable = + bad $ DupDefinitionError name + | otherwise = + return $ M.insert name (ret, params) ftable + +initialFtable :: + Checkable lore => + TypeM lore (M.Map Name (FunBinding lore)) +initialFtable = fmap M.fromList $ mapM addBuiltin $ M.toList builtInFunctions + where + addBuiltin (fname, (t, ts)) = do + ps <- mapM (primFParam name) ts + return (fname, ([primRetType t], ps)) + name = VName (nameFromString "x") 0 + +checkFun :: + Checkable lore => + FunDef (Aliases lore) -> + TypeM lore () +checkFun (FunDef _ _ fname rettype params body) = + context ("In function " ++ nameToString fname) $ + checkFun' + ( fname, + map declExtTypeOf rettype, + funParamsToNameInfos params + ) + consumable + $ do + checkFunParams params + checkRetType rettype + context "When checking function body" $ checkFunBody rettype body + where + consumable = + [ (paramName param, mempty) + | param <- params, + unique $ paramDeclType param + ] + +funParamsToNameInfos :: + [FParam lore] -> + [(VName, NameInfo (Aliases lore))] +funParamsToNameInfos = map nameTypeAndLore + where + nameTypeAndLore fparam = + ( paramName fparam, + FParamName $ paramDec fparam + ) + +checkFunParams :: + Checkable lore => + [FParam lore] -> + TypeM lore () +checkFunParams = mapM_ $ \param -> + context ("In function parameter " ++ pretty param) $ + checkFParamLore (paramName param) (paramDec param) + +checkLambdaParams :: + Checkable lore => + [LParam lore] -> + TypeM lore () +checkLambdaParams = mapM_ $ \param -> + context ("In lambda parameter " ++ pretty param) $ + checkLParamLore (paramName param) (paramDec param) + +checkFun' :: + Checkable lore => + ( Name, + [DeclExtType], + [(VName, NameInfo (Aliases lore))] + ) -> + [(VName, Names)] -> + TypeM lore [Names] -> + TypeM lore () +checkFun' (fname, rettype, params) consumable check = do + checkNoDuplicateParams + binding (M.fromList params) $ + consumeOnlyParams consumable $ do + body_aliases <- check + scope <- askScope + let isArray = maybe False ((> 0) . arrayRank . typeOf) . (`M.lookup` scope) + context + ( "When checking the body aliases: " + ++ pretty (map namesToList body_aliases) + ) + $ checkReturnAlias $ map (namesFromList . filter isArray . namesToList) body_aliases + where + param_names = map fst params + + checkNoDuplicateParams = foldM_ expand [] param_names + + expand seen pname + | Just _ <- find (== pname) seen = + bad $ DupParamError fname pname + | otherwise = + return $ pname : seen + checkReturnAlias = + foldM_ checkReturnAlias' mempty . returnAliasing rettype + + checkReturnAlias' seen (Unique, names) + | any (`S.member` S.map fst seen) $ namesToList names = + bad $ UniqueReturnAliased fname + | otherwise = do + consume names + return $ seen <> tag Unique names + checkReturnAlias' seen (Nonunique, names) + | any (`S.member` seen) $ tag Unique names = + bad $ UniqueReturnAliased fname + | otherwise = return $ seen <> tag Nonunique names + + tag u = S.fromList . map (,u) . namesToList + + returnAliasing expected got = + reverse $ + zip (reverse (map uniqueness expected) ++ repeat Nonunique) $ + reverse got + +checkSubExp :: Checkable lore => SubExp -> TypeM lore Type +checkSubExp (Constant val) = + return $ Prim $ primValueType val +checkSubExp (Var ident) = context ("In subexp " ++ pretty ident) $ do + observe ident + lookupType ident + +checkStms :: + Checkable lore => + Stms (Aliases lore) -> + TypeM lore a -> + TypeM lore a +checkStms origbnds m = delve $ stmsToList origbnds + where + delve (stm@(Let pat _ e) : bnds) = do + context ("In expression of statement " ++ pretty pat) $ + checkExp e + checkStm stm $ + delve bnds + delve [] = + m + +checkResult :: + Checkable lore => + Result -> + TypeM lore () +checkResult = mapM_ checkSubExp + +checkFunBody :: + Checkable lore => + [RetType lore] -> + Body (Aliases lore) -> + TypeM lore [Names] +checkFunBody rt (Body (_, lore) bnds res) = do + checkBodyLore lore + checkStms bnds $ do + context "When checking body result" $ checkResult res + context "When matching declared return type to result of body" $ + matchReturnType rt res + map (`namesSubtract` bound_here) <$> mapM subExpAliasesM res + where + bound_here = namesFromList $ M.keys $ scopeOf bnds + +checkLambdaBody :: + Checkable lore => + [Type] -> + Body (Aliases lore) -> + TypeM lore [Names] +checkLambdaBody ret (Body (_, lore) bnds res) = do + checkBodyLore lore + checkStms bnds $ do + checkLambdaResult ret res + map (`namesSubtract` bound_here) <$> mapM subExpAliasesM res + where + bound_here = namesFromList $ M.keys $ scopeOf bnds + +checkLambdaResult :: + Checkable lore => + [Type] -> + Result -> + TypeM lore () +checkLambdaResult ts es + | length ts /= length es = + bad $ + TypeError $ + "Lambda has return type " ++ prettyTuple ts + ++ " describing " + ++ show (length ts) + ++ " values, but body returns " + ++ show (length es) + ++ " values: " + ++ prettyTuple es + | otherwise = forM_ (zip ts es) $ \(t, e) -> do + et <- checkSubExp e + unless (et == t) $ + bad $ + TypeError $ + "Subexpression " ++ pretty e ++ " has type " ++ pretty et + ++ " but expected " + ++ pretty t + +checkBody :: + Checkable lore => + Body (Aliases lore) -> + TypeM lore [Names] +checkBody (Body (_, lore) bnds res) = do + checkBodyLore lore + checkStms bnds $ do + checkResult res + map (`namesSubtract` bound_here) <$> mapM subExpAliasesM res + where + bound_here = namesFromList $ M.keys $ scopeOf bnds + +checkBasicOp :: Checkable lore => BasicOp -> TypeM lore () +checkBasicOp (SubExp es) = + void $ checkSubExp es +checkBasicOp (Opaque es) = + void $ checkSubExp es +checkBasicOp (ArrayLit [] _) = + return () +checkBasicOp (ArrayLit (e : es') t) = do + let check elemt eleme = do + elemet <- checkSubExp eleme + unless (elemet == elemt) $ + bad $ + TypeError $ + pretty elemet + ++ " is not of expected type " + ++ pretty elemt + ++ "." + et <- checkSubExp e + + -- Compare that type with the one given for the array literal. + checkAnnotation "array-element" t et + + mapM_ (check et) es' +checkBasicOp (UnOp op e) = require [Prim $ unOpType op] e +checkBasicOp (BinOp op e1 e2) = checkBinOpArgs (binOpType op) e1 e2 +checkBasicOp (CmpOp op e1 e2) = checkCmpOp op e1 e2 +checkBasicOp (ConvOp op e) = require [Prim $ fst $ convOpType op] e +checkBasicOp (Index ident idxes) = do + vt <- lookupType ident + observe ident + when (arrayRank vt /= length idxes) $ + bad $ SlicingError (arrayRank vt) (length idxes) + mapM_ checkDimIndex idxes +checkBasicOp (Update src idxes se) = do + src_t <- checkArrIdent src + when (arrayRank src_t /= length idxes) $ + bad $ SlicingError (arrayRank src_t) (length idxes) + + se_aliases <- subExpAliasesM se + when (src `nameIn` se_aliases) $ + bad $ TypeError "The target of an Update must not alias the value to be written." + + mapM_ checkDimIndex idxes + require [Prim (elemType src_t) `arrayOfShape` Shape (sliceDims idxes)] se + consume =<< lookupAliases src +checkBasicOp (Iota e x s et) = do + require [Prim int64] e + require [Prim $ IntType et] x + require [Prim $ IntType et] s +checkBasicOp (Replicate (Shape dims) valexp) = do + mapM_ (require [Prim int64]) dims + void $ checkSubExp valexp +checkBasicOp (Scratch _ shape) = + mapM_ checkSubExp shape +checkBasicOp (Reshape newshape arrexp) = do + rank <- arrayRank <$> checkArrIdent arrexp + mapM_ (require [Prim int64] . newDim) newshape + zipWithM_ (checkDimChange rank) newshape [0 ..] + where + checkDimChange _ (DimNew _) _ = + return () + checkDimChange rank (DimCoercion se) i + | i >= rank = + bad $ + TypeError $ + "Asked to coerce dimension " ++ show i ++ " to " ++ pretty se + ++ ", but array " + ++ pretty arrexp + ++ " has only " + ++ pretty rank + ++ " dimensions" + | otherwise = + return () +checkBasicOp (Rearrange perm arr) = do + arrt <- lookupType arr + let rank = arrayRank arrt + when (length perm /= rank || sort perm /= [0 .. rank -1]) $ + bad $ PermutationError perm rank $ Just arr +checkBasicOp (Rotate rots arr) = do + arrt <- lookupType arr + let rank = arrayRank arrt + mapM_ (require [Prim int64]) rots + when (length rots /= rank) $ + bad $ + TypeError $ + "Cannot rotate " ++ show (length rots) + ++ " dimensions of " + ++ show rank + ++ "-dimensional array." +checkBasicOp (Concat i arr1exp arr2exps ressize) = do + arr1t <- checkArrIdent arr1exp + arr2ts <- mapM checkArrIdent arr2exps + let success = + all + ( (== dropAt i 1 (arrayDims arr1t)) + . dropAt i 1 + . arrayDims + ) + arr2ts + unless success $ + bad $ + TypeError $ + "Types of arguments to concat do not match. Got " + ++ pretty arr1t + ++ " and " + ++ intercalate ", " (map pretty arr2ts) + require [Prim int64] ressize +checkBasicOp (Copy e) = + void $ checkArrIdent e +checkBasicOp (Manifest perm arr) = + checkBasicOp $ Rearrange perm arr -- Basically same thing! +checkBasicOp (Assert e _ _) = + require [Prim Bool] e + +matchLoopResultExt :: + Checkable lore => + [Param DeclType] -> + [Param DeclType] -> + [SubExp] -> + TypeM lore () +matchLoopResultExt ctx val loopres = do + let rettype_ext = + existentialiseExtTypes (map paramName ctx) $ + staticShapes $ map typeOf $ ctx ++ val + + bodyt <- mapM subExpType loopres + + case instantiateShapes (`maybeNth` loopres) rettype_ext of + Nothing -> + bad $ + ReturnTypeError + (nameFromString "") + rettype_ext + (staticShapes bodyt) + Just rettype' -> + unless (bodyt `subtypesOf` rettype') $ + bad $ + ReturnTypeError + (nameFromString "") + (staticShapes rettype') + (staticShapes bodyt) + +checkExp :: + Checkable lore => + Exp (Aliases lore) -> + TypeM lore () +checkExp (BasicOp op) = checkBasicOp op +checkExp (If e1 e2 e3 info) = do + require [Prim Bool] e1 + _ <- checkBody e2 `alternative` checkBody e3 + context "in true branch" $ matchBranchType (ifReturns info) e2 + context "in false branch" $ matchBranchType (ifReturns info) e3 +checkExp (Apply fname args rettype_annot _) = do + (rettype_derived, paramtypes) <- lookupFun fname $ map fst args + argflows <- mapM (checkArg . fst) args + when (rettype_derived /= rettype_annot) $ + bad $ + TypeError $ + "Expected apply result type " ++ pretty rettype_derived + ++ " but annotation is " + ++ pretty rettype_annot + checkFuncall (Just fname) paramtypes argflows +checkExp (DoLoop ctxmerge valmerge form loopbody) = do + let merge = ctxmerge ++ valmerge + (mergepat, mergeexps) = unzip merge + mergeargs <- mapM checkArg mergeexps + + let val_free = freeIn $ map fst valmerge + usedInVal p = paramName p `nameIn` val_free + case find (not . usedInVal . fst) ctxmerge of + Just p -> + bad $ TypeError $ "Loop context parameter " ++ pretty p ++ " unused." + Nothing -> + return () + + binding (scopeOf form) $ do + case form of + ForLoop loopvar it boundexp loopvars -> do + iparam <- primFParam loopvar $ IntType it + let funparams = iparam : mergepat + paramts = map paramDeclType funparams + + forM_ loopvars $ \(p, a) -> do + a_t <- lookupType a + observe a + case peelArray 1 a_t of + Just a_t_r -> do + checkLParamLore (paramName p) $ paramDec p + unless (a_t_r `subtypeOf` typeOf (paramDec p)) $ + bad $ + TypeError $ + "Loop parameter " ++ pretty p + ++ " not valid for element of " + ++ pretty a + ++ ", which has row type " + ++ pretty a_t_r + _ -> + bad $ + TypeError $ + "Cannot loop over " ++ pretty a + ++ " of type " + ++ pretty a_t + + boundarg <- checkArg boundexp + checkFuncall Nothing paramts $ boundarg : mergeargs + WhileLoop cond -> do + case find ((== cond) . paramName . fst) merge of + Just (condparam, _) -> + unless (paramType condparam == Prim Bool) $ + bad $ + TypeError $ + "Conditional '" ++ pretty cond ++ "' of while-loop is not boolean, but " + ++ pretty (paramType condparam) + ++ "." + Nothing -> + bad $ + TypeError $ + "Conditional '" ++ pretty cond ++ "' of while-loop is not a merge variable." + let funparams = mergepat + paramts = map paramDeclType funparams + checkFuncall Nothing paramts mergeargs + + let rettype = map paramDeclType mergepat + consumable = + [ (paramName param, mempty) + | param <- mergepat, + unique $ paramDeclType param + ] + + context "Inside the loop body" $ + checkFun' + ( nameFromString "", + staticShapes rettype, + funParamsToNameInfos mergepat + ) + consumable + $ do + checkFunParams mergepat + checkBodyLore $ snd $ bodyDec loopbody + + checkStms (bodyStms loopbody) $ do + checkResult $ bodyResult loopbody + + context "When matching result of body with loop parameters" $ + matchLoopResult (map fst ctxmerge) (map fst valmerge) $ + bodyResult loopbody + + let bound_here = + namesFromList $ + M.keys $ + scopeOf $ bodyStms loopbody + map (`namesSubtract` bound_here) + <$> mapM subExpAliasesM (bodyResult loopbody) +checkExp (Op op) = do + checker <- asks envCheckOp + checker op + +checkSOACArrayArgs :: + Checkable lore => + SubExp -> + [VName] -> + TypeM lore [Arg] +checkSOACArrayArgs width vs = + forM vs $ \v -> do + (vt, v') <- checkSOACArrayArg v + let argSize = arraySize 0 vt + unless (argSize == width) $ + bad $ + TypeError $ + "SOAC argument " ++ pretty v ++ " has outer size " + ++ pretty argSize + ++ ", but width of SOAC is " + ++ pretty width + return v' + where + checkSOACArrayArg ident = do + (t, als) <- checkArg $ Var ident + case peelArray 1 t of + Nothing -> + bad $ + TypeError $ + "SOAC argument " ++ pretty ident ++ " is not an array" + Just rt -> return (t, (rt, als)) + +checkType :: + Checkable lore => + TypeBase Shape u -> + TypeM lore () +checkType (Mem (ScalarSpace d _)) = mapM_ (require [Prim int64]) d +checkType t = mapM_ checkSubExp $ arrayDims t + +checkExtType :: + Checkable lore => + TypeBase ExtShape u -> + TypeM lore () +checkExtType = mapM_ checkExtDim . shapeDims . arrayShape + where + checkExtDim (Free se) = void $ checkSubExp se + checkExtDim (Ext _) = return () + +checkCmpOp :: + Checkable lore => + CmpOp -> + SubExp -> + SubExp -> + TypeM lore () +checkCmpOp (CmpEq t) x y = do + require [Prim t] x + require [Prim t] y +checkCmpOp (CmpUlt t) x y = checkBinOpArgs (IntType t) x y +checkCmpOp (CmpUle t) x y = checkBinOpArgs (IntType t) x y +checkCmpOp (CmpSlt t) x y = checkBinOpArgs (IntType t) x y +checkCmpOp (CmpSle t) x y = checkBinOpArgs (IntType t) x y +checkCmpOp (FCmpLt t) x y = checkBinOpArgs (FloatType t) x y +checkCmpOp (FCmpLe t) x y = checkBinOpArgs (FloatType t) x y +checkCmpOp CmpLlt x y = checkBinOpArgs Bool x y +checkCmpOp CmpLle x y = checkBinOpArgs Bool x y + +checkBinOpArgs :: + Checkable lore => + PrimType -> + SubExp -> + SubExp -> + TypeM lore () +checkBinOpArgs t e1 e2 = do + require [Prim t] e1 + require [Prim t] e2 + +checkPatElem :: + Checkable lore => + PatElemT (LetDec lore) -> + TypeM lore () +checkPatElem (PatElem name dec) = + context ("When checking pattern element " ++ pretty name) $ + checkLetBoundLore name dec + +checkDimIndex :: + Checkable lore => + DimIndex SubExp -> + TypeM lore () +checkDimIndex (DimFix i) = require [Prim int64] i +checkDimIndex (DimSlice i n s) = mapM_ (require [Prim int64]) [i, n, s] + +checkStm :: + Checkable lore => + Stm (Aliases lore) -> + TypeM lore a -> + TypeM lore a +checkStm stm@(Let pat (StmAux (Certificates cs) _ (_, dec)) e) m = do + context "When checking certificates" $ mapM_ (requireI [Prim Cert]) cs + context "When checking expression annotation" $ checkExpLore dec + context ("When matching\n" ++ message " " pat ++ "\nwith\n" ++ message " " e) $ + matchPattern pat e + binding (maybeWithoutAliases $ scopeOf stm) $ do + mapM_ checkPatElem (patternElements $ removePatternAliases pat) + m + where + -- FIXME: this is wrong. However, the core language type system + -- is not strong enough to fully capture the aliases we want (see + -- issue #803). Since we eventually inline everything anyway, and + -- our intra-procedural alias analysis is much simpler and + -- correct, I could not justify spending time on improving the + -- inter-procedural alias analysis. If we ever stop inlining + -- everything, probably we need to go back and refine this. + maybeWithoutAliases = + case stmExp stm of + Apply {} -> M.map withoutAliases + _ -> id + withoutAliases (LetName (_, ldec)) = LetName (mempty, ldec) + withoutAliases info = info + +matchExtPattern :: + Checkable lore => + Pattern (Aliases lore) -> + [ExtType] -> + TypeM lore () +matchExtPattern pat ts = + unless (expExtTypesFromPattern pat == ts) $ + bad $ InvalidPatternError pat ts Nothing + +matchExtReturnType :: + Checkable lore => + [ExtType] -> + Result -> + TypeM lore () +matchExtReturnType rettype res = do + ts <- mapM subExpType res + matchExtReturns rettype res ts + +matchExtBranchType :: + Checkable lore => + [ExtType] -> + Body (Aliases lore) -> + TypeM lore () +matchExtBranchType rettype (Body _ stms res) = do + ts <- extendedScope (traverse subExpType res) stmscope + matchExtReturns rettype res ts + where + stmscope = scopeOf stms + +matchExtReturns :: [ExtType] -> Result -> [Type] -> TypeM lore () +matchExtReturns rettype res ts = do + let problem :: TypeM lore a + problem = + bad $ + TypeError $ + unlines + [ "Type annotation is", + " " ++ prettyTuple rettype, + "But result returns type", + " " ++ prettyTuple ts + ] + + let (ctx_res, val_res) = splitFromEnd (length rettype) res + (ctx_ts, val_ts) = splitFromEnd (length rettype) ts + + unless (length val_res == length rettype) problem + + let num_exts = + length $ + S.fromList $ + concatMap (mapMaybe isExt . arrayExtDims) rettype + unless (num_exts == length ctx_res) $ + bad $ + TypeError $ + "Number of context results does not match number of existentials in the return type.\n" + ++ "Type:\n " + ++ prettyTuple rettype + ++ "\ncannot match context parameters:\n " + ++ prettyTuple ctx_res + + let ctx_vals = zip ctx_res ctx_ts + instantiateExt i = case maybeNth i ctx_vals of + Just (se, Prim (IntType Int64)) -> return se + _ -> problem + + rettype' <- instantiateShapes instantiateExt rettype + + unless (rettype' == val_ts) problem + +validApply :: + ArrayShape shape => + [TypeBase shape Uniqueness] -> + [TypeBase shape NoUniqueness] -> + Bool +validApply expected got = + length got == length expected + && and + ( zipWith + subtypeOf + (map rankShaped got) + (map (fromDecl . rankShaped) expected) + ) + +type Arg = (Type, Names) + +argType :: Arg -> Type +argType (t, _) = t + +-- | Remove all aliases from the 'Arg'. +argAliases :: Arg -> Names +argAliases (_, als) = als + +noArgAliases :: Arg -> Arg +noArgAliases (t, _) = (t, mempty) + +checkArg :: + Checkable lore => + SubExp -> + TypeM lore Arg +checkArg arg = do + argt <- checkSubExp arg + als <- subExpAliasesM arg + return (argt, als) + +checkFuncall :: + Maybe Name -> + [DeclType] -> + [Arg] -> + TypeM lore () +checkFuncall fname paramts args = do + let argts = map argType args + unless (validApply paramts argts) $ + bad $ + ParameterMismatch + fname + (map fromDecl paramts) + $ map argType args + forM_ (zip (map diet paramts) args) $ \(d, (_, als)) -> + occur [consumption (consumeArg als d)] + where + consumeArg als Consume = als + consumeArg _ _ = mempty + +checkLambda :: + Checkable lore => + Lambda (Aliases lore) -> + [Arg] -> + TypeM lore () +checkLambda (Lambda params body rettype) args = do + let fname = nameFromString "" + if length params == length args + then do + checkFuncall + Nothing + (map ((`toDecl` Nonunique) . paramType) params) + args + let consumable = zip (map paramName params) (map argAliases args) + checkFun' + ( fname, + staticShapes $ map (`toDecl` Nonunique) rettype, + [ ( paramName param, + LParamName $ paramDec param + ) + | param <- params + ] + ) + consumable + $ do + checkLambdaParams params + mapM_ checkType rettype + checkLambdaBody rettype body + else bad $ TypeError $ "Anonymous function defined with " ++ show (length params) ++ " parameters, but expected to take " ++ show (length args) ++ " arguments." + +checkPrimExp :: Checkable lore => PrimExp VName -> TypeM lore () +checkPrimExp ValueExp {} = return () +checkPrimExp (LeafExp v pt) = requireI [Prim pt] v +checkPrimExp (BinOpExp op x y) = do + requirePrimExp (binOpType op) x + requirePrimExp (binOpType op) y +checkPrimExp (CmpOpExp op x y) = do + requirePrimExp (cmpOpType op) x + requirePrimExp (cmpOpType op) y +checkPrimExp (UnOpExp op x) = requirePrimExp (unOpType op) x +checkPrimExp (ConvOpExp op x) = requirePrimExp (fst $ convOpType op) x +checkPrimExp (FunExp h args t) = do + (h_ts, h_ret, _) <- + maybe + (bad $ TypeError $ "Unknown function: " ++ h) + return + $ M.lookup h primFuns + when (length h_ts /= length args) $ + bad $ + TypeError $ + "Function expects " ++ show (length h_ts) + ++ " parameters, but given " + ++ show (length args) + ++ " arguments." + when (h_ret /= t) $ + bad $ + TypeError $ + "Function return annotation is " ++ pretty t + ++ ", but expected " + ++ pretty h_ret + zipWithM_ requirePrimExp h_ts args + +requirePrimExp :: Checkable lore => PrimType -> PrimExp VName -> TypeM lore () +requirePrimExp t e = context ("in PrimExp " ++ pretty e) $ do + checkPrimExp e + unless (primExpType e == t) $ + bad $ + TypeError $ + pretty e ++ " must have type " ++ pretty t + +class ASTLore lore => CheckableOp lore where + checkOp :: OpWithAliases (Op lore) -> TypeM lore () + -- ^ Used at top level; can be locally changed with 'checkOpWith'. + +-- | The class of lores that can be type-checked. +class (ASTLore lore, CanBeAliased (Op lore), CheckableOp lore) => Checkable lore where + checkExpLore :: ExpDec lore -> TypeM lore () + checkBodyLore :: BodyDec lore -> TypeM lore () + checkFParamLore :: VName -> FParamInfo lore -> TypeM lore () + checkLParamLore :: VName -> LParamInfo lore -> TypeM lore () + checkLetBoundLore :: VName -> LetDec lore -> TypeM lore () + checkRetType :: [RetType lore] -> TypeM lore () + matchPattern :: Pattern (Aliases lore) -> Exp (Aliases lore) -> TypeM lore () + primFParam :: VName -> PrimType -> TypeM lore (FParam (Aliases lore)) + matchReturnType :: [RetType lore] -> Result -> TypeM lore () + matchBranchType :: [BranchType lore] -> Body (Aliases lore) -> TypeM lore () + matchLoopResult :: + [FParam (Aliases lore)] -> + [FParam (Aliases lore)] -> + [SubExp] -> + TypeM lore () + + default checkExpLore :: ExpDec lore ~ () => ExpDec lore -> TypeM lore () + checkExpLore = return + + default checkBodyLore :: BodyDec lore ~ () => BodyDec lore -> TypeM lore () + checkBodyLore = return + + default checkFParamLore :: FParamInfo lore ~ DeclType => VName -> FParamInfo lore -> TypeM lore () + checkFParamLore _ = checkType + + default checkLParamLore :: LParamInfo lore ~ Type => VName -> LParamInfo lore -> TypeM lore () + checkLParamLore _ = checkType + + default checkLetBoundLore :: LetDec lore ~ Type => VName -> LetDec lore -> TypeM lore () + checkLetBoundLore _ = checkType + + default checkRetType :: RetType lore ~ DeclExtType => [RetType lore] -> TypeM lore () + checkRetType = mapM_ $ checkExtType . declExtTypeOf + + default matchPattern :: Pattern (Aliases lore) -> Exp (Aliases lore) -> TypeM lore () + matchPattern pat = matchExtPattern pat <=< expExtType + + default primFParam :: FParamInfo lore ~ DeclType => VName -> PrimType -> TypeM lore (FParam (Aliases lore)) + primFParam name t = return $ Param name (Prim t) + + default matchReturnType :: RetType lore ~ DeclExtType => [RetType lore] -> Result -> TypeM lore () + matchReturnType = matchExtReturnType . map fromDecl + + default matchBranchType :: BranchType lore ~ ExtType => [BranchType lore] -> Body (Aliases lore) -> TypeM lore () + matchBranchType = matchExtBranchType + + default matchLoopResult :: + FParamInfo lore ~ DeclType => + [FParam (Aliases lore)] -> + [FParam (Aliases lore)] -> + [SubExp] -> + TypeM lore () matchLoopResult = matchLoopResultExt diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs --- a/src/Futhark/Util.hs +++ b/src/Futhark/Util.hs @@ -1,4 +1,5 @@ {-# LANGUAGE Trustworthy #-} + -- | Non-Futhark-specific utilities. If you find yourself writing -- general functions on generic data structures, consider putting them -- here. @@ -8,66 +9,77 @@ -- note where you got it from (and make sure that the license is -- compatible). module Futhark.Util - (mapAccumLM, - maxinum, - chunk, - chunks, - dropAt, - takeLast, - dropLast, - mapEither, - maybeNth, - maybeHead, - splitFromEnd, - splitAt3, - focusNth, - unixEnvironment, - isEnvVarSet, - fancyTerminal, - runProgramWithExitCode, - directoryContents, - roundFloat, ceilFloat, floorFloat, - roundDouble, ceilDouble, floorDouble, - lgamma, lgammaf, tgamma, tgammaf, - fromPOSIX, - toPOSIX, - trim, - pmapIO, - UserString, - EncodedString, - zEncodeString - ) - where + ( mapAccumLM, + maxinum, + chunk, + chunks, + dropAt, + takeLast, + dropLast, + mapEither, + maybeNth, + maybeHead, + splitFromEnd, + splitAt3, + focusNth, + unixEnvironment, + isEnvVarSet, + fancyTerminal, + runProgramWithExitCode, + directoryContents, + roundFloat, + ceilFloat, + floorFloat, + roundDouble, + ceilDouble, + floorDouble, + lgamma, + lgammaf, + tgamma, + tgammaf, + fromPOSIX, + toPOSIX, + trim, + pmapIO, + UserString, + EncodedString, + zEncodeString, + ) +where -import Numeric import Control.Concurrent import Control.Exception import Control.Monad import qualified Data.ByteString as BS -import qualified Data.Text as T -import qualified Data.Text.Encoding as T -import qualified Data.Text.Encoding.Error as T import Data.Char -import Data.List (foldl', genericDrop, genericSplitAt) import Data.Either +import Data.List (foldl', genericDrop, genericSplitAt) import Data.Maybe +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import qualified Data.Text.Encoding.Error as T +import Numeric +import qualified System.Directory.Tree as Dir import System.Environment +import System.Exit +import qualified System.FilePath as Native +import qualified System.FilePath.Posix as Posix import System.IO (hIsTerminalDevice, stdout) import System.IO.Unsafe -import qualified System.Directory.Tree as Dir import System.Process.ByteString -import System.Exit -import qualified System.FilePath.Posix as Posix -import qualified System.FilePath as Native -- | Like 'Data.Traversable.mapAccumL', but monadic. -mapAccumLM :: Monad m => - (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y]) +mapAccumLM :: + Monad m => + (acc -> x -> m (acc, y)) -> + acc -> + [x] -> + m (acc, [y]) mapAccumLM _ acc [] = return (acc, []) -mapAccumLM f acc (x:xs) = do +mapAccumLM f acc (x : xs) = do (acc', x') <- f acc x (acc'', xs') <- mapAccumLM f acc' xs - return (acc'', x':xs') + return (acc'', x' : xs') -- | @chunk n a@ splits @a@ into @n@-size-chunks. If the length of -- @a@ is not divisible by @n@, the last chunk will have fewer than @@ -75,17 +87,17 @@ chunk :: Int -> [a] -> [[a]] chunk _ [] = [] chunk n xs = - let (bef,aft) = splitAt n xs - in bef : chunk n aft + let (bef, aft) = splitAt n xs + in bef : chunk n aft -- | @chunks ns a@ splits @a@ into chunks determined by the elements -- of @ns@. It must hold that @sum ns == length a@, or the resulting -- list may contain too few chunks, or not all elements of @a@. chunks :: [Int] -> [a] -> [[a]] chunks [] _ = [] -chunks (n:ns) xs = - let (bef,aft) = splitAt n xs - in bef : chunks ns aft +chunks (n : ns) xs = + let (bef, aft) = splitAt n xs + in bef : chunks ns aft -- | Like 'maximum', but returns zero for an empty list. maxinum :: (Num a, Ord a, Foldable f) => f a -> a @@ -93,7 +105,7 @@ -- | @dropAt i n@ drops @n@ elements starting at element @i@. dropAt :: Int -> Int -> [a] -> [a] -dropAt i n xs = take i xs ++ drop (i+n) xs +dropAt i n xs = take i xs ++ drop (i + n) xs -- | @takeLast n l@ takes the last @n@ elements of @l@. takeLast :: Int -> [a] -> [a] @@ -110,13 +122,13 @@ -- | Return the list element at the given index, if the index is valid. maybeNth :: Integral int => int -> [a] -> Maybe a maybeNth i l - | i >= 0, v:_ <- genericDrop i l = Just v - | otherwise = Nothing + | i >= 0, v : _ <- genericDrop i l = Just v + | otherwise = Nothing -- | Return the first element of the list, if it exists. maybeHead :: [a] -> Maybe a maybeHead [] = Nothing -maybeHead (x:_) = Just x +maybeHead (x : _) = Just x -- | Like 'splitAt', but from the end. splitFromEnd :: Int -> [a] -> ([a], [a]) @@ -127,18 +139,19 @@ splitAt3 n m l = let (xs, l') = splitAt n l (ys, zs) = splitAt m l' - in (xs, ys, zs) + in (xs, ys, zs) -- | Return the list element at the given index, if the index is -- valid, along with the elements before and after. focusNth :: Integral int => int -> [a] -> Maybe ([a], a, [a]) focusNth i xs - | (bef, x:aft) <- genericSplitAt i xs = Just (bef, x, aft) - | otherwise = Nothing + | (bef, x : aft) <- genericSplitAt i xs = Just (bef, x, aft) + | otherwise = Nothing {-# NOINLINE unixEnvironment #-} + -- | The Unix environment when the Futhark compiler started. -unixEnvironment :: [(String,String)] +unixEnvironment :: [(String, String)] unixEnvironment = unsafePerformIO getEnvironment -- | Is an environment variable set to 0 or 1? If 0, return False; if @@ -152,6 +165,7 @@ _ -> Nothing {-# NOINLINE fancyTerminal #-} + -- | Are we running in a terminal capable of fancy commands and -- visualisation? fancyTerminal :: Bool @@ -163,14 +177,18 @@ -- | Like 'readProcessWithExitCode', but also wraps exceptions when -- the indicated binary cannot be launched, or some other exception is -- thrown. Also does shenanigans to handle improperly encoded outputs. -runProgramWithExitCode :: FilePath -> [String] -> BS.ByteString - -> IO (Either IOException (ExitCode, String, String)) +runProgramWithExitCode :: + FilePath -> + [String] -> + BS.ByteString -> + IO (Either IOException (ExitCode, String, String)) runProgramWithExitCode exe args inp = (Right . postprocess <$> readProcessWithExitCode exe args inp) - `catch` \e -> return (Left e) - where decode = T.unpack . T.decodeUtf8With T.lenientDecode - postprocess (code, out, err) = - (code, decode out, decode err) + `catch` \e -> return (Left e) + where + decode = T.unpack . T.decodeUtf8With T.lenientDecode + postprocess (code, out, err) = + (code, decode out, decode err) -- | Every non-directory file contained in a directory tree. directoryContents :: FilePath -> IO [FilePath] @@ -179,14 +197,20 @@ case Dir.failures tree of Dir.Failed _ err : _ -> throw err _ -> return $ mapMaybe isFile $ Dir.flattenDir tree - where isFile (Dir.File _ path) = Just path - isFile _ = Nothing + where + isFile (Dir.File _ path) = Just path + isFile _ = Nothing foreign import ccall "nearbyint" c_nearbyint :: Double -> Double + foreign import ccall "nearbyintf" c_nearbyintf :: Float -> Float + foreign import ccall "ceil" c_ceil :: Double -> Double + foreign import ccall "ceilf" c_ceilf :: Float -> Float + foreign import ccall "floor" c_floor :: Double -> Double + foreign import ccall "floorf" c_floorf :: Float -> Float -- | Round a single-precision floating point number correctly. @@ -214,8 +238,11 @@ floorDouble = c_floor foreign import ccall "lgamma" c_lgamma :: Double -> Double + foreign import ccall "lgammaf" c_lgammaf :: Float -> Float + foreign import ccall "tgamma" c_tgamma :: Double -> Double + foreign import ccall "tgammaf" c_tgammaf :: Float -> Float -- | The system-level @lgamma()@ function. @@ -250,10 +277,13 @@ trim = reverse . dropWhile isSpace . reverse . dropWhile isSpace fork :: (a -> IO b) -> a -> IO (MVar b) -fork f x = do cell <- newEmptyMVar - void $ forkIO $ do result <- f x - putMVar cell result - return cell +fork f x = do + cell <- newEmptyMVar + void $ + forkIO $ do + result <- f x + putMVar cell result + return cell -- | Run various 'IO' actions concurrently, possibly with a bound on -- the number of threads. @@ -263,8 +293,8 @@ go [] res = return res go xs res = do numThreads <- maybe getNumCapabilities pure concurrency - let (e,es) = splitAt numThreads xs - mvars <- mapM (fork f') e + let (e, es) = splitAt numThreads xs + mvars <- mapM (fork f') e result <- mapM takeMVar mvars case sequence result of Left err -> throw (err :: SomeException) @@ -290,56 +320,61 @@ -- programming languages. zEncodeString :: UserString -> EncodedString zEncodeString "" = "" -zEncodeString (c:cs) = encodeDigitChar c ++ concatMap encodeChar cs +zEncodeString (c : cs) = encodeDigitChar c ++ concatMap encodeChar cs -unencodedChar :: Char -> Bool -- True for chars that don't need encoding +unencodedChar :: Char -> Bool -- True for chars that don't need encoding unencodedChar 'Z' = False unencodedChar 'z' = False unencodedChar '_' = True -unencodedChar c = isAsciiLower c - || isAsciiUpper c - || isDigit c +unencodedChar c = + isAsciiLower c + || isAsciiUpper c + || isDigit c -- If a digit is at the start of a symbol then we need to encode it. -- Otherwise names like 9pH-0.1 give linker errors. encodeDigitChar :: Char -> EncodedString -encodeDigitChar c | isDigit c = encodeAsUnicodeCharar c - | otherwise = encodeChar c +encodeDigitChar c + | isDigit c = encodeAsUnicodeCharar c + | otherwise = encodeChar c encodeChar :: Char -> EncodedString -encodeChar c | unencodedChar c = [c] -- Common case first +encodeChar c | unencodedChar c = [c] -- Common case first -- Constructors -encodeChar '(' = "ZL" -- Needed for things like (,), and (->) -encodeChar ')' = "ZR" -- For symmetry with ( -encodeChar '[' = "ZM" -encodeChar ']' = "ZN" -encodeChar ':' = "ZC" -encodeChar 'Z' = "ZZ" - +encodeChar '(' = "ZL" -- Needed for things like (,), and (->) +encodeChar ')' = "ZR" -- For symmetry with ( +encodeChar '[' = "ZM" +encodeChar ']' = "ZN" +encodeChar ':' = "ZC" +encodeChar 'Z' = "ZZ" -- Variables -encodeChar 'z' = "zz" -encodeChar '&' = "za" -encodeChar '|' = "zb" -encodeChar '^' = "zc" -encodeChar '$' = "zd" -encodeChar '=' = "ze" -encodeChar '>' = "zg" -encodeChar '#' = "zh" -encodeChar '.' = "zi" -encodeChar '<' = "zl" -encodeChar '-' = "zm" -encodeChar '!' = "zn" -encodeChar '+' = "zp" +encodeChar 'z' = "zz" +encodeChar '&' = "za" +encodeChar '|' = "zb" +encodeChar '^' = "zc" +encodeChar '$' = "zd" +encodeChar '=' = "ze" +encodeChar '>' = "zg" +encodeChar '#' = "zh" +encodeChar '.' = "zi" +encodeChar '<' = "zl" +encodeChar '-' = "zm" +encodeChar '!' = "zn" +encodeChar '+' = "zp" encodeChar '\'' = "zq" encodeChar '\\' = "zr" -encodeChar '/' = "zs" -encodeChar '*' = "zt" -encodeChar '_' = "zu" -encodeChar '%' = "zv" -encodeChar c = encodeAsUnicodeCharar c +encodeChar '/' = "zs" +encodeChar '*' = "zt" +encodeChar '_' = "zu" +encodeChar '%' = "zv" +encodeChar c = encodeAsUnicodeCharar c encodeAsUnicodeCharar :: Char -> EncodedString -encodeAsUnicodeCharar c = 'z' : if isDigit (head hex_str) then hex_str - else '0':hex_str - where hex_str = showHex (ord c) "U" +encodeAsUnicodeCharar c = + 'z' : + if isDigit (head hex_str) + then hex_str + else '0' : hex_str + where + hex_str = showHex (ord c) "U" diff --git a/src/Futhark/Util/Console.hs b/src/Futhark/Util/Console.hs --- a/src/Futhark/Util/Console.hs +++ b/src/Futhark/Util/Console.hs @@ -1,11 +1,11 @@ -- | Some utility functions for working with pretty console output. module Futhark.Util.Console - ( color - , inRed - , inGreen - , inBold - ) - where + ( color, + inRed, + inGreen, + inBold, + ) +where import System.Console.ANSI diff --git a/src/Futhark/Util/IntegralExp.hs b/src/Futhark/Util/IntegralExp.hs --- a/src/Futhark/Util/IntegralExp.hs +++ b/src/Futhark/Util/IntegralExp.hs @@ -15,10 +15,10 @@ -- typeclasses that have been modified to make generic functions -- slightly easier to write. module Futhark.Util.IntegralExp - ( IntegralExp (..) - , Wrapped (..) - ) - where + ( IntegralExp (..), + Wrapped (..), + ) +where import Data.Int import Prelude @@ -38,22 +38,22 @@ divUp x y = (x + y - 1) `Futhark.Util.IntegralExp.div` y - fromInt8 :: Int8 -> e - fromInt16 :: Int16 -> e - fromInt32 :: Int32 -> e - fromInt64 :: Int64 -> e - -- | This wrapper allows you to use a type that is an instance of the -- true class whenever the simile class is required. -newtype Wrapped a = Wrapped { wrappedValue :: a } - deriving (Eq, Ord, Show) +newtype Wrapped a = Wrapped {wrappedValue :: a} + deriving (Eq, Ord, Show) -liftOp :: (a -> a) - -> Wrapped a -> Wrapped a +liftOp :: + (a -> a) -> + Wrapped a -> + Wrapped a liftOp op (Wrapped x) = Wrapped $ op x -liftOp2 :: (a -> a -> a) - -> Wrapped a -> Wrapped a -> Wrapped a +liftOp2 :: + (a -> a -> a) -> + Wrapped a -> + Wrapped a -> + Wrapped a liftOp2 op (Wrapped x) (Wrapped y) = Wrapped $ x `op` y instance Num a => Num (Wrapped a) where @@ -71,8 +71,3 @@ div = liftOp2 Prelude.div mod = liftOp2 Prelude.mod sgn = Just . fromIntegral . signum . toInteger . wrappedValue - - fromInt8 = fromInteger . toInteger - fromInt16 = fromInteger . toInteger - fromInt32 = fromInteger . toInteger - fromInt64 = fromInteger . toInteger diff --git a/src/Futhark/Util/Loc.hs b/src/Futhark/Util/Loc.hs --- a/src/Futhark/Util/Loc.hs +++ b/src/Futhark/Util/Loc.hs @@ -1,5 +1,6 @@ {-# LANGUAGE Trustworthy #-} + -- | A Safe Haskell-trusted re-export of the @srcloc@ package. -module Futhark.Util.Loc ( module Data.Loc ) where +module Futhark.Util.Loc (module Data.Loc) where import Data.Loc diff --git a/src/Futhark/Util/Log.hs b/src/Futhark/Util/Log.hs --- a/src/Futhark/Util/Log.hs +++ b/src/Futhark/Util/Log.hs @@ -1,25 +1,25 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE OverloadedStrings #-} + -- | Opaque type for an operations log that provides fast O(1) -- appends. module Futhark.Util.Log - ( Log - , toText - , ToLog (..) - , MonadLogger (..) - ) - + ( Log, + toText, + ToLog (..), + MonadLogger (..), + ) where -import Control.Monad.Writer -import qualified Control.Monad.RWS.Strict import qualified Control.Monad.RWS.Lazy -import qualified Data.Text as T +import qualified Control.Monad.RWS.Strict +import Control.Monad.Writer import qualified Data.DList as DL +import qualified Data.Text as T -- | An efficiently catenable sequence of log entries. -newtype Log = Log { unLog :: DL.DList T.Text } +newtype Log = Log {unLog :: DL.DList T.Text} instance Semigroup Log where Log l1 <> Log l2 = Log $ l1 <> l2 diff --git a/src/Futhark/Util/Options.hs b/src/Futhark/Util/Options.hs --- a/src/Futhark/Util/Options.hs +++ b/src/Futhark/Util/Options.hs @@ -1,16 +1,16 @@ -- | Common code for parsing command line options based on getopt. module Futhark.Util.Options - ( FunOptDescr - , mainWithOptions - , commonOptions - ) where + ( FunOptDescr, + mainWithOptions, + commonOptions, + ) +where import Control.Monad.IO.Class -import System.IO -import System.Exit -import System.Console.GetOpt - import Futhark.Version +import System.Console.GetOpt +import System.Exit +import System.IO -- | A command line option that either purely updates a configuration, -- or performs an IO action (and stops). @@ -18,30 +18,34 @@ -- | Generate a main action that parses the given command line options -- (while always adding 'commonOptions'). -mainWithOptions :: cfg - -> [FunOptDescr cfg] - -> String - -> ([String] -> cfg -> Maybe (IO ())) - -> String - -> [String] - -> IO () +mainWithOptions :: + cfg -> + [FunOptDescr cfg] -> + String -> + ([String] -> cfg -> Maybe (IO ())) -> + String -> + [String] -> + IO () mainWithOptions emptyConfig commandLineOptions usage f prog args = case getOpt' Permute commandLineOptions' args of (opts, nonopts, [], []) -> case applyOpts opts of Right config | Just m <- f nonopts config -> m - | otherwise -> invalid nonopts [] [] - Left m -> m + | otherwise -> invalid nonopts [] [] + Left m -> m (_, nonopts, unrecs, errs) -> invalid nonopts unrecs errs - where applyOpts opts = do fs <- sequence opts - return $ foldl (.) id (reverse fs) emptyConfig + where + applyOpts opts = do + fs <- sequence opts + return $ foldl (.) id (reverse fs) emptyConfig - invalid nonopts unrecs errs = do help <- helpStr prog usage commandLineOptions' - badOptions help nonopts errs unrecs + invalid nonopts unrecs errs = do + help <- helpStr prog usage commandLineOptions' + badOptions help nonopts errs unrecs - commandLineOptions' = - commonOptions prog usage commandLineOptions ++ commandLineOptions + commandLineOptions' = + commonOptions prog usage commandLineOptions ++ commandLineOptions helpStr :: String -> String -> [OptDescr a] -> IO String helpStr prog usage opts = do @@ -63,20 +67,30 @@ -- options. commonOptions :: String -> String -> [FunOptDescr cfg] -> [FunOptDescr cfg] commonOptions prog usage options = - [ Option "V" ["version"] - (NoArg $ Left $ do header - exitSuccess) - "Print version information and exit." - - , Option "h" ["help"] - (NoArg $ Left $ do header - putStrLn "" - putStrLn =<< helpStr prog usage (commonOptions prog usage [] ++ options) - exitSuccess) - "Print help and exit." + [ Option + "V" + ["version"] + ( NoArg $ + Left $ do + header + exitSuccess + ) + "Print version information and exit.", + Option + "h" + ["help"] + ( NoArg $ + Left $ do + header + putStrLn "" + putStrLn =<< helpStr prog usage (commonOptions prog usage [] ++ options) + exitSuccess + ) + "Print help and exit." ] - where header = do - putStrLn $ "Futhark " ++ versionString - putStrLn "Copyright (C) DIKU, University of Copenhagen, released under the ISC license." - putStrLn "This is free software: you are free to change and redistribute it." - putStrLn "There is NO WARRANTY, to the extent permitted by law." + where + header = do + putStrLn $ "Futhark " ++ versionString + putStrLn "Copyright (C) DIKU, University of Copenhagen, released under the ISC license." + putStrLn "This is free software: you are free to change and redistribute it." + putStrLn "There is NO WARRANTY, to the extent permitted by law." diff --git a/src/Futhark/Util/Pretty.hs b/src/Futhark/Util/Pretty.hs --- a/src/Futhark/Util/Pretty.hs +++ b/src/Futhark/Util/Pretty.hs @@ -1,30 +1,29 @@ {-# LANGUAGE Trustworthy #-} {-# OPTIONS_GHC -fno-warn-orphans #-} + -- | A re-export of the prettyprinting library, along with some convenience functions. module Futhark.Util.Pretty - ( module Text.PrettyPrint.Mainland - , module Text.PrettyPrint.Mainland.Class - , pretty - , prettyDoc - , prettyTuple - , prettyText - , prettyOneLine - - , apply - , oneLine - , annot - , nestedBlock - , textwrap - , shorten - ) - where + ( module Text.PrettyPrint.Mainland, + module Text.PrettyPrint.Mainland.Class, + pretty, + prettyDoc, + prettyTuple, + prettyText, + prettyOneLine, + apply, + oneLine, + annot, + nestedBlock, + textwrap, + shorten, + ) +where import Data.Text (Text) import qualified Data.Text.Lazy as LT - import Text.PrettyPrint.Mainland hiding (pretty) -import Text.PrettyPrint.Mainland.Class import qualified Text.PrettyPrint.Mainland as PP +import Text.PrettyPrint.Mainland.Class -- | Prettyprint a value, wrapped to 80 characters. pretty :: Pretty a => a -> String @@ -36,7 +35,7 @@ -- | Prettyprint a value without any width restriction. prettyOneLine :: Pretty a => a -> String -prettyOneLine = ($"") . displayS . renderCompact . oneLine . ppr +prettyOneLine = ($ "") . displayS . renderCompact . oneLine . ppr -- | Re-export of 'PP.pretty'. prettyDoc :: Int -> Doc -> String @@ -72,14 +71,17 @@ -- | Surround the given document with enclosers and add linebreaks and -- indents. nestedBlock :: String -> String -> Doc -> Doc -nestedBlock pre post body = text pre - PP.indent 2 body - text post +nestedBlock pre post body = + text pre + PP.indent 2 body + text post -- | Prettyprint on a single line up to at most some appropriate -- number of characters, with trailing ... if necessary. Used for -- error messages. shorten :: Pretty a => a -> Doc -shorten a | length s > 70 = text (take 70 s) <> text "..." - | otherwise = text s - where s = pretty a +shorten a + | length s > 70 = text (take 70 s) <> text "..." + | otherwise = text s + where + s = pretty a diff --git a/src/Futhark/Util/Table.hs b/src/Futhark/Util/Table.hs --- a/src/Futhark/Util/Table.hs +++ b/src/Futhark/Util/Table.hs @@ -1,15 +1,15 @@ -- | Basic table building for prettier futhark-test output. module Futhark.Util.Table - ( buildTable - , mkEntry - , Entry - ) where + ( buildTable, + mkEntry, + Entry, + ) +where import Data.List (intercalate, transpose) -import System.Console.ANSI - import Futhark.Util (maxinum) import Futhark.Util.Console (color) +import System.Console.ANSI data RowTemplate = RowTemplate [Int] Int deriving (Show) @@ -22,32 +22,35 @@ mkEntry s = (s, []) buildRowTemplate :: [[Entry]] -> Int -> RowTemplate - buildRowTemplate rows = RowTemplate widths - where widths = map (maxinum . map (length . fst)) . transpose $ rows + where + widths = map (maxinum . map (length . fst)) . transpose $ rows buildRow :: RowTemplate -> [Entry] -> String buildRow (RowTemplate widths pad) entries = cells ++ "\n" - where bar = "\x2502" - cells = concatMap buildCell (zip entries widths) ++ bar - buildCell ((entry, sgr), width) = - let padding = width - length entry + pad - in bar ++ " " ++ color sgr entry ++ replicate padding ' ' + where + bar = "\x2502" + cells = concatMap buildCell (zip entries widths) ++ bar + buildCell ((entry, sgr), width) = + let padding = width - length entry + pad + in bar ++ " " ++ color sgr entry ++ replicate padding ' ' buildSep :: Char -> Char -> Char -> RowTemplate -> String buildSep lCorner rCorner sep (RowTemplate widths pad) = corners . concatMap cellFloor $ widths - where cellFloor width = replicate (width + pad + 1) '\x2500' ++ [sep] - corners [] = "" - corners s = [lCorner] ++ init s ++ [rCorner] + where + cellFloor width = replicate (width + pad + 1) '\x2500' ++ [sep] + corners [] = "" + corners s = [lCorner] ++ init s ++ [rCorner] -- | Builds a table from a list of entries and a padding amount that -- determines padding from the right side of the widest entry in each column. buildTable :: [[Entry]] -> Int -> String buildTable rows pad = buildTop template ++ sepRows ++ buildBottom template - where sepRows = intercalate (buildFloor template) builtRows - builtRows = map (buildRow template) rows - template = buildRowTemplate rows pad - buildTop rt = buildSep '\x250C' '\x2510' '\x252C' rt ++ "\n" - buildFloor rt = buildSep '\x251C' '\x2524' '\x253C' rt ++ "\n" - buildBottom = buildSep '\x2514' '\x2518' '\x2534' + where + sepRows = intercalate (buildFloor template) builtRows + builtRows = map (buildRow template) rows + template = buildRowTemplate rows pad + buildTop rt = buildSep '\x250C' '\x2510' '\x252C' rt ++ "\n" + buildFloor rt = buildSep '\x251C' '\x2524' '\x253C' rt ++ "\n" + buildBottom = buildSep '\x2514' '\x2518' '\x2534' diff --git a/src/Futhark/Version.hs b/src/Futhark/Version.hs --- a/src/Futhark/Version.hs +++ b/src/Futhark/Version.hs @@ -1,17 +1,16 @@ {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE Trustworthy #-} + -- | This module exports version information about the Futhark -- compiler. module Futhark.Version - ( - version - , versionString - ) - where + ( version, + versionString, + ) +where import Data.Version import Development.GitRev - import qualified Paths_futhark -- | The version of Futhark that we are using. This is equivalent to @@ -21,20 +20,27 @@ -- | The version of Futhark that we are using, as a 'String' versionString :: String -versionString = showVersion version ++ - if used_hash /= "UNKNOWN" - then "\n" ++ gitversion - else "" +versionString = + showVersion version + ++ if used_hash /= "UNKNOWN" + then "\n" ++ gitversion + else "" where used_hash = take 7 $(gitHash) - gitversion = concat ["git: " - , branch - , used_hash - , " (", $(gitCommitDate), ")" - , dirty - ] - branch | $(gitBranch) == "master" = "" - | otherwise = $(gitBranch) ++ " @ " - dirty | $(gitDirtyTracked) = " [modified]" - | otherwise = "" + gitversion = + concat + [ "git: ", + branch, + used_hash, + " (", + $(gitCommitDate), + ")", + dirty + ] + branch + | $(gitBranch) == "master" = "" + | otherwise = $(gitBranch) ++ " @ " + dirty + | $(gitDirtyTracked) = " [modified]" + | otherwise = "" diff --git a/src/Language/Futhark.hs b/src/Language/Futhark.hs --- a/src/Language/Futhark.hs +++ b/src/Language/Futhark.hs @@ -1,22 +1,35 @@ {-# LANGUAGE Safe #-} + -- | Re-export the external Futhark modules for convenience. module Language.Futhark - ( module Language.Futhark.Syntax - , module Language.Futhark.Prop - , module Language.Futhark.Pretty - - , Ident, DimIndex, Exp, Pattern - , ModExp, ModParam, SigExp, ModBind, SigBind - , ValBind, Dec, Spec, Prog - , TypeBind, TypeDecl - , StructTypeArg, ScalarType - , TypeParam, Case + ( module Language.Futhark.Syntax, + module Language.Futhark.Prop, + module Language.Futhark.Pretty, + Ident, + DimIndex, + Exp, + Pattern, + ModExp, + ModParam, + SigExp, + ModBind, + SigBind, + ValBind, + Dec, + Spec, + Prog, + TypeBind, + TypeDecl, + StructTypeArg, + ScalarType, + TypeParam, + Case, ) - where +where -import Language.Futhark.Syntax -import Language.Futhark.Prop import Language.Futhark.Pretty +import Language.Futhark.Prop +import Language.Futhark.Syntax -- | An identifier with type- and aliasing information. type Ident = IdentBase Info VName diff --git a/src/Language/Futhark/Core.hs b/src/Language/Futhark/Core.hs --- a/src/Language/Futhark/Core.hs +++ b/src/Language/Futhark/Core.hs @@ -1,61 +1,84 @@ -{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Trustworthy #-} + -- | This module contains very basic definitions for Futhark - so basic, -- that they can be shared between the internal and external -- representation. module Language.Futhark.Core - ( Uniqueness(..) - , Commutativity(..) + ( Uniqueness (..), + Commutativity (..), - -- * Location utilities - , SrcLoc - , Loc - , Located(..) - , srclocOf - , locStr - , locStrRel - , prettyStacktrace + -- * Location utilities + SrcLoc, + Loc, + Located (..), + srclocOf, + locStr, + locStrRel, + prettyStacktrace, - -- * Name handling - , Name - , nameToString - , nameFromString - , nameToText - , nameFromText - , VName(..) - , baseTag - , baseName - , baseString - , pretty - , quote - , pquote + -- * Name handling + Name, + nameToString, + nameFromString, + nameToText, + nameFromText, + VName (..), + baseTag, + baseName, + baseString, + pretty, + quote, + pquote, - -- * Special identifiers - , defaultEntryPoint + -- * Special identifiers + defaultEntryPoint, -- * Integer re-export - , Int8, Int16, Int32, Int64 - , Word8, Word16, Word32, Word64 + Int8, + Int16, + Int32, + Int64, + Word8, + Word16, + Word32, + Word64, ) - where -import Data.Int (Int8, Int16, Int32, Int64) +import Control.Category +import Data.Int (Int16, Int32, Int64, Int8) import Data.String -import Data.Word (Word8, Word16, Word32, Word64) import qualified Data.Text as T - -import Futhark.Util.Pretty +import Data.Word (Word16, Word32, Word64, Word8) import Futhark.Util.Loc +import Futhark.Util.Pretty +import GHC.Generics +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Text.Read +import Prelude hiding (id, (.)) -- | The uniqueness attribute of a type. This essentially indicates -- whether or not in-place modifications are acceptable. With respect -- to ordering, 'Unique' is greater than 'Nonunique'. -data Uniqueness = Nonunique -- ^ May have references outside current function. - | Unique -- ^ No references outside current function. - deriving (Eq, Ord, Show) +data Uniqueness + = -- | May have references outside current function. + Nonunique + | -- | No references outside current function. + Unique + deriving (Eq, Ord, Show, Generic) +instance SexpIso Uniqueness where + sexpIso = + match $ + With (. Sexp.sym "nonunique") $ + With + (. Sexp.sym "unique") + End + instance Semigroup Uniqueness where (<>) = min @@ -68,10 +91,19 @@ -- | Whether some operator is commutative or not. The 'Monoid' -- instance returns the least commutative of its arguments. -data Commutativity = Noncommutative - | Commutative - deriving (Eq, Ord, Show) +data Commutativity + = Noncommutative + | Commutative + deriving (Eq, Ord, Show, Generic) +instance SexpIso Commutativity where + sexpIso = + match $ + With (. Sexp.sym "noncommutative") $ + With + (. Sexp.sym "commutative") + End + instance Semigroup Commutativity where (<>) = min @@ -86,8 +118,11 @@ -- compiler. 'String's, being lists of characters, are very slow, -- while 'T.Text's are based on byte-arrays. newtype Name = Name T.Text - deriving (Show, Eq, Ord, IsString, Semigroup) + deriving (Show, Eq, Ord, IsString, Semigroup, Generic) +instance SexpIso Name where + sexpIso = with (. symbol) + instance Pretty Name where ppr = text . nameToString @@ -119,13 +154,13 @@ case locOf a of NoLoc -> "unknown location" Loc (Pos file line1 col1 _) (Pos _ line2 col2 _) - -- Do not show line2 if it is identical to line1. + -- Do not show line2 if it is identical to line1. | line1 == line2 -> - first_part ++ "-" ++ show col2 + first_part ++ "-" ++ show col2 | otherwise -> - first_part ++ "-" ++ show line2 ++ ":" ++ show col2 - where first_part = file ++ ":" ++ show line1 ++ ":" ++ show col1 - + first_part ++ "-" ++ show line2 ++ ":" ++ show col2 + where + first_part = file ++ ":" ++ show line1 ++ ":" ++ show col1 -- | Like 'locStr', but @locStrRel prev now@ prints the location @now@ -- with the file name left out if the same as @prev@. This is useful @@ -138,10 +173,11 @@ (Loc (Pos a_file _ _ _) _, Loc (Pos b_file line1 col1 _) (Pos _ line2 col2 _)) | a_file == b_file, line1 == line2 -> - first_part ++ "-" ++ show col2 + first_part ++ "-" ++ show col2 | a_file == b_file -> - first_part ++ "-" ++ show line2 ++ ":" ++ show col2 - where first_part = show line1 ++ ":" ++ show col1 + first_part ++ "-" ++ show line2 ++ ":" ++ show col2 + where + first_part = show line1 ++ ":" ++ show col1 _ -> locStr b -- | Given a list of strings representing entries in the stack trace @@ -150,18 +186,39 @@ -- should also be preceded by a newline. The most recent stack frame -- must come first in the list. prettyStacktrace :: Int -> [String] -> String -prettyStacktrace cur = unlines . zipWith f [(0::Int)..] - where -- Formatting hack: assume no stack is deeper than 100 - -- elements. Since Futhark does not support recursion, going - -- beyond that would require a truly perverse program. - f i x = (if cur == i then "-> " else " ") ++ - '#' : show i ++ - (if i > 9 then "" else " ") ++ " " ++ x +prettyStacktrace cur = unlines . zipWith f [(0 :: Int) ..] + where + -- Formatting hack: assume no stack is deeper than 100 + -- elements. Since Futhark does not support recursion, going + -- beyond that would require a truly perverse program. + f i x = + (if cur == i then "-> " else " ") + ++ '#' : + show i + ++ (if i > 9 then "" else " ") + ++ " " + ++ x -- | A name tagged with some integer. Only the integer is used in -- comparisons, no matter the type of @vn@. data VName = VName !Name !Int - deriving (Show) + deriving (Show, Generic) + +instance SexpIso VName where + sexpIso = with $ \vname -> + Sexp.symbol + >>> flipped + ( pair + >>> partialIso + (\(nm, i) -> T.pack $ nameToString nm ++ "_" ++ show i) + ( \s -> + let (nm, i) = T.breakOnEnd "_" s + in case readMaybe $ T.unpack i of + Just i' -> Right (nameFromString $ T.unpack $ T.init nm, i') + Nothing -> Left $ expected "Couldn't parse int of vname" + ) + ) + >>> vname -- | Return the tag contained in the 'VName'. baseTag :: VName -> Int 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 @@ -1,1535 +1,1713 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE DeriveTraversable #-} -{-# LANGUAGE OverloadedStrings #-} --- | An interpreter operating on type-checked source Futhark terms. --- Relatively slow. -module Language.Futhark.Interpreter - ( Ctx(..) - , Env - , InterpreterError - , initialCtx - , interpretExp - , interpretDec - , interpretImport - , interpretFunction - , ExtOp(..) - , BreakReason(..) - , StackFrame(..) - , typeCheckerEnv - , Value (ValuePrim, ValueArray, ValueRecord) - , fromTuple - , isEmptyArray - , prettyEmptyArray - ) where - -import Control.Monad.Trans.Maybe -import Control.Monad.Free.Church -import Control.Monad.Except -import Control.Monad.State -import Control.Monad.Reader -import Data.Array -import Data.Bifunctor (first) -import Data.List - (transpose, genericLength, isPrefixOf, foldl', find, intercalate) -import Data.Maybe -import qualified Data.Map as M -import qualified Data.List.NonEmpty as NE -import Data.Monoid hiding (Sum) - -import Language.Futhark hiding (Value, matchDims) -import qualified Language.Futhark as F -import Futhark.IR.Primitive (intValue, floatValue) -import qualified Futhark.IR.Primitive as P -import qualified Language.Futhark.Semantic as T - -import Futhark.Util.Pretty hiding (apply, bool) -import Futhark.Util (chunk, splitFromEnd, maybeHead) -import Futhark.Util.Loc - -import Prelude hiding (mod, break) - -data StackFrame = StackFrame { stackFrameLoc :: Loc - , stackFrameCtx :: Ctx - } - -instance Located StackFrame where - locOf = stackFrameLoc - --- | What is the reason for this break point? -data BreakReason - = BreakPoint -- ^ An explicit breakpoint in the program. - | BreakNaN -- ^ A - -data ExtOp a = ExtOpTrace Loc String a - | ExtOpBreak BreakReason (NE.NonEmpty StackFrame) a - | ExtOpError InterpreterError - -instance Functor ExtOp where - fmap f (ExtOpTrace w s x) = ExtOpTrace w s $ f x - fmap f (ExtOpBreak why backtrace x) = ExtOpBreak why backtrace $ f x - fmap _ (ExtOpError err) = ExtOpError err - -type Stack = [StackFrame] - -type Sizes = M.Map VName Int32 - --- | The monad in which evaluation takes place. -newtype EvalM a = EvalM (ReaderT (Stack, M.Map FilePath Env) - (StateT Sizes (F ExtOp)) a) - deriving (Monad, Applicative, Functor, - MonadFree ExtOp, - MonadReader (Stack, M.Map FilePath Env), - MonadState Sizes) - -runEvalM :: M.Map FilePath Env -> EvalM a -> F ExtOp a -runEvalM imports (EvalM m) = evalStateT (runReaderT m (mempty, imports)) mempty - -stacking :: SrcLoc -> Env -> EvalM a -> EvalM a -stacking loc env = local $ \(ss, imports) -> - if isNoLoc loc - then (ss, imports) - else let s = StackFrame (locOf loc) (Ctx env imports) - in (s:ss, imports) - where isNoLoc :: SrcLoc -> Bool - isNoLoc = (==NoLoc) . locOf - -stacktrace :: EvalM [Loc] -stacktrace = asks $ map stackFrameLoc . fst - -lookupImport :: FilePath -> EvalM (Maybe Env) -lookupImport f = asks $ M.lookup f . snd - -putExtSize :: VName -> Int32 -> EvalM () -putExtSize v x = modify $ M.insert v x - -getSizes :: EvalM Sizes -getSizes = get - -extSizeEnv :: EvalM Env -extSizeEnv = i32Env <$> getSizes - -prettyRecord :: Pretty a => M.Map Name a -> Doc -prettyRecord m - | Just vs <- areTupleFields m = - parens $ commasep $ map ppr vs - | otherwise = - braces $ commasep $ map field $ M.toList m - where field (k, v) = ppr k <+> equals <+> ppr v - -valueStructType :: ValueType -> StructType -valueStructType = first (ConstDim . fromIntegral) - --- | A shape is a tree to accomodate the case of records. It is --- parameterised over the representation of dimensions. -data Shape d = ShapeDim d (Shape d) - | ShapeLeaf - | ShapeRecord (M.Map Name (Shape d)) - | ShapeSum (M.Map Name [Shape d]) - deriving (Eq, Show, Functor, Foldable, Traversable) - -type ValueShape = Shape Int32 - -instance Pretty d => Pretty (Shape d) where - ppr ShapeLeaf = mempty - ppr (ShapeDim d s) = brackets (ppr d) <> ppr s - ppr (ShapeRecord m) = prettyRecord m - ppr (ShapeSum cs) = - mconcat (punctuate (text " | ") cs') - where ppConstr (name, fs) = sep $ (text "#" <> ppr name) : map ppr fs - cs' = map ppConstr $ M.toList cs - -emptyShape :: ValueShape -> Bool -emptyShape (ShapeDim d s) = d == 0 || emptyShape s -emptyShape _ = False - -typeShape :: M.Map VName (Shape d) -> TypeBase d () -> Shape d -typeShape shapes = go - where go (Array _ _ et shape) = - foldr ShapeDim (go (Scalar et)) $ shapeDims shape - go (Scalar (Record fs)) = - ShapeRecord $ M.map go fs - go (Scalar (Sum cs)) = - ShapeSum $ M.map (map go) cs - go (Scalar (TypeVar _ _ (TypeName [] v) [])) - | Just shape <- M.lookup v shapes = - shape - go _ = - ShapeLeaf - -structTypeShape :: M.Map VName ValueShape -> StructType -> Shape (Maybe Int32) -structTypeShape shapes = fmap dim . typeShape shapes' - where dim (ConstDim d) = Just $ fromIntegral d - dim _ = Nothing - shapes' = M.map (fmap $ ConstDim . fromIntegral) shapes - -resolveTypeParams :: [VName] -> StructType -> StructType -> Env -resolveTypeParams names = match - where match (Scalar (TypeVar _ _ tn _)) t - | typeLeaf tn `elem` names = - typeEnv $ M.singleton (typeLeaf tn) t - match (Scalar (Record poly_fields)) (Scalar (Record fields)) = - mconcat $ M.elems $ - M.intersectionWith match poly_fields fields - match (Scalar (Sum poly_fields)) (Scalar (Sum fields)) = - mconcat $ map mconcat $ M.elems $ - M.intersectionWith (zipWith match) poly_fields fields - match (Scalar (Arrow _ _ poly_t1 poly_t2)) (Scalar (Arrow _ _ t1 t2)) = - match poly_t1 t1 <> match poly_t2 t2 - match poly_t t - | d1 : _ <- shapeDims (arrayShape poly_t), - d2 : _ <- shapeDims (arrayShape t) = - matchDims d1 d2 <> match (stripArray 1 poly_t) (stripArray 1 t) - match _ _ = mempty - - matchDims (NamedDim (QualName _ d1)) (ConstDim d2) - | d1 `elem` names = - i32Env $ M.singleton d1 $ fromIntegral d2 - matchDims _ _ = mempty - -resolveExistentials :: [VName] -> StructType -> ValueShape -> M.Map VName Int32 -resolveExistentials names = match - where match (Scalar (Record poly_fields)) (ShapeRecord fields) = - mconcat $ M.elems $ - M.intersectionWith match poly_fields fields - match (Scalar (Sum poly_fields)) (ShapeSum fields) = - mconcat $ map mconcat $ M.elems $ - M.intersectionWith (zipWith match) poly_fields fields - match poly_t (ShapeDim d2 rowshape) - | d1 : _ <- shapeDims (arrayShape poly_t) = - matchDims d1 d2 <> match (stripArray 1 poly_t) rowshape - match _ _ = mempty - - matchDims (NamedDim (QualName _ d1)) d2 - | d1 `elem` names = M.singleton d1 d2 - matchDims _ _ = mempty - --- | A fully evaluated Futhark value. -data Value = ValuePrim !PrimValue - | ValueArray ValueShape !(Array Int Value) - -- Stores the full shape. - | ValueRecord (M.Map Name Value) - | ValueFun (Value -> EvalM Value) - | ValueSum ValueShape Name [Value] - -- Stores the full shape. - -instance Eq Value where - ValuePrim x == ValuePrim y = x == y - ValueArray _ x == ValueArray _ y = x == y - ValueRecord x == ValueRecord y = x == y - ValueSum _ n1 vs1 == ValueSum _ n2 vs2 = n1 == n2 && vs1 == vs2 - _ == _ = False - -instance Pretty Value where - ppr (ValuePrim v) = ppr v - ppr (ValueArray _ a) = - let elements = elems a -- [Value] - (x:_) = elements - separator = case x of - ValueArray _ _ -> comma <> line - _ -> comma <> space - in brackets $ cat $ punctuate separator (map ppr elements) - - ppr (ValueRecord m) = prettyRecord m - ppr ValueFun{} = text "#" - ppr (ValueSum _ n vs) = text "#" <> sep (ppr n : map ppr vs) - -valueShape :: Value -> ValueShape -valueShape (ValueArray shape _) = shape -valueShape (ValueRecord fs) = ShapeRecord $ M.map valueShape fs -valueShape (ValueSum shape _ _) = shape -valueShape _ = ShapeLeaf - -checkShape :: Shape (Maybe Int32) -> ValueShape -> Maybe ValueShape -checkShape (ShapeDim Nothing shape1) (ShapeDim d2 shape2) = - ShapeDim d2 <$> checkShape shape1 shape2 -checkShape (ShapeDim (Just d1) shape1) (ShapeDim d2 shape2) = do - guard $ d1 == d2 - ShapeDim d2 <$> checkShape shape1 shape2 -checkShape (ShapeDim d1 shape1) ShapeLeaf = - -- This case is for handling polymorphism, when a function doesn't - -- know that the array it produced actually has more dimensions. - ShapeDim (fromMaybe 0 d1) <$> checkShape shape1 ShapeLeaf -checkShape (ShapeRecord shapes1) (ShapeRecord shapes2) = - ShapeRecord <$> sequence (M.intersectionWith checkShape shapes1 shapes2) -checkShape (ShapeRecord shapes1) ShapeLeaf = - Just $ fromMaybe 0 <$> ShapeRecord shapes1 -checkShape (ShapeSum shapes1) (ShapeSum shapes2) = - ShapeSum <$> sequence (M.intersectionWith (zipWithM checkShape) shapes1 shapes2) -checkShape (ShapeSum shapes1) ShapeLeaf = - Just $ fromMaybe 0 <$> ShapeSum shapes1 -checkShape _ shape2 = - Just shape2 - --- | Does the value correspond to an empty array? -isEmptyArray :: Value -> Bool -isEmptyArray = emptyShape . valueShape - --- | String representation of an empty array with the provided element --- type. This is pretty ad-hoc - don't expect good results unless the --- element type is a primitive. -prettyEmptyArray :: TypeBase () () -> Value -> String -prettyEmptyArray t v = - "empty(" ++ dims (valueShape v) ++ pretty t' ++ ")" - where t' = stripArray (arrayRank t) t - dims (ShapeDim n rowshape) = - "[" ++ pretty n ++ "]" ++ dims rowshape - dims _ = "" - --- | Create an array value; failing if that would result in an --- irregular array. -mkArray :: TypeBase Int32 () -> [Value] -> Maybe Value -mkArray t [] = - return $ toArray (typeShape mempty t) [] -mkArray _ (v:vs) = do - let v_shape = valueShape v - guard $ all ((==v_shape) . valueShape) vs - return $ toArray' v_shape $ v:vs - -arrayLength :: Integral int => Array Int Value -> int -arrayLength = fromIntegral . (+1) . snd . bounds - -toTuple :: [Value] -> Value -toTuple = ValueRecord . M.fromList . zip tupleFieldNames - -fromTuple :: Value -> Maybe [Value] -fromTuple (ValueRecord m) = areTupleFields m -fromTuple _ = Nothing - -asInteger :: Value -> Integer -asInteger (ValuePrim (SignedValue v)) = P.valueIntegral v -asInteger (ValuePrim (UnsignedValue v)) = - toInteger (P.valueIntegral (P.doZExt v Int64) :: Word64) -asInteger v = error $ "Unexpectedly not an integer: " ++ pretty v - -asInt :: Value -> Int -asInt = fromIntegral . asInteger - -asSigned :: Value -> IntValue -asSigned (ValuePrim (SignedValue v)) = v -asSigned v = error $ "Unexpected not a signed integer: " ++ pretty v - -asInt32 :: Value -> Int32 -asInt32 = fromIntegral . asInteger - -asBool :: Value -> Bool -asBool (ValuePrim (BoolValue x)) = x -asBool v = error $ "Unexpectedly not a boolean: " ++ pretty v - -lookupInEnv :: (Env -> M.Map VName x) - -> QualName VName -> Env -> Maybe x -lookupInEnv onEnv qv env = f env $ qualQuals qv - where f m (q:qs) = - case M.lookup q $ envTerm m of - Just (TermModule (Module mod)) -> f mod qs - _ -> Nothing - f m [] = M.lookup (qualLeaf qv) $ onEnv m - -lookupVar :: QualName VName -> Env -> Maybe TermBinding -lookupVar = lookupInEnv envTerm - -lookupType :: QualName VName -> Env -> Maybe T.TypeBinding -lookupType = lookupInEnv envType - --- | A TermValue with a 'Nothing' type annotation is an intrinsic. -data TermBinding = TermValue (Maybe T.BoundV) Value - | TermPoly (Maybe T.BoundV) (StructType -> EvalM Value) - -- ^ A polymorphic value that must be instantiated. - | TermModule Module - -data Module = Module Env - | ModuleFun (Module -> EvalM Module) - --- | The actual type- and value environment. -data Env = Env { envTerm :: M.Map VName TermBinding - , envType :: M.Map VName T.TypeBinding - , envShapes :: M.Map VName ValueShape - -- ^ A mapping from type parameters to the shapes of - -- the value to which they were initially bound. - } - -instance Monoid Env where - mempty = Env mempty mempty mempty - -instance Semigroup Env where - Env vm1 tm1 sm1 <> Env vm2 tm2 sm2 = - Env (vm1 <> vm2) (tm1 <> tm2) (sm1 <> sm2) - --- | An error occurred during interpretation due to an error in the --- user program. Actual interpreter errors will be signaled with an --- IO exception ('error'). -newtype InterpreterError = InterpreterError String - -valEnv :: M.Map VName (Maybe T.BoundV, Value) -> Env -valEnv m = Env { envTerm = M.map (uncurry TermValue) m - , envType = mempty - , envShapes = mempty - } - -modEnv :: M.Map VName Module -> Env -modEnv m = Env { envTerm = M.map TermModule m - , envType = mempty - , envShapes = mempty - } - -typeEnv :: M.Map VName StructType -> Env -typeEnv m = Env { envTerm = mempty - , envType = M.map tbind m - , envShapes = mempty - } - where tbind = T.TypeAbbr Unlifted [] - -i32Env :: M.Map VName Int32 -> Env -i32Env = valEnv . M.map f - where f x = (Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int32, - ValuePrim $ SignedValue $ Int32Value x) - -instance Show InterpreterError where - show (InterpreterError s) = s - -bad :: SrcLoc -> Env -> String -> EvalM a -bad loc env s = stacking loc env $ do - ss <- map (locStr . srclocOf) <$> stacktrace - liftF $ ExtOpError $ InterpreterError $ "Error at\n" ++ prettyStacktrace 0 ss ++ s - -trace :: Value -> EvalM () -trace v = do - -- We take the second-to-top element of the stack, because any - -- actual call to 'implicits.trace' is going to be in the trace - -- function in the prelude, which is not interesting. - top <- fromMaybe noLoc . maybeHead . drop 1 <$> stacktrace - liftF $ ExtOpTrace top (prettyOneLine v) () - -typeCheckerEnv :: Env -> T.Env -typeCheckerEnv env = - -- FIXME: some shadowing issues are probably not right here. - let valMap (TermValue (Just t) _) = Just t - valMap _ = Nothing - vtable = M.mapMaybe valMap $ envTerm env - nameMap k | k `M.member` vtable = Just ((T.Term, baseName k), qualName k) - | otherwise = Nothing - in mempty { T.envNameMap = M.fromList $ mapMaybe nameMap $ M.keys $ envTerm env - , T.envVtable = vtable } - -break :: EvalM () -break = do - -- We don't want the env of the function that is calling - -- intrinsics.break, since that is just going to be the boring - -- wrapper function (intrinsics are never called directly). - -- This is why we go a step up the stack. - backtrace <- asks $ drop 1 . fst - case NE.nonEmpty backtrace of - Nothing -> return () - Just backtrace' -> liftF $ ExtOpBreak BreakPoint backtrace' () - -fromArray :: Value -> (ValueShape, [Value]) -fromArray (ValueArray shape as) = (shape, elems as) -fromArray v = error $ "Expected array value, but found: " ++ pretty v - -toArray :: ValueShape -> [Value] -> Value -toArray shape vs = ValueArray shape (listArray (0, length vs - 1) vs) - -toArray' :: ValueShape -> [Value] -> Value -toArray' rowshape vs = ValueArray shape (listArray (0, length vs - 1) vs) - where shape = ShapeDim (genericLength vs) rowshape - -apply :: SrcLoc -> Env -> Value -> Value -> EvalM Value -apply loc env (ValueFun f) v = stacking loc env (f v) -apply _ _ f _ = error $ "Cannot apply non-function: " ++ pretty f - -apply2 :: SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value -apply2 loc env f x y = stacking loc env $ do f' <- apply noLoc mempty f x - apply noLoc mempty f' y - -matchPattern :: Env -> Pattern -> Value -> EvalM Env -matchPattern env p v = do - m <- runMaybeT $ patternMatch env p v - case m of - Nothing -> error $ "matchPattern: missing case for " ++ pretty p ++ " and " ++ pretty v - Just env' -> return env' - -patternMatch :: Env -> Pattern -> Value -> MaybeT EvalM Env -patternMatch env (Id v (Info t) _) val = - lift $ pure $ - valEnv (M.singleton v (Just $ T.BoundV [] $ toStruct t, val)) <> env -patternMatch env Wildcard{} _ = - lift $ pure env -patternMatch env (TuplePattern ps _) (ValueRecord vs) = - foldM (\env' (p,v) -> patternMatch env' p v) env $ - zip ps (map snd $ sortFields vs) -patternMatch env (RecordPattern ps _) (ValueRecord vs) = - foldM (\env' (p,v) -> patternMatch env' p v) env $ - M.intersectionWith (,) (M.fromList ps) vs -patternMatch env (PatternParens p _) v = patternMatch env p v -patternMatch env (PatternAscription p _ _) v = - patternMatch env p v -patternMatch env (PatternLit e _ _) v = do - v' <- lift $ eval env e - if v == v' - then pure env - else mzero -patternMatch env (PatternConstr n _ ps _) (ValueSum _ n' vs) - | n == n' = - foldM (\env' (p,v) -> patternMatch env' p v) env $ zip ps vs -patternMatch _ _ _ = mzero - -data Indexing = IndexingFix Int32 - | IndexingSlice (Maybe Int32) (Maybe Int32) (Maybe Int32) - -instance Pretty Indexing where - ppr (IndexingFix i) = ppr i - ppr (IndexingSlice i j (Just s)) = - maybe mempty ppr i <> text ":" <> - maybe mempty ppr j <> text ":" <> - ppr s - ppr (IndexingSlice i (Just j) s) = - maybe mempty ppr i <> text ":" <> - ppr j <> - maybe mempty ((text ":" <>) . ppr) s - ppr (IndexingSlice i Nothing Nothing) = - maybe mempty ppr i <> text ":" - -indexesFor :: Maybe Int32 -> Maybe Int32 -> Maybe Int32 - -> Int32 -> Maybe [Int] -indexesFor start end stride n - | (start', end', stride') <- slice, - end' == start' || signum' (end' - start') == signum' stride', - stride' /= 0, - is <- [start', start'+stride' .. end'-signum stride'], - all inBounds is = - Just $ map fromIntegral is - | otherwise = - Nothing - where inBounds i = i >= 0 && i < n - - slice = - case (start, end, stride) of - (Just start', _, _) -> - let end' = fromMaybe n end - in (start', end', fromMaybe 1 stride) - (Nothing, Just end', _) -> - let start' = 0 - in (start', end', fromMaybe 1 stride) - (Nothing, Nothing, Just stride') -> - (if stride' > 0 then 0 else n-1, - if stride' > 0 then n else -1, - stride') - (Nothing, Nothing, Nothing) -> - (0, n, 1) - --- | 'signum', but with 0 as 1. -signum' :: (Eq p, Num p) => p -> p -signum' 0 = 1 -signum' x = signum x - -indexShape :: [Indexing] -> ValueShape -> ValueShape -indexShape (IndexingFix{}:is) (ShapeDim _ shape) = - indexShape is shape -indexShape (IndexingSlice start end stride:is) (ShapeDim d shape) = - ShapeDim n $ indexShape is shape - where n = maybe 0 genericLength $ indexesFor start end stride d -indexShape _ shape = - shape - -indexArray :: [Indexing] -> Value -> Maybe Value -indexArray (IndexingFix i:is) (ValueArray _ arr) - | i >= 0, i < n = - indexArray is $ arr ! fromIntegral i - | otherwise = - Nothing - where n = arrayLength arr -indexArray (IndexingSlice start end stride:is) (ValueArray (ShapeDim _ rowshape) arr) = do - js <- indexesFor start end stride $ arrayLength arr - toArray' (indexShape is rowshape) <$> mapM (indexArray is . (arr!)) js -indexArray _ v = Just v - -updateArray :: [Indexing] -> Value -> Value -> Maybe Value -updateArray (IndexingFix i:is) (ValueArray shape arr) v - | i >= 0, i < n = do - v' <- updateArray is (arr ! i') v - Just $ ValueArray shape $ arr // [(i', v')] - | otherwise = - Nothing - where n = arrayLength arr - i' = fromIntegral i -updateArray (IndexingSlice start end stride:is) (ValueArray shape arr) (ValueArray _ v) = do - arr_is <- indexesFor start end stride $ arrayLength arr - guard $ length arr_is == arrayLength v - let update arr' (i, v') = do - x <- updateArray is (arr!i) v' - return $ arr' // [(i, x)] - fmap (ValueArray shape) $ foldM update arr $ zip arr_is $ elems v -updateArray _ _ v = Just v - -evalDimIndex :: Env -> DimIndex -> EvalM Indexing -evalDimIndex env (DimFix x) = - IndexingFix . asInt32 <$> eval env x -evalDimIndex env (DimSlice start end stride) = - IndexingSlice <$> traverse (fmap asInt32 . eval env) start - <*> traverse (fmap asInt32 . eval env) end - <*> traverse (fmap asInt32 . eval env) stride - -evalIndex :: SrcLoc -> Env -> [Indexing] -> Value -> EvalM Value -evalIndex loc env is arr = do - let oob = bad loc env $ "Index [" <> intercalate ", " (map pretty is) <> - "] out of bounds for array of shape " <> - pretty (valueShape arr) <> "." - maybe oob return $ indexArray is arr - --- | Expand type based on information that was not available at --- type-checking time (the structure of abstract types). -evalType :: Env -> StructType -> StructType -evalType _ (Scalar (Prim pt)) = Scalar $ Prim pt -evalType env (Scalar (Record fs)) = Scalar $ Record $ fmap (evalType env) fs -evalType env (Scalar (Arrow () p t1 t2)) = - Scalar $ Arrow () p (evalType env t1) (evalType env t2) -evalType env t@(Array _ u _ shape) = - let et = stripArray (shapeRank shape) t - et' = evalType env et - shape' = fmap evalDim shape - in arrayOf et' shape' u - where evalDim (NamedDim qn) - | Just (TermValue _ (ValuePrim (SignedValue (Int32Value x)))) <- - lookupVar qn env = - ConstDim $ fromIntegral x - evalDim d = d -evalType env t@(Scalar (TypeVar () _ tn args)) = - case lookupType (qualNameFromTypeName tn) env of - Just (T.TypeAbbr _ ps t') -> - let (substs, types) = mconcat $ zipWith matchPtoA ps args - onDim (NamedDim v) = fromMaybe (NamedDim v) $ M.lookup (qualLeaf v) substs - onDim d = d - in if null ps then first onDim t' - else evalType (Env mempty types mempty <> env) $ first onDim t' - Nothing -> t - - where matchPtoA (TypeParamDim p _) (TypeArgDim (NamedDim qv) _) = - (M.singleton p $ NamedDim qv, mempty) - matchPtoA (TypeParamDim p _) (TypeArgDim (ConstDim k) _) = - (M.singleton p $ ConstDim k, mempty) - matchPtoA (TypeParamType l p _) (TypeArgType t' _) = - let t'' = evalType env t' - in (mempty, M.singleton p $ T.TypeAbbr l [] t'') - matchPtoA _ _ = mempty -evalType env (Scalar (Sum cs)) = Scalar $ Sum $ (fmap . fmap) (evalType env) cs - -evalTermVar :: Env -> QualName VName -> StructType -> EvalM Value -evalTermVar env qv t = - case lookupVar qv env of - Just (TermPoly _ v) -> do size_env <- extSizeEnv - v $ evalType (size_env <> env) t - Just (TermValue _ v) -> return v - _ -> error $ "`" <> pretty qv <> "` is not bound to a value." - -typeValueShape :: Env -> StructType -> EvalM ValueShape -typeValueShape env t = do - size_env <- extSizeEnv - let t' = evalType (size_env <> env) t - case traverse dim $ typeShape mempty t' of - Nothing -> error $ "typeValueShape: failed to fully evaluate type " ++ pretty t' - Just shape -> return shape - where dim (ConstDim x) = Just $ fromIntegral x - dim _ = Nothing - -evalFunction :: Env -> [VName] -> [Pattern] -> Exp -> StructType -> EvalM Value - --- We treat zero-parameter lambdas as simply an expression to --- evaluate immediately. Note that this is *not* the same as a lambda --- that takes an empty tuple '()' as argument! Zero-parameter lambdas --- can never occur in a well-formed Futhark program, but they are --- convenient in the interpreter. -evalFunction env _ [] body rettype = - -- Eta-expand the rest to make any sizes visible. - etaExpand [] env rettype - where etaExpand vs env' (Scalar (Arrow _ _ pt rt)) = - return $ ValueFun $ \v -> do - env'' <- matchPattern env' (Wildcard (Info $ fromStruct pt) noLoc) v - etaExpand (v:vs) env'' rt - etaExpand vs env' _ = do - f <- eval env' body - foldM (apply noLoc mempty) f $ reverse vs - -evalFunction env missing_sizes (p:ps) body rettype = - return $ ValueFun $ \v -> do - env' <- matchPattern env p v - -- Fix up the last sizes, if any. - let env'' | null missing_sizes = env' - | otherwise = env' <> - i32Env (resolveExistentials missing_sizes - (patternStructType p) (valueShape v)) - evalFunction env'' missing_sizes ps body rettype - -evalFunctionBinding :: Env - -> [TypeParam] -> [Pattern] -> StructType -> [VName] -> Exp - -> EvalM TermBinding -evalFunctionBinding env tparams ps ret retext fbody = do - let ret' = evalType env ret - arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt - ftype = foldr (arrow . patternParam) ret' ps - - -- Distinguish polymorphic and non-polymorphic bindings here. - if null tparams - then TermValue (Just $ T.BoundV [] ftype) <$> - (returned env ret retext =<< evalFunction env [] ps fbody ret') - else return $ TermPoly (Just $ T.BoundV [] ftype) $ \ftype' -> do - let tparam_names = map typeParamName tparams - env' = resolveTypeParams tparam_names ftype ftype' <> env - - -- In some cases (abstract lifted types) there may be - -- missing sizes that were not fixed by the type - -- instantiation. These will have to be set by looking - -- at the actual function arguments. - missing_sizes = - filter (`M.notMember` envTerm env') $ - map typeParamName (filter isSizeParam tparams) - returned env ret retext =<< evalFunction env' missing_sizes ps fbody ret' - -evalArg :: Env -> Exp -> Maybe VName -> EvalM Value -evalArg env e ext = do - v <- eval env e - case ext of Just ext' -> putExtSize ext' $ asInt32 v - Nothing -> return () - return v - -returned :: Env -> TypeBase (DimDecl VName) als -> [VName] -> Value -> EvalM Value -returned _ _ [] v = return v -returned env ret retext v = do - mapM_ (uncurry putExtSize) $ M.toList $ - resolveExistentials retext (evalType env $ toStruct ret) $ valueShape v - return v - -eval :: Env -> Exp -> EvalM Value - -eval _ (Literal v _) = return $ ValuePrim v - -eval env (Parens e _ ) = eval env e - -eval env (QualParens (qv, _) e loc) = do - m <- evalModuleVar env qv - case m of - ModuleFun{} -> error $ "Local open of module function at " ++ locStr loc - Module m' -> eval (m'<>env) e - -eval env (TupLit vs _) = toTuple <$> mapM (eval env) vs - -eval env (RecordLit fields _) = - ValueRecord . M.fromList <$> mapM evalField fields - where evalField (RecordFieldExplicit k e _) = do - v <- eval env e - return (k, v) - evalField (RecordFieldImplicit k t loc) = do - v <- eval env $ Var (qualName k) t loc - return (baseName k, v) - -eval _ (StringLit vs _) = - return $ toArray' ShapeLeaf $ - map (ValuePrim . UnsignedValue . Int8Value . fromIntegral) vs - -eval env (ArrayLit [] (Info t) _) = do - t' <- typeValueShape env $ toStruct t - return $ toArray t' [] - -eval env (ArrayLit (v:vs) _ _) = do - v' <- eval env v - vs' <- mapM (eval env) vs - return $ toArray' (valueShape v') (v':vs') - -eval env (Range start maybe_second end (Info t, Info retext) loc) = do - start' <- asInteger <$> eval env start - maybe_second' <- traverse (fmap asInteger . eval env) maybe_second - end' <- traverse (fmap asInteger . eval env) end - - let (end_adj, step, ok) = - case (end', maybe_second') of - (DownToExclusive end'', Nothing) -> - (end'' + 1, -1, start' >= end'') - (DownToExclusive end'', Just second') -> - (end'' + 1, second' - start', start' >= end'' && second' < start') - - (ToInclusive end'', Nothing) -> - (end'', 1, start' <= end'') - (ToInclusive end'', Just second') - | second' > start' -> - (end'', second' - start', start' <= end'') - | otherwise -> - (end'', second' - start', start' >= end'' && second' /= start') - - (UpToExclusive x, Nothing) -> - (x-1, 1, start' <= x) - (UpToExclusive x, Just second') -> - (x-1, second' - start', start' <= x && second' > start') - - if ok - then returned env t retext $ - toArray' ShapeLeaf $ map toInt [start',start'+step..end_adj] - else bad loc env $ badRange start' maybe_second' end' - - where toInt = - case stripArray 1 t of - Scalar (Prim (Signed t')) -> - ValuePrim . SignedValue . intValue t' - Scalar (Prim (Unsigned t')) -> - ValuePrim . UnsignedValue . intValue t' - _ -> error $ "Nonsensical range type: " ++ show t - - badRange start' maybe_second' end' = - "Range " ++ pretty start' ++ - (case maybe_second' of - Nothing -> "" - Just second' -> ".." ++ pretty second') ++ - (case end' of - DownToExclusive x -> "..>" ++ pretty x - ToInclusive x -> "..." ++ pretty x - UpToExclusive x -> "..<"++ pretty x) ++ - " is invalid." - -eval env (Var qv (Info t) _) = evalTermVar env qv (toStruct t) - -eval env (Ascript e _ _ ) = eval env e - -eval env (Coerce e td (Info ret, Info retext) loc) = do - v <- returned env ret retext =<< eval env e - let t = evalType env $ unInfo $ expandedType td - case checkShape (structTypeShape (envShapes env) t) (valueShape v) of - Just _ -> return v - Nothing -> - bad loc env $ "Value `" <> pretty v <> "` of shape `" ++ pretty (valueShape v) ++ - "` cannot match shape of type `" <> - pretty (declaredType td) <> "` (`" <> pretty t <> "`)" - -eval env (LetPat p e body (Info ret, Info retext) _) = do - v <- eval env e - env' <- matchPattern env p v - returned env ret retext =<< eval env' body - -eval env (LetFun f (tparams, ps, _, Info ret, fbody) body _ _) = do - binding <- evalFunctionBinding env tparams ps ret [] fbody - eval (env { envTerm = M.insert f binding $ envTerm env }) body - -eval _ (IntLit v (Info t) _) = - case t of - Scalar (Prim (Signed it)) -> - return $ ValuePrim $ SignedValue $ intValue it v - Scalar (Prim (Unsigned it)) -> - return $ ValuePrim $ UnsignedValue $ intValue it v - Scalar (Prim (FloatType ft)) -> - return $ ValuePrim $ FloatValue $ floatValue ft v - _ -> error $ "eval: nonsensical type for integer literal: " ++ pretty t - -eval _ (FloatLit v (Info t) _) = - case t of - Scalar (Prim (FloatType ft)) -> - return $ ValuePrim $ FloatValue $ floatValue ft v - _ -> error $ "eval: nonsensical type for float literal: " ++ pretty t - -eval env (BinOp (op, _) op_t - (x, Info (_, xext)) (y, Info (_, yext)) - (Info t) (Info retext) loc) - | baseString (qualLeaf op) == "&&" = do - x' <- asBool <$> eval env x - if x' - then eval env y - else return $ ValuePrim $ BoolValue False - | baseString (qualLeaf op) == "||" = do - x' <- asBool <$> eval env x - if x' - then return $ ValuePrim $ BoolValue True - else eval env y - | otherwise = do - op' <- eval env $ Var op op_t loc - x' <- evalArg env x xext - y' <- evalArg env y yext - returned env t retext =<< apply2 loc env op' x' y' - -eval env (If cond e1 e2 (Info ret, Info retext) _) = do - cond' <- asBool <$> eval env cond - returned env ret retext =<< - if cond' then eval env e1 else eval env e2 - -eval env (Apply f x (Info (_, ext)) (Info t, Info retext) loc) = do - -- It is important that 'x' is evaluated first in order to bring any - -- sizes into scope that may be used in the type of 'f'. - x' <- evalArg env x ext - f' <- eval env f - returned env t retext =<< apply loc env f' x' - -eval env (Negate e _) = do - ev <- eval env e - ValuePrim <$> case ev of - ValuePrim (SignedValue (Int8Value v)) -> return $ SignedValue $ Int8Value (-v) - ValuePrim (SignedValue (Int16Value v)) -> return $ SignedValue $ Int16Value (-v) - ValuePrim (SignedValue (Int32Value v)) -> return $ SignedValue $ Int32Value (-v) - ValuePrim (SignedValue (Int64Value v)) -> return $ SignedValue $ Int64Value (-v) - ValuePrim (UnsignedValue (Int8Value v)) -> return $ UnsignedValue $ Int8Value (-v) - ValuePrim (UnsignedValue (Int16Value v)) -> return $ UnsignedValue $ Int16Value (-v) - ValuePrim (UnsignedValue (Int32Value v)) -> return $ UnsignedValue $ Int32Value (-v) - ValuePrim (UnsignedValue (Int64Value v)) -> return $ UnsignedValue $ Int64Value (-v) - ValuePrim (FloatValue (Float32Value v)) -> return $ FloatValue $ Float32Value (-v) - ValuePrim (FloatValue (Float64Value v)) -> return $ FloatValue $ Float64Value (-v) - _ -> error $ "Cannot negate " ++ pretty ev - -eval env (Index e is (Info t, Info retext) loc) = do - is' <- mapM (evalDimIndex env) is - arr <- eval env e - returned env t retext =<< evalIndex loc env is' arr - -eval env (Update src is v loc) = - maybe oob return =<< - updateArray <$> mapM (evalDimIndex env) is <*> eval env src <*> eval env v - where oob = bad loc env "Bad update" - -eval env (RecordUpdate src all_fs v _ _) = - update <$> eval env src <*> pure all_fs <*> eval env v - where update _ [] v' = v' - update (ValueRecord src') (f:fs) v' - | Just f_v <- M.lookup f src' = - ValueRecord $ M.insert f (update f_v fs v') src' - update _ _ _ = error "eval RecordUpdate: invalid value." - -eval env (LetWith dest src is v body _ loc) = do - let Ident src_vn (Info src_t) _ = src - dest' <- maybe oob return =<< - updateArray <$> mapM (evalDimIndex env) is <*> - evalTermVar env (qualName src_vn) (toStruct src_t) <*> eval env v - let t = T.BoundV [] $ toStruct $ unInfo $ identType dest - eval (valEnv (M.singleton (identName dest) (Just t, dest')) <> env) body - where oob = bad loc env "Bad update" - --- We treat zero-parameter lambdas as simply an expression to --- evaluate immediately. Note that this is *not* the same as a lambda --- that takes an empty tuple '()' as argument! Zero-parameter lambdas --- can never occur in a well-formed Futhark program, but they are --- convenient in the interpreter. -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 - 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 - y <- evalArg env e argext - return $ ValueFun $ \x -> do - f <- evalTermVar env qv $ toStruct t - apply2 loc env f x y - -eval env (IndexSection is _ loc) = do - is' <- mapM (evalDimIndex env) is - return $ ValueFun $ evalIndex loc env is' - -eval _ (ProjectSection ks _ _) = - return $ ValueFun $ flip (foldM walk) ks - where walk (ValueRecord fs) f - | Just v' <- M.lookup f fs = return v' - walk _ _ = error "Value does not have expected field." - -eval env (DoLoop sparams pat init_e form body (Info (ret, retext)) _) = do - init_v <- eval env init_e - returned env ret retext =<< - case form of For iv bound -> do - bound' <- asSigned <$> eval env bound - forLoop (identName iv) bound' (zero bound') init_v - ForIn in_pat in_e -> do - (_, in_vs) <- fromArray <$> eval env in_e - foldM (forInLoop in_pat) init_v in_vs - While cond -> - whileLoop cond init_v - where withLoopParams v = - let sparams' = - resolveExistentials sparams - (patternStructType pat) (valueShape v) - in matchPattern (i32Env sparams' <> env) pat v - - inc = (`P.doAdd` Int64Value 1) - zero = (`P.doMul` Int64Value 0) - - forLoop iv bound i v - | i >= bound = return v - | otherwise = do - env' <- withLoopParams v - forLoop iv bound (inc i) =<< - eval (valEnv (M.singleton iv (Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int32, - ValuePrim (SignedValue i))) <> env') body - - whileLoop cond v = do - env' <- withLoopParams v - continue <- asBool <$> eval env' cond - if continue - then whileLoop cond =<< eval env' body - else return v - - forInLoop in_pat v in_v = do - env' <- withLoopParams v - env'' <- matchPattern env' in_pat in_v - eval env'' body - -eval env (Project f e _ _) = do - v <- eval env e - case v of - ValueRecord fs | Just v' <- M.lookup f fs -> return v' - _ -> error "Value does not have expected field." - -eval env (Assert what e (Info s) loc) = do - cond <- asBool <$> eval env what - unless cond $ bad loc env s - eval env e - -eval env (Constr c es (Info t) _) = do - vs <- mapM (eval env) es - shape <- typeValueShape env $ toStruct t - return $ ValueSum shape c vs - -eval env (Match e cs (Info ret, Info retext) _) = do - v <- eval env e - returned env ret retext =<< match v (NE.toList cs) - where match _ [] = - error "Pattern match failure." - match v (c:cs') = do - c' <- evalCase v env c - case c' of - Just v' -> return v' - Nothing -> match v cs' - -eval env (Attr _ e _) = eval env e - -evalCase :: Value -> Env -> CaseBase Info VName - -> EvalM (Maybe Value) -evalCase v env (CasePat p cExp _) = runMaybeT $ do - env' <- patternMatch env p v - lift $ eval env' cExp - -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 - replaceM f m = M.fromList $ do - (k, v) <- M.toList m - return (replace 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) - onTerm (TermValue t v) = TermValue t v - onTerm (TermPoly t v) = TermPoly t v - onTerm (TermModule m) = TermModule $ onModule m - onType (T.TypeAbbr l ps t) = T.TypeAbbr l ps $ first onDim t - onDim (NamedDim v) = NamedDim $ replaceQ v - 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 - Just (TermModule m) -> return m - _ -> error $ quote (pretty qv) <> " is not bound to a module." - -evalModExp :: Env -> ModExp -> EvalM Module - -evalModExp _ (ModImport _ (Info f) _) = do - f' <- lookupImport f - case f' of Nothing -> error $ "Unknown import " ++ show f - Just m -> return $ Module m - -evalModExp env (ModDecs ds _) = do - Env terms types _ <- foldM evalDec env ds - -- Remove everything that was present in the original Env. - return $ Module $ Env (terms `M.difference` envTerm env) - (types `M.difference` envType env) - mempty - -evalModExp env (ModVar qv _) = - evalModuleVar env qv - -evalModExp env (ModAscript me _ (Info substs) _) = - substituteInModule substs <$> evalModExp env me - -evalModExp env (ModParens me _) = evalModExp env me - -evalModExp env (ModLambda p ret e loc) = - return $ ModuleFun $ \am -> do - let env' = env { envTerm = M.insert (modParamName p) (TermModule am) $ envTerm env } - evalModExp env' $ case ret of - Nothing -> e - Just (se, rsubsts) -> ModAscript e se rsubsts loc - -evalModExp env (ModApply f e (Info psubst) (Info rsubst) _) = do - f' <- evalModExp env f - case f' of - ModuleFun f'' -> do - e' <- evalModExp env e - substituteInModule rsubst <$> f'' (substituteInModule psubst e') - _ -> error "Expected ModuleFun." - -evalDec :: Env -> Dec -> EvalM Env - -evalDec env (ValDec (ValBind _ v _ (Info (ret, retext)) tparams ps fbody _ _ _)) = do - binding <- evalFunctionBinding env tparams ps ret retext fbody - return $ env { envTerm = M.insert v binding $ envTerm env } - -evalDec env (OpenDec me _) = do - me' <- evalModExp env me - case me' of - Module me'' -> return $ me'' <> env - _ -> error "Expected Module" - -evalDec env (ImportDec name name' loc) = - evalDec env $ LocalDec (OpenDec (ModImport name name' loc) loc) loc - -evalDec env (LocalDec d _) = evalDec env d -evalDec env SigDec{} = return env -evalDec env (TypeDec (TypeBind v l ps t _ _)) = do - let abbr = T.TypeAbbr l ps $ - evalType env $ unInfo $ expandedType t - return env { envType = M.insert v abbr $ envType env } -evalDec env (ModDec (ModBind v ps ret body _ loc)) = do - mod <- evalModExp env $ wrapInLambda ps - return $ modEnv (M.singleton v mod) <> env - where wrapInLambda [] = case ret of - Just (se, substs) -> ModAscript body se substs loc - Nothing -> body - wrapInLambda [p] = ModLambda p ret body loc - wrapInLambda (p:ps') = ModLambda p Nothing (wrapInLambda ps') loc - --- | The interpreter context. All evaluation takes place with respect --- to a context, and it can be extended with more definitions, which --- is how the REPL works. -data Ctx = Ctx { ctxEnv :: Env - , ctxImports :: M.Map FilePath Env - } - -nanValue :: PrimValue -> Bool -nanValue (FloatValue v) = - case v of Float32Value x -> isNaN x - Float64Value x -> isNaN x -nanValue _ = False - -breakOnNaN :: [PrimValue] -> PrimValue -> EvalM () -breakOnNaN inputs result - | not (any nanValue inputs) && nanValue result = do - backtrace <- asks fst - case NE.nonEmpty backtrace of - Nothing -> return () - Just backtrace' -> liftF $ ExtOpBreak BreakNaN backtrace' () -breakOnNaN _ _ = - return () - --- | The initial environment contains definitions of the various intrinsic functions. -initialCtx :: Ctx -initialCtx = - Ctx (Env (M.insert (VName (nameFromString "intrinsics") 0) - (TermModule (Module $ Env terms types mempty)) terms) - types mempty) - mempty - where - terms = M.mapMaybeWithKey (const . def . baseString) intrinsics - types = M.mapMaybeWithKey (const . tdef . baseString) intrinsics - - sintOp f = [ (getS, putS, P.doBinOp (f Int8)) - , (getS, putS, P.doBinOp (f Int16)) - , (getS, putS, P.doBinOp (f Int32)) - , (getS, putS, P.doBinOp (f Int64))] - uintOp f = [ (getU, putU, P.doBinOp (f Int8)) - , (getU, putU, P.doBinOp (f Int16)) - , (getU, putU, P.doBinOp (f Int32)) - , (getU, putU, P.doBinOp (f Int64))] - intOp f = sintOp f ++ uintOp f - floatOp f = [ (getF, putF, P.doBinOp (f Float32)) - , (getF, putF, P.doBinOp (f Float64))] - arithOp f g = Just $ bopDef $ intOp f ++ floatOp g - - flipCmps = map (\(f, g, h) -> (f, g, flip h)) - sintCmp f = [ (getS, Just . BoolValue, P.doCmpOp (f Int8)) - , (getS, Just . BoolValue, P.doCmpOp (f Int16)) - , (getS, Just . BoolValue, P.doCmpOp (f Int32)) - , (getS, Just . BoolValue, P.doCmpOp (f Int64))] - uintCmp f = [ (getU, Just . BoolValue, P.doCmpOp (f Int8)) - , (getU, Just . BoolValue, P.doCmpOp (f Int16)) - , (getU, Just . BoolValue, P.doCmpOp (f Int32)) - , (getU, Just . BoolValue, P.doCmpOp (f Int64))] - floatCmp f = [ (getF, Just . BoolValue, P.doCmpOp (f Float32)) - , (getF, Just . BoolValue, P.doCmpOp (f Float64))] - boolCmp f = [ (getB, Just . BoolValue, P.doCmpOp f) ] - - getV (SignedValue x) = Just $ P.IntValue x - getV (UnsignedValue x) = Just $ P.IntValue x - getV (FloatValue x) = Just $ P.FloatValue x - getV (BoolValue x) = Just $ P.BoolValue x - putV (P.IntValue x) = SignedValue x - putV (P.FloatValue x) = FloatValue x - putV (P.BoolValue x) = BoolValue x - putV P.Checked = BoolValue True - - getS (SignedValue x) = Just $ P.IntValue x - getS _ = Nothing - putS (P.IntValue x) = Just $ SignedValue x - putS _ = Nothing - - getU (UnsignedValue x) = Just $ P.IntValue x - getU _ = Nothing - putU (P.IntValue x) = Just $ UnsignedValue x - putU _ = Nothing - - getF (FloatValue x) = Just $ P.FloatValue x - getF _ = Nothing - putF (P.FloatValue x) = Just $ FloatValue x - putF _ = Nothing - - getB (BoolValue x) = Just $ P.BoolValue x - getB _ = Nothing - putB (P.BoolValue x) = Just $ BoolValue x - putB _ = Nothing - - fun1 f = - TermValue Nothing $ ValueFun $ \x -> f x - fun2 f = - TermValue Nothing $ ValueFun $ \x -> - return $ ValueFun $ \y -> f x y - fun2t f = - TermValue Nothing $ ValueFun $ \v -> - case fromTuple v of Just [x,y] -> f x y - _ -> error $ "Expected pair; got: " ++ pretty v - fun3t f = - TermValue Nothing $ ValueFun $ \v -> - case fromTuple v of Just [x,y,z] -> f x y z - _ -> error $ "Expected triple; got: " ++ pretty v - - fun6t f = - TermValue Nothing $ ValueFun $ \v -> - case fromTuple v of Just [x,y,z,a,b,c] -> f x y z a b c - _ -> error $ "Expected sextuple; got: " ++ pretty v - - bopDef fs = fun2 $ \x y -> - case (x, y) of - (ValuePrim x', ValuePrim y') - | Just z <- msum $ map (`bopDef'` (x', y')) fs -> do - breakOnNaN [x', y'] z - return $ ValuePrim z - _ -> - bad noLoc mempty $ "Cannot apply operator to arguments " <> - quote (pretty x) <> " and " <> quote (pretty y) <> "." - where bopDef' (valf, retf, op) (x, y) = do - x' <- valf x - y' <- valf y - retf =<< op x' y' - - unopDef fs = fun1 $ \x -> - case x of - (ValuePrim x') - | Just r <- msum $ map (`unopDef'` x') fs -> do - breakOnNaN [x'] r - return $ ValuePrim r - _ -> - bad noLoc mempty $ "Cannot apply function to argument " <> - quote (pretty x) <> "." - where unopDef' (valf, retf, op) x = do - x' <- valf x - retf =<< op x' - - tbopDef f = fun1 $ \v -> - case fromTuple v of - Just [ValuePrim x, ValuePrim y] - | Just x' <- getV x, - Just y' <- getV y, - Just z <- putV <$> f x' y' -> do - breakOnNaN [x, y] z - return $ ValuePrim z - _ -> - bad noLoc mempty $ "Cannot apply operator to argument " <> - quote (pretty v) <> "." - - def "!" = Just $ unopDef [ (getS, putS, P.doUnOp $ P.Complement Int8) - , (getS, putS, P.doUnOp $ P.Complement Int16) - , (getS, putS, P.doUnOp $ P.Complement Int32) - , (getS, putS, P.doUnOp $ P.Complement Int64) - , (getU, putU, P.doUnOp $ P.Complement Int8) - , (getU, putU, P.doUnOp $ P.Complement Int16) - , (getU, putU, P.doUnOp $ P.Complement Int32) - , (getU, putU, P.doUnOp $ P.Complement Int64) - , (getB, putB, P.doUnOp P.Not) ] - - def "+" = arithOp (`P.Add` P.OverflowWrap) P.FAdd - def "-" = arithOp (`P.Sub` P.OverflowWrap) P.FSub - def "*" = arithOp (`P.Mul` P.OverflowWrap) P.FMul - def "**" = arithOp P.Pow P.FPow - def "/" = Just $ bopDef $ - sintOp (`P.SDiv` P.Unsafe) ++ - uintOp (`P.UDiv` P.Unsafe) ++ - floatOp P.FDiv - def "%" = Just $ bopDef $ - sintOp (`P.SMod` P.Unsafe) ++ - uintOp (`P.UMod` P.Unsafe) ++ - floatOp P.FMod - def "//" = Just $ bopDef $ - sintOp (`P.SQuot` P.Unsafe) ++ - uintOp (`P.UDiv` P.Unsafe) - def "%%" = Just $ bopDef $ - sintOp (`P.SRem` P.Unsafe) ++ - uintOp (`P.UMod` P.Unsafe) - def "^" = Just $ bopDef $ intOp P.Xor - def "&" = Just $ bopDef $ intOp P.And - def "|" = Just $ bopDef $ intOp P.Or - def ">>" = Just $ bopDef $ sintOp P.AShr ++ uintOp P.LShr - def "<<" = Just $ bopDef $ intOp P.Shl - def ">>>" = Just $ bopDef $ sintOp P.LShr ++ uintOp P.LShr - def "==" = Just $ fun2 $ - \xs ys -> return $ ValuePrim $ BoolValue $ xs == ys - def "!=" = Just $ fun2 $ - \xs ys -> return $ ValuePrim $ BoolValue $ xs /= ys - - -- The short-circuiting is handled directly in 'eval'; these cases - -- are only used when partially applying and such. - def "&&" = Just $ fun2 $ \x y -> - return $ ValuePrim $ BoolValue $ asBool x && asBool y - def "||" = Just $ fun2 $ \x y -> - return $ ValuePrim $ BoolValue $ asBool x || asBool y - - def "<" = Just $ bopDef $ - sintCmp P.CmpSlt ++ uintCmp P.CmpUlt ++ - floatCmp P.FCmpLt ++ boolCmp P.CmpLlt - def ">" = Just $ bopDef $ flipCmps $ - sintCmp P.CmpSlt ++ uintCmp P.CmpUlt ++ - floatCmp P.FCmpLt ++ boolCmp P.CmpLlt - def "<=" = Just $ bopDef $ - sintCmp P.CmpSle ++ uintCmp P.CmpUle ++ - floatCmp P.FCmpLe ++ boolCmp P.CmpLle - def ">=" = Just $ bopDef $ flipCmps $ - sintCmp P.CmpSle ++ uintCmp P.CmpUle ++ - floatCmp P.FCmpLe ++ boolCmp P.CmpLle - - def s - | Just bop <- find ((s==) . pretty) P.allBinOps = - Just $ tbopDef $ P.doBinOp bop - | Just unop <- find ((s==) . pretty) P.allCmpOps = - Just $ tbopDef $ \x y -> P.BoolValue <$> P.doCmpOp unop x y - | Just cop <- find ((s==) . pretty) P.allConvOps = - Just $ unopDef [(getV, Just . putV, P.doConvOp cop)] - | Just unop <- find ((s==) . pretty) P.allUnOps = - Just $ unopDef [(getV, Just . putV, P.doUnOp unop)] - - | Just (pts, _, f) <- M.lookup s P.primFuns = - case length pts of - 1 -> Just $ unopDef [(getV, Just . putV, f . pure)] - _ -> Just $ fun1 $ \x -> do - let getV' (ValuePrim v) = Just v - getV' _ = Nothing - case mapM getV' =<< fromTuple x of - Just vs - | Just res <- fmap putV . f =<< mapM getV vs -> do - breakOnNaN vs res - return $ ValuePrim res - _ -> - error $ "Cannot apply " ++ pretty s ++ " to " ++ pretty x - - | "sign_" `isPrefixOf` s = - Just $ fun1 $ \x -> - case x of (ValuePrim (UnsignedValue x')) -> - return $ ValuePrim $ SignedValue x' - _ -> error $ "Cannot sign: " ++ pretty x - | "unsign_" `isPrefixOf` s = - Just $ fun1 $ \x -> - case x of (ValuePrim (SignedValue x')) -> - return $ ValuePrim $ UnsignedValue x' - _ -> error $ "Cannot unsign: " ++ pretty x - - def s | "map_stream" `isPrefixOf` s = - Just $ fun2t stream - - def s | "reduce_stream" `isPrefixOf` s = - Just $ fun3t $ \_ f arg -> stream f arg - - def "map" = Just $ TermPoly Nothing $ \t -> return $ ValueFun $ \v -> - case (fromTuple v, unfoldFunType t) of - (Just [f, xs], ([_], ret_t)) - | Just rowshape <- typeRowShape ret_t -> - toArray' rowshape <$> mapM (apply noLoc mempty f) (snd $ fromArray xs) - | otherwise -> - error $ "Bad return type: " ++ pretty ret_t - _ -> - error $ "Invalid arguments to map intrinsic:\n" ++ - unlines [pretty t, pretty v] - where typeRowShape = sequenceA . structTypeShape mempty . stripArray 1 - - def s | "reduce" `isPrefixOf` s = Just $ fun3t $ \f ne xs -> - foldM (apply2 noLoc mempty f) ne $ snd $ fromArray xs - - def "scan" = Just $ fun3t $ \f ne xs -> do - let next (out, acc) x = do - x' <- apply2 noLoc mempty f acc x - return (x':out, x') - toArray' (valueShape ne) . reverse . fst <$> - foldM next ([], ne) (snd $ fromArray xs) - - def "scatter" = Just $ fun3t $ \arr is vs -> - case arr of - ValueArray shape arr' -> - return $ ValueArray shape $ foldl' update arr' $ - zip (map asInt $ snd $ fromArray is) (snd $ fromArray vs) - _ -> - error $ "scatter expects array, but got: " ++ pretty arr - where update arr' (i, v) = - if i >= 0 && i < arrayLength arr' - then arr' // [(i, v)] else arr' - - def "hist" = Just $ fun6t $ \_ arr fun _ is vs -> - case arr of - ValueArray shape arr' -> - ValueArray shape <$> foldM (update fun) arr' - (zip (map asInt $ snd $ fromArray is) (snd $ fromArray vs)) - _ -> - error $ "hist expects array, but got: " ++ pretty arr - where update fun arr' (i, v) = - if i >= 0 && i < arrayLength arr' - then do - v' <- apply2 noLoc mempty fun (arr' ! i) v - return $ arr' // [(i, v')] - else return arr' - - def "partition" = Just $ fun3t $ \k f xs -> do - let (ShapeDim _ rowshape, xs') = fromArray xs - - next outs x = do - i <- asInt <$> apply noLoc mempty f x - return $ insertAt i x outs - pack parts = - toTuple [toArray' rowshape $ concat parts, - toArray' rowshape $ - map (ValuePrim . SignedValue . Int32Value . genericLength) parts] - - pack . map reverse <$> - foldM next (replicate (asInt k) []) xs' - where insertAt 0 x (l:ls) = (x:l):ls - insertAt i x (l:ls) = l:insertAt (i-1) x ls - insertAt _ _ ls = ls - - def "unzip" = Just $ fun1 $ \x -> do - let ShapeDim _ (ShapeRecord fs) = valueShape x - Just [xs_shape, ys_shape] = areTupleFields fs - listPair (xs, ys) = - [toArray' xs_shape xs, toArray' ys_shape ys] - - return $ toTuple $ listPair $ unzip $ map (fromPair . fromTuple) $ snd $ fromArray x - where fromPair (Just [x,y]) = (x,y) - fromPair l = error $ "Not a pair: " ++ pretty l - - def "zip" = Just $ fun2t $ \xs ys -> do - let ShapeDim _ xs_rowshape = valueShape xs - ShapeDim _ ys_rowshape = valueShape ys - return $ toArray' (ShapeRecord (tupleFields [xs_rowshape, ys_rowshape])) $ - map toTuple $ transpose [snd $ fromArray xs, snd $ fromArray ys] - - def "concat" = Just $ fun2t $ \xs ys -> do - let (ShapeDim _ rowshape, xs') = fromArray xs - (_, ys') = fromArray ys - return $ toArray' rowshape $ xs' ++ ys' - - def "transpose" = Just $ fun1 $ \xs -> do - let (ShapeDim n (ShapeDim m shape), xs') = fromArray xs - return $ toArray (ShapeDim m (ShapeDim n shape)) $ - map (toArray (ShapeDim n shape)) $ transpose $ map (snd . fromArray) xs' - - def "rotate" = Just $ fun2t $ \i xs -> do - let (shape, xs') = fromArray xs - return $ - if asInt i > 0 - then let (bef, aft) = splitAt (asInt i) xs' - in toArray shape $ aft ++ bef - else let (bef, aft) = splitFromEnd (-asInt i) xs' - in toArray shape $ aft ++ bef - - def "flatten" = Just $ fun1 $ \xs -> do - let (ShapeDim n (ShapeDim m shape), xs') = fromArray xs - return $ toArray (ShapeDim (n*m) shape) $ concatMap (snd . fromArray) xs' - - def "unflatten" = Just $ fun3t $ \n m xs -> do - let (ShapeDim _ innershape, xs') = fromArray xs - rowshape = ShapeDim (asInt32 m) innershape - shape = ShapeDim (asInt32 n) rowshape - return $ toArray shape $ map (toArray rowshape) $ chunk (asInt m) xs' - - def "opaque" = Just $ fun1 return - - def "trace" = Just $ fun1 $ \v -> trace v >> return v - - def "break" = Just $ fun1 $ \v -> do - break - return v - - def s | nameFromString s `M.member` namesToPrimTypes = Nothing - - def s = error $ "Missing intrinsic: " ++ s - - tdef s = do - t <- nameFromString s `M.lookup` namesToPrimTypes - return $ T.TypeAbbr Unlifted [] $ Scalar $ Prim t - - stream f arg@(ValueArray _ xs) = - let n = ValuePrim $ SignedValue $ Int32Value $ arrayLength xs - in apply2 noLoc mempty f n arg - stream _ arg = error $ "Cannot stream: " ++ pretty arg - - -interpretExp :: Ctx -> Exp -> F ExtOp Value -interpretExp ctx e = runEvalM (ctxImports ctx) $ eval (ctxEnv ctx) e - -interpretDec :: Ctx -> Dec -> F ExtOp Ctx -interpretDec ctx d = do - env <- runEvalM (ctxImports ctx) $ evalDec (ctxEnv ctx) d - return ctx { ctxEnv = env } - -interpretImport :: Ctx -> (FilePath, Prog) -> F ExtOp Ctx -interpretImport ctx (fp, prog) = do - env <- runEvalM (ctxImports ctx) $ foldM evalDec (ctxEnv ctx) $ progDecs prog - return ctx { ctxImports = M.insert fp env $ ctxImports ctx } - -checkEntryArgs :: VName -> [F.Value] -> StructType -> Either String () -checkEntryArgs entry args entry_t - | args_ts == param_ts = - return () - | otherwise = - Left $ pretty $ expected - "Got input of types" - indent 2 (stack (map ppr args_ts)) - where (param_ts, _) = unfoldFunType entry_t - args_ts = map (valueStructType . valueType) args - expected - | null param_ts = - "Entry point " <> pquote (pprName entry) <> " is not a function." - | otherwise = - "Entry point " <> pquote (pprName entry) <> " expects input of type(s)" - indent 2 (stack (map ppr param_ts)) - --- | Execute the named function on the given arguments; may fail --- horribly if these are ill-typed. -interpretFunction :: Ctx -> VName -> [F.Value] -> Either String (F ExtOp Value) -interpretFunction ctx fname vs = do - ft <- case lookupVar (qualName fname) $ ctxEnv ctx of - Just (TermValue (Just (T.BoundV _ t)) _) -> - Right $ updateType (map valueType vs) t - Just (TermPoly (Just (T.BoundV _ t)) _) -> - Right $ updateType (map valueType vs) t - _ -> - Left $ "Unknown function `" <> prettyName fname <> "`." - - vs' <- case mapM convertValue vs of - Just vs' -> Right vs' - Nothing -> Left "Invalid input: irregular array." - - checkEntryArgs fname vs ft - - Right $ runEvalM (ctxImports ctx) $ do - f <- evalTermVar (ctxEnv ctx) (qualName fname) ft - foldM (apply noLoc mempty) f vs' - - where updateType (vt:vts) (Scalar (Arrow als u _ rt)) = - Scalar $ Arrow als u (valueStructType vt) $ updateType vts rt - updateType _ t = t - - convertValue (F.PrimValue p) = Just $ ValuePrim p - convertValue (F.ArrayValue arr t) = mkArray t =<< mapM convertValue (elems arr) +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | An interpreter operating on type-checked source Futhark terms. +-- Relatively slow. +module Language.Futhark.Interpreter + ( Ctx (..), + Env, + InterpreterError, + initialCtx, + interpretExp, + interpretDec, + interpretImport, + interpretFunction, + ExtOp (..), + BreakReason (..), + StackFrame (..), + typeCheckerEnv, + Value (ValuePrim, ValueArray, ValueRecord), + fromTuple, + isEmptyArray, + prettyEmptyArray, + ) +where + +import Control.Monad.Except +import Control.Monad.Free.Church +import Control.Monad.Reader +import Control.Monad.State +import Control.Monad.Trans.Maybe +import Data.Array +import Data.Bifunctor (first) +import Data.List + ( find, + foldl', + genericLength, + intercalate, + isPrefixOf, + transpose, + ) +import qualified Data.List.NonEmpty as NE +import qualified Data.Map as M +import Data.Maybe +import Data.Monoid hiding (Sum) +import Futhark.IR.Primitive (floatValue, intValue) +import qualified Futhark.IR.Primitive as P +import Futhark.Util (chunk, maybeHead, splitFromEnd) +import Futhark.Util.Loc +import Futhark.Util.Pretty hiding (apply, bool) +import Language.Futhark hiding (Value, matchDims) +import qualified Language.Futhark as F +import qualified Language.Futhark.Semantic as T +import Prelude hiding (break, mod) + +data StackFrame = StackFrame + { stackFrameLoc :: Loc, + stackFrameCtx :: Ctx + } + +instance Located StackFrame where + locOf = stackFrameLoc + +-- | What is the reason for this break point? +data BreakReason + = -- | An explicit breakpoint in the program. + BreakPoint + | -- | A + BreakNaN + +data ExtOp a + = ExtOpTrace Loc String a + | ExtOpBreak BreakReason (NE.NonEmpty StackFrame) a + | ExtOpError InterpreterError + +instance Functor ExtOp where + fmap f (ExtOpTrace w s x) = ExtOpTrace w s $ f x + fmap f (ExtOpBreak why backtrace x) = ExtOpBreak why backtrace $ f x + fmap _ (ExtOpError err) = ExtOpError err + +type Stack = [StackFrame] + +type Sizes = M.Map VName Int64 + +-- | The monad in which evaluation takes place. +newtype EvalM a + = EvalM + ( ReaderT + (Stack, M.Map FilePath Env) + (StateT Sizes (F ExtOp)) + a + ) + deriving + ( Monad, + Applicative, + Functor, + MonadFree ExtOp, + MonadReader (Stack, M.Map FilePath Env), + MonadState Sizes + ) + +runEvalM :: M.Map FilePath Env -> EvalM a -> F ExtOp a +runEvalM imports (EvalM m) = evalStateT (runReaderT m (mempty, imports)) mempty + +stacking :: SrcLoc -> Env -> EvalM a -> EvalM a +stacking loc env = local $ \(ss, imports) -> + if isNoLoc loc + then (ss, imports) + else + let s = StackFrame (locOf loc) (Ctx env imports) + in (s : ss, imports) + where + isNoLoc :: SrcLoc -> Bool + isNoLoc = (== NoLoc) . locOf + +stacktrace :: EvalM [Loc] +stacktrace = asks $ map stackFrameLoc . fst + +lookupImport :: FilePath -> EvalM (Maybe Env) +lookupImport f = asks $ M.lookup f . snd + +putExtSize :: VName -> Int64 -> EvalM () +putExtSize v x = modify $ M.insert v x + +getSizes :: EvalM Sizes +getSizes = get + +extSizeEnv :: EvalM Env +extSizeEnv = i64Env <$> getSizes + +prettyRecord :: Pretty a => M.Map Name a -> Doc +prettyRecord m + | Just vs <- areTupleFields m = + parens $ commasep $ map ppr vs + | otherwise = + braces $ commasep $ map field $ M.toList m + where + field (k, v) = ppr k <+> equals <+> ppr v + +valueStructType :: ValueType -> StructType +valueStructType = first (ConstDim . fromIntegral) + +-- | A shape is a tree to accomodate the case of records. It is +-- parameterised over the representation of dimensions. +data Shape d + = ShapeDim d (Shape d) + | ShapeLeaf + | ShapeRecord (M.Map Name (Shape d)) + | ShapeSum (M.Map Name [Shape d]) + deriving (Eq, Show, Functor, Foldable, Traversable) + +type ValueShape = Shape Int64 + +instance Pretty d => Pretty (Shape d) where + ppr ShapeLeaf = mempty + ppr (ShapeDim d s) = brackets (ppr d) <> ppr s + ppr (ShapeRecord m) = prettyRecord m + ppr (ShapeSum cs) = + mconcat (punctuate (text " | ") cs') + where + ppConstr (name, fs) = sep $ (text "#" <> ppr name) : map ppr fs + cs' = map ppConstr $ M.toList cs + +emptyShape :: ValueShape -> Bool +emptyShape (ShapeDim d s) = d == 0 || emptyShape s +emptyShape _ = False + +typeShape :: M.Map VName (Shape d) -> TypeBase d () -> Shape d +typeShape shapes = go + where + go (Array _ _ et shape) = + foldr ShapeDim (go (Scalar et)) $ shapeDims shape + go (Scalar (Record fs)) = + ShapeRecord $ M.map go fs + go (Scalar (Sum cs)) = + ShapeSum $ M.map (map go) cs + go (Scalar (TypeVar _ _ (TypeName [] v) [])) + | Just shape <- M.lookup v shapes = + shape + go _ = + ShapeLeaf + +structTypeShape :: M.Map VName ValueShape -> StructType -> Shape (Maybe Int64) +structTypeShape shapes = fmap dim . typeShape shapes' + where + dim (ConstDim d) = Just $ fromIntegral d + dim _ = Nothing + shapes' = M.map (fmap $ ConstDim . fromIntegral) shapes + +resolveTypeParams :: [VName] -> StructType -> StructType -> Env +resolveTypeParams names = match + where + match (Scalar (TypeVar _ _ tn _)) t + | typeLeaf tn `elem` names = + typeEnv $ M.singleton (typeLeaf tn) t + match (Scalar (Record poly_fields)) (Scalar (Record fields)) = + mconcat $ + M.elems $ + M.intersectionWith match poly_fields fields + match (Scalar (Sum poly_fields)) (Scalar (Sum fields)) = + mconcat $ + map mconcat $ + M.elems $ + M.intersectionWith (zipWith match) poly_fields fields + match (Scalar (Arrow _ _ poly_t1 poly_t2)) (Scalar (Arrow _ _ t1 t2)) = + match poly_t1 t1 <> match poly_t2 t2 + match poly_t t + | d1 : _ <- shapeDims (arrayShape poly_t), + d2 : _ <- shapeDims (arrayShape t) = + matchDims d1 d2 <> match (stripArray 1 poly_t) (stripArray 1 t) + match _ _ = mempty + + matchDims (NamedDim (QualName _ d1)) (ConstDim d2) + | d1 `elem` names = + i64Env $ M.singleton d1 $ fromIntegral d2 + matchDims _ _ = mempty + +resolveExistentials :: [VName] -> StructType -> ValueShape -> M.Map VName Int64 +resolveExistentials names = match + where + match (Scalar (Record poly_fields)) (ShapeRecord fields) = + mconcat $ + M.elems $ + M.intersectionWith match poly_fields fields + match (Scalar (Sum poly_fields)) (ShapeSum fields) = + mconcat $ + map mconcat $ + M.elems $ + M.intersectionWith (zipWith match) poly_fields fields + match poly_t (ShapeDim d2 rowshape) + | d1 : _ <- shapeDims (arrayShape poly_t) = + matchDims d1 d2 <> match (stripArray 1 poly_t) rowshape + match _ _ = mempty + + matchDims (NamedDim (QualName _ d1)) d2 + | d1 `elem` names = M.singleton d1 d2 + matchDims _ _ = mempty + +-- | A fully evaluated Futhark value. +data Value + = ValuePrim !PrimValue + | ValueArray ValueShape !(Array Int Value) + | -- Stores the full shape. + ValueRecord (M.Map Name Value) + | ValueFun (Value -> EvalM Value) + | ValueSum ValueShape Name [Value] + +-- Stores the full shape. + +instance Eq Value where + ValuePrim x == ValuePrim y = x == y + ValueArray _ x == ValueArray _ y = x == y + ValueRecord x == ValueRecord y = x == y + ValueSum _ n1 vs1 == ValueSum _ n2 vs2 = n1 == n2 && vs1 == vs2 + _ == _ = False + +instance Pretty Value where + ppr (ValuePrim v) = ppr v + ppr (ValueArray _ a) = + let elements = elems a -- [Value] + (x : _) = elements + separator = case x of + ValueArray _ _ -> comma <> line + _ -> comma <> space + in brackets $ cat $ punctuate separator (map ppr elements) + ppr (ValueRecord m) = prettyRecord m + ppr ValueFun {} = text "#" + ppr (ValueSum _ n vs) = text "#" <> sep (ppr n : map ppr vs) + +valueShape :: Value -> ValueShape +valueShape (ValueArray shape _) = shape +valueShape (ValueRecord fs) = ShapeRecord $ M.map valueShape fs +valueShape (ValueSum shape _ _) = shape +valueShape _ = ShapeLeaf + +checkShape :: Shape (Maybe Int64) -> ValueShape -> Maybe ValueShape +checkShape (ShapeDim Nothing shape1) (ShapeDim d2 shape2) = + ShapeDim d2 <$> checkShape shape1 shape2 +checkShape (ShapeDim (Just d1) shape1) (ShapeDim d2 shape2) = do + guard $ d1 == d2 + ShapeDim d2 <$> checkShape shape1 shape2 +checkShape (ShapeDim d1 shape1) ShapeLeaf = + -- This case is for handling polymorphism, when a function doesn't + -- know that the array it produced actually has more dimensions. + ShapeDim (fromMaybe 0 d1) <$> checkShape shape1 ShapeLeaf +checkShape (ShapeRecord shapes1) (ShapeRecord shapes2) = + ShapeRecord <$> sequence (M.intersectionWith checkShape shapes1 shapes2) +checkShape (ShapeRecord shapes1) ShapeLeaf = + Just $ fromMaybe 0 <$> ShapeRecord shapes1 +checkShape (ShapeSum shapes1) (ShapeSum shapes2) = + ShapeSum <$> sequence (M.intersectionWith (zipWithM checkShape) shapes1 shapes2) +checkShape (ShapeSum shapes1) ShapeLeaf = + Just $ fromMaybe 0 <$> ShapeSum shapes1 +checkShape _ shape2 = + Just shape2 + +-- | Does the value correspond to an empty array? +isEmptyArray :: Value -> Bool +isEmptyArray = emptyShape . valueShape + +-- | String representation of an empty array with the provided element +-- type. This is pretty ad-hoc - don't expect good results unless the +-- element type is a primitive. +prettyEmptyArray :: TypeBase () () -> Value -> String +prettyEmptyArray t v = + "empty(" ++ dims (valueShape v) ++ pretty t' ++ ")" + where + t' = stripArray (arrayRank t) t + dims (ShapeDim n rowshape) = + "[" ++ pretty n ++ "]" ++ dims rowshape + dims _ = "" + +-- | Create an array value; failing if that would result in an +-- irregular array. +mkArray :: TypeBase Int64 () -> [Value] -> Maybe Value +mkArray t [] = + return $ toArray (typeShape mempty t) [] +mkArray _ (v : vs) = do + let v_shape = valueShape v + guard $ all ((== v_shape) . valueShape) vs + return $ toArray' v_shape $ v : vs + +arrayLength :: Integral int => Array Int Value -> int +arrayLength = fromIntegral . (+ 1) . snd . bounds + +toTuple :: [Value] -> Value +toTuple = ValueRecord . M.fromList . zip tupleFieldNames + +fromTuple :: Value -> Maybe [Value] +fromTuple (ValueRecord m) = areTupleFields m +fromTuple _ = Nothing + +asInteger :: Value -> Integer +asInteger (ValuePrim (SignedValue v)) = P.valueIntegral v +asInteger (ValuePrim (UnsignedValue v)) = + toInteger (P.valueIntegral (P.doZExt v Int64) :: Word64) +asInteger v = error $ "Unexpectedly not an integer: " ++ pretty v + +asInt :: Value -> Int +asInt = fromIntegral . asInteger + +asSigned :: Value -> IntValue +asSigned (ValuePrim (SignedValue v)) = v +asSigned v = error $ "Unexpected not a signed integer: " ++ pretty v + +asInt64 :: Value -> Int64 +asInt64 = fromIntegral . asInteger + +asBool :: Value -> Bool +asBool (ValuePrim (BoolValue x)) = x +asBool v = error $ "Unexpectedly not a boolean: " ++ pretty v + +lookupInEnv :: + (Env -> M.Map VName x) -> + QualName VName -> + Env -> + Maybe x +lookupInEnv onEnv qv env = f env $ qualQuals qv + where + f m (q : qs) = + case M.lookup q $ envTerm m of + Just (TermModule (Module mod)) -> f mod qs + _ -> Nothing + f m [] = M.lookup (qualLeaf qv) $ onEnv m + +lookupVar :: QualName VName -> Env -> Maybe TermBinding +lookupVar = lookupInEnv envTerm + +lookupType :: QualName VName -> Env -> Maybe T.TypeBinding +lookupType = lookupInEnv envType + +-- | A TermValue with a 'Nothing' type annotation is an intrinsic. +data TermBinding + = TermValue (Maybe T.BoundV) Value + | -- | A polymorphic value that must be instantiated. + TermPoly (Maybe T.BoundV) (StructType -> EvalM Value) + | TermModule Module + +data Module + = Module Env + | ModuleFun (Module -> EvalM Module) + +-- | The actual type- and value environment. +data Env = Env + { envTerm :: M.Map VName TermBinding, + envType :: M.Map VName T.TypeBinding, + -- | A mapping from type parameters to the shapes of + -- the value to which they were initially bound. + envShapes :: M.Map VName ValueShape + } + +instance Monoid Env where + mempty = Env mempty mempty mempty + +instance Semigroup Env where + Env vm1 tm1 sm1 <> Env vm2 tm2 sm2 = + Env (vm1 <> vm2) (tm1 <> tm2) (sm1 <> sm2) + +-- | An error occurred during interpretation due to an error in the +-- user program. Actual interpreter errors will be signaled with an +-- IO exception ('error'). +newtype InterpreterError = InterpreterError String + +valEnv :: M.Map VName (Maybe T.BoundV, Value) -> Env +valEnv m = + Env + { envTerm = M.map (uncurry TermValue) m, + envType = mempty, + envShapes = mempty + } + +modEnv :: M.Map VName Module -> Env +modEnv m = + Env + { envTerm = M.map TermModule m, + envType = mempty, + envShapes = mempty + } + +typeEnv :: M.Map VName StructType -> Env +typeEnv m = + Env + { envTerm = mempty, + envType = M.map tbind m, + envShapes = mempty + } + where + tbind = T.TypeAbbr Unlifted [] + +i64Env :: M.Map VName Int64 -> Env +i64Env = valEnv . M.map f + where + f x = + ( Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int64, + ValuePrim $ SignedValue $ Int64Value x + ) + +instance Show InterpreterError where + show (InterpreterError s) = s + +bad :: SrcLoc -> Env -> String -> EvalM a +bad loc env s = stacking loc env $ do + ss <- map (locStr . srclocOf) <$> stacktrace + liftF $ ExtOpError $ InterpreterError $ "Error at\n" ++ prettyStacktrace 0 ss ++ s + +trace :: Value -> EvalM () +trace v = do + -- We take the second-to-top element of the stack, because any + -- actual call to 'implicits.trace' is going to be in the trace + -- function in the prelude, which is not interesting. + top <- fromMaybe noLoc . maybeHead . drop 1 <$> stacktrace + liftF $ ExtOpTrace top (prettyOneLine v) () + +typeCheckerEnv :: Env -> T.Env +typeCheckerEnv env = + -- FIXME: some shadowing issues are probably not right here. + let valMap (TermValue (Just t) _) = Just t + valMap _ = Nothing + vtable = M.mapMaybe valMap $ envTerm env + nameMap k + | k `M.member` vtable = Just ((T.Term, baseName k), qualName k) + | otherwise = Nothing + in mempty + { T.envNameMap = M.fromList $ mapMaybe nameMap $ M.keys $ envTerm env, + T.envVtable = vtable + } + +break :: EvalM () +break = do + -- We don't want the env of the function that is calling + -- intrinsics.break, since that is just going to be the boring + -- wrapper function (intrinsics are never called directly). + -- This is why we go a step up the stack. + backtrace <- asks $ drop 1 . fst + case NE.nonEmpty backtrace of + Nothing -> return () + Just backtrace' -> liftF $ ExtOpBreak BreakPoint backtrace' () + +fromArray :: Value -> (ValueShape, [Value]) +fromArray (ValueArray shape as) = (shape, elems as) +fromArray v = error $ "Expected array value, but found: " ++ pretty v + +toArray :: ValueShape -> [Value] -> Value +toArray shape vs = ValueArray shape (listArray (0, length vs - 1) vs) + +toArray' :: ValueShape -> [Value] -> Value +toArray' rowshape vs = ValueArray shape (listArray (0, length vs - 1) vs) + where + shape = ShapeDim (genericLength vs) rowshape + +apply :: SrcLoc -> Env -> Value -> Value -> EvalM Value +apply loc env (ValueFun f) v = stacking loc env (f v) +apply _ _ f _ = error $ "Cannot apply non-function: " ++ pretty f + +apply2 :: SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value +apply2 loc env f x y = stacking loc env $ do + f' <- apply noLoc mempty f x + apply noLoc mempty f' y + +matchPattern :: Env -> Pattern -> Value -> EvalM Env +matchPattern env p v = do + m <- runMaybeT $ patternMatch env p v + case m of + Nothing -> error $ "matchPattern: missing case for " ++ pretty p ++ " and " ++ pretty v + Just env' -> return env' + +patternMatch :: Env -> Pattern -> Value -> MaybeT EvalM Env +patternMatch env (Id v (Info t) _) val = + lift $ + pure $ + valEnv (M.singleton v (Just $ T.BoundV [] $ toStruct t, val)) <> env +patternMatch env Wildcard {} _ = + lift $ pure env +patternMatch env (TuplePattern ps _) (ValueRecord vs) = + foldM (\env' (p, v) -> patternMatch env' p v) env $ + zip ps (map snd $ sortFields vs) +patternMatch env (RecordPattern ps _) (ValueRecord vs) = + foldM (\env' (p, v) -> patternMatch env' p v) env $ + M.intersectionWith (,) (M.fromList ps) vs +patternMatch env (PatternParens p _) v = patternMatch env p v +patternMatch env (PatternAscription p _ _) v = + patternMatch env p v +patternMatch env (PatternLit e _ _) v = do + v' <- lift $ eval env e + if v == v' + then pure env + else mzero +patternMatch env (PatternConstr n _ ps _) (ValueSum _ n' vs) + | n == n' = + foldM (\env' (p, v) -> patternMatch env' p v) env $ zip ps vs +patternMatch _ _ _ = mzero + +data Indexing + = IndexingFix Int64 + | IndexingSlice (Maybe Int64) (Maybe Int64) (Maybe Int64) + +instance Pretty Indexing where + ppr (IndexingFix i) = ppr i + ppr (IndexingSlice i j (Just s)) = + maybe mempty ppr i <> text ":" + <> maybe mempty ppr j + <> text ":" + <> ppr s + ppr (IndexingSlice i (Just j) s) = + maybe mempty ppr i <> text ":" + <> ppr j + <> maybe mempty ((text ":" <>) . ppr) s + ppr (IndexingSlice i Nothing Nothing) = + maybe mempty ppr i <> text ":" + +indexesFor :: + Maybe Int64 -> + Maybe Int64 -> + Maybe Int64 -> + Int64 -> + Maybe [Int] +indexesFor start end stride n + | (start', end', stride') <- slice, + end' == start' || signum' (end' - start') == signum' stride', + stride' /= 0, + is <- [start', start' + stride' .. end' - signum stride'], + all inBounds is = + Just $ map fromIntegral is + | otherwise = + Nothing + where + inBounds i = i >= 0 && i < n + + slice = + case (start, end, stride) of + (Just start', _, _) -> + let end' = fromMaybe n end + in (start', end', fromMaybe 1 stride) + (Nothing, Just end', _) -> + let start' = 0 + in (start', end', fromMaybe 1 stride) + (Nothing, Nothing, Just stride') -> + ( if stride' > 0 then 0 else n -1, + if stride' > 0 then n else -1, + stride' + ) + (Nothing, Nothing, Nothing) -> + (0, n, 1) + +-- | 'signum', but with 0 as 1. +signum' :: (Eq p, Num p) => p -> p +signum' 0 = 1 +signum' x = signum x + +indexShape :: [Indexing] -> ValueShape -> ValueShape +indexShape (IndexingFix {} : is) (ShapeDim _ shape) = + indexShape is shape +indexShape (IndexingSlice start end stride : is) (ShapeDim d shape) = + ShapeDim n $ indexShape is shape + where + n = maybe 0 genericLength $ indexesFor start end stride d +indexShape _ shape = + shape + +indexArray :: [Indexing] -> Value -> Maybe Value +indexArray (IndexingFix i : is) (ValueArray _ arr) + | i >= 0, + i < n = + indexArray is $ arr ! fromIntegral i + | otherwise = + Nothing + where + n = arrayLength arr +indexArray (IndexingSlice start end stride : is) (ValueArray (ShapeDim _ rowshape) arr) = do + js <- indexesFor start end stride $ arrayLength arr + toArray' (indexShape is rowshape) <$> mapM (indexArray is . (arr !)) js +indexArray _ v = Just v + +updateArray :: [Indexing] -> Value -> Value -> Maybe Value +updateArray (IndexingFix i : is) (ValueArray shape arr) v + | i >= 0, + i < n = do + v' <- updateArray is (arr ! i') v + Just $ ValueArray shape $ arr // [(i', v')] + | otherwise = + Nothing + where + n = arrayLength arr + i' = fromIntegral i +updateArray (IndexingSlice start end stride : is) (ValueArray shape arr) (ValueArray _ v) = do + arr_is <- indexesFor start end stride $ arrayLength arr + guard $ length arr_is == arrayLength v + let update arr' (i, v') = do + x <- updateArray is (arr ! i) v' + return $ arr' // [(i, x)] + fmap (ValueArray shape) $ foldM update arr $ zip arr_is $ elems v +updateArray _ _ v = Just v + +evalDimIndex :: Env -> DimIndex -> EvalM Indexing +evalDimIndex env (DimFix x) = + IndexingFix . asInt64 <$> eval env x +evalDimIndex env (DimSlice start end stride) = + IndexingSlice <$> traverse (fmap asInt64 . eval env) start + <*> traverse (fmap asInt64 . eval env) end + <*> traverse (fmap asInt64 . eval env) stride + +evalIndex :: SrcLoc -> Env -> [Indexing] -> Value -> EvalM Value +evalIndex loc env is arr = do + let oob = + bad loc env $ + "Index [" <> intercalate ", " (map pretty is) + <> "] out of bounds for array of shape " + <> pretty (valueShape arr) + <> "." + maybe oob return $ indexArray is arr + +-- | Expand type based on information that was not available at +-- type-checking time (the structure of abstract types). +evalType :: Env -> StructType -> StructType +evalType _ (Scalar (Prim pt)) = Scalar $ Prim pt +evalType env (Scalar (Record fs)) = Scalar $ Record $ fmap (evalType env) fs +evalType env (Scalar (Arrow () p t1 t2)) = + Scalar $ Arrow () p (evalType env t1) (evalType env t2) +evalType env t@(Array _ u _ shape) = + let et = stripArray (shapeRank shape) t + et' = evalType env et + shape' = fmap evalDim shape + in arrayOf et' shape' u + where + evalDim (NamedDim qn) + | Just (TermValue _ (ValuePrim (SignedValue (Int64Value x)))) <- + lookupVar qn env = + ConstDim $ fromIntegral x + evalDim d = d +evalType env t@(Scalar (TypeVar () _ tn args)) = + case lookupType (qualNameFromTypeName tn) env of + Just (T.TypeAbbr _ ps t') -> + let (substs, types) = mconcat $ zipWith matchPtoA ps args + onDim (NamedDim v) = fromMaybe (NamedDim v) $ M.lookup (qualLeaf v) substs + onDim d = d + in if null ps + then first onDim t' + else evalType (Env mempty types mempty <> env) $ first onDim t' + Nothing -> t + where + matchPtoA (TypeParamDim p _) (TypeArgDim (NamedDim qv) _) = + (M.singleton p $ NamedDim qv, mempty) + matchPtoA (TypeParamDim p _) (TypeArgDim (ConstDim k) _) = + (M.singleton p $ ConstDim k, mempty) + matchPtoA (TypeParamType l p _) (TypeArgType t' _) = + let t'' = evalType env t' + in (mempty, M.singleton p $ T.TypeAbbr l [] t'') + matchPtoA _ _ = mempty +evalType env (Scalar (Sum cs)) = Scalar $ Sum $ (fmap . fmap) (evalType env) cs + +evalTermVar :: Env -> QualName VName -> StructType -> EvalM Value +evalTermVar env qv t = + case lookupVar qv env of + Just (TermPoly _ v) -> do + size_env <- extSizeEnv + v $ evalType (size_env <> env) t + Just (TermValue _ v) -> return v + _ -> error $ "`" <> pretty qv <> "` is not bound to a value." + +typeValueShape :: Env -> StructType -> EvalM ValueShape +typeValueShape env t = do + size_env <- extSizeEnv + let t' = evalType (size_env <> env) t + case traverse dim $ typeShape mempty t' of + Nothing -> error $ "typeValueShape: failed to fully evaluate type " ++ pretty t' + Just shape -> return shape + where + dim (ConstDim x) = Just $ fromIntegral x + dim _ = Nothing + +evalFunction :: Env -> [VName] -> [Pattern] -> Exp -> StructType -> EvalM Value +-- We treat zero-parameter lambdas as simply an expression to +-- evaluate immediately. Note that this is *not* the same as a lambda +-- that takes an empty tuple '()' as argument! Zero-parameter lambdas +-- can never occur in a well-formed Futhark program, but they are +-- convenient in the interpreter. +evalFunction env _ [] body rettype = + -- Eta-expand the rest to make any sizes visible. + etaExpand [] env rettype + where + etaExpand vs env' (Scalar (Arrow _ _ pt rt)) = + return $ + ValueFun $ \v -> do + env'' <- matchPattern env' (Wildcard (Info $ fromStruct pt) noLoc) v + etaExpand (v : vs) env'' rt + etaExpand vs env' _ = do + f <- eval env' body + foldM (apply noLoc mempty) f $ reverse vs +evalFunction env missing_sizes (p : ps) body rettype = + return $ + ValueFun $ \v -> do + env' <- matchPattern env p v + -- Fix up the last sizes, if any. + let env'' + | null missing_sizes = env' + | otherwise = + env' + <> i64Env + ( resolveExistentials + missing_sizes + (patternStructType p) + (valueShape v) + ) + evalFunction env'' missing_sizes ps body rettype + +evalFunctionBinding :: + Env -> + [TypeParam] -> + [Pattern] -> + StructType -> + [VName] -> + Exp -> + EvalM TermBinding +evalFunctionBinding env tparams ps ret retext fbody = do + let ret' = evalType env ret + arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt + ftype = foldr (arrow . patternParam) ret' ps + + -- Distinguish polymorphic and non-polymorphic bindings here. + if null tparams + then + TermValue (Just $ T.BoundV [] ftype) + <$> (returned env ret retext =<< evalFunction env [] ps fbody ret') + else return $ + TermPoly (Just $ T.BoundV [] ftype) $ \ftype' -> do + let tparam_names = map typeParamName tparams + env' = resolveTypeParams tparam_names ftype ftype' <> env + + -- In some cases (abstract lifted types) there may be + -- missing sizes that were not fixed by the type + -- instantiation. These will have to be set by looking + -- at the actual function arguments. + missing_sizes = + filter (`M.notMember` envTerm env') $ + map typeParamName (filter isSizeParam tparams) + returned env ret retext =<< evalFunction env' missing_sizes ps fbody ret' + +evalArg :: Env -> Exp -> Maybe VName -> EvalM Value +evalArg env e ext = do + v <- eval env e + case ext of + Just ext' -> putExtSize ext' $ asInt64 v + Nothing -> return () + return v + +returned :: Env -> TypeBase (DimDecl VName) als -> [VName] -> Value -> EvalM Value +returned _ _ [] v = return v +returned env ret retext v = do + mapM_ (uncurry putExtSize) $ + M.toList $ + resolveExistentials retext (evalType env $ toStruct ret) $ valueShape v + return v + +eval :: Env -> Exp -> EvalM Value +eval _ (Literal v _) = return $ ValuePrim v +eval env (Parens e _) = eval env e +eval env (QualParens (qv, _) e loc) = do + m <- evalModuleVar env qv + case m of + ModuleFun {} -> error $ "Local open of module function at " ++ locStr loc + Module m' -> eval (m' <> env) e +eval env (TupLit vs _) = toTuple <$> mapM (eval env) vs +eval env (RecordLit fields _) = + ValueRecord . M.fromList <$> mapM evalField fields + where + evalField (RecordFieldExplicit k e _) = do + v <- eval env e + return (k, v) + evalField (RecordFieldImplicit k t loc) = do + v <- eval env $ Var (qualName k) t loc + return (baseName k, v) +eval _ (StringLit vs _) = + return $ + toArray' ShapeLeaf $ + map (ValuePrim . UnsignedValue . Int8Value . fromIntegral) vs +eval env (ArrayLit [] (Info t) _) = do + t' <- typeValueShape env $ toStruct t + return $ toArray t' [] +eval env (ArrayLit (v : vs) _ _) = do + v' <- eval env v + vs' <- mapM (eval env) vs + return $ toArray' (valueShape v') (v' : vs') +eval env (Range start maybe_second end (Info t, Info retext) loc) = do + start' <- asInteger <$> eval env start + maybe_second' <- traverse (fmap asInteger . eval env) maybe_second + end' <- traverse (fmap asInteger . eval env) end + + let (end_adj, step, ok) = + case (end', maybe_second') of + (DownToExclusive end'', Nothing) -> + (end'' + 1, -1, start' >= end'') + (DownToExclusive end'', Just second') -> + (end'' + 1, second' - start', start' >= end'' && second' < start') + (ToInclusive end'', Nothing) -> + (end'', 1, start' <= end'') + (ToInclusive end'', Just second') + | second' > start' -> + (end'', second' - start', start' <= end'') + | otherwise -> + (end'', second' - start', start' >= end'' && second' /= start') + (UpToExclusive x, Nothing) -> + (x -1, 1, start' <= x) + (UpToExclusive x, Just second') -> + (x -1, second' - start', start' <= x && second' > start') + + if ok + then + returned env t retext $ + toArray' ShapeLeaf $ map toInt [start', start' + step .. end_adj] + else bad loc env $ badRange start' maybe_second' end' + where + toInt = + case stripArray 1 t of + Scalar (Prim (Signed t')) -> + ValuePrim . SignedValue . intValue t' + Scalar (Prim (Unsigned t')) -> + ValuePrim . UnsignedValue . intValue t' + _ -> error $ "Nonsensical range type: " ++ show t + + badRange start' maybe_second' end' = + "Range " ++ pretty start' + ++ ( case maybe_second' of + Nothing -> "" + Just second' -> ".." ++ pretty second' + ) + ++ ( case end' of + DownToExclusive x -> "..>" ++ pretty x + ToInclusive x -> "..." ++ pretty x + UpToExclusive x -> "..<" ++ pretty x + ) + ++ " is invalid." +eval env (Var qv (Info t) _) = evalTermVar env qv (toStruct t) +eval env (Ascript e _ _) = eval env e +eval env (Coerce e td (Info ret, Info retext) loc) = do + v <- returned env ret retext =<< eval env e + let t = evalType env $ unInfo $ expandedType td + case checkShape (structTypeShape (envShapes env) t) (valueShape v) of + Just _ -> return v + Nothing -> + bad loc env $ + "Value `" <> pretty v <> "` of shape `" ++ pretty (valueShape v) + ++ "` cannot match shape of type `" + <> pretty (declaredType td) + <> "` (`" + <> pretty t + <> "`)" +eval env (LetPat p e body (Info ret, Info retext) _) = do + v <- eval env e + env' <- matchPattern env p v + returned env ret retext =<< eval env' body +eval env (LetFun f (tparams, ps, _, Info ret, fbody) body _ _) = do + binding <- evalFunctionBinding env tparams ps ret [] fbody + eval (env {envTerm = M.insert f binding $ envTerm env}) body +eval _ (IntLit v (Info t) _) = + case t of + Scalar (Prim (Signed it)) -> + return $ ValuePrim $ SignedValue $ intValue it v + Scalar (Prim (Unsigned it)) -> + return $ ValuePrim $ UnsignedValue $ intValue it v + Scalar (Prim (FloatType ft)) -> + return $ ValuePrim $ FloatValue $ floatValue ft v + _ -> error $ "eval: nonsensical type for integer literal: " ++ pretty t +eval _ (FloatLit v (Info t) _) = + case t of + Scalar (Prim (FloatType ft)) -> + return $ ValuePrim $ FloatValue $ floatValue ft v + _ -> error $ "eval: nonsensical type for float literal: " ++ pretty t +eval + env + ( BinOp + (op, _) + op_t + (x, Info (_, xext)) + (y, Info (_, yext)) + (Info t) + (Info retext) + loc + ) + | baseString (qualLeaf op) == "&&" = do + x' <- asBool <$> eval env x + if x' + then eval env y + else return $ ValuePrim $ BoolValue False + | baseString (qualLeaf op) == "||" = do + x' <- asBool <$> eval env x + if x' + then return $ ValuePrim $ BoolValue True + else eval env y + | otherwise = do + op' <- eval env $ Var op op_t loc + x' <- evalArg env x xext + y' <- evalArg env y yext + returned env t retext =<< apply2 loc env op' x' y' +eval env (If cond e1 e2 (Info ret, Info retext) _) = do + cond' <- asBool <$> eval env cond + returned env ret retext + =<< if cond' then eval env e1 else eval env e2 +eval env (Apply f x (Info (_, ext)) (Info t, Info retext) loc) = do + -- It is important that 'x' is evaluated first in order to bring any + -- sizes into scope that may be used in the type of 'f'. + x' <- evalArg env x ext + f' <- eval env f + returned env t retext =<< apply loc env f' x' +eval env (Negate e _) = do + ev <- eval env e + ValuePrim <$> case ev of + ValuePrim (SignedValue (Int8Value v)) -> return $ SignedValue $ Int8Value (- v) + ValuePrim (SignedValue (Int16Value v)) -> return $ SignedValue $ Int16Value (- v) + ValuePrim (SignedValue (Int32Value v)) -> return $ SignedValue $ Int32Value (- v) + ValuePrim (SignedValue (Int64Value v)) -> return $ SignedValue $ Int64Value (- v) + ValuePrim (UnsignedValue (Int8Value v)) -> return $ UnsignedValue $ Int8Value (- v) + ValuePrim (UnsignedValue (Int16Value v)) -> return $ UnsignedValue $ Int16Value (- v) + ValuePrim (UnsignedValue (Int32Value v)) -> return $ UnsignedValue $ Int32Value (- v) + ValuePrim (UnsignedValue (Int64Value v)) -> return $ UnsignedValue $ Int64Value (- v) + ValuePrim (FloatValue (Float32Value v)) -> return $ FloatValue $ Float32Value (- v) + ValuePrim (FloatValue (Float64Value v)) -> return $ FloatValue $ Float64Value (- v) + _ -> error $ "Cannot negate " ++ pretty ev +eval env (Index e is (Info t, Info retext) loc) = do + is' <- mapM (evalDimIndex env) is + arr <- eval env e + returned env t retext =<< evalIndex loc env is' arr +eval env (Update src is v loc) = + maybe oob return + =<< updateArray <$> mapM (evalDimIndex env) is <*> eval env src <*> eval env v + where + oob = bad loc env "Bad update" +eval env (RecordUpdate src all_fs v _ _) = + update <$> eval env src <*> pure all_fs <*> eval env v + where + update _ [] v' = v' + update (ValueRecord src') (f : fs) v' + | Just f_v <- M.lookup f src' = + ValueRecord $ M.insert f (update f_v fs v') src' + update _ _ _ = error "eval RecordUpdate: invalid value." +eval env (LetWith dest src is v body _ loc) = do + let Ident src_vn (Info src_t) _ = src + dest' <- + maybe oob return + =<< updateArray <$> mapM (evalDimIndex env) is + <*> evalTermVar env (qualName src_vn) (toStruct src_t) + <*> eval env v + let t = T.BoundV [] $ toStruct $ unInfo $ identType dest + eval (valEnv (M.singleton (identName dest) (Just t, dest')) <> env) body + where + oob = bad loc env "Bad update" + +-- We treat zero-parameter lambdas as simply an expression to +-- evaluate immediately. Note that this is *not* the same as a lambda +-- that takes an empty tuple '()' as argument! Zero-parameter lambdas +-- can never occur in a well-formed Futhark program, but they are +-- convenient in the interpreter. +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 + 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 + y <- evalArg env e argext + return $ + ValueFun $ \x -> do + f <- evalTermVar env qv $ toStruct t + apply2 loc env f x y +eval env (IndexSection is _ loc) = do + is' <- mapM (evalDimIndex env) is + return $ ValueFun $ evalIndex loc env is' +eval _ (ProjectSection ks _ _) = + return $ ValueFun $ flip (foldM walk) ks + where + walk (ValueRecord fs) f + | Just v' <- M.lookup f fs = return v' + walk _ _ = error "Value does not have expected field." +eval env (DoLoop sparams pat init_e form body (Info (ret, retext)) _) = do + init_v <- eval env init_e + returned env ret retext + =<< case form of + For iv bound -> do + bound' <- asSigned <$> eval env bound + forLoop (identName iv) bound' (zero bound') init_v + ForIn in_pat in_e -> do + (_, in_vs) <- fromArray <$> eval env in_e + foldM (forInLoop in_pat) init_v in_vs + While cond -> + whileLoop cond init_v + where + withLoopParams v = + let sparams' = + resolveExistentials + sparams + (patternStructType pat) + (valueShape v) + in matchPattern (i64Env sparams' <> env) pat v + + inc = (`P.doAdd` Int64Value 1) + zero = (`P.doMul` Int64Value 0) + + forLoop iv bound i v + | i >= bound = return v + | otherwise = do + env' <- withLoopParams v + forLoop iv bound (inc i) + =<< eval + ( valEnv + ( M.singleton + iv + ( Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int64, + ValuePrim (SignedValue i) + ) + ) + <> env' + ) + body + + whileLoop cond v = do + env' <- withLoopParams v + continue <- asBool <$> eval env' cond + if continue + then whileLoop cond =<< eval env' body + else return v + + forInLoop in_pat v in_v = do + env' <- withLoopParams v + env'' <- matchPattern env' in_pat in_v + eval env'' body +eval env (Project f e _ _) = do + v <- eval env e + case v of + ValueRecord fs | Just v' <- M.lookup f fs -> return v' + _ -> error "Value does not have expected field." +eval env (Assert what e (Info s) loc) = do + cond <- asBool <$> eval env what + unless cond $ bad loc env s + eval env e +eval env (Constr c es (Info t) _) = do + vs <- mapM (eval env) es + shape <- typeValueShape env $ toStruct t + return $ ValueSum shape c vs +eval env (Match e cs (Info ret, Info retext) _) = do + v <- eval env e + returned env ret retext =<< match v (NE.toList cs) + where + match _ [] = + error "Pattern match failure." + match v (c : cs') = do + c' <- evalCase v env c + case c' of + Just v' -> return v' + Nothing -> match v cs' +eval env (Attr _ e _) = eval env e + +evalCase :: + Value -> + Env -> + CaseBase Info VName -> + EvalM (Maybe Value) +evalCase v env (CasePat p cExp _) = runMaybeT $ do + env' <- patternMatch env p v + lift $ eval env' cExp + +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 + replaceM f m = M.fromList $ do + (k, v) <- M.toList m + return (replace 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) + onTerm (TermValue t v) = TermValue t v + onTerm (TermPoly t v) = TermPoly t v + onTerm (TermModule m) = TermModule $ onModule m + onType (T.TypeAbbr l ps t) = T.TypeAbbr l ps $ first onDim t + onDim (NamedDim v) = NamedDim $ replaceQ v + 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 + Just (TermModule m) -> return m + _ -> error $ quote (pretty qv) <> " is not bound to a module." + +evalModExp :: Env -> ModExp -> EvalM Module +evalModExp _ (ModImport _ (Info f) _) = do + f' <- lookupImport f + case f' of + Nothing -> error $ "Unknown import " ++ show f + Just m -> return $ Module m +evalModExp env (ModDecs ds _) = do + Env terms types _ <- foldM evalDec env ds + -- Remove everything that was present in the original Env. + return $ + Module $ + Env + (terms `M.difference` envTerm env) + (types `M.difference` envType env) + mempty +evalModExp env (ModVar qv _) = + evalModuleVar env qv +evalModExp env (ModAscript me _ (Info substs) _) = + substituteInModule substs <$> evalModExp env me +evalModExp env (ModParens me _) = evalModExp env me +evalModExp env (ModLambda p ret e loc) = + return $ + ModuleFun $ \am -> do + let env' = env {envTerm = M.insert (modParamName p) (TermModule am) $ envTerm env} + evalModExp env' $ case ret of + Nothing -> e + Just (se, rsubsts) -> ModAscript e se rsubsts loc +evalModExp env (ModApply f e (Info psubst) (Info rsubst) _) = do + f' <- evalModExp env f + case f' of + ModuleFun f'' -> do + e' <- evalModExp env e + substituteInModule rsubst <$> f'' (substituteInModule psubst e') + _ -> error "Expected ModuleFun." + +evalDec :: Env -> Dec -> EvalM Env +evalDec env (ValDec (ValBind _ v _ (Info (ret, retext)) tparams ps fbody _ _ _)) = do + binding <- evalFunctionBinding env tparams ps ret retext fbody + return $ env {envTerm = M.insert v binding $ envTerm env} +evalDec env (OpenDec me _) = do + me' <- evalModExp env me + case me' of + Module me'' -> return $ me'' <> env + _ -> error "Expected Module" +evalDec env (ImportDec name name' loc) = + evalDec env $ LocalDec (OpenDec (ModImport name name' loc) loc) loc +evalDec env (LocalDec d _) = evalDec env d +evalDec env SigDec {} = return env +evalDec env (TypeDec (TypeBind v l ps t _ _)) = do + let abbr = + T.TypeAbbr l ps $ + evalType env $ unInfo $ expandedType t + return env {envType = M.insert v abbr $ envType env} +evalDec env (ModDec (ModBind v ps ret body _ loc)) = do + mod <- evalModExp env $ wrapInLambda ps + return $ modEnv (M.singleton v mod) <> env + where + wrapInLambda [] = case ret of + Just (se, substs) -> ModAscript body se substs loc + Nothing -> body + wrapInLambda [p] = ModLambda p ret body loc + wrapInLambda (p : ps') = ModLambda p Nothing (wrapInLambda ps') loc + +-- | The interpreter context. All evaluation takes place with respect +-- to a context, and it can be extended with more definitions, which +-- is how the REPL works. +data Ctx = Ctx + { ctxEnv :: Env, + ctxImports :: M.Map FilePath Env + } + +nanValue :: PrimValue -> Bool +nanValue (FloatValue v) = + case v of + Float32Value x -> isNaN x + Float64Value x -> isNaN x +nanValue _ = False + +breakOnNaN :: [PrimValue] -> PrimValue -> EvalM () +breakOnNaN inputs result + | not (any nanValue inputs) && nanValue result = do + backtrace <- asks fst + case NE.nonEmpty backtrace of + Nothing -> return () + Just backtrace' -> liftF $ ExtOpBreak BreakNaN backtrace' () +breakOnNaN _ _ = + return () + +-- | The initial environment contains definitions of the various intrinsic functions. +initialCtx :: Ctx +initialCtx = + Ctx + ( Env + ( M.insert + (VName (nameFromString "intrinsics") 0) + (TermModule (Module $ Env terms types mempty)) + terms + ) + types + mempty + ) + mempty + where + terms = M.mapMaybeWithKey (const . def . baseString) intrinsics + types = M.mapMaybeWithKey (const . tdef . baseString) intrinsics + + sintOp f = + [ (getS, putS, P.doBinOp (f Int8)), + (getS, putS, P.doBinOp (f Int16)), + (getS, putS, P.doBinOp (f Int32)), + (getS, putS, P.doBinOp (f Int64)) + ] + uintOp f = + [ (getU, putU, P.doBinOp (f Int8)), + (getU, putU, P.doBinOp (f Int16)), + (getU, putU, P.doBinOp (f Int32)), + (getU, putU, P.doBinOp (f Int64)) + ] + intOp f = sintOp f ++ uintOp f + floatOp f = + [ (getF, putF, P.doBinOp (f Float32)), + (getF, putF, P.doBinOp (f Float64)) + ] + arithOp f g = Just $ bopDef $ intOp f ++ floatOp g + + flipCmps = map (\(f, g, h) -> (f, g, flip h)) + sintCmp f = + [ (getS, Just . BoolValue, P.doCmpOp (f Int8)), + (getS, Just . BoolValue, P.doCmpOp (f Int16)), + (getS, Just . BoolValue, P.doCmpOp (f Int32)), + (getS, Just . BoolValue, P.doCmpOp (f Int64)) + ] + uintCmp f = + [ (getU, Just . BoolValue, P.doCmpOp (f Int8)), + (getU, Just . BoolValue, P.doCmpOp (f Int16)), + (getU, Just . BoolValue, P.doCmpOp (f Int32)), + (getU, Just . BoolValue, P.doCmpOp (f Int64)) + ] + floatCmp f = + [ (getF, Just . BoolValue, P.doCmpOp (f Float32)), + (getF, Just . BoolValue, P.doCmpOp (f Float64)) + ] + boolCmp f = [(getB, Just . BoolValue, P.doCmpOp f)] + + getV (SignedValue x) = Just $ P.IntValue x + getV (UnsignedValue x) = Just $ P.IntValue x + getV (FloatValue x) = Just $ P.FloatValue x + getV (BoolValue x) = Just $ P.BoolValue x + putV (P.IntValue x) = SignedValue x + putV (P.FloatValue x) = FloatValue x + putV (P.BoolValue x) = BoolValue x + putV P.Checked = BoolValue True + + getS (SignedValue x) = Just $ P.IntValue x + getS _ = Nothing + putS (P.IntValue x) = Just $ SignedValue x + putS _ = Nothing + + getU (UnsignedValue x) = Just $ P.IntValue x + getU _ = Nothing + putU (P.IntValue x) = Just $ UnsignedValue x + putU _ = Nothing + + getF (FloatValue x) = Just $ P.FloatValue x + getF _ = Nothing + putF (P.FloatValue x) = Just $ FloatValue x + putF _ = Nothing + + getB (BoolValue x) = Just $ P.BoolValue x + getB _ = Nothing + putB (P.BoolValue x) = Just $ BoolValue x + putB _ = Nothing + + fun1 f = + TermValue Nothing $ ValueFun $ \x -> f x + fun2 f = + TermValue Nothing $ + ValueFun $ \x -> + return $ ValueFun $ \y -> f x y + fun2t f = + TermValue Nothing $ + ValueFun $ \v -> + case fromTuple v of + Just [x, y] -> f x y + _ -> error $ "Expected pair; got: " ++ pretty v + fun3t f = + TermValue Nothing $ + ValueFun $ \v -> + case fromTuple v of + Just [x, y, z] -> f x y z + _ -> error $ "Expected triple; got: " ++ pretty v + + fun6t f = + TermValue Nothing $ + ValueFun $ \v -> + case fromTuple v of + Just [x, y, z, a, b, c] -> f x y z a b c + _ -> error $ "Expected sextuple; got: " ++ pretty v + + bopDef fs = fun2 $ \x y -> + case (x, y) of + (ValuePrim x', ValuePrim y') + | Just z <- msum $ map (`bopDef'` (x', y')) fs -> do + breakOnNaN [x', y'] z + return $ ValuePrim z + _ -> + bad noLoc mempty $ + "Cannot apply operator to arguments " + <> quote (pretty x) + <> " and " + <> quote (pretty y) + <> "." + where + bopDef' (valf, retf, op) (x, y) = do + x' <- valf x + y' <- valf y + retf =<< op x' y' + + unopDef fs = fun1 $ \x -> + case x of + (ValuePrim x') + | Just r <- msum $ map (`unopDef'` x') fs -> do + breakOnNaN [x'] r + return $ ValuePrim r + _ -> + bad noLoc mempty $ + "Cannot apply function to argument " + <> quote (pretty x) + <> "." + where + unopDef' (valf, retf, op) x = do + x' <- valf x + retf =<< op x' + + tbopDef f = fun1 $ \v -> + case fromTuple v of + Just [ValuePrim x, ValuePrim y] + | Just x' <- getV x, + Just y' <- getV y, + Just z <- putV <$> f x' y' -> do + breakOnNaN [x, y] z + return $ ValuePrim z + _ -> + bad noLoc mempty $ + "Cannot apply operator to argument " + <> quote (pretty v) + <> "." + + def "!" = + Just $ + unopDef + [ (getS, putS, P.doUnOp $ P.Complement Int8), + (getS, putS, P.doUnOp $ P.Complement Int16), + (getS, putS, P.doUnOp $ P.Complement Int32), + (getS, putS, P.doUnOp $ P.Complement Int64), + (getU, putU, P.doUnOp $ P.Complement Int8), + (getU, putU, P.doUnOp $ P.Complement Int16), + (getU, putU, P.doUnOp $ P.Complement Int32), + (getU, putU, P.doUnOp $ P.Complement Int64), + (getB, putB, P.doUnOp P.Not) + ] + def "+" = arithOp (`P.Add` P.OverflowWrap) P.FAdd + def "-" = arithOp (`P.Sub` P.OverflowWrap) P.FSub + def "*" = arithOp (`P.Mul` P.OverflowWrap) P.FMul + def "**" = arithOp P.Pow P.FPow + def "/" = + Just $ + bopDef $ + sintOp (`P.SDiv` P.Unsafe) + ++ uintOp (`P.UDiv` P.Unsafe) + ++ floatOp P.FDiv + def "%" = + Just $ + bopDef $ + sintOp (`P.SMod` P.Unsafe) + ++ uintOp (`P.UMod` P.Unsafe) + ++ floatOp P.FMod + def "//" = + Just $ + bopDef $ + sintOp (`P.SQuot` P.Unsafe) + ++ uintOp (`P.UDiv` P.Unsafe) + def "%%" = + Just $ + bopDef $ + sintOp (`P.SRem` P.Unsafe) + ++ uintOp (`P.UMod` P.Unsafe) + def "^" = Just $ bopDef $ intOp P.Xor + def "&" = Just $ bopDef $ intOp P.And + def "|" = Just $ bopDef $ intOp P.Or + def ">>" = Just $ bopDef $ sintOp P.AShr ++ uintOp P.LShr + def "<<" = Just $ bopDef $ intOp P.Shl + def ">>>" = Just $ bopDef $ sintOp P.LShr ++ uintOp P.LShr + def "==" = Just $ + fun2 $ + \xs ys -> return $ ValuePrim $ BoolValue $ xs == ys + def "!=" = Just $ + fun2 $ + \xs ys -> return $ ValuePrim $ BoolValue $ xs /= ys + -- The short-circuiting is handled directly in 'eval'; these cases + -- are only used when partially applying and such. + def "&&" = Just $ + fun2 $ \x y -> + return $ ValuePrim $ BoolValue $ asBool x && asBool y + def "||" = Just $ + fun2 $ \x y -> + return $ ValuePrim $ BoolValue $ asBool x || asBool y + def "<" = + Just $ + bopDef $ + sintCmp P.CmpSlt ++ uintCmp P.CmpUlt + ++ floatCmp P.FCmpLt + ++ boolCmp P.CmpLlt + def ">" = + Just $ + bopDef $ + flipCmps $ + sintCmp P.CmpSlt ++ uintCmp P.CmpUlt + ++ floatCmp P.FCmpLt + ++ boolCmp P.CmpLlt + def "<=" = + Just $ + bopDef $ + sintCmp P.CmpSle ++ uintCmp P.CmpUle + ++ floatCmp P.FCmpLe + ++ boolCmp P.CmpLle + def ">=" = + Just $ + bopDef $ + flipCmps $ + sintCmp P.CmpSle ++ uintCmp P.CmpUle + ++ floatCmp P.FCmpLe + ++ boolCmp P.CmpLle + def s + | Just bop <- find ((s ==) . pretty) P.allBinOps = + Just $ tbopDef $ P.doBinOp bop + | Just unop <- find ((s ==) . pretty) P.allCmpOps = + Just $ tbopDef $ \x y -> P.BoolValue <$> P.doCmpOp unop x y + | Just cop <- find ((s ==) . pretty) P.allConvOps = + Just $ unopDef [(getV, Just . putV, P.doConvOp cop)] + | Just unop <- find ((s ==) . pretty) P.allUnOps = + Just $ unopDef [(getV, Just . putV, P.doUnOp unop)] + | Just (pts, _, f) <- M.lookup s P.primFuns = + case length pts of + 1 -> Just $ unopDef [(getV, Just . putV, f . pure)] + _ -> Just $ + fun1 $ \x -> do + let getV' (ValuePrim v) = Just v + getV' _ = Nothing + case mapM getV' =<< fromTuple x of + Just vs + | Just res <- fmap putV . f =<< mapM getV vs -> do + breakOnNaN vs res + return $ ValuePrim res + _ -> + error $ "Cannot apply " ++ pretty s ++ " to " ++ pretty x + | "sign_" `isPrefixOf` s = + Just $ + fun1 $ \x -> + case x of + (ValuePrim (UnsignedValue x')) -> + return $ ValuePrim $ SignedValue x' + _ -> error $ "Cannot sign: " ++ pretty x + | "unsign_" `isPrefixOf` s = + Just $ + fun1 $ \x -> + case x of + (ValuePrim (SignedValue x')) -> + return $ ValuePrim $ UnsignedValue x' + _ -> error $ "Cannot unsign: " ++ pretty x + def s + | "map_stream" `isPrefixOf` s = + Just $ fun2t stream + def s | "reduce_stream" `isPrefixOf` s = + Just $ fun3t $ \_ f arg -> stream f arg + def "map" = Just $ + TermPoly Nothing $ \t -> return $ + ValueFun $ \v -> + case (fromTuple v, unfoldFunType t) of + (Just [f, xs], ([_], ret_t)) + | Just rowshape <- typeRowShape ret_t -> + toArray' rowshape <$> mapM (apply noLoc mempty f) (snd $ fromArray xs) + | otherwise -> + error $ "Bad return type: " ++ pretty ret_t + _ -> + error $ + "Invalid arguments to map intrinsic:\n" + ++ unlines [pretty t, pretty v] + where + typeRowShape = sequenceA . structTypeShape mempty . stripArray 1 + def s | "reduce" `isPrefixOf` s = Just $ + fun3t $ \f ne xs -> + foldM (apply2 noLoc mempty f) ne $ snd $ fromArray xs + def "scan" = Just $ + fun3t $ \f ne xs -> do + let next (out, acc) x = do + x' <- apply2 noLoc mempty f acc x + return (x' : out, x') + toArray' (valueShape ne) . reverse . fst + <$> foldM next ([], ne) (snd $ fromArray xs) + def "scatter" = Just $ + fun3t $ \arr is vs -> + case arr of + ValueArray shape arr' -> + return $ + ValueArray shape $ + foldl' update arr' $ + zip (map asInt $ snd $ fromArray is) (snd $ fromArray vs) + _ -> + error $ "scatter expects array, but got: " ++ pretty arr + where + update arr' (i, v) = + if i >= 0 && i < arrayLength arr' + then arr' // [(i, v)] + else arr' + def "hist" = Just $ + fun6t $ \_ arr fun _ is vs -> + case arr of + ValueArray shape arr' -> + ValueArray shape + <$> foldM + (update fun) + arr' + (zip (map asInt $ snd $ fromArray is) (snd $ fromArray vs)) + _ -> + error $ "hist expects array, but got: " ++ pretty arr + where + update fun arr' (i, v) = + if i >= 0 && i < arrayLength arr' + then do + v' <- apply2 noLoc mempty fun (arr' ! i) v + return $ arr' // [(i, v')] + else return arr' + def "partition" = Just $ + fun3t $ \k f xs -> do + let (ShapeDim _ rowshape, xs') = fromArray xs + + next outs x = do + i <- asInt <$> apply noLoc mempty f x + return $ insertAt i x outs + pack parts = + toTuple + [ toArray' rowshape $ concat parts, + toArray' rowshape $ + map (ValuePrim . SignedValue . Int64Value . genericLength) parts + ] + + pack . map reverse + <$> foldM next (replicate (asInt k) []) xs' + where + insertAt 0 x (l : ls) = (x : l) : ls + insertAt i x (l : ls) = l : insertAt (i -1) x ls + insertAt _ _ ls = ls + def "unzip" = Just $ + fun1 $ \x -> do + let ShapeDim _ (ShapeRecord fs) = valueShape x + Just [xs_shape, ys_shape] = areTupleFields fs + listPair (xs, ys) = + [toArray' xs_shape xs, toArray' ys_shape ys] + + return $ toTuple $ listPair $ unzip $ map (fromPair . fromTuple) $ snd $ fromArray x + where + fromPair (Just [x, y]) = (x, y) + fromPair l = error $ "Not a pair: " ++ pretty l + def "zip" = Just $ + fun2t $ \xs ys -> do + let ShapeDim _ xs_rowshape = valueShape xs + ShapeDim _ ys_rowshape = valueShape ys + return $ + toArray' (ShapeRecord (tupleFields [xs_rowshape, ys_rowshape])) $ + map toTuple $ transpose [snd $ fromArray xs, snd $ fromArray ys] + def "concat" = Just $ + fun2t $ \xs ys -> do + let (ShapeDim _ rowshape, xs') = fromArray xs + (_, ys') = fromArray ys + return $ toArray' rowshape $ xs' ++ ys' + def "transpose" = Just $ + fun1 $ \xs -> do + let (ShapeDim n (ShapeDim m shape), xs') = fromArray xs + return $ + toArray (ShapeDim m (ShapeDim n shape)) $ + map (toArray (ShapeDim n shape)) $ transpose $ map (snd . fromArray) xs' + def "rotate" = Just $ + fun2t $ \i xs -> do + let (shape, xs') = fromArray xs + return $ + if asInt i > 0 + then + let (bef, aft) = splitAt (asInt i) xs' + in toArray shape $ aft ++ bef + else + let (bef, aft) = splitFromEnd (- asInt i) xs' + in toArray shape $ aft ++ bef + def "flatten" = Just $ + fun1 $ \xs -> do + let (ShapeDim n (ShapeDim m shape), xs') = fromArray xs + return $ toArray (ShapeDim (n * m) shape) $ concatMap (snd . fromArray) xs' + def "unflatten" = Just $ + fun3t $ \n m xs -> do + let (ShapeDim _ innershape, xs') = fromArray xs + rowshape = ShapeDim (asInt64 m) innershape + shape = ShapeDim (asInt64 n) rowshape + return $ toArray shape $ map (toArray rowshape) $ chunk (asInt m) xs' + def "opaque" = Just $ fun1 return + def "trace" = Just $ fun1 $ \v -> trace v >> return v + def "break" = Just $ + fun1 $ \v -> do + break + return v + def s | nameFromString s `M.member` namesToPrimTypes = Nothing + def s = error $ "Missing intrinsic: " ++ s + + tdef s = do + t <- nameFromString s `M.lookup` namesToPrimTypes + return $ T.TypeAbbr Unlifted [] $ Scalar $ Prim t + + stream f arg@(ValueArray _ xs) = + let n = ValuePrim $ SignedValue $ Int64Value $ arrayLength xs + in apply2 noLoc mempty f n arg + stream _ arg = error $ "Cannot stream: " ++ pretty arg + +interpretExp :: Ctx -> Exp -> F ExtOp Value +interpretExp ctx e = runEvalM (ctxImports ctx) $ eval (ctxEnv ctx) e + +interpretDec :: Ctx -> Dec -> F ExtOp Ctx +interpretDec ctx d = do + env <- runEvalM (ctxImports ctx) $ evalDec (ctxEnv ctx) d + return ctx {ctxEnv = env} + +interpretImport :: Ctx -> (FilePath, Prog) -> F ExtOp Ctx +interpretImport ctx (fp, prog) = do + env <- runEvalM (ctxImports ctx) $ foldM evalDec (ctxEnv ctx) $ progDecs prog + return ctx {ctxImports = M.insert fp env $ ctxImports ctx} + +checkEntryArgs :: VName -> [F.Value] -> StructType -> Either String () +checkEntryArgs entry args entry_t + | args_ts == param_ts = + return () + | otherwise = + Left $ + pretty $ + expected + "Got input of types" + indent 2 (stack (map ppr args_ts)) + where + (param_ts, _) = unfoldFunType entry_t + args_ts = map (valueStructType . valueType) args + expected + | null param_ts = + "Entry point " <> pquote (pprName entry) <> " is not a function." + | otherwise = + "Entry point " <> pquote (pprName entry) <> " expects input of type(s)" + indent 2 (stack (map ppr param_ts)) + +-- | Execute the named function on the given arguments; may fail +-- horribly if these are ill-typed. +interpretFunction :: Ctx -> VName -> [F.Value] -> Either String (F ExtOp Value) +interpretFunction ctx fname vs = do + ft <- case lookupVar (qualName fname) $ ctxEnv ctx of + Just (TermValue (Just (T.BoundV _ t)) _) -> + Right $ updateType (map valueType vs) t + Just (TermPoly (Just (T.BoundV _ t)) _) -> + Right $ updateType (map valueType vs) t + _ -> + Left $ "Unknown function `" <> prettyName fname <> "`." + + vs' <- case mapM convertValue vs of + Just vs' -> Right vs' + Nothing -> Left "Invalid input: irregular array." + + checkEntryArgs fname vs ft + + Right $ + runEvalM (ctxImports ctx) $ do + f <- evalTermVar (ctxEnv ctx) (qualName fname) ft + foldM (apply noLoc mempty) f vs' + where + updateType (vt : vts) (Scalar (Arrow als u _ rt)) = + Scalar $ Arrow als u (valueStructType vt) $ updateType vts rt + updateType _ t = t + + convertValue (F.PrimValue p) = Just $ ValuePrim p + convertValue (F.ArrayValue arr t) = mkArray t =<< mapM convertValue (elems arr) diff --git a/src/Language/Futhark/Parser.hs b/src/Language/Futhark/Parser.hs --- a/src/Language/Futhark/Parser.hs +++ b/src/Language/Futhark/Parser.hs @@ -1,64 +1,72 @@ {-# LANGUAGE Safe #-} + -- | Interface to the Futhark parser. module Language.Futhark.Parser - ( parseFuthark - , parseExp - , parseModExp - , parseType - - , parseValue - , parseValues - - , parseDecOrExpIncrM - - , ParseError (..) - - , scanTokensText - , L(..) - , Token(..) + ( parseFuthark, + parseExp, + parseModExp, + parseType, + parseValue, + parseValues, + parseDecOrExpIncrM, + ParseError (..), + scanTokensText, + L (..), + Token (..), ) - where +where import qualified Data.Text as T - -import Language.Futhark.Syntax -import Language.Futhark.Prop -import Language.Futhark.Parser.Parser import Language.Futhark.Parser.Lexer +import Language.Futhark.Parser.Parser +import Language.Futhark.Prop +import Language.Futhark.Syntax -- | Parse an entire Futhark program from the given 'T.Text', using -- the 'FilePath' as the source name for error messages. -parseFuthark :: FilePath -> T.Text - -> Either ParseError UncheckedProg +parseFuthark :: + FilePath -> + T.Text -> + Either ParseError UncheckedProg parseFuthark = parse prog -- | Parse an Futhark expression from the given 'String', using the -- 'FilePath' as the source name for error messages. -parseExp :: FilePath -> T.Text - -> Either ParseError UncheckedExp +parseExp :: + FilePath -> + T.Text -> + Either ParseError UncheckedExp parseExp = parse expression -- | Parse a Futhark module expression from the given 'String', using the -- 'FilePath' as the source name for error messages. -parseModExp :: FilePath -> T.Text - -> Either ParseError (ModExpBase NoInfo Name) +parseModExp :: + FilePath -> + T.Text -> + Either ParseError (ModExpBase NoInfo Name) parseModExp = parse modExpression -- | Parse an Futhark type from the given 'String', using the -- 'FilePath' as the source name for error messages. -parseType :: FilePath -> T.Text - -> Either ParseError UncheckedTypeExp +parseType :: + FilePath -> + T.Text -> + Either ParseError UncheckedTypeExp parseType = parse futharkType -- | Parse any Futhark value from the given 'String', using the 'FilePath' -- as the source name for error messages. -parseValue :: FilePath -> T.Text - -> Either ParseError Value +parseValue :: + FilePath -> + T.Text -> + Either ParseError Value parseValue = parse anyValue -- | Parse several Futhark values (separated by anything) from the given -- 'String', using the 'FilePath' as the source name for error -- messages. -parseValues :: FilePath -> T.Text - -> Either ParseError [Value] +parseValues :: + FilePath -> + T.Text -> + Either ParseError [Value] parseValues = parse anyValues diff --git a/src/Language/Futhark/Parser/Parser.y b/src/Language/Futhark/Parser/Parser.y --- a/src/Language/Futhark/Parser/Parser.y +++ b/src/Language/Futhark/Parser/Parser.y @@ -974,7 +974,7 @@ | '[' ']' {% emptyArrayError $1 } -Dim :: { Int32 } +Dim :: { Int64 } Dim : intlit { let L _ (INTLIT num) = $1 in fromInteger num } ValueType :: { ValueType } diff --git a/src/Language/Futhark/Prelude.hs b/src/Language/Futhark/Prelude.hs --- a/src/Language/Futhark/Prelude.hs +++ b/src/Language/Futhark/Prelude.hs @@ -1,7 +1,8 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE QuasiQuotes #-} -{-# LANGUAGE Trustworthy #-} {-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE Trustworthy #-} + -- | The Futhark Prelude Library embedded embedded as strings read -- during compilation of the Futhark compiler. The advantage is that -- the prelude can be accessed without reading it from disk, thus @@ -11,12 +12,12 @@ import Data.FileEmbed import qualified Data.Text as T import qualified Data.Text.Encoding as T -import qualified System.FilePath.Posix as Posix - import Futhark.Util (toPOSIX) +import qualified System.FilePath.Posix as Posix -- | Prelude embedded as 'T.Text' values, one for every file. prelude :: [(Posix.FilePath, T.Text)] prelude = map fixup prelude_bs - where prelude_bs = $(embedDir "prelude") - fixup (path, s) = ("/prelude" Posix. toPOSIX path, T.decodeUtf8 s) + where + prelude_bs = $(embedDir "prelude") + fixup (path, s) = ("/prelude" Posix. toPOSIX path, T.decodeUtf8 s) 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 @@ -1,38 +1,36 @@ -{-# OPTIONS_GHC -fno-warn-orphans #-} -{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + -- | Futhark prettyprinter. This module defines 'Pretty' instances -- for the AST defined in "Language.Futhark.Syntax". module Language.Futhark.Pretty - ( pretty - , prettyTuple - , leadingOperator - , IsName(..) - , prettyName - , Annot(..) + ( pretty, + prettyTuple, + leadingOperator, + IsName (..), + prettyName, + Annot (..), ) where -import Control.Monad -import Codec.Binary.UTF8.String (decode) -import Data.Array -import Data.Functor -import qualified Data.Map.Strict as M -import Data.List (intersperse) -import qualified Data.List.NonEmpty as NE -import Data.Maybe -import Data.Monoid hiding (Sum) -import Data.Ord -import Data.Word - -import Prelude - -import Futhark.Util.Pretty -import Futhark.Util - -import Language.Futhark.Syntax -import Language.Futhark.Prop +import Codec.Binary.UTF8.String (decode) +import Control.Monad +import Data.Array +import Data.Functor +import Data.List (intersperse) +import qualified Data.List.NonEmpty as NE +import qualified Data.Map.Strict as M +import Data.Maybe +import Data.Monoid hiding (Sum) +import Data.Ord +import Data.Word +import Futhark.Util +import Futhark.Util.Pretty +import Language.Futhark.Prop +import Language.Futhark.Syntax +import Prelude commastack :: [Doc] -> Doc commastack = align . stack . punctuate comma @@ -52,9 +50,10 @@ -- VNames are printed as either the name with an internal tag, or just -- the base name. instance IsName VName where - pprName | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False = - \(VName vn i) -> ppr vn <> text "_" <> text (show i) - | otherwise = ppr . baseName + pprName + | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False = + \(VName vn i) -> ppr vn <> text "_" <> text (show i) + | otherwise = ppr . baseName instance IsName Name where pprName = ppr @@ -83,30 +82,30 @@ ppr (PrimValue bv) = ppr bv ppr (ArrayValue a t) | [] <- elems a = text "empty" <> parens (ppr t) - | Array{} <- t = brackets $ commastack $ map ppr $ elems a - | otherwise = brackets $ commasep $ map ppr $ elems a + | Array {} <- t = brackets $ commastack $ map ppr $ elems a + | otherwise = brackets $ commasep $ map ppr $ elems a instance Pretty PrimValue where ppr (UnsignedValue (Int8Value v)) = - text (show (fromIntegral v::Word8)) <> text "u8" + text (show (fromIntegral v :: Word8)) <> text "u8" ppr (UnsignedValue (Int16Value v)) = - text (show (fromIntegral v::Word16)) <> text "u16" + text (show (fromIntegral v :: Word16)) <> text "u16" ppr (UnsignedValue (Int32Value v)) = - text (show (fromIntegral v::Word32)) <> text "u32" + text (show (fromIntegral v :: Word32)) <> text "u32" ppr (UnsignedValue (Int64Value v)) = - text (show (fromIntegral v::Word64)) <> text "u64" + text (show (fromIntegral v :: Word64)) <> text "u64" ppr (SignedValue v) = ppr v ppr (BoolValue True) = text "true" ppr (BoolValue False) = text "false" ppr (FloatValue v) = ppr v instance IsName vn => Pretty (DimDecl vn) where - ppr AnyDim = mempty + ppr AnyDim = mempty ppr (NamedDim v) = ppr v ppr (ConstDim n) = ppr n instance IsName vn => Pretty (DimExp vn) where - ppr DimExpAny = mempty + ppr DimExpAny = mempty ppr (DimExpNamed v _) = ppr v ppr (DimExpConst n _) = ppr n @@ -116,7 +115,7 @@ instance Pretty (ShapeDecl ()) where ppr (ShapeDecl ds) = mconcat $ replicate (length ds) $ text "[]" -instance Pretty (ShapeDecl Int32) where +instance Pretty (ShapeDecl Int64) where ppr (ShapeDecl ds) = mconcat (map (brackets . ppr) ds) instance Pretty (ShapeDecl Bool) where @@ -127,27 +126,29 @@ pprPrec _ (Prim et) = ppr et pprPrec p (TypeVar _ u et targs) = parensIf (not (null targs) && p > 3) $ - ppr u <> ppr (qualNameFromTypeName et) <+> spread (map (pprPrec 3) targs) + ppr u <> ppr (qualNameFromTypeName et) <+> spread (map (pprPrec 3) targs) pprPrec _ (Record fs) | Just ts <- areTupleFields fs = - oneLine (parens $ commasep $ map ppr ts) + oneLine (parens $ commasep $ map ppr ts) <|> parens (align $ mconcat $ punctuate (text "," <> line) $ map ppr ts) | otherwise = - oneLine (braces $ commasep fs') + oneLine (braces $ commasep fs') <|> braces (align $ mconcat $ punctuate (text "," <> line) fs') - where ppField (name, t) = text (nameToString name) <> colon <+> align (ppr t) - fs' = map ppField $ M.toList fs + where + ppField (name, t) = text (nameToString name) <> colon <+> align (ppr t) + fs' = map ppField $ M.toList fs pprPrec p (Arrow _ (Named v) t1 t2) = parensIf (p > 1) $ - parens (pprName v <> colon <+> align (ppr t1)) <+/> text "->" <+> pprPrec 1 t2 + parens (pprName v <> colon <+> align (ppr t1)) <+/> text "->" <+> pprPrec 1 t2 pprPrec p (Arrow _ Unnamed t1 t2) = parensIf (p > 1) $ pprPrec 2 t1 <+/> text "->" <+> pprPrec 1 t2 pprPrec p (Sum cs) = parensIf (p > 0) $ - oneLine (mconcat $ punctuate (text " | ") cs') - <|> align (mconcat $ punctuate (text " |" <> line) cs') - where ppConstr (name, fs) = sep $ (text "#" <> ppr name) : map (pprPrec 1) fs - cs' = map ppConstr $ M.toList cs + oneLine (mconcat $ punctuate (text " | ") cs') + <|> align (mconcat $ punctuate (text " |" <> line) cs') + where + ppConstr (name, fs) = sep $ (text "#" <> ppr name) : map (pprPrec 1) fs + cs' = map ppConstr $ M.toList cs instance Pretty (ShapeDecl dim) => Pretty (TypeBase dim as) where ppr = pprPrec 0 @@ -164,15 +165,18 @@ ppr (TEArray at d _) = brackets (ppr d) <> ppr at ppr (TETuple ts _) = parens $ commasep $ map ppr ts ppr (TERecord fs _) = braces $ commasep $ map ppField fs - where ppField (name, t) = text (nameToString name) <> colon <+> ppr t + where + ppField (name, t) = text (nameToString name) <> colon <+> ppr t ppr (TEVar name _) = ppr name ppr (TEApply t arg _) = ppr t <+> ppr arg ppr (TEArrow (Just v) t1 t2 _) = parens v' <+> text "->" <+> ppr t2 - where v' = pprName v <> colon <+> ppr t1 + where + v' = pprName v <> colon <+> ppr t1 ppr (TEArrow Nothing t1 t2 _) = ppr t1 <+> text "->" <+> ppr t2 ppr (TESum cs _) = align $ cat $ punctuate (text " |" <> softline) $ map ppConstr cs - where ppConstr (name, fs) = text "#" <> ppr name <+> sep (map ppr fs) + where + ppConstr (name, fs) = text "#" <> ppr name <+> sep (map ppr fs) instance (Eq vn, IsName vn) => Pretty (TypeArgExp vn) where ppr (TypeArgExpDim d _) = brackets $ ppr d @@ -189,36 +193,38 @@ ppr = pprName . identName hasArrayLit :: ExpBase ty vn -> Bool -hasArrayLit ArrayLit{} = True +hasArrayLit ArrayLit {} = True hasArrayLit (TupLit es2 _) = any hasArrayLit es2 -hasArrayLit _ = False +hasArrayLit _ = False instance (Eq vn, IsName vn, Annot f) => Pretty (DimIndexBase f vn) where - ppr (DimFix e) = ppr e + ppr (DimFix e) = ppr e ppr (DimSlice i j (Just s)) = - maybe mempty ppr i <> text ":" <> - maybe mempty ppr j <> text ":" <> - ppr s + maybe mempty ppr i <> text ":" + <> maybe mempty ppr j + <> text ":" + <> ppr s ppr (DimSlice i (Just j) s) = - maybe mempty ppr i <> text ":" <> - ppr j <> - maybe mempty ((text ":" <>) . ppr) s + maybe mempty ppr i <> text ":" + <> ppr j + <> maybe mempty ((text ":" <>) . ppr) s ppr (DimSlice i Nothing Nothing) = maybe mempty ppr i <> text ":" letBody :: (Eq vn, IsName vn, Annot f) => ExpBase f vn -> Doc -letBody body@LetPat{} = ppr body -letBody body@LetFun{} = ppr body -letBody body = text "in" <+> align (ppr body) +letBody body@LetPat {} = ppr body +letBody body@LetFun {} = ppr body +letBody body = text "in" <+> align (ppr body) instance (Eq vn, IsName vn, Annot f) => Pretty (ExpBase f vn) where ppr = pprPrec (-1) pprPrec _ (Var name t _) = ppr name <> inst - where inst = case unAnnot t of - Just t' - | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False -> - text "@" <> parens (align $ ppr t') - _ -> mempty + where + inst = case unAnnot t of + Just t' + | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False -> + text "@" <> parens (align $ ppr t') + _ -> mempty pprPrec _ (Parens e _) = align $ parens $ ppr e pprPrec _ (QualParens (v, _) e _) = ppr v <> text "." <> align (parens $ ppr e) pprPrec p (Ascript e t _) = @@ -230,84 +236,98 @@ pprPrec _ (FloatLit v _ _) = ppr v pprPrec _ (TupLit es _) | any hasArrayLit es = parens $ commastack $ map ppr es - | otherwise = parens $ commasep $ map ppr es + | otherwise = parens $ commasep $ map ppr es pprPrec _ (RecordLit fs _) | any fieldArray fs = braces $ commastack $ map ppr fs - | otherwise = braces $ commasep $ map ppr fs - where fieldArray (RecordFieldExplicit _ e _) = hasArrayLit e - fieldArray RecordFieldImplicit{} = False + | otherwise = braces $ commasep $ map ppr fs + where + fieldArray (RecordFieldExplicit _ e _) = hasArrayLit e + fieldArray RecordFieldImplicit {} = False pprPrec _ (ArrayLit es info _) = brackets (commasep $ map ppr es) <> info' - where info' = case unAnnot info of - Just t - | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False -> - text "@" <> parens (align $ ppr t) - _ -> mempty + where + info' = case unAnnot info of + Just t + | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False -> + text "@" <> parens (align $ ppr t) + _ -> mempty pprPrec _ (StringLit s _) = text $ show $ decode s pprPrec p (Range start maybe_step end _ _) = - parensIf (p /= -1) $ ppr start <> - maybe mempty ((text ".." <>) . ppr) maybe_step <> - case end of - DownToExclusive end' -> text "..>" <> ppr end' - ToInclusive end' -> text "..." <> ppr end' - UpToExclusive end' -> text "..<" <> ppr end' - pprPrec p (BinOp (bop,_) _ (x,_) (y,_) _ _ _) = prettyBinOp p bop x y + parensIf (p /= -1) $ + ppr start + <> maybe mempty ((text ".." <>) . ppr) maybe_step + <> case end of + DownToExclusive end' -> text "..>" <> ppr end' + ToInclusive end' -> text "..." <> ppr end' + UpToExclusive end' -> text "..<" <> ppr end' + pprPrec p (BinOp (bop, _) _ (x, _) (y, _) _ _ _) = prettyBinOp p bop x y pprPrec _ (Project k e _ _) = ppr e <> text "." <> ppr k - pprPrec _ (If c t f _ _) = text "if" <+> ppr c - text "then" <+> align (ppr t) - text "else" <+> align (ppr f) + pprPrec _ (If c t f _ _) = + text "if" <+> ppr c + text "then" <+> align (ppr t) + text "else" <+> align (ppr f) pprPrec p (Apply f arg _ _ _) = parensIf (p >= 10) $ pprPrec 0 f <+/> pprPrec 10 arg pprPrec _ (Negate e _) = text "-" <> ppr e pprPrec p (LetPat pat e body _ _) = - parensIf (p /= -1) $ align $ - text "let" <+> align (ppr pat) <+> - (if linebreak - then equals indent 2 (ppr e) - else equals <+> align (ppr e)) - letBody body - where linebreak = case e of - DoLoop{} -> True - LetPat{} -> True - LetWith{} -> True - If{} -> True - Match{} -> True - Attr{} -> True - ArrayLit{} -> False - _ -> hasArrayLit e + parensIf (p /= -1) $ + align $ + text "let" <+> align (ppr pat) + <+> ( if linebreak + then equals indent 2 (ppr e) + else equals <+> align (ppr e) + ) + letBody body + where + linebreak = case e of + DoLoop {} -> True + LetPat {} -> True + LetWith {} -> True + If {} -> True + Match {} -> True + Attr {} -> True + ArrayLit {} -> False + _ -> hasArrayLit e pprPrec _ (LetFun fname (tparams, params, retdecl, rettype, e) body _ _) = - text "let" <+> pprName fname <+> spread (map ppr tparams ++ map ppr params) <> - retdecl' <+> equals indent 2 (ppr e) - letBody body - where retdecl' = case (ppr <$> unAnnot rettype) `mplus` (ppr <$> retdecl) of - Just rettype' -> colon <+> align rettype' - Nothing -> mempty + text "let" <+> pprName fname <+> spread (map ppr tparams ++ map ppr params) + <> retdecl' <+> equals + indent 2 (ppr e) + letBody body + where + retdecl' = case (ppr <$> unAnnot rettype) `mplus` (ppr <$> retdecl) of + Just rettype' -> colon <+> align rettype' + Nothing -> mempty pprPrec _ (LetWith dest src idxs ve body _ _) | dest == src = - text "let" <+> ppr dest <> list (map ppr idxs) <+> - equals <+> align (ppr ve) - letBody body + text "let" <+> ppr dest <> list (map ppr idxs) + <+> equals + <+> align (ppr ve) + letBody body | otherwise = - text "let" <+> ppr dest <+> equals <+> ppr src <+> - text "with" <+> brackets (commasep (map ppr idxs)) <+> - text "=" <+> align (ppr ve) - letBody body + text "let" <+> ppr dest <+> equals <+> ppr src + <+> text "with" + <+> brackets (commasep (map ppr idxs)) + <+> text "=" + <+> align (ppr ve) + letBody body pprPrec _ (Update src idxs ve _) = - ppr src <+> text "with" <+> - brackets (commasep (map ppr idxs)) <+> - text "=" <+> align (ppr ve) + ppr src <+> text "with" + <+> brackets (commasep (map ppr idxs)) + <+> text "=" + <+> align (ppr ve) pprPrec _ (RecordUpdate src fs ve _ _) = - ppr src <+> text "with" <+> - mconcat (intersperse (text ".") (map ppr fs)) <+> - text "=" <+> align (ppr ve) + ppr src <+> text "with" + <+> mconcat (intersperse (text ".") (map ppr fs)) + <+> text "=" + <+> align (ppr ve) pprPrec _ (Index e idxs _ _) = pprPrec 9 e <> brackets (commasep (map ppr idxs)) pprPrec _ (Assert e1 e2 _ _) = text "assert" <+> pprPrec 10 e1 <+> pprPrec 10 e2 pprPrec p (Lambda params body rettype _ _) = parensIf (p /= -1) $ - text "\\" <> spread (map ppr params) <> ppAscription rettype <+> - text "->" indent 2 (ppr body) + text "\\" <> spread (map ppr params) <> ppAscription rettype + <+> text "->" indent 2 (ppr body) pprPrec _ (OpSection binop _ _) = parens $ ppr binop pprPrec _ (OpSectionLeft binop _ x _ _ _) = @@ -316,14 +336,19 @@ parens $ ppr binop <+> ppr x pprPrec _ (ProjectSection fields _ _) = parens $ mconcat $ map p fields - where p name = text "." <> ppr name + where + p name = text "." <> ppr name pprPrec _ (IndexSection idxs _ _) = parens $ text "." <> brackets (commasep (map ppr idxs)) pprPrec _ (DoLoop sizeparams pat initexp form loopbody _ _) = - text "loop" <+> - align (spread (map (brackets . pprName) sizeparams) <+/> - ppr pat <+> equals <+/> ppr initexp <+/> ppr form <+> text "do") - indent 2 (ppr loopbody) + text "loop" + <+> align + ( spread (map (brackets . pprName) sizeparams) + <+/> ppr pat <+> equals + <+/> ppr initexp + <+/> ppr form <+> text "do" + ) + indent 2 (ppr loopbody) pprPrec _ (Constr n cs _ _) = text "#" <> ppr n <+> sep (map ppr cs) pprPrec _ (Match e cs _ _) = text "match" <+> ppr e (stack . map ppr) (NE.toList cs) pprPrec _ (Attr attr e _) = @@ -350,33 +375,34 @@ instance (Eq vn, IsName vn, Annot f) => Pretty (PatternBase f vn) where ppr (PatternAscription p t _) = ppr p <> colon <+> align (ppr t) - ppr (PatternParens p _) = parens $ ppr p - ppr (Id v t _) = case unAnnot t of - Just t' -> parens $ pprName v <> colon <+> align (ppr t') - Nothing -> pprName v - ppr (TuplePattern pats _) = parens $ commasep $ map ppr pats - ppr (RecordPattern fs _) = braces $ commasep $ map ppField fs - where ppField (name, t) = text (nameToString name) <> equals <> ppr t - ppr (Wildcard t _) = case unAnnot t of - Just t' -> parens $ text "_" <> colon <+> ppr t' - Nothing -> text "_" - ppr (PatternLit e _ _) = ppr e - ppr (PatternConstr n _ ps _) = text "#" <> ppr n <+> sep (map ppr ps) + ppr (PatternParens p _) = parens $ ppr p + ppr (Id v t _) = case unAnnot t of + Just t' -> parens $ pprName v <> colon <+> align (ppr t') + Nothing -> pprName v + ppr (TuplePattern pats _) = parens $ commasep $ map ppr pats + ppr (RecordPattern fs _) = braces $ commasep $ map ppField fs + where + ppField (name, t) = text (nameToString name) <> equals <> ppr t + ppr (Wildcard t _) = case unAnnot t of + Just t' -> parens $ text "_" <> colon <+> ppr t' + Nothing -> text "_" + ppr (PatternLit e _ _) = ppr e + ppr (PatternConstr n _ ps _) = text "#" <> ppr n <+> sep (map ppr ps) ppAscription :: Pretty t => Maybe t -> Doc -ppAscription Nothing = mempty +ppAscription Nothing = mempty ppAscription (Just t) = colon <> align (ppr t) instance (Eq vn, IsName vn, Annot f) => Pretty (ProgBase f vn) where ppr = stack . punctuate line . map ppr . progDecs instance (Eq vn, IsName vn, Annot f) => Pretty (DecBase f vn) where - ppr (ValDec dec) = ppr dec - ppr (TypeDec dec) = ppr dec - ppr (SigDec sig) = ppr sig - ppr (ModDec sd) = ppr sd - ppr (OpenDec x _) = text "open" <+> ppr x - ppr (LocalDec dec _) = text "local" <+> ppr dec + ppr (ValDec dec) = ppr dec + ppr (TypeDec dec) = ppr dec + ppr (SigDec sig) = ppr sig + ppr (ModDec sd) = ppr sd + ppr (OpenDec x _) = text "open" <+> ppr x + ppr (LocalDec dec _) = text "local" <+> ppr dec ppr (ImportDec x _ _) = text "import" <+> ppr x instance (Eq vn, IsName vn, Annot f) => Pretty (ModExpBase f vn) where @@ -387,10 +413,12 @@ ppr (ModApply f a _ _ _) = parens $ ppr f <+> parens (ppr a) ppr (ModAscript me se _ _) = ppr me <> colon <+> ppr se ppr (ModLambda param maybe_sig body _) = - text "\\" <> ppr param <> maybe_sig' <+> - text "->" indent 2 (ppr body) - where maybe_sig' = case maybe_sig of Nothing -> mempty - Just (sig, _) -> colon <+> ppr sig + text "\\" <> ppr param <> maybe_sig' + <+> text "->" indent 2 (ppr body) + where + maybe_sig' = case maybe_sig of + Nothing -> mempty + Just (sig, _) -> colon <+> ppr sig instance Pretty Liftedness where ppr Unlifted = text "" @@ -399,8 +427,10 @@ instance (Eq vn, IsName vn, Annot f) => Pretty (TypeBindBase f vn) where ppr (TypeBind name l params usertype _ _) = - text "type" <> ppr l <+> pprName name <+> - spread (map ppr params) <+> equals <+> ppr usertype + text "type" <> ppr l <+> pprName name + <+> spread (map ppr params) + <+> equals + <+> ppr usertype instance (Eq vn, IsName vn) => Pretty (TypeParamBase vn) where ppr (TypeParamDim name _) = brackets $ pprName name @@ -408,15 +438,20 @@ instance (Eq vn, IsName vn, Annot f) => Pretty (ValBindBase f vn) where ppr (ValBind entry name retdecl rettype tparams args body _ attrs _) = - mconcat (map ((<> line) . ppr) attrs) <> - text fun <+> pprName name <+> - align (sep (map ppr tparams ++ map ppr args)) <> retdecl' <> text " =" - indent 2 (ppr body) - where fun | isJust entry = "entry" - | otherwise = "let" - retdecl' = case (ppr . fst <$> unAnnot rettype) `mplus` (ppr <$> retdecl) of - Just rettype' -> colon <+> align rettype' - Nothing -> mempty + mconcat (map ((<> line) . ppr) attrs) + <> text fun + <+> pprName name + <+> align (sep (map ppr tparams ++ map ppr args)) + <> retdecl' + <> text " =" + indent 2 (ppr body) + where + fun + | isJust entry = "entry" + | otherwise = "let" + retdecl' = case (ppr . fst <$> unAnnot rettype) `mplus` (ppr <$> retdecl) of + Just rettype' -> colon <+> align rettype' + Nothing -> mempty instance (Eq vn, IsName vn, Annot f) => Pretty (SpecBase f vn) where ppr (TypeAbbrSpec tpsig) = ppr tpsig @@ -451,44 +486,54 @@ instance (Eq vn, IsName vn, Annot f) => Pretty (ModBindBase f vn) where ppr (ModBind name ps sig e _ _) = text "module" <+> pprName name <+> spread (map ppr ps) <+> sig' <> text " =" <+> ppr e - where sig' = case sig of Nothing -> mempty - Just (s,_) -> colon <+> ppr s <> text " " + where + sig' = case sig of + Nothing -> mempty + Just (s, _) -> colon <+> ppr s <> text " " -prettyBinOp :: (Eq vn, IsName vn, Annot f) => - Int -> QualName vn -> ExpBase f vn -> ExpBase f vn -> Doc -prettyBinOp p bop x y = parensIf (p > symPrecedence) $ - pprPrec symPrecedence x <+/> - bop' <+> - pprPrec symRPrecedence y - where bop' = case leading of Backtick -> text "`" <> ppr bop <> text "`" - _ -> ppr bop - leading = leadingOperator $ nameFromString $ pretty $ pprName $ qualLeaf bop - symPrecedence = precedence leading - symRPrecedence = rprecedence leading - precedence PipeRight = -1 - precedence PipeLeft = -1 - precedence LogAnd = 0 - precedence LogOr = 0 - precedence Band = 1 - precedence Bor = 1 - precedence Xor = 1 - precedence Equal = 2 - precedence NotEqual = 2 - precedence Less = 2 - precedence Leq = 2 - precedence Greater = 2 - precedence Geq = 2 - precedence ShiftL = 3 - precedence ShiftR = 3 - precedence Plus = 4 - precedence Minus = 4 - precedence Times = 5 - precedence Divide = 5 - precedence Mod = 5 - precedence Quot = 5 - precedence Rem = 5 - precedence Pow = 6 - precedence Backtick = 9 - rprecedence Minus = 10 - rprecedence Divide = 10 - rprecedence op = precedence op +prettyBinOp :: + (Eq vn, IsName vn, Annot f) => + Int -> + QualName vn -> + ExpBase f vn -> + ExpBase f vn -> + Doc +prettyBinOp p bop x y = + parensIf (p > symPrecedence) $ + pprPrec symPrecedence x + <+/> bop' + <+> pprPrec symRPrecedence y + where + bop' = case leading of + Backtick -> text "`" <> ppr bop <> text "`" + _ -> ppr bop + leading = leadingOperator $ nameFromString $ pretty $ pprName $ qualLeaf bop + symPrecedence = precedence leading + symRPrecedence = rprecedence leading + precedence PipeRight = -1 + precedence PipeLeft = -1 + precedence LogAnd = 0 + precedence LogOr = 0 + precedence Band = 1 + precedence Bor = 1 + precedence Xor = 1 + precedence Equal = 2 + precedence NotEqual = 2 + precedence Less = 2 + precedence Leq = 2 + precedence Greater = 2 + precedence Geq = 2 + precedence ShiftL = 3 + precedence ShiftR = 3 + precedence Plus = 4 + precedence Minus = 4 + precedence Times = 5 + precedence Divide = 5 + precedence Mod = 5 + precedence Quot = 5 + precedence Rem = 5 + precedence Pow = 6 + precedence Backtick = 9 + rprecedence Minus = 10 + rprecedence Divide = 10 + rprecedence op = precedence op 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 @@ -1,1033 +1,1128 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE ScopedTypeVariables #-} --- | This module provides various simple ways to query and manipulate --- fundamental Futhark terms, such as types and values. The intent is to --- keep "Futhark.Language.Syntax" simple, and put whatever embellishments --- we need here. -module Language.Futhark.Prop - ( - -- * Various - Intrinsic(..) - , intrinsics - , maxIntrinsicTag - , namesToPrimTypes - , qualName - , qualify - , typeName - , valueType - , primValueType - , leadingOperator - , progImports - , decImports - , progModuleTypes - , identifierReference - , prettyStacktrace - - -- * Queries on expressions - , typeOf - - -- * Queries on patterns and params - , patternIdents - , patternNames - , patternMap - , patternType - , patternStructType - , patternParam - , patternOrderZero - , patternDimNames - - -- * Queries on types - , uniqueness - , unique - , aliases - , diet - , arrayRank - , arrayShape - , nestedDims - , orderZero - , unfoldFunType - , foldFunType - , typeVars - , typeDimNames - , primByteSize - - -- * Operations on types - , rank - , peelArray - , stripArray - , arrayOf - , toStructural - , toStruct - , fromStruct - , setAliases - , addAliases - , setUniqueness - , noSizes - , addSizes - , anySizes - , traverseDims - , DimPos(..) - , mustBeExplicit - , mustBeExplicitInType - , tupleRecord - , isTupleRecord - , areTupleFields - , tupleFields - , tupleFieldNames - , sortFields - , sortConstrs - , isTypeParam - , isSizeParam - , combineTypeShapes - , matchDims - , unscopeType - , onRecordField - - -- | Values of these types are produces by the parser. They use - -- unadorned names and have no type information, apart from that - -- which is syntactically required. - , NoInfo(..) - , UncheckedType - , UncheckedTypeExp - , UncheckedIdent - , UncheckedTypeDecl - , UncheckedDimIndex - , UncheckedExp - , UncheckedModExp - , UncheckedSigExp - , UncheckedTypeParam - , UncheckedPattern - , UncheckedValBind - , UncheckedDec - , UncheckedSpec - , UncheckedProg - , UncheckedCase - ) - where - -import Control.Monad.State -import Data.Char -import Data.Foldable -import qualified Data.Map.Strict as M -import qualified Data.Set as S -import Data.List (sortOn, genericLength, isPrefixOf, nub) -import Data.Maybe -import Data.Ord -import Data.Bifunctor -import Data.Bifoldable -import Data.Bitraversable (bitraverse) - -import Futhark.Util (maxinum) -import Futhark.Util.Pretty - -import Language.Futhark.Syntax -import qualified Futhark.IR.Primitive as Primitive - --- | Return the dimensionality of a type. For non-arrays, this is --- zero. For a one-dimensional array it is one, for a two-dimensional --- it is two, and so forth. -arrayRank :: TypeBase dim as -> Int -arrayRank = shapeRank . arrayShape - --- | Return the shape of a type - for non-arrays, this is 'mempty'. -arrayShape :: TypeBase dim as -> ShapeDecl dim -arrayShape (Array _ _ _ ds) = ds -arrayShape _ = mempty - --- | Return any shape declarations in the type, with duplicates --- removed. -nestedDims :: TypeBase (DimDecl VName) as -> [DimDecl VName] -nestedDims t = - case t of Array _ _ a ds -> - nub $ nestedDims (Scalar a) <> shapeDims ds - Scalar (Record fs) -> - nub $ foldMap nestedDims fs - Scalar Prim{} -> - mempty - Scalar (Sum cs) -> - nub $ foldMap (foldMap nestedDims) cs - Scalar (Arrow _ v t1 t2) -> - filter (notV v) $ nestedDims t1 <> nestedDims t2 - Scalar (TypeVar _ _ _ targs) -> - concatMap typeArgDims targs - - where typeArgDims (TypeArgDim d _) = [d] - typeArgDims (TypeArgType at _) = nestedDims at - - notV Unnamed = const True - notV (Named v) = (/=NamedDim (qualName v)) - --- | Change the shape of a type to be just the rank. -noSizes :: TypeBase (DimDecl vn) as -> TypeBase () as -noSizes = first $ const () - --- | Add size annotations that are all 'AnyDim'. -addSizes :: TypeBase () as -> TypeBase (DimDecl vn) as -addSizes = first $ const AnyDim - --- | Change all size annotations to be 'AnyDim'. -anySizes :: TypeBase (DimDecl vn) as -> TypeBase (DimDecl vn) as -anySizes = first $ const AnyDim - --- | Where does this dimension occur? -data DimPos - = PosImmediate - -- ^ Immediately in the argument to 'traverseDims'. - | PosParam - -- ^ In a function parameter type. - | PosReturn - -- ^ In a function return type. - deriving (Eq, Ord, Show) - --- | Perform a traversal (possibly including replacement) on sizes --- that are parameters in a function type, but also including the type --- immediately passed to the function. Also passes along a set of the --- parameter names inside the type that have come in scope at the --- occurrence of the dimension. -traverseDims :: forall f fdim tdim als. - Applicative f => - (S.Set VName -> DimPos -> fdim -> f tdim) - -> TypeBase fdim als - -> f (TypeBase tdim als) -traverseDims f = go mempty PosImmediate - where go :: forall als'. - S.Set VName -> DimPos -> TypeBase fdim als' - -> f (TypeBase tdim als') - go bound b t@Array{} = - bitraverse (f bound b) pure t - go bound b (Scalar (Record fields)) = - Scalar . Record <$> traverse (go bound b) fields - go bound b (Scalar (TypeVar as u tn targs)) = - Scalar <$> (TypeVar as u tn <$> traverse (onTypeArg bound b) targs) - go bound b (Scalar (Sum cs)) = - Scalar . Sum <$> traverse (traverse (go bound b)) cs - go _ _ (Scalar (Prim t)) = - pure $ Scalar $ Prim t - go bound _ (Scalar (Arrow als p t1 t2)) = - Scalar <$> (Arrow als p <$> go bound' PosParam t1 <*> go bound' PosReturn t2) - where bound' = case p of Named p' -> S.insert p' bound - Unnamed -> bound - - onTypeArg bound b (TypeArgDim d loc) = - TypeArgDim <$> f bound b d <*> pure loc - onTypeArg bound b (TypeArgType t loc) = - TypeArgType <$> go bound b t <*> pure loc - -mustBeExplicitAux :: StructType -> M.Map VName Bool -mustBeExplicitAux t = - execState (traverseDims onDim t) mempty - where onDim bound _ (NamedDim d) - | qualLeaf d `S.member` bound = - modify $ \s -> M.insertWith (&&) (qualLeaf d) False s - onDim _ PosImmediate (NamedDim d) = - modify $ \s -> M.insertWith (&&) (qualLeaf d) False s - onDim _ _ (NamedDim d) = - modify $ M.insertWith (&&) (qualLeaf d) True - onDim _ _ _ = - return () - --- | Figure out which of the sizes in a parameter type must be passed --- explicitly, because their first use is as something else than just --- an array dimension. 'mustBeExplicit' is like this function, but --- first decomposes into parameter types. -mustBeExplicitInType :: StructType -> S.Set VName -mustBeExplicitInType t = - S.fromList $ M.keys $ M.filter id $ mustBeExplicitAux t - --- | Figure out which of the sizes in a binding type must be passed --- explicitly, because their first use is as something else than just --- an array dimension. -mustBeExplicit :: StructType -> S.Set VName -mustBeExplicit bind_t = - let (ts, ret) = unfoldFunType bind_t - alsoRet = M.unionWith (&&) $ - M.fromList $ zip (S.toList $ typeDimNames ret) $ repeat True - in S.fromList $ M.keys $ M.filter id $ alsoRet $ foldl' onType mempty ts - where onType uses t = uses <> mustBeExplicitAux t -- Left-biased union. - --- | Return the uniqueness of a type. -uniqueness :: TypeBase shape as -> Uniqueness -uniqueness (Array _ u _ _) = u -uniqueness (Scalar (TypeVar _ u _ _)) = u -uniqueness (Scalar (Sum ts)) = foldMap (foldMap uniqueness) $ M.elems ts -uniqueness (Scalar (Record fs)) = foldMap uniqueness $ M.elems fs -uniqueness _ = Nonunique - --- | @unique t@ is 'True' if the type of the argument is unique. -unique :: TypeBase shape as -> Bool -unique = (==Unique) . uniqueness - --- | Return the set of all variables mentioned in the aliasing of a --- type. -aliases :: Monoid as => TypeBase shape as -> as -aliases = bifoldMap (const mempty) id - --- | @diet t@ returns a description of how a function parameter of --- type @t@ might consume its argument. -diet :: TypeBase shape as -> Diet -diet (Scalar (Record ets)) = RecordDiet $ fmap diet ets -diet (Scalar (Prim _)) = Observe -diet (Scalar (Arrow _ _ t1 t2)) = FuncDiet (diet t1) (diet t2) -diet (Array _ Unique _ _) = Consume -diet (Array _ Nonunique _ _) = Observe -diet (Scalar (TypeVar _ Unique _ _)) = Consume -diet (Scalar (TypeVar _ Nonunique _ _)) = Observe -diet (Scalar Sum{}) = Observe - --- | Convert any type to one that has rank information, no alias --- information, and no embedded names. -toStructural :: TypeBase dim as - -> TypeBase () () -toStructural = flip setAliases () . first (const ()) - --- | Remove aliasing information from a type. -toStruct :: TypeBase dim as - -> TypeBase dim () -toStruct t = t `setAliases` () - --- | Replace no aliasing with an empty alias set. -fromStruct :: TypeBase dim as - -> TypeBase dim Aliasing -fromStruct t = t `setAliases` S.empty - --- | @peelArray n t@ returns the type resulting from peeling the first --- @n@ array dimensions from @t@. Returns @Nothing@ if @t@ has less --- than @n@ dimensions. -peelArray :: Int -> TypeBase dim as -> Maybe (TypeBase dim as) -peelArray n (Array als u t shape) - | shapeRank shape == n = - Just $ Scalar t `addAliases` const als - | otherwise = - Array als u t <$> stripDims n shape -peelArray _ _ = Nothing - --- | @arrayOf t s u@ constructs an array type. The convenience --- compared to using the 'Array' constructor directly is that @t@ can --- itself be an array. If @t@ is an @n@-dimensional array, and @s@ is --- a list of length @n@, the resulting type is of an @n+m@ dimensions. --- The uniqueness of the new array will be @u@, no matter the --- uniqueness of @t@. -arrayOf :: Monoid as => - TypeBase dim as - -> ShapeDecl dim - -> Uniqueness - -> TypeBase dim as -arrayOf t = arrayOfWithAliases (t `setUniqueness` Nonunique) mempty - -arrayOfWithAliases :: Monoid as => - TypeBase dim as - -> as - -> ShapeDecl dim - -> Uniqueness - -> TypeBase dim as -arrayOfWithAliases (Array as1 _ et shape1) as2 shape2 u = - Array (as1<>as2) u et (shape2 <> shape1) -arrayOfWithAliases (Scalar t) as shape u = - Array as u (second (const ()) t) shape - --- | @stripArray n t@ removes the @n@ outermost layers of the array. --- Essentially, it is the type of indexing an array of type @t@ with --- @n@ indexes. -stripArray :: Int -> TypeBase dim as -> TypeBase dim as -stripArray n (Array als u et shape) - | Just shape' <- stripDims n shape = - Array als u et shape' - | otherwise = - Scalar et `setUniqueness` u `setAliases` als -stripArray _ t = t - --- | Create a record type corresponding to a tuple with the given --- element types. -tupleRecord :: [TypeBase dim as] -> TypeBase dim as -tupleRecord = Scalar . Record . M.fromList . zip tupleFieldNames - --- | Does this type corespond to a tuple? If so, return the elements --- of that tuple. -isTupleRecord :: TypeBase dim as -> Maybe [TypeBase dim as] -isTupleRecord (Scalar (Record fs)) = areTupleFields fs -isTupleRecord _ = Nothing - --- | Does this record map correspond to a tuple? -areTupleFields :: M.Map Name a -> Maybe [a] -areTupleFields fs = - let fs' = sortFields fs - in if and $ zipWith (==) (map fst fs') tupleFieldNames - then Just $ map snd fs' - else Nothing - --- | Construct a record map corresponding to a tuple. -tupleFields :: [a] -> M.Map Name a -tupleFields as = M.fromList $ zip tupleFieldNames as - --- | Increasing field names for a tuple (starts at 0). -tupleFieldNames :: [Name] -tupleFieldNames = map (nameFromString . show) [(0::Int)..] - --- | Sort fields by their name; taking care to sort numeric fields by --- their numeric value. This ensures that tuples and tuple-like --- records match. -sortFields :: M.Map Name a -> [(Name,a)] -sortFields l = map snd $ sortOn fst $ zip (map (fieldish . fst) l') l' - where l' = M.toList l - fieldish s = case reads $ nameToString s of - [(x, "")] -> Left (x::Int) - _ -> Right s - --- | Sort the constructors of a sum type in some well-defined (but not --- otherwise significant) manner. -sortConstrs :: M.Map Name a -> [(Name, a)] -sortConstrs cs = sortOn fst $ M.toList cs - --- | Is this a 'TypeParamType'? -isTypeParam :: TypeParamBase vn -> Bool -isTypeParam TypeParamType{} = True -isTypeParam TypeParamDim{} = False - --- | Is this a 'TypeParamDim'? -isSizeParam :: TypeParamBase vn -> Bool -isSizeParam = not . isTypeParam - --- | Combine the shape information of types as much as possible. The first --- argument is the orignal type and the second is the type of the transformed --- expression. This is necessary since the original type may contain additional --- information (e.g., shape restrictions) from the user given annotation. -combineTypeShapes :: (Monoid as, ArrayDim dim) => - TypeBase dim as -> TypeBase dim as -> TypeBase dim as -combineTypeShapes (Scalar (Record ts1)) (Scalar (Record ts2)) - | M.keys ts1 == M.keys ts2 = - Scalar $ Record $ M.map (uncurry combineTypeShapes) (M.intersectionWith (,) ts1 ts2) -combineTypeShapes (Scalar (Arrow als1 p1 a1 b1)) (Scalar (Arrow als2 _p2 a2 b2)) = - Scalar $ Arrow (als1<>als2) p1 (combineTypeShapes a1 a2) (combineTypeShapes b1 b2) -combineTypeShapes (Scalar (TypeVar als1 u1 v targs1)) (Scalar (TypeVar als2 _ _ targs2)) = - Scalar $ TypeVar (als1<>als2) u1 v $ zipWith f targs1 targs2 - where f (TypeArgType t1 loc) (TypeArgType t2 _) = - TypeArgType (combineTypeShapes t1 t2) loc - f targ _ = targ -combineTypeShapes (Array als1 u1 et1 shape1) (Array als2 _u2 et2 _shape2) = - arrayOfWithAliases (combineTypeShapes (Scalar et1) (Scalar et2) - `setAliases` mempty) - (als1<>als2) shape1 u1 -combineTypeShapes _ new_tp = new_tp - --- | Match the dimensions of otherwise assumed-equal types. -matchDims :: (Monoid as, Monad m) => - (d1 -> d2 -> m d1) - -> TypeBase d1 as -> TypeBase d2 as - -> m (TypeBase d1 as) -matchDims onDims t1 t2 = - case (t1, t2) of - (Array als1 u1 et1 shape1, Array als2 u2 et2 shape2) -> - flip setAliases (als1<>als2) <$> - (arrayOf <$> - matchDims onDims (Scalar et1) (Scalar et2) <*> - onShapes shape1 shape2 <*> pure (min u1 u2)) - (Scalar (Record f1), Scalar (Record f2)) -> - Scalar . Record <$> - traverse (uncurry (matchDims onDims)) (M.intersectionWith (,) f1 f2) - (Scalar (Sum cs1), Scalar (Sum cs2)) -> - Scalar . Sum <$> - traverse (traverse (uncurry (matchDims onDims))) - (M.intersectionWith zip cs1 cs2) - (Scalar (Arrow als1 p1 a1 b1), Scalar (Arrow als2 _p2 a2 b2)) -> - Scalar <$> - (Arrow (als1 <> als2) p1 <$> matchDims onDims a1 a2 <*> matchDims onDims b1 b2) - (Scalar (TypeVar als1 u v targs1), - Scalar (TypeVar als2 _ _ targs2)) -> - Scalar . TypeVar (als1 <> als2) u v <$> zipWithM matchTypeArg targs1 targs2 - _ -> return t1 - - where matchTypeArg ta@TypeArgType{} _ = return ta - matchTypeArg a _ = return a - - onShapes shape1 shape2 = - ShapeDecl <$> zipWithM onDims (shapeDims shape1) (shapeDims shape2) - - --- | Set the uniqueness attribute of a type. If the type is a record --- or sum type, the uniqueness of its components will be modified. -setUniqueness :: TypeBase dim as -> Uniqueness -> TypeBase dim as -setUniqueness (Array als _ et shape) u = - Array als u et shape -setUniqueness (Scalar (TypeVar als _ t targs)) u = - Scalar $ TypeVar als u t targs -setUniqueness (Scalar (Record ets)) u = - Scalar $ Record $ fmap (`setUniqueness` u) ets -setUniqueness (Scalar (Sum ets)) u = - Scalar $ Sum $ fmap (map (`setUniqueness` u)) ets -setUniqueness t _ = t - --- | @t \`setAliases\` als@ returns @t@, but with @als@ substituted for --- any already present aliasing. -setAliases :: TypeBase dim asf -> ast -> TypeBase dim ast -setAliases t = addAliases t . const - --- | @t \`addAliases\` f@ returns @t@, but with any already present --- aliasing replaced by @f@ applied to that aliasing. -addAliases :: TypeBase dim asf -> (asf -> ast) - -> TypeBase dim ast -addAliases = flip second - -intValueType :: IntValue -> IntType -intValueType Int8Value{} = Int8 -intValueType Int16Value{} = Int16 -intValueType Int32Value{} = Int32 -intValueType Int64Value{} = Int64 - -floatValueType :: FloatValue -> FloatType -floatValueType Float32Value{} = Float32 -floatValueType Float64Value{} = Float64 - --- | The type of a basic value. -primValueType :: PrimValue -> PrimType -primValueType (SignedValue v) = Signed $ intValueType v -primValueType (UnsignedValue v) = Unsigned $ intValueType v -primValueType (FloatValue v) = FloatType $ floatValueType v -primValueType BoolValue{} = Bool - --- | The type of the value. -valueType :: Value -> ValueType -valueType (PrimValue bv) = Scalar $ Prim $ primValueType bv -valueType (ArrayValue _ t) = t - --- | The size of values of this type, in bytes. -primByteSize :: Num a => PrimType -> a -primByteSize (Signed it) = Primitive.intByteSize it -primByteSize (Unsigned it) = Primitive.intByteSize it -primByteSize (FloatType ft) = Primitive.floatByteSize ft -primByteSize Bool = 1 - --- | Construct a 'ShapeDecl' with the given number of 'AnyDim' --- dimensions. -rank :: Int -> ShapeDecl (DimDecl VName) -rank n = ShapeDecl $ replicate n AnyDim - --- | The type is leaving a scope, so clean up any aliases that --- reference the bound variables, and turn any dimensions that name --- them into AnyDim instead. -unscopeType :: S.Set VName -> PatternType -> PatternType -unscopeType bound_here t = first onDim $ t `addAliases` S.map unbind - where unbind (AliasBound v) | v `S.member` bound_here = AliasFree v - unbind a = a - onDim (NamedDim qn) | qualLeaf qn `S.member` bound_here = AnyDim - onDim d = d - --- | Perform some operation on a given record field. Returns --- 'Nothing' if that field does not exist. -onRecordField :: (TypeBase dim als -> TypeBase dim als) - -> [Name] - -> TypeBase dim als -> Maybe (TypeBase dim als) -onRecordField f [] t = Just $ f t -onRecordField f (k:ks) (Scalar (Record m)) = do - t <- onRecordField f ks =<< M.lookup k m - Just $ Scalar $ Record $ M.insert k t m -onRecordField _ _ _ = Nothing - --- | The type of an Futhark term. The aliasing will refer to itself, if --- the term is a non-tuple-typed variable. -typeOf :: ExpBase Info VName -> PatternType -typeOf (Literal val _) = Scalar $ Prim $ primValueType val -typeOf (IntLit _ (Info t) _) = t -typeOf (FloatLit _ (Info t) _) = t -typeOf (Parens e _) = typeOf e -typeOf (QualParens _ e _) = typeOf e -typeOf (TupLit es _) = tupleRecord $ map typeOf es -typeOf (RecordLit fs _) = - -- Reverse, because M.unions is biased to the left. - Scalar $ Record $ M.unions $ reverse $ map record fs - where record (RecordFieldExplicit name e _) = M.singleton name $ typeOf e - record (RecordFieldImplicit name (Info t) _) = - M.singleton (baseName name) $ t - `addAliases` S.insert (AliasBound name) -typeOf (ArrayLit _ (Info t) _) = t -typeOf (StringLit vs _) = - Array mempty Unique (Prim (Unsigned Int8)) - (ShapeDecl [ConstDim $ genericLength vs]) -typeOf (Range _ _ _ (Info t, _) _) = t -typeOf (BinOp _ _ _ _ (Info t) _ _) = t -typeOf (Project _ _ (Info t) _) = t -typeOf (If _ _ _ (Info t, _) _) = t -typeOf (Var _ (Info t) _) = t -typeOf (Ascript e _ _) = typeOf e -typeOf (Coerce _ _ (Info t, _) _) = t -typeOf (Apply _ _ _ (Info t, _) _) = t -typeOf (Negate e _) = typeOf e -typeOf (LetPat _ _ _ (Info t, _) _) = t -typeOf (LetFun _ _ _ (Info t) _) = t -typeOf (LetWith _ _ _ _ _ (Info t) _) = t -typeOf (Index _ _ (Info t, _) _) = t -typeOf (Update e _ _ _) = typeOf e `setAliases` mempty -typeOf (RecordUpdate _ _ _ (Info t) _) = t -typeOf (Assert _ e _ _) = typeOf e -typeOf (DoLoop _ _ _ _ _ (Info (t, _)) _) = t -typeOf (Lambda params _ _ (Info (als, t)) _) = - unscopeType bound_here $ foldr (arrow . patternParam) t params `setAliases` als - where bound_here = S.map identName (mconcat $ map patternIdents params) `S.difference` - S.fromList (mapMaybe (named . patternParam) params) - arrow (px, tx) y = Scalar $ Arrow () px tx y - named (Named x, _) = Just x - 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 (ProjectSection _ (Info t) _) = t -typeOf (IndexSection _ (Info t) _) = t -typeOf (Constr _ _ (Info t) _) = t -typeOf (Match _ cs (Info t, _) _) = - unscopeType (foldMap unscopeSet cs) t - where unscopeSet (CasePat p _ _) = S.map identName $ patternIdents p -typeOf (Attr _ e _) = typeOf e - --- | @foldFunType ts ret@ creates a function type ('Arrow') that takes --- @ts@ as parameters and returns @ret@. -foldFunType :: Monoid as => [TypeBase dim as] -> TypeBase dim as -> TypeBase dim as -foldFunType ps ret = foldr arrow ret ps - where arrow t1 t2 = Scalar $ Arrow mempty Unnamed t1 t2 - --- | Extract the parameter types and return type from a type. --- If the type is not an arrow type, the list of parameter types is empty. -unfoldFunType :: TypeBase dim as -> ([TypeBase dim as], TypeBase dim as) -unfoldFunType (Scalar (Arrow _ _ t1 t2)) = - let (ps, r) = unfoldFunType t2 - in (t1 : ps, r) -unfoldFunType t = ([], t) - --- | The type names mentioned in a type. -typeVars :: Monoid as => TypeBase dim as -> S.Set VName -typeVars t = - case t of - Scalar Prim{} -> mempty - Scalar (TypeVar _ _ tn targs) -> - mconcat $ typeVarFree tn : map typeArgFree targs - Scalar (Arrow _ _ t1 t2) -> typeVars t1 <> typeVars t2 - Scalar (Record fields) -> foldMap typeVars fields - Scalar (Sum cs) -> mconcat $ (foldMap . fmap) typeVars cs - Array _ _ rt _ -> typeVars $ Scalar rt - where typeVarFree = S.singleton . typeLeaf - typeArgFree (TypeArgType ta _) = typeVars ta - typeArgFree TypeArgDim{} = mempty - --- | @orderZero t@ is 'True' if the argument type has order 0, i.e., it is not --- a function type, does not contain a function type as a subcomponent, and may --- not be instantiated with a function type. -orderZero :: TypeBase dim as -> Bool -orderZero Array{} = True -orderZero (Scalar (Prim _)) = True -orderZero (Scalar (Record fs)) = all orderZero $ M.elems fs -orderZero (Scalar TypeVar{}) = True -orderZero (Scalar Arrow{}) = False -orderZero (Scalar (Sum cs)) = all (all orderZero) cs - --- | Extract all the shape names that occur in a given pattern. -patternDimNames :: PatternBase Info VName -> S.Set VName -patternDimNames (TuplePattern ps _) = foldMap patternDimNames ps -patternDimNames (RecordPattern fs _) = foldMap (patternDimNames . snd) fs -patternDimNames (PatternParens p _) = patternDimNames p -patternDimNames (Id _ (Info tp) _) = typeDimNames tp -patternDimNames (Wildcard (Info tp) _) = typeDimNames tp -patternDimNames (PatternAscription p (TypeDecl _ (Info t)) _) = - patternDimNames p <> typeDimNames t -patternDimNames (PatternLit _ (Info tp) _) = typeDimNames tp -patternDimNames (PatternConstr _ _ ps _) = foldMap patternDimNames ps - --- | Extract all the shape names that occur in a given type. -typeDimNames :: TypeBase (DimDecl VName) als -> S.Set VName -typeDimNames = foldMap dimName . nestedDims - where dimName :: DimDecl VName -> S.Set VName - dimName (NamedDim qn) = S.singleton $ qualLeaf qn - dimName _ = mempty - --- | @patternOrderZero pat@ is 'True' if all of the types in the given pattern --- have order 0. -patternOrderZero :: PatternBase Info vn -> Bool -patternOrderZero pat = case pat of - TuplePattern ps _ -> all patternOrderZero ps - RecordPattern fs _ -> all (patternOrderZero . snd) fs - PatternParens p _ -> patternOrderZero p - Id _ (Info t) _ -> orderZero t - Wildcard (Info t) _ -> orderZero t - PatternAscription p _ _ -> patternOrderZero p - PatternLit _ (Info t) _ -> orderZero t - PatternConstr _ _ ps _ -> all patternOrderZero ps - --- | The set of identifiers bound in a pattern. -patternIdents :: (Functor f, Ord vn) => PatternBase f vn -> S.Set (IdentBase f vn) -patternIdents (Id v t loc) = S.singleton $ Ident v t loc -patternIdents (PatternParens p _) = patternIdents p -patternIdents (TuplePattern pats _) = mconcat $ map patternIdents pats -patternIdents (RecordPattern fs _) = mconcat $ map (patternIdents . snd) fs -patternIdents Wildcard{} = mempty -patternIdents (PatternAscription p _ _) = patternIdents p -patternIdents PatternLit{} = mempty -patternIdents (PatternConstr _ _ ps _ ) = mconcat $ map patternIdents ps - --- | The set of names bound in a pattern. -patternNames :: (Functor f, Ord vn) => PatternBase f vn -> S.Set vn -patternNames (Id v _ _) = S.singleton v -patternNames (PatternParens p _) = patternNames p -patternNames (TuplePattern pats _) = mconcat $ map patternNames pats -patternNames (RecordPattern fs _) = mconcat $ map (patternNames . snd) fs -patternNames Wildcard{} = mempty -patternNames (PatternAscription p _ _) = patternNames p -patternNames PatternLit{} = mempty -patternNames (PatternConstr _ _ ps _ ) = mconcat $ map patternNames ps - --- | A mapping from names bound in a map to their identifier. -patternMap :: (Functor f) => PatternBase f VName -> M.Map VName (IdentBase f VName) -patternMap pat = - M.fromList $ zip (map identName idents) idents - where idents = S.toList $ patternIdents pat - --- | The type of values bound by the pattern. -patternType :: PatternBase Info VName -> PatternType -patternType (Wildcard (Info t) _) = t -patternType (PatternParens p _) = patternType p -patternType (Id _ (Info t) _) = t -patternType (TuplePattern pats _) = tupleRecord $ map patternType pats -patternType (RecordPattern fs _) = Scalar $ Record $ patternType <$> M.fromList fs -patternType (PatternAscription p _ _) = patternType p -patternType (PatternLit _ (Info t) _) = t -patternType (PatternConstr _ (Info t) _ _) = t - --- | The type matched by the pattern, including shape declarations if present. -patternStructType :: PatternBase Info VName -> StructType -patternStructType = toStruct . patternType - --- | When viewed as a function parameter, does this pattern correspond --- to a named parameter of some type? -patternParam :: PatternBase Info VName -> (PName, StructType) -patternParam (PatternParens p _) = - patternParam p -patternParam (PatternAscription (Id v _ _) td _) = - (Named v, unInfo $ expandedType td) -patternParam (Id v (Info t) _) = - (Named v, toStruct t) -patternParam p = - (Unnamed, patternStructType p) - --- | Names of primitive types to types. This is only valid if no --- shadowing is going on, but useful for tools. -namesToPrimTypes :: M.Map Name PrimType -namesToPrimTypes = M.fromList - [ (nameFromString $ pretty t, t) | - t <- Bool : - map Signed [minBound..maxBound] ++ - map Unsigned [minBound..maxBound] ++ - map FloatType [minBound..maxBound] ] - --- | The nature of something predefined. These can either be --- monomorphic or overloaded. An overloaded builtin is a list valid --- types it can be instantiated with, to the parameter and result --- type, with 'Nothing' representing the overloaded parameter type. -data Intrinsic = IntrinsicMonoFun [PrimType] PrimType - | IntrinsicOverloadedFun [PrimType] [Maybe PrimType] (Maybe PrimType) - | IntrinsicPolyFun [TypeParamBase VName] [StructType] StructType - | IntrinsicType PrimType - | IntrinsicEquality -- Special cased. - --- | A map of all built-ins. -intrinsics :: M.Map VName Intrinsic -intrinsics = M.fromList $ zipWith namify [10..] $ - - map primFun (M.toList Primitive.primFuns) ++ - - [("opaque", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a)] ++ - - map unOpFun Primitive.allUnOps ++ - - map binOpFun Primitive.allBinOps ++ - - map cmpOpFun Primitive.allCmpOps ++ - - map convOpFun Primitive.allConvOps ++ - - map signFun Primitive.allIntTypes ++ - - map unsignFun Primitive.allIntTypes ++ - - map intrinsicType (map Signed [minBound..maxBound] ++ - map Unsigned [minBound..maxBound] ++ - map FloatType [minBound..maxBound] ++ - [Bool]) ++ - - -- This overrides the ! from Primitive. - [ ("!", IntrinsicOverloadedFun - (map Signed [minBound..maxBound] ++ - map Unsigned [minBound..maxBound] ++ - [Bool]) - [Nothing] Nothing) ] ++ - - -- The reason for the loop formulation is to ensure that we - -- get a missing case warning if we forget a case. - mapMaybe mkIntrinsicBinOp [minBound..maxBound] ++ - - [("flatten", IntrinsicPolyFun [tp_a] - [Array () Nonunique t_a (rank 2)] $ - Array () Nonunique t_a (rank 1)), - ("unflatten", IntrinsicPolyFun [tp_a] - [Scalar $ Prim $ Signed Int32, - Scalar $ Prim $ Signed Int32, - Array () Nonunique t_a (rank 1)] $ - Array () Nonunique t_a (rank 2)), - - ("concat", IntrinsicPolyFun [tp_a] - [arr_a, arr_a] uarr_a), - ("rotate", IntrinsicPolyFun [tp_a] - [Scalar $ Prim $ Signed Int32, arr_a] arr_a), - ("transpose", IntrinsicPolyFun [tp_a] [arr_2d_a] arr_2d_a), - - ("scatter", IntrinsicPolyFun [tp_a] - [Array () Unique t_a (rank 1), - Array () Nonunique (Prim $ Signed Int32) (rank 1), - Array () Nonunique t_a (rank 1)] $ - Array () Unique t_a (rank 1)), - - ("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", IntrinsicPolyFun [tp_a] - [Scalar $ Prim $ Signed Int32, - uarr_a, - Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), - Scalar t_a, - Array () Nonunique (Prim $ Signed Int32) (rank 1), - arr_a] - uarr_a), - - ("map", IntrinsicPolyFun [tp_a, tp_b] [Scalar t_a `arr` Scalar t_b, arr_a] uarr_b), - - ("reduce", IntrinsicPolyFun [tp_a] - [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a] $ - Scalar t_a), - - ("reduce_comm", IntrinsicPolyFun [tp_a] - [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a] $ - Scalar t_a), - - ("scan", IntrinsicPolyFun [tp_a] - [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a] uarr_a), - - ("partition", - IntrinsicPolyFun [tp_a] - [Scalar (Prim $ Signed Int32), - Scalar t_a `arr` Scalar (Prim $ Signed Int32), arr_a] $ - tupleRecord [uarr_a, Array () Unique (Prim $ Signed Int32) (rank 1)]), - - ("map_stream", - IntrinsicPolyFun [tp_a, tp_b] - [Scalar (Prim $ Signed Int32) `karr` (arr_ka `arr` arr_kb), arr_a] - uarr_b), - - ("map_stream_per", - IntrinsicPolyFun [tp_a, tp_b] - [Scalar (Prim $ Signed Int32) `karr` (arr_ka `arr` arr_kb), arr_a] - uarr_b), - - ("reduce_stream", - IntrinsicPolyFun [tp_a, tp_b] - [Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b), - Scalar (Prim $ Signed Int32) `karr` (arr_ka `arr` Scalar t_b), - arr_a] $ - Scalar t_b), - - ("reduce_stream_per", - IntrinsicPolyFun [tp_a, tp_b] - [Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b), - Scalar (Prim $ Signed Int32) `karr` (arr_ka `arr` Scalar t_b), - arr_a] $ - Scalar t_b), - - ("trace", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a), - ("break", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a)] - - where tv_a = VName (nameFromString "a") 0 - 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_a = Array () Unique t_a (rank 1) - tp_a = TypeParamType Unlifted tv_a mempty - - tv_b = VName (nameFromString "b") 1 - t_b = TypeVar () Nonunique (typeName tv_b) [] - arr_b = Array () Nonunique t_b (rank 1) - uarr_b = Array () Unique t_b (rank 1) - tp_b = TypeParamType Unlifted tv_b mempty - - arr_a_b = Array () Nonunique - (Record (M.fromList $ zip tupleFieldNames [Scalar t_a, Scalar t_b])) - (rank 1) - t_arr_a_arr_b = Scalar $ Record $ M.fromList $ zip tupleFieldNames [arr_a, arr_b] - - arr x y = Scalar $ Arrow mempty Unnamed x y - - kv = VName (nameFromString "k") 2 - arr_ka = Array () Nonunique t_a (ShapeDecl [NamedDim $ qualName kv]) - arr_kb = Array () Nonunique t_b (ShapeDecl [NamedDim $ qualName kv]) - karr x y = Scalar $ Arrow mempty (Named kv) x y - - namify i (k,v) = (VName (nameFromString k) i, v) - - primFun (name, (ts,t, _)) = - (name, IntrinsicMonoFun (map unPrim ts) $ unPrim t) - - unOpFun bop = (pretty bop, IntrinsicMonoFun [t] t) - where t = unPrim $ Primitive.unOpType bop - - binOpFun bop = (pretty bop, IntrinsicMonoFun [t, t] t) - where t = unPrim $ Primitive.binOpType bop - - cmpOpFun bop = (pretty bop, IntrinsicMonoFun [t, t] Bool) - where t = unPrim $ Primitive.cmpOpType bop - - convOpFun cop = (pretty cop, IntrinsicMonoFun [unPrim ft] $ unPrim tt) - where (ft, tt) = Primitive.convOpType cop - - signFun t = ("sign_" ++ pretty t, IntrinsicMonoFun [Unsigned t] $ Signed t) - - unsignFun t = ("unsign_" ++ pretty t, IntrinsicMonoFun [Signed t] $ Unsigned t) - - unPrim (Primitive.IntType t) = Signed t - unPrim (Primitive.FloatType t) = FloatType t - unPrim Primitive.Bool = Bool - unPrim Primitive.Cert = Bool - - intrinsicType t = (pretty t, IntrinsicType t) - - anyIntType = map Signed [minBound..maxBound] ++ - map Unsigned [minBound..maxBound] - anyNumberType = anyIntType ++ - map FloatType [minBound..maxBound] - anyPrimType = Bool : anyNumberType - - mkIntrinsicBinOp :: BinOp -> Maybe (String, Intrinsic) - mkIntrinsicBinOp op = do op' <- intrinsicBinOp op - return (pretty op, op') - - binOp ts = Just $ IntrinsicOverloadedFun ts [Nothing, Nothing] Nothing - ordering = Just $ IntrinsicOverloadedFun anyPrimType [Nothing, Nothing] (Just Bool) - - intrinsicBinOp Plus = binOp anyNumberType - intrinsicBinOp Minus = binOp anyNumberType - intrinsicBinOp Pow = binOp anyNumberType - intrinsicBinOp Times = binOp anyNumberType - intrinsicBinOp Divide = binOp anyNumberType - intrinsicBinOp Mod = binOp anyNumberType - intrinsicBinOp Quot = binOp anyIntType - intrinsicBinOp Rem = binOp anyIntType - intrinsicBinOp ShiftR = binOp anyIntType - intrinsicBinOp ShiftL = binOp anyIntType - intrinsicBinOp Band = binOp anyIntType - intrinsicBinOp Xor = binOp anyIntType - intrinsicBinOp Bor = binOp anyIntType - intrinsicBinOp LogAnd = binOp [Bool] - intrinsicBinOp LogOr = binOp [Bool] - intrinsicBinOp Equal = Just IntrinsicEquality - intrinsicBinOp NotEqual = Just IntrinsicEquality - intrinsicBinOp Less = ordering - intrinsicBinOp Leq = ordering - intrinsicBinOp Greater = ordering - intrinsicBinOp Geq = ordering - intrinsicBinOp _ = Nothing - --- | 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. -maxIntrinsicTag :: Int -maxIntrinsicTag = maxinum $ map baseTag $ M.keys intrinsics - --- | Create a name with no qualifiers from a name. -qualName :: v -> QualName v -qualName = QualName [] - --- | Add another qualifier (at the head) to a qualified name. -qualify :: v -> QualName v -> QualName v -qualify k (QualName ks v) = QualName (k:ks) v - --- | Create a type name name with no qualifiers from a 'VName'. -typeName :: VName -> TypeName -typeName = typeNameFromQualName . qualName - --- | The modules imported by a Futhark program. -progImports :: ProgBase f vn -> [(String,SrcLoc)] -progImports = concatMap decImports . progDecs - --- | The modules imported by a single declaration. -decImports :: DecBase f vn -> [(String,SrcLoc)] -decImports (OpenDec x _) = modExpImports x -decImports (ModDec md) = modExpImports $ modExp md -decImports SigDec{} = [] -decImports TypeDec{} = [] -decImports ValDec{} = [] -decImports (LocalDec d _) = decImports d -decImports (ImportDec x _ loc) = [(x, loc)] - -modExpImports :: ModExpBase f vn -> [(String,SrcLoc)] -modExpImports ModVar{} = [] -modExpImports (ModParens p _) = modExpImports p -modExpImports (ModImport f _ loc) = [(f,loc)] -modExpImports (ModDecs ds _) = concatMap decImports ds -modExpImports (ModApply _ me _ _ _) = modExpImports me -modExpImports (ModAscript me _ _ _) = modExpImports me -modExpImports ModLambda{} = [] - --- | The set of module types used in any exported (non-local) --- declaration. -progModuleTypes :: Ord vn => ProgBase f vn -> S.Set vn -progModuleTypes = mconcat . map onDec . progDecs - where onDec (OpenDec x _) = onModExp x - onDec (ModDec md) = - maybe mempty (onSigExp . fst) (modSignature md) <> onModExp (modExp md) - onDec SigDec{} = mempty - onDec TypeDec{} = mempty - onDec ValDec{} = mempty - onDec LocalDec{} = mempty - onDec ImportDec{} = mempty - - onModExp ModVar{} = mempty - onModExp (ModParens p _) = onModExp p - onModExp ModImport {} = mempty - onModExp (ModDecs ds _) = mconcat $ map onDec ds - onModExp (ModApply me1 me2 _ _ _) = onModExp me1 <> onModExp me2 - onModExp (ModAscript me se _ _) = onModExp me <> onSigExp se - onModExp (ModLambda p r me _) = - onModParam p <> maybe mempty (onSigExp . fst) r <> onModExp me - - onModParam = onSigExp . modParamType - - onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v - onSigExp (SigParens e _) = onSigExp e - onSigExp SigSpecs{} = mempty - onSigExp (SigWith e _ _) = onSigExp e - onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2 - --- | Extract a leading @((name, namespace, file), remainder)@ from a --- documentation comment string. These are formatted as --- \`name\`\@namespace[\@file]. Let us hope that this pattern does not occur --- anywhere else. -identifierReference :: String -> Maybe ((String, String, Maybe FilePath), String) -identifierReference ('`' : s) - | (identifier, '`' : '@' : s') <- break (=='`') s, - (namespace, s'') <- span isAlpha s', - not $ null namespace = - case s'' of - '@' : '"' : s''' - | (file, '"' : s'''') <- span (/= '"') s''' -> - Just ((identifier, namespace, Just file), s'''') - _ -> Just ((identifier, namespace, Nothing), s'') - -identifierReference _ = Nothing - --- | Given an operator name, return the operator that determines its --- syntactical properties. -leadingOperator :: Name -> BinOp -leadingOperator s = maybe Backtick snd $ find ((`isPrefixOf` s') . fst) $ - sortOn (Down . length . fst) $ - zip (map pretty operators) operators - where s' = nameToString s - operators :: [BinOp] - operators = [minBound..maxBound::BinOp] +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | This module provides various simple ways to query and manipulate +-- fundamental Futhark terms, such as types and values. The intent is to +-- keep "Futhark.Language.Syntax" simple, and put whatever embellishments +-- we need here. +module Language.Futhark.Prop + ( -- * Various + Intrinsic (..), + intrinsics, + maxIntrinsicTag, + namesToPrimTypes, + qualName, + qualify, + typeName, + valueType, + primValueType, + leadingOperator, + progImports, + decImports, + progModuleTypes, + identifierReference, + prettyStacktrace, + + -- * Queries on expressions + typeOf, + + -- * Queries on patterns and params + patternIdents, + patternNames, + patternMap, + patternType, + patternStructType, + patternParam, + patternOrderZero, + patternDimNames, + + -- * Queries on types + uniqueness, + unique, + aliases, + diet, + arrayRank, + arrayShape, + nestedDims, + orderZero, + unfoldFunType, + foldFunType, + typeVars, + typeDimNames, + primByteSize, + + -- * Operations on types + rank, + peelArray, + stripArray, + arrayOf, + toStructural, + toStruct, + fromStruct, + setAliases, + addAliases, + setUniqueness, + noSizes, + addSizes, + anySizes, + traverseDims, + DimPos (..), + mustBeExplicit, + mustBeExplicitInType, + tupleRecord, + isTupleRecord, + areTupleFields, + tupleFields, + tupleFieldNames, + sortFields, + sortConstrs, + isTypeParam, + isSizeParam, + combineTypeShapes, + matchDims, + unscopeType, + onRecordField, + -- | Values of these types are produces by the parser. They use + -- unadorned names and have no type information, apart from that + -- which is syntactically required. + NoInfo (..), + UncheckedType, + UncheckedTypeExp, + UncheckedIdent, + UncheckedTypeDecl, + UncheckedDimIndex, + UncheckedExp, + UncheckedModExp, + UncheckedSigExp, + UncheckedTypeParam, + UncheckedPattern, + UncheckedValBind, + UncheckedDec, + UncheckedSpec, + UncheckedProg, + UncheckedCase, + ) +where + +import Control.Monad.State +import Data.Bifoldable +import Data.Bifunctor +import Data.Bitraversable (bitraverse) +import Data.Char +import Data.Foldable +import Data.List (genericLength, isPrefixOf, nub, sortOn) +import qualified Data.Map.Strict as M +import Data.Maybe +import Data.Ord +import qualified Data.Set as S +import qualified Futhark.IR.Primitive as Primitive +import Futhark.Util (maxinum) +import Futhark.Util.Pretty +import Language.Futhark.Syntax + +-- | Return the dimensionality of a type. For non-arrays, this is +-- zero. For a one-dimensional array it is one, for a two-dimensional +-- it is two, and so forth. +arrayRank :: TypeBase dim as -> Int +arrayRank = shapeRank . arrayShape + +-- | Return the shape of a type - for non-arrays, this is 'mempty'. +arrayShape :: TypeBase dim as -> ShapeDecl dim +arrayShape (Array _ _ _ ds) = ds +arrayShape _ = mempty + +-- | Return any shape declarations in the type, with duplicates +-- removed. +nestedDims :: TypeBase (DimDecl VName) as -> [DimDecl VName] +nestedDims t = + case t of + Array _ _ a ds -> + nub $ nestedDims (Scalar a) <> shapeDims ds + Scalar (Record fs) -> + nub $ foldMap nestedDims fs + Scalar Prim {} -> + mempty + Scalar (Sum cs) -> + nub $ foldMap (foldMap nestedDims) cs + Scalar (Arrow _ v t1 t2) -> + filter (notV v) $ nestedDims t1 <> nestedDims t2 + Scalar (TypeVar _ _ _ targs) -> + concatMap typeArgDims targs + where + typeArgDims (TypeArgDim d _) = [d] + typeArgDims (TypeArgType at _) = nestedDims at + + notV Unnamed = const True + notV (Named v) = (/= NamedDim (qualName v)) + +-- | Change the shape of a type to be just the rank. +noSizes :: TypeBase (DimDecl vn) as -> TypeBase () as +noSizes = first $ const () + +-- | Add size annotations that are all 'AnyDim'. +addSizes :: TypeBase () as -> TypeBase (DimDecl vn) as +addSizes = first $ const AnyDim + +-- | Change all size annotations to be 'AnyDim'. +anySizes :: TypeBase (DimDecl vn) as -> TypeBase (DimDecl vn) as +anySizes = first $ const AnyDim + +-- | Where does this dimension occur? +data DimPos + = -- | Immediately in the argument to 'traverseDims'. + PosImmediate + | -- | In a function parameter type. + PosParam + | -- | In a function return type. + PosReturn + deriving (Eq, Ord, Show) + +-- | Perform a traversal (possibly including replacement) on sizes +-- that are parameters in a function type, but also including the type +-- immediately passed to the function. Also passes along a set of the +-- parameter names inside the type that have come in scope at the +-- occurrence of the dimension. +traverseDims :: + forall f fdim tdim als. + Applicative f => + (S.Set VName -> DimPos -> fdim -> f tdim) -> + TypeBase fdim als -> + f (TypeBase tdim als) +traverseDims f = go mempty PosImmediate + where + go :: + forall als'. + S.Set VName -> + DimPos -> + TypeBase fdim als' -> + f (TypeBase tdim als') + go bound b t@Array {} = + bitraverse (f bound b) pure t + go bound b (Scalar (Record fields)) = + Scalar . Record <$> traverse (go bound b) fields + go bound b (Scalar (TypeVar as u tn targs)) = + Scalar <$> (TypeVar as u tn <$> traverse (onTypeArg bound b) targs) + go bound b (Scalar (Sum cs)) = + Scalar . Sum <$> traverse (traverse (go bound b)) cs + go _ _ (Scalar (Prim t)) = + pure $ Scalar $ Prim t + go bound _ (Scalar (Arrow als p t1 t2)) = + Scalar <$> (Arrow als p <$> go bound' PosParam t1 <*> go bound' PosReturn t2) + where + bound' = case p of + Named p' -> S.insert p' bound + Unnamed -> bound + + onTypeArg bound b (TypeArgDim d loc) = + TypeArgDim <$> f bound b d <*> pure loc + onTypeArg bound b (TypeArgType t loc) = + TypeArgType <$> go bound b t <*> pure loc + +mustBeExplicitAux :: StructType -> M.Map VName Bool +mustBeExplicitAux t = + execState (traverseDims onDim t) mempty + where + onDim bound _ (NamedDim d) + | qualLeaf d `S.member` bound = + modify $ \s -> M.insertWith (&&) (qualLeaf d) False s + onDim _ PosImmediate (NamedDim d) = + modify $ \s -> M.insertWith (&&) (qualLeaf d) False s + onDim _ _ (NamedDim d) = + modify $ M.insertWith (&&) (qualLeaf d) True + onDim _ _ _ = + return () + +-- | Figure out which of the sizes in a parameter type must be passed +-- explicitly, because their first use is as something else than just +-- an array dimension. 'mustBeExplicit' is like this function, but +-- first decomposes into parameter types. +mustBeExplicitInType :: StructType -> S.Set VName +mustBeExplicitInType t = + S.fromList $ M.keys $ M.filter id $ mustBeExplicitAux t + +-- | Figure out which of the sizes in a binding type must be passed +-- explicitly, because their first use is as something else than just +-- an array dimension. +mustBeExplicit :: StructType -> S.Set VName +mustBeExplicit bind_t = + let (ts, ret) = unfoldFunType bind_t + alsoRet = + M.unionWith (&&) $ + M.fromList $ zip (S.toList $ typeDimNames ret) $ repeat True + in S.fromList $ M.keys $ M.filter id $ alsoRet $ foldl' onType mempty ts + where + onType uses t = uses <> mustBeExplicitAux t -- Left-biased union. + +-- | Return the uniqueness of a type. +uniqueness :: TypeBase shape as -> Uniqueness +uniqueness (Array _ u _ _) = u +uniqueness (Scalar (TypeVar _ u _ _)) = u +uniqueness (Scalar (Sum ts)) = foldMap (foldMap uniqueness) $ M.elems ts +uniqueness (Scalar (Record fs)) = foldMap uniqueness $ M.elems fs +uniqueness _ = Nonunique + +-- | @unique t@ is 'True' if the type of the argument is unique. +unique :: TypeBase shape as -> Bool +unique = (== Unique) . uniqueness + +-- | Return the set of all variables mentioned in the aliasing of a +-- type. +aliases :: Monoid as => TypeBase shape as -> as +aliases = bifoldMap (const mempty) id + +-- | @diet t@ returns a description of how a function parameter of +-- type @t@ might consume its argument. +diet :: TypeBase shape as -> Diet +diet (Scalar (Record ets)) = RecordDiet $ fmap diet ets +diet (Scalar (Prim _)) = Observe +diet (Scalar (Arrow _ _ t1 t2)) = FuncDiet (diet t1) (diet t2) +diet (Array _ Unique _ _) = Consume +diet (Array _ Nonunique _ _) = Observe +diet (Scalar (TypeVar _ Unique _ _)) = Consume +diet (Scalar (TypeVar _ Nonunique _ _)) = Observe +diet (Scalar Sum {}) = Observe + +-- | Convert any type to one that has rank information, no alias +-- information, and no embedded names. +toStructural :: + TypeBase dim as -> + TypeBase () () +toStructural = flip setAliases () . first (const ()) + +-- | Remove aliasing information from a type. +toStruct :: + TypeBase dim as -> + TypeBase dim () +toStruct t = t `setAliases` () + +-- | Replace no aliasing with an empty alias set. +fromStruct :: + TypeBase dim as -> + TypeBase dim Aliasing +fromStruct t = t `setAliases` S.empty + +-- | @peelArray n t@ returns the type resulting from peeling the first +-- @n@ array dimensions from @t@. Returns @Nothing@ if @t@ has less +-- than @n@ dimensions. +peelArray :: Int -> TypeBase dim as -> Maybe (TypeBase dim as) +peelArray n (Array als u t shape) + | shapeRank shape == n = + Just $ Scalar t `addAliases` const als + | otherwise = + Array als u t <$> stripDims n shape +peelArray _ _ = Nothing + +-- | @arrayOf t s u@ constructs an array type. The convenience +-- compared to using the 'Array' constructor directly is that @t@ can +-- itself be an array. If @t@ is an @n@-dimensional array, and @s@ is +-- a list of length @n@, the resulting type is of an @n+m@ dimensions. +-- The uniqueness of the new array will be @u@, no matter the +-- uniqueness of @t@. +arrayOf :: + Monoid as => + TypeBase dim as -> + ShapeDecl dim -> + Uniqueness -> + TypeBase dim as +arrayOf t = arrayOfWithAliases (t `setUniqueness` Nonunique) mempty + +arrayOfWithAliases :: + Monoid as => + TypeBase dim as -> + as -> + ShapeDecl dim -> + Uniqueness -> + TypeBase dim as +arrayOfWithAliases (Array as1 _ et shape1) as2 shape2 u = + Array (as1 <> as2) u et (shape2 <> shape1) +arrayOfWithAliases (Scalar t) as shape u = + Array as u (second (const ()) t) shape + +-- | @stripArray n t@ removes the @n@ outermost layers of the array. +-- Essentially, it is the type of indexing an array of type @t@ with +-- @n@ indexes. +stripArray :: Int -> TypeBase dim as -> TypeBase dim as +stripArray n (Array als u et shape) + | Just shape' <- stripDims n shape = + Array als u et shape' + | otherwise = + Scalar et `setUniqueness` u `setAliases` als +stripArray _ t = t + +-- | Create a record type corresponding to a tuple with the given +-- element types. +tupleRecord :: [TypeBase dim as] -> TypeBase dim as +tupleRecord = Scalar . Record . M.fromList . zip tupleFieldNames + +-- | Does this type corespond to a tuple? If so, return the elements +-- of that tuple. +isTupleRecord :: TypeBase dim as -> Maybe [TypeBase dim as] +isTupleRecord (Scalar (Record fs)) = areTupleFields fs +isTupleRecord _ = Nothing + +-- | Does this record map correspond to a tuple? +areTupleFields :: M.Map Name a -> Maybe [a] +areTupleFields fs = + let fs' = sortFields fs + in if and $ zipWith (==) (map fst fs') tupleFieldNames + then Just $ map snd fs' + else Nothing + +-- | Construct a record map corresponding to a tuple. +tupleFields :: [a] -> M.Map Name a +tupleFields as = M.fromList $ zip tupleFieldNames as + +-- | Increasing field names for a tuple (starts at 0). +tupleFieldNames :: [Name] +tupleFieldNames = map (nameFromString . show) [(0 :: Int) ..] + +-- | Sort fields by their name; taking care to sort numeric fields by +-- their numeric value. This ensures that tuples and tuple-like +-- records match. +sortFields :: M.Map Name a -> [(Name, a)] +sortFields l = map snd $ sortOn fst $ zip (map (fieldish . fst) l') l' + where + l' = M.toList l + fieldish s = case reads $ nameToString s of + [(x, "")] -> Left (x :: Int) + _ -> Right s + +-- | Sort the constructors of a sum type in some well-defined (but not +-- otherwise significant) manner. +sortConstrs :: M.Map Name a -> [(Name, a)] +sortConstrs cs = sortOn fst $ M.toList cs + +-- | Is this a 'TypeParamType'? +isTypeParam :: TypeParamBase vn -> Bool +isTypeParam TypeParamType {} = True +isTypeParam TypeParamDim {} = False + +-- | Is this a 'TypeParamDim'? +isSizeParam :: TypeParamBase vn -> Bool +isSizeParam = not . isTypeParam + +-- | Combine the shape information of types as much as possible. The first +-- argument is the orignal type and the second is the type of the transformed +-- expression. This is necessary since the original type may contain additional +-- information (e.g., shape restrictions) from the user given annotation. +combineTypeShapes :: + (Monoid as, ArrayDim dim) => + TypeBase dim as -> + TypeBase dim as -> + TypeBase dim as +combineTypeShapes (Scalar (Record ts1)) (Scalar (Record ts2)) + | M.keys ts1 == M.keys ts2 = + Scalar $ Record $ M.map (uncurry combineTypeShapes) (M.intersectionWith (,) ts1 ts2) +combineTypeShapes (Scalar (Arrow als1 p1 a1 b1)) (Scalar (Arrow als2 _p2 a2 b2)) = + Scalar $ Arrow (als1 <> als2) p1 (combineTypeShapes a1 a2) (combineTypeShapes b1 b2) +combineTypeShapes (Scalar (TypeVar als1 u1 v targs1)) (Scalar (TypeVar als2 _ _ targs2)) = + Scalar $ TypeVar (als1 <> als2) u1 v $ zipWith f targs1 targs2 + where + f (TypeArgType t1 loc) (TypeArgType t2 _) = + TypeArgType (combineTypeShapes t1 t2) loc + f targ _ = targ +combineTypeShapes (Array als1 u1 et1 shape1) (Array als2 _u2 et2 _shape2) = + arrayOfWithAliases + ( combineTypeShapes (Scalar et1) (Scalar et2) + `setAliases` mempty + ) + (als1 <> als2) + shape1 + u1 +combineTypeShapes _ new_tp = new_tp + +-- | Match the dimensions of otherwise assumed-equal types. +matchDims :: + (Monoid as, Monad m) => + (d1 -> d2 -> m d1) -> + TypeBase d1 as -> + TypeBase d2 as -> + m (TypeBase d1 as) +matchDims onDims t1 t2 = + case (t1, t2) of + (Array als1 u1 et1 shape1, Array als2 u2 et2 shape2) -> + flip setAliases (als1 <> als2) + <$> ( arrayOf + <$> matchDims onDims (Scalar et1) (Scalar et2) + <*> onShapes shape1 shape2 + <*> pure (min u1 u2) + ) + (Scalar (Record f1), Scalar (Record f2)) -> + Scalar . Record + <$> traverse (uncurry (matchDims onDims)) (M.intersectionWith (,) f1 f2) + (Scalar (Sum cs1), Scalar (Sum cs2)) -> + Scalar . Sum + <$> traverse + (traverse (uncurry (matchDims onDims))) + (M.intersectionWith zip cs1 cs2) + (Scalar (Arrow als1 p1 a1 b1), Scalar (Arrow als2 _p2 a2 b2)) -> + Scalar + <$> (Arrow (als1 <> als2) p1 <$> matchDims onDims a1 a2 <*> matchDims onDims b1 b2) + ( Scalar (TypeVar als1 u v targs1), + Scalar (TypeVar als2 _ _ targs2) + ) -> + Scalar . TypeVar (als1 <> als2) u v <$> zipWithM matchTypeArg targs1 targs2 + _ -> return t1 + where + matchTypeArg ta@TypeArgType {} _ = return ta + matchTypeArg a _ = return a + + onShapes shape1 shape2 = + ShapeDecl <$> zipWithM onDims (shapeDims shape1) (shapeDims shape2) + +-- | Set the uniqueness attribute of a type. If the type is a record +-- or sum type, the uniqueness of its components will be modified. +setUniqueness :: TypeBase dim as -> Uniqueness -> TypeBase dim as +setUniqueness (Array als _ et shape) u = + Array als u et shape +setUniqueness (Scalar (TypeVar als _ t targs)) u = + Scalar $ TypeVar als u t targs +setUniqueness (Scalar (Record ets)) u = + Scalar $ Record $ fmap (`setUniqueness` u) ets +setUniqueness (Scalar (Sum ets)) u = + Scalar $ Sum $ fmap (map (`setUniqueness` u)) ets +setUniqueness t _ = t + +-- | @t \`setAliases\` als@ returns @t@, but with @als@ substituted for +-- any already present aliasing. +setAliases :: TypeBase dim asf -> ast -> TypeBase dim ast +setAliases t = addAliases t . const + +-- | @t \`addAliases\` f@ returns @t@, but with any already present +-- aliasing replaced by @f@ applied to that aliasing. +addAliases :: + TypeBase dim asf -> + (asf -> ast) -> + TypeBase dim ast +addAliases = flip second + +intValueType :: IntValue -> IntType +intValueType Int8Value {} = Int8 +intValueType Int16Value {} = Int16 +intValueType Int32Value {} = Int32 +intValueType Int64Value {} = Int64 + +floatValueType :: FloatValue -> FloatType +floatValueType Float32Value {} = Float32 +floatValueType Float64Value {} = Float64 + +-- | The type of a basic value. +primValueType :: PrimValue -> PrimType +primValueType (SignedValue v) = Signed $ intValueType v +primValueType (UnsignedValue v) = Unsigned $ intValueType v +primValueType (FloatValue v) = FloatType $ floatValueType v +primValueType BoolValue {} = Bool + +-- | The type of the value. +valueType :: Value -> ValueType +valueType (PrimValue bv) = Scalar $ Prim $ primValueType bv +valueType (ArrayValue _ t) = t + +-- | The size of values of this type, in bytes. +primByteSize :: Num a => PrimType -> a +primByteSize (Signed it) = Primitive.intByteSize it +primByteSize (Unsigned it) = Primitive.intByteSize it +primByteSize (FloatType ft) = Primitive.floatByteSize ft +primByteSize Bool = 1 + +-- | Construct a 'ShapeDecl' with the given number of 'AnyDim' +-- dimensions. +rank :: Int -> ShapeDecl (DimDecl VName) +rank n = ShapeDecl $ replicate n AnyDim + +-- | The type is leaving a scope, so clean up any aliases that +-- reference the bound variables, and turn any dimensions that name +-- them into AnyDim instead. +unscopeType :: S.Set VName -> PatternType -> PatternType +unscopeType bound_here t = first onDim $ t `addAliases` S.map unbind + where + unbind (AliasBound v) | v `S.member` bound_here = AliasFree v + unbind a = a + onDim (NamedDim qn) | qualLeaf qn `S.member` bound_here = AnyDim + onDim d = d + +-- | Perform some operation on a given record field. Returns +-- 'Nothing' if that field does not exist. +onRecordField :: + (TypeBase dim als -> TypeBase dim als) -> + [Name] -> + TypeBase dim als -> + Maybe (TypeBase dim als) +onRecordField f [] t = Just $ f t +onRecordField f (k : ks) (Scalar (Record m)) = do + t <- onRecordField f ks =<< M.lookup k m + Just $ Scalar $ Record $ M.insert k t m +onRecordField _ _ _ = Nothing + +-- | The type of an Futhark term. The aliasing will refer to itself, if +-- the term is a non-tuple-typed variable. +typeOf :: ExpBase Info VName -> PatternType +typeOf (Literal val _) = Scalar $ Prim $ primValueType val +typeOf (IntLit _ (Info t) _) = t +typeOf (FloatLit _ (Info t) _) = t +typeOf (Parens e _) = typeOf e +typeOf (QualParens _ e _) = typeOf e +typeOf (TupLit es _) = tupleRecord $ map typeOf es +typeOf (RecordLit fs _) = + -- Reverse, because M.unions is biased to the left. + Scalar $ Record $ M.unions $ reverse $ map record fs + where + record (RecordFieldExplicit name e _) = M.singleton name $ typeOf e + record (RecordFieldImplicit name (Info t) _) = + M.singleton (baseName name) $ + t + `addAliases` S.insert (AliasBound name) +typeOf (ArrayLit _ (Info t) _) = t +typeOf (StringLit vs _) = + Array + mempty + Unique + (Prim (Unsigned Int8)) + (ShapeDecl [ConstDim $ genericLength vs]) +typeOf (Range _ _ _ (Info t, _) _) = t +typeOf (BinOp _ _ _ _ (Info t) _ _) = t +typeOf (Project _ _ (Info t) _) = t +typeOf (If _ _ _ (Info t, _) _) = t +typeOf (Var _ (Info t) _) = t +typeOf (Ascript e _ _) = typeOf e +typeOf (Coerce _ _ (Info t, _) _) = t +typeOf (Apply _ _ _ (Info t, _) _) = t +typeOf (Negate e _) = typeOf e +typeOf (LetPat _ _ _ (Info t, _) _) = t +typeOf (LetFun _ _ _ (Info t) _) = t +typeOf (LetWith _ _ _ _ _ (Info t) _) = t +typeOf (Index _ _ (Info t, _) _) = t +typeOf (Update e _ _ _) = typeOf e `setAliases` mempty +typeOf (RecordUpdate _ _ _ (Info t) _) = t +typeOf (Assert _ e _ _) = typeOf e +typeOf (DoLoop _ _ _ _ _ (Info (t, _)) _) = t +typeOf (Lambda params _ _ (Info (als, t)) _) = + unscopeType bound_here $ foldr (arrow . patternParam) t params `setAliases` als + where + bound_here = + S.map identName (mconcat $ map patternIdents params) + `S.difference` S.fromList (mapMaybe (named . patternParam) params) + arrow (px, tx) y = Scalar $ Arrow () px tx y + named (Named x, _) = Just x + 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 (ProjectSection _ (Info t) _) = t +typeOf (IndexSection _ (Info t) _) = t +typeOf (Constr _ _ (Info t) _) = t +typeOf (Match _ cs (Info t, _) _) = + unscopeType (foldMap unscopeSet cs) t + where + unscopeSet (CasePat p _ _) = S.map identName $ patternIdents p +typeOf (Attr _ e _) = typeOf e + +-- | @foldFunType ts ret@ creates a function type ('Arrow') that takes +-- @ts@ as parameters and returns @ret@. +foldFunType :: Monoid as => [TypeBase dim as] -> TypeBase dim as -> TypeBase dim as +foldFunType ps ret = foldr arrow ret ps + where + arrow t1 t2 = Scalar $ Arrow mempty Unnamed t1 t2 + +-- | Extract the parameter types and return type from a type. +-- If the type is not an arrow type, the list of parameter types is empty. +unfoldFunType :: TypeBase dim as -> ([TypeBase dim as], TypeBase dim as) +unfoldFunType (Scalar (Arrow _ _ t1 t2)) = + let (ps, r) = unfoldFunType t2 + in (t1 : ps, r) +unfoldFunType t = ([], t) + +-- | The type names mentioned in a type. +typeVars :: Monoid as => TypeBase dim as -> S.Set VName +typeVars t = + case t of + Scalar Prim {} -> mempty + Scalar (TypeVar _ _ tn targs) -> + mconcat $ typeVarFree tn : map typeArgFree targs + Scalar (Arrow _ _ t1 t2) -> typeVars t1 <> typeVars t2 + Scalar (Record fields) -> foldMap typeVars fields + Scalar (Sum cs) -> mconcat $ (foldMap . fmap) typeVars cs + Array _ _ rt _ -> typeVars $ Scalar rt + where + typeVarFree = S.singleton . typeLeaf + typeArgFree (TypeArgType ta _) = typeVars ta + typeArgFree TypeArgDim {} = mempty + +-- | @orderZero t@ is 'True' if the argument type has order 0, i.e., it is not +-- a function type, does not contain a function type as a subcomponent, and may +-- not be instantiated with a function type. +orderZero :: TypeBase dim as -> Bool +orderZero Array {} = True +orderZero (Scalar (Prim _)) = True +orderZero (Scalar (Record fs)) = all orderZero $ M.elems fs +orderZero (Scalar TypeVar {}) = True +orderZero (Scalar Arrow {}) = False +orderZero (Scalar (Sum cs)) = all (all orderZero) cs + +-- | Extract all the shape names that occur in a given pattern. +patternDimNames :: PatternBase Info VName -> S.Set VName +patternDimNames (TuplePattern ps _) = foldMap patternDimNames ps +patternDimNames (RecordPattern fs _) = foldMap (patternDimNames . snd) fs +patternDimNames (PatternParens p _) = patternDimNames p +patternDimNames (Id _ (Info tp) _) = typeDimNames tp +patternDimNames (Wildcard (Info tp) _) = typeDimNames tp +patternDimNames (PatternAscription p (TypeDecl _ (Info t)) _) = + patternDimNames p <> typeDimNames t +patternDimNames (PatternLit _ (Info tp) _) = typeDimNames tp +patternDimNames (PatternConstr _ _ ps _) = foldMap patternDimNames ps + +-- | Extract all the shape names that occur in a given type. +typeDimNames :: TypeBase (DimDecl VName) als -> S.Set VName +typeDimNames = foldMap dimName . nestedDims + where + dimName :: DimDecl VName -> S.Set VName + dimName (NamedDim qn) = S.singleton $ qualLeaf qn + dimName _ = mempty + +-- | @patternOrderZero pat@ is 'True' if all of the types in the given pattern +-- have order 0. +patternOrderZero :: PatternBase Info vn -> Bool +patternOrderZero pat = case pat of + TuplePattern ps _ -> all patternOrderZero ps + RecordPattern fs _ -> all (patternOrderZero . snd) fs + PatternParens p _ -> patternOrderZero p + Id _ (Info t) _ -> orderZero t + Wildcard (Info t) _ -> orderZero t + PatternAscription p _ _ -> patternOrderZero p + PatternLit _ (Info t) _ -> orderZero t + PatternConstr _ _ ps _ -> all patternOrderZero ps + +-- | The set of identifiers bound in a pattern. +patternIdents :: (Functor f, Ord vn) => PatternBase f vn -> S.Set (IdentBase f vn) +patternIdents (Id v t loc) = S.singleton $ Ident v t loc +patternIdents (PatternParens p _) = patternIdents p +patternIdents (TuplePattern pats _) = mconcat $ map patternIdents pats +patternIdents (RecordPattern fs _) = mconcat $ map (patternIdents . snd) fs +patternIdents Wildcard {} = mempty +patternIdents (PatternAscription p _ _) = patternIdents p +patternIdents PatternLit {} = mempty +patternIdents (PatternConstr _ _ ps _) = mconcat $ map patternIdents ps + +-- | The set of names bound in a pattern. +patternNames :: (Functor f, Ord vn) => PatternBase f vn -> S.Set vn +patternNames (Id v _ _) = S.singleton v +patternNames (PatternParens p _) = patternNames p +patternNames (TuplePattern pats _) = mconcat $ map patternNames pats +patternNames (RecordPattern fs _) = mconcat $ map (patternNames . snd) fs +patternNames Wildcard {} = mempty +patternNames (PatternAscription p _ _) = patternNames p +patternNames PatternLit {} = mempty +patternNames (PatternConstr _ _ ps _) = mconcat $ map patternNames ps + +-- | A mapping from names bound in a map to their identifier. +patternMap :: (Functor f) => PatternBase f VName -> M.Map VName (IdentBase f VName) +patternMap pat = + M.fromList $ zip (map identName idents) idents + where + idents = S.toList $ patternIdents pat + +-- | The type of values bound by the pattern. +patternType :: PatternBase Info VName -> PatternType +patternType (Wildcard (Info t) _) = t +patternType (PatternParens p _) = patternType p +patternType (Id _ (Info t) _) = t +patternType (TuplePattern pats _) = tupleRecord $ map patternType pats +patternType (RecordPattern fs _) = Scalar $ Record $ patternType <$> M.fromList fs +patternType (PatternAscription p _ _) = patternType p +patternType (PatternLit _ (Info t) _) = t +patternType (PatternConstr _ (Info t) _ _) = t + +-- | The type matched by the pattern, including shape declarations if present. +patternStructType :: PatternBase Info VName -> StructType +patternStructType = toStruct . patternType + +-- | When viewed as a function parameter, does this pattern correspond +-- to a named parameter of some type? +patternParam :: PatternBase Info VName -> (PName, StructType) +patternParam (PatternParens p _) = + patternParam p +patternParam (PatternAscription (Id v _ _) td _) = + (Named v, unInfo $ expandedType td) +patternParam (Id v (Info t) _) = + (Named v, toStruct t) +patternParam p = + (Unnamed, patternStructType p) + +-- | Names of primitive types to types. This is only valid if no +-- shadowing is going on, but useful for tools. +namesToPrimTypes :: M.Map Name PrimType +namesToPrimTypes = + M.fromList + [ (nameFromString $ pretty t, t) + | t <- + Bool : + map Signed [minBound .. maxBound] + ++ map Unsigned [minBound .. maxBound] + ++ map FloatType [minBound .. maxBound] + ] + +-- | The nature of something predefined. These can either be +-- monomorphic or overloaded. An overloaded builtin is a list valid +-- types it can be instantiated with, to the parameter and result +-- type, with 'Nothing' representing the overloaded parameter type. +data Intrinsic + = IntrinsicMonoFun [PrimType] PrimType + | IntrinsicOverloadedFun [PrimType] [Maybe PrimType] (Maybe PrimType) + | IntrinsicPolyFun [TypeParamBase VName] [StructType] StructType + | IntrinsicType PrimType + | IntrinsicEquality -- Special cased. + +-- | A map of all built-ins. +intrinsics :: M.Map VName Intrinsic +intrinsics = + M.fromList $ + zipWith namify [10 ..] $ + map primFun (M.toList Primitive.primFuns) + ++ [("opaque", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a)] + ++ map unOpFun Primitive.allUnOps + ++ map binOpFun Primitive.allBinOps + ++ map cmpOpFun Primitive.allCmpOps + ++ map convOpFun Primitive.allConvOps + ++ map signFun Primitive.allIntTypes + ++ map unsignFun Primitive.allIntTypes + ++ map + intrinsicType + ( map Signed [minBound .. maxBound] + ++ map Unsigned [minBound .. maxBound] + ++ map FloatType [minBound .. maxBound] + ++ [Bool] + ) + ++ + -- This overrides the ! from Primitive. + [ ( "!", + IntrinsicOverloadedFun + ( map Signed [minBound .. maxBound] + ++ map Unsigned [minBound .. maxBound] + ++ [Bool] + ) + [Nothing] + Nothing + ) + ] + ++ + -- The reason for the loop formulation is to ensure that we + -- get a missing case warning if we forget a case. + mapMaybe mkIntrinsicBinOp [minBound .. maxBound] + ++ [ ( "flatten", + IntrinsicPolyFun + [tp_a] + [Array () Nonunique t_a (rank 2)] + $ Array () Nonunique t_a (rank 1) + ), + ( "unflatten", + IntrinsicPolyFun + [tp_a] + [ Scalar $ Prim $ Signed Int64, + Scalar $ Prim $ Signed Int64, + Array () Nonunique t_a (rank 1) + ] + $ Array () Nonunique t_a (rank 2) + ), + ( "concat", + IntrinsicPolyFun + [tp_a] + [arr_a, arr_a] + uarr_a + ), + ( "rotate", + IntrinsicPolyFun + [tp_a] + [Scalar $ Prim $ Signed Int64, arr_a] + arr_a + ), + ("transpose", IntrinsicPolyFun [tp_a] [arr_2d_a] arr_2d_a), + ( "scatter", + IntrinsicPolyFun + [tp_a] + [ Array () Unique t_a (rank 1), + Array () Nonunique (Prim $ Signed Int64) (rank 1), + Array () Nonunique t_a (rank 1) + ] + $ Array () Unique t_a (rank 1) + ), + ("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", + IntrinsicPolyFun + [tp_a] + [ Scalar $ Prim $ Signed Int64, + uarr_a, + Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), + Scalar t_a, + Array () Nonunique (Prim $ Signed Int64) (rank 1), + arr_a + ] + uarr_a + ), + ("map", IntrinsicPolyFun [tp_a, tp_b] [Scalar t_a `arr` Scalar t_b, arr_a] uarr_b), + ( "reduce", + IntrinsicPolyFun + [tp_a] + [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a] + $ Scalar t_a + ), + ( "reduce_comm", + IntrinsicPolyFun + [tp_a] + [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a] + $ Scalar t_a + ), + ( "scan", + IntrinsicPolyFun + [tp_a] + [Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, arr_a] + uarr_a + ), + ( "partition", + IntrinsicPolyFun + [tp_a] + [ Scalar (Prim $ Signed Int32), + Scalar t_a `arr` Scalar (Prim $ Signed Int64), + arr_a + ] + $ tupleRecord [uarr_a, Array () Unique (Prim $ Signed Int64) (rank 1)] + ), + ( "map_stream", + IntrinsicPolyFun + [tp_a, tp_b] + [Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` arr_kb), arr_a] + uarr_b + ), + ( "map_stream_per", + IntrinsicPolyFun + [tp_a, tp_b] + [Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` arr_kb), arr_a] + uarr_b + ), + ( "reduce_stream", + IntrinsicPolyFun + [tp_a, tp_b] + [ Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b), + Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` Scalar t_b), + arr_a + ] + $ Scalar t_b + ), + ( "reduce_stream_per", + IntrinsicPolyFun + [tp_a, tp_b] + [ Scalar t_b `arr` (Scalar t_b `arr` Scalar t_b), + Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` Scalar t_b), + arr_a + ] + $ Scalar t_b + ), + ("trace", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a), + ("break", IntrinsicPolyFun [tp_a] [Scalar t_a] $ Scalar t_a) + ] + where + tv_a = VName (nameFromString "a") 0 + 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_a = Array () Unique t_a (rank 1) + tp_a = TypeParamType Unlifted tv_a mempty + + tv_b = VName (nameFromString "b") 1 + t_b = TypeVar () Nonunique (typeName tv_b) [] + arr_b = Array () Nonunique t_b (rank 1) + uarr_b = Array () Unique t_b (rank 1) + tp_b = TypeParamType Unlifted tv_b mempty + + arr_a_b = + Array + () + Nonunique + (Record (M.fromList $ zip tupleFieldNames [Scalar t_a, Scalar t_b])) + (rank 1) + t_arr_a_arr_b = Scalar $ Record $ M.fromList $ zip tupleFieldNames [arr_a, arr_b] + + arr x y = Scalar $ Arrow mempty Unnamed x y + + kv = VName (nameFromString "k") 2 + arr_ka = Array () Nonunique t_a (ShapeDecl [NamedDim $ qualName kv]) + arr_kb = Array () Nonunique t_b (ShapeDecl [NamedDim $ qualName kv]) + karr x y = Scalar $ Arrow mempty (Named kv) x y + + namify i (k, v) = (VName (nameFromString k) i, v) + + primFun (name, (ts, t, _)) = + (name, IntrinsicMonoFun (map unPrim ts) $ unPrim t) + + unOpFun bop = (pretty bop, IntrinsicMonoFun [t] t) + where + t = unPrim $ Primitive.unOpType bop + + binOpFun bop = (pretty bop, IntrinsicMonoFun [t, t] t) + where + t = unPrim $ Primitive.binOpType bop + + cmpOpFun bop = (pretty bop, IntrinsicMonoFun [t, t] Bool) + where + t = unPrim $ Primitive.cmpOpType bop + + convOpFun cop = (pretty cop, IntrinsicMonoFun [unPrim ft] $ unPrim tt) + where + (ft, tt) = Primitive.convOpType cop + + signFun t = ("sign_" ++ pretty t, IntrinsicMonoFun [Unsigned t] $ Signed t) + + unsignFun t = ("unsign_" ++ pretty t, IntrinsicMonoFun [Signed t] $ Unsigned t) + + unPrim (Primitive.IntType t) = Signed t + unPrim (Primitive.FloatType t) = FloatType t + unPrim Primitive.Bool = Bool + unPrim Primitive.Cert = Bool + + intrinsicType t = (pretty t, IntrinsicType t) + + anyIntType = + map Signed [minBound .. maxBound] + ++ map Unsigned [minBound .. maxBound] + anyNumberType = + anyIntType + ++ map FloatType [minBound .. maxBound] + anyPrimType = Bool : anyNumberType + + mkIntrinsicBinOp :: BinOp -> Maybe (String, Intrinsic) + mkIntrinsicBinOp op = do + op' <- intrinsicBinOp op + return (pretty op, op') + + binOp ts = Just $ IntrinsicOverloadedFun ts [Nothing, Nothing] Nothing + ordering = Just $ IntrinsicOverloadedFun anyPrimType [Nothing, Nothing] (Just Bool) + + intrinsicBinOp Plus = binOp anyNumberType + intrinsicBinOp Minus = binOp anyNumberType + intrinsicBinOp Pow = binOp anyNumberType + intrinsicBinOp Times = binOp anyNumberType + intrinsicBinOp Divide = binOp anyNumberType + intrinsicBinOp Mod = binOp anyNumberType + intrinsicBinOp Quot = binOp anyIntType + intrinsicBinOp Rem = binOp anyIntType + intrinsicBinOp ShiftR = binOp anyIntType + intrinsicBinOp ShiftL = binOp anyIntType + intrinsicBinOp Band = binOp anyIntType + intrinsicBinOp Xor = binOp anyIntType + intrinsicBinOp Bor = binOp anyIntType + intrinsicBinOp LogAnd = binOp [Bool] + intrinsicBinOp LogOr = binOp [Bool] + intrinsicBinOp Equal = Just IntrinsicEquality + intrinsicBinOp NotEqual = Just IntrinsicEquality + intrinsicBinOp Less = ordering + intrinsicBinOp Leq = ordering + intrinsicBinOp Greater = ordering + intrinsicBinOp Geq = ordering + intrinsicBinOp _ = Nothing + +-- | 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. +maxIntrinsicTag :: Int +maxIntrinsicTag = maxinum $ map baseTag $ M.keys intrinsics + +-- | Create a name with no qualifiers from a name. +qualName :: v -> QualName v +qualName = QualName [] + +-- | Add another qualifier (at the head) to a qualified name. +qualify :: v -> QualName v -> QualName v +qualify k (QualName ks v) = QualName (k : ks) v + +-- | Create a type name name with no qualifiers from a 'VName'. +typeName :: VName -> TypeName +typeName = typeNameFromQualName . qualName + +-- | The modules imported by a Futhark program. +progImports :: ProgBase f vn -> [(String, SrcLoc)] +progImports = concatMap decImports . progDecs + +-- | The modules imported by a single declaration. +decImports :: DecBase f vn -> [(String, SrcLoc)] +decImports (OpenDec x _) = modExpImports x +decImports (ModDec md) = modExpImports $ modExp md +decImports SigDec {} = [] +decImports TypeDec {} = [] +decImports ValDec {} = [] +decImports (LocalDec d _) = decImports d +decImports (ImportDec x _ loc) = [(x, loc)] + +modExpImports :: ModExpBase f vn -> [(String, SrcLoc)] +modExpImports ModVar {} = [] +modExpImports (ModParens p _) = modExpImports p +modExpImports (ModImport f _ loc) = [(f, loc)] +modExpImports (ModDecs ds _) = concatMap decImports ds +modExpImports (ModApply _ me _ _ _) = modExpImports me +modExpImports (ModAscript me _ _ _) = modExpImports me +modExpImports ModLambda {} = [] + +-- | The set of module types used in any exported (non-local) +-- declaration. +progModuleTypes :: Ord vn => ProgBase f vn -> S.Set vn +progModuleTypes = mconcat . map onDec . progDecs + where + onDec (OpenDec x _) = onModExp x + onDec (ModDec md) = + maybe mempty (onSigExp . fst) (modSignature md) <> onModExp (modExp md) + onDec SigDec {} = mempty + onDec TypeDec {} = mempty + onDec ValDec {} = mempty + onDec LocalDec {} = mempty + onDec ImportDec {} = mempty + + onModExp ModVar {} = mempty + onModExp (ModParens p _) = onModExp p + onModExp ModImport {} = mempty + onModExp (ModDecs ds _) = mconcat $ map onDec ds + onModExp (ModApply me1 me2 _ _ _) = onModExp me1 <> onModExp me2 + onModExp (ModAscript me se _ _) = onModExp me <> onSigExp se + onModExp (ModLambda p r me _) = + onModParam p <> maybe mempty (onSigExp . fst) r <> onModExp me + + onModParam = onSigExp . modParamType + + onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v + onSigExp (SigParens e _) = onSigExp e + onSigExp SigSpecs {} = mempty + onSigExp (SigWith e _ _) = onSigExp e + onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2 + +-- | Extract a leading @((name, namespace, file), remainder)@ from a +-- documentation comment string. These are formatted as +-- \`name\`\@namespace[\@file]. Let us hope that this pattern does not occur +-- anywhere else. +identifierReference :: String -> Maybe ((String, String, Maybe FilePath), String) +identifierReference ('`' : s) + | (identifier, '`' : '@' : s') <- break (== '`') s, + (namespace, s'') <- span isAlpha s', + not $ null namespace = + case s'' of + '@' : '"' : s''' + | (file, '"' : s'''') <- span (/= '"') s''' -> + Just ((identifier, namespace, Just file), s'''') + _ -> Just ((identifier, namespace, Nothing), s'') +identifierReference _ = Nothing + +-- | Given an operator name, return the operator that determines its +-- syntactical properties. +leadingOperator :: Name -> BinOp +leadingOperator s = + maybe Backtick snd $ + find ((`isPrefixOf` s') . fst) $ + sortOn (Down . length . fst) $ + zip (map pretty operators) operators + where + s' = nameToString s + operators :: [BinOp] + operators = [minBound .. maxBound :: BinOp] -- | A type with no aliasing information but shape annotations. type UncheckedType = TypeBase (ShapeDecl Name) () diff --git a/src/Language/Futhark/Query.hs b/src/Language/Futhark/Query.hs --- a/src/Language/Futhark/Query.hs +++ b/src/Language/Futhark/Query.hs @@ -1,33 +1,34 @@ {-# LANGUAGE FlexibleContexts #-} + -- | Facilities for answering queries about a program, such as "what -- appears at this source location", or "where is this name bound". -- The intent is that this is used as a building block for IDE-like -- functionality. module Language.Futhark.Query - ( BoundTo(..) - , boundLoc - , AtPos(..) - , atPos - , Pos(..) + ( BoundTo (..), + boundLoc, + AtPos (..), + atPos, + Pos (..), ) - where +where import Control.Monad import Control.Monad.State import Data.List (find) import qualified Data.Map as M -import qualified System.FilePath.Posix as Posix - -import Futhark.Util.Loc (Pos(..), Loc(..)) +import Futhark.Util.Loc (Loc (..), Pos (..)) import Language.Futhark import Language.Futhark.Semantic import Language.Futhark.Traversals +import qualified System.FilePath.Posix as Posix -- | What a name is bound to. -data BoundTo = BoundTerm StructType Loc - | BoundModule Loc - | BoundModuleType Loc - | BoundType Loc +data BoundTo + = BoundTerm StructType Loc + | BoundModule Loc + | BoundModuleType Loc + | BoundType Loc deriving (Eq, Show) data Def = DefBound BoundTo | DefIndirect VName @@ -52,8 +53,8 @@ mconcat $ map (patternDefs . snd) fields patternDefs (PatternParens pat _) = patternDefs pat -patternDefs Wildcard{} = mempty -patternDefs PatternLit{} = mempty +patternDefs Wildcard {} = mempty +patternDefs PatternLit {} = mempty patternDefs (PatternAscription pat _ _) = patternDefs pat patternDefs (PatternConstr _ _ pats _) = @@ -68,50 +69,54 @@ expDefs :: Exp -> Defs expDefs e = execState (astMap mapper e) extra - where mapper = ASTMapper { mapOnExp = onExp - , mapOnName = pure - , mapOnQualName = pure - , mapOnStructType = pure - , mapOnPatternType = pure - } - onExp e' = do - modify (<>expDefs e') - return e' + where + mapper = + ASTMapper + { mapOnExp = onExp, + mapOnName = pure, + mapOnQualName = pure, + mapOnStructType = pure, + mapOnPatternType = pure + } + onExp e' = do + modify (<> expDefs e') + return e' - identDefs (Ident v (Info vt) vloc) = - M.singleton v $ DefBound $ BoundTerm (toStruct vt) $ locOf vloc + identDefs (Ident v (Info vt) vloc) = + M.singleton v $ DefBound $ BoundTerm (toStruct vt) $ locOf vloc - extra = - case e of - LetPat pat _ _ _ _ -> - patternDefs pat - Lambda params _ _ _ _ -> - mconcat (map patternDefs params) - LetFun name (tparams, params, _, Info ret, _) _ _ loc -> - let name_t = foldFunType (map patternStructType params) ret - in M.singleton name (DefBound $ BoundTerm name_t (locOf loc)) <> - mconcat (map typeParamDefs tparams) <> - mconcat (map patternDefs params) - LetWith v _ _ _ _ _ _ -> - identDefs v - DoLoop _ merge _ form _ _ _ -> - patternDefs merge <> - case form of - For i _ -> identDefs i - ForIn pat _ -> patternDefs pat - While{} -> mempty - _ -> - mempty + extra = + case e of + LetPat pat _ _ _ _ -> + patternDefs pat + Lambda params _ _ _ _ -> + mconcat (map patternDefs params) + LetFun name (tparams, params, _, Info ret, _) _ _ loc -> + let name_t = foldFunType (map patternStructType params) ret + in M.singleton name (DefBound $ BoundTerm name_t (locOf loc)) + <> mconcat (map typeParamDefs tparams) + <> mconcat (map patternDefs params) + LetWith v _ _ _ _ _ _ -> + identDefs v + DoLoop _ merge _ form _ _ _ -> + patternDefs merge + <> case form of + For i _ -> identDefs i + ForIn pat _ -> patternDefs pat + While {} -> mempty + _ -> + mempty valBindDefs :: ValBind -> Defs valBindDefs vbind = M.insert (valBindName vbind) (DefBound $ BoundTerm vbind_t (locOf vbind)) $ - mconcat (map typeParamDefs (valBindTypeParams vbind)) <> - mconcat (map patternDefs (valBindParams vbind)) <> - expDefs (valBindBody vbind) - where vbind_t = - foldFunType (map patternStructType (valBindParams vbind)) $ - fst $ unInfo $ valBindRetType vbind + mconcat (map typeParamDefs (valBindTypeParams vbind)) + <> mconcat (map patternDefs (valBindParams vbind)) + <> expDefs (valBindBody vbind) + where + vbind_t = + foldFunType (map patternStructType (valBindParams vbind)) $ + fst $ unInfo $ valBindRetType vbind typeBindDefs :: TypeBind -> Defs typeBindDefs tbind = @@ -119,15 +124,15 @@ modParamDefs :: ModParam -> Defs modParamDefs (ModParam p se _ loc) = - M.singleton p (DefBound $ BoundModule $ locOf loc) <> - sigExpDefs se + M.singleton p (DefBound $ BoundModule $ locOf loc) + <> sigExpDefs se modExpDefs :: ModExp -> Defs -modExpDefs ModVar{} = +modExpDefs ModVar {} = mempty modExpDefs (ModParens me _) = modExpDefs me -modExpDefs ModImport{} = +modExpDefs ModImport {} = mempty modExpDefs (ModDecs decs _) = mconcat $ map decDefs decs @@ -140,26 +145,26 @@ modBindDefs :: ModBind -> Defs modBindDefs mbind = - M.singleton (modName mbind) (DefBound $ BoundModule $ locOf mbind) <> - mconcat (map modParamDefs (modParams mbind)) <> - modExpDefs (modExp mbind) <> - case modSignature mbind of - Nothing -> mempty - Just (_, Info substs) -> - M.map DefIndirect substs + M.singleton (modName mbind) (DefBound $ BoundModule $ locOf mbind) + <> mconcat (map modParamDefs (modParams mbind)) + <> modExpDefs (modExp mbind) + <> case modSignature mbind of + Nothing -> mempty + Just (_, Info substs) -> + M.map DefIndirect substs specDefs :: Spec -> Defs specDefs spec = case spec of ValSpec v tparams tdecl _ loc -> let vdef = DefBound $ BoundTerm (unInfo $ expandedType tdecl) (locOf loc) - in M.insert v vdef $ mconcat (map typeParamDefs tparams) + in M.insert v vdef $ mconcat (map typeParamDefs tparams) TypeAbbrSpec tbind -> typeBindDefs tbind TypeSpec _ v _ _ loc -> M.singleton v $ DefBound $ BoundType $ locOf loc ModSpec v se _ loc -> - M.singleton v (DefBound $ BoundModuleType $ locOf loc) <> - sigExpDefs se + M.singleton v (DefBound $ BoundModuleType $ locOf loc) + <> sigExpDefs se IncludeSpec se _ -> sigExpDefs se sigExpDefs :: SigExp -> Defs @@ -173,8 +178,8 @@ sigBindDefs :: SigBind -> Defs sigBindDefs sbind = - M.singleton (sigName sbind) (DefBound $ BoundModuleType $ locOf sbind) <> - sigExpDefs (sigExp sbind) + M.singleton (sigName sbind) (DefBound $ BoundModuleType $ locOf sbind) + <> sigExpDefs (sigExp sbind) decDefs :: Dec -> Defs decDefs (ValDec vbind) = valBindDefs vbind @@ -183,7 +188,7 @@ decDefs (SigDec mbind) = sigBindDefs mbind decDefs (OpenDec me _) = modExpDefs me decDefs (LocalDec dec _) = decDefs dec -decDefs ImportDec{} = mempty +decDefs ImportDec {} = mempty -- | All bindings of everything in the program. progDefs :: Prog -> Defs @@ -191,9 +196,10 @@ allBindings :: Imports -> M.Map VName BoundTo allBindings imports = M.mapMaybe forward defs - where defs = mconcat $ map (progDefs . fileProg . snd) imports - forward (DefBound x) = Just x - forward (DefIndirect v) = forward =<< M.lookup v defs + where + defs = mconcat $ map (progDefs . fileProg . snd) imports + forward (DefBound x) = Just x + forward (DefIndirect v) = forward =<< M.lookup v defs data RawAtPos = RawAtName (QualName VName) Loc @@ -223,12 +229,13 @@ atPosInTypeExp e1 pos `mplus` atPosInTypeExp e2 pos TESum cs _ -> msum $ map (`atPosInTypeExp` pos) $ concatMap snd cs - where inArg (TypeArgExpDim dim _) = inDim dim - inArg (TypeArgExpType e2) = atPosInTypeExp e2 pos - inDim (DimExpNamed qn loc) = do - guard $ loc `contains` pos - Just $ RawAtName qn $ locOf loc - inDim _ = Nothing + where + inArg (TypeArgExpDim dim _) = inDim dim + inArg (TypeArgExpType e2) = atPosInTypeExp e2 pos + inDim (DimExpNamed qn loc) = do + guard $ loc `contains` pos + Just $ RawAtName qn $ locOf loc + inDim _ = Nothing atPosInPattern :: Pattern -> Pos -> Maybe RawAtPos atPosInPattern (Id vn _ loc) pos = do @@ -244,8 +251,8 @@ atPosInPattern pat pos `mplus` atPosInTypeExp (declaredType tdecl) pos atPosInPattern (PatternConstr _ _ pats _) pos = msum $ map (`atPosInPattern` pos) pats -atPosInPattern PatternLit{} _ = Nothing -atPosInPattern Wildcard{} _ = Nothing +atPosInPattern PatternLit {} _ = Nothing +atPosInPattern Wildcard {} _ = Nothing atPosInExp :: Exp -> Pos -> Maybe RawAtPos atPosInExp (Var qn _ loc) pos = do @@ -253,28 +260,21 @@ Just $ RawAtName qn $ locOf loc atPosInExp (QualParens (qn, loc) _ _) pos | loc `contains` pos = Just $ RawAtName qn $ locOf loc - -- All the value cases are TODO - we need another RawAtPos constructor. -atPosInExp Literal{} _ = Nothing -atPosInExp IntLit{} _ = Nothing -atPosInExp FloatLit{} _ = Nothing - +atPosInExp Literal {} _ = Nothing +atPosInExp IntLit {} _ = Nothing +atPosInExp FloatLit {} _ = Nothing atPosInExp (LetPat pat _ _ _ _) pos | pat `contains` pos = atPosInPattern pat pos - atPosInExp (LetWith a b _ _ _ _ _) pos | a `contains` pos = Just $ RawAtName (qualName $ identName a) (locOf a) | b `contains` pos = Just $ RawAtName (qualName $ identName b) (locOf b) - atPosInExp (DoLoop _ merge _ _ _ _ _) pos | merge `contains` pos = atPosInPattern merge pos - atPosInExp (Ascript _ tdecl _) pos | tdecl `contains` pos = atPosInTypeExp (declaredType tdecl) pos - atPosInExp (Coerce _ tdecl _ _) pos | tdecl `contains` pos = atPosInTypeExp (declaredType tdecl) pos - atPosInExp e pos = do guard $ e `contains` pos -- Use the Either monad for short-circuiting for efficiency reasons. @@ -282,16 +282,19 @@ case astMap mapper e of Left atpos -> Just atpos Right _ -> Nothing - where mapper = ASTMapper { mapOnExp = onExp - , mapOnName = pure - , mapOnQualName = pure - , mapOnStructType = pure - , mapOnPatternType = pure - } - onExp e' = - case atPosInExp e' pos of - Just atpos -> Left atpos - Nothing -> Right e' + where + mapper = + ASTMapper + { mapOnExp = onExp, + mapOnName = pure, + mapOnQualName = pure, + mapOnStructType = pure, + mapOnPatternType = pure + } + onExp e' = + case atPosInExp e' pos of + Just atpos -> Left atpos + Nothing -> Right e' atPosInModExp :: ModExp -> Pos -> Maybe RawAtPos atPosInModExp (ModVar qn loc) pos = do @@ -299,7 +302,7 @@ Just $ RawAtName qn $ locOf loc atPosInModExp (ModParens me _) pos = atPosInModExp me pos -atPosInModExp ModImport{} _ = +atPosInModExp ModImport {} _ = Nothing atPosInModExp (ModDecs decs _) pos = msum $ map (`atPosInDec` pos) decs @@ -315,15 +318,16 @@ case spec of ValSpec _ _ tdecl _ _ -> atPosInTypeExp (declaredType tdecl) pos TypeAbbrSpec tbind -> atPosInTypeBind tbind pos - TypeSpec{} -> Nothing + TypeSpec {} -> Nothing ModSpec _ se _ _ -> atPosInSigExp se pos IncludeSpec se _ -> atPosInSigExp se pos atPosInSigExp :: SigExp -> Pos -> Maybe RawAtPos atPosInSigExp se pos = case se of - SigVar qn _ loc -> do guard $ loc `contains` pos - Just $ RawAtName qn $ locOf loc + SigVar qn _ loc -> do + guard $ loc `contains` pos + Just $ RawAtName qn $ locOf loc SigParens e _ -> atPosInSigExp e pos SigSpecs specs _ -> msum $ map (`atPosInSpec` pos) specs SigWith e _ _ -> atPosInSigExp e pos @@ -331,20 +335,22 @@ atPosInValBind :: ValBind -> Pos -> Maybe RawAtPos atPosInValBind vbind pos = - msum (map (`atPosInPattern` pos) (valBindParams vbind)) `mplus` - atPosInExp (valBindBody vbind) pos `mplus` - join (atPosInTypeExp <$> valBindRetDecl vbind <*> pure pos) + msum (map (`atPosInPattern` pos) (valBindParams vbind)) + `mplus` atPosInExp (valBindBody vbind) pos + `mplus` join (atPosInTypeExp <$> valBindRetDecl vbind <*> pure pos) atPosInTypeBind :: TypeBind -> Pos -> Maybe RawAtPos atPosInTypeBind = atPosInTypeExp . declaredType . typeExp atPosInModBind :: ModBind -> Pos -> Maybe RawAtPos atPosInModBind (ModBind _ params sig e _ _) pos = - msum (map inParam params) `mplus` - atPosInModExp e pos `mplus` - case sig of Nothing -> Nothing - Just (se, _) -> atPosInSigExp se pos - where inParam (ModParam _ se _ _) = atPosInSigExp se pos + msum (map inParam params) + `mplus` atPosInModExp e pos + `mplus` case sig of + Nothing -> Nothing + Just (se, _) -> atPosInSigExp se pos + where + inParam (ModParam _ se _ _) = atPosInSigExp se pos atPosInSigBind :: SigBind -> Pos -> Maybe RawAtPos atPosInSigBind = atPosInSigExp . sigExp @@ -359,7 +365,7 @@ SigDec sbind -> atPosInSigBind sbind pos OpenDec e _ -> atPosInModExp e pos LocalDec dec' _ -> atPosInDec dec' pos - ImportDec{} -> Nothing + ImportDec {} -> Nothing atPosInProg :: Prog -> Pos -> Maybe RawAtPos atPosInProg prog pos = @@ -367,9 +373,12 @@ containingModule :: Imports -> Pos -> Maybe FileModule containingModule imports (Pos file _ _ _) = - snd <$> find ((==file') . fst) imports - where file' = includeToString $ mkInitialImport $ - fst $ Posix.splitExtension file + snd <$> find ((== file') . fst) imports + where + file' = + includeToString $ + mkInitialImport $ + fst $ Posix.splitExtension file -- | Information about what is at the given source location. data AtPos = AtName (QualName VName) (Maybe BoundTo) Loc diff --git a/src/Language/Futhark/Semantic.hs b/src/Language/Futhark/Semantic.hs --- a/src/Language/Futhark/Semantic.hs +++ b/src/Language/Futhark/Semantic.hs @@ -1,39 +1,36 @@ {-# LANGUAGE Safe #-} + -- | Definitions of various semantic objects (*not* the Futhark -- semantics themselves). module Language.Futhark.Semantic - ( ImportName - , mkInitialImport - , mkImportFrom - , includeToFilePath - , includeToString - - , FileModule(..) - , Imports - - , Namespace(..) - , Env(..) - , TySet - , FunSig(..) - , NameMap - , BoundV(..) - , Mod(..) - , TypeBinding(..) - , MTy(..) + ( ImportName, + mkInitialImport, + mkImportFrom, + includeToFilePath, + includeToString, + FileModule (..), + Imports, + Namespace (..), + Env (..), + TySet, + FunSig (..), + NameMap, + BoundV (..), + Mod (..), + TypeBinding (..), + MTy (..), ) where import qualified Data.Map.Strict as M -import qualified System.FilePath.Posix as Posix +import Futhark.Util (dropLast, fromPOSIX, toPOSIX) +import Futhark.Util.Loc +import Futhark.Util.Pretty +import Language.Futhark import qualified System.FilePath as Native - +import qualified System.FilePath.Posix as Posix import Prelude hiding (mod) -import Language.Futhark -import Futhark.Util (dropLast, toPOSIX, fromPOSIX) -import Futhark.Util.Pretty -import Futhark.Util.Loc - -- | Canonical reference to a Futhark code file. Does not include the -- @.fut@ extension. This is most often a path relative to the -- current working directory of the compiler. @@ -55,13 +52,14 @@ mkImportFrom (ImportName includer _) includee | Posix.isAbsolute includee = ImportName includee | otherwise = ImportName $ Posix.normalise $ Posix.joinPath $ includer' ++ includee' - where (dotdots, includee') = span ("../"==) $ Posix.splitPath includee - includer_parts = init $ Posix.splitPath includer - includer' - | length dotdots > length includer_parts = - replicate (length dotdots - length includer_parts) "../" - | otherwise = - dropLast (length dotdots) includer_parts + where + (dotdots, includee') = span ("../" ==) $ Posix.splitPath includee + includer_parts = init $ Posix.splitPath includer + includer' + | length dotdots > length includer_parts = + replicate (length dotdots - length includer_parts) "../" + | otherwise = + dropLast (length dotdots) includer_parts -- | Create a @.fut@ file corresponding to an 'ImportName'. includeToFilePath :: ImportName -> Native.FilePath @@ -74,70 +72,79 @@ -- | The result of type checking some file. Can be passed to further -- invocations of the type checker. -data FileModule = FileModule { fileAbs :: TySet -- ^ Abstract types. - , fileEnv :: Env - , fileProg :: Prog - } +data FileModule = FileModule + { -- | Abstract types. + fileAbs :: TySet, + fileEnv :: Env, + fileProg :: Prog + } -- | A mapping from import names to imports. The ordering is significant. type Imports = [(String, FileModule)] -- | The space inhabited by a name. -data Namespace = Term -- ^ Functions and values. - | Type - | Signature - deriving (Eq, Ord, Show, Enum) +data Namespace + = -- | Functions and values. + Term + | Type + | Signature + deriving (Eq, Ord, Show, Enum) -- | A mapping of abstract types to their liftedness. type TySet = M.Map (QualName VName) Liftedness -- | Representation of a module, which is either a plain environment, -- or a parametric module ("functor" in SML). -data Mod = ModEnv Env - | ModFun FunSig - deriving (Show) +data Mod + = ModEnv Env + | ModFun FunSig + deriving (Show) -- | A parametric functor consists of a set of abstract types, the -- environment of its parameter, and the resulting module type. -data FunSig = FunSig { funSigAbs :: TySet - , funSigMod :: Mod - , funSigMty :: MTy - } - deriving (Show) +data FunSig = FunSig + { funSigAbs :: TySet, + funSigMod :: Mod, + funSigMty :: MTy + } + deriving (Show) -- | Representation of a module type. -data MTy = MTy { mtyAbs :: TySet - -- ^ Abstract types in the module type. - , mtyMod :: Mod - } - deriving (Show) +data MTy = MTy + { -- | Abstract types in the module type. + mtyAbs :: TySet, + mtyMod :: Mod + } + deriving (Show) -- | A binding from a name to its definition as a type. data TypeBinding = TypeAbbr Liftedness [TypeParam] StructType - deriving (Eq, Show) + deriving (Eq, Show) -- | Type parameters, list of parameter types (optinally named), and -- return type. The type parameters are in scope in both parameter -- types and the return type. Non-functional values have only a -- return type. data BoundV = BoundV [TypeParam] StructType - deriving (Show) + deriving (Show) -- | A mapping from names (which always exist in some namespace) to a -- unique (tagged) name. type NameMap = M.Map (Namespace, Name) (QualName VName) -- | Modules produces environment with this representation. -data Env = Env { envVtable :: M.Map VName BoundV - , envTypeTable :: M.Map VName TypeBinding - , envSigTable :: M.Map VName MTy - , envModTable :: M.Map VName Mod - , envNameMap :: NameMap - } deriving (Show) +data Env = Env + { envVtable :: M.Map VName BoundV, + envTypeTable :: M.Map VName TypeBinding, + envSigTable :: M.Map VName MTy, + envModTable :: M.Map VName Mod, + envNameMap :: NameMap + } + deriving (Show) instance Semigroup Env where Env vt1 tt1 st1 mt1 nt1 <> Env vt2 tt2 st2 mt2 nt2 = - Env (vt1<>vt2) (tt1<>tt2) (st1<>st2) (mt1<>mt2) (nt1<>nt2) + Env (vt1 <> vt2) (tt1 <> tt2) (st1 <> st2) (mt1 <> mt2) (nt1 <> nt2) instance Pretty Namespace where ppr Term = text "name" @@ -156,21 +163,27 @@ instance Pretty Env where ppr (Env vtable ttable sigtable modtable _) = - nestedBlock "{" "}" $ stack $ punctuate line $ concat - [map renderTypeBind (M.toList ttable), - map renderValBind (M.toList vtable), - map renderModType (M.toList sigtable), - map renderMod (M.toList modtable)] - where renderTypeBind (name, TypeAbbr l tps tp) = - p l <+> pprName name <> mconcat (map ((text " "<>) . ppr) tps) <> - text " =" <+> ppr tp - where p Lifted = text "type^" - p SizeLifted = text "type~" - p Unlifted = text "type" - renderValBind (name, BoundV tps t) = - text "val" <+> pprName name <> mconcat (map ((text " "<>) . ppr) tps) <> - text " =" <+> ppr t - renderModType (name, _sig) = - text "module type" <+> pprName name - renderMod (name, mod) = - text "module" <+> pprName name <> text " =" <+> ppr mod + nestedBlock "{" "}" $ + stack $ + punctuate line $ + concat + [ map renderTypeBind (M.toList ttable), + map renderValBind (M.toList vtable), + map renderModType (M.toList sigtable), + map renderMod (M.toList modtable) + ] + where + renderTypeBind (name, TypeAbbr l tps tp) = + p l <+> pprName name <> mconcat (map ((text " " <>) . ppr) tps) + <> text " =" <+> ppr tp + where + p Lifted = text "type^" + p SizeLifted = text "type~" + p Unlifted = text "type" + renderValBind (name, BoundV tps t) = + text "val" <+> pprName name <> mconcat (map ((text " " <>) . ppr) tps) + <> text " =" <+> ppr t + renderModType (name, _sig) = + text "module type" <+> pprName name + renderMod (name, mod) = + text "module" <+> pprName name <> text " =" <+> ppr mod 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 @@ -1,1105 +1,1219 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE Safe #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE Strict #-} --- | The Futhark source language AST definition. Many types, such as --- 'ExpBase'@, are parametrised by type and name representation. See --- the @https://futhark.readthedocs.org@ for a language reference, or --- this module may be a little hard to understand. -module Language.Futhark.Syntax - ( - module Language.Futhark.Core - - -- * Types - , Uniqueness(..) - , IntType(..) - , FloatType(..) - , PrimType(..) - , ArrayDim (..) - , DimDecl (..) - , ShapeDecl (..) - , shapeRank - , stripDims - , unifyShapes - , TypeName(..) - , typeNameFromQualName - , qualNameFromTypeName - , TypeBase(..) - , TypeArg(..) - , DimExp(..) - , TypeExp(..) - , TypeArgExp(..) - , PName(..) - , ScalarTypeBase(..) - , PatternType - , StructType - , ValueType - , Diet(..) - , TypeDeclBase (..) - - -- * Values - , IntValue(..) - , FloatValue(..) - , PrimValue(..) - , IsPrimValue(..) - , Value(..) - - -- * Abstract syntax tree - , AttrInfo(..) - , BinOp (..) - , IdentBase (..) - , Inclusiveness(..) - , DimIndexBase(..) - , ExpBase(..) - , FieldBase(..) - , CaseBase(..) - , LoopFormBase (..) - , PatternBase(..) - - -- * Module language - , SpecBase(..) - , SigExpBase(..) - , TypeRefBase(..) - , SigBindBase(..) - , ModExpBase(..) - , ModBindBase(..) - , ModParamBase(..) - - -- * Definitions - , DocComment(..) - , ValBindBase(..) - , EntryPoint(..) - , EntryType(..) - , Liftedness(..) - , TypeBindBase(..) - , TypeParamBase(..) - , typeParamName - , ProgBase(..) - , DecBase(..) - - -- * Miscellaneous - , Showable - , NoInfo(..) - , Info(..) - , Alias(..) - , Aliasing - , QualName(..) - ) - where - -import Control.Applicative -import Control.Monad -import Data.Array -import Data.Bifoldable -import Data.Bifunctor -import Data.Bitraversable -import Data.Foldable -import qualified Data.Map.Strict as M -import Data.Monoid hiding (Sum) -import Data.Ord -import qualified Data.Set as S -import Data.Traversable -import qualified Data.List.NonEmpty as NE -import Prelude - -import Futhark.IR.Primitive (FloatType (..), - FloatValue (..), - IntType (..), IntValue (..)) -import Futhark.Util.Pretty -import Futhark.Util.Loc -import Language.Futhark.Core - --- | Convenience class for deriving 'Show' instances for the AST. -class (Show vn, - Show (f VName), - Show (f (Diet, Maybe VName)), - Show (f String), - Show (f [VName]), - Show (f ([VName], [VName])), - Show (f PatternType), - Show (f (PatternType, [VName])), - Show (f (StructType, [VName])), - Show (f EntryPoint), - Show (f Int), - Show (f StructType), - Show (f (StructType, Maybe VName)), - Show (f (Aliasing, StructType)), - Show (f (M.Map VName VName)), - Show (f Uniqueness)) => Showable f vn where - --- | No information functor. Usually used for placeholder type- or --- aliasing information. -data NoInfo a = NoInfo - deriving (Eq, Ord, Show) - -instance Show vn => Showable NoInfo vn where -instance Functor NoInfo where - fmap _ NoInfo = NoInfo -instance Foldable NoInfo where - foldr _ b NoInfo = b -instance Traversable NoInfo where - traverse _ NoInfo = pure NoInfo - --- | Some information. The dual to 'NoInfo' -newtype Info a = Info { unInfo :: a } - deriving (Eq, Ord, Show) - -instance Show vn => Showable Info vn where -instance Functor Info where - fmap f (Info x) = Info $ f x -instance Foldable Info where - foldr f b (Info x) = f x b -instance Traversable Info where - traverse f (Info x) = Info <$> f x - --- | Low-level primitive types. -data PrimType = Signed IntType - | Unsigned IntType - | FloatType FloatType - | Bool - deriving (Eq, Ord, Show) - --- | Non-array values. -data PrimValue = SignedValue !IntValue - | UnsignedValue !IntValue - | FloatValue !FloatValue - | BoolValue !Bool - deriving (Eq, Ord, Show) - --- | A class for converting ordinary Haskell values to primitive --- Futhark values. -class IsPrimValue v where - primValue :: v -> PrimValue - -instance IsPrimValue Int where - primValue = SignedValue . Int32Value . fromIntegral - -instance IsPrimValue Int8 where - primValue = SignedValue . Int8Value -instance IsPrimValue Int16 where - primValue = SignedValue . Int16Value -instance IsPrimValue Int32 where - primValue = SignedValue . Int32Value -instance IsPrimValue Int64 where - primValue = SignedValue . Int64Value - -instance IsPrimValue Word8 where - primValue = UnsignedValue . Int8Value . fromIntegral -instance IsPrimValue Word16 where - primValue = UnsignedValue . Int16Value . fromIntegral -instance IsPrimValue Word32 where - primValue = UnsignedValue . Int32Value . fromIntegral -instance IsPrimValue Word64 where - primValue = UnsignedValue . Int64Value . fromIntegral - -instance IsPrimValue Float where - primValue = FloatValue . Float32Value - -instance IsPrimValue Double where - primValue = FloatValue . Float64Value - -instance IsPrimValue Bool where - primValue = BoolValue - --- | The payload of an attribute. -data AttrInfo - = AttrAtom Name - | AttrComp Name [AttrInfo] - deriving (Eq, Ord, Show) - --- | A type class for things that can be array dimensions. -class Eq dim => ArrayDim dim where - -- | @unifyDims x y@ combines @x@ and @y@ to contain their maximum - -- common information, and fails if they conflict. - unifyDims :: dim -> dim -> Maybe dim - -instance ArrayDim () where - unifyDims () () = Just () - --- | Declaration of a dimension size. -data DimDecl vn = NamedDim (QualName vn) - -- ^ The size of the dimension is this name, which - -- must be in scope. In a return type, this will - -- give rise to an assertion. - | ConstDim Int - -- ^ The size is a constant. - | AnyDim - -- ^ No dimension declaration. - deriving Show -deriving instance Eq (DimDecl Name) -deriving instance Eq (DimDecl VName) -deriving instance Ord (DimDecl Name) -deriving instance Ord (DimDecl VName) - -instance Functor DimDecl where - fmap = fmapDefault - -instance Foldable DimDecl where - foldMap = foldMapDefault - -instance Traversable DimDecl where - traverse f (NamedDim qn) = NamedDim <$> traverse f qn - traverse _ (ConstDim x) = pure $ ConstDim x - traverse _ AnyDim = pure AnyDim - --- Note that the notion of unifyDims here is intentionally not what we --- use when we do real type unification in the type checker. -instance ArrayDim (DimDecl VName) where - unifyDims AnyDim y = Just y - unifyDims x AnyDim = Just x - unifyDims (NamedDim x) (NamedDim y) | x == y = Just $ NamedDim x - unifyDims (ConstDim x) (ConstDim y) | x == y = Just $ ConstDim x - unifyDims _ _ = Nothing - --- | The size of an array type is a list of its dimension sizes. If --- 'Nothing', that dimension is of a (statically) unknown size. -newtype ShapeDecl dim = ShapeDecl { shapeDims :: [dim] } - deriving (Eq, Ord, Show) - -instance Foldable ShapeDecl where - foldr f x (ShapeDecl ds) = foldr f x ds - -instance Traversable ShapeDecl where - traverse f (ShapeDecl ds) = ShapeDecl <$> traverse f ds - -instance Functor ShapeDecl where - fmap f (ShapeDecl ds) = ShapeDecl $ map f ds - -instance Semigroup (ShapeDecl dim) where - ShapeDecl l1 <> ShapeDecl l2 = ShapeDecl $ l1 ++ l2 - -instance Monoid (ShapeDecl dim) where - mempty = ShapeDecl [] - --- | The number of dimensions contained in a shape. -shapeRank :: ShapeDecl dim -> Int -shapeRank = length . shapeDims - --- | @stripDims n shape@ strips the outer @n@ dimensions from --- @shape@, returning 'Nothing' if this would result in zero or --- fewer dimensions. -stripDims :: Int -> ShapeDecl dim -> Maybe (ShapeDecl dim) -stripDims i (ShapeDecl l) - | i < length l = Just $ ShapeDecl $ drop i l - | otherwise = Nothing - - --- | @unifyShapes x y@ combines @x@ and @y@ to contain their maximum --- common information, and fails if they conflict. -unifyShapes :: ArrayDim dim => ShapeDecl dim -> ShapeDecl dim -> Maybe (ShapeDecl dim) -unifyShapes (ShapeDecl xs) (ShapeDecl ys) = do - guard $ length xs == length ys - ShapeDecl <$> zipWithM unifyDims xs ys - --- | A type name consists of qualifiers (for error messages) and a --- 'VName' (for equality checking). -data TypeName = TypeName { typeQuals :: [VName], typeLeaf :: VName } - deriving (Show) - -instance Eq TypeName where - TypeName _ x == TypeName _ y = x == y - -instance Ord TypeName where - TypeName _ x `compare` TypeName _ y = x `compare` y - --- | Convert a 'QualName' to a 'TypeName'. -typeNameFromQualName :: QualName VName -> TypeName -typeNameFromQualName (QualName qs x) = TypeName qs x - --- | Convert a 'TypeName' to a 'QualName'. -qualNameFromTypeName :: TypeName -> QualName VName -qualNameFromTypeName (TypeName qs x) = QualName qs x - --- | The name (if any) of a function parameter. The 'Eq' and 'Ord' --- instances always compare values of this type equal. -data PName = Named VName | Unnamed - deriving (Show) - -instance Eq PName where - _ == _ = True - -instance Ord PName where - _ <= _ = True - --- | Types that can be elements of arrays. This representation does --- allow arrays of records of functions, which is nonsensical, but it --- convolutes the code too much if we try to statically rule it out. -data ScalarTypeBase dim as - = Prim PrimType - | TypeVar as Uniqueness TypeName [TypeArg dim] - | Record (M.Map Name (TypeBase dim as)) - | Sum (M.Map Name [TypeBase dim as]) - | Arrow as PName (TypeBase dim as) (TypeBase dim as) - -- ^ The aliasing corresponds to the lexical - -- closure of the function. - deriving (Eq, Ord, Show) - -instance Bitraversable ScalarTypeBase where - bitraverse _ _ (Prim t) = pure $ Prim t - bitraverse f g (Record fs) = Record <$> traverse (bitraverse f g) fs - bitraverse f g (TypeVar als u t args) = - TypeVar <$> g als <*> pure u <*> pure t <*> traverse (traverse f) args - bitraverse f g (Arrow als v t1 t2) = - Arrow <$> g als <*> pure v <*> bitraverse f g t1 <*> bitraverse f g t2 - bitraverse f g (Sum cs) = Sum <$> (traverse . traverse) (bitraverse f g) cs - -instance Bifunctor ScalarTypeBase where - bimap = bimapDefault - -instance Bifoldable ScalarTypeBase where - bifoldMap = bifoldMapDefault - --- | An expanded Futhark type is either an array, or something that --- can be an element of an array. When comparing types for equality, --- function parameter names are ignored. This representation permits --- some malformed types (arrays of functions), but importantly rules --- out arrays-of-arrays. -data TypeBase dim as - = Scalar (ScalarTypeBase dim as) - | Array as Uniqueness (ScalarTypeBase dim ()) (ShapeDecl dim) - deriving (Eq, Ord, Show) - -instance Bitraversable TypeBase where - bitraverse f g (Scalar t) = Scalar <$> bitraverse f g t - bitraverse f g (Array a u t shape) = - Array <$> g a <*> pure u <*> bitraverse f pure t <*> traverse f shape - -instance Bifunctor TypeBase where - bimap = bimapDefault - -instance Bifoldable TypeBase where - bifoldMap = bifoldMapDefault - --- | An argument passed to a type constructor. -data TypeArg dim = TypeArgDim dim SrcLoc - | TypeArgType (TypeBase dim ()) SrcLoc - deriving (Eq, Ord, Show) - -instance Traversable TypeArg where - traverse f (TypeArgDim v loc) = TypeArgDim <$> f v <*> pure loc - traverse f (TypeArgType t loc) = TypeArgType <$> bitraverse f pure t <*> pure loc - -instance Functor TypeArg where - fmap = fmapDefault - -instance Foldable TypeArg where - foldMap = foldMapDefault - --- | A variable that is aliased. Can be still in-scope, or have gone --- out of scope and be free. In the latter case, it behaves more like --- an equivalence class. See uniqueness-error18.fut for an example of --- why this is necessary. -data Alias = AliasBound { aliasVar :: VName } - | AliasFree { aliasVar :: VName } - deriving (Eq, Ord, Show) - --- | Aliasing for a type, which is a set of the variables that are --- aliased. -type Aliasing = S.Set Alias - --- | A type with aliasing information and shape annotations, used for --- describing the type patterns and expressions. -type PatternType = TypeBase (DimDecl VName) Aliasing - --- | A "structural" type with shape annotations and no aliasing --- information, used for declarations. -type StructType = TypeBase (DimDecl VName) () - --- | A value type contains full, manifest size information. -type ValueType = TypeBase Int32 () - --- | A dimension declaration expression for use in a 'TypeExp'. -data DimExp vn = DimExpNamed (QualName vn) SrcLoc - -- ^ The size of the dimension is this name, which - -- must be in scope. - | DimExpConst Int SrcLoc - -- ^ The size is a constant. - | DimExpAny - -- ^ No dimension declaration. - deriving Show -deriving instance Eq (DimExp Name) -deriving instance Eq (DimExp VName) -deriving instance Ord (DimExp Name) -deriving instance Ord (DimExp VName) - --- | An unstructured type with type variables and possibly shape --- declarations - this is what the user types in the source program. --- These are used to construct 'TypeBase's in the type checker. -data TypeExp vn = TEVar (QualName vn) SrcLoc - | TETuple [TypeExp vn] SrcLoc - | TERecord [(Name, TypeExp vn)] SrcLoc - | TEArray (TypeExp vn) (DimExp vn) SrcLoc - | TEUnique (TypeExp vn) SrcLoc - | TEApply (TypeExp vn) (TypeArgExp vn) SrcLoc - | TEArrow (Maybe vn) (TypeExp vn) (TypeExp vn) SrcLoc - | TESum [(Name, [TypeExp vn])] SrcLoc - deriving (Show) -deriving instance Eq (TypeExp Name) -deriving instance Eq (TypeExp VName) -deriving instance Ord (TypeExp Name) -deriving instance Ord (TypeExp VName) - -instance Located (TypeExp vn) where - locOf (TEArray _ _ loc) = locOf loc - locOf (TETuple _ loc) = locOf loc - locOf (TERecord _ loc) = locOf loc - locOf (TEVar _ loc) = locOf loc - locOf (TEUnique _ loc) = locOf loc - locOf (TEApply _ _ loc) = locOf loc - locOf (TEArrow _ _ _ loc) = locOf loc - locOf (TESum _ loc) = locOf loc - --- | A type argument expression passed to a type constructor. -data TypeArgExp vn = TypeArgExpDim (DimExp vn) SrcLoc - | TypeArgExpType (TypeExp vn) - deriving (Show) -deriving instance Eq (TypeArgExp Name) -deriving instance Eq (TypeArgExp VName) -deriving instance Ord (TypeArgExp Name) -deriving instance Ord (TypeArgExp VName) - -instance Located (TypeArgExp vn) where - locOf (TypeArgExpDim _ loc) = locOf loc - locOf (TypeArgExpType t) = locOf t - --- | A declaration of the type of something. -data TypeDeclBase f vn = - TypeDecl { declaredType :: TypeExp vn - -- ^ The type declared by the user. - , expandedType :: f StructType - -- ^ The type deduced by the type checker. - } -deriving instance Showable f vn => Show (TypeDeclBase f vn) -deriving instance Eq (TypeDeclBase NoInfo VName) -deriving instance Ord (TypeDeclBase NoInfo VName) - -instance Located (TypeDeclBase f vn) where - locOf = locOf . declaredType - --- | Information about which parts of a value/type are consumed. -data Diet = RecordDiet (M.Map Name Diet) -- ^ Consumes these fields in the record. - | FuncDiet Diet Diet - -- ^ A function that consumes its argument(s) like this. - -- The final 'Diet' should always be 'Observe', as there - -- is no way for a function to consume its return value. - | Consume -- ^ Consumes this value. - | Observe -- ^ Only observes value in this position, does - -- not consume. - deriving (Eq, Show) - --- | Simple Futhark values. Values are fully evaluated and their type --- is always unambiguous. -data Value = PrimValue !PrimValue - | ArrayValue !(Array Int Value) ValueType - -- ^ It is assumed that the array is 0-indexed. The type - -- is the full type. - deriving (Eq, Show) - --- | An identifier consists of its name and the type of the value --- bound to the identifier. -data IdentBase f vn = Ident { identName :: vn - , identType :: f PatternType - , identSrcLoc :: SrcLoc - } -deriving instance Showable f vn => Show (IdentBase f vn) - -instance Eq vn => Eq (IdentBase ty vn) where - x == y = identName x == identName y - -instance Ord vn => Ord (IdentBase ty vn) where - compare = comparing identName - -instance Located (IdentBase ty vn) where - locOf = locOf . identSrcLoc - --- | Default binary operators. -data BinOp = Backtick - -- ^ A pseudo-operator standing in for any normal - -- identifier used as an operator (they all have the - -- same fixity). - -- Binary Ops for Numbers - | Plus - | Minus - | Pow - | Times - | Divide - | Mod - | Quot - | Rem - | ShiftR - | ShiftL - | Band - | Xor - | Bor - | LogAnd - | LogOr - -- Relational Ops for all primitive types at least - | Equal - | NotEqual - | Less - | Leq - | Greater - | Geq - -- Some functional ops. - | PipeRight -- ^ @|>@ - | PipeLeft -- ^ @<|@ - -- Misc - deriving (Eq, Ord, Show, Enum, Bounded) - --- | Whether a bound for an end-point of a 'DimSlice' or a range --- literal is inclusive or exclusive. -data Inclusiveness a = DownToExclusive a - | ToInclusive a -- ^ May be "down to" if step is negative. - | UpToExclusive a - deriving (Eq, Ord, Show) - -instance Located a => Located (Inclusiveness a) where - locOf (DownToExclusive x) = locOf x - locOf (ToInclusive x) = locOf x - locOf (UpToExclusive x) = locOf x - -instance Functor Inclusiveness where - fmap = fmapDefault - -instance Foldable Inclusiveness where - foldMap = foldMapDefault - -instance Traversable Inclusiveness where - traverse f (DownToExclusive x) = DownToExclusive <$> f x - traverse f (ToInclusive x) = ToInclusive <$> f x - traverse f (UpToExclusive x) = UpToExclusive <$> f x - --- | An indexing of a single dimension. -data DimIndexBase f vn = DimFix (ExpBase f vn) - | DimSlice (Maybe (ExpBase f vn)) - (Maybe (ExpBase f vn)) - (Maybe (ExpBase f vn)) -deriving instance Showable f vn => Show (DimIndexBase f vn) -deriving instance Eq (DimIndexBase NoInfo VName) -deriving instance Ord (DimIndexBase NoInfo VName) - --- | A name qualified with a breadcrumb of module accesses. -data QualName vn = QualName { qualQuals :: ![vn] - , qualLeaf :: !vn - } - deriving (Show) - -instance Eq (QualName Name) where - QualName qs1 v1 == QualName qs2 v2 = qs1 == qs2 && v1 == v2 - -instance Eq (QualName VName) where - QualName _ v1 == QualName _ v2 = v1 == v2 - -instance Ord (QualName Name) where - QualName qs1 v1 `compare` QualName qs2 v2 = compare (qs1, v1) (qs2, v2) - -instance Ord (QualName VName) where - QualName _ v1 `compare` QualName _ v2 = compare v1 v2 - -instance Functor QualName where - fmap = fmapDefault - -instance Foldable QualName where - foldMap = foldMapDefault - -instance Traversable QualName where - traverse f (QualName qs v) = QualName <$> traverse f qs <*> f v - --- | The Futhark expression language. --- --- In a value of type @Exp f vn@, annotations are wrapped in the --- functor @f@, and all names are of type @vn@. --- --- This allows us to encode whether or not the expression has been --- type-checked in the Haskell type of the expression. Specifically, --- the parser will produce expressions of type @Exp 'NoInfo' 'Name'@, --- and the type checker will convert these to @Exp 'Info' 'VName'@, in --- which type information is always present and all names are unique. -data ExpBase f vn = - Literal PrimValue SrcLoc - - | IntLit Integer (f PatternType) SrcLoc - -- ^ A polymorphic integral literal. - - | FloatLit Double (f PatternType) SrcLoc - -- ^ A polymorphic decimal literal. - - | StringLit [Word8] SrcLoc - -- ^ A string literal is just a fancy syntax for an array - -- of bytes. - - | Parens (ExpBase f vn) SrcLoc - -- ^ A parenthesized expression. - - | QualParens (QualName vn, SrcLoc) (ExpBase f vn) SrcLoc - - | TupLit [ExpBase f vn] SrcLoc - -- ^ Tuple literals, e.g., @{1+3, {x, y+z}}@. - - | RecordLit [FieldBase f vn] SrcLoc - -- ^ Record literals, e.g. @{x=2,y=3,z}@. - - | ArrayLit [ExpBase f vn] (f PatternType) SrcLoc - -- ^ Array literals, e.g., @[ [1+x, 3], [2, 1+4] ]@. - -- Second arg is the row type of the rows of the array. - - | Range (ExpBase f vn) (Maybe (ExpBase f vn)) (Inclusiveness (ExpBase f vn)) - (f PatternType, f [VName]) SrcLoc - - | Var (QualName vn) (f PatternType) SrcLoc - - | Ascript (ExpBase f vn) (TypeDeclBase f vn) SrcLoc - -- ^ Type ascription: @e : t@. - - | Coerce (ExpBase f vn) (TypeDeclBase f vn) (f PatternType, f [VName]) SrcLoc - -- ^ Size coercion: @e :> t@. - - | LetPat (PatternBase f vn) (ExpBase f vn) (ExpBase f vn) - (f PatternType, f [VName]) SrcLoc - - | LetFun vn ([TypeParamBase vn], - [PatternBase f vn], - Maybe (TypeExp vn), - f StructType, - ExpBase f vn) - (ExpBase f vn) (f PatternType) SrcLoc - - | If (ExpBase f vn) (ExpBase f vn) (ExpBase f vn) (f PatternType, f [VName]) SrcLoc - - | Apply (ExpBase f vn) (ExpBase f vn) - (f (Diet, Maybe VName)) (f PatternType, f [VName]) SrcLoc - -- ^ The @Maybe VName@ is a possible existential size - -- that is instantiated by this argument.. - -- - -- The @[VName]@ are the existential sizes that come - -- into being at this call site. - - | Negate (ExpBase f vn) SrcLoc - -- ^ Numeric negation (ugly special case; Haskell did it first). - - | Lambda [PatternBase f vn] (ExpBase f vn) - (Maybe (TypeExp vn)) (f (Aliasing, StructType)) SrcLoc - - | OpSection (QualName vn) (f PatternType) SrcLoc - -- ^ @+@; first two types are operands, third is result. - | OpSectionLeft (QualName vn) (f PatternType) (ExpBase f vn) - (f (StructType, Maybe VName), f StructType) (f PatternType, f [VName]) SrcLoc - -- ^ @2+@; first type is operand, second is result. - | OpSectionRight (QualName vn) (f PatternType) (ExpBase f vn) - (f StructType, f (StructType, Maybe VName)) (f PatternType) SrcLoc - -- ^ @+2@; first type is operand, second is result. - | ProjectSection [Name] (f PatternType) SrcLoc - -- ^ Field projection as a section: @(.x.y.z)@. - | IndexSection [DimIndexBase f vn] (f PatternType) SrcLoc - -- ^ Array indexing as a section: @(.[i,j])@. - - | DoLoop - [VName] -- Size parameters. - (PatternBase f vn) -- Merge variable pattern. - (ExpBase f vn) -- Initial values of merge variables. - (LoopFormBase f vn) -- Do or while loop. - (ExpBase f vn) -- Loop body. - (f (PatternType, [VName])) -- Return type. - SrcLoc - - | BinOp (QualName vn, SrcLoc) (f PatternType) - (ExpBase f vn, f (StructType, Maybe VName)) - (ExpBase f vn, f (StructType, Maybe VName)) - (f PatternType) (f [VName]) SrcLoc - - | Project Name (ExpBase f vn) (f PatternType) SrcLoc - - -- Primitive array operations - | LetWith (IdentBase f vn) (IdentBase f vn) - [DimIndexBase f vn] (ExpBase f vn) - (ExpBase f vn) (f PatternType) SrcLoc - - | Index (ExpBase f vn) [DimIndexBase f vn] (f PatternType, f [VName]) SrcLoc - - | Update (ExpBase f vn) [DimIndexBase f vn] (ExpBase f vn) SrcLoc - - | RecordUpdate (ExpBase f vn) [Name] (ExpBase f vn) (f PatternType) SrcLoc - - | Assert (ExpBase f vn) (ExpBase f vn) (f String) SrcLoc - -- ^ Fail if the first expression does not return true, - -- and return the value of the second expression if it - -- does. - - | Constr Name [ExpBase f vn] (f PatternType) SrcLoc - -- ^ An n-ary value constructor. - - | Match (ExpBase f vn) (NE.NonEmpty (CaseBase f vn)) - (f PatternType, f [VName]) SrcLoc - -- ^ A match expression. - - | Attr AttrInfo (ExpBase f vn) SrcLoc - -- ^ An attribute applied to the following expression. - -deriving instance Showable f vn => Show (ExpBase f vn) -deriving instance Eq (ExpBase NoInfo VName) -deriving instance Ord (ExpBase NoInfo VName) - -instance Located (ExpBase f vn) where - locOf (Literal _ loc) = locOf loc - locOf (IntLit _ _ loc) = locOf loc - locOf (FloatLit _ _ loc) = locOf loc - locOf (Parens _ loc) = locOf loc - locOf (QualParens _ _ loc) = locOf loc - locOf (TupLit _ pos) = locOf pos - locOf (RecordLit _ pos) = locOf pos - locOf (Project _ _ _ pos) = locOf pos - locOf (ArrayLit _ _ pos) = locOf pos - locOf (StringLit _ loc) = locOf loc - locOf (Range _ _ _ _ pos) = locOf pos - locOf (BinOp _ _ _ _ _ _ loc) = locOf loc - locOf (If _ _ _ _ pos) = locOf pos - locOf (Var _ _ loc) = locOf loc - locOf (Ascript _ _ loc) = locOf loc - locOf (Coerce _ _ _ loc) = locOf loc - locOf (Negate _ pos) = locOf pos - locOf (Apply _ _ _ _ loc) = locOf loc - locOf (LetPat _ _ _ _ loc) = locOf loc - locOf (LetFun _ _ _ _ loc) = locOf loc - locOf (LetWith _ _ _ _ _ _ loc) = locOf loc - locOf (Index _ _ _ loc) = locOf loc - locOf (Update _ _ _ pos) = locOf pos - locOf (RecordUpdate _ _ _ _ pos) = locOf pos - locOf (Lambda _ _ _ _ loc) = locOf loc - locOf (OpSection _ _ loc) = locOf loc - locOf (OpSectionLeft _ _ _ _ _ loc) = locOf loc - locOf (OpSectionRight _ _ _ _ _ loc) = locOf loc - locOf (ProjectSection _ _ loc) = locOf loc - locOf (IndexSection _ _ loc) = locOf loc - locOf (DoLoop _ _ _ _ _ _ loc) = locOf loc - locOf (Assert _ _ _ loc) = locOf loc - locOf (Constr _ _ _ loc) = locOf loc - locOf (Match _ _ _ loc) = locOf loc - locOf (Attr _ _ loc) = locOf loc - --- | An entry in a record literal. -data FieldBase f vn = RecordFieldExplicit Name (ExpBase f vn) SrcLoc - | RecordFieldImplicit vn (f PatternType) SrcLoc -deriving instance Showable f vn => Show (FieldBase f vn) -deriving instance Eq (FieldBase NoInfo VName) -deriving instance Ord (FieldBase NoInfo VName) - -instance Located (FieldBase f vn) where - locOf (RecordFieldExplicit _ _ loc) = locOf loc - locOf (RecordFieldImplicit _ _ loc) = locOf loc - --- | A case in a match expression. -data CaseBase f vn = CasePat (PatternBase f vn) (ExpBase f vn) SrcLoc -deriving instance Showable f vn => Show (CaseBase f vn) -deriving instance Eq (CaseBase NoInfo VName) -deriving instance Ord (CaseBase NoInfo VName) - -instance Located (CaseBase f vn) where - locOf (CasePat _ _ loc) = locOf loc - --- | Whether the loop is a @for@-loop or a @while@-loop. -data LoopFormBase f vn = For (IdentBase f vn) (ExpBase f vn) - | ForIn (PatternBase f vn) (ExpBase f vn) - | While (ExpBase f vn) -deriving instance Showable f vn => Show (LoopFormBase f vn) -deriving instance Eq (LoopFormBase NoInfo VName) -deriving instance Ord (LoopFormBase NoInfo VName) - --- | A pattern as used most places where variables are bound (function --- parameters, @let@ expressions, etc). -data PatternBase f vn = TuplePattern [PatternBase f vn] SrcLoc - | RecordPattern [(Name, PatternBase f vn)] SrcLoc - | PatternParens (PatternBase f vn) SrcLoc - | Id vn (f PatternType) SrcLoc - | Wildcard (f PatternType) SrcLoc -- Nothing, i.e. underscore. - | PatternAscription (PatternBase f vn) (TypeDeclBase f vn) SrcLoc - | PatternLit (ExpBase f vn) (f PatternType) SrcLoc - | PatternConstr Name (f PatternType) [PatternBase f vn] SrcLoc -deriving instance Showable f vn => Show (PatternBase f vn) -deriving instance Eq (PatternBase NoInfo VName) -deriving instance Ord (PatternBase NoInfo VName) - -instance Located (PatternBase f vn) where - locOf (TuplePattern _ loc) = locOf loc - locOf (RecordPattern _ loc) = locOf loc - locOf (PatternParens _ loc) = locOf loc - locOf (Id _ _ loc) = locOf loc - locOf (Wildcard _ loc) = locOf loc - locOf (PatternAscription _ _ loc) = locOf loc - locOf (PatternLit _ _ loc) = locOf loc - locOf (PatternConstr _ _ _ loc) = locOf loc - --- | Documentation strings, including source location. -data DocComment = DocComment String SrcLoc - deriving (Show) - -instance Located DocComment where - locOf (DocComment _ loc) = locOf loc - --- | Part of the type of an entry point. Has an actual type, and --- maybe also an ascribed type expression. -data EntryType = - EntryType { entryType :: StructType - , entryAscribed :: Maybe (TypeExp VName) - } - deriving (Show) - --- | Information about the external interface exposed by an entry --- point. The important thing is that that we remember the original --- source-language types, without desugaring them at all. The --- annoying thing is that we do not require type annotations on entry --- points, so the types can be either ascribed or inferred. -data EntryPoint = - EntryPoint { entryParams :: [EntryType] - , entryReturn :: EntryType - } - deriving (Show) - --- | Function Declarations -data ValBindBase f vn = ValBind - { valBindEntryPoint :: Maybe (f EntryPoint) - -- ^ Just if this function is an entry point. If so, it also - -- contains the externally visible interface. Note that this may not - -- strictly be well-typed after some desugaring operations, as it - -- may refer to abstract types that are no longer in scope. - , valBindName :: vn - , valBindRetDecl :: Maybe (TypeExp vn) - , valBindRetType :: f (StructType, [VName]) - , valBindTypeParams :: [TypeParamBase vn] - , valBindParams :: [PatternBase f vn] - , valBindBody :: ExpBase f vn - , valBindDoc :: Maybe DocComment - , valBindAttrs :: [AttrInfo] - , valBindLocation :: SrcLoc - } -deriving instance Showable f vn => Show (ValBindBase f vn) - -instance Located (ValBindBase f vn) where - locOf = locOf . valBindLocation - --- | Type Declarations -data TypeBindBase f vn = TypeBind { typeAlias :: vn - , typeLiftedness :: Liftedness - , typeParams :: [TypeParamBase vn] - , typeExp :: TypeDeclBase f vn - , typeDoc :: Maybe DocComment - , typeBindLocation :: SrcLoc - } -deriving instance Showable f vn => Show (TypeBindBase f vn) - -instance Located (TypeBindBase f vn) where - locOf = locOf . typeBindLocation - --- | The liftedness of a type parameter. By the @Ord@ instance, --- @Unlifted < SizeLifted < Lifted@. -data Liftedness - = Unlifted - -- ^ May only be instantiated with a zero-order type of (possibly - -- symbolically) known size. - | SizeLifted - -- ^ May only be instantiated with a zero-order type, but the size - -- can be varying. - | Lifted - -- ^ May be instantiated with a functional type. - deriving (Eq, Ord, Show) - --- | A type parameter. -data TypeParamBase vn - = TypeParamDim vn SrcLoc - -- ^ A type parameter that must be a size. - | TypeParamType Liftedness vn SrcLoc - -- ^ A type parameter that must be a type. - deriving (Eq, Ord, Show) - -instance Functor TypeParamBase where - fmap = fmapDefault - -instance Foldable TypeParamBase where - foldMap = foldMapDefault - -instance Traversable TypeParamBase where - traverse f (TypeParamDim v loc) = TypeParamDim <$> f v <*> pure loc - traverse f (TypeParamType l v loc) = TypeParamType l <$> f v <*> pure loc - -instance Located (TypeParamBase vn) where - locOf (TypeParamDim _ loc) = locOf loc - locOf (TypeParamType _ _ loc) = locOf loc - --- | The name of a type parameter. -typeParamName :: TypeParamBase vn -> vn -typeParamName (TypeParamDim v _) = v -typeParamName (TypeParamType _ v _) = v - --- | A spec is a component of a module type. -data SpecBase f vn - = ValSpec { specName :: vn - , specTypeParams :: [TypeParamBase vn] - , specType :: TypeDeclBase f vn - , specDoc :: Maybe DocComment - , specLocation :: SrcLoc - } - | TypeAbbrSpec (TypeBindBase f vn) - | TypeSpec Liftedness vn [TypeParamBase vn] (Maybe DocComment) SrcLoc - -- ^ Abstract type. - | ModSpec vn (SigExpBase f vn) (Maybe DocComment) SrcLoc - | IncludeSpec (SigExpBase f vn) SrcLoc -deriving instance Showable f vn => Show (SpecBase f vn) - -instance Located (SpecBase f vn) where - locOf (ValSpec _ _ _ _ loc) = locOf loc - locOf (TypeAbbrSpec tbind) = locOf tbind - locOf (TypeSpec _ _ _ _ loc) = locOf loc - locOf (ModSpec _ _ _ loc) = locOf loc - locOf (IncludeSpec _ loc) = locOf loc - --- | A module type expression. -data SigExpBase f vn - = SigVar (QualName vn) (f (M.Map VName VName)) SrcLoc - | SigParens (SigExpBase f vn) SrcLoc - | SigSpecs [SpecBase f vn] SrcLoc - | SigWith (SigExpBase f vn) (TypeRefBase f vn) SrcLoc - | SigArrow (Maybe vn) (SigExpBase f vn) (SigExpBase f vn) SrcLoc -deriving instance Showable f vn => Show (SigExpBase f vn) - --- | A type refinement. -data TypeRefBase f vn = TypeRef (QualName vn) [TypeParamBase vn] (TypeDeclBase f vn) SrcLoc -deriving instance Showable f vn => Show (TypeRefBase f vn) - -instance Located (TypeRefBase f vn) where - locOf (TypeRef _ _ _ loc) = locOf loc - -instance Located (SigExpBase f vn) where - locOf (SigVar _ _ loc) = locOf loc - locOf (SigParens _ loc) = locOf loc - locOf (SigSpecs _ loc) = locOf loc - locOf (SigWith _ _ loc) = locOf loc - locOf (SigArrow _ _ _ loc) = locOf loc - --- | Module type binding. -data SigBindBase f vn = SigBind { sigName :: vn - , sigExp :: SigExpBase f vn - , sigDoc :: Maybe DocComment - , sigLoc :: SrcLoc - } -deriving instance Showable f vn => Show (SigBindBase f vn) - -instance Located (SigBindBase f vn) where - locOf = locOf . sigLoc - --- | Module expression. -data ModExpBase f vn - = ModVar (QualName vn) SrcLoc - | ModParens (ModExpBase f vn) SrcLoc - | ModImport FilePath (f FilePath) SrcLoc - -- ^ The contents of another file as a module. - | ModDecs [DecBase f vn] SrcLoc - | ModApply (ModExpBase f vn) (ModExpBase f vn) - (f (M.Map VName VName)) (f (M.Map VName VName)) SrcLoc - -- ^ Functor application. The first mapping is from parameter - -- names to argument names, while the second maps names in the - -- constructed module to the names inside the functor. - | ModAscript (ModExpBase f vn) (SigExpBase f vn) (f (M.Map VName VName)) SrcLoc - | ModLambda (ModParamBase f vn) - (Maybe (SigExpBase f vn, f (M.Map VName VName))) - (ModExpBase f vn) - SrcLoc -deriving instance Showable f vn => Show (ModExpBase f vn) - -instance Located (ModExpBase f vn) where - locOf (ModVar _ loc) = locOf loc - locOf (ModParens _ loc) = locOf loc - locOf (ModImport _ _ loc) = locOf loc - locOf (ModDecs _ loc) = locOf loc - locOf (ModApply _ _ _ _ loc) = locOf loc - locOf (ModAscript _ _ _ loc) = locOf loc - locOf (ModLambda _ _ _ loc) = locOf loc - --- | A module binding. -data ModBindBase f vn = - ModBind { modName :: vn - , modParams :: [ModParamBase f vn] - , modSignature :: Maybe (SigExpBase f vn, f (M.Map VName VName)) - , modExp :: ModExpBase f vn - , modDoc :: Maybe DocComment - , modLocation :: SrcLoc - } -deriving instance Showable f vn => Show (ModBindBase f vn) - -instance Located (ModBindBase f vn) where - locOf = locOf . modLocation - --- | A module parameter. -data ModParamBase f vn = ModParam { modParamName :: vn - , modParamType :: SigExpBase f vn - , modParamAbs :: f [VName] - , modParamLocation :: SrcLoc - } -deriving instance Showable f vn => Show (ModParamBase f vn) - -instance Located (ModParamBase f vn) where - locOf = locOf . modParamLocation - --- | A top-level binding. -data DecBase f vn = ValDec (ValBindBase f vn) - | TypeDec (TypeBindBase f vn) - | SigDec (SigBindBase f vn) - | ModDec (ModBindBase f vn) - | OpenDec (ModExpBase f vn) SrcLoc - | LocalDec (DecBase f vn) SrcLoc - | ImportDec FilePath (f FilePath) SrcLoc -deriving instance Showable f vn => Show (DecBase f vn) - -instance Located (DecBase f vn) where - locOf (ValDec d) = locOf d - locOf (TypeDec d) = locOf d - locOf (SigDec d) = locOf d - locOf (ModDec d) = locOf d - locOf (OpenDec _ loc) = locOf loc - locOf (LocalDec _ loc) = locOf loc - locOf (ImportDec _ _ loc) = locOf loc - --- | The program described by a single Futhark file. May depend on --- other files. -data ProgBase f vn = Prog { progDoc :: Maybe DocComment - , progDecs :: [DecBase f vn] - } -deriving instance Showable f vn => Show (ProgBase f vn) - ---- Some prettyprinting definitions are here because we need them in ---- the Attributes module. - -instance Pretty PrimType where - ppr (Unsigned Int8) = text "u8" - ppr (Unsigned Int16) = text "u16" - ppr (Unsigned Int32) = text "u32" - ppr (Unsigned Int64) = text "u64" - ppr (Signed t) = ppr t - ppr (FloatType t) = ppr t - ppr Bool = text "bool" - -instance Pretty BinOp where - ppr Backtick = text "``" - ppr Plus = text "+" - ppr Minus = text "-" - ppr Pow = text "**" - ppr Times = text "*" - ppr Divide = text "/" - ppr Mod = text "%" - ppr Quot = text "//" - ppr Rem = text "%%" - ppr ShiftR = text ">>" - ppr ShiftL = text "<<" - ppr Band = text "&" - ppr Xor = text "^" - ppr Bor = text "|" - ppr LogAnd = text "&&" - ppr LogOr = text "||" - ppr Equal = text "==" - ppr NotEqual = text "!=" - ppr Less = text "<" - ppr Leq = text "<=" - ppr Greater = text ">" - ppr Geq = text ">=" - ppr PipeLeft = text "<|" +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE Safe #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE Strict #-} + +-- | The Futhark source language AST definition. Many types, such as +-- 'ExpBase'@, are parametrised by type and name representation. See +-- the @https://futhark.readthedocs.org@ for a language reference, or +-- this module may be a little hard to understand. +module Language.Futhark.Syntax + ( module Language.Futhark.Core, + + -- * Types + Uniqueness (..), + IntType (..), + FloatType (..), + PrimType (..), + ArrayDim (..), + DimDecl (..), + ShapeDecl (..), + shapeRank, + stripDims, + unifyShapes, + TypeName (..), + typeNameFromQualName, + qualNameFromTypeName, + TypeBase (..), + TypeArg (..), + DimExp (..), + TypeExp (..), + TypeArgExp (..), + PName (..), + ScalarTypeBase (..), + PatternType, + StructType, + ValueType, + Diet (..), + TypeDeclBase (..), + + -- * Values + IntValue (..), + FloatValue (..), + PrimValue (..), + IsPrimValue (..), + Value (..), + + -- * Abstract syntax tree + AttrInfo (..), + BinOp (..), + IdentBase (..), + Inclusiveness (..), + DimIndexBase (..), + ExpBase (..), + FieldBase (..), + CaseBase (..), + LoopFormBase (..), + PatternBase (..), + + -- * Module language + SpecBase (..), + SigExpBase (..), + TypeRefBase (..), + SigBindBase (..), + ModExpBase (..), + ModBindBase (..), + ModParamBase (..), + + -- * Definitions + DocComment (..), + ValBindBase (..), + EntryPoint (..), + EntryType (..), + Liftedness (..), + TypeBindBase (..), + TypeParamBase (..), + typeParamName, + ProgBase (..), + DecBase (..), + + -- * Miscellaneous + Showable, + NoInfo (..), + Info (..), + Alias (..), + Aliasing, + QualName (..), + ) +where + +import Control.Applicative +import Control.Monad +import Data.Array +import Data.Bifoldable +import Data.Bifunctor +import Data.Bitraversable +import Data.Foldable +import qualified Data.List.NonEmpty as NE +import qualified Data.Map.Strict as M +import Data.Monoid hiding (Sum) +import Data.Ord +import qualified Data.Set as S +import Data.Traversable +import Futhark.IR.Primitive + ( FloatType (..), + FloatValue (..), + IntType (..), + IntValue (..), + ) +import Futhark.Util.Loc +import Futhark.Util.Pretty +import Language.Futhark.Core +import Prelude + +-- | Convenience class for deriving 'Show' instances for the AST. +class + ( Show vn, + Show (f VName), + Show (f (Diet, Maybe VName)), + Show (f String), + Show (f [VName]), + Show (f ([VName], [VName])), + Show (f PatternType), + Show (f (PatternType, [VName])), + Show (f (StructType, [VName])), + Show (f EntryPoint), + Show (f Int), + Show (f StructType), + Show (f (StructType, Maybe VName)), + Show (f (Aliasing, StructType)), + Show (f (M.Map VName VName)), + Show (f Uniqueness) + ) => + Showable f vn + +-- | No information functor. Usually used for placeholder type- or +-- aliasing information. +data NoInfo a = NoInfo + deriving (Eq, Ord, Show) + +instance Show vn => Showable NoInfo vn + +instance Functor NoInfo where + fmap _ NoInfo = NoInfo + +instance Foldable NoInfo where + foldr _ b NoInfo = b + +instance Traversable NoInfo where + traverse _ NoInfo = pure NoInfo + +-- | Some information. The dual to 'NoInfo' +newtype Info a = Info {unInfo :: a} + deriving (Eq, Ord, Show) + +instance Show vn => Showable Info vn + +instance Functor Info where + fmap f (Info x) = Info $ f x + +instance Foldable Info where + foldr f b (Info x) = f x b + +instance Traversable Info where + traverse f (Info x) = Info <$> f x + +-- | Low-level primitive types. +data PrimType + = Signed IntType + | Unsigned IntType + | FloatType FloatType + | Bool + deriving (Eq, Ord, Show) + +-- | Non-array values. +data PrimValue + = SignedValue !IntValue + | UnsignedValue !IntValue + | FloatValue !FloatValue + | BoolValue !Bool + deriving (Eq, Ord, Show) + +-- | A class for converting ordinary Haskell values to primitive +-- Futhark values. +class IsPrimValue v where + primValue :: v -> PrimValue + +instance IsPrimValue Int where + primValue = SignedValue . Int32Value . fromIntegral + +instance IsPrimValue Int8 where + primValue = SignedValue . Int8Value + +instance IsPrimValue Int16 where + primValue = SignedValue . Int16Value + +instance IsPrimValue Int32 where + primValue = SignedValue . Int32Value + +instance IsPrimValue Int64 where + primValue = SignedValue . Int64Value + +instance IsPrimValue Word8 where + primValue = UnsignedValue . Int8Value . fromIntegral + +instance IsPrimValue Word16 where + primValue = UnsignedValue . Int16Value . fromIntegral + +instance IsPrimValue Word32 where + primValue = UnsignedValue . Int32Value . fromIntegral + +instance IsPrimValue Word64 where + primValue = UnsignedValue . Int64Value . fromIntegral + +instance IsPrimValue Float where + primValue = FloatValue . Float32Value + +instance IsPrimValue Double where + primValue = FloatValue . Float64Value + +instance IsPrimValue Bool where + primValue = BoolValue + +-- | The payload of an attribute. +data AttrInfo + = AttrAtom Name + | AttrComp Name [AttrInfo] + deriving (Eq, Ord, Show) + +-- | A type class for things that can be array dimensions. +class Eq dim => ArrayDim dim where + -- | @unifyDims x y@ combines @x@ and @y@ to contain their maximum + -- common information, and fails if they conflict. + unifyDims :: dim -> dim -> Maybe dim + +instance ArrayDim () where + unifyDims () () = Just () + +-- | Declaration of a dimension size. +data DimDecl vn + = -- | The size of the dimension is this name, which + -- must be in scope. In a return type, this will + -- give rise to an assertion. + NamedDim (QualName vn) + | -- | The size is a constant. + ConstDim Int + | -- | No dimension declaration. + AnyDim + deriving (Show) + +deriving instance Eq (DimDecl Name) + +deriving instance Eq (DimDecl VName) + +deriving instance Ord (DimDecl Name) + +deriving instance Ord (DimDecl VName) + +instance Functor DimDecl where + fmap = fmapDefault + +instance Foldable DimDecl where + foldMap = foldMapDefault + +instance Traversable DimDecl where + traverse f (NamedDim qn) = NamedDim <$> traverse f qn + traverse _ (ConstDim x) = pure $ ConstDim x + traverse _ AnyDim = pure AnyDim + +-- Note that the notion of unifyDims here is intentionally not what we +-- use when we do real type unification in the type checker. +instance ArrayDim (DimDecl VName) where + unifyDims AnyDim y = Just y + unifyDims x AnyDim = Just x + unifyDims (NamedDim x) (NamedDim y) | x == y = Just $ NamedDim x + unifyDims (ConstDim x) (ConstDim y) | x == y = Just $ ConstDim x + unifyDims _ _ = Nothing + +-- | The size of an array type is a list of its dimension sizes. If +-- 'Nothing', that dimension is of a (statically) unknown size. +newtype ShapeDecl dim = ShapeDecl {shapeDims :: [dim]} + deriving (Eq, Ord, Show) + +instance Foldable ShapeDecl where + foldr f x (ShapeDecl ds) = foldr f x ds + +instance Traversable ShapeDecl where + traverse f (ShapeDecl ds) = ShapeDecl <$> traverse f ds + +instance Functor ShapeDecl where + fmap f (ShapeDecl ds) = ShapeDecl $ map f ds + +instance Semigroup (ShapeDecl dim) where + ShapeDecl l1 <> ShapeDecl l2 = ShapeDecl $ l1 ++ l2 + +instance Monoid (ShapeDecl dim) where + mempty = ShapeDecl [] + +-- | The number of dimensions contained in a shape. +shapeRank :: ShapeDecl dim -> Int +shapeRank = length . shapeDims + +-- | @stripDims n shape@ strips the outer @n@ dimensions from +-- @shape@, returning 'Nothing' if this would result in zero or +-- fewer dimensions. +stripDims :: Int -> ShapeDecl dim -> Maybe (ShapeDecl dim) +stripDims i (ShapeDecl l) + | i < length l = Just $ ShapeDecl $ drop i l + | otherwise = Nothing + +-- | @unifyShapes x y@ combines @x@ and @y@ to contain their maximum +-- common information, and fails if they conflict. +unifyShapes :: ArrayDim dim => ShapeDecl dim -> ShapeDecl dim -> Maybe (ShapeDecl dim) +unifyShapes (ShapeDecl xs) (ShapeDecl ys) = do + guard $ length xs == length ys + ShapeDecl <$> zipWithM unifyDims xs ys + +-- | A type name consists of qualifiers (for error messages) and a +-- 'VName' (for equality checking). +data TypeName = TypeName {typeQuals :: [VName], typeLeaf :: VName} + deriving (Show) + +instance Eq TypeName where + TypeName _ x == TypeName _ y = x == y + +instance Ord TypeName where + TypeName _ x `compare` TypeName _ y = x `compare` y + +-- | Convert a 'QualName' to a 'TypeName'. +typeNameFromQualName :: QualName VName -> TypeName +typeNameFromQualName (QualName qs x) = TypeName qs x + +-- | Convert a 'TypeName' to a 'QualName'. +qualNameFromTypeName :: TypeName -> QualName VName +qualNameFromTypeName (TypeName qs x) = QualName qs x + +-- | The name (if any) of a function parameter. The 'Eq' and 'Ord' +-- instances always compare values of this type equal. +data PName = Named VName | Unnamed + deriving (Show) + +instance Eq PName where + _ == _ = True + +instance Ord PName where + _ <= _ = True + +-- | Types that can be elements of arrays. This representation does +-- allow arrays of records of functions, which is nonsensical, but it +-- convolutes the code too much if we try to statically rule it out. +data ScalarTypeBase dim as + = Prim PrimType + | TypeVar as Uniqueness TypeName [TypeArg dim] + | Record (M.Map Name (TypeBase dim as)) + | Sum (M.Map Name [TypeBase dim as]) + | -- | The aliasing corresponds to the lexical + -- closure of the function. + Arrow as PName (TypeBase dim as) (TypeBase dim as) + deriving (Eq, Ord, Show) + +instance Bitraversable ScalarTypeBase where + bitraverse _ _ (Prim t) = pure $ Prim t + bitraverse f g (Record fs) = Record <$> traverse (bitraverse f g) fs + bitraverse f g (TypeVar als u t args) = + TypeVar <$> g als <*> pure u <*> pure t <*> traverse (traverse f) args + bitraverse f g (Arrow als v t1 t2) = + Arrow <$> g als <*> pure v <*> bitraverse f g t1 <*> bitraverse f g t2 + bitraverse f g (Sum cs) = Sum <$> (traverse . traverse) (bitraverse f g) cs + +instance Bifunctor ScalarTypeBase where + bimap = bimapDefault + +instance Bifoldable ScalarTypeBase where + bifoldMap = bifoldMapDefault + +-- | An expanded Futhark type is either an array, or something that +-- can be an element of an array. When comparing types for equality, +-- function parameter names are ignored. This representation permits +-- some malformed types (arrays of functions), but importantly rules +-- out arrays-of-arrays. +data TypeBase dim as + = Scalar (ScalarTypeBase dim as) + | Array as Uniqueness (ScalarTypeBase dim ()) (ShapeDecl dim) + deriving (Eq, Ord, Show) + +instance Bitraversable TypeBase where + bitraverse f g (Scalar t) = Scalar <$> bitraverse f g t + bitraverse f g (Array a u t shape) = + Array <$> g a <*> pure u <*> bitraverse f pure t <*> traverse f shape + +instance Bifunctor TypeBase where + bimap = bimapDefault + +instance Bifoldable TypeBase where + bifoldMap = bifoldMapDefault + +-- | An argument passed to a type constructor. +data TypeArg dim + = TypeArgDim dim SrcLoc + | TypeArgType (TypeBase dim ()) SrcLoc + deriving (Eq, Ord, Show) + +instance Traversable TypeArg where + traverse f (TypeArgDim v loc) = TypeArgDim <$> f v <*> pure loc + traverse f (TypeArgType t loc) = TypeArgType <$> bitraverse f pure t <*> pure loc + +instance Functor TypeArg where + fmap = fmapDefault + +instance Foldable TypeArg where + foldMap = foldMapDefault + +-- | A variable that is aliased. Can be still in-scope, or have gone +-- out of scope and be free. In the latter case, it behaves more like +-- an equivalence class. See uniqueness-error18.fut for an example of +-- why this is necessary. +data Alias + = AliasBound {aliasVar :: VName} + | AliasFree {aliasVar :: VName} + deriving (Eq, Ord, Show) + +-- | Aliasing for a type, which is a set of the variables that are +-- aliased. +type Aliasing = S.Set Alias + +-- | A type with aliasing information and shape annotations, used for +-- describing the type patterns and expressions. +type PatternType = TypeBase (DimDecl VName) Aliasing + +-- | A "structural" type with shape annotations and no aliasing +-- information, used for declarations. +type StructType = TypeBase (DimDecl VName) () + +-- | A value type contains full, manifest size information. +type ValueType = TypeBase Int64 () + +-- | A dimension declaration expression for use in a 'TypeExp'. +data DimExp vn + = -- | The size of the dimension is this name, which + -- must be in scope. + DimExpNamed (QualName vn) SrcLoc + | -- | The size is a constant. + DimExpConst Int SrcLoc + | -- | No dimension declaration. + DimExpAny + deriving (Show) + +deriving instance Eq (DimExp Name) + +deriving instance Eq (DimExp VName) + +deriving instance Ord (DimExp Name) + +deriving instance Ord (DimExp VName) + +-- | An unstructured type with type variables and possibly shape +-- declarations - this is what the user types in the source program. +-- These are used to construct 'TypeBase's in the type checker. +data TypeExp vn + = TEVar (QualName vn) SrcLoc + | TETuple [TypeExp vn] SrcLoc + | TERecord [(Name, TypeExp vn)] SrcLoc + | TEArray (TypeExp vn) (DimExp vn) SrcLoc + | TEUnique (TypeExp vn) SrcLoc + | TEApply (TypeExp vn) (TypeArgExp vn) SrcLoc + | TEArrow (Maybe vn) (TypeExp vn) (TypeExp vn) SrcLoc + | TESum [(Name, [TypeExp vn])] SrcLoc + deriving (Show) + +deriving instance Eq (TypeExp Name) + +deriving instance Eq (TypeExp VName) + +deriving instance Ord (TypeExp Name) + +deriving instance Ord (TypeExp VName) + +instance Located (TypeExp vn) where + locOf (TEArray _ _ loc) = locOf loc + locOf (TETuple _ loc) = locOf loc + locOf (TERecord _ loc) = locOf loc + locOf (TEVar _ loc) = locOf loc + locOf (TEUnique _ loc) = locOf loc + locOf (TEApply _ _ loc) = locOf loc + locOf (TEArrow _ _ _ loc) = locOf loc + locOf (TESum _ loc) = locOf loc + +-- | A type argument expression passed to a type constructor. +data TypeArgExp vn + = TypeArgExpDim (DimExp vn) SrcLoc + | TypeArgExpType (TypeExp vn) + deriving (Show) + +deriving instance Eq (TypeArgExp Name) + +deriving instance Eq (TypeArgExp VName) + +deriving instance Ord (TypeArgExp Name) + +deriving instance Ord (TypeArgExp VName) + +instance Located (TypeArgExp vn) where + locOf (TypeArgExpDim _ loc) = locOf loc + locOf (TypeArgExpType t) = locOf t + +-- | A declaration of the type of something. +data TypeDeclBase f vn = TypeDecl + { -- | The type declared by the user. + declaredType :: TypeExp vn, + -- | The type deduced by the type checker. + expandedType :: f StructType + } + +deriving instance Showable f vn => Show (TypeDeclBase f vn) + +deriving instance Eq (TypeDeclBase NoInfo VName) + +deriving instance Ord (TypeDeclBase NoInfo VName) + +instance Located (TypeDeclBase f vn) where + locOf = locOf . declaredType + +-- | Information about which parts of a value/type are consumed. +data Diet + = -- | Consumes these fields in the record. + RecordDiet (M.Map Name Diet) + | -- | A function that consumes its argument(s) like this. + -- The final 'Diet' should always be 'Observe', as there + -- is no way for a function to consume its return value. + FuncDiet Diet Diet + | -- | Consumes this value. + Consume + | -- | Only observes value in this position, does + -- not consume. + Observe + deriving (Eq, Show) + +-- | Simple Futhark values. Values are fully evaluated and their type +-- is always unambiguous. +data Value + = PrimValue !PrimValue + | -- | It is assumed that the array is 0-indexed. The type + -- is the full type. + ArrayValue !(Array Int Value) ValueType + deriving (Eq, Show) + +-- | An identifier consists of its name and the type of the value +-- bound to the identifier. +data IdentBase f vn = Ident + { identName :: vn, + identType :: f PatternType, + identSrcLoc :: SrcLoc + } + +deriving instance Showable f vn => Show (IdentBase f vn) + +instance Eq vn => Eq (IdentBase ty vn) where + x == y = identName x == identName y + +instance Ord vn => Ord (IdentBase ty vn) where + compare = comparing identName + +instance Located (IdentBase ty vn) where + locOf = locOf . identSrcLoc + +-- | Default binary operators. +data BinOp + = -- | A pseudo-operator standing in for any normal + -- identifier used as an operator (they all have the + -- same fixity). + -- Binary Ops for Numbers + Backtick + | Plus + | Minus + | Pow + | Times + | Divide + | Mod + | Quot + | Rem + | ShiftR + | ShiftL + | Band + | Xor + | Bor + | LogAnd + | LogOr + | -- Relational Ops for all primitive types at least + Equal + | NotEqual + | Less + | Leq + | Greater + | Geq + | -- Some functional ops. + + -- | @|>@ + PipeRight + | -- | @<|@ + -- Misc + PipeLeft + deriving (Eq, Ord, Show, Enum, Bounded) + +-- | Whether a bound for an end-point of a 'DimSlice' or a range +-- literal is inclusive or exclusive. +data Inclusiveness a + = DownToExclusive a + | -- | May be "down to" if step is negative. + ToInclusive a + | UpToExclusive a + deriving (Eq, Ord, Show) + +instance Located a => Located (Inclusiveness a) where + locOf (DownToExclusive x) = locOf x + locOf (ToInclusive x) = locOf x + locOf (UpToExclusive x) = locOf x + +instance Functor Inclusiveness where + fmap = fmapDefault + +instance Foldable Inclusiveness where + foldMap = foldMapDefault + +instance Traversable Inclusiveness where + traverse f (DownToExclusive x) = DownToExclusive <$> f x + traverse f (ToInclusive x) = ToInclusive <$> f x + traverse f (UpToExclusive x) = UpToExclusive <$> f x + +-- | An indexing of a single dimension. +data DimIndexBase f vn + = DimFix (ExpBase f vn) + | DimSlice + (Maybe (ExpBase f vn)) + (Maybe (ExpBase f vn)) + (Maybe (ExpBase f vn)) + +deriving instance Showable f vn => Show (DimIndexBase f vn) + +deriving instance Eq (DimIndexBase NoInfo VName) + +deriving instance Ord (DimIndexBase NoInfo VName) + +-- | A name qualified with a breadcrumb of module accesses. +data QualName vn = QualName + { qualQuals :: ![vn], + qualLeaf :: !vn + } + deriving (Show) + +instance Eq (QualName Name) where + QualName qs1 v1 == QualName qs2 v2 = qs1 == qs2 && v1 == v2 + +instance Eq (QualName VName) where + QualName _ v1 == QualName _ v2 = v1 == v2 + +instance Ord (QualName Name) where + QualName qs1 v1 `compare` QualName qs2 v2 = compare (qs1, v1) (qs2, v2) + +instance Ord (QualName VName) where + QualName _ v1 `compare` QualName _ v2 = compare v1 v2 + +instance Functor QualName where + fmap = fmapDefault + +instance Foldable QualName where + foldMap = foldMapDefault + +instance Traversable QualName where + traverse f (QualName qs v) = QualName <$> traverse f qs <*> f v + +-- | The Futhark expression language. +-- +-- In a value of type @Exp f vn@, annotations are wrapped in the +-- functor @f@, and all names are of type @vn@. +-- +-- This allows us to encode whether or not the expression has been +-- type-checked in the Haskell type of the expression. Specifically, +-- the parser will produce expressions of type @Exp 'NoInfo' 'Name'@, +-- and the type checker will convert these to @Exp 'Info' 'VName'@, in +-- which type information is always present and all names are unique. +data ExpBase f vn + = Literal PrimValue SrcLoc + | -- | A polymorphic integral literal. + IntLit Integer (f PatternType) SrcLoc + | -- | A polymorphic decimal literal. + FloatLit Double (f PatternType) SrcLoc + | -- | A string literal is just a fancy syntax for an array + -- of bytes. + StringLit [Word8] SrcLoc + | -- | A parenthesized expression. + Parens (ExpBase f vn) SrcLoc + | QualParens (QualName vn, SrcLoc) (ExpBase f vn) SrcLoc + | -- | Tuple literals, e.g., @{1+3, {x, y+z}}@. + TupLit [ExpBase f vn] SrcLoc + | -- | Record literals, e.g. @{x=2,y=3,z}@. + RecordLit [FieldBase f vn] SrcLoc + | -- | Array literals, e.g., @[ [1+x, 3], [2, 1+4] ]@. + -- Second arg is the row type of the rows of the array. + ArrayLit [ExpBase f vn] (f PatternType) SrcLoc + | Range + (ExpBase f vn) + (Maybe (ExpBase f vn)) + (Inclusiveness (ExpBase f vn)) + (f PatternType, f [VName]) + SrcLoc + | Var (QualName vn) (f PatternType) SrcLoc + | -- | Type ascription: @e : t@. + Ascript (ExpBase f vn) (TypeDeclBase f vn) SrcLoc + | -- | Size coercion: @e :> t@. + Coerce (ExpBase f vn) (TypeDeclBase f vn) (f PatternType, f [VName]) SrcLoc + | LetPat + (PatternBase f vn) + (ExpBase f vn) + (ExpBase f vn) + (f PatternType, f [VName]) + SrcLoc + | LetFun + vn + ( [TypeParamBase vn], + [PatternBase f vn], + Maybe (TypeExp vn), + f StructType, + ExpBase f vn + ) + (ExpBase f vn) + (f PatternType) + SrcLoc + | If (ExpBase f vn) (ExpBase f vn) (ExpBase f vn) (f PatternType, f [VName]) SrcLoc + | -- | The @Maybe VName@ is a possible existential size + -- that is instantiated by this argument.. + -- + -- The @[VName]@ are the existential sizes that come + -- into being at this call site. + Apply + (ExpBase f vn) + (ExpBase f vn) + (f (Diet, Maybe VName)) + (f PatternType, f [VName]) + SrcLoc + | -- | Numeric negation (ugly special case; Haskell did it first). + Negate (ExpBase f vn) SrcLoc + | Lambda + [PatternBase f vn] + (ExpBase f vn) + (Maybe (TypeExp vn)) + (f (Aliasing, StructType)) + SrcLoc + | -- | @+@; first two types are operands, third is result. + OpSection (QualName vn) (f PatternType) SrcLoc + | -- | @2+@; first type is operand, second is result. + OpSectionLeft + (QualName vn) + (f PatternType) + (ExpBase f vn) + (f (StructType, Maybe VName), f StructType) + (f PatternType, f [VName]) + SrcLoc + | -- | @+2@; first type is operand, second is result. + OpSectionRight + (QualName vn) + (f PatternType) + (ExpBase f vn) + (f StructType, f (StructType, Maybe VName)) + (f PatternType) + SrcLoc + | -- | Field projection as a section: @(.x.y.z)@. + ProjectSection [Name] (f PatternType) SrcLoc + | -- | Array indexing as a section: @(.[i,j])@. + IndexSection [DimIndexBase f vn] (f PatternType) SrcLoc + | DoLoop + [VName] -- Size parameters. + (PatternBase f vn) -- Merge variable pattern. + (ExpBase f vn) -- Initial values of merge variables. + (LoopFormBase f vn) -- Do or while loop. + (ExpBase f vn) -- Loop body. + (f (PatternType, [VName])) -- Return type. + SrcLoc + | BinOp + (QualName vn, SrcLoc) + (f PatternType) + (ExpBase f vn, f (StructType, Maybe VName)) + (ExpBase f vn, f (StructType, Maybe VName)) + (f PatternType) + (f [VName]) + SrcLoc + | Project Name (ExpBase f vn) (f PatternType) SrcLoc + | -- Primitive array operations + LetWith + (IdentBase f vn) + (IdentBase f vn) + [DimIndexBase f vn] + (ExpBase f vn) + (ExpBase f vn) + (f PatternType) + SrcLoc + | Index (ExpBase f vn) [DimIndexBase f vn] (f PatternType, f [VName]) SrcLoc + | Update (ExpBase f vn) [DimIndexBase f vn] (ExpBase f vn) SrcLoc + | RecordUpdate (ExpBase f vn) [Name] (ExpBase f vn) (f PatternType) SrcLoc + | -- | Fail if the first expression does not return true, + -- and return the value of the second expression if it + -- does. + Assert (ExpBase f vn) (ExpBase f vn) (f String) SrcLoc + | -- | An n-ary value constructor. + Constr Name [ExpBase f vn] (f PatternType) SrcLoc + | -- | A match expression. + Match + (ExpBase f vn) + (NE.NonEmpty (CaseBase f vn)) + (f PatternType, f [VName]) + SrcLoc + | -- | An attribute applied to the following expression. + Attr AttrInfo (ExpBase f vn) SrcLoc + +deriving instance Showable f vn => Show (ExpBase f vn) + +deriving instance Eq (ExpBase NoInfo VName) + +deriving instance Ord (ExpBase NoInfo VName) + +instance Located (ExpBase f vn) where + locOf (Literal _ loc) = locOf loc + locOf (IntLit _ _ loc) = locOf loc + locOf (FloatLit _ _ loc) = locOf loc + locOf (Parens _ loc) = locOf loc + locOf (QualParens _ _ loc) = locOf loc + locOf (TupLit _ pos) = locOf pos + locOf (RecordLit _ pos) = locOf pos + locOf (Project _ _ _ pos) = locOf pos + locOf (ArrayLit _ _ pos) = locOf pos + locOf (StringLit _ loc) = locOf loc + locOf (Range _ _ _ _ pos) = locOf pos + locOf (BinOp _ _ _ _ _ _ loc) = locOf loc + locOf (If _ _ _ _ pos) = locOf pos + locOf (Var _ _ loc) = locOf loc + locOf (Ascript _ _ loc) = locOf loc + locOf (Coerce _ _ _ loc) = locOf loc + locOf (Negate _ pos) = locOf pos + locOf (Apply _ _ _ _ loc) = locOf loc + locOf (LetPat _ _ _ _ loc) = locOf loc + locOf (LetFun _ _ _ _ loc) = locOf loc + locOf (LetWith _ _ _ _ _ _ loc) = locOf loc + locOf (Index _ _ _ loc) = locOf loc + locOf (Update _ _ _ pos) = locOf pos + locOf (RecordUpdate _ _ _ _ pos) = locOf pos + locOf (Lambda _ _ _ _ loc) = locOf loc + locOf (OpSection _ _ loc) = locOf loc + locOf (OpSectionLeft _ _ _ _ _ loc) = locOf loc + locOf (OpSectionRight _ _ _ _ _ loc) = locOf loc + locOf (ProjectSection _ _ loc) = locOf loc + locOf (IndexSection _ _ loc) = locOf loc + locOf (DoLoop _ _ _ _ _ _ loc) = locOf loc + locOf (Assert _ _ _ loc) = locOf loc + locOf (Constr _ _ _ loc) = locOf loc + locOf (Match _ _ _ loc) = locOf loc + locOf (Attr _ _ loc) = locOf loc + +-- | An entry in a record literal. +data FieldBase f vn + = RecordFieldExplicit Name (ExpBase f vn) SrcLoc + | RecordFieldImplicit vn (f PatternType) SrcLoc + +deriving instance Showable f vn => Show (FieldBase f vn) + +deriving instance Eq (FieldBase NoInfo VName) + +deriving instance Ord (FieldBase NoInfo VName) + +instance Located (FieldBase f vn) where + locOf (RecordFieldExplicit _ _ loc) = locOf loc + locOf (RecordFieldImplicit _ _ loc) = locOf loc + +-- | A case in a match expression. +data CaseBase f vn = CasePat (PatternBase f vn) (ExpBase f vn) SrcLoc + +deriving instance Showable f vn => Show (CaseBase f vn) + +deriving instance Eq (CaseBase NoInfo VName) + +deriving instance Ord (CaseBase NoInfo VName) + +instance Located (CaseBase f vn) where + locOf (CasePat _ _ loc) = locOf loc + +-- | Whether the loop is a @for@-loop or a @while@-loop. +data LoopFormBase f vn + = For (IdentBase f vn) (ExpBase f vn) + | ForIn (PatternBase f vn) (ExpBase f vn) + | While (ExpBase f vn) + +deriving instance Showable f vn => Show (LoopFormBase f vn) + +deriving instance Eq (LoopFormBase NoInfo VName) + +deriving instance Ord (LoopFormBase NoInfo VName) + +-- | A pattern as used most places where variables are bound (function +-- parameters, @let@ expressions, etc). +data PatternBase f vn + = TuplePattern [PatternBase f vn] SrcLoc + | RecordPattern [(Name, PatternBase f vn)] SrcLoc + | PatternParens (PatternBase f vn) SrcLoc + | Id vn (f PatternType) SrcLoc + | Wildcard (f PatternType) SrcLoc -- Nothing, i.e. underscore. + | PatternAscription (PatternBase f vn) (TypeDeclBase f vn) SrcLoc + | PatternLit (ExpBase f vn) (f PatternType) SrcLoc + | PatternConstr Name (f PatternType) [PatternBase f vn] SrcLoc + +deriving instance Showable f vn => Show (PatternBase f vn) + +deriving instance Eq (PatternBase NoInfo VName) + +deriving instance Ord (PatternBase NoInfo VName) + +instance Located (PatternBase f vn) where + locOf (TuplePattern _ loc) = locOf loc + locOf (RecordPattern _ loc) = locOf loc + locOf (PatternParens _ loc) = locOf loc + locOf (Id _ _ loc) = locOf loc + locOf (Wildcard _ loc) = locOf loc + locOf (PatternAscription _ _ loc) = locOf loc + locOf (PatternLit _ _ loc) = locOf loc + locOf (PatternConstr _ _ _ loc) = locOf loc + +-- | Documentation strings, including source location. +data DocComment = DocComment String SrcLoc + deriving (Show) + +instance Located DocComment where + locOf (DocComment _ loc) = locOf loc + +-- | Part of the type of an entry point. Has an actual type, and +-- maybe also an ascribed type expression. +data EntryType = EntryType + { entryType :: StructType, + entryAscribed :: Maybe (TypeExp VName) + } + deriving (Show) + +-- | Information about the external interface exposed by an entry +-- point. The important thing is that that we remember the original +-- source-language types, without desugaring them at all. The +-- annoying thing is that we do not require type annotations on entry +-- points, so the types can be either ascribed or inferred. +data EntryPoint = EntryPoint + { entryParams :: [EntryType], + entryReturn :: EntryType + } + deriving (Show) + +-- | Function Declarations +data ValBindBase f vn = ValBind + { -- | Just if this function is an entry point. If so, it also + -- contains the externally visible interface. Note that this may not + -- strictly be well-typed after some desugaring operations, as it + -- may refer to abstract types that are no longer in scope. + valBindEntryPoint :: Maybe (f EntryPoint), + valBindName :: vn, + valBindRetDecl :: Maybe (TypeExp vn), + valBindRetType :: f (StructType, [VName]), + valBindTypeParams :: [TypeParamBase vn], + valBindParams :: [PatternBase f vn], + valBindBody :: ExpBase f vn, + valBindDoc :: Maybe DocComment, + valBindAttrs :: [AttrInfo], + valBindLocation :: SrcLoc + } + +deriving instance Showable f vn => Show (ValBindBase f vn) + +instance Located (ValBindBase f vn) where + locOf = locOf . valBindLocation + +-- | Type Declarations +data TypeBindBase f vn = TypeBind + { typeAlias :: vn, + typeLiftedness :: Liftedness, + typeParams :: [TypeParamBase vn], + typeExp :: TypeDeclBase f vn, + typeDoc :: Maybe DocComment, + typeBindLocation :: SrcLoc + } + +deriving instance Showable f vn => Show (TypeBindBase f vn) + +instance Located (TypeBindBase f vn) where + locOf = locOf . typeBindLocation + +-- | The liftedness of a type parameter. By the @Ord@ instance, +-- @Unlifted < SizeLifted < Lifted@. +data Liftedness + = -- | May only be instantiated with a zero-order type of (possibly + -- symbolically) known size. + Unlifted + | -- | May only be instantiated with a zero-order type, but the size + -- can be varying. + SizeLifted + | -- | May be instantiated with a functional type. + Lifted + deriving (Eq, Ord, Show) + +-- | A type parameter. +data TypeParamBase vn + = -- | A type parameter that must be a size. + TypeParamDim vn SrcLoc + | -- | A type parameter that must be a type. + TypeParamType Liftedness vn SrcLoc + deriving (Eq, Ord, Show) + +instance Functor TypeParamBase where + fmap = fmapDefault + +instance Foldable TypeParamBase where + foldMap = foldMapDefault + +instance Traversable TypeParamBase where + traverse f (TypeParamDim v loc) = TypeParamDim <$> f v <*> pure loc + traverse f (TypeParamType l v loc) = TypeParamType l <$> f v <*> pure loc + +instance Located (TypeParamBase vn) where + locOf (TypeParamDim _ loc) = locOf loc + locOf (TypeParamType _ _ loc) = locOf loc + +-- | The name of a type parameter. +typeParamName :: TypeParamBase vn -> vn +typeParamName (TypeParamDim v _) = v +typeParamName (TypeParamType _ v _) = v + +-- | A spec is a component of a module type. +data SpecBase f vn + = ValSpec + { specName :: vn, + specTypeParams :: [TypeParamBase vn], + specType :: TypeDeclBase f vn, + specDoc :: Maybe DocComment, + specLocation :: SrcLoc + } + | TypeAbbrSpec (TypeBindBase f vn) + | -- | Abstract type. + TypeSpec Liftedness vn [TypeParamBase vn] (Maybe DocComment) SrcLoc + | ModSpec vn (SigExpBase f vn) (Maybe DocComment) SrcLoc + | IncludeSpec (SigExpBase f vn) SrcLoc + +deriving instance Showable f vn => Show (SpecBase f vn) + +instance Located (SpecBase f vn) where + locOf (ValSpec _ _ _ _ loc) = locOf loc + locOf (TypeAbbrSpec tbind) = locOf tbind + locOf (TypeSpec _ _ _ _ loc) = locOf loc + locOf (ModSpec _ _ _ loc) = locOf loc + locOf (IncludeSpec _ loc) = locOf loc + +-- | A module type expression. +data SigExpBase f vn + = SigVar (QualName vn) (f (M.Map VName VName)) SrcLoc + | SigParens (SigExpBase f vn) SrcLoc + | SigSpecs [SpecBase f vn] SrcLoc + | SigWith (SigExpBase f vn) (TypeRefBase f vn) SrcLoc + | SigArrow (Maybe vn) (SigExpBase f vn) (SigExpBase f vn) SrcLoc + +deriving instance Showable f vn => Show (SigExpBase f vn) + +-- | A type refinement. +data TypeRefBase f vn = TypeRef (QualName vn) [TypeParamBase vn] (TypeDeclBase f vn) SrcLoc + +deriving instance Showable f vn => Show (TypeRefBase f vn) + +instance Located (TypeRefBase f vn) where + locOf (TypeRef _ _ _ loc) = locOf loc + +instance Located (SigExpBase f vn) where + locOf (SigVar _ _ loc) = locOf loc + locOf (SigParens _ loc) = locOf loc + locOf (SigSpecs _ loc) = locOf loc + locOf (SigWith _ _ loc) = locOf loc + locOf (SigArrow _ _ _ loc) = locOf loc + +-- | Module type binding. +data SigBindBase f vn = SigBind + { sigName :: vn, + sigExp :: SigExpBase f vn, + sigDoc :: Maybe DocComment, + sigLoc :: SrcLoc + } + +deriving instance Showable f vn => Show (SigBindBase f vn) + +instance Located (SigBindBase f vn) where + locOf = locOf . sigLoc + +-- | Module expression. +data ModExpBase f vn + = ModVar (QualName vn) SrcLoc + | ModParens (ModExpBase f vn) SrcLoc + | -- | The contents of another file as a module. + ModImport FilePath (f FilePath) SrcLoc + | ModDecs [DecBase f vn] SrcLoc + | -- | Functor application. The first mapping is from parameter + -- names to argument names, while the second maps names in the + -- constructed module to the names inside the functor. + ModApply + (ModExpBase f vn) + (ModExpBase f vn) + (f (M.Map VName VName)) + (f (M.Map VName VName)) + SrcLoc + | ModAscript (ModExpBase f vn) (SigExpBase f vn) (f (M.Map VName VName)) SrcLoc + | ModLambda + (ModParamBase f vn) + (Maybe (SigExpBase f vn, f (M.Map VName VName))) + (ModExpBase f vn) + SrcLoc + +deriving instance Showable f vn => Show (ModExpBase f vn) + +instance Located (ModExpBase f vn) where + locOf (ModVar _ loc) = locOf loc + locOf (ModParens _ loc) = locOf loc + locOf (ModImport _ _ loc) = locOf loc + locOf (ModDecs _ loc) = locOf loc + locOf (ModApply _ _ _ _ loc) = locOf loc + locOf (ModAscript _ _ _ loc) = locOf loc + locOf (ModLambda _ _ _ loc) = locOf loc + +-- | A module binding. +data ModBindBase f vn = ModBind + { modName :: vn, + modParams :: [ModParamBase f vn], + modSignature :: Maybe (SigExpBase f vn, f (M.Map VName VName)), + modExp :: ModExpBase f vn, + modDoc :: Maybe DocComment, + modLocation :: SrcLoc + } + +deriving instance Showable f vn => Show (ModBindBase f vn) + +instance Located (ModBindBase f vn) where + locOf = locOf . modLocation + +-- | A module parameter. +data ModParamBase f vn = ModParam + { modParamName :: vn, + modParamType :: SigExpBase f vn, + modParamAbs :: f [VName], + modParamLocation :: SrcLoc + } + +deriving instance Showable f vn => Show (ModParamBase f vn) + +instance Located (ModParamBase f vn) where + locOf = locOf . modParamLocation + +-- | A top-level binding. +data DecBase f vn + = ValDec (ValBindBase f vn) + | TypeDec (TypeBindBase f vn) + | SigDec (SigBindBase f vn) + | ModDec (ModBindBase f vn) + | OpenDec (ModExpBase f vn) SrcLoc + | LocalDec (DecBase f vn) SrcLoc + | ImportDec FilePath (f FilePath) SrcLoc + +deriving instance Showable f vn => Show (DecBase f vn) + +instance Located (DecBase f vn) where + locOf (ValDec d) = locOf d + locOf (TypeDec d) = locOf d + locOf (SigDec d) = locOf d + locOf (ModDec d) = locOf d + locOf (OpenDec _ loc) = locOf loc + locOf (LocalDec _ loc) = locOf loc + locOf (ImportDec _ _ loc) = locOf loc + +-- | The program described by a single Futhark file. May depend on +-- other files. +data ProgBase f vn = Prog + { progDoc :: Maybe DocComment, + progDecs :: [DecBase f vn] + } + +deriving instance Showable f vn => Show (ProgBase f vn) + +--- Some prettyprinting definitions are here because we need them in +--- the Attributes module. + +instance Pretty PrimType where + ppr (Unsigned Int8) = text "u8" + ppr (Unsigned Int16) = text "u16" + ppr (Unsigned Int32) = text "u32" + ppr (Unsigned Int64) = text "u64" + ppr (Signed t) = ppr t + ppr (FloatType t) = ppr t + ppr Bool = text "bool" + +instance Pretty BinOp where + ppr Backtick = text "``" + ppr Plus = text "+" + ppr Minus = text "-" + ppr Pow = text "**" + ppr Times = text "*" + ppr Divide = text "/" + ppr Mod = text "%" + ppr Quot = text "//" + ppr Rem = text "%%" + ppr ShiftR = text ">>" + ppr ShiftL = text "<<" + ppr Band = text "&" + ppr Xor = text "^" + ppr Bor = text "|" + ppr LogAnd = text "&&" + ppr LogOr = text "||" + ppr Equal = text "==" + ppr NotEqual = text "!=" + ppr Less = text "<" + ppr Leq = text "<=" + ppr Greater = text ">" + ppr Geq = text ">=" + ppr PipeLeft = text "<|" ppr PipeRight = text "|>" 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 #-} + -- | -- -- Functions for generic traversals across Futhark syntax trees. The @@ -20,36 +21,38 @@ -- functions expressing the operations to be performed on the various -- types of nodes. module Language.Futhark.Traversals - ( ASTMapper(..) - , ASTMappable(..) - , identityMapper - , bareExp - ) where - -import qualified Data.Set as S -import qualified Data.List.NonEmpty as NE + ( ASTMapper (..), + ASTMappable (..), + identityMapper, + bareExp, + ) +where -import Language.Futhark.Syntax +import qualified Data.List.NonEmpty as NE +import qualified Data.Set as S +import Language.Futhark.Syntax -- | Express a monad mapping operation on a syntax node. Each element -- of this structure expresses the operation to be performed on a -- given child. -data ASTMapper m = ASTMapper { - mapOnExp :: ExpBase Info VName -> m (ExpBase Info VName) - , mapOnName :: VName -> m VName - , mapOnQualName :: QualName VName -> m (QualName VName) - , mapOnStructType :: StructType -> m StructType - , mapOnPatternType :: PatternType -> m PatternType +data ASTMapper m = ASTMapper + { mapOnExp :: ExpBase Info VName -> m (ExpBase Info VName), + mapOnName :: VName -> m VName, + mapOnQualName :: QualName VName -> m (QualName VName), + mapOnStructType :: StructType -> m StructType, + mapOnPatternType :: PatternType -> m PatternType } -- | An 'ASTMapper' that just leaves its input unchanged. identityMapper :: Monad m => ASTMapper m -identityMapper = ASTMapper { mapOnExp = return - , mapOnName = return - , mapOnQualName = return - , mapOnStructType = return - , mapOnPatternType = return - } +identityMapper = + ASTMapper + { mapOnExp = return, + mapOnName = return, + mapOnQualName = return, + mapOnStructType = return, + mapOnPatternType = return + } -- | The class of things that we can map an 'ASTMapper' across. class ASTMappable x where @@ -61,8 +64,8 @@ instance ASTMappable (ExpBase Info VName) where astMap tv (Var name t loc) = - Var <$> mapOnQualName tv name <*> traverse (mapOnPatternType tv) t <*> - pure loc + Var <$> mapOnQualName tv name <*> traverse (mapOnPatternType tv) t + <*> pure loc astMap _ (Literal val loc) = pure $ Literal val loc astMap _ (StringLit vs loc) = @@ -74,8 +77,9 @@ astMap tv (Parens e loc) = Parens <$> mapOnExp tv e <*> pure loc astMap tv (QualParens (name, nameloc) e loc) = - QualParens <$> ((,) <$> mapOnQualName tv name <*> pure nameloc) <*> - mapOnExp tv e <*> pure loc + QualParens <$> ((,) <$> mapOnQualName tv name <*> pure nameloc) + <*> mapOnExp tv e + <*> pure loc astMap tv (TupLit els loc) = TupLit <$> mapM (mapOnExp tv) els <*> pure loc astMap tv (RecordLit fields loc) = @@ -83,99 +87,134 @@ astMap tv (ArrayLit els t loc) = ArrayLit <$> mapM (mapOnExp tv) els <*> traverse (mapOnPatternType tv) t <*> pure loc astMap tv (Range start next end (t, ext) loc) = - Range <$> mapOnExp tv start <*> traverse (mapOnExp tv) next <*> - traverse (mapOnExp tv) end <*> - ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) <*> pure loc + Range <$> mapOnExp tv start <*> traverse (mapOnExp tv) next + <*> traverse (mapOnExp tv) end + <*> ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) + <*> pure loc astMap tv (Ascript e tdecl loc) = Ascript <$> mapOnExp tv e <*> astMap tv tdecl <*> pure loc astMap tv (Coerce e tdecl (t, ext) loc) = - Coerce <$> mapOnExp tv e <*> astMap tv tdecl <*> - ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) <*> pure loc - astMap tv (BinOp (fname, fname_loc) t (x,Info(xt,xext)) (y,Info(yt,yext)) (Info rt) ext loc) = - BinOp <$> ((,) <$> mapOnQualName tv fname <*> pure fname_loc) <*> - traverse (mapOnPatternType tv) t <*> - ((,) <$> mapOnExp tv x <*> - (Info <$> ((,) <$> mapOnStructType tv xt <*> pure xext))) <*> - ((,) <$> mapOnExp tv y <*> - (Info <$> ((,) <$> mapOnStructType tv yt <*> pure yext))) <*> - (Info <$> mapOnPatternType tv rt) <*> pure ext <*> pure loc + Coerce <$> mapOnExp tv e <*> astMap tv tdecl + <*> ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) + <*> pure loc + astMap tv (BinOp (fname, fname_loc) t (x, Info (xt, xext)) (y, Info (yt, yext)) (Info rt) ext loc) = + BinOp <$> ((,) <$> mapOnQualName tv fname <*> pure fname_loc) + <*> traverse (mapOnPatternType tv) t + <*> ( (,) <$> mapOnExp tv x + <*> (Info <$> ((,) <$> mapOnStructType tv xt <*> pure xext)) + ) + <*> ( (,) <$> mapOnExp tv y + <*> (Info <$> ((,) <$> mapOnStructType tv yt <*> pure yext)) + ) + <*> (Info <$> mapOnPatternType tv rt) + <*> pure ext + <*> pure loc astMap tv (Negate x loc) = Negate <$> mapOnExp tv x <*> pure loc astMap tv (If c texp fexp (t, ext) loc) = - If <$> mapOnExp tv c <*> mapOnExp tv texp <*> mapOnExp tv fexp <*> - ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) <*> pure loc + If <$> mapOnExp tv c <*> mapOnExp tv texp <*> mapOnExp tv fexp + <*> ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) + <*> pure loc astMap tv (Apply f arg d (Info t, ext) loc) = - Apply <$> mapOnExp tv f <*> mapOnExp tv arg <*> pure d <*> - ((,) <$> (Info <$> mapOnPatternType tv t) <*> pure ext) <*> pure loc + Apply <$> mapOnExp tv f <*> mapOnExp tv arg <*> pure d + <*> ((,) <$> (Info <$> mapOnPatternType tv t) <*> pure ext) + <*> pure loc astMap tv (LetPat pat e body (t, ext) loc) = - LetPat <$> astMap tv pat <*> mapOnExp tv e <*> mapOnExp tv body <*> - ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) <*> pure loc + LetPat <$> astMap tv pat <*> mapOnExp tv e <*> mapOnExp tv body + <*> ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) + <*> pure loc astMap tv (LetFun name (fparams, params, ret, t, e) body body_t loc) = - LetFun <$> mapOnName tv name <*> - ((,,,,) <$> mapM (astMap tv) fparams <*> mapM (astMap tv) params <*> - traverse (astMap tv) ret <*> traverse (mapOnStructType tv) t <*> - mapOnExp tv e) <*> - mapOnExp tv body <*> traverse (mapOnPatternType tv) body_t <*> pure loc + LetFun <$> mapOnName tv name + <*> ( (,,,,) <$> mapM (astMap tv) fparams <*> mapM (astMap tv) params + <*> traverse (astMap tv) ret + <*> traverse (mapOnStructType tv) t + <*> mapOnExp tv e + ) + <*> mapOnExp tv body + <*> traverse (mapOnPatternType tv) body_t + <*> pure loc astMap tv (LetWith dest src idxexps vexp body t loc) = - LetWith <$> - astMap tv dest <*> astMap tv src <*> - mapM (astMap tv) idxexps <*> mapOnExp tv vexp <*> - mapOnExp tv body <*> traverse (mapOnPatternType tv) t <*> pure loc + LetWith + <$> astMap tv dest + <*> astMap tv src + <*> mapM (astMap tv) idxexps + <*> mapOnExp tv vexp + <*> mapOnExp tv body + <*> traverse (mapOnPatternType tv) t + <*> pure loc astMap tv (Update src slice v loc) = - Update <$> mapOnExp tv src <*> mapM (astMap tv) slice <*> - mapOnExp tv v <*> pure loc + Update <$> mapOnExp tv src <*> mapM (astMap tv) slice + <*> mapOnExp tv v + <*> pure loc astMap tv (RecordUpdate src fs v (Info t) loc) = - RecordUpdate <$> mapOnExp tv src <*> pure fs <*> - mapOnExp tv v <*> (Info <$> mapOnPatternType tv t) <*> pure loc + RecordUpdate <$> mapOnExp tv src <*> pure fs + <*> mapOnExp tv v + <*> (Info <$> mapOnPatternType tv t) + <*> pure loc astMap tv (Project field e t loc) = Project field <$> mapOnExp tv e <*> traverse (mapOnPatternType tv) t <*> pure loc astMap tv (Index arr idxexps (t, ext) loc) = - Index <$> mapOnExp tv arr <*> mapM (astMap tv) idxexps <*> - ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) <*> pure loc + Index <$> mapOnExp tv arr <*> mapM (astMap tv) idxexps + <*> ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) + <*> pure loc astMap tv (Assert e1 e2 desc loc) = Assert <$> mapOnExp tv e1 <*> mapOnExp tv e2 <*> pure desc <*> pure loc astMap tv (Lambda params body ret t loc) = - Lambda <$> mapM (astMap tv) params <*> - mapOnExp tv body <*> traverse (astMap tv) ret <*> - traverse (traverse $ mapOnStructType tv) t <*> pure loc + Lambda <$> mapM (astMap tv) params + <*> mapOnExp tv body + <*> traverse (astMap tv) ret + <*> traverse (traverse $ mapOnStructType tv) t + <*> pure loc astMap tv (OpSection name t loc) = - OpSection <$> mapOnQualName tv name <*> - traverse (mapOnPatternType tv) t <*> pure loc + OpSection <$> mapOnQualName tv name + <*> traverse (mapOnPatternType tv) t + <*> pure loc astMap tv (OpSectionLeft name t arg (Info (t1a, argext), 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) <*> - ((,) <$> traverse (mapOnPatternType tv) t2 <*> pure retext) <*> pure loc - astMap tv (OpSectionRight name t arg (t1a, Info (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))) <*> - traverse (mapOnPatternType tv) t2 <*> pure loc + OpSectionLeft <$> mapOnQualName tv name + <*> traverse (mapOnPatternType tv) t + <*> mapOnExp tv arg + <*> ( (,) + <$> (Info <$> ((,) <$> mapOnStructType tv t1a <*> pure argext)) + <*> traverse (mapOnStructType tv) t1b + ) + <*> ((,) <$> traverse (mapOnPatternType tv) t2 <*> pure retext) + <*> pure loc + astMap tv (OpSectionRight name t arg (t1a, Info (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)) + ) + <*> traverse (mapOnPatternType tv) t2 + <*> pure loc astMap tv (ProjectSection fields t loc) = ProjectSection fields <$> traverse (mapOnPatternType tv) t <*> pure loc astMap tv (IndexSection idxs t loc) = - IndexSection <$> mapM (astMap tv) idxs <*> - traverse (mapOnPatternType tv) t <*> pure loc + IndexSection <$> mapM (astMap tv) idxs + <*> traverse (mapOnPatternType tv) t + <*> pure loc astMap tv (DoLoop sparams mergepat mergeexp form loopbody (Info (ret, ext)) loc) = - DoLoop <$> mapM (mapOnName tv) sparams <*> astMap tv mergepat <*> - mapOnExp tv mergeexp <*> astMap tv form <*> mapOnExp tv loopbody <*> - (Info <$> ((,) <$> mapOnPatternType tv ret <*> pure ext)) <*> pure loc + DoLoop <$> mapM (mapOnName tv) sparams <*> astMap tv mergepat + <*> mapOnExp tv mergeexp + <*> astMap tv form + <*> mapOnExp tv loopbody + <*> (Info <$> ((,) <$> mapOnPatternType tv ret <*> pure ext)) + <*> pure loc astMap tv (Constr name es ts loc) = Constr name <$> traverse (mapOnExp tv) es <*> traverse (mapOnPatternType tv) ts <*> pure loc astMap tv (Match e cases (t, ext) loc) = Match <$> mapOnExp tv e <*> astMap tv cases - <*> ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) <*> pure loc + <*> ((,) <$> traverse (mapOnPatternType tv) t <*> pure ext) + <*> pure loc astMap tv (Attr attr e loc) = Attr attr <$> mapOnExp tv e <*> pure loc instance ASTMappable (LoopFormBase Info VName) where astMap tv (For i bound) = For <$> astMap tv i <*> astMap tv bound astMap tv (ForIn pat e) = ForIn <$> astMap tv pat <*> mapOnExp tv e - astMap tv (While e) = While <$> mapOnExp tv e + astMap tv (While e) = While <$> mapOnExp tv e instance ASTMappable (TypeExp VName) where astMap tv (TEVar qn loc) = TEVar <$> mapOnQualName tv qn <*> pure loc @@ -206,8 +245,8 @@ instance ASTMappable (DimDecl VName) where astMap tv (NamedDim vn) = NamedDim <$> mapOnQualName tv vn - astMap _ (ConstDim k) = pure $ ConstDim k - astMap _ AnyDim = pure AnyDim + astMap _ (ConstDim k) = pure $ ConstDim k + astMap _ AnyDim = pure AnyDim instance ASTMappable (TypeParamBase VName) where astMap = traverse . mapOnName @@ -215,10 +254,10 @@ instance ASTMappable (DimIndexBase Info VName) where astMap tv (DimFix j) = DimFix <$> astMap tv j astMap tv (DimSlice i j stride) = - DimSlice <$> - maybe (return Nothing) (fmap Just . astMap tv) i <*> - maybe (return Nothing) (fmap Just . astMap tv) j <*> - maybe (return Nothing) (fmap Just . astMap tv) stride + DimSlice + <$> maybe (return Nothing) (fmap Just . astMap tv) i + <*> maybe (return Nothing) (fmap Just . astMap tv) j + <*> maybe (return Nothing) (fmap Just . astMap tv) stride instance ASTMappable Alias where astMap tv (AliasBound v) = AliasBound <$> mapOnName tv v @@ -228,11 +267,15 @@ astMap tv = fmap S.fromList . traverse (astMap tv) . S.toList type TypeTraverser f t dim1 als1 dim2 als2 = - (TypeName -> f TypeName) -> (dim1 -> f dim2) -> (als1 -> f als2) -> - t dim1 als1 -> f (t dim2 als2) + (TypeName -> f TypeName) -> + (dim1 -> f dim2) -> + (als1 -> f als2) -> + t dim1 als1 -> + f (t dim2 als2) -traverseScalarType :: Applicative f => - TypeTraverser f ScalarTypeBase dim1 als1 dims als2 +traverseScalarType :: + Applicative f => + TypeTraverser f ScalarTypeBase dim1 als1 dims als2 traverseScalarType _ _ _ (Prim t) = pure $ Prim t traverseScalarType f g h (Record fs) = Record <$> traverse (traverseType f g h) fs traverseScalarType f g h (TypeVar als u t args) = @@ -242,17 +285,22 @@ traverseScalarType f g h (Sum cs) = Sum <$> (traverse . traverse) (traverseType f g h) cs -traverseType :: Applicative f => - TypeTraverser f TypeBase dim1 als1 dims als2 +traverseType :: + Applicative f => + TypeTraverser f TypeBase dim1 als1 dims als2 traverseType f g h (Array als u et shape) = - Array <$> h als <*> pure u <*> - traverseScalarType f g pure et <*> traverse g shape + Array <$> h als <*> pure u + <*> traverseScalarType f g pure et + <*> traverse g shape traverseType f g h (Scalar t) = Scalar <$> traverseScalarType f g h t -traverseTypeArg :: Applicative f => - (TypeName -> f TypeName) -> (dim1 -> f dim2) - -> TypeArg dim1 -> f (TypeArg dim2) +traverseTypeArg :: + Applicative f => + (TypeName -> f TypeName) -> + (dim1 -> f dim2) -> + TypeArg dim1 -> + f (TypeArg dim2) traverseTypeArg _ g (TypeArgDim d loc) = TypeArgDim <$> g d <*> pure loc traverseTypeArg f g (TypeArgType t loc) = @@ -260,11 +308,13 @@ instance ASTMappable StructType where astMap tv = traverseType f (astMap tv) pure - where f = fmap typeNameFromQualName . mapOnQualName tv . qualNameFromTypeName + where + f = fmap typeNameFromQualName . mapOnQualName tv . qualNameFromTypeName instance ASTMappable PatternType where astMap tv = traverseType f (astMap tv) (astMap tv) - where f = fmap typeNameFromQualName . mapOnQualName tv . qualNameFromTypeName + where + f = fmap typeNameFromQualName . mapOnQualName tv . qualNameFromTypeName instance ASTMappable (TypeDeclBase Info VName) where astMap tv (TypeDecl dt (Info et)) = @@ -288,7 +338,7 @@ astMap tv (Wildcard (Info t) loc) = Wildcard <$> (Info <$> mapOnPatternType tv t) <*> pure loc astMap tv (PatternLit e (Info t) loc) = - PatternLit <$> astMap tv e <*> (Info <$> mapOnPatternType tv t) <*> pure loc + PatternLit <$> astMap tv e <*> (Info <$> mapOnPatternType tv t) <*> pure loc astMap tv (PatternConstr n (Info t) ps loc) = PatternConstr n <$> (Info <$> mapOnPatternType tv t) <*> mapM (astMap tv) ps <*> pure loc @@ -297,7 +347,8 @@ RecordFieldExplicit name <$> mapOnExp tv e <*> pure loc astMap tv (RecordFieldImplicit name t loc) = RecordFieldImplicit <$> mapOnName tv name - <*> traverse (mapOnPatternType tv) t <*> pure loc + <*> traverse (mapOnPatternType tv) t + <*> pure loc instance ASTMappable (CaseBase Info VName) where astMap tv (CasePat pat e loc) = @@ -312,11 +363,11 @@ instance ASTMappable a => ASTMappable (NE.NonEmpty a) where astMap tv = traverse $ astMap tv -instance (ASTMappable a, ASTMappable b) => ASTMappable (a,b) where - astMap tv (x,y) = (,) <$> astMap tv x <*> astMap tv y +instance (ASTMappable a, ASTMappable b) => ASTMappable (a, b) where + astMap tv (x, y) = (,) <$> astMap tv x <*> astMap tv y -instance (ASTMappable a, ASTMappable b, ASTMappable c) => ASTMappable (a,b,c) where - astMap tv (x,y,z) = (,,) <$> astMap tv x <*> astMap tv y <*> astMap tv z +instance (ASTMappable a, ASTMappable b, ASTMappable c) => ASTMappable (a, b, c) where + astMap tv (x, y, z) = (,,) <$> astMap tv x <*> astMap tv y <*> astMap tv z -- It would be lovely if the following code would be written in terms -- of ASTMappable, but unfortunately it involves changing the Info @@ -373,13 +424,17 @@ bareExp (RecordLit fields loc) = RecordLit (map bareField fields) loc bareExp (ArrayLit els _ loc) = ArrayLit (map bareExp els) NoInfo loc bareExp (Range start next end _ loc) = - Range (bareExp start) (fmap bareExp next) - (fmap bareExp end) (NoInfo, NoInfo) loc + Range + (bareExp start) + (fmap bareExp next) + (fmap bareExp end) + (NoInfo, NoInfo) + loc bareExp (Ascript e tdecl loc) = Ascript (bareExp e) (bareTypeDecl tdecl) loc bareExp (Coerce e tdecl _ loc) = Coerce (bareExp e) (bareTypeDecl tdecl) (NoInfo, NoInfo) loc -bareExp (BinOp fname _ (x,_) (y,_) _ _ loc) = +bareExp (BinOp fname _ (x, _) (y, _) _ _ loc) = BinOp fname NoInfo (bareExp x, NoInfo) (bareExp y, NoInfo) NoInfo NoInfo loc bareExp (Negate x loc) = Negate (bareExp x) loc bareExp (If c texp fexp _ loc) = @@ -391,9 +446,14 @@ bareExp (LetFun name (fparams, params, ret, _, e) body _ loc) = LetFun name (fparams, map barePat params, ret, NoInfo, bareExp e) (bareExp body) NoInfo loc bareExp (LetWith (Ident dest _ destloc) (Ident src _ srcloc) idxexps vexp body _ loc) = - LetWith (Ident dest NoInfo destloc) (Ident src NoInfo srcloc) - (map bareDimIndex idxexps) (bareExp vexp) - (bareExp body) NoInfo loc + LetWith + (Ident dest NoInfo destloc) + (Ident src NoInfo srcloc) + (map bareDimIndex idxexps) + (bareExp vexp) + (bareExp body) + NoInfo + loc bareExp (Update src slice v loc) = Update (bareExp src) (map bareDimIndex slice) (bareExp v) loc bareExp (RecordUpdate src fs v _ loc) = @@ -413,11 +473,17 @@ bareExp (IndexSection slice _ loc) = IndexSection (map bareDimIndex slice) NoInfo loc bareExp (DoLoop _ mergepat mergeexp form loopbody _ loc) = - DoLoop [] (barePat mergepat) (bareExp mergeexp) (bareLoopForm form) - (bareExp loopbody) NoInfo loc + DoLoop + [] + (barePat mergepat) + (bareExp mergeexp) + (bareLoopForm form) + (bareExp loopbody) + NoInfo + loc bareExp (Constr name es _ loc) = Constr name (map bareExp es) NoInfo loc bareExp (Match e cases _ loc) = - Match (bareExp e) (fmap bareCase cases) (NoInfo,NoInfo) loc + Match (bareExp e) (fmap bareCase cases) (NoInfo, NoInfo) loc bareExp (Attr attr e loc) = Attr attr (bareExp e) loc diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs --- a/src/Language/Futhark/TypeChecker.hs +++ b/src/Language/Futhark/TypeChecker.hs @@ -1,41 +1,41 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Safe #-} + -- | The type checker checks whether the program is type-consistent -- and adds type annotations and various other elaborations. The -- program does not need to have any particular properties for the -- type checker to function; in particular it does not need unique -- names. module Language.Futhark.TypeChecker - ( checkProg - , checkExp - , checkDec - , checkModExp - , TypeError - , Warnings - , initialEnv + ( checkProg, + checkExp, + checkDec, + checkModExp, + TypeError, + Warnings, + initialEnv, ) - where +where import Control.Monad.Except import Control.Monad.Writer hiding (Sum) +import Data.Char (isAlpha, isAlphaNum) +import Data.Either import Data.List (isPrefixOf) +import qualified Data.Map.Strict as M import Data.Maybe -import Data.Either import Data.Ord -import qualified Data.Map.Strict as M import qualified Data.Set as S - -import Prelude hiding (abs, mod) - +import Futhark.FreshNames hiding (newName) +import Futhark.Util.Pretty hiding (space) import Language.Futhark import Language.Futhark.Semantic -import Futhark.FreshNames hiding (newName) -import Language.Futhark.TypeChecker.Monad import Language.Futhark.TypeChecker.Modules +import Language.Futhark.TypeChecker.Monad import Language.Futhark.TypeChecker.Terms import Language.Futhark.TypeChecker.Types -import Futhark.Util.Pretty hiding (space) +import Prelude hiding (abs, mod) --- The main checker @@ -44,82 +44,93 @@ -- Accepts a mapping from file names (excluding extension) to -- previously type checker results. The 'FilePath' is used to resolve -- relative @import@s. -checkProg :: Imports - -> VNameSource - -> ImportName - -> UncheckedProg - -> Either TypeError (FileModule, Warnings, VNameSource) +checkProg :: + Imports -> + VNameSource -> + ImportName -> + UncheckedProg -> + Either TypeError (FileModule, Warnings, VNameSource) checkProg files src name prog = runTypeM initialEnv files' name src $ checkProgM prog - where files' = M.map fileEnv $ M.fromList files + where + files' = M.map fileEnv $ M.fromList files -- | Type check a single expression containing no type information, -- yielding either a type error or the same expression annotated with -- type information. Also returns a list of type parameters, which -- will be nonempty if the expression is polymorphic. See also -- 'checkProg'. -checkExp :: Imports - -> VNameSource - -> Env - -> UncheckedExp - -> Either TypeError ([TypeParam], Exp) +checkExp :: + Imports -> + VNameSource -> + Env -> + UncheckedExp -> + Either TypeError ([TypeParam], Exp) checkExp files src env e = do (e', _, _) <- runTypeM env files' (mkInitialImport "") src $ checkOneExp e return e' - where files' = M.map fileEnv $ M.fromList files + where + files' = M.map fileEnv $ M.fromList files -- | Type check a single declaration containing no type information, -- yielding either a type error or the same declaration annotated with -- type information along the Env produced by that declaration. See -- also 'checkProg'. -checkDec :: Imports - -> VNameSource - -> Env - -> ImportName - -> UncheckedDec - -> Either TypeError (Env, Dec, VNameSource) +checkDec :: + Imports -> + VNameSource -> + Env -> + ImportName -> + UncheckedDec -> + Either TypeError (Env, Dec, VNameSource) checkDec files src env name d = do ((env', d'), _, src') <- runTypeM env files' name src $ do (_, env', d') <- checkOneDec d return (env' <> env, d') return (env', d', src') - where files' = M.map fileEnv $ M.fromList files + where + files' = M.map fileEnv $ M.fromList files -- | Type check a single module expression containing no type information, -- yielding either a type error or the same expression annotated with -- type information along the Env produced by that declaration. See -- also 'checkProg'. -checkModExp :: Imports - -> VNameSource - -> Env - -> ModExpBase NoInfo Name - -> Either TypeError (MTy, ModExpBase Info VName) +checkModExp :: + Imports -> + VNameSource -> + Env -> + ModExpBase NoInfo Name -> + Either TypeError (MTy, ModExpBase Info VName) checkModExp files src env me = do (x, _, _) <- runTypeM env files' (mkInitialImport "") src $ checkOneModExp me return x - where files' = M.map fileEnv $ M.fromList files + where + files' = M.map fileEnv $ M.fromList files -- | An initial environment for the type checker, containing -- intrinsics and such. initialEnv :: Env -initialEnv = intrinsicsModule - { envModTable = initialModTable - , envNameMap = M.insert - (Term, nameFromString "intrinsics") - (qualName intrinsics_v) - topLevelNameMap - } - where initialTypeTable = M.fromList $ mapMaybe addIntrinsicT $ M.toList intrinsics - initialModTable = M.singleton intrinsics_v (ModEnv intrinsicsModule) +initialEnv = + intrinsicsModule + { envModTable = initialModTable, + envNameMap = + M.insert + (Term, nameFromString "intrinsics") + (qualName intrinsics_v) + topLevelNameMap + } + where + initialTypeTable = M.fromList $ mapMaybe addIntrinsicT $ M.toList intrinsics + initialModTable = M.singleton intrinsics_v (ModEnv intrinsicsModule) - intrinsics_v = VName (nameFromString "intrinsics") 0 + intrinsics_v = VName (nameFromString "intrinsics") 0 - intrinsicsModule = Env mempty initialTypeTable mempty mempty intrinsicsNameMap + intrinsicsModule = Env mempty initialTypeTable mempty mempty intrinsicsNameMap - addIntrinsicT (name, IntrinsicType t) = - Just (name, TypeAbbr Unlifted [] $ Scalar $ Prim t) - addIntrinsicT _ = - Nothing + addIntrinsicT (name, IntrinsicType t) = + Just (name, TypeAbbr Unlifted [] $ Scalar $ Prim t) + addIntrinsicT _ = + Nothing checkProgM :: UncheckedProg -> TypeM FileModule checkProgM (Prog doc decs) = do @@ -127,54 +138,64 @@ (abs, env, decs') <- checkDecs decs return (FileModule abs env $ Prog doc decs') -dupDefinitionError :: MonadTypeChecker m => - Namespace -> Name -> SrcLoc -> SrcLoc -> m a +dupDefinitionError :: + MonadTypeChecker m => + Namespace -> + Name -> + SrcLoc -> + SrcLoc -> + m a dupDefinitionError space name loc1 loc2 = typeError loc1 mempty $ - "Duplicate definition of" <+> ppr space <+> - pprName name <> ". Previously defined at" <+> text (locStr loc2) <> "." + "Duplicate definition of" <+> ppr space + <+> pprName name <> ". Previously defined at" + <+> text (locStr loc2) <> "." checkForDuplicateDecs :: [DecBase NoInfo Name] -> TypeM () checkForDuplicateDecs = foldM_ (flip f) mempty - where check namespace name loc known = - case M.lookup (namespace, name) known of - Just loc' -> - dupDefinitionError namespace name loc loc' - _ -> return $ M.insert (namespace, name) loc known - - f (ValDec vb) = - check Term (valBindName vb) (srclocOf vb) - - f (TypeDec (TypeBind name _ _ _ _ loc)) = - check Type name loc - - f (SigDec (SigBind name _ _ loc)) = - check Signature name loc - - f (ModDec (ModBind name _ _ _ _ loc)) = - check Term name loc + where + check namespace name loc known = + case M.lookup (namespace, name) known of + Just loc' -> + dupDefinitionError namespace name loc loc' + _ -> return $ M.insert (namespace, name) loc known - f OpenDec{} = return - f LocalDec{} = return - f ImportDec{} = return + f (ValDec vb) = + check Term (valBindName vb) (srclocOf vb) + f (TypeDec (TypeBind name _ _ _ _ loc)) = + check Type name loc + f (SigDec (SigBind name _ _ loc)) = + check Signature name loc + f (ModDec (ModBind name _ _ _ _ loc)) = + check Term name loc + f OpenDec {} = return + f LocalDec {} = return + f ImportDec {} = return bindingTypeParams :: [TypeParam] -> TypeM a -> TypeM a bindingTypeParams tparams = localEnv env - where env = mconcat $ map typeParamEnv tparams + where + env = mconcat $ map typeParamEnv tparams - typeParamEnv (TypeParamDim v _) = - mempty { envVtable = - M.singleton v $ BoundV [] (Scalar $ Prim $ Signed Int32) } - typeParamEnv (TypeParamType l v _) = - mempty { envTypeTable = - M.singleton v $ TypeAbbr l [] $ - Scalar $ TypeVar () Nonunique (typeName v) [] } + typeParamEnv (TypeParamDim v _) = + mempty + { envVtable = + M.singleton v $ BoundV [] (Scalar $ Prim $ Signed Int64) + } + typeParamEnv (TypeParamType l v _) = + mempty + { envTypeTable = + M.singleton v $ + TypeAbbr l [] $ + Scalar $ TypeVar () Nonunique (typeName v) [] + } emptyDimParam :: StructType -> Bool emptyDimParam = isNothing . traverseDims onDim - where onDim _ pos AnyDim | pos `elem` [PosImmediate, PosParam] = Nothing - onDim _ _ d = Just d + where + onDim _ pos AnyDim | pos `elem` [PosImmediate, PosParam] = Nothing + onDim _ _ d = Just d -- In this function, after the recursion, we add the Env of the -- current Spec *after* the one that is returned from the recursive @@ -183,9 +204,7 @@ -- the specific structure of substitutions in case some module type is -- redundantly imported multiple times). checkSpecs :: [SpecBase NoInfo Name] -> TypeM (TySet, Env, [SpecBase Info VName]) - checkSpecs [] = return (mempty, mempty, []) - checkSpecs (ValSpec name tparams vtype doc loc : specs) = bindSpaced [(Term, name)] $ do name' <- checkName Term name loc @@ -196,75 +215,90 @@ when (emptyDimParam $ unInfo $ expandedType vtype') $ typeError loc mempty $ - "All function parameters must have non-anonymous sizes." - "Hint: add size parameters to" <+> pquote (pprName name') <> "." + "All function parameters must have non-anonymous sizes." + "Hint: add size parameters to" <+> pquote (pprName name') <> "." let (params, _) = unfoldFunType $ unInfo $ expandedType vtype' when (null params && any isSizeParam tparams) $ - typeError loc mempty - "Size parameters are only allowed on bindings that also have value parameters." + typeError + loc + mempty + "Size parameters are only allowed on bindings that also have value parameters." let binding = BoundV tparams' $ unInfo $ expandedType vtype' valenv = - mempty { envVtable = M.singleton name' binding - , envNameMap = M.singleton (Term, name) $ qualName name' - } + mempty + { envVtable = M.singleton name' binding, + envNameMap = M.singleton (Term, name) $ qualName name' + } (abstypes, env, specs') <- localEnv valenv $ checkSpecs specs - return (abstypes, - env <> valenv, - ValSpec name' tparams' vtype' doc loc : specs') - + return + ( abstypes, + env <> valenv, + ValSpec name' tparams' vtype' doc loc : specs' + ) checkSpecs (TypeAbbrSpec tdec : specs) = bindSpaced [(Type, typeAlias tdec)] $ do (tenv, tdec') <- checkTypeBind tdec (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs - return (abstypes, - env <> tenv, - TypeAbbrSpec tdec' : specs') - + return + ( abstypes, + env <> tenv, + TypeAbbrSpec tdec' : specs' + ) checkSpecs (TypeSpec l name ps doc loc : specs) = checkTypeParams ps $ \ps' -> - bindSpaced [(Type, name)] $ do - name' <- checkName Type name loc - let tenv = mempty - { envNameMap = - M.singleton (Type, name) $ qualName name' - , envTypeTable = - M.singleton name' $ TypeAbbr l ps' $ - Scalar $ TypeVar () Nonunique (typeName name') $ - map typeParamToArg ps' - } - (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs - return (M.insert (qualName name') l abstypes, - env <> tenv, - TypeSpec l name' ps' doc loc : specs') - + bindSpaced [(Type, name)] $ do + name' <- checkName Type name loc + let tenv = + mempty + { envNameMap = + M.singleton (Type, name) $ qualName name', + envTypeTable = + M.singleton name' $ + TypeAbbr l ps' $ + Scalar $ + TypeVar () Nonunique (typeName name') $ + map typeParamToArg ps' + } + (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs + return + ( M.insert (qualName name') l abstypes, + env <> tenv, + TypeSpec l name' ps' doc loc : specs' + ) checkSpecs (ModSpec name sig doc loc : specs) = bindSpaced [(Term, name)] $ do name' <- checkName Term name loc (mty, sig') <- checkSigExp sig - let senv = mempty { envNameMap = M.singleton (Term, name) $ qualName name' - , envModTable = M.singleton name' $ mtyMod mty - } + let senv = + mempty + { envNameMap = M.singleton (Term, name) $ qualName name', + envModTable = M.singleton name' $ mtyMod mty + } (abstypes, env, specs') <- localEnv senv $ checkSpecs specs - return (M.mapKeys (qualify name') (mtyAbs mty) <> abstypes, - env <> senv, - ModSpec name' sig' doc loc : specs') - + return + ( M.mapKeys (qualify name') (mtyAbs mty) <> abstypes, + env <> senv, + ModSpec name' sig' doc loc : specs' + ) checkSpecs (IncludeSpec e loc : specs) = do (e_abs, e_env, e') <- checkSigExpToEnv e mapM_ (warnIfShadowing . fmap baseName) $ M.keys e_abs (abstypes, env, specs') <- localEnv e_env $ checkSpecs specs - return (abstypes <> e_abs, - env <> e_env, - IncludeSpec e' loc : specs') - where warnIfShadowing qn = - (lookupType loc qn >> warnAbout qn) - `catchError` \_ -> return () - warnAbout qn = - warn loc $ "Inclusion shadows type " ++ quote (pretty qn) ++ "." + return + ( abstypes <> e_abs, + env <> e_env, + IncludeSpec e' loc : specs' + ) + where + warnIfShadowing qn = + (lookupType loc qn >> warnAbout qn) + `catchError` \_ -> return () + warnAbout qn = + warn loc $ "Inclusion shadows type " ++ quote (pretty qn) ++ "." checkSigExp :: SigExpBase NoInfo Name -> TypeM (MTy, SigExpBase Info VName) checkSigExp (SigParens e loc) = do @@ -290,32 +324,40 @@ case maybe_pname of Just pname -> bindSpaced [(Term, pname)] $ do pname' <- checkName Term pname loc - return (mempty { envNameMap = M.singleton (Term, pname) $ qualName pname' - , envModTable = M.singleton pname' e1_mod - }, - Just pname') + return + ( mempty + { envNameMap = M.singleton (Term, pname) $ qualName pname', + envModTable = M.singleton pname' e1_mod + }, + Just pname' + ) Nothing -> return (mempty, Nothing) (e2_mod, e2') <- localEnv env_for_e2 $ checkSigExp e2 - return (MTy mempty $ ModFun $ FunSig s_abs e1_mod e2_mod, - SigArrow maybe_pname' e1' e2' loc) + return + ( MTy mempty $ ModFun $ FunSig s_abs e1_mod e2_mod, + SigArrow maybe_pname' e1' e2' loc + ) checkSigExpToEnv :: SigExpBase NoInfo Name -> TypeM (TySet, Env, SigExpBase Info VName) checkSigExpToEnv e = do (MTy abs mod, e') <- checkSigExp e case mod of ModEnv env -> return (abs, env, e') - ModFun{} -> unappliedFunctor $ srclocOf e + ModFun {} -> unappliedFunctor $ srclocOf e checkSigBind :: SigBindBase NoInfo Name -> TypeM (Env, SigBindBase Info VName) checkSigBind (SigBind name e doc loc) = do (env, e') <- checkSigExp e bindSpaced [(Signature, name)] $ do name' <- checkName Signature name loc - return (mempty { envSigTable = M.singleton name' env - , envNameMap = M.singleton (Signature, name) (qualName name') - }, - SigBind name' e' doc loc) + return + ( mempty + { envSigTable = M.singleton name' env, + envNameMap = M.singleton (Signature, name) (qualName name') + }, + SigBind name' e' doc loc + ) checkOneModExp :: ModExpBase NoInfo Name -> TypeM (MTy, ModExpBase Info VName) checkOneModExp (ModParens e loc) = do @@ -324,18 +366,24 @@ checkOneModExp (ModDecs decs loc) = do checkForDuplicateDecs decs (abstypes, env, decs') <- checkDecs decs - return (MTy abstypes $ ModEnv env, - ModDecs decs' loc) + return + ( MTy abstypes $ ModEnv env, + ModDecs decs' loc + ) checkOneModExp (ModVar v loc) = do (v', env) <- lookupMod loc v - when (baseName (qualLeaf v') == nameFromString "intrinsics" && - baseTag (qualLeaf v') <= maxIntrinsicTag) $ - typeError loc mempty "The 'intrinsics' module may not be used in module expressions." + when + ( baseName (qualLeaf v') == nameFromString "intrinsics" + && baseTag (qualLeaf v') <= maxIntrinsicTag + ) + $ typeError loc mempty "The 'intrinsics' module may not be used in module expressions." return (MTy mempty env, ModVar v' loc) checkOneModExp (ModImport name NoInfo loc) = do (name', env) <- lookupImport loc name - return (MTy mempty $ ModEnv env, - ModImport name (Info name') loc) + return + ( MTy mempty $ ModEnv env, + ModImport name (Info name') loc + ) checkOneModExp (ModApply f e NoInfo NoInfo loc) = do (f_mty, f') <- checkOneModExp f case mtyMod f_mty of @@ -352,41 +400,49 @@ return (se_mty, ModAscript me' se' (Info match_subst) loc) checkOneModExp (ModLambda param maybe_fsig_e body_e loc) = withModParam param $ \param' param_abs param_mod -> do - (maybe_fsig_e', body_e', mty) <- checkModBody (fst <$> maybe_fsig_e) body_e loc - return (MTy mempty $ ModFun $ FunSig param_abs param_mod mty, - ModLambda param' maybe_fsig_e' body_e' loc) + (maybe_fsig_e', body_e', mty) <- checkModBody (fst <$> maybe_fsig_e) body_e loc + return + ( MTy mempty $ ModFun $ FunSig param_abs param_mod mty, + ModLambda param' maybe_fsig_e' body_e' loc + ) checkOneModExpToEnv :: ModExpBase NoInfo Name -> TypeM (TySet, Env, ModExpBase Info VName) checkOneModExpToEnv e = do (MTy abs mod, e') <- checkOneModExp e case mod of ModEnv env -> return (abs, env, e') - ModFun{} -> unappliedFunctor $ srclocOf e + ModFun {} -> unappliedFunctor $ srclocOf e -withModParam :: ModParamBase NoInfo Name - -> (ModParamBase Info VName -> TySet -> Mod -> TypeM a) - -> TypeM a +withModParam :: + ModParamBase NoInfo Name -> + (ModParamBase Info VName -> TySet -> Mod -> TypeM a) -> + TypeM a withModParam (ModParam pname psig_e NoInfo loc) m = do (MTy p_abs p_mod, psig_e') <- checkSigExp psig_e bindSpaced [(Term, pname)] $ do pname' <- checkName Term pname loc - let in_body_env = mempty { envModTable = M.singleton pname' p_mod } + let in_body_env = mempty {envModTable = M.singleton pname' p_mod} localEnv in_body_env $ m (ModParam pname' psig_e' (Info $ map qualLeaf $ M.keys p_abs) loc) p_abs p_mod -withModParams :: [ModParamBase NoInfo Name] - -> ([(ModParamBase Info VName, TySet, Mod)] -> TypeM a) - -> TypeM a +withModParams :: + [ModParamBase NoInfo Name] -> + ([(ModParamBase Info VName, TySet, Mod)] -> TypeM a) -> + TypeM a withModParams [] m = m [] -withModParams (p:ps) m = +withModParams (p : ps) m = withModParam p $ \p' pabs pmod -> - withModParams ps $ \ps' -> m $ (p',pabs,pmod) : ps' + withModParams ps $ \ps' -> m $ (p', pabs, pmod) : ps' -checkModBody :: Maybe (SigExpBase NoInfo Name) - -> ModExpBase NoInfo Name - -> SrcLoc - -> TypeM (Maybe (SigExp, Info (M.Map VName VName)), - ModExp, MTy) +checkModBody :: + Maybe (SigExpBase NoInfo Name) -> + ModExpBase NoInfo Name -> + SrcLoc -> + TypeM + ( Maybe (SigExp, Info (M.Map VName VName)), + ModExp, + MTy + ) checkModBody maybe_fsig_e body_e loc = do (body_mty, body_e') <- checkOneModExp body_e case maybe_fsig_e of @@ -402,118 +458,137 @@ (maybe_fsig_e', e', mty) <- checkModBody (fst <$> maybe_fsig_e) e loc bindSpaced [(Term, name)] $ do name' <- checkName Term name loc - return (mtyAbs mty, - mempty { envModTable = M.singleton name' $ mtyMod mty - , envNameMap = M.singleton (Term, name) $ qualName name' - }, - ModBind name' [] maybe_fsig_e' e' doc loc) -checkModBind (ModBind name (p:ps) maybe_fsig_e body_e doc loc) = do + return + ( mtyAbs mty, + mempty + { envModTable = M.singleton name' $ mtyMod mty, + envNameMap = M.singleton (Term, name) $ qualName name' + }, + ModBind name' [] maybe_fsig_e' e' doc loc + ) +checkModBind (ModBind name (p : ps) maybe_fsig_e body_e doc loc) = do (params', maybe_fsig_e', body_e', funsig) <- withModParam p $ \p' p_abs p_mod -> - withModParams ps $ \params_stuff -> do - let (ps', ps_abs, ps_mod) = unzip3 params_stuff - (maybe_fsig_e', body_e', mty) <- checkModBody (fst <$> maybe_fsig_e) body_e loc - let addParam (x,y) mty' = MTy mempty $ ModFun $ FunSig x y mty' - return (p' : ps', maybe_fsig_e', body_e', - FunSig p_abs p_mod $ foldr addParam mty $ zip ps_abs ps_mod) + withModParams ps $ \params_stuff -> do + let (ps', ps_abs, ps_mod) = unzip3 params_stuff + (maybe_fsig_e', body_e', mty) <- checkModBody (fst <$> maybe_fsig_e) body_e loc + let addParam (x, y) mty' = MTy mempty $ ModFun $ FunSig x y mty' + return + ( p' : ps', + maybe_fsig_e', + body_e', + FunSig p_abs p_mod $ foldr addParam mty $ zip ps_abs ps_mod + ) bindSpaced [(Term, name)] $ do name' <- checkName Term name loc - return (mempty, - mempty { envModTable = - M.singleton name' $ ModFun funsig - , envNameMap = - M.singleton (Term, name) $ qualName name' - }, - ModBind name' params' maybe_fsig_e' body_e' doc loc) + return + ( mempty, + mempty + { envModTable = + M.singleton name' $ ModFun funsig, + envNameMap = + M.singleton (Term, name) $ qualName name' + }, + ModBind name' params' maybe_fsig_e' body_e' doc loc + ) checkForDuplicateSpecs :: [SpecBase NoInfo Name] -> TypeM () checkForDuplicateSpecs = foldM_ (flip f) mempty - where check namespace name loc known = - case M.lookup (namespace, name) known of - Just loc' -> - dupDefinitionError namespace name loc loc' - _ -> return $ M.insert (namespace, name) loc known - - f (ValSpec name _ _ _ loc) = - check Term name loc - - f (TypeAbbrSpec (TypeBind name _ _ _ _ loc)) = - check Type name loc - - f (TypeSpec _ name _ _ loc) = - check Type name loc - - f (ModSpec name _ _ loc) = - check Term name loc + where + check namespace name loc known = + case M.lookup (namespace, name) known of + Just loc' -> + dupDefinitionError namespace name loc loc' + _ -> return $ M.insert (namespace, name) loc known - f IncludeSpec{} = - return + f (ValSpec name _ _ _ loc) = + check Term name loc + f (TypeAbbrSpec (TypeBind name _ _ _ _ loc)) = + check Type name loc + f (TypeSpec _ name _ _ loc) = + check Type name loc + f (ModSpec name _ _ loc) = + check Term name loc + f IncludeSpec {} = + return -checkTypeBind :: TypeBindBase NoInfo Name - -> TypeM (Env, TypeBindBase Info VName) +checkTypeBind :: + TypeBindBase NoInfo Name -> + TypeM (Env, TypeBindBase Info VName) checkTypeBind (TypeBind name l tps td doc loc) = checkTypeParams tps $ \tps' -> do (td', l') <- bindingTypeParams tps' $ checkTypeDecl td let used_dims = typeDimNames $ unInfo $ expandedType td' case filter ((`S.notMember` used_dims) . typeParamName) $ - filter isSizeParam tps' of + filter isSizeParam tps' of [] -> return () - tp:_ -> typeError loc mempty $ - "Size parameter" <+> pquote (ppr tp) <+> "unused." + tp : _ -> + typeError loc mempty $ + "Size parameter" <+> pquote (ppr tp) <+> "unused." case (l, l') of (_, Lifted) | l < Lifted -> typeError loc mempty $ - "Non-lifted type abbreviations may not contain functions." - "Hint: consider using 'type^'." + "Non-lifted type abbreviations may not contain functions." + "Hint: consider using 'type^'." (_, SizeLifted) | l < SizeLifted -> typeError loc mempty $ - "Non-size-lifted type abbreviations may not contain size-lifted types." - "Hint: consider using 'type~'." + "Non-size-lifted type abbreviations may not contain size-lifted types." + "Hint: consider using 'type~'." (Unlifted, _) | emptyDimParam $ unInfo $ expandedType td' -> - typeError loc mempty $ - "Non-lifted type abbreviations may not use anonymous sizes in their definition." - "Hint: use 'type~' or add size parameters to" <+> - pquote (pprName name) <> "." + typeError loc mempty $ + "Non-lifted type abbreviations may not use anonymous sizes in their definition." + "Hint: use 'type~' or add size parameters to" + <+> pquote (pprName name) <> "." _ -> return () bindSpaced [(Type, name)] $ do name' <- checkName Type name loc - return (mempty { envTypeTable = - M.singleton name' $ TypeAbbr l tps' $ unInfo $ expandedType td', - envNameMap = - M.singleton (Type, name) $ qualName name' - }, - TypeBind name' l tps' td' doc loc) - + return + ( mempty + { envTypeTable = + M.singleton name' $ TypeAbbr l tps' $ unInfo $ expandedType td', + envNameMap = + M.singleton (Type, name) $ qualName name' + }, + TypeBind name' l tps' td' doc loc + ) entryPoint :: [Pattern] -> Maybe (TypeExp VName) -> StructType -> EntryPoint entryPoint params orig_ret_te orig_ret = EntryPoint (map patternEntry params ++ more_params) rettype' - where (more_params, rettype') = - onRetType orig_ret_te orig_ret + where + (more_params, rettype') = + onRetType orig_ret_te orig_ret - patternEntry (PatternParens p _) = - patternEntry p - patternEntry (PatternAscription _ tdecl _) = - EntryType (unInfo (expandedType tdecl)) (Just (declaredType tdecl)) - patternEntry p = - EntryType (patternStructType p) Nothing + patternEntry (PatternParens p _) = + patternEntry p + patternEntry (PatternAscription _ tdecl _) = + EntryType (unInfo (expandedType tdecl)) (Just (declaredType tdecl)) + patternEntry p = + EntryType (patternStructType p) Nothing - onRetType (Just (TEArrow _ t1_te t2_te _)) (Scalar (Arrow _ _ t1 t2)) = - let (xs, y) = onRetType (Just t2_te) t2 - in (EntryType t1 (Just t1_te) : xs, y) - onRetType _ (Scalar (Arrow _ _ t1 t2)) = - let (xs, y) = onRetType Nothing t2 - in (EntryType t1 Nothing : xs, y) - onRetType te t = - ([], EntryType t te) + onRetType (Just (TEArrow _ t1_te t2_te _)) (Scalar (Arrow _ _ t1 t2)) = + let (xs, y) = onRetType (Just t2_te) t2 + in (EntryType t1 (Just t1_te) : xs, y) + onRetType _ (Scalar (Arrow _ _ t1 t2)) = + let (xs, y) = onRetType Nothing t2 + in (EntryType t1 Nothing : xs, y) + onRetType te t = + ([], EntryType t te) +entryPointNameIsAcceptable :: Name -> Bool +entryPointNameIsAcceptable = check . nameToString + where + check [] = True -- academic + check (c : cs) = isAlpha c && all constituent cs + constituent c = isAlphaNum c || c == '_' + checkValBind :: ValBindBase NoInfo Name -> TypeM (Env, ValBind) checkValBind (ValBind entry fname maybe_tdecl NoInfo tparams params body doc attrs loc) = do (fname', tparams', params', maybe_tdecl', rettype, retext, body') <- @@ -524,73 +599,79 @@ case entry' of Just _ + | not $ entryPointNameIsAcceptable fname -> + typeError loc mempty "Entry point names must start with a letter and contain only letters, digits, and underscores." | any isTypeParam tparams' -> - typeError loc mempty "Entry point functions may not be polymorphic." - + typeError loc mempty "Entry point functions may not be polymorphic." | not (all patternOrderZero params') - || not (all orderZero rettype_params) - || not (orderZero rettype') -> - typeError loc mempty "Entry point functions may not be higher-order." - + || not (all orderZero rettype_params) + || not (orderZero rettype') -> + typeError loc mempty "Entry point functions may not be higher-order." | sizes_only_in_ret <- - S.fromList (map typeParamName tparams') `S.intersection` - typeDimNames rettype' `S.difference` - foldMap typeDimNames (map patternStructType params' ++ rettype_params), + S.fromList (map typeParamName tparams') + `S.intersection` typeDimNames rettype' + `S.difference` foldMap typeDimNames (map patternStructType params' ++ rettype_params), not $ S.null sizes_only_in_ret -> - typeError loc mempty "Entry point functions must not be size-polymorphic in their return type." - + typeError loc mempty "Entry point functions must not be size-polymorphic in their return type." | p : _ <- filter nastyParameter params' -> - warn loc $ pretty $ "Entry point parameter\n" - indent 2 (ppr p) - "\nwill have an opaque type, so the entry point will likely not be callable." - + warn loc $ + pretty $ + "Entry point parameter\n" + indent 2 (ppr p) + "\nwill have an opaque type, so the entry point will likely not be callable." | nastyReturnType maybe_tdecl' rettype -> - warn loc $ pretty $ "Entry point return type\n" - indent 2 (ppr rettype) - "\nwill have an opaque type, so the result will likely not be usable." - + warn loc $ + pretty $ + "Entry point return type\n" + indent 2 (ppr rettype) + "\nwill have an opaque type, so the result will likely not be usable." _ -> return () let arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt - return (mempty { envVtable = - M.singleton fname' $ - BoundV tparams' $ foldr (arrow . patternParam) rettype params' - , envNameMap = - M.singleton (Term, fname) $ qualName fname' - }, - ValBind entry' fname' maybe_tdecl' (Info (rettype, retext)) tparams' params' body' doc attrs loc) + return + ( mempty + { envVtable = + M.singleton fname' $ + BoundV tparams' $ foldr (arrow . patternParam) rettype params', + envNameMap = + M.singleton (Term, fname) $ qualName fname' + }, + ValBind entry' fname' maybe_tdecl' (Info (rettype, retext)) tparams' params' body' doc attrs loc + ) nastyType :: Monoid als => TypeBase dim als -> Bool -nastyType (Scalar Prim{}) = False -nastyType t@Array{} = nastyType $ stripArray 1 t +nastyType (Scalar Prim {}) = False +nastyType t@Array {} = nastyType $ stripArray 1 t nastyType _ = True nastyReturnType :: Monoid als => Maybe (TypeExp VName) -> TypeBase dim als -> Bool nastyReturnType Nothing (Scalar (Arrow _ _ t1 t2)) = nastyType t1 || nastyReturnType Nothing t2 nastyReturnType (Just (TEArrow _ te1 te2 _)) (Scalar (Arrow _ _ t1 t2)) = - (not (niceTypeExp te1) && nastyType t1) || - nastyReturnType (Just te2) t2 + (not (niceTypeExp te1) && nastyType t1) + || nastyReturnType (Just te2) t2 nastyReturnType (Just te) _ | niceTypeExp te = False nastyReturnType te t | Just ts <- isTupleRecord t = - case te of - Just (TETuple tes _) -> or $ zipWith nastyType' (map Just tes) ts - _ -> any nastyType ts + case te of + Just (TETuple tes _) -> or $ zipWith nastyType' (map Just tes) ts + _ -> any nastyType ts | otherwise = nastyType' te t - where nastyType' (Just te') _ | niceTypeExp te' = False - nastyType' _ t' = nastyType t' + where + nastyType' (Just te') _ | niceTypeExp te' = False + nastyType' _ t' = nastyType t' nastyParameter :: Pattern -> Bool nastyParameter p = nastyType (patternType p) && not (ascripted p) - where ascripted (PatternAscription _ (TypeDecl te _) _) = niceTypeExp te - ascripted (PatternParens p' _) = ascripted p' - ascripted _ = False + where + ascripted (PatternAscription _ (TypeDecl te _) _) = niceTypeExp te + ascripted (PatternParens p' _) = ascripted p' + ascripted _ = False niceTypeExp :: TypeExp VName -> Bool niceTypeExp (TEVar (QualName [] _) _) = True -niceTypeExp (TEApply te TypeArgExpDim{} _) = niceTypeExp te +niceTypeExp (TEApply te TypeArgExpDim {} _) = niceTypeExp te niceTypeExp (TEArray te _ _) = niceTypeExp te niceTypeExp _ = False @@ -598,47 +679,43 @@ checkOneDec (ModDec struct) = do (abs, modenv, struct') <- checkModBind struct return (abs, modenv, ModDec struct') - checkOneDec (SigDec sig) = do (sigenv, sig') <- checkSigBind sig return (mempty, sigenv, SigDec sig') - checkOneDec (TypeDec tdec) = do (tenv, tdec') <- checkTypeBind tdec return (mempty, tenv, TypeDec tdec') - checkOneDec (OpenDec x loc) = do (x_abs, x_env, x') <- checkOneModExpToEnv x return (x_abs, x_env, OpenDec x' loc) - checkOneDec (LocalDec d loc) = do (abstypes, env, d') <- checkOneDec d return (abstypes, env, LocalDec d' loc) - checkOneDec (ImportDec name NoInfo loc) = do (name', env) <- lookupImport loc name when ("/prelude" `isPrefixOf` name) $ typeError loc mempty $ ppr name <+> "may not be explicitly imported." return (mempty, env, ImportDec name (Info name') loc) - checkOneDec (ValDec vb) = do (env, vb') <- checkValBind vb return (mempty, env, ValDec vb') checkDecs :: [DecBase NoInfo Name] -> TypeM (TySet, Env, [DecBase Info VName]) -checkDecs (LocalDec d loc:ds) = do +checkDecs (LocalDec d loc : ds) = do (d_abstypes, d_env, d') <- checkOneDec d (ds_abstypes, ds_env, ds') <- localEnv d_env $ checkDecs ds - return (d_abstypes <> ds_abstypes, - ds_env, - LocalDec d' loc : ds') - -checkDecs (d:ds) = do + return + ( d_abstypes <> ds_abstypes, + ds_env, + LocalDec d' loc : ds' + ) +checkDecs (d : ds) = do (d_abstypes, d_env, d') <- checkOneDec d (ds_abstypes, ds_env, ds') <- localEnv d_env $ checkDecs ds - return (d_abstypes <> ds_abstypes, - ds_env <> d_env, - d' : ds') - + return + ( d_abstypes <> ds_abstypes, + ds_env <> d_env, + d' : ds' + ) checkDecs [] = return (mempty, mempty, []) diff --git a/src/Language/Futhark/TypeChecker/Modules.hs b/src/Language/Futhark/TypeChecker/Modules.hs --- a/src/Language/Futhark/TypeChecker/Modules.hs +++ b/src/Language/Futhark/TypeChecker/Modules.hs @@ -1,32 +1,32 @@ -{-# LANGUAGE TupleSections #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Safe #-} +{-# LANGUAGE TupleSections #-} + -- | Implementation of the Futhark module system (at least most of it; -- some is scattered elsewhere in the type checker). module Language.Futhark.TypeChecker.Modules - ( matchMTys - , newNamesForMTy - , refineEnv - , applyFunctor - ) where + ( matchMTys, + newNamesForMTy, + refineEnv, + applyFunctor, + ) +where import Control.Monad.Except import Control.Monad.Writer hiding (Sum) +import Data.Either import Data.List (intersect) +import qualified Data.Map.Strict as M import Data.Maybe -import Data.Either import Data.Ord -import qualified Data.Map.Strict as M import qualified Data.Set as S - -import Prelude hiding (abs, mod) - +import Futhark.Util.Pretty import Language.Futhark import Language.Futhark.Semantic import Language.Futhark.TypeChecker.Monad -import Language.Futhark.TypeChecker.Unify (doUnification) import Language.Futhark.TypeChecker.Types -import Futhark.Util.Pretty +import Language.Futhark.TypeChecker.Unify (doUnification) +import Prelude hiding (abs, mod) substituteTypesInMod :: TypeSubs -> Mod -> Mod substituteTypesInMod substs (ModEnv e) = @@ -39,13 +39,15 @@ substituteTypesInEnv :: TypeSubs -> Env -> Env substituteTypesInEnv substs env = - env { envVtable = M.map (substituteTypesInBoundV substs) $ envVtable env - , envTypeTable = M.mapWithKey subT $ envTypeTable env - , envModTable = M.map (substituteTypesInMod substs) $ envModTable env - } - where subT name _ - | Just (TypeSub (TypeAbbr l ps t)) <- M.lookup name substs = TypeAbbr l ps t - subT _ (TypeAbbr l ps t) = TypeAbbr l ps $ substituteTypes substs t + env + { envVtable = M.map (substituteTypesInBoundV substs) $ envVtable env, + envTypeTable = M.mapWithKey subT $ envTypeTable env, + envModTable = M.map (substituteTypesInMod substs) $ envModTable env + } + where + subT name _ + | Just (TypeSub (TypeAbbr l ps t)) <- M.lookup name substs = TypeAbbr l ps t + subT _ (TypeAbbr l ps t) = TypeAbbr l ps $ substituteTypes substs t substituteTypesInBoundV :: TypeSubs -> BoundV -> BoundV substituteTypesInBoundV substs (BoundV tps t) = @@ -54,16 +56,22 @@ -- | All names defined anywhere in the 'Env'. allNamesInEnv :: Env -> S.Set VName allNamesInEnv (Env vtable ttable stable modtable _names) = - S.fromList (M.keys vtable ++ M.keys ttable ++ - M.keys stable ++ M.keys modtable) <> - mconcat (map allNamesInMTy (M.elems stable) ++ - map allNamesInMod (M.elems modtable) ++ - map allNamesInType (M.elems ttable)) - where allNamesInType (TypeAbbr _ ps _) = S.fromList $ map typeParamName ps + S.fromList + ( M.keys vtable ++ M.keys ttable + ++ M.keys stable + ++ M.keys modtable + ) + <> mconcat + ( map allNamesInMTy (M.elems stable) + ++ map allNamesInMod (M.elems modtable) + ++ map allNamesInType (M.elems ttable) + ) + where + allNamesInType (TypeAbbr _ ps _) = S.fromList $ map typeParamName ps allNamesInMod :: Mod -> S.Set VName allNamesInMod (ModEnv env) = allNamesInEnv env -allNamesInMod ModFun{} = mempty +allNamesInMod ModFun {} = mempty allNamesInMTy :: MTy -> S.Set VName allNamesInMTy (MTy abs mod) = @@ -80,7 +88,6 @@ rev_substs = M.fromList $ map (uncurry $ flip (,)) pairs return (substituteInMTy substs orig_mty, rev_substs) - where substituteInMTy :: M.Map VName VName -> MTy -> MTy substituteInMTy substs (MTy mty_abs mty_mod) = @@ -90,21 +97,23 @@ let vtable' = substituteInMap substituteInBinding vtable ttable' = substituteInMap substituteInTypeBinding ttable mtable' = substituteInMap substituteInMod modtable - in Env { envVtable = vtable' - , envTypeTable = ttable' - , envSigTable = mempty - , envModTable = mtable' - , envNameMap = M.map (fmap substitute) names - } + in Env + { envVtable = vtable', + envTypeTable = ttable', + envSigTable = mempty, + envModTable = mtable', + envNameMap = M.map (fmap substitute) names + } substitute v = fromMaybe v $ M.lookup v substs substituteInMap f m = let (ks, vs) = unzip $ M.toList m - in M.fromList $ - zip (map (\k -> fromMaybe k $ M.lookup k substs) ks) - (map f vs) + in M.fromList $ + zip + (map (\k -> fromMaybe k $ M.lookup k substs) ks) + (map f vs) substituteInBinding (BoundV ps t) = BoundV (map substituteInTypeParam ps) (substituteInType t) @@ -115,8 +124,10 @@ ModFun $ substituteInFunSig funsig substituteInFunSig (FunSig abs mod mty) = - FunSig (M.mapKeys (fmap substitute) abs) - (substituteInMod mod) (substituteInMTy substs mty) + FunSig + (M.mapKeys (fmap substitute) abs) + (substituteInMod mod) + (substituteInMTy substs mty) substituteInTypeBinding (TypeAbbr l ps t) = TypeAbbr l (map substituteInTypeParam ps) $ substituteInType t @@ -128,8 +139,9 @@ substituteInType :: StructType -> StructType substituteInType (Scalar (TypeVar () u (TypeName qs v) targs)) = - Scalar $ TypeVar () u (TypeName (map substitute qs) $ substitute v) $ - map substituteInTypeArg targs + Scalar $ + TypeVar () u (TypeName (map substitute qs) $ substitute v) $ + map substituteInTypeArg targs substituteInType (Scalar (Prim t)) = Scalar $ Prim t substituteInType (Scalar (Record ts)) = @@ -167,230 +179,307 @@ envTypeAbbrs :: Env -> M.Map VName TypeBinding envTypeAbbrs env = - envTypeTable env <> - (mconcat . map modTypeAbbrs . M.elems . envModTable) env + envTypeTable env + <> (mconcat . map modTypeAbbrs . M.elems . envModTable) env -- | Refine the given type name in the given env. -refineEnv :: SrcLoc -> TySet -> Env -> QualName Name -> [TypeParam] -> StructType - -> TypeM (QualName VName, TySet, Env) +refineEnv :: + SrcLoc -> + TySet -> + Env -> + QualName Name -> + [TypeParam] -> + StructType -> + TypeM (QualName VName, TySet, Env) refineEnv loc tset env tname ps t | Just (tname', TypeAbbr l cur_ps (Scalar (TypeVar () _ (TypeName qs v) _))) <- findTypeDef tname (ModEnv env), QualName (qualQuals tname') v `M.member` tset = - if paramsMatch cur_ps ps then - return (tname', - QualName qs v `M.delete` tset, - substituteTypesInEnv - (M.fromList [(qualLeaf tname', - TypeSub $ TypeAbbr l cur_ps t), - (v, TypeSub $ TypeAbbr l ps t)]) - env) - else typeError loc mempty $ "Cannot refine a type having" <+> - tpMsg ps <> " with a type having " <> tpMsg cur_ps <> "." + if paramsMatch cur_ps ps + then + return + ( tname', + QualName qs v `M.delete` tset, + substituteTypesInEnv + ( M.fromList + [ ( qualLeaf tname', + TypeSub $ TypeAbbr l cur_ps t + ), + (v, TypeSub $ TypeAbbr l ps t) + ] + ) + env + ) + else + typeError loc mempty $ + "Cannot refine a type having" + <+> tpMsg ps <> " with a type having " <> tpMsg cur_ps <> "." | otherwise = - typeError loc mempty $ ppr tname <+> "is not an abstract type in the module type." - where tpMsg [] = "no type parameters" - tpMsg xs = "type parameters" <+> spread (map ppr xs) + typeError loc mempty $ ppr tname <+> "is not an abstract type in the module type." + where + tpMsg [] = "no type parameters" + tpMsg xs = "type parameters" <+> spread (map ppr xs) paramsMatch :: [TypeParam] -> [TypeParam] -> Bool paramsMatch ps1 ps2 = length ps1 == length ps2 && all match (zip ps1 ps2) - where match (TypeParamType l1 _ _, TypeParamType l2 _ _) = l1 <= l2 - match (TypeParamDim _ _, TypeParamDim _ _) = True - match _ = False + where + match (TypeParamType l1 _ _, TypeParamType l2 _ _) = l1 <= l2 + match (TypeParamDim _ _, TypeParamDim _ _) = True + match _ = False -findBinding :: (Env -> M.Map VName v) - -> Namespace -> Name - -> Env - -> Maybe (VName, v) +findBinding :: + (Env -> M.Map VName v) -> + Namespace -> + Name -> + Env -> + Maybe (VName, v) findBinding table namespace name the_env = do QualName _ name' <- M.lookup (namespace, name) $ envNameMap the_env (name',) <$> M.lookup name' (table the_env) findTypeDef :: QualName Name -> Mod -> Maybe (QualName VName, TypeBinding) -findTypeDef _ ModFun{} = Nothing +findTypeDef _ ModFun {} = Nothing findTypeDef (QualName [] name) (ModEnv the_env) = do (name', tb) <- findBinding envTypeTable Type name the_env return (qualName name', tb) -findTypeDef (QualName (q:qs) name) (ModEnv the_env) = do +findTypeDef (QualName (q : qs) name) (ModEnv the_env) = do (q', q_mod) <- findBinding envModTable Term q the_env (QualName qs' name', tb) <- findTypeDef (QualName qs name) q_mod - return (QualName (q':qs') name', tb) + return (QualName (q' : qs') name', tb) -resolveAbsTypes :: TySet -> Mod -> TySet -> SrcLoc - -> Either TypeError (M.Map VName (QualName VName, TypeBinding)) +resolveAbsTypes :: + TySet -> + Mod -> + TySet -> + SrcLoc -> + Either TypeError (M.Map VName (QualName VName, TypeBinding)) resolveAbsTypes mod_abs mod sig_abs loc = do - let abs_mapping = M.fromList $ zip - (map (fmap baseName . fst) $ M.toList mod_abs) (M.toList mod_abs) - fmap M.fromList $ forM (M.toList sig_abs) $ \(name, name_l) -> - case findTypeDef (fmap baseName name) mod of - Just (name', TypeAbbr mod_l ps t) - | mod_l > name_l -> - mismatchedLiftedness name_l - (map qualLeaf $ M.keys mod_abs) (qualLeaf name) (mod_l, ps, t) - | name_l < SizeLifted, - emptyDims t -> - anonymousSizes (map qualLeaf $ M.keys mod_abs) - (qualLeaf name) (mod_l, ps, t) - | Just (abs_name, _) <- M.lookup (fmap baseName name) abs_mapping -> + let abs_mapping = + M.fromList $ + zip + (map (fmap baseName . fst) $ M.toList mod_abs) + (M.toList mod_abs) + fmap M.fromList $ + forM (M.toList sig_abs) $ \(name, name_l) -> + case findTypeDef (fmap baseName name) mod of + Just (name', TypeAbbr mod_l ps t) + | mod_l > name_l -> + mismatchedLiftedness + name_l + (map qualLeaf $ M.keys mod_abs) + (qualLeaf name) + (mod_l, ps, t) + | name_l < SizeLifted, + emptyDims t -> + anonymousSizes + (map qualLeaf $ M.keys mod_abs) + (qualLeaf name) + (mod_l, ps, t) + | Just (abs_name, _) <- M.lookup (fmap baseName name) abs_mapping -> return (qualLeaf name, (abs_name, TypeAbbr name_l ps t)) - | otherwise -> + | otherwise -> return (qualLeaf name, (name', TypeAbbr name_l ps t)) - _ -> - missingType loc $ fmap baseName name - where mismatchedLiftedness name_l abs name mod_t = - Left $ TypeError loc mempty $ - "Module defines" - indent 2 (ppTypeAbbr abs name mod_t) - "but module type requires" <+> text what <> "." - where what = case name_l of Unlifted -> "a non-lifted type" - SizeLifted -> "a size-lifted type" - Lifted -> "a lifted type" + _ -> + missingType loc $ fmap baseName name + where + mismatchedLiftedness name_l abs name mod_t = + Left $ + TypeError loc mempty $ + "Module defines" + indent 2 (ppTypeAbbr abs name mod_t) + "but module type requires" <+> text what <> "." + where + what = case name_l of + Unlifted -> "a non-lifted type" + SizeLifted -> "a size-lifted type" + Lifted -> "a lifted type" - anonymousSizes abs name mod_t = - Left $ TypeError loc mempty $ - "Module defines" - indent 2 (ppTypeAbbr abs name mod_t) - "which contains anonymous sizes, but module type requires non-lifted type." + anonymousSizes abs name mod_t = + Left $ + TypeError loc mempty $ + "Module defines" + indent 2 (ppTypeAbbr abs name mod_t) + "which contains anonymous sizes, but module type requires non-lifted type." - emptyDims :: StructType -> Bool - emptyDims = isNothing . traverseDims onDim - where onDim _ PosImmediate AnyDim = Nothing - onDim _ _ d = Just d + emptyDims :: StructType -> Bool + emptyDims = isNothing . traverseDims onDim + where + onDim _ PosImmediate AnyDim = Nothing + onDim _ _ d = Just d -resolveMTyNames :: MTy -> MTy - -> M.Map VName (QualName VName) +resolveMTyNames :: + MTy -> + MTy -> + M.Map VName (QualName VName) resolveMTyNames = resolveMTyNames' - where resolveMTyNames' (MTy _mod_abs mod) (MTy _sig_abs sig) = - resolveModNames mod sig + where + resolveMTyNames' (MTy _mod_abs mod) (MTy _sig_abs sig) = + resolveModNames mod sig - resolveModNames (ModEnv mod_env) (ModEnv sig_env) = - resolveEnvNames mod_env sig_env - resolveModNames (ModFun mod_fun) (ModFun sig_fun) = - resolveModNames (funSigMod mod_fun) (funSigMod sig_fun) <> - resolveMTyNames' (funSigMty mod_fun) (funSigMty sig_fun) - resolveModNames _ _ = - mempty + resolveModNames (ModEnv mod_env) (ModEnv sig_env) = + resolveEnvNames mod_env sig_env + resolveModNames (ModFun mod_fun) (ModFun sig_fun) = + resolveModNames (funSigMod mod_fun) (funSigMod sig_fun) + <> resolveMTyNames' (funSigMty mod_fun) (funSigMty sig_fun) + resolveModNames _ _ = + mempty - resolveEnvNames mod_env sig_env = - let mod_substs = resolve Term mod_env $ envModTable sig_env - onMod (modname, mod_env_mod) = - case M.lookup modname mod_substs of - Just (QualName _ modname') - | Just sig_env_mod <- - M.lookup modname' $ envModTable mod_env -> - resolveModNames sig_env_mod mod_env_mod - _ -> mempty - in mconcat [ resolve Term mod_env $ envVtable sig_env - , resolve Type mod_env $ envVtable sig_env - , resolve Signature mod_env $ envVtable sig_env - , mod_substs - , mconcat $ map onMod $ M.toList $ envModTable sig_env - ] + resolveEnvNames mod_env sig_env = + let mod_substs = resolve Term mod_env $ envModTable sig_env + onMod (modname, mod_env_mod) = + case M.lookup modname mod_substs of + Just (QualName _ modname') + | Just sig_env_mod <- + M.lookup modname' $ envModTable mod_env -> + resolveModNames sig_env_mod mod_env_mod + _ -> mempty + in mconcat + [ resolve Term mod_env $ envVtable sig_env, + resolve Type mod_env $ envVtable sig_env, + resolve Signature mod_env $ envVtable sig_env, + mod_substs, + mconcat $ map onMod $ M.toList $ envModTable sig_env + ] - resolve namespace mod_env = M.mapMaybeWithKey resolve' - where resolve' name _ = - M.lookup (namespace, baseName name) $ envNameMap mod_env + resolve namespace mod_env = M.mapMaybeWithKey resolve' + where + resolve' name _ = + M.lookup (namespace, baseName name) $ envNameMap mod_env missingType :: Pretty a => SrcLoc -> a -> Either TypeError b missingType loc name = - Left $ TypeError loc mempty $ - "Module does not define a type named" <+> ppr name <> "." + Left $ + TypeError loc mempty $ + "Module does not define a type named" <+> ppr name <> "." missingVal :: Pretty a => SrcLoc -> a -> Either TypeError b missingVal loc name = - Left $ TypeError loc mempty $ - "Module does not define a value named" <+> ppr name <> "." + Left $ + TypeError loc mempty $ + "Module does not define a value named" <+> ppr name <> "." missingMod :: Pretty a => SrcLoc -> a -> Either TypeError b missingMod loc name = - Left $ TypeError loc mempty $ - "Module does not define a module named" <+> ppr name <> "." + Left $ + TypeError loc mempty $ + "Module does not define a module named" <+> ppr name <> "." -mismatchedType :: SrcLoc - -> [VName] - -> VName - -> (Liftedness, [TypeParam], StructType) - -> (Liftedness, [TypeParam], StructType) - -> Either TypeError b +mismatchedType :: + SrcLoc -> + [VName] -> + VName -> + (Liftedness, [TypeParam], StructType) -> + (Liftedness, [TypeParam], StructType) -> + Either TypeError b mismatchedType loc abs name spec_t env_t = - Left $ TypeError loc mempty $ - "Module defines" - indent 2 (ppTypeAbbr abs name env_t) - "but module type requires" - indent 2 (ppTypeAbbr abs name spec_t) + Left $ + TypeError loc mempty $ + "Module defines" + indent 2 (ppTypeAbbr abs name env_t) + "but module type requires" + indent 2 (ppTypeAbbr abs name spec_t) ppTypeAbbr :: [VName] -> VName -> (Liftedness, [TypeParam], StructType) -> Doc ppTypeAbbr abs name (l, ps, Scalar (TypeVar () _ tn args)) | typeLeaf tn `elem` abs, map typeParamToArg ps == args = - "type" <> ppr l <+> pprName name <+> - spread (map ppr ps) + "type" <> ppr l <+> pprName name + <+> spread (map ppr ps) ppTypeAbbr _ name (l, ps, t) = - "type" <> ppr l <+> pprName name <+> - spread (map ppr ps) <+> equals <+/> - nest 2 (align (ppr t)) + "type" <> ppr l <+> pprName name + <+> spread (map ppr ps) + <+> equals + <+/> nest 2 (align (ppr t)) -- | Return new renamed/abstracted env, as well as a mapping from -- names in the signature to names in the new env. This is used for -- functor application. The first env is the module env, and the -- second the env it must match. -matchMTys :: MTy -> MTy -> SrcLoc - -> Either TypeError (M.Map VName VName) +matchMTys :: + MTy -> + MTy -> + SrcLoc -> + Either TypeError (M.Map VName VName) matchMTys orig_mty orig_mty_sig = - matchMTys' (M.map (DimSub . NamedDim) $ - resolveMTyNames orig_mty orig_mty_sig) - orig_mty orig_mty_sig + matchMTys' + ( M.map (DimSub . NamedDim) $ + resolveMTyNames orig_mty orig_mty_sig + ) + orig_mty + orig_mty_sig where - matchMTys' :: TypeSubs -> MTy -> MTy -> SrcLoc - -> Either TypeError (M.Map VName VName) - - matchMTys' _ (MTy _ ModFun{}) (MTy _ ModEnv{}) loc = - Left $ TypeError loc mempty - "Cannot match parametric module with non-parametric module type." - - matchMTys' _ (MTy _ ModEnv{}) (MTy _ ModFun{}) loc = - Left $ TypeError loc mempty - "Cannot match non-parametric module with paramatric module type." + matchMTys' :: + TypeSubs -> + MTy -> + MTy -> + SrcLoc -> + Either TypeError (M.Map VName VName) + matchMTys' _ (MTy _ ModFun {}) (MTy _ ModEnv {}) loc = + Left $ + TypeError + loc + mempty + "Cannot match parametric module with non-parametric module type." + matchMTys' _ (MTy _ ModEnv {}) (MTy _ ModFun {}) loc = + Left $ + TypeError + loc + mempty + "Cannot match non-parametric module with paramatric module type." matchMTys' old_abs_subst_to_type (MTy mod_abs mod) (MTy sig_abs sig) loc = do -- Check that abstract types in 'sig' have an implementation in -- 'mod'. This also gives us a substitution that we use to check -- the types of values. abs_substs <- resolveAbsTypes mod_abs mod sig_abs loc - let abs_subst_to_type = old_abs_subst_to_type <> - M.map (TypeSub . snd) abs_substs - abs_name_substs = M.map (qualLeaf . fst) abs_substs + let abs_subst_to_type = + old_abs_subst_to_type + <> M.map (TypeSub . snd) abs_substs + abs_name_substs = M.map (qualLeaf . fst) abs_substs substs <- matchMods abs_subst_to_type mod sig loc return (substs <> abs_name_substs) - matchMods :: TypeSubs -> Mod -> Mod -> SrcLoc - -> Either TypeError (M.Map VName VName) - matchMods _ ModEnv{} ModFun{} loc = - Left $ TypeError loc mempty - "Cannot match non-parametric module with parametric module type." - matchMods _ ModFun{} ModEnv{} loc = - Left $ TypeError loc mempty - "Cannot match parametric module with non-parametric module type." - + matchMods :: + TypeSubs -> + Mod -> + Mod -> + SrcLoc -> + Either TypeError (M.Map VName VName) + matchMods _ ModEnv {} ModFun {} loc = + Left $ + TypeError + loc + mempty + "Cannot match non-parametric module with parametric module type." + matchMods _ ModFun {} ModEnv {} loc = + Left $ + TypeError + loc + mempty + "Cannot match parametric module with non-parametric module type." matchMods abs_subst_to_type (ModEnv mod) (ModEnv sig) loc = matchEnvs abs_subst_to_type mod sig loc - - matchMods old_abs_subst_to_type - (ModFun (FunSig mod_abs mod_pmod mod_mod)) - (ModFun (FunSig sig_abs sig_pmod sig_mod)) - loc = do - abs_substs <- resolveAbsTypes mod_abs mod_pmod sig_abs loc - let abs_subst_to_type = old_abs_subst_to_type <> - M.map (TypeSub . snd) abs_substs - abs_name_substs = M.map (qualLeaf . fst) abs_substs - pmod_substs <- matchMods abs_subst_to_type mod_pmod sig_pmod loc - mod_substs <- matchMTys' abs_subst_to_type mod_mod sig_mod loc - return (pmod_substs <> mod_substs <> abs_name_substs) + matchMods + old_abs_subst_to_type + (ModFun (FunSig mod_abs mod_pmod mod_mod)) + (ModFun (FunSig sig_abs sig_pmod sig_mod)) + loc = do + abs_substs <- resolveAbsTypes mod_abs mod_pmod sig_abs loc + let abs_subst_to_type = + old_abs_subst_to_type + <> M.map (TypeSub . snd) abs_substs + abs_name_substs = M.map (qualLeaf . fst) abs_substs + pmod_substs <- matchMods abs_subst_to_type mod_pmod sig_pmod loc + mod_substs <- matchMTys' abs_subst_to_type mod_mod sig_mod loc + return (pmod_substs <> mod_substs <> abs_name_substs) - matchEnvs :: TypeSubs - -> Env -> Env -> SrcLoc - -> Either TypeError (M.Map VName VName) + matchEnvs :: + TypeSubs -> + Env -> + Env -> + SrcLoc -> + Either TypeError (M.Map VName VName) matchEnvs abs_subst_to_type env sig loc = do -- XXX: we only want to create substitutions for visible names. -- This must be wrong in some cases. Probably we need to @@ -400,35 +489,49 @@ -- Check that all type abbreviations are correctly defined. abbr_name_substs <- fmap M.fromList $ - forM (filter (isVisible . fst) $ M.toList $ - envTypeTable sig) $ \(name, TypeAbbr spec_l spec_ps spec_t) -> - case findBinding envTypeTable Type (baseName name) env of - Just (name', TypeAbbr l ps t) -> - matchTypeAbbr loc abs_subst_to_type name spec_l spec_ps spec_t name' l ps t - Nothing -> missingType loc $ baseName name + forM + ( filter (isVisible . fst) $ + M.toList $ + envTypeTable sig + ) + $ \(name, TypeAbbr spec_l spec_ps spec_t) -> + case findBinding envTypeTable Type (baseName name) env of + Just (name', TypeAbbr l ps t) -> + matchTypeAbbr loc abs_subst_to_type name spec_l spec_ps spec_t name' l ps t + Nothing -> missingType loc $ baseName name -- Check that all values are defined correctly, substituting the -- abstract types first. - val_substs <- fmap M.fromList $ forM (M.toList $ envVtable sig) $ \(name, spec_bv) -> do - let spec_bv' = substituteTypesInBoundV abs_subst_to_type spec_bv - case findBinding envVtable Term (baseName name) env of - Just (name', bv) -> matchVal loc name spec_bv' name' bv - _ -> missingVal loc (baseName name) + val_substs <- fmap M.fromList $ + forM (M.toList $ envVtable sig) $ \(name, spec_bv) -> do + let spec_bv' = substituteTypesInBoundV abs_subst_to_type spec_bv + case findBinding envVtable Term (baseName name) env of + Just (name', bv) -> matchVal loc name spec_bv' name' bv + _ -> missingVal loc (baseName name) -- Check for correct modules. - mod_substs <- fmap M.unions $ forM (M.toList $ envModTable sig) $ \(name, modspec) -> - case findBinding envModTable Term (baseName name) env of - Just (name', mod) -> - M.insert name name' <$> matchMods abs_subst_to_type mod modspec loc - Nothing -> - missingMod loc $ baseName name + mod_substs <- fmap M.unions $ + forM (M.toList $ envModTable sig) $ \(name, modspec) -> + case findBinding envModTable Term (baseName name) env of + Just (name', mod) -> + M.insert name name' <$> matchMods abs_subst_to_type mod modspec loc + Nothing -> + missingMod loc $ baseName name return $ val_substs <> mod_substs <> abbr_name_substs - matchTypeAbbr :: SrcLoc -> TypeSubs - -> VName -> Liftedness -> [TypeParam] -> StructType - -> VName -> Liftedness -> [TypeParam] -> StructType - -> Either TypeError (VName, VName) + matchTypeAbbr :: + SrcLoc -> + TypeSubs -> + VName -> + Liftedness -> + [TypeParam] -> + StructType -> + VName -> + Liftedness -> + [TypeParam] -> + StructType -> + Either TypeError (VName, VName) matchTypeAbbr loc abs_subst_to_type spec_name spec_l spec_ps spec_t name l ps t = do -- We have to create substitutions for the type parameters, too. unless (length spec_ps == length ps) $ nomatch spec_t @@ -440,44 +543,63 @@ when (M.member spec_name abs_subst_to_type) $ case S.toList (mustBeExplicitInType t) `intersect` map typeParamName ps of [] -> return () - d:_ -> Left $ TypeError loc mempty $ - "Type" - indent 2 (ppTypeAbbr [] name (l, ps, t)) - textwrap "cannot be made abstract because size parameter" <+/> pquote (pprName d) <+/> - textwrap "is not used as an array size in the definition." + d : _ -> + Left $ + TypeError loc mempty $ + "Type" + indent 2 (ppTypeAbbr [] name (l, ps, t)) + textwrap "cannot be made abstract because size parameter" + <+/> pquote (pprName d) + <+/> textwrap "is not used as an array size in the definition." - let spec_t' = substituteTypes (param_substs<>abs_subst_to_type) spec_t + let spec_t' = substituteTypes (param_substs <> abs_subst_to_type) spec_t if spec_t' == t then return (spec_name, name) else nomatch spec_t' - where nomatch spec_t' = mismatchedType loc (M.keys abs_subst_to_type) - spec_name (spec_l, spec_ps, spec_t') (l, ps, t) + where + nomatch spec_t' = + mismatchedType + loc + (M.keys abs_subst_to_type) + spec_name + (spec_l, spec_ps, spec_t') + (l, ps, t) - matchTypeParam (TypeParamDim x _) (TypeParamDim y _) = - pure $ M.singleton x $ DimSub $ NamedDim $ qualName y - matchTypeParam (TypeParamType Unlifted x _) (TypeParamType Unlifted y _) = - pure $ M.singleton x $ TypeSub $ TypeAbbr Unlifted [] $ - Scalar $ TypeVar () Nonunique (typeName y) [] - matchTypeParam (TypeParamType _ x _) (TypeParamType Lifted y _) = - pure $ M.singleton x $ TypeSub $ TypeAbbr Lifted [] $ - Scalar $ TypeVar () Nonunique (typeName y) [] - matchTypeParam _ _ = - nomatch spec_t + matchTypeParam (TypeParamDim x _) (TypeParamDim y _) = + pure $ M.singleton x $ DimSub $ NamedDim $ qualName y + matchTypeParam (TypeParamType Unlifted x _) (TypeParamType Unlifted y _) = + pure $ + M.singleton x $ + TypeSub $ + TypeAbbr Unlifted [] $ + Scalar $ TypeVar () Nonunique (typeName y) [] + matchTypeParam (TypeParamType _ x _) (TypeParamType Lifted y _) = + pure $ + M.singleton x $ + TypeSub $ + TypeAbbr Lifted [] $ + Scalar $ TypeVar () Nonunique (typeName y) [] + matchTypeParam _ _ = + nomatch spec_t - matchVal :: SrcLoc - -> VName -> BoundV - -> VName -> BoundV - -> Either TypeError (VName, VName) + matchVal :: + SrcLoc -> + VName -> + BoundV -> + VName -> + BoundV -> + Either TypeError (VName, VName) matchVal loc spec_name spec_v name v = case matchValBinding loc spec_v v of Nothing -> return (spec_name, name) Just problem -> - Left $ TypeError loc mempty $ - "Module type specifies" - indent 2 (ppValBind spec_name spec_v) - "but module provides" - indent 2 (ppValBind spec_name v) - fromMaybe mempty problem + Left $ + TypeError loc mempty $ + "Module type specifies" + indent 2 (ppValBind spec_name spec_v) + "but module provides" + indent 2 (ppValBind spec_name v) + fromMaybe mempty problem matchValBinding :: SrcLoc -> BoundV -> BoundV -> Maybe (Maybe Doc) matchValBinding loc (BoundV _ orig_spec_t) (BoundV tps orig_t) = @@ -486,27 +608,34 @@ Just $ Just $ msg <> ppr notes -- Even if they unify, we still have to verify the uniqueness -- properties. - Right t | noSizes t `subtypeOf` - noSizes orig_spec_t -> Nothing - | otherwise -> Just Nothing + Right t + | noSizes t + `subtypeOf` noSizes orig_spec_t -> + Nothing + | otherwise -> Just Nothing ppValBind v (BoundV tps t) = - "val" <+> pprName v <+> spread (map ppr tps) <+> colon - indent 2 (align (ppr t)) + "val" <+> pprName v <+> spread (map ppr tps) <+> colon + indent 2 (align (ppr t)) -- | Apply a parametric module to an argument. -applyFunctor :: SrcLoc -> FunSig -> MTy - -> TypeM (MTy, - M.Map VName VName, - M.Map VName VName) +applyFunctor :: + SrcLoc -> + FunSig -> + MTy -> + TypeM + ( MTy, + M.Map VName VName, + M.Map VName VName + ) applyFunctor applyloc (FunSig p_abs p_mod body_mty) a_mty = do p_subst <- badOnLeft $ matchMTys a_mty (MTy p_abs p_mod) applyloc -- Apply type abbreviations from a_mty to body_mty. let a_abbrs = mtyTypeAbbrs a_mty isSub v = case M.lookup v a_abbrs of - Just abbr -> Just $ TypeSub abbr - _ -> Just $ DimSub $ NamedDim $ qualName v + Just abbr -> Just $ TypeSub abbr + _ -> Just $ DimSub $ NamedDim $ qualName v type_subst = M.mapMaybe isSub p_subst body_mty' = substituteTypesInMTy type_subst body_mty (body_mty'', body_subst) <- newNamesForMTy body_mty' diff --git a/src/Language/Futhark/TypeChecker/Monad.hs b/src/Language/Futhark/TypeChecker/Monad.hs --- a/src/Language/Futhark/TypeChecker/Monad.hs +++ b/src/Language/Futhark/TypeChecker/Monad.hs @@ -1,78 +1,67 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TupleSections #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TupleSections #-} + -- | Main monad in which the type checker runs, as well as ancillary -- data definitions. module Language.Futhark.TypeChecker.Monad - ( TypeM - , runTypeM - , askEnv - , askImportName - , bindSpaced - , qualifyTypeVars - , lookupMTy - , lookupImport - , localEnv - - , TypeError(..) - , unappliedFunctor - , unknownVariable - , unknownType - , underscoreUse - , Notes - , aNote - - , MonadTypeChecker(..) - , checkName - , badOnLeft - - , module Language.Futhark.Warnings - - , Env(..) - , TySet - , FunSig(..) - , ImportTable - , NameMap - , BoundV(..) - , Mod(..) - , TypeBinding(..) - , MTy(..) - - , anySignedType - , anyUnsignedType - , anyIntType - , anyFloatType - , anyNumberType - , anyPrimType - - , Namespace(..) - , intrinsicsNameMap - , topLevelNameMap + ( TypeM, + runTypeM, + askEnv, + askImportName, + bindSpaced, + qualifyTypeVars, + lookupMTy, + lookupImport, + localEnv, + TypeError (..), + unappliedFunctor, + unknownVariable, + unknownType, + underscoreUse, + Notes, + aNote, + MonadTypeChecker (..), + checkName, + badOnLeft, + module Language.Futhark.Warnings, + Env (..), + TySet, + FunSig (..), + ImportTable, + NameMap, + BoundV (..), + Mod (..), + TypeBinding (..), + MTy (..), + anySignedType, + anyUnsignedType, + anyFloatType, + anyNumberType, + anyPrimType, + Namespace (..), + intrinsicsNameMap, + topLevelNameMap, ) where import Control.Monad.Except -import Control.Monad.Reader -import Control.Monad.Writer hiding (Sum) -import Control.Monad.State import Control.Monad.RWS.Strict hiding (Sum) -import Data.List (isPrefixOf, find) -import Data.Maybe import Data.Either +import Data.List (find, isPrefixOf) import qualified Data.Map.Strict as M +import Data.Maybe import qualified Data.Set as S - -import Prelude hiding (mapM, mod) - -import Language.Futhark -import Language.Futhark.Semantic -import Language.Futhark.Warnings import Futhark.FreshNames hiding (newName) import qualified Futhark.FreshNames -import Futhark.Util.Pretty hiding (space) import Futhark.Util.Console +import Futhark.Util.Pretty hiding (space) +import Language.Futhark +import Language.Futhark.Semantic +import Language.Futhark.Warnings +import Prelude hiding (mapM, mod) -- | A note with extra information regarding a type error. newtype Note = Note Doc @@ -85,7 +74,7 @@ ppr (Note msg) = "Note:" <+> align msg instance Pretty Notes where - ppr (Notes notes) = foldMap (((line<>line)<>) . ppr) notes + ppr (Notes notes) = foldMap (((line <> line) <>) . ppr) notes -- | A single note. aNote :: Pretty a => a -> Notes @@ -96,8 +85,8 @@ instance Pretty TypeError where ppr (TypeError loc notes msg) = - text (inRed $ "Error at " <> locStr loc <> ":") - msg <> ppr notes + text (inRed $ "Error at " <> locStr loc <> ":") + msg <> ppr notes -- | An unexpected functor appeared! unappliedFunctor :: MonadTypeChecker m => SrcLoc -> m a @@ -105,11 +94,15 @@ typeError loc mempty "Cannot have parametric module here." -- | An unknown variable was referenced. -unknownVariable :: MonadTypeChecker m => - Namespace -> QualName Name -> SrcLoc -> m a +unknownVariable :: + MonadTypeChecker m => + Namespace -> + QualName Name -> + SrcLoc -> + m a unknownVariable space name loc = typeError loc mempty $ - "Unknown" <+> ppr space <+> pquote (ppr name) + "Unknown" <+> ppr space <+> pquote (ppr name) -- | An unknown type was referenced. unknownType :: MonadTypeChecker m => SrcLoc -> QualName Name -> m a @@ -117,38 +110,54 @@ typeError loc mempty $ "Unknown type" <+> ppr name <> "." -- | A name prefixed with an underscore was used. -underscoreUse :: MonadTypeChecker m => - SrcLoc -> QualName Name -> m a +underscoreUse :: + MonadTypeChecker m => + SrcLoc -> + QualName Name -> + m a underscoreUse loc name = - typeError loc mempty $ "Use of" <+> pquote (ppr name) <> - ": variables prefixed with underscore may not be accessed." + typeError loc mempty $ + "Use of" <+> pquote (ppr name) + <> ": variables prefixed with underscore may not be accessed." -- | A mapping from import strings to 'Env's. This is used to resolve -- @import@ declarations. type ImportTable = M.Map String Env -data Context = Context { contextEnv :: Env - , contextImportTable :: ImportTable - , contextImportName :: ImportName - } +data Context = Context + { contextEnv :: Env, + contextImportTable :: ImportTable, + contextImportName :: ImportName + } -- | The type checker runs in this monad. -newtype TypeM a = TypeM (RWST - Context -- Reader - Warnings -- Writer - VNameSource -- State - (Except TypeError) -- Inner monad - a) - deriving (Monad, Functor, Applicative, - MonadReader Context, - MonadWriter Warnings, - MonadState VNameSource, - MonadError TypeError) +newtype TypeM a + = TypeM + ( RWST + Context -- Reader + Warnings -- Writer + VNameSource -- State + (Except TypeError) -- Inner monad + a + ) + deriving + ( Monad, + Functor, + Applicative, + MonadReader Context, + MonadWriter Warnings, + MonadState VNameSource, + MonadError TypeError + ) -- | Run a 'TypeM' computation. -runTypeM :: Env -> ImportTable -> ImportName -> VNameSource - -> TypeM a - -> Either TypeError (a, Warnings, VNameSource) +runTypeM :: + Env -> + ImportTable -> + ImportName -> + VNameSource -> + TypeM a -> + Either TypeError (a, Warnings, VNameSource) runTypeM env imports fpath src (TypeM m) = do (x, src', ws) <- runExcept $ runRWST m (Context env imports fpath) src return (x, ws, src') @@ -166,7 +175,8 @@ lookupMTy loc qn = do (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Signature qn loc (qn',) <$> maybe explode return (M.lookup name $ envSigTable scope) - where explode = unknownVariable Signature qn loc + where + explode = unknownVariable Signature qn loc -- | Look up an import. lookupImport :: SrcLoc -> FilePath -> TypeM (FilePath, Env) @@ -175,9 +185,10 @@ my_path <- asks contextImportName let canonical_import = includeToString $ mkImportFrom my_path file loc case M.lookup canonical_import imports of - Nothing -> typeError loc mempty $ - "Unknown import" <+> dquotes (text canonical_import) - "Known:" <+> commasep (map text (M.keys imports)) + Nothing -> + typeError loc mempty $ + "Unknown import" <+> dquotes (text canonical_import) + "Known:" <+> commasep (map text (M.keys imports)) Just scope -> return (canonical_import, scope) -- | Evaluate a 'TypeM' computation within an extended (/not/ @@ -185,7 +196,7 @@ localEnv :: Env -> TypeM a -> TypeM a localEnv env = local $ \ctx -> let env' = env <> contextEnv ctx - in ctx { contextEnv = env' } + in ctx {contextEnv = env'} -- | Monads that support type checking. The reason we have this -- internal interface is because we use distinct monads for checking @@ -209,9 +220,10 @@ checkNamedDim loc v = do (v', t) <- lookupVar loc v case t of - Scalar (Prim (Signed Int32)) -> return v' - _ -> typeError loc mempty $ - "Dimension declaration" <+> ppr v <+> "should be of type i32." + Scalar (Prim (Signed Int64)) -> return v' + _ -> + typeError loc mempty $ + "Dimension declaration" <+> ppr v <+> "should be of type i64." typeError :: Located loc => loc -> Notes -> Doc -> m a @@ -231,20 +243,25 @@ instance MonadTypeChecker TypeM where warn loc problem = tell $ singleWarning (srclocOf loc) problem - newName s = do src <- get - let (s', src') = Futhark.FreshNames.newName src s - put src' - return s' + newName s = do + src <- get + let (s', src') = Futhark.FreshNames.newName src s + put src' + return s' newID s = newName $ VName s 0 bindNameMap m = local $ \ctx -> let env = contextEnv ctx - in ctx { contextEnv = env { envNameMap = m <> envNameMap env } } + in ctx {contextEnv = env {envNameMap = m <> envNameMap env}} bindVal v t = local $ \ctx -> - ctx { contextEnv = (contextEnv ctx) - { envVtable = M.insert v t $ envVtable $ contextEnv ctx } } + ctx + { contextEnv = + (contextEnv ctx) + { envVtable = M.insert v t $ envVtable $ contextEnv ctx + } + } checkQualName space name loc = snd <$> checkQualNameWithEnv space name loc @@ -259,7 +276,7 @@ (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Term qn loc case M.lookup name $ envModTable scope of Nothing -> unknownVariable Term qn loc - Just m -> return (qn', m) + Just m -> return (qn', m) lookupVar loc qn = do outer_env <- askEnv @@ -269,19 +286,26 @@ Just (BoundV _ t) | "_" `isPrefixOf` baseString name -> underscoreUse loc qn | otherwise -> - case getType t of - Left{} -> typeError loc mempty $ - "Attempt to use function" <+> pprName name <+> "as value." - Right t' -> return (qn', fromStruct $ - qualifyTypeVars outer_env mempty qs t') + case getType t of + Left {} -> + typeError loc mempty $ + "Attempt to use function" <+> pprName name <+> "as value." + Right t' -> + return + ( qn', + fromStruct $ + qualifyTypeVars outer_env mempty qs t' + ) typeError loc notes s = throwError $ TypeError (srclocOf loc) notes s -- | Extract from a type either a function type comprising a list of -- parameter types and a return type, or a first-order type. -getType :: TypeBase dim as - -> Either ([(PName, TypeBase dim as)], TypeBase dim as) - (TypeBase dim as) +getType :: + TypeBase dim as -> + Either + ([(PName, TypeBase dim as)], TypeBase dim as) + (TypeBase dim as) getType (Scalar (Arrow _ v t1 t2)) = case getType t2 of Left (ps, r) -> Left ((v, t1) : ps, r) @@ -292,79 +316,89 @@ checkQualNameWithEnv space qn@(QualName quals name) loc = do env <- askEnv descend env quals - where descend scope [] - | Just name' <- M.lookup (space, name) $ envNameMap scope = - return (scope, name') - | otherwise = - unknownVariable space qn loc - - descend scope (q:qs) - | Just (QualName _ q') <- M.lookup (Term, q) $ envNameMap scope, - Just res <- M.lookup q' $ envModTable scope = - case res of - ModEnv q_scope -> do - (scope', QualName qs' name') <- descend q_scope qs - return (scope', QualName (q':qs') name') - ModFun{} -> unappliedFunctor loc - | otherwise = - unknownVariable space qn loc + where + descend scope [] + | Just name' <- M.lookup (space, name) $ envNameMap scope = + return (scope, name') + | otherwise = + unknownVariable space qn loc + descend scope (q : qs) + | Just (QualName _ q') <- M.lookup (Term, q) $ envNameMap scope, + Just res <- M.lookup q' $ envModTable scope = + case res of + ModEnv q_scope -> do + (scope', QualName qs' name') <- descend q_scope qs + return (scope', QualName (q' : qs') name') + ModFun {} -> unappliedFunctor loc + | otherwise = + unknownVariable space qn loc -- | Try to prepend qualifiers to the type names such that they -- represent how to access the type in some scope. -qualifyTypeVars :: Env -> [VName] -> [VName] -> TypeBase (DimDecl VName) as - -> TypeBase (DimDecl VName) as +qualifyTypeVars :: + Env -> + [VName] -> + [VName] -> + TypeBase (DimDecl VName) as -> + TypeBase (DimDecl VName) as qualifyTypeVars outer_env orig_except ref_qs = onType (S.fromList orig_except) - where onType :: S.Set VName -> TypeBase (DimDecl VName) as - -> TypeBase (DimDecl VName) as - onType except (Array as u et shape) = - Array as u (onScalar except et) (fmap (onDim except) shape) - onType except (Scalar t) = - Scalar $ onScalar except t - - onScalar _ (Prim t) = Prim t - onScalar except (TypeVar as u tn targs) = - TypeVar as u tn' $ map (onTypeArg except) targs - where tn' = typeNameFromQualName $ qual except $ qualNameFromTypeName tn - onScalar except (Record m) = - Record $ M.map (onType except) m - onScalar except (Sum m) = - Sum $ M.map (map $ onType except) m - onScalar except (Arrow as p t1 t2) = - Arrow as p (onType except' t1) (onType except' t2) - where except' = case p of Named p' -> S.insert p' except - Unnamed -> except + where + onType :: + S.Set VName -> + TypeBase (DimDecl VName) as -> + TypeBase (DimDecl VName) as + onType except (Array as u et shape) = + Array as u (onScalar except et) (fmap (onDim except) shape) + onType except (Scalar t) = + Scalar $ onScalar except t - onTypeArg except (TypeArgDim d loc) = - TypeArgDim (onDim except d) loc - onTypeArg except (TypeArgType t loc) = - TypeArgType (onType except t) loc + onScalar _ (Prim t) = Prim t + onScalar except (TypeVar as u tn targs) = + TypeVar as u tn' $ map (onTypeArg except) targs + where + tn' = typeNameFromQualName $ qual except $ qualNameFromTypeName tn + onScalar except (Record m) = + Record $ M.map (onType except) m + onScalar except (Sum m) = + Sum $ M.map (map $ onType except) m + onScalar except (Arrow as p t1 t2) = + Arrow as p (onType except' t1) (onType except' t2) + where + except' = case p of + Named p' -> S.insert p' except + Unnamed -> except - onDim except (NamedDim qn) = NamedDim $ qual except qn - onDim _ d = d + onTypeArg except (TypeArgDim d loc) = + TypeArgDim (onDim except d) loc + onTypeArg except (TypeArgType t loc) = + TypeArgType (onType except t) loc - qual except (QualName orig_qs name) - | name `elem` except || reachable orig_qs name outer_env = - QualName orig_qs name - | otherwise = - prependAsNecessary [] ref_qs $ QualName orig_qs name + onDim except (NamedDim qn) = NamedDim $ qual except qn + onDim _ d = d - prependAsNecessary qs rem_qs (QualName orig_qs name) - | reachable (qs++orig_qs) name outer_env = QualName (qs++orig_qs) name - | otherwise = case rem_qs of - q:rem_qs' -> prependAsNecessary (qs++[q]) rem_qs' (QualName orig_qs name) - [] -> QualName (qs++orig_qs) name + qual except (QualName orig_qs name) + | name `elem` except || reachable orig_qs name outer_env = + QualName orig_qs name + | otherwise = + prependAsNecessary [] ref_qs $ QualName orig_qs name - reachable [] name env = - name `M.member` envVtable env || - isJust (find matches $ M.elems (envTypeTable env)) - where matches (TypeAbbr _ _ (Scalar (TypeVar _ _ (TypeName x_qs name') _))) = - null x_qs && name == name' - matches _ = False + prependAsNecessary qs rem_qs (QualName orig_qs name) + | reachable (qs ++ orig_qs) name outer_env = QualName (qs ++ orig_qs) name + | otherwise = case rem_qs of + q : rem_qs' -> prependAsNecessary (qs ++ [q]) rem_qs' (QualName orig_qs name) + [] -> QualName (qs ++ orig_qs) name - reachable (q:qs') name env - | Just (ModEnv env') <- M.lookup q $ envModTable env = - reachable qs' name env' - | otherwise = False + reachable [] name env = + name `M.member` envVtable env + || isJust (find matches $ M.elems (envTypeTable env)) + where + matches (TypeAbbr _ _ (Scalar (TypeVar _ _ (TypeName x_qs name') _))) = + null x_qs && name == name' + matches _ = False + reachable (q : qs') name env + | Just (ModEnv env') <- M.lookup q $ envModTable env = + reachable qs' name env' + | otherwise = False -- | Turn a 'Left' 'TypeError' into an actual error. badOnLeft :: Either TypeError a -> TypeM a @@ -399,18 +433,24 @@ -- | The 'NameMap' corresponding to the intrinsics module. intrinsicsNameMap :: NameMap intrinsicsNameMap = M.fromList $ map mapping $ M.toList intrinsics - where mapping (v, IntrinsicType{}) = ((Type, baseName v), QualName [] v) - mapping (v, _) = ((Term, baseName v), QualName [] v) + where + mapping (v, IntrinsicType {}) = ((Type, baseName v), QualName [] v) + mapping (v, _) = ((Term, baseName v), QualName [] v) -- | The names that are available in the initial environment. topLevelNameMap :: NameMap topLevelNameMap = M.filterWithKey (\k _ -> atTopLevel k) intrinsicsNameMap - where atTopLevel :: (Namespace, Name) -> Bool - atTopLevel (Type, _) = True - atTopLevel (Term, v) = v `S.member` (type_names <> binop_names <> unop_names <> fun_names) - where type_names = S.fromList $ map (nameFromString . pretty) anyPrimType - binop_names = S.fromList $ map (nameFromString . pretty) - [minBound..(maxBound::BinOp)] - unop_names = S.fromList $ map nameFromString ["!"] - fun_names = S.fromList $ map nameFromString ["shape"] - atTopLevel _ = False + where + atTopLevel :: (Namespace, Name) -> Bool + atTopLevel (Type, _) = True + atTopLevel (Term, v) = v `S.member` (type_names <> binop_names <> unop_names <> fun_names) + where + type_names = S.fromList $ map (nameFromString . pretty) anyPrimType + binop_names = + S.fromList $ + map + (nameFromString . pretty) + [minBound .. (maxBound :: BinOp)] + unop_names = S.fromList $ map nameFromString ["!"] + fun_names = S.fromList $ map nameFromString ["shape"] + atTopLevel _ = False 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 @@ -1,3024 +1,3311 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances, DeriveFunctor #-} -{-# Language TupleSections #-} -{-# Language OverloadedStrings #-} -{-# LANGUAGE Trustworthy #-} --- | Facilities for type-checking Futhark terms. Checking a term --- requires a little more context to track uniqueness and such. --- --- Type inference is implemented through a variation of --- Hindley-Milner. The main complication is supporting the rich --- number of built-in language constructs, as well as uniqueness --- types. This is mostly done in an ad hoc way, and many programs --- will require the programmer to fall back on type annotations. -module Language.Futhark.TypeChecker.Terms - ( checkOneExp - , checkFunDef - ) -where - -import Control.Monad.Identity -import Control.Monad.Except -import Control.Monad.State -import Control.Monad.RWS hiding (Sum) -import Control.Monad.Writer hiding (Sum) -import Data.Bifunctor -import Data.Char (isAscii) -import Data.Either -import Data.List (isPrefixOf, foldl', find, (\\), nub, transpose, sort, group) -import qualified Data.List.NonEmpty as NE -import Data.Maybe -import qualified Data.Map.Strict as M -import qualified Data.Set as S - -import Prelude hiding (mod) - -import Language.Futhark hiding (unscopeType) -import Language.Futhark.Semantic (includeToString) -import Language.Futhark.Traversals -import Language.Futhark.TypeChecker.Monad hiding (BoundV) -import Language.Futhark.TypeChecker.Types hiding (checkTypeDecl) -import Language.Futhark.TypeChecker.Unify hiding (Usage) -import qualified Language.Futhark.TypeChecker.Types as Types -import qualified Language.Futhark.TypeChecker.Monad as TypeM -import Futhark.IR.Primitive (intByteSize) -import Futhark.Util.Pretty hiding (space, bool, group) - ---- Uniqueness - -data Usage = Consumed SrcLoc - | Observed SrcLoc - deriving (Eq, Ord, Show) - -type Names = S.Set VName - --- | The consumption set is a Maybe so we can distinguish whether a --- consumption took place, but the variable went out of scope since, --- or no consumption at all took place. -data Occurence = Occurence { observed :: Names - , consumed :: Maybe Names - , location :: SrcLoc - } - deriving (Eq, Show) - -instance Located Occurence where - locOf = locOf . location - -observation :: Aliasing -> SrcLoc -> Occurence -observation = flip Occurence Nothing . S.map aliasVar - -consumption :: Aliasing -> SrcLoc -> Occurence -consumption = Occurence S.empty . Just . S.map aliasVar - --- | A null occurence is one that we can remove without affecting --- anything. -nullOccurence :: Occurence -> Bool -nullOccurence occ = S.null (observed occ) && isNothing (consumed occ) - --- | A seminull occurence is one that does not contain references to --- any variables in scope. The big difference is that a seminull --- occurence may denote a consumption, as long as the array that was --- consumed is now out of scope. -seminullOccurence :: Occurence -> Bool -seminullOccurence occ = S.null (observed occ) && maybe True S.null (consumed occ) - -type Occurences = [Occurence] - -type UsageMap = M.Map VName [Usage] - -usageMap :: Occurences -> UsageMap -usageMap = foldl comb M.empty - where comb m (Occurence obs cons loc) = - let m' = S.foldl' (ins $ Observed loc) m obs - in S.foldl' (ins $ Consumed loc) m' $ fromMaybe mempty cons - ins v m k = M.insertWith (++) k [v] m - -combineOccurences :: VName -> Usage -> Usage -> TermTypeM Usage -combineOccurences _ (Observed loc) (Observed _) = return $ Observed loc -combineOccurences name (Consumed wloc) (Observed rloc) = - useAfterConsume (baseName name) rloc wloc -combineOccurences name (Observed rloc) (Consumed wloc) = - useAfterConsume (baseName name) rloc wloc -combineOccurences name (Consumed loc1) (Consumed loc2) = - consumeAfterConsume (baseName name) (max loc1 loc2) (min loc1 loc2) - -checkOccurences :: Occurences -> TermTypeM () -checkOccurences = void . M.traverseWithKey comb . usageMap - where comb _ [] = return () - comb name (u:us) = foldM_ (combineOccurences name) u us - -allObserved :: Occurences -> Names -allObserved = S.unions . map observed - -allConsumed :: Occurences -> Names -allConsumed = S.unions . map (fromMaybe mempty . consumed) - -allOccuring :: Occurences -> Names -allOccuring occs = allConsumed occs <> allObserved occs - -anyConsumption :: Occurences -> Maybe Occurence -anyConsumption = find (isJust . consumed) - -seqOccurences :: Occurences -> Occurences -> Occurences -seqOccurences occurs1 occurs2 = - filter (not . nullOccurence) $ map filt occurs1 ++ occurs2 - where filt occ = - occ { observed = observed occ `S.difference` postcons } - postcons = allConsumed occurs2 - -altOccurences :: Occurences -> Occurences -> Occurences -altOccurences occurs1 occurs2 = - filter (not . nullOccurence) $ map filt1 occurs1 ++ map filt2 occurs2 - where filt1 occ = - occ { consumed = S.difference <$> consumed occ <*> pure cons2 - , observed = observed occ `S.difference` cons2 } - filt2 occ = - occ { consumed = consumed occ - , observed = observed occ `S.difference` cons1 } - cons1 = allConsumed occurs1 - cons2 = allConsumed occurs2 - ---- Scope management - -data Checking - = CheckingApply (Maybe (QualName VName)) Exp StructType StructType - | CheckingReturn StructType StructType - | CheckingAscription StructType StructType - | CheckingLetGeneralise Name - | CheckingParams (Maybe Name) - | CheckingPattern UncheckedPattern InferredType - | CheckingLoopBody StructType StructType - | CheckingLoopInitial StructType StructType - | CheckingRecordUpdate [Name] StructType StructType - | CheckingRequired [StructType] StructType - | CheckingBranches StructType StructType - -instance Pretty Checking where - ppr (CheckingApply f e expected actual) = - header - "Expected:" <+> align (ppr expected) - "Actual: " <+> align (ppr actual) - where header = - case f of - Nothing -> - "Cannot apply function to" <+> - pquote (shorten $ pretty $ flatten $ ppr e) <> " (invalid type)." - Just fname -> - "Cannot apply" <+> pquote (ppr fname) <+> "to" <+> - pquote (shorten $ pretty $ flatten $ ppr e) <> " (invalid type)." - - ppr (CheckingReturn expected actual) = - "Function body does not have expected type." - "Expected:" <+> align (ppr expected) - "Actual: " <+> align (ppr actual) - - ppr (CheckingAscription expected actual) = - "Expression does not have expected type from explicit ascription." - "Expected:" <+> align (ppr expected) - "Actual: " <+> align (ppr actual) - - ppr (CheckingLetGeneralise fname) = - "Cannot generalise type of" <+> pquote (ppr fname) <> "." - - ppr (CheckingParams fname) = - "Invalid use of parameters in" <+> pquote fname' <> "." - where fname' = maybe "anonymous function" ppr fname - - ppr (CheckingPattern pat NoneInferred) = - "Invalid pattern" <+> pquote (ppr pat) <> "." - - ppr (CheckingPattern pat (Ascribed t)) = - "Pattern" <+> pquote (ppr pat) <+> - "cannot match value of type" - indent 2 (ppr t) - - ppr (CheckingLoopBody expected actual) = - "Loop body does not have expected type." - "Expected:" <+> align (ppr expected) - "Actual: " <+> align (ppr actual) - - ppr (CheckingLoopInitial expected actual) = - "Initial loop values do not have expected type." - "Expected:" <+> align (ppr expected) - "Actual: " <+> align (ppr actual) - - ppr (CheckingRecordUpdate fs expected actual) = - "Type mismatch when updating record field" <+> pquote fs' <> "." - "Existing:" <+> align (ppr expected) - "New: " <+> align (ppr actual) - where fs' = mconcat $ punctuate "." $ map ppr fs - - ppr (CheckingRequired [expected] actual) = - "Expression must must have type" <+> ppr expected <> "." - "Actual type:" <+> align (ppr actual) - - ppr (CheckingRequired expected actual) = - "Type of expression must must be one of " <+> expected' <> "." - "Actual type:" <+> align (ppr actual) - where expected' = commasep (map ppr expected) - - ppr (CheckingBranches t1 t2) = - "Conditional branches differ in type." - "Former:" <+> ppr t1 - "Latter:" <+> ppr t2 - --- | Whether something is a global or a local variable. -data Locality = Local | Global - deriving (Show) - -data ValBinding = BoundV Locality [TypeParam] PatternType - -- ^ Aliases in parameters indicate the lexical - -- closure. - | OverloadedF [PrimType] [Maybe PrimType] (Maybe PrimType) - | EqualityF - | WasConsumed SrcLoc - deriving (Show) - --- | Type checking happens with access to this environment. The --- 'TermScope' will be extended during type-checking as bindings come into --- scope. -data TermEnv = TermEnv { termScope :: TermScope - , termChecking :: Maybe Checking - , termLevel :: Level - } - -data TermScope = TermScope { scopeVtable :: M.Map VName ValBinding - , scopeTypeTable :: M.Map VName TypeBinding - , scopeModTable :: M.Map VName Mod - , scopeNameMap :: NameMap - } deriving (Show) - -instance Semigroup TermScope where - TermScope vt1 tt1 mt1 nt1 <> TermScope vt2 tt2 mt2 nt2 = - TermScope (vt2 `M.union` vt1) (tt2 `M.union` tt1) (mt1 `M.union` mt2) (nt2 `M.union` nt1) - -envToTermScope :: Env -> TermScope -envToTermScope env = TermScope { scopeVtable = vtable - , scopeTypeTable = envTypeTable env - , scopeNameMap = envNameMap env - , scopeModTable = envModTable env - } - where vtable = M.mapWithKey valBinding $ envVtable env - valBinding k (TypeM.BoundV tps v) = - BoundV Global tps $ v `setAliases` - (if arrayRank v > 0 then S.singleton (AliasBound k) else mempty) - -withEnv :: TermEnv -> Env -> TermEnv -withEnv tenv env = tenv { termScope = termScope tenv <> envToTermScope env } - -overloadedTypeVars :: Constraints -> Names -overloadedTypeVars = mconcat . map f . M.elems - where f (_, HasFields fs _) = mconcat $ map typeVars $ M.elems fs - f _ = mempty - --- | Get the type of an expression, with top level type variables --- substituted. Never call 'typeOf' directly (except in a few --- carefully inspected locations)! -expType :: Exp -> TermTypeM PatternType -expType = normPatternType . typeOf - --- | Get the type of an expression, with all type variables --- substituted. Slower than 'expType', but sometimes necessary. --- Never call 'typeOf' directly (except in a few carefully inspected --- locations)! -expTypeFully :: Exp -> TermTypeM PatternType -expTypeFully = normTypeFully . typeOf - --- Wrap a function name to give it a vacuous Eq instance for SizeSource. -newtype FName = FName (Maybe (QualName VName)) - deriving (Show) - -instance Eq FName where - _ == _ = True - -instance Ord FName where - compare _ _ = EQ - --- | What was the source of some existential size? This is used for --- using the same existential variable if the same source is --- encountered in multiple locations. -data SizeSource = SourceArg FName (ExpBase NoInfo VName) - | SourceBound (ExpBase NoInfo VName) - | SourceSlice - (Maybe (DimDecl VName)) - (Maybe (ExpBase NoInfo VName)) - (Maybe (ExpBase NoInfo VName)) - (Maybe (ExpBase NoInfo VName)) - deriving (Eq, Ord, Show) - --- | The state is a set of constraints and a counter for generating --- type names. This is distinct from the usual counter we use for --- generating unique names, as these will be user-visible. -data TermTypeState = TermTypeState - { stateConstraints :: Constraints - , stateCounter :: !Int - , stateDimTable :: M.Map SizeSource VName - -- ^ Mapping function arguments encountered to - -- the sizes they ended up generating (when - -- they could not be substituted directly). - -- This happens for function arguments that are - -- not constants or names. - } - -newtype TermTypeM a = TermTypeM (RWST - TermEnv - Occurences - TermTypeState - TypeM - a) - deriving (Monad, Functor, Applicative, - MonadReader TermEnv, - MonadWriter Occurences, - MonadState TermTypeState, - MonadError TypeError) - -instance MonadUnify TermTypeM where - getConstraints = gets stateConstraints - putConstraints x = modify $ \s -> s { stateConstraints = x } - - newTypeVar loc desc = do - i <- incCounter - v <- newID $ mkTypeVarName desc i - constrain v $ NoConstraint Lifted $ mkUsage' loc - return $ Scalar $ TypeVar mempty Nonunique (typeName v) [] - - curLevel = asks termLevel - - newDimVar loc rigidity name = do - i <- incCounter - dim <- newID $ mkTypeVarName name i - case rigidity of - Rigid rsrc -> constrain dim $ UnknowableSize loc rsrc - Nonrigid -> constrain dim $ Size Nothing $ mkUsage' loc - return dim - - unifyError loc notes bcs doc = do - checking <- asks termChecking - case checking of - Just checking' -> - throwError $ TypeError (srclocOf loc) notes $ - ppr checking' <> line doc <> ppr bcs - Nothing -> - throwError $ TypeError (srclocOf loc) notes $ doc <> ppr bcs - - matchError loc notes bcs t1 t2 = do - checking <- asks termChecking - case checking of - Just checking' - | hasNoBreadCrumbs bcs -> - throwError $ TypeError (srclocOf loc) notes $ - ppr checking' - | otherwise -> - throwError $ TypeError (srclocOf loc) notes $ - ppr checking' <> line doc <> ppr bcs - Nothing -> - throwError $ TypeError (srclocOf loc) notes $ doc <> ppr bcs - where doc = "Types" - indent 2 (ppr t1) - "and" - indent 2 (ppr t2) - "do not match." - -onFailure :: Checking -> TermTypeM a -> TermTypeM a -onFailure c = local $ \env -> env { termChecking = Just c } - -runTermTypeM :: TermTypeM a -> TypeM (a, Occurences) -runTermTypeM (TermTypeM m) = do - initial_scope <- (initialTermScope <>) . envToTermScope <$> askEnv - let initial_tenv = TermEnv { termScope = initial_scope - , termChecking = Nothing - , termLevel = 0 - } - evalRWST m initial_tenv $ TermTypeState mempty 0 mempty - -liftTypeM :: TypeM a -> TermTypeM a -liftTypeM = TermTypeM . lift - -localScope :: (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a -localScope f = local $ \tenv -> tenv { termScope = f $ termScope tenv } - -incCounter :: TermTypeM Int -incCounter = do s <- get - put s { stateCounter = stateCounter s + 1 } - return $ stateCounter s - -extSize :: SrcLoc -> SizeSource -> TermTypeM (DimDecl VName, Maybe VName) -extSize loc e = do - prev <- gets $ M.lookup e . stateDimTable - case prev of - Nothing -> do - let rsrc = case e of - SourceArg (FName fname) e' -> - RigidArg fname $ prettyOneLine e' - SourceBound e' -> - RigidBound $ prettyOneLine e' - SourceSlice d i j s -> - RigidSlice d $ prettyOneLine $ DimSlice i j s - d <- newDimVar loc (Rigid rsrc) "argdim" - modify $ \s -> s { stateDimTable = M.insert e d $ stateDimTable s } - return (NamedDim $ qualName d, - Just d) - Just d -> return (NamedDim $ qualName d, - Nothing) - --- Any argument sizes created with 'extSize' inside the given action --- will be removed once the action finishes. This is to ensure that --- just because e.g. @n+1@ appears as a size in one branch of a --- conditional, that doesn't mean it's also available in the other branch. -noSizeEscape :: TermTypeM a -> TermTypeM a -noSizeEscape m = do - dimtable <- gets stateDimTable - x <- m - modify $ \s -> s { stateDimTable = dimtable } - return x - -constrain :: VName -> Constraint -> TermTypeM () -constrain v c = do - lvl <- curLevel - modifyConstraints $ M.insert v (lvl, c) - -incLevel :: TermTypeM a -> TermTypeM a -incLevel = local $ \env -> env { termLevel = termLevel env + 1 } - -initialTermScope :: TermScope -initialTermScope = TermScope { scopeVtable = initialVtable - , scopeTypeTable = mempty - , scopeNameMap = topLevelNameMap - , scopeModTable = mempty - } - where initialVtable = M.fromList $ mapMaybe addIntrinsicF $ M.toList intrinsics - - prim = Scalar . Prim - arrow x y = Scalar $ Arrow mempty Unnamed x y - - addIntrinsicF (name, IntrinsicMonoFun pts t) = - Just (name, BoundV Global [] $ arrow pts' $ prim t) - where pts' = case pts of [pt] -> prim pt - _ -> tupleRecord $ map prim pts - - addIntrinsicF (name, IntrinsicOverloadedFun ts pts rts) = - Just (name, OverloadedF ts pts rts) - addIntrinsicF (name, IntrinsicPolyFun tvs pts rt) = - Just (name, BoundV Global tvs $ - fromStruct $ Scalar $ Arrow mempty Unnamed pts' rt) - where pts' = case pts of [pt] -> pt - _ -> tupleRecord pts - addIntrinsicF (name, IntrinsicEquality) = - Just (name, EqualityF) - addIntrinsicF _ = Nothing - -instance MonadTypeChecker TermTypeM where - warn loc problem = liftTypeM $ warn loc problem - newName = liftTypeM . newName - newID = liftTypeM . newID - - checkQualName space name loc = snd <$> checkQualNameWithEnv space name loc - - bindNameMap m = localScope $ \scope -> - scope { scopeNameMap = m <> scopeNameMap scope } - - bindVal v (TypeM.BoundV tps t) = localScope $ \scope -> - scope { scopeVtable = M.insert v vb $ scopeVtable scope } - where vb = BoundV Local tps $ fromStruct t - - lookupType loc qn = do - outer_env <- liftTypeM askEnv - (scope, qn'@(QualName qs name)) <- checkQualNameWithEnv Type qn loc - case M.lookup name $ scopeTypeTable scope of - Nothing -> unknownType loc qn - Just (TypeAbbr l ps def) -> - return (qn', ps, qualifyTypeVars outer_env (map typeParamName ps) qs def, l) - - lookupMod loc qn = do - (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Term qn loc - case M.lookup name $ scopeModTable scope of - Nothing -> unknownVariable Term qn loc - Just m -> return (qn', m) - - lookupVar loc qn = do - outer_env <- liftTypeM askEnv - (scope, qn'@(QualName qs name)) <- checkQualNameWithEnv Term qn loc - let usage = mkUsage loc $ "use of " ++ quote (pretty qn) - - t <- case M.lookup name $ scopeVtable scope of - Nothing -> typeError loc mempty $ - "Unknown variable" <+> pquote (ppr qn) <> "." - - Just (WasConsumed wloc) -> useAfterConsume (baseName name) loc wloc - - Just (BoundV _ tparams t) - | "_" `isPrefixOf` baseString name -> underscoreUse loc qn - | otherwise -> do - (tnames, t') <- instantiateTypeScheme loc tparams t - return $ qualifyTypeVars outer_env tnames qs t' - - Just EqualityF -> do - argtype <- newTypeVar loc "t" - equalityType usage argtype - return $ - Scalar $ Arrow mempty Unnamed argtype $ - Scalar $ Arrow mempty Unnamed argtype $ Scalar $ Prim Bool - - Just (OverloadedF ts pts rt) -> do - argtype <- newTypeVar loc "t" - mustBeOneOf ts usage argtype - let (pts', rt') = instOverloaded argtype pts rt - arrow xt yt = Scalar $ Arrow mempty Unnamed xt yt - return $ fromStruct $ foldr arrow rt' pts' - - observe $ Ident name (Info t) loc - return (qn', t) - - where instOverloaded argtype pts rt = - (map (maybe (toStruct argtype) (Scalar . Prim)) pts, - maybe (toStruct argtype) (Scalar . Prim) rt) - - checkNamedDim loc v = do - (v', t) <- lookupVar loc v - onFailure (CheckingRequired [Scalar $ Prim $ Signed Int32] (toStruct t)) $ - unify (mkUsage loc "use as array size") (toStruct t) $ - Scalar $ Prim $ Signed Int32 - return v' - - typeError loc notes s = do - checking <- asks termChecking - case checking of - Just checking' -> - throwError $ TypeError (srclocOf loc) notes (ppr checking' <> line s) - Nothing -> - throwError $ TypeError (srclocOf loc) notes s - -checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) -checkQualNameWithEnv space qn@(QualName quals name) loc = do - scope <- asks termScope - descend scope quals - where descend scope [] - | Just name' <- M.lookup (space, name) $ scopeNameMap scope = - return (scope, name') - | otherwise = - unknownVariable space qn loc - - descend scope (q:qs) - | Just (QualName _ q') <- M.lookup (Term, q) $ scopeNameMap scope, - Just res <- M.lookup q' $ scopeModTable scope = - case res of - -- Check if we are referring to the magical intrinsics - -- module. - _ | baseTag q' <= maxIntrinsicTag -> - checkIntrinsic space qn loc - ModEnv q_scope -> do - (scope', QualName qs' name') <- descend (envToTermScope q_scope) qs - return (scope', QualName (q':qs') name') - ModFun{} -> unappliedFunctor loc - | otherwise = - unknownVariable space qn loc - -checkIntrinsic :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) -checkIntrinsic space qn@(QualName _ name) loc - | Just v <- M.lookup (space, name) intrinsicsNameMap = do - me <- liftTypeM askImportName - unless ("/prelude" `isPrefixOf` includeToString me) $ - warn loc "Using intrinsic functions directly can easily crash the compiler or result in wrong code generation." - scope <- asks termScope - return (scope, v) - | otherwise = - unknownVariable space qn loc - --- | Wrap 'Types.checkTypeDecl' to also perform an observation of --- every size in the type. -checkTypeDecl :: TypeDeclBase NoInfo Name -> TermTypeM (TypeDeclBase Info VName) -checkTypeDecl tdecl = do - (tdecl', _) <- Types.checkTypeDecl tdecl - mapM_ observeDim $ nestedDims $ unInfo $ expandedType tdecl' - return tdecl' - where observeDim (NamedDim v) = - observe $ Ident (qualLeaf v) (Info $ Scalar $ Prim $ Signed Int32) mempty - observeDim _ = return () - --- | Instantiate a type scheme with fresh type variables for its type --- parameters. Returns the names of the fresh type variables, the --- instance list, and the instantiated type. -instantiateTypeScheme :: SrcLoc -> [TypeParam] -> PatternType - -> TermTypeM ([VName], PatternType) -instantiateTypeScheme loc tparams t = do - let tnames = map typeParamName tparams - (tparam_names, tparam_substs) <- unzip <$> mapM (instantiateTypeParam loc) tparams - let substs = M.fromList $ zip tnames tparam_substs - t' = applySubst (`M.lookup` substs) t - return (tparam_names, t') - --- | Create a new type name and insert it (unconstrained) in the --- substitution map. -instantiateTypeParam :: Monoid as => SrcLoc -> TypeParam -> TermTypeM (VName, Subst (TypeBase dim as)) -instantiateTypeParam loc tparam = do - i <- incCounter - v <- newID $ mkTypeVarName (takeWhile isAscii (baseString (typeParamName tparam))) i - case tparam of TypeParamType x _ _ -> do - constrain v $ NoConstraint x $ mkUsage' loc - return (v, Subst $ Scalar $ TypeVar mempty Nonunique (typeName v) []) - TypeParamDim{} -> do - constrain v $ Size Nothing $ mkUsage' loc - return (v, SizeSubst $ NamedDim $ qualName v) - -newArrayType :: SrcLoc -> String -> Int -> TermTypeM (StructType, StructType) -newArrayType loc desc r = do - v <- newID $ nameFromString desc - constrain v $ NoConstraint Unlifted $ mkUsage' loc - dims <- replicateM r $ newDimVar loc Nonrigid "dim" - let rowt = TypeVar () Nonunique (typeName v) [] - return (Array () Nonunique rowt (ShapeDecl $ map (NamedDim . qualName) dims), - Scalar rowt) - ---- Errors - -useAfterConsume :: Name -> SrcLoc -> SrcLoc -> TermTypeM a -useAfterConsume name rloc wloc = - typeError rloc mempty $ - "Variable" <+> pquote (pprName name) <+> "previously consumed at" <+> - text (locStrRel rloc wloc) <> ". (Possibly through aliasing.)" - -consumeAfterConsume :: Name -> SrcLoc -> SrcLoc -> TermTypeM a -consumeAfterConsume name loc1 loc2 = - typeError loc2 mempty $ - "Variable" <+> pprName name <+> "previously consumed at" <+> - text (locStrRel loc2 loc1) <> "." - -badLetWithValue :: SrcLoc -> TermTypeM a -badLetWithValue loc = - typeError loc mempty - "New value for elements in let-with shares data with source array. This is illegal, as it prevents in-place modification." - -returnAliased :: Name -> Name -> SrcLoc -> TermTypeM () -returnAliased fname name loc = - typeError loc mempty $ - "Unique return value of" <+> pquote (pprName fname) <+> - "is aliased to" <+> pquote (pprName name) <> ", which is not consumed." - -uniqueReturnAliased :: Name -> SrcLoc -> TermTypeM a -uniqueReturnAliased fname loc = - typeError loc mempty $ - "A unique tuple element of return value of" <+> - pquote (pprName fname) <+> "is aliased to some other tuple component." - -unexpectedType :: MonadTypeChecker m => SrcLoc -> StructType -> [StructType] -> m a -unexpectedType loc _ [] = - typeError loc mempty $ - "Type of expression at" <+> text (locStr loc) <+> - "cannot have any type - possibly a bug in the type checker." -unexpectedType loc t ts = - typeError loc mempty $ - "Type of expression at" <+> text (locStr loc) <+> "must be one of" <+> - commasep (map ppr ts) <> ", but is" <+> - ppr t <> "." - ---- Basic checking - --- | Determine if the two types of identical, ignoring uniqueness. --- Mismatched dimensions are turned into fresh rigid type variables. --- Causes a 'TypeError' if they fail to match, and otherwise returns --- one of them. -unifyBranchTypes :: SrcLoc -> PatternType -> PatternType -> TermTypeM (PatternType, [VName]) -unifyBranchTypes loc t1 t2 = - onFailure (CheckingBranches (toStruct t1) (toStruct t2)) $ - unifyMostCommon (mkUsage loc "unification of branch results") t1 t2 - -unifyBranches :: SrcLoc -> Exp -> Exp -> TermTypeM (PatternType, [VName]) -unifyBranches loc e1 e2 = do - e1_t <- expTypeFully e1 - e2_t <- expTypeFully e2 - unifyBranchTypes loc e1_t e2_t - ---- General binding. - -doNotShadow :: [String] -doNotShadow = ["&&", "||"] - -data InferredType = NoneInferred - | Ascribed PatternType - - -checkPattern' :: UncheckedPattern -> InferredType - -> TermTypeM Pattern - -checkPattern' (PatternParens p loc) t = - PatternParens <$> checkPattern' p t <*> pure loc - -checkPattern' (Id name _ loc) _ - | name' `elem` doNotShadow = - typeError loc mempty $ "The" <+> text name' <+> "operator may not be redefined." - where name' = nameToString name - -checkPattern' (Id name NoInfo loc) (Ascribed t) = do - name' <- newID name - return $ Id name' (Info t) loc -checkPattern' (Id name NoInfo loc) NoneInferred = do - name' <- newID name - t <- newTypeVar loc "t" - return $ Id name' (Info t) loc - -checkPattern' (Wildcard _ loc) (Ascribed t) = - return $ Wildcard (Info $ t `setUniqueness` Nonunique) loc -checkPattern' (Wildcard NoInfo loc) NoneInferred = do - t <- newTypeVar loc "t" - return $ Wildcard (Info t) loc - -checkPattern' (TuplePattern ps loc) (Ascribed t) - | Just ts <- isTupleRecord t, length ts == length ps = - TuplePattern <$> zipWithM checkPattern' ps (map Ascribed ts) <*> pure loc -checkPattern' p@(TuplePattern ps loc) (Ascribed t) = do - ps_t <- replicateM (length ps) (newTypeVar loc "t") - unify (mkUsage loc "matching a tuple pattern") (tupleRecord ps_t) $ toStruct t - t' <- normTypeFully t - checkPattern' p $ Ascribed t' -checkPattern' (TuplePattern ps loc) NoneInferred = - TuplePattern <$> mapM (`checkPattern'` NoneInferred) ps <*> pure loc - -checkPattern' (RecordPattern p_fs _) _ - | Just (f, fp) <- find (("_" `isPrefixOf`) . nameToString . fst) p_fs = - typeError fp mempty $ - "Underscore-prefixed fields are not allowed." - "Did you mean" <> dquotes (text (drop 1 (nameToString f)) <> "=_") <> "?" - -checkPattern' (RecordPattern p_fs loc) (Ascribed (Scalar (Record t_fs))) - | sort (map fst p_fs) == sort (M.keys t_fs) = - RecordPattern . M.toList <$> check <*> pure loc - where check = traverse (uncurry checkPattern') $ M.intersectionWith (,) - (M.fromList p_fs) (fmap Ascribed t_fs) -checkPattern' p@(RecordPattern fields loc) (Ascribed t) = do - fields' <- traverse (const $ newTypeVar loc "t") $ M.fromList fields - - when (sort (M.keys fields') /= sort (map fst fields)) $ - typeError loc mempty $ "Duplicate fields in record pattern" <+> ppr p <> "." - - unify (mkUsage loc "matching a record pattern") (Scalar (Record fields')) $ toStruct t - t' <- normTypeFully t - checkPattern' p $ Ascribed t' -checkPattern' (RecordPattern fs loc) NoneInferred = - RecordPattern . M.toList <$> traverse (`checkPattern'` NoneInferred) (M.fromList fs) <*> pure loc - -checkPattern' (PatternAscription p (TypeDecl t NoInfo) loc) maybe_outer_t = do - (t', st_nodims, _) <- checkTypeExp t - (st, _) <- instantiateEmptyArrayDims loc "impl" Nonrigid st_nodims - - let st' = fromStruct st - case maybe_outer_t of - Ascribed outer_t -> do - unify (mkUsage loc "explicit type ascription") (toStruct st) (toStruct outer_t) - - -- We also have to make sure that uniqueness matches. This is - -- done explicitly, because it is ignored by unification. - st'' <- normTypeFully st' - outer_t' <- normTypeFully outer_t - case unifyTypesU unifyUniqueness st'' outer_t' of - Just outer_t'' -> - PatternAscription <$> checkPattern' p (Ascribed outer_t'') <*> - pure (TypeDecl t' (Info st)) <*> pure loc - Nothing -> - typeError loc mempty $ - "Cannot match type" <+> pquote (ppr outer_t') <+> "with expected type" <+> - pquote (ppr st'') <> "." - - NoneInferred -> - PatternAscription <$> checkPattern' p (Ascribed st') <*> - pure (TypeDecl t' (Info st)) <*> pure loc - where unifyUniqueness u1 u2 = if u2 `subuniqueOf` u1 then Just u1 else Nothing - -checkPattern' (PatternLit e NoInfo loc) (Ascribed t) = do - e' <- checkExp e - t' <- expTypeFully e' - unify (mkUsage loc "matching against literal") (toStruct t') (toStruct t) - return $ PatternLit e' (Info t') loc - -checkPattern' (PatternLit e NoInfo loc) NoneInferred = do - e' <- checkExp e - t' <- expTypeFully e' - return $ PatternLit e' (Info t') loc - -checkPattern' (PatternConstr n NoInfo ps loc) (Ascribed (Scalar (Sum cs))) - | Just ts <- M.lookup n cs = do - ps' <- zipWithM checkPattern' ps $ map Ascribed ts - return $ PatternConstr n (Info (Scalar (Sum cs))) ps' loc - -checkPattern' (PatternConstr n NoInfo ps loc) (Ascribed t) = do - t' <- newTypeVar loc "t" - ps' <- mapM (`checkPattern'` NoneInferred) ps - mustHaveConstr usage n t' (patternStructType <$> ps') - unify usage t' (toStruct t) - t'' <- normTypeFully t - return $ PatternConstr n (Info t'') ps' loc - where usage = mkUsage loc "matching against constructor" - -checkPattern' (PatternConstr n NoInfo ps loc) NoneInferred = do - ps' <- mapM (`checkPattern'` NoneInferred) ps - t <- newTypeVar loc "t" - mustHaveConstr usage n t (patternStructType <$> ps') - return $ PatternConstr n (Info $ fromStruct t) ps' loc - where usage = mkUsage loc "matching against constructor" - -patternNameMap :: Pattern -> NameMap -patternNameMap = M.fromList . map asTerm . S.toList . patternNames - where asTerm v = ((Term, baseName v), qualName v) - -checkPattern :: UncheckedPattern -> InferredType -> (Pattern -> TermTypeM a) - -> TermTypeM a -checkPattern p t m = do - checkForDuplicateNames [p] - p' <- onFailure (CheckingPattern p t) $ checkPattern' p t - bindNameMap (patternNameMap p') $ m p' - -binding :: [Ident] -> TermTypeM a -> TermTypeM a -binding bnds = check . handleVars - where handleVars m = - localScope (`bindVars` bnds) $ do - - -- Those identifiers that can potentially also be sizes are - -- added as type constraints. This is necessary so that we - -- can properly detect scope violations during unification. - -- We do this for *all* identifiers, not just those that are - -- integers, because they may become integers later due to - -- inference... - forM_ bnds $ \ident -> - constrain (identName ident) $ ParamSize $ srclocOf ident - m - - bindVars :: TermScope -> [Ident] -> TermScope - bindVars = foldl bindVar - - bindVar :: TermScope -> Ident -> TermScope - bindVar scope (Ident name (Info tp) _) = - let inedges = boundAliases $ aliases tp - update (BoundV l tparams in_t) - -- If 'name' is record or sum-typed, don't alias the - -- components to 'name', because these no identity - -- beyond their components. - | Array{} <- tp = BoundV l tparams (in_t `addAliases` S.insert (AliasBound name)) - | otherwise = BoundV l tparams in_t - update b = b - - tp' = tp `addAliases` S.insert (AliasBound name) - in scope { scopeVtable = M.insert name (BoundV Local [] tp') $ - adjustSeveral update inedges $ - scopeVtable scope - } - - adjustSeveral f = flip $ foldl $ flip $ M.adjust f - - -- Check whether the bound variables have been used correctly - -- within their scope. - check m = do - (a, usages) <- collectBindingsOccurences m - checkOccurences usages - - mapM_ (checkIfUsed usages) bnds - - return a - - -- Collect and remove all occurences in @bnds@. This relies - -- on the fact that no variables shadow any other. - collectBindingsOccurences m = pass $ do - (x, usage) <- listen m - let (relevant, rest) = split usage - return ((x, relevant), const rest) - where split = unzip . - map (\occ -> - let (obs1, obs2) = divide $ observed occ - occ_cons = divide <$> consumed occ - con1 = fst <$> occ_cons - con2 = snd <$> occ_cons - in (occ { observed = obs1, consumed = con1 }, - occ { observed = obs2, consumed = con2 })) - names = S.fromList $ map identName bnds - divide s = (s `S.intersection` names, s `S.difference` names) - -bindingTypes :: [Either (VName, TypeBinding) (VName, Constraint)] - -> TermTypeM a -> TermTypeM a -bindingTypes types m = do - lvl <- curLevel - modifyConstraints (<>M.map (lvl,) (M.fromList constraints)) - localScope extend m - where (tbinds, constraints) = partitionEithers types - extend scope = scope { - scopeTypeTable = M.fromList tbinds <> scopeTypeTable scope - } - -bindingTypeParams :: [TypeParam] -> TermTypeM a -> TermTypeM a -bindingTypeParams tparams = binding (mapMaybe typeParamIdent tparams) . - bindingTypes (concatMap typeParamType tparams) - where typeParamType (TypeParamType l v loc) = - [ Left (v, TypeAbbr l [] (Scalar (TypeVar () Nonunique (typeName v) []))) - , Right (v, ParamType l loc) ] - typeParamType (TypeParamDim v loc) = - [ Right (v, ParamSize loc) ] - -typeParamIdent :: TypeParam -> Maybe Ident -typeParamIdent (TypeParamDim v loc) = - Just $ Ident v (Info $ Scalar $ Prim $ Signed Int32) loc -typeParamIdent _ = Nothing - -bindingIdent :: IdentBase NoInfo Name -> PatternType -> (Ident -> TermTypeM a) - -> TermTypeM a -bindingIdent (Ident v NoInfo vloc) t m = - bindSpaced [(Term, v)] $ do - v' <- checkName Term v vloc - let ident = Ident v' (Info t) vloc - binding [ident] $ m ident - -bindingParams :: [UncheckedTypeParam] - -> [UncheckedPattern] - -> ([TypeParam] -> [Pattern] -> TermTypeM a) -> TermTypeM a -bindingParams tps orig_ps m = do - checkForDuplicateNames orig_ps - checkTypeParams tps $ \tps' -> bindingTypeParams tps' $ do - let descend ps' (p:ps) = - checkPattern p NoneInferred $ \p' -> - binding (S.toList $ patternIdents p') $ descend (p':ps') ps - descend ps' [] = do - -- Perform an observation of every type parameter. This - -- prevents unused-name warnings for otherwise unused - -- dimensions. - mapM_ observe $ mapMaybe typeParamIdent tps' - m tps' $ reverse ps' - - descend [] orig_ps - -bindingPattern :: PatternBase NoInfo Name -> InferredType - -> (Pattern -> TermTypeM a) -> TermTypeM a -bindingPattern p t m = do - checkForDuplicateNames [p] - checkPattern p t $ \p' -> binding (S.toList $ patternIdents p') $ do - -- Perform an observation of every declared dimension. This - -- prevents unused-name warnings for otherwise unused dimensions. - mapM_ observe $ patternDims p' - - m p' - -patternDims :: Pattern -> [Ident] -patternDims (PatternParens p _) = patternDims p -patternDims (TuplePattern pats _) = concatMap patternDims pats -patternDims (PatternAscription p (TypeDecl _ (Info t)) _) = - patternDims p <> mapMaybe (dimIdent (srclocOf p)) (nestedDims t) - where dimIdent _ AnyDim = Nothing - dimIdent _ (ConstDim _) = Nothing - dimIdent _ NamedDim{} = Nothing -patternDims _ = [] - -sliceShape :: Maybe (SrcLoc, Rigidity) -> [DimIndex] -> TypeBase (DimDecl VName) as - -> TermTypeM (TypeBase (DimDecl VName) as, [VName]) -sliceShape r slice t@(Array als u et (ShapeDecl orig_dims)) = - runWriterT $ setDims <$> adjustDims slice orig_dims - where setDims [] = stripArray (length orig_dims) t - setDims dims' = Array als u et $ ShapeDecl dims' - - -- If the result is supposed to be AnyDim or a nonrigid size - -- variable, then don't bother trying to create - -- non-existential sizes. This is necessary to make programs - -- type-check without too much ceremony; see - -- e.g. tests/inplace5.fut. - isRigid Rigid{} = True - isRigid _ = False - refine_sizes = maybe False (isRigid . snd) r - - sliceSize orig_d i j stride = - case r of - Just (loc, Rigid _) -> do - (d, ext) <- - lift $ extSize loc $ - SourceSlice orig_d' (bareExp <$> i) (bareExp <$> j) (bareExp <$> stride) - tell $ maybeToList ext - return d - Just (loc, Nonrigid) -> - lift $ NamedDim . qualName <$> newDimVar loc Nonrigid "slice_dim" - Nothing -> - pure AnyDim - where - -- The original size does not matter if the slice is fully specified. - orig_d' | isJust i, isJust j = Nothing - | otherwise = Just orig_d - - adjustDims (DimFix{} : idxes') (_:dims) = - adjustDims idxes' dims - - -- Pattern match some known slices to be non-existential. - adjustDims (DimSlice i j stride : idxes') (_:dims) - | refine_sizes, - maybe True ((==Just 0) . isInt32) i, - Just j' <- maybeDimFromExp =<< j, - maybe True ((==Just 1) . isInt32) stride = - (j':) <$> adjustDims idxes' dims - - adjustDims (DimSlice Nothing Nothing stride : idxes') (d:dims) - | refine_sizes, - maybe True (maybe False ((==1) . abs) . isInt32) stride = - (d:) <$> adjustDims idxes' dims - - adjustDims (DimSlice i j stride : idxes') (d:dims) = - (:) <$> sliceSize d i j stride <*> adjustDims idxes' dims - - adjustDims _ dims = - pure dims - -sliceShape _ _ t = pure (t, []) - ---- Main checkers - --- | @require ts e@ causes a 'TypeError' if @expType e@ is not one of --- the types in @ts@. Otherwise, simply returns @e@. -require :: String -> [PrimType] -> Exp -> TermTypeM Exp -require why ts e = do mustBeOneOf ts (mkUsage (srclocOf e) why) . toStruct =<< expType e - return e - -unifies :: String -> StructType -> Exp -> TermTypeM Exp -unifies why t e = do - unify (mkUsage (srclocOf e) why) t . toStruct =<< expType e - return e - --- The closure of a lambda or local function are those variables that --- it references, and which local to the current top-level function. -lexicalClosure :: [Pattern] -> Occurences -> TermTypeM Aliasing -lexicalClosure params closure = do - vtable <- asks $ scopeVtable . termScope - let isLocal v = case v `M.lookup` vtable of - Just (BoundV Local _ _) -> True - _ -> False - return $ S.map AliasBound $ S.filter isLocal $ - allOccuring closure S.\\ mconcat (map patternNames params) - -noAliasesIfOverloaded :: PatternType -> TermTypeM PatternType -noAliasesIfOverloaded t@(Scalar (TypeVar _ u tn [])) = do - subst <- fmap snd . M.lookup (typeLeaf tn) <$> getConstraints - case subst of - Just Overloaded{} -> return $ Scalar $ TypeVar mempty u tn [] - _ -> return t -noAliasesIfOverloaded t = - return t - --- Check the common parts of ascription and coercion. -checkAscript :: SrcLoc - -> UncheckedTypeDecl - -> UncheckedExp - -> (StructType -> StructType) - -> TermTypeM (TypeDecl, Exp) -checkAscript loc decl e shapef = do - decl' <- checkTypeDecl decl - e' <- checkExp e - t <- expTypeFully e' - - (decl_t_nonrigid, _) <- - instantiateEmptyArrayDims loc "impl" Nonrigid $ shapef $ - unInfo $ expandedType decl' - - onFailure (CheckingAscription (unInfo $ expandedType decl') (toStruct t)) $ - unify (mkUsage loc "type ascription") decl_t_nonrigid (toStruct t) - - -- We also have to make sure that uniqueness matches. This is done - -- explicitly, because uniqueness is ignored by unification. - t' <- normTypeFully t - decl_t' <- normTypeFully $ unInfo $ expandedType decl' - unless (t' `subtypeOf` anySizes decl_t') $ - typeError loc mempty $ "Type" <+> pquote (ppr t') <+> "is not a subtype of" <+> - pquote (ppr decl_t') <> "." - - return (decl', e') - -unscopeType :: SrcLoc - -> M.Map VName Ident - -> PatternType - -> TermTypeM (PatternType, [VName]) -unscopeType tloc unscoped t = do - (t', m) <- runStateT (traverseDims onDim t) mempty - return (t' `addAliases` S.map unAlias, M.elems m) - where onDim _ p (NamedDim d) - | Just loc <- srclocOf <$> M.lookup (qualLeaf d) unscoped = - if p == PosImmediate || p == PosParam - then inst loc $ qualLeaf d - else return AnyDim - onDim _ _ d = return d - - inst loc d = do - prev <- gets $ M.lookup d - case prev of - Just d' -> return $ NamedDim $ qualName d' - Nothing -> do - d' <- lift $ newDimVar tloc (Rigid $ RigidOutOfScope loc d) "d" - modify $ M.insert d d' - return $ NamedDim $ qualName d' - - unAlias (AliasBound v) | v `M.member` unscoped = AliasFree v - unAlias a = a - --- 'checkApplyExp' is like 'checkExp', but tries to find the "root --- function", for better error messages. -checkApplyExp :: UncheckedExp -> TermTypeM (Exp, ApplyOp) - -checkApplyExp (Apply e1 e2 _ _ loc) = do - (e1', (fname, i)) <- checkApplyExp e1 - arg <- checkArg e2 - t <- expType e1' - (t1, rt, argext, exts) <- checkApply loc (fname, i) t arg - return (Apply e1' (argExp arg) (Info (diet t1, argext)) (Info rt, Info exts) loc, - (fname, i+1)) - -checkApplyExp e = do - e' <- checkExp e - return (e', - (case e' of Var qn _ _ -> Just qn - _ -> Nothing, - 0)) - -checkExp :: UncheckedExp -> TermTypeM Exp - -checkExp (Literal val loc) = - return $ Literal val loc - -checkExp (StringLit vs loc) = - return $ StringLit vs loc - -checkExp (IntLit val NoInfo loc) = do - t <- newTypeVar loc "t" - mustBeOneOf anyNumberType (mkUsage loc "integer literal") t - return $ IntLit val (Info $ fromStruct t) loc - -checkExp (FloatLit val NoInfo loc) = do - t <- newTypeVar loc "t" - mustBeOneOf anyFloatType (mkUsage loc "float literal") t - return $ FloatLit val (Info $ fromStruct t) loc - -checkExp (TupLit es loc) = - TupLit <$> mapM checkExp es <*> pure loc - -checkExp (RecordLit fs loc) = do - fs' <- evalStateT (mapM checkField fs) mempty - - return $ RecordLit fs' loc - where checkField (RecordFieldExplicit f e rloc) = do - errIfAlreadySet f rloc - modify $ M.insert f rloc - RecordFieldExplicit f <$> lift (checkExp e) <*> pure rloc - checkField (RecordFieldImplicit name NoInfo rloc) = do - errIfAlreadySet name rloc - (QualName _ name', t) <- lift $ lookupVar rloc $ qualName name - modify $ M.insert name rloc - return $ RecordFieldImplicit name' (Info t) rloc - - errIfAlreadySet f rloc = do - maybe_sloc <- gets $ M.lookup f - case maybe_sloc of - Just sloc -> - lift $ typeError rloc mempty $ "Field" <+> pquote (ppr f) <+> - "previously defined at" <+> text (locStrRel rloc sloc) <> "." - Nothing -> return () - -checkExp (ArrayLit all_es _ loc) = - -- Construct the result type and unify all elements with it. We - -- only create a type variable for empty arrays; otherwise we use - -- the type of the first element. This significantly cuts down on - -- the number of type variables generated for pathologically large - -- multidimensional array literals. - case all_es of - [] -> do et <- newTypeVar loc "t" - t <- arrayOfM loc et (ShapeDecl [ConstDim 0]) Unique - return $ ArrayLit [] (Info t) loc - e:es -> do - e' <- checkExp e - et <- expType e' - es' <- mapM (unifies "type of first array element" (toStruct et) <=< checkExp) es - et' <- normTypeFully et - t <- arrayOfM loc et' (ShapeDecl [ConstDim $ length all_es]) Unique - return $ ArrayLit (e':es') (Info t) loc - -checkExp (Range start maybe_step end _ loc) = do - start' <- require "use in range expression" anyIntType =<< checkExp start - start_t <- toStruct <$> expTypeFully start' - maybe_step' <- case maybe_step of - Nothing -> return Nothing - Just step -> do - let warning = warn loc "First and second element of range are identical, this will produce an empty array." - case (start, step) of - (Literal x _, Literal y _) -> when (x == y) warning - (Var x_name _ _, Var y_name _ _) -> when (x_name == y_name) warning - _ -> return () - Just <$> (unifies "use in range expression" start_t =<< checkExp step) - - let unifyRange e = unifies "use in range expression" start_t =<< checkExp e - end' <- case end of - DownToExclusive e -> DownToExclusive <$> unifyRange e - UpToExclusive e -> UpToExclusive <$> unifyRange e - ToInclusive e -> ToInclusive <$> unifyRange e - - -- Special case some ranges to give them a known size. - let dimFromBound = dimFromExp (SourceBound . bareExp) - (dim, retext) <- - case (isInt32 start', isInt32 <$> maybe_step', end') of - (Just 0, Just (Just 1), UpToExclusive end'') -> - dimFromBound end'' - (Just 0, Nothing, UpToExclusive end'') -> - dimFromBound end'' - (Just 1, Just (Just 2), ToInclusive end'') -> - dimFromBound end'' - _ -> do - d <- newDimVar loc (Rigid RigidRange) "range_dim" - return (NamedDim $ qualName d, Just d) - - t <- arrayOfM loc start_t (ShapeDecl [dim]) Unique - let ret = (Info (t `setAliases` mempty), Info $ maybeToList retext) - - return $ Range start' maybe_step' end' ret loc - -checkExp (Ascript e decl loc) = do - (decl', e') <- checkAscript loc decl e id - return $ Ascript e' decl' loc - -checkExp (Coerce e decl _ loc) = do - -- We instantiate the declared types with all dimensions as nonrigid - -- fresh type variables, which we then use to unify with the type of - -- 'e'. This lets 'e' have whatever sizes it wants, but the overall - -- type must still match. Eventually we will throw away those sizes - -- (they will end up being unified with various sizes in 'e', which - -- is fine). - (decl', e') <- checkAscript loc decl e anySizes - - -- Now we instantiate the declared type again, but this time we keep - -- around the sizes as existentials. This is the result of the - -- ascription as a whole. We use matchDims to obtain the aliasing - -- of 'e'. - (decl_t_rigid, ext) <- - instantiateDimsInReturnType loc Nothing $ unInfo $ expandedType decl' - - t <- expTypeFully e' - - t' <- matchDims (const pure) t $ fromStruct decl_t_rigid - - return $ Coerce e' decl' (Info t', Info ext) loc - -checkExp (BinOp (op, oploc) NoInfo (e1,_) (e2,_) NoInfo NoInfo loc) = do - (op', ftype) <- lookupVar oploc op - e1_arg <- checkArg e1 - e2_arg <- checkArg e2 - - -- Note that the application to the first operand cannot fix any - -- existential sizes, because it must by necessity be a function. - (p1_t, rt, p1_ext, _) <- checkApply loc (Just op', 0) ftype e1_arg - (p2_t, rt', p2_ext, retext) <- checkApply loc (Just op', 1) rt e2_arg - - return $ BinOp (op', oploc) (Info ftype) - (argExp e1_arg, Info (toStruct p1_t, p1_ext)) - (argExp e2_arg, Info (toStruct p2_t, p2_ext)) - (Info rt') (Info retext) loc - -checkExp (Project k e NoInfo loc) = do - e' <- checkExp e - t <- expType e' - kt <- mustHaveField (mkUsage loc $ "projection of field " ++ quote (pretty k)) k t - return $ Project k e' (Info kt) loc - -checkExp (If e1 e2 e3 _ loc) = - sequentially checkCond $ \e1' _ -> do - ((e2', e3'), dflow) <- tapOccurences $ checkExp e2 `alternative` checkExp e3 - - (brancht, retext) <- unifyBranches loc e2' e3' - let t' = addAliases brancht (`S.difference` S.map AliasBound (allConsumed dflow)) - - zeroOrderType (mkUsage loc "returning value of this type from 'if' expression") - "type returned from branch" t' - - return $ If e1' e2' e3' (Info t', Info retext) loc - - where checkCond = do - e1' <- checkExp e1 - let bool = Scalar $ Prim Bool - e1_t <- toStruct <$> expType e1' - onFailure (CheckingRequired [bool] e1_t) $ - unify (mkUsage (srclocOf e1') "use as 'if' condition") bool e1_t - return e1' - -checkExp (Parens e loc) = - Parens <$> checkExp e <*> pure loc - -checkExp (QualParens (modname, modnameloc) e loc) = do - (modname',mod) <- lookupMod loc modname - case mod of - ModEnv env -> local (`withEnv` qualifyEnv modname' env) $ do - e' <- checkExp e - return $ QualParens (modname', modnameloc) e' loc - ModFun{} -> - typeError loc mempty $ "Module" <+> ppr modname <+> " is a parametric module." - where qualifyEnv modname' env = - env { envNameMap = M.map (qualify' modname') $ envNameMap env } - qualify' modname' (QualName qs name) = - QualName (qualQuals modname' ++ [qualLeaf modname'] ++ qs) name - -checkExp (Var qn NoInfo loc) = do - -- The qualifiers of a variable is divided into two parts: first a - -- possibly-empty sequence of module qualifiers, followed by a - -- possible-empty sequence of record field accesses. We use scope - -- information to perform the split, by taking qualifiers off the - -- end until we find a module. - - (qn', t, fields) <- findRootVar (qualQuals qn) (qualLeaf qn) - - foldM checkField (Var qn' (Info t) loc) fields - - where findRootVar qs name = - (whenFound <$> lookupVar loc (QualName qs name)) `catchError` notFound qs name - - whenFound (qn', t) = (qn', t, []) - - notFound qs name err - | null qs = throwError err - | otherwise = do - (qn', t, fields) <- findRootVar (init qs) (last qs) `catchError` - const (throwError err) - return (qn', t, fields++[name]) - - checkField e k = do - t <- expType e - let usage = mkUsage loc $ "projection of field " ++ quote (pretty k) - kt <- mustHaveField usage k t - return $ Project k e (Info kt) loc - -checkExp (Negate arg loc) = do - arg' <- require "numeric negation" anyNumberType =<< checkExp arg - return $ Negate arg' loc - -checkExp e@Apply{} = fst <$> checkApplyExp e - -checkExp (LetPat pat e body _ loc) = - sequentially (checkExp e) $ \e' e_occs -> do - -- Not technically an ascription, but we want the pattern to have - -- exactly the type of 'e'. - t <- expType e' - case anyConsumption e_occs of - Just c -> - let msg = "type computed with consumption at " ++ locStr (location c) - in zeroOrderType (mkUsage loc "consumption in right-hand side of 'let'-binding") msg t - _ -> return () - - incLevel $ bindingPattern pat (Ascribed t) $ \pat' -> do - body' <- checkExp body - (body_t, retext) <- - unscopeType loc (patternMap pat') =<< expTypeFully body' - - return $ LetPat pat' e' body' (Info body_t, Info retext) loc - -checkExp (LetFun name (tparams, params, maybe_retdecl, NoInfo, e) body NoInfo loc) = - sequentially (checkBinding (name, maybe_retdecl, tparams, params, e, loc)) $ - \(tparams', params', maybe_retdecl', rettype, _, e') closure -> do - - closure' <- lexicalClosure params' closure - - bindSpaced [(Term, name)] $ do - name' <- checkName Term name loc - - let arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt - ftype = foldr (arrow . patternParam) rettype params' - entry = BoundV Local tparams' $ ftype `setAliases` closure' - bindF scope = scope { scopeVtable = - M.insert name' entry $ scopeVtable scope - , scopeNameMap = - M.insert (Term, name) (qualName name') $ - scopeNameMap scope } - body' <- localScope bindF $ checkExp body - - -- We fake an ident here, but it's OK as it can't be a size - -- anyway. - let fake_ident = Ident name' (Info $ fromStruct ftype) mempty - (body_t, _) <- - unscopeType loc (M.singleton name' fake_ident) =<< - expTypeFully body' - - return $ LetFun name' (tparams', params', maybe_retdecl', Info rettype, e') - body' (Info body_t) loc - -checkExp (LetWith dest src idxes ve body NoInfo loc) = - sequentially (checkIdent src) $ \src' _ -> do - (t, _) <- newArrayType (srclocOf src) "src" $ length idxes - unify (mkUsage loc "type of target array") t $ toStruct $ unInfo $ identType src' - - -- Need the fully normalised type here to get the proper aliasing information. - src_t <- normTypeFully $ unInfo $ identType src' - - idxes' <- mapM checkDimIndex idxes - (elemt, _) <- sliceShape (Just (loc, Nonrigid)) idxes' =<< normTypeFully t - - unless (unique src_t) $ - typeError loc mempty $ "Source" <+> pquote (pprName (identName src)) <+> - "has type" <+> ppr src_t <> ", which is not unique." - vtable <- asks $ scopeVtable . termScope - forM_ (aliases src_t) $ \v -> - case aliasVar v `M.lookup` vtable of - Just (BoundV Local _ v_t) - | not $ unique v_t -> - typeError loc mempty $ "Source" <+> pquote (pprName (identName src)) <+> - "aliases" <+> pquote (pprName (aliasVar v)) <> ", which is not consumable." - _ -> return () - - sequentially (unifies "type of target array" (toStruct elemt) =<< checkExp ve) $ \ve' _ -> do - ve_t <- expTypeFully ve' - when (AliasBound (identName src') `S.member` aliases ve_t) $ - badLetWithValue loc - - bindingIdent dest (src_t `setAliases` S.empty) $ \dest' -> do - body' <- consuming src' $ checkExp body - (body_t, _) <- - unscopeType loc (M.singleton (identName dest') dest') =<< - expTypeFully body' - return $ LetWith dest' src' idxes' ve' body' (Info body_t) loc - -checkExp (Update src idxes ve loc) = do - (t, _) <- newArrayType (srclocOf src) "src" $ length idxes - idxes' <- mapM checkDimIndex idxes - (elemt, _) <- sliceShape (Just (loc, Nonrigid)) idxes' =<< normTypeFully t - - sequentially (checkExp ve >>= unifies "type of target array" elemt) $ \ve' _ -> - sequentially (checkExp src >>= unifies "type of target array" t) $ \src' _ -> do - - src_t <- expTypeFully src' - unless (unique src_t) $ - typeError loc mempty $ "Source" <+> pquote (ppr src) <+> - "has type" <+> ppr src_t <> ", which is not unique." - - let src_als = aliases src_t - ve_t <- expTypeFully ve' - unless (S.null $ src_als `S.intersection` aliases ve_t) $ badLetWithValue loc - - consume loc src_als - return $ Update src' idxes' ve' loc - --- Record updates are a bit hacky, because we do not have row typing --- (yet?). For now, we only permit record updates where we know the --- full type up to the field we are updating. -checkExp (RecordUpdate src fields ve NoInfo loc) = do - src' <- checkExp src - ve' <- checkExp ve - a <- expTypeFully src' - let usage = mkUsage loc "record update" - r <- foldM (flip $ mustHaveField usage) a fields - ve_t <- expType ve' - let r' = anySizes $ toStruct r - ve_t' = anySizes $ toStruct ve_t - onFailure (CheckingRecordUpdate fields r' ve_t') $ - unify usage r' ve_t' - maybe_a' <- onRecordField (const ve_t) fields <$> expTypeFully src' - case maybe_a' of - Just a' -> return $ RecordUpdate src' fields ve' (Info a') loc - Nothing -> typeError loc mempty $ - "Full type of" - indent 2 (ppr src) - textwrap " is not known at this point. Add a size annotation to the original record to disambiguate." - -checkExp (Index e idxes _ loc) = do - (t, _) <- newArrayType loc "e" $ length idxes - e' <- unifies "being indexed at" t =<< checkExp e - idxes' <- mapM checkDimIndex idxes - -- XXX, the RigidSlice here will be overridden in sliceShape with a proper value. - (t', retext) <- - sliceShape (Just (loc, Rigid (RigidSlice Nothing ""))) idxes' =<< - expTypeFully e' - - -- Remove aliases if the result is an overloaded type, because that - -- will certainly not be aliased. - t'' <- noAliasesIfOverloaded t' - - return $ Index e' idxes' (Info t'', Info retext) loc - -checkExp (Assert e1 e2 NoInfo loc) = do - e1' <- require "being asserted" [Bool] =<< checkExp e1 - e2' <- checkExp e2 - return $ Assert e1' e2' (Info (pretty e1)) loc - -checkExp (Lambda params body rettype_te NoInfo loc) = - removeSeminullOccurences $ noUnique $ incLevel $ - bindingParams [] params $ \_ params' -> do - rettype_checked <- traverse checkTypeExp rettype_te - let declared_rettype = - case rettype_checked of Just (_, st, _) -> Just st - Nothing -> Nothing - (body', closure) <- - tapOccurences $ checkFunBody params' body declared_rettype loc - body_t <- expTypeFully body' - - params'' <- mapM updateTypes params' - - (rettype', rettype_st) <- - case rettype_checked of - Just (te, st, _) -> - return (Just te, st) - Nothing -> do - ret <- inferReturnSizes params'' $ toStruct $ - inferReturnUniqueness params'' body_t - return (Nothing, ret) - - checkGlobalAliases params' body_t loc - verifyFunctionParams Nothing params' - - closure' <- lexicalClosure params'' closure - - return $ Lambda params'' body' rettype' (Info (closure', rettype_st)) loc - - where - -- Inferring the sizes of the return type of a lambda is a lot - -- like let-generalisation. We wish to remove any rigid sizes - -- that were created when checking the body, except for those that - -- are visible in types that existed before we entered the body, - -- are parameters, or are used in parameters. - inferReturnSizes params' ret = do - cur_lvl <- curLevel - let named (Named x, _) = Just x - named (Unnamed, _) = Nothing - param_names = mapMaybe (named . patternParam) params' - pos_sizes = - typeDimNamesPos (foldFunType (map patternStructType params') ret) - hide k (lvl, _) = - lvl >= cur_lvl && k `notElem` param_names && k `S.notMember` pos_sizes - - hidden_sizes <- - S.fromList . M.keys . M.filterWithKey hide <$> getConstraints - - let onDim (NamedDim name) - | not (qualLeaf name `S.member` hidden_sizes) = NamedDim name - | otherwise = AnyDim - onDim d = d - - return $ first onDim ret - -checkExp (OpSection op _ loc) = do - (op', ftype) <- lookupVar loc op - return $ OpSection op' (Info ftype) loc - -checkExp (OpSectionLeft op _ e _ _ loc) = do - (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) -> - return $ OpSectionLeft op' (Info ftype) (argExp e_arg) - (Info (toStruct t1, argext), Info $ toStruct t2) (Info rettype, Info retext) loc - _ -> typeError loc mempty $ - "Operator section with invalid operator of type" <+> ppr ftype - -checkExp (OpSectionRight op _ e _ NoInfo loc) = do - (op', ftype) <- lookupVar loc op - e_arg <- checkArg e - case ftype of - Scalar (Arrow as1 m1 t1 (Scalar (Arrow as2 m2 t2 ret))) -> do - (t2', ret', argext, _) <- - checkApply loc (Just op', 1) - (Scalar $ Arrow as2 m2 t2 $ Scalar $ Arrow as1 m1 t1 ret) e_arg - return $ OpSectionRight op' (Info ftype) (argExp e_arg) - (Info $ toStruct t1, Info (toStruct t2', argext)) - (Info $ addAliases ret (<>aliases ret')) loc - _ -> typeError loc mempty $ - "Operator section with invalid operator of type" <+> ppr ftype - -checkExp (ProjectSection fields NoInfo loc) = do - a <- newTypeVar loc "a" - let usage = mkUsage loc "projection at" - b <- foldM (flip $ mustHaveField usage) a fields - return $ ProjectSection fields (Info $ Scalar $ Arrow mempty Unnamed a b) loc - -checkExp (IndexSection idxes NoInfo loc) = do - (t, _) <- newArrayType loc "e" $ length idxes - idxes' <- mapM checkDimIndex idxes - (t', _) <- sliceShape Nothing idxes' t - return $ IndexSection idxes' (Info $ fromStruct $ Scalar $ Arrow mempty Unnamed t t') loc - -checkExp (DoLoop _ mergepat mergeexp form loopbody NoInfo loc) = - sequentially (checkExp mergeexp) $ \mergeexp' _ -> do - - zeroOrderType (mkUsage (srclocOf mergeexp) "use as loop variable") - "type used as loop variable" =<< expTypeFully mergeexp' - - -- The handling of dimension sizes is a bit intricate, but very - -- similar to checking a function, followed by checking a call to - -- it. The overall procedure is as follows: - -- - -- (1) All empty dimensions in the merge pattern are instantiated - -- with nonrigid size variables. All explicitly specified - -- dimensions are preserved. - -- - -- (2) The body of the loop is type-checked. The result type is - -- combined with the merge pattern type to determine which sizes are - -- variant, and these are turned into size parameters for the merge - -- pattern. - -- - -- (3) We now conceptually have a function parameter type and return - -- type. We check that it can be called with the initial merge - -- values as argument. The result of this is the type of the loop - -- as a whole. - -- - -- (There is also a convergence loop for inferring uniqueness, but - -- that's orthogonal to the size handling.) - - (merge_t, new_dims) <- - instantiateEmptyArrayDims loc "loop" Nonrigid . -- dim handling (1) - anySizes - =<< expTypeFully mergeexp' - - -- dim handling (2) - let checkLoopReturnSize mergepat' loopbody' = do - loopbody_t <- expTypeFully loopbody' - pat_t <- normTypeFully $ patternType mergepat' - -- We are ignoring the dimensions here, because any mismatches - -- should be turned into fresh size variables. - onFailure (CheckingLoopBody (toStruct (anySizes pat_t)) (toStruct loopbody_t)) $ - expect (mkUsage (srclocOf loopbody) "matching loop body to loop pattern") - (toStruct (anySizes pat_t)) - (toStruct loopbody_t) - pat_t' <- normTypeFully pat_t - loopbody_t' <- normTypeFully loopbody_t - - -- For each new_dims, figure out what they are instantiated - -- with in the initial value. This is used to determine - -- whether a size is invariant because it always matches the - -- initial instantiation of that size. - let initSubst (NamedDim v, d) = Just (v, d) - initSubst _ = Nothing - init_substs <- M.fromList . mapMaybe initSubst . snd . - anyDimOnMismatch pat_t' <$> - expTypeFully mergeexp' - - -- Figure out which of the 'new_dims' dimensions are variant. - -- This works because we know that each dimension from - -- new_dims in the pattern is unique and distinct. - -- - -- Our logic here is a bit reversed: the *mismatches* (from - -- new_dims) are what we want to extract and turn into size - -- parameters. - let mismatchSubst (NamedDim v, d) - | qualLeaf v `elem` new_dims = - case M.lookup v init_substs of - Just d' - | d' == d -> - return $ Just (qualLeaf v, SizeSubst d) - _ -> do tell [qualLeaf v] - return Nothing - mismatchSubst _ = return Nothing - - (init_substs', sparams) = - runWriter $ M.fromList . catMaybes <$> mapM mismatchSubst - (snd $ anyDimOnMismatch pat_t' loopbody_t') - - -- Make sure that any of new_dims that are invariant will be - -- replaced with the invariant size in the loop body. Failure - -- to do this can cause type annotations to still refer to - -- new_dims. - let dimToInit (v, SizeSubst d) = - constrain v $ Size (Just d) (mkUsage loc "size of loop parameter") - dimToInit _ = - return () - mapM_ dimToInit $ M.toList init_substs' - - mergepat'' <- applySubst (`M.lookup` init_substs') <$> updateTypes mergepat' - return (nub sparams, mergepat'') - - -- First we do a basic check of the loop body to figure out which of - -- the merge parameters are being consumed. For this, we first need - -- to check the merge pattern, which requires the (initial) merge - -- expression. - -- - -- Play a little with occurences to ensure it does not look like - -- none of the merge variables are being used. - ((sparams, mergepat', form', loopbody'), bodyflow) <- - case form of - For i uboundexp -> do - uboundexp' <- require "being the bound in a 'for' loop" anySignedType =<< checkExp uboundexp - bound_t <- expTypeFully uboundexp' - bindingIdent i bound_t $ \i' -> - noUnique $ bindingPattern mergepat (Ascribed merge_t) $ - \mergepat' -> onlySelfAliasing $ tapOccurences $ do - loopbody' <- noSizeEscape $ checkExp loopbody - (sparams, mergepat'') <- checkLoopReturnSize mergepat' loopbody' - return (sparams, - mergepat'', - For i' uboundexp', - loopbody') - - ForIn xpat e -> do - (arr_t, _) <- newArrayType (srclocOf e) "e" 1 - e' <- unifies "being iterated in a 'for-in' loop" arr_t =<< checkExp e - t <- expTypeFully e' - case t of - _ | Just t' <- peelArray 1 t -> - bindingPattern xpat (Ascribed t') $ \xpat' -> - noUnique $ bindingPattern mergepat (Ascribed merge_t) $ - \mergepat' -> onlySelfAliasing $ tapOccurences $ do - loopbody' <- noSizeEscape $ checkExp loopbody - (sparams, mergepat'') <- checkLoopReturnSize mergepat' loopbody' - return (sparams, - mergepat'', - ForIn xpat' e', - loopbody') - | otherwise -> - typeError (srclocOf e) mempty $ - "Iteratee of a for-in loop must be an array, but expression has type" <+> - ppr t - - While cond -> - noUnique $ bindingPattern mergepat (Ascribed merge_t) $ \mergepat' -> - onlySelfAliasing $ tapOccurences $ - sequentially (checkExp cond >>= - unifies "being the condition of a 'while' loop" (Scalar $ Prim Bool)) $ \cond' _ -> do - loopbody' <- noSizeEscape $ checkExp loopbody - (sparams, mergepat'') <- checkLoopReturnSize mergepat' loopbody' - return (sparams, - mergepat'', - While cond', - loopbody') - - mergepat'' <- do - loopbody_t <- expTypeFully loopbody' - convergePattern mergepat' (allConsumed bodyflow) loopbody_t $ - mkUsage (srclocOf loopbody') "being (part of) the result of the loop body" - - let consumeMerge (Id _ (Info pt) ploc) mt - | unique pt = consume ploc $ aliases mt - consumeMerge (TuplePattern pats _) t | Just ts <- isTupleRecord t = - zipWithM_ consumeMerge pats ts - consumeMerge (PatternParens pat _) t = - consumeMerge pat t - consumeMerge (PatternAscription pat _ _) t = - consumeMerge pat t - consumeMerge _ _ = - return () - consumeMerge mergepat'' =<< expTypeFully mergeexp' - - -- dim handling (3) - let sparams_anydim = M.fromList $ zip sparams $ repeat $ SizeSubst AnyDim - loopt_anydims = applySubst (`M.lookup` sparams_anydim) $ - patternType mergepat'' - (merge_t', _) <- - instantiateEmptyArrayDims loc "loopres" Nonrigid $ toStruct loopt_anydims - mergeexp_t <- toStruct <$> expTypeFully mergeexp' - onFailure (CheckingLoopInitial (toStruct loopt_anydims) mergeexp_t) $ - unify (mkUsage (srclocOf mergeexp') "matching initial loop values to pattern") - merge_t' mergeexp_t - - (loopt, retext) <- instantiateDimsInType loc RigidLoop loopt_anydims - -- We set all of the uniqueness to be unique. This is intentional, - -- and matches what happens for function calls. Those arrays that - -- really *cannot* be consumed will alias something unconsumable, - -- and will be caught that way. - let bound_here = patternNames mergepat'' <> S.fromList sparams <> form_bound - form_bound = - case form' of - For v _ -> S.singleton $ identName v - ForIn forpat _ -> patternNames forpat - While{} -> mempty - loopt' = second (`S.difference` S.map AliasBound bound_here) $ - loopt `setUniqueness` Unique - - - -- Eliminate those new_dims that turned into sparams so it won't - -- look like we have ambiguous sizes lying around. - modifyConstraints $ M.filterWithKey $ \k _ -> k `notElem` sparams - - return $ DoLoop sparams mergepat'' mergeexp' form' loopbody' (Info (loopt', retext)) loc - - where - convergePattern pat body_cons body_t body_loc = do - let consumed_merge = patternNames pat `S.intersection` body_cons - - uniquePat (Wildcard (Info t) wloc) = - Wildcard (Info $ t `setUniqueness` Nonunique) wloc - uniquePat (PatternParens p ploc) = - PatternParens (uniquePat p) ploc - uniquePat (Id name (Info t) iloc) - | name `S.member` consumed_merge = - let t' = t `setUniqueness` Unique `setAliases` mempty - in Id name (Info t') iloc - | otherwise = - let t' = t `setUniqueness` Nonunique - in Id name (Info t') iloc - uniquePat (TuplePattern pats ploc) = - TuplePattern (map uniquePat pats) ploc - uniquePat (RecordPattern fs ploc) = - RecordPattern (map (fmap uniquePat) fs) ploc - uniquePat (PatternAscription p t ploc) = - PatternAscription p t ploc - uniquePat p@PatternLit{} = p - uniquePat (PatternConstr n t ps ploc) = - PatternConstr n t (map uniquePat ps) ploc - - -- Make the pattern unique where needed. - pat' = uniquePat pat - - pat_t <- normTypeFully $ patternType pat' - unless (toStructural body_t `subtypeOf` toStructural pat_t) $ - unexpectedType (srclocOf body_loc) (toStruct body_t) [toStruct pat_t] - - -- Check that the new values of consumed merge parameters do not - -- alias something bound outside the loop, AND that anything - -- returned for a unique merge parameter does not alias anything - -- else returned. We also update the aliases for the pattern. - bound_outside <- asks $ S.fromList . M.keys . scopeVtable . termScope - let combAliases t1 t2 = - case t1 of Scalar Record{} -> t1 - _ -> t1 `addAliases` (<>aliases t2) - - checkMergeReturn (Id pat_v (Info pat_v_t) patloc) t - | unique pat_v_t, - v:_ <- S.toList $ - S.map aliasVar (aliases t) `S.intersection` bound_outside = - lift $ typeError loc mempty $ - "Return value for loop parameter" <+> - pquote (pprName pat_v) <+> "aliases" <+> pprName v <> "." - - | otherwise = do - (cons,obs) <- get - unless (S.null $ aliases t `S.intersection` cons) $ - lift $ typeError loc mempty $ - "Return value for loop parameter" <+> - pquote (pprName pat_v) <+> - "aliases other consumed loop parameter." - when (unique pat_v_t && - not (S.null (aliases t `S.intersection` (cons<>obs)))) $ - lift $ typeError loc mempty $ - "Return value for consuming loop parameter" <+> - pquote (pprName pat_v) <+> "aliases previously returned value." - if unique pat_v_t - then put (cons<>aliases t, obs) - else put (cons, obs<>aliases t) - - return $ Id pat_v (Info (combAliases pat_v_t t)) patloc - - checkMergeReturn (Wildcard (Info pat_v_t) patloc) t = - return $ Wildcard (Info (combAliases pat_v_t t)) patloc - - checkMergeReturn (PatternParens p _) t = - checkMergeReturn p t - - checkMergeReturn (PatternAscription p _ _) t = - checkMergeReturn p t - - checkMergeReturn (RecordPattern pfs patloc) (Scalar (Record tfs)) = - RecordPattern . M.toList <$> sequence pfs' <*> pure patloc - where pfs' = M.intersectionWith checkMergeReturn - (M.fromList pfs) tfs - - checkMergeReturn (TuplePattern pats patloc) t - | Just ts <- isTupleRecord t = - TuplePattern - <$> zipWithM checkMergeReturn pats ts - <*> pure patloc - - checkMergeReturn p _ = - return p - - (pat'', (pat_cons, _)) <- - runStateT (checkMergeReturn pat' body_t) (mempty, mempty) - - let body_cons' = body_cons <> S.map aliasVar pat_cons - if body_cons' == body_cons && patternType pat'' == patternType pat - then return pat' - else convergePattern pat'' body_cons' body_t body_loc - -checkExp (Constr name es NoInfo loc) = do - t <- newTypeVar loc "t" - es' <- mapM checkExp es - ets <- mapM expTypeFully es' - mustHaveConstr (mkUsage loc "use of constructor") name t (toStruct <$> ets) - -- A sum value aliases *anything* that went into its construction. - let als = foldMap aliases ets - return $ Constr name es' (Info $ fromStruct t `addAliases` (<>als)) loc - -checkExp (Match e cs _ loc) = - sequentially (checkExp e) $ \e' _ -> do - mt <- expTypeFully e' - (cs', t, retext) <- checkCases mt cs - zeroOrderType (mkUsage loc "being returned 'match'") - "type returned from pattern match" t - return $ Match e' cs' (Info t, Info retext) loc - -checkExp (Attr info e loc) = - Attr info <$> checkExp e <*> pure loc - -checkCases :: PatternType - -> NE.NonEmpty (CaseBase NoInfo Name) - -> TermTypeM (NE.NonEmpty (CaseBase Info VName), PatternType, [VName]) -checkCases mt rest_cs = - case NE.uncons rest_cs of - (c, Nothing) -> do - (c', t, retext) <- checkCase mt c - return (c' NE.:| [], t, retext) - (c, Just cs) -> do - (((c', c_t, _), (cs', cs_t, _)), dflow) <- - tapOccurences $ checkCase mt c `alternative` checkCases mt cs - (brancht, retext) <- unifyBranchTypes (srclocOf c) c_t cs_t - let t = addAliases brancht - (`S.difference` S.map AliasBound (allConsumed dflow)) - return (NE.cons c' cs', t, retext) - -checkCase :: PatternType -> CaseBase NoInfo Name - -> TermTypeM (CaseBase Info VName, PatternType, [VName]) -checkCase mt (CasePat p e loc) = - bindingPattern p (Ascribed mt) $ \p' -> do - e' <- checkExp e - (t, retext) <- unscopeType loc (patternMap p') =<< expTypeFully e' - return (CasePat p' e' loc, t, retext) - --- | An unmatched pattern. Used in in the generation of --- unmatched pattern warnings by the type checker. -data Unmatched p = UnmatchedNum p [ExpBase Info VName] - | UnmatchedBool p - | UnmatchedConstr p - | Unmatched p - deriving (Functor, Show) - -instance Pretty (Unmatched (PatternBase Info VName)) where - ppr um = case um of - (UnmatchedNum p nums) -> ppr' p <+> "where p is not one of" <+> ppr nums - (UnmatchedBool p) -> ppr' p - (UnmatchedConstr p) -> ppr' p - (Unmatched p) -> ppr' p - where - ppr' (PatternAscription p t _) = ppr p <> ":" <+> ppr t - ppr' (PatternParens p _) = parens $ ppr' p - ppr' (Id v _ _) = pprName v - ppr' (TuplePattern pats _) = parens $ commasep $ map ppr' pats - ppr' (RecordPattern fs _) = braces $ commasep $ map ppField fs - where ppField (name, t) = text (nameToString name) <> equals <> ppr' t - ppr' Wildcard{} = "_" - ppr' (PatternLit e _ _) = ppr e - ppr' (PatternConstr n _ ps _) = "#" <> ppr n <+> sep (map ppr' ps) - -unpackPat :: Pattern -> [Maybe Pattern] -unpackPat Wildcard{} = [Nothing] -unpackPat (PatternParens p _) = unpackPat p -unpackPat Id{} = [Nothing] -unpackPat (TuplePattern ps _) = Just <$> ps -unpackPat (RecordPattern fs _) = Just . snd <$> sortFields (M.fromList fs) -unpackPat (PatternAscription p _ _) = unpackPat p -unpackPat p@PatternLit{} = [Just p] -unpackPat p@PatternConstr{} = [Just p] - -wildPattern :: Pattern -> Int -> Unmatched Pattern -> Unmatched Pattern -wildPattern (TuplePattern ps loc) pos um = wildTuple <$> um - where wildTuple p = TuplePattern (take (pos - 1) ps' ++ [p] ++ drop pos ps') loc - ps' = map wildOut ps - wildOut p = Wildcard (Info (patternType p)) (srclocOf p) -wildPattern (RecordPattern fs loc) pos um = wildRecord <$> um - where wildRecord p = - RecordPattern (take (pos - 1) fs' ++ [(fst (fs!!(pos - 1)), p)] ++ drop pos fs') loc - fs' = map wildOut fs - wildOut (f,p) = (f, Wildcard (Info (patternType p)) (srclocOf p)) -wildPattern (PatternAscription p _ _) pos um = wildPattern p pos um -wildPattern (PatternParens p _) pos um = wildPattern p pos um -wildPattern (PatternConstr n t ps loc) pos um = wildConstr <$> um - where wildConstr p = PatternConstr n t (take (pos - 1) ps' ++ [p] ++ drop pos ps') loc - ps' = map wildOut ps - wildOut p = Wildcard (Info (patternType p)) (srclocOf p) -wildPattern _ _ um = um - -checkUnmatched :: Exp -> TermTypeM () -checkUnmatched e = void $ checkUnmatched' e >> astMap tv e - where checkUnmatched' (Match _ cs _ loc) = - let ps = fmap (\(CasePat p _ _) -> p) cs - in case unmatched id $ NE.toList ps of - [] -> return () - ps' -> typeError loc mempty $ - "Unmatched cases in match expression:" - indent 2 (stack (map ppr ps')) - checkUnmatched' _ = return () - tv = ASTMapper { mapOnExp = - \e' -> checkUnmatched' e' >> return e' - , mapOnName = pure - , mapOnQualName = pure - , mapOnStructType = pure - , mapOnPatternType = pure - } - --- | A data type for constructor patterns. This is used to make the --- code for detecting unmatched constructors cleaner, by separating --- the constructor-pattern cases from other cases. -data ConstrPat = ConstrPat { constrName :: Name - , constrType :: PatternType - , constrPayload :: [Pattern] - , constrSrcLoc :: SrcLoc - } - --- Be aware of these fishy equality instances! - -instance Eq ConstrPat where - ConstrPat c1 _ _ _ == ConstrPat c2 _ _ _ = c1 == c2 - -instance Ord ConstrPat where - ConstrPat c1 _ _ _ `compare` ConstrPat c2 _ _ _ = c1 `compare` c2 - -unmatched :: (Unmatched Pattern -> Unmatched Pattern) -> [Pattern] -> [Unmatched Pattern] -unmatched hole orig_ps - | p:_ <- orig_ps, - sameStructure labeledCols = do - (i, cols) <- labeledCols - let hole' = if isConstr p then hole else hole . wildPattern p i - case sequence cols of - Nothing -> [] - Just cs - | all isPatternLit cs -> map hole' $ localUnmatched cs - | otherwise -> unmatched hole' cs - | otherwise = [] - - where labeledCols = zip [1..] $ transpose $ map unpackPat orig_ps - - localUnmatched :: [Pattern] -> [Unmatched Pattern] - localUnmatched [] = [] - localUnmatched ps'@(p':_) = - case patternType p' of - Scalar (Sum cs'') -> - -- We now know that we are matching a sum type, and thus - -- that all patterns ps' are constructors (checked by - -- 'all isPatternLit' before this function is called). - let constrs = M.keys cs'' - matched = mapMaybe constr ps' - unmatched' = map (UnmatchedConstr . buildConstr cs'') $ - constrs \\ map constrName matched - in case unmatched' of - [] -> - let constrGroups = group (sort matched) - removedConstrs = mapMaybe stripConstrs constrGroups - transposed = (fmap . fmap) transpose removedConstrs - findUnmatched (pc, trans) = do - col <- trans - case col of - [] -> [] - ((i, _):_) -> unmatched (wilder i pc) (map snd col) - wilder i pc s = (`PatternParens` mempty) <$> wildPattern pc i s - in concatMap findUnmatched transposed - _ -> unmatched' - Scalar (Prim t) | not (any idOrWild ps') -> - -- We now know that we are matching a sum type, and thus - -- that all patterns ps' are literals (checked by 'all - -- isPatternLit' before this function is called). - case t of - Bool -> - let matched = nub $ mapMaybe (pExp >=> bool) $ filter isPatternLit ps' - in map (UnmatchedBool . buildBool (Scalar (Prim t))) $ [True, False] \\ matched - _ -> - let matched = mapMaybe pExp $ filter isPatternLit ps' - in [UnmatchedNum (buildId (Info $ Scalar $ Prim t) "p") matched] - _ -> [] - - isConstr PatternConstr{} = True - isConstr (PatternParens p _) = isConstr p - isConstr _ = False - - - stripConstrs :: [ConstrPat] -> Maybe (Pattern, [[(Int, Pattern)]]) - stripConstrs (pc@ConstrPat{} : cs') = Just (unConstr pc, stripConstr pc : map stripConstr cs') - stripConstrs [] = Nothing - - stripConstr :: ConstrPat -> [(Int, Pattern)] - stripConstr (ConstrPat _ _ ps' _) = zip [1..] ps' - - sameStructure [] = True - sameStructure (x:xs) = all (\y -> length y == length x' ) xs' - where (x':xs') = map snd (x:xs) - - pExp (PatternLit e' _ _) = Just e' - pExp _ = Nothing - - constr (PatternConstr c (Info t) ps loc) = Just $ ConstrPat c t ps loc - constr (PatternParens p _) = constr p - constr (PatternAscription p' _ _) = constr p' - constr _ = Nothing - - unConstr p = - PatternConstr (constrName p) (Info $ constrType p) (constrPayload p) (constrSrcLoc p) - - isPatternLit PatternLit{} = True - isPatternLit (PatternAscription p' _ _) = isPatternLit p' - isPatternLit (PatternParens p' _) = isPatternLit p' - isPatternLit PatternConstr{} = True - isPatternLit _ = False - - idOrWild Id{} = True - idOrWild Wildcard{} = True - idOrWild (PatternAscription p' _ _) = idOrWild p' - idOrWild (PatternParens p' _) = idOrWild p' - idOrWild _ = False - - bool (Literal (BoolValue b) _ ) = Just b - bool _ = Nothing - - buildConstr m c = - let t = Scalar $ Sum m - cs = m M.! c - wildCS = map (\ct -> Wildcard (Info ct) mempty) cs - in if null wildCS - then PatternConstr c (Info t) [] mempty - else PatternParens (PatternConstr c (Info t) wildCS mempty) mempty - buildBool t b = - PatternLit (Literal (BoolValue b) mempty) (Info (addSizes t)) mempty - buildId t n = - -- The VName tag here will never be used since the value - -- exists exclusively for printing warnings. - Id (VName (nameFromString n) (-1)) t mempty - -checkIdent :: IdentBase NoInfo Name -> TermTypeM Ident -checkIdent (Ident name _ loc) = do - (QualName _ name', vt) <- lookupVar loc (qualName name) - return $ Ident name' (Info vt) loc - -checkDimIndex :: DimIndexBase NoInfo Name -> TermTypeM DimIndex -checkDimIndex (DimFix i) = - DimFix <$> (unifies "use as index" (Scalar $ Prim $ Signed Int32) =<< checkExp i) -checkDimIndex (DimSlice i j s) = - DimSlice <$> check i <*> check j <*> check s - where check = maybe (return Nothing) $ - fmap Just . unifies "use as index" (Scalar $ Prim $ Signed Int32) <=< checkExp - -sequentially :: TermTypeM a -> (a -> Occurences -> TermTypeM b) -> TermTypeM b -sequentially m1 m2 = do - (a, m1flow) <- collectOccurences m1 - (b, m2flow) <- collectOccurences $ m2 a m1flow - occur $ m1flow `seqOccurences` m2flow - return b - -type Arg = (Exp, PatternType, Occurences, SrcLoc) - -argExp :: Arg -> Exp -argExp (e, _, _, _) = e - -argType :: Arg -> PatternType -argType (_, t, _, _) = t - -checkArg :: UncheckedExp -> TermTypeM Arg -checkArg arg = do - (arg', dflow) <- collectOccurences $ checkExp arg - arg_t <- expType arg' - return (arg', arg_t, dflow, srclocOf arg') - -instantiateDimsInType :: SrcLoc -> RigidSource - -> TypeBase (DimDecl VName) als - -> TermTypeM (TypeBase (DimDecl VName) als, [VName]) -instantiateDimsInType tloc rsrc = - instantiateEmptyArrayDims tloc "d" $ Rigid rsrc - -instantiateDimsInReturnType :: SrcLoc -> Maybe (QualName VName) - -> TypeBase (DimDecl VName) als - -> TermTypeM (TypeBase (DimDecl VName) als, [VName]) -instantiateDimsInReturnType tloc fname = - instantiateEmptyArrayDims tloc "ret" $ Rigid $ RigidRet fname - --- Some information about the function/operator we are trying to --- apply, and how many arguments it has previously accepted. Used for --- generating nicer type errors. -type ApplyOp = (Maybe (QualName VName), Int) - -checkApply :: SrcLoc -> ApplyOp -> PatternType -> Arg - -> TermTypeM (PatternType, PatternType, Maybe VName, [VName]) -checkApply loc (fname, _) - (Scalar (Arrow as pname tp1 tp2)) - (argexp, argtype, dflow, argloc) = - onFailure (CheckingApply fname argexp (toStruct tp1) (toStruct argtype)) $ do - expect (mkUsage argloc "use as function argument") (toStruct tp1) (toStruct argtype) - - -- Perform substitutions of instantiated variables in the types. - tp1' <- normTypeFully tp1 - (tp2', ext) <- instantiateDimsInReturnType loc fname =<< normTypeFully tp2 - argtype' <- normTypeFully argtype - - -- Check whether this would produce an impossible return type. - let (_, tp2_paramdims, _) = dimUses $ toStruct tp2' - case filter (`S.member` tp2_paramdims) ext of - [] -> return () - ext_paramdims -> do - let onDim (NamedDim qn) - | qualLeaf qn `elem` ext_paramdims = AnyDim - onDim d = d - typeError loc mempty $ - "Anonymous size would appear in function parameter of return type:" - indent 2 (ppr (first onDim tp2')) - textwrap "This is usually because a higher-order function is used with functional arguments that return anonymous sizes, which are then used as parameters of other function arguments." - - occur [observation as loc] - - checkOccurences dflow - - case anyConsumption dflow of - Just c -> - let msg = "type of expression with consumption at " ++ locStr (location c) - in zeroOrderType (mkUsage argloc "potential consumption in expression") msg tp1 - _ -> return () - - occurs <- (dflow `seqOccurences`) <$> consumeArg argloc argtype' (diet tp1') - - checkIfConsumable loc $ S.map AliasBound $ allConsumed occurs - occur occurs - - (argext, parsubst) <- - case pname of - Named pname' -> do - (d, argext) <- sizeSubst tp1' argexp - return (argext, - (`M.lookup` M.singleton pname' (SizeSubst d))) - _ -> return (Nothing, const Nothing) - let tp2'' = applySubst parsubst $ returnType tp2' (diet tp1') argtype' - - return (tp1', tp2'', argext, ext) - where sizeSubst (Scalar (Prim (Signed Int32))) e = dimFromArg fname e - sizeSubst _ _ = return (AnyDim, Nothing) - -checkApply loc fname tfun@(Scalar TypeVar{}) arg = do - tv <- newTypeVar loc "b" - unify (mkUsage loc "use as function") (toStruct tfun) $ - Scalar $ Arrow mempty Unnamed (toStruct (argType arg)) tv - tfun' <- normPatternType tfun - checkApply loc fname tfun' arg - -checkApply loc (fname, prev_applied) ftype (argexp, _, _, _) = do - let fname' = maybe "expression" (pquote . ppr) fname - - typeError loc mempty $ - if prev_applied == 0 - then "Cannot apply" <+> fname' <+> "as function, as it has type:" - indent 2 (ppr ftype) - else "Cannot apply" <+> fname' <+> "to argument #" <> ppr (prev_applied+1) <+> - pquote (shorten $ pretty $ flatten $ ppr argexp) <> "," <+/> - "as" <+> fname' <+> "only takes" <+> ppr prev_applied <+> - arguments <> "." - where arguments | prev_applied == 1 = "argument" - | otherwise = "arguments" - -isInt32 :: Exp -> Maybe Int32 -isInt32 (Literal (SignedValue (Int32Value k')) _) = Just $ fromIntegral k' -isInt32 (IntLit k' _ _) = Just $ fromInteger k' -isInt32 (Negate x _) = negate <$> isInt32 x -isInt32 _ = Nothing - -maybeDimFromExp :: Exp -> Maybe (DimDecl VName) -maybeDimFromExp (Var v _ _) = Just $ NamedDim v -maybeDimFromExp (Parens e _) = maybeDimFromExp e -maybeDimFromExp (QualParens _ e _) = maybeDimFromExp e -maybeDimFromExp e = ConstDim . fromIntegral <$> isInt32 e - -dimFromExp :: (Exp -> SizeSource) -> Exp -> TermTypeM (DimDecl VName, Maybe VName) -dimFromExp rf (Parens e _) = dimFromExp rf e -dimFromExp rf (QualParens _ e _) = dimFromExp rf e -dimFromExp rf e - | Just d <- maybeDimFromExp e = - return (d, Nothing) - | otherwise = - extSize (srclocOf e) $ rf e - -dimFromArg :: Maybe (QualName VName) -> Exp -> TermTypeM (DimDecl VName, Maybe VName) -dimFromArg fname = dimFromExp $ SourceArg (FName fname) . bareExp - --- | @returnType ret_type arg_diet arg_type@ gives result of applying --- an argument the given types to a function with the given return --- type, consuming the argument with the given diet. -returnType :: PatternType - -> Diet - -> PatternType - -> PatternType -returnType (Array _ Unique et shape) _ _ = - Array mempty Unique et shape -returnType (Array als Nonunique et shape) d arg = - Array (als<>arg_als) Unique et shape -- Intentional! - where arg_als = aliases $ maskAliases arg d -returnType (Scalar (Record fs)) d arg = - Scalar $ Record $ fmap (\et -> returnType et d arg) fs -returnType (Scalar (Prim t)) _ _ = - Scalar $ Prim t -returnType (Scalar (TypeVar _ Unique t targs)) _ _ = - Scalar $ TypeVar mempty Unique t targs -returnType (Scalar (TypeVar als Nonunique t targs)) d arg = - Scalar $ TypeVar (als<>arg_als) Unique t targs -- Intentional! - where arg_als = aliases $ maskAliases arg d -returnType (Scalar (Arrow old_als v t1 t2)) d arg = - Scalar $ Arrow als v (t1 `setAliases` mempty) (t2 `setAliases` als) - -- Make sure to propagate the aliases of an existing closure. - where als = old_als <> aliases (maskAliases arg d) -returnType (Scalar (Sum cs)) d arg = - Scalar $ Sum $ (fmap . fmap) (\et -> returnType et d arg) cs - --- | @t `maskAliases` d@ removes aliases (sets them to 'mempty') from --- the parts of @t@ that are denoted as consumed by the 'Diet' @d@. -maskAliases :: Monoid as => - TypeBase shape as - -> Diet - -> TypeBase shape as -maskAliases t Consume = t `setAliases` mempty -maskAliases t Observe = t -maskAliases (Scalar (Record ets)) (RecordDiet ds) = - Scalar $ Record $ M.intersectionWith maskAliases ets ds -maskAliases t FuncDiet{} = t -maskAliases _ _ = error "Invalid arguments passed to maskAliases." - -consumeArg :: SrcLoc -> PatternType -> Diet -> TermTypeM [Occurence] -consumeArg loc (Scalar (Record ets)) (RecordDiet ds) = - concat . M.elems <$> traverse (uncurry $ consumeArg loc) (M.intersectionWith (,) ets ds) -consumeArg loc (Array _ Nonunique _ _) Consume = - typeError loc mempty "Consuming parameter passed non-unique argument." -consumeArg loc (Scalar (TypeVar _ Nonunique _ _)) Consume = - typeError loc mempty "Consuming parameter passed non-unique argument." -consumeArg loc (Scalar (Arrow _ _ t1 _)) (FuncDiet d _) - | not $ contravariantArg t1 d = - typeError loc mempty "Non-consuming higher-order parameter passed consuming argument." - where contravariantArg (Array _ Unique _ _) Observe = - False - contravariantArg (Scalar (TypeVar _ Unique _ _)) Observe = - False - contravariantArg (Scalar (Record ets)) (RecordDiet ds) = - and (M.intersectionWith contravariantArg ets ds) - contravariantArg (Scalar (Arrow _ _ tp tr)) (FuncDiet dp dr) = - contravariantArg tp dp && contravariantArg tr dr - contravariantArg _ _ = - True -consumeArg loc (Scalar (Arrow _ _ _ t2)) (FuncDiet _ pd) = - consumeArg loc t2 pd -consumeArg loc at Consume = return [consumption (aliases at) loc] -consumeArg loc at _ = return [observation (aliases at) loc] - --- | Type-check a single expression in isolation. This expression may --- turn out to be polymorphic, in which case the list of type --- parameters will be non-empty. -checkOneExp :: UncheckedExp -> TypeM ([TypeParam], Exp) -checkOneExp e = fmap fst . runTermTypeM $ do - e' <- checkExp e - let t = toStruct $ typeOf e' - (tparams, _, _, _) <- - letGeneralise (nameFromString "") (srclocOf e) [] [] t - fixOverloadedTypes - e'' <- updateTypes e' - checkUnmatched e'' - causalityCheck e'' - literalOverflowCheck e'' - return (tparams, e'') - --- Verify that all sum type constructors and empty array literals have --- a size that is known (rigid or a type parameter). This is to --- ensure that we can actually determine their shape at run-time. -causalityCheck :: Exp -> TermTypeM () -causalityCheck binding_body = do - constraints <- getConstraints - - let checkCausality what known t loc - | (d,dloc):_ <- mapMaybe (unknown constraints known) $ - S.toList $ typeDimNames $ toStruct t = - Just $ lift $ causality what loc d dloc t - | otherwise = Nothing - - checkParamCausality known p = - checkCausality (ppr p) known (patternType p) (srclocOf p) - - onExp :: S.Set VName -> Exp - -> StateT (S.Set VName) (Either TypeError) Exp - - onExp known (Var v (Info t) loc) - | Just bad <- checkCausality (pquote (ppr v)) known t loc = - bad - - onExp known (ArrayLit [] (Info t) loc) - | Just bad <- checkCausality "empty array" known t loc = - bad - - onExp known (Lambda params _ _ _ _) - | bad : _ <- mapMaybe (checkParamCausality known) params = - bad - - onExp known e@(Coerce what _ (_, Info ext) _) = do - modify (S.fromList ext<>) - void $ onExp known what - return e - - onExp known e@(LetPat _ bindee_e body_e (_, Info ext) _) = do - sequencePoint known bindee_e body_e ext - return e - - onExp known e@(Apply f arg (Info (_, p)) (_, Info ext) _) = do - sequencePoint known arg f $ maybeToList p ++ ext - return e - - onExp known e@(BinOp (f, floc) ft - (x, Info (_, xp)) (y, Info (_, yp)) _ (Info ext) _) = do - args_known <- lift $ - execStateT (sequencePoint known x y $ catMaybes [xp, yp]) mempty - void $ onExp (args_known<>known) (Var f ft floc) - modify ((args_known<>S.fromList ext)<>) - return e - - onExp known e = do - recurse known e - - case e of - DoLoop _ _ _ _ _ (Info (_, ext)) _ -> - modify (<>S.fromList ext) - If _ _ _ (_, Info ext) _ -> - modify (<>S.fromList ext) - Index _ _ (_, Info ext) _ -> - modify (<>S.fromList ext) - Match _ _ (_, Info ext) _ -> - modify (<>S.fromList ext) - Range _ _ _ (_, Info ext) _ -> - modify (<>S.fromList ext) - _ -> - return () - - return e - - recurse known = void . astMap mapper - where mapper = identityMapper { mapOnExp = onExp known } - - sequencePoint known x y ext = do - new_known <- lift $ execStateT (onExp known x) mempty - void $ onExp (new_known<>known) y - modify ((new_known<>S.fromList ext)<>) - - either throwError (const $ return ()) $ - evalStateT (onExp mempty binding_body) mempty - where unknown constraints known v = do - guard $ v `S.notMember` known - loc <- unknowable constraints v - return (v,loc) - - unknowable constraints v = - case snd <$> M.lookup v constraints of - Just (UnknowableSize loc _) -> Just loc - _ -> Nothing - - causality what loc d dloc t = - Left $ TypeError loc mempty $ - "Causality check: size" <+/> pquote (pprName d) <+/> - "needed for type of" <+> what <> colon - indent 2 (ppr t) - "But" <+> pquote (pprName d) <+> "is computed at" <+/> - text (locStrRel loc dloc) <> "." - "" - "Hint:" <+> - align (textwrap "Bind the expression producing" <+> pquote (pprName d) <+> - "with 'let' beforehand.") - --- | Traverse the expression, emitting warnings if any of the literals overflow --- their inferred types --- --- Note: currently unable to detect float underflow (such as 1e-400 -> 0) -literalOverflowCheck :: Exp -> TermTypeM () -literalOverflowCheck = void . check - where check e@(IntLit x ty loc) = e <$ case ty of - Info (Scalar (Prim t)) -> warnBounds (inBoundsI x t) x t loc - _ -> error "Inferred type of int literal is not a number" - check e@(FloatLit x ty loc) = e <$ case ty of - Info (Scalar (Prim (FloatType t))) -> warnBounds (inBoundsF x t) x t loc - _ -> error "Inferred type of float literal is not a float" - check e@(Negate (IntLit x ty loc1) loc2) = e <$ case ty of - Info (Scalar (Prim t)) -> warnBounds (inBoundsI (-x) t) (-x) t (loc1 <> loc2) - _ -> error "Inferred type of int literal is not a number" - check e = astMap identityMapper{mapOnExp = check} e - bitWidth ty = 8 * intByteSize ty :: Int - inBoundsI x (Signed t) = x >= -2^(bitWidth t - 1) && x < 2^(bitWidth t - 1) - inBoundsI x (Unsigned t) = x >= 0 && x < 2^bitWidth t - inBoundsI x (FloatType Float32) = not $ isInfinite (fromIntegral x :: Float) - inBoundsI x (FloatType Float64) = not $ isInfinite (fromIntegral x :: Double) - inBoundsI _ Bool = error "Inferred type of int literal is not a number" - inBoundsF x Float32 = not $ isInfinite (realToFrac x :: Float) - inBoundsF x Float64 = not $ isInfinite x - warnBounds inBounds x ty loc = unless inBounds - $ typeError loc mempty $ "Literal " <> ppr x <> - " out of bounds for inferred type " <> ppr ty <> "." - --- | Type-check a top-level (or module-level) function definition. --- Despite the name, this is also used for checking constant --- definitions, by treating them as 0-ary functions. -checkFunDef :: (Name, Maybe UncheckedTypeExp, - [UncheckedTypeParam], [UncheckedPattern], - UncheckedExp, SrcLoc) - -> TypeM (VName, [TypeParam], [Pattern], Maybe (TypeExp VName), - StructType, [VName], Exp) -checkFunDef (fname, maybe_retdecl, tparams, params, body, loc) = - fmap fst $ runTermTypeM $ do - (tparams', params', maybe_retdecl', rettype', retext, body') <- - checkBinding (fname, maybe_retdecl, tparams, params, body, loc) - - -- Since this is a top-level function, we also resolve overloaded - -- types, using either defaults or complaining about ambiguities. - fixOverloadedTypes - - -- Then replace all inferred types in the body and parameters. - body'' <- updateTypes body' - params'' <- updateTypes params' - maybe_retdecl'' <- traverse updateTypes maybe_retdecl' - rettype'' <- normTypeFully rettype' - - -- Check if pattern matches are exhaustive and yield - -- errors if not. - checkUnmatched body'' - - -- Check if the function body can actually be evaluated. - causalityCheck body'' - - literalOverflowCheck body'' - - bindSpaced [(Term, fname)] $ do - fname' <- checkName Term fname loc - when (nameToString fname `elem` doNotShadow) $ - typeError loc mempty $ - "The" <+> pprName fname <+> "operator may not be redefined." - - return (fname', tparams', params'', maybe_retdecl'', rettype'', retext, body'') - --- | This is "fixing" as in "setting them", not "correcting them". We --- only make very conservative fixing. -fixOverloadedTypes :: TermTypeM () -fixOverloadedTypes = getConstraints >>= mapM_ fixOverloaded . M.toList . M.map snd - where fixOverloaded (v, Overloaded ots usage) - | Signed Int32 `elem` ots = do - unify usage (Scalar (TypeVar () Nonunique (typeName v) [])) $ - Scalar $ Prim $ Signed Int32 - warn usage "Defaulting ambiguous type to i32." - | FloatType Float64 `elem` ots = do - unify usage (Scalar (TypeVar () Nonunique (typeName v) [])) $ - Scalar $ Prim $ FloatType Float64 - warn usage "Defaulting ambiguous type to f64." - | otherwise = - typeError usage mempty $ - "Type is ambiguous (could be one of" <+> commasep (map ppr ots) <> ")." - "Add a type annotation to disambiguate the type." - - fixOverloaded (_, NoConstraint _ usage) = - typeError usage mempty $ - "Type of expression is ambiguous." - "Add a type annotation to disambiguate the type." - - fixOverloaded (_, Equality usage) = - typeError usage mempty $ - "Type is ambiguous (must be equality type)." - "Add a type annotation to disambiguate the type." - - fixOverloaded (_, HasFields fs usage) = - typeError usage mempty $ - "Type is ambiguous. Must be record with fields:" - indent 2 (stack $ map field $ M.toList fs) - "Add a type annotation to disambiguate the type." - where field (l, t) = ppr l <> colon <+> align (ppr t) - - fixOverloaded (_, HasConstrs cs usage) = - typeError usage mempty $ - "Type is ambiguous (must be a sum type with constructors:" <+> - ppr (Sum cs) <> ")." - "Add a type annotation to disambiguate the type." - - fixOverloaded (_, Size Nothing usage) = - typeError usage mempty "Size is ambiguous." - - fixOverloaded _ = return () - -hiddenParamNames :: [Pattern] -> Names -hiddenParamNames params = hidden - where param_all_names = mconcat $ map patternNames params - named (Named x, _) = Just x - named (Unnamed, _) = Nothing - param_names = - S.fromList $ mapMaybe (named . patternParam) params - hidden = param_all_names `S.difference` param_names - -inferredReturnType :: SrcLoc -> [Pattern] -> PatternType -> TermTypeM StructType -inferredReturnType loc params t = - -- The inferred type may refer to names that are bound by the - -- parameter patterns, but which will not be visible in the type. - -- These we must turn into fresh type variables, which will be - -- existential in the return type. - fmap (toStruct . fst) $ - unscopeType loc - (M.filterWithKey (const . (`S.member` hidden)) $ foldMap patternMap params) $ - inferReturnUniqueness params t - where hidden = hiddenParamNames params - -checkBinding :: (Name, Maybe UncheckedTypeExp, - [UncheckedTypeParam], [UncheckedPattern], - UncheckedExp, SrcLoc) - -> TermTypeM ([TypeParam], [Pattern], Maybe (TypeExp VName), - StructType, [VName], Exp) -checkBinding (fname, maybe_retdecl, tparams, params, body, loc) = - noUnique $ incLevel $ bindingParams tparams params $ \tparams' params' -> do - when (null params && any isSizeParam tparams) $ - typeError loc mempty - "Size parameters are only allowed on bindings that also have value parameters." - - maybe_retdecl' <- forM maybe_retdecl $ \retdecl -> do - (retdecl', ret_nodims, _) <- checkTypeExp retdecl - (ret, _) <- instantiateEmptyArrayDims loc "funret" Nonrigid ret_nodims - return (retdecl', ret) - - body' <- checkFunBody params' body - (snd <$> maybe_retdecl') - (maybe loc srclocOf maybe_retdecl) - - params'' <- mapM updateTypes params' - body_t <- expTypeFully body' - - (maybe_retdecl'', rettype) <- case maybe_retdecl' of - Just (retdecl', ret) -> do - let rettype_structural = toStructural ret - checkReturnAlias rettype_structural params'' body_t - - when (null params) $ nothingMustBeUnique loc rettype_structural - - ret' <- normTypeFully ret - - return (Just retdecl', ret') - - Nothing - | null params -> - return (Nothing, toStruct $ body_t `setUniqueness` Nonunique) - | otherwise -> do - body_t' <- inferredReturnType loc params'' body_t - return (Nothing, body_t') - - verifyFunctionParams (Just fname) params'' - - (tparams'', params''', rettype'', retext) <- - letGeneralise fname loc tparams' params'' rettype - - checkGlobalAliases params'' body_t loc - - return (tparams'', params''', maybe_retdecl'', rettype'', retext, body') - - where -- | Check that unique return values do not alias a - -- non-consumed parameter. - checkReturnAlias rettp params' = - foldM_ (checkReturnAlias' params') S.empty . returnAliasing rettp - checkReturnAlias' params' seen (Unique, names) - | any (`S.member` S.map snd seen) $ S.toList names = - uniqueReturnAliased fname loc - | otherwise = do - notAliasingParam params' names - return $ seen `S.union` tag Unique names - checkReturnAlias' _ seen (Nonunique, names) - | any (`S.member` seen) $ S.toList $ tag Unique names = - uniqueReturnAliased fname loc - | otherwise = return $ seen `S.union` tag Nonunique names - - notAliasingParam params' names = - forM_ params' $ \p -> - let consumedNonunique p' = - not (unique $ unInfo $ identType p') && (identName p' `S.member` names) - in case find consumedNonunique $ S.toList $ patternIdents p of - Just p' -> - returnAliased fname (baseName $ identName p') loc - Nothing -> - return () - - tag u = S.map (u,) - - returnAliasing (Scalar (Record ets1)) (Scalar (Record ets2)) = - concat $ M.elems $ M.intersectionWith returnAliasing ets1 ets2 - returnAliasing expected got = - [(uniqueness expected, S.map aliasVar $ aliases got)] - --- | Extract all the shape names that occur in positive position --- (roughly, left side of an arrow) in a given type. -typeDimNamesPos :: TypeBase (DimDecl VName) als -> S.Set VName -typeDimNamesPos (Scalar (Arrow _ _ t1 t2)) = onParam t1 <> typeDimNamesPos t2 - where onParam :: TypeBase (DimDecl VName) als -> S.Set VName - onParam (Scalar Arrow{}) = mempty - onParam (Scalar (Record fs)) = mconcat $ map onParam $ M.elems fs - onParam (Scalar (TypeVar _ _ _ targs)) = mconcat $ map onTypeArg targs - onParam t = typeDimNames t - onTypeArg (TypeArgDim (NamedDim d) _) = S.singleton $ qualLeaf d - onTypeArg (TypeArgDim _ _) = mempty - onTypeArg (TypeArgType t _) = onParam t -typeDimNamesPos _ = mempty - -checkGlobalAliases :: [Pattern] -> PatternType -> SrcLoc -> TermTypeM () -checkGlobalAliases params body_t loc = do - vtable <- asks $ scopeVtable . termScope - let isLocal v = case v `M.lookup` vtable of - Just (BoundV Local _ _) -> True - _ -> False - let als = filter (not . isLocal) $ S.toList $ - boundArrayAliases body_t `S.difference` - foldMap patternNames params - case als of - v:_ | not $ null params -> - typeError loc mempty $ - "Function result aliases the free variable " <> - pquote (pprName v) <> "." - "Use" <+> pquote "copy" <+> "to break the aliasing." - _ -> - return () - -inferReturnUniqueness :: [Pattern] -> PatternType -> PatternType -inferReturnUniqueness params t = - let forbidden = aliasesMultipleTimes t - uniques = uniqueParamNames params - delve (Scalar (Record fs)) = - Scalar $ Record $ M.map delve fs - delve t' - | all (`S.member` uniques) (boundArrayAliases t'), - not $ any ((`S.member` forbidden) . aliasVar) (aliases t') = - t' - | otherwise = - t' `setUniqueness` Nonunique - in delve t - --- An alias inhibits uniqueness if it is used in disjoint values. -aliasesMultipleTimes :: PatternType -> Names -aliasesMultipleTimes = S.fromList . map fst . filter ((>1) . snd) . M.toList . delve - where delve (Scalar (Record fs)) = - foldl' (M.unionWith (+)) mempty $ map delve $ M.elems fs - delve t = - M.fromList $ zip (map aliasVar $ S.toList (aliases t)) $ repeat (1::Int) - -uniqueParamNames :: [Pattern] -> Names -uniqueParamNames = - S.map identName - . S.filter (unique . unInfo . identType) - . foldMap patternIdents - -boundArrayAliases :: PatternType -> S.Set VName -boundArrayAliases (Array als _ _ _) = boundAliases als -boundArrayAliases (Scalar Prim{}) = mempty -boundArrayAliases (Scalar (Record fs)) = foldMap boundArrayAliases fs -boundArrayAliases (Scalar (TypeVar als _ _ _)) = boundAliases als -boundArrayAliases (Scalar Arrow{}) = mempty -boundArrayAliases (Scalar (Sum fs)) = - mconcat $ concatMap (map boundArrayAliases) $ M.elems fs - --- | The set of in-scope variables that are being aliased. -boundAliases :: Aliasing -> S.Set VName -boundAliases = S.map aliasVar . S.filter bound - where bound AliasBound{} = True - bound AliasFree{} = False - -nothingMustBeUnique :: SrcLoc -> TypeBase () () -> TermTypeM () -nothingMustBeUnique loc = check - where check (Array _ Unique _ _) = bad - check (Scalar (TypeVar _ Unique _ _)) = bad - check (Scalar (Record fs)) = mapM_ check fs - check (Scalar (Sum fs)) = mapM_ (mapM_ check) fs - check _ = return () - bad = typeError loc mempty "A top-level constant cannot have a unique type." - --- | Verify certain restrictions on function parameters, and bail out --- on dubious constructions. --- --- These restrictions apply to all functions (anonymous or otherwise). --- Top-level functions have further restrictions that are checked --- during let-generalisation. -verifyFunctionParams :: Maybe Name -> [Pattern] -> TermTypeM () -verifyFunctionParams fname params = - onFailure (CheckingParams fname) $ - verifyParams (foldMap patternNames params) =<< mapM updateTypes params - where - verifyParams forbidden (p:ps) - | d:_ <- S.toList $ patternDimNames p `S.intersection` forbidden = - typeError p mempty $ - "Parameter" <+> pquote (ppr p) <+/> - "refers to size" <+> pquote (pprName d) <> comma <+/> - textwrap "which will not be accessible to the caller" <> comma <+/> - textwrap "possibly because it is nested in a tuple or record." <+/> - textwrap "Consider ascribing an explicit type that does not reference " <> - pquote (pprName d) <> "." - | otherwise = verifyParams forbidden' ps - where forbidden' = - case patternParam p of - (Named v, _) -> forbidden `S.difference` S.singleton v - _ -> forbidden - - verifyParams _ [] = return () - --- Returns the sizes of the immediate type produced, --- the sizes of parameter types, and the sizes of return types. -dimUses :: StructType -> (Names, Names, Names) -dimUses = execWriter . traverseDims f - where f _ PosImmediate (NamedDim v) = tell (S.singleton (qualLeaf v), mempty, mempty) - f _ PosParam (NamedDim v) = tell (mempty, S.singleton (qualLeaf v), mempty) - f _ PosReturn (NamedDim v) = tell (mempty, mempty, S.singleton (qualLeaf v)) - f _ _ _ = return () - --- | Find at all type variables in the given type that are covered by --- the constraints, and produce type parameters that close over them. --- --- The passed-in list of type parameters is always prepended to the --- produced list of type parameters. -closeOverTypes :: Name -> SrcLoc - -> [TypeParam] -> [StructType] -> StructType - -> Constraints -> TermTypeM ([TypeParam], StructType, [VName]) -closeOverTypes defname defloc tparams paramts ret substs = do - (more_tparams, retext) <- partitionEithers . catMaybes <$> - mapM closeOver (M.toList $ M.map snd to_close_over) - let retToAnyDim v = do guard $ v `S.member` ret_sizes - UnknowableSize{} <- snd <$> M.lookup v substs - Just $ SizeSubst AnyDim - return (tparams ++ more_tparams, - applySubst retToAnyDim ret, - retext) - where t = foldFunType paramts ret - to_close_over = M.filterWithKey (\k _ -> k `S.member` visible) substs - visible = typeVars t <> typeDimNames t - - (produced_sizes, param_sizes, ret_sizes) = dimUses t - - -- Avoid duplicate type parameters. - closeOver (k, _) - | k `elem` map typeParamName tparams = - return Nothing - closeOver (k, NoConstraint l usage) = - return $ Just $ Left $ TypeParamType l k $ srclocOf usage - closeOver (k, ParamType l loc) = - return $ Just $ Left $ TypeParamType l k loc - closeOver (k, Size Nothing usage) = - return $ Just $ Left $ TypeParamDim k $ srclocOf usage - closeOver (k, UnknowableSize _ _) - | k `S.member` param_sizes = do - notes <- dimNotes defloc $ NamedDim $ qualName k - typeError defloc notes $ - "Unknowable size" <+> pquote (pprName k) <+> - "imposes constraint on type of" <+> - pquote (pprName defname) <> - ", which is inferred as:" - indent 2 (ppr t) - | k `S.member` produced_sizes = - return $ Just $ Right k - closeOver (_, _) = - return Nothing - -letGeneralise :: Name -> SrcLoc - -> [TypeParam] -> [Pattern] -> StructType - -> TermTypeM ([TypeParam], [Pattern], StructType, [VName]) -letGeneralise defname defloc tparams params rettype = - onFailure (CheckingLetGeneralise defname) $ do - now_substs <- getConstraints - - -- Candidates for let-generalisation are those type variables that - -- - -- (1) were not known before we checked this function, and - -- - -- (2) are not used in the (new) definition of any type variables - -- known before we checked this function. - -- - -- (3) are not referenced from an overloaded type (for example, - -- are the element types of an incompletely resolved record type). - -- This is a bit more restrictive than I'd like, and SML for - -- example does not have this restriction. - -- - -- Criteria (1) and (2) is implemented by looking at the binding - -- level of the type variables. - let keep_type_vars = overloadedTypeVars now_substs - - cur_lvl <- curLevel - let candidate k (lvl, _) = (k `S.notMember` keep_type_vars) && lvl >= cur_lvl - new_substs = M.filterWithKey candidate now_substs - - (tparams', rettype', retext) <- - closeOverTypes defname defloc tparams - (map patternStructType params) rettype new_substs - - rettype'' <- updateTypes rettype' - - let used_sizes = foldMap typeDimNames $ - rettype'' : map patternStructType params - case filter ((`S.notMember` used_sizes) . typeParamName) $ - filter isSizeParam tparams' of - [] -> return () - tp:_ -> typeError defloc mempty $ - "Size parameter" <+> pquote (ppr tp) <+> "unused." - - -- We keep those type variables that were not closed over by - -- let-generalisation. - modifyConstraints $ M.filterWithKey $ \k _ -> k `notElem` map typeParamName tparams' - - return (tparams', params, rettype'', retext) - -checkFunBody :: [Pattern] - -> UncheckedExp - -> Maybe StructType - -> SrcLoc - -> TermTypeM Exp -checkFunBody params body maybe_rettype loc = do - body' <- noSizeEscape $ checkExp body - - -- Unify body return type with return annotation, if one exists. - case maybe_rettype of - Just rettype -> do - (rettype_withdims, _) <- instantiateEmptyArrayDims loc "impl" Nonrigid rettype - - body_t <- expTypeFully body' - -- We need to turn any sizes provided by "hidden" parameter - -- names into existential sizes instead. - let hidden = hiddenParamNames params - (body_t', _) <- unscopeType loc - (M.filterWithKey (const . (`S.member` hidden)) $ - foldMap patternMap params) - body_t - - let usage = mkUsage (srclocOf body) "return type annotation" - onFailure (CheckingReturn rettype (toStruct body_t')) $ - expect usage rettype_withdims $ toStruct body_t' - - -- We also have to make sure that uniqueness matches. This is done - -- explicitly, because uniqueness is ignored by unification. - rettype' <- normTypeFully rettype - body_t'' <- normTypeFully rettype -- Substs may have changed. - unless (body_t'' `subtypeOf` anySizes rettype') $ - typeError (srclocOf body) mempty $ - "Body type" indent 2 (ppr body_t'') - "is not a subtype of annotated type" - indent 2 (ppr rettype') - - Nothing -> return () - - return body' - ---- Consumption - -occur :: Occurences -> TermTypeM () -occur = tell - --- | Proclaim that we have made read-only use of the given variable. -observe :: Ident -> TermTypeM () -observe (Ident nm (Info t) loc) = - let als = AliasBound nm `S.insert` aliases t - in occur [observation als loc] - -checkIfConsumable :: SrcLoc -> Aliasing -> TermTypeM () -checkIfConsumable loc als = do - vtable <- asks $ scopeVtable . termScope - let consumable v = case M.lookup v vtable of - Just (BoundV Local _ t) - | arrayRank t > 0 -> unique t - | Scalar TypeVar{} <- t -> unique t - | otherwise -> True - _ -> False - case filter (not . consumable) $ map aliasVar $ S.toList als of - v:_ -> typeError loc mempty $ - "Would consume variable" <+> pquote (pprName v) - <> ", which is not allowed." - [] -> return () - --- | Proclaim that we have written to the given variable. -consume :: SrcLoc -> Aliasing -> TermTypeM () -consume loc als = do - checkIfConsumable loc als - occur [consumption als loc] - --- | Proclaim that we have written to the given variable, and mark --- accesses to it and all of its aliases as invalid inside the given --- computation. -consuming :: Ident -> TermTypeM a -> TermTypeM a -consuming (Ident name (Info t) loc) m = do - consume loc $ AliasBound name `S.insert` aliases t - localScope consume' m - where consume' scope = - scope { scopeVtable = M.insert name (WasConsumed loc) $ scopeVtable scope } - -collectOccurences :: TermTypeM a -> TermTypeM (a, Occurences) -collectOccurences m = pass $ do - (x, dataflow) <- listen m - return ((x, dataflow), const mempty) - -tapOccurences :: TermTypeM a -> TermTypeM (a, Occurences) -tapOccurences = listen - -removeSeminullOccurences :: TermTypeM a -> TermTypeM a -removeSeminullOccurences = censor $ filter $ not . seminullOccurence - -checkIfUsed :: Occurences -> Ident -> TermTypeM () -checkIfUsed occs v - | not $ identName v `S.member` allOccuring occs, - not $ "_" `isPrefixOf` prettyName (identName v) = - warn (srclocOf v) $ "Unused variable " ++ quote (pretty $ baseName $ identName v) ++ "." - | otherwise = - return () - -alternative :: TermTypeM a -> TermTypeM b -> TermTypeM (a,b) -alternative m1 m2 = pass $ do - (x, occurs1) <- listen $ noSizeEscape m1 - (y, occurs2) <- listen $ noSizeEscape m2 - checkOccurences occurs1 - checkOccurences occurs2 - let usage = occurs1 `altOccurences` occurs2 - return ((x, y), const usage) - --- | Make all bindings nonunique. -noUnique :: TermTypeM a -> TermTypeM a -noUnique = localScope (\scope -> scope { scopeVtable = M.map set $ scopeVtable scope}) - where set (BoundV l tparams t) = BoundV l tparams $ t `setUniqueness` Nonunique - set (OverloadedF ts pts rt) = OverloadedF ts pts rt - set EqualityF = EqualityF - set (WasConsumed loc) = WasConsumed loc - -onlySelfAliasing :: TermTypeM a -> TermTypeM a -onlySelfAliasing = localScope (\scope -> scope { scopeVtable = M.mapWithKey set $ scopeVtable scope}) - where set k (BoundV l tparams t) = BoundV l tparams $ - t `addAliases` S.intersection (S.singleton (AliasBound k)) - set _ (OverloadedF ts pts rt) = OverloadedF ts pts rt - set _ EqualityF = EqualityF - set _ (WasConsumed loc) = WasConsumed loc - -arrayOfM :: (Pretty (ShapeDecl dim), Monoid as) => - SrcLoc - -> TypeBase dim as -> ShapeDecl dim -> Uniqueness - -> TermTypeM (TypeBase dim as) -arrayOfM loc t shape u = do - zeroOrderType (mkUsage loc "use as array element") "type used in array" t - return $ arrayOf t shape u - -updateTypes :: ASTMappable e => e -> TermTypeM e -updateTypes = astMap tv - where tv = ASTMapper { mapOnExp = astMap tv - , mapOnName = pure - , mapOnQualName = pure - , mapOnStructType = normTypeFully - , mapOnPatternType = normTypeFully - } +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TupleSections #-} + +-- | Facilities for type-checking Futhark terms. Checking a term +-- requires a little more context to track uniqueness and such. +-- +-- Type inference is implemented through a variation of +-- Hindley-Milner. The main complication is supporting the rich +-- number of built-in language constructs, as well as uniqueness +-- types. This is mostly done in an ad hoc way, and many programs +-- will require the programmer to fall back on type annotations. +module Language.Futhark.TypeChecker.Terms + ( checkOneExp, + checkFunDef, + ) +where + +import Control.Monad.Except +import Control.Monad.RWS hiding (Sum) +import Control.Monad.State +import Control.Monad.Writer hiding (Sum) +import Data.Bifunctor +import Data.Char (isAscii) +import Data.Either +import Data.List (find, foldl', group, isPrefixOf, nub, sort, transpose, (\\)) +import qualified Data.List.NonEmpty as NE +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import Futhark.IR.Primitive (intByteSize) +import Futhark.Util.Pretty hiding (bool, group, space) +import Language.Futhark hiding (unscopeType) +import Language.Futhark.Semantic (includeToString) +import Language.Futhark.Traversals +import Language.Futhark.TypeChecker.Monad hiding (BoundV) +import qualified Language.Futhark.TypeChecker.Monad as TypeM +import Language.Futhark.TypeChecker.Types hiding (checkTypeDecl) +import qualified Language.Futhark.TypeChecker.Types as Types +import Language.Futhark.TypeChecker.Unify hiding (Usage) +import Prelude hiding (mod) + +--- Uniqueness + +data Usage + = Consumed SrcLoc + | Observed SrcLoc + deriving (Eq, Ord, Show) + +type Names = S.Set VName + +-- | The consumption set is a Maybe so we can distinguish whether a +-- consumption took place, but the variable went out of scope since, +-- or no consumption at all took place. +data Occurence = Occurence + { observed :: Names, + consumed :: Maybe Names, + location :: SrcLoc + } + deriving (Eq, Show) + +instance Located Occurence where + locOf = locOf . location + +observation :: Aliasing -> SrcLoc -> Occurence +observation = flip Occurence Nothing . S.map aliasVar + +consumption :: Aliasing -> SrcLoc -> Occurence +consumption = Occurence S.empty . Just . S.map aliasVar + +-- | A null occurence is one that we can remove without affecting +-- anything. +nullOccurence :: Occurence -> Bool +nullOccurence occ = S.null (observed occ) && isNothing (consumed occ) + +-- | A seminull occurence is one that does not contain references to +-- any variables in scope. The big difference is that a seminull +-- occurence may denote a consumption, as long as the array that was +-- consumed is now out of scope. +seminullOccurence :: Occurence -> Bool +seminullOccurence occ = S.null (observed occ) && maybe True S.null (consumed occ) + +type Occurences = [Occurence] + +type UsageMap = M.Map VName [Usage] + +usageMap :: Occurences -> UsageMap +usageMap = foldl comb M.empty + where + comb m (Occurence obs cons loc) = + let m' = S.foldl' (ins $ Observed loc) m obs + in S.foldl' (ins $ Consumed loc) m' $ fromMaybe mempty cons + ins v m k = M.insertWith (++) k [v] m + +combineOccurences :: VName -> Usage -> Usage -> TermTypeM Usage +combineOccurences _ (Observed loc) (Observed _) = return $ Observed loc +combineOccurences name (Consumed wloc) (Observed rloc) = + useAfterConsume (baseName name) rloc wloc +combineOccurences name (Observed rloc) (Consumed wloc) = + useAfterConsume (baseName name) rloc wloc +combineOccurences name (Consumed loc1) (Consumed loc2) = + consumeAfterConsume (baseName name) (max loc1 loc2) (min loc1 loc2) + +checkOccurences :: Occurences -> TermTypeM () +checkOccurences = void . M.traverseWithKey comb . usageMap + where + comb _ [] = return () + comb name (u : us) = foldM_ (combineOccurences name) u us + +allObserved :: Occurences -> Names +allObserved = S.unions . map observed + +allConsumed :: Occurences -> Names +allConsumed = S.unions . map (fromMaybe mempty . consumed) + +allOccuring :: Occurences -> Names +allOccuring occs = allConsumed occs <> allObserved occs + +anyConsumption :: Occurences -> Maybe Occurence +anyConsumption = find (isJust . consumed) + +seqOccurences :: Occurences -> Occurences -> Occurences +seqOccurences occurs1 occurs2 = + filter (not . nullOccurence) $ map filt occurs1 ++ occurs2 + where + filt occ = + occ {observed = observed occ `S.difference` postcons} + postcons = allConsumed occurs2 + +altOccurences :: Occurences -> Occurences -> Occurences +altOccurences occurs1 occurs2 = + filter (not . nullOccurence) $ map filt1 occurs1 ++ map filt2 occurs2 + where + filt1 occ = + occ + { consumed = S.difference <$> consumed occ <*> pure cons2, + observed = observed occ `S.difference` cons2 + } + filt2 occ = + occ + { consumed = consumed occ, + observed = observed occ `S.difference` cons1 + } + cons1 = allConsumed occurs1 + cons2 = allConsumed occurs2 + +--- Scope management + +data Checking + = CheckingApply (Maybe (QualName VName)) Exp StructType StructType + | CheckingReturn StructType StructType + | CheckingAscription StructType StructType + | CheckingLetGeneralise Name + | CheckingParams (Maybe Name) + | CheckingPattern UncheckedPattern InferredType + | CheckingLoopBody StructType StructType + | CheckingLoopInitial StructType StructType + | CheckingRecordUpdate [Name] StructType StructType + | CheckingRequired [StructType] StructType + | CheckingBranches StructType StructType + +instance Pretty Checking where + ppr (CheckingApply f e expected actual) = + header + "Expected:" <+> align (ppr expected) + "Actual: " <+> align (ppr actual) + where + header = + case f of + Nothing -> + "Cannot apply function to" + <+> pquote (shorten $ pretty $ flatten $ ppr e) <> " (invalid type)." + Just fname -> + "Cannot apply" <+> pquote (ppr fname) <+> "to" + <+> pquote (shorten $ pretty $ flatten $ ppr e) <> " (invalid type)." + ppr (CheckingReturn expected actual) = + "Function body does not have expected type." + "Expected:" <+> align (ppr expected) + "Actual: " <+> align (ppr actual) + ppr (CheckingAscription expected actual) = + "Expression does not have expected type from explicit ascription." + "Expected:" <+> align (ppr expected) + "Actual: " <+> align (ppr actual) + ppr (CheckingLetGeneralise fname) = + "Cannot generalise type of" <+> pquote (ppr fname) <> "." + ppr (CheckingParams fname) = + "Invalid use of parameters in" <+> pquote fname' <> "." + where + fname' = maybe "anonymous function" ppr fname + ppr (CheckingPattern pat NoneInferred) = + "Invalid pattern" <+> pquote (ppr pat) <> "." + ppr (CheckingPattern pat (Ascribed t)) = + "Pattern" <+> pquote (ppr pat) + <+> "cannot match value of type" + indent 2 (ppr t) + ppr (CheckingLoopBody expected actual) = + "Loop body does not have expected type." + "Expected:" <+> align (ppr expected) + "Actual: " <+> align (ppr actual) + ppr (CheckingLoopInitial expected actual) = + "Initial loop values do not have expected type." + "Expected:" <+> align (ppr expected) + "Actual: " <+> align (ppr actual) + ppr (CheckingRecordUpdate fs expected actual) = + "Type mismatch when updating record field" <+> pquote fs' <> "." + "Existing:" <+> align (ppr expected) + "New: " <+> align (ppr actual) + where + fs' = mconcat $ punctuate "." $ map ppr fs + ppr (CheckingRequired [expected] actual) = + "Expression must must have type" <+> ppr expected <> "." + "Actual type:" <+> align (ppr actual) + ppr (CheckingRequired expected actual) = + "Type of expression must must be one of " <+> expected' <> "." + "Actual type:" <+> align (ppr actual) + where + expected' = commasep (map ppr expected) + ppr (CheckingBranches t1 t2) = + "Conditional branches differ in type." + "Former:" <+> ppr t1 + "Latter:" <+> ppr t2 + +-- | Whether something is a global or a local variable. +data Locality = Local | Global + deriving (Show) + +data ValBinding + = -- | Aliases in parameters indicate the lexical + -- closure. + BoundV Locality [TypeParam] PatternType + | OverloadedF [PrimType] [Maybe PrimType] (Maybe PrimType) + | EqualityF + | WasConsumed SrcLoc + deriving (Show) + +-- | Type checking happens with access to this environment. The +-- 'TermScope' will be extended during type-checking as bindings come into +-- scope. +data TermEnv = TermEnv + { termScope :: TermScope, + termChecking :: Maybe Checking, + termLevel :: Level + } + +data TermScope = TermScope + { scopeVtable :: M.Map VName ValBinding, + scopeTypeTable :: M.Map VName TypeBinding, + scopeModTable :: M.Map VName Mod, + scopeNameMap :: NameMap + } + deriving (Show) + +instance Semigroup TermScope where + TermScope vt1 tt1 mt1 nt1 <> TermScope vt2 tt2 mt2 nt2 = + TermScope (vt2 `M.union` vt1) (tt2 `M.union` tt1) (mt1 `M.union` mt2) (nt2 `M.union` nt1) + +envToTermScope :: Env -> TermScope +envToTermScope env = + TermScope + { scopeVtable = vtable, + scopeTypeTable = envTypeTable env, + scopeNameMap = envNameMap env, + scopeModTable = envModTable env + } + where + vtable = M.mapWithKey valBinding $ envVtable env + valBinding k (TypeM.BoundV tps v) = + BoundV Global tps $ + v + `setAliases` (if arrayRank v > 0 then S.singleton (AliasBound k) else mempty) + +withEnv :: TermEnv -> Env -> TermEnv +withEnv tenv env = tenv {termScope = termScope tenv <> envToTermScope env} + +overloadedTypeVars :: Constraints -> Names +overloadedTypeVars = mconcat . map f . M.elems + where + f (_, HasFields fs _) = mconcat $ map typeVars $ M.elems fs + f _ = mempty + +-- | Get the type of an expression, with top level type variables +-- substituted. Never call 'typeOf' directly (except in a few +-- carefully inspected locations)! +expType :: Exp -> TermTypeM PatternType +expType = normPatternType . typeOf + +-- | Get the type of an expression, with all type variables +-- substituted. Slower than 'expType', but sometimes necessary. +-- Never call 'typeOf' directly (except in a few carefully inspected +-- locations)! +expTypeFully :: Exp -> TermTypeM PatternType +expTypeFully = normTypeFully . typeOf + +-- Wrap a function name to give it a vacuous Eq instance for SizeSource. +newtype FName = FName (Maybe (QualName VName)) + deriving (Show) + +instance Eq FName where + _ == _ = True + +instance Ord FName where + compare _ _ = EQ + +-- | What was the source of some existential size? This is used for +-- using the same existential variable if the same source is +-- encountered in multiple locations. +data SizeSource + = SourceArg FName (ExpBase NoInfo VName) + | SourceBound (ExpBase NoInfo VName) + | SourceSlice + (Maybe (DimDecl VName)) + (Maybe (ExpBase NoInfo VName)) + (Maybe (ExpBase NoInfo VName)) + (Maybe (ExpBase NoInfo VName)) + deriving (Eq, Ord, Show) + +-- | The state is a set of constraints and a counter for generating +-- type names. This is distinct from the usual counter we use for +-- generating unique names, as these will be user-visible. +data TermTypeState = TermTypeState + { stateConstraints :: Constraints, + stateCounter :: !Int, + -- | Mapping function arguments encountered to + -- the sizes they ended up generating (when + -- they could not be substituted directly). + -- This happens for function arguments that are + -- not constants or names. + stateDimTable :: M.Map SizeSource VName + } + +newtype TermTypeM a + = TermTypeM + ( RWST + TermEnv + Occurences + TermTypeState + TypeM + a + ) + deriving + ( Monad, + Functor, + Applicative, + MonadReader TermEnv, + MonadWriter Occurences, + MonadState TermTypeState, + MonadError TypeError + ) + +instance MonadUnify TermTypeM where + getConstraints = gets stateConstraints + putConstraints x = modify $ \s -> s {stateConstraints = x} + + newTypeVar loc desc = do + i <- incCounter + v <- newID $ mkTypeVarName desc i + constrain v $ NoConstraint Lifted $ mkUsage' loc + return $ Scalar $ TypeVar mempty Nonunique (typeName v) [] + + curLevel = asks termLevel + + newDimVar loc rigidity name = do + i <- incCounter + dim <- newID $ mkTypeVarName name i + case rigidity of + Rigid rsrc -> constrain dim $ UnknowableSize loc rsrc + Nonrigid -> constrain dim $ Size Nothing $ mkUsage' loc + return dim + + unifyError loc notes bcs doc = do + checking <- asks termChecking + case checking of + Just checking' -> + throwError $ + TypeError (srclocOf loc) notes $ + ppr checking' <> line doc <> ppr bcs + Nothing -> + throwError $ TypeError (srclocOf loc) notes $ doc <> ppr bcs + + matchError loc notes bcs t1 t2 = do + checking <- asks termChecking + case checking of + Just checking' + | hasNoBreadCrumbs bcs -> + throwError $ + TypeError (srclocOf loc) notes $ + ppr checking' + | otherwise -> + throwError $ + TypeError (srclocOf loc) notes $ + ppr checking' <> line doc <> ppr bcs + Nothing -> + throwError $ TypeError (srclocOf loc) notes $ doc <> ppr bcs + where + doc = + "Types" + indent 2 (ppr t1) + "and" + indent 2 (ppr t2) + "do not match." + +onFailure :: Checking -> TermTypeM a -> TermTypeM a +onFailure c = local $ \env -> env {termChecking = Just c} + +runTermTypeM :: TermTypeM a -> TypeM (a, Occurences) +runTermTypeM (TermTypeM m) = do + initial_scope <- (initialTermScope <>) . envToTermScope <$> askEnv + let initial_tenv = + TermEnv + { termScope = initial_scope, + termChecking = Nothing, + termLevel = 0 + } + evalRWST m initial_tenv $ TermTypeState mempty 0 mempty + +liftTypeM :: TypeM a -> TermTypeM a +liftTypeM = TermTypeM . lift + +localScope :: (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a +localScope f = local $ \tenv -> tenv {termScope = f $ termScope tenv} + +incCounter :: TermTypeM Int +incCounter = do + s <- get + put s {stateCounter = stateCounter s + 1} + return $ stateCounter s + +extSize :: SrcLoc -> SizeSource -> TermTypeM (DimDecl VName, Maybe VName) +extSize loc e = do + prev <- gets $ M.lookup e . stateDimTable + case prev of + Nothing -> do + let rsrc = case e of + SourceArg (FName fname) e' -> + RigidArg fname $ prettyOneLine e' + SourceBound e' -> + RigidBound $ prettyOneLine e' + SourceSlice d i j s -> + RigidSlice d $ prettyOneLine $ DimSlice i j s + d <- newDimVar loc (Rigid rsrc) "argdim" + modify $ \s -> s {stateDimTable = M.insert e d $ stateDimTable s} + return + ( NamedDim $ qualName d, + Just d + ) + Just d -> + return + ( NamedDim $ qualName d, + Nothing + ) + +-- Any argument sizes created with 'extSize' inside the given action +-- will be removed once the action finishes. This is to ensure that +-- just because e.g. @n+1@ appears as a size in one branch of a +-- conditional, that doesn't mean it's also available in the other branch. +noSizeEscape :: TermTypeM a -> TermTypeM a +noSizeEscape m = do + dimtable <- gets stateDimTable + x <- m + modify $ \s -> s {stateDimTable = dimtable} + return x + +constrain :: VName -> Constraint -> TermTypeM () +constrain v c = do + lvl <- curLevel + modifyConstraints $ M.insert v (lvl, c) + +incLevel :: TermTypeM a -> TermTypeM a +incLevel = local $ \env -> env {termLevel = termLevel env + 1} + +initialTermScope :: TermScope +initialTermScope = + TermScope + { scopeVtable = initialVtable, + scopeTypeTable = mempty, + scopeNameMap = topLevelNameMap, + scopeModTable = mempty + } + where + initialVtable = M.fromList $ mapMaybe addIntrinsicF $ M.toList intrinsics + + prim = Scalar . Prim + arrow x y = Scalar $ Arrow mempty Unnamed x y + + addIntrinsicF (name, IntrinsicMonoFun pts t) = + Just (name, BoundV Global [] $ arrow pts' $ prim t) + where + pts' = case pts of + [pt] -> prim pt + _ -> tupleRecord $ map prim pts + addIntrinsicF (name, IntrinsicOverloadedFun ts pts rts) = + Just (name, OverloadedF ts pts rts) + addIntrinsicF (name, IntrinsicPolyFun tvs pts rt) = + Just + ( name, + BoundV Global tvs $ + fromStruct $ Scalar $ Arrow mempty Unnamed pts' rt + ) + where + pts' = case pts of + [pt] -> pt + _ -> tupleRecord pts + addIntrinsicF (name, IntrinsicEquality) = + Just (name, EqualityF) + addIntrinsicF _ = Nothing + +instance MonadTypeChecker TermTypeM where + warn loc problem = liftTypeM $ warn loc problem + newName = liftTypeM . newName + newID = liftTypeM . newID + + checkQualName space name loc = snd <$> checkQualNameWithEnv space name loc + + bindNameMap m = localScope $ \scope -> + scope {scopeNameMap = m <> scopeNameMap scope} + + bindVal v (TypeM.BoundV tps t) = localScope $ \scope -> + scope {scopeVtable = M.insert v vb $ scopeVtable scope} + where + vb = BoundV Local tps $ fromStruct t + + lookupType loc qn = do + outer_env <- liftTypeM askEnv + (scope, qn'@(QualName qs name)) <- checkQualNameWithEnv Type qn loc + case M.lookup name $ scopeTypeTable scope of + Nothing -> unknownType loc qn + Just (TypeAbbr l ps def) -> + return (qn', ps, qualifyTypeVars outer_env (map typeParamName ps) qs def, l) + + lookupMod loc qn = do + (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Term qn loc + case M.lookup name $ scopeModTable scope of + Nothing -> unknownVariable Term qn loc + Just m -> return (qn', m) + + lookupVar loc qn = do + outer_env <- liftTypeM askEnv + (scope, qn'@(QualName qs name)) <- checkQualNameWithEnv Term qn loc + let usage = mkUsage loc $ "use of " ++ quote (pretty qn) + + t <- case M.lookup name $ scopeVtable scope of + Nothing -> + typeError loc mempty $ + "Unknown variable" <+> pquote (ppr qn) <> "." + Just (WasConsumed wloc) -> useAfterConsume (baseName name) loc wloc + Just (BoundV _ tparams t) + | "_" `isPrefixOf` baseString name -> underscoreUse loc qn + | otherwise -> do + (tnames, t') <- instantiateTypeScheme loc tparams t + return $ qualifyTypeVars outer_env tnames qs t' + Just EqualityF -> do + argtype <- newTypeVar loc "t" + equalityType usage argtype + return $ + Scalar $ + Arrow mempty Unnamed argtype $ + Scalar $ Arrow mempty Unnamed argtype $ Scalar $ Prim Bool + Just (OverloadedF ts pts rt) -> do + argtype <- newTypeVar loc "t" + mustBeOneOf ts usage argtype + let (pts', rt') = instOverloaded argtype pts rt + arrow xt yt = Scalar $ Arrow mempty Unnamed xt yt + return $ fromStruct $ foldr arrow rt' pts' + + observe $ Ident name (Info t) loc + return (qn', t) + where + instOverloaded argtype pts rt = + ( map (maybe (toStruct argtype) (Scalar . Prim)) pts, + maybe (toStruct argtype) (Scalar . Prim) rt + ) + + checkNamedDim loc v = do + (v', t) <- lookupVar loc v + onFailure (CheckingRequired [Scalar $ Prim $ Signed Int64] (toStruct t)) $ + unify (mkUsage loc "use as array size") (toStruct t) $ + Scalar $ Prim $ Signed Int64 + return v' + + typeError loc notes s = do + checking <- asks termChecking + case checking of + Just checking' -> + throwError $ TypeError (srclocOf loc) notes (ppr checking' <> line s) + Nothing -> + throwError $ TypeError (srclocOf loc) notes s + +checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) +checkQualNameWithEnv space qn@(QualName quals name) loc = do + scope <- asks termScope + descend scope quals + where + descend scope [] + | Just name' <- M.lookup (space, name) $ scopeNameMap scope = + return (scope, name') + | otherwise = + unknownVariable space qn loc + descend scope (q : qs) + | Just (QualName _ q') <- M.lookup (Term, q) $ scopeNameMap scope, + Just res <- M.lookup q' $ scopeModTable scope = + case res of + -- Check if we are referring to the magical intrinsics + -- module. + _ + | baseTag q' <= maxIntrinsicTag -> + checkIntrinsic space qn loc + ModEnv q_scope -> do + (scope', QualName qs' name') <- descend (envToTermScope q_scope) qs + return (scope', QualName (q' : qs') name') + ModFun {} -> unappliedFunctor loc + | otherwise = + unknownVariable space qn loc + +checkIntrinsic :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) +checkIntrinsic space qn@(QualName _ name) loc + | Just v <- M.lookup (space, name) intrinsicsNameMap = do + me <- liftTypeM askImportName + unless ("/prelude" `isPrefixOf` includeToString me) $ + warn loc "Using intrinsic functions directly can easily crash the compiler or result in wrong code generation." + scope <- asks termScope + return (scope, v) + | otherwise = + unknownVariable space qn loc + +-- | Wrap 'Types.checkTypeDecl' to also perform an observation of +-- every size in the type. +checkTypeDecl :: TypeDeclBase NoInfo Name -> TermTypeM (TypeDeclBase Info VName) +checkTypeDecl tdecl = do + (tdecl', _) <- Types.checkTypeDecl tdecl + mapM_ observeDim $ nestedDims $ unInfo $ expandedType tdecl' + return tdecl' + where + observeDim (NamedDim v) = + observe $ Ident (qualLeaf v) (Info $ Scalar $ Prim $ Signed Int64) mempty + observeDim _ = return () + +-- | Instantiate a type scheme with fresh type variables for its type +-- parameters. Returns the names of the fresh type variables, the +-- instance list, and the instantiated type. +instantiateTypeScheme :: + SrcLoc -> + [TypeParam] -> + PatternType -> + TermTypeM ([VName], PatternType) +instantiateTypeScheme loc tparams t = do + let tnames = map typeParamName tparams + (tparam_names, tparam_substs) <- unzip <$> mapM (instantiateTypeParam loc) tparams + let substs = M.fromList $ zip tnames tparam_substs + t' = applySubst (`M.lookup` substs) t + return (tparam_names, t') + +-- | Create a new type name and insert it (unconstrained) in the +-- substitution map. +instantiateTypeParam :: Monoid as => SrcLoc -> TypeParam -> TermTypeM (VName, Subst (TypeBase dim as)) +instantiateTypeParam loc tparam = do + i <- incCounter + v <- newID $ mkTypeVarName (takeWhile isAscii (baseString (typeParamName tparam))) i + case tparam of + TypeParamType x _ _ -> do + constrain v $ NoConstraint x $ mkUsage' loc + return (v, Subst $ Scalar $ TypeVar mempty Nonunique (typeName v) []) + TypeParamDim {} -> do + constrain v $ Size Nothing $ mkUsage' loc + return (v, SizeSubst $ NamedDim $ qualName v) + +newArrayType :: SrcLoc -> String -> Int -> TermTypeM (StructType, StructType) +newArrayType loc desc r = do + v <- newID $ nameFromString desc + constrain v $ NoConstraint Unlifted $ mkUsage' loc + dims <- replicateM r $ newDimVar loc Nonrigid "dim" + let rowt = TypeVar () Nonunique (typeName v) [] + return + ( Array () Nonunique rowt (ShapeDecl $ map (NamedDim . qualName) dims), + Scalar rowt + ) + +--- Errors + +useAfterConsume :: Name -> SrcLoc -> SrcLoc -> TermTypeM a +useAfterConsume name rloc wloc = + typeError rloc mempty $ + "Variable" <+> pquote (pprName name) <+> "previously consumed at" + <+> text (locStrRel rloc wloc) <> ". (Possibly through aliasing.)" + +consumeAfterConsume :: Name -> SrcLoc -> SrcLoc -> TermTypeM a +consumeAfterConsume name loc1 loc2 = + typeError loc2 mempty $ + "Variable" <+> pprName name <+> "previously consumed at" + <+> text (locStrRel loc2 loc1) <> "." + +badLetWithValue :: SrcLoc -> TermTypeM a +badLetWithValue loc = + typeError + loc + mempty + "New value for elements in let-with shares data with source array. This is illegal, as it prevents in-place modification." + +returnAliased :: Name -> Name -> SrcLoc -> TermTypeM () +returnAliased fname name loc = + typeError loc mempty $ + "Unique return value of" <+> pquote (pprName fname) + <+> "is aliased to" + <+> pquote (pprName name) <> ", which is not consumed." + +uniqueReturnAliased :: Name -> SrcLoc -> TermTypeM a +uniqueReturnAliased fname loc = + typeError loc mempty $ + "A unique tuple element of return value of" + <+> pquote (pprName fname) + <+> "is aliased to some other tuple component." + +unexpectedType :: MonadTypeChecker m => SrcLoc -> StructType -> [StructType] -> m a +unexpectedType loc _ [] = + typeError loc mempty $ + "Type of expression at" <+> text (locStr loc) + <+> "cannot have any type - possibly a bug in the type checker." +unexpectedType loc t ts = + typeError loc mempty $ + "Type of expression at" <+> text (locStr loc) <+> "must be one of" + <+> commasep (map ppr ts) <> ", but is" + <+> ppr t <> "." + +--- Basic checking + +-- | Determine if the two types of identical, ignoring uniqueness. +-- Mismatched dimensions are turned into fresh rigid type variables. +-- Causes a 'TypeError' if they fail to match, and otherwise returns +-- one of them. +unifyBranchTypes :: SrcLoc -> PatternType -> PatternType -> TermTypeM (PatternType, [VName]) +unifyBranchTypes loc t1 t2 = + onFailure (CheckingBranches (toStruct t1) (toStruct t2)) $ + unifyMostCommon (mkUsage loc "unification of branch results") t1 t2 + +unifyBranches :: SrcLoc -> Exp -> Exp -> TermTypeM (PatternType, [VName]) +unifyBranches loc e1 e2 = do + e1_t <- expTypeFully e1 + e2_t <- expTypeFully e2 + unifyBranchTypes loc e1_t e2_t + +--- General binding. + +doNotShadow :: [String] +doNotShadow = ["&&", "||"] + +data InferredType + = NoneInferred + | Ascribed PatternType + +checkPattern' :: + UncheckedPattern -> + InferredType -> + TermTypeM Pattern +checkPattern' (PatternParens p loc) t = + PatternParens <$> checkPattern' p t <*> pure loc +checkPattern' (Id name _ loc) _ + | name' `elem` doNotShadow = + typeError loc mempty $ "The" <+> text name' <+> "operator may not be redefined." + where + name' = nameToString name +checkPattern' (Id name NoInfo loc) (Ascribed t) = do + name' <- newID name + return $ Id name' (Info t) loc +checkPattern' (Id name NoInfo loc) NoneInferred = do + name' <- newID name + t <- newTypeVar loc "t" + return $ Id name' (Info t) loc +checkPattern' (Wildcard _ loc) (Ascribed t) = + return $ Wildcard (Info $ t `setUniqueness` Nonunique) loc +checkPattern' (Wildcard NoInfo loc) NoneInferred = do + t <- newTypeVar loc "t" + return $ Wildcard (Info t) loc +checkPattern' (TuplePattern ps loc) (Ascribed t) + | Just ts <- isTupleRecord t, + length ts == length ps = + TuplePattern <$> zipWithM checkPattern' ps (map Ascribed ts) <*> pure loc +checkPattern' p@(TuplePattern ps loc) (Ascribed t) = do + ps_t <- replicateM (length ps) (newTypeVar loc "t") + unify (mkUsage loc "matching a tuple pattern") (tupleRecord ps_t) $ toStruct t + t' <- normTypeFully t + checkPattern' p $ Ascribed t' +checkPattern' (TuplePattern ps loc) NoneInferred = + TuplePattern <$> mapM (`checkPattern'` NoneInferred) ps <*> pure loc +checkPattern' (RecordPattern p_fs _) _ + | Just (f, fp) <- find (("_" `isPrefixOf`) . nameToString . fst) p_fs = + typeError fp mempty $ + "Underscore-prefixed fields are not allowed." + "Did you mean" <> dquotes (text (drop 1 (nameToString f)) <> "=_") <> "?" +checkPattern' (RecordPattern p_fs loc) (Ascribed (Scalar (Record t_fs))) + | sort (map fst p_fs) == sort (M.keys t_fs) = + RecordPattern . M.toList <$> check <*> pure loc + where + check = + traverse (uncurry checkPattern') $ + M.intersectionWith + (,) + (M.fromList p_fs) + (fmap Ascribed t_fs) +checkPattern' p@(RecordPattern fields loc) (Ascribed t) = do + fields' <- traverse (const $ newTypeVar loc "t") $ M.fromList fields + + when (sort (M.keys fields') /= sort (map fst fields)) $ + typeError loc mempty $ "Duplicate fields in record pattern" <+> ppr p <> "." + + unify (mkUsage loc "matching a record pattern") (Scalar (Record fields')) $ toStruct t + t' <- normTypeFully t + checkPattern' p $ Ascribed t' +checkPattern' (RecordPattern fs loc) NoneInferred = + RecordPattern . M.toList <$> traverse (`checkPattern'` NoneInferred) (M.fromList fs) <*> pure loc +checkPattern' (PatternAscription p (TypeDecl t NoInfo) loc) maybe_outer_t = do + (t', st_nodims, _) <- checkTypeExp t + (st, _) <- instantiateEmptyArrayDims loc "impl" Nonrigid st_nodims + + let st' = fromStruct st + case maybe_outer_t of + Ascribed outer_t -> do + unify (mkUsage loc "explicit type ascription") (toStruct st) (toStruct outer_t) + + -- We also have to make sure that uniqueness matches. This is + -- done explicitly, because it is ignored by unification. + st'' <- normTypeFully st' + outer_t' <- normTypeFully outer_t + case unifyTypesU unifyUniqueness st'' outer_t' of + Just outer_t'' -> + PatternAscription <$> checkPattern' p (Ascribed outer_t'') + <*> pure (TypeDecl t' (Info st)) + <*> pure loc + Nothing -> + typeError loc mempty $ + "Cannot match type" <+> pquote (ppr outer_t') <+> "with expected type" + <+> pquote (ppr st'') <> "." + NoneInferred -> + PatternAscription <$> checkPattern' p (Ascribed st') + <*> pure (TypeDecl t' (Info st)) + <*> pure loc + where + unifyUniqueness u1 u2 = if u2 `subuniqueOf` u1 then Just u1 else Nothing +checkPattern' (PatternLit e NoInfo loc) (Ascribed t) = do + e' <- checkExp e + t' <- expTypeFully e' + unify (mkUsage loc "matching against literal") (toStruct t') (toStruct t) + return $ PatternLit e' (Info t') loc +checkPattern' (PatternLit e NoInfo loc) NoneInferred = do + e' <- checkExp e + t' <- expTypeFully e' + return $ PatternLit e' (Info t') loc +checkPattern' (PatternConstr n NoInfo ps loc) (Ascribed (Scalar (Sum cs))) + | Just ts <- M.lookup n cs = do + ps' <- zipWithM checkPattern' ps $ map Ascribed ts + return $ PatternConstr n (Info (Scalar (Sum cs))) ps' loc +checkPattern' (PatternConstr n NoInfo ps loc) (Ascribed t) = do + t' <- newTypeVar loc "t" + ps' <- mapM (`checkPattern'` NoneInferred) ps + mustHaveConstr usage n t' (patternStructType <$> ps') + unify usage t' (toStruct t) + t'' <- normTypeFully t + return $ PatternConstr n (Info t'') ps' loc + where + usage = mkUsage loc "matching against constructor" +checkPattern' (PatternConstr n NoInfo ps loc) NoneInferred = do + ps' <- mapM (`checkPattern'` NoneInferred) ps + t <- newTypeVar loc "t" + mustHaveConstr usage n t (patternStructType <$> ps') + return $ PatternConstr n (Info $ fromStruct t) ps' loc + where + usage = mkUsage loc "matching against constructor" + +patternNameMap :: Pattern -> NameMap +patternNameMap = M.fromList . map asTerm . S.toList . patternNames + where + asTerm v = ((Term, baseName v), qualName v) + +checkPattern :: + UncheckedPattern -> + InferredType -> + (Pattern -> TermTypeM a) -> + TermTypeM a +checkPattern p t m = do + checkForDuplicateNames [p] + p' <- onFailure (CheckingPattern p t) $ checkPattern' p t + bindNameMap (patternNameMap p') $ m p' + +binding :: [Ident] -> TermTypeM a -> TermTypeM a +binding bnds = check . handleVars + where + handleVars m = + localScope (`bindVars` bnds) $ do + -- Those identifiers that can potentially also be sizes are + -- added as type constraints. This is necessary so that we + -- can properly detect scope violations during unification. + -- We do this for *all* identifiers, not just those that are + -- integers, because they may become integers later due to + -- inference... + forM_ bnds $ \ident -> + constrain (identName ident) $ ParamSize $ srclocOf ident + m + + bindVars :: TermScope -> [Ident] -> TermScope + bindVars = foldl bindVar + + bindVar :: TermScope -> Ident -> TermScope + bindVar scope (Ident name (Info tp) _) = + let inedges = boundAliases $ aliases tp + update (BoundV l tparams in_t) + -- If 'name' is record or sum-typed, don't alias the + -- components to 'name', because these no identity + -- beyond their components. + | Array {} <- tp = BoundV l tparams (in_t `addAliases` S.insert (AliasBound name)) + | otherwise = BoundV l tparams in_t + update b = b + + tp' = tp `addAliases` S.insert (AliasBound name) + in scope + { scopeVtable = + M.insert name (BoundV Local [] tp') $ + adjustSeveral update inedges $ + scopeVtable scope + } + + adjustSeveral f = flip $ foldl $ flip $ M.adjust f + + -- Check whether the bound variables have been used correctly + -- within their scope. + check m = do + (a, usages) <- collectBindingsOccurences m + checkOccurences usages + + mapM_ (checkIfUsed usages) bnds + + return a + + -- Collect and remove all occurences in @bnds@. This relies + -- on the fact that no variables shadow any other. + collectBindingsOccurences m = pass $ do + (x, usage) <- listen m + let (relevant, rest) = split usage + return ((x, relevant), const rest) + where + split = + unzip + . map + ( \occ -> + let (obs1, obs2) = divide $ observed occ + occ_cons = divide <$> consumed occ + con1 = fst <$> occ_cons + con2 = snd <$> occ_cons + in ( occ {observed = obs1, consumed = con1}, + occ {observed = obs2, consumed = con2} + ) + ) + names = S.fromList $ map identName bnds + divide s = (s `S.intersection` names, s `S.difference` names) + +bindingTypes :: + [Either (VName, TypeBinding) (VName, Constraint)] -> + TermTypeM a -> + TermTypeM a +bindingTypes types m = do + lvl <- curLevel + modifyConstraints (<> M.map (lvl,) (M.fromList constraints)) + localScope extend m + where + (tbinds, constraints) = partitionEithers types + extend scope = + scope + { scopeTypeTable = M.fromList tbinds <> scopeTypeTable scope + } + +bindingTypeParams :: [TypeParam] -> TermTypeM a -> TermTypeM a +bindingTypeParams tparams = + binding (mapMaybe typeParamIdent tparams) + . bindingTypes (concatMap typeParamType tparams) + where + typeParamType (TypeParamType l v loc) = + [ Left (v, TypeAbbr l [] (Scalar (TypeVar () Nonunique (typeName v) []))), + Right (v, ParamType l loc) + ] + typeParamType (TypeParamDim v loc) = + [Right (v, ParamSize loc)] + +typeParamIdent :: TypeParam -> Maybe Ident +typeParamIdent (TypeParamDim v loc) = + Just $ Ident v (Info $ Scalar $ Prim $ Signed Int64) loc +typeParamIdent _ = Nothing + +bindingIdent :: + IdentBase NoInfo Name -> + PatternType -> + (Ident -> TermTypeM a) -> + TermTypeM a +bindingIdent (Ident v NoInfo vloc) t m = + bindSpaced [(Term, v)] $ do + v' <- checkName Term v vloc + let ident = Ident v' (Info t) vloc + binding [ident] $ m ident + +bindingParams :: + [UncheckedTypeParam] -> + [UncheckedPattern] -> + ([TypeParam] -> [Pattern] -> TermTypeM a) -> + TermTypeM a +bindingParams tps orig_ps m = do + checkForDuplicateNames orig_ps + checkTypeParams tps $ \tps' -> bindingTypeParams tps' $ do + let descend ps' (p : ps) = + checkPattern p NoneInferred $ \p' -> + binding (S.toList $ patternIdents p') $ descend (p' : ps') ps + descend ps' [] = do + -- Perform an observation of every type parameter. This + -- prevents unused-name warnings for otherwise unused + -- dimensions. + mapM_ observe $ mapMaybe typeParamIdent tps' + m tps' $ reverse ps' + + descend [] orig_ps + +bindingPattern :: + PatternBase NoInfo Name -> + InferredType -> + (Pattern -> TermTypeM a) -> + TermTypeM a +bindingPattern p t m = do + checkForDuplicateNames [p] + checkPattern p t $ \p' -> binding (S.toList $ patternIdents p') $ do + -- Perform an observation of every declared dimension. This + -- prevents unused-name warnings for otherwise unused dimensions. + mapM_ observe $ patternDims p' + + m p' + +patternDims :: Pattern -> [Ident] +patternDims (PatternParens p _) = patternDims p +patternDims (TuplePattern pats _) = concatMap patternDims pats +patternDims (PatternAscription p (TypeDecl _ (Info t)) _) = + patternDims p <> mapMaybe (dimIdent (srclocOf p)) (nestedDims t) + where + dimIdent _ AnyDim = Nothing + dimIdent _ (ConstDim _) = Nothing + dimIdent _ NamedDim {} = Nothing +patternDims _ = [] + +sliceShape :: + Maybe (SrcLoc, Rigidity) -> + [DimIndex] -> + TypeBase (DimDecl VName) as -> + TermTypeM (TypeBase (DimDecl VName) as, [VName]) +sliceShape r slice t@(Array als u et (ShapeDecl orig_dims)) = + runWriterT $ setDims <$> adjustDims slice orig_dims + where + setDims [] = stripArray (length orig_dims) t + setDims dims' = Array als u et $ ShapeDecl dims' + + -- If the result is supposed to be AnyDim or a nonrigid size + -- variable, then don't bother trying to create + -- non-existential sizes. This is necessary to make programs + -- type-check without too much ceremony; see + -- e.g. tests/inplace5.fut. + isRigid Rigid {} = True + isRigid _ = False + refine_sizes = maybe False (isRigid . snd) r + + sliceSize orig_d i j stride = + case r of + Just (loc, Rigid _) -> do + (d, ext) <- + lift $ + extSize loc $ + SourceSlice orig_d' (bareExp <$> i) (bareExp <$> j) (bareExp <$> stride) + tell $ maybeToList ext + return d + Just (loc, Nonrigid) -> + lift $ NamedDim . qualName <$> newDimVar loc Nonrigid "slice_dim" + Nothing -> + pure AnyDim + where + -- The original size does not matter if the slice is fully specified. + orig_d' + | isJust i, isJust j = Nothing + | otherwise = Just orig_d + + adjustDims (DimFix {} : idxes') (_ : dims) = + adjustDims idxes' dims + -- Pattern match some known slices to be non-existential. + adjustDims (DimSlice i j stride : idxes') (_ : dims) + | refine_sizes, + maybe True ((== Just 0) . isInt64) i, + Just j' <- maybeDimFromExp =<< j, + maybe True ((== Just 1) . isInt64) stride = + (j' :) <$> adjustDims idxes' dims + adjustDims (DimSlice Nothing Nothing stride : idxes') (d : dims) + | refine_sizes, + maybe True (maybe False ((== 1) . abs) . isInt64) stride = + (d :) <$> adjustDims idxes' dims + adjustDims (DimSlice i j stride : idxes') (d : dims) = + (:) <$> sliceSize d i j stride <*> adjustDims idxes' dims + adjustDims _ dims = + pure dims +sliceShape _ _ t = pure (t, []) + +--- Main checkers + +-- | @require ts e@ causes a 'TypeError' if @expType e@ is not one of +-- the types in @ts@. Otherwise, simply returns @e@. +require :: String -> [PrimType] -> Exp -> TermTypeM Exp +require why ts e = do + mustBeOneOf ts (mkUsage (srclocOf e) why) . toStruct =<< expType e + return e + +unifies :: String -> StructType -> Exp -> TermTypeM Exp +unifies why t e = do + unify (mkUsage (srclocOf e) why) t . toStruct =<< expType e + return e + +-- The closure of a lambda or local function are those variables that +-- it references, and which local to the current top-level function. +lexicalClosure :: [Pattern] -> Occurences -> TermTypeM Aliasing +lexicalClosure params closure = do + vtable <- asks $ scopeVtable . termScope + let isLocal v = case v `M.lookup` vtable of + Just (BoundV Local _ _) -> True + _ -> False + return $ + S.map AliasBound $ + S.filter isLocal $ + allOccuring closure S.\\ mconcat (map patternNames params) + +noAliasesIfOverloaded :: PatternType -> TermTypeM PatternType +noAliasesIfOverloaded t@(Scalar (TypeVar _ u tn [])) = do + subst <- fmap snd . M.lookup (typeLeaf tn) <$> getConstraints + case subst of + Just Overloaded {} -> return $ Scalar $ TypeVar mempty u tn [] + _ -> return t +noAliasesIfOverloaded t = + return t + +-- Check the common parts of ascription and coercion. +checkAscript :: + SrcLoc -> + UncheckedTypeDecl -> + UncheckedExp -> + (StructType -> StructType) -> + TermTypeM (TypeDecl, Exp) +checkAscript loc decl e shapef = do + decl' <- checkTypeDecl decl + e' <- checkExp e + t <- expTypeFully e' + + (decl_t_nonrigid, _) <- + instantiateEmptyArrayDims loc "impl" Nonrigid $ + shapef $ + unInfo $ expandedType decl' + + onFailure (CheckingAscription (unInfo $ expandedType decl') (toStruct t)) $ + unify (mkUsage loc "type ascription") decl_t_nonrigid (toStruct t) + + -- We also have to make sure that uniqueness matches. This is done + -- explicitly, because uniqueness is ignored by unification. + t' <- normTypeFully t + decl_t' <- normTypeFully $ unInfo $ expandedType decl' + unless (t' `subtypeOf` anySizes decl_t') $ + typeError loc mempty $ + "Type" <+> pquote (ppr t') <+> "is not a subtype of" + <+> pquote (ppr decl_t') <> "." + + return (decl', e') + +unscopeType :: + SrcLoc -> + M.Map VName Ident -> + PatternType -> + TermTypeM (PatternType, [VName]) +unscopeType tloc unscoped t = do + (t', m) <- runStateT (traverseDims onDim t) mempty + return (t' `addAliases` S.map unAlias, M.elems m) + where + onDim _ p (NamedDim d) + | Just loc <- srclocOf <$> M.lookup (qualLeaf d) unscoped = + if p == PosImmediate || p == PosParam + then inst loc $ qualLeaf d + else return AnyDim + onDim _ _ d = return d + + inst loc d = do + prev <- gets $ M.lookup d + case prev of + Just d' -> return $ NamedDim $ qualName d' + Nothing -> do + d' <- lift $ newDimVar tloc (Rigid $ RigidOutOfScope loc d) "d" + modify $ M.insert d d' + return $ NamedDim $ qualName d' + + unAlias (AliasBound v) | v `M.member` unscoped = AliasFree v + unAlias a = a + +-- 'checkApplyExp' is like 'checkExp', but tries to find the "root +-- function", for better error messages. +checkApplyExp :: UncheckedExp -> TermTypeM (Exp, ApplyOp) +checkApplyExp (Apply e1 e2 _ _ loc) = do + (e1', (fname, i)) <- checkApplyExp e1 + arg <- checkArg e2 + t <- expType e1' + (t1, rt, argext, exts) <- checkApply loc (fname, i) t arg + return + ( Apply e1' (argExp arg) (Info (diet t1, argext)) (Info rt, Info exts) loc, + (fname, i + 1) + ) +checkApplyExp e = do + e' <- checkExp e + return + ( e', + ( case e' of + Var qn _ _ -> Just qn + _ -> Nothing, + 0 + ) + ) + +checkExp :: UncheckedExp -> TermTypeM Exp +checkExp (Literal val loc) = + return $ Literal val loc +checkExp (StringLit vs loc) = + return $ StringLit vs loc +checkExp (IntLit val NoInfo loc) = do + t <- newTypeVar loc "t" + mustBeOneOf anyNumberType (mkUsage loc "integer literal") t + return $ IntLit val (Info $ fromStruct t) loc +checkExp (FloatLit val NoInfo loc) = do + t <- newTypeVar loc "t" + mustBeOneOf anyFloatType (mkUsage loc "float literal") t + return $ FloatLit val (Info $ fromStruct t) loc +checkExp (TupLit es loc) = + TupLit <$> mapM checkExp es <*> pure loc +checkExp (RecordLit fs loc) = do + fs' <- evalStateT (mapM checkField fs) mempty + + return $ RecordLit fs' loc + where + checkField (RecordFieldExplicit f e rloc) = do + errIfAlreadySet f rloc + modify $ M.insert f rloc + RecordFieldExplicit f <$> lift (checkExp e) <*> pure rloc + checkField (RecordFieldImplicit name NoInfo rloc) = do + errIfAlreadySet name rloc + (QualName _ name', t) <- lift $ lookupVar rloc $ qualName name + modify $ M.insert name rloc + return $ RecordFieldImplicit name' (Info t) rloc + + errIfAlreadySet f rloc = do + maybe_sloc <- gets $ M.lookup f + case maybe_sloc of + Just sloc -> + lift $ + typeError rloc mempty $ + "Field" <+> pquote (ppr f) + <+> "previously defined at" + <+> text (locStrRel rloc sloc) <> "." + Nothing -> return () +checkExp (ArrayLit all_es _ loc) = + -- Construct the result type and unify all elements with it. We + -- only create a type variable for empty arrays; otherwise we use + -- the type of the first element. This significantly cuts down on + -- the number of type variables generated for pathologically large + -- multidimensional array literals. + case all_es of + [] -> do + et <- newTypeVar loc "t" + t <- arrayOfM loc et (ShapeDecl [ConstDim 0]) Unique + return $ ArrayLit [] (Info t) loc + e : es -> do + e' <- checkExp e + et <- expType e' + es' <- mapM (unifies "type of first array element" (toStruct et) <=< checkExp) es + et' <- normTypeFully et + t <- arrayOfM loc et' (ShapeDecl [ConstDim $ length all_es]) Unique + return $ ArrayLit (e' : es') (Info t) loc +checkExp (Range start maybe_step end _ loc) = do + start' <- require "use in range expression" anySignedType =<< checkExp start + start_t <- toStruct <$> expTypeFully start' + maybe_step' <- case maybe_step of + Nothing -> return Nothing + Just step -> do + let warning = warn loc "First and second element of range are identical, this will produce an empty array." + case (start, step) of + (Literal x _, Literal y _) -> when (x == y) warning + (Var x_name _ _, Var y_name _ _) -> when (x_name == y_name) warning + _ -> return () + Just <$> (unifies "use in range expression" start_t =<< checkExp step) + + let unifyRange e = unifies "use in range expression" start_t =<< checkExp e + end' <- traverse unifyRange end + + end_t <- case end' of + DownToExclusive e -> expType e + ToInclusive e -> expType e + UpToExclusive e -> expType e + + -- Special case some ranges to give them a known size. + let dimFromBound = dimFromExp (SourceBound . bareExp) + (dim, retext) <- + case (isInt64 start', isInt64 <$> maybe_step', end') of + (Just 0, Just (Just 1), UpToExclusive end'') + | Scalar (Prim (Signed Int64)) <- end_t -> + dimFromBound end'' + (Just 0, Nothing, UpToExclusive end'') + | Scalar (Prim (Signed Int64)) <- end_t -> + dimFromBound end'' + (Just 1, Just (Just 2), ToInclusive end'') + | Scalar (Prim (Signed Int64)) <- end_t -> + dimFromBound end'' + _ -> do + d <- newDimVar loc (Rigid RigidRange) "range_dim" + return (NamedDim $ qualName d, Just d) + + t <- arrayOfM loc start_t (ShapeDecl [dim]) Unique + let ret = (Info (t `setAliases` mempty), Info $ maybeToList retext) + + return $ Range start' maybe_step' end' ret loc +checkExp (Ascript e decl loc) = do + (decl', e') <- checkAscript loc decl e id + return $ Ascript e' decl' loc +checkExp (Coerce e decl _ loc) = do + -- We instantiate the declared types with all dimensions as nonrigid + -- fresh type variables, which we then use to unify with the type of + -- 'e'. This lets 'e' have whatever sizes it wants, but the overall + -- type must still match. Eventually we will throw away those sizes + -- (they will end up being unified with various sizes in 'e', which + -- is fine). + (decl', e') <- checkAscript loc decl e anySizes + + -- Now we instantiate the declared type again, but this time we keep + -- around the sizes as existentials. This is the result of the + -- ascription as a whole. We use matchDims to obtain the aliasing + -- of 'e'. + (decl_t_rigid, ext) <- + instantiateDimsInReturnType loc Nothing $ unInfo $ expandedType decl' + + t <- expTypeFully e' + + t' <- matchDims (const pure) t $ fromStruct decl_t_rigid + + return $ Coerce e' decl' (Info t', Info ext) loc +checkExp (BinOp (op, oploc) NoInfo (e1, _) (e2, _) NoInfo NoInfo loc) = do + (op', ftype) <- lookupVar oploc op + e1_arg <- checkArg e1 + e2_arg <- checkArg e2 + + -- Note that the application to the first operand cannot fix any + -- existential sizes, because it must by necessity be a function. + (p1_t, rt, p1_ext, _) <- checkApply loc (Just op', 0) ftype e1_arg + (p2_t, rt', p2_ext, retext) <- checkApply loc (Just op', 1) rt e2_arg + + return $ + BinOp + (op', oploc) + (Info ftype) + (argExp e1_arg, Info (toStruct p1_t, p1_ext)) + (argExp e2_arg, Info (toStruct p2_t, p2_ext)) + (Info rt') + (Info retext) + loc +checkExp (Project k e NoInfo loc) = do + e' <- checkExp e + t <- expType e' + kt <- mustHaveField (mkUsage loc $ "projection of field " ++ quote (pretty k)) k t + return $ Project k e' (Info kt) loc +checkExp (If e1 e2 e3 _ loc) = + sequentially checkCond $ \e1' _ -> do + ((e2', e3'), dflow) <- tapOccurences $ checkExp e2 `alternative` checkExp e3 + + (brancht, retext) <- unifyBranches loc e2' e3' + let t' = addAliases brancht (`S.difference` S.map AliasBound (allConsumed dflow)) + + zeroOrderType + (mkUsage loc "returning value of this type from 'if' expression") + "type returned from branch" + t' + + return $ If e1' e2' e3' (Info t', Info retext) loc + where + checkCond = do + e1' <- checkExp e1 + let bool = Scalar $ Prim Bool + e1_t <- toStruct <$> expType e1' + onFailure (CheckingRequired [bool] e1_t) $ + unify (mkUsage (srclocOf e1') "use as 'if' condition") bool e1_t + return e1' +checkExp (Parens e loc) = + Parens <$> checkExp e <*> pure loc +checkExp (QualParens (modname, modnameloc) e loc) = do + (modname', mod) <- lookupMod loc modname + case mod of + ModEnv env -> local (`withEnv` qualifyEnv modname' env) $ do + e' <- checkExp e + return $ QualParens (modname', modnameloc) e' loc + ModFun {} -> + typeError loc mempty $ "Module" <+> ppr modname <+> " is a parametric module." + where + qualifyEnv modname' env = + env {envNameMap = M.map (qualify' modname') $ envNameMap env} + qualify' modname' (QualName qs name) = + QualName (qualQuals modname' ++ [qualLeaf modname'] ++ qs) name +checkExp (Var qn NoInfo loc) = do + -- The qualifiers of a variable is divided into two parts: first a + -- possibly-empty sequence of module qualifiers, followed by a + -- possible-empty sequence of record field accesses. We use scope + -- information to perform the split, by taking qualifiers off the + -- end until we find a module. + + (qn', t, fields) <- findRootVar (qualQuals qn) (qualLeaf qn) + + foldM checkField (Var qn' (Info t) loc) fields + where + findRootVar qs name = + (whenFound <$> lookupVar loc (QualName qs name)) `catchError` notFound qs name + + whenFound (qn', t) = (qn', t, []) + + notFound qs name err + | null qs = throwError err + | otherwise = do + (qn', t, fields) <- + findRootVar (init qs) (last qs) + `catchError` const (throwError err) + return (qn', t, fields ++ [name]) + + checkField e k = do + t <- expType e + let usage = mkUsage loc $ "projection of field " ++ quote (pretty k) + kt <- mustHaveField usage k t + return $ Project k e (Info kt) loc +checkExp (Negate arg loc) = do + arg' <- require "numeric negation" anyNumberType =<< checkExp arg + return $ Negate arg' loc +checkExp e@Apply {} = fst <$> checkApplyExp e +checkExp (LetPat pat e body _ loc) = + sequentially (checkExp e) $ \e' e_occs -> do + -- Not technically an ascription, but we want the pattern to have + -- exactly the type of 'e'. + t <- expType e' + case anyConsumption e_occs of + Just c -> + let msg = "type computed with consumption at " ++ locStr (location c) + in zeroOrderType (mkUsage loc "consumption in right-hand side of 'let'-binding") msg t + _ -> return () + + incLevel $ + bindingPattern pat (Ascribed t) $ \pat' -> do + body' <- checkExp body + (body_t, retext) <- + unscopeType loc (patternMap pat') =<< expTypeFully body' + + return $ LetPat pat' e' body' (Info body_t, Info retext) loc +checkExp (LetFun name (tparams, params, maybe_retdecl, NoInfo, e) body NoInfo loc) = + sequentially (checkBinding (name, maybe_retdecl, tparams, params, e, loc)) $ + \(tparams', params', maybe_retdecl', rettype, _, e') closure -> do + closure' <- lexicalClosure params' closure + + bindSpaced [(Term, name)] $ do + name' <- checkName Term name loc + + let arrow (xp, xt) yt = Scalar $ Arrow () xp xt yt + ftype = foldr (arrow . patternParam) rettype params' + entry = BoundV Local tparams' $ ftype `setAliases` closure' + bindF scope = + scope + { scopeVtable = + M.insert name' entry $ scopeVtable scope, + scopeNameMap = + M.insert (Term, name) (qualName name') $ + scopeNameMap scope + } + body' <- localScope bindF $ checkExp body + + -- We fake an ident here, but it's OK as it can't be a size + -- anyway. + let fake_ident = Ident name' (Info $ fromStruct ftype) mempty + (body_t, _) <- + unscopeType loc (M.singleton name' fake_ident) + =<< expTypeFully body' + + return $ + LetFun + name' + (tparams', params', maybe_retdecl', Info rettype, e') + body' + (Info body_t) + loc +checkExp (LetWith dest src idxes ve body NoInfo loc) = + sequentially (checkIdent src) $ \src' _ -> do + (t, _) <- newArrayType (srclocOf src) "src" $ length idxes + unify (mkUsage loc "type of target array") t $ toStruct $ unInfo $ identType src' + + -- Need the fully normalised type here to get the proper aliasing information. + src_t <- normTypeFully $ unInfo $ identType src' + + idxes' <- mapM checkDimIndex idxes + (elemt, _) <- sliceShape (Just (loc, Nonrigid)) idxes' =<< normTypeFully t + + unless (unique src_t) $ + typeError loc mempty $ + "Source" <+> pquote (pprName (identName src)) + <+> "has type" + <+> ppr src_t <> ", which is not unique." + vtable <- asks $ scopeVtable . termScope + forM_ (aliases src_t) $ \v -> + case aliasVar v `M.lookup` vtable of + Just (BoundV Local _ v_t) + | not $ unique v_t -> + typeError loc mempty $ + "Source" <+> pquote (pprName (identName src)) + <+> "aliases" + <+> pquote (pprName (aliasVar v)) <> ", which is not consumable." + _ -> return () + + sequentially (unifies "type of target array" (toStruct elemt) =<< checkExp ve) $ \ve' _ -> do + ve_t <- expTypeFully ve' + when (AliasBound (identName src') `S.member` aliases ve_t) $ + badLetWithValue loc + + bindingIdent dest (src_t `setAliases` S.empty) $ \dest' -> do + body' <- consuming src' $ checkExp body + (body_t, _) <- + unscopeType loc (M.singleton (identName dest') dest') + =<< expTypeFully body' + return $ LetWith dest' src' idxes' ve' body' (Info body_t) loc +checkExp (Update src idxes ve loc) = do + (t, _) <- newArrayType (srclocOf src) "src" $ length idxes + idxes' <- mapM checkDimIndex idxes + (elemt, _) <- sliceShape (Just (loc, Nonrigid)) idxes' =<< normTypeFully t + + sequentially (checkExp ve >>= unifies "type of target array" elemt) $ \ve' _ -> + sequentially (checkExp src >>= unifies "type of target array" t) $ \src' _ -> do + src_t <- expTypeFully src' + unless (unique src_t) $ + typeError loc mempty $ + "Source" <+> pquote (ppr src) + <+> "has type" + <+> ppr src_t <> ", which is not unique." + + let src_als = aliases src_t + ve_t <- expTypeFully ve' + unless (S.null $ src_als `S.intersection` aliases ve_t) $ badLetWithValue loc + + consume loc src_als + return $ Update src' idxes' ve' loc + +-- Record updates are a bit hacky, because we do not have row typing +-- (yet?). For now, we only permit record updates where we know the +-- full type up to the field we are updating. +checkExp (RecordUpdate src fields ve NoInfo loc) = do + src' <- checkExp src + ve' <- checkExp ve + a <- expTypeFully src' + let usage = mkUsage loc "record update" + r <- foldM (flip $ mustHaveField usage) a fields + ve_t <- expType ve' + let r' = anySizes $ toStruct r + ve_t' = anySizes $ toStruct ve_t + onFailure (CheckingRecordUpdate fields r' ve_t') $ + unify usage r' ve_t' + maybe_a' <- onRecordField (const ve_t) fields <$> expTypeFully src' + case maybe_a' of + Just a' -> return $ RecordUpdate src' fields ve' (Info a') loc + Nothing -> + typeError loc mempty $ + "Full type of" + indent 2 (ppr src) + textwrap " is not known at this point. Add a size annotation to the original record to disambiguate." +checkExp (Index e idxes _ loc) = do + (t, _) <- newArrayType loc "e" $ length idxes + e' <- unifies "being indexed at" t =<< checkExp e + idxes' <- mapM checkDimIndex idxes + -- XXX, the RigidSlice here will be overridden in sliceShape with a proper value. + (t', retext) <- + sliceShape (Just (loc, Rigid (RigidSlice Nothing ""))) idxes' + =<< expTypeFully e' + + -- Remove aliases if the result is an overloaded type, because that + -- will certainly not be aliased. + t'' <- noAliasesIfOverloaded t' + + return $ Index e' idxes' (Info t'', Info retext) loc +checkExp (Assert e1 e2 NoInfo loc) = do + e1' <- require "being asserted" [Bool] =<< checkExp e1 + e2' <- checkExp e2 + return $ Assert e1' e2' (Info (pretty e1)) loc +checkExp (Lambda params body rettype_te NoInfo loc) = + removeSeminullOccurences $ + noUnique $ + incLevel $ + bindingParams [] params $ \_ params' -> do + rettype_checked <- traverse checkTypeExp rettype_te + let declared_rettype = + case rettype_checked of + Just (_, st, _) -> Just st + Nothing -> Nothing + (body', closure) <- + tapOccurences $ checkFunBody params' body declared_rettype loc + body_t <- expTypeFully body' + + params'' <- mapM updateTypes params' + + (rettype', rettype_st) <- + case rettype_checked of + Just (te, st, _) -> + return (Just te, st) + Nothing -> do + ret <- + inferReturnSizes params'' $ + toStruct $ + inferReturnUniqueness params'' body_t + return (Nothing, ret) + + checkGlobalAliases params' body_t loc + verifyFunctionParams Nothing params' + + closure' <- lexicalClosure params'' closure + + return $ Lambda params'' body' rettype' (Info (closure', rettype_st)) loc + where + -- Inferring the sizes of the return type of a lambda is a lot + -- like let-generalisation. We wish to remove any rigid sizes + -- that were created when checking the body, except for those that + -- are visible in types that existed before we entered the body, + -- are parameters, or are used in parameters. + inferReturnSizes params' ret = do + cur_lvl <- curLevel + let named (Named x, _) = Just x + named (Unnamed, _) = Nothing + param_names = mapMaybe (named . patternParam) params' + pos_sizes = + typeDimNamesPos (foldFunType (map patternStructType params') ret) + hide k (lvl, _) = + lvl >= cur_lvl && k `notElem` param_names && k `S.notMember` pos_sizes + + hidden_sizes <- + S.fromList . M.keys . M.filterWithKey hide <$> getConstraints + + let onDim (NamedDim name) + | not (qualLeaf name `S.member` hidden_sizes) = NamedDim name + | otherwise = AnyDim + onDim d = d + + return $ first onDim ret +checkExp (OpSection op _ loc) = do + (op', ftype) <- lookupVar loc op + return $ OpSection op' (Info ftype) loc +checkExp (OpSectionLeft op _ e _ _ loc) = do + (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) -> + return $ + OpSectionLeft + op' + (Info ftype) + (argExp e_arg) + (Info (toStruct t1, argext), Info $ toStruct t2) + (Info rettype, Info retext) + loc + _ -> + typeError loc mempty $ + "Operator section with invalid operator of type" <+> ppr ftype +checkExp (OpSectionRight op _ e _ NoInfo loc) = do + (op', ftype) <- lookupVar loc op + e_arg <- checkArg e + case ftype of + Scalar (Arrow as1 m1 t1 (Scalar (Arrow as2 m2 t2 ret))) -> do + (t2', ret', argext, _) <- + checkApply + loc + (Just op', 1) + (Scalar $ Arrow as2 m2 t2 $ Scalar $ Arrow as1 m1 t1 ret) + e_arg + return $ + OpSectionRight + op' + (Info ftype) + (argExp e_arg) + (Info $ toStruct t1, Info (toStruct t2', argext)) + (Info $ addAliases ret (<> aliases ret')) + loc + _ -> + typeError loc mempty $ + "Operator section with invalid operator of type" <+> ppr ftype +checkExp (ProjectSection fields NoInfo loc) = do + a <- newTypeVar loc "a" + let usage = mkUsage loc "projection at" + b <- foldM (flip $ mustHaveField usage) a fields + return $ ProjectSection fields (Info $ Scalar $ Arrow mempty Unnamed a b) loc +checkExp (IndexSection idxes NoInfo loc) = do + (t, _) <- newArrayType loc "e" $ length idxes + idxes' <- mapM checkDimIndex idxes + (t', _) <- sliceShape Nothing idxes' t + return $ IndexSection idxes' (Info $ fromStruct $ Scalar $ Arrow mempty Unnamed t t') loc +checkExp (DoLoop _ mergepat mergeexp form loopbody NoInfo loc) = + sequentially (checkExp mergeexp) $ \mergeexp' _ -> do + zeroOrderType + (mkUsage (srclocOf mergeexp) "use as loop variable") + "type used as loop variable" + =<< expTypeFully mergeexp' + + -- The handling of dimension sizes is a bit intricate, but very + -- similar to checking a function, followed by checking a call to + -- it. The overall procedure is as follows: + -- + -- (1) All empty dimensions in the merge pattern are instantiated + -- with nonrigid size variables. All explicitly specified + -- dimensions are preserved. + -- + -- (2) The body of the loop is type-checked. The result type is + -- combined with the merge pattern type to determine which sizes are + -- variant, and these are turned into size parameters for the merge + -- pattern. + -- + -- (3) We now conceptually have a function parameter type and return + -- type. We check that it can be called with the initial merge + -- values as argument. The result of this is the type of the loop + -- as a whole. + -- + -- (There is also a convergence loop for inferring uniqueness, but + -- that's orthogonal to the size handling.) + + (merge_t, new_dims) <- + instantiateEmptyArrayDims loc "loop" Nonrigid + . anySizes -- dim handling (1) + =<< expTypeFully mergeexp' + + -- dim handling (2) + let checkLoopReturnSize mergepat' loopbody' = do + loopbody_t <- expTypeFully loopbody' + pat_t <- normTypeFully $ patternType mergepat' + -- We are ignoring the dimensions here, because any mismatches + -- should be turned into fresh size variables. + onFailure (CheckingLoopBody (toStruct (anySizes pat_t)) (toStruct loopbody_t)) $ + expect + (mkUsage (srclocOf loopbody) "matching loop body to loop pattern") + (toStruct (anySizes pat_t)) + (toStruct loopbody_t) + pat_t' <- normTypeFully pat_t + loopbody_t' <- normTypeFully loopbody_t + + -- For each new_dims, figure out what they are instantiated + -- with in the initial value. This is used to determine + -- whether a size is invariant because it always matches the + -- initial instantiation of that size. + let initSubst (NamedDim v, d) = Just (v, d) + initSubst _ = Nothing + init_substs <- + M.fromList . mapMaybe initSubst . snd + . anyDimOnMismatch pat_t' + <$> expTypeFully mergeexp' + + -- Figure out which of the 'new_dims' dimensions are variant. + -- This works because we know that each dimension from + -- new_dims in the pattern is unique and distinct. + -- + -- Our logic here is a bit reversed: the *mismatches* (from + -- new_dims) are what we want to extract and turn into size + -- parameters. + let mismatchSubst (NamedDim v, d) + | qualLeaf v `elem` new_dims = + case M.lookup v init_substs of + Just d' + | d' == d -> + return $ Just (qualLeaf v, SizeSubst d) + _ -> do + tell [qualLeaf v] + return Nothing + mismatchSubst _ = return Nothing + + (init_substs', sparams) = + runWriter $ + M.fromList . catMaybes + <$> mapM + mismatchSubst + (snd $ anyDimOnMismatch pat_t' loopbody_t') + + -- Make sure that any of new_dims that are invariant will be + -- replaced with the invariant size in the loop body. Failure + -- to do this can cause type annotations to still refer to + -- new_dims. + let dimToInit (v, SizeSubst d) = + constrain v $ Size (Just d) (mkUsage loc "size of loop parameter") + dimToInit _ = + return () + mapM_ dimToInit $ M.toList init_substs' + + mergepat'' <- applySubst (`M.lookup` init_substs') <$> updateTypes mergepat' + return (nub sparams, mergepat'') + + -- First we do a basic check of the loop body to figure out which of + -- the merge parameters are being consumed. For this, we first need + -- to check the merge pattern, which requires the (initial) merge + -- expression. + -- + -- Play a little with occurences to ensure it does not look like + -- none of the merge variables are being used. + ((sparams, mergepat', form', loopbody'), bodyflow) <- + case form of + For i uboundexp -> do + uboundexp' <- require "being the bound in a 'for' loop" anySignedType =<< checkExp uboundexp + bound_t <- expTypeFully uboundexp' + bindingIdent i bound_t $ \i' -> + noUnique $ + bindingPattern mergepat (Ascribed merge_t) $ + \mergepat' -> onlySelfAliasing $ + tapOccurences $ do + loopbody' <- noSizeEscape $ checkExp loopbody + (sparams, mergepat'') <- checkLoopReturnSize mergepat' loopbody' + return + ( sparams, + mergepat'', + For i' uboundexp', + loopbody' + ) + ForIn xpat e -> do + (arr_t, _) <- newArrayType (srclocOf e) "e" 1 + e' <- unifies "being iterated in a 'for-in' loop" arr_t =<< checkExp e + t <- expTypeFully e' + case t of + _ + | Just t' <- peelArray 1 t -> + bindingPattern xpat (Ascribed t') $ \xpat' -> + noUnique $ + bindingPattern mergepat (Ascribed merge_t) $ + \mergepat' -> onlySelfAliasing $ + tapOccurences $ do + loopbody' <- noSizeEscape $ checkExp loopbody + (sparams, mergepat'') <- checkLoopReturnSize mergepat' loopbody' + return + ( sparams, + mergepat'', + ForIn xpat' e', + loopbody' + ) + | otherwise -> + typeError (srclocOf e) mempty $ + "Iteratee of a for-in loop must be an array, but expression has type" + <+> ppr t + While cond -> + noUnique $ + bindingPattern mergepat (Ascribed merge_t) $ \mergepat' -> + onlySelfAliasing $ + tapOccurences $ + sequentially + ( checkExp cond + >>= unifies "being the condition of a 'while' loop" (Scalar $ Prim Bool) + ) + $ \cond' _ -> do + loopbody' <- noSizeEscape $ checkExp loopbody + (sparams, mergepat'') <- checkLoopReturnSize mergepat' loopbody' + return + ( sparams, + mergepat'', + While cond', + loopbody' + ) + + mergepat'' <- do + loopbody_t <- expTypeFully loopbody' + convergePattern mergepat' (allConsumed bodyflow) loopbody_t $ + mkUsage (srclocOf loopbody') "being (part of) the result of the loop body" + + let consumeMerge (Id _ (Info pt) ploc) mt + | unique pt = consume ploc $ aliases mt + consumeMerge (TuplePattern pats _) t + | Just ts <- isTupleRecord t = + zipWithM_ consumeMerge pats ts + consumeMerge (PatternParens pat _) t = + consumeMerge pat t + consumeMerge (PatternAscription pat _ _) t = + consumeMerge pat t + consumeMerge _ _ = + return () + consumeMerge mergepat'' =<< expTypeFully mergeexp' + + -- dim handling (3) + let sparams_anydim = M.fromList $ zip sparams $ repeat $ SizeSubst AnyDim + loopt_anydims = + applySubst (`M.lookup` sparams_anydim) $ + patternType mergepat'' + (merge_t', _) <- + instantiateEmptyArrayDims loc "loopres" Nonrigid $ toStruct loopt_anydims + mergeexp_t <- toStruct <$> expTypeFully mergeexp' + onFailure (CheckingLoopInitial (toStruct loopt_anydims) mergeexp_t) $ + unify + (mkUsage (srclocOf mergeexp') "matching initial loop values to pattern") + merge_t' + mergeexp_t + + (loopt, retext) <- instantiateDimsInType loc RigidLoop loopt_anydims + -- We set all of the uniqueness to be unique. This is intentional, + -- and matches what happens for function calls. Those arrays that + -- really *cannot* be consumed will alias something unconsumable, + -- and will be caught that way. + let bound_here = patternNames mergepat'' <> S.fromList sparams <> form_bound + form_bound = + case form' of + For v _ -> S.singleton $ identName v + ForIn forpat _ -> patternNames forpat + While {} -> mempty + loopt' = + second (`S.difference` S.map AliasBound bound_here) $ + loopt `setUniqueness` Unique + + -- Eliminate those new_dims that turned into sparams so it won't + -- look like we have ambiguous sizes lying around. + modifyConstraints $ M.filterWithKey $ \k _ -> k `notElem` sparams + + return $ DoLoop sparams mergepat'' mergeexp' form' loopbody' (Info (loopt', retext)) loc + where + convergePattern pat body_cons body_t body_loc = do + let consumed_merge = patternNames pat `S.intersection` body_cons + + uniquePat (Wildcard (Info t) wloc) = + Wildcard (Info $ t `setUniqueness` Nonunique) wloc + uniquePat (PatternParens p ploc) = + PatternParens (uniquePat p) ploc + uniquePat (Id name (Info t) iloc) + | name `S.member` consumed_merge = + let t' = t `setUniqueness` Unique `setAliases` mempty + in Id name (Info t') iloc + | otherwise = + let t' = t `setUniqueness` Nonunique + in Id name (Info t') iloc + uniquePat (TuplePattern pats ploc) = + TuplePattern (map uniquePat pats) ploc + uniquePat (RecordPattern fs ploc) = + RecordPattern (map (fmap uniquePat) fs) ploc + uniquePat (PatternAscription p t ploc) = + PatternAscription p t ploc + uniquePat p@PatternLit {} = p + uniquePat (PatternConstr n t ps ploc) = + PatternConstr n t (map uniquePat ps) ploc + + -- Make the pattern unique where needed. + pat' = uniquePat pat + + pat_t <- normTypeFully $ patternType pat' + unless (toStructural body_t `subtypeOf` toStructural pat_t) $ + unexpectedType (srclocOf body_loc) (toStruct body_t) [toStruct pat_t] + + -- Check that the new values of consumed merge parameters do not + -- alias something bound outside the loop, AND that anything + -- returned for a unique merge parameter does not alias anything + -- else returned. We also update the aliases for the pattern. + bound_outside <- asks $ S.fromList . M.keys . scopeVtable . termScope + let combAliases t1 t2 = + case t1 of + Scalar Record {} -> t1 + _ -> t1 `addAliases` (<> aliases t2) + + checkMergeReturn (Id pat_v (Info pat_v_t) patloc) t + | unique pat_v_t, + v : _ <- + S.toList $ + S.map aliasVar (aliases t) `S.intersection` bound_outside = + lift $ + typeError loc mempty $ + "Return value for loop parameter" + <+> pquote (pprName pat_v) + <+> "aliases" + <+> pprName v <> "." + | otherwise = do + (cons, obs) <- get + unless (S.null $ aliases t `S.intersection` cons) $ + lift $ + typeError loc mempty $ + "Return value for loop parameter" + <+> pquote (pprName pat_v) + <+> "aliases other consumed loop parameter." + when + ( unique pat_v_t + && not (S.null (aliases t `S.intersection` (cons <> obs))) + ) + $ lift $ + typeError loc mempty $ + "Return value for consuming loop parameter" + <+> pquote (pprName pat_v) + <+> "aliases previously returned value." + if unique pat_v_t + then put (cons <> aliases t, obs) + else put (cons, obs <> aliases t) + + return $ Id pat_v (Info (combAliases pat_v_t t)) patloc + checkMergeReturn (Wildcard (Info pat_v_t) patloc) t = + return $ Wildcard (Info (combAliases pat_v_t t)) patloc + checkMergeReturn (PatternParens p _) t = + checkMergeReturn p t + checkMergeReturn (PatternAscription p _ _) t = + checkMergeReturn p t + checkMergeReturn (RecordPattern pfs patloc) (Scalar (Record tfs)) = + RecordPattern . M.toList <$> sequence pfs' <*> pure patloc + where + pfs' = + M.intersectionWith + checkMergeReturn + (M.fromList pfs) + tfs + checkMergeReturn (TuplePattern pats patloc) t + | Just ts <- isTupleRecord t = + TuplePattern + <$> zipWithM checkMergeReturn pats ts + <*> pure patloc + checkMergeReturn p _ = + return p + + (pat'', (pat_cons, _)) <- + runStateT (checkMergeReturn pat' body_t) (mempty, mempty) + + let body_cons' = body_cons <> S.map aliasVar pat_cons + if body_cons' == body_cons && patternType pat'' == patternType pat + then return pat' + else convergePattern pat'' body_cons' body_t body_loc +checkExp (Constr name es NoInfo loc) = do + t <- newTypeVar loc "t" + es' <- mapM checkExp es + ets <- mapM expTypeFully es' + mustHaveConstr (mkUsage loc "use of constructor") name t (toStruct <$> ets) + -- A sum value aliases *anything* that went into its construction. + let als = foldMap aliases ets + return $ Constr name es' (Info $ fromStruct t `addAliases` (<> als)) loc +checkExp (Match e cs _ loc) = + sequentially (checkExp e) $ \e' _ -> do + mt <- expTypeFully e' + (cs', t, retext) <- checkCases mt cs + zeroOrderType + (mkUsage loc "being returned 'match'") + "type returned from pattern match" + t + return $ Match e' cs' (Info t, Info retext) loc +checkExp (Attr info e loc) = + Attr info <$> checkExp e <*> pure loc + +checkCases :: + PatternType -> + NE.NonEmpty (CaseBase NoInfo Name) -> + TermTypeM (NE.NonEmpty (CaseBase Info VName), PatternType, [VName]) +checkCases mt rest_cs = + case NE.uncons rest_cs of + (c, Nothing) -> do + (c', t, retext) <- checkCase mt c + return (c' NE.:| [], t, retext) + (c, Just cs) -> do + (((c', c_t, _), (cs', cs_t, _)), dflow) <- + tapOccurences $ checkCase mt c `alternative` checkCases mt cs + (brancht, retext) <- unifyBranchTypes (srclocOf c) c_t cs_t + let t = + addAliases + brancht + (`S.difference` S.map AliasBound (allConsumed dflow)) + return (NE.cons c' cs', t, retext) + +checkCase :: + PatternType -> + CaseBase NoInfo Name -> + TermTypeM (CaseBase Info VName, PatternType, [VName]) +checkCase mt (CasePat p e loc) = + bindingPattern p (Ascribed mt) $ \p' -> do + e' <- checkExp e + (t, retext) <- unscopeType loc (patternMap p') =<< expTypeFully e' + return (CasePat p' e' loc, t, retext) + +-- | An unmatched pattern. Used in in the generation of +-- unmatched pattern warnings by the type checker. +data Unmatched p + = UnmatchedNum p [ExpBase Info VName] + | UnmatchedBool p + | UnmatchedConstr p + | Unmatched p + deriving (Functor, Show) + +instance Pretty (Unmatched (PatternBase Info VName)) where + ppr um = case um of + (UnmatchedNum p nums) -> ppr' p <+> "where p is not one of" <+> ppr nums + (UnmatchedBool p) -> ppr' p + (UnmatchedConstr p) -> ppr' p + (Unmatched p) -> ppr' p + where + ppr' (PatternAscription p t _) = ppr p <> ":" <+> ppr t + ppr' (PatternParens p _) = parens $ ppr' p + ppr' (Id v _ _) = pprName v + ppr' (TuplePattern pats _) = parens $ commasep $ map ppr' pats + ppr' (RecordPattern fs _) = braces $ commasep $ map ppField fs + where + ppField (name, t) = text (nameToString name) <> equals <> ppr' t + ppr' Wildcard {} = "_" + ppr' (PatternLit e _ _) = ppr e + ppr' (PatternConstr n _ ps _) = "#" <> ppr n <+> sep (map ppr' ps) + +unpackPat :: Pattern -> [Maybe Pattern] +unpackPat Wildcard {} = [Nothing] +unpackPat (PatternParens p _) = unpackPat p +unpackPat Id {} = [Nothing] +unpackPat (TuplePattern ps _) = Just <$> ps +unpackPat (RecordPattern fs _) = Just . snd <$> sortFields (M.fromList fs) +unpackPat (PatternAscription p _ _) = unpackPat p +unpackPat p@PatternLit {} = [Just p] +unpackPat p@PatternConstr {} = [Just p] + +wildPattern :: Pattern -> Int -> Unmatched Pattern -> Unmatched Pattern +wildPattern (TuplePattern ps loc) pos um = wildTuple <$> um + where + wildTuple p = TuplePattern (take (pos - 1) ps' ++ [p] ++ drop pos ps') loc + ps' = map wildOut ps + wildOut p = Wildcard (Info (patternType p)) (srclocOf p) +wildPattern (RecordPattern fs loc) pos um = wildRecord <$> um + where + wildRecord p = + RecordPattern (take (pos - 1) fs' ++ [(fst (fs !! (pos - 1)), p)] ++ drop pos fs') loc + fs' = map wildOut fs + wildOut (f, p) = (f, Wildcard (Info (patternType p)) (srclocOf p)) +wildPattern (PatternAscription p _ _) pos um = wildPattern p pos um +wildPattern (PatternParens p _) pos um = wildPattern p pos um +wildPattern (PatternConstr n t ps loc) pos um = wildConstr <$> um + where + wildConstr p = PatternConstr n t (take (pos - 1) ps' ++ [p] ++ drop pos ps') loc + ps' = map wildOut ps + wildOut p = Wildcard (Info (patternType p)) (srclocOf p) +wildPattern _ _ um = um + +checkUnmatched :: Exp -> TermTypeM () +checkUnmatched e = void $ checkUnmatched' e >> astMap tv e + where + checkUnmatched' (Match _ cs _ loc) = + let ps = fmap (\(CasePat p _ _) -> p) cs + in case unmatched id $ NE.toList ps of + [] -> return () + ps' -> + typeError loc mempty $ + "Unmatched cases in match expression:" + indent 2 (stack (map ppr ps')) + checkUnmatched' _ = return () + tv = + ASTMapper + { mapOnExp = + \e' -> checkUnmatched' e' >> return e', + mapOnName = pure, + mapOnQualName = pure, + mapOnStructType = pure, + mapOnPatternType = pure + } + +-- | A data type for constructor patterns. This is used to make the +-- code for detecting unmatched constructors cleaner, by separating +-- the constructor-pattern cases from other cases. +data ConstrPat = ConstrPat + { constrName :: Name, + constrType :: PatternType, + constrPayload :: [Pattern], + constrSrcLoc :: SrcLoc + } + +-- Be aware of these fishy equality instances! + +instance Eq ConstrPat where + ConstrPat c1 _ _ _ == ConstrPat c2 _ _ _ = c1 == c2 + +instance Ord ConstrPat where + ConstrPat c1 _ _ _ `compare` ConstrPat c2 _ _ _ = c1 `compare` c2 + +unmatched :: (Unmatched Pattern -> Unmatched Pattern) -> [Pattern] -> [Unmatched Pattern] +unmatched hole orig_ps + | p : _ <- orig_ps, + sameStructure labeledCols = do + (i, cols) <- labeledCols + let hole' = if isConstr p then hole else hole . wildPattern p i + case sequence cols of + Nothing -> [] + Just cs + | all isPatternLit cs -> map hole' $ localUnmatched cs + | otherwise -> unmatched hole' cs + | otherwise = [] + where + labeledCols = zip [1 ..] $ transpose $ map unpackPat orig_ps + + localUnmatched :: [Pattern] -> [Unmatched Pattern] + localUnmatched [] = [] + localUnmatched ps'@(p' : _) = + case patternType p' of + Scalar (Sum cs'') -> + -- We now know that we are matching a sum type, and thus + -- that all patterns ps' are constructors (checked by + -- 'all isPatternLit' before this function is called). + let constrs = M.keys cs'' + matched = mapMaybe constr ps' + unmatched' = + map (UnmatchedConstr . buildConstr cs'') $ + constrs \\ map constrName matched + in case unmatched' of + [] -> + let constrGroups = group (sort matched) + removedConstrs = mapMaybe stripConstrs constrGroups + transposed = (fmap . fmap) transpose removedConstrs + findUnmatched (pc, trans) = do + col <- trans + case col of + [] -> [] + ((i, _) : _) -> unmatched (wilder i pc) (map snd col) + wilder i pc s = (`PatternParens` mempty) <$> wildPattern pc i s + in concatMap findUnmatched transposed + _ -> unmatched' + Scalar (Prim t) | not (any idOrWild ps') -> + -- We now know that we are matching a sum type, and thus + -- that all patterns ps' are literals (checked by 'all + -- isPatternLit' before this function is called). + case t of + Bool -> + let matched = nub $ mapMaybe (pExp >=> bool) $ filter isPatternLit ps' + in map (UnmatchedBool . buildBool (Scalar (Prim t))) $ [True, False] \\ matched + _ -> + let matched = mapMaybe pExp $ filter isPatternLit ps' + in [UnmatchedNum (buildId (Info $ Scalar $ Prim t) "p") matched] + _ -> [] + + isConstr PatternConstr {} = True + isConstr (PatternParens p _) = isConstr p + isConstr _ = False + + stripConstrs :: [ConstrPat] -> Maybe (Pattern, [[(Int, Pattern)]]) + stripConstrs (pc@ConstrPat {} : cs') = Just (unConstr pc, stripConstr pc : map stripConstr cs') + stripConstrs [] = Nothing + + stripConstr :: ConstrPat -> [(Int, Pattern)] + stripConstr (ConstrPat _ _ ps' _) = zip [1 ..] ps' + + sameStructure [] = True + sameStructure (x : xs) = all (\y -> length y == length x') xs' + where + (x' : xs') = map snd (x : xs) + + pExp (PatternLit e' _ _) = Just e' + pExp _ = Nothing + + constr (PatternConstr c (Info t) ps loc) = Just $ ConstrPat c t ps loc + constr (PatternParens p _) = constr p + constr (PatternAscription p' _ _) = constr p' + constr _ = Nothing + + unConstr p = + PatternConstr (constrName p) (Info $ constrType p) (constrPayload p) (constrSrcLoc p) + + isPatternLit PatternLit {} = True + isPatternLit (PatternAscription p' _ _) = isPatternLit p' + isPatternLit (PatternParens p' _) = isPatternLit p' + isPatternLit PatternConstr {} = True + isPatternLit _ = False + + idOrWild Id {} = True + idOrWild Wildcard {} = True + idOrWild (PatternAscription p' _ _) = idOrWild p' + idOrWild (PatternParens p' _) = idOrWild p' + idOrWild _ = False + + bool (Literal (BoolValue b) _) = Just b + bool _ = Nothing + + buildConstr m c = + let t = Scalar $ Sum m + cs = m M.! c + wildCS = map (\ct -> Wildcard (Info ct) mempty) cs + in if null wildCS + then PatternConstr c (Info t) [] mempty + else PatternParens (PatternConstr c (Info t) wildCS mempty) mempty + buildBool t b = + PatternLit (Literal (BoolValue b) mempty) (Info (addSizes t)) mempty + buildId t n = + -- The VName tag here will never be used since the value + -- exists exclusively for printing warnings. + Id (VName (nameFromString n) (-1)) t mempty + +checkIdent :: IdentBase NoInfo Name -> TermTypeM Ident +checkIdent (Ident name _ loc) = do + (QualName _ name', vt) <- lookupVar loc (qualName name) + return $ Ident name' (Info vt) loc + +checkDimIndex :: DimIndexBase NoInfo Name -> TermTypeM DimIndex +checkDimIndex (DimFix i) = + DimFix <$> (require "use as index" anySignedType =<< checkExp i) +checkDimIndex (DimSlice i j s) = + DimSlice <$> check i <*> check j <*> check s + where + check = + maybe (return Nothing) $ + fmap Just . unifies "use as index" (Scalar $ Prim $ Signed Int64) <=< checkExp + +sequentially :: TermTypeM a -> (a -> Occurences -> TermTypeM b) -> TermTypeM b +sequentially m1 m2 = do + (a, m1flow) <- collectOccurences m1 + (b, m2flow) <- collectOccurences $ m2 a m1flow + occur $ m1flow `seqOccurences` m2flow + return b + +type Arg = (Exp, PatternType, Occurences, SrcLoc) + +argExp :: Arg -> Exp +argExp (e, _, _, _) = e + +argType :: Arg -> PatternType +argType (_, t, _, _) = t + +checkArg :: UncheckedExp -> TermTypeM Arg +checkArg arg = do + (arg', dflow) <- collectOccurences $ checkExp arg + arg_t <- expType arg' + return (arg', arg_t, dflow, srclocOf arg') + +instantiateDimsInType :: + SrcLoc -> + RigidSource -> + TypeBase (DimDecl VName) als -> + TermTypeM (TypeBase (DimDecl VName) als, [VName]) +instantiateDimsInType tloc rsrc = + instantiateEmptyArrayDims tloc "d" $ Rigid rsrc + +instantiateDimsInReturnType :: + SrcLoc -> + Maybe (QualName VName) -> + TypeBase (DimDecl VName) als -> + TermTypeM (TypeBase (DimDecl VName) als, [VName]) +instantiateDimsInReturnType tloc fname = + instantiateEmptyArrayDims tloc "ret" $ Rigid $ RigidRet fname + +-- Some information about the function/operator we are trying to +-- apply, and how many arguments it has previously accepted. Used for +-- generating nicer type errors. +type ApplyOp = (Maybe (QualName VName), Int) + +checkApply :: + SrcLoc -> + ApplyOp -> + PatternType -> + Arg -> + TermTypeM (PatternType, PatternType, Maybe VName, [VName]) +checkApply + loc + (fname, _) + (Scalar (Arrow as pname tp1 tp2)) + (argexp, argtype, dflow, argloc) = + onFailure (CheckingApply fname argexp (toStruct tp1) (toStruct argtype)) $ do + expect (mkUsage argloc "use as function argument") (toStruct tp1) (toStruct argtype) + + -- Perform substitutions of instantiated variables in the types. + tp1' <- normTypeFully tp1 + (tp2', ext) <- instantiateDimsInReturnType loc fname =<< normTypeFully tp2 + argtype' <- normTypeFully argtype + + -- Check whether this would produce an impossible return type. + let (_, tp2_paramdims, _) = dimUses $ toStruct tp2' + case filter (`S.member` tp2_paramdims) ext of + [] -> return () + ext_paramdims -> do + let onDim (NamedDim qn) + | qualLeaf qn `elem` ext_paramdims = AnyDim + onDim d = d + typeError loc mempty $ + "Anonymous size would appear in function parameter of return type:" + indent 2 (ppr (first onDim tp2')) + textwrap "This is usually because a higher-order function is used with functional arguments that return anonymous sizes, which are then used as parameters of other function arguments." + + occur [observation as loc] + + checkOccurences dflow + + case anyConsumption dflow of + Just c -> + let msg = "type of expression with consumption at " ++ locStr (location c) + in zeroOrderType (mkUsage argloc "potential consumption in expression") msg tp1 + _ -> return () + + occurs <- (dflow `seqOccurences`) <$> consumeArg argloc argtype' (diet tp1') + + checkIfConsumable loc $ S.map AliasBound $ allConsumed occurs + occur occurs + + (argext, parsubst) <- + case pname of + Named pname' -> do + (d, argext) <- sizeSubst tp1' argexp + return + ( argext, + (`M.lookup` M.singleton pname' (SizeSubst d)) + ) + _ -> return (Nothing, const Nothing) + let tp2'' = applySubst parsubst $ returnType tp2' (diet tp1') argtype' + + return (tp1', tp2'', argext, ext) + where + sizeSubst (Scalar (Prim (Signed Int64))) e = dimFromArg fname e + sizeSubst _ _ = return (AnyDim, Nothing) +checkApply loc fname tfun@(Scalar TypeVar {}) arg = do + tv <- newTypeVar loc "b" + unify (mkUsage loc "use as function") (toStruct tfun) $ + Scalar $ Arrow mempty Unnamed (toStruct (argType arg)) tv + tfun' <- normPatternType tfun + checkApply loc fname tfun' arg +checkApply loc (fname, prev_applied) ftype (argexp, _, _, _) = do + let fname' = maybe "expression" (pquote . ppr) fname + + typeError loc mempty $ + if prev_applied == 0 + then + "Cannot apply" <+> fname' <+> "as function, as it has type:" + indent 2 (ppr ftype) + else + "Cannot apply" <+> fname' <+> "to argument #" <> ppr (prev_applied + 1) + <+> pquote (shorten $ pretty $ flatten $ ppr argexp) <> "," + <+/> "as" + <+> fname' + <+> "only takes" + <+> ppr prev_applied + <+> arguments <> "." + where + arguments + | prev_applied == 1 = "argument" + | otherwise = "arguments" + +isInt64 :: Exp -> Maybe Int64 +isInt64 (Literal (SignedValue (Int64Value k')) _) = Just $ fromIntegral k' +isInt64 (IntLit k' _ _) = Just $ fromInteger k' +isInt64 (Negate x _) = negate <$> isInt64 x +isInt64 _ = Nothing + +maybeDimFromExp :: Exp -> Maybe (DimDecl VName) +maybeDimFromExp (Var v _ _) = Just $ NamedDim v +maybeDimFromExp (Parens e _) = maybeDimFromExp e +maybeDimFromExp (QualParens _ e _) = maybeDimFromExp e +maybeDimFromExp e = ConstDim . fromIntegral <$> isInt64 e + +dimFromExp :: (Exp -> SizeSource) -> Exp -> TermTypeM (DimDecl VName, Maybe VName) +dimFromExp rf (Parens e _) = dimFromExp rf e +dimFromExp rf (QualParens _ e _) = dimFromExp rf e +dimFromExp rf e + | Just d <- maybeDimFromExp e = + return (d, Nothing) + | otherwise = + extSize (srclocOf e) $ rf e + +dimFromArg :: Maybe (QualName VName) -> Exp -> TermTypeM (DimDecl VName, Maybe VName) +dimFromArg fname = dimFromExp $ SourceArg (FName fname) . bareExp + +-- | @returnType ret_type arg_diet arg_type@ gives result of applying +-- an argument the given types to a function with the given return +-- type, consuming the argument with the given diet. +returnType :: + PatternType -> + Diet -> + PatternType -> + PatternType +returnType (Array _ Unique et shape) _ _ = + Array mempty Unique et shape +returnType (Array als Nonunique et shape) d arg = + Array (als <> arg_als) Unique et shape -- Intentional! + where + arg_als = aliases $ maskAliases arg d +returnType (Scalar (Record fs)) d arg = + Scalar $ Record $ fmap (\et -> returnType et d arg) fs +returnType (Scalar (Prim t)) _ _ = + Scalar $ Prim t +returnType (Scalar (TypeVar _ Unique t targs)) _ _ = + Scalar $ TypeVar mempty Unique t targs +returnType (Scalar (TypeVar als Nonunique t targs)) d arg = + Scalar $ TypeVar (als <> arg_als) Unique t targs -- Intentional! + where + arg_als = aliases $ maskAliases arg d +returnType (Scalar (Arrow old_als v t1 t2)) d arg = + Scalar $ Arrow als v (t1 `setAliases` mempty) (t2 `setAliases` als) + where + -- Make sure to propagate the aliases of an existing closure. + als = old_als <> aliases (maskAliases arg d) +returnType (Scalar (Sum cs)) d arg = + Scalar $ Sum $ (fmap . fmap) (\et -> returnType et d arg) cs + +-- | @t `maskAliases` d@ removes aliases (sets them to 'mempty') from +-- the parts of @t@ that are denoted as consumed by the 'Diet' @d@. +maskAliases :: + Monoid as => + TypeBase shape as -> + Diet -> + TypeBase shape as +maskAliases t Consume = t `setAliases` mempty +maskAliases t Observe = t +maskAliases (Scalar (Record ets)) (RecordDiet ds) = + Scalar $ Record $ M.intersectionWith maskAliases ets ds +maskAliases t FuncDiet {} = t +maskAliases _ _ = error "Invalid arguments passed to maskAliases." + +consumeArg :: SrcLoc -> PatternType -> Diet -> TermTypeM [Occurence] +consumeArg loc (Scalar (Record ets)) (RecordDiet ds) = + concat . M.elems <$> traverse (uncurry $ consumeArg loc) (M.intersectionWith (,) ets ds) +consumeArg loc (Array _ Nonunique _ _) Consume = + typeError loc mempty "Consuming parameter passed non-unique argument." +consumeArg loc (Scalar (TypeVar _ Nonunique _ _)) Consume = + typeError loc mempty "Consuming parameter passed non-unique argument." +consumeArg loc (Scalar (Arrow _ _ t1 _)) (FuncDiet d _) + | not $ contravariantArg t1 d = + typeError loc mempty "Non-consuming higher-order parameter passed consuming argument." + where + contravariantArg (Array _ Unique _ _) Observe = + False + contravariantArg (Scalar (TypeVar _ Unique _ _)) Observe = + False + contravariantArg (Scalar (Record ets)) (RecordDiet ds) = + and (M.intersectionWith contravariantArg ets ds) + contravariantArg (Scalar (Arrow _ _ tp tr)) (FuncDiet dp dr) = + contravariantArg tp dp && contravariantArg tr dr + contravariantArg _ _ = + True +consumeArg loc (Scalar (Arrow _ _ _ t2)) (FuncDiet _ pd) = + consumeArg loc t2 pd +consumeArg loc at Consume = return [consumption (aliases at) loc] +consumeArg loc at _ = return [observation (aliases at) loc] + +-- | Type-check a single expression in isolation. This expression may +-- turn out to be polymorphic, in which case the list of type +-- parameters will be non-empty. +checkOneExp :: UncheckedExp -> TypeM ([TypeParam], Exp) +checkOneExp e = fmap fst . runTermTypeM $ do + e' <- checkExp e + let t = toStruct $ typeOf e' + (tparams, _, _, _) <- + letGeneralise (nameFromString "") (srclocOf e) [] [] t + fixOverloadedTypes $ typeVars t + e'' <- updateTypes e' + checkUnmatched e'' + causalityCheck e'' + literalOverflowCheck e'' + return (tparams, e'') + +-- Verify that all sum type constructors and empty array literals have +-- a size that is known (rigid or a type parameter). This is to +-- ensure that we can actually determine their shape at run-time. +causalityCheck :: Exp -> TermTypeM () +causalityCheck binding_body = do + constraints <- getConstraints + + let checkCausality what known t loc + | (d, dloc) : _ <- + mapMaybe (unknown constraints known) $ + S.toList $ typeDimNames $ toStruct t = + Just $ lift $ causality what loc d dloc t + | otherwise = Nothing + + checkParamCausality known p = + checkCausality (ppr p) known (patternType p) (srclocOf p) + + onExp :: + S.Set VName -> + Exp -> + StateT (S.Set VName) (Either TypeError) Exp + + onExp known (Var v (Info t) loc) + | Just bad <- checkCausality (pquote (ppr v)) known t loc = + bad + onExp known (ArrayLit [] (Info t) loc) + | Just bad <- checkCausality "empty array" known t loc = + bad + onExp known (Lambda params _ _ _ _) + | bad : _ <- mapMaybe (checkParamCausality known) params = + bad + onExp known e@(Coerce what _ (_, Info ext) _) = do + modify (S.fromList ext <>) + void $ onExp known what + return e + onExp known e@(LetPat _ bindee_e body_e (_, Info ext) _) = do + sequencePoint known bindee_e body_e ext + return e + onExp known e@(Apply f arg (Info (_, p)) (_, Info ext) _) = do + sequencePoint known arg f $ maybeToList p ++ ext + return e + onExp + known + e@( BinOp + (f, floc) + ft + (x, Info (_, xp)) + (y, Info (_, yp)) + _ + (Info ext) + _ + ) = do + args_known <- + lift $ + execStateT (sequencePoint known x y $ catMaybes [xp, yp]) mempty + void $ onExp (args_known <> known) (Var f ft floc) + modify ((args_known <> S.fromList ext) <>) + return e + onExp known e = do + recurse known e + + case e of + DoLoop _ _ _ _ _ (Info (_, ext)) _ -> + modify (<> S.fromList ext) + If _ _ _ (_, Info ext) _ -> + modify (<> S.fromList ext) + Index _ _ (_, Info ext) _ -> + modify (<> S.fromList ext) + Match _ _ (_, Info ext) _ -> + modify (<> S.fromList ext) + Range _ _ _ (_, Info ext) _ -> + modify (<> S.fromList ext) + _ -> + return () + + return e + + recurse known = void . astMap mapper + where + mapper = identityMapper {mapOnExp = onExp known} + + sequencePoint known x y ext = do + new_known <- lift $ execStateT (onExp known x) mempty + void $ onExp (new_known <> known) y + modify ((new_known <> S.fromList ext) <>) + + either throwError (const $ return ()) $ + evalStateT (onExp mempty binding_body) mempty + where + unknown constraints known v = do + guard $ v `S.notMember` known + loc <- unknowable constraints v + return (v, loc) + + unknowable constraints v = + case snd <$> M.lookup v constraints of + Just (UnknowableSize loc _) -> Just loc + _ -> Nothing + + causality what loc d dloc t = + Left $ + TypeError loc mempty $ + "Causality check: size" <+/> pquote (pprName d) + <+/> "needed for type of" + <+> what <> colon + indent 2 (ppr t) + "But" + <+> pquote (pprName d) + <+> "is computed at" + <+/> text (locStrRel loc dloc) <> "." + "" + "Hint:" + <+> align + ( textwrap "Bind the expression producing" <+> pquote (pprName d) + <+> "with 'let' beforehand." + ) + +-- | Traverse the expression, emitting warnings if any of the literals overflow +-- their inferred types +-- +-- Note: currently unable to detect float underflow (such as 1e-400 -> 0) +literalOverflowCheck :: Exp -> TermTypeM () +literalOverflowCheck = void . check + where + check e@(IntLit x ty loc) = + e <$ case ty of + Info (Scalar (Prim t)) -> warnBounds (inBoundsI x t) x t loc + _ -> error "Inferred type of int literal is not a number" + check e@(FloatLit x ty loc) = + e <$ case ty of + Info (Scalar (Prim (FloatType t))) -> warnBounds (inBoundsF x t) x t loc + _ -> error "Inferred type of float literal is not a float" + check e@(Negate (IntLit x ty loc1) loc2) = + e <$ case ty of + Info (Scalar (Prim t)) -> warnBounds (inBoundsI (- x) t) (- x) t (loc1 <> loc2) + _ -> error "Inferred type of int literal is not a number" + check e = astMap identityMapper {mapOnExp = check} e + bitWidth ty = 8 * intByteSize ty :: Int + inBoundsI x (Signed t) = x >= -2 ^ (bitWidth t - 1) && x < 2 ^ (bitWidth t - 1) + inBoundsI x (Unsigned t) = x >= 0 && x < 2 ^ bitWidth t + inBoundsI x (FloatType Float32) = not $ isInfinite (fromIntegral x :: Float) + inBoundsI x (FloatType Float64) = not $ isInfinite (fromIntegral x :: Double) + inBoundsI _ Bool = error "Inferred type of int literal is not a number" + inBoundsF x Float32 = not $ isInfinite (realToFrac x :: Float) + inBoundsF x Float64 = not $ isInfinite x + warnBounds inBounds x ty loc = + unless inBounds $ + typeError loc mempty $ + "Literal " <> ppr x + <> " out of bounds for inferred type " + <> ppr ty + <> "." + +-- | Type-check a top-level (or module-level) function definition. +-- Despite the name, this is also used for checking constant +-- definitions, by treating them as 0-ary functions. +checkFunDef :: + ( Name, + Maybe UncheckedTypeExp, + [UncheckedTypeParam], + [UncheckedPattern], + UncheckedExp, + SrcLoc + ) -> + TypeM + ( VName, + [TypeParam], + [Pattern], + Maybe (TypeExp VName), + StructType, + [VName], + Exp + ) +checkFunDef (fname, maybe_retdecl, tparams, params, body, loc) = + fmap fst $ + runTermTypeM $ do + (tparams', params', maybe_retdecl', rettype', retext, body') <- + checkBinding (fname, maybe_retdecl, tparams, params, body, loc) + + -- Since this is a top-level function, we also resolve overloaded + -- types, using either defaults or complaining about ambiguities. + fixOverloadedTypes $ + typeVars rettype' <> foldMap (typeVars . patternType) params' + + -- Then replace all inferred types in the body and parameters. + body'' <- updateTypes body' + params'' <- updateTypes params' + maybe_retdecl'' <- traverse updateTypes maybe_retdecl' + rettype'' <- normTypeFully rettype' + + -- Check if pattern matches are exhaustive and yield + -- errors if not. + checkUnmatched body'' + + -- Check if the function body can actually be evaluated. + causalityCheck body'' + + literalOverflowCheck body'' + + bindSpaced [(Term, fname)] $ do + fname' <- checkName Term fname loc + when (nameToString fname `elem` doNotShadow) $ + typeError loc mempty $ + "The" <+> pprName fname <+> "operator may not be redefined." + + return (fname', tparams', params'', maybe_retdecl'', rettype'', retext, body'') + +-- | This is "fixing" as in "setting them", not "correcting them". We +-- only make very conservative fixing. +fixOverloadedTypes :: Names -> TermTypeM () +fixOverloadedTypes tyvars_at_toplevel = + getConstraints >>= mapM_ fixOverloaded . M.toList . M.map snd + where + fixOverloaded (v, Overloaded ots usage) + | Signed Int32 `elem` ots = do + unify usage (Scalar (TypeVar () Nonunique (typeName v) [])) $ + Scalar $ Prim $ Signed Int32 + when (v `S.member` tyvars_at_toplevel) $ + warn usage "Defaulting ambiguous type to i32." + | FloatType Float64 `elem` ots = do + unify usage (Scalar (TypeVar () Nonunique (typeName v) [])) $ + Scalar $ Prim $ FloatType Float64 + when (v `S.member` tyvars_at_toplevel) $ + warn usage "Defaulting ambiguous type to f64." + | otherwise = + typeError usage mempty $ + "Type is ambiguous (could be one of" <+> commasep (map ppr ots) <> ")." + "Add a type annotation to disambiguate the type." + fixOverloaded (_, NoConstraint _ usage) = + typeError usage mempty $ + "Type of expression is ambiguous." + "Add a type annotation to disambiguate the type." + fixOverloaded (_, Equality usage) = + typeError usage mempty $ + "Type is ambiguous (must be equality type)." + "Add a type annotation to disambiguate the type." + fixOverloaded (_, HasFields fs usage) = + typeError usage mempty $ + "Type is ambiguous. Must be record with fields:" + indent 2 (stack $ map field $ M.toList fs) + "Add a type annotation to disambiguate the type." + where + field (l, t) = ppr l <> colon <+> align (ppr t) + fixOverloaded (_, HasConstrs cs usage) = + typeError usage mempty $ + "Type is ambiguous (must be a sum type with constructors:" + <+> ppr (Sum cs) <> ")." + "Add a type annotation to disambiguate the type." + fixOverloaded (_, Size Nothing usage) = + typeError usage mempty "Size is ambiguous." + fixOverloaded _ = return () + +hiddenParamNames :: [Pattern] -> Names +hiddenParamNames params = hidden + where + param_all_names = mconcat $ map patternNames params + named (Named x, _) = Just x + named (Unnamed, _) = Nothing + param_names = + S.fromList $ mapMaybe (named . patternParam) params + hidden = param_all_names `S.difference` param_names + +inferredReturnType :: SrcLoc -> [Pattern] -> PatternType -> TermTypeM StructType +inferredReturnType loc params t = + -- The inferred type may refer to names that are bound by the + -- parameter patterns, but which will not be visible in the type. + -- These we must turn into fresh type variables, which will be + -- existential in the return type. + fmap (toStruct . fst) $ + unscopeType + loc + (M.filterWithKey (const . (`S.member` hidden)) $ foldMap patternMap params) + $ inferReturnUniqueness params t + where + hidden = hiddenParamNames params + +checkBinding :: + ( Name, + Maybe UncheckedTypeExp, + [UncheckedTypeParam], + [UncheckedPattern], + UncheckedExp, + SrcLoc + ) -> + TermTypeM + ( [TypeParam], + [Pattern], + Maybe (TypeExp VName), + StructType, + [VName], + Exp + ) +checkBinding (fname, maybe_retdecl, tparams, params, body, loc) = + noUnique $ + incLevel $ + bindingParams tparams params $ \tparams' params' -> do + when (null params && any isSizeParam tparams) $ + typeError + loc + mempty + "Size parameters are only allowed on bindings that also have value parameters." + + maybe_retdecl' <- forM maybe_retdecl $ \retdecl -> do + (retdecl', ret_nodims, _) <- checkTypeExp retdecl + (ret, _) <- instantiateEmptyArrayDims loc "funret" Nonrigid ret_nodims + return (retdecl', ret) + + body' <- + checkFunBody + params' + body + (snd <$> maybe_retdecl') + (maybe loc srclocOf maybe_retdecl) + + params'' <- mapM updateTypes params' + body_t <- expTypeFully body' + + (maybe_retdecl'', rettype) <- case maybe_retdecl' of + Just (retdecl', ret) -> do + let rettype_structural = toStructural ret + checkReturnAlias rettype_structural params'' body_t + + when (null params) $ nothingMustBeUnique loc rettype_structural + + ret' <- normTypeFully ret + + return (Just retdecl', ret') + Nothing + | null params -> + return (Nothing, toStruct $ body_t `setUniqueness` Nonunique) + | otherwise -> do + body_t' <- inferredReturnType loc params'' body_t + return (Nothing, body_t') + + verifyFunctionParams (Just fname) params'' + + (tparams'', params''', rettype'', retext) <- + letGeneralise fname loc tparams' params'' rettype + + checkGlobalAliases params'' body_t loc + + return (tparams'', params''', maybe_retdecl'', rettype'', retext, body') + where + checkReturnAlias rettp params' = + foldM_ (checkReturnAlias' params') S.empty . returnAliasing rettp + checkReturnAlias' params' seen (Unique, names) + | any (`S.member` S.map snd seen) $ S.toList names = + uniqueReturnAliased fname loc + | otherwise = do + notAliasingParam params' names + return $ seen `S.union` tag Unique names + checkReturnAlias' _ seen (Nonunique, names) + | any (`S.member` seen) $ S.toList $ tag Unique names = + uniqueReturnAliased fname loc + | otherwise = return $ seen `S.union` tag Nonunique names + + notAliasingParam params' names = + forM_ params' $ \p -> + let consumedNonunique p' = + not (unique $ unInfo $ identType p') && (identName p' `S.member` names) + in case find consumedNonunique $ S.toList $ patternIdents p of + Just p' -> + returnAliased fname (baseName $ identName p') loc + Nothing -> + return () + + tag u = S.map (u,) + + returnAliasing (Scalar (Record ets1)) (Scalar (Record ets2)) = + concat $ M.elems $ M.intersectionWith returnAliasing ets1 ets2 + returnAliasing expected got = + [(uniqueness expected, S.map aliasVar $ aliases got)] + +-- | Extract all the shape names that occur in positive position +-- (roughly, left side of an arrow) in a given type. +typeDimNamesPos :: TypeBase (DimDecl VName) als -> S.Set VName +typeDimNamesPos (Scalar (Arrow _ _ t1 t2)) = onParam t1 <> typeDimNamesPos t2 + where + onParam :: TypeBase (DimDecl VName) als -> S.Set VName + onParam (Scalar Arrow {}) = mempty + onParam (Scalar (Record fs)) = mconcat $ map onParam $ M.elems fs + onParam (Scalar (TypeVar _ _ _ targs)) = mconcat $ map onTypeArg targs + onParam t = typeDimNames t + onTypeArg (TypeArgDim (NamedDim d) _) = S.singleton $ qualLeaf d + onTypeArg (TypeArgDim _ _) = mempty + onTypeArg (TypeArgType t _) = onParam t +typeDimNamesPos _ = mempty + +checkGlobalAliases :: [Pattern] -> PatternType -> SrcLoc -> TermTypeM () +checkGlobalAliases params body_t loc = do + vtable <- asks $ scopeVtable . termScope + let isLocal v = case v `M.lookup` vtable of + Just (BoundV Local _ _) -> True + _ -> False + let als = + filter (not . isLocal) $ + S.toList $ + boundArrayAliases body_t + `S.difference` foldMap patternNames params + case als of + v : _ + | not $ null params -> + typeError loc mempty $ + "Function result aliases the free variable " + <> pquote (pprName v) + <> "." + "Use" <+> pquote "copy" <+> "to break the aliasing." + _ -> + return () + +inferReturnUniqueness :: [Pattern] -> PatternType -> PatternType +inferReturnUniqueness params t = + let forbidden = aliasesMultipleTimes t + uniques = uniqueParamNames params + delve (Scalar (Record fs)) = + Scalar $ Record $ M.map delve fs + delve t' + | all (`S.member` uniques) (boundArrayAliases t'), + not $ any ((`S.member` forbidden) . aliasVar) (aliases t') = + t' + | otherwise = + t' `setUniqueness` Nonunique + in delve t + +-- An alias inhibits uniqueness if it is used in disjoint values. +aliasesMultipleTimes :: PatternType -> Names +aliasesMultipleTimes = S.fromList . map fst . filter ((> 1) . snd) . M.toList . delve + where + delve (Scalar (Record fs)) = + foldl' (M.unionWith (+)) mempty $ map delve $ M.elems fs + delve t = + M.fromList $ zip (map aliasVar $ S.toList (aliases t)) $ repeat (1 :: Int) + +uniqueParamNames :: [Pattern] -> Names +uniqueParamNames = + S.map identName + . S.filter (unique . unInfo . identType) + . foldMap patternIdents + +boundArrayAliases :: PatternType -> S.Set VName +boundArrayAliases (Array als _ _ _) = boundAliases als +boundArrayAliases (Scalar Prim {}) = mempty +boundArrayAliases (Scalar (Record fs)) = foldMap boundArrayAliases fs +boundArrayAliases (Scalar (TypeVar als _ _ _)) = boundAliases als +boundArrayAliases (Scalar Arrow {}) = mempty +boundArrayAliases (Scalar (Sum fs)) = + mconcat $ concatMap (map boundArrayAliases) $ M.elems fs + +-- | The set of in-scope variables that are being aliased. +boundAliases :: Aliasing -> S.Set VName +boundAliases = S.map aliasVar . S.filter bound + where + bound AliasBound {} = True + bound AliasFree {} = False + +nothingMustBeUnique :: SrcLoc -> TypeBase () () -> TermTypeM () +nothingMustBeUnique loc = check + where + check (Array _ Unique _ _) = bad + check (Scalar (TypeVar _ Unique _ _)) = bad + check (Scalar (Record fs)) = mapM_ check fs + check (Scalar (Sum fs)) = mapM_ (mapM_ check) fs + check _ = return () + bad = typeError loc mempty "A top-level constant cannot have a unique type." + +-- | Verify certain restrictions on function parameters, and bail out +-- on dubious constructions. +-- +-- These restrictions apply to all functions (anonymous or otherwise). +-- Top-level functions have further restrictions that are checked +-- during let-generalisation. +verifyFunctionParams :: Maybe Name -> [Pattern] -> TermTypeM () +verifyFunctionParams fname params = + onFailure (CheckingParams fname) $ + verifyParams (foldMap patternNames params) =<< mapM updateTypes params + where + verifyParams forbidden (p : ps) + | d : _ <- S.toList $ patternDimNames p `S.intersection` forbidden = + typeError p mempty $ + "Parameter" <+> pquote (ppr p) + <+/> "refers to size" <+> pquote (pprName d) + <> comma + <+/> textwrap "which will not be accessible to the caller" + <> comma + <+/> textwrap "possibly because it is nested in a tuple or record." + <+/> textwrap "Consider ascribing an explicit type that does not reference " + <> pquote (pprName d) + <> "." + | otherwise = verifyParams forbidden' ps + where + forbidden' = + case patternParam p of + (Named v, _) -> forbidden `S.difference` S.singleton v + _ -> forbidden + verifyParams _ [] = return () + +-- Returns the sizes of the immediate type produced, +-- the sizes of parameter types, and the sizes of return types. +dimUses :: StructType -> (Names, Names, Names) +dimUses = execWriter . traverseDims f + where + f _ PosImmediate (NamedDim v) = tell (S.singleton (qualLeaf v), mempty, mempty) + f _ PosParam (NamedDim v) = tell (mempty, S.singleton (qualLeaf v), mempty) + f _ PosReturn (NamedDim v) = tell (mempty, mempty, S.singleton (qualLeaf v)) + f _ _ _ = return () + +-- | Find at all type variables in the given type that are covered by +-- the constraints, and produce type parameters that close over them. +-- +-- The passed-in list of type parameters is always prepended to the +-- produced list of type parameters. +closeOverTypes :: + Name -> + SrcLoc -> + [TypeParam] -> + [StructType] -> + StructType -> + Constraints -> + TermTypeM ([TypeParam], StructType, [VName]) +closeOverTypes defname defloc tparams paramts ret substs = do + (more_tparams, retext) <- + partitionEithers . catMaybes + <$> mapM closeOver (M.toList $ M.map snd to_close_over) + let retToAnyDim v = do + guard $ v `S.member` ret_sizes + UnknowableSize {} <- snd <$> M.lookup v substs + Just $ SizeSubst AnyDim + return + ( tparams ++ more_tparams, + applySubst retToAnyDim ret, + retext + ) + where + t = foldFunType paramts ret + to_close_over = M.filterWithKey (\k _ -> k `S.member` visible) substs + visible = typeVars t <> typeDimNames t + + (produced_sizes, param_sizes, ret_sizes) = dimUses t + + -- Avoid duplicate type parameters. + closeOver (k, _) + | k `elem` map typeParamName tparams = + return Nothing + closeOver (k, NoConstraint l usage) = + return $ Just $ Left $ TypeParamType l k $ srclocOf usage + closeOver (k, ParamType l loc) = + return $ Just $ Left $ TypeParamType l k loc + closeOver (k, Size Nothing usage) = + return $ Just $ Left $ TypeParamDim k $ srclocOf usage + closeOver (k, UnknowableSize _ _) + | k `S.member` param_sizes = do + notes <- dimNotes defloc $ NamedDim $ qualName k + typeError defloc notes $ + "Unknowable size" <+> pquote (pprName k) + <+> "imposes constraint on type of" + <+> pquote (pprName defname) + <> ", which is inferred as:" + indent 2 (ppr t) + | k `S.member` produced_sizes = + return $ Just $ Right k + closeOver (_, _) = + return Nothing + +letGeneralise :: + Name -> + SrcLoc -> + [TypeParam] -> + [Pattern] -> + StructType -> + TermTypeM ([TypeParam], [Pattern], StructType, [VName]) +letGeneralise defname defloc tparams params rettype = + onFailure (CheckingLetGeneralise defname) $ do + now_substs <- getConstraints + + -- Candidates for let-generalisation are those type variables that + -- + -- (1) were not known before we checked this function, and + -- + -- (2) are not used in the (new) definition of any type variables + -- known before we checked this function. + -- + -- (3) are not referenced from an overloaded type (for example, + -- are the element types of an incompletely resolved record type). + -- This is a bit more restrictive than I'd like, and SML for + -- example does not have this restriction. + -- + -- Criteria (1) and (2) is implemented by looking at the binding + -- level of the type variables. + let keep_type_vars = overloadedTypeVars now_substs + + cur_lvl <- curLevel + let candidate k (lvl, _) = (k `S.notMember` keep_type_vars) && lvl >= cur_lvl + new_substs = M.filterWithKey candidate now_substs + + (tparams', rettype', retext) <- + closeOverTypes + defname + defloc + tparams + (map patternStructType params) + rettype + new_substs + + rettype'' <- updateTypes rettype' + + let used_sizes = + foldMap typeDimNames $ + rettype'' : map patternStructType params + case filter ((`S.notMember` used_sizes) . typeParamName) $ + filter isSizeParam tparams' of + [] -> return () + tp : _ -> + typeError defloc mempty $ + "Size parameter" <+> pquote (ppr tp) <+> "unused." + + -- We keep those type variables that were not closed over by + -- let-generalisation. + modifyConstraints $ M.filterWithKey $ \k _ -> k `notElem` map typeParamName tparams' + + return (tparams', params, rettype'', retext) + +checkFunBody :: + [Pattern] -> + UncheckedExp -> + Maybe StructType -> + SrcLoc -> + TermTypeM Exp +checkFunBody params body maybe_rettype loc = do + body' <- noSizeEscape $ checkExp body + + -- Unify body return type with return annotation, if one exists. + case maybe_rettype of + Just rettype -> do + (rettype_withdims, _) <- instantiateEmptyArrayDims loc "impl" Nonrigid rettype + + body_t <- expTypeFully body' + -- We need to turn any sizes provided by "hidden" parameter + -- names into existential sizes instead. + let hidden = hiddenParamNames params + (body_t', _) <- + unscopeType + loc + ( M.filterWithKey (const . (`S.member` hidden)) $ + foldMap patternMap params + ) + body_t + + let usage = mkUsage (srclocOf body) "return type annotation" + onFailure (CheckingReturn rettype (toStruct body_t')) $ + expect usage rettype_withdims $ toStruct body_t' + + -- We also have to make sure that uniqueness matches. This is done + -- explicitly, because uniqueness is ignored by unification. + rettype' <- normTypeFully rettype + body_t'' <- normTypeFully rettype -- Substs may have changed. + unless (body_t'' `subtypeOf` anySizes rettype') $ + typeError (srclocOf body) mempty $ + "Body type" indent 2 (ppr body_t'') + "is not a subtype of annotated type" + indent 2 (ppr rettype') + Nothing -> return () + + return body' + +--- Consumption + +occur :: Occurences -> TermTypeM () +occur = tell + +-- | Proclaim that we have made read-only use of the given variable. +observe :: Ident -> TermTypeM () +observe (Ident nm (Info t) loc) = + let als = AliasBound nm `S.insert` aliases t + in occur [observation als loc] + +checkIfConsumable :: SrcLoc -> Aliasing -> TermTypeM () +checkIfConsumable loc als = do + vtable <- asks $ scopeVtable . termScope + let consumable v = case M.lookup v vtable of + Just (BoundV Local _ t) + | arrayRank t > 0 -> unique t + | Scalar TypeVar {} <- t -> unique t + | otherwise -> True + _ -> False + case filter (not . consumable) $ map aliasVar $ S.toList als of + v : _ -> + typeError loc mempty $ + "Would consume variable" <+> pquote (pprName v) + <> ", which is not allowed." + [] -> return () + +-- | Proclaim that we have written to the given variable. +consume :: SrcLoc -> Aliasing -> TermTypeM () +consume loc als = do + checkIfConsumable loc als + occur [consumption als loc] + +-- | Proclaim that we have written to the given variable, and mark +-- accesses to it and all of its aliases as invalid inside the given +-- computation. +consuming :: Ident -> TermTypeM a -> TermTypeM a +consuming (Ident name (Info t) loc) m = do + consume loc $ AliasBound name `S.insert` aliases t + localScope consume' m + where + consume' scope = + scope {scopeVtable = M.insert name (WasConsumed loc) $ scopeVtable scope} + +collectOccurences :: TermTypeM a -> TermTypeM (a, Occurences) +collectOccurences m = pass $ do + (x, dataflow) <- listen m + return ((x, dataflow), const mempty) + +tapOccurences :: TermTypeM a -> TermTypeM (a, Occurences) +tapOccurences = listen + +removeSeminullOccurences :: TermTypeM a -> TermTypeM a +removeSeminullOccurences = censor $ filter $ not . seminullOccurence + +checkIfUsed :: Occurences -> Ident -> TermTypeM () +checkIfUsed occs v + | not $ identName v `S.member` allOccuring occs, + not $ "_" `isPrefixOf` prettyName (identName v) = + warn (srclocOf v) $ "Unused variable " ++ quote (pretty $ baseName $ identName v) ++ "." + | otherwise = + return () + +alternative :: TermTypeM a -> TermTypeM b -> TermTypeM (a, b) +alternative m1 m2 = pass $ do + (x, occurs1) <- listen $ noSizeEscape m1 + (y, occurs2) <- listen $ noSizeEscape m2 + checkOccurences occurs1 + checkOccurences occurs2 + let usage = occurs1 `altOccurences` occurs2 + return ((x, y), const usage) + +-- | Make all bindings nonunique. +noUnique :: TermTypeM a -> TermTypeM a +noUnique = localScope (\scope -> scope {scopeVtable = M.map set $ scopeVtable scope}) + where + set (BoundV l tparams t) = BoundV l tparams $ t `setUniqueness` Nonunique + set (OverloadedF ts pts rt) = OverloadedF ts pts rt + set EqualityF = EqualityF + set (WasConsumed loc) = WasConsumed loc + +onlySelfAliasing :: TermTypeM a -> TermTypeM a +onlySelfAliasing = localScope (\scope -> scope {scopeVtable = M.mapWithKey set $ scopeVtable scope}) + where + set k (BoundV l tparams t) = + BoundV l tparams $ + t `addAliases` S.intersection (S.singleton (AliasBound k)) + set _ (OverloadedF ts pts rt) = OverloadedF ts pts rt + set _ EqualityF = EqualityF + set _ (WasConsumed loc) = WasConsumed loc + +arrayOfM :: + (Pretty (ShapeDecl dim), Monoid as) => + SrcLoc -> + TypeBase dim as -> + ShapeDecl dim -> + Uniqueness -> + TermTypeM (TypeBase dim as) +arrayOfM loc t shape u = do + zeroOrderType (mkUsage loc "use as array element") "type used in array" t + return $ arrayOf t shape u + +updateTypes :: ASTMappable e => e -> TermTypeM e +updateTypes = astMap tv + where + tv = + ASTMapper + { mapOnExp = astMap tv, + mapOnName = pure, + mapOnQualName = pure, + mapOnStructType = normTypeFully, + mapOnPatternType = normTypeFully + } diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs --- a/src/Language/Futhark/TypeChecker/Types.hs +++ b/src/Language/Futhark/TypeChecker/Types.hs @@ -2,26 +2,23 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Safe #-} + -- | Type checker building blocks that do not involve unification. module Language.Futhark.TypeChecker.Types - ( checkTypeExp - , checkTypeDecl - - , unifyTypesU - , subtypeOf - , subuniqueOf - - , checkForDuplicateNames - , checkTypeParams - , typeParamToArg - - , TypeSub(..) - , TypeSubs - , substituteTypes - - , Subst(..) - , Substitutable(..) - , substTypesAny + ( checkTypeExp, + checkTypeDecl, + unifyTypesU, + subtypeOf, + subuniqueOf, + checkForDuplicateNames, + checkTypeParams, + typeParamToArg, + TypeSub (..), + TypeSubs, + substituteTypes, + Subst (..), + Substitutable (..), + substTypesAny, ) where @@ -29,57 +26,70 @@ import Control.Monad.Reader import Control.Monad.State import Data.Bifunctor -import Data.List (foldl', sort, nub) -import Data.Maybe +import Data.List (foldl', nub, sort) import qualified Data.Map.Strict as M - +import Data.Maybe +import Futhark.Util.Pretty import Language.Futhark -import Language.Futhark.TypeChecker.Monad import Language.Futhark.Traversals -import Futhark.Util.Pretty +import Language.Futhark.TypeChecker.Monad -- | @unifyTypes uf t1 t2@ attempts to unify @t1@ and @t2@. If -- unification cannot happen, 'Nothing' is returned, otherwise a type -- that combines the aliasing of @t1@ and @t2@ is returned. -- Uniqueness is unified with @uf@. -unifyTypesU :: (Monoid als, ArrayDim dim) => - (Uniqueness -> Uniqueness -> Maybe Uniqueness) - -> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als) +unifyTypesU :: + (Monoid als, ArrayDim dim) => + (Uniqueness -> Uniqueness -> Maybe Uniqueness) -> + TypeBase dim als -> + TypeBase dim als -> + Maybe (TypeBase dim als) unifyTypesU uf (Array als1 u1 et1 shape1) (Array als2 u2 et2 shape2) = Array (als1 <> als2) <$> uf u1 u2 - <*> unifyScalarTypes uf et1 et2 <*> unifyShapes shape1 shape2 + <*> unifyScalarTypes uf et1 et2 + <*> unifyShapes shape1 shape2 unifyTypesU uf (Scalar t1) (Scalar t2) = Scalar <$> unifyScalarTypes uf t1 t2 unifyTypesU _ _ _ = Nothing -unifyScalarTypes :: (Monoid als, ArrayDim dim) => - (Uniqueness -> Uniqueness -> Maybe Uniqueness) - -> ScalarTypeBase dim als -> ScalarTypeBase dim als -> Maybe (ScalarTypeBase dim als) +unifyScalarTypes :: + (Monoid als, ArrayDim dim) => + (Uniqueness -> Uniqueness -> Maybe Uniqueness) -> + ScalarTypeBase dim als -> + ScalarTypeBase dim als -> + Maybe (ScalarTypeBase dim als) unifyScalarTypes _ (Prim t1) (Prim t2) - | t1 == t2 = Just $ Prim t1 + | t1 == t2 = Just $ Prim t1 | otherwise = Nothing unifyScalarTypes uf (TypeVar als1 u1 t1 targs1) (TypeVar als2 u2 t2 targs2) | t1 == t2 = do - u3 <- uf u1 u2 - targs3 <- zipWithM (unifyTypeArgs uf) targs1 targs2 - Just $ TypeVar (als1 <> als2) u3 t1 targs3 + u3 <- uf u1 u2 + targs3 <- zipWithM (unifyTypeArgs uf) targs1 targs2 + Just $ TypeVar (als1 <> als2) u3 t1 targs3 | otherwise = Nothing unifyScalarTypes uf (Record ts1) (Record ts2) | length ts1 == length ts2, sort (M.keys ts1) == sort (M.keys ts2) = - Record <$> traverse (uncurry (unifyTypesU uf)) - (M.intersectionWith (,) ts1 ts2) + Record + <$> traverse + (uncurry (unifyTypesU uf)) + (M.intersectionWith (,) ts1 ts2) unifyScalarTypes uf (Arrow as1 mn1 t1 t1') (Arrow as2 _ t2 t2') = Arrow (as1 <> as2) mn1 <$> unifyTypesU (flip uf) t1 t2 <*> unifyTypesU uf t1' t2' unifyScalarTypes uf (Sum cs1) (Sum cs2) | length cs1 == length cs2, sort (M.keys cs1) == sort (M.keys cs2) = - Sum <$> traverse (uncurry (zipWithM (unifyTypesU uf))) - (M.intersectionWith (,) cs1 cs2) + Sum + <$> traverse + (uncurry (zipWithM (unifyTypesU uf))) + (M.intersectionWith (,) cs1 cs2) unifyScalarTypes _ _ _ = Nothing -unifyTypeArgs :: (ArrayDim dim) => - (Uniqueness -> Uniqueness -> Maybe Uniqueness) - -> TypeArg dim -> TypeArg dim -> Maybe (TypeArg dim) +unifyTypeArgs :: + (ArrayDim dim) => + (Uniqueness -> Uniqueness -> Maybe Uniqueness) -> + TypeArg dim -> + TypeArg dim -> + Maybe (TypeArg dim) unifyTypeArgs _ (TypeArgDim d1 loc) (TypeArgDim d2 _) = TypeArgDim <$> unifyDims d1 d2 <*> pure loc unifyTypeArgs uf (TypeArgType t1 loc) (TypeArgType t2 _) = @@ -89,20 +99,25 @@ -- | @x \`subtypeOf\` y@ is true if @x@ is a subtype of @y@ (or equal to -- @y@), meaning @x@ is valid whenever @y@ is. -subtypeOf :: ArrayDim dim => - TypeBase dim as1 -> TypeBase dim as2 -> Bool +subtypeOf :: + ArrayDim dim => + TypeBase dim as1 -> + TypeBase dim as2 -> + Bool subtypeOf t1 t2 = isJust $ unifyTypesU unifyUniqueness (toStruct t1) (toStruct t2) - where unifyUniqueness u2 u1 = if u2 `subuniqueOf` u1 then Just u1 else Nothing + where + unifyUniqueness u2 u1 = if u2 `subuniqueOf` u1 then Just u1 else Nothing -- | @x `subuniqueOf` y@ is true if @x@ is not less unique than @y@. subuniqueOf :: Uniqueness -> Uniqueness -> Bool subuniqueOf Nonunique Unique = False -subuniqueOf _ _ = True +subuniqueOf _ _ = True -- | Use 'checkTypeExp' to check a type declaration. -checkTypeDecl :: MonadTypeChecker m => - TypeDeclBase NoInfo Name - -> m (TypeDeclBase Info VName, Liftedness) +checkTypeDecl :: + MonadTypeChecker m => + TypeDeclBase NoInfo Name -> + m (TypeDeclBase Info VName, Liftedness) checkTypeDecl (TypeDecl t NoInfo) = do checkForDuplicateNamesInType t (t', st, l) <- checkTypeExp t @@ -111,16 +126,18 @@ -- | Type-check a single 'TypeExp', returning the checked 'TypeExp', -- its fully expanded type (modulo yet-unelaborated type variables), -- and whether it is potentially higher-order. -checkTypeExp :: MonadTypeChecker m => - TypeExp Name - -> m (TypeExp VName, StructType, Liftedness) +checkTypeExp :: + MonadTypeChecker m => + TypeExp Name -> + m (TypeExp VName, StructType, Liftedness) checkTypeExp (TEVar name loc) = do (name', ps, t, l) <- lookupType loc name case ps of [] -> return (TEVar name' loc, t, l) - _ -> typeError loc mempty $ - "Type constructor" <+> pquote (spread (ppr name : map ppr ps)) <+> - "used without any arguments." + _ -> + typeError loc mempty $ + "Type constructor" <+> pquote (spread (ppr name : map ppr ps)) + <+> "used without any arguments." checkTypeExp (TETuple ts loc) = do (ts', ts_s, ls) <- unzip3 <$> mapM checkTypeExp ts return (TETuple ts' loc, tupleRecord ts_s, foldl' max Unlifted ls) @@ -131,12 +148,14 @@ typeError loc mempty $ "Duplicate record fields in" <+> ppr t <> "." fs_ts_ls <- traverse checkTypeExp $ M.fromList fs - let fs' = fmap (\(x,_,_) -> x) fs_ts_ls - ts_s = fmap (\(_,y,_) -> y) fs_ts_ls - ls = fmap (\(_,_,z) -> z) fs_ts_ls - return (TERecord (M.toList fs') loc, - Scalar $ Record ts_s, - foldl' max Unlifted ls) + let fs' = fmap (\(x, _, _) -> x) fs_ts_ls + ts_s = fmap (\(_, y, _) -> y) fs_ts_ls + ls = fmap (\(_, _, z) -> z) fs_ts_ls + return + ( TERecord (M.toList fs') loc, + Scalar $ Record ts_s, + foldl' max Unlifted ls + ) checkTypeExp (TEArray t d loc) = do (t', st, l) <- checkTypeExp t (d', d'') <- checkDimExp d @@ -144,86 +163,106 @@ (Unlifted, st') -> return (TEArray t' d' loc, st', Unlifted) (SizeLifted, _) -> typeError loc mempty $ - "Cannot create array with elements of size-lifted type" <+> pquote (ppr t) <+/> - "(might cause irregular array)." + "Cannot create array with elements of size-lifted type" <+> pquote (ppr t) + <+/> "(might cause irregular array)." (Lifted, _) -> typeError loc mempty $ - "Cannot create array with elements of lifted type" <+> pquote (ppr t) <+/> - "(might contain function)." - where checkDimExp DimExpAny = - return (DimExpAny, AnyDim) - checkDimExp (DimExpConst k dloc) = - return (DimExpConst k dloc, ConstDim k) - checkDimExp (DimExpNamed v dloc) = do - v' <- checkNamedDim loc v - return (DimExpNamed v' dloc, NamedDim v') + "Cannot create array with elements of lifted type" <+> pquote (ppr t) + <+/> "(might contain function)." + where + checkDimExp DimExpAny = + return (DimExpAny, AnyDim) + checkDimExp (DimExpConst k dloc) = + return (DimExpConst k dloc, ConstDim k) + checkDimExp (DimExpNamed v dloc) = do + v' <- checkNamedDim loc v + return (DimExpNamed v' dloc, NamedDim v') checkTypeExp (TEUnique t loc) = do (t', st, l) <- checkTypeExp t unless (mayContainArray st) $ warn loc $ "Declaring " <> quote (pretty st) <> " as unique has no effect." return (TEUnique t' loc, st `setUniqueness` Unique, l) - where mayContainArray (Scalar Prim{}) = False - mayContainArray Array{} = True - mayContainArray (Scalar (Record fs)) = any mayContainArray fs - mayContainArray (Scalar TypeVar{}) = True - mayContainArray (Scalar Arrow{}) = False - mayContainArray (Scalar (Sum cs)) = (any . any) mayContainArray cs + where + mayContainArray (Scalar Prim {}) = False + mayContainArray Array {} = True + mayContainArray (Scalar (Record fs)) = any mayContainArray fs + mayContainArray (Scalar TypeVar {}) = True + mayContainArray (Scalar Arrow {}) = False + mayContainArray (Scalar (Sum cs)) = (any . any) mayContainArray cs checkTypeExp (TEArrow (Just v) t1 t2 loc) = do (t1', st1, _) <- checkTypeExp t1 bindSpaced [(Term, v)] $ do v' <- checkName Term v loc bindVal v' (BoundV [] st1) $ do (t2', st2, _) <- checkTypeExp t2 - return (TEArrow (Just v') t1' t2' loc, - Scalar $ Arrow mempty (Named v') st1 st2, - Lifted) + return + ( TEArrow (Just v') t1' t2' loc, + Scalar $ Arrow mempty (Named v') st1 st2, + Lifted + ) checkTypeExp (TEArrow Nothing t1 t2 loc) = do (t1', st1, _) <- checkTypeExp t1 (t2', st2, _) <- checkTypeExp t2 - return (TEArrow Nothing t1' t2' loc, - Scalar $ Arrow mempty Unnamed st1 st2, - Lifted) -checkTypeExp ote@TEApply{} = do + return + ( TEArrow Nothing t1' t2' loc, + Scalar $ Arrow mempty Unnamed st1 st2, + Lifted + ) +checkTypeExp ote@TEApply {} = do (tname, tname_loc, targs) <- rootAndArgs ote (tname', ps, t, l) <- lookupType tloc tname if length ps /= length targs - then typeError tloc mempty $ - "Type constructor" <+> pquote (ppr tname) <+> "requires" <+> ppr (length ps) <+> - "arguments, but provided" <+> ppr (length targs) <> "." - else do - (targs', substs) <- unzip <$> zipWithM checkArgApply ps targs - return (foldl (\x y -> TEApply x y tloc) (TEVar tname' tname_loc) targs', - substituteTypes (mconcat substs) t, - l) - where tloc = srclocOf ote - - rootAndArgs :: MonadTypeChecker m => TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name]) - rootAndArgs (TEVar qn loc) = return (qn, loc, []) - rootAndArgs (TEApply op arg _) = do (op', loc, args) <- rootAndArgs op - return (op', loc, args++[arg]) - rootAndArgs te' = typeError (srclocOf te') mempty $ - "Type" <+> pquote (ppr te') <+> "is not a type constructor." - - checkArgApply (TypeParamDim pv _) (TypeArgExpDim (DimExpNamed v dloc) loc) = do - v' <- checkNamedDim loc v - return (TypeArgExpDim (DimExpNamed v' dloc) loc, - M.singleton pv $ DimSub $ NamedDim v') - checkArgApply (TypeParamDim pv _) (TypeArgExpDim (DimExpConst x dloc) loc) = - return (TypeArgExpDim (DimExpConst x dloc) loc, - M.singleton pv $ DimSub $ ConstDim x) - checkArgApply (TypeParamDim pv _) (TypeArgExpDim DimExpAny loc) = - return (TypeArgExpDim DimExpAny loc, - M.singleton pv $ DimSub AnyDim) - - checkArgApply (TypeParamType l pv _) (TypeArgExpType te) = do - (te', st, _) <- checkTypeExp te - return (TypeArgExpType te', - M.singleton pv $ TypeSub $ TypeAbbr l [] st) + then + typeError tloc mempty $ + "Type constructor" <+> pquote (ppr tname) <+> "requires" <+> ppr (length ps) + <+> "arguments, but provided" + <+> ppr (length targs) <> "." + else do + (targs', substs) <- unzip <$> zipWithM checkArgApply ps targs + return + ( foldl (\x y -> TEApply x y tloc) (TEVar tname' tname_loc) targs', + substituteTypes (mconcat substs) t, + l + ) + where + tloc = srclocOf ote - checkArgApply p a = - typeError tloc mempty $ "Type argument" <+> ppr a <+> - "not valid for a type parameter" <+> ppr p <> "." + rootAndArgs :: MonadTypeChecker m => TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name]) + rootAndArgs (TEVar qn loc) = return (qn, loc, []) + rootAndArgs (TEApply op arg _) = do + (op', loc, args) <- rootAndArgs op + return (op', loc, args ++ [arg]) + rootAndArgs te' = + typeError (srclocOf te') mempty $ + "Type" <+> pquote (ppr te') <+> "is not a type constructor." + checkArgApply (TypeParamDim pv _) (TypeArgExpDim (DimExpNamed v dloc) loc) = do + v' <- checkNamedDim loc v + return + ( TypeArgExpDim (DimExpNamed v' dloc) loc, + M.singleton pv $ DimSub $ NamedDim v' + ) + checkArgApply (TypeParamDim pv _) (TypeArgExpDim (DimExpConst x dloc) loc) = + return + ( TypeArgExpDim (DimExpConst x dloc) loc, + M.singleton pv $ DimSub $ ConstDim x + ) + checkArgApply (TypeParamDim pv _) (TypeArgExpDim DimExpAny loc) = + return + ( TypeArgExpDim DimExpAny loc, + M.singleton pv $ DimSub AnyDim + ) + checkArgApply (TypeParamType l pv _) (TypeArgExpType te) = do + (te', st, _) <- checkTypeExp te + return + ( TypeArgExpType te', + M.singleton pv $ TypeSub $ TypeAbbr l [] st + ) + checkArgApply p a = + typeError tloc mempty $ + "Type argument" <+> ppr a + <+> "not valid for a type parameter" + <+> ppr p <> "." checkTypeExp t@(TESum cs loc) = do let constructors = map fst cs unless (sort constructors == sort (nub constructors)) $ @@ -233,94 +272,109 @@ typeError loc mempty "Sum types must have less than 256 constructors." cs_ts_ls <- (traverse . traverse) checkTypeExp $ M.fromList cs - let cs' = (fmap . fmap) (\(x,_,_) -> x) cs_ts_ls + let cs' = (fmap . fmap) (\(x, _, _) -> x) cs_ts_ls ts_s = (fmap . fmap) (\(_, y, _) -> y) cs_ts_ls - ls = (concatMap . fmap) (\(_, _, z) -> z) cs_ts_ls - return (TESum (M.toList cs') loc, - Scalar $ Sum ts_s, - foldl' max Unlifted ls) + ls = (concatMap . fmap) (\(_, _, z) -> z) cs_ts_ls + return + ( TESum (M.toList cs') loc, + Scalar $ Sum ts_s, + foldl' max Unlifted ls + ) -- | Check for duplication of names inside a pattern group. Produces -- a description of all names used in the pattern group. -checkForDuplicateNames :: MonadTypeChecker m => - [UncheckedPattern] -> m () +checkForDuplicateNames :: + MonadTypeChecker m => + [UncheckedPattern] -> + m () checkForDuplicateNames = (`evalStateT` mempty) . mapM_ check - where check (Id v _ loc) = seen v loc - check (PatternParens p _) = check p - check Wildcard{} = return () - check (TuplePattern ps _) = mapM_ check ps - check (RecordPattern fs _) = mapM_ (check . snd) fs - check (PatternAscription p _ _) = check p - check PatternLit{} = return () - check (PatternConstr _ _ ps _) = mapM_ check ps + where + check (Id v _ loc) = seen v loc + check (PatternParens p _) = check p + check Wildcard {} = return () + check (TuplePattern ps _) = mapM_ check ps + check (RecordPattern fs _) = mapM_ (check . snd) fs + check (PatternAscription p _ _) = check p + check PatternLit {} = return () + check (PatternConstr _ _ ps _) = mapM_ check ps - seen v loc = do - already <- gets $ M.lookup v - case already of - Just prev_loc -> - lift $ typeError loc mempty $ - "Name" <+> pquote (ppr v) <+> "also bound at" <+> - text (locStr prev_loc) <> "." - Nothing -> - modify $ M.insert v loc + seen v loc = do + already <- gets $ M.lookup v + case already of + Just prev_loc -> + lift $ + typeError loc mempty $ + "Name" <+> pquote (ppr v) <+> "also bound at" + <+> text (locStr prev_loc) <> "." + Nothing -> + modify $ M.insert v loc -- | Check whether the type contains arrow types that define the same -- parameter. These might also exist further down, but that's not -- really a problem - we mostly do this checking to help the user, -- since it is likely an error, but it's easy to assign a semantics to -- it (normal name shadowing). -checkForDuplicateNamesInType :: MonadTypeChecker m => - TypeExp Name -> m () +checkForDuplicateNamesInType :: + MonadTypeChecker m => + TypeExp Name -> + m () checkForDuplicateNamesInType = check mempty - where check seen (TEArrow (Just v) t1 t2 loc) - | Just prev_loc <- M.lookup v seen = - typeError loc mempty $ - text "Name" <+> pquote (ppr v) <+> - "also bound at" <+> text (locStr prev_loc) <> "." - | otherwise = - check seen' t1 >> check seen' t2 - where seen' = M.insert v loc seen - check seen (TEArrow Nothing t1 t2 _) = - check seen t1 >> check seen t2 - check seen (TETuple ts _) = mapM_ (check seen) ts - check seen (TERecord fs _) = mapM_ (check seen . snd) fs - check seen (TEUnique t _) = check seen t - check seen (TESum cs _) = mapM_ (mapM (check seen) . snd) cs - check seen (TEApply t1 (TypeArgExpType t2) _) = - check seen t1 >> check seen t2 - check seen (TEApply t1 TypeArgExpDim{} _) = - check seen t1 - check _ TEArray{} = return () - check _ TEVar{} = return () + where + check seen (TEArrow (Just v) t1 t2 loc) + | Just prev_loc <- M.lookup v seen = + typeError loc mempty $ + text "Name" <+> pquote (ppr v) + <+> "also bound at" + <+> text (locStr prev_loc) <> "." + | otherwise = + check seen' t1 >> check seen' t2 + where + seen' = M.insert v loc seen + check seen (TEArrow Nothing t1 t2 _) = + check seen t1 >> check seen t2 + check seen (TETuple ts _) = mapM_ (check seen) ts + check seen (TERecord fs _) = mapM_ (check seen . snd) fs + check seen (TEUnique t _) = check seen t + check seen (TESum cs _) = mapM_ (mapM (check seen) . snd) cs + check seen (TEApply t1 (TypeArgExpType t2) _) = + check seen t1 >> check seen t2 + check seen (TEApply t1 TypeArgExpDim {} _) = + check seen t1 + check _ TEArray {} = return () + check _ TEVar {} = return () -- | @checkTypeParams ps m@ checks the type parameters @ps@, then -- invokes the continuation @m@ with the checked parameters, while -- extending the monadic name map with @ps@. -checkTypeParams :: MonadTypeChecker m => - [TypeParamBase Name] - -> ([TypeParamBase VName] -> m a) - -> m a +checkTypeParams :: + MonadTypeChecker m => + [TypeParamBase Name] -> + ([TypeParamBase VName] -> m a) -> + m a checkTypeParams ps m = bindSpaced (map typeParamSpace ps) $ - m =<< evalStateT (mapM checkTypeParam ps) mempty - where typeParamSpace (TypeParamDim pv _) = (Term, pv) - typeParamSpace (TypeParamType _ pv _) = (Type, pv) + m =<< evalStateT (mapM checkTypeParam ps) mempty + where + typeParamSpace (TypeParamDim pv _) = (Term, pv) + typeParamSpace (TypeParamType _ pv _) = (Type, pv) - checkParamName ns v loc = do - seen <- gets $ M.lookup (ns,v) - case seen of - Just prev -> - lift $ typeError loc mempty $ - text "Type parameter" <+> pquote (ppr v) <+> - "previously defined at" <+> text (locStr prev) <> "." - Nothing -> do - modify $ M.insert (ns,v) loc - lift $ checkName ns v loc + checkParamName ns v loc = do + seen <- gets $ M.lookup (ns, v) + case seen of + Just prev -> + lift $ + typeError loc mempty $ + text "Type parameter" <+> pquote (ppr v) + <+> "previously defined at" + <+> text (locStr prev) <> "." + Nothing -> do + modify $ M.insert (ns, v) loc + lift $ checkName ns v loc - checkTypeParam (TypeParamDim pv loc) = - TypeParamDim <$> checkParamName Term pv loc <*> pure loc - checkTypeParam (TypeParamType l pv loc) = - TypeParamType l <$> checkParamName Type pv loc <*> pure loc + checkTypeParam (TypeParamDim pv loc) = + TypeParamDim <$> checkParamName Term pv loc <*> pure loc + checkTypeParam (TypeParamType l pv loc) = + TypeParamType l <$> checkParamName Type pv loc <*> pure loc -- | Construct a type argument corresponding to a type parameter. typeParamToArg :: TypeParam -> StructTypeArg @@ -330,9 +384,10 @@ TypeArgType (Scalar $ TypeVar () Nonunique (typeName v) []) ploc -- | A substitution for when using 'substituteTypes'. -data TypeSub = TypeSub TypeBinding - | DimSub (DimDecl VName) - deriving (Show) +data TypeSub + = TypeSub TypeBinding + | DimSub (DimDecl VName) + deriving (Show) -- | A collection of type substitutions. type TypeSubs = M.Map VName TypeSub @@ -341,14 +396,17 @@ substituteTypes :: Monoid als => TypeSubs -> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als substituteTypes substs ot = case ot of Array als u at shape -> - arrayOf (substituteTypes substs (Scalar at) `setAliases` mempty) - (substituteInShape shape) u `addAliases` (<>als) + arrayOf + (substituteTypes substs (Scalar at) `setAliases` mempty) + (substituteInShape shape) + u + `addAliases` (<> als) Scalar (Prim t) -> Scalar $ Prim t Scalar (TypeVar als u v targs) | Just (TypeSub (TypeAbbr _ ps t)) <- M.lookup (qualLeaf (qualNameFromTypeName v)) substs -> - applyType ps (t `setAliases` mempty) (map substituteInTypeArg targs) - `setUniqueness` u `addAliases` (<>als) + applyType ps (t `setAliases` mempty) (map substituteInTypeArg targs) + `setUniqueness` u `addAliases` (<> als) | otherwise -> Scalar $ TypeVar als u v $ map substituteInTypeArg targs Scalar (Record ts) -> Scalar $ Record $ fmap (substituteTypes substs) ts @@ -356,34 +414,40 @@ Scalar $ Arrow als v (substituteTypes substs t1) (substituteTypes substs t2) Scalar (Sum cs) -> Scalar $ Sum $ (fmap . fmap) (substituteTypes substs) cs - where substituteInTypeArg (TypeArgDim d loc) = - TypeArgDim (substituteInDim d) loc - substituteInTypeArg (TypeArgType t loc) = - TypeArgType (substituteTypes substs t) loc + where + substituteInTypeArg (TypeArgDim d loc) = + TypeArgDim (substituteInDim d) loc + substituteInTypeArg (TypeArgType t loc) = + TypeArgType (substituteTypes substs t) loc - substituteInShape (ShapeDecl ds) = - ShapeDecl $ map substituteInDim ds + substituteInShape (ShapeDecl ds) = + ShapeDecl $ map substituteInDim ds - substituteInDim (NamedDim v) - | Just (DimSub d) <- M.lookup (qualLeaf v) substs = d - substituteInDim d = d + substituteInDim (NamedDim v) + | Just (DimSub d) <- M.lookup (qualLeaf v) substs = d + substituteInDim d = d -applyType :: Monoid als => - [TypeParam] -> TypeBase (DimDecl VName) als -> [StructTypeArg] -> TypeBase (DimDecl VName) als +applyType :: + Monoid als => + [TypeParam] -> + TypeBase (DimDecl VName) als -> + [StructTypeArg] -> + TypeBase (DimDecl VName) als applyType ps t args = substituteTypes substs t - where substs = M.fromList $ zipWith mkSubst ps args - -- We are assuming everything has already been type-checked for correctness. - mkSubst (TypeParamDim pv _) (TypeArgDim (NamedDim v) _) = - (pv, DimSub $ NamedDim v) - mkSubst (TypeParamDim pv _) (TypeArgDim (ConstDim x) _) = - (pv, DimSub $ ConstDim x) - mkSubst (TypeParamDim pv _) (TypeArgDim AnyDim _) = - (pv, DimSub AnyDim) - mkSubst (TypeParamType l pv _) (TypeArgType at _) = - (pv, TypeSub $ TypeAbbr l [] at) - mkSubst p a = - error $ "applyType mkSubst: cannot substitute " ++ pretty a ++ " for " ++ pretty p + where + substs = M.fromList $ zipWith mkSubst ps args + -- We are assuming everything has already been type-checked for correctness. + mkSubst (TypeParamDim pv _) (TypeArgDim (NamedDim v) _) = + (pv, DimSub $ NamedDim v) + mkSubst (TypeParamDim pv _) (TypeArgDim (ConstDim x) _) = + (pv, DimSub $ ConstDim x) + mkSubst (TypeParamDim pv _) (TypeArgDim AnyDim _) = + (pv, DimSub AnyDim) + mkSubst (TypeParamType l pv _) (TypeArgType at _) = + (pv, TypeSub $ TypeAbbr l [] at) + mkSubst p a = + error $ "applyType mkSubst: cannot substitute " ++ pretty a ++ " for " ++ pretty p -- | A type substituion may be a substitution or a yet-unknown -- substitution (but which is certainly an overloaded primitive @@ -406,7 +470,7 @@ applySubst = substTypesAny instance Substitutable (TypeBase (DimDecl VName) Aliasing) where - applySubst = substTypesAny . (fmap (fmap fromStruct).) + applySubst = substTypesAny . (fmap (fmap fromStruct) .) instance Substitutable (DimDecl VName) where applySubst f (NamedDim (QualName _ v)) @@ -418,28 +482,36 @@ instance Substitutable Pattern where applySubst f = runIdentity . astMap mapper - where mapper = ASTMapper { mapOnExp = return - , mapOnName = return - , mapOnQualName = return - , mapOnStructType = return . applySubst f - , mapOnPatternType = return . applySubst f - } + where + mapper = + ASTMapper + { mapOnExp = return, + mapOnName = return, + mapOnQualName = return, + mapOnStructType = return . applySubst f, + mapOnPatternType = return . applySubst f + } -- | Perform substitutions, from type names to types, on a type. Works -- regardless of what shape and uniqueness information is attached to the type. -substTypesAny :: Monoid as => - (VName -> Maybe (Subst (TypeBase (DimDecl VName) as))) - -> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as +substTypesAny :: + Monoid as => + (VName -> Maybe (Subst (TypeBase (DimDecl VName) as))) -> + TypeBase (DimDecl VName) as -> + TypeBase (DimDecl VName) as substTypesAny lookupSubst ot = case ot of Array als u et shape -> - arrayOf (substTypesAny lookupSubst' (Scalar et)) - (applySubst lookupSubst' shape) u `setAliases` als + arrayOf + (substTypesAny lookupSubst' (Scalar et)) + (applySubst lookupSubst' shape) + u + `setAliases` als Scalar (Prim t) -> Scalar $ Prim t -- We only substitute for a type variable with no arguments, since -- type parameters cannot have higher kind. Scalar (TypeVar als u v targs) -> case lookupSubst $ qualLeaf (qualNameFromTypeName v) of - Just (Subst t) -> substTypesAny lookupSubst $ t `setUniqueness` u `addAliases` (<>als) + Just (Subst t) -> substTypesAny lookupSubst $ t `setUniqueness` u `addAliases` (<> als) Just PrimSubst -> Scalar $ TypeVar mempty u v $ map subsTypeArg targs _ -> Scalar $ TypeVar als u v $ map subsTypeArg targs Scalar (Record ts) -> Scalar $ Record $ fmap (substTypesAny lookupSubst) ts @@ -447,10 +519,10 @@ Scalar $ Arrow als v (substTypesAny lookupSubst t1) (substTypesAny lookupSubst t2) Scalar (Sum ts) -> Scalar $ Sum $ fmap (fmap $ substTypesAny lookupSubst) ts - - where subsTypeArg (TypeArgType t loc) = - TypeArgType (substTypesAny lookupSubst' t) loc - subsTypeArg (TypeArgDim v loc) = - TypeArgDim (applySubst lookupSubst' v) loc + where + subsTypeArg (TypeArgType t loc) = + TypeArgType (substTypesAny lookupSubst' t) loc + subsTypeArg (TypeArgDim v loc) = + TypeArgDim (applySubst lookupSubst' v) loc - lookupSubst' = fmap (fmap $ second (const ())) . lookupSubst + lookupSubst' = fmap (fmap $ second (const ())) . lookupSubst diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs --- a/src/Language/Futhark/TypeChecker/Unify.hs +++ b/src/Language/Futhark/TypeChecker/Unify.hs @@ -1,980 +1,1146 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE Trustworthy #-} --- | Implementation of unification and other core type system building --- blocks. -module Language.Futhark.TypeChecker.Unify - ( Constraint(..) - , Usage - , mkUsage - , mkUsage' - , Level - , Constraints - , MonadUnify(..) - , Rigidity(..) - , RigidSource(..) - , BreadCrumbs - , noBreadCrumbs - , hasNoBreadCrumbs - , dimNotes - , mkTypeVarName - - , zeroOrderType - , mustHaveConstr - , mustHaveField - , mustBeOneOf - , equalityType - , normType - , normPatternType - , normTypeFully - , instantiateEmptyArrayDims - - , unify - , expect - , unifyMostCommon - , anyDimOnMismatch - , doUnification - ) -where - -import Control.Monad.Except -import Control.Monad.Writer hiding (Sum) -import Control.Monad.RWS.Strict hiding (Sum) -import Control.Monad.State -import Data.Bifoldable (biany) -import Data.List (intersect) -import Data.Maybe -import qualified Data.Map.Strict as M -import qualified Data.Set as S - -import Language.Futhark hiding (unifyDims) -import Language.Futhark.TypeChecker.Monad hiding (BoundV) -import Language.Futhark.TypeChecker.Types -import Futhark.Util.Pretty hiding (empty) - --- | A piece of information that describes what process the type --- checker currently performing. This is used to give better error --- messages for unification errors. -data BreadCrumb = MatchingTypes StructType StructType - | MatchingFields [Name] - | MatchingConstructor Name - | Matching Doc - -instance Pretty BreadCrumb where - ppr (MatchingTypes t1 t2) = - "When matching type" indent 2 (ppr t1) - "with" indent 2 (ppr t2) - ppr (MatchingFields fields) = - "When matching types of record field" <+> - pquote (mconcat $ punctuate "." $ map ppr fields) <> dot - ppr (MatchingConstructor c) = - "When matching types of constructor" <+> pquote (ppr c) <> dot - ppr (Matching s) = - s - --- | Unification failures can occur deep down inside complicated types --- (consider nested records). We leave breadcrumbs behind us so we --- can report the path we took to find the mismatch. -newtype BreadCrumbs = BreadCrumbs [BreadCrumb] - --- | An empty path. -noBreadCrumbs :: BreadCrumbs -noBreadCrumbs = BreadCrumbs [] - --- | Is the path empty? -hasNoBreadCrumbs :: BreadCrumbs -> Bool -hasNoBreadCrumbs (BreadCrumbs xs) = null xs - --- | Drop a breadcrumb on the path behind you. -breadCrumb :: BreadCrumb -> BreadCrumbs -> BreadCrumbs -breadCrumb (MatchingFields xs) (BreadCrumbs (MatchingFields ys : bcs)) = - BreadCrumbs $ MatchingFields (ys++xs) : bcs -breadCrumb bc (BreadCrumbs bcs) = - BreadCrumbs $ bc : bcs - -instance Pretty BreadCrumbs where - ppr (BreadCrumbs []) = mempty - ppr (BreadCrumbs bcs) = line <> stack (map ppr bcs) - --- | A usage that caused a type constraint. -data Usage = Usage (Maybe String) SrcLoc - deriving (Show) - --- | Construct a 'Usage' from a location and a description. -mkUsage :: SrcLoc -> String -> Usage -mkUsage = flip (Usage . Just) - --- | Construct a 'Usage' that has just a location, but no particular --- description. -mkUsage' :: SrcLoc -> Usage -mkUsage' = Usage Nothing - -instance Pretty Usage where - ppr (Usage Nothing loc) = "use at " <> textwrap (locStr loc) - ppr (Usage (Just s) loc) = textwrap s <+/> "at" <+> textwrap (locStr loc) - -instance Located Usage where - locOf (Usage _ loc) = locOf loc - --- | The level at which a type variable is bound. Higher means --- deeper. We can only unify a type variable at level @i@ with a type --- @t@ if all type names that occur in @t@ are at most at level @i@. -type Level = Int - --- | A constraint on a yet-ambiguous type variable. -data Constraint = NoConstraint Liftedness Usage - | ParamType Liftedness SrcLoc - | Constraint StructType Usage - | Overloaded [PrimType] Usage - | HasFields (M.Map Name StructType) Usage - | Equality Usage - | HasConstrs (M.Map Name [StructType]) Usage - | ParamSize SrcLoc - | Size (Maybe (DimDecl VName)) Usage - -- ^ Is not actually a type, but a term-level size, - -- possibly already set to something specific. - | UnknowableSize SrcLoc RigidSource - -- ^ A size that does not unify with anything - - -- created from the result of applying a function - -- whose return size is existential, or otherwise - -- hiding a size. - deriving Show - -instance Located Constraint where - locOf (NoConstraint _ usage) = locOf usage - locOf (ParamType _ usage) = locOf usage - locOf (Constraint _ usage) = locOf usage - locOf (Overloaded _ usage) = locOf usage - locOf (HasFields _ usage) = locOf usage - locOf (Equality usage) = locOf usage - locOf (HasConstrs _ usage) = locOf usage - locOf (ParamSize loc) = locOf loc - locOf (Size _ usage) = locOf usage - locOf (UnknowableSize loc _) = locOf loc - --- | Mapping from fresh type variables, instantiated from the type --- schemes of polymorphic functions, to (possibly) specific types as --- determined on application and the location of that application, or --- a partial constraint on their type. -type Constraints = M.Map VName (Level, Constraint) - -lookupSubst :: VName -> Constraints -> Maybe (Subst StructType) -lookupSubst v constraints = case snd <$> M.lookup v constraints of - Just (Constraint t _) -> Just $ Subst t - Just Overloaded{} -> Just PrimSubst - Just (Size (Just d) _) -> - Just $ SizeSubst $ applySubst (`lookupSubst` constraints) d - _ -> Nothing - --- | The source of a rigid size. -data RigidSource - = RigidArg (Maybe (QualName VName)) String - -- ^ A function argument that is not a constant or variable name. - | RigidRet (Maybe (QualName VName)) - -- ^ An existential return size. - | RigidLoop - | RigidSlice (Maybe (DimDecl VName)) String - -- ^ Produced by a complicated slice expression. - | RigidRange - -- ^ Produced by a complicated range expression. - | RigidBound String - -- ^ Produced by a range expression with this bound. - | RigidCond StructType StructType - -- ^ Mismatch in branches. - | RigidUnify - -- ^ Invented during unification. - | RigidOutOfScope SrcLoc VName - deriving (Eq, Ord, Show) - --- | The ridigity of a size variable. All rigid sizes are tagged with --- information about how they were generated. -data Rigidity = Rigid RigidSource | Nonrigid - deriving (Eq, Ord, Show) - -prettySource :: SrcLoc -> SrcLoc -> RigidSource -> Doc - -prettySource ctx loc (RigidRet Nothing) = - "is unknown size returned by function at" <+> - text (locStrRel ctx loc) <> "." - -prettySource ctx loc (RigidRet (Just fname)) = - "is unknown size returned by" <+> pquote (ppr fname) <+> - "at" <+> text (locStrRel ctx loc) <> "." - -prettySource ctx loc (RigidArg fname arg) = - "is value of argument" - indent 2 (shorten arg) - "passed to" <+> fname' <+> "at" <+> text (locStrRel ctx loc) <> "." - where fname' = maybe "function" (pquote . ppr) fname - -prettySource ctx loc (RigidSlice d slice) = - "is size produced by slice" - indent 2 (shorten slice) - d_desc <> "at" <+> text (locStrRel ctx loc) <> "." - where d_desc = case d of - Just d' -> "of dimension of size " <> pquote (ppr d') <> " " - Nothing -> mempty - -prettySource ctx loc RigidLoop = - "is unknown size of value returned at" <+> text (locStrRel ctx loc) <> "." - -prettySource ctx loc RigidRange = - "is unknown length of range at" <+> text (locStrRel ctx loc) <> "." - -prettySource ctx loc (RigidBound bound) = - "generated from expression" - indent 2 (shorten bound) - "used in range at " <> text (locStrRel ctx loc) <> "." - -prettySource ctx loc (RigidOutOfScope boundloc v) = - "is an unknown size arising from " <> pquote (pprName v) <> - " going out of scope at " <> text (locStrRel ctx loc) <> "." - "Originally bound at " <> text (locStrRel ctx boundloc) <> "." - -prettySource _ _ RigidUnify = - "is an artificial size invented during unification of functions with anonymous sizes" - -prettySource ctx loc (RigidCond t1 t2) = - "is unknown due to conditional expression at " <> - text (locStrRel ctx loc) <> "." - "One branch returns array of type: " <> align (ppr t1) - "The other an array of type: " <> align (ppr t2) - --- | Retrieve notes describing the purpose or origin of the given --- 'DimDecl'. The location is used as the *current* location, for the --- purpose of reporting relative locations. -dimNotes :: (Located a, MonadUnify m) => a -> DimDecl VName -> m Notes -dimNotes ctx (NamedDim d) = do - c <- M.lookup (qualLeaf d) <$> getConstraints - case c of - Just (_, UnknowableSize loc rsrc) -> - return $ aNote $ pretty $ - pquote (ppr d) <+> prettySource (srclocOf ctx) loc rsrc - _ -> return mempty -dimNotes _ _ = return mempty - -typeNotes :: (Located a, MonadUnify m) => a -> StructType -> m Notes -typeNotes ctx = - fmap mconcat . mapM (dimNotes ctx . NamedDim . qualName) . - S.toList . typeDimNames - --- | Monads that which to perform unification must implement this type --- class. -class Monad m => MonadUnify m where - getConstraints :: m Constraints - putConstraints :: Constraints -> m () - modifyConstraints :: (Constraints -> Constraints) -> m () - modifyConstraints f = do - x <- getConstraints - putConstraints $ f x - - newTypeVar :: Monoid als => SrcLoc -> String -> m (TypeBase dim als) - newDimVar :: SrcLoc -> Rigidity -> String -> m VName - - curLevel :: m Level - - matchError :: Located loc => loc -> Notes -> BreadCrumbs - -> StructType -> StructType -> m a - - unifyError :: Located loc => loc -> Notes -> BreadCrumbs - -> Doc -> m a - --- | Replace all type variables with their substitution. -normTypeFully :: (Substitutable a, MonadUnify m) => a -> m a -normTypeFully t = do constraints <- getConstraints - return $ applySubst (`lookupSubst` constraints) t - --- | Replace any top-level type variable with its substitution. -normType :: MonadUnify m => StructType -> m StructType -normType t@(Scalar (TypeVar _ _ (TypeName [] v) [])) = do - constraints <- getConstraints - case snd <$> M.lookup v constraints of - Just (Constraint t' _) -> normType t' - _ -> return t -normType t = return t - --- | Replace any top-level type variable with its substitution. -normPatternType :: MonadUnify m => PatternType -> m PatternType -normPatternType t@(Scalar (TypeVar als u (TypeName [] v) [])) = do - constraints <- getConstraints - case snd <$> M.lookup v constraints of - Just (Constraint t' _) -> - normPatternType $ t' `setUniqueness` u `setAliases` als - _ -> return t -normPatternType t = return t - -rigidConstraint :: Constraint -> Bool -rigidConstraint ParamType{} = True -rigidConstraint ParamSize{} = True -rigidConstraint UnknowableSize{} = True -rigidConstraint _ = False - --- | Replace 'AnyDim' dimensions that occur as 'PosImmediate' or --- 'PosParam' with a fresh 'NamedDim'. -instantiateEmptyArrayDims :: MonadUnify m => - SrcLoc -> String -> Rigidity - -> TypeBase (DimDecl VName) als - -> m (TypeBase (DimDecl VName) als, [VName]) -instantiateEmptyArrayDims tloc desc r = runWriterT . traverseDims onDim - where onDim _ PosImmediate AnyDim = inst - onDim _ PosParam AnyDim = inst - onDim _ _ d = return d - inst = do - dim <- lift $ newDimVar tloc r desc - tell [dim] - return $ NamedDim $ qualName dim - --- | Is the given type variable the name of an abstract type or type --- parameter, which we cannot substitute? -isRigid :: VName -> Constraints -> Bool -isRigid v constraints = - maybe True (rigidConstraint . snd) $ M.lookup v constraints - --- | If the given type variable is nonrigid, what is its level? -isNonRigid :: VName -> Constraints -> Maybe Level -isNonRigid v constraints = do - (lvl, c) <- M.lookup v constraints - guard $ not $ rigidConstraint c - return lvl - -type UnifyDims m = - BreadCrumbs -> [VName] -> (VName -> Maybe Int) -> DimDecl VName -> DimDecl VName -> m () - -flipUnifyDims :: UnifyDims m -> UnifyDims m -flipUnifyDims onDims bcs bound nonrigid t1 t2 = - onDims bcs bound nonrigid t2 t1 - -unifyWith :: MonadUnify m => - UnifyDims m -> Usage -> BreadCrumbs - -> StructType -> StructType -> m () -unifyWith onDims usage = subunify False mempty - where - swap True x y = (y, x) - swap False x y = (x, y) - - subunify ord bound bcs t1 t2 = do - constraints <- getConstraints - - t1' <- normType t1 - t2' <- normType t2 - - let nonrigid v = isNonRigid v constraints - - failure = matchError (srclocOf usage) mempty bcs t1' t2' - - -- Remove any of the intermediate dimensions we added just - -- for unification purposes. - unbound = applySubst f - where f d | d `elem` bound = Just $ SizeSubst AnyDim - | otherwise = Nothing - - link ord' v lvl = - linkVarToType linkDims usage bcs v lvl . unbound - where -- We may have to flip the order of future calls to - -- onDims inside linkVarToType. - linkDims | ord' = flipUnifyDims onDims - | otherwise = onDims - - unifyTypeArg bcs' (TypeArgDim d1 _) (TypeArgDim d2 _) = - onDims' bcs' (swap ord d1 d2) - unifyTypeArg bcs' (TypeArgType t _) (TypeArgType arg_t _) = - subunify ord bound bcs' t arg_t - unifyTypeArg bcs' _ _ = unifyError usage mempty bcs' - "Cannot unify a type argument with a dimension argument (or vice versa)." - - onDims' bcs' (d1, d2) = - onDims bcs' bound nonrigid - (applySubst (`lookupSubst` constraints) d1) - (applySubst (`lookupSubst` constraints) d2) - - case (t1', t2') of - (Scalar (Record fs), - Scalar (Record arg_fs)) - | M.keys fs == M.keys arg_fs -> - forM_ (M.toList $ M.intersectionWith (,) fs arg_fs) $ \(k, (k_t1, k_t2)) -> do - let bcs' = breadCrumb (MatchingFields [k]) bcs - subunify ord bound bcs' k_t1 k_t2 - | otherwise -> do - let missing = filter (`notElem` M.keys arg_fs) (M.keys fs) ++ - filter (`notElem` M.keys fs) (M.keys arg_fs) - unifyError usage mempty bcs $ - "Unshared fields:" <+> commasep (map ppr missing) <> "." - - (Scalar (TypeVar _ _ (TypeName _ tn) targs), - Scalar (TypeVar _ _ (TypeName _ arg_tn) arg_targs)) - | tn == arg_tn, length targs == length arg_targs -> do - let bcs' = breadCrumb (Matching "When matching type arguments.") bcs - zipWithM_ (unifyTypeArg bcs') targs arg_targs - - (Scalar (TypeVar _ _ (TypeName [] v1) []), - Scalar (TypeVar _ _ (TypeName [] v2) [])) -> - case (nonrigid v1, nonrigid v2) of - (Nothing, Nothing) -> failure - (Just lvl1, Nothing) -> link ord v1 lvl1 t2' - (Nothing, Just lvl2) -> link (not ord) v2 lvl2 t1' - (Just lvl1, Just lvl2) - | lvl1 <= lvl2 -> link ord v1 lvl1 t2' - | otherwise -> link (not ord) v2 lvl2 t1' - - (Scalar (TypeVar _ _ (TypeName [] v1) []), _) - | Just lvl <- nonrigid v1 -> - link ord v1 lvl t2' - (_, Scalar (TypeVar _ _ (TypeName [] v2) [])) - | Just lvl <- nonrigid v2 -> - link (not ord) v2 lvl t1' - - (Scalar (Arrow _ p1 a1 b1), - Scalar (Arrow _ p2 a2 b2)) -> do - let (r1, r2) = swap ord (Rigid RigidUnify) Nonrigid - (a1', a1_dims) <- instantiateEmptyArrayDims (srclocOf usage) "anonymous" r1 a1 - (a2', a2_dims) <- instantiateEmptyArrayDims (srclocOf usage) "anonymous" r2 a2 - let bound' = bound <> mapMaybe pname [p1, p2] <> a1_dims <> a2_dims - subunify (not ord) bound - (breadCrumb (Matching "When matching parameter types.") bcs) - a1' a2' - subunify ord bound' - (breadCrumb (Matching "When matching return types.") bcs) - b1' b2' - where (b1', b2') = - -- Replace one parameter name with the other in the - -- return type, in case of dependent types. I.e., - -- we want type '(n: i32) -> [n]i32' to unify with - -- type '(x: i32) -> [x]i32'. - case (p1, p2) of - (Named p1', Named p2') -> - let f v | v == p2' = Just $ SizeSubst $ NamedDim $ qualName p1' - | otherwise = Nothing - in (b1, applySubst f b2) - - (_, _) -> - (b1, b2) - - pname (Named x) = Just x - pname Unnamed = Nothing - - (Array{}, Array{}) - | ShapeDecl (t1_d : _) <- arrayShape t1', - ShapeDecl (t2_d : _) <- arrayShape t2', - Just t1'' <- peelArray 1 t1', - Just t2'' <- peelArray 1 t2' -> do - onDims' bcs (swap ord t1_d t2_d) - subunify ord bound bcs t1'' t2'' - - (Scalar (Sum cs), - Scalar (Sum arg_cs)) - | M.keys cs == M.keys arg_cs -> - unifySharedConstructors onDims usage bcs - (map unbound <$> cs) (map unbound <$> arg_cs) - | otherwise -> do - let missing = filter (`notElem` M.keys arg_cs) (M.keys cs) ++ - filter (`notElem` M.keys cs) (M.keys arg_cs) - unifyError usage mempty bcs $ - "Unshared constructors:" <+> commasep (map (("#"<>) . ppr) missing) <> "." - - _ | t1' == t2' -> return () - | otherwise -> failure - -unifyDims :: MonadUnify m => Usage -> UnifyDims m -unifyDims _ _ _ _ d1 d2 - | d1 == d2 = return () -unifyDims usage bcs _ nonrigid (NamedDim (QualName _ d1)) d2 - | Just lvl1 <- nonrigid d1 = - linkVarToDim usage bcs d1 lvl1 d2 -unifyDims usage bcs _ nonrigid d1 (NamedDim (QualName _ d2)) - | Just lvl2 <- nonrigid d2 = - linkVarToDim usage bcs d2 lvl2 d1 -unifyDims usage bcs _ _ d1 d2 = do - notes <- (<>) <$> dimNotes usage d1 <*> dimNotes usage d2 - unifyError usage notes bcs $ - "Dimensions" <+> pquote (ppr d1) <+> - "and" <+> pquote (ppr d2) <+> "do not match." - --- | Unifies two types. -unify :: MonadUnify m => Usage -> StructType -> StructType -> m () -unify usage = unifyWith (unifyDims usage) usage noBreadCrumbs - --- | @expect super sub@ checks that @sub@ is a subtype of @super@. -expect :: MonadUnify m => Usage -> StructType -> StructType -> m () -expect usage = unifyWith onDims usage noBreadCrumbs - where onDims _ _ _ AnyDim _ = return () - onDims _ _ _ d1 d2 - | d1 == d2 = return () - onDims bcs bound nonrigid (NamedDim (QualName _ d1)) d2 - | Just lvl1 <- nonrigid d1, d2 /= AnyDim, not $ boundParam bound d2 = - linkVarToDim usage bcs d1 lvl1 d2 - onDims bcs bound nonrigid d1 (NamedDim (QualName _ d2)) - | Just lvl2 <- nonrigid d2, not $ boundParam bound d1 = - linkVarToDim usage bcs d2 lvl2 d1 - onDims bcs _ _ d1 d2 = do - notes <- (<>) <$> dimNotes usage d1 <*> dimNotes usage d2 - unifyError usage notes bcs $ "Dimensions" <+> pquote (ppr d1) <+> - "and" <+> pquote (ppr d2) <+> "do not match." - - boundParam bound (NamedDim (QualName _ d)) = d `elem` bound - boundParam _ _ = False - -hasEmptyDims :: StructType -> Bool -hasEmptyDims = biany empty (const False) - where empty AnyDim = True - empty _ = False - -occursCheck :: MonadUnify m => - Usage -> BreadCrumbs - -> VName -> StructType -> m () -occursCheck usage bcs vn tp = - when (vn `S.member` typeVars tp) $ - unifyError usage mempty bcs $ "Occurs check: cannot instantiate" <+> - pprName vn <+> "with" <+> ppr tp <> "." - -scopeCheck :: MonadUnify m => - Usage -> BreadCrumbs - -> VName -> Level -> StructType -> m () -scopeCheck usage bcs vn max_lvl tp = do - constraints <- getConstraints - checkType constraints tp - where checkType constraints t = - mapM_ (check constraints) $ typeVars t <> typeDimNames t - - check constraints v - | Just (lvl, c) <- M.lookup v constraints, - lvl > max_lvl = - if rigidConstraint c - then scopeViolation v - else modifyConstraints $ M.insert v (max_lvl, c) - - | otherwise = - return () - - scopeViolation v = do - notes <- typeNotes usage tp - unifyError usage notes bcs $ "Cannot unify type" - indent 2 (ppr tp) - "with" <+> pquote (pprName vn) <+> "(scope violation)." - "This is because" <+> pquote (pprName v) <+> - "is rigidly bound in a deeper scope." - -linkVarToType :: MonadUnify m => - UnifyDims m -> Usage -> BreadCrumbs - -> VName -> Level -> StructType -> m () -linkVarToType onDims usage bcs vn lvl tp = do - occursCheck usage bcs vn tp - scopeCheck usage bcs vn lvl tp - - constraints <- getConstraints - let tp' = removeUniqueness tp - modifyConstraints $ M.insert vn (lvl, Constraint tp' usage) - case snd <$> M.lookup vn constraints of - - Just (NoConstraint Unlifted unlift_usage) -> do - let bcs' = breadCrumb - (Matching $ "When verifying that" <+> pquote (pprName vn) <+> - textwrap "is not instantiated with a function type, due to" <+> - ppr unlift_usage) - bcs - zeroOrderTypeWith usage bcs' tp' - - when (hasEmptyDims tp') $ - unifyError usage mempty bcs $ "Type variable" <+> pprName vn <+> - "cannot be instantiated with type containing anonymous sizes:" - indent 2 (ppr tp) - textwrap "This is usually because the size of an array returned by a higher-order function argument cannot be determined statically. This can also be due to the return size being a value parameter. Add type annotation to clarify." - - Just (Equality _) -> - equalityType usage tp' - - Just (Overloaded ts old_usage) - | tp `notElem` map (Scalar . Prim) ts -> - case tp' of - Scalar (TypeVar _ _ (TypeName [] v) []) - | not $ isRigid v constraints -> - linkVarToTypes usage v ts - _ -> - unifyError usage mempty bcs $ "Cannot instantiate" <+> pquote (pprName vn) <+> - "with type" indent 2 (ppr tp) "as" <+> - pquote (pprName vn) <+> "must be one of" <+> - commasep (map ppr ts) <+/> - "due to" <+/> ppr old_usage <> "." - - Just (HasFields required_fields old_usage) -> - case tp of - Scalar (Record tp_fields) - | all (`M.member` tp_fields) $ M.keys required_fields -> do - required_fields' <- mapM normTypeFully required_fields - let bcs' = - breadCrumb - (Matching $ pprName vn <+> - "must be a record with at least the fields:" - indent 2 (ppr (Record required_fields')) - "due to" <+> ppr old_usage <> ".") - bcs - mapM_ (uncurry $ unifyWith onDims usage bcs') $ M.elems $ - M.intersectionWith (,) required_fields tp_fields - Scalar (TypeVar _ _ (TypeName [] v) []) - | not $ isRigid v constraints -> - modifyConstraints $ M.insert v - (lvl, HasFields required_fields old_usage) - _ -> - unifyError usage mempty bcs $ - "Cannot instantiate" <+> pquote (pprName vn) <+> "with type" - indent 2 (ppr tp) - "as" <+> pquote (pprName vn) <+> "must be a record with fields" - indent 2 (ppr (Record required_fields)) - "due to" <+> ppr old_usage <> "." - - Just (HasConstrs required_cs old_usage) -> - case tp of - Scalar (Sum ts) - | all (`M.member` ts) $ M.keys required_cs -> - unifySharedConstructors onDims usage bcs required_cs ts - Scalar (TypeVar _ _ (TypeName [] v) []) - | not $ isRigid v constraints -> do - case M.lookup v constraints of - Just (_, HasConstrs v_cs _) -> - unifySharedConstructors onDims usage bcs required_cs v_cs - _ -> return () - modifyConstraints $ M.insertWith combineConstrs v - (lvl, HasConstrs required_cs old_usage) - where combineConstrs (_, HasConstrs cs1 usage1) (_, HasConstrs cs2 _) = - (lvl, HasConstrs (M.union cs1 cs2) usage1) - combineConstrs hasCs _ = hasCs - _ -> noSumType - - _ -> return () - - where noSumType = unifyError usage mempty bcs - "Cannot unify a sum type with a non-sum type" - -linkVarToDim :: MonadUnify m => - Usage -> BreadCrumbs - -> VName -> Level -> DimDecl VName -> m () -linkVarToDim usage bcs vn lvl dim = do - constraints <- getConstraints - - case dim of - NamedDim dim' - | Just (dim_lvl, c) <- qualLeaf dim' `M.lookup` constraints, - dim_lvl > lvl -> - case c of - ParamSize{} -> do - notes <- dimNotes usage dim - unifyError usage notes bcs $ - "Cannot unify size variable" <+> pquote (ppr dim') <+> - "with" <+> pquote (pprName vn) <+> "(scope violation)." - "This is because" <+> pquote (ppr dim') <+> - "is rigidly bound in a deeper scope." - _ -> modifyConstraints $ M.insert (qualLeaf dim') (lvl, c) - _ -> return () - - modifyConstraints $ M.insert vn (lvl, Size (Just dim) usage) - -removeUniqueness :: TypeBase dim as -> TypeBase dim as -removeUniqueness (Scalar (Record ets)) = - Scalar $ Record $ fmap removeUniqueness ets -removeUniqueness (Scalar (Arrow als p t1 t2)) = - Scalar $ Arrow als p (removeUniqueness t1) (removeUniqueness t2) -removeUniqueness (Scalar (Sum cs)) = - Scalar $ Sum $ (fmap . fmap) removeUniqueness cs -removeUniqueness t = t `setUniqueness` Nonunique - --- | Assert that this type must be one of the given primitive types. -mustBeOneOf :: MonadUnify m => [PrimType] -> Usage -> StructType -> m () -mustBeOneOf [req_t] usage t = unify usage (Scalar (Prim req_t)) t -mustBeOneOf ts usage t = do - t' <- normType t - constraints <- getConstraints - let isRigid' v = isRigid v constraints - - case t' of - Scalar (TypeVar _ _ (TypeName [] v) []) - | not $ isRigid' v -> linkVarToTypes usage v ts - - Scalar (Prim pt) | pt `elem` ts -> return () - - _ -> failure - - where failure = unifyError usage mempty noBreadCrumbs $ - text "Cannot unify type" <+> pquote (ppr t) <+> - "with any of " <> commasep (map ppr ts) <> "." - -linkVarToTypes :: MonadUnify m => Usage -> VName -> [PrimType] -> m () -linkVarToTypes usage vn ts = do - vn_constraint <- M.lookup vn <$> getConstraints - case vn_constraint of - Just (lvl, Overloaded vn_ts vn_usage) -> - case ts `intersect` vn_ts of - [] -> unifyError usage mempty noBreadCrumbs $ - "Type constrained to one of" <+> - commasep (map ppr ts) <+> "but also one of" <+> - commasep (map ppr vn_ts) <+> "due to" <+> ppr vn_usage <> "." - ts' -> modifyConstraints $ M.insert vn (lvl, Overloaded ts' usage) - - Just (_, HasConstrs _ vn_usage) -> - unifyError usage mempty noBreadCrumbs $ - "Type constrained to one of" <+> commasep (map ppr ts) <> - ", but also inferred to be sum type due to" <+> ppr vn_usage <> "." - - Just (_, HasFields _ vn_usage) -> - unifyError usage mempty noBreadCrumbs $ - "Type constrained to one of" <+> commasep (map ppr ts) <> - ", but also inferred to be record due to" <+> ppr vn_usage <> "." - - Just (lvl, _) -> modifyConstraints $ M.insert vn (lvl, Overloaded ts usage) - - Nothing -> - unifyError usage mempty noBreadCrumbs $ - "Cannot constrain type to one of" <+> commasep (map ppr ts) - --- | Assert that this type must support equality. -equalityType :: (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => - Usage -> TypeBase dim as -> m () -equalityType usage t = do - unless (orderZero t) $ - unifyError usage mempty noBreadCrumbs $ - "Type " <+> pquote (ppr t) <+> "does not support equality (is higher-order)." - mapM_ mustBeEquality $ typeVars t - where mustBeEquality vn = do - constraints <- getConstraints - case M.lookup vn constraints of - Just (_, Constraint (Scalar (TypeVar _ _ (TypeName [] vn') [])) _) -> - mustBeEquality vn' - Just (_, Constraint vn_t cusage) - | not $ orderZero vn_t -> - unifyError usage mempty noBreadCrumbs $ - "Type" <+> pquote (ppr t) <+> "does not support equality." - "Constrained to be higher-order due to" <+> ppr cusage <+> "." - | otherwise -> return () - Just (lvl, NoConstraint _ _) -> - modifyConstraints $ M.insert vn (lvl, Equality usage) - Just (_, Overloaded _ _) -> - return () -- All primtypes support equality. - Just (_, Equality{}) -> - return () - Just (_, HasConstrs cs _) -> - mapM_ (equalityType usage) $ concat $ M.elems cs - _ -> - unifyError usage mempty noBreadCrumbs $ - "Type" <+> pprName vn <+> "does not support equality." - -zeroOrderTypeWith :: (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => - Usage -> BreadCrumbs -> TypeBase dim as -> m () -zeroOrderTypeWith usage bcs t = do - unless (orderZero t) $ - unifyError usage mempty bcs $ - "Type" indent 2 (ppr t) "found to be functional." - mapM_ mustBeZeroOrder . S.toList . typeVars $ t - where mustBeZeroOrder vn = do - constraints <- getConstraints - case M.lookup vn constraints of - Just (lvl, NoConstraint _ _) -> - modifyConstraints $ M.insert vn (lvl, NoConstraint Unlifted usage) - Just (_, ParamType Lifted ploc) -> - unifyError usage mempty bcs $ "Type parameter" <+> - pquote (pprName vn) <+> "at" <+> - text (locStr ploc) <+> "may be a function." - _ -> return () - --- | Assert that this type must be zero-order. -zeroOrderType :: (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => - Usage -> String -> TypeBase dim as -> m () -zeroOrderType usage desc = - zeroOrderTypeWith usage $ breadCrumb bc noBreadCrumbs - where bc = Matching $ "When checking" <+> textwrap desc - -unifySharedConstructors :: MonadUnify m => - UnifyDims m -> Usage -> BreadCrumbs - -> M.Map Name [StructType] - -> M.Map Name [StructType] - -> m () -unifySharedConstructors onDims usage bcs cs1 cs2 = - forM_ (M.toList $ M.intersectionWith (,) cs1 cs2) $ \(c, (f1, f2)) -> - unifyConstructor c f1 f2 - where unifyConstructor c f1 f2 - | length f1 == length f2 = do - let bcs' = breadCrumb (MatchingConstructor c) bcs - zipWithM_ (unifyWith onDims usage bcs') f1 f2 - | otherwise = - unifyError usage mempty bcs $ - "Cannot unify constructor" <+> pquote (pprName c) <> "." - --- | In @mustHaveConstr usage c t fs@, the type @t@ must have a --- constructor named @c@ that takes arguments of types @ts@. -mustHaveConstr :: MonadUnify m => - Usage -> Name -> StructType -> [StructType] -> m () -mustHaveConstr usage c t fs = do - constraints <- getConstraints - case t of - Scalar (TypeVar _ _ (TypeName _ tn) []) - | Just (lvl, NoConstraint{}) <- M.lookup tn constraints -> do - mapM_ (scopeCheck usage noBreadCrumbs tn lvl) fs - modifyConstraints $ M.insert tn (lvl, HasConstrs (M.singleton c fs) usage) - | Just (lvl, HasConstrs cs _) <- M.lookup tn constraints -> - case M.lookup c cs of - Nothing -> modifyConstraints $ M.insert tn (lvl, HasConstrs (M.insert c fs cs) usage) - Just fs' - | length fs == length fs' -> zipWithM_ (unify usage) fs fs' - | otherwise -> - unifyError usage mempty noBreadCrumbs $ - "Different arity for constructor" <+> pquote (ppr c) <> "." - - Scalar (Sum cs) -> - case M.lookup c cs of - Nothing -> - unifyError usage mempty noBreadCrumbs $ - "Constuctor" <+> pquote (ppr c) <+> "not present in type." - Just fs' - | length fs == length fs' -> zipWithM_ (unify usage) fs fs' - | otherwise -> - unifyError usage mempty noBreadCrumbs $ - "Different arity for constructor" <+> pquote (ppr c) <+> "." - - _ -> do unify usage t $ Scalar $ Sum $ M.singleton c fs - return () - -mustHaveFieldWith :: MonadUnify m => - UnifyDims m -> Usage -> BreadCrumbs - -> Name -> PatternType -> m PatternType -mustHaveFieldWith onDims usage bcs l t = do - constraints <- getConstraints - l_type <- newTypeVar (srclocOf usage) "t" - let l_type' = toStruct l_type - case t of - Scalar (TypeVar _ _ (TypeName _ tn) []) - | Just (lvl, NoConstraint{}) <- M.lookup tn constraints -> do - scopeCheck usage bcs tn lvl l_type' - modifyConstraints $ M.insert tn (lvl, HasFields (M.singleton l l_type') usage) - return l_type - | Just (lvl, HasFields fields _) <- M.lookup tn constraints -> do - case M.lookup l fields of - Just t' -> unifyWith onDims usage bcs l_type' t' - Nothing -> modifyConstraints $ M.insert tn - (lvl, HasFields (M.insert l l_type' fields) usage) - return l_type - Scalar (Record fields) - | Just t' <- M.lookup l fields -> do - unify usage l_type' $ toStruct t' - return t' - | otherwise -> - unifyError usage mempty bcs $ - "Attempt to access field" <+> pquote (ppr l) <+> " of value of type" <+> - ppr (toStructural t) <> "." - _ -> do unify usage (toStruct t) $ Scalar $ Record $ M.singleton l l_type' - return l_type - --- | Assert that some type must have a field with this name and type. -mustHaveField :: MonadUnify m => - Usage -> Name -> PatternType -> m PatternType -mustHaveField usage = mustHaveFieldWith (unifyDims usage) usage noBreadCrumbs - --- | Replace dimension mismatches with AnyDim. -anyDimOnMismatch :: Monoid as => - TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as - -> (TypeBase (DimDecl VName) as, [(DimDecl VName, DimDecl VName)]) -anyDimOnMismatch t1 t2 = runWriter $ matchDims onDims t1 t2 - where onDims d1 d2 - | d1 == d2 = return d1 - | otherwise = do tell [(d1, d2)] - return AnyDim - -newDimOnMismatch :: (Monoid as, MonadUnify m) => - SrcLoc -> TypeBase (DimDecl VName) as -> TypeBase (DimDecl VName) as - -> m (TypeBase (DimDecl VName) as, [VName]) -newDimOnMismatch loc t1 t2 = do - (t, seen) <- runStateT (matchDims onDims t1 t2) mempty - return (t, M.elems seen) - where r = Rigid $ RigidCond (toStruct t1) (toStruct t2) - onDims d1 d2 - | d1 == d2 = return d1 - | otherwise = do - -- Remember mismatches we have seen before and reuse the - -- same new size. - maybe_d <- gets $ M.lookup (d1, d2) - case maybe_d of - Just d -> return $ NamedDim $ qualName d - Nothing -> do - d <- lift $ newDimVar loc r "differ" - modify $ M.insert (d1, d2) d - return $ NamedDim $ qualName d - --- | Like unification, but creates new size variables where mismatches --- occur. Returns the new dimensions thus created. -unifyMostCommon :: MonadUnify m => - Usage -> PatternType -> PatternType -> m (PatternType, [VName]) -unifyMostCommon usage t1 t2 = do - -- We are ignoring the dimensions here, because any mismatches - -- should be turned into fresh size variables. - let allOK _ _ _ _ _ = return () - unifyWith allOK usage noBreadCrumbs (toStruct t1) (toStruct t2) - t1' <- normTypeFully t1 - t2' <- normTypeFully t2 - newDimOnMismatch (srclocOf usage) t1' t2' - --- Simple MonadUnify implementation. - -type UnifyMState = (Constraints, Int) - -newtype UnifyM a = UnifyM (StateT UnifyMState (Except TypeError) a) - deriving (Monad, Functor, Applicative, - MonadState UnifyMState, - MonadError TypeError) - -newVar :: String -> UnifyM VName -newVar name = do - (x, i) <- get - put (x, i+1) - return $ VName (mkTypeVarName name i) i - -instance MonadUnify UnifyM where - getConstraints = gets fst - putConstraints x = modify $ \(_, i) -> (x, i) - - newTypeVar loc name = do - v <- newVar name - modifyConstraints $ M.insert v (0, NoConstraint Lifted $ Usage Nothing loc) - return $ Scalar $ TypeVar mempty Nonunique (typeName v) [] - - newDimVar loc rigidity name = do - dim <- newVar name - case rigidity of - Rigid src -> modifyConstraints $ M.insert dim (0, UnknowableSize loc src) - Nonrigid -> modifyConstraints $ M.insert dim (0, Size Nothing $ Usage Nothing loc) - return dim - - curLevel = pure 0 - - unifyError loc notes bcs doc = - throwError $ TypeError (srclocOf loc) notes $ doc <> ppr bcs - - matchError loc notes bcs t1 t2 = - throwError $ TypeError (srclocOf loc) notes $ doc <> ppr bcs - where doc = "Types" - indent 2 (ppr t1) - "and" - indent 2 (ppr t2) - "do not match." - --- | Construct a the name of a new type variable given a base --- description and a tag number (note that this is distinct from --- actually constructing a VName; the tag here is intended for human --- consumption but the machine does not care). -mkTypeVarName :: String -> Int -> Name -mkTypeVarName desc i = - nameFromString $ desc ++ mapMaybe subscript (show i) - where subscript = flip lookup $ zip "0123456789" "₀₁₂₃₄₅₆₇₈₉" - -runUnifyM :: [TypeParam] -> UnifyM a -> Either TypeError a -runUnifyM tparams (UnifyM m) = runExcept $ evalStateT m (constraints, 0) - where constraints = M.fromList $ map f tparams - f (TypeParamDim p loc) = (p, (0, Size Nothing $ Usage Nothing loc)) - f (TypeParamType l p loc) = (p, (0, NoConstraint l $ Usage Nothing loc)) - --- | Perform a unification of two types outside a monadic context. --- The type parameters are allowed to be instantiated; all other types --- are considered rigid. -doUnification :: SrcLoc -> [TypeParam] - -> StructType -> StructType - -> Either TypeError StructType +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE Trustworthy #-} + +-- | Implementation of unification and other core type system building +-- blocks. +module Language.Futhark.TypeChecker.Unify + ( Constraint (..), + Usage, + mkUsage, + mkUsage', + Level, + Constraints, + MonadUnify (..), + Rigidity (..), + RigidSource (..), + BreadCrumbs, + noBreadCrumbs, + hasNoBreadCrumbs, + dimNotes, + mkTypeVarName, + zeroOrderType, + mustHaveConstr, + mustHaveField, + mustBeOneOf, + equalityType, + normType, + normPatternType, + normTypeFully, + instantiateEmptyArrayDims, + unify, + expect, + unifyMostCommon, + anyDimOnMismatch, + doUnification, + ) +where + +import Control.Monad.Except +import Control.Monad.RWS.Strict hiding (Sum) +import Control.Monad.State +import Control.Monad.Writer hiding (Sum) +import Data.Bifoldable (biany) +import Data.List (intersect) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import Futhark.Util.Pretty hiding (empty) +import Language.Futhark hiding (unifyDims) +import Language.Futhark.TypeChecker.Monad hiding (BoundV) +import Language.Futhark.TypeChecker.Types + +-- | A piece of information that describes what process the type +-- checker currently performing. This is used to give better error +-- messages for unification errors. +data BreadCrumb + = MatchingTypes StructType StructType + | MatchingFields [Name] + | MatchingConstructor Name + | Matching Doc + +instance Pretty BreadCrumb where + ppr (MatchingTypes t1 t2) = + "When matching type" indent 2 (ppr t1) + "with" + indent 2 (ppr t2) + ppr (MatchingFields fields) = + "When matching types of record field" + <+> pquote (mconcat $ punctuate "." $ map ppr fields) <> dot + ppr (MatchingConstructor c) = + "When matching types of constructor" <+> pquote (ppr c) <> dot + ppr (Matching s) = + s + +-- | Unification failures can occur deep down inside complicated types +-- (consider nested records). We leave breadcrumbs behind us so we +-- can report the path we took to find the mismatch. +newtype BreadCrumbs = BreadCrumbs [BreadCrumb] + +-- | An empty path. +noBreadCrumbs :: BreadCrumbs +noBreadCrumbs = BreadCrumbs [] + +-- | Is the path empty? +hasNoBreadCrumbs :: BreadCrumbs -> Bool +hasNoBreadCrumbs (BreadCrumbs xs) = null xs + +-- | Drop a breadcrumb on the path behind you. +breadCrumb :: BreadCrumb -> BreadCrumbs -> BreadCrumbs +breadCrumb (MatchingFields xs) (BreadCrumbs (MatchingFields ys : bcs)) = + BreadCrumbs $ MatchingFields (ys ++ xs) : bcs +breadCrumb bc (BreadCrumbs bcs) = + BreadCrumbs $ bc : bcs + +instance Pretty BreadCrumbs where + ppr (BreadCrumbs []) = mempty + ppr (BreadCrumbs bcs) = line <> stack (map ppr bcs) + +-- | A usage that caused a type constraint. +data Usage = Usage (Maybe String) SrcLoc + deriving (Show) + +-- | Construct a 'Usage' from a location and a description. +mkUsage :: SrcLoc -> String -> Usage +mkUsage = flip (Usage . Just) + +-- | Construct a 'Usage' that has just a location, but no particular +-- description. +mkUsage' :: SrcLoc -> Usage +mkUsage' = Usage Nothing + +instance Pretty Usage where + ppr (Usage Nothing loc) = "use at " <> textwrap (locStr loc) + ppr (Usage (Just s) loc) = textwrap s <+/> "at" <+> textwrap (locStr loc) + +instance Located Usage where + locOf (Usage _ loc) = locOf loc + +-- | The level at which a type variable is bound. Higher means +-- deeper. We can only unify a type variable at level @i@ with a type +-- @t@ if all type names that occur in @t@ are at most at level @i@. +type Level = Int + +-- | A constraint on a yet-ambiguous type variable. +data Constraint + = NoConstraint Liftedness Usage + | ParamType Liftedness SrcLoc + | Constraint StructType Usage + | Overloaded [PrimType] Usage + | HasFields (M.Map Name StructType) Usage + | Equality Usage + | HasConstrs (M.Map Name [StructType]) Usage + | ParamSize SrcLoc + | -- | Is not actually a type, but a term-level size, + -- possibly already set to something specific. + Size (Maybe (DimDecl VName)) Usage + | -- | A size that does not unify with anything - + -- created from the result of applying a function + -- whose return size is existential, or otherwise + -- hiding a size. + UnknowableSize SrcLoc RigidSource + deriving (Show) + +instance Located Constraint where + locOf (NoConstraint _ usage) = locOf usage + locOf (ParamType _ usage) = locOf usage + locOf (Constraint _ usage) = locOf usage + locOf (Overloaded _ usage) = locOf usage + locOf (HasFields _ usage) = locOf usage + locOf (Equality usage) = locOf usage + locOf (HasConstrs _ usage) = locOf usage + locOf (ParamSize loc) = locOf loc + locOf (Size _ usage) = locOf usage + locOf (UnknowableSize loc _) = locOf loc + +-- | Mapping from fresh type variables, instantiated from the type +-- schemes of polymorphic functions, to (possibly) specific types as +-- determined on application and the location of that application, or +-- a partial constraint on their type. +type Constraints = M.Map VName (Level, Constraint) + +lookupSubst :: VName -> Constraints -> Maybe (Subst StructType) +lookupSubst v constraints = case snd <$> M.lookup v constraints of + Just (Constraint t _) -> Just $ Subst t + Just Overloaded {} -> Just PrimSubst + Just (Size (Just d) _) -> + Just $ SizeSubst $ applySubst (`lookupSubst` constraints) d + _ -> Nothing + +-- | The source of a rigid size. +data RigidSource + = -- | A function argument that is not a constant or variable name. + RigidArg (Maybe (QualName VName)) String + | -- | An existential return size. + RigidRet (Maybe (QualName VName)) + | RigidLoop + | -- | Produced by a complicated slice expression. + RigidSlice (Maybe (DimDecl VName)) String + | -- | Produced by a complicated range expression. + RigidRange + | -- | Produced by a range expression with this bound. + RigidBound String + | -- | Mismatch in branches. + RigidCond StructType StructType + | -- | Invented during unification. + RigidUnify + | RigidOutOfScope SrcLoc VName + deriving (Eq, Ord, Show) + +-- | The ridigity of a size variable. All rigid sizes are tagged with +-- information about how they were generated. +data Rigidity = Rigid RigidSource | Nonrigid + deriving (Eq, Ord, Show) + +prettySource :: SrcLoc -> SrcLoc -> RigidSource -> Doc +prettySource ctx loc (RigidRet Nothing) = + "is unknown size returned by function at" + <+> text (locStrRel ctx loc) <> "." +prettySource ctx loc (RigidRet (Just fname)) = + "is unknown size returned by" <+> pquote (ppr fname) + <+> "at" + <+> text (locStrRel ctx loc) <> "." +prettySource ctx loc (RigidArg fname arg) = + "is value of argument" + indent 2 (shorten arg) + "passed to" <+> fname' <+> "at" <+> text (locStrRel ctx loc) <> "." + where + fname' = maybe "function" (pquote . ppr) fname +prettySource ctx loc (RigidSlice d slice) = + "is size produced by slice" + indent 2 (shorten slice) + d_desc <> "at" <+> text (locStrRel ctx loc) <> "." + where + d_desc = case d of + Just d' -> "of dimension of size " <> pquote (ppr d') <> " " + Nothing -> mempty +prettySource ctx loc RigidLoop = + "is unknown size of value returned at" <+> text (locStrRel ctx loc) <> "." +prettySource ctx loc RigidRange = + "is unknown length of range at" <+> text (locStrRel ctx loc) <> "." +prettySource ctx loc (RigidBound bound) = + "generated from expression" + indent 2 (shorten bound) + "used in range at " <> text (locStrRel ctx loc) <> "." +prettySource ctx loc (RigidOutOfScope boundloc v) = + "is an unknown size arising from " <> pquote (pprName v) + <> " going out of scope at " + <> text (locStrRel ctx loc) + <> "." + "Originally bound at " + <> text (locStrRel ctx boundloc) + <> "." +prettySource _ _ RigidUnify = + "is an artificial size invented during unification of functions with anonymous sizes" +prettySource ctx loc (RigidCond t1 t2) = + "is unknown due to conditional expression at " + <> text (locStrRel ctx loc) + <> "." + "One branch returns array of type: " + <> align (ppr t1) + "The other an array of type: " + <> align (ppr t2) + +-- | Retrieve notes describing the purpose or origin of the given +-- 'DimDecl'. The location is used as the *current* location, for the +-- purpose of reporting relative locations. +dimNotes :: (Located a, MonadUnify m) => a -> DimDecl VName -> m Notes +dimNotes ctx (NamedDim d) = do + c <- M.lookup (qualLeaf d) <$> getConstraints + case c of + Just (_, UnknowableSize loc rsrc) -> + return $ + aNote $ + pretty $ + pquote (ppr d) <+> prettySource (srclocOf ctx) loc rsrc + _ -> return mempty +dimNotes _ _ = return mempty + +typeNotes :: (Located a, MonadUnify m) => a -> StructType -> m Notes +typeNotes ctx = + fmap mconcat . mapM (dimNotes ctx . NamedDim . qualName) + . S.toList + . typeDimNames + +-- | Monads that which to perform unification must implement this type +-- class. +class Monad m => MonadUnify m where + getConstraints :: m Constraints + putConstraints :: Constraints -> m () + modifyConstraints :: (Constraints -> Constraints) -> m () + modifyConstraints f = do + x <- getConstraints + putConstraints $ f x + + newTypeVar :: Monoid als => SrcLoc -> String -> m (TypeBase dim als) + newDimVar :: SrcLoc -> Rigidity -> String -> m VName + + curLevel :: m Level + + matchError :: + Located loc => + loc -> + Notes -> + BreadCrumbs -> + StructType -> + StructType -> + m a + + unifyError :: + Located loc => + loc -> + Notes -> + BreadCrumbs -> + Doc -> + m a + +-- | Replace all type variables with their substitution. +normTypeFully :: (Substitutable a, MonadUnify m) => a -> m a +normTypeFully t = do + constraints <- getConstraints + return $ applySubst (`lookupSubst` constraints) t + +-- | Replace any top-level type variable with its substitution. +normType :: MonadUnify m => StructType -> m StructType +normType t@(Scalar (TypeVar _ _ (TypeName [] v) [])) = do + constraints <- getConstraints + case snd <$> M.lookup v constraints of + Just (Constraint t' _) -> normType t' + _ -> return t +normType t = return t + +-- | Replace any top-level type variable with its substitution. +normPatternType :: MonadUnify m => PatternType -> m PatternType +normPatternType t@(Scalar (TypeVar als u (TypeName [] v) [])) = do + constraints <- getConstraints + case snd <$> M.lookup v constraints of + Just (Constraint t' _) -> + normPatternType $ t' `setUniqueness` u `setAliases` als + _ -> return t +normPatternType t = return t + +rigidConstraint :: Constraint -> Bool +rigidConstraint ParamType {} = True +rigidConstraint ParamSize {} = True +rigidConstraint UnknowableSize {} = True +rigidConstraint _ = False + +-- | Replace 'AnyDim' dimensions that occur as 'PosImmediate' or +-- 'PosParam' with a fresh 'NamedDim'. +instantiateEmptyArrayDims :: + MonadUnify m => + SrcLoc -> + String -> + Rigidity -> + TypeBase (DimDecl VName) als -> + m (TypeBase (DimDecl VName) als, [VName]) +instantiateEmptyArrayDims tloc desc r = runWriterT . traverseDims onDim + where + onDim _ PosImmediate AnyDim = inst + onDim _ PosParam AnyDim = inst + onDim _ _ d = return d + inst = do + dim <- lift $ newDimVar tloc r desc + tell [dim] + return $ NamedDim $ qualName dim + +-- | Is the given type variable the name of an abstract type or type +-- parameter, which we cannot substitute? +isRigid :: VName -> Constraints -> Bool +isRigid v constraints = + maybe True (rigidConstraint . snd) $ M.lookup v constraints + +-- | If the given type variable is nonrigid, what is its level? +isNonRigid :: VName -> Constraints -> Maybe Level +isNonRigid v constraints = do + (lvl, c) <- M.lookup v constraints + guard $ not $ rigidConstraint c + return lvl + +type UnifyDims m = + BreadCrumbs -> [VName] -> (VName -> Maybe Int) -> DimDecl VName -> DimDecl VName -> m () + +flipUnifyDims :: UnifyDims m -> UnifyDims m +flipUnifyDims onDims bcs bound nonrigid t1 t2 = + onDims bcs bound nonrigid t2 t1 + +unifyWith :: + MonadUnify m => + UnifyDims m -> + Usage -> + BreadCrumbs -> + StructType -> + StructType -> + m () +unifyWith onDims usage = subunify False mempty + where + swap True x y = (y, x) + swap False x y = (x, y) + + subunify ord bound bcs t1 t2 = do + constraints <- getConstraints + + t1' <- normType t1 + t2' <- normType t2 + + let nonrigid v = isNonRigid v constraints + + failure = matchError (srclocOf usage) mempty bcs t1' t2' + + -- Remove any of the intermediate dimensions we added just + -- for unification purposes. + unbound = applySubst f + where + f d + | d `elem` bound = Just $ SizeSubst AnyDim + | otherwise = Nothing + + link ord' v lvl = + linkVarToType linkDims usage bcs v lvl . unbound + where + -- We may have to flip the order of future calls to + -- onDims inside linkVarToType. + linkDims + | ord' = flipUnifyDims onDims + | otherwise = onDims + + unifyTypeArg bcs' (TypeArgDim d1 _) (TypeArgDim d2 _) = + onDims' bcs' (swap ord d1 d2) + unifyTypeArg bcs' (TypeArgType t _) (TypeArgType arg_t _) = + subunify ord bound bcs' t arg_t + unifyTypeArg bcs' _ _ = + unifyError + usage + mempty + bcs' + "Cannot unify a type argument with a dimension argument (or vice versa)." + + onDims' bcs' (d1, d2) = + onDims + bcs' + bound + nonrigid + (applySubst (`lookupSubst` constraints) d1) + (applySubst (`lookupSubst` constraints) d2) + + case (t1', t2') of + ( Scalar (Record fs), + Scalar (Record arg_fs) + ) + | M.keys fs == M.keys arg_fs -> + forM_ (M.toList $ M.intersectionWith (,) fs arg_fs) $ \(k, (k_t1, k_t2)) -> do + let bcs' = breadCrumb (MatchingFields [k]) bcs + subunify ord bound bcs' k_t1 k_t2 + | otherwise -> do + let missing = + filter (`notElem` M.keys arg_fs) (M.keys fs) + ++ filter (`notElem` M.keys fs) (M.keys arg_fs) + unifyError usage mempty bcs $ + "Unshared fields:" <+> commasep (map ppr missing) <> "." + ( Scalar (TypeVar _ _ (TypeName _ tn) targs), + Scalar (TypeVar _ _ (TypeName _ arg_tn) arg_targs) + ) + | tn == arg_tn, + length targs == length arg_targs -> do + let bcs' = breadCrumb (Matching "When matching type arguments.") bcs + zipWithM_ (unifyTypeArg bcs') targs arg_targs + ( Scalar (TypeVar _ _ (TypeName [] v1) []), + Scalar (TypeVar _ _ (TypeName [] v2) []) + ) -> + case (nonrigid v1, nonrigid v2) of + (Nothing, Nothing) -> failure + (Just lvl1, Nothing) -> link ord v1 lvl1 t2' + (Nothing, Just lvl2) -> link (not ord) v2 lvl2 t1' + (Just lvl1, Just lvl2) + | lvl1 <= lvl2 -> link ord v1 lvl1 t2' + | otherwise -> link (not ord) v2 lvl2 t1' + (Scalar (TypeVar _ _ (TypeName [] v1) []), _) + | Just lvl <- nonrigid v1 -> + link ord v1 lvl t2' + (_, Scalar (TypeVar _ _ (TypeName [] v2) [])) + | Just lvl <- nonrigid v2 -> + link (not ord) v2 lvl t1' + ( Scalar (Arrow _ p1 a1 b1), + Scalar (Arrow _ p2 a2 b2) + ) -> do + let (r1, r2) = swap ord (Rigid RigidUnify) Nonrigid + (a1', a1_dims) <- instantiateEmptyArrayDims (srclocOf usage) "anonymous" r1 a1 + (a2', a2_dims) <- instantiateEmptyArrayDims (srclocOf usage) "anonymous" r2 a2 + let bound' = bound <> mapMaybe pname [p1, p2] <> a1_dims <> a2_dims + subunify + (not ord) + bound + (breadCrumb (Matching "When matching parameter types.") bcs) + a1' + a2' + subunify + ord + bound' + (breadCrumb (Matching "When matching return types.") bcs) + b1' + b2' + where + (b1', b2') = + -- Replace one parameter name with the other in the + -- return type, in case of dependent types. I.e., + -- we want type '(n: i32) -> [n]i32' to unify with + -- type '(x: i32) -> [x]i32'. + case (p1, p2) of + (Named p1', Named p2') -> + let f v + | v == p2' = Just $ SizeSubst $ NamedDim $ qualName p1' + | otherwise = Nothing + in (b1, applySubst f b2) + (_, _) -> + (b1, b2) + + pname (Named x) = Just x + pname Unnamed = Nothing + (Array {}, Array {}) + | ShapeDecl (t1_d : _) <- arrayShape t1', + ShapeDecl (t2_d : _) <- arrayShape t2', + Just t1'' <- peelArray 1 t1', + Just t2'' <- peelArray 1 t2' -> do + onDims' bcs (swap ord t1_d t2_d) + subunify ord bound bcs t1'' t2'' + ( Scalar (Sum cs), + Scalar (Sum arg_cs) + ) + | M.keys cs == M.keys arg_cs -> + unifySharedConstructors + onDims + usage + bcs + (map unbound <$> cs) + (map unbound <$> arg_cs) + | otherwise -> do + let missing = + filter (`notElem` M.keys arg_cs) (M.keys cs) + ++ filter (`notElem` M.keys cs) (M.keys arg_cs) + unifyError usage mempty bcs $ + "Unshared constructors:" <+> commasep (map (("#" <>) . ppr) missing) <> "." + _ + | t1' == t2' -> return () + | otherwise -> failure + +unifyDims :: MonadUnify m => Usage -> UnifyDims m +unifyDims _ _ _ _ d1 d2 + | d1 == d2 = return () +unifyDims usage bcs _ nonrigid (NamedDim (QualName _ d1)) d2 + | Just lvl1 <- nonrigid d1 = + linkVarToDim usage bcs d1 lvl1 d2 +unifyDims usage bcs _ nonrigid d1 (NamedDim (QualName _ d2)) + | Just lvl2 <- nonrigid d2 = + linkVarToDim usage bcs d2 lvl2 d1 +unifyDims usage bcs _ _ d1 d2 = do + notes <- (<>) <$> dimNotes usage d1 <*> dimNotes usage d2 + unifyError usage notes bcs $ + "Dimensions" <+> pquote (ppr d1) + <+> "and" + <+> pquote (ppr d2) + <+> "do not match." + +-- | Unifies two types. +unify :: MonadUnify m => Usage -> StructType -> StructType -> m () +unify usage = unifyWith (unifyDims usage) usage noBreadCrumbs + +-- | @expect super sub@ checks that @sub@ is a subtype of @super@. +expect :: MonadUnify m => Usage -> StructType -> StructType -> m () +expect usage = unifyWith onDims usage noBreadCrumbs + where + onDims _ _ _ AnyDim _ = return () + onDims _ _ _ d1 d2 + | d1 == d2 = return () + onDims bcs bound nonrigid (NamedDim (QualName _ d1)) d2 + | Just lvl1 <- nonrigid d1, + d2 /= AnyDim, + not $ boundParam bound d2 = + linkVarToDim usage bcs d1 lvl1 d2 + onDims bcs bound nonrigid d1 (NamedDim (QualName _ d2)) + | Just lvl2 <- nonrigid d2, + not $ boundParam bound d1 = + linkVarToDim usage bcs d2 lvl2 d1 + onDims bcs _ _ d1 d2 = do + notes <- (<>) <$> dimNotes usage d1 <*> dimNotes usage d2 + unifyError usage notes bcs $ + "Dimensions" <+> pquote (ppr d1) + <+> "and" + <+> pquote (ppr d2) + <+> "do not match." + + boundParam bound (NamedDim (QualName _ d)) = d `elem` bound + boundParam _ _ = False + +hasEmptyDims :: StructType -> Bool +hasEmptyDims = biany empty (const False) + where + empty AnyDim = True + empty _ = False + +occursCheck :: + MonadUnify m => + Usage -> + BreadCrumbs -> + VName -> + StructType -> + m () +occursCheck usage bcs vn tp = + when (vn `S.member` typeVars tp) $ + unifyError usage mempty bcs $ + "Occurs check: cannot instantiate" + <+> pprName vn + <+> "with" + <+> ppr tp <> "." + +scopeCheck :: + MonadUnify m => + Usage -> + BreadCrumbs -> + VName -> + Level -> + StructType -> + m () +scopeCheck usage bcs vn max_lvl tp = do + constraints <- getConstraints + checkType constraints tp + where + checkType constraints t = + mapM_ (check constraints) $ typeVars t <> typeDimNames t + + check constraints v + | Just (lvl, c) <- M.lookup v constraints, + lvl > max_lvl = + if rigidConstraint c + then scopeViolation v + else modifyConstraints $ M.insert v (max_lvl, c) + | otherwise = + return () + + scopeViolation v = do + notes <- typeNotes usage tp + unifyError usage notes bcs $ + "Cannot unify type" + indent 2 (ppr tp) + "with" + <+> pquote (pprName vn) + <+> "(scope violation)." + "This is because" + <+> pquote (pprName v) + <+> "is rigidly bound in a deeper scope." + +linkVarToType :: + MonadUnify m => + UnifyDims m -> + Usage -> + BreadCrumbs -> + VName -> + Level -> + StructType -> + m () +linkVarToType onDims usage bcs vn lvl tp = do + occursCheck usage bcs vn tp + scopeCheck usage bcs vn lvl tp + + constraints <- getConstraints + let tp' = removeUniqueness tp + modifyConstraints $ M.insert vn (lvl, Constraint tp' usage) + case snd <$> M.lookup vn constraints of + Just (NoConstraint Unlifted unlift_usage) -> do + let bcs' = + breadCrumb + ( Matching $ + "When verifying that" <+> pquote (pprName vn) + <+> textwrap "is not instantiated with a function type, due to" + <+> ppr unlift_usage + ) + bcs + zeroOrderTypeWith usage bcs' tp' + + when (hasEmptyDims tp') $ + unifyError usage mempty bcs $ + "Type variable" <+> pprName vn + <+> "cannot be instantiated with type containing anonymous sizes:" + indent 2 (ppr tp) + textwrap "This is usually because the size of an array returned by a higher-order function argument cannot be determined statically. This can also be due to the return size being a value parameter. Add type annotation to clarify." + Just (Equality _) -> + equalityType usage tp' + Just (Overloaded ts old_usage) + | tp `notElem` map (Scalar . Prim) ts -> + case tp' of + Scalar (TypeVar _ _ (TypeName [] v) []) + | not $ isRigid v constraints -> + linkVarToTypes usage v ts + _ -> + unifyError usage mempty bcs $ + "Cannot instantiate" <+> pquote (pprName vn) + <+> "with type" indent 2 (ppr tp) "as" + <+> pquote (pprName vn) + <+> "must be one of" + <+> commasep (map ppr ts) + <+/> "due to" + <+/> ppr old_usage <> "." + Just (HasFields required_fields old_usage) -> + case tp of + Scalar (Record tp_fields) + | all (`M.member` tp_fields) $ M.keys required_fields -> do + required_fields' <- mapM normTypeFully required_fields + let bcs' = + breadCrumb + ( Matching $ + pprName vn + <+> "must be a record with at least the fields:" + indent 2 (ppr (Record required_fields')) + "due to" + <+> ppr old_usage <> "." + ) + bcs + mapM_ (uncurry $ unifyWith onDims usage bcs') $ + M.elems $ + M.intersectionWith (,) required_fields tp_fields + Scalar (TypeVar _ _ (TypeName [] v) []) + | not $ isRigid v constraints -> + modifyConstraints $ + M.insert + v + (lvl, HasFields required_fields old_usage) + _ -> + unifyError usage mempty bcs $ + "Cannot instantiate" <+> pquote (pprName vn) <+> "with type" + indent 2 (ppr tp) + "as" <+> pquote (pprName vn) <+> "must be a record with fields" + indent 2 (ppr (Record required_fields)) + "due to" <+> ppr old_usage <> "." + Just (HasConstrs required_cs old_usage) -> + case tp of + Scalar (Sum ts) + | all (`M.member` ts) $ M.keys required_cs -> + unifySharedConstructors onDims usage bcs required_cs ts + Scalar (TypeVar _ _ (TypeName [] v) []) + | not $ isRigid v constraints -> do + case M.lookup v constraints of + Just (_, HasConstrs v_cs _) -> + unifySharedConstructors onDims usage bcs required_cs v_cs + _ -> return () + modifyConstraints $ + M.insertWith + combineConstrs + v + (lvl, HasConstrs required_cs old_usage) + where + combineConstrs (_, HasConstrs cs1 usage1) (_, HasConstrs cs2 _) = + (lvl, HasConstrs (M.union cs1 cs2) usage1) + combineConstrs hasCs _ = hasCs + _ -> noSumType + _ -> return () + where + noSumType = + unifyError + usage + mempty + bcs + "Cannot unify a sum type with a non-sum type" + +linkVarToDim :: + MonadUnify m => + Usage -> + BreadCrumbs -> + VName -> + Level -> + DimDecl VName -> + m () +linkVarToDim usage bcs vn lvl dim = do + constraints <- getConstraints + + case dim of + NamedDim dim' + | Just (dim_lvl, c) <- qualLeaf dim' `M.lookup` constraints, + dim_lvl > lvl -> + case c of + ParamSize {} -> do + notes <- dimNotes usage dim + unifyError usage notes bcs $ + "Cannot unify size variable" <+> pquote (ppr dim') + <+> "with" + <+> pquote (pprName vn) + <+> "(scope violation)." + "This is because" + <+> pquote (ppr dim') + <+> "is rigidly bound in a deeper scope." + _ -> modifyConstraints $ M.insert (qualLeaf dim') (lvl, c) + _ -> return () + + modifyConstraints $ M.insert vn (lvl, Size (Just dim) usage) + +removeUniqueness :: TypeBase dim as -> TypeBase dim as +removeUniqueness (Scalar (Record ets)) = + Scalar $ Record $ fmap removeUniqueness ets +removeUniqueness (Scalar (Arrow als p t1 t2)) = + Scalar $ Arrow als p (removeUniqueness t1) (removeUniqueness t2) +removeUniqueness (Scalar (Sum cs)) = + Scalar $ Sum $ (fmap . fmap) removeUniqueness cs +removeUniqueness t = t `setUniqueness` Nonunique + +-- | Assert that this type must be one of the given primitive types. +mustBeOneOf :: MonadUnify m => [PrimType] -> Usage -> StructType -> m () +mustBeOneOf [req_t] usage t = unify usage (Scalar (Prim req_t)) t +mustBeOneOf ts usage t = do + t' <- normType t + constraints <- getConstraints + let isRigid' v = isRigid v constraints + + case t' of + Scalar (TypeVar _ _ (TypeName [] v) []) + | not $ isRigid' v -> linkVarToTypes usage v ts + Scalar (Prim pt) | pt `elem` ts -> return () + _ -> failure + where + failure = + unifyError usage mempty noBreadCrumbs $ + text "Cannot unify type" <+> pquote (ppr t) + <+> "with any of " <> commasep (map ppr ts) <> "." + +linkVarToTypes :: MonadUnify m => Usage -> VName -> [PrimType] -> m () +linkVarToTypes usage vn ts = do + vn_constraint <- M.lookup vn <$> getConstraints + case vn_constraint of + Just (lvl, Overloaded vn_ts vn_usage) -> + case ts `intersect` vn_ts of + [] -> + unifyError usage mempty noBreadCrumbs $ + "Type constrained to one of" + <+> commasep (map ppr ts) + <+> "but also one of" + <+> commasep (map ppr vn_ts) + <+> "due to" + <+> ppr vn_usage <> "." + ts' -> modifyConstraints $ M.insert vn (lvl, Overloaded ts' usage) + Just (_, HasConstrs _ vn_usage) -> + unifyError usage mempty noBreadCrumbs $ + "Type constrained to one of" <+> commasep (map ppr ts) + <> ", but also inferred to be sum type due to" <+> ppr vn_usage + <> "." + Just (_, HasFields _ vn_usage) -> + unifyError usage mempty noBreadCrumbs $ + "Type constrained to one of" <+> commasep (map ppr ts) + <> ", but also inferred to be record due to" <+> ppr vn_usage + <> "." + Just (lvl, _) -> modifyConstraints $ M.insert vn (lvl, Overloaded ts usage) + Nothing -> + unifyError usage mempty noBreadCrumbs $ + "Cannot constrain type to one of" <+> commasep (map ppr ts) + +-- | Assert that this type must support equality. +equalityType :: + (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => + Usage -> + TypeBase dim as -> + m () +equalityType usage t = do + unless (orderZero t) $ + unifyError usage mempty noBreadCrumbs $ + "Type " <+> pquote (ppr t) <+> "does not support equality (is higher-order)." + mapM_ mustBeEquality $ typeVars t + where + mustBeEquality vn = do + constraints <- getConstraints + case M.lookup vn constraints of + Just (_, Constraint (Scalar (TypeVar _ _ (TypeName [] vn') [])) _) -> + mustBeEquality vn' + Just (_, Constraint vn_t cusage) + | not $ orderZero vn_t -> + unifyError usage mempty noBreadCrumbs $ + "Type" <+> pquote (ppr t) <+> "does not support equality." + "Constrained to be higher-order due to" <+> ppr cusage <+> "." + | otherwise -> return () + Just (lvl, NoConstraint _ _) -> + modifyConstraints $ M.insert vn (lvl, Equality usage) + Just (_, Overloaded _ _) -> + return () -- All primtypes support equality. + Just (_, Equality {}) -> + return () + Just (_, HasConstrs cs _) -> + mapM_ (equalityType usage) $ concat $ M.elems cs + _ -> + unifyError usage mempty noBreadCrumbs $ + "Type" <+> pprName vn <+> "does not support equality." + +zeroOrderTypeWith :: + (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => + Usage -> + BreadCrumbs -> + TypeBase dim as -> + m () +zeroOrderTypeWith usage bcs t = do + unless (orderZero t) $ + unifyError usage mempty bcs $ + "Type" indent 2 (ppr t) "found to be functional." + mapM_ mustBeZeroOrder . S.toList . typeVars $ t + where + mustBeZeroOrder vn = do + constraints <- getConstraints + case M.lookup vn constraints of + Just (lvl, NoConstraint _ _) -> + modifyConstraints $ M.insert vn (lvl, NoConstraint Unlifted usage) + Just (_, ParamType Lifted ploc) -> + unifyError usage mempty bcs $ + "Type parameter" + <+> pquote (pprName vn) + <+> "at" + <+> text (locStr ploc) + <+> "may be a function." + _ -> return () + +-- | Assert that this type must be zero-order. +zeroOrderType :: + (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => + Usage -> + String -> + TypeBase dim as -> + m () +zeroOrderType usage desc = + zeroOrderTypeWith usage $ breadCrumb bc noBreadCrumbs + where + bc = Matching $ "When checking" <+> textwrap desc + +unifySharedConstructors :: + MonadUnify m => + UnifyDims m -> + Usage -> + BreadCrumbs -> + M.Map Name [StructType] -> + M.Map Name [StructType] -> + m () +unifySharedConstructors onDims usage bcs cs1 cs2 = + forM_ (M.toList $ M.intersectionWith (,) cs1 cs2) $ \(c, (f1, f2)) -> + unifyConstructor c f1 f2 + where + unifyConstructor c f1 f2 + | length f1 == length f2 = do + let bcs' = breadCrumb (MatchingConstructor c) bcs + zipWithM_ (unifyWith onDims usage bcs') f1 f2 + | otherwise = + unifyError usage mempty bcs $ + "Cannot unify constructor" <+> pquote (pprName c) <> "." + +-- | In @mustHaveConstr usage c t fs@, the type @t@ must have a +-- constructor named @c@ that takes arguments of types @ts@. +mustHaveConstr :: + MonadUnify m => + Usage -> + Name -> + StructType -> + [StructType] -> + m () +mustHaveConstr usage c t fs = do + constraints <- getConstraints + case t of + Scalar (TypeVar _ _ (TypeName _ tn) []) + | Just (lvl, NoConstraint {}) <- M.lookup tn constraints -> do + mapM_ (scopeCheck usage noBreadCrumbs tn lvl) fs + modifyConstraints $ M.insert tn (lvl, HasConstrs (M.singleton c fs) usage) + | Just (lvl, HasConstrs cs _) <- M.lookup tn constraints -> + case M.lookup c cs of + Nothing -> modifyConstraints $ M.insert tn (lvl, HasConstrs (M.insert c fs cs) usage) + Just fs' + | length fs == length fs' -> zipWithM_ (unify usage) fs fs' + | otherwise -> + unifyError usage mempty noBreadCrumbs $ + "Different arity for constructor" <+> pquote (ppr c) <> "." + Scalar (Sum cs) -> + case M.lookup c cs of + Nothing -> + unifyError usage mempty noBreadCrumbs $ + "Constuctor" <+> pquote (ppr c) <+> "not present in type." + Just fs' + | length fs == length fs' -> zipWithM_ (unify usage) fs fs' + | otherwise -> + unifyError usage mempty noBreadCrumbs $ + "Different arity for constructor" <+> pquote (ppr c) <+> "." + _ -> do + unify usage t $ Scalar $ Sum $ M.singleton c fs + return () + +mustHaveFieldWith :: + MonadUnify m => + UnifyDims m -> + Usage -> + BreadCrumbs -> + Name -> + PatternType -> + m PatternType +mustHaveFieldWith onDims usage bcs l t = do + constraints <- getConstraints + l_type <- newTypeVar (srclocOf usage) "t" + let l_type' = toStruct l_type + case t of + Scalar (TypeVar _ _ (TypeName _ tn) []) + | Just (lvl, NoConstraint {}) <- M.lookup tn constraints -> do + scopeCheck usage bcs tn lvl l_type' + modifyConstraints $ M.insert tn (lvl, HasFields (M.singleton l l_type') usage) + return l_type + | Just (lvl, HasFields fields _) <- M.lookup tn constraints -> do + case M.lookup l fields of + Just t' -> unifyWith onDims usage bcs l_type' t' + Nothing -> + modifyConstraints $ + M.insert + tn + (lvl, HasFields (M.insert l l_type' fields) usage) + return l_type + Scalar (Record fields) + | Just t' <- M.lookup l fields -> do + unify usage l_type' $ toStruct t' + return t' + | otherwise -> + unifyError usage mempty bcs $ + "Attempt to access field" <+> pquote (ppr l) <+> " of value of type" + <+> ppr (toStructural t) <> "." + _ -> do + unify usage (toStruct t) $ Scalar $ Record $ M.singleton l l_type' + return l_type + +-- | Assert that some type must have a field with this name and type. +mustHaveField :: + MonadUnify m => + Usage -> + Name -> + PatternType -> + m PatternType +mustHaveField usage = mustHaveFieldWith (unifyDims usage) usage noBreadCrumbs + +-- | Replace dimension mismatches with AnyDim. +anyDimOnMismatch :: + Monoid as => + TypeBase (DimDecl VName) as -> + TypeBase (DimDecl VName) as -> + (TypeBase (DimDecl VName) as, [(DimDecl VName, DimDecl VName)]) +anyDimOnMismatch t1 t2 = runWriter $ matchDims onDims t1 t2 + where + onDims d1 d2 + | d1 == d2 = return d1 + | otherwise = do + tell [(d1, d2)] + return AnyDim + +newDimOnMismatch :: + (Monoid as, MonadUnify m) => + SrcLoc -> + TypeBase (DimDecl VName) as -> + TypeBase (DimDecl VName) as -> + m (TypeBase (DimDecl VName) as, [VName]) +newDimOnMismatch loc t1 t2 = do + (t, seen) <- runStateT (matchDims onDims t1 t2) mempty + return (t, M.elems seen) + where + r = Rigid $ RigidCond (toStruct t1) (toStruct t2) + onDims d1 d2 + | d1 == d2 = return d1 + | otherwise = do + -- Remember mismatches we have seen before and reuse the + -- same new size. + maybe_d <- gets $ M.lookup (d1, d2) + case maybe_d of + Just d -> return $ NamedDim $ qualName d + Nothing -> do + d <- lift $ newDimVar loc r "differ" + modify $ M.insert (d1, d2) d + return $ NamedDim $ qualName d + +-- | Like unification, but creates new size variables where mismatches +-- occur. Returns the new dimensions thus created. +unifyMostCommon :: + MonadUnify m => + Usage -> + PatternType -> + PatternType -> + m (PatternType, [VName]) +unifyMostCommon usage t1 t2 = do + -- We are ignoring the dimensions here, because any mismatches + -- should be turned into fresh size variables. + let allOK _ _ _ _ _ = return () + unifyWith allOK usage noBreadCrumbs (toStruct t1) (toStruct t2) + t1' <- normTypeFully t1 + t2' <- normTypeFully t2 + newDimOnMismatch (srclocOf usage) t1' t2' + +-- Simple MonadUnify implementation. + +type UnifyMState = (Constraints, Int) + +newtype UnifyM a = UnifyM (StateT UnifyMState (Except TypeError) a) + deriving + ( Monad, + Functor, + Applicative, + MonadState UnifyMState, + MonadError TypeError + ) + +newVar :: String -> UnifyM VName +newVar name = do + (x, i) <- get + put (x, i + 1) + return $ VName (mkTypeVarName name i) i + +instance MonadUnify UnifyM where + getConstraints = gets fst + putConstraints x = modify $ \(_, i) -> (x, i) + + newTypeVar loc name = do + v <- newVar name + modifyConstraints $ M.insert v (0, NoConstraint Lifted $ Usage Nothing loc) + return $ Scalar $ TypeVar mempty Nonunique (typeName v) [] + + newDimVar loc rigidity name = do + dim <- newVar name + case rigidity of + Rigid src -> modifyConstraints $ M.insert dim (0, UnknowableSize loc src) + Nonrigid -> modifyConstraints $ M.insert dim (0, Size Nothing $ Usage Nothing loc) + return dim + + curLevel = pure 0 + + unifyError loc notes bcs doc = + throwError $ TypeError (srclocOf loc) notes $ doc <> ppr bcs + + matchError loc notes bcs t1 t2 = + throwError $ TypeError (srclocOf loc) notes $ doc <> ppr bcs + where + doc = + "Types" + indent 2 (ppr t1) + "and" + indent 2 (ppr t2) + "do not match." + +-- | Construct a the name of a new type variable given a base +-- description and a tag number (note that this is distinct from +-- actually constructing a VName; the tag here is intended for human +-- consumption but the machine does not care). +mkTypeVarName :: String -> Int -> Name +mkTypeVarName desc i = + nameFromString $ desc ++ mapMaybe subscript (show i) + where + subscript = flip lookup $ zip "0123456789" "₀₁₂₃₄₅₆₇₈₉" + +runUnifyM :: [TypeParam] -> UnifyM a -> Either TypeError a +runUnifyM tparams (UnifyM m) = runExcept $ evalStateT m (constraints, 0) + where + constraints = M.fromList $ map f tparams + f (TypeParamDim p loc) = (p, (0, Size Nothing $ Usage Nothing loc)) + f (TypeParamType l p loc) = (p, (0, NoConstraint l $ Usage Nothing loc)) + +-- | Perform a unification of two types outside a monadic context. +-- The type parameters are allowed to be instantiated; all other types +-- are considered rigid. +doUnification :: + SrcLoc -> + [TypeParam] -> + StructType -> + StructType -> + Either TypeError StructType doUnification loc tparams t1 t2 = runUnifyM tparams $ do let rsrc = RigidUnify (t1', _) <- instantiateEmptyArrayDims loc "n" (Rigid rsrc) t1 diff --git a/src/Language/Futhark/Warnings.hs b/src/Language/Futhark/Warnings.hs --- a/src/Language/Futhark/Warnings.hs +++ b/src/Language/Futhark/Warnings.hs @@ -1,21 +1,21 @@ {-# LANGUAGE Safe #-} + -- | A very simple representation of collections of warnings. -- Warnings have a position (so they can be ordered), and their -- 'Show'-instance produces a human-readable string. module Language.Futhark.Warnings - ( Warnings - , singleWarning - , singleWarning' - ) where + ( Warnings, + singleWarning, + singleWarning', + ) +where +import Data.List (intercalate, sortOn) import Data.Monoid -import Data.List (sortOn, intercalate) - -import Prelude - -import Language.Futhark.Core (locStr, prettyStacktrace) import Futhark.Util.Console (inRed) import Futhark.Util.Loc +import Language.Futhark.Core (locStr, prettyStacktrace) +import Prelude -- | The warnings produced by the compiler. The 'Show' instance -- produces a human-readable description. @@ -31,17 +31,20 @@ show (Warnings []) = "" show (Warnings ws) = intercalate "\n\n" ws' ++ "\n" - where ws' = map showWarning $ sortOn (rep . wloc) ws - wloc (x, _, _) = locOf x - rep NoLoc = ("", 0) - rep (Loc p _) = (posFile p, posCoff p) - showWarning (loc, [], w) = - inRed ("Warning at " ++ locStr loc ++ ":") ++ "\n" ++ - intercalate "\n" (map (" "<>) $ lines w) - showWarning (loc, locs, w) = - inRed ("Warning at\n" ++ - prettyStacktrace 0 (map locStr (loc:locs))) ++ - intercalate "\n" (map (" "<>) $ lines w) + where + ws' = map showWarning $ sortOn (rep . wloc) ws + wloc (x, _, _) = locOf x + rep NoLoc = ("", 0) + rep (Loc p _) = (posFile p, posCoff p) + showWarning (loc, [], w) = + inRed ("Warning at " ++ locStr loc ++ ":") ++ "\n" + ++ intercalate "\n" (map (" " <>) $ lines w) + showWarning (loc, locs, w) = + inRed + ( "Warning at\n" + ++ prettyStacktrace 0 (map locStr (loc : locs)) + ) + ++ intercalate "\n" (map (" " <>) $ lines w) -- | A single warning at the given location. singleWarning :: SrcLoc -> String -> Warnings diff --git a/src/futhark.hs b/src/futhark.hs --- a/src/futhark.hs +++ b/src/futhark.hs @@ -1,108 +1,105 @@ {-# LANGUAGE OverloadedStrings #-} + -- | The futhark command line tool. module Main (main) where -import Data.Maybe import Control.Exception import Control.Monad import Data.List (sortOn) +import Data.Maybe import qualified Data.Text as T import qualified Data.Text.IO as T -import GHC.IO.Encoding (setLocaleEncoding) -import System.IO -import System.Exit -import System.Environment - -import Prelude - -import Futhark.Error -import Futhark.Util.Options -import Futhark.Util (maxinum) - -import qualified Futhark.CLI.Dev as Dev +import qualified Futhark.CLI.Autotune as Autotune +import qualified Futhark.CLI.Bench as Bench import qualified Futhark.CLI.C as C -import qualified Futhark.CLI.OpenCL as OpenCL import qualified Futhark.CLI.CUDA as CCUDA -import qualified Futhark.CLI.Python as Python -import qualified Futhark.CLI.PyOpenCL as PyOpenCL -import qualified Futhark.CLI.Test as Test -import qualified Futhark.CLI.Bench as Bench import qualified Futhark.CLI.Check as Check -import qualified Futhark.CLI.Dataset as Dataset import qualified Futhark.CLI.Datacmp as Datacmp -import qualified Futhark.CLI.Pkg as Pkg +import qualified Futhark.CLI.Dataset as Dataset +import qualified Futhark.CLI.Dev as Dev import qualified Futhark.CLI.Doc as Doc -import qualified Futhark.CLI.REPL as REPL -import qualified Futhark.CLI.Run as Run import qualified Futhark.CLI.Misc as Misc -import qualified Futhark.CLI.Autotune as Autotune +import qualified Futhark.CLI.OpenCL as OpenCL +import qualified Futhark.CLI.Pkg as Pkg +import qualified Futhark.CLI.PyOpenCL as PyOpenCL +import qualified Futhark.CLI.Python as Python import qualified Futhark.CLI.Query as Query +import qualified Futhark.CLI.REPL as REPL +import qualified Futhark.CLI.Run as Run +import qualified Futhark.CLI.Test as Test +import Futhark.Error +import Futhark.Util (maxinum) +import Futhark.Util.Options +import GHC.IO.Encoding (setLocaleEncoding) +import System.Environment +import System.Exit +import System.IO +import Prelude type Command = String -> [String] -> IO () commands :: [(String, (Command, String))] -commands = sortOn fst - [ ("dev", (Dev.main, "Run compiler passes directly.")) - - , ("repl", (REPL.main, "Run interactive Read-Eval-Print-Loop.")) - , ("run", (Run.main, "Run a program through the (slow!) interpreter.")) - - , ("c", (C.main, "Compile to sequential C.")) - , ("opencl", (OpenCL.main, "Compile to C calling OpenCL.")) - , ("cuda", (CCUDA.main, "Compile to C calling CUDA.")) - - , ("python", (Python.main, "Compile to sequential Python.")) - , ("pyopencl", (PyOpenCL.main, "Compile to Python calling PyOpenCL.")) - - , ("test", (Test.main, "Test Futhark programs.")) - , ("bench", (Bench.main, "Benchmark Futhark programs.")) - - , ("dataset", (Dataset.main, "Generate random test data.")) - , ("datacmp", (Datacmp.main, "Compare Futhark data files for equality.")) - , ("dataget", (Misc.mainDataget, "Extract test data.")) - - , ("doc", (Doc.main, "Generate documentation for Futhark code.")) - , ("pkg", (Pkg.main, "Manage local packages.")) - - , ("check", (Check.main, "Type check a program.")) - , ("imports", (Misc.mainImports, "Print all non-builtin imported Futhark files.")) - , ("autotune", (Autotune.main, "Autotune threshold parameters.")) - , ("query", (Query.main, "Query semantic information about program.")) - ] +commands = + sortOn + fst + [ ("dev", (Dev.main, "Run compiler passes directly.")), + ("repl", (REPL.main, "Run interactive Read-Eval-Print-Loop.")), + ("run", (Run.main, "Run a program through the (slow!) interpreter.")), + ("c", (C.main, "Compile to sequential C.")), + ("opencl", (OpenCL.main, "Compile to C calling OpenCL.")), + ("cuda", (CCUDA.main, "Compile to C calling CUDA.")), + ("python", (Python.main, "Compile to sequential Python.")), + ("pyopencl", (PyOpenCL.main, "Compile to Python calling PyOpenCL.")), + ("test", (Test.main, "Test Futhark programs.")), + ("bench", (Bench.main, "Benchmark Futhark programs.")), + ("dataset", (Dataset.main, "Generate random test data.")), + ("datacmp", (Datacmp.main, "Compare Futhark data files for equality.")), + ("dataget", (Misc.mainDataget, "Extract test data.")), + ("doc", (Doc.main, "Generate documentation for Futhark code.")), + ("pkg", (Pkg.main, "Manage local packages.")), + ("check", (Check.main, "Type check a program.")), + ("imports", (Misc.mainImports, "Print all non-builtin imported Futhark files.")), + ("autotune", (Autotune.main, "Autotune threshold parameters.")), + ("query", (Query.main, "Query semantic information about program.")) + ] msg :: String -msg = unlines $ - [" options...", "Commands:", ""] ++ - [ " " <> cmd <> replicate (k - length cmd) ' ' <> desc - | (cmd, (_, desc)) <- commands ] - where k = maxinum (map (length . fst) commands) + 3 +msg = + unlines $ + [" options...", "Commands:", ""] + ++ [ " " <> cmd <> replicate (k - length cmd) ' ' <> desc + | (cmd, (_, desc)) <- commands + ] + where + k = maxinum (map (length . fst) commands) + 3 -- | Catch all IO exceptions and print a better error message if they -- happen. reportingIOErrors :: IO () -> IO () reportingIOErrors = flip catches [Handler onExit, Handler onICE, Handler onError] - where onExit :: ExitCode -> IO () - onExit = throwIO + where + onExit :: ExitCode -> IO () + onExit = throwIO - onICE :: InternalError -> IO () - onICE (Error CompilerLimitation s) = do - T.hPutStrLn stderr "Known compiler limitation encountered. Sorry." - T.hPutStrLn stderr "Revise your program or try a different Futhark compiler." - T.hPutStrLn stderr s - onICE (Error CompilerBug s) = do - T.hPutStrLn stderr "Internal compiler error." - T.hPutStrLn stderr "Please report this at https://github.com/diku-dk/futhark/issues." - T.hPutStrLn stderr s + onICE :: InternalError -> IO () + onICE (Error CompilerLimitation s) = do + T.hPutStrLn stderr "Known compiler limitation encountered. Sorry." + T.hPutStrLn stderr "Revise your program or try a different Futhark compiler." + T.hPutStrLn stderr s + onICE (Error CompilerBug s) = do + T.hPutStrLn stderr "Internal compiler error." + T.hPutStrLn stderr "Please report this at https://github.com/diku-dk/futhark/issues." + T.hPutStrLn stderr s - onError :: SomeException -> IO () - onError e - | Just UserInterrupt <- asyncExceptionFromException e = - return () -- This corresponds to CTRL-C, which is not an error. - | otherwise = do - T.hPutStrLn stderr "Internal compiler error (unhandled IO exception)." - T.hPutStrLn stderr "Please report this at https://github.com/diku-dk/futhark/issues" - T.hPutStrLn stderr $ T.pack $ show e - exitWith $ ExitFailure 1 + onError :: SomeException -> IO () + onError e + | Just UserInterrupt <- asyncExceptionFromException e = + return () -- This corresponds to CTRL-C, which is not an error. + | otherwise = do + T.hPutStrLn stderr "Internal compiler error (unhandled IO exception)." + T.hPutStrLn stderr "Please report this at https://github.com/diku-dk/futhark/issues" + T.hPutStrLn stderr $ T.pack $ show e + exitWith $ ExitFailure 1 main :: IO () main = reportingIOErrors $ do @@ -112,6 +109,6 @@ args <- getArgs prog <- getProgName case args of - cmd:args' + cmd : args' | Just (m, _) <- lookup cmd commands -> m (unwords [prog, cmd]) args' _ -> mainWithOptions () [] msg (const . const Nothing) prog args diff --git a/unittests/Futhark/BenchTests.hs b/unittests/Futhark/BenchTests.hs --- a/unittests/Futhark/BenchTests.hs +++ b/unittests/Futhark/BenchTests.hs @@ -1,13 +1,12 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} + module Futhark.BenchTests (tests) where import qualified Data.Text as T - +import Futhark.Bench import Test.Tasty import Test.Tasty.QuickCheck -import Futhark.Bench - instance Arbitrary RunResult where arbitrary = RunResult . getPositive <$> arbitrary @@ -15,11 +14,15 @@ printable = getPrintableString <$> arbitrary instance Arbitrary DataResult where - arbitrary = DataResult - <$> printable - <*> oneof [Left <$> arbText, - Right <$> ((,) <$> arbitrary <*> arbText)] - where arbText = T.pack <$> printable + arbitrary = + DataResult + <$> printable + <*> oneof + [ Left <$> arbText, + Right <$> ((,) <$> arbitrary <*> arbText) + ] + where + arbText = T.pack <$> printable -- XXX: we restrict this generator to single datasets to we don't have -- to worry about duplicates. @@ -28,8 +31,9 @@ encodeDecodeJSON :: TestTree encodeDecodeJSON = testProperty "encoding and decoding are inverse" prop - where prop :: BenchResult -> Bool - prop brs = decodeBenchResults (encodeBenchResults [brs]) == Right [brs] + where + prop :: BenchResult -> Bool + prop brs = decodeBenchResults (encodeBenchResults [brs]) == Right [brs] tests :: TestTree tests = testGroup "Futhark.BenchTests" [encodeDecodeJSON] diff --git a/unittests/Futhark/IR/Mem/IxFun/Alg.hs b/unittests/Futhark/IR/Mem/IxFun/Alg.hs --- a/unittests/Futhark/IR/Mem/IxFun/Alg.hs +++ b/unittests/Futhark/IR/Mem/IxFun/Alg.hs @@ -1,40 +1,48 @@ -- | A simple index operation representation. Every operation corresponds to a -- constructor. module Futhark.IR.Mem.IxFun.Alg - ( IxFun(..) - , iota - , offsetIndex - , permute - , rotate - , reshape - , slice - , rebase - , shape - , index + ( IxFun (..), + iota, + offsetIndex, + permute, + rotate, + reshape, + slice, + rebase, + shape, + index, ) where -import Prelude hiding (mod) - -import Futhark.IR.Syntax - (ShapeChange, DimChange(..), Slice, sliceDims, DimIndex(..), unitSlice) -import Futhark.IR.Prop import Futhark.IR.Pretty () +import Futhark.IR.Prop +import Futhark.IR.Syntax + ( DimChange (..), + DimIndex (..), + ShapeChange, + Slice, + sliceDims, + unitSlice, + ) import Futhark.Util.IntegralExp import Futhark.Util.Pretty +import Prelude hiding (mod) type Shape num = [num] + type Indices num = [num] + type Permutation = [Int] -data IxFun num = Direct (Shape num) - | Permute (IxFun num) Permutation - | Rotate (IxFun num) (Indices num) - | Index (IxFun num) (Slice num) - | Reshape (IxFun num) (ShapeChange num) - | OffsetIndex (IxFun num) num - | Rebase (IxFun num) (IxFun num) - deriving (Eq, Show) +data IxFun num + = Direct (Shape num) + | Permute (IxFun num) Permutation + | Rotate (IxFun num) (Indices num) + | Index (IxFun num) (Slice num) + | Reshape (IxFun num) (ShapeChange num) + | OffsetIndex (IxFun num) num + | Rebase (IxFun num) (IxFun num) + deriving (Eq, Show) instance Pretty num => Pretty (IxFun num) where ppr (Direct dims) = @@ -43,14 +51,13 @@ ppr (Rotate fun offsets) = ppr fun <> brackets (commasep $ map ((text "+" <>) . ppr) offsets) ppr (Index fun is) = ppr fun <> brackets (commasep $ map ppr is) ppr (Reshape fun oldshape) = - ppr fun <> text "->reshape" <> - parens (commasep (map ppr oldshape)) + ppr fun <> text "->reshape" + <> parens (commasep (map ppr oldshape)) ppr (OffsetIndex fun i) = ppr fun <> text "->offset_index" <> parens (ppr i) ppr (Rebase new_base fun) = text "rebase(" <> ppr new_base <> text ", " <> ppr fun <> text ")" - iota :: Shape num -> IxFun num iota = Direct @@ -72,8 +79,10 @@ reshape :: IxFun num -> ShapeChange num -> IxFun num reshape = Reshape -shape :: IntegralExp num => - IxFun num -> Shape num +shape :: + IntegralExp num => + IxFun num -> + Shape num shape (Direct dims) = dims shape (Permute ixfun perm) = @@ -89,46 +98,53 @@ shape (Rebase _ ixfun) = shape ixfun -index :: (IntegralExp num, Eq num) => - IxFun num -> Indices num -> num +index :: + (IntegralExp num, Eq num) => + IxFun num -> + Indices num -> + num index (Direct dims) is = sum $ zipWith (*) is slicesizes - where slicesizes = drop 1 $ sliceSizes dims + where + slicesizes = drop 1 $ sliceSizes dims index (Permute fun perm) is_new = index fun is_old - where is_old = rearrangeShape (rearrangeInverse perm) is_new + where + is_old = rearrangeShape (rearrangeInverse perm) is_new index (Rotate fun offsets) is = index fun $ zipWith mod (zipWith (+) is offsets) dims - where dims = shape fun + where + dims = shape fun index (Index fun js) is = index fun (adjust js is) - where adjust (DimFix j:js') is' = j : adjust js' is' - adjust (DimSlice j _ s:js') (i:is') = j + i * s : adjust js' is' - adjust _ _ = [] + where + adjust (DimFix j : js') is' = j : adjust js' is' + adjust (DimSlice j _ s : js') (i : is') = j + i * s : adjust js' is' + adjust _ _ = [] index (Reshape fun newshape) is = let new_indices = reshapeIndex (shape fun) (newDims newshape) is - in index fun new_indices + in index fun new_indices index (OffsetIndex fun i) is = case shape fun of d : ds -> - index (Index fun (DimSlice i (d-i) 1 : map (unitSlice 0) ds)) is + index (Index fun (DimSlice i (d - i) 1 : map (unitSlice 0) ds)) is [] -> error "index: OffsetIndex: underlying index function has rank zero" index (Rebase new_base fun) is = let fun' = case fun of - Direct old_shape -> - if old_shape == shape new_base - then new_base - else reshape new_base $ map DimCoercion old_shape - Permute ixfun perm -> - permute (rebase new_base ixfun) perm - Rotate ixfun offsets -> - rotate (rebase new_base ixfun) offsets - Index ixfun iis -> - slice (rebase new_base ixfun) iis - Reshape ixfun new_shape -> - reshape (rebase new_base ixfun) new_shape - OffsetIndex ixfun s -> - offsetIndex (rebase new_base ixfun) s - r@Rebase{} -> - r - in index fun' is + Direct old_shape -> + if old_shape == shape new_base + then new_base + else reshape new_base $ map DimCoercion old_shape + Permute ixfun perm -> + permute (rebase new_base ixfun) perm + Rotate ixfun offsets -> + rotate (rebase new_base ixfun) offsets + Index ixfun iis -> + slice (rebase new_base ixfun) iis + Reshape ixfun new_shape -> + reshape (rebase new_base ixfun) new_shape + OffsetIndex ixfun s -> + offsetIndex (rebase new_base ixfun) s + r@Rebase {} -> + r + in index fun' is diff --git a/unittests/Futhark/IR/Mem/IxFunTests.hs b/unittests/Futhark/IR/Mem/IxFunTests.hs --- a/unittests/Futhark/IR/Mem/IxFunTests.hs +++ b/unittests/Futhark/IR/Mem/IxFunTests.hs @@ -1,49 +1,46 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} + module Futhark.IR.Mem.IxFunTests - ( tests + ( tests, ) where -import Prelude hiding (span) -import qualified Prelude as P import qualified Data.List as DL - -import Test.Tasty -import Test.Tasty.HUnit - -import Futhark.IR.Syntax -import Futhark.IR.Syntax.Core() -import qualified Futhark.Util.Pretty as PR -import qualified Futhark.Util.IntegralExp as IE import qualified Futhark.IR.Mem.IxFun as IxFunLMAD import qualified Futhark.IR.Mem.IxFun.Alg as IxFunAlg - -import qualified Futhark.IR.Mem.IxFunWrapper as IxFunWrap import Futhark.IR.Mem.IxFunWrapper - +import qualified Futhark.IR.Mem.IxFunWrapper as IxFunWrap +import Futhark.IR.Syntax +import Futhark.IR.Syntax.Core () +import qualified Futhark.Util.IntegralExp as IE +import qualified Futhark.Util.Pretty as PR +import Test.Tasty +import Test.Tasty.HUnit +import Prelude hiding (span) +import qualified Prelude as P instance IE.IntegralExp Int where quot = P.quot - rem = P.rem - div = P.div - mod = P.mod - sgn = Just . P.signum - - fromInt8 = fromInteger . toInteger - fromInt16 = fromInteger . toInteger - fromInt32 = fromInteger . toInteger - fromInt64 = fromInteger . toInteger + rem = P.rem + div = P.div + mod = P.mod + sgn = Just . P.signum allPoints :: [Int] -> [[Int]] allPoints dims = - let total = product dims - strides = drop 1 $ DL.reverse $ scanl (*) 1 $ DL.reverse dims - in map (unflatInd strides) [0..total-1] - where unflatInd :: [Int] -> Int -> [Int] - unflatInd strides x = fst $ - foldl (\(res, acc) span -> - (res ++ [acc `P.div` span], acc `P.mod` span)) - ([], x) strides + let total = product dims + strides = drop 1 $ DL.reverse $ scanl (*) 1 $ DL.reverse dims + in map (unflatInd strides) [0 .. total -1] + where + unflatInd :: [Int] -> Int -> [Int] + unflatInd strides x = + fst $ + foldl + ( \(res, acc) span -> + (res ++ [acc `P.div` span], acc `P.mod` span) + ) + ([], x) + strides compareIxFuns :: IxFunLMAD.IxFun Int -> IxFunAlg.IxFun Int -> Assertion compareIxFuns ixfunLMAD ixfunAlg = @@ -52,15 +49,29 @@ points = allPoints lmadShape resLMAD = map (IxFunLMAD.index ixfunLMAD) points resAlg = map (IxFunAlg.index ixfunAlg) points - errorMessage = "lmad ixfun: " ++ PR.pretty ixfunLMAD ++ "\n" ++ - "alg ixfun: " ++ PR.pretty ixfunAlg ++ "\n" ++ - "lmad shape: " ++ show lmadShape ++ "\n" ++ - "alg shape: " ++ show algShape ++ "\n" ++ - "lmad points length: " ++ show (length resLMAD) ++ "\n" ++ - "alg points length: " ++ show (length resAlg) ++ "\n" ++ - "lmad points: " ++ show resLMAD ++ "\n" ++ - "alg points: " ++ show resAlg - in (lmadShape == algShape && resLMAD == resAlg) @? errorMessage + errorMessage = + "lmad ixfun: " ++ PR.pretty ixfunLMAD ++ "\n" + ++ "alg ixfun: " + ++ PR.pretty ixfunAlg + ++ "\n" + ++ "lmad shape: " + ++ show lmadShape + ++ "\n" + ++ "alg shape: " + ++ show algShape + ++ "\n" + ++ "lmad points length: " + ++ show (length resLMAD) + ++ "\n" + ++ "alg points length: " + ++ show (length resAlg) + ++ "\n" + ++ "lmad points: " + ++ show resLMAD + ++ "\n" + ++ "alg points: " + ++ show resAlg + in (lmadShape == algShape && resLMAD == resAlg) @? errorMessage compareOps :: IxFunWrap.IxFun Int -> Assertion compareOps (ixfunLMAD, ixfunAlg) = compareIxFuns ixfunLMAD ixfunAlg @@ -68,254 +79,317 @@ -- XXX: Clean this up. n :: Int n = 19 -slice3 :: [DimIndex Int] -slice3 = [ DimSlice 2 (n `P.div` 3) 3 - , DimFix (n `P.div` 2) - , DimSlice 1 (n `P.div` 2) 2 - ] +slice3 :: [DimIndex Int] +slice3 = + [ DimSlice 2 (n `P.div` 3) 3, + DimFix (n `P.div` 2), + DimSlice 1 (n `P.div` 2) 2 + ] -- Actual tests. tests :: TestTree -tests = testGroup "IxFunTests" - $ concat - [ test_iota - , test_slice_iota - , test_reshape_slice_iota1 - , test_permute_slice_iota - , test_rotate_rotate_permute_slice_iota - , test_slice_rotate_permute_slice_iota1 - , test_slice_rotate_permute_slice_iota2 - , test_slice_rotate_permute_slice_iota3 - , test_permute_rotate_slice_permute_slice_iota - , test_reshape_rotate_iota - , test_reshape_permute_iota - , test_reshape_slice_iota2 - , test_reshape_slice_iota3 - , test_complex1 - , test_complex2 - , test_rebase1 - , test_rebase2 - , test_rebase3 - , test_rebase4_5 - ] +tests = + testGroup "IxFunTests" $ + concat + [ test_iota, + test_slice_iota, + test_reshape_slice_iota1, + test_permute_slice_iota, + test_rotate_rotate_permute_slice_iota, + test_slice_rotate_permute_slice_iota1, + test_slice_rotate_permute_slice_iota2, + test_slice_rotate_permute_slice_iota3, + test_permute_rotate_slice_permute_slice_iota, + test_reshape_rotate_iota, + test_reshape_permute_iota, + test_reshape_slice_iota2, + test_reshape_slice_iota3, + test_complex1, + test_complex2, + test_rebase1, + test_rebase2, + test_rebase3, + test_rebase4_5 + ] singleton :: TestTree -> [TestTree] singleton = (: []) test_iota :: [TestTree] -test_iota = singleton $ testCase "iota" $ compareOps $ - iota [n] +test_iota = + singleton $ + testCase "iota" $ + compareOps $ + iota [n] test_slice_iota :: [TestTree] -test_slice_iota = singleton $ testCase "slice . iota" $ compareOps $ - slice (iota [n, n, n]) slice3 +test_slice_iota = + singleton $ + testCase "slice . iota" $ + compareOps $ + slice (iota [n, n, n]) slice3 test_reshape_slice_iota1 :: [TestTree] -test_reshape_slice_iota1 = singleton $ testCase "reshape . slice . iota 1" $ compareOps $ - reshape (slice (iota [n, n, n]) slice3) - [DimNew (n `P.div` 2), DimNew (n `P.div` 3)] +test_reshape_slice_iota1 = + singleton $ + testCase "reshape . slice . iota 1" $ + compareOps $ + reshape + (slice (iota [n, n, n]) slice3) + [DimNew (n `P.div` 2), DimNew (n `P.div` 3)] test_permute_slice_iota :: [TestTree] -test_permute_slice_iota = singleton $ testCase "permute . slice . iota" $ compareOps $ - permute (slice (iota [n, n, n]) slice3) [1, 0] +test_permute_slice_iota = + singleton $ + testCase "permute . slice . iota" $ + compareOps $ + permute (slice (iota [n, n, n]) slice3) [1, 0] test_rotate_rotate_permute_slice_iota :: [TestTree] test_rotate_rotate_permute_slice_iota = - singleton $ testCase "rotate . rotate . permute . slice . iota" $ compareOps $ - let ixfun = permute (slice (iota [n, n, n]) slice3) [1, 0] - in rotate (rotate ixfun [2, 1]) [1, 2] + singleton $ + testCase "rotate . rotate . permute . slice . iota" $ + compareOps $ + let ixfun = permute (slice (iota [n, n, n]) slice3) [1, 0] + in rotate (rotate ixfun [2, 1]) [1, 2] test_slice_rotate_permute_slice_iota1 :: [TestTree] test_slice_rotate_permute_slice_iota1 = - singleton $ testCase "slice . rotate . permute . slice . iota 1" $ compareOps $ - let slice2 = [ DimSlice 0 n 1 - , DimSlice 1 (n `P.div` 2) 2 - , DimSlice 0 n 1 - ] - slice13 = [ DimSlice 2 (n `P.div` 3) 3 - , DimSlice 0 (n `P.div` 2) 1 - , DimSlice 1 (n `P.div` 2) 2 - ] - ixfun = permute (slice (iota [n, n, n]) slice2) [2, 1, 0] - ixfun' = slice (rotate ixfun [3, 1, 2]) slice13 - in ixfun' + singleton $ + testCase "slice . rotate . permute . slice . iota 1" $ + compareOps $ + let slice2 = + [ DimSlice 0 n 1, + DimSlice 1 (n `P.div` 2) 2, + DimSlice 0 n 1 + ] + slice13 = + [ DimSlice 2 (n `P.div` 3) 3, + DimSlice 0 (n `P.div` 2) 1, + DimSlice 1 (n `P.div` 2) 2 + ] + ixfun = permute (slice (iota [n, n, n]) slice2) [2, 1, 0] + ixfun' = slice (rotate ixfun [3, 1, 2]) slice13 + in ixfun' test_slice_rotate_permute_slice_iota2 :: [TestTree] test_slice_rotate_permute_slice_iota2 = - singleton $ testCase "slice . rotate . permute . slice . iota 2" $ compareOps $ - let slice2 = [ DimSlice 0 (n `P.div` 2) 1 - , DimFix (n `P.div` 2) - , DimSlice 0 (n `P.div` 3) 1 - ] - slice13 = [ DimSlice 2 (n `P.div` 3) 3 - , DimSlice 0 n 1 - , DimSlice 1 (n `P.div` 2) 2 - ] - ixfun = permute (slice (iota [n, n, n]) slice13) [2, 1, 0] - ixfun' = slice (rotate ixfun [3, 1, 2]) slice2 - in ixfun' + singleton $ + testCase "slice . rotate . permute . slice . iota 2" $ + compareOps $ + let slice2 = + [ DimSlice 0 (n `P.div` 2) 1, + DimFix (n `P.div` 2), + DimSlice 0 (n `P.div` 3) 1 + ] + slice13 = + [ DimSlice 2 (n `P.div` 3) 3, + DimSlice 0 n 1, + DimSlice 1 (n `P.div` 2) 2 + ] + ixfun = permute (slice (iota [n, n, n]) slice13) [2, 1, 0] + ixfun' = slice (rotate ixfun [3, 1, 2]) slice2 + in ixfun' test_slice_rotate_permute_slice_iota3 :: [TestTree] test_slice_rotate_permute_slice_iota3 = - singleton $ testCase "slice . rotate . permute . slice . iota 3" $ compareOps $ - -- full-slice of (-1) stride - let ixfun = permute (slice (iota [n, n, n]) slice3) [1, 0] - ixfun' = rotate ixfun [2, 1] + singleton $ + testCase "slice . rotate . permute . slice . iota 3" $ + compareOps $ + -- full-slice of (-1) stride + let ixfun = permute (slice (iota [n, n, n]) slice3) [1, 0] + ixfun' = rotate ixfun [2, 1] - (n1, m1) = case IxFunLMAD.shape (fst ixfun') of - [a, b] -> (a, b) - _ -> error "expecting 2 dimensions at this point!" - negslice = [DimSlice 0 n1 1, DimSlice (m1 - 1) m1 (-1)] - ixfun'' = rotate (slice ixfun' negslice) [1,2] - in ixfun'' + (n1, m1) = case IxFunLMAD.shape (fst ixfun') of + [a, b] -> (a, b) + _ -> error "expecting 2 dimensions at this point!" + negslice = [DimSlice 0 n1 1, DimSlice (m1 - 1) m1 (-1)] + ixfun'' = rotate (slice ixfun' negslice) [1, 2] + in ixfun'' test_permute_rotate_slice_permute_slice_iota :: [TestTree] test_permute_rotate_slice_permute_slice_iota = - singleton $ testCase "permute . rotate . slice . permute . slice . iota" $ compareOps $ - -- contiguousness - let slice33 = [ DimFix (n `P.div` 2) - , DimSlice (n - 1) (n `P.div` 3) (-1) - , DimSlice 0 n 1 - ] - ixfun = permute (slice (iota [n, n, n]) slice33) [1, 0] - m = n `P.div` 3 - slice1 = [DimSlice (n - 1) n (-1), DimSlice 2 (m - 2) 1] - ixfun' = permute (rotate (slice ixfun slice1) [1, 2]) [1, 0] - in ixfun' + singleton $ + testCase "permute . rotate . slice . permute . slice . iota" $ + compareOps $ + -- contiguousness + let slice33 = + [ DimFix (n `P.div` 2), + DimSlice (n - 1) (n `P.div` 3) (-1), + DimSlice 0 n 1 + ] + ixfun = permute (slice (iota [n, n, n]) slice33) [1, 0] + m = n `P.div` 3 + slice1 = [DimSlice (n - 1) n (-1), DimSlice 2 (m - 2) 1] + ixfun' = permute (rotate (slice ixfun slice1) [1, 2]) [1, 0] + in ixfun' test_reshape_rotate_iota :: [TestTree] test_reshape_rotate_iota = -- negative reshape test - singleton $ testCase "reshape . rotate . iota" $ compareOps $ - let newdims = [DimNew (n * n), DimCoercion n] - in reshape (rotate (iota [n, n, n]) [1, 0, 0]) newdims + singleton $ + testCase "reshape . rotate . iota" $ + compareOps $ + let newdims = [DimNew (n * n), DimCoercion n] + in reshape (rotate (iota [n, n, n]) [1, 0, 0]) newdims test_reshape_permute_iota :: [TestTree] test_reshape_permute_iota = -- negative reshape test - singleton $ testCase "reshape . permute . iota" $ compareOps $ - let newdims = [DimNew (n * n), DimCoercion n] - in reshape (permute (iota [n, n, n]) [1, 2, 0]) newdims + singleton $ + testCase "reshape . permute . iota" $ + compareOps $ + let newdims = [DimNew (n * n), DimCoercion n] + in reshape (permute (iota [n, n, n]) [1, 2, 0]) newdims test_reshape_slice_iota2 :: [TestTree] test_reshape_slice_iota2 = -- negative reshape test - singleton $ testCase "reshape . slice . iota 2" $ compareOps $ - let newdims = [DimNew (n*n), DimCoercion n] - slc = [ DimFix (n `P.div` 2) - , DimSlice (n-1) n (-1) - , DimSlice 0 n 1 - , DimSlice (n-1) n (-1) - ] - in reshape (slice (iota [n, n, n, n]) slc) newdims + singleton $ + testCase "reshape . slice . iota 2" $ + compareOps $ + let newdims = [DimNew (n * n), DimCoercion n] + slc = + [ DimFix (n `P.div` 2), + DimSlice (n -1) n (-1), + DimSlice 0 n 1, + DimSlice (n -1) n (-1) + ] + in reshape (slice (iota [n, n, n, n]) slc) newdims test_reshape_slice_iota3 :: [TestTree] test_reshape_slice_iota3 = -- negative reshape test - singleton $ testCase "reshape . slice . iota 3" $ compareOps $ - let newdims = [DimNew (n*n), DimCoercion n] - slc = [ DimFix (n `P.div` 2) - , DimSlice 0 n 1 - , DimSlice 0 (n `P.div` 2) 1 - , DimSlice 0 n 1 - ] - in reshape (slice (iota [n, n, n, n]) slc) newdims + singleton $ + testCase "reshape . slice . iota 3" $ + compareOps $ + let newdims = [DimNew (n * n), DimCoercion n] + slc = + [ DimFix (n `P.div` 2), + DimSlice 0 n 1, + DimSlice 0 (n `P.div` 2) 1, + DimSlice 0 n 1 + ] + in reshape (slice (iota [n, n, n, n]) slc) newdims test_complex1 :: [TestTree] test_complex1 = - singleton $ testCase "reshape . permute . rotate . slice . permute . slice . iota 1" $ compareOps $ - let newdims = [ DimCoercion n - , DimCoercion n - , DimNew n - , DimCoercion ((n `P.div` 3) - 2) - ] - slice33 = [ DimSlice (n-1) (n `P.div` 3) (-1) - , DimSlice (n-1) n (-1) - , DimSlice (n-1) n (-1) - , DimSlice 0 n 1 - ] - ixfun = permute (slice (iota [n, n, n, n, n]) slice33) [3, 1, 2, 0] - m = n `P.div` 3 - slice1 = [DimSlice 0 n 1, DimSlice (n-1) n (-1), DimSlice (n-1) n (-1), DimSlice 1 (m-2) (-1)] - ixfun' = reshape (rotate (slice ixfun slice1) [1, 2, 3, 4]) newdims - in ixfun' + singleton $ + testCase "reshape . permute . rotate . slice . permute . slice . iota 1" $ + compareOps $ + let newdims = + [ DimCoercion n, + DimCoercion n, + DimNew n, + DimCoercion ((n `P.div` 3) - 2) + ] + slice33 = + [ DimSlice (n -1) (n `P.div` 3) (-1), + DimSlice (n -1) n (-1), + DimSlice (n -1) n (-1), + DimSlice 0 n 1 + ] + ixfun = permute (slice (iota [n, n, n, n, n]) slice33) [3, 1, 2, 0] + m = n `P.div` 3 + slice1 = [DimSlice 0 n 1, DimSlice (n -1) n (-1), DimSlice (n -1) n (-1), DimSlice 1 (m -2) (-1)] + ixfun' = reshape (rotate (slice ixfun slice1) [1, 2, 3, 4]) newdims + in ixfun' test_complex2 :: [TestTree] test_complex2 = - singleton $ testCase "reshape . permute . rotate . slice . permute . slice . iota 2" $ compareOps $ - let newdims = [ DimCoercion n - , DimNew (n*n) - , DimCoercion ((n `P.div` 3) - 2)] - slc2 = [ DimFix (n `P.div` 2) - , DimSlice (n-1) (n `P.div` 3) (-1) - , DimSlice (n-1) n (-1) - , DimSlice (n-1) n (-1) - , DimSlice 0 n 1 - ] - ixfun = permute (slice (iota [n, n, n, n, n]) slc2) [3, 1, 2, 0] - m = n `P.div` 3 - slice1 = [DimSlice 0 n 1, DimSlice (n-1) n (-1), DimSlice (n-1) n (-1), DimSlice 1 (m-2) (-1)] - ixfun' = reshape (rotate (slice ixfun slice1) [1, 0, 0, 2]) newdims - in ixfun' + singleton $ + testCase "reshape . permute . rotate . slice . permute . slice . iota 2" $ + compareOps $ + let newdims = + [ DimCoercion n, + DimNew (n * n), + DimCoercion ((n `P.div` 3) - 2) + ] + slc2 = + [ DimFix (n `P.div` 2), + DimSlice (n -1) (n `P.div` 3) (-1), + DimSlice (n -1) n (-1), + DimSlice (n -1) n (-1), + DimSlice 0 n 1 + ] + ixfun = permute (slice (iota [n, n, n, n, n]) slc2) [3, 1, 2, 0] + m = n `P.div` 3 + slice1 = [DimSlice 0 n 1, DimSlice (n -1) n (-1), DimSlice (n -1) n (-1), DimSlice 1 (m -2) (-1)] + ixfun' = reshape (rotate (slice ixfun slice1) [1, 0, 0, 2]) newdims + in ixfun' test_rebase1 :: [TestTree] test_rebase1 = - singleton $ testCase "rebase 1" $ compareOps $ - let slice_base = [ DimFix (n `P.div` 2) - , DimSlice 2 (n-2) 1 - , DimSlice 3 (n-3) 1 - ] - ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] - ixfn_orig = rotate (permute (iota [n-3, n-2]) [1, 0]) [1, 2] - ixfn_rebase = rebase ixfn_base ixfn_orig - in ixfn_rebase + singleton $ + testCase "rebase 1" $ + compareOps $ + let slice_base = + [ DimFix (n `P.div` 2), + DimSlice 2 (n -2) 1, + DimSlice 3 (n -3) 1 + ] + ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] + ixfn_orig = rotate (permute (iota [n -3, n -2]) [1, 0]) [1, 2] + ixfn_rebase = rebase ixfn_base ixfn_orig + in ixfn_rebase test_rebase2 :: [TestTree] test_rebase2 = - singleton $ testCase "rebase 2" $ compareOps $ - let slice_base = [ DimFix (n `P.div` 2) - , DimSlice (n-1) (n-2) (-1) - , DimSlice (n-1) (n-3) (-1) - ] - slice_orig = [ DimSlice (n-4) (n-3) (-1) - , DimSlice (n-3) (n-2) (-1) - ] - ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] - ixfn_orig = rotate (permute (slice (iota [n-3, n-2]) slice_orig) [1, 0]) [1, 2] - ixfn_rebase = rebase ixfn_base ixfn_orig - in ixfn_rebase + singleton $ + testCase "rebase 2" $ + compareOps $ + let slice_base = + [ DimFix (n `P.div` 2), + DimSlice (n -1) (n -2) (-1), + DimSlice (n -1) (n -3) (-1) + ] + slice_orig = + [ DimSlice (n -4) (n -3) (-1), + DimSlice (n -3) (n -2) (-1) + ] + ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] + ixfn_orig = rotate (permute (slice (iota [n -3, n -2]) slice_orig) [1, 0]) [1, 2] + ixfn_rebase = rebase ixfn_base ixfn_orig + in ixfn_rebase test_rebase3 :: [TestTree] test_rebase3 = - singleton $ testCase "rebase full orig but not monotonic" $ compareOps $ - let n2 = (n-2) `P.div` 3 - n3 = (n-3) `P.div` 2 - slice_base = [ DimFix (n `P.div` 2) - , DimSlice (n-1) n2 (-3) - , DimSlice (n-1) n3 (-2) - ] - slice_orig = [ DimSlice (n3-1) n3 (-1) - , DimSlice (n2-1) n2 (-1) - ] - ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] - ixfn_orig = rotate (permute (slice (iota [n3, n2]) slice_orig) [1, 0]) [1, 2] - ixfn_rebase = rebase ixfn_base ixfn_orig - in ixfn_rebase + singleton $ + testCase "rebase full orig but not monotonic" $ + compareOps $ + let n2 = (n -2) `P.div` 3 + n3 = (n -3) `P.div` 2 + slice_base = + [ DimFix (n `P.div` 2), + DimSlice (n -1) n2 (-3), + DimSlice (n -1) n3 (-2) + ] + slice_orig = + [ DimSlice (n3 -1) n3 (-1), + DimSlice (n2 -1) n2 (-1) + ] + ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] + ixfn_orig = rotate (permute (slice (iota [n3, n2]) slice_orig) [1, 0]) [1, 2] + ixfn_rebase = rebase ixfn_base ixfn_orig + in ixfn_rebase test_rebase4_5 :: [TestTree] test_rebase4_5 = - let n2 = (n-2) `P.div` 3 - n3 = (n-3) `P.div` 2 - slice_base = [ DimFix (n `P.div` 2) - , DimSlice (n-1) n2 (-3) - , DimSlice 3 n3 2 - ] - slice_orig = [ DimSlice (n3-1) n3 (-1) - , DimSlice 0 n2 1 - ] + let n2 = (n -2) `P.div` 3 + n3 = (n -3) `P.div` 2 + slice_base = + [ DimFix (n `P.div` 2), + DimSlice (n -1) n2 (-3), + DimSlice 3 n3 2 + ] + slice_orig = + [ DimSlice (n3 -1) n3 (-1), + DimSlice 0 n2 1 + ] ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] ixfn_orig = rotate (permute (slice (iota [n3, n2]) slice_orig) [1, 0]) [1, 2] - in [ testCase "rebase mixed monotonicities" $ compareOps $ - rebase ixfn_base ixfn_orig - ] + in [ testCase "rebase mixed monotonicities" $ + compareOps $ + rebase ixfn_base ixfn_orig + ] diff --git a/unittests/Futhark/IR/Mem/IxFunWrapper.hs b/unittests/Futhark/IR/Mem/IxFunWrapper.hs --- a/unittests/Futhark/IR/Mem/IxFunWrapper.hs +++ b/unittests/Futhark/IR/Mem/IxFunWrapper.hs @@ -1,48 +1,66 @@ -- | Perform index function operations in both algebraic and LMAD -- representations. module Futhark.IR.Mem.IxFunWrapper - ( IxFun - , iota - , permute - , rotate - , reshape - , slice - , rebase + ( IxFun, + iota, + permute, + rotate, + reshape, + slice, + rebase, ) where -import Futhark.Util.IntegralExp -import Futhark.IR.Syntax (ShapeChange, Slice) import qualified Futhark.IR.Mem.IxFun as I import qualified Futhark.IR.Mem.IxFun.Alg as IA - +import Futhark.IR.Syntax (ShapeChange, Slice) +import Futhark.Util.IntegralExp type Shape num = [num] + type Indices num = [num] + type Permutation = [Int] type IxFun num = (I.IxFun num, IA.IxFun num) -iota :: IntegralExp num => - Shape num -> IxFun num +iota :: + IntegralExp num => + Shape num -> + IxFun num iota x = (I.iota x, IA.iota x) -permute :: IntegralExp num => - IxFun num -> Permutation -> IxFun num +permute :: + IntegralExp num => + IxFun num -> + Permutation -> + IxFun num permute (l, a) x = (I.permute l x, IA.permute a x) -rotate :: (Eq num, IntegralExp num) => - IxFun num -> Indices num -> IxFun num +rotate :: + (Eq num, IntegralExp num) => + IxFun num -> + Indices num -> + IxFun num rotate (l, a) x = (I.rotate l x, IA.rotate a x) -reshape :: (Eq num, IntegralExp num) => - IxFun num -> ShapeChange num -> IxFun num +reshape :: + (Eq num, IntegralExp num) => + IxFun num -> + ShapeChange num -> + IxFun num reshape (l, a) x = (I.reshape l x, IA.reshape a x) -slice :: (Eq num, IntegralExp num) => - IxFun num -> Slice num -> IxFun num +slice :: + (Eq num, IntegralExp num) => + IxFun num -> + Slice num -> + IxFun num slice (l, a) x = (I.slice l x, IA.slice a x) -rebase :: (Eq num, IntegralExp num) => - IxFun num -> IxFun num -> IxFun num +rebase :: + (Eq num, IntegralExp num) => + IxFun num -> + IxFun num -> + IxFun num rebase (l, a) (l1, a1) = (I.rebase l l1, IA.rebase a a1) diff --git a/unittests/Futhark/IR/PrimitiveTests.hs b/unittests/Futhark/IR/PrimitiveTests.hs --- a/unittests/Futhark/IR/PrimitiveTests.hs +++ b/unittests/Futhark/IR/PrimitiveTests.hs @@ -1,54 +1,77 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} {-# OPTIONS_GHC -fno-warn-orphans #-} + module Futhark.IR.PrimitiveTests - ( tests - , arbitraryPrimValOfType - ) - where + ( tests, + arbitraryPrimValOfType, + ) +where import Control.Applicative - +import Futhark.IR.Primitive +import Language.SexpGrammar import Test.QuickCheck import Test.Tasty import Test.Tasty.HUnit - import Prelude -import Futhark.IR.Primitive - tests :: TestTree -tests = testGroup "PrimitiveTests" propPrimValuesHaveRightType +tests = testGroup "PrimitiveTests" [propPrimValuesHaveRightType, doUnOpTests] -propPrimValuesHaveRightType :: [TestTree] -propPrimValuesHaveRightType = [ testCase (show t ++ " has blank of right type") $ - primValueType (blankPrimValue t) @?= t - | t <- [minBound..maxBound] - ] +propPrimValuesHaveRightType :: TestTree +propPrimValuesHaveRightType = + testGroup + "propPrimValuesHaveRightTypes" + [ testCase (show t ++ " has blank of right type") $ + primValueType (blankPrimValue t) @?= t + | t <- [minBound .. maxBound] + ] +doUnOpTests :: TestTree +doUnOpTests = + testGroup + "doUnOp" + [ testCase "not" $ + let unop = decode @UnOp "(complement i32)" + val = decode @PrimValue "42i32" + res = doUnOp <$> unop <*> val + in res @?= Right (Just (IntValue (Int32Value (-43)))) + ] + instance Arbitrary IntType where - arbitrary = elements [minBound..maxBound] + arbitrary = elements [minBound .. maxBound] instance Arbitrary FloatType where - arbitrary = elements [minBound..maxBound] + arbitrary = elements [minBound .. maxBound] instance Arbitrary PrimType where - arbitrary = elements [minBound..maxBound] + arbitrary = elements [minBound .. maxBound] instance Arbitrary IntValue where - arbitrary = oneof [ Int8Value <$> arbitrary - , Int16Value <$> arbitrary - , Int32Value <$> arbitrary - , Int64Value <$> arbitrary ] + arbitrary = + oneof + [ Int8Value <$> arbitrary, + Int16Value <$> arbitrary, + Int32Value <$> arbitrary, + Int64Value <$> arbitrary + ] instance Arbitrary FloatValue where - arbitrary = oneof [ Float32Value <$> arbitrary - , Float64Value <$> arbitrary ] + arbitrary = + oneof + [ Float32Value <$> arbitrary, + Float64Value <$> arbitrary + ] instance Arbitrary PrimValue where - arbitrary = oneof [ IntValue <$> arbitrary - , FloatValue <$> arbitrary - , BoolValue <$> arbitrary - , pure Checked - ] + arbitrary = + oneof + [ IntValue <$> arbitrary, + FloatValue <$> arbitrary, + BoolValue <$> arbitrary, + pure Checked + ] arbitraryPrimValOfType :: PrimType -> Gen PrimValue arbitraryPrimValOfType (IntType Int8) = IntValue . Int8Value <$> arbitrary diff --git a/unittests/Futhark/IR/Prop/RearrangeTests.hs b/unittests/Futhark/IR/Prop/RearrangeTests.hs --- a/unittests/Futhark/IR/Prop/RearrangeTests.hs +++ b/unittests/Futhark/IR/Prop/RearrangeTests.hs @@ -1,55 +1,54 @@ -module Futhark.IR.Prop.RearrangeTests - ( tests ) - where +module Futhark.IR.Prop.RearrangeTests (tests) where import Control.Applicative - +import Futhark.IR.Prop.Rearrange import Test.Tasty import Test.Tasty.HUnit import Test.Tasty.QuickCheck - import Prelude -import Futhark.IR.Prop.Rearrange - tests :: TestTree -tests = testGroup "RearrangeTests" $ - isMapTransposeTests ++ - [isMapTransposeProp] +tests = + testGroup "RearrangeTests" $ + isMapTransposeTests + ++ [isMapTransposeProp] isMapTransposeTests :: [TestTree] isMapTransposeTests = [ testCase (unwords ["isMapTranspose", show perm, "==", show dres]) $ - isMapTranspose perm @?= dres - | (perm, dres) <- [ ([0,1,4,5,2,3], Just (2,2,2)) - , ([1,0,4,5,2,3], Nothing) - , ([1,0], Just (0, 1, 1)) - , ([0,2,1], Just (1, 1, 1)) - , ([0,1,2], Nothing) - , ([1,0,2], Nothing) - ] + isMapTranspose perm @?= dres + | (perm, dres) <- + [ ([0, 1, 4, 5, 2, 3], Just (2, 2, 2)), + ([1, 0, 4, 5, 2, 3], Nothing), + ([1, 0], Just (0, 1, 1)), + ([0, 2, 1], Just (1, 1, 1)), + ([0, 1, 2], Nothing), + ([1, 0, 2], Nothing) + ] ] newtype Permutation = Permutation [Int] - deriving (Eq, Ord, Show) + deriving (Eq, Ord, Show) instance Arbitrary Permutation where arbitrary = do Positive n <- arbitrary - Permutation <$> shuffle [0..n-1] + Permutation <$> shuffle [0 .. n -1] isMapTransposeProp :: TestTree isMapTransposeProp = testProperty "isMapTranspose corresponds to a map of transpose" prop - where prop :: Permutation -> Bool - prop (Permutation perm) = - case isMapTranspose perm of - Nothing -> True - Just (r1, r2, r3) -> - and [r1 >= 0, - r2 > 0, - r3 > 0, - r1 + r2 + r3 == length perm, - let (mapped, notmapped) =splitAt r1 perm - (pretrans, posttrans) = splitAt r2 notmapped - in mapped ++ posttrans ++ pretrans == [0..length perm-1] - ] + where + prop :: Permutation -> Bool + prop (Permutation perm) = + case isMapTranspose perm of + Nothing -> True + Just (r1, r2, r3) -> + and + [ r1 >= 0, + r2 > 0, + r3 > 0, + r1 + r2 + r3 == length perm, + let (mapped, notmapped) = splitAt r1 perm + (pretrans, posttrans) = splitAt r2 notmapped + in mapped ++ posttrans ++ pretrans == [0 .. length perm -1] + ] diff --git a/unittests/Futhark/IR/Prop/ReshapeTests.hs b/unittests/Futhark/IR/Prop/ReshapeTests.hs --- a/unittests/Futhark/IR/Prop/ReshapeTests.hs +++ b/unittests/Futhark/IR/Prop/ReshapeTests.hs @@ -1,74 +1,74 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} + module Futhark.IR.Prop.ReshapeTests - ( tests - ) - where + ( tests, + ) +where import Control.Applicative - +import Futhark.IR.Prop.Constants +import Futhark.IR.Prop.Reshape +import Futhark.IR.Syntax import Test.Tasty import Test.Tasty.HUnit import Test.Tasty.QuickCheck - import Prelude -import Futhark.IR.Prop.Reshape -import Futhark.IR.Syntax -import Futhark.IR.Prop.Constants - tests :: TestTree -tests = testGroup "ReshapeTests" $ - fuseReshapeTests ++ - informReshapeTests ++ - reshapeOuterTests ++ - reshapeInnerTests ++ - [ fuseReshapeProp - , informReshapeProp - ] +tests = + testGroup "ReshapeTests" $ + fuseReshapeTests + ++ informReshapeTests + ++ reshapeOuterTests + ++ reshapeInnerTests + ++ [ fuseReshapeProp, + informReshapeProp + ] fuseReshapeTests :: [TestTree] fuseReshapeTests = [ testCase (unwords ["fuseReshape ", show d1, show d2]) $ - fuseReshape (d1 :: ShapeChange Int) d2 @?= dres -- type signature to avoid warning - | (d1, d2, dres) <- [ ([DimCoercion 1], [DimNew 1], [DimCoercion 1]) - , ([DimNew 1], [DimCoercion 1], [DimNew 1]) - , ([DimCoercion 1, DimNew 2], [DimNew 1, DimNew 2], [DimCoercion 1, DimNew 2]) - , ([DimNew 1, DimNew 2], [DimCoercion 1, DimNew 2], [DimNew 1, DimNew 2]) - ] + fuseReshape (d1 :: ShapeChange Int) d2 @?= dres -- type signature to avoid warning + | (d1, d2, dres) <- + [ ([DimCoercion 1], [DimNew 1], [DimCoercion 1]), + ([DimNew 1], [DimCoercion 1], [DimNew 1]), + ([DimCoercion 1, DimNew 2], [DimNew 1, DimNew 2], [DimCoercion 1, DimNew 2]), + ([DimNew 1, DimNew 2], [DimCoercion 1, DimNew 2], [DimNew 1, DimNew 2]) + ] ] informReshapeTests :: [TestTree] informReshapeTests = [ testCase (unwords ["informReshape ", show shape, show sc, show sc_res]) $ - informReshape (shape :: [Int]) sc @?= sc_res -- type signature to avoid warning - | (shape, sc, sc_res) <- - [ ([1, 2], [DimNew 1, DimNew 3], [DimCoercion 1, DimNew 3]) - , ([2, 2], [DimNew 1, DimNew 3], [DimNew 1, DimNew 3]) - ] + informReshape (shape :: [Int]) sc @?= sc_res -- type signature to avoid warning + | (shape, sc, sc_res) <- + [ ([1, 2], [DimNew 1, DimNew 3], [DimCoercion 1, DimNew 3]), + ([2, 2], [DimNew 1, DimNew 3], [DimNew 1, DimNew 3]) + ] ] reshapeOuterTests :: [TestTree] reshapeOuterTests = [ testCase (unwords ["reshapeOuter", show sc, show n, show shape, "==", show sc_res]) $ - reshapeOuter (intShapeChange sc) n (intShape shape) @?= intShapeChange sc_res - | (sc, n, shape, sc_res) <- - [ ([DimNew 1], 1, [4, 3], [DimNew 1, DimCoercion 3]) - , ([DimNew 1], 2, [4, 3], [DimNew 1]) - , ([DimNew 2, DimNew 2], 1, [4, 3], [DimNew 2, DimNew 2, DimNew 3]) - , ([DimNew 2, DimNew 2], 2, [4, 3], [DimNew 2, DimNew 2]) - ] + reshapeOuter (intShapeChange sc) n (intShape shape) @?= intShapeChange sc_res + | (sc, n, shape, sc_res) <- + [ ([DimNew 1], 1, [4, 3], [DimNew 1, DimCoercion 3]), + ([DimNew 1], 2, [4, 3], [DimNew 1]), + ([DimNew 2, DimNew 2], 1, [4, 3], [DimNew 2, DimNew 2, DimNew 3]), + ([DimNew 2, DimNew 2], 2, [4, 3], [DimNew 2, DimNew 2]) + ] ] reshapeInnerTests :: [TestTree] reshapeInnerTests = [ testCase (unwords ["reshapeInner", show sc, show n, show shape, "==", show sc_res]) $ - reshapeInner (intShapeChange sc) n (intShape shape) @?= intShapeChange sc_res - | (sc, n, shape, sc_res) <- - [ ([DimNew 1], 1, [4, 3], [DimCoercion 4, DimNew 1]) - , ([DimNew 1], 0, [4, 3], [DimNew 1]) - , ([DimNew 2, DimNew 2], 1, [4, 3], [DimNew 4, DimNew 2, DimNew 2]) - , ([DimNew 2, DimNew 2], 0, [4, 3], [DimNew 2, DimNew 2]) - ] + reshapeInner (intShapeChange sc) n (intShape shape) @?= intShapeChange sc_res + | (sc, n, shape, sc_res) <- + [ ([DimNew 1], 1, [4, 3], [DimCoercion 4, DimNew 1]), + ([DimNew 1], 0, [4, 3], [DimNew 1]), + ([DimNew 2, DimNew 2], 1, [4, 3], [DimNew 4, DimNew 2, DimNew 2]), + ([DimNew 2, DimNew 2], 0, [4, 3], [DimNew 2, DimNew 2]) + ] ] intShape :: [Int] -> Shape @@ -79,15 +79,19 @@ fuseReshapeProp :: TestTree fuseReshapeProp = testProperty "fuseReshape result matches second argument" prop - where prop :: ShapeChange Int -> ShapeChange Int -> Bool - prop sc1 sc2 = map newDim (fuseReshape sc1 sc2) == map newDim sc2 + where + prop :: ShapeChange Int -> ShapeChange Int -> Bool + prop sc1 sc2 = map newDim (fuseReshape sc1 sc2) == map newDim sc2 informReshapeProp :: TestTree informReshapeProp = testProperty "informReshape result matches second argument" prop - where prop :: [Int] -> ShapeChange Int -> Bool - prop sc1 sc2 = map newDim (informReshape sc1 sc2) == map newDim sc2 + where + prop :: [Int] -> ShapeChange Int -> Bool + prop sc1 sc2 = map newDim (informReshape sc1 sc2) == map newDim sc2 instance Arbitrary d => Arbitrary (DimChange d) where - arbitrary = oneof [ DimNew <$> arbitrary - , DimCoercion <$> arbitrary - ] + arbitrary = + oneof + [ DimNew <$> arbitrary, + DimCoercion <$> arbitrary + ] diff --git a/unittests/Futhark/IR/PropTests.hs b/unittests/Futhark/IR/PropTests.hs --- a/unittests/Futhark/IR/PropTests.hs +++ b/unittests/Futhark/IR/PropTests.hs @@ -1,16 +1,19 @@ {-# LANGUAGE FlexibleInstances #-} {-# OPTIONS_GHC -fno-warn-orphans #-} + module Futhark.IR.PropTests - ( tests + ( tests, ) where -import Test.Tasty - -import qualified Futhark.IR.Prop.ReshapeTests import qualified Futhark.IR.Prop.RearrangeTests +import qualified Futhark.IR.Prop.ReshapeTests +import Test.Tasty tests :: TestTree -tests = testGroup "PropTests" - [Futhark.IR.Prop.ReshapeTests.tests, - Futhark.IR.Prop.RearrangeTests.tests] +tests = + testGroup + "PropTests" + [ Futhark.IR.Prop.ReshapeTests.tests, + Futhark.IR.Prop.RearrangeTests.tests + ] diff --git a/unittests/Futhark/IR/Syntax/CoreTests.hs b/unittests/Futhark/IR/Syntax/CoreTests.hs --- a/unittests/Futhark/IR/Syntax/CoreTests.hs +++ b/unittests/Futhark/IR/Syntax/CoreTests.hs @@ -1,67 +1,69 @@ {-# LANGUAGE FlexibleInstances #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -module Futhark.IR.Syntax.CoreTests - ( tests ) - where -import Control.Applicative +module Futhark.IR.Syntax.CoreTests (tests) where +import Control.Applicative +import Futhark.IR.Pretty () +import Futhark.IR.PrimitiveTests () +import Futhark.IR.Syntax.Core +import Language.Futhark.CoreTests () +import Test.QuickCheck import Test.Tasty import Test.Tasty.HUnit -import Test.QuickCheck - import Prelude -import Language.Futhark.CoreTests () -import Futhark.IR.PrimitiveTests() -import Futhark.IR.Syntax.Core -import Futhark.IR.Pretty () - tests :: TestTree tests = testGroup "Internal CoreTests" subShapeTests subShapeTests :: [TestTree] subShapeTests = - [ shape [free 1, free 2] `isSubShapeOf` shape [free 1, free 2] - , shape [free 1, free 3] `isNotSubShapeOf` shape [free 1, free 2] - , shape [free 1] `isNotSubShapeOf` shape [free 1, free 2] - , shape [free 1, free 2] `isSubShapeOf` shape [free 1, Ext 3] - , shape [Ext 1, Ext 2] `isNotSubShapeOf` shape [Ext 1, Ext 1] - , shape [Ext 1, Ext 1] `isSubShapeOf` shape [Ext 1, Ext 2] + [ shape [free 1, free 2] `isSubShapeOf` shape [free 1, free 2], + shape [free 1, free 3] `isNotSubShapeOf` shape [free 1, free 2], + shape [free 1] `isNotSubShapeOf` shape [free 1, free 2], + shape [free 1, free 2] `isSubShapeOf` shape [free 1, Ext 3], + shape [Ext 1, Ext 2] `isNotSubShapeOf` shape [Ext 1, Ext 1], + shape [Ext 1, Ext 1] `isSubShapeOf` shape [Ext 1, Ext 2] ] - where shape :: [ExtSize] -> ExtShape - shape = Shape + where + shape :: [ExtSize] -> ExtShape + shape = Shape - free :: Int -> ExtSize - free = Free . Constant . IntValue . Int32Value . fromIntegral + free :: Int -> ExtSize + free = Free . Constant . IntValue . Int32Value . fromIntegral - isSubShapeOf shape1 shape2 = - subShapeTest shape1 shape2 True - isNotSubShapeOf shape1 shape2 = - subShapeTest shape1 shape2 False + isSubShapeOf shape1 shape2 = + subShapeTest shape1 shape2 True + isNotSubShapeOf shape1 shape2 = + subShapeTest shape1 shape2 False - subShapeTest :: ExtShape -> ExtShape -> Bool -> TestTree - subShapeTest shape1 shape2 expected = - testCase ("subshapeOf " ++ pretty shape1 ++ " " ++ - pretty shape2 ++ " == " ++ - show expected) $ - shape1 `subShapeOf` shape2 @?= expected + subShapeTest :: ExtShape -> ExtShape -> Bool -> TestTree + subShapeTest shape1 shape2 expected = + testCase + ( "subshapeOf " ++ pretty shape1 ++ " " + ++ pretty shape2 + ++ " == " + ++ show expected + ) + $ shape1 `subShapeOf` shape2 @?= expected instance Arbitrary NoUniqueness where arbitrary = pure NoUniqueness instance (Arbitrary shape, Arbitrary u) => Arbitrary (TypeBase shape u) where arbitrary = - oneof [ Prim <$> arbitrary - , Array <$> arbitrary <*> arbitrary <*> arbitrary - ] + oneof + [ Prim <$> arbitrary, + Array <$> arbitrary <*> arbitrary <*> arbitrary + ] instance Arbitrary Ident where arbitrary = Ident <$> arbitrary <*> arbitrary instance Arbitrary Rank where - arbitrary = Rank <$> elements [1..9] + arbitrary = Rank <$> elements [1 .. 9] instance Arbitrary Shape where - arbitrary = Shape . map intconst <$> listOf1 (elements [1..9]) - where intconst = Constant . IntValue . Int32Value + arbitrary = Shape . map intconst <$> listOf1 (elements [1 .. 9]) + where + intconst = Constant . IntValue . Int32Value diff --git a/unittests/Futhark/IR/SyntaxTests.hs b/unittests/Futhark/IR/SyntaxTests.hs --- a/unittests/Futhark/IR/SyntaxTests.hs +++ b/unittests/Futhark/IR/SyntaxTests.hs @@ -1,7 +1,6 @@ {-# OPTIONS_GHC -fno-warn-orphans #-} -module Futhark.IR.SyntaxTests - () -where + +module Futhark.IR.SyntaxTests () where -- There isn't anything to test in this module. At some point, maybe -- we can put some Arbitrary instances here. diff --git a/unittests/Futhark/Pkg/SolveTests.hs b/unittests/Futhark/Pkg/SolveTests.hs --- a/unittests/Futhark/Pkg/SolveTests.hs +++ b/unittests/Futhark/Pkg/SolveTests.hs @@ -1,59 +1,88 @@ {-# LANGUAGE OverloadedStrings #-} + module Futhark.Pkg.SolveTests (tests) where import qualified Data.Map as M -import qualified Data.Text as T import Data.Monoid - +import qualified Data.Text as T +import Futhark.Pkg.Solve +import Futhark.Pkg.Types import Test.Tasty import Test.Tasty.HUnit - -import Futhark.Pkg.Types -import Futhark.Pkg.Solve - import Prelude semverE :: T.Text -> SemVer semverE s = case parseVersion s of - Left err -> error $ T.unpack s <> - " is not a valid version number: " <> - errorBundlePretty err - Right x -> x + Left err -> + error $ + T.unpack s + <> " is not a valid version number: " + <> errorBundlePretty err + Right x -> x -- | A world of packages and interdependencies for testing the solver -- without touching the outside world. testEnv :: PkgRevDepInfo -testEnv = M.fromList $ concatMap frob - [ ("athas", [ ("foo", [ ("0.1.0", []) - , ("0.2.0", [("athas/bar", "1.0.0")]) - , ("0.3.0", [])]) - , ("foo@v2", [ ("2.0.0", [("athas/quux", "0.1.0")])]) - , ("bar", [ ("1.0.0", [])]) - , ("baz", [ ("0.1.0", [("athas/foo", "0.3.0")])]) - , ("quux", [ ("0.1.0", [ ("athas/foo", "0.2.0") - , ("athas/baz", "0.1.0") ])]) - , ("quux_perm", [ ("0.1.0", [ ("athas/baz", "0.1.0") - , ("athas/foo", "0.2.0")])]) - , ("x_bar", [ ("1.0.0", [("athas/bar", "1.0.0")])]) - , ("x_foo", [ ("1.0.0", [("athas/foo", "0.3.0")])]) - , ("tricky", [ ("1.0.0", [ ("athas/foo", "0.2.0") - , ("athas/x_foo", "1.0.0")])]) - ]) - - -- Some mutually recursive packages. - , ("nasty", [ ("foo", [ ("1.0.0", [("nasty/bar", "1.0.0")])]) - , ("bar", [ ("1.0.0", [("nasty/foo", "1.0.0")])])]) - ] - where frob (user, repos) = do - (repo, repo_revs) <- repos - (rev, deps) <- repo_revs - let rev' = semverE rev - onDep (dp, dv) = (dp, (semverE dv, Nothing)) - deps' = PkgRevDeps $ M.fromList $ map onDep deps - return ((user <> "/" <> repo, rev'), deps') +testEnv = + M.fromList $ + concatMap + frob + [ ( "athas", + [ ( "foo", + [ ("0.1.0", []), + ("0.2.0", [("athas/bar", "1.0.0")]), + ("0.3.0", []) + ] + ), + ("foo@v2", [("2.0.0", [("athas/quux", "0.1.0")])]), + ("bar", [("1.0.0", [])]), + ("baz", [("0.1.0", [("athas/foo", "0.3.0")])]), + ( "quux", + [ ( "0.1.0", + [ ("athas/foo", "0.2.0"), + ("athas/baz", "0.1.0") + ] + ) + ] + ), + ( "quux_perm", + [ ( "0.1.0", + [ ("athas/baz", "0.1.0"), + ("athas/foo", "0.2.0") + ] + ) + ] + ), + ("x_bar", [("1.0.0", [("athas/bar", "1.0.0")])]), + ("x_foo", [("1.0.0", [("athas/foo", "0.3.0")])]), + ( "tricky", + [ ( "1.0.0", + [ ("athas/foo", "0.2.0"), + ("athas/x_foo", "1.0.0") + ] + ) + ] + ) + ] + ), + -- Some mutually recursive packages. + ( "nasty", + [ ("foo", [("1.0.0", [("nasty/bar", "1.0.0")])]), + ("bar", [("1.0.0", [("nasty/foo", "1.0.0")])]) + ] + ) + ] + where + frob (user, repos) = do + (repo, repo_revs) <- repos + (rev, deps) <- repo_revs + let rev' = semverE rev + onDep (dp, dv) = (dp, (semverE dv, Nothing)) + deps' = PkgRevDeps $ M.fromList $ map onDep deps + return ((user <> "/" <> repo, rev'), deps') newtype SolverRes = SolverRes BuildList - deriving (Eq) + deriving (Eq) instance Show SolverRes where show (SolverRes bl) = T.unpack $ prettyBuildList bl @@ -61,49 +90,55 @@ solverTest :: PkgPath -> T.Text -> Either T.Text [(PkgPath, T.Text)] -> TestTree solverTest p v expected = testCase (T.unpack $ p <> "-" <> prettySemVer v') $ - fmap SolverRes (solveDepsPure testEnv target) - @?= expected' - where target = PkgRevDeps $ M.singleton p (v', Nothing) - v' = semverE v - expected' = SolverRes . BuildList . M.fromList . map onRes <$> expected - onRes (dp, dv) = (dp, semverE dv) + fmap SolverRes (solveDepsPure testEnv target) + @?= expected' + where + target = PkgRevDeps $ M.singleton p (v', Nothing) + v' = semverE v + expected' = SolverRes . BuildList . M.fromList . map onRes <$> expected + onRes (dp, dv) = (dp, semverE dv) tests :: TestTree -tests = testGroup "SolveTests" - [ - solverTest "athas/foo" "0.1.0" $ - Right [ ("athas/foo", "0.1.0")] - - , solverTest "athas/foo" "0.2.0" $ - Right [ ("athas/foo", "0.2.0") - , ("athas/bar", "1.0.0")] - - , solverTest "athas/quux" "0.1.0" $ - Right [ ("athas/quux", "0.1.0") - , ("athas/foo", "0.3.0") - , ("athas/baz", "0.1.0")] - - , solverTest "athas/quux_perm" "0.1.0" $ - Right [ ("athas/quux_perm", "0.1.0") - , ("athas/foo", "0.3.0") - , ("athas/baz", "0.1.0")] - - , solverTest "athas/foo@v2" "2.0.0" $ - Right [ ("athas/foo@v2", "2.0.0") - , ("athas/quux", "0.1.0") - , ("athas/foo", "0.3.0") - , ("athas/baz", "0.1.0") +tests = + testGroup + "SolveTests" + [ solverTest "athas/foo" "0.1.0" $ + Right [("athas/foo", "0.1.0")], + solverTest "athas/foo" "0.2.0" $ + Right + [ ("athas/foo", "0.2.0"), + ("athas/bar", "1.0.0") + ], + solverTest "athas/quux" "0.1.0" $ + Right + [ ("athas/quux", "0.1.0"), + ("athas/foo", "0.3.0"), + ("athas/baz", "0.1.0") + ], + solverTest "athas/quux_perm" "0.1.0" $ + Right + [ ("athas/quux_perm", "0.1.0"), + ("athas/foo", "0.3.0"), + ("athas/baz", "0.1.0") + ], + solverTest "athas/foo@v2" "2.0.0" $ + Right + [ ("athas/foo@v2", "2.0.0"), + ("athas/quux", "0.1.0"), + ("athas/foo", "0.3.0"), + ("athas/baz", "0.1.0") + ], + solverTest "athas/foo@v3" "3.0.0" $ + Left "Unknown package/version: athas/foo@v3-3.0.0", + solverTest "nasty/foo" "1.0.0" $ + Right + [ ("nasty/foo", "1.0.0"), + ("nasty/bar", "1.0.0") + ], + solverTest "athas/tricky" "1.0.0" $ + Right + [ ("athas/tricky", "1.0.0"), + ("athas/foo", "0.3.0"), + ("athas/x_foo", "1.0.0") ] - - , solverTest "athas/foo@v3" "3.0.0" $ - Left "Unknown package/version: athas/foo@v3-3.0.0" - - , solverTest "nasty/foo" "1.0.0" $ - Right [ ("nasty/foo", "1.0.0") - , ("nasty/bar", "1.0.0")] - - , solverTest "athas/tricky" "1.0.0" $ - Right [ ("athas/tricky", "1.0.0") - , ("athas/foo", "0.3.0") - , ("athas/x_foo", "1.0.0")] - ] + ] diff --git a/unittests/Language/Futhark/CoreTests.hs b/unittests/Language/Futhark/CoreTests.hs --- a/unittests/Language/Futhark/CoreTests.hs +++ b/unittests/Language/Futhark/CoreTests.hs @@ -1,15 +1,14 @@ {-# LANGUAGE FlexibleInstances #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -module Language.Futhark.CoreTests () -where -import Test.QuickCheck +module Language.Futhark.CoreTests () where +import Futhark.IR.PrimitiveTests () import Language.Futhark.Core -import Futhark.IR.PrimitiveTests() +import Test.QuickCheck instance Arbitrary Name where - arbitrary = nameFromString <$> listOf1 (elements ['a'..'z']) + arbitrary = nameFromString <$> listOf1 (elements ['a' .. 'z']) instance Arbitrary VName where arbitrary = VName <$> arbitrary <*> arbitrary diff --git a/unittests/Language/Futhark/SyntaxTests.hs b/unittests/Language/Futhark/SyntaxTests.hs --- a/unittests/Language/Futhark/SyntaxTests.hs +++ b/unittests/Language/Futhark/SyntaxTests.hs @@ -1,38 +1,38 @@ -{-# OPTIONS_GHC -fno-warn-orphans #-} {-# LANGUAGE FlexibleInstances #-} -module Language.Futhark.SyntaxTests (tests) -where - -import Control.Applicative +{-# OPTIONS_GHC -fno-warn-orphans #-} -import Prelude +module Language.Futhark.SyntaxTests (tests) where +import Control.Applicative +import Futhark.IR.PrimitiveTests () +import Language.Futhark.Syntax import Test.QuickCheck import Test.Tasty - -import Language.Futhark.Syntax - -import Futhark.IR.PrimitiveTests() +import Prelude tests :: TestTree tests = testGroup "Source SyntaxTests" [] instance Arbitrary BinOp where - arbitrary = elements [minBound..maxBound] + arbitrary = elements [minBound .. maxBound] instance Arbitrary Uniqueness where arbitrary = elements [Unique, Nonunique] instance Arbitrary PrimType where - arbitrary = oneof [ Signed <$> arbitrary - , Unsigned <$> arbitrary - , FloatType <$> arbitrary - , pure Bool - ] + arbitrary = + oneof + [ Signed <$> arbitrary, + Unsigned <$> arbitrary, + FloatType <$> arbitrary, + pure Bool + ] instance Arbitrary PrimValue where - arbitrary = oneof [ SignedValue <$> arbitrary - , UnsignedValue <$> arbitrary - , FloatValue <$> arbitrary - , BoolValue <$> arbitrary - ] + arbitrary = + oneof + [ SignedValue <$> arbitrary, + UnsignedValue <$> arbitrary, + FloatValue <$> arbitrary, + BoolValue <$> arbitrary + ] diff --git a/unittests/futhark_tests.hs b/unittests/futhark_tests.hs --- a/unittests/futhark_tests.hs +++ b/unittests/futhark_tests.hs @@ -1,26 +1,26 @@ module Main (main) where -import qualified Language.Futhark.SyntaxTests import qualified Futhark.BenchTests -import qualified Futhark.IR.Syntax.CoreTests -import qualified Futhark.IR.PropTests import qualified Futhark.IR.Mem.IxFunTests -import qualified Futhark.Pkg.SolveTests import qualified Futhark.IR.PrimitiveTests - +import qualified Futhark.IR.PropTests +import qualified Futhark.IR.Syntax.CoreTests +import qualified Futhark.Pkg.SolveTests +import qualified Language.Futhark.SyntaxTests import Test.Tasty allTests :: TestTree allTests = - testGroup "" - [ Language.Futhark.SyntaxTests.tests - , Futhark.BenchTests.tests - , Futhark.IR.PropTests.tests - , Futhark.IR.Syntax.CoreTests.tests - , Futhark.Pkg.SolveTests.tests - , Futhark.IR.Mem.IxFunTests.tests - , Futhark.IR.PrimitiveTests.tests - ] + testGroup + "" + [ Language.Futhark.SyntaxTests.tests, + Futhark.BenchTests.tests, + Futhark.IR.PropTests.tests, + Futhark.IR.Syntax.CoreTests.tests, + Futhark.Pkg.SolveTests.tests, + Futhark.IR.Mem.IxFunTests.tests, + Futhark.IR.PrimitiveTests.tests + ] main :: IO () main = defaultMain allTests