6 files changed, 76 insertions, 9 deletions
diff --git a/docs/MIRLangRef.rst b/docs/MIRLangRef.rst
index e217b792e62..562b11a6d18 100644
--- a/docs/MIRLangRef.rst
+++ b/docs/MIRLangRef.rst
@@ -430,6 +430,35 @@ immediate machine operand ``-42``:
 
     %eax = MOV32ri -42
 
+An immediate operand is also used to represent a subregister index when the
+machine instruction has one of the following opcodes:
+
+- ``EXTRACT_SUBREG``
+
+- ``INSERT_SUBREG``
+
+- ``REG_SEQUENCE``
+
+- ``SUBREG_TO_REG``
+
+In case this is true, the Machine Operand is printed according to the target.
+
+For example:
+
+In AArch64RegisterInfo.td:
+
+.. code-block:: text
+
+  def sub_32 : SubRegIndex<32>;
+
+If the third operand is an immediate with the value ``15`` (target-dependent
+value), based on the instruction's opcode and the operand's index the operand
+will be printed as ``%subreg.sub_32``:
+
+.. code-block:: text
+
+    %1:gpr64 = SUBREG_TO_REG 0, %0, %subreg.sub_32
+
 For integers > 64bit, we use a special machine operand, ``MO_CImmediate``,
 which stores the immediate in a ``ConstantInt`` using an ``APInt`` (LLVM's
 arbitrary precision integers).
diff --git a/include/llvm/CodeGen/MachineOperand.h b/include/llvm/CodeGen/MachineOperand.h
index 757de85f158..a7043ea90e3 100644
--- a/include/llvm/CodeGen/MachineOperand.h
+++ b/include/llvm/CodeGen/MachineOperand.h
@@ -227,6 +227,13 @@ public:
   ///
   void clearParent() { ParentMI = nullptr; }
 
+  /// Print a subreg index operand.
+  /// MO_Immediate operands can also be subreg idices. If it's the case, the
+  /// subreg index name will be printed. MachineInstr::isOperandSubregIdx can be
+  /// called to check this.
+  static void printSubregIdx(raw_ostream &OS, uint64_t Index,
+                             const TargetRegisterInfo *TRI);
+
   /// Print the MachineOperand to \p os.
   /// Providing a valid \p TRI and \p IntrinsicInfo results in a more
   /// target-specific printing. If \p TRI and \p IntrinsicInfo are null, the
diff --git a/lib/CodeGen/MIRPrinter.cpp b/lib/CodeGen/MIRPrinter.cpp
index 1250e3588bc..ccec5b4348d 100644
--- a/lib/CodeGen/MIRPrinter.cpp
+++ b/lib/CodeGen/MIRPrinter.cpp
@@ -854,23 +854,23 @@ void MIPrinter::print(const MachineInstr &MI, unsigned OpIdx,
   const MachineOperand &Op = MI.getOperand(OpIdx);
   printTargetFlags(Op);
   switch (Op.getType()) {
+  case MachineOperand::MO_Immediate:
+    if (MI.isOperandSubregIdx(OpIdx)) {
+      MachineOperand::printSubregIdx(OS, Op.getImm(), TRI);
+      break;
+    }
+    LLVM_FALLTHROUGH;
   case MachineOperand::MO_Register:
   case MachineOperand::MO_CImmediate:
   case MachineOperand::MO_MachineBasicBlock: {
     unsigned TiedOperandIdx = 0;
-    if (ShouldPrintRegisterTies && Op.isTied() && !Op.isDef())
+    if (ShouldPrintRegisterTies && Op.isReg() && Op.isTied() && !Op.isDef())
       TiedOperandIdx = Op.getParent()->findTiedOperandIdx(OpIdx);
     const TargetIntrinsicInfo *TII = MI.getMF()->getTarget().getIntrinsicInfo();
     Op.print(OS, MST, TypeToPrint, PrintDef, ShouldPrintRegisterTies,
              TiedOperandIdx, TRI, TII);
     break;
   }
-  case MachineOperand::MO_Immediate:
-    if (MI.isOperandSubregIdx(OpIdx))
-      OS << "%subreg." << TRI->getSubRegIndexName(Op.getImm());
-    else
-      OS << Op.getImm();
-    break;
   case MachineOperand::MO_FPImmediate:
     Op.getFPImm()->printAsOperand(OS, /*PrintType=*/true, MST);
     break;
diff --git a/lib/CodeGen/MachineInstr.cpp b/lib/CodeGen/MachineInstr.cpp
index fb0b82c348c..96722b26ee8 100644
--- a/lib/CodeGen/MachineInstr.cpp
+++ b/lib/CodeGen/MachineInstr.cpp
@@ -1405,8 +1405,11 @@ void MachineInstr::print(raw_ostream &OS, ModuleSlotTracker &MST,
     } else {
       LLT TypeToPrint = MRI ? getTypeToPrint(i, PrintedTypes, *MRI) : LLT{};
       unsigned TiedOperandIdx = getTiedOperandIdx(i);
-      MO.print(OS, MST, TypeToPrint, /*PrintDef=*/true, ShouldPrintRegisterTies,
-               TiedOperandIdx, TRI, IntrinsicInfo);
+      if (MO.isImm() && isOperandSubregIdx(i))
+        MachineOperand::printSubregIdx(OS, MO.getImm(), TRI);
+      else
+        MO.print(OS, MST, TypeToPrint, /*PrintDef=*/true,
+                 ShouldPrintRegisterTies, TiedOperandIdx, TRI, IntrinsicInfo);
     }
   }
 
diff --git a/lib/CodeGen/MachineOperand.cpp b/lib/CodeGen/MachineOperand.cpp
index 0cbcb65a99a..8bd6a7a965b 100644
--- a/lib/CodeGen/MachineOperand.cpp
+++ b/lib/CodeGen/MachineOperand.cpp
@@ -345,6 +345,15 @@ static void tryToGetTargetInfo(const MachineOperand &MO,
   }
 }
 
+void MachineOperand::printSubregIdx(raw_ostream &OS, uint64_t Index,
+                                    const TargetRegisterInfo *TRI) {
+  OS << "%subreg.";
+  if (TRI)
+    OS << TRI->getSubRegIndexName(Index);
+  else
+    OS << Index;
+}
+
 void MachineOperand::print(raw_ostream &OS, const TargetRegisterInfo *TRI,
                            const TargetIntrinsicInfo *IntrinsicInfo) const {
   tryToGetTargetInfo(*this, TRI, IntrinsicInfo);
diff --git a/unittests/CodeGen/MachineOperandTest.cpp b/unittests/CodeGen/MachineOperandTest.cpp
index 4884d374521..24a7f5563ff 100644
--- a/unittests/CodeGen/MachineOperandTest.cpp
+++ b/unittests/CodeGen/MachineOperandTest.cpp
@@ -11,6 +11,7 @@
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/ModuleSlotTracker.h"
 #include "llvm/Support/raw_ostream.h"
 #include "gtest/gtest.h"
 
@@ -100,4 +101,22 @@ TEST(MachineOperandTest, PrintCImm) {
   ASSERT_TRUE(OS.str() == "i128 18446744073709551616");
 }
 
+TEST(MachineOperandTest, PrintSubRegIndex) {
+  // Create a MachineOperand with an immediate and print it as a subreg index.
+  MachineOperand MO = MachineOperand::CreateImm(3);
+
+  // Checking some preconditions on the newly created
+  // MachineOperand.
+  ASSERT_TRUE(MO.isImm());
+  ASSERT_TRUE(MO.getImm() == 3);
+
+  // Print a MachineOperand containing a SubRegIdx. Here we check that without a
+  // TRI and IntrinsicInfo we can print the operand as a subreg index.
+  std::string str;
+  raw_string_ostream OS(str);
+  ModuleSlotTracker DummyMST(nullptr);
+  MachineOperand::printSubregIdx(OS, MO.getImm(), nullptr);
+  ASSERT_TRUE(OS.str() == "%subreg.3");
+}
+
 } // end namespace