[Performance] GSpMM regression with FP16
- Dominant language
- Python
- Stars
- 14.3k
- Forks
- 3.1k
- PR merge metrics
- No merged PRs in 30d
Description
I spent some time profiling the GAT example with AMP in https://docs.dgl.ai/en/0.9.x/guide/mixed_precision.html and want to know why we didn't obtain performance gain from FP16. I observed regression in both the forward and backward phases.
## Performance on A100
| | AMP | FP32 |
|--------------|--------|--------|
| forward | 254 ms | 160 ms |
| backward | 482 ms | 369 ms |
| step | 738 ms | 529 ms |
Take the forward phase for example, there are three `EdgeSoftmax-GSpMM` pairs.
- AMP

- FP32

Elapsed times (ms) of the `EdgeSoftmax-GSpMM` pairs x3.
| | AMP | FP32 |
|--------------|--------|--------|
| EdgeSoftmax1 | 20.896 | 22.055 |
| GSpMM1 | 73.851 | 31.890 |
| EdgeSoftmax2 | 20.865 | 22.080 |
| GSpMM2 | 73.797 | 31.882 |
| EdgeSoftmax3 | 20.869 | 22.154 |
| GSpMM3 | 27.528 | 13.422 |
We can see that the `EdgeSoftmax` is a bit faster with AMP while it's `GSpMM` that results in the regression.
Dig it a bit deeper, `GSpMM` with FP32 will go through the `cusparse::csrmm` code path while adopting `dgl::aten::cuda::SpMMCsrKernel` with FP16 inputs. I'll look into the performance issue on cusparse with FP16.
### ENV
- CUDA_VERSION: 11.7.0.022
- CUDA_DRIVER_VERSION: 515.43.04
- CUSPARSE_VERSION: 11.7.3.50
- GPU: NVIDIA A100 80GB PCIe
- PyTorch: 1.12.0a0+8a1a93a
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.