upgrade offload dyn_ptr handling from C++ to mostly safe Rust

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -45,7 +45,7 @@ use crate::errors::{
 use crate::llvm::diagnostic::OptimizationDiagnosticKind::*;
 use crate::llvm::{self, DiagnosticInfo};
 use crate::type_::llvm_type_ptr;
-use crate::{LlvmCodegenBackend, ModuleLlvm, base, common, llvm_util};
+use crate::{LlvmCodegenBackend, ModuleLlvm, SimpleCx, base, common, llvm_util};
 
 pub(crate) fn llvm_err<'a>(dcx: DiagCtxtHandle<'_>, err: LlvmError<'a>) -> ! {
     match llvm::last_error() {
@@ -658,19 +658,70 @@ pub(crate) unsafe fn llvm_optimize(
         None
     };
 
-    fn handle_offload(m: &llvm::Module, llcx: &llvm::Context, old_fn: &llvm::Value) {
-        unsafe { llvm::LLVMRustOffloadWrapper(m, old_fn) };
+    fn handle_offload<'ll>(cx: &'ll SimpleCx<'_>, old_fn: &llvm::Value) {
+        let old_fn_ty = cx.get_type_of_global(old_fn);
+        let old_param_types = cx.func_params_types(old_fn_ty);
+        let old_param_count = old_param_types.len();
+        if old_param_count == 0 {
+            return;
+        }
+
+        let first_param = llvm::get_param(old_fn, 0);
+        let c_name = llvm::get_value_name(first_param);
+        let first_arg_name = str::from_utf8(&c_name).unwrap();
+        // We might call llvm_optimize (and thus this code) multiple times on the same IR,
+        // but we shouldn't add this helper ptr multiple times.
+        if first_arg_name == "dyn_ptr" {
+            return;
+        }
+
+        // Create the new parameter list, with ptr as the first argument
+        let mut new_param_types = Vec::with_capacity(old_param_count as usize + 1);
+        new_param_types.push(cx.type_ptr());
+        for old_param in old_param_types {
+            new_param_types.push(old_param);
+        }
+
+        // Create the new function type
+        let ret_ty = unsafe { llvm::LLVMGetReturnType(old_fn_ty) };
+        let new_fn_ty = cx.type_func(&new_param_types, ret_ty);
+
+        // Create the new function, with a temporary .offload name to avoid a name collision.
+        let old_fn_name = String::from_utf8(llvm::get_value_name(old_fn)).unwrap();
+        let new_fn_name = format!("{}.offload", &old_fn_name);
+        let new_fn = cx.add_func(&new_fn_name, new_fn_ty);
+        let a0 = llvm::get_param(new_fn, 0);
+        llvm::set_value_name(a0, CString::new("dyn_ptr").unwrap().as_bytes());
+
+        // Here we map the old arguments to the new arguments, with an offset of 1 to make sure
+        // that we don't use the newly added `%dyn_ptr`.
+        unsafe {
+            llvm::LLVMRustOffloadMapper(cx.llmod(), old_fn, new_fn);
+        }
+
+        llvm::set_linkage(new_fn, llvm::get_linkage(old_fn));
+        llvm::set_visibility(new_fn, llvm::get_visibility(old_fn));
+
+        // Replace all uses of old_fn with new_fn (RAUW)
+        unsafe {
+            llvm::LLVMReplaceAllUsesWith(old_fn, new_fn);
+        }
+        let name = llvm::get_value_name(old_fn);
+        unsafe {
+            llvm::LLVMDeleteFunction(old_fn);
+        }
+        // Now we can re-use the old name, without name collision.
+        llvm::set_value_name(new_fn, &name);
     }
 
     let consider_offload = config.offload.contains(&config::Offload::Enable);
     if consider_offload && (cgcx.target_arch == "amdgpu" || cgcx.target_arch == "nvptx64") {
+        let cx =
+            SimpleCx::new(module.module_llvm.llmod(), module.module_llvm.llcx, cgcx.pointer_size);
         for num in 0..9 {
             let name = format!("kernel_{num}");
-            let c_name = CString::new(name).unwrap();
-            if let Some(kernel) =
-                unsafe { llvm::LLVMGetNamedFunction(module.module_llvm.llmod(), c_name.as_ptr()) }
-            {
-                handle_offload(module.module_llvm.llmod(), module.module_llvm.llcx, kernel);
+            if let Some(kernel) = cx.get_function(&name) {
+                handle_offload(&cx, kernel);
             }
         }
     }

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1140,7 +1140,6 @@ unsafe extern "C" {
 
     // Operations on basic blocks
     pub(crate) fn LLVMGetBasicBlockParent(BB: &BasicBlock) -> &Value;
-    pub(crate) fn LLVMAppendExistingBasicBlock<'a>(Fn: &'a Value, BB: &BasicBlock);
     pub(crate) fn LLVMAppendBasicBlockInContext<'a>(
         C: &'a Context,
         Fn: &'a Value,
@@ -2013,7 +2012,7 @@ unsafe extern "C" {
     ) -> &Attribute;
 
     // Operations on functions
-    pub(crate) fn LLVMRustOffloadWrapper<'a>(M: &'a Module, Fn: &'a Value);
+    pub(crate) fn LLVMRustOffloadMapper<'a>(M: &'a Module, Fn: &'a Value, Fn: &'a Value);
     pub(crate) fn LLVMRustGetOrInsertFunction<'a>(
         M: &'a Module,
         Name: *const c_char,

compiler/rustc_codegen_llvm/src/type_.rs

@@ -68,6 +68,11 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         unsafe { llvm::LLVMVectorType(ty, len as c_uint) }
     }
 
+    pub(crate) fn add_func(&self, name: &str, ty: &'ll Type) -> &'ll Value {
+        let name = SmallCStr::new(name);
+        unsafe { llvm::LLVMAddFunction(self.llmod(), name.as_ptr(), ty) }
+    }
+
     pub(crate) fn func_params_types(&self, ty: &'ll Type) -> Vec<&'ll Type> {
         unsafe {
             let n_args = llvm::LLVMCountParamTypes(ty) as usize;

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -144,31 +144,10 @@ extern "C" void LLVMRustPrintStatistics(RustStringRef OutBuf) {
   llvm::PrintStatistics(OS);
 }
 
-extern "C" void LLVMRustOffloadWrapper(LLVMModuleRef M, LLVMValueRef Fn) {
+extern "C" void LLVMRustOffloadMapper(LLVMModuleRef M, LLVMValueRef OldFn, LLVMValueRef NewFn) {
   llvm::Module *module = llvm::unwrap(M);
-  llvm::Function *oldFn = llvm::unwrap<llvm::Function>(Fn);
-
-  if (oldFn->arg_size() > 0 && oldFn->getArg(0)->getName() == "dyn_ptr") {
-    return;
-  }
-
-  // 1. Create new function type with the leading extra %dyn_ptr arg which llvm
-  // offload requries.
-  llvm::LLVMContext &ctx = module->getContext();
-  llvm::Type *dynPtrType = llvm::PointerType::get(ctx, 0);
-  std::vector<llvm::Type *> argTypes;
-  argTypes.push_back(dynPtrType);
-
-  for (auto &arg : oldFn->args()) {
-    argTypes.push_back(arg.getType());
-  }
-
-  llvm::FunctionType *newFnType = llvm::FunctionType::get(
-      oldFn->getReturnType(), argTypes, oldFn->isVarArg());
-
-  // use a temporary .offload appendix to avoid name clashes
-  llvm::Function *newFn = llvm::Function::Create(
-      newFnType, oldFn->getLinkage(), oldFn->getName() + ".offload", module);
+  llvm::Function *oldFn = llvm::unwrap<llvm::Function>(OldFn);
+  llvm::Function *newFn = llvm::unwrap<llvm::Function>(NewFn);
 
   // Map old arguments to new arguments. We skip the first dyn_ptr argument,
   // since it can't be used directly by user code.
@@ -184,14 +163,6 @@ extern "C" void LLVMRustOffloadWrapper(LLVMModuleRef M, LLVMValueRef Fn) {
   llvm::CloneFunctionInto(newFn, oldFn, vmap,
                           llvm::CloneFunctionChangeType::LocalChangesOnly,
                           returns);
-  newFn->setLinkage(oldFn->getLinkage());
-  newFn->setVisibility(oldFn->getVisibility());
-
-  // Replace uses, delete old function, and reset name to the original one.
-  oldFn->replaceAllUsesWith(newFn);
-  auto name = oldFn->getName();
-  oldFn->eraseFromParent();
-  newFn->setName(name);
 }
 
 extern "C" LLVMValueRef LLVMRustGetNamedValue(LLVMModuleRef M, const char *Name,