Repository navigation
Conversation
The Arm quantizer never annotated aten.argmax, so an argmax fed by a float tensor or by a dequantized producer stayed in fp32 and the INT profile partitioner rejected it with "TOSA ARGMAX requires PRO-FP for torch.float32 input". In a logits -> argmax -> embedding chain this left the argmax on the CPU between two NPU delegates. Annotate only the argmax input, since the TOSA ARGMAX output is int32 and carries no quantization, and let the argmax support check look through a leading dequantize node so it judges the dtype the folded graph will present. A Linear -> argmax -> Embedding model now lowers to a single Ethos-U85 command stream: conv, ARGMAX_X pool, gather DMA. Tests cover float input on the TOSA INT pipeline and on the Ethos-U85 flow up to the Vela compile. Authored with assistance from Claude Code (Claude Fable 5.1). Signed-off-by: Paul Wilk <paul@randomproc.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24394
Note: Links to docs will display an error until the docs builds have been completed.
|
|
Hi @pcwilk! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
One or more co-authors of this pull request were not found. You must specify co-authors in commit message trailer via: Supported
Alternatively, if the co-author should not be included, remove the Please update your commit message(s) by doing |
|
|
|
@pytorchbot label "release notes: arm" |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
meta-cla
Bot
added
the
CLA Signed
executorch-triage
Bot
added
the
community: contribution
Summary
The Arm quantizer never annotated
aten.argmax, so an argmax fed by a float tensor, or by a dequantized producer such as a quantized linear head, stayed in fp32 and the INT-profile partitioner rejected it with "TOSA ARGMAX requires PRO-FP for torch.float32 input". In a logits → argmax → embedding chain this left the argmax on the CPU between two NPU delegates.This change annotates only the argmax input (the TOSA ARGMAX output is int32 and carries no quantization) and lets the argmax support check look through a leading dequantize node, so it judges the dtype the folded graph will actually present. A
Linear(64, 4096)→argmax→Embedding(4096, 64)model now lowers to a single Ethos-U85 command stream:NPU_OP_CONV, anARGMAX_Xpool over the 4096 logits, and a gather DMA, with zero CPU operators reported by Vela.Suggested label:
Release notes: arm.Test plan
Two tests are added to
backends/arm/test/ops/test_argmax.py: float input on the TOSA INT pipeline (asserting exactly oneARGMAXop), and the same through the Ethos-U85 flow up to the Vela compile.On an editable build of main at 05092c1 (torch 2.14.0+cpu, Vela 5.2.0): lint clean; 24 passed, 12 skipped (no model-converter for VGF), 3 xfailed (no Corstone-300 runner), 0 failed. The U85 subset was additionally run on the Corstone-320 FVP with the semihosting runner: 9 passed.
🤖 Generated with Claude Code
cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani @nil-is-all