Codegen Overloaded LLVM intrinsics based on their name

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1237,12 +1237,123 @@ fn autocast<'ll>(
     }
 }
 
+fn parse_integer(string: &mut &[u8]) -> Option<u64> {
+    let mut number = 0;
+    let mut position = 0;
+    while let Some(&digit @ b'0'..=b'9') = string.get(position) {
+        number = (10 * number) + (digit - b'0') as u64;
+        position += 1;
+    }
+
+    if position != number.checked_ilog10().unwrap_or(0) as usize + 1 {
+        return None;
+    }
+
+    *string = &string[position..];
+    Some(number)
+}
+
+fn strip_off_prefix(slice: &mut &[u8], prefix: &[u8]) -> bool {
+    slice.strip_prefix(prefix).map(|remainder| *slice = remainder).is_some()
+}
+
+fn demangle_type_str<'ll>(cx: &CodegenCx<'ll, '_>, slice: &mut &[u8]) -> Option<&'ll Type> {
+    Some(if strip_off_prefix(slice, b"isVoid") {
+        cx.type_void()
+    } else if strip_off_prefix(slice, b"f16") {
+        cx.type_f16()
+    } else if strip_off_prefix(slice, b"bf16") {
+        cx.type_bf16()
+    } else if strip_off_prefix(slice, b"f32") {
+        cx.type_f32()
+    } else if strip_off_prefix(slice, b"f64") {
+        cx.type_f64()
+    } else if strip_off_prefix(slice, b"f128") {
+        cx.type_f128()
+    } else if strip_off_prefix(slice, b"i") {
+        let width = parse_integer(slice)?;
+        cx.type_ix(width)
+    } else if strip_off_prefix(slice, b"p") {
+        let address_space = parse_integer(slice)?;
+        cx.type_ptr_ext(AddressSpace(address_space as u32))
+    } else if strip_off_prefix(slice, b"v") {
+        let length = parse_integer(slice)?;
+        let element_type = demangle_type_str(cx, slice)?;
+        cx.type_vector(element_type, length)
+    } else if strip_off_prefix(slice, b"nxv") {
+        let length = parse_integer(slice)?;
+        let element_type = demangle_type_str(cx, slice)?;
+        cx.type_scalable_vector(element_type, length)
+    } else if strip_off_prefix(slice, b"a") {
+        let length = parse_integer(slice)?;
+        let element_type = demangle_type_str(cx, slice)?;
+        cx.type_array(element_type, length)
+    } else if strip_off_prefix(slice, b"sl_") {
+        let mut elements = Vec::new();
+
+        loop {
+            if let Some(remainder) = slice.strip_prefix(b"s")
+                && !remainder.starts_with(b"_")
+                && !remainder.starts_with(b"l_")
+            {
+                *slice = remainder;
+                break cx.type_struct(&elements, true);
+            }
+            elements.push(demangle_type_str(cx, slice)?);
+        }
+    } else if strip_off_prefix(slice, b"f_") {
+        let return_type = demangle_type_str(cx, slice)?;
+        let mut arguments = Vec::new();
+
+        loop {
+            if let Some(remainder) = slice.strip_prefix(b"f")
+                && !remainder.starts_with(b"_")
+            {
+                *slice = remainder;
+                break cx.type_func(&arguments, return_type);
+            }
+            if strip_off_prefix(slice, b"varargf") {
+                break cx.type_variadic_func(&arguments, return_type);
+            }
+            arguments.push(demangle_type_str(cx, slice)?);
+        }
+    } else {
+        return None;
+    })
+}
+
+fn parse_type_parameters<'ll, 'tcx>(
+    cx: &CodegenCx<'ll, 'tcx>,
+    intrinsic: llvm::Intrinsic,
+    name: &str,
+) -> Option<Vec<&'ll Type>> {
+    let base_name: &'ll [u8] = intrinsic.base_name();
+
+    let slice = &mut name.as_bytes().strip_prefix(base_name).unwrap();
+
+    if !intrinsic.is_overloaded() {
+        return slice.is_empty().then(|| Vec::new());
+    }
+
+    let mut type_params = Vec::new();
+
+    while !slice.is_empty() {
+        if !strip_off_prefix(slice, b".") {
+            return None;
+        }
+
+        type_params.push(demangle_type_str(cx, slice)?);
+    }
+
+    Some(type_params)
+}
+
 fn intrinsic_fn<'ll, 'tcx>(
     bx: &Builder<'_, 'll, 'tcx>,
     name: &str,
     rust_return_ty: &'ll Type,
     rust_argument_tys: Vec<&'ll Type>,
-    instance: ty::Instance<'tcx>,
+    instance: Instance<'tcx>,
 ) -> &'ll Value {
     let tcx = bx.tcx;
 
@@ -1268,10 +1379,9 @@ fn intrinsic_fn<'ll, 'tcx>(
     }
 
     if let Some(intrinsic) = intrinsic
-        && !intrinsic.is_overloaded()
+        && let Some(type_params) = parse_type_parameters(bx.cx, intrinsic, name)
     {
-        // FIXME: also do this for overloaded intrinsics
-        let llfn = intrinsic.get_declaration(bx.llmod, &[]);
+        let llfn = intrinsic.get_declaration(bx.llmod, &type_params);
         let llvm_fn_ty = bx.get_type_of_global(llfn);
 
         let llvm_return_ty = bx.get_return_type(llvm_fn_ty);

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1124,6 +1124,10 @@ unsafe extern "C" {
         NewFn: &mut Option<&'a Value>,
     ) -> bool;
     pub(crate) fn LLVMRustIsTargetIntrinsic(ID: NonZero<c_uint>) -> bool;
+    pub(crate) fn LLVMRustIntrinsicGetBaseName(
+        ID: NonZero<c_uint>,
+        NameLength: &mut size_t,
+    ) -> *const c_char;
 
     // Operations on parameters
     pub(crate) fn LLVMIsAArgument(Val: &Value) -> Option<&Value>;

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -2,8 +2,8 @@
 
 use std::ffi::{CStr, CString};
 use std::num::NonZero;
-use std::ptr;
 use std::string::FromUtf8Error;
+use std::{ptr, slice};
 
 use libc::c_uint;
 use rustc_abi::{AddressSpace, Align, Size, WrappingRange};
@@ -340,6 +340,12 @@ impl Intrinsic {
             LLVMGetIntrinsicDeclaration(llmod, self.id, type_params.as_ptr(), type_params.len())
         }
     }
+
+    pub(crate) fn base_name<'ll>(self) -> &'ll [u8] {
+        let mut length = 0;
+        let ptr = unsafe { LLVMRustIntrinsicGetBaseName(self.id, &mut length) };
+        unsafe { slice::from_raw_parts(ptr.cast(), length) }
+    }
 }
 
 /// Safe wrapper for `LLVMSetValueName2` from a byte slice

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -130,7 +130,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     let start = dest.val.llval;
                     let size = bx.const_usize(dest.layout.size.bytes());
 
-                    // Use llvm.memset.p0i8.* to initialize all same byte arrays
+                    // Use llvm.memset.p0.* to initialize all same byte arrays
                     if let Some(int) = bx.cx().const_to_opt_u128(v, false)
                         && let bytes = &int.to_le_bytes()[..cg_elem.layout.size.bytes_usize()]
                         && let Ok(&byte) = bytes.iter().all_equal_value()
@@ -140,7 +140,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         return true;
                     }
 
-                    // Use llvm.memset.p0i8.* to initialize byte arrays
+                    // Use llvm.memset.p0.* to initialize byte arrays
                     let v = bx.from_immediate(v);
                     if bx.cx().val_ty(v) == bx.cx().type_i8() {
                         bx.memset(start, v, size, dest.val.align, MemFlags::empty());

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -1841,6 +1841,13 @@ extern "C" bool LLVMRustIsTargetIntrinsic(unsigned ID) {
   return Intrinsic::isTargetIntrinsic(ID);
 }
 
+extern "C" const char *LLVMRustIntrinsicGetBaseName(unsigned ID,
+                                                    size_t *NameLength) {
+  auto baseName = Intrinsic::getBaseName(ID);
+  *NameLength = baseName.size();
+  return baseName.data();
+}
+
 // Statically assert that the fixed metadata kind IDs declared in
 // `metadata_kind.rs` match the ones actually used by LLVM.
 #define FIXED_MD_KIND(VARIANT, VALUE)                                          \

library/core/src/intrinsics/mod.rs

@@ -953,7 +953,7 @@ pub const unsafe fn slice_get_unchecked<
 #[rustc_intrinsic]
 pub fn ptr_mask<T>(ptr: *const T, mask: usize) -> *const T;
 
-/// Equivalent to the appropriate `llvm.memcpy.p0i8.0i8.*` intrinsic, with
+/// Equivalent to the appropriate `llvm.memcpy.p0.p0.*` intrinsic, with
 /// a size of `count` * `size_of::<T>()` and an alignment of `align_of::<T>()`.
 ///
 /// This intrinsic does not have a stable counterpart.
@@ -967,7 +967,7 @@ pub fn ptr_mask<T>(ptr: *const T, mask: usize) -> *const T;
 #[rustc_intrinsic]
 #[rustc_nounwind]
 pub unsafe fn volatile_copy_nonoverlapping_memory<T>(dst: *mut T, src: *const T, count: usize);
-/// Equivalent to the appropriate `llvm.memmove.p0i8.0i8.*` intrinsic, with
+/// Equivalent to the appropriate `llvm.memmove.p0.p0.*` intrinsic, with
 /// a size of `count * size_of::<T>()` and an alignment of `align_of::<T>()`.
 ///
 /// The volatile parameter is set to `true`, so it will not be optimized out
@@ -977,7 +977,7 @@ pub unsafe fn volatile_copy_nonoverlapping_memory<T>(dst: *mut T, src: *const T,
 #[rustc_intrinsic]
 #[rustc_nounwind]
 pub unsafe fn volatile_copy_memory<T>(dst: *mut T, src: *const T, count: usize);
-/// Equivalent to the appropriate `llvm.memset.p0i8.*` intrinsic, with a
+/// Equivalent to the appropriate `llvm.memset.p0.*` intrinsic, with a
 /// size of `count * size_of::<T>()` and an alignment of `align_of::<T>()`.
 ///
 /// This intrinsic does not have a stable counterpart.

src/tools/rustfmt/tests/target/format_strings/issue-202.rs

@@ -6,7 +6,7 @@ fn compile_empty_program() {
     let expected = "; ModuleID = \'foo\'
 
 ; Function Attrs: nounwind
-declare void @llvm.memset.p0i8.i32(i8* nocapture, i8, i32, i32, i1) #0
+declare void @llvm.memset.p0.i32(i8* nocapture, i8, i32, i32, i1) #0
 
 declare i32 @write(i32, i8*, i32)
 

src/tools/rustfmt/tests/target/format_strings/issue-3263.rs

@@ -7,7 +7,7 @@ fn compile_empty_program() {
     let expected = "; ModuleID = \'foo\'
 
 ; Function Attrs: nounwind
-declare void @llvm.memset.p0i8.i32(i8* nocapture, i8, i32, i32, i1) #0
+declare void @llvm.memset.p0.i32(i8* nocapture, i8, i32, i32, i1) #0
 
 declare i32 @write(i32, i8*, i32)
 

src/tools/rustfmt/tests/target/string-lit-custom.rs

@@ -2,7 +2,7 @@ fn main() {
     let expected = "; ModuleID = \'foo\'
 
 ; Function Attrs: nounwind
-declare void @llvm.memset.p0i8.i32(i8* nocapture, i8, i32, i32, i1) #0
+declare void @llvm.memset.p0.i32(i8* nocapture, i8, i32, i32, i1) #0
 
 declare i32 @write(i32, i8*, i32)
 

tests/codegen-llvm/inject-autocast.rs

@@ -103,6 +103,20 @@ pub unsafe fn amx_autocast(m: u16, n: u16, k: u16, a: Tile, b: Tile, c: Tile) ->
     foo(m, n, k, a, b, c)
 }
 
+// CHECK-LABEL: @overloaded_bf16_autocast
+#[no_mangle]
+pub unsafe fn overloaded_bf16_autocast(a: i16x8) -> i16x8 {
+    extern "unadjusted" {
+        #[link_name = "llvm.sqrt.v8bf16"]
+        fn foo(a: i16x8) -> i16x8;
+    }
+
+    // CHECK: [[A:%[0-9]+]] = bitcast <8 x i16> {{.*}} to <8 x bfloat>
+    // CHECK: [[B:%[0-9]+]] = call <8 x bfloat> @llvm.sqrt.v8bf16(<8 x bfloat> [[A]])
+    // CHECK: bitcast <8 x bfloat> [[B]] to <8 x i16>
+    foo(a)
+}
+
 // CHECK: declare { i32, <2 x i64>, <2 x i64>, <2 x i64>, <2 x i64>, <2 x i64>, <2 x i64> } @llvm.x86.encodekey128(i32, <2 x i64>)
 
 // CHECK: declare { <2 x i1>, <2 x i1> } @llvm.x86.avx512.vp2intersect.q.128(<2 x i64>, <2 x i64>)
@@ -116,3 +130,5 @@ pub unsafe fn amx_autocast(m: u16, n: u16, k: u16, a: Tile, b: Tile, c: Tile) ->
 // CHECK: declare x86_amx @llvm.x86.cast.vector.to.tile.v1024i8(<1024 x i8>)
 
 // CHECK: declare <1024 x i8> @llvm.x86.cast.tile.to.vector.v1024i8(x86_amx)
+
+// CHECK: declare <8 x bfloat> @llvm.sqrt.v8bf16(<8 x bfloat>)

tests/codegen-llvm/intrinsics/mask.rs

@@ -7,6 +7,6 @@
 #[no_mangle]
 pub fn mask_ptr(ptr: *const u16, mask: usize) -> *const u16 {
     // CHECK: call
-    // CHECK-SAME: @llvm.ptrmask.{{p0|p0i8}}.[[WORD]](ptr {{%ptr|%1}}, [[WORD]] %mask)
+    // CHECK-SAME: @llvm.ptrmask.p0.[[WORD]](ptr {{%ptr|%1}}, [[WORD]] %mask)
     core::intrinsics::ptr_mask(ptr, mask)
 }

tests/ui/abi/arm-unadjusted-intrinsic.rs

@@ -34,8 +34,8 @@ impl Copy for int8x16x4_t {}
 pub unsafe fn vld1q_s8_x4(a: *const i8) -> int8x16x4_t {
     #[allow(improper_ctypes)]
     extern "unadjusted" {
-        #[cfg_attr(target_arch = "arm", link_name = "llvm.arm.neon.vld1x4.v16i8.p0i8")]
-        #[cfg_attr(target_arch = "aarch64", link_name = "llvm.aarch64.neon.ld1x4.v16i8.p0i8")]
+        #[cfg_attr(target_arch = "arm", link_name = "llvm.arm.neon.vld1x4.v16i8.p0")]
+        #[cfg_attr(target_arch = "aarch64", link_name = "llvm.aarch64.neon.ld1x4.v16i8.p0")]
         fn vld1q_s8_x4_(a: *const i8) -> int8x16x4_t;
     }
     vld1q_s8_x4_(a)