From d028b1721e8fdfd21b1b0581b0e2e1b1f5f7726d Mon Sep 17 00:00:00 2001
From: Roman Levenstein <rlevenstein@meta.com>
Date: Sun, 28 Apr 2024 21:54:10 -0700
Subject: [PATCH] Fix how mangled C++ functions are handled

Reviewed By: jaybean-dev

Differential Revision: D56683446

fbshipit-source-id: d9cfec00f5d9260a241b6deaebbe4c168216586c
---
 lib/LLVMIRCodeGen/LLVMIRGen.cpp | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/lib/LLVMIRCodeGen/LLVMIRGen.cpp b/lib/LLVMIRCodeGen/LLVMIRGen.cpp
index dd7b80582..20a1831c6 100644
--- a/lib/LLVMIRCodeGen/LLVMIRGen.cpp
+++ b/lib/LLVMIRCodeGen/LLVMIRGen.cpp
@@ -795,7 +795,7 @@ static std::string createName(const std::string &name, ElemKind elemTy) {
 void LLVMIRGen::initLLVMFunctionNameToMangledNameMap() {
   CHECK(llvmFunctionNameToMangledName_.empty());
   constexpr size_t maxFnBaseNameLen = 4096;
-  char *fnNameBuf = static_cast<char *>(std::malloc(maxFnBaseNameLen));
+  char fnNameBuf[maxFnBaseNameLen];
   // Build a map from names to the list of matching mangled names.
   for (llvm::Function &F : getModule()) {
     auto mangledName = F.getName().str();
@@ -805,6 +805,17 @@ void LLVMIRGen::initLLVMFunctionNameToMangledNameMap() {
       continue;
     }
     size_t fnNameLen = maxFnBaseNameLen;
+    size_t fnContextLen = maxFnBaseNameLen;
+    // Skip C++ functions that have names like a::b::c. It helps to avoid name
+    // conflicts with kernels that may be called just c and conflict with C++
+    // functions.
+    fnNameBuf[0] = '\0';
+    char *contextNamePtr =
+        Mangler.getFunctionDeclContextName(fnNameBuf, &fnContextLen);
+    if (contextNamePtr && fnContextLen != 0 && contextNamePtr[0]) {
+      continue;
+    }
+    fnNameBuf[0] = '\0';
     char *demangledNamePtr = Mangler.getFunctionBaseName(fnNameBuf, &fnNameLen);
     if (!demangledNamePtr || fnNameLen == 0) {
       continue;
@@ -817,10 +828,6 @@ void LLVMIRGen::initLLVMFunctionNameToMangledNameMap() {
     }
     llvmFunctionNameToMangledName_[demangledFnName].push_back(mangledName);
   }
-  // Free up the memory.
-  if (fnNameBuf) {
-    free(fnNameBuf);
-  }
   DEBUG_GLOW({
     // Dump the map for debugging purposes.
     llvm::dbgs() << "Mapping between function names and matching LLVM function "