Optimize atomic float on NVPTX

taiki-e · Jun 4, 2023 · 02adc3a · 02adc3a
1 parent 8fcbe7c
commit 02adc3a
Show file tree

Hide file tree

Showing 5 changed files with 455 additions and 2 deletions.
diff --git a/build.rs b/build.rs
@@ -155,7 +155,8 @@ fn main() {
         // https://github.com/rust-lang/rust/pull/111331 merged in Rust 1.71 (nightly-2023-05-09).
         if !no_asm
             && (target_arch == "powerpc64" && version.probe(60, 2022, 2, 13)
-                || target_arch == "s390x" && version.probe(71, 2023, 5, 8))
+                || target_arch == "s390x" && version.probe(71, 2023, 5, 8)
+                || target_arch == "nvptx64")
             && is_allowed_feature("asm_experimental_arch")
         {
             println!("cargo:rustc-cfg=portable_atomic_unstable_asm_experimental_arch");
@@ -318,6 +319,36 @@ fn main() {
                 false,
             );
         }
+        "nvptx64" => {
+            let mut has_sm_70 = false;
+            if let Some(rustflags) = env::var_os("CARGO_ENCODED_RUSTFLAGS") {
+                for mut flag in rustflags.to_string_lossy().split('\x1f') {
+                    flag = strip_prefix(flag, "-C").unwrap_or(flag);
+                    if let Some(flag) = strip_prefix(flag, "target-feature=") {
+                        for s in flag.split(',') {
+                            // TODO: Handles cases where a specific target feature
+                            // implicitly enables another target feature.
+                            match (s.as_bytes().first(), s.get(1..)) {
+                                (Some(b'+'), Some(f)) => {
+                                    if let Some(sm) = strip_prefix(f, "sm_") {
+                                        if let Ok(sm) = sm.parse::<u32>() {
+                                            if sm >= 70 {
+                                                has_sm_70 = true;
+                                            }
+                                        }
+                                    }
+                                }
+                                (Some(b'-'), Some(_f)) => {
+                                    // TODO
+                                }
+                                _ => {}
+                            }
+                        }
+                    }
+                }
+            }
+            target_feature_if("sm_70", has_sm_70, &version, None, false);
+        }
         _ => {}
     }
 }

diff --git a/src/imp/mod.rs b/src/imp/mod.rs
@@ -117,6 +117,15 @@ mod arm_linux;
 #[cfg(target_arch = "msp430")]
 pub(crate) mod msp430;
 
+#[cfg(portable_atomic_unstable_asm_experimental_arch)]
+#[cfg(feature = "float")]
+#[cfg(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    portable_atomic_unstable_asm_experimental_arch,
+))]
+pub(crate) mod nvptx;
+
 #[cfg(any(test, not(feature = "critical-section")))]
 #[cfg_attr(portable_atomic_no_cfg_target_has_atomic, cfg(any(test, portable_atomic_no_atomic_cas)))]
 #[cfg_attr(
@@ -199,8 +208,21 @@ mod interrupt;
 // Atomic float implementations
 
 #[cfg(feature = "float")]
+#[cfg(not(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    portable_atomic_unstable_asm_experimental_arch,
+)))]
 pub(crate) mod float;
 
+#[cfg(feature = "float")]
+#[cfg(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    portable_atomic_unstable_asm_experimental_arch,
+))]
+pub(crate) use nvptx as float;
+
 // -----------------------------------------------------------------------------
 
 // Atomic{Isize,Usize,Bool,Ptr}, Atomic{I,U}{8,16}