[llvm] [SandboxIR] Functions to find vectorizor-relevant properties (PR #109221)

[via llvm-commits]

https://github.com/Sterling-Augustine created https://github.com/llvm/llvm-project/pull/109221

When vectorizing, the destintation type and value of stores is more relevant than the type of the instruction itself. Similarly for return instructions. These functions provide a convenient way to do that without special-casing them everywhere, and avoids the need for friending any class that needs access to Value::LLVMTy to calculate it.

Open to better naming.

>From 6ae7ad9ad77da43002e03a75a45ec383653d4983 Mon Sep 17 00:00:00 2001
From: Sterling Augustine <saugustine at google.com>
Date: Wed, 18 Sep 2024 16:46:11 -0700
Subject: [PATCH] [SandboxIR] Functions to find vectorizor-relevant instruction
 properties.

When vectorizing, the destintation type and value of stores is more
relevant than the type of the instruction itself. Similarly for return
instructions. These functions provide a convenient way to do that
without special-casing them everywhere, and avoids the need for
friending any class that needs access to Value::LLVMTy to calculate it.

Open to better naming.
---
 llvm/include/llvm/SandboxIR/SandboxIR.h    | 10 +++++++
 llvm/lib/SandboxIR/SandboxIR.cpp           | 20 +++++++++++++
 llvm/unittests/SandboxIR/SandboxIRTest.cpp | 34 ++++++++++++++++++++++
 3 files changed, 64 insertions(+)

diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index 2fdbbbd094650f..4772f95a900dd1 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -411,6 +411,11 @@ class Value {
 
   Type *getType() const;
 
+  /// \Returns the expected type of \p V. For most Values this is equivalent
+  /// to getType, but for stores returns the stored type, rather than void,
+  /// and for ReturnInsts returns the returned type.
+  Type *getExpectedType() const;
+
   Context &getContext() const { return Ctx; }
 
   void replaceUsesWithIf(Value *OtherV,
@@ -1422,6 +1427,11 @@ class Instruction : public sandboxir::User {
   /// LangRef.html for the meaning of these flags.
   void copyFastMathFlags(FastMathFlags FMF);
 
+  /// \Returns the expected Value for this instruction. For most instructions,
+  /// this is the instruction it self, but for stores returns the stored
+  /// operand, and for ReturnInstructions returns the returned value.
+  Value *getExpectedValue() const;
+
 #ifndef NDEBUG
   void dumpOS(raw_ostream &OS) const override;
 #endif
diff --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
index 18fdcda15a1a91..8852aba7318d35 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -137,6 +137,18 @@ unsigned Value::getNumUses() const { return range_size(Val->users()); }
 
 Type *Value::getType() const { return Ctx.getType(Val->getType()); }
 
+Type *Value::getExpectedType() const {
+  if (auto *I = dyn_cast<sandboxir::Instruction>(this)) {
+    // A Return's value operand can be null if it returns void.
+    if (auto *RI = dyn_cast<sandboxir::ReturnInst>(I)) {
+      if (RI->getReturnValue() == nullptr)
+        return RI->getType();
+    }
+    return I->getExpectedValue()->getType();
+  }
+  return getType();
+}
+
 void Value::replaceUsesWithIf(
     Value *OtherV, llvm::function_ref<bool(const Use &)> ShouldReplace) {
   assert(getType() == OtherV->getType() && "Can't replace with different type");
@@ -571,6 +583,14 @@ void Instruction::setHasApproxFunc(bool B) {
   cast<llvm::Instruction>(Val)->setHasApproxFunc(B);
 }
 
+llvm::sandboxir::Value *Instruction::getExpectedValue() const {
+  if (auto *SI = dyn_cast<StoreInst>(this))
+    return SI->getValueOperand();
+  if (auto *RI = dyn_cast<ReturnInst>(this))
+    return RI->getReturnValue();
+  return const_cast<Instruction *>(this);
+}
+
 #ifndef NDEBUG
 void Instruction::dumpOS(raw_ostream &OS) const {
   OS << "Unimplemented! Please override dump().";
diff --git a/llvm/unittests/SandboxIR/SandboxIRTest.cpp b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
index b76d24dc297b96..055b09eed61d5c 100644
--- a/llvm/unittests/SandboxIR/SandboxIRTest.cpp
+++ b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
@@ -875,6 +875,8 @@ OperandNo: 0
   EXPECT_TRUE(I0->hasNUses(1u));
   EXPECT_FALSE(I0->hasNUses(2u));
 
+  // Check Value.getExpectedType
+
   // Check User.setOperand().
   Ret->setOperand(0, Arg0);
   EXPECT_EQ(Ret->getOperand(0), Arg0);
@@ -947,6 +949,38 @@ define i32 @foo(i32 %arg0, i32 %arg1) {
   EXPECT_EQ(Glob0->getOperand(0), Glob1);
 }
 
+TEST_F(SandboxIRTest, GetExpected) {
+  parseIR(C, R"IR(
+define float @foo(float %v, ptr %ptr) {
+  %add = fadd float %v, %v
+  store float %v, ptr %ptr
+  ret float %v
+}
+)IR");
+  llvm::Function &LLVMF = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  llvm::BasicBlock *LLVMBB = &*LLVMF.begin();
+
+  Ctx.createFunction(&LLVMF);
+  auto *BB = cast<sandboxir::BasicBlock>(Ctx.getValue(LLVMBB));
+  auto It = BB->begin();
+  auto Add = cast<sandboxir::Instruction>(&*It++);
+  auto *S0 = cast<sandboxir::Instruction>(&*It++);
+  auto *Ret = cast<sandboxir::Instruction>(&*It++);
+  // Instruction.getExpectedValue
+  EXPECT_EQ(Add->getExpectedValue(), Add);
+  EXPECT_EQ(S0->getExpectedValue(),
+            cast<sandboxir::StoreInst>(S0)->getValueOperand());
+  EXPECT_EQ(Ret->getExpectedValue(),
+            cast<sandboxir::ReturnInst>(Ret)->getReturnValue());
+  // Value.getExpectedType
+  EXPECT_EQ(Add->getExpectedType(), Add->getType());
+  EXPECT_EQ(S0->getExpectedType(),
+            cast<sandboxir::StoreInst>(S0)->getValueOperand()->getType());
+  EXPECT_EQ(Ret->getExpectedType(),
+            cast<sandboxir::ReturnInst>(Ret)->getReturnValue()->getType());
+}
+
 TEST_F(SandboxIRTest, RAUW_RUWIf) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr) {