[llvm] [AArch64][PeepholeOpt] Use correct register class for subreg uses fro… (PR #226595)

[David Green via llvm-commits]

https://github.com/davemgreen created https://github.com/llvm/llvm-project/pull/226595

isCoalescableExtInstr from an instruction like SBFMXri can reuse the def of the instruction. For some targets this is expected to be a gpr32->gpr64 extend, but for AArch64 and PPC is gpr64->gpr64 and we are UseSrcSubIdx and are expecting to generate a subreg copy. The regclass used was for the full register though, not the subreg, generating an invalid copy and INSERT_SUBREG operand (in this case).

Fixes #226352

>From 8cb3bc266b5196af75baa79522204162ffc5a4ea Mon Sep 17 00:00:00 2001
From: David Green <david.green at arm.com>
Date: Fri, 25 Sep 2026 21:34:46 +0100
Subject: [PATCH] [AArch64][PeepholeOpt] Use correct register class for subreg
 uses from optimizeExtInstr.

isCoalescableExtInstr from an instruction like SBFMXri can reuse the def of the
instruction. For some targets this is expected to be a gpr32->gpr64 extend, but
for AArch64 and PPC is gpr64->gpr64 and we are UseSrcSubIdx and are expecting
to generate a subreg copy. The regclass used was for the full register though,
not the subreg, generating an invalid copy and INSERT_SUBREG operand (in this
case).

Fixes #226352
---
 llvm/lib/CodeGen/PeepholeOptimizer.cpp        | 16 ++++--
 .../CodeGen/AArch64/coalescableextinstr.mir   | 51 +++++++++++++++++++
 .../CodeGen/PowerPC/peephole-subreg-def.mir   |  2 +-
 3 files changed, 63 insertions(+), 6 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/coalescableextinstr.mir

diff --git a/llvm/lib/CodeGen/PeepholeOptimizer.cpp b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
index 5edfc97bab30cb..aeea9ca581d310 100644
--- a/llvm/lib/CodeGen/PeepholeOptimizer.cpp
+++ b/llvm/lib/CodeGen/PeepholeOptimizer.cpp
@@ -922,17 +922,23 @@ bool PeepholeOptimizer::optimizeExtInstr(
       // %6:gprc_and_gprc_nor0 = COPY %1.sub_32:g8rc_and_g8rc_nox0
       // %3:gprc_and_gprc_nor0 = COPY %6:gprc_and_gprc_nor0
       //
-      if (UseSrcSubIdx)
-        RC = MRI->getRegClass(UseMI->getOperand(0).getReg());
+      if (UseSrcSubIdx) {
+        RC = MRI->getRegClass(UseMO->getReg());
+        if (UseMO->getSubReg())
+          RC = TRI->getSubRegisterClass(RC, UseMO->getSubReg());
+      }
 
       Register NewVR = MRI->createVirtualRegister(RC);
-      BuildMI(*UseMBB, UseMI, UseMI->getDebugLoc(),
-              TII->get(TargetOpcode::COPY), NewVR)
-          .addReg(DstReg, {}, SubIdx);
+      [[maybe_unused]] auto Copy = BuildMI(*UseMBB, UseMI, UseMI->getDebugLoc(),
+                                           TII->get(TargetOpcode::COPY), NewVR)
+                                       .addReg(DstReg, {}, SubIdx);
+      LLVM_DEBUG(dbgs() << "  Build new copy: " << *Copy);
+      LLVM_DEBUG(dbgs() << "  Changing: " << *UseMI);
       if (UseSrcSubIdx)
         UseMO->setSubReg(0);
 
       UseMO->setReg(NewVR);
+      LLVM_DEBUG(dbgs() << "        to: " << *UseMI);
       ++NumReuse;
       Changed = true;
     }
diff --git a/llvm/test/CodeGen/AArch64/coalescableextinstr.mir b/llvm/test/CodeGen/AArch64/coalescableextinstr.mir
new file mode 100644
index 00000000000000..9d972b80ba3fb0
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/coalescableextinstr.mir
@@ -0,0 +1,51 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 6
+# RUN: llc -mtriple=aarch64 -run-pass=peephole-opt -verify-machineinstrs -o - %s | FileCheck %s
+
+---
+name:            test
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: test
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT:   liveins: $x0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   %0:gpr32 = COPY $x0
+  ; CHECK-NEXT:   B %bb.2
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   %1:gpr64all = IMPLICIT_DEF
+  ; CHECK-NEXT:   %7:gpr32 = COPY %5.sub_32
+  ; CHECK-NEXT:   %2:gpr64 = INSERT_SUBREG %1, %7, %subreg.sub_32
+  ; CHECK-NEXT:   %4:gpr64common = ANDXri %5, 4096
+  ; CHECK-NEXT:   $w0 = COPY %2
+  ; CHECK-NEXT:   RET_ReallyLR implicit $w0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   $x0 = COPY %0
+  ; CHECK-NEXT:   %6:gpr64 = COPY $x0
+  ; CHECK-NEXT:   %3:gpr32all = COPY %6.sub_32
+  ; CHECK-NEXT:   %5:gpr64 = SBFMXri %6, 0, 31
+  ; CHECK-NEXT:   B %bb.1
+  bb.0:
+    liveins: $x0
+    %6:gpr32 = COPY $x0
+
+    B %bb.2
+
+  bb.1:
+    %3:gpr64all = IMPLICIT_DEF
+    %4:gpr64 = INSERT_SUBREG %3, %1, %subreg.sub_32
+    %5:gpr64common = ANDXri %2, 4096
+    $w0 = COPY %4
+    RET_ReallyLR implicit $w0
+
+  bb.2:
+    $x0 = COPY %6
+    %0:gpr64 = COPY $x0
+    %1:gpr32all = COPY %0.sub_32
+    %2:gpr64 = SBFMXri %0, 0, 31
+    B %bb.1
+...
diff --git a/llvm/test/CodeGen/PowerPC/peephole-subreg-def.mir b/llvm/test/CodeGen/PowerPC/peephole-subreg-def.mir
index 2a234e83f8e9a1..42dec448afc7e2 100644
--- a/llvm/test/CodeGen/PowerPC/peephole-subreg-def.mir
+++ b/llvm/test/CodeGen/PowerPC/peephole-subreg-def.mir
@@ -24,7 +24,7 @@ body:             |
     ; CHECK-NEXT: [[EXTSW:%[0-9]+]]:g8rc_and_g8rc_nox0 = EXTSW [[ADDI8_]]
     ; CHECK-NEXT: [[LI8_:%[0-9]+]]:g8rc = LI8 0
     ; CHECK-NEXT: STB8 [[LI8_]], 0, [[EXTSW]]
-    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gprc_and_gprc_nor0 = COPY [[EXTSW]].sub_32
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gprc = COPY [[EXTSW]].sub_32
     ; CHECK-NEXT: [[COPY2:%[0-9]+]]:gprc_and_gprc_nor0 = COPY [[EXTSW]].sub_32
     ; CHECK-NEXT: [[ADDI:%[0-9]+]]:gprc = ADDI killed [[COPY2]], 1
     ; CHECK-NEXT: [[EXTSW_32_64_:%[0-9]+]]:g8rc_and_g8rc_nox0 = EXTSW_32_64 killed [[ADDI]]