Repository navigation
Add RVV vectorized FP32 kernel to CPU Adam - #8817
zhangjian29 wants to merge 3 commits into
Conversation
DeepSpeed CPU Adam falls back to the scalar Step_1 tail on riscv64 because neither the builder's simd_width() nor the kernel dispatch recognizes the RISC-V vector extension. Add the FP32 RVV tier, mirroring the SVE backend: - csrc/includes/rvv.h: FP32 Adam kernel using RVV 1.0 intrinsics. The hot loop processes one m4 block per iteration, so at VLEN=128 each block covers a full 64-byte cache line per stream and the loop runs entirely at the loop-invariant vlmax. A single vsetvl truncates the trailing block, so any tensor length is covered with no scalar tail. The operation order matches the scalar and SVE kernels, keeping all tiers within the existing test tolerance. Scalar hyperparameters ride the vf-form intrinsics, so of all the loop constants only epsilon occupies a vector register. - cpu_adam.h / cpu_adam_impl.cpp: Step_RVV dispatch alongside Step_AVX and Step_SVE; only the FP32 param/FP32 state pair is vectorized, other dtype pairs keep the scalar path (same policy as SVE). - op_builder/builder.py: detect the V extension from the /proc/cpuinfo isa line with a per-character token test (a substring test misreads rv64imafdcvh as having V) and build with -D__RVV__ -march=rv64gcv on riscv64 (rv64gc without V, so non-vector hardware is unaffected). Also normalize py-cpuinfo's "RISCV_64" arch string, without which the riscv64 branch never fires and the build silently falls back to -march=native. - docs/_tutorials/zero-offload.md: document the automatic RVV tier. Measured on a SpacemiT 60-core server (VLEN=128, gcc-14), single core: 1B step 3.15s, 2B step 6.31s (scalar baseline 13.07s/26.13s, ~4.2x). pytest unit/ops/adam/test_cpu_adam.py: 48 passed / 1 skipped / 0 failed. Signed-off-by: 张健10355098 <zhang.jian60@zte.com.cn>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8b7461f168
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
Addressed the P1 on The base ISA segment is now tokenized into per-extension tokens ( |
An ISA string like rv64imafdcv0p7 (draft V 0.7) passed the single-letter set test, so cpu_arch() built with -march=rv64gcv and simd_width() set __RVV__, emitting ratified V 1.0 encodings such a CPU cannot execute. Tokenize the base segment into per-extension version tokens and only enable the vector path for an unversioned V or an explicit version >= 1.0, keeping the scalar rv64gc build otherwise. Signed-off-by: 张健10355098 <zhang.jian60@zte.com.cn>
3d382b6 to
2f8a2af
Compare
There was a problem hiding this comment.
🟡 Changes recommended
Unversioned draft RVV hardware can be misdetected as RVV 1.0, producing incompatible instructions.
1 open finding
What changed in this PR
Adds RVV 1.0 FP32 acceleration to CPU Adam on supported RISC-V systems.
Changes:
- Implements and dispatches an RVV-vectorized Adam kernel.
- Adds RISC-V vector detection and tests.
- Documents automatic RVV optimization.
| File | Description |
|---|---|
csrc/includes/rvv.h |
Implements the RVV FP32 Adam kernel. |
csrc/includes/cpu_adam.h |
Integrates the RVV update path. |
csrc/adam/cpu_adam_impl.cpp |
Adds RVV dispatch and capability reporting. |
op_builder/builder.py |
Detects RISC-V V and configures compilation. |
tests/unit/ops/test_op_builder.py |
Tests ISA-string parsing. |
docs/_tutorials/zero-offload.md |
Documents RVV behavior. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
Follow-up to the P1 on the bare The probe now reads AT_HWCAP from |
An unversioned "v" in the ISA line is not evidence of ratified V 1.0: T-Head C906 (Allwinner D1) systems report bare rv64imafdcv while implementing incompatible draft V 0.7.1, so the string parse alone can still select -march=rv64gcv and emit encodings the CPU cannot execute. Read AT_HWCAP from /proc/self/auxv (the interface the kernel's RISC-V vector documentation defines for this question) and treat its COMPAT_HWCAP_ISA_V bit as the decision signal; keep the explicitly versioned ISA token as an authoritative override (v1p0 yes, v0p7 never) so pre-HWCAP kernels with precise devicetrees still work, and let a bare v without HWCAP fall back to the scalar rv64gc build. Signed-off-by: 张健10355098 <zhang.jian60@zte.com.cn>
43b1f40 to
976ff52
Compare
|
@zhangjian29 thanks for your contribution. I think the RISCV specific test should be skipped on non-riscv platform rather than mock RISCV behaviors. They could run on RISC platform either manually or on a RISCV runner if that is part of the plan. Thanks! |
Description
On riscv64 hosts with the RISC-V Vector (RVV 1.0) extension, DeepSpeed CPU Adam currently falls back to the scalar
Step_1tail because neither the op builder'ssimd_width()nor the kernel dispatch recognizes the V extension. This PR adds an RVV FP32 tier to CPU Adam, mirroring the existing SVE backend, so RVV-capable CPUs get a ~4x single-core speedup on the optimizer step.csrc/includes/rvv.h: an RVV 1.0 FP32 Adam kernel. The hot loop processes onem4block per iteration — at VLEN=128 each block covers a full 64-byte cache line per stream — and runs entirely at the loop-invariant vlmax; a singlevsetvltruncates the trailing block, so any tensor length is covered with no scalar tail. The arithmetic order matches the scalar and SVE kernels, keeping every tier within the existing test tolerance. Scalar hyperparameters ride the vf-form intrinsics, so only epsilon occupies a vector register.cpu_adam.h/cpu_adam_impl.cpp:Step_RVVdispatch alongsideStep_AVXandStep_SVE; only the FP32 param/FP32 state pair is vectorized, other dtype pairs keep the scalar path (same policy as SVE).op_builder/builder.py: detect the V extension from the/proc/cpuinfoisa line with a per-character token test (a substring test misreadsrv64imafdcvhas having V), build with-D__RVV__ -march=rv64gcvon riscv64, and normalize py-cpuinfo'sRISCV_64arch string so the riscv64 branch actually fires. Hosts without V are unaffected (rv64gc).docs/_tutorials/zero-offload.md: document the automatic RVV tier.Validation
Hardware/software: SG2044 60-core RISC-V server, VLEN=128, gcc-14.3, torch 2.15 CPU-only build.
Correctness:
pytest tests/unit/ops/adam/test_cpu_adam.py→ 48 passed, 1 skipped, 0 failed. The suite compares CPU Adam against the PyTorch fused-Adam reference across fp32/fp16 and multiple sizes; the fp16 equality tests pass (this tier keeps non-fp32 dtype pairs on the scalar path).Key log:
Performance (single core,
OMP_NUM_THREADS=1 taskset -c 0, 100 steps):Key log: