[llvm] [SandboxIR] Implement getNumBits for Instructions (PR #110748)

[via llvm-commits]

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

This also moves the other getNumbits test into UtilsTest.cc where it belongs.

>From c7a4448aec897b98369f9566d911ecaae53de960 Mon Sep 17 00:00:00 2001
From: Sterling Augustine <saugustine at google.com>
Date: Tue, 1 Oct 2024 14:45:40 -0700
Subject: [PATCH] [SandboxIR] Implement getNumBits for Instructions

---
 llvm/include/llvm/SandboxIR/Instruction.h  |  3 ++
 llvm/include/llvm/SandboxIR/Utils.h        |  6 ++++
 llvm/unittests/SandboxIR/SandboxIRTest.cpp | 20 -------------
 llvm/unittests/SandboxIR/UtilsTest.cpp     | 34 ++++++++++++++++++++++
 4 files changed, 43 insertions(+), 20 deletions(-)

diff --git a/llvm/include/llvm/SandboxIR/Instruction.h b/llvm/include/llvm/SandboxIR/Instruction.h
index f5f5bb5c4443c2..1ab65831fbc8b7 100644
--- a/llvm/include/llvm/SandboxIR/Instruction.h
+++ b/llvm/include/llvm/SandboxIR/Instruction.h
@@ -97,6 +97,9 @@ class Instruction : public User {
 
   const char *getOpcodeName() const { return getOpcodeName(Opc); }
 
+  const DataLayout &getDataLayout() const {
+    return cast<llvm::Instruction>(Val)->getModule()->getDataLayout();
+  }
   // Note that these functions below are calling into llvm::Instruction.
   // A sandbox IR instruction could introduce a new opcode that could change the
   // behavior of one of these functions. It is better that these functions are
diff --git a/llvm/include/llvm/SandboxIR/Utils.h b/llvm/include/llvm/SandboxIR/Utils.h
index 17fc837f555b8e..c25edd7dc74e0f 100644
--- a/llvm/include/llvm/SandboxIR/Utils.h
+++ b/llvm/include/llvm/SandboxIR/Utils.h
@@ -56,6 +56,12 @@ class Utils {
     return DL.getTypeSizeInBits(Ty->LLVMTy);
   }
 
+  /// \Returns the number of bits required to represent the operands or
+  /// return value of \p I.
+  static unsigned getNumBits(Instruction *I) {
+    return I->getDataLayout().getTypeSizeInBits(getExpectedType(I)->LLVMTy);
+  }
+
   /// Equivalent to MemoryLocation::getOrNone(I).
   static std::optional<llvm::MemoryLocation>
   memoryLocationGetOrNone(const Instruction *I) {
diff --git a/llvm/unittests/SandboxIR/SandboxIRTest.cpp b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
index 3bd520f3174c26..885ed9c44775d8 100644
--- a/llvm/unittests/SandboxIR/SandboxIRTest.cpp
+++ b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
@@ -1512,26 +1512,6 @@ define void @bar(float %v, ptr %ptr) {
   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) {
diff --git a/llvm/unittests/SandboxIR/UtilsTest.cpp b/llvm/unittests/SandboxIR/UtilsTest.cpp
index 7ce8f40f8985bf..44b23d866cec5b 100644
--- a/llvm/unittests/SandboxIR/UtilsTest.cpp
+++ b/llvm/unittests/SandboxIR/UtilsTest.cpp
@@ -135,3 +135,37 @@ define void @foo(ptr %ptr) {
   EXPECT_FALSE(sandboxir::Utils::atLowerAddress(L1, L0, SE, DL));
   EXPECT_FALSE(sandboxir::Utils::atLowerAddress(L3, V3L3, SE, DL));
 }
+
+TEST_F(UtilsTest, 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 via the Value overload
+  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);
+
+  // getNumBits for scalars via the Instruction overload
+  EXPECT_EQ(
+      sandboxir::Utils::getNumBits(cast<sandboxir::Instruction>(F->getArg(0))),
+      DL.getTypeSizeInBits(Type::getFloatTy(C)));
+  EXPECT_EQ(
+      sandboxir::Utils::getNumBits(cast<sandboxir::Instruction>(F->getArg(1))),
+      DL.getTypeSizeInBits(Type::getDoubleTy(C)));
+  EXPECT_EQ(
+      sandboxir::Utils::getNumBits(cast<sandboxir::Instruction>(F->getArg(2))),
+      8u);
+  EXPECT_EQ(
+      sandboxir::Utils::getNumBits(cast<sandboxir::Instruction>(F->getArg(3))),
+      64u);
+}