From 759b2b09dd0ed20d3da2f6579baf231eeb75e8bd Mon Sep 17 00:00:00 2001
From: fzi-hielscher <47524191+fzi-hielscher@users.noreply.github.com>
Date: Thu, 17 Oct 2024 21:54:45 +0200
Subject: [PATCH] [CombFolds] Preserve two-state attribute in
 `narrowOperationWidth` (#7712)

We're already doing a best-effort preservation of discardable attributes through `narrowOperationWidth`. I cannot think of a reason why the inherent two-state attribute should not be carried over.
---
 lib/Dialect/Comb/CombFolds.cpp          | 18 +++++++++++-------
 test/Dialect/Comb/canonicalization.mlir | 14 ++++++++++++++
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/lib/Dialect/Comb/CombFolds.cpp b/lib/Dialect/Comb/CombFolds.cpp
index ececa63a30f7..92aa494dc789 100644
--- a/lib/Dialect/Comb/CombFolds.cpp
+++ b/lib/Dialect/Comb/CombFolds.cpp
@@ -271,20 +271,24 @@ static bool narrowOperationWidth(OpTy op, bool narrowTrailingBits,
       args.push_back(rewriter.createOrFold<ExtractOp>(inop.getLoc(), newType,
                                                       inop, range.first));
   }
-  Value newop = rewriter.createOrFold<OpTy>(op.getLoc(), newType, args);
-  newop.getDefiningOp()->setDialectAttrs(op->getDialectAttrs());
+  auto newop = rewriter.create<OpTy>(op.getLoc(), newType, args);
+  newop->setDialectAttrs(op->getDialectAttrs());
+  if (op.getTwoState())
+    newop.setTwoState(true);
+
+  Value newResult = newop.getResult();
   if (range.first)
-    newop = rewriter.createOrFold<ConcatOp>(
-        op.getLoc(), newop,
+    newResult = rewriter.createOrFold<ConcatOp>(
+        op.getLoc(), newResult,
         rewriter.create<hw::ConstantOp>(op.getLoc(),
                                         APInt::getZero(range.first)));
   if (range.second + 1 < opType.getWidth())
-    newop = rewriter.createOrFold<ConcatOp>(
+    newResult = rewriter.createOrFold<ConcatOp>(
         op.getLoc(),
         rewriter.create<hw::ConstantOp>(
             op.getLoc(), APInt::getZero(opType.getWidth() - range.second - 1)),
-        newop);
-  rewriter.replaceOp(op, newop);
+        newResult);
+  rewriter.replaceOp(op, newResult);
   return true;
 }
 
diff --git a/test/Dialect/Comb/canonicalization.mlir b/test/Dialect/Comb/canonicalization.mlir
index 77d0eadbc9ab..8c1a3f4345dc 100644
--- a/test/Dialect/Comb/canonicalization.mlir
+++ b/test/Dialect/Comb/canonicalization.mlir
@@ -569,6 +569,20 @@ hw.module @narrowAdditionToDirectAddition(in %x: i8, in %y: i8, out z1: i8) {
   hw.output %3 : i8
 }
 
+// Validates that narrowing preserves the two-state attribute.
+// CHECK-LABEL: hw.module @narrowAdditionToDirectAdditionTwoState
+hw.module @narrowAdditionToDirectAdditionTwoState(in %x: i8, in %y: i8, out z1: i8) {
+  // CHECK-NEXT: [[RESULT:%.+]] = comb.add bin %x, %y : i8
+  // CHECK-NEXT: hw.output [[RESULT]]
+
+  %false = hw.constant false
+  %0 = comb.concat %x, %x : i8, i8
+  %1 = comb.concat %y, %y : i8, i8
+  %2 = comb.add bin %0, %1 : i16
+  %3 = comb.extract %2 from 0 : (i16) -> i8
+  hw.output %3 : i8
+}
+
 // Validates that addition narrow to the widest extract
 // CHECK-LABEL: hw.module @narrowAdditionToWidestExtract
 hw.module @narrowAdditionToWidestExtract(in %x: i8, in %y: i8, out z1: i3, out z2: i4) {