Use llvm.memset.p0i8.* to initialize all same-bytes arrays

rust-lang · Jan 8, 2025 · 3ca023e · 3ca023e
1 parent f278858
commit 3ca023e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -97,11 +97,15 @@ 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 zero arrays
-                    if bx.cx().const_to_opt_u128(v, false) == Some(0) {
-                        let fill = bx.cx().const_u8(0);
-                        bx.memset(start, fill, size, dest.val.align, MemFlags::empty());
-                        return true;
+                    // Use llvm.memset.p0i8.* 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 first = bytes[0];
+                        if bytes[1..].iter().all(|&b| b == first) {
+                            let fill = bx.cx().const_u8(first);
+                            bx.memset(start, fill, size, dest.val.align, MemFlags::empty());
+                            return true;
+                        }
                     }
 
                     // Use llvm.memset.p0i8.* to initialize byte arrays

diff --git a/tests/codegen/slice-init.rs b/tests/codegen/slice-init.rs
@@ -71,10 +71,10 @@ const N: usize = 100;
 #[no_mangle]
 pub fn u16_init_one_bytes() -> [u16; N] {
     // CHECK-NOT: select
-    // CHECK: br
+    // CHECK-NOT: br
     // CHECK-NOT: switch
-    // CHECK: icmp
-    // CHECK-NOT: call void @llvm.memset.p0
+    // CHECK-NOT: icmp
+    // CHECK: call void @llvm.memset.p0
     [const { u16::from_be_bytes([1, 1]) }; N]
 }