[RFC] TVM Object Schema DSL

Introduction

TVM Object system provides a convenient and decent way to share objects between backend (C++) and frontend (Python/Java/Rust/etc.). For example, one can construct a variable in Python and pass it to functions written in C++, and vice versa.

However, adding one object node into TVM stack requires manually adding lines of code to different places in both Python and C++. For example, here’s how tvm::tir::IntImm is implemented and registered,

This RFC advocates the approach to generate C++ implement directly from Python class definition and registry. Moreover, as we still allow users to write C++ code manually in order to bring in more complex features, the object transpiler will provide basic validation for these manually written C++ code.

Here is an example of how an object can be described in Python and how the generated C++ code looks like:

@declare
class BaseExprNode(Object):
    """
    Base type of all the expression.

    See Also
    --------
    BaseExpr
    """
    type_key = "BaseExpr"
    default_visit_attrs = False
    default_sequal_reduce = False
    default_shash_reduce = False

@declare
class IntImmNode(PrimExprNode):
    """
    Constant integer literals in the program.

    See Also
    --------
    IntImm

    Attributes
    ----------
    value
        The internal value.
    """
    type_key = "IntImm"
    value: ty.int64_t

/*!
 * \brief Base type of all the expressions.
 * \sa Expr
 */
class BaseExprNode : public Object {
 public:
  TVM_DECLARE_BASE_OBJECT_INFO(BaseExprNode, Object);
};

/*!
 * \brief Managed reference to BaseExprNode.
 * \sa BaseExprNode
 */
class BaseExpr : public ObjectRef {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(BaseExpr, ObjectRef, BaseExprNode);
};

/*!
 * \brief Constant integer literals in the program.
 */
class IntImmNode : public PrimExprNode {
 public:
  /*! \brief The internal value. */
  int64_t value;
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("value", &value);
  }
  void SEqualReduce(const IntImmNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(value, other->value)
  }
  void SHashReduce(SHashReducer hash_reducer) const {
    hash_reducer(dtype);
    hash_reducer(value);
  }
  static constexpr const char* _type_key = "IntImm";
  TVM_DECLARE_BASE_OBJECT_INFO(IntImmNode, PrimExprNode);
};

/*!
 * \brief Managed reference class to IntImmNode.
 *
 * \sa IntImmNode
 */
class IntImm : public PrimExpr {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(IntImm, PrimExpr, IntImmNode);
};

We name it as TVM Object Schema DSL, or tschema. In summary, tschema will bring several benefits for the TVM architecture:

  • Reduce boilerplate code;
  • Verify to avoid missing some definition like TVM_REGISTER(...);
  • Enable deployment on all kinds environment even without C++;
  • Fields like type_child_slots can be automatically generate for optimizing;
  • Allow users to define Objects in Python, build and export them to a .o/.so file;
  • Have more type information during runtime, enable some optimizations in TIR compilation;

High-level Object Compilation Pipeline

  • Define TVM Object in Python. This object definition Python file is in a seperate directory (which will not be a part of PYTHONPATH) other than python/tvm/
  • Run python class parser to generate related .h, .cc files. This step can be triggered manually or via cmake. The generated files will be checked into the code base so that code completion tools can locate them.
  • Compile TVM using cmake as usual.

Notice that the second step happens during (or before) compiling TVM itself. We provide a standalone tool to parse the Python code.

TSchema DSL

Take IntImm as an example, the

@declare
class IntImmNode(PrimExprNode):
    """
    Constant integer literals in the program.

    See Also
    --------
    IntImm

    Attributes
    ----------
    value
        The internal value.
    """
    type_key = "IntImm"
    value: ty.int64_t

There are several things require to be parsed,

  • Object name. In the above example it is IntImmNode, therefore class IntImmNode (extends Object) will be generated.
  • Type key. In the above example it is IntImm, therefore class IntImm (extends ObjectRef) will be generated.
  • Parent class. In the above example it is PrimExprNode
  • Member variables. In the above example they are,
    • value and its type annotation int64_t
  • The constructor arguments in C++ will be generated as the same order of the arguments in Python class definition.
  • We also will generate default VisitAttrs, SEqualReduce, SHashReduce methods unless user specify default_visit_attrs as False.

Inplace C++ Source File Modification

As we mentioned before, there are cases where users need to implement complex functions manually. To leverage the convenience of Python declaration and automatic code generation in such cases, we provide an option to modify the C++ source file in-place, and give users the control to specify which part of the file can be modified.

We provide comment parser for .h and .cc file, in which users can wrap the auto-generated section by comments, e.g.,

// tschema: ObjectName

The lines between tschema: ObjectName and tschema: end
will be manipulated by tschema

// tschema: custom-begin

User can also mark sections which should be left unchanged by objgen
This section will be inserted at the end of the class definition,
right before the close brace

// tschema: custom-end
// tschema: end

Here is also an example for it:

Before generation

// tschema: GlobalVarNode
// tschema: custom-begin
bool SEqualReduce(const GlobalVarNode* other, SEqualReducer equal) const {
  return equal(name_hint, other->name_hint) && equal.FreeVarEqualImpl(this, other)
}
bool SHashReduce(SHashReducer hash_reducer) const {
  hash_reduce(name_hint);
  hash_reduce.FreeVarHashImpl(this);
}
// tschema: custom-end
// tschema: end

TSchema Definition

@declare
class GlobalVarNode(RelayExprNode):
    """
    Global variable that lives in the top-level module.

    A GlobalVar only refers to function definitions.
    This is used to enable recursive calls between function.

    See Also
    --------
    GlobalVarNode

    Attributes
    ----------
    name_hint
        The name of the variable, this only acts as a hint.
    """
    type_key = "GlobalVar"
    default_sequal_reduce = False
    default_shash_reduce = False
    name_hint: ty.String

Generated Code

// tschema: GlobalVarNode
class GlobalVarNode : public RelayExprNode {
 public:
  String name_hint;
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("span", &span);
    v->Visit("checked_type_", &checked_type_);
    v->Visit("name_hint", &name_hint);
  }
  static constexpr const char* _type_key = "GlobalVar";
  TVM_DECLARE_BASE_OBJECT_INFO(GlobalVarNode, RelayExprNode);
  // tschema: custom-begin
  bool SEqualReduce(const GlobalVarNode* other, SEqualReducer equal) const {
    return equal(name_hint, other->name_hint) && equal.FreeVarEqualImpl(this, other)
  }
  bool SHashReduce(SHashReducer hash_reducer) const {
    hash_reduce(name_hint);
    hash_reduce.FreeVarHashImpl(this);
  }
  // tschema: custom-end
};
// tschema: end

@tqchen @yzhliu @jwfromm @jroesch @junrushao1994 , also thanks Yizhi for the initial idea and RFC writing.

6 Likes

Thanks for the RFC and it looks super useful! I have two questions from the RFC:

  1. Are the generated .h and .cc files supposed to be tracked in the repo, or they are more like the build files?

  2. For the in-place modification, I am a bit confused about the example of customized C++ code (the example of Before generation). I imagine the TSchema definition is a standalone Python file. Then where should this piece of C++ code be specified?

Thanks.

Hi @comaniac,

They will be tracked in the repo. But user should write the tschema for objects instead of writing the cpp files directly except the tschema-custom part.

Sorry that I did not make it clear, there actually no such “before generation” files. Finally we will keep the generated code in our codebase and normal users will build them directly. I just use the code snippets to explain what’s tschema’s job.

Thanks for the clarification :slight_smile: So this is more like a tool to facilitate the development process. After the C++ code has been generated, we can continue working on it as always.

1 Like

Thanks, this is really an interesting work! For those who have a requirement to add their own modifications to use TVM, it will be very helpful!

I’m just thinking about how frequently will this new feature be used. IMHO, advanced users who may benefit from it are more likely to write their C++ code directly, while other users may not really have a requirement on this.

Another problem is I guess it will be hard for a IDE or editor(e.g. VSCode, Vim with CTags) to track the code and provide navigation?

Yeah, this could be a useful tool to generate the generic templates or the code with the fixed pattern which is actually the major part of a node. For some other members, e.g. SEqualReduce and SHashReduce, we may still need users to manually check/add since they are not always Equal(this->a, other->a) && Equal(this->b, other->b);

Hi @jcf94,

First, this is not only for C++ code generation. In the future, we will extend it for Python/Rust code generation, which is helpful for unifying object definitions between different languages.

Second, some object fields is hard to fill in even for advanced users, e.g, type_child_slots, which is the number of object’s children.

And last but not least, by defining objects with tschema, we will have more in-memory information about the object itself. For example, the type hierarchy between objects, the memory layout of an object, etc. This will enable more compilation optimization in the TIR and help us improve TIR’s type system (my next step).

Since we will keep the generated C++ code in the codebase, it will not make any difference with current code in terms of code navigation.

@zhiics Yep, we have an option to turn off the default method generation and allow user to fill their customized code snippets.

Hey @ziheng! I think this is a great idea. As someone who is pushing on the Rust bindings right now (along with @jroesch), I love the idea of deduplicating work.

One design choice I see is whether to centralize or decentralize code-generation. It seems like your original design is leaning towards centralizing it. It would like to start a little discussion on why/if this is the right idea.

Decentralizing code generation could have some benefits, here’s how I see it looking. The schemas themselves live in some central location like /schema, and they are defined as simply data (perhaps JSON). Each “backend”, including C++ and Python, is then responsible for reading this data and generating code for itself. The downside is that there may be some duplicated logic in the code generation. But on the upside, each backend gets to use different tooling to do the codegen; for example, it would be nice to use Rust (using syn and quote) to generate the Rust code. This could also simplify the story for implementing additional code for the methods: each backend just handles it itself, no need to toggle or parse anything.

Here’s a example on what the JSON could look like:

[
    {
        "name": "IntImmNode",
        "key": "IntImm",
        "ref": "IntImm",
        "parent": "PrimExprNode",
        "fields": { "value": "int64" }
    },
    ...
]

You could imaging grouping these schema into namespaces or something too, if you want.

On the topic of checking the generated code in, I’m not sure why that is necessary. As long as the files are generated by the build system, shouldn’t autocomplete and stuff work fine?

Hi @mwillsey, The decentralizing code generation sounds a good idea technically! We choose Python mainly for user-friendly. I would also like to know @tqchen’s opinion here.

We can make an automated build pipeline, but checking in the code directly will make the project codebase more clear. After all, not all the user need to know those details.

I like the idea of using rust to generate rust side. In the meantime, a python syntax for data structure setup can be useful in the future when we want to design custom data types from python side. One potential solution is we keep the python schema frontend, and generate a json exchange format that the rust generator can take. Essentially a separation of frontend, ir repr and backend.