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

[via llvm-commits]

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

>From 24c57a538031f82aa5dbeda25b0b360dd803f903 Mon Sep 17 00:00:00 2001
From: chengzhiwei <chengzhiwei6 at huawei.com>
Date: Thu, 18 Jul 2024 15:11:33 +0800
Subject: [PATCH] [AArch64] Fix MatchDup Lane Out Of Range In AArch64

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.
---
 .../GISel/AArch64PostLegalizerLowering.cpp    | 10 +++++--
 .../postlegalizer-lowering-shuffle-splat.mir  | 29 +++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp b/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp
index 4a1977ba1a00f..afea8bdec1973 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64PostLegalizerLowering.cpp
@@ -288,10 +288,16 @@ 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);
+  auto *BuildVecMI =
+      getOpcodeDef(TargetOpcode::G_BUILD_VECTOR,
+                   MI.getOperand(Lane < NumElements ? 1 : 2).getReg(), MRI);
+  // If Lane >= NumElements then it is point to RHS, just check from RHS
+  if (NumElements <= Lane)
+    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..71094825e42f3 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,32 @@ 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, $w4
+    ; 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, $w4
+    ; CHECK: %lane_1:_(s32) = COPY $w1
+    ; CHECK: %shuf:_(<4 x s32>) = G_DUP %lane_1(s32)
+    ; CHECK: $q0 = COPY %shuf(<4 x s32>)
+    ; CHECK: RET_ReallyLR implicit $q0
+    %lane_0:_(s32) = COPY $w0
+    %lane_1:_(s32) = COPY $w1
+    %b:_(s32) = COPY $w2
+    %c:_(s32) = COPY $w3
+    %d:_(s32) = COPY $w4
+    %buildvec0:_(<4 x s32>) = G_BUILD_VECTOR %lane_0(s32), %b(s32), %c(s32), %d(s32)
+    %buildvec1:_(<4 x s32>) = G_BUILD_VECTOR %lane_1(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