[FIRRTL] Subword assignment support via rewriting to read-modify-write

lib/Dialect/FIRRTL/FIRRTLFolds.cpp

@@ -1096,14 +1096,51 @@ OpFoldResult BitsPrimOp::fold(ArrayRef<Attribute> operands) {
 LogicalResult BitsPrimOp::canonicalize(BitsPrimOp op,
                                        PatternRewriter &rewriter) {
   auto *inputOp = op.getInput().getDefiningOp();
+  if (!inputOp)
+    return failure();
   // bits(bits(x, ...), ...) -> bits(x, ...).
-  if (auto innerBits = dyn_cast_or_null<BitsPrimOp>(inputOp)) {
+  if (auto innerBits = dyn_cast<BitsPrimOp>(inputOp)) {
     auto newLo = op.getLo() + innerBits.getLo();
     auto newHi = newLo + op.getHi() - op.getLo();
     replaceOpWithNewOpAndCopyName<BitsPrimOp>(
         rewriter, op, innerBits.getInput(), newHi, newLo);
     return success();
   }
+  // bits(mux(sel, a, b), ...) -> mux(sel, bits(a, ...), bits(b, ...)).
+  if (auto mux = dyn_cast<MuxPrimOp>(inputOp)) {
+    auto bitsHigh = rewriter.create<BitsPrimOp>(op->getLoc(), mux.getHigh(),
+                                                op.getHi(), op.getLo());
+    auto bitsLow = rewriter.create<BitsPrimOp>(op->getLoc(), mux.getLow(),
+                                               op.getHi(), op.getLo());
+    replaceOpWithNewOpAndCopyName<MuxPrimOp>(rewriter, op, mux.getSel(),
+                                             bitsHigh, bitsLow);
+    return success();
+  }
+  // bits(cat(a, b), ...) -> bits(a, ...), or bits(b, ...), or cat(bits(a, ...),
+  // bits(b, ...)).
+  if (auto cat = dyn_cast<CatPrimOp>(inputOp)) {
+    auto rhsT = cat.getRhs().getType().cast<IntType>();
+    if (!rhsT.hasWidth())
+      return failure();
+    uint32_t lhsLo = rhsT.getWidth().value();
+    uint32_t rhsHi = lhsLo - 1;
+    if (op.getLo() >= lhsLo) {
+      // Only indexing the lhs.
+      replaceOpWithNewOpAndCopyName<BitsPrimOp>(
+          rewriter, op, cat.getLhs(), op.getHi() - lhsLo, op.getLo() - lhsLo);
+    } else if (op.getHi() <= rhsHi) {
+      // Only indexing the rhs.
+      replaceOpWithNewOpAndCopyName<BitsPrimOp>(rewriter, op, cat.getRhs(),
+                                                op.getHi(), op.getLo());
+    } else {
+      auto bitsLhs = rewriter.create<BitsPrimOp>(op->getLoc(), cat.getLhs(),
+                                                 op.getHi() - lhsLo, 0);
+      auto bitsRhs = rewriter.create<BitsPrimOp>(op->getLoc(), cat.getRhs(),
+                                                 rhsHi, op.getLo());
+      replaceOpWithNewOpAndCopyName<CatPrimOp>(rewriter, op, bitsLhs, bitsRhs);
+    }
+    return success();
+  }
   return failure();
 }
 

lib/Dialect/FIRRTL/FIRRTLOps.cpp

@@ -128,7 +128,7 @@ Flow firrtl::foldFlow(Value val, Flow accumulatedFlow) {
         return foldFlow(op.getInput(),
                         op.isFieldFlipped() ? swap() : accumulatedFlow);
       })
-      .Case<SubindexOp, SubaccessOp>(
+      .Case<SubindexOp, SubaccessOp, BitsPrimOp>(
           [&](auto op) { return foldFlow(op.getInput(), accumulatedFlow); })
       // Registers, Wires, and behavioral memory ports are always Duplex.
       .Case<RegOp, RegResetOp, WireOp, MemoryPortOp>(
@@ -153,7 +153,7 @@ DeclKind firrtl::getDeclarationKind(Value val) {
 
   return TypeSwitch<Operation *, DeclKind>(op)
       .Case<InstanceOp>([](auto) { return DeclKind::Instance; })
-      .Case<SubfieldOp, SubindexOp, SubaccessOp>(
+      .Case<SubfieldOp, SubindexOp, SubaccessOp, BitsPrimOp>(
           [](auto op) { return getDeclarationKind(op.getInput()); })
       .Default([](auto) { return DeclKind::Other; });
 }

lib/Dialect/FIRRTL/Import/FIRParser.cpp

@@ -1400,6 +1400,7 @@ ParseResult FIRStmtParser::parseExpImpl(Value &result, const Twine &message,
 ///
 ///  exp ::= exp '.' fieldId
 ///      ::= exp '[' intLit ']'
+///      ::= exp '[' intLit ':' intLit ']'
 /// XX   ::= exp '.' DoubleLit // TODO Workaround for #470
 ///      ::= exp '[' exp ']'
 ParseResult FIRStmtParser::parseOptionalExpPostscript(Value &result) {
@@ -1414,7 +1415,9 @@ ParseResult FIRStmtParser::parseOptionalExpPostscript(Value &result) {
       continue;
     }
 
-    // Subindex: exp ::= exp '[' intLit ']' | exp '[' exp ']'
+    // Subindex: exp ::= exp '[' intLit ']'
+    //               ::= exp '[' intLit ':' intLit ']'
+    //               ::= exp '[' exp ']'
     if (consumeIf(FIRToken::l_square)) {
       if (getToken().isAny(FIRToken::integer, FIRToken::string)) {
         if (parsePostFixIntSubscript(result))
@@ -1481,51 +1484,79 @@ ParseResult FIRStmtParser::parsePostFixFieldId(Value &result) {
   return success();
 }
 
-/// exp ::= exp '[' intLit ']'
+/// exp ::= exp '[' intLit ']' | exp '[' intLit ':' intLit ']'
 ///
 /// The "exp '['" part of the production has already been parsed.
 ///
 ParseResult FIRStmtParser::parsePostFixIntSubscript(Value &result) {
   auto loc = getToken().getLoc();
   int32_t indexNo;
-  if (parseIntLit(indexNo, "expected index") ||
-      parseToken(FIRToken::r_square, "expected ']'"))
+  if (parseIntLit(indexNo, "expected index"))
     return failure();
-
   if (indexNo < 0)
     return emitError(loc, "invalid index specifier"), failure();
 
   // Make sure the index expression is valid and compute the result type for the
   // expression.
   // TODO: This should ideally be folded into a `tryCreate` method on the
   // builder (https://llvm.discourse.group/t/3504).
-  NamedAttribute attrs = {getConstants().indexIdentifier,
-                          builder.getI32IntegerAttr(indexNo)};
-  auto resultType = SubindexOp::inferReturnType({result}, attrs, {});
-  if (!resultType) {
-    // Emit the error at the right location.  translateLocation is expensive.
-    (void)SubindexOp::inferReturnType({result}, attrs, translateLocation(loc));
-    return failure();
-  }
 
-  // Check if we already have created this Subindex op.
-  auto &value = moduleContext.getCachedSubaccess(result, indexNo);
-  if (value) {
-    result = value;
+  if (result.getType().isa<IntType>()) {
+    NamedAttrList bitAttrs;
+    bitAttrs.append(getConstants().hiIdentifier,
+                    builder.getI32IntegerAttr(indexNo));
+    if (consumeIf(FIRToken::colon)) {
+      if (parseIntLit(indexNo, "expected low index"))
+        return failure();
+    }
+    if (parseToken(FIRToken::r_square, "expected ']'"))
+      return failure();
+    bitAttrs.append(getConstants().loIdentifier,
+                    builder.getI32IntegerAttr(indexNo));
+
+    auto resultType = BitsPrimOp::inferReturnType({result}, bitAttrs, {});
+    if (!resultType) {
+      (void)BitsPrimOp::inferReturnType({result}, bitAttrs,
+                                        translateLocation(loc));
+      return failure();
+    }
+    locationProcessor.setLoc(loc);
+    OpBuilder::InsertionGuard guard(builder);
+    builder.setInsertionPointAfterValue(result);
+    auto op = builder.create<BitsPrimOp>(resultType, result, bitAttrs);
+    result = op.getResult();
     return success();
-  }
+  } else {
+    if (parseToken(FIRToken::r_square, "expected ']'"))
+      return failure();
+    NamedAttribute attrs = {getConstants().indexIdentifier,
+                            builder.getI32IntegerAttr(indexNo)};
+    auto resultType = SubindexOp::inferReturnType({result}, attrs, {});
+    if (!resultType) {
+      (void)SubindexOp::inferReturnType({result}, attrs,
+                                        translateLocation(loc));
+      return failure();
+    }
+    // Check if we already have created this Subindex op.
+    auto &value = moduleContext.getCachedSubaccess(result, indexNo);
+    if (value) {
+      result = value;
+      return success();
+    }
 
-  // Create the result operation, inserting at the location of the declaration.
-  // This will cache the subindex operation for further uses.
-  locationProcessor.setLoc(loc);
-  OpBuilder::InsertionGuard guard(builder);
-  builder.setInsertionPointAfterValue(result);
-  auto op = builder.create<SubindexOp>(resultType, result, attrs);
+    // Create the result operation, inserting at the location of the
+    // declaration. This will cache the subindex operation for further uses.
+    locationProcessor.setLoc(loc);
+    OpBuilder::InsertionGuard guard(builder);
+    builder.setInsertionPointAfterValue(result);
 
-  // Insert the newly creatd operation into the cache.
-  value = op.getResult();
+    auto op = builder.create<SubindexOp>(resultType, result, attrs);
+
+    // Insert the newly created operation into the cache.
+    value = op.getResult();
+    result = value;
+  }
 
-  result = value;
   return success();
 }