adding_ops docs

docs/adding_ops.md

@@ -0,0 +1,80 @@
+## Adding Ops
+In order to implement an Op we need to 1) create a python facing interface for the Op via pybind11, and 2) use the torch-mlir C-API to emit the correct Op when this binding is invoked.
+
+### [Example: Overloading a method signature](https://github.com/nod-ai/PI/pull/30)
+Many Op implementations are already autogenerated in `cpp_ext/TorchOps.impls.cpp` hence if we only need to overload the python binding (say, so that we can accept a different method signature) then we can just reuse the implementation in `TorchOps.impls.cpp` and create another python binding for the corresponding operation.
+
+For example a relatively common case is that we have to accept either a list of integers or a parameter pack that can have multiple integers in order to define the dimensions along which an operation will act or the shape of an output tensor. Take for example `torch.Tensor.view` whose `shape` parameter can either be a list or a sequence of integers (i.e. both of the following are legal: `x.view([3,3,-1])` and `x.view(3,3,-1)`). Luckily, this has been implemented already in `TorchOps.impls.cpp` for the case where the shape is a list, so all we need to do to bind this method signature to our own `Tensor` class is to create a python binding that can take in the arguments we expect and call the original function
+
+```cpp
+// @overload view(self, *size _int) -> Tensor  
+c.def(  
+	"view",  
+	[](PyAnyTorchTensorValue &self, const py::args &size,  
+	DefaultingPyLocation &loc, const DefaultingPyInsertionPoint &ip) {  
+		return view(self, PyAnyTorchListOfTorchIntValue(size), loc.get(),  
+	ip.get());  
+	},  
+	// When combining *args or **kwargs with Keyword arguments you should not  
+	// include py::arg tags for the py::args and py::kwargs arguments.  
+	py::kw_only(), "loc"_a = py::none(), "ip"_a = py::none());
+```
+Note that torch-mlir expects certain types for input operands to torch-mlir Ops, thus attempting to emit the view Op by using the `createOperation` function with different operand types would result in a failure when lowering to linalg IR. We can check the expected operand and result types for each operation by inspecting the relevant tablegen, either in [TorchOps.td](https://github.com/llvm/torch-mlir/blob/3ca35b4f3cf10b08d6c3f2e19285899f440b3ac2/include/torch-mlir/Dialect/Torch/IR/TorchOps.td#L1) or [GeneratedTorchOps.td](https://github.com/llvm/torch-mlir/blob/3ca35b4f3cf10b08d6c3f2e19285899f440b3ac2/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td#L1)
+### Example: Implementing from scratch
+Sometimes we don't have an autogenerated Op definition in our C++ extension, in this case we have to implement the logic that emits the Op. This mainly consists of copy and pasting some boiler plate to emit the operation and creating a python binding that allows us to call the Op from our python frontend. For example if we want to implement `torch.abs` we first have to emit the Op via some logic in `TorchOps.cpp`
+
+```cpp
+// prim::abs.Scalar : (Scalar) -> (Scalar)  
+PyAnyTorchScalarValue abs(const PyAnyTorchScalarValue &a, PyLocation *loc,  
+PyInsertionPoint *ip) {  
+	auto resultType = py::cast(mlirValueGetType(a)).cast<PyType>();  
+	PyOperationRef opRef =  
+	createOperation("torch.prim.abs.Scalar", {resultType}, {a},  
+	/*attributes=*/{}, loc, ip);  
+	return {opRef, mlirOperationGetResult(opRef->get(), 0)};  
+}
+```
+The operation name `torch.prim.abs.Scalar` can be retrieved in many ways, but the simplest is to check either `_torch_ops_gen.py` or `TorchOps.td` in torch-mlir.
+
+We then create a binding for the python interface at the bottom of `TorchOps.cpp`
+
+```cpp
+// prim::abs.Scalar : (Scalar) -> (Scalar)  
+m.def(  
+"abs",  
+[](const PyAnyTorchScalarValue &x, const DefaultingPyLocation &loc,  
+const DefaultingPyInsertionPoint &ip) {  
+	return abs(x, loc.get(), ip.get());  
+},  
+"x"_a, py::kw_only(), "loc"_a = py::none(), "ip"_a = py::none());
+```
+Often if there is still an autogenerated stub or incorrect python binding in `TorchOps.impls.cpp`, `TorchOps.pybinds.cpp` or `TorchTensor.pybinds.cpp` then we also need to prevent these stubs from being autogenerated in the future, to do this we need to edit the autogeneration logic.
+
+### Editing the autogeneration logic
+When need to rewrite the autogenerated definition of an Op or its autogenerated pybind, we need to prevent the autogeneration script from regenerating that Op - otherwise the autogenerated version will conflict with our own definition.
+
+First we run `scripts/generate_stuff/dump_torch_mlir_ods_json_and_initpyi.sh` to generate a `torch.json` file, once the file is generated we must edit `scripts/generate_stuff/generate_torch_mlir_bindings_from_torch_json.py`. In this file there are two python dictionaries we can edit: `SKIP_OPS` and `SKIP_TENSOR_BINDS`. These correspond to operations generated at the module level (`pi.abs()`, `pi.add()`, etc.) or at the class level for the `Tensor` object (`x.view()`, `x.mean()`, etc.). To prevent an Op from being autogenerated we add its signature to whichever one of these dictionaries is appropriate.
+
+For example if we wanted to prevent the `constant.str` Op from being generated we could [lookup its name in the tablegen](https://github.com/llvm/torch-mlir/blob/c9add6b7d8d8880b707fefe655eadbb4f08338f5/include/torch-mlir/Dialect/Torch/IR/TorchOps.td#L675) and add it to the `SKIP_OPS` dictionary:
+
+```python
+SKIP_OPS = {
+			...
+			"Torch_ConstantStrOp",
+			...
+			}
+```
+
+Similarly if we want to prevent autogeneration for a binding from the `Tensor` class we can use its python signature to do so:
+
+```python
+SKIP_TENSOR_BINDS = {
+			...
+			"@overload view(self, *size: _int) -> Tensor",
+			...
+			}
+```
+
+Finally we should run `scripts/generate_stuff/generate_torch_mlir_bindings_from_torch_json.py` to remove the autogenerated items.
+
+[Example PR](https://github.com/nod-ai/PI/pull/59)