llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update E2E links

pull/2566/head snapshot-20231110.1018
James Newling 2023-11-09 18:28:05 +00:00 committed by Ramiro Leal-Cavazos
parent d082310bd8
commit 98ee7fe548
1 changed files with 9 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
docs/Torch-ops-E2E-implementation.md
View File

@ -9,29 +9,27 @@ https://github.com/llvm/torch-mlir/pull/294
### Step 1. Add an end-to-end test to iterate on ### Step 1. Add an end-to-end test to iterate on
Add an end-to-end test to the [end-to-end test suite](https://github.com/llvm/torch-mlir/tree/main/python/torch_mlir_e2e_test/test_suite). Ideally there is an existing file that your op fits into. If not, you can create a new file. Add an end-to-end test to the [end-to-end test suite](https://github.com/llvm/torch-mlir/blob/main/docs/adding_an_e2e_test.md)). 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) 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. The end-to-end test is important to check the correctness of the other steps.
### Step 2. Update ods ### Step 2. Update ods
Update [torch_ods_gen.py](https://github.com/llvm/torch-mlir/blob/main/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py) with the new op and run [update_torch_ods.sh](https://github.com/llvm/torch-mlir/blob/main/build_tools/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 Update [torch_ods_gen.py](https://github.com/llvm/torch-mlir/blob/main/projects/pt1/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py) with the new op and run [update_torch_ods.sh](https://github.com/llvm/torch-mlir/blob/main/build_tools/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.
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](https://github.com/llvm/torch-mlir/blob/main/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](https://github.com/llvm/llvm-project/blob/main/mlir/utils/generate-test-checks.py) to generate [FileCheck](https://llvm.org/docs/CommandGuide/FileCheck.html) statements.
### Step 4. Add a shape function to the shape library ### Step 3. Propagate types
It’s essential to make sure the new op implements shape and dtype inference. See [abstract_interp_lib](https://github.com/llvm/torch-mlir/blob/main/docs/abstract_interp_lib.md) for information on adding shape and dtype inference.
See the documentation in [Adding a Shape Function](https://github.com/llvm/torch-mlir/blob/main/docs/adding_a_shape_function.md). ### Step 4. Torch ops lowering
### Step 5. Torch ops lowering
#### Decompose #### Decompose
If your op can be decomposed into other supported ops, then you can add a pattern into [DecomposeComplexOps](https://github.com/llvm/torch-mlir/blob/8d3ca887df5ac5126fa3fc2ec3546c6322a4d066/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp#L1). If your op can be decomposed into other supported ops, then you can add a pattern into [DecomposeComplexOps](https://github.com/llvm/torch-mlir/blob/8d3ca887df5ac5126fa3fc2ec3546c6322a4d066/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp#L1).
You can find an [example PR here](https://github.com/llvm/torch-mlir/pull/1453). You can find example PRs [here](https://github.com/llvm/torch-mlir/pull/2550) and [here](https://github.com/llvm/torch-mlir/pull/2553).
#### Lower to Linalg #### Lower to Linalg