[LLVM] Codegen subroutine call when CallNode::op is GlobalVar (#14901)

* [CodeGen][LLVM] Codegen to generate internal functions

Previously, `CodeGenLLVM` required all TIR PrimFuncs to have the
`kGlobalSymbol` attribute, using its value as the externally-visible
symbol in the generated library.  This commit relaxes that
requirement, using the presence of `kGlobalSymbol` to indicate whether
a function should be exposed externally.  If `kGlobalSymbol` is not
defined, then the symbol name is generated from the name of the
`tvm::GlobalVar` with the prefix `"_internal_"`, and the symbol is not
exposed externally.

Since this does not change the codegen behavior for any function that
was previously supported, this is not a breaking change.

* [Codegen][LLVM] Handle callsite for internal functions

* [UnitTest][LLVM] Added test for LLVM codegen for subroutine

Author: Lunderberg

src/target/llvm/codegen_amdgpu.cc

@@ -87,9 +87,9 @@ class CodeGenAMDGPU : public CodeGenLLVM {
   CodeGenAMDGPU() = default;
   virtual ~CodeGenAMDGPU() = default;
 
-  void AddFunction(const PrimFunc& f) final {
+  void AddFunction(const GlobalVar& gvar, const PrimFunc& f) final {
     // add function as void return value
-    CodeGenLLVM::AddFunctionInternal(f, true);
+    CodeGenLLVM::AddFunctionInternal(gvar, f, true);
     function_->setCallingConv(llvm::CallingConv::AMDGPU_KERNEL);
     std::ostringstream attr;
     attr << "1," << DetectROCMmaxThreadsPerBlock();
@@ -262,11 +262,7 @@ runtime::Module BuildAMDGPU(IRModule mod, Target target) {
 
   cg->Init("TVMAMDGPUModule", llvm_target.get(), NullOpt, false, false);
 
-  cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end(), [](auto& kv) {
-    ICHECK(kv.second->template IsInstance<PrimFuncNode>())
-        << "Can only lower IR Module with PrimFuncs";
-    return Downcast<PrimFunc>(kv.second);
-  });
+  cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end());
 
   llvm::TargetMachine* tm = llvm_target->GetOrCreateTargetMachine();
   const auto* find_rocm_bitcodes = tvm::runtime::Registry::Get("tvm_callback_rocm_bitcode_path");

src/target/llvm/codegen_cpu.cc

@@ -220,18 +220,17 @@ llvm::DISubprogram* CodeGenCPU::CreateDebugFunction(const PrimFunc& f) {
 #endif
 }
 
-void CodeGenCPU::AddFunction(const PrimFunc& f) {
+void CodeGenCPU::AddFunction(const GlobalVar& gvar, const PrimFunc& f) {
 #if TVM_LLVM_VERSION >= 50
   di_subprogram_ = CreateDebugFunction(f);
 #endif
   EmitDebugLocation(f->span);
-  CodeGenLLVM::AddFunction(f);
+  CodeGenLLVM::AddFunction(gvar, f);
   if (f_tvm_register_system_symbol_ != nullptr) {
-    auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
-    ICHECK(global_symbol.defined())
-        << "CodeGenLLVM: Expect PrimFunc to have the global_symbol attribute";
-    export_system_symbols_.emplace_back(
-        std::make_pair(global_symbol.value().operator std::string(), function_));
+    if (auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol)) {
+      export_system_symbols_.emplace_back(
+          std::make_pair(global_symbol.value().operator std::string(), function_));
+    }
   }
   AddDebugInformation(f, function_);
 }

src/target/llvm/codegen_cpu.h

@@ -67,7 +67,7 @@ class CodeGenCPU : public CodeGenLLVM {
   void Init(const std::string& module_name, LLVMTarget* llvm_target,
             Optional<String> system_lib_prefix, bool dynamic_lookup,
             bool target_c_runtime) override;
-  void AddFunction(const PrimFunc& f) override;
+  void AddFunction(const GlobalVar& gvar, const PrimFunc& f) override;
   void AddMainFunction(const std::string& entry_func_name) override;
   std::unique_ptr<llvm::Module> Finish() override;
   void VisitStmt_(const AssertStmtNode* op) override;

src/target/llvm/codegen_hexagon.cc

@@ -547,7 +547,6 @@ runtime::Module BuildHexagon(IRModule mod, Target target) {
 
   auto cg = std::make_unique<CodeGenHexagon>();
 
-  std::vector<PrimFunc> funcs;
   std::string entry_func;
 
   for (auto kv : mod->functions) {
@@ -562,11 +561,10 @@ runtime::Module BuildHexagon(IRModule mod, Target target) {
       ICHECK(global_symbol.defined());
       entry_func = global_symbol.value();
     }
-    funcs.emplace_back(f);
   }
 
   cg->Init("TVMHexagonModule", llvm_target.get(), NullOpt, false, false);
-  cg->AddFunctionsOrdered(funcs.begin(), funcs.end());
+  cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end());
   if (entry_func.length() != 0) {
     cg->AddMainFunction(entry_func);
   }

src/target/llvm/codegen_llvm.cc

@@ -224,7 +224,13 @@ void CodeGenLLVM::InitTarget() {
 #endif  // TVM_LLVM_VERSION >= 60
 }
 
-void CodeGenLLVM::AddFunction(const PrimFunc& f) { this->AddFunctionInternal(f, false); }
+llvm::Function* CodeGenLLVM::DeclareFunction(const GlobalVar& gvar, const PrimFunc& f) {
+  return this->DeclareFunctionInternal(gvar, f, false);
+}
+
+void CodeGenLLVM::AddFunction(const GlobalVar& gvar, const PrimFunc& f) {
+  this->AddFunctionInternal(gvar, f, false);
+}
 
 void CodeGenLLVM::InitFuncState() {
   var_map_.clear();
@@ -234,15 +240,34 @@ void CodeGenLLVM::InitFuncState() {
   analyzer_.reset(new arith::Analyzer());
 }
 
-void CodeGenLLVM::AddFunctionInternal(const PrimFunc& f, bool ret_void) {
-  this->InitFuncState();
+std::tuple<std::string, llvm::Function::LinkageTypes> CodeGenLLVM::GetLinkage(
+    const GlobalVar& gvar, const PrimFunc& func) {
+  if (auto global_symbol = func->GetAttr<String>(tvm::attr::kGlobalSymbol)) {
+    return {global_symbol.value(), llvm::Function::ExternalLinkage};
+  }
+
+  std::string symbol_name = [&]() {
+    std::stringstream ss;
+    ss << "_internal_";
+    ss << gvar->name_hint;
+    return ss.str();
+  }();
+
+  return {symbol_name, llvm::Function::PrivateLinkage};
+}
+
+llvm::Function* CodeGenLLVM::DeclareFunctionInternal(const GlobalVar& gvar, const PrimFunc& func,
+                                                     bool ret_void) {
+  if (auto it = functions_.find(gvar.get()); it != functions_.end()) {
+    return it->second;
+  }
 
-  ICHECK_EQ(f->buffer_map.size(), 0U)
+  ICHECK_EQ(func->buffer_map.size(), 0U)
       << "Cannot codegen function with buffer_map, please lower them first";
 
   std::vector<llvm::Type*> param_types;
-  is_restricted_ = f->HasNonzeroAttr(tir::attr::kNoAlias);
-  for (Var param : f->params) {
+  is_restricted_ = func->HasNonzeroAttr(tir::attr::kNoAlias);
+  for (Var param : func->params) {
     param_types.push_back(GetLLVMType(param));
     if (!is_restricted_ && param.dtype().is_handle()) {
       alias_var_set_.insert(param.get());
@@ -254,17 +279,26 @@ void CodeGenLLVM::AddFunctionInternal(const PrimFunc& f, bool ret_void) {
   llvm::FunctionType* ftype =
       llvm::FunctionType::get(ret_void ? t_void_ : t_int_, param_types, false);
 
-  auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
-  ICHECK(global_symbol.defined())
-      << "CodeGenLLVM: Expect PrimFunc to have the global_symbol attribute";
-  function_ = module_->getFunction(MakeStringRef(global_symbol.value()));
-  if (function_ == nullptr) {
-    function_ = llvm::Function::Create(ftype, llvm::Function::ExternalLinkage,
-                                       MakeStringRef(global_symbol.value()), module_.get());
+  auto [symbol_name, linkage_type] = GetLinkage(gvar, func);
+
+  auto function = module_->getFunction(MakeStringRef(symbol_name));
+  if (function == nullptr) {
+    function =
+        llvm::Function::Create(ftype, linkage_type, MakeStringRef(symbol_name), module_.get());
   }
-  function_->setCallingConv(llvm::CallingConv::C);
-  function_->setDLLStorageClass(llvm::GlobalValue::DLLStorageClassTypes::DLLExportStorageClass);
-  SetTargetAttributes(function_);
+  function->setCallingConv(llvm::CallingConv::C);
+  function->setDLLStorageClass(llvm::GlobalValue::DLLStorageClassTypes::DLLExportStorageClass);
+  SetTargetAttributes(function);
+
+  functions_[gvar.get()] = function;
+
+  return function;
+}
+
+void CodeGenLLVM::AddFunctionInternal(const GlobalVar& gvar, const PrimFunc& f, bool ret_void) {
+  this->InitFuncState();
+
+  function_ = DeclareFunctionInternal(gvar, f, ret_void);
 
   // set var map and align information
   auto arg_it = function_->arg_begin();
@@ -1747,9 +1781,19 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const CallNode* op) {
       VLOG(2) << "CreateIntrinsic done";
       return x;
     }
+  } else if (auto* ptr_gvar = op->op.as<GlobalVarNode>()) {
+    auto gvar = GetRef<GlobalVar>(ptr_gvar);
+    auto it = functions_.find(ptr_gvar);
+    ICHECK(it != functions_.end()) << "Call to undefined GlobalVar \"" << gvar << "\"";
+    llvm::Function* callee = it->second;
+    std::vector<llvm::Value*> arg_value;
+    for (const auto& arg : op->args) {
+      arg_value.push_back(MakeValue(arg));
+    }
+    return builder_->CreateCall(callee, arg_value);
+
   } else {
-    ICHECK(op->op.as<GlobalVarNode>());
-    LOG(FATAL) << "Do not yet support cross function call";
+    LOG(FATAL) << "Unsupported operation in CallNode: " << op->op;
   }
 }
 

src/target/llvm/codegen_llvm.h

@@ -63,6 +63,7 @@
 #include <algorithm>
 #include <memory>
 #include <string>
+#include <tuple>
 #include <unordered_map>
 #include <unordered_set>
 #include <utility>
@@ -132,11 +133,17 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
    */
   void SetFastMathFlags(llvm::FastMathFlags fmf);
 
+  virtual llvm::Function* DeclareFunction(const GlobalVar& gvar, const PrimFunc& f);
+
   /*!
    * \brief Compile and add function f to the current module.
+   *
+   * \param gvar The GlobalVar which may be used to may internal calls
+   * to this function from elsewhere in the module.
+   *
    * \param f The function to be added.
    */
-  virtual void AddFunction(const PrimFunc& f);
+  virtual void AddFunction(const GlobalVar& gvar, const PrimFunc& f);
   /*!
    * \brief Add main function as the entry name
    * \param entry_func_name The name of entry function to be added.
@@ -356,7 +363,28 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
   virtual int NativeVectorBits(const runtime::StorageScope& storage_scope) const;
   // Get correct address space depending on the backend
   virtual unsigned GetGlobalAddressSpace() const;
-  void AddFunctionInternal(const PrimFunc& f, bool ret_void);
+
+  /*! \brief Get the linkage parameters for the function
+   *
+   * Returns a tuple whose first element is the name of the function
+   * and whose second element is the linkage type to be used
+   * (e.g. llvm::Function::ExternalLinkage or
+   * llvm::Function::PrivateLinkage)
+   *
+   * \param func The PrimFunc whose symbol name and linkage type
+   * should be returned
+   *
+   * \param gvar The GlobalVar to be used when generating the symbol
+   * name.  Only used for internal functions, for which the
+   * kGlobalSymbol attribute is not defined.
+   */
+  std::tuple<std::string, llvm::Function::LinkageTypes> GetLinkage(const GlobalVar& gvar,
+                                                                   const PrimFunc& func);
+
+  llvm::Function* DeclareFunctionInternal(const GlobalVar& gvar, const PrimFunc& f, bool ret_void);
+
+  void AddFunctionInternal(const GlobalVar& gvar, const PrimFunc& f, bool ret_void);
+
   // Create extern call
   llvm::CallInst* CreateCallExtern(llvm::Type* ret, const std::string& name,
                                    const std::vector<llvm::Value*>& value);
@@ -517,6 +545,11 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
   std::unordered_map<const VarNode*, llvm::Value*> var_map_;
   // global strings
   std::unordered_map<std::string, llvm::Constant*> str_map_;
+
+  // Map from TVM's GlobalVar to the llvm::Function that represents
+  // that function.
+  std::unordered_map<const GlobalVarNode*, llvm::Function*> functions_;
+
   // Whether current function is restricted
   bool is_restricted_{true};
   // The analyzer information
@@ -569,18 +602,26 @@ inline int CodeGenLLVM::GetVectorNumElements(llvm::Value* vec) {
 
 template <typename IterType, typename ConvType>
 void CodeGenLLVM::AddFunctionsOrdered(IterType begin, IterType end, ConvType pfunc) {
-  std::vector<PrimFunc> funcs;
+  std::vector<std::tuple<GlobalVar, PrimFunc>> funcs;
   for (auto it = begin; it != end; ++it) {
-    funcs.push_back(pfunc(*it));
+    auto [gvar, func] = *it;
+    auto converted = pfunc(func);
+    funcs.push_back({gvar, Downcast<PrimFunc>(converted)});
   }
-  std::sort(funcs.begin(), funcs.end(), [](PrimFunc func_a, PrimFunc func_b) {
-    std::string name_a = func_a->GetAttr<String>(tvm::attr::kGlobalSymbol).value();
-    std::string name_b = func_b->GetAttr<String>(tvm::attr::kGlobalSymbol).value();
+  std::sort(funcs.begin(), funcs.end(), [this](const auto& pair_a, const auto& pair_b) {
+    const auto& [gvar_a, func_a] = pair_a;
+    std::string name_a = std::get<std::string>(GetLinkage(gvar_a, func_a));
+
+    const auto& [gvar_b, func_b] = pair_b;
+    std::string name_b = std::get<std::string>(GetLinkage(gvar_b, func_b));
     return name_a < name_b;
   });
-  for (auto& f : funcs) {
-    auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
-    AddFunction(f);
+
+  for (const auto& [gvar, func] : funcs) {
+    DeclareFunction(gvar, func);
+  }
+  for (const auto& [gvar, func] : funcs) {
+    AddFunction(gvar, func);
   }
 }
 

src/target/llvm/codegen_nvptx.cc

@@ -66,9 +66,13 @@ namespace codegen {
 // NVPTX code generator.
 class CodeGenNVPTX : public CodeGenLLVM {
  public:
-  void AddFunction(const PrimFunc& f) final {
+  llvm::Function* DeclareFunction(const GlobalVar& gvar, const PrimFunc& f) final {
     // add function as void return value
-    CodeGenLLVM::AddFunctionInternal(f, true);
+    return CodeGenLLVM::DeclareFunctionInternal(gvar, f, true);
+  }
+  void AddFunction(const GlobalVar& gvar, const PrimFunc& f) final {
+    // add function as void return value
+    CodeGenLLVM::AddFunctionInternal(gvar, f, true);
     // annotate as kernel function
     llvm::LLVMContext* ctx = llvm_target_->GetContext();
     module_->getOrInsertNamedMetadata("nvvm.annotations")
@@ -311,11 +315,7 @@ runtime::Module BuildNVPTX(IRModule mod, Target target) {
 
   cg->Init("TVMPTXModule", llvm_target.get(), NullOpt, false, false);
 
-  cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end(), [](auto& kv) {
-    ICHECK(kv.second->template IsInstance<PrimFuncNode>())
-        << "Can only lower IR Module with PrimFuncs";
-    return Downcast<PrimFunc>(kv.second);
-  });
+  cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end());
 
   llvm::TargetMachine* tm = llvm_target->GetOrCreateTargetMachine();
   const auto* flibdevice_path = tvm::runtime::Registry::Get("tvm_callback_libdevice_path");

src/target/llvm/llvm_module.cc

@@ -125,7 +125,7 @@ class LLVMModuleNode final : public runtime::ModuleNode {
   // The unique_ptr owning the module. This becomes empty once JIT has been initialized
   // (EngineBuilder takes ownership of the module).
   std::unique_ptr<llvm::Module> module_owning_ptr_;
-  /* \brief names of the functions declared in this module */
+  /* \brief names of the external functions declared in this module */
   Array<String> function_names_;
 };
 
@@ -295,7 +295,6 @@ void LLVMModuleNode::Init(const IRModule& mod, const Target& target) {
   llvm::TargetMachine* tm = llvm_target->GetOrCreateTargetMachine();
   std::unique_ptr<CodeGenLLVM> cg = CodeGenLLVM::Create(llvm_target.get());
 
-  std::vector<PrimFunc> funcs;
   std::string entry_func;
   relay::Runtime runtime =
       mod->GetAttr<relay::Runtime>(tvm::attr::kRuntime).value_or(relay::Runtime::Create("cpp"));
@@ -315,12 +314,16 @@ void LLVMModuleNode::Init(const IRModule& mod, const Target& target) {
     }
     auto f = Downcast<PrimFunc>(kv.second);
     auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
-    ICHECK(global_symbol.defined());
-    function_names_.push_back(global_symbol.value());
-    if (f->HasNonzeroAttr(tir::attr::kIsEntryFunc)) {
-      entry_func = global_symbol.value();
+    bool is_entry_func = f->HasNonzeroAttr(tir::attr::kIsEntryFunc);
+
+    ICHECK(global_symbol || !is_entry_func) << "The entry func must be exposed externally.";
+
+    if (global_symbol) {
+      function_names_.push_back(global_symbol.value());
+      if (is_entry_func) {
+        entry_func = global_symbol.value();
+      }
     }
-    funcs.push_back(f);
   }
   // TODO(@jroesch): follow up on this condition.
   // ICHECK(funcs.size() > 0);
@@ -330,7 +333,7 @@ void LLVMModuleNode::Init(const IRModule& mod, const Target& target) {
            target_c_runtime);
   cg->SetFastMathFlags(llvm_target->GetFastMathFlags());
 
-  cg->AddFunctionsOrdered(funcs.begin(), funcs.end());
+  cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end());
   if (entry_func.length() != 0) {
     cg->AddMainFunction(entry_func);
   }