[RFC] More PackedFunc metadata

(Moved from #2981)

Proposal: add type signature metadata to PackedFunc

Each PackedFunc would have an optional field describing what arguments it takes and what type it returns. This field would be automatically populated during conversion from a TypedPackedFunc.

Once we have these type signatures, we could use them to generate bindings in new languages (i.e., Rust) and get coverage of a good chunk of TVM’s API for free.

This would result in much easier-to-maintain language bindings (not only for the runtime, for the compiler stack too) and might allow us to eventually rip out a lot of the manually-written stuff in e.g. the Python bindings.

The downside of this idea is that it would result in some fairly hairy codegen, which can be difficult to maintain. It would also add small amounts of overhead to the tvm runtime system; we could reduce that issue by adding a #define to disable the type metadata. Also, as @tqchen pointed out, the generated code might be lower-quality than handwritten bindings.

A few extensions of this idea:

• Add more compile-time metadata to the Node heirarchy, allowing codegen to access their methods / attributes.
• Add docstrings to PackedFuncs to allow auto-generation of documentation.
• Allow std::variant or equivalent in TypedPackedFunc signatures. Lots of PackedFuncs have arguments that can be one of a few types (e.g. Int or Expr); a simple extension to the PackedFunc system + runtime type system would allow these to be described automatically.

[tqchen]

One thing I would like to mention is that manual wrapping has its merit 😃 Take the following wrapper code for scan as an example. By manually wrapping it in python, we can benefit from:

• A clear signature that users can look up
• Keyword argument and default value
• Python-specific documents,
  • python type signature
  • useful code python examples

def scan(init, update, state_placeholder, inputs=None, name="scan", tag="", attrs=None):
    """Construct new tensors by scanning over axis.

    Parameters
    ----------
    init: Tensor or list of Tensor
        The initial condition of first init.shape[0] timestamps
    update: Tensor or list of Tensor
        The update rule of the scan given by symbolic tensor.
    state_placeholder: Tensor or list of Tensor
        The placeholder variables used by update.
    inputs: Tensor or list of Tensor, optional
        The list of inputs to the scan. This is not required, but can
        be useful for the compiler to detect scan body faster.
    name: str, optional
        The name hint of the tensor
    tag: str, optional
        Additonal tag information about the compute.
    attrs: dict, optional
        The additional auxiliary attributes about the compute.

    Returns
    -------
    tensor: Tensor or list of Tensors
        The created tensor or tuple of tensors it it contains multiple outputs.

    Example
    -------
    .. code-block:: python

      # The following code is equivalent to numpy.cumsum
      m = tvm.var("m")
      n = tvm.var("n")
      X = tvm.placeholder((m, n), name="X")
      s_state = tvm.placeholder((m, n))
      s_init = tvm.compute((1, n), lambda _, i: X[0, i])
      s_update = tvm.compute((m, n), lambda t, i: s_state[t-1, i] + X[t, i])
      res = tvm.scan(s_init, s_update, s_state, X)

While it is possible to have a system that codegen these components, that might mean we have to write the python example documents directly in somewhere else, which is less natural.

Of course, one drawback of doing the manual wrapping is that one wrapper has to be created for each language. This may not be a bad thing, especially we want to think clearly about the language specific features, and write good docs that are language specific

I do think there is some merit to do have good metadata for the node system and automatically generate some of the accessors.

[kazimuth]

NodeRef is just a type-erased version of NodePtr<T> right?

I think the main question is if we want to allow nodes to be implemented in languages besides C++. If we require nodes to be implemented as Node subclasses, then they have to have C++ impls for inheritance to work. If we design a Node C API that doesn’t have this requirement, then we can mix-and-match Node implementations between languages. That’s really only useful for adding temporary node types for e.g. passes implemented in python though. (Of course, we can still implement all the C++ nodes as Node subclasses.)