diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,36 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). +## [0.25.37] + +### Added + +* `futhark fmt` prints module types better. + +* New server protocol commands: `rank`, `elemtype`, `new_array`, `set`. + +* `futhark lsp` provides inlay hints, they show type ascriptions for + inferred types of bindings, by VegOwOtenks. (#2398) + +* Scan chunk size is now exposed as tuning parameter in `cuda` and `hip` + backends. + +* `futhark lsp` offers code actions, they insert type ascriptions for + inferred types. Every named binding has an action. + +* `futhark bench` and `futhark test` can now handle entry points that return + opaque values, as long as there is no expected result. + +* Better fusion for `scan` SOACs. + +* New prelude function: `exscan`, an exclusive scan. + +### Fixed + +* `i64.set_bit`/`u64.set_bit` would produce wrong results in C-based backends. (#2396) + +* Some uses of higher order modules could cause infinite loops. (#2407) + ## [0.25.36] ### Added diff --git a/Setup.hs b/Setup.hs --- a/Setup.hs +++ b/Setup.hs @@ -4,4 +4,5 @@ main :: IO () main = defaultMainWithHooks myHooks - where myHooks = simpleUserHooks + where + myHooks = simpleUserHooks diff --git a/docs/man/futhark-lsp.rst b/docs/man/futhark-lsp.rst --- a/docs/man/futhark-lsp.rst +++ b/docs/man/futhark-lsp.rst @@ -44,8 +44,21 @@ Activating the code lens will load the file into the interpreter, evaluate the expression and write the result below. - The evaluation will be aborted after 15 seconds or when it has allocated + The evaluation will be aborted after 15 seconds or when it has allocated 100 GB in total, only the last 100 debugging traces will be retained. + +Inlay Hints + + Provides virtual text hints to visualize the results of type-checking. + Inferred types of bindings will be shown for e.g. lambda arguments, function + arguments, let bindings or loop bindings. + +Code Actions + + Every name binding with an inlay type hint has an associated code action that + inserts exactly the type ascription the virtual text shows. + If the type contains any inferred type variables or sizes, they will also be + introduced at the appropriate position. SEE ALSO ======== diff --git a/docs/server-protocol.rst b/docs/server-protocol.rst --- a/docs/server-protocol.rst +++ b/docs/server-protocol.rst @@ -186,15 +186,34 @@ Array Commands ~~~~~~~~~~~~~~ -``elemtype`` *v* -.................... +``rank`` *t* +............ -Print the typename of the elements of array-typed variable *v*. +Print the rank of array type *t*. +``elemtype`` *t* +................ + +Print the typename of the elements of array type *t*. + +``new_array`` *v0* *t* *s0* ... *sN-1* *v1* ... *vM* +.................................................... + +Create a new variable *v0* of type *t*, which must be an array type of rank *N*. +The size of each dimension of the array is given by *s0* ... *sN-1*, and the +values by *v1* ... *vM* in row-major order, where *M* is the product of the +dimension sizes. + +``set`` *v0* *v1* *i0* ... *iN-1* +................................. + +Perform an in-place replacement on array-typed variable *v0* of rank *N* at +indices *i0* ... *iN-1* with the value of variable *v1*. + ``shape`` *v* -.................... +............. -Print the shape of array-typed variable *v* as space-separated integers. +Print the shape of array-typed variable *v*, with one integer per line. ``index`` *v0* *v1* *i0* ... *iN-1* ................................... diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: futhark -version: 0.25.36 +version: 0.25.37 synopsis: An optimising compiler for a functional, array-oriented language. description: Futhark is a small programming language designed to be compiled to @@ -305,14 +305,17 @@ Futhark.Internalise.Monomorphise Futhark.Internalise.ReplaceRecords Futhark.Internalise.TypesValues + Futhark.LSP.CodeAction Futhark.LSP.CodeLens Futhark.LSP.CommandType Futhark.LSP.Compile Futhark.LSP.Diagnostic Futhark.LSP.Handlers - Futhark.LSP.Tool - Futhark.LSP.State + Futhark.LSP.InlayHint Futhark.LSP.PositionMapping + Futhark.LSP.State + Futhark.LSP.Tool + Futhark.LSP.TypeAscription Futhark.MonadFreshNames Futhark.Optimise.BlkRegTiling Futhark.Optimise.CSE @@ -322,6 +325,7 @@ Futhark.Optimise.Fusion.Composing Futhark.Optimise.Fusion.GraphRep Futhark.Optimise.Fusion.RulesWithAccs + Futhark.Optimise.Fusion.Screma Futhark.Optimise.Fusion.TryFusion Futhark.Optimise.GenRedOpt Futhark.Optimise.HistAccs @@ -474,7 +478,7 @@ , filepath >=1.4.1.1 , free >=5.1.10 , futhark-data >= 1.1.3.0 - , futhark-server >= 1.3.0.0 + , futhark-server >= 1.3.3.0 , futhark-manifest == 1.7.0.0 , githash >=0.1.6.1 , half >= 0.3 @@ -537,10 +541,13 @@ Futhark.IR.Syntax.CoreTests Futhark.IR.SyntaxTests Futhark.Internalise.TypesValuesTests + Futhark.LspTests Futhark.Optimise.MemoryBlockMerging.GreedyColoringTests Futhark.Optimise.ArrayLayout.AnalyseTests Futhark.Optimise.ArrayLayout.LayoutTests Futhark.Optimise.ArrayLayoutTests + Futhark.Optimise.Fusion.ScremaTests + Futhark.Optimise.FusionTests Futhark.Pkg.SolveTests Futhark.ProfileTests Language.Futhark.CoreTests @@ -551,6 +558,7 @@ Language.Futhark.SyntaxTests Language.Futhark.TypeChecker.TypesTests Language.Futhark.TypeChecker.ConsumptionTests + Language.Futhark.TypeChecker.ModulesTests Language.Futhark.TypeCheckerTests build-depends: QuickCheck >=2.8 @@ -558,9 +566,15 @@ , mtl >=2.2.1 , base , containers + , co-log-core , free , futhark + , lens + , lsp + , lsp-test , megaparsec + , neat-interpolation + , process , srcloc , tasty , tasty-hunit diff --git a/prelude/math.fut b/prelude/math.fut --- a/prelude/math.fut +++ b/prelude/math.fut @@ -637,7 +637,7 @@ def get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) def set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (!(1 intrinsics.<< bit))) | intrinsics.zext_i32_i64 (b intrinsics.<< bit)) + (x & (!(1i64 << i32 bit))) | (i32 b << i32 bit) def popc = intrinsics.popc64 def mul_hi a b = intrinsics.smul_hi64 (i64 a, i64 b) @@ -949,7 +949,7 @@ def get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) def set_bit (bit: i32) (x: t) (b: i32) = - ((x & i32 (!(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + (x & (!(1u64 << i32 bit))) | (i32 b << i32 bit) def popc x = intrinsics.popc64 (sign x) def mul_hi a b = unsign (intrinsics.umul_hi64 (sign a, sign b)) diff --git a/prelude/soacs.fut b/prelude/soacs.fut --- a/prelude/soacs.fut +++ b/prelude/soacs.fut @@ -162,30 +162,6 @@ def scan [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a) : *[n]a = intrinsics.scan op ne as --- | Split an array into those elements that satisfy the given --- predicate, and those that do not. --- --- **Work:** *O(n ✕ W(p))* --- --- **Span:** *O(log(n) ✕ W(p))* -def partition [n] 'a (p: a -> bool) (as: [n]a) : ?[k].([k]a, [n - k]a) = - let p' x = if p x then 0 else 1 - let (as', is) = intrinsics.partition 2 p' as - in (as'[0:is[0]], as'[is[0]:n]) - --- | Split an array by two predicates, producing three arrays. --- --- **Work:** *O(n ✕ (W(p1) + W(p2)))* --- --- **Span:** *O(log(n) ✕ (W(p1) + W(p2)))* -def partition2 [n] 'a (p1: a -> bool) (p2: a -> bool) (as: [n]a) : ?[k][l].([k]a, [l]a, [n - k - l]a) = - let p' x = if p1 x then 0 else if p2 x then 1 else 2 - let (as', is) = intrinsics.partition 3 p' as - in ( as'[0:is[0]] - , as'[is[0]:is[0] + is[1]] :> [is[1]]a - , as'[is[0] + is[1]:n] :> [n - is[0] - is[1]]a - ) - -- | Return `true` if the given function returns `true` for all -- elements in the array. -- @@ -227,6 +203,17 @@ def scatter 't [k] [n] (dest: *[k]t) (is: [n]i64) (vs: [n]t) : *[k]t = intrinsics.scatter dest is vs +-- | Exclusive prefix scan. Has the same caveats with respect to +-- associativity and complexity as `scan`. +-- +-- **Work:** *O(n ✕ W(op))* +-- +-- **Span:** *O(log(n) ✕ W(op))* +def exscan [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a) : *[n]a = + scatter (map (\_ -> ne) (0..1.. bool) (as: [n]a) : *[]a = - let flags = map (\x -> if p x then 1 else 0) as - let offsets = scan (+) 0 flags - let m = if n == 0 then 0 else offsets[n - 1] - in scatter (#[scratch] map (\x -> x) as[:m]) - (map2 (\f o -> if f == 1 then o - 1 else -1) flags offsets) - as + let flags = map p as + let offsets = scan (+) 0 (map intrinsics.btoi_bool_i64 flags) + let flags' = map p as + let is = map2 (\f o -> if f then o - 1 else -1) flags' offsets + -- This following is carefully written such that the two scatters will be + -- fused horisontally, which allows the entire thing to become a single + -- kernel. + let result = scatter (#[scratch] [as][0]) is as + let count = + scatter [n] + (map (\j -> if j == n - 1 then 0 else -1) (0..1.. bool) (as: [n]a) : ?[k].([k]a, [n - k]a) = + let offset = + reduce_comm (+) 0 (map (\x -> intrinsics.btoi_bool_i64 (p x)) as) + let add2 (a0, b0) (a1, b1) = (a0 + a1, b0 + b1) + let to_index f (o0, o1) = if f then o0 - 1i64 else offset + o1 - 1 + let t_flags = map p as + let f_flags = map (\x -> !x) t_flags + let flags = + map2 (\x y -> + ( intrinsics.btoi_bool_i64 x + , intrinsics.btoi_bool_i64 y + )) + t_flags + f_flags + let offsets = scan add2 (0, 0) flags + let idxs = map2 to_index (map p as) offsets + let res = scatter (#[scratch] [as][0]) idxs as + in (res[0:offset], res[offset:n]) + +-- | Split an array by two predicates, producing three arrays. +-- +-- **Work:** *O(n ✕ (W(p1) + W(p2)))* +-- +-- **Span:** *O(log(n) ✕ (W(p1) + W(p2)))* +def partition2 [n] 'a (p1: a -> bool) (p2: a -> bool) (as: [n]a) : ?[k][l].([k]a, [l]a, [n - k - l]a) = + let t1_flags = map p1 as + let t2_flags = map2 (\t a -> !t && p2 a) t1_flags as + let offset0 = reduce_comm (+) 0 (map intrinsics.btoi_bool_i64 t1_flags) + let offset1 = reduce_comm (+) 0 (map intrinsics.btoi_bool_i64 t2_flags) + let add3 (a0, b0, c0) (a1, b1, c1) = (a0 + a1, b0 + b1, c0 + c1) + let to_index f0 f1 (o0, o1, o2) = + if f0 + then o0 - 1i64 + else if f1 + then offset0 + o1 - 1 + else offset0 + offset1 + o2 - 1 + let t1_flags' = map p1 as + let t2_flags' = map2 (\t a -> !t && p2 a) t1_flags' as + let t3_flags' = map2 (\x y -> !(x || y)) t1_flags' t2_flags' + let flags' = + map3 (\x y z -> + ( intrinsics.btoi_bool_i64 x + , intrinsics.btoi_bool_i64 y + , intrinsics.btoi_bool_i64 z + )) + t1_flags' + t2_flags' + t3_flags' + let offsets = scan add3 (0, 0, 0) flags' + let idxs = map3 to_index t1_flags' t2_flags' offsets + let res = scatter (#[scratch] [as][0]) idxs as + in ( res[0:offset0] + , res[offset0:offset0 + offset1] :> [offset1]a + , res[offset0 + offset1:n] :> [n - offset0 - offset1]a + ) diff --git a/rts/c/backends/hip.h b/rts/c/backends/hip.h --- a/rts/c/backends/hip.h +++ b/rts/c/backends/hip.h @@ -309,6 +309,12 @@ (device_query(ctx->dev, hipDeviceAttributePhysicalMultiProcessorCount) * device_query(ctx->dev, hipDeviceAttributeMaxThreadsPerMultiProcessor)) / cfg->gpu.default_block_size; + + // XXX: this is a hack due to the inability of two-pass scan to handle a + // grid size that is larger than the maximum block size. + if (cfg->gpu.default_grid_size > ctx->max_thread_block_size) { + cfg->gpu.default_grid_size = ctx->max_thread_block_size; + } } for (int i = 0; i < NUM_TUNING_PARAMS; i++) { diff --git a/rts/c/backends/opencl.h b/rts/c/backends/opencl.h --- a/rts/c/backends/opencl.h +++ b/rts/c/backends/opencl.h @@ -811,6 +811,12 @@ ctx->max_shared_memory = max_shared_memory; } + // XXX: this is a hack due to the inability of two-pass scan to handle a + // grid size that is larger than the maximum block size. + if (ctx->cfg->gpu.default_grid_size > max_thread_block_size) { + ctx->cfg->gpu.default_grid_size = max_thread_block_size; + } + // Now we go through all the sizes, clamp them to the valid range, // or set them to the default. for (int i = 0; i < NUM_TUNING_PARAMS; i++) { diff --git a/rts/c/server.h b/rts/c/server.h --- a/rts/c/server.h +++ b/rts/c/server.h @@ -17,8 +17,9 @@ typedef int (*restore_fn)(const void*, FILE*, struct futhark_context*, void*); typedef void (*store_fn)(const void*, FILE*, struct futhark_context*, void*); typedef int (*free_fn)(const void*, struct futhark_context*, void*); -typedef void* (*array_new_fn)(struct futhark_context *, const void*, const int64_t*); -typedef const int64_t* (*array_shape_fn)(struct futhark_context*, void*); +typedef int (*array_new_fn)(struct futhark_context *, void**, const void*, const int64_t*); +typedef int (*array_set_fn)(struct futhark_context *, const void*, const void*, const int64_t*); +typedef const int64_t* (*array_shape_fn)(struct futhark_context*, const void*); typedef int (*array_index_fn)(struct futhark_context*, void*, const void*, const int64_t*); typedef int (*project_fn)(struct futhark_context*, void*, const void*); typedef int (*variant_fn)(struct futhark_context*, const void*); @@ -37,6 +38,7 @@ int rank; const struct type *element_type; array_new_fn new; + array_set_fn set; array_shape_fn shape; array_index_fn index; }; @@ -713,9 +715,8 @@ const int64_t *shape = a->shape(s->ctx, v->value.value.v_ptr); for (int i = 0; i < a->rank; ++i) { - printf("%lld ", shape[i]); + printf("%lld\n", (long long)shape[i]); } - printf("\n"); } void cmd_elemtype(struct server_state *s, const char *args[]) { @@ -733,6 +734,164 @@ printf("%s\n", a->element_type->name); } +void cmd_rank(struct server_state *s, const char *args[]) { + const char *type = get_arg(args, 0); + const struct type *t = get_type(s, type); + + if (t->kind != ARRAY) { + failure(); + printf("Not an array type\n"); + return; + } + + const struct array *a = t->info; + printf("%d\n", a->rank); +} + +void cmd_new_array(struct server_state *s, const char *args[]) { + const char *to_name = get_arg(args, 0); + const char *type_name = get_arg(args, 1); + const struct type *type = get_type(s, type_name); + struct variable *to = create_variable(s, to_name, type); + + if (to == NULL) { + failure(); + printf("Variable already exists: %s\n", to_name); + return; + } + + if (type->kind != ARRAY) { + failure(); + printf("Not an array type\n"); + return; + } + + const struct array *a = type->info; + + int num_args = 0; + for (int i = 2; arg_exists(args, i); i++) { + num_args++; + } + + if (num_args < a->rank) { + failure(); + printf("Expected %d dimensions, but got %d.\n", a->rank, num_args); + return; + } + + int64_t* dims = alloca(a->rank * sizeof(int64_t)); + int64_t n_values = 1; + + for (int i = 0; i < a->rank; ++i) { + const char *size_arg = get_arg(args, 2+i); + char* end; + errno = 0; + int64_t size = strtoll(size_arg, &end, 10); + + if (errno == ERANGE || *end != '\0' || size < 0) { + failure(); + printf("Invalid size `%s` of dimension %d.\n", size_arg, i+1); + return; + } + + dims[i] = size; + n_values *= size; + } + + if (num_args - a->rank != n_values) { + failure(); + printf("Expected %d values, but got %d.\n", (int)n_values, num_args - a->rank); + return; + } + + const void** value_ptrs = alloca(n_values * sizeof(void*)); + + for (int64_t i = 0; i < n_values; i++) { + struct variable* v = get_variable(s, args[2+a->rank+i]); + + if (v == NULL) { + failure(); + printf("Unknown variable: %s\n", args[2+a->rank+i]); + return; + } + + if (strcmp(v->value.type->name, a->element_type->name) != 0) { + failure(); + printf("Value %d mismatch: expected type %s, got %s\n", + (int)i, a->element_type->name, v->value.type->name); + return; + } + + value_ptrs[i] = value_ptr(&v->value); + } + + a->new(s->ctx, value_ptr(&to->value), value_ptrs, dims); +} + +void cmd_set(struct server_state *s, const char *args[]) { + const char *arr_name = get_arg(args, 0); + const char *val_name = get_arg(args, 1); + struct variable* arr = get_variable(s, arr_name); + struct variable* val = get_variable(s, val_name); + + if (arr == NULL) { + failure(); + printf("Unknown variable: %s\n", arr_name); + return; + } + if (val == NULL) { + failure(); + printf("Unknown variable: %s\n", val_name); + return; + } + + if (arr->value.type->kind != ARRAY) { + failure(); + printf("Not an array type\n"); + return; + } + + const struct array *a = arr->value.type->info; + + if (strcmp(val->value.type->name, a->element_type->name) != 0) { + failure(); + printf("Type mismatch: expected element of type %s, got %s\n", + a->element_type->name, val->value.type->name); + return; + } + + for (int i = 0; ; ++i) { + if (!arg_exists(args, 2+i)) { + if (i != a->rank) { + failure(); + printf("%d indices expected but %d values provided.\n", a->rank, i); + return; + } + break; + } + } + + const int64_t *shape = a->shape(s->ctx, arr->value.value.v_ptr); + int64_t* indices = alloca(a->rank * sizeof(int64_t)); + + for (int i = 0; i < a->rank; ++i) { + const char *idx_arg = get_arg(args, 2+i); + char* end; + errno = 0; + int64_t idx = strtoll(idx_arg, &end, 10); + + if (errno == ERANGE || *end != '\0' || idx < 0 || idx >= shape[i]) { + failure(); + printf("Invalid index `%s` on dimension %d.\n", idx_arg, i+1); + return; + } + + indices[i] = idx; + } + + a->set(s->ctx, arr->value.value.v_ptr, value_ptr(&val->value), indices); +} + void cmd_index(struct server_state *s, const char *args[]) { const char *to_name = get_arg(args, 0); const char *from_name = get_arg(args, 1); @@ -752,8 +911,6 @@ const struct array *a = from->value.type->info; - const int64_t *shape = a->shape(s->ctx, from->value.value.v_ptr); - int64_t* indices = alloca(a->rank * sizeof(int64_t)); for (int i = 0; ; ++i) { if (!arg_exists(args, 2+i)) { if (i != a->rank) { @@ -764,6 +921,10 @@ break; } } + + const int64_t *shape = a->shape(s->ctx, from->value.value.v_ptr); + int64_t* indices = alloca(a->rank * sizeof(int64_t)); + for (int i = 0; i < a->rank; ++i) { const char *idx_arg = get_arg(args, 2+i); char* end; @@ -1182,6 +1343,12 @@ cmd_shape(s, tokens+1); } else if (strcmp(command, "elemtype") == 0) { cmd_elemtype(s, tokens+1); + } else if (strcmp(command, "rank") == 0) { + cmd_rank(s, tokens+1); + } else if (strcmp(command, "new_array") == 0) { + cmd_new_array(s, tokens+1); + } else if (strcmp(command, "set") == 0) { + cmd_set(s, tokens+1); } else if (strcmp(command, "index") == 0) { cmd_index(s, tokens+1); } else if (strcmp(command, "fields") == 0) { @@ -1242,17 +1409,20 @@ // The aux struct lets us write generic method implementations without // code duplication. -typedef int (*array_values_fn)(struct futhark_context*, void*, void*); -typedef int (*array_free_fn)(struct futhark_context*, void*); +typedef void* (*aux_array_new_fn)(struct futhark_context*, const void**, const int64_t*); +typedef const int64_t* (*aux_array_shape_fn)(struct futhark_context*, void*); +typedef int (*aux_array_index_fn)(struct futhark_context*, void*, const void*, const int64_t*); +typedef int (*aux_array_values_fn)(struct futhark_context*, void*, void*); +typedef int (*aux_array_free_fn)(struct futhark_context*, void*); struct array_aux { int rank; const struct primtype_info_t* info; const char *name; - array_new_fn new; - array_shape_fn shape; - array_values_fn values; - array_free_fn free; + aux_array_new_fn new; + aux_array_shape_fn shape; + aux_array_values_fn values; + aux_array_free_fn free; }; int restore_array(const struct array_aux *aux, FILE *f, diff --git a/rts/javascript/server.js b/rts/javascript/server.js --- a/rts/javascript/server.js +++ b/rts/javascript/server.js @@ -16,7 +16,10 @@ 'pause_profiling', 'unpause_profiling', 'report', - 'rename' + 'rename', + 'types', + 'fields', + 'project' ]; } @@ -97,6 +100,63 @@ delete this._types[oldname]; } + _cmd_types(args) { + var types = this.ctx.get_types(); + for (var t in types) { + console.log(t); + } + } + + _cmd_fields(args) { + var type_name = this._get_arg(args, 0); + var types = this.ctx.get_types(); + var type_info = types[type_name]; + if (!type_info || type_info[0] !== "record") { + throw "Not a record type: " + type_name; + } + var fields = type_info[1]; + for (var i = 0; i < fields.length; i++) { + console.log(fields[i][0] + " " + fields[i][1]); + } + } + + _cmd_project(args) { + var to_name = this._get_arg(args, 0); + var from_name = this._get_arg(args, 1); + var field_name = this._get_arg(args, 2); + + if (to_name in this._vars) { + throw "Variable already exists: " + to_name; + } + + var from_val = this._get_var(from_name); + var from_type = this._get_type(from_name); + + var types = this.ctx.get_types(); + var type_info = types[from_type]; + if (!type_info || type_info[0] !== "record") { + throw "Not a record type: " + from_type; + } + + var fields = type_info[1]; + var field_info = null; + for (var i = 0; i < fields.length; i++) { + if (fields[i][0] === field_name) { + field_info = fields[i]; + break; + } + } + + if (field_info === null) { + throw "No such field: " + field_name; + } + + var field_type = field_info[1]; + var project_fn = field_info[2]; + var result = this.ctx[project_fn](from_val); + this._set_var(to_name, result, field_type); + } + _cmd_call(args) { var entry = this._get_entry_point(this._get_arg(args, 0)); var num_ins = entry[1].length; @@ -221,6 +281,9 @@ case 'unpause_profiling': this._cmd_dummy(args); break case 'report': this._cmd_dummy(args); break case 'rename': this._cmd_rename(args); break + case 'types': this._cmd_types(args); break + case 'fields': this._cmd_fields(args); break + case 'project': this._cmd_project(args); break } } else { throw "Unknown command: " + cmd; diff --git a/src-testing/Futhark/IR/SyntaxTests.hs b/src-testing/Futhark/IR/SyntaxTests.hs --- a/src-testing/Futhark/IR/SyntaxTests.hs +++ b/src-testing/Futhark/IR/SyntaxTests.hs @@ -31,3 +31,6 @@ instance IsString SubExpRes where fromString = parseString "SubExpRes" parseSubExpRes + +instance IsString Ident where + fromString = parseString "Ident" parseIdent diff --git a/src-testing/Futhark/LspTests.hs b/src-testing/Futhark/LspTests.hs new file mode 100644 --- /dev/null +++ b/src-testing/Futhark/LspTests.hs @@ -0,0 +1,365 @@ +{-# LANGUAGE DuplicateRecordFields #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE RecordWildCards #-} + +module Futhark.LspTests (tests, main) where + +import Colog.Core (LogAction (LogAction)) +import Control.Concurrent (forkIO) +import Control.Lens ((^.)) +import Control.Monad (zipWithM_) +import Control.Monad.IO.Class (liftIO) +import Data.Foldable (for_) +import Data.Functor (void) +import Data.Maybe (mapMaybe) +import Data.Text (Text) +import Data.Text qualified as T +import Futhark.CLI.LSP (serverDefinition) +import Futhark.Fmt.Printer (fmtToText) +import Language.Futhark.Parser.Monad (SyntaxError (..)) +import Language.LSP.Protocol.Lens (uri) +import Language.LSP.Protocol.Message (SMethod (SMethod_WorkspaceApplyEdit)) +import Language.LSP.Protocol.Types + ( CodeAction (..), + CodeLens (CodeLens), + Definition (Definition), + FormattingOptions (FormattingOptions), + Hover (Hover), + InlayHint (..), + InlayHintKind (InlayHintKind_Type), + LanguageKind (LanguageKind_Custom), + Location (..), + MarkupContent (..), + MarkupKind (MarkupKind_PlainText), + Position (..), + Range (..), + TextDocumentIdentifier, + type (|?) (..), + ) +import Language.LSP.Server (runServerWithHandles) +import Language.LSP.Test + ( Session, + createDoc, + defaultConfig, + documentContents, + executeCodeAction, + executeCommand, + formatDoc, + fullLatestClientCaps, + getAndResolveCodeLenses, + getAndResolveInlayHints, + getCodeActions, + getDefinitions, + getHover, + message, + runSessionWithHandles, + ) +import NeatInterpolation (text) +import System.IO (hClose) +import System.Process (createPipe) +import Test.Tasty (TestName, TestTree, defaultMain, testGroup) +import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) + +runServerSessionPair :: Session a -> IO a +runServerSessionPair session = do + (serverIn, clientOut) <- createPipe + (clientIn, serverOut) <- createPipe + futharkServer <- serverDefinition + let ignoreLog :: (Applicative m) => LogAction m msg + ignoreLog = LogAction . const $ pure () + + _serverTid <- + forkIO . void $ + runServerWithHandles ignoreLog ignoreLog serverIn serverOut futharkServer + + result <- runSessionWithHandles clientOut clientIn defaultConfig fullLatestClientCaps "." session + + for_ [serverIn, clientOut, clientIn, serverOut] hClose + + pure result + +createMainDoc :: Text -> Session TextDocumentIdentifier +createMainDoc = createDoc "main.fut" futharkLanguage + where + futharkLanguage :: LanguageKind + futharkLanguage = LanguageKind_Custom "Futhark" + +serverTestCase :: TestName -> Session () -> TestTree +serverTestCase testName = testCase testName . runServerSessionPair + +testHoverInformation :: TestTree +testHoverInformation = serverTestCase "Hover" $ do + docIdent <- createMainDoc "def main = let xyz = 42i32 in xyz" + hoverInfo <- getHover docIdent $ Position {_character = 31, _line = 0} + case hoverInfo of + Nothing -> liftIO $ assertFailure "Did not receive any hover information" + Just (Hover content range) -> do + liftIO $ range @?= Just expectedRange + liftIO $ content @?= InL expectedMarkup + where + expectedMarkup = + MarkupContent + { _value = "i32", + _kind = MarkupKind_PlainText + } + expectedRange = + Range + { _start = Position {_line = 0, _character = 30}, + _end = Position {_line = 0, _character = 34} + } + +testDefinition :: TestTree +testDefinition = serverTestCase "Go To Definition" $ do + docIdent <- createMainDoc mainContents + definition <- getDefinitions docIdent fooBodyPosition + + let expectedDefinition = + Definition $ + InL $ + Location + { _uri = docIdent ^. uri, + _range = + Range + { _end = Position 0 15, + _start = Position 0 0 + } + } + + liftIO $ definition @?= InL expectedDefinition + where + mainContents = + [text| + def foo = 0i32 + def bar = foo + |] + + fooBodyPosition :: Position + fooBodyPosition = + Position + { _line = 1, + _character = 12 + } + +testFormatting :: TestTree +testFormatting = serverTestCase "Formatting" $ do + docIdent <- createMainDoc mainContents + formatDoc docIdent formattingOptions + lspFormattedDoc <- documentContents docIdent + formattedDoc <- case fmtToText "main.fut" mainContents of + Left (SyntaxError loc msg) -> + liftIO . assertFailure $ + "Formatting failed: " <> show loc <> ".\n" <> show msg + Right d -> pure d + + liftIO $ lspFormattedDoc @?= formattedDoc + where + -- these are ignored by the formatter anyway + formattingOptions = FormattingOptions 0 False Nothing Nothing Nothing + mainContents = + [text| + -- this is where all the lines start + def main = + 0i32 + |] + +main :: IO () +main = defaultMain testEvaluationComment + +testEvaluationComment :: TestTree +testEvaluationComment = serverTestCase "Evaluation Comment" $ do + docIdent <- createMainDoc mainContents + lensCommand <- + getAndResolveCodeLenses docIdent >>= \case + [CodeLens _ (Just command) _] -> pure command + bad -> liftIO $ assertFailure $ "Unexpected Code Lenses: " <> show bad + + executeCommand lensCommand + + -- this is important, it crashes with an IOError otherwise + _workspaceEdit <- message SMethod_WorkspaceApplyEdit + + newContents <- documentContents docIdent + liftIO $ newContents @?= expectedContents + where + expectedContents = + mainContents + <> "-- 47\n" + mainContents = + [text| + def x = 42i32 + -- >>> x + 5 + |] + +fullRange :: Range +fullRange = + Range + { _start = Position minBound minBound, + _end = Position maxBound maxBound + } + +testInlayTypeHint :: TestTree +testInlayTypeHint = + testGroup + "Inlay type hint" + [ defParamHint, + letHint, + noSizeHint, + loopHint, + defReturnHint, + defConstantReturnHint, + lambdaArgHint, + typeArgHint, + defNestedHint + ] + where + expectHint InlayHint {..} l p = do + _position @?= p + _label @?= InL l + _kind @?= Just InlayHintKind_Type + hintTestCase name mainDoc expectedHints = serverTestCase name $ do + docIdent <- createMainDoc mainDoc + hints <- getAndResolveInlayHints docIdent fullRange + liftIO $ + if length hints /= length expectedHints + then assertFailure $ "Unexpected inlay hints: " ++ show hints + else zipWithM_ (\h (p, l) -> expectHint h l p) hints expectedHints + typeArgHint = + hintTestCase + "type argument hint" + "def identity = let f 'a (x: a): a = x in f" + [(Position 0 12, " 'a\8320"), (Position 0 12, " : (x: a\8320) -> a\8320")] + lambdaArgHint = + hintTestCase + "lambda argument hint" + "def lambda : i32 -> i32 = \\ x -> x + 0i32" + [(Position 0 28, "("), (Position 0 29, ": i32)")] + defConstantReturnHint = + hintTestCase + "def constant hint" + "def pi = 3" + [(Position 0 6, " : i32")] + defReturnHint = + hintTestCase + "def return hint" + "def twice (x: i32) = x + x" + [(Position 0 18, " : i32")] + loopHint = + hintTestCase + "loop hint" + "def factorial (n: i32) : i32 = loop result = 1 for i < n do result * (i + 1)" + [(Position 0 36, "("), (Position 0 42, ": i32)")] + -- this is on purpose, all sizes have type i64 + noSizeHint = + hintTestCase + "no size hint" + "def sz [n] 'a (_: [n]a) : i64 = n" + [] + letHint = + hintTestCase + "let hint" + "def foo : bool = let f = false in true && f" + [(Position 0 22, ": bool")] + defParamHint = + hintTestCase + "def param hint" + "def plus5 x : i32 = x + 5i32" + [(Position 0 10, "("), (Position 0 11, ": i32)")] + defNestedHint = + hintTestCase + "def nested hint" + "def bar 'a (f: a -> a -> a) (x : a, y) : a = f x y" + [(Position 0 37, ": a")] + +testTypeCompletionCodeAction :: TestTree +testTypeCompletionCodeAction = + testGroup + "Type completion code action" + [ testDefReturn, + testDefParam, + testDefParamTypeArg, + testNestedTypeArg, + testLetType, + testLetArg, + testDefDim, + testDefConstantReturnPosition + ] + where + -- Retrieves code actions, selects the one containing the title phrase + -- and checks that the new document looks like expected + actionTestCase name contents actionTitleLike newContents = + serverTestCase name $ do + docIdent <- createMainDoc contents + actions <- getCodeActions docIdent fullRange + let filterAction = \case + InL _ -> Nothing + InR c@(CodeAction {..}) + | actionTitleLike `T.isInfixOf` _title -> Just c + | otherwise -> Nothing + selectedAction <- liftIO $ case mapMaybe filterAction actions of + [a] -> pure a + rest -> assertFailure $ "No single action found: " ++ show rest + + executeCodeAction selectedAction + documentContents docIdent >>= liftIO . (@?= newContents) + testDefReturn = + actionTestCase + "def return" + "def pi = 3i32" + "return type" + "def pi : i32 = 3i32" + testDefParam = + actionTestCase + "param type" + "def identity 'a x : a = x" + "Insert type" + "def identity 'a (x: a) : a = x" + testDefParamTypeArg = + actionTestCase + "param type arg" + "def identity x = x" + "Insert type" + "def identity '^t0 (x: t0) = x" + testNestedTypeArg = + actionTestCase + "nested type arg" + "def identity x = let y = x in y" + "Insert type: `y" + "def identity '^t0 x = let y: t0 = x in y" + testLetType = + actionTestCase + "let type" + "def inc x = let y = x + 1i32 in y" + "Insert type: `y" + "def inc x = let y: i32 = x + 1i32 in y" + testLetArg = + actionTestCase + "let arg" + "def inc x = let f y = y + 1i32 in f x" + "Insert type: `(y" + "def inc x = let f (y: i32) = y + 1i32 in f x" + testDefDim = + actionTestCase + "def dim" + "def id'arr 'a (xs: []a) = xs" + "return type" + "def id'arr [d0] 'a (xs: []a) : [d0]a = xs" + testDefConstantReturnPosition = + actionTestCase + "def constant return position" + "def it '^t = \\(x: t) -> x" + "return type" + "def it '^t : (x: t) -> t = \\(x: t) -> x" + +tests :: TestTree +tests = + testGroup + "Futhark.LspTests" + [ testHoverInformation, + testDefinition, + testFormatting, + testEvaluationComment, + testInlayTypeHint, + testTypeCompletionCodeAction + ] diff --git a/src-testing/Futhark/Optimise/Fusion/ScremaTests.hs b/src-testing/Futhark/Optimise/Fusion/ScremaTests.hs new file mode 100644 --- /dev/null +++ b/src-testing/Futhark/Optimise/Fusion/ScremaTests.hs @@ -0,0 +1,1270 @@ +module Futhark.Optimise.Fusion.ScremaTests (tests) where + +import Control.Monad.Reader +import Control.Monad.State +import Data.String (fromString) +import Futhark.Analysis.HORep.SOAC as SOAC +import Futhark.FreshNames +import Futhark.IR.SOACS +import Futhark.IR.SOACSTests () +import Futhark.Optimise.Fusion.Screma +import Test.Tasty +import Test.Tasty.HUnit + +withFreshNamesScopeError :: ReaderT (Scope SOACS) (StateT VNameSource Maybe) a -> Maybe a +withFreshNamesScopeError m = + evalStateT (runReaderT m mempty) (newNameSource 10000) + +withFreshNames :: State VNameSource a -> a +withFreshNames m = + evalState m (newNameSource 10000) + +fromLines :: [String] -> Lambda SOACS +fromLines = fromString . unlines + +-- | A wrapper that makes 'show' behave like 'prettyString'. +newtype P a = P a + deriving (Eq, Ord) + +instance (Pretty a) => Show (P a) where + show (P x) = prettyString x + +splitLambdaByParTester :: [VName] -> Lambda SOACS -> P (Lambda SOACS, Lambda SOACS) +splitLambdaByParTester names lam = P (lam_x', lam_y') + where + Just ((_, lam_x', _), (_, lam_y', _)) = + withFreshNamesScopeError $ + splitLambdaByPar names (lambdaParams lam) lam (lambdaReturnType lam) + +splitLambdaByParTests :: TestTree +splitLambdaByParTests = + testGroup + "splitLambdaByPar" + [ testCase "keeps params and result" $ + let lam = "\\{x_0 : i32, x_1 : i32} : {i32, i32} -> {x_0, x_1}" + lam_x = "\\{x_0 : i32} : {i32} -> {x_0}" + lam_y = "\\{x_1 : i32} : {i32} -> {x_1}" + names = ["x_0"] + in splitLambdaByParTester names lam @?= P (lam_x, lam_y), + testCase "keeps computation in first lambda" $ + let lam = + fromLines + [ "\\{x_0 : i32, x_1 : i32} : {i32, i32} ->", + " let {x_2 : i32} = add32(x_0, x_1)", + " in {x_0, x_2}" + ] + lam_x = + fromLines + [ "\\{x_0 : i32, x_1 : i32} : {i32} ->", + " let {x_2 : i32} = add32(x_0, x_1)", + " in {x_2}" + ] + lam_y = "\\{x_0 : i32} : {i32} -> {x_0}" + names = ["x_1"] + in splitLambdaByParTester names lam @?= P (lam_x, lam_y), + testCase "keeps computations in both lambdas" $ + let lam = + fromLines + [ "\\{x_0 : i32, x_1 : i32} : {i32, i32} -> ", + " let {x_2 : i32} = add32(x_0, x_1) ", + " let {x_3 : i32} = add32(1i32, x_0) ", + " in {x_3, x_2}" + ] + lam_x = + fromLines + [ "\\{x_0 : i32, x_1 : i32} : {i32} -> ", + " let {x_2 : i32} = add32(x_0, x_1) ", + " in {x_2}" + ] + lam_y = + fromLines + [ "\\{x_0 : i32} : {i32} -> ", + " let {x_3 : i32} = add32(1i32, x_0) ", + " in {x_3}" + ] + names = ["x_1"] + in splitLambdaByParTester names lam @?= P (lam_x, lam_y), + testCase "keeps line order" $ + let lam = + fromLines + [ "\\{x_0 : i32, x_1 : i32} : {i32, i32} -> ", + " let {x_3 : i32} = add32(1i32, x_0) ", + " let {x_2 : i32} = add32(x_0, x_1) ", + " let {x_4 : i32} = add32(1i32, x_3) ", + " in {x_4, x_2}" + ] + lam_x = + fromLines + [ "\\{x_0 : i32, x_1 : i32} : {i32} -> ", + " let {x_2 : i32} = add32(x_0, x_1) ", + " in {x_2}" + ] + lam_y = + fromLines + [ "\\{x_0 : i32} : {i32} -> ", + " let {x_3 : i32} = add32(1i32, x_0) ", + " let {x_4 : i32} = add32(1i32, x_3) ", + " in {x_4}" + ] + names = ["x_1"] + in splitLambdaByParTester names lam @?= P (lam_x, lam_y), + testCase "does redundant work" $ + let lam = + fromLines + [ "\\{x_0 : i32, x_1 : i32} : {i32, i32} -> ", + " let {x_2 : i32} = add32(1i32, x_0) ", + " let {x_3 : i32} = add32(x_1, x_2) ", + " in {x_3, x_2}" + ] + lam_x = + fromLines + [ "\\{x_0 : i32, x_1 : i32} : {i32} -> ", + " let {x_2 : i32} = add32(1i32, x_0) ", + " let {x_3 : i32} = add32(x_1, x_2) ", + " in {x_3}" + ] + lam_y = + fromLines + [ "\\{x_0 : i32} : {i32} -> ", + " let {x_2 : i32} = add32(1i32, x_0) ", + " in {x_2}" + ] + names = ["x_1"] + in splitLambdaByParTester names lam @?= P (lam_x, lam_y) + ] + +fuseSuperScremaTests :: TestTree +fuseSuperScremaTests = + testGroup + "fuseSuperScrema" + [ testCase "map-scan (vertical)" $ + let scan_op = + Scan + ( fromLines + [ "\\ {eta_p_5571 : i32, eta_p_5572 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_5571, eta_p_5572)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + ident_a = "input_a_5565 : [d_5537]i32" + ident_b = "input_b_5538 : [d_5537]i32" + input_a = SOAC.identInput ident_a + input_b = SOAC.identInput ident_b + in P + ( withFreshNames + ( fuseSuperScrema + "d_5537" + [input_a] + ( ScremaForm + ( fromLines + [ "\\ {eta_p_5566 : i32} : {i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567}" + ] + ) + [] + [] + "\\ {x_5568 : i32} : {i32} -> {x_5568}" + ) + [identName ident_b] + [input_b] + ( ScremaForm + "\\ {x_5570 : i32} : {i32} -> {x_5570}" + [scan_op] + [] + "\\ {x_5574 : i32} : {i32} -> {x_5574}" + ) + ["defunc_0_scan_res_5569"] + ) + ) + @?= P + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {eta_p_5566 : i32}: {i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5568 : i32}: {i32, i32} -> ", + "let {x_5570 : i32} = x_5568", + "in {x_5570, x_5568}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_5574 : i32, x_10000 : i32}: {i32, i32} ->", + "{x_5574, x_10000}" + ] + ), + ["defunc_0_scan_res_5569", identName ident_b] + ), + testCase "map-scan (horizontal)" $ + let scan_op = + Scan + ( fromLines + [ "\\ {eta_p_5571 : i32, eta_p_5572 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_5571, eta_p_5572)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + ident_a = "input_a_5565 : [d_5537]i32" + ident_b = "input_b_5538 : [d_5537]i32" + input_a = SOAC.identInput ident_a + input_b = SOAC.identInput ident_b + out_a = "out_a_5564145" + out_b = "out_b_5534156" + in P + ( withFreshNames + ( fuseSuperScrema + "d_5537" + [input_a] + ( ScremaForm + ( fromLines + [ "\\ {eta_p_5566 : i32} : {i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567}" + ] + ) + [] + [] + "\\ {x_5568 : i32} : {i32} -> {x_5568}" + ) + [out_a] + [input_b] + ( ScremaForm + "\\ {x_5570 : i32} : {i32} -> {x_5570}" + [scan_op] + [] + "\\ {x_5574 : i32} : {i32} -> {x_5574}" + ) + [out_b] + ) + ) + @?= P + ( SuperScrema + "d_5537" + [input_a, input_b] + ( fromLines + [ "\\ {eta_p_5566 : i32, x_10000 : i32}: {i32, i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567, x_10000}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5568 : i32, x_5570 : i32}: {i32, i32} -> ", + "{x_5570, x_5568}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_5574 : i32, x_10001 : i32}: {i32, i32} ->", + "{x_5574, x_10001}" + ] + ), + [out_b, out_a] + ), + testCase "map-scan (vertical) with reduce (horizontal)" $ + let scan_op = + Scan + ( fromLines + [ "\\ {eta_p_5571 : i32, eta_p_5572 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_5571, eta_p_5572)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + reduce_op = + Reduce + Commutative + ( fromLines + [ "\\ {eta_p_55720 : i32, eta_p_557201 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_55720, eta_p_557201)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + + ident_a = "input_a_5565 : [d_5537]i32" + ident_b = "input_b_5538 : [d_5537]i32" + input_a = SOAC.identInput ident_a + input_b = SOAC.identInput ident_b + in P + ( withFreshNames + ( fuseSuperScrema + "d_5537" + [input_a] + ( ScremaForm + ( fromLines + [ "\\ {eta_p_5566 : i32} : {i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567, lifted_lambda_res_5567}" + ] + ) + [] + [reduce_op] + "\\ {x_5568 : i32} : {i32} -> {x_5568}" + ) + ["red_out_543532", identName ident_b] + [input_b] + ( ScremaForm + "\\ {x_5570 : i32} : {i32} -> {x_5570}" + [scan_op] + [] + ( fromLines + [ "\\ {x_5574 : i32} : {i32} ->", + "let {y_5567 : i32} = add32(2i32, x_5574)", + "in {y_5567}" + ] + ) + ) + ["defunc_0_scan_res_5569"] + ) + ) + @?= P + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {eta_p_5566 : i32}: {i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567, lifted_lambda_res_5567}" + ] + ) + [] + [reduce_op] + ( fromLines + [ "\\ {x_5568 : i32}: {i32, i32} -> ", + "let {x_5570 : i32} = x_5568", + "in {x_5570, x_5568}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_5574 : i32, x_10000 : i32}: {i32, i32} ->", + "let {y_5567 : i32} = add32(2i32, x_5574)", + "in {y_5567, x_10000}" + ] + ), + ["red_out_543532", "defunc_0_scan_res_5569", identName ident_b] + ), + testCase "map-map (vertical)" $ + let ident_a = "input_a_5565 : [d_5537]i32" + ident_b = "input_b_5538 : [d_5537]i32" + input_a = SOAC.identInput ident_a + input_b = SOAC.identInput ident_b + in P + ( withFreshNames + ( fuseSuperScrema + "d_5537" + [input_a] + ( ScremaForm + ( fromLines + [ "\\ {eta_p_5566 : i32} : {i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567, lifted_lambda_res_5567}" + ] + ) + [] + [] + "\\ {x_5568 : i32} : {i32} -> {x_5568}" + ) + [identName ident_b] + [input_b] + ( ScremaForm + ( fromLines + [ "\\ {x_5574 : i32} : {i32} ->", + "let {y_5567 : i32} = add32(3i32, x_5574)", + "in {y_5567}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5570 : i32} : {i32} -> {x_5570}" + ] + ) + ) + ["defunc_0_scan_res_5569"] + ) + ) + @?= P + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {eta_p_5566 : i32}: {i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567, lifted_lambda_res_5567}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5568 : i32}: {i32, i32} -> ", + "let {x_5574 : i32} = x_5568", + "let {y_5567 : i32} = add32(3i32, x_5574)", + "in {y_5567, x_5568}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5570 : i32, x_10000 : i32}: {i32, i32} ->", + "in {x_5570, x_10000}" + ] + ), + ["defunc_0_scan_res_5569", identName ident_b] + ), + testCase "map-scan-map (vertical)" $ + let scan_op = + Scan + ( fromLines + [ "\\ {eta_p_5571 : i32, eta_p_5572 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_5571, eta_p_5572)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + ident_a = "input_a_5565 : [d_5537]i32" + ident_b = "input_b_5538 : [d_5537]i32" + input_a = SOAC.identInput ident_a + input_b = SOAC.identInput ident_b + in P + ( withFreshNames + ( fuseSuperScrema + "d_5537" + [input_a] + ( ScremaForm + ( fromLines + [ "\\ {eta_p_5566 : i32} : {i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567}" + ] + ) + [] + [] + "\\ {x_5568 : i32} : {i32} -> {x_5568}" + ) + [identName ident_b] + [input_b] + ( ScremaForm + "\\ {x_5570 : i32} : {i32} -> {x_5570}" + [scan_op] + [] + ( fromLines + [ "\\ {x_5574 : i32} : {i32} ->", + "let {y_6363: i32} = add32(3i32, x_5574)", + "in {y_6363}" + ] + ) + ) + ["defunc_0_scan_res_5569"] + ) + ) + @?= P + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {eta_p_5566 : i32}: {i32} ->", + "let {lifted_lambda_res_5567 : i32} = add32(2i32, eta_p_5566)", + "in {lifted_lambda_res_5567}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5568 : i32}: {i32, i32} -> ", + "let {x_5570 : i32} = x_5568", + "in {x_5570, x_5568}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_5574 : i32, x_10000 : i32}: {i32, i32} ->", + "let {y_6363 : i32} = add32(3i32, x_5574)", + "in {y_6363, x_10000}" + ] + ), + ["defunc_0_scan_res_5569", identName ident_b] + ), + testCase "red-red fusion" $ + let reduce_op' = + Reduce + Commutative + ( fromLines + [ "\\ {a_0 : i32, b_1 : i32} : {i32} ->", + "let {c_2 : i32} = add32(a_0, b_1)", + "in {c_2}" + ] + ) + ["0i32"] + reduce_op = + Reduce + Commutative + ( fromLines + [ "\\ {a_4 : f32, b_5 : f32} : {f32} ->", + "let {c_6 : f32} = fadd32(a_4, b_5)", + "in {c_6}" + ] + ) + ["0.0f32"] + ident_a = "input_a_7 : [d_9]i32" + input_a = SOAC.identInput ident_a + ident_b = "input_b_8 : [d_9]f32" + input_b = SOAC.identInput ident_b + in P + ( withFreshNames + ( fuseSuperScrema + "d_9" + [input_a] + ( ScremaForm + ( fromLines + [ "\\ {x_11 : i32} : {i32} ->", + "in {x_11}" + ] + ) + [] + [reduce_op'] + (fromLines ["nilFn"]) + ) + ["out_a_12"] + [input_b] + ( ScremaForm + ( fromLines + [ "\\ {x_13 : f32} : {f32} ->", + "in {x_13}" + ] + ) + [] + [reduce_op] + (fromLines ["nilFn"]) + ) + ["out_b_14"] + ) + ) + @?= P + ( SuperScrema + "d_9" + [input_a, input_b] + ( fromLines + [ "\\ {x_11 : i32, x_10000 : f32}: {i32, f32} ->", + "{x_11, x_10000}" + ] + ) + [] + [reduce_op'] + ( fromLines + [ "\\ {x_13 : f32} : {f32} ->", + "{x_13}" + ] + ) + [] + [reduce_op] + (fromLines ["nilFn"]), + ["out_a_12", "out_b_14"] + ), + testCase "scan-scan fusion" $ + let scan_op' = + Scan + ( fromLines + [ "\\ {a_0 : i32, b_1 : i32} : {i32} ->", + "let {c_2 : i32} = add32(a_0, b_1)", + "in {c_2}" + ] + ) + ["0i32"] + scan_op = + Scan + ( fromLines + [ "\\ {a_4 : f32, b_5 : f32} : {f32} ->", + "let {c_6 : f32} = fadd32(a_4, b_5)", + "in {c_6}" + ] + ) + ["0.0f32"] + ident_a = "input_a_7 : [d_9]i32" + input_a = SOAC.identInput ident_a + ident_b = "input_b_8 : [d_9]f32" + input_b = SOAC.identInput ident_b + in P + ( withFreshNames + ( fuseSuperScrema + "d_9" + [input_a] + ( ScremaForm + ( fromLines + [ "\\ {x_11 : i32} : {i32} ->", + "in {x_11}" + ] + ) + [scan_op'] + [] + ( fromLines + [ "\\ {x_12 : i32} : {i32} ->", + "in {x_12}" + ] + ) + ) + ["out_a_16"] + [input_b] + ( ScremaForm + ( fromLines + [ "\\ {x_13 : f32} : {f32} ->", + "in {x_13}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_14 : f32} : {f32} ->", + "in {x_14}" + ] + ) + ) + ["out_b_15"] + ) + ) + @?= P + ( SuperScrema + "d_9" + [input_a, input_b] + ( fromLines + [ "\\ {x_11 : i32, x_10000 : f32}: {i32, f32} ->", + "{x_11, x_10000}" + ] + ) + [scan_op'] + [] + ( fromLines + [ "\\ {x_12 : i32, x_13 : f32} : {f32, i32} ->", + "{x_13, x_12}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_14 : f32, x_10001 : i32} : {f32, i32} ->", + "{x_14, x_10001}" + ] + ), + ["out_b_15", "out_a_16"] + ), + testCase "scan,red-scan,red (horizontal)" $ + let reduce_op' = + Reduce + Commutative + ( fromLines + [ "\\ {a_0 : i32, b_1 : i32} : {i32} ->", + "let {c_2 : i32} = add32(a_0, b_1)", + "in {c_2}" + ] + ) + ["0i64"] + reduce_op = + Reduce + Commutative + ( fromLines + [ "\\ {a_4 : f32, b_5 : f32} : {f32} ->", + "let {c_6 : f32} = fadd32(a_4, b_5)", + "in {c_6}" + ] + ) + ["0.0f64"] + scan_op' = + Scan + ( fromLines + [ "\\ {a_7 : i64, b_8 : i64} : {i64} ->", + "let {c_9 : i64} = add64(a_7, b_8)", + "in {c_9}" + ] + ) + ["0i32"] + scan_op = + Scan + ( fromLines + [ "\\ {a_10 : f64, b_11 : f64} : {f64} ->", + "let {c_12 : f64} = fadd64(a_10, b_11)", + "in {c_12}" + ] + ) + ["0.0f32"] + ident_a = "input_a_13 : [d_27]i64" + input_a = SOAC.identInput ident_a + ident_b = "input_b_15 : [d_27]i32" + input_b = SOAC.identInput ident_b + ident_c = "input_c_14 : [d_27]f64" + input_c = SOAC.identInput ident_c + ident_d = "input_d_16 : [d_27]f32" + input_d = SOAC.identInput ident_d + out_a = "out_a_20" + out_b = "out_b_21" + out_c = "out_c_25" + out_d = "out_d_26" + in P + ( withFreshNames + ( fuseSuperScrema + "d_27" + [input_a, input_b] + ( ScremaForm + ( fromLines + [ "\\ {x_17 : i64, x_18 : i32} : {i64, i32} ->", + "in {x_17, x_18}" + ] + ) + [scan_op'] + [reduce_op'] + ( fromLines + [ "\\ {x_19 : i64} : {i64} ->", + "in {x_19}" + ] + ) + ) + [out_b, out_a] + [input_c, input_d] + ( ScremaForm + ( fromLines + [ "\\ {x_22 : f64, x_23 : f32} : {f64, f32} ->", + "in {x_22, x_23}" + ] + ) + [scan_op] + [reduce_op] + ( fromLines + [ "\\ {x_24 : f64} : {f64} ->", + "in {x_24}" + ] + ) + ) + [out_d, out_c] + ) + ) + @?= P + ( SuperScrema + "d_27" + [input_a, input_b, input_c, input_d] + ( fromLines + [ "\\ {x_17 : i64, x_18 : i32, x_10000 : f64, x_10001 : f32}: {i64, i32, f64, f32} ->", + "{x_17, x_18, x_10000, x_10001}" + ] + ) + [scan_op'] + [reduce_op'] + ( fromLines + [ "\\ {x_19 : i64, x_22 : f64, x_23 : f32}: {f64, f32, i64} ->", + "{x_22, x_23, x_19}" + ] + ) + [scan_op] + [reduce_op] + ( fromLines + [ "\\ {x_24 : f64, x_10002 : i64}: {f64, i64} ->", + "{x_24, x_10002}" + ] + ), + [out_b, out_d, out_c, out_a] + ), + testCase "map-map (vertical) duplicate input" $ + let ident_a = "arr_5649 : [d_5648]f64" + input_a = SOAC.identInput ident_a + ident_out_a = "defunc_0_map_res_5769 : [d_5648]f64" + input_out_a = SOAC.identInput ident_out_a + out_a = identName ident_out_a + in P + ( withFreshNames + ( fuseSuperScrema + "d_5648" + [input_a] + ( ScremaForm + ( fromLines + [ "\\ {eta_p_5705 : f64} : {f64} ->", + "let {f_res_5706 : f64} = fadd64(3.0f64, eta_p_5705)", + "in {f_res_5706}" + ] + ) + [] + [] + "\\ {x_5707 : f64} : {f64} -> {x_5707}" + ) + [out_a] + [input_out_a, input_out_a] + ( ScremaForm + ( fromLines + [ "\\ {eta_p_5766 : f64, eta_p_5767 : f64} : {f64, f64} ->", + "let {f_res_5768 : f64} = fadd64(3.0f64, eta_p_5766)", + "in {f_res_5768, eta_p_5767}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5711 : f64, eta_p_5744 : f64} : {f64} ->", + "let {eta_p_5745 : f64} = x_5711", + "let {g_res_5746 : f64} = fmul64(3.0f64, eta_p_5744)", + "let {x_5754 : f64} = g_res_5746", + "let {eta_p_5755 : f64} = eta_p_5745", + "let {x_5756 : f64} = x_5711", + "let {eta_p_5757 : f64} = x_5754", + "let {-_lhs_5758 : f64} = fmul64(eta_p_5755, eta_p_5757)", + "let {-_rhs_5759 : f64} = fadd64(eta_p_5755, eta_p_5757)", + "let {h_res_5760 : f64} = fsub64(-_lhs_5758, -_rhs_5759)", + "let {x_5761 : f64} = h_res_5760", + "let {x_5762 : f64} = x_5754", + "in {x_5761}" + ] + ) + ) + ["defunc_0_map_res_5770"] + ) + ) + @?= P + ( SuperScrema + "d_5648" + [input_a] + ( fromLines + [ "\\ {eta_p_5705 : f64} : {f64} ->", + "let {f_res_5706 : f64} = fadd64(3.0f64, eta_p_5705)", + "in {f_res_5706}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5707 : f64} : {f64, f64, f64} ->", + "let {eta_p_5766 : f64} = x_5707", + "let {eta_p_5767 : f64} = eta_p_5766", + "let {f_res_5768 : f64} = fadd64(3.0f64, eta_p_5766)", + "in {f_res_5768, eta_p_5767, x_5707}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5711 : f64, eta_p_5744 : f64, x_10000 : f64} : {f64, f64} ->", + "let {eta_p_5745 : f64} = x_5711", + "let {g_res_5746 : f64} = fmul64(3.0f64, eta_p_5744)", + "let {x_5754 : f64} = g_res_5746", + "let {eta_p_5755 : f64} = eta_p_5745", + "let {x_5756 : f64} = x_5711", + "let {eta_p_5757 : f64} = x_5754", + "let {-_lhs_5758 : f64} = fmul64(eta_p_5755, eta_p_5757)", + "let {-_rhs_5759 : f64} = fadd64(eta_p_5755, eta_p_5757)", + "let {h_res_5760 : f64} = fsub64(-_lhs_5758, -_rhs_5759)", + "let {x_5761 : f64} = h_res_5760", + "let {x_5762 : f64} = x_5754", + "in {x_5761, x_10000}" + ] + ), + ["defunc_0_map_res_5770", out_a] + ) + ] + +moveRedScanSuperScremaTests :: TestTree +moveRedScanSuperScremaTests = + testGroup + "moveRedScanSuperScrema" + [ testCase "Only map" $ + let ident_a = "input_a_5565 : [d_5537]i32" + input_a = SOAC.identInput ident_a + in P + ( withFreshNamesScopeError + ( moveRedScanSuperScrema + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_0 : i32}: {i32} ->", + "let {y_1 : i32} = add32(1i32, x_0)", + "in {y_1}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_2 : i32}: {i32} -> ", + "let {y_3 : i32} = add32(2i32, x_2)", + "in {y_3}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_4 : i32}: {i32} -> ", + "let {y_5 : i32} = add32(3i32, x_4)", + "in {y_5}" + ] + ) + ) + ) + ) + @?= P + ( Just + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_0 : i32}: {i32} ->", + "let {y_1 : i32} = add32(1i32, x_0)", + "in {y_1}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_2 : i32}: {i32} -> ", + "let {y_3 : i32} = add32(2i32, x_2)", + "in {y_3}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_4 : i32}: {i32} -> ", + "let {y_5 : i32} = add32(3i32, x_4)", + "in {y_5}" + ] + ) + ) + ), + testCase "map-map-scan-map" $ + let scan_op = + Scan + ( fromLines + [ "\\ {eta_p_5571 : i32, eta_p_5572 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_5571, eta_p_5572)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + ident_a = "input_a_5565 : [d_5537]i32" + input_a = SOAC.identInput ident_a + in P + ( withFreshNamesScopeError + ( moveRedScanSuperScrema + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_0 : i32}: {i32} ->", + "let {y_1 : i32} = add32(2i32, x_0)", + "in {y_1}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_2 : i32}: {i32} -> ", + "let {y_3 : i32} = add32(2i32, x_2)", + "in {y_3}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_4 : i32}: {i32} -> ", + "let {y_5 : i32} = add32(2i32, x_4)", + "in {y_5}" + ] + ) + ) + ) + ) + @?= P + ( Just + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_0 : i32}: {i32, i32} ->", + "let {y_1 : i32} = add32(2i32, x_0)", + "let {x_10000 : i32} = y_1", + "let {y_10001 : i32} = add32(2i32, x_10000)", + "in {y_10001, y_1}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_10002 : i32, x_2 : i32}: {i32} -> ", + "in {x_10002}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_4 : i32}: {i32} -> ", + "let {y_5 : i32} = add32(2i32, x_4)", + "in {y_5}" + ] + ) + ) + ), + testCase "map-scan-map-map" $ + let scan_op = + Scan + ( fromLines + [ "\\ {eta_p_5571 : i32, eta_p_5572 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_5571, eta_p_5572)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + ident_a = "input_a_5565 : [d_5537]i32" + input_a = SOAC.identInput ident_a + in P + ( withFreshNamesScopeError + ( moveRedScanSuperScrema + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_5566 : i32}: {i32} ->", + "let {y_5567 : i32} = add32(1i32, x_5566)", + "in {y_5567}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_5568 : i32}: {i32} -> ", + "let {y_5570 : i32} = add32(2i32, x_5568)", + "in {y_5570}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5571 : i32}: {i32} -> ", + "let {y_5572 : i32} = add32(3i32, x_5571)", + "in {y_5572}" + ] + ) + ) + ) + ) + @?= P + ( Just + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_5566 : i32}: {i32} ->", + "let {y_5567 : i32} = add32(1i32, x_5566)", + "in {y_5567}" + ] + ) + [scan_op] + [] + ( fromLines + [ "\\ {x_5568 : i32}: {i32} -> ", + "let {y_5570 : i32} = add32(2i32, x_5568)", + "in {y_5570}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5571 : i32}: {i32} -> ", + "let {y_5572 : i32} = add32(3i32, x_5571)", + "in {y_5572}" + ] + ) + ) + ), + testCase "map-map-scan,reduce-map" $ + let scan_op = + Scan + ( fromLines + [ "\\ {eta_p_5571 : i32, eta_p_5572 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_5571, eta_p_5572)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + reduce_op = + Reduce + Commutative + ( fromLines + [ "\\ {eta_p_55720 : i32, eta_p_557201 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_55720, eta_p_557201)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + ident_a = "input_a_5565 : [d_5537]i32" + input_a = SOAC.identInput ident_a + in P + ( withFreshNamesScopeError + ( moveRedScanSuperScrema + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_5566 : i32}: {i32} ->", + "let {y_5567 : i32} = add32(2i32, x_5566)", + "in {y_5567}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5568 : i32}: {i32, i32} -> ", + "let {y_5570 : i32} = add32(2i32, x_5568)", + "in {y_5570, y_5570}" + ] + ) + [scan_op] + [reduce_op] + ( fromLines + [ "\\ {x_5571 : i32}: {i32} -> ", + "let {y_5572 : i32} = add32(2i32, x_5571)", + "in {y_5572}" + ] + ) + ) + ) + ) + @?= P + ( Just + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_5566 : i32}: {i32, i32, i32} ->", + "let {y_5567 : i32} = add32(2i32, x_5566)", + "let {x_10000 : i32} = y_5567", + "let {y_10001 : i32} = add32(2i32, x_10000)", + "in {y_10001, y_10001, y_5567}" + ] + ) + [scan_op] + [reduce_op] + ( fromLines + [ "\\ {x_10002 : i32, x_5568 : i32}: {i32} -> ", + "in {x_10002}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5571 : i32}: {i32} -> ", + "let {y_5572 : i32} = add32(2i32, x_5571)", + "in {y_5572}" + ] + ) + ) + ), + testCase "map-reduce-map-scan,reduce-map" $ + let scan_op = + Scan + ( fromLines + [ "\\ {eta_p_5571 : i32, eta_p_5572 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_5571, eta_p_5572)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + reduce_op = + Reduce + Commutative + ( fromLines + [ "\\ {eta_p_55720 : i32, eta_p_557201 : i32} : {i32} ->", + "let {defunc_0_op_res_5573 : i32} = add32(eta_p_55720, eta_p_557201)", + "in {defunc_0_op_res_5573}" + ] + ) + ["0i32"] + reduce_op' = + Reduce + Commutative + ( fromLines + [ "\\ {eta_p_0 : i32, eta_p_1 : i32} : {i32} ->", + "let {defunc_0_op_res_3 : i32} = add32(eta_p_0, eta_p_1)", + "in {defunc_0_op_res_3}" + ] + ) + ["0i32"] + ident_a = "input_a_5565 : [d_5537]i32" + input_a = SOAC.identInput ident_a + in P + ( withFreshNamesScopeError + ( moveRedScanSuperScrema + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_5566 : i32}: {i32, i32} ->", + "let {y_5567 : i32} = add32(1i32, x_5566)", + "in {x_5566, y_5567}" + ] + ) + [] + [reduce_op'] + ( fromLines + [ "\\ {x_6777 : i32}: {i32, i32} -> ", + "{x_6777, x_6777}" + ] + ) + [scan_op] + [reduce_op] + ( fromLines + [ "\\ {x_5571 : i32}: {i32} -> ", + "let {y_5572 : i32} = add32(3i32, x_5571)", + "in {y_5572}" + ] + ) + ) + ) + ) + @?= P + ( Just + ( SuperScrema + "d_5537" + [input_a] + ( fromLines + [ "\\ {x_5566 : i32}: {i32, i32, i32, i32} ->", + "let {y_5567 : i32} = add32(1i32, x_5566)", + "let {x_10000 : i32} = y_5567", + "in {x_10000, x_5566, x_10000, y_5567}" + ] + ) + [scan_op] + [reduce_op', reduce_op] + ( fromLines + [ "\\ {x_10001 : i32, x_6777 : i32}: {i32} -> ", + "in {x_10001}" + ] + ) + [] + [] + ( fromLines + [ "\\ {x_5571 : i32}: {i32} -> ", + "let {y_5572 : i32} = add32(3i32, x_5571)", + "in {y_5572}" + ] + ) + ) + ) + ] + +tests :: TestTree +tests = + testGroup + "ScremaTests" + [ splitLambdaByParTests, + fuseSuperScremaTests, + moveRedScanSuperScremaTests + ] diff --git a/src-testing/Futhark/Optimise/FusionTests.hs b/src-testing/Futhark/Optimise/FusionTests.hs new file mode 100644 --- /dev/null +++ b/src-testing/Futhark/Optimise/FusionTests.hs @@ -0,0 +1,10 @@ +module Futhark.Optimise.FusionTests (tests) where + +import Futhark.Optimise.Fusion.ScremaTests qualified +import Test.Tasty + +tests :: TestTree +tests = + testGroup + "OptimiseFusionTests" + [Futhark.Optimise.Fusion.ScremaTests.tests] diff --git a/src-testing/Language/Futhark/SyntaxTests.hs b/src-testing/Language/Futhark/SyntaxTests.hs --- a/src-testing/Language/Futhark/SyntaxTests.hs +++ b/src-testing/Language/Futhark/SyntaxTests.hs @@ -4,7 +4,7 @@ import Control.Applicative hiding (many, some) import Data.Bifunctor -import Data.Char (isAlpha) +import Data.Char (isAlphaNum) import Data.Functor import Data.Map qualified as M import Data.String @@ -12,7 +12,7 @@ import Data.Void import Language.Futhark import Language.Futhark.Parser (SyntaxError (syntaxErrorMsg), parseExp, parseType) -import Language.Futhark.Primitive.Parse (constituent, keyword, lexeme) +import Language.Futhark.Primitive.Parse (keyword, lexeme) import Language.Futhark.PrimitiveTests () import Test.QuickCheck import Test.Tasty @@ -54,9 +54,22 @@ let (s', '_' : tag) = span (/= '_') s in VName (fromString s') (read tag) -instance (IsString v) => IsString (QualName v) where - fromString = QualName [] . fromString +pQualVName :: Parser (QualName VName) +pQualVName = do + xs <- pVName `sepBy1` "." + pure $ QualName (init xs) (last xs) +pQualName :: Parser (QualName Name) +pQualName = do + xs <- pName `sepBy1` "." + pure $ QualName (init xs) (last xs) + +instance IsString (QualName VName) where + fromString = fromStringParse pQualVName "QualName VName" + +instance IsString (QualName Name) where + fromString = fromStringParse pQualName "QualName Name" + instance IsString UncheckedTypeExp where fromString = either (error . T.unpack . syntaxErrorMsg) id @@ -70,25 +83,22 @@ brackets = between (lexeme "[") (lexeme "]") parens = between (lexeme "(") (lexeme ")") +constituent :: Char -> Bool +constituent c = isAlphaNum c || (c `elem` ("_/'+-=!&^<>*|%" :: String)) + pName :: Parser Name pName = - lexeme . fmap nameFromString $ - (:) <$> satisfy isAlpha <*> many (satisfy constituent) + lexeme . fmap nameFromString $ some $ satisfy constituent pVName :: Parser VName pVName = lexeme $ do (s, tag) <- - satisfy constituent - `manyTill_` try pTag - "variable name" + satisfy constituent `manyTill_` try pTag "variable name" pure $ VName (nameFromString s) tag where pTag = "_" *> L.decimal <* notFollowedBy (satisfy constituent) -pQualName :: Parser (QualName VName) -pQualName = QualName [] <$> pVName - pPrimType :: Parser PrimType pPrimType = choice $ @@ -117,7 +127,7 @@ brackets $ choice [ flip sizeFromInteger mempty <$> lexeme L.decimal, - flip sizeFromName mempty <$> pQualName + flip sizeFromName mempty <$> pQualVName ] pScalarNonFun :: Parser (ScalarTypeBase Size Uniqueness) @@ -130,7 +140,7 @@ ] where pField = (,) <$> pName <* lexeme ":" <*> pType - pTypeVar = TypeVar <$> pUniqueness <*> pQualName <*> many pTypeArg + pTypeVar = TypeVar <$> pUniqueness <*> pQualVName <*> many pTypeArg pTypeArg = choice [ TypeArgDim <$> pSize, diff --git a/src-testing/Language/Futhark/TypeChecker/ModulesTests.hs b/src-testing/Language/Futhark/TypeChecker/ModulesTests.hs new file mode 100644 --- /dev/null +++ b/src-testing/Language/Futhark/TypeChecker/ModulesTests.hs @@ -0,0 +1,100 @@ +module Language.Futhark.TypeChecker.ModulesTests + ( tests, + ) +where + +import Data.Map qualified as M +import Futhark.Util.Pretty (docString) +import Language.Futhark.Semantic +import Language.Futhark.Syntax +import Language.Futhark.SyntaxTests () +import Language.Futhark.TypeChecker +import Language.Futhark.TypeChecker.Modules +import Test.Tasty +import Test.Tasty.HUnit + +higherOrderModules :: TestTree +higherOrderModules = + testCase "higher order modules" $ + case matchMTys mty_a mty_b mempty of + Left err -> assertFailure $ docString $ prettyTypeErrorNoLoc err + Right m -> m @?= M.fromList [("t_4654", "t_4658"), ("f_4656", "f_4661")] + where + mty_a = + MTy mempty . ModFun $ + FunModType + { funModTypeAbs = + M.fromList [("t_4658", Unlifted)], + funModTypeMod = + ModEnv $ + mempty + { envTypeTable = + M.fromList + [("t_4658", TypeAbbr Unlifted [] "t_4658")], + envNameMap = M.fromList [((Type, "t"), "t_4658")] + }, + funModTypeMty = + MTy + { mtyAbs = mempty, + mtyMod = + ModEnv $ + mempty + { envVtable = + M.fromList + [ ( "f_4661", + BoundV + { boundValTParams = [], + boundValType = "(x_4660: t_4658) -> t_4659.t_4658" + } + ) + ], + envNameMap = M.fromList [((Term, "f"), "f_4661")] + } + } + } + mty_b = + MTy mempty . ModFun $ + FunModType + { funModTypeAbs = M.fromList [("t_4654", Unlifted)], + funModTypeMod = + ModEnv $ + mempty + { envTypeTable = + M.fromList + [ ( VName "t" 4654, + TypeAbbr Unlifted [] "t_4654" + ) + ], + envNameMap = + M.fromList [((Type, "t"), "t_4654")] + }, + funModTypeMty = + MTy + { mtyAbs = mempty, + mtyMod = + ModEnv $ + mempty + { envVtable = + M.fromList + [ ( "f_4656", + BoundV + { boundValTParams = [], + boundValType = "V_4655.t_4654 -> V_4655.t_4654" + } + ) + ], + envNameMap = + M.fromList [((Term, "f"), "f_4656")] + } + } + } + +tests :: TestTree +tests = + testGroup + "ModulesTests" + [ testGroup + "matchMTys" + [ higherOrderModules + ] + ] diff --git a/src-testing/Language/Futhark/TypeCheckerTests.hs b/src-testing/Language/Futhark/TypeCheckerTests.hs --- a/src-testing/Language/Futhark/TypeCheckerTests.hs +++ b/src-testing/Language/Futhark/TypeCheckerTests.hs @@ -1,6 +1,7 @@ module Language.Futhark.TypeCheckerTests (tests) where import Language.Futhark.TypeChecker.ConsumptionTests qualified +import Language.Futhark.TypeChecker.ModulesTests qualified import Language.Futhark.TypeChecker.TypesTests qualified import Test.Tasty @@ -9,5 +10,6 @@ testGroup "Source type checker tests" [ Language.Futhark.TypeChecker.TypesTests.tests, - Language.Futhark.TypeChecker.ConsumptionTests.tests + Language.Futhark.TypeChecker.ConsumptionTests.tests, + Language.Futhark.TypeChecker.ModulesTests.tests ] diff --git a/src-testing/futhark_tests.hs b/src-testing/futhark_tests.hs --- a/src-testing/futhark_tests.hs +++ b/src-testing/futhark_tests.hs @@ -9,7 +9,9 @@ import Futhark.IR.PropTests qualified import Futhark.IR.Syntax.CoreTests qualified import Futhark.Internalise.TypesValuesTests qualified +import Futhark.LspTests qualified import Futhark.Optimise.ArrayLayoutTests qualified +import Futhark.Optimise.FusionTests qualified import Futhark.Optimise.MemoryBlockMerging.GreedyColoringTests qualified import Futhark.Pkg.SolveTests qualified import Language.Futhark.PrettyTests qualified @@ -39,7 +41,9 @@ Futhark.Analysis.DataDependenciesTests.tests, Language.Futhark.TypeCheckerTests.tests, Language.Futhark.SemanticTests.tests, - Futhark.Optimise.ArrayLayoutTests.tests + Futhark.Optimise.ArrayLayoutTests.tests, + Futhark.Optimise.FusionTests.tests, + Futhark.LspTests.tests ] main :: IO () diff --git a/src/Futhark/AD/Fwd.hs b/src/Futhark/AD/Fwd.hs --- a/src/Futhark/AD/Fwd.hs +++ b/src/Futhark/AD/Fwd.hs @@ -264,13 +264,14 @@ letExp "zero" $ zeroExp t fwdSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM () -fwdSOAC pat aux (Screma size xs (ScremaForm f scs reds)) = do +fwdSOAC pat aux (Screma size xs (ScremaForm f scs reds post_lam)) = do pat' <- bundleNewPat pat xs' <- bundleTangents xs f' <- fwdLambda f scs' <- mapM fwdScan scs reds' <- mapM fwdRed reds - addStm $ Let pat' aux $ Op $ Screma size xs' $ ScremaForm f' scs' reds' + post_lam' <- fwdLambda post_lam + addStm $ Let pat' aux $ Op $ Screma size xs' $ ScremaForm f' scs' reds' post_lam' where fwdScan :: Scan SOACS -> ADM (Scan SOACS) fwdScan sc = do diff --git a/src/Futhark/AD/Rev/Hist.hs b/src/Futhark/AD/Rev/Hist.hs --- a/src/Futhark/AD/Rev/Hist.hs +++ b/src/Futhark/AD/Rev/Hist.hs @@ -71,8 +71,8 @@ params <- traverse (\tp -> newParam "x" $ Array tp (Shape s) NoUniqueness) pt body <- nestedmap r pt lam mkLambda params $ - fmap varsRes . letTupExp "res" . Op $ - Screma h (map paramName params) (mapSOAC body) + fmap varsRes . letTupExp "res" . Op . Screma h (map paramName params) + =<< mapSOAC body -- \ds hs -> map2 lam ds hs mkF' :: Lambda SOACS -> [Type] -> SubExp -> ADM ([VName], [VName], Lambda SOACS) @@ -85,8 +85,8 @@ let hs_pars = fmap paramName hs_params lam_map <- mkLambda (ds_params <> hs_params) $ - fmap varsRes . letTupExp "map_f'" . Op $ - Screma n (ds_pars <> hs_pars) (mapSOAC lam') + fmap varsRes . letTupExp "map_f'" . Op . Screma n (ds_pars <> hs_pars) + =<< mapSOAC lam' pure (ds_pars, hs_pars, lam_map) @@ -109,11 +109,12 @@ rs_params <- traverse (newParam "rs_param") tps let map_params = ls_params <> as_params <> rs_params lam_map <- - mkLambda map_params $ - fmap varsRes . letTupExp "map_f" $ - Op $ - Screma n (map paramName map_params) $ - mapSOAC lam' + mkLambda map_params + . fmap varsRes + . letTupExp "map_f" + . Op + . Screma n (map paramName map_params) + =<< mapSOAC lam' pure (map paramName as_params, lam_map) @@ -128,7 +129,7 @@ (eBody $ pure $ eParam par_is) (eBody $ pure $ eSubExp w) - letExp "is'" $ Op $ Screma n (pure is) $ mapSOAC is'_lam + letExp "is'" . Op . Screma n (pure is) =<< mapSOAC is'_lam multiScatter :: [VName] -> VName -> [VName] -> ADM [VName] multiScatter dst is vs = @@ -228,7 +229,7 @@ fmap varsRes . letTupExp "res" =<< do pure $ BasicOp $ Replicate (Shape inner_dims) $ Var $ paramName i - letExp "res" $ Op $ Screma n [iota_n] $ mapSOAC lam + letExp "res" . Op . Screma n [iota_n] =<< mapSOAC lam let hist_op = HistOp (Shape [w]) rf [dst_cpy, dst_minus_ones] [ne, if nr_dims == 1 then intConst Int64 (-1) else ne_minus_ones] hist_lam f' <- mkIdentityLambda [Prim int64, rowType vs_type, rowType $ Array int64 (Shape vs_dims) NoUniqueness] @@ -254,8 +255,8 @@ dst_lam <- nestedmap inner_dims [int64, vs_elm_type] dst_lam_inner dst_bar <- - letExp (baseName dst <> "_bar") . Op $ - Screma w [x_inds, x_bar] (mapSOAC dst_lam) + letExp (baseName dst <> "_bar") . Op . Screma w [x_inds, x_bar] + =<< mapSOAC dst_lam updateAdj dst dst_bar @@ -283,8 +284,8 @@ vs_lam <- nestedmap inner_dims (vs_elm_type : replicate nr_dims int64) vs_lam_inner vs_bar_p <- - letExp (baseName vs <> "_partial") . Op $ - Screma w (x_bar : inds) (mapSOAC vs_lam) + letExp (baseName vs <> "_partial") . Op . Screma w (x_bar : inds) + =<< mapSOAC vs_lam q <- letSubExp "q" @@ -324,7 +325,7 @@ mk_indices dims $ iotas ++ [Var $ paramName i_param] - letTupExp "res" $ Op $ Screma d [iota_d] $ mapSOAC lam + letTupExp "res" . Op . Screma d [iota_d] =<< mapSOAC lam -- -- special case of histogram with multiplication as operator. @@ -467,7 +468,8 @@ (eBody $ pure $ pure $ zeroExp $ rowType dst_type) vs_bar <- - letExp (baseName vs <> "_bar") $ Op $ Screma n [is, vs] $ mapSOAC lam_vsbar + letExp (baseName vs <> "_bar") . Op . Screma n [is, vs] + =<< mapSOAC lam_vsbar updateAdj vs vs_bar @@ -515,7 +517,9 @@ (eBody $ pure $ pure $ BasicOp $ Index x_bar $ fullSlice x_type [DimFix $ Var i]) (eBody $ pure $ eSubExp ne) - vs_bar <- letExp (baseName vs <> "_bar") $ Op $ Screma n [is] $ mapSOAC lam_vsbar + vs_bar <- + letExp (baseName vs <> "_bar") . Op . Screma n [is] + =<< mapSOAC lam_vsbar updateAdj vs vs_bar -- Special case for vectorised combining operator. Rewrite @@ -569,9 +573,8 @@ [HistOp (Shape [w]) rf [dst_col_cpy] [Var $ paramName ne] op] f histT <- - letExp "histT" . Op $ - Screma (arraySize 0 t_dstT) [dstT, vssT, nes] $ - mapSOAC map_lam + letExp "histT" . Op . Screma (arraySize 0 t_dstT) [dstT, vssT, nes] + =<< mapSOAC map_lam auxing aux . letBindNames [x] . BasicOp $ Rearrange histT dims foldr (vjpStm ops) m stms @@ -625,7 +628,7 @@ ) ) - bins <- letExp "bins" $ Op $ Screma n [is] $ mapSOAC num_lam + bins <- letExp "bins" . Op . Screma n [is] =<< mapSOAC num_lam flag_param <- newParam "flag" $ Prim int64 flag_lam <- mkLambda [flag_param] $ @@ -635,7 +638,9 @@ (map ((,) (eParam flag_param) . iConst) [0 .. 2]) (map (eBody . fmap iConst . (\i -> map (\j -> if i == j then 1 else 0) [0 .. 3])) ([0 .. 3] :: [Integer])) - flags <- letTupExp "flags" $ Op $ Screma n [bins] $ mapSOAC flag_lam + flags <- + letTupExp "flags" . Op . Screma n [bins] + =<< mapSOAC flag_lam scan_params <- traverse (flip newParam $ Prim int64) ["a1", "b1", "c1", "d1", "a2", "b2", "c2", "d2"] scan_lam <- @@ -670,7 +675,7 @@ (tail map_params) ) - nis <- letExp "nis" $ Op $ Screma n (bins : offsets) $ mapSOAC map_lam + nis <- letExp "nis" . Op . Screma n (bins : offsets) =<< mapSOAC map_lam scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) tps multiScatter scatter_dst nis xs @@ -738,7 +743,9 @@ let slice = [DimFix $ Var $ paramName i_param] map_lam <- mkLambda [i_param] $ varsRes <$> multiIndex (tail xs) slice - sorted <- letTupExp "sorted" $ Op $ Screma n [iota'] $ mapSOAC map_lam + sorted <- + letTupExp "sorted" . Op . Screma n [iota'] + =<< mapSOAC map_lam pure $ iota' : is' : sorted -- @@ -784,8 +791,8 @@ Hist n as [HistOp (Shape w) rf nes ne lam0] h_map lam0' <- renameLambda lam0 - auxing aux . letBindNames xs . Op $ - Screma (head w) (dst <> h_part) (mapSOAC lam0') + auxing aux . letBindNames xs . Op . Screma (head w) (dst <> h_part) + =<< mapSOAC lam0' m @@ -817,7 +824,7 @@ par_i <- newParam "i" $ Prim int64 flag_lam <- mkFlagLam par_i sis - flag <- letExp "flag" $ Op $ Screma n [iota_n] $ mapSOAC flag_lam + flag <- letExp "flag" . Op . Screma n [iota_n] =<< mapSOAC flag_lam -- map (\i -> (if flag[i] then (true,ne) else (false,vs[i-1]), if i==0 || flag[n-i] then (true,ne) else (false,vs[n-i]))) (iota n) par_i' <- newParam "i" $ Prim int64 @@ -887,8 +894,8 @@ let ne' = Constant (BoolValue False) : ne scansres <- - letTupExp "adj_ctrb_scan" . Op $ - Screma n [iota_n] (scanomapSOAC (map (`Scan` ne') scan_lams) g_lam) + letTupExp "adj_ctrb_scan" . Op . Screma n [iota_n] + =<< scanomapSOAC (map (`Scan` ne') scan_lams) g_lam let (_ : ls_arr, _ : rs_arr_rev) = splitAt (length ne + 1) scansres @@ -905,7 +912,7 @@ map (\t -> pure $ BasicOp $ Replicate (Shape $ tail $ arrayDims t) (Constant $ blankPrimValue $ elemType t)) as_type ) - f_bar <- letTupExp "f_bar" $ Op $ Screma n [sis] $ mapSOAC map_lam + f_bar <- letTupExp "f_bar" . Op . Screma n [sis] =<< mapSOAC map_lam (as_params, f) <- mkF lam0 as_type n f_adj <- vjpLambda ops (map adjFromVar f_bar) as_params f @@ -917,7 +924,9 @@ nmim1 <- letSubExp "n_i_1" =<< toExp (pe64 n - le64 i''' - 1) varsRes <$> multiIndex rs_arr_rev [DimFix nmim1] - rs_arr <- letTupExp "rs_arr" $ Op $ Screma n [iota_n] $ mapSOAC rev_lam + rs_arr <- + letTupExp "rs_arr" . Op . Screma n [iota_n] + =<< mapSOAC rev_lam sas_bar <- bindSubExpRes "sas_bar" diff --git a/src/Futhark/AD/Rev/Map.hs b/src/Futhark/AD/Rev/Map.hs --- a/src/Futhark/AD/Rev/Map.hs +++ b/src/Futhark/AD/Rev/Map.hs @@ -183,8 +183,11 @@ (param_contribs, free_contribs) <- fmap (splitAt (length (lambdaParams map_lam'))) $ - auxing aux . letTupExp "map_adjs" . Op $ - Screma w (as ++ pat_adj_vals ++ free_adjs) (mapSOAC lam_rev) + auxing aux + . letTupExp "map_adjs" + . Op + . Screma w (as ++ pat_adj_vals ++ free_adjs) + =<< mapSOAC lam_rev -- Crucial that we handle the free contribs first in case 'free' -- and 'as' intersect. diff --git a/src/Futhark/AD/Rev/Monad.hs b/src/Futhark/AD/Rev/Monad.hs --- a/src/Futhark/AD/Rev/Monad.hs +++ b/src/Futhark/AD/Rev/Monad.hs @@ -336,7 +336,7 @@ ret <- mapM lookupType res pure (ret, varsRes res) let lam = Lambda (iparam : params) ret (Body () stms res) - letTupExp "tab" $ Op $ Screma w (iota : vs) (mapSOAC lam) + letTupExp "tab" . Op . Screma w (iota : vs) =<< mapSOAC lam -- | Construct a lambda for adding two values of the given type. addLambda :: Type -> ADM (Lambda SOACS) @@ -347,8 +347,10 @@ lam <- addLambda $ rowType t body <- insertStmsM $ do res <- - letSubExp "lam_map" . Op $ - Screma (arraySize 0 t) [paramName xs_p, paramName ys_p] (mapSOAC lam) + letSubExp "lam_map" + . Op + . Screma (arraySize 0 t) [paramName xs_p, paramName ys_p] + =<< mapSOAC lam pure $ resultBody [res] pure Lambda @@ -368,7 +370,7 @@ pure $ BasicOp $ BinOp (addBinOp pt) (Var x) (Var y) Array {} -> do lam <- addLambda $ rowType x_t - pure $ Op $ Screma (arraySize 0 x_t) [x, y] (mapSOAC lam) + Op . Screma (arraySize 0 x_t) [x, y] <$> mapSOAC lam _ -> error $ "addExp: unexpected type: " ++ prettyString x_t diff --git a/src/Futhark/AD/Rev/Reduce.hs b/src/Futhark/AD/Rev/Reduce.hs --- a/src/Futhark/AD/Rev/Reduce.hs +++ b/src/Futhark/AD/Rev/Reduce.hs @@ -57,7 +57,7 @@ (resultBodyM $ scanNeutral scan) (eBody $ map (`eIndex` [prev]) res_incl) - letTupExp desc $ Op $ Screma w [iota] (mapSOAC lam) + letTupExp desc . Op . Screma w [iota] =<< mapSOAC lam mkF :: Lambda SOACS -> ADM ([VName], Lambda SOACS) mkF lam = do @@ -117,7 +117,9 @@ f_adj <- vjpLambda ops (map adjFromVar pat_adj) as_params f - as_adj <- letTupExp "adjs" $ Op $ Screma w (ls ++ as ++ rs) (mapSOAC f_adj) + as_adj <- + letTupExp "adjs" . Op . Screma w (ls ++ as ++ rs) + =<< mapSOAC f_adj zipWithM_ updateAdj as as_adj where @@ -222,7 +224,8 @@ fmap varsRes . letTupExp "idx_res" $ Op $ Screma w [paramName as_param] reduce_form - addStm $ Let x aux $ Op $ Screma (arraySize 0 ts) [tran_as, ne] $ mapSOAC map_lam + addStm . Let x aux . Op . Screma (arraySize 0 ts) [tran_as, ne] + =<< mapSOAC map_lam foldr (vjpStm ops) m stms @@ -264,10 +267,8 @@ ps <- newVName "ps" zs <- newVName "zs" auxing aux $ - letBindNames [ps, zs] $ - Op $ - Screma w [as] $ - mapSOAC map_lam + letBindNames [ps, zs] . Op . Screma w [as] + =<< mapSOAC map_lam red_lam_mul <- binOpLambda mul t red_lam_add <- binOpLambda (Add Int64 OverflowUndef) int64 @@ -322,7 +323,9 @@ (eBody $ pure const_zero) ) - as_adjup <- letExp "adjs" $ Op $ Screma w [as] $ mapSOAC map_lam_rev + as_adjup <- + letExp "adjs" . Op . Screma w [as] + =<< mapSOAC map_lam_rev updateAdj as as_adjup where diff --git a/src/Futhark/AD/Rev/SOAC.hs b/src/Futhark/AD/Rev/SOAC.hs --- a/src/Futhark/AD/Rev/SOAC.hs +++ b/src/Futhark/AD/Rev/SOAC.hs @@ -53,7 +53,7 @@ histomapToMapAndHist :: Pat Type -> (SubExp, [HistOp SOACS], Lambda SOACS, [VName]) -> ADM (Stm SOACS, Stm SOACS) histomapToMapAndHist (Pat pes) (w, histops, map_lam, as) = do map_pat <- traverse accMapPatElem $ lambdaReturnType map_lam - let map_stm = mkLet map_pat $ Op $ Screma w as $ mapSOAC map_lam + map_stm <- mkLet map_pat . Op . Screma w as <$> mapSOAC map_lam new_lam <- mkIdentityLambda $ lambdaReturnType map_lam let hist_stm = Let (Pat pes) (defAux ()) $ Op $ Hist w (map identName map_pat) histops new_lam pure (map_stm, hist_stm) @@ -137,6 +137,12 @@ scanomapToMapAndScan pat (w, scans, map_lam, as) vjpStm ops mapstm $ vjpStm ops scanstm m +-- Get rid of post-lambdas +vjpSOAC ops pat aux (Screma w as form) m + | not $ isIdentityLambda $ scremaPostLambda form = do + stms <- collectStms_ $ auxing aux $ extractPostLambda pat w as form + foldr (vjpStm ops) m stms + -- Differentiating Histograms vjpSOAC ops pat aux (Hist n as histops f) m | isIdentityLambda f, @@ -218,7 +224,8 @@ cs == mempty, [map_stm] <- stmsToList (bodyStms lam_body), (Let (Pat [pe]) _ (Op scrm)) <- map_stm, - (Screma _ [a1, a2] (ScremaForm map_lam [] [])) <- scrm, + (Screma _ [a1, a2] (ScremaForm map_lam [] [] post_lam)) <- scrm, + isIdentityLambda post_lam, (a1 == paramName pa1 && a2 == paramName pa2) || (a1 == paramName pa2 && a2 == paramName pa1), r == Var (patElemName pe) = Just map_lam diff --git a/src/Futhark/AD/Rev/Scan.hs b/src/Futhark/AD/Rev/Scan.hs --- a/src/Futhark/AD/Rev/Scan.hs +++ b/src/Futhark/AD/Rev/Scan.hs @@ -184,7 +184,7 @@ ) iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64 - letTupExp "scan_contribs" $ Op $ Screma w (iota : xs) $ mapSOAC map_lam + letTupExp "scan_contribs" . Op . Screma w (iota : xs) =<< mapSOAC map_lam -- | Scan special cases. data SpecialCase = ZeroQuadrant | MatrixMul deriving (Show) @@ -286,11 +286,11 @@ map_scan <- revArrLam as -- perform the scan scan_res <- - letTupExp "adj_ctrb_scan" . Op . Screma w [iota] $ - scanomapSOAC [rev_scan] map_scan + letTupExp "adj_ctrb_scan" . Op . Screma w [iota] + =<< scanomapSOAC [rev_scan] map_scan -- flip the output array again rev_lam <- revArrLam scan_res - letTupExp "reverse_scan_result" $ Op $ Screma w [iota] $ mapSOAC rev_lam + letTupExp "reverse_scan_result" . Op . Screma w [iota] =<< mapSOAC rev_lam where revArrLam :: [VName] -> ADM (Lambda SOACS) revArrLam arrs = do @@ -348,7 +348,7 @@ pure $ varRes a_lift iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64 - letTupExp "as_lift" $ Op $ Screma w [iota] $ mapSOAC lmb + letTupExp "as_lift" . Op . Screma w [iota] =<< mapSOAC lmb ysRightPPAD :: [VName] -> SubExp -> [SubExp] -> ADM [VName] ysRightPPAD ys w e = do @@ -364,7 +364,7 @@ pure $ varRes a_lift iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64 - letTupExp "ys_right" $ Op $ Screma w [iota] $ mapSOAC lmb + letTupExp "ys_right" . Op . Screma w [iota] =<< mapSOAC lmb finalMapPPAD :: VjpOps -> [VName] -> Scan SOACS -> ADM (Lambda SOACS) finalMapPPAD ops as scan = do @@ -402,7 +402,8 @@ ys_right <- ysRightPPAD ys w e final_lmb <- finalMapPPAD ops as scan - letTupExp "as_bar" $ Op $ Screma w (ys_right ++ as ++ rs_adj) $ mapSOAC final_lmb + letTupExp "as_bar" . Op . Screma w (ys_right ++ as ++ rs_adj) + =<< mapSOAC final_lmb GenericIFL23 sc -> do -- IFL23 map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj sc d @@ -411,8 +412,8 @@ iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64 r_scan <- - letTupExp "adj_ctrb_scan" . Op . Screma w [iota] $ - scanomapSOAC scan_lams map1_lam + letTupExp "adj_ctrb_scan" . Op . Screma w [iota] + =<< scanomapSOAC scan_lams map1_lam mkScanFinalMap ops w (scanLambda scan) as ys (splitScanRes sc r_scan d) -- Goal: calculate as_contribs in new way -- zipWithM_ updateAdj as as_contribs -- as_bar += new adjoint @@ -458,8 +459,8 @@ Screma w (map paramName as_par) scan_form transp_ys <- - letTupExp "trans_ys" . Op $ - Screma n (transp_as ++ subExpVars ne) (mapSOAC map_lam) + letTupExp "trans_ys" . Op . Screma n (transp_as ++ subExpVars ne) + =<< mapSOAC map_lam forM (zip ys transp_ys) $ \(y, x) -> auxing aux $ letBindNames [y] $ BasicOp $ Rearrange x rear @@ -477,10 +478,11 @@ letExp "iota" $ BasicOp $ Iota n (intConst Int64 0) (intConst Int64 1) Int64 scan_res <- - letExp "res_rev" $ Op $ Screma n [iota] $ scanomapSOAC [Scan lam [ne]] map_scan + letExp "res_rev" . Op . Screma n [iota] + =<< scanomapSOAC [Scan lam [ne]] map_scan rev_lam <- rev_arr_lam scan_res - contrb <- letExp "contrb" $ Op $ Screma n [iota] $ mapSOAC rev_lam + contrb <- letExp "contrb" . Op . Screma n [iota] =<< mapSOAC rev_lam updateAdj as contrb where diff --git a/src/Futhark/Actions.hs b/src/Futhark/Actions.hs --- a/src/Futhark/Actions.hs +++ b/src/Futhark/Actions.hs @@ -536,6 +536,9 @@ else [] ) ++ ["-s", "WASM_BIGINT"] + ++ [ "-s", + "EXPORTED_RUNTIME_METHODS=['HEAP8','HEAP16','HEAP32','HEAP64','HEAPU8','HEAPU16','HEAPU32','HEAPF32','HEAPF64']" + ] ++ cmdCFLAGS cflags_def ++ cmdEMCFLAGS [""] ++ [ "-s", diff --git a/src/Futhark/Analysis/DataDependencies.hs b/src/Futhark/Analysis/DataDependencies.hs --- a/src/Futhark/Analysis/DataDependencies.hs +++ b/src/Futhark/Analysis/DataDependencies.hs @@ -2,13 +2,16 @@ module Futhark.Analysis.DataDependencies ( Dependencies, dataDependencies, + dataDependencies', depsOf, depsOf', + depsOfVar, depsOfArrays, depsOfShape, lambdaDependencies, reductionDependencies, findNecessaryForReturned, + depsOfNames, ) where @@ -23,13 +26,13 @@ type Dependencies = M.Map VName Names -- | Compute the data dependencies for an entire body. -dataDependencies :: (ASTRep rep) => Body rep -> Dependencies +dataDependencies :: (ASTRep rep) => GBody rep res -> Dependencies dataDependencies = dataDependencies' M.empty dataDependencies' :: (ASTRep rep) => Dependencies -> - Body rep -> + GBody rep res -> Dependencies dataDependencies' startdeps = foldl grow startdeps . bodyStms where diff --git a/src/Futhark/Analysis/HORep/MapNest.hs b/src/Futhark/Analysis/HORep/MapNest.hs --- a/src/Futhark/Analysis/HORep/MapNest.hs +++ b/src/Futhark/Analysis/HORep/MapNest.hs @@ -22,7 +22,7 @@ import Futhark.Analysis.HORep.SOAC qualified as SOAC import Futhark.Construct import Futhark.IR hiding (typeOf) -import Futhark.IR.SOACS (SOACS) +import Futhark.IR.SOACS (SOACS, isMapSOAC) import Futhark.IR.SOACS.SOAC qualified as Futhark import Futhark.Transform.Substitute @@ -70,32 +70,36 @@ setDepth n nw = n {nestingWidth = nw} pushIntoMapLambda :: + (MonadFreshNames m) => Stms SOACS -> Stm SOACS -> - Maybe (Stm SOACS) + m (Maybe (Stm SOACS)) pushIntoMapLambda stms (Let pat aux (Op (Futhark.Screma w inps form))) | Just map_lam <- Futhark.isMapSOAC form, - not $ any ((`namesIntersect` bound_by_stms) . freeIn) inps = + not $ any ((`namesIntersect` bound_by_stms) . freeIn) inps = do let lam_body = lambdaBody map_lam map_lam' = map_lam {lambdaBody = lam_body {bodyStms = stms <> bodyStms lam_body}} - form' = Futhark.mapSOAC map_lam' - in Just $ Let pat aux (Op (Futhark.Screma w inps form')) + form' <- Futhark.mapSOAC map_lam' + pure $ Just $ Let pat aux (Op (Futhark.Screma w inps form')) where bound_by_stms = namesFromList $ foldMap (patNames . stmPat) stms -pushIntoMapLambda _ _ = Nothing +pushIntoMapLambda _ _ = pure Nothing -massage :: SOAC SOACS -> SOAC SOACS -massage (SOAC.Screma w inps form) +massage :: (MonadFreshNames m) => SOAC SOACS -> m (SOAC SOACS) +massage soac@(SOAC.Screma w inps form) | Just lam <- Futhark.isMapSOAC form, Just (init_stms, last_stm) <- stmsLast $ bodyStms $ lambdaBody lam, all (cheap . stmExp) init_stms, all (`notNameIn` freeIn (bodyResult (lambdaBody lam))) $ - foldMap (patNames . stmPat) init_stms, - Just last_stm' <- pushIntoMapLambda init_stms last_stm = - let lam' = - lam {lambdaBody = (lambdaBody lam) {bodyStms = oneStm last_stm'}} - in SOAC.Screma w inps (Futhark.mapSOAC lam') + foldMap (patNames . stmPat) init_stms = do + maybe_last_stm <- pushIntoMapLambda init_stms last_stm + case maybe_last_stm of + Just last_stm' -> do + let lam' = + lam {lambdaBody = (lambdaBody lam) {bodyStms = oneStm last_stm'}} + SOAC.Screma w inps <$> Futhark.mapSOAC lam' + Nothing -> pure soac where cheap (BasicOp BinOp {}) = True cheap (BasicOp SubExp {}) = True @@ -103,15 +107,18 @@ cheap (BasicOp ConvOp {}) = True cheap (BasicOp UnOp {}) = True cheap _ = False -massage soac = soac +massage soac = pure soac fromSOAC' :: (MonadFreshNames m, LocalScope SOACS m) => [Ident] -> SOAC SOACS -> m (Maybe MapNest) -fromSOAC' bound soac - | SOAC.Screma w inps (SOAC.ScremaForm lam [] []) <- massage soac = do +fromSOAC' bound soac = do + screma <- massage soac + + case screma of + SOAC.Screma w inps form | Just lam <- isMapSOAC form -> do let bound' = bound <> map paramIdent (lambdaParams lam) maybenest <- case ( stmsToList $ bodyStms $ lambdaBody lam, @@ -121,7 +128,11 @@ | map resSubExp res == map Var (patNames pat) -> localScope (scopeOfLParams $ lambdaParams lam) $ SOAC.fromExp e - >>= either (pure . Left) (fmap (Right . fmap (pat,)) . fromSOAC' bound') + >>= either + (pure . Left) + ( fmap (Right . fmap (pat,)) + . fromSOAC' bound' + ) _ -> pure $ Right Nothing @@ -169,14 +180,14 @@ ++ [Param mempty name t | Ident name t <- newParams] } pure $ Just $ MapNest w lam' [] inps' -fromSOAC' _ _ = pure Nothing + _any -> pure Nothing fromSOAC :: (MonadFreshNames m, LocalScope SOACS m) => SOAC SOACS -> m (Maybe MapNest) fromSOAC = fromSOAC' mempty toSOAC :: (MonadFreshNames m, HasScope SOACS m) => MapNest -> m (SOAC SOACS) toSOAC (MapNest w lam [] inps) = - pure $ SOAC.Screma w inps (Futhark.mapSOAC lam) + SOAC.Screma w inps <$> Futhark.mapSOAC lam toSOAC (MapNest w lam (Nesting npnames nres nrettype nw : ns) inps) = do let nparams = zipWith (Param mempty) npnames $ map SOAC.inputRowType inps body <- runBodyBuilder $ @@ -191,7 +202,7 @@ lambdaBody = body, lambdaReturnType = nrettype } - pure $ SOAC.Screma w inps (Futhark.mapSOAC outerlam) + SOAC.Screma w inps <$> Futhark.mapSOAC outerlam fixInputs :: (MonadFreshNames m) => 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 @@ -424,15 +424,15 @@ -- | The lambda used in a given SOAC. lambda :: SOAC rep -> Lambda rep lambda (Stream _ _ _ lam) = lam -lambda (Screma _ _ (ScremaForm lam _ _)) = lam +lambda (Screma _ _ (ScremaForm lam _ _ _)) = lam lambda (Hist _ _ _ lam) = lam -- | Set the lambda used in the SOAC. setLambda :: Lambda rep -> SOAC rep -> SOAC rep setLambda lam (Stream w arrs nes _) = Stream w arrs nes lam -setLambda lam (Screma w arrs (ScremaForm _ scan red)) = - Screma w arrs (ScremaForm lam scan red) +setLambda lam (Screma w arrs (ScremaForm _ scan red post_lam)) = + Screma w arrs (ScremaForm lam scan red post_lam) setLambda lam (Hist w ops inps _) = Hist w ops inps lam @@ -529,10 +529,10 @@ -- -- array result and input IDs of the stream's lambda strm_resids <- mapM (newIdent "res") loutps - let insoac = - Futhark.Screma chvar (map paramName strm_inpids) $ - Futhark.mapSOAC lam' - insstm = mkLet strm_resids $ Op insoac + insoac <- + Futhark.Screma chvar (map paramName strm_inpids) + <$> Futhark.mapSOAC lam' + let insstm = mkLet strm_resids $ Op insoac strmbdy = mkBody (oneStm insstm) $ map (subExpRes . Var . identName) strm_resids strmpar = chunk_param : strm_inpids strmlam = Lambda strmpar loutps strmbdy @@ -560,9 +560,9 @@ strmlam <- fmap fst . runBuilder . mkLambda strmpar $ do -- 1. let (scan0_ids,map_resids) = scanomap(scan_lam,nes,map_lam,a_ch) (scan0_ids, map_resids) <- - fmap (splitAt (length scan_arr_ts)) . letTupExp "scan" . Op $ - Futhark.Screma chvar (map paramName strm_inpids) $ - Futhark.scanomapSOAC [Futhark.Scan scan_lam nes] lam' + (fmap (splitAt (length scan_arr_ts)) . letTupExp "scan" . Op) + . Futhark.Screma chvar (map paramName strm_inpids) + =<< Futhark.scanomapSOAC [Futhark.Scan scan_lam nes] lam' -- 2. let outerszm1id = chunksize - 1 outszm1id <- letSubExp "outszm1" . BasicOp $ @@ -589,8 +589,8 @@ let (addlelstm, addlelres) = (bodyStms addlelbdy, bodyResult addlelbdy) -- 4. let strm_resids = map (acc `+`,nes, scan0_ids) strm_resids <- - letTupExp "strm_res" . Op $ - Futhark.Screma chvar scan0_ids (Futhark.mapSOAC maplam) + letTupExp "strm_res" . Op . Futhark.Screma chvar scan0_ids + =<< Futhark.mapSOAC maplam -- 5. let acc' = acc + lasteel_ids addStms addlelstm pure $ addlelres ++ map (subExpRes . Var) (strm_resids ++ map_resids) @@ -615,17 +615,14 @@ inpacc_ids <- mapM (newParam "inpacc") accrtps acc0_ids <- mapM (newIdent "acc0") accrtps -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in - let insoac = - Futhark.Screma - chvar - (map paramName strm_inpids) - $ Futhark.redomapSOAC [Futhark.Reduce comm lamin nes] foldlam - insstm = mkLet (acc0_ids ++ strm_resids) $ Op insoac + insoac <- + Futhark.Screma chvar (map paramName strm_inpids) + <$> Futhark.redomapSOAC [Futhark.Reduce comm lamin nes] foldlam + let insstm = mkLet (acc0_ids ++ strm_resids) $ Op insoac -- 2. let acc' = acc + acc0_ids in addaccbdy <- - mkPlusBnds lamin $ - map Var $ - map paramName inpacc_ids ++ map identName acc0_ids + mkPlusBnds lamin . map Var $ + map paramName inpacc_ids ++ map identName acc0_ids -- Construct the stream let (addaccstm, addaccres) = (bodyStms addaccbdy, bodyResult addaccbdy) strmbdy = diff --git a/src/Futhark/Analysis/Interference.hs b/src/Futhark/Analysis/Interference.hs --- a/src/Futhark/Analysis/Interference.hs +++ b/src/Futhark/Analysis/Interference.hs @@ -161,8 +161,10 @@ analyseKernelBody lumap inuse body analyseSegOp lumap inuse (SegRed _ _ _ body binops) = segWithBinOps lumap inuse binops body -analyseSegOp lumap inuse (SegScan _ _ _ body binops) = do - segWithBinOps lumap inuse binops body +analyseSegOp lumap inuse (SegScan _ _ _ body binops post_op) = do + (inuse', lus', graph') <- segWithBinOps lumap inuse binops body + (inuse'', lus'', graph'') <- analyseSegPostOp lumap inuse' post_op + pure (inuse'', lus' <> lus'', graph' <> graph'') analyseSegOp lumap inuse (SegHist _ _ _ body histops) = do (inuse', lus', graph) <- analyseKernelBody lumap inuse body (inuse'', lus'', graph') <- mconcat <$> mapM (analyseHistOp lumap inuse') histops @@ -191,6 +193,15 @@ SegBinOp GPUMem -> m (InUse, LastUsed, Graph VName) analyseSegBinOp lumap inuse (SegBinOp _ lambda _ _) = + analyseLambda lumap inuse lambda + +analyseSegPostOp :: + (LocalScope GPUMem m) => + LUTabFun -> + InUse -> + SegPostOp GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseSegPostOp lumap inuse (SegPostOp lambda) = analyseLambda lumap inuse lambda analyseHistOp :: diff --git a/src/Futhark/Analysis/LastUse.hs b/src/Futhark/Analysis/LastUse.hs --- a/src/Futhark/Analysis/LastUse.hs +++ b/src/Futhark/Analysis/LastUse.hs @@ -317,11 +317,13 @@ (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms') pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'') -lastUseSegOp (SegScan _ _ tps kbody sbos) used_nms = do +lastUseSegOp (SegScan _ _ tps kbody sbos post_op) used_nms = do + (lutab_spo, lu_vars_spo, used_nms_spo) <- lastUseSegPostOp post_op used_nms + (used_nms', lu_vars) <- lastUsedInNames used_nms_spo $ freeIn tps (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseSegBinOp sbos used_nms - (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps - (body_lutab, used_nms'') <- lastUseKernelBody kbody (mempty, used_nms') - pure (M.union lutab_sbo body_lutab, lu_vars <> lu_vars_sbo, used_nms_sbo <> used_nms' <> used_nms'') + (used_nms'', lu_vars') <- lastUsedInNames used_nms_spo $ freeIn tps + (body_lutab, used_nms''') <- lastUseKernelBody kbody (mempty, used_nms') + pure (M.unions [lutab_spo, lutab_sbo, body_lutab], lu_vars <> lu_vars' <> lu_vars_sbo <> lu_vars_spo, used_nms_spo <> used_nms_sbo <> used_nms' <> used_nms'' <> used_nms''') lastUseSegOp (SegHist _ _ tps kbody hos) used_nms = do (lutab_sbo, lu_vars_sbo, used_nms_sbo) <- lastUseHistOp hos used_nms (used_nms', lu_vars) <- lastUsedInNames used_nms_sbo $ freeIn tps @@ -382,6 +384,17 @@ (used_nms', lu_vars) <- lastUsedInNames used_nms $ freeIn shp <> freeIn rf <> freeIn dest <> freeIn neutral <> freeIn shp' (body_lutab, used_nms'') <- lastUseBody body (mempty, used_nms') pure (body_lutab, lu_vars, used_nms'') + +lastUseSegPostOp :: + (Constraints rep) => + SegPostOp (Aliases rep) -> + Names -> + LastUseM rep (LUTabFun, Names, Names) +lastUseSegPostOp (SegPostOp l@(Lambda _ _ body)) used_nms = + inScopeOf l $ do + (body_lutab, used_nms') <- lastUseBody body (mempty, used_nms) + -- NOTE: Assume this should be mempty? + pure (body_lutab, mempty, used_nms') lastUseSeqOp :: Op (Aliases SeqMem) -> Names -> LastUseM SeqMem (LUTabFun, Names, Names) lastUseSeqOp (Alloc se sp) used_nms = do diff --git a/src/Futhark/Analysis/MemAlias.hs b/src/Futhark/Analysis/MemAlias.hs --- a/src/Futhark/Analysis/MemAlias.hs +++ b/src/Futhark/Analysis/MemAlias.hs @@ -70,7 +70,7 @@ analyzeStms (bodyStms kbody) m analyzeHostOp m (SegOp (SegRed _ _ _ kbody _)) = analyzeStms (bodyStms kbody) m -analyzeHostOp m (SegOp (SegScan _ _ _ kbody _)) = +analyzeHostOp m (SegOp (SegScan _ _ _ kbody _ _)) = analyzeStms (bodyStms kbody) m analyzeHostOp m (SegOp (SegHist _ _ _ kbody _)) = analyzeStms (bodyStms kbody) m diff --git a/src/Futhark/Bench.hs b/src/Futhark/Bench.hs --- a/src/Futhark/Bench.hs +++ b/src/Futhark/Bench.hs @@ -351,11 +351,9 @@ cmdMaybe . liftIO $ cmdPauseProfiling server pure $ Just profile_log - vs <- readResults server outs <* freeOuts - - pure (vs, DL.toList xs, profile_log) + pure (DL.toList xs, profile_log) - (vs, call_logs, profile_log) <- case maybe_call_logs of + (call_logs, profile_log) <- case maybe_call_logs of Nothing -> throwError . T.pack $ "Execution exceeded " ++ show (runTimeout opts) ++ " seconds." @@ -373,8 +371,11 @@ liftIO $ maybe (pure Nothing) (fmap Just . getExpectedValues) expected_spec case maybe_expected of - Just expected -> checkResult program expected vs - Nothing -> pure () + Just expected -> do + vs <- readResults server outs <* freeOuts + checkResult program expected vs + Nothing -> + freeOuts pure ( map fst call_logs, diff --git a/src/Futhark/CLI/LSP.hs b/src/Futhark/CLI/LSP.hs --- a/src/Futhark/CLI/LSP.hs +++ b/src/Futhark/CLI/LSP.hs @@ -1,7 +1,7 @@ {-# LANGUAGE ExplicitNamespaces #-} -- | @futhark lsp@ -module Futhark.CLI.LSP (main) where +module Futhark.CLI.LSP (main, serverDefinition) where import Control.Monad.IO.Class (MonadIO (liftIO)) import Data.IORef (newIORef) @@ -33,36 +33,37 @@ runServer, type (<~>) (Iso), ) -import System.Exit -import System.IO (BufferMode (LineBuffering), hSetBuffering, stderr) +import System.Exit (ExitCode (ExitFailure), exitSuccess, exitWith) -- | Run @futhark lsp@ main :: String -> [String] -> IO () main _prog _args = do - state_mvar <- newIORef emptyState - -- makes the lines appear in full in the logs - hSetBuffering stderr LineBuffering code <- - runServer $ - ServerDefinition - { onConfigChange = const $ pure (), - configSection = "Futhark", - parseConfig = const . const $ Right (), - defaultConfig = (), - doInitialize = \env _req -> pure $ Right env, - staticHandlers = handlers state_mvar, - interpretHandler = \env -> Iso (runLspT env) liftIO, - options = - defaultOptions - { optTextDocumentSync = - Just syncOptions, - optExecuteCommandCommands = - Just $ map showText [minBound :: CommandType .. maxBound] - } - } + runServer =<< serverDefinition case code of 0 -> exitSuccess _ -> exitWith $ ExitFailure code + +serverDefinition :: IO (ServerDefinition ()) +serverDefinition = do + state_mvar <- newIORef emptyState + pure $ + ServerDefinition + { onConfigChange = const $ pure (), + configSection = "Futhark", + parseConfig = const . const $ Right (), + defaultConfig = (), + doInitialize = \env _req -> pure $ Right env, + staticHandlers = handlers state_mvar, + interpretHandler = \env -> Iso (runLspT env) liftIO, + options = + defaultOptions + { optTextDocumentSync = + Just syncOptions, + optExecuteCommandCommands = + Just $ map showText [minBound :: CommandType .. maxBound] + } + } syncOptions :: TextDocumentSyncOptions syncOptions = diff --git a/src/Futhark/CLI/Query.hs b/src/Futhark/CLI/Query.hs --- a/src/Futhark/CLI/Query.hs +++ b/src/Futhark/CLI/Query.hs @@ -25,7 +25,8 @@ putStrLn $ "Position: " ++ locStr (srclocOf loc) case def of Nothing -> pure () - Just (BoundTerm t defloc) -> do + Just (BoundTerm term defloc) -> do + let t = termBindingType term putStrLn $ "Type: " ++ prettyString t putStrLn $ "Definition: " ++ locStr (srclocOf defloc) Just (BoundType defloc) -> 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 @@ -144,9 +144,9 @@ instance Ord TestCase where x `compare` y = testCaseProgram x `compare` testCaseProgram y -data RunResult +data RunResult m = ErrorResult T.Text - | SuccessResult [Value] + | SuccessResult (m [Value]) (m ()) progNotFound :: T.Text -> T.Text progNotFound s = s <> ": command not found" @@ -247,8 +247,8 @@ throwError $ progNotFound $ T.pack futhark _ -> liftExcept $ - compareResult entry index program expectedResult' - =<< runResult program code output err + compareResult entry index program expectedResult' $ + runResult program code output err in accErrors_ $ map runInterpretedCase run_cases runTestCase :: TestCase -> TestM () @@ -366,13 +366,13 @@ call_r <- liftIO $ cmdCall server entry outs ins liftCommand $ cmdFree server ins - res <- case call_r of - Left (CmdFailure _ err) -> - pure $ ErrorResult $ T.unlines err - Right _ -> - SuccessResult - <$> readResults server outs - <* liftCommand (cmdFree server outs) + let res = case call_r of + Left (CmdFailure _ err) -> + ErrorResult $ T.unlines err + Right _ -> + SuccessResult + (readResults server outs <* liftCommand (cmdFree server outs)) + (liftCommand (cmdFree server outs)) compareResult entry index program expected res @@ -393,16 +393,18 @@ ExitCode -> SBS.ByteString -> SBS.ByteString -> - m RunResult + RunResult m runResult program ExitSuccess stdout_s _ = - case valuesFromByteString "stdout" $ LBS.fromStrict stdout_s of - Left e -> do - let actualf = program `addExtension` "actual" - liftIO $ SBS.writeFile actualf stdout_s - E.throwError $ T.pack e <> "\n(See " <> T.pack actualf <> ")" - Right vs -> pure $ SuccessResult vs + SuccessResult fetch (pure ()) + where + fetch = case valuesFromByteString "stdout" $ LBS.fromStrict stdout_s of + Left e -> do + let actualf = program `addExtension` "actual" + liftIO $ SBS.writeFile actualf stdout_s + E.throwError $ T.pack e <> "\n(See " <> T.pack actualf <> ")" + Right vs -> pure vs runResult _ (ExitFailure _) _ stderr_s = - pure $ ErrorResult $ T.decodeUtf8 stderr_s + ErrorResult $ T.decodeUtf8 stderr_s compileTestProgram :: [String] -> FutharkExe -> String -> FilePath -> [WarningTest] -> TestM () compileTestProgram extra_options futhark backend program warnings = do @@ -417,11 +419,12 @@ Int -> FilePath -> ExpectedResult [Value] -> - RunResult -> + RunResult m -> m () -compareResult _ _ _ (Succeeds Nothing) SuccessResult {} = - pure () -compareResult entry index program (Succeeds (Just expected_vs)) (SuccessResult actual_vs) = +compareResult _ _ _ (Succeeds Nothing) (SuccessResult _ free_vs) = + free_vs +compareResult entry index program (Succeeds (Just expected_vs)) (SuccessResult get_actual_vs _) = do + actual_vs <- get_actual_vs checkResult (program <.> T.unpack entry <.> show index) expected_vs @@ -430,7 +433,8 @@ checkError expectedError actualError compareResult _ _ _ (Succeeds _) (ErrorResult err) = E.throwError $ "Function failed with error:\n" <> err -compareResult _ _ _ (RunTimeFailure f) (SuccessResult _) = +compareResult _ _ _ (RunTimeFailure f) (SuccessResult _ free_vs) = do + free_vs E.throwError $ "Program succeeded, but expected failure:\n " <> showText f --- diff --git a/src/Futhark/CodeGen/Backends/GPU.hs b/src/Futhark/CodeGen/Backends/GPU.hs --- a/src/Futhark/CodeGen/Backends/GPU.hs +++ b/src/Futhark/CodeGen/Backends/GPU.hs @@ -79,7 +79,7 @@ where field = "max_" <> prettyString size_class kernelConstToExp (SizeUserParam name def) = - [C.cexp|$exp:name'.set ? $exp:name'.val : $exp:def|] + [C.cexp|$exp:name'.set ? $exp:name'.val : $exp:(compileConstExp def)|] where name' = getParamByKey name diff --git a/src/Futhark/CodeGen/Backends/GenericC/Server.hs b/src/Futhark/CodeGen/Backends/GenericC/Server.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Server.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Server.hs @@ -142,24 +142,55 @@ -- First component is forward declaration so we don't have to worry -- about ordering. typeBoilerplate :: Manifest -> (T.Text, Type) -> (C.Definition, C.Initializer, [C.Definition]) -typeBoilerplate _ (tname, TypeArray c_type_name et rank ops) = +typeBoilerplate manifest (tname, TypeArray c_type_name et rank ops) = let element_type_name = typeStructName et + element_c_type = cType manifest et type_name = typeStructName tname array_name = type_name <> "_array" aux_name = type_name <> "_aux" info_name = et <> "_info" + aux_array_new_wrap = arrayNew ops <> "_aux_wrap" array_new_wrap = arrayNew ops <> "_wrap" + array_set = arrayNew ops <> "_set" array_index_wrap = arrayIndex ops <> "_wrap" shape_args = [[C.cexp|shape[$int:i]|] | i <- [0 .. rank - 1]] is_args = [[C.cexp|is[$int:i]|] | i <- [0 .. rank - 1]] in ( [C.cedecl|const struct type $id:type_name;|], [C.cinit|&$id:type_name|], [C.cunit| - void* $id:array_new_wrap(struct futhark_context *ctx, - const void* p, - const typename int64_t* shape) { + void* $id:aux_array_new_wrap(struct futhark_context *ctx, + const void* p, + const typename int64_t* shape) { return $id:(arrayNew ops)(ctx, p, $args:shape_args); } + int $id:array_new_wrap(struct futhark_context* ctx, + typename $id:c_type_name* outp, + $ty:element_c_type *ps[], + const typename int64_t* shape) { + typename int64_t n_values = 1; + for (int i = 0; i < $int:rank; ++i) { + n_values *= shape[i]; + } + $ty:element_c_type *values = alloca(n_values * sizeof($ty:element_c_type)); + for (typename int64_t i = 0; i < n_values; ++i) { + values[i] = *ps[i]; + } + *outp = $id:(arrayNew ops)(ctx, values, $args:shape_args); + return 0; + } + int $id:array_set(struct futhark_context *ctx, + typename $id:c_type_name arr, + $ty:element_c_type *val, + const typename int64_t *is) { + const typename int64_t *shape = $id:(arrayShape ops)(ctx, arr); + typename uint64_t idx = is[0]; + for (int i = 1; i < $int:rank; ++i) { + idx *= shape[i-1]; + idx += is[i]; + } + (($ty:element_c_type*)$id:(arrayValuesRaw ops)(ctx, arr))[idx] = *val; + return 0; + } int $id:array_index_wrap(struct futhark_context *ctx, void *dest, typename $id:c_type_name arr, @@ -170,6 +201,7 @@ .rank = $int:rank, .element_type = &$id:element_type_name, .new = (typename array_new_fn)$id:array_new_wrap, + .set = (typename array_set_fn)$id:array_set, .shape = (typename array_shape_fn)$id:(arrayShape ops), .index = (typename array_index_fn)$id:array_index_wrap, }; @@ -177,10 +209,10 @@ .name = $string:(T.unpack tname), .rank = $int:rank, .info = &$id:info_name, - .new = (typename array_new_fn)$id:array_new_wrap, - .free = (typename array_free_fn)$id:(arrayFree ops), - .shape = (typename array_shape_fn)$id:(arrayShape ops), - .values = (typename array_values_fn)$id:(arrayValues ops) + .new = (typename aux_array_new_fn)$id:aux_array_new_wrap, + .free = (typename aux_array_free_fn)$id:(arrayFree ops), + .shape = (typename aux_array_shape_fn)$id:(arrayShape ops), + .values = (typename aux_array_values_fn)$id:(arrayValues ops) }; const struct type $id:type_name = { .name = $string:(T.unpack tname), @@ -303,26 +335,51 @@ [C.cinit|&$id:sum_name|], Sum ) - transparentDefs type_name (Just (OpaqueArray ops')) = opaqueArrayDefs type_name (opaqueArrayRank ops') (opaqueArrayElemType ops') (opaqueArrayShape ops') (opaqueArrayIndex ops') - transparentDefs type_name (Just (OpaqueRecordArray ops')) = opaqueArrayDefs type_name (recordArrayRank ops') (recordArrayElemType ops') (recordArrayShape ops') (recordArrayIndex ops') + transparentDefs type_name (Just (OpaqueArray ops')) = opaqueArrayDefs type_name (opaqueArrayRank ops') (opaqueArrayElemType ops') (opaqueArrayNew ops') (opaqueArraySet ops') (opaqueArrayShape ops') (opaqueArrayIndex ops') + transparentDefs type_name (Just (OpaqueRecordArray ops')) = opaqueArrayDefs type_name (recordArrayRank ops') (recordArrayElemType ops') (recordArrayNew ops') (recordArraySet ops') (recordArrayShape ops') (recordArrayIndex ops') transparentDefs _ _ = ([], [C.cinit|NULL|], Opaque) - opaqueArrayDefs type_name rank et shape index = + opaqueArrayDefs type_name rank et new set shape index = let array_name = type_name <> "_array" element_type_name = typeStructName et + element_c_type = cType manifest et + new_wrap = new <> "_wrap" + set_wrap = set <> "_wrap" index_wrap = index <> "_wrap" + shape_args = [[C.cexp|shape[$int:i]|] | i <- [0 .. rank - 1]] is_args = [[C.cexp|is[$int:i]|] | i <- [0 .. rank - 1]] in ( [C.cunit| + int $id:new_wrap(struct futhark_context *ctx, + typename $id:c_type_name *outp, + $ty:element_c_type *ps[], + const typename int64_t shape[]) { + typename int64_t n_values = 1; + for (int i = 0; i < $int:rank; ++i) { + n_values *= shape[i]; + } + $ty:element_c_type *values = alloca(n_values * sizeof($ty:element_c_type)); + for (typename int64_t i = 0; i < n_values; ++i) { + values[i] = *ps[i]; + } + return $id:new(ctx, outp, values, $args:shape_args); + } + int $id:set_wrap(struct futhark_context *ctx, + typename $id:c_type_name arr, + $ty:element_c_type *val, + const typename int64_t *is) { + return $id:set(ctx, arr, *val, $args:is_args); + } int $id:index_wrap(struct futhark_context *ctx, - void *dest, - typename $id:c_type_name arr, - const typename int64_t *is) { + void *dest, + typename $id:c_type_name arr, + const typename int64_t *is) { return $id:index(ctx, dest, arr, $args:is_args); } const struct array $id:array_name = { .rank = $int:rank, .element_type = &$id:element_type_name, - .new = NULL, + .new = (typename array_new_fn)$id:new_wrap, + .set = (typename array_set_fn)$id:set_wrap, .shape = (typename array_shape_fn)$id:shape, .index = (typename array_index_fn)$id:index_wrap, };|], diff --git a/src/Futhark/CodeGen/Backends/GenericWASM.hs b/src/Futhark/CodeGen/Backends/GenericWASM.hs --- a/src/Futhark/CodeGen/Backends/GenericWASM.hs +++ b/src/Futhark/CodeGen/Backends/GenericWASM.hs @@ -7,9 +7,12 @@ GC.asServer, EntryPointType, JSEntryPoint (..), + JSRecordField (..), + JSOpaqueType (..), emccExportNames, javascriptWrapper, extToString, + opaqueToJS, runServer, libraryExports, ) @@ -18,10 +21,11 @@ import Data.List (intercalate) import Data.Text qualified as T import Futhark.CodeGen.Backends.GenericC qualified as GC -import Futhark.CodeGen.Backends.SimpleRep (opaqueName) +import Futhark.CodeGen.Backends.SimpleRep (isValidCName, opaqueName) import Futhark.CodeGen.ImpCode.Sequential qualified as Imp import Futhark.CodeGen.RTS.JavaScript -import Futhark.Util (nubOrd, showText) +import Futhark.Util (nubOrd, showText, zEncodeText) +import Language.Futhark.Core (nameToText) import Language.Futhark.Primitive import NeatInterpolation (text) @@ -51,8 +55,71 @@ ret :: [EntryPointType] } -emccExportNames :: [JSEntryPoint] -> [String] -emccExportNames jses = +-- | A field in a JavaScript record opaque type. +data JSRecordField = JSRecordField + { -- | Original field name (used in the server protocol). + jsrfName :: String, + -- | JavaScript type string (e.g. @"i32"@, @"[]f64"@, @"opaque_foo"@). + jsrfType :: EntryPointType, + -- | JavaScript method name on @FutharkContext@ (e.g. @"project_opaque_foo_x"@). + jsrfProjectFn :: String + } + +-- | How an opaque type is represented for the JavaScript server. +data JSOpaqueType + = -- | A record type with named, projectable fields. + JSOpaqueRecord [JSRecordField] + | -- | Any other opaque type (no project support). + JSOpaqueOther + +-- | Convert the opaque-type table from an ImpCode program into the +-- JavaScript-oriented representation used by the code generator. +opaqueToJS :: Imp.OpaqueTypes -> [(String, JSOpaqueType)] +opaqueToJS (Imp.OpaqueTypes types) = map convertOne types + where + convertOne (desc, Imp.OpaqueRecord fields) = + (T.unpack (opaqueName desc), JSOpaqueRecord (map (convertField desc) fields)) + convertOne (desc, _) = + (T.unpack (opaqueName desc), JSOpaqueOther) + + convertField desc (fname, etype) = + JSRecordField + { jsrfName = T.unpack (nameToText fname), + jsrfType = entryTypeToJSString etype, + jsrfProjectFn = + "project_" + ++ T.unpack (opaqueName desc) + ++ "_" + ++ T.unpack f' + } + where + f' + | isValidCName (opaqueName desc <> "_" <> nameToText fname) = + nameToText fname + | otherwise = zEncodeText (nameToText fname) + +-- | Convert an 'Imp.EntryPointType' to its JavaScript string representation. +entryTypeToJSString :: Imp.EntryPointType -> String +entryTypeToJSString (Imp.TypeOpaque desc) = T.unpack $ opaqueName desc +entryTypeToJSString (Imp.TypeTransparent (Imp.ValueType sign (Imp.Rank rank) pt)) = + concat (replicate rank "[]") ++ primToJSString sign pt + where + primToJSString _ (FloatType Float16) = "f16" + primToJSString _ (FloatType Float32) = "f32" + primToJSString _ (FloatType Float64) = "f64" + primToJSString Imp.Signed (IntType Int8) = "i8" + primToJSString Imp.Signed (IntType Int16) = "i16" + primToJSString Imp.Signed (IntType Int32) = "i32" + primToJSString Imp.Signed (IntType Int64) = "i64" + primToJSString Imp.Unsigned (IntType Int8) = "u8" + primToJSString Imp.Unsigned (IntType Int16) = "u16" + primToJSString Imp.Unsigned (IntType Int32) = "u32" + primToJSString Imp.Unsigned (IntType Int64) = "u64" + primToJSString _ Bool = "bool" + primToJSString _ Unit = error "entryTypeToJSString: Unit" + +emccExportNames :: [JSEntryPoint] -> [(String, JSOpaqueType)] -> [String] +emccExportNames jses opaqueTypes = map (\jse -> "'_futhark_entry_" ++ T.unpack (GC.escapeName (T.pack (name jse))) ++ "'") jses ++ map (\arg -> "'" ++ gfn "new" arg ++ "'") arrays ++ map (\arg -> "'" ++ gfn "free" arg ++ "'") arrays @@ -60,6 +127,7 @@ ++ map (\arg -> "'" ++ gfn "values_raw" arg ++ "'") arrays ++ map (\arg -> "'" ++ gfn "values" arg ++ "'") arrays ++ map (\arg -> "'" ++ "_futhark_free_" ++ arg ++ "'") opaques + ++ map (\rf -> "'_futhark_" ++ jsrfProjectFn rf ++ "'") allRecordFields ++ [ "_futhark_context_config_new", "_futhark_context_config_free", "_futhark_context_new", @@ -67,29 +135,35 @@ "_futhark_context_get_error" ] where - arrays = filter isArray typs - opaques = filter isOpaque typs - typs = nubOrd $ concatMap (\jse -> parameters jse ++ ret jse) jses + -- Include array types from both entry points and record fields. + arrays = nubOrd $ filter isArray (entryPointTypes ++ recordFieldTypes) + -- Include opaque types from both entry points and record fields. + opaques = nubOrd $ filter isOpaque (entryPointTypes ++ recordFieldTypes) + entryPointTypes = concatMap (\jse -> parameters jse ++ ret jse) jses + recordFieldTypes = [jsrfType rf | (_, JSOpaqueRecord fields) <- opaqueTypes, rf <- fields] gfn typ str = "_futhark_" ++ typ ++ "_" ++ baseType str ++ "_" ++ show (dim str) ++ "d" + allRecordFields = [rf | (_, JSOpaqueRecord fields) <- opaqueTypes, rf <- fields] -javascriptWrapper :: [JSEntryPoint] -> T.Text -javascriptWrapper entryPoints = +javascriptWrapper :: [JSEntryPoint] -> [(String, JSOpaqueType)] -> T.Text +javascriptWrapper entryPoints opaqueTypes = T.unlines [ serverJs, valuesJs, wrapperclassesJs, - classFutharkContext entryPoints + classFutharkContext entryPoints opaqueTypes ] -classFutharkContext :: [JSEntryPoint] -> T.Text -classFutharkContext entryPoints = +classFutharkContext :: [JSEntryPoint] -> [(String, JSOpaqueType)] -> T.Text +classFutharkContext entryPoints opaqueTypes = T.unlines [ "class FutharkContext {", - constructor entryPoints, + constructor entryPoints opaqueTypes, getFreeFun, getEntryPointsFun, + getTypesFun, getErrorFun, T.unlines $ map toFutharkArray arrays, + T.unlines $ concatMap (generateProjectMethods . snd) opaqueTypes, T.unlines $ map jsWrapEntryPoint entryPoints, "}", [text| @@ -100,11 +174,14 @@ |] ] where - arrays = filter isArray typs - typs = nubOrd $ concatMap (\jse -> parameters jse ++ ret jse) entryPoints + -- Collect array types from entry points AND from record fields so the + -- array constructors are always available when returning projected values. + arrays = nubOrd $ filter isArray (entryPointTypes ++ recordFieldTypes) + entryPointTypes = concatMap (\jse -> parameters jse ++ ret jse) entryPoints + recordFieldTypes = [jsrfType rf | (_, JSOpaqueRecord fields) <- opaqueTypes, rf <- fields] -constructor :: [JSEntryPoint] -> T.Text -constructor jses = +constructor :: [JSEntryPoint] -> [(String, JSOpaqueType)] -> T.Text +constructor jses opaqueTypes = [text| constructor(wasm, num_threads) { this.wasm = wasm; @@ -114,10 +191,14 @@ this.entry_points = { ${entries} }; + this.types = { + ${type_entries} + }; } |] where entries = T.intercalate "," $ map dicEntry jses + type_entries = T.intercalate "," $ map dicTypeEntry opaqueTypes getFreeFun :: T.Text getFreeFun = @@ -136,13 +217,21 @@ } |] +getTypesFun :: T.Text +getTypesFun = + [text| + get_types() { + return this.types; + } + |] + getErrorFun :: T.Text getErrorFun = [text| get_error() { var ptr = this.wasm._futhark_context_get_error(this.ctx); - var len = HEAP8.subarray(ptr).indexOf(0); - var str = String.fromCharCode(...HEAP8.subarray(ptr, ptr + len)); + var len = this.wasm.HEAP8.subarray(ptr).indexOf(0); + var str = String.fromCharCode(...this.wasm.HEAP8.subarray(ptr, ptr + len)); this.wasm._free(ptr); return str; } @@ -158,6 +247,54 @@ ename = T.pack $ name jse params = showText $ parameters jse rets = showText $ ret jse + +-- | Generate the @types@ dictionary entry for a single opaque type. +dicTypeEntry :: (String, JSOpaqueType) -> T.Text +dicTypeEntry (tname, JSOpaqueRecord fields) = + T.pack $ + show tname + ++ ": [\"record\", [" + ++ intercalate ", " (map fieldEntry fields) + ++ "]]" + where + fieldEntry rf = + "[" + ++ show (jsrfName rf) + ++ ", " + ++ show (jsrfType rf) + ++ ", \"" + ++ jsrfProjectFn rf + ++ "\"]" +dicTypeEntry (tname, JSOpaqueOther) = + T.pack $ show tname ++ ": [\"opaque\"]" + +-- | Generate @project_*@ wrapper methods for a record opaque type. +generateProjectMethods :: JSOpaqueType -> [T.Text] +generateProjectMethods (JSOpaqueRecord fields) = map generateProjectMethod fields +generateProjectMethods JSOpaqueOther = [] + +generateProjectMethod :: JSRecordField -> T.Text +generateProjectMethod rf = + T.pack $ + unlines + [ " " ++ jsrfProjectFn rf ++ "(obj) {", + " var out = this.wasm._malloc(" ++ show (typeSize ftype) ++ ");", + " this.wasm._futhark_" ++ jsrfProjectFn rf ++ "(this.ctx, out, obj.ptr);", + " var result = " ++ readResult ++ ";", + " this.wasm._free(out);", + " return result;", + " }" + ] + where + ftype = jsrfType rf + readout = typeHeap ftype ++ "[out >> " ++ show (typeShift ftype) ++ "]" + readResult + | isArray ftype = + "this.new_" ++ baseType ftype ++ "_" ++ show (dim ftype) ++ "d_from_ptr(" ++ readout ++ ")" + | isOpaque ftype = + "new FutharkOpaque(this, " ++ readout ++ ", this.wasm._futhark_free_" ++ ftype ++ ")" + | ftype == "bool" = readout ++ "!==0" + | otherwise = readout jsWrapEntryPoint :: JSEntryPoint -> T.Text jsWrapEntryPoint jse = diff --git a/src/Futhark/CodeGen/Backends/MulticoreWASM.hs b/src/Futhark/CodeGen/Backends/MulticoreWASM.hs --- a/src/Futhark/CodeGen/Backends/MulticoreWASM.hs +++ b/src/Futhark/CodeGen/Backends/MulticoreWASM.hs @@ -58,8 +58,8 @@ pure ( ws, ( prog'', - javascriptWrapper (fRepMyRep prog'), - "_futhark_context_config_set_num_threads" : emccExportNames (fRepMyRep prog') + javascriptWrapper (fRepMyRep prog') (opaqueToJS (Imp.defTypes prog')), + "_futhark_context_config_set_num_threads" : emccExportNames (fRepMyRep prog') (opaqueToJS (Imp.defTypes prog')) ) ) diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL.hs b/src/Futhark/CodeGen/Backends/PyOpenCL.hs --- a/src/Futhark/CodeGen/Backends/PyOpenCL.hs +++ b/src/Futhark/CodeGen/Backends/PyOpenCL.hs @@ -212,7 +212,7 @@ Call (Field (Var "self.user_params") "get") [ Arg $ String (nameToText name), - Arg $ Var $ compileName def + Arg $ compileConstExp def ] compileConstExp :: Imp.KernelConstExp -> PyExp 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 @@ -46,7 +46,7 @@ Call (Field (Var "self.user_params") "get") [ Arg $ String (nameToText name), - Arg $ Var $ Py.compileName def + Arg $ compileConstExp def ] -- | Python code (as a string) that calls the diff --git a/src/Futhark/CodeGen/Backends/SequentialWASM.hs b/src/Futhark/CodeGen/Backends/SequentialWASM.hs --- a/src/Futhark/CodeGen/Backends/SequentialWASM.hs +++ b/src/Futhark/CodeGen/Backends/SequentialWASM.hs @@ -49,7 +49,7 @@ (DefaultSpace, [DefaultSpace]) [] prog' - pure (ws, (prog'', javascriptWrapper (fRepMyRep prog'), emccExportNames (fRepMyRep prog'))) + pure (ws, (prog'', javascriptWrapper (fRepMyRep prog') (opaqueToJS (Imp.defTypes prog')), emccExportNames (fRepMyRep prog') (opaqueToJS (Imp.defTypes prog')))) where operations :: GC.Operations Imp.Sequential () operations = diff --git a/src/Futhark/CodeGen/ImpCode/GPU.hs b/src/Futhark/CodeGen/ImpCode/GPU.hs --- a/src/Futhark/CodeGen/ImpCode/GPU.hs +++ b/src/Futhark/CodeGen/ImpCode/GPU.hs @@ -36,7 +36,7 @@ data KernelConst = SizeConst Name SizeClass | SizeMaxConst SizeClass - | SizeUserParam Name VName + | SizeUserParam Name KernelConstExp deriving (Eq, Ord, Show) -- | An expression whose variables are kernel constants. diff --git a/src/Futhark/CodeGen/ImpGen/GPU.hs b/src/Futhark/CodeGen/ImpGen/GPU.hs --- a/src/Futhark/CodeGen/ImpGen/GPU.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU.hs @@ -178,8 +178,8 @@ compileSegMap pat lvl space kbody segOpCompiler pat (SegRed lvl@(SegThread _ _) space _ kbody reds) = compileSegRed pat lvl space reds kbody -segOpCompiler pat (SegScan lvl@(SegThread _ _) space _ kbody scans) = - compileSegScan pat lvl space scans kbody +segOpCompiler pat (SegScan lvl@(SegThread _ _) space ts kbody scans post_op) = + compileSegScan pat lvl space ts kbody scans post_op segOpCompiler pat (SegHist lvl@(SegThread _ _) space _ kbody ops) = compileSegHist pat lvl space ops kbody segOpCompiler pat segop = diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs @@ -33,8 +33,8 @@ genZeroes, isPrimParam, kernelConstToExp, - getChunkSize, getSize, + isConstExp, -- * Host-level bulk operations sReplicate, @@ -295,31 +295,9 @@ pure v f (Imp.SizeUserParam name def) = do v <- dPrimS (nameFromText $ prettyText name) int64 - emit $ Imp.GetUserParam v name $ le64 def + def' <- kernelConstToExp def + emit $ Imp.GetUserParam v name $ isInt64 def' pure v - --- | Given available register and a list of parameter types, compute --- the largest available chunk size given the parameters for which we --- want chunking and the available resources. Used in --- 'SegScan.SinglePass.compileSegScan', and 'SegRed.compileSegRed' --- (with primitive non-commutative operators only). -getChunkSize :: [Type] -> Imp.KernelConstExp -getChunkSize types = do - let max_tblock_size = Imp.SizeMaxConst SizeThreadBlock - max_block_mem = Imp.SizeMaxConst SizeSharedMemory - max_block_reg = Imp.SizeMaxConst SizeRegisters - k_mem = le64 max_block_mem `quot` le64 max_tblock_size - k_reg = le64 max_block_reg `quot` le64 max_tblock_size - types' = map elemType $ filter primType types - sizes = map primByteSize types' - - sum_sizes = sum sizes - sum_sizes' = sum (map (sMax64 4 . primByteSize) types') `quot` 4 - max_size = maximum sizes - - mem_constraint = max k_mem sum_sizes `quot` max_size - reg_constraint = (k_reg - 1 - sum_sizes') `quot` (2 * sum_sizes') - untyped $ sMax64 1 $ sMin64 mem_constraint reg_constraint inChunkScan :: KernelConstants -> diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Block.hs b/src/Futhark/CodeGen/ImpGen/GPU/Block.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/Block.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/Block.hs @@ -365,24 +365,47 @@ zipWithM_ (compileThreadResult space) (patElems pat) $ bodyResult body sOp $ Imp.ErrorSync Imp.FenceLocal -compileBlockOp pat (Inner (SegOp (SegScan lvl space _ body scans))) = do +compileBlockOp pat (Inner (SegOp (SegScan lvl space _ body scans post_op))) = do compileFlatId space let (ltids, dims) = unzip $ unSegSpace space dims' = map pe64 dims + [scan] = scans + shpT op = (segBinOpShape op,) <$> lambdaReturnType (segBinOpLambda op) + scan_ts = concatMap shpT scans + shpOfT t s = + arrayShape $ + foldr (flip arrayOfRow) (arrayOfShape t s) $ + segSpaceDims space + (scan_pars, map_pars) = + splitAt (length $ segBinOpNeutral scan) $ lambdaParams $ segPostOpLambda post_op + (body_scan_res, body_map_res) = + splitAt (length $ segBinOpNeutral scan) $ bodyResult body + scan_out <- forM scan_ts $ \(s, t) -> + sAllocArray "scan_out" (elemType t) (shpOfT t s) $ Space "shared" + + dScope Nothing $ scopeOfLParams $ lambdaParams $ segPostOpLambda post_op + blockCoverSegSpace (segVirt lvl) space $ - compileStms mempty (bodyStms body) $ - forM_ (zip (patNames pat) $ bodyResult body) $ \(dest, res) -> + compileStms mempty (bodyStms body) $ do + forM_ (zip scan_out body_scan_res) $ \(dest, res) -> copyDWIMFix dest (map Imp.le64 ltids) (kernelResultSubExp res) [] - fence <- fenceForArrays $ patNames pat - sOp $ Imp.ErrorSync fence + sComment "bind map results to post lambda params" $ + forM_ (zip map_pars body_map_res) $ \(par, res) -> + copyDWIMFix + (paramName par) + [] + (kernelResultSubExp res) + [] + sOp $ Imp.ErrorSync Imp.FenceLocal + let segment_size = last dims' crossesSegment from to = (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size) @@ -391,12 +414,8 @@ -- array of scan elements, so we invent some new flattened arrays -- here. dims_flat <- dPrimV "dims_flat" $ product dims' - let scan = head scans - num_scan_results = length $ segBinOpNeutral scan arrs_flat <- - mapM (flattenArray (length dims') dims_flat) $ - take num_scan_results $ - patNames pat + mapM (flattenArray (length dims') dims_flat) scan_out case segVirt lvl of SegVirt -> @@ -412,6 +431,22 @@ (product dims') (segBinOpLambda scan) arrs_flat + + -- FIXME: we actually need something like blockCoverSegSpacee here, although in + -- practice we currently do not generate virtualised scans. + sWhen (isActive $ zip ltids dims) . localOps threadOperations $ do + sComment "bind scan results to post lambda params" $ + forM_ (zip scan_pars scan_out) $ \(par, acc) -> + copyDWIMFix (paramName par) [] (Var acc) (map Imp.le64 ltids) + + let res = fmap resSubExp $ bodyResult $ lambdaBody $ segPostOpLambda post_op + sComment "compute post op." $ + compileStms mempty (bodyStms $ lambdaBody $ segPostOpLambda post_op) $ do + sComment "write values" $ + forM_ (zip (patElems pat) res) $ \(pe, subexp) -> + copyDWIMFix (patElemName pe) (map le64 ltids) subexp [] + + sOp $ Imp.Barrier Imp.FenceGlobal compileBlockOp pat (Inner (SegOp (SegRed lvl space _ body ops))) = do compileFlatId space diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs @@ -73,6 +73,27 @@ forM2_ :: (Monad m) => [a] -> [b] -> (a -> b -> m c) -> m () forM2_ xs ys f = forM_ (zip xs ys) (uncurry f) +-- | Given available register and a list of parameter types, compute +-- the largest available chunk size given the parameters for which we +-- want chunking and the available resources. +getRedChunkSize :: [Type] -> Imp.KernelConstExp +getRedChunkSize types = do + let max_tblock_size = Imp.SizeMaxConst SizeThreadBlock + max_block_mem = Imp.SizeMaxConst SizeSharedMemory + max_block_reg = Imp.SizeMaxConst SizeRegisters + k_mem = le64 max_block_mem `quot` le64 max_tblock_size + k_reg = le64 max_block_reg `quot` le64 max_tblock_size + types' = map elemType $ filter primType types + sizes = map primByteSize types' + + sum_sizes = sum sizes + sum_sizes' = sum (map (sMax64 4 . primByteSize) types') `quot` 4 + max_size = maximum sizes + + mem_constraint = max k_mem sum_sizes `quot` max_size + reg_constraint = (k_reg - 1 - sum_sizes') `quot` (2 * sum_sizes') + untyped $ sMax64 1 $ sMin64 mem_constraint reg_constraint + -- | The maximum number of operators we support in a single SegRed. -- This limit arises out of the static allocation of counters. maxNumOps :: Int @@ -174,7 +195,7 @@ chunk_const = if Noncommutative `elem` map segBinOpComm segbinops && all isPrimSegBinOp segbinops - then getChunkSize param_types + then getRedChunkSize param_types else Imp.ValueExp $ IntValue $ intValue Int64 (1 :: Int64) -- | Prepare intermediate arrays for the reduction. Prim-typed diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs @@ -67,19 +67,21 @@ Pat LetDecMem -> SegLevel -> SegSpace -> - [SegBinOp GPUMem] -> + [Type] -> KernelBody GPUMem -> + [SegBinOp GPUMem] -> + SegPostOp GPUMem -> CallKernelGen () -compileSegScan pat lvl space scan_ops map_kbody = +compileSegScan pat lvl space ts map_kbody scan_ops post_op = sWhen (0 .<. n) $ do emit $ Imp.DebugPrint "\n# SegScan" Nothing target <- hostTarget <$> askEnv case (targetSupportsSinglePass target, canBeSinglePass scan_ops) of (True, Just scan_ops') -> - SinglePass.compileSegScan pat lvl space scan_ops' map_kbody + SinglePass.compileSegScan pat lvl space ts scan_ops' map_kbody post_op _ -> - TwoPass.compileSegScan pat lvl space scan_ops map_kbody + TwoPass.compileSegScan pat lvl space ts scan_ops map_kbody post_op emit $ Imp.DebugPrint "" Nothing where n = product $ map pe64 $ segSpaceDims space diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs @@ -6,7 +6,7 @@ module Futhark.CodeGen.ImpGen.GPU.SegScan.SinglePass (compileSegScan) where import Control.Monad -import Data.List (zip4, zip7) +import Data.List (zip4, zip5, zip7) import Data.Map qualified as M import Data.Maybe import Futhark.CodeGen.ImpCode.GPU qualified as Imp @@ -15,7 +15,7 @@ import Futhark.IR.GPUMem import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.Transform.Rename -import Futhark.Util (mapAccumLM, takeLast) +import Futhark.Util (mapAccumLM) import Futhark.Util.IntegralExp (IntegralExp (mod, rem), divUp, nextMul, quot) import Prelude hiding (mod, quot, rem) @@ -25,6 +25,28 @@ yParams scan = drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) +-- | Given available register, thread block size, scan parameter +-- types, and map parameter types, compute the largest available chunk +-- size given the parameters for which we want chunking and the +-- available resources. +getScanChunkSize :: [Type] -> Imp.KernelConstExp +getScanChunkSize types = + let max_tblock_size = Imp.SizeMaxConst SizeThreadBlock + max_block_mem = Imp.SizeMaxConst SizeSharedMemory + max_block_reg = Imp.SizeMaxConst SizeRegisters + k_mem = le64 max_block_mem `quot` le64 max_tblock_size + k_reg = le64 max_block_reg `quot` le64 max_tblock_size + types' = map elemType $ filter primType types + sizes = map primByteSize types' + + sum_sizes = sum sizes + sum_sizes' = sum (map (sMax64 4 . primByteSize) types') `quot` 4 + max_size = maximum sizes + + mem_constraint = max k_mem sum_sizes `quot` max_size + reg_constraint = (k_reg - 1 - sum_sizes') `quot` (2 * sum_sizes') + in untyped $ sMax64 1 $ sMin64 mem_constraint reg_constraint + createLocalArrays :: Count BlockSize SubExp -> SubExp -> @@ -211,34 +233,134 @@ copyDWIMFix arr [ltid] (Var $ paramName x) [] copyDWIM (paramName y) [] (Var $ paramName x) [] +-- | Calculate the number of u64 words needed to store n bits +bitArrayWords :: Imp.KernelConstExp -> Imp.KernelConstExp +bitArrayWords n = untyped $ isInt64 n `divUp` 64 + +-- | Set a bit in a bit array stored as u64 words +setBitInBitArray :: Imp.TExp Int64 -> VName -> Imp.TExp Int64 -> Imp.TExp Bool -> InKernelGen () +setBitInBitArray chunk bit_array idx bool_val = do + word_idx <- dPrimV "word_idx" (0 :: Imp.TExp Int64) + bit_idx <- dPrimV "bit_idx" idx + + sIf + (chunk .<=. 64) + ( do + word_idx <-- 0 + bit_idx <-- idx + ) + ( do + word_idx <-- idx `quot` 64 + bit_idx <-- idx `rem` 64 + ) + + -- Read current word + current_word <- dPrimV "current_word" (0 :: Imp.TExp Int64) + copyDWIMFix (tvVar current_word) [] (Var bit_array) [tvExp word_idx] + + bool_as_int <- dPrimVE "bool_as_int" $ fromBoolExp bool_val + + -- Create mask and update word + let bit_mask = 1 .<<. tvExp bit_idx + cleared_word = tvExp current_word .&. (bit_mask .^. (-1)) + set_bit = (bool_as_int .&. 1) .<<. tvExp bit_idx + new_word = cleared_word .|. set_bit + + new_word_var <- dPrimV "new_word" new_word + copyDWIMFix bit_array [tvExp word_idx] (Var $ tvVar new_word_var) [] + +-- | Get a bit from a bit array stored as u64 words, storing result in destination +getBitFromBitArray :: Imp.TExp Int64 -> VName -> VName -> Imp.TExp Int64 -> InKernelGen () +getBitFromBitArray chunk dest bit_array idx = do + word_idx <- dPrimV "word_idx" (0 :: Imp.TExp Int64) + bit_idx <- dPrimV "bit_idx" idx + + sIf + (chunk .<=. 64) + ( do + word_idx <-- 0 + bit_idx <-- idx + ) + ( do + word_idx <-- idx `quot` 64 + bit_idx <-- idx `rem` 64 + ) + + -- Read word containing the bit + word <- dPrimV "bit_word_read" (0 :: Imp.TExp Int64) + copyDWIMFix (tvVar word) [] (Var bit_array) [tvExp word_idx] + + -- Extract bit: (word >> bitIdx) & 1 + let extracted_bit = (tvExp word .>>. tvExp bit_idx) .&. 1 + bool_val = extracted_bit .==. 1 + + bool_var <- dPrimV "bool_val" bool_val + copyDWIMFix dest [] (Var $ tvVar bool_var) [] + +-- | Helper to determine if a type should use bit array representation +shouldUseBitArray :: Type -> Bool +shouldUseBitArray t = primType t && elemType t == Bool + +earlyReturn :: + Pat LetDecMem -> + SegBinOp GPUMem -> + SegPostOp GPUMem -> + ([[PatElem LetDecMem]], [Bool]) +earlyReturn pat scan_op post_op = (early_write_pats, res_flags) + where + post_lam = segPostOpLambda post_op + num_scan_res = length $ segBinOpNeutral scan_op + earlyWrite i par = + if num_scan_res < i && paramName par `notNameIn` free_in_post + then + filter ((== Just (paramName par)) . subExpResVName . fst) + . zip (bodyResult $ lambdaBody post_lam) + $ patElems pat + else [] + early_writes = zipWith earlyWrite [1 ..] $ lambdaParams post_lam + res_flags = (`elem` mconcat early_write_pats) <$> patElems pat + early_write_pats = drop num_scan_res $ map snd <$> early_writes + free_in_post = freeIn $ bodyStms $ lambdaBody post_lam + -- | Compile 'SegScan' instance to host-level code with calls to a -- single-pass kernel. compileSegScan :: Pat LetDecMem -> SegLevel -> SegSpace -> + [Type] -> SegBinOp GPUMem -> KernelBody GPUMem -> + SegPostOp GPUMem -> CallKernelGen () -compileSegScan pat lvl space scan_op map_kbody = do +compileSegScan pat lvl space ts scan_op map_kbody post_op = do attrs <- lvlKernelAttrs lvl - let Pat all_pes = pat - - scanop_nes = segBinOpNeutral scan_op + let scanop_nes = segBinOpNeutral scan_op n = product $ map pe64 $ segSpaceDims space - tys' = lambdaReturnType $ segBinOpLambda scan_op + scan_tys' = lambdaReturnType $ segBinOpLambda scan_op + map_tys' = drop (length $ segBinOpNeutral scan_op) ts - tys = map elemType tys' + scan_tys = map elemType scan_tys' tblock_size_e = pe64 $ unCount $ kAttrBlockSize attrs num_phys_blocks_e = pe64 $ unCount $ kAttrNumBlocks attrs + (early_write_pats, res_flags) = earlyReturn pat scan_op post_op - let chunk_const = getChunkSize tys' - chunk_v <- dPrimV "chunk_size" . isInt64 =<< kernelConstToExp chunk_const - let chunk = tvExp chunk_v + let chunk_const = getScanChunkSize scan_tys' + chunk_v <- dPrim "chunk_size" + let chunk_name = nameFromText $ prettyText $ tvVar chunk_v + addTuningParam chunk_name Nothing + emit . Imp.GetUserParam (tvVar chunk_v) chunk_name . isInt64 + =<< kernelConstToExp chunk_const + let chunk_constexp = LeafExp (Imp.SizeUserParam chunk_name chunk_const) int64 + chunk = tvExp chunk_v + let num_words_const = bitArrayWords chunk_const + num_words <- + dPrimV "num_bit_words" . isInt64 =<< kernelConstToExp num_words_const + num_virt_blocks <- tvSize <$> dPrimV "num_virt_blocks" (n `divUp` (tblock_size_e * chunk)) let num_virt_blocks_e = pe64 num_virt_blocks @@ -259,15 +381,39 @@ (aggregateArrays, incprefixArrays) <- fmap unzip $ - forM tys $ \ty -> + forM scan_tys $ \ty -> (,) <$> sAllocArray "aggregates" ty (Shape [num_virt_blocks]) (Space "device") <*> sAllocArray "incprefixes" ty (Shape [num_virt_blocks]) (Space "device") global_id <- genZeroes "global_dynid" 1 - let attrs' = attrs {kAttrConstExps = M.singleton (tvVar chunk_v) chunk_const} + let attrs' = + attrs + { kAttrConstExps = + M.fromList + [ (tvVar chunk_v, chunk_constexp), + (tvVar num_words, num_words_const) + ] + } + map_global_chunks <- + forM (zip map_tys' early_write_pats) $ \(t, pats) -> + if isAcc t || primType t || not (null pats) + then pure Nothing + else + Just + <$> sAllocArray + "global" + (elemType t) + ( Shape + [ unCount $ kAttrNumBlocks attrs, + unCount $ kAttrBlockSize attrs, + tvSize chunk_v + ] + <> arrayShape t + ) + (Space "device") sKernelThread "segscan" (segFlat space) attrs' $ do chunk32 <- dPrimVE "chunk_size_32b" $ sExt32 $ tvExp chunk_v @@ -277,7 +423,7 @@ ltid = sExt64 ltid32 (sharedId, transposedArrays, prefixArrays, warpscan, exchanges) <- - createLocalArrays (kAttrBlockSize attrs) (tvSize chunk_v) tys + createLocalArrays (kAttrBlockSize attrs) (tvSize chunk_v) scan_tys -- We wrap the entire kernel body in a virtualisation loop to -- handle the case where we do not have enough thread blocks to @@ -292,8 +438,7 @@ sOp $ Imp.GetBlockId (tvVar phys_block_id) 0 iters <- dPrimVE "virtloop_bound" $ - (num_virt_blocks_e - tvExp phys_block_id) - `divUp` num_phys_blocks_e + (num_virt_blocks_e - tvExp phys_block_id) `divUp` num_phys_blocks_e sFor "virtloop_i" iters $ const $ do dyn_id <- dPrim "dynamic_id" @@ -335,14 +480,35 @@ dPrimVE "segsize_compact" $ sExt32 $ sMin64 (chunk * tblock_size_e) segment_size - private_chunks <- - forM tys $ \ty -> + scan_private_chunks <- + forM scan_tys $ \ty -> sAllocArray "private" ty (Shape [tvSize chunk_v]) (ScalarSpace [tvSize chunk_v] ty) + map_private_chunks <- + forM (zip map_tys' early_write_pats) $ \(t, pats) -> + if isAcc t || not (primType t) || not (null pats) + then pure Nothing + else + if shouldUseBitArray t + then do + Just + <$> sAllocArray + "private_bits" + int64 + (Shape [tvSize num_words]) + (ScalarSpace [tvSize num_words] int64) + else + Just + <$> sAllocArray + "private" + (elemType t) + (Shape [tvSize chunk_v]) + (ScalarSpace [tvSize chunk_v] (elemType t)) + thd_offset <- dPrimVE "thd_offset" $ block_offset + ltid sComment "Load and map" $ @@ -356,23 +522,36 @@ let (all_scan_res, map_res) = splitAt (segBinOpResults [scan_op]) $ bodyResult map_kbody - -- Write map results to their global memory destinations - forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(dest, src) -> - copyDWIMFix (patElemName dest) (map Imp.le64 gtids) (kernelResultSubExp src) [] + -- Write map results to memory. + forM_ (zip5 map_private_chunks map_global_chunks (map kernelResultSubExp map_res) map_tys' early_write_pats) $ + \(priv_dest, glob_dest, src, ty, pats) -> do + case priv_dest of + Just d + | shouldUseBitArray ty -> + setBitInBitArray chunk d i $ isBool $ toExp' Bool src + Just d -> + copyDWIMFix d [i] src [] + Nothing -> pure () + maybe (pure ()) (\d -> copyDWIMFix d [tvExp phys_block_id, ltid, i] src []) glob_dest + + if null pats + then pure () + else forM_ pats $ \pe -> copyDWIMFix (patElemName pe) (map le64 gtids) src [] + -- Write to-scan results to private memory. - forM_ (zip private_chunks $ map kernelResultSubExp all_scan_res) $ \(dest, src) -> + forM_ (zip scan_private_chunks $ map kernelResultSubExp all_scan_res) $ \(dest, src) -> copyDWIMFix dest [i] src [] out_of_bounds = - forM_ (zip private_chunks scanop_nes) $ \(dest, ne) -> + forM_ (zip scan_private_chunks scanop_nes) $ \(dest, ne) -> copyDWIMFix dest [i] ne [] sIf (virt_tid .<. n) in_bounds out_of_bounds sOp $ Imp.ErrorSync Imp.FenceLocal sComment "Transpose scan inputs" $ do - forM_ (zip transposedArrays private_chunks) $ \(trans, priv) -> do + forM_ (zip transposedArrays scan_private_chunks) $ \(trans, priv) -> do sFor "i" chunk $ \i -> do sharedIdx <- dPrimVE "sharedIdx" $ ltid + i * tblock_size_e copyDWIMFix trans [sharedIdx] (Var priv) [i] @@ -397,18 +576,18 @@ else pure false -- skip scan of first element in segment sUnless new_sgm $ do - forM_ (zip4 private_chunks xs ys tys) $ \(src, x, y, ty) -> do + forM_ (zip4 scan_private_chunks xs ys scan_tys) $ \(src, x, y, ty) -> do dPrim_ x ty dPrim_ y ty copyDWIMFix x [] (Var src) [i] copyDWIMFix y [] (Var src) [i + 1] compileStms mempty (bodyStms $ lambdaBody $ segBinOpLambda scan_op) $ - forM_ (zip private_chunks $ map resSubExp $ bodyResult $ lambdaBody $ segBinOpLambda scan_op) $ \(dest, res) -> + forM_ (zip scan_private_chunks $ map resSubExp $ bodyResult $ lambdaBody $ segBinOpLambda scan_op) $ \(dest, res) -> copyDWIMFix dest [i + 1] res [] sComment "Publish results in shared memory" $ do - forM_ (zip prefixArrays private_chunks) $ \(dest, src) -> + forM_ (zip prefixArrays scan_private_chunks) $ \(dest, src) -> copyDWIMFix dest [ltid] (Var src) [chunk - 1] sOp local_barrier @@ -421,7 +600,7 @@ scan_op1 <- renameLambda $ segBinOpLambda scan_op - accs <- mapM (dPrimSV "acc") tys + accs <- mapM (dPrimSV "acc") scan_tys sComment "Scan results (with warp scan)" $ do blockScan crossesSegment @@ -443,7 +622,7 @@ sOp local_barrier - prefixes <- forM (zip scanop_nes tys) $ \(ne, ty) -> + prefixes <- forM (zip scanop_nes scan_tys) $ \(ne, ty) -> dPrimV "prefix" $ TPrimExp $ toExp' ty ne blockNewSgm <- dPrimVE "block_new_sgm" $ sgm_idx .==. 0 sComment "Perform lookback" $ do @@ -507,7 +686,7 @@ | otherwise = true sWhile (tvExp readOffset .>. loopStop) $ do readI <- dPrimV "read_i" $ tvExp readOffset + ltid32 - aggrs <- forM (zip scanop_nes tys) $ \(ne, ty) -> + aggrs <- forM (zip scanop_nes scan_tys) $ \(ne, ty) -> dPrimV "aggr" $ TPrimExp $ toExp' ty ne flag <- dPrimV "flag" (statusX :: Imp.TExp Int8) everythingVolatile . sWhen (tvExp readI .>=. 0) $ do @@ -559,11 +738,11 @@ -- update prefix if flag different than STATUS_X: sWhen (tvExp flag .>. statusX) $ do lam <- renameLambda scan_op1 - let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams lam + let (xs, ys) = splitAt (length scan_tys) $ map paramName $ lambdaParams lam forM_ (zip xs aggrs) $ \(x, aggr) -> dPrimV_ x (tvExp aggr) forM_ (zip ys prefixes) $ \(y, prefix) -> dPrimV_ y (tvExp prefix) compileStms mempty (bodyStms $ lambdaBody lam) $ - forM_ (zip3 prefixes tys $ map resSubExp $ bodyResult $ lambdaBody lam) $ + forM_ (zip3 prefixes scan_tys $ map resSubExp $ bodyResult $ lambdaBody lam) $ \(prefix, ty, res) -> prefix <-- TPrimExp (toExp' ty res) sOp local_fence ) @@ -572,8 +751,8 @@ -- end sIf sWhen (ltid32 .==. 0) $ do scan_op2 <- renameLambda scan_op1 - let xs = map paramName $ take (length tys) $ lambdaParams scan_op2 - ys = map paramName $ drop (length tys) $ lambdaParams scan_op2 + let xs = map paramName $ take (length scan_tys) $ lambdaParams scan_op2 + ys = map paramName $ drop (length scan_tys) $ lambdaParams scan_op2 sWhen (boundary .==. sExt32 (tblock_size_e * chunk)) $ do forM_ (zip xs prefixes) $ \(x, prefix) -> dPrimV_ x $ tvExp prefix forM_ (zip ys accs) $ \(y, acc) -> dPrimV_ y $ tvExp acc @@ -585,7 +764,7 @@ everythingVolatile $ copyDWIMFix statusFlags [tvExp dyn_id] (intConst Int8 statusP) [] forM_ (zip exchanges prefixes) $ \(exchange, prefix) -> copyDWIMFix exchange [0] (tvSize prefix) [] - forM_ (zip3 accs tys scanop_nes) $ \(acc, ty, ne) -> + forM_ (zip3 accs scan_tys scanop_nes) $ \(acc, ty, ne) -> tvVar acc <~~ toExp' ty ne -- end sWhen -- end sWhen @@ -602,10 +781,10 @@ scan_op4 <- renameLambda scan_op1 sComment "Distribute results" $ do - let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams scan_op3 - (xs', ys') = splitAt (length tys) $ map paramName $ lambdaParams scan_op4 + let (xs, ys) = splitAt (length scan_tys) $ map paramName $ lambdaParams scan_op3 + (xs', ys') = splitAt (length scan_tys) $ map paramName $ lambdaParams scan_op4 - forM_ (zip7 prefixes accs xs xs' ys ys' tys) $ + forM_ (zip7 prefixes accs xs xs' ys ys' scan_tys) $ \(prefix, acc, x, x', y, y', ty) -> do dPrim_ x ty dPrim_ y ty @@ -615,7 +794,7 @@ sIf (ltid32 * chunk32 .<. boundary .&&. bNot blockNewSgm) ( compileStms mempty (bodyStms $ lambdaBody scan_op4) $ - forM_ (zip3 xs tys $ map resSubExp $ bodyResult $ lambdaBody scan_op4) $ + forM_ (zip3 xs scan_tys $ map resSubExp $ bodyResult $ lambdaBody scan_op4) $ \(x, ty, res) -> x <~~ toExp' ty res ) (forM_ (zip xs accs) $ \(x, acc) -> copyDWIMFix x [] (Var $ tvVar acc) []) @@ -626,17 +805,16 @@ segsize_compact - (ltid32 * chunk32 - 1 + segsize_compact - boundary) `rem` segsize_compact sFor "i" chunk $ \i -> do sWhen (sExt32 i .<. stop - 1) $ do - forM_ (zip private_chunks ys) $ \(src, y) -> + forM_ (zip scan_private_chunks ys) $ \(src, y) -> -- only include prefix for the first segment part per thread copyDWIMFix y [] (Var src) [i] compileStms mempty (bodyStms $ lambdaBody scan_op3) $ - forM_ (zip private_chunks $ map resSubExp $ bodyResult $ lambdaBody scan_op3) $ + forM_ (zip scan_private_chunks $ map resSubExp $ bodyResult $ lambdaBody scan_op3) $ \(dest, res) -> copyDWIMFix dest [i] res [] - sComment "Transpose scan output and Write it to global memory in coalesced fashion" $ do - forM_ (zip3 transposedArrays private_chunks $ map patElemName all_pes) $ \(locmem, priv, dest) -> do - -- sOp local_barrier + sComment "Transpose scan output and to write it later in coalesced fashion to global memory" $ do + forM_ (zip transposedArrays scan_private_chunks) $ \(locmem, priv) -> do sFor "i" chunk $ \i -> do sharedIdx <- dPrimV "sharedIdx" $ @@ -648,9 +826,54 @@ dIndexSpace (zip gtids dims') flat_idx sWhen (flat_idx .<. n) $ do copyDWIMFix - dest - (map Imp.le64 gtids) + priv + [i] (Var locmem) [sExt64 $ flat_idx - block_offset] sOp local_barrier + + let (scan_pars, map_pars) = + splitAt (length $ segBinOpNeutral scan_op) $ + map paramName $ + lambdaParams $ + segPostOpLambda post_op + + sComment "Compute post op and write to global memory." $ do + sFor "i" chunk $ \i -> do + flat_idx <- dPrimVE "flat_idx" $ thd_offset + i * tblock_size_e + sWhen (flat_idx .<. n) $ do + dIndexSpace (zip gtids dims') flat_idx + + dScope Nothing $ + scopeOfLParams $ + lambdaParams $ + segPostOpLambda post_op + + sComment "bind scan results to post lambda params" $ do + forM_ (zip scan_pars scan_private_chunks) $ \(par, priv) -> + copyDWIMFix par [] (Var priv) [i] + + sComment "bind map results to post lamda params" $ + forM_ (zip4 map_pars map_private_chunks map_global_chunks map_tys') $ + \(par, priv, glob, ty) -> do + case priv of + Just p + | shouldUseBitArray ty -> + getBitFromBitArray chunk par p i + Just p -> + copyDWIMFix par [] (Var p) [i] + Nothing -> pure () + + maybe (pure ()) (\g -> copyDWIMFix par [] (Var g) [tvExp phys_block_id, ltid, i]) glob + + let res = fmap resSubExp $ bodyResult $ lambdaBody $ segPostOpLambda post_op + sComment "compute post op." $ + compileStms mempty (bodyStms $ lambdaBody $ segPostOpLambda post_op) $ + sComment "write mapped values" $ + forM_ (zip3 (patElems pat) res res_flags) $ \(pe, subexp, flag) -> + if flag + then pure () + else copyDWIMFix (patElemName pe) (map le64 gtids) subexp [] + + sOp local_barrier {-# NOINLINE compileSegScan #-} diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs @@ -14,7 +14,6 @@ import Futhark.IR.GPUMem import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.Transform.Rename -import Futhark.Util (takeLast) import Futhark.Util.IntegralExp (divUp, quot, rem) import Prelude hiding (quot, rem) @@ -91,13 +90,13 @@ writeToScanValues :: [VName] -> - ([PatElem LetDecMem], SegBinOp GPUMem, [KernelResult]) -> + ([VName], SegBinOp GPUMem, [KernelResult]) -> InKernelGen () writeToScanValues gtids (pes, scan, scan_res) | shapeRank (segBinOpShape scan) > 0 = forM_ (zip pes scan_res) $ \(pe, res) -> copyDWIMFix - (patElemName pe) + pe (map Imp.le64 gtids) (kernelResultSubExp res) [] @@ -107,13 +106,13 @@ readToScanValues :: [Imp.TExp Int64] -> - [PatElem LetDecMem] -> + [VName] -> SegBinOp GPUMem -> InKernelGen () readToScanValues is pes scan | shapeRank (segBinOpShape scan) > 0 = forM_ (zip (yParams scan) pes) $ \(p, pe) -> - copyDWIMFix (paramName p) [] (Var (patElemName pe)) is + copyDWIMFix (paramName p) [] (Var pe) is | otherwise = pure () @@ -122,7 +121,7 @@ Imp.TExp Int64 -> [Imp.TExp Int64] -> [Imp.TExp Int64] -> - [PatElem LetDecMem] -> + [VName] -> SegBinOp GPUMem -> InKernelGen () readCarries chunk_id chunk_offset dims' vec_is pes scan @@ -135,7 +134,7 @@ ( do let is = unflattenIndex dims' $ chunk_offset - 1 forM_ (zip (xParams scan) pes) $ \(p, pe) -> - copyDWIMFix (paramName p) [] (Var (patElemName pe)) (is ++ vec_is) + copyDWIMFix (paramName p) [] (Var pe) (is ++ vec_is) ) ( forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> copyDWIMFix (paramName p) [] ne [] @@ -144,15 +143,19 @@ pure () -- | Produce partially scanned intervals; one per threadblock. +-- +-- The @Maybe VName@ for the result is to cater to accumulators, which become +-- @Nothing@. scanStage1 :: - Pat LetDecMem -> + [VName] -> + [Maybe VName] -> Count NumBlocks SubExp -> Count BlockSize SubExp -> SegSpace -> [SegBinOp GPUMem] -> KernelBody GPUMem -> CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment) -scanStage1 (Pat all_pes) num_tblocks tblock_size space scans kbody = do +scanStage1 scan_out map_out num_tblocks tblock_size space scans kbody = do let num_tblocks' = fmap pe64 num_tblocks tblock_size' = fmap pe64 tblock_size num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_tblocks' * unCount tblock_size' @@ -191,28 +194,30 @@ -- Construct segment indices. zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx - let per_scan_pes = segBinOpChunks scans all_pes + let per_scan_pes = segBinOpChunks scans scan_out in_bounds = foldl1 (.&&.) $ zipWith (.<.) (map Imp.le64 gtids) dims' - when_in_bounds = compileStms mempty (bodyStms kbody) $ do - let (all_scan_res, map_res) = - splitAt (segBinOpResults scans) $ bodyResult kbody - per_scan_res = - segBinOpChunks scans all_scan_res + when_in_bounds = + compileStms mempty (bodyStms kbody) $ do + let (all_scan_res, map_res) = + splitAt (segBinOpResults scans) $ bodyResult kbody + per_scan_res = + segBinOpChunks scans all_scan_res - sComment "write to-scan values to parameters" $ - mapM_ (writeToScanValues gtids) $ - zip3 per_scan_pes scans per_scan_res + sComment "write to-scan values to parameters" $ + mapM_ (writeToScanValues gtids) $ + zip3 per_scan_pes scans per_scan_res - sComment "write mapped values results to global memory" $ - forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(pe, se) -> - copyDWIMFix - (patElemName pe) - (map Imp.le64 gtids) - (kernelResultSubExp se) - [] + sComment "write mapped values results to global memory" $ + forM_ (zip map_out map_res) $ \(pe, se) -> + forM pe $ \p -> + copyDWIMFix + p + (map Imp.le64 gtids) + (kernelResultSubExp se) + [] sComment "threads in bounds read input" $ sWhen in_bounds when_in_bounds @@ -270,7 +275,7 @@ sWhen in_bounds $ forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> copyDWIMFix - (patElemName pe) + pe (map Imp.le64 gtids ++ vec_is) (Var arr) [localArrayIndex constants t] @@ -319,7 +324,7 @@ pure (num_threads, elems_per_group, crossesSegment) scanStage2 :: - Pat LetDecMem -> + [VName] -> TV Int32 -> Imp.TExp Int64 -> Count NumBlocks SubExp -> @@ -327,7 +332,7 @@ SegSpace -> [SegBinOp GPUMem] -> CallKernelGen () -scanStage2 (Pat all_pes) stage1_num_threads elems_per_group num_tblocks crossesSegment space scans = do +scanStage2 scan_out stage1_num_threads elems_per_group num_tblocks crossesSegment space scans = do let (gtids, dims) = unzip $ unSegSpace space dims' = map pe64 dims @@ -345,7 +350,7 @@ constants <- kernelConstants <$> askEnv per_scan_local_arrs <- makeLocalArrays tblock_size (tvSize stage1_num_threads) scans let per_scan_rets = map (lambdaReturnType . segBinOpLambda) scans - per_scan_pes = segBinOpChunks scans all_pes + per_scan_pes = segBinOpChunks scans scan_out flat_idx <- dPrimV "flat_idx" $ @@ -365,7 +370,7 @@ copyDWIMFix arr [localArrayIndex constants t] - (Var $ patElemName pe) + (Var pe) glob_is when_out_of_bounds = forM_ (zip3 rets local_arrs nes) $ \(t, arr, ne) -> @@ -389,21 +394,24 @@ sWhen in_bounds $ forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> copyDWIMFix - (patElemName pe) + pe glob_is (Var arr) [localArrayIndex constants t] scanStage3 :: Pat LetDecMem -> + [VName] -> + [Maybe VName] -> Count NumBlocks SubExp -> Count BlockSize SubExp -> Imp.TExp Int64 -> CrossesSegment -> SegSpace -> [SegBinOp GPUMem] -> + SegPostOp GPUMem -> CallKernelGen () -scanStage3 (Pat all_pes) num_tblocks tblock_size elems_per_group crossesSegment space scans = do +scanStage3 pat scan_out map_out num_tblocks tblock_size elems_per_group crossesSegment space scans post_op = do let tblock_size' = fmap pe64 tblock_size (gtids, dims) = unzip $ unSegSpace space dims' = map pe64 dims @@ -446,8 +454,8 @@ is_a_carry = flat_idx .==. (tvExp orig_group + 1) * elems_per_group - 1 no_carry_in = tvExp orig_group .==. 0 .||. is_a_carry .||. crosses_segment - let per_scan_pes = segBinOpChunks scans all_pes - sWhen in_bounds $ + let per_scan_pes = segBinOpChunks scans scan_out + sWhen in_bounds $ do sUnless no_carry_in $ forM_ (zip per_scan_pes scans) $ \(pes, SegBinOp _ scan_op nes vec_shape) -> do @@ -460,35 +468,78 @@ copyDWIMFix (paramName p) [] - (Var $ patElemName pe) + (Var pe) (carry_in_idx ++ vec_is) forM_ (zip scan_y_params pes) $ \(p, pe) -> copyDWIMFix (paramName p) [] - (Var $ patElemName pe) + (Var pe) (map Imp.le64 gtids ++ vec_is) compileBody' scan_x_params $ lambdaBody scan_op forM_ (zip scan_x_params pes) $ \(p, pe) -> copyDWIMFix - (patElemName pe) + pe (map Imp.le64 gtids ++ vec_is) (Var $ paramName p) [] + sOp $ Imp.Barrier Imp.FenceLocal + + if isIdentityPostOp post_op + then pure () + else do + sWhen in_bounds $ do + let (scan_pars, map_pars) = splitAt (segBinOpResults scans) $ lambdaParams $ segPostOpLambda post_op + dScope Nothing $ + scopeOfLParams $ + lambdaParams $ + segPostOpLambda post_op + + sComment "bind scan results to post lambda params" $ do + forM_ (zip scan_pars scan_out) $ \(par, acc) -> + copyDWIMFix (paramName par) [] (Var acc) (map Imp.le64 gtids) + + sComment "bind map results to post lamda params" $ + forM_ (zip map_pars map_out) $ \(par, out) -> do + maybe + (pure ()) + ( \o -> + copyDWIMFix (paramName par) [] (Var o) (map Imp.le64 gtids) + ) + out + + let res = fmap resSubExp $ bodyResult $ lambdaBody $ segPostOpLambda post_op + sComment "compute post op." $ + compileStms mempty (bodyStms $ lambdaBody $ segPostOpLambda post_op) $ + sComment "write values" $ + forM_ (zip (patElems pat) res) $ \(pe, subexp) -> + copyDWIMFix (patElemName pe) (map Imp.le64 gtids) subexp [] + + sOp $ Imp.Barrier Imp.FenceGlobal + +isIdentityPostOp :: SegPostOp rep -> Bool +isIdentityPostOp post_op = + map resSubExp (bodyResult (lambdaBody lam)) + == map (Var . paramName) (lambdaParams lam) + where + lam = segPostOpLambda post_op + -- | Compile 'SegScan' instance to host-level code with calls to -- various kernels. compileSegScan :: Pat LetDecMem -> SegLevel -> SegSpace -> + [Type] -> [SegBinOp GPUMem] -> KernelBody GPUMem -> + SegPostOp GPUMem -> CallKernelGen () -compileSegScan pat lvl space scans kbody = do +compileSegScan pat lvl space ts scans kbody post_op = do attrs <- lvlKernelAttrs lvl -- Since stage 2 involves a group size equal to the number of groups @@ -504,10 +555,40 @@ pe64 . Imp.unCount . kAttrNumBlocks $ attrs + let shpT op = (segBinOpShape op,) <$> lambdaReturnType (segBinOpLambda op) + scan_ts = concatMap shpT scans + shpOfT t s = + arrayShape $ + foldr (flip arrayOfRow) (arrayOfShape t s) $ + segSpaceDims space + + (scan_out, map_out) <- + if isIdentityPostOp post_op + then + let (scan_out, map_out) = splitAt (segBinOpResults scans) $ patElemName <$> patElems pat + map_out' = + zipWith + ( \t n -> + if isAcc t then Nothing else Just n + ) + (drop (segBinOpResults scans) ts) + map_out + in pure (scan_out, map_out') + else do + scan_out <- forM scan_ts $ \(s, t) -> + sAllocArray "scan_out" (elemType t) (shpOfT t s) (Space "device") + + map_out <- forM (drop (segBinOpResults scans) ts) $ \t -> + if isAcc t + then pure Nothing + else Just <$> sAllocArray "map_out" (elemType t) (shpOfT t mempty) (Space "device") + + pure (scan_out, map_out) + (stage1_num_threads, elems_per_group, crossesSegment) <- - scanStage1 pat stage1_num_tblocks (kAttrBlockSize attrs) space scans kbody + scanStage1 scan_out map_out stage1_num_tblocks (kAttrBlockSize attrs) space scans kbody emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group - scanStage2 pat stage1_num_threads elems_per_group stage1_num_tblocks crossesSegment space scans - scanStage3 pat (kAttrNumBlocks attrs) (kAttrBlockSize attrs) elems_per_group crossesSegment space scans + scanStage2 scan_out stage1_num_threads elems_per_group stage1_num_tblocks crossesSegment space scans + scanStage3 pat scan_out map_out (kAttrNumBlocks attrs) (kAttrBlockSize attrs) elems_per_group crossesSegment space scans post_op diff --git a/src/Futhark/CodeGen/ImpGen/Multicore.hs b/src/Futhark/CodeGen/ImpGen/Multicore.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore.hs @@ -169,8 +169,8 @@ ImpM MCMem HostEnv Imp.Multicore () compileSegOp pat (SegHist _ space _ kbody histops) ntasks = compileSegHist pat space histops kbody ntasks -compileSegOp pat (SegScan _ space _ kbody scans) ntasks = - compileSegScan pat space scans kbody ntasks +compileSegOp pat (SegScan _ space ts kbody scans post_op) ntasks = + compileSegScan pat space ts kbody scans post_op ntasks compileSegOp pat (SegRed _ space _ kbody reds) ntasks = compileSegRed pat space reds kbody ntasks compileSegOp pat (SegMap _ space _ kbody) _ = diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs @@ -81,7 +81,7 @@ getSpace :: SegOp () MCMem -> SegSpace getSpace (SegHist _ space _ _ _) = space getSpace (SegRed _ space _ _ _) = space -getSpace (SegScan _ space _ _ _) = space +getSpace (SegScan _ space _ _ _ _) = space getSpace (SegMap _ space _ _) = space getLoopBounds :: MulticoreGen (Imp.TExp Int64, Imp.TExp Int64) @@ -492,6 +492,7 @@ _ -> (id, id) int + | primBitSize t == 8 = int8 | primBitSize t == 16 = int16 | primBitSize t == 32 = int32 | otherwise = int64 diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs @@ -165,7 +165,8 @@ copyDWIMFix prefix_arr (tvExp block_idx : vec_is) se [] seqScanFastPath :: - Pat LetDecMem -> + [[VName]] -> + [Maybe VName] -> VName -> [SegBinOp MCMem] -> KernelBody MCMem -> @@ -174,8 +175,9 @@ TV Int64 -> [[VName]] -> TV Int64 -> + TV Int64 -> MulticoreGen () -seqScanFastPath pat i scan_ops kbody per_op_prefixes_var start chunk_length per_op_prefix_arr block_idx = do +seqScanFastPath scan_out map_out i scan_ops kbody per_op_prefixes_var start chunk_length per_op_prefix_arr block_idx task_id = do kbody_renamed <- renameBody kbody scan_ops_renamed <- renameSegBinOp scan_ops genBinOpParams scan_ops_renamed @@ -186,7 +188,6 @@ let (all_scan_res, map_res) = splitAt n_scan results let per_scan_res = segBinOpChunks scan_ops_renamed all_scan_res - let per_scan_pes = segBinOpChunks scan_ops_renamed $ patElems pat forM_ (zip3 scan_ops_renamed per_op_prefixes_var per_op_local_accum) $ \(scan_op, prefix_vars, local_accums) -> sLoopNest (segBinOpShape scan_op) $ \vec_is -> do @@ -196,15 +197,12 @@ z <- dPrimV "z" (0 :: Imp.TExp Int64) sWhile (tvExp z .<. tvExp chunk_length) $ do dPrimV_ i (tvExp start + tvExp z) - compileStms mempty (bodyStms kbody_renamed) $ do - let map_arrs = drop (segBinOpResults scan_ops_renamed) $ patElems pat - sComment "write mapped values results to memory" $ - forM_ (zip (map patElemName map_arrs) (map kernelResultSubExp map_res)) $ \(arr, res) -> - copyDWIMFix arr [tvExp start + tvExp z] res [] + forM_ (zip map_out (map kernelResultSubExp map_res)) $ \(marr, res) -> + forM_ marr $ \arr -> copyDWIMFix arr [tvExp task_id, tvExp z] res [] - forM_ (zip4 per_scan_pes scan_ops_renamed per_scan_res per_op_local_accum) $ \(pes, scan_op, scan_res, local_accums) -> + forM_ (zip4 scan_out scan_ops_renamed per_scan_res per_op_local_accum) $ \(pes, scan_op, scan_res, local_accums) -> sLoopNest (segBinOpShape scan_op) $ \vec_is -> do forM_ (zip (xParams scan_op) local_accums) $ \(p, acc) -> copyDWIMFix (paramName p) [] (Var acc) vec_is @@ -212,7 +210,7 @@ copyDWIMFix (paramName py) [] (kernelResultSubExp kr) vec_is compileStms mempty (bodyStms $ lamBody scan_op) $ forM_ (zip3 local_accums (map resSubExp $ bodyResult $ lamBody scan_op) pes) $ \(acc, se, pe) -> do - copyDWIMFix (patElemName pe) ((tvExp start + tvExp z) : vec_is) se [] + copyDWIMFix pe (tvExp task_id : tvExp z : vec_is) se [] copyDWIMFix acc vec_is se [] z <-- tvExp z + 1 @@ -223,15 +221,16 @@ copyDWIMFix prefix_arr (tvExp block_idx : vec_is) (Var acc) vec_is seqScanLB :: - Pat LetDecMem -> + [[VName]] -> VName -> [SegBinOp MCMem] -> KernelBody MCMem -> [[VName]] -> TV Int64 -> TV Int64 -> + TV Int64 -> MulticoreGen () -seqScanLB pat i scan_ops kbody per_op_prefixes_var start chunk_length = do +seqScanLB scan_out i scan_ops kbody per_op_prefixes_var start chunk_length task_id = do kbody_renamed <- renameBody kbody scan_ops_renamed <- renameSegBinOp scan_ops genBinOpParams scan_ops_renamed @@ -242,7 +241,6 @@ let (all_scan_res, _) = splitAt n_scan results let per_scan_res = segBinOpChunks scan_ops_renamed all_scan_res - let per_scan_pes = segBinOpChunks scan_ops_renamed $ patElems pat forM_ (zip3 scan_ops_renamed per_op_prefixes_var per_op_local_accum) $ \(scan_op, prefix_vars, local_accums) -> sLoopNest (segBinOpShape scan_op) $ \vec_is -> @@ -255,7 +253,7 @@ if shouldRecompute kbody_renamed then do compileStms mempty (bodyStms kbody_renamed) $ do - forM_ (zip4 per_scan_pes scan_ops_renamed per_scan_res per_op_local_accum) $ \(pes, scan_op, scan_res, local_accums) -> + forM_ (zip4 scan_out scan_ops_renamed per_scan_res per_op_local_accum) $ \(pes, scan_op, scan_res, local_accums) -> sLoopNest (segBinOpShape scan_op) $ \vec_is -> do forM_ (zip (xParams scan_op) local_accums) $ \(px, acc) -> copyDWIMFix (paramName px) [] (Var acc) vec_is @@ -264,23 +262,23 @@ compileStms mempty (bodyStms $ lamBody scan_op) $ forM_ (zip3 (map resSubExp $ bodyResult $ lamBody scan_op) pes local_accums) $ \(se, pe, acc) -> do copyDWIMFix acc vec_is se [] - copyDWIMFix (patElemName pe) ((tvExp start + tvExp z) : vec_is) se [] + copyDWIMFix pe (tvExp task_id : tvExp z : vec_is) se [] z <-- tvExp z + 1 else do - forM_ (zip4 per_scan_pes scan_ops_renamed per_scan_res per_op_local_accum) $ \(pes, scan_op, _, local_accums) -> + forM_ (zip4 scan_out scan_ops_renamed per_scan_res per_op_local_accum) $ \(pes, scan_op, _, local_accums) -> sLoopNest (segBinOpShape scan_op) $ \vec_is -> do forM_ (zip (xParams scan_op) local_accums) $ \(px, acc) -> copyDWIMFix (paramName px) [] (Var acc) vec_is forM_ (zip (yParams scan_op) pes) $ \(py, pe) -> - -- reading from output array - copyDWIMFix (paramName py) [] (Var (patElemName pe)) ((tvExp start + tvExp z) : vec_is) + copyDWIMFix (paramName py) [] (Var pe) (tvExp task_id : tvExp z : vec_is) compileStms mempty (bodyStms $ lamBody scan_op) $ do forM_ (zip (map resSubExp $ bodyResult $ lamBody scan_op) pes) $ \(se, pe) -> - copyDWIMFix (patElemName pe) ((tvExp start + tvExp z) : vec_is) se [] + copyDWIMFix pe (tvExp task_id : tvExp z : vec_is) se [] z <-- tvExp z + 1 seqAggregate :: - Pat LetDecMem -> + [[VName]] -> + [Maybe VName] -> VName -> [SegBinOp MCMem] -> KernelBody MCMem -> @@ -288,8 +286,9 @@ TV Int64 -> [[VName]] -> TV Int64 -> + TV Int64 -> MulticoreGen () -seqAggregate pat i scan_ops kbody start chunk_length per_op_aggr_arrs block_idx = do +seqAggregate scan_out map_out i scan_ops kbody start chunk_length per_op_aggr_arrs block_idx task_id = do scan_ops_renamed <- renameSegBinOp scan_ops kbody_renamed <- renameBody kbody genBinOpParams scan_ops_renamed @@ -297,7 +296,6 @@ let n_scan = segBinOpResults scan_ops_renamed let (all_scan_res, map_res) = splitAt n_scan results let per_scan_res = segBinOpChunks scan_ops_renamed all_scan_res - let per_scan_pes = segBinOpChunks scan_ops_renamed $ patElems pat per_op_local_accum <- genLocalArray scan_ops_renamed @@ -305,23 +303,22 @@ sWhile (tvExp j .<. tvExp chunk_length) $ do dPrimV_ i (tvExp start + tvExp j) compileStms mempty (bodyStms kbody_renamed) $ do - let map_arrs = drop (segBinOpResults scan_ops_renamed) $ patElems pat sComment "write mapped values results to memory" $ - forM_ (zip (map patElemName map_arrs) (map kernelResultSubExp map_res)) $ \(arr, res) -> - copyDWIMFix arr [tvExp start + tvExp j] res [] + forM_ (zip map_out (map kernelResultSubExp map_res)) $ \(marr, res) -> + forM_ marr $ \arr -> copyDWIMFix arr [tvExp task_id, tvExp j] res [] sIf (tvExp j .==. 0) - ( forM_ (zip4 scan_ops_renamed per_scan_res per_op_local_accum per_scan_pes) $ + ( forM_ (zip4 scan_ops_renamed per_scan_res per_op_local_accum scan_out) $ \(scan_op, scan_res_op, local_accums, pes) -> do let shape = segBinOpShape scan_op sLoopNest shape $ \vec_is -> do forM_ (zip3 scan_res_op local_accums pes) $ \(kr, acc, pe) -> do copyDWIMFix acc vec_is (kernelResultSubExp kr) vec_is unless (shouldRecompute kbody_renamed) $ - copyDWIMFix (patElemName pe) ((tvExp start + tvExp j) : vec_is) (kernelResultSubExp kr) vec_is + copyDWIMFix pe (tvExp task_id : tvExp j : vec_is) (kernelResultSubExp kr) vec_is ) - ( forM_ (zip4 scan_ops_renamed per_scan_res per_op_local_accum per_scan_pes) $ + ( forM_ (zip4 scan_ops_renamed per_scan_res per_op_local_accum scan_out) $ \(scan_op, scan_res, local_accums, pes) -> sLoopNest (segBinOpShape scan_op) $ \vec_is -> do forM_ (zip (xParams scan_op) local_accums) $ \(px, acc) -> @@ -332,7 +329,7 @@ compileStms mempty (bodyStms $ lamBody scan_op) $ do forM_ (zip3 (map resSubExp $ bodyResult $ lamBody scan_op) local_accums pes) $ \(se, acc, pe) -> do copyDWIMFix acc vec_is se [] - unless (shouldRecompute kbody_renamed) $ copyDWIMFix (patElemName pe) ((tvExp start + tvExp j) : vec_is) se [] + unless (shouldRecompute kbody_renamed) $ copyDWIMFix pe (tvExp task_id : tvExp j : vec_is) se [] ) j <-- tvExp j + 1 @@ -364,19 +361,66 @@ MulticoreGen () add64 v arr i x = sOp $ Imp.Atomic $ Imp.AtomicAdd Int64 (tvVar v) arr i (untyped x) +applyPostOp :: + Pat LetDecMem -> + [[VName]] -> + [Maybe VName] -> + [SegBinOp MCMem] -> + SegPostOp MCMem -> + TV Int64 -> + TV Int64 -> + TV Int64 -> + MulticoreGen () +applyPostOp pat scan_out map_out scan_ops post_op start chunk_length task_id = do + z <- dPrimV "z" (0 :: Imp.TExp Int64) + let (scan_pars, map_pars) = splitAt (segBinOpResults scan_ops) $ lambdaParams $ segPostOpLambda post_op + dScope Nothing $ + scopeOfLParams $ + lambdaParams $ + segPostOpLambda post_op + + sWhile (tvExp z .<. tvExp chunk_length) $ do + sComment "bind scan results to post lambda params" $ do + forM_ (zip scan_pars $ mconcat scan_out) $ \(par, acc) -> + copyDWIMFix (paramName par) [] (Var acc) [tvExp task_id, tvExp z] + + sComment "bind map results to post lamda params" $ + forM_ (zip map_pars map_out) $ \(par, out) -> do + forM_ out $ \o -> + copyDWIMFix (paramName par) [] (Var o) [tvExp task_id, tvExp z] + + let res = fmap resSubExp $ bodyResult $ lambdaBody $ segPostOpLambda post_op + sComment "compute post op." $ + compileStms mempty (bodyStms $ lambdaBody $ segPostOpLambda post_op) $ + sComment "write values" $ + forM_ (zip (patElems pat) res) $ \(pe, subexp) -> + copyDWIMFix (patElemName pe) [tvExp start + tvExp z] subexp [] + + z <-- tvExp z + 1 + +getTaskId :: MulticoreGen (TV Int64) +getTaskId = do + task_id <- dPrim "task_id" + sOp $ Imp.GetTaskId (tvVar task_id) + pure task_id + nonsegmentedScan :: Pat LetDecMem -> SegSpace -> + [Type] -> [SegBinOp MCMem] -> KernelBody MCMem -> + SegPostOp MCMem -> TV Int32 -> MulticoreGen () nonsegmentedScan pat (SegSpace fid [(i, n)]) + ts scan_ops kbody - _nsubtasks = do + post_op + nsubtasks = do let multiplier = 1 -- For playing with. blockSize = cacheSize `divUp` (totalBytes scan_ops * multiplier) @@ -385,12 +429,24 @@ -- allocate flags/aggr/prefix arrays of length nblocks flagsArr <- sAllocArray "scan_flags" int64 (Shape [Var (tvVar block_no)]) DefaultSpace + nsubtasks_i64 <- dPrimV "nsubtasks_i64" $ sExt64 (tvExp nsubtasks) + + blockSize_var <- dPrimV "block_size" blockSize + let block_shape = Shape [Var (tvVar block_no)] + nsubtasks_shape = Shape [Var (tvVar nsubtasks_i64), Var (tvVar blockSize_var)] aggrArrs <- genArrays scan_ops "scan_aggr" block_shape prefArrs <- genArrays scan_ops "scan_pref" block_shape + scan_out <- genArrays scan_ops "scan_out" nsubtasks_shape + + map_out <- forM (drop (segBinOpResults scan_ops) ts) $ \t -> + if isAcc t + then pure Nothing + else Just <$> sAllocArray "map_out" (elemType t) (nsubtasks_shape <> arrayShape t) DefaultSpace + work_index <- sAllocArray "work_index" int64 (Shape [intConst Int64 1]) DefaultSpace sFor "init" (tvExp block_no) $ \j -> do @@ -426,6 +482,8 @@ let flag_loc_name = memLocName flags_loc let block_idx_32 = sExt32 (tvExp block_idx) + task_id <- getTaskId + sWhen (tvExp seq_flag .==. true) ( sIf @@ -444,12 +502,12 @@ sIf (tvExp seq_flag .==. true) ( do - seqScanFastPath pat i scan_ops kbody prefix_seqs start chunk_length prefArrs block_idx + seqScanFastPath scan_out map_out i scan_ops kbody prefix_seqs start chunk_length prefArrs block_idx task_id store64 flag_loc_name (Imp.elements block_idx_32) 2 ) ( do - seqAggregate pat i scan_ops kbody start chunk_length aggrArrs block_idx + seqAggregate scan_out map_out i scan_ops kbody start chunk_length aggrArrs block_idx task_id -- write flag as 1 store64 flag_loc_name (Imp.elements block_idx_32) 1 @@ -497,9 +555,11 @@ store64 flag_loc_name (Imp.elements block_idx_32) 2 - seqScanLB pat i scan_ops kbody prefix_vars start chunk_length + seqScanLB scan_out i scan_ops kbody prefix_vars start chunk_length task_id ) + applyPostOp pat scan_out map_out scan_ops post_op start chunk_length task_id + add64 block_idx work_index_loc_name (Imp.elements 0) 1 free_params <- freeParams fbody @@ -509,12 +569,14 @@ compileSegScan :: Pat LetDecMem -> SegSpace -> - [SegBinOp MCMem] -> + [Type] -> KernelBody MCMem -> + [SegBinOp MCMem] -> + SegPostOp MCMem -> TV Int32 -> MulticoreGen () -compileSegScan pat space reds kbody nsubtasks +compileSegScan pat space ts kbody reds post_op nsubtasks | [_] <- unSegSpace space = - nonsegmentedScan pat space reds kbody nsubtasks + nonsegmentedScan pat space ts reds kbody post_op nsubtasks | otherwise = error "only nonsegmented scans for now" diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs --- a/src/Futhark/Doc/Generator.hs +++ b/src/Futhark/Doc/Generator.hs @@ -378,7 +378,7 @@ pure $ specRow lhs (mhs <> " : ") rhs valBindHtml :: Html -> ValBind -> DocM (Html, Html, Html) -valBindHtml name (ValBind _ _ retdecl (Info rettype) tparams params _ _ _ _) = do +valBindHtml name (ValBind _ _ _ retdecl (Info rettype) tparams params _ _ _ _) = do tparams' <- mconcat <$> mapM (fmap (" " <>) . typeParamHtml) tparams let noLink' = noLink $ map typeParamName tparams <> foldMap patNames params rettype' <- noLink' $ maybe (retTypeHtml rettype) typeExpHtml retdecl diff --git a/src/Futhark/Fmt/Printer.hs b/src/Futhark/Fmt/Printer.hs --- a/src/Futhark/Fmt/Printer.hs +++ b/src/Futhark/Fmt/Printer.hs @@ -474,7 +474,7 @@ leading = leadingOperator $ toName $ qualLeaf bop instance Format UncheckedValBind where - fmt (ValBind entry name retdecl _rettype tparams args body docs attrs loc) = + fmt (ValBind entry name _ retdecl _rettype tparams args body docs attrs loc) = addComments loc $ fmt docs <> attrs' @@ -558,7 +558,12 @@ ps_op = if null ps then (<>) else (<+>) fmt (ModTypeArrow (Just v) te0 te1 loc) = addComments loc $ - parens (fmtName bindingStyle v <> ":" <+> fmt te0) <+> align ("->" fmt te1) + parens (fmtName bindingStyle v <> ":" <+> fmt te0) + "->" + <+> case te1 of + ModTypeArrow {} -> fmt te1 + ModTypeSpecs {} -> fmt te1 + _ -> align (fmt te1) fmt (ModTypeArrow Nothing te0 te1 loc) = addComments loc $ fmt te0 <+> "->" <+> fmt te1 diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs --- a/src/Futhark/IR/Parse.hs +++ b/src/Futhark/IR/Parse.hs @@ -16,6 +16,7 @@ parseVName, parseSubExp, parseSubExpRes, + parseIdent, -- * Representation-specific fragments parseLambdaSOACS, @@ -745,6 +746,9 @@ pStateAndProg :: PR rep -> Parser (VNameSource, Prog rep) pStateAndProg pr = (,) <$> (pVNameSource <|> pure (VNameSource 0)) <*> pProg pr +pIdent :: Parser Ident +pIdent = Ident <$> pVName <* pColon <*> pType + pSOAC :: PR rep -> Parser (SOAC.SOAC rep) pSOAC pr = choice @@ -773,20 +777,33 @@ <*> braces (pScan pr `sepBy` pComma) <* pComma <*> braces (pReduce pr `sepBy` pComma) + <* pComma + <*> pLambda pr pRedomapForm = SOAC.ScremaForm <$> pLambda pr <*> pure [] <* pComma <*> braces (pReduce pr `sepBy` pComma) + <* pComma + -- NOTE: This is dumb, but it also seems weird to have + -- multiple waus of parsing a screma? but it is human readable. + <*> pLambda pr pScanomapForm = SOAC.ScremaForm <$> pLambda pr <* pComma <*> braces (pScan pr `sepBy` pComma) <*> pure [] + <* pComma + <*> pLambda pr pMapForm = - SOAC.ScremaForm <$> pLambda pr <*> pure mempty <*> pure mempty + SOAC.ScremaForm + <$> pLambda pr + <*> pure mempty + <*> pure mempty + <* pComma + <*> pLambda pr pHist = keyword "hist" *> parens @@ -951,6 +968,9 @@ comm <- pComm lam <- pLambda pr pure $ SegOp.SegBinOp comm lam nes shape + pSegPostOp = + SegOp.SegPostOp + <$> pLambda pr pHistOp = SegOp.HistOp <$> pShape @@ -966,7 +986,7 @@ <*> pLambda pr pSegMap = pSegOp' SegOp.SegMap pSegRed = pSegOp' SegOp.SegRed <*> parens (pSegBinOp `sepBy` pComma) - pSegScan = pSegOp' SegOp.SegScan <*> parens (pSegBinOp `sepBy` pComma) + pSegScan = pSegOp' SegOp.SegScan <*> parens (pSegBinOp `sepBy` pComma) <*> pSegPostOp pSegHist = pSegOp' SegOp.SegHist <*> parens (pHistOp `sepBy` pComma) pSegLevel :: Parser GPU.SegLevel @@ -1202,6 +1222,9 @@ parseSubExpRes :: FilePath -> T.Text -> Either T.Text SubExpRes parseSubExpRes = parseFull pSubExpRes + +parseIdent :: FilePath -> T.Text -> Either T.Text Ident +parseIdent = parseFull pIdent -- Rep-specific fragment parsers diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs --- a/src/Futhark/IR/SOACS.hs +++ b/src/Futhark/IR/SOACS.hs @@ -55,10 +55,11 @@ expUsesAD (Op JVP {}) = True expUsesAD (Op VJP {}) = True expUsesAD (Op (Stream _ _ _ lam)) = lamUsesAD lam - expUsesAD (Op (Screma _ _ (ScremaForm lam scans reds))) = + expUsesAD (Op (Screma _ _ (ScremaForm lam scans reds post_lam))) = lamUsesAD lam || any (lamUsesAD . scanLambda) scans || any (lamUsesAD . redLambda) reds + || lamUsesAD post_lam expUsesAD (Op (Hist _ _ ops lam)) = lamUsesAD lam || any (lamUsesAD . histOp) ops expUsesAD (Match _ cases def_case _) = 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 @@ -15,17 +15,21 @@ singleReduce, -- * Utility + composeBinds, scremaType, soacType, typeCheckSOAC, mkIdentityLambda, isIdentityLambda, + isNilLambda, nilFn, + maposcanomapSOAC, scanomapSOAC, redomapSOAC, scanSOAC, reduceSOAC, mapSOAC, + isMaposcanomapSOAC, isScanomapSOAC, isRedomapSOAC, isScanSOAC, @@ -107,7 +111,8 @@ -- the input arrays, however. scremaLambda :: Lambda rep, scremaScans :: [Scan rep], - scremaReduces :: [Reduce rep] + scremaReduces :: [Reduce rep], + scremaPostLambda :: Lambda rep } deriving (Eq, Ord, Show) @@ -173,15 +178,61 @@ -- | The types produced by a single 'Screma', given the size of the -- input array. scremaType :: SubExp -> ScremaForm rep -> [Type] -scremaType w (ScremaForm map_lam scans reds) = - scan_tps ++ red_tps ++ map (`arrayOfRow` w) map_tps +scremaType w (ScremaForm _map_lam _scans reds post_lam) = + red_tps <> fmap (`arrayOfRow` w) (lambdaReturnType post_lam) where - scan_tps = - map (`arrayOfRow` w) $ - concatMap (lambdaReturnType . scanLambda) scans red_tps = concatMap (lambdaReturnType . redLambda) reds - map_tps = drop (length scan_tps + length red_tps) $ lambdaReturnType map_lam +-- | Creates let-bindings to compose two lambda functions (producer → +-- consumer). +-- +-- When composing two operations where the outputs of one lambda +-- (producer) flow into the inputs of another lambda (consumer), this +-- function generates the necessary let-bindings that connect matching +-- outputs to inputs. +-- +-- The function looks for outputs from the producer that correspond to +-- inputs expected by the consumer, and creates bindings like: let +-- consumer_param = producer_result +-- +-- This allows the two lambdas to be composed into a operation. +-- +-- Returns: Statements containing let-bindings for each matched +-- producer output → consumer input pair. Preserves certificates from +-- the producer results. +-- +-- Example: If out_p[i] == inp_c[j], then producer's result[i] is +-- bound to consumer's parameter[j]. Unmatched outputs (producer +-- results not consumed) are omitted. +composeBinds :: + (Buildable rep, Ord a) => + -- | Producer lambda + Lambda rep -> + -- | Producer outputs to match + [a] -> + -- | Consumer inputs to match + [a] -> + -- | Consumer lambda + Lambda rep -> + -- | Let-bindings connecting them + Stms rep +composeBinds lam_p out_p inp_c lam_c = + stmsFromList . mapMaybe bindResToPar $ zip3 out_p res_p ts_p + where + ts_p = lambdaReturnType lam_p + res_p = bodyResult $ lambdaBody lam_p + + inp_c_map = + M.fromList . zip inp_c $ paramName <$> lambdaParams lam_c + + bindResToPar (out, res, t) = + case M.lookup out inp_c_map of + Just name -> + Just $ certify cs $ mkLet [Ident name t] $ BasicOp $ SubExp e + where + SubExpRes cs e = res + Nothing -> Nothing + -- | Construct a lambda that takes parameters of the given types and -- simply returns them unchanged. mkIdentityLambda :: @@ -203,19 +254,64 @@ map resSubExp (bodyResult (lambdaBody lam)) == map (Var . paramName) (lambdaParams lam) +-- | Is the given lambda a nil lambda? +isNilLambda :: Lambda rep -> Bool +isNilLambda lam = + null (lambdaParams lam) + && isIdentityLambda lam + -- | A lambda with no parameters that returns no values. nilFn :: (Buildable rep) => Lambda rep nilFn = Lambda mempty mempty (mkBody mempty mempty) --- | Construct a Screma with possibly multiple scans, and --- the given map function. -scanomapSOAC :: [Scan rep] -> Lambda rep -> ScremaForm rep -scanomapSOAC scans lam = ScremaForm lam scans [] +-- | Construct a Screma with possibly multiple scans, and the given +-- map function. +scanomapSOAC :: + (Buildable rep, MonadFreshNames m) => + [Scan rep] -> + Lambda rep -> + m (ScremaForm rep) +scanomapSOAC scans lam = + ScremaForm lam scans [] <$> mkIdentityLambda (lambdaReturnType lam) +-- | Construct a Screma with possibly multiple scans, +-- the given map function, and a given post lambda. +maposcanomapSOAC :: + (Buildable rep, MonadFreshNames m) => + Lambda rep -> + [Scan rep] -> + Lambda rep -> + m (ScremaForm rep) +maposcanomapSOAC pre_lam [] post_lam = do + new_post_lam <- mkIdentityLambda $ lambdaReturnType post_lam + let new_pre_lam = + Lambda + { lambdaParams = lambdaParams pre_lam, + lambdaReturnType = lambdaReturnType post_lam, + lambdaBody = mkBody new_stms new_res + } + pure $ ScremaForm new_pre_lam [] [] new_post_lam + where + new_res = bodyResult $ lambdaBody post_lam + stmsFromLam = bodyStms . lambdaBody + deps = [0 .. length $ lambdaReturnType pre_lam] + new_stms = + stmsFromLam pre_lam + <> composeBinds pre_lam deps deps post_lam + <> stmsFromLam post_lam +maposcanomapSOAC lam scans post_lam = + pure $ ScremaForm lam scans [] post_lam + -- | Construct a Screma with possibly multiple reductions, and -- the given map function. -redomapSOAC :: [Reduce rep] -> Lambda rep -> ScremaForm rep -redomapSOAC reds lam = ScremaForm lam [] reds +redomapSOAC :: + (Buildable rep, MonadFreshNames m) => + [Reduce rep] -> + Lambda rep -> + m (ScremaForm rep) +redomapSOAC reds lam = ScremaForm lam [] reds <$> mkIdentityLambda map_ts + where + map_ts = drop (redResults reds) $ lambdaReturnType lam -- | Construct a Screma with possibly multiple scans, and identity map -- function. @@ -223,7 +319,7 @@ (Buildable rep, MonadFreshNames m) => [Scan rep] -> m (ScremaForm rep) -scanSOAC scans = scanomapSOAC scans <$> mkIdentityLambda ts +scanSOAC scans = scanomapSOAC scans =<< mkIdentityLambda ts where ts = concatMap (lambdaReturnType . scanLambda) scans @@ -233,33 +329,47 @@ (Buildable rep, MonadFreshNames m) => [Reduce rep] -> m (ScremaForm rep) -reduceSOAC reds = redomapSOAC reds <$> mkIdentityLambda ts +reduceSOAC reds = redomapSOAC reds =<< mkIdentityLambda ts where ts = concatMap (lambdaReturnType . redLambda) reds -- | Construct a Screma corresponding to a map. -mapSOAC :: Lambda rep -> ScremaForm rep -mapSOAC lam = ScremaForm lam [] [] +mapSOAC :: + (Buildable rep, MonadFreshNames m) => + Lambda rep -> + m (ScremaForm rep) +mapSOAC lam = do + post_lam <- mkIdentityLambda $ lambdaReturnType lam + pure $ ScremaForm lam [] [] post_lam -- | Does this Screma correspond to a scan-map composition? isScanomapSOAC :: ScremaForm rep -> Maybe ([Scan rep], Lambda rep) -isScanomapSOAC (ScremaForm map_lam scans reds) = do +isScanomapSOAC (ScremaForm map_lam scans reds post_lam) = do guard $ null reds guard $ not $ null scans + guard $ isIdentityLambda post_lam pure (scans, map_lam) +isMaposcanomapSOAC :: ScremaForm rep -> Maybe (Lambda rep, [Scan rep], Lambda rep) +isMaposcanomapSOAC (ScremaForm map_lam scans reds post_lam) = do + guard $ null reds + guard $ not $ null scans + pure (post_lam, scans, map_lam) + -- | Does this Screma correspond to pure scan? isScanSOAC :: ScremaForm rep -> Maybe [Scan rep] isScanSOAC form = do (scans, map_lam) <- isScanomapSOAC form guard $ isIdentityLambda map_lam + guard $ length (lambdaReturnType map_lam) == scanResults scans pure scans -- | Does this Screma correspond to a reduce-map composition? isRedomapSOAC :: ScremaForm rep -> Maybe ([Reduce rep], Lambda rep) -isRedomapSOAC (ScremaForm map_lam scans reds) = do +isRedomapSOAC (ScremaForm map_lam scans reds post_lam) = do guard $ null scans guard $ not $ null reds + guard $ isIdentityLambda post_lam pure (reds, map_lam) -- | Does this Screma correspond to a pure reduce? @@ -267,14 +377,16 @@ isReduceSOAC form = do (reds, map_lam) <- isRedomapSOAC form guard $ isIdentityLambda map_lam + guard $ length (lambdaReturnType map_lam) == redResults reds pure reds -- | Does this Screma correspond to a simple map, without any -- reduction or scan results? isMapSOAC :: ScremaForm rep -> Maybe (Lambda rep) -isMapSOAC (ScremaForm map_lam scans reds) = do +isMapSOAC (ScremaForm map_lam scans reds post_lam) = do guard $ null scans guard $ null reds + guard $ isIdentityLambda post_lam pure map_lam -- | Like 'Mapper', but just for 'SOAC's. @@ -332,28 +444,29 @@ ) ops <*> mapOnSOACLambda tv bucket_fun -mapSOACM tv (Screma w arrs (ScremaForm map_lam scans reds)) = +mapSOACM tv (Screma w arrs (ScremaForm map_lam scans reds post_lam)) = Screma <$> mapOnSOACSubExp tv w <*> mapM (mapOnSOACVName tv) arrs <*> ( ScremaForm <$> mapOnSOACLambda tv map_lam - <*> forM - scans - ( \(Scan red_lam red_nes) -> - Scan - <$> mapOnSOACLambda tv red_lam - <*> mapM (mapOnSOACSubExp tv) red_nes - ) - <*> forM - reds - ( \(Reduce comm red_lam red_nes) -> - Reduce comm - <$> mapOnSOACLambda tv red_lam - <*> mapM (mapOnSOACSubExp tv) red_nes - ) + <*> mapM (mapOnSOACScan tv) scans + <*> mapM (mapOnSOACReduce tv) reds + <*> mapOnSOACLambda tv post_lam ) +mapOnSOACScan :: (Monad m) => SOACMapper frep trep m -> Scan frep -> m (Scan trep) +mapOnSOACScan tv (Scan red_lam red_nes) = + Scan + <$> mapOnSOACLambda tv red_lam + <*> mapM (mapOnSOACSubExp tv) red_nes + +mapOnSOACReduce :: (Monad m) => SOACMapper frep trep m -> Reduce frep -> m (Reduce trep) +mapOnSOACReduce tv (Reduce comm red_lam red_nes) = + Reduce comm + <$> mapOnSOACLambda tv red_lam + <*> mapM (mapOnSOACSubExp tv) red_nes + -- | A helper for defining 'TraverseOpStms'. traverseSOACStms :: (Monad m) => OpStmsTraverser m (SOAC rep) rep traverseSOACStms f = mapSOACM mapper @@ -367,8 +480,8 @@ freeIn' (Reduce _ lam ne) = freeIn' lam <> freeIn' ne instance (ASTRep rep) => FreeIn (ScremaForm rep) where - freeIn' (ScremaForm scans reds lam) = - freeIn' scans <> freeIn' reds <> freeIn' lam + freeIn' (ScremaForm scans reds lam post_lam) = + freeIn' scans <> freeIn' reds <> freeIn' lam <> freeIn' post_lam instance (ASTRep rep) => FreeIn (HistOp rep) where freeIn' (HistOp w rf dests nes lam) = @@ -429,7 +542,7 @@ consumedInOp VJP {} = mempty -- Only map functions can consume anything. The operands to scan -- and reduce functions are always considered "fresh". - consumedInOp (Screma _ arrs (ScremaForm map_lam _ _)) = + consumedInOp (Screma _ arrs (ScremaForm map_lam _ _ _)) = mapNames consumedArray $ consumedByLambda map_lam where consumedArray v = fromMaybe v $ lookup v params_to_arrs @@ -464,12 +577,13 @@ arrs (map (mapHistOp (Alias.analyseLambda aliases)) ops) (Alias.analyseLambda aliases bucket_fun) - addOpAliases aliases (Screma w arrs (ScremaForm map_lam scans reds)) = + addOpAliases aliases (Screma w arrs (ScremaForm map_lam scans reds post_lam)) = Screma w arrs $ ScremaForm (Alias.analyseLambda aliases map_lam) (map onScan scans) (map onRed reds) + (Alias.analyseLambda aliases post_lam) where onRed red = red {redLambda = Alias.analyseLambda aliases $ redLambda red} onScan scan = scan {scanLambda = Alias.analyseLambda aliases $ scanLambda scan} @@ -514,7 +628,7 @@ lam (zipWith (<>) (map depsOf' args) (map depsOf' vec)) <> map (const $ freeIn args <> freeIn lam) (lambdaParams lam) - opDependencies (Screma w arrs (ScremaForm map_lam scans reds)) = + opDependencies (Screma w arrs (ScremaForm map_lam scans reds post_lam)) = let (scans_in, reds_in, map_deps) = splitAt3 (scanResults scans) (redResults reds) $ lambdaDependencies mempty map_lam (depsOfArrays w arrs) @@ -522,7 +636,7 @@ concatMap depsOfScan (zip scans $ chunks (scanSizes scans) scans_in) reds_deps = concatMap depsOfRed (zip reds $ chunks (redSizes reds) reds_in) - in scans_deps <> reds_deps <> map_deps + in reds_deps <> lambdaDependencies mempty post_lam (scans_deps <> map_deps) where depsOfScan (Scan lam nes, deps_in) = reductionDependencies mempty lam nes deps_in @@ -554,7 +668,9 @@ SubExpRes _ (Var v) -> uncurry (flip ST.Indexed) <$> M.lookup v arr_indexes' _ -> Nothing where - lambdaAndSubExp (Screma _ arrs (ScremaForm map_lam scans reds)) = + lambdaAndSubExp (Screma _ arrs (ScremaForm map_lam scans reds post_lam)) = do + -- UNSURE_IF_CORRECT + guard $ isIdentityLambda post_lam nthMapOut (scanResults scans + redResults reds) map_lam arrs lambdaAndSubExp _ = Nothing @@ -665,49 +781,57 @@ <> prettyTuple (lambdaReturnType bucket_fun) <> " but should have type " <> prettyTuple bucket_ret_t -typeCheckSOAC (Screma w arrs (ScremaForm map_lam scans reds)) = do +typeCheckSOAC (Screma w arrs (ScremaForm map_lam scans reds post_lam)) = do TC.require (Prim int64) w arrs' <- TC.checkSOACArrayArgs w arrs TC.checkLambda map_lam arrs' - - scan_nes' <- fmap concat $ - forM scans $ \(Scan scan_lam scan_nes) -> do - scan_nes' <- mapM TC.checkArg scan_nes - let scan_t = map TC.argType scan_nes' - TC.checkLambda scan_lam $ map TC.noArgAliases $ scan_nes' ++ scan_nes' - unless (scan_t == lambdaReturnType scan_lam) $ - TC.bad . TC.TypeError $ - "Scan function returns type " - <> prettyTuple (lambdaReturnType scan_lam) - <> " but neutral element has type " - <> prettyTuple scan_t - pure scan_nes' - - red_nes' <- fmap concat $ - forM reds $ \(Reduce _ red_lam red_nes) -> do - red_nes' <- mapM TC.checkArg red_nes - let red_t = map TC.argType red_nes' - TC.checkLambda red_lam $ map TC.noArgAliases $ red_nes' ++ red_nes' - unless (red_t == lambdaReturnType red_lam) $ - TC.bad . TC.TypeError $ - "Reduce function returns type " - <> prettyTuple (lambdaReturnType red_lam) - <> " but neutral element has type " - <> prettyTuple red_t - pure red_nes' - + scan_nes' <- concat <$> mapM typeCheckScan scans + red_nes' <- concat <$> mapM typeCheckReduce reds let map_lam_ts = lambdaReturnType map_lam - unless ( take (length scan_nes' + length red_nes') map_lam_ts == map TC.argType (scan_nes' ++ red_nes') ) . TC.bad . TC.TypeError - $ "Map function return type " + $ "Pre-lambda function return type " <> prettyTuple map_lam_ts <> " wrong for given scan and reduction functions." + let (scan_ts, _, map_ts) = + splitAt3 (length scan_nes') (length red_nes') map_lam_ts + post_lam_args = map (,mempty) $ scan_ts <> map_ts + TC.checkLambda post_lam post_lam_args + when (null scans && not (isIdentityLambda post_lam)) $ + TC.bad $ + TC.TypeError "Screma has post-lambda but no scan operations." + +typeCheckScan :: (TC.Checkable rep) => Scan (Aliases rep) -> TC.TypeM rep [(Type, Names)] +typeCheckScan (Scan scan_lam scan_nes) = do + scan_nes' <- mapM TC.checkArg scan_nes + let scan_t = map TC.argType scan_nes' + TC.checkLambda scan_lam $ map TC.noArgAliases $ scan_nes' ++ scan_nes' + unless (scan_t == lambdaReturnType scan_lam) $ + TC.bad . TC.TypeError $ + "Scan function returns type " + <> prettyTuple (lambdaReturnType scan_lam) + <> " but neutral element has type " + <> prettyTuple scan_t + pure scan_nes' + +typeCheckReduce :: (TC.Checkable rep) => Reduce (Aliases rep) -> TC.TypeM rep [(Type, Names)] +typeCheckReduce (Reduce _ red_lam red_nes) = do + red_nes' <- mapM TC.checkArg red_nes + let red_t = map TC.argType red_nes' + TC.checkLambda red_lam $ map TC.noArgAliases $ red_nes' ++ red_nes' + unless (red_t == lambdaReturnType red_lam) $ + TC.bad . TC.TypeError $ + "Reduce function returns type " + <> prettyTuple (lambdaReturnType red_lam) + <> " but neutral element has type " + <> prettyTuple red_t + pure red_nes' + instance RephraseOp SOAC where rephraseInOp r (VJP args vec lam) = VJP args vec <$> rephraseLambda r lam @@ -720,17 +844,23 @@ where onOp (HistOp dest_shape rf dests nes op) = HistOp dest_shape rf dests nes <$> rephraseLambda r op - rephraseInOp r (Screma w arrs (ScremaForm lam scans red)) = + rephraseInOp r (Screma w arrs (ScremaForm lam scans red post_lam)) = Screma w arrs <$> ( ScremaForm <$> rephraseLambda r lam - <*> mapM onScan scans - <*> mapM onRed red + <*> mapM (rephraseScan r) scans + <*> mapM (rephraseRed r) red + <*> rephraseLambda r post_lam ) - where - onScan (Scan op nes) = Scan <$> rephraseLambda r op <*> pure nes - onRed (Reduce comm op nes) = Reduce comm <$> rephraseLambda r op <*> pure nes +rephraseRed :: (Monad m) => Rephraser m from to -> Reduce from -> m (Reduce to) +rephraseRed r (Reduce comm op nes) = + Reduce comm <$> rephraseLambda r op <*> pure nes + +rephraseScan :: (Monad m) => Rephraser m from to -> Scan from -> m (Scan to) +rephraseScan r (Scan op nes) = + Scan <$> rephraseLambda r op <*> pure nes + instance (OpMetrics (Op rep)) => OpMetrics (SOAC rep) where opMetrics (VJP _ _ lam) = inside "VJP" $ lambdaMetrics lam @@ -740,11 +870,12 @@ inside "Stream" $ lambdaMetrics lam opMetrics (Hist _ _ ops bucket_fun) = inside "Hist" $ mapM_ (lambdaMetrics . histOp) ops >> lambdaMetrics bucket_fun - opMetrics (Screma _ _ (ScremaForm map_lam scans reds)) = + opMetrics (Screma _ _ (ScremaForm map_lam scans reds post_lam)) = inside "Screma" $ do lambdaMetrics map_lam mapM_ (lambdaMetrics . scanLambda) scans mapM_ (lambdaMetrics . redLambda) reds + lambdaMetrics post_lam instance (PrettyRep rep) => PP.Pretty (SOAC rep) where pretty (VJP args vec lam) = @@ -767,16 +898,16 @@ ppStream size arrs acc lam pretty (Hist w arrs ops bucket_fun) = ppHist w arrs ops bucket_fun - pretty (Screma w arrs (ScremaForm map_lam scans reds)) - | null scans, - null reds = + pretty (Screma w arrs screma) + | Just map_lam <- isMapSOAC screma = "map" <> (parens . align) ( pretty w <> comma ppTuple' (map pretty arrs) <> comma pretty map_lam + <> comma pretty (scremaPostLambda screma) ) - | null scans = + | Just (reds, map_lam) <- isRedomapSOAC screma = "redomap" <> (parens . align) ( pretty w @@ -784,8 +915,9 @@ <> comma pretty map_lam <> comma PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty reds) + <> comma pretty (scremaPostLambda screma) ) - | null reds = + | Just (scans, map_lam) <- isScanomapSOAC screma = "scanomap" <> (parens . align) ( pretty w @@ -794,13 +926,14 @@ <> comma PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty scans) + <> comma pretty (scremaPostLambda screma) ) pretty (Screma w arrs form) = ppScrema w arrs form -- | Prettyprint the given Screma. ppScrema :: (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> ScremaForm rep -> Doc ann -ppScrema w arrs (ScremaForm map_lam scans reds) = +ppScrema w arrs (ScremaForm map_lam scans reds post_lam) = "screma" <> (parens . align) ( pretty w @@ -810,6 +943,7 @@ PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty scans) <> comma PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty reds) + <> comma pretty post_lam ) -- | Prettyprint the given Stream. @@ -872,3 +1006,15 @@ ppTuple' (map pretty nes) <> comma pretty op + +instance (PrettyRep rep) => PP.Pretty (ScremaForm rep) where + pretty (ScremaForm pre_lam scans reds post_lam) = + "screma" + <> (parens . align) + ( pretty pre_lam + <> comma + PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty scans) + <> comma + PP.braces (mconcat $ intersperse (comma <> PP.line) $ map pretty reds) + <> comma pretty post_lam + ) diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs --- a/src/Futhark/IR/SOACS/Simplify.hs +++ b/src/Futhark/IR/SOACS/Simplify.hs @@ -17,6 +17,10 @@ liftIdentityMapping, simplifyMapIota, SOACS, + eliminate, + eliminateByRes, + prunePreLambdaResults, + dedupInput, ) where @@ -27,7 +31,9 @@ import Data.Bifunctor import Data.Either import Data.Foldable +import Data.Function (on) import Data.List (partition, transpose, unzip4) +import Data.List qualified as L import Data.List.NonEmpty (NonEmpty (..)) import Data.Map.Strict qualified as M import Data.Maybe @@ -114,29 +120,36 @@ imgs' <- mapM Engine.simplify imgs (bfun', bfun_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty bfun pure (Hist w' imgs' ops' bfun', mconcat hoisted <> bfun_hoisted) -simplifySOAC (Screma w arrs (ScremaForm map_lam scans reds)) = do - (scans', scans_hoisted) <- fmap unzip $ - forM scans $ \(Scan lam nes) -> do - (lam', hoisted) <- Engine.simplifyLambda mempty lam - nes' <- Engine.simplify nes - pure (Scan lam' nes', hoisted) - - (reds', reds_hoisted) <- fmap unzip $ - forM reds $ \(Reduce comm lam nes) -> do - (lam', hoisted) <- Engine.simplifyLambda mempty lam - nes' <- Engine.simplify nes - pure (Reduce comm lam' nes', hoisted) - +simplifySOAC (Screma w arrs (ScremaForm map_lam scans reds post_lam)) = do + (scans', scans_hoisted) <- mapAndUnzipM simplifyScan scans + (reds', reds_hoisted) <- mapAndUnzipM simplifyReduce reds (map_lam', map_lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty map_lam + (post_lam', post_lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty post_lam (,) <$> ( Screma <$> Engine.simplify w <*> Engine.simplify arrs - <*> pure (ScremaForm map_lam' scans' reds') + <*> pure (ScremaForm map_lam' scans' reds' post_lam') ) - <*> pure (mconcat scans_hoisted <> mconcat reds_hoisted <> map_lam_hoisted) + <*> pure (mconcat scans_hoisted <> mconcat reds_hoisted <> map_lam_hoisted <> post_lam_hoisted) +simplifyScan :: + (Simplify.SimplifiableRep rep) => + Simplify.SimplifyOp rep (Scan (Wise rep)) +simplifyScan (Scan lam nes) = do + (lam', hoisted) <- Engine.simplifyLambda mempty lam + nes' <- Engine.simplify nes + pure (Scan lam' nes', hoisted) + +simplifyReduce :: + (Simplify.SimplifiableRep rep) => + Simplify.SimplifyOp rep (Reduce (Wise rep)) +simplifyReduce (Reduce comm lam nes) = do + (lam', hoisted) <- Engine.simplifyLambda mempty lam + nes' <- Engine.simplify nes + pure (Reduce comm lam' nes', hoisted) + instance BuilderOps (Wise SOACS) instance TraverseOpStms (Wise SOACS) where @@ -178,6 +191,7 @@ lam_body' = lam_body {bodyResult = keep' $ bodyResult lam_body} ret = keep' $ lambdaReturnType lam +{-# NOINLINE soacRules #-} soacRules :: RuleBook (Wise SOACS) soacRules = standardRules <> ruleBook topDownRules bottomUpRules @@ -203,12 +217,15 @@ RuleOp fuseConcatScatter, RuleOp simplifyMapIota, RuleOp moveTransformToInput, - RuleOp moveTransformToOutput + RuleOp moveTransformToOutput, + RuleOp removeUnusedPreLamResult, + RuleOp removeDeadScan, + RuleOp removeDuplicateInput ] bottomUpRules :: [BottomUpRule (Wise SOACS)] bottomUpRules = - [ RuleOp removeDeadMapping, + [ RuleOp removeDeadResult, RuleOp removeDeadReduction, RuleBasicOp removeUnnecessaryCopy, RuleOp liftIdentityStreaming, @@ -292,8 +309,8 @@ } auxing aux $ do mapM_ (uncurry letBind) invariant - letBindNames (map patElemName pat') . Op $ - soacOp (Screma w arrs (mapSOAC fun')) + letBindNames (map patElemName pat') . Op . soacOp . Screma w arrs + =<< mapSOAC fun' liftIdentityMapping _ _ _ _ = Skip liftIdentityStreaming :: BottomUpRuleOp (Wise SOACS) @@ -331,14 +348,15 @@ -- | Remove all arguments to the map that are simply replicates. -- These can be turned into free variables instead. removeReplicateMapping :: - (Aliased rep, BuilderOps rep, HasSOAC rep) => + (Aliased rep, Buildable rep, BuilderOps rep, HasSOAC rep) => TopDownRuleOp rep removeReplicateMapping vtable pat aux op | Just (Screma w arrs form) <- asSOAC op, Just fun <- isMapSOAC form, Just (stms, fun', arrs') <- removeReplicateInput vtable fun arrs = Simplify $ do forM_ stms $ \(vs, cs, e) -> certifying cs $ letBindNames vs e - auxing aux $ letBind pat $ Op $ soacOp $ Screma w arrs' $ mapSOAC fun' + auxing aux . letBind pat . Op . soacOp . Screma w arrs' + =<< mapSOAC fun' removeReplicateMapping _ _ _ _ = Skip removeReplicateInput :: @@ -385,10 +403,10 @@ TopDownRuleOp rep removeUnusedSOACInput _ pat aux op | Just (Screma w arrs form :: SOAC rep) <- asSOAC op, - ScremaForm map_lam scan reduce <- form, + ScremaForm map_lam scan reduce post_lam <- form, Just (used_arrs, map_lam') <- remove map_lam arrs = Simplify . auxing aux . letBind pat . Op $ - soacOp (Screma w used_arrs (ScremaForm map_lam' scan reduce)) + soacOp (Screma w used_arrs (ScremaForm map_lam' scan reduce post_lam)) where used_in_body map_lam = freeIn $ lambdaBody map_lam usedInput map_lam (param, _) = paramName param `nameIn` used_in_body map_lam @@ -399,29 +417,68 @@ in if null unused then Nothing else Just (used_arrs, map_lam') removeUnusedSOACInput _ _ _ _ = Skip -removeDeadMapping :: BottomUpRuleOp (Wise SOACS) -removeDeadMapping (_, used) (Pat pes) aux (Screma w arrs (ScremaForm lam scans reds)) - | (nonmap_pes, map_pes) <- splitAt num_nonmap_res pes, - not $ null map_pes = - let (nonmap_res, map_res) = splitAt num_nonmap_res $ bodyResult $ lambdaBody lam - (nonmap_ts, map_ts) = splitAt num_nonmap_res $ lambdaReturnType lam +removeDeadResult :: BottomUpRuleOp (Wise SOACS) +removeDeadResult (_, used) (Pat pes) aux (Screma w arrs (ScremaForm lam scans reds post_lam)) + | (red_pes, post_pes) <- splitAt (redResults reds) pes, + not $ null post_pes = + let res = bodyResult $ lambdaBody post_lam + ts = lambdaReturnType post_lam isUsed (bindee, _, _) = (`UT.used` used) $ patElemName bindee - (map_pes', map_res', map_ts') = - unzip3 $ filter isUsed $ zip3 map_pes map_res map_ts - lam' = - lam - { lambdaBody = (lambdaBody lam) {bodyResult = nonmap_res <> map_res'}, - lambdaReturnType = nonmap_ts <> map_ts' + (post_pes', res', ts') = + unzip3 $ filter isUsed $ zip3 post_pes res ts + post_lam' = + post_lam + { lambdaBody = mkBody (bodyStms (lambdaBody post_lam)) res', + lambdaReturnType = ts' } - in if map_pes /= map_pes' + in if post_pes /= post_pes' then + -- It is possible that we produce a Screma here that has a + -- non-identity post_lam, but that will be cleaned up by + -- removeUnusedPreLamResult. Simplify . auxing aux $ - letBind (Pat $ nonmap_pes <> map_pes') . Op $ - Screma w arrs (ScremaForm lam' scans reds) + letBind (Pat $ red_pes <> post_pes') . Op $ + Screma w arrs (ScremaForm lam scans reds post_lam') else Skip +removeDeadResult _ _ _ _ = Skip + +-- | If we have pre-lambda result that is passed to the post-lambda, but not +-- actually used for anything in the post-lambda, then get rid of it. +removeUnusedPreLamResult :: TopDownRuleOp (Wise SOACS) +removeUnusedPreLamResult _ pat aux (Screma w arrs (ScremaForm pre_lam scans reds post_lam)) + | not $ and used_mask = Simplify $ do + let pre_lam_new_res = + keep + (replicate num_scanred_results True ++ used_mask) + (bodyResult (lambdaBody pre_lam)) + pre_lam_new_ts = + keep + (replicate num_scanred_results True ++ used_mask) + (lambdaReturnType pre_lam) + pre_lam' = + pre_lam + { lambdaBody = (lambdaBody pre_lam) {bodyResult = pre_lam_new_res}, + lambdaReturnType = pre_lam_new_ts + } + post_lam' = + post_lam + { lambdaParams = + keep + (replicate (scanResults scans) True ++ used_mask) + (lambdaParams post_lam) + } + auxing aux . letBind pat $ + Op (Screma w arrs $ ScremaForm pre_lam' scans reds post_lam') where - num_nonmap_res = scanResults scans + redResults reds -removeDeadMapping _ _ _ _ = Skip + keep mask xs = map snd $ filter fst $ zip mask xs + num_scanred_results = scanResults scans + redResults reds + -- This mask covers only the results/parameters directly passed from prelam + -- to postlam. + used_mask = + map + ((`nameIn` freeIn (lambdaBody post_lam)) . paramName) + (drop (scanResults scans) (lambdaParams post_lam)) +removeUnusedPreLamResult _ _ _ _ = Skip removeDuplicateMapOutput :: TopDownRuleOp (Wise SOACS) removeDuplicateMapOutput _ (Pat pes) aux (Screma w arrs form) @@ -441,7 +498,7 @@ lambdaReturnType = ts' } auxing aux $ do - letBind (Pat pes') $ Op $ Screma w arrs $ mapSOAC fun' + letBind (Pat pes') . Op . Screma w arrs =<< mapSOAC fun' forM_ copies $ \(from, to) -> letBind (Pat [to]) $ BasicOp $ Replicate mempty $ Var $ patElemName from where @@ -453,6 +510,18 @@ | otherwise = (ses_ts_pes' ++ [(se, t, pe)], copies) removeDuplicateMapOutput _ _ _ _ = Skip +removeDuplicateInput :: TopDownRuleOp (Wise SOACS) +removeDuplicateInput _ pat aux (Screma w arrs form) + | length arrs /= length (nubOrd arrs) = Simplify $ do + let (new_arrs, new_form) = dedupInput arrs form + auxing aux + . letBind pat + . Op + . Screma w new_arrs + $ new_form + | otherwise = Skip +removeDuplicateInput _ _ _ _ = Skip + reshapeInner :: SubExp -> NewShape SubExp -> NewShape SubExp reshapeInner w new_shape = reshapeCoerce outer <> newshapeInner outer new_shape @@ -503,62 +572,94 @@ Just (map_pe, stmAuxCerts aux2, w, e', lambdaParams map_lam, arrs) | otherwise = Nothing +-- Shared checks for removeDeadReduction/removeDeadScan. +deadRedScanCheck :: + (p -> Bool) -> + Lambda (Wise SOACS) -> + [SubExp] -> + [p] -> + ([p], [p], [SubExp], [Bool]) +deadRedScanCheck needed oplam nes postlam_params = + let (postlam_scanparams, postlam_mapparams) = + splitAt (length nes) postlam_params + oplam_deps = dataDependencies $ lambdaBody oplam + oplam_res = bodyResult $ lambdaBody oplam + oplam_params = lambdaParams oplam + (oplam_xparams, oplam_yparams) = + splitAt (length nes) oplam_params + used_in_postlam = + map snd $ + filter (needed . fst) $ + zip postlam_scanparams oplam_params + necessary = + findNecessaryForReturned + (`elem` used_in_postlam) + (zip oplam_params $ map resSubExp $ oplam_res <> oplam_res) + oplam_deps + alive_mask = + zipWith + (||) + (map ((`nameIn` necessary) . paramName) oplam_xparams) + (map ((`nameIn` necessary) . paramName) oplam_yparams) + (used_postlam_scanparams, used_nes) = + unzip . map snd . filter fst $ zip alive_mask $ zip postlam_scanparams nes + in (used_postlam_scanparams, postlam_mapparams, used_nes, alive_mask) + -- | Some of the results of a reduction (or really: Redomap) may be -- dead. We remove them here. The trick is that we need to look at -- the data dependencies to see that the "dead" result is not -- actually used for computing one of the live ones. removeDeadReduction :: BottomUpRuleOp (Wise SOACS) -removeDeadReduction (_, used) pat aux (Screma w arrs form) = - case isRedomapSOAC form of - Just ([Reduce comm redlam rednes], maplam) -> - let mkOp lam nes' = redomapSOAC [Reduce comm lam nes'] - in removeDeadReduction' redlam rednes maplam mkOp - _ -> - case isScanomapSOAC form of - Just ([Scan scanlam nes], maplam) -> - let mkOp lam nes' = scanomapSOAC [Scan lam nes'] - in removeDeadReduction' scanlam nes maplam mkOp - _ -> Skip - where - removeDeadReduction' redlam nes maplam mkOp - | not $ all (`UT.used` used) $ patNames pat, -- Quick/cheap check - let (red_pes, map_pes) = splitAt (length nes) $ patElems pat, - let redlam_deps = dataDependencies $ lambdaBody redlam, - let redlam_res = bodyResult $ lambdaBody redlam, - let redlam_params = lambdaParams redlam, - let (redlam_xparams, redlam_yparams) = - splitAt (length nes) redlam_params, - let used_after = - map snd . filter ((`UT.used` used) . patElemName . fst) $ - zip (red_pes <> red_pes) redlam_params, - let necessary = - findNecessaryForReturned - (`elem` used_after) - (zip redlam_params $ map resSubExp $ redlam_res <> redlam_res) - redlam_deps, - let alive_mask = - zipWith - (||) - (map ((`nameIn` necessary) . paramName) redlam_xparams) - (map ((`nameIn` necessary) . paramName) redlam_yparams), - not $ and alive_mask = Simplify $ do - let fixDeadToNeutral lives ne = if lives then Nothing else Just ne - dead_fix = zipWith fixDeadToNeutral alive_mask nes - (used_red_pes, used_nes) = - unzip . map snd . filter fst $ zip alive_mask $ zip red_pes nes - - when (used_nes == nes) cannotSimplify +removeDeadReduction (_, used) pat aux (Screma w arrs form) + | Just ([Reduce comm redlam nes], maplam) <- isRedomapSOAC form, + not $ all (`UT.used` used) $ patNames pat, -- Quick/cheap check + (used_red_pes, map_pes, used_nes, alive_mask) <- + deadRedScanCheck ((`UT.used` used) . patElemName) redlam nes (patElems pat), + used_nes /= nes = Simplify $ do + let fixDeadToNeutral lives ne = if lives then Nothing else Just ne + dead_fix = zipWith fixDeadToNeutral alive_mask nes - let maplam' = removeLambdaResults alive_mask maplam - redlam' <- - removeLambdaResults alive_mask - <$> fixLambdaParams redlam (dead_fix ++ dead_fix) + let maplam' = removeLambdaResults alive_mask maplam + redlam' <- + removeLambdaResults alive_mask + <$> fixLambdaParams redlam (dead_fix ++ dead_fix) - auxing aux . letBind (Pat $ used_red_pes ++ map_pes) . Op $ - Screma w arrs (mkOp redlam' used_nes maplam') - removeDeadReduction' _ _ _ _ = Skip + auxing aux + . letBind (Pat $ used_red_pes ++ map_pes) + . Op + . Screma w arrs + =<< redomapSOAC [Reduce comm redlam' used_nes] maplam' removeDeadReduction _ _ _ _ = Skip +{-# NOINLINE removeDeadScan #-} +removeDeadScan :: TopDownRuleOp (Wise SOACS) +removeDeadScan _ pat aux (Screma w arrs form) + | ScremaForm prelam [Scan scanlam nes] [] postlam <- form, + -- Quick/cheap check + not $ all ((`nameIn` freeIn postlam) . paramName) (lambdaParams postlam), + let used = (`nameIn` freeIn (lambdaBody postlam)) . paramName, + (used_postlam_scanparams, postlam_mapparams, used_nes, alive_mask) <- + deadRedScanCheck used scanlam nes (lambdaParams postlam), + used_nes /= nes = Simplify $ do + let fixDeadToNeutral lives ne = if lives then Nothing else Just ne + dead_fix = zipWith fixDeadToNeutral alive_mask nes + + let prelam' = removeLambdaResults alive_mask prelam + scanlam' <- + removeLambdaResults alive_mask + <$> fixLambdaParams scanlam (dead_fix ++ dead_fix) + + postlam' <- + runLambdaBuilder (used_postlam_scanparams <> postlam_mapparams) $ + bodyBind (lambdaBody postlam) + + auxing aux + . letBind pat + . Op + . Screma w arrs + $ ScremaForm prelam' [Scan scanlam' used_nes] [] postlam' +removeDeadScan _ _ _ _ = Skip + {-# NOINLINE fuseConcatScatter #-} fuseConcatScatter :: TopDownRuleOp (Wise SOACS) fuseConcatScatter vtable pat aux (Screma _ arrs form) @@ -578,7 +679,8 @@ lam' <- mkLambda (acc_params <> input_params) $ subExpsRes <$> recurse (map (Var . paramName) acc_params) lams - letBind pat $ Op $ Screma w' (map snd accs <> concat xivs) (mapSOAC lam') + letBind pat . Op . Screma w' (map snd accs <> concat xivs) + =<< mapSOAC lam' where recurse accs [] = pure accs recurse accs (lam : lams) = do @@ -622,13 +724,12 @@ (Buildable rep, BuilderOps rep, HasSOAC rep, Alias.AliasableRep rep) => TopDownRuleOp rep simplifyKnownIterationSOAC _ pat _ op - | Just (Screma (Constant k) arrs (ScremaForm map_lam scans reds)) <- asSOAC op, + | Just (Screma (Constant k) arrs (ScremaForm map_lam scans reds post_lam)) <- asSOAC op, oneIsh k = Simplify $ do let (Reduce _ red_lam red_nes) = singleReduce reds (Scan scan_lam scan_nes) = singleScan scans - (scan_pes, red_pes, map_pes) = - splitAt3 (length scan_nes) (length red_nes) $ - patElems pat + (red_pes, post_pes) = + splitAt (length red_nes) $ patElems pat bindMapParam p a = do a_t <- lookupType a letBindNames [paramName p] . BasicOp $ @@ -641,8 +742,8 @@ ArrayLit [se] $ rowType $ patElemType pe - bindResult pe (SubExpRes cs se) = - certifying cs $ letBindNames [patElemName pe] $ BasicOp $ SubExp se + bindResult name (SubExpRes cs se) = + certifying cs $ letBindNames [name] $ BasicOp $ SubExp se zipWithM_ bindMapParam (lambdaParams map_lam) arrs (to_scan, to_red, map_res) <- @@ -651,9 +752,12 @@ scan_res <- eLambda scan_lam $ map eSubExp $ scan_nes ++ map resSubExp to_scan red_res <- eLambda red_lam $ map eSubExp $ red_nes ++ map resSubExp to_red - zipWithM_ bindArrayResult scan_pes scan_res - zipWithM_ bindResult red_pes red_res - zipWithM_ bindArrayResult map_pes map_res + zipWithM_ bindResult (patElemName <$> red_pes) red_res + + let par_names = paramName <$> lambdaParams post_lam + zipWithM_ bindResult par_names (scan_res <> map_res) + + zipWithM_ bindArrayResult post_pes =<< bodyBind (lambdaBody post_lam) simplifyKnownIterationSOAC _ pat _ op | Just (Stream (Constant k) arrs nes fold_lam) <- asSOAC op, oneIsh k = Simplify $ do @@ -818,12 +922,14 @@ -- case - if you find yourself planning to extend it to handle more -- complex situations (rotate or whatnot), consider turning it into a -- separate compiler pass instead. +-- +-- NOTE: Mybe this should also be handled in the post lambda. simplifyMapIota :: forall rep. (Buildable rep, BuilderOps rep, HasSOAC rep) => TopDownRuleOp rep simplifyMapIota vtable screma_pat aux op - | Just (Screma w arrs (ScremaForm map_lam scan reduce) :: SOAC rep) <- asSOAC op, + | Just (Screma w arrs (ScremaForm map_lam scan reduce post_lam) :: SOAC rep) <- asSOAC op, Just (p, _) <- find isIota (zip (lambdaParams map_lam) arrs), indexings <- mapMaybe (indexesWith (paramName p)) . S.toList $ @@ -842,7 +948,7 @@ } auxing aux . letBind screma_pat . Op . soacOp $ - Screma w (arrs <> more_arrs) (ScremaForm map_lam' scan reduce) + Screma w (arrs <> more_arrs) (ScremaForm map_lam' scan reduce post_lam) where isIota (_, arr) = case ST.lookupBasicOp arr vtable of Just (Iota _ (Constant o) (Constant s) _, _) -> @@ -895,7 +1001,7 @@ -- corresponding to that transformation performed on the rows of the -- full array. moveTransformToInput :: TopDownRuleOp (Wise SOACS) -moveTransformToInput vtable screma_pat aux soac@(Screma w arrs (ScremaForm map_lam scan reduce)) +moveTransformToInput vtable screma_pat aux soac@(Screma w arrs (ScremaForm map_lam scan reduce post_lam)) | ops <- filter arrayIsMapParam $ S.toList $ arrayOps mempty $ lambdaBody map_lam, not $ null ops = Simplify $ do (more_arrs, more_params, replacements) <- @@ -910,7 +1016,7 @@ } auxing aux . letBind screma_pat . Op $ - Screma w (arrs <> more_arrs) (ScremaForm map_lam' scan reduce) + Screma w (arrs <> more_arrs) (ScremaForm map_lam' scan reduce post_lam) where -- It is not safe to move the transform if the root array is being -- consumed by the Screma. This is a bit too conservative - it's @@ -1000,22 +1106,27 @@ -- TODO: currently we only handle reshapes here, but the principle -- should actually hold for any ArrayTransform. moveTransformToOutput :: TopDownRuleOp (Wise SOACS) -moveTransformToOutput vtable screma_pat screma_aux (Screma w arrs (ScremaForm map_lam scan reduce)) +moveTransformToOutput vtable screma_pat screma_aux (Screma w arrs (ScremaForm map_lam scan reduce post_lam)) | (transformed, map_infos, stms') <- foldl' onStm ([], zip3 map_res map_rets map_pes, mempty) $ bodyStms $ lambdaBody map_lam, (map_res', map_rets', map_pes') <- unzip3 map_infos, - not $ null transformed = Simplify $ do - (tr_res, tr_rets, tr_names, post) <- unzip4 <$> mapM mkTransformed transformed - let map_lam' = - map_lam - { lambdaBody = mkBody stms' $ nonmap_res <> map_res' <> tr_res, - lambdaReturnType = nonmap_rets <> map_rets' <> tr_rets - } - pat_names = map patElemName (nonmap_pes <> map_pes') <> tr_names - auxing screma_aux . letBindNames pat_names . Op $ - Screma w arrs (ScremaForm map_lam' scan reduce) - sequence_ post + not $ null transformed, + -- Should probably account for non-identity lambda. + isIdentityLambda post_lam = + Simplify $ do + (tr_res, tr_rets, tr_names, post) <- unzip4 <$> mapM mkTransformed transformed + let map_lam' = + map_lam + { lambdaBody = mkBody stms' $ nonmap_res <> map_res' <> tr_res, + lambdaReturnType = nonmap_rets <> map_rets' <> tr_rets + } + pat_names = map patElemName (nonmap_pes <> map_pes') <> tr_names + post_lam' <- mkIdentityLambda (scan_ts <> map_rets' <> tr_rets) + auxing screma_aux . letBindNames pat_names . Op $ + Screma w arrs (ScremaForm map_lam' scan reduce post_lam') + sequence_ post where + scan_ts = concatMap (lambdaReturnType . scanLambda) scan num_nonmap_res = scanResults scan + redResults reduce (nonmap_pes, map_pes) = splitAt num_nonmap_res $ patElems screma_pat @@ -1046,3 +1157,194 @@ pure (SubExpRes cs (Var arr), t, v, bind) moveTransformToOutput _ _ _ _ = Skip + +-- | Eliminate statements if it is not an dependency used to form the +-- names given. +eliminate :: (Buildable rep) => Names -> Stms rep -> Stms rep +eliminate = auxiliary (stmsFromList []) + where + auxiliary stms' deps stms + | Just (stms'', stm@(Let v aux e)) <- stmsLast stms = + if namesIntersect deps $ namesFromList $ patNames v + then + auxiliary (oneStm stm <> stms') (freeIn (aux, e) <> deps) stms'' + else + auxiliary stms' deps stms'' + | otherwise = stms' + +-- | Eliminate statements inside a lambda if they are not used to +-- compute the result. +eliminateByRes :: (Buildable rep) => Lambda rep -> Lambda rep +eliminateByRes lam = lam {lambdaBody = mkBody new_stms res} + where + res = bodyResult $ lambdaBody lam + stms = bodyStms $ lambdaBody lam + new_stms = eliminate (freeIn res) stms + +-- | Prunes unused map results from the pre-lambda in a ScremaForm. +-- +-- A ScremaForm contains pre-lambda and post-lambda functions where +-- results from the pre-lambda are passed as arguments to the +-- post-lambda. This function eliminates map-related results from the +-- pre-lambda that don't contribute to the post-lambda's +-- +-- results: +-- 1. Identifies which post-lambda parameters are never used (dead) +-- 2. Removes those parameters from the post-lambda +-- 3. Removes the corresponding results from the pre-lambda +-- 4. Removes the corresponding return types from the pre-lambda +-- +-- Only affects map results; scan and reduction results are preserved. +-- +-- Returns: A ScremaForm with unused pre-lambda map results +-- eliminated. +prunePreLambdaMapResults :: (Buildable rep) => ScremaForm rep -> ScremaForm rep +prunePreLambdaMapResults (ScremaForm pre_lam scan red post_lam) = + ScremaForm new_pre_lam scan red new_post_lam + where + (rest_res_p, map_res_p) = + splitAt (scanResults scan + redResults red) . bodyResult $ lambdaBody pre_lam + (rest_ts_p, map_ts_p) = + splitAt (scanResults scan + redResults red) $ lambdaReturnType pre_lam + (scan_pars_c, map_pars_c) = + splitAt (scanResults scan) $ lambdaParams temp_post_lam + new_post_lam = temp_post_lam {lambdaParams = scan_pars_c <> new_map_pars_c} + new_pre_lam = + eliminateByRes $ + pre_lam + { lambdaBody = + mkBody + (bodyStms $ lambdaBody pre_lam) + (rest_res_p <> new_map_res_p), + lambdaReturnType = rest_ts_p <> new_map_ts_p + } + (new_map_res_p, new_map_ts_p, new_map_pars_c) = + unzip3 + . filter (\(_, _, p) -> paramName p `nameIn` deps) + $ zip3 map_res_p map_ts_p map_pars_c + temp_post_lam = eliminateByRes post_lam + deps = freeIn $ lambdaBody temp_post_lam + +-- | Prunes unused scan results from the pre-lambda in a ScremaForm. +-- +-- Similar to 'prunePreLambdaMapResults', but for scan +-- operations. Removes entire scan operations when none of their +-- pre-lambda results are used to produce the post-lambda's results. +-- +-- For each scan operation (which produces multiple results in the +-- pre-lambda), checks if ANY of the corresponding post-lambda +-- parameters are used. If not, the entire scan and its results are +-- eliminated. +-- +-- Returns: A ScremaForm with unused pre-lambda scan results +-- eliminated. +prunePreLambdaScanResults :: (Buildable rep) => ScremaForm rep -> ScremaForm rep +prunePreLambdaScanResults (ScremaForm pre_lam scan red post_lam) = + ScremaForm new_pre_lam new_scan red new_post_lam + where + (scan_res_p, rest_res_p) = + splitAt (scanResults scan) . bodyResult $ lambdaBody pre_lam + (scan_ts_p, rest_ts_p) = + splitAt (scanResults scan) $ lambdaReturnType pre_lam + (scan_pars_c, map_pars_c) = + splitAt (scanResults scan) $ lambdaParams temp_post_lam + new_post_lam = temp_post_lam {lambdaParams = mconcat new_scan_pars_c <> map_pars_c} + new_pre_lam = + eliminateByRes $ + pre_lam + { lambdaBody = + mkBody + (bodyStms $ lambdaBody pre_lam) + (mconcat new_scan_res_p <> rest_res_p), + lambdaReturnType = mconcat new_scan_ts_p <> rest_ts_p + } + + chunkByScan :: [a] -> [[a]] + chunkByScan = chunks (map (length . scanNeutral) scan) + + chunked_scan_res_p = chunkByScan scan_res_p + chunked_scan_ts_p = chunkByScan scan_ts_p + chunked_scan_pars_c = chunkByScan scan_pars_c + (new_scan, new_scan_res_p, new_scan_ts_p, new_scan_pars_c) = + L.unzip4 + . filter (\(_, _, _, ps) -> any ((`nameIn` deps) . paramName) ps) + $ L.zip4 scan chunked_scan_res_p chunked_scan_ts_p chunked_scan_pars_c + temp_post_lam = eliminateByRes post_lam + deps = freeIn $ lambdaBody temp_post_lam + +-- | Prunes all unused results from the pre-lambda in a ScremaForm +-- (fixed-point). +-- +-- Repeatedly prunes unused scan and map results until no further +-- changes occur. This is necessary because eliminating some results +-- may make other results unused. +-- +-- Example: If a map result is only used by a scan operation, and that +-- scan's results don't contribute to the post-lambda's results, then +-- the first pass removes the scan results, and the second pass can +-- then remove the map result. +-- +-- Returns: A ScremaForm with all transitively unused pre-lambda +-- results eliminated. +prunePreLambdaResults :: (Buildable rep) => ScremaForm rep -> ScremaForm rep +prunePreLambdaResults form = + if form == form' then form' else prunePreLambdaResults form' + where + form' = prunePreLambdaScanResults $ prunePreLambdaMapResults form + +-- | Removes duplicate inputs from a ScremaForm's lambda parameters. +-- +-- When the same input appears multiple times in the input list (with +-- corresponding duplicate lambda parameters), this function: 1. Keeps +-- only one copy of each unique input 2. Creates let-bindings in the +-- lambda body to alias the duplicates +-- +-- Example: If inputs [x, y, x] map to lambda params [a, b, c], the +-- result will have inputs [x, y] with params [a, b], and a +-- let-binding c = a. +-- +-- Arguments: +-- * Input list that corresponds 1:1 with the lambda's parameters +-- * ScremaForm containing the lambda to transform +-- +-- Returns: +-- * Deduplicated input list +-- * Modified ScremaForm with updated lambda (fewer params, +-- additional bindings) +dedupInput :: + (Buildable rep, Ord a) => + -- | Inputs + [a] -> + -- | Screma + ScremaForm rep -> + -- | Deduplicated inputs and new screma + ([a], ScremaForm rep) +dedupInput inp form = (new_inp, form {scremaLambda = new_lam}) + where + lam = scremaLambda form + body = lambdaBody lam + new_body = mkBody (binds <> bodyStms body) $ bodyResult body + new_lam = lam {lambdaParams = new_pars, lambdaBody = new_body} + pars = lambdaParams lam + auxiliary [] = Nothing + auxiliary (x : xs) = Just (x, xs) + pairs = zip inp pars + (new_inp, new_pars) = unzip $ filter (`elem` keep_pairs) pairs + keep_pairs = map fst bind_pairs + bind_pairs = + mapMaybe auxiliary + . L.groupBy ((==) `on` fst) + $ L.sortOn fst pairs + par_bind_pairs = bimap snd (map snd) <$> bind_pairs + binds = foldMap mkBinds par_bind_pairs + mkBinds (par_name, names) = + stmsFromList $ + map + ( \name -> + mkLet [Ident (paramName name) (paramType par_name)] + . BasicOp + . SubExp + . Var + $ paramName par_name + ) + names 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 @@ -26,6 +26,7 @@ KernelResult (..), kernelResultCerts, kernelResultSubExp, + SegPostOp (..), -- ** Generic traversal SegOpMapper (..), @@ -60,9 +61,11 @@ unzip4, zip4, ) +import Data.List qualified as L import Data.Map.Strict qualified as M import Data.Maybe import Futhark.Analysis.Alias qualified as Alias +import Futhark.Analysis.DataDependencies import Futhark.Analysis.Metrics import Futhark.Analysis.PrimExp.Convert import Futhark.Analysis.SymbolTable qualified as ST @@ -145,6 +148,12 @@ } deriving (Eq, Ord, Show) +newtype SegPostOp rep = SegPostOp + { -- | The function applied the the result of some SegBinOp. + segPostOpLambda :: Lambda rep + } + deriving (Eq, Ord, Show) + -- | How many reduction results are produced by these 'SegBinOp's? segBinOpResults :: [SegBinOp rep] -> Int segBinOpResults = sum . map (length . segBinOpNeutral) @@ -360,7 +369,7 @@ | -- | The KernelSpace must always have at least two dimensions, -- implying that the result of a SegRed is always an array. SegRed lvl SegSpace [Type] (KernelBody rep) [SegBinOp rep] - | SegScan lvl SegSpace [Type] (KernelBody rep) [SegBinOp rep] + | SegScan lvl SegSpace [Type] (KernelBody rep) [SegBinOp rep] (SegPostOp rep) | SegHist lvl SegSpace [Type] (KernelBody rep) [HistOp rep] deriving (Eq, Ord, Show) @@ -368,14 +377,14 @@ segLevel :: SegOp lvl rep -> lvl segLevel (SegMap lvl _ _ _) = lvl segLevel (SegRed lvl _ _ _ _) = lvl -segLevel (SegScan lvl _ _ _ _) = lvl +segLevel (SegScan lvl _ _ _ _ _) = lvl segLevel (SegHist lvl _ _ _ _) = lvl -- | The space of a 'SegOp'. segSpace :: SegOp lvl rep -> SegSpace segSpace (SegMap _ lvl _ _) = lvl segSpace (SegRed _ lvl _ _ _) = lvl -segSpace (SegScan _ lvl _ _ _) = lvl +segSpace (SegScan _ lvl _ _ _ _) = lvl segSpace (SegHist _ lvl _ _ _) = lvl -- | The body of a 'SegOp'. @@ -384,7 +393,7 @@ case segop of SegMap _ _ _ body -> body SegRed _ _ _ body _ -> body - SegScan _ _ _ body _ -> body + SegScan _ _ _ body _ _ -> body SegHist _ _ _ body _ -> body segResultShape :: SegSpace -> Type -> KernelResult -> Type @@ -412,18 +421,11 @@ op <- reds let shape = Shape segment_dims <> segBinOpShape op map (`arrayOfShape` shape) (lambdaReturnType $ segBinOpLambda op) -segOpType (SegScan _ space ts kbody scans) = - scan_ts - ++ zipWith - (segResultShape space) - map_ts - (drop (length scan_ts) $ bodyResult kbody) +segOpType (SegScan _ space _ts _kbody _scans post_op) = + flip (foldr (flip arrayOfRow)) segment_dims <$> lam_ts where - map_ts = drop (length scan_ts) ts - scan_ts = do - op <- scans - let shape = Shape (segSpaceDims space) <> segBinOpShape op - map (`arrayOfShape` shape) (lambdaReturnType $ segBinOpLambda op) + lam_ts = lambdaReturnType $ segPostOpLambda post_op + segment_dims = segSpaceDims space segOpType (SegHist _ space _ _ ops) = do op <- ops let shape = Shape segment_dims <> histShape op <> histOpShape op @@ -442,8 +444,10 @@ consumedInBody kbody consumedInOp (SegRed _ _ _ kbody _) = consumedInBody kbody - consumedInOp (SegScan _ _ _ kbody _) = - consumedInBody kbody + consumedInOp (SegScan _ _ _ kbody _ post_op) = + consumed_lam <> consumedInBody kbody + where + consumed_lam = consumedByLambda $ segPostOpLambda post_op consumedInOp (SegHist _ _ _ kbody ops) = namesFromList (concatMap histDest ops) <> consumedInBody kbody @@ -455,19 +459,24 @@ TC.TypeM rep () typeCheckSegOp checkLvl (SegMap lvl space ts kbody) = do checkLvl lvl - checkScanRed space [] ts kbody + checkSegSpace space + TC.binding (scopeOfSegSpace space) $ checkScanRed [] ts kbody typeCheckSegOp checkLvl (SegRed lvl space ts body reds) = do checkLvl lvl - checkScanRed space reds' ts body + checkSegSpace space + TC.binding (scopeOfSegSpace space) $ checkScanRed reds' ts body where reds' = zip3 (map segBinOpLambda reds) (map segBinOpNeutral reds) (map segBinOpShape reds) -typeCheckSegOp checkLvl (SegScan lvl space ts body scans) = do +typeCheckSegOp checkLvl (SegScan lvl space ts body scans post_op) = do checkLvl lvl - checkScanRed space scans' ts body + checkSegSpace space + TC.binding (scopeOfSegSpace space) $ do + checkScanRed scans' ts body + checkSegPostOp post_op scans' ts where scans' = zip3 @@ -523,49 +532,62 @@ where segment_dims = init $ segSpaceDims space +checkSegPostOp :: + (TC.Checkable rep) => + SegPostOp (Aliases rep) -> + [(Lambda (Aliases rep), [SubExp], Shape)] -> + [Type] -> + TC.TypeM rep () +checkSegPostOp (SegPostOp lam) ops ts = do + let (shps, nes) = unzip $ concatMap (\(_, a, shp) -> (shp,) <$> a) ops + nes' <- mapM TC.checkArg nes + let kbody' = (,mempty) <$> drop (length nes') ts + + TC.checkLambda lam $ + map TC.noArgAliases $ + zipWith (\shp (t, e) -> (arrayOfShape t shp, e)) shps nes' <> kbody' + checkScanRed :: (TC.Checkable rep) => - SegSpace -> [(Lambda (Aliases rep), [SubExp], Shape)] -> [Type] -> KernelBody (Aliases rep) -> TC.TypeM rep () -checkScanRed space ops ts kbody = do - checkSegSpace space +checkScanRed ops ts kbody = do mapM_ TC.checkType ts - TC.binding (scopeOfSegSpace space) $ do - ne_ts <- forM ops $ \(lam, nes, shape) -> do - mapM_ (TC.require (Prim int64)) $ shapeDims shape - nes' <- mapM TC.checkArg nes + ne_ts <- forM ops $ \(lam, nes, shape) -> do + mapM_ (TC.require (Prim int64)) $ shapeDims shape + nes' <- mapM TC.checkArg nes - -- Operator type must match the type of neutral elements. - TC.checkLambda lam $ map TC.noArgAliases $ nes' ++ nes' - let nes_t = map TC.argType nes' + -- Operator type must match the type of neutral elements. + TC.checkLambda lam $ map TC.noArgAliases $ nes' ++ nes' + let nes_t = map TC.argType nes' - unless (lambdaReturnType lam == nes_t) $ - TC.bad $ - TC.TypeError "wrong type for operator or neutral elements." + unless (lambdaReturnType lam == nes_t) $ + TC.bad $ + TC.TypeError "wrong type for operator or neutral elements." - pure $ map (`arrayOfShape` shape) nes_t + pure $ map (`arrayOfShape` shape) nes_t - let expecting = concat ne_ts - got = take (length expecting) ts - unless (expecting == got) $ - TC.bad $ - TC.TypeError $ - "Wrong return for body (does not match neutral elements; expected " - <> prettyText expecting - <> "; found " - <> prettyText got - <> ")" + let expecting = concat ne_ts + got = take (length expecting) ts + unless (expecting == got) $ + TC.bad $ + TC.TypeError $ + "Wrong return for body (does not match neutral elements; expected " + <> prettyText expecting + <> "; found " + <> prettyText got + <> ")" - checkKernelBody ts kbody + checkKernelBody ts kbody -- | Like 'Mapper', but just for 'SegOp's. data SegOpMapper lvl frep trep m = SegOpMapper { mapOnSegOpSubExp :: SubExp -> m SubExp, - mapOnSegOpLambda :: Lambda frep -> m (Lambda trep), + mapOnSegBinOpLambda :: Lambda frep -> m (Lambda trep), + mapOnSegPostOpLambda :: Lambda frep -> m (Lambda trep), mapOnSegOpBody :: KernelBody frep -> m (KernelBody trep), mapOnSegOpVName :: VName -> m VName, mapOnSegOpLevel :: lvl -> m lvl @@ -576,7 +598,8 @@ identitySegOpMapper = SegOpMapper { mapOnSegOpSubExp = pure, - mapOnSegOpLambda = pure, + mapOnSegBinOpLambda = pure, + mapOnSegPostOpLambda = pure, mapOnSegOpBody = pure, mapOnSegOpVName = pure, mapOnSegOpLevel = pure @@ -596,10 +619,18 @@ m (SegBinOp trep) mapSegBinOp tv (SegBinOp comm red_op nes shape) = SegBinOp comm - <$> mapOnSegOpLambda tv red_op + <$> mapOnSegBinOpLambda tv red_op <*> mapM (mapOnSegOpSubExp tv) nes <*> (Shape <$> mapM (mapOnSegOpSubExp tv) (shapeDims shape)) +mapSegPostOp :: + (Monad m) => + SegOpMapper lvl frep trep m -> + SegPostOp frep -> + m (SegPostOp trep) +mapSegPostOp tv (SegPostOp lam) = + SegPostOp <$> mapOnSegPostOpLambda tv lam + -- | Apply a 'SegOpMapper' to the given 'SegOp'. mapSegOpM :: (Monad m) => @@ -619,13 +650,14 @@ <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts <*> mapOnSegOpBody tv body <*> mapM (mapSegBinOp tv) reds -mapSegOpM tv (SegScan lvl space ts body scans) = +mapSegOpM tv (SegScan lvl space ts body scans post_op) = SegScan <$> mapOnSegOpLevel tv lvl <*> mapOnSegSpace tv space <*> mapM (mapOnType $ mapOnSegOpSubExp tv) ts <*> mapOnSegOpBody tv body <*> mapM (mapSegBinOp tv) scans + <*> mapSegPostOp tv post_op mapSegOpM tv (SegHist lvl space ts body ops) = SegHist <$> mapOnSegOpLevel tv lvl @@ -641,7 +673,7 @@ <*> mapM (mapOnSegOpVName tv) arrs <*> mapM (mapOnSegOpSubExp tv) nes <*> (Shape <$> mapM (mapOnSegOpSubExp tv) (shapeDims shape)) - <*> mapOnSegOpLambda tv op + <*> mapOnSegBinOpLambda tv op mapOnSegOpType :: (Monad m) => @@ -667,6 +699,14 @@ rephraseBinOp r (SegBinOp comm lam nes shape) = SegBinOp comm <$> rephraseLambda r lam <*> pure nes <*> pure shape +rephrasePostOp :: + (Monad f) => + Rephraser f from rep -> + SegPostOp from -> + f (SegPostOp rep) +rephrasePostOp r (SegPostOp lam) = + SegPostOp <$> rephraseLambda r lam + instance RephraseOp (SegOp lvl) where rephraseInOp r (SegMap lvl space ts body) = SegMap lvl space ts <$> rephraseBody r body @@ -674,10 +714,11 @@ SegRed lvl space ts <$> rephraseBody r body <*> mapM (rephraseBinOp r) reds - rephraseInOp r (SegScan lvl space ts body scans) = + rephraseInOp r (SegScan lvl space ts body scans post_op) = SegScan lvl space ts <$> rephraseBody r body <*> mapM (rephraseBinOp r) scans + <*> rephrasePostOp r post_op rephraseInOp r (SegHist lvl space ts body hists) = SegHist lvl space ts <$> rephraseBody r body @@ -694,7 +735,7 @@ f' scope = f (seg_scope <> scope) mapper = identitySegOpMapper - { mapOnSegOpLambda = traverseLambdaStms f', + { mapOnSegBinOpLambda = traverseLambdaStms f', mapOnSegOpBody = onBody } onBody (Body dec stms res) = @@ -709,7 +750,8 @@ substitute = SegOpMapper { mapOnSegOpSubExp = pure . substituteNames subst, - mapOnSegOpLambda = pure . substituteNames subst, + mapOnSegBinOpLambda = pure . substituteNames subst, + mapOnSegPostOpLambda = pure . substituteNames subst, mapOnSegOpBody = pure . substituteNames subst, mapOnSegOpVName = pure . substituteNames subst, mapOnSegOpLevel = pure . substituteNames subst @@ -719,7 +761,7 @@ rename op = renameBound (M.keys (scopeOfSegSpace (segSpace op))) $ mapSegOpM renamer op where - renamer = SegOpMapper rename rename rename rename rename + renamer = SegOpMapper rename rename rename rename rename rename instance (ASTRep rep, FreeIn lvl) => FreeIn (SegOp lvl rep) where freeIn' e = @@ -731,7 +773,8 @@ free = SegOpMapper { mapOnSegOpSubExp = walk freeIn', - mapOnSegOpLambda = walk freeIn', + mapOnSegBinOpLambda = walk freeIn', + mapOnSegPostOpLambda = walk freeIn', mapOnSegOpBody = walk freeIn', mapOnSegOpVName = walk freeIn', mapOnSegOpLevel = walk freeIn' @@ -744,9 +787,10 @@ inside "SegRed" $ do mapM_ (inside "SegBinOp" . lambdaMetrics . segBinOpLambda) reds bodyMetrics body - opMetrics (SegScan _ _ _ body scans) = + opMetrics (SegScan _ _ _ body scans post_op) = inside "SegScan" $ do mapM_ (inside "SegBinOp" . lambdaMetrics . segBinOpLambda) scans + inside "SegBinOp" $ lambdaMetrics $ segPostOpLambda post_op bodyMetrics body opMetrics (SegHist _ _ _ body ops) = inside "SegHist" $ do @@ -775,6 +819,10 @@ Commutative -> "commutative " Noncommutative -> mempty +instance (PrettyRep rep) => Pretty (SegPostOp rep) where + pretty (SegPostOp lam) = + pretty lam + instance (PrettyRep rep, PP.Pretty lvl) => PP.Pretty (SegOp lvl rep) where pretty (SegMap lvl space ts body) = "segmap" @@ -791,7 +839,7 @@ <+> ppTuple' (map pretty ts) <+> PP.nestedBlock (pretty body) PP.parens (mconcat $ intersperse (PP.comma <> PP.line) $ map pretty reds) - pretty (SegScan lvl space ts body scans) = + pretty (SegScan lvl space ts body scans post_op) = "segscan" <> pretty lvl PP.align (pretty space) @@ -799,6 +847,7 @@ <+> ppTuple' (map pretty ts) <+> PP.nestedBlock (pretty body) PP.parens (mconcat $ intersperse (PP.comma <> PP.line) $ map pretty scans) + pretty post_op pretty (SegHist lvl space ts body ops) = "seghist" <> pretty lvl @@ -828,6 +877,7 @@ SegOpMapper pure (pure . Alias.analyseLambda aliases) + (pure . Alias.analyseLambda aliases) (pure . Alias.analyseBody aliases) pure pure @@ -836,7 +886,7 @@ addOpWisdom = runIdentity . mapSegOpM add where add = - SegOpMapper pure (pure . informLambda) (pure . informBody) pure pure + SegOpMapper pure (pure . informLambda) (pure . informLambda) (pure . informBody) pure pure instance (ASTRep rep) => ST.IndexOp (SegOp lvl rep) where indexOp vtable k (SegMap _ space _ kbody) is = do @@ -904,21 +954,28 @@ <*> Engine.simplify dims_n_tiles <*> Engine.simplify what +segOpBlocker :: + (Engine.SimplifiableRep rep) => + SegSpace -> Engine.SimpleM rep (Engine.BlockPred (Wise rep)) +segOpBlocker space = do + par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers + pure $ + Engine.hasFree bound_here + `Engine.orIf` Engine.isOp + `Engine.orIf` par_blocker + `Engine.orIf` Engine.isConsumed + `Engine.orIf` Engine.isConsuming + `Engine.orIf` Engine.isDeviceMigrated + where + bound_here = namesFromList $ M.keys $ scopeOfSegSpace space + simplifyKernelBody :: (Engine.SimplifiableRep rep, BodyDec rep ~ ()) => SegSpace -> KernelBody (Wise rep) -> Engine.SimpleM rep (KernelBody (Wise rep), Stms (Wise rep)) simplifyKernelBody space (Body _ stms res) = do - par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers - - let blocker = - Engine.hasFree bound_here - `Engine.orIf` Engine.isOp - `Engine.orIf` par_blocker - `Engine.orIf` Engine.isConsumed - `Engine.orIf` Engine.isConsuming - `Engine.orIf` Engine.isDeviceMigrated + blocker <- segOpBlocker space -- Ensure we do not try to use anything that is consumed in the result. (body_res, body_stms, hoisted) <- @@ -933,8 +990,6 @@ pure (res', UT.usages $ freeIn res') pure (mkWiseBody () body_stms body_res, hoisted) - where - bound_here = namesFromList $ M.keys $ scopeOfSegSpace space simplifyLambda :: (Engine.SimplifiableRep rep) => @@ -962,6 +1017,19 @@ nes' <- mapM Engine.simplify nes pure (SegBinOp comm lam' nes' shape', hoisted) +simplifySegPostOp :: + (Engine.SimplifiableRep rep) => + SegSpace -> + SegPostOp (Wise rep) -> + Engine.SimpleM rep (SegPostOp (Wise rep), Stms (Wise rep)) +simplifySegPostOp space (SegPostOp lam) = do + (lam', hoisted) <- + Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $ + simplifyLambda bound_here lam + pure (SegPostOp lam', hoisted) + where + bound_here = namesFromList $ M.keys $ scopeOfSegSpace space + -- | Simplify the given 'SegOp'. simplifySegOp :: ( Engine.SimplifiableRep rep, @@ -990,16 +1058,19 @@ ) where scope = scopeOfSegSpace space -simplifySegOp (SegScan lvl space ts kbody scans) = do +simplifySegOp (SegScan lvl space ts kbody scans post_op) = do (lvl', space', ts') <- Engine.simplify (lvl, space, ts) (scans', scans_hoisted) <- Engine.localVtable (ST.insertScope scope) $ mapAndUnzipM (simplifySegBinOp (segFlat space)) scans (kbody', body_hoisted) <- simplifyKernelBody space kbody + (post_op', post_op_hoisted) <- + Engine.localVtable (ST.insertScope scope) $ + simplifySegPostOp space post_op pure - ( SegScan lvl' space' ts' kbody' scans', - mconcat scans_hoisted <> body_hoisted + ( SegScan lvl' space' ts' kbody' scans' post_op', + mconcat scans_hoisted <> body_hoisted <> post_op_hoisted ) where scope = scopeOfSegSpace space @@ -1064,6 +1135,19 @@ | otherwise = Skip +-- | Only handle Returns cases. +depsOfRes :: Dependencies -> KernelResult -> Names +depsOfRes deps (Returns _ cs se) = depsOf deps se <> depsOfNames deps (freeIn cs) +depsOfRes _ _ = mempty + +kernelBodyDependencies :: (ASTRep rep) => Dependencies -> KernelBody rep -> [Names] +kernelBodyDependencies deps kbody = + let names_in_scope = freeIn kbody + deps' = dataDependencies' deps kbody + in map + (flip namesSubtract names_in_scope . depsOfRes deps') + (bodyResult kbody) + topDownSegOp :: (HasSegOp rep, BuilderOps rep, Buildable rep) => ST.SymbolTable rep -> @@ -1153,6 +1237,71 @@ }, op1_aux ++ op2_aux ) + +-- Remove unused kernelBody result in SegScan. +topDownSegOp _ pat aux op + | SegScan lvl space ts kbody seg_op post_op <- op, + -- Figure out which of the names in 'pat' are used... + Just (new_kbody, new_ts, m_new_post_op) <- + newKbodyPostOp kbody ts seg_op post_op = Simplify $ do + new_post_op <- m_new_post_op + auxing aux + . letBind pat + . Op + . segOp + $ SegScan lvl space new_ts new_kbody seg_op new_post_op + where + newKbodyPostOp kbody ts seg_op post_op = + if null sub_map_res_ts_pars + then Nothing + else Just (new_kbody, new_ts, new_post_op) + where + res = bodyResult kbody + post_lam = segPostOpLambda post_op + pars = lambdaParams post_lam + + mkBind t p r = + mkLet + [Ident (paramName p) t] + (BasicOp $ SubExp $ kernelResultSubExp r) + + new_kbody = kbody {bodyResult = new_res} + new_post_op = do + let sub_res = map (\(r, _, _, _) -> r) sub_map_res_ts_pars + temp_body <- renameBody $ kbody {bodyResult = sub_res} + let new_binds = + stmsFromList + $ zipWith + (\(_, t, p, _) r -> mkBind t p r) + sub_map_res_ts_pars + $ bodyResult temp_body + pure $ + SegPostOp $ + Lambda + { lambdaParams = new_pars, + lambdaBody = + mkBody + (bodyStms temp_body <> new_binds <> bodyStms (lambdaBody post_lam)) + (bodyResult (lambdaBody post_lam)), + lambdaReturnType = lambdaReturnType post_lam + } + + scan_deps = mconcat $ (\(_, _, _, d) -> d) <$> scan_res_ts_pars + deps = kernelBodyDependencies mempty kbody + + isNotReturns (Returns {}) = False + isNotReturns _ = True + + (new_res, new_ts, new_pars, _) = + L.unzip4 $ scan_res_ts_pars <> new_map_res_ts_pars + (new_map_res_ts_pars, sub_map_res_ts_pars) = + L.partition + ( \(r, t, _, d) -> + isNotReturns r || d `namesIntersect` scan_deps || isAcc t + ) + map_res_ts_pars + (scan_res_ts_pars, map_res_ts_pars) = + splitAt (segBinOpResults seg_op) $ L.zip4 res ts pars deps topDownSegOp _ _ _ _ = Skip -- A convenient way of operating on the type and body of a SegOp, @@ -1166,8 +1315,8 @@ ) segOpGuts (SegMap lvl space kts body) = (kts, body, 0, SegMap lvl space) -segOpGuts (SegScan lvl space kts body ops) = - (kts, body, segBinOpResults ops, \t b -> SegScan lvl space t b ops) +segOpGuts (SegScan lvl space kts body ops post_op) = + (kts, body, 0, \t b -> SegScan lvl space t b ops post_op) segOpGuts (SegRed lvl space kts body ops) = (kts, body, segBinOpResults ops, \t b -> SegRed lvl space t b ops) segOpGuts (SegHist lvl space kts body ops) = @@ -1180,11 +1329,22 @@ StmAux (ExpDec rep) -> SegOp (SegOpLevel rep) rep -> Rule rep --- Some SegOp results can be moved outside the SegOp, which can --- simplify further analysis. +-- Because of the postop we have to handle SegScan specially. +bottomUpSegOp (_vtable, used) (Pat kpes) aux (SegScan lvl space kts body ops (SegPostOp postlam)) + | -- Remove dead results from the postlam. + (kpes', res') <- + unzip $ filter keep $ zip kpes (bodyResult (lambdaBody postlam)), + kpes /= kpes' = Simplify $ do + postlam' <- mkLambda (lambdaParams postlam) $ do + addStms $ bodyStms $ lambdaBody postlam + pure res' + addStm $ Let (Pat kpes') aux $ Op $ segOp $ SegScan lvl space kts body ops (SegPostOp postlam') + | otherwise = Skip + where + keep (pe, _) = patElemName pe `UT.used` used bottomUpSegOp (_vtable, used) (Pat kpes) dec segop - -- Remove dead results. This is a bit tricky to do with scan/red - -- results, so we only deal with map results for now. + -- Remove dead results. This is a bit tricky to do with reduction results, so + -- we only deal with map results for now. | (_, kpes', kts', kres') <- unzip4 $ filter keep $ zip4 [0 ..] kpes kts kres, kpes' /= kpes = Simplify $ do kbody' <- localScope (scopeOfSegSpace space) $ mkBodyM kstms kres' diff --git a/src/Futhark/Internalise/ApplyTypeAbbrs.hs b/src/Futhark/Internalise/ApplyTypeAbbrs.hs --- a/src/Futhark/Internalise/ApplyTypeAbbrs.hs +++ b/src/Futhark/Internalise/ApplyTypeAbbrs.hs @@ -54,7 +54,7 @@ -- Remove all type variables and type abbreviations from a value binding. removeTypeVariables :: Types -> ValBind -> ValBind removeTypeVariables types valbind = do - let (ValBind entry _ _ (Info (RetType dims rettype)) _ pats body _ _ _) = valbind + let (ValBind entry _ _ _ (Info (RetType dims rettype)) _ pats body _ _ _) = valbind mapper = ASTMapper { mapOnExp = onExp, 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 @@ -222,6 +222,7 @@ ValBind { valBindEntryPoint = Nothing, valBindName = fname, + valBindNameLoc = mempty, valBindRetDecl = Nothing, valBindRetType = Info rettype_st, valBindTypeParams = dims', @@ -1245,13 +1246,14 @@ -- boolean is true if the function is a 'DynamicFun'. defuncValBind :: ValBind -> DefM (ValBind, Env) -- Eta-expand entry points with a functional return type. -defuncValBind (ValBind entry name _ (Info rettype) tparams params body _ attrs loc) +defuncValBind (ValBind entry name name_loc _ (Info rettype) tparams params body _ attrs loc) | Scalar Arrow {} <- retType rettype = do (body_pats, body', rettype') <- etaExpand (second (const mempty) rettype) body defuncValBind $ ValBind entry name + name_loc Nothing (Info rettype') tparams @@ -1260,7 +1262,7 @@ Nothing attrs loc -defuncValBind valbind@(ValBind _ name retdecl (Info (RetType ret_dims rettype)) tparams params body _ _ _) = do +defuncValBind valbind@(ValBind _ name _ retdecl (Info (RetType ret_dims rettype)) tparams params body _ _ _) = do when (any isTypeParam tparams) $ error $ show name diff --git a/src/Futhark/Internalise/Defunctorise.hs b/src/Futhark/Internalise/Defunctorise.hs --- a/src/Futhark/Internalise/Defunctorise.hs +++ b/src/Futhark/Internalise/Defunctorise.hs @@ -284,14 +284,14 @@ EntryType <$> transformStructType t <*> pure te transformValBind :: ValBind -> TransformM () -transformValBind (ValBind entry name tdecl (Info (RetType dims t)) tparams params e doc attrs loc) = do +transformValBind (ValBind entry name nameloc tdecl (Info (RetType dims t)) tparams params e doc attrs loc) = do entry' <- traverse (traverse transformEntry) entry name' <- transformName name tdecl' <- traverse transformTypeExp tdecl t' <- transformResType t e' <- transformExp e params' <- traverse transformNames params - emit $ ValDec $ ValBind entry' name' tdecl' (Info (RetType dims t')) tparams params' e' doc attrs loc + emit $ ValDec $ ValBind entry' name' nameloc tdecl' (Info (RetType dims t')) tparams params' e' doc attrs loc transformTypeBind :: TypeBind -> TransformM () transformTypeBind (TypeBind name l tparams te (Info (RetType dims t)) doc loc) = do diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs --- a/src/Futhark/Internalise/Exps.hs +++ b/src/Futhark/Internalise/Exps.hs @@ -47,7 +47,7 @@ shiftRetAls d (RetAls pals rals) = RetAls pals $ map (+ d) rals internaliseValBind :: VisibleTypes -> E.ValBind -> InternaliseM () -internaliseValBind types fb@(E.ValBind entry fname _ (Info rettype) tparams params body _ attrs _) = do +internaliseValBind types fb@(E.ValBind entry fname _ _ (Info rettype) tparams params body _ attrs _) = do bindingFParams tparams params $ \shapeparams params' -> do let shapenames = map I.paramName shapeparams all_params = map pure shapeparams ++ concat params' @@ -111,7 +111,7 @@ generateEntryPoint :: VisibleTypes -> E.EntryPoint -> E.ValBind -> InternaliseM () generateEntryPoint types (E.EntryPoint e_params e_rettype) vb = do - let (E.ValBind _ ofname _ (Info rettype) tparams params _ _ attrs _) = vb + let (E.ValBind _ ofname _ _ (Info rettype) tparams params _ _ attrs _) = vb bindingFParams tparams params $ \shapeparams params' -> do let all_params = map pure shapeparams ++ concat params' (entry_rettype, retals) = @@ -1389,9 +1389,11 @@ bs_ts <- mapM lookupType bs' lam' <- internaliseLambdaCoerce lam $ map rowType $ paramType acc_p : bs_ts let w = arraysSize 0 bs_ts - fmap subExpsRes . letValExp' "acc_res" $ - I.Op $ - I.Screma w (I.paramName acc_p : bs') (I.mapSOAC lam') + fmap subExpsRes + . letValExp' "acc_res" + . I.Op + . I.Screma w (I.paramName acc_p : bs') + =<< I.mapSOAC lam' op' <- case op of @@ -1714,9 +1716,10 @@ -- Compare the elements. cmp_lam <- cmpOpLambda $ I.CmpEq (elemType x_t) cmps <- - letExp "cmps" $ - I.Op $ - I.Screma x_num_elems [x_flat, y_flat] (I.mapSOAC cmp_lam) + letExp "cmps" + . I.Op + . I.Screma x_num_elems [x_flat, y_flat] + =<< I.mapSOAC cmp_lam -- Check that all were equal. and_lam <- binOpLambda I.LogAnd I.Bool @@ -1799,7 +1802,7 @@ arr_ts <- mapM lookupType arr' lam' <- internaliseLambdaCoerce lam $ map rowType arr_ts let w = arraysSize 0 arr_ts - letTupExp' desc $ I.Op $ I.Screma w arr' (I.mapSOAC lam') + letTupExp' desc . I.Op . I.Screma w arr' =<< I.mapSOAC lam' handleSOACs [k, lam, arr] "partition" = do k' <- fromIntegral <$> fromInt32 k Just $ \_desc -> do @@ -2157,7 +2160,11 @@ partitionWithSOACS k lam arrs = do arr_ts <- mapM lookupType arrs let w = arraysSize 0 arr_ts - classes_and_increments <- letTupExp "increments" $ I.Op $ I.Screma w arrs (mapSOAC lam) + classes_and_increments <- + letTupExp "increments" + . I.Op + . I.Screma w arrs + =<< mapSOAC lam (classes, increments) <- case classes_and_increments of classes : increments -> pure (classes, take k increments) _ -> error "partitionWithSOACS" diff --git a/src/Futhark/Internalise/LiftLambdas.hs b/src/Futhark/Internalise/LiftLambdas.hs --- a/src/Futhark/Internalise/LiftLambdas.hs +++ b/src/Futhark/Internalise/LiftLambdas.hs @@ -113,6 +113,7 @@ addValBind $ ValBind { valBindName = fname, + valBindNameLoc = mempty, valBindTypeParams = tparams, valBindParams = map mkParam free_ts ++ params, valBindRetDecl = Nothing, 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 @@ -1110,6 +1110,7 @@ ValBind { valBindEntryPoint = Info <$> entry, valBindName = name', + valBindNameLoc = mempty, valBindRetType = Info rettype', valBindRetDecl = Nothing, valBindTypeParams = tparams', @@ -1200,7 +1201,7 @@ PatConstr n (Info tp) ps loc -> PatConstr n (Info $ f tp) ps loc toPolyBinding :: ValBind -> PolyBinding -toPolyBinding (ValBind entry name _ (Info rettype) tparams params body _ attrs loc) = +toPolyBinding (ValBind entry name _ _ (Info rettype) tparams params body _ attrs loc) = PolyBinding (unInfo <$> entry, name, tparams, params, rettype, body, attrs, loc) transformValBind :: ValBind -> MonoM Env diff --git a/src/Futhark/LSP/CodeAction.hs b/src/Futhark/LSP/CodeAction.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/LSP/CodeAction.hs @@ -0,0 +1,141 @@ +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE QuasiQuotes #-} + +module Futhark.LSP.CodeAction (getCodeActions) where + +import Data.Function ((&)) +import Data.Loc (Pos (Pos), locOf) +import Data.Map qualified as M +import Data.Maybe (fromMaybe, mapMaybe) +import Data.Set qualified as S +import Data.Text (Text) +import Data.Text qualified as T +import Futhark.LSP.State (State, getStaleMapping) +import Futhark.LSP.Tool (bindingsInState, filterByLspRange) +import Futhark.LSP.TypeAscription (TypeAscription (TypeAscLet, TypeAscParam, TypeAscReturn, TypeAscType), missingAscriptions, missingTypeParameters) +import Futhark.Util.Pretty (prettyText) +import Language.Futhark.Core (VName, baseText) +import Language.Futhark.FreeVars (freeInType, fvVars) +import Language.Futhark.Prop (typeVars) +import Language.Futhark.Syntax (RetTypeBase (RetType), TypeParamBase) +import Language.LSP.Protocol.Types (CodeAction (..), CodeActionKind (CodeActionKind_Custom), Command, Position (Position), Range (Range), TextEdit (..), Uri, WorkspaceEdit (..), type (|?) (InR)) +import NeatInterpolation qualified as NI + +getCodeActions :: Uri -> Range -> State -> FilePath -> [Command |? CodeAction] +getCodeActions file_uri range state filepath = fromMaybe [] $ do + allBindings <- bindingsInState state + let inferredTypes = foldMap missingTypeParameters allBindings + let mapping = getStaleMapping state filepath + bindings <- + filterByLspRange mapping range filepath (locOf . snd) $ + M.assocs allBindings + + pure $ + concatMap + ( mapMaybe (uncurry $ codeActions inferredTypes) + . (\(n, b) -> (n,) <$> missingAscriptions b) + ) + bindings + & map InR + where + codeActions :: + M.Map VName (Maybe Pos, TypeParamBase VName) -> + VName -> + TypeAscription -> + Maybe CodeAction + codeActions types name (TypeAscLet typ (Pos _ line column _)) = + Just + $ mkAction + [NI.text|Insert type: `$nameText$typeText`|] + $ (typeText, line, column) : typeVarEdits types (typesAndVars typ) + where + typeText = ": " <> prettyText typ + nameText = baseText name + codeActions types name (TypeAscParam openPos typ pos) = + Just + $ mkAction + [NI.text|Insert type: `($nameText: $typeText)`|] + $ posEdit openPos "(" + : posEdit pos (": " <> typeText <> ")") + : typeVarEdits types (typesAndVars typ) + where + posEdit (Pos _ line column _) t = (t, line, column) + typeText = prettyText typ + nameText = baseText name + codeActions types name (TypeAscReturn typ (Pos _ line column _)) = + Just + $ mkAction + [NI.text|Insert return type: `$nameText$typeText`|] + $ (typeText, line, column) : typeVarEdits types retTypeVars + where + typeText = " : " <> prettyText typ + nameText = baseText name + retTypeVars = + S.fromList dims `S.union` typesAndVars innerType + where + RetType dims innerType = typ + codeActions _ _ (TypeAscType _ _) = Nothing + + typesAndVars t = typeVars t `S.union` fvVars (freeInType t) + + -- Text edit information for referenced types that are not yet + -- syntactically present + typeVarEdits :: + M.Map VName (Maybe Pos, TypeParamBase VName) -> + S.Set VName -> + [(Text, Int, Int)] + typeVarEdits availableTypes requiredTypes = + S.toList requiredTypes + & mapMaybe (availableTypes M.!?) + & mapMaybe typeTriple + where + typeTriple :: (Maybe Pos, TypeParamBase VName) -> Maybe (Text, Int, Int) + typeTriple (pos, typ) = do + Pos _ line col _ <- pos + pure (" " <> prettyText typ, line, col) + + mkAction title edits = + CodeAction + { _isPreferred = Nothing, + _disabled = Nothing, + _diagnostics = Nothing, + _edit = + Just $ + WorkspaceEdit + { _documentChanges = Nothing, + _changeAnnotations = Nothing, + _changes = + Just $ + M.singleton + file_uri + (map mkEdit edits) + }, + _title = title, + _command = Nothing, + _data_ = Nothing, + _kind = Just $ CodeActionKind_Custom "TypeAscription" + } + + mkEdit (text, line, column) = + TextEdit + { _range = + let toLsp = fromIntegral . pred + insertPos = Position (toLsp line) (toLsp column) + in Range insertPos insertPos, + _newText = normalizeNames text + } + where + -- replace subscript characters + normalizeNames = T.map $ \case + '\8320' -> '0' + '\8321' -> '1' + '\8322' -> '2' + '\8323' -> '3' + '\8324' -> '4' + '\8325' -> '5' + '\8326' -> '6' + '\8327' -> '7' + '\8328' -> '8' + '\8329' -> '9' + x -> x diff --git a/src/Futhark/LSP/CodeLens.hs b/src/Futhark/LSP/CodeLens.hs --- a/src/Futhark/LSP/CodeLens.hs +++ b/src/Futhark/LSP/CodeLens.hs @@ -9,7 +9,7 @@ import Control.Exception (AllocationLimitExceeded (AllocationLimitExceeded), handle) import Control.Lens ((^.)) import Control.Monad (void) -import Control.Monad.Except (Except, ExceptT, MonadError (throwError), runExceptT) +import Control.Monad.Except (Except, MonadError (throwError), runExceptT) import Control.Monad.Trans (MonadIO (liftIO), lift) import Data.Aeson (FromJSON (parseJSON)) import Data.Aeson qualified as Aeson @@ -28,11 +28,11 @@ import Futhark.Compiler.Program (VFS) import Futhark.Eval (Evaluation (abort), InterpreterConfig (InterpreterConfig), newFutharkiState, runEvalRecordRef, runExpr) import Futhark.LSP.CommandType qualified as CommandType -import Futhark.LSP.Tool (transformVFS) +import Futhark.LSP.Tool (Execute, transformVFS) import Futhark.Util (showText) import Futhark.Util.Pretty (docText, plural, pretty, vcat) import Language.LSP.Protocol.Message (SMethod (SMethod_WorkspaceApplyEdit)) -import Language.LSP.Protocol.Types (ApplyWorkspaceEditParams (..), CodeLens (..), Command (..), ErrorCodes (ErrorCodes_InvalidParams), LSPErrorCodes, Position (..), Range (..), TextEdit (..), UInt, Uri, WorkspaceEdit (..), fromNormalizedFilePath, toNormalizedUri, uriToNormalizedFilePath, type (|?) (..)) +import Language.LSP.Protocol.Types (ApplyWorkspaceEditParams (..), CodeLens (..), Command (..), ErrorCodes (ErrorCodes_InvalidParams), Position (..), Range (..), TextEdit (..), UInt, Uri, WorkspaceEdit (..), fromNormalizedFilePath, toNormalizedUri, uriToNormalizedFilePath, type (|?) (..)) import Language.LSP.Server (LspT, getVirtualFile, getVirtualFiles, sendRequest) import Language.LSP.VFS (file_text) import Prettyprinter (Doc) @@ -143,8 +143,6 @@ _command = showText CommandType.CodeLens } } - -type Execute a = ExceptT (Text, LSPErrorCodes |? ErrorCodes) (LspT () IO) a -- | Decode the arguments execute :: Maybe [Aeson.Value] -> Execute () diff --git a/src/Futhark/LSP/Compile.hs b/src/Futhark/LSP/Compile.hs --- a/src/Futhark/LSP/Compile.hs +++ b/src/Futhark/LSP/Compile.hs @@ -4,14 +4,15 @@ -- the old state around. module Futhark.LSP.Compile (tryTakeStateFromIORef, tryReCompile) where -import Colog.Core (logStringStderr, (<&)) +import Colog.Core (Severity (Debug, Warning), (<&)) import Control.Monad.IO.Class (MonadIO (liftIO)) import Data.IORef (IORef, readIORef, writeIORef) import Data.Maybe (fromMaybe) import Futhark.Compiler.Program (LoadedProg, lpFilePaths, lpWarnings, noLoadedProg, reloadProg) import Futhark.LSP.Diagnostic (diagnosticSource, maxDiagnostic, publishErrorDiagnostics, publishWarningDiagnostics) import Futhark.LSP.State (State (..), emptyState, updateStaleContent, updateStaleMapping) -import Futhark.LSP.Tool (computeMapping, transformVFS) +import Futhark.LSP.Tool (computeMapping, logWithSeverity, transformVFS) +import Futhark.Util (showText) import Language.Futhark.Warnings (listWarnings) import Language.LSP.Protocol.Types ( filePathToUri, @@ -35,8 +36,8 @@ state <- case file_path of Just file_path' | file_path' `notElem` files -> do - logStringStderr <& ("File not part of program: " <> show file_path') - logStringStderr <& ("Program contains: " <> show files) + logWithSeverity Warning <& ("File not part of program: " <> showText file_path') + logWithSeverity Warning <& ("Program contains: " <> showText files) tryCompile old_state file_path noLoadedProg _ -> pure old_state liftIO $ writeIORef state_mvar state @@ -45,18 +46,18 @@ -- | Try to (re)-compile, replace old state if successful. tryReCompile :: IORef State -> Maybe FilePath -> LspT () IO () tryReCompile state_mvar file_path = do - logStringStderr <& "(Re)-compiling ..." + logWithSeverity Debug <& "(Re)-compiling ..." old_state <- liftIO $ readIORef state_mvar let loaded_prog = getLoadedProg old_state new_state <- tryCompile old_state file_path loaded_prog case stateProgram new_state of Nothing -> do - logStringStderr <& "Failed to (re)-compile, using old state or Nothing" - logStringStderr <& "Computing PositionMapping for: " <> show file_path + logWithSeverity Debug <& "Failed to (re)-compile, using old state or Nothing" + logWithSeverity Debug <& "Computing PositionMapping for: " <> showText file_path mapping <- computeMapping old_state file_path liftIO $ writeIORef state_mvar $ updateStaleMapping file_path mapping old_state Just _ -> do - logStringStderr <& "(Re)-compile successful" + logWithSeverity Debug <& "(Re)-compile successful" liftIO $ writeIORef state_mvar new_state -- | Try to compile, publish diagnostics on warnings and errors, return newly compiled state. @@ -64,12 +65,13 @@ tryCompile :: State -> Maybe FilePath -> LoadedProg -> LspT () IO State tryCompile _ Nothing _ = pure emptyState tryCompile state (Just path) old_loaded_prog = do - logStringStderr <& "Reloading program from " <> show path + logWithSeverity Debug <& "Reloading program from " <> showText path vfs <- getVirtualFiles res <- liftIO $ reloadProg old_loaded_prog [path] (transformVFS vfs) -- NOTE: vfs only keeps track of current opened files flushDiagnosticsBySource maxDiagnostic diagnosticSource case res of Right new_loaded_prog -> do + logWithSeverity Debug <& "Successfully compiled, publishing warnings" publishWarningDiagnostics $ listWarnings $ lpWarnings new_loaded_prog maybe_virtual_file <- getVirtualFile $ toNormalizedUri $ filePathToUri path case maybe_virtual_file of @@ -80,7 +82,7 @@ -- But still might need an update on re-compile logic, don't discard all state afterwards, -- try to compile from root file, if there is a depencency relatetion, improve performance and provide more dignostic. Left prog_error -> do - logStringStderr <& "Compilation failed, publishing diagnostics" + logWithSeverity Debug <& "Compilation failed, publishing diagnostics" publishErrorDiagnostics prog_error pure emptyState diff --git a/src/Futhark/LSP/Diagnostic.hs b/src/Futhark/LSP/Diagnostic.hs --- a/src/Futhark/LSP/Diagnostic.hs +++ b/src/Futhark/LSP/Diagnostic.hs @@ -8,14 +8,15 @@ ) where -import Colog.Core (logStringStderr, (<&)) +import Colog.Core (Severity (Debug), (<&)) import Control.Lens ((^.)) import Data.Foldable (for_) import Data.List.NonEmpty qualified as NE import Data.Map qualified as M import Data.Text qualified as T import Futhark.Compiler.Program (ProgError (..)) -import Futhark.LSP.Tool (posToUri, rangeFromLoc) +import Futhark.LSP.Tool (logWithSeverity, posToUri, rangeFromLoc) +import Futhark.Util (showText) import Futhark.Util.Loc (Loc (..)) import Futhark.Util.Pretty (Doc, docText) import Language.LSP.Diagnostics (partitionBySource) @@ -40,8 +41,8 @@ publish :: [(Uri, [Diagnostic])] -> LspT () IO () publish uri_diags_map = for_ uri_diags_map $ \(uri, diags) -> do doc <- getVersionedTextDoc $ TextDocumentIdentifier uri - logStringStderr - <& ("Publishing diagnostics for " ++ show uri ++ " Version: " ++ show (doc ^. version)) + logWithSeverity Debug + <& ("Publishing diagnostics for " <> showText uri <> " Version: " <> showText (doc ^. version)) publishDiagnostics maxDiagnostic (toNormalizedUri uri) diff --git a/src/Futhark/LSP/Handlers.hs b/src/Futhark/LSP/Handlers.hs --- a/src/Futhark/LSP/Handlers.hs +++ b/src/Futhark/LSP/Handlers.hs @@ -5,7 +5,7 @@ -- | The handlers exposed by the language server. module Futhark.LSP.Handlers (handlers) where -import Colog.Core (logStringStderr, (<&)) +import Colog.Core (Severity (Debug, Info), (<&)) import Control.Lens ((^.)) import Control.Monad.Except (ExceptT, MonadError (throwError), liftEither, throwError) import Control.Monad.Trans (lift) @@ -21,16 +21,18 @@ import Data.Text.Mixed.Rope qualified as R import Data.Vector qualified as V import Futhark.Fmt.Printer (fmtToText) +import Futhark.LSP.CodeAction (getCodeActions) import Futhark.LSP.CodeLens qualified as CodeLens import Futhark.LSP.CommandType (CommandType (CodeLens)) import Futhark.LSP.Compile (tryReCompile, tryTakeStateFromIORef) +import Futhark.LSP.InlayHint (getInlayHints) import Futhark.LSP.State (State (..)) -import Futhark.LSP.Tool (findDefinitionRange, getHoverInfoFromState) +import Futhark.LSP.Tool (findDefinitionRange, getHoverInfoFromState, logWithSeverity) import Futhark.Util (showText) import Futhark.Util.Pretty (prettyText) import Language.Futhark.Core (locText) import Language.Futhark.Parser.Monad (SyntaxError (SyntaxError)) -import Language.LSP.Protocol.Lens (HasTextDocument (textDocument), HasUri (uri), arguments, command, line, params, range, start) +import Language.LSP.Protocol.Lens (arguments, command, line, params, range, start, textDocument, uri) import Language.LSP.Protocol.Message ( Method (..), SMethod (..), @@ -64,7 +66,7 @@ onInitializeHandler :: Handlers (LspM ()) onInitializeHandler = notificationHandler SMethod_Initialized $ \_msg -> - logStringStderr <& "Initialized" + logWithSeverity Info <& "Initialized" onHoverHandler :: IORef State -> Handlers (LspM ()) onHoverHandler state_mvar = @@ -72,14 +74,14 @@ let TRequestMessage _ _ _ (HoverParams doc pos _workDone) = req Position l c = pos file_path = uriToFilePath $ doc ^. uri - logStringStderr <& ("Got hover request: " <> show (file_path, pos)) + logWithSeverity Debug <& "Got hover request: " <> showText (file_path, pos) state <- tryTakeStateFromIORef state_mvar file_path responder $ Right $ maybe (InR Null) InL $ getHoverInfoFromState state file_path (fromEnum l + 1) (fromEnum c + 1) onDocumentFocusHandler :: IORef State -> Handlers (LspM ()) onDocumentFocusHandler state_mvar = notificationHandler (SMethod_CustomMethod (Proxy @"custom/onFocusTextDocument")) $ \msg -> do - logStringStderr <& "Got custom request: onFocusTextDocument" + logWithSeverity Debug <& "Got custom request: onFocusTextDocument" let TNotificationMessage _ _ (Array vector_param) = msg String focused_uri = V.head vector_param -- only one parameter passed from the client tryReCompile state_mvar (uriToFilePath (Uri focused_uri)) @@ -90,7 +92,8 @@ let TRequestMessage _ _ _ (DefinitionParams doc pos _workDone _partial) = req Position l c = pos file_path = uriToFilePath $ doc ^. uri - logStringStderr <& ("Got goto definition: " <> show (file_path, pos)) + logWithSeverity Debug + <& ("Got goto definition: " <> showText (file_path, pos)) state <- tryTakeStateFromIORef state_mvar file_path case findDefinitionRange state file_path (fromEnum l + 1) (fromEnum c + 1) of Nothing -> responder $ Right $ InR $ InR Null @@ -101,7 +104,7 @@ notificationHandler SMethod_TextDocumentDidSave $ \msg -> do let TNotificationMessage _ _ (DidSaveTextDocumentParams doc _text) = msg file_path = uriToFilePath $ doc ^. uri - logStringStderr <& ("Saved document: " ++ show doc) + logWithSeverity Debug <& ("Saved document: " <> showText doc) tryReCompile state_mvar file_path onDocumentChangeHandler :: IORef State -> Handlers (LspM ()) @@ -124,27 +127,26 @@ onWorkspaceDidChangeConfiguration :: IORef State -> Handlers (LspM ()) onWorkspaceDidChangeConfiguration _state_mvar = notificationHandler SMethod_WorkspaceDidChangeConfiguration $ \_ -> - logStringStderr <& "WorkspaceDidChangeConfiguration" + logWithSeverity Debug <& "WorkspaceDidChangeConfiguration" onDocumentFormattingHandler :: Handlers (LspM ()) onDocumentFormattingHandler = - requestHandler SMethod_TextDocumentFormatting $ \message report -> + requestHandler SMethod_TextDocumentFormatting $ \message report -> do let TRequestMessage _ _ _ formattingParams = message DocumentFormattingParams _progressToken textDoc _opts = formattingParams fileUri = textDoc ^. uri - in do - logStringStderr <& ("Formatting: " ++ show (textDoc ^. uri)) - result <- runExceptT $ do - virtualFile <- getVirtualFile' fileUri - let fileText = R.toText $ virtualFile ^. file_text - formattedText <- fmtToText' (show fileUri) fileText - pure $ - if formattedText == fileText - then InR Null - else InL [fullTextEdit formattedText] + logWithSeverity Debug <& "Formatting: " <> showText (textDoc ^. uri) + result <- runExceptT $ do + virtualFile <- getVirtualFile' fileUri + let fileText = R.toText $ virtualFile ^. file_text + formattedText <- fmtToText' (show fileUri) fileText + pure $ + if formattedText == fileText + then InR Null + else InL [fullTextEdit formattedText] - logStringStderr <& show result - report result + logWithSeverity Debug <& showText result + report result where fullTextEdit newText = TextEdit @@ -192,12 +194,12 @@ onDocumentCodeLenses :: Handlers (LspM ()) onDocumentCodeLenses = - requestHandler SMethod_TextDocumentCodeLens $ \request respond -> + requestHandler SMethod_TextDocumentCodeLens $ \request respond -> do let textDocUri = request ^. params . textDocument . uri - in do - logStringStderr <& ("textDocument/CodeLens for " ++ show textDocUri) - eitherLenses <- CodeLens.evalLensesFor textDocUri - respond $ bimap failure success eitherLenses + logWithSeverity Debug + <& ("textDocument/CodeLens for " <> showText textDocUri) + eitherLenses <- CodeLens.evalLensesFor textDocUri + respond $ bimap failure success eitherLenses where success :: [CodeLens] -> [CodeLens] |? Null success = InL @@ -211,13 +213,13 @@ onDocumentCodeLensResolve :: Handlers (LspM ()) onDocumentCodeLensResolve = - requestHandler SMethod_CodeLensResolve $ \request respond -> + requestHandler SMethod_CodeLensResolve $ \request respond -> do let codeLens = request ^. params codeLensLine = codeLens ^. (range . start . line) - in do - logStringStderr <& ("Resolving code lens on line " ++ show codeLensLine) - let result = runExcept $ CodeLens.resolve codeLens - respond . first failure $ result + logWithSeverity Debug + <& ("Resolving code lens on line " <> showText codeLensLine) + let result = runExcept $ CodeLens.resolve codeLens + respond . first failure $ result where failure :: Text -> TResponseError Method_CodeLensResolve failure text = @@ -242,13 +244,12 @@ onWorkspaceExecuteCommandHandler :: Handlers (LspM ()) onWorkspaceExecuteCommandHandler = - requestHandler SMethod_WorkspaceExecuteCommand $ \request respond -> + requestHandler SMethod_WorkspaceExecuteCommand $ \request respond -> do let parameters = request ^. params - in do - let commandName = parameters ^. command - let commandArgs = parameters ^. arguments - result <- runExceptT $ executeCommand commandName commandArgs - respond $ bimap (uncurry failure) (const $ InR Null) result + commandName = parameters ^. command + commandArgs = parameters ^. arguments + result <- runExceptT $ executeCommand commandName commandArgs + respond $ bimap (uncurry failure) (const $ InR Null) result where failure message err = TResponseError @@ -257,6 +258,33 @@ _code = err } +onTextDocumentInlayHint :: IORef State -> Handlers (LspM ()) +onTextDocumentInlayHint state_ref = + requestHandler SMethod_TextDocumentInlayHint $ \request respond -> do + let parameters = request ^. params + filepath = uriToFilePath $ parameters ^. (textDocument . uri) + textRange = parameters ^. range + logWithSeverity Debug <& "Inlay hints request for range: " <> showText textRange + + state <- tryTakeStateFromIORef state_ref filepath + let result = maybe [] (getInlayHints textRange state) filepath + + respond . Right $ InL result + +onTextDocumentCodeAction :: IORef State -> Handlers (LspM ()) +onTextDocumentCodeAction state_ref = + requestHandler SMethod_TextDocumentCodeAction $ \request respond -> do + let parameters = request ^. params + file_uri = parameters ^. (textDocument . uri) + filepath = uriToFilePath $ parameters ^. (textDocument . uri) + textRange = parameters ^. range + logWithSeverity Debug <& "Code action request for range: " <> showText textRange + + state <- tryTakeStateFromIORef state_ref filepath + let result = maybe [] (getCodeActions file_uri textRange state) filepath + + respond $ Right $ InL result + -- | Given an 'IORef' tracking the state, produce a set of handlers. -- When we want to add more features to the language server, this is -- the thing to change. @@ -269,6 +297,8 @@ onDocumentCodeLenses, onDocumentCodeLensResolve, onDocumentFormattingHandler, + onTextDocumentInlayHint state_mvar, + onTextDocumentCodeAction state_mvar, onDocumentSaveHandler state_mvar, onDocumentChangeHandler state_mvar, onDocumentFocusHandler state_mvar, diff --git a/src/Futhark/LSP/InlayHint.hs b/src/Futhark/LSP/InlayHint.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/LSP/InlayHint.hs @@ -0,0 +1,65 @@ +{-# LANGUAGE ExplicitNamespaces #-} + +module Futhark.LSP.InlayHint (getInlayHints) where + +import Data.Function ((&)) +import Data.Loc (Pos (Pos)) +import Data.Text (Text) +import Futhark.LSP.State (State) +import Futhark.LSP.Tool (bindingsInRange) +import Futhark.LSP.TypeAscription (TypeAscription (..), missingAscriptions) +import Futhark.Util.Pretty (prettyText) +import Language.LSP.Protocol.Types + ( InlayHint (..), + InlayHintKind (InlayHintKind_Type), + Position (..), + Range, + type (|?) (InL), + ) + +getInlayHints :: Range -> State -> FilePath -> [InlayHint] +getInlayHints range state filepath = + maybe + [] + (map snd) + (bindingsInRange range state filepath) + & concatMap missingAscriptions + & concatMap inlayHint + where + inlayHint :: TypeAscription -> [InlayHint] + inlayHint (TypeAscLet typName pos) = + [bareHint (": " <> prettyText typName) pos] + inlayHint (TypeAscParam s tname pos) = + [startHint s, bareHint (": " <> prettyText tname <> ")") pos] + inlayHint (TypeAscReturn typName pos) = + [bareHint (" : " <> prettyText typName) pos] + inlayHint (TypeAscType typName pos) = + [bareHint (" " <> prettyText typName) pos] + + startHint :: Pos -> InlayHint + startHint (Pos _ l c _) = + InlayHint + { _position = + Position (fromIntegral l - 1) (fromIntegral c - 1), + _label = InL "(", + _kind = Just InlayHintKind_Type, + _textEdits = Nothing, + _tooltip = Nothing, + _paddingLeft = Nothing, + _paddingRight = Nothing, + _data_ = Nothing + } + + bareHint :: Text -> Pos -> InlayHint + bareHint tname (Pos _ l c _) = + InlayHint + { _position = + Position (fromIntegral l - 1) (fromIntegral c - 1), + _label = InL tname, + _kind = Just InlayHintKind_Type, + _textEdits = Nothing, + _tooltip = Nothing, + _paddingLeft = Nothing, + _paddingRight = Nothing, + _data_ = Nothing + } diff --git a/src/Futhark/LSP/PositionMapping.hs b/src/Futhark/LSP/PositionMapping.hs --- a/src/Futhark/LSP/PositionMapping.hs +++ b/src/Futhark/LSP/PositionMapping.hs @@ -56,7 +56,7 @@ toStalePos (Just (PositionMapping _ current_to_stale)) pos = if l > Prelude.length current_to_stale then Nothing - else Just $ Pos file (V.unsafeIndex current_to_stale (l - 1) + 1) c o + else Just $ Pos file ((current_to_stale V.! (l - 1)) + 1) c o where Pos file l c o = pos toStalePos Nothing pos = Just pos @@ -66,7 +66,7 @@ toCurrentPos (Just (PositionMapping stale_to_current _)) pos = if l > Prelude.length stale_to_current then Nothing - else Just $ Pos file (V.unsafeIndex stale_to_current (l - 1) + 1) c o + else Just $ Pos file ((stale_to_current V.! (l - 1)) + 1) c o where Pos file l c o = pos toCurrentPos Nothing pos = Just pos diff --git a/src/Futhark/LSP/Tool.hs b/src/Futhark/LSP/Tool.hs --- a/src/Futhark/LSP/Tool.hs +++ b/src/Futhark/LSP/Tool.hs @@ -1,4 +1,6 @@ +{-# LANGUAGE ExplicitNamespaces #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PartialTypeSignatures #-} -- | Generally useful definition used in various places in the -- language server implementation. @@ -9,11 +11,22 @@ posToUri, computeMapping, transformVFS, + logWithSeverity, + Execute, + bindingsInRange, + bindingsInState, + filterByLspRange, ) where +import Colog.Core (LogAction, Severity, WithSeverity (WithSeverity)) import Control.Lens.Getter ((^.)) +import Control.Monad.Except (ExceptT) +import Data.Foldable qualified as F +import Data.Function ((&)) +import Data.Functor.Contravariant (contramap) import Data.Map qualified as M +import Data.Text (Text) import Data.Text qualified as T import Futhark.Compiler.Program (lpImports) import Futhark.LSP.PositionMapping @@ -23,20 +36,42 @@ toStalePos, ) import Futhark.LSP.State (State (..), getStaleContent, getStaleMapping) -import Futhark.Util.Loc (Loc (Loc, NoLoc), Pos (Pos)) +import Futhark.Util.Loc (Loc (Loc, NoLoc), Pos (Pos), contains, locOf) import Futhark.Util.Pretty (prettyText) -import Language.Futhark.Core (isBuiltinLoc) +import Language.Futhark.Core (VName, isBuiltinLoc) import Language.Futhark.Query ( AtPos (AtName), BoundTo (..), + allBindings, atPos, boundLoc, + termBindingType, ) +import Language.LSP.Logging (logToLogMessage) import Language.LSP.Protocol.Types -import Language.LSP.Server (LspM, getVirtualFile) + ( ErrorCodes, + Hover (Hover), + LSPErrorCodes, + Location (Location), + MarkupContent (MarkupContent), + MarkupKind (MarkupKind_PlainText), + Position (Position), + Range (Range), + UInt, + Uri, + filePathToUri, + fromNormalizedFilePath, + toNormalizedUri, + uriToNormalizedFilePath, + type (|?) (InL), + ) +import Language.LSP.Server (LspM, LspT, MonadLsp, getVirtualFile) import Language.LSP.VFS (VFS, VirtualFile, vfsMap, virtualFileText, virtualFileVersion) import Language.LSP.VFS qualified as VFS +-- | Request Handler code usually runs in this monad +type Execute a = ExceptT (Text, LSPErrorCodes |? ErrorCodes) (LspT () IO) a + -- | Retrieve hover info for the definition referenced at the given -- file at the given line and column number (the two 'Int's). getHoverInfoFromState :: State -> Maybe FilePath -> Int -> Int -> Maybe Hover @@ -44,7 +79,7 @@ AtName _ (Just def) loc <- queryAtPos state $ Pos path l c 0 let msg = case def of - BoundTerm t _ -> prettyText t + BoundTerm term _ -> prettyText $ termBindingType term BoundModule {} -> "module" BoundModuleType {} -> "module type" BoundType {} -> "type" @@ -74,31 +109,31 @@ loaded_prog <- stateProgram state stale_pos <- toStalePos mapping pos query_result <- atPos (lpImports loaded_prog) stale_pos - updateAtPos mapping query_result - where - -- Update the 'AtPos' with the current mapping. - updateAtPos :: Maybe PositionMapping -> AtPos -> Maybe AtPos - updateAtPos mapping (AtName qn (Just def) loc) = do - let def_loc = boundLoc def - Loc (Pos def_file _ _ _) _ = def_loc - Pos current_file _ _ _ = pos - current_loc <- toCurrentLoc mapping loc - if def_file == current_file - then do - current_def_loc <- toCurrentLoc mapping def_loc - Just $ AtName qn (Just (updateBoundLoc def current_def_loc)) current_loc - else do - -- Defined in another file, get the corresponding PositionMapping. - let def_mapping = getStaleMapping state def_file - current_def_loc <- toCurrentLoc def_mapping def_loc - Just $ AtName qn (Just (updateBoundLoc def current_def_loc)) current_loc - updateAtPos _ _ = Nothing + updateAtPos mapping query_result pos state +-- Update the 'AtPos' with the current mapping. +updateAtPos :: Maybe PositionMapping -> AtPos -> Pos -> State -> Maybe AtPos +updateAtPos mapping (AtName qn (Just def) loc) pos state = do + let def_loc = boundLoc def + Loc (Pos def_file _ _ _) _ = def_loc + Pos current_file _ _ _ = pos + current_loc <- toCurrentLoc mapping loc + if def_file == current_file + then do + current_def_loc <- toCurrentLoc mapping def_loc + Just $ AtName qn (Just (updateBoundLoc def current_def_loc)) current_loc + else do + -- Defined in another file, get the corresponding PositionMapping. + let def_mapping = getStaleMapping state def_file + current_def_loc <- toCurrentLoc def_mapping def_loc + Just $ AtName qn (Just (updateBoundLoc def current_def_loc)) current_loc + where updateBoundLoc :: BoundTo -> Loc -> BoundTo updateBoundLoc (BoundTerm t _loc) current_loc = BoundTerm t current_loc updateBoundLoc (BoundModule _loc) current_loc = BoundModule current_loc updateBoundLoc (BoundModuleType _loc) current_loc = BoundModuleType current_loc updateBoundLoc (BoundType _loc) current_loc = BoundType current_loc +updateAtPos _ _ _ _ = Nothing -- | Entry point for computing PositionMapping. computeMapping :: State -> Maybe FilePath -> LspM () (Maybe PositionMapping) @@ -150,3 +185,41 @@ keepOpenFile = \case VFS.Open file -> Just file VFS.Closed _ -> Nothing + +logWithSeverity :: (MonadLsp c m) => Severity -> LogAction m Text +logWithSeverity severity = contramap (`WithSeverity` severity) logToLogMessage + +bindingsInState :: State -> Maybe (M.Map VName BoundTo) +bindingsInState state = allBindings . lpImports <$> stateProgram state + +filterByLspRange :: + (Foldable t) => + Maybe PositionMapping -> + Range -> + FilePath -> + (a -> Loc) -> + t a -> + Maybe [a] +filterByLspRange mapping range filepath f t = do + let (Range (Position l1 c1) (Position l2 c2)) = range + posStart <- toStalePos mapping $ mkPos l1 c1 + posEnd <- toStalePos mapping $ mkPos l2 c2 + pure $ + F.toList t + & filter (isInRange (Loc posStart posEnd) . f) + where + isInRange locRange pos = case pos of + NoLoc -> False + Loc start _ -> locRange `contains` start + -- increment by one: Pos counts from one onwards, LSP starts at zero + mkPos :: UInt -> UInt -> Pos + mkPos l c = Pos filepath (1 + fromIntegral l) (1 + fromIntegral c) (-1) + +bindingsInRange :: Range -> State -> FilePath -> Maybe [(VName, BoundTo)] +bindingsInRange range state filepath = do + let mapping = getStaleMapping state filepath + bindings <- bindingsInState state + + bindings + & M.assocs + & filterByLspRange mapping range filepath (locOf . snd) diff --git a/src/Futhark/LSP/TypeAscription.hs b/src/Futhark/LSP/TypeAscription.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/LSP/TypeAscription.hs @@ -0,0 +1,88 @@ +module Futhark.LSP.TypeAscription + ( TypeAscription (..), + missingAscriptions, + missingTypeParameters, + ) +where + +import Control.Arrow ((&&&)) +import Data.Foldable (asum) +import Data.Function ((&)) +import Data.Loc (Loc (..), Pos, locOf) +import Data.Map qualified as M +import Language.Futhark.Query + ( BoundTo (..), + TermBindSrc (..), + TermBinding (..), + TermFunData (..), + ) +import Language.Futhark.Syntax + ( ResRetType, + StructType, + TypeParamBase, + VName, + typeParamName, + ) + +data TypeAscription + = -- | type hint for a let-binding or nested parameter binding + TypeAscLet StructType Pos + | -- | type hint for a parameter + TypeAscParam Pos StructType Pos + | -- | type hint for a return type + TypeAscReturn ResRetType Pos + | -- | type hint for inferred type parameters + TypeAscType (TypeParamBase VName) Pos + +-- techincally, this name is not entirely correct since it covers type +-- parameters as well +missingAscriptions :: BoundTo -> [TypeAscription] +missingAscriptions (BoundTerm term (Loc termStart termEnd)) = + case term of + TermSize -> [] + TermVar _ _ (Just _) -> [] + TermVar src inferredType Nothing -> + let bareHints = [TypeAscLet inferredType termEnd] + in case src of + TermBindId -> [] + TermBindLet -> bareHints + TermBindPat -> + [TypeAscParam termStart inferredType termEnd] + TermBindNested -> bareHints + TermFun tfData -> + -- ordering is relevant, see the lsp documentation for inlay hints + let isSynthesized typeParam = case locOf typeParam of + NoLoc -> True + _ -> False + inferredTypeParams = filter isSynthesized (termFunTypeParams tfData) + paramInfo p = case termFunNameEnd tfData of + Nothing -> id + Just pos -> (TypeAscType p pos :) + in foldr (\p f hs -> paramInfo p $ f hs) id inferredTypeParams $ + let retType = termFunRetType tfData + retTypePos = + asum $ + map + ($ tfData) + [ termFunArgEnd, + termFunTypeArgEnd, + termFunNameEnd + ] + in case (termFunAscription tfData, retTypePos) of + (Just _, _) -> [] + (Nothing, Just pos) -> [TypeAscReturn retType pos] + _ -> [] +missingAscriptions _ = [] + +missingTypeParameters :: BoundTo -> M.Map VName (Maybe Pos, TypeParamBase VName) +missingTypeParameters (BoundTerm (TermFun tfData) _) = + M.fromList inferredTypeParams + where + inferredTypeParams = + termFunTypeParams tfData + & filter keepNoLoc + & map (typeParamName &&& (termFunNameEnd tfData,)) + keepNoLoc located = case locOf located of + NoLoc -> True + _ -> False +missingTypeParameters _ = M.empty diff --git a/src/Futhark/Optimise/ArrayShortCircuiting.hs b/src/Futhark/Optimise/ArrayShortCircuiting.hs --- a/src/Futhark/Optimise/ArrayShortCircuiting.hs +++ b/src/Futhark/Optimise/ArrayShortCircuiting.hs @@ -169,9 +169,9 @@ replaceInSegOp (SegRed lvl sp tps body binops) = do stms <- updateStms $ bodyStms body pure $ SegRed lvl sp tps (body {bodyStms = stms}) binops -replaceInSegOp (SegScan lvl sp tps body binops) = do +replaceInSegOp (SegScan lvl sp tps body binops post_op) = do stms <- updateStms $ bodyStms body - pure $ SegScan lvl sp tps (body {bodyStms = stms}) binops + pure $ SegScan lvl sp tps (body {bodyStms = stms}) binops post_op replaceInSegOp (SegHist lvl sp tps body hist_ops) = do stms <- updateStms $ bodyStms body pure $ SegHist lvl sp tps (body {bodyStms = stms}) hist_ops diff --git a/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs --- a/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs +++ b/src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs @@ -218,10 +218,11 @@ op <- binops let shp = Shape segment_dims <> segBinOpShape op map (`arrayOfShape` shp) (lambdaReturnType $ segBinOpLambda op) -shortCircuitSegOp lvlOK lutab pat pat_certs (SegScan lvl space _ kernel_body binops) td_env bu_env = +shortCircuitSegOp lvlOK lutab pat pat_certs (SegScan lvl space _ kernel_body binops post_op) td_env bu_env = -- Like in the handling of 'SegRed', we do not want to coalesce anything that - -- is used in the 'SegBinOp' - let to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` foldMap (freeIn . segBinOpLambda) binops) $ activeCoals bu_env + -- is used in the 'SegBinOp'. We do not coalesce anything that is using in SegPostOp either. + let free_in_lams = freeIn (segPostOpLambda post_op) <> foldMap (freeIn . segBinOpLambda) binops + to_fail = M.filter (\entry -> namesFromList (M.keys $ vartab entry) `namesIntersect` free_in_lams) $ activeCoals bu_env (active, inh) = foldl markFailedCoal (activeCoals bu_env, inhibit bu_env) $ M.keys to_fail bu_env' = bu_env {activeCoals = active, inhibit = inh} in shortCircuitSegOpHelper 0 lvlOK lvl lutab pat pat_certs space kernel_body td_env bu_env' 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 @@ -312,7 +312,7 @@ cseInOp = subCSE . GPU.mapSegOpM - (GPU.SegOpMapper pure cseInLambda cseInKernelBody pure pure) + (GPU.SegOpMapper pure cseInLambda cseInLambda cseInKernelBody pure pure) cseInKernelBody :: (Aliased rep, CSEInOp (Op rep)) => 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 @@ -155,7 +155,8 @@ where mapper = identitySegOpMapper - { mapOnSegOpLambda = optimiseLambda, + { mapOnSegBinOpLambda = optimiseLambda, + mapOnSegPostOpLambda = optimiseLambda, mapOnSegOpBody = optimiseKernelBody } inSegOp env = env {envOptimiseLoop = optimiseLoop} @@ -169,7 +170,8 @@ where mapper = identitySegOpMapper - { mapOnSegOpLambda = optimiseLambda, + { mapOnSegBinOpLambda = optimiseLambda, + mapOnSegPostOpLambda = optimiseLambda, mapOnSegOpBody = optimiseKernelBody } inSegOp env = env {envOptimiseLoop = optimiseLoop} 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 @@ -193,13 +193,27 @@ NodeT -> m NodeT makeCopiesOfFusedExcept noCopy (SoacNode ots pats soac aux) = do - let lam = H.lambda soac - localScope (scopeOf lam) $ do - fused_inner <- - filterM (fmap (not . isAcc) . lookupType) . namesToList . consumedByLambda $ - Alias.analyseLambda mempty lam - lam' <- makeCopiesInLambda (fused_inner L.\\ noCopy) lam - pure $ SoacNode ots pats (H.setLambda lam' soac) aux + case soac of + H.Screma w arrs (ScremaForm lam scans reduces postlam) -> do + localScope (scopeOf lam <> scopeOf postlam) $ do + fused_in_main <- + filterM (fmap (not . isAcc) . lookupType) . namesToList . consumedByLambda $ + Alias.analyseLambda mempty lam + fused_in_post <- + filterM (fmap (not . isAcc) . lookupType) . namesToList . consumedByLambda $ + Alias.analyseLambda mempty postlam + lam' <- makeCopiesInLambda (fused_in_main L.\\ noCopy) lam + postlam' <- makeCopiesInLambda (fused_in_post L.\\ noCopy) postlam + let form' = ScremaForm lam' scans reduces postlam' + pure $ SoacNode ots pats (H.Screma w arrs form') aux + _any -> do + let lam = H.lambda soac + localScope (scopeOf lam) $ do + fused_inner <- + filterM (fmap (not . isAcc) . lookupType) . namesToList . consumedByLambda $ + Alias.analyseLambda mempty lam + lam' <- makeCopiesInLambda (fused_inner L.\\ noCopy) lam + pure $ SoacNode ots pats (H.setLambda lam' soac) aux makeCopiesOfFusedExcept _ nodeT = pure nodeT makeCopiesInLambda :: @@ -483,38 +497,92 @@ WithAcc w3_inps lam3 hFuseNodeT _ _ = pure Nothing -removeOutputsExcept :: [VName] -> NodeT -> NodeT +removeOutputsExcept :: [VName] -> NodeT -> FusionM NodeT removeOutputsExcept toKeep s = case s of - SoacNode ots (Pat pats1) soac@(H.Screma _ _ (ScremaForm lam_1 scans_1 red_1)) aux1 -> - SoacNode ots (Pat $ pats_unchanged <> pats_new) (H.setLambda lam_new soac) aux1 + SoacNode ots (Pat pats) (H.Screma w inp (ScremaForm pre_lam [] red post_lam)) aux1 -> do + pre_lam' <- if changed then simplifyLambda new_pre else pure new_pre + post_lam' <- if changed then simplifyLambda new_post else pure new_post + pure $ + SoacNode + ots + (Pat $ red_pats <> new_pats) + (H.Screma w inp (ScremaForm pre_lam' [] red post_lam')) + aux1 where - scan_output_size = Futhark.scanResults scans_1 - red_output_size = Futhark.redResults red_1 + (pre_red_res, pre_map_res) = + splitAt (redResults red) $ resFromLambda pre_lam + (pre_red_ts, pre_map_ts) = + splitAt (redResults red) $ lambdaReturnType pre_lam - (pats_unchanged, pats_toChange) = splitAt (scan_output_size + red_output_size) pats1 - (res_unchanged, res_toChange) = splitAt (scan_output_size + red_output_size) (zip (resFromLambda lam_1) (lambdaReturnType lam_1)) + to_change = + L.zip5 + pre_map_res + pre_map_ts + (lambdaParams post_lam) + (resFromLambda post_lam) + (lambdaReturnType post_lam) - (pats_new, other) = unzip $ filter (\(x, _) -> patElemName x `elem` toKeep) (zip pats_toChange res_toChange) - (results, types) = unzip (res_unchanged ++ other) - lam_new = - lam_1 + (red_pats, map_pats) = splitAt (redResults red) pats + + changed = new_post /= post_lam || new_pre /= pre_lam + + (new_pats, new) = + unzip $ + filter (\(x, _) -> patElemName x `elem` toKeep) (zip map_pats to_change) + ( new_pre_map_res, + new_pre_map_ts, + new_post_pars, + new_post_res, + new_post_ts + ) = L.unzip5 new + new_post = + Lambda + { lambdaParams = new_post_pars, + lambdaReturnType = new_post_ts, + lambdaBody = (lambdaBody post_lam) {bodyResult = new_post_res} + } + new_pre = + pre_lam + { lambdaReturnType = pre_red_ts <> new_pre_map_ts, + lambdaBody = + (lambdaBody pre_lam) + { bodyResult = pre_red_res <> new_pre_map_res + } + } + SoacNode ots (Pat pats1) (H.Screma w inp (ScremaForm pre_lam scan red post_lam)) aux1 -> do + post_lam' <- if changed then simplifyLambda new_post else pure new_post + pure $ + SoacNode + ots + (Pat $ pats_unchanged <> pats_new) + (H.Screma w inp (ScremaForm pre_lam scan red post_lam')) + aux1 + where + red_output_size = Futhark.redResults red + + (pats_unchanged, pats_toChange) = splitAt red_output_size pats1 + res_toChange = zip (resFromLambda post_lam) (lambdaReturnType post_lam) + + changed = new_post /= post_lam + + (pats_new, res_new) = + unzip $ filter (\(x, _) -> patElemName x `elem` toKeep) (zip pats_toChange res_toChange) + (results, types) = unzip res_new + new_post = + post_lam { lambdaReturnType = types, - lambdaBody = (lambdaBody lam_1) {bodyResult = results} + lambdaBody = (lambdaBody post_lam) {bodyResult = results} } - node -> node + node -> pure node vNameFromAdj :: G.Node -> (EdgeT, G.Node) -> VName vNameFromAdj n1 (edge, n2) = depsFromEdge (n2, n1, edge) -removeUnusedOutputsFromContext :: DepContext -> FusionM DepContext -removeUnusedOutputsFromContext (incoming, n1, nodeT, outgoing) = - pure (incoming, n1, nodeT', outgoing) - where - toKeep = map (vNameFromAdj n1) incoming - nodeT' = removeOutputsExcept toKeep nodeT - removeUnusedOutputs :: DepGraphAug FusionM -removeUnusedOutputs = mapAcross removeUnusedOutputsFromContext +removeUnusedOutputs = mapAcross $ \(incoming, n1, nodeT, outgoing) -> do + let toKeep = map (vNameFromAdj n1) incoming + nodeT' <- removeOutputsExcept toKeep nodeT + pure (incoming, n1, nodeT', outgoing) tryFuseNodeInGraph :: DepNode -> DepGraphAug FusionM tryFuseNodeInGraph node_to_fuse dg@DepGraph {dgGraph = g} @@ -582,12 +650,10 @@ doAllFusion :: DepGraphAug FusionM doAllFusion = - applyAugs - [ keepTrying . applyAugs $ - [ doVerticalFusion, - doHorizontalFusion, - doInnerFusion - ], + keepTrying . applyAugs $ + [ doVerticalFusion, + doHorizontalFusion, + doInnerFusion, removeUnusedOutputs ] @@ -614,12 +680,13 @@ soac' <- case soac of H.Stream w inputs accs lam -> H.Stream w inputs accs <$> dontFuseScans (onLambda lam) - H.Screma w inputs (ScremaForm lam scans reds) -> + H.Screma w inputs (ScremaForm lam scans reds post_lam) -> H.Screma w inputs <$> ( ScremaForm <$> doFuseScans (onLambda lam) <*> mapM onScan scans <*> mapM onRed reds + <*> doFuseScans (onLambda post_lam) ) _ -> H.setLambda <$> doFuseScans (onLambda (H.lambda soac)) <*> pure soac diff --git a/src/Futhark/Optimise/Fusion/Screma.hs b/src/Futhark/Optimise/Fusion/Screma.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Fusion/Screma.hs @@ -0,0 +1,495 @@ +module Futhark.Optimise.Fusion.Screma + ( splitLambdaByPar, + fuseScrema, + fuseSuperScrema, + SuperScrema (..), + moveRedScanSuperScrema, + moveLastSuperScrema, + fusible, + toScrema, + ) +where + +import Control.Monad +import Data.Bifunctor +import Data.Function +import Data.List qualified as L +import Data.Map.Strict qualified as M +import Data.Maybe +import Futhark.Analysis.Alias +import Futhark.Analysis.DataDependencies +import Futhark.Analysis.HORep.SOAC qualified as SOAC +import Futhark.IR +import Futhark.IR.Prop.Aliases +import Futhark.IR.SOACS +import Futhark.IR.SOACS.Simplify +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Transform.Rename +import Futhark.Util (splitAt3) +import Futhark.Util.Pretty + +-- | The fused representation of two scremas. +data SuperScrema rep + = SuperScrema + SubExp + [SOAC.Input] + (Lambda rep) + [Scan rep] + [Reduce rep] + (Lambda rep) + [Scan rep] + [Reduce rep] + (Lambda rep) + deriving (Eq, Ord, Show) + +instance (PrettyRep rep) => Pretty (SuperScrema rep) where + pretty (SuperScrema w inps lam1 scans1 reds1 lam2 scans2 reds2 lam3) = + "superscrema" + <> (parens . align) + ( (pretty w <> comma) + ppTuple' (map pretty inps) + pretty lam1 + <> comma + p' scans1 + <> comma + p' reds1 + <> comma + pretty lam2 + <> comma + p' scans2 + <> comma + p' reds2 + <> comma + pretty lam3 + ) + where + p' xs = braces (mconcat $ L.intersperse (comma <> line) $ map pretty xs) + +-- | Using a boolean mask pick elements in the +pick :: [Bool] -> [a] -> [a] +pick bs xs = map snd $ filter fst $ zip bs xs + +-- | Check that all inputs and outputs are fusible for a producer and +-- consumer. +fuseIsVarish :: [SOAC.Input] -> [VName] -> Bool +fuseIsVarish inp_c = + all (maybe True (isJust . SOAC.isVarishInput) . flip M.lookup name_to_inp_c) + where + name_to_inp_c = M.fromList $ zip (SOAC.inputArray <$> inp_c) inp_c + +-- | Given a list of parameter names and a lambda, split the lambda +-- into two where the first function will depend on the given +-- parameters given if they exist plus additional parameters if they +-- are need for some computation. The other lambda will have all the +-- other parameters and do the remaining computational work. Some +-- statements may be present in both lambdas. +-- +-- Further, each parameter is associated with some additional +-- information @a@, and each result with some additional information +-- @b@. This is also partitioned and returned appropriately. +splitLambdaByPar :: + (MonadFail m) => + [VName] -> + [inp] -> + Lambda SOACS -> + [out] -> + m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out])) +splitLambdaByPar names inp lam out = do + when + (consumedOverlap new_lam new_lam') + (fail "Can not fuse due to overlap consumption in pre and post.") + when + (parAccsOverlap new_lam new_lam') + (fail "Can not fuse due to overlap in parameter accumalators.") + when + (resAccsOverlap new_lam new_lam') + (fail "Can not fuse due to overlap in result accumalators.") + pure + ( (new_inp, new_lam, new_out), + (new_inp', new_lam', new_out') + ) + where + new_lam = eliminateByRes $ Lambda new_params new_ts (mkBody stms new_res) + new_lam' = eliminateByRes $ Lambda new_params' new_ts' (mkBody stms new_res') + pars = lambdaParams lam + m = M.fromList $ zip pars inp + par_deps = lambdaDependencies mempty lam (oneName . paramName <$> pars) + body = lambdaBody lam + stms = bodyStms body + new_inp = (m M.!) <$> new_params + new_inp' = (m M.!) <$> new_params' + new_params = filter ((`nameIn` deps) . paramName) pars + new_params' = filter ((`nameIn` deps') . paramName) pars + auxiliary = (\(a, b, c, d) -> (mconcat a, b, c, d)) . L.unzip4 + ((deps, new_res, new_ts, new_out), (deps', new_res', new_ts', new_out')) = + bimap auxiliary auxiliary + . L.partition (namesIntersect (namesFromList names) . (\(a, _, _, _) -> a)) + $ L.zip4 par_deps (bodyResult body) (lambdaReturnType lam) out + +-- | Check that two scremas are fusible if they are give back the +-- producer scremas post lambda that has been split into the scan +-- lambda and map lambda. It is not fusible if inputs and outputs are +-- being transformed and if the producers scan result is used for the +-- consumers scans or reduces. +fusible :: + (MonadFail m) => + [SOAC.Input] -> + ScremaForm SOACS -> + [VName] -> + [SOAC.Input] -> + ScremaForm SOACS -> + [VName] -> + m () +fusible inp_p form_p out_p inp_c form_c out_c = do + ((_, post_scan_p, _), _) <- + splitLambdaByPar post_scan_pars_p inp_p post_p out_c + let post_scan_res_p = bodyResult $ lambdaBody post_scan_p + forbidden_p = namesFromList $ resToOut out_p post_p <$> post_scan_res_p + is_fusible = + fuseIsVarish inp_c out_p + && not (forbidden_c `namesIntersect` forbidden_p) + unless is_fusible (fail "Scremas are not fusible.") + where + pre_pars_c = oneName . paramName <$> lambdaParams pre_c + (pre_scan_deps_c, pre_red_deps_c, _) = + splitAt3 num_scan_c num_red_c $ + lambdaDependencies mempty pre_c pre_pars_c + forbidden_c = + namesFromList + . mapMaybe (fmap SOAC.inputArray . parToInp inp_c pre_c) + . namesToList + $ mconcat (pre_scan_deps_c <> pre_red_deps_c) + pre_c = scremaLambda form_c + post_p = scremaPostLambda form_p + post_scan_pars_p = take num_scan_p $ paramName <$> lambdaParams post_p + num_scan_c = scanResults $ scremaScans form_c + num_red_c = redResults $ scremaReduces form_c + num_scan_p = scanResults $ scremaScans form_p + +-- | Given two scremas that are fusible, fuse them into a super +-- screma. This is fused but work will have to be moved around in the +-- super screma for it to become a screma. +fuseSuperScrema :: + (MonadFreshNames m) => + SubExp -> + [SOAC.Input] -> + ScremaForm SOACS -> + [VName] -> + [SOAC.Input] -> + ScremaForm SOACS -> + [VName] -> + m (SuperScrema SOACS, [VName]) +fuseSuperScrema w inp_p' form_p' out_p inp_c' form_c' out_c = do + let (inp_p, form_p) = dedupInput inp_p' form_p' + (inp_c, form_c) = dedupInput inp_c' form_c' + inp_c_real_map = map (not . inputFromOutput) inp_c + inp_c_real = pick inp_c_real_map inp_c + inp_r = inp_p <> inp_c_real + + (out_red_p, out_post_p) = + splitAt (redResults $ scremaReduces form_p) out_p + (out_red_c, out_post_c) = + splitAt (redResults $ scremaReduces form_c) out_c + + forward_params <- forM (pick inp_c_real_map (lambdaParams (scremaLambda form_c))) $ \p -> + newParam (baseName (paramName p)) (paramType p) + + let lam1 = + Lambda + { lambdaParams = + lambdaParams (scremaLambda form_p) <> forward_params, + lambdaReturnType = + lambdaReturnType (scremaLambda form_p) + <> map paramType forward_params, + lambdaBody = + mkBody + (bodyStms (lambdaBody (scremaLambda form_p))) + ( bodyResult (lambdaBody (scremaLambda form_p)) + <> varsRes (map paramName forward_params) + ) + } + + let lam2 = + Lambda + { lambdaParams = + lambdaParams (scremaPostLambda form_p) + <> pick inp_c_real_map (lambdaParams (scremaLambda form_c)), + lambdaReturnType = + lambdaReturnType (scremaLambda form_c) + <> lambdaReturnType (scremaPostLambda form_p), + lambdaBody = + mkBody + ( bodyStms (lambdaBody (scremaPostLambda form_p)) + <> composeBinds + (scremaPostLambda form_p) + out_post_p + (SOAC.inputArray <$> inp_c) + (scremaLambda form_c) + <> bodyStms (lambdaBody (scremaLambda form_c)) + ) + ( bodyResult (lambdaBody (scremaLambda form_c)) + <> bodyResult (lambdaBody (scremaPostLambda form_p)) + ) + } + + post_forward_params <- forM + ( zip + (bodyResult (lambdaBody (scremaPostLambda form_p))) + (lambdaReturnType (scremaPostLambda form_p)) + ) + $ \(res, t) -> + newParam (maybe (nameFromString "x") baseName (subExpResVName res)) t + + let lam3 = + Lambda + { lambdaParams = + lambdaParams (scremaPostLambda form_c) <> post_forward_params, + lambdaReturnType = + lambdaReturnType (scremaPostLambda form_c) + <> map paramType post_forward_params, + lambdaBody = + mkBody + (bodyStms (lambdaBody (scremaPostLambda form_c))) + ( bodyResult (lambdaBody (scremaPostLambda form_c)) + <> varsRes (map paramName post_forward_params) + ) + } + pure + ( SuperScrema + w + inp_r + lam1 + (scremaScans form_p) + (scremaReduces form_p) + lam2 + (scremaScans form_c) + (scremaReduces form_c) + lam3, + out_red_p <> out_red_c <> out_post_c <> out_post_p + ) + where + inputFromOutput inp = SOAC.inputArray inp `elem` out_p + +-- | Moves the last scan and reduce from the super screma to the top +-- of the super screma. +moveRedScanSuperScrema :: + (MonadFail m, MonadFreshNames m) => + SuperScrema SOACS -> + m (SuperScrema SOACS) +moveRedScanSuperScrema super_screma = do + ((scan_red_inp_c, scan_red_lam', _), (_, map_lam', _)) <- + splitAtLambdaByRes + (scanResults scan' + redResults red') + inp_c + lam' + (replicate (length $ lambdaReturnType lam') ()) + renamed_scan_red_lam' <- renameLambda scan_red_lam' + let (scan_res', red_res') = + splitAt (scanResults scan') + . bodyResult + $ lambdaBody renamed_scan_red_lam' + (scan_ts', red_ts') = + splitAt (scanResults scan') $ + lambdaReturnType renamed_scan_red_lam' + binds = composeBinds lam out_p scan_red_inp_c renamed_scan_red_lam' + stms' = bodyStms $ lambdaBody renamed_scan_red_lam' + new_scan = scan <> scan' + new_red = red <> red' + new_ts = scan_ts <> scan_ts' <> red_ts <> red_ts' <> map_ts + new_pars = lambdaParams lam + new_res = scan_res <> scan_res' <> red_res <> red_res' <> map_res + new_body = mkBody (stms <> binds <> stms') new_res + new_lam = eliminateByRes $ Lambda new_pars new_ts new_body + (scan_pars', map_pars') = + splitAt (scanResults scan) (lambdaParams lam') + + extra_scan_pars' <- mapM (newParam "x") scan_ts' + + let new_pars' = scan_pars' <> extra_scan_pars' <> map_pars' + new_ts' = scan_ts' <> lambdaReturnType map_lam' + new_stms' = bodyStms $ lambdaBody map_lam' + new_res' = + varsRes (map paramName extra_scan_pars') + <> bodyResult (lambdaBody map_lam') + new_body' = mkBody new_stms' new_res' + new_lam' = eliminateByRes $ Lambda new_pars' new_ts' new_body' + + pure $ + SuperScrema w inp new_lam new_scan new_red new_lam' [] [] lam'' + where + stms = bodyStms $ lambdaBody lam + out_p = [0 .. length (bodyResult $ lambdaBody lam) - 1] + (scan_out, _, map_out) = + splitAt3 + (scanResults scan) + (redResults red) + [0 .. length (bodyResult $ lambdaBody lam) - 1] + (scan_res, red_res, map_res) = + splitAt3 + (scanResults scan) + (redResults red) + (bodyResult $ lambdaBody lam) + (scan_ts, red_ts, map_ts) = + splitAt3 + (scanResults scan) + (redResults red) + (lambdaReturnType lam) + inp_c = scan_out <> map_out + + SuperScrema w inp lam scan red lam' scan' red' lam'' = + super_screma + +-- | Moves all work done in the last lambda to middle lambda of the +-- super screma. +moveLastSuperScrema :: + (MonadFreshNames m) => + SuperScrema SOACS -> + m (SuperScrema SOACS) +moveLastSuperScrema (SuperScrema w inp lam scan red lam' [] [] lam'') = do + temp_lam'' <- renameLambda lam'' + let new_pars = lambdaParams lam' + new_ts = lambdaReturnType lam'' + out_p = [0 .. length (bodyResult $ lambdaBody lam') - 1] + inp_c = out_p + binds = composeBinds lam' out_p inp_c temp_lam'' + stms' = bodyStms $ lambdaBody lam' + stms'' = bodyStms $ lambdaBody temp_lam'' + new_stms = stms' <> binds <> stms'' + new_res = bodyResult $ lambdaBody temp_lam'' + new_body = mkBody new_stms new_res + new_lam' = eliminateByRes $ Lambda new_pars new_ts new_body + + new_lam'' <- mkIdentityLambda $ lambdaReturnType lam'' + pure $ + SuperScrema w inp lam scan red new_lam' [] [] new_lam'' +moveLastSuperScrema _ = + error "moveLastSuperScrema must not have any scans or reduce operation in the end." + +-- | Find all Accumulator parameters. +parAccs :: Lambda SOACS -> [Type] +parAccs = filter isAcc . map typeOf . lambdaParams + +-- | Find all Accumulator results. +resAccs :: Lambda SOACS -> [Type] +resAccs = filter isAcc . lambdaReturnType + +-- | Check if the lambda parameters have overlapping accumulators. +parAccsOverlap :: Lambda SOACS -> Lambda SOACS -> Bool +parAccsOverlap lam = any (`elem` accs) . parAccs + where + accs = parAccs lam + +-- | Check if the lambdas result have overlapping accumulators. +resAccsOverlap :: Lambda SOACS -> Lambda SOACS -> Bool +resAccsOverlap lam = any (`elem` accs) . resAccs + where + accs = resAccs lam + +-- | Check if the lambdas have parameters that overlap due to +-- consumption. +consumedOverlap :: Lambda SOACS -> Lambda SOACS -> Bool +consumedOverlap = + on namesIntersect (consumedByLambda . analyseLambda mempty) + +-- | Split a lambda at some index for the result and construct two +-- lambda which do not have overlapping results but may have +-- overlapping parameters. +splitAtLambdaByRes :: + (MonadFail m) => + Int -> + [inp] -> + Lambda SOACS -> + [out] -> + m (([inp], Lambda SOACS, [out]), ([inp], Lambda SOACS, [out])) +splitAtLambdaByRes i inp lam out = do + when + (consumedOverlap new_lam new_lam') + (fail "Can not fuse due to overlap consumption in pre and post.") + when + (parAccsOverlap new_lam new_lam') + (fail "Can not fuse due to overlap in parameter accumalators.") + pure ((new_inp, new_lam, new_out), (new_inp', new_lam', new_out')) + where + new_lam = Lambda new_pars new_ts new_body + new_lam' = Lambda new_pars' new_ts' new_body' + pars = lambdaParams lam + stms = bodyStms $ lambdaBody lam + new_body = mkBody (eliminate (freeIn new_res) stms) new_res + new_body' = mkBody (eliminate (freeIn new_res') stms) new_res' + inBody body = (`nameIn` freeIn body) . paramName . fst + removePars body = unzip $ filter (inBody body) $ zip pars inp + (new_pars, new_inp) = removePars new_body + (new_pars', new_inp') = removePars new_body' + (new_res, new_res') = splitAt i $ bodyResult $ lambdaBody lam + (new_ts, new_ts') = splitAt i $ lambdaReturnType lam + (new_out, new_out') = splitAt i out + +-- | Create a mapping from lambda results expression to their output +-- array. +resToOut :: [VName] -> Lambda SOACS -> SubExpRes -> VName +resToOut out lam = (m M.!) + where + m = M.fromList $ flip zip out $ bodyResult $ lambdaBody lam + +-- | Create a mapping from lambda parameter names to their input +-- array. +parToInp :: [SOAC.Input] -> Lambda SOACS -> VName -> Maybe SOAC.Input +parToInp inp lam = flip M.lookup m + where + m = M.fromList $ flip zip inp $ paramName <$> lambdaParams lam + +-- | Turn a SuperScrema into Screma. +toScrema :: + SuperScrema SOACS -> + ([SOAC.Input], ScremaForm SOACS) +toScrema (SuperScrema _ inp lam scan red lam' _ _ _) = + (inp, ScremaForm lam scan red lam') + +-- | Try to force the post lambda to be an indentity lambda. +tryIdentityPost :: + (MonadFreshNames m) => ScremaForm SOACS -> m (ScremaForm SOACS) +tryIdentityPost (ScremaForm pre_lam [] reds post_lam) + | not $ isIdentityLambda post_lam = do + new_post_lam <- mkIdentityLambda $ lambdaReturnType post_lam + pure (ScremaForm new_pre_lam [] reds new_post_lam) + where + out_p = [0 .. length (bodyResult $ lambdaBody pre_lam) - 1] + inp_c = drop (redResults reds) out_p + binds = composeBinds pre_lam out_p inp_c post_lam + pre_stms = bodyStms $ lambdaBody pre_lam + post_stms = bodyStms $ lambdaBody post_lam + post_res = bodyResult $ lambdaBody post_lam + pre_res = bodyResult $ lambdaBody pre_lam + post_ts = lambdaReturnType post_lam + pre_ts = lambdaReturnType pre_lam + res = take (redResults reds) pre_res <> post_res + ts = take (redResults reds) pre_ts <> post_ts + new_pre_lam = + Lambda + { lambdaParams = lambdaParams pre_lam, + lambdaReturnType = ts, + lambdaBody = mkBody (pre_stms <> binds <> post_stms) res + } +tryIdentityPost form = pure form + +-- Tries to fuse two Scremas into one. +fuseScrema :: + (MonadFail m, MonadFreshNames m, HasScope SOACS m) => + SubExp -> + [SOAC.Input] -> + ScremaForm SOACS -> + [VName] -> + [SOAC.Input] -> + ScremaForm SOACS -> + [VName] -> + m ([SOAC.Input], ScremaForm SOACS, [VName]) +fuseScrema w inp_p form_p out_p inp_c form_c out_c = do + fusible inp_p form_p out_p inp_c form_c out_c + (super_screma, new_out) <- fuseSuperScrema w inp_p form_p out_p inp_c form_c out_c + (new_inp, form') <- + fmap (second prunePreLambdaResults . toScrema) $ + moveRedScanSuperScrema super_screma + >>= moveLastSuperScrema + form <- tryIdentityPost form' + pure (new_inp, form, new_out) diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs --- a/src/Futhark/Optimise/Fusion/TryFusion.hs +++ b/src/Futhark/Optimise/Fusion/TryFusion.hs @@ -18,7 +18,7 @@ import Control.Monad import Control.Monad.Reader import Control.Monad.State -import Data.List (find, (\\)) +import Data.List (find) import Data.Map.Strict qualified as M import Data.Maybe import Futhark.Analysis.HORep.MapNest (MapNest) @@ -28,10 +28,10 @@ import Futhark.IR.SOACS hiding (SOAC (..)) import Futhark.IR.SOACS qualified as Futhark import Futhark.Optimise.Fusion.Composing +import Futhark.Optimise.Fusion.Screma import Futhark.Pass.ExtractKernels.ISRWIM (rwimPossible) import Futhark.Transform.Rename (renameLambda) import Futhark.Transform.Substitute -import Futhark.Util (splitAt3) newtype TryFusion a = TryFusion @@ -187,11 +187,6 @@ scope <- askScope tryFusion (applyFusionRules mode unfus_nms outVars soac ker) scope --- | Check that the consumer does not use any scan or reduce results. -scremaFusionOK :: ([VName], [VName]) -> FusedSOAC -> Bool -scremaFusionOK (nonmap_outs, _map_outs) ker = - all (`notElem` nonmap_outs) $ mapMaybe SOAC.isVarishInput (inputs ker) - -- | Check that the consumer uses all the outputs of the producer unmodified. mapWriteFusionOK :: [VName] -> FusedSOAC -> Bool mapWriteFusionOK outVars ker = all (`elem` inpIds) outVars @@ -212,6 +207,8 @@ let soac_c = fsSOAC ker inp_p_arr = SOAC.inputs soac_p inp_c_arr = SOAC.inputs soac_c + out_p = outVars + out_c = fsOutNames ker lam_p = SOAC.lambda soac_p lam_c = SOAC.lambda soac_c w = SOAC.width soac_p @@ -260,53 +257,31 @@ (_, _, Vertical) | unfus_set /= mempty, not (SOAC.nullTransforms $ fsOutputTransform ker) -> - fail "Cannot perform diagonal fusion in the presence of output transforms." - ( SOAC.Screma _ _ (ScremaForm _ scans_c reds_c), - SOAC.Screma _ _ (ScremaForm _ scans_p reds_p), + fail + "Cannot perform diagonal fusion in the presence of output transforms." + ( SOAC.Screma _ inp_c form_c, + SOAC.Screma _ inp_p form_p, _ - ) - | scremaFusionOK - ( splitAt - ( Futhark.scanResults scans_p - + Futhark.redResults reds_p - ) - outVars - ) - ker -> do - let red_nes_p = concatMap redNeutral reds_p - red_nes_c = concatMap redNeutral reds_c - scan_nes_p = concatMap scanNeutral scans_p - scan_nes_c = concatMap scanNeutral scans_c - (res_lam', new_inp) = - fuseRedomap - unfus_set - outVars - lam_p - scan_nes_p - red_nes_p - inp_p_arr - outPairs - lam_c - scan_nes_c - red_nes_c - inp_c_arr - (soac_p_scanout, soac_p_redout, _soac_p_mapout) = - splitAt3 (length scan_nes_p) (length red_nes_p) outVars - (soac_c_scanout, soac_c_redout, soac_c_mapout) = - splitAt3 (length scan_nes_c) (length red_nes_c) $ fsOutNames ker - unfus_arrs = returned_outvars \\ (soac_p_scanout ++ soac_p_redout) - success - ( soac_p_scanout - ++ soac_c_scanout - ++ soac_p_redout - ++ soac_c_redout - ++ soac_c_mapout - ++ unfus_arrs - ) - $ SOAC.Screma - w - new_inp - (ScremaForm res_lam' (scans_p ++ scans_c) (reds_p ++ reds_c)) + ) -> + do + (inp, form, out) <- fuseScrema w inp_p form_p out_p inp_c form_c out_c + success out $ SOAC.Screma w inp form + <|> do + -- If these two scremas cannot be fused, then see what happens if we + -- turn the producer into a stream. The most common (only?) case of + -- this mattering is when the producer is a scan and the consumer is + -- a reduction. + Just _ <- pure $ Futhark.isScanomapSOAC form_p + (soac_p', newacc_ids) <- SOAC.soacToStream soac_p + if soac_p' /= soac_p + then + fuseSOACwithKer + mode + (namesFromList (map identName newacc_ids) <> unfus_set) + (map identName newacc_ids ++ outVars) + soac_p' + ker + else fail "SOAC could not be turned into stream." -- Map-Hist fusion. -- @@ -388,20 +363,6 @@ (map identName newacc_ids ++ outVars) soac_p' ker - (_, SOAC.Screma _ _ form, _) | Just _ <- Futhark.isScanomapSOAC form -> do - -- A Scan soac can be currently only fused as a (sequential) stream, - -- hence it is first translated to a (sequential) Stream and then - -- fusion with a kernel is attempted. - (soac_p', newacc_ids) <- SOAC.soacToStream soac_p - if soac_p' /= soac_p - then - fuseSOACwithKer - mode - (namesFromList (map identName newacc_ids) <> unfus_set) - (map identName newacc_ids ++ outVars) - soac_p' - ker - else fail "SOAC could not be turned into stream." (_, SOAC.Stream {}, _) -> do -- If it reached this case then soac_c is NOT a Stream kernel, -- hence transform the kernel's soac to a stream and attempt @@ -419,11 +380,6 @@ $ ker {fsSOAC = soac_c', fsOutNames = map identName newacc_ids ++ fsOutNames ker} else fail "SOAC could not be turned into stream." - --------------------------------- - --- DEFAULT, CANNOT FUSE CASE --- - --------------------------------- - _ -> fail "Cannot fuse" - fuseStreamHelper :: [VName] -> Names -> @@ -544,10 +500,9 @@ [] -> [] t : _ -> 1 : 0 : [2 .. arrayRank t] - pure - ( SOAC.Screma map_w map_arrs' (mapSOAC map_fun'), - ots SOAC.|> SOAC.Rearrange map_aux perm - ) + (,ots SOAC.|> SOAC.Rearrange map_aux perm) + . SOAC.Screma map_w map_arrs' + <$> mapSOAC map_fun' iswim _ _ _ = fail "ISWIM does not apply." @@ -651,7 +606,7 @@ else runLambdaBuilder (lambdaParams lam) $ mapM sliceRes =<< bodyBind (lambdaBody lam) - pure (SOAC.Screma w' (map sliceInput inps) (mapSOAC lam'), ots') + (,ots') . SOAC.Screma w' (map sliceInput inps) <$> mapSOAC lam' pullIndex _ _ = fail "Cannot pull index" pushRearrange :: diff --git a/src/Futhark/Optimise/GenRedOpt.hs b/src/Futhark/Optimise/GenRedOpt.hs --- a/src/Futhark/Optimise/GenRedOpt.hs +++ b/src/Futhark/Optimise/GenRedOpt.hs @@ -159,7 +159,7 @@ map_lam <- renameLambda map_lam0 (k1_res, ker1_stms) <- runBuilderT' $ do iota <- letExp "iota" $ BasicOp $ Iota inv_dim_len (intConst Int64 0) (intConst Int64 1) Int64 - let op_exp = Op (OtherOp (Screma inv_dim_len [iota] (ScremaForm map_lam [] [red]))) + let op_exp = Op (OtherOp (Screma inv_dim_len [iota] (ScremaForm map_lam [] [red] nilFn))) res_redmap <- letTupExp "res_mapred" op_exp letSubExp (baseName pat_acc_nm <> "_big_update") $ BasicOp (UpdateAcc safety acc_nm acc_inds $ map Var res_redmap) diff --git a/src/Futhark/Optimise/MemoryBlockMerging.hs b/src/Futhark/Optimise/MemoryBlockMerging.hs --- a/src/Futhark/Optimise/MemoryBlockMerging.hs +++ b/src/Futhark/Optimise/MemoryBlockMerging.hs @@ -39,7 +39,7 @@ foldMap getAllocsStm (bodyStms body) getAllocsSegOp (SegRed _ _ _ body _) = foldMap getAllocsStm (bodyStms body) -getAllocsSegOp (SegScan _ _ _ body _) = +getAllocsSegOp (SegScan _ _ _ body _ _) = foldMap getAllocsStm (bodyStms body) getAllocsSegOp (SegHist _ _ _ body _) = foldMap getAllocsStm (bodyStms body) @@ -73,8 +73,8 @@ SegRed lvl sp tps body' ops where body' = body {bodyStms = setAllocsStm m <$> bodyStms body} -setAllocsSegOp m (SegScan lvl sp tps body ops) = - SegScan lvl sp tps body' ops +setAllocsSegOp m (SegScan lvl sp tps body ops post_op) = + SegScan lvl sp tps body' ops post_op where body' = body {bodyStms = setAllocsStm m <$> bodyStms body} setAllocsSegOp m (SegHist lvl sp tps body ops) = @@ -111,9 +111,9 @@ onKernelBodyStms (SegRed lvl space ts body binops) f = do stms <- f $ bodyStms body pure $ SegRed lvl space ts (body {bodyStms = stms}) binops -onKernelBodyStms (SegScan lvl space ts body binops) f = do +onKernelBodyStms (SegScan lvl space ts body binops post_op) f = do stms <- f $ bodyStms body - pure $ SegScan lvl space ts (body {bodyStms = stms}) binops + pure $ SegScan lvl space ts (body {bodyStms = stms}) binops post_op onKernelBodyStms (SegHist lvl space ts body binops) f = do stms <- f $ bodyStms body pure $ SegHist lvl space ts (body {bodyStms = stms}) binops @@ -155,8 +155,8 @@ SegRed lvl sp tps body' ops where body' = body {bodyStms = maxstms <> stms <> bodyStms body} - SegScan lvl sp tps body ops -> - SegScan lvl sp tps body' ops + SegScan lvl sp tps body ops post_op -> + SegScan lvl sp tps body' ops post_op where body' = body {bodyStms = maxstms <> stms <> bodyStms body} SegHist lvl sp tps body ops -> diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs.hs b/src/Futhark/Optimise/ReduceDeviceSyncs.hs --- a/src/Futhark/Optimise/ReduceDeviceSyncs.hs +++ b/src/Futhark/Optimise/ReduceDeviceSyncs.hs @@ -330,10 +330,11 @@ ops' <- mapM addReadsToSegBinOp ops kbody' <- addReadsToBody kbody pure . SegOp $ SegRed lvl space types kbody' ops' -optimizeHostOp (SegOp (SegScan lvl space types kbody ops)) = do +optimizeHostOp (SegOp (SegScan lvl space types kbody ops post_op)) = do ops' <- mapM addReadsToSegBinOp ops + post_op' <- addReadsToSegPostOp post_op kbody' <- addReadsToBody kbody - pure . SegOp $ SegScan lvl space types kbody' ops' + pure . SegOp $ SegScan lvl space types kbody' ops' post_op' optimizeHostOp (SegOp (SegHist lvl space types kbody ops)) = do ops' <- mapM addReadsToHistOp ops kbody' <- addReadsToBody kbody @@ -643,6 +644,12 @@ addReadsToSegBinOp op = do f' <- addReadsToLambda (segBinOpLambda op) pure (op {segBinOpLambda = f'}) + +-- | Rewrite 'SegPostOp' dependencies to scalars that have been migrated. +addReadsToSegPostOp :: SegPostOp GPU -> ReduceM (SegPostOp GPU) +addReadsToSegPostOp op = do + f' <- addReadsToLambda (segPostOpLambda op) + pure (op {segPostOpLambda = f'}) -- | Rewrite 'HistOp' dependencies to scalars that have been migrated. addReadsToHistOp :: HistOp GPU -> ReduceM (HistOp GPU) diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs --- a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs +++ b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs @@ -1081,7 +1081,7 @@ collectSegLevel lvl collectSegSpace sp mapM_ collectSegBinOp ops - collectHostOp (SegOp (SegScan lvl sp _ _ ops)) = do + collectHostOp (SegOp (SegScan lvl sp _ _ ops _)) = do collectSegLevel lvl collectSegSpace sp mapM_ collectSegBinOp ops 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 @@ -221,7 +221,13 @@ where opMapper scope = identitySegOpMapper - { mapOnSegOpLambda = \lam -> do + { mapOnSegBinOpLambda = \lam -> do + let (body, sunk) = + optimiseBody onOp op_vtable sinking $ + lambdaBody lam + modify (<> sunk) + pure lam {lambdaBody = body}, + mapOnSegPostOpLambda = \lam -> do let (body, sunk) = optimiseBody onOp op_vtable sinking $ lambdaBody lam 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 @@ -837,7 +837,7 @@ tiles' <- mapM sliceTile tiles - let form' = redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam + form' <- redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam fmap varsRes $ letTupExp "acc" =<< eIf @@ -1091,9 +1091,9 @@ -- point). thread_accs <- forM accs $ \acc -> letSubExp "acc" $ BasicOp $ Index acc $ Slice [DimFix $ Var ltid_x, DimFix $ Var ltid_y] - let form' = redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam + form' <- redomapSOAC [Reduce red_comm red_lam thread_accs] map_lam - sliceTile (InputUntiled arr) = + let sliceTile (InputUntiled arr) = sliceUntiled arr tile_id tile_size actual_tile_size sliceTile (InputTiled perm tile) = do tile_t <- lookupType tile 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 @@ -116,7 +116,8 @@ optimise = identitySegOpMapper { mapOnSegOpBody = optimiseBody onOp, - mapOnSegOpLambda = optimiseLambda onOp + mapOnSegBinOpLambda = optimiseLambda onOp, + mapOnSegPostOpLambda = optimiseLambda onOp } onMCOp :: Stage -> OnOp MC 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 @@ -125,13 +125,14 @@ ( alloc_stms, Op $ Inner $ SegOp $ SegRed lvl' space ts kbody' reds' ) -transformExp (Op (Inner (SegOp (SegScan lvl space ts kbody scans)))) = do - (alloc_stms, (lvl', lams, kbody')) <- - transformScanRed lvl space (map segBinOpLambda scans) kbody +transformExp (Op (Inner (SegOp (SegScan lvl space ts kbody scans post_op)))) = do + (alloc_stms, (lvl', lams', kbody')) <- + transformScanRed lvl space (segPostOpLambda post_op : map segBinOpLambda scans) kbody + let post_op_lam : lams = lams' let scans' = zipWith (\red lam -> red {segBinOpLambda = lam}) scans lams pure ( alloc_stms, - Op $ Inner $ SegOp $ SegScan lvl' space ts kbody' scans' + Op $ Inner $ SegOp $ SegScan lvl' space ts kbody' scans' (post_op {segPostOpLambda = post_op_lam}) ) transformExp (Op (Inner (SegOp (SegHist lvl space ts kbody ops)))) = do (alloc_stms, (lvl', lams', kbody')) <- transformScanRed lvl space lams kbody @@ -455,7 +456,8 @@ opMapper user' = identitySegOpMapper - { mapOnSegOpLambda = onLambda user', + { mapOnSegBinOpLambda = onLambda user', + mapOnSegPostOpLambda = onLambda user', mapOnSegOpBody = onKernelBody user' } @@ -812,7 +814,8 @@ segOpMapper = identitySegOpMapper { mapOnSegOpBody = offsetMemoryInKernelBody offsets, - mapOnSegOpLambda = offsetMemoryInLambda offsets + mapOnSegBinOpLambda = offsetMemoryInLambda offsets, + mapOnSegPostOpLambda = offsetMemoryInLambda offsets } onOp op = pure op @@ -900,7 +903,8 @@ where mapper = identitySegOpMapper - { mapOnSegOpLambda = unAllocLambda, + { mapOnSegBinOpLambda = unAllocLambda, + mapOnSegPostOpLambda = unAllocLambda, mapOnSegOpBody = unAllocKernelBody } 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 @@ -16,6 +16,8 @@ AllocEnv (..), SizeSubst (..), allocInStms, + allocInLambda, + allocInLParams, allocForArray, simplifiable, mkLetNamesB', @@ -1131,3 +1133,28 @@ defaultExpHints :: (ASTRep rep, HasScope rep m) => Exp rep -> m [ExpHint] defaultExpHints e = map (const NoHint) <$> expExtType e + +-- I have no Idea if this is correct +allocInLParams :: + (Allocable fromrep torep inner) => + SubExp -> + TPrimExp Int64 VName -> + [LParam fromrep] -> + AllocM fromrep torep [LParam torep] +allocInLParams num_threads idxs = mapM alloc + where + alloc x = + case paramType x of + Array pt shape u -> do + let t = paramType x `arrayOfRow` num_threads + mem <- allocForArray t =<< askDefaultSpace + let base_dims = map pe64 $ arrayDims t + lmad_base = LMAD.iota 0 base_dims + lmad_x = + LMAD.slice lmad_base $ + fullSliceNum base_dims [DimFix idxs] + pure $ x {paramDec = MemArray pt shape u $ ArrayIn mem lmad_x} + Prim bt -> pure $ x {paramDec = MemPrim bt} + Mem space -> pure $ x {paramDec = MemMem space} + -- This next case will never happen. + Acc acc ispace ts u -> pure $ x {paramDec = MemAcc acc ispace ts u} diff --git a/src/Futhark/Pass/ExplicitAllocations/GPU.hs b/src/Futhark/Pass/ExplicitAllocations/GPU.hs --- a/src/Futhark/Pass/ExplicitAllocations/GPU.hs +++ b/src/Futhark/Pass/ExplicitAllocations/GPU.hs @@ -72,9 +72,12 @@ identitySegOpMapper { mapOnSegOpBody = localScope scope . local f . allocInKernelBody, - mapOnSegOpLambda = + mapOnSegBinOpLambda = local inThread - . allocInBinOpLambda num_threads (segSpace op) + . allocInBinOpLambda num_threads (segSpace op), + mapOnSegPostOpLambda = + local inThread + . allocInPostOpLambda num_threads (segSpace op) } f = case segLevel op of SegThread {} -> inThread diff --git a/src/Futhark/Pass/ExplicitAllocations/MC.hs b/src/Futhark/Pass/ExplicitAllocations/MC.hs --- a/src/Futhark/Pass/ExplicitAllocations/MC.hs +++ b/src/Futhark/Pass/ExplicitAllocations/MC.hs @@ -21,8 +21,10 @@ identitySegOpMapper { mapOnSegOpBody = localScope scope . allocInKernelBody, - mapOnSegOpLambda = - allocInBinOpLambda num_threads (segSpace op) + mapOnSegBinOpLambda = + allocInBinOpLambda num_threads (segSpace op), + mapOnSegPostOpLambda = + allocInPostOpLambda num_threads (segSpace op) } handleMCOp :: Op MC -> AllocM MC MCMem (Op MCMem) 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 @@ -4,6 +4,7 @@ module Futhark.Pass.ExplicitAllocations.SegOp ( allocInKernelBody, allocInBinOpLambda, + allocInPostOpLambda, ) where @@ -21,16 +22,6 @@ allocInKernelBody (Body () stms res) = uncurry (flip (Body ())) <$> collectStms (allocInStms stms (pure res)) -allocInLambda :: - (Allocable fromrep torep inner) => - [LParam torep] -> - Body fromrep -> - AllocM fromrep torep (Lambda torep) -allocInLambda params body = - mkLambda params . allocInStms (bodyStms body) $ - pure $ - bodyResult body - allocInBinOpParams :: (Allocable fromrep torep inner) => SubExp -> @@ -96,3 +87,16 @@ allocInBinOpParams num_threads index_x index_y acc_params arr_params allocInLambda (acc_params' ++ arr_params') (lambdaBody lam) + +allocInPostOpLambda :: + (Allocable fromrep torep inner) => + SubExp -> + SegSpace -> + Lambda fromrep -> + AllocM fromrep torep (Lambda torep) +allocInPostOpLambda num_threads (SegSpace flat _) lam = do + let arr_params = lambdaParams lam + index_x = TPrimExp $ LeafExp flat int64 + arr_params' <- allocInLParams num_threads index_x arr_params + + allocInLambda arr_params' (lambdaBody lam) 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 @@ -398,27 +398,35 @@ Just do_iswim <- iswim pat w scan_lam $ zip nes arrs = do types <- asksScope scopeForSOACs transformStms path . stmsToList . snd =<< runBuilderT (certifying cs do_iswim) types - | Just (scans, map_lam) <- isScanomapSOAC form = do + | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form = do let paralleliseOuter = runBuilder_ $ do scan_ops <- forM scans $ \(Scan scan_lam nes) -> do (scan_lam', nes', shape) <- determineReduceOp scan_lam nes let scan_lam'' = soacsLambdaToGPU scan_lam' pure $ SegBinOp Noncommutative scan_lam'' nes' shape let map_lam_sequential = soacsLambdaToGPU map_lam + post_op = SegPostOp $ soacsLambdaToGPU post_lam lvl <- segThreadCapped [w] "segscan" $ NoRecommendation SegNoVirt addStms . fmap (certify cs) - =<< segScan lvl pat mempty w scan_ops map_lam_sequential arrs [] [] + =<< segScan lvl pat mempty w scan_ops map_lam_sequential post_op arrs [] [] outerParallelBody = renameBody =<< (mkBody <$> paralleliseOuter <*> pure (varsRes (patNames pat))) paralleliseInner path' = do - (mapstm, scanstm) <- - scanomapToMapAndScan pat (w, scans, map_lam, arrs) + (mapstm, scanstm, poststm) <- + maposcanomapToMapScanAndMap pat (w, post_lam, scans, map_lam, arrs) types <- asksScope scopeForSOACs transformStms path' . stmsToList <=< (`runBuilderT_` types) $ - addStms =<< simplifyStms (stmsFromList [certify cs mapstm, certify cs scanstm]) + addStms + =<< simplifyStms + ( stmsFromList + [ certify cs mapstm, + certify cs scanstm, + certify cs poststm + ] + ) innerParallelBody path' = renameBody @@ -525,7 +533,7 @@ max (bodyInterest defbody) (map (bodyInterest . caseBody) cases) - | Op (Screma w _ (ScremaForm lam' _ _)) <- stmExp stm = + | Op (Screma w _ (ScremaForm lam' _ _ _)) <- stmExp stm = zeroIfTooSmall w + bodyInterest (lambdaBody lam') | Op (Stream _ _ _ lam') <- stmExp stm = bodyInterest $ lambdaBody lam' @@ -556,7 +564,7 @@ interest depth stm | "sequential" `inAttrs` attrs = 0 :: Int - | Op (Screma _ _ form@(ScremaForm lam' _ _)) <- stmExp stm, + | Op (Screma _ _ form@(ScremaForm lam' _ _ _)) <- stmExp stm, isJust $ isMapSOAC form = if sequential_inner then 0 @@ -565,7 +573,7 @@ bodyInterest (depth + 1) body * 10 | WithAcc _ withacc_lam <- stmExp stm = bodyInterest (depth + 1) (lambdaBody withacc_lam) - | Op (Screma _ _ form@(ScremaForm lam' _ _)) <- stmExp stm = + | Op (Screma _ _ form@(ScremaForm lam' _ _ _)) <- stmExp stm = 1 + bodyInterest (depth + 1) (lambdaBody lam') + @@ -809,9 +817,9 @@ onInnerMap path maploop@(MapLoop pat aux w lam arrs) acc | unbalancedLambda lam, lambdaContainsParallelism lam = - addStmToAcc (mapLoopStm maploop) acc + flip addStmToAcc acc =<< mapLoopStm maploop | otherwise = - distributeSingleStm acc (mapLoopStm maploop) >>= \case + (distributeSingleStm acc =<< mapLoopStm maploop) >>= \case Just (post_kernels, res, nest, acc') | Just (perm, pat', lam') <- removeUnusedMapResults pat res lam -> do addPostStms post_kernels 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 @@ -85,10 +85,8 @@ m (Stms rep) segRed lvl pat cs w ops map_lam arrs ispace inps = runBuilder_ $ do (kspace, kbody) <- prepareRedOrScan cs w map_lam arrs ispace inps - letBind pat $ - Op $ - segOp $ - SegRed lvl kspace (lambdaReturnType map_lam) kbody ops + letBind pat . Op . segOp $ + SegRed lvl kspace (lambdaReturnType map_lam) kbody ops segScan :: (MonadFreshNames m, DistRep rep, HasScope rep m) => @@ -98,16 +96,19 @@ SubExp -> -- segment size [SegBinOp rep] -> Lambda rep -> + SegPostOp rep -> [VName] -> [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this scan [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace m (Stms rep) -segScan lvl pat cs w ops map_lam arrs ispace inps = runBuilder_ $ do +segScan lvl pat cs w ops map_lam post_op arrs ispace inps = runBuilder_ $ do + let SegPostOp post_lam = post_op (kspace, kbody) <- prepareRedOrScan cs w map_lam arrs ispace inps - letBind pat $ - Op $ - segOp $ - SegScan lvl kspace (lambdaReturnType map_lam) kbody ops + post_lam' <- runLambdaBuilder (lambdaParams post_lam) $ do + mapM_ readKernelInput inps + bodyBind $ lambdaBody post_lam + letBind pat . Op . segOp $ + SegScan lvl kspace (lambdaReturnType map_lam) kbody ops (SegPostOp post_lam') segMap :: (MonadFreshNames m, DistRep rep, HasScope rep m) => @@ -121,10 +122,8 @@ m (Stms rep) segMap lvl pat w map_lam arrs ispace inps = runBuilder_ $ do (kspace, kbody) <- prepareRedOrScan mempty w map_lam arrs ispace inps - letBind pat $ - Op $ - segOp $ - SegMap lvl kspace (lambdaReturnType map_lam) kbody + letBind pat . Op . segOp $ + SegMap lvl kspace (lambdaReturnType map_lam) kbody dummyDim :: (MonadBuilder m) => 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 @@ -60,9 +60,9 @@ data MapLoop = MapLoop (Pat Type) (StmAux ()) SubExp (Lambda SOACS) [VName] -mapLoopStm :: MapLoop -> Stm SOACS +mapLoopStm :: (MonadFreshNames m) => MapLoop -> m (Stm SOACS) mapLoopStm (MapLoop pat aux w lam arrs) = - Let pat aux $ Op $ Screma w arrs $ mapSOAC lam + Let pat aux . Op . Screma w arrs <$> mapSOAC lam data DistEnv rep m = DistEnv { distNest :: Nestings, @@ -243,9 +243,11 @@ lambdaReturnType = map rowType $ patTypes pat } stms <- - runBuilder_ . auxing aux . FOT.transformSOAC pat $ - Screma w used_arrs $ - mapSOAC lam' + runBuilder_ + . auxing aux + . FOT.transformSOAC pat + . Screma w used_arrs + =<< mapSOAC lam' pure $ acc {distTargets = newtargets, distStms = stms} | otherwise -> do @@ -495,8 +497,8 @@ -- If the scan cannot be distributed by itself, it will be -- sequentialised in the default case for this function. maybeDistributeStm stm@(Let pat aux (Op (Screma w arrs form))) acc - | Just (scans, map_lam) <- isScanomapSOAC form, - Scan lam nes <- singleScan scans = + | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form, + Scan op_lam nes <- singleScan scans = distributeSingleStm acc stm >>= \case Just (kernels, res, nest, acc') | Just (perm, pat_unused) <- permutationAndMissing pat res -> @@ -506,7 +508,7 @@ nest' <- expandKernelNest pat_unused nest map_lam' <- soacsLambda map_lam localScope (typeEnvFromDistAcc acc') $ - segmentedScanomapKernel nest' perm (stmAuxCerts aux) w lam map_lam' nes arrs + segmentedScanomapKernel nest' perm (stmAuxCerts aux) w op_lam map_lam' post_lam nes arrs >>= kernelOrNot mempty stm acc kernels acc' _ -> addStmToAcc stm acc @@ -950,22 +952,25 @@ SubExp -> Lambda SOACS -> Lambda rep -> + Lambda SOACS -> [SubExp] -> [VName] -> DistNestT rep m (Maybe (Stms rep)) -segmentedScanomapKernel nest perm cs segment_size lam map_lam nes arrs = do +segmentedScanomapKernel nest perm cs segment_size op_lam map_lam post_lam nes arrs = do mk_lvl <- asks distSegLevel onLambda <- asks distOnSOACSLambda let onLambda' = fmap fst . runBuilder . onLambda - isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $ + isSegmentedOp nest perm (freeIn op_lam) (freeIn map_lam) nes [] $ \pat ispace inps nes' _ -> do - (lam', nes'', shape) <- determineReduceOp lam nes' - lam'' <- onLambda' lam' - let scan_op = SegBinOp Noncommutative lam'' nes'' shape + (op_lam', nes'', shape) <- determineReduceOp op_lam nes' + op_lam'' <- onLambda' op_lam' + let scan_op = SegBinOp Noncommutative op_lam'' nes'' shape + post_lam' <- onLambda' post_lam + let post_op = SegPostOp post_lam' lvl <- mk_lvl (segment_size : map snd ispace) "segscan" $ NoRecommendation SegNoVirt addStms =<< traverse renameStm - =<< segScan lvl pat cs segment_size [scan_op] map_lam arrs ispace inps + =<< segScan lvl pat cs segment_size [scan_op] map_lam post_op arrs ispace inps regularSegmentedRedomapKernel :: (MonadFreshNames m, LocalScope rep m, DistRep rep) => 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 @@ -65,8 +65,8 @@ mapM (newIdent' (<> "_transposed") . transposeIdentType) $ patIdents res_pat - addStm . Let res_pat' map_aux . Op $ - Screma map_w map_arrs' (mapSOAC map_fun') + addStm . Let res_pat' map_aux . Op . Screma map_w map_arrs' + =<< mapSOAC map_fun' forM_ (zip (patIdents res_pat) (patIdents res_pat')) $ \(to, from) -> do let perm = [1, 0] ++ [2 .. arrayRank (identType from) - 1] @@ -132,8 +132,8 @@ let map_fun' = Lambda map_params map_rettype map_body - addStm . Let res_pat map_aux . Op . Screma map_w arrs' $ - mapSOAC map_fun' + addStm . Let res_pat map_aux . Op . Screma map_w arrs' + =<< mapSOAC map_fun' | otherwise = Nothing -- | Does this reduce operator contain an inner map, and if so, what diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs --- a/src/Futhark/Pass/ExtractKernels/Interchange.hs +++ b/src/Futhark/Pass/ExtractKernels/Interchange.hs @@ -75,11 +75,10 @@ unzip . catMaybes <$> mapM copyOrRemoveParam params_and_arrs let lam = Lambda (params' <> new_params) rettype body - map_stm = - Let loop_pat_expanded aux $ - Op $ - Screma w (arrs' <> new_arrs) (mapSOAC lam) - res = varsRes $ patNames loop_pat_expanded + map_stm <- + Let loop_pat_expanded aux . Op . Screma w (arrs' <> new_arrs) + <$> mapSOAC lam + let res = varsRes $ patNames loop_pat_expanded pat' = Pat $ rearrangeShape perm $ patElems pat pure $ @@ -138,22 +137,29 @@ f (p, arg) = pure (p, arg) -manifestMaps :: [LoopNesting] -> [VName] -> Stms SOACS -> ([VName], Stms SOACS) -manifestMaps [] res stms = (res, stms) -manifestMaps (n : ns) res stms = - let (res', stms') = manifestMaps ns res stms - (params, arrs) = unzip $ loopNestingParamsAndArrs n +manifestMaps :: + (MonadFreshNames m) => + [LoopNesting] -> + [VName] -> + Stms SOACS -> + m ([VName], Stms SOACS) +manifestMaps [] res stms = pure (res, stms) +manifestMaps (n : ns) res stms = do + (res', stms') <- manifestMaps ns res stms + let (params, arrs) = unzip $ loopNestingParamsAndArrs n lam = Lambda params (map rowType $ patTypes (loopNestingPat n)) (mkBody stms' $ varsRes res') - in ( patNames $ loopNestingPat n, - oneStm $ - Let (loopNestingPat n) (loopNestingAux n) $ - Op $ - Screma (loopNestingWidth n) arrs (mapSOAC lam) - ) + st <- + oneStm + . Let (loopNestingPat n) (loopNestingAux n) + . Op + . Screma (loopNestingWidth n) arrs + <$> mapSOAC lam + pure + (patNames $ loopNestingPat n, st) -- | Given a (parallel) map nesting and an inner sequential loop, move -- the maps inside the sequential loop. The result is several @@ -185,7 +191,7 @@ else let loop_stm = seqLoopStm loop' names = rearrangeShape (loopPerm loop') (patNames (stmPat loop_stm)) - in pure $ snd $ manifestMaps ns names $ stms <> oneStm loop_stm + in snd <$> manifestMaps ns names (stms <> oneStm loop_stm) | otherwise = pure $ oneStm $ seqLoopStm loop -- | An encoding of a branch with alongside its result pattern. @@ -215,7 +221,7 @@ mkBranch branch = (renameBody =<<) $ runBodyBuilder $ do let lam = Lambda params lam_ret branch - addStm $ Let branch_pat' aux $ Op $ Screma w arrs $ mapSOAC lam + addStm . Let branch_pat' aux . Op . Screma w arrs =<< mapSOAC lam pure $ varsRes $ patNames branch_pat' cases' <- mapM (traverse mkBranch) cases @@ -262,9 +268,12 @@ let (params, arrs) = unzip params_and_arrs maplam_ret = lambdaReturnType acc_lam maplam = Lambda (iota_p : orig_acc_params ++ params) maplam_ret (lambdaBody acc_lam) - auxing map_aux . fmap subExpsRes . letTupExp' "withacc_inter" $ - Op $ - Screma w (iota_w : map paramName acc_params ++ arrs) (mapSOAC maplam) + auxing map_aux + . fmap subExpsRes + . letTupExp' "withacc_inter" + . Op + . Screma w (iota_w : map paramName acc_params ++ arrs) + =<< mapSOAC maplam let pat = Pat $ rearrangeShape perm $ patElems map_pat pure $ WithAccStm perm pat inputs' acc_lam' where diff --git a/src/Futhark/Pass/ExtractKernels/Intrablock.hs b/src/Futhark/Pass/ExtractKernels/Intrablock.hs --- a/src/Futhark/Pass/ExtractKernels/Intrablock.hs +++ b/src/Futhark/Pass/ExtractKernels/Intrablock.hs @@ -251,14 +251,15 @@ addStms =<< runDistNestT env (distributeMapBodyStms acc (bodyStms $ lambdaBody lam)) Op (Screma w arrs form) - | Just (scans, mapfun) <- isScanomapSOAC form, + | Just (post_lam, scans, mapfun) <- isMaposcanomapSOAC form, -- FIXME: Futhark.CodeGen.ImpGen.GPU.Block.compileGroupOp -- cannot handle multiple scan operators yet. Scan scanfun nes <- singleScan scans -> do let scanfun' = soacsLambdaToGPU scanfun mapfun' = soacsLambdaToGPU mapfun + post_op = SegPostOp $ soacsLambdaToGPU post_lam certifying (stmAuxCerts aux) $ - addStms =<< segScan lvl pat mempty w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' arrs [] [] + addStms =<< segScan lvl pat mempty w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' post_op arrs [] [] parallelMin [w] Op (Screma w arrs form) | Just (reds, map_lam) <- isRedomapSOAC form -> do diff --git a/src/Futhark/Pass/ExtractMulticore.hs b/src/Futhark/Pass/ExtractMulticore.hs --- a/src/Futhark/Pass/ExtractMulticore.hs +++ b/src/Futhark/Pass/ExtractMulticore.hs @@ -235,17 +235,18 @@ pure $ mconcat seq_reds_stms <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op) - | Just (scans, map_lam) <- isScanomapSOAC form = do + | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form = do (gtid, space) <- mkSegSpace w kbody <- mapLambdaToKernelBody transformBody gtid map_lam arrs (scans_stms, scans') <- mapAndUnzipM scanToSegBinOp scans + post_op <- SegPostOp <$> transformLambda post_lam pure $ mconcat scans_stms <> oneStm ( Let pat (defAux ()) $ Op $ ParOp Nothing $ - SegScan () space (lambdaReturnType map_lam) kbody scans' + SegScan () space (lambdaReturnType map_lam) kbody scans' post_op ) | otherwise = do -- This screma is too complicated for us to immediately do diff --git a/src/Futhark/Pass/LiftAllocations.hs b/src/Futhark/Pass/LiftAllocations.hs --- a/src/Futhark/Pass/LiftAllocations.hs +++ b/src/Futhark/Pass/LiftAllocations.hs @@ -137,9 +137,9 @@ liftAllocationsInSegOp (SegRed lvl sp tps body binops) = do stms <- liftAllocationsInStms (bodyStms body) pure $ SegRed lvl sp tps (body {bodyStms = stms}) binops -liftAllocationsInSegOp (SegScan lvl sp tps body binops) = do +liftAllocationsInSegOp (SegScan lvl sp tps body binops post_op) = do stms <- liftAllocationsInStms (bodyStms body) - pure $ SegScan lvl sp tps (body {bodyStms = stms}) binops + pure $ SegScan lvl sp tps (body {bodyStms = stms}) binops post_op liftAllocationsInSegOp (SegHist lvl sp tps body histops) = do stms <- liftAllocationsInStms (bodyStms body) pure $ SegHist lvl sp tps (body {bodyStms = stms}) histops diff --git a/src/Futhark/Pass/LowerAllocations.hs b/src/Futhark/Pass/LowerAllocations.hs --- a/src/Futhark/Pass/LowerAllocations.hs +++ b/src/Futhark/Pass/LowerAllocations.hs @@ -114,9 +114,9 @@ lowerAllocationsInSegOp (SegRed lvl sp tps body binops) = do stms <- lowerAllocationsInStms (bodyStms body) mempty mempty pure $ SegRed lvl sp tps (body {bodyStms = stms}) binops -lowerAllocationsInSegOp (SegScan lvl sp tps body binops) = do +lowerAllocationsInSegOp (SegScan lvl sp tps body binops post_op) = do stms <- lowerAllocationsInStms (bodyStms body) mempty mempty - pure $ SegScan lvl sp tps (body {bodyStms = stms}) binops + pure $ SegScan lvl sp tps (body {bodyStms = stms}) binops post_op lowerAllocationsInSegOp (SegHist lvl sp tps body histops) = do stms <- lowerAllocationsInStms (bodyStms body) mempty mempty pure $ SegHist lvl sp tps (body {bodyStms = stms}) histops diff --git a/src/Futhark/Script.hs b/src/Futhark/Script.hs --- a/src/Futhark/Script.hs +++ b/src/Futhark/Script.hs @@ -52,7 +52,7 @@ import Data.Void import Data.Word (Word8) import Futhark.Data.Parser qualified as V -import Futhark.Server +import Futhark.Server hiding (Record) import Futhark.Server.Values (getValue, putValue) import Futhark.Test.Values qualified as V import Futhark.Util (nubOrd) diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs --- a/src/Futhark/Tools.hs +++ b/src/Futhark/Tools.hs @@ -6,7 +6,9 @@ ( module Futhark.Construct, redomapToMapAndReduce, scanomapToMapAndScan, + maposcanomapToMapScanAndMap, dissectScrema, + extractPostLambda, sequentialStreamWholeArray, partitionChunkedFoldParameters, withAcc, @@ -22,7 +24,6 @@ import Futhark.Construct import Futhark.IR import Futhark.IR.SOACS.SOAC -import Futhark.Util -- | Turns a binding of a @redomap@ into two seperate bindings, a -- @map@ binding and a @reduce@ binding (returned in that order). @@ -45,7 +46,7 @@ redomapToMapAndReduce (Pat pes) (w, reds, map_lam, arrs) = do (map_pat, red_pat, red_arrs) <- splitScanOrRedomap pes w map_lam $ map redNeutral reds - let map_stm = mkLet map_pat $ Op $ Screma w arrs (mapSOAC map_lam) + map_stm <- mkLet map_pat . Op . Screma w arrs <$> mapSOAC map_lam red_stm <- Let red_pat (defAux ()) . Op <$> (Screma w red_arrs <$> reduceSOAC reds) @@ -67,12 +68,38 @@ scanomapToMapAndScan (Pat pes) (w, scans, map_lam, arrs) = do (map_pat, scan_pat, scan_arrs) <- splitScanOrRedomap pes w map_lam $ map scanNeutral scans - let map_stm = mkLet map_pat $ Op $ Screma w arrs (mapSOAC map_lam) + map_stm <- mkLet map_pat . Op . Screma w arrs <$> mapSOAC map_lam scan_stm <- Let scan_pat (defAux ()) . Op <$> (Screma w scan_arrs <$> scanSOAC scans) pure (map_stm, scan_stm) +maposcanomapToMapScanAndMap :: + ( MonadFreshNames m, + Buildable rep, + Op rep ~ SOAC rep + ) => + Pat (LetDec rep) -> + ( SubExp, + Lambda rep, + [Scan rep], + Lambda rep, + [VName] + ) -> + m (Stm rep, Stm rep, Stm rep) +maposcanomapToMapScanAndMap (Pat pes) (w, post_lam, scans, map_lam, arrs) = do + map_res <- mapM tempRes $ lambdaReturnType map_lam + map_stm <- mkLet map_res . Op . Screma w arrs <$> mapSOAC map_lam + scan_res <- mapM tempRes $ take (scanResults scans) $ lambdaReturnType map_lam + let (scan_arrs, map_arrs) = splitAt (scanResults scans) $ map identName map_res + scan_stm <- mkLet scan_res . Op . Screma w scan_arrs <$> scanSOAC scans + let post_arrs = map identName scan_res <> map_arrs + res = patElemIdent <$> pes + post_stm <- mkLet res . Op . Screma w post_arrs <$> mapSOAC post_lam + pure (map_stm, scan_stm, post_stm) + where + tempRes res = newIdent "temp_res" $ res `arrayOfRow` w + splitScanOrRedomap :: (Typed dec, MonadFreshNames m) => [PatElem dec] -> @@ -94,7 +121,7 @@ newIdent (baseName (patElemName pe) <> "_map_acc") $ acc_t `arrayOfRow` w arrMapPatElem = pure . patElemIdent --- | Turn a Screma into a Scanomap (possibly with mapout parts) and a +-- | Turn a Screma into a maposcanomap (possibly with mapout parts) and a -- Redomap. This is used to handle Scremas that are so complicated -- that we cannot directly generate efficient parallel code for them. -- In essense, what happens is the opposite of horisontal fusion. @@ -108,20 +135,58 @@ ScremaForm (Rep m) -> [VName] -> m () -dissectScrema pat w (ScremaForm map_lam scans reds) arrs = do +dissectScrema pat w (ScremaForm map_lam scans reds post_lam) arrs = do let num_reds = redResults reds num_scans = scanResults scans - (scan_res, red_res, map_res) = splitAt3 num_scans num_reds $ patNames pat + reds_ts = concatMap (lambdaReturnType . redLambda) reds + (red_res, scan_map_res) = splitAt num_reds $ patNames pat + (scan_pars, map_pars) = splitAt num_scans $ lambdaParams post_lam + post_res = bodyResult $ lambdaBody post_lam to_red <- replicateM num_reds $ newVName "to_red" + red_pars <- mapM (newParam "x") reds_ts + let red_post_res = paramName <$> red_pars - let scanomap = scanomapSOAC scans map_lam - letBindNames (scan_res <> to_red <> map_res) $ - Op (Screma w arrs scanomap) + let post_lam' = + post_lam + { lambdaParams = scan_pars <> red_pars <> map_pars, + lambdaBody = + (lambdaBody post_lam) + { bodyResult = varsRes red_post_res <> post_res + }, + lambdaReturnType = reds_ts <> lambdaReturnType post_lam + } + maposcanomap <- maposcanomapSOAC map_lam scans post_lam' + letBindNames (to_red <> scan_map_res) $ Op (Screma w arrs maposcanomap) + reduce <- reduceSOAC reds letBindNames red_res $ Op $ Screma w to_red reduce +-- | Remove the post lambda from a screma, producing the screma with an identity +-- post-lambda, and a new map screma that is just the post-lambda. You can apply +-- this indefinitely in case the post-lambda is already an identity lambda, so +-- be careful. +extractPostLambda :: + ( MonadBuilder m, + Op (Rep m) ~ SOAC (Rep m), + Buildable (Rep m) + ) => + Pat (LetDec (Rep m)) -> + SubExp -> + [VName] -> + ScremaForm (Rep m) -> + m () +extractPostLambda pat w arrs (ScremaForm pre_lam scans reds post_lam) = do + tmp_names <- + mapM (newVName . (<> "_extract") . baseName . paramName) (lambdaParams post_lam) + id_lam <- mkIdentityLambda $ map paramType $ lambdaParams post_lam + letBindNames (map patElemName red_res <> tmp_names) $ + Op (Screma w arrs $ ScremaForm pre_lam scans reds id_lam) + letBind (Pat nonred_res) . Op . Screma w tmp_names =<< mapSOAC post_lam + where + (red_res, nonred_res) = splitAt (redResults reds) (patElems pat) + -- | Turn a stream SOAC into statements that apply the stream lambda -- to the entire input. sequentialStreamWholeArray :: @@ -231,8 +296,11 @@ fmap subExpsRes $ forM (zip acc_ps_inner vs) $ \(acc_p_inner, v) -> letSubExp "scatter_acc" . BasicOp $ UpdateAcc Safe (paramName acc_p_inner) is [v] + let w = arraysSize 0 arrs_ts - fmap varsRes . letTupExp "acc_res" . Op $ - Screma w (map paramName acc_ps <> arrs) (mapSOAC map_lam) + (fmap varsRes . letTupExp "acc_res") + . Op + . Screma w (map paramName acc_ps <> arrs) + =<< mapSOAC map_lam letTupExp desc $ WithAcc [(acc_shape, [v], Nothing) | v <- dest] withacc_lam 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 @@ -122,26 +122,24 @@ [VName] -> ScremaForm (Rep m) -> m () -transformScrema pat w arrs form@(ScremaForm map_lam scans reds) = do +transformScrema pat w arrs form@(ScremaForm map_lam scans reds post_lam) = do -- See Note [Translation of Screma]. -- -- Start by combining all the reduction and scan parts into a single -- operator let Reduce _ red_lam red_nes = singleReduce reds Scan scan_lam scan_nes = singleScan scans - (scan_arr_ts, _red_ts, map_arr_ts) = - splitAt3 (length scan_nes) (length red_nes) $ scremaType w form + (_red_ts, post_ts) = + splitAt (length red_nes) $ scremaType w form - scan_arrs <- resultArray [] scan_arr_ts - map_arrs <- resultArray arrs map_arr_ts + post_arrs <- resultArray arrs post_ts scanacc_params <- mapM (newParam "scanacc" . flip toDecl Nonunique) $ lambdaReturnType scan_lam - scanout_params <- mapM (newParam "scanout" . flip toDecl Unique) scan_arr_ts redout_params <- mapM (newParam "redout" . flip toDecl Nonunique) $ lambdaReturnType red_lam - mapout_params <- mapM (newParam "mapout" . flip toDecl Unique) map_arr_ts + out_params <- mapM (newParam "out" . flip toDecl Unique) post_ts arr_ts <- mapM lookupType arrs - let paramForAcc (Acc c _ _ _) = find (f . paramType) mapout_params + let paramForAcc (Acc c _ _ _) = find (f . paramType) out_params where f (Acc c2 _ _ _) = c == c2 f _ = False @@ -150,9 +148,8 @@ let merge = concat [ zip scanacc_params scan_nes, - zip scanout_params $ map Var scan_arrs, zip redout_params red_nes, - zip mapout_params $ map Var map_arrs + zip out_params $ map Var post_arrs ] i <- newVName "i" let loopform = ForLoop i Int64 w @@ -198,23 +195,25 @@ map (pure . BasicOp . SubExp) $ map (Var . paramName) redout_params ++ map resSubExp red_res - -- Write the scan accumulator to the scan result arrays. - scan_outarrs <- - certifying (foldMap resCerts scan_res) $ - letwith (map paramName scanout_params) (Var i) $ - map resSubExp scan_res' + let res = scan_res' <> map_res + param_bind = resSubExp <$> res + certs = resCerts <$> res + forM_ (zip3 (paramName <$> lambdaParams post_lam) param_bind certs) $ + \(par, v, cs) -> do + certifying cs $ letBindNames [par] $ BasicOp $ SubExp v - -- Write the map results to the map result arrays. - map_outarrs <- - certifying (foldMap resCerts map_res) $ - letwith (map paramName mapout_params) (Var i) $ - map resSubExp map_res + mapM_ addStm $ bodyStms $ lambdaBody post_lam + let post_res = bodyResult $ lambdaBody post_lam + outarrs <- + certifying (foldMap resCerts post_res) $ + letwith (map paramName out_params) (Var i) $ + map resSubExp post_res + pure . concat $ [ scan_res', - varsRes scan_outarrs, red_res', - varsRes map_outarrs + varsRes outarrs ] -- We need to discard the final scan accumulators, as they are not diff --git a/src/Futhark/Util/Loc.hs b/src/Futhark/Util/Loc.hs --- a/src/Futhark/Util/Loc.hs +++ b/src/Futhark/Util/Loc.hs @@ -1,4 +1,10 @@ -- | A Safe Haskell-trusted re-export of the @srcloc@ package. -module Futhark.Util.Loc (module Data.Loc) where +module Futhark.Util.Loc (module Data.Loc, contains) where import Data.Loc + +contains :: (Located a) => a -> Pos -> Bool +contains a pos = + case locOf a of + Loc start end -> pos >= start && pos <= end + NoLoc -> False 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 @@ -1234,7 +1234,7 @@ _ -> error "Expected ModuleFun." evalDec :: Env -> Dec -> EvalM Env -evalDec env (ValDec (ValBind _ v _ (Info ret) tparams ps fbody _ _ _)) = localExts $ do +evalDec env (ValDec (ValBind _ v _ _ (Info ret) tparams ps fbody _ _ _)) = localExts $ do binding <- evalValBinding env tparams ps ret fbody sizes <- extEnv pure $ mempty {envTerm = M.singleton v binding} <> sizes 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 @@ -302,9 +302,9 @@ { let L loc (ID name) = $2 in ValSpec name $3 $5 NoInfo Nothing (srcspan $1 $>) } | val BindingBinOp TypeParams ':' TypeExp - { ValSpec $2 $3 $5 NoInfo Nothing (srcspan $1 $>) } + { ValSpec (fst $2) $3 $5 NoInfo Nothing (srcspan $1 $>) } | val '(' BindingBinOp ')' TypeParams ':' TypeExp - { ValSpec $3 $5 $7 NoInfo Nothing (srcspan $1 $>) } + { ValSpec (fst $3) $5 $7 NoInfo Nothing (srcspan $1 $>) } | TypeAbbr { TypeAbbrSpec $1 } @@ -395,43 +395,43 @@ | '=...' { binOpName $1 } | '`' QualName '`' { $2 } -BindingBinOp :: { Name } +BindingBinOp :: { (Name, SrcLoc) } : BinOp {% let (QualName qs name, loc) = $1 in do unless (null qs) $ parseErrorAt loc $ Just "Cannot use a qualified name in binding position." - pure name } - | '-' { nameFromString "-" } + pure (name, srclocOf loc) } + | '-' { (nameFromString "-", srclocOf $1) } | '!' {% parseErrorAt $1 $ Just $ "'!' is a prefix operator and cannot be used as infix operator." } -BindingId :: { (Name, Loc) } - : id { let L loc (ID name) = $1 in (name, loc) } - | '(' BindingBinOp ')' { ($2, $1) } +BindingId :: { (Name, SrcLoc) } + : id { let L loc (ID name) = $1 in (name, srclocOf loc) } + | '(' BindingBinOp ')' { $2 } Val :: { ValBindBase NoInfo Name } Val : def BindingId TypeParams FunParams maybeAscription(TypeExp) '=' Exp - { let (name, _) = $2 - in ValBind Nothing name $5 NoInfo + { let (name, loc) = $2 + in ValBind Nothing name loc $5 NoInfo $3 $4 $7 Nothing mempty (srcspan $1 $>) } | entry BindingId TypeParams FunParams maybeAscription(TypeExp) '=' Exp { let (name, loc) = $2 - in ValBind (Just NoInfo) name $5 NoInfo + in ValBind (Just NoInfo) name loc $5 NoInfo $3 $4 $7 Nothing mempty (srcspan $1 $>) } | def FunParam BindingBinOp FunParam maybeAscription(TypeExp) '=' Exp - { ValBind Nothing $3 $5 NoInfo [] [$2,$4] $7 + { ValBind Nothing (fst $3) (snd $3) $5 NoInfo [] [$2,$4] $7 Nothing mempty (srcspan $1 $>) } -- The next two for backwards compatibility. | let BindingId TypeParams FunParams maybeAscription(TypeExp) '=' Exp - { let (name, _) = $2 - in ValBind Nothing name $5 NoInfo + { let (name, loc) = $2 + in ValBind Nothing name loc $5 NoInfo $3 $4 $7 Nothing mempty (srcspan $1 $>) } | let FunParam BindingBinOp FunParam maybeAscription(TypeExp) '=' Exp - { ValBind Nothing $3 $5 NoInfo [] [$2,$4] $7 + { ValBind Nothing (fst $3) (snd $3) $5 NoInfo [] [$2,$4] $7 Nothing mempty (srcspan $1 $>) } @@ -824,7 +824,7 @@ -- Parameter patterns are slightly restricted; see #2017. ParamPat :: { PatBase NoInfo Name StructType } : id { let L loc (ID name) = $1 in Id name NoInfo (srclocOf loc) } - | '(' BindingBinOp ')' { Id $2 NoInfo (srcspan $1 $>) } + | '(' BindingBinOp ')' { Id (fst $2) NoInfo (srcspan $1 $>) } | '_' { Wildcard NoInfo (srclocOf $1) } | '(' ')' { TuplePat [] (srcspan $1 $>) } | '(' Pat ')' { PatParens $2 (srcspan $1 $>) } @@ -841,7 +841,7 @@ InnerPat :: { PatBase NoInfo Name StructType } : id { let L loc (ID name) = $1 in Id name NoInfo (srclocOf loc) } - | '(' BindingBinOp ')' { Id $2 NoInfo (srcspan $1 $>) } + | '(' BindingBinOp ')' { Id (fst $2) NoInfo (srcspan $1 $>) } | '_' { Wildcard NoInfo (srclocOf $1) } | '(' ')' { TuplePat [] (srcspan $1 $>) } | '(' Pat ')' { PatParens $2 (srcspan $1 $>) } 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 @@ -534,7 +534,7 @@ pretty (TypeParamType l name _) = "'" <> pretty l <> prettyName name instance (IsName vn, Annot f) => Pretty (ValBindBase f vn) where - pretty (ValBind entry name retdecl rettype tparams args body _ attrs _) = + pretty (ValBind entry name _ retdecl rettype tparams args body _ attrs _) = mconcat (map ((<> line) . prettyAttr) attrs) <> fun <+> align diff --git a/src/Language/Futhark/Query.hs b/src/Language/Futhark/Query.hs --- a/src/Language/Futhark/Query.hs +++ b/src/Language/Futhark/Query.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE LambdaCase #-} + -- | Facilities for answering queries about a program, such as "what -- appears at this source location", or "where is this name bound". -- The intent is that this is used as a building block for IDE-like @@ -8,27 +10,78 @@ AtPos (..), atPos, Pos (..), + allBindings, + termBindingType, + TermBinding (..), + TermBindSrc (..), + TermFunData (..), ) where import Control.Monad import Control.Monad.State -import Data.List (find) +import Data.List (find, unsnoc) import Data.Map qualified as M -import Futhark.Util.Loc (Loc (..), Pos (..)) +import Futhark.Util.Loc (Loc (..), Pos (..), contains) import Language.Futhark import Language.Futhark.Semantic import Language.Futhark.Traversals import System.FilePath.Posix qualified as Posix +type TypeAscription = TypeExp (ExpBase Info VName) VName + +data TermBindSrc + = -- | Term was bound where a pattern is allowed + TermBindPat + | -- | Term was bound with a let + TermBindLet + | -- | term was bound where only an identifier is allowed + TermBindId + | -- | term was bound inside another pattern + TermBindNested + deriving (Eq, Show) + +data TermFunData + = TermFunData + { termFunType :: StructType, + termFunRetType :: ResRetType, + termFunAscription :: Maybe TypeAscription, + -- I wanted to remove the @Maybe@ but @Loc@ always includes a @NoLoc@ case + termFunArgEnd :: Maybe Pos, + termFunTypeArgEnd :: Maybe Pos, + termFunNameEnd :: Maybe Pos, + termFunTypeParams :: [TypeParamBase VName] + } + deriving (Eq, Show) + +data TermBinding + = -- | The bound term is a size type + TermSize + | -- | Inferred Type, Ascripted type + TermVar TermBindSrc StructType (Maybe TypeAscription) + | -- | The bound term is a function + TermFun TermFunData + deriving (Eq, Show) + +termBindingType :: TermBinding -> TypeBase Size NoUniqueness +termBindingType = \case + TermSize -> Scalar (Prim (Signed Int64)) + TermVar _ t _ -> t + TermFun tf_data -> termFunType tf_data + -- | What a name is bound to. data BoundTo - = BoundTerm StructType Loc + = -- | inferred type, optional ascription from source + BoundTerm TermBinding Loc | BoundModule Loc | BoundModuleType Loc | BoundType Loc deriving (Eq, Show) +instance Located BoundTo where + locOf :: BoundTo -> Loc + locOf = boundLoc + data Def = DefBound BoundTo | DefIndirect VName deriving (Eq, Show) @@ -44,27 +97,36 @@ sizeDefs :: SizeBinder VName -> Defs sizeDefs (SizeBinder v loc) = - M.singleton v $ DefBound $ BoundTerm (Scalar (Prim (Signed Int64))) (locOf loc) + M.singleton v $ DefBound $ BoundTerm TermSize (locOf loc) -patternDefs :: Pat (TypeBase Size u) -> Defs -patternDefs (Id vn (Info t) loc) = - M.singleton vn $ DefBound $ BoundTerm (toStruct t) (locOf loc) -patternDefs (TuplePat pats _) = - mconcat $ map patternDefs pats -patternDefs (RecordPat fields _) = - mconcat $ map (patternDefs . snd) fields -patternDefs (PatParens pat _) = patternDefs pat -patternDefs (PatAttr _ pat _) = patternDefs pat -patternDefs Wildcard {} = mempty -patternDefs PatLit {} = mempty -patternDefs (PatAscription pat _ _) = - patternDefs pat -patternDefs (PatConstr _ _ pats _) = - mconcat $ map patternDefs pats +patternDefs :: + TermBindSrc -> + Pat (TypeBase Size u) -> + Defs +patternDefs bindSrc (Id vn (Info t) loc) = + M.singleton vn $ DefBound $ BoundTerm tvar (locOf loc) + where + tvar = TermVar bindSrc (toStruct t) Nothing +patternDefs _ (TuplePat pats _) = + mconcat $ map (patternDefs TermBindNested) pats +patternDefs _ (RecordPat fields _) = + mconcat $ map (patternDefs TermBindNested . snd) fields +patternDefs bindSrc (PatParens pat _) = patternDefs bindSrc pat +patternDefs bindSrc (PatAttr _ pat _) = patternDefs bindSrc pat +patternDefs _ Wildcard {} = mempty +patternDefs _ PatLit {} = mempty +patternDefs bindSrc (PatAscription (Id vn (Info t) idLoc) texp _) = + M.singleton vn $ DefBound $ BoundTerm tvar (locOf idLoc) + where + tvar = TermVar bindSrc (toStruct t) (Just texp) +patternDefs bindSrc (PatAscription pat _ _) = + patternDefs bindSrc pat +patternDefs bindSrc (PatConstr _ _ pats _) = + mconcat $ map (patternDefs bindSrc) pats typeParamDefs :: TypeParamBase VName -> Defs typeParamDefs (TypeParamDim vn loc) = - M.singleton vn $ DefBound $ BoundTerm (Scalar $ Prim $ Signed Int64) (locOf loc) + M.singleton vn $ DefBound $ BoundTerm TermSize (locOf loc) typeParamDefs (TypeParamType _ vn loc) = M.singleton vn $ DefBound $ BoundType $ locOf loc @@ -79,39 +141,81 @@ pure e' identDefs (Ident v (Info vt) vloc) = - M.singleton v $ DefBound $ BoundTerm (toStruct vt) $ locOf vloc + M.singleton v $ + DefBound $ + BoundTerm (TermVar TermBindId (toStruct vt) Nothing) $ + locOf vloc extra = case e of AppExp (LetPat sizes pat _ _ _) _ -> - foldMap sizeDefs sizes <> patternDefs pat + foldMap sizeDefs sizes <> patternDefs TermBindLet pat Lambda params _ _ _ _ -> - mconcat (map patternDefs params) - AppExp (LetFun (name, _) (tparams, params, _, Info ret, _) _ loc) _ -> + mconcat $ map (patternDefs TermBindPat) params + AppExp (LetFun (name, name_loc) (tparams, params, tasc, Info ret, _) _ loc) _ -> let name_t = funType params ret - in M.singleton name (DefBound $ BoundTerm name_t (locOf loc)) + tfun = + TermFun $ + TermFunData + { termFunType = name_t, + termFunRetType = ret, + termFunAscription = tasc, + termFunArgEnd = start_pos, + termFunNameEnd = name_end, + termFunTypeParams = tparams, + termFunTypeArgEnd = do + (_, last_type_arg) <- unsnoc tparams + case locOf last_type_arg of + NoLoc -> Nothing + Loc _ end -> Just end + } + name_end = case locOf name_loc of + NoLoc -> Nothing + Loc _ end_pos -> Just end_pos + start_pos = do + (_, last_arg) <- unsnoc params + case locOf last_arg of + NoLoc -> Nothing + Loc _ end -> Just end + in M.singleton name (DefBound $ BoundTerm tfun (locOf loc)) <> mconcat (map typeParamDefs tparams) - <> mconcat (map patternDefs params) + <> mconcat (map (patternDefs TermBindPat) params) AppExp (LetWith v _ _ _ _ _) _ -> identDefs v AppExp (Loop _ merge _ form _ _) _ -> - patternDefs merge + patternDefs TermBindPat merge <> case form of For i _ -> identDefs i - ForIn pat _ -> patternDefs pat + ForIn pat _ -> patternDefs TermBindLet pat While {} -> mempty _ -> mempty valBindDefs :: ValBind -> Defs valBindDefs vbind = - M.insert (valBindName vbind) (DefBound $ BoundTerm vbind_t (locOf vbind)) $ + M.insert (valBindName vbind) (DefBound $ BoundTerm term_fun (locOf vbind)) $ mconcat (map typeParamDefs (valBindTypeParams vbind)) - <> mconcat (map patternDefs (valBindParams vbind)) + <> mconcat (map (patternDefs TermBindPat) (valBindParams vbind)) <> expDefs (valBindBody vbind) where - vbind_t = - funType (valBindParams vbind) $ unInfo $ valBindRetType vbind + term_fun = + TermFun $ + TermFunData + { termFunType = + funType (valBindParams vbind) $ unInfo $ valBindRetType vbind, + termFunRetType = unInfo $ valBindRetType vbind, + termFunAscription = valBindRetDecl vbind, + termFunArgEnd = args_end, + termFunNameEnd = name_end_pos, + termFunTypeParams = valBindTypeParams vbind, + termFunTypeArgEnd = do + (_, last_type_arg) <- unsnoc $ valBindTypeParams vbind + locPos $ locOf last_type_arg + } + args_end = locPos . locOf . snd =<< unsnoc (valBindParams vbind) + name_end_pos = locPos . locOf . valBindNameLoc $ vbind + locPos NoLoc = Nothing + locPos (Loc _ e) = Just e typeBindDefs :: TypeBind -> Defs typeBindDefs tbind = @@ -151,8 +255,9 @@ specDefs :: Spec -> Defs specDefs spec = case spec of - ValSpec v tparams _ (Info t) _ loc -> - let vdef = DefBound $ BoundTerm t (locOf loc) + ValSpec v tparams texp (Info t) _ loc -> + let vdef = DefBound $ BoundTerm tval (locOf loc) + tval = TermVar TermBindId t (Just texp) in M.insert v vdef $ mconcat (map typeParamDefs tparams) TypeAbbrSpec tbind -> typeBindDefs tbind TypeSpec _ v _ _ loc -> @@ -197,12 +302,6 @@ forward (DefIndirect v) = forward =<< M.lookup v defs data RawAtPos = RawAtName (QualName VName) Loc - -contains :: (Located a) => a -> Pos -> Bool -contains a pos = - case locOf a of - Loc start end -> pos >= start && pos <= end - NoLoc -> False atPosInTypeExp :: TypeExp Exp VName -> Pos -> Maybe RawAtPos atPosInTypeExp te pos = 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 @@ -1104,6 +1104,8 @@ -- may refer to abstract types that are no longer in scope. valBindEntryPoint :: Maybe (f EntryPoint), valBindName :: vn, + -- | Location of the name of this binding itself. + valBindNameLoc :: SrcLoc, valBindRetDecl :: Maybe (TypeExp (ExpBase f vn) vn), -- | If 'valBindParams' is null, then the 'retDims' are brought -- into scope at this point. 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 @@ -675,7 +675,7 @@ checkValBind :: ValBindBase NoInfo Name -> TypeM (Env, ValBind) checkValBind vb = do - (ValBind entry fname maybe_tdecl NoInfo tparams params body doc attrs loc) <- + (ValBind entry fname fname_loc maybe_tdecl NoInfo tparams params body doc attrs loc) <- resolveValBind vb top_level <- atTopLevel @@ -693,7 +693,7 @@ Just _ -> checkEntryPoint loc tparams' params' rettype _ -> pure () - let vb' = ValBind entry' fname maybe_tdecl' (Info rettype) tparams' params' body' doc attrs' loc + let vb' = ValBind entry' fname fname_loc maybe_tdecl' (Info rettype) tparams' params' body' doc attrs' loc pure ( mempty { envVtable = diff --git a/src/Language/Futhark/TypeChecker/Modules.hs b/src/Language/Futhark/TypeChecker/Modules.hs --- a/src/Language/Futhark/TypeChecker/Modules.hs +++ b/src/Language/Futhark/TypeChecker/Modules.hs @@ -500,7 +500,12 @@ abs_name_substs = M.map (qualLeaf . fst) abs_substs pmod_substs <- matchMods p_abs_subst_to_type quals sig_pmod mod_pmod loc mod_substs <- matchMTys' abs_subst_to_type quals mod_mod sig_mod loc - pure (pmod_substs <> mod_substs <> abs_name_substs) + pure $ + -- Avoid including substitutions that refer to abstract types, as this + -- results in circular substitutions (#2407). + M.filterWithKey (\k _ -> qualName k `M.notMember` mod_abs) pmod_substs + <> mod_substs + <> abs_name_substs matchEnvs :: M.Map VName (Subst StructRetType) -> diff --git a/src/Language/Futhark/TypeChecker/Names.hs b/src/Language/Futhark/TypeChecker/Names.hs --- a/src/Language/Futhark/TypeChecker/Names.hs +++ b/src/Language/Futhark/TypeChecker/Names.hs @@ -494,7 +494,7 @@ -- | Resolve names in a value binding. If this succeeds, then it is -- guaranteed that all names references things that are in scope. resolveValBind :: ValBindBase NoInfo Name -> TypeM (ValBindBase NoInfo VName) -resolveValBind (ValBind entry fname ret NoInfo tparams params body doc attrs loc) = do +resolveValBind (ValBind entry fname fname_loc ret NoInfo tparams params body doc attrs loc) = do attrs' <- mapM resolveAttrInfo attrs checkForDuplicateNames tparams params checkDoNotShadow loc fname @@ -504,4 +504,4 @@ body' <- resolveExp body bindSpaced1 Term fname loc $ \fname' -> do usedName fname' - pure $ ValBind entry fname' ret' NoInfo tparams' params' body' doc attrs' loc + pure $ ValBind entry fname' fname_loc ret' NoInfo tparams' params' body' doc attrs' loc diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -764,22 +764,24 @@ StructType -> StructType -> TermTypeM StructType -updateFieldPath src [] ve_t src_leaf_t = do - (src_leaf_t', _) <- allDimsFreshInType usage Nonrigid "any" src_leaf_t - onFailure (CheckingRecordUpdate [] src_leaf_t' ve_t) $ - unify usage src_leaf_t' ve_t - pure ve_t +updateFieldPath src all_fs ve_t = recurse [] all_fs where - usage = mkUsage (locOf src) "record update" -updateFieldPath src (f : fs) ve_t (Scalar (Record m)) - | Just f_t <- M.lookup f m = do - f_t' <- updateFieldPath src fs ve_t f_t - pure $ Scalar $ Record $ M.insert f f_t' m -updateFieldPath src _ _ _ = - typeError (locOf src) mempty . withIndexLink "record-type-not-known" $ - "Full type of" - indent 2 (pretty src) - textwrap " is not known at this point. Add a type annotation to the original record to disambiguate." + recurse seen [] t = do + (t', _) <- allDimsFreshInType usage Nonrigid "any" t + onFailure (CheckingRecordUpdate seen t' ve_t) $ + unify usage t' ve_t + pure ve_t + where + usage = mkUsage (locOf src) "record update" + recurse seen (f : fs) (Scalar (Record m)) + | Just f_t <- M.lookup f m = do + f_t' <- recurse (seen ++ [f]) fs f_t + pure $ Scalar $ Record $ M.insert f f_t' m + recurse _ _ _ = + typeError (locOf src) mempty . withIndexLink "record-type-not-known" $ + "Full type of" + indent 2 (pretty src) + textwrap " is not known at this point. Add a type annotation to the original record to disambiguate." checkUpdateSteps :: SrcLoc -> @@ -1573,12 +1575,12 @@ closeOver (k, _) | k `elem` map typeParamName tparams = pure Nothing - closeOver (k, NoConstraint l usage) = - pure $ Just $ Left $ TypeParamType l k $ srclocOf usage - closeOver (k, ParamType l loc) = - pure $ Just $ Left $ TypeParamType l k $ srclocOf loc - closeOver (k, Size Nothing usage) = - pure $ Just $ Left $ TypeParamDim k $ srclocOf usage + closeOver (k, NoConstraint l _) = + pure $ Just $ Left $ TypeParamType l k mempty + closeOver (k, ParamType l _) = + pure $ Just $ Left $ TypeParamType l k mempty + closeOver (k, Size Nothing _) = + pure $ Just $ Left $ TypeParamDim k mempty closeOver (k, UnknownSize _ _) | k `S.member` param_sizes, k `S.notMember` produced_sizes = do