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

When vectorizing, the destination 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.

Author: Sterling-Augustine

llvm/include/llvm/SandboxIR/Type.h

@@ -56,6 +56,7 @@ class Type {
   friend class ConstantVector;     // For LLVMTy.
   friend class CmpInst;            // For LLVMTy. TODO: Cleanup after
                                    // sandboxir::VectorType is more complete.
+  friend class Utils;              // for LLVMTy
 
   // Friend all instruction classes because `create()` functions use LLVMTy.
 #define DEF_INSTR(ID, OPCODE, CLASS) friend class CLASS;

llvm/include/llvm/SandboxIR/Utils.h

@@ -0,0 +1,54 @@
+//===- Utils.h --------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Collector for SandboxIR related convenience functions that don't belong in
+// other classes.
+
+#ifndef LLVM_SANDBOXIR_UTILS_H
+#define LLVM_SANDBOXIR_UTILS_H
+
+namespace llvm::sandboxir {
+
+class Utils {
+public:
+  /// \Returns the expected type of \p Value 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.
+  static Type *getExpectedType(const Value *V) {
+    if (auto *I = dyn_cast<Instruction>(V)) {
+      // A Return's value operand can be null if it returns void.
+      if (auto *RI = dyn_cast<ReturnInst>(I)) {
+        if (RI->getReturnValue() == nullptr)
+          return RI->getType();
+      }
+      return getExpectedValue(I)->getType();
+    }
+    return V->getType();
+  }
+
+  /// \Returns the expected Value for this instruction. For most instructions,
+  /// this is the instruction itself, but for stores returns the stored
+  /// operand, and for ReturnInstructions returns the returned value.
+  static Value *getExpectedValue(const Instruction *I) {
+    if (auto *SI = dyn_cast<StoreInst>(I))
+      return SI->getValueOperand();
+    if (auto *RI = dyn_cast<ReturnInst>(I))
+      return RI->getReturnValue();
+    return const_cast<Instruction *>(I);
+  }
+
+  /// \Returns the number of bits required to represent the operands or return
+  /// value of \p V in \p DL.
+  static unsigned getNumBits(Value *V, const DataLayout &DL) {
+    Type *Ty = getExpectedType(V);
+    return DL.getTypeSizeInBits(Ty->LLVMTy);
+  }
+};
+} // namespace llvm::sandboxir
+
+#endif // LLVM_SANDBOXIR_UTILS_H

llvm/unittests/SandboxIR/SandboxIRTest.cpp

@@ -14,6 +14,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Module.h"
+#include "llvm/SandboxIR/Utils.h"
 #include "llvm/Support/SourceMgr.h"
 #include "gmock/gmock-matchers.h"
 #include "gtest/gtest.h"
@@ -1373,6 +1374,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);
@@ -1436,7 +1439,6 @@ define i32 @foo(i32 %arg0, i32 %arg1) {
   Replaced = Ret->replaceUsesOfWith(I0, Arg0);
   EXPECT_TRUE(Replaced);
   EXPECT_EQ(Ret->getOperand(0), Arg0);
-
   // Check RAUW on constant.
   auto *Glob0 = cast<sandboxir::Constant>(I1->getOperand(0));
   auto *Glob1 = cast<sandboxir::Constant>(I2->getOperand(0));
@@ -1445,6 +1447,68 @@ 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
+}
+define void @bar(float %v, ptr %ptr) {
+  ret void
+}
+)IR");
+  llvm::Function &Foo = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+
+  Ctx.createFunction(&Foo);
+  auto *FooBB = cast<sandboxir::BasicBlock>(Ctx.getValue(&*Foo.begin()));
+  auto FooIt = FooBB->begin();
+  auto Add = cast<sandboxir::Instruction>(&*FooIt++);
+  auto *S0 = cast<sandboxir::Instruction>(&*FooIt++);
+  auto *RetF = cast<sandboxir::Instruction>(&*FooIt++);
+  // getExpectedValue
+  EXPECT_EQ(sandboxir::Utils::getExpectedValue(Add), Add);
+  EXPECT_EQ(sandboxir::Utils::getExpectedValue(S0),
+            cast<sandboxir::StoreInst>(S0)->getValueOperand());
+  EXPECT_EQ(sandboxir::Utils::getExpectedValue(RetF),
+            cast<sandboxir::ReturnInst>(RetF)->getReturnValue());
+  // getExpectedType
+  EXPECT_EQ(sandboxir::Utils::getExpectedType(Add), Add->getType());
+  EXPECT_EQ(sandboxir::Utils::getExpectedType(S0),
+            cast<sandboxir::StoreInst>(S0)->getValueOperand()->getType());
+  EXPECT_EQ(sandboxir::Utils::getExpectedType(RetF),
+            cast<sandboxir::ReturnInst>(RetF)->getReturnValue()->getType());
+
+  // getExpectedValue for void returns
+  llvm::Function &Bar = *M->getFunction("bar");
+  Ctx.createFunction(&Bar);
+  auto *BarBB = cast<sandboxir::BasicBlock>(Ctx.getValue(&*Bar.begin()));
+  auto BarIt = BarBB->begin();
+  auto *RetV = cast<sandboxir::Instruction>(&*BarIt++);
+  EXPECT_EQ(sandboxir::Utils::getExpectedValue(RetV), nullptr);
+}
+
+TEST_F(SandboxIRTest, GetNumBits) {
+  parseIR(C, R"IR(
+define void @foo(float %arg0, double %arg1, i8 %arg2, i64 %arg3) {
+bb0:
+  ret void
+}
+)IR");
+  llvm::Function &Foo = *M->getFunction("foo");
+  sandboxir::Context Ctx(C);
+  sandboxir::Function *F = Ctx.createFunction(&Foo);
+  const DataLayout &DL = M->getDataLayout();
+  // getNumBits for scalars
+  EXPECT_EQ(sandboxir::Utils::getNumBits(F->getArg(0), DL),
+            DL.getTypeSizeInBits(Type::getFloatTy(C)));
+  EXPECT_EQ(sandboxir::Utils::getNumBits(F->getArg(1), DL),
+            DL.getTypeSizeInBits(Type::getDoubleTy(C)));
+  EXPECT_EQ(sandboxir::Utils::getNumBits(F->getArg(2), DL), 8u);
+  EXPECT_EQ(sandboxir::Utils::getNumBits(F->getArg(3), DL), 64u);
+}
+
 TEST_F(SandboxIRTest, RAUW_RUWIf) {
   parseIR(C, R"IR(
 define void @foo(ptr %ptr) {