[llvm] [AMDGPU][GlobalISel] Fix G_UNMERGE_VALUES combine (PR #141812)

[Maksim Shelegov via llvm-commits]

https://github.com/mshelego updated https://github.com/llvm/llvm-project/pull/141812

>From 92c4ecdb338d8a28c543b30767a3461e8d52b42b Mon Sep 17 00:00:00 2001
From: "Shelegov, Maksim" <maksim.shelegov at intel.com>
Date: Wed, 28 May 2025 19:37:08 +0200
Subject: [PATCH] [AMDGPU][GlobalISel] Fix G_UNMERGE_VALUES combine

Fixes #139791.

When trying to combine two G_UNMERGE_VALUES with a COPY between them,
a crash occurs in tryCombineUnmergeValues() because getDefIndex() tries
to find the index of the original source register in the def found by
getDefIgnoringCopies(). In the case of a COPY in between, this register
is not present in the def, only in the COPY.
---
 .../GlobalISel/LegalizationArtifactCombiner.h |  8 +++++--
 .../GlobalISel/legalize-unmerge-values.mir    | 21 +++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/GlobalISel/LegalizationArtifactCombiner.h b/llvm/include/llvm/CodeGen/GlobalISel/LegalizationArtifactCombiner.h
index 3712a7fa06d9a..676b69e04c34e 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/LegalizationArtifactCombiner.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/LegalizationArtifactCombiner.h
@@ -1064,13 +1064,17 @@ class LegalizationArtifactCombiner {
                                GISelChangeObserver &Observer) {
     unsigned NumDefs = MI.getNumDefs();
     Register SrcReg = MI.getSourceReg();
-    MachineInstr *SrcDef = getDefIgnoringCopies(SrcReg, MRI);
+    std::optional<DefinitionAndSourceRegister> DefSrcReg =
+        getDefSrcRegIgnoringCopies(SrcReg, MRI);
+    if (!DefSrcReg)
+      return false;
+    MachineInstr *SrcDef = DefSrcReg->MI;
     if (!SrcDef)
       return false;
 
     LLT OpTy = MRI.getType(SrcReg);
     LLT DestTy = MRI.getType(MI.getReg(0));
-    unsigned SrcDefIdx = getDefIndex(*SrcDef, SrcReg);
+    unsigned SrcDefIdx = getDefIndex(*SrcDef, DefSrcReg->Reg);
 
     Builder.setInstrAndDebugLoc(MI);
 
diff --git a/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-unmerge-values.mir b/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-unmerge-values.mir
index c231aa8334d45..ee57b72a8b4e9 100644
--- a/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-unmerge-values.mir
+++ b/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-unmerge-values.mir
@@ -1090,3 +1090,24 @@ body: |
     $vgpr9_vgpr10_vgpr11 = COPY %8
 
 ...
+
+---
+name: test_unmerge_through_copy
+body: |
+  bb.0:
+    liveins: $vgpr0
+
+    ; CHECK-LABEL: name: test_unmerge_through_copy
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(s32) = COPY $vgpr0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 255
+    ; CHECK-NEXT: [[AND:%[0-9]+]]:_(s32) = G_AND [[COPY]], [[C]]
+    ; CHECK-NEXT: $vgpr0 = COPY [[AND]](s32)
+    %0:_(s32) = COPY $vgpr0
+    %1:_(s16), %2:_(s16) = G_UNMERGE_VALUES %0:_(s32)
+    %3:_(s16) = COPY %1:_(s16)
+    %4:_(s8), %5:_(s8) = G_UNMERGE_VALUES %3:_(s16)
+    %6:_(s32) = G_ZEXT %4:_(s8)
+    $vgpr0 = COPY %6:_(s32)
+...