[Mlir-commits] [mlir] Adding masked operations to OpenMP Dialect (PR #96022)

[Anchu Rajendran S]

https://github.com/anchuraj created https://github.com/llvm/llvm-project/pull/96022

Adding MLIR Op support for omp masked. Omp masked is introduced in 5.2 standard and allows a parallel region to be executed by threads specified by a programmer.  This is achieved with the help of filter clause which helps to specify thread id expected to execute the region.

>From 069de8fc4f0c73fccef2eb17027aaf9bf0d86150 Mon Sep 17 00:00:00 2001
From: Anchu Rajendran <asudhaku at amd.com>
Date: Sat, 15 Jun 2024 00:03:59 -0500
Subject: [PATCH] Adding masked operations to OpenMP Dialect

---
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 21 +++++++++++++++++++
 mlir/test/Dialect/OpenMP/ops.mlir             | 15 +++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 122abbe7cc975..08a89e18510a7 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -2179,4 +2179,25 @@ def ReductionOp : OpenMP_Op<"reduction"> {
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// [Spec 5.2] 10.5 masked Construct
+//===----------------------------------------------------------------------===//
+def MaskedOp : OpenMP_Op<"masked"> {
+  let summary = "masked construct";
+  let description = [{
+    Masked construct allows to specify a structured block to be executed by a subset of 
+    threads of the current team. Filter clause allows to select the threads expected to
+    execute the region
+  }];
+
+  let arguments = (ins Optional<I32>:$filteredThreadId);
+  let regions = (region AnyRegion:$region);
+
+  let assemblyFormat = [{
+    oilist(
+      `filter` `(` $filteredThreadId `:` type($filteredThreadId) `)`
+    ) $region attr-dict
+  }];
+}
+
 #endif // OPENMP_OPS
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 0d5fd9383a92f..fe24b6484f292 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -16,6 +16,21 @@ func.func @omp_master() -> () {
   return
 }
 
+// CHECK-LABEL: omp_masked
+func.func @omp_masked(%filtered_thread_id : i32) -> () {
+
+
+  // CHECK: omp.masked filter(%{{.*}} : i32)
+  "omp.masked" (%filtered_thread_id) ({
+    omp.terminator
+  }) : (i32) -> ()
+
+  // CHECK: omp.masked
+  "omp.masked" () ({
+    omp.terminator
+  }) : () -> ()
+  return
+}
 func.func @omp_taskwait() -> () {
   // CHECK: omp.taskwait
   omp.taskwait