diff --git a/compiler/rustc_middle/src/ty/offload_meta.rs b/compiler/rustc_middle/src/ty/offload_meta.rs index a58e517e05f61..c689c67397b2b 100644 --- a/compiler/rustc_middle/src/ty/offload_meta.rs +++ b/compiler/rustc_middle/src/ty/offload_meta.rs @@ -1,5 +1,6 @@ use bitflags::bitflags; use rustc_abi::{BackendRepr, TyAbiInterface}; +use rustc_span::sym; use rustc_target::callconv::ArgAbi; use crate::ty::{self, PseudoCanonicalInput, Ty, TyCtxt, TypingEnv}; @@ -75,6 +76,12 @@ impl OffloadMetadata { where Ty<'tcx>: TyAbiInterface<'tcx, C>, { + if let Some(elem_ty) = region_element_ty(tcx, ty) { + let ptr = OffloadMetadata::from_ty(tcx, Ty::new_slice(tcx, elem_ty)); + let len = OffloadMetadata::from_ty(tcx, tcx.types.usize); + return vec![(ptr, Ty::new_mut_ptr(tcx, elem_ty)), (len, tcx.types.usize)]; + } + match arg_abi.layout.backend_repr { BackendRepr::ScalarPair { a: _, b: _, b_offset: _ } => (0..2) .map(|i| { @@ -87,6 +94,16 @@ impl OffloadMetadata { } } +fn region_element_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Option> { + if let ty::Adt(def, args) = ty.kind() + && Some(def.did()) == tcx.get_diagnostic_item(sym::offload_region) + { + Some(args.type_at(1)) + } else { + None + } +} + // FIXME(Sa4dUs): implement a solid logic to determine the payload size fn get_payload_size<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> OffloadSize { match ty.kind() { diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index 45df107bf7469..625251321e609 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1469,6 +1469,7 @@ symbols! { off, offload, offload_kernel, + offload_region, offset, offset_of, offset_of_enum, diff --git a/library/core/src/offload/mod.rs b/library/core/src/offload/mod.rs index 17ff74f0bbfbb..549a44ed9a803 100644 --- a/library/core/src/offload/mod.rs +++ b/library/core/src/offload/mod.rs @@ -101,3 +101,95 @@ macro_rules! offload { (@value (SOME $val:expr)) => { $val }; (@value ($val:expr)) => { $val }; } + +// Region & Partitioning Strategy + +/// Defines how execution units access memory regions. +/// +/// # Safety +/// +/// Implementations must guarantee that generated views are disjoint. +#[unstable(feature = "offload", issue = "124509")] +pub unsafe trait PartitioningStrategy { + /// Read-only view type for the partitioned memory region. + type View<'a, T: 'a>; + + /// Mutable view type for the partitioned memory region. + type ViewMut<'a, T: 'a>; + + /// Returns the execution index of the current unit. + fn index() -> usize; + + /// Returns a read-only view of the region for the current execution context. + /// + /// # Safety + /// + /// `ptr` must point to `len` valid, initialized elements of type `T`. + /// The memory must stay valid for lifetime `'a`. + unsafe fn get<'a, T>(ptr: *const T, len: usize) -> Option>; + + /// Returns a mutable view of the region for the current execution context. + /// + /// # Safety + /// + /// `ptr` must point to `len` valid, initialized elements of type `T`. + /// The memory must stay valid for lifetime `'a`. + /// The returned view must be disjoint from all other active views. + unsafe fn get_mut<'a, T>(ptr: *mut T, len: usize) -> Option>; +} + +/// A memory region bound to a partitioning strategy. +#[derive(Copy, Clone, Debug)] +#[unstable(feature = "offload", issue = "124509")] +#[rustc_diagnostic_item = "offload_region"] +pub struct Region<'a, T, S: PartitioningStrategy> { + ptr: *mut T, + len: usize, + _marker: core::marker::PhantomData<(&'a mut [T], S)>, +} + +/// Raw representation used to build a [`Region`] from common aggregate types. +#[derive(Debug)] +#[unstable(feature = "offload", issue = "124509")] +pub struct RawRegion<'a, T> { + ptr: *mut T, + len: usize, + _marker: core::marker::PhantomData<&'a mut [T]>, +} + +impl<'a, T> From<&'a mut [T]> for RawRegion<'a, T> { + fn from(data: &'a mut [T]) -> Self { + Self { ptr: data.as_mut_ptr(), len: data.len(), _marker: core::marker::PhantomData } + } +} + +impl<'a, T, const N: usize> From<&'a mut [T; N]> for RawRegion<'a, T> { + fn from(data: &'a mut [T; N]) -> Self { + Self { ptr: data.as_mut_ptr(), len: N, _marker: core::marker::PhantomData } + } +} + +#[unstable(feature = "offload", issue = "124509")] +impl<'a, T, S: PartitioningStrategy> Region<'a, T, S> { + /// Creates a new partitioned region from data convertible into a [`RawRegion`]. + pub fn new(data: D) -> Self + where + D: Into>, + { + let raw = data.into(); + Self { ptr: raw.ptr, len: raw.len, _marker: core::marker::PhantomData } + } + + /// Returns a read-only view for the current execution context. + pub fn get(&self) -> Option> { + // SAFETY: `self.ptr` points to `self.len` valid elements for lifetime `'a`. + unsafe { S::get(self.ptr as *const T, self.len) } + } + + /// Returns a mutable view for the current execution context. + pub fn get_mut(&mut self) -> Option> { + // SAFETY: `self.ptr` points to `self.len` valid elements for lifetime `'a`. + // The strategy guarantees that the returned view is disjoint. + unsafe { S::get_mut(self.ptr, self.len) } + } +} diff --git a/tests/codegen-llvm/gpu_offload/region_host.rs b/tests/codegen-llvm/gpu_offload/region_host.rs new file mode 100644 index 0000000000000..602dc247f6540 --- /dev/null +++ b/tests/codegen-llvm/gpu_offload/region_host.rs @@ -0,0 +1,65 @@ +//@ compile-flags: -Zoffload=Test -Zunstable-options -C opt-level=1 -Clto=fat +//@ no-prefer-dynamic +//@ needs-offload + +// This test verifies that a `Region` kernel argument is mapped like a slice. +#![feature(abi_gpu_kernel)] +#![feature(core_intrinsics)] +#![feature(gpu_offload)] +#![feature(offload)] +#![feature(rustc_attrs)] +#![no_main] + +extern crate core; + +use core::offload::{PartitioningStrategy, Region}; + +struct Dummy; + +unsafe impl PartitioningStrategy for Dummy { + type View<'a, T: 'a> = &'a T; + type ViewMut<'a, T: 'a> = &'a mut T; + + fn index() -> usize { + 0 + } + + unsafe fn get<'a, T>(_ptr: *const T, _len: usize) -> Option> { + None + } + + unsafe fn get_mut<'a, T>(_ptr: *mut T, _len: usize) -> Option> { + None + } +} + +// CHECK: @anon.[[ID:.*]].0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1 + +// CHECK-DAG: @.offload_sizes.[[K:[^ ]*foo]] = private unnamed_addr constant [2 x i64] [i64 0, i64 8] +// CHECK-DAG: @.offload_maptypes.[[K]].begin = private unnamed_addr constant [2 x i64] [i64 1, i64 768] +// CHECK-DAG: @.offload_maptypes.[[K]].kernel = private unnamed_addr constant [2 x i64] [i64 32, i64 800] +// CHECK-DAG: @.offload_maptypes.[[K]].end = private unnamed_addr constant [2 x i64] [i64 2, i64 0] + +// CHECK: define{{( dso_local)?}} void @main() +// CHECK: %.offload_sizes = alloca [2 x i64], align 8 +// CHECK: call void @llvm.memcpy.p0.p0.i64(ptr {{.*}} %.offload_sizes, ptr {{.*}} @.offload_sizes.[[K]], i64 16, i1 false) +// CHECK: store i64 16, ptr %.offload_sizes, align 8 +// CHECK: call void @__tgt_target_data_begin_mapper(ptr nonnull @anon.[[ID]].1, i64 -1, i32 2, ptr nonnull %.offload_baseptrs, ptr nonnull %.offload_ptrs, ptr nonnull %.offload_sizes, ptr nonnull @.offload_maptypes.[[K]].begin, ptr null, ptr null) +// CHECK: call i32 @__tgt_target_kernel(ptr nonnull @anon.[[ID]].1, i64 -1, i32 1, i32 1, ptr nonnull @.[[K]].region_id, ptr nonnull %kernel_args) +// CHECK-NEXT: call void @__tgt_target_data_end_mapper(ptr nonnull @anon.[[ID]].1, i64 -1, i32 2, ptr nonnull %.offload_baseptrs, ptr nonnull %.offload_ptrs, ptr nonnull %.offload_sizes, ptr nonnull @.offload_maptypes.[[K]].end, ptr null, ptr null) + +#[unsafe(no_mangle)] +fn main() { + let mut x = [0.0f32; 4]; + core::intrinsics::offload::<_, _, ()>( + foo, + [1, 1, 1], + [1, 1, 1], + 0, + (Region::::new(&mut x as &mut [f32]),), + ); +} + +fn foo(region: Region<'_, f32, Dummy>) { + unreachable!(); +}