[Comb] Add canonicalization for comb.truth_table

Author: MrRoy09

include/circt/Dialect/Comb/Combinational.td

@@ -330,6 +330,7 @@ def TruthTableOp : CombOp<"truth_table", [Pure]> {
   }];
 
   let hasVerifier = 1;
+  let hasCanonicalizeMethod = 1;
 }
 
 def ReverseOp : CombOp<"reverse", [

lib/Dialect/Comb/CombFolds.cpp

@@ -3291,3 +3291,89 @@ LogicalResult ICmpOp::canonicalize(ICmpOp op, PatternRewriter &rewriter) {
 
   return failure();
 }
+
+//===----------------------------------------------------------------------===//
+// TruthTableOp
+//===----------------------------------------------------------------------===//
+
+// Canonicalize truth tables by folding to constant when output does not depend
+// on any of the inputs or simplifying when the output depends on only a single
+// input.
+LogicalResult TruthTableOp::canonicalize(TruthTableOp op,
+                                         PatternRewriter &rewriter) {
+  if (isOpTriviallyRecursive(op))
+    return failure();
+
+  const auto inputs = op.getInputs();
+  const auto table = op.getLookupTable();
+  size_t numInputs = inputs.size();
+  size_t tableSize = table.size();
+
+  if (numInputs == 0 || tableSize == 0 || tableSize != (1ull << numInputs))
+    return failure();
+
+  // Check if the table can be folded to just a constant.
+  bool firstValue = cast<BoolAttr>(table[0]).getValue();
+  bool allSame = llvm::all_of(table, [firstValue](Attribute attr) {
+    return cast<BoolAttr>(attr).getValue() == firstValue;
+  });
+  if (allSame) {
+    auto constOp =
+        hw::ConstantOp::create(rewriter, op.getLoc(), APInt(1, firstValue));
+    replaceOpAndCopyNamehint(rewriter, op, constOp);
+    return success();
+  }
+
+  // Detect if the truth table depends only on one of the inputs.
+  // For each input bit, we test whether flipping only that input bit changes
+  // the output value of the truth table at any point. For this, iterate over
+  // all table entries and compute the index of the entry where just that input
+  // bit is inverted.
+  SmallVector<bool> dependsOn(numInputs, false);
+  int dependentInput = -1;
+  unsigned numDependencies = 0;
+
+  for (size_t idx = 0; idx < tableSize; ++idx) {
+    bool currentValue = cast<BoolAttr>(table[idx]).getValue();
+
+    for (size_t bitPos = 0; bitPos < numInputs; ++bitPos) {
+      // Skip if we already know this input matters.
+      if (dependsOn[bitPos])
+        continue;
+
+      // Calculate the index of the entry with the bit in question flipped.
+      size_t bitPositionInTable = numInputs - 1 - bitPos;
+      size_t flippedIdx = idx ^ (1ull << bitPositionInTable);
+      bool flippedValue = cast<BoolAttr>(table[flippedIdx]).getValue();
+
+      // If flipping this bit changes the output, this input is a dependency.
+      if (currentValue != flippedValue) {
+        dependsOn[bitPos] = true;
+        dependentInput = bitPos;
+        numDependencies++;
+
+        // Exit early if we already found more than one dependency.
+        if (numDependencies > 1)
+          break;
+      }
+    }
+
+    // Exit early from outer loop if we found more than one dependency.
+    if (numDependencies > 1)
+      break;
+  }
+
+  if (numDependencies != 1)
+    return failure();
+
+  // Determine if the truth table is identity or inverted by checking the output
+  // when the dependent input is 1 (all other inputs at 0).
+  size_t idxWhen1 = 1ull << (numInputs - 1 - dependentInput);
+  bool isIdentity = cast<BoolAttr>(table[idxWhen1]).getValue();
+
+  // Replace with the input or its negation.
+  Value input = inputs[dependentInput];
+  Value replacement = isIdentity ? input : createOrFoldNot(op.getLoc(), input, rewriter);
+  replaceOpAndCopyNamehint(rewriter, op, replacement);
+  return success();
+}

test/Dialect/Comb/canonicalization.mlir

@@ -1960,3 +1960,103 @@ hw.module @issue9403(in %sel: i1, out out1: ui1) {
   %mux1 = comb.mux %sel, %true, %false : ui1
   hw.output %mux1 : ui1
 }
+
+
+// CHECK-LABEL: @truth_table_constant_true
+hw.module @truth_table_constant_true(in %a: i1, in %b: i1, out out: i1) {
+  // Truth table that is always true (all ones)
+  // CHECK-NEXT: [[TRUE:%.+]] = hw.constant true
+  // CHECK-NEXT: hw.output [[TRUE]]
+  %0 = comb.truth_table %a, %b -> [true, true, true, true]
+  hw.output %0 : i1
+}
+
+// CHECK-LABEL: @truth_table_constant_false
+hw.module @truth_table_constant_false(in %a: i1, in %b: i1, out out: i1) {
+  // Truth table that is always false (all zeros)
+  // CHECK-NEXT: [[FALSE:%.+]] = hw.constant false
+  // CHECK-NEXT: hw.output [[FALSE]]
+  %0 = comb.truth_table %a, %b -> [false, false, false, false]
+  hw.output %0 : i1
+}
+
+// CHECK-LABEL: @truth_table_identity
+hw.module @truth_table_identity(in %a: i1, in %b: i1, in %c: i1, out out: i1) {
+  // Truth table that depends only on %a 
+  // Pattern: [0,0,0,0,1,1,1,1] means output follows first input
+  // CHECK-NEXT: hw.output %a
+  %0 = comb.truth_table %a, %b, %c -> [false, false, false, false, true, true, true, true]
+  hw.output %0 : i1
+}
+
+// CHECK-LABEL: @truth_table_inverted
+hw.module @truth_table_inverted(in %a: i1, in %b: i1, in %c: i1, out out: i1) {
+  // Truth table that depends only on %a (inverted)
+  // Pattern: [1,1,1,1,0,0,0,0] means output is NOT of first input
+  // CHECK-NEXT: %true = hw.constant true
+  // CHECK-NEXT: [[NOT:%.+]] = comb.xor %a, %true
+  // CHECK-NEXT: hw.output [[NOT]]
+  %0 = comb.truth_table %a, %b, %c -> [true, true, true, true, false, false, false, false]
+  hw.output %0 : i1
+}
+
+// CHECK-LABEL: @truth_table_middle_input_identity
+hw.module @truth_table_middle_input_identity(in %a: i1, in %b: i1, in %c: i1, out out: i1) {
+  // Truth table that depends only on %b (middle input, identity)
+  // Pattern: [0,0,1,1,0,0,1,1] means output follows second input
+  // CHECK-NEXT: hw.output %b
+  %0 = comb.truth_table %a, %b, %c -> [false, false, true, true, false, false, true, true]
+  hw.output %0 : i1
+}
+
+// CHECK-LABEL: @truth_table_middle_input_inverted
+hw.module @truth_table_middle_input_inverted(in %a: i1, in %b: i1, in %c: i1, out out: i1) {
+  // Truth table that depends only on %b (middle input, inverted)
+  // Pattern: [1,1,0,0,1,1,0,0] means output is NOT of second input
+  // CHECK-NEXT: %true = hw.constant true
+  // CHECK-NEXT: [[NOT:%.+]] = comb.xor %b, %true
+  // CHECK-NEXT: hw.output [[NOT]]
+  %0 = comb.truth_table %a, %b, %c -> [true, true, false, false, true, true, false, false]
+  hw.output %0 : i1
+}
+
+// CHECK-LABEL: @truth_table_last_input_identity
+hw.module @truth_table_last_input_identity(in %a: i1, in %b: i1, in %c: i1, out out: i1) {
+  // Truth table that depends only on %c (last input, identity)
+  // Pattern: [0,1,0,1,0,1,0,1] means output follows third input
+  // CHECK-NEXT: hw.output %c
+  %0 = comb.truth_table %a, %b, %c -> [false, true, false, true, false, true, false, true]
+  hw.output %0 : i1
+}
+
+// CHECK-LABEL: @truth_table_last_input_inverted
+hw.module @truth_table_last_input_inverted(in %a: i1, in %b: i1, in %c: i1, out out: i1) {
+  // Truth table that depends only on %c (last input, inverted)
+  // Pattern: [1,0,1,0,1,0,1,0] means output is NOT of third input
+  // CHECK-NEXT: %true = hw.constant true
+  // CHECK-NEXT: [[NOT:%.+]] = comb.xor %c, %true
+  // CHECK-NEXT: hw.output [[NOT]]
+  %0 = comb.truth_table %a, %b, %c -> [true, false, true, false, true, false, true, false]
+  hw.output %0 : i1
+}
+
+// CHECK-LABEL: @truth_table_two_input_non_foldable
+hw.module @truth_table_two_input_non_foldable(in %a: i1, in %b: i1, out out: i1) {
+  // Truth table depends on both inputs, Should not be canonicalized
+  // CHECK-NEXT: %0 = comb.truth_table %a, %b -> [false, false, false, true]
+  // CHECK-NEXT: hw.output %0
+  %0 = comb.truth_table %a, %b -> [false, false, false, true]
+  hw.output %0 : i1
+}
+
+// CHECK-LABEL: @truth_table_with_extract_operations
+hw.module @truth_table_with_extract_operations(in %c: i3, out out: i1) {
+  // Truth table depends only on first input (%2 = LSB of %c)
+  // CHECK: [[TMP:%.+]] = comb.extract %c from 0
+  // CHECK: hw.output [[TMP]]
+  %0 = comb.extract %c from 2 : (i3) -> i1
+  %1 = comb.extract %c from 1 : (i3) -> i1
+  %2 = comb.extract %c from 0 : (i3) -> i1
+  %3 = comb.truth_table %2, %0, %1 -> [false, false, false, false, true, true, true, true]
+  hw.output %3 : i1
+}