Repository: rust-lang/rust
File: compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs (Rust)

Code before:
    cx: &CodegenCx<'ll, '_>,
    metadata: &[OffloadMetadata],
    symbol: String,
    offload_globals: &OffloadGlobals<'ll>,
) -> OffloadKernelGlobals<'ll> {
    if let Some(entry) = cx.offload_kernel_cache.borrow().get(&symbol) {
        return *entry;
    }

    let offload_entry_ty = offload_globals.offload_entry_ty;

    let (sizes, transfer): (Vec<_>, Vec<_>) =
        metadata.iter().map(|m| (m.payload_size, m.mode)).unzip();
    // Our begin mapper should only see simplified information about which args have to be
    // transferred to the device, the end mapper only about which args should be transferred back.
    // Any information beyond that makes it harder for LLVM's opt pass to evaluate whether it can
    // safely move (=optimize) the LLVM-IR location of this data transfer. Only the mapping types
    // mentioned below are handled, so make sure that we don't generate any other ones.
    let handled_mappings = MappingFlags::TO
        | MappingFlags::FROM
        | MappingFlags::TARGET_PARAM
        | MappingFlags::LITERAL
        | MappingFlags::IMPLICIT;
    for arg in &transfer {
        debug_assert!(!arg.contains_unknown_bits());
        debug_assert!(arg.difference(handled_mappings).is_empty());
    }

    let valid_begin_mappings = MappingFlags::TO | MappingFlags::LITERAL | MappingFlags::IMPLICIT;
    let transfer_to: Vec<u64> =
        transfer.iter().map(|m| m.intersection(valid_begin_mappings).bits()).collect();
    let transfer_from: Vec<u64> =
        transfer.iter().map(|m| m.intersection(MappingFlags::FROM).bits()).collect();
    // FIXME(offload): add `OMP_MAP_TARGET_PARAM = 0x20` only if necessary
    let transfer_kernel = vec![MappingFlags::TARGET_PARAM.bits(); transfer_to.len()];

    let offload_sizes = add_priv_unnamed_arr(&cx, &format!(".offload_sizes.{symbol}"), &sizes);
    let memtransfer_begin =
        add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.begin"), &transfer_to);
    let memtransfer_kernel =
        add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.kernel"), &transfer_kernel);
    let memtransfer_end =
        add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.end"), &transfer_from);

    // Next: For each function, generate these three entries. A weak constant,
    // the llvm.rodata entry name, and  the llvm_offload_entries value

    let name = format!(".{symbol}.region_id");
    let initializer = cx.get_const_i8(0);
    let region_id = add_global(&cx, &name, initializer, WeakAnyLinkage);


Review comment (suggestion):
Is there a `subset` operation for this type? That would be a little easier to read.

Code after:
    cx: &CodegenCx<'ll, '_>,
    metadata: &[OffloadMetadata],
    symbol: String,
    offload_globals: &OffloadGlobals<'ll>,
) -> OffloadKernelGlobals<'ll> {
    if let Some(entry) = cx.offload_kernel_cache.borrow().get(&symbol) {
        return *entry;
    }

    let offload_entry_ty = offload_globals.offload_entry_ty;

    let (sizes, transfer): (Vec<_>, Vec<_>) =
        metadata.iter().map(|m| (m.payload_size, m.mode)).unzip();
    // Our begin mapper should only see simplified information about which args have to be
    // transferred to the device, the end mapper only about which args should be transferred back.
    // Any information beyond that makes it harder for LLVM's opt pass to evaluate whether it can
    // safely move (=optimize) the LLVM-IR location of this data transfer. Only the mapping types
    // mentioned below are handled, so make sure that we don't generate any other ones.
    let handled_mappings = MappingFlags::TO
        | MappingFlags::FROM
        | MappingFlags::TARGET_PARAM
        | MappingFlags::LITERAL
        | MappingFlags::IMPLICIT;
    for arg in &transfer {
        debug_assert!(!arg.contains_unknown_bits());
        debug_assert!(handled_mappings.contains(*arg));
    }

    let valid_begin_mappings = MappingFlags::TO | MappingFlags::LITERAL | MappingFlags::IMPLICIT;
    let transfer_to: Vec<u64> =
        transfer.iter().map(|m| m.intersection(valid_begin_mappings).bits()).collect();
    let transfer_from: Vec<u64> =
        transfer.iter().map(|m| m.intersection(MappingFlags::FROM).bits()).collect();
    // FIXME(offload): add `OMP_MAP_TARGET_PARAM = 0x20` only if necessary
    let transfer_kernel = vec![MappingFlags::TARGET_PARAM.bits(); transfer_to.len()];

    let offload_sizes = add_priv_unnamed_arr(&cx, &format!(".offload_sizes.{symbol}"), &sizes);
    let memtransfer_begin =
        add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.begin"), &transfer_to);
    let memtransfer_kernel =
        add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.kernel"), &transfer_kernel);
    let memtransfer_end =
        add_priv_unnamed_arr(&cx, &format!(".offload_maptypes.{symbol}.end"), &transfer_from);

    // Next: For each function, generate these three entries. A weak constant,
    // the llvm.rodata entry name, and  the llvm_offload_entries value

    let name = format!(".{symbol}.region_id");
    let initializer = cx.get_const_i8(0);
    let region_id = add_global(&cx, &name, initializer, WeakAnyLinkage);