Skip to content

[torch] Add decomposition for aten._foreach_lerp.List#4502

Draft
BruceXinXin wants to merge 1 commit intollvm:mainfrom
BruceXinXin:bruce_add_foreach_lerp_lowering
Draft

[torch] Add decomposition for aten._foreach_lerp.List#4502
BruceXinXin wants to merge 1 commit intollvm:mainfrom
BruceXinXin:bruce_add_foreach_lerp_lowering

Conversation

@BruceXinXin
Copy link

[torch] Add decomposition for aten._foreach_lerp.List

Summary

This PR adds support for decomposing aten._foreach_lerp.List into element-wise aten.lerp.Tensor operations.

aten::_foreach_lerp.List takes three tensor lists (self, tensors1, weights) and computes lerp(self[i], tensors1[i], weights[i]) for each element. The decomposition unpacks the lists via PrimListConstructOp, applies aten.lerp.Tensor to each corresponding triple, and repacks the results.

Changes

  • GeneratedTorchOps.td: Add Torch_Aten_ForeachLerpListOp ODS definition with HasValueSemantics and ReadOnly traits.
  • torch_ods_gen.py: Register aten::_foreach_lerp.List : (Tensor[], Tensor[], Tensor[]) -> (Tensor[]).
  • DecomposeComplexOps.cpp: Add DecomposeAten_ForeachLerpListOp rewrite pattern that decomposes the foreach op into per-element aten.lerp.Tensor calls, with validation that all three input lists have matching sizes.
  • decompose-complex-ops.mlir: Add lit test verifying the decomposition produces the expected sub/mul/add sequence (from lerp.Tensor lowering) and final prim.ListConstruct.

Decompose `aten._foreach_lerp.List` into element-wise `aten.lerp.Tensor`
operations. The decomposition unpacks the three input tensor lists (self,
tensors1, weights), applies `aten.lerp.Tensor` to each corresponding
triple, and repacks the results into a list.

Changes:
- Add ODS definition for `Aten_ForeachLerpListOp`
- Add `torch_ods_gen.py` entry for `aten::_foreach_lerp.List`
- Add `DecomposeAten_ForeachLerpListOp` pattern in DecomposeComplexOps
- Add lit test for the decomposition
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.

1 participant