Introduce arch::aarch64 and use it for rint{,f}

5 files changed, 50 insertions, 2 deletions

diff --git a/library/compiler-builtins/libm/etc/function-definitions.json b/library/compiler-builtins/libm/etc/function-definitions.json
index f60a7e5673c..39b6c97029c 100644
--- a/library/compiler-builtins/libm/etc/function-definitions.json
+++ b/library/compiler-builtins/libm/etc/function-definitions.json
@@ -604,6 +604,7 @@
     "rint": {
         "sources": [
             "src/libm_helper.rs",
+            "src/math/arch/aarch64.rs",
             "src/math/arch/wasm32.rs",
             "src/math/rint.rs"
         ],
@@ -611,6 +612,7 @@
     },
     "rintf": {
         "sources": [
+            "src/math/arch/aarch64.rs",
             "src/math/arch/wasm32.rs",
             "src/math/rintf.rs"
         ],
diff --git a/library/compiler-builtins/libm/src/math/arch/aarch64.rs b/library/compiler-builtins/libm/src/math/arch/aarch64.rs
new file mode 100644
index 00000000000..374ec11bfec
--- /dev/null
+++ b/library/compiler-builtins/libm/src/math/arch/aarch64.rs
@@ -0,0 +1,33 @@
+use core::arch::aarch64::{
+    float32x2_t, float64x1_t, vdup_n_f32, vdup_n_f64, vget_lane_f32, vget_lane_f64, vrndn_f32,
+    vrndn_f64,
+};
+
+pub fn rint(x: f64) -> f64 {
+    // SAFETY: only requires target_feature=neon, ensured by `cfg_if` in parent module.
+    let x_vec: float64x1_t = unsafe { vdup_n_f64(x) };
+
+    // SAFETY: only requires target_feature=neon, ensured by `cfg_if` in parent module.
+    let result_vec: float64x1_t = unsafe { vrndn_f64(x_vec) };
+
+    // SAFETY: only requires target_feature=neon, ensured by `cfg_if` in parent module.
+    let result: f64 = unsafe { vget_lane_f64::<0>(result_vec) };
+
+    result
+}
+
+pub fn rintf(x: f32) -> f32 {
+    // There's a scalar form of this instruction (FRINTN) but core::arch doesn't expose it, so we
+    // have to use the vector form and drop the other lanes afterwards.
+
+    // SAFETY: only requires target_feature=neon, ensured by `cfg_if` in parent module.
+    let x_vec: float32x2_t = unsafe { vdup_n_f32(x) };
+
+    // SAFETY: only requires target_feature=neon, ensured by `cfg_if` in parent module.
+    let result_vec: float32x2_t = unsafe { vrndn_f32(x_vec) };
+
+    // SAFETY: only requires target_feature=neon, ensured by `cfg_if` in parent module.
+    let result: f32 = unsafe { vget_lane_f32::<0>(result_vec) };
+
+    result
+}
diff --git a/library/compiler-builtins/libm/src/math/arch/mod.rs b/library/compiler-builtins/libm/src/math/arch/mod.rs
index 3992419cbaa..091d7650a5a 100644
--- a/library/compiler-builtins/libm/src/math/arch/mod.rs
+++ b/library/compiler-builtins/libm/src/math/arch/mod.rs
@@ -17,6 +17,13 @@ cfg_if! {
     } else if #[cfg(target_feature = "sse2")] {
         mod i686;
         pub use i686::{sqrt, sqrtf};
+    } else if #[cfg(all(
+        target_arch = "aarch64", // TODO: also arm64ec?
+        target_feature = "neon",
+        target_endian = "little", // see https://github.com/rust-lang/stdarch/issues/1484
+    ))] {
+        mod aarch64;
+        pub use aarch64::{rint, rintf};
     }
 }
 
diff --git a/library/compiler-builtins/libm/src/math/rint.rs b/library/compiler-builtins/libm/src/math/rint.rs
index 50192ffdf03..c9ea6402ec7 100644
--- a/library/compiler-builtins/libm/src/math/rint.rs
+++ b/library/compiler-builtins/libm/src/math/rint.rs
@@ -2,7 +2,10 @@
 pub fn rint(x: f64) -> f64 {
     select_implementation! {
         name: rint,
-        use_arch: all(target_arch = "wasm32", intrinsics_enabled),
+        use_arch: any(
+            all(target_arch = "wasm32", intrinsics_enabled),
+            all(target_arch = "aarch64", target_feature = "neon", target_endian = "little"),
+        ),
         args: x,
     }
 
diff --git a/library/compiler-builtins/libm/src/math/rintf.rs b/library/compiler-builtins/libm/src/math/rintf.rs
index 64968b6be3a..33b5b3ddebf 100644
--- a/library/compiler-builtins/libm/src/math/rintf.rs
+++ b/library/compiler-builtins/libm/src/math/rintf.rs
@@ -2,7 +2,10 @@
 pub fn rintf(x: f32) -> f32 {
     select_implementation! {
         name: rintf,
-        use_arch: all(target_arch = "wasm32", intrinsics_enabled),
+        use_arch: any(
+            all(target_arch = "wasm32", intrinsics_enabled),
+            all(target_arch = "aarch64", target_feature = "neon", target_endian = "little"),
+        ),
         args: x,
     }