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 3aadba0b3a5edf2cad4c5e08e2d52c9ebde1fde9 Mon Sep 17 00:00:00 2001 From: wanghongyan Date: Thu, 24 Sep 2026 15:47:12 +0800 Subject: [PATCH 3/3] [RISC-V] Vectorize col2im accumulation in convolution backward Co-authored-by: Yang Wang Co-authored-by: Yuansheng --- src/operator/nn/im2col.h | 103 ++++++++++++++++++++++++++++++++++----- 1 file changed, 91 insertions(+), 12 deletions(-) diff --git a/src/operator/nn/im2col.h b/src/operator/nn/im2col.h index 82acc1d2ce..c31eb7b4f9 100644 --- a/src/operator/nn/im2col.h +++ b/src/operator/nn/im2col.h @@ -62,10 +62,19 @@ #include #include +#include #include +#include #include #include "../mxnet_op.h" +#if defined(__riscv_v_intrinsic) || defined(__riscv_vector) +#include +#define MXNET_IM2COL_USE_RVV 1 +#else +#define MXNET_IM2COL_USE_RVV 0 +#endif + namespace mxnet { namespace op { @@ -79,6 +88,53 @@ inline bool is_a_ge_zero_and_a_lt_b(int a, int b) { return static_cast(a) < static_cast(b); } +/*! + * \brief Accumulate a fixed-stride run of `n` column elements into `dst` + * (`dst[i * stride] += src[i]`). On RISC-V targets that expose the + * vector extension this uses RVV unit-/strided loads and stores; + * on every other target it degrades to the plain scalar loop. + * + * `stride` is expressed in elements and must be >= 1. Because stride >= 1, + * the lane addresses are pairwise distinct, so the vector accumulate is + * element-wise identical to the scalar loop: every `dst` element still + * receives the same per-tap `+=` in the same order (no cross-lane + * reassociation and no change of the FP accumulation semantics). + */ +template +inline void col2im_accumulate_strided_row(const DType* src, DType* dst, int stride, int n) { +#if MXNET_IM2COL_USE_RVV + if (std::is_same::value) { + const float* s = reinterpret_cast(src); + float* d = reinterpret_cast(dst); + int i = 0; + if (stride == 1) { + while (i < n) { + size_t vl = __riscv_vsetvl_e32m4(static_cast(n - i)); + vfloat32m4_t vc = __riscv_vle32_v_f32m4(s + i, vl); + vfloat32m4_t vd = __riscv_vle32_v_f32m4(d + i, vl); + __riscv_vse32_v_f32m4(d + i, __riscv_vfadd_vv_f32m4(vd, vc, vl), vl); + i += static_cast(vl); + } + } else { + const ptrdiff_t byte_stride = + static_cast(stride) * static_cast(sizeof(float)); + while (i < n) { + const ptrdiff_t offset = static_cast(i) * stride; + size_t vl = __riscv_vsetvl_e32m4(static_cast(n - i)); + vfloat32m4_t vc = __riscv_vle32_v_f32m4(s + i, vl); + vfloat32m4_t vd = __riscv_vlse32_v_f32m4(d + offset, byte_stride, vl); + __riscv_vsse32_v_f32m4(d + offset, byte_stride, __riscv_vfadd_vv_f32m4(vd, vc, vl), vl); + i += static_cast(vl); + } + } + return; + } +#endif + for (int i = 0; i < n; ++i) { + dst[static_cast(i) * stride] += src[i]; + } +} + /*! * \brief im2col 2D cpu version. * DO NOT call this function directly. @@ -312,20 +368,43 @@ inline void col2im_cpu(const DType* data_col, for (int channel = channels; channel--; data_im += channel_size) { for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) { for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) { - int input_row = -pad_h + kernel_row * dilation_h; + // The input column touched at output position `o` is + // input_col(o) = input_col0 + o * stride_w. + // Only positions with 0 <= input_col(o) < width contribute to the + // image; because stride_w >= 1 that set is a single contiguous + // interval [o_start, o_end]. Padding positions contribute nothing, so + // they are skipped by only advancing the column pointer (no + // out-of-bounds pointer construction and no per-element bounds + // branch). The in-range run is accumulated with RVV loads/stores + // (see col2im_accumulate_strided_row). + const int input_col0 = -pad_w + kernel_col * dilation_w; + int o_start = 0; + if (input_col0 < 0) { + o_start = (-input_col0 + stride_w - 1) / stride_w; + } + int o_end = output_w - 1; + if (input_col0 + (output_w - 1) * stride_w >= width) { + const int last_valid = width - 1 - input_col0; + o_end = (last_valid < 0) ? (o_start - 1) : (last_valid / stride_w); + if (o_end > output_w - 1) { + o_end = output_w - 1; + } + } + const int n_valid = (o_end >= o_start) ? (o_end - o_start + 1) : 0; + int input_row = -pad_h + kernel_row * dilation_h; for (int output_rows = output_h; output_rows; output_rows--) { - if (!is_a_ge_zero_and_a_lt_b(input_row, height)) { - data_col += output_w; - } else { - int input_col = -pad_w + kernel_col * dilation_w; - for (int output_col = output_w; output_col; output_col--) { - if (is_a_ge_zero_and_a_lt_b(input_col, width)) { - data_im[input_row * width + input_col] += *data_col; - } - data_col++; - input_col += stride_w; - } + if (n_valid > 0 && is_a_ge_zero_and_a_lt_b(input_row, height)) { + col2im_accumulate_strided_row( + data_col + o_start, + data_im + input_row * width + input_col0 + o_start * stride_w, + stride_w, + n_valid); } + // The column buffer is laid out [kernel_col][output_row][output_col]; + // every output position consumes exactly one column entry, so the + // pointer advance is independent of whether the position was + // in-range or padding. + data_col += output_w; input_row += stride_h; } }