diff --git a/docs/conf.py b/docs/conf.py --- a/docs/conf.py +++ b/docs/conf.py @@ -262,6 +262,7 @@ man_pages = [ ('man/futhark', 'futhark', 'a parallel functional array language', [], 1), ('man/futhark-c', 'futhark-c', 'compile Futhark to sequential C', [], 1), + ('man/futhark-multicore', 'futhark-multicore', 'compile Futhark to multithreaded C', [], 1), ('man/futhark-opencl', 'futhark-opencl', 'compile Futhark to OpenCL', [], 1), ('man/futhark-cuda', 'futhark-cuda', 'compile Futhark to CUDA', [], 1), ('man/futhark-python', 'futhark-python', 'compile Futhark to sequential Python', [], 1), diff --git a/docs/index.rst b/docs/index.rst --- a/docs/index.rst +++ b/docs/index.rst @@ -45,6 +45,7 @@ man/futhark-autotune.rst man/futhark-bench.rst man/futhark-c.rst + man/futhark-multicore.rst man/futhark-cuda.rst man/futhark-dataset.rst man/futhark-doc.rst diff --git a/docs/man/futhark-multicore.rst b/docs/man/futhark-multicore.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-multicore.rst @@ -0,0 +1,109 @@ +.. role:: ref(emphasis) + +.. _futhark-multicore(1): + +================= +futhark-multicore +================= + +SYNOPSIS +======== + +futhark multicore [options...] + +DESCRIPTION +=========== + +``futhark multicore`` translates a Futhark program to multithreaded C +code, and either compiles that C code with gcc(1) to an executable +binary program, or produces a ``.h`` and ``.c`` file that can be +linked with other code.. The standard Futhark optimisation pipeline +is used, and GCC is invoked with ``-O3 -lm -std=c11 -pthread``. + +The resulting program will read the arguments to the entry point +(``main`` by default) from standard input and print its return value +on standard output. The arguments are read and printed in Futhark +syntax. + +OPTIONS +======= + +-h + Print help text to standard output and exit. + +--library + Generate a library instead of an executable. Appends ``.c``/``.h`` + to the name indicated by the ``-o`` option to determine output + file names. + +-o outfile + Where to write the result. If the source program is named + ``foo.fut``, this defaults to ``foo``. + +--safe + Ignore ``unsafe`` in program and perform safety checks unconditionally. + +-v verbose + Enable debugging output. If compilation fails due to a compiler + error, the result of the last successful compiler step will be + printed to standard error. + +-V + Print version information on standard output and exit. + +-W + Do not print any warnings. + +--Werror + Treat warnings as errors. + +EXECUTABLE OPTIONS +================== + +The following options are accepted by executables generated by +``futhark multicore``. + +-h, --help + + Print help text to standard output and exit. + +-b, --binary-output + + Print the program result in the binary output format. The default + is human-readable text, which is very slow. + +-D, --debugging + + Perform possibly expensive internal correctness checks and verbose + logging. Implies ``-L``. + +-e, --entry-point=FUN + + The entry point to run. Defaults to ``main``. + +-L, --log + + Print various low-overhead logging information to stderr while + running. + +-r, --runs=NUM + + Perform NUM runs of the program. With ``-t``, the runtime for each + individual run will be printed. Additionally, a single leading + warmup run will be performed (not counted). Only the final run will + have its result written to stdout. + +-t, --write-runtime-to=FILE + + Print the time taken to execute the program to the indicated file, an + integral number of microseconds. + +BUGS +==== + +Currently works only on Unix-like systems. + +SEE ALSO +======== + +:ref:`futhark-c(1)`, :ref:`futhark-test(1)` diff --git a/docs/usage.rst b/docs/usage.rst --- a/docs/usage.rst +++ b/docs/usage.rst @@ -107,8 +107,8 @@ (:ref:`binary-data-format`). For large outputs, this is significantly faster and takes up less space. -Parallel Options -~~~~~~~~~~~~~~~~ +GPU Options +~~~~~~~~~~~ The following options are supported by executables generated with the GPU backends (``opencl``, ``pyopencl``, and ``cuda``). 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.17.3 +version: 0.18.1 synopsis: An optimising compiler for a functional, array-oriented language. description: Futhark is a small programming language designed to be compiled to @@ -79,6 +79,7 @@ Futhark.CLI.Dev Futhark.CLI.Doc Futhark.CLI.Misc + Futhark.CLI.Multicore Futhark.CLI.OpenCL Futhark.CLI.Pkg Futhark.CLI.PyOpenCL @@ -97,6 +98,7 @@ Futhark.CodeGen.Backends.GenericPython.AST Futhark.CodeGen.Backends.GenericPython.Definitions Futhark.CodeGen.Backends.GenericPython.Options + Futhark.CodeGen.Backends.MulticoreC Futhark.CodeGen.Backends.PyOpenCL Futhark.CodeGen.Backends.PyOpenCL.Boilerplate Futhark.CodeGen.Backends.SequentialC @@ -104,6 +106,7 @@ Futhark.CodeGen.Backends.SimpleRep Futhark.CodeGen.ImpCode Futhark.CodeGen.ImpCode.Kernels + Futhark.CodeGen.ImpCode.Multicore Futhark.CodeGen.ImpCode.OpenCL Futhark.CodeGen.ImpCode.Sequential Futhark.CodeGen.ImpGen @@ -116,6 +119,12 @@ Futhark.CodeGen.ImpGen.Kernels.SegScan Futhark.CodeGen.ImpGen.Kernels.ToOpenCL Futhark.CodeGen.ImpGen.Kernels.Transpose + Futhark.CodeGen.ImpGen.Multicore + Futhark.CodeGen.ImpGen.Multicore.Base + Futhark.CodeGen.ImpGen.Multicore.SegHist + Futhark.CodeGen.ImpGen.Multicore.SegMap + Futhark.CodeGen.ImpGen.Multicore.SegRed + Futhark.CodeGen.ImpGen.Multicore.SegScan Futhark.CodeGen.ImpGen.OpenCL Futhark.CodeGen.ImpGen.Sequential Futhark.CodeGen.ImpGen.Transpose @@ -136,6 +145,9 @@ Futhark.IR.Kernels.Simplify Futhark.IR.Kernels.Sizes Futhark.IR.KernelsMem + Futhark.IR.MC + Futhark.IR.MC.Op + Futhark.IR.MCMem Futhark.IR.Mem Futhark.IR.Mem.IxFun Futhark.IR.Mem.Simplify @@ -195,6 +207,7 @@ Futhark.Pass.ExplicitAllocations.Kernels Futhark.Pass.ExplicitAllocations.SegOp Futhark.Pass.ExplicitAllocations.Seq + Futhark.Pass.ExplicitAllocations.MC Futhark.Pass.ExtractKernels Futhark.Pass.ExtractKernels.BlockedKernel Futhark.Pass.ExtractKernels.DistributeNests @@ -204,6 +217,7 @@ Futhark.Pass.ExtractKernels.Intragroup Futhark.Pass.ExtractKernels.StreamKernel Futhark.Pass.ExtractKernels.ToKernels + Futhark.Pass.ExtractMulticore Futhark.Pass.FirstOrderTransform Futhark.Pass.KernelBabysitting Futhark.Pass.Simplify 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 = +let iota (n: i64): *[n]i64 = 0..1.. x) (iota n) +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 = { @@ -122,8 +120,7 @@ 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 @@ -852,8 +849,7 @@ let bool (x: bool) = if x then 1f64 else 0f64 - let from_fraction (x: i32) (y: i32) = i32 x / i32 y - let to_i32 (x: f64) = intrinsics.fptosi_f64_i32 x + let from_fraction (x: i64) (y: i64) = i64 x / i64 y let to_i64 (x: f64) = intrinsics.fptosi_f64_i64 x let to_f64 (x: f64) = x @@ -960,8 +956,7 @@ let bool (x: bool) = if x then 1f32 else 0f32 - let from_fraction (x: i32) (y: i32) = i32 x / i32 y - let to_i32 (x: f32) = intrinsics.fptosi_f32_i32 x + let from_fraction (x: i64) (y: i64) = i64 x / i64 y let to_i64 (x: f32) = intrinsics.fptosi_f32_i64 x let to_f64 (x: f32) = intrinsics.fpconv_f32_f64 x diff --git a/prelude/soacs.fut b/prelude/soacs.fut --- a/prelude/soacs.fut +++ b/prelude/soacs.fut @@ -118,7 +118,7 @@ -- -- In practice, the *O(n)* behaviour only occurs if *m* is also very -- large. -let reduce_by_index 'a [m] [n] (dest : *[m]a) (f : a -> 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/chaselev.h b/rts/c/chaselev.h new file mode 100644 --- /dev/null +++ b/rts/c/chaselev.h @@ -0,0 +1,180 @@ +// Start of chaselev.h + +#ifndef _CHASELEV_H_ +#define _CHASELEV_H_ + +/* Implementation of Chase-lev's concurrent lock-free deque + from ``Dynamic Circular Work-Stealing Deque`` (2005) + This implementation was ported from + https://github.com/deepsea-inria/heartbeat + + !!! + This implementation leaks memory, + if the circular array is grown + as we don't maintain a list of the old buffers. + However, we can't safely free it either as a stealing thread might + be reading from it. + !!! + */ + +#if defined(MCCHASELEV) +#include +#include +#include + + +static struct subtask* const STEAL_RES_EMPTY = (struct subtask*) 0; +static struct subtask* const STEAL_RES_ABORT = (struct subtask*) 1; + +static const int strong = 0; +static const int backoff_nb_cycles = 1l << 10; + + +static inline struct subtask* cb_get(struct subtask **buf, int64_t capacity, int64_t i) { + return (struct subtask*)__atomic_load_n(&buf[i % capacity], __ATOMIC_RELAXED); +} + +static inline void cb_put (struct subtask **buf, int64_t capacity, int64_t i, struct subtask* x) { + __atomic_store_n(&buf[i % capacity], x, __ATOMIC_RELAXED); +} + +struct deque_buffer* grow(struct subtask **old_array, + int64_t old_capacity, + int64_t new_capacity, + int64_t b, + int64_t t) +{ + struct deque_buffer* new_deque_buffer = malloc(sizeof(struct deque_buffer)); + new_deque_buffer->size = new_capacity; + new_deque_buffer->array = calloc(new_capacity, sizeof(struct subtask*)); + + for (int64_t i = t; i < b; i++) { + cb_put(new_deque_buffer->array, new_capacity, i, cb_get(old_array, old_capacity, i)); + } + return new_deque_buffer; +} + +static inline int deque_init(struct deque *q, int64_t capacity) { + assert(q != NULL); + memset(q, 0, sizeof(struct deque)); + + q->buffer = malloc(sizeof(struct deque_buffer)); + q->buffer->array = calloc(capacity, sizeof(struct subtask*)); + q->buffer->size = capacity; + + q->dead = 0; + + if (q->buffer->array == NULL) { + return -1; + } + + if (q->buffer == NULL) { + return -1; + } + return 0; +} + +static inline void deque_destroy(struct deque* q) +{ + q->dead = 1; + free(q->buffer->array); + free(q->buffer); +} + +static inline int cas_top (struct deque *q, int64_t old_val, int64_t new_val) { + int64_t ov = old_val; + if(__atomic_compare_exchange_n(&q->top, &ov, new_val, strong, + __ATOMIC_SEQ_CST, __ATOMIC_RELAXED)) { + return 1; + } + spin_for(backoff_nb_cycles); + return 0; +} + + +void push_back(struct deque *q, struct subtask*subtask) +{ + assert(subtask != NULL); + assert(q != NULL); + + int64_t b = __atomic_load_n(&q->bottom, __ATOMIC_RELAXED); // load atomically + int64_t t = __atomic_load_n(&q->top, __ATOMIC_ACQUIRE); // load atomically + struct deque_buffer *buffer = __atomic_load_n(&q->buffer, __ATOMIC_RELAXED); + if (b-t >= (buffer->size - 1)) { + // grow_queue + struct subtask **old_array = buffer->array; + int64_t old_capacity = __atomic_load_n(&buffer->size, __ATOMIC_RELAXED); + int64_t new_capacity = old_capacity * 2; + struct deque_buffer *new_buffer = grow(old_array, old_capacity, new_capacity, b, t); + __atomic_store_n(&q->buffer, new_buffer, __ATOMIC_RELEASE); + buffer = __atomic_load_n(&q->buffer, __ATOMIC_RELAXED); + memset(old_array, 0, sizeof(struct subtask*) * old_capacity); + // free(old_array); Not safe!! + } + + cb_put(buffer->array, buffer->size, b, subtask); + __atomic_thread_fence(__ATOMIC_RELEASE); + __atomic_store_n(&q->bottom, b+1, __ATOMIC_RELAXED); + return; +} + + +struct subtask * pop_back(struct deque *q) +{ + int64_t b = __atomic_load_n(&q->bottom, __ATOMIC_RELAXED) - 1; // load atomically + struct deque_buffer *buffer = __atomic_load_n(&q->buffer, __ATOMIC_RELAXED); + __atomic_store_n(&q->bottom, b, __ATOMIC_RELAXED); + __atomic_thread_fence(__ATOMIC_SEQ_CST); + int64_t t = __atomic_load_n(&q->top, __ATOMIC_RELAXED); + if (b < t) { + __atomic_store_n(&q->bottom, t, __ATOMIC_RELAXED); + return NULL; + } + struct subtask* item = cb_get(buffer->array, buffer->size, b); + if (b > t) { + return item; + } + + // else there's only one item left + // Did we win the race? + if (!cas_top(q, t, t + 1)) { + item = NULL; + } + __atomic_store_n(&q->bottom, t+1, __ATOMIC_RELAXED); + return item; +} + +struct subtask* steal(struct deque *q) +{ + assert(q != NULL); + + int64_t t = __atomic_load_n(&q->top, __ATOMIC_ACQUIRE); // load atomically + __atomic_thread_fence(__ATOMIC_SEQ_CST); + int64_t b = __atomic_load_n(&q->bottom, __ATOMIC_ACQUIRE); // load atomically + if (t >= b) { + return STEAL_RES_EMPTY; + } + + struct deque_buffer *buffer = __atomic_load_n(&q->buffer, __ATOMIC_CONSUME); + struct subtask* item = cb_get(buffer->array, buffer->size, t); + if (!cas_top(q, t, t + 1)) { + return STEAL_RES_ABORT; + } + + return item; +} + + +static inline size_t nb_subtasks(struct deque *q) +{ + return (size_t)__atomic_load_n(&q->bottom, __ATOMIC_RELAXED) - __atomic_load_n(&q->top, __ATOMIC_RELAXED); +} + +static inline int empty(struct deque *q) +{ + return nb_subtasks(q) < 1; +} + +#endif +#endif +// end of chaselev.h diff --git a/rts/c/multicore_defs.h b/rts/c/multicore_defs.h new file mode 100644 --- /dev/null +++ b/rts/c/multicore_defs.h @@ -0,0 +1,108 @@ +// start of multicore_defs.h + +#ifndef MULTICORE_DEFS +#define MULTICORE_DEFS + +#include + +/* #define MCPROFILE */ + +// Which queue implementation to use +#define MCJOBQUEUE +// NOTE! MCCHASELEV has been removed from multicore branch +// Switch to multicore_deque branch to use chase-lev deque +/* #define MCCHASELEV */ + + +#if defined(_WIN32) +#include +#elif defined(__APPLE__) +#include +// For getting cpu usage of threads +#include +#include +#elif defined(__linux__) +#include +#include +#include +#endif + + +// Forward declarations +// Scheduler definitions +struct scheduler; +struct scheduler_info; +struct scheduler_subtask; +struct scheduler_task; + + +struct subtask_queue { + int capacity; // Size of the buffer. + int first; // Index of the start of the ring buffer. + int num_used; // Number of used elements in the buffer. + struct subtask **buffer; + + pthread_mutex_t mutex; // Mutex used for synchronisation. + pthread_cond_t cond; // Condition variable used for synchronisation. + int dead; + +#if defined(MCPROFILE) + /* Profiling fields */ + uint64_t time_enqueue; + uint64_t time_dequeue; + uint64_t n_dequeues; + uint64_t n_enqueues; +#endif +}; + + + +// Function definitions +typedef int (*segop_fn)(void* args, int64_t iterations, int tid, struct scheduler_info info); +typedef int (*parloop_fn)(void* args, int64_t start, int64_t end, int subtask_id, int tid); + + +/* A subtask that can be executed by a worker */ +struct subtask { + /* The parloop function */ + parloop_fn fn; + /* Execution parameters */ + void* args; + int64_t start, end; + int id; + + /* Dynamic scheduling parameters */ + int chunkable; + int64_t chunk_size; + + /* Shared variables across subtasks */ + volatile int *counter; // Counter for ongoing subtasks + // Shared task timers and iterators + int64_t *task_time; + int64_t *task_iter; + + /* For debugging */ + const char *name; +}; + + +struct worker { + pthread_t thread; + struct scheduler *scheduler; /* Reference to the scheduler struct the worker belongs to*/ + struct subtask_queue q; + int dead; + int tid; /* Just a thread id */ + + /* "thread local" time fields used for online algorithm */ + uint64_t timer; + uint64_t total; + int nested; /* How nested the current computation is */ + + // Profiling fields + int output_usage; /* Whether to dump thread usage */ + uint64_t time_spent_working; /* Time spent in parloop functions */ +}; + +#endif + +// end of multicore_defs.h diff --git a/rts/c/multicore_util.h b/rts/c/multicore_util.h new file mode 100644 --- /dev/null +++ b/rts/c/multicore_util.h @@ -0,0 +1,103 @@ +// start of multicore_util.h + +/* Multicore Utility functions */ + +#ifndef _MULTICORE_UTIL_H_ +#define _MULTICORE_UTIL_H_ + +/* A wrapper for getting rusage on Linux and MacOS */ +/* TODO maybe figure out this for windows */ +static inline int getrusage_thread(struct rusage *rusage) +{ + int err = -1; +#if defined(__APPLE__) + thread_basic_info_data_t info = { 0 }; + mach_msg_type_number_t info_count = THREAD_BASIC_INFO_COUNT; + kern_return_t kern_err; + + kern_err = thread_info(mach_thread_self(), + THREAD_BASIC_INFO, + (thread_info_t)&info, + &info_count); + if (kern_err == KERN_SUCCESS) { + memset(rusage, 0, sizeof(struct rusage)); + rusage->ru_utime.tv_sec = info.user_time.seconds; + rusage->ru_utime.tv_usec = info.user_time.microseconds; + rusage->ru_stime.tv_sec = info.system_time.seconds; + rusage->ru_stime.tv_usec = info.system_time.microseconds; + err = 0; + } else { + errno = EINVAL; + } +#elif defined(__linux__) + err = getrusage(RUSAGE_THREAD, rusage); +#endif + return err; +} + +/* returns the number of logical cores */ +static int num_processors() +{ +#if defined(_WIN32) +/* https://docs.microsoft.com/en-us/windows/win32/api/sysinfoapi/ns-sysinfoapi-system_info */ + SYSTEM_INFO sysinfo; + GetSystemInfo(&sysinfo); + int ncores = sysinfo.dwNumberOfProcessors; + fprintf(stderr, "Found %d cores on your Windows machine\n Is that correct?\n", ncores); + return ncores; +#elif defined(__APPLE__) + int ncores; + size_t ncores_size = sizeof(ncores); + CHECK_ERRNO(sysctlbyname("hw.logicalcpu", &ncores, &ncores_size, NULL, 0), + "sysctlbyname (hw.logicalcpu)"); + return ncores; +#elif defined(__linux__) + return get_nprocs(); +#else + fprintf(stderr, "operating system not recognised\n"); + return -1; +#endif +} + +static inline void output_thread_usage(struct worker *worker) +{ + struct rusage usage; + CHECK_ERRNO(getrusage_thread(&usage), "getrusage_thread"); + struct timeval user_cpu_time = usage.ru_utime; + struct timeval sys_cpu_time = usage.ru_stime; + fprintf(stderr, "tid: %2d - work time %10llu us - user time: %10llu us - sys: %10llu us\n", + worker->tid, + (long long unsigned)worker->time_spent_working / 1000, + (long long unsigned)(user_cpu_time.tv_sec * 1000000 + user_cpu_time.tv_usec), + (long long unsigned)(sys_cpu_time.tv_sec * 1000000 + sys_cpu_time.tv_usec)); +} + + +static unsigned int g_seed; + +// Used to seed the generator. +static inline void fast_srand(unsigned int seed) { + g_seed = seed; +} + +// Compute a pseudorandom integer. +// Output value in range [0, 32767] +static inline unsigned int fast_rand(void) { + g_seed = (214013*g_seed+2531011); + return (g_seed>>16)&0x7FFF; +} + + +int64_t min_int64(int64_t x, int64_t y) +{ + return x < y ? x : y; +} + +int64_t max_int64(int64_t x, int64_t y) +{ + return x > y ? x : y; +} + + +#endif +// end of multicore_util.h diff --git a/rts/c/scheduler.h b/rts/c/scheduler.h new file mode 100644 --- /dev/null +++ b/rts/c/scheduler.h @@ -0,0 +1,278 @@ +// start of scheduler.h +#ifndef _SCHEDULER_H_ +#define _SCHEDULER_H_ + + +static int dummy_counter = 0; +static int64_t dummy_timer = 0; +static int64_t dummy_iter = 0; + +static int dummy_fn(void *args, int64_t start, int64_t end, int subtask_id, int tid) { + (void)args; + (void)start; + (void)end; + (void)subtask_id; + (void)tid; + return 0; +} + +// Wake up threads, who are blocking by pushing a dummy task +// onto their queue +static inline void wake_up_threads(struct scheduler *scheduler, int start_tid, int end_tid) { + +#if defined(MCDEBUG) + assert(start_tid >= 1); + assert(end_tid <= scheduler->num_threads); +#endif + for (int i = start_tid; i < end_tid; i++) { + struct subtask *subtask = create_subtask(dummy_fn, NULL, "dummy_fn", + &dummy_counter, + &dummy_timer, &dummy_iter, + 0, 0, + 0, 0, + 0); + CHECK_ERR(subtask_queue_enqueue(&scheduler->workers[i], subtask), "subtask_queue_enqueue"); + } +} + +static inline int is_finished(struct worker *worker) { + return worker->dead && subtask_queue_is_empty(&worker->q); +} + +// Try to steal from a random queue +static inline int steal_from_random_worker(struct worker* worker) +{ + int my_id = worker->tid; + struct scheduler* scheduler = worker->scheduler; + int k = random_other_worker(scheduler, my_id); + struct worker *worker_k = &scheduler->workers[k]; + struct subtask* subtask = NULL; + int retval = subtask_queue_steal(worker_k, &subtask); + if (retval == 0) { + subtask_queue_enqueue(worker, subtask); + return 1; + } + return 0; +} + + +static inline void *scheduler_worker(void* args) +{ + struct worker *worker = (struct worker*) args; + worker_local = worker; + struct subtask * subtask = NULL; + while(!is_finished(worker)) { + if (!subtask_queue_is_empty(&worker->q)) { + int retval = subtask_queue_dequeue(worker, &subtask, 0); + if (retval == 0) { + assert(subtask != NULL); + CHECK_ERR(run_subtask(worker, subtask), "run_subtask"); + } // else someone stole our work + + } else if (active_work) { /* steal */ + while (!is_finished(worker) && active_work) { + if (steal_from_random_worker(worker)) { + break; + } + } + } else { /* go back to sleep and wait for work */ + int retval = subtask_queue_dequeue(worker, &subtask, 1); + if (retval == 0) { + assert(subtask != NULL); + CHECK_ERR(run_subtask(worker, subtask), "run_subtask"); + } + } + } + + assert(subtask_queue_is_empty(&worker->q)); + __atomic_fetch_sub(&num_workers, 1, __ATOMIC_RELAXED); +#if defined(MCPROFILE) + if (worker->output_usage) + output_thread_usage(worker); +#endif + return NULL; +} + + +static inline int scheduler_execute_parloop(struct scheduler *scheduler, + struct scheduler_parloop *task, + int64_t *timer) +{ + + struct worker * worker = worker_local; + + struct scheduler_info info = task->info; + int64_t iter_pr_subtask = info.iter_pr_subtask; + int64_t remainder = info.remainder; + int nsubtasks = info.nsubtasks; + volatile int join_counter = nsubtasks; + + // Shared timer used to sum up all + // sequential work from each subtask + int64_t task_timer = 0; + int64_t task_iter = 0; + + enum scheduling sched = info.sched; + /* If each subtasks should be processed in chunks */ + int chunkable = sched == STATIC ? 0 : 1; + int64_t chunk_size = 1; // The initial chunk size when no info is avaliable + + + if (info.wake_up_threads || sched == DYNAMIC) + __atomic_add_fetch(&active_work, nsubtasks, __ATOMIC_RELAXED); + + int64_t start = 0; + int64_t end = iter_pr_subtask + (int64_t)(remainder != 0); + for (int subtask_id = 0; subtask_id < nsubtasks; subtask_id++) { + struct subtask *subtask = create_subtask(task->fn, task->args, task->name, + &join_counter, + &task_timer, &task_iter, + start, end, + chunkable, chunk_size, + subtask_id); + assert(subtask != NULL); + if (worker->nested){ + CHECK_ERR(subtask_queue_enqueue(&scheduler->workers[worker->tid], subtask), + "subtask_queue_enqueue"); + } else { + CHECK_ERR(subtask_queue_enqueue(&scheduler->workers[subtask_id], subtask), + "subtask_queue_enqueue"); + } + // Update range params + start = end; + end += iter_pr_subtask + ((subtask_id + 1) < remainder); + } + + if (info.wake_up_threads) { + wake_up_threads(scheduler, nsubtasks, scheduler->num_threads); + } + + // Join (wait for subtasks to finish) + while(join_counter != 0) { + if (!subtask_queue_is_empty(&worker->q)) { + struct subtask *subtask = NULL; + int err = subtask_queue_dequeue(worker, &subtask, 0); + if (err == 0 ) { + CHECK_ERR(run_subtask(worker, subtask), "run_subtask"); + } + } else { + if (steal_from_random_worker(worker)) { + struct subtask *subtask = NULL; + int err = subtask_queue_dequeue(worker, &subtask, 0); + if (err == 0) { + CHECK_ERR(run_subtask(worker, subtask), "run_subtask"); + } + } + } + } + + + if (info.wake_up_threads || sched == DYNAMIC) { + __atomic_sub_fetch(&active_work, nsubtasks, __ATOMIC_RELAXED); + } + + // Write back timing results of all sequential work + (*timer) += task_timer; + return scheduler_error; +} + + +static inline int scheduler_execute_task(struct scheduler *scheduler, + struct scheduler_parloop *task) +{ + + struct worker *worker = worker_local; + + int err = 0; + if (task->iterations == 0) { + return err; + } + + // How much sequential work was performed by the task + int64_t task_timer = 0; + + /* Execute task sequential or parallel based on decision made earlier */ + if (task->info.nsubtasks == 1) { + int64_t start = get_wall_time_ns(); + err = task->fn(task->args, 0, task->iterations, 0, worker->tid); + int64_t end = get_wall_time_ns(); + task_timer = end - start; + worker->time_spent_working += task_timer; + // Report time measurements + // TODO the update of both of these should really be a single atomic!! + __atomic_fetch_add(task->info.task_time, task_timer, __ATOMIC_RELAXED); + __atomic_fetch_add(task->info.task_iter, task->iterations, __ATOMIC_RELAXED); + } else { + // Add "before" time if we already are inside a task + int64_t time_before = 0; + if (worker->nested > 0) { + time_before = total_now(worker->total, worker->timer); + } + + err = scheduler_execute_parloop(scheduler, task, &task_timer); + + // Report time measurements + // TODO the update of both of these should really be a single atomic!! + __atomic_fetch_add(task->info.task_time, task_timer, __ATOMIC_RELAXED); + __atomic_fetch_add(task->info.task_iter, task->iterations, __ATOMIC_RELAXED); + + // Update timers to account for new timings + worker->total = time_before + task_timer; + worker->timer = get_wall_time_ns(); + } + + + return err; +} + +/* Decide on how schedule the incoming task i.e. how many subtasks and + to run sequential or (potentially nested) parallel code body */ +static inline int scheduler_prepare_task(struct scheduler* scheduler, + struct scheduler_segop *task) +{ + assert(task != NULL); + + struct worker *worker = worker_local; + struct scheduler_info info; + info.task_time = task->task_time; + info.task_iter = task->task_iter; + + int nsubtasks; + // Decide if task should be scheduled sequentially + if (is_small(task, scheduler->num_threads, &nsubtasks)) { + info.iter_pr_subtask = task->iterations; + info.remainder = 0; + info.nsubtasks = nsubtasks; + return task->top_level_fn(task->args, task->iterations, worker->tid, info); + } else { + info.iter_pr_subtask = task->iterations / nsubtasks; + info.remainder = task->iterations % nsubtasks; + info.sched = task->sched; + switch (task->sched) { + case STATIC: + info.nsubtasks = info.iter_pr_subtask == 0 ? info.remainder : ((task->iterations - info.remainder) / info.iter_pr_subtask); + break; + case DYNAMIC: + // As any thread can take any subtasks, we are being safe with using + // an upper bound on the number of tasks such that the task allocate enough memory + info.nsubtasks = info.iter_pr_subtask == 0 ? info.remainder : nsubtasks; + break; + default: + assert(!"Got unknown scheduling"); + } + } + + info.wake_up_threads = 0; + // We only use the nested parallel segop function if we can't exchaust all cores + // using the outer most level + if (task->nested_fn != NULL && info.nsubtasks < scheduler->num_threads && info.nsubtasks == task->iterations) { + if (worker->nested == 0) + info.wake_up_threads = 1; + return task->nested_fn(task->args, task->iterations, worker->tid, info); + } + + return task->top_level_fn(task->args, task->iterations, worker->tid, info); +} + +#endif +// End of scheduler.h diff --git a/rts/c/scheduler_common.h b/rts/c/scheduler_common.h new file mode 100644 --- /dev/null +++ b/rts/c/scheduler_common.h @@ -0,0 +1,244 @@ +// start of scheduler_common.h + +#ifndef _SCHEDULER_COMMON_H_ +#define _SCHEDULER_COMMON_H_ + +#include + +/* Scheduler definitions */ +enum scheduling { + DYNAMIC, + STATIC +}; + +/* How a given task should be executed */ +/* Filled out by the scheduler + and passed to the segop function +*/ +struct scheduler_info { + int64_t iter_pr_subtask; + int64_t remainder; + int nsubtasks; + enum scheduling sched; + int wake_up_threads; + + int64_t *task_time; + int64_t *task_iter; +}; + +struct scheduler { + struct worker *workers; + int num_threads; +}; + +/* A parallel parloop task */ +struct scheduler_parloop { + const char* name; + parloop_fn fn; + void* args; + int64_t iterations; + struct scheduler_info info; +}; + + +/* A task for the scheduler to execute */ +struct scheduler_segop { + void *args; + segop_fn top_level_fn; + segop_fn nested_fn; + int64_t iterations; + enum scheduling sched; + + // Pointers to timer and iter associated with the task + int64_t *task_time; + int64_t *task_iter; + + // For debugging + const char* name; +}; + +// If there is work to steal => active_work > 0 +static volatile int active_work = 0; +// Number of alive workers +static volatile sig_atomic_t num_workers; + +// Thread local variable worker struct +// Note that, accesses to tls variables are expensive +// Minimize direct references to this variable +__thread struct worker* worker_local = NULL; + +/* Only one error can be returned at the time now + Maybe we can provide a stack like structure for pushing errors onto + if we wish to backpropagte multiple errors */ +static volatile sig_atomic_t scheduler_error = 0; + +// kappa time unit in nanoseconds +static double kappa = 5.1f * 1000; + +int64_t total_now(int64_t total, int64_t time) { + return total + (get_wall_time_ns() - time); +} + +int random_other_worker(struct scheduler *scheduler, int my_id) { + (void)scheduler; + int my_num_workers = __atomic_load_n(&num_workers, __ATOMIC_RELAXED); + assert(my_num_workers != 1); + int i = fast_rand() % (my_num_workers - 1); + if (i >= my_id) { + i++; + } +#ifdef MCDEBUG + assert(i >= 0); + assert(i < my_num_workers); + assert(i != my_id); +#endif + + return i; +} + + +static inline int64_t compute_chunk_size(struct subtask* subtask) +{ + double C = (double)*subtask->task_time / (double)*subtask->task_iter; + if (C == 0.0F) C += DBL_EPSILON; + return max_int64((int64_t)(kappa / C), 1); +} + +/* Takes a chunk from subtask and enqueues the remaining iterations onto the worker's queue */ +/* A no-op if the subtask is not chunkable */ +static inline struct subtask* chunk_subtask(struct worker* worker, struct subtask *subtask) +{ + if (subtask->chunkable) { + // Do we have information from previous runs avaliable + if (*subtask->task_iter > 0) { + subtask->chunk_size = compute_chunk_size(subtask); + assert(subtask->chunk_size > 0); + } + int64_t remaining_iter = subtask->end - subtask->start; + assert(remaining_iter > 0); + if (remaining_iter > subtask->chunk_size) { + struct subtask *new_subtask = malloc(sizeof(struct subtask)); + *new_subtask = *subtask; + // increment the subtask join counter to account for new subtask + __atomic_fetch_add(subtask->counter, 1, __ATOMIC_RELAXED); + // Update range parameters + subtask->end = subtask->start + subtask->chunk_size; + new_subtask->start = subtask->end; + subtask_queue_enqueue(worker, new_subtask); + } + } + return subtask; +} + +static inline int run_subtask(struct worker* worker, struct subtask* subtask) +{ + assert(subtask != NULL); + assert(worker != NULL); + + subtask = chunk_subtask(worker, subtask); + worker->total = 0; + worker->timer = get_wall_time_ns(); +#if defined(MCPROFILE) + int64_t start = worker->timer; +#endif + worker->nested++; + int err = subtask->fn(subtask->args, subtask->start, subtask->end, + subtask->chunkable ? worker->tid : subtask->id, + worker->tid); + worker->nested--; + // Some error occured during some other subtask + // so we just clean-up and return + if (scheduler_error != 0) { + // Even a failed task counts as finished. + __atomic_fetch_sub(subtask->counter, 1, __ATOMIC_RELAXED); + free(subtask); + return 0; + } + if (err != 0) { + __atomic_store_n(&scheduler_error, err, __ATOMIC_RELAXED); + } + // Total sequential time spent + int64_t time_elapsed = total_now(worker->total, worker->timer); +#if defined(MCPROFILE) + worker->time_spent_working += get_wall_time_ns() - start; +#endif + int64_t iter = subtask->end - subtask->start; + // report measurements + // These updates should really be done using a single atomic CAS operation + __atomic_fetch_add(subtask->task_time, time_elapsed, __ATOMIC_RELAXED); + __atomic_fetch_add(subtask->task_iter, iter, __ATOMIC_RELAXED); + // We need a fence here, since if the counter is decremented before either + // of the two above are updated bad things can happen, e.g. if they are stack-allocated + __atomic_thread_fence(__ATOMIC_SEQ_CST); + __atomic_fetch_sub(subtask->counter, 1, __ATOMIC_RELAXED); + free(subtask); + return 0; +} + + +static inline int is_small(struct scheduler_segop *task, int nthreads, int *nsubtasks) +{ + int64_t time = *task->task_time; + int64_t iter = *task->task_iter; + + if (task->sched == DYNAMIC || iter == 0) { + *nsubtasks = nthreads; + return 0; + } + + // Estimate the constant C + double C = (double)time / (double)iter; + double cur_task_iter = (double) task->iterations; + + // Returns true if the task is small i.e. + // if the number of iterations times C is smaller + // than the overhead of subtask creation + if (C == 0.0F || C * cur_task_iter < kappa) { + *nsubtasks = 1; + return 1; + } + + // Else compute how many subtasks this tasks should create + int64_t min_iter_pr_subtask = max_int64((int64_t)(kappa / C), 1); + *nsubtasks = (int)min_int64(max_int64(task->iterations / min_iter_pr_subtask, 1), nthreads); + + return 0; +} + +// TODO make this prettier +static inline struct subtask* create_subtask(parloop_fn fn, + void* args, + const char* name, + volatile int* counter, + int64_t *timer, + int64_t *iter, + int64_t start, int64_t end, + int chunkable, + int64_t chunk_size, + int id) +{ + struct subtask* subtask = malloc(sizeof(struct subtask)); + if (subtask == NULL) { + assert(!"malloc failed in create_subtask"); + return NULL; + } + subtask->fn = fn; + subtask->args = args; + + subtask->counter = counter; + subtask->task_time = timer; + subtask->task_iter = iter; + + subtask->start = start; + subtask->end = end; + subtask->id = id; + subtask->chunkable = chunkable; + subtask->chunk_size = chunk_size; + + subtask->name = name; + return subtask; +} + + +#endif +// end of scheduler_common.h diff --git a/rts/c/scheduler_tune.h b/rts/c/scheduler_tune.h new file mode 100644 --- /dev/null +++ b/rts/c/scheduler_tune.h @@ -0,0 +1,127 @@ +/* The self-tuning program to estimate $\kappa$ */ + +struct futhark_mc_segred_stage_1_struct { + struct futhark_context *ctx; + int32_t *free_tuning_res; + int32_t *array; +}; + +/* Reduction function over an integer array */ +int futhark_mc_tuning_segred_stage_1(void *args, int64_t start, int64_t end, + int flat_tid, int tid) { + (void)flat_tid; + (void)tid; + + int err = 0; + struct futhark_mc_segred_stage_1_struct *futhark_mc_segred_stage_1_struct = (struct futhark_mc_segred_stage_1_struct *) args; + struct futhark_context *ctx = futhark_mc_segred_stage_1_struct->ctx; + int32_t *array = futhark_mc_segred_stage_1_struct->array; + int32_t *tuning_res = futhark_mc_segred_stage_1_struct->free_tuning_res; + + int32_t sum = 0; + for (int i = start; i < end; i++) { + int32_t y = array[i]; + sum = add32(sum, y); + } + *tuning_res = sum; + return err; +} + +/* The main entry point for the tuning process */ +/* Sets the global variable ``kappa`` */ +int futhark_segred_tuning_program(struct futhark_context *ctx) +{ + int err = 0; + + int64_t iterations = 100000000; + int64_t tuning_time = 0; + int64_t tuning_iter = 0; + + int32_t *array = malloc(sizeof(int32_t) * iterations); + for (int64_t i = 0; i < iterations; i++) { + array[i] = fast_rand(); + } + + int64_t start_tuning = get_wall_time_ns(); + /* **************************** */ + /* Run sequential reduce first' */ + /* **************************** */ + int64_t tuning_sequentiual_start = get_wall_time_ns(); + struct futhark_mc_segred_stage_1_struct futhark_mc_segred_stage_1_struct; + int32_t tuning_res; + futhark_mc_segred_stage_1_struct.ctx = ctx; + futhark_mc_segred_stage_1_struct.free_tuning_res = &tuning_res; + futhark_mc_segred_stage_1_struct.array = array; + + err = futhark_mc_tuning_segred_stage_1(&futhark_mc_segred_stage_1_struct, 0, iterations, 0, 0); + int64_t tuning_sequentiual_end = get_wall_time_ns(); + int64_t sequential_elapsed = tuning_sequentiual_end - tuning_sequentiual_start; + + double C = (double)sequential_elapsed / (double)iterations; + fprintf(stderr, " Time for sequential run is %lld - Found C %f\n", (long long)sequential_elapsed, C); + + /* ********************** */ + /* Now run tuning process */ + /* ********************** */ + // Setup a scheduler with a single worker + int num_threads = ctx->scheduler.num_threads; + ctx->scheduler.num_threads = 1; + ctx->scheduler.workers = malloc(sizeof(struct worker)); + worker_local = &ctx->scheduler.workers[0]; + worker_local->tid = 0; + CHECK_ERR(subtask_queue_init(&ctx->scheduler.workers[0].q, 1024), "failed to init queue for worker %d\n", 0); + + // Start tuning for kappa + double kappa_tune = 1000; // Initial kappa is 1 us + double ratio; + int64_t time_elapsed; + while(1) { + int64_t min_iter_pr_subtask = (int64_t) (kappa_tune / C) == 0 ? 1 : (kappa_tune / C); + int nsubtasks = iterations / min_iter_pr_subtask; + struct scheduler_info info; + info.iter_pr_subtask = min_iter_pr_subtask; + + info.nsubtasks = iterations / min_iter_pr_subtask; + info.remainder = iterations % min_iter_pr_subtask; + info.task_time = &tuning_time; + info.task_iter = &tuning_iter; + info.sched = STATIC; + + struct scheduler_parloop futhark_segred_tuning_scheduler_parloop; + futhark_segred_tuning_scheduler_parloop.name = "futhark_mc_tuning_segred_stage_1"; + futhark_segred_tuning_scheduler_parloop.fn = futhark_mc_tuning_segred_stage_1; + futhark_segred_tuning_scheduler_parloop.args = &futhark_mc_segred_stage_1_struct; + futhark_segred_tuning_scheduler_parloop.iterations = iterations; + futhark_segred_tuning_scheduler_parloop.info = info; + + int64_t tuning_chunked_start = get_wall_time_ns(); + int futhark_segred_tuning_program_err = + scheduler_execute_task(&ctx->scheduler, + &futhark_segred_tuning_scheduler_parloop); + assert(futhark_segred_tuning_program_err == 0); + int64_t tuning_chunked_end = get_wall_time_ns(); + time_elapsed = tuning_chunked_end - tuning_chunked_start; + + ratio = (double)time_elapsed / (double)sequential_elapsed; + if (ratio < 1.055) { + break; + } + kappa_tune += 100; // Increase by 100 ns at the time + fprintf(stderr, "nsubtask %d - kappa %f - ratio %f\n", nsubtasks, kappa_tune, ratio); + } + + int64_t end_tuning = get_wall_time_ns(); + fprintf(stderr, "tuning took %lld ns and found kappa %f - time %lld - ratio %f\n", + (long long)end_tuning - start_tuning, + kappa_tune, + (long long)time_elapsed, + ratio); + kappa = kappa_tune; + + // Clean-up + CHECK_ERR(subtask_queue_destroy(&ctx->scheduler.workers[0].q), "failed to destroy queue"); + free(array); + free(ctx->scheduler.workers); + ctx->scheduler.num_threads = num_threads; + return err; +} diff --git a/rts/c/subtask_queue.h b/rts/c/subtask_queue.h new file mode 100644 --- /dev/null +++ b/rts/c/subtask_queue.h @@ -0,0 +1,247 @@ +// start of subtask_queue.h + +#ifndef SUBTASK_QUEUE_H +#define SUBTASK_QUEUE_H + +/* Doubles the size of the queue */ +static inline int subtask_queue_grow_queue(struct subtask_queue *subtask_queue) { + + int new_capacity = 2 * subtask_queue->capacity; +#ifdef MCDEBUG + fprintf(stderr, "Growing queue to %d\n", subtask_queue->capacity * 2); +#endif + + struct subtask **new_buffer = calloc(new_capacity, sizeof(struct subtask*)); + for (int i = 0; i < subtask_queue->num_used; i++) { + new_buffer[i] = subtask_queue->buffer[(subtask_queue->first + i) % subtask_queue->capacity]; + } + + free(subtask_queue->buffer); + subtask_queue->buffer = new_buffer; + subtask_queue->capacity = new_capacity; + subtask_queue->first = 0; + + return 0; +} + +// Initialise a job queue with the given capacity. The queue starts out +// empty. Returns non-zero on error. +static inline int subtask_queue_init(struct subtask_queue *subtask_queue, int capacity) +{ + assert(subtask_queue != NULL); + memset(subtask_queue, 0, sizeof(struct subtask_queue)); + + subtask_queue->capacity = capacity; + subtask_queue->buffer = calloc(capacity, sizeof(struct subtask*)); + if (subtask_queue->buffer == NULL) { + return -1; + } + + CHECK_ERRNO(pthread_mutex_init(&subtask_queue->mutex, NULL), "pthread_mutex_init"); + CHECK_ERRNO(pthread_cond_init(&subtask_queue->cond, NULL), "pthread_cond_init"); + + return 0; +} + +// Destroy the job queue. Blocks until the queue is empty before it +// is destroyed. +static inline int subtask_queue_destroy(struct subtask_queue *subtask_queue) +{ + assert(subtask_queue != NULL); + + CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock"); + + while (subtask_queue->num_used != 0) { + CHECK_ERR(pthread_cond_wait(&subtask_queue->cond, &subtask_queue->mutex), "pthread_cond_wait"); + } + + // Queue is now empty. Let's kill it! + subtask_queue->dead = 1; + free(subtask_queue->buffer); + CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast"); + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); + + return 0; +} + +static inline void dump_queue(struct worker *worker) +{ + struct subtask_queue *subtask_queue = &worker->q; + CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock"); + for (int i = 0; i < subtask_queue->num_used; i++) { + struct subtask * subtask = subtask_queue->buffer[(subtask_queue->first + i) % subtask_queue->capacity]; + printf("queue tid %d with %d task %s\n", worker->tid, i, subtask->name); + } + CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast"); + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); +} + +// Push an element onto the end of the job queue. Blocks if the +// subtask_queue is full (its size is equal to its capacity). Returns +// non-zero on error. It is an error to push a job onto a queue that +// has been destroyed. +static inline int subtask_queue_enqueue(struct worker *worker, struct subtask *subtask ) +{ + assert(worker != NULL); + struct subtask_queue *subtask_queue = &worker->q; + +#ifdef MCPROFILE + uint64_t start = get_wall_time(); +#endif + + CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock"); + // Wait until there is room in the subtask_queue. + while (subtask_queue->num_used == subtask_queue->capacity && !subtask_queue->dead) { + if (subtask_queue->num_used == subtask_queue->capacity) { + CHECK_ERR(subtask_queue_grow_queue(subtask_queue), "subtask_queue_grow_queue"); + continue; + } + CHECK_ERR(pthread_cond_wait(&subtask_queue->cond, &subtask_queue->mutex), "pthread_cond_wait"); + } + + if (subtask_queue->dead) { + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); + return -1; + } + + // If we made it past the loop, there is room in the subtask_queue. + subtask_queue->buffer[(subtask_queue->first + subtask_queue->num_used) % subtask_queue->capacity] = subtask; + subtask_queue->num_used++; + +#ifdef MCPROFILE + uint64_t end = get_wall_time(); + subtask_queue->time_enqueue += (end - start); + subtask_queue->n_enqueues++; +#endif + // Broadcast a reader (if any) that there is now an element. + CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast"); + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); + + return 0; +} + + +/* Like subtask_queue_dequeue, but with two differences: + 1) the subtask is stolen from the __front__ of the queue + 2) returns immediately if there is no subtasks queued, + as we dont' want to block on another workers queue and +*/ +static inline int subtask_queue_steal(struct worker *worker, + struct subtask **subtask) +{ + struct subtask_queue *subtask_queue = &worker->q; + assert(subtask_queue != NULL); + +#ifdef MCPROFILE + uint64_t start = get_wall_time(); +#endif + CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock"); + + if (subtask_queue->num_used == 0) { + CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast"); + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); + return 1; + } + + if (subtask_queue->dead) { + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); + return -1; + } + + // Tasks gets stolen from the "front" + struct subtask *cur_back = subtask_queue->buffer[subtask_queue->first]; + struct subtask *new_subtask = NULL; + int remaining_iter = cur_back->end - cur_back->start; + // If subtask is chunkable, we steal half of the iterations + if (cur_back->chunkable && remaining_iter > 1) { + int64_t half = remaining_iter / 2; + new_subtask = malloc(sizeof(struct subtask)); + *new_subtask = *cur_back; + new_subtask->start = cur_back->end - half; + cur_back->end = new_subtask->start; + __atomic_fetch_add(cur_back->counter, 1, __ATOMIC_RELAXED); + } else { + new_subtask = cur_back; + subtask_queue->num_used--; + subtask_queue->first = (subtask_queue->first + 1) % subtask_queue->capacity; + } + *subtask = new_subtask; + + if (*subtask == NULL) { + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthred_mutex_unlock"); + return 1; + } + +#ifdef MCPROFILE + uint64_t end = get_wall_time(); + subtask_queue->time_dequeue += (end - start); + subtask_queue->n_dequeues++; +#endif + + // Broadcast a writer (if any) that there is now room for more. + CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast"); + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); + + return 0; +} + + +// Pop an element from the back of the job queue. +// Optional argument can be provided to block or not +static inline int subtask_queue_dequeue(struct worker *worker, + struct subtask **subtask, int blocking) +{ + assert(worker != NULL); + struct subtask_queue *subtask_queue = &worker->q; + +#ifdef MCPROFILE + uint64_t start = get_wall_time(); +#endif + + CHECK_ERR(pthread_mutex_lock(&subtask_queue->mutex), "pthread_mutex_lock"); + if (subtask_queue->num_used == 0 && !blocking) { + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); + return 1; + } + // Try to steal some work while the subtask_queue is empty + while (subtask_queue->num_used == 0 && !subtask_queue->dead) { + pthread_cond_wait(&subtask_queue->cond, &subtask_queue->mutex); + } + + if (subtask_queue->dead) { + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); + return -1; + } + + // dequeue pops from the back + *subtask = subtask_queue->buffer[(subtask_queue->first + subtask_queue->num_used - 1) % subtask_queue->capacity]; + subtask_queue->num_used--; + + if (*subtask == NULL) { + assert(!"got NULL ptr"); + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthred_mutex_unlock"); + return -1; + } + +#ifdef MCPROFILE + uint64_t end = get_wall_time(); + subtask_queue->time_dequeue += (end - start); + subtask_queue->n_dequeues++; +#endif + + // Broadcast a writer (if any) that there is now room for more. + CHECK_ERR(pthread_cond_broadcast(&subtask_queue->cond), "pthread_cond_broadcast"); + CHECK_ERR(pthread_mutex_unlock(&subtask_queue->mutex), "pthread_mutex_unlock"); + + return 0; +} + +static inline int subtask_queue_is_empty(struct subtask_queue *subtask_queue) +{ + return subtask_queue->num_used == 0; +} + + +#endif + +// End of subtask_queue.h diff --git a/rts/c/timing.h b/rts/c/timing.h --- a/rts/c/timing.h +++ b/rts/c/timing.h @@ -26,6 +26,30 @@ return time.tv_sec * 1000000 + time.tv_usec; } +static int64_t get_wall_time_ns(void) { + struct timespec time; + assert(clock_gettime(CLOCK_REALTIME, &time) == 0); + return time.tv_sec * 1000000000 + time.tv_nsec; +} + + +static inline uint64_t rdtsc() { + unsigned int hi, lo; + __asm__ __volatile__("rdtsc" : "=a"(lo), "=d"(hi)); + return ((uint64_t) lo) | (((uint64_t) hi) << 32); +} + +static inline void rdtsc_wait(uint64_t n) { + const uint64_t start = rdtsc(); + while (rdtsc() < (start + n)) { + __asm__("PAUSE"); + } +} +static inline void spin_for(uint64_t nb_cycles) { + rdtsc_wait(nb_cycles); +} + + #endif // End of timing.h. diff --git a/rts/c/util.h b/rts/c/util.h --- a/rts/c/util.h +++ b/rts/c/util.h @@ -2,6 +2,9 @@ // // Various helper functions that are useful in all generated C code. +#include +#include + static const char *fut_progname = "(embedded Futhark)"; static void futhark_panic(int eval, const char *fmt, ...) { @@ -24,6 +27,24 @@ vsnprintf(buffer, needed, s, vl); return buffer; } + + +static inline void check_err(int errval, int sets_errno, const char *fun, int line, + const char *msg, ...) +{ + if (errval) { + char str[256]; + char errnum[10]; + sprintf(errnum, "%d", errval); + sprintf(str, "ERROR: %s in %s() at line %d with error code %s\n", msg, fun, line, + sets_errno ? strerror(errno) : errnum); + fprintf(stderr, "%s", str); + exit(errval); + } +} + +#define CHECK_ERR(err, msg...) check_err(err, 0, __func__, __LINE__, msg) +#define CHECK_ERRNO(err, msg...) check_err(err, 1, __func__, __LINE__, msg) // Read a file into a NUL-terminated string; returns NULL on error. static void* slurp_file(const char *filename, size_t *size) { diff --git a/rts/python/opencl.py b/rts/python/opencl.py --- a/rts/python/opencl.py +++ b/rts/python/opencl.py @@ -122,7 +122,7 @@ 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) @@ -225,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 @@ -6,10 +6,12 @@ ( printAction, impCodeGenAction, kernelImpCodeGenAction, + multicoreImpCodeGenAction, metricsAction, compileCAction, compileOpenCLAction, compileCUDAAction, + compileMulticoreAction, sexpAction, ) where @@ -21,12 +23,15 @@ 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.MulticoreC as MulticoreC import qualified Futhark.CodeGen.Backends.SequentialC as SequentialC import qualified Futhark.CodeGen.ImpGen.Kernels as ImpGenKernels +import qualified Futhark.CodeGen.ImpGen.Multicore as ImpGenMulticore import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGenSequential import Futhark.Compiler.CLI import Futhark.IR import Futhark.IR.KernelsMem (KernelsMem) +import Futhark.IR.MCMem (MCMem) import Futhark.IR.Prop.Aliases import Futhark.IR.SeqMem (SeqMem) import Futhark.Util (runProgramWithExitCode) @@ -71,6 +76,14 @@ actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenKernels.compileProgOpenCL } +multicoreImpCodeGenAction :: Action MCMem +multicoreImpCodeGenAction = + Action + { actionName = "Compile to imperative multicore", + actionDescription = "Translate program into imperative multicore IL and write it on standard output.", + actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenMulticore.compileProg + } + -- | Print metrics about AST node counts to stdout. sexpAction :: ASTLore lore => Action lore sexpAction = @@ -207,6 +220,46 @@ [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 + Right (ExitSuccess, _, _) -> + return () + +-- | The @futhark multicore@ action. +compileMulticoreAction :: FutharkConfig -> CompilerMode -> FilePath -> Action MCMem +compileMulticoreAction fcfg mode outpath = + Action + { actionName = "Compile to multicore", + actionDescription = "Compile to multicore", + actionProcedure = helper + } + where + helper prog = do + cprog <- handleWarnings fcfg $ MulticoreC.compileProg prog + let cpath = outpath `addExtension` "c" + hpath = outpath `addExtension` "h" + + case mode of + ToLibrary -> do + let (header, impl) = MulticoreC.asLibrary cprog + liftIO $ writeFile hpath header + liftIO $ writeFile cpath impl + ToExecutable -> do + liftIO $ writeFile cpath $ MulticoreC.asExecutable cprog + -- let debug_flags = ["-g", "-fno-omit-frame-pointer", "-fsanitize=address", "-fsanitize=integer", "-fsanitize=undefined", "-fno-sanitize-recover=null"] + ret <- + liftIO $ + runProgramWithExitCode + "gcc" + [cpath, "-O3", "-pthread", "-std=c11", "-lm", "-o", outpath] + mempty case ret of Left err -> externalErrorS $ "Failed to run gcc: " ++ show err 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 @@ -526,7 +526,7 @@ SOAC lore -> m (SOAC lore, [Ident]) soacToStream soac = do - chunk_param <- newParam "chunk" $ Prim int32 + chunk_param <- newParam "chunk" $ Prim int64 let chvar = Futhark.Var $ paramName chunk_param (lam, inps) = (lambda soac, inputs soac) w = width soac @@ -579,7 +579,7 @@ 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 + 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) $ @@ -591,17 +591,17 @@ mkLet [] [outszm1id] $ BasicOp $ BinOp - (Sub Int32 OverflowUndef) + (Sub Int64 OverflowUndef) (Futhark.Var $ paramName chunk_param) - (constant (1 :: Int32)) + (constant (1 :: Int64)) -- 3. let lasteel_ids = ... empty_arr_bnd = mkLet [] [empty_arr] $ BasicOp $ CmpOp - (CmpSlt Int32) + (CmpSlt Int64) (Futhark.Var $ identName outszm1id) - (constant (0 :: Int32)) + (constant (0 :: Int64)) leltmpbnds = zipWith ( \lid arrid -> 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 @@ -48,6 +48,10 @@ class OpMetrics op where opMetrics :: op -> MetricsM () +instance OpMetrics a => OpMetrics (Maybe a) where + opMetrics Nothing = return () + opMetrics (Just x) = opMetrics x + instance OpMetrics () where opMetrics () = return () 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 @@ -7,6 +7,8 @@ primExpFromSubExp, pe32, le32, + pe64, + le64, primExpFromSubExpM, replaceInPrimExp, replaceInPrimExpM, @@ -92,6 +94,14 @@ 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 => @@ -133,9 +143,9 @@ fromMaybe (LeafExp v t) $ M.lookup v tab -- | Convert a 'SubExp' slice to a 'PrimExp' slice. -primExpSlice :: Slice SubExp -> Slice (TPrimExp Int32 VName) -primExpSlice = map $ fmap $ isInt32 . 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 (TPrimExp Int32 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/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs --- a/src/Futhark/Analysis/SymbolTable.hs +++ b/src/Futhark/Analysis/SymbolTable.hs @@ -111,7 +111,7 @@ Indexed Certificates (PrimExp VName) | -- | The indexing corresponds to another (perhaps more -- advantageous) array. - IndexedArray Certificates VName [TPrimExp Int32 VName] + IndexedArray Certificates VName [TPrimExp Int64 VName] indexedAddCerts :: Certificates -> Indexed -> Indexed indexedAddCerts cs1 (Indexed cs2 v) = Indexed (cs1 <> cs2) v @@ -122,7 +122,7 @@ freeIn' (IndexedArray cs arr v) = freeIn' cs <> freeIn' arr <> freeIn' v -- | Indexing a delayed array if possible. -type IndexArray = [TPrimExp Int32 VName] -> Maybe Indexed +type IndexArray = [TPrimExp Int64 VName] -> Maybe Indexed data Entry lore = Entry { -- | True if consumed. @@ -265,7 +265,7 @@ index' :: VName -> - [TPrimExp Int32 VName] -> + [TPrimExp Int64 VName] -> SymbolTable lore -> Maybe Indexed index' name is vtable = do @@ -288,7 +288,7 @@ SymbolTable lore -> Int -> op -> - [TPrimExp Int32 VName] -> + [TPrimExp Int64 VName] -> Maybe Indexed indexOp _ _ _ _ = Nothing @@ -322,18 +322,18 @@ | Just oldshape <- arrayDims <$> lookupType v table = let is' = reshapeIndex - (map pe32 oldshape) - (map pe32 $ newDims newshape) + (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' = - pe32 j : adjust js' is' + pe64 j : adjust js' is' adjust (DimSlice j _ s : js') (i : is') = - let i_t_s = i * pe32 s - j_p_i_t_s = pe32 j + i_t_s + 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 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 @@ -18,6 +18,8 @@ import Futhark.IR (ASTLore, Op, Prog, pretty) import qualified Futhark.IR.Kernels as Kernels import qualified Futhark.IR.KernelsMem as KernelsMem +import qualified Futhark.IR.MC as MC +import qualified Futhark.IR.MCMem as MCMem import Futhark.IR.Prop.Aliases (CanBeAliased) import qualified Futhark.IR.SOACS as SOACS import qualified Futhark.IR.Seq as Seq @@ -38,6 +40,7 @@ import qualified Futhark.Pass.ExplicitAllocations.Kernels as Kernels import qualified Futhark.Pass.ExplicitAllocations.Seq as Seq import Futhark.Pass.ExtractKernels +import Futhark.Pass.ExtractMulticore import Futhark.Pass.FirstOrderTransform import Futhark.Pass.KernelBabysitting import Futhark.Pass.Simplify @@ -91,8 +94,10 @@ data UntypedPassState = SOACS (Prog SOACS.SOACS) | Kernels (Prog Kernels.Kernels) + | MC (Prog MC.MC) | Seq (Prog Seq.Seq) | KernelsMem (Prog KernelsMem.KernelsMem) + | MCMem (Prog MCMem.MCMem) | SeqMem (Prog SeqMem.SeqMem) getSOACSProg :: UntypedPassState -> Maybe (Prog SOACS.SOACS) @@ -107,15 +112,19 @@ instance Representation UntypedPassState where representation (SOACS _) = "SOACS" representation (Kernels _) = "Kernels" + representation (MC _) = "MC" representation (Seq _) = "Seq" representation (KernelsMem _) = "KernelsMem" + representation (MCMem _) = "MCMem" representation (SeqMem _) = "SeqMEm" instance PP.Pretty UntypedPassState where ppr (SOACS prog) = PP.ppr prog ppr (Kernels prog) = PP.ppr prog + ppr (MC prog) = PP.ppr prog ppr (Seq prog) = PP.ppr prog ppr (SeqMem prog) = PP.ppr prog + ppr (MCMem prog) = PP.ppr prog ppr (KernelsMem prog) = PP.ppr prog newtype UntypedPass @@ -129,6 +138,7 @@ = SOACSAction (Action SOACS.SOACS) | KernelsAction (Action Kernels.Kernels) | KernelsMemAction (FilePath -> Action KernelsMem.KernelsMem) + | MCMemAction (FilePath -> Action MCMem.MCMem) | SeqMemAction (FilePath -> Action SeqMem.SeqMem) | PolyAction ( forall lore. @@ -144,12 +154,14 @@ untypedActionName (KernelsAction a) = actionName a untypedActionName (SeqMemAction a) = actionName $ a "" untypedActionName (KernelsMemAction a) = actionName $ a "" +untypedActionName (MCMemAction a) = actionName $ a "" untypedActionName (PolyAction a) = actionName (a :: Action SOACS.SOACS) instance Representation UntypedAction where representation (SOACSAction _) = "SOACS" representation (KernelsAction _) = "Kernels" representation (KernelsMemAction _) = "KernelsMem" + representation (MCMemAction _) = "MCMem" representation (SeqMemAction _) = "SeqMem" representation PolyAction {} = "" @@ -247,12 +259,16 @@ SOACS <$> runPipeline (onePass simplifySOACS) config prog perform (Kernels prog) config = Kernels <$> runPipeline (onePass simplifyKernels) config prog + perform (MC prog) config = + MC <$> runPipeline (onePass simplifyMC) 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 + perform (MCMem prog) config = + MCMem <$> runPipeline (onePass simplifyMCMem) config prog long = [passLongOption pass] pass = simplifySOACS @@ -299,12 +315,16 @@ SOACS <$> runPipeline (onePass $ performCSE True) config prog perform (Kernels prog) config = Kernels <$> runPipeline (onePass $ performCSE True) config prog + perform (MC prog) config = + MC <$> 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 + perform (MCMem prog) config = + MCMem <$> runPipeline (onePass $ performCSE False) config prog long = [passLongOption pass] pass = performCSE True :: Pass SOACS.SOACS SOACS.SOACS @@ -384,6 +404,14 @@ "Translate program into the imperative IL with kernels and write it on standard output.", Option [] + ["compile-imperative-multicore"] + ( NoArg $ + Right $ \opts -> + opts {futharkAction = MCMemAction $ const multicoreImpCodeGenAction} + ) + "Translate program into the imperative IL with kernels and write it on standard output.", + Option + [] ["compile-opencl"] ( NoArg $ Right $ \opts -> @@ -443,12 +471,13 @@ soacsPassOption inlineFunctions [], kernelsPassOption babysitKernels [], kernelsPassOption tileLoops [], - kernelsPassOption unstream [], - kernelsPassOption sink [], + kernelsPassOption unstreamKernels [], + kernelsPassOption sinkKernels [], typedPassOption soacsProg Kernels extractKernels [], + typedPassOption soacsProg MC extractMulticore [], iplOption [], allocateOption "a", - kernelsMemPassOption doubleBuffer [], + kernelsMemPassOption doubleBufferKernels [], kernelsMemPassOption expandAllocations [], cseOption [], simplifyOption "e", @@ -480,7 +509,15 @@ "Run the sequential CPU compilation pipeline" sequentialCpuPipeline [] - ["cpu"] + ["cpu"], + pipelineOption + getSOACSProg + "MCMem" + MCMem + "Run the multicore compilation pipeline" + multicorePipeline + [] + ["multicore"] ] incVerbosity :: Maybe FilePath -> FutharkConfig -> FutharkConfig @@ -593,15 +630,21 @@ actionProcedure (action base) prog (KernelsMem prog, KernelsMemAction action) -> actionProcedure (action base) prog + (MCMem prog, MCMemAction action) -> + actionProcedure (action base) prog (SOACS soacs_prog, PolyAction acs) -> actionProcedure acs soacs_prog (Kernels kernels_prog, PolyAction acs) -> actionProcedure acs kernels_prog + (MC mc_prog, PolyAction acs) -> + actionProcedure acs mc_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 + (MCMem mem_prog, PolyAction acs) -> actionProcedure acs mem_prog (_, action) -> externalErrorS $ diff --git a/src/Futhark/CLI/Multicore.hs b/src/Futhark/CLI/Multicore.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CLI/Multicore.hs @@ -0,0 +1,17 @@ +{-# LANGUAGE FlexibleContexts #-} + +module Futhark.CLI.Multicore (main) where + +import Futhark.Actions (compileMulticoreAction) +import Futhark.Compiler.CLI +import Futhark.Passes (multicorePipeline) + +main :: String -> [String] -> IO () +main = compilerMain + () + [] + "Compile to multicore C" + "Generate multicore C code from optimised Futhark program." + multicorePipeline + $ \fcfg () mode outpath prog -> + actionProcedure (compileMulticoreAction fcfg mode outpath) prog 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 @@ -579,7 +579,7 @@ { configBackend = "c", configFuthark = Nothing, configRunner = "", - configExtraOptions = [], + configExtraOptions = ["-b"], configExtraCompilerOptions = [], configTuning = Just "tuning" }, 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 @@ -392,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 @@ -442,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, 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 @@ -41,7 +41,8 @@ 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]] @@ -375,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. @@ -472,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, 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 @@ -33,12 +33,13 @@ -- * Monadic compiler interface CompilerM, - CompilerState (compUserState), + CompilerState (compUserState, compNameSrc), getUserState, modifyUserState, contextContents, contextFinalInits, runCompilerM, + inNewFunction, cachingMemory, blockScope, compileFun, @@ -64,9 +65,16 @@ publicName, contextType, contextField, + memToCType, + cacheMem, + fatMemory, + rawMemCType, + cproduct, + fatMemType, -- * Building Blocks primTypeToCType, + intTypeToCType, copyMemoryDefaultSpace, ) where @@ -209,6 +217,7 @@ 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 @@ -403,6 +412,22 @@ const w {accItems = mempty} ) +-- | Used when we, inside an existing 'CompilerM' action, want to +-- generate code for a new function. Use this so that the compiler +-- understands that previously declared memory doesn't need to be +-- freed inside this action. +inNewFunction :: Bool -> CompilerM op s a -> CompilerM op s a +inNewFunction keep_cached m = do + old_mem <- gets compDeclaredMem + modify $ \s -> s {compDeclaredMem = mempty} + x <- local noCached m + modify $ \s -> s {compDeclaredMem = old_mem} + return x + where + noCached env + | keep_cached = env + | otherwise = env {envCachedMem = mempty} + item :: C.BlockItem -> CompilerM op s () item x = tell $ mempty {accItems = DL.singleton x} @@ -675,7 +700,9 @@ 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))++; + if (rhs->references != NULL) { + (*(rhs->references))++; + } *lhs = *rhs; return ret; } @@ -1617,6 +1644,11 @@ let headerdefs = [C.cunit| $esc:("// Headers\n") +/* We need to define _GNU_SOURCE before + _any_ headers files are imported to get + the usage statistics of a thread (i.e. have RUSAGE_THREAD) on GNU/Linux + https://manpages.courier-mta.org/htmlman2/getrusage.2.html */ +$esc:("#define _GNU_SOURCE") $esc:("#include ") $esc:("#include ") $esc:("#include ") @@ -2097,7 +2129,7 @@ iexp' <- compileExp $ untyped iexp return [C.cexp|$id:src[$exp:iexp']|] compileLeaf (SizeOf t) = - return [C.cexp|(typename int64_t)sizeof($ty:t')|] + return [C.cexp|(typename int32_t)sizeof($ty:t')|] where t' = primTypeToCType t 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 @@ -1132,6 +1132,7 @@ 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 diff --git a/src/Futhark/CodeGen/Backends/MulticoreC.hs b/src/Futhark/CodeGen/Backends/MulticoreC.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/MulticoreC.hs @@ -0,0 +1,703 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} + +-- | C code generator. This module can convert a correct ImpCode +-- program to an equivalent C program. +module Futhark.CodeGen.Backends.MulticoreC + ( compileProg, + GC.CParts (..), + GC.asLibrary, + GC.asExecutable, + ) +where + +import Control.Monad +import Data.FileEmbed +import qualified Data.Map as M +import Data.Maybe +import qualified Futhark.CodeGen.Backends.GenericC as GC +import Futhark.CodeGen.Backends.GenericC.Options +import Futhark.CodeGen.Backends.SimpleRep +import Futhark.CodeGen.ImpCode.Multicore +import qualified Futhark.CodeGen.ImpGen.Multicore as ImpGen +import Futhark.IR.MCMem (MCMem, Prog) +import Futhark.MonadFreshNames +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C + +compileProg :: + MonadFreshNames m => + Prog MCMem -> + m (ImpGen.Warnings, GC.CParts) +compileProg = + traverse + ( GC.compileProg + "multicore" + operations + generateContext + "" + [DefaultSpace] + cliOptions + ) + <=< ImpGen.compileProg + where + generateContext = do + let multicore_defs_h = $(embedStringFile "rts/c/multicore_defs.h") + multicore_util_h = $(embedStringFile "rts/c/multicore_util.h") + subtask_queue_h = $(embedStringFile "rts/c/subtask_queue.h") + scheduler_common_h = $(embedStringFile "rts/c/scheduler_common.h") + scheduler_h = $(embedStringFile "rts/c/scheduler.h") + scheduler_tune_h = $(embedStringFile "rts/c/scheduler_tune.h") + + mapM_ + GC.earlyDecl + [C.cunit| + $esc:multicore_defs_h + $esc:multicore_util_h + $esc:subtask_queue_h + $esc:scheduler_common_h + $esc:scheduler_h + |] + + mapM_ GC.earlyDecl [C.cunit|int futhark_segred_tuning_program(struct futhark_context *ctx);|] + mapM_ GC.libDecl [C.cunit|$esc:scheduler_tune_h|] + + cfg <- GC.publicDef "context_config" GC.InitDecl $ \s -> + ( [C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { int debugging; int profiling; };|] + ) + + 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; + cfg->profiling = 0; + return cfg; + }|] + ) + + GC.publicDef_ "context_config_free" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:cfg* cfg);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg) { + free(cfg); + }|] + ) + + GC.publicDef_ "context_config_set_debugging" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int detail) { + cfg->debugging = detail; + }|] + ) + + GC.publicDef_ "context_config_set_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) { + 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 detail) { + /* Does nothing for this backend. */ + (void)cfg; (void)detail; + }|] + ) + + (fields, init_fields) <- GC.contextContents + + ctx <- GC.publicDef "context" GC.InitDecl $ \s -> + ( [C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { + struct scheduler scheduler; + int detail_memory; + int debugging; + int profiling; + int profiling_paused; + typename lock_t lock; + char *error; + int total_runs; + long int total_runtime; + $sdecls:fields + + // Tuning parameters + typename int64_t tuning_timing; + typename int64_t tuning_iter; + };|] + ) + + 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; + } + + // Initialize rand() + fast_srand(time(0)); + ctx->detail_memory = cfg->debugging; + ctx->debugging = cfg->debugging; + ctx->profiling = cfg->profiling; + ctx->profiling_paused = 0; + ctx->error = NULL; + create_lock(&ctx->lock); + ctx->scheduler.num_threads = num_processors(); + if (ctx->scheduler.num_threads < 1) return NULL; + + $stms:init_fields + + // futhark_segred_tuning_program(ctx); + + ctx->scheduler.workers = calloc(ctx->scheduler.num_threads, sizeof(struct worker)); + if (ctx->scheduler.workers == NULL) return NULL; + num_workers = ctx->scheduler.num_threads; + + worker_local = &ctx->scheduler.workers[0]; + worker_local->tid = 0; + worker_local->scheduler = &ctx->scheduler; + CHECK_ERR(subtask_queue_init(&worker_local->q, 1024), "failed to init queue for worker %d\n", 0); + + + for (int i = 1; i < ctx->scheduler.num_threads; i++) { + struct worker *cur_worker = &ctx->scheduler.workers[i]; + memset(cur_worker, 0, sizeof(struct worker)); + cur_worker->tid = i; + cur_worker->output_usage = 0; + cur_worker->scheduler = &ctx->scheduler; + CHECK_ERR(subtask_queue_init(&cur_worker->q, 1024), "failed to init queue for worker %d\n", i); + + CHECK_ERR(pthread_create(&cur_worker->thread, NULL, &scheduler_worker, + cur_worker), + "Failed to create worker %d\n", i); + } + + init_constants(ctx); + + + return ctx; + }|] + ) + + GC.publicDef_ "context_free" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:ctx* ctx);|], + [C.cedecl|void $id:s(struct $id:ctx* ctx) { + free_constants(ctx); + + // output_thread_usage(worker_local); + for (int i = 1; i < ctx->scheduler.num_threads; i++) + { + struct worker *cur_worker = &ctx->scheduler.workers[i]; + cur_worker->dead = 1; + subtask_queue_destroy(&cur_worker->q); + CHECK_ERR(pthread_join(ctx->scheduler.workers[i].thread, NULL), "pthread_join"); + } + + + free(ctx->scheduler.workers); + free_lock(&ctx->lock); + free(ctx); + }|] + ) + + GC.publicDef_ "context_sync" GC.InitDecl $ \s -> + ( [C.cedecl|int $id:s(struct $id:ctx* ctx);|], + [C.cedecl|int $id:s(struct $id:ctx* ctx) { + (void)ctx; + return 0; + }|] + ) + +cliOptions :: [Option] +cliOptions = + [ Option + { optionLongName = "profile", + optionShortName = Just 'P', + optionArgument = NoArgument, + optionAction = [C.cstm|futhark_context_config_set_profiling(cfg, 1);|], + optionDescription = "Gather profiling information." + } + ] + +operations :: GC.Operations Multicore () +operations = + GC.defaultOperations + { GC.opsCompiler = compileOp, + GC.opsCopy = copyMulticoreMemory, + GC.opsCritical = + -- The thread entering an API function is always considered + -- the "first worker" - note that this might differ from the + -- thread that created the context! This likely only matters + -- for entry points, since they are the only API functions + -- that contain parallel operations. + ( [C.citems|worker_local = &ctx->scheduler.workers[0];|], + [] + ) + } + +copyMulticoreMemory :: GC.Copy Multicore () +copyMulticoreMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes = + GC.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes +copyMulticoreMemory _ _ destspace _ _ srcspace _ = + error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace + +closureFreeStructField :: VName -> Name +closureFreeStructField v = + nameFromString "free_" <> nameFromString (pretty v) + +closureRetvalStructField :: VName -> Name +closureRetvalStructField v = + nameFromString "retval_" <> nameFromString (pretty v) + +data ValueType = Prim | MemBlock | RawMem + +compileFreeStructFields :: [VName] -> [(C.Type, ValueType)] -> [C.FieldGroup] +compileFreeStructFields = zipWith field + where + field name (ty, Prim) = + [C.csdecl|$ty:ty $id:(closureFreeStructField name);|] + field name (_, _) = + [C.csdecl|$ty:defaultMemBlockType $id:(closureFreeStructField name);|] + +compileRetvalStructFields :: [VName] -> [(C.Type, ValueType)] -> [C.FieldGroup] +compileRetvalStructFields = zipWith field + where + field name (ty, Prim) = + [C.csdecl|$ty:ty *$id:(closureRetvalStructField name);|] + field name (_, _) = + [C.csdecl|$ty:defaultMemBlockType $id:(closureRetvalStructField name);|] + +compileSetStructValues :: + C.ToIdent a => + a -> + [VName] -> + [(C.Type, ValueType)] -> + [C.Stm] +compileSetStructValues struct = zipWith field + where + field name (_, Prim) = + [C.cstm|$id:struct.$id:(closureFreeStructField name)=$id:name;|] + field name (_, MemBlock) = + [C.cstm|$id:struct.$id:(closureFreeStructField name)=$id:name.mem;|] + field name (_, RawMem) = + [C.cstm|$id:struct.$id:(closureFreeStructField name)=$id:name;|] + +compileSetRetvalStructValues :: + C.ToIdent a => + a -> + [VName] -> + [(C.Type, ValueType)] -> + [C.Stm] +compileSetRetvalStructValues struct = zipWith field + where + field name (_, Prim) = + [C.cstm|$id:struct.$id:(closureRetvalStructField name)=&$id:name;|] + field name (_, MemBlock) = + [C.cstm|$id:struct.$id:(closureRetvalStructField name)=$id:name.mem;|] + field name (_, RawMem) = + [C.cstm|$id:struct.$id:(closureRetvalStructField name)=$id:name;|] + +compileGetRetvalStructVals :: C.ToIdent a => a -> [VName] -> [(C.Type, ValueType)] -> [C.InitGroup] +compileGetRetvalStructVals struct = zipWith field + where + field name (ty, Prim) = + [C.cdecl|$ty:ty $id:name = *$id:struct->$id:(closureRetvalStructField name);|] + field name (ty, _) = + [C.cdecl|$ty:ty $id:name = + {.desc = $string:(pretty name), + .mem = $id:struct->$id:(closureRetvalStructField name), + .size = 0, .references = NULL};|] + +compileGetStructVals :: + C.ToIdent a => + a -> + [VName] -> + [(C.Type, ValueType)] -> + [C.InitGroup] +compileGetStructVals struct = zipWith field + where + field name (ty, Prim) = + [C.cdecl|$ty:ty $id:name = $id:struct->$id:(closureFreeStructField name);|] + field name (ty, _) = + [C.cdecl|$ty:ty $id:name = + {.desc = $string:(pretty name), + .mem = $id:struct->$id:(closureFreeStructField name), + .size = 0, .references = NULL};|] + +compileWriteBackResVals :: C.ToIdent a => a -> [VName] -> [(C.Type, ValueType)] -> [C.Stm] +compileWriteBackResVals struct = zipWith field + where + field name (_, Prim) = + [C.cstm|*$id:struct->$id:(closureRetvalStructField name) = $id:name;|] + field name (_, _) = + [C.cstm|$id:struct->$id:(closureRetvalStructField name) = $id:name.mem;|] + +paramToCType :: Param -> GC.CompilerM op s (C.Type, ValueType) +paramToCType (ScalarParam _ pt) = do + let t = GC.primTypeToCType pt + return (t, Prim) +paramToCType (MemParam name space') = mcMemToCType name space' + +mcMemToCType :: VName -> Space -> GC.CompilerM op s (C.Type, ValueType) +mcMemToCType v space = do + refcount <- GC.fatMemory space + cached <- isJust <$> GC.cacheMem v + return + ( GC.fatMemType space, + if refcount && not cached + then MemBlock + else RawMem + ) + +functionRuntime :: Name -> C.Id +functionRuntime = (`C.toIdent` mempty) . (<> "_total_runtime") + +functionRuns :: Name -> C.Id +functionRuns = (`C.toIdent` mempty) . (<> "_runs") + +functionIter :: Name -> C.Id +functionIter = (`C.toIdent` mempty) . (<> "_iter") + +multiCoreReport :: [(Name, Bool)] -> [C.BlockItem] +multiCoreReport names = report_kernels + where + report_kernels = concatMap reportKernel names + max_name_len_pad = 40 + format_string name True = + let name_s = nameToString name + padding = replicate (max_name_len_pad - length name_s) ' ' + in unwords ["tid %2d -", name_s ++ padding, "ran %10d times; avg: %10ldus; total: %10ldus; time pr. iter %9.6f; iters %9ld; avg %ld\n"] + format_string name False = + let name_s = nameToString name + padding = replicate (max_name_len_pad - length name_s) ' ' + in unwords [" ", name_s ++ padding, "ran %10d times; avg: %10ldus; total: %10ldus; time pr. iter %9.6f; iters %9ld; avg %ld\n"] + reportKernel (name, is_array) = + let runs = functionRuns name + total_runtime = functionRuntime name + iters = functionIter name + in if is_array + then + [ [C.citem| + for (int i = 0; i < ctx->scheduler.num_threads; i++) { + fprintf(stderr, + $string:(format_string name is_array), + i, + ctx->$id:runs[i], + (long int) ctx->$id:total_runtime[i] / (ctx->$id:runs[i] != 0 ? ctx->$id:runs[i] : 1), + (long int) ctx->$id:total_runtime[i], + (double) ctx->$id:total_runtime[i] / (ctx->$id:iters[i] == 0 ? 1 : (double)ctx->$id:iters[i]), + (long int) (ctx->$id:iters[i]), + (long int) (ctx->$id:iters[i]) / (ctx->$id:runs[i] != 0 ? ctx->$id:runs[i] : 1) + ); + } + |] + ] + else + [ [C.citem| + fprintf(stderr, + $string:(format_string name is_array), + ctx->$id:runs, + (long int) ctx->$id:total_runtime / (ctx->$id:runs != 0 ? ctx->$id:runs : 1), + (long int) ctx->$id:total_runtime, + (double) ctx->$id:total_runtime / (ctx->$id:iters == 0 ? 1 : (double)ctx->$id:iters), + (long int) (ctx->$id:iters), + (long int) (ctx->$id:iters) / (ctx->$id:runs != 0 ? ctx->$id:runs : 1)); + |], + [C.citem|ctx->total_runtime += ctx->$id:total_runtime;|], + [C.citem|ctx->total_runs += ctx->$id:runs;|] + ] + +addBenchmarkFields :: Name -> Maybe VName -> GC.CompilerM op s () +addBenchmarkFields name (Just _) = do + GC.contextField (functionRuntime name) [C.cty|typename int64_t*|] $ Just [C.cexp|calloc(sizeof(typename int64_t), ctx->scheduler.num_threads)|] + GC.contextField (functionRuns name) [C.cty|int*|] $ Just [C.cexp|calloc(sizeof(int), ctx->scheduler.num_threads)|] + GC.contextField (functionIter name) [C.cty|typename int64_t*|] $ Just [C.cexp|calloc(sizeof(sizeof(typename int64_t)), ctx->scheduler.num_threads)|] +addBenchmarkFields name Nothing = do + GC.contextField (functionRuntime name) [C.cty|typename int64_t|] $ Just [C.cexp|0|] + GC.contextField (functionRuns name) [C.cty|int|] $ Just [C.cexp|0|] + GC.contextField (functionIter name) [C.cty|typename int64_t|] $ Just [C.cexp|0|] + +benchmarkCode :: Name -> Maybe VName -> [C.BlockItem] -> GC.CompilerM op s [C.BlockItem] +benchmarkCode name tid code = do + addBenchmarkFields name tid + return + [C.citems| + typename uint64_t $id:start; + if (ctx->profiling && !ctx->profiling_paused) { + $id:start = get_wall_time(); + } + $items:code + if (ctx->profiling && !ctx->profiling_paused) { + typename uint64_t $id:end = get_wall_time(); + typename uint64_t elapsed = $id:end - $id:start; + $items:(updateFields tid) + } + |] + where + start = name <> "_start" + end = name <> "_end" + updateFields Nothing = + [C.citems|__atomic_fetch_add(&ctx->$id:(functionRuns name), 1, __ATOMIC_RELAXED); + __atomic_fetch_add(&ctx->$id:(functionRuntime name), elapsed, __ATOMIC_RELAXED); + __atomic_fetch_add(&ctx->$id:(functionIter name), iterations, __ATOMIC_RELAXED);|] + updateFields (Just _tid') = + [C.citems|ctx->$id:(functionRuns name)[tid]++; + ctx->$id:(functionRuntime name)[tid] += elapsed; + ctx->$id:(functionIter name)[tid] += iterations;|] + +functionTiming :: Name -> C.Id +functionTiming = (`C.toIdent` mempty) . (<> "_total_time") + +functionIterations :: Name -> C.Id +functionIterations = (`C.toIdent` mempty) . (<> "_total_iter") + +addTimingFields :: Name -> GC.CompilerM op s () +addTimingFields name = do + GC.contextField (functionTiming name) [C.cty|typename int64_t|] $ Just [C.cexp|0|] + GC.contextField (functionIterations name) [C.cty|typename int64_t|] $ Just [C.cexp|0|] + +multicoreName :: String -> GC.CompilerM op s Name +multicoreName s = do + s' <- newVName ("futhark_mc_" ++ s) + return $ nameFromString $ baseString s' ++ "_" ++ show (baseTag s') + +multicoreDef :: String -> (Name -> GC.CompilerM op s C.Definition) -> GC.CompilerM op s Name +multicoreDef s f = do + s' <- multicoreName s + GC.libDecl =<< f s' + return s' + +generateFunction :: + C.ToIdent a => + M.Map VName Space -> + String -> + Code -> + a -> + [(VName, (C.Type, ValueType))] -> + [(VName, (C.Type, ValueType))] -> + VName -> + VName -> + GC.CompilerM Multicore s Name +generateFunction lexical basename code fstruct free retval tid ntasks = do + let (fargs, fctypes) = unzip free + let (retval_args, retval_ctypes) = unzip retval + multicoreDef basename $ \s -> do + fbody <- benchmarkCode s (Just tid) <=< GC.inNewFunction False $ + GC.cachingMemory lexical $ + \decl_cached free_cached -> GC.blockScope $ do + mapM_ GC.item [C.citems|$decls:(compileGetStructVals fstruct fargs fctypes)|] + mapM_ GC.item [C.citems|$decls:(compileGetRetvalStructVals fstruct retval_args retval_ctypes)|] + mapM_ GC.item decl_cached + code' <- GC.blockScope $ GC.compileCode code + mapM_ GC.item [C.citems|$items:code'|] + mapM_ GC.stm free_cached + return + [C.cedecl|int $id:s(void *args, typename int64_t iterations, int tid, struct scheduler_info info) { + int err = 0; + int $id:tid = tid; + int $id:ntasks = info.nsubtasks; + struct $id:fstruct *$id:fstruct = (struct $id:fstruct*) args; + struct futhark_context *ctx = $id:fstruct->ctx; + $items:fbody + $stms:(compileWriteBackResVals fstruct retval_args retval_ctypes) + cleanup: {} + return err; + }|] + +prepareTaskStruct :: + String -> + [VName] -> + [(C.Type, ValueType)] -> + [VName] -> + [(C.Type, ValueType)] -> + GC.CompilerM Multicore s Name +prepareTaskStruct name free_args free_ctypes retval_args retval_ctypes = do + fstruct <- multicoreDef name $ \s -> + return + [C.cedecl|struct $id:s { + struct futhark_context *ctx; + $sdecls:(compileFreeStructFields free_args free_ctypes) + $sdecls:(compileRetvalStructFields retval_args retval_ctypes) + };|] + GC.decl [C.cdecl|struct $id:fstruct $id:fstruct;|] + GC.stm [C.cstm|$id:fstruct.ctx = ctx;|] + GC.stms [C.cstms|$stms:(compileSetStructValues fstruct free_args free_ctypes)|] + GC.stms [C.cstms|$stms:(compileSetRetvalStructValues fstruct retval_args retval_ctypes)|] + return fstruct + +-- Generate a segop function for top_level and potentially nested SegOp code +compileOp :: GC.OpCompiler Multicore () +compileOp (Segop name params seq_task par_task retvals (SchedulerInfo nsubtask e sched)) = do + let (ParallelTask seq_code tid) = seq_task + free_ctypes <- mapM paramToCType params + retval_ctypes <- mapM paramToCType retvals + let free_args = map paramName params + retval_args = map paramName retvals + free = zip free_args free_ctypes + retval = zip retval_args retval_ctypes + + e' <- GC.compileExp e + + let lexical = lexicalMemoryUsage $ Function False [] params seq_code [] [] + + fstruct <- + prepareTaskStruct "task" free_args free_ctypes retval_args retval_ctypes + + fpar_task <- generateFunction lexical (name ++ "_task") seq_code fstruct free retval tid nsubtask + addTimingFields fpar_task + + let ftask_name = fstruct <> "_task" + GC.decl [C.cdecl|struct scheduler_segop $id:ftask_name;|] + GC.stm [C.cstm|$id:ftask_name.args = &$id:fstruct;|] + GC.stm [C.cstm|$id:ftask_name.top_level_fn = $id:fpar_task;|] + GC.stm [C.cstm|$id:ftask_name.name = $string:(nameToString fpar_task);|] + GC.stm [C.cstm|$id:ftask_name.iterations = $exp:e';|] + -- Create the timing fields for the task + GC.stm [C.cstm|$id:ftask_name.task_time = &ctx->$id:(functionTiming fpar_task);|] + GC.stm [C.cstm|$id:ftask_name.task_iter = &ctx->$id:(functionIterations fpar_task);|] + + case sched of + Dynamic -> GC.stm [C.cstm|$id:ftask_name.sched = DYNAMIC;|] + Static -> GC.stm [C.cstm|$id:ftask_name.sched = STATIC;|] + + -- Generate the nested segop function if available + fnpar_task <- case par_task of + Just (ParallelTask nested_code nested_tid) -> do + let lexical_nested = lexicalMemoryUsage $ Function False [] params nested_code [] [] + fnpar_task <- generateFunction lexical_nested (name ++ "_nested_task") nested_code fstruct free retval nested_tid nsubtask + GC.stm [C.cstm|$id:ftask_name.nested_fn = $id:fnpar_task;|] + return $ zip [fnpar_task] [True] + Nothing -> do + GC.stm [C.cstm|$id:ftask_name.nested_fn=NULL;|] + return mempty + + let ftask_err = fpar_task <> "_err" + let code = + [C.citems|int $id:ftask_err = scheduler_prepare_task(&ctx->scheduler, &$id:ftask_name); + if ($id:ftask_err != 0) { + err = 1; goto cleanup; + }|] + + mapM_ GC.item code + + -- Add profile fields for -P option + mapM_ GC.profileReport $ multiCoreReport $ (fpar_task, True) : fnpar_task +compileOp (ParLoop s' i prebody body postbody free tid) = do + free_ctypes <- mapM paramToCType free + let free_args = map paramName free + + let lexical = + lexicalMemoryUsage $ + Function False [] free (prebody <> body) [] [] + + fstruct <- + prepareTaskStruct (s' ++ "_parloop_struct") free_args free_ctypes mempty mempty + + ftask <- multicoreDef (s' ++ "_parloop") $ \s -> do + fbody <- benchmarkCode s (Just tid) + <=< GC.inNewFunction False + $ GC.cachingMemory lexical $ + \decl_cached free_cached -> GC.blockScope $ do + mapM_ + GC.item + [C.citems|$decls:(compileGetStructVals fstruct free_args free_ctypes)|] + + mapM_ GC.item decl_cached + + GC.decl [C.cdecl|typename int64_t iterations = end - start;|] + GC.decl [C.cdecl|typename int64_t $id:i = start;|] + GC.compileCode prebody + body' <- GC.blockScope $ GC.compileCode body + GC.stm + [C.cstm|for (; $id:i < end; $id:i++) { + $items:body' + }|] + GC.compileCode postbody + GC.stm [C.cstm|cleanup: {}|] + mapM_ GC.stm free_cached + + return + [C.cedecl|int $id:s(void *args, typename int64_t start, typename int64_t end, int $id:tid, int tid) { + int err = 0; + struct $id:fstruct *$id:fstruct = (struct $id:fstruct*) args; + struct futhark_context *ctx = $id:fstruct->ctx; + $items:fbody + return err; + }|] + + let ftask_name = ftask <> "_task" + GC.decl [C.cdecl|struct scheduler_parloop $id:ftask_name;|] + GC.stm [C.cstm|$id:ftask_name.name = $string:(nameToString ftask);|] + GC.stm [C.cstm|$id:ftask_name.fn = $id:ftask;|] + GC.stm [C.cstm|$id:ftask_name.args = &$id:fstruct;|] + GC.stm [C.cstm|$id:ftask_name.iterations = iterations;|] + GC.stm [C.cstm|$id:ftask_name.info = info;|] + + let ftask_err = ftask <> "_err" + ftask_total = ftask <> "_total" + code' <- + benchmarkCode + ftask_total + Nothing + [C.citems|int $id:ftask_err = scheduler_execute_task(&ctx->scheduler, &$id:ftask_name); + if ($id:ftask_err != 0) { + err = 1; + goto cleanup; + }|] + + mapM_ GC.item code' + mapM_ GC.profileReport $ multiCoreReport $ zip [ftask, ftask_total] [True, False] +compileOp (Atomic aop) = + atomicOps aop + +doAtomic :: + (C.ToIdent a1, C.ToIdent a2) => + a1 -> + a2 -> + Count u (TExp Int32) -> + Exp -> + String -> + C.Type -> + GC.CompilerM op s () +doAtomic old arr ind val op ty = do + ind' <- GC.compileExp $ untyped $ unCount ind + val' <- GC.compileExp val + GC.stm [C.cstm|$id:old = $id:op(&(($ty:ty*)$id:arr.mem)[$exp:ind'], ($ty:ty) $exp:val', __ATOMIC_RELAXED);|] + +atomicOps :: AtomicOp -> GC.CompilerM op s () +atomicOps (AtomicCmpXchg t old arr ind res val) = do + ind' <- GC.compileExp $ untyped $ unCount ind + new_val' <- GC.compileExp val + let cast = [C.cty|$ty:(GC.primTypeToCType t)*|] + GC.stm + [C.cstm|$id:res = $id:op(&(($ty:cast)$id:arr.mem)[$exp:ind'], + ($ty:cast)&$id:old, + $exp:new_val', + 0, __ATOMIC_SEQ_CST, __ATOMIC_RELAXED);|] + where + op :: String + op = "__atomic_compare_exchange_n" +atomicOps (AtomicXchg t old arr ind val) = do + ind' <- GC.compileExp $ untyped $ unCount ind + val' <- GC.compileExp val + let cast = [C.cty|$ty:(GC.primTypeToCType t)*|] + GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast)$id:arr.mem)[$exp:ind'], $exp:val', __ATOMIC_SEQ_CST);|] + where + op :: String + op = "__atomic_exchange_n" +atomicOps (AtomicAdd t old arr ind val) = + doAtomic old arr ind val "__atomic_fetch_add" [C.cty|$ty:(GC.intTypeToCType t)|] +atomicOps (AtomicSub t old arr ind val) = + doAtomic old arr ind val "__atomic_fetch_sub" [C.cty|$ty:(GC.intTypeToCType t)|] +atomicOps (AtomicAnd t old arr ind val) = + doAtomic old arr ind val "__atomic_fetch_and" [C.cty|$ty:(GC.intTypeToCType t)|] +atomicOps (AtomicOr t old arr ind val) = + doAtomic old arr ind val "__atomic_fetch_or" [C.cty|$ty:(GC.intTypeToCType t)|] +atomicOps (AtomicXor t old arr ind val) = + doAtomic old arr ind val "__atomic_fetch_xor" [C.cty|$ty:(GC.intTypeToCType t)|] 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 @@ -82,6 +82,7 @@ onPart (ErrorString s) = formatEscape s onPart ErrorInt32 {} = "{}" + onPart ErrorInt64 {} = "{}" sizeClassesToPython :: M.Map Name SizeClass -> PyExp sizeClassesToPython = Dict . map f . M.toList 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 @@ -6,6 +6,7 @@ ( tupleField, funName, defaultMemBlockType, + intTypeToCType, primTypeToCType, signedPrimTypeToCType, 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 @@ -41,6 +41,7 @@ ErrorMsgPart (..), errorMsgArgTypes, ArrayContents (..), + declaredIn, lexicalMemoryUsage, calledFuncs, @@ -302,7 +303,8 @@ go f (Comment _ x) = f x go _ _ = mempty - declared (DeclareMem mem space) = M.singleton mem space + declared (DeclareMem mem space) = + M.singleton mem space declared x = go declared x set (SetMem x y _) = namesFromList [x, y] @@ -364,7 +366,7 @@ -- | Convert a count of elements into a count of bytes, given the -- per-element size. -withElemType :: Count Elements (TExp Int32) -> PrimType -> Count Bytes (TExp Int64) +withElemType :: Count Elements (TExp Int64) -> PrimType -> Count Bytes (TExp Int64) withElemType (Count e) t = bytes $ sExt64 e * isInt64 (LeafExp (SizeOf t) (IntType Int64)) 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 @@ -165,17 +165,17 @@ -- 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.TExp Int32)) Exp - | AtomicFAdd FloatType VName VName (Count Elements (Imp.TExp Int32)) Exp - | AtomicSMax IntType VName VName (Count Elements (Imp.TExp Int32)) Exp - | AtomicSMin IntType VName VName (Count Elements (Imp.TExp Int32)) Exp - | AtomicUMax IntType VName VName (Count Elements (Imp.TExp Int32)) Exp - | AtomicUMin IntType VName VName (Count Elements (Imp.TExp Int32)) Exp - | AtomicAnd IntType VName VName (Count Elements (Imp.TExp Int32)) Exp - | AtomicOr IntType VName VName (Count Elements (Imp.TExp Int32)) Exp - | AtomicXor IntType VName VName (Count Elements (Imp.TExp Int32)) Exp - | AtomicCmpXchg PrimType VName VName (Count Elements (Imp.TExp Int32)) Exp Exp - | AtomicXchg PrimType VName VName (Count Elements (Imp.TExp Int32)) Exp + = 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 diff --git a/src/Futhark/CodeGen/ImpCode/Multicore.hs b/src/Futhark/CodeGen/ImpCode/Multicore.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpCode/Multicore.hs @@ -0,0 +1,126 @@ +-- | Multicore imperative code. +module Futhark.CodeGen.ImpCode.Multicore + ( Program, + Function, + FunctionT (Function), + Code, + Multicore (..), + Scheduling (..), + SchedulerInfo (..), + AtomicOp (..), + ParallelTask (..), + module Futhark.CodeGen.ImpCode, + ) +where + +import Futhark.CodeGen.ImpCode hiding (Code, Function) +import qualified Futhark.CodeGen.ImpCode as Imp +import Futhark.Util.Pretty + +-- | An imperative program. +type Program = Imp.Functions Multicore + +-- | An imperative function. +type Function = Imp.Function Multicore + +-- | A piece of imperative code, with multicore operations inside. +type Code = Imp.Code Multicore + +-- | A multicore operation. +data Multicore + = Segop String [Param] ParallelTask (Maybe ParallelTask) [Param] SchedulerInfo + | ParLoop String VName Code Code Code [Param] VName + | Atomic AtomicOp + +-- | 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.TExp Int32)) Exp + | AtomicSub IntType VName VName (Count Elements (Imp.TExp Int32)) Exp + | AtomicAnd IntType VName VName (Count Elements (Imp.TExp Int32)) Exp + | AtomicOr IntType VName VName (Count Elements (Imp.TExp Int32)) Exp + | AtomicXor IntType VName VName (Count Elements (Imp.TExp Int32)) Exp + | AtomicXchg PrimType VName VName (Count Elements (Imp.TExp Int32)) Exp + | AtomicCmpXchg PrimType VName VName (Count Elements (Imp.TExp Int32)) VName Exp + deriving (Show) + +instance FreeIn AtomicOp where + freeIn' (AtomicAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicSub _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicAnd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicOr _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicXor _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicCmpXchg _ _ arr i retval x) = freeIn' arr <> freeIn' i <> freeIn' x <> freeIn' retval + freeIn' (AtomicXchg _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + +data SchedulerInfo = SchedulerInfo + { nsubtasks :: VName, -- The variable that describes how many subtasks the scheduler created + iterations :: Imp.Exp, -- The number of total iterations for a task + scheduling :: Scheduling -- The type scheduling for the task + } + +data ParallelTask = ParallelTask + { task_code :: Code, + flatTid :: VName -- The variable for the thread id execution the code + } + +-- | Whether the Scheduler should schedule the tasks as Dynamic +-- or it is restainted to Static +data Scheduling + = Dynamic + | Static + +instance Pretty Scheduling where + ppr Dynamic = text "Dynamic" + ppr Static = text "Static" + +-- TODO fix all of this! +instance Pretty SchedulerInfo where + ppr (SchedulerInfo nsubtask i sched) = + text "SchedulingInfo" + <+> text "number of subtasks" + <+> ppr nsubtask + <+> text "scheduling" + <+> ppr sched + <+> text "iter" + <+> ppr i + +instance Pretty ParallelTask where + ppr (ParallelTask code _) = + ppr code + +instance Pretty Multicore where + ppr (Segop s free _par_code seq_code retval scheduler) = + text "parfor" + <+> ppr scheduler + <+> ppr free + <+> text s + <+> text "seq_code" + <+> nestedBlock "{" "}" (ppr seq_code) + <+> text "retvals" + <+> ppr retval + ppr (ParLoop s i prebody body postbody params info) = + text "parloop" <+> ppr s <+> ppr i + <+> ppr prebody + <+> ppr params + <+> ppr info + <+> langle + <+> nestedBlock "{" "}" (ppr body) + <+> ppr postbody + ppr (Atomic _) = text "AtomicOp" + +instance FreeIn SchedulerInfo where + freeIn' (SchedulerInfo nsubtask iter _) = + freeIn' iter <> freeIn' nsubtask + +instance FreeIn ParallelTask where + freeIn' (ParallelTask code _) = + freeIn' code + +instance FreeIn Multicore where + freeIn' (Segop _ _ par_code seq_code _ info) = + freeIn' par_code <> freeIn' seq_code <> freeIn' info + freeIn' (ParLoop _ _ prebody body postbody _ _) = + freeIn' prebody <> fvBind (Imp.declaredIn prebody) (freeIn' $ body <> postbody) + freeIn' (Atomic aop) = freeIn' aop 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 @@ -156,9 +156,9 @@ type CopyCompiler lore r op = PrimType -> MemLocation -> - Slice (Imp.TExp Int32) -> + Slice (Imp.TExp Int64) -> MemLocation -> - Slice (Imp.TExp Int32) -> + Slice (Imp.TExp Int64) -> ImpM lore r op () -- | An alternate way of compiling an allocation. @@ -191,7 +191,7 @@ data MemLocation = MemLocation { memLocationName :: VName, memLocationShape :: [Imp.DimSize], - memLocationIxFun :: IxFun.IxFun (Imp.TExp Int32) + memLocationIxFun :: IxFun.IxFun (Imp.TExp Int64) } deriving (Eq, Show) @@ -621,7 +621,7 @@ Nothing -> do out <- imp $ newVName "out_arrsize" tell - ( [Imp.ScalarParam out int32], + ( [Imp.ScalarParam out int64], M.singleton x $ ScalarDestination out ) put (memseen, M.insert x out arrseen) @@ -773,7 +773,7 @@ ForLoop i _ bound loopvars -> do let setLoopParam (p, a) | Prim _ <- paramType p = - copyDWIM (paramName p) [] (Var a) [DimFix $ Imp.vi32 i] + copyDWIM (paramName p) [] (Var a) [DimFix $ Imp.vi64 i] | otherwise = return () @@ -828,22 +828,22 @@ 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 . toInt32Exp) idxs + copyDWIM (patElemName pe) [] (Var src) $ map (DimFix . toInt64Exp) idxs defCompileBasicOp _ Index {} = return () defCompileBasicOp (Pattern _ [pe]) (Update _ slice se) = - sUpdate (patElemName pe) (map (fmap toInt32Exp) 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.vi32) is) se [] + 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" (toInt32Exp n) $ \i -> do + sFor "i" (toInt64Exp n) $ \i -> do let i' = sExt it $ untyped i x <- dPrimV "x" $ @@ -856,16 +856,16 @@ defCompileBasicOp (Pattern _ [pe]) (Manifest _ src) = copyDWIM (patElemName pe) [] (Var src) [] defCompileBasicOp (Pattern _ [pe]) (Concat i x ys _) = do - offs_glb <- dPrimV "tmp_offs" (0 :: Imp.TExp Int32) + 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 : _ -> toInt32Exp r + r : _ -> toInt64Exp r skip_dims = take i y_dims sliceAllDim d = DimSlice 0 d 1 - skip_slices = map (sliceAllDim . toInt32Exp) skip_dims + 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 @@ -877,7 +877,7 @@ 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] $ + MemLocation static_array [intConst Int64 $ fromIntegral $ length es] $ IxFun.iota [fromIntegral $ length es] entry = MemVar Nothing $ MemEntry dest_space addVar static_array entry @@ -1216,7 +1216,7 @@ fullyIndexArray :: VName -> - [Imp.TExp Int32] -> + [Imp.TExp Int64] -> ImpM lore r op (VName, Imp.Space, Count Elements (Imp.TExp Int64)) fullyIndexArray name indices = do arr <- lookupArray name @@ -1224,7 +1224,7 @@ fullyIndexArray' :: MemLocation -> - [Imp.TExp Int32] -> + [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 @@ -1233,13 +1233,10 @@ let (zero_is, is) = splitFromEnd (length ds) indices in map (const 0) zero_is ++ is _ -> indices - - ixfun64 = fmap sExt64 ixfun - indices64 = fmap sExt64 indices' return ( mem, space, - elements $ IxFun.index ixfun64 indices64 + elements $ IxFun.index ixfun indices' ) -- More complicated read/write operations that use index functions. @@ -1253,15 +1250,15 @@ isMapTransposeCopy :: PrimType -> MemLocation -> - Slice (Imp.TExp Int32) -> + Slice (Imp.TExp Int64) -> MemLocation -> - Slice (Imp.TExp Int32) -> + Slice (Imp.TExp Int64) -> Maybe - ( Imp.TExp Int32, - Imp.TExp Int32, - Imp.TExp Int32, - Imp.TExp Int32, - Imp.TExp Int32 + ( Imp.TExp Int64, + Imp.TExp Int64, + Imp.TExp Int64, + Imp.TExp Int64, + Imp.TExp Int64 ) isMapTransposeCopy bt @@ -1334,16 +1331,16 @@ $ transposeArgs pt destmem - (bytes $ sExt64 destoffset) + (bytes destoffset) srcmem - (bytes $ sExt64 srcoffset) - (sExt64 num_arrays) - (sExt64 size_x) - (sExt64 size_y) + (bytes srcoffset) + num_arrays + size_x + size_y | Just destoffset <- - IxFun.linearWithOffset (IxFun.slice dest_ixfun64 destslice64) pt_size, + IxFun.linearWithOffset (IxFun.slice dest_ixfun destslice) pt_size, Just srcoffset <- - IxFun.linearWithOffset (IxFun.slice src_ixfun64 srcslice64) pt_size = do + IxFun.linearWithOffset (IxFun.slice src_ixfun srcslice) pt_size = do srcspace <- entryMemSpace <$> lookupMemory srcmem destspace <- entryMemSpace <$> lookupMemory destmem if isScalarSpace srcspace || isScalarSpace destspace @@ -1367,11 +1364,6 @@ MemLocation destmem _ dest_ixfun = dest MemLocation srcmem _ src_ixfun = src - dest_ixfun64 = fmap sExt64 dest_ixfun - destslice64 = map (fmap sExt64) destslice - src_ixfun64 = fmap sExt64 src_ixfun - srcslice64 = map (fmap sExt64) srcslice - isScalarSpace ScalarSpace {} = True isScalarSpace _ = False @@ -1379,7 +1371,7 @@ copyElementWise bt dest destslice src srcslice = do let bounds = sliceDims srcslice is <- replicateM (length bounds) (newVName "i") - let ivars = map Imp.vi32 is + let ivars = map Imp.vi64 is (destmem, destspace, destidx) <- fullyIndexArray' dest $ fixSlice destslice ivars (srcmem, srcspace, srcidx) <- @@ -1395,9 +1387,9 @@ copyArrayDWIM :: PrimType -> MemLocation -> - [DimIndex (Imp.TExp Int32)] -> + [DimIndex (Imp.TExp Int64)] -> MemLocation -> - [DimIndex (Imp.TExp Int32)] -> + [DimIndex (Imp.TExp Int64)] -> ImpM lore r op (Imp.Code op) copyArrayDWIM bt @@ -1419,9 +1411,9 @@ Imp.index srcmem srcoffset bt srcspace vol | otherwise = do let destslice' = - fullSliceNum (map toInt32Exp destshape) destslice + fullSliceNum (map toInt64Exp destshape) destslice srcslice' = - fullSliceNum (map toInt32Exp srcshape) srcslice + fullSliceNum (map toInt64Exp srcshape) srcslice destrank = length $ sliceDims destslice' srcrank = length $ sliceDims srcslice' if destrank /= srcrank @@ -1445,9 +1437,9 @@ -- instead of a variable name. copyDWIMDest :: ValueDestination -> - [DimIndex (Imp.TExp Int32)] -> + [DimIndex (Imp.TExp Int64)] -> SubExp -> - [DimIndex (Imp.TExp Int32)] -> + [DimIndex (Imp.TExp Int64)] -> ImpM lore r op () copyDWIMDest _ _ (Constant v) (_ : _) = error $ @@ -1539,9 +1531,9 @@ -- Thing. Both destination and source must be in scope. copyDWIM :: VName -> - [DimIndex (Imp.TExp Int32)] -> + [DimIndex (Imp.TExp Int64)] -> SubExp -> - [DimIndex (Imp.TExp Int32)] -> + [DimIndex (Imp.TExp Int64)] -> ImpM lore r op () copyDWIM dest dest_slice src src_slice = do dest_entry <- lookupVar dest @@ -1558,9 +1550,9 @@ -- | As 'copyDWIM', but implicitly 'DimFix'es the indexes. copyDWIMFix :: VName -> - [Imp.TExp Int32] -> + [Imp.TExp Int64] -> SubExp -> - [Imp.TExp Int32] -> + [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) @@ -1589,7 +1581,7 @@ typeSize t = Imp.bytes $ isInt64 (Imp.LeafExp (Imp.SizeOf $ elemType t) int64) - * product (map (sExt64 . toInt32Exp) (arrayDims t)) + * product (map toInt64Exp (arrayDims t)) --- Building blocks for constructing code. @@ -1664,14 +1656,14 @@ sArrayInMem name pt shape mem = sArray name pt shape $ ArrayIn mem $ - IxFun.iota $ map (isInt32 . primExpFromSubExp int32) $ shapeDims shape + 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 (isInt32 . primExpFromSubExp int32) permuted_dims + let iota_ixfun = IxFun.iota $ map (isInt64 . primExpFromSubExp int64) permuted_dims sArray name pt shape $ ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm @@ -1686,30 +1678,30 @@ let num_elems = case vs of Imp.ArrayValues vs' -> length vs' Imp.ArrayZeros n -> fromIntegral n - shape = Shape [intConst Int32 $ toInteger num_elems] + 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 Int32] -> Imp.Exp -> ImpM lore r op () +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 Int32) -> SubExp -> ImpM lore r op () +sUpdate :: VName -> Slice (Imp.TExp Int64) -> SubExp -> ImpM lore r op () sUpdate arr slice v = copyDWIM arr slice v [] sLoopNest :: Shape -> - ([Imp.TExp Int32] -> ImpM lore r op ()) -> + ([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" (toInt32Exp d) $ \i -> sLoopNest' (i : is) ds f + sFor "nest_i" (toInt64Exp d) $ \i -> sLoopNest' (i : is) ds f -- | Untyped assignment. (<~~) :: VName -> Imp.Exp -> ImpM lore r op () 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 @@ -188,7 +188,7 @@ x' <- toExp x s' <- toExp s - sIota (patElemName pe) (toInt32Exp 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. @@ -243,7 +243,7 @@ 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 toInt32Exp srcshape + let num_elems = Imp.elements $ product $ map toInt64Exp srcshape srcspace <- entryMemSpace <$> lookupMemory srcmem destspace <- entryMemSpace <$> lookupMemory destmem emit $ 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 @@ -68,8 +68,8 @@ kernelGlobalThreadIdVar :: VName, kernelLocalThreadIdVar :: VName, kernelGroupIdVar :: VName, - kernelNumGroups :: Imp.TExp Int32, - kernelGroupSize :: Imp.TExp Int32, + kernelNumGroups :: Imp.TExp Int64, + kernelGroupSize :: Imp.TExp Int64, kernelNumThreads :: Imp.TExp Int32, kernelWaveSize :: Imp.TExp Int32, kernelThreadActive :: Imp.TExp Bool, @@ -102,7 +102,7 @@ localEnv f m where mkMap ltid dims = do - let dims' = map toInt32Exp dims + let dims' = map (sExt32 . toInt64Exp) dims ids' <- mapM (dPrimVE "ltid_pre") $ unflattenIndex dims' ltid return (dims, ids') @@ -140,16 +140,16 @@ ImpM lore r op () splitSpace (Pattern [] [size]) o w i elems_per_thread = do num_elements <- Imp.elements . TPrimExp <$> toExp w - let i' = toInt32Exp i + 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) int32) + 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 :: Int32)] e [] + copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] compileThreadExp dest e = defCompileExp dest e @@ -179,13 +179,13 @@ -- passed-in function is invoked with the (symbolic) iteration. For -- multidimensional loops, use 'groupCoverSpace'. groupLoop :: - Imp.TExp Int32 -> - (Imp.TExp Int32 -> InKernelGen ()) -> + Imp.TExp Int64 -> + (Imp.TExp Int64 -> InKernelGen ()) -> InKernelGen () groupLoop n f = do constants <- kernelConstants <$> askEnv kernelLoop - (kernelLocalThreadId constants) + (sExt64 $ kernelLocalThreadId constants) (kernelGroupSize constants) n f @@ -194,8 +194,8 @@ -- all threads in the group participate. The passed-in function is -- invoked with a (symbolic) point in the index space. groupCoverSpace :: - [Imp.TExp Int32] -> - ([Imp.TExp Int32] -> InKernelGen ()) -> + [Imp.TExp Int64] -> + ([Imp.TExp Int64] -> InKernelGen ()) -> InKernelGen () groupCoverSpace ds f = groupLoop (product ds) $ f . unflattenIndex ds @@ -204,9 +204,9 @@ -- 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 [] + copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] compileGroupExp (Pattern _ [dest]) (BasicOp (Replicate ds se)) = do - let ds' = map toInt32Exp $ shapeDims ds + let ds' = map toInt64Exp $ shapeDims ds groupCoverSpace ds' $ \is -> copyDWIMFix (patElemName dest) is se (drop (shapeRank ds) is) sOp $ Imp.Barrier Imp.FenceLocal @@ -232,7 +232,7 @@ sOp $ Imp.Barrier Imp.FenceLocal ltid <- kernelLocalThreadId . kernelConstants <$> askEnv sWhen (ltid .==. 0) $ - copyDWIM (patElemName pe) (map (fmap toInt32Exp) slice) se [] + copyDWIM (patElemName pe) (map (fmap toInt64Exp) slice) se [] sOp $ Imp.Barrier Imp.FenceLocal compileGroupExp dest e = defCompileExp dest e @@ -242,11 +242,11 @@ sanityCheckLevel SegGroup {} = error "compileGroupOp: unexpected group-level SegOp." -localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int32] +localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64] localThreadIDs dims = do - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - let dims' = map toInt32Exp dims - fromMaybe (unflattenIndex dims' ltid) + ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv + let dims' = map toInt64Exp dims + maybe (unflattenIndex dims' ltid) (map sExt64) . M.lookup dims . kernelLocalIdMap . kernelConstants @@ -264,7 +264,7 @@ prepareIntraGroupSegHist :: Count GroupSize SubExp -> [HistOp KernelsMem] -> - InKernelGen [[Imp.TExp Int32] -> InKernelGen ()] + InKernelGen [[Imp.TExp Int64] -> InKernelGen ()] prepareIntraGroupSegHist group_size = fmap snd . mapAccumLM onOp Nothing where @@ -281,8 +281,8 @@ (Nothing, AtomicLocking f) -> do locks <- newVName "locks" - let num_locks = toInt32Exp $ unCount group_size - dims = map toInt32Exp $ shapeDims (histShape op) ++ [histWidth op] + 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 @@ -290,7 +290,7 @@ dArray locks int32 (arrayShape locks_t) $ ArrayIn locks_mem $ IxFun.iota $ - map pe32 $ arrayDims locks_t + map pe64 $ arrayDims locks_t sComment "All locks start out unlocked" $ groupCoverSpace [kernelGroupSize constants] $ \is -> @@ -321,21 +321,22 @@ compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do compileGroupSpace lvl space let (ltids, dims) = unzip $ unSegSpace space - dims' = map toInt32Exp dims + dims' = map toInt64Exp dims whenActive lvl space $ compileStms mempty (kernelBodyStms body) $ forM_ (zip (patternNames pat) $ kernelBodyResult body) $ \(dest, res) -> copyDWIMFix dest - (map Imp.vi32 ltids) + (map Imp.vi64 ltids) (kernelResultSubExp res) [] sOp $ Imp.ErrorSync Imp.FenceLocal let segment_size = last dims' - crossesSegment from to = (to - from) .>. (to `rem` segment_size) + 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 @@ -351,7 +352,7 @@ (baseString (patElemName pe) ++ "_flat") (elemType pe_t) (Shape arr_dims) - $ ArrayIn mem $ IxFun.iota $ map pe32 arr_dims + $ ArrayIn mem $ IxFun.iota $ map pe64 arr_dims num_scan_results = sum $ map (length . segBinOpNeutral) scans @@ -367,7 +368,7 @@ (red_pes, map_pes) = splitAt (segBinOpResults ops) $ patternElements pat - dims' = map toInt32Exp dims + dims' = map toInt64Exp dims mkTempArr t = sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local" @@ -380,7 +381,7 @@ let (red_res, map_res) = splitAt (segBinOpResults ops) $ kernelBodyResult body forM_ (zip tmp_arrs red_res) $ \(dest, res) -> - copyDWIMFix dest (map Imp.vi32 ltids) (kernelResultSubExp res) [] + copyDWIMFix dest (map Imp.vi64 ltids) (kernelResultSubExp res) [] zipWithM_ (compileThreadResult space) map_pes map_res sOp $ Imp.ErrorSync Imp.FenceLocal @@ -390,7 +391,7 @@ -- handle directly with a group-level reduction. [dim'] -> do forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> - groupReduce dim' (segBinOpLambda op) tmps + groupReduce (sExt32 dim') (segBinOpLambda op) tmps sOp $ Imp.ErrorSync Imp.FenceLocal @@ -413,10 +414,11 @@ drop (length ltids) (memLocationShape arr_loc) sArray "red_arr_flat" pt flat_shape $ ArrayIn (memLocationName arr_loc) $ - IxFun.iota $ map pe32 $ shapeDims flat_shape + IxFun.iota $ map pe64 $ shapeDims flat_shape let segment_size = last dims' - crossesSegment from to = (to - from) .>. (to `rem` segment_size) + 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 @@ -463,10 +465,10 @@ forM_ (zip4 red_is vs_per_op ops' ops) $ \(bin, op_vs, do_op, HistOp dest_w _ _ _ shape lam) -> do - let bin' = toInt32Exp bin - dest_w' = toInt32Exp dest_w + let bin' = toInt64Exp bin + dest_w' = toInt64Exp dest_w bin_in_bounds = 0 .<=. bin' .&&. bin' .<. dest_w' - bin_is = map Imp.vi32 (init ltids) ++ [bin'] + bin_is = map Imp.vi64 (init ltids) ++ [bin'] vs_params = takeLast (length op_vs) $ lambdaParams lam sComment "perform atomic updates" $ @@ -502,13 +504,13 @@ -- | 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 Int32] -> [Imp.TExp Int32] + 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 Int32] -> ImpM lore r Imp.KernelOp () + 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 @@ -524,7 +526,7 @@ -- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? type AtomicBinOp = BinOp -> - Maybe (VName -> VName -> Count Imp.Elements (Imp.TExp Int32) -> Imp.Exp -> Imp.AtomicOp) + Maybe (VName -> VName -> Count Imp.Elements (Imp.TExp Int64) -> Imp.Exp -> Imp.AtomicOp) -- | Do an atomic update corresponding to a binary operator lambda. atomicUpdateLocking :: @@ -546,7 +548,7 @@ (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket - case opHasAtomicSupport space (tvVar old) arr' (sExt32 <$> bucket_offset) op of + 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 $ @@ -588,7 +590,7 @@ int32 (tvVar old) locks' - (sExt32 <$> locks_offset) + locks_offset (untyped $ lockingIsUnlocked locking) (untyped $ lockingToLock locking) lock_acquired = tvExp old .==. lockingIsUnlocked locking @@ -601,7 +603,7 @@ int32 (tvVar old) locks' - (sExt32 <$> locks_offset) + locks_offset (untyped $ lockingToLock locking) (untyped $ lockingToUnlock locking) break_loop = continue <-- false @@ -656,7 +658,7 @@ PrimType -> VName -> VName -> - [Imp.TExp Int32] -> + [Imp.TExp Int64] -> VName -> InKernelGen () -> InKernelGen () @@ -698,7 +700,7 @@ int32 (tvVar old_bits) arr' - (sExt32 <$> bucket_offset) + bucket_offset (toBits (Imp.var assumed t)) (toBits (Imp.var x t)) old <~~ fromBits (untyped $ tvExp old_bits) @@ -773,16 +775,16 @@ computeThreadChunkSize :: SplitOrdering -> - Imp.TExp Int32 -> - Imp.Count Imp.Elements (Imp.TExp Int32) -> - Imp.Count Imp.Elements (Imp.TExp Int32) -> - TV Int32 -> + 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 - <-- sMin32 + <-- sMin64 (Imp.unCount elements_per_thread) - ((Imp.unCount num_elements - thread_index) `divUp` toInt32Exp stride) + ((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" $ @@ -796,7 +798,7 @@ sIf (no_remaining_elements .||. beyond_bounds) - (chunk_var <-- (0 :: Imp.TExp Int32)) + (chunk_var <-- 0) ( sIf is_last_thread (chunk_var <-- Imp.unCount last_thread_elements) @@ -810,8 +812,8 @@ .<. (thread_index + 1) * Imp.unCount elements_per_thread kernelInitialisationSimple :: - Count NumGroups (Imp.TExp Int32) -> - Count GroupSize (Imp.TExp Int32) -> + 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" @@ -829,7 +831,7 @@ group_id num_groups group_size - (group_size * num_groups) + (sExt32 (group_size * num_groups)) (Imp.vi32 wave_size) true mempty @@ -837,7 +839,7 @@ let set_constants = do dPrim_ global_tid int32 dPrim_ local_tid int32 - dPrim_ inner_group_size int32 + dPrim_ inner_group_size int64 dPrim_ wave_size int32 dPrim_ group_id int32 @@ -855,8 +857,8 @@ x : xs -> foldl (.&&.) x xs where (is, ws) = unzip limit - actives = zipWith active is $ map toInt32Exp ws - active i = (Imp.vi32 i .<.) + 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 @@ -901,20 +903,20 @@ readReduceArgument param arr | Prim _ <- paramType param = do let i = local_tid + tvExp offset - copyDWIMFix (paramName param) [] (Var arr) [i] + copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] | otherwise = do let i = global_tid + tvExp offset - copyDWIMFix (paramName param) [] (Var arr) [i] + copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] writeReduceOpResult param arr | Prim _ <- paramType param = - copyDWIMFix arr [local_tid] (Var $ paramName 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 <- dPrim "skip_waves" int32 + skip_waves <- dPrimV "skip_waves" (1 :: Imp.TExp Int32) dLParams $ lambdaParams lam offset <-- (0 :: Imp.TExp Int32) @@ -936,7 +938,7 @@ 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 + num_waves = (sExt32 group_size + wave_size - 1) `quot` wave_size arg_in_bounds = local_tid + tvExp offset .<. w doing_in_wave_reductions = @@ -959,8 +961,7 @@ (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 = do - skip_waves <-- (1 :: Imp.TExp Int32) + cross_wave_reductions = sWhile doing_cross_wave_reductions $ do barrier offset <-- tvExp skip_waves * wave_size @@ -974,8 +975,8 @@ groupScan :: Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> - Imp.TExp Int32 -> - Imp.TExp Int32 -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> Lambda KernelsMem -> [VName] -> InKernelGen () @@ -983,11 +984,14 @@ constants <- kernelConstants <$> askEnv renamed_lam <- renameLambda lam - let ltid = kernelLocalThreadId constants + 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 @@ -1000,8 +1004,8 @@ -- it were a runtime parameter. Some day. let block_size = 32 simd_width = kernelWaveSize constants - block_id = ltid `quot` block_size - in_block_id = ltid - block_id * block_size + block_id = ltid32 `quot` block_size + in_block_id = ltid32 - block_id * block_size doInBlockScan seg_flag' active = inBlockScan constants @@ -1012,7 +1016,6 @@ active arrs barrier - ltid_in_bounds = ltid .<. w array_scan = not $ all primType $ lambdaReturnType lam barrier | array_scan = @@ -1020,19 +1023,19 @@ | otherwise = sOp $ Imp.Barrier Imp.FenceLocal - group_offset = kernelGroupId constants * kernelGroupSize constants + group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants writeBlockResult p arr | primType $ paramType p = - copyDWIM arr [DimFix block_id] (Var $ paramName p) [] + copyDWIM arr [DimFix $ sExt64 block_id] (Var $ paramName p) [] | otherwise = - copyDWIM arr [DimFix $ group_offset + block_id] (Var $ paramName p) [] + copyDWIM arr [DimFix $ group_offset + sExt64 block_id] (Var $ paramName p) [] readPrevBlockResult p arr | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix $ block_id - 1] + copyDWIM (paramName p) [] (Var arr) [DimFix $ sExt64 block_id - 1] | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + block_id - 1] + copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + sExt64 block_id - 1] doInBlockScan seg_flag ltid_in_bounds lam barrier @@ -1043,7 +1046,7 @@ 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) [] + copyDWIM arr [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) [] barrier @@ -1074,7 +1077,7 @@ arr [DimFix $ arrs_full_size + group_offset + ltid] (Var arr) - [DimFix $ arrs_full_size + group_offset + block_size + ltid] + [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] barrier @@ -1092,7 +1095,7 @@ compileBody' x_params $ lambdaBody lam | Just flag_true <- seg_flag = do inactive <- - dPrimVE "inactive" $ flag_true (block_id * block_size -1) ltid + 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 @@ -1122,7 +1125,7 @@ inBlockScan :: KernelConstants -> Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> - Imp.TExp Int32 -> + Imp.TExp Int64 -> Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool -> @@ -1158,7 +1161,7 @@ | Just flag_true <- seg_flag = do inactive <- dPrimVE "inactive" $ - flag_true (ltid - tvExp skip_threads) ltid + flag_true (ltid32 - tvExp skip_threads) ltid32 sWhen inactive y_to_x when array_scan barrier sUnless inactive $ compileBody' x_params $ lambdaBody scan_lam @@ -1169,11 +1172,11 @@ barrier sComment "in-block scan (hopefully no barriers needed)" $ do - skip_threads <-- (1 :: Imp.TExp Int32) + skip_threads <-- 1 sWhile (tvExp skip_threads .<. block_size) $ do sWhen (in_block_thread_active .&&. active) $ do sComment "read operands" $ - zipWithM_ (readParam (tvExp skip_threads)) x_params arrs + zipWithM_ (readParam (sExt64 $ tvExp skip_threads)) x_params arrs sComment "perform operation" op_to_x maybeBarrier @@ -1186,10 +1189,11 @@ skip_threads <-- tvExp skip_threads * 2 where - block_id = ltid `quot` block_size - in_block_id = ltid - block_id * block_size - ltid = kernelLocalThreadId constants - gtid = kernelGlobalThreadId constants + 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 @@ -1211,13 +1215,13 @@ | otherwise = copyDWIM (paramName y) [] (Var $ paramName x) [] -computeMapKernelGroups :: Imp.TExp Int64 -> CallKernelGen (Imp.TExp Int64, Imp.TExp Int32) +computeMapKernelGroups :: Imp.TExp Int64 -> CallKernelGen (Imp.TExp Int64, Imp.TExp Int64) computeMapKernelGroups kernel_size = do - group_size <- dPrim "group_size" int32 + 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` sExt64 (tvExp group_size) + num_groups <- dPrimV "num_groups" $ kernel_size `divUp` tvExp group_size return (tvExp num_groups, tvExp group_size) simpleKernelConstants :: @@ -1245,9 +1249,9 @@ thread_gtid thread_ltid group_id - (sExt32 num_groups) + num_groups group_size - (group_size * sExt32 num_groups) + (sExt32 (group_size * num_groups)) 0 (Imp.vi64 thread_gtid .<. kernel_size) mempty, @@ -1272,13 +1276,13 @@ sOp $ Imp.GetGroupId (tvVar phys_group_id) 0 let iterations = (required_groups - tvExp phys_group_id) - `divUp` kernelNumGroups constants + `divUp` sExt32 (kernelNumGroups constants) sFor "i" iterations $ \i -> do m . tvExp =<< dPrimV "virt_group_id" - (tvExp phys_group_id + i * kernelNumGroups constants) + (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 @@ -1288,8 +1292,8 @@ sKernelThread :: String -> - Count NumGroups (Imp.TExp Int32) -> - Count GroupSize (Imp.TExp Int32) -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> VName -> InKernelGen () -> CallKernelGen () @@ -1297,8 +1301,8 @@ sKernelGroup :: String -> - Count NumGroups (Imp.TExp Int32) -> - Count GroupSize (Imp.TExp Int32) -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> VName -> InKernelGen () -> CallKernelGen () @@ -1331,8 +1335,8 @@ Operations KernelsMem KernelEnv Imp.KernelOp -> (KernelConstants -> Imp.TExp Int32) -> String -> - Count NumGroups (Imp.TExp Int32) -> - Count GroupSize (Imp.TExp Int32) -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> VName -> InKernelGen () -> CallKernelGen () @@ -1392,7 +1396,7 @@ t <- subExpType se ds <- dropLast (arrayRank t) . arrayDims <$> lookupType arr - let dims = map toInt32Exp $ ds ++ arrayDims t + let dims = map toInt64Exp $ ds ++ arrayDims t (constants, set_constants) <- simpleKernelConstants (product $ map sExt64 dims) "replicate" @@ -1401,7 +1405,7 @@ keyWithEntryPoint fname $ nameFromString $ "replicate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - is' = unflattenIndex dims $ kernelGlobalThreadId constants + is' = unflattenIndex dims $ sExt64 $ kernelGlobalThreadId constants sKernelFailureTolerant True threadOperations constants name $ do set_constants @@ -1432,7 +1436,7 @@ sArray "arr" bt shape $ ArrayIn mem $ IxFun.iota $ - map pe32 $ shapeDims shape + map pe64 $ shapeDims shape sReplicateKernel arr $ Var val return fname @@ -1451,7 +1455,7 @@ [] fname [ Imp.MemArg arr_mem, - Imp.ExpArg $ untyped $ product $ map toInt32Exp arr_shape, + Imp.ExpArg $ untyped $ product $ map toInt64Exp arr_shape, Imp.ExpArg $ toExp' v_t' v ] _ -> return Nothing @@ -1488,7 +1492,7 @@ sKernelFailureTolerant True threadOperations constants name $ do set_constants - let gtid = kernelGlobalThreadId constants + let gtid = sExt64 $ kernelGlobalThreadId constants sWhen (kernelThreadActive constants) $ do (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] @@ -1520,7 +1524,7 @@ Imp.ScalarParam s $ IntType bt ] shape = Shape [Var n] - n' = Imp.vi32 n + n' = Imp.vi64 n x' = Imp.var x $ IntType bt s' = Imp.var s $ IntType bt @@ -1529,7 +1533,7 @@ sArray "arr" (IntType bt) shape $ ArrayIn mem $ IxFun.iota $ - map pe32 $ shapeDims shape + map pe64 $ shapeDims shape sIotaKernel arr (sExt64 n') x' s' bt return fname @@ -1537,7 +1541,7 @@ -- | Perform an Iota with a kernel. sIota :: VName -> - Imp.TExp Int32 -> + Imp.TExp Int64 -> Imp.Exp -> Imp.Exp -> IntType -> @@ -1552,7 +1556,7 @@ [] fname [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s] - else sIotaKernel arr (sExt64 n) x s et + else sIotaKernel arr n x s et sCopy :: CopyCompiler KernelsMem HostEnv Imp.HostOp sCopy @@ -1565,7 +1569,7 @@ -- Note that the shape of the destination and the source are -- necessarily the same. let shape = sliceDims srcslice - kernel_size = product $ map sExt64 shape + kernel_size = product shape (constants, set_constants) <- simpleKernelConstants kernel_size "copy" @@ -1578,7 +1582,7 @@ sKernelFailureTolerant True threadOperations constants name $ do set_constants - let gtid = kernelGlobalThreadId constants + let gtid = sExt64 $ kernelGlobalThreadId constants dest_is = unflattenIndex shape gtid src_is = dest_is @@ -1587,7 +1591,7 @@ (_, srcspace, srcidx) <- fullyIndexArray' srcloc $ fixSlice srcslice src_is - sWhen (gtid .<. sExt32 kernel_size) $ + sWhen (gtid .<. kernel_size) $ emit $ Imp.Write destmem destidx bt destspace Imp.Nonvolatile $ Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile @@ -1598,26 +1602,29 @@ KernelResult -> InKernelGen () compileGroupResult _ pe (TileReturns [(w, per_group_elems)] what) = do - n <- toInt32Exp . arraySize 0 <$> lookupType what + n <- toInt64Exp . arraySize 0 <$> lookupType what constants <- kernelConstants <$> askEnv - let ltid = kernelLocalThreadId constants - offset = toInt32Exp per_group_elems * kernelGroupId constants + 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 toInt32Exp per_group_elems == kernelGroupSize constants + if toInt64Exp per_group_elems == kernelGroupSize constants then - sWhen (offset + ltid .<. toInt32Exp w) $ + 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 .<. n) $ copyDWIMFix (patElemName pe) [j + offset] (Var what) [j] + 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 (toInt32Exp . snd) dims - group_is = zipWith (*) (map Imp.vi32 gids) out_tile_sizes + 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") $ @@ -1629,7 +1636,7 @@ compileGroupResult space pe (Returns _ what) = do constants <- kernelConstants <$> askEnv in_local_memory <- arrayInLocalMemory what - let gids = map (Imp.vi32 . fst) $ unSegSpace space + let gids = map (Imp.vi64 . fst) $ unSegSpace space if not in_local_memory then @@ -1652,22 +1659,24 @@ KernelResult -> InKernelGen () compileThreadResult space pe (Returns _ what) = do - let is = map (Imp.vi32 . fst) $ unSegSpace space + 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 = toInt32Exp per_thread_elems * kernelGlobalThreadId constants - n <- toInt32Exp . arraySize 0 <$> lookupType what + 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 <- kernelGlobalThreadId . kernelConstants <$> askEnv - n <- toInt32Exp . arraySize 0 <$> lookupType what - copyDWIM (patElemName pe) [DimSlice offset n $ toInt32Exp stride] (Var what) [] + 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 toInt32Exp rws + let rws' = map toInt64Exp rws forM_ dests $ \(slice, e) -> do - let slice' = map (fmap toInt32Exp) slice + let slice' = map (fmap toInt64Exp) slice condInBounds (DimFix i) rw = 0 .<=. i .&&. i .<. rw condInBounds (DimSlice i n s) rw = 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 @@ -62,23 +62,22 @@ data SegHistSlug = SegHistSlug { slugOp :: HistOp KernelsMem, - slugNumSubhistos :: TV Int32, + slugNumSubhistos :: TV Int64, slugSubhistos :: [SubhistosInfo], slugAtomicUpdate :: AtomicUpdate KernelsMem KernelEnv } histoSpaceUsage :: HistOp KernelsMem -> - Imp.Count Imp.Bytes (Imp.TExp Int32) + Imp.Count Imp.Bytes (Imp.TExp Int64) histoSpaceUsage op = - fmap sExt32 $ - sum $ - map - ( typeSize - . (`arrayOfRow` histWidth op) - . (`arrayOfShape` histShape op) - ) - $ lambdaReturnType $ histOp 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 @@ -87,8 +86,8 @@ SegSpace -> HistOp KernelsMem -> CallKernelGen - ( Imp.Count Imp.Bytes (Imp.TExp Int32), - Imp.Count Imp.Bytes (Imp.TExp Int32), + ( Imp.Count Imp.Bytes (Imp.TExp Int64), + Imp.Count Imp.Bytes (Imp.TExp Int64), SegHistSlug ) computeHistoUsage space op = do @@ -111,7 +110,7 @@ subhistos_membind = ArrayIn subhistos_mem $ IxFun.iota $ - map pe32 $ shapeDims subhistos_shape + map pe64 $ shapeDims subhistos_shape subhistos <- sArray (baseString dest ++ "_subhistos") @@ -128,8 +127,8 @@ multiHistoCase = do let num_elems = - foldl' (*) (tvExp num_subhistos) $ - map toInt32Exp $ arrayDims dest_t + foldl' (*) (sExt64 $ tvExp num_subhistos) $ + map toInt64Exp $ arrayDims dest_t let subhistos_mem_size = Imp.bytes $ @@ -139,15 +138,15 @@ sReplicate subhistos ne subhistos_t <- lookupType subhistos let slice = - fullSliceNum (map toInt32Exp $ arrayDims subhistos_t) $ - map (unitSlice 0 . toInt32Exp . snd) segment_dims + 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 . toInt32Exp) $ init $ segSpaceDims space) + segmented_h = h * product (map (Imp.bytes . toInt64Exp) $ init $ segSpaceDims space) atomics <- hostAtomics <$> askEnv @@ -164,7 +163,7 @@ SegHistSlug -> CallKernelGen ( Maybe Locking, - [Imp.TExp Int32] -> InKernelGen () + [Imp.TExp Int64] -> InKernelGen () ) prepareAtomicUpdateGlobal l dests slug = -- We need a separate lock array if the operators are not all of a @@ -183,7 +182,7 @@ -- algorithm to ensure good distribution of locks. let num_locks = 100151 dims = - map toInt32Exp $ + map toInt64Exp $ shapeDims (histShape (slugOp slug)) ++ [ tvSize (slugNumSubhistos slug), histWidth (slugOp slug) @@ -208,11 +207,11 @@ prepareIntermediateArraysGlobal :: Passage -> Imp.TExp Int32 -> - Imp.TExp Int32 -> + Imp.TExp Int64 -> [SegHistSlug] -> CallKernelGen ( Imp.TExp Int32, - [[Imp.TExp Int32] -> InKernelGen ()] + [[Imp.TExp Int64] -> InKernelGen ()] ) prepareIntermediateArraysGlobal passage hist_T hist_N slugs = do -- The paper formulae assume there is only one histogram, but in our @@ -223,11 +222,11 @@ -- paper. -- The sum of all Hs. - hist_H <- dPrimVE "hist_H" $ sum $ map (toInt32Exp . histWidth . slugOp) slugs + hist_H <- dPrimVE "hist_H" $ sum $ map (toInt64Exp . histWidth . slugOp) slugs hist_RF <- dPrimVE "hist_RF" $ - sum (map (r64 . toInt32Exp . histRaceFactor . slugOp) slugs) + sum (map (r64 . toInt64Exp . histRaceFactor . slugOp) slugs) / genericLength slugs hist_el_size <- dPrimVE "hist_el_size" $ sum $ map slugElAvgSize slugs @@ -238,7 +237,7 @@ hist_M_min <- dPrimVE "hist_M_min" $ - sMax32 1 $ t64 $ r64 hist_T / hist_C_max + 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. @@ -268,8 +267,9 @@ $ hist_S <-- case passage of MayBeMultiPass -> - (hist_M_min * hist_H * hist_el_size) - `divUp` t64 (hist_F_L2 * r64 (tvExp hist_L2) * hist_RACE_exp) + 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 @@ -289,7 +289,7 @@ hist_k_RF = 0.75 -- Chosen experimentally hist_F_L2 = 0.4 -- Chosen experimentally r64 = isF64 . ConvOpExp (SIToFP Int32 Float64) . untyped - t64 = isInt32 . ConvOpExp (FPToSI Float64 Int32) . untyped + t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped -- "Average element size" as computed by a formula that also takes -- locking into account. @@ -319,9 +319,9 @@ 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 = toInt32Exp $ histWidth op + hist_H = toInt64Exp $ histWidth op - hist_H_chk <- dPrimVE "hist_H_chk" $ hist_H `divUp` hist_S + 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 @@ -345,14 +345,14 @@ hist_M <- dPrimVE "hist_M" $ case slugAtomicUpdate slug of AtomicPrim {} -> 1 - _ -> sMax32 hist_M_min $ t64 $ r64 hist_T / hist_C + _ -> 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 <-- hist_M + num_subhistos <-- sExt64 hist_M -- Initialise sub-histograms. -- @@ -384,22 +384,22 @@ histKernelGlobalPass :: [PatElem KernelsMem] -> - Count NumGroups (Imp.TExp Int32) -> - Count GroupSize (Imp.TExp Int32) -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> SegSpace -> [SegHistSlug] -> KernelBody KernelsMem -> - [[Imp.TExp Int32] -> InKernelGen ()] -> + [[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 . toInt32Exp) space_sizes + 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" $ toInt32Exp w `divUp` hist_S + dPrimVE "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S sKernelThread "seghist_global" num_groups group_size (segFlat space) $ do constants <- kernelConstants <$> askEnv @@ -408,7 +408,9 @@ subhisto_inds <- forM slugs $ \slug -> dPrimVE "subhisto_ind" $ kernelGlobalThreadId constants - `quot` (kernelNumThreads constants `divUp` tvExp (slugNumSubhistos slug)) + `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, @@ -434,7 +436,7 @@ forM_ (zip map_pes map_res) $ \(pe, res) -> copyDWIMFix (patElemName pe) - (map (Imp.vi32 . fst) $ unSegSpace space) + (map (Imp.vi64 . fst) $ unSegSpace space) (kernelResultSubExp res) [] @@ -450,9 +452,9 @@ subhisto_ind, hist_H_chk ) -> do - let chk_beg = chk_i * hist_H_chk - bucket' = toInt32Exp $ kernelResultSubExp bucket - dest_w' = toInt32Exp dest_w + 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) @@ -461,8 +463,8 @@ sWhen bucket_in_bounds $ do let bucket_is = - map Imp.vi32 (init space_is) - ++ [subhisto_ind, bucket'] + map Imp.vi64 (init space_is) + ++ [sExt64 subhisto_ind, bucket'] dLParams $ lambdaParams lam sLoopNest shape $ \is -> do forM_ (zip vs_params vs') $ \(p, res) -> @@ -478,10 +480,10 @@ KernelBody KernelsMem -> CallKernelGen () histKernelGlobal map_pes num_groups group_size space slugs kbody = do - let num_groups' = fmap toInt32Exp num_groups - group_size' = fmap toInt32Exp group_size + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size let (_space_is, space_sizes) = unzip $ unSegSpace space - num_threads = unCount num_groups' * unCount group_size' + num_threads = sExt32 $ unCount num_groups' * unCount group_size' emit $ Imp.DebugPrint "## Using global memory" Nothing @@ -489,7 +491,7 @@ prepareIntermediateArraysGlobal (bodyPassage kbody) num_threads - (toInt32Exp $ last space_sizes) + (toInt64Exp $ last space_sizes) slugs sFor "chk_i" hist_S $ \chk_i -> @@ -509,25 +511,25 @@ SubExp -> InKernelGen ( [VName], - [Imp.TExp Int32] -> InKernelGen () + [Imp.TExp Int64] -> InKernelGen () ) ) ] prepareIntermediateArraysLocal :: TV Int32 -> - Count NumGroups (Imp.TExp 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 (toInt32Exp . snd) $ init $ unSegSpace space + 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 <-- unCount groups_per_segment * num_segments + num_subhistos <-- sExt64 (unCount groups_per_segment) * num_segments emit $ Imp.DebugPrint "Number of subhistograms in global memory" $ @@ -544,7 +546,7 @@ shapeDims (histShape op) ++ [hist_H_chk] - let dims = map toInt32Exp $ shapeDims lock_shape + let dims = map toInt64Exp $ shapeDims lock_shape locks <- sAllocArray "locks" int32 lock_shape $ Space "local" @@ -581,10 +583,10 @@ histKernelLocalPass :: TV Int32 -> - Count NumGroups (Imp.TExp Int32) -> + Count NumGroups (Imp.TExp Int64) -> [PatElem KernelsMem] -> - Count NumGroups (Imp.TExp Int32) -> - Count GroupSize (Imp.TExp Int32) -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> SegSpace -> [SegHistSlug] -> KernelBody KernelsMem -> @@ -609,33 +611,46 @@ 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' = toInt32Exp segment_size + segment_size' = toInt64Exp segment_size num_segments <- dPrimVE "num_segments" $ - product $ map toInt32Exp segment_dims + product $ map toInt64Exp segment_dims hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w -> - dPrimV "hist_H_chk" $ toInt32Exp w `divUp` hist_S + dPrimV "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S + histo_sizes <- forM (zip slugs hist_H_chks) $ \(slug, 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 + let group_hists_size = + sExt64 num_subhistos_per_group * histo_size + init_per_thread <- + dPrimVE "init_per_thread" $ sExt32 $ group_hists_size `divUp` unCount group_size + return (histo_dims, histo_size, init_per_thread) + sKernelThread "seghist_local" num_groups group_size (segFlat space) $ - virtualiseGroups SegVirt (unCount groups_per_segment * num_segments) $ \group_id -> do + 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` unCount groups_per_segment - gid_in_segment <- dPrimVE "gid_in_segment" $ group_id `rem` unCount groups_per_segment + 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 * kernelGroupSize constants + kernelLocalThreadId constants + gid_in_segment * sExt32 (kernelGroupSize constants) + + kernelLocalThreadId constants threads_per_segment <- dPrimVE "threads_per_segment" $ - unCount groups_per_segment * kernelGroupSize constants + sExt32 $ unCount groups_per_segment * kernelGroupSize constants -- Set segment indices. zipWithM_ dPrimV_ segment_is $ - unflattenIndex (map toInt32Exp segment_dims) flat_segment_id + 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 @@ -649,38 +664,32 @@ 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 toInt32Exp (shapeDims (histShape (slugOp slug))) - histo_size <- dPrimVE "histo_size" $ product histo_dims - f slug dests (tvExp hist_H_chk) histo_dims histo_size + let onSlugs f = + forM_ (zip3 slugs histograms histo_sizes) $ + \(slug, (dests, hist_H_chk, _), (histo_dims, histo_size, init_per_thread)) -> + f slug dests (tvExp hist_H_chk) histo_dims histo_size init_per_thread 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 + onSlugs $ \slug dests hist_H_chk histo_dims histo_size init_per_thread -> do + let group_hists_size = num_subhistos_per_group * sExt32 histo_size 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 + i * sExt32 (kernelGroupSize constants) + kernelLocalThreadId constants j_offset <- dPrimVE "j_offset" $ - num_subhistos_per_group * histo_size * gid_in_segment + j + num_subhistos_per_group * sExt32 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 + 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 + chk_i * hist_H_chk : + head local_bucket_is + sExt64 chk_i * hist_H_chk : tail local_bucket_is - global_subhisto_i <- dPrimVE "global_subhisto_i" $ j_offset `quot` histo_size + global_subhisto_i <- dPrimVE "global_subhisto_i" $ j_offset `quot` sExt32 histo_size sWhen (j .<. group_hists_size) $ f @@ -696,8 +705,8 @@ 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 + 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) @@ -707,7 +716,7 @@ sOp $ Imp.Barrier Imp.FenceLocal - kernelLoop pgtid_in_segment threads_per_segment segment_size' $ \ie -> do + 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 @@ -726,7 +735,7 @@ forM_ (zip map_pes map_res) $ \(pe, se) -> copyDWIMFix (patElemName pe) - (map Imp.vi32 space_is) + (map Imp.vi64 space_is) se [] @@ -736,14 +745,14 @@ bucket, vs' ) -> do - let chk_beg = chk_i * tvExp hist_H_chk - bucket' = toInt32Exp bucket - dest_w' = toInt32Exp dest_w + 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 = [thread_local_subhisto_i, bucket' - chk_beg] + bucket_is = [sExt64 thread_local_subhisto_i, bucket' - chk_beg] vs_params = takeLast (length vs') $ lambdaParams lam sComment "perform atomic updates" $ @@ -757,30 +766,28 @@ 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 - + onSlugs $ \slug dests hist_H_chk histo_dims _histo_size bins_per_thread -> do trunc_H <- dPrimV "trunc_H" $ - sMin32 hist_H_chk $ - toInt32Exp (histWidth (slugOp slug)) - chk_i * head histo_dims + sMin64 hist_H_chk $ + toInt64Exp (histWidth (slugOp slug)) + - sExt64 chk_i * head histo_dims let trunc_histo_dims = tvExp trunc_H : - map toInt32Exp (shapeDims (histShape (slugOp slug))) - trunc_histo_size <- dPrimVE "histo_size" $ product trunc_histo_dims + map toInt64Exp (shapeDims (histShape (slugOp slug))) + trunc_histo_size <- dPrimVE "histo_size" $ sExt32 $ product trunc_histo_dims sFor "local_i" bins_per_thread $ \i -> do j <- dPrimVE "j" $ - i * kernelGroupSize constants + kernelLocalThreadId constants + i * sExt32 (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 + let local_bucket_is = unflattenIndex histo_dims $ sExt64 j global_bucket_is = - head local_bucket_is + chk_i * hist_H_chk : + 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 @@ -803,20 +810,20 @@ (paramName yp) [] (Var subhisto) - (subhisto_id + 1 : local_bucket_is) + (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.vi32 segment_is - ++ [group_id `rem` unCount groups_per_segment] + 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 Int32) -> + Count NumGroups (Imp.TExp Int64) -> [PatElem KernelsMem] -> Count NumGroups SubExp -> Count GroupSize SubExp -> @@ -826,8 +833,8 @@ 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 toInt32Exp num_groups - group_size' = fmap toInt32Exp group_size + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size num_subhistos_per_group = tvExp num_subhistos_per_group_var emit $ @@ -864,9 +871,9 @@ [PatElem KernelsMem] -> Imp.TExp Int32 -> SegSpace -> - Imp.TExp Int32 -> - Imp.TExp Int32 -> - Imp.TExp Int32 -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> Imp.TExp Int32 -> [SegHistSlug] -> KernelBody KernelsMem -> @@ -885,20 +892,20 @@ num_groups <- fmap (Imp.Count . tvSize) $ dPrimV "num_groups" $ - hist_T `divUp` toInt32Exp (unCount group_size) - let num_groups' = toInt32Exp <$> num_groups - group_size' = toInt32Exp <$> group_size + sExt64 hist_T `divUp` toInt64Exp (unCount group_size) + let num_groups' = toInt64Exp <$> num_groups + group_size' = toInt64Exp <$> group_size - let r64 = isF64 . ConvOpExp (SIToFP Int32 Float64) . untyped - t64 = isInt32 . ConvOpExp (FPToSI Float64 Int32) . untyped + let r64 = isF64 . ConvOpExp (SIToFP Int64 Float64) . untyped + t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped -- M approximation. hist_m' <- dPrimVE "hist_m_prime" $ r64 - ( sMin32 - (tvExp hist_L `quot` hist_el_size) - (hist_N `divUp` unCount num_groups') + ( sMin64 + (sExt64 (tvExp hist_L `quot` hist_el_size)) + (hist_N `divUp` sExt64 (unCount num_groups')) ) / r64 hist_H @@ -907,15 +914,15 @@ -- M in the paper, but not adjusted for asymptotic efficiency. hist_M0 <- dPrimVE "hist_M0" $ - sMax32 1 $ sMin32 (t64 hist_m') hist_B + 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 toInt32Exp segment_dims + hist_Nout <- dPrimVE "hist_Nout" $ product $ map toInt64Exp segment_dims - hist_Nin <- dPrimVE "hist_Nin" $ toInt32Exp $ last space_sizes + hist_Nin <- dPrimVE "hist_Nin" $ toInt64Exp $ last space_sizes -- Maximum M for work efficiency. work_asymp_M_max <- @@ -928,9 +935,9 @@ `divUp` sExt64 hist_Nout -- Number of groups, rounded up. - let r = hist_T_hist_min `divUp` hist_B + let r = hist_T_hist_min `divUp` sExt32 hist_B - dPrimVE "work_asymp_M_max" $ hist_Nin `quot` (r * hist_H) + dPrimVE "work_asymp_M_max" $ hist_Nin `quot` (sExt64 r * hist_H) else dPrimVE "work_asymp_M_max" $ (hist_Nout * hist_N) @@ -939,7 +946,7 @@ ) -- Number of subhistograms per result histogram. - hist_M <- dPrimV "hist_M" $ sMin32 hist_M0 work_asymp_M_max + 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. @@ -949,7 +956,7 @@ -- working on the same (sub)histogram. hist_C <- dPrimVE "hist_C" $ - hist_B `divUp` hist_M_nonzero + 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 @@ -958,18 +965,30 @@ emit $ Imp.DebugPrint "local M" $ Just $ untyped $ tvExp hist_M emit $ Imp.DebugPrint "local memory needed" $ - Just $ untyped $ hist_H * hist_el_size * tvExp hist_M + 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 * tvExp hist_M - hist_S <- dPrimVE "hist_S" $ (hist_H * local_mem_needed) `divUp` tvExp hist_L + 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 + groups_per_segment <- + if segmented + then + fmap Count $ + dPrimVE "groups_per_segment" $ unCount num_groups' `divUp` hist_Nout + else pure num_groups' + -- 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 @@ -981,10 +1000,6 @@ .&&. 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 @@ -1020,9 +1035,9 @@ -- 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 toInt32Exp num_groups - group_size' = fmap toInt32Exp group_size - dims = map toInt32Exp $ segSpaceDims space + 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 @@ -1038,7 +1053,7 @@ let hist_B = unCount group_size' -- Size of a histogram. - hist_H <- dPrimVE "hist_H" $ sum $ map (toInt32Exp . histWidth) ops + 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 @@ -1060,7 +1075,7 @@ sum (map (toInt32Exp . histRaceFactor . slugOp) slugs) `quot` genericLength slugs - let hist_T = unCount num_groups' * unCount group_size' + 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 @@ -1068,7 +1083,7 @@ emit $ Imp.DebugPrint "Input elements per histogram (N)" $ Just $ untyped hist_N emit $ Imp.DebugPrint "Number of segments" $ - Just $ untyped $ product $ map (toInt32Exp . snd) segment_dims + 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 @@ -1126,7 +1141,7 @@ red_cont $ flip map subhistos $ \subhisto -> ( Var subhisto, - map Imp.vi32 $ + map Imp.vi64 $ map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids ) where 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 @@ -24,14 +24,15 @@ CallKernelGen () compileSegMap pat lvl space kbody = do let (is, dims) = unzip $ unSegSpace space - dims' = map toInt32Exp dims - num_groups' = toInt32Exp <$> segNumGroups lvl - group_size' = toInt32Exp <$> segGroupSize lvl + dims' = map toInt64Exp dims + num_groups' = toInt64Exp <$> segNumGroups lvl + group_size' = toInt64Exp <$> segGroupSize lvl case lvl of 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 @@ -40,7 +41,7 @@ + sExt64 local_tid zipWithM_ dPrimV_ is $ - map sExt32 $ unflattenIndex (map sExt64 dims') global_tid + map sExt64 $ unflattenIndex (map sExt64 dims') global_tid sWhen (isActive $ unSegSpace space) $ compileStms mempty (kernelBodyStms kbody) $ @@ -48,10 +49,10 @@ kernelBodyResult kbody SegGroup {} -> sKernelGroup "segmap_intragroup" num_groups' group_size' (segFlat space) $ do - let virt_num_groups = product dims' + let virt_num_groups = sExt32 $ product dims' precomputeSegOpIDs (kernelBodyStms kbody) $ virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do - zipWithM_ dPrimV_ is $ unflattenIndex dims' group_id + zipWithM_ dPrimV_ is $ unflattenIndex dims' $ sExt64 group_id compileStms mempty (kernelBodyStms kbody) $ zipWithM_ (compileGroupResult space) (patternElements pat) $ 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 @@ -72,7 +72,7 @@ -- 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.TExp Int32])] -> InKernelGen ()) -> InKernelGen () +type DoSegBody = ([(SubExp, [Imp.TExp Int64])] -> InKernelGen ()) -> InKernelGen () -- | Compile 'SegRed' instance to host-level code with calls to -- various kernels. @@ -106,11 +106,11 @@ | genericLength reds > maxNumOps = compilerLimitationS $ "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported." - | [(_, Constant (IntValue (Int32Value 1))), _] <- unSegSpace space = + | [(_, Constant (IntValue (Int64Value 1))), _] <- unSegSpace space = nonsegmentedReduction pat num_groups group_size space reds body | otherwise = do - let group_size' = toInt32Exp $ unCount group_size - segment_size = toInt32Exp $ last $ segSpaceDims space + let group_size' = toInt64Exp $ unCount group_size + segment_size = toInt64Exp $ last $ segSpaceDims space use_small_segments = segment_size * 2 .<. group_size' sIf use_small_segments @@ -139,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 pe32 $ shapeDims shape' + ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape' _ -> do let pt = elemType $ paramType p shape = Shape [group_size] @@ -176,9 +176,9 @@ CallKernelGen () nonsegmentedReduction segred_pat num_groups group_size space reds body = do let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt32Exp dims - num_groups' = fmap toInt32Exp num_groups - group_size' = fmap toInt32Exp group_size + 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' @@ -201,10 +201,12 @@ -- Since this is the nonsegmented case, all outer segment IDs must -- necessarily be 0. - forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int32) + forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int64) 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 @@ -253,7 +255,7 @@ 0 [0] 0 - (kernelNumGroups constants) + (sExt64 $ kernelNumGroups constants) slug red_x_params red_y_params @@ -276,19 +278,19 @@ CallKernelGen () smallSegmentsReduction (Pattern _ segred_pes) num_groups group_size space reds body = do let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt32Exp dims + dims' = map toInt64Exp dims segment_size = last dims' -- Careful to avoid division by zero now. segment_size_nonzero <- - dPrimVE "segment_size_nonzero" $ sMax32 1 segment_size + dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size - let num_groups' = fmap toInt32Exp num_groups - group_size' = fmap toInt32Exp 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 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 $ untyped num_segments @@ -307,8 +309,10 @@ -- 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 @@ -336,13 +340,14 @@ 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) - (tvExp num_threads) + (sExt64 $ tvExp num_threads) (segment_size * segments_per_group) red_op red_arrs @@ -351,13 +356,15 @@ sComment "save final values of segments" $ sWhen - ( group_id' * segments_per_group + ltid .<. num_segments + ( 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 = group_id' * segments_per_group + ltid - gtids' = unflattenIndex (init dims') flat_segment_index + let flat_segment_index = + sExt64 group_id' * segments_per_group + sExt64 ltid + gtids' = + unflattenIndex (init dims') flat_segment_index copyDWIMFix (patElemName pe) gtids' @@ -378,11 +385,11 @@ CallKernelGen () largeSegmentsReduction segred_pat num_groups group_size space reds body = do let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt32Exp dims + dims' = map toInt64Exp dims num_segments = product $ init dims' segment_size = last dims' - num_groups' = fmap toInt32Exp num_groups - group_size' = fmap toInt32Exp group_size + num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size (groups_per_segment, elems_per_thread) <- groupsPerSegmentAndElementsPerThread @@ -436,26 +443,26 @@ -- 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 (tvExp virt_num_groups) $ \group_id -> do + virtualiseGroups SegVirt (sExt32 (tvExp virt_num_groups)) $ \group_id -> do let segment_gtids = init gtids w = last dims local_tid = kernelLocalThreadId constants flat_segment_id <- dPrimVE "flat_segment_id" $ - group_id `quot` groups_per_segment + 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) + (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 - let num_elements = Imp.elements $ toInt32Exp w + 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) $ @@ -465,7 +472,7 @@ constants (zip gtids dims') num_elements - global_tid + (sExt32 global_tid) elems_per_thread (tvVar threads_per_segment) slugs @@ -501,8 +508,8 @@ pes group_id flat_segment_id - (map Imp.vi32 segment_gtids) - first_group_for_segment + (map Imp.vi64 segment_gtids) + (sExt64 first_group_for_segment) groups_per_segment slug red_x_params @@ -521,25 +528,25 @@ 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.vi32 segment_gtids) (Var 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.TExp Int32 -> - Imp.TExp Int32 -> - Count NumGroups (Imp.TExp Int32) -> - Count GroupSize (Imp.TExp Int32) -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> CallKernelGen - ( Imp.TExp Int32, - Imp.Count Imp.Elements (Imp.TExp Int32) + ( 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` sMax32 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) @@ -552,7 +559,7 @@ -- (either local or global memory). slugArrs :: [VName], -- | Places to store accumulator in stage 1 reduction. - slugAccs :: [(VName, [Imp.TExp Int32])] + slugAccs :: [(VName, [Imp.TExp Int64])] } slugBody :: SegBinOpSlug -> Body KernelsMem @@ -585,29 +592,29 @@ acc <- dPrim (baseString (paramName p) <> "_acc") t return (tvVar acc, []) | otherwise = - return (param_arr, [local_tid, group_id]) + return (param_arr, [sExt64 local_tid, sExt64 group_id]) reductionStageZero :: KernelConstants -> - [(VName, Imp.TExp Int32)] -> - Imp.Count Imp.Elements (Imp.TExp Int32) -> + [(VName, Imp.TExp Int64)] -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> Imp.TExp Int32 -> - Imp.Count Imp.Elements (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 = mkTV (last gtids) int32 - 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 + computeThreadChunkSize ordering (sExt64 global_tid) elems_per_thread num_elements chunk_size dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs @@ -631,7 +638,7 @@ copyDWIMFix arr [local_tid] (Var $ paramName p) [] sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier. - groupReduce (kernelGroupSize constants) slug_op_renamed (slugArrs slug) + groupReduce (sExt32 (kernelGroupSize constants)) slug_op_renamed (slugArrs slug) sOp $ Imp.Barrier Imp.FenceLocal @@ -656,13 +663,13 @@ gtid <-- case comm of Commutative -> - global_tid - + Imp.vi32 threads_per_segment * 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 + 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) check_bounds $ sComment "apply map function" $ @@ -704,10 +711,10 @@ reductionStageOne :: KernelConstants -> - [(VName, Imp.TExp Int32)] -> - Imp.Count Imp.Elements (Imp.TExp Int32) -> + [(VName, Imp.TExp Int64)] -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> Imp.TExp Int32 -> - Imp.Count Imp.Elements (Imp.TExp Int32) -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> VName -> [SegBinOpSlug] -> DoSegBody -> @@ -730,9 +737,9 @@ [PatElem KernelsMem] -> Imp.TExp Int32 -> Imp.TExp Int32 -> - [Imp.TExp Int32] -> - Imp.TExp Int32 -> - Imp.TExp Int32 -> + [Imp.TExp Int64] -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> SegBinOpSlug -> [LParam KernelsMem] -> [LParam KernelsMem] -> @@ -770,13 +777,14 @@ (counter_mem, _, counter_offset) <- fullyIndexArray counter - [ counter_i * num_counters - + flat_segment_id `rem` num_counters + [ 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, group_id] (Var 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. @@ -786,7 +794,7 @@ Int32 (tvVar old_counter) counter_mem - (sExt32 <$> counter_offset) + 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. @@ -806,7 +814,7 @@ sWhen (local_tid .==. 0) $ sOp $ Imp.Atomic DefaultSpace $ - Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem (sExt32 <$> counter_offset) $ + Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $ untyped $ negate groups_per_segment sLoopNest (slugShape slug) $ \vec_is -> do @@ -818,7 +826,7 @@ comment "read in the per-group-results" $ do read_per_thread <- dPrimVE "read_per_thread" $ - groups_per_segment `divUp` group_size + groups_per_segment `divUp` sExt64 group_size forM_ (zip red_x_params nes) $ \(p, ne) -> copyDWIMFix (paramName p) [] ne [] @@ -826,7 +834,7 @@ sFor "i" read_per_thread $ \i -> do group_res_id <- dPrimVE "group_res_id" $ - local_tid * read_per_thread + i + sExt64 local_tid * read_per_thread + i index_of_group_res <- dPrimVE "index_of_group_res" $ first_group_for_segment + group_res_id @@ -846,12 +854,12 @@ forM_ (zip red_x_params red_arrs) $ \(p, arr) -> when (primType $ paramType p) $ - copyDWIMFix arr [local_tid] (Var $ paramName p) [] + copyDWIMFix arr [sExt64 local_tid] (Var $ paramName p) [] sOp $ Imp.Barrier Imp.FenceLocal sComment "reduce the per-group results" $ do - groupReduce group_size red_op_renamed red_arrs + groupReduce (sExt32 group_size) red_op_renamed red_arrs sComment "and back to memory with the final result" $ sWhen (local_tid .==. 0) $ 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 @@ -44,7 +44,7 @@ arr <- lift $ sArray "scan_arr" pt shape' $ - ArrayIn mem $ IxFun.iota $ map pe32 $ shapeDims shape' + ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape' return (arr, []) _ -> do let pt = elemType $ paramType p @@ -69,13 +69,13 @@ mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "local" return ([size], mem) -type CrossesSegment = Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) +type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool) -localArrayIndex :: KernelConstants -> Type -> Imp.TExp Int32 +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) @@ -100,7 +100,7 @@ forM_ (zip pes scan_res) $ \(pe, res) -> copyDWIMFix (patElemName pe) - (map Imp.vi32 gtids) + (map Imp.vi64 gtids) (kernelResultSubExp res) [] | otherwise = @@ -108,7 +108,7 @@ copyDWIMFix (paramName p) [] (kernelResultSubExp res) [] readToScanValues :: - [Imp.TExp Int32] -> + [Imp.TExp Int64] -> [PatElem KernelsMem] -> SegBinOp KernelsMem -> InKernelGen () @@ -120,9 +120,9 @@ return () readCarries :: - Imp.TExp Int32 -> - [Imp.TExp Int32] -> - [Imp.TExp Int32] -> + Imp.TExp Int64 -> + [Imp.TExp Int64] -> + [Imp.TExp Int64] -> [PatElem KernelsMem] -> SegBinOp KernelsMem -> InKernelGen () @@ -152,16 +152,16 @@ SegSpace -> [SegBinOp KernelsMem] -> KernelBody KernelsMem -> - CallKernelGen (TV Int32, Imp.TExp Int32, CrossesSegment) + CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment) scanStage1 (Pattern _ all_pes) num_groups group_size space scans kbody = do - let num_groups' = fmap toInt32Exp num_groups - group_size' = fmap toInt32Exp 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' = map toInt32Exp dims + dims' = map toInt64Exp dims let num_elements = product dims' - elems_per_thread = num_elements `divUp` tvExp num_threads + elems_per_thread = num_elements `divUp` sExt64 (tvExp num_threads) elems_per_group = unCount group_size' * elems_per_thread let crossesSegment = @@ -184,18 +184,18 @@ sFor "j" elems_per_thread $ \j -> do chunk_offset <- dPrimV "chunk_offset" $ - kernelGroupSize constants * j - + kernelGroupId constants * elems_per_group + sExt64 (kernelGroupSize constants) * j + + sExt64 (kernelGroupId constants) * elems_per_group flat_idx <- dPrimV "flat_idx" $ - tvExp chunk_offset + kernelLocalThreadId constants + tvExp chunk_offset + sExt64 (kernelLocalThreadId constants) -- Construct segment indices. 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) = @@ -211,7 +211,7 @@ forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(pe, se) -> copyDWIMFix (patElemName pe) - (map Imp.vi32 gtids) + (map Imp.vi64 gtids) (kernelResultSubExp se) [] @@ -232,7 +232,7 @@ sIf in_bounds ( do - readToScanValues (map Imp.vi32 gtids ++ vec_is) pes scan + 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) -> @@ -242,13 +242,14 @@ 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 [] + \(t, arr, se) -> + copyDWIMFix arr [localArrayIndex constants t] se [] let crossesSegment' = do f <- crossesSegment Just $ \from to -> - let from' = from + tvExp chunk_offset - to' = to + tvExp chunk_offset + let from' = sExt64 from + tvExp chunk_offset + to' = sExt64 to + tvExp chunk_offset in f from' to' sOp $ Imp.ErrorSync fence @@ -257,8 +258,8 @@ scan_op_renamed <- renameLambda scan_op groupScan crossesSegment' - (tvExp num_threads) - (kernelGroupSize constants) + (sExt64 $ tvExp num_threads) + (sExt64 $ kernelGroupSize constants) scan_op_renamed local_arrs @@ -267,7 +268,7 @@ forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> copyDWIMFix (patElemName pe) - (map Imp.vi32 gtids ++ vec_is) + (map Imp.vi64 gtids ++ vec_is) (Var arr) [localArrayIndex constants t] @@ -280,8 +281,10 @@ [] (Var arr) [ if primType $ paramType p - then kernelGroupSize constants - 1 - else (kernelGroupId constants + 1) * kernelGroupSize constants - 1 + then sExt64 (kernelGroupSize constants) - 1 + else + (sExt64 (kernelGroupId constants) + 1) + * sExt64 (kernelGroupSize constants) - 1 ] load_neutral = forM_ (zip nes scan_x_params) $ \(ne, p) -> @@ -294,10 +297,10 @@ Just f -> f ( tvExp chunk_offset - + kernelGroupSize constants -1 + + sExt64 (kernelGroupSize constants) -1 ) ( tvExp chunk_offset - + kernelGroupSize constants + + sExt64 (kernelGroupSize constants) ) should_load_carry <- dPrimVE "should_load_carry" $ @@ -313,7 +316,7 @@ scanStage2 :: Pattern KernelsMem -> TV Int32 -> - Imp.TExp Int32 -> + Imp.TExp Int64 -> Count NumGroups SubExp -> CrossesSegment -> SegSpace -> @@ -321,16 +324,18 @@ CallKernelGen () scanStage2 (Pattern _ all_pes) stage1_num_threads elems_per_group num_groups crossesSegment space scans = do let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt32Exp 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' = fmap toInt32Exp 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 constants <- kernelConstants <$> askEnv @@ -340,17 +345,17 @@ flat_idx <- dPrimV "flat_idx" $ - (kernelLocalThreadId constants + 1) * elems_per_group - 1 + (sExt64 (kernelLocalThreadId constants) + 1) * elems_per_group - 1 -- Construct segment indices. 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' + foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) -> copyDWIMFix @@ -371,8 +376,8 @@ groupScan crossesSegment' - (tvExp stage1_num_threads) - (kernelGroupSize constants) + (sExt64 $ tvExp stage1_num_threads) + (sExt64 $ kernelGroupSize constants) scan_op local_arrs @@ -389,19 +394,19 @@ Pattern KernelsMem -> Count NumGroups SubExp -> Count GroupSize SubExp -> - Imp.TExp Int32 -> + Imp.TExp Int64 -> CrossesSegment -> SegSpace -> [SegBinOp KernelsMem] -> CallKernelGen () scanStage3 (Pattern _ all_pes) num_groups group_size elems_per_group crossesSegment space scans = do - let num_groups' = fmap toInt32Exp num_groups - group_size' = fmap toInt32Exp group_size + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt32Exp dims + dims' = map toInt64Exp dims required_groups <- dPrimVE "required_groups" $ - product dims' `divUp` unCount group_size' + 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 @@ -410,8 +415,8 @@ -- Compute our logical index. flat_idx <- dPrimVE "flat_idx" $ - virt_group_id * unCount group_size' - + kernelLocalThreadId constants + 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. @@ -428,7 +433,7 @@ -- 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' + foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' crosses_segment = fromMaybe false $ crossesSegment @@ -459,14 +464,14 @@ (paramName p) [] (Var $ patElemName pe) - (map Imp.vi32 gtids ++ vec_is) + (map Imp.vi64 gtids ++ vec_is) compileBody' scan_x_params $ lambdaBody scan_op forM_ (zip scan_x_params pes) $ \(p, pe) -> copyDWIMFix (patElemName pe) - (map Imp.vi32 gtids ++ vec_is) + (map Imp.vi64 gtids ++ vec_is) (Var $ paramName p) [] @@ -485,14 +490,14 @@ -- 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 + stage1_max_num_groups <- dPrim "stage1_max_num_groups" int64 sOp $ Imp.GetSizeMax (tvVar stage1_max_num_groups) SizeGroup stage1_num_groups <- fmap (Imp.Count . tvSize) $ dPrimV "stage1_num_groups" $ - sMin32 (tvExp stage1_max_num_groups) $ - toInt32Exp $ Imp.unCount $ segNumGroups lvl + sMin64 (tvExp stage1_max_num_groups) $ + toInt64Exp $ Imp.unCount $ segNumGroups lvl (stage1_num_threads, elems_per_group, crossesSegment) <- scanStage1 pat stage1_num_groups (segGroupSize lvl) space scans kbody @@ -502,4 +507,4 @@ 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 toInt32Exp $ segSpaceDims space + n = product $ map toInt64Exp $ 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 @@ -180,7 +180,7 @@ let params = [ [C.cparam|__global int *global_failure|], - [C.cparam|__global int *global_failure_args|] + [C.cparam|__global typename int64_t *global_failure_args|] ] (func, cstate) = genGPUCode FunMode (functionBody device_func) failures $ @@ -312,7 +312,7 @@ 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 typename int64_t *global_failure_args|] ] params = @@ -780,6 +780,10 @@ 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 diff --git a/src/Futhark/CodeGen/ImpGen/Multicore.hs b/src/Futhark/CodeGen/ImpGen/Multicore.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Multicore.hs @@ -0,0 +1,101 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +module Futhark.CodeGen.ImpGen.Multicore + ( Futhark.CodeGen.ImpGen.Multicore.compileProg, + Warnings, + ) +where + +import qualified Futhark.CodeGen.ImpCode.Multicore as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.Multicore.Base +import Futhark.CodeGen.ImpGen.Multicore.SegHist +import Futhark.CodeGen.ImpGen.Multicore.SegMap +import Futhark.CodeGen.ImpGen.Multicore.SegRed +import Futhark.CodeGen.ImpGen.Multicore.SegScan +import Futhark.IR.MCMem +import Futhark.MonadFreshNames +import Prelude hiding (quot, rem) + +-- GCC supported primitve atomic Operations +-- TODO: Add support for 1, 2, and 16 bytes too +gccAtomics :: AtomicBinOp +gccAtomics = flip lookup cpu + where + cpu = + [ (Add Int32 OverflowUndef, Imp.AtomicAdd Int32), + (Sub Int32 OverflowUndef, Imp.AtomicSub Int32), + (And Int32, Imp.AtomicAnd Int32), + (Xor Int32, Imp.AtomicXor Int32), + (Or Int32, Imp.AtomicOr Int32), + (Add Int64 OverflowUndef, Imp.AtomicAdd Int64), + (Sub Int64 OverflowUndef, Imp.AtomicSub Int64), + (And Int64, Imp.AtomicAnd Int64), + (Xor Int64, Imp.AtomicXor Int64), + (Or Int64, Imp.AtomicOr Int64) + ] + +compileProg :: + MonadFreshNames m => + Prog MCMem -> + m (Warnings, Imp.Definitions Imp.Multicore) +compileProg = Futhark.CodeGen.ImpGen.compileProg (HostEnv gccAtomics) ops Imp.DefaultSpace + where + ops = defaultOperations opCompiler + opCompiler dest (Alloc e space) = compileAlloc dest e space + opCompiler dest (Inner op) = compileMCOp dest op + +compileMCOp :: + Pattern MCMem -> + MCOp MCMem () -> + ImpM MCMem HostEnv Imp.Multicore () +compileMCOp _ (OtherOp ()) = pure () +compileMCOp pat (ParOp par_op op) = do + let space = getSpace op + dPrimV_ (segFlat space) (0 :: Imp.TExp Int32) + iterations <- getIterationDomain op space + nsubtasks <- dPrim "num_tasks" $ IntType Int32 + seq_code <- compileSegOp pat op nsubtasks + retvals <- getReturnParams pat op + + let scheduling_info = Imp.SchedulerInfo (tvVar nsubtasks) (untyped iterations) + + par_code <- case par_op of + Just nested_op -> do + let space' = getSpace nested_op + dPrimV_ (segFlat space') (0 :: Imp.TExp Int32) + compileSegOp pat nested_op nsubtasks + Nothing -> return mempty + + let par_task = case par_op of + Just nested_op -> Just $ Imp.ParallelTask par_code $ segFlat $ getSpace nested_op + Nothing -> Nothing + + let non_free = + ( [segFlat space, tvVar nsubtasks] + ++ map Imp.paramName retvals + ) + ++ case par_op of + Just nested_op -> + [segFlat $ getSpace nested_op] + Nothing -> [] + + s <- segOpString op + free_params <- freeParams (par_code <> seq_code) non_free + let seq_task = Imp.ParallelTask seq_code (segFlat space) + emit $ Imp.Op $ Imp.Segop s free_params seq_task par_task retvals $ scheduling_info (decideScheduling' op seq_code) + +compileSegOp :: + Pattern MCMem -> + SegOp () MCMem -> + TV Int32 -> + ImpM MCMem HostEnv Imp.Multicore Imp.Code +compileSegOp pat (SegHist _ space histops _ kbody) ntasks = + compileSegHist pat space histops kbody ntasks +compileSegOp pat (SegScan _ space scans _ kbody) ntasks = + compileSegScan pat space scans kbody ntasks +compileSegOp pat (SegRed _ space reds _ kbody) ntasks = + compileSegRed pat space reds kbody ntasks +compileSegOp pat (SegMap _ space _ kbody) _ = + compileSegMap pat space kbody diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs @@ -0,0 +1,458 @@ +module Futhark.CodeGen.ImpGen.Multicore.Base + ( toParam, + compileKBody, + extractAllocations, + compileThreadResult, + HostEnv (..), + AtomicBinOp, + MulticoreGen, + decideScheduling, + decideScheduling', + groupResultArrays, + renameSegBinOp, + resultArrays, + freeParams, + renameHistOpLambda, + atomicUpdateLocking, + AtomicUpdate (..), + Locking (..), + getSpace, + getIterationDomain, + getReturnParams, + segOpString, + ) +where + +import Control.Monad +import Data.Bifunctor +import Data.List (elemIndex, find) +import Data.Maybe +import qualified Futhark.CodeGen.ImpCode.Multicore as Imp +import Futhark.CodeGen.ImpGen +import Futhark.Error +import Futhark.IR.MCMem +import Futhark.Transform.Rename +import Futhark.Util (maybeNth) +import Prelude hiding (quot, rem) + +-- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? +type AtomicBinOp = + BinOp -> + Maybe (VName -> VName -> Imp.Count Imp.Elements (Imp.TExp Int32) -> Imp.Exp -> Imp.AtomicOp) + +newtype HostEnv = HostEnv + {hostAtomics :: AtomicBinOp} + +type MulticoreGen = ImpM MCMem HostEnv Imp.Multicore + +segOpString :: SegOp () MCMem -> MulticoreGen String +segOpString SegMap {} = return "segmap" +segOpString SegRed {} = return "segred" +segOpString SegScan {} = return "segscan" +segOpString SegHist {} = return "seghist" + +toParam :: VName -> TypeBase shape u -> MulticoreGen Imp.Param +toParam name (Prim pt) = return $ Imp.ScalarParam name pt +toParam name (Mem space) = return $ Imp.MemParam name space +toParam name Array {} = do + name_entry <- lookupVar name + case name_entry of + ArrayVar _ (ArrayEntry (MemLocation mem _ _) _) -> + return $ Imp.MemParam mem DefaultSpace + _ -> error $ "[toParam] Could not handle array for " ++ show name + +getSpace :: SegOp () MCMem -> SegSpace +getSpace (SegHist _ space _ _ _) = space +getSpace (SegRed _ space _ _ _) = space +getSpace (SegScan _ space _ _ _) = space +getSpace (SegMap _ space _ _) = space + +getIterationDomain :: SegOp () MCMem -> SegSpace -> MulticoreGen (Imp.TExp Int64) +getIterationDomain SegMap {} space = do + let ns = map snd $ unSegSpace space + ns_64 = map (sExt64 . toInt32Exp) ns + return $ product ns_64 +getIterationDomain _ space = do + let ns = map snd $ unSegSpace space + ns_64 = map (sExt64 . toInt32Exp) ns + case unSegSpace space of + [_] -> return $ product ns_64 + -- A segmented SegOp is over the segments + -- so we drop the last dimension, which is + -- executed sequentially + _ -> return $ product $ init ns_64 + +-- When the SegRed's return value is a scalar +-- we perform a call by value-result in the segop function +getReturnParams :: Pattern MCMem -> SegOp () MCMem -> MulticoreGen [Imp.Param] +getReturnParams pat SegRed {} = do + let retvals = map patElemName $ patternElements pat + retvals_ts <- mapM lookupType retvals + zipWithM toParam retvals retvals_ts +getReturnParams _ _ = return mempty + +renameSegBinOp :: [SegBinOp MCMem] -> MulticoreGen [SegBinOp MCMem] +renameSegBinOp segbinops = + forM segbinops $ \(SegBinOp comm lam ne shape) -> do + lam' <- renameLambda lam + return $ SegBinOp comm lam' ne shape + +compileKBody :: + KernelBody MCMem -> + ([(SubExp, [Imp.Exp])] -> ImpM MCMem () Imp.Multicore ()) -> + ImpM MCMem () Imp.Multicore () +compileKBody kbody red_cont = + compileStms (freeIn $ kernelBodyResult kbody) (kernelBodyStms kbody) $ do + let red_res = kernelBodyResult kbody + red_cont $ zip (map kernelResultSubExp red_res) $ repeat [] + +compileThreadResult :: + SegSpace -> + PatElem MCMem -> + KernelResult -> + MulticoreGen () +compileThreadResult space pe (Returns _ what) = do + let is = map (Imp.vi64 . fst) $ unSegSpace space + copyDWIMFix (patElemName pe) is what [] +compileThreadResult _ _ ConcatReturns {} = + compilerBugS "compileThreadResult: ConcatReturn unhandled." +compileThreadResult _ _ WriteReturns {} = + compilerBugS "compileThreadResult: WriteReturns unhandled." +compileThreadResult _ _ TileReturns {} = + compilerBugS "compileThreadResult: TileReturns unhandled." + +freeVariables :: Imp.Code -> [VName] -> [VName] +freeVariables code names = + namesToList $ freeIn code `namesSubtract` namesFromList names + +freeParams :: Imp.Code -> [VName] -> MulticoreGen [Imp.Param] +freeParams code names = do + let freeVars = freeVariables code names + ts <- mapM lookupType freeVars + zipWithM toParam freeVars ts + +-- | Arrays for storing group results. +resultArrays :: String -> [SegBinOp MCMem] -> MulticoreGen [[VName]] +resultArrays s segops = + forM segops $ \(SegBinOp _ lam _ shape) -> + forM (lambdaReturnType lam) $ \t -> do + let pt = elemType t + full_shape = shape <> arrayShape t + sAllocArray s pt full_shape DefaultSpace + +-- | Arrays for storing group results shared between threads +groupResultArrays :: + String -> + SubExp -> + [SegBinOp MCMem] -> + MulticoreGen [[VName]] +groupResultArrays s num_threads reds = + forM reds $ \(SegBinOp _ lam _ shape) -> + forM (lambdaReturnType lam) $ \t -> do + let pt = elemType t + full_shape = Shape [num_threads] <> shape <> arrayShape t + sAllocArray s pt full_shape DefaultSpace + +isLoadBalanced :: Imp.Code -> Bool +isLoadBalanced (a Imp.:>>: b) = isLoadBalanced a && isLoadBalanced b +isLoadBalanced (Imp.For _ _ a) = isLoadBalanced a +isLoadBalanced (Imp.If _ a b) = isLoadBalanced a && isLoadBalanced b +isLoadBalanced (Imp.Comment _ a) = isLoadBalanced a +isLoadBalanced Imp.While {} = False +isLoadBalanced (Imp.Op (Imp.ParLoop _ _ _ code _ _ _)) = isLoadBalanced code +isLoadBalanced _ = True + +segBinOpComm' :: [SegBinOp lore] -> Commutativity +segBinOpComm' = mconcat . map segBinOpComm + +decideScheduling' :: SegOp () lore -> Imp.Code -> Imp.Scheduling +decideScheduling' SegHist {} _ = Imp.Static +decideScheduling' SegScan {} _ = Imp.Static +decideScheduling' (SegRed _ _ reds _ _) code = + case segBinOpComm' reds of + Commutative -> decideScheduling code + Noncommutative -> Imp.Static +decideScheduling' SegMap {} code = decideScheduling code + +decideScheduling :: Imp.Code -> Imp.Scheduling +decideScheduling code = + if isLoadBalanced code + then Imp.Static + else Imp.Dynamic + +-- | Try to extract invariant allocations. If we assume that the +-- given 'Code' is the body of a 'SegOp', then it is always safe to +-- move the immediate allocations to the prebody. +extractAllocations :: Imp.Code -> (Imp.Code, Imp.Code) +extractAllocations segop_code = f segop_code + where + declared = Imp.declaredIn segop_code + f (Imp.DeclareMem name space) = + -- Hoisting declarations out is always safe. + (Imp.DeclareMem name space, mempty) + f (Imp.Allocate name size space) + | not $ freeIn size `namesIntersect` declared = + (Imp.Allocate name size space, mempty) + f (x Imp.:>>: y) = f x <> f y + f (Imp.While cond body) = + (mempty, Imp.While cond body) + f (Imp.For i bound body) = + (mempty, Imp.For i bound body) + f (Imp.Comment s code) = + second (Imp.Comment s) (f code) + f Imp.Free {} = + mempty + f (Imp.If cond tcode fcode) = + let (ta, tcode') = f tcode + (fa, fcode') = f fcode + in (ta <> fa, Imp.If cond tcode' fcode') + f (Imp.Op (Imp.ParLoop s i prebody body postbody free info)) = + let (body_allocs, body') = extractAllocations body + (free_allocs, here_allocs) = f body_allocs + free' = + filter + ( not + . (`nameIn` Imp.declaredIn body_allocs) + . Imp.paramName + ) + free + in ( free_allocs, + here_allocs + <> Imp.Op (Imp.ParLoop s i prebody body' postbody free' info) + ) + f code = + (mempty, code) + +------------------------------- +------- SegHist helpers ------- +------------------------------- +renameHistOpLambda :: [HistOp MCMem] -> MulticoreGen [HistOp MCMem] +renameHistOpLambda hist_ops = + forM hist_ops $ \(HistOp w rf dest neutral shape lam) -> do + lam' <- renameLambda lam + return $ HistOp w rf dest neutral shape lam' + +-- | 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 = + [VName] -> [Imp.TExp Int64] -> MulticoreGen () + +-- | 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) + | -- | Can be done by efficient swaps. + AtomicCAS (DoAtomicUpdate lore r) + | -- | Requires explicit locking. + AtomicLocking (Locking -> DoAtomicUpdate lore r) + +atomicUpdateLocking :: + AtomicBinOp -> + Lambda MCMem -> + AtomicUpdate MCMem () +atomicUpdateLocking atomicBinOp lam + | Just ops_and_ts <- splitOp lam, + all (\(_, t, _, _) -> supportedPrims $ primBitSize t) ops_and_ts = + primOrCas ops_and_ts $ \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 (tvVar old) arr' (sExt32 <$> bucket_offset) op of + Just f -> sOp $ f $ Imp.var y t + Nothing -> + atomicUpdateCAS t a (tvVar old) bucket x $ + x <~~ Imp.BinOpExp op (Imp.var x t) (Imp.var y t) + where + opHasAtomicSupport old arr' bucket' bop = do + let atomic f = Imp.Atomic . f old arr' bucket' + atomic <$> atomicBinOp bop + + primOrCas ops + | all isPrim ops = AtomicPrim + | otherwise = AtomicCAS + + isPrim (op, _, _, _) = isJust $ atomicBinOp op +atomicUpdateLocking _ op + | [Prim t] <- lambdaReturnType op, + [xp, _] <- lambdaParams op, + supportedPrims (primBitSize t) = AtomicCAS $ \[arr] bucket -> do + old <- dPrim "old" t + atomicUpdateCAS t arr (tvVar old) bucket (paramName xp) $ + compileBody' [xp] $ lambdaBody op +atomicUpdateLocking _ op = AtomicLocking $ \locking arrs bucket -> do + old <- dPrim "old" int32 + continue <- dPrimVol "continue" int32 (0 :: Imp.TExp Int32) + + -- Correctly index into locks. + (locks', _locks_space, locks_offset) <- + fullyIndexArray (lockingArray locking) $ lockingMapping locking bucket + + -- Critical section + let try_acquire_lock = do + old <-- (0 :: Imp.TExp Int32) + sOp $ + Imp.Atomic $ + Imp.AtomicCmpXchg + int32 + (tvVar old) + locks' + (sExt32 <$> locks_offset) + (tvVar continue) + (untyped (lockingToLock locking)) + lock_acquired = tvExp continue + -- Even the releasing is done with an atomic rather than a + -- simple write, for memory coherency reasons. + release_lock = do + old <-- lockingToLock locking + sOp $ + Imp.Atomic $ + Imp.AtomicCmpXchg + int32 + (tvVar old) + locks' + (sExt32 <$> locks_offset) + (tvVar continue) + (untyped (lockingToUnlock locking)) + + -- Preparing parameters. It is assumed that the caller has already + -- filled the arr_params. We copy the current value to the + -- accumulator parameters. + 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 + + -- While-loop: Try to insert your value + sWhile (tvExp continue .==. 0) $ do + try_acquire_lock + sUnless (lock_acquired .==. 0) $ do + dLParams acc_params + bind_acc_params + op_body + do_hist + release_lock + where + writeArray bucket arr val = copyDWIMFix arr bucket val [] + +atomicUpdateCAS :: + PrimType -> + VName -> + VName -> + [Imp.TExp Int64] -> + VName -> + MulticoreGen () -> + MulticoreGen () +atomicUpdateCAS 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); + run_loop <- dPrimV "run_loop" (0 :: Imp.TExp Int32) + everythingVolatile $ copyDWIMFix old [] (Var arr) bucket + (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket + + bytes <- toIntegral $ primBitSize t + (to, from) <- getBitConvertFunc $ primBitSize t + -- While-loop: Try to insert your value + let (toBits, _fromBits) = + case t of + FloatType _ -> + ( \v -> Imp.FunExp to [v] bytes, + \v -> Imp.FunExp from [v] t + ) + _ -> (id, id) + + sWhile (tvExp run_loop .==. 0) $ do + x <~~ Imp.var old t + do_op -- Writes result into x + sOp $ + Imp.Atomic $ + Imp.AtomicCmpXchg + bytes + old + arr' + (sExt32 <$> bucket_offset) + (tvVar run_loop) + (toBits (Imp.var x t)) + +-- | 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 + +-- TODO for supporting 8 and 16 bits (and 128) +-- we need a functions for converting to and from bits +getBitConvertFunc :: Int -> MulticoreGen (String, String) +-- getBitConvertFunc 8 = return $ ("to_bits8, from_bits8") +-- getBitConvertFunc 16 = return $ ("to_bits8, from_bits8") +getBitConvertFunc 32 = return ("to_bits32", "from_bits32") +getBitConvertFunc 64 = return ("to_bits64", "from_bits64") +getBitConvertFunc b = error $ "number of bytes is not supported " ++ pretty b + +supportedPrims :: Int -> Bool +supportedPrims 8 = True +supportedPrims 16 = True +supportedPrims 32 = True +supportedPrims 64 = True +supportedPrims _ = False + +-- Supported bytes lengths by GCC (and clang) compiler +toIntegral :: Int -> MulticoreGen PrimType +toIntegral 8 = return int8 +toIntegral 16 = return int16 +toIntegral 32 = return int32 +toIntegral 64 = return int64 +toIntegral b = error $ "number of bytes is not supported for CAS - " ++ pretty b diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs @@ -0,0 +1,359 @@ +module Futhark.CodeGen.ImpGen.Multicore.SegHist + ( compileSegHist, + ) +where + +import Control.Monad +import Data.List (zip4, zip5) +import qualified Futhark.CodeGen.ImpCode.Multicore as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.Multicore.Base +import Futhark.CodeGen.ImpGen.Multicore.SegRed (compileSegRed') +import Futhark.IR.MCMem +import Futhark.MonadFreshNames +import Futhark.Util (chunks, splitFromEnd, takeLast) +import Futhark.Util.IntegralExp (rem) +import Prelude hiding (quot, rem) + +compileSegHist :: + Pattern MCMem -> + SegSpace -> + [HistOp MCMem] -> + KernelBody MCMem -> + TV Int32 -> + MulticoreGen Imp.Code +compileSegHist pat space histops kbody nsubtasks + | [_] <- unSegSpace space = + nonsegmentedHist pat space histops kbody nsubtasks + | otherwise = + segmentedHist pat space histops kbody + +-- | Split some list into chunks equal to the number of values +-- returned by each 'SegBinOp' +segHistOpChunks :: [HistOp lore] -> [a] -> [[a]] +segHistOpChunks = chunks . map (length . histNeutral) + +nonsegmentedHist :: + Pattern MCMem -> + SegSpace -> + [HistOp MCMem] -> + KernelBody MCMem -> + TV Int32 -> + MulticoreGen Imp.Code +nonsegmentedHist pat space histops kbody num_histos = do + let ns = map snd $ unSegSpace space + ns_64 = map toInt64Exp ns + num_histos' = tvExp num_histos + hist_width = toInt64Exp $ histWidth $ head histops + use_subhistogram = sExt64 num_histos' * hist_width .<=. product ns_64 + + histops' <- renameHistOpLambda histops + + collect $ do + flat_idx <- dPrim "iter" int64 + sIf + use_subhistogram + (subHistogram pat flat_idx space histops num_histos kbody) + (atomicHistogram pat flat_idx space histops' kbody) + +-- | +-- Atomic Histogram approach +-- The implementation has three sub-strategies depending on the +-- type of the operator +-- 1. If values are integral scalars, a direct-supported atomic update is used. +-- 2. If values are on one memory location, e.g. a float, then a +-- CAS operation is used to perform the update, where the float is +-- casted to an integral scalar. +-- 1. and 2. currently only works for 32-bit and 64-bit types, +-- but GCC has support for 8-, 16- and 128- bit types as well. +-- 3. Otherwise a locking based approach is used +onOpAtomic :: HistOp MCMem -> MulticoreGen ([VName] -> [Imp.TExp Int64] -> MulticoreGen ()) +onOpAtomic op = do + atomics <- hostAtomics <$> askEnv + let lambda = histOp op + do_op = atomicUpdateLocking atomics lambda + case do_op of + AtomicPrim f -> return f + AtomicCAS f -> return f + AtomicLocking f -> do + -- Allocate a static array of locks + -- as in the GPU backend + let num_locks = 100151 -- This number is taken from the GPU backend + dims = + map toInt64Exp $ + shapeDims (histShape op) ++ [histWidth op] + locks <- + sStaticArray "hist_locks" DefaultSpace int32 $ + Imp.ArrayZeros num_locks + let l' = Locking locks 0 1 0 (pure . (`rem` fromIntegral num_locks) . flattenIndex dims) + return $ f l' + +atomicHistogram :: + Pattern MCMem -> + TV Int64 -> + SegSpace -> + [HistOp MCMem] -> + KernelBody MCMem -> + MulticoreGen () +atomicHistogram pat flat_idx space histops kbody = do + let (is, ns) = unzip $ unSegSpace space + ns_64 = map toInt64Exp ns + let num_red_res = length histops + sum (map (length . histNeutral) histops) + (all_red_pes, map_pes) = splitAt num_red_res $ patternValueElements pat + + atomicOps <- mapM onOpAtomic histops + + body <- collect $ do + zipWithM_ dPrimV_ is $ unflattenIndex ns_64 $ tvExp flat_idx + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody + perOp = chunks $ map (length . histDest) histops + (buckets, vs) = splitAt (length histops) red_res + + let pes_per_op = chunks (map (length . histDest) histops) all_red_pes + forM_ (zip5 histops (perOp vs) buckets atomicOps pes_per_op) $ + \(HistOp dest_w _ _ _ shape lam, vs', bucket, do_op, dest_res) -> do + let (_is_params, vs_params) = splitAt (length vs') $ lambdaParams lam + dest_w' = toInt64Exp dest_w + bucket' = toInt64Exp $ kernelResultSubExp bucket + bucket_in_bounds = bucket' .<. dest_w' .&&. 0 .<=. bucket' + + sComment "save map-out results" $ + forM_ (zip map_pes map_res) $ \(pe, res) -> + copyDWIMFix (patElemName pe) (map Imp.vi64 is) (kernelResultSubExp res) [] + + sComment "perform updates" $ + sWhen bucket_in_bounds $ do + let bucket_is = map Imp.vi64 (init is) ++ [bucket'] + dLParams $ lambdaParams lam + sLoopNest shape $ \is' -> do + forM_ (zip vs_params vs') $ \(p, res) -> + copyDWIMFix (paramName p) [] (kernelResultSubExp res) is' + do_op (map patElemName dest_res) (bucket_is ++ is') + + free_params <- freeParams body (segFlat space : [tvVar flat_idx]) + emit $ Imp.Op $ Imp.ParLoop "atomic_seg_hist" (tvVar flat_idx) mempty body mempty free_params $ segFlat space + +updateHisto :: HistOp MCMem -> [VName] -> [Imp.TExp Int64] -> MulticoreGen () +updateHisto op arrs bucket = do + let acc_params = take (length arrs) $ lambdaParams $ histOp op + bind_acc_params = + forM_ (zip acc_params arrs) $ \(acc_p, arr) -> + copyDWIMFix (paramName acc_p) [] (Var arr) bucket + op_body = compileBody' [] $ lambdaBody $ histOp op + writeArray arr val = copyDWIMFix arr bucket val [] + do_hist = zipWithM_ writeArray arrs $ bodyResult $ lambdaBody $ histOp op + + sComment "Start of body" $ do + dLParams acc_params + bind_acc_params + op_body + do_hist + +-- Generates num_histos sub-histograms of the size +-- of the destination histogram +-- Then for each chunk of the input each subhistogram +-- is computed and finally combined through a segmented reduction +-- across the histogram indicies. +-- This is expected to be fast if len(histDest) is small +subHistogram :: + Pattern MCMem -> + TV Int64 -> + SegSpace -> + [HistOp MCMem] -> + TV Int32 -> + KernelBody MCMem -> + MulticoreGen () +subHistogram pat flat_idx space histops num_histos kbody = do + emit $ Imp.DebugPrint "subHistogram segHist" Nothing + + let (is, ns) = unzip $ unSegSpace space + ns_64 = map toInt64Exp ns + + let pes = patternElements pat + num_red_res = length histops + sum (map (length . histNeutral) histops) + map_pes = drop num_red_res pes + per_red_pes = segHistOpChunks histops $ patternValueElements pat + + -- Allocate array of subhistograms in the calling thread. Each + -- tasks will work in its own private allocations (to avoid false + -- sharing), but this is where they will ultimately copy their + -- results. + global_subhistograms <- forM histops $ \histop -> + forM (histType histop) $ \t -> do + let shape = Shape [tvSize num_histos] <> arrayShape t + sAllocArray "subhistogram" (elemType t) shape DefaultSpace + + let tid' = Imp.vi64 $ segFlat space + flat_idx' = tvExp flat_idx + + (local_subhistograms, prebody) <- collect' $ do + zipWithM_ dPrimV_ is $ unflattenIndex ns_64 $ sExt64 flat_idx' + + forM (zip per_red_pes histops) $ \(pes', histop) -> do + op_local_subhistograms <- forM (histType histop) $ \t -> + sAllocArray "subhistogram" (elemType t) (arrayShape t) DefaultSpace + + forM_ (zip3 pes' op_local_subhistograms (histNeutral histop)) $ \(pe, hist, ne) -> + -- First thread initializes histogram with dest vals. Others + -- initialize with neutral element + sIf + (tid' .==. 0) + (copyDWIMFix hist [] (Var $ patElemName pe) []) + ( sFor "i" (toInt64Exp $ histWidth histop) $ \i -> + sLoopNest (histShape histop) $ \vec_is -> + copyDWIMFix hist (i : vec_is) ne [] + ) + + return op_local_subhistograms + + -- Generate loop body of parallel function + body <- collect $ do + zipWithM_ dPrimV_ is $ unflattenIndex ns_64 $ sExt64 flat_idx' + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody + (buckets, vs) = splitAt (length histops) red_res + perOp = chunks $ map (length . histDest) histops + + sComment "save map-out results" $ + forM_ (zip map_pes map_res) $ \(pe, res) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 is) + (kernelResultSubExp res) + [] + + forM_ (zip4 histops local_subhistograms buckets (perOp vs)) $ + \( histop@(HistOp dest_w _ _ _ shape lam), + histop_subhistograms, + bucket, + vs' + ) -> do + let bucket' = toInt64Exp $ kernelResultSubExp bucket + dest_w' = toInt64Exp dest_w + bucket_in_bounds = bucket' .<. dest_w' .&&. 0 .<=. bucket' + vs_params = takeLast (length vs') $ lambdaParams lam + bucket_is = [bucket'] + + sComment "perform updates" $ + sWhen bucket_in_bounds $ do + dLParams $ lambdaParams lam + sLoopNest shape $ \is' -> do + forM_ (zip vs_params vs') $ \(p, res) -> + copyDWIMFix (paramName p) [] (kernelResultSubExp res) is' + updateHisto histop histop_subhistograms (bucket_is ++ is') + + -- Copy the task-local subhistograms to the global subhistograms, + -- where they will be combined. + postbody <- collect $ + forM_ (zip (concat global_subhistograms) (concat local_subhistograms)) $ + \(global, local) -> copyDWIMFix global [tid'] (Var local) [] + + free_params <- freeParams (prebody <> body <> postbody) (segFlat space : [tvVar flat_idx]) + let (body_allocs, body') = extractAllocations body + emit $ Imp.Op $ Imp.ParLoop "seghist_stage_1" (tvVar flat_idx) (body_allocs <> prebody) body' postbody free_params $ segFlat space + + -- Perform a segmented reduction over the subhistograms + forM_ (zip3 per_red_pes global_subhistograms histops) $ \(red_pes, hists, op) -> do + bucket_id <- newVName "bucket_id" + subhistogram_id <- newVName "subhistogram_id" + + let num_buckets = histWidth op + segred_space = + SegSpace (segFlat space) $ + segment_dims + ++ [(bucket_id, num_buckets)] + ++ [(subhistogram_id, tvSize num_histos)] + + segred_op = SegBinOp Noncommutative (histOp op) (histNeutral op) (histShape op) + + nsubtasks_red <- dPrim "num_tasks" $ IntType Int32 + red_code <- compileSegRed' (Pattern [] red_pes) segred_space [segred_op] nsubtasks_red $ \red_cont -> + red_cont $ + flip map hists $ \subhisto -> + ( Var subhisto, + map Imp.vi64 $ + map fst segment_dims ++ [subhistogram_id, bucket_id] + ) + + let ns_red = map (toInt64Exp . snd) $ unSegSpace segred_space + iterations = product $ init ns_red -- The segmented reduction is sequential over the inner most dimension + scheduler_info = Imp.SchedulerInfo (tvVar nsubtasks_red) (untyped iterations) Imp.Static + red_task = Imp.ParallelTask red_code $ segFlat space + free_params_red <- freeParams red_code [segFlat space, tvVar nsubtasks_red] + emit $ Imp.Op $ Imp.Segop "seghist_red" free_params_red red_task Nothing mempty scheduler_info + where + segment_dims = init $ unSegSpace space + +-- This implementation for a Segmented Hist only +-- parallelize over the segments, +-- where each segment is updated sequentially. +segmentedHist :: + Pattern MCMem -> + SegSpace -> + [HistOp MCMem] -> + KernelBody MCMem -> + MulticoreGen Imp.Code +segmentedHist pat space histops kbody = do + emit $ Imp.DebugPrint "Segmented segHist" Nothing + -- Iteration variable over the segments + segments_i <- dPrim "segment_iter" $ IntType Int64 + collect $ do + par_body <- compileSegHistBody (tvExp segments_i) pat space histops kbody + free_params <- freeParams par_body [segFlat space, tvVar segments_i] + let (body_allocs, body') = extractAllocations par_body + emit $ Imp.Op $ Imp.ParLoop "segmented_hist" (tvVar segments_i) body_allocs body' mempty free_params $ segFlat space + +compileSegHistBody :: + Imp.TExp Int64 -> + Pattern MCMem -> + SegSpace -> + [HistOp MCMem] -> + KernelBody MCMem -> + MulticoreGen Imp.Code +compileSegHistBody idx pat space histops kbody = do + let (is, ns) = unzip $ unSegSpace space + ns_64 = map toInt64Exp ns + + let num_red_res = length histops + sum (map (length . histNeutral) histops) + map_pes = drop num_red_res $ patternValueElements pat + per_red_pes = segHistOpChunks histops $ patternValueElements pat + + collect $ do + let inner_bound = last ns_64 + sFor "i" inner_bound $ \i -> do + zipWithM_ dPrimV_ (init is) $ unflattenIndex (init ns_64) idx + dPrimV_ (last is) i + + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = + splitFromEnd (length map_pes) $ + map kernelResultSubExp $ kernelBodyResult kbody + (buckets, vs) = splitAt (length histops) red_res + perOp = chunks $ map (length . histDest) histops + + forM_ (zip4 per_red_pes histops (perOp vs) buckets) $ + \(red_pes, HistOp dest_w _ _ _ shape lam, vs', bucket) -> do + let (is_params, vs_params) = splitAt (length vs') $ lambdaParams lam + bucket' = toInt64Exp bucket + dest_w' = toInt64Exp dest_w + bucket_in_bounds = bucket' .<. dest_w' .&&. 0 .<=. bucket' + + sComment "save map-out results" $ + forM_ (zip map_pes map_res) $ \(pe, res) -> + copyDWIMFix (patElemName pe) (map Imp.vi64 is) res [] + + sComment "perform updates" $ + sWhen bucket_in_bounds $ do + dLParams $ lambdaParams lam + sLoopNest shape $ \vec_is -> do + -- Index + let buck = toInt64Exp bucket + forM_ (zip red_pes is_params) $ \(pe, p) -> + copyDWIMFix (paramName p) [] (Var $ patElemName pe) (map Imp.vi64 (init is) ++ [buck] ++ vec_is) + -- Value at index + forM_ (zip vs_params vs') $ \(p, v) -> + copyDWIMFix (paramName p) [] v vec_is + compileStms mempty (bodyStms $ lambdaBody lam) $ + forM_ (zip red_pes $ bodyResult $ lambdaBody lam) $ + \(pe, se) -> copyDWIMFix (patElemName pe) (map Imp.vi64 (init is) ++ [buck] ++ vec_is) se [] diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs @@ -0,0 +1,62 @@ +module Futhark.CodeGen.ImpGen.Multicore.SegMap + ( compileSegMap, + ) +where + +import Control.Monad +import qualified Futhark.CodeGen.ImpCode.Multicore as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.Multicore.Base +import Futhark.IR.MCMem +import Futhark.Transform.Rename + +writeResult :: + [VName] -> + PatElemT dec -> + KernelResult -> + MulticoreGen () +writeResult is pe (Returns _ se) = + copyDWIMFix (patElemName pe) (map Imp.vi64 is) se [] +writeResult _ pe (WriteReturns rws _ idx_vals) = do + let (iss, vs) = unzip idx_vals + rws' = map toInt64Exp rws + forM_ (zip iss vs) $ \(slice, v) -> 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 + in_bounds = foldl1 (.&&.) $ zipWith condInBounds slice' rws' + when_in_bounds = copyDWIM (patElemName pe) slice' v [] + sWhen in_bounds when_in_bounds +writeResult _ _ res = + error $ "writeResult: cannot handle " ++ pretty res + +compileSegMapBody :: + TV Int64 -> + Pattern MCMem -> + SegSpace -> + KernelBody MCMem -> + MulticoreGen Imp.Code +compileSegMapBody flat_idx pat space (KernelBody _ kstms kres) = do + let (is, ns) = unzip $ unSegSpace space + ns' = map toInt64Exp ns + kstms' <- mapM renameStm kstms + collect $ do + emit $ Imp.DebugPrint "SegMap fbody" Nothing + zipWithM_ dPrimV_ is $ map sExt64 $ unflattenIndex ns' $ tvExp flat_idx + compileStms (freeIn kres) kstms' $ + zipWithM_ (writeResult is) (patternElements pat) kres + +compileSegMap :: + Pattern MCMem -> + SegSpace -> + KernelBody MCMem -> + MulticoreGen Imp.Code +compileSegMap pat space kbody = + collect $ do + flat_par_idx <- dPrim "iter" int64 + body <- compileSegMapBody flat_par_idx pat space kbody + free_params <- freeParams body [segFlat space, tvVar flat_par_idx] + let (body_allocs, body') = extractAllocations body + emit $ Imp.Op $ Imp.ParLoop "segmap" (tvVar flat_par_idx) body_allocs body' mempty free_params $ segFlat space diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs @@ -0,0 +1,267 @@ +module Futhark.CodeGen.ImpGen.Multicore.SegRed + ( compileSegRed, + compileSegRed', + ) +where + +import Control.Monad +import qualified Futhark.CodeGen.ImpCode.Multicore as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.Multicore.Base +import Futhark.IR.MCMem +import Futhark.Util (chunks) +import Prelude hiding (quot, rem) + +type DoSegBody = (([(SubExp, [Imp.TExp Int64])] -> MulticoreGen ()) -> MulticoreGen ()) + +-- | Generate code for a SegRed construct +compileSegRed :: + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + KernelBody MCMem -> + TV Int32 -> + MulticoreGen Imp.Code +compileSegRed pat space reds kbody nsubtasks = + compileSegRed' pat space reds nsubtasks $ \red_cont -> + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = splitAt (segBinOpResults reds) $ kernelBodyResult kbody + + 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 [] + +-- | Like 'compileSegRed', but where the body is a monadic action. +compileSegRed' :: + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + TV Int32 -> + DoSegBody -> + MulticoreGen Imp.Code +compileSegRed' pat space reds nsubtasks kbody + | [_] <- unSegSpace space = + nonsegmentedReduction pat space reds nsubtasks kbody + | otherwise = + segmentedReduction pat space reds kbody + +-- | A SegBinOp with auxiliary information. +data SegBinOpSlug = SegBinOpSlug + { slugOp :: SegBinOp MCMem, + -- | The array in which we write the intermediate results, indexed + -- by the flat/physical thread ID. + slugResArrs :: [VName] + } + +slugBody :: SegBinOpSlug -> Body MCMem +slugBody = lambdaBody . segBinOpLambda . slugOp + +slugParams :: SegBinOpSlug -> [LParam MCMem] +slugParams = lambdaParams . segBinOpLambda . slugOp + +slugNeutral :: SegBinOpSlug -> [SubExp] +slugNeutral = segBinOpNeutral . slugOp + +slugShape :: SegBinOpSlug -> Shape +slugShape = segBinOpShape . slugOp + +accParams, nextParams :: SegBinOpSlug -> [LParam MCMem] +accParams slug = take (length (slugNeutral slug)) $ slugParams slug +nextParams slug = drop (length (slugNeutral slug)) $ slugParams slug + +nonsegmentedReduction :: + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + TV Int32 -> + DoSegBody -> + MulticoreGen Imp.Code +nonsegmentedReduction pat space reds nsubtasks kbody = collect $ do + thread_res_arrs <- groupResultArrays "reduce_stage_1_tid_res_arr" (tvSize nsubtasks) reds + let slugs1 = zipWith SegBinOpSlug reds thread_res_arrs + nsubtasks' = tvExp nsubtasks + + reductionStage1 space slugs1 kbody + reds2 <- renameSegBinOp reds + let slugs2 = zipWith SegBinOpSlug reds2 thread_res_arrs + reductionStage2 pat space nsubtasks' slugs2 + +reductionStage1 :: + SegSpace -> + [SegBinOpSlug] -> + DoSegBody -> + MulticoreGen () +reductionStage1 space slugs kbody = do + let (is, ns) = unzip $ unSegSpace space + ns' = map (sExt64 . toInt32Exp) ns + flat_idx <- dPrim "iter" int64 + + -- Create local accumulator variables in which we carry out the + -- sequential reduction of this function. If we are dealing with + -- vectorised operators, then this implies a private allocation. If + -- the original operand type of the reduction is a memory block, + -- then our hands are unfortunately tied, and we have to use exactly + -- that memory. This is likely to be slow. + + (slug_local_accs, prebody) <- collect' $ do + dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs + + forM slugs $ \slug -> do + let shape = segBinOpShape $ slugOp slug + + forM (zip (accParams slug) (slugNeutral slug)) $ \(p, ne) -> do + -- Declare accumulator variable. + acc <- + case paramType p of + Prim pt + | shape == mempty -> + tvVar <$> dPrim "local_acc" pt + | otherwise -> + sAllocArray "local_acc" pt shape DefaultSpace + _ -> + pure $ paramName p + + -- Now neutral-initialise the accumulator. + sLoopNest (slugShape slug) $ \vec_is -> + copyDWIMFix acc vec_is ne [] + + pure acc + + fbody <- collect $ do + zipWithM_ dPrimV_ is $ unflattenIndex ns' $ tvExp flat_idx + kbody $ \all_red_res -> do + let all_red_res' = segBinOpChunks (map slugOp slugs) all_red_res + forM_ (zip3 all_red_res' slugs slug_local_accs) $ \(red_res, slug, local_accs) -> + sLoopNest (slugShape slug) $ \vec_is -> do + let lamtypes = lambdaReturnType $ segBinOpLambda $ slugOp slug + -- Load accum params + sComment "Load accum params" $ + forM_ (zip3 (accParams slug) local_accs lamtypes) $ + \(p, local_acc, t) -> + when (primType t) $ + copyDWIMFix (paramName p) [] (Var local_acc) vec_is + + sComment "Load next params" $ + forM_ (zip (nextParams slug) red_res) $ \(p, (res, res_is)) -> + copyDWIMFix (paramName p) [] res (res_is ++ vec_is) + + sComment "Red body" $ + compileStms mempty (bodyStms $ slugBody slug) $ + forM_ (zip local_accs (bodyResult $ slugBody slug)) $ + \(local_acc, se) -> + copyDWIMFix local_acc vec_is se [] + + postbody <- collect $ + forM_ (zip slugs slug_local_accs) $ \(slug, local_accs) -> + forM (zip (slugResArrs slug) local_accs) $ \(acc, local_acc) -> + copyDWIMFix acc [Imp.vi64 $ segFlat space] (Var local_acc) [] + + free_params <- freeParams (prebody <> fbody <> postbody) (segFlat space : [tvVar flat_idx]) + let (body_allocs, fbody') = extractAllocations fbody + emit $ Imp.Op $ Imp.ParLoop "segred_stage_1" (tvVar flat_idx) (body_allocs <> prebody) fbody' postbody free_params $ segFlat space + +reductionStage2 :: + Pattern MCMem -> + SegSpace -> + Imp.TExp Int32 -> + [SegBinOpSlug] -> + MulticoreGen () +reductionStage2 pat space nsubtasks slugs = do + let per_red_pes = segBinOpChunks (map slugOp slugs) $ patternValueElements pat + phys_id = Imp.vi64 (segFlat space) + sComment "neutral-initialise the output" $ + forM_ (zip (map slugOp slugs) per_red_pes) $ \(red, red_res) -> + forM_ (zip red_res $ segBinOpNeutral red) $ \(pe, ne) -> + sLoopNest (segBinOpShape red) $ \vec_is -> + copyDWIMFix (patElemName pe) vec_is ne [] + + dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs + + sFor "i" nsubtasks $ \i' -> do + mkTV (segFlat space) int64 <-- i' + sComment "Apply main thread reduction" $ + forM_ (zip slugs per_red_pes) $ \(slug, red_res) -> + sLoopNest (slugShape slug) $ \vec_is -> do + sComment "load acc params" $ + forM_ (zip (accParams slug) red_res) $ \(p, pe) -> + copyDWIMFix (paramName p) [] (Var $ patElemName pe) vec_is + sComment "load next params" $ + forM_ (zip (nextParams slug) (slugResArrs slug)) $ \(p, acc) -> + copyDWIMFix (paramName p) [] (Var acc) (phys_id : vec_is) + sComment "red body" $ + compileStms mempty (bodyStms $ slugBody slug) $ + forM_ (zip red_res (bodyResult $ slugBody slug)) $ + \(pe, se') -> copyDWIMFix (patElemName pe) vec_is se' [] + +-- Each thread reduces over the number of segments +-- each of which is done sequentially +-- Maybe we should select the work of the inner loop +-- based on n_segments and dimensions etc. +segmentedReduction :: + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + DoSegBody -> + MulticoreGen Imp.Code +segmentedReduction pat space reds kbody = + collect $ do + n_par_segments <- dPrim "segment_iter" $ IntType Int64 + body <- compileSegRedBody n_par_segments pat space reds kbody + free_params <- freeParams body (segFlat space : [tvVar n_par_segments]) + let (body_allocs, body') = extractAllocations body + emit $ Imp.Op $ Imp.ParLoop "segmented_segred" (tvVar n_par_segments) body_allocs body' mempty free_params $ segFlat space + +compileSegRedBody :: + TV Int64 -> + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + DoSegBody -> + MulticoreGen Imp.Code +compileSegRedBody n_segments pat space reds kbody = do + let (is, ns) = unzip $ unSegSpace space + ns_64 = map toInt64Exp ns + inner_bound = last ns_64 + n_segments' = tvExp n_segments + + let per_red_pes = segBinOpChunks reds $ patternValueElements pat + -- Perform sequential reduce on inner most dimension + collect $ do + flat_idx <- dPrimVE "flat_idx" $ n_segments' * inner_bound + zipWithM_ dPrimV_ is $ unflattenIndex ns_64 flat_idx + sComment "neutral-initialise the accumulators" $ + forM_ (zip per_red_pes reds) $ \(pes, red) -> + forM_ (zip pes (segBinOpNeutral red)) $ \(pe, ne) -> + sLoopNest (segBinOpShape red) $ \vec_is -> + copyDWIMFix (patElemName pe) (map Imp.vi64 (init is) ++ vec_is) ne [] + + sComment "main body" $ do + dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) reds + sFor "i" inner_bound $ \i -> do + zipWithM_ + (<--) + (map (`mkTV` int64) $ init is) + (unflattenIndex (init ns_64) (sExt64 n_segments')) + dPrimV_ (last is) i + kbody $ \all_red_res -> do + let red_res' = chunks (map (length . segBinOpNeutral) reds) all_red_res + forM_ (zip3 per_red_pes reds red_res') $ \(pes, red, res') -> + sLoopNest (segBinOpShape red) $ \vec_is -> do + sComment "load accum" $ do + let acc_params = take (length (segBinOpNeutral red)) $ (lambdaParams . segBinOpLambda) red + forM_ (zip acc_params pes) $ \(p, pe) -> + copyDWIMFix (paramName p) [] (Var $ patElemName pe) (map Imp.vi64 (init is) ++ vec_is) + + sComment "load new val" $ do + let next_params = drop (length (segBinOpNeutral red)) $ (lambdaParams . segBinOpLambda) red + forM_ (zip next_params res') $ \(p, (res, res_is)) -> + copyDWIMFix (paramName p) [] res (res_is ++ vec_is) + + sComment "apply reduction" $ do + let lbody = (lambdaBody . segBinOpLambda) red + compileStms mempty (bodyStms lbody) $ + sComment "write back to res" $ + forM_ (zip pes (bodyResult lbody)) $ + \(pe, se') -> copyDWIMFix (patElemName pe) (map Imp.vi64 (init is) ++ vec_is) se' [] diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs @@ -0,0 +1,298 @@ +module Futhark.CodeGen.ImpGen.Multicore.SegScan + ( compileSegScan, + ) +where + +import Control.Monad +import Data.List (zip4) +import qualified Futhark.CodeGen.ImpCode.Multicore as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.Multicore.Base +import Futhark.IR.MCMem +import Futhark.Util.IntegralExp (quot, rem) +import Prelude hiding (quot, rem) + +-- Compile a SegScan construct +compileSegScan :: + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + KernelBody MCMem -> + TV Int32 -> + MulticoreGen Imp.Code +compileSegScan pat space reds kbody nsubtasks + | [_] <- unSegSpace space = + nonsegmentedScan pat space reds kbody nsubtasks + | otherwise = + segmentedScan pat space reds kbody + +xParams, yParams :: SegBinOp MCMem -> [LParam MCMem] +xParams scan = + take (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) +yParams scan = + drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) + +lamBody :: SegBinOp MCMem -> Body MCMem +lamBody = lambdaBody . segBinOpLambda + +nonsegmentedScan :: + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + KernelBody MCMem -> + TV Int32 -> + MulticoreGen Imp.Code +nonsegmentedScan pat space scan_ops kbody nsubtasks = do + emit $ Imp.DebugPrint "nonsegmented segScan" Nothing + collect $ do + scanStage1 pat space scan_ops kbody + + let nsubtasks' = tvExp nsubtasks + sWhen (nsubtasks' .>. 1) $ do + scan_ops2 <- renameSegBinOp scan_ops + scanStage2 pat nsubtasks space scan_ops2 kbody + scan_ops3 <- renameSegBinOp scan_ops + scanStage3 pat space scan_ops3 kbody + +scanStage1 :: + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + KernelBody MCMem -> + MulticoreGen () +scanStage1 pat space scan_ops kbody = do + let (all_scan_res, map_res) = splitAt (segBinOpResults scan_ops) $ kernelBodyResult kbody + per_scan_res = segBinOpChunks scan_ops all_scan_res + per_scan_pes = segBinOpChunks scan_ops $ patternValueElements pat + let (is, ns) = unzip $ unSegSpace space + ns' = map toInt64Exp ns + iter <- dPrim "iter" $ IntType Int32 + + -- Stage 1 : each thread partially scans a chunk of the input + -- Writes directly to the resulting array + (local_accs, prebody) <- collect' $ do + dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) scan_ops + forM scan_ops $ \scan_op -> do + let shape = segBinOpShape scan_op + ts = lambdaReturnType $ segBinOpLambda scan_op + forM (zip3 (xParams scan_op) (segBinOpNeutral scan_op) ts) $ \(p, ne, t) -> do + acc <- + case shapeDims shape of + [] -> pure $ paramName p + _ -> do + let pt = elemType t + sAllocArray "local_acc" pt (shape <> arrayShape t) DefaultSpace + + -- Now neutral-initialise the accumulator. + sLoopNest (segBinOpShape scan_op) $ \vec_is -> + copyDWIMFix acc vec_is ne [] + + pure acc + + body <- collect $ do + zipWithM_ dPrimV_ is $ unflattenIndex ns' $ tvExp iter + sComment "stage 1 scan body" $ + compileStms mempty (kernelBodyStms kbody) $ do + sComment "write mapped values results to memory" $ do + let map_arrs = drop (segBinOpResults scan_ops) $ patternElements pat + zipWithM_ (compileThreadResult space) map_arrs map_res + + forM_ (zip4 per_scan_pes scan_ops per_scan_res local_accs) $ \(pes, scan_op, scan_res, acc) -> + sLoopNest (segBinOpShape scan_op) $ \vec_is -> do + -- Read accum value + forM_ (zip (xParams scan_op) acc) $ \(p, acc') -> + copyDWIMFix (paramName p) [] (Var acc') vec_is + + -- Read next value + sComment "Read next values" $ + forM_ (zip (yParams scan_op) scan_res) $ \(p, se) -> + copyDWIMFix (paramName p) [] (kernelResultSubExp se) vec_is + + compileStms mempty (bodyStms $ lamBody scan_op) $ + forM_ (zip3 acc pes (bodyResult $ lamBody scan_op)) $ + \(acc', pe, se) -> do + copyDWIMFix (patElemName pe) (map Imp.vi64 is ++ vec_is) se [] + copyDWIMFix acc' vec_is se [] + + free_params <- freeParams (prebody <> body) (segFlat space : [tvVar iter]) + let (body_allocs, body') = extractAllocations body + emit $ Imp.Op $ Imp.ParLoop "scan_stage_1" (tvVar iter) (body_allocs <> prebody) body' mempty free_params $ segFlat space + +scanStage2 :: + Pattern MCMem -> + TV Int32 -> + SegSpace -> + [SegBinOp MCMem] -> + KernelBody MCMem -> + MulticoreGen () +scanStage2 pat nsubtasks space scan_ops kbody = do + emit $ Imp.DebugPrint "nonsegmentedScan stage 2" Nothing + let (is, ns) = unzip $ unSegSpace space + ns_64 = map toInt64Exp ns + per_scan_pes = segBinOpChunks scan_ops $ patternValueElements pat + nsubtasks' = tvExp nsubtasks + + dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) scan_ops + offset <- dPrimV "offset" (0 :: Imp.TExp Int64) + let offset' = tvExp offset + offset_index <- dPrimV "offset_index" (0 :: Imp.TExp Int64) + let offset_index' = tvExp offset_index + + -- Parameters used to find the chunk sizes + -- Perhaps get this information from ``scheduling information`` + -- instead of computing it manually here. + let iter_pr_subtask = product ns_64 `quot` sExt64 nsubtasks' + remainder = product ns_64 `rem` sExt64 nsubtasks' + + accs <- resultArrays "scan_stage_2_accum" scan_ops + forM_ (zip scan_ops accs) $ \(scan_op, acc) -> + sLoopNest (segBinOpShape scan_op) $ \vec_is -> + forM_ (zip acc $ segBinOpNeutral scan_op) $ \(acc', ne) -> + copyDWIMFix acc' vec_is ne [] + + -- Perform sequential scan over the last element of each chunk + sFor "i" (nsubtasks' - 1) $ \i -> do + offset <-- iter_pr_subtask + sWhen (sExt64 i .<. remainder) (offset <-- offset' + 1) + offset_index <-- offset_index' + offset' + zipWithM_ dPrimV_ is $ unflattenIndex ns_64 $ sExt64 offset_index' + + compileStms mempty (kernelBodyStms kbody) $ + forM_ (zip3 per_scan_pes scan_ops accs) $ \(pes, scan_op, acc) -> + sLoopNest (segBinOpShape scan_op) $ \vec_is -> do + sComment "Read carry in" $ + forM_ (zip (xParams scan_op) acc) $ \(p, acc') -> + copyDWIMFix (paramName p) [] (Var acc') vec_is + + sComment "Read next values" $ + forM_ (zip (yParams scan_op) pes) $ \(p, pe) -> + copyDWIMFix (paramName p) [] (Var $ patElemName pe) ((offset_index' - 1) : vec_is) + + compileStms mempty (bodyStms $ lamBody scan_op) $ + forM_ (zip3 acc pes (bodyResult $ lamBody scan_op)) $ + \(acc', pe, se) -> do + copyDWIMFix (patElemName pe) ((offset_index' - 1) : vec_is) se [] + copyDWIMFix acc' vec_is se [] + +-- Stage 3 : Finally each thread partially scans a chunk of the input +-- reading its corresponding carry-in +scanStage3 :: + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + KernelBody MCMem -> + MulticoreGen () +scanStage3 pat space scan_ops kbody = do + let (is, ns) = unzip $ unSegSpace space + all_scan_res = take (segBinOpResults scan_ops) $ kernelBodyResult kbody + per_scan_res = segBinOpChunks scan_ops all_scan_res + per_scan_pes = segBinOpChunks scan_ops $ patternValueElements pat + ns' = map toInt64Exp ns + + iter <- dPrimV "iter" (0 :: Imp.TExp Int64) + let iter' = tvExp iter + + (local_accs, prebody) <- collect' $ do + dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) scan_ops + forM (zip scan_ops per_scan_pes) $ \(scan_op, pes) -> do + let shape = segBinOpShape scan_op + ts = lambdaReturnType $ segBinOpLambda scan_op + forM (zip4 (xParams scan_op) pes ts $ segBinOpNeutral scan_op) $ \(p, pe, t, ne) -> do + acc <- + case shapeDims shape of + [] -> pure $ paramName p + _ -> do + let pt = elemType t + sAllocArray "local_acc" pt (shape <> arrayShape t) DefaultSpace + + -- Initialise the accumulator with neutral from previous chunk. + -- or read neutral if first ``iter`` + sLoopNest (segBinOpShape scan_op) $ \vec_is -> do + let read_carry_in = + copyDWIMFix acc vec_is (Var $ patElemName pe) (iter' - 1 : vec_is) + read_neutral = + copyDWIMFix acc vec_is ne [] + sIf (iter' .==. 0) read_neutral read_carry_in + pure acc + + body <- collect $ do + zipWithM_ dPrimV_ is $ unflattenIndex ns' iter' + sComment "stage 3 scan body" $ + compileStms mempty (kernelBodyStms kbody) $ + forM_ (zip4 per_scan_pes scan_ops per_scan_res local_accs) $ \(pes, scan_op, scan_res, acc) -> + sLoopNest (segBinOpShape scan_op) $ \vec_is -> do + forM_ (zip (xParams scan_op) acc) $ \(p, acc') -> + copyDWIMFix (paramName p) [] (Var acc') vec_is + + -- Read next value + forM_ (zip (yParams scan_op) scan_res) $ \(p, se) -> + copyDWIMFix (paramName p) [] (kernelResultSubExp se) vec_is + + compileStms mempty (bodyStms $ lamBody scan_op) $ + forM_ (zip3 pes (bodyResult $ lamBody scan_op) acc) $ + \(pe, se, acc') -> do + copyDWIMFix (patElemName pe) (map Imp.vi64 is ++ vec_is) se [] + copyDWIMFix acc' vec_is se [] + + free_params' <- freeParams (prebody <> body) (segFlat space : [tvVar iter]) + let (body_allocs, body') = extractAllocations body + emit $ Imp.Op $ Imp.ParLoop "scan_stage_3" (tvVar iter) (body_allocs <> prebody) body' mempty free_params' $ segFlat space + +-- This implementation for a Segmented scan only +-- parallelize over the segments and each segment is +-- scanned sequentially. +segmentedScan :: + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + KernelBody MCMem -> + MulticoreGen Imp.Code +segmentedScan pat space scan_ops kbody = do + emit $ Imp.DebugPrint "segmented segScan" Nothing + collect $ do + segment_iter <- dPrim "segment_iter" $ IntType Int64 + body <- compileSegScanBody (tvExp segment_iter) pat space scan_ops kbody + free_params <- freeParams body (segFlat space : [tvVar segment_iter]) + let (body_allocs, body') = extractAllocations body + emit $ Imp.Op $ Imp.ParLoop "seg_scan" (tvVar segment_iter) body_allocs body' mempty free_params $ segFlat space + +compileSegScanBody :: + Imp.TExp Int64 -> + Pattern MCMem -> + SegSpace -> + [SegBinOp MCMem] -> + KernelBody MCMem -> + MulticoreGen Imp.Code +compileSegScanBody segment_i pat space scan_ops kbody = do + let (is, ns) = unzip $ unSegSpace space + ns_64 = map toInt64Exp ns + + let per_scan_pes = segBinOpChunks scan_ops $ patternValueElements pat + collect $ + forM_ (zip scan_ops per_scan_pes) $ \(scan_op, scan_pes) -> do + dScope Nothing $ scopeOfLParams $ lambdaParams $ segBinOpLambda scan_op + let (scan_x_params, scan_y_params) = splitAt (length $ segBinOpNeutral scan_op) $ (lambdaParams . segBinOpLambda) scan_op + + forM_ (zip scan_x_params $ segBinOpNeutral scan_op) $ \(p, ne) -> + copyDWIMFix (paramName p) [] ne [] + + let inner_bound = last ns_64 + -- Perform a sequential scan over the segment ``segment_i`` + sFor "i" inner_bound $ \i -> do + zipWithM_ dPrimV_ (init is) $ unflattenIndex (init ns_64) segment_i + dPrimV_ (last is) i + compileStms mempty (kernelBodyStms kbody) $ do + let (scan_res, map_res) = splitAt (length $ segBinOpNeutral scan_op) $ kernelBodyResult kbody + sComment "write to-scan values to parameters" $ + forM_ (zip scan_y_params scan_res) $ \(p, se) -> + copyDWIMFix (paramName p) [] (kernelResultSubExp se) [] + + sComment "write mapped values results to memory" $ + forM_ (zip (drop (length $ segBinOpNeutral scan_op) $ patternElements pat) map_res) $ \(pe, se) -> + copyDWIMFix (patElemName pe) (map Imp.vi64 is) (kernelResultSubExp se) [] + + sComment "combine with carry and write to memory" $ + compileStms mempty (bodyStms $ lambdaBody $ segBinOpLambda scan_op) $ + forM_ (zip3 scan_x_params scan_pes (bodyResult $ lambdaBody $ segBinOpLambda scan_op)) $ \(p, pe, se) -> do + copyDWIMFix (patElemName pe) (map Imp.vi64 is) se [] + copyDWIMFix (paramName p) [] se [] diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -330,12 +330,12 @@ 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)) + slice j m = DimSlice j m (constant (1 :: Int64)) -- | Are these indexes out-of-bounds for the array? eOutOfBounds :: @@ -350,10 +350,10 @@ let checkDim w i = do less_than_zero <- letSubExp "less_than_zero" $ - BasicOp $ CmpOp (CmpSlt Int32) i (constant (0 :: Int32)) + BasicOp $ CmpOp (CmpSlt Int64) i (constant (0 :: Int64)) greater_than_size <- letSubExp "greater_than_size" $ - BasicOp $ CmpOp (CmpSle Int32) w i + 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' @@ -479,7 +479,7 @@ -- | 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 @@ -579,7 +579,7 @@ runWriterT $ instantiateShapes instantiate ts where instantiate _ = do - v <- lift $ newIdent "size" $ Prim int32 + v <- lift $ newIdent "size" $ Prim int64 tell [v] return $ Var $ identName v 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 @@ -204,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 SplitSpace {} = pure [Prim int64] + opType (GetSize _ _) = pure [Prim int64] + opType (GetSizeMax _) = pure [Prim int64] opType CmpSizeLe {} = pure [Prim Bool] - opType CalcNumGroups {} = pure [Prim int32] + opType CalcNumGroups {} = pure [Prim int64] instance AliasedOp SizeOp where opAliases _ = [mempty] @@ -251,14 +251,14 @@ 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] + 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 int32] x +typeCheckSizeOp (CmpSizeLe _ _ x) = TC.require [Prim int64] x typeCheckSizeOp (CalcNumGroups w _ group_size) = do TC.require [Prim int64] w - TC.require [Prim int32] group_size + TC.require [Prim int64] group_size -- | A host-level operation; parameterised by what else it can do. data HostOp lore op @@ -357,8 +357,8 @@ SegLevel -> TC.TypeM lore () checkSegLevel Nothing lvl = do - TC.require [Prim int32] $ unCount $ segNumGroups lvl - TC.require [Prim int32] $ unCount $ segGroupSize lvl + 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 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 @@ -17,7 +17,7 @@ where import Control.Category -import Data.Int (Int32) +import Data.Int (Int64) import Data.Traversable import Futhark.IR.Prop.Names (FreeIn) import Futhark.Transform.Substitute @@ -37,7 +37,7 @@ -- impose constraints on the valid values. data SizeClass = -- | A threshold with an optional default. - SizeThreshold KernelPath (Maybe Int32) + SizeThreshold KernelPath (Maybe Int64) | SizeGroup | SizeNumGroups | SizeTile @@ -45,7 +45,7 @@ -- maximum can be handy. SizeLocalMemory | -- | A bespoke size with a default. - SizeBespoke Name Int32 + SizeBespoke Name Int64 deriving (Eq, Ord, Show, Generic) instance SexpIso SizeClass where @@ -72,7 +72,7 @@ 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 diff --git a/src/Futhark/IR/MC.hs b/src/Futhark/IR/MC.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/MC.hs @@ -0,0 +1,82 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} + +-- | A representation for multicore CPU parallelism. +module Futhark.IR.MC + ( -- * The Lore definition + MC, + + -- * Simplification + simplifyProg, + + -- * Module re-exports + module Futhark.IR.Prop, + module Futhark.IR.Traversals, + module Futhark.IR.Pretty, + module Futhark.IR.Syntax, + module Futhark.IR.SegOp, + module Futhark.IR.SOACS.SOAC, + module Futhark.IR.MC.Op, + ) +where + +import Futhark.Binder +import Futhark.Construct +import Futhark.IR.MC.Op +import Futhark.IR.Pretty +import Futhark.IR.Prop +import Futhark.IR.SOACS.SOAC hiding (HistOp (..)) +import qualified Futhark.IR.SOACS.Simplify as SOAC +import Futhark.IR.SegOp +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 + +data MC + +instance Decorations MC where + type Op MC = MCOp MC (SOAC MC) + +instance ASTLore MC where + expTypesFromPattern = return . expExtTypesFromPattern + +instance TypeCheck.CheckableOp MC where + checkOp = typeCheckMCOp typeCheckSOAC + +instance TypeCheck.Checkable MC + +instance Bindable MC where + mkBody = Body () + mkExpPat ctx val _ = basicPattern ctx val + mkExpDec _ _ = () + mkLetNames = simpleMkLetNames + +instance BinderOps MC + +instance BinderOps (Engine.Wise MC) + +instance PrettyLore MC + +simpleMC :: Simplify.SimpleOps MC +simpleMC = Simplify.bindableSimpleOps $ simplifyMCOp SOAC.simplifySOAC + +simplifyProg :: Prog MC -> PassM (Prog MC) +simplifyProg = Simplify.simplifyProg simpleMC rules blockers + where + blockers = Engine.noExtraHoistBlockers + rules = standardRules <> segOpRules + +instance HasSegOp MC where + type SegOpLevel MC = () + asSegOp = const Nothing + segOp = ParOp Nothing + +instance HasSegOp (Engine.Wise MC) where + type SegOpLevel (Engine.Wise MC) = () + asSegOp = const Nothing + segOp = ParOp Nothing diff --git a/src/Futhark/IR/MC/Op.hs b/src/Futhark/IR/MC/Op.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/MC/Op.hs @@ -0,0 +1,175 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | Definitions for multicore operations. +-- +-- Most of the interesting stuff is in "Futhark.IR.SegOp", which is +-- also re-exported from here. +module Futhark.IR.MC.Op + ( MCOp (..), + typeCheckMCOp, + simplifyMCOp, + module Futhark.IR.SegOp, + ) +where + +import Control.Category +import Data.Bifunctor (first) +import Futhark.Analysis.Metrics +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.IR +import Futhark.IR.Aliases (Aliases) +import Futhark.IR.Prop.Aliases +import Futhark.IR.SegOp +import qualified Futhark.Optimise.Simplify as Simplify +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.Pretty + ( Pretty, + nestedBlock, + ppr, + (<+>), + (), + ) +import GHC.Generics (Generic) +import Language.SexpGrammar as Sexp +import Language.SexpGrammar.Generic +import Prelude hiding (id, (.)) + +-- | An operation for the multicore representation. Feel free to +-- extend this on an ad hoc basis as needed. Parameterised with some +-- other operation. +data MCOp lore op + = -- | The first 'SegOp' (if it exists) contains nested parallelism, + -- while the second one has a fully sequential body. They are + -- semantically fully equivalent. + ParOp + (Maybe (SegOp () lore)) + (SegOp () lore) + | -- | Something else (in practice often a SOAC). + OtherOp op + deriving (Eq, Ord, Show, Generic) + +instance (Decorations lore, SexpIso op) => SexpIso (MCOp lore op) where + sexpIso = + match $ + With (. Sexp.list (Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el sexpIso)) + End + +instance (ASTLore lore, Substitute op) => Substitute (MCOp lore op) where + substituteNames substs (ParOp par_op op) = + ParOp (substituteNames substs <$> par_op) (substituteNames substs op) + substituteNames substs (OtherOp op) = + OtherOp $ substituteNames substs op + +instance (ASTLore lore, Rename op) => Rename (MCOp lore op) where + rename (ParOp par_op op) = ParOp <$> rename par_op <*> rename op + rename (OtherOp op) = OtherOp <$> rename op + +instance (ASTLore lore, FreeIn op) => FreeIn (MCOp lore op) where + freeIn' (ParOp par_op op) = freeIn' par_op <> freeIn' op + freeIn' (OtherOp op) = freeIn' op + +instance (ASTLore lore, IsOp op) => IsOp (MCOp lore op) where + safeOp (ParOp _ op) = safeOp op + safeOp (OtherOp op) = safeOp op + + cheapOp (ParOp _ op) = cheapOp op + cheapOp (OtherOp op) = cheapOp op + +instance TypedOp op => TypedOp (MCOp lore op) where + opType (ParOp _ op) = opType op + opType (OtherOp op) = opType op + +instance + (Aliased lore, AliasedOp op, ASTLore lore) => + AliasedOp (MCOp lore op) + where + opAliases (ParOp _ op) = opAliases op + opAliases (OtherOp op) = opAliases op + + consumedInOp (ParOp _ op) = consumedInOp op + consumedInOp (OtherOp op) = consumedInOp op + +instance + (CanBeAliased (Op lore), CanBeAliased op, ASTLore lore) => + CanBeAliased (MCOp lore op) + where + type OpWithAliases (MCOp lore op) = MCOp (Aliases lore) (OpWithAliases op) + + addOpAliases (ParOp par_op op) = + ParOp (addOpAliases <$> par_op) (addOpAliases op) + addOpAliases (OtherOp op) = + OtherOp $ addOpAliases op + + removeOpAliases (ParOp par_op op) = + ParOp (removeOpAliases <$> par_op) (removeOpAliases op) + removeOpAliases (OtherOp op) = + OtherOp $ removeOpAliases op + +instance + (CanBeWise (Op lore), CanBeWise op, ASTLore lore) => + CanBeWise (MCOp lore op) + where + type OpWithWisdom (MCOp lore op) = MCOp (Wise lore) (OpWithWisdom op) + + removeOpWisdom (ParOp par_op op) = + ParOp (removeOpWisdom <$> par_op) (removeOpWisdom op) + removeOpWisdom (OtherOp op) = + OtherOp $ removeOpWisdom op + +instance (ASTLore lore, ST.IndexOp op) => ST.IndexOp (MCOp lore op) where + indexOp vtable k (ParOp _ op) is = ST.indexOp vtable k op is + indexOp vtable k (OtherOp op) is = ST.indexOp vtable k op is + +instance (PrettyLore lore, Pretty op) => Pretty (MCOp lore op) where + ppr (ParOp Nothing op) = ppr op + ppr (ParOp (Just par_op) op) = + "par" <+> nestedBlock "{" "}" (ppr par_op) + "seq" <+> nestedBlock "{" "}" (ppr op) + ppr (OtherOp op) = ppr op + +instance (OpMetrics (Op lore), OpMetrics op) => OpMetrics (MCOp lore op) where + opMetrics (ParOp par_op op) = opMetrics par_op >> opMetrics op + opMetrics (OtherOp op) = opMetrics op + +typeCheckMCOp :: + TC.Checkable lore => + (op -> TC.TypeM lore ()) -> + MCOp (Aliases lore) op -> + TC.TypeM lore () +typeCheckMCOp _ (ParOp (Just par_op) op) = do + -- It is valid for the same array to be consumed in both par_op and op. + _ <- typeCheckSegOp return par_op `TC.alternative` typeCheckSegOp return op + return () +typeCheckMCOp _ (ParOp Nothing op) = + typeCheckSegOp return op +typeCheckMCOp f (OtherOp op) = f op + +simplifyMCOp :: + ( Engine.SimplifiableLore lore, + BodyDec lore ~ () + ) => + Simplify.SimplifyOp lore op -> + MCOp lore op -> + Engine.SimpleM lore (MCOp (Wise lore) (OpWithWisdom op), Stms (Wise lore)) +simplifyMCOp f (OtherOp op) = do + (op', stms) <- f op + return (OtherOp op', stms) +simplifyMCOp _ (ParOp par_op op) = do + (par_op', par_op_hoisted) <- + case par_op of + Nothing -> return (Nothing, mempty) + Just x -> first Just <$> simplifySegOp x + + (op', op_hoisted) <- simplifySegOp op + + return (ParOp par_op' op', par_op_hoisted <> op_hoisted) diff --git a/src/Futhark/IR/MCMem.hs b/src/Futhark/IR/MCMem.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/MCMem.hs @@ -0,0 +1,84 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} + +module Futhark.IR.MCMem + ( MCMem, + + -- * Simplification + simplifyProg, + + -- * Module re-exports + module Futhark.IR.Mem, + module Futhark.IR.SegOp, + module Futhark.IR.MC.Op, + ) +where + +import Futhark.Analysis.PrimExp.Convert +import Futhark.IR.MC.Op +import Futhark.IR.Mem +import Futhark.IR.Mem.Simplify +import Futhark.IR.SegOp +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 MCMem + +instance Decorations MCMem where + type LetDec MCMem = LetDecMem + type FParamInfo MCMem = FParamMem + type LParamInfo MCMem = LParamMem + type RetType MCMem = RetTypeMem + type BranchType MCMem = BranchTypeMem + type Op MCMem = MemOp (MCOp MCMem ()) + +instance ASTLore MCMem where + expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern + +instance OpReturns MCMem where + opReturns (Alloc _ space) = return [MemMem space] + opReturns (Inner (ParOp _ op)) = segOpReturns op + opReturns (Inner (OtherOp ())) = pure [] + +instance PrettyLore MCMem + +instance TC.CheckableOp MCMem where + checkOp = typeCheckMemoryOp + where + typeCheckMemoryOp (Alloc size _) = + TC.require [Prim int64] size + typeCheckMemoryOp (Inner op) = + typeCheckMCOp pure op + +instance TC.Checkable MCMem where + checkFParamLore = checkMemInfo + checkLParamLore = checkMemInfo + checkLetBoundLore = checkMemInfo + checkRetType = mapM_ (TC.checkExtType . declExtTypeOf) + primFParam name t = return $ Param name (MemPrim t) + matchPattern = matchPatternToExp + matchReturnType = matchFunctionReturnType + matchBranchType = matchBranchReturnType + matchLoopResult = matchLoopResultMem + +instance BinderOps MCMem where + mkExpDecB _ _ = return () + mkBodyB stms res = return $ Body () stms res + mkLetNamesB = mkLetNamesB' () + +instance BinderOps (Engine.Wise MCMem) where + mkExpDecB pat e = return $ Engine.mkWiseExpDec pat () e + mkBodyB stms res = return $ Engine.mkWiseBody () stms res + mkLetNamesB = mkLetNamesB'' + +simplifyProg :: Prog MCMem -> PassM (Prog MCMem) +simplifyProg = simplifyProgGeneric simpleMCMem + +simpleMCMem :: Engine.SimpleOps MCMem +simpleMCMem = + simpleGeneric (const mempty) $ simplifyMCOp $ const $ return ((), mempty) 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 @@ -248,10 +248,10 @@ indexOp _ _ _ _ = Nothing -- | The index function representation used for memory annotations. -type IxFun = IxFun.IxFun (TPrimExp Int32 VName) +type IxFun = IxFun.IxFun (TPrimExp Int64 VName) -- | An index function that may contain existential variables. -type ExtIxFun = IxFun.IxFun (TPrimExp Int32 (Ext VName)) +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 @@ -333,13 +333,13 @@ Engine.SimplifiableLore lore => IxFun -> Engine.SimpleM lore IxFun -simplifyIxFun = traverse $ fmap isInt32 . simplifyPrimExp . untyped +simplifyIxFun = traverse $ fmap isInt64 . simplifyPrimExp . untyped simplifyExtIxFun :: Engine.SimplifiableLore lore => ExtIxFun -> Engine.SimpleM lore ExtIxFun -simplifyExtIxFun = traverse $ fmap isInt32 . simplifyExtPrimExp . untyped +simplifyExtIxFun = traverse $ fmap isInt64 . simplifyExtPrimExp . untyped isStaticIxFun :: ExtIxFun -> Maybe IxFun isStaticIxFun = traverse $ traverse inst @@ -467,22 +467,22 @@ ReturnsInBlock v $ fixExtIxFun i - (primExpFromSubExp int32 (Var v)) + (primExpFromSubExp int64 (Var v)) ixfun fixExt i se (ReturnsNewBlock space j ixfun) = ReturnsNewBlock space j' - (fixExtIxFun i (primExpFromSubExp int32 se) ixfun) + (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 int32 se) ixfun) + ReturnsInBlock mem (fixExtIxFun i (primExpFromSubExp int64 se) ixfun) fixExtIxFun :: Int -> PrimExp VName -> ExtIxFun -> ExtIxFun -fixExtIxFun i e = fmap $ isInt32 . replaceInPrimExp update . untyped +fixExtIxFun i e = fmap $ isInt64 . replaceInPrimExp update . untyped where update (Ext j) t | j > i = LeafExp (Ext $ j - 1) t @@ -490,8 +490,8 @@ | otherwise = LeafExp (Ext j) t update (Free x) t = LeafExp (Free x) t -leafExp :: Int -> TPrimExp Int32 (Ext a) -leafExp i = isInt32 $ LeafExp (Ext i) int32 +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) @@ -657,15 +657,15 @@ -- occurs. getExtMaps :: [(VName, Int)] -> - ( M.Map (Ext VName) (TPrimExp Int32 (Ext VName)), - M.Map (Ext VName) (TPrimExp Int32 (Ext VName)) + ( 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 -> isInt32 $ LeafExp (Ext i) int32) + ( fmap (\i -> isInt64 $ LeafExp (Ext i) int64) . (`lookup` ctx_lst_ids) ) . uncurry (flip (,)) @@ -928,7 +928,7 @@ lookupArraySummary :: (Mem lore, HasScope lore m, Monad m) => VName -> - m (VName, IxFun.IxFun (TPrimExp Int32 VName)) + m (VName, IxFun.IxFun (TPrimExp Int64 VName)) lookupArraySummary name = do summary <- lookupMemInfo name case summary of @@ -943,7 +943,7 @@ MemInfo SubExp u MemBind -> TC.TypeM lore () checkMemInfo _ (MemPrim _) = return () -checkMemInfo _ (MemMem (ScalarSpace d _)) = mapM_ (TC.require [Prim int32]) d +checkMemInfo _ (MemMem (ScalarSpace d _)) = mapM_ (TC.require [Prim int64]) d checkMemInfo _ (MemMem _) = return () checkMemInfo name (MemArray _ shape _ (ArrayIn v ixfun)) = do t <- lookupType v @@ -959,7 +959,7 @@ ++ "." TC.context ("in index function " ++ pretty ixfun) $ do - traverse_ (TC.requirePrimExp int32 . untyped) ixfun + traverse_ (TC.requirePrimExp int64 . untyped) ixfun let ixfun_rank = IxFun.rank ixfun ident_rank = shapeRank shape unless (ixfun_rank == ident_rank) $ @@ -1044,8 +1044,8 @@ IxFun.iota $ map convert $ shapeDims shape | otherwise = return $ MemArray bt shape u Nothing - convert (Ext i) = le32 (Ext i) - convert (Free v) = Free <$> pe32 v + convert (Ext i) = le64 (Ext i) + convert (Free v) = Free <$> pe64 v arrayVarReturns :: (HasScope lore m, Monad m, Mem lore) => @@ -1095,7 +1095,7 @@ Just $ ReturnsInBlock mem $ existentialiseIxFun [] $ - IxFun.reshape ixfun $ map (fmap pe32) newshape + IxFun.reshape ixfun $ map (fmap pe64) newshape ] expReturns (BasicOp (Rearrange perm v)) = do (et, Shape dims, mem, ixfun) <- arrayVarReturns v @@ -1107,7 +1107,7 @@ ] expReturns (BasicOp (Rotate offsets v)) = do (et, Shape dims, mem, ixfun) <- arrayVarReturns v - let offsets' = map pe32 offsets + let offsets' = map pe64 offsets ixfun' = IxFun.rotate ixfun offsets' return [ MemArray et (Shape $ map Free dims) NoUniqueness $ @@ -1176,7 +1176,7 @@ ArrayIn mem $ IxFun.slice ixfun - (map (fmap (isInt32 . primExpFromSubExp int32)) slice) + (map (fmap (isInt64 . primExpFromSubExp int64)) slice) class TypedOp (Op lore) => OpReturns lore where opReturns :: 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 @@ -237,6 +237,7 @@ 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)) = 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 @@ -66,7 +66,7 @@ primOpType (ArrayLit es rt) = pure [arrayOf rt (Shape [n]) NoUniqueness] where - n = intConst Int32 $ toInteger $ length es + n = intConst Int64 $ toInteger $ length es primOpType (BinOp bop _ _) = pure [Prim $ binOpType bop] primOpType (UnOp _ x) = @@ -147,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@, 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 @@ -246,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. @@ -267,7 +267,7 @@ -- 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 _ [] = constant (0 :: Int64) arraysSize i (t : _) = arraySize i t -- | Return the immediate row-type of an array. For @[[int]]@, this 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 @@ -659,13 +659,13 @@ 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 + inttp = Prim int64 lamarrs' = map (`setOuterSize` Var (paramName chunk)) arrargs let acc_len = length accexps let lamrtp = take acc_len $ lambdaReturnType lam @@ -698,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 -- @@ -712,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 @@ -727,12 +727,12 @@ -- 2. forM_ rtsI $ \rtI -> - unless (Prim int32 == rtI) $ - TC.bad $ TC.TypeError "Scatter: Index return type must be i32." + 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 @@ -744,13 +744,13 @@ 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' @@ -775,7 +775,7 @@ -- 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 $ @@ -784,7 +784,7 @@ ++ " 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' 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 @@ -517,7 +517,7 @@ Simplify $ certifying (stmAuxCerts aux1 <> cs) $ letBind pat $ - BasicOp $ Rotate (intConst Int32 0 : rots) arr + BasicOp $ Rotate (intConst Int64 0 : rots) arr mapOpToOp _ _ _ _ = Skip isMapWithOp :: @@ -680,7 +680,7 @@ bindMapParam p a = do a_t <- lookupType a letBindNames [paramName p] $ - BasicOp $ Index a $ fullSlice a_t [DimFix $ constant (0 :: Int32)] + BasicOp $ Index a $ fullSlice a_t [DimFix $ constant (0 :: Int64)] bindArrayResult pe se = letBindNames [patElemName pe] $ BasicOp $ ArrayLit [se] $ rowType $ patElemType pe @@ -705,7 +705,7 @@ partitionChunkedFoldParameters (length nes) (lambdaParams fold_lam) letBindNames [paramName chunk_param] $ - BasicOp $ SubExp $ intConst Int32 1 + BasicOp $ SubExp $ intConst Int64 1 forM_ (zip acc_params nes) $ \(p, ne) -> letBindNames [paramName p] $ BasicOp $ SubExp ne @@ -858,7 +858,7 @@ letExp (baseString arr ++ "_prefix") $ BasicOp $ Index arr $ - fullSlice arr_t [DimSlice (intConst Int32 0) w (intConst Int32 1)] + fullSlice arr_t [DimSlice (intConst Int64 0) w (intConst Int64 1)] return $ Just ( arr', @@ -920,7 +920,7 @@ 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) + 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 @@ -929,7 +929,7 @@ ArrayRearrange _ _ perm -> BasicOp $ Rearrange (0 : map (+ 1) perm) arr ArrayRotate _ _ rots -> - BasicOp $ Rotate (intConst Int32 0 : rots) arr + BasicOp $ Rotate (intConst Int64 0 : rots) arr ArrayVar {} -> BasicOp $ SubExp $ Var arr arr_transformed_t <- lookupType arr_transformed 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 @@ -395,10 +395,10 @@ checkKernelResult (Returns _ what) t = TC.require [t] what checkKernelResult (WriteReturns rws arr res) t = do - mapM_ (TC.require [Prim int32]) rws + mapM_ (TC.require [Prim int64]) rws arr_t <- lookupType arr forM_ res $ \(slice, e) -> do - mapM_ (traverse $ TC.require [Prim int32]) slice + mapM_ (traverse $ TC.require [Prim int64]) slice TC.require [t] e unless (arr_t == t `arrayOfShape` Shape rws) $ TC.bad $ @@ -415,16 +415,16 @@ 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 + 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 int32] dim - TC.require [Prim int32] tile + 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 @@ -514,11 +514,11 @@ -- this 'SegSpace'. scopeOfSegSpace :: SegSpace -> Scope lore scopeOfSegSpace (SegSpace phys space) = - M.fromList $ zip (phys : map fst space) $ repeat $ IndexName Int32 + 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 int32] . snd) 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, @@ -662,10 +662,10 @@ 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 + TC.require [Prim int64] dest_w + TC.require [Prim int64] rf nes' <- mapM TC.checkArg nes - mapM_ (TC.require [Prim int32]) $ shapeDims shape + mapM_ (TC.require [Prim int64]) $ shapeDims shape -- Operator type must match the type of neutral elements. let stripVecDims = stripArray $ shapeRank shape @@ -691,7 +691,7 @@ -- 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 + let bucket_ret_t = replicate (length ops) (Prim int64) ++ concat nes_ts unless (bucket_ret_t == ts) $ TC.bad $ TC.TypeError $ @@ -715,7 +715,7 @@ TC.binding (scopeOfSegSpace space) $ do ne_ts <- forM ops $ \(lam, nes, shape) -> do - mapM_ (TC.require [Prim int32]) $ shapeDims shape + mapM_ (TC.require [Prim int64]) $ shapeDims shape nes' <- mapM TC.checkArg nes -- Operator type must match the type of neutral elements. @@ -1018,7 +1018,7 @@ ST.IndexedArray (stmCerts stm <> cs) arr - (fixSlice (map (fmap isInt32) slice') excess_is) + (fixSlice (map (fmap isInt64) slice') excess_is) in M.insert v idx table | otherwise = table @@ -1119,9 +1119,9 @@ 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 + foldl' f (ST.fromScope $ M.singleton flat $ IndexName Int64) gtids_and_dims where - f vtable (gtid, dim) = ST.insertLoopVar gtid Int32 dim vtable + f vtable (gtid, dim) = ST.insertLoopVar gtid Int64 dim vtable simplifySegBinOp :: Engine.SimplifiableLore lore => @@ -1385,9 +1385,9 @@ map ( \d -> DimSlice - (constant (0 :: Int32)) + (constant (0 :: Int64)) d - (constant (1 :: Int32)) + (constant (1 :: Int64)) ) $ segSpaceDims space index kpe' = 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 @@ -498,15 +498,18 @@ 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)) - End + 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 @@ -523,14 +526,17 @@ 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 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? @@ -539,3 +545,4 @@ where onPart ErrorString {} = Nothing onPart ErrorInt32 {} = Just $ IntType Int32 + onPart ErrorInt64 {} = Just $ IntType Int64 diff --git a/src/Futhark/Internalise.hs b/src/Futhark/Internalise.hs --- a/src/Futhark/Internalise.hs +++ b/src/Futhark/Internalise.hs @@ -106,7 +106,7 @@ return $ Param v $ toDecl v_t Nonunique let free_shape_params = - map (`Param` I.Prim int32) $ + 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' @@ -397,7 +397,7 @@ flat_arr_t <- lookupType flat_arr let new_shape' = reshapeOuter - (map (DimNew . intConst Int32 . toInteger) new_shape) + (map (DimNew . intConst Int64 . toInteger) new_shape) 1 $ I.arrayShape flat_arr_t letSubExp desc $ I.BasicOp $ I.Reshape new_shape' flat_arr @@ -453,25 +453,25 @@ -- Construct an error message in case the range is invalid. let conv = case E.typeOf start of - E.Scalar (E.Prim (E.Unsigned _)) -> asIntS Int32 - _ -> asIntS Int32 - start'_i32 <- conv start' - end'_i32 <- conv end' - maybe_second'_i32 <- traverse conv maybe_second' + 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 "] - ++ [ErrorInt32 start'_i32] - ++ ( case maybe_second'_i32 of + ++ [ErrorInt64 start'_i64] + ++ ( case maybe_second'_i64 of Nothing -> [] - Just second_i32 -> ["..", ErrorInt32 second_i32] + Just second_i64 -> ["..", ErrorInt64 second_i64] ) ++ ( case end of DownToExclusive {} -> ["..>"] ToInclusive {} -> ["..."] UpToExclusive {} -> ["..<"] ) - ++ [ErrorInt32 end'_i32, " is invalid."] + ++ [ErrorInt64 end'_i64, " is invalid."] (it, le_op, lt_op) <- case E.typeOf start of @@ -497,7 +497,7 @@ return (default_step, constant False) step_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum it) step - step_sign_i32 <- asIntS Int32 step_sign + step_sign_i64 <- asIntS Int64 step_sign bounds_invalid_downwards <- letSubExp "bounds_invalid_downwards" $ @@ -514,15 +514,15 @@ distance <- letSubExp "distance" $ I.BasicOp $ I.BinOp (Sub it I.OverflowWrap) start' end' - distance_i32 <- asIntS Int32 distance - return (distance_i32, step_wrong_dir, bounds_invalid_downwards) + 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_i32 <- asIntS Int32 distance - return (distance_i32, step_wrong_dir, bounds_invalid_upwards) + distance_i64 <- asIntS Int64 distance + return (distance_i64, step_wrong_dir, bounds_invalid_upwards) ToInclusive {} -> do downwards <- letSubExp "downwards" $ @@ -548,14 +548,14 @@ (resultBody [distance_downwards_exclusive]) (resultBody [distance_upwards_exclusive]) $ ifCommon [I.Prim $ IntType it] - distance_exclusive_i32 <- asIntS Int32 distance_exclusive + distance_exclusive_i64 <- asIntS Int64 distance_exclusive distance <- letSubExp "distance" $ I.BasicOp $ I.BinOp - (Add Int32 I.OverflowWrap) - distance_exclusive_i32 - (intConst Int32 1) + (Add Int64 I.OverflowWrap) + distance_exclusive_i64 + (intConst Int64 1) return (distance, constant False, bounds_invalid) step_invalid <- @@ -568,15 +568,15 @@ valid <- letSubExp "valid" $ I.BasicOp $ I.UnOp I.Not invalid cs <- assert "range_valid_c" valid errmsg loc - step_i32 <- asIntS Int32 step + step_i64 <- asIntS Int64 step pos_step <- letSubExp "pos_step" $ - I.BasicOp $ I.BinOp (Mul Int32 I.OverflowWrap) step_i32 step_sign_i32 + 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 Int32 I.Unsafe) distance pos_step + 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] @@ -592,7 +592,7 @@ dims <- arrayDims <$> subExpType e' let parts = ["Value of (core language) shape ("] - ++ intersperse ", " (map ErrorInt32 dims) + ++ intersperse ", " (map ErrorInt64 dims) ++ [") cannot match shape of type `"] ++ dt' ++ ["`."] @@ -721,7 +721,7 @@ 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 + I.ForLoop i Int64 w loopvars handleForm mergeinit (E.For i num_iterations) = do num_iterations' <- internaliseExp1 "upper_bound" num_iterations i' <- internaliseIdent i @@ -858,7 +858,7 @@ (ts, constr_map) <- internaliseSumType $ M.map (map E.toStruct) fs es' <- concat <$> mapM (internaliseExp "payload") es - let noExt _ = return $ intConst Int32 0 + let noExt _ = return $ intConst Int64 0 ts' <- instantiateShapes noExt $ map fromDecl ts case M.lookup c constr_map of @@ -1081,7 +1081,7 @@ errorMsg $ ["Index ["] ++ intercalate [", "] parts ++ ["] out of bounds for array of shape ["] - ++ intersperse "][" (map ErrorInt32 $ take (length idxs) dims) + ++ intersperse "][" (map ErrorInt64 $ take (length idxs) dims) ++ ["]."] c <- assert "index_certs" ok msg loc return (idxs', c) @@ -1094,12 +1094,12 @@ (i', _) <- internaliseDimExp "i" i let lowerBound = I.BasicOp $ - I.CmpOp (I.CmpSle I.Int32) (I.constant (0 :: I.Int32)) i' + I.CmpOp (I.CmpSle I.Int64) (I.constant (0 :: I.Int64)) i' upperBound = I.BasicOp $ - I.CmpOp (I.CmpSlt I.Int32) i' w + I.CmpOp (I.CmpSlt I.Int64) i' w ok <- letSubExp "bounds_check" =<< eBinOp I.LogAnd (pure lowerBound) (pure upperBound) - return (I.DimFix i', ok, [ErrorInt32 i']) + return (I.DimFix i', ok, [ErrorInt64 i']) -- Special-case an important common case that otherwise leads to horrible code. internaliseDimIndex @@ -1111,45 +1111,45 @@ ) = do w_minus_1 <- letSubExp "w_minus_1" $ - BasicOp $ I.BinOp (Sub Int32 I.OverflowWrap) w one + BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) w one return - ( I.DimSlice w_minus_1 w $ intConst Int32 (-1), + ( I.DimSlice w_minus_1 w $ intConst Int64 (-1), constant True, mempty ) where - one = constant (1 :: Int32) + 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 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 + 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 int32] + $ ifCommon [I.Prim int64] j_def = letSubExp "j_def" $ I.If backwards (resultBody [negone]) (resultBody [w]) - $ ifCommon [I.Prim int32] + $ 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 Int32 I.OverflowWrap) j' i' + 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 Int32 Unsafe) + (SQuot Int64 Unsafe) ( eBinOp - (Add Int32 I.OverflowWrap) + (Add Int64 I.OverflowWrap) x - (eBinOp (Sub Int32 I.OverflowWrap) y (eSignum $ toExp s')) + (eBinOp (Sub Int64 I.OverflowWrap) y (eSignum $ toExp s')) ) y n <- letSubExp "n" =<< divRounding (toExp j_m_i) (toExp s') @@ -1158,29 +1158,29 @@ -- 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 + empty_slice <- letSubExp "empty_slice" $ I.BasicOp $ I.CmpOp (CmpEq int64) 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 + 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 Int32) zero 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 Int32) i_p_m_t_s 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 Int32) i_p_m_t_s 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 Int32) zero i' - i_lte_j <- letSubExp "i_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) i' j' + 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 Int32) negone j' - j_lte_i <- letSubExp "j_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) j' i' + 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 @@ -1199,25 +1199,25 @@ let parts = case (i, j, s) of (_, _, Just {}) -> - [ maybe "" (const $ ErrorInt32 i') i, + [ maybe "" (const $ ErrorInt64 i') i, ":", - maybe "" (const $ ErrorInt32 j') j, + maybe "" (const $ ErrorInt64 j') j, ":", - ErrorInt32 s' + ErrorInt64 s' ] (_, Just {}, _) -> - [ maybe "" (const $ ErrorInt32 i') i, + [ maybe "" (const $ ErrorInt64 i') i, ":", - ErrorInt32 j' + ErrorInt64 j' ] - ++ maybe mempty (const [":", ErrorInt32 s']) s + ++ maybe mempty (const [":", ErrorInt64 s']) s (_, Nothing, Nothing) -> - [ErrorInt32 i', ":"] + [ErrorInt64 i', ":"] return (I.DimSlice i' n s', ok_or_empty, parts) where - zero = constant (0 :: Int32) - negone = constant (-1 :: Int32) - one = constant (1 :: Int32) + zero = constant (0 :: Int64) + negone = constant (-1 :: Int64) + one = constant (1 :: Int64) internaliseScanOrReduce :: String -> @@ -1276,10 +1276,10 @@ -- reshape return type of bucket function to have same size as neutral element -- (modulo the index) - bucket_param <- newParam "bucket_p" $ I.Prim int32 + 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 int32 : ne_ts + rettype = I.Prim int64 : ne_ts body = mkBody mempty $ map (I.Var . paramName) params body' <- localScope (scopeOfLParams params) $ @@ -1297,7 +1297,7 @@ -- 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 + cmp <- letSubExp "bucket_cmp" $ I.BasicOp $ I.CmpOp (I.CmpEq I.int64) b_w w_img c <- assert "bucket_cert" @@ -1345,7 +1345,7 @@ -- Synthesize neutral elements by applying the fold function -- to an empty chunk. letBindNames [I.paramName chunk_param] $ - I.BasicOp $ I.SubExp $ constant (0 :: Int32) + I.BasicOp $ I.SubExp $ constant (0 :: Int64) forM_ lam_val_params $ \p -> letBindNames [I.paramName p] $ I.BasicOp $ @@ -1410,7 +1410,7 @@ 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 (Signed it)) -> (,it) <$> asIntS Int64 e' _ -> error "internaliseDimExp: bad type" internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName] @@ -1709,13 +1709,13 @@ 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 + 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 Int32 I.OverflowUndef) (constant (1 :: Int32)) x_dims + =<< 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' @@ -1760,7 +1760,7 @@ Just $ \_desc -> do arrs <- internaliseExpToVars "partition_input" arr lam' <- internalisePartitionLambda internaliseLambda k' lam $ map I.Var arrs - uncurry (++) <$> partitionWithSOACS k' lam' 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' @@ -1808,8 +1808,8 @@ dim_ok <- letSubExp "dim_ok" =<< eCmpOp - (I.CmpEq I.int32) - (eBinOp (I.Mul Int32 I.OverflowUndef) (eSubExp n') (eSubExp m')) + (I.CmpEq I.int64) + (eBinOp (I.Mul Int64 I.OverflowUndef) (eSubExp n') (eSubExp m')) (eSubExp old_dim) dim_ok_cert <- assert @@ -1829,7 +1829,7 @@ 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 + 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' @@ -1840,7 +1840,7 @@ let sumdims xsize ysize = letSubExp "conc_tmp" $ I.BasicOp $ - I.BinOp (I.Add I.Int32 I.OverflowUndef) xsize ysize + I.BinOp (I.Add I.Int64 I.OverflowUndef) xsize ysize ressize <- foldM sumdims outer_size =<< mapM (fmap (arraysSize 0) . mapM lookupType) [ys] @@ -1852,7 +1852,7 @@ offset' <- internaliseExp1 "rotation_offset" offset internaliseOperation desc e $ \v -> do r <- I.arrayRank <$> lookupType v - let zero = intConst Int32 0 + let zero = intConst Int64 0 offsets = offset' : replicate (r -1) zero return $ I.Rotate offsets v handleRest [e] "transpose" = Just $ \desc -> @@ -1932,7 +1932,7 @@ cmp <- letSubExp "write_cmp" $ I.BasicOp $ - I.CmpOp (I.CmpEq I.int32) si_w sv_w + I.CmpOp (I.CmpEq I.int64) si_w sv_w c <- assert "write_cert" @@ -2061,9 +2061,9 @@ _ -> error "partitionWithSOACS" add_lam_x_params <- - replicateM k $ I.Param <$> newVName "x" <*> pure (I.Prim int32) + replicateM k $ I.Param <$> newVName "x" <*> pure (I.Prim int64) add_lam_y_params <- - replicateM k $ I.Param <$> newVName "y" <*> pure (I.Prim int32) + 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 $ @@ -2071,16 +2071,16 @@ letSubExp "z" $ I.BasicOp $ I.BinOp - (I.Add Int32 I.OverflowUndef) + (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 int32 + I.lambdaReturnType = replicate k $ I.Prim int64 } - nes = replicate (length increments) $ constant (0 :: Int32) + 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 @@ -2088,17 +2088,17 @@ -- 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) + 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 :: Int32) - is_empty <- letSubExp "is_empty" $ I.BasicOp $ I.CmpOp (CmpEq int32) w $ constant (0 :: Int32) + 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 int32 + ifCommon $ replicate k $ I.Prim int64 -- The total size of all partitions must necessarily be equal to the -- size of the input array. @@ -2111,8 +2111,8 @@ -- 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) + 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) <- @@ -2126,7 +2126,7 @@ I.Lambda { I.lambdaParams = c_param : offset_params ++ value_params, I.lambdaReturnType = - replicate (length arr_ts) (I.Prim int32) + replicate (length arr_ts) (I.Prim int64) ++ map I.rowType arr_ts, I.lambdaBody = mkBody offset_stms $ @@ -2144,7 +2144,7 @@ sizes' <- letSubExp "partition_sizes" $ I.BasicOp $ - I.ArrayLit (map I.Var sizes) $ I.Prim int32 + I.ArrayLit (map I.Var sizes) $ I.Prim int64 return (map I.Var results, [sizes']) where mkOffsetLambdaBody :: @@ -2154,26 +2154,26 @@ [I.LParam] -> InternaliseM SubExp mkOffsetLambdaBody _ _ _ [] = - return $ constant (-1 :: Int32) + return $ constant (-1 :: Int64) mkOffsetLambdaBody sizes c i (p : ps) = do is_this_one <- letSubExp "is_this_one" $ I.BasicOp $ - I.CmpOp (CmpEq int32) c $ - intConst Int32 $ toInteger i + I.CmpOp (CmpEq int64) c $ + intConst Int64 $ toInteger i next_one <- mkOffsetLambdaBody sizes c (i + 1) ps this_one <- letSubExp "this_offset" =<< foldBinOp - (Add Int32 OverflowUndef) - (constant (-1 :: Int32)) + (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 int32] + $ ifCommon [I.Prim int64] typeExpForError :: E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp] typeExpForError (E.TEVar qn _) = @@ -2217,7 +2217,7 @@ d' <- case substs of Just [v] -> return v _ -> return $ I.Var $ E.qualLeaf d - return $ ErrorInt32 d' + return $ ErrorInt64 d' dimExpForError (DimExpConst d _) = return $ ErrorString $ pretty d dimExpForError DimExpAny = return "" 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 @@ -47,7 +47,7 @@ let addShape name = case M.lookup name mapping of Just se -> se - _ -> intConst Int32 0 -- FIXME: we only need this because + _ -> intConst Int64 0 -- FIXME: we only need this because -- the defunctionaliser throws away -- sizes. return $ map addShape shapes @@ -156,4 +156,4 @@ | otherwise = return v where checkDim desired has = - letSubExp "dim_match" $ BasicOp $ CmpOp (CmpEq int32) 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 @@ -32,7 +32,7 @@ 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] + 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) $ @@ -49,7 +49,7 @@ pat_idents <- flattenPattern pat pat_ts <- internaliseLoopParamType (E.patternStructType pat) - let shape_params = [I.Param v $ I.Prim I.int32 | E.TypeParamDim v _ <- tparams] + 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 -> 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 @@ -126,7 +126,7 @@ | baseTag x <= maxIntrinsicTag -> return IntrinsicSV | otherwise -> -- Anything not in scope is going to be an -- existential size. - return $ Dynamic $ Scalar $ Prim $ Signed Int32 + return $ Dynamic $ Scalar $ Prim $ Signed Int64 | otherwise -> error $ "Variable " ++ pretty x ++ " at " @@ -842,7 +842,7 @@ ++ "." envFromDimNames :: [VName] -> Env -envFromDimNames = M.fromList . flip zip (repeat $ Dynamic $ Scalar $ Prim $ Signed Int32) +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. 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 @@ -44,12 +44,12 @@ 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 @@ -96,11 +96,11 @@ 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 @@ -126,19 +126,19 @@ lambdaWithIncrement body return $ I.Lambda params body' rettype where - rettype = replicate (k + 2) $ I.Prim int32 + rettype = replicate (k + 2) $ I.Prim int64 result i = map constant $ - (fromIntegral i :: Int32) : - (replicate i 0 ++ [1 :: Int32] ++ replicate (k - i) 0) + 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 $ - intConst Int32 $ toInteger i + CmpOp (CmpEq int64) eq_class $ + intConst Int64 $ toInteger i fmap (map I.Var) . letTupExp "part_res" =<< eIf (eSubExp is_i) 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 @@ -46,8 +46,8 @@ 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 @@ -223,7 +223,7 @@ f size_arg (Info (Observe, Nothing)) - (Info (foldFunType (replicate i i32) (fromStruct t)), Info []) + (Info (foldFunType (replicate i i64) (fromStruct t)), Info []) loc ) @@ -236,7 +236,7 @@ (qualName fname') ( Info ( foldFunType - (map (const i32) size_args) + (map (const i64) size_args) (fromStruct t') ) ) @@ -593,7 +593,7 @@ noticeDims :: TypeBase (DimDecl VName) as -> MonoM () noticeDims = mapM_ notice . nestedDims where - notice (NamedDim v) = void $ transformFName mempty v i32 + notice (NamedDim v) = void $ transformFName mempty v i64 notice _ = return () -- Convert a collection of 'ValBind's to a nested sequence of let-bound, @@ -670,9 +670,9 @@ tparamArg dinst tp = case M.lookup (typeParamName tp) dinst of Just (NamedDim d) -> - Just $ Var d (Info i32) mempty + Just $ Var d (Info i64) mempty Just (ConstDim x) -> - Just $ Literal (SignedValue $ Int32Value $ fromIntegral x) mempty + Just $ Literal (SignedValue $ Int64Value $ fromIntegral x) mempty _ -> Nothing @@ -768,7 +768,7 @@ 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 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 @@ -102,7 +102,7 @@ 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 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 @@ -47,6 +47,7 @@ removeStmAliases, ) import qualified Futhark.IR.Kernels.Kernel as Kernel +import qualified Futhark.IR.MC as MC import qualified Futhark.IR.Mem as Memory import Futhark.IR.Prop.Aliases import qualified Futhark.IR.SOACS.SOAC as SOAC @@ -281,6 +282,19 @@ 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 op + ) => + CSEInOp (MC.MCOp lore op) + where + cseInOp (MC.ParOp par_op op) = + MC.ParOp <$> traverse cseInOp par_op <*> cseInOp op + cseInOp (MC.OtherOp op) = + MC.OtherOp <$> cseInOp op instance (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => 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 @@ -22,7 +22,7 @@ -- 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 (doubleBufferKernels, doubleBufferMC) where import Control.Monad.Reader import Control.Monad.State @@ -31,17 +31,25 @@ import qualified Data.Map.Strict as M import Data.Maybe import Futhark.Construct -import Futhark.IR -import Futhark.IR.KernelsMem +import Futhark.IR.KernelsMem as Kernels +import Futhark.IR.MCMem as MC import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.Pass import Futhark.Pass.ExplicitAllocations (arraySizeInBytesExp) import Futhark.Pass.ExplicitAllocations.Kernels () import Futhark.Util (maybeHead) +-- | The pass for GPU kernels. +doubleBufferKernels :: Pass KernelsMem KernelsMem +doubleBufferKernels = doubleBuffer optimiseKernelsOp + +-- | The pass for multicore +doubleBufferMC :: Pass MCMem MCMem +doubleBufferMC = doubleBuffer optimiseMCOp + -- | The double buffering pass definition. -doubleBuffer :: Pass KernelsMem KernelsMem -doubleBuffer = +doubleBuffer :: Mem lore => OptimiseOp lore -> Pass lore lore +doubleBuffer onOp = Pass { passName = "Double buffer", passDescription = "Perform double buffering for merge parameters of sequential loops.", @@ -55,36 +63,54 @@ fmap stmsFromList $ optimiseStms $ stmsToList stms in runState (runReaderT m env) src - env = Env mempty doNotTouchLoop + env = Env mempty doNotTouchLoop onOp doNotTouchLoop ctx val body = return (mempty, ctx, val, body) -data Env = Env - { envScope :: Scope KernelsMem, - envOptimiseLoop :: OptimiseLoop +type OptimiseLoop lore = + [(FParam lore, SubExp)] -> + [(FParam lore, SubExp)] -> + Body lore -> + DoubleBufferM + lore + ( [Stm lore], + [(FParam lore, SubExp)], + [(FParam lore, SubExp)], + Body lore + ) + +type OptimiseOp lore = + Op lore -> DoubleBufferM lore (Op lore) + +data Env lore = Env + { envScope :: Scope lore, + envOptimiseLoop :: OptimiseLoop lore, + envOptimiseOp :: OptimiseOp lore } -newtype DoubleBufferM a = DoubleBufferM {runDoubleBufferM :: ReaderT Env (State VNameSource) a} - deriving (Functor, Applicative, Monad, MonadReader Env, MonadFreshNames) +newtype DoubleBufferM lore a = DoubleBufferM + { runDoubleBufferM :: ReaderT (Env lore) (State VNameSource) a + } + deriving (Functor, Applicative, Monad, MonadReader (Env lore), MonadFreshNames) -instance HasScope KernelsMem DoubleBufferM where +instance ASTLore lore => HasScope lore (DoubleBufferM lore) where askScope = asks envScope -instance LocalScope KernelsMem DoubleBufferM where +instance ASTLore lore => LocalScope lore (DoubleBufferM lore) where localScope scope = local $ \env -> env {envScope = envScope env <> scope} -optimiseBody :: Body KernelsMem -> DoubleBufferM (Body KernelsMem) +optimiseBody :: ASTLore lore => Body lore -> DoubleBufferM lore (Body lore) optimiseBody body = do bnds' <- optimiseStms $ stmsToList $ bodyStms body return $ body {bodyStms = stmsFromList bnds'} -optimiseStms :: [Stm KernelsMem] -> DoubleBufferM [Stm KernelsMem] +optimiseStms :: ASTLore lore => [Stm lore] -> DoubleBufferM lore [Stm lore] optimiseStms [] = return [] 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 :: forall lore. ASTLore lore => Stm lore -> DoubleBufferM lore [Stm lore] optimiseStm (Let pat aux (DoLoop ctx val form body)) = do body' <- localScope (scopeOf form <> scopeOfFParams (map fst $ ctx ++ val)) $ @@ -92,20 +118,19 @@ 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 +optimiseStm (Let pat aux e) = do + onOp <- asks envOptimiseOp + pure . Let pat aux <$> mapExpM (optimise onOp) e where - optimise = + optimise onOp = identityMapper { mapOnBody = \_ x -> - optimiseBody x :: DoubleBufferM (Body KernelsMem), - mapOnOp = optimiseOp + optimiseBody x :: DoubleBufferM lore (Body lore), + mapOnOp = onOp } -optimiseOp :: - Op KernelsMem -> - DoubleBufferM (Op KernelsMem) -optimiseOp (Inner (SegOp op)) = +optimiseKernelsOp :: OptimiseOp KernelsMem +optimiseKernelsOp (Inner (SegOp op)) = local inSegOp $ Inner . SegOp <$> mapSegOpM mapper op where mapper = @@ -114,32 +139,51 @@ mapOnSegOpBody = optimiseKernelBody } inSegOp env = env {envOptimiseLoop = optimiseLoop} -optimiseOp op = return op +optimiseKernelsOp op = return op +optimiseMCOp :: OptimiseOp MCMem +optimiseMCOp (Inner (ParOp par_op op)) = + local inSegOp $ + Inner + <$> (ParOp <$> traverse (mapSegOpM mapper) par_op <*> mapSegOpM mapper op) + where + mapper = + identitySegOpMapper + { mapOnSegOpLambda = optimiseLambda, + mapOnSegOpBody = optimiseKernelBody + } + inSegOp env = env {envOptimiseLoop = optimiseLoop} +optimiseMCOp op = return op + optimiseKernelBody :: - KernelBody KernelsMem -> - DoubleBufferM (KernelBody KernelsMem) + ASTLore lore => + KernelBody lore -> + DoubleBufferM lore (KernelBody lore) optimiseKernelBody kbody = do stms' <- optimiseStms $ stmsToList $ kernelBodyStms kbody return $ kbody {kernelBodyStms = stmsFromList stms'} -optimiseLambda :: Lambda KernelsMem -> DoubleBufferM (Lambda KernelsMem) +optimiseLambda :: + ASTLore lore => + Lambda lore -> + DoubleBufferM lore (Lambda lore) optimiseLambda lam = do body <- localScope (castScope $ scopeOf lam) $ optimiseBody $ lambdaBody lam 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 - ) +type Constraints lore = + ( ASTLore lore, + FParamInfo lore ~ FParamMem, + LParamInfo lore ~ LParamMem, + RetType lore ~ RetTypeMem, + LetDec lore ~ LetDecMem, + BranchType lore ~ BranchTypeMem, + ExpDec lore ~ (), + BodyDec lore ~ (), + OpReturns lore + ) -optimiseLoop :: OptimiseLoop +optimiseLoop :: (Constraints lore, Op lore ~ MemOp inner, BinderOps lore) => OptimiseLoop lore optimiseLoop ctx val body = do -- We start out by figuring out which of the merge variables should -- be double-buffered. @@ -171,8 +215,8 @@ doubleBufferMergeParams :: MonadFreshNames m => - [(FParam KernelsMem, SubExp)] -> - [FParam KernelsMem] -> + [(Param FParamMem, SubExp)] -> + [Param FParamMem] -> Names -> m [DoubleBuffer] doubleBufferMergeParams ctx_and_res val_params bound_in_loop = @@ -229,9 +273,10 @@ _ -> return NoBuffer allocStms :: - [(FParam KernelsMem, SubExp)] -> + (Constraints lore, Op lore ~ MemOp inner, BinderOps lore) => + [(FParam lore, SubExp)] -> [DoubleBuffer] -> - DoubleBufferM ([(FParam KernelsMem, SubExp)], [Stm KernelsMem]) + DoubleBufferM lore ([(FParam lore, SubExp)], [Stm lore]) allocStms merge = runWriterT . zipWithM allocation merge where allocation m@(Param pname _, _) (BufferAlloc name size space b) = do @@ -265,11 +310,12 @@ return (f, se) doubleBufferResult :: - [FParam KernelsMem] -> + (Constraints lore) => + [FParam lore] -> [DoubleBuffer] -> - Body KernelsMem -> - Body KernelsMem -doubleBufferResult valparams buffered (Body () bnds res) = + Body lore -> + Body lore +doubleBufferResult valparams buffered (Body _ bnds res) = let (ctx_res, val_res) = splitAt (length res - length valparams) res (copybnds, val_res') = unzip $ zipWith3 buffer valparams buffered val_res 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 @@ -690,7 +690,7 @@ (loop_params, loop_arrs) = unzip loop_vars chunk_size <- newVName "chunk_size" offset <- newVName "offset" - let chunk_param = Param chunk_size $ Prim int32 + let chunk_param = Param chunk_size $ Prim int64 offset_param = Param offset $ Prim $ IntType it acc_params <- forM merge_params $ \p -> @@ -719,7 +719,7 @@ [ 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) + BasicOp $ BinOp (Add Int64 OverflowUndef) (Futhark.Var offset) (Futhark.Var chunk_size) ] let lam = Lambda @@ -733,7 +733,7 @@ -- 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 + let discard_pe = PatElem discard $ Prim int64 fusionGatherStms fres @@ -805,8 +805,8 @@ 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 + ForLoop i it _ loopvars -> + Ident i (Prim (IntType it)) : map (paramIdent . fst) loopvars WhileLoop {} -> [] new_res <- 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 @@ -442,7 +442,7 @@ { lambdaParams = lambdaParams lam_c ++ lambdaParams lam_p, lambdaBody = body', lambdaReturnType = - replicate (c_num_buckets + p_num_buckets) (Prim int32) + replicate (c_num_buckets + p_num_buckets) (Prim int64) ++ drop c_num_buckets (lambdaReturnType lam_c) ++ drop p_num_buckets (lambdaReturnType lam_p) } @@ -844,7 +844,7 @@ SOAC.Reshape cs shape SOAC.:< ots' <- SOAC.viewf ots, all primType $ lambdaReturnType maplam = do let mapw' = case reverse $ newDims shape of - [] -> intConst Int32 0 + [] -> intConst Int64 0 d : _ -> d inputs' = map (SOAC.addTransform $ SOAC.ReshapeOuter cs shape) inps inputTypes = map SOAC.inputType inputs' 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 @@ -64,6 +64,7 @@ module Futhark.Optimise.InPlaceLowering ( inPlaceLoweringKernels, inPlaceLoweringSeq, + inPlaceLoweringMC, ) where @@ -73,6 +74,7 @@ import Futhark.Binder import Futhark.IR.Aliases import Futhark.IR.Kernels +import Futhark.IR.MC import Futhark.IR.Seq (Seq) import Futhark.Optimise.InPlaceLowering.LowerIntoStm import Futhark.Pass @@ -86,6 +88,10 @@ inPlaceLoweringSeq = inPlaceLowering pure lowerUpdate -- | Apply the in-place lowering optimisation to the given program. +inPlaceLoweringMC :: Pass MC MC +inPlaceLoweringMC = inPlaceLowering onMCOp lowerUpdate + +-- | Apply the in-place lowering optimisation to the given program. inPlaceLowering :: Constraints lore => OnOp lore -> @@ -192,8 +198,11 @@ { mapOnBody = const optimiseBody } -onKernelOp :: OnOp Kernels -onKernelOp (SegOp op) = +onSegOp :: + (Bindable lore, CanBeAliased (Op lore)) => + SegOp lvl (Aliases lore) -> + ForwardingM lore (SegOp lvl (Aliases lore)) +onSegOp op = bindingScope (scopeOfSegSpace (segSpace op)) $ do let mapper = identitySegOpMapper {mapOnSegOpBody = onKernelBody} onKernelBody kbody = do @@ -202,7 +211,14 @@ optimiseStms (stmsToList (kernelBodyStms kbody)) $ mapM_ seenVar $ namesToList $ freeIn $ kernelBodyResult kbody return kbody {kernelBodyStms = stmsFromList stms} - SegOp <$> mapSegOpM mapper op + mapSegOpM mapper op + +onMCOp :: OnOp MC +onMCOp (ParOp par_op op) = ParOp <$> traverse onSegOp par_op <*> onSegOp op +onMCOp op = return op + +onKernelOp :: OnOp Kernels +onKernelOp (SegOp op) = SegOp <$> onSegOp op onKernelOp op = return op data Entry lore = Entry 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 @@ -62,14 +62,14 @@ (patternNames pat) inputsize mempty - Int32 + 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) + BasicOp $ CmpOp (CmpEq int64) inputsize (intConst Int64 0) letBind pat =<< ( If (Var isEmpty) <$> resultBodyM accs @@ -183,7 +183,7 @@ | v `nameIn` nonFree = M.lookup v knownBnds asFreeSubExp se = Just se - properIntSize Int32 = Just $ return size + properIntSize Int64 = Just $ return size properIntSize t = Just $ letSubExp "converted_size" $ 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 @@ -340,7 +340,7 @@ letExp "for_in_partial" $ BasicOp $ Index arr' $ - DimSlice (intConst Int32 0) w (intConst Int32 1) : slice' + 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 @@ -355,16 +355,15 @@ notIndex _ = True simplifyLoopVariables _ _ _ _ = Skip --- If a for-loop with no loop variables has a counter of a large --- integer type, 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 it n [], body) - | Just (n', it', cs) <- smallerType, - it' < it = +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' [] @@ -409,7 +408,7 @@ letBindNames [paramName p] $ BasicOp $ Index arr $ - DimFix (intConst Int32 i) : fullSlice (paramType p) [] + 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. @@ -753,7 +752,7 @@ `add` primExpFromSubExp (IntType to_it) i_offset' i_stride'' <- letSubExp "iota_offset" $ - BasicOp $ BinOp (Mul Int32 OverflowWrap) s i_stride' + BasicOp $ BinOp (Mul Int64 OverflowWrap) s i_stride' fmap (SubExpResult cs) $ letSubExp "slice_iota" $ BasicOp $ Iota i_n i_offset'' i_stride'' to_it @@ -763,8 +762,8 @@ | 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) + 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) = @@ -791,7 +790,7 @@ return $ IndexResult cs arr $ ds_inds' ++ rest_inds where index DimFix {} = Nothing - index (DimSlice _ n s) = Just (n, DimSlice (constant (0 :: Int32)) n s) + 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 @@ -836,7 +835,7 @@ 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 + 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 @@ -844,9 +843,9 @@ 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 + cmp <- letSubExp "index_concat_cmp" $ BasicOp $ CmpOp (CmpSle Int64) start i (thisres, thisbnds) <- collectStms $ do - i' <- letSubExp "index_concat_i" $ BasicOp $ BinOp (Sub Int32 OverflowWrap) i start + 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' @@ -856,7 +855,7 @@ IfDec [primBodyType res_t] IfNormal SubExpResult cs <$> mkBranch xs_and_starts Just (ArrayLit ses _, cs) - | DimFix (Constant (IntValue (Int32Value i))) : inds' <- inds, + | DimFix (Constant (IntValue (Int64Value i))) : inds' <- inds, Just se <- maybeNth i ses -> case inds' of [] -> Just $ pure $ SubExpResult cs se @@ -871,7 +870,7 @@ Just $ pure $ IndexResult mempty idd $ - DimFix (constant (0 :: Int32)) : inds' + DimFix (constant (0 :: Int64)) : inds' _ -> Nothing where defOf v = do @@ -920,7 +919,7 @@ 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 pe32 ws) + 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 @@ -1241,7 +1240,7 @@ 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 = + Just (ST.Indexed cs e) <- ST.index v [intConst Int64 0] vtable = Simplify $ do e' <- toSubExp "update_elem" e auxing aux $ @@ -1330,7 +1329,7 @@ ruleBasicOp _ pat _ (ArrayLit (se : ses) _) | all (== se) ses = Simplify $ - let n = constant (fromIntegral (length ses) + 1 :: Int32) + 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, @@ -1347,9 +1346,9 @@ oldshape <- arrayDims <$> lookupType idd2 let new_inds = reshapeIndex - (map pe32 oldshape) - (map pe32 $ newDims newshape) - (map pe32 inds) + (map pe64 oldshape) + (map pe64 $ newDims newshape) + (map pe64 inds) new_inds' <- mapM (toSubExp "new_index") new_inds certifying idd_cs $ @@ -1400,7 +1399,7 @@ | 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 + 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 @@ -1415,7 +1414,7 @@ auxing aux $ letBind pat $ BasicOp $ Rotate offsets v2 where - add x y = letSubExp "offset" $ BasicOp $ BinOp (Add Int32 OverflowWrap) x y + 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 @@ -1430,8 +1429,8 @@ 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)] + 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 $ @@ -1439,7 +1438,7 @@ -- 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, + | Constant (IntValue (Int64Value 0)) <- x, Var v <- y, Just _ <- ST.lookupLoopVar v vtable = Simplify $ auxing aux $ letBind pat $ BasicOp $ SubExp $ constant True 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,3 +1,4 @@ +{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE TypeFamilies #-} @@ -42,9 +43,10 @@ -- 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 +module Futhark.Optimise.Sink (sinkKernels, sinkMC) where import Control.Monad.State +import Data.Bifunctor import Data.List (foldl') import qualified Data.Map as M import qualified Data.Set as S @@ -52,20 +54,27 @@ import qualified Futhark.Analysis.SymbolTable as ST import Futhark.IR.Aliases import Futhark.IR.Kernels +import Futhark.IR.MC import Futhark.Pass -type SinkLore = Aliases Kernels - -type SymbolTable = ST.SymbolTable SinkLore +type SymbolTable lore = ST.SymbolTable lore -type Sinking = M.Map VName (Stm SinkLore) +type Sinking lore = M.Map VName (Stm lore) type Sunk = S.Set VName +type Sinker lore a = SymbolTable lore -> Sinking lore -> a -> (a, Sunk) + +type Constraints lore = + ( ASTLore lore, + Aliased lore, + ST.IndexOp (Op lore) + ) + -- | Given a statement, compute how often each of its free variables -- are used. Not accurate: what we care about are only 1, and greater -- than 1. -multiplicity :: Stm SinkLore -> M.Map VName Int +multiplicity :: Constraints lore => Stm lore -> M.Map VName Int multiplicity stm = case stmExp stm of If cond tbranch fbranch _ -> @@ -78,12 +87,11 @@ comb = M.unionWith (+) 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 + Constraints lore => + Sinker lore (Op lore) -> + Sinker lore (Body lore) +optimiseBranch onOp vtable sinking (Body dec stms res) = + let (stms', stms_sunk) = optimiseStms onOp vtable sinking' stms $ freeIn res in ( Body dec (sunk_stms <> stms') res, sunk <> stms_sunk ) @@ -97,12 +105,14 @@ sunk = S.fromList $ concatMap (patternNames . stmPattern) sunk_stms optimiseStms :: - SymbolTable -> - Sinking -> - Stms SinkLore -> + Constraints lore => + Sinker lore (Op lore) -> + SymbolTable lore -> + Sinking lore -> + Stms lore -> Names -> - (Stms SinkLore, Sunk) -optimiseStms init_vtable init_sinking all_stms free_in_res = + (Stms lore, Sunk) +optimiseStms onOp 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) @@ -125,17 +135,16 @@ 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 + let (tbranch', tsunk) = optimiseBranch onOp vtable sinking tbranch + (fbranch', fsunk) = optimiseBranch onOp 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 + | Op op <- stmExp stm = + let (op', op_sunk) = onOp vtable sinking op (stms', stms_sunk) = optimiseStms' vtable' sinking stms - (op', op_sunk) = runState (mapSegOpM (opMapper scope) op) mempty - in ( stm {stmExp = Op (SegOp op')} : stms', + in ( stm {stmExp = Op op'} : stms', stms_sunk <> op_sunk ) | otherwise = @@ -150,49 +159,66 @@ identityMapper { mapOnBody = \scope body -> do let (body', sunk) = - optimiseBody (ST.fromScope scope <> vtable) sinking body + optimiseBody + onOp + (ST.fromScope scope <> vtable) + sinking + body modify (<> sunk) return body' } - 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 - 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) + Constraints lore => + Sinker lore (Op lore) -> + Sinker lore (Body lore) +optimiseBody onOp vtable sinking (Body attr stms res) = + let (stms', sunk) = optimiseStms onOp vtable sinking stms $ freeIn res + in (Body attr stms' res, 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) + Constraints lore => + Sinker lore (Op lore) -> + Sinker lore (KernelBody lore) +optimiseKernelBody onOp vtable sinking (KernelBody attr stms res) = + let (stms', sunk) = optimiseStms onOp vtable sinking stms $ freeIn res + in (KernelBody attr stms' res, sunk) --- | The pass definition. -sink :: Pass Kernels Kernels -sink = +optimiseSegOp :: + Constraints lore => + Sinker lore (Op lore) -> + Sinker lore (SegOp lvl lore) +optimiseSegOp onOp vtable sinking op = + let scope = scopeOfSegSpace $ segSpace op + in runState (mapSegOpM (opMapper scope) op) mempty + where + opMapper scope = + identitySegOpMapper + { mapOnSegOpLambda = \lam -> do + let (body, sunk) = + optimiseBody onOp op_vtable sinking $ + lambdaBody lam + modify (<> sunk) + return lam {lambdaBody = body}, + mapOnSegOpBody = \body -> do + let (body', sunk) = + optimiseKernelBody onOp op_vtable sinking body + modify (<> sunk) + return body' + } + where + op_vtable = ST.fromScope scope <> vtable + +type SinkLore lore = Aliases lore + +sink :: + ( ASTLore lore, + CanBeAliased (Op lore), + ST.IndexOp (OpWithAliases (Op lore)) + ) => + Sinker (SinkLore lore) (Op (SinkLore lore)) -> + Pass lore lore +sink onOp = Pass "sink" "move memory loads closer to their uses" $ fmap removeProgAliases . intraproceduralTransformationWithConsts onConsts onFun @@ -200,11 +226,35 @@ where onFun _ fd = do let vtable = ST.insertFParams (funDefParams fd) mempty - (body, _) = optimiseBody vtable mempty $ funDefBody fd + (body, _) = optimiseBody onOp vtable mempty $ funDefBody fd return fd {funDefBody = body} onConsts consts = pure $ fst $ - optimiseStms mempty mempty consts $ + optimiseStms onOp mempty mempty consts $ namesFromList $ M.keys $ scopeOf consts + +-- | Sinking in GPU kernels. +sinkKernels :: Pass Kernels Kernels +sinkKernels = sink onHostOp + where + onHostOp :: Sinker (SinkLore Kernels) (Op (SinkLore Kernels)) + onHostOp vtable sinking (SegOp op) = + first SegOp $ optimiseSegOp onHostOp vtable sinking op + onHostOp _ _ op = (op, mempty) + +-- | Sinking for multicore. +sinkMC :: Pass MC MC +sinkMC = sink onHostOp + where + onHostOp :: Sinker (SinkLore MC) (Op (SinkLore MC)) + onHostOp vtable sinking (ParOp par_op op) = + let (par_op', par_sunk) = + maybe + (Nothing, mempty) + (first Just . optimiseSegOp onHostOp vtable sinking) + par_op + (op', sunk) = optimiseSegOp onHostOp vtable sinking op + in (ParOp par_op' op', par_sunk <> sunk) + onHostOp _ _ op = (op, mempty) 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 @@ -614,7 +614,7 @@ <*> pure (Var mergeinit) tile_id <- newVName "tile_id" - let loopform = ForLoop tile_id Int32 num_whole_tiles [] + let loopform = ForLoop tile_id Int64 num_whole_tiles [] loopbody <- renameBody <=< runBodyBinder $ inScopeOf loopform $ localScope (scopeOfFParams $ map fst merge) $ do @@ -664,7 +664,7 @@ 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) + let unit_dims = replicate (length dims_on_top) (intConst Int64 1) arr' <- if null dims_on_top then return arr @@ -697,9 +697,6 @@ SegOp $ SegMap lvl space ts $ KernelBody () stms' $ map (Returns manifest) res' -v32 :: VName -> TPrimExp Int32 VName -v32 v = TPrimExp $ LeafExp v int32 - reconstructGtids1D :: Count GroupSize SubExp -> VName -> @@ -708,7 +705,7 @@ Binder Kernels () reconstructGtids1D group_size gtid gid ltid = letBindNames [gtid] - =<< toExp (v32 gid * pe32 (unCount group_size) + v32 ltid) + =<< toExp (le64 gid * pe64 (unCount group_size) + le64 ltid) readTile1D :: SubExp -> @@ -734,7 +731,7 @@ segMap1D "full_tile" (SegThread num_groups group_size SegNoVirt) ResultNoSimplify $ \ltid -> do j <- letSubExp "j" - =<< toExp (pe32 tile_id * pe32 tile_size + v32 ltid) + =<< toExp (pe64 tile_id * pe64 tile_size + le64 ltid) reconstructGtids1D group_size gtid gid ltid addPrivStms [DimFix $ Var ltid] privstms @@ -752,7 +749,7 @@ TilePartial -> letTupExp "pre" =<< eIf - (toExp $ pe32 j .<. pe32 w) + (toExp $ pe64 j .<. pe64 w) (resultBody <$> mapM (fmap Var . readTileElem) arrs) (eBody $ map eBlank tile_ts) TileFull -> @@ -801,7 +798,7 @@ fmap (map Var) $ letTupExp "acc" =<< eIf - (toExp $ v32 gtid .<. pe32 kdim) + (toExp $ le64 gtid .<. pe64 kdim) (eBody [pure $ Op $ OtherOp $ Screma tile_size form' tile]) (resultBodyM thread_accs) @@ -840,11 +837,11 @@ -- the whole tiles. residual_input <- letSubExp "residual_input" $ - BasicOp $ BinOp (SRem Int32 Unsafe) w tile_size + BasicOp $ BinOp (SRem Int64 Unsafe) w tile_size letTupExp "acc_after_residual" =<< eIf - (toExp $ pe32 residual_input .==. 0) + (toExp $ pe64 residual_input .==. 0) (resultBodyM $ map Var accs) (nonemptyTile residual_input) where @@ -867,7 +864,7 @@ BasicOp $ Index tile - [DimSlice (intConst Int32 0) residual_input (intConst Int32 1)] + [DimSlice (intConst Int64 0) residual_input (intConst Int64 1)] -- Now each thread performs a traversal of the tile and -- updates its accumulator. @@ -901,16 +898,16 @@ else do group_size <- letSubExp "computed_group_size" $ - BasicOp $ BinOp (SMin Int32) (unCount (segGroupSize initial_lvl)) kdim + 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 Int32 Unsafe) kdim group_size + BasicOp $ BinOp (SDivUp Int64 Unsafe) kdim group_size num_groups <- letSubExp "computed_num_groups" - =<< foldBinOp (Mul Int32 OverflowUndef) ldim (map snd dims_on_top) + =<< foldBinOp (Mul Int64 OverflowUndef) ldim (map snd dims_on_top) return ( SegGroup (Count num_groups) (Count group_size) SegNoVirt, @@ -922,8 +919,8 @@ Tiling { tilingSegMap = \desc lvl' manifest f -> segMap1D desc lvl' manifest $ \ltid -> do letBindNames [gtid] - =<< toExp (v32 gid * pe32 tile_size + v32 ltid) - f (untyped $ v32 gtid .<. pe32 kdim) [DimFix $ Var ltid], + =<< 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 = @@ -934,7 +931,7 @@ tilingTileShape = Shape [tile_size], tilingNumWholeTiles = letSubExp "num_whole_tiles" $ - BasicOp $ BinOp (SQuot Int32 Unsafe) w tile_size, + BasicOp $ BinOp (SQuot Int64 Unsafe) w tile_size, tilingLevel = lvl, tilingSpace = space } @@ -990,9 +987,9 @@ 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 (v32 gid_x * pe32 tile_size + v32 ltid_x) + =<< toExp (le64 gid_x * pe64 tile_size + le64 ltid_x) letBindNames [gtid_y] - =<< toExp (v32 gid_y * pe32 tile_size + v32 ltid_y) + =<< toExp (le64 gid_y * pe64 tile_size + le64 ltid_y) readTile2D :: (SubExp, SubExp) -> @@ -1015,10 +1012,10 @@ $ \(ltid_x, ltid_y) -> do i <- letSubExp "i" - =<< toExp (pe32 tile_id * pe32 tile_size + v32 ltid_x) + =<< toExp (pe64 tile_id * pe64 tile_size + le64 ltid_x) j <- letSubExp "j" - =<< toExp (pe32 tile_id * pe32 tile_size + v32 ltid_y) + =<< 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 @@ -1041,11 +1038,11 @@ last $ rearrangeShape perm - [ isInt32 (LeafExp gtid_y int32) .<. pe32 kdim_y, - isInt32 (LeafExp gtid_x int32) .<. pe32 kdim_x + [ le64 gtid_y .<. pe64 kdim_y, + le64 gtid_x .<. pe64 kdim_x ] eIf - (toExp $ pe32 idx .<. pe32 w .&&. othercheck) + (toExp $ pe64 idx .<. pe64 w .&&. othercheck) (eBody [return $ BasicOp $ Index arr [DimFix idx]]) (eBody [eBlank tile_t]) @@ -1116,9 +1113,7 @@ fmap (map Var) $ letTupExp "acc" =<< eIf - ( toExp $ - isInt32 (LeafExp gtid_x int32) .<. pe32 kdim_x - .&&. isInt32 (LeafExp gtid_y int32) .<. pe32 kdim_y + ( 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) @@ -1158,11 +1153,11 @@ -- the whole tiles. residual_input <- letSubExp "residual_input" $ - BasicOp $ BinOp (SRem Int32 Unsafe) w tile_size + BasicOp $ BinOp (SRem Int64 Unsafe) w tile_size letTupExp "acc_after_residual" =<< eIf - (toExp $ pe32 residual_input .==. 0) + (toExp $ pe64 residual_input .==. 0) (resultBodyM $ map Var accs) (nonemptyTile residual_input) where @@ -1187,8 +1182,8 @@ BasicOp $ Index tile - [ DimSlice (intConst Int32 0) residual_input (intConst Int32 1), - DimSlice (intConst Int32 0) residual_input (intConst Int32 1) + [ 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 @@ -1215,19 +1210,19 @@ 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 + group_size <- letSubExp "group_size" $ BasicOp $ BinOp (Mul Int64 OverflowUndef) tile_size tile_size num_groups_x <- letSubExp "num_groups_x" $ - BasicOp $ BinOp (SDivUp Int32 Unsafe) kdim_x tile_size + BasicOp $ BinOp (SDivUp Int64 Unsafe) kdim_x tile_size num_groups_y <- letSubExp "num_groups_y" $ - BasicOp $ BinOp (SDivUp Int32 Unsafe) kdim_y tile_size + BasicOp $ BinOp (SDivUp Int64 Unsafe) kdim_y tile_size num_groups <- letSubExp "num_groups_top" =<< foldBinOp - (Mul Int32 OverflowUndef) + (Mul Int64 OverflowUndef) num_groups_x (num_groups_y : map snd dims_on_top) @@ -1244,8 +1239,8 @@ reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) f ( untyped $ - isInt32 (LeafExp gtid_x int32) .<. pe32 kdim_x - .&&. isInt32 (LeafExp gtid_y int32) .<. pe32 kdim_y + 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), @@ -1255,7 +1250,7 @@ tilingTileShape = Shape [tile_size, tile_size], tilingNumWholeTiles = letSubExp "num_whole_tiles" $ - BasicOp $ BinOp (SQuot Int32 Unsafe) w tile_size, + BasicOp $ BinOp (SQuot Int64 Unsafe) w tile_size, tilingLevel = lvl, tilingSpace = space } 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 @@ -18,78 +18,158 @@ -- simplified away, but only *before* they are turned into loops. In -- principle this pass could be split into multiple, but for now it is -- kept together. -module Futhark.Optimise.Unstream (unstream) where +module Futhark.Optimise.Unstream (unstreamKernels, unstreamMC) where import Control.Monad.Reader import Control.Monad.State import Futhark.IR.Kernels +import qualified Futhark.IR.Kernels as Kernels import Futhark.IR.Kernels.Simplify (simplifyKernels) +import Futhark.IR.MC +import qualified Futhark.IR.MC as MC import Futhark.MonadFreshNames import Futhark.Pass import Futhark.Tools import qualified Futhark.Transform.FirstOrderTransform as FOT +-- | The pass for GPU kernels. +unstreamKernels :: Pass Kernels Kernels +unstreamKernels = unstream onHostOp simplifyKernels + +-- | The pass for multicore. +unstreamMC :: Pass MC MC +unstreamMC = unstream onMCOp MC.simplifyProg + data Stage = SeqStreams | SeqAll --- | The pass definition. -unstream :: Pass Kernels Kernels -unstream = +unstream :: + ASTLore lore => + (Stage -> OnOp lore) -> + (Prog lore -> PassM (Prog lore)) -> + Pass lore lore +unstream onOp simplify = Pass "unstream" "sequentialise remaining SOACs" $ intraproceduralTransformation (optimise SeqStreams) - >=> simplifyKernels + >=> simplify >=> intraproceduralTransformation (optimise SeqAll) where optimise stage scope stms = - modifyNameSource $ runState $ runReaderT (optimiseStms stage stms) scope + modifyNameSource $ + runState $ + runReaderT (optimiseStms (onOp stage) stms) scope -type UnstreamM = ReaderT (Scope Kernels) (State VNameSource) +type UnstreamM lore = ReaderT (Scope lore) (State VNameSource) -optimiseStms :: Stage -> Stms Kernels -> UnstreamM (Stms Kernels) -optimiseStms stage stms = +type OnOp lore = + Pattern lore -> StmAux (ExpDec lore) -> Op lore -> UnstreamM lore [Stm lore] + +optimiseStms :: + ASTLore lore => + OnOp lore -> + Stms lore -> + UnstreamM lore (Stms lore) +optimiseStms onOp stms = localScope (scopeOf stms) $ - stmsFromList . concat <$> mapM (optimiseStm stage) (stmsToList stms) + stmsFromList . concat <$> mapM (optimiseStm onOp) (stmsToList stms) -optimiseBody :: Stage -> Body Kernels -> UnstreamM (Body Kernels) -optimiseBody stage (Body () stms res) = - Body () <$> optimiseStms stage stms <*> pure res +optimiseBody :: + ASTLore lore => + OnOp lore -> + Body lore -> + UnstreamM lore (Body lore) +optimiseBody onOp (Body aux stms res) = + Body aux <$> optimiseStms onOp stms <*> pure res -optimiseKernelBody :: Stage -> KernelBody Kernels -> UnstreamM (KernelBody Kernels) -optimiseKernelBody stage (KernelBody () stms res) = +optimiseKernelBody :: + ASTLore lore => + OnOp lore -> + KernelBody lore -> + UnstreamM lore (KernelBody lore) +optimiseKernelBody onOp (KernelBody attr stms res) = localScope (scopeOf stms) $ - KernelBody () - <$> (stmsFromList . concat <$> mapM (optimiseStm stage) (stmsToList stms)) + KernelBody attr + <$> (stmsFromList . concat <$> mapM (optimiseStm onOp) (stmsToList stms)) <*> pure res -optimiseLambda :: Stage -> Lambda Kernels -> UnstreamM (Lambda Kernels) -optimiseLambda stage lam = localScope (scopeOfLParams $ lambdaParams lam) $ do - body <- optimiseBody stage $ lambdaBody lam +optimiseLambda :: + ASTLore lore => + OnOp lore -> + Lambda lore -> + UnstreamM lore (Lambda lore) +optimiseLambda onOp lam = localScope (scopeOfLParams $ lambdaParams lam) $ do + body <- optimiseBody onOp $ lambdaBody lam return lam {lambdaBody = body} -sequentialise :: Stage -> SOAC Kernels -> Bool +optimiseStm :: + ASTLore lore => + OnOp lore -> + Stm lore -> + UnstreamM lore [Stm lore] +optimiseStm onOp (Let pat aux (Op op)) = + onOp pat aux op +optimiseStm onOp (Let pat aux e) = + pure <$> (Let pat aux <$> mapExpM optimise e) + where + optimise = + identityMapper + { mapOnBody = \scope -> + localScope scope . optimiseBody onOp + } + +optimiseSegOp :: + ASTLore lore => + OnOp lore -> + SegOp lvl lore -> + UnstreamM lore (SegOp lvl lore) +optimiseSegOp onOp op = + localScope (scopeOfSegSpace $ segSpace op) $ mapSegOpM optimise op + where + optimise = + identitySegOpMapper + { mapOnSegOpBody = optimiseKernelBody onOp, + mapOnSegOpLambda = optimiseLambda onOp + } + +onMCOp :: Stage -> OnOp MC +onMCOp stage pat aux (ParOp par_op op) = do + par_op' <- traverse (optimiseSegOp (onMCOp stage)) par_op + op' <- optimiseSegOp (onMCOp stage) op + pure [Let pat aux $ Op $ ParOp par_op' op'] +onMCOp stage pat aux (MC.OtherOp soac) + | sequentialise stage soac = do + stms <- runBinder_ $ FOT.transformSOAC pat soac + fmap concat $ + localScope (scopeOf stms) $ + mapM (optimiseStm (onMCOp stage)) $ stmsToList stms + | otherwise = + -- Still sequentialise whatever's inside. + pure <$> (Let pat aux . Op . MC.OtherOp <$> mapSOACM optimise soac) + where + optimise = + identitySOACMapper + { mapOnSOACLambda = optimiseLambda (onMCOp stage) + } + +sequentialise :: Stage -> SOAC lore -> Bool 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))) +onHostOp :: Stage -> OnOp Kernels +onHostOp stage pat aux (Kernels.OtherOp soac) | sequentialise stage soac = do stms <- runBinder_ $ FOT.transformSOAC pat soac - fmap concat $ localScope (scopeOf stms) $ mapM (optimiseStm stage) $ stmsToList stms - | otherwise = do + fmap concat $ + localScope (scopeOf stms) $ + mapM (optimiseStm (onHostOp 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) + pure <$> (Let pat aux . Op . Kernels.OtherOp <$> mapSOACM optimise soac) where optimise = - identitySegOpMapper - { mapOnSegOpBody = optimiseKernelBody stage, - mapOnSegOpLambda = optimiseLambda stage + identitySOACMapper + { mapOnSOACLambda = optimiseLambda (onHostOp stage) } -optimiseStm stage (Let pat aux e) = - pure <$> (Let pat aux <$> mapExpM optimise e) - where - optimise = identityMapper {mapOnBody = \scope -> localScope scope . optimiseBody stage} +onHostOp stage pat aux (SegOp op) = + pure <$> (Let pat aux . Op . SegOp <$> optimiseSegOp (onHostOp stage) op) +onHostOp _ pat aux op = return [Let pat aux $ Op op] 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 @@ -212,24 +212,19 @@ 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 <- + (num_threads, num_threads_stms) <- + runBinder $ letSubExp "num_threads" $ BasicOp $ BinOp - (Mul Int32 OverflowUndef) + (Mul Int64 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) (invariant_alloc_stms, invariant_alloc_offsets) <- inScopeOf num_threads_stms $ expandedInvariantAllocations - (num_threads64, num_groups64, segNumGroups lvl, segGroupSize lvl) + (num_threads, segNumGroups lvl, segGroupSize lvl) space invariant_allocs @@ -356,7 +351,6 @@ expandedInvariantAllocations :: ( SubExp, - SubExp, Count NumGroups SubExp, Count GroupSize SubExp ) -> @@ -364,8 +358,7 @@ Extraction -> ExpandM (Stms KernelsMem, RebaseMap) expandedInvariantAllocations - ( num_threads64, - num_groups64, + ( num_threads, Count num_groups, Count group_size ) @@ -382,8 +375,8 @@ 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 + SegThread {} -> num_threads + SegGroup {} -> num_groups return ( stmsFromList [ Let sizepat (defAux ()) $ @@ -402,21 +395,20 @@ root_ixfun = IxFun.iota ( old_shape - ++ [ pe32 num_groups - * pe32 group_size + ++ [ pe64 num_groups * pe64 group_size ] ) permuted_ixfun = IxFun.permute root_ixfun perm offset_ixfun = IxFun.slice permuted_ixfun $ - DimFix (le32 (segFlat segspace)) : + DimFix (le64 (segFlat segspace)) : map untouched old_shape in offset_ixfun newBase SegGroup {} (old_shape, _) = - let root_ixfun = IxFun.iota (pe32 num_groups : old_shape) + let root_ixfun = IxFun.iota (pe64 num_groups : old_shape) offset_ixfun = IxFun.slice root_ixfun $ - DimFix (le32 (segFlat segspace)) : + DimFix (le64 (segFlat segspace)) : map untouched old_shape in offset_ixfun @@ -463,15 +455,14 @@ M.singleton mem $ newBase offset ) - num_threads' = pe32 num_threads - gtid = isInt32 $ 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 elems_per_thread = - isInt32 (sExt Int32 (primExpFromSubExp int64 size_per_thread)) - `quot` primByteSize pt + pe64 size_per_thread `quot` primByteSize pt root_ixfun = IxFun.iota [elems_per_thread, num_threads'] offset_ixfun = IxFun.slice @@ -486,7 +477,7 @@ in IxFun.reshape offset_ixfun shapechange -- | A map from memory block names to new index function bases. -type RebaseMap = M.Map VName (([TPrimExp Int32 VName], PrimType) -> IxFun) +type RebaseMap = M.Map VName (([TPrimExp Int64 VName], PrimType) -> IxFun) newtype OffsetM a = OffsetM @@ -511,7 +502,7 @@ askRebaseMap :: OffsetM RebaseMap askRebaseMap = OffsetM $ lift ask -lookupNewBase :: VName -> ([TPrimExp Int32 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 @@ -754,7 +745,7 @@ 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) @@ -769,8 +760,8 @@ let (kspace_gtids, kspace_dims) = unzip $ unSegSpace space new_inds = unflattenIndex - (map pe32 kspace_dims) - (pe32 $ Var $ paramName flat_gtid_lparam) + (map pe64 kspace_dims) + (pe64 $ Var $ paramName flat_gtid_lparam) zipWithM_ letBindNames (map pure kspace_gtids) =<< mapM toExp new_inds mapM_ addStm kstms' @@ -780,10 +771,6 @@ 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 @@ -792,12 +779,12 @@ w <- letSubExp "size_slice_w" - =<< foldBinOp (Mul Int32 OverflowUndef) (intConst Int32 1) (segSpaceDims space) + =<< 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 + Iota w (intConst Int64 0) (intConst Int64 1) Int64 let red_op = SegBinOp Commutative @@ -811,7 +798,7 @@ 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 @@ -273,14 +273,14 @@ arraySizeInBytesExp :: Type -> PrimExp VName arraySizeInBytesExp t = - untyped $ foldl' (*) (elemSize t) $ map (sExt64 . pe32) (arrayDims 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_i32 = product $ map (sExt64 . pe32) dims + let dim_prod_i64 = product $ map pe64 dims elm_size_i64 = primByteSize $ elemType t - return $ untyped $ dim_prod_i32 * elm_size_i64 + return $ untyped $ dim_prod_i64 * elm_size_i64 arraySizeInBytes :: Allocator lore m => Type -> m SubExp arraySizeInBytes = computeSize "bytes" <=< arraySizeInBytesExpM @@ -330,7 +330,7 @@ [PatElem lore] ) allocsForPattern sizeidents validents rts hints = do - let sizes' = [PatElem size $ MemPrim int32 | size <- map identName sizeidents] + let sizes' = [PatElem size $ MemPrim int64 | size <- map identName sizeidents] (vals, (exts, mems)) <- runWriterT $ forM (zip3 validents rts hints) $ \(ident, rt, hint) -> do @@ -414,7 +414,7 @@ size_exts sizeidents substs = M.fromList $ new_substs <> size_substs - ixfn <- instantiateIxFun $ IxFun.substituteInIxFun (fmap isInt32 substs) ext_ixfn + ixfn <- instantiateIxFun $ IxFun.substituteInIxFun (fmap isInt64 substs) ext_ixfn return (patels, ixfn) @@ -446,8 +446,8 @@ computeSize "bytes" $ untyped $ product - [ product $ map sExt64 $ IxFun.base ixfun, - fromIntegral (primByteSize (elemType t) :: Int64) + [ product $ IxFun.base ixfun, + primByteSize (elemType t) ] m <- allocateMemory "mem" bytes space return $ MemArray bt (arrayShape t) NoUniqueness $ ArrayIn m ixfun @@ -461,7 +461,7 @@ directIxFun :: PrimType -> Shape -> u -> VName -> Type -> MemBound u directIxFun bt shape u mem t = - let ixf = IxFun.iota $ map pe32 $ arrayDims t + let ixf = IxFun.iota $ map pe64 $ arrayDims t in MemArray bt shape u $ ArrayIn mem ixf allocInFParams :: @@ -488,7 +488,7 @@ case paramDeclType param of Array bt shape u -> do let memname = baseString (paramName param) <> "_mem" - ixfun = IxFun.iota $ map pe32 $ shapeDims shape + 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} @@ -541,8 +541,8 @@ ( \_ -> do vname <- lift $ newVName "ctx_param_ext" return - ( Param vname $ MemPrim int32, - fmap Free $ pe32 $ Var vname + ( Param vname $ MemPrim int64, + fmap Free $ pe64 $ Var vname ) ) substs @@ -573,7 +573,7 @@ (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => Space -> VName -> - AllocM fromlore tolore (SubExp, ExtIxFun, [TPrimExp Int32 VName], VName) + AllocM fromlore tolore (SubExp, ExtIxFun, [TPrimExp Int64 VName], VName) existentializeArray space v = do (mem', ixfun) <- lookupArraySummary v sp <- lookupMemSpace mem' @@ -604,7 +604,7 @@ <$> mapM ( \s -> do vname <- lift $ letExp "ctx_val" =<< toExp s - return (Var vname, fmap Free $ primExpFromSubExp int32 $ Var vname) + return (Var vname, fmap Free $ primExpFromSubExp int64 $ Var vname) ) substs @@ -726,8 +726,8 @@ ReturnsNewBlock DefaultSpace i $ IxFun.iota $ map convert $ shapeDims shape - convert (Ext i) = le32 $ Ext i - convert (Free v) = Free <$> pe32 v + convert (Ext i) = le64 $ Ext i + convert (Free v) = Free <$> pe64 v startOfFreeIDRange :: [TypeBase ExtShape u] -> Int startOfFreeIDRange = S.size . shapeContext @@ -877,7 +877,7 @@ generalize :: (Maybe Space, Maybe IxFun) -> (Maybe Space, Maybe IxFun) -> - (Maybe Space, Maybe (ExtIxFun, [(TPrimExp Int32 VName, TPrimExp Int32 VName)])) + (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) @@ -938,7 +938,7 @@ [ExtType] -> Body tolore -> [Maybe Space] -> - [Maybe (ExtIxFun, [TPrimExp Int32 VName])] -> + [Maybe (ExtIxFun, [TPrimExp Int64 VName])] -> AllocM fromlore tolore (Body tolore, [BodyReturns]) addResCtxInIfBody ifrets (Body _ bnds res) spaces substs = do let num_vals = length ifrets @@ -1006,8 +1006,8 @@ inspect (Prim pt) _ = MemPrim pt inspect (Mem space) _ = MemMem space - convert (Ext i) = le32 (Ext i) - convert (Free v) = Free <$> pe32 v + convert (Ext i) = le64 (Ext i) + convert (Free v) = Free <$> pe64 v adjustExtV :: Int -> Ext VName -> Ext VName adjustExtV _ (Free v) = Free v @@ -1050,10 +1050,10 @@ (mem, ixfun) <- lookupArraySummary a case paramType p of Array bt shape u -> do - dims <- map pe32 . arrayDims <$> lookupType a + dims <- map pe64 . arrayDims <$> lookupType a let ixfun' = IxFun.slice ixfun $ - fullSliceNum dims [DimFix $ le32 i] + fullSliceNum dims [DimFix $ le64 i] return (p {paramDec = MemArray bt shape u $ ArrayIn mem ixfun'}, a) Prim bt -> return (p {paramDec = MemPrim bt}, a) 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 @@ -27,9 +27,6 @@ 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 @@ -49,7 +46,7 @@ letSubExp "num_threads" $ BasicOp $ BinOp - (Mul Int32 OverflowUndef) + (Mul Int64 OverflowUndef) (unCount (segNumGroups lvl)) (unCount (segGroupSize lvl)) allocAtLevel lvl $ mapSegOpM (mapper num_threads) op @@ -85,7 +82,7 @@ dims <- arrayDims <$> lookupType v let perm_inv = rearrangeInverse perm dims' = rearrangeShape perm dims - ixfun = IxFun.permute (IxFun.iota $ map pe32 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)))) = zipWithM (mapResultHint lvl space) ts $ kernelBodyResult body @@ -107,12 +104,12 @@ mapResultHint lvl space = hint where num_threads = - pe32 (unCount $ segNumGroups lvl) * pe32 (unCount $ segGroupSize lvl) + pe64 (unCount $ segNumGroups lvl) * pe64 (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 bs [Constant (IntValue (Int64Value d))] = bs * d > 4 coalesceReturnOfShape _ _ = True hint t Returns {} @@ -124,9 +121,9 @@ 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 [num_threads, pe32 elems_per_thread] + 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 . pe32) [w] + ixfun = IxFun.reshape ixfun_tr $ map (DimNew . pe64) [w] return $ Hint ixfun DefaultSpace hint _ _ = return NoHint @@ -139,7 +136,7 @@ ++ [0 .. length space_dims -1] perm_inv = rearrangeInverse perm dims_perm = rearrangeShape perm dims - ixfun_base = IxFun.iota $ map pe32 dims_perm + ixfun_base = IxFun.iota $ map pe64 dims_perm ixfun_rearranged = IxFun.permute ixfun_base perm_inv in ixfun_rearranged @@ -156,8 +153,8 @@ return $ if private r && all (semiStatic consts) (arrayDims t) then - let seg_dims = map pe32 $ segSpaceDims space - dims = seg_dims ++ map pe32 (arrayDims t) + 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) $ @@ -178,7 +175,7 @@ maybePrivate consts t | Just (Array pt shape _) <- hasStaticShape t, all (semiStatic consts) $ shapeDims shape = do - let ixfun = IxFun.iota $ map pe32 $ shapeDims shape + let ixfun = IxFun.iota $ map pe64 $ shapeDims shape return $ Hint ixfun $ ScalarSpace (shapeDims shape) pt | otherwise = return NoHint diff --git a/src/Futhark/Pass/ExplicitAllocations/MC.hs b/src/Futhark/Pass/ExplicitAllocations/MC.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExplicitAllocations/MC.hs @@ -0,0 +1,37 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Futhark.Pass.ExplicitAllocations.MC (explicitAllocations) where + +import Futhark.IR.MC +import Futhark.IR.MCMem +import Futhark.Pass.ExplicitAllocations +import Futhark.Pass.ExplicitAllocations.SegOp + +instance SizeSubst (MCOp lore op) where + opSizeSubst _ _ = mempty + +handleSegOp :: SegOp () MC -> AllocM MC MCMem (SegOp () MCMem) +handleSegOp op = do + let num_threads = intConst Int64 256 -- FIXME + mapSegOpM (mapper num_threads) op + where + scope = scopeOfSegSpace $ segSpace op + mapper num_threads = + identitySegOpMapper + { mapOnSegOpBody = + localScope scope . allocInKernelBody, + mapOnSegOpLambda = + allocInBinOpLambda num_threads (segSpace op) + } + +handleMCOp :: Op MC -> AllocM MC MCMem (Op MCMem) +handleMCOp (ParOp par_op op) = + Inner <$> (ParOp <$> traverse handleSegOp par_op <*> handleSegOp op) +handleMCOp (OtherOp soac) = + error $ "Cannot allocate memory in SOAC: " ++ pretty soac + +explicitAllocations :: Pass MC MCMem +explicitAllocations = explicitAllocationsGeneric handleMCOp defaultExpHints 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,6 +1,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} module Futhark.Pass.ExplicitAllocations.SegOp ( allocInKernelBody, @@ -12,6 +13,9 @@ import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.Pass.ExplicitAllocations +instance SizeSubst (SegOp lvl lore) where + opSizeSubst _ _ = mempty + allocInKernelBody :: Allocable fromlore tolore => KernelBody fromlore -> @@ -34,8 +38,8 @@ allocInBinOpParams :: Allocable fromlore tolore => SubExp -> - TPrimExp Int32 VName -> - TPrimExp Int32 VName -> + TPrimExp Int64 VName -> + TPrimExp Int64 VName -> [LParam fromlore] -> [LParam fromlore] -> AllocM fromlore tolore ([LParam tolore], [LParam tolore]) @@ -46,12 +50,12 @@ Array bt shape u -> do twice_num_threads <- letSubExp "twice_num_threads" $ - BasicOp $ BinOp (Mul Int32 OverflowUndef) num_threads $ intConst Int32 2 + 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 pe32 $ arrayDims t + let base_dims = map pe64 $ arrayDims t ixfun_base = IxFun.iota base_dims ixfun_x = IxFun.slice ixfun_base $ @@ -83,8 +87,8 @@ allocInBinOpLambda num_threads (SegSpace flat _) lam = do let (acc_params, arr_params) = splitAt (length (lambdaParams lam) `div` 2) $ lambdaParams lam - index_x = TPrimExp $ LeafExp flat int32 - index_y = index_x + pe32 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 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,6 +1,5 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} module Futhark.Pass.ExplicitAllocations.Seq ( explicitAllocations, 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 @@ -315,7 +315,7 @@ runBinder $ do to_what' <- letSubExp "comparatee" - =<< foldBinOp (Mul Int32 OverflowUndef) (intConst Int32 1) to_what + =<< 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) @@ -594,7 +594,7 @@ String -> [SubExp] -> KernelPath -> - Maybe Int32 -> + Maybe Int64 -> DistribM ((SubExp, Name), Out.Stms Out.Kernels) sufficientParallelism desc ws path def = cmpSizeLe desc (Out.SizeThreshold path def) ws @@ -733,7 +733,7 @@ -- 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' :: @@ -796,7 +796,7 @@ fits <- letSubExp "fits" $ BasicOp $ - CmpOp (CmpSle Int32) group_size max_group_size + CmpOp (CmpSle Int64) group_size max_group_size addStms check_suff_stms 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 @@ -135,10 +135,10 @@ -- 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', @@ -148,7 +148,7 @@ letBindNames [to] $ BasicOp $ Index from $ - fullSlice from_t [DimFix $ intConst Int32 0] + fullSlice from_t [DimFix $ intConst Int64 0] ) nonSegRed :: 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 @@ -580,7 +580,7 @@ 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 + 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 = @@ -614,10 +614,10 @@ lam = Lambda { lambdaParams = [], - lambdaReturnType = [Prim int32, et], + lambdaReturnType = [Prim int64, et], lambdaBody = mkBody mempty [i, v] } - maybeDistributeStm (Let pat aux $ Op $ Scatter (intConst Int32 1) lam [] [(w, 1, arr)]) acc + maybeDistributeStm (Let pat aux $ Op $ Scatter (intConst Int64 1) lam [] [(w, 1, arr)]) acc where amortises DoLoop {} = True amortises Op {} = True @@ -839,7 +839,7 @@ letSubExp "v" $ BasicOp $ Index v $ map (DimFix . Var) slice_gtids slice_is <- traverse (toSubExp "index") $ - fixSlice (map (fmap pe32) slice) $ map (pe32 . Var) slice_gtids + fixSlice (map (fmap pe64) slice) $ map (pe64 . Var) slice_gtids let write_is = map (Var . fst) base_ispace ++ slice_is arr' = @@ -991,7 +991,7 @@ BasicOp $ Index ne_v $ fullSlice ne_v_t $ - replicate (shapeRank shape) $ DimFix $ intConst Int32 0 + replicate (shapeRank shape) $ DimFix $ intConst Int64 0 return (lam', nes', shape) Nothing -> return (lam, nes, mempty) 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 @@ -256,7 +256,7 @@ constructKernel mk_lvl kernel_nest inner_body = runBinderT' $ do (ispace, inps) <- flatKernel kernel_nest let aux = loopNestingAux first_nest - ispace_scope = M.fromList $ zip (map fst ispace) $ repeat $ IndexName Int32 + ispace_scope = M.fromList $ zip (map fst ispace) $ repeat $ IndexName Int64 pat = loopNestingPattern first_nest rts = map (stripArray (length ispace)) $ patternTypes pat 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 @@ -103,7 +103,7 @@ letSubExp "acc" $ BasicOp $ Index v $ - fullSlice v_t [DimFix $ intConst Int32 0] + fullSlice v_t [DimFix $ intConst Int64 0] indexAcc Constant {} = error "irwim: array accumulator is a constant." accs' <- mapM indexAcc accs 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 @@ -59,7 +59,7 @@ lift $ runBinder $ letSubExp "intra_num_groups" - =<< foldBinOp (Mul Int32 OverflowUndef) (intConst Int32 1) (map snd ispace) + =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) (map snd ispace) let body = lambdaBody lam @@ -82,18 +82,18 @@ ((intra_avail_par, kspace, read_input_stms), prelude_stms) <- lift $ runBinder $ do - let foldBinOp' _ [] = eSubExp $ intConst Int32 0 + let foldBinOp' _ [] = eSubExp $ intConst Int64 0 foldBinOp' bop (x : xs) = foldBinOp bop x xs ws_min <- - mapM (letSubExp "one_intra_par_min" <=< foldBinOp' (Mul Int32 OverflowUndef)) $ + 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 Int32 OverflowUndef)) $ + 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 + 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 @@ -102,10 +102,10 @@ =<< if null ws_min then eBinOp - (SMin Int32) + (SMin Int64) (eSubExp =<< letSubExp "max_group_size" (Op $ SizeOp $ Out.GetSizeMax Out.SizeGroup)) (eSubExp intra_avail_par) - else foldBinOp' (SMax Int32) ws_min + else foldBinOp' (SMax Int64) ws_min let inputIsUsed input = kernelInputName input `nameIn` freeIn body used_inps = filter inputIsUsed inps 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 @@ -48,12 +48,14 @@ SubExp -> SubExp -> m (SubExp, SubExp) -numberOfGroups desc w64 group_size = do +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 + 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 :: @@ -64,12 +66,11 @@ 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) + =<< eBinOp (SDivUp Int64 Unsafe) (eSubExp w) (eSubExp num_threads) return $ KernelSize per_thread_elements num_threads @@ -87,13 +88,13 @@ 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 + 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))] + 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 @@ -132,7 +133,7 @@ red_ts = take num_nonconcat $ lambdaReturnType lam map_ts = map rowType $ drop num_nonconcat $ lambdaReturnType lam - per_thread <- asIntS Int32 $ kernelElementsPerThread kernel_size + per_thread <- asIntS Int64 $ kernelElementsPerThread kernel_size splitArrays (paramName chunk_size) (map paramName arr_params) @@ -175,7 +176,7 @@ m (Lambda Kernels) kerneliseLambda nes lam = do thread_index <- newVName "thread_index" - let thread_index_param = Param thread_index $ Prim int32 + let thread_index_param = Param thread_index $ Prim int64 (fold_chunk_param, fold_acc_params, fold_inp_params) = partitionChunkedFoldParameters (length nes) $ lambdaParams lam @@ -214,8 +215,6 @@ 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 ()))) $ @@ -224,7 +223,7 @@ (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 + ConcatReturns split_ordering w elems_per_thread $ patElemName pe return ( map (Returns ResultMaySimplify . Var . patElemName) chunk_red_pes ++ map concatReturns chunk_map_pes @@ -304,24 +303,20 @@ -- array. segThreadCapped :: MonadFreshNames m => MkSegLevel Kernels m segThreadCapped ws desc r = do - w64 <- + w <- letSubExp "nest_size" - =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) - =<< mapM (asIntS Int64) ws + =<< 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 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 @@ -9,6 +9,7 @@ soacsStmToKernels, scopeForKernels, scopeForSOACs, + injectSOACS, ) where diff --git a/src/Futhark/Pass/ExtractMulticore.hs b/src/Futhark/Pass/ExtractMulticore.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractMulticore.hs @@ -0,0 +1,406 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} + +module Futhark.Pass.ExtractMulticore (extractMulticore) where + +import Control.Monad.Identity +import Control.Monad.Reader +import Control.Monad.State +import Futhark.Analysis.Rephrase +import Futhark.IR +import Futhark.IR.MC +import qualified Futhark.IR.MC as MC +import Futhark.IR.SOACS hiding + ( Body, + Exp, + LParam, + Lambda, + Pattern, + Stm, + ) +import qualified Futhark.IR.SOACS as SOACS +import qualified Futhark.IR.SOACS.Simplify as SOACS +import Futhark.Pass +import Futhark.Pass.ExtractKernels.DistributeNests +import Futhark.Pass.ExtractKernels.ToKernels (injectSOACS) +import Futhark.Tools +import qualified Futhark.Transform.FirstOrderTransform as FOT +import Futhark.Transform.Rename (Rename, renameSomething) +import Futhark.Util (chunks, takeLast) +import Futhark.Util.Log + +newtype ExtractM a = ExtractM (ReaderT (Scope MC) (State VNameSource) a) + deriving + ( Functor, + Applicative, + Monad, + HasScope MC, + LocalScope MC, + MonadFreshNames + ) + +-- XXX: throwing away the log here... +instance MonadLogger ExtractM where + addLog _ = pure () + +indexArray :: VName -> LParam SOACS -> VName -> Stm MC +indexArray i (Param p t) arr = + Let (Pattern [] [PatElem p t]) (defAux ()) $ + BasicOp $ Index arr $ DimFix (Var i) : map sliceDim (arrayDims t) + +mapLambdaToBody :: + (Body SOACS -> ExtractM (Body MC)) -> + VName -> + Lambda SOACS -> + [VName] -> + ExtractM (Body MC) +mapLambdaToBody onBody i lam arrs = do + let indexings = zipWith (indexArray i) (lambdaParams lam) arrs + Body () stms res <- inScopeOf indexings $ onBody $ lambdaBody lam + return $ Body () (stmsFromList indexings <> stms) res + +mapLambdaToKernelBody :: + (Body SOACS -> ExtractM (Body MC)) -> + VName -> + Lambda SOACS -> + [VName] -> + ExtractM (KernelBody MC) +mapLambdaToKernelBody onBody i lam arrs = do + Body () stms res <- mapLambdaToBody onBody i lam arrs + return $ KernelBody () stms $ map (Returns ResultMaySimplify) res + +reduceToSegBinOp :: Reduce SOACS -> ExtractM (Stms MC, SegBinOp MC) +reduceToSegBinOp (Reduce comm lam nes) = do + ((lam', nes', shape), stms) <- runBinder $ determineReduceOp lam nes + lam'' <- transformLambda lam' + return (stms, SegBinOp comm lam'' nes' shape) + +scanToSegBinOp :: Scan SOACS -> ExtractM (Stms MC, SegBinOp MC) +scanToSegBinOp (Scan lam nes) = do + ((lam', nes', shape), stms) <- runBinder $ determineReduceOp lam nes + lam'' <- transformLambda lam' + return (stms, SegBinOp Noncommutative lam'' nes' shape) + +histToSegBinOp :: SOACS.HistOp SOACS -> ExtractM (Stms MC, MC.HistOp MC) +histToSegBinOp (SOACS.HistOp num_bins rf dests nes op) = do + ((op', nes', shape), stms) <- runBinder $ determineReduceOp op nes + op'' <- transformLambda op' + return (stms, MC.HistOp num_bins rf dests nes' shape op'') + +mkSegSpace :: MonadFreshNames m => SubExp -> m (VName, SegSpace) +mkSegSpace w = do + flat <- newVName "flat_tid" + gtid <- newVName "gtid" + let space = SegSpace flat [(gtid, w)] + return (gtid, space) + +transformLoopForm :: LoopForm SOACS -> LoopForm MC +transformLoopForm (WhileLoop cond) = WhileLoop cond +transformLoopForm (ForLoop i it bound params) = ForLoop i it bound params + +transformStm :: Stm SOACS -> ExtractM (Stms MC) +transformStm (Let pat aux (BasicOp op)) = + pure $ oneStm $ Let pat aux $ BasicOp op +transformStm (Let pat aux (Apply f args ret info)) = + pure $ oneStm $ Let pat aux $ Apply f args ret info +transformStm (Let pat aux (DoLoop ctx val form body)) = do + let form' = transformLoopForm form + body' <- + localScope + ( scopeOfFParams (map fst ctx) + <> scopeOfFParams (map fst val) + <> scopeOf form' + ) + $ transformBody body + return $ oneStm $ Let pat aux $ DoLoop ctx val form' body' +transformStm (Let pat aux (If cond tbranch fbranch ret)) = + oneStm . Let pat aux + <$> (If cond <$> transformBody tbranch <*> transformBody fbranch <*> pure ret) +transformStm (Let pat aux (Op op)) = + fmap (certify (stmAuxCerts aux)) <$> transformSOAC pat (stmAuxAttrs aux) op + +transformLambda :: Lambda SOACS -> ExtractM (Lambda MC) +transformLambda (Lambda params body ret) = + Lambda params + <$> localScope (scopeOfLParams params) (transformBody body) + <*> pure ret + +transformStms :: Stms SOACS -> ExtractM (Stms MC) +transformStms stms = + case stmsHead stms of + Nothing -> return mempty + Just (stm, stms') -> do + stm_stms <- transformStm stm + inScopeOf stm_stms $ (stm_stms <>) <$> transformStms stms' + +transformBody :: Body SOACS -> ExtractM (Body MC) +transformBody (Body () stms res) = + Body () <$> transformStms stms <*> pure res + +sequentialiseBody :: Body SOACS -> ExtractM (Body MC) +sequentialiseBody = pure . runIdentity . rephraseBody toMC + where + toMC = injectSOACS OtherOp + +transformFunDef :: FunDef SOACS -> ExtractM (FunDef MC) +transformFunDef (FunDef entry attrs name rettype params body) = do + body' <- localScope (scopeOfFParams params) $ transformBody body + return $ FunDef entry attrs name rettype params body' + +-- Sets the chunk size to one. +unstreamLambda :: Attrs -> [SubExp] -> Lambda SOACS -> ExtractM (Lambda SOACS) +unstreamLambda attrs nes lam = do + let (chunk_param, acc_params, slice_params) = + partitionChunkedFoldParameters (length nes) (lambdaParams lam) + + inp_params <- forM slice_params $ \(Param p t) -> + newParam (baseString p) (rowType t) + + body <- runBodyBinder $ + localScope (scopeOfLParams inp_params) $ do + 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 inp_params) $ \(slice, v) -> + letBindNames [paramName slice] $ + BasicOp $ ArrayLit [Var $ paramName v] (paramType v) + + (red_res, map_res) <- splitAt (length nes) <$> bodyBind (lambdaBody lam) + + map_res' <- forM map_res $ \se -> do + v <- letExp "map_res" $ BasicOp $ SubExp se + v_t <- lookupType v + letSubExp "chunk" $ + BasicOp $ + Index v $ + fullSlice v_t [DimFix $ intConst Int64 0] + + pure $ resultBody $ red_res <> map_res' + + let (red_ts, map_ts) = splitAt (length nes) $ lambdaReturnType lam + map_lam = + Lambda + { lambdaReturnType = red_ts ++ map rowType map_ts, + lambdaParams = inp_params, + lambdaBody = body + } + + soacs_scope <- castScope <$> askScope + map_lam' <- runReaderT (SOACS.simplifyLambda map_lam) soacs_scope + + if "sequential_inner" `inAttrs` attrs + then FOT.transformLambda map_lam' + else return map_lam' + +-- Code generation for each parallel basic block is parameterised over +-- how we handle parallelism in the body (whether it's sequentialised +-- by keeping it as SOACs, or turned into SegOps). + +data NeedsRename = DoRename | DoNotRename + +renameIfNeeded :: Rename a => NeedsRename -> a -> ExtractM a +renameIfNeeded DoRename = renameSomething +renameIfNeeded DoNotRename = pure + +transformMap :: + NeedsRename -> + (Body SOACS -> ExtractM (Body MC)) -> + SubExp -> + Lambda SOACS -> + [VName] -> + ExtractM (SegOp () MC) +transformMap rename onBody w map_lam arrs = do + (gtid, space) <- mkSegSpace w + kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs + renameIfNeeded rename $ + SegMap () space (lambdaReturnType map_lam) kbody + +transformRedomap :: + NeedsRename -> + (Body SOACS -> ExtractM (Body MC)) -> + SubExp -> + [Reduce SOACS] -> + Lambda SOACS -> + [VName] -> + ExtractM ([Stms MC], SegOp () MC) +transformRedomap rename onBody w reds map_lam arrs = do + (gtid, space) <- mkSegSpace w + kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs + (reds_stms, reds') <- unzip <$> mapM reduceToSegBinOp reds + op' <- + renameIfNeeded rename $ + SegRed () space reds' (lambdaReturnType map_lam) kbody + return (reds_stms, op') + +transformHist :: + NeedsRename -> + (Body SOACS -> ExtractM (Body MC)) -> + SubExp -> + [SOACS.HistOp SOACS] -> + Lambda SOACS -> + [VName] -> + ExtractM ([Stms MC], SegOp () MC) +transformHist rename onBody w hists map_lam arrs = do + (gtid, space) <- mkSegSpace w + kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs + (hists_stms, hists') <- unzip <$> mapM histToSegBinOp hists + op' <- + renameIfNeeded rename $ + SegHist () space hists' (lambdaReturnType map_lam) kbody + return (hists_stms, op') + +transformParStream :: + NeedsRename -> + (Body SOACS -> ExtractM (Body MC)) -> + SubExp -> + Commutativity -> + Lambda SOACS -> + [SubExp] -> + Lambda SOACS -> + [VName] -> + ExtractM (Stms MC, SegOp () MC) +transformParStream rename onBody w comm red_lam red_nes map_lam arrs = do + (gtid, space) <- mkSegSpace w + kbody <- mapLambdaToKernelBody onBody gtid map_lam arrs + (red_stms, red) <- reduceToSegBinOp $ Reduce comm red_lam red_nes + op <- + renameIfNeeded rename $ + SegRed () space [red] (lambdaReturnType map_lam) kbody + return (red_stms, op) + +transformSOAC :: Pattern SOACS -> Attrs -> SOAC SOACS -> ExtractM (Stms MC) +transformSOAC pat _ (Screma w form arrs) + | Just lam <- isMapSOAC form = do + seq_op <- transformMap DoNotRename sequentialiseBody w lam arrs + if lambdaContainsParallelism lam + then do + par_op <- transformMap DoRename transformBody w lam arrs + return $ oneStm (Let pat (defAux ()) $ Op $ ParOp (Just par_op) seq_op) + else return $ oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op) + | Just (reds, map_lam) <- isRedomapSOAC form = do + (seq_reds_stms, seq_op) <- + transformRedomap DoNotRename sequentialiseBody w reds map_lam arrs + if lambdaContainsParallelism map_lam + then do + (par_reds_stms, par_op) <- + transformRedomap DoRename transformBody w reds map_lam arrs + return $ + mconcat (seq_reds_stms <> par_reds_stms) + <> oneStm (Let pat (defAux ()) $ Op $ ParOp (Just par_op) seq_op) + else + return $ + mconcat seq_reds_stms + <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op) + | Just (scans, map_lam) <- isScanomapSOAC form = do + (gtid, space) <- mkSegSpace w + kbody <- mapLambdaToKernelBody transformBody gtid map_lam arrs + (scans_stms, scans') <- unzip <$> mapM scanToSegBinOp scans + return $ + mconcat scans_stms + <> oneStm + ( Let pat (defAux ()) $ + Op $ + ParOp Nothing $ + SegScan () space scans' (lambdaReturnType map_lam) kbody + ) + | otherwise = do + -- This screma is too complicated for us to immediately do + -- anything, so split it up and try again. + scope <- castScope <$> askScope + transformStms =<< runBinderT_ (dissectScrema pat w form arrs) scope +transformSOAC pat _ (Scatter w lam ivs dests) = do + (gtid, space) <- mkSegSpace w + + Body () kstms res <- mapLambdaToBody transformBody gtid lam ivs + + let (dests_ws, dests_ns, dests_vs) = unzip3 dests + (i_res, v_res) = splitAt (sum dests_ns) res + rets = takeLast (length dests) $ lambdaReturnType lam + kres = do + (a_w, a, is_vs) <- + zip3 dests_ws dests_vs $ + chunks dests_ns $ zip i_res v_res + return $ WriteReturns [a_w] a [([DimFix i], v) | (i, v) <- is_vs] + kbody = KernelBody () kstms kres + return $ + oneStm $ + Let pat (defAux ()) $ + Op $ + ParOp Nothing $ + SegMap () space rets kbody +transformSOAC pat _ (Hist w hists map_lam arrs) = do + (seq_hist_stms, seq_op) <- + transformHist DoNotRename sequentialiseBody w hists map_lam arrs + + if lambdaContainsParallelism map_lam + then do + (par_hist_stms, par_op) <- + transformHist DoRename transformBody w hists map_lam arrs + return $ + mconcat (seq_hist_stms <> par_hist_stms) + <> oneStm (Let pat (defAux ()) $ Op $ ParOp (Just par_op) seq_op) + else + return $ + mconcat seq_hist_stms + <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op) +transformSOAC pat attrs (Stream w (Parallel _ comm red_lam red_nes) fold_lam arrs) + | not $ null red_nes = do + map_lam <- unstreamLambda attrs red_nes fold_lam + (seq_red_stms, seq_op) <- + transformParStream + DoNotRename + sequentialiseBody + w + comm + red_lam + red_nes + map_lam + arrs + + if lambdaContainsParallelism map_lam + then do + (par_red_stms, par_op) <- + transformParStream + DoRename + transformBody + w + comm + red_lam + red_nes + map_lam + arrs + return $ + seq_red_stms <> par_red_stms + <> oneStm (Let pat (defAux ()) $ Op $ ParOp (Just par_op) seq_op) + else + return $ + seq_red_stms + <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op) +transformSOAC pat _ (Stream w form lam arrs) = do + -- Just remove the stream and transform the resulting stms. + soacs_scope <- castScope <$> askScope + stream_stms <- + flip runBinderT_ soacs_scope $ + sequentialStreamWholeArray pat w (getStreamAccums form) lam arrs + transformStms stream_stms + +transformProg :: Prog SOACS -> PassM (Prog MC) +transformProg (Prog consts funs) = + modifyNameSource $ runState (runReaderT m mempty) + where + ExtractM m = do + consts' <- transformStms consts + funs' <- inScopeOf consts' $ mapM transformFunDef funs + return $ Prog consts' funs' + +extractMulticore :: Pass SOACS MC +extractMulticore = + Pass + { passName = "extract multicore parallelism", + passDescription = "Extract multicore parallelism", + passFunction = transformProg + } 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 @@ -118,7 +118,7 @@ letSubExp "num_threads" $ BasicOp $ BinOp - (Mul Int32 OverflowUndef) + (Mul Int64 OverflowUndef) (unCount $ segNumGroups lvl) (unCount $ segGroupSize lvl) evalStateT @@ -310,11 +310,10 @@ if null is then untyped $ pe32 num_threads else - coerceIntPrimExp Int32 $ - untyped $ - product $ - map pe32 $ - drop (length is) thread_gdims + 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 @@ -456,7 +455,7 @@ per_chunk <- letSubExp "per_chunk" $ - BasicOp $ BinOp (SQuot Int32 Unsafe) w_padded num_chunks' + 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 @@ -489,7 +488,7 @@ (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) + =<< eSliceArray d arr_inv_tr (eSubExp $ constant (0 :: Int64)) (eSubExp w) paddedScanReduceInput :: MonadBinder m => @@ -499,8 +498,8 @@ 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 + =<< eRoundToMultipleOf Int64 (eSubExp w) (eSubExp stride) + padding <- letSubExp "padding" $ BasicOp $ BinOp (Sub Int64 OverflowUndef) w_padded w return (w_padded, padding) --- Computing variance. 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 @@ -4,14 +4,18 @@ ( simplify, simplifySOACS, simplifySeq, + simplifyMC, simplifyKernels, simplifyKernelsMem, simplifySeqMem, + simplifyMCMem, ) where import qualified Futhark.IR.Kernels.Simplify as Kernels import qualified Futhark.IR.KernelsMem as KernelsMem +import qualified Futhark.IR.MC as MC +import qualified Futhark.IR.MCMem as MCMem import qualified Futhark.IR.SOACS.Simplify as SOACS import qualified Futhark.IR.Seq as Seq import qualified Futhark.IR.SeqMem as SeqMem @@ -32,8 +36,14 @@ simplifySeq :: Pass Seq.Seq Seq.Seq simplifySeq = simplify Seq.simplifyProg +simplifyMC :: Pass MC.MC MC.MC +simplifyMC = simplify MC.simplifyProg + simplifyKernelsMem :: Pass KernelsMem.KernelsMem KernelsMem.KernelsMem simplifyKernelsMem = simplify KernelsMem.simplifyProg simplifySeqMem :: Pass SeqMem.SeqMem SeqMem.SeqMem simplifySeqMem = simplify SeqMem.simplifyProg + +simplifyMCMem :: Pass MCMem.MCMem MCMem.MCMem +simplifyMCMem = simplify MCMem.simplifyProg diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs --- a/src/Futhark/Passes.hs +++ b/src/Futhark/Passes.hs @@ -7,12 +7,16 @@ kernelsPipeline, sequentialCpuPipeline, gpuPipeline, + mcPipeline, + multicorePipeline, ) where import Control.Category ((>>>)) import Futhark.IR.Kernels (Kernels) import Futhark.IR.KernelsMem (KernelsMem) +import Futhark.IR.MC (MC) +import Futhark.IR.MCMem (MCMem) import Futhark.IR.SOACS (SOACS) import Futhark.IR.Seq (Seq) import Futhark.IR.SeqMem (SeqMem) @@ -26,8 +30,10 @@ import Futhark.Optimise.Unstream import Futhark.Pass.ExpandAllocations import qualified Futhark.Pass.ExplicitAllocations.Kernels as Kernels +import qualified Futhark.Pass.ExplicitAllocations.MC as MC import qualified Futhark.Pass.ExplicitAllocations.Seq as Seq import Futhark.Pass.ExtractKernels +import Futhark.Pass.ExtractMulticore import Futhark.Pass.FirstOrderTransform import Futhark.Pass.KernelBabysitting import Futhark.Pass.Simplify @@ -59,10 +65,10 @@ [ simplifyKernels, babysitKernels, tileLoops, - unstream, + unstreamKernels, performCSE True, simplifyKernels, - sink, + sinkKernels, inPlaceLoweringKernels ] @@ -93,8 +99,33 @@ [ simplifyKernelsMem, performCSE False, simplifyKernelsMem, - doubleBuffer, + doubleBufferKernels, simplifyKernelsMem, expandAllocations, simplifyKernelsMem + ] + +mcPipeline :: Pipeline SOACS MC +mcPipeline = + standardPipeline + >>> onePass extractMulticore + >>> passes + [ simplifyMC, + unstreamMC, + performCSE True, + simplifyMC, + sinkMC, + inPlaceLoweringMC + ] + +multicorePipeline :: Pipeline SOACS MCMem +multicorePipeline = + mcPipeline + >>> onePass MC.explicitAllocations + >>> passes + [ simplifyMCMem, + performCSE False, + simplifyMCMem, + doubleBufferMC, + simplifyMCMem ] 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 @@ -142,7 +142,7 @@ 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) $ @@ -220,10 +220,10 @@ 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 @@ -232,7 +232,7 @@ ) $ do let slice = - [DimSlice (Var i) (Var (paramName chunk_size_param)) (intConst Int32 1)] + [DimSlice (Var i) (Var (paramName chunk_size_param)) (intConst Int64 1)] forM_ (zip chunk_params arrs) $ \(p, arr) -> letBindNames [paramName p] $ BasicOp $ @@ -265,7 +265,7 @@ let merge = loopMerge asOuts $ map Var as_vs loopBody <- runBodyBinder $ localScope - ( M.insert iter (IndexName Int32) $ + ( M.insert iter (IndexName Int64) $ scopeOfFParams $ map fst merge ) $ do @@ -283,7 +283,7 @@ foldM saveInArray arr $ zip indexes' values' return $ resultBody (map Var ress) - letBind pat $ DoLoop [] merge (ForLoop iter Int32 len []) loopBody + letBind pat $ DoLoop [] merge (ForLoop iter Int64 len []) loopBody transformSOAC pat (Hist len ops bucket_fun imgs) = do iter <- newVName "iter" @@ -295,7 +295,7 @@ -- Bind lambda-bodies for operators. loopBody <- runBodyBinder $ localScope - ( M.insert iter (IndexName Int32) $ + ( M.insert iter (IndexName Int64) $ scopeOfFParams $ map fst merge ) $ do @@ -345,7 +345,7 @@ 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 :: 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 @@ -20,6 +20,7 @@ renameBody, renameLambda, renamePattern, + renameSomething, -- * Renaming annotations RenameM, @@ -111,6 +112,13 @@ renamePattern = modifyNameSource . runRenamer . rename' where rename' pat = bind (patternNames pat) $ rename pat + +-- | Rename the bound variables in something (does not affect free variables). +renameSomething :: + (Rename a, MonadFreshNames m) => + a -> + m a +renameSomething = modifyNameSource . runRenamer . rename newtype RenameEnv = RenameEnv {envNameMap :: M.Map VName VName} diff --git a/src/Futhark/TypeCheck.hs b/src/Futhark/TypeCheck.hs --- a/src/Futhark/TypeCheck.hs +++ b/src/Futhark/TypeCheck.hs @@ -49,6 +49,7 @@ consume, consumeOnlyParams, binding, + alternative, ) where @@ -810,17 +811,17 @@ 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 int64] e require [Prim $ IntType et] x require [Prim $ IntType et] s checkBasicOp (Replicate (Shape dims) valexp) = do - mapM_ (require [Prim int32]) dims + 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 int32] . newDim) newshape + mapM_ (require [Prim int64] . newDim) newshape zipWithM_ (checkDimChange rank) newshape [0 ..] where checkDimChange _ (DimNew _) _ = @@ -845,7 +846,7 @@ checkBasicOp (Rotate rots arr) = do arrt <- lookupType arr let rank = arrayRank arrt - mapM_ (require [Prim int32]) rots + mapM_ (require [Prim int64]) rots when (length rots /= rank) $ bad $ TypeError $ @@ -870,7 +871,7 @@ ++ pretty arr1t ++ " and " ++ intercalate ", " (map pretty arr2ts) - require [Prim int32] ressize + require [Prim int64] ressize checkBasicOp (Copy e) = void $ checkArrIdent e checkBasicOp (Manifest perm arr) = @@ -1054,7 +1055,7 @@ Checkable lore => TypeBase Shape u -> TypeM lore () -checkType (Mem (ScalarSpace d _)) = mapM_ (require [Prim int32]) d +checkType (Mem (ScalarSpace d _)) = mapM_ (require [Prim int64]) d checkType t = mapM_ checkSubExp $ arrayDims t checkExtType :: @@ -1106,8 +1107,8 @@ 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] +checkDimIndex (DimFix i) = require [Prim int64] i +checkDimIndex (DimSlice i n s) = mapM_ (require [Prim int64]) [i, n, s] checkStm :: Checkable lore => @@ -1199,7 +1200,7 @@ let ctx_vals = zip ctx_res ctx_ts instantiateExt i = case maybeNth i ctx_vals of - Just (se, Prim (IntType Int32)) -> return se + Just (se, Prim (IntType Int64)) -> return se _ -> problem rettype' <- instantiateShapes instantiateExt rettype 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 @@ -80,7 +80,7 @@ type Stack = [StackFrame] -type Sizes = M.Map VName Int32 +type Sizes = M.Map VName Int64 -- | The monad in which evaluation takes place. newtype EvalM a @@ -119,14 +119,14 @@ lookupImport :: FilePath -> EvalM (Maybe Env) lookupImport f = asks $ M.lookup f . snd -putExtSize :: VName -> Int32 -> EvalM () +putExtSize :: VName -> Int64 -> EvalM () putExtSize v x = modify $ M.insert v x getSizes :: EvalM Sizes getSizes = get extSizeEnv :: EvalM Env -extSizeEnv = i32Env <$> getSizes +extSizeEnv = i64Env <$> getSizes prettyRecord :: Pretty a => M.Map Name a -> Doc prettyRecord m @@ -149,7 +149,7 @@ | ShapeSum (M.Map Name [Shape d]) deriving (Eq, Show, Functor, Foldable, Traversable) -type ValueShape = Shape Int32 +type ValueShape = Shape Int64 instance Pretty d => Pretty (Shape d) where ppr ShapeLeaf = mempty @@ -180,7 +180,7 @@ go _ = ShapeLeaf -structTypeShape :: M.Map VName ValueShape -> StructType -> Shape (Maybe Int32) +structTypeShape :: M.Map VName ValueShape -> StructType -> Shape (Maybe Int64) structTypeShape shapes = fmap dim . typeShape shapes' where dim (ConstDim d) = Just $ fromIntegral d @@ -212,10 +212,10 @@ matchDims (NamedDim (QualName _ d1)) (ConstDim d2) | d1 `elem` names = - i32Env $ M.singleton d1 $ fromIntegral d2 + i64Env $ M.singleton d1 $ fromIntegral d2 matchDims _ _ = mempty -resolveExistentials :: [VName] -> StructType -> ValueShape -> M.Map VName Int32 +resolveExistentials :: [VName] -> StructType -> ValueShape -> M.Map VName Int64 resolveExistentials names = match where match (Scalar (Record poly_fields)) (ShapeRecord fields) = @@ -280,7 +280,7 @@ valueShape (ValueSum shape _ _) = shape valueShape _ = ShapeLeaf -checkShape :: Shape (Maybe Int32) -> ValueShape -> Maybe ValueShape +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 @@ -319,7 +319,7 @@ -- | Create an array value; failing if that would result in an -- irregular array. -mkArray :: TypeBase Int32 () -> [Value] -> Maybe Value +mkArray :: TypeBase Int64 () -> [Value] -> Maybe Value mkArray t [] = return $ toArray (typeShape mempty t) [] mkArray _ (v : vs) = do @@ -350,8 +350,8 @@ asSigned (ValuePrim (SignedValue v)) = v asSigned v = error $ "Unexpected not a signed integer: " ++ pretty v -asInt32 :: Value -> Int32 -asInt32 = fromIntegral . asInteger +asInt64 :: Value -> Int64 +asInt64 = fromIntegral . asInteger asBool :: Value -> Bool asBool (ValuePrim (BoolValue x)) = x @@ -434,12 +434,12 @@ where tbind = T.TypeAbbr Unlifted [] -i32Env :: M.Map VName Int32 -> Env -i32Env = valEnv . M.map f +i64Env :: M.Map VName Int64 -> Env +i64Env = valEnv . M.map f where f x = - ( Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int32, - ValuePrim $ SignedValue $ Int32Value x + ( Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int64, + ValuePrim $ SignedValue $ Int64Value x ) instance Show InterpreterError where @@ -538,8 +538,8 @@ patternMatch _ _ _ = mzero data Indexing - = IndexingFix Int32 - | IndexingSlice (Maybe Int32) (Maybe Int32) (Maybe Int32) + = IndexingFix Int64 + | IndexingSlice (Maybe Int64) (Maybe Int64) (Maybe Int64) instance Pretty Indexing where ppr (IndexingFix i) = ppr i @@ -556,10 +556,10 @@ maybe mempty ppr i <> text ":" indexesFor :: - Maybe Int32 -> - Maybe Int32 -> - Maybe Int32 -> - Int32 -> + Maybe Int64 -> + Maybe Int64 -> + Maybe Int64 -> + Int64 -> Maybe [Int] indexesFor start end stride n | (start', end', stride') <- slice, @@ -640,11 +640,11 @@ evalDimIndex :: Env -> DimIndex -> EvalM Indexing evalDimIndex env (DimFix x) = - IndexingFix . asInt32 <$> eval env x + IndexingFix . asInt64 <$> 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 + 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 @@ -670,7 +670,7 @@ in arrayOf et' shape' u where evalDim (NamedDim qn) - | Just (TermValue _ (ValuePrim (SignedValue (Int32Value x)))) <- + | Just (TermValue _ (ValuePrim (SignedValue (Int64Value x)))) <- lookupVar qn env = ConstDim $ fromIntegral x evalDim d = d @@ -742,7 +742,7 @@ | null missing_sizes = env' | otherwise = env' - <> i32Env + <> i64Env ( resolveExistentials missing_sizes (patternStructType p) @@ -786,7 +786,7 @@ evalArg env e ext = do v <- eval env e case ext of - Just ext' -> putExtSize ext' $ asInt32 v + Just ext' -> putExtSize ext' $ asInt64 v Nothing -> return () return v @@ -1037,7 +1037,7 @@ sparams (patternStructType pat) (valueShape v) - in matchPattern (i32Env sparams' <> env) pat v + in matchPattern (i64Env sparams' <> env) pat v inc = (`P.doAdd` Int64Value 1) zero = (`P.doMul` Int64Value 0) @@ -1051,7 +1051,7 @@ ( valEnv ( M.singleton iv - ( Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int32, + ( Just $ T.BoundV [] $ Scalar $ Prim $ Signed Int64, ValuePrim (SignedValue i) ) ) @@ -1579,7 +1579,7 @@ toTuple [ toArray' rowshape $ concat parts, toArray' rowshape $ - map (ValuePrim . SignedValue . Int32Value . genericLength) parts + map (ValuePrim . SignedValue . Int64Value . genericLength) parts ] pack . map reverse @@ -1635,8 +1635,8 @@ def "unflatten" = Just $ fun3t $ \n m xs -> do let (ShapeDim _ innershape, xs') = fromArray xs - rowshape = ShapeDim (asInt32 m) innershape - shape = ShapeDim (asInt32 n) rowshape + 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 @@ -1652,7 +1652,7 @@ return $ T.TypeAbbr Unlifted [] $ Scalar $ Prim t stream f arg@(ValueArray _ xs) = - let n = ValuePrim $ SignedValue $ Int32Value $ arrayLength xs + let n = ValuePrim $ SignedValue $ Int64Value $ arrayLength xs in apply2 noLoc mempty f n arg stream _ arg = error $ "Cannot stream: " ++ pretty arg 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/Pretty.hs b/src/Language/Futhark/Pretty.hs --- a/src/Language/Futhark/Pretty.hs +++ b/src/Language/Futhark/Pretty.hs @@ -115,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 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 @@ -832,8 +832,8 @@ ( "unflatten", IntrinsicPolyFun [tp_a] - [ Scalar $ Prim $ Signed Int32, - Scalar $ Prim $ Signed Int32, + [ Scalar $ Prim $ Signed Int64, + Scalar $ Prim $ Signed Int64, Array () Nonunique t_a (rank 1) ] $ Array () Nonunique t_a (rank 2) @@ -847,7 +847,7 @@ ( "rotate", IntrinsicPolyFun [tp_a] - [Scalar $ Prim $ Signed Int32, arr_a] + [Scalar $ Prim $ Signed Int64, arr_a] arr_a ), ("transpose", IntrinsicPolyFun [tp_a] [arr_2d_a] arr_2d_a), @@ -855,7 +855,7 @@ IntrinsicPolyFun [tp_a] [ Array () Unique t_a (rank 1), - Array () Nonunique (Prim $ Signed Int32) (rank 1), + Array () Nonunique (Prim $ Signed Int64) (rank 1), Array () Nonunique t_a (rank 1) ] $ Array () Unique t_a (rank 1) @@ -865,11 +865,11 @@ ( "hist", IntrinsicPolyFun [tp_a] - [ Scalar $ Prim $ Signed Int32, + [ Scalar $ Prim $ Signed Int64, uarr_a, Scalar t_a `arr` (Scalar t_a `arr` Scalar t_a), Scalar t_a, - Array () Nonunique (Prim $ Signed Int32) (rank 1), + Array () Nonunique (Prim $ Signed Int64) (rank 1), arr_a ] uarr_a @@ -897,28 +897,28 @@ IntrinsicPolyFun [tp_a] [ Scalar (Prim $ Signed Int32), - Scalar t_a `arr` Scalar (Prim $ Signed Int32), + Scalar t_a `arr` Scalar (Prim $ Signed Int64), arr_a ] - $ tupleRecord [uarr_a, Array () Unique (Prim $ Signed Int32) (rank 1)] + $ tupleRecord [uarr_a, Array () Unique (Prim $ Signed Int64) (rank 1)] ), ( "map_stream", IntrinsicPolyFun [tp_a, tp_b] - [Scalar (Prim $ Signed Int32) `karr` (arr_ka `arr` arr_kb), arr_a] + [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 Int32) `karr` (arr_ka `arr` arr_kb), arr_a] + [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 Int32) `karr` (arr_ka `arr` Scalar t_b), + Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` Scalar t_b), arr_a ] $ Scalar t_b @@ -927,7 +927,7 @@ 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), + Scalar (Prim $ Signed Int64) `karr` (arr_ka `arr` Scalar t_b), arr_a ] $ Scalar t_b 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 @@ -433,7 +433,7 @@ type StructType = TypeBase (DimDecl VName) () -- | A value type contains full, manifest size information. -type ValueType = TypeBase Int32 () +type ValueType = TypeBase Int64 () -- | A dimension declaration expression for use in a 'TypeExp'. data DimExp vn 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 @@ -181,7 +181,7 @@ typeParamEnv (TypeParamDim v _) = mempty { envVtable = - M.singleton v $ BoundV [] (Scalar $ Prim $ Signed Int32) + M.singleton v $ BoundV [] (Scalar $ Prim $ Signed Int64) } typeParamEnv (TypeParamType l v _) = mempty 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 @@ -220,10 +220,10 @@ checkNamedDim loc v = do (v', t) <- lookupVar loc v case t of - Scalar (Prim (Signed Int32)) -> return v' + Scalar (Prim (Signed Int64)) -> return v' _ -> typeError loc mempty $ - "Dimension declaration" <+> ppr v <+> "should be of type i32." + "Dimension declaration" <+> ppr v <+> "should be of type i64." typeError :: Located loc => loc -> Notes -> Doc -> m a diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -576,9 +576,9 @@ checkNamedDim loc v = do (v', t) <- lookupVar loc v - onFailure (CheckingRequired [Scalar $ Prim $ Signed Int32] (toStruct t)) $ + onFailure (CheckingRequired [Scalar $ Prim $ Signed Int64] (toStruct t)) $ unify (mkUsage loc "use as array size") (toStruct t) $ - Scalar $ Prim $ Signed Int32 + Scalar $ Prim $ Signed Int64 return v' typeError loc notes s = do @@ -635,7 +635,7 @@ return tdecl' where observeDim (NamedDim v) = - observe $ Ident (qualLeaf v) (Info $ Scalar $ Prim $ Signed Int32) mempty + observe $ Ident (qualLeaf v) (Info $ Scalar $ Prim $ Signed Int64) mempty observeDim _ = return () -- | Instantiate a type scheme with fresh type variables for its type @@ -983,7 +983,7 @@ typeParamIdent :: TypeParam -> Maybe Ident typeParamIdent (TypeParamDim v loc) = - Just $ Ident v (Info $ Scalar $ Prim $ Signed Int32) loc + Just $ Ident v (Info $ Scalar $ Prim $ Signed Int64) loc typeParamIdent _ = Nothing bindingIdent :: @@ -1086,13 +1086,13 @@ -- Pattern match some known slices to be non-existential. adjustDims (DimSlice i j stride : idxes') (_ : dims) | refine_sizes, - maybe True ((== Just 0) . isInt32) i, + maybe True ((== Just 0) . isInt64) i, Just j' <- maybeDimFromExp =<< j, - maybe True ((== Just 1) . isInt32) stride = + 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) . isInt32) stride = + 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 @@ -1290,21 +1290,26 @@ 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 + 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 (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'' + 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) @@ -2282,7 +2287,7 @@ where check = maybe (return Nothing) $ - fmap Just . unifies "use as index" (Scalar $ Prim $ Signed Int32) <=< checkExp + fmap Just . unifies "use as index" (Scalar $ Prim $ Signed Int64) <=< checkExp sequentially :: TermTypeM a -> (a -> Occurences -> TermTypeM b) -> TermTypeM b sequentially m1 m2 = do @@ -2386,7 +2391,7 @@ return (tp1', tp2'', argext, ext) where - sizeSubst (Scalar (Prim (Signed Int32))) e = dimFromArg fname e + 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" @@ -2415,17 +2420,17 @@ | 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 +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 <$> isInt32 e +maybeDimFromExp e = ConstDim . fromIntegral <$> isInt64 e dimFromExp :: (Exp -> SizeSource) -> Exp -> TermTypeM (DimDecl VName, Maybe VName) dimFromExp rf (Parens e _) = dimFromExp rf e diff --git a/src/futhark.hs b/src/futhark.hs --- a/src/futhark.hs +++ b/src/futhark.hs @@ -19,6 +19,7 @@ import qualified Futhark.CLI.Dev as Dev import qualified Futhark.CLI.Doc as Doc import qualified Futhark.CLI.Misc as Misc +import qualified Futhark.CLI.Multicore as Multicore import qualified Futhark.CLI.OpenCL as OpenCL import qualified Futhark.CLI.Pkg as Pkg import qualified Futhark.CLI.PyOpenCL as PyOpenCL @@ -48,6 +49,7 @@ ("c", (C.main, "Compile to sequential C.")), ("opencl", (OpenCL.main, "Compile to C calling OpenCL.")), ("cuda", (CCUDA.main, "Compile to C calling CUDA.")), + ("multicore", (Multicore.main, "Compile to multicore C.")), ("python", (Python.main, "Compile to sequential Python.")), ("pyopencl", (PyOpenCL.main, "Compile to Python calling PyOpenCL.")), ("test", (Test.main, "Test Futhark programs.")),