From ae89f724aec6d5c2cbe8f3588eea314b63f88f07 Mon Sep 17 00:00:00 2001 From: wanghongyan Date: Wed, 23 Sep 2026 11:12:26 +0800 Subject: [PATCH 1/3] broadcast_reduce: strength-reduce the reduce index walk Co-authored-by: Yang Wang Co-authored-by: Yuansheng --- src/operator/tensor/broadcast_reduce-inl.h | 41 +++++++++++++++++++--- 1 file changed, 37 insertions(+), 4 deletions(-) diff --git a/src/operator/tensor/broadcast_reduce-inl.h b/src/operator/tensor/broadcast_reduce-inl.h index f9834d05b6..bcfe1d0c91 100644 --- a/src/operator/tensor/broadcast_reduce-inl.h +++ b/src/operator/tensor/broadcast_reduce-inl.h @@ -329,12 +329,31 @@ MSHADOW_XINLINE std::pair seq_reduce_assign_block(size_t start, Shape coord; AType val, residual{}; Reducer::SetInitValue(val, residual); + // Loop-level strength reduction: seed the coordinate/offset walk once for the + // whole block instead of recomputing mxnet_op::unravel(k, rshape) and + // mxnet_op::dot(coord, rstride) for every reduce element. Seeding is skipped + // for empty blocks so that a zero-sized reduce shape (len == 0, where rshape + // may contain a 0) never triggers a division by zero in unravel. + index_t off = 0; + if (len > 0) { + coord = mxnet_op::unravel(start, rshape); + off = mxnet_op::dot(coord, rstride); + } for (size_t k = start; k < start + len; ++k) { - coord = mxnet_op::unravel(k, rshape); - AType temp = OP::Map(big[j + mxnet_op::dot(coord, rstride)]); + AType temp = OP::Map(big[j + off]); if (IndexOP::do_op) IndexOP::Op(&temp, k); Reducer::Reduce(val, temp, residual); + // Advance to the coordinate of k + 1 (mixed-radix carry from the fastest + // dimension), keeping `off` in sync with dot(coord, rstride). + for (int i = ndim - 1; i >= 0; --i) { + if (++coord[i] < rshape[i]) { + off += rstride[i]; + break; + } + coord[i] = 0; + off -= (rshape[i] - 1) * rstride[i]; + } } return std::make_pair(val, residual); } @@ -361,13 +380,27 @@ MSHADOW_XINLINE void seq_reduce_assign(const index_t idx, AType val, residual; Reducer::SetInitValue(val, residual); if (!use_omp) { + // Loop-level strength reduction: walk the flat offset into `big` + // incrementally (mixed-radix carry starting from the fastest dimension) + // instead of recomputing mxnet_op::unravel(k, rshape) and + // mxnet_op::dot(coord, rstride) for every reduce element. + index_t offset = 0; + for (int i = 0; i < ndim; ++i) + coord[i] = 0; for (size_t k = 0; k < M; ++k) { - coord = mxnet_op::unravel(k, rshape); - AType temp = OP::Map(big[j + mxnet_op::dot(coord, rstride)]); + AType temp = OP::Map(big[j + offset]); // argmin/max, set IndexedNum.idx if (IndexOP::do_op) IndexOP::Op(&temp, k); Reducer::Reduce(val, temp, residual); + for (int i = ndim - 1; i >= 0; --i) { + if (++coord[i] < rshape[i]) { + offset += rstride[i]; + break; + } + coord[i] = 0; + offset -= (rshape[i] - 1) * rstride[i]; + } } } else { const int thread_count = engine::OpenMP::Get()->GetRecommendedOMPThreadCount(); From 722f8b0d9b57062c151d0b8886afc823199c5bfc Mon Sep 17 00:00:00 2001 From: wanghongyan Date: Wed, 23 Sep 2026 11:15:57 +0800 Subject: [PATCH 2/3] Revert "broadcast_reduce: strength-reduce the reduce index walk" This reverts commit ae89f724aec6d5c2cbe8f3588eea314b63f88f07. --- src/operator/tensor/broadcast_reduce-inl.h | 41 +++------------------- 1 file changed, 4 insertions(+), 37 deletions(-) diff --git a/src/operator/tensor/broadcast_reduce-inl.h b/src/operator/tensor/broadcast_reduce-inl.h index bcfe1d0c91..f9834d05b6 100644 --- a/src/operator/tensor/broadcast_reduce-inl.h +++ b/src/operator/tensor/broadcast_reduce-inl.h @@ -329,31 +329,12 @@ MSHADOW_XINLINE std::pair seq_reduce_assign_block(size_t start, Shape coord; AType val, residual{}; Reducer::SetInitValue(val, residual); - // Loop-level strength reduction: seed the coordinate/offset walk once for the - // whole block instead of recomputing mxnet_op::unravel(k, rshape) and - // mxnet_op::dot(coord, rstride) for every reduce element. Seeding is skipped - // for empty blocks so that a zero-sized reduce shape (len == 0, where rshape - // may contain a 0) never triggers a division by zero in unravel. - index_t off = 0; - if (len > 0) { - coord = mxnet_op::unravel(start, rshape); - off = mxnet_op::dot(coord, rstride); - } for (size_t k = start; k < start + len; ++k) { - AType temp = OP::Map(big[j + off]); + coord = mxnet_op::unravel(k, rshape); + AType temp = OP::Map(big[j + mxnet_op::dot(coord, rstride)]); if (IndexOP::do_op) IndexOP::Op(&temp, k); Reducer::Reduce(val, temp, residual); - // Advance to the coordinate of k + 1 (mixed-radix carry from the fastest - // dimension), keeping `off` in sync with dot(coord, rstride). - for (int i = ndim - 1; i >= 0; --i) { - if (++coord[i] < rshape[i]) { - off += rstride[i]; - break; - } - coord[i] = 0; - off -= (rshape[i] - 1) * rstride[i]; - } } return std::make_pair(val, residual); } @@ -380,27 +361,13 @@ MSHADOW_XINLINE void seq_reduce_assign(const index_t idx, AType val, residual; Reducer::SetInitValue(val, residual); if (!use_omp) { - // Loop-level strength reduction: walk the flat offset into `big` - // incrementally (mixed-radix carry starting from the fastest dimension) - // instead of recomputing mxnet_op::unravel(k, rshape) and - // mxnet_op::dot(coord, rstride) for every reduce element. - index_t offset = 0; - for (int i = 0; i < ndim; ++i) - coord[i] = 0; for (size_t k = 0; k < M; ++k) { - AType temp = OP::Map(big[j + offset]); + coord = mxnet_op::unravel(k, rshape); + AType temp = OP::Map(big[j + mxnet_op::dot(coord, rstride)]); // argmin/max, set IndexedNum.idx if (IndexOP::do_op) IndexOP::Op(&temp, k); Reducer::Reduce(val, temp, residual); - for (int i = ndim - 1; i >= 0; --i) { - if (++coord[i] < rshape[i]) { - offset += rstride[i]; - break; - } - coord[i] = 0; - offset -= (rshape[i] - 1) * rstride[i]; - } } } else { const int thread_count = engine::OpenMP::Get()->GetRecommendedOMPThreadCount(); From a3f99605807262f8034d0c7c272cb2f9af31b1d9 Mon Sep 17 00:00:00 2001 From: wanghongyan Date: Thu, 24 Sep 2026 15:50:26 +0800 Subject: [PATCH 3/3] [RISC-V] Vectorize batched cumsum over the middle axis Co-authored-by: Yang Wang Co-authored-by: Yuansheng --- src/operator/numpy/np_cumsum-inl.h | 84 ++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) 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); });