Add decorator for custom op and inductor decomp registration #434
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/434
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 320b846 with merge base c2cf973 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
jerryzh168
requested review from
supriyar,
msaroufim,
zou3519,
HDCharles and
kimishpatel
jerryzh168
force-pushed
the
executorch-ir2
branch
from
5186214
to
fd6dbc7
Compare
jerryzh168
force-pushed
the
executorch-ir2
branch
from
fd6dbc7
to
4383e08
Compare
| quant_min: Optional[int] = None, | ||
| quant_max: Optional[int] = None, | ||
| zero_point_domain: str = "INT", | ||
| *, |
There was a problem hiding this comment.
Remove the *, it's not in the schema
| if TORCH_VERSION_AFTER_2_5: | ||
| # TODO: change order | ||
| lib_namespace = lib.ns | ||
| op_name = schema.split("(")[0] |
There was a problem hiding this comment.
Maybe construct schema object from string and query op name? I thought such a functionality existed, but not sure
There was a problem hiding this comment.
oh, not sure if this is possible, cc @zou3519 is there a better way to get op_name here?
There was a problem hiding this comment.
torch._C.parse_schema will give you a FunctionSchema object
There was a problem hiding this comment.
I just used fn.__name__[1:] for now
| after_export = model(x) | ||
| self.assertTrue(torch.equal(after_export, ref)) | ||
| if api is _int8da_int8w_api: |
There was a problem hiding this comment.
this is because right now we will only see these ops for int8da_int8w quantization, other types of quant (e.g. int4 weight only) will call into the efficient kernels directly
we should probably figure out a path for executorch, I think we could abstract this with "layout", what would be a good name here?
jerryzh168
force-pushed
the
executorch-ir2
branch
4 times, most recently
from
f2ca7f2
to
3c45fdf
Compare
jerryzh168
force-pushed
the
executorch-ir2
branch
2 times, most recently
from
604f69c
to
cd8d0c1
Compare
| # expecting fn.__name__ starts with `_` and we want to take the rest | ||
| # to be the name of the custom op | ||
| assert fn.__name__[0] == "_", f"Expecting function name starts with `_`, got {fn.__name__}" | ||
| op_name = fn.__name__[1:] |
There was a problem hiding this comment.
Can you assert there is no "." or "<" or ">" in fn.name? this can happen with lambdas or local functions
Summary: This PR adds a decorator to register custom op and also an inductor dcomposition. The goal is for torch.export path to be able to see high level ops like quantize_affine instead of breaking down the op, this is because some backends like xnnpack wants to work with these higher level ops. This is a redo for pytorch#408, difference is we can preserve the enums on the python side in this PR Test Plan: regression tests: python test/quantization/test_quant_api.py python test/integration/test_integration.py also need to check performance with python tutorials/quantize_vit/run_vit_b_quant.py Reviewers: Subscribers: Tasks: Tags:
jerryzh168
force-pushed
the
executorch-ir2
branch
from
cd8d0c1
to
320b846
Compare
…#434) Summary: This PR adds a decorator to register custom op and also an inductor dcomposition. The goal is for torch.export path to be able to see high level ops like quantize_affine instead of breaking down the op, this is because some backends like xnnpack wants to work with these higher level ops. This is a redo for pytorch#408, difference is we can preserve the enums on the python side in this PR Test Plan: regression tests: python test/quantization/test_quant_api.py python test/integration/test_integration.py also need to check performance with python tutorials/quantize_vit/run_vit_b_quant.py Reviewers: Subscribers: Tasks: Tags:
Summary:
This PR adds a decorator to register custom op and also an inductor dcomposition.
The goal is for torch.export path to be able to see high level ops like quantize_affine instead of breaking down the op, this is because some backends like xnnpack wants to work with these higher level ops.
This is a redo for #408, difference is we can preserve the enums on the python side in this PR
Test Plan:
regression tests:
python test/quantization/test_quant_api.py
python test/integration/test_integration.py
also need to check performance with python tutorials/quantize_vit/run_vit_b_quant.py
Reviewers:
Subscribers:
Tasks:
Tags: