[llvm] [AArch64] Fix MatchDup Lane Out Of Range In AArch64 (PR #101275)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-globalisel

Author: None (cceerczw)

<details>
<summary>Changes</summary>

The original code is intended to process pattern which ISEL generated. The purpose of this pattern is duplicate a scalar value to vector register which behavior like AArch64's Dup Inst.

See Url https://reviews.llvm.org/D81979 && https://reviews.llvm.org/D81221

The current code considers only the preceding situation which just duplicate from Shuffle's LHS but does not consider the user code. RHS should be considered.

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


2 Files Affected:

- (modified) llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp (+7) 
- (modified) llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-lowering-shuffle-splat.mir (+28) 


``````````diff
diff --git a/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp b/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp
index 4a1977ba1a00f..a727e7f14062a 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp
@@ -288,10 +288,17 @@ bool matchDupFromBuildVector(int Lane, MachineInstr &MI,
                              MachineRegisterInfo &MRI,
                              ShuffleVectorPseudo &MatchInfo) {
   assert(Lane >= 0 && "Expected positive lane?");
+  int NumElements = MRI.getType(MI.getOperand(1).getReg()).getNumElements();
   // Test if the LHS is a BUILD_VECTOR. If it is, then we can just reference the
   // lane's definition directly.
   auto *BuildVecMI = getOpcodeDef(TargetOpcode::G_BUILD_VECTOR,
                                   MI.getOperand(1).getReg(), MRI);
+  // If Lane >= NumElements then it is point to RHS, just check from RHS
+  if (NumElements <= Lane) {
+    BuildVecMI = getOpcodeDef(TargetOpcode::G_BUILD_VECTOR,
+                              MI.getOperand(2).getReg(), MRI);
+    Lane -= NumElements;
+  }
   if (!BuildVecMI)
     return false;
   Register Reg = BuildVecMI->getOperand(Lane + 1).getReg();
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-lowering-shuffle-splat.mir b/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-lowering-shuffle-splat.mir
index 9d12c3c32c7f8..ef2160df92f2b 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-lowering-shuffle-splat.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-lowering-shuffle-splat.mir
@@ -367,3 +367,31 @@ body:             |
     %shuf:_(<4 x s32>) = G_SHUFFLE_VECTOR %buildvec(<4 x s32>), %undef, shufflemask(0, 0, 0, 0)
     $q0 = COPY %shuf(<4 x s32>)
     RET_ReallyLR implicit $q0
+ 
+...
+---
+name:            build_vector_rhs
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.1.entry:
+    liveins: $w0, $w1, $w2, $w3
+    ; The G_SHUFFLE_VECTOR is fed by a G_BUILD_VECTOR, and the 0th input
+    ; operand is not a constant. We should get a G_DUP.
+    ;
+    ; CHECK-LABEL: name: build_vector
+    ; CHECK: liveins: $w0, $w1, $w2, $w3
+    ; CHECK: %lane:_(s32) = COPY $w0
+    ; CHECK: %shuf:_(<4 x s32>) = G_DUP %lane(s32)
+    ; CHECK: $q0 = COPY %shuf(<4 x s32>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %lane:_(s32) = COPY $w0
+    %b:_(s32) = COPY $w1
+    %c:_(s32) = COPY $w2
+    %d:_(s32) = COPY $w3
+    %buildvec0:_(<4 x s32>) = G_BUILD_VECTOR %lane(s32), %b(s32), %c(s32), %d(s32)
+    %buildvec1:_(<4 x s32>) = G_BUILD_VECTOR %lane(s32), %b(s32), %c(s32), %d(s32)
+    %shuf:_(<4 x s32>) = G_SHUFFLE_VECTOR %buildvec0(<4 x s32>), %buildvec1, shufflemask(4, 4, 4, 4)
+    $q0 = COPY %shuf(<4 x s32>)
+    RET_ReallyLR implicit $q0

``````````

</details>


https://github.com/llvm/llvm-project/pull/101275