[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?

[michalpaszkowski]

@mshelego Thanks for the PR and sorry for waiting this long! LGTM! I can merge the change by the end of this week. @arsenm or anyone else working on GlobalISel and/or AMDGPU, please let me know if anything should be addressed before merging.

[github-actions]

@mshelego Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!