[TableGen][GISel] Extract common function for determining MI's regclass (llvm#120135)

Add some comments that hopefully clarify a few things.

This was supposed to be NFC, but there is a difference in the inferred
register class for EXTRACT_SUBREG.

Pull Request: llvm#120135

Author: s-barannikov

llvm/test/TableGen/GlobalISelEmitter-nested-subregs.td

@@ -34,7 +34,7 @@ def A0  : RegisterClass<"MyTarget", [i32], 32, (add a0)>;
 // CHECK-NEXT: GIM_CheckOpcode, /*MI*/0, GIMT_Encode2(TargetOpcode::G_ANYEXT),
 // CHECK-NEXT: // MIs[0] DstI[dst]
 // CHECK-NEXT: GIM_RootCheckType, /*Op*/0, /*Type*/GILLT_s16,
-// CHECK-NEXT: GIM_RootCheckRegBankForClass, /*Op*/0, /*RC*/GIMT_Encode2(MyTarget::A0RegClassID),
+// CHECK-NEXT: GIM_RootCheckRegBankForClass, /*Op*/0, /*RC*/GIMT_Encode2(MyTarget::A0wRegClassID),
 // CHECK-NEXT: // MIs[0] src
 // CHECK-NEXT: GIM_RootCheckType, /*Op*/1, /*Type*/GILLT_s8,
 // CHECK-NEXT: // (anyext:{ *:[i16] } i8:{ *:[i8] }:$src)  =>  (EXTRACT_SUBREG:{ *:[i16] } (INSERT_SUBREG:{ *:[i32] } (IMPLICIT_DEF:{ *:[i32] }), A0b:{ *:[i8] }:$src, lo8:{ *:[i32] }), lo16:{ *:[i32] })

llvm/utils/TableGen/GlobalISelEmitter.cpp

@@ -459,6 +459,10 @@ class GlobalISelEmitter final : public GlobalISelMatchTableExecutorEmitter {
   const CodeGenRegisterClass *
   inferRegClassFromPattern(const TreePatternNode &N);
 
+  const CodeGenRegisterClass *
+  inferRegClassFromInstructionPattern(const TreePatternNode &N,
+                                      unsigned ResIdx);
+
   Error constrainOperands(action_iterator InsertPt, RuleMatcher &M,
                           unsigned InsnID, const TreePatternNode &Dst);
 
@@ -1856,46 +1860,85 @@ GlobalISelEmitter::inferRegClassFromPattern(const TreePatternNode &N) {
 
   // Don't want to try and infer things when there could potentially be more
   // than one candidate register class.
-  auto &Inst = Target.getInstruction(OpRec);
+  return inferRegClassFromInstructionPattern(N, /*ResIdx=*/0);
+}
+
+const CodeGenRegisterClass *
+GlobalISelEmitter::inferRegClassFromInstructionPattern(const TreePatternNode &N,
+                                                       unsigned ResIdx) {
+  const CodeGenInstruction &Inst = Target.getInstruction(N.getOperator());
+  assert(ResIdx < Inst.Operands.NumDefs &&
+         "Can only infer register class for explicit defs");
 
   // Handle any special-case instructions which we can safely infer register
   // classes from.
   StringRef InstName = Inst.TheDef->getName();
-  bool IsRegSequence = InstName == "REG_SEQUENCE";
-  if (IsRegSequence || InstName == "COPY_TO_REGCLASS") {
-    // If we have a COPY_TO_REGCLASS, then we need to handle it specially. It
-    // has the desired register class as the first child.
-    const TreePatternNode &RCChild = N.getChild(IsRegSequence ? 0 : 1);
+  if (InstName == "REG_SEQUENCE") {
+    // (outs $super_dst), (ins $dst_regclass, variable_ops)
+    // Destination register class is explicitly specified by the first operand.
+    const TreePatternNode &RCChild = N.getChild(0);
     if (!RCChild.isLeaf())
       return nullptr;
     return getRegClassFromLeaf(RCChild);
   }
+
+  if (InstName == "COPY_TO_REGCLASS") {
+    // (outs $dst), (ins $src, $dst_regclass)
+    // Destination register class is explicitly specified by the second operand.
+    const TreePatternNode &RCChild = N.getChild(1);
+    if (!RCChild.isLeaf())
+      return nullptr;
+    return getRegClassFromLeaf(RCChild);
+  }
+
   if (InstName == "INSERT_SUBREG") {
+    // (outs $super_dst), (ins $super_src, $sub_src, $sub_idx);
+    // If we can infer the register class for the first operand, use that.
+    // Otherwise, find a register class that supports both the specified
+    // sub-register index and the type of the instruction's result.
     const TreePatternNode &Child0 = N.getChild(0);
     assert(Child0.getNumTypes() == 1 && "Unexpected number of types!");
-    const TypeSetByHwMode &VTy = Child0.getExtType(0);
-    return inferSuperRegisterClassForNode(VTy, Child0, N.getChild(2));
+    return inferSuperRegisterClassForNode(N.getExtType(0), Child0,
+                                          N.getChild(2));
   }
+
   if (InstName == "EXTRACT_SUBREG") {
-    assert(N.getNumTypes() == 1 && "Unexpected number of types!");
-    const TypeSetByHwMode &VTy = N.getExtType(0);
-    return inferSuperRegisterClass(VTy, N.getChild(1));
+    // (outs $sub_dst), (ins $super_src, $sub_idx)
+    // Find a register class that can be used for a sub-register copy from
+    // the specified source at the specified sub-register index.
+    const CodeGenRegisterClass *SuperRC =
+        inferRegClassFromPattern(N.getChild(0));
+    if (!SuperRC)
+      return nullptr;
+
+    const CodeGenSubRegIndex *SubIdx = inferSubRegIndexForNode(N.getChild(1));
+    if (!SubIdx)
+      return nullptr;
+
+    const auto SubRCAndSubRegRC =
+        SuperRC->getMatchingSubClassWithSubRegs(CGRegs, SubIdx);
+    if (!SubRCAndSubRegRC)
+      return nullptr;
+
+    return SubRCAndSubRegRC->second;
+  }
+
+  if (InstName == "SUBREG_TO_REG") {
+    // (outs $super_dst), (ins $super_src, $sub_src, $sub_idx)
+    // Find a register class that supports both the specified sub-register
+    // index and the type of the instruction's result.
+    return inferSuperRegisterClass(N.getExtType(0), N.getChild(2));
   }
 
   // Handle destination record types that we can safely infer a register class
   // from.
-  const auto &DstIOperand = Inst.Operands[0];
+  const auto &DstIOperand = Inst.Operands[ResIdx];
   const Record *DstIOpRec = DstIOperand.Rec;
-  if (DstIOpRec->isSubClassOf("RegisterOperand")) {
-    DstIOpRec = DstIOpRec->getValueAsDef("RegClass");
-    const CodeGenRegisterClass &RC = Target.getRegisterClass(DstIOpRec);
-    return &RC;
-  }
+  if (DstIOpRec->isSubClassOf("RegisterOperand"))
+    return &Target.getRegisterClass(DstIOpRec->getValueAsDef("RegClass"));
 
-  if (DstIOpRec->isSubClassOf("RegisterClass")) {
-    const CodeGenRegisterClass &RC = Target.getRegisterClass(DstIOpRec);
-    return &RC;
-  }
+  if (DstIOpRec->isSubClassOf("RegisterClass"))
+    return &Target.getRegisterClass(DstIOpRec);
 
   return nullptr;
 }
@@ -2043,8 +2086,7 @@ Expected<RuleMatcher> GlobalISelEmitter::runOnPattern(const PatternToMatch &P) {
   if (!DstOp->isSubClassOf("Instruction"))
     return failedImport("Pattern operator isn't an instruction");
 
-  auto &DstI = Target.getInstruction(DstOp);
-  StringRef DstIName = DstI.TheDef->getName();
+  const CodeGenInstruction &DstI = Target.getInstruction(DstOp);
 
   // Count both implicit and explicit defs in the dst instruction.
   // This avoids errors importing patterns that have inherent implicit defs.
@@ -2070,68 +2112,24 @@ Expected<RuleMatcher> GlobalISelEmitter::runOnPattern(const PatternToMatch &P) {
 
   // The root of the match also has constraints on the register bank so that it
   // matches the result instruction.
-  unsigned OpIdx = 0;
   unsigned N = std::min(DstExpDefs, SrcNumDefs);
   for (unsigned I = 0; I < N; ++I) {
-    const TypeSetByHwMode &VTy = Src.getExtType(I);
+    const auto &DstIOperand = DstI.Operands[I];
 
-    const auto &DstIOperand = DstI.Operands[OpIdx];
-    PointerUnion<const Record *, const CodeGenRegisterClass *> MatchedRC =
-        DstIOperand.Rec;
-    if (DstIName == "COPY_TO_REGCLASS") {
-      MatchedRC = getInitValueAsRegClass(Dst.getChild(1).getLeafValue());
-
-      if (MatchedRC.isNull())
-        return failedImport(
-            "COPY_TO_REGCLASS operand #1 isn't a register class");
-    } else if (DstIName == "REG_SEQUENCE") {
-      MatchedRC = getInitValueAsRegClass(Dst.getChild(0).getLeafValue());
-      if (MatchedRC.isNull())
-        return failedImport("REG_SEQUENCE operand #0 isn't a register class");
-    } else if (DstIName == "EXTRACT_SUBREG") {
-      const CodeGenRegisterClass *InferredClass =
-          inferRegClassFromPattern(Dst.getChild(0));
-      if (!InferredClass)
-        return failedImport(
-            "Could not infer class for EXTRACT_SUBREG operand #0");
-
-      // We can assume that a subregister is in the same bank as it's super
-      // register.
-      MatchedRC = InferredClass->getDef();
-    } else if (DstIName == "INSERT_SUBREG") {
-      const CodeGenRegisterClass *MaybeSuperClass =
-          inferSuperRegisterClassForNode(VTy, Dst.getChild(0), Dst.getChild(2));
-      if (!MaybeSuperClass)
-        return failedImport(
-            "Cannot infer register class for INSERT_SUBREG operand #0");
-      // Move to the next pattern here, because the register class we found
-      // doesn't necessarily have a record associated with it. So, we can't
-      // set DstIOpRec using this.
-      MatchedRC = MaybeSuperClass;
-    } else if (DstIName == "SUBREG_TO_REG") {
-      const CodeGenRegisterClass *MaybeRegClass =
-          inferSuperRegisterClass(VTy, Dst.getChild(2));
-      if (!MaybeRegClass)
-        return failedImport(
-            "Cannot infer register class for SUBREG_TO_REG operand #0");
-      MatchedRC = MaybeRegClass;
-    } else if (cast<const Record *>(MatchedRC)->isSubClassOf("RegisterOperand"))
-      MatchedRC = cast<const Record *>(MatchedRC)->getValueAsDef("RegClass");
-    else if (!cast<const Record *>(MatchedRC)->isSubClassOf("RegisterClass"))
-      return failedImport("Dst MI def isn't a register class" + to_string(Dst));
-
-    OperandMatcher &OM = InsnMatcher.getOperand(OpIdx);
+    OperandMatcher &OM = InsnMatcher.getOperand(I);
     // The operand names declared in the DstI instruction are unrelated to
     // those used in pattern's source and destination DAGs, so mangle the
     // former to prevent implicitly adding unexpected
     // GIM_CheckIsSameOperand predicates by the defineOperand method.
     OM.setSymbolicName(getMangledRootDefName(DstIOperand.Name));
     M.defineOperand(OM.getSymbolicName(), OM);
-    if (auto *R = dyn_cast<const Record *>(MatchedRC))
-      MatchedRC = &Target.getRegisterClass(R);
-    OM.addPredicate<RegisterBankOperandMatcher>(
-        *cast<const CodeGenRegisterClass *>(MatchedRC));
-    ++OpIdx;
+
+    const CodeGenRegisterClass *RC =
+        inferRegClassFromInstructionPattern(Dst, I);
+    if (!RC)
+      return failedImport("Could not infer register class for result #" +
+                          Twine(I) + " from pattern " + to_string(Dst));
+    OM.addPredicate<RegisterBankOperandMatcher>(*RC);
   }
 
   auto DstMIBuilderOrError =