From 2caafc76bb1f35a92868c2c77fed114f40b9ffd7 Mon Sep 17 00:00:00 2001
From: Prithayan Barua <prithayan@gmail.com>
Date: Mon, 18 Dec 2023 13:42:46 -0500
Subject: [PATCH] [LowerToHW] Fix output port index mapping (#6530)

The logic that was moving the symbol from the port to the wire was
 incorrectly assuming that all the output ports are ordered after the input
  ports, this resulted in the port index mapping being incorrect.
This was causing the incorrect port symbol being moved to a temporary wire.
This commit fixes the port index calculation logic.
---
 lib/Conversion/FIRRTLToHW/LowerToHW.cpp     | 22 +++++++++++----------
 test/Conversion/FIRRTLToHW/lower-to-hw.mlir | 19 ++++++++++++++++++
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/lib/Conversion/FIRRTLToHW/LowerToHW.cpp b/lib/Conversion/FIRRTLToHW/LowerToHW.cpp
index dd1f1eb536ad..32ccdb6a272d 100644
--- a/lib/Conversion/FIRRTLToHW/LowerToHW.cpp
+++ b/lib/Conversion/FIRRTLToHW/LowerToHW.cpp
@@ -1232,30 +1232,33 @@ LogicalResult FIRRTLModuleLowering::lowerModulePortsAndMoveBody(
   bodyBuilder.setInsertionPoint(cursor);
 
   // Insert argument casts, and re-vector users in the old body to use them.
-  SmallVector<PortInfo> ports = oldModule.getPorts();
-  assert(oldModule.getBody().getNumArguments() == ports.size() &&
+  SmallVector<PortInfo> firrtlPorts = oldModule.getPorts();
+  SmallVector<hw::PortInfo> hwPorts = newModule.getPortList();
+  assert(oldModule.getBody().getNumArguments() == firrtlPorts.size() &&
          "port count mismatch");
 
-  size_t nextNewArg = 0;
-  size_t firrtlArg = 0;
   SmallVector<Value, 4> outputs;
 
   // This is the terminator in the new module.
   auto outputOp = newModule.getBodyBlock()->getTerminator();
   ImplicitLocOpBuilder outputBuilder(oldModule.getLoc(), outputOp);
 
-  for (auto &port : ports) {
+  unsigned nextHWInputArg = 0;
+  int hwPortIndex = -1;
+  for (auto [firrtlPortIndex, port] : llvm::enumerate(firrtlPorts)) {
     // Inputs and outputs are both modeled as arguments in the FIRRTL level.
-    auto oldArg = oldModule.getBody().getArgument(firrtlArg++);
+    auto oldArg = oldModule.getBody().getArgument(firrtlPortIndex);
 
     bool isZeroWidth =
         type_isa<FIRRTLBaseType>(port.type) &&
         type_cast<FIRRTLBaseType>(port.type).getBitWidthOrSentinel() == 0;
+    if (!isZeroWidth)
+      ++hwPortIndex;
 
     if (!port.isOutput() && !isZeroWidth) {
       // Inputs and InOuts are modeled as arguments in the result, so we can
       // just map them over.  We model zero bit outputs as inouts.
-      Value newArg = newModule.getBody().getArgument(nextNewArg++);
+      Value newArg = newModule.getBody().getArgument(nextHWInputArg++);
 
       // Cast the argument to the old type, reintroducing sign information in
       // the hw.module body.
@@ -1298,13 +1301,12 @@ LogicalResult FIRRTLModuleLowering::lowerModulePortsAndMoveBody(
     if (!resultHWType.isInteger(0)) {
       auto output =
           castFromFIRRTLType(newArg.getResult(), resultHWType, outputBuilder);
-      auto idx = newModule.getNumInputPorts() + outputs.size();
       outputs.push_back(output);
 
       // If output port has symbol, move it to this wire.
-      if (auto sym = newModule.getPortList()[idx].getSym()) {
+      if (auto sym = hwPorts[hwPortIndex].getSym()) {
         newArg.setInnerSymAttr(sym);
-        newModule.setPortSymbolAttr(idx, {});
+        newModule.setPortSymbolAttr(hwPortIndex, {});
       }
     }
   }
diff --git a/test/Conversion/FIRRTLToHW/lower-to-hw.mlir b/test/Conversion/FIRRTLToHW/lower-to-hw.mlir
index f5554a29f018..5ec14b4c08a0 100644
--- a/test/Conversion/FIRRTLToHW/lower-to-hw.mlir
+++ b/test/Conversion/FIRRTLToHW/lower-to-hw.mlir
@@ -1567,3 +1567,22 @@ firrtl.circuit "ZeroWidthForeignOperand" {
     dbg.variable "v2", %a : !firrtl.uint<0>
   }
 }
+
+// -----
+
+// Check correct symbol mapping from output port to wire happens.
+firrtl.circuit "PortSym" {
+  firrtl.extmodule private @Blackbox(out bar: !firrtl.uint<1>)
+  // CHECK-LABEL: module @PortSym(
+  // CHECK-SAME: out a : i1 {hw.exportPort = #hw<innerSym@out_a_m>}
+  // CHECK-SAME: out out : i5, in %c : i1)
+  firrtl.module @PortSym(out %a: !firrtl.uint<1> sym @out_a_m, in %b: !firrtl.uint<0>, out %out: !firrtl.uint<5> sym @out_sym, in %c: !firrtl.uint<1>) attributes {convention = #firrtl<convention scalarized>} {
+    // CHECK: %[[OUT:.+]] = hw.wire %{{.+}} sym @out_sym
+    %c1_ui1 = firrtl.constant 1 : !firrtl.uint<1>
+    %c1_ui5 = firrtl.constant 1 : !firrtl.uint<5>
+    firrtl.strictconnect %out, %c1_ui5 : !firrtl.uint<5>
+    %e_a = firrtl.instance sub1 @Blackbox(out bar: !firrtl.uint<1>)
+    firrtl.strictconnect %a, %e_a : !firrtl.uint<1>
+    %0 = firrtl.eq %out, %c1_ui5 : (!firrtl.uint<5>, !firrtl.uint<5>) -> !firrtl.uint<1>
+  }
+}