cg_llvm: sve_cast intrinsic

Abstract over the existing `simd_cast` intrinsic to implement a new
`sve_cast` intrinsic - this is better than allowing scalable vectors to
be used with all of the generic `simd_*` intrinsics.

Author: davidtwco

compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs

@@ -3,7 +3,7 @@ use std::fmt::{self, Write};
 use std::hash::{Hash, Hasher};
 use std::path::PathBuf;
 use std::sync::Arc;
-use std::{iter, ptr};
+use std::{assert_matches, iter, ptr};
 
 use libc::{c_longlong, c_uint};
 use rustc_abi::{Align, Layout, NumScalableVectors, Size};

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -597,6 +597,27 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 self.pointercast(val, self.type_ptr())
             }
 
+            sym::sve_cast => {
+                let Some((in_cnt, in_elem, in_num_vecs)) =
+                    args[0].layout.ty.scalable_vector_parts(self.cx.tcx)
+                else {
+                    bug!("input parameter to `sve_cast` was not scalable vector");
+                };
+                let out_layout = self.layout_of(fn_args.type_at(1));
+                let Some((out_cnt, out_elem, out_num_vecs)) =
+                    out_layout.ty.scalable_vector_parts(self.cx.tcx)
+                else {
+                    bug!("output parameter to `sve_cast` was not scalable vector");
+                };
+                assert_eq!(in_cnt, out_cnt);
+                assert_eq!(in_num_vecs, out_num_vecs);
+                let out_llty = self.backend_type(out_layout);
+                match simd_cast(self, sym::simd_cast, args, out_llty, in_elem, out_elem) {
+                    Some(val) => val,
+                    _ => bug!("could not cast scalable vectors"),
+                }
+            }
+
             sym::sve_tuple_create2 => {
                 assert_matches!(
                     self.layout_of(fn_args.type_at(0)).backend_repr,
@@ -2763,96 +2784,17 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
                 out_len
             }
         );
-        // casting cares about nominal type, not just structural type
-        if in_elem == out_elem {
-            return Ok(args[0].immediate());
-        }
-
-        #[derive(Copy, Clone)]
-        enum Sign {
-            Unsigned,
-            Signed,
-        }
-        use Sign::*;
-
-        enum Style {
-            Float,
-            Int(Sign),
-            Unsupported,
-        }
-
-        let (in_style, in_width) = match in_elem.kind() {
-            // vectors of pointer-sized integers should've been
-            // disallowed before here, so this unwrap is safe.
-            ty::Int(i) => (
-                Style::Int(Signed),
-                i.normalize(bx.tcx().sess.target.pointer_width).bit_width().unwrap(),
-            ),
-            ty::Uint(u) => (
-                Style::Int(Unsigned),
-                u.normalize(bx.tcx().sess.target.pointer_width).bit_width().unwrap(),
-            ),
-            ty::Float(f) => (Style::Float, f.bit_width()),
-            _ => (Style::Unsupported, 0),
-        };
-        let (out_style, out_width) = match out_elem.kind() {
-            ty::Int(i) => (
-                Style::Int(Signed),
-                i.normalize(bx.tcx().sess.target.pointer_width).bit_width().unwrap(),
-            ),
-            ty::Uint(u) => (
-                Style::Int(Unsigned),
-                u.normalize(bx.tcx().sess.target.pointer_width).bit_width().unwrap(),
-            ),
-            ty::Float(f) => (Style::Float, f.bit_width()),
-            _ => (Style::Unsupported, 0),
-        };
-
-        match (in_style, out_style) {
-            (Style::Int(sign), Style::Int(_)) => {
-                return Ok(match in_width.cmp(&out_width) {
-                    Ordering::Greater => bx.trunc(args[0].immediate(), llret_ty),
-                    Ordering::Equal => args[0].immediate(),
-                    Ordering::Less => match sign {
-                        Sign::Signed => bx.sext(args[0].immediate(), llret_ty),
-                        Sign::Unsigned => bx.zext(args[0].immediate(), llret_ty),
-                    },
-                });
-            }
-            (Style::Int(Sign::Signed), Style::Float) => {
-                return Ok(bx.sitofp(args[0].immediate(), llret_ty));
-            }
-            (Style::Int(Sign::Unsigned), Style::Float) => {
-                return Ok(bx.uitofp(args[0].immediate(), llret_ty));
-            }
-            (Style::Float, Style::Int(sign)) => {
-                return Ok(match (sign, name == sym::simd_as) {
-                    (Sign::Unsigned, false) => bx.fptoui(args[0].immediate(), llret_ty),
-                    (Sign::Signed, false) => bx.fptosi(args[0].immediate(), llret_ty),
-                    (_, true) => bx.cast_float_to_int(
-                        matches!(sign, Sign::Signed),
-                        args[0].immediate(),
-                        llret_ty,
-                    ),
-                });
-            }
-            (Style::Float, Style::Float) => {
-                return Ok(match in_width.cmp(&out_width) {
-                    Ordering::Greater => bx.fptrunc(args[0].immediate(), llret_ty),
-                    Ordering::Equal => args[0].immediate(),
-                    Ordering::Less => bx.fpext(args[0].immediate(), llret_ty),
-                });
-            }
-            _ => { /* Unsupported. Fallthrough. */ }
+        match simd_cast(bx, name, args, llret_ty, in_elem, out_elem) {
+            Some(val) => return Ok(val),
+            None => return_error!(InvalidMonomorphization::UnsupportedCast {
+                span,
+                name,
+                in_ty,
+                in_elem,
+                ret_ty,
+                out_elem
+            }),
         }
-        return_error!(InvalidMonomorphization::UnsupportedCast {
-            span,
-            name,
-            in_ty,
-            in_elem,
-            ret_ty,
-            out_elem
-        });
     }
     macro_rules! arith_binary {
         ($($name: ident: $($($p: ident),* => $call: ident),*;)*) => {
@@ -3026,3 +2968,86 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
 
     span_bug!(span, "unknown SIMD intrinsic");
 }
+
+/// Implementation of `core::intrinsics::simd_cast`, re-used by `core::scalable::sve_cast`.
+fn simd_cast<'ll, 'tcx>(
+    bx: &mut Builder<'_, 'll, 'tcx>,
+    name: Symbol,
+    args: &[OperandRef<'tcx, &'ll Value>],
+    llret_ty: &'ll Type,
+    in_elem: Ty<'tcx>,
+    out_elem: Ty<'tcx>,
+) -> Option<&'ll Value> {
+    // Casting cares about nominal type, not just structural type
+    if in_elem == out_elem {
+        return Some(args[0].immediate());
+    }
+
+    #[derive(Copy, Clone)]
+    enum Sign {
+        Unsigned,
+        Signed,
+    }
+    use Sign::*;
+
+    enum Style {
+        Float,
+        Int(Sign),
+        Unsupported,
+    }
+
+    let (in_style, in_width) = match in_elem.kind() {
+        // vectors of pointer-sized integers should've been
+        // disallowed before here, so this unwrap is safe.
+        ty::Int(i) => (
+            Style::Int(Signed),
+            i.normalize(bx.tcx().sess.target.pointer_width).bit_width().unwrap(),
+        ),
+        ty::Uint(u) => (
+            Style::Int(Unsigned),
+            u.normalize(bx.tcx().sess.target.pointer_width).bit_width().unwrap(),
+        ),
+        ty::Float(f) => (Style::Float, f.bit_width()),
+        _ => (Style::Unsupported, 0),
+    };
+    let (out_style, out_width) = match out_elem.kind() {
+        ty::Int(i) => (
+            Style::Int(Signed),
+            i.normalize(bx.tcx().sess.target.pointer_width).bit_width().unwrap(),
+        ),
+        ty::Uint(u) => (
+            Style::Int(Unsigned),
+            u.normalize(bx.tcx().sess.target.pointer_width).bit_width().unwrap(),
+        ),
+        ty::Float(f) => (Style::Float, f.bit_width()),
+        _ => (Style::Unsupported, 0),
+    };
+
+    match (in_style, out_style) {
+        (Style::Int(sign), Style::Int(_)) => Some(match in_width.cmp(&out_width) {
+            Ordering::Greater => bx.trunc(args[0].immediate(), llret_ty),
+            Ordering::Equal => args[0].immediate(),
+            Ordering::Less => match sign {
+                Sign::Signed => bx.sext(args[0].immediate(), llret_ty),
+                Sign::Unsigned => bx.zext(args[0].immediate(), llret_ty),
+            },
+        }),
+        (Style::Int(Sign::Signed), Style::Float) => Some(bx.sitofp(args[0].immediate(), llret_ty)),
+        (Style::Int(Sign::Unsigned), Style::Float) => {
+            Some(bx.uitofp(args[0].immediate(), llret_ty))
+        }
+        (Style::Float, Style::Int(sign)) => Some(match (sign, name == sym::simd_as) {
+            (Sign::Unsigned, false) => bx.fptoui(args[0].immediate(), llret_ty),
+            (Sign::Signed, false) => bx.fptosi(args[0].immediate(), llret_ty),
+            (_, true) => {
+                bx.cast_float_to_int(matches!(sign, Sign::Signed), args[0].immediate(), llret_ty)
+            }
+        }),
+        (Style::Float, Style::Float) => Some(match in_width.cmp(&out_width) {
+            Ordering::Greater => bx.fptrunc(args[0].immediate(), llret_ty),
+            Ordering::Equal => args[0].immediate(),
+            Ordering::Less => bx.fpext(args[0].immediate(), llret_ty),
+        }),
+        _ => None,
+    }
+}

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -789,6 +789,7 @@ pub(crate) fn check_intrinsic_type(
         sym::simd_shuffle => (3, 0, vec![param(0), param(0), param(1)], param(2)),
         sym::simd_shuffle_const_generic => (2, 1, vec![param(0), param(0)], param(1)),
 
+        sym::sve_cast => (2, 0, vec![param(0)], param(1)),
         sym::sve_tuple_create2 => (2, 0, vec![param(0), param(0)], param(1)),
         sym::sve_tuple_create3 => (2, 0, vec![param(0), param(0), param(0)], param(1)),
         sym::sve_tuple_create4 => (2, 0, vec![param(0), param(0), param(0), param(0)], param(1)),

compiler/rustc_span/src/symbol.rs

@@ -1974,6 +1974,7 @@ symbols! {
         suggestion,
         super_let,
         supertrait_item_shadowing,
+        sve_cast,
         sve_tuple_create2,
         sve_tuple_create3,
         sve_tuple_create4,

library/core/src/intrinsics/simd/scalable.rs

@@ -2,6 +2,29 @@
 //!
 //! In this module, a "vector" is any `#[rustc_scalable_vector]`-annotated type.
 
+/// Numerically casts a vector, elementwise.
+///
+/// `T` and `U` must be vectors of integers or floats, and must have the same length.
+///
+/// When casting floats to integers, the result is truncated. Out-of-bounds result lead to UB.
+/// When casting integers to floats, the result is rounded.
+/// Otherwise, truncates or extends the value, maintaining the sign for signed integers.
+///
+/// # Safety
+/// Casting from integer types is always safe.
+/// Casting between two float types is also always safe.
+///
+/// Casting floats to integers truncates, following the same rules as `to_int_unchecked`.
+/// Specifically, each element must:
+/// * Not be `NaN`
+/// * Not be infinite
+/// * Be representable in the return type, after truncating off its fractional part
+#[cfg(target_arch = "aarch64")]
+#[rustc_intrinsic]
+#[rustc_nounwind]
+#[target_feature(enable = "sve")]
+pub unsafe fn sve_cast<T, U>(x: T) -> U;
+
 /// Create a tuple of two vectors.
 ///
 /// `SVecTup` must be a scalable vector tuple (`#[rustc_scalable_vector]`) and `SVec` must be a

tests/ui/scalable-vectors/cast-intrinsic.rs

@@ -0,0 +1,65 @@
+//@ check-pass
+//@ only-aarch64
+#![crate_type = "lib"]
+#![allow(incomplete_features, internal_features, improper_ctypes)]
+#![feature(abi_unadjusted, core_intrinsics, link_llvm_intrinsics, rustc_attrs)]
+
+use std::intrinsics::simd::scalable::sve_cast;
+
+#[derive(Copy, Clone)]
+#[rustc_scalable_vector(16)]
+#[allow(non_camel_case_types)]
+pub struct svbool_t(bool);
+
+#[derive(Copy, Clone)]
+#[rustc_scalable_vector(2)]
+#[allow(non_camel_case_types)]
+pub struct svbool2_t(bool);
+
+#[derive(Copy, Clone)]
+#[rustc_scalable_vector(2)]
+#[allow(non_camel_case_types)]
+pub struct svint64_t(i64);
+
+#[derive(Copy, Clone)]
+#[rustc_scalable_vector(2)]
+#[allow(non_camel_case_types)]
+pub struct nxv2i16(i16);
+
+pub trait SveInto<T>: Sized {
+    unsafe fn sve_into(self) -> T;
+}
+
+impl SveInto<svbool2_t> for svbool_t {
+    #[target_feature(enable = "sve")]
+    unsafe fn sve_into(self) -> svbool2_t {
+        unsafe extern "C" {
+            #[cfg_attr(
+                target_arch = "aarch64",
+                link_name = concat!("llvm.aarch64.sve.convert.from.svbool.nxv2i1")
+            )]
+            fn convert_from_svbool(b: svbool_t) -> svbool2_t;
+        }
+        unsafe { convert_from_svbool(self) }
+    }
+}
+
+#[target_feature(enable = "sve")]
+pub unsafe fn svld1sh_gather_s64offset_s64(
+    pg: svbool_t,
+    base: *const i16,
+    offsets: svint64_t,
+) -> svint64_t {
+    unsafe extern "unadjusted" {
+        #[cfg_attr(
+            target_arch = "aarch64",
+            link_name = "llvm.aarch64.sve.ld1.gather.nxv2i16"
+        )]
+        fn _svld1sh_gather_s64offset_s64(
+            pg: svbool2_t,
+            base: *const i16,
+            offsets: svint64_t,
+        ) -> nxv2i16;
+    }
+    sve_cast(_svld1sh_gather_s64offset_s64(pg.sve_into(), base, offsets))
+}