Skip to content

Commit d318646

Browse files
committed
refactor(cpu): simplify NEON where_select mask construction
Replace multi-step mask expansion (u8 -> u16 -> u32) with direct construction of full-width lane masks. For f32, build u32x4 directly from condition bytes. For f64, use vld1q_u64 instead of vcombine_u64. This reduces instruction count while maintaining identical behavior.
1 parent 10456f9 commit d318646

1 file changed

Lines changed: 11 additions & 15 deletions

File tree

  • src/runtime/cpu/kernels/simd/where_select/aarch64

src/runtime/cpu/kernels/simd/where_select/aarch64/neon.rs

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,6 @@ pub unsafe fn where_f32(cond: *const u8, x: *const f32, y: *const f32, out: *mut
2828
let chunks = len / F32_LANES;
2929
let remainder = len % F32_LANES;
3030

31-
let zero_u8 = vdup_n_u8(0);
32-
3331
for i in 0..chunks {
3432
let offset = i * F32_LANES;
3533

@@ -39,20 +37,18 @@ pub unsafe fn where_f32(cond: *const u8, x: *const f32, y: *const f32, out: *mut
3937
let c2 = *cond.add(offset + 2);
4038
let c3 = *cond.add(offset + 3);
4139

42-
// Create u8x8 vector with condition bytes in first 4 positions
43-
let cond_bytes = vcreate_u8(
44-
(c0 as u64) | ((c1 as u64) << 8) | ((c2 as u64) << 16) | ((c3 as u64) << 24),
40+
// Build full-width lane masks required by vbsl:
41+
// 0xFFFF_FFFF selects from x, 0x0000_0000 selects from y.
42+
let cond_mask_u32 = vld1q_u32(
43+
[
44+
if c0 != 0 { !0u32 } else { 0u32 },
45+
if c1 != 0 { !0u32 } else { 0u32 },
46+
if c2 != 0 { !0u32 } else { 0u32 },
47+
if c3 != 0 { !0u32 } else { 0u32 },
48+
]
49+
.as_ptr(),
4550
);
4651

47-
// Compare with zero to get mask: 0xFF for non-zero, 0x00 for zero
48-
let cond_mask_u8 = vcgt_u8(cond_bytes, zero_u8);
49-
50-
// Expand u8 mask to u32 mask (4 elements)
51-
// First expand u8x8 -> u16x8
52-
let cond_mask_u16 = vmovl_u8(cond_mask_u8);
53-
// Then expand low u16x4 -> u32x4
54-
let cond_mask_u32 = vmovl_u16(vget_low_u16(cond_mask_u16));
55-
5652
// Load x and y vectors
5753
let vx = vld1q_f32(x.add(offset));
5854
let vy = vld1q_f32(y.add(offset));
@@ -98,7 +94,7 @@ pub unsafe fn where_f64(cond: *const u8, x: *const f64, y: *const f64, out: *mut
9894
let m0: u64 = if c0 != 0 { !0u64 } else { 0u64 };
9995
let m1: u64 = if c1 != 0 { !0u64 } else { 0u64 };
10096

101-
let mask = vcombine_u64(vcreate_u64(m0), vcreate_u64(m1));
97+
let mask = vld1q_u64([m0, m1].as_ptr());
10298

10399
// Load x and y vectors
104100
let vx = vld1q_f64(x.add(offset));

0 commit comments

Comments
 (0)