Add simd_masked_store platform-intrinsic

This maps to llvm.masked.store

Author: farnoy

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1498,7 +1498,7 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
         // * T: type of the element to load
         // * M: any integer width is supported, will be truncated to i1
 
-        // The first argument is
+        // The first argument is a passthrough vector providing values for disabled lanes
         let (_, element_ty0) = require_simd!(in_ty, SimdFirst);
 
         // The element type of the second argument must be a signed integer type of any width:
@@ -1586,6 +1586,107 @@ fn generic_simd_intrinsic<'ll, 'tcx>(
         return Ok(v);
     }
 
+    if name == sym::simd_masked_store {
+        // simd_masked_store(values: <N x T>, pointer: *mut T,
+        //                   mask: <N x i{M}>) -> ()
+        // * N: number of elements in the input vectors
+        // * T: type of the element to load
+        // * M: any integer width is supported, will be truncated to i1
+
+        // The first argument is a passthrough vector providing values for disabled lanes
+        let (element_len1, element_ty0) = require_simd!(in_ty, SimdFirst);
+
+        // The second argument must be a simd vector with an element type that's a pointer
+        // to the element type of the first argument
+        let (mask_len, element_ty2) = require_simd!(arg_tys[2], SimdSecond);
+
+        // Of the same length:
+        require!(
+            in_len == element_len1,
+            InvalidMonomorphization::SecondArgumentLength {
+                span,
+                name,
+                in_len,
+                in_ty,
+                arg_ty: arg_tys[1],
+                out_len: element_len1
+            }
+        );
+        require!(
+            in_len == mask_len,
+            InvalidMonomorphization::ThirdArgumentLength {
+                span,
+                name,
+                in_len,
+                in_ty,
+                arg_ty: arg_tys[2],
+                out_len: mask_len
+            }
+        );
+
+        // Pointer type must match the element type
+        require!(
+            matches!(
+                arg_tys[1].kind(),
+                ty::RawPtr(p) if p.ty == in_elem && p.ty.kind() == element_ty0.kind() && p.mutbl.is_mut()
+            ),
+            InvalidMonomorphization::ExpectedElementType {
+                span,
+                name,
+                expected_element: in_elem,
+                second_arg: arg_tys[1],
+                in_elem,
+                in_ty,
+                mutability: ExpectedPointerMutability::Mut,
+            }
+        );
+
+        // The element type of the third argument must be a signed integer type of any width:
+        match element_ty2.kind() {
+            ty::Int(_) => (),
+            _ => {
+                return_error!(InvalidMonomorphization::ThirdArgElementType {
+                    span,
+                    name,
+                    expected_element: element_ty2,
+                    third_arg: arg_tys[2]
+                });
+            }
+        }
+
+        // Alignment of T, must be a constant integer value:
+        let alignment_ty = bx.type_i32();
+        let alignment = bx.const_i32(bx.align_of(in_elem).bytes() as i32);
+
+        // Truncate the mask vector to a vector of i1s:
+        let (mask, mask_ty) = {
+            let i1 = bx.type_i1();
+            let i1xn = bx.type_vector(i1, in_len);
+            (bx.trunc(args[2].immediate(), i1xn), i1xn)
+        };
+
+        let ret_t = bx.type_void();
+
+        let llvm_pointer = bx.type_ptr();
+
+        // Type of the vector of elements:
+        let llvm_elem_vec_ty = llvm_vector_ty(bx, element_ty0, in_len);
+        let llvm_elem_vec_str = llvm_vector_str(bx, in_elem, in_len);
+
+        let llvm_intrinsic = format!("llvm.masked.store.{llvm_elem_vec_str}.p0");
+        let fn_ty = bx.type_func(&[llvm_elem_vec_ty, llvm_pointer, alignment_ty, mask_ty], ret_t);
+        let f = bx.declare_cfn(&llvm_intrinsic, llvm::UnnamedAddr::No, fn_ty);
+        let v = bx.call(
+            fn_ty,
+            None,
+            None,
+            f,
+            &[args[0].immediate(), args[1].immediate(), alignment, mask],
+            None,
+        );
+        return Ok(v);
+    }
+
     if name == sym::simd_scatter {
         // simd_scatter(values: <N x T>, pointers: <N x *mut T>,
         //             mask: <N x i{M}>) -> ()

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -541,6 +541,7 @@ pub fn check_platform_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>)
         sym::simd_fma => (1, 0, vec![param(0), param(0), param(0)], param(0)),
         sym::simd_gather => (3, 0, vec![param(0), param(1), param(2)], param(0)),
         sym::simd_masked_load => (3, 0, vec![param(0), param(1), param(2)], param(0)),
+        sym::simd_masked_store => (3, 0, vec![param(0), param(1), param(2)], Ty::new_unit(tcx)),
         sym::simd_scatter => (3, 0, vec![param(0), param(1), param(2)], Ty::new_unit(tcx)),
         sym::simd_insert => (2, 0, vec![param(0), tcx.types.u32, param(1)], param(0)),
         sym::simd_extract => (2, 0, vec![param(0), tcx.types.u32], param(1)),

compiler/rustc_span/src/symbol.rs

@@ -1533,6 +1533,7 @@ symbols! {
         simd_le,
         simd_lt,
         simd_masked_load,
+        simd_masked_store,
         simd_mul,
         simd_ne,
         simd_neg,