[Mlir-commits] [mlir] 57ef633 - [mlir][OpenMP] Add omp.atomic.update canonicalization

Shraiysh Vaishay llvmlistbot at llvm.org
Wed Jul 27 11:49:27 PDT 2022


Author: Shraiysh Vaishay
Date: 2022-07-28T00:19:18+05:30
New Revision: 57ef6332597bc27e60f940b07eb63d1473476707

URL: https://github.com/llvm/llvm-project/commit/57ef6332597bc27e60f940b07eb63d1473476707
DIFF: https://github.com/llvm/llvm-project/commit/57ef6332597bc27e60f940b07eb63d1473476707.diff

LOG: [mlir][OpenMP] Add omp.atomic.update canonicalization

This patch adds canonicalization conditions for omp.atomic.update thus
eliminating it when it becomes just a write or a no-op due to other
changes during canonicalization.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D126531

Added: 
    mlir/test/Dialect/OpenMP/canonicalize.mlir

Modified: 
    mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
    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/OpenMPDialect.h b/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
index 05d7637d52d7e..6e8baac4ba5c3 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h
@@ -16,6 +16,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"

diff  --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 7caf1cc99c988..cc657d723a332 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1112,7 +1112,8 @@ def AtomicWriteOp : OpenMP_Op<"atomic.write"> {
 }
 
 def AtomicUpdateOp : OpenMP_Op<"atomic.update",
-                               [SingleBlockImplicitTerminator<"YieldOp">]> {
+                               [SingleBlockImplicitTerminator<"YieldOp">,
+                                RecursiveSideEffects]> {
 
   let summary = "performs an atomic update";
 
@@ -1145,7 +1146,9 @@ def AtomicUpdateOp : OpenMP_Op<"atomic.update",
     operations.
   }];
 
-  let arguments = (ins OpenMP_PointerLikeType:$x,
+  let arguments = (ins Arg<OpenMP_PointerLikeType,
+                           "Address of variable to be updated",
+                           [MemRead, MemWrite]>:$x,
                        DefaultValuedAttr<I64Attr, "0">:$hint_val,
                        OptionalAttr<MemoryOrderKindAttr>:$memory_order_val);
   let regions = (region SizedRegion<1>:$region);
@@ -1156,10 +1159,19 @@ def AtomicUpdateOp : OpenMP_Op<"atomic.update",
   }];
   let hasVerifier = 1;
   let hasRegionVerifier = 1;
+  let hasCanonicalizeMethod = 1;
   let extraClassDeclaration = [{
     Operation* getFirstOp() {
       return &getRegion().front().getOperations().front();
     }
+
+    /// Returns true if the new value is same as old value and the operation is
+    /// a no-op, false otherwise.
+    bool isNoOp();
+
+    /// Returns the new value if the operation is equivalent to just a write
+    /// operation. Otherwise, returns nullptr.
+    Value getWriteOpVal();
   }];
 }
 

diff  --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 96ff6b1c1414e..0e709de347a11 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -836,6 +836,34 @@ LogicalResult AtomicWriteOp::verify() {
 // Verifier for AtomicUpdateOp
 //===----------------------------------------------------------------------===//
 
+bool AtomicUpdateOp::isNoOp() {
+  YieldOp yieldOp = dyn_cast<omp::YieldOp>(getFirstOp());
+  return (yieldOp &&
+          yieldOp.results().front() == getRegion().front().getArgument(0));
+}
+
+Value AtomicUpdateOp::getWriteOpVal() {
+  YieldOp yieldOp = dyn_cast<omp::YieldOp>(getFirstOp());
+  if (yieldOp &&
+      yieldOp.results().front() != getRegion().front().getArgument(0))
+    return yieldOp.results().front();
+  return nullptr;
+}
+
+LogicalResult AtomicUpdateOp::canonicalize(AtomicUpdateOp op,
+                                           PatternRewriter &rewriter) {
+  if (op.isNoOp()) {
+    rewriter.eraseOp(op);
+    return success();
+  }
+  if (Value writeVal = op.getWriteOpVal()) {
+    rewriter.replaceOpWithNewOp<AtomicWriteOp>(
+        op, op.x(), writeVal, op.hint_valAttr(), op.memory_order_valAttr());
+    return success();
+  }
+  return failure();
+}
+
 LogicalResult AtomicUpdateOp::verify() {
   if (auto mo = memory_order_val()) {
     if (*mo == ClauseMemoryOrderKind::Acq_rel ||
@@ -845,6 +873,9 @@ LogicalResult AtomicUpdateOp::verify() {
     }
   }
 
+  if (region().getNumArguments() != 1)
+    return emitError("the region must accept exactly one argument");
+
   if (x().getType().cast<PointerLikeType>().getElementType() !=
       region().getArgument(0).getType()) {
     return emitError("the type of the operand must be a pointer type whose "
@@ -855,12 +886,6 @@ LogicalResult AtomicUpdateOp::verify() {
 }
 
 LogicalResult AtomicUpdateOp::verifyRegions() {
-  if (region().getNumArguments() != 1)
-    return emitError("the region must accept exactly one argument");
-
-  if (region().front().getOperations().size() < 2)
-    return emitError() << "the update region must have at least two operations "
-                          "(binop and terminator)";
 
   YieldOp yieldOp = *region().getOps<YieldOp>().begin();
 

diff  --git a/mlir/test/Dialect/OpenMP/canonicalize.mlir b/mlir/test/Dialect/OpenMP/canonicalize.mlir
new file mode 100644
index 0000000000000..61dbd81cffd91
--- /dev/null
+++ b/mlir/test/Dialect/OpenMP/canonicalize.mlir
@@ -0,0 +1,74 @@
+// RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s
+
+func.func @update_no_op(%x : memref<i32>) {
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval : i32):
+    omp.yield(%xval : i32)
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @update_no_op
+// CHECK-NOT: omp.atomic.update
+
+// -----
+
+func.func @update_write_op(%x : memref<i32>, %value: i32) {
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval : i32):
+    omp.yield(%value : i32)
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @update_write_op
+// CHECK-SAME:            (%[[X:.+]]: memref<i32>, %[[VALUE:.+]]: i32)
+// CHECK: omp.atomic.write %[[X]] = %[[VALUE]] : memref<i32>, i32
+// CHECK-NOT: omp.atomic.update
+
+// -----
+
+func.func @update_normal(%x : memref<i32>, %value: i32) {
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval : i32):
+    %newval = arith.addi %xval, %value : i32
+    omp.yield(%newval : i32)
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @update_normal
+// CHECK: omp.atomic.update
+// CHECK: arith.addi
+// CHECK: omp.yield
+
+// -----
+
+func.func @update_unnecessary_computations(%x: memref<i32>) {
+  %c0 = arith.constant 0 : i32
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = arith.addi %xval, %c0 : i32
+    omp.yield(%newval: i32)
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @update_unnecessary_computations
+// CHECK-NOT: omp.atomic.update
+
+// -----
+
+func.func @update_unnecessary_computations(%x: memref<i32>) {
+  %c0 = arith.constant 0 : i32
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval: i32):
+    %newval = arith.muli %xval, %c0 : i32
+    omp.yield(%newval: i32)
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @update_unnecessary_computations
+// CHECK-NOT: omp.atomic.update
+// CHECK: omp.atomic.write

diff  --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index f1401a7b31954..8408a5ac82d44 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -705,17 +705,6 @@ func.func @omp_atomic_update8(%x: memref<i32>, %expr: i32) {
 
 // -----
 
-func.func @omp_atomic_update9(%x: memref<i32>, %expr: i32) {
-  // expected-error @below {{the update region must have at least two operations (binop and terminator)}}
-  omp.atomic.update %x : memref<i32> {
-  ^bb0(%xval: i32):
-    omp.yield (%xval : i32)
-  }
-  return
-}
-
-// -----
-
 func.func @omp_atomic_update(%x: memref<i32>, %expr: i32) {
   // expected-error @below {{the hints omp_sync_hint_uncontended and omp_sync_hint_contended cannot be combined}}
   omp.atomic.update hint(uncontended, contended) %x : memref<i32> {

diff  --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 393378f70b7e9..4511d116de959 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -757,6 +757,25 @@ func.func @omp_atomic_update(%x : memref<i32>, %expr : i32, %xBool : memref<i1>,
     omp.yield(%newval : i1)
   }
 
+  // CHECK: omp.atomic.update %[[X]] : memref<i32> {
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   omp.yield(%[[XVAL]] : i32)
+  // CHECK-NEXT: }
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval:i32):
+    omp.yield(%xval:i32)
+  }
+
+  // CHECK: omp.atomic.update %[[X]] : memref<i32> {
+  // CHECK-NEXT: (%[[XVAL:.*]]: i32):
+  // CHECK-NEXT:   omp.yield(%{{.+}} : i32)
+  // CHECK-NEXT: }
+  %const = arith.constant 42 : i32
+  omp.atomic.update %x : memref<i32> {
+  ^bb0(%xval:i32):
+    omp.yield(%const:i32)
+  }
+
   // CHECK: omp.atomic.update hint(none) %[[X]] : memref<i32>
   // CHECK-NEXT: (%[[XVAL:.*]]: i32):
   // CHECK-NEXT:   %[[NEWVAL:.*]] = llvm.add %[[XVAL]], %[[EXPR]] : i32


        


More information about the Mlir-commits mailing list