ENH: Add SIMD implementation for heaviside - #19780

@seiko2plus I am wondering should I take NaN and Inf into account? In the current implementation, I ensure the output of a NaN input is NaN. However, in the tests, I got error RuntimeWarning: invalid value encountered in heaviside.
Should we reject NaN and Inf as input value?

@howjmay,

I am wondering should I take NaN and Inf into account?
yes, we should and your final implementation should be equivalent to the scalar function:

numpy/numpy/core/src/npymath/npy_math_internal.h.src

Lines 580 to 600 in fae6fa4

	/**begin repeat
	* #type = npy_float, npy_double, npy_longdouble#
	* #c = f, ,l#
	* #C = F, ,L#
	*/
	@type@ npy_heaviside@c@(@type@ x, @type@ h0)
	{
	if (npy_isnan(x)) {
	return (@type@) NPY_NAN;
	}
	else if (x == 0) {
	return h0;
	}
	else if (x < 0) {
	return (@type@) 0.0;
	}
	else {
	return (@type@) 1.0;
	}
	}

I got error RuntimeWarning: invalid value encountered in heaviside.
Should we reject NaN and Inf as input value?

I think this error caused by the scalar function, try to clear FP barrier by adding the following code after the scalar loop:

    npy_clear_floatstatus_barrier((char*)dimensions);

regards to your current impl of heaviside, I think it need to get improved. Try something like that:

    const npyv_f32 zero = npyv_zero_f32();
    const npyv_f32 one  = npyv_setall_f32(1.0f);
    npyv_b32 mask_nnan = npyv_notnan_f32(x);
    npyv_b32 mask_zero = npyv_cmpeq_f32(x, zero);
    // arithmetic right shift to the edge of the exponent
    npyv_f32 shift_exp  = npyv_reinterpret_f32_s32(
        npyv_shri_s32(npyv_reinterpret_s32_f32(x), 8)
    );
    // x < 0 ? 0 : 1, since the fraction of 1.0 is zero.
    // npyv_f32 zero_or_one = npyv_andc_f32(one, shift_exp); better to have intrin for `AND With Complement`, on x86 named `andnot`
    npyv_f32 zero_or_one = npyv_and_f32(one, npyv_not_f32(shift_exp));
    return npyv_select_f32(mask_zero, h, zero_or_one);

@howjmay, try to use _simd module during the design stage, it can your reduce time/efforts.

from numpy.core import _simd as simd
v = simd.baseline # for AVX2 use simd.FMA3__AVX2 or simd.targets['FMA3__AVX2']
print("the available targets ", simd.targets)

def heaviside_f32(x, h):
   zero = v.zero_f32();
   one  = v.setall_f32(1.0);
   mask_nnan = v.notnan_f32(x);
   mask_zero = v.cmpeq_f32(x, zero);
   # ....
# test your kernel
heaviside_f32(v.setall_f32(1), v.setall_f32(2))