[llvm] [GlobalISel] Clamp out-of-range G_EXTRACT_VECTOR_ELT constant indices when converting them into loads. (PR #82460)

[Owen Anderson via llvm-commits]

https://github.com/resistor updated https://github.com/llvm/llvm-project/pull/82460

>From b5db7ffcf64e78501ca1c91532914127174dbbfa Mon Sep 17 00:00:00 2001
From: Owen Anderson <resistor at mac.com>
Date: Tue, 20 Feb 2024 23:43:42 -0500
Subject: [PATCH 1/2] [GlobalISel] Clamp out-of-range G_EXTRACT_VECTOR_ELT
 constant indices when converting them into loads.

This avoid turning a poison value into a segfault, and fixes
https://github.com/llvm/llvm-project/issues/78383
---
 .../CodeGen/GlobalISel/LegalizerHelper.cpp    | 19 ++++++----
 .../AArch64/extractvector-oob-load.mir        | 38 +++++++++++++++++++
 2 files changed, 50 insertions(+), 7 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/extractvector-oob-load.mir

diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index e5b229fcd54f58..dacb13d685f33e 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -3971,14 +3971,19 @@ LegalizerHelper::createStackTemporary(TypeSize Bytes, Align Alignment,
   return MIRBuilder.buildFrameIndex(FramePtrTy, FrameIdx);
 }
 
-static Register clampDynamicVectorIndex(MachineIRBuilder &B, Register IdxReg,
-                                        LLT VecTy) {
-  int64_t IdxVal;
-  if (mi_match(IdxReg, *B.getMRI(), m_ICst(IdxVal)))
-    return IdxReg;
-
+static Register clampVectorIndex(MachineIRBuilder &B, Register IdxReg,
+                                 LLT VecTy) {
   LLT IdxTy = B.getMRI()->getType(IdxReg);
   unsigned NElts = VecTy.getNumElements();
+
+  int64_t IdxVal;
+  if (mi_match(IdxReg, *B.getMRI(), m_ICst(IdxVal))) {
+    if (IdxVal < VecTy.getNumElements()) {
+      return IdxReg;
+    }
+    // If a constant index would be out of bounds, clamp it as well.
+  }
+
   if (isPowerOf2_32(NElts)) {
     APInt Imm = APInt::getLowBitsSet(IdxTy.getSizeInBits(), Log2_32(NElts));
     return B.buildAnd(IdxTy, IdxReg, B.buildConstant(IdxTy, Imm)).getReg(0);
@@ -3997,7 +4002,7 @@ Register LegalizerHelper::getVectorElementPointer(Register VecPtr, LLT VecTy,
   assert(EltSize * 8 == EltTy.getSizeInBits() &&
          "Converting bits to bytes lost precision");
 
-  Index = clampDynamicVectorIndex(MIRBuilder, Index, VecTy);
+  Index = clampVectorIndex(MIRBuilder, Index, VecTy);
 
   LLT IdxTy = MRI.getType(Index);
   auto Mul = MIRBuilder.buildMul(IdxTy, Index,
diff --git a/llvm/test/CodeGen/AArch64/extractvector-oob-load.mir b/llvm/test/CodeGen/AArch64/extractvector-oob-load.mir
new file mode 100644
index 00000000000000..e8c5819e75e090
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/extractvector-oob-load.mir
@@ -0,0 +1,38 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
+# RUN: llc -mtriple=aarch64-linux-gnu -run-pass=aarch64-prelegalizer-combiner -verify-machineinstrs %s -o - | FileCheck %s
+
+---
+name:            f
+alignment:       4
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: _ }
+  - { id: 1, class: _ }
+  - { id: 2, class: _ }
+  - { id: 3, class: _ }
+liveins:
+  - { reg: '$x0' }
+frameInfo:
+  maxAlignment:    1
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $x0
+
+    ; CHECK-LABEL: name: f
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(p0) = COPY $x0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 16
+    ; CHECK-NEXT: [[PTR_ADD:%[0-9]+]]:_(p0) = G_PTR_ADD [[COPY]], [[C]](s64)
+    ; CHECK-NEXT: [[LOAD:%[0-9]+]]:_(s64) = G_LOAD [[PTR_ADD]](p0) :: (load (s64))
+    ; CHECK-NEXT: $x0 = COPY [[LOAD]](s64)
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %0:_(p0) = COPY $x0
+    %3:_(s64) = G_CONSTANT i64 224567957
+    %1:_(<3 x s64>) = G_LOAD %0(p0) :: (load (<3 x s64>), align 32)
+    %2:_(s64) = G_EXTRACT_VECTOR_ELT %1(<3 x s64>), %3(s64)
+    $x0 = COPY %2(s64)
+    RET_ReallyLR implicit $x0
+
+...

>From 4982a4897da315a7a8cbaf58ae9e54600b8ead2e Mon Sep 17 00:00:00 2001
From: Owen Anderson <resistor at mac.com>
Date: Wed, 21 Feb 2024 00:07:33 -0500
Subject: [PATCH 2/2] Fix brace nit.

---
 llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index dacb13d685f33e..044cd3d2d426ec 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -3978,9 +3978,8 @@ static Register clampVectorIndex(MachineIRBuilder &B, Register IdxReg,
 
   int64_t IdxVal;
   if (mi_match(IdxReg, *B.getMRI(), m_ICst(IdxVal))) {
-    if (IdxVal < VecTy.getNumElements()) {
+    if (IdxVal < VecTy.getNumElements())
       return IdxReg;
-    }
     // If a constant index would be out of bounds, clamp it as well.
   }