compiler/rustc_middle/src/ty/offload_meta.rs

use bitflags::bitflags;
use rustc_abi::{BackendRepr, TyAbiInterface};
use rustc_target::callconv::ArgAbi;

use crate::ty::{self, PseudoCanonicalInput, Ty, TyCtxt, TypingEnv};

#[derive(Debug, Copy, Clone)]
pub struct OffloadMetadata {
    pub payload_size: OffloadSize,
    pub mode: MappingFlags,
}

#[derive(Debug, Copy, Clone)]
pub enum OffloadSize {
    Static(u64),
    Slice { element_size: u64 },
}

bitflags! {
    /// Mirrors `OpenMPOffloadMappingFlags` from Clang/OpenMP.
    #[derive(Debug, Copy, Clone, PartialEq, Eq)]
    #[repr(transparent)]
    pub struct MappingFlags: u64 {
        /// No flags.
        const NONE           = 0x0;
        /// Allocate memory on the device and move data from host to device.
        const TO             = 0x01;
        /// Allocate memory on the device and move data from device to host.
        const FROM           = 0x02;
        /// Always perform the requested mapping action, even if already mapped.
        const ALWAYS         = 0x04;
        /// Delete the element from the device environment, ignoring ref count.
        const DELETE         = 0x08;
        /// The element being mapped is a pointer-pointee pair.
        const PTR_AND_OBJ    = 0x10;
        /// The base address should be passed to the target kernel as argument.
        const TARGET_PARAM   = 0x20;
        /// The runtime must return the device pointer.
        const RETURN_PARAM   = 0x40;
        /// The reference being passed is a pointer to private data.
        const PRIVATE        = 0x80;
        /// Pass the element by value.
        const LITERAL        = 0x100;
        /// Implicit map (generated by compiler, not explicit in code).
        const IMPLICIT       = 0x200;
        /// Hint to allocate memory close to the target device.
        const CLOSE          = 0x400;
        /// Reserved (0x800 in OpenMP for XLC compatibility).
        const RESERVED       = 0x800;
        /// Require that the data is already allocated on the device.
        const PRESENT        = 0x1000;
        /// Increment/decrement a separate ref counter (OpenACC compatibility).
        const OMPX_HOLD      = 0x2000;
        /// Used for non-contiguous list items in target update.
        const NON_CONTIG     = 0x100000000000;
        /// 16 MSBs indicate membership in a struct.
        const MEMBER_OF      = 0xffff000000000000;
    }
}

impl OffloadMetadata {
    pub fn from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Self {
        OffloadMetadata {
            payload_size: get_payload_size(tcx, ty),
            mode: MappingFlags::from_ty(tcx, ty),
        }
    }

    pub fn handle_abi<'tcx, C>(
        cx: &C,
        tcx: TyCtxt<'tcx>,
        ty: Ty<'tcx>,
        arg_abi: &ArgAbi<'tcx, Ty<'tcx>>,
    ) -> Vec<(Self, Ty<'tcx>)>
    where
        Ty<'tcx>: TyAbiInterface<'tcx, C>,
    {
        match arg_abi.layout.backend_repr {
            BackendRepr::ScalarPair { a: _, b: _, b_offset: _ } => (0..2)
                .map(|i| {
                    let ty = arg_abi.layout.field(cx, i).ty;
                    (OffloadMetadata::from_ty(tcx, ty), ty)
                })
                .collect(),
            _ => vec![(OffloadMetadata::from_ty(tcx, ty), ty)],
        }
    }
}

// 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() {
        ty::Slice(elem_ty) => {
            let layout = tcx.layout_of(PseudoCanonicalInput {
                typing_env: TypingEnv::fully_monomorphized(),
                value: *elem_ty,
            });
            OffloadSize::Slice { element_size: layout.unwrap().size.bytes() }
        }
        ty::RawPtr(inner, _) | ty::Ref(_, inner, _) => get_payload_size(tcx, *inner),
        _ => OffloadSize::Static(
            tcx.layout_of(PseudoCanonicalInput {
                typing_env: TypingEnv::fully_monomorphized(),
                value: ty,
            })
            .unwrap()
            .size
            .bytes(),
        ),
    }
}

impl MappingFlags {
    fn from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Self {
        use rustc_ast::Mutability::*;

        match ty.kind() {
            ty::Bool | ty::Char | ty::Int(_) | ty::Uint(_) | ty::Float(_) => {
                MappingFlags::LITERAL | MappingFlags::IMPLICIT
            }

            // FIXME: This should not treat aliases this way.
            ty::Adt(_, _) | ty::Tuple(_) | ty::Array(_, _) | ty::Alias(_, _) | ty::Param(_) => {
                MappingFlags::TO
            }

            ty::RawPtr(_, Not) | ty::Ref(_, _, Not) => MappingFlags::TO,

            ty::RawPtr(_, Mut) | ty::Ref(_, _, Mut) => MappingFlags::TO | MappingFlags::FROM,

            ty::Slice(_) | ty::Str | ty::Dynamic(_, _) => MappingFlags::TO | MappingFlags::FROM,

            ty::Foreign(_) | ty::Pat(_, _) | ty::UnsafeBinder(_) => {
                MappingFlags::TO | MappingFlags::FROM
            }

            ty::FnDef(_, _)
            | ty::FnPtr(_, _)
            | ty::Closure(_, _)
            | ty::CoroutineClosure(_, _)
            | ty::Coroutine(_, _)
            | ty::CoroutineWitness(_, _)
            | ty::Never
            | ty::Bound(_, _)
            | ty::Placeholder(_)
            | ty::Infer(_)
            | ty::Error(_) => {
                tcx.dcx()
                    .span_err(rustc_span::DUMMY_SP, format!("type `{ty:?}` cannot be offloaded"));
                MappingFlags::empty()
            }
        }
    }
}