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 dtypes with RefineTypes pass

The RefineTypes pass propagates refined tensor dtypes across the entire program. Each visit function infers the output tensor dtype based on the input. It’s necessary to make sure the new op is handled correctly by this pass. If existing helpers can’t be reused and new code logic is added, unit tests like those in test/Dialect/Torch/refine-types.mlir are needed. The unit tests use LLVM’s FileCheck and MLIR provides a script mlir/utils/generate-test-checks.py to generate FileCheck statements.

Step 4. Add a shape function to the shape library

See the documentation in Adding a Shape Function.

Step 5. Torch ops lowering

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.

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.