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

[shraiysh]

Summary

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.

[shraiysh]

Ping for review.

[shraiysh]

Ping for review.

[shraiysh]

Ping for review!

[shraiysh]

Ping for review.

[kiranchandramohan]

Thanks @shraiysh for this patch.

"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."

Could you expand this a bit more? Is it write with the same value?

[kiranchandramohan]

Could you also check whether it is OK to accept atomic updates like the following?

func.func @update_unnecessary_computations(%arg0: memref<i32>) {
  %c0_i32 = arith.constant 0 : i32
  omp.atomic.update   %arg0 : memref<i32> {
  ^bb0(%arg1: i32):
    %0 = arith.cmpi ne, %arg1, %c0_i32 : i32
    %1 = arith.extui %0 : i1 to i32
    %2 = arith.addi %arg1, %1 : i32
    omp.yield(%2 : i32)
  }
  return
}

[shraiysh]

In D126531#3641724, @kiranchandramohan wrote:

Could you also check whether it is OK to accept atomic updates like the following?

func.func @update_unnecessary_computations(%arg0: memref<i32>) {
  %c0_i32 = arith.constant 0 : i32
  omp.atomic.update   %arg0 : memref<i32> {
  ^bb0(%arg1: i32):
    %0 = arith.cmpi ne, %arg1, %c0_i32 : i32
    %1 = arith.extui %0 : i1 to i32
    %2 = arith.addi %arg1, %1 : i32
    omp.yield(%2 : i32)
  }
  return
}

This will work as expected - there will be an atomic update to the location with the new value (lowering to LLVM IR will also work as expected). However, I understand that we might not be able to reach this situation because C/C++/Fortran semantic checks will stop it.

mlir/test/Dialect/OpenMP/canonicalize.mlir

17–23

Could you expand this a bit more? Is it write with the same value?

@kiranchandramohan here is an example where an atomic update operation is essentially an atomic write operation. Due to other simplifications during canonicalization of the ops under atomic update region (possibly unrelated to the openmp dialect) like elimination unused instructions etc. - if the IR reaches this state after after them, then this will be replaced with an atomic write operation writing the same value (as can be seen in CHECK lines under this) to the memory.

[shraiysh]

Ping for review. (This is not urgent for the drop on July 26th, and patches related to those are high priority. This is just a gentle reminder, if anyone is willing to look into this patch. :) )

[mehdi_amini]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

841

Is there a case for AtomicUpdateOp where yieldOp can have more or less than 1 operand?

846

getFirstOp()?

[shraiysh]

Addressed comments.

[shraiysh]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

841

No, removed the check from here. It is already there in the verifier.

mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_
 #define MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_
 
 #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"
 
 #include "mlir/Dialect/OpenMP/OpenMPOpsDialect.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPOpsInterfaces.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPTypeInterfaces.h.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.h.inc"
 
 #endif // MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

let extraClassDeclaration = [{
     Value getVariableOperand(unsigned i) {
       assert(i < 2 && "invalid index position for an operand");
       return i == 0 ? address() : value();
     }
   }];
 }
 
 def AtomicUpdateOp : OpenMP_Op<"atomic.update",
-                               [SingleBlockImplicitTerminator<"YieldOp">]> {
+                               [SingleBlockImplicitTerminator<"YieldOp">,
+                                RecursiveSideEffects]> {
 
   let summary = "performs an atomic update";
 
   let description = [{
     This operation performs an atomic update.
 
     The operand `x` is exactly the same as the operand `x` in the OpenMP
     Standard (OpenMP 5.0, section 2.17.7). It is the address of the variable
let description = [{
     appropriately update a value. For example, while generating LLVM IR, if
     there are no special `atomicrmw` instructions for the operation-type
     combination in atomic update, a compare-exchange loop is generated, where
     the core update operation is directly translated like regular operations by
     the host dialect. The front-end must handle semantic checks for allowed
     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);
   let assemblyFormat = [{
     oilist( `memory_order` `(` custom<ClauseAttr>($memory_order_val) `)`
           | `hint` `(` custom<SynchronizationHint>($hint_val) `)`)
     $x `:` type($x) $region attr-dict
   }];
   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();
   }];
 }
 
 def AtomicCaptureOp : OpenMP_Op<"atomic.capture",
     [SingleBlockImplicitTerminator<"TerminatorOp">]> {
   let summary = "performs an atomic capture";
   let description = [{
     This operation performs an atomic capture.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

if (address().getType().cast<PointerLikeType>().getElementType() !=
     return emitError("address must dereference to value type");
   return verifySynchronizationHint(*this, hint_val());
 }
 
 //===----------------------------------------------------------------------===//
 // 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 ||
         *mo == ClauseMemoryOrderKind::Acquire) {
       return emitError(
           "memory-order must not be acq_rel or acquire for atomic updates");
     }
   }
 
+  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 "
                      "element type is the same as that of the region argument");
   }
 
   return verifySynchronizationHint(*this, hint_val());
 }
 
 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();
 
   if (yieldOp.results().size() != 1)
     return emitError("only updated value must be returned");
   if (yieldOp.results().front().getType() != region().getArgument(0).getType())
     return emitError("input and yielded value must have the same type");
   return success();

mlir/test/Dialect/OpenMP/canonicalize.mlir

+// 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

mlir/test/Dialect/OpenMP/invalid.mlir

^bb0(%xval: i32, %tmp: i32):
     %newval = llvm.add %xval, %expr : i32
     omp.yield (%newval : i32)
   }
   return
 }
 
 // -----
 
-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> {
   ^bb0(%xval: i32):
     %newval = llvm.add %xval, %expr : i32
     omp.yield (%newval : i32)
   }
   return

mlir/test/Dialect/OpenMP/ops.mlir

func.func @omp_atomic_update(%x : memref<i32>, %expr : i32, %xBool : memref<i1>, %exprBool : i1) {
   // CHECK-NEXT:   omp.yield(%[[NEWVAL]] : i1)
   // }
   omp.atomic.update %xBool : memref<i1> {
   ^bb0(%xval: i1):
     %newval = llvm.icmp "eq" %xval, %exprBool : 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
   // CHECK-NEXT:   omp.yield(%[[NEWVAL]] : i32)
   omp.atomic.update hint(none) %x : memref<i32> {
   ^bb0(%xval: i32):
     %newval = llvm.add %xval, %expr : i32
     omp.yield(%newval : i32)