[Comb] Let parent op decide if cross-block canonicalizations should be allowed

include/circt/Dialect/Comb/Comb.td

@@ -38,6 +38,8 @@ def CombDialect : Dialect {
 class CombOp<string mnemonic, list<Trait> traits = []> :
     Op<CombDialect, mnemonic, traits>;
 
+def NonCombSemanticPreservingRegion : NativeOpTrait<"NonCombSemanticPreservingRegion">;
+
 include "circt/Dialect/HW/HWTypes.td"
 include "circt/Dialect/Comb/Combinational.td"
 

include/circt/Dialect/Comb/CombOps.h

@@ -37,6 +37,18 @@ class PatternRewriter;
 namespace circt {
 namespace comb {
 
+/// This trait declares all regions of the operation it is attached to as
+/// non-comb-semantic-perserving. This means, the Comb dialect will not attempt
+/// any folding, canonicalization, and optimization across region and block
+/// boundaries. More precisely, if a comb operation is defined in block#1 inside
+/// region#0 of an operation with this trait, comb canonicalizers will not
+/// consider def-use edges coming from outside region#0 as well as outside
+/// block#1.
+template <typename ConcreteType>
+class NonCombSemanticPreservingRegion
+    : public mlir::OpTrait::TraitBase<ConcreteType,
+                                      NonCombSemanticPreservingRegion> {};
+
 using llvm::KnownBits;
 
 /// Compute "known bits" information about the specified value - the set of bits

lib/Dialect/Comb/CombFolds.cpp

@@ -21,18 +21,30 @@ using namespace circt;
 using namespace comb;
 using namespace matchers;
 
-/// In comb, we assume no knowledge of the semantics of cross-block dataflow. As
-/// such, cross-block dataflow is interpreted as a canonicalization barrier.
-/// This is a conservative approach which:
-/// 1. still allows for efficient canonicalization for the common CIRCT usecase
-///    of comb (comb logic nested inside single-block hw.module's)
-/// 2. allows comb operations to be used in non-HW container ops - that may use
-///    MLIR blocks and regions to represent various forms of hierarchical
-///    abstractions, thus allowing comb to compose with other dialects.
+/// Check if folding or canonicalizing across blocks should be considered
+/// illegal.
 static bool hasOperandsOutsideOfBlock(Operation *op) {
-  Block *thisBlock = op->getBlock();
   return llvm::any_of(op->getOperands(), [&](Value operand) {
-    return operand.getParentBlock() != thisBlock;
+    Operation *parentOp = op;
+    while (parentOp->getParentRegion() != operand.getParentRegion()) {
+      parentOp = parentOp->getParentOp();
+      // If the parent block or region is not (yet) attached to a parent
+      // region/op, conservatively assume it will be attached to a
+      // non-semantic-perserving operation.
+      if (!parentOp || parentOp->hasTrait<NonCombSemanticPreservingRegion>())
+        return true;
+    }
+
+    if (operand.getParentBlock() == parentOp->getBlock())
+      return false;
+
+    // If they aren't defined in the same block, we also need to check if the
+    // trait is added to the regions parent op
+    Operation *parentsParent = parentOp->getParentOp();
+    if (!parentsParent)
+      return true;
+
+    return parentsParent->hasTrait<NonCombSemanticPreservingRegion>();
   });
 }
 

test/Dialect/Calyx/remove-comb-groups.mlir

@@ -12,8 +12,7 @@ calyx.component @main(%go: i1 {go}, %clk: i1 {clk}, %reset: i1 {reset}) -> (%don
 // CHECK:   calyx.assign %eq_reg.write_en = %true : i1
 // CHECK:   calyx.assign %eq.left = %true : i1
 // CHECK:   calyx.assign %eq.right = %true : i1
-// CHECK:    %0 = comb.and %eq_reg.done : i1
-// CHECK:   calyx.group_done %0 ? %true : i1
+// CHECK:   calyx.group_done %eq_reg.done ? %true : i1
     calyx.comb_group @Cond {
       calyx.assign %eq.left =  %c1_1 : i1
       calyx.assign %eq.right = %c1_1 : i1
@@ -60,8 +59,7 @@ calyx.component @main(%go: i1 {go}, %clk: i1 {clk}, %reset: i1 {reset}) -> (%don
 // CHECK:   calyx.assign %eq_reg.write_en = %true : i1
 // CHECK:   calyx.assign %eq.left = %true : i1
 // CHECK:   calyx.assign %eq.right = %true : i1
-// CHECK:   %0 = comb.and %eq_reg.done : i1
-// CHECK:   calyx.group_done %0 ? %true : i1
+// CHECK:   calyx.group_done %eq_reg.done ? %true : i1
 // CHECK: }
     calyx.comb_group @Cond1 {
       calyx.assign %eq.left =  %c1_1 : i1
@@ -74,8 +72,7 @@ calyx.component @main(%go: i1 {go}, %clk: i1 {clk}, %reset: i1 {reset}) -> (%don
 // CHECK:   calyx.assign %eq_reg.write_en = %true : i1
 // CHECK:   calyx.assign %eq.left = %true : i1
 // CHECK:   calyx.assign %eq.right = %true : i1
-// CHECK:   %0 = comb.and %eq_reg.done : i1
-// CHECK:   calyx.group_done %0 ? %true : i1
+// CHECK:   calyx.group_done %eq_reg.done ? %true : i1
 // CHECK: }
     calyx.comb_group @Cond2 {
       calyx.assign %eq.left =  %c1_1 : i1

test/Dialect/Comb/canonicalization.mlir

@@ -1534,60 +1534,21 @@ hw.module @OrMuxSameTrueValueAndZero(in %tag_0: i1, in %tag_1: i1, in %tag_2: i1
   hw.output %4 : i4
 }
 
-// CHECK-LABEL:   "test.acrossBlockCanonicalizationBarrier"() ({
-// CHECK:         ^bb0(%[[A:.*]]: i32, %[[B:.*]]: i32, %[[C:.*]]: i32):
-// CHECK:           %[[ADD1:.*]] = comb.add %[[A]], %[[B]] : i32
-// CHECK:           "terminator"() : () -> ()
-// CHECK:         ^bb1:
-// CHECK:           %[[ADD2:.*]] = comb.add %[[ADD1]], %[[C]] : i32
-// CHECK:           "terminator"(%[[ADD2]]) : (i32) -> ()
-// CHECK:         }) : () -> ()
-"test.acrossBlockCanonicalizationBarrier"() ({
-  ^bb0(%a : i32, %b : i32, %c : i32):
-    %add1 = comb.add %a, %b : i32
-    "terminator"() : () -> ()
+// CHECK-LABEL: @acrossBlockCanonicalizationConstant
+// CHECK-NEXT: %c-1_i32 = hw.constant -1
+// CHECK-NEXT: hw.output %c-1_i32, %c-1_i32 : i32, i32
+hw.module @acrossBlockCanonicalizationConstant(in %arg0: i32, out out0 : i32, out out1 : i32) {
+  %0 = hw.constant 0 : i32
+  %res:2 = scf.execute_region -> (i32, i32) {
+    %1 = hw.constant -1 : i32
+    %2 = comb.or %0, %1, %1 : i32
+    cf.br ^bb1
   ^bb1:
-    // Canonicalization should _not_ pull %add1 into a combined add op when %add1
-    // is defined in another block.
-    %add2 = comb.add %add1, %c : i32
-    "terminator"(%add2) : (i32) -> ()
-}) : () -> ()
-
-// CHECK-LABEL: "test.acrossBlockCanonicalizationBarrierFlattenAndIdem"
-// CHECK: ^bb1:
-// CHECK-NEXT: %[[OUT:.+]] = comb.or %0, %1, %2 : i32
-// CHECK-NEXT: "terminator"(%[[OUT]]) : (i32) -> ()
-"test.acrossBlockCanonicalizationBarrierFlattenAndIdem"() ({
-^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
-  %0 = comb.or %arg0, %arg1 : i32
-  %1 = comb.or %arg1, %arg2 : i32
-  %2 = comb.or %arg0, %arg2 : i32
-  "terminator"() : () -> ()
-^bb1:  // no predecessors
-  // Flatten and unique, but not across blocks.
-  %3 = comb.or %0, %1 : i32
-  %4 = comb.or %1, %2 : i32
-  %5 = comb.or %3, %4, %0, %1, %1, %2 : i32
-
-  "terminator"(%5) : (i32) -> ()
-}) : () -> ()
-
-// CHECK-LABEL: "test.acrossBlockCanonicalizationBarrierIdem"
-// CHECK: ^bb1:
-// CHECK-NEXT: %[[OUT1:.+]] = comb.or %0, %1 : i32
-// CHECK-NEXT: %[[OUT2:.+]] = comb.or %[[OUT1]], %arg0 : i32
-// CHECK-NEXT: "terminator"(%[[OUT1]], %[[OUT2]]) : (i32, i32) -> ()
-"test.acrossBlockCanonicalizationBarrierIdem"() ({
-^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
-  %0 = comb.or %arg0, %arg1 : i32
-  %1 = comb.or %arg1, %arg2 : i32
-  "terminator"() : () -> ()
-^bb1:  // no predecessors
-  %2 = comb.or %0, %1, %1 : i32
-  %3 = comb.or %2, %0, %1, %arg0 : i32
-
-  "terminator"(%2, %3) : (i32, i32) -> ()
-}) : () -> ()
+    %3 = comb.or %2, %0, %1, %arg0 : i32
+    scf.yield %2, %3 : i32, i32
+  }
+  hw.output %res#0, %res#1 : i32, i32
+}
 
 // Check multi-operation idempotent operand deduplication.
 // CHECK-LABEL: @IdemTwoState