Rollup merge of rust-lang#135047 - Flakebi:amdgpu-kernel-cc, r=workin…

…gjubilee Add gpu-kernel calling convention The amdgpu-kernel calling convention was reverted in commit f6b21e9 (rust-lang#120495 and rust-lang/rust-analyzer#16463) due to inactivity in the amdgpu target. Introduce a `gpu-kernel` calling convention that translates to `ptx_kernel` or `amdgpu_kernel`, depending on the target that rust compiles for. Tracking issue: rust-lang#135467 amdgpu target tracking issue: rust-lang#135024
matthiaskrgr · Jan 14, 2025 · 75029dc · 75029dc
2 parents 603fdb9 + 68b2639
commit 75029dc
Show file tree

Hide file tree

Showing 29 changed files with 440 additions and 175 deletions.
diff --git a/compiler/rustc_abi/src/extern_abi/mod.rs b/compiler/rustc_abi/src/extern_abi/mod.rs
@@ -45,6 +45,9 @@ pub enum ExternAbi {
     PtxKernel,
     Msp430Interrupt,
     X86Interrupt,
+    /// An entry-point function called by the GPU's host
+    // FIXME: should not be callable from Rust on GPU targets, is for host's use only
+    GpuKernel,
     EfiApi,
     AvrInterrupt,
     AvrNonBlockingInterrupt,
@@ -122,6 +125,7 @@ const AbiDatas: &[AbiData] = &[
     AbiData { abi: Abi::PtxKernel, name: "ptx-kernel" },
     AbiData { abi: Abi::Msp430Interrupt, name: "msp430-interrupt" },
     AbiData { abi: Abi::X86Interrupt, name: "x86-interrupt" },
+    AbiData { abi: Abi::GpuKernel, name: "gpu-kernel" },
     AbiData { abi: Abi::EfiApi, name: "efiapi" },
     AbiData { abi: Abi::AvrInterrupt, name: "avr-interrupt" },
     AbiData { abi: Abi::AvrNonBlockingInterrupt, name: "avr-non-blocking-interrupt" },
@@ -239,6 +243,10 @@ pub fn is_stable(name: &str) -> Result<(), AbiDisabled> {
             feature: sym::abi_x86_interrupt,
             explain: "x86-interrupt ABI is experimental and subject to change",
         }),
+        "gpu-kernel" => Err(AbiDisabled::Unstable {
+            feature: sym::abi_gpu_kernel,
+            explain: "gpu-kernel ABI is experimental and subject to change",
+        }),
         "avr-interrupt" | "avr-non-blocking-interrupt" => Err(AbiDisabled::Unstable {
             feature: sym::abi_avr_interrupt,
             explain: "avr-interrupt and avr-non-blocking-interrupt ABIs are experimental and subject to change",
@@ -293,20 +301,21 @@ impl Abi {
             PtxKernel => 19,
             Msp430Interrupt => 20,
             X86Interrupt => 21,
-            EfiApi => 22,
-            AvrInterrupt => 23,
-            AvrNonBlockingInterrupt => 24,
-            CCmseNonSecureCall => 25,
-            CCmseNonSecureEntry => 26,
+            GpuKernel => 22,
+            EfiApi => 23,
+            AvrInterrupt => 24,
+            AvrNonBlockingInterrupt => 25,
+            CCmseNonSecureCall => 26,
+            CCmseNonSecureEntry => 27,
             // Cross-platform ABIs
-            System { unwind: false } => 27,
-            System { unwind: true } => 28,
-            RustIntrinsic => 29,
-            RustCall => 30,
-            Unadjusted => 31,
-            RustCold => 32,
-            RiscvInterruptM => 33,
-            RiscvInterruptS => 34,
+            System { unwind: false } => 28,
+            System { unwind: true } => 29,
+            RustIntrinsic => 30,
+            RustCall => 31,
+            Unadjusted => 32,
+            RustCold => 33,
+            RiscvInterruptM => 34,
+            RiscvInterruptS => 35,
         };
         debug_assert!(
             AbiDatas

diff --git a/compiler/rustc_codegen_cranelift/src/abi/mod.rs b/compiler/rustc_codegen_cranelift/src/abi/mod.rs
@@ -65,7 +65,11 @@ pub(crate) fn conv_to_call_conv(sess: &Session, c: Conv, default_call_conv: Call
             sess.dcx().fatal("C-cmse-nonsecure-entry call conv is not yet implemented");
         }
 
-        Conv::Msp430Intr | Conv::PtxKernel | Conv::AvrInterrupt | Conv::AvrNonBlockingInterrupt => {
+        Conv::Msp430Intr
+        | Conv::PtxKernel
+        | Conv::GpuKernel
+        | Conv::AvrInterrupt
+        | Conv::AvrNonBlockingInterrupt => {
             unreachable!("tried to use {c:?} call conv which only exists on an unsupported target");
         }
     }

diff --git a/compiler/rustc_codegen_llvm/src/abi.rs b/compiler/rustc_codegen_llvm/src/abi.rs
@@ -1,3 +1,4 @@
+use std::borrow::Borrow;
 use std::cmp;
 
 use libc::c_uint;
@@ -312,7 +313,7 @@ impl<'ll, 'tcx> ArgAbiBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
 pub(crate) trait FnAbiLlvmExt<'ll, 'tcx> {
     fn llvm_type(&self, cx: &CodegenCx<'ll, 'tcx>) -> &'ll Type;
     fn ptr_to_llvm_type(&self, cx: &CodegenCx<'ll, 'tcx>) -> &'ll Type;
-    fn llvm_cconv(&self) -> llvm::CallConv;
+    fn llvm_cconv(&self, cx: &CodegenCx<'ll, 'tcx>) -> llvm::CallConv;
 
     /// Apply attributes to a function declaration/definition.
     fn apply_attrs_llfn(
@@ -404,8 +405,8 @@ impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
         cx.type_ptr_ext(cx.data_layout().instruction_address_space)
     }
 
-    fn llvm_cconv(&self) -> llvm::CallConv {
-        self.conv.into()
+    fn llvm_cconv(&self, cx: &CodegenCx<'ll, 'tcx>) -> llvm::CallConv {
+        llvm::CallConv::from_conv(self.conv, cx.tcx.sess.target.arch.borrow())
     }
 
     fn apply_attrs_llfn(
@@ -617,7 +618,7 @@ impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
             }
         }
 
-        let cconv = self.llvm_cconv();
+        let cconv = self.llvm_cconv(&bx.cx);
         if cconv != llvm::CCallConv {
             llvm::SetInstructionCallConv(callsite, cconv);
         }
@@ -655,8 +656,8 @@ impl<'tcx> AbiBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
     }
 }
 
-impl From<Conv> for llvm::CallConv {
-    fn from(conv: Conv) -> Self {
+impl llvm::CallConv {
+    pub fn from_conv(conv: Conv, arch: &str) -> Self {
         match conv {
             Conv::C
             | Conv::Rust
@@ -666,6 +667,15 @@ impl From<Conv> for llvm::CallConv {
             Conv::Cold => llvm::ColdCallConv,
             Conv::PreserveMost => llvm::PreserveMost,
             Conv::PreserveAll => llvm::PreserveAll,
+            Conv::GpuKernel => {
+                if arch == "amdgpu" {
+                    llvm::AmdgpuKernel
+                } else if arch == "nvptx64" {
+                    llvm::PtxKernel
+                } else {
+                    panic!("Architecture {arch} does not support GpuKernel calling convention");
+                }
+            }
             Conv::AvrInterrupt => llvm::AvrInterrupt,
             Conv::AvrNonBlockingInterrupt => llvm::AvrNonBlockingInterrupt,
             Conv::ArmAapcs => llvm::ArmAapcsCallConv,

diff --git a/compiler/rustc_codegen_llvm/src/context.rs b/compiler/rustc_codegen_llvm/src/context.rs
@@ -741,7 +741,10 @@ impl<'ll, 'tcx> MiscCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
         if self.get_declared_value(entry_name).is_none() {
             Some(self.declare_entry_fn(
                 entry_name,
-                self.sess().target.entry_abi.into(),
+                llvm::CallConv::from_conv(
+                    self.sess().target.entry_abi,
+                    self.sess().target.arch.borrow(),
+                ),
                 llvm::UnnamedAddr::Global,
                 fn_type,
             ))

diff --git a/compiler/rustc_codegen_llvm/src/declare.rs b/compiler/rustc_codegen_llvm/src/declare.rs
@@ -125,7 +125,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
         let llfn = declare_raw_fn(
             self,
             name,
-            fn_abi.llvm_cconv(),
+            fn_abi.llvm_cconv(self),
             llvm::UnnamedAddr::Global,
             llvm::Visibility::Default,
             fn_abi.llvm_type(self),

diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -120,6 +120,7 @@ pub enum CallConv {
     X86_Intr = 83,
     AvrNonBlockingInterrupt = 84,
     AvrInterrupt = 85,
+    AmdgpuKernel = 91,
 }
 
 /// Must match the layout of `LLVMLinkage`.

diff --git a/compiler/rustc_feature/src/unstable.rs b/compiler/rustc_feature/src/unstable.rs
@@ -361,6 +361,8 @@ declare_features! (
     (unstable, abi_avr_interrupt, "1.45.0", Some(69664)),
     /// Allows `extern "C-cmse-nonsecure-call" fn()`.
     (unstable, abi_c_cmse_nonsecure_call, "1.51.0", Some(81391)),
+    /// Allows `extern "gpu-kernel" fn()`.
+    (unstable, abi_gpu_kernel, "CURRENT_RUSTC_VERSION", Some(135467)),
     /// Allows `extern "msp430-interrupt" fn()`.
     (unstable, abi_msp430_interrupt, "1.16.0", Some(38487)),
     /// Allows `extern "ptx-*" fn()`.

diff --git a/compiler/rustc_middle/src/ty/layout.rs b/compiler/rustc_middle/src/ty/layout.rs
@@ -1241,6 +1241,7 @@ pub fn fn_can_unwind(tcx: TyCtxt<'_>, fn_def_id: Option<DefId>, abi: ExternAbi)
         PtxKernel
         | Msp430Interrupt
         | X86Interrupt
+        | GpuKernel
         | EfiApi
         | AvrInterrupt
         | AvrNonBlockingInterrupt

diff --git a/compiler/rustc_smir/src/rustc_internal/internal.rs b/compiler/rustc_smir/src/rustc_internal/internal.rs
@@ -472,6 +472,7 @@ impl RustcInternal for Abi {
             Abi::PtxKernel => rustc_abi::ExternAbi::PtxKernel,
             Abi::Msp430Interrupt => rustc_abi::ExternAbi::Msp430Interrupt,
             Abi::X86Interrupt => rustc_abi::ExternAbi::X86Interrupt,
+            Abi::GpuKernel => rustc_abi::ExternAbi::GpuKernel,
             Abi::EfiApi => rustc_abi::ExternAbi::EfiApi,
             Abi::AvrInterrupt => rustc_abi::ExternAbi::AvrInterrupt,
             Abi::AvrNonBlockingInterrupt => rustc_abi::ExternAbi::AvrNonBlockingInterrupt,

diff --git a/compiler/rustc_smir/src/rustc_smir/convert/abi.rs b/compiler/rustc_smir/src/rustc_smir/convert/abi.rs
@@ -113,6 +113,7 @@ impl<'tcx> Stable<'tcx> for callconv::Conv {
             Conv::X86VectorCall => CallConvention::X86VectorCall,
             Conv::X86_64SysV => CallConvention::X86_64SysV,
             Conv::X86_64Win64 => CallConvention::X86_64Win64,
+            Conv::GpuKernel => CallConvention::GpuKernel,
             Conv::AvrInterrupt => CallConvention::AvrInterrupt,
             Conv::AvrNonBlockingInterrupt => CallConvention::AvrNonBlockingInterrupt,
             Conv::RiscvInterrupt { .. } => CallConvention::RiscvInterrupt,

diff --git a/compiler/rustc_smir/src/rustc_smir/convert/ty.rs b/compiler/rustc_smir/src/rustc_smir/convert/ty.rs
@@ -911,6 +911,7 @@ impl<'tcx> Stable<'tcx> for rustc_abi::ExternAbi {
             ExternAbi::Win64 { unwind } => Abi::Win64 { unwind },
             ExternAbi::SysV64 { unwind } => Abi::SysV64 { unwind },
             ExternAbi::PtxKernel => Abi::PtxKernel,
+            ExternAbi::GpuKernel => Abi::GpuKernel,
             ExternAbi::Msp430Interrupt => Abi::Msp430Interrupt,
             ExternAbi::X86Interrupt => Abi::X86Interrupt,
             ExternAbi::EfiApi => Abi::EfiApi,

diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
@@ -379,6 +379,7 @@ symbols! {
         abi_avr_interrupt,
         abi_c_cmse_nonsecure_call,
         abi_efiapi,
+        abi_gpu_kernel,
         abi_msp430_interrupt,
         abi_ptx,
         abi_riscv_interrupt,

diff --git a/compiler/rustc_target/src/callconv/mod.rs b/compiler/rustc_target/src/callconv/mod.rs
@@ -546,6 +546,8 @@ pub enum Conv {
 
     PtxKernel,
 
+    GpuKernel,
+
     X86Fastcall,
     X86Intr,
     X86Stdcall,
@@ -865,6 +867,7 @@ impl FromStr for Conv {
             "X86VectorCall" => Ok(Conv::X86VectorCall),
             "X86_64SysV" => Ok(Conv::X86_64SysV),
             "X86_64Win64" => Ok(Conv::X86_64Win64),
+            "GpuKernel" => Ok(Conv::GpuKernel),
             "AvrInterrupt" => Ok(Conv::AvrInterrupt),
             "AvrNonBlockingInterrupt" => Ok(Conv::AvrNonBlockingInterrupt),
             "RiscvInterrupt(machine)" => {

diff --git a/compiler/rustc_target/src/json.rs b/compiler/rustc_target/src/json.rs
@@ -113,6 +113,7 @@ impl ToJson for crate::abi::call::Conv {
             Self::X86VectorCall => "X86VectorCall",
             Self::X86_64SysV => "X86_64SysV",
             Self::X86_64Win64 => "X86_64Win64",
+            Self::GpuKernel => "GpuKernel",
             Self::AvrInterrupt => "AvrInterrupt",
             Self::AvrNonBlockingInterrupt => "AvrNonBlockingInterrupt",
             Self::RiscvInterrupt { kind } => {

diff --git a/compiler/rustc_target/src/spec/mod.rs b/compiler/rustc_target/src/spec/mod.rs
@@ -2860,6 +2860,7 @@ impl Target {
             }
             Win64 { .. } | SysV64 { .. } => self.arch == "x86_64",
             PtxKernel => self.arch == "nvptx64",
+            GpuKernel => ["amdgpu", "nvptx64"].contains(&&self.arch[..]),
             Msp430Interrupt => self.arch == "msp430",
             RiscvInterruptM | RiscvInterruptS => ["riscv32", "riscv64"].contains(&&self.arch[..]),
             AvrInterrupt | AvrNonBlockingInterrupt => self.arch == "avr",

diff --git a/compiler/rustc_ty_utils/src/abi.rs b/compiler/rustc_ty_utils/src/abi.rs
@@ -293,6 +293,7 @@ fn conv_from_spec_abi(tcx: TyCtxt<'_>, abi: ExternAbi, c_variadic: bool) -> Conv
         PtxKernel => Conv::PtxKernel,
         Msp430Interrupt => Conv::Msp430Intr,
         X86Interrupt => Conv::X86Intr,
+        GpuKernel => Conv::GpuKernel,
         AvrInterrupt => Conv::AvrInterrupt,
         AvrNonBlockingInterrupt => Conv::AvrNonBlockingInterrupt,
         RiscvInterruptM => Conv::RiscvInterrupt { kind: RiscvInterruptKind::Machine },

diff --git a/compiler/stable_mir/src/abi.rs b/compiler/stable_mir/src/abi.rs
@@ -442,6 +442,8 @@ pub enum CallConvention {
 
     PtxKernel,
 
+    GpuKernel,
+
     X86Fastcall,
     X86Intr,
     X86Stdcall,

diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
@@ -1077,6 +1077,7 @@ pub enum Abi {
     PtxKernel,
     Msp430Interrupt,
     X86Interrupt,
+    GpuKernel,
     EfiApi,
     AvrInterrupt,
     AvrNonBlockingInterrupt,

diff --git a/tests/codegen/gpu-kernel-abi.rs b/tests/codegen/gpu-kernel-abi.rs
@@ -0,0 +1,23 @@
+// Checks that the gpu-kernel calling convention correctly translates to LLVM calling conventions.
+
+//@ revisions: amdgpu nvptx
+//@ [amdgpu] compile-flags: --crate-type=rlib --target=amdgcn-amd-amdhsa -Ctarget-cpu=gfx900
+//@ [amdgpu] needs-llvm-components: amdgpu
+// amdgpu target is not yet merged
+//@ [amdgpu] should-fail
+//@ [nvptx] compile-flags: --crate-type=rlib --target=nvptx64-nvidia-cuda
+//@ [nvptx] needs-llvm-components: nvptx
+#![feature(no_core, lang_items, abi_gpu_kernel)]
+#![no_core]
+
+#[lang = "sized"]
+trait Sized {}
+#[lang = "freeze"]
+trait Freeze {}
+#[lang = "copy"]
+trait Copy {}
+
+// amdgpu: define amdgpu_kernel void @fun(i32
+// nvptx: define ptx_kernel void @fun(i32
+#[no_mangle]
+pub extern "gpu-kernel" fn fun(_: i32) {}