[llvm] [llubi] Fix GEP return type on the fast path (PR #214235)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/214235

Closes https://github.com/llvm/llvm-project/issues/214144.


>From f2ae8305b1dc395af81f2559c683c138511b7831 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Wed, 5 Aug 2026 21:57:25 +0800
Subject: [PATCH] [llubi] Fix GEP return type on the fast path

---
 llvm/test/tools/llubi/gep.ll     |  3 +++
 llvm/tools/llubi/lib/Context.cpp | 18 ++++++++++++------
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/llvm/test/tools/llubi/gep.ll b/llvm/test/tools/llubi/gep.ll
index 8cfa63b105b4b..16daa878e6880 100644
--- a/llvm/test/tools/llubi/gep.ll
+++ b/llvm/test/tools/llubi/gep.ll
@@ -90,6 +90,8 @@ define void @main() {
   %gep_nuw_valid4 = getelementptr nuw i32, ptr %alloc, i64 1073741821
   %gep_nuw_invalid4 = getelementptr nuw i32, ptr %alloc, i64 1073741822
 
+  %gep_splat_zero = getelementptr i8, ptr %alloc, <2 x i64> zeroinitializer
+
   ret void
 }
 ; CHECK: Entering function: main
@@ -142,5 +144,6 @@ define void @main() {
 ; CHECK-NEXT:   %gep_nuw_invalid3 = getelementptr nuw [2 x i16], ptr null, i32 1073741823, i32 2 => poison
 ; CHECK-NEXT:   %gep_nuw_valid4 = getelementptr nuw i32, ptr %alloc, i64 1073741821 => ptr 0xFFFFFFFC [alloc + 4294967284]
 ; CHECK-NEXT:   %gep_nuw_invalid4 = getelementptr nuw i32, ptr %alloc, i64 1073741822 => poison
+; CHECK-NEXT:   %gep_splat_zero = getelementptr i8, ptr %alloc, <2 x i64> zeroinitializer => { ptr 0x8 [alloc], ptr 0x8 [alloc] }
 ; CHECK-NEXT:   ret void
 ; CHECK-NEXT: Exiting function: main
diff --git a/llvm/tools/llubi/lib/Context.cpp b/llvm/tools/llubi/lib/Context.cpp
index 5055b8d1af462..f83f7fbbd8a47 100644
--- a/llvm/tools/llubi/lib/Context.cpp
+++ b/llvm/tools/llubi/lib/Context.cpp
@@ -1011,12 +1011,8 @@ Context::computeGEP(GEPOperator &GEP,
     AccumulatedOffset =
         AnyValue::getVectorSplat(AccumulatedOffset, Res.asAggregate().size());
   auto ApplyScaledOffset = [&](const AnyValue &Index, const APInt &Scale) {
-    if (Index.isAggregate() && !Res.isAggregate()) {
-      Res = AnyValue::getVectorSplat(Res, Index.asAggregate().size());
-      AccumulatedOffset = AnyValue::getVectorSplat(AccumulatedOffset,
-                                                   Index.asAggregate().size());
-    }
-    if (Index.isAggregate() && Res.isAggregate()) {
+    if (Index.isAggregate()) {
+      assert(Res.isAggregate() && "Res must be splatted before.");
       for (auto &&[ResElem, IndexElem, OffsetElem] :
            zip(Res.asAggregate(), Index.asAggregate(),
                AccumulatedOffset.asAggregate()))
@@ -1041,6 +1037,16 @@ Context::computeGEP(GEPOperator &GEP,
        GTI != GTE; ++GTI) {
     Value *V = GTI.getOperand();
 
+    // If the index is a vector, make sure the accumulated pointer is also a
+    // vector. Otherwise, make it a vector splat.
+    if (auto *VTy = dyn_cast<VectorType>(V->getType());
+        VTy && !Res.isAggregate()) {
+      uint32_t NumElements = getEVL(VTy->getElementCount());
+      Res = AnyValue::getVectorSplat(Res, NumElements);
+      AccumulatedOffset =
+          AnyValue::getVectorSplat(AccumulatedOffset, NumElements);
+    }
+
     // Fast path for zero offsets.
     if (auto *CI = dyn_cast<ConstantInt>(V)) {
       if (CI->isZero())