[mlir][acc] Erase empty kernel_environment ops during canonicalization

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

@@ -2184,6 +2184,8 @@ def OpenACC_KernelEnvironmentOp : OpenACC_Op<"kernel_environment",
     )
     $region attr-dict
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -1042,6 +1042,37 @@ struct RemoveConstantIfConditionWithRegion : public OpRewritePattern<OpTy> {
   }
 };
 
+/// Remove empty acc.kernel_environment operations. If the operation has wait
+/// operands, create a acc.wait operation to preserve synchronization.
+struct RemoveEmptyKernelEnvironment
+    : public OpRewritePattern<acc::KernelEnvironmentOp> {
+  using OpRewritePattern<acc::KernelEnvironmentOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(acc::KernelEnvironmentOp op,
+                                PatternRewriter &rewriter) const override {
+    assert(op->getNumRegions() == 1 && "expected op to have one region");
+
+    Block &block = op.getRegion().front();
+    if (!block.empty())
+      return failure();
+
+    // Remove empty kernel environment
+    // preserve synchronization by creating acc.wait operation if needed
+    if (!op.getWaitOperands().empty()) {
+      rewriter.replaceOpWithNewOp<acc::WaitOp>(
+          op,
+          /*waitOperands=*/op.getWaitOperands(),
+          /*asyncOperand=*/Value(),
+          /*waitDevnum=*/Value(),
+          /*async=*/nullptr,
+          /*ifCond=*/Value());
+    } else
+      rewriter.eraseOp(op);
+
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // Recipe Region Helpers
 //===----------------------------------------------------------------------===//
@@ -2690,6 +2721,15 @@ void acc::HostDataOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<RemoveConstantIfConditionWithRegion<HostDataOp>>(context);
 }
 
+//===----------------------------------------------------------------------===//
+// KernelEnvironmentOp
+//===----------------------------------------------------------------------===//
+
+void acc::KernelEnvironmentOp::getCanonicalizationPatterns(
+    RewritePatternSet &results, MLIRContext *context) {
+  results.add<RemoveEmptyKernelEnvironment>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // LoopOp
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/OpenACC/canonicalize.mlir

@@ -219,3 +219,29 @@ func.func @update_unnecessary_computations(%x: memref<i32>) {
 // CHECK-LABEL: func.func @update_unnecessary_computations
 // CHECK-NOT: acc.atomic.update
 // CHECK: acc.atomic.write
+
+// -----
+
+func.func @remove_empty_kernel_environment() {
+  acc.kernel_environment {
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @remove_empty_kernel_environment
+// CHECK-NOT: acc.kernel_environment
+// CHECK: return
+
+// -----
+
+func.func @kernel_environment_with_wait(%q1: i32, %q2: i32) {
+  acc.kernel_environment wait({%q1 : i32, %q2 : i32}) {
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @kernel_environment_with_wait
+// CHECK-SAME: ([[Q1:%.*]]: i32, [[Q2:%.*]]: i32)
+// CHECK-NOT: acc.kernel_environment
+// CHECK: acc.wait([[Q1]], [[Q2]] : i32, i32)
+// CHECK: return