TVM stack uses the PackedFunc and TypedPackedFun extensively to expose
functions to the frontend. They serve as a foundation of the runtime system and
FFI.
A PackedFunc call passes an ObjectRef type by the pointer value of the internal
object pointer. The corresponding `Object*` values can be viewed as a lvalue
reference to the object. The callee will need to increase the ref counter to
get another strong reference to the same object Iin the function body.
### Immutable Objects and Need for Move Semantics
While the current PackedFunc calling convention served as quite well so far, it
have one potential problem, which is explained below.
We make most IR objects to be **immutable** in the compiler infrastructure. The
immutable design brings various benefit in terms of easier reasoning of
correctness, thread safety and caching (we can cache hash value and other
functions computed on immutable objects).
However, the immutable nature means that we have to copy a data structure
whenever we want to transform its content, although the copy is usually
shallow, and only with respect to the node to be mutated, we also need to copy
all the parents that references the node. Such copy happens quite often in
transformation passes and simple actions such as attaching a new attribute to a
function.
The solution to the problem is to introduce the **copy on write** optimization
to most of the immutable objects. In particular, if we want to mutate a unique
reference to the IR node(whose ref count equals 1), we do not have to copy and
can write inplace (since there is no other reference to the same node).
```c++
PrimFunc Transform(PrimFunc f) {
// create a new body
auto update = AddAttr(PrimFunc f) {
// if there f is an unique reference, no copy will be performed.
f.CopyOnWrite()->body = new_body;
return WithAttr(std::move(f), tir::attr::NoAlias, Integer(1));
};
return update(std::move(f));
}
void Run() {
PrimFunc f = CreateFunc();
f = Transform(std::move(f));
f = Transform(std::move(f));
}
```
The above code-block gives an example of copy on write pattern. In order to
make sure the code always keep an unique reference to f, we use move
extensively to pass these values by rvalue reference. In C++ we use `std::move`
to pass values by r-value reference. Notably, passing by r-value is the default
calling convention in Rust. By using r-value passing and copy on write, we get
the benefit of immutable objects without the need to do extensive copies
during transformations.
```c++
f = WithAttr(std::move(f), attr0, val0);
f = WithAttr(std::move(f), attr1, val1);
f = WithAttr(std::move(f), attr2, val2);
```
One interesting thing that worth noting is copy on write optimization can
easily chains together. The above code block will trigger copy of f for at
most once. Even if f is not the unique refernce in the first call to WithAttr.
The copy triggered in the first call will return an unique reference, and the
subsequent calls to WithAttr won't copy.
However, due to the current limitation of the PackedFunc calling convention
(which can only pass objects as l-value), we can no longer use the move + copy
on write pattern in a PackedFunc.
```c++
PrimFunc Transform(PrimFunc f) {
// create a new body
auto update = AddAttr(PrimFunc f) {
// we won't have a unique copy of f under the current calling convention
return WithAttr(std::move(f), tir::attr::NoAlias, Integer(1));
};
return TypedPackedFunc<PrimFunc(PrimfFunc)>(update)(std::move(f));
}
```
When we pass an object to the TypedPackedFunc, there is a reference in the
caller side. Additonally, the callee need to create anothe reference in the
function body — this means a object argument in a TypedPackedFunc is never
going to be unique. This limitation removes the possibility of the copy on
write optimization in passes since we are using TypedPackedFunc as basic
building blocks for pass functions.
### Support Object RValue Reference in PackedFunc Calls
We propose to introduce RValue reference support the PackedFunc calling
convention to address the above issue. Specifically, when we find that an
argument is a r-value reference, we will use a assign a different type
code(`kObjectRValueRefArg`), and pass `Object**` (the address to the Object
pointer) instead through the values array. The callee can choose to move out
this Object pointer and set the original Object pointer from the caller side to
be nullptr.
```c++
class ObjectPtr {
private:
/*!
* \brief Move an ObjectPtr from an RValueRef argument.
* \param ref The rvalue reference.
* \return the moved result.
*/
static ObjectPtr<T> MoveFromRValueRefArg(Object** ref) {
ObjectPtr<T> ptr;
ptr.data_ = *ref;
*ref = nullptr;
return ptr;
}
friend class PackedFunc;
};
```
The enhanced calling convention allows a single reference will be moved from
the caller side to the callee side, enabling the copy on write optimization
when necessary.
### Move a Python Object
We face the same copy-on-write multiple reference problem in the python side.
Because the caller in the python side retains the reference to the original
object. We could resolve the issue by introducing a move semantics for tvm
python objects. The main idea is to support a function(`move`) which indicate
that an object is moved and can be passed to a PackedFunc. Note that we cannot
use the object after it is being moved. The following code snippet provide an
example.
```python
def test(f):
# will result in a copy, because f is retained.
f = attach_attr(f, "tir.noalias", 1)
# will not result in a copy due to Copy on write
f = attach_attr(f.move(), "tir.noalias", 1)
```
There are multiple API design choices. It would be great to get feedback from
everyone about their opinions about the design.
**Choice of Move API**
- M0: make move a member function`f.move()`
- M1: make move a global function `tvm.move(f)`
- M2: Do not introduce move support to python side yet.
**Choice of API that involves Move**
- A0: allow a move attribute as part of the API function, indicating we want to
move the original argument.
```python
f = f.with_attr("tir.noalias", True, move=True)
```
- A1: only allow global functions, and allow pass in a moved parameter.
```python
f = tvm.with_attr(f.move(), "tir.noalias", True)
```
- A2: Create a separate API with move in it (naming suggestion more than
welcomed)
```python
f = f.move_with_attr("tir.noalias", True)
```
---
[Visit
Topic](https://discuss.tvm.ai/t/rfc-support-rvalue-reference-passing-in-typedpackedfunc/6264/1)
to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.ai/email/unsubscribe/0d20f8fe5352faf5d8063540859131afa53192b5c47cfc88ae014e9408e0a6d1).
