[Mlir-commits] [mlir] 70fe6ad - [MLIR][OpenMP] Make omp.distribute into a loop wrapper (#87239)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Apr 16 02:24:26 PDT 2024
Author: Sergio Afonso
Date: 2024-04-16T10:24:23+01:00
New Revision: 70fe6ad535365de20c3a960dcfe97bd5cf68abb5
URL: https://github.com/llvm/llvm-project/commit/70fe6ad535365de20c3a960dcfe97bd5cf68abb5
DIFF: https://github.com/llvm/llvm-project/commit/70fe6ad535365de20c3a960dcfe97bd5cf68abb5.diff
LOG: [MLIR][OpenMP] Make omp.distribute into a loop wrapper (#87239)
This patch updates the definition of `omp.distribute` to enforce the
restrictions of a wrapper operation.
Added:
Modified:
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
mlir/test/Dialect/OpenMP/invalid.mlir
mlir/test/Dialect/OpenMP/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 3abdbe3adfd0be..7e677c9839f602 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -840,7 +840,8 @@ def YieldOp : OpenMP_Op<"yield",
//===----------------------------------------------------------------------===//
def DistributeOp : OpenMP_Op<"distribute", [AttrSizedOperandSegments,
DeclareOpInterfaceMethods<LoopWrapperInterface>,
- RecursiveMemoryEffects]> {
+ RecursiveMemoryEffects,
+ SingleBlockImplicitTerminator<"TerminatorOp">]> {
let summary = "distribute construct";
let description = [{
The distribute construct specifies that the iterations of one or more loops
@@ -855,15 +856,28 @@ def DistributeOp : OpenMP_Op<"distribute", [AttrSizedOperandSegments,
The distribute loop construct specifies that the iterations of the loop(s)
will be executed in parallel by threads in the current context. These
iterations are spread across threads that already exist in the enclosing
- region. The lower and upper bounds specify a half-open range: the
- range includes the lower bound but does not include the upper bound. If the
- `inclusive` attribute is specified then the upper bound is also included.
+ region.
+
+ The body region can contain a single block which must contain a single
+ operation and a terminator. The operation must be another compatible loop
+ wrapper or an `omp.loop_nest`.
The `dist_schedule_static` attribute specifies the schedule for this
loop, determining how the loop is distributed across the parallel threads.
The optional `schedule_chunk` associated with this determines further
controls this distribution.
+ ```mlir
+ omp.distribute <clauses> {
+ omp.loop_nest (%i1, %i2) : index = (%c0, %c0) to (%c10, %c10) step (%c1, %c1) {
+ %a = load %arrA[%i1, %i2] : memref<?x?xf32>
+ %b = load %arrB[%i1, %i2] : memref<?x?xf32>
+ %sum = arith.addf %a, %b : f32
+ store %sum, %arrC[%i1, %i2] : memref<?x?xf32>
+ omp.yield
+ }
+ }
+ ```
// TODO: private_var, firstprivate_var, lastprivate_var, collapse
}];
let arguments = (ins
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 90b49b2528b790..365b3c4246bfd9 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1656,6 +1656,17 @@ LogicalResult DistributeOp::verify() {
return emitError(
"expected equal sizes for allocate and allocator variables");
+ if (!isWrapper())
+ return emitOpError() << "must be a loop wrapper";
+
+ if (LoopWrapperInterface nested = getNestedWrapper()) {
+ // Check for the allowed leaf constructs that may appear in a composite
+ // construct directly after DISTRIBUTE.
+ if (!isa<ParallelOp, SimdLoopOp>(nested))
+ return emitError() << "only supported nested wrappers are 'omp.parallel' "
+ "and 'omp.simdloop'";
+ }
+
return success();
}
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 88dca1b85ee5f7..f8739887214bfd 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1866,7 +1866,16 @@ func.func @omp_target_depend(%data_var: memref<i32>) {
// -----
-func.func @omp_distribute(%data_var : memref<i32>) -> () {
+func.func @omp_distribute_schedule(%chunk_size : i32) -> () {
+ // expected-error @below {{op chunk size set without dist_schedule_static being present}}
+ "omp.distribute"(%chunk_size) <{operandSegmentSizes = array<i32: 1, 0, 0>}> ({
+ "omp.terminator"() : () -> ()
+ }) : (i32) -> ()
+}
+
+// -----
+
+func.func @omp_distribute_allocate(%data_var : memref<i32>) -> () {
// expected-error @below {{expected equal sizes for allocate and allocator variables}}
"omp.distribute"(%data_var) <{operandSegmentSizes = array<i32: 0, 1, 0>}> ({
"omp.terminator"() : () -> ()
@@ -1875,6 +1884,29 @@ func.func @omp_distribute(%data_var : memref<i32>) -> () {
// -----
+func.func @omp_distribute_wrapper() -> () {
+ // expected-error @below {{op must be a loop wrapper}}
+ "omp.distribute"() ({
+ %0 = arith.constant 0 : i32
+ "omp.terminator"() : () -> ()
+ }) : () -> ()
+}
+
+// -----
+
+func.func @omp_distribute_nested_wrapper(%data_var : memref<i32>) -> () {
+ // expected-error @below {{only supported nested wrappers are 'omp.parallel' and 'omp.simdloop'}}
+ "omp.distribute"() ({
+ "omp.wsloop"() ({
+ %0 = arith.constant 0 : i32
+ "omp.terminator"() : () -> ()
+ }) : () -> ()
+ "omp.terminator"() : () -> ()
+ }) : () -> ()
+}
+
+// -----
+
omp.private {type = private} @x.privatizer : i32 alloc {
^bb0(%arg0: i32):
%0 = arith.constant 0.0 : f32
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 851d44ad984eef..018b82e995d7d3 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -559,30 +559,54 @@ func.func @omp_simdloop_pretty_multiple(%lb1 : index, %ub1 : index, %step1 : ind
}
// CHECK-LABEL: omp_distribute
-func.func @omp_distribute(%chunk_size : i32, %data_var : memref<i32>) -> () {
+func.func @omp_distribute(%chunk_size : i32, %data_var : memref<i32>, %arg0 : i32) -> () {
// CHECK: omp.distribute
"omp.distribute" () ({
- omp.terminator
+ "omp.loop_nest" (%arg0, %arg0, %arg0) ({
+ ^bb0(%iv: i32):
+ "omp.yield"() : () -> ()
+ }) : (i32, i32, i32) -> ()
+ "omp.terminator"() : () -> ()
}) {} : () -> ()
// CHECK: omp.distribute
omp.distribute {
- omp.terminator
+ omp.loop_nest (%iv) : i32 = (%arg0) to (%arg0) step (%arg0) {
+ omp.yield
+ }
}
// CHECK: omp.distribute dist_schedule_static
omp.distribute dist_schedule_static {
- omp.terminator
+ omp.loop_nest (%iv) : i32 = (%arg0) to (%arg0) step (%arg0) {
+ omp.yield
+ }
}
// CHECK: omp.distribute dist_schedule_static chunk_size(%{{.+}} : i32)
omp.distribute dist_schedule_static chunk_size(%chunk_size : i32) {
- omp.terminator
+ omp.loop_nest (%iv) : i32 = (%arg0) to (%arg0) step (%arg0) {
+ omp.yield
+ }
}
// CHECK: omp.distribute order(concurrent)
omp.distribute order(concurrent) {
- omp.terminator
+ omp.loop_nest (%iv) : i32 = (%arg0) to (%arg0) step (%arg0) {
+ omp.yield
+ }
}
// CHECK: omp.distribute allocate(%{{.+}} : memref<i32> -> %{{.+}} : memref<i32>)
omp.distribute allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
- omp.terminator
+ omp.loop_nest (%iv) : i32 = (%arg0) to (%arg0) step (%arg0) {
+ omp.yield
+ }
+ }
+ // CHECK: omp.distribute
+ omp.distribute {
+ // TODO Remove induction variables from omp.simdloop.
+ omp.simdloop for (%iv) : i32 = (%arg0) to (%arg0) step (%arg0) {
+ omp.loop_nest (%iv2) : i32 = (%arg0) to (%arg0) step (%arg0) {
+ omp.yield
+ }
+ omp.yield
+ }
}
return
}
More information about the Mlir-commits
mailing list