[llvm] [SandboxIR] createFunction() should always create a function (PR #124665)

[via llvm-commits]

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

This patch removes the assertion that checks for an existing function. If one exists it will remove it and create a new one. This helps remove a crash when a function declaration object already exists and we are about to create a SandboxIR object for the definition.

>From 72e78332b0c6f0023e220dcc55596ce23bfbc006 Mon Sep 17 00:00:00 2001
From: Vasileios Porpodas <vporpodas at google.com>
Date: Mon, 27 Jan 2025 17:01:33 -0800
Subject: [PATCH] [SandboxIR] createFunction() should always create a function

This patch removes the assertion that checks for an existing function.
If one exists it will remove it and create a new one.
This helps remove a crash when a function declaration object already exists
and we are about to create a SandboxIR object for the definition.
---
 llvm/lib/SandboxIR/Context.cpp                |  5 ++++-
 .../SandboxVectorizer/SandboxVectorizer.cpp   |  5 ++++-
 llvm/unittests/SandboxIR/SandboxIRTest.cpp    | 20 +++++++++++++++++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/SandboxIR/Context.cpp b/llvm/lib/SandboxIR/Context.cpp
index 42ca456881fd0d..440210f5a1bf7b 100644
--- a/llvm/lib/SandboxIR/Context.cpp
+++ b/llvm/lib/SandboxIR/Context.cpp
@@ -628,12 +628,15 @@ Module *Context::getOrCreateModule(llvm::Module *LLVMM) {
 }
 
 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());
 
+  // There may be a function declaration already defined. Regardless destroy it.
+  if (Function *ExistingF = cast_or_null<Function>(getValue(F)))
+    detach(ExistingF);
+
   auto NewFPtr = std::unique_ptr<Function>(new Function(F, *this));
   auto *SBF = cast<Function>(registerValue(std::move(NewFPtr)));
   // Create arguments.
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
index a6e2b40000529a..542fcde71e83c6 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
@@ -87,5 +87,8 @@ bool SandboxVectorizerPass::runImpl(Function &LLVMF) {
   // Create SandboxIR for LLVMF and run BottomUpVec on it.
   sandboxir::Function &F = *Ctx->createFunction(&LLVMF);
   sandboxir::Analyses A(*AA, *SE, *TTI);
-  return FPM.runOnFunction(F, A);
+  bool Change = FPM.runOnFunction(F, A);
+  // TODO: This is a function pass, so we won't be needing the function-level
+  // Sandbox IR objects in the future. So we should clear them.
+  return Change;
 }
diff --git a/llvm/unittests/SandboxIR/SandboxIRTest.cpp b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
index 9d1c86a9b9c727..9eeac9b60372f9 100644
--- a/llvm/unittests/SandboxIR/SandboxIRTest.cpp
+++ b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
@@ -6080,3 +6080,23 @@ TEST_F(SandboxIRTest, InstructionCallbacks) {
   EXPECT_THAT(Removed, testing::IsEmpty());
   EXPECT_THAT(Moved, testing::IsEmpty());
 }
+
+TEST_F(SandboxIRTest, FunctionObjectAlreadyExists) {
+  parseIR(C, R"IR(
+define void @foo() {
+  call void @bar()
+  ret void
+}
+define void @bar() {
+  ret void
+}
+)IR");
+  Function &LLVMFoo = *M->getFunction("foo");
+  Function &LLVMBar = *M->getFunction("bar");
+  sandboxir::Context Ctx(C);
+  // This will create a Function object for @bar().
+  Ctx.createFunction(&LLVMFoo);
+  EXPECT_NE(Ctx.getValue(&LLVMBar), nullptr);
+  // This should not crash, even though there is already a value for LLVMBar.
+  Ctx.createFunction(&LLVMBar);
+}