Torch ops E2E implementation

Example PR

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

Major steps

Step 1. 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 that file.

Step 2. Propagate shapes with RefineTypes pass

The RefineTypes pass propagates refined tensor types across the entire program. Each visit function infers the output tensor shape and element type 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 3. 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. 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.

Tutorials

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