[TableGen][GISel] Extract helper function for constraining operands

As a side effect, this fixes COPY_TO_REGCLASS not being constrained
if it is not the top-level instruction in the destination pattern.

Author: s-barannikov

llvm/test/CodeGen/AArch64/GlobalISel/select-bitcast-bigendian.mir

@@ -1,3 +1,4 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
 # RUN: llc -O0 -mtriple=arm64eb-- -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s
 ---
 name:            bitcast_v2f32_to_s64
@@ -9,9 +10,13 @@ body:             |
     liveins: $x0
 
     ; CHECK-LABEL: name: bitcast_v2f32_to_s64
-    ; CHECK: [[COPY:%[0-9]+]]:fpr64 = COPY $x0
-    ; CHECK: [[REV:%[0-9]+]]:fpr64 = REV64v2i32 [[COPY]]
-    ; CHECK: $x0 = COPY [[REV]]
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:fpr64 = COPY $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr64 = COPY [[COPY]]
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:fpr64 = COPY [[COPY1]]
+    ; CHECK-NEXT: [[REV64v2i32_:%[0-9]+]]:fpr64 = REV64v2i32 [[COPY2]]
+    ; CHECK-NEXT: $x0 = COPY [[REV64v2i32_]]
     %0:fpr(<2 x s32>) = COPY $x0
     %1:fpr(s64) = G_BITCAST %0
     $x0 = COPY %1(s64)

llvm/test/CodeGen/AMDGPU/GlobalISel/inst-select-sext.mir

@@ -107,7 +107,7 @@ body: |
     ; GCN-NEXT: [[S_MOV_B32_:%[0-9]+]]:sreg_32 = S_MOV_B32 31
     ; GCN-NEXT: [[S_SEXT_I32_I16_:%[0-9]+]]:sreg_32 = S_SEXT_I32_I16 [[COPY]]
     ; GCN-NEXT: [[S_ASHR_I32_:%[0-9]+]]:sreg_32 = S_ASHR_I32 [[S_SEXT_I32_I16_]], [[S_MOV_B32_]], implicit-def dead $scc
-    ; GCN-NEXT: [[COPY1:%[0-9]+]]:sreg_32_xm0 = COPY [[S_ASHR_I32_]]
+    ; GCN-NEXT: [[COPY1:%[0-9]+]]:sgpr_32 = COPY [[S_ASHR_I32_]]
     ; GCN-NEXT: [[S_SEXT_I32_I16_1:%[0-9]+]]:sreg_32_xexec_hi_and_sreg_32_xm0 = S_SEXT_I32_I16 [[COPY]]
     ; GCN-NEXT: [[REG_SEQUENCE:%[0-9]+]]:sreg_64 = REG_SEQUENCE [[S_SEXT_I32_I16_1]], %subreg.sub0, [[COPY1]], %subreg.sub1
     ; GCN-NEXT: $sgpr0_sgpr1 = COPY [[REG_SEQUENCE]]

llvm/test/CodeGen/AMDGPU/GlobalISel/inst-select-zext.mir

@@ -107,7 +107,7 @@ body: |
     ; GCN-NEXT: [[S_MOV_B32_:%[0-9]+]]:sreg_32_xm0 = S_MOV_B32 0
     ; GCN-NEXT: [[S_MOV_B32_1:%[0-9]+]]:sreg_32 = S_MOV_B32 65535
     ; GCN-NEXT: [[S_AND_B32_:%[0-9]+]]:sreg_32 = S_AND_B32 [[COPY]], [[S_MOV_B32_1]], implicit-def dead $scc
-    ; GCN-NEXT: [[COPY1:%[0-9]+]]:sreg_32_xexec_hi_and_sreg_32_xm0 = COPY [[S_AND_B32_]]
+    ; GCN-NEXT: [[COPY1:%[0-9]+]]:sgpr_32 = COPY [[S_AND_B32_]]
     ; GCN-NEXT: [[REG_SEQUENCE:%[0-9]+]]:sreg_64 = REG_SEQUENCE [[COPY1]], %subreg.sub0, [[S_MOV_B32_]], %subreg.sub1
     ; GCN-NEXT: $sgpr0_sgpr1 = COPY [[REG_SEQUENCE]]
     %0:sgpr(s32) = COPY $sgpr0

llvm/utils/TableGen/GlobalISelEmitter.cpp

@@ -460,6 +460,9 @@ class GlobalISelEmitter final : public GlobalISelMatchTableExecutorEmitter {
   std::optional<const CodeGenRegisterClass *>
   inferRegClassFromPattern(const TreePatternNode &N);
 
+  Error constrainOperands(action_iterator InsertPt, RuleMatcher &M,
+                          unsigned InsnID, const TreePatternNode &Dst);
+
   /// Return the size of the MemoryVT in this predicate, if possible.
   std::optional<unsigned>
   getMemSizeBitsFromPredicate(const TreePredicateFn &Predicate);
@@ -1409,103 +1412,10 @@ GlobalISelEmitter::createAndImportSubInstructionRenderer(
   if (auto Error = InsertPtOrError.takeError())
     return std::move(Error);
 
-  // We need to make sure that when we import an INSERT_SUBREG as a
-  // subinstruction that it ends up being constrained to the correct super
-  // register and subregister classes.
-  auto OpName = Target.getInstruction(Dst.getOperator()).TheDef->getName();
-  if (OpName == "INSERT_SUBREG") {
-    auto SubClass = inferRegClassFromPattern(Dst.getChild(1));
-    if (!SubClass)
-      return failedImport(
-          "Cannot infer register class from INSERT_SUBREG operand #1");
-    std::optional<const CodeGenRegisterClass *> SuperClass =
-        inferSuperRegisterClassForNode(Dst.getExtType(0), Dst.getChild(0),
-                                       Dst.getChild(2));
-    if (!SuperClass)
-      return failedImport(
-          "Cannot infer register class for INSERT_SUBREG operand #0");
-    // The destination and the super register source of an INSERT_SUBREG must
-    // be the same register class.
-    M.insertAction<ConstrainOperandToRegClassAction>(
-        InsertPt, DstMIBuilder.getInsnID(), 0, **SuperClass);
-    M.insertAction<ConstrainOperandToRegClassAction>(
-        InsertPt, DstMIBuilder.getInsnID(), 1, **SuperClass);
-    M.insertAction<ConstrainOperandToRegClassAction>(
-        InsertPt, DstMIBuilder.getInsnID(), 2, **SubClass);
-    return InsertPtOrError.get();
-  }
-
-  if (OpName == "EXTRACT_SUBREG") {
-    // EXTRACT_SUBREG selects into a subregister COPY but unlike most
-    // instructions, the result register class is controlled by the
-    // subregisters of the operand. As a result, we must constrain the result
-    // class rather than check that it's already the right one.
-    auto SuperClass = inferRegClassFromPattern(Dst.getChild(0));
-    if (!SuperClass)
-      return failedImport(
-          "Cannot infer register class from EXTRACT_SUBREG operand #0");
-
-    auto SubIdx = inferSubRegIndexForNode(Dst.getChild(1));
-    if (!SubIdx)
-      return failedImport("EXTRACT_SUBREG child #1 is not a subreg index");
-
-    const auto SrcRCDstRCPair =
-        (*SuperClass)->getMatchingSubClassWithSubRegs(CGRegs, *SubIdx);
-    assert(SrcRCDstRCPair->second && "Couldn't find a matching subclass");
-    M.insertAction<ConstrainOperandToRegClassAction>(
-        InsertPt, DstMIBuilder.getInsnID(), 0, *SrcRCDstRCPair->second);
-    M.insertAction<ConstrainOperandToRegClassAction>(
-        InsertPt, DstMIBuilder.getInsnID(), 1, *SrcRCDstRCPair->first);
-
-    // We're done with this pattern!  It's eligible for GISel emission; return
-    // it.
-    return InsertPtOrError.get();
-  }
-
-  // Similar to INSERT_SUBREG, we also have to handle SUBREG_TO_REG as a
-  // subinstruction.
-  if (OpName == "SUBREG_TO_REG") {
-    auto SubClass = inferRegClassFromPattern(Dst.getChild(1));
-    if (!SubClass)
-      return failedImport(
-          "Cannot infer register class from SUBREG_TO_REG child #1");
-    auto SuperClass =
-        inferSuperRegisterClass(Dst.getExtType(0), Dst.getChild(2));
-    if (!SuperClass)
-      return failedImport(
-          "Cannot infer register class for SUBREG_TO_REG operand #0");
-    M.insertAction<ConstrainOperandToRegClassAction>(
-        InsertPt, DstMIBuilder.getInsnID(), 0, **SuperClass);
-    M.insertAction<ConstrainOperandToRegClassAction>(
-        InsertPt, DstMIBuilder.getInsnID(), 2, **SubClass);
-    return InsertPtOrError.get();
-  }
-
-  if (OpName == "REG_SEQUENCE") {
-    auto SuperClass = inferRegClassFromPattern(Dst.getChild(0));
-    M.insertAction<ConstrainOperandToRegClassAction>(
-        InsertPt, DstMIBuilder.getInsnID(), 0, **SuperClass);
-
-    unsigned Num = Dst.getNumChildren();
-    for (unsigned I = 1; I != Num; I += 2) {
-      const TreePatternNode &SubRegChild = Dst.getChild(I + 1);
-
-      auto SubIdx = inferSubRegIndexForNode(SubRegChild);
-      if (!SubIdx)
-        return failedImport("REG_SEQUENCE child is not a subreg index");
-
-      const auto SrcRCDstRCPair =
-          (*SuperClass)->getMatchingSubClassWithSubRegs(CGRegs, *SubIdx);
-      assert(SrcRCDstRCPair->second && "Couldn't find a matching subclass");
-      M.insertAction<ConstrainOperandToRegClassAction>(
-          InsertPt, DstMIBuilder.getInsnID(), I, *SrcRCDstRCPair->second);
-    }
-
-    return InsertPtOrError.get();
-  }
+  if (auto Error =
+          constrainOperands(InsertPt, M, DstMIBuilder.getInsnID(), Dst))
+    return std::move(Error);
 
-  M.insertAction<ConstrainOperandsToDefinitionAction>(InsertPt,
-                                                      DstMIBuilder.getInsnID());
   return InsertPtOrError.get();
 }
 
@@ -1794,6 +1704,116 @@ Error GlobalISelEmitter::importImplicitDefRenderers(
   return Error::success();
 }
 
+Error GlobalISelEmitter::constrainOperands(action_iterator InsertPt,
+                                           RuleMatcher &M, unsigned InsnID,
+                                           const TreePatternNode &Dst) {
+  const Record *DstOp = Dst.getOperator();
+  const CodeGenInstruction &DstI = Target.getInstruction(DstOp);
+  StringRef DstIName = DstI.TheDef->getName();
+
+  if (DstIName == "COPY_TO_REGCLASS") {
+    // COPY_TO_REGCLASS does not provide operand constraints itself but the
+    // result is constrained to the class given by the second child.
+    const Record *DstIOpRec =
+        getInitValueAsRegClass(Dst.getChild(1).getLeafValue());
+
+    if (DstIOpRec == nullptr)
+      return failedImport("COPY_TO_REGCLASS operand #1 isn't a register class");
+
+    M.insertAction<ConstrainOperandToRegClassAction>(
+        InsertPt, InsnID, 0, Target.getRegisterClass(DstIOpRec));
+  } else if (DstIName == "EXTRACT_SUBREG") {
+    auto SuperClass = inferRegClassFromPattern(Dst.getChild(0));
+    if (!SuperClass)
+      return failedImport(
+          "Cannot infer register class from EXTRACT_SUBREG operand #0");
+
+    auto SubIdx = inferSubRegIndexForNode(Dst.getChild(1));
+    if (!SubIdx)
+      return failedImport("EXTRACT_SUBREG child #1 is not a subreg index");
+
+    // It would be nice to leave this constraint implicit but we're required
+    // to pick a register class so constrain the result to a register class
+    // that can hold the correct MVT.
+    //
+    // FIXME: This may introduce an extra copy if the chosen class doesn't
+    //        actually contain the subregisters.
+    const auto SrcRCDstRCPair =
+        (*SuperClass)->getMatchingSubClassWithSubRegs(CGRegs, *SubIdx);
+    if (!SrcRCDstRCPair) {
+      return failedImport("subreg index is incompatible "
+                          "with inferred reg class");
+    }
+
+    assert(SrcRCDstRCPair->second && "Couldn't find a matching subclass");
+    M.insertAction<ConstrainOperandToRegClassAction>(InsertPt, InsnID, 0,
+                                                     *SrcRCDstRCPair->second);
+    M.insertAction<ConstrainOperandToRegClassAction>(InsertPt, InsnID, 1,
+                                                     *SrcRCDstRCPair->first);
+  } else if (DstIName == "INSERT_SUBREG") {
+    // We need to constrain the destination, a super regsister source, and a
+    // subregister source.
+    auto SubClass = inferRegClassFromPattern(Dst.getChild(1));
+    if (!SubClass)
+      return failedImport(
+          "Cannot infer register class from INSERT_SUBREG operand #1");
+    auto SuperClass = inferSuperRegisterClassForNode(
+        Dst.getExtType(0), Dst.getChild(0), Dst.getChild(2));
+    if (!SuperClass)
+      return failedImport(
+          "Cannot infer register class for INSERT_SUBREG operand #0");
+    M.insertAction<ConstrainOperandToRegClassAction>(InsertPt, InsnID, 0,
+                                                     **SuperClass);
+    M.insertAction<ConstrainOperandToRegClassAction>(InsertPt, InsnID, 1,
+                                                     **SuperClass);
+    M.insertAction<ConstrainOperandToRegClassAction>(InsertPt, InsnID, 2,
+                                                     **SubClass);
+  } else if (DstIName == "SUBREG_TO_REG") {
+    // We need to constrain the destination and subregister source.
+    // Attempt to infer the subregister source from the first child. If it has
+    // an explicitly given register class, we'll use that. Otherwise, we will
+    // fail.
+    auto SubClass = inferRegClassFromPattern(Dst.getChild(1));
+    if (!SubClass)
+      return failedImport(
+          "Cannot infer register class from SUBREG_TO_REG child #1");
+    // We don't have a child to look at that might have a super register node.
+    auto SuperClass =
+        inferSuperRegisterClass(Dst.getExtType(0), Dst.getChild(2));
+    if (!SuperClass)
+      return failedImport(
+          "Cannot infer register class for SUBREG_TO_REG operand #0");
+    M.insertAction<ConstrainOperandToRegClassAction>(InsertPt, InsnID, 0,
+                                                     **SuperClass);
+    M.insertAction<ConstrainOperandToRegClassAction>(InsertPt, InsnID, 2,
+                                                     **SubClass);
+  } else if (DstIName == "REG_SEQUENCE") {
+    auto SuperClass = inferRegClassFromPattern(Dst.getChild(0));
+
+    M.insertAction<ConstrainOperandToRegClassAction>(InsertPt, InsnID, 0,
+                                                     **SuperClass);
+
+    unsigned Num = Dst.getNumChildren();
+    for (unsigned I = 1; I != Num; I += 2) {
+      const TreePatternNode &SubRegChild = Dst.getChild(I + 1);
+
+      auto SubIdx = inferSubRegIndexForNode(SubRegChild);
+      if (!SubIdx)
+        return failedImport("REG_SEQUENCE child is not a subreg index");
+
+      const auto SrcRCDstRCPair =
+          (*SuperClass)->getMatchingSubClassWithSubRegs(CGRegs, *SubIdx);
+
+      M.insertAction<ConstrainOperandToRegClassAction>(InsertPt, InsnID, I,
+                                                       *SrcRCDstRCPair->second);
+    }
+  } else {
+    M.insertAction<ConstrainOperandsToDefinitionAction>(InsertPt, InsnID);
+  }
+
+  return Error::success();
+}
+
 std::optional<const CodeGenRegisterClass *>
 GlobalISelEmitter::getRegClassFromLeaf(const TreePatternNode &Leaf) {
   assert(Leaf.isLeaf() && "Expected leaf?");
@@ -2123,106 +2143,9 @@ Expected<RuleMatcher> GlobalISelEmitter::runOnPattern(const PatternToMatch &P) {
 
   // Constrain the registers to classes. This is normally derived from the
   // emitted instruction but a few instructions require special handling.
-  if (DstIName == "COPY_TO_REGCLASS") {
-    // COPY_TO_REGCLASS does not provide operand constraints itself but the
-    // result is constrained to the class given by the second child.
-    const Record *DstIOpRec =
-        getInitValueAsRegClass(Dst.getChild(1).getLeafValue());
-
-    if (DstIOpRec == nullptr)
-      return failedImport("COPY_TO_REGCLASS operand #1 isn't a register class");
-
-    M.addAction<ConstrainOperandToRegClassAction>(
-        0, 0, Target.getRegisterClass(DstIOpRec));
-  } else if (DstIName == "EXTRACT_SUBREG") {
-    auto SuperClass = inferRegClassFromPattern(Dst.getChild(0));
-    if (!SuperClass)
-      return failedImport(
-          "Cannot infer register class from EXTRACT_SUBREG operand #0");
-
-    auto SubIdx = inferSubRegIndexForNode(Dst.getChild(1));
-    if (!SubIdx)
-      return failedImport("EXTRACT_SUBREG child #1 is not a subreg index");
-
-    // It would be nice to leave this constraint implicit but we're required
-    // to pick a register class so constrain the result to a register class
-    // that can hold the correct MVT.
-    //
-    // FIXME: This may introduce an extra copy if the chosen class doesn't
-    //        actually contain the subregisters.
-    assert(Src.getExtTypes().size() == 1 &&
-           "Expected Src of EXTRACT_SUBREG to have one result type");
-
-    const auto SrcRCDstRCPair =
-        (*SuperClass)->getMatchingSubClassWithSubRegs(CGRegs, *SubIdx);
-    if (!SrcRCDstRCPair) {
-      return failedImport("subreg index is incompatible "
-                          "with inferred reg class");
-    }
-
-    assert(SrcRCDstRCPair->second && "Couldn't find a matching subclass");
-    M.addAction<ConstrainOperandToRegClassAction>(0, 0,
-                                                  *SrcRCDstRCPair->second);
-    M.addAction<ConstrainOperandToRegClassAction>(0, 1, *SrcRCDstRCPair->first);
-  } else if (DstIName == "INSERT_SUBREG") {
-    assert(Src.getExtTypes().size() == 1 &&
-           "Expected Src of INSERT_SUBREG to have one result type");
-    // We need to constrain the destination, a super regsister source, and a
-    // subregister source.
-    auto SubClass = inferRegClassFromPattern(Dst.getChild(1));
-    if (!SubClass)
-      return failedImport(
-          "Cannot infer register class from INSERT_SUBREG operand #1");
-    auto SuperClass = inferSuperRegisterClassForNode(
-        Src.getExtType(0), Dst.getChild(0), Dst.getChild(2));
-    if (!SuperClass)
-      return failedImport(
-          "Cannot infer register class for INSERT_SUBREG operand #0");
-    M.addAction<ConstrainOperandToRegClassAction>(0, 0, **SuperClass);
-    M.addAction<ConstrainOperandToRegClassAction>(0, 1, **SuperClass);
-    M.addAction<ConstrainOperandToRegClassAction>(0, 2, **SubClass);
-  } else if (DstIName == "SUBREG_TO_REG") {
-    // We need to constrain the destination and subregister source.
-    assert(Src.getExtTypes().size() == 1 &&
-           "Expected Src of SUBREG_TO_REG to have one result type");
-
-    // Attempt to infer the subregister source from the first child. If it has
-    // an explicitly given register class, we'll use that. Otherwise, we will
-    // fail.
-    auto SubClass = inferRegClassFromPattern(Dst.getChild(1));
-    if (!SubClass)
-      return failedImport(
-          "Cannot infer register class from SUBREG_TO_REG child #1");
-    // We don't have a child to look at that might have a super register node.
-    auto SuperClass =
-        inferSuperRegisterClass(Src.getExtType(0), Dst.getChild(2));
-    if (!SuperClass)
-      return failedImport(
-          "Cannot infer register class for SUBREG_TO_REG operand #0");
-    M.addAction<ConstrainOperandToRegClassAction>(0, 0, **SuperClass);
-    M.addAction<ConstrainOperandToRegClassAction>(0, 2, **SubClass);
-  } else if (DstIName == "REG_SEQUENCE") {
-    auto SuperClass = inferRegClassFromPattern(Dst.getChild(0));
-
-    M.addAction<ConstrainOperandToRegClassAction>(0, 0, **SuperClass);
-
-    unsigned Num = Dst.getNumChildren();
-    for (unsigned I = 1; I != Num; I += 2) {
-      TreePatternNode &SubRegChild = Dst.getChild(I + 1);
-
-      auto SubIdx = inferSubRegIndexForNode(SubRegChild);
-      if (!SubIdx)
-        return failedImport("REG_SEQUENCE child is not a subreg index");
-
-      const auto SrcRCDstRCPair =
-          (*SuperClass)->getMatchingSubClassWithSubRegs(CGRegs, *SubIdx);
-
-      M.addAction<ConstrainOperandToRegClassAction>(0, I,
-                                                    *SrcRCDstRCPair->second);
-    }
-  } else {
-    M.addAction<ConstrainOperandsToDefinitionAction>(0);
-  }
+  if (auto Error =
+          constrainOperands(M.actions_end(), M, DstMIBuilder.getInsnID(), Dst))
+    return std::move(Error);
 
   // Erase the root.
   unsigned RootInsnID = M.getInsnVarID(InsnMatcher);