Optimize atomic float on NVPTX

build.rs

@@ -48,7 +48,7 @@ fn main() {
 
     if version.minor >= 80 {
         println!(
-            r#"cargo:rustc-check-cfg=cfg(target_feature,values("lsfe","fast-serialization","load-store-on-cond","distinct-ops","miscellaneous-extensions-3"))"#
+            r#"cargo:rustc-check-cfg=cfg(target_feature,values("lsfe","sm_70","fast-serialization","load-store-on-cond","distinct-ops","miscellaneous-extensions-3"))"#
         );
 
         // Custom cfgs set by build script. Not public API.
@@ -59,7 +59,7 @@ fn main() {
         // TODO: handle multi-line target_feature_fallback
         // grep -F 'target_feature_fallback("' build.rs | grep -Ev '^ *//' | sed -E 's/^.*target_feature_fallback\(//; s/",.*$/"/' | LC_ALL=C sort -u | tr '\n' ',' | sed -E 's/,$/\n/'
         println!(
-            r#"cargo:rustc-check-cfg=cfg(portable_atomic_target_feature,values("cmpxchg16b","distinct-ops","fast-serialization","load-store-on-cond","lse","lse128","lse2","lsfe","mclass","miscellaneous-extensions-3","quadword-atomics","rcpc3","v6","zaamo","zabha","zacas"))"#
+            r#"cargo:rustc-check-cfg=cfg(portable_atomic_target_feature,values("cmpxchg16b","distinct-ops","fast-serialization","load-store-on-cond","lse","lse128","lse2","lsfe","mclass","miscellaneous-extensions-3","quadword-atomics","rcpc3","sm_70","v6","zaamo","zabha","zacas"))"#
         );
     }
 
@@ -176,6 +176,11 @@ fn main() {
                     println!("cargo:rustc-cfg=portable_atomic_unstable_asm_experimental_arch");
                 }
             }
+            "nvptx64" => {
+                if version.nightly && is_allowed_feature("asm_experimental_arch") {
+                    println!("cargo:rustc-cfg=portable_atomic_unstable_asm_experimental_arch");
+                }
+            }
             _ => {}
         }
     }
@@ -443,6 +448,36 @@ fn main() {
             // nand (nnr{,g}k), select (sel{,g}r), etc.
             target_feature_fallback("miscellaneous-extensions-3", arch13_features);
         }
+        "nvptx64" => {
+            let mut 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 {
+                                                sm_70 = true;
+                                            }
+                                        }
+                                    }
+                                }
+                                (Some(b'-'), Some(_f)) => {
+                                    // TODO
+                                }
+                                _ => {}
+                            }
+                        }
+                    }
+                }
+            }
+            target_feature_fallback("sm_70", sm_70);
+        }
         _ => {}
     }
 }

src/imp/float/int.rs

@@ -10,8 +10,7 @@ CAS loops, which can be slower than equivalent operations of atomic integers.
 
 AArch64 with FEAT_LSFE and GPU targets have atomic instructions for float.
 See aarch64.rs for AArch64 with FEAT_LSFE.
-GPU targets will also use architecture-specific implementations instead of this implementation in the
-future: https://github.com/taiki-e/portable-atomic/issues/34 / https://github.com/taiki-e/portable-atomic/pull/45
+See nvptx.rs for NVPTX.
 */
 
 // TODO: fetch_{minimum,maximum}* https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2024/p3008r2.html / https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2024/p0493r5.pdf
@@ -239,9 +238,21 @@ macro_rules! atomic_float {
 cfg_has_atomic_16! {
     atomic_float!(AtomicF16, f16, AtomicU16, u16, 2);
 }
+#[cfg(not(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    not(any(miri, portable_atomic_sanitize_thread)),
+    portable_atomic_unstable_asm_experimental_arch,
+)))]
 cfg_has_atomic_32! {
     atomic_float!(AtomicF32, f32, AtomicU32, u32, 4);
 }
+#[cfg(not(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    not(any(miri, portable_atomic_sanitize_thread)),
+    portable_atomic_unstable_asm_experimental_arch,
+)))]
 cfg_has_atomic_64! {
     atomic_float!(AtomicF64, f64, AtomicU64, u64, 8);
 }

src/imp/float/mod.rs

@@ -6,6 +6,13 @@ Atomic float implementations
 
 #![allow(clippy::float_arithmetic)]
 
+#[cfg(not(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    not(any(miri, portable_atomic_sanitize_thread)),
+    portable_atomic_unstable_asm_experimental_arch,
+    not(any(portable_atomic_unstable_f16, portable_atomic_unstable_f128)),
+)))]
 mod int;
 
 #[cfg(all(
@@ -17,13 +24,40 @@ mod int;
 ))]
 mod aarch64;
 
+#[cfg(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    not(any(miri, portable_atomic_sanitize_thread)),
+    portable_atomic_unstable_asm_experimental_arch,
+))]
+mod nvptx;
+
 #[cfg(portable_atomic_unstable_f16)]
 cfg_has_atomic_16! {
     pub(crate) use self::int::AtomicF16;
 }
+#[cfg(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    not(any(miri, portable_atomic_sanitize_thread)),
+    portable_atomic_unstable_asm_experimental_arch,
+))]
+pub(crate) use self::nvptx::{AtomicF32, AtomicF64};
+#[cfg(not(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    not(any(miri, portable_atomic_sanitize_thread)),
+    portable_atomic_unstable_asm_experimental_arch,
+)))]
 cfg_has_atomic_32! {
     pub(crate) use self::int::AtomicF32;
 }
+#[cfg(not(all(
+    target_arch = "nvptx64",
+    any(target_feature = "sm_70", portable_atomic_target_feature = "sm_70"),
+    not(any(miri, portable_atomic_sanitize_thread)),
+    portable_atomic_unstable_asm_experimental_arch,
+)))]
 cfg_has_atomic_64! {
     pub(crate) use self::int::AtomicF64;
 }