You are viewing a plain text version of this content. The canonical link for it is here.
Posted to dev@tvm.apache.org by Lite Ye <no...@github.com.INVALID> on 2022/06/16 17:35:40 UTC
[apache/tvm-rfcs] [RFC] Add TVMScript Metaprogramming RFC (PR #79)
[Rendered](https://github.com/yelite/tvm-rfcs/blob/tvmscript-metaprogramming/rfcs/0079-tvmscript-metaprogramming.md)
You can view, comment on, or merge this pull request online at:
https://github.com/apache/tvm-rfcs/pull/79
-- Commit Summary --
* Add TVMScript Metaprogramming RFC
-- File Changes --
A rfcs/0079-tvmscript-metaprogramming.md (398)
-- Patch Links --
https://github.com/apache/tvm-rfcs/pull/79.patch
https://github.com/apache/tvm-rfcs/pull/79.diff
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79
You are receiving this because you are subscribed to this thread.
Message ID: <apache/tvm-rfcs/pull/79@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Lite Ye <no...@github.com.INVALID>.
Cross post from https://discuss.tvm.apache.org/t/rfc-tvmscript-metaprogramming/12969/2 for more visibility.
From @Johnson9009,
> Thanks for the work, for F4，do you mean we can write some impreative code that can be execute by Python?
Yes, that’s right. Actually you can write imperative code to construct IR graph now, by calling IR node constructors from tvm.tir.stmt|expr manually. The new things behind F4 are:
1. IRBuilder API, which can be used to construct IR nodes in a more concise way compared to calling node constructor directly.
2. The ability to embed these functions inside TVMScript. This means you can use imperative code to construct small fragments inside a large T.prim_func.
I update the RFC to clarify when and where the imperative code gets evaluated.
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1159467109
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Tristan Konolige <no...@github.com.INVALID>.
I'm still a bit confused on 2. The example you give for a single class is one that has all static members, but what if instead it was just a regular class. It would instantiate a single instance of said class to do parsing. This is what relax is doing (https://github.com/tlc-pack/relax/blob/25cb42133130fb8ec75d2cc021c6c347e99f0b28/python/tvm/script/relax/parser.py#L107). Instantiating a class would also give the benefit of being able to maintain some state while parsing which is not possible with either free functions or a class with static methods. Here is how I would imagine it (this is just an example):
```python
# To clearly define the interface
class AbstractParser:
def parse(func: ast.Function):
raise NotImplemented()
@register("relax")
class RelaxParser(AbstractParser):
def __init__():
self.tir_parser = TIRParser()
def parse(func: ast.Function):
self.inputs.push(parse_args(func.args))
self.tir_parser.add_variables(self.inputs)
parse_stmt(func.body)
def parse_stmt(stmt):
if stmt == ast.Let:
assert stmt.var in self.inputs
if stmt == ast.Function:
# parse nested functions as TIR
self.tir_parser.parse(stmt)
```
If we want to consider future dialects, this also has the added benefit of being able to add dialects by subclassing one of the existing parsers.
Another issue I see with the proposed registration method is that there is no way to handle different parser rules based on context. As an example, what if I want to treat the first line of a function body differently than every other line (say it is a metadata dictionary):
```python
def foo():
{"version": "0.1"} # This should become a `MetadataNode`, only strings allowed as values.
my_dictionary = {"hi", 1} # This should become a `DictLiteralNode`, only integers allowed as values.
```
With the proposed registration method there is no way to distinguish between what should be a `DictLiteralNode` and a `MetadataNode`:
```python
@dispatch.register(token="mylang", type_name="Dict")
def visit_dict(self: Parser, node: doc.Dict) -> None:
for key_ast, value_ast in node.items:
key = visit(key_ast)
assert isinstance(key, StringLiteral)
value = visit(value_ast)
assert isinstance(value, StringLiteral) # or wait, should it be IntLiteral? There's no way to tell if this is the first line or not.
```
(Maybe I am misunderstanding how this approach works though. Can we put arbitrary state into the the `self: Parser` object? If so then we can fix this issue. But it would be a complicated and confusing solution as the parser would have to maintain a stack of rules takes or flags or some such.)
But with the class based approach I proposed above:
```python
@register("mylang")
class MyLangParser:
def parse_function(func):
# assuming func.body is a list of statements, also applies if body is a recursive set of statements
parse_metadata(func.body[0])
for stmt in func.body[1:]:
parse_statement(stmt)
def parse_metadata(stmt):
for key_ast, value_ast in node.items:
key = parse_string(key_ast)
value = parse_string(value_ast)
def parse_statement(stmt):
for key_ast, value_ast in node.items:
key = parse_string(key_ast)
value = parse_int(value_ast)
```
This issue is caused by registering visitors at the python node ast level. When visiting each ast node there is no context as to where we may be in the production rules. If we instead my proposed approach, the parser can control which production rules can be used at each location.
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1199837603
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Junru Shao <no...@github.com.INVALID>.
To follow up with our latest discussion with @tkonolige @areusch @csullivan @jwfromm et al.
The following questions are raised in our discussion:
1. Move discussion of vendor IR to tradeoffs / benefits section rather than core motivation.
2. (Section 2) Parser registration example is a little confusing (dispatch.register), probably just needs more detail, less …
3. (Section 1) Add tradeoffs section noting that this isnt a typical parser, gives much more pythonic interface, easier testing, more flexibility.
4. (Section 3) More discussion around what lives in host language. Maybe need script.host_var and script.quote in part of future work, might be worth highlighting in a future work section.
5. (Section 2) tradeoff on registration approach.
Our response:
### 1. Tradeoffs: TVMScript parser is not a typical parser
Unlike a typical parser that converts a token stream to an AST, which in our case, converts python source code to TIR AST, the TVMScript parser is technically speaking a transpiler that transforms a Python AST to TIR AST instead of parsing raw python.
Pros:
1. Python’s AST package provides accurate parsing functionality, so that all issues around ambiguity, grammar and parser performance is a non-issue.
2. Compared with the existing monolithic TVMScript parser which is 1.5k lines of code, the transpiler provides much easier testing, more pythonic interface and more flexibility.
Cons:
1. This means we are depending on Python’s AST package - which could come with breaking changes across python versions
### 2. Registration Logic
The registration logic is in fact quite straightforward. For example, to support TIR’s for-loop syntax like:
```python
for i, *many_j, k in T.grid(...): # or anything like T.serial/parallel/vectorized/unroll/thread_binding
...
```
The registration logic is as simple as calling into our IR-builder:
```python
@dispatch.register(token="tir", type_name="For")
def visit_for(self: Parser, node: doc.For) -> None:
for_frame = self.eval_expr(node.iter)
if not isinstance(for_frame, T.ForFrame):
self.report_error(
node.iter,
"Expect the for loop to be one of the following: "
"range, T.serial, T.grid, T.parallel, T.vectorized, T.unroll, T.thread_binding",
)
with self.var_table.with_frame():
with for_frame as iters:
self.eval_assign(target=node.target, source=iters, bind_value=bind_value)
self.visit_body(node.body)
```
There is an alternative proposal that registration should happen at class-level instead of method-level, e.g.
```python
## Our RFC
@dispatch.register(token="tir", type_name="For")
def registered_method(self: Parser, node: doc.For) -> None: ...
## Alternative
@dispatch.register(token="tir"):
class DispatchOfTIR:
@staticmethod
def visit_For(self: Parser, node: doc.For) -> None: ...
```
The advantage of the alternative proposal is that it limits the users so that they have to put all the logic inside a class, while the disadvantage is that the class itself doesn’t mean anything other than a collection of static methods, which could bring some confusion if developers attempt to instantiate the class.
To this end, “putting all logic inside a class” is equivalent to “putting all logic inside a file” because the class only serves as a namespace. Therefore, we believe the best design should be as minimal as possible, i.e. without a class.
**Drawback**. Registry is a necessary design when it comes to supporting per-node tooling and multiple IRs at scale. However, less carefully designed registries, for example, relay operator strategy, where definitions span in multiple folders and tens of files, would lead to worse discoverability and debugging experience. On the other hand, in our particular case, because all dispatches for a certain IR is confined in a single file (e.g. [tir.py](http://tir.py/)), it would not be a particular concern.
### 3. Metaprogramming syntax
A question has been raised on supporting quotations for metaprogramming in languages like MetaOCaml [1].
This feature is very cool, and it could be an important future work for well motivated use-cases. In terms of syntax, we believe the following could be like:
```python
with script.quote(): ## <===== hand back control to python interpreter
arbitrary python program
with script.unquote(): ## <===== gain control from python interpreter
parsable TVMScript
```
[1] MetaOCaml -- an OCaml dialect for multi-stage programming [https://okmij.org/ftp/ML/MetaOCaml.html](https://okmij.org/ftp/ML/MetaOCaml.html)
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1198582332
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Junru Shao <no...@github.com.INVALID>.
Merged #79 into main.
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#event-7130575684
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by "Steven S. Lyubomirsky" <no...@github.com.INVALID>.
So if you define a variable in a quoted portion, you should be able to reference it in the quoted portion?
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1198624317
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by xqdan <no...@github.com.INVALID>.
@yelite It's a great RFC，and this is what we need right now.
the requirements we need:
1) For compute fusion. With TE compute, it's easy to concate TE computes with producer-comsuer relation to get a fused compute. for example, conv + elemwise ops fusion. We should have similar function in TVM script. Which thread is related to this requirement?
2) For conditional lowering. We may have some attributtes in graph/relay level, which will further decide how to lowering into different tir. With old ir builder/TE compute, we can do that. F4 in this RFC will ensure this，correct？
3) For reducing boilerplate code. F3 is a good idea. Another one is we define a tir function (with or without host python code), and we reuse it other place. We see this in F4 which foucus on conditional lowering, however I think we should define/declare it as standalone Fearture.
Looking forward to see this RFC in upstream!
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1181184802
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Junru Shao <no...@github.com.INVALID>.
I believe we are all aware the RFC is to support general IRBuilder-based metaprogramming, and with this context, I would love to address your concerns as below.
> there is no way to handle different parser rules based on context
Our design handles context-dependent parsing in a unified approach - both TIR and Relax are context-dependent, and therefore, it is a hard requirement for any parser.
Here is an example of context-dependent parsing in TIR:
```python
@T.prim_func
def f(a: T.handle):
A = T.match_buffer(a, shape=(128, 128), dtype="float32")
^^^^^^^^^^^^^^
# `match_buffer` here interprets buffer's data pointer to tvm.tir.Buffer
@T.prim_func
def f(...):
# An example from tensorization:
# https://github.com/apache/tvm/blob/c2ec95616ed52e8377c42c5f3e15621841924059/tests/python/unittest/test_tir_schedule_tensorize.py#L199-L203
with T.block(...):
A_sub = T.match_buffer(A[vi * 16 : vi * 16 + 16, vk * 16 : vk * 16 + 16], shape=(16, 16))
^^^^^^^^^^^^^^
# `match_buffer` here corresponds to `MatchBufferRegion` in TIR
# Link: https://github.com/apache/tvm/blob/c2ec95616ed52e8377c42c5f3e15621841924059/include/tvm/tir/stmt.h#L1215-L1216
```
> The example you give for a single class is one that has all static members, but what if instead it was just a regular class. It would instantiate a single instance of said class to do parsing.
> Instantiating a class would also give the benefit of being able to maintain some state while parsing which is not possible with either free functions or a class with static methods.
Thanks for bringing up the case where some state needs maintaining during parsing, and to be clear, that’s the exactly what an IRBuilder is capable of handling.
For example, considering the IRBuilder program below:
```python
def example():
with Builder():
with T.block("block"):
T.block_attr({"some_attr": some_value})
^^^^^^^^^^^^
`block_attr` depends on the `T.Block` above
```
To handle this, like every IRBuilder, the IRBuilder proposed here keeps a stack of contexts as its internal state so that information could be stored there.
The proposed parser, as a thin wrapper of the IRBuilder, does not store any extra context-dependent data other than symbol table and those already stored in the IRBuilder. As an example,
```python
@T.prim_func
def example():
...
with T.block("block"):
^^^^^^^
calls IRBuilder via `T.block`
T.block_attr({"some_attr": some_value})
^^^^^^^^^^^^
calls IRBuilder via `T.block_attr`
```
> If we want to consider future dialects, this also has the added benefit of being able to add dialects by subclassing one of the existing parsers.
Agreed that we need to consider future dialects and glad we are on the same page about this. In our RFC, IRBuilder is IR-agnostic and extensible to multiple dialects. Contextual information is stored in the IRBuilder so that metaprogramming with IRBuilder is possible.
> self.tir_parser = TIRParser()
> self.tir_parser.parse(stmt)
The problem of this design is that
- To support with IRBuilder-based metaprogramming, extra work with similar logic has to be done in the IRBuilder, which is less desirable and could cause divergence between parser and IRBuilder;
- Here is an assumption that there is a big `RelaxParser` which knows there is one and only one language “TIR” exists, which is less desirable, as you mentioned, when we want to “add dialects”. Hope this [snippet](https://github.com/tlc-pack/relax/blob/25cb42133130fb8ec75d2cc021c6c347e99f0b28/python/tvm/script/parser.py#L1372-L1391) that comes from the exact commit you provided reveals how many underlying assumptions there are with this approach.
It would be great if you could refer to Section [Parsing-time evaluation](https://github.com/yelite/tvm-rfcs/blob/tvmscript-metaprogramming/rfcs/0079-tvmscript-metaprogramming.md#parse-time-evaluation) to understand how the new parser is a thin wrap of IRBuilder.
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1203374706
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Lite Ye <no...@github.com.INVALID>.
@xqdan Thanks! For your questions,
> For compute fusion. With TE compute, it's easy to concate TE computes with producer-comsuer relation to get a fused compute. for example, conv + elemwise ops fusion. We should have similar function in TVM script. Which thread is related to this requirement?
This will be supported naturally by the F3 (TE compute). F3 will be implemented by calling the IRBuilder APIs during parse-time evaluation. So one can write multiple TE compute side by side, like
```
x = T.compute(...)
y = T.compute(...)
```
and the actual fusion is done in MetaSchedule and TIR.
> For conditional lowering. We may have some attributtes in graph/relay level, which will further decide how to lower into different tir. With old ir builder/TE compute, we can do that. F4 in this RFC will ensure this，correct？
Yes, this is correct.
> For reducing boilerplate code. F3 is a good idea. Another one is we define a tir function (with or without host python code), and we reuse it other place. We see this in F4 which foucus on conditional lowering, however I think we should define/declare it as standalone Fearture.
Mutual function call will be supported in the new TVMScript parser as well. We didn't include it in this RFC because it doesn't belong to the metaprogramming category.
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1183563695
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Junru Shao <no...@github.com.INVALID>.
> So if you define a variable in a quoted portion, you should be able to reference it in the unquoted portion
- From quoted to unquoted: Yes, that's correct.
- From unquoted to quoted: For security concern, accessing values from unquoted portion will require explicit specification if the values are not basic atomic ones (int, float, str, None, etc)
> I am mostly curious about how the interactions between quoting and unquoting will be defined and also how it would be implemented.
The mechanism here is that the core logic is in IRBuilder, and the parser is only a thin wrapper of the IRBuilder. More specifically, any `T.xxx` is a call into the IRBuilder, in both quoted and unquoted portion, and the only difference is it's run by the interpreter or the parser.
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1198677227
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Tristan Konolige <no...@github.com.INVALID>.
Ok, I think I understand things a little bit better now. Thanks for the explanation! I can see how if the IRBuilder handles all state then there is not really much of a point to having classes for the parser. It might be worthwhile to mention having a stateful parser class as an alternative in the Alternatives section (with the mentioned disadvantages around duplicating logic).
Two more questions:
Given that IRBuilder maintains all state, how does it determine what context to use for evaluating expressions. Say I have
```python
@T.prim_func
def f(a: T.Buffer[2, "float32"]):
...
```
And I want this buffer to be a `TIRBuffer` when the language is tir and a `RelaxBuffer` when the language is relax. Is this possible? From my reading, the parser and IRBuilder are only called for statements and for function calls. Is this correct? Or can you register a parser for any node in the ast (instead of just statements)?
Anther question is how does this IRBuilder approach handle error catching and reporting? Lets say we have
```python
@T.prim_func
def f(a: T.handle):
with T.block(...):
A = T.match_buffer(a, ["a string", 1], "float32")
# ^^^^^^^^^^
# An error should be reported here because the second arguments to match_buffer should be a `Sequence[Union[int, T.Var]]`
```
If I understand correctly, the `T.match_buffer` call is `eval`uated within a python context. This context calls `match_buffer` on the IRBuilder. Then the IRBuilder constructs a `MatchBuffer` object which in turn throws an error because its second argument has the wrong type. How is this error and error location then bubbled back up to the user?
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1204583687
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Junru Shao <no...@github.com.INVALID>.
Thanks for following up!
> the parser and IRBuilder are only called for statements and for function calls. Is this correct?
In fact we allow registering visitors for any Python AST constructs. For example, we could specify the behavior when visiting type annotations, function arguments, etc.
> And I want this buffer to be a `TIRBuffer` when the language is tir and a `RelaxBuffer` when the language is relax.
It’s allowed with our infrastructure, but in reality, we instead prefer to use prefix, i.e. `T.Buffer` and `R.Buffer` which is more friendly to python tooling.
> how does this IRBuilder approach handle error catching and reporting
It handles error reporting in a consistent approach with the current parser, i.e. using `DiagnosticContext` and its rendering.
In your specific case, where the error comes from a type mismatch of a particular parameter of a function call, despite being not well supported in the current parser (only supported in a small subset of buffer-related intrinsics: [`match_buffer`](https://github.com/apache/tvm/blob/df29e826290a3eba606a93b59f640e025b5cbd4d/python/tvm/script/tir/special_stmt.py#L144-L147), [`buffer_decl`](https://github.com/apache/tvm/blob/df29e826290a3eba606a93b59f640e025b5cbd4d/python/tvm/script/tir/special_stmt.py#L217-L220), [`alloc_buffer`](https://github.com/apache/tvm/blob/df29e826290a3eba606a93b59f640e025b5cbd4d/python/tvm/script/tir/special_stmt.py#L277-L279)) and the mechanism is not general enough, it is in fact one of our design consideration before starting to work on this project, as well as overall tooling of the entire TVM project. A general mechanism is that we interpret expressions in finer-grained level and do runtime type checking.
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1204646078
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Lite Ye <no...@github.com.INVALID>.
Thanks @tkonolige and @areusch for the valuable feedback! They are very helpful on getting me to better clarify the motivation behind this work and also the design detail. I updated the RFC to address them. Please take another look and let me know if you have more questions.
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1178167734
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by Junru Shao <no...@github.com.INVALID>.
@slyubomirsky Sure! Please see F1 and F2 for existing meta-programming capability (https://github.com/yelite/tvm-rfcs/blob/tvmscript-metaprogramming/rfcs/0079-tvmscript-metaprogramming.md#f1-template-metaprogramming), and see F4 for interleaving python interpreter with the parser. The quotation thing is stated in Section 3 of my response above
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1198594637
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>
Re: [apache/tvm-rfcs] [RFC] TVMScript Metaprogramming (PR #79)
Posted by "Steven S. Lyubomirsky" <no...@github.com.INVALID>.
Could you give an example of the metaprogramming? E.g., one of creating some value in ordinary Python and then referencing it from TVMScript
--
Reply to this email directly or view it on GitHub:
https://github.com/apache/tvm-rfcs/pull/79#issuecomment-1198590164
You are receiving this because you are subscribed to this thread.
Message ID: <ap...@github.com>