diff --git a/src/operator/numpy/np_cumsum-inl.h b/src/operator/numpy/np_cumsum-inl.h index 5d58b0254d..66cd088837 100644 --- a/src/operator/numpy/np_cumsum-inl.h +++ b/src/operator/numpy/np_cumsum-inl.h @@ -27,6 +27,8 @@ #include #include +#include +#include #include #include #include "../mxnet_op.h" @@ -34,6 +36,13 @@ #include "../elemwise_op_common.h" #include "../../api/operator/op_utils.h" +#if defined(__riscv_v_intrinsic) || defined(__riscv_vector) +#include +#define MXNET_NP_CUMSUM_USE_RVV 1 +#else +#define MXNET_NP_CUMSUM_USE_RVV 0 +#endif + namespace mxnet { namespace op { @@ -91,6 +100,61 @@ struct cumsum_forward { } }; +#if MXNET_NP_CUMSUM_USE_RVV +/*! + * \brief RVV cross-lane batched inclusive prefix-sum (cumsum) forward kernel. + * + * The scalar `cumsum_forward::Map` scans a single lane along `middle` with an + * element stride of `trailing`. Lanes that share the same `left` group (i.e. + * consecutive `right` indices) are stored at consecutive addresses, so up to + * VLMAX of them can be advanced together: a single vector accumulator holds + * the running prefix of every lane in the batch. Per lane the additions still + * happen strictly in `j = 1 .. middle-1` order (`acc = acc + x`), which is the + * exact order used by the scalar reference, so the results are bit-identical + * (no reassociation). + * + * Each `Map` call handles one contiguous chunk of at most VLMAX lanes inside + * one `left` group; `Kernel::Launch` then spreads the chunks across + * OpenMP threads exactly like the scalar path. The scan axis itself is never + * re-associated, so `axis=None` / `trailing == 1` shapes keep using the scalar + * kernel (there is no second axis to vectorize over). + */ +struct cumsum_forward_rvv { + MSHADOW_XINLINE static void Map(index_t chunk, + float* out, + const float* in, + const index_t middle, + const index_t trailing, + const index_t chunks_per_group) { + const size_t vlmax = __riscv_vsetvlmax_e32m4(); + const index_t left = chunk / chunks_per_group; + const index_t within = chunk % chunks_per_group; + const index_t right = within * static_cast(vlmax); + index_t right_end = right + static_cast(vlmax); + if (right_end > trailing) { + right_end = trailing; + } + const index_t base = left * middle * trailing + right; + const float* in0 = in + base; + float* out0 = out + base; + index_t i = right; + while (i < right_end) { + const size_t vl = __riscv_vsetvl_e32m4(static_cast(right_end - i)); + const index_t off = i - right; + vfloat32m4_t acc = __riscv_vle32_v_f32m4(in0 + off, vl); + __riscv_vse32_v_f32m4(out0 + off, acc, vl); + for (index_t j = 1; j < middle; ++j) { + const ptrdiff_t stride = static_cast(j) * static_cast(trailing); + vfloat32m4_t x = __riscv_vle32_v_f32m4(in0 + off + stride, vl); + acc = __riscv_vfadd_vv_f32m4(acc, x, vl); + __riscv_vse32_v_f32m4(out0 + off + stride, acc, vl); + } + i += static_cast(vl); + } + } +}; +#endif // MXNET_NP_CUMSUM_USE_RVV + template void CumsumForwardImpl(const OpContext& ctx, const TBlob& in, @@ -116,6 +180,26 @@ void CumsumForwardImpl(const OpContext& ctx, Stream* s = ctx.get_stream(); MSHADOW_TYPE_SWITCH_WITH_BOOL(in.type_flag_, IType, { MSHADOW_TYPE_SWITCH(out.type_flag_, OType, { +#if MXNET_NP_CUMSUM_USE_RVV + if constexpr (std::is_same::value && std::is_same::value && + std::is_same::value) { + if (trailing > 1) { + const size_t nlanes = out.Size() / middle; + const size_t vlmax = __riscv_vsetvlmax_e32m4(); + const index_t chunks_per_group = + static_cast((trailing + vlmax - 1) / vlmax); + Kernel::Launch( + s, + nlanes / trailing * static_cast(chunks_per_group), + out.dptr(), + in.dptr(), + middle, + trailing, + chunks_per_group); + return; + } + } +#endif Kernel::Launch( s, out.Size() / middle, out.dptr(), in.dptr(), middle, trailing); });