Feat: add libm overriding

Author: RDambrosio016

crates/cuda_builder/src/lib.rs

@@ -107,6 +107,16 @@ pub struct CudaBuilder {
     ///
     /// `false` by default.
     pub optix: bool,
+    /// Whether to override calls to [`libm`](https://docs.rs/libm/latest/libm/) with calls to libdevice intrinsics.
+    ///
+    /// Libm is used by no_std crates for functions such as sin, cos, fabs, etc. However, CUDA provides
+    /// extremely fast GPU-specific implementations of such functions through `libdevice`. Therefore, the codegen
+    /// exposes the option to automatically override any calls to libm functions with calls to libdevice functions.
+    /// However, this means the overriden functions are likely to not be deterministic, so if you rely on strict
+    /// determinism in things like `rapier`, then it may be helpful to disable such a feature.
+    ///
+    /// `true` by default.
+    pub override_libm: bool,
 }
 
 impl CudaBuilder {
@@ -125,6 +135,7 @@ impl CudaBuilder {
             fma_contraction: true,
             emit: None,
             optix: false,
+            override_libm: true,
         }
     }
 
@@ -233,6 +244,18 @@ impl CudaBuilder {
         self
     }
 
+    /// Whether to override calls to [`libm`](https://docs.rs/libm/latest/libm/) with calls to libdevice intrinsics.
+    ///
+    /// Libm is used by no_std crates for functions such as sin, cos, fabs, etc. However, CUDA provides
+    /// extremely fast GPU-specific implementations of such functions through `libdevice`. Therefore, the codegen
+    /// exposes the option to automatically override any calls to libm functions with calls to libdevice functions.
+    /// However, this means the overriden functions are likely to not be deterministic, so if you rely on strict
+    /// determinism in things like `rapier`, then it may be helpful to disable such a feature.
+    pub fn override_libm(mut self, override_libm: bool) -> Self {
+        self.override_libm = override_libm;
+        self
+    }
+
     /// Runs rustc to build the codegen and codegens the gpu crate, returning the path of the final
     /// ptx file. If [`ptx_file_copy_path`](Self::ptx_file_copy_path) is set, this returns the copied path.
     pub fn build(self) -> Result<PathBuf, CudaBuilderError> {
@@ -351,6 +374,10 @@ fn invoke_rustc(builder: &CudaBuilder) -> Result<PathBuf, CudaBuilderError> {
         llvm_args.push("-fma=0".to_string());
     }
 
+    if builder.override_libm {
+        llvm_args.push("--override-libm".to_string());
+    }
+
     let llvm_args = llvm_args.join(" ");
     if !llvm_args.is_empty() {
         rustflags.push(["-Cllvm-args=", &llvm_args].concat());

crates/cust/CHANGELOG.md

@@ -14,3 +14,4 @@ at some point, but now we reconsidered that it may be the wrong choice.
 - Remove `GpuBox::as_device_ptr_mut` and `GpuBuffer::as_device_ptr_mut`.
 - Change `GpuBox::as_device_ptr` and `GpuBuffer::as_device_ptr` to take `&self` instead of `&mut self`.
 - Remove accidentally added `vek` default feature.
+- `vek` feature now uses `default-features = false`, this also means `Rgb` and `Rgba` no longer implement `DeviceCopy`.

crates/cust/Cargo.toml

@@ -11,7 +11,7 @@ description = "High level bindings to the CUDA Driver API"
 cust_raw = { path = "../cust_raw", version = "0.11.2"}
 bitflags = "1.2"
 cust_derive = { path = "../cust_derive", version = "0.1" }
-vek = { version = "0.15.1", optional = true }
+vek = { version = "0.15.1", optional = true, default-features = false }
 
 [build-dependencies]
 find_cuda_helper = { path = "../find_cuda_helper", version = "0.1" }

crates/cust/src/memory/mod.rs

@@ -281,7 +281,7 @@ use vek::*;
 
 #[cfg(feature = "vek")]
 impl_device_copy_vek! {
-    Vec2, Vec3, Vec4, Extent2, Extent3, Rgb, Rgba,
+    Vec2, Vec3, Vec4, Extent2, Extent3,
     Mat2, Mat3, Mat4,
     CubicBezier2, CubicBezier3,
     Quaternion,

crates/rustc_codegen_nvvm/CHANGELOG.md

@@ -11,7 +11,22 @@ method of codegen for the final steps of the codegen. We no longer lazily-load m
 we instead merge all the modules into one then run global DCE on it before giving it to libnvvm.
 
 This means all of the dead code is gone before it gets to the libnvvm stage, drastically lowering the size of 
-the built PTX and improving codegen performance.
+the built PTX and improving codegen performance. `cuda_std` also has a macro `#[externally_visible]` which can
+be used if you want to keep a function around for things like linking multiple PTX files together.
+
+### Libm override
+
+The codegen now has the ability to override [`libm`](https://docs.rs/libm/latest/libm/) functions with 
+[`libdevice`](https://docs.nvidia.com/cuda/libdevice-users-guide/introduction.html#introduction) intrinsics.
+
+Libdevice is a bitcode library shipped with every CUDA SDK installation which provides float routines that
+are optimized for the GPU and for specific GPU architectures. However, these routines are hard to use automatically because
+no_std math crates typically use libm for float things. So users often ended up with needlessly slow or large PTX files
+because they used "emulated" routines.
+
+Now, by default (can be disabled in cuda_builder) the codegen will override libm functions with calls to libdevice automatically.
+However, if you rely on libm for determinism, you must disable the overriding, since libdevice is not strictly deterministic.
+This also makes PTX much smaller generally, in our example path tracer, it slimmed the PTX file from about `3800` LoC to `2300` LoC.
 
 - Trace-level debug is compiled out for release now, decreasing the size of the codegen dll and improving compile times.
 

crates/rustc_codegen_nvvm/src/back.rs

@@ -1,4 +1,5 @@
 use crate::llvm::{self};
+use crate::override_fns::define_or_override_fn;
 use crate::{builder::Builder, context::CodegenCx, lto::ThinBuffer, LlvmMod, NvvmCodegenBackend};
 use libc::{c_char, size_t};
 use rustc_codegen_ssa::back::write::{TargetMachineFactoryConfig, TargetMachineFactoryFn};
@@ -14,6 +15,7 @@ use rustc_data_structures::small_c_str::SmallCStr;
 use rustc_errors::{FatalError, Handler};
 use rustc_fs_util::path_to_c_string;
 use rustc_middle::bug;
+use rustc_middle::mir::mono::MonoItem;
 use rustc_middle::{dep_graph, ty::TyCtxt};
 use rustc_session::config::{self, DebugInfo, OutputType};
 use rustc_session::Session;
@@ -263,7 +265,11 @@ pub fn compile_codegen_unit(tcx: TyCtxt<'_>, cgu_name: Symbol) -> (ModuleCodegen
 
             // ... and now that we have everything pre-defined, fill out those definitions.
             for &(mono_item, _) in &mono_items {
-                mono_item.define::<Builder<'_, '_, '_>>(&cx);
+                if let MonoItem::Fn(func) = mono_item {
+                    define_or_override_fn(func, &cx);
+                } else {
+                    mono_item.define::<Builder<'_, '_, '_>>(&cx);
+                }
             }
 
             // a main function for gpu kernels really makes no sense but

crates/rustc_codegen_nvvm/src/context.rs

@@ -92,13 +92,7 @@ pub(crate) struct CodegenCx<'ll, 'tcx> {
     eh_personality: &'ll Value,
 
     pub symbols: Symbols,
-
-    // we do not currently use codegen_args before linking, and during linking we reparse
-    // them because codegencx is not available at link time. However, we keep this so
-    // it is easier to use them in the future and add args we want to use before linking.
-    #[allow(dead_code)]
     pub codegen_args: CodegenArgs,
-
     // the value of the last call instruction. Needed for return type remapping.
     pub last_call_llfn: Cell<Option<&'ll Value>>,
 }
@@ -505,25 +499,31 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
     }
 }
 
+#[derive(Default, Clone)]
 pub struct CodegenArgs {
     pub nvvm_options: Vec<NvvmOption>,
+    pub override_libm: bool,
 }
 
 impl CodegenArgs {
     pub fn from_session(sess: &Session) -> Self {
-        match Self::parse(&sess.opts.cg.llvm_args) {
-            Ok(x) => x,
-            Err(err) => sess.fatal(&format!("Failed to parse codegen args: {}", err)),
-        }
+        Self::parse(&sess.opts.cg.llvm_args)
     }
 
     // we may want to use rustc's own option parsing facilities to have better errors in the future.
-    pub fn parse(args: &[String]) -> Result<Self, &'static str> {
-        let nvvm_options = args
-            .iter()
-            .map(|x| NvvmOption::from_str(x))
-            .collect::<Result<Vec<_>, _>>()?;
-        Ok(Self { nvvm_options })
+    pub fn parse(args: &[String]) -> Self {
+        // TODO: replace this with a "proper" arg parser.
+        let mut cg_args = Self::default();
+
+        for arg in args {
+            if let Ok(flag) = NvvmOption::from_str(arg) {
+                cg_args.nvvm_options.push(flag);
+            } else if arg == "--override-libm" {
+                cg_args.override_libm = true;
+            }
+        }
+
+        cg_args
     }
 }
 

crates/rustc_codegen_nvvm/src/ctx_intrinsics.rs

@@ -262,8 +262,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
             "__nv_tgamma" |
             "__nv_trunc" |
             "__nv_y0" |
-            "__nv_y1" |
-            "__nv_yn",
+            "__nv_y1",
             fn(t_f64) -> t_f64
         );
 
@@ -316,8 +315,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
             "__nv_tgammaf" |
             "__nv_truncf" |
             "__nv_y0f" |
-            "__nv_y1f" |
-            "__nv_ynf",
+            "__nv_y1f",
             fn(t_f32) -> t_f32
         );
 
@@ -408,5 +406,17 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
             "__nv_fmaf",
             fn(t_f32, t_f32, t_f32) -> t_f32
         );
+
+        ifn!(
+            map,
+            "__nv_yn",
+            fn(t_i32, t_f64) -> t_f64
+        );
+
+        ifn!(
+            map,
+            "__nv_ynf",
+            fn(t_i32, t_f32) -> t_f32
+        );
     }
 }

crates/rustc_codegen_nvvm/src/lib.rs

@@ -44,6 +44,7 @@ mod llvm;
 mod lto;
 mod mono_item;
 mod nvvm;
+mod override_fns;
 mod target;
 mod ty;
 

crates/rustc_codegen_nvvm/src/llvm.rs

@@ -119,6 +119,17 @@ pub(crate) fn get_param(llfn: &Value, index: c_uint) -> &Value {
     }
 }
 
+/// Safe wrapper around `LLVMGetParams`.
+pub(crate) fn get_params(llfn: &Value) -> Vec<&Value> {
+    unsafe {
+        let count = LLVMCountParams(llfn) as usize;
+        let mut params = Vec::with_capacity(count);
+        LLVMGetParams(llfn, params.as_mut_ptr());
+        params.set_len(count);
+        params
+    }
+}
+
 /// Safe wrapper for `LLVMGetValueName2` into a byte slice
 pub(crate) fn get_value_name(value: &Value) -> &[u8] {
     unsafe {