[mlir][openacc] Add canonicalization for standalone data operations for if condition

[clementval]

Summary

This patch add canonicalization for the standalone data operation with constant if condition.
It is extracted from this patch D103325.

[ftynse]

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

174

This needs to be wrapped in rewriter.updateRootInPlace.

[clementval]

Wrap ifCond removal in updateRootinPlace.

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

let assemblyFormat = [{
     ( `wait` `(` $waitOperands^ `:` type($waitOperands) `)` )?
     ( `copyin` `(` $copyinOperands^ `:` type($copyinOperands) `)` )?
     ( `create` `(` $createOperands^ `:` type($createOperands) `)` )?
     ( `create_zero` `(` $createZeroOperands^ `:`
         type($createZeroOperands) `)` )?
     ( `attach` `(` $attachOperands^ `:` type($attachOperands) `)` )?
     attr-dict-with-keyword
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // 2.6.6 Exit Data Directive
 //===----------------------------------------------------------------------===//
 
 def OpenACC_ExitDataOp : OpenACC_Op<"exit_data", [AttrSizedOperandSegments]> {
   let summary = "exit data operation";
let assemblyFormat = [{
     ( `async` `(` $asyncOperand^ `:` type($asyncOperand) `)` )?
     ( `wait_devnum` `(` $waitDevnum^ `:` type($waitDevnum) `)` )?
     ( `wait` `(` $waitOperands^ `:` type($waitOperands) `)` )?
     ( `copyout` `(` $copyoutOperands^ `:` type($copyoutOperands) `)` )?
     ( `delete` `(` $deleteOperands^ `:` type($deleteOperands) `)` )?
     ( `detach` `(` $detachOperands^ `:` type($detachOperands) `)` )?
     attr-dict-with-keyword
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // 2.9 loop Construct
 //===----------------------------------------------------------------------===//
 
 def OpenACC_LoopOp : OpenACC_Op<"loop",
       [AttrSizedOperandSegments,
let assemblyFormat = [{
     ( `wait_devnum` `(` $waitDevnum^ `:` type($waitDevnum) `)` )?
     ( `device_type` `(` $deviceTypeOperands^ `:`
         type($deviceTypeOperands) `)` )?
     ( `wait` `(` $waitOperands^ `:` type($waitOperands) `)` )?
     ( `host` `(` $hostOperands^ `:` type($hostOperands) `)` )?
     ( `device` `(` $deviceOperands^ `:` type($deviceOperands) `)` )?
     attr-dict-with-keyword
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // 2.16.3. Wait Directive
 //===----------------------------------------------------------------------===//
 
 def OpenACC_WaitOp : OpenACC_Op<"wait", [AttrSizedOperandSegments]> {
   let summary = "wait operation";

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

 //===- OpenACC.cpp - OpenACC MLIR Operations ------------------------------===//
 //
 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 // =============================================================================
 
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenACC/OpenACCOpsEnums.cpp.inc"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace acc;
 
 //===----------------------------------------------------------------------===//
 // OpenACC operations
 //===----------------------------------------------------------------------===//
 
static OptionalParseResult parserOptionalOperandAndTypeWithPrefix(
   }
   return llvm::None;
 }
 
 static bool isComputeOperation(Operation *op) {
   return isa<acc::ParallelOp>(op) || isa<acc::LoopOp>(op);
 }
 
+namespace {
+/// Pattern to remove operation without region that have constant false `ifCond`
+/// and remove the condition from the operation if the `ifCond` is a true
+/// constant.
+template <typename OpTy>
+struct RemoveConstantIfCondition : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    // Early return if there is no condition.
+    if (!op.ifCond())
+      return success();
+
+    auto constOp = op.ifCond().template getDefiningOp<ConstantOp>();
+    if (constOp && constOp.getValue().template cast<IntegerAttr>().getInt())
+      rewriter.updateRootInPlace(op, [&]() { op.ifCondMutable().erase(0); });
+    else if (constOp)
+      rewriter.eraseOp(op);
+
+    return success();
+  }
+};
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 /// Parse acc.parallel operation
 /// operation := `acc.parallel` `async` `(` index `)`?
 ///                             `wait` `(` index-list `)`?
 ///                             `num_gangs` `(` value `)`?

 
 Value ExitDataOp::getDataOperand(unsigned i) {
   unsigned numOptional = ifCond() ? 1 : 0;
   numOptional += asyncOperand() ? 1 : 0;
   numOptional += waitDevnum() ? 1 : 0;
   return getOperand(waitOperands().size() + numOptional + i);
 }
 
+void ExitDataOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                             MLIRContext *context) {
+  results.add<RemoveConstantIfCondition<ExitDataOp>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // EnterDataOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(acc::EnterDataOp op) {
   // 2.6.6. Data Enter Directive restriction
   // At least one copyin, create, or attach clause must appear on an enter data
   // directive.

 
 Value EnterDataOp::getDataOperand(unsigned i) {
   unsigned numOptional = ifCond() ? 1 : 0;
   numOptional += asyncOperand() ? 1 : 0;
   numOptional += waitDevnum() ? 1 : 0;
   return getOperand(waitOperands().size() + numOptional + i);
 }
 
+void EnterDataOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                              MLIRContext *context) {
+  results.add<RemoveConstantIfCondition<EnterDataOp>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // InitOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(acc::InitOp initOp) {
   Operation *currOp = initOp;
   while ((currOp = currOp->getParentOp())) {
     if (isComputeOperation(currOp))

 Value UpdateOp::getDataOperand(unsigned i) {
   unsigned numOptional = asyncOperand() ? 1 : 0;
   numOptional += waitDevnum() ? 1 : 0;
   numOptional += ifCond() ? 1 : 0;
   return getOperand(waitOperands().size() + deviceTypeOperands().size() +
                     numOptional + i);
 }
 
+void UpdateOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                           MLIRContext *context) {
+  results.add<RemoveConstantIfCondition<UpdateOp>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // WaitOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(acc::WaitOp waitOp) {
   // The async attribute represent the async clause without value. Therefore the
   // attribute and operand cannot appear at the same time.
   if (waitOp.asyncOperand() && waitOp.async())

mlir/test/Dialect/OpenACC/canonicalize.mlir

+// RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s
+
+func @testenterdataop(%a: memref<10xf32>) -> () {
+  %ifCond = constant true
+  acc.enter_data if(%ifCond) create(%a: memref<10xf32>)
+  return
+}
+
+// CHECK: acc.enter_data create(%{{.*}} : memref<10xf32>)
+
+// -----
+
+func @testenterdataop(%a: memref<10xf32>) -> () {
+  %ifCond = constant false
+  acc.enter_data if(%ifCond) create(%a: memref<10xf32>)
+  return
+}
+
+// CHECK: func @testenterdataop
+// CHECK-NOT: acc.enter_data
+
+// -----
+
+func @testexitdataop(%a: memref<10xf32>) -> () {
+  %ifCond = constant true
+  acc.exit_data if(%ifCond) delete(%a: memref<10xf32>)
+  return
+}
+
+// CHECK: acc.exit_data delete(%{{.*}} : memref<10xf32>)
+
+// -----
+
+func @testexitdataop(%a: memref<10xf32>) -> () {
+  %ifCond = constant false
+  acc.exit_data if(%ifCond) delete(%a: memref<10xf32>)
+  return
+}
+
+// CHECK: func @testexitdataop
+// CHECK-NOT: acc.exit_data
+
+// -----
+
+func @testupdateop(%a: memref<10xf32>) -> () {
+  %ifCond = constant true
+  acc.update if(%ifCond) host(%a: memref<10xf32>)
+  return
+}
+
+// CHECK: acc.update host(%{{.*}} : memref<10xf32>)
+
+// -----
+
+func @testupdateop(%a: memref<10xf32>) -> () {
+  %ifCond = constant false
+  acc.update if(%ifCond) host(%a: memref<10xf32>)
+  return
+}
+
+// CHECK: func @testupdateop
+// CHECK-NOT: acc.update
+
+// ----
+
+func @testenterdataop(%a: memref<10xf32>, %ifCond: i1) -> () {
+  acc.enter_data if(%ifCond) create(%a: memref<10xf32>)
+  return
+}
+
+// CHECK:  func @testenterdataop(%{{.*}}: memref<10xf32>, [[IFCOND:%.*]]: i1)
+// CHECK:    acc.enter_data if(%{{.*}}) create(%{{.*}} : memref<10xf32>)
+
+// -----
+
+func @testexitdataop(%a: memref<10xf32>, %ifCond: i1) -> () {
+  acc.exit_data if(%ifCond) delete(%a: memref<10xf32>)
+  return
+}
+
+// CHECK: func @testexitdataop(%{{.*}}: memref<10xf32>, [[IFCOND:%.*]]: i1)
+// CHECK:   acc.exit_data if(%{{.*}}) delete(%{{.*}} : memref<10xf32>)
+
+// -----
+
+func @testupdateop(%a: memref<10xf32>, %ifCond: i1) -> () {
+  acc.update if(%ifCond) host(%a: memref<10xf32>)
+  return
+}
+
+// CHECK:  func @testupdateop(%{{.*}}: memref<10xf32>, [[IFCOND:%.*]]: i1)
+// CHECK:    acc.update if(%{{.*}}) host(%{{.*}} : memref<10xf32>)