[AIE] Fix sub-register index composition in split instruction rewriter

llvm/lib/Target/AIE/AIESplitInstructionRewriter.cpp

@@ -47,7 +47,8 @@ class AIESplitInstrBuilder : public MachineFunctionPass {
 private:
   /// Add the atomic sub-regs of \p InitialMO as new operands to \p MIB.
   void addRegOperands(MachineInstrBuilder &MIB, const MachineOperand &InitialMO,
-                      ArrayRef<SubRegSplit> NewSubRegs);
+                      ArrayRef<SubRegSplit> NewSubRegs,
+                      const TargetRegisterInfo &TRI);
 
   /// Replace \p MI with a new instruction where the tuple operands are
   /// rewritten into multiple operands for their different sub-registers
@@ -79,11 +80,18 @@ bool AIESplitInstrBuilder::runOnMachineFunction(MachineFunction &MF) {
 
 void AIESplitInstrBuilder::addRegOperands(MachineInstrBuilder &MIB,
                                           const MachineOperand &InitialMO,
-                                          ArrayRef<SubRegSplit> NewSubRegs) {
+                                          ArrayRef<SubRegSplit> NewSubRegs,
+                                          const TargetRegisterInfo &TRI) {
   assert(!NewSubRegs.empty());
   for (const SubRegSplit &SRS : NewSubRegs) {
     MachineOperand NewMO = InitialMO;
-    NewMO.setSubReg(SRS.SubReg);
+    unsigned NewSubReg = SRS.SubReg;
+    // If the original operand already has a sub-register index, compose it
+    // with the new sub-register. E.g. if the operand is %R.sub_hi_dim and
+    // we split into sub_mod, the result should be sub_hi_dim_then_sub_mod.
+    if (unsigned ExistingSubReg = InitialMO.getSubReg())
+      NewSubReg = TRI.composeSubRegIndices(ExistingSubReg, NewSubReg);
+    NewMO.setSubReg(NewSubReg);
     if (SRS.IsUndef)
       NewMO.setIsUndef();
     MIB->addOperand(NewMO);
@@ -98,6 +106,7 @@ void AIESplitInstrBuilder::rewriteInstruction(
   MachineBasicBlock &MBB = *MI.getParent();
   MachineFunction &MF = *MBB.getParent();
   auto *TII = MF.getSubtarget().getInstrInfo();
+  auto *TRI = MF.getSubtarget().getRegisterInfo();
 
   unsigned NumInitialOps = MI.getNumOperands();
   MachineInstrBuilder MIB =
@@ -108,7 +117,7 @@ void AIESplitInstrBuilder::rewriteInstruction(
     const MachineOperand &MO = MI.getOperand(OpIdx);
     if (const OperandSubRegMapping *SRM = OperandsInfo.findOperandInfo(OpIdx);
         SRM && !SRM->SubRegsSplit.empty()) {
-      addRegOperands(MIB, MO, SRM->SubRegsSplit);
+      addRegOperands(MIB, MO, SRM->SubRegsSplit, *TRI);
     } else if (!MO.isImplicit()) {
       // Note implicit ops are added by default by BuildMI.
       MIB->addOperand(MO);

llvm/test/CodeGen/AIE/aie2ps/ra/split-instrs-create-subreg-compose.mir

@@ -0,0 +1,33 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+#
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates
+# RUN: llc -O2 -mtriple=aie2ps -run-pass=aie-split-instrs-create %s -o - | FileCheck %s
+
+# Verify that when a tuple operand already has a sub-register index,
+# the split correctly composes the indices. E.g., %R.sub_hi_dim split
+# into sub_mod should produce sub_hi_dim_then_sub_mod, not sub_mod.
+
+---
+name:            test_split_2D_with_subreg_hi_dim
+alignment:       16
+legalized:       true
+regBankSelected: true
+selected:        true
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $p0, $d1_3d
+    ; CHECK-LABEL: name: test_split_2D_with_subreg_hi_dim
+    ; CHECK: liveins: $p0, $d1_3d
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:ep = COPY $p0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:eds = COPY $d1_3d
+    ; CHECK-NEXT: dead [[COPY:%[0-9]+]]:ep, dead [[COPY1:%[0-9]+]].sub_hi_dim_then_sub_dim_count:eds = PADD_2D_pseudo_split killed [[COPY]], killed [[COPY1]].sub_hi_dim_then_sub_mod, killed [[COPY1]].sub_hi_dim_then_sub_dim_size, killed [[COPY1]].sub_hi_dim_then_sub_dim_stride, killed [[COPY1]].sub_hi_dim_then_sub_dim_count
+    %20:ep = COPY $p0
+    %100:eds = COPY $d1_3d
+    dead %20:ep, dead %100.sub_hi_dim_then_sub_dim_count:eds = PADD_2D_pseudo killed %20, killed %100.sub_hi_dim
+...