Skip to content

[TORCH] Add support for aten.absolute op via torch decomposition#4624

Open
AiyyappanMR wants to merge 2 commits into
llvm:mainfrom
AiyyappanMR:aten-absolute-decomp
Open

[TORCH] Add support for aten.absolute op via torch decomposition#4624
AiyyappanMR wants to merge 2 commits into
llvm:mainfrom
AiyyappanMR:aten-absolute-decomp

Conversation

@AiyyappanMR

Copy link
Copy Markdown

torch.absolute is an alias for torch.abs but is not registered in the torch dialect. It arrives as an unregistered torch.operator "torch.aten.absolute" which fails legalization when lowering to backend contract.

On testing via fx_importer:
./projects/pt1/tools/e2e_test.sh --config=fx_importer -f ElementwiseAbsolute

TorchFX IR:
%0 = "torch.operator"(%arg0) <{name = "torch.aten.absolute"}> : (!torch.vtensor<[1,1,3],f32>) -> !torch.vtensor<[1,1,3],f32>

error: failed to legalize operation 'torch.operator' that was explicitly marked illegal

This PR adds support by:

  • Registering aten::absolute in torch_ods_gen.py and regenerating GeneratedTorchOps.td
  • Adding shape/dtype inference in abstract_interp_lib_gen.py and regenerating AbstractInterpLibrary.cpp
  • Adding DecomposeAtenAbsoluteOp in DecomposeComplexOps.cpp which decomposes to AtenAbsOp
  • Adding e2e tests for float and int variants

@AiyyappanMR

Copy link
Copy Markdown
Author

Hi @zjgarvey and @penguin-wwy, I'm a first-time contributor. Could either of you trigger the CI pipeline for this PR, or point me to someone who can? Also happy to get a code review whenever someone has time. Thanks!

@penguin-wwy

Copy link
Copy Markdown
Collaborator

please resolve conflicts

@AiyyappanMR

Copy link
Copy Markdown
Author

Hi @penguin-wwy, I've resolved the merge conflicts and pushed the update. Please let me know if anything else is needed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants