[Seq] Canonicalize firreg with a preset (#7350)

The canonicalization was ignoring any registers with a preset, this commit fixes that.
Add fold and canonicalization patterns for registers with a preset.

include/circt/Dialect/Seq/SeqOps.td

@@ -235,7 +235,8 @@ def FirRegOp : SeqOp<"firreg",
   let builders = [
     OpBuilder<(ins "Value":$next, "Value":$clk,
                    "StringAttr":$name,
-                   CArg<"hw::InnerSymAttr", "{}">:$innerSym)>,
+                   CArg<"hw::InnerSymAttr", "{}">:$innerSym,
+                   CArg<"Attribute","{}">:$preset)>,
     OpBuilder<(ins "Value":$next, "Value":$clk,
                    "StringAttr":$name,
                    "Value":$reset, "Value":$resetValue,

lib/Dialect/Seq/SeqOps.cpp

@@ -359,7 +359,8 @@ LogicalResult ShiftRegOp::verify() {
 //===----------------------------------------------------------------------===//
 
 void FirRegOp::build(OpBuilder &builder, OperationState &result, Value input,
-                     Value clk, StringAttr name, hw::InnerSymAttr innerSym) {
+                     Value clk, StringAttr name, hw::InnerSymAttr innerSym,
+                     Attribute preset) {
 
   OpBuilder::InsertionGuard guard(builder);
 
@@ -371,6 +372,9 @@ void FirRegOp::build(OpBuilder &builder, OperationState &result, Value input,
   if (innerSym)
     result.addAttribute(getInnerSymAttrName(result.name), innerSym);
 
+  if (preset)
+    result.addAttribute(getPresetAttrName(result.name), preset);
+
   result.addTypes(input.getType());
 }
 
@@ -545,19 +549,15 @@ void FirRegOp::getAsmResultNames(OpAsmSetValueNameFn setNameFn) {
 std::optional<size_t> FirRegOp::getTargetResultIndex() { return 0; }
 
 LogicalResult FirRegOp::canonicalize(FirRegOp op, PatternRewriter &rewriter) {
-  // Don't canonicalize if there is a preset value for now.
-  // TODO: Handle a preset value.
-  if (op.getPresetAttr())
-    return failure();
 
   // If the register has a constant zero reset, drop the reset and reset value
-  // altogether.
+  // altogether (And preserve the PresetAttr).
   if (auto reset = op.getReset()) {
     if (auto constOp = reset.getDefiningOp<hw::ConstantOp>()) {
       if (constOp.getValue().isZero()) {
-        rewriter.replaceOpWithNewOp<FirRegOp>(op, op.getNext(), op.getClk(),
+        rewriter.replaceOpWithNewOp<FirRegOp>(
-                                              op.getNameAttr(),
+            op, op.getNext(), op.getClk(), op.getNameAttr(),
-                                              op.getInnerSymAttr());
+            op.getInnerSymAttr(), op.getPresetAttr());
         return success();
       }
     }
@@ -579,7 +579,13 @@ LogicalResult FirRegOp::canonicalize(FirRegOp op, PatternRewriter &rewriter) {
     return false;
   };
 
-  if (isConstant() && !op.getResetValue()) {
+  // Preset can block canonicalization only if it is non-zero.
+  bool replaceWithConstZero = true;
+  if (auto preset = op.getPresetAttr())
+    if (!preset.getValue().isZero())
+      replaceWithConstZero = false;
+
+  if (isConstant() && !op.getResetValue() && replaceWithConstZero) {
     if (isa<seq::ClockType>(op.getType())) {
       rewriter.replaceOpWithNewOp<seq::ConstClockOp>(
           op, seq::ClockConstAttr::get(rewriter.getContext(), ClockConst::Low));
@@ -597,7 +603,7 @@ LogicalResult FirRegOp::canonicalize(FirRegOp op, PatternRewriter &rewriter) {
   // initialized. If we don't enable aggregate preservation, `r_0` is replaced
   // with `0`. Hence this canonicalization replaces 0th element of the next
   // value with zero to match the behaviour.
-  if (!op.getReset()) {
+  if (!op.getReset() && !op.getPresetAttr()) {
     if (auto arrayCreate = op.getNext().getDefiningOp<hw::ArrayCreateOp>()) {
       // For now only support 1d arrays.
       // TODO: Support nested arrays and bundles.
@@ -651,20 +657,23 @@ LogicalResult FirRegOp::canonicalize(FirRegOp op, PatternRewriter &rewriter) {
 OpFoldResult FirRegOp::fold(FoldAdaptor adaptor) {
   // If the register has a symbol or preset value, we can't optimize it away.
   // TODO: Handle a preset value.
-  if (getInnerSymAttr() || getPresetAttr())
+  if (getInnerSymAttr())
     return {};
 
+  auto presetAttr = getPresetAttr();
+
   // If the register is held in permanent reset, replace it with its reset
   // value. This works trivially if the reset is asynchronous and therefore
   // level-sensitive, in which case it will always immediately assume the reset
   // value in silicon. If it is synchronous, the register value is undefined
   // until the first clock edge at which point it becomes the reset value, in
   // which case we simply define the initial value to already be the reset
-  // value.
+  // value. Works only if no preset.
-  if (auto reset = getReset())
+  if (!presetAttr)
-    if (auto constOp = reset.getDefiningOp<hw::ConstantOp>())
+    if (auto reset = getReset())
-      if (constOp.getValue().isOne())
+      if (auto constOp = reset.getDefiningOp<hw::ConstantOp>())
-        return getResetValue();
+        if (constOp.getValue().isOne())
+          return getResetValue();
 
   // If the register's next value is trivially it's current value, or the
   // register is never clocked, we can replace the register with a constant
@@ -675,12 +684,16 @@ OpFoldResult FirRegOp::fold(FoldAdaptor adaptor) {
   if (!isTrivialFeedback && !isNeverClocked)
     return {};
 
-  // If the register has a const reset value, we can replace it with that.
+  // If the register has a const reset value, and no preset, we can replace it
-  // We cannot replace it with a non-constant reset value.
+  // with the const reset. We cannot replace it with a non-constant reset value.
   if (auto resetValue = getResetValue()) {
-    if (auto *op = resetValue.getDefiningOp())
+    if (auto *op = resetValue.getDefiningOp()) {
-      if (op->hasTrait<OpTrait::ConstantLike>())
+      if (op->hasTrait<OpTrait::ConstantLike>() && !presetAttr)
         return resetValue;
+      if (auto constOp = dyn_cast<hw::ConstantOp>(op))
+        if (presetAttr.getValue() == constOp.getValue())
+          return resetValue;
+    }
     return {};
   }
 
@@ -689,6 +702,9 @@ OpFoldResult FirRegOp::fold(FoldAdaptor adaptor) {
   auto intType = dyn_cast<IntegerType>(getType());
   if (!intType)
     return {};
+  // If preset present, then replace with preset.
+  if (presetAttr)
+    return presetAttr;
   return IntegerAttr::get(intType, 0);
 }
 

test/Dialect/Seq/canonicalization.mlir

@@ -35,6 +35,7 @@ hw.module @FirRegSymbol(in %clk : !seq.clock, out out : i32) {
 
 // CHECK-LABEL: @FirRegReset
 hw.module @FirRegReset(in %clk : !seq.clock, in %in : i32, in %r : i1, in %v : i32) {
+  %c3_i32 = hw.constant 3 : i32
   %false = hw.constant false
   %true = hw.constant true
 
@@ -79,6 +80,25 @@ hw.module @FirRegReset(in %clk : !seq.clock, in %in : i32, in %r : i1, in %v : i
   // A register with preset value is not folded right now
   // CHECK: %reg_preset = seq.firreg
   %reg_preset = seq.firreg %reg_preset clock %clk reset sync %r, %c0_i32 preset 3: i32
+
+  // A register with 0 preset value is folded.
+  %reg_preset_0 = seq.firreg %reg_preset_0 clock %clk preset 0: i32
+  hw.instance "reg_preset_0" @Observe(x: %reg_preset_0: i32) -> ()
+  // CHECK-NEXT: hw.instance "reg_preset_0" @Observe(x: %c0_i32: i32) -> ()
+
+  // A register with const false reset and 0 preset value is folded.
+  %reg_preset_1 = seq.firreg %reg_preset_1 clock %clk reset sync %false, %c0_i32 preset 0: i32
+  // CHECK-NEXT: hw.instance "reg_preset_1" @Observe(x: %c0_i32: i32) -> ()
+  hw.instance "reg_preset_1" @Observe(x: %reg_preset_1: i32) -> ()
+
+  // A register with const false reset and 0 preset value is folded.
+  %reg_preset_2 = seq.firreg %reg_preset_2 clock %clk reset sync %false, %c0_i32 preset 3: i32
+  // CHECK-NEXT: hw.instance "reg_preset_2" @Observe(x: %c3_i32: i32) -> ()
+  hw.instance "reg_preset_2" @Observe(x: %reg_preset_2: i32) -> ()
+
+  %reg_preset_3 = seq.firreg %reg_preset_3 clock %clk reset sync %r, %c3_i32 preset 3: i32
+  // CHECK-NEXT: hw.instance "reg_preset_3" @Observe(x: %c3_i32: i32) -> ()
+  hw.instance "reg_preset_3" @Observe(x: %reg_preset_3: i32) -> ()
 }
 
 // CHECK-LABEL: @FirRegAggregate