alibaba / alibaba/BladeDISC

Initial quantization support in BladeDISC

Open
#567 0 comments 5 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
933
Forks
169
PR merge metrics
No merged PRs in 30d

Description

We're now planing to start support quantization in BladeDISC. The basic idea is:
- introduce a fake_quant op. The op will be used to pass quantization related info (e.g. scale/zero_point) from frontend to BladeDISC.
- The input of BladeDISC is still a FP32 graph, with some fake_quant ops to annotate which set of ops need to be quantized.
- It's the responsibility of BladeDISC to replace the fake_quant op and the corresponding FP32 ops (e.g. Gemm/Conv) to the quantized versions.

The task list is shown as follow:
- [x] TensorFlow/PyTorch fake_quant custom ops and their corresponding implementation in TF/PyTorch Dialect.
- [x] mhlo_disc fake_quantize/quantize/de-quantize IR definition.
- [x] mhlo_disc QCONV/QMATMUL IR defintion.
- [x] TF custom fake-quant -> disc fake quant
- [ ] Torch custom fake-quant -> disc fake quant
- [x] Fake_quantize + Conv + Fake_quantize -> QConv pattern.
- [x] Fake_quantize + GEMM + Fake_quantize -> QGEMM pattern.
- [x] Fake_quantize lowering.
- [x] Quantize op lowering patterns.
- [x] de-quantize op lowering patterns.
- [x] Placement rule for Qconv/Qmatmul.
- [ ] (low priority) QConv/QMatmul layout optimization.
- [x] QConv shape inference implementation.
- [x] QMatmul shape inference implementation.
- [x] QConv lowering to library call implementation.
- [x] QMatmul lowering to library call implementation.
- [x] QConv kernel and e2e UT for arm.
- [x] QMatMul kernel and e2e UT for arm.
- [x] QConv kernel and e2e UT for cuda.
- [x] QMatMul kernel and e2e UT for cuda.
- [ ] QConv kernel and e2e UT for x86.
- [ ] QMatMul kernel and e2e UT for x86.
- [ ] e2e demo model for ARM
- [ ] e2e demo model for X86
- [ ] e2e demo model for CUDA

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reviewing the unchecked task list, especially Torch custom fake-quant to Disc fake quant, the x86 QConv and QMatMul kernels with end-to-end tests, and the ARM, x86, and CUDA demo models. Done means completing a narrowly selected item with its corresponding implementation and validation.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, pytorch, tensorflow
Domain
compilers, machine-learning, testing-qa
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
20/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.