[llvm] c8032c8 - [SandboxIR] Implement helper getLoadStoreAddressSpace() (#194934)

[via llvm-commits]

Author: vporpo
Date: 2026-04-29T12:58:46-07:00
New Revision: c8032c8862af9349899bb5021408c7c6d0960201

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

LOG: [SandboxIR] Implement helper getLoadStoreAddressSpace() (#194934)

This is similar to LLVM IR's getLoadStoreAddressSpace(), but with an `Instruction *` argument instead of a `Value *`.

Added: 
    

Modified: 
    llvm/include/llvm/SandboxIR/Instruction.h
    llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
    llvm/unittests/SandboxIR/SandboxIRTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/SandboxIR/Instruction.h b/llvm/include/llvm/SandboxIR/Instruction.h
index 72c5b16289f1a..9cd446b48cf31 100644
--- a/llvm/include/llvm/SandboxIR/Instruction.h
+++ b/llvm/include/llvm/SandboxIR/Instruction.h
@@ -2746,6 +2746,20 @@ class OpaqueInst : public SingleLLVMInstructionImpl<llvm::Instruction> {
   }
 };
 
+//===----------------------------------------------------------------------===//
+//                          Helper functions
+//===----------------------------------------------------------------------===//
+
+/// A helper function that returns the address space of the pointer operand of
+/// load or store instruction.
+inline unsigned getLoadStoreAddressSpace(const Instruction *I) {
+  assert((isa<LoadInst>(I) || isa<StoreInst>(I)) &&
+         "Expected Load or Store instruction");
+  if (auto *LI = dyn_cast<LoadInst>(I))
+    return LI->getPointerAddressSpace();
+  return cast<StoreInst>(I)->getPointerAddressSpace();
+}
+
 } // namespace llvm::sandboxir
 
 #endif // LLVM_SANDBOXIR_INSTRUCTION_H

diff  --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
index 0ffa88e4fe661..54cda0eaa6a8b 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
@@ -48,14 +48,6 @@ SeedCollection::SeedCollection(StringRef Pipeline, StringRef AuxArg)
   }
 }
 
-static unsigned getAddressSpace(const Instruction *I) {
-  if (auto *LI = dyn_cast<LoadInst>(I))
-    return LI->getPointerAddressSpace();
-  if (auto *SI = dyn_cast<StoreInst>(I))
-    return SI->getPointerAddressSpace();
-  return 0;
-}
-
 bool SeedCollection::runOnFunction(Function &F, const Analyses &A) {
   bool Change = false;
   const auto &DL = F.getParent()->getDataLayout();
@@ -71,7 +63,7 @@ bool SeedCollection::runOnFunction(Function &F, const Analyses &A) {
           Utils::getNumBits(VecUtils::getElementType(Utils::getExpectedType(
                                 Seeds[Seeds.getFirstUnusedElementIdx()])),
                             DL);
-      unsigned AS = getAddressSpace(Seeds[0]);
+      unsigned AS = getLoadStoreAddressSpace(Seeds[0]);
       unsigned VecRegBits = OverrideVecRegBits != 0
                                 ? OverrideVecRegBits
                                 : A.getTTI().getLoadStoreVecRegBitWidth(AS);

diff  --git a/llvm/unittests/SandboxIR/SandboxIRTest.cpp b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
index 109599c1bbec3..8a035b00f794d 100644
--- a/llvm/unittests/SandboxIR/SandboxIRTest.cpp
+++ b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
@@ -3229,6 +3229,8 @@ define void @foo(ptr %arg0, ptr %arg1) {
   // Check getPointerAddressSpace()
   EXPECT_EQ(NewLd->getPointerAddressSpace(),
             Arg1->getType()->getPointerAddressSpace());
+  // Check helper function getLoadStoreAddressSpace()
+  EXPECT_EQ(getLoadStoreAddressSpace(NewLd), NewLd->getPointerAddressSpace());
   EXPECT_EQ(NewLd->getAlign(), 8);
   EXPECT_EQ(NewLd->getName(), "NewLd");
   // Check create(InsertBefore, IsVolatile=true)
@@ -3299,6 +3301,15 @@ define void @foo(i8 %val, ptr %ptr) {
   // Check getPointerAddressSpace()
   EXPECT_EQ(St->getPointerAddressSpace(),
             Ptr->getType()->getPointerAddressSpace());
+  // Check helper function getLoadStoreAddressSpace(St)
+  EXPECT_EQ(getLoadStoreAddressSpace(St), St->getPointerAddressSpace());
+  EXPECT_EQ(
+      getLoadStoreAddressSpace(const_cast<const sandboxir::StoreInst *>(St)),
+      St->getPointerAddressSpace());
+#ifndef NDEBUG
+  // Check the assertion in getLoadStoreAddressSpace(Ret) if not a load or store
+  EXPECT_DEATH(getLoadStoreAddressSpace(Ret), ".*Expected.*");
+#endif
   // Check getAlign()
   EXPECT_EQ(St->getAlign(), 64);
   // Check create(InsertBefore)