Skip to content

Commit

Permalink
Refactor implements to support secondary dispatch
Browse files Browse the repository at this point in the history
Summary:
Some popular ops like linear will get a lot of implementations based on
the different characteristics of input and weight, e.g. int8 act + int8 weight, int8 act + int4 weight
etc. For `AffineQuantizedTensor` rigth now all of these implementations are added to the main body of the implementation of linear dispatch, this makes the code hard to read and extend. In this PR we supported
a secondary dispatch condition check for `implements` function:

```
def disptch_condition(func, types, args, kwargs):
    ...

@implements(torch.nn.functional.linear, dispatch_condition)
def _(func, types, args, kwargs):
    # implementation for inputs that passes the dispatch_condition
    ...

@implements(torch.nn.functional.linear)
def _(func, types, args, kwargs):
    ...
```

Test Plan:
regression tests
python test/quantization/test_quant_api.py
python test/integration/test_integration.py

python tutorials/quantize_vit/run_vit_b_quant.py

Reviewers:

Subscribers:

Tasks:

Tags:
  • Loading branch information
jerryzh168 committed Aug 13, 2024
1 parent 433cd14 commit b6787b2
Showing 1 changed file with 245 additions and 192 deletions.
Loading

0 comments on commit b6787b2

Please sign in to comment.