apache / apache/tvm

[Bug] [RISC-V RVV] softmax operator shows suboptimal vectorization

Open
#18,569 0 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

### Issue: [RISC-V RVV] softmax operator shows suboptimal vectorization

#### Description
The softmax operator performs worse with the RISC‑V Vector (RVV) extension, achieving only 0.745× the performance of the scalar implementation. This suggests inefficient vectorization for softmax operations.

#### Steps to Reproduce
1. Generate the softmax operator with the following configuration:
```python
params = {
"dtype": "float32",
"batch": 14,
"features": 185
}
```

2. Export the operator to two targets:
- **RV target** (scalar, without vector extension):
```
llvm -mtriple=riscv64-linux-gnu -mcpu=generic-rv64 -mabi=lp64d -mattr=+64bit,+m,+a,+f,+d,+c
```
- **RVV target** (with vector extension):
```
llvm -mtriple=riscv64-linux-gnu -mcpu=generic-rv64 -mabi=lp64d -mattr=+64bit,+m,+a,+f,+d,+c,+v
```

3. Run performance measurement on both targets.

Operator definition code:
```python
def export_softmax(params, set_dir=None, platform="rv"):
data = relay.var("data", shape=(params["batch"], params["features"]),
dtype=params["dtype"])
softmax = relay.nn.softmax(data)
export_op(softmax, params["op_name"], [data], params, set_dir=set_dir)
```

#### Performance Data
- **RV execution time**: 1.831500 ms
- **RVV execution time**: 2.457040 ms
- **Acceleration ratio (RV/RVV)**: 0.745 (RVV is ~1.34× slower)

#### Environment Information
- **TVM version**: 0.19.0
- **LLVM version**: [Please provide: `llvm-config --version`]
- **Hardware**: Spacemit K1‑X bit‑brick board
- **CPU**: Spacemit X60 (8 cores, 1.6 GHz)
- **ISA**: rv64imafdcv (with vector extensions)
- **Memory**: 7.6 GB
- **OS**: Bianbu 2.2, Linux kernel 6.6.63
- **Operation**: Softmax on a 2D tensor of shape (14, 185)

#### Expected Behavior
RVV vectorization should provide a performance improvement over the scalar RV baseline for softmax operations, which involve reduction and elementwise operations that can be vectorized.

#### Additional Context
- The softmax operation is applied to a 2D tensor of shape (14, 185), which is a relatively small tensor compared to other operators tested.
- The performance regression, though less severe than some other operators, still indicates that the vectorized implementation of softmax may have inefficiencies in the reduction and exponentiation steps.
- This is part of a broader pattern where multiple operators show performance degradation with RVV, suggesting potential issues with vectorization strategies for reduction and elementwise operations.

Contributor guide

No contributing guide indexed for this repository

Research direction

Reproduce the softmax export using the stated float32, batch-14, feature-185 parameters and compare the scalar RV and RVV timings. No source file or test is named; completion would require identifying the RVV vectorization bottleneck and demonstrating improved performance against the reported baseline.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.