[llvm] [SandboxVec][SeedCollector] Change vector bit width calculation (PR #194090)

[via llvm-commits]

https://github.com/vporpo updated https://github.com/llvm/llvm-project/pull/194090

>From 040305ba972d758e927bb7e0eb6da55cd809a0fe Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vasileios.porpodas at amd.com>
Date: Fri, 24 Apr 2026 23:16:34 +0000
Subject: [PATCH 1/2] [SandboxVec][SeedCollector] Change vector bit width
 calculation

On GPUs TTI::getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector)
returns the elemnt bit width, not the whole vector size as it does on the CPU.
So this patch changes this to a call to .getLoadStoreVecRegBitWidth(), which
depends on the address space, so it also moves the calculation per seed.

This patch also adds an AMDGPU lit test directory with a simple test.
---
 .../Passes/SeedCollection.cpp                 | 18 ++++++++++-----
 .../SandboxVectorizer/AMDGPU/basic.ll         | 22 +++++++++++++++++++
 .../SandboxVectorizer/AMDGPU/lit.local.cfg    |  2 ++
 3 files changed, 36 insertions(+), 6 deletions(-)
 create mode 100644 llvm/test/Transforms/SandboxVectorizer/AMDGPU/basic.ll
 create mode 100644 llvm/test/Transforms/SandboxVectorizer/AMDGPU/lit.local.cfg

diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
index 0a6c086be259a..32bfb3c815e6d 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
@@ -48,15 +48,17 @@ SeedCollection::SeedCollection(StringRef Pipeline, StringRef AuxArg)
   }
 }
 
+static unsigned getAddressSpace(Instruction *I) {
+  if (auto *LI = dyn_cast<LoadInst>(I))
+    return LI->getPointerOperand()->getType()->getPointerAddressSpace();
+  if (auto *SI = dyn_cast<StoreInst>(I))
+    return SI->getPointerOperand()->getType()->getPointerAddressSpace();
+  return 0;
+}
+
 bool SeedCollection::runOnFunction(Function &F, const Analyses &A) {
   bool Change = false;
   const auto &DL = F.getParent()->getDataLayout();
-  unsigned VecRegBits =
-      OverrideVecRegBits != 0
-          ? OverrideVecRegBits
-          : A.getTTI()
-                .getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector)
-                .getFixedValue();
   bool CollectStores = CollectSeeds.find(StoreSeedsDef) != std::string::npos;
   bool CollectLoads = CollectSeeds.find(LoadSeedsDef) != std::string::npos;
 
@@ -69,6 +71,10 @@ bool SeedCollection::runOnFunction(Function &F, const Analyses &A) {
           Utils::getNumBits(VecUtils::getElementType(Utils::getExpectedType(
                                 Seeds[Seeds.getFirstUnusedElementIdx()])),
                             DL);
+      unsigned AS = getAddressSpace(Seeds[0]);
+      unsigned VecRegBits = OverrideVecRegBits != 0
+                                ? OverrideVecRegBits
+                                : A.getTTI().getLoadStoreVecRegBitWidth(AS);
 
       auto DivideBy2 = [](unsigned Num) {
         auto Floor = VecUtils::getFloorPowerOf2(Num);
diff --git a/llvm/test/Transforms/SandboxVectorizer/AMDGPU/basic.ll b/llvm/test/Transforms/SandboxVectorizer/AMDGPU/basic.ll
new file mode 100644
index 0000000000000..e6f75fcc48c6f
--- /dev/null
+++ b/llvm/test/Transforms/SandboxVectorizer/AMDGPU/basic.ll
@@ -0,0 +1,22 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -passes=sandbox-vectorizer -sbvec-passes="seed-collection<load-store-vec>" -mtriple=amdgcn--amdhsa -mcpu=gfx900 %s -S | FileCheck %s
+
+define void @basic(ptr %ptr0) {
+; CHECK-LABEL: define void @basic(
+; CHECK-SAME: ptr [[PTR0:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[VECIINITL:%.*]] = load <2 x float>, ptr [[PTR0]], align 1, !sandboxvec [[META0:![0-9]+]]
+; CHECK-NEXT:    store <2 x float> [[VECIINITL]], ptr [[PTR0]], align 1, !sandboxvec [[META0]]
+; CHECK-NEXT:    ret void
+;
+  %ptr1 = getelementptr inbounds i8, ptr %ptr0, i64 4
+
+  %ld0 = load float, ptr %ptr0
+  %ld1 = load float, ptr %ptr1
+
+  store float %ld0, ptr %ptr0
+  store float %ld1, ptr %ptr1
+  ret void
+}
+;.
+; CHECK: [[META0]] = distinct !{!"sandboxregion"}
+;.
diff --git a/llvm/test/Transforms/SandboxVectorizer/AMDGPU/lit.local.cfg b/llvm/test/Transforms/SandboxVectorizer/AMDGPU/lit.local.cfg
new file mode 100644
index 0000000000000..7c492428aec76
--- /dev/null
+++ b/llvm/test/Transforms/SandboxVectorizer/AMDGPU/lit.local.cfg
@@ -0,0 +1,2 @@
+if not "AMDGPU" in config.root.targets:
+    config.unsupported = True

>From a2579bd763eccf74faf70bd276c904dd42bccfcc Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vasileios.porpodas at amd.com>
Date: Tue, 28 Apr 2026 23:11:11 +0000
Subject: [PATCH 2/2] fixup! [SandboxVec][SeedCollector] Change vector bit
 width calculation

---
 .../Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp     | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
index 32bfb3c815e6d..fc28c5bdf8217 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
@@ -50,9 +50,9 @@ SeedCollection::SeedCollection(StringRef Pipeline, StringRef AuxArg)
 
 static unsigned getAddressSpace(Instruction *I) {
   if (auto *LI = dyn_cast<LoadInst>(I))
-    return LI->getPointerOperand()->getType()->getPointerAddressSpace();
+    return LI->getPointerAddressSpace();
   if (auto *SI = dyn_cast<StoreInst>(I))
-    return SI->getPointerOperand()->getType()->getPointerAddressSpace();
+    return SI->getPointerAddressSpace();
   return 0;
 }