Torch ops E2E implementation

Tutorials

• Linalg.generic op introduction here
• Basic E2E debugging walk-through here

Example PR

https://github.com/llvm/torch-mlir/pull/294

Major steps

Step 1. Add an end-to-end test to iterate on

Add an end-to-end test to the end-to-end test suite). Ideally there is an existing file that your op fits into. If not, you can create a new file.

We generally recommend testing by invoking torch.ops.aten.someop from Python -- that gives a very precise test for the individual Torch operator you are implementing (calling torch.ops.aten.someop from Python always lowers into the MLIR torch.aten.someop operation)

The end-to-end test is important to check the correctness of the other steps.

Step 2. Update ods

Update torch_ods_gen.py with the new op and run update_torch_ods.sh to generate the ods. Running update_torch_ods.sh would dump all the operators with schema into JITOperatorRegistryDump.txt. It’s convenient to look for ops signatures and operands names in this file.

Step 3. Propagate types

It’s essential to make sure the new op implements shape and dtype inference. See abstract_interp_lib for information on adding shape and dtype inference.

Step 4. Torch ops lowering

Decompose

If your op can be decomposed into other supported ops, then you can add a pattern into DecomposeComplexOps.

You can find example PRs here and here.

Lower to Linalg

The Torch dialect needs to be lowered to Linalg dialect which can be used as input IR of backends. Here is a high level introduction about Linalg ops and here is a video explaining linalg.generic op. The building block is the linalg.generic op which consists of indexing maps, iterator types, input/output tensors and a compute payload. You would want to get familiar with the concept of affine map. The linalg.generic op anatomy tutorial covers the basics of linalg.generic from a user's perspective.

You can find an example PR here.

Delivering Code

1. The codebase follows the LLVM’s coding conventions.The following items might be the most frequently used rules:

• use-early-exits-and-continue-to-simplify-code
• don-t-use-else-after-a-return
• use-auto-type-deduction-to-make-code-more-readable
• anonymous-namespaces
• avoid-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

2. Try to refactor and reuse existing code/helpers when working on RefineTypes and TorchToLinalg lowering for easier maintenance, testing and better readability. Try not to copy & paste existing code.
3. Squash all the commits into one, including the commits addressing review comments.
4. Use git clang-format HEAD~1 to automatically format your commit.
5. Rebase on HEAD before delivering.