[llvm] [ValueTracking] Fix Overflow with i1 Constant GEPs (PR #125470)

[Pierre van Houtryve via llvm-commits]

https://github.com/Pierre-vh created https://github.com/llvm/llvm-project/pull/125470

The LoadStoreVectorizer can cause ValueTracking to crash with I1 GEPs. ValueTracking creates a 1 bit APInt and then tries to multiply it.

This changes the minimum width of those APInts to 8 bits to avoid the issue.

Fixes SWDEV-507697

>From c6404254c1a13bd4bc821f5a4128805b17c46710 Mon Sep 17 00:00:00 2001
From: pvanhout <pierre.vanhoutryve at amd.com>
Date: Mon, 3 Feb 2025 10:21:36 +0100
Subject: [PATCH] [ValueTracking] Fix Overflow with i1 Constant GEPs

The LoadStoreVectorizer can cause ValueTracking to crash with I1 GEPs.
ValueTracking creates a 1 bit APInt and then tries
to multiply it.

This changes the minimum width of those APInts to 8 bits to avoid the issue.

Fixes SWDEV-507697
---
 llvm/lib/Analysis/ValueTracking.cpp           |  7 +++++--
 .../AMDGPU/knownbits-gep-i1.ll                | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/LoadStoreVectorizer/AMDGPU/knownbits-gep-i1.ll

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index 6b61a3546e8b7c5..b76afbf0a7249bd 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -1477,8 +1477,11 @@ static void computeKnownBitsFromOperator(const Operator *I,
         // that this is a multiple of the minimum size.
         ScalingFactor.Zero.setLowBits(llvm::countr_zero(TypeSizeInBytes));
       } else if (IndexBits.isConstant()) {
-        APInt IndexConst = IndexBits.getConstant();
-        APInt ScalingFactor(IndexBitWidth, TypeSizeInBytes);
+        // i1 is a valid GEP index, ensure we have enough space to do the
+        // computation in that case.
+        unsigned CalcBitWidth = std::max(IndexBitWidth, 8u);
+        APInt IndexConst = IndexBits.getConstant().zext(CalcBitWidth);
+        APInt ScalingFactor(CalcBitWidth, TypeSizeInBytes);
         IndexConst *= ScalingFactor;
         AccConstIndices += IndexConst.sextOrTrunc(BitWidth);
         continue;
diff --git a/llvm/test/Transforms/LoadStoreVectorizer/AMDGPU/knownbits-gep-i1.ll b/llvm/test/Transforms/LoadStoreVectorizer/AMDGPU/knownbits-gep-i1.ll
new file mode 100644
index 000000000000000..a2dc00fbb700b3e
--- /dev/null
+++ b/llvm/test/Transforms/LoadStoreVectorizer/AMDGPU/knownbits-gep-i1.ll
@@ -0,0 +1,19 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -mcpu=gfx942 -passes=load-store-vectorizer -S -o - %s | FileCheck %s
+
+define amdgpu_kernel void @simple_users_scores() {
+; CHECK-LABEL: define amdgpu_kernel void @simple_users_scores(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[SIMPLEUSER:%.*]] = alloca [4 x i64], i32 0, align 4, addrspace(5)
+; CHECK-NEXT:    [[G:%.*]] = getelementptr i32, ptr addrspace(5) [[SIMPLEUSER]], i1 true
+; CHECK-NEXT:    store <2 x i32> zeroinitializer, ptr addrspace(5) [[G]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %simpleuser = alloca [4 x i64], i32 0, align 4, addrspace(5)
+  store i32 0, ptr addrspace(5) %simpleuser, align 4
+  %G = getelementptr i32, ptr addrspace(5) %simpleuser, i1 true
+  store i32 0, ptr addrspace(5) %G, align 4
+  ret void
+}