Properly implement the gpu-kernel ABI for amdgpu

compiler/rustc_abi/src/layout/ty.rs

@@ -253,8 +253,8 @@ impl<'a, Ty> TyAndLayout<'a, Ty> {
     }
 
     /// If this method returns `true`, then this type should always have a `PassMode` of
-    /// `Indirect { on_stack: false, .. }` when being used as the argument type of a function with a
-    /// non-Rustic ABI (this is true for structs annotated with the
+    /// `Indirect { mode: IndirectMode::Pointer, .. }` when being used as the argument type of a
+    /// function with a non-Rustic ABI (this is true for structs annotated with the
     /// `#[rustc_pass_indirectly_in_non_rustic_abis]` attribute).
     ///
     /// This is used to replicate some of the behaviour of C array-to-pointer decay; however unlike

compiler/rustc_abi/src/lib.rs

@@ -1772,6 +1772,10 @@ pub struct AddressSpace(pub u32);
 impl AddressSpace {
     /// LLVM's `0` address space.
     pub const ZERO: Self = AddressSpace(0);
+    /// The address space for constant memory on nvptx and amdgpu.
+    /// This address space is used e.g. for kernel arguments that are constant throughout the
+    /// execution.
+    pub const GPU_CONSTANT: Self = AddressSpace(4);
     /// The address space for workgroup memory on nvptx and amdgpu.
     /// See e.g. the `gpu_launch_sized_workgroup_mem` intrinsic for details.
     pub const GPU_WORKGROUP: Self = AddressSpace(3);

compiler/rustc_codegen_cranelift/src/abi/pass_mode.rs

@@ -3,7 +3,7 @@
 use cranelift_codegen::ir::ArgumentPurpose;
 use rustc_abi::{Reg, RegKind};
 use rustc_target::callconv::{
-    ArgAbi, ArgAttributes, ArgExtension as RustcArgExtension, CastTarget, PassMode,
+    ArgAbi, ArgAttributes, ArgExtension as RustcArgExtension, CastTarget, IndirectMode, PassMode,
 };
 use smallvec::{SmallVec, smallvec};
 
@@ -126,8 +126,12 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                 assert_eq!(pad_i32_count, 0, "padding support not yet implemented");
                 cast_target_to_abi_params(cast).into_iter().map(|(_, param)| param).collect()
             }
-            PassMode::Indirect { attrs, meta_attrs: None, on_stack } => {
-                if on_stack {
+            PassMode::Indirect { attrs, meta_attrs: None, address_space: _, mode } => {
+                assert!(
+                    mode != IndirectMode::AmdgpuKernelArg,
+                    "unsupported amdgpu kernel argument"
+                );
+                if mode == IndirectMode::OnStack {
                     // Abi requires aligning struct size to pointer size
                     let size = self.layout.size.align_to(tcx.data_layout.pointer_align().abi);
                     let size = u32::try_from(size.bytes()).unwrap();
@@ -139,8 +143,8 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                     smallvec![apply_attrs_to_abi_param(AbiParam::new(pointer_ty(tcx)), attrs)]
                 }
             }
-            PassMode::Indirect { attrs, meta_attrs: Some(meta_attrs), on_stack } => {
-                assert!(!on_stack);
+            PassMode::Indirect { attrs, meta_attrs: Some(meta_attrs), address_space: _, mode } => {
+                assert!(mode == IndirectMode::Pointer);
                 smallvec![
                     apply_attrs_to_abi_param(AbiParam::new(pointer_ty(tcx)), attrs),
                     apply_attrs_to_abi_param(AbiParam::new(pointer_ty(tcx)), meta_attrs),
@@ -184,8 +188,8 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                 None,
                 cast_target_to_abi_params(cast).into_iter().map(|(_, param)| param).collect(),
             ),
-            PassMode::Indirect { attrs, meta_attrs: None, on_stack } => {
-                assert!(!on_stack);
+            PassMode::Indirect { attrs, meta_attrs: None, address_space: _, mode } => {
+                assert!(mode == IndirectMode::Pointer);
                 (
                     Some(apply_attrs_to_abi_param(
                         AbiParam::special(pointer_ty(tcx), ArgumentPurpose::StructReturn),
@@ -194,7 +198,7 @@ impl<'tcx> ArgAbiExt<'tcx> for ArgAbi<'tcx, Ty<'tcx>> {
                     vec![],
                 )
             }
-            PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+            PassMode::Indirect { attrs: _, meta_attrs: Some(_), address_space: _, mode: _ } => {
                 unreachable!("unsized return value")
             }
         }
@@ -324,7 +328,7 @@ pub(super) fn cvalue_for_param<'tcx>(
         PassMode::Cast { ref cast, .. } => {
             from_casted_value(fx, &block_params, arg_abi.layout, cast)
         }
-        PassMode::Indirect { attrs, meta_attrs: None, on_stack: _ } => {
+        PassMode::Indirect { attrs, meta_attrs: None, address_space: _, mode: _ } => {
             assert_eq!(block_params.len(), 1, "{:?}", block_params);
             if let Some(pointee_align) = attrs.pointee_align
                 && pointee_align < arg_abi.layout.align.abi
@@ -342,7 +346,7 @@ pub(super) fn cvalue_for_param<'tcx>(
                 CValue::by_ref(Pointer::new(block_params[0]), arg_abi.layout)
             }
         }
-        PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+        PassMode::Indirect { attrs: _, meta_attrs: Some(_), address_space: _, mode: _ } => {
             assert_eq!(block_params.len(), 2, "{:?}", block_params);
             CValue::by_ref_unsized(Pointer::new(block_params[0]), block_params[1], arg_abi.layout)
         }

compiler/rustc_codegen_cranelift/src/abi/returning.rs

@@ -17,12 +17,12 @@ pub(super) fn codegen_return_param<'tcx>(
             let is_ssa = ssa_analyzed[RETURN_PLACE].is_ssa(fx, fx.fn_abi.ret.layout.ty);
             (super::make_local_place(fx, RETURN_PLACE, fx.fn_abi.ret.layout, is_ssa), smallvec![])
         }
-        PassMode::Indirect { attrs: _, meta_attrs: None, on_stack: _ } => {
+        PassMode::Indirect { attrs: _, meta_attrs: None, address_space: _, mode: _ } => {
             let ret_param = block_params_iter.next().unwrap();
             assert_eq!(fx.bcx.func.dfg.value_type(ret_param), fx.pointer_type);
             (CPlace::for_ptr(Pointer::new(ret_param), fx.fn_abi.ret.layout), smallvec![ret_param])
         }
-        PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+        PassMode::Indirect { attrs: _, meta_attrs: Some(_), address_space: _, mode: _ } => {
             unreachable!("unsized return value")
         }
     };
@@ -50,7 +50,7 @@ pub(super) fn codegen_with_call_return_arg<'tcx>(
 ) {
     let (ret_temp_place, return_ptr) = match ret_arg_abi.mode {
         PassMode::Ignore => (None, None),
-        PassMode::Indirect { attrs: _, meta_attrs: None, on_stack: _ } => {
+        PassMode::Indirect { attrs: _, meta_attrs: None, address_space: _, mode: _ } => {
             if let Some(ret_ptr) = ret_place.try_to_ptr() {
                 // This is an optimization to prevent unnecessary copies of the return value when
                 // the return place is already a memory place as opposed to a register.
@@ -61,7 +61,7 @@ pub(super) fn codegen_with_call_return_arg<'tcx>(
                 (Some(place), Some(place.to_ptr().get_addr(fx)))
             }
         }
-        PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+        PassMode::Indirect { attrs: _, meta_attrs: Some(_), address_space: _, mode: _ } => {
             unreachable!("unsized return value")
         }
         PassMode::Direct(_) | PassMode::Pair(_, _) | PassMode::Cast { .. } => (None, None),
@@ -86,14 +86,14 @@ pub(super) fn codegen_with_call_return_arg<'tcx>(
                 super::pass_mode::from_casted_value(fx, &results, ret_place.layout(), cast);
             ret_place.write_cvalue(fx, result);
         }
-        PassMode::Indirect { attrs: _, meta_attrs: None, on_stack: _ } => {
+        PassMode::Indirect { attrs: _, meta_attrs: None, address_space: _, mode: _ } => {
             if let Some(ret_temp_place) = ret_temp_place {
                 // If ret_temp_place is None, it is not necessary to copy the return value.
                 let ret_temp_value = ret_temp_place.to_cvalue(fx);
                 ret_place.write_cvalue(fx, ret_temp_value);
             }
         }
-        PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+        PassMode::Indirect { attrs: _, meta_attrs: Some(_), address_space: _, mode: _ } => {
             unreachable!("unsized return value")
         }
     }
@@ -102,10 +102,11 @@ pub(super) fn codegen_with_call_return_arg<'tcx>(
 /// Codegen a return instruction with the right return value(s) if any.
 pub(crate) fn codegen_return(fx: &mut FunctionCx<'_, '_, '_>) {
     match fx.fn_abi.ret.mode {
-        PassMode::Ignore | PassMode::Indirect { attrs: _, meta_attrs: None, on_stack: _ } => {
+        PassMode::Ignore
+        | PassMode::Indirect { attrs: _, meta_attrs: None, address_space: _, mode: _ } => {
             fx.bcx.ins().return_(&[]);
         }
-        PassMode::Indirect { attrs: _, meta_attrs: Some(_), on_stack: _ } => {
+        PassMode::Indirect { attrs: _, meta_attrs: Some(_), address_space: _, mode: _ } => {
             unreachable!("unsized return value")
         }
         PassMode::Direct(_) => {

compiler/rustc_codegen_gcc/src/abi.rs

@@ -11,7 +11,7 @@ use rustc_middle::ty::layout::LayoutOf;
 #[cfg(feature = "master")]
 use rustc_session::{Session, config};
 use rustc_span::bug;
-use rustc_target::callconv::{ArgAttributes, CastTarget, FnAbi, PassMode};
+use rustc_target::callconv::{ArgAttributes, CastTarget, FnAbi, IndirectMode, PassMode};
 #[cfg(feature = "master")]
 use rustc_target::spec::Arch;
 
@@ -178,19 +178,42 @@ impl<'gcc, 'tcx> FnAbiGccExt<'gcc, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
                     let ty = cast.gcc_type(cx);
                     apply_attrs(ty, &cast.attrs, argument_tys.len())
                 }
-                PassMode::Indirect { attrs: _, meta_attrs: None, on_stack: true } => {
+                PassMode::Indirect {
+                    attrs: _,
+                    meta_attrs: None,
+                    address_space: _,
+                    mode: IndirectMode::OnStack,
+                } => {
                     // This is a "byval" argument, so we don't apply the `restrict` attribute on it.
                     on_stack_param_indices.insert(argument_tys.len());
                     arg.layout.gcc_type(cx)
                 }
+                PassMode::Indirect {
+                    attrs: _,
+                    meta_attrs: None,
+                    address_space: _,
+                    mode: IndirectMode::AmdgpuKernelArg,
+                } => {
+                    unimplemented!("unsupported amdgpu kernel argument")
+                }
                 PassMode::Direct(attrs) => {
                     apply_attrs(arg.layout.immediate_gcc_type(cx), &attrs, argument_tys.len())
                 }
-                PassMode::Indirect { attrs, meta_attrs: None, on_stack: false } => {
+                PassMode::Indirect {
+                    attrs,
+                    meta_attrs: None,
+                    address_space: _,
+                    mode: IndirectMode::Pointer,
+                } => {
                     apply_attrs(cx.type_ptr_to(arg.layout.gcc_type(cx)), &attrs, argument_tys.len())
                 }
-                PassMode::Indirect { attrs, meta_attrs: Some(meta_attrs), on_stack } => {
-                    assert!(!on_stack);
+                PassMode::Indirect {
+                    attrs,
+                    meta_attrs: Some(meta_attrs),
+                    address_space: _,
+                    mode,
+                } => {
+                    assert!(mode == IndirectMode::Pointer);
                     // Construct the type of a (wide) pointer to `ty`, and pass its two fields.
                     // Any two ABI-compatible unsized types have the same metadata type and
                     // moreover the same metadata value leads to the same dynamic size and