[mlir][openacc] Add canonicalization pattern for acc.host_data

[clementval]

Summary

Add if condition removal pattern for acc.host_data in a same way as
acc.enter_data, acc.exit_data and acc.update.

The condition is removed from the op if it is a true constant. If
it is a false constant then the region is inlined before the op
and the op is removed.

[vzakhari]

LGTM. Thank you!

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

328

nit: the braces are not needed.

[clementval]

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

328

Right. Thanks for catching. I had more code in there in a previous version. I'll update that before pushing.

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

let assemblyFormat = [{
     oilist(
         `if` `(` $ifCond `)`
       | `dataOperands` `(` $dataOperands `:` type($dataOperands) `)`
     )
     $region attr-dict-with-keyword
   }];
 
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // 2.9 loop Construct
 //===----------------------------------------------------------------------===//
 
 def OpenACC_LoopOp : OpenACC_Op<"loop",
     [AttrSizedOperandSegments, RecursiveMemoryEffects]> {

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

LogicalResult matchAndRewrite(OpTy op,
     if (constAttr.getInt())
       rewriter.updateRootInPlace(op, [&]() { op.getIfCondMutable().erase(0); });
     else
       rewriter.eraseOp(op);
 
     return success();
   }
 };
+
+/// Replaces the given op with the contents of the given single-block region,
+/// using the operands of the block terminator to replace operation results.
+static void replaceOpWithRegion(PatternRewriter &rewriter, Operation *op,
+                                Region &region, ValueRange blockArgs = {}) {
+  assert(llvm::hasSingleElement(region) && "expected single-region block");
+  Block *block = &region.front();
+  Operation *terminator = block->getTerminator();
+  ValueRange results = terminator->getOperands();
+  rewriter.inlineBlockBefore(block, op, blockArgs);
+  rewriter.replaceOp(op, results);
+  rewriter.eraseOp(terminator);
+}
+
+/// Pattern to remove operation with region that have constant false `ifCond`
+/// and remove the condition from the operation if the `ifCond` is constant
+/// true.
+template <typename OpTy>
+struct RemoveConstantIfConditionWithRegion : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    // Early return if there is no condition.
+    Value ifCond = op.getIfCond();
+    if (!ifCond)
+      return failure();
+
+    IntegerAttr constAttr;
+    if (!matchPattern(ifCond, m_Constant(&constAttr)))
+      return failure();
+    if (constAttr.getInt())
+      rewriter.updateRootInPlace(op, [&]() { op.getIfCondMutable().erase(0); });
+    else
+      replaceOpWithRegion(rewriter, op, op.getRegion());
+
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 /// Check dataOperands for acc.parallel, acc.serial and acc.kernels.
 template <typename Op>
return emitError("at least one operand must appear on the host_data "
                      "operation");
 
   for (mlir::Value operand : getDataOperands())
     if (!mlir::isa<acc::UseDeviceOp>(operand.getDefiningOp()))
       return emitError("expect data entry operation as defining op");
   return success();
 }
 
+void acc::HostDataOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                  MLIRContext *context) {
+  results.add<RemoveConstantIfConditionWithRegion<HostDataOp>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // LoopOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult
 parseGangClause(OpAsmParser &parser,
                 std::optional<OpAsmParser::UnresolvedOperand> &gangNum,
                 Type &gangNumType,

mlir/test/Dialect/OpenACC/canonicalize.mlir

func.func @testupdateop(%a: memref<f32>, %ifCond: i1) -> () {
   %0 = acc.getdeviceptr varPtr(%a : memref<f32>) -> memref<f32>
   acc.update_host accPtr(%0 : memref<f32>) to varPtr(%a : memref<f32>)
   acc.update if(%ifCond) dataOperands(%0: memref<f32>)
   return
 }
 
 // CHECK:  func @testupdateop(%{{.*}}: memref<f32>, [[IFCOND:%.*]]: i1)
 // CHECK:    acc.update if(%{{.*}}) dataOperands(%{{.*}} : memref<f32>)
+
+// -----
+
+func.func @testhostdataop(%a: memref<f32>, %ifCond: i1) -> () {
+  %0 = acc.use_device varPtr(%a : memref<f32>) -> memref<f32>
+  %false = arith.constant false
+  acc.host_data dataOperands(%0 : memref<f32>) if(%false) {
+    acc.loop {
+      acc.yield
+    }
+    acc.loop {
+      acc.yield
+    }
+    acc.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @testhostdataop
+// CHECK-NOT: acc.host_data
+// CHECK: acc.loop
+// CHECK: acc.yield
+// CHECK: acc.loop
+// CHECK: acc.yield
+
+// -----
+
+func.func @testhostdataop(%a: memref<f32>, %ifCond: i1) -> () {
+  %0 = acc.use_device varPtr(%a : memref<f32>) -> memref<f32>
+  %true = arith.constant true
+  acc.host_data dataOperands(%0 : memref<f32>) if(%true) {
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @testhostdataop
+// CHECK: acc.host_data dataOperands(%{{.*}} : memref<f32>) {