Add decorator for custom op and inductor decomp registration #408

jerryzh168 · 2024-06-20T22:33:26Z

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.

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:

pytorch-bot · 2024-06-20T22:33:28Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/408

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit ac2e283 with merge base bc8599f ():

NEW FAILURES - The following jobs have failed:

Run Regression Tests / test (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/whl/nightl... / linux-job (gh)
RuntimeError: Command docker exec -t 13416d5670052772475f929ce2f83b49c411cf8b715456ae86eeeb5e4cf54f76 /exec failed with exit code 2
Run Regression Tests / test (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://download.pytorc... / linux-job (gh)
RuntimeError: Command docker exec -t d68d9c69365ae5f657f9d9d41320c95d5a7b597d7c1b5de7359e3ed4cf37363c /exec failed with exit code 2

This comment was automatically generated by Dr. CI and updates every 15 minutes.

supriyar · 2024-06-21T17:38:31Z

torchao/quantization/quant_primitives.py

@@ -151,12 +144,12 @@ def quantize_affine(
 output_dtype (torch.dtype): requested dtype (e.g. torch.uint8) for output Tensor
 quant_min (Optional[int]): minimum quantized value for output Tensor, if not specified, it will be derived from dtype
 quant_max (Optional[int]): maximum quantized value for output Tensor, if not specified, it will be derived from dtype
- zero_point_domain (ZeroPointDomain): the domain that zero_point is in, should be eitehr integer or float
+ zero_point_domain (str): the domain that zero_point is in, should be eitehr "int" for "float"


why did you change these from Enum to str?

oh it's because custom op API doesn't support enum right now

@zou3519 any plans to add support for this in the API?

No, the set of types supported by operators is intentionally limited so that people don't have a difficult time working with it (for example, writing graph passes on FX graphs with operators). You can wrap the call to the operator in a python function that does support enums

You can wrap the call to the operator in a python function that does support enums

oh, I see, maybe that's what we should be doing here

supriyar · 2024-06-21T17:39:17Z

torchao/quantization/quant_primitives.py

@@ -205,15 +203,16 @@ def quantize_affine(

 return quant

+@register_custom_op("quant::dequantize_affine")


does this make the default behavior to preserve the higher level op when we run export? What about compile?

yes, and in compile we'll register a inductor decomposition to decompose the op

jerryzh168 · 2024-06-21T18:26:11Z

torchao/quantization/quant_primitives.py

+ if TORCH_VERSION_AFTER_2_5:
+ opdef = torch.library.custom_op(name, mutates_args=())(fn)
+ opdef.register_fake(fn)
+ register_decomposition([opdef._opoverload])(fn)


@supriyar here we register a decomp for inductor, so in torch.compile these ops will still be decomposed

jerryzh168 · 2024-06-21T18:27:10Z

Requires pytorch/pytorch#129179 and pytorch/pytorch#129189 to be landed, and the nightly version got updated until the CI can pass

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. 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 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 · 2024-06-25T20:35:11Z

recreated the PR in #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 #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 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:

…#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:

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 20, 2024

jerryzh168 force-pushed the executorch-ir branch 3 times, most recently from cf3234c to 2ba121f Compare June 21, 2024 16:57

supriyar reviewed Jun 21, 2024

View reviewed changes

jerryzh168 commented Jun 21, 2024

View reviewed changes

jerryzh168 force-pushed the executorch-ir branch from 2ba121f to ac2e283 Compare June 21, 2024 22:25

jerryzh168 mentioned this pull request Jun 25, 2024

Add decorator for custom op and inductor decomp registration #434

Merged

jerryzh168 closed this Jun 25, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add decorator for custom op and inductor decomp registration #408

Add decorator for custom op and inductor decomp registration #408

jerryzh168 commented Jun 20, 2024

pytorch-bot bot commented Jun 20, 2024 •

edited

Loading

supriyar Jun 21, 2024

jerryzh168 Jun 21, 2024

supriyar Jun 21, 2024

zou3519 Jun 23, 2024

jerryzh168 Jun 24, 2024

supriyar Jun 21, 2024

jerryzh168 Jun 21, 2024

jerryzh168 Jun 21, 2024

jerryzh168 commented Jun 21, 2024

jerryzh168 commented Jun 25, 2024

		@@ -205,15 +203,16 @@ def quantize_affine(

		return quant

		@register_custom_op("quant::dequantize_affine")

Add decorator for custom op and inductor decomp registration #408

Add decorator for custom op and inductor decomp registration #408

Conversation

jerryzh168 commented Jun 20, 2024

pytorch-bot bot commented Jun 20, 2024 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/408

❌ 2 New Failures

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jerryzh168 commented Jun 21, 2024

jerryzh168 commented Jun 25, 2024

pytorch-bot bot commented Jun 20, 2024 •

edited

Loading