fix device code generation

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -613,7 +613,7 @@ pub(crate) fn run_pass_manager(
         write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage);
     }
 
-    if enable_gpu && !thin {
+    if enable_gpu && !thin && !(cgcx.target_arch == "nvptx64" || cgcx.target_arch == "amdgpu") {
         let cx =
             SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);
         crate::builder::gpu_offload::handle_gpu_code(cgcx, &cx);

compiler/rustc_codegen_llvm/src/back/write.rs

@@ -658,6 +658,23 @@ 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) };
+    }
+
+    let consider_offload = config.offload.contains(&config::Offload::Enable);
+    if consider_offload && (cgcx.target_arch == "amdgpu" || cgcx.target_arch == "nvptx64") {
+        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);
+            }
+        }
+    }
+
     let mut llvm_profiler = cgcx
         .prof
         .llvm_recording_enabled()

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1117,6 +1117,11 @@ unsafe extern "C" {
 
     // Operations on functions
     pub(crate) fn LLVMSetFunctionCallConv(Fn: &Value, CC: c_uint);
+    pub(crate) fn LLVMAddFunction<'a>(
+        Mod: &'a Module,
+        Name: *const c_char,
+        FunctionTy: &'a Type,
+    ) -> &'a Value;
     pub(crate) fn LLVMDeleteFunction(Fn: &Value);
 
     // Operations about llvm intrinsics
@@ -1135,6 +1140,7 @@ 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,
@@ -2007,6 +2013,7 @@ unsafe extern "C" {
     ) -> &Attribute;
 
     // Operations on functions
+    pub(crate) fn LLVMRustOffloadWrapper<'a>(M: &'a Module, Fn: &'a Value);
     pub(crate) fn LLVMRustGetOrInsertFunction<'a>(
         M: &'a Module,
         Name: *const c_char,

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -35,6 +35,8 @@
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Timer.h"
 #include "llvm/Support/ToolOutputFile.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
 #include <iostream>
 
 // for raw `write` in the bad-alloc handler
@@ -142,6 +144,56 @@ extern "C" void LLVMRustPrintStatistics(RustStringRef OutBuf) {
   llvm::PrintStatistics(OS);
 }
 
+extern "C" void LLVMRustOffloadWrapper(LLVMModuleRef M, LLVMValueRef Fn) {
+  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);
+
+  // Map old arguments to new arguments. We skip the first dyn_ptr argument,
+  // since it can't be used directly by user code.
+  llvm::ValueToValueMapTy vmap;
+  auto newArgIt = newFn->arg_begin();
+  newArgIt->setName("dyn_ptr");
+  ++newArgIt; // skip %dyn_ptr
+  for (auto &oldArg : oldFn->args()) {
+    vmap[&oldArg] = &*newArgIt++;
+  }
+
+  llvm::SmallVector<llvm::ReturnInst *, 8> returns;
+  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,
                                               size_t NameLen) {
   return wrap(unwrap(M)->getNamedValue(StringRef(Name, NameLen)));

compiler/rustc_target/src/callconv/mod.rs

@@ -577,6 +577,7 @@ impl RiscvInterruptKind {
 ///
 /// The signature represented by this type may not match the MIR function signature.
 /// Certain attributes, like `#[track_caller]` can introduce additional arguments, which are present in [`FnAbi`], but not in `FnSig`.
+/// The std::offload module also adds an addition dyn_ptr argument to the GpuKernel ABI.
 /// While this difference is rarely relevant, it should still be kept in mind.
 ///
 /// I will do my best to describe this structure, but these