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

[via llvm-commits]

Author: Owen Anderson
Date: 2024-02-21T00:42:22-05:00
New Revision: 44b717df4d837ce4e8d76b00cee2e122ae6ad28c

URL: https://github.com/llvm/llvm-project/commit/44b717df4d837ce4e8d76b00cee2e122ae6ad28c
DIFF: https://github.com/llvm/llvm-project/commit/44b717df4d837ce4e8d76b00cee2e122ae6ad28c.diff

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

This avoid turning a poison value into a segfault, and fixes
https://github.com/llvm/llvm-project/issues/78383

Added: 
    llvm/test/CodeGen/AArch64/extractvector-oob-load.mir

Modified: 
    llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index e5b229fcd54f58..044cd3d2d426ec 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -3971,14 +3971,18 @@ 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 +4001,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
+
+...