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 a129431c2884e4e0b2d8484850e2b9a936de13ff Mon Sep 17 00:00:00 2001 From: wanghongyan Date: Thu, 24 Sep 2026 15:43:47 +0800 Subject: [PATCH 3/3] [RISC-V] Vectorize float32 minimum reductions with RVV Co-authored-by: Yang Wang Co-authored-by: Yuansheng --- src/operator/tensor/broadcast_reduce-inl.h | 124 +++++++++++++++++++-- 1 file changed, 117 insertions(+), 7 deletions(-) diff --git a/src/operator/tensor/broadcast_reduce-inl.h b/src/operator/tensor/broadcast_reduce-inl.h index f9834d05b6..993489cb09 100644 --- a/src/operator/tensor/broadcast_reduce-inl.h +++ b/src/operator/tensor/broadcast_reduce-inl.h @@ -32,6 +32,9 @@ #include "../mshadow_op.h" #include "../mxnet_op.h" #include "../operator_common.h" +#if defined(__riscv_v_intrinsic) || defined(__riscv_vector) +#include +#endif namespace mxnet { namespace op { @@ -298,6 +301,88 @@ MSHADOW_XINLINE void assign(DType* dst, const bool addto, const DType src) { } } +namespace rvv_reduce { + +/*! + * \brief RVV kernel selector for the seq_reduce_assign inner k reduction. + * + * The primary template is disabled, so any reducer without a vector + * implementation keeps the scalar loop. Reducers that do have one specialize + * this template (minimum below; maximum/nrm2/product can be added alongside + * without touching the dispatch site). `Reduce` returns the reduced scalar for + * an arithmetic progression of `len` elements starting at `base` with element + * stride `stride`, continuing from the scalar accumulator `seed`. + */ +template +struct SeqReduceKernel { + static constexpr bool enabled = false; + static MSHADOW_XINLINE size_t MinLength() { return 0; } + static MSHADOW_XINLINE AType Reduce(const DType* /*base*/, + ptrdiff_t /*stride*/, + size_t /*len*/, + AType seed) { + return seed; + } +}; + +#if defined(__riscv_v_intrinsic) || defined(__riscv_vector) + +/*! + * \brief RVV 1.0 float32 minimum reduction matching mshadow::red::minimum. + * + * Reproduces the scalar `Reduce` contract bit-for-bit: + * - NaN propagates: the first NaN in the sequence becomes the result + * (the scalar reducer updates only while `dst` is not NaN); + * - ties keep the earlier element, so for (+0, -0) the sign of the first zero + * is preserved, whereas vfredmin's minNum semantics would always return -0. + * The per-chunk minimum combined with its first matching lane is exactly the + * first element attaining the running strict minimum, i.e. the scalar result. + */ +inline float mshadow_rvv_min_f32(const float* base, ptrdiff_t stride, size_t len, float seed) { + float best = seed; + size_t off = 0; + while (off < len) { + const size_t vl = __riscv_vsetvl_e32m1(len - off); + const float* p = base + off * stride; + vfloat32m1_t v = (stride == 1) ? + __riscv_vle32_v_f32m1(p, vl) : + __riscv_vlse32_v_f32m1(p, stride * static_cast(sizeof(float)), vl); + vbool32_t non_nan = __riscv_vmfeq_vv_f32m1_b32(v, v, vl); + const long nan_lane = __riscv_vfirst_m_b32(__riscv_vmnot_m_b32(non_nan, vl), vl); + if (nan_lane >= 0) { + // First NaN in the sequence: the scalar reducer stops updating afterwards. + best = base[(off + static_cast(nan_lane)) * stride]; + return best; + } + vfloat32m1_t vmin = __riscv_vfredmin_vs_f32m1_f32m1(v, v, vl); + const float chunk_min = __riscv_vfmv_f_s_f32m1_f32(vmin); + vbool32_t is_min = __riscv_vmfeq_vv_f32m1_b32(v, __riscv_vfmv_v_f_f32m1(chunk_min, vl), vl); + const long first = __riscv_vfirst_m_b32(is_min, vl); + const float candidate = base[(off + static_cast(first)) * stride]; + if (candidate < best) { + best = candidate; + } + off += vl; + } + return best; +} + +template <> +struct SeqReduceKernel { + static constexpr bool enabled = true; + static MSHADOW_XINLINE size_t MinLength() { return __riscv_vsetvlmax_e32m1(); } + static MSHADOW_XINLINE float Reduce(const float* base, + ptrdiff_t stride, + size_t len, + float seed) { + return mshadow_rvv_min_f32(base, stride, len, seed); + } +}; + +#endif // defined(__riscv_v_intrinsic) || defined(__riscv_vector) + +} // namespace rvv_reduce + template MSHADOW_XINLINE void binary_broadcast_assign(const index_t idx, const bool addto, @@ -361,13 +446,38 @@ MSHADOW_XINLINE void seq_reduce_assign(const index_t idx, AType val, residual; Reducer::SetInitValue(val, residual); if (!use_omp) { - 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)]); - // argmin/max, set IndexedNum.idx - if (IndexOP::do_op) - IndexOP::Op(&temp, k); - Reducer::Reduce(val, temp, residual); + // RVV fast path: only for reducers that provide a vector kernel (minimum + // today), for the identity elementwise op, and for a value-only reduction. + // It additionally requires the k offsets to form one arithmetic progression + // (exactly one reduced axis), which is the shape produced by + // BroadcastReduceShapeCompact for the 2-D cases handled here. + bool use_rvv = false; + if (!IndexOP::do_op && std::is_same::value && + rvv_reduce::SeqReduceKernel::enabled) { + int red_dim = -1; + int red_count = 0; + for (int d = 0; d < ndim; ++d) { + if (rshape[d] != 1) { + ++red_count; + red_dim = d; + } + } + use_rvv = (red_count == 1) && + (M >= rvv_reduce::SeqReduceKernel::MinLength()); + if (use_rvv) { + val = rvv_reduce::SeqReduceKernel::Reduce( + big + j, static_cast(rstride[red_dim]), M, val); + } + } + if (!use_rvv) { + 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)]); + // argmin/max, set IndexedNum.idx + if (IndexOP::do_op) + IndexOP::Op(&temp, k); + Reducer::Reduce(val, temp, residual); + } } } else { const int thread_count = engine::OpenMP::Get()->GetRecommendedOMPThreadCount();