Feat: add address spaces for user defined statics

Author: RDambrosio016

crates/cuda_std_macros/src/lib.rs

@@ -208,3 +208,31 @@ pub fn externally_visible(
 
     func.into_token_stream().into()
 }
+
+/// Notifies the codegen to put a `static`/`static mut` inside of a specific memory address space.
+/// This is mostly for internal use and/or advanced users, as the codegen and `cuda_std` handle address space placement
+/// implicitly. **Improper use of this macro could yield weird or undefined behavior**.
+///
+/// This macro takes a single argument which can either be `global`, `shared`, `constant`, or `local`.
+///
+/// This macro does nothing on the CPU.
+#[proc_macro_attribute]
+pub fn address_space(attr: proc_macro::TokenStream, item: proc_macro::TokenStream) -> TokenStream {
+    let mut global = syn::parse_macro_input!(item as syn::ItemStatic);
+    let input = syn::parse_macro_input!(attr as Ident);
+
+    let addrspace_num = match input.to_string().as_str() {
+        "global" => 1,
+        // what did you do to address space 2 libnvvm??
+        "shared" => 3,
+        "constant" => 4,
+        "local" => 5,
+        addr => panic!("Invalid address space `{}`", addr),
+    };
+
+    let new_attr =
+        parse_quote!(#[cfg_attr(target_os = "cuda", nvvm_internal(addrspace(#addrspace_num)))]);
+    global.attrs.push(new_attr);
+
+    global.into_token_stream().into()
+}

crates/rustc_codegen_nvvm/CHANGELOG.md

@@ -4,6 +4,17 @@ Notable changes to this project will be documented in this file.
 
 ## Unreleased
 
+### Address Spaces 
+
+CUDA Address Spaces have been mostly implemented. Statics that are not mut statics and do not rely on
+interior mutability (are "freeze" types) are placed in constant memory (`__constant__` in CUDA C++), otherwise
+they are placed in global memory (`__global__`). Currently this only happens for user-defined statics, not for
+codegen-internal globals such as intermediate alloc globals.
+
+An `#[address_space(...)]` macro has been added to cuda_std to complement this change. However, this macro
+is mostly just for advanced users and internal support for things like shared memory. Improper use can 
+cause undefined behavior, so its use is generally discouraged.
+
 ### Dead Code Elimination 
 
 PTX files no longer include useless functions and globals, we have switched to an alternative

crates/rustc_codegen_nvvm/rustc_llvm_wrapper/RustWrapper.cpp

@@ -187,22 +187,20 @@ extern "C" LLVMValueRef LLVMRustGetOrInsertFunction(LLVMModuleRef M,
 }
 
 extern "C" LLVMValueRef
-LLVMRustGetOrInsertGlobal(LLVMModuleRef M, const char *Name, LLVMTypeRef Ty, unsigned AddressSpace)
-{
-  // Module *Mod = unwrap(M);
-  // GlobalVariable *GV = dyn_cast_or_null<GlobalVariable>(Mod->getNamedValue(Name));
-  // if (!GV)
-  // {
-  //   GV = new GlobalVariable(unwrap(Ty), false, GlobalValue::ExternalLinkage,
-  //                           nullptr, Name, GlobalValue::NotThreadLocal, AddressSpace);
-  // }
-  // Type *GVTy = GV->getType();
-  // PointerType *PTy = PointerType::get(unwrap(Ty), GVTy->getPointerAddressSpace());
-  // if (GVTy != PTy)
-  //   return wrap(ConstantExpr::getBitCast(GV, PTy));
-
-  return wrap(unwrap(M)->getOrInsertGlobal(Name, unwrap(Ty)));
-  // return wrap(GV);
+LLVMRustGetOrInsertGlobal(LLVMModuleRef M, const char *Name, size_t NameLen, LLVMTypeRef Ty, unsigned AddressSpace)
+{
+  Module *Mod = unwrap(M);
+  StringRef NameRef(Name, NameLen);
+
+  // We don't use Module::getOrInsertGlobal because that returns a Constant*,
+  // which may either be the real GlobalVariable*, or a constant bitcast of it
+  // if our type doesn't match the original declaration. We always want the
+  // GlobalVariable* so we can access linkage, visibility, etc.
+  GlobalVariable *GV = Mod->getGlobalVariable(NameRef, true);
+  if (!GV)
+    GV = new GlobalVariable(*Mod, unwrap(Ty), false,
+                            GlobalValue::ExternalLinkage, nullptr, NameRef, nullptr, GlobalValue::NotThreadLocal, AddressSpace);
+  return wrap(GV);
 }
 
 extern "C" LLVMTypeRef LLVMRustMetadataTypeInContext(LLVMContextRef C)

crates/rustc_codegen_nvvm/src/attributes.rs

@@ -1,5 +1,5 @@
 use crate::llvm::{self, AttributePlace::*, Value};
-use rustc_ast::Attribute;
+use rustc_ast::{Attribute, Lit, LitKind};
 use rustc_attr::{InlineAttr, OptimizeAttr};
 use rustc_middle::{middle::codegen_fn_attrs::CodegenFnAttrFlags, ty};
 use rustc_session::{config::OptLevel, Session};
@@ -93,13 +93,15 @@ pub(crate) fn from_fn_attrs<'ll, 'tcx>(
 pub struct Symbols {
     pub nvvm_internal: Symbol,
     pub kernel: Symbol,
+    pub addrspace: Symbol,
 }
 
 // inspired by rust-gpu's attribute handling
 #[derive(Default, Clone, PartialEq)]
 pub(crate) struct NvvmAttributes {
     pub kernel: bool,
     pub used: bool,
+    pub addrspace: Option<u8>,
 }
 
 impl NvvmAttributes {
@@ -116,6 +118,21 @@ impl NvvmAttributes {
                     if arg.has_name(sym::used) {
                         nvvm_attrs.used = true;
                     }
+                    if arg.has_name(cx.symbols.addrspace) {
+                        let args = arg.meta_item_list().unwrap_or_default();
+                        if let Some(arg) = args.first() {
+                            let lit = arg.literal();
+                            if let Some(Lit {
+                                kind: LitKind::Int(val, _),
+                                ..
+                            }) = lit
+                            {
+                                nvvm_attrs.addrspace = Some(*val as u8);
+                            } else {
+                                panic!();
+                            }
+                        }
+                    }
                 }
             }
         }

crates/rustc_codegen_nvvm/src/builder.rs

@@ -17,7 +17,7 @@ use rustc_middle::ty::layout::{
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_span::Span;
 use rustc_target::abi::call::FnAbi;
-use rustc_target::abi::{self, Align, Size, WrappingRange};
+use rustc_target::abi::{self, AddressSpace, Align, Size, WrappingRange};
 use rustc_target::spec::{HasTargetSpec, Target};
 use std::borrow::Cow;
 use std::ffi::{CStr, CString};
@@ -870,6 +870,21 @@ impl<'ll, 'tcx, 'a> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
     fn bitcast(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {
         trace!("Bitcast `{:?}` to ty `{:?}`", val, dest_ty);
         unsafe {
+            let ty = llvm::LLVMRustGetValueType(val);
+            let kind = llvm::LLVMRustGetTypeKind(ty);
+            if kind == llvm::TypeKind::Pointer {
+                let element = llvm::LLVMGetElementType(ty);
+                let addrspace = llvm::LLVMGetPointerAddressSpace(ty);
+                let new_ty = self.type_ptr_to_ext(element, AddressSpace::DATA);
+                if addrspace != 0 {
+                    return llvm::LLVMBuildAddrSpaceCast(
+                        *self.llbuilder.lock().unwrap(),
+                        val,
+                        new_ty,
+                        unnamed(),
+                    );
+                }
+            }
             llvm::LLVMBuildBitCast(&mut self.llbuilder.lock().unwrap(), val, dest_ty, unnamed())
         }
     }
@@ -1212,8 +1227,19 @@ impl<'ll, 'tcx, 'a> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
 
 impl<'a, 'll, 'tcx> StaticBuilderMethods for Builder<'a, 'll, 'tcx> {
     fn get_static(&mut self, def_id: DefId) -> &'ll Value {
-        // Forward to the `get_static` method of `CodegenCx`
-        self.cx().get_static(def_id)
+        unsafe {
+            let mut g = self.cx.get_static(def_id);
+            let llty = llvm::LLVMRustGetValueType(g);
+            let addrspace = AddressSpace(llvm::LLVMGetPointerAddressSpace(llty));
+            if addrspace != AddressSpace::DATA {
+                trace!("Remapping global address space of global {:?}", g);
+                let llty = llvm::LLVMGetElementType(llty);
+                let ty = self.type_ptr_to_ext(llty, AddressSpace::DATA);
+                let builder = &*self.llbuilder.lock().unwrap();
+                g = llvm::LLVMBuildAddrSpaceCast(builder, g, ty, unnamed());
+            }
+            g
+        }
     }
 }
 

crates/rustc_codegen_nvvm/src/const_ty.rs

@@ -110,7 +110,12 @@ impl<'ll, 'tcx> ConstMethods<'tcx> for CodegenCx<'ll, 'tcx> {
         }
     }
 
-    fn scalar_to_backend(&self, cv: Scalar, layout: abi::Scalar, llty: &'ll Type) -> &'ll Value {
+    fn scalar_to_backend(
+        &self,
+        cv: Scalar,
+        layout: abi::Scalar,
+        mut llty: &'ll Type,
+    ) -> &'ll Value {
         trace!("Scalar to backend `{:?}`, `{:?}`, `{:?}`", cv, layout, llty);
         let bitsize = if layout.is_bool() {
             1
@@ -152,7 +157,11 @@ impl<'ll, 'tcx> ConstMethods<'tcx> for CodegenCx<'ll, 'tcx> {
                     GlobalAlloc::Static(def_id) => {
                         assert!(self.tcx.is_static(def_id));
                         assert!(!self.tcx.is_thread_local_static(def_id));
-                        (self.get_static(def_id), AddressSpace::DATA)
+                        let val = self.get_static(def_id);
+                        let addrspace = unsafe {
+                            llvm::LLVMGetPointerAddressSpace(llvm::LLVMRustGetValueType(val))
+                        };
+                        (self.get_static(def_id), AddressSpace(addrspace))
                     }
                 };
                 let llval = unsafe {
@@ -162,15 +171,23 @@ impl<'ll, 'tcx> ConstMethods<'tcx> for CodegenCx<'ll, 'tcx> {
                         1,
                     )
                 };
+
                 if layout.value != Pointer {
                     unsafe { llvm::LLVMConstPtrToInt(llval, llty) }
                 } else {
+                    if base_addr_space != AddressSpace::DATA {
+                        unsafe {
+                            let element = llvm::LLVMGetElementType(llty);
+                            llty = self.type_ptr_to_ext(element, base_addr_space);
+                        }
+                    }
                     self.const_bitcast(llval, llty)
                 }
             }
         };
 
         trace!("...Scalar to backend: `{:?}`", val);
+        trace!("{:?}", std::backtrace::Backtrace::force_capture());
 
         val
     }

crates/rustc_codegen_nvvm/src/consts.rs

@@ -213,7 +213,9 @@ fn check_and_apply_linkage<'ll, 'tcx>(
     ty: Ty<'tcx>,
     sym: &str,
     span_def_id: DefId,
+    instance: Instance<'tcx>,
 ) -> &'ll Value {
+    let addrspace = cx.static_addrspace(instance);
     let llty = cx.layout_of(ty).llvm_type(cx);
     if let Some(linkage) = attrs.linkage {
         // https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html#linkage-types-nvvm
@@ -239,7 +241,7 @@ fn check_and_apply_linkage<'ll, 'tcx>(
         };
         unsafe {
             // Declare a symbol `foo` with the desired linkage.
-            let g1 = cx.declare_global(sym, llty2, AddressSpace::DATA);
+            let g1 = cx.declare_global(sym, llty2, addrspace);
             llvm::LLVMRustSetLinkage(g1, linkage_to_llvm(linkage));
 
             // Declare an internal global `extern_with_linkage_foo` which
@@ -251,7 +253,7 @@ fn check_and_apply_linkage<'ll, 'tcx>(
             let mut real_name = "_rust_extern_with_linkage_".to_string();
             real_name.push_str(sym);
             let g2 = cx
-                .define_global(&real_name, llty, AddressSpace::DATA)
+                .define_global(&real_name, llty, addrspace)
                 .unwrap_or_else(|| {
                     cx.sess().span_fatal(
                         cx.tcx.def_span(span_def_id),
@@ -263,7 +265,7 @@ fn check_and_apply_linkage<'ll, 'tcx>(
             g2
         }
     } else {
-        cx.declare_global(sym, llty, AddressSpace::DATA)
+        cx.declare_global(sym, llty, addrspace)
     }
 }
 
@@ -322,7 +324,8 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
                 }
             }
 
-            let g = self.declare_global(sym, llty, AddressSpace::DATA);
+            let addrspace = self.static_addrspace(instance);
+            let g = self.declare_global(sym, llty, addrspace);
 
             if !self.tcx.is_reachable_non_generic(def_id) {
                 unsafe {
@@ -332,7 +335,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
 
             g
         } else {
-            check_and_apply_linkage(self, fn_attrs, ty, sym, def_id)
+            check_and_apply_linkage(self, fn_attrs, ty, sym, def_id, instance)
         };
 
         if fn_attrs.flags.contains(CodegenFnAttrFlags::THREAD_LOCAL) {
@@ -406,12 +409,13 @@ impl<'ll, 'tcx> StaticMethods for CodegenCx<'ll, 'tcx> {
                 let linkage = llvm::LLVMRustGetLinkage(g);
                 let visibility = llvm::LLVMRustGetVisibility(g);
 
+                let addrspace = self.static_addrspace(instance);
                 let new_g = llvm::LLVMRustGetOrInsertGlobal(
                     self.llmod,
                     name.as_ptr().cast(),
                     name.len(),
                     val_llty,
-                    AddressSpace::DATA.0,
+                    addrspace.0,
                 );
 
                 llvm::LLVMRustSetLinkage(new_g, linkage);
@@ -434,7 +438,7 @@ impl<'ll, 'tcx> StaticMethods for CodegenCx<'ll, 'tcx> {
             if !is_mutable && self.type_is_freeze(ty) {
                 // TODO(RDambrosio016): is this the same as putting this in
                 // the __constant__ addrspace for nvvm? should we set this addrspace explicitly?
-                llvm::LLVMSetGlobalConstant(g, llvm::True);
+                // llvm::LLVMSetGlobalConstant(g, llvm::True);
             }
 
             debug_info::create_global_var_metadata(self, def_id, g);
@@ -446,6 +450,7 @@ impl<'ll, 'tcx> StaticMethods for CodegenCx<'ll, 'tcx> {
             if attrs.flags.contains(CodegenFnAttrFlags::USED) {
                 self.add_used_global(g);
             }
+            trace!("Codegen static `{:?}`", g);
         }
     }
 

crates/rustc_codegen_nvvm/src/context.rs

@@ -1,16 +1,16 @@
 use crate::abi::FnAbiLlvmExt;
-use crate::attributes::{self, Symbols};
+use crate::attributes::{self, NvvmAttributes, Symbols};
 use crate::debug_info::{self, compile_unit_metadata, CrateDebugContext};
 use crate::llvm::{self, BasicBlock, Type, Value};
 use crate::{target, LlvmMod};
 use nvvm::NvvmOption;
-use rustc_codegen_ssa::traits::ConstMethods;
 use rustc_codegen_ssa::traits::{BackendTypes, BaseTypeMethods, CoverageInfoMethods, MiscMethods};
+use rustc_codegen_ssa::traits::{ConstMethods, DerivedTypeMethods};
 use rustc_data_structures::base_n;
 use rustc_hash::FxHashMap;
 use rustc_middle::dep_graph::DepContext;
 use rustc_middle::ty::layout::{
-    FnAbiError, FnAbiOf, FnAbiRequest, HasParamEnv, LayoutError, TyAndLayout,
+    FnAbiError, FnAbiOf, FnAbiRequest, HasParamEnv, HasTyCtxt, LayoutError, TyAndLayout,
 };
 use rustc_middle::ty::layout::{FnAbiOfHelpers, LayoutOfHelpers};
 use rustc_middle::ty::{Ty, TypeFoldable};
@@ -158,6 +158,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
             symbols: Symbols {
                 nvvm_internal: Symbol::intern("nvvm_internal"),
                 kernel: Symbol::intern("kernel"),
+                addrspace: Symbol::intern("addrspace"),
             },
             dbg_cx,
             codegen_args: CodegenArgs::from_session(tcx.sess()),
@@ -261,6 +262,24 @@ impl<'ll, 'tcx> MiscMethods<'tcx> for CodegenCx<'ll, 'tcx> {
 }
 
 impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
+    /// Computes the address space for a static.
+    pub fn static_addrspace(&self, instance: Instance<'tcx>) -> AddressSpace {
+        let ty = instance.ty(self.tcx, ty::ParamEnv::reveal_all());
+        let is_mutable = self.tcx().is_mutable_static(instance.def_id());
+        let attrs = self.tcx.get_attrs(instance.def_id());
+        let nvvm_attrs = NvvmAttributes::parse(self, attrs);
+
+        if let Some(addr) = nvvm_attrs.addrspace {
+            return AddressSpace(addr as u32);
+        }
+
+        if !is_mutable && self.type_is_freeze(ty) {
+            AddressSpace(4)
+        } else {
+            AddressSpace::DATA
+        }
+    }
+
     /// Declare a global value, returns the existing value if it was already declared.
     pub fn declare_global(
         &self,

crates/rustc_codegen_nvvm/src/link.rs

@@ -227,22 +227,10 @@ fn codegen_into_ptx_file(
     // in this crate. We must unpack them and devour their bitcode to link in.
     for rlib in rlibs {
         let mut cgus = Vec::with_capacity(16);
-        // just pick the first cgu name as the overall name for now.
-        let mut name = String::new();
         for entry in Archive::new(File::open(rlib)?).entries()? {
             let mut entry = entry?;
             // metadata is where rustc puts rlib metadata, so its not a cgu we are interested in.
             if entry.path().unwrap() != Path::new(".metadata") {
-                if name == String::new() {
-                    name = entry
-                        .path()
-                        .unwrap()
-                        .file_name()
-                        .unwrap()
-                        .to_str()
-                        .unwrap()
-                        .to_string();
-                }
                 // std::fs::read adds 1 to the size, so do the same here - see comment:
                 // https://github.com/rust-lang/rust/blob/72868e017bdade60603a25889e253f556305f996/library/std/src/fs.rs#L200-L202
                 let mut bitcode = Vec::with_capacity(entry.size() as usize + 1);