[AMDGPU][GlobalISel] Fix G_UNMERGE_VALUES combine

[mshelego]

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.

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-llvm-globalisel

@llvm/pr-subscribers-backend-amdgpu

Author: Maksim Shelegov (mshelego)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/141812.diff

2 Files Affected:

• (modified) llvm/include/llvm/CodeGen/GlobalISel/LegalizationArtifactCombiner.h (+6-2)
• (modified) llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-unmerge-values.mir (+17)

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..ce50040b2e726 100644
--- a/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-unmerge-values.mir
+++ b/llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-unmerge-values.mir
@@ -1090,3 +1090,20 @@ 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: {{  $}}
+    %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)
+...

[mshelego]

@arsenm could you please review?