From 1bab8fce685c6ed865b9d7b78a39be20ef57751d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E7=BB=B4=E6=94=BF?=
 <15849716+huang-weizheng19@user.noreply.gitee.com>
Date: Thu, 22 May 2025 22:02:58 +0800
Subject: [PATCH] [AArch64][GISel] Fix MatchDup Lane Out Of Range In AArch64
 Signed-off-by: huang-weizheng19 <huangweizheng4@h-partners.com>

---
 .../GISel/AArch64PostLegalizerLowering.cpp    | 11 ++++++++
 .../postlegalizer-lowering-shuffle-splat.mir  | 28 +++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp b/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp
index eab1de94e9c8..b8736667471e 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp
@@ -341,10 +341,21 @@ static 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 7ceb2e2dbe30..7079681b16e9 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-lowering-shuffle-splat.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-lowering-shuffle-splat.mir
@@ -356,3 +356,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
-- 
Gitee