[wip] Add intrinsic for shared memory

TODO
- Docs for the intrinsic
- Allow only for amdgpu and nvptx
- Re-use the same global in a module
- Add tests

Author: Flakebi

compiler/rustc_abi/src/lib.rs

@@ -1394,6 +1394,8 @@ pub struct AddressSpace(pub u32);
 impl AddressSpace {
     /// The default address space, corresponding to data space.
     pub const DATA: Self = AddressSpace(0);
+    /// The address space for shared memory on nvptx and amdgpu.
+    pub const SHARED: Self = AddressSpace(3);
 }
 
 /// The way we represent values to the backend

compiler/rustc_codegen_llvm/src/declare.rs

@@ -12,6 +12,7 @@
 //! * When in doubt, define.
 
 use itertools::Itertools;
+use rustc_abi::AddressSpace;
 use rustc_codegen_ssa::traits::TypeMembershipCodegenMethods;
 use rustc_data_structures::fx::FxIndexSet;
 use rustc_middle::ty::{Instance, Ty};
@@ -90,6 +91,28 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
         unsafe { llvm::LLVMRustGetOrInsertGlobal(self.llmod, name.as_c_char_ptr(), name.len(), ty) }
     }
 
+    /// Declare a global value in a specific address space.
+    ///
+    /// If there’s 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.llmod,
+                name.as_c_char_ptr(),
+                name.len(),
+                ty,
+                addr_space.0,
+            )
+        }
+    }
+
     /// Declare a C ABI function.
     ///
     /// Only use this for foreign function ABIs and glue. For Rust functions use

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1,7 +1,7 @@
 use std::assert_matches::assert_matches;
 use std::cmp::Ordering;
 
-use rustc_abi::{self as abi, Align, Float, HasDataLayout, Primitive, Size};
+use rustc_abi::{self as abi, AddressSpace, Align, Float, HasDataLayout, Primitive, Size};
 use rustc_codegen_ssa::base::{compare_simd_types, wants_msvc_seh, wants_wasm_eh};
 use rustc_codegen_ssa::common::{IntPredicate, TypeKind};
 use rustc_codegen_ssa::errors::{ExpectedPointerMutability, InvalidMonomorphization};
@@ -551,6 +551,25 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 return Ok(());
             }
 
+            sym::dynamic_shared_memory => {
+                // TODO Allow only on nvptx and amdgpu
+                // TODO create only one global per module, take the maximum alignment of all calls
+                let global = self.declare_global_in_addrspace(
+                    "",
+                    self.type_array(self.type_i8(), 0),
+                    AddressSpace::SHARED,
+                );
+                // TODO Set alignment
+                let ty::RawPtr(inner_ty, _) = ret_ty.kind() else { unreachable!() };
+                let alignment = self.align_of(*inner_ty).bytes() as u32;
+                unsafe {
+                    if alignment > llvm::LLVMGetAlignment(global) {
+                        llvm::LLVMSetAlignment(global, alignment);
+                    }
+                }
+                self.cx().const_pointercast(global, self.type_ptr())
+            }
+
             _ if name.as_str().starts_with("simd_") => {
                 // Unpack non-power-of-2 #[repr(packed, simd)] arguments.
                 // This gives them the expected layout of a regular #[repr(simd)] vector.

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1777,6 +1777,13 @@ unsafe extern "C" {
         NameLen: size_t,
         T: &'a Type,
     ) -> &'a Value;
+    pub fn LLVMRustGetOrInsertGlobalInAddrspace<'a>(
+        M: &'a Module,
+        Name: *const c_char,
+        NameLen: size_t,
+        T: &'a Type,
+        AddressSpace: c_uint,
+    ) -> &'a Value;
     pub fn LLVMRustInsertPrivateGlobal<'a>(M: &'a Module, T: &'a Type) -> &'a Value;
     pub fn LLVMRustGetNamedValue(
         M: &Module,

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -92,6 +92,7 @@ pub fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -
         | sym::min_align_of
         | sym::needs_drop
         | sym::caller_location
+        | sym::dynamic_shared_memory
         | sym::add_with_overflow
         | sym::sub_with_overflow
         | sym::mul_with_overflow
@@ -251,6 +252,7 @@ pub fn check_intrinsic_type(
             }
             sym::rustc_peek => (1, 0, vec![param(0)], param(0)),
             sym::caller_location => (0, 0, vec![], tcx.caller_location_ty()),
+            sym::dynamic_shared_memory => (1, 0, vec![], Ty::new_mut_ptr(tcx, param(0))),
             sym::assert_inhabited
             | sym::assert_zero_valid
             | sym::assert_mem_uninitialized_valid => (1, 0, vec![], tcx.types.unit),

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -205,10 +205,11 @@ extern "C" LLVMValueRef LLVMRustGetOrInsertFunction(LLVMModuleRef M,
                   .getCallee());
 }
 
-extern "C" LLVMValueRef LLVMRustGetOrInsertGlobal(LLVMModuleRef M,
-                                                  const char *Name,
-                                                  size_t NameLen,
-                                                  LLVMTypeRef Ty) {
+extern "C" LLVMValueRef LLVMRustGetOrInsertGlobalInAddrspace(LLVMModuleRef M,
+                                                             const char *Name,
+                                                             size_t NameLen,
+                                                             LLVMTypeRef Ty,
+                                                             unsigned AddressSpace) {
   Module *Mod = unwrap(M);
   auto NameRef = StringRef(Name, NameLen);
 
@@ -219,10 +220,20 @@ extern "C" LLVMValueRef LLVMRustGetOrInsertGlobal(LLVMModuleRef M,
   GlobalVariable *GV = Mod->getGlobalVariable(NameRef, true);
   if (!GV)
     GV = new GlobalVariable(*Mod, unwrap(Ty), false,
-                            GlobalValue::ExternalLinkage, nullptr, NameRef);
+                            GlobalValue::ExternalLinkage, nullptr, NameRef,
+                            nullptr, GlobalValue::NotThreadLocal, AddressSpace);
   return wrap(GV);
 }
 
+extern "C" LLVMValueRef LLVMRustGetOrInsertGlobal(LLVMModuleRef M,
+                                                  const char *Name,
+                                                  size_t NameLen,
+                                                  LLVMTypeRef Ty) {
+  Module *Mod = unwrap(M);
+  unsigned AddressSpace = Mod->getDataLayout().getDefaultGlobalsAddressSpace();
+  return LLVMRustGetOrInsertGlobalInAddrspace(M, Name, NameLen, Ty, AddressSpace);
+}
+
 extern "C" LLVMValueRef LLVMRustInsertPrivateGlobal(LLVMModuleRef M,
                                                     LLVMTypeRef Ty) {
   return wrap(new GlobalVariable(*unwrap(M), unwrap(Ty), false,

compiler/rustc_span/src/symbol.rs

@@ -825,6 +825,7 @@ symbols! {
         dyn_star,
         dyn_trait,
         dynamic_no_pic: "dynamic-no-pic",
+        dynamic_shared_memory,
         e,
         edition_panic,
         effects,

library/core/src/intrinsics/mod.rs

@@ -4234,6 +4234,24 @@ pub const unsafe fn min_align_of_val<T: ?Sized>(_ptr: *const T) -> usize {
     unreachable!()
 }
 
+/// Returns a pointer to dynamic shared memory
+/// The return pointer is aligned by properly.
+/// All pointers returned by `dynamic_shared_memory` point to the same address,
+/// so alias the same memory.
+///
+/// # Other APIs
+///
+/// CUDA and HIP call this shared memory.
+/// OpenCL and SYCL call this local memory.
+#[cfg(not(bootstrap))]
+#[rustc_nounwind]
+#[unstable(feature = "dynamic_shared_memory", issue = "135513")]
+#[rustc_intrinsic]
+#[rustc_intrinsic_must_be_overridden]
+pub fn dynamic_shared_memory<T: ?Sized>() -> *mut T {
+    unreachable!()
+}
+
 /// Gets a static string slice containing the name of a type.
 ///
 /// Note that, unlike most intrinsics, this is safe to call;