[llvm] [SandboxIR] Implement Module (PR #109716)

via llvm-commits llvm-commits at lists.llvm.org
Mon Sep 23 13:42:56 PDT 2024


https://github.com/vporpo created https://github.com/llvm/llvm-project/pull/109716

This patch implements sandboxir::Module.
It provides access to globals.

>From ad271ac883693d81e29e4c9b94c2a7368d09baab Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Fri, 23 Aug 2024 16:36:59 -0700
Subject: [PATCH] [SandboxIR] Implement Module

This patch implements sandboxir::Module.
It provides access to globals.
---
 llvm/include/llvm/SandboxIR/Module.h       | 92 ++++++++++++++++++++++
 llvm/include/llvm/SandboxIR/SandboxIR.h    | 24 +++++-
 llvm/include/llvm/SandboxIR/Tracker.h      |  3 +
 llvm/include/llvm/SandboxIR/Type.h         |  3 +
 llvm/lib/SandboxIR/CMakeLists.txt          |  1 +
 llvm/lib/SandboxIR/Module.cpp              | 40 ++++++++++
 llvm/lib/SandboxIR/SandboxIR.cpp           | 39 +++++++++
 llvm/unittests/SandboxIR/SandboxIRTest.cpp | 58 ++++++++++++++
 8 files changed, 259 insertions(+), 1 deletion(-)
 create mode 100644 llvm/include/llvm/SandboxIR/Module.h
 create mode 100644 llvm/lib/SandboxIR/Module.cpp

diff --git a/llvm/include/llvm/SandboxIR/Module.h b/llvm/include/llvm/SandboxIR/Module.h
new file mode 100644
index 00000000000000..429bb04539bcbb
--- /dev/null
+++ b/llvm/include/llvm/SandboxIR/Module.h
@@ -0,0 +1,92 @@
+//===- Module.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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SANDBOXIR_MODULE_H
+#define LLVM_SANDBOXIR_MODULE_H
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/IR/Module.h"
+#include <string>
+
+namespace llvm {
+
+class DataLayout;
+
+namespace sandboxir {
+
+class Context;
+class Function;
+class GlobalVariable;
+class Type;
+class Constant;
+class GlobalAlias;
+class GlobalIFunc;
+
+/// In SandboxIR the Module is mainly used to access the list of global objects.
+class Module {
+  llvm::Module &LLVMM;
+  Context &Ctx;
+
+  Module(llvm::Module &LLVMM, Context &Ctx) : LLVMM(LLVMM), Ctx(Ctx) {}
+  friend class Context; // For constructor.
+
+public:
+  Context &getContext() const { return Ctx; }
+
+  Function *getFunction(StringRef Name) const;
+
+  const DataLayout &getDataLayout() const { return LLVMM.getDataLayout(); }
+
+  const std::string &getSourceFileName() const {
+    return LLVMM.getSourceFileName();
+  }
+
+  /// Look up the specified global variable in the module symbol table. If it
+  /// does not exist, return null. If AllowInternal is set to true, this
+  /// function will return types that have InternalLinkage. By default, these
+  /// types are not returned.
+  GlobalVariable *getGlobalVariable(StringRef Name, bool AllowInternal) const;
+  GlobalVariable *getGlobalVariable(StringRef Name) const {
+    return getGlobalVariable(Name, /*AllowInternal=*/false);
+  }
+  /// Return the global variable in the module with the specified name, of
+  /// arbitrary type. This method returns null if a global with the specified
+  /// name is not found.
+  GlobalVariable *getNamedGlobal(StringRef Name) const {
+    return getGlobalVariable(Name, true);
+  }
+
+  // TODO: missing getOrInsertGlobal().
+
+  /// Return the global alias in the module with the specified name, of
+  /// arbitrary type. This method returns null if a global with the specified
+  /// name is not found.
+  GlobalAlias *getNamedAlias(StringRef Name) const;
+
+  /// Return the global ifunc in the module with the specified name, of
+  /// arbitrary type. This method returns null if a global with the specified
+  /// name is not found.
+  GlobalIFunc *getNamedIFunc(StringRef Name) const;
+
+  // TODO: Missing removeGlobalVariable() eraseGlobalVariable(),
+  // insertGlobalVariable()
+
+  // TODO: Missing global_begin(), global_end(), globals().
+
+  // TODO: Missing many other functions.
+
+#ifndef NDEBUG
+  void dumpOS(raw_ostream &OS) const;
+  LLVM_DUMP_METHOD void dump() const;
+#endif // NDEBUG
+};
+
+} // namespace sandboxir
+} // namespace llvm
+
+#endif // LLVM_SANDBOXIR_MODULE_H
diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index 4e751a75196ce1..306fdfcdd05ab3 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -109,6 +109,7 @@
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
+#include "llvm/SandboxIR/Module.h"
 #include "llvm/SandboxIR/Tracker.h"
 #include "llvm/SandboxIR/Type.h"
 #include "llvm/SandboxIR/Use.h"
@@ -138,6 +139,7 @@ class ConstantPtrAuth;
 class ConstantExpr;
 class Context;
 class Function;
+class Module;
 class Instruction;
 class VAArgInst;
 class FreezeInst;
@@ -346,6 +348,7 @@ class Value {
   friend class NoCFIValue;            // For `Val`.
   friend class ConstantPtrAuth;       // For `Val`.
   friend class ConstantExpr;          // For `Val`.
+  friend class Module;                // For `Val`.
 
   /// All values point to the context.
   Context &Ctx;
@@ -1317,7 +1320,10 @@ class GlobalWithNodeAPI : public ParentT {
   GlobalWithNodeAPI(Value::ClassID ID, LLVMParentT *C, Context &Ctx)
       : ParentT(ID, C, Ctx) {}
 
-  // TODO: Missing getParent(). Should be added once Module is available.
+  Module *getParent() const {
+    llvm::Module *LLVMM = cast<LLVMGlobalT>(this->Val)->getParent();
+    return this->Ctx.getModule(LLVMM);
+  }
 
   using iterator = mapped_iterator<
       decltype(static_cast<LLVMGlobalT *>(nullptr)->getIterator()), LLVMGVToGV>;
@@ -4478,6 +4484,9 @@ class Context {
   DenseMap<llvm::Value *, std::unique_ptr<sandboxir::Value>>
       LLVMValueToValueMap;
 
+  /// Maps an LLVM Module to the corresponding sandboxir::Module.
+  DenseMap<llvm::Module *, std::unique_ptr<Module>> LLVMModuleToModuleMap;
+
   /// Type has a protected destructor to prohibit the user from managing the
   /// lifetime of the Type objects. Context is friend of Type, and this custom
   /// deleter can destroy Type.
@@ -4621,6 +4630,10 @@ class Context {
     return getValue(const_cast<llvm::Value *>(V));
   }
 
+  Module *getModule(llvm::Module *LLVMM) const;
+
+  Module *getOrCreateModule(llvm::Module *LLVMM);
+
   Type *getType(llvm::Type *LLVMTy) {
     if (LLVMTy == nullptr)
       return nullptr;
@@ -4634,8 +4647,13 @@ class Context {
   /// Create a sandboxir::Function for an existing LLVM IR \p F, including all
   /// blocks and instructions.
   /// This is the main API function for creating Sandbox IR.
+  /// Note: this will not fully populate its parent module. The only globals
+  /// that will be available are those used within the function.
   Function *createFunction(llvm::Function *F);
 
+  /// Create a sandboxir::Module corresponding to \p LLVMM.
+  Module *createModule(llvm::Module *LLVMM);
+
   /// \Returns the number of values registered with Context.
   size_t getNumValues() const { return LLVMValueToValueMap.size(); }
 };
@@ -4661,6 +4679,10 @@ class Function : public GlobalWithNodeAPI<Function, llvm::Function,
     return From->getSubclassID() == ClassID::Function;
   }
 
+  Module *getParent() {
+    return Ctx.getModule(cast<llvm::Function>(Val)->getParent());
+  }
+
   Argument *getArg(unsigned Idx) const {
     llvm::Argument *Arg = cast<llvm::Function>(Val)->getArg(Idx);
     return cast<Argument>(Ctx.getValue(Arg));
diff --git a/llvm/include/llvm/SandboxIR/Tracker.h b/llvm/include/llvm/SandboxIR/Tracker.h
index 5fc43db82bd707..3e3e539a8c7c16 100644
--- a/llvm/include/llvm/SandboxIR/Tracker.h
+++ b/llvm/include/llvm/SandboxIR/Tracker.h
@@ -64,6 +64,9 @@ class SwitchInst;
 class ConstantInt;
 class ShuffleVectorInst;
 class CmpInst;
+class Module;
+class GlobalVariable;
+
 /// The base class for IR Change classes.
 class IRChangeBase {
 protected:
diff --git a/llvm/include/llvm/SandboxIR/Type.h b/llvm/include/llvm/SandboxIR/Type.h
index f99f80967797c1..829c9f3c72125f 100644
--- a/llvm/include/llvm/SandboxIR/Type.h
+++ b/llvm/include/llvm/SandboxIR/Type.h
@@ -31,6 +31,7 @@ class IntegerType;
 class FunctionType;
 class ArrayType;
 class StructType;
+class Module;
 #define DEF_INSTR(ID, OPCODE, CLASS) class CLASS;
 #define DEF_CONST(ID, CLASS) class CLASS;
 #include "llvm/SandboxIR/SandboxIRValues.def"
@@ -57,6 +58,8 @@ class Type {
   friend class CmpInst;            // For LLVMTy. TODO: Cleanup after
                                    // sandboxir::VectorType is more complete.
   friend class Utils;              // for LLVMTy
+  friend class TargetExtType;      // For LLVMTy.
+  friend class Module;             // For LLVMTy.
 
   // Friend all instruction classes because `create()` functions use LLVMTy.
 #define DEF_INSTR(ID, OPCODE, CLASS) friend class CLASS;
diff --git a/llvm/lib/SandboxIR/CMakeLists.txt b/llvm/lib/SandboxIR/CMakeLists.txt
index 03474be0c7b80a..c8c055da6a0c96 100644
--- a/llvm/lib/SandboxIR/CMakeLists.txt
+++ b/llvm/lib/SandboxIR/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_llvm_component_library(LLVMSandboxIR
+  Module.cpp
   Pass.cpp
   PassManager.cpp
   SandboxIR.cpp
diff --git a/llvm/lib/SandboxIR/Module.cpp b/llvm/lib/SandboxIR/Module.cpp
new file mode 100644
index 00000000000000..7510f621556d41
--- /dev/null
+++ b/llvm/lib/SandboxIR/Module.cpp
@@ -0,0 +1,40 @@
+//===- Module.cpp ---------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/SandboxIR/Module.h"
+#include "llvm/SandboxIR/SandboxIR.h"
+
+using namespace llvm::sandboxir;
+
+Function *Module::getFunction(StringRef Name) const {
+  llvm::Function *LLVMF = LLVMM.getFunction(Name);
+  return cast_or_null<Function>(Ctx.getValue(LLVMF));
+}
+
+GlobalVariable *Module::getGlobalVariable(StringRef Name,
+                                          bool AllowInternal) const {
+  return cast_or_null<GlobalVariable>(
+      Ctx.getValue(LLVMM.getGlobalVariable(Name, AllowInternal)));
+}
+
+GlobalAlias *Module::getNamedAlias(StringRef Name) const {
+  return cast_or_null<GlobalAlias>(Ctx.getValue(LLVMM.getNamedAlias(Name)));
+}
+
+GlobalIFunc *Module::getNamedIFunc(StringRef Name) const {
+  return cast_or_null<GlobalIFunc>(Ctx.getValue(LLVMM.getNamedIFunc(Name)));
+}
+
+#ifndef NDEBUG
+void Module::dumpOS(raw_ostream &OS) const { OS << LLVMM; }
+
+void Module::dump() const {
+  dumpOS(dbgs());
+  dbgs() << "\n";
+}
+#endif // NDEBUG
diff --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
index b1ed72be1030d6..84b73a92a8dcd4 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -3400,8 +3400,29 @@ Value *Context::getValue(llvm::Value *V) const {
   return nullptr;
 }
 
+Module *Context::getModule(llvm::Module *LLVMM) const {
+  auto It = LLVMModuleToModuleMap.find(LLVMM);
+  if (It != LLVMModuleToModuleMap.end())
+    return It->second.get();
+  return nullptr;
+}
+
+Module *Context::getOrCreateModule(llvm::Module *LLVMM) {
+  auto Pair = LLVMModuleToModuleMap.insert({LLVMM, nullptr});
+  auto It = Pair.first;
+  if (!Pair.second)
+    return It->second.get();
+  It->second = std::unique_ptr<Module>(new Module(*LLVMM, *this));
+  return It->second.get();
+}
+
 Function *Context::createFunction(llvm::Function *F) {
   assert(getValue(F) == nullptr && "Already exists!");
+  // Create the module if needed before we create the new sandboxir::Function.
+  // Note: this won't fully populate the module. The only globals that will be
+  // available will be the ones being used within the function.
+  getOrCreateModule(F->getParent());
+
   auto NewFPtr = std::unique_ptr<Function>(new Function(F, *this));
   auto *SBF = cast<Function>(registerValue(std::move(NewFPtr)));
   // Create arguments.
@@ -3413,6 +3434,24 @@ Function *Context::createFunction(llvm::Function *F) {
   return SBF;
 }
 
+Module *Context::createModule(llvm::Module *LLVMM) {
+  auto *M = getOrCreateModule(LLVMM);
+  // Create the functions.
+  for (auto &LLVMF : *LLVMM)
+    createFunction(&LLVMF);
+  // Create globals.
+  for (auto &Global : LLVMM->globals())
+    getOrCreateValue(&Global);
+  // Create aliases.
+  for (auto &Alias : LLVMM->aliases())
+    getOrCreateValue(&Alias);
+  // Create ifuncs.
+  for (auto &IFunc : LLVMM->ifuncs())
+    getOrCreateValue(&IFunc);
+
+  return M;
+}
+
 Function *BasicBlock::getParent() const {
   auto *BB = cast<llvm::BasicBlock>(Val);
   auto *F = BB->getParent();
diff --git a/llvm/unittests/SandboxIR/SandboxIRTest.cpp b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
index ae143fd5b74ea1..22a68452ec9777 100644
--- a/llvm/unittests/SandboxIR/SandboxIRTest.cpp
+++ b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
@@ -1685,6 +1685,64 @@ void @foo0(i32 %arg0, i32 %arg1) {
 #endif // NDEBUG
 }
 
+TEST_F(SandboxIRTest, Module) {
+  parseIR(C, R"IR(
+ at glob0 = global i32 42
+ at glob1 = global i32 43
+ at internal0 = internal global i32 42
+ at const0 = constant i32 42
+ at alias0 = dso_local alias void(), ptr @foo
+ at ifunc = ifunc void(), ptr @foo
+define void @foo() {
+  ret void
+}
+define void @bar() {
+  ret void
+}
+)IR");
+  llvm::Module *LLVMM = &*M;
+  llvm::Function *LLVMFFoo = &*M->getFunction("foo");
+  llvm::Function *LLVMFBar = &*M->getFunction("bar");
+
+  sandboxir::Context Ctx(C);
+  auto *M = Ctx.createModule(LLVMM);
+  // Check getContext().
+  EXPECT_EQ(&M->getContext(), &Ctx);
+  // Check getFunction().
+  auto *FFoo = M->getFunction("foo");
+  auto *FBar = M->getFunction("bar");
+  EXPECT_EQ(FFoo, Ctx.getValue(LLVMFFoo));
+  EXPECT_EQ(FBar, Ctx.getValue(LLVMFBar));
+  // Check getDataLayout().
+  EXPECT_EQ(&M->getDataLayout(), &LLVMM->getDataLayout());
+  // Check getSourceFileName().
+  EXPECT_EQ(M->getSourceFileName(), LLVMM->getSourceFileName());
+  // Check getGlobalVariable().
+  for (const char *Name : {"global0", "global1", "internal0"})
+    EXPECT_EQ(M->getGlobalVariable(Name),
+              Ctx.getValue(LLVMM->getGlobalVariable(Name)));
+  // Check getGlobalVariable(AllowInternal).
+  {
+    auto *Internal0 = M->getGlobalVariable("internal0", /*AllowInternal=*/true);
+    EXPECT_TRUE(Internal0 != nullptr);
+    EXPECT_EQ(Internal0, Ctx.getValue(LLVMM->getNamedGlobal("internal0")));
+  }
+  // Check getNamedGlobal().
+  {
+    auto *Internal = M->getNamedGlobal("internal0");
+    EXPECT_TRUE(Internal != nullptr);
+    EXPECT_EQ(Internal, Ctx.getValue(LLVMM->getNamedGlobal("internal0")));
+  }
+  // Check getNamedAlias().
+  auto *Alias0 = M->getNamedAlias("alias0");
+  EXPECT_EQ(Alias0, Ctx.getValue(LLVMM->getNamedAlias("alias0")));
+  EXPECT_EQ(M->getNamedAlias("aliasFOO"), nullptr);
+  // Check getNamedIFunc().
+  auto *IFunc0 = M->getNamedIFunc("ifunc0");
+  EXPECT_EQ(IFunc0, Ctx.getValue(LLVMM->getNamedAlias("ifunc0")));
+  EXPECT_EQ(M->getNamedIFunc("ifuncFOO"), nullptr);
+}
+
 TEST_F(SandboxIRTest, BasicBlock) {
   parseIR(C, R"IR(
 define void @foo(i32 %v1) {



More information about the llvm-commits mailing list