RISC-V Vector Extension (RVV) support for ATen by zhangfeiv0 · Pull Request #46 · RuyiAI-Stack/pytorch

zhangfeiv0 · 2026-08-31T06:19:43Z

嗨, 我是软件所的张飞, 我之前有做过 pytorch 相关工作:
pytorch#127867
pytorch#143979
pytorch#166602
pytorch#135570
其中, pytorch#135570 的工作并没有受上游太重视且审核缓慢, 希望这个 PR 内容能够合并到 RuyiAI, , 希望随着你们的 ci 等推进一并合并到上游。

变更说明 / Summary

为 pytorch vec 后端添加了 rvv 后端支持, 且该 PR 已适配到最新的上游代码

关联 Issue / Related issue

下面是提交给上游的pr(均未合并), 因长时间未改动已被关闭, 其中一个是联合RISE提交的:
pytorch#175746
pytorch#135570

验证方式 / Validation

gcc: gcc-15
平台: sg2044
vec库功能测试: ./build/bin/vec_test_all_types_RVV, 结果均已通过:

[----------] Global test environment tear-down
[==========] 355 tests from 134 test suites ran. (38973 ms total)
[  PASSED  ] 355 tests.

检查清单 / Checklist

The change is focused and contains no unrelated modifications.
I have followed the target repository's contribution guidelines.
I have added or updated tests where applicable.
I have updated related documentation where applicable.
I have run the relevant formatting, lint, build, and test checks.
I have described known limitations or compatibility impact.

* Add RISC-V 64 BLOCK_LIST * Skip long time testcase Co-authored-by: Cursor <cursoragent@cursor.com>

* Add riscv64 ci with PR

⭐ Run Main Diff base and head Push to riscv From https://github.com/RuyiAI-Stack/pytorch * branch riscv -> FETCH_HEAD fatal: Not a valid object name origin/main Error: ❌ Failure - Main Diff base and head Error: exit status 128

* mklnn is unavailable on RISC-V * Remove test_cpu_select_algorithm from block_list * Fix block list format

These cases are too slow on riscv64, adding them to here simply Drop test_torch from the list because it is one core case

bytes_to_scalar previously round-tripped raw bytes through Python float/complex values (via ctypes) before constructing the tensor. This loses NaN bit patterns on architectures (such as RISC-V) that canonicalize NaNs in floating-point loads/conversions, causing test_bytes_to_scalar_cpu_{float32,float64,complex64,complex128} to fail with mismatched storage bytes. Construct the scalar tensor by writing the raw bytes directly into its untyped storage so all input bit patterns (including NaN payloads) are preserved exactly.

RISC-V converts non-finite floats to integers by saturating: -inf -> min, inf/nan -> max for wider int types. Add IS_RISCV64 flag and RISC-V-specific reference values. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

for action, these yaml must be merged first then take effect, so merge it skipping ci

…orch#40)

Signed-off-by: Fei Zhang <zhangfei@iscas.ac.cn> Co-authored-by: Chung-Lin Tang <cltang@baylibre.com>

zhangfeiv0 · 2026-09-02T07:59:38Z

@zhanghb97 后面该pr合并的要求是啥？后续我这里还要做啥？

yuzibo · 2026-09-03T03:26:13Z

Hi, 感谢贡献 PR!

我个人先从 CI 的角度介绍下, 我们现在的 CI job 只开启了 core test, 前他的全量测试集正在开发测试中.

好消息是这个 PR 看起来在 core test pass.

XYenChi · 2026-09-03T16:59:51Z

@zhangfeiv0 Hi 张老师，感谢您的工作！注意到 FindRVV.cmake 固定使用 -march=rv64gcv_zvl128b，且算子优化时默认VLEN=128。在 VLEN > 128 的机器上，部分算子优化性能不如编译器 RVV 自动向量化。是否考虑探测 VLEN ？

zhangfeiv0 · 2026-09-04T01:31:20Z

@zhangfeiv0 Hi 张老师，感谢您的工作！注意到 FindRVV.cmake 固定使用 -march=rv64gcv_zvl128b，且算子优化时默认VLEN=128。在 VLEN > 128 的机器上，部分算子优化性能不如编译器 RVV 自动向量化。是否考虑探测 VLEN ？

之所以这么做，是因为pytorch vec后端框架的历史遗留问题导致的，其设计不适合vector length agnostic，目前上游 SVE 和本 PR增添的rvv都是 vls设计，相关讨论可见：
pytorch#153471
pytorch#175746 (comment)

好像目前上游对VLA的推动比较缓慢？这个推进可能会影响更多的架构和需要详细的验证。

XYenChi · 2026-09-07T06:21:47Z

注意到您向上游提交patch时附带了性能测试结果
能否提供patch优化与编译器自动向量化的性能数据？

zhangfeiv0 · 2026-09-09T03:49:07Z

注意到您向上游提交patch时附带了性能测试结果能否提供patch优化与编译器自动向量化的性能数据？

测试平台sg2044，编译器gcc-15，以下是我的测试结果：

测试内容覆盖与总体对比

类别	内容
pointwise	31 种 FP32 算子 × 64/4096/262144/4194304 元素
dtype	float64/float16/bfloat16/int32/int64; 算术、abs、sum 和适用的数学函数
layout	stride2、转置、广播、offset1; add/mul/exp/ReLU
tail	65539 元素非整齐长度
stream	16M 元素 add/mul/copy
reduction	sum/mean/max/min/var/norm、按行/列归约
normalization	softmax/log_softmax/layer_norm/rms_norm
nn	batch_norm、池化; NCHW/channels_last
linear_algebra	128/512 方阵 GEMM、BMM、conv2d
autograd	GELU/LayerNorm/Softmax/逐元素平方误差 forward+backward
optimizer	Adam 状态更新
quantization	qint8/qint8/qint32 量化/反量化及 qadd/qmul/qrelu
model	MLP 512→1024→512，简易 2 层 CNN，1 层 TransformerEncoder；另测 MLP 梯度

类别	1 线程几何平均	8 线程几何平均
pointwise	1.106×	1.066×
dtype	1.066×	1.045×
layout	1.364×	1.237×
tail	1.178×	1.176×
stream	1.208×	1.304×
reduction	0.875×	0.920×
normalization	1.057×	1.069×
nn	1.174×	1.207×
linear_algebra	1.025×	1.317×
autograd	1.076×	1.112×
optimizer	1.028×	1.127×
quantization	1.388×	1.333×
model	1.009×	1.218×

分类几何平均加速比；超过 1 表示 VEC RVV 更快。

总体效果可能会掩盖部分算子的情况，以下是更进一步的测试：

逐元素操作性能对比

模型与训练

配置	线程	auto (ms)	VEC RVV (ms)	加速比
cnn_inference 1x3x64x64	1	3.578	3.543	1.010×
cnn_inference 8x3x64x64	1	31.016	31.188	0.994×
mlp_inference 1x512	1	2.892	2.932	0.986×
mlp_inference 32x512	1	7.065	7.012	1.008×
mlp_training 32x512	1	20.845	20.376	1.023×
transformer_inference 2x128x256	1	45.039	43.539	1.034×
cnn_inference 1x3x64x64	8	79.919	39.964	2.000×
cnn_inference 8x3x64x64	8	739.973	669.952	1.105×
mlp_inference 1x512	8	3.280	3.327	0.986×
mlp_inference 32x512	8	19.996	19.996	1.000×
mlp_training 32x512	8	59.956	39.950	1.501×
transformer_inference 2x128x256	8	79.956	79.962	1.000×

可以看到还是有许多退化的用例，由于我这个pr写于几年前，在此期间RISE也没有进一步优化，因此这需要后期不断的进行优化，我也很乐意RuyiAI团队能够在此基础上持续优化；同时vec rvv后端支持也是必须的，如果自动矢量化更优，那么其他架构在vec后端支持也毫无意义。

XYenChi · 2026-09-10T13:34:12Z

能否改写成用户自行配置使用自动向量化还是使用intrinsic的分发逻辑？

zhangfeiv0 · 2026-09-11T05:58:03Z

能否改写成用户自行配置使用自动向量化还是使用intrinsic的分发逻辑？

关于这个问题，我认为当前代码就支持用户自行配置编译选项和使用default默认后端，理由如下：

1.pytorch目前支持"同一份 kernel 编译多个版本，再由 DispatchStub 选择"的方式，上游说明也记录了这种机制(https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cpu/README.md)；

例如，同一份 UnaryOpsKernel.cpp 分别生成 UnaryOpsKernel.cpp.DEFAULT.cpp 和 UnaryOpsKernel.cpp.RVV.cpp,UnaryOpsKernel.cpp.DEFAULT.cpp会走默认后端实现,所以 DEFAULT 的实现可由 GCC 自动向量化, UnaryOpsKernel.cpp.RVV.cpp会走vec rvv intrinsics实现;

因此对于用户来说即使编译使用了此版本的RVV wheel，仍能控制其走不同的实现，使用方式如下:

ATEN_CPU_CAPABILITY=default python your_script.py
ATEN_CPU_CAPABILITY=rvv python your_script.py

也可以在 Python 程序开头设置：

import os
os.environ["ATEN_CPU_CAPABILITY"] = "default"

import torch
print(torch.backends.cpu.get_cpu_capability())  # DEFAULT

具体实际分发矩阵如下：

wheel / 场景	capability	neg / abs / mul 实际分发
RVV wheel，未设置变量	RVV	全部 RVV
RVV wheel，启动前设置 `default`	DEFAULT	全部 DEFAULT
RVV wheel，启动前设置 `rvv`	RVV	全部 RVV
RVV wheel，Python 中 import torch 前设置 `default`	DEFAULT	全部 DEFAULT
RVV wheel，import torch 后、首次 capability 查询前设置 `default`	DEFAULT	全部 DEFAULT，本 wheel 的实测结果
RVV wheel，先查询 capability 并执行 neg，再改成 `default`	RVV	neg 保持 RVV，随后首次执行的 abs/mul 也走 RVV
自动向量化 wheel，未设置变量	DEFAULT	全部 DEFAULT，未注册 RVV
自动向量化 wheel，设置 `rvv`	DEFAULT	警告该值无效，全部 DEFAULT

2.假如用户用 -march=rv64gcv_zvfh 构建 DEFAULT，再通过 ATEN_CPU_CAPABILITY=default 运行它，当前代码支持。但 VEC库 RVV intrinsic 编译版本会被固定的 -march 覆盖，FindRVV.cmake 设置的参数最终只追加到 RVV 编译单元，不会给 DEFAULT 追加同样的参数：

用户构建时的 `CMAKE_CXX_FLAGS`	DEFAULT 最终生效的 ISA 参数	VEC RVV 最终生效的 ISA 参数
`-march=rv64gcv`	`-march=rv64gcv`	`-march=rv64gcv_zvl128b`
`-march=rv64gcv_zvfh`	`-march=rv64gcv_zvfh`	`-march=rv64gcv_zvl128b`

3.除此之外，我在最新的提交中引入了官方提供的宏USE_CPU_VECTORIZATION来支持只生成DEFAULT，因此“只构建通用 vec，允许按用户指定 ISA 自动向量化”，可以这样配置后构建：

export USE_CPU_VECTORIZATION=0
export CMAKE_CXX_FLAGS="-march=rv64gcv_zvfh"

综上，主要有两个开关可供使用：
USE_CPU_VECTORIZATION=0：构建时不生成额外的 ATen CPU kernel 版本。
ATEN_CPU_CAPABILITY=default：运行时从已有的多个版本中选择 DEFAULT。

XYenChi and others added 21 commits August 31, 2026 03:00

Add RISC-V Blocklist (pytorch#1)

0479f8b

* Add RISC-V 64 BLOCK_LIST * Skip long time testcase Co-authored-by: Cursor <cursoragent@cursor.com>

add riscv64 ci

fb0bd0a

Test ci with PR (pytorch#8)

98c8c8a

* Add riscv64 ci with PR

Fix no main brach issue

26aab04

⭐ Run Main Diff base and head Push to riscv From https://github.com/RuyiAI-Stack/pytorch * branch riscv -> FETCH_HEAD fatal: Not a valid object name origin/main Error: ❌ Failure - Main Diff base and head Error: exit status 128

move the patch to dest

72843ab

Fix block list format and remove test_cpu_select_algorithm (pytorch#4)

7338bc0

* mklnn is unavailable on RISC-V * Remove test_cpu_select_algorithm from block_list * Fix block list format

move the patch to dest

7444a79

[blacklist]: update it (pytorch#11)

c8cb84b

These cases are too slow on riscv64, adding them to here simply Drop test_torch from the list because it is one core case

Use commit sha on PR workflow (pytorch#12)

abf25cc

Fix test_float_to_int_conversion_nonfinite for RISC-V

782e92a

RISC-V converts non-finite floats to integers by saturating: -inf -> min, inf/nan -> max for wider int types. Add IS_RISCV64 flag and RISC-V-specific reference values. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Skip test if no qengine

12e0d70

Replace offical cpuinfo repo (pytorch#18)

01db43b

[ci] split core and full ci (pytorch#20)

ddad2c3

[ci] fix ci workflow and other issue (pytorch#21)

54a71e9

for action, these yaml must be merged first then take effect, so merge it skipping ci

[ci]: continue parsing Jenkins outputs after Jenkins job failure (pyt…

2653522

…orch#40)

[CI] workaround to judge Jenkins job success (pytorch#41)

1571898

[SyncBots] Integrate PyTorch at 1154a55

5b326fe

[SyncBots] Integrate PyTorch at 2760264

55d456e

[test]Scope linear cross entropy ULP tolerance to RISC-V (pytorch#42)

6e72cda

RISC-V Vector Extension (RVV) support for ATen

91cc85c

Signed-off-by: Fei Zhang <zhangfei@iscas.ac.cn> Co-authored-by: Chung-Lin Tang <cltang@baylibre.com>

zhangfeiv0 had a problem deploying to riscv64 August 31, 2026 06:19 — with GitHub Actions Failure

zhangfeiv0 had a problem deploying to riscv64 August 31, 2026 07:24 — with GitHub Actions Failure

zhangfeiv0 temporarily deployed to riscv64 August 31, 2026 10:41 — with GitHub Actions Inactive

WuXintong123 force-pushed the riscv branch 2 times, most recently from e4b937e to ae1e2fb Compare September 2, 2026 01:46

WuXintong123 force-pushed the riscv branch from ae1e2fb to bf72fb4 Compare September 3, 2026 01:04

WuXintong123 force-pushed the riscv branch from bf72fb4 to 313ad9a Compare September 4, 2026 00:26

WuXintong123 force-pushed the riscv branch 3 times, most recently from cfb52d8 to b8e4232 Compare September 7, 2026 00:10

WuXintong123 force-pushed the riscv branch 2 times, most recently from 0ca63b6 to 2349b8f Compare September 9, 2026 00:50

WuXintong123 force-pushed the riscv branch from 2349b8f to 9560b3f Compare September 10, 2026 00:54

WuXintong123 force-pushed the riscv branch from 9560b3f to b6464b9 Compare September 11, 2026 00:55

cpu: rv64: add USE_CPU_VECTORIZATION for default build

9470a8a

zhangfeiv0 had a problem deploying to riscv64 September 11, 2026 05:53 — with GitHub Actions Failure

WuXintong123 force-pushed the riscv branch from b6464b9 to 976d5b9 Compare September 12, 2026 00:54

Navigation Menu

Uh oh!

RISC-V Vector Extension (RVV) support for ATen - #46

RISC-V Vector Extension (RVV) support for ATen#46
zhangfeiv0 wants to merge 22 commits into
RuyiAI-Stack:riscvfrom
zhangfeiv0:riscv

zhangfeiv0 commented Aug 31, 2026

Uh oh!

zhangfeiv0 commented Sep 2, 2026

Uh oh!

yuzibo commented Sep 3, 2026

Uh oh!

XYenChi commented Sep 3, 2026

Uh oh!

zhangfeiv0 commented Sep 4, 2026

Uh oh!

XYenChi commented Sep 7, 2026

Uh oh!

zhangfeiv0 commented Sep 9, 2026

Uh oh!

XYenChi commented Sep 10, 2026

Uh oh!

zhangfeiv0 commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

zhangfeiv0 commented Aug 31, 2026

变更说明 / Summary

关联 Issue / Related issue

验证方式 / Validation

检查清单 / Checklist

Uh oh!

zhangfeiv0 commented Sep 2, 2026

Uh oh!

yuzibo commented Sep 3, 2026

Uh oh!

XYenChi commented Sep 3, 2026

Uh oh!

zhangfeiv0 commented Sep 4, 2026

Uh oh!

XYenChi commented Sep 7, 2026

Uh oh!

zhangfeiv0 commented Sep 9, 2026

测试内容覆盖与总体对比

逐元素操作性能对比

模型与训练

Uh oh!

XYenChi commented Sep 10, 2026

Uh oh!

zhangfeiv0 commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants