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 build_tools/update_torch_ods.sh to generate the ods. Running build_tools/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 Lowing torch dialect to Linalg dialect. Here is a high level introduction about Linalg ops. The most commonly used op is linalg.generic. Constructing a linalg.generic op requires indexing maps, iterator types, input/output tensors and a compute payload. Among those components, the indexing map is probably the hardest to grasp. It is a list of affine maps defining how each operand tensor is iterated. The Linalg.generic op anatomy tutorial covers the basics of this concept.

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.