From bff72e693f9cc85d45444811a26bd87460203b08 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 14 Feb 2025 16:09:10 +0000
Subject: [PATCH] Implement vpmaxq_u8 on aarch64

---
 src/shims/aarch64.rs       | 69 ++++++++++++++++++++++++++++++++++++++
 src/shims/foreign_items.rs | 18 +++-------
 src/shims/mod.rs           |  1 +
 3 files changed, 75 insertions(+), 13 deletions(-)
 create mode 100644 src/shims/aarch64.rs

diff --git a/src/shims/aarch64.rs b/src/shims/aarch64.rs
new file mode 100644
index 0000000000..1394cc9d84
--- /dev/null
+++ b/src/shims/aarch64.rs
@@ -0,0 +1,69 @@
+use std::cmp;
+
+use rustc_middle::ty::Ty;
+use rustc_span::Symbol;
+use rustc_target::callconv::{Conv, FnAbi};
+
+use self::helpers::bool_to_simd_element;
+use crate::*;
+
+impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
+pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
+    fn emulate_aarch64_intrinsic(
+        &mut self,
+        link_name: Symbol,
+        abi: &FnAbi<'tcx, Ty<'tcx>>,
+        args: &[OpTy<'tcx>],
+        dest: &MPlaceTy<'tcx>,
+    ) -> InterpResult<'tcx, EmulateItemResult> {
+        let this = self.eval_context_mut();
+        // Prefix should have already been checked.
+        let unprefixed_name = link_name.as_str().strip_prefix("llvm.aarch64.").unwrap();
+        match unprefixed_name {
+            "isb" => {
+                let [arg] = this.check_shim(abi, Conv::C, link_name, args)?;
+                let arg = this.read_scalar(arg)?.to_i32()?;
+                match arg {
+                    // SY ("full system scope")
+                    15 => {
+                        this.yield_active_thread();
+                    }
+                    _ => {
+                        throw_unsup_format!("unsupported llvm.aarch64.isb argument {}", arg);
+                    }
+                }
+            }
+
+            "neon.umaxp.v16i8" => {
+                let [left, right] = this.check_shim(abi, Conv::C, link_name, args)?;
+
+                let (left, left_len) = this.project_to_simd(left)?;
+                let (right, right_len) = this.project_to_simd(right)?;
+                let (dest, lane_count) = this.project_to_simd(dest)?;
+                assert_eq!(left_len, right_len);
+                assert_eq!(lane_count, left_len);
+
+                for lane_idx in 0..lane_count {
+                    let src = if lane_idx < (lane_count / 2) { &left } else { &right };
+                    let src_idx = lane_idx % (lane_count / 2);
+
+                    let lhs_lane = this.read_scalar(&this.project_index(src, src_idx * 2)?)?;
+                    let rhs_lane = this.read_scalar(&this.project_index(src, src_idx * 2 + 1)?)?;
+
+                    let res_lane = match unprefixed_name {
+                        "neon.umaxp.v16i8" => {
+                            Scalar::from_i8(cmp::max(lhs_lane.to_i8()?, rhs_lane.to_i8()?))
+                        }
+                        _ => unreachable!(),
+                    };
+
+                    let dest = this.project_index(&dest, lane_idx)?;
+                    this.write_scalar(res_lane, &dest)?;
+                }
+            }
+
+            _ => return interp_ok(EmulateItemResult::NotSupported),
+        }
+        interp_ok(EmulateItemResult::NeedsReturn)
+    }
+}
diff --git a/src/shims/foreign_items.rs b/src/shims/foreign_items.rs
index 97bfb04f1f..011eeac87b 100644
--- a/src/shims/foreign_items.rs
+++ b/src/shims/foreign_items.rs
@@ -939,20 +939,12 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                     this, link_name, abi, args, dest,
                 );
             }
-            // FIXME: Move these to an `arm` submodule.
-            "llvm.aarch64.isb" if this.tcx.sess.target.arch == "aarch64" => {
-                let [arg] = this.check_shim(abi, Conv::C, link_name, args)?;
-                let arg = this.read_scalar(arg)?.to_i32()?;
-                match arg {
-                    // SY ("full system scope")
-                    15 => {
-                        this.yield_active_thread();
-                    }
-                    _ => {
-                        throw_unsup_format!("unsupported llvm.aarch64.isb argument {}", arg);
-                    }
-                }
+            name if name.starts_with("llvm.aarch64.") && this.tcx.sess.target.arch == "aarch64" => {
+                return shims::aarch64::EvalContextExt::emulate_aarch64_intrinsic(
+                    this, link_name, abi, args, dest,
+                );
             }
+            // FIXME: Move this to an `arm` submodule.
             "llvm.arm.hint" if this.tcx.sess.target.arch == "arm" => {
                 let [arg] = this.check_shim(abi, Conv::C, link_name, args)?;
                 let arg = this.read_scalar(arg)?.to_i32()?;
diff --git a/src/shims/mod.rs b/src/shims/mod.rs
index 61681edcf7..b498551ace 100644
--- a/src/shims/mod.rs
+++ b/src/shims/mod.rs
@@ -1,5 +1,6 @@
 #![warn(clippy::arithmetic_side_effects)]
 
+mod aarch64;
 mod alloc;
 mod backtrace;
 mod files;