diff --git a/mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.h b/mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.h --- a/mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.h +++ b/mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.h @@ -33,12 +33,12 @@ namespace transform { namespace gpu { -/// Searches `scf.foreach_thread` ops nested under `target` and maps each such +/// Searches `scf.forall` ops nested under `target` and maps each such /// op to GPU threads. Mapping is one-to-one and the induction variables of -/// `scf.foreach_thread` are rewritten to gpu.thread_id according to the -/// thread_dim_apping attribute. Sibling `scf.foreach_thread` are supported in +/// `scf.forall` are rewritten to gpu.thread_id according to the +/// thread_dim_apping attribute. Sibling `scf.forall` are supported in /// which case, the union of the number of threads is computed and may result in -/// predication. Dynamic, `scf.foreach_thread` trip counts are currently not +/// predication. Dynamic, `scf.forall` trip counts are currently not /// supported. Dynamic block dim sizes are currently not supported. DiagnosedSilenceableFailure mapNestedForeachToThreadsImpl( RewriterBase &rewriter, Operation *target, @@ -46,24 +46,22 @@ std::optional transformOp, const ArrayRef &threadMappingAttributes); -/// Maps the top level `scf.foreach_thread` op to GPU Thread Blocks. Mapping is -/// one-to-one and the induction variables of `scf.foreach_thread` are rewritten +/// Maps the top level `scf.forall` op to GPU Thread Blocks. Mapping is +/// one-to-one and the induction variables of `scf.forall` are rewritten /// to gpu.block_id according to the thread_dim_apping attribute. Dynamic, -/// `scf.foreach_thread` trip counts are currently not supported. Dynamic block +/// `scf.forall` trip counts are currently not supported. Dynamic block /// dim sizes are currently not supported. DiagnosedSilenceableFailure mapForeachToBlocksImpl( - RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp, - function_ref &)> + RewriterBase &rewriter, scf::ForallOp foreachThreadOp, + function_ref &)> blockIdGenerator, SmallVectorImpl &gridDims, TransformOpInterface transformOp, const ArrayRef &mappingAttributes); -/// Finds the top level scf::ForeachThreadOp of given target. +/// Finds the top level scf::ForallOp of given target. DiagnosedSilenceableFailure -findTopLevelForeachThreadOp(Operation *target, - scf::ForeachThreadOp &topLevelForeachThreadOp, - TransformOpInterface transformOp); +findTopLevelForallOp(Operation *target, scf::ForallOp &topLevelForallOp, + TransformOpInterface transformOp); } // namespace gpu } // namespace transform diff --git a/mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.td b/mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.td --- a/mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.td +++ b/mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.td @@ -22,30 +22,30 @@ TransformEachOpTrait, TransformOpInterface]> { let description = [{ - Target the `gpu.launch op` and rewrite all `scf.foreach_thread` + Target the `gpu.launch op` and rewrite all `scf.forall` nested in it to distributed `gpu.thread_id` attribute. - The operation searches for `scf.foreach_thread` ops nested under `target` + The operation searches for `scf.forall` ops nested under `target` and maps each such op to GPU threads. Mapping is one-to-one and the - induction variables of `scf.foreach_thread` are rewritten to + induction variables of `scf.forall` are rewritten to `gpu.thread_id` according to the `mapping` attribute. - Sibling `scf.foreach_thread` are supported in which case, the union of + Sibling `scf.forall` are supported in which case, the union of the number of threads is computed and may result in predication. - Multiple scf.foreach_thread are supported per `gpu.launch` in which case, + Multiple scf.forall are supported per `gpu.launch` in which case, the max of all the threads is computed and taken for the global - `gpu.thread_id`. If necessary, `scf.foreach_thread` that do not use the + `gpu.thread_id`. If necessary, `scf.forall` that do not use the whole thread range result in predicated computations. - Dynamic `scf.foreach_thread` trip counts are currently not supported. + Dynamic `scf.forall` trip counts are currently not supported. Dynamic block dim sizes are currently not supported. - Only **bufferized** `scf.foreach_thread` are currently supported. - Only `scf.foreach_thread` distributed to **at most 3 dimensions** are + Only **bufferized** `scf.forall` are currently supported. + Only `scf.forall` distributed to **at most 3 dimensions** are currently supported. - Barriers are inserted after each scf.foreach_thread op for now. + Barriers are inserted after each scf.forall op for now. The operation alters the block size of the given gpu_launch using blockDim argument. @@ -54,15 +54,15 @@ This operation ignores non-gpu_launch ops and drops them in the return. - If any scf.foreach_thread with tensors is found, the transform definitely + If any scf.forall with tensors is found, the transform definitely fails. - If all the scf.foreach_thread operations with gpu.thread mapping contained + If all the scf.forall operations with gpu.thread mapping contained within the LaunchOp referred to by the `target` PDLOperation lower to GPU properly, the transform succeeds. Otherwise the transform definitely fails. - scf.foreach_thread operations with mappings other than gpu.thread are + scf.forall operations with mappings other than gpu.thread are ignored. The returned handle points to the same LaunchOp operand, consuming it and @@ -74,10 +74,10 @@ ``` gpu.launch blocks(%bx, %by, %bz) in (%x = %0, %y = %1, %z = %2) threads(%tx, %ty, %tz) in (%tx = %3, %ty = %4, %tz = %5) { - scf.foreach_thread (%i, %j) in (7, 9) { + scf.forall (%i, %j) in (7, 9) { ... // body 1 } {mapping = [#gpu.thread, #gpu.thread, #gpu.thread]} - scf.foreach_thread (%i) in (12) { + scf.forall (%i) in (12) { ... // body 2 } {mapping = [#gpu.thread]} gpu.terminator @@ -125,21 +125,21 @@ TransformOpInterface, TransformEachOpTrait]> { let description = [{ - Target the gpu_launch op and rewrite the top level `scf.foreach_thread` + Target the gpu_launch op and rewrite the top level `scf.forall` to distributed gpu.block_id attribute. If `generate_gpu_launch` attribute is set, then first generates `gpu_launch` and moves the top level - `scf.foreach_thread` inside. + `scf.forall` inside. - The operation searches top level `scf.foreach_thread` ops under + The operation searches top level `scf.forall` ops under `gpu_launch` and maps each such op to GPU blocks. Mapping is - one-to-one and the induction variables of `scf.foreach_thread` are + one-to-one and the induction variables of `scf.forall` are rewritten to gpu.block_id according to the `thread_dim_mapping` attribute. - Dynamic, `scf.foreach_thread` trip counts are currently not supported. + Dynamic, `scf.forall` trip counts are currently not supported. Dynamic block dim sizes are currently not supported. - Only **bufferized** scf.foreach_thread are currently supported. - Only scf.foreach_thread distributed to **at most 3 dimensions** are + Only **bufferized** scf.forall are currently supported. + Only scf.forall distributed to **at most 3 dimensions** are currently supported. The operation alters the block size of the given gpu_launch using @@ -149,10 +149,10 @@ This operation ignores non-gpu_launch ops and drops them in the return. - If any scf.foreach_thread with tensors is found, the transform definitely + If any scf.forall with tensors is found, the transform definitely fails. - If all the scf.foreach_thread operations contained within the LaunchOp + If all the scf.forall operations contained within the LaunchOp referred to by the `target` PDLOperation lower to GPU properly, the transform succeeds. Otherwise the transform definitely fails. diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h --- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h +++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h @@ -78,8 +78,8 @@ /// Return true if `linalgOp` contains an embedded gemm subcomputation. bool containsMostMinorGemm(linalg::LinalgOp linalgOp); -/// Implementation of tiling operations using `scf.foreach_thread`. -DiagnosedSilenceableFailure tileToForeachThreadOpImpl( +/// Implementation of tiling operations using `scf.forall`. +DiagnosedSilenceableFailure tileToForallOpImpl( RewriterBase &rewriter, transform::TransformState &state, TransformOpInterface transformOp, ArrayRef targets, ArrayRef mixedNumThreads, diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td --- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td +++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td @@ -196,8 +196,8 @@ [DenseArrayNonNegative]>:$iterator_interchange); let results = (outs PDL_Operation:$transformed); - let assemblyFormat = [{ - $target + let assemblyFormat = [{ + $target (`iterator_interchange` `=` $iterator_interchange^)? attr-dict }]; let hasVerifier = 1; @@ -256,7 +256,7 @@ TransformEachOpTrait, TransformOpInterface]> { let description = [{ - Lower a tensor.unpack into empty + linalg.transpose + tensor.collapse_shape + + Lower a tensor.unpack into empty + linalg.transpose + tensor.collapse_shape + tensor.extract_slice. #### Return modes @@ -274,7 +274,7 @@ Transform_ConcreteOpType<"linalg.transpose">:$transpose_op, Transform_ConcreteOpType<"tensor.collapse_shape">:$collapse_shape_op, Transform_ConcreteOpType<"tensor.extract_slice">:$extract_slice_op); - let assemblyFormat = [{ + let assemblyFormat = [{ $target attr-dict `:` functional-type(operands, results) }]; @@ -445,7 +445,7 @@ let description = [{ Pack a LinalgOp by applying a data tiling transformation on the op and packing the operands according to the `packed_sizes` specification. - + Iterator dimensions are tiled in their canonical order in the op spec. Operands are packed according to the same canonical order of the op iterator dimensions. @@ -480,15 +480,15 @@ // affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d4, d5)> // M N m n // affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)> - %0 = linalg.generic_representing_some_higher_d_matmul + %0 = linalg.generic_representing_some_higher_d_matmul ins(%A, %B: tensor, tensor) outs( %C: tensor) ``` In particular, note that the second operand `B` has shape `KxNxnxk` (and not `KxNxkxn` as one could expect by looking **only** at the operand). - Other layouts can be obtained unsurprisingly from this canonical - transformation by composing the resulting operation with a (future) + Other layouts can be obtained unsurprisingly from this canonical + transformation by composing the resulting operation with a (future) `transform.structured.pack_transpose` op. This composition allows separating concerns and composes better compared to adding additional permutation attributes to this transform op. @@ -507,7 +507,7 @@ DefaultValuedAttr:$static_packed_sizes); let results = (outs TransformHandleTypeInterface:$packed_op); let assemblyFormat = [{ - $target + $target `packed_sizes` `=` custom($packed_sizes, $static_packed_sizes) attr-dict @@ -610,7 +610,7 @@ MemoryEffectsOpInterface, DeclareOpInterfaceMethods]> { let description = [{ - Apply a transposition to a single `tensor.pack` (resp. `tensor.unpack`) and + Apply a transposition to a single `tensor.pack` (resp. `tensor.unpack`) and update the `linalg.generic` op that consumes (resp. produces) the operation. This transform allows composing a simple `structured.pack` with additional @@ -622,7 +622,7 @@ the specified `tensor.pack` or `tensor.unpack` op. If the `target` of this op is a `tensor.pack` then a new `tensor.empty` will - be created along with transposed versions of the `tensor.pack` and the + be created along with transposed versions of the `tensor.pack` and the consuming `linalg.generic`, which is expected to be the sole consumer. If the `target` of this op is a `tensor.unpack` then the whole pack / compute @@ -642,7 +642,7 @@ This operation returns 3 handles, one to the transformed LinalgOp, one to the transformed `tensor.pack` and one to the transformed `tensor.unpack`. - The last handle for `tensor.unpack` is empty if `target_pack_or_unpack_op` + The last handle for `tensor.unpack` is empty if `target_pack_or_unpack_op` was not itself a `tensor.unpack`. }]; @@ -1106,7 +1106,7 @@ ``` }]; - // TODO: support mixed static-dynamic (see TileToForeachThreadOp). + // TODO: support mixed static-dynamic (see TileToForallOp). let arguments = (ins PDL_Operation:$target, DefaultValuedAttr:$tile_sizes); let results = (outs PDL_Operation:$for_op, @@ -1134,31 +1134,31 @@ } //===----------------------------------------------------------------------===// -// TileReductionUsingForeachThreadOp +// TileReductionUsingForallOp //===----------------------------------------------------------------------===// -def TileReductionUsingForeachThreadOp : - Op { let description = [{ - Tile a PartialReductionOpInterface op to a tiled `scf.foreach_thread` doing + Tile a PartialReductionOpInterface op to a tiled `scf.forall` doing partial reduction. This transformation tiles the `target` along the reduction dimensions. It creates a tensor initialized with the identity value. Then it creates a - `scf.foreach_thread` loops with the number threads given by `num_threads`. + `scf.forall` loops with the number threads given by `num_threads`. The op is tiled op with a size equal to `floordiv(size, num_threads)`. All the partial reduction value is are parallel inserted to create a new tensor. After the loop a merge operation is created to do a final reduction with the partial reductions tensor. If an extra `tile_sizes` parameter is passed the tiles are cyclically - distributed on the threads of the `scf.foreach_threads` loop. + distributed on the threads of the `scf.foralls` loop. #### Return modes This 4 returned handles point to: - - the parent foreach_thread op, + - the parent forall op, - the fill op used to initialize the neutral element, - the parallel tiled op and - the result-combining op. @@ -1183,7 +1183,7 @@ ``` %0 = tensor.empty(%dim_1) : tensor %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor) -> tensor - %2 = scf.foreach_thread (%arg2) in (%c5) shared_outs(%arg3 = %1) -> (tensor) { + %2 = scf.forall (%arg2) in (%c5) shared_outs(%arg3 = %1) -> (tensor) { %4 = affine.min #map(%arg2)[%dim_0] %5 = affine.max #map1(%4) %extracted_slice = tensor.extract_slice %arg3[0, %arg2] [%dim, 1] [1, 1] : tensor to tensor @@ -1195,7 +1195,7 @@ %9 = arith.addf %in, %out : f32 linalg.yield %9 : f32 } -> tensor - scf.foreach_thread.perform_concurrently { + scf.forall.perform_concurrently { tensor.parallel_insert_slice %7 into %arg3[0, %arg2] [%dim, 1] [1, 1] : tensor into tensor } } {mapping = []} @@ -1207,12 +1207,12 @@ ``` }]; - // TODO: support mixed static-dynamic (see TileToForeachThreadOp). + // TODO: support mixed static-dynamic (see TileToForallOp). let arguments = (ins PDL_Operation:$target, DefaultValuedAttr:$num_threads, DefaultValuedAttr:$tile_sizes, OptionalAttr:$mapping); - let results = (outs PDL_Operation:$foreach_thread_op, + let results = (outs PDL_Operation:$forall_op, PDL_Operation:$fill_op, PDL_Operation:$split_linalg_op, PDL_Operation:$combining_linalg_op); @@ -1321,16 +1321,16 @@ } //===----------------------------------------------------------------------===// -// TileToForeachThreadOp +// TileToForallOp //===----------------------------------------------------------------------===// -def TileToForeachThreadOp : - Op, TransformOpInterface]> { let description = [{ - Tile a TilingInterface op to a tiled `scf.foreach_thread`. + Tile a TilingInterface op to a tiled `scf.forall`. Tiling is applied by either specifying `num_threads` or `tile_size`. If `num_threads` is specified, then the tile size for each dimension `i` is @@ -1347,7 +1347,7 @@ e.g. in the Linalg case). If non-empty, the `mapping` is added as an attribute to the - resulting `scf.foreach_thread`. + resulting `scf.forall`. Note: `tile_sizes` and `num_threads` are variadic. Each tile size/number of threads can be an index attribute or a transform handle that is mapped to @@ -1366,14 +1366,14 @@ tiled operations, which can all be empty. These two returned handles point to: - - the new scf.foreach_thread op, + - the new scf.forall op, - the tiled op that implements TilingInterface. #### Example using `num_threads` ``` %0 = pdl_match @match_matmul in %arg1 - %3:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 20] + %3:2 = transform.structured.tile_to_forall_op %0 num_threads [10, 20] ``` #### Example using `tile_sizes` @@ -1381,7 +1381,7 @@ ``` %0 = pdl_match @match_matmul in %arg1 %sz = pdl_match @match_size_op in %arg1 - %3:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [0, %sz, 20] + %3:2 = transform.structured.tile_to_forall_op %0 tile_sizes [0, %sz, 20] ``` }]; @@ -1393,7 +1393,7 @@ DefaultValuedOptionalAttr:$static_num_threads, DefaultValuedOptionalAttr:$static_tile_sizes, OptionalAttr:$mapping); - let results = (outs PDL_Operation:$foreach_thread_op, + let results = (outs PDL_Operation:$forall_op, PDL_Operation:$tiled_op); let builders = [ diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h --- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h @@ -478,10 +478,10 @@ computeStaticMultiTileSizes(LinalgOp op, unsigned dimension, int64_t targetSize, int64_t divisor); -/// Rewrite a TilingInterface `op` to a tiled `scf.foreach_thread`, applying +/// Rewrite a TilingInterface `op` to a tiled `scf.forall`, applying /// tiling by `numThreads`. /// If non-empty, the `mapping` is added as an attribute to the -/// resulting `scf.foreach_thread`. +/// resulting `scf.forall`. /// Zero tile sizes indicate that the dimension is not tiled, and can be /// thought of as tiling by the full size of data. It is the user's /// responsibility to ensure that `numThreads` is a valid tiling specification @@ -491,16 +491,16 @@ Operation *tiledOp; }; FailureOr -tileToForeachThreadOp(RewriterBase &builder, TilingInterface op, - ArrayRef numThreads, - std::optional mapping); +tileToForallOp(RewriterBase &builder, TilingInterface op, + ArrayRef numThreads, + std::optional mapping); -/// Same as `tileToForeachThreadOp`, but calculate the number of threads +/// Same as `tileToForallOp`, but calculate the number of threads /// required using the given tileSizes. FailureOr -tileToForeachThreadOpUsingTileSizes(RewriterBase &builder, TilingInterface op, - ArrayRef tileSizes, - std::optional mapping); +tileToForallOpUsingTileSizes(RewriterBase &builder, TilingInterface op, + ArrayRef tileSizes, + std::optional mapping); /// Transformation information returned after reduction tiling. struct ForeachThreadReductionTilingResult { @@ -510,8 +510,8 @@ Operation *mergeOp; /// The op initializing the tensor used for partial reductions. Operation *initialOp; - /// The `scf.foreach_thread` operation that iterate over the tiles. - scf::ForeachThreadOp loops; + /// The `scf.forall` operation that iterate over the tiles. + scf::ForallOp loops; }; /// Method to tile a reduction to parallel iterations computing partial @@ -527,7 +527,7 @@ /// /// ```mlir /// %0 = linalg.fill ... : tensor<7x4xf32> -/// %1 = scf.foreach_thread (%iv) in (%c4) shared_outs(%arg0 = %0) +/// %1 = scf.forall (%iv) in (%c4) shared_outs(%arg0 = %0) /// -> (tensor<7x4xf32>) { /// %2 = tensor.extract_slice %arg3 : tensor<7x4xf32> to tensor<7xf32> /// %3 = tensor.extract_slice %in : tensor<7x9xf32> -> tensor<7x?xf32> diff --git a/mlir/include/mlir/Dialect/SCF/IR/DeviceMappingInterface.td b/mlir/include/mlir/Dialect/SCF/IR/DeviceMappingInterface.td --- a/mlir/include/mlir/Dialect/SCF/IR/DeviceMappingInterface.td +++ b/mlir/include/mlir/Dialect/SCF/IR/DeviceMappingInterface.td @@ -30,7 +30,7 @@ can be used by the device-specific code generators and the desired regions can be connected to the given processing unit. - Currently, `scf.foreach_thread` uses this interface to express the mapping + Currently, `scf.forall` uses this interface to express the mapping of the loops it contains to the GPU's parallelism units such as threads and thread blocks. }]; diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCF.h b/mlir/include/mlir/Dialect/SCF/IR/SCF.h --- a/mlir/include/mlir/Dialect/SCF/IR/SCF.h +++ b/mlir/include/mlir/Dialect/SCF/IR/SCF.h @@ -52,9 +52,9 @@ /// value is not an induction variable, then return nullptr. ParallelOp getParallelForInductionVarOwner(Value val); -/// Returns the ForeachThreadOp parent of an thread index variable. +/// Returns the ForallOp parent of an thread index variable. /// If the provided value is not a thread index variable, then return nullptr. -ForeachThreadOp getForeachThreadOpThreadIndexOwner(Value val); +ForallOp getForallOpThreadIndexOwner(Value val); /// Return true if ops a and b (or their ancestors) are in mutually exclusive /// regions/blocks of an IfOp. diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td --- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td +++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td @@ -348,18 +348,18 @@ } //===----------------------------------------------------------------------===// -// ForeachThreadOp +// ForallOp //===----------------------------------------------------------------------===// -def ForeachThreadOp : SCF_Op<"foreach_thread", [ +def ForallOp : SCF_Op<"forall", [ AttrSizedOperandSegments, - SingleBlockImplicitTerminator<"scf::PerformConcurrentlyOp">, + SingleBlockImplicitTerminator<"scf::InParallelOp">, RecursiveMemoryEffects, AutomaticAllocationScope, ]> { let summary = "evaluate a block multiple times in parallel"; let description = [{ - `scf.foreach_thread` is a target-independent multi-dimensional parallel + `scf.forall` is a target-independent multi-dimensional parallel region application operation. It has exactly one block that represents the parallel body and it takes index operands that indicate how many parallel instances of that function are created. @@ -389,17 +389,17 @@ the op is lowered to, or to ignore it when the specification is ill-formed or unsupported for a particular target. - The only allowed terminator is `scf.foreach_thread.perform_concurrently`. - `scf.foreach_thread` returns one value per `shared_out` operand. The - actions of the `perform_concurrently` terminators specify how to combine the + The only allowed terminator is `scf.forall.in_parallel`. + `scf.forall` returns one value per `shared_out` operand. The + actions of the `in_parallel` terminators specify how to combine the partial results of all parallel invocations into a full value, in some unspecified order. The "destination" of each such op must be a `shared_out` - block argument of the `scf.foreach_thread` op. + block argument of the `scf.forall` op. The actions involved in constructing the return values are further described by `tensor.parallel_insert_slice`. - `scf.foreach_thread` acts as an implicit synchronization point. + `scf.forall` acts as an implicit synchronization point. When the parallel function body has side effects, their order is unspecified across threads. @@ -410,7 +410,7 @@ // // Sequential context. // - %matmul_and_pointwise:2 = scf.foreach_thread (%thread_id_1, %thread_id_2) in + %matmul_and_pointwise:2 = scf.forall (%thread_id_1, %thread_id_2) in (%num_threads_1, %numthread_id_2) shared_outs(%o1 = %C, %o2 = %pointwise) -> (tensor, tensor) { // @@ -429,11 +429,11 @@ tensor to tensor %sE = add ins(%spointwise) outs(%sD) - scf.foreach_thread.perform_concurrently { - scf.foreach_thread.parallel_insert_slice %sD into %o1[h((%thread_id_1, %thread_id_2))]: + scf.forall.in_parallel { + scf.forall.parallel_insert_slice %sD into %o1[h((%thread_id_1, %thread_id_2))]: tensor into tensor - scf.foreach_thread.parallel_insert_slice %spointwise into %o2[i((%thread_id_1, %thread_id_2))]: + scf.forall.parallel_insert_slice %spointwise into %o2[i((%thread_id_1, %thread_id_2))]: tensor into tensor } } @@ -449,14 +449,14 @@ // Sequential context. Here `mapping` is expressed as GPU thread mapping // attributes // - %matmul_and_pointwise:2 = scf.foreach_thread (%thread_id_1, %thread_id_2) in + %matmul_and_pointwise:2 = scf.forall (%thread_id_1, %thread_id_2) in (%num_threads_1, %numthread_id_2) shared_outs(...) -> (tensor, tensor) { // // Parallel context, each thread with id = **(%thread_id_2, %thread_id_1)** // runs its version of the code. // - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { ... } } { mapping = [#gpu.thread, #gpu.thread] } @@ -470,9 +470,9 @@ ```mlir %t0 = ... %t1 = ... - %r = scf.foreach_thread ... shared_outs(%o = t0) -> tensor { + %r = scf.forall ... shared_outs(%o = t0) -> tensor { // %t0 and %t1 are privatized. %t0 is definitely copied for each thread - // because the scf.foreach_thread op's %t0 use bufferizes to a memory + // because the scf.forall op's %t0 use bufferizes to a memory // write. In the absence of other conflicts, %t1 is copied only if there // are uses of %t1 in the body that bufferize to a memory read and to a // memory write. @@ -481,9 +481,12 @@ } ``` }]; - let arguments = (ins Variadic:$num_threads, - Variadic:$outputs, - OptionalAttr:$mapping); + let arguments = (ins + Variadic:$lowerBound, + Variadic:$upperBound, + Variadic:$step, + Variadic:$outputs, + OptionalAttr:$mapping); let results = (outs Variadic:$results); let regions = (region SizedRegion<1>:$region); @@ -496,57 +499,67 @@ let skipDefaultBuilders = 1; let builders = [ // Bodyless builder, outputs must be specified. - OpBuilder<(ins "ValueRange":$outputs, "ValueRange":$num_threads, - "std::optional":$mapping)>, + OpBuilder<(ins "ValueRange":$lbs, "ValueRange":$ubs, "ValueRange":$steps, + "ValueRange":$outputs, "std::optional":$mapping)>, + // Builder that takes a bodyBuilder lambda. - OpBuilder<(ins "ValueRange":$outputs, "ValueRange":$num_threads, - "ArrayRef":$mapping, + OpBuilder<(ins "ValueRange":$lbs, "ValueRange":$ubs, "ValueRange":$steps, + "ValueRange":$outputs, "ArrayRef":$mapping, "function_ref":$bodyBuilder)> ]; let extraClassDeclaration = [{ - int64_t getRank() { return getNumThreads().size(); } + int64_t getRank() { return getLowerBound().size(); } + + /// Number of operands controlling the loop: lbs, ubs, steps + unsigned getNumControlOperands() { return 3 * getRank(); } OpResult getTiedOpResult(OpOperand *opOperand) { - assert(opOperand->getOperandNumber() >= getRank() && "invalid operand"); + assert(opOperand->getOperandNumber() >= getNumControlOperands() && + "invalid operand"); return getOperation()->getOpResult( - opOperand->getOperandNumber() - getRank()); + opOperand->getOperandNumber() - getNumControlOperands()); } /// Return the num_threads operand that is tied to the given thread id /// block argument. OpOperand *getTiedOpOperand(BlockArgument bbArg) { assert(bbArg.getArgNumber() >= getRank() && "invalid bbArg"); - return &getOperation()->getOpOperand(bbArg.getArgNumber()); + return &getOperation()->getOpOperand(bbArg.getArgNumber() + 2 * getRank()); } /// Return the shared_outs operand that is tied to the given OpResult. OpOperand *getTiedOpOperand(OpResult opResult) { assert(opResult.getDefiningOp() == getOperation() && "invalid OpResult"); return &getOperation()->getOpOperand( - opResult.getResultNumber() + getRank()); + opResult.getResultNumber() + getNumControlOperands()); } BlockArgument getTiedBlockArgument(OpOperand *opOperand) { - assert(opOperand->getOperandNumber() >= getRank() && "invalid operand"); - return getBody()->getArgument(opOperand->getOperandNumber()); + assert(opOperand->getOperandNumber() >= getNumControlOperands() + && "invalid operand"); + return getBody()->getArgument( + opOperand->getOperandNumber() - 2 * getRank()); } ArrayRef getOutputBlockArguments() { return getBody()->getArguments().drop_front(getRank()); } - ::mlir::ValueRange getThreadIndices() { + ::mlir::ValueRange getInductionVars() { return getBody()->getArguments().take_front(getRank()); } - ::mlir::Value getThreadIndex(int64_t idx) { - return getThreadIndices()[idx]; + ::mlir::Value getInductionVar(int64_t idx) { + return getInductionVars()[idx]; } ::mlir::Block::BlockArgListType getRegionOutArgs() { return getBody()->getArguments().drop_front(getRank()); } + /// Checks if the lbs are zeros and steps are ones. + bool isNormalized(); + /// Helper to sort `values` according to matching `keys`. /// Take a custom `compare` binary comparator which returns true if the first /// element is smaller than the second (i.e. compatible with std::sort). @@ -561,28 +574,28 @@ // well-formed. We override it here to ensure that we do the right thing. static void ensureTerminator(Region ®ion, OpBuilder &builder, Location loc); - PerformConcurrentlyOp getTerminator(); + InParallelOp getTerminator(); }]; } //===----------------------------------------------------------------------===// -// PerformConcurrentlyOp +// InParallelOp //===----------------------------------------------------------------------===// -def PerformConcurrentlyOp : SCF_Op<"foreach_thread.perform_concurrently", [ +def InParallelOp : SCF_Op<"forall.in_parallel", [ Pure, Terminator, DeclareOpInterfaceMethods, - HasParent<"ForeachThreadOp">, + HasParent<"ForallOp">, ] # GraphRegionNoTerminator.traits> { - let summary = "terminates a `foreach_thread` block"; + let summary = "terminates a `forall` block"; let description = [{ - `scf.foreach_thread.perform_concurrently` is a designated terminator for - the `scf.foreach_thread` operation. + `scf.forall.in_parallel` is a designated terminator for + the `scf.forall` operation. It has a single region with a single block that contains a flat list of ops. Each such op participates in the aggregate formation of a single result of - the enclosing `scf.foreach_thread`. + the enclosing `scf.forall`. The result number corresponds to the position of the op in the terminator. }]; @@ -597,8 +610,8 @@ OpBuilder<(ins)>, ]; - // TODO: Add a `PerformConcurrentlyOpInterface` interface for ops that can - // appear inside perform_concurrently. + // TODO: Add a `InParallelOpInterface` interface for ops that can + // appear inside in_parallel. let extraClassDeclaration = [{ ::llvm::SmallVector<::mlir::BlockArgument> getDests(); ::llvm::iterator_range<::mlir::Block::iterator> getYieldingOps(); diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td --- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td @@ -1326,7 +1326,7 @@ // ParallelInsertSliceOp //===----------------------------------------------------------------------===// -// TODO: Implement PerformConcurrentlyOpInterface. +// TODO: Implement InParallelOpInterface. def Tensor_ParallelInsertSliceOp : Tensor_Op<"parallel_insert_slice", [ AttrSizedOperandSegments, OffsetSizeAndStrideOpInterface, @@ -1442,7 +1442,7 @@ /// `strides` operands. static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; } - /// Return the OpResult of the enclosing ForeachThreadOp that is + /// Return the OpResult of the enclosing ForallOp that is /// corresponding to this ParallelInsertSliceOp. OpResult getTiedOpResult(); }]; diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp --- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp +++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp @@ -169,37 +169,44 @@ //===----------------------------------------------------------------------===// DiagnosedSilenceableFailure mlir::transform::gpu::mapForeachToBlocksImpl( - RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp, - function_ref &)> + RewriterBase &rewriter, scf::ForallOp forallOp, + function_ref &)> blockIdGenerator, SmallVectorImpl &gridDims, TransformOpInterface transformOp, const ArrayRef &mappingAttributes) { // Step 0. Target-specific verifications. There is no good place to anchor - // those right now: the ForeachThreadOp is target-independent and the - // transform op does not apply to individual ForeachThreadOp. - Location loc = foreachThreadOp->getLoc(); + // those right now: the ForallOp is target-independent and the + // transform op does not apply to individual ForallOp. + Location loc = forallOp->getLoc(); - if (foreachThreadOp.getNumResults() > 0) + if (!forallOp.isNormalized()) return transformOp.emitSilenceableError() - << "only bufferized scf.foreach_thread lowers to " + << "unsupported non-normalized loops"; + if (forallOp.getNumResults() > 0) + return transformOp.emitSilenceableError() + << "only bufferized scf.forall lowers to " "gpu.block_id"; - if (foreachThreadOp.getNumThreads().size() > 3) + if (forallOp.getRank() > 3) return transformOp.emitSilenceableError() - << "scf.foreach_thread with rank > 3 does not lower to " + << "scf.forall with rank > 3 does not lower to " "gpu.block_id"; - if (llvm::any_of(foreachThreadOp.getNumThreads(), [](Value v) { + if (llvm::any_of(forallOp.getLowerBound(), [](Value v) { + return !v.getDefiningOp(); + })) { + return transformOp.emitSilenceableError() + << "unsupported dynamic griddim size"; + } + if (llvm::any_of(forallOp.getUpperBound(), [](Value v) { return !v.getDefiningOp(); })) { return transformOp.emitSilenceableError() << "unsupported dynamic griddim size"; } SmallVector blockMapping = - llvm::to_vector(foreachThreadOp.getMapping()->getValue()); + llvm::to_vector(forallOp.getMapping()->getValue()); // Step 1. Complete the blockMapping to a full mapping (with 1s) if necessary. - SmallVector numBlocks = - llvm::to_vector(foreachThreadOp.getNumThreads()); + SmallVector numBlocks = llvm::to_vector(forallOp.getUpperBound()); // Ensure we have 3 block sizes, one for each id. Value one; for (auto attr : mappingAttributes) { @@ -216,68 +223,68 @@ DeviceMappingAttrInterface b) -> bool { return a.getMappingId() < b.getMappingId(); }; - SmallVector gridDimValues = scf::ForeachThreadOp::getValuesSortedByKey( - blockMapping, numBlocks, comparator); + SmallVector gridDimValues = + scf::ForallOp::getValuesSortedByKey(blockMapping, numBlocks, comparator); for (Value v : gridDimValues) gridDims.push_back(v.getDefiningOp().value()); // Step 3. Generate the blockIds using the provided generator and map the // induction variables to the newly created ops. SmallVector blockOps; - blockIdGenerator(rewriter, foreachThreadOp, blockOps); + blockIdGenerator(rewriter, forallOp, blockOps); IRMapping bvm; for (auto [blockIdx, blockDim] : - llvm::zip(foreachThreadOp.getThreadIndices(), blockMapping)) { + llvm::zip(forallOp.getInductionVars(), blockMapping)) { bvm.map(blockIdx, blockOps[static_cast( blockDim.cast().getMappingId())]); } - // Step 4. Move the body of foreachThreadOp. + // Step 4. Move the body of forallOp. // Erase the terminator first, it will not be used since we are on buffers. - rewriter.eraseOp(foreachThreadOp.getTerminator()); - Block *targetBlock = foreachThreadOp->getBlock(); - Block::iterator insertionPoint = Block::iterator(foreachThreadOp); - Block &sourceBlock = foreachThreadOp.getRegion().front(); + rewriter.eraseOp(forallOp.getTerminator()); + Block *targetBlock = forallOp->getBlock(); + Block::iterator insertionPoint = Block::iterator(forallOp); + Block &sourceBlock = forallOp.getRegion().front(); targetBlock->getOperations().splice(insertionPoint, sourceBlock.getOperations()); // Step 5. RAUW thread indices to thread ops. - for (Value loopIndex : foreachThreadOp.getThreadIndices()) { + for (Value loopIndex : forallOp.getInductionVars()) { Value blockIdx = bvm.lookup(loopIndex); rewriter.replaceAllUsesWith(loopIndex, blockIdx); } // Step 6. Erase old op. - rewriter.eraseOp(foreachThreadOp); + rewriter.eraseOp(forallOp); return DiagnosedSilenceableFailure::success(); } -DiagnosedSilenceableFailure mlir::transform::gpu::findTopLevelForeachThreadOp( - Operation *target, scf::ForeachThreadOp &topLevelForeachThreadOp, - TransformOpInterface transformOp) { - auto walkResult = target->walk([&](scf::ForeachThreadOp foreachThreadOp) { - if (foreachThreadOp->getParentOfType()) +DiagnosedSilenceableFailure +mlir::transform::gpu::findTopLevelForallOp(Operation *target, + scf::ForallOp &topLevelForallOp, + TransformOpInterface transformOp) { + auto walkResult = target->walk([&](scf::ForallOp forallOp) { + if (forallOp->getParentOfType()) return WalkResult::advance(); - if (topLevelForeachThreadOp) + if (topLevelForallOp) // TODO: Handle multiple foreach if there is no dependences between them return WalkResult::interrupt(); - topLevelForeachThreadOp = foreachThreadOp; + topLevelForallOp = forallOp; return WalkResult::advance(); }); if (walkResult.wasInterrupted()) return transformOp.emitSilenceableError() - << "could not find a unique topLevel scf.foreach_thread"; + << "could not find a unique topLevel scf.forall"; return DiagnosedSilenceableFailure::success(); } /// This is a helper that is only used in /// rewriteTopLevelForeachThreadToGpuBlocks. It generates GPU dialects /// block_id. -static void generateGpuBlockIds(RewriterBase &rewriter, - scf::ForeachThreadOp foreachOp, +static void generateGpuBlockIds(RewriterBase &rewriter, scf::ForallOp foreachOp, SmallVectorImpl &blockOps) { Location loc = foreachOp->getLoc(); OpBuilder::InsertionGuard guard(rewriter); @@ -306,17 +313,16 @@ return diag; } - scf::ForeachThreadOp topLevelForeachThreadOp; - DiagnosedSilenceableFailure diag = - mlir::transform::gpu::findTopLevelForeachThreadOp( - target, topLevelForeachThreadOp, transformOp); + scf::ForallOp topLevelForallOp; + DiagnosedSilenceableFailure diag = mlir::transform::gpu::findTopLevelForallOp( + target, topLevelForallOp, transformOp); if (!diag.succeeded()) { diag.attachNote(target->getLoc()) << "when applied to this payload op"; return diag; } OpBuilder::InsertionGuard guard(rewriter); - rewriter.setInsertionPoint(topLevelForeachThreadOp); + rewriter.setInsertionPoint(topLevelForallOp); // Generate gpu launch here and move the foreach_thread inside if (getGenerateGpuLaunch()) { @@ -326,9 +332,9 @@ return diag; } rewriter.setInsertionPointToStart(&gpuLaunch.getBody().front()); - Operation *newForeachThreadOp = rewriter.clone(*topLevelForeachThreadOp); - rewriter.eraseOp(topLevelForeachThreadOp); - topLevelForeachThreadOp = cast(newForeachThreadOp); + Operation *newForallOp = rewriter.clone(*topLevelForallOp); + rewriter.eraseOp(topLevelForallOp); + topLevelForallOp = cast(newForallOp); } SmallVector gridDim = extractFromI64ArrayAttr(getGridDim()); @@ -338,11 +344,11 @@ GPUBlockMappingAttr::get(getContext(), Blocks::DimZ)}; diag = checkAttributeType(blockMappingAttributes, - topLevelForeachThreadOp.getMapping(), transformOp); + topLevelForallOp.getMapping(), transformOp); if (diag.succeeded()) diag = mlir::transform::gpu::mapForeachToBlocksImpl( - rewriter, topLevelForeachThreadOp, generateGpuBlockIds, gridDim, - transformOp, blockMappingAttributes); + rewriter, topLevelForallOp, generateGpuBlockIds, gridDim, transformOp, + blockMappingAttributes); if (diag.succeeded()) { diag = alterGpuLaunch(rewriter, gpuLaunch, cast(getOperation()), @@ -357,50 +363,50 @@ // MapNestedForeachToThreads //===----------------------------------------------------------------------===// -/// Searches `scf.foreach_thread` ops nested under `target` and maps each such +/// Searches `scf.forall` ops nested under `target` and maps each such /// op to GPU threads. Mapping is one-to-one and the induction variables of -/// `scf.foreach_thread` are rewritten to gpu.thread_id according to the -/// thread_dim_mapping attribute. Sibling `scf.foreach_thread` are supported in +/// `scf.forall` are rewritten to gpu.thread_id according to the +/// thread_dim_mapping attribute. Sibling `scf.forall` are supported in /// which case, the union of the number of threads is computed and may result -/// in predication. Dynamic, `scf.foreach_thread` trip counts are currently +/// in predication. Dynamic, `scf.forall` trip counts are currently /// not supported. Dynamic block dim sizes are currently not supported. static DiagnosedSilenceableFailure rewriteOneForeachThreadToGpuThreads( - RewriterBase &rewriter, scf::ForeachThreadOp foreachThreadOp, + RewriterBase &rewriter, scf::ForallOp forallOp, const SmallVectorImpl &globalBlockDims, const SmallVectorImpl &threadOps, bool syncAfterDistribute, std::optional transformOp, const ArrayRef &threadMappingAttributes) { // Step 0. Target-specific verifications. There is no good place to anchor - // those right now: the ForeachThreadOp is target-independent and the - // transform op does not apply to individual ForeachThreadOp. + // those right now: the ForallOp is target-independent and the + // transform op does not apply to individual ForallOp. auto failureHelper = [&](const Twine &message) -> DiagnosedSilenceableFailure { if (transformOp.has_value()) { return transformOp->emitSilenceableError() << message; } - return emitDefiniteFailure(foreachThreadOp, message); + return emitDefiniteFailure(forallOp, message); }; - Location loc = foreachThreadOp->getLoc(); - if (foreachThreadOp.getNumResults() > 0) - return failureHelper( - "only bufferized scf.foreach_thread lowers to gpu.thread_id"); - if (foreachThreadOp.getNumThreads().size() > 3) + Location loc = forallOp->getLoc(); + if (!forallOp.isNormalized()) + return failureHelper("unsupported non-normalized loops"); + if (forallOp.getNumResults() > 0) + return failureHelper("only bufferized scf.forall lowers to gpu.thread_id"); + if (forallOp.getRank() > 3) return failureHelper( - "scf.foreach_thread with rank > 3 does not lower to gpu.thread_id"); - if (llvm::any_of(foreachThreadOp.getNumThreads(), [](Value v) { + "scf.forall with rank > 3 does not lower to gpu.thread_id"); + if (llvm::any_of(forallOp.getUpperBound(), [](Value v) { return !v.getDefiningOp(); })) { return failureHelper("unsupported dynamic blockdim size"); } - if (!foreachThreadOp.getMapping().has_value()) + if (!forallOp.getMapping().has_value()) return failureHelper("mapping must be present"); SmallVector threadMapping = - llvm::to_vector(foreachThreadOp.getMapping()->getValue()); + llvm::to_vector(forallOp.getMapping()->getValue()); // Step 1. Complete the threadMapping to a full mapping (with 1s) if // necessary. - SmallVector numThreads = - llvm::to_vector(foreachThreadOp.getNumThreads()); + SmallVector numThreads = llvm::to_vector(forallOp.getUpperBound()); // Ensure we have 3 block sizes, one for each id. Value one; for (auto attr : threadMappingAttributes) { @@ -417,9 +423,8 @@ DeviceMappingAttrInterface b) -> bool { return a.getMappingId() < b.getMappingId(); }; - SmallVector blockDimValues = - scf::ForeachThreadOp::getValuesSortedByKey(threadMapping, numThreads, - comparator); + SmallVector blockDimValues = scf::ForallOp::getValuesSortedByKey( + threadMapping, numThreads, comparator); SmallVector blockDims = llvm::to_vector(llvm::map_range(blockDimValues, [](Value v) { return v.getDefiningOp().value(); @@ -437,7 +442,7 @@ } IRMapping bvm; for (auto [blockIdx, blockDim] : - llvm::zip(foreachThreadOp.getThreadIndices(), threadMapping)) { + llvm::zip(forallOp.getInductionVars(), threadMapping)) { bvm.map(blockIdx, threadOpsUpdated[blockDim.cast() .getMappingId()]); @@ -450,7 +455,7 @@ if (blockDim > globalBlockDim) { return failureHelper( "The requested GPU threads are fewer than the number of loop trip " - "counts. Try to tile scf.foreach_thread before mapping or set " + "counts. Try to tile scf.forall before mapping or set " "small blockDim."); } if (blockDim == globalBlockDim) @@ -463,9 +468,9 @@ : tmpPredicate; } - // Step 5. Move the body of foreachThreadOp. + // Step 5. Move the body of forallOp. // Erase the terminator first, it will not be used. - rewriter.eraseOp(foreachThreadOp.getTerminator()); + rewriter.eraseOp(forallOp.getTerminator()); Block *targetBlock; Block::iterator insertionPoint; if (predicate) { @@ -475,16 +480,16 @@ targetBlock = ifOp.thenBlock(); insertionPoint = ifOp.thenBlock()->begin(); } else { - // Step 5.b. Otherwise, move inline just before foreachThreadOp. - targetBlock = foreachThreadOp->getBlock(); - insertionPoint = Block::iterator(foreachThreadOp); + // Step 5.b. Otherwise, move inline just before forallOp. + targetBlock = forallOp->getBlock(); + insertionPoint = Block::iterator(forallOp); } - Block &sourceBlock = foreachThreadOp.getRegion().front(); + Block &sourceBlock = forallOp.getRegion().front(); targetBlock->getOperations().splice(insertionPoint, sourceBlock.getOperations()); // Step 6. RAUW thread indices to thread ops. - for (Value loopIndex : foreachThreadOp.getThreadIndices()) { + for (Value loopIndex : forallOp.getInductionVars()) { Value threadIdx = bvm.lookup(loopIndex); rewriter.replaceAllUsesWith(loopIndex, threadIdx); } @@ -495,7 +500,7 @@ rewriter.create(loc); // Step 8. Erase old op. - rewriter.eraseOp(foreachThreadOp); + rewriter.eraseOp(forallOp); return DiagnosedSilenceableFailure::success(); } @@ -506,27 +511,27 @@ std::optional transformOp, const ArrayRef &threadMappingAttributes) { DiagnosedSilenceableFailure diag = DiagnosedSilenceableFailure::success(); - target->walk([&](scf::ForeachThreadOp foreachThreadOp) { + target->walk([&](scf::ForallOp forallOp) { // Ignore cases with different attributes. - for (Attribute map : foreachThreadOp.getMapping()->getValue()) { + for (Attribute map : forallOp.getMapping()->getValue()) { if (!llvm::is_contained(threadMappingAttributes, map)) { return WalkResult::skip(); } } - diag = checkAttributeType(threadMappingAttributes, - foreachThreadOp.getMapping(), transformOp); + diag = checkAttributeType(threadMappingAttributes, forallOp.getMapping(), + transformOp); if (diag.succeeded()) { - rewriter.setInsertionPoint(foreachThreadOp); + rewriter.setInsertionPoint(forallOp); IndexType indexType = rewriter.getIndexType(); SmallVector threadOps{ - rewriter.create(foreachThreadOp.getLoc(), indexType, + rewriter.create(forallOp.getLoc(), indexType, Dimension::x), - rewriter.create(foreachThreadOp.getLoc(), indexType, + rewriter.create(forallOp.getLoc(), indexType, Dimension::y), - rewriter.create(foreachThreadOp.getLoc(), indexType, + rewriter.create(forallOp.getLoc(), indexType, Dimension::z)}; diag = rewriteOneForeachThreadToGpuThreads( - rewriter, foreachThreadOp, blockDim, threadOps, syncAfterDistribute, + rewriter, forallOp, blockDim, threadOps, syncAfterDistribute, transformOp, threadMappingAttributes); } return diag.succeeded() ? WalkResult::advance() : WalkResult::interrupt(); diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp --- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp +++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp @@ -457,7 +457,7 @@ return fusedOp; } -/// First, find the first "scf::ForeachThreadOp" user of `producerOp` and ensure +/// First, find the first "scf::ForallOp" user of `producerOp` and ensure /// it is exactly the `containingOp`, otherwise bail. /// Then, find the first "extract" user of the tied block argument and tile it /// right before its "extract" use. The tiled op is fused under the @@ -475,11 +475,11 @@ return nullptr; } - // Search the first use by a "scf::ForeachThreadOp" user. - scf::ForeachThreadOp foreachThreadOp; + // Search the first use by a "scf::ForallOp" user. + scf::ForallOp foreachThreadOp; auto itProducerUses = llvm::find_if(tileableProducer->getUses(), [&](OpOperand &use) { - foreachThreadOp = dyn_cast(use.getOwner()); + foreachThreadOp = dyn_cast(use.getOwner()); return foreachThreadOp; }); // If it's not from the containing op, return. @@ -2137,7 +2137,7 @@ // This is future-proof re mixed static-dynamic and setting up the proper // operands segment sizes attributes for multiple variadic operands. // In the absence of this, horrible bugs ensue. - // TODO: support mixed static-dynamic (see TileToForeachThreadOp). + // TODO: support mixed static-dynamic (see TileToForallOp). MLIRContext *ctx = builder.getContext(); auto opTy = pdl::OperationType::get(ctx); auto staticTileSizesAttr = builder.getDenseI64ArrayAttr(staticTileSizes); @@ -2166,10 +2166,10 @@ } //===----------------------------------------------------------------------===// -// TileReductionUsingForeachThreadOp +// TileReductionUsingForallOp //===----------------------------------------------------------------------===// -void transform::TileReductionUsingForeachThreadOp::build( +void transform::TileReductionUsingForallOp::build( OpBuilder &builder, OperationState &result, Value target, ArrayRef staticNumThreads, ArrayRef staticTileSizes, ArrayAttr mapping) { @@ -2177,7 +2177,7 @@ // This is future-proof re mixed static-dynamic and setting up the proper // operands segment sizes attributes for multiple variadic operands. // In the absence of this, horrible bugs ensue. - // TODO: support mixed static-dynamic (see TileToForeachThreadOp). + // TODO: support mixed static-dynamic (see TileToForallOp). MLIRContext *ctx = builder.getContext(); auto opTy = pdl::OperationType::get(ctx); auto staticNumThreadsAttr = builder.getDenseI64ArrayAttr(staticNumThreads); @@ -2190,8 +2190,7 @@ /*mapping=*/mapping); } -DiagnosedSilenceableFailure -transform::TileReductionUsingForeachThreadOp::applyToOne( +DiagnosedSilenceableFailure transform::TileReductionUsingForallOp::applyToOne( LinalgOp target, transform::ApplyToEachResultList &results, transform::TransformState &state) { TrivialPatternRewriter rewriter(getContext()); @@ -2502,15 +2501,14 @@ } //===----------------------------------------------------------------------===// -// TileToForeachThreadOp +// TileToForallOp //===----------------------------------------------------------------------===// -void transform::TileToForeachThreadOp::build(OpBuilder &builder, - OperationState &result, - Value target, - ArrayRef staticTileSizes, - transform::TileSizesSpec, - ArrayAttr mapping) { +void transform::TileToForallOp::build(OpBuilder &builder, + OperationState &result, Value target, + ArrayRef staticTileSizes, + transform::TileSizesSpec, + ArrayAttr mapping) { return build(builder, result, /*target=*/target, /*mixedTileSizes=*/ @@ -2519,10 +2517,11 @@ /*mapping=*/mapping); } -void transform::TileToForeachThreadOp::build( - OpBuilder &builder, OperationState &result, Value target, - ArrayRef mixedTileSizes, transform::TileSizesSpec, - ArrayAttr mapping) { +void transform::TileToForallOp::build(OpBuilder &builder, + OperationState &result, Value target, + ArrayRef mixedTileSizes, + transform::TileSizesSpec, + ArrayAttr mapping) { SmallVector staticTileSizes; SmallVector dynamicTileSizes; dispatchIndexOpFoldResults(mixedTileSizes, dynamicTileSizes, staticTileSizes); @@ -2544,21 +2543,21 @@ /*mapping=*/mapping); } -void transform::TileToForeachThreadOp::build(OpBuilder &builder, - OperationState &result, - Value target, - ArrayRef staticNumThreads, - transform::NumThreadsSpec, - ArrayAttr mapping) { +void transform::TileToForallOp::build(OpBuilder &builder, + OperationState &result, Value target, + ArrayRef staticNumThreads, + transform::NumThreadsSpec, + ArrayAttr mapping) { return build(builder, result, target, getAsOpFoldResult(builder.getI64ArrayAttr(staticNumThreads)), NumThreadsSpec(), mapping); } -void transform::TileToForeachThreadOp::build( - OpBuilder &builder, OperationState &result, Value target, - ArrayRef mixedNumThreads, transform::NumThreadsSpec, - ArrayAttr mapping) { +void transform::TileToForallOp::build(OpBuilder &builder, + OperationState &result, Value target, + ArrayRef mixedNumThreads, + transform::NumThreadsSpec, + ArrayAttr mapping) { SmallVector staticNumThreads; SmallVector dynamicNumThreads; dispatchIndexOpFoldResults(mixedNumThreads, dynamicNumThreads, @@ -2581,7 +2580,7 @@ /*mapping=*/mapping); } -DiagnosedSilenceableFailure transform::tileToForeachThreadOpImpl( +DiagnosedSilenceableFailure transform::tileToForallOpImpl( RewriterBase &rewriter, transform::TransformState &state, TransformOpInterface transformOp, ArrayRef targets, ArrayRef mixedNumThreads, @@ -2603,10 +2602,10 @@ rewriter.setInsertionPoint(tileableOp); FailureOr tilingResult = failure(); if (!mixedNumThreads.empty()) { - tilingResult = linalg::tileToForeachThreadOp(rewriter, tileableOp, - mixedNumThreads, mapping); + tilingResult = linalg::tileToForallOp(rewriter, tileableOp, + mixedNumThreads, mapping); } else { - tilingResult = linalg::tileToForeachThreadOpUsingTileSizes( + tilingResult = linalg::tileToForallOpUsingTileSizes( rewriter, tileableOp, mixedTileSizes, mapping); } @@ -2620,9 +2619,9 @@ return DiagnosedSilenceableFailure::success(); } -DiagnosedSilenceableFailure transform::TileToForeachThreadOp::apply( - transform::TransformResults &transformResults, - transform::TransformState &state) { +DiagnosedSilenceableFailure +transform::TileToForallOp::apply(transform::TransformResults &transformResults, + transform::TransformState &state) { IRRewriter rewriter(getContext()); auto transformOp = cast(getOperation()); ArrayRef targets = state.getPayloadOps(getTarget()); @@ -2650,20 +2649,20 @@ if (!status.succeeded()) return status; - DiagnosedSilenceableFailure diag = tileToForeachThreadOpImpl( - rewriter, state, transformOp, targets, mixedNumThreads, mixedTileSizes, - getMapping(), tileOps, tiledOps); + DiagnosedSilenceableFailure diag = + tileToForallOpImpl(rewriter, state, transformOp, targets, mixedNumThreads, + mixedTileSizes, getMapping(), tileOps, tiledOps); if (!diag.succeeded()) return diag; - transformResults.set(getForeachThreadOp().cast(), tileOps); + transformResults.set(getForallOp().cast(), tileOps); transformResults.set(getTiledOp().cast(), tiledOps); return DiagnosedSilenceableFailure::success(); } -void transform::TileToForeachThreadOp::getEffects( +void transform::TileToForallOp::getEffects( SmallVectorImpl &effects) { consumesHandle(getTarget(), effects); onlyReadsHandle(getTileSizes(), effects); @@ -2674,17 +2673,17 @@ modifiesPayload(effects); } -SmallVector TileToForeachThreadOp::getMixedNumThreads() { +SmallVector TileToForallOp::getMixedNumThreads() { Builder b(getContext()); return getMixedValues(getStaticNumThreads(), getNumThreads(), b); } -SmallVector TileToForeachThreadOp::getMixedTileSizes() { +SmallVector TileToForallOp::getMixedTileSizes() { Builder b(getContext()); return getMixedValues(getStaticTileSizes(), getTileSizes(), b); } -LogicalResult TileToForeachThreadOp::verify() { +LogicalResult TileToForallOp::verify() { int numThreadsSpec = static_cast(!getMixedNumThreads().empty()) + static_cast(getPackedNumThreads() != Value()); if (numThreadsSpec > 1) diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp @@ -244,16 +244,16 @@ /// Fill out the `tiledOffsets` and `tiledSizes` to be used to tile to a given /// number of threads. static void calculateTileOffsetsAndSizes( - RewriterBase &b, Location loc, scf::ForeachThreadOp foreachThreadOp, + RewriterBase &b, Location loc, scf::ForallOp forallOp, ArrayRef numThreads, SmallVector loopRanges, bool omitTileOffsetBoundsCheck, std::optional> nominalTileSizes, SmallVector &tiledOffsets, SmallVector &tiledSizes) { OpBuilder::InsertionGuard g(b); - b.setInsertionPointToStart(foreachThreadOp.getBody(0)); + b.setInsertionPointToStart(forallOp.getBody(0)); - ValueRange threadIds = foreachThreadOp.getThreadIndices(); + ValueRange threadIds = forallOp.getInductionVars(); SmallVector nonZeroNumThreads = llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) { return !isConstantIntValue(ofr, 0); @@ -314,19 +314,19 @@ } } -/// Rewrite a TilingInterface `op` to a tiled `scf.foreach_thread`. The +/// Rewrite a TilingInterface `op` to a tiled `scf.forall`. The /// tiling is specified by the number of tiles/threads `numThreads` and the /// optional nominal tile size `nominalTileSizes`. If `nominalTilSizes` is /// not specified, then it is derived from `numThreads` as `ceilDiv(dimSize[i], /// numThreads[i])`. If non-empty, the `mapping` is added as an -/// attribute to the resulting `scf.foreach_thread`. A zero tile sizes indicate +/// attribute to the resulting `scf.forall`. A zero tile sizes indicate /// that the dimension is not tiled, and can be thought of as tiling by the full /// size of data. /// It is the user's responsibility to ensure that `numThreads` is a valid /// tiling specification (i.e. that only tiles parallel dimensions, e.g. in the /// Linalg case). If `omitTileOffsetBoundsCheck` is true, then the function will /// assume that `tileSize[i] * (numThread[i] -1) <= dimSize[i]` holds. -static FailureOr tileToForeachThreadOpImpl( +static FailureOr tileToForallOpImpl( RewriterBase &b, TilingInterface op, ArrayRef numThreads, std::optional> nominalTileSizes, std::optional mapping, bool omitTileOffsetBoundsCheck) { @@ -356,26 +356,29 @@ Operation *tiledOp = nullptr; - // 1. Create the ForeachThreadOp. We don't use the lambda body-builder + // 1. Create the ForallOp. We don't use the lambda body-builder // version because we require the use of RewriterBase in the body, so we // manually move the insertion point to the body below. - scf::ForeachThreadOp foreachThreadOp = b.create( - loc, dest, ValueRange(materializedNonZeroNumThreads), mapping); - - // 2. Fill out the ForeachThreadOp body. + unsigned numLoops = materializedNonZeroNumThreads.size(); + SmallVector lbs(numLoops, b.create(loc, 0)); + SmallVector steps(numLoops, b.create(loc, 1)); + scf::ForallOp forallOp = b.create( + loc, lbs, ValueRange(materializedNonZeroNumThreads), steps, dest, + mapping); + + // 2. Fill out the ForallOp body. SmallVector tiledOffsets, tiledSizes; - calculateTileOffsetsAndSizes(b, loc, foreachThreadOp, numThreads, loopRanges, + calculateTileOffsetsAndSizes(b, loc, forallOp, numThreads, loopRanges, omitTileOffsetBoundsCheck, nominalTileSizes, tiledOffsets, tiledSizes); // 3. Clone the tileable op and update its destination operands to use the - // output bbArgs of the ForeachThreadOp. - ArrayRef destBbArgs = - foreachThreadOp.getOutputBlockArguments(); + // output bbArgs of the ForallOp. + ArrayRef destBbArgs = forallOp.getOutputBlockArguments(); { - // 3.a. RAII guard, inserting within foreachThreadOp, before terminator. + // 3.a. RAII guard, inserting within forallOp, before terminator. OpBuilder::InsertionGuard g(b); - b.setInsertionPoint(foreachThreadOp.getTerminator()); + b.setInsertionPoint(forallOp.getTerminator()); Operation *clonedOp = b.clone(*op.getOperation()); auto destinationStyleOp = dyn_cast(clonedOp); if (destinationStyleOp) { @@ -404,7 +407,7 @@ tilingInterfaceOp->getResults(), destBbArgs)) { // 5.a. Partial subset information is inserted just before the terminator. OpBuilder::InsertionGuard g(b); - b.setInsertionPoint(foreachThreadOp.getTerminator()); + b.setInsertionPoint(forallOp.getTerminator()); SmallVector resultOffsets, resultSizes; if (failed(op.getResultTilePosition(b, std::get<0>(it), tiledOffsets, @@ -415,27 +418,27 @@ // 5.b. Parallel insertions are inserted at the end of the combining // terminator. - b.setInsertionPointToEnd(foreachThreadOp.getTerminator().getBody()); + b.setInsertionPointToEnd(forallOp.getTerminator().getBody()); b.create(loc, std::get<1>(it), std::get<2>(it), resultOffsets, resultSizes, strides); } - return ForeachThreadTilingResult{foreachThreadOp, tiledOp}; + return ForeachThreadTilingResult{forallOp, tiledOp}; } FailureOr -linalg::tileToForeachThreadOp(RewriterBase &b, TilingInterface op, - ArrayRef numThreads, - std::optional mapping) { - return tileToForeachThreadOpImpl(b, op, numThreads, - /*nominalTileSizes=*/std::nullopt, mapping, - /*omitTileOffsetBoundsCheck=*/false); +linalg::tileToForallOp(RewriterBase &b, TilingInterface op, + ArrayRef numThreads, + std::optional mapping) { + return tileToForallOpImpl(b, op, numThreads, + /*nominalTileSizes=*/std::nullopt, mapping, + /*omitTileOffsetBoundsCheck=*/false); } FailureOr -linalg::tileToForeachThreadOpUsingTileSizes(RewriterBase &b, TilingInterface op, - ArrayRef tileSizes, - std::optional mapping) { +linalg::tileToForallOpUsingTileSizes(RewriterBase &b, TilingInterface op, + ArrayRef tileSizes, + std::optional mapping) { SmallVector loopRanges = op.getIterationDomain(b); unsigned nLoops = loopRanges.size(); SmallVector numThreads; @@ -450,9 +453,9 @@ b, op.getLoc(), divExpr, {std::get<1>(it).size, std::get<0>(it)}); numThreads.push_back(numTiles); } - return tileToForeachThreadOpImpl(b, op, numThreads, - /*nominalTileSizes=*/tileSizes, mapping, - /*omitTileOffsetBoundsCheck=*/true); + return tileToForallOpImpl(b, op, numThreads, + /*nominalTileSizes=*/tileSizes, mapping, + /*omitTileOffsetBoundsCheck=*/true); } template @@ -679,28 +682,30 @@ SmallVector materializedNonZeroNumThreads = getAsValues(b, loc, nonZeroNumThreads); - // 2. Create the ForeachThreadOp with an empty region. - scf::ForeachThreadOp foreachThreadOp = b.create( - loc, (*identityTensor)->getResults(), - ValueRange(materializedNonZeroNumThreads), mapping); + // 2. Create the ForallOp with an empty region. + unsigned numLoops = materializedNonZeroNumThreads.size(); + SmallVector lbs(numLoops, b.create(loc, 0)); + SmallVector steps(numLoops, b.create(loc, 1)); + scf::ForallOp forallOp = b.create( + loc, lbs, ValueRange(materializedNonZeroNumThreads), steps, + (*identityTensor)->getResults(), mapping); // 3. Calculate the tile offsets and sizes for the subsequent loop that will - // be nested under `foreachThreadOp`. + // be nested under `forallOp`. SmallVector tiledOffsets, tiledSizes; - calculateTileOffsetsAndSizes( - b, loc, foreachThreadOp, numThreads, iterationDomain, - /*omitTileOffsetBoundsCheck =*/false, - /*nominalTileSizes=*/std::nullopt, tiledOffsets, tiledSizes); + calculateTileOffsetsAndSizes(b, loc, forallOp, numThreads, iterationDomain, + /*omitTileOffsetBoundsCheck =*/false, + /*nominalTileSizes=*/std::nullopt, tiledOffsets, + tiledSizes); // 4. Clone the tileable op and update its destination operands to use the - // output bbArgs of the ForeachThreadOp. + // output bbArgs of the ForallOp. ValueRange tilingResults; - ArrayRef destBbArgs = - foreachThreadOp.getOutputBlockArguments(); + ArrayRef destBbArgs = forallOp.getOutputBlockArguments(); { - // 4.a. RAII guard, inserting within foreachThreadOp, before terminator. + // 4.a. RAII guard, inserting within forallOp, before terminator. OpBuilder::InsertionGuard g(b); - b.setInsertionPoint(foreachThreadOp.getTerminator()); + b.setInsertionPoint(forallOp.getTerminator()); SmallVector tiledDpsInitOperands; for (OpOperand *initOperand : destinationStyleOp.getDpsInitOperands()) { @@ -712,7 +717,7 @@ b.getIndexAttr(0)); SmallVector sizes = tiledSizes; sizes[reductionDim] = b.getIndexAttr(1); - outOffsets[reductionDim] = foreachThreadOp.getThreadIndices().front(); + outOffsets[reductionDim] = forallOp.getInductionVars().front(); // TODO: use SubsetExtractOpInterface once it is available. tiledDpsInitOperands.push_back(b.create( loc, initOperand->get().getType().cast(), @@ -746,7 +751,7 @@ if (failed(maybeTiled)) return b.notifyMatchFailure(op, "failed tileLinalgOpImpl"); - SmallVector ids = foreachThreadOp.getThreadIndices(); + SmallVector ids = forallOp.getInductionVars(); mapLoopToProcessorIds(cast(maybeTiled->loops.back()), ids, materializedNonZeroNumThreads); assert(maybeTiled->loops.size() == 1 && @@ -763,7 +768,7 @@ llvm::seq(0, dest.size()), tilingResults, destBbArgs)) { // 6.a. Partial subset information is inserted just before the terminator. OpBuilder::InsertionGuard g(b); - b.setInsertionPoint(foreachThreadOp.getTerminator()); + b.setInsertionPoint(forallOp.getTerminator()); SmallVector resultOffsets, resultSizes; if (failed(tilingInterfaceOp.getResultTilePosition( @@ -774,7 +779,7 @@ int64_t sizeIdx = 0; for (int64_t i = 0, e = numThreads.size(); i < e; ++i) { if (i == reductionDim) { - resultOffsetsRank.push_back(foreachThreadOp.getThreadIndices().front()); + resultOffsetsRank.push_back(forallOp.getInductionVars().front()); resultSizesRank.push_back(b.getIndexAttr(1)); continue; } @@ -786,21 +791,21 @@ // 6.b. Parallel insertions are inserted at the end of the combining // terminator. - b.setInsertionPointToEnd(foreachThreadOp.getTerminator().getBody()); + b.setInsertionPointToEnd(forallOp.getTerminator().getBody()); b.create( loc, result, bbArg, resultOffsetsRank, resultSizesRank, strides); } // 7. Merge the partial reductions. - b.setInsertionPointAfter(foreachThreadOp); + b.setInsertionPointAfter(forallOp); Operation *mergeOp = - op.mergeReductions(b, loc, foreachThreadOp->getResults(), reductionDim); + op.mergeReductions(b, loc, forallOp->getResults(), reductionDim); b.replaceOp(op, mergeOp->getResults()); // 8. Return. ForeachThreadReductionTilingResult results; results.initialOp = *identityTensor; - results.loops = foreachThreadOp; + results.loops = forallOp; results.parallelTiledOp = tiledOp; results.mergeOp = mergeOp; return results; diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -1106,10 +1106,10 @@ } //===----------------------------------------------------------------------===// -// ForeachThreadOp +// ForallOp //===----------------------------------------------------------------------===// -LogicalResult ForeachThreadOp::verify() { +LogicalResult ForallOp::verify() { // Check number of outputs. if (getNumResults() != getOutputs().size()) return emitOpError("produces ") @@ -1141,12 +1141,10 @@ return success(); } -void ForeachThreadOp::print(OpAsmPrinter &p) { - p << " ("; - llvm::interleaveComma(getThreadIndices(), p); - p << ") in ("; - llvm::interleaveComma(getNumThreads(), p); - p << ")"; +void ForallOp::print(OpAsmPrinter &p) { + p << " (" << getInductionVars() << ") = (" << getLowerBound() << ") to (" + << getUpperBound() << ") step (" << getStep() << ") "; + printInitializationList(p, getRegionOutArgs(), getOutputs(), " shared_outs"); p << " "; if (!getRegionOutArgs().empty()) @@ -1158,23 +1156,35 @@ {"operand_segment_sizes"}); } -ParseResult ForeachThreadOp::parse(OpAsmParser &parser, - OperationState &result) { +ParseResult ForallOp::parse(OpAsmParser &parser, OperationState &result) { auto &builder = parser.getBuilder(); // Parse an opening `(` followed by thread index variables followed by `)` // TODO: when we can refer to such "induction variable"-like handles from the // declarative assembly format, we can implement the parser as a custom hook. - SmallVector threadIndices; - if (parser.parseArgumentList(threadIndices, OpAsmParser::Delimiter::Paren)) + SmallVector ivs; + if (parser.parseArgumentList(ivs, OpAsmParser::Delimiter::Paren)) + return failure(); + + // Parse lower bounds. + SmallVector lbs; + if (parser.parseEqual() || + parser.parseOperandList(lbs, ivs.size(), OpAsmParser::Delimiter::Paren) || + parser.resolveOperands(lbs, builder.getIndexType(), result.operands)) return failure(); - // Parse `in` threadNums. - SmallVector threadNums; - if (parser.parseKeyword("in") || - parser.parseOperandList(threadNums, threadIndices.size(), + // Parse upper bounds. + SmallVector ubs; + if (parser.parseKeyword("to") || + parser.parseOperandList(ubs, ivs.size(), OpAsmParser::Delimiter::Paren) || + parser.resolveOperands(ubs, builder.getIndexType(), result.operands)) + return failure(); + + // Parse step values. + SmallVector steps; + if (parser.parseKeyword("step") || + parser.parseOperandList(steps, ivs.size(), OpAsmParser::Delimiter::Paren) || - parser.resolveOperands(threadNums, builder.getIndexType(), - result.operands)) + parser.resolveOperands(steps, builder.getIndexType(), result.operands)) return failure(); // Parse out operands and results. @@ -1195,9 +1205,9 @@ // Parse region. SmallVector regionArgs; std::unique_ptr region = std::make_unique(); - for (auto &idx : threadIndices) { - idx.type = builder.getIndexType(); - regionArgs.push_back(idx); + for (auto &iv : ivs) { + iv.type = builder.getIndexType(); + regionArgs.push_back(iv); } for (const auto &it : llvm::enumerate(regionOutArgs)) { auto &out = it.value(); @@ -1209,34 +1219,39 @@ // Ensure terminator and move region. OpBuilder b(builder.getContext()); - ForeachThreadOp::ensureTerminator(*region, b, result.location); + ForallOp::ensureTerminator(*region, b, result.location); result.addRegion(std::move(region)); // Parse the optional attribute list. if (parser.parseOptionalAttrDict(result.attributes)) return failure(); - result.addAttribute("operand_segment_sizes", - parser.getBuilder().getDenseI32ArrayAttr( - {static_cast(threadNums.size()), - static_cast(outOperands.size())})); + result.addAttribute( + "operand_segment_sizes", + parser.getBuilder().getDenseI32ArrayAttr( + {static_cast(lbs.size()), static_cast(ubs.size()), + static_cast(steps.size()), + static_cast(outOperands.size())})); return success(); } // Bodyless builder, outputs must be specified. -void ForeachThreadOp::build(mlir::OpBuilder &builder, - mlir::OperationState &result, ValueRange outputs, - ValueRange numThreads, - std::optional mapping) { - result.addOperands(numThreads); +void ForallOp::build(mlir::OpBuilder &builder, mlir::OperationState &result, + ValueRange lbs, ValueRange ubs, ValueRange steps, + ValueRange outputs, std::optional mapping) { + result.addOperands(lbs); + result.addOperands(ubs); + result.addOperands(steps); result.addOperands(outputs); if (mapping.has_value()) { - result.addAttribute(ForeachThreadOp::getMappingAttrName(result.name), + result.addAttribute(ForallOp::getMappingAttrName(result.name), mapping.value()); } result.addAttribute( "operand_segment_sizes", - builder.getDenseI32ArrayAttr({static_cast(numThreads.size()), + builder.getDenseI32ArrayAttr({static_cast(lbs.size()), + static_cast(ubs.size()), + static_cast(steps.size()), static_cast(outputs.size())})); result.addTypes(TypeRange(outputs)); @@ -1248,27 +1263,31 @@ builder.createBlock(bodyRegion); Block &bodyBlock = bodyRegion->front(); // Add block arguments for indices and outputs. - bodyBlock.addArguments( - SmallVector(numThreads.size(), builder.getIndexType()), - SmallVector(numThreads.size(), result.location)); + bodyBlock.addArguments(SmallVector(lbs.size(), builder.getIndexType()), + SmallVector(lbs.size(), result.location)); bodyBlock.addArguments( TypeRange(outputs), SmallVector(outputs.size(), result.location)); - ForeachThreadOp::ensureTerminator(*bodyRegion, builder, result.location); + ForallOp::ensureTerminator(*bodyRegion, builder, result.location); } // Builder that takes a bodyBuilder lambda. -void ForeachThreadOp::build( - mlir::OpBuilder &builder, mlir::OperationState &result, ValueRange outputs, - ValueRange numThreads, ArrayRef mapping, +void ForallOp::build( + mlir::OpBuilder &builder, mlir::OperationState &result, ValueRange lbs, + ValueRange ubs, ValueRange steps, ValueRange outputs, + ArrayRef mapping, function_ref bodyBuilder) { - result.addOperands(numThreads); + result.addOperands(lbs); + result.addOperands(ubs); + result.addOperands(steps); result.addOperands(outputs); - result.addAttribute(ForeachThreadOp::getMappingAttrName(result.name), + result.addAttribute(ForallOp::getMappingAttrName(result.name), builder.getArrayAttr(mapping)); result.addAttribute( "operand_segment_sizes", - builder.getDenseI32ArrayAttr({static_cast(numThreads.size()), + builder.getDenseI32ArrayAttr({static_cast(lbs.size()), + static_cast(ubs.size()), + static_cast(steps.size()), static_cast(outputs.size())})); result.addTypes(TypeRange(outputs)); @@ -1277,9 +1296,8 @@ builder.createBlock(bodyRegion); Block &bodyBlock = bodyRegion->front(); // Add block arguments for indices and outputs. - bodyBlock.addArguments( - SmallVector(numThreads.size(), builder.getIndexType()), - SmallVector(numThreads.size(), result.location)); + bodyBlock.addArguments(SmallVector(lbs.size(), builder.getIndexType()), + SmallVector(lbs.size(), result.location)); bodyBlock.addArguments( TypeRange(outputs), SmallVector(outputs.size(), result.location)); @@ -1287,32 +1305,40 @@ builder.setInsertionPointToStart(&bodyBlock); bodyBuilder(builder, result.location, bodyBlock.getArguments()); #ifndef NDEBUG - auto terminator = - llvm::dyn_cast(bodyBlock.getTerminator()); + auto terminator = llvm::dyn_cast(bodyBlock.getTerminator()); assert(terminator && - "expected bodyBuilder to create PerformConcurrentlyOp terminator"); + "expected bodyBuilder to create InParallelOp terminator"); #endif // NDEBUG } +// Checks if the lbs are zeros and steps are ones. +bool ForallOp::isNormalized() { + return llvm::all_of( + getLowerBound(), + [](Value v) { return matchPattern(v, mlir::m_Zero()); }) && + llvm::all_of(getStep(), + [](Value v) { return matchPattern(v, m_One()); }); +} + // The ensureTerminator method generated by SingleBlockImplicitTerminator is // unaware of the fact that our terminator also needs a region to be // well-formed. We override it here to ensure that we do the right thing. -void ForeachThreadOp::ensureTerminator(Region ®ion, OpBuilder &builder, - Location loc) { - OpTrait::SingleBlockImplicitTerminator::Impl< - ForeachThreadOp>::ensureTerminator(region, builder, loc); +void ForallOp::ensureTerminator(Region ®ion, OpBuilder &builder, + Location loc) { + OpTrait::SingleBlockImplicitTerminator::Impl< + ForallOp>::ensureTerminator(region, builder, loc); auto terminator = - llvm::dyn_cast(region.front().getTerminator()); + llvm::dyn_cast(region.front().getTerminator()); if (terminator.getRegion().empty()) builder.createBlock(&terminator.getRegion()); } -PerformConcurrentlyOp ForeachThreadOp::getTerminator() { - return cast(getBody()->getTerminator()); +InParallelOp ForallOp::getTerminator() { + return cast(getBody()->getTerminator()); } /// Helper to sort `values` according to matching `keys`. -SmallVector ForeachThreadOp::getValuesSortedByKey( +SmallVector ForallOp::getValuesSortedByKey( ArrayRef keys, ValueRange values, llvm::function_ref compare) { if (keys.empty()) @@ -1328,28 +1354,27 @@ return res; } -ForeachThreadOp mlir::scf::getForeachThreadOpThreadIndexOwner(Value val) { +ForallOp mlir::scf::getForallOpThreadIndexOwner(Value val) { auto tidxArg = val.dyn_cast(); if (!tidxArg) - return ForeachThreadOp(); + return ForallOp(); assert(tidxArg.getOwner() && "unlinked block argument"); auto *containingOp = tidxArg.getOwner()->getParentOp(); - return dyn_cast(containingOp); + return dyn_cast(containingOp); } namespace { -/// Fold tensor.dim(foreach_thread shared_outs(... = %t)) to tensor.dim(%t). -struct DimOfForeachThreadOp : public OpRewritePattern { +/// Fold tensor.dim(forall shared_outs(... = %t)) to tensor.dim(%t). +struct DimOfForallOp : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(tensor::DimOp dimOp, PatternRewriter &rewriter) const final { - auto foreachThreadOp = dimOp.getSource().getDefiningOp(); - if (!foreachThreadOp) + auto forallOp = dimOp.getSource().getDefiningOp(); + if (!forallOp) return failure(); Value sharedOut = - foreachThreadOp.getTiedOpOperand(dimOp.getSource().cast()) - ->get(); + forallOp.getTiedOpOperand(dimOp.getSource().cast())->get(); rewriter.updateRootInPlace( dimOp, [&]() { dimOp.getSourceMutable().assign(sharedOut); }); return success(); @@ -1357,29 +1382,29 @@ }; } // namespace -void ForeachThreadOp::getCanonicalizationPatterns(RewritePatternSet &results, - MLIRContext *context) { - results.add(context); +void ForallOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.add(context); } //===----------------------------------------------------------------------===// -// PerformConcurrentlyOp +// InParallelOp //===----------------------------------------------------------------------===// -// Build a PerformConcurrentlyOp with mixed static and dynamic entries. -void PerformConcurrentlyOp::build(OpBuilder &b, OperationState &result) { +// Build a InParallelOp with mixed static and dynamic entries. +void InParallelOp::build(OpBuilder &b, OperationState &result) { OpBuilder::InsertionGuard g(b); Region *bodyRegion = result.addRegion(); b.createBlock(bodyRegion); } -LogicalResult PerformConcurrentlyOp::verify() { - scf::ForeachThreadOp foreachThreadOp = - dyn_cast(getOperation()->getParentOp()); - if (!foreachThreadOp) - return this->emitOpError("expected foreach_thread op parent"); +LogicalResult InParallelOp::verify() { + scf::ForallOp forallOp = + dyn_cast(getOperation()->getParentOp()); + if (!forallOp) + return this->emitOpError("expected forall op parent"); - // TODO: PerformConcurrentlyOpInterface. + // TODO: InParallelOpInterface. for (Operation &op : getRegion().front().getOperations()) { if (!isa(op)) { return this->emitOpError("expected only ") @@ -1388,14 +1413,14 @@ // Verify that inserts are into out block arguments. Value dest = cast(op).getDest(); - ArrayRef regionOutArgs = foreachThreadOp.getRegionOutArgs(); + ArrayRef regionOutArgs = forallOp.getRegionOutArgs(); if (!llvm::is_contained(regionOutArgs, dest)) return op.emitOpError("may only insert into an output block argument"); } return success(); } -void PerformConcurrentlyOp::print(OpAsmPrinter &p) { +void InParallelOp::print(OpAsmPrinter &p) { p << " "; p.printRegion(getRegion(), /*printEntryBlockArgs=*/false, @@ -1403,8 +1428,7 @@ p.printOptionalAttrDict(getOperation()->getAttrs()); } -ParseResult PerformConcurrentlyOp::parse(OpAsmParser &parser, - OperationState &result) { +ParseResult InParallelOp::parse(OpAsmParser &parser, OperationState &result) { auto &builder = parser.getBuilder(); SmallVector regionOperands; @@ -1422,11 +1446,11 @@ return success(); } -OpResult PerformConcurrentlyOp::getParentResult(int64_t idx) { +OpResult InParallelOp::getParentResult(int64_t idx) { return getOperation()->getParentOp()->getResult(idx); } -SmallVector PerformConcurrentlyOp::getDests() { +SmallVector InParallelOp::getDests() { return llvm::to_vector<4>( llvm::map_range(getYieldingOps(), [](Operation &op) { // Add new ops here as needed. @@ -1435,7 +1459,7 @@ })); } -llvm::iterator_range PerformConcurrentlyOp::getYieldingOps() { +llvm::iterator_range InParallelOp::getYieldingOps() { return getRegion().front().getOperations(); } diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp --- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp @@ -1003,51 +1003,51 @@ }; /// Return `true` if the given loop may have 0 iterations. -bool mayHaveZeroIterations(scf::ForeachThreadOp foreachThreadOp) { - int64_t p = 1; - for (Value v : foreachThreadOp.getNumThreads()) { - if (std::optional c = getConstantIntValue(v)) { - p *= *c; - } else { +bool mayHaveZeroIterations(scf::ForallOp forallOp) { + for (auto [lb, ub] : + llvm::zip(forallOp.getLowerBound(), forallOp.getUpperBound())) { + std::optional lbConst = getConstantIntValue(lb); + std::optional ubConst = getConstantIntValue(ub); + if (!lbConst.has_value() || !ubConst.has_value() || *lbConst >= *ubConst) { return true; } } - return p == 0; + return false; } -/// Bufferization of ForeachThreadOp. This also bufferizes the terminator of the -/// region. There are op interfaces for the terminators (PerformConcurrentlyOp +/// Bufferization of ForallOp. This also bufferizes the terminator of the +/// region. There are op interfaces for the terminators (InParallelOp /// and ParallelInsertSliceOp), but these are only used during analysis. Not /// for bufferization. -struct ForeachThreadOpInterface - : public BufferizableOpInterface::ExternalModel { +struct ForallOpInterface + : public BufferizableOpInterface::ExternalModel { bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, const AnalysisState &state) const { - auto foreachThreadOp = cast(op); + auto forallOp = cast(op); // If the loop has zero iterations, the results of the op are their // corresponding shared_outs, meaning that the shared_outs bufferize to a // read. - if (mayHaveZeroIterations(foreachThreadOp)) + if (mayHaveZeroIterations(forallOp)) return true; - // scf::ForeachThreadOp alone doesn't bufferize to a memory read, one of the + // scf::ForallOp alone doesn't bufferize to a memory read, one of the // uses of its matching bbArg may. - return state.isValueRead(foreachThreadOp.getTiedBlockArgument(&opOperand)); + return state.isValueRead(forallOp.getTiedBlockArgument(&opOperand)); } bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand, const AnalysisState &state) const { - // Outputs of scf::ForeachThreadOps are always considered as a write. + // Outputs of scf::ForallOps are always considered as a write. return true; } AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand, const AnalysisState &state) const { - auto foreachThreadOp = cast(op); - return {{{foreachThreadOp.getTiedOpResult(&opOperand), - BufferRelation::Equivalent}}}; + auto forallOp = cast(op); + return { + {{forallOp.getTiedOpResult(&opOperand), BufferRelation::Equivalent}}}; } bool isWritable(Operation *op, Value value, @@ -1058,12 +1058,12 @@ LogicalResult bufferize(Operation *op, RewriterBase &rewriter, const BufferizationOptions &options) const { OpBuilder::InsertionGuard guard(rewriter); - auto foreachThreadOp = cast(op); - int64_t rank = foreachThreadOp.getRank(); + auto forallOp = cast(op); + int64_t rank = forallOp.getRank(); // Get buffers for all output operands. SmallVector buffers; - for (Value out : foreachThreadOp.getOutputs()) { + for (Value out : forallOp.getOutputs()) { FailureOr buffer = getBuffer(rewriter, out, options); if (failed(buffer)) return failure(); @@ -1071,35 +1071,33 @@ } // Use buffers instead of block arguments. - rewriter.setInsertionPointToStart(foreachThreadOp.getBody()); - for (const auto &it : - llvm::zip(foreachThreadOp.getBody()->getArguments().drop_front(rank), - buffers)) { + rewriter.setInsertionPointToStart(forallOp.getBody()); + for (const auto &it : llvm::zip( + forallOp.getBody()->getArguments().drop_front(rank), buffers)) { BlockArgument bbArg = std::get<0>(it); Value buffer = std::get<1>(it); Value bufferAsTensor = - rewriter.create(foreachThreadOp.getLoc(), buffer); + rewriter.create(forallOp.getLoc(), buffer); bbArg.replaceAllUsesWith(bufferAsTensor); } - // Create new ForeachThreadOp without any results and drop the automatically + // Create new ForallOp without any results and drop the automatically // introduced terminator. - rewriter.setInsertionPoint(foreachThreadOp); - ForeachThreadOp newForeachThreadOp; - newForeachThreadOp = rewriter.create( - foreachThreadOp.getLoc(), /*outputs=*/ValueRange(), - foreachThreadOp.getNumThreads(), foreachThreadOp.getMapping()); + rewriter.setInsertionPoint(forallOp); + ForallOp newForallOp; + newForallOp = rewriter.create( + forallOp.getLoc(), forallOp.getLowerBound(), forallOp.getUpperBound(), + forallOp.getStep(), /*outputs=*/ValueRange(), forallOp.getMapping()); - newForeachThreadOp.getBody()->getTerminator()->erase(); + newForallOp.getBody()->getTerminator()->erase(); // Move over block contents of the old op. SmallVector replacementBbArgs; - replacementBbArgs.append( - newForeachThreadOp.getBody()->getArguments().begin(), - newForeachThreadOp.getBody()->getArguments().end()); - replacementBbArgs.append(foreachThreadOp.getOutputs().size(), Value()); - rewriter.mergeBlocks(foreachThreadOp.getBody(), - newForeachThreadOp.getBody(), replacementBbArgs); + replacementBbArgs.append(newForallOp.getBody()->getArguments().begin(), + newForallOp.getBody()->getArguments().end()); + replacementBbArgs.append(forallOp.getOutputs().size(), Value()); + rewriter.mergeBlocks(forallOp.getBody(), newForallOp.getBody(), + replacementBbArgs); // Remove the old op and replace all of its uses. replaceOpWithBufferizedValues(rewriter, op, buffers); @@ -1110,34 +1108,51 @@ FailureOr getBufferType(Operation *op, Value value, const BufferizationOptions &options, const DenseMap &fixedTypes) const { - auto foreachThreadOp = cast(op); + auto forallOp = cast(op); if (auto bbArg = value.dyn_cast()) // A tensor block argument has the same bufferized type as the // corresponding output operand. return bufferization::getBufferType( - foreachThreadOp.getTiedOpOperand(bbArg)->get(), options, fixedTypes); + forallOp.getTiedOpOperand(bbArg)->get(), options, fixedTypes); // The bufferized result type is the same as the bufferized type of the // corresponding output operand. return bufferization::getBufferType( - foreachThreadOp.getOutputs()[value.cast().getResultNumber()], + forallOp.getOutputs()[value.cast().getResultNumber()], options, fixedTypes); } bool isRepetitiveRegion(Operation *op, unsigned index) const { - auto foreachThreadOp = cast(op); - // This op is not repetitive if it has just a single thread. - return !llvm::all_of(foreachThreadOp.getNumThreads(), [](Value v) { - return getConstantIntValue(v) == static_cast(1); - }); + auto forallOp = cast(op); + + // This op is not repetitive if it has just a single step. + for (auto [lb, ub, step] : + llvm::zip(forallOp.getLowerBound(), forallOp.getUpperBound(), + forallOp.getStep())) { + std::optional lbConstant = getConstantIntValue(lb); + if (!lbConstant) + return true; + + std::optional ubConstant = getConstantIntValue(ub); + if (!ubConstant) + return true; + + std::optional stepConstant = getConstantIntValue(step); + if (!stepConstant) + return true; + + if (*lbConstant + *stepConstant < *ubConstant) + return true; + } + return false; } }; -/// Nothing to do for PerformConcurrentlyOp. -struct PerformConcurrentlyOpInterface - : public BufferizableOpInterface::ExternalModel< - PerformConcurrentlyOpInterface, PerformConcurrentlyOp> { +/// Nothing to do for InParallelOp. +struct InParallelOpInterface + : public BufferizableOpInterface::ExternalModel { LogicalResult bufferize(Operation *op, RewriterBase &b, const BufferizationOptions &options) const { llvm_unreachable("op does not have any tensor OpOperands / OpResults"); @@ -1156,9 +1171,8 @@ ExecuteRegionOp::attachInterface(*ctx); ForOp::attachInterface(*ctx); IfOp::attachInterface(*ctx); - ForeachThreadOp::attachInterface(*ctx); - PerformConcurrentlyOp::attachInterface( - *ctx); + ForallOp::attachInterface(*ctx); + InParallelOp::attachInterface(*ctx); WhileOp::attachInterface(*ctx); YieldOp::attachInterface(*ctx); }); diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp --- a/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp +++ b/mlir/lib/Dialect/SCF/Transforms/LoopCanonicalization.cpp @@ -177,13 +177,12 @@ } return failure(); } - if (scf::ForeachThreadOp foreachThreadOp = - scf::getForeachThreadOpThreadIndexOwner(iv)) { - for (int64_t idx = 0; idx < foreachThreadOp.getRank(); ++idx) { - if (foreachThreadOp.getThreadIndices()[idx] == iv) { - lb = OpBuilder(iv.getContext()).getIndexAttr(0); - ub = foreachThreadOp.getNumThreads()[idx]; - step = OpBuilder(iv.getContext()).getIndexAttr(1); + if (scf::ForallOp forallOp = scf::getForallOpThreadIndexOwner(iv)) { + for (int64_t idx = 0; idx < forallOp.getRank(); ++idx) { + if (forallOp.getInductionVar(idx) == iv) { + lb = forallOp.getLowerBound()[idx]; + ub = forallOp.getUpperBound()[idx]; + step = forallOp.getStep()[idx]; return success(); } } diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir @@ -144,14 +144,14 @@ // CHECK-SAME: %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref // CHECK-SAME: %[[sz:[0-9a-zA-Z]*]]: index func.func @parallel_insert_slice( - %t: tensor {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true}, - %sz: index) - -> (tensor) -{ + %t: tensor {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true}, + %sz: index) -> (tensor) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %f0 = arith.constant 0.0: f32 %c512 = arith.constant 512 : index - %r1 = scf.foreach_thread (%iv) in (%c512) shared_outs(%o = %t) -> (tensor) { + %r1 = scf.forall (%iv) = (%c0) to (%c512) step (%c1) shared_outs(%o = %t) -> (tensor) { // tensor.empty itself does not alloc but forwards to the insert_slice. // EmptyTensorOpElimination replaces the tensor.empty with an inplace // extract_slice. @@ -162,7 +162,7 @@ %f = linalg.fill ins(%f0 : f32) outs(%a : tensor) -> tensor // Self-copy canonicalizes away later. - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %f into %o[42][%sz][1]: tensor into tensor } } diff --git a/mlir/test/Dialect/GPU/transform-gpu-failing.mlir b/mlir/test/Dialect/GPU/transform-gpu-failing.mlir --- a/mlir/test/Dialect/GPU/transform-gpu-failing.mlir +++ b/mlir/test/Dialect/GPU/transform-gpu-failing.mlir @@ -14,14 +14,15 @@ // ----- func.func @map_nested_foreach_to_threads_excessive_threads(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> { - %one = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c900 = arith.constant 900 : index %c9 = arith.constant 9 : index %c7 = arith.constant 7 : index - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { - scf.foreach_thread (%i, %j) in (%c7, %c900) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c900) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 @@ -30,10 +31,10 @@ gpu.terminator } - %name2 = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name2 = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { - scf.foreach_thread (%i, %j) in (%c7, %c9) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c9) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 @@ -55,14 +56,15 @@ // ----- func.func @map_nested_foreach_to_threads_fewer_threads(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> { - %one = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c900 = arith.constant 900 : index %c9 = arith.constant 9 : index %c7 = arith.constant 7 : index - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { - scf.foreach_thread (%i, %j) in (%c7, %c900) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c900) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 @@ -71,10 +73,10 @@ gpu.terminator } - %name2 = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name2 = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { - scf.foreach_thread (%i, %j) in (%c7, %c9) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c9) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 @@ -89,19 +91,20 @@ transform.sequence failures(propagate) { ^bb1(%arg0: !pdl.operation): %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!pdl.operation) -> !pdl.operation - // expected-error @below {{The requested GPU threads are fewer than the number of loop trip counts. Try to tile scf.foreach_thread before mapping or set small blockDim.}} + // expected-error @below {{The requested GPU threads are fewer than the number of loop trip counts. Try to tile scf.forall before mapping or set small blockDim.}} transform.gpu.map_nested_foreach_to_threads %funcop { blockDim = [128, 4, 1] } } // ----- func.func @map_nested_foreach_to_threads_dynamic_trip_count(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token, %c9 : index, %c7 : index) -> memref<2 x 32 x f32> { - %one = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c900 = arith.constant 900 : index - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { - scf.foreach_thread (%i, %j) in (%c7, %c900) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c900) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 @@ -122,9 +125,9 @@ // ----- func.func @map_nested_foreach_to_threads_not_buffer(%x: tensor<32x32xf32>, %y: tensor<32x32xf32>, %z: tensor<32x32xf32>, %stream : !gpu.async.token) { - %one = arith.constant 1 : index - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %c1 = arith.constant 1 : index + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { %t = linalg.matmul ins(%x, %y: tensor<32x32xf32>, tensor<32x32xf32>) outs(%z : tensor<32x32xf32>) -> tensor<32x32xf32> gpu.terminator @@ -135,9 +138,9 @@ transform.sequence failures(propagate) { ^bb1(%arg0: !pdl.operation): %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0 : (!pdl.operation) -> !pdl.operation - %foreach, %tiled = transform.structured.tile_to_foreach_thread_op %matmul num_threads [10, 20, 30] (mapping = [ #gpu.thread, #gpu.thread, #gpu.thread ] ) + %foreach, %tiled = transform.structured.tile_to_forall_op %matmul num_threads [10, 20, 30] (mapping = [ #gpu.thread, #gpu.thread, #gpu.thread ] ) %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!pdl.operation) -> !pdl.operation - // expected-error @below {{only bufferized scf.foreach_thread lowers to gpu.thread_id}} + // expected-error @below {{only bufferized scf.forall lowers to gpu.thread_id}} transform.gpu.map_nested_foreach_to_threads %funcop { blockDim = [128, 4, 1] } } @@ -159,22 +162,23 @@ // ----- func.func @map_foreach_to_blocks_not_unique(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> { - %one = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c900 = arith.constant 900 : index %c9 = arith.constant 9 : index %c7 = arith.constant 7 : index // expected-note @below {{when applied to this payload op}} - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { - scf.foreach_thread (%i, %j) in (%c7, %c900) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c900) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 memref.store %6, %y[%i, %j] : memref<2 x 32 x f32> } { mapping = [#gpu.thread, #gpu.thread] } - scf.foreach_thread (%i, %j) in (%c7, %c9) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c9) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 @@ -189,7 +193,7 @@ transform.sequence failures(propagate) { ^bb0(%arg0: !pdl.operation): %funcop = transform.structured.match ops{["gpu.launch"]} in %arg0 : (!pdl.operation) -> !pdl.operation - // expected-error @below {{could not find a unique topLevel scf.foreach_thread}} + // expected-error @below {{could not find a unique topLevel scf.forall}} %1 = transform.gpu.map_foreach_to_blocks %funcop } @@ -197,19 +201,20 @@ // expected-note @below {{when applied to this payload op}} func.func @map_foreach_to_blocks_large_loop(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> { - %one = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c65537 = arith.constant 65536 : index %c9 = arith.constant 9 : index %c7 = arith.constant 7 : index - scf.foreach_thread (%i, %j) in (%c7, %c65537) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c65537) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 memref.store %6, %y[%i, %j] : memref<2 x 32 x f32> } { mapping = [#gpu.thread, #gpu.thread] } - scf.foreach_thread (%i, %j) in (%c7, %c9) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c9) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 @@ -222,16 +227,17 @@ transform.sequence failures(propagate) { ^bb0(%arg0: !pdl.operation): %funcop = transform.structured.match ops{["func.func"]} in %arg0 : (!pdl.operation) -> !pdl.operation - // expected-error @below {{could not find a unique topLevel scf.foreach_thread}} + // expected-error @below {{could not find a unique topLevel scf.forall}} %1 = transform.gpu.map_foreach_to_blocks %funcop { generate_gpu_launch } } // ----- func.func @map_foreach_to_blocks_large_loop(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> { - %one = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c65535 = arith.constant 65535 : index - scf.foreach_thread (%i, %j) in (%c65535, %c65535) { + scf.forall (%i, %j) = (%c0, %c0) to (%c65535, %c65535) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 @@ -251,12 +257,13 @@ !type = memref<32x32xf32> func.func @saxpy2d_singleloop(%x: !type, %y: !type, %stream : !gpu.async.token) -> !type { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c32 = arith.constant 32 : index - %one = arith.constant 1 : index - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { - scf.foreach_thread (%i, %j) in (%c32, %c32) { + scf.forall (%i, %j) = (%c0, %c0) to (%c32, %c32) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : !type %5 = memref.load %y[%i, %j] : !type %6 = arith.mulf %4, %5 : f32 @@ -277,9 +284,9 @@ // ----- func.func @tiling_buffer_semantic_op(%x: memref<32x32xf32>, %y: memref<32x32xf32>, %stream : !gpu.async.token) { - %one = arith.constant 1 : index - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %c1 = arith.constant 1 : index + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { // expected-error @below {{'linalg.generic' op must have "tensor semantic" for tiling}} // expected-note @below {{when applied to this op}} @@ -300,6 +307,6 @@ transform.sequence failures(propagate) { ^bb1(%arg0: !pdl.operation): %matmul = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!pdl.operation) -> !pdl.operation - // expected-error @below {{transform.structured.tile_to_foreach_thread_op failed to apply}} - %foreach, %tiled = transform.structured.tile_to_foreach_thread_op %matmul num_threads [10, 20, 30] (mapping = [ #gpu.thread, #gpu.thread, #gpu.thread ] ) + // expected-error @below {{transform.structured.tile_to_forall_op failed to apply}} + %foreach, %tiled = transform.structured.tile_to_forall_op %matmul num_threads [10, 20, 30] (mapping = [ #gpu.thread, #gpu.thread, #gpu.thread ] ) } diff --git a/mlir/test/Dialect/GPU/transform-gpu.mlir b/mlir/test/Dialect/GPU/transform-gpu.mlir --- a/mlir/test/Dialect/GPU/transform-gpu.mlir +++ b/mlir/test/Dialect/GPU/transform-gpu.mlir @@ -10,6 +10,7 @@ func.func @saxpy2dblock(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { %c9 = arith.constant 9 : index %c7 = arith.constant 7 : index + %zero = arith.constant 0 : index %one = arith.constant 1 : index // CHECK: gpu.launch // CHECK: %[[BLKX:.*]] = gpu.block_id x @@ -19,7 +20,7 @@ %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) { - scf.foreach_thread (%i, %j) in (%c7, %c9) { + scf.forall (%i, %j) = (%zero, %zero) to (%c7, %c9) step (%one, %one) { %4 = memref.load %x[%i, %j] : !type %5 = memref.load %y[%i, %j] : !type %6 = math.fma %alpha, %4, %5 : f32 @@ -46,6 +47,7 @@ // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<2x32xf32> // CHECK-SAME: %[[ARGT:[0-9a-z]+]]: memref<32xf32> func.func @saxpy2d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { + %zero = arith.constant 0 : index %one = arith.constant 1 : index %c12 = arith.constant 12 : index %c9 = arith.constant 9 : index @@ -68,13 +70,13 @@ %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) { - scf.foreach_thread (%i, %j) in (%c7, %c9) { + scf.forall (%i, %j) = (%zero, %zero) to (%c7, %c9) step (%one, %one) { %4 = memref.load %x[%i, %j] : !type %5 = memref.load %y[%i, %j] : !type %6 = math.fma %alpha, %4, %5 : f32 memref.store %6, %y[%i, %j] : !type } { mapping = [#gpu.thread, #gpu.thread]} - scf.foreach_thread (%i) in (%c12) { + scf.forall (%i) = (%zero) to (%c12) step (%one) { %7 = memref.load %t[%i] : !type1d %8 = arith.addf %alpha, %7 : f32 memref.store %8, %t[%i] : !type1d @@ -98,22 +100,24 @@ // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<32x64x4x32xf32> // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<32x64x4x32xf32> func.func @saxpy4d(%x: !type4d, %y: !type4d, %alpha : f32) -> !type4d { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c4 = arith.constant 4 : index %c32 = arith.constant 32 : index %c64 = arith.constant 64 : index - %c4 = arith.constant 4 : index -// CHECK: %[[C32:.*]] = arith.constant 32 : index -// CHECK: %[[C64:.*]] = arith.constant 64 : index -// CHECK: %[[C4:.*]] = arith.constant 4 : index -// CHECK: %[[C1:.*]] = arith.constant 1 : index -// CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C32]], %{{.*}} = %[[C64]], %{{.*}} = %[[C1]]) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C32]], %{{.*}} = %[[C4]], %{{.*}} = %[[C1]]) -// CHECK: %[[BLKX:.*]] = gpu.block_id x -// CHECK: %[[BLKY:.*]] = gpu.block_id y -// CHECK: %[[TIDX:.*]] = gpu.thread_id x -// CHECK: %[[TIDY:.*]] = gpu.thread_id y -// CHECK: memref.load %[[ARGX]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]] -// CHECK: memref.load %[[ARGY]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]] - scf.foreach_thread (%i, %j) in (%c32, %c64) { - scf.foreach_thread (%k, %l) in (%c4, %c32) { +// CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index +// CHECK-DAG: %[[C64:.*]] = arith.constant 64 : index +// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index +// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index +// CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C32]], %{{.*}} = %[[C64]], %{{.*}} = %[[C1]]) threads(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C32]], %{{.*}} = %[[C4]], %{{.*}} = %[[C1]]) +// CHECK: %[[BLKX:.*]] = gpu.block_id x +// CHECK: %[[BLKY:.*]] = gpu.block_id y +// CHECK: %[[TIDX:.*]] = gpu.thread_id x +// CHECK: %[[TIDY:.*]] = gpu.thread_id y +// CHECK: memref.load %[[ARGX]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]] +// CHECK: memref.load %[[ARGY]][%[[BLKX]], %[[BLKY]], %[[TIDY]], %[[TIDX]]] + scf.forall (%i, %j) = (%c0, %c0) to (%c32, %c64) step (%c1, %c1) { + scf.forall (%k, %l) = (%c0, %c0) to (%c4, %c32) step (%c1, %c1) { %4 = memref.load %x[%i, %j, %k, %l] : !type4d %5 = memref.load %y[%i, %j, %k, %l] : !type4d %6 = math.fma %alpha, %4, %5 : f32 @@ -137,6 +141,7 @@ // CHECK-LABEL: func.func @saxpy2d_no_barrier( func.func @saxpy2d_no_barrier(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { + %zero = arith.constant 0 : index %one = arith.constant 1 : index %c12 = arith.constant 12 : index %c9 = arith.constant 9 : index @@ -146,7 +151,7 @@ %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) { - scf.foreach_thread (%i, %j) in (%c7, %c9) { + scf.forall (%i, %j) = (%zero, %zero) to (%c7, %c9) step (%one, %one) { %4 = memref.load %x[%i, %j] : !type %5 = memref.load %y[%i, %j] : !type %6 = math.fma %alpha, %4, %5 : f32 @@ -170,15 +175,16 @@ // CHECK-SAME: %[[ARGX:[0-9a-z]+]]: memref<32x32xf32> // CHECK-SAME: %[[ARGY:[0-9a-z]+]]: memref<32x32xf32> func.func @saxpy2d_singleloop(%x: !type, %y: !type, %stream : !gpu.async.token) -> !type { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c32 = arith.constant 32 : index - %one = arith.constant 1 : index - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { // CHECK: %[[TIDX:.*]] = gpu.thread_id x // CHECK: memref.load %[[ARGX]][%[[TIDX]], %[[TIDX]]] // CHECK: memref.load %[[ARGY]][%[[TIDX]], %[[TIDX]]] - scf.foreach_thread (%i) in (%c32) { + scf.forall (%i) = (%c0) to (%c32) step (%c1) { %4 = memref.load %x[%i, %i] : !type %5 = memref.load %y[%i, %i] : !type %6 = arith.mulf %4, %5 : f32 @@ -202,16 +208,17 @@ // CHECK-LABEL: func.func @saxpy3d_fold_id_z( func.func @saxpy3d_fold_id_z(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { - %one = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c12 = arith.constant 12 : index %c9 = arith.constant 9 : index %c7 = arith.constant 7 : index // CHECK: %[[C0:.+]] = arith.constant 0 : index // CHECK-NOT: gpu.thread_id z - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { - scf.foreach_thread (%i, %j, %k) in (%one, %c7, %c9) { + scf.forall (%i, %j, %k) = (%c0, %c0, %c0) to (%c1, %c7, %c9) step (%c1, %c1, %c1) { // CHECK: memref.load %{{.*}}[%[[C0]], // CHECK: memref.load %{{.*}}[%[[C0]], %4 = memref.load %x[%i, %j, %k] : !type @@ -238,23 +245,24 @@ // CHECK-LABEL: func.func @map_multi_level( func.func @map_multi_level(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type { - %one = arith.constant 1 : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c12 = arith.constant 12 : index %c9 = arith.constant 9 : index %c7 = arith.constant 7 : index // check that the thread level got distributed but not the warp level. // CHECK-NOT: {mapping = #gpu.thread // CHECK: {mapping = [#gpu.warp]} - %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) - threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) + %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) + threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) { - scf.foreach_thread (%i, %j) in (%c7, %c9) { + scf.forall (%i, %j) = (%c0, %c0) to (%c7, %c9) step (%c1, %c1) { %4 = memref.load %x[%i, %j] : !type %5 = memref.load %y[%i, %j] : !type %6 = math.fma %alpha, %4, %5 : f32 memref.store %6, %y[%i, %j] : !type } { mapping = [#gpu.thread, #gpu.thread]} - scf.foreach_thread (%i) in (%c12) { + scf.forall (%i) = (%c0) to (%c12) step (%c1) { %7 = memref.load %t[%i] : !type1d %8 = arith.addf %alpha, %7 : f32 memref.store %8, %t[%i] : !type1d diff --git a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir --- a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir +++ b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir @@ -876,14 +876,16 @@ // ----- func.func @reduce_dispatch_0() -> tensor<4x2xf32> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index %c4 = arith.constant 4 : index %cst = arith.constant 0.000000e+00 : f32 %0 = tensor.empty() : tensor<4x2xf32> - %res = scf.foreach_thread (%arg0, %arg1) in (%c4, %c2) shared_outs(%o = %0) -> (tensor<4x2xf32>) { + %res = scf.forall (%arg0, %arg1) = (%c0, %c0) to (%c4, %c2) step (%c1, %c1) shared_outs(%o = %0) -> (tensor<4x2xf32>) { %1 = tensor.empty() : tensor<1x1xf32> %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x1xf32>) -> tensor<1x1xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { // CHECK: tensor.parallel_insert_slice %{{[0-9a-z]*}} into %{{[0-9a-z]*}} // CHECK-SAME: [%{{.*}}, %{{.*}}] [1, 1] [1, 1] : tensor into tensor<4x2xf32> tensor.parallel_insert_slice %2 into %o[%arg0, %arg1] [1, 1] [1, 1] : diff --git a/mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir b/mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir --- a/mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir +++ b/mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir @@ -15,14 +15,16 @@ func.func @matmul(%A: tensor, %B: tensor, %C: tensor) -> tensor { // CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index // CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index - // CHECK: scf.foreach_thread ({{.*}}) in (%[[C10]], %[[C20]]) shared_outs(%[[C_BLK:.*]] = %[[C]]) -> (tensor) { + // CHECK: scf.forall + // CHECK-SAME: to (%[[C10]], %[[C20]]) + // CHECK-SAME: shared_outs(%[[C_BLK:.*]] = %[[C]]) -> (tensor) { // CHECK: %[[tA:.*]] = tensor.extract_slice %[[A]]{{.*}} : tensor to tensor // CHECK: %[[tB:.*]] = tensor.extract_slice %[[B]]{{.*}} : tensor to tensor // CHECK: %[[tC:.*]] = tensor.extract_slice %[[C_BLK]]{{.*}} : tensor to tensor // CHECK: %[[RES:.*]] = linalg.matmul // CHECK-SAME: ins(%[[tA]], %[[tB]] : tensor, tensor) // CHECK-SAME: outs(%[[tC]] : tensor) -> tensor - // CHECK: scf.foreach_thread.perform_concurrently { + // CHECK: scf.forall.in_parallel { // CHECK-NEXT: tensor.parallel_insert_slice %[[RES]] into %[[C_BLK]]{{.*}} : // CHECK-SAME: tensor into tensor // CHECK-NEXT: } @@ -35,7 +37,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 20] (mapping = [ #gpu.thread, #gpu.thread ] ) + %1:2 = transform.structured.tile_to_forall_op %0 num_threads [10, 20] (mapping = [ #gpu.thread, #gpu.thread ] ) } } @@ -60,12 +62,14 @@ // CHECK-DAG: %[[N:.+]] = tensor.dim %[[B]], %c1 : // CHECK-DAG: %[[NT0:.+]] = affine.apply #[[$map0]]()[%[[M]], %[[tile_size_1]]] // CHECK-DAG: %[[NT1:.+]] = affine.apply #[[$map0]]()[%[[N]], %[[tile_size_2]]] - // CHECK: scf.foreach_thread (%[[IV0:.+]], %[[IV1:.+]]) in (%[[NT0]], %[[NT1]]) shared_outs(%[[C_BLK:.*]] = %[[C]]) + // CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) = + // CHECK-SAME: to (%[[NT0]], %[[NT1]]) + // CHECK-SAME: shared_outs(%[[C_BLK:.*]] = %[[C]]) // CHECK: tensor.extract_slice %[[A]] // CHECK: tensor.extract_slice %[[B]] // CHECK: tensor.extract_slice %[[C_BLK]] // CHECK: linalg.matmul - // CHECK: scf.foreach_thread.perform_concurrently + // CHECK: scf.forall.in_parallel // CHECK-NEXT: tensor.parallel_insert_slice %tile_size_1 = "test.dummy"() : () -> (index) %tile_size_2 = "test.dummy"() : () -> (index) @@ -78,7 +82,7 @@ ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation %sz = transform.structured.match ops{["test.dummy"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes %sz + %1:2 = transform.structured.tile_to_forall_op %0 tile_sizes %sz } // ----- @@ -97,7 +101,9 @@ func.func @matmul_static(%A: tensor<100x200xf32>, %B: tensor<200x300xf32>, %C: tensor<100x300xf32>) -> tensor<100x300xf32> { // CHECK-DAG: %[[c10:.+]] = arith.constant 10 : index // CHECK-DAG: %[[c21:.+]] = arith.constant 21 : index - // CHECK: scf.foreach_thread (%[[IV0:.+]], %[[IV1:.+]]) in (%[[c10]], %[[c21]]) shared_outs(%[[C_BLK:.*]] = %[[C]]) + // CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) = + // CHECK-SAME: to (%[[c10]], %[[c21]]) + // CHECK-SAME: shared_outs(%[[C_BLK:.*]] = %[[C]]) // CHECK: %[[TSMIN:.+]] = affine.min #[[$map0]](%[[IV1]]) // CHECK: %[[TS:.+]] = affine.max #[[$map1]](%[[TSMIN]]) // CHECK-NOT: affine.min @@ -108,7 +114,7 @@ // CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] : // CHECK: %[[tC:.+]] = tensor.extract_slice %[[C_BLK]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] : // CHECK: linalg.matmul - // CHECK: scf.foreach_thread.perform_concurrently + // CHECK: scf.forall.in_parallel // CHECK-NEXT: tensor.parallel_insert_slice %0 = linalg.matmul ins(%A, %B : tensor<100x200xf32>, tensor<200x300xf32>) outs(%C : tensor<100x300xf32>) -> (tensor<100x300xf32>) @@ -118,7 +124,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 21] + %1:2 = transform.structured.tile_to_forall_op %0 num_threads [10, 21] } @@ -140,7 +146,9 @@ // CHECK: %[[N:.+]] = tensor.dim %[[B]], %c1 : // CHECK: %[[NT0:.+]] = affine.apply #map()[%[[M]]] // CHECK: %[[NT1:.+]] = affine.apply #map1()[%[[N]]] - // CHECK: scf.foreach_thread (%[[IV0:.+]], %[[IV1:.+]]) in (%[[NT0]], %[[NT1]]) shared_outs(%[[C_BLK:.*]] = %[[C]]) + // CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) = + // CHECK-SAME: to (%[[NT0]], %[[NT1]]) + // CHECK-SAME: shared_outs(%[[C_BLK:.*]] = %[[C]]) // CHECK: %[[TS0:.+]] = affine.min #[[$map2]](%[[IV0]])[%[[M]]] // CHECK: %[[TS1:.+]] = affine.min #[[$map4]](%[[IV1]])[%[[N]]] // CHECK: %[[LB0:.+]] = affine.apply #[[$map5]](%[[IV0]]) @@ -149,7 +157,7 @@ // CHECK: tensor.extract_slice %[[B]] // CHECK: tensor.extract_slice %[[C_BLK]] // CHECK: linalg.matmul - // CHECK: scf.foreach_thread.perform_concurrently + // CHECK: scf.forall.in_parallel // CHECK-NEXT: tensor.parallel_insert_slice %0 = linalg.matmul ins(%A, %B : tensor, tensor) outs(%C : tensor) -> (tensor) @@ -159,7 +167,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [10, 20] + %1:2 = transform.structured.tile_to_forall_op %0 tile_sizes [10, 20] } // ----- @@ -177,7 +185,9 @@ func.func @matmul_tile_size_static(%A: tensor<100x200xf32>, %B: tensor<200x300xf32>, %C: tensor<100x300xf32>) -> tensor<100x300xf32> { // CHECK-DAG: %[[c10:.+]] = arith.constant 10 : // CHECK-DAG: %[[c15:.+]] = arith.constant 15 : - // CHECK: scf.foreach_thread (%[[IV0:.+]], %[[IV1:.+]]) in (%[[c10]], %[[c15]]) shared_outs(%[[C_BLK:.*]] = %[[C]]) + // CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) = + // CHECK-SAME: to (%[[c10]], %[[c15]]) + // CHECK-SAME: shared_outs(%[[C_BLK:.*]] = %[[C]]) // CHECK: %[[TS:.+]] = affine.min #[[$map0]](%[[IV1]]) // CHECK-NOT: affine.max // CHECK-NOT: affine.min @@ -187,7 +197,7 @@ // CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] : // CHECK: %[[tC:.+]] = tensor.extract_slice %[[C_BLK]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] : // CHECK: linalg.matmul - // CHECK: scf.foreach_thread.perform_concurrently + // CHECK: scf.forall.in_parallel // CHECK-NEXT: tensor.parallel_insert_slice %0 = linalg.matmul ins(%A, %B : tensor<100x200xf32>, tensor<200x300xf32>) outs(%C : tensor<100x300xf32>) -> (tensor<100x300xf32>) @@ -197,7 +207,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [10, 21] + %1:2 = transform.structured.tile_to_forall_op %0 tile_sizes [10, 21] } // ----- @@ -219,16 +229,18 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [2] ( mapping = [#gpu.thread]) + %1:2 = transform.structured.tile_to_forall_op %0 num_threads [2] ( mapping = [#gpu.thread]) } } // CHECK-DAG: #[[$map0:.+]] = affine_map<(d0) -> (d0 * 2)> // CHECK-LABEL: extract_source( // CHECK: %[[C2:.*]] = arith.constant 2 : index -// CHECK: scf.foreach_thread (%[[ARG:.*]]) in (%[[C2]]) shared_outs(%{{.*}} = %{{.*}}) -> (tensor<4xf32>) { +// CHECK: scf.forall (%[[ARG:.*]]) = +// CHECK-SAME: to (%[[C2]]) +// CHECK-SAME: shared_outs(%{{.*}} = %{{.*}}) -> (tensor<4xf32>) { // CHECK: %[[OFF:.*]] = affine.apply #[[$map0]](%[[ARG]]) -// CHECK: scf.foreach_thread.perform_concurrently { +// CHECK: scf.forall.in_parallel { // CHECK: tensor.parallel_insert_slice %{{.*}} into %{{.*}}[%[[OFF]]] [2] [1] : tensor<2xf32> into tensor<4xf32> // ----- @@ -254,12 +266,14 @@ // CHECK-DAG: %[[N:.+]] = tensor.dim %[[B]], %c1 : // CHECK-DAG: %[[NT0:.+]] = affine.apply #[[$map0]]()[%[[M]], %[[tile_size]]] // CHECK-DAG: %[[NT1:.+]] = affine.apply #[[$map1]]()[%[[N]]] - // CHECK: scf.foreach_thread (%[[IV0:.+]], %[[IV1:.+]]) in (%[[NT0]], %[[NT1]]) shared_outs(%[[C_BLK:.*]] = %[[C]]) + // CHECK: scf.forall (%[[IV0:.+]], %[[IV1:.+]]) = + // CHECK-SAME: to (%[[NT0]], %[[NT1]]) + // CHECK-SAME: shared_outs(%[[C_BLK:.*]] = %[[C]]) // CHECK: tensor.extract_slice %[[A]] // CHECK: tensor.extract_slice %[[B]] // CHECK: tensor.extract_slice %[[C_BLK]] // CHECK: linalg.matmul - // CHECK: scf.foreach_thread.perform_concurrently + // CHECK: scf.forall.in_parallel // CHECK-NEXT: tensor.parallel_insert_slice %tile_size = "test.dummy"() : () -> (index) %0 = linalg.matmul ins(%A, %B : tensor, tensor) @@ -271,7 +285,7 @@ ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation %sz = transform.structured.match ops{["test.dummy"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [%sz, 20] + %1:2 = transform.structured.tile_to_forall_op %0 tile_sizes [%sz, 20] } // ----- @@ -290,7 +304,9 @@ %OUT1: tensor<100xf32>, %OUT2: tensor<100xf32>) -> (tensor<100xf32>, tensor<100xf32>) { // CHECK-DAG: %[[c0:.+]] = arith.constant 7 : -// CHECK: scf.foreach_thread (%[[IV0:.+]]) in (%[[c0]]) shared_outs(%[[OUT1:[0-9a-z]+]] = %[[ORGOUT1]], %[[OUT2:[0-9a-z]+]] = %[[ORGOUT2]]) +// CHECK: scf.forall (%[[IV0:.+]]) = +// CHECK-SAME: to (%[[c0]]) +// CHECK-SAME: shared_outs(%[[OUT1:[0-9a-z]+]] = %[[ORGOUT1]], %[[OUT2:[0-9a-z]+]] = %[[ORGOUT2]]) // CHECK: %[[TSMIN:.+]] = affine.min #[[$map0]](%[[IV0]]) // CHECK: %[[TS:.+]] = affine.max #[[$map1]](%[[TSMIN]]) // CHECK-NOT: affine.min @@ -301,7 +317,7 @@ // CHECK: %[[tOUT1:.+]] = tensor.extract_slice %[[OUT1]][%[[LB]]] [%[[TS]]] [1] : // CHECK: %[[tOUT2:.+]] = tensor.extract_slice %[[OUT2]][%[[LB]]] [%[[TS]]] [1] : // CHECK: %[[RES1:[0-9]+]]:[[RES2:[0-9]+]] = linalg.generic -// CHECK: scf.foreach_thread.perform_concurrently +// CHECK: scf.forall.in_parallel // CHECK-NEXT: tensor.parallel_insert_slice %[[RES1]]#0 into %[[OUT1]][%[[LB]]] [%[[TS]]] [1] : // CHECK-NEXT: tensor.parallel_insert_slice %[[RES1]]#1 into %[[OUT2]][%[[LB]]] [%[[TS]]] [1] : %res1, %res2 = linalg.generic @@ -325,7 +341,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %foreach_thread, %tiled_generic = transform.structured.tile_to_foreach_thread_op %0 num_threads [7] + %forall, %tiled_generic = transform.structured.tile_to_forall_op %0 num_threads [7] } // ----- @@ -346,14 +362,16 @@ %OUT1: tensor<300x100xf32>, %OUT2: tensor<300xf32>) -> (tensor<300x100xf32>, tensor<300xf32>) { // CHECK-DAG: %[[c0:.+]] = arith.constant 4 : -// CHECK: scf.foreach_thread (%[[IV0:.+]]) in (%[[c0]]) shared_outs(%[[OUT1:[0-9a-z]+]] = %[[ORGOUT1]], %[[OUT2:[0-9a-z]+]] = %[[ORGOUT2]]) +// CHECK: scf.forall (%[[IV0:.+]]) = +// CHECK-SAME: to (%[[c0]]) +// CHECK-SAME: shared_outs(%[[OUT1:[0-9a-z]+]] = %[[ORGOUT1]], %[[OUT2:[0-9a-z]+]] = %[[ORGOUT2]]) // CHECK: %[[LB:.+]] = affine.apply #[[$map0]](%[[IV0]]) // CHECK: %[[tIN1:.+]] = tensor.extract_slice %[[IN2]][0, %[[LB]]] [100, 75] // CHECK: %[[tIN2:.+]] = tensor.extract_slice %[[IN3]][%[[LB]]] [75] // CHECK: %[[tOUT1:.+]] = tensor.extract_slice %[[OUT1]][%[[LB]], 0] [75, 100] // CHECK: %[[tOUT2:.+]] = tensor.extract_slice %[[OUT2]][%[[LB]]] [75] // CHECK: %[[RES1:[0-9]+]]:[[RES2:[0-9]+]] = linalg.generic -// CHECK: scf.foreach_thread.perform_concurrently +// CHECK: scf.forall.in_parallel // CHECK-NEXT: tensor.parallel_insert_slice %[[RES1]]#0 into %[[OUT1]][%[[LB]], 0] [75, 100] // CHECK-NEXT: tensor.parallel_insert_slice %[[RES1]]#1 into %[[OUT2]][%[[LB]]] [75] %res2, %res3 = linalg.generic { @@ -379,6 +397,6 @@ transform.sequence failures(propagate) { ^bb1(%IN_MAT2: !pdl.operation): %0 = transform.structured.match ops{["linalg.generic"]} in %IN_MAT2 : (!pdl.operation) -> !pdl.operation - %foreach_thread, %tiled_generic = transform.structured.tile_to_foreach_thread_op %0 num_threads [4] + %forall, %tiled_generic = transform.structured.tile_to_forall_op %0 num_threads [4] } diff --git a/mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir b/mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir --- a/mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir +++ b/mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir @@ -12,12 +12,13 @@ func.func @fuse_tileable_op(%arg0: index, %arg1: tensor, %arg2: tensor) -> tensor { %cst = arith.constant 4.200000e+01 : f32 %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %0 = linalg.fill ins(%cst : f32) outs(%arg1 : tensor) -> tensor %d0 = tensor.dim %arg1, %c0 : tensor %1 = affine.apply #map0()[%d0, %arg0] - // CHECK: scf.foreach_thread {{.*}} { - %2 = scf.foreach_thread (%arg3) in (%1) shared_outs(%o = %arg2) -> (tensor) { + // CHECK: scf.forall {{.*}} { + %2 = scf.forall (%arg3) = (%c0) to (%1) step (%c1) shared_outs(%o = %arg2) -> (tensor) { %3 = affine.apply #map1(%arg3)[%arg0] %4 = affine.min #map2(%arg3)[%d0, %arg0] %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor to tensor @@ -28,7 +29,7 @@ // CHECK: %[[T2:.*]] = linalg.elemwise_unary ins(%[[T1]] %7 = linalg.elemwise_unary ins(%6 : tensor) outs(%5 : tensor) -> tensor - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor into tensor } } @@ -44,7 +45,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1 = transform.structured.match ops{["scf.foreach_thread"]} in %arg1 : (!pdl.operation) -> !pdl.operation + %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!pdl.operation) -> !pdl.operation // linalg.fill is tileable. The op is tiled and fused. transform.structured.fuse_into_containing_op %0 into %1 @@ -63,11 +64,13 @@ // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<64xf32> // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<64xf32> func.func @fuse_untileable_op(%arg0: index, %arg1: tensor<64xf32>, %arg2: tensor<64xf32>) -> tensor<64xf32> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %0 = tensor.empty(%arg0) : tensor %1 = affine.apply #map0()[%arg0] - // CHECK: scf.foreach_thread {{.*}} { - %2 = scf.foreach_thread (%arg3) in (%1) shared_outs(%o = %arg2) -> (tensor<64xf32>) { + // CHECK: scf.forall {{.*}} { + %2 = scf.forall (%arg3) = (%c0) to (%1) step (%c1) shared_outs(%o = %arg2) -> (tensor<64xf32>) { // CHECK: %[[INIT_TENSOR:.*]] = tensor.empty %3 = affine.apply #map1(%arg3)[%arg0] %4 = affine.min #map2(%arg3)[%arg0] @@ -75,7 +78,7 @@ // CHECK: %[[T2:.*]] = linalg.elemwise_unary ins(%[[INIT_TENSOR]] %7 = linalg.elemwise_unary ins(%0 : tensor) outs(%5 : tensor) -> tensor - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor into tensor<64xf32> } } @@ -87,7 +90,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["tensor.empty"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1 = transform.structured.match ops{["scf.foreach_thread"]} in %arg1 : (!pdl.operation) -> !pdl.operation + %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!pdl.operation) -> !pdl.operation // tensor.empty is not tileable. The op is cloned and fused. transform.structured.fuse_into_containing_op %0 into %1 @@ -108,20 +111,21 @@ func.func @fuse_tileable_op_rank_reducing(%arg0: index, %arg1: tensor, %arg2: tensor) -> tensor { %cst = arith.constant 4.200000e+01 : f32 %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %0 = linalg.fill ins(%cst : f32) outs(%arg2 : tensor) -> tensor %d0 = tensor.dim %arg1, %c0 : tensor - // CHECK: scf.foreach_thread {{.*}} -> (tensor) { - %2 = scf.foreach_thread (%arg3) in (%d0) shared_outs(%o = %0) -> (tensor) { + // CHECK: scf.forall {{.*}} -> (tensor) { + %2 = scf.forall (%arg3) = (%c0) to (%d0) step (%c1) shared_outs(%o = %0) -> (tensor) { %5 = tensor.extract_slice %o[%arg3] [1] [1] : tensor to tensor - + // CHECK: tensor.extract_slice %{{.*}}[%{{.*}}] [1] [1] : tensor to tensor<1xf32> // CHECK: linalg.fill ins(%{{.*}} : f32) outs(%{{.*}} : tensor<1xf32>) -> tensor<1xf32> // CHECK: tensor.extract_slice %{{.*}}[0] [1] [1] : tensor<1xf32> to tensor // CHECK: func.call @foo(%{{.*}}) : (tensor) -> tensor %7 = func.call @foo(%5) : (tensor) -> tensor - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { // CHECK: tensor.parallel_insert_slice %{{.*}} into %{{.*}}[%{{.*}}] [1] [1] : tensor into tensor tensor.parallel_insert_slice %7 into %o[%arg3] [1] [1] : tensor into tensor } @@ -133,7 +137,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1 = transform.structured.match ops{["scf.foreach_thread"]} in %arg1 : (!pdl.operation) -> !pdl.operation + %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!pdl.operation) -> !pdl.operation // linalg.fill is tileable. The op is tiled and fused. transform.structured.fuse_into_containing_op %0 into %1 @@ -154,12 +158,14 @@ func.func @fuse_tileable_op_through_bbarg(%arg0: index, %arg1: tensor, %arg2: tensor) -> tensor { %cst = arith.constant 4.200000e+01 : f32 %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %0 = linalg.fill ins(%cst : f32) outs(%arg2 : tensor) -> tensor %d0 = tensor.dim %arg1, %c0 : tensor %1 = affine.apply #map0()[%d0, %arg0] - // CHECK: scf.foreach_thread {{.*}} shared_outs(%[[BBARGOUT:.*]] = %[[OUT]]) -> (tensor) { - %2 = scf.foreach_thread (%arg3) in (%1) shared_outs(%o = %0) -> (tensor) { + // CHECK: scf.forall + // CHECK-SAME: shared_outs(%[[BBARGOUT:.*]] = %[[OUT]]) -> (tensor) { + %2 = scf.forall (%arg3) = (%c0) to (%1) step (%c1) shared_outs(%o = %0) -> (tensor) { %3 = affine.apply #map1(%arg3)[%arg0] %4 = affine.min #map2(%arg3)[%d0, %arg0] %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor to tensor @@ -170,7 +176,7 @@ // CHECK: %[[T2:.*]] = linalg.elemwise_unary {{.*}} outs(%[[T1]] %7 = linalg.elemwise_unary ins(%6 : tensor) outs(%5 : tensor) -> tensor - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor into tensor } } @@ -181,7 +187,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1 = transform.structured.match ops{["scf.foreach_thread"]} in %arg1 : (!pdl.operation) -> !pdl.operation + %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!pdl.operation) -> !pdl.operation // linalg.fill is tileable. The op is tiled and fused. transform.structured.fuse_into_containing_op %0 into %1 @@ -204,6 +210,7 @@ func.func @fuse_tileable_multi_output_op(%idx: index, %in: tensor, %out_1: tensor, %out_2: tensor, %out_3: tensor) -> tensor { %cst = arith.constant 4.200000e+01 : f32 %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %0:2 = linalg.generic { indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], @@ -218,8 +225,8 @@ %1 = affine.apply #map0()[%d0, %idx] - // CHECK: scf.foreach_thread {{.*}} { - %2 = scf.foreach_thread (%i) in (%1) shared_outs(%o = %out_2) -> (tensor) { + // CHECK: scf.forall {{.*}} { + %2 = scf.forall (%i) = (%c0) to (%1) step (%c1) shared_outs(%o = %out_2) -> (tensor) { %3 = affine.apply #map1(%i)[%idx] %4 = affine.min #map2(%i)[%d0, %idx] %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor to tensor @@ -230,7 +237,7 @@ // CHECK: %[[T2:.*]] = linalg.elemwise_unary ins(%[[T1]]#0 %7 = linalg.elemwise_unary ins(%6 : tensor) outs(%5 : tensor) -> tensor - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor into tensor } } @@ -241,7 +248,7 @@ transform.sequence failures(propagate) { ^bb1(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %1 = transform.structured.match ops{["scf.foreach_thread"]} in %arg1 : (!pdl.operation) -> !pdl.operation + %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!pdl.operation) -> !pdl.operation // linalg.generic is tileable. The op is tiled and fused. transform.structured.fuse_into_containing_op %0 into %1 diff --git a/mlir/test/Dialect/Linalg/transform-tile-and-fuse.mlir b/mlir/test/Dialect/Linalg/transform-tile-and-fuse.mlir --- a/mlir/test/Dialect/Linalg/transform-tile-and-fuse.mlir +++ b/mlir/test/Dialect/Linalg/transform-tile-and-fuse.mlir @@ -4,7 +4,7 @@ module { // CHECK: func @foo - // CHECK: scf.foreach_thread {{.*}} { + // CHECK: scf.forall {{.*}} { // CHECK: linalg.fill // CHECK: linalg.matmul // CHECK: linalg.generic @@ -47,7 +47,7 @@ %producers = transform.structured.match attributes{"__producer__"} in %arg1 : (!pdl.operation) -> !pdl.operation // Tile the root. - %foreach_thread_op, %tiled_op = transform.structured.tile_to_foreach_thread_op %root num_threads [10, 20] + %foreach_thread_op, %tiled_op = transform.structured.tile_to_forall_op %root num_threads [10, 20] // Fuse all producers. transform.structured.fuse_into_containing_op %producers into %foreach_thread_op @@ -56,12 +56,12 @@ // ----- -// Inverse the order of the payload ops passed to the tile_to_foreach_thread_op +// Inverse the order of the payload ops passed to the tile_to_forall_op // op. Fusion should still work. module { // CHECK: func @foo - // CHECK: scf.foreach_thread {{.*}} { + // CHECK: scf.forall {{.*}} { // CHECK: linalg.fill // CHECK: linalg.matmul // CHECK: linalg.generic @@ -105,7 +105,7 @@ %reversed_producers = transform.test_reverse_payload_ops %producers // Tile the root. - %foreach_thread_op, %tiled_op = transform.structured.tile_to_foreach_thread_op %root num_threads [10, 20] + %foreach_thread_op, %tiled_op = transform.structured.tile_to_forall_op %root num_threads [10, 20] // Fuse all producers. transform.structured.fuse_into_containing_op %reversed_producers into %foreach_thread_op diff --git a/mlir/test/Dialect/Linalg/transform-tile-reduction.mlir b/mlir/test/Dialect/Linalg/transform-tile-reduction.mlir --- a/mlir/test/Dialect/Linalg/transform-tile-reduction.mlir +++ b/mlir/test/Dialect/Linalg/transform-tile-reduction.mlir @@ -109,7 +109,7 @@ transform.sequence failures(propagate) { ^bb0(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %loop, %1, %2, %3 = transform.structured.tile_reduction_using_foreach_thread %0 + %loop, %1, %2, %3 = transform.structured.tile_reduction_using_forall %0 by num_threads = [0, 5], tile_sizes = [] } @@ -128,7 +128,9 @@ // CHECK-DAG: %[[D2:.*]] = tensor.dim %[[ARG1]], %[[C0]] : tensor // CHECK: %[[E:.*]] = tensor.empty(%[[D2]]) : tensor // CHECK: %[[F:.*]] = linalg.fill ins(%[[I]] : f32) outs(%[[E]] : tensor) -> tensor -// CHECK: %[[L:.*]] = scf.foreach_thread (%[[IV:.+]]) in (%[[C5]]) shared_outs(%[[ARG3:.+]] = %[[F]]) -> (tensor) { +// CHECK: %[[L:.*]] = scf.forall (%[[IV:.+]]) = +//CHECK-SAME: to (%[[C5]]) +//CHECK-SAME: shared_outs(%[[ARG3:.+]] = %[[F]]) -> (tensor) { // CHECK-DAG: %[[TS0:.+]] = affine.min #[[MAP0]](%[[IV]])[%[[D1]]] // CHECK-DAG: %[[TS1:.+]] = affine.max #[[MAP1]](%[[TS0]]) // CHECK-DAG: %[[ET:.+]] = tensor.extract_slice %[[ARG3:.+]][0, %[[IV]]] [%[[D0]], 1] [1, 1] : tensor to tensor @@ -140,7 +142,7 @@ // CHECK: arith.addf // CHECK: linalg.yield // CHECK: } -> tensor -// CHECK: scf.foreach_thread.perform_concurrently { +// CHECK: scf.forall.in_parallel { // CHECK: tensor.parallel_insert_slice %[[PARTIAL]] into %[[ARG3]][0, %[[IV]]] [%[[D0]], 1] [1, 1] : tensor into tensor // CHECK: } // CHECK: } @@ -162,7 +164,7 @@ transform.sequence failures(propagate) { ^bb0(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %loop, %1, %2, %3 = transform.structured.tile_reduction_using_foreach_thread %0 + %loop, %1, %2, %3 = transform.structured.tile_reduction_using_forall %0 by num_threads = [0, 0, 5], tile_sizes = [] } @@ -183,7 +185,9 @@ // CHECK-DAG: %[[D4:.*]] = tensor.dim %[[ARG2]], %[[C1]] : tensor // CHECK: %[[E:.*]] = tensor.empty(%[[D3]], %[[D4]]) : tensor // CHECK: %[[F:.*]] = linalg.fill ins(%[[I]] : f32) outs(%[[E]] : tensor) -> tensor -// CHECK: %[[L:.*]] = scf.foreach_thread (%[[IV:.+]]) in (%[[C5]]) shared_outs(%[[ARG3:.+]] = %[[F]]) -> (tensor) { +// CHECK: %[[L:.*]] = scf.forall (%[[IV:.+]]) = +//CHECK-SAME: to (%[[C5]]) +//CHECK-SAME: shared_outs(%[[ARG3:.+]] = %[[F]]) -> (tensor) { // CHECK-DAG: %[[TS0:.+]] = affine.min #[[MAP0]](%[[IV]])[%[[D1]]] // CHECK-DAG: %[[TS1:.+]] = affine.max #[[MAP1]](%[[TS0]]) // CHECK-DAG: %[[ET:.+]] = tensor.extract_slice %[[ARG3:.+]][0, 0, %[[IV]]] [%[[D0]], %[[D2]], 1] [1, 1, 1] : tensor to tensor @@ -192,7 +196,7 @@ // CHECK: %[[INCHUNKB:.+]] = tensor.extract_slice %[[ARG1]][%[[TINDEX]], 0] [%[[TS1]], %[[D2]]] [1, 1] : tensor to tensor // CHECK: %[[TEMPEXT:.+]] = tensor.extract_slice %[[ET]][0, 0] [%[[D0]], %[[D2]]] [1, 1] : tensor to tensor // CHECK: %[[PARTIAL:.+]] = linalg.matmul ins(%[[INCHUNKA]], %[[INCHUNKB]] : tensor, tensor) outs(%[[TEMPEXT]] : tensor) -> tensor -// CHECK: scf.foreach_thread.perform_concurrently { +// CHECK: scf.forall.in_parallel { // CHECK: tensor.parallel_insert_slice %[[PARTIAL]] into %[[ARG3]][0, 0, %[[IV]]] [%[[D0]], %[[D2]], 1] [1, 1, 1] : tensor into tensor // CHECK: } // CHECK: } @@ -222,7 +226,7 @@ transform.sequence failures(propagate) { ^bb0(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %loop, %1, %2, %3 = transform.structured.tile_reduction_using_foreach_thread %0 + %loop, %1, %2, %3 = transform.structured.tile_reduction_using_forall %0 by num_threads = [0, 5], tile_sizes = [0, 3], mapping = [#gpu.thread] } @@ -241,7 +245,9 @@ // CHECK-DAG: %[[D2:.*]] = tensor.dim %[[ARG1]], %[[C0]] : tensor // CHECK: %[[E:.*]] = tensor.empty(%[[D2]]) : tensor // CHECK: %[[F:.*]] = linalg.fill ins(%[[I]] : f32) outs(%[[E]] : tensor) -> tensor -// CHECK: %[[L:.*]] = scf.foreach_thread (%[[IV:.+]]) in (%[[C5]]) shared_outs(%[[ARG3:.+]] = %[[F]]) -> (tensor) { +// CHECK: %[[L:.*]] = scf.forall (%[[IV:.+]]) = +//CHECK-SAME: to (%[[C5]]) +//CHECK-SAME: shared_outs(%[[ARG3:.+]] = %[[F]]) -> (tensor) { // CHECK: %[[ET:.+]] = tensor.extract_slice %[[ARG3:.+]][0, %[[IV]]] [%[[D0]], 1] [1, 1] : tensor to tensor // CHECK: %[[D1:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor // CHECK: %[[LB:.+]] = affine.apply #[[MAP0]]()[%[[IV]]] @@ -258,7 +264,7 @@ // CHECK: %[[INS:.+]] = tensor.insert_slice %[[PARTIAL]] into %[[ACC]][0] [%[[D3]]] [1] : tensor into tensor // CHECK: scf.yield %[[INS]] : tensor // CHECK: } -// CHECK: scf.foreach_thread.perform_concurrently { +// CHECK: scf.forall.in_parallel { // CHECK: tensor.parallel_insert_slice %[[CARRY]] into %[[ARG3]][0, %[[IV]]] [%[[D0]], 1] [1, 1] : tensor into tensor // CHECK: } // CHECK: } @@ -288,7 +294,7 @@ transform.sequence failures(propagate) { ^bb0(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!pdl.operation) -> !pdl.operation - %loop, %1, %2, %3 = transform.structured.tile_reduction_using_foreach_thread %0 + %loop, %1, %2, %3 = transform.structured.tile_reduction_using_forall %0 by num_threads = [0, 5], tile_sizes = [0, 3], mapping = [#gpu.thread] // CHECK: expecting fill @@ -326,7 +332,7 @@ ^bb0(%arg1: !pdl.operation): %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!pdl.operation) -> !pdl.operation // expected-error @below {{could not tile reduction}} - %loop, %1, %2, %3 = transform.structured.tile_reduction_using_foreach_thread %0 + %loop, %1, %2, %3 = transform.structured.tile_reduction_using_forall %0 by num_threads = [5], tile_sizes = [3], mapping = [#gpu.thread] } diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir --- a/mlir/test/Dialect/SCF/canonicalize.mlir +++ b/mlir/test/Dialect/SCF/canonicalize.mlir @@ -1484,10 +1484,11 @@ %arg0 : tensor<1x5xf32>, %arg1: tensor, %num_threads : index) -> index { // CHECK: %[[c1:.*]] = arith.constant 1 : index + %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index - %2 = scf.foreach_thread (%tidx) in (%num_threads) shared_outs(%o = %arg1) -> (tensor) { - scf.foreach_thread.perform_concurrently { + %2 = scf.forall (%tidx) = (%c0) to (%num_threads) step (%c1) shared_outs(%o = %arg1) -> (tensor) { + scf.forall.in_parallel { tensor.parallel_insert_slice %arg0 into %o[%tidx, 0] [1, 5] [1, 1] : tensor<1x5xf32> into tensor } } diff --git a/mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir b/mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir --- a/mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir +++ b/mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir @@ -2,13 +2,15 @@ func.func @reduce() { // CHECK: %[[C64:.*]] = arith.constant 64 : index + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index %cst_0 = arith.constant -0.000000e+00 : f32 %0 = memref.alloc() : memref<128x384xf32> linalg.fill ins(%cst_0 : f32) outs(%0 : memref<128x384xf32>) %2 = memref.alloc() : memref<128xf32> linalg.fill ins(%cst_0 : f32) outs(%2 : memref<128xf32>) - scf.foreach_thread (%arg0) in (%c2) { + scf.forall (%arg0) = (%c0) to (%c2) step (%c1) { %7 = affine.min affine_map<(d0) -> (d0 * -64 + 128, 64)>(%arg0) %8 = affine.max affine_map<(d0) -> (0, d0)>(%7) %9 = affine.apply affine_map<(d0) -> (d0 * 64)>(%arg0) diff --git a/mlir/test/Dialect/SCF/invalid.mlir b/mlir/test/Dialect/SCF/invalid.mlir --- a/mlir/test/Dialect/SCF/invalid.mlir +++ b/mlir/test/Dialect/SCF/invalid.mlir @@ -544,13 +544,14 @@ // ----- func.func @wrong_num_results(%in: tensor<100xf32>, %out: tensor<100xf32>) { + %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %num_threads = arith.constant 100 : index // expected-error @+1 {{1 operands present, but expected 2}} - %result:2 = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> (tensor<100xf32>, tensor<100xf32>) { + %result:2 = scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) shared_outs(%o = %out) -> (tensor<100xf32>, tensor<100xf32>) { %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : tensor<1xf32> into tensor<100xf32> } @@ -561,12 +562,13 @@ // ----- func.func @invalid_insert_dest(%in: tensor<100xf32>, %out: tensor<100xf32>) { + %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %num_threads = arith.constant 100 : index - %result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> (tensor<100xf32>) { + %result = scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) shared_outs(%o = %out) -> (tensor<100xf32>) { %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { // expected-error @+1 {{may only insert into an output block argument}} tensor.parallel_insert_slice %1 into %out[%thread_idx][1][1] : tensor<1xf32> into tensor<100xf32> @@ -578,13 +580,14 @@ // ----- func.func @wrong_terminator_op(%in: tensor<100xf32>, %out: tensor<100xf32>) { + %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %num_threads = arith.constant 100 : index - %result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> (tensor<100xf32>) { + %result = scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) shared_outs(%o = %out) -> (tensor<100xf32>) { %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> // expected-error @+1 {{expected only tensor.parallel_insert_slice ops}} - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : tensor<1xf32> into tensor<100xf32> %0 = arith.constant 1: index @@ -596,10 +599,11 @@ // ----- func.func @mismatched_mapping(%x: memref<2 x 32 x f32>, %y: memref<2 x 32 x f32>, %t: memref<32 x f32>, %alpha : f32, %stream : !gpu.async.token) -> memref<2 x 32 x f32> { + %zero = arith.constant 0 : index %one = arith.constant 1 : index %c65535 = arith.constant 65535 : index - // expected-error @below {{'scf.foreach_thread' op mapping attribute size must match op rank}} - scf.foreach_thread (%i, %j) in (%c65535, %c65535) { + // expected-error @below {{'scf.forall' op mapping attribute size must match op rank}} + scf.forall (%i, %j) = (%zero, %zero) to (%c65535, %c65535) step (%one, %one) { %4 = memref.load %x[%i, %j] : memref<2 x 32 x f32> %5 = memref.load %y[%i, %j] : memref<2 x 32 x f32> %6 = math.fma %alpha, %4, %5 : f32 diff --git a/mlir/test/Dialect/SCF/one-shot-bufferize-analysis.mlir b/mlir/test/Dialect/SCF/one-shot-bufferize-analysis.mlir --- a/mlir/test/Dialect/SCF/one-shot-bufferize-analysis.mlir +++ b/mlir/test/Dialect/SCF/one-shot-bufferize-analysis.mlir @@ -612,9 +612,10 @@ -> tensor<320xf32> { %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %cst = arith.constant -0.000000e+00 : f32 %c320 = arith.constant 320 : index - %4 = scf.foreach_thread (%arg0) in (%c320) shared_outs(%arg1 = %2) -> (tensor<320xf32>) { + %4 = scf.forall (%arg0) = (%c0) to (%c320) step (%c1) shared_outs(%arg1 = %2) -> (tensor<320xf32>) { // CHECK: tensor.extract_slice {{.*}} {__inplace_operands_attr__ = ["true", "none"]} %5 = tensor.extract_slice %3[%arg0, 0] [1, 10240] [1, 1] : tensor<320x10240xf32> to tensor<1x10240xf32> // CHECK: tensor.extract_slice {{.*}} {__inplace_operands_attr__ = ["true", "none"]} @@ -624,7 +625,7 @@ // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]} %8 = linalg.fill ins(%cst : f32) outs(%7 : tensor<1xf32>) -> tensor<1xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { // CHECK: tensor.parallel_insert_slice {{.*}} {__inplace_operands_attr__ = ["true", "true", "none"]} tensor.parallel_insert_slice %8 into %arg1[%arg0] [1] [1] : tensor<1xf32> into tensor<320xf32> } diff --git a/mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir b/mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir --- a/mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir +++ b/mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir @@ -108,23 +108,24 @@ // ----- -// CHECK-LABEL: func @scf_foreach_thread_out_of_place( +// CHECK-LABEL: func @scf.forall_out_of_place( // CHECK-SAME: %[[arg0:.*]]: tensor<100xf32>, %[[arg1:.*]]: tensor<100xf32> -// CHECK-FUNC-LABEL: func @scf_foreach_thread_out_of_place( -func.func @scf_foreach_thread_out_of_place(%in: tensor<100xf32>, +// CHECK-FUNC-LABEL: func @scf.forall_out_of_place( +func.func @scf.forall_out_of_place(%in: tensor<100xf32>, %out: tensor<100xf32>) { + %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %num_threads = arith.constant 100 : index // CHECK-FUNC-NOT: alloc_tensor // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[arg1]]) {bufferization.escape = [false]} : tensor<100xf32> - // CHECK: scf.foreach_thread {{.*}} shared_outs(%[[o:.*]] = %[[alloc]]) - %result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> tensor<100xf32> { + // CHECK: scf.forall {{.*}} shared_outs(%[[o:.*]] = %[[alloc]]) + %result = scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) shared_outs(%o = %out) -> tensor<100xf32> { // CHECK: tensor.extract_slice - // CHECK: scf.foreach_thread.perform_concurrently + // CHECK: scf.forall.in_parallel // CHECK: tensor.parallel_insert_slice %{{.*}} into %[[o]] %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : tensor<1xf32> into tensor<100xf32> } diff --git a/mlir/test/Dialect/SCF/one-shot-bufferize.mlir b/mlir/test/Dialect/SCF/one-shot-bufferize.mlir --- a/mlir/test/Dialect/SCF/one-shot-bufferize.mlir +++ b/mlir/test/Dialect/SCF/one-shot-bufferize.mlir @@ -543,8 +543,8 @@ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index - // CHECK: scf.foreach_thread (%[[tidx:.*]]) in (%[[idx2]]) - %2 = scf.foreach_thread (%arg3) in (%idx2) shared_outs(%o = %arg2) -> (tensor) { + // CHECK: scf.forall (%[[tidx:.*]]) to (%[[idx2]]) + %2 = scf.forall (%arg3) = (%c0) to (%idx2) step (%c1) shared_outs(%o = %arg2) -> (tensor) { // CHECK: %[[subview:.*]] = memref.subview %[[arg2]][5] [%[[idx]]] [1] %6 = tensor.extract_slice %o[5] [%idx] [%c1] : tensor to tensor // CHECK: linalg.fill ins(%{{.*}}) outs(%[[subview]] : memref into tensor } @@ -589,8 +589,8 @@ // CHECK: %[[alloc1:.*]] = memref.alloc // CHECK: memref.copy %[[arg2]], %[[alloc1]] - // CHECK: scf.foreach_thread (%[[tidx:.*]]) in (%[[idx2]]) - %2 = scf.foreach_thread (%arg3) in (%idx2) shared_outs(%o = %arg2) -> (tensor) { + // CHECK: scf.forall (%[[tidx:.*]]) to (%[[idx2]]) + %2 = scf.forall (%arg3) = (%c0) to (%idx2) step (%c1) shared_outs(%o = %arg2) -> (tensor) { // CHECK: %[[subview1:.*]] = memref.subview %[[alloc1]][5] [%[[idx]]] [1] %6 = tensor.extract_slice %o[5] [%idx] [%c1] : tensor to tensor @@ -601,9 +601,9 @@ // CHECK: memref.copy %[[subview1]], %[[subview1]] // Empty terminator is elided from pretty-printing. - // CHECK-NOT: scf.foreach_thread.perform_concurrently + // CHECK-NOT: scf.forall.in_parallel // CHECK-NOT: parallel_insert_slice - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %8 into %o[5] [%idx] [%c1] : tensor into tensor } @@ -625,12 +625,15 @@ #map1 = affine_map<(d0) -> (d0 * 2)> // CHECK-LABEL: func.func @matmul -func.func @matmul(%arg0: tensor<8x8xf32>, %arg1: tensor<8x8xf32>, %arg2: tensor<8x8xf32> {bufferization.writable = true}) -> tensor<8x8xf32> { +func.func @matmul(%arg0: tensor<8x8xf32>, %arg1: tensor<8x8xf32>, + %arg2: tensor<8x8xf32> {bufferization.writable = true}) -> tensor<8x8xf32> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index %c4 = arith.constant 4 : index - // CHECK: scf.foreach_thread {{.*}} - %0 = scf.foreach_thread (%arg3, %arg4) in (%c2, %c4) shared_outs(%o = %arg2) -> (tensor<8x8xf32>) { + // CHECK: scf.forall + %0 = scf.forall (%arg3, %arg4) = (%c0, %c0) to (%c2, %c4) step (%c1, %c1) shared_outs(%o = %arg2) -> (tensor<8x8xf32>) { %1 = affine.apply #map0(%arg3) %3 = tensor.extract_slice %arg0[%1, 0] [4, 8] [1, 1] : tensor<8x8xf32> to tensor<4x8xf32> %4 = affine.apply #map1(%arg4) @@ -639,7 +642,7 @@ // CHECK: linalg.matmul ins({{.*}}memref<4x8xf32, strided<[?, ?], offset: ?>>, memref<8x4xf32, strided<[?, ?], offset: ?>>) outs({{.*}} : memref<4x4xf32, strided<[?, ?], offset: ?>>) %8 = linalg.matmul ins(%3, %6 : tensor<4x8xf32>, tensor<8x4xf32>) outs(%7 : tensor<4x4xf32>) -> tensor<4x4xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %8 into %o[%1, %4] [4, 4] [1, 1] : tensor<4x4xf32> into tensor<8x8xf32> } } @@ -651,6 +654,8 @@ // CHECK-LABEL: func @scf_foreach_private_var( // CHECK-SAME: %[[t:.*]]: memref<10xf32 func.func @scf_foreach_private_var(%t: tensor<10xf32>) -> f32 { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index %c5 = arith.constant 5 : index @@ -658,19 +663,19 @@ // CHECK: %[[t_copy:.*]] = memref.alloc() {{.*}} : memref<10xf32> // CHECK: memref.copy %[[t]], %[[t_copy]] - // CHECK: scf.foreach_thread (%{{.*}}) in (%{{.*}}) { + // CHECK: scf.forall // Load from the copy and store into the shared output. // CHECK: %[[subview:.*]] = memref.subview %[[t]] // CHECK: memref.load %[[t_copy]] // CHECK: memref.store %{{.*}}, %[[subview]] - %0 = scf.foreach_thread (%tid) in (%c2) shared_outs(%o = %t) -> tensor<10xf32> { + %0 = scf.forall (%tid) = (%c0) to (%c2) step (%c1) shared_outs(%o = %t) -> tensor<10xf32> { %offset = arith.muli %c5, %tid : index %slice = tensor.extract_slice %o[%offset] [5] [1] : tensor<10xf32> to tensor<5xf32> %r2 = tensor.extract %t[%tid] : tensor<10xf32> %i = tensor.insert %r2 into %slice[%c2] : tensor<5xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %i into %o[%offset] [5] [1] : tensor<5xf32> into tensor<10xf32> } @@ -686,13 +691,15 @@ // CHECK-SAME: %[[t0:.*]]: memref<10xf32, {{.*}}>, %[[t1:.*]]: memref<10xf32 func.func @scf_foreach_privatized_but_not_copied( %t0: tensor<10xf32>, %t1: tensor<10xf32>) -> f32 { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index %c5 = arith.constant 5 : index // CHECK-NOT: memref.alloc // CHECK-NOT: memref.copy - // CHECK: scf.foreach_thread {{.*}} { - %0 = scf.foreach_thread (%tid) in (%c2) shared_outs(%o = %t0) -> tensor<10xf32> { + // CHECK: scf.forall + %0 = scf.forall (%tid) = (%c0) to (%c2) step (%c1) shared_outs(%o = %t0) -> tensor<10xf32> { %offset = arith.muli %c5, %tid : index %slice = tensor.extract_slice %o[%offset] [5] [1] : tensor<10xf32> to tensor<5xf32> @@ -701,7 +708,7 @@ // CHECK: memref.load %[[t1]] %r2 = tensor.extract %t1[%tid] : tensor<10xf32> %i = tensor.insert %r2 into %slice[%c2] : tensor<5xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %i into %o[%offset] [5] [1] : tensor<5xf32> into tensor<10xf32> } diff --git a/mlir/test/Dialect/SCF/ops.mlir b/mlir/test/Dialect/SCF/ops.mlir --- a/mlir/test/Dialect/SCF/ops.mlir +++ b/mlir/test/Dialect/SCF/ops.mlir @@ -313,19 +313,20 @@ // CHECK-LABEL: func.func @simple_example func.func @simple_example(%in: tensor<100xf32>, %out: tensor<100xf32>) { + %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %num_threads = arith.constant 100 : index - // CHECK: scf.foreach_thread + // CHECK: scf.forall // CHECK-NEXT: tensor.extract_slice - // CHECK-NEXT: scf.foreach_thread.perform_concurrently + // CHECK-NEXT: scf.forall.in_parallel // CHECK-NEXT: tensor.parallel_insert_slice // CHECK-NEXT: } // CHECK-NEXT: } // CHECK-NEXT: return - %result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> tensor<100xf32> { + %result = scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) shared_outs(%o = %out) -> tensor<100xf32> { %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : tensor<1xf32> into tensor<100xf32> } @@ -335,13 +336,15 @@ // CHECK-LABEL: func.func @elide_terminator func.func @elide_terminator() -> () { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %num_threads = arith.constant 100 : index - // CHECK: scf.foreach_thread + // CHECK: scf.forall // CHECK-NEXT: } {mapping = [#gpu.thread]} // CHECK-NEXT: return - scf.foreach_thread (%thread_idx) in (%num_threads) { - scf.foreach_thread.perform_concurrently { + scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) { + scf.forall.in_parallel { } } {mapping = [#gpu.thread]} return diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir --- a/mlir/test/Dialect/Tensor/canonicalize.mlir +++ b/mlir/test/Dialect/Tensor/canonicalize.mlir @@ -1531,12 +1531,14 @@ %c1 = arith.constant 1 : index // CHECK-NOT: tensor.cast - // CHECK: scf.foreach_thread (%[[tidx:[0-9a-z]*]]) in (%[[num_threads]]) shared_outs(%[[o:.*]] = %[[arg1]]) -> (tensor) { - // CHECK-NEXT: scf.foreach_thread.perform_concurrently { + // CHECK: scf.forall (%[[tidx:[0-9a-z]*]]) = (%{{[0-9a-z]*}}) + // CHECK-SAME: to (%[[num_threads]]) step (%{{[0-9a-z]*}}) + // CHECK-SAME: shared_outs(%[[o:.*]] = %[[arg1]]) -> (tensor) { + // CHECK-NEXT: scf.forall.in_parallel { // CHECK-NEXT: tensor.parallel_insert_slice %[[arg0]] into %[[o]][%[[tidx]], 0] [1, 5] [1, 1] - %2 = scf.foreach_thread (%tidx) in (%num_threads) shared_outs(%o = %arg1) -> (tensor) { + %2 = scf.forall (%tidx) = (%c0) to (%num_threads) step (%c1) shared_outs(%o = %arg1) -> (tensor) { %3 = tensor.cast %arg0 : tensor<1x5xf32> to tensor - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %3 into %o[%tidx, %c0] [%c1, 5] [%c1, %c1] : tensor into tensor } } @@ -1553,11 +1555,11 @@ { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index - // CHECK: scf.foreach_thread () in () shared_outs(%[[o:.*]] = %[[arg1]]) -> (tensor<1x5xf32>) { - // CHECK-NEXT: scf.foreach_thread.perform_concurrently { + // CHECK: scf.forall () = () to () step () shared_outs(%[[o:.*]] = %[[arg1]]) -> (tensor<1x5xf32>) { + // CHECK-NEXT: scf.forall.in_parallel { // CHECK-NEXT: tensor.parallel_insert_slice %[[arg0]] into %[[o]][0, 0] [1, 5] [1, 1] : tensor<1x5xf32> into tensor<1x5xf32> - %2 = scf.foreach_thread () in () shared_outs(%o = %arg1) -> (tensor<1x5xf32>) { - scf.foreach_thread.perform_concurrently { + %2 = scf.forall () = () to () step () shared_outs(%o = %arg1) -> (tensor<1x5xf32>) { + scf.forall.in_parallel { tensor.parallel_insert_slice %arg0 into %o[%c0, %c0] [1, 5] [%c1, %c1] : tensor<1x5xf32> into tensor<1x5xf32> } } diff --git a/mlir/test/Dialect/Tensor/extract-slice-from-collapse-shape.mlir b/mlir/test/Dialect/Tensor/extract-slice-from-collapse-shape.mlir --- a/mlir/test/Dialect/Tensor/extract-slice-from-collapse-shape.mlir +++ b/mlir/test/Dialect/Tensor/extract-slice-from-collapse-shape.mlir @@ -29,11 +29,13 @@ // FOREACH-DAG: %[[c5:.+]] = arith.constant 5 : index // FOREACH-DAG: %[[c7:.+]] = arith.constant 7 : index // FOREACH-DAG: %[[init:.+]] = tensor.empty() : tensor<20x11xf32> -// FOREACH: %[[tile:.+]] = scf.foreach_thread (%[[iv:.+]]) in (%[[c20]]) shared_outs(%[[dest:.+]] = %[[init]]) +// FOREACH: %[[tile:.+]] = scf.forall (%[[iv:.+]]) = +//FOREACH-SAME: to (%[[c20]]) +//FOREACH-SAME: shared_outs(%[[dest:.+]] = %[[init]]) // FOREACH: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[iv]] into (%[[c3]], %[[c5]], %[[c7]] // FOREACH: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 11] [1, 1, 1, 1] : // FOREACH: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} : -// FOREACH: perform_concurrently +// FOREACH: in_parallel // FOREACH-NEXT: tensor.parallel_insert_slice %[[sliceFlat]] into %[[dest]][%[[iv]], 0] [1, 11] [1, 1] : // FOREACH: return %[[tile]] @@ -137,14 +139,16 @@ // FOREACH-DAG: %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] : // FOREACH-DAG: %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] : // FOREACH-DAG: %[[d4:.+]] = tensor.dim %[[arg0]], %[[c4]] : -// FOREACH: %[[tile1:.+]] = scf.foreach_thread (%[[tid1:.+]], %[[tid2:.+]]) in (%[[sz1]], %[[sz2]]) shared_outs(%[[dest:.+]] = %[[init]]) +// FOREACH: %[[tile1:.+]] = scf.forall (%[[tid1:.+]], %[[tid2:.+]]) = +//FOREACH-SAME: to (%[[sz1]], %[[sz2]]) +//FOREACH-SAME: shared_outs(%[[dest:.+]] = %[[init]]) // FOREACH-DAG: %[[iv1:.+]] = affine.apply #[[map1]](%[[tid1]])[%[[lb1]]] // FOREACH: %[[multiIndex1:.+]]:3 = affine.delinearize_index %[[iv1]] into (%[[c3]], %[[d1]], %[[d2]]) : // FOREACH-DAG: %[[iv2:.+]] = affine.apply #[[map1]](%[[tid2]])[%[[lb2]]] // FOREACH: %[[multiIndex2:.+]]:2 = affine.delinearize_index %[[iv2]] into (%[[c11]], %[[d4]]) : // FOREACH: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex1]]#0, %[[multiIndex1]]#1, %[[multiIndex1]]#2, %[[multiIndex2]]#0, %[[multiIndex2]]#1] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] : // FOREACH: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3, 4]{{\]}} : -// FOREACH: perform_concurrently +// FOREACH: in_parallel //FOREACH-NEXT: tensor.parallel_insert_slice %[[sliceFlat]] into %[[dest]][%[[tid1]], %[[tid2]]] [1, 1] [1, 1] : // ----- diff --git a/mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir b/mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir --- a/mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir +++ b/mlir/test/Dialect/Tensor/fold-consecutive-insert-extract-slice.mlir @@ -88,11 +88,12 @@ // CHECK-NOT: tensor.insert_slice // CHECK: tensor.parallel_insert_slice %{{.*}} into %{{.*}}[%{{.*}}, %{{.*}}] [1, 1] [1, 1] : tensor into tensor<1x2xf32> func.func @parallel_insert_slice(%t0: tensor<1x2xf32>, %t1: tensor, %t2: tensor<1x1xf32>) -> tensor<1x2xf32> { + %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index - %r = scf.foreach_thread (%arg2, %arg3) in (%c1, %c2) shared_outs(%arg4 = %t0) -> (tensor<1x2xf32>) { + %r = scf.forall (%arg2, %arg3) = (%c0, %c0) to (%c1, %c2) step (%c1, %c1) shared_outs(%arg4 = %t0) -> (tensor<1x2xf32>) { %inserted_slice = tensor.insert_slice %t1 into %t2[0, 0] [1, 1] [1, 1] : tensor into tensor<1x1xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { tensor.parallel_insert_slice %inserted_slice into %arg4[%arg2, %arg3] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x2xf32> } } diff --git a/mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir b/mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir --- a/mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir +++ b/mlir/test/Dialect/Tensor/fold-reassociative-reshapes.mlir @@ -42,10 +42,12 @@ func.func @rank_reducing_parallel_insert_of_collapse_shape( %t: tensor, %d: tensor, %sz: index, %thr: index) -> tensor { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index %0 = tensor.collapse_shape %t [[0, 1], [2], [3]] : tensor into tensor - %1 = scf.foreach_thread (%iv) in (%thr) shared_outs(%o = %d) -> (tensor) { - scf.foreach_thread.perform_concurrently { + %1 = scf.forall (%iv) = (%c0) to (%thr) step (%c1) shared_outs(%o = %d) -> (tensor) { + scf.forall.in_parallel { tensor.parallel_insert_slice %0 into %o[0, 0, 0, 0][%sz, 1, 1, 5][1, 1, 1, 1] : tensor into tensor } diff --git a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir --- a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir +++ b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir @@ -193,13 +193,14 @@ // CHECK-LABEL: func.func @rank_reducing_parallel_insert_slice func.func @rank_reducing_parallel_insert_slice(%in: tensor<100xf32>, %out: tensor<200x100xf32>) { + %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %num_threads = arith.constant 100 : index - // CHECK: scf.foreach_thread {{.*}} { - %result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs (%o = %out) -> tensor<200x100xf32> { + // CHECK: scf.forall {{.*}} { + %result = scf.forall (%thread_idx) = (%c0) to (%num_threads) step (%c1) shared_outs (%o = %out) -> tensor<200x100xf32> { %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> - scf.foreach_thread.perform_concurrently { + scf.forall.in_parallel { // CHECK: memref.subview %{{.*}}[%{{.*}}] [1] [1] : memref<100xf32, strided<[?], offset: ?>> to memref<1xf32, strided<[?], offset: ?>> // CHECK: memref.subview %{{.*}}[1, %{{.*}}] [1, 1] [1, 1] : memref<200x100xf32, strided<[?, ?], offset: ?>> to memref<1xf32, strided<[?], offset: ?>> tensor.parallel_insert_slice %1 into %o[1, %thread_idx][1, 1][1, 1] : @@ -335,7 +336,7 @@ // CHECK-LABEL: func @dim_not_reading( // CHECK-SAME: %[[t:.*]]: memref, %f: f32, %pos: index) +func.func @dim_not_reading(%t: tensor, %f: f32, %pos: index) -> (tensor, index) { %c0 = arith.constant 0 : index diff --git a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp --- a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp +++ b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp @@ -82,7 +82,7 @@ Option useForeach{ *this, "use-foreach", llvm::cl::desc( - "Use the scf.foreach_thread operation when generating loop nests for " + "Use the scf.forall operation when generating loop nests for " "the extract_slice of collapse_shape pattern"), llvm::cl::init(false)}; @@ -247,8 +247,14 @@ tensor::ExtractSliceFromCollapseHelper &helper, PatternRewriter &rewriter) const override { Location loc = op.getLoc(); - auto foreachOp = rewriter.create( - loc, /*outputs=*/dest, /*numThreads=*/helper.getIterationSpaceSizes(), + unsigned numLoops = helper.getIterationSpaceSizes().size(); + SmallVector lbs(numLoops, + rewriter.create(loc, 0)); + SmallVector steps(numLoops, + rewriter.create(loc, 1)); + auto forallOp = rewriter.create( + loc, lbs, /*numThreads=*/helper.getIterationSpaceSizes(), steps, + /*outputs=*/dest, /*mapping=*/ArrayRef{}, [&](OpBuilder &nestedBuilder, Location loc, ValueRange regionArgs) { unsigned numThreadIdRegionArgs = @@ -262,12 +268,12 @@ auto [tile, insertParams] = helper.emitLoopNestBody(nestedBuilder, loc, outputIvs); // Insert the slice into the destination. - auto term = nestedBuilder.create(loc); + auto term = nestedBuilder.create(loc); nestedBuilder.setInsertionPointToStart(term.getBody()); nestedBuilder.create( loc, tile, outputArgs[0], insertParams); }); - rewriter.replaceOp(op, foreachOp->getResult(0)); + rewriter.replaceOp(op, forallOp->getResult(0)); return success(); } }; diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py --- a/mlir/test/python/dialects/transform_structured_ext.py +++ b/mlir/test/python/dialects/transform_structured_ext.py @@ -181,14 +181,13 @@ # CHECK-COUNT-3: !transform.op<"scf.for"> - @run def testTileExplicitLoopTypeAll(): sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], transform.AnyOpType.get()) types = [ transform.OperationType.get(x) - for x in ["scf.for", "scf.parallel", "scf.foreach_thread"] + for x in ["scf.for", "scf.parallel", "scf.forall"] ] with InsertionPoint(sequence.body): structured.TileOp(types, sequence.bodyTarget, sizes=[2, 3, 4]) @@ -196,7 +195,7 @@ # CHECK-LABEL: TEST: testTileExplicitLoopTypeAll # CHECK: = transform.structured.tile # CHECK-SAME : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, - # CHECK-SAME: !transform.op<"scf.parallel">, !transform.op<"scf.foreach_thread"> + # CHECK-SAME: !transform.op<"scf.parallel">, !transform.op<"scf.forall"> @run def testVectorize():