Repository navigation
Add RVV (RISC-V Vector Extension) optimized convolution and pooling kernels for the NCHWc blocked format in MLAS - #28411
Conversation
|
Hi Hariharan Seshadri (@hariharans29) |
There was a problem hiding this comment.
Pull request overview
This PR extends MLAS’ riscv64 RVV support by wiring up optimized float32 convolution and NCHWc pooling kernels, enabling the NCHWc blocked-format fast paths (BlockSize=16) on RVV-capable systems and replacing the previous generic depthwise implementation.
Changes:
- Adds new riscv64 RVV kernel implementations for direct NCHW/NCHWc conv, depthwise/pointwise NCHWc conv, and max/avg pooling in NCHWc format.
- Wires the new kernels into
MLAS_PLATFORMinitialization for riscv64 when RVV is available, and enables NCHWc fast-path selection for RISCV64+RVV. - Updates build configuration to compile the new RVV sources and swap out the previous depthwise kernel source for RISCV64 builds.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/core/mlas/lib/snchwc.cpp | Enables RISCV64+RVV to use platform-selected NCHWc conv/pool kernels and block size. |
| onnxruntime/core/mlas/lib/riscv64/sconv_nchwc_kernel_rvv.cpp | New RVV implementations for NCHW/NCHWc conv, depthwise/pointwise NCHWc conv, and NCHWc pooling. |
| onnxruntime/core/mlas/lib/riscv64/sconv_depthwise_kernel_rvv.cpp | New RVV 3x3 s1 depthwise CHW kernel implementation for the multiplier-1 path. |
| onnxruntime/core/mlas/lib/platform.cpp | Registers RVV NCHWc conv/pool kernels and sets NCHWc block size for riscv64 when RVV is present. |
| onnxruntime/core/mlas/lib/mlasi.h | Declares new RVV kernel entry points and adds RISCV64+RVV NCHWc members to MLAS_PLATFORM. |
| onnxruntime/core/mlas/lib/convolve.cpp | Enables the depthwise-direct algorithm path on RISCV64 and updates stride restrictions comment/logic. |
| onnxruntime/core/mlas/inc/mlas.h | Exposes MlasConvAlgorithmDepthwise in the enum for RISCV64 builds. |
| cmake/onnxruntime_mlas.cmake | Adds new RVV kernel sources and adjusts which depthwise source is compiled for RISCV64/RVV. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
Both comments point to the same concern: "What if In practice, however, this is not an issue because the minimum VLEN on current RISC-V CPUs is 128 bits. Therefore, this case is effectively covered (though I have implemented changes regardless). The calculation is as follows:
|
Thanks for the clarification. It looks good to me. Let me get one last opinion from Copilot. |
|
Copilot generated a lot more comments in this round - can you please take a look ? |
I've finished the modifications, mainly adding some asserts. Thanks a lot! |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
onnxruntime/core/mlas/lib/convolve.cpp:1572
- On riscv64 builds that enable MLAS_USE_RVV, the depthwise fast-path (MlasConvAlgorithmDepthwise) can still be selected purely based on compile-time macros. In this configuration the depthwise implementation is RVV-only (and the scalar/generic depthwise source is removed in CMake), so selecting this algorithm when RVV is not available at runtime (or when ORT_MLAS_RISCV_FORCE_SCALAR disables RVV dispatch) will execute vector instructions and can fault (illegal instruction). Consider gating depthwise algorithm selection on runtime RVV availability (e.g., a platform dispatch pointer/flag) or keeping a non-RVV fallback implementation compiled in and dispatching accordingly.
#if defined(MLAS_TARGET_WASM_SCALAR) || defined(MLAS_TARGET_ARM64) || defined(MLAS_TARGET_RISCV64)
// Scalar (WASM_SCALAR) / vectorized (ARM64/RISCV64) direct conv for depthwise convolution.
// Currently only support 3x3 kernel with padding <=1 and dilations = 1
// and on ARM64/RISCV64, it is further restricted to strides = 1.
// TODO: support more general depthwise convolution.
#if defined(MLAS_TARGET_ARM64) || defined(MLAS_TARGET_RISCV64)
bool depthwise_conv_stride_support_check = Parameters->StrideShape[0] == 1 && Parameters->StrideShape[1] == 1;
#else
bool depthwise_conv_stride_support_check = true;
#endif
if (Dimensions == 2
&& Parameters->FilterCount == 1 && Parameters->InputChannels == 1
&& Parameters->KernelShape[0] == 3 && Parameters->KernelShape[1] == 3
&& Parameters->Padding[0] <= 1 && Parameters->Padding[1] <= 1
&& Parameters->Padding[2] <= 1 && Parameters->Padding[3] <= 1
&& depthwise_conv_stride_support_check
&& Parameters->DilationShape[0] == 1 && Parameters->DilationShape[1] == 1) {
|
Could you please re-run the CI? |
0e72188
into
microsoft:main
PR microsoft#28261 installed the riscv64 RVV kernels only after asking the auxiliary vector whether the part implements the vector extension, and offered ORT_MLAS_RISCV_FORCE_SCALAR to turn that answer off again. PR microsoft#28411 replaced that block with a test of the vector register width, which removed the check, the override, and their supporting code. platform.cpp is one of the portable MLAS sources, so it is compiled without the vector extension. That is deliberate rather than an oversight: this is the translation unit that decides whether the part implements the extension, so it has to be able to run where it does not. Reading the register width is a vector builtin, so it cannot be done from here. cmake/riscv64.toolchain.cmake passes no -march, so whether this file compiles is decided by the toolchain's default architecture, and on a toolchain whose default omits the extension it does not compile at all: error: built-in function '__riscv_vlenb()' requires the 'xtheadvector or v' ISA extension Where the default does include the extension, every kernel is installed on the strength of the build flag alone, the first vector instruction reached being the one in this file, during platform construction. The check and the override are restored. The register width test is dropped: the extension the check requires implies VLEN>=128, which is exactly the width the NCHWc channel block of sixteen floats needs in one LMUL=4 register group, so the block size is the fixed value it always was and no vector register has to be read to confirm it.
## Description ### Motivation PR #28261 installed the MLAS RVV kernels only when the runtime reported the RISC-V V extension, and provided `ORT_MLAS_RISCV_FORCE_SCALAR` as an escape hatch. PR #28411 replaced that gate with an unconditional `__riscv_vlenb()` call and removed the override. This causes two problems: - `platform.cpp` is compiled without requiring `-march=...v`, but `__riscv_vlenb()` requires the vector extension at compile time. Toolchains whose default architecture does not include V therefore reject the build. - An RVV-enabled build can install and execute vector kernels on a riscv64 system that does not provide the V extension. ### Changes - Restore the Linux HWCAP runtime check before installing MLAS RVV kernels. - Restore the `ORT_MLAS_RISCV_FORCE_SCALAR` override. - Remove the `__riscv_vlenb()` check from `platform.cpp`. The V extension guarantees VLEN >= 128, which is sufficient for the existing sixteen-float NCHWc block. - Keep the change confined to `onnxruntime/core/mlas/lib/platform.cpp`. ### Testing - Cross-compiled the RVV-enabled Release build with the XuanTie toolchain under `-Werror`: 1490 targets completed without errors. - Ran `onnxruntime_mlas_test` on a SpacemiT K3 (X100, VLEN=256): no failures. - Ran the relevant MLAS suites on a XuanTie A210 (C920V2, VLEN=128): 2344 passed and 0 failed with RVV enabled; 1800 passed and 0 failed with `ORT_MLAS_RISCV_FORCE_SCALAR=1`. - The forced-scalar run reduced the NCHWc-related test count from 549 to 5, confirming that the runtime fallback disables the NCHWc RVV path. A physical riscv64 system without the V extension was not available for testing; the scalar fallback was validated through the environment override.
### Description Improve FP32 inference in the MLAS riscv64 RVV backend with five commits: 1. **Softmax:** accumulate across each row, reduce once at the end, and use LMUL 4. 2. **NCHWc convolution:** compute four adjacent outputs with independent accumulators and shared filter loads. 3. **NCHWc register groups:** select LMUL 1/2/4 from the vector length, preserving the sixteen-float block layout. 4. **Depthwise convolution:** extend the existing `sconv_depthwise_kernel_rvv.cpp` to additional shapes and strides, with unit dilation and kernel width up to sixteen. Enable the wider route only when RVV kernels are installed. 5. **Activation:** add an RVV `MlasActivation` routine, including the clamp used by ReLU6. Changes are limited to MLAS and its CMake source list. Public interfaces remain unchanged. **Performance** A210/C920V2: VLEN=128; SpacemiT K3/X100: VLEN=256. **Baseline** is unmodified upstream `main` at `9f913ae52`, which already includes #32406. **Optimized** is that baseline with this PR's five RVV kernel changes. Both use the CPU Execution Provider, `ORT_ENABLE_ALL`, and RVV kernels enabled, with the same toolchain and build configuration. Runs use one pinned core, one thread, warmup, fixed frequency, and best of three. Speedup is baseline latency divided by optimized latency. These measurements precede the final depthwise file consolidation and latest rebase. | Model | Input shape | C920V2 baseline (ms) | C920V2 optimized (ms) | Speedup | X100 baseline (ms) | X100 optimized (ms) | Speedup | |---|---|---|---|---|---|---|---| | [retinaface](https://github.com/serengil/retinaface) | 1x640x640x3 NHWC | 22300.2 | 7404.1 | **3.01x** | 11327.5 | 4037.0 | **2.81x** | | [resnet50-v1-7](https://github.com/onnx/models/blob/main/validated/vision/classification/resnet/model/resnet50-v1-7.onnx) | 1x3x224x224 NCHW | 1398.1 | 567.3 | **2.46x** | 816.3 | 271.6 | **3.01x** | | [mobilenetv2-12](https://github.com/onnx/models/blob/main/validated/vision/classification/mobilenet/model/mobilenetv2-12.onnx) | 1x3x224x224 NCHW | 148.2 | 70.7 | **2.10x** | 93.9 | 30.4 | **3.09x** | | [mobilevit-s](https://github.com/andestech/ModelZoo/blob/master/MobileViT-Small/Model/mobilevit-s.onnx) | 1x3x256x256 NCHW | 845.8 | 419.7 | 2.02x | 356.8 | 200.2 | 1.78x | | [shufflenet-v2-10](https://github.com/onnx/models/blob/main/validated/vision/classification/shufflenet/model/shufflenet-v2-10.onnx) | 1x3x224x224 NCHW | 68.5 | 47.2 | 1.45x | 41.9 | 20.3 | 2.06x | | [efficientnet-b0](https://huggingface.co/qualcomm/EfficientNet-B0/blob/main/EfficientNet-B0.onnx) | 1x3x224x224 NCHW | 244.1 | 189.5 | 1.29x | 124.3 | 67.1 | 1.85x | RetinaFace uses the official serengil TensorFlow/ResNet50 implementation and weights, exported with tf2onnx at opset 13. **Validation** Before the latest rebase: XuanTie cross-compilation passed under `-Werror`. After depthwise consolidation, the K3 convolution/activation subset passed **1109 tests with RVV** and **565 with forced scalar**, with zero failures and three HalfConv skips in each mode. Softmax/log-softmax differential checks against a double-precision reference passed on both boards. ### Motivation and Context Follow-up to #28261 and #28411, building on the runtime gate restored in #32406. The existing NCHWc kernels reached only 0.20x–0.84x of the alternative convolution path's throughput across seven measured shapes. Independent accumulators and appropriately sized register groups address these bottlenecks while preserving the blocked layout. Broader depthwise routing avoids im2col plus a single-row GEMM per channel. It benefits convolutions that cannot use NCHWc, including some ShuffleNet layers at the default optimization level, as well as `ORT_ENABLE_EXTENDED`. Fewer softmax reductions and vectorized activation address the remaining overhead.
Description
New kernel files:
Motivation and Context
Following #28261, Optimize more MLAS kernels using RISC-V Vector (RVV) extensions.
Please Note:
On the K3 (SpacemiT X60), VLEN=256. With LMUL=4 and e32, the hardware can hold (256/32) * 4 = 32 floats per vector register group — but we only request 16. So we're using half the available vector width.
The reason is that BlockSize=16 is baked into the NCHWc data layout across the whole framework (matching ARM64 NEON). Changing it to 32 would require a different NCHWc format and is not a localized change.
Benchmark ((SpacemiT K3, VLEN=256, 8-core))
All tests pass with zero numerical error.