[CombFolds] Preserve two-state attribute in narrowOperationWidth (l…

…lvm#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

@@ -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;
 }
 

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) {