[CIR] Canonicalization: leverage MLIR traits and folding (#173305)

Replace custom rewrite patterns with dedicated fold implementations for
ScopeOp, or rely on DCE in cases of effect-less SwitchOp.

Author: xlauko

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -842,7 +842,7 @@ defvar CIR_YieldableScopes = [
 ];
 
 def CIR_YieldOp : CIR_Op<"yield", [
-  ReturnLike, Terminator, ParentOneOf<CIR_YieldableScopes>
+  ReturnLike, Terminator, ParentOneOf<CIR_YieldableScopes>, NoMemoryEffect
 ]> {
   let summary = "Represents the default branching behaviour of a region";
   let description = [{
@@ -999,7 +999,8 @@ def CIR_ResumeFlatOp : CIR_Op<"resume.flat", [
 
 def CIR_ScopeOp : CIR_Op<"scope", [
   DeclareOpInterfaceMethods<RegionBranchOpInterface>,
-  RecursivelySpeculatable, AutomaticAllocationScope, NoRegionArguments
+  RecursivelySpeculatable, AutomaticAllocationScope, NoRegionArguments,
+  RecursiveMemoryEffects
 ]> {
   let summary = "Represents a C/C++ scope";
   let description = [{
@@ -1026,6 +1027,7 @@ def CIR_ScopeOp : CIR_Op<"scope", [
   let results = (outs Optional<CIR_AnyType>:$results);
   let regions = (region AnyRegion:$scopeRegion);
 
+  let hasFolder = 1;
   let hasVerifier = 1;
   let skipDefaultBuilders = 1;
   let assemblyFormat = [{
@@ -1104,7 +1106,8 @@ def CIR_CaseOp : CIR_Op<"case", [
 def CIR_SwitchOp : CIR_Op<"switch", [
   SameVariadicOperandSize,
   DeclareOpInterfaceMethods<RegionBranchOpInterface>,
-  RecursivelySpeculatable, AutomaticAllocationScope, NoRegionArguments
+  RecursivelySpeculatable, AutomaticAllocationScope, NoRegionArguments,
+  RecursiveMemoryEffects
 ]> {
   let summary = "Switch operation";
   let description = [{

clang/lib/CIR/Dialect/IR/CIRDialect.cpp

@@ -1240,6 +1240,27 @@ LogicalResult cir::ScopeOp::verify() {
   return success();
 }
 
+LogicalResult cir::ScopeOp::fold(FoldAdaptor /*adaptor*/,
+                                 SmallVectorImpl<OpFoldResult> &results) {
+  // Only fold "trivial" scopes: a single block containing only a `cir.yield`.
+  if (!getRegion().hasOneBlock())
+    return failure();
+  Block &block = getRegion().front();
+  if (block.getOperations().size() != 1)
+    return failure();
+
+  auto yield = dyn_cast<cir::YieldOp>(block.front());
+  if (!yield)
+    return failure();
+
+  // Only fold when the scope produces a value.
+  if (getNumResults() != 1 || yield.getNumOperands() != 1)
+    return failure();
+
+  results.push_back(yield.getOperand(0));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // BrOp
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp

@@ -66,42 +66,6 @@ struct RemoveRedundantBranches : public OpRewritePattern<BrOp> {
   }
 };
 
-struct RemoveEmptyScope : public OpRewritePattern<ScopeOp> {
-  using OpRewritePattern<ScopeOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(ScopeOp op,
-                                PatternRewriter &rewriter) const final {
-    // TODO: Remove this logic once CIR uses MLIR infrastructure to remove
-    // trivially dead operations
-    if (op.isEmpty()) {
-      rewriter.eraseOp(op);
-      return success();
-    }
-
-    Region &region = op.getScopeRegion();
-    if (region.getBlocks().front().getOperations().size() == 1 &&
-        isa<YieldOp>(region.getBlocks().front().front())) {
-      rewriter.eraseOp(op);
-      return success();
-    }
-
-    return failure();
-  }
-};
-
-struct RemoveEmptySwitch : public OpRewritePattern<SwitchOp> {
-  using OpRewritePattern<SwitchOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(SwitchOp op,
-                                PatternRewriter &rewriter) const final {
-    if (!(op.getBody().empty() || isa<YieldOp>(op.getBody().front().front())))
-      return failure();
-
-    rewriter.eraseOp(op);
-    return success();
-  }
-};
-
 //===----------------------------------------------------------------------===//
 // CIRCanonicalizePass
 //===----------------------------------------------------------------------===//
@@ -124,8 +88,7 @@ struct CIRCanonicalizePass
 void populateCIRCanonicalizePatterns(RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
-    RemoveRedundantBranches,
-    RemoveEmptyScope
+    RemoveRedundantBranches
   >(patterns.getContext());
   // clang-format on
 }
@@ -138,7 +101,6 @@ void CIRCanonicalizePass::runOnOperation() {
   // Collect operations to apply patterns.
   llvm::SmallVector<Operation *, 16> ops;
   getOperation()->walk([&](Operation *op) {
-    assert(!cir::MissingFeatures::switchOp());
     assert(!cir::MissingFeatures::tryOp());
     assert(!cir::MissingFeatures::callOp());
 

clang/test/CIR/Transforms/canonicalize.cir

@@ -22,6 +22,19 @@ module {
   // CHECK-NEXT:   cir.return
   // CHECK-NEXT: }
 
+  cir.func @scope_yield_value_fold() -> !u32i {
+    %0 = cir.const #cir.int<7> : !u32i
+    %1 = cir.scope {
+      cir.yield %0 : !u32i
+    } : !u32i
+    cir.return %1 : !u32i
+  }
+  // CHECK:      cir.func{{.*}} @scope_yield_value_fold() -> !u32i {
+  // CHECK-NEXT:   %[[C:.*]] = cir.const #cir.int<7> : !u32i
+  // CHECK-NOT:    cir.scope
+  // CHECK:        cir.return %[[C]] : !u32i
+  // CHECK-NEXT: }
+
   cir.func @empty_scope() {
     cir.scope {
     }
@@ -31,6 +44,19 @@ module {
   // CHECK-NEXT:   cir.return
   // CHECK-NEXT: }
 
+  cir.func @dead_switch() {
+    %0 = cir.const #cir.int<0> : !s32i
+    cir.switch (%0 : !s32i) {
+      cir.yield
+    }
+    cir.return
+  }
+  // CHECK:      cir.func{{.*}} @dead_switch() {
+  // CHECK-NOT:   %[[Z:.*]] = cir.const #cir.int<0> : !s32i
+  // CHECK-NOT:    cir.switch
+  // CHECK-NEXT:   cir.return
+  // CHECK-NEXT: }
+
   cir.func @unary_not(%arg0: !cir.bool) -> !cir.bool {
     %0 = cir.unary(not, %arg0) : !cir.bool, !cir.bool
     %1 = cir.unary(not, %0) : !cir.bool, !cir.bool