[SandboxIR] Add pointer-diff utility function (#110343)

This will be used when gathering seeds to calculate the lane an
instruction in the bundle uses.

This is a retry of PR 110176, which somehow got messed up while
resolving a merge conflict.

Author: Sterling-Augustine

llvm/include/llvm/SandboxIR/Utils.h

@@ -12,8 +12,12 @@
 #ifndef LLVM_SANDBOXIR_UTILS_H
 #define LLVM_SANDBOXIR_UTILS_H
 
+#include "llvm/Analysis/LoopAccessAnalysis.h"
 #include "llvm/Analysis/MemoryLocation.h"
+#include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/SandboxIR/Instruction.h"
+#include <optional>
 
 namespace llvm::sandboxir {
 
@@ -57,6 +61,38 @@ class Utils {
   memoryLocationGetOrNone(const Instruction *I) {
     return llvm::MemoryLocation::getOrNone(cast<llvm::Instruction>(I->Val));
   }
+
+  /// \Returns the gap between the memory locations accessed by \p I0 and
+  /// \p I1 in bytes.
+  template <typename LoadOrStoreT>
+  static std::optional<int>
+  getPointerDiffInBytes(LoadOrStoreT *I0, LoadOrStoreT *I1, ScalarEvolution &SE,
+                        const DataLayout &DL) {
+    static_assert(std::is_same_v<LoadOrStoreT, LoadInst> ||
+                      std::is_same_v<LoadOrStoreT, StoreInst>,
+                  "Expected sandboxir::Load or sandboxir::Store!");
+    llvm::Value *Opnd0 = I0->getPointerOperand()->Val;
+    llvm::Value *Opnd1 = I1->getPointerOperand()->Val;
+    llvm::Value *Ptr0 = getUnderlyingObject(Opnd0);
+    llvm::Value *Ptr1 = getUnderlyingObject(Opnd1);
+    if (Ptr0 != Ptr1)
+      return false;
+    llvm::Type *ElemTy = llvm::Type::getInt8Ty(SE.getContext());
+    return getPointersDiff(ElemTy, Opnd0, ElemTy, Opnd1, DL, SE,
+                           /*StrictCheck=*/false, /*CheckType=*/false);
+  }
+
+  /// \Returns true if \p I0 accesses a memory location lower than \p I1.
+  /// Returns false if the difference cannot be determined, if the memory
+  /// locations are equal, or if I1 accesses a memory location greater than I0.
+  template <typename LoadOrStoreT>
+  static bool atLowerAddress(LoadOrStoreT *I0, LoadOrStoreT *I1,
+                             ScalarEvolution &SE, const DataLayout &DL) {
+    auto Diff = getPointerDiffInBytes(I0, I1, SE, DL);
+    if (!Diff)
+      return false;
+    return *Diff > 0;
+  }
 };
 } // namespace llvm::sandboxir
 

llvm/unittests/SandboxIR/UtilsTest.cpp

@@ -7,9 +7,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/SandboxIR/Utils.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/BasicAliasAnalysis.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Module.h"
@@ -55,3 +60,77 @@ define void @foo(ptr %arg0) {
   EXPECT_EQ(sandboxir::Utils::memoryLocationGetOrNone(Ld),
             MemoryLocation::getOrNone(LLVMLd));
 }
+
+TEST_F(UtilsTest, GetPointerDiffInBytes) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptr) {
+  %gep0 = getelementptr inbounds float, ptr %ptr, i64 0
+  %gep1 = getelementptr inbounds float, ptr %ptr, i64 1
+  %gep2 = getelementptr inbounds float, ptr %ptr, i64 2
+  %gep3 = getelementptr inbounds float, ptr %ptr, i64 3
+
+  %ld0 = load float, ptr %gep0
+  %ld1 = load float, ptr %gep1
+  %ld2 = load float, ptr %gep2
+  %ld3 = load float, ptr %gep3
+
+  %v2ld0 = load <2 x float>, ptr %gep0
+  %v2ld1 = load <2 x float>, ptr %gep1
+  %v2ld2 = load <2 x float>, ptr %gep2
+  %v2ld3 = load <2 x float>, ptr %gep3
+
+  %v3ld0 = load <3 x float>, ptr %gep0
+  %v3ld1 = load <3 x float>, ptr %gep1
+  %v3ld2 = load <3 x float>, ptr %gep2
+  %v3ld3 = load <3 x float>, ptr %gep3
+  ret void
+}
+)IR");
+  llvm::Function &LLVMF = *M->getFunction("foo");
+  DominatorTree DT(LLVMF);
+  TargetLibraryInfoImpl TLII;
+  TargetLibraryInfo TLI(TLII);
+  DataLayout DL(M->getDataLayout());
+  AssumptionCache AC(LLVMF);
+  BasicAAResult BAA(DL, LLVMF, TLI, AC, &DT);
+  AAResults AA(TLI);
+  AA.addAAResult(BAA);
+  LoopInfo LI(DT);
+  ScalarEvolution SE(LLVMF, TLI, AC, DT, LI);
+  sandboxir::Context Ctx(C);
+
+  auto &F = *Ctx.createFunction(&LLVMF);
+  auto &BB = *F.begin();
+  auto It = std::next(BB.begin(), 4);
+  auto *L0 = cast<sandboxir::LoadInst>(&*It++);
+  auto *L1 = cast<sandboxir::LoadInst>(&*It++);
+  auto *L2 = cast<sandboxir::LoadInst>(&*It++);
+  [[maybe_unused]] auto *L3 = cast<sandboxir::LoadInst>(&*It++);
+
+  auto *V2L0 = cast<sandboxir::LoadInst>(&*It++);
+  auto *V2L1 = cast<sandboxir::LoadInst>(&*It++);
+  auto *V2L2 = cast<sandboxir::LoadInst>(&*It++);
+  auto *V2L3 = cast<sandboxir::LoadInst>(&*It++);
+
+  [[maybe_unused]] auto *V3L0 = cast<sandboxir::LoadInst>(&*It++);
+  auto *V3L1 = cast<sandboxir::LoadInst>(&*It++);
+  [[maybe_unused]] auto *V3L2 = cast<sandboxir::LoadInst>(&*It++);
+  [[maybe_unused]] auto *V3L3 = cast<sandboxir::LoadInst>(&*It++);
+
+  // getPointerDiffInBytes
+  EXPECT_EQ(*sandboxir::Utils::getPointerDiffInBytes(L0, L1, SE, DL), 4);
+  EXPECT_EQ(*sandboxir::Utils::getPointerDiffInBytes(L0, L2, SE, DL), 8);
+  EXPECT_EQ(*sandboxir::Utils::getPointerDiffInBytes(L1, L0, SE, DL), -4);
+  EXPECT_EQ(*sandboxir::Utils::getPointerDiffInBytes(L0, V2L0, SE, DL), 0);
+
+  EXPECT_EQ(*sandboxir::Utils::getPointerDiffInBytes(L0, V2L1, SE, DL), 4);
+  EXPECT_EQ(*sandboxir::Utils::getPointerDiffInBytes(L0, V3L1, SE, DL), 4);
+  EXPECT_EQ(*sandboxir::Utils::getPointerDiffInBytes(V2L0, V2L2, SE, DL), 8);
+  EXPECT_EQ(*sandboxir::Utils::getPointerDiffInBytes(V2L0, V2L3, SE, DL), 12);
+  EXPECT_EQ(*sandboxir::Utils::getPointerDiffInBytes(V2L3, V2L0, SE, DL), -12);
+
+  // atLowerAddress
+  EXPECT_TRUE(sandboxir::Utils::atLowerAddress(L0, L1, SE, DL));
+  EXPECT_FALSE(sandboxir::Utils::atLowerAddress(L1, L0, SE, DL));
+  EXPECT_FALSE(sandboxir::Utils::atLowerAddress(L3, V3L3, SE, DL));
+}