[TableGen] Combine the two separate OperandMapping loops in PseudoLoweringEmitter. (#136007)

Previously we had one loop over the DAG for immediates and registers and
another loop over the destination operands for mapping from the source.

Now we have a single loop over the destination operands that handles immediates,
registers, and named operands. A helper method is added so we can handle
operands and sub-operands specified by a sub-dag.

My goal is to allow a named operand to appear in a sub-dag which wasn't
supported before. This will allow the destination instruction to have an
operand with sub-operands when the source does not have sub operands.

For RISC-V, I'm looking into using an operand with sub-operands to
represent an reg+offset memory address. I need to be able to lower a
pseudo instruction that only has a register operand to an instruction
that has a reg+offset operand. The offset will be filled in with 0
during expansion and the register will be copied from the source.

The expansion would look like this:
def PseudoCALLIndirect : Pseudo<(outs), (ins GPRJALR:$rs1),
                                [(riscv_call GPRJALR:$rs1)]>,
PseudoInstExpansion<(JALR X1, (ops GPR:$rs1, 0))>;

Author: topperc

llvm/utils/TableGen/PseudoLoweringEmitter.cpp

@@ -51,10 +51,11 @@ class PseudoLoweringEmitter {
 
   SmallVector<PseudoExpansion, 64> Expansions;
 
-  unsigned addDagOperandMapping(const Record *Rec, const DagInit *Dag,
-                                const CodeGenInstruction &Insn,
-                                IndexedMap<OpData> &OperandMap,
-                                unsigned BaseIdx);
+  void addOperandMapping(unsigned MIOpNo, unsigned NumOps, const Record *Rec,
+                         const DagInit *Dag, unsigned DagIdx,
+                         const Record *OpRec, IndexedMap<OpData> &OperandMap,
+                         const StringMap<unsigned> &SourceOperands,
+                         const CodeGenInstruction &SourceInsn);
   void evaluateExpansion(const Record *Pseudo);
   void emitLoweringEmitter(raw_ostream &o);
 
@@ -66,64 +67,67 @@ class PseudoLoweringEmitter {
 };
 } // End anonymous namespace
 
-// FIXME: This pass currently can only expand a pseudo to a single instruction.
-//        The pseudo expansion really should take a list of dags, not just
-//        a single dag, so we can do fancier things.
-unsigned PseudoLoweringEmitter::addDagOperandMapping(
-    const Record *Rec, const DagInit *Dag, const CodeGenInstruction &Insn,
-    IndexedMap<OpData> &OperandMap, unsigned BaseIdx) {
-  unsigned OpsAdded = 0;
-  for (unsigned i = 0, e = Dag->getNumArgs(); i != e; ++i) {
-    if (const DefInit *DI = dyn_cast<DefInit>(Dag->getArg(i))) {
-      // Physical register reference. Explicit check for the special case
-      // "zero_reg" definition.
-      if (DI->getDef()->isSubClassOf("Register") ||
-          DI->getDef()->getName() == "zero_reg") {
-        auto &Entry = OperandMap[BaseIdx + i];
-        Entry.Kind = OpData::Reg;
-        Entry.Data.Reg = DI->getDef();
-        ++OpsAdded;
-        continue;
-      }
+void PseudoLoweringEmitter::addOperandMapping(
+    unsigned MIOpNo, unsigned NumOps, const Record *Rec, const DagInit *Dag,
+    unsigned DagIdx, const Record *OpRec, IndexedMap<OpData> &OperandMap,
+    const StringMap<unsigned> &SourceOperands,
+    const CodeGenInstruction &SourceInsn) {
+  const Init *DagArg = Dag->getArg(DagIdx);
+  if (const DefInit *DI = dyn_cast<DefInit>(DagArg)) {
+    // Physical register reference. Explicit check for the special case
+    // "zero_reg" definition.
+    if (DI->getDef()->isSubClassOf("Register") ||
+        DI->getDef()->getName() == "zero_reg") {
+      auto &Entry = OperandMap[MIOpNo];
+      Entry.Kind = OpData::Reg;
+      Entry.Data.Reg = DI->getDef();
+      return;
+    }
 
-      // Normal operands should always have the same type, or we have a
-      // problem.
-      // FIXME: We probably shouldn't ever get a non-zero BaseIdx here.
-      assert(BaseIdx == 0 && "Named subargument in pseudo expansion?!");
-      if (DI->getDef() != Insn.Operands[BaseIdx + i].Rec)
-        PrintFatalError(Rec, "In pseudo instruction '" + Rec->getName() +
-                                 "', operand type '" + DI->getDef()->getName() +
-                                 "' does not match expansion operand type '" +
-                                 Insn.Operands[BaseIdx + i].Rec->getName() +
-                                 "'");
-      // Source operand maps to destination operand. The Data element
-      // will be filled in later, just set the Kind for now. Do it
-      // for each corresponding MachineInstr operand, not just the first.
-      for (unsigned I = 0, E = Insn.Operands[i].MINumOperands; I != E; ++I)
-        OperandMap[BaseIdx + i + I].Kind = OpData::Operand;
-      OpsAdded += Insn.Operands[i].MINumOperands;
-    } else if (const IntInit *II = dyn_cast<IntInit>(Dag->getArg(i))) {
-      auto &Entry = OperandMap[BaseIdx + i];
-      Entry.Kind = OpData::Imm;
-      Entry.Data.Imm = II->getValue();
-      ++OpsAdded;
-    } else if (const auto *BI = dyn_cast<BitsInit>(Dag->getArg(i))) {
-      auto &Entry = OperandMap[BaseIdx + i];
-      Entry.Kind = OpData::Imm;
-      Entry.Data.Imm = *BI->convertInitializerToInt();
-      ++OpsAdded;
-    } else if (const DagInit *SubDag = dyn_cast<DagInit>(Dag->getArg(i))) {
-      // Just add the operands recursively. This is almost certainly
-      // a constant value for a complex operand (> 1 MI operand).
-      unsigned NewOps =
-          addDagOperandMapping(Rec, SubDag, Insn, OperandMap, BaseIdx + i);
-      OpsAdded += NewOps;
-      // Since we added more than one, we also need to adjust the base.
-      BaseIdx += NewOps - 1;
-    } else
-      llvm_unreachable("Unhandled pseudo-expansion argument type!");
-  }
-  return OpsAdded;
+    if (DI->getDef() != OpRec)
+      PrintFatalError(Rec, "In pseudo instruction '" + Rec->getName() +
+                               "', operand type '" + DI->getDef()->getName() +
+                               "' does not match expansion operand type '" +
+                               OpRec->getName() + "'");
+
+    StringMap<unsigned>::const_iterator SourceOp =
+        SourceOperands.find(Dag->getArgNameStr(DagIdx));
+    if (SourceOp == SourceOperands.end())
+      PrintFatalError(Rec, "In pseudo instruction '" + Rec->getName() +
+                               "', output operand '" +
+                               Dag->getArgNameStr(DagIdx) +
+                               "' has no matching source operand");
+    const auto &SrcOpnd = SourceInsn.Operands[SourceOp->getValue()];
+    if (NumOps != SrcOpnd.MINumOperands)
+      PrintFatalError(
+          Rec,
+          "In pseudo instruction '" + Rec->getName() + "', output operand '" +
+              OpRec->getName() +
+              "' has a different number of sub operands than source operand '" +
+              SrcOpnd.Rec->getName() + "'");
+
+    // Source operand maps to destination operand. Do it for each corresponding
+    // MachineInstr operand, not just the first.
+    for (unsigned I = 0, E = NumOps; I != E; ++I) {
+      auto &Entry = OperandMap[MIOpNo + I];
+      Entry.Kind = OpData::Operand;
+      Entry.Data.Operand = SrcOpnd.MIOperandNo + I;
+    }
+
+    LLVM_DEBUG(dbgs() << "    " << SourceOp->getValue() << " ==> " << DagIdx
+                      << "\n");
+  } else if (const auto *II = dyn_cast<IntInit>(DagArg)) {
+    assert(NumOps == 1);
+    auto &Entry = OperandMap[MIOpNo];
+    Entry.Kind = OpData::Imm;
+    Entry.Data.Imm = II->getValue();
+  } else if (const auto *BI = dyn_cast<BitsInit>(DagArg)) {
+    assert(NumOps == 1);
+    auto &Entry = OperandMap[MIOpNo];
+    Entry.Kind = OpData::Imm;
+    Entry.Data.Imm = *BI->convertInitializerToInt();
+  } else
+    llvm_unreachable("Unhandled pseudo-expansion argument type!");
 }
 
 void PseudoLoweringEmitter::evaluateExpansion(const Record *Rec) {
@@ -157,14 +161,6 @@ void PseudoLoweringEmitter::evaluateExpansion(const Record *Rec) {
                              "', result operator '" + Operator->getName() +
                              "' has the wrong number of operands");
 
-  unsigned NumMIOperands = 0;
-  for (const auto &Op : Insn.Operands)
-    NumMIOperands += Op.MINumOperands;
-  IndexedMap<OpData> OperandMap;
-  OperandMap.grow(NumMIOperands);
-
-  addDagOperandMapping(Rec, Dag, Insn, OperandMap, 0);
-
   // If there are more operands that weren't in the DAG, they have to
   // be operands that have default values, or we have an error. Currently,
   // Operands that are a subclass of OperandWithDefaultOp have default values.
@@ -180,37 +176,43 @@ void PseudoLoweringEmitter::evaluateExpansion(const Record *Rec) {
   for (const auto &[Idx, SrcOp] : enumerate(SourceInsn.Operands))
     SourceOperands[SrcOp.Name] = Idx;
 
-  LLVM_DEBUG(dbgs() << "  Operand mapping:\n");
-  for (const auto &[Idx, Opnd] : enumerate(Insn.Operands)) {
-    // We've already handled constant values. Just map instruction operands
-    // here.
-    if (OperandMap[Opnd.MIOperandNo].Kind != OpData::Operand)
-      continue;
-    StringMap<unsigned>::iterator SourceOp =
-        SourceOperands.find(Dag->getArgNameStr(Idx));
-    if (SourceOp == SourceOperands.end())
-      PrintFatalError(Rec, "In pseudo instruction '" + Rec->getName() +
-                               "', output operand '" + Dag->getArgNameStr(Idx) +
-                               "' has no matching source operand");
-    const auto &SrcOpnd = SourceInsn.Operands[SourceOp->getValue()];
-    if (Opnd.MINumOperands != SrcOpnd.MINumOperands)
-      PrintFatalError(
-          Rec,
-          "In pseudo instruction '" + Rec->getName() + "', output operand '" +
-              Opnd.Rec->getName() +
-              "' has a different number of sub operands than source operand '" +
-              SrcOpnd.Rec->getName() + "'");
+  unsigned NumMIOperands = 0;
+  for (const auto &Op : Insn.Operands)
+    NumMIOperands += Op.MINumOperands;
+  IndexedMap<OpData> OperandMap;
+  OperandMap.grow(NumMIOperands);
 
-    // Map the source operand to the destination operand index for each
-    // MachineInstr operand.
-    for (unsigned I = 0, E = Opnd.MINumOperands; I != E; ++I)
-      OperandMap[Opnd.MIOperandNo + I].Data.Operand = SrcOpnd.MIOperandNo + I;
+  // FIXME: This pass currently can only expand a pseudo to a single
+  // instruction. The pseudo expansion really should take a list of dags, not
+  // just a single dag, so we can do fancier things.
+  LLVM_DEBUG(dbgs() << "  Operand mapping:\n");
+  for (const auto &[Idx, DstOp] : enumerate(Insn.Operands)) {
+    unsigned MIOpNo = DstOp.MIOperandNo;
 
-    LLVM_DEBUG(dbgs() << "    " << SourceOp->getValue() << " ==> " << Idx
-                      << "\n");
+    if (const auto *SubDag = dyn_cast<DagInit>(Dag->getArg(Idx))) {
+      if (!DstOp.MIOperandInfo || DstOp.MIOperandInfo->getNumArgs() == 0)
+        PrintFatalError(Rec, "In pseudo instruction '" + Rec->getName() +
+                                 "', operand '" + DstOp.Rec->getName() +
+                                 "' does not have suboperands");
+      if (DstOp.MINumOperands != SubDag->getNumArgs()) {
+        PrintFatalError(
+            Rec, "In pseudo instruction '" + Rec->getName() + "', '" +
+                     SubDag->getAsString() +
+                     "' has wrong number of operands for operand type '" +
+                     DstOp.Rec->getName() + "'");
+      }
+      for (unsigned I = 0, E = DstOp.MINumOperands; I != E; ++I) {
+        auto *OpndRec = cast<DefInit>(DstOp.MIOperandInfo->getArg(I))->getDef();
+        addOperandMapping(MIOpNo + I, 1, Rec, SubDag, I, OpndRec, OperandMap,
+                          SourceOperands, SourceInsn);
+      }
+    } else {
+      addOperandMapping(MIOpNo, DstOp.MINumOperands, Rec, Dag, Idx, DstOp.Rec,
+                        OperandMap, SourceOperands, SourceInsn);
+    }
   }
 
-  Expansions.push_back(PseudoExpansion(SourceInsn, Insn, OperandMap));
+  Expansions.emplace_back(SourceInsn, Insn, OperandMap);
 }
 
 void PseudoLoweringEmitter::emitLoweringEmitter(raw_ostream &o) {