Files
rust/compiler/rustc_codegen_llvm/src/declare.rs
T
Flakebi 13ec3de673 Add intrinsic for launch-sized workgroup memory on GPUs
Workgroup memory is a memory region that is shared between all
threads in a workgroup on GPUs. Workgroup memory can be allocated
statically or after compilation, when launching a gpu-kernel.
The intrinsic added here returns the pointer to the memory that is
allocated at launch-time.

# Interface

With this change, workgroup memory can be accessed in Rust by
calling the new `gpu_launch_sized_workgroup_mem<T>() -> *mut T`
intrinsic.

It returns the pointer to workgroup memory guaranteeing that it is
aligned to at least the alignment of `T`.
The pointer is dereferencable for the size specified when launching the
current gpu-kernel (which may be the size of `T` but can also be larger
or smaller or zero).

All calls to this intrinsic return a pointer to the same address.

See the intrinsic documentation for more details.

## Alternative Interfaces

It was also considered to expose dynamic workgroup memory as extern
static variables in Rust, like they are represented in LLVM IR.
However, due to the pointer not being guaranteed to be dereferencable
(that depends on the allocated size at runtime), such a global must be
zero-sized, which makes global variables a bad fit.

# Implementation Details

Workgroup memory in amdgpu and nvptx lives in address space 3.
Workgroup memory from a launch is implemented by creating an
external global variable in address space 3. The global is declared with
size 0, as the actual size is only known at runtime. It is defined
behavior in LLVM to access an external global outside the defined size.

There is no similar way to get the allocated size of launch-sized
workgroup memory on amdgpu an nvptx, so users have to pass this
out-of-band or rely on target specific ways for now.
2026-04-24 10:03:45 +02:00

278 lines
9.9 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! Declare various LLVM values.
//!
//! Prefer using functions and methods from this module rather than calling LLVM
//! functions directly. These functions do some additional work to ensure we do
//! the right thing given the preconceptions of codegen.
//!
//! Some useful guidelines:
//!
//! * Use declare_* family of methods if you are declaring, but are not
//! interested in defining the Value they return.
//! * Use define_* family of methods when you might be defining the Value.
//! * When in doubt, define.
use std::borrow::Borrow;
use itertools::Itertools;
use rustc_abi::AddressSpace;
use rustc_codegen_ssa::traits::{MiscCodegenMethods, TypeMembershipCodegenMethods};
use rustc_data_structures::fx::FxIndexSet;
use rustc_middle::ty::{Instance, Ty};
use rustc_sanitizers::{cfi, kcfi};
use rustc_target::callconv::FnAbi;
use smallvec::SmallVec;
use tracing::debug;
use crate::abi::FnAbiLlvmExt;
use crate::attributes;
use crate::common::AsCCharPtr;
use crate::context::{CodegenCx, GenericCx, SCx, SimpleCx};
use crate::llvm::AttributePlace::Function;
use crate::llvm::{self, FromGeneric, Type, Value, Visibility};
/// Declare a function with a SimpleCx.
///
/// If theres a value with the same name already declared, the function will
/// update the declaration and return existing Value instead.
pub(crate) fn declare_simple_fn<'ll>(
cx: &SimpleCx<'ll>,
name: &str,
callconv: llvm::CallConv,
unnamed: llvm::UnnamedAddr,
visibility: llvm::Visibility,
ty: &'ll Type,
) -> &'ll Value {
debug!("declare_simple_fn(name={:?}, ty={:?})", name, ty);
let llfn = unsafe {
llvm::LLVMRustGetOrInsertFunction(cx.llmod, name.as_c_char_ptr(), name.len(), ty)
};
llvm::SetFunctionCallConv(llfn, callconv);
llvm::set_unnamed_address(llfn, unnamed);
llvm::set_visibility(llfn, visibility);
llfn
}
/// Declare a function.
///
/// If theres a value with the same name already declared, the function will
/// update the declaration and return existing Value instead.
pub(crate) fn declare_raw_fn<'ll, 'tcx>(
cx: &CodegenCx<'ll, 'tcx>,
name: &str,
callconv: llvm::CallConv,
unnamed: llvm::UnnamedAddr,
visibility: llvm::Visibility,
ty: &'ll Type,
) -> &'ll Value {
debug!("declare_raw_fn(name={:?}, ty={:?})", name, ty);
let llfn = declare_simple_fn(cx, name, callconv, unnamed, visibility, ty);
let mut attrs = SmallVec::<[_; 4]>::new();
if cx.sess().target.is_like_gpu {
// Conservatively apply convergent to all functions in case they may call
// a convergent function. Rely on LLVM to optimize away the unnecessary
// convergent attributes.
attrs.push(llvm::AttributeKind::Convergent.create_attr(cx.llcx));
}
if cx.tcx.sess.opts.cg.no_redzone.unwrap_or(cx.tcx.sess.target.disable_redzone) {
attrs.push(llvm::AttributeKind::NoRedZone.create_attr(cx.llcx));
}
attrs.extend(attributes::non_lazy_bind_attr(cx, cx.tcx.sess));
attributes::apply_to_llfn(llfn, Function, &attrs);
llfn
}
impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
/// Declare a global value.
///
/// If theres a value with the same name already declared, the function will
/// return its Value instead.
pub(crate) fn declare_global(&self, name: &str, ty: &'ll Type) -> &'ll Value {
debug!("declare_global(name={:?})", name);
unsafe {
llvm::LLVMRustGetOrInsertGlobal(
(**self).borrow().llmod,
name.as_c_char_ptr(),
name.len(),
ty,
)
}
}
/// Declare a global value in a specific address space.
///
/// If theres a value with the same name already declared, the function will
/// return its Value instead.
pub(crate) fn declare_global_in_addrspace(
&self,
name: &str,
ty: &'ll Type,
addr_space: AddressSpace,
) -> &'ll Value {
debug!("declare_global(name={name:?}, addrspace={addr_space:?})");
unsafe {
llvm::LLVMRustGetOrInsertGlobalInAddrspace(
(**self).borrow().llmod,
name.as_c_char_ptr(),
name.len(),
ty,
addr_space.0,
)
}
}
}
impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
/// Declare a C ABI function.
///
/// Only use this for foreign function ABIs and glue. For Rust functions use
/// `declare_fn` instead.
///
/// If theres a value with the same name already declared, the function will
/// update the declaration and return existing Value instead.
pub(crate) fn declare_cfn(
&self,
name: &str,
unnamed: llvm::UnnamedAddr,
fn_type: &'ll Type,
) -> &'ll Value {
// Visibility should always be default for declarations, otherwise the linker may report an
// error.
declare_raw_fn(self, name, llvm::CCallConv, unnamed, Visibility::Default, fn_type)
}
/// Declare an entry Function
///
/// The ABI of this function can change depending on the target (although for now the same as
/// `declare_cfn`)
///
/// If theres a value with the same name already declared, the function will
/// update the declaration and return existing Value instead.
pub(crate) fn declare_entry_fn(
&self,
name: &str,
callconv: llvm::CallConv,
unnamed: llvm::UnnamedAddr,
fn_type: &'ll Type,
) -> &'ll Value {
let visibility = Visibility::from_generic(self.tcx.sess.default_visibility());
declare_raw_fn(self, name, callconv, unnamed, visibility, fn_type)
}
/// Declare a Rust function.
///
/// If theres a value with the same name already declared, the function will
/// update the declaration and return existing Value instead.
pub(crate) fn declare_fn(
&self,
name: &str,
fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
instance: Option<Instance<'tcx>>,
) -> &'ll Value {
debug!("declare_rust_fn(name={:?}, fn_abi={:?})", name, fn_abi);
// Function addresses in Rust are never significant, allowing functions to
// be merged.
let llfn = declare_raw_fn(
self,
name,
fn_abi.llvm_cconv(self),
llvm::UnnamedAddr::Global,
llvm::Visibility::Default,
fn_abi.llvm_type(self),
);
fn_abi.apply_attrs_llfn(self, llfn, instance);
if self.tcx.sess.is_sanitizer_cfi_enabled() {
if let Some(instance) = instance {
let mut typeids = FxIndexSet::default();
for options in [
cfi::TypeIdOptions::GENERALIZE_POINTERS,
cfi::TypeIdOptions::NORMALIZE_INTEGERS,
cfi::TypeIdOptions::USE_CONCRETE_SELF,
]
.into_iter()
.powerset()
.map(cfi::TypeIdOptions::from_iter)
{
let typeid = cfi::typeid_for_instance(self.tcx, instance, options);
if typeids.insert(typeid.clone()) {
self.add_type_metadata(llfn, typeid.as_bytes());
}
}
} else {
for options in [
cfi::TypeIdOptions::GENERALIZE_POINTERS,
cfi::TypeIdOptions::NORMALIZE_INTEGERS,
]
.into_iter()
.powerset()
.map(cfi::TypeIdOptions::from_iter)
{
let typeid = cfi::typeid_for_fnabi(self.tcx, fn_abi, options);
self.add_type_metadata(llfn, typeid.as_bytes());
}
}
}
if self.tcx.sess.is_sanitizer_kcfi_enabled() {
// LLVM KCFI does not support multiple !kcfi_type attachments
let mut options = kcfi::TypeIdOptions::empty();
if self.tcx.sess.is_sanitizer_cfi_generalize_pointers_enabled() {
options.insert(kcfi::TypeIdOptions::GENERALIZE_POINTERS);
}
if self.tcx.sess.is_sanitizer_cfi_normalize_integers_enabled() {
options.insert(kcfi::TypeIdOptions::NORMALIZE_INTEGERS);
}
if let Some(instance) = instance {
let kcfi_typeid = kcfi::typeid_for_instance(self.tcx, instance, options);
self.set_kcfi_type_metadata(llfn, kcfi_typeid);
} else {
let kcfi_typeid = kcfi::typeid_for_fnabi(self.tcx, fn_abi, options);
self.set_kcfi_type_metadata(llfn, kcfi_typeid);
}
}
llfn
}
}
impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
/// Declare a global with an intention to define it.
///
/// Use this function when you intend to define a global. This function will
/// return `None` if the name already has a definition associated with it. In that
/// case an error should be reported to the user, because it usually happens due
/// to users fault (e.g., misuse of `#[no_mangle]` or `#[export_name]` attributes).
pub(crate) fn define_global(&self, name: &str, ty: &'ll Type) -> Option<&'ll Value> {
if self.get_defined_value(name).is_some() {
None
} else {
Some(self.declare_global(name, ty))
}
}
/// Gets declared value by name.
pub(crate) fn get_declared_value(&self, name: &str) -> Option<&'ll Value> {
debug!("get_declared_value(name={:?})", name);
unsafe { llvm::LLVMRustGetNamedValue(self.llmod(), name.as_c_char_ptr(), name.len()) }
}
/// Gets defined or externally defined (AvailableExternally linkage) value by
/// name.
pub(crate) fn get_defined_value(&self, name: &str) -> Option<&'ll Value> {
self.get_declared_value(name).and_then(|val| {
let declaration = llvm::is_declaration(val);
if !declaration { Some(val) } else { None }
})
}
}