[InferRW] Remove dependence of write-mode on enable for memory

lib/Dialect/FIRRTL/Transforms/InferReadWrite.cpp

@@ -45,6 +45,7 @@ struct InferReadWritePass : public InferReadWriteBase<InferReadWritePass> {
     for (MemOp memOp : llvm::make_early_inc_range(
              getOperation().getBodyBlock()->getOps<MemOp>())) {
       inferUnmasked(memOp, opsToErase);
+      simplifyWmode(memOp);
       size_t nReads, nWrites, nRWs, nDbgs;
       memOp.getNumPorts(nReads, nWrites, nRWs, nDbgs);
       // Run the analysis only for Seq memories (latency=1) and a single read
@@ -209,6 +210,7 @@ struct InferReadWritePass : public InferReadWriteBase<InferReadWritePass> {
             opsToErase.push_back(sf);
           }
       }
+      simplifyWmode(rwMem);
       // All uses for all results of mem removed, now erase the memOp.
       opsToErase.push_back(memOp);
     }
@@ -312,6 +314,73 @@ struct InferReadWritePass : public InferReadWriteBase<InferReadWritePass> {
     return {};
   }
 
+  // Remove redundant dependence of wmode on the enable signal. wmode can assume
+  // the enable signal be true.
+  void simplifyWmode(MemOp &memOp) {
+
+    // Iterate over all results, and find the enable and wmode fields of the RW
+    // port.
+    for (const auto &portIt : llvm::enumerate(memOp.getResults())) {
+      auto portKind = memOp.getPortKind(portIt.index());
+      if (portKind != MemOp::PortKind::ReadWrite)
+        continue;
+      Value enableDriver, wmodeDriver;
+      Value portVal = portIt.value();
+      // Iterate over all users of the rw port.
+      for (Operation *u : portVal.getUsers())
+        if (auto sf = dyn_cast<SubfieldOp>(u)) {
+          // Get the field name.
+          auto fName =
+              sf.getInput().getType().base().getElementName(sf.getFieldIndex());
+          // Record the enable and wmode fields.
+          if (fName.contains("en"))
+            enableDriver = getConnectSrc(sf.getResult());
+          if (fName.contains("wmode"))
+            wmodeDriver = getConnectSrc(sf.getResult());
+        }
+
+      if (enableDriver && wmodeDriver) {
+        ImplicitLocOpBuilder builder(memOp.getLoc(), memOp);
+        auto constOne = builder.create<ConstantOp>(
+            UIntType::get(builder.getContext(), 1), APInt(1, 1));
+        setEnable(enableDriver, wmodeDriver, constOne);
+      }
+    }
+  }
+
+  // Replace any occurence of enable on the expression tree of wmode with a
+  // constant one.
+  void setEnable(Value enableDriver, Value wmodeDriver, Value constOne) {
+    auto getDriverOp = [&](Value dst) -> Operation * {
+      // Look through one level of wire to get the driver op.
+      auto *defOp = dst.getDefiningOp();
+      if (defOp) {
+        if (isa<WireOp>(defOp))
+          dst = getConnectSrc(dst);
+        if (dst)
+          defOp = dst.getDefiningOp();
+      }
+      return defOp;
+    };
+    SmallVector<Value> stack;
+    llvm::SmallDenseSet<Value> visited;
+    stack.push_back(wmodeDriver);
+    while (!stack.empty()) {
+      auto driver = stack.pop_back_val();
+      if (!visited.insert(driver).second)
+        continue;
+      auto *defOp = getDriverOp(driver);
+      if (!defOp)
+        continue;
+      for (auto operand : llvm::enumerate(defOp->getOperands())) {
+        if (operand.value() == enableDriver)
+          defOp->setOperand(operand.index(), constOne);
+        else
+          stack.push_back(operand.value());
+      }
+    }
+  }
+
   void inferUnmasked(MemOp &memOp, SmallVector<Operation *> &opsToErase) {
     bool isMasked = true;
 

test/Dialect/FIRRTL/inferRW.mlir

@@ -282,5 +282,25 @@ firrtl.circuit "TLRAM" {
       firrtl.connect %auto_0, %11 : !firrtl.uint<8>, !firrtl.uint<8>
 
     }
+
+    // CHECK-LABEL: firrtl.module @SimplifyWMODE
+    firrtl.module @SimplifyWMODE(in %rwPort_enable: !firrtl.uint<1>, in %rwPort_isWrite: !firrtl.uint<1>) attributes {} {
+      %mem_rwPort_readData_rw = firrtl.mem Undefined {depth = 64 : i64, name = "t", portNames = ["rw"], prefix = "", readLatency = 1 : i32, writeLatency = 1 : i32} : !firrtl.bundle<addr: uint<6>, en: uint<1>, clk: clock, rdata flip: uint<10>, wmode: uint<1>, wdata: uint<10>, wmask: uint<5>>
+      %mem_rwPort_readData_rw_wmode = firrtl.wire : !firrtl.uint<1>
+      %0 = firrtl.subfield %mem_rwPort_readData_rw[addr] : !firrtl.bundle<addr: uint<6>, en: uint<1>, clk: clock, rdata flip: uint<10>, wmode: uint<1>, wdata: uint<10>, wmask: uint<5>>
+      %1 = firrtl.subfield %mem_rwPort_readData_rw[en] : !firrtl.bundle<addr: uint<6>, en: uint<1>, clk: clock, rdata flip: uint<10>, wmode: uint<1>, wdata: uint<10>, wmask: uint<5>>
+      firrtl.strictconnect %1, %rwPort_enable : !firrtl.uint<1>
+      %2 = firrtl.subfield %mem_rwPort_readData_rw[clk] : !firrtl.bundle<addr: uint<6>, en: uint<1>, clk: clock, rdata flip: uint<10>, wmode: uint<1>, wdata: uint<10>, wmask: uint<5>>
+      %3 = firrtl.subfield %mem_rwPort_readData_rw[rdata] : !firrtl.bundle<addr: uint<6>, en: uint<1>, clk: clock, rdata flip: uint<10>, wmode: uint<1>, wdata: uint<10>, wmask: uint<5>>
+      %6 = firrtl.subfield %mem_rwPort_readData_rw[wmode] : !firrtl.bundle<addr: uint<6>, en: uint<1>, clk: clock, rdata flip: uint<10>, wmode: uint<1>, wdata: uint<10>, wmask: uint<5>>
+      firrtl.strictconnect %6, %mem_rwPort_readData_rw_wmode : !firrtl.uint<1>
+      %7 = firrtl.subfield %mem_rwPort_readData_rw[wdata] : !firrtl.bundle<addr: uint<6>, en: uint<1>, clk: clock, rdata flip: uint<10>, wmode: uint<1>, wdata: uint<10>, wmask: uint<5>>
+      %9 = firrtl.subfield %mem_rwPort_readData_rw[wmask] : !firrtl.bundle<addr: uint<6>, en: uint<1>, clk: clock, rdata flip: uint<10>, wmode: uint<1>, wdata: uint<10>, wmask: uint<5>>
+      %c0_ui1 = firrtl.constant 0 : !firrtl.uint<1> 
+      %18 = firrtl.mux(%rwPort_enable, %rwPort_isWrite, %c0_ui1) : (!firrtl.uint<1>, !firrtl.uint<1>, !firrtl.uint<1>) -> !firrtl.uint<1>
+      // CHECK: %[[c1_ui1:.+]] = firrtl.constant 1 : !firrtl.uint<1>
+      // CHECK: %[[v7:.+]] = firrtl.mux(%[[c1_ui1]], %rwPort_isWrite, %c0_ui1)
+      firrtl.strictconnect %mem_rwPort_readData_rw_wmode, %18 : !firrtl.uint<1>
+    }
 }