[RFC] Lowering `aten.index_add_` through the TM Tensor Dialect
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 1.9k
- Forks
- 736
- Avg merge
- 5d 22h
- Merged PRs (30d)
- 15
Description
Hey @ramiro050 and @cathyzhyi,
Trying to move `aten.index_add_` so that it goes through the TM Tensor Dialect
Inspired by [this PR lowering aten.bincount](https://github.com/llvm/torch-mlir/pull/574/commits/f818ce22f3999e28b05455b188aaa24dc9ee1798), here's my current plan:
1. Add `AtenIndexAdd` to `e2e_testing/torchscript/basic.py`
2. Define `Toch_AtenIndexAddOp` in `include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td`
3. Create a `AtenIndexAddOp` class in `lib/Conversion/TorchToTMTensor/TorchToTMTensor.cpp`
-- I'm also going to the op under `runOnOperation` within the `ConvertTorchToTmTensor` class
4. Add the op to the conditional statement in `lib/Dialect/Torch/Transforms/RefineTypes.cpp`
5. Emit the op under `python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py`
Would love any feedback! Am I thinking about the sequence correctly?
Thank you in advance,
Lorenzo
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with e2e_testing/torchscript/basic.py and the referenced aten.bincount lowering, then inspect the listed Torch dialect, TorchToTMTensor, RefineTypes.cpp, and torch_ods_gen.py locations. Done means aten.index_add_ is represented and lowered through the TM Tensor Dialect, with the end-to-end test entry added.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 28/100