Add --use-constant-memory-space flag, off by default

Author: LegNeato

crates/cuda_builder/src/lib.rs

@@ -130,6 +130,18 @@ pub struct CudaBuilder {
     ///
     /// `true` by default.
     pub override_libm: bool,
+    /// If `true`, the codegen will attempt to place `static` variables in CUDA's
+    /// constant memory, which is fast but limited in size (~64KB total across all
+    /// statics). The codegen avoids placing any single item too large, but it does not
+    /// track cumulative size. Exceeding the limit may cause `IllegalAddress` runtime
+    /// errors (CUDA error code: `700`).
+    ///
+    /// The default is `false`, which places all statics in global memory. This avoids
+    /// such errors but may reduce performance and use more general memory.
+    ///
+    /// Future versions may support smarter placement and user-controlled
+    /// packing/spilling strategies.
+    pub use_constant_memory_space: bool,
     /// Whether to generate any debug info and what level of info to generate.
     pub debug: DebugInfo,
     /// Additional arguments passed to cargo during `cargo build`.
@@ -155,6 +167,7 @@ impl CudaBuilder {
             emit: None,
             optix: false,
             override_libm: true,
+            use_constant_memory_space: false,
             debug: DebugInfo::None,
             build_args: vec![],
             final_module_path: None,
@@ -284,6 +297,22 @@ impl CudaBuilder {
         self
     }
 
+    /// If `true`, the codegen will attempt to place `static` variables in CUDA's
+    /// constant memory, which is fast but limited in size (~64KB total across all
+    /// statics). The codegen avoids placing any single item too large, but it does not
+    /// track cumulative size. Exceeding the limit may cause `IllegalAddress` runtime
+    /// errors (CUDA error code: `700`).
+    ///
+    /// If `false`, all statics are placed in global memory. This avoids such errors but
+    /// may reduce performance and use more general memory.
+    ///
+    /// Future versions may support smarter placement and user-controlled
+    /// packing/spilling strategies.
+    pub fn use_constant_memory_space(mut self, use_constant_memory_space: bool) -> Self {
+        self.use_constant_memory_space = use_constant_memory_space;
+        self
+    }
+
     /// An optional path where to dump LLVM IR of the final output the codegen will feed to libnvvm. Usually
     /// used for debugging.
     pub fn final_module_path(mut self, path: impl AsRef<Path>) -> Self {

crates/rustc_codegen_nvvm/src/context.rs

@@ -271,17 +271,30 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
         }
 
         if !is_mutable && self.type_is_freeze(ty) {
-            let layout = self.layout_of(ty);
-            if layout.size.bytes() > CONSTANT_MEMORY_SIZE_LIMIT_BYTES {
-                self.tcx.sess.dcx().warn(format!(
-                    "static `{}` exceeds the constant-memory limit; placing in global memory (performance may be reduced)",
-                    instance
-                ));
-                // Global memory
+            if !self.codegen_args.use_constant_memory_space {
+                // We aren't using constant memory, so put the instance in global memory.
                 AddressSpace(1)
             } else {
-                // Constant memory
-                AddressSpace(4)
+                // We are using constant memory, see if the instance will fit.
+                //
+                // FIXME(@LegNeato) ideally we keep track of what we have put into
+                // constant memory and when it is filled up spill instead of only
+                // spilling when a static is big. We'll probably want some packing
+                // strategy controlled by the user...for example, if you have one large
+                // static and many small ones, you might want the small ones to all be
+                // in constant memory or just the big one depending on your workload.
+                let layout = self.layout_of(ty);
+                if layout.size.bytes() > CONSTANT_MEMORY_SIZE_LIMIT_BYTES {
+                    self.tcx.sess.dcx().warn(format!(
+                    "static `{}` exceeds the constant memory limit; placing in global memory (performance may be reduced)",
+                    instance
+                ));
+                    // Place instance in global memory if it is too big for constant memory.
+                    AddressSpace(1)
+                } else {
+                    // Place instance in constant memory if it fits.
+                    AddressSpace(4)
+                }
             }
         } else {
             AddressSpace::DATA
@@ -534,6 +547,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
 pub struct CodegenArgs {
     pub nvvm_options: Vec<NvvmOption>,
     pub override_libm: bool,
+    pub use_constant_memory_space: bool,
     pub final_module_path: Option<PathBuf>,
 }
 
@@ -552,6 +566,8 @@ impl CodegenArgs {
                 cg_args.nvvm_options.push(flag);
             } else if arg == "--override-libm" {
                 cg_args.override_libm = true;
+            } else if arg == "--use-constant-memory-space" {
+                cg_args.use_constant_memory_space = true;
             } else if arg == "--final-module-path" {
                 cg_args.final_module_path = Some(PathBuf::from(
                     args.get(idx + 1).expect("No path for --final-module-path"),