Skip to content

Commit 18669b3

Browse files
renxidarenxida
authored
Create add_ops.md (#2770)
1 parent 3b85c70 commit 18669b3

File tree

1 file changed

+159
-0
lines changed

1 file changed

+159
-0
lines changed

‎docs/add_ops.md‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
# How to Add Ops to Torch-Mlir
2+
3+
Collected links and contacts for how to add ops to torch-mlir.
4+
5+
6+
<details>
7+
<summary>Turbine Camp: Start Here</summary>
8+
This document was previously known as `turbine-camp.md` to Nod.ai. "Turbine Camp" is part of Nod.ai's onboarding process. Welcome to turbine camp. This document originated at Nod.ai as a part of onboardding process, where new nod-ai folks learn about the architecture of our work by adding support for 2 ops to torch-mlir. I decided to put this into torch mlir because a lot of this is about torch-mlir.
9+
10+
Written & maintained by @renxida
11+
12+
Guides by other folks that were used during the creation of this document:
13+
- [Chi Liu](https://gist.github.com/AmosLewis/dd31ab37517977b1c499d06495b4adc2)
14+
- [Sunsoon](https://docs.google.com/document/d/1H79DwW_wnVzUU81EogwY5ueXgnl-QzKet1p2lnqPar4/edit?pli=1)
15+
16+
## Before you begin...
17+
18+
Nod-ai maintains the pipeline below, which allows us to take a ML model from e.g. huggingface, and compile it to a variety of devices including llvm-cpu, rocm and cuda and more as an optimized `vmfb` binary.
19+
20+
1. The pipeline begins with a huggingface model, or some other supported source like llama.cpp.
21+
2. [nod-ai/SHARK-Turbine](https://github.com/nod-ai/SHARK-Turbine) takes a huggingface model and exports a `.mlir` file.
22+
3. **[llvm/torch-mlir](https://github.com/llvm/torch-mlir)**, which you will be working on in turbine-camp, will lower torchscript, torch dialect, and torch aten ops further into a mixture `linalg` or `math` MLIR dialects (with occasionally other dialects in the mix)
23+
4. [IREE](https://github.com/openxla/iree) converts the final `.mlir` file into a binary (typically `.vmfb`) for running on a device (llvm-cpu, rocm, vulcan, cuda, etc).
24+
25+
The details of how we do it and helpful commands to help you set up each repo is in [Sungsoon's Shark Getting Started Google Doc](https://docs.google.com/document/d/1H79DwW_wnVzUU81EogwY5ueXgnl-QzKet1p2lnqPar4/edit?pli=1)
26+
27+
PS: IREE is pronounced Eerie, and hence the ghost icon.
28+
29+
## How to begin
30+
1. You will start by adding support for 2 ops in torch-mlir, to get you familiar with the center of our pipeline. Begin by reading [torch-mlir's documentation on how to implement a new torch op](https://github.com/llvm/torch-mlir/blob/main/docs/Torch-ops-E2E-implementation.md), and set up `llvm/torch_mlir` using https://github.com/llvm/torch-mlir/blob/main/docs/development.md
31+
2. Pick 1 of the yet-unimplemented from the following. You should choose something that looks easy to you. **Make sure you create an issue by clicking the little "target" icon to the right of the op, thereby marking the op as yours**
32+
- [TorchToLinalg ops tracking issue](https://github.com/nod-ai/SHARK-Turbine/issues/347)
33+
- [TorchOnnnxToTorch ops tracking issue](https://github.com/nod-ai/SHARK-Turbine/issues/215)
34+
3. Implement it. For torch -> linalg, see the how to torchop section below. For Onnx ops, see how to onnx below.
35+
5. Make a pull request and reference your issue. When the pull request is closed, also close your issue to mark the op as done
36+
37+
</details>
38+
39+
### How to TorchToLinalg
40+
41+
You will need to do 4 things:
42+
- make sure the op exists in `torch_ods_gen.py`, and then run `build_tools/update_torch_ods.sh`, and then build. This generates `GeneratedTorchOps.td`, which is used to generate the cpp and h files where ops function signatures are defined.
43+
- Reference [torch op registry](https://github.com/pytorch/pytorch/blob/7451dd058564b5398af79bfc1e2669d75f9ecfa2/torch/csrc/jit/passes/utils/op_registry.cpp#L21)
44+
- make sure the op exists in `abstract_interp_lib_gen.py`, and then run `build_tools/update_abstract_interp_lib.sh`, and then build. This generates `AbstractInterpLib.cpp`, which is used to generate the cpp and h files where ops function signatures are defined.
45+
- Reference [torch shape functions](https://github.com/pytorch/pytorch/blob/7451dd058564b5398af79bfc1e2669d75f9ecfa2/torch/jit/_shape_functions.py#L1311)
46+
- write test cases. They live in `projects/pt1`. See the [Dec 2023 example](https://github.com/llvm/torch-mlir/pull/2640/files).
47+
- implement the op in one of the `lib/Conversion/TorchToLinalg/*.cpp` files
48+
49+
Reference Examples
50+
- [A Dec 2023 example with the most up to date lowering](https://github.com/llvm/torch-mlir/pull/2640/files)
51+
- [Chi's simple example of adding op lowering](https://github.com/llvm/torch-mlir/pull/1454) useful instructions and referring links for you to understand the op lowering pipeline in torch-mlir in the comments
52+
53+
Resources:
54+
- how to set up torch-mlir: [https://github.com/llvm/torch-mlir/blob/main/docs/development.md](https://github.com/llvm/torch-mlir/blob/main/docs/development.md#checkout-and-build-from-source)
55+
- torch-mlir doc on how to debug and test: [ttps://github.com/llvm/torch-mlir/blob/main/docs/development.md#testing](https://github.com/llvm/torch-mlir/blob/main/docs/development.md#testing)
56+
- [torch op registry](https://github.com/pytorch/pytorch/blob/7451dd058564b5398af79bfc1e2669d75f9ecfa2/torch/csrc/jit/passes/utils/op_registry.cpp#L21)
57+
- [torch shape functions](https://github.com/pytorch/pytorch/blob/7451dd058564b5398af79bfc1e2669d75f9ecfa2/torch/jit/_shape_functions.py#L1311)
58+
59+
### How to TorchOnnxToTorch
60+
0. Generate the big folder of ONNX IR. Use https://github.com/llvm/torch-mlir/blob/main/test/python/onnx_importer/import_smoke_test.py . Alternatively, if you're trying to support a certain model, convert that model to onnx IR with
61+
```
62+
optimum-cli export onnx --model facebook/opt-125M fb-opt
63+
python -m torch_mlir.tools.import_onnx fb-opt/model.onnx -o fb-opt-125m.onnx.mlir
64+
```
65+
2. Find an instance of the Op that you're trying to implement inside the smoke tests folder or the generated model IR, and write a test case. Later you will save it to one of the files in `torch-mlir/test/Conversion/TorchOnnxToTorch`, but for now feel free to put it anywhere.
66+
3. Implement the op in `lib/Conversion/TorchOnnxToTorch/something.cpp`.
67+
4. Test the conversion by running `./build/bin/torch-mlir-opt -split-input-file -verify-diagnostics -convert-torch-onnx-to-torch your_mlir_file.mlir`. For more details, see https://github.com/llvm/torch-mlir/blob/main/docs/development.md#testing . Xida usually creates a separate MLIR file to test it to his satisfaction before integrating it into one of the files at `torch-mlir/test/Conversion/TorchOnnxToTorch`.
68+
69+
Helpful examples:
70+
- [A Dec 2023 example where an ONNX op is implemented](https://github.com/llvm/torch-mlir/pull/2641/files#diff-b584b152020af6d2e5dbf62a08b2f25ed5afc2c299228383b9651d22d44b5af4R493)
71+
- [Vivek's example of ONNX op lowering](https://github.com/llvm/torch-mlir/commit/dc9ea08db5ac295b4b3f91fc776fef6a702900b9)
72+
73+
## Contacts
74+
People who've worked on this for a while
75+
- Vivek (@vivek97 on discord)
76+
- Chi.Liu@amd.com
77+
78+
Recent Turbine Camp Attendees, from recent to less recent
79+
- Xida.ren@amd.com (@xida_ren on discord)
80+
- Sungsoon.Cho@amd.com
81+
82+
## Links
83+
84+
- Tutorials
85+
- [Sungsoon's Shark Getting Started Google Doc](https://docs.google.com/document/d/1H79DwW_wnVzUU81EogwY5ueXgnl-QzKet1p2lnqPar4/edit?pli=1)
86+
- This document contains commands that would help you set up shark and run demos
87+
- [How to implement ONNX op lowering](https://github.com/llvm/torch-mlir/blob/main/docs/importers/onnx_importer.md)
88+
- Examples
89+
- [A Dec 2023 example with the most up to date lowering](https://github.com/llvm/torch-mlir/pull/2640/files)
90+
- Chi's Example Lowering
91+
- Github issue and code detailing how to implement the lowring of an OP.
92+
- [Chi's simple example of adding op lowering](https://github.com/llvm/torch-mlir/pull/1454) useful instructions and referring links for you to understand the op lowering pipeline in torch-mlir in the comments
93+
- If you have questions, reach out to [Chi on Discord](https://discordapp.com/channels/973663919757492264/1104195883307892837/1180233875058868224)
94+
- [Vivek's example of ONNX op lowering](https://github.com/llvm/torch-mlir/commit/dc9ea08db5ac295b4b3f91fc776fef6a702900b9)
95+
- Find Ops To Lower
96+
- [Torch MLIR + ONNX Unimplemented Ops on Sharepoint](https://amdcloud-my.sharepoint.com/:x:/r/personal/esaimana_amd_com/Documents/Torch%20MLIR%20+%20ONNX%20Unimplemented%20Ops.xlsx?d=w438f26fac8fd44eeafb89bc99e2c563b&csf=1&web=1&e=Qd4eHm)
97+
- If you don't have access yet, request it.
98+
- nod-ai/SHARK-Turbine ssues tracking op support
99+
- [Model and Op Support](https://github.com/nod-ai/SHARK-Turbine/issues/119)
100+
- [ONNX op support](https://github.com/nod-ai/SHARK-Turbine/issues/215)
101+
102+
103+
## Chi's useful commands for debugging torch mlir
104+
105+
https://gist.github.com/AmosLewis/dd31ab37517977b1c499d06495b4adc2
106+
107+
## How to write test cases and test your new op
108+
109+
https://github.com/llvm/torch-mlir/blob/main/docs/development.md#testing
110+
111+
112+
113+
## How to set up vs code and intellisence for [torch-mlir]
114+
Xida: This is optional. If you're using VS code like me, you might want to set it up so you can use the jump to definition / references, auto fix, and other features.
115+
116+
Feel free to contact me on discord if you have trouble figuring this out.
117+
118+
You may need to write something like this into your
119+
120+
```.vscode/settings.json```
121+
122+
under `torch-mlir`
123+
124+
```json
125+
{
126+
"files.associations": {
127+
"*.inc": "cpp",
128+
"ranges": "cpp",
129+
"regex": "cpp",
130+
"functional": "cpp",
131+
"chrono": "cpp",
132+
"__functional_03": "cpp",
133+
"target": "cpp"
134+
},
135+
"cmake.sourceDirectory": ["/home/xida/torch-mlir/externals/llvm-project/llvm"],
136+
"cmake.buildDirectory": "${workspaceFolder}/build",
137+
"cmake.generator": "Ninja",
138+
"cmake.configureArgs": [
139+
"-DLLVM_ENABLE_PROJECTS=mlir",
140+
"-DLLVM_EXTERNAL_PROJECTS=\"torch-mlir\"",
141+
"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR=\"/home/xida/torch-mlir\"",
142+
"-DCMAKE_BUILD_TYPE=Release",
143+
"-DCMAKE_C_COMPILER_LAUNCHER=ccache",
144+
"-DCMAKE_CXX_COMPILER_LAUNCHER=ccache",
145+
"-DLLVM_ENABLE_PROJECTS=mlir",
146+
"-DLLVM_EXTERNAL_PROJECTS=torch-mlir",
147+
"-DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR=${workspaceFolder}",
148+
"-DMLIR_ENABLE_BINDINGS_PYTHON=ON",
149+
"-DLLVM_ENABLE_ASSERTIONS=ON",
150+
"-DLLVM_TARGETS_TO_BUILD=host",
151+
],
152+
"C_Cpp.default.configurationProvider": "ms-vscode.cmake-tools",
153+
"cmake.configureEnvironment": {
154+
"PATH": "/home/xida/miniconda/envs/torch-mlir/bin:/home/xida/miniconda/condabin:/home/xida/miniconda/bin:/home/xida/miniconda/bin:/home/xida/miniconda/condabin:/home/xida/miniconda/bin:/home/xida/miniconda/bin:/home/xida/miniconda/condabin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin"
155+
},
156+
"cmake.cmakePath": "/home/xida/miniconda/envs/torch-mlir/bin/cmake", // make sure this is a cmake that knows where your python is
157+
}
158+
```
159+
The important things to note are the `cmake.configureArgs`, which specify the location of your torch mlir, and the `cmake.sourceDirectory`, which indicates that CMAKE should not build from the current directory and should instead build from `externals/llvm-project/llvm`

0 commit comments

Comments
 (0)