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).
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.
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.
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.
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.
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.
f = f.with_attr("tir.noalias", True, move=True)
- A1: only allow global functions, and allow pass in a moved parameter.
f = tvm.with_attr(f.move(), "tir.noalias", True)
- A2: Create a separate API with move in it (naming suggestion more than welcomed)
f = f.move_with_attr("tir.noalias", True)
Python Move Behavior
- P0: Move ways means move the object, will affect all python reference to the same object
f2 = f
# f2 is also affected, because they points to the same object
f = attach_attr(f.move(), "tir.noalias", 1)
- P1: Only creates R-value reference if there is a single reference count of the f
f2 = f
# move checks the python refcount
# because there is another reference to the object, f.move()
# simply returns f and we don't do rvalue copy
f = attach_attr(f.move(), "tir.noalias", 1)