[Mlir-commits] [mlir] [MLIR][OpenMP] Verify loop wrapper properties of omp.parallel (PR #88722)

[Sergio Afonso]

https://github.com/skatrak updated https://github.com/llvm/llvm-project/pull/88722

>From 74c3fc4cf870bced8daa8b22e919b525496adbbc Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Mon, 15 Apr 2024 13:17:56 +0100
Subject: [PATCH] [MLIR][OpenMP] Verify loop wrapper properties of omp.parallel

This patch extends verification of the `omp.parallel` operation to check it is
correctly defined when taking a loop wrapper role.

In OpenMP, a PARALLEL construct can be either a (potenially combined) block
construct or a loop construct, when appearing as part of a composite
construct. This is currently the case for the DISTRIBUTE PARALLEL DO/FOR and
DISTRIBUTE PARALLEL DO/FOR SIMD exclusively.

When used to represent the PARALLEL leaf of a composite construct, it must
follow the rules of a wrapper loop operation in MLIR, and this is what this
patch ensures. No additional restrictions are introduced for PARALLEL block
constructs.
---
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 16 ++++++
 mlir/test/Dialect/OpenMP/invalid.mlir        | 53 ++++++++++++++++++++
 mlir/test/Dialect/OpenMP/ops.mlir            | 20 +++++++-
 3 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 90b49b2528b790..4456e2e4c2c9b3 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1344,6 +1344,22 @@ static LogicalResult verifyPrivateVarList(OpType &op) {
 }
 
 LogicalResult ParallelOp::verify() {
+  // Check that it is a valid loop wrapper if it's taking that role.
+  if (isa<DistributeOp>((*this)->getParentOp())) {
+    if (!isWrapper())
+      return emitOpError() << "must take a loop wrapper role if nested inside "
+                              "of 'omp.distribute'";
+
+    if (LoopWrapperInterface nested = getNestedWrapper()) {
+      // Check for the allowed leaf constructs that may appear in a composite
+      // construct directly after PARALLEL.
+      if (!isa<WsloopOp>(nested))
+        return emitError() << "only supported nested wrapper is 'omp.wsloop'";
+    } else {
+      return emitOpError() << "must not wrap an 'omp.loop_nest' directly";
+    }
+  }
+
   if (getAllocateVars().size() != getAllocatorsVars().size())
     return emitError(
         "expected equal sizes for allocate and allocator variables");
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 88dca1b85ee5f7..6f90c035a6f0b3 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -10,6 +10,59 @@ func.func @unknown_clause() {
 
 // -----
 
+func.func @not_wrapper() {
+  omp.distribute {
+    // expected-error at +1 {{op must take a loop wrapper role if nested inside of 'omp.distribute'}}
+    omp.parallel {
+      %0 = arith.constant 0 : i32
+      omp.terminator
+    }
+    omp.terminator
+  }
+
+  return
+}
+
+// -----
+
+func.func @invalid_nested_wrapper(%lb : index, %ub : index, %step : index) {
+  omp.distribute {
+    // expected-error at +1 {{only supported nested wrapper is 'omp.wsloop'}}
+    omp.parallel {
+      // TODO Remove induction variables from omp.simdloop.
+      omp.simdloop for (%iv) : index = (%lb) to (%ub) step (%step) {
+        omp.loop_nest (%iv2) : index = (%lb) to (%ub) step (%step) {
+          omp.yield
+        }
+        omp.terminator
+      }
+      omp.terminator
+    }
+    omp.terminator
+  }
+
+  return
+}
+
+// -----
+
+func.func @no_nested_wrapper(%lb : index, %ub : index, %step : index) {
+  omp.distribute {
+    // expected-error at +1 {{op must not wrap an 'omp.loop_nest' directly}}
+    omp.parallel {
+      omp.loop_nest (%iv) : index = (%lb) to (%ub) step (%step) {
+        omp.yield
+      }
+      omp.terminator
+    }
+    omp.terminator
+  }
+
+  return
+}
+
+// -----
+
 func.func @if_once(%n : i1) {
   // expected-error at +1 {{`if` clause can appear at most once in the expansion of the oilist directive}}
   omp.parallel if(%n : i1) if(%n : i1) {
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 851d44ad984eef..0840ca17d523e5 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -51,7 +51,7 @@ func.func @omp_terminator() -> () {
   omp.terminator
 }
 
-func.func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : i32) -> () {
+func.func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : i32, %idx : index) -> () {
   // CHECK: omp.parallel if(%{{.*}}) num_threads(%{{.*}} : i32) allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>)
   "omp.parallel" (%if_cond, %num_threads, %data_var, %data_var) ({
 
@@ -85,6 +85,24 @@ func.func @omp_parallel(%data_var : memref<i32>, %if_cond : i1, %num_threads : i
     omp.terminator
   }) {operandSegmentSizes = array<i32: 0,0,1,1,0,0>} : (memref<i32>, memref<i32>) -> ()
 
+  // CHECK: omp.distribute
+  omp.distribute {
+    // CHECK-NEXT: omp.parallel
+    omp.parallel {
+      // CHECK-NEXT: omp.wsloop
+      // TODO Remove induction variables from omp.wsloop.
+      omp.wsloop for (%iv) : index = (%idx) to (%idx) step (%idx) {
+        // CHECK-NEXT: omp.loop_nest
+        omp.loop_nest (%iv2) : index = (%idx) to (%idx) step (%idx) {
+          omp.yield
+        }
+        omp.terminator
+      }
+      omp.terminator
+    }
+    omp.terminator
+  }
+
   return
 }