Yeah, to fully replace IRBuilder, there are still some missing elements of TVM script: meta programming and hygiene macros. Let’s consider the support after the upstreaming is done 
CC @tqchen
@junrushao If by “meta programming” you mean an ability to call python function from script to generate other code, and embed the generated code into the calling context, then YES, we absolutely need this! I think it’s called “splicing” or “unquote” in the literature.
The lack of such feature is what turned me away from hybrid script. Since there is basically no abstraction to compose multiple code to generate bigger code, I had to manually duplicate a lot of code. The resulting code is super ugly (just look at various shape functions). For example, tvm/_transform.py at 8e23806d2d522b71979d0a2730b38cc5c3bf6185 · apache/tvm · GitHub (this one I wrote)
For examples of kind of programming style we’ve developed for ir builder, see topi/cuda/sort.py or topi/cuda/nms.py. cc @mbrookhart
1 Like
@masahi Yeah. We should allow to embed IR fragments, functions that produce IR fragments, or replace some tokens with caller specified IR fragments into the script
2 Likes
Just curious when will such a big feature mainline, is there a initial planning on it? Thanks, can’t wait for long to use it. Haha.
The upstreaming is in progress. You can track it at [RFC] TensorIR Scheduling Tracking Issue · Issue #7527 · apache/tvm · GitHub.
As you mentioned that the project is huge, it still needs some time (maybe 2-3 months) to finish. But we are trying our best.
4 Likes
One issue in old schedule ops is we can not get the accurate bouds with inferbound, what will it be like in new schedule system? thanks.
Is this for merging two reduction stages into one stage? Thank you very much!
It’s not. merge_reduction is designed for merging init part and update part into one block with if branch
Hi, I am a learner at compiler stack, I have a few questions of Tensor IR
- instead of TE lower to TIR, what is the process from Relay directly to TIR? it will be better if you have example or tutorial
- you mentioned that “Now, there is no stage during the schedule. Rather than lowering the schedule into TIR, we directly mutate the TIR itself.”. is this means there won’t have concept of “lower”， since no stage?
- do you have any details of difference between new Tensor IR and former TIR(like TVM 0.6) I am looking forward for your reply. Thank you!
1 Like
@crazyjackie1993 Thanks for the questions!
Terminologies: let call:
- S-TIR: The schedulable TIR proposed in this RFC
- NS-TIR: the original non-scheduleable TIR
what is the process from Relay directly to TIR
With this proposal landed, There are two ways of lowering Relay:
- Relay => TE compute w/ schedule => S-TIR => NS-TIR
- Relay => S-TIR => NS-TIR
Effectively, TE compute becomes a syntactic sugar to build S-TIR
is this means there won’t have concept of “lower”, since no stage?
To make sure I understand the question, are you asking why there is no “lower”?
That’s correct. There is no indirection like the TE, where we need to lower the schedule tree into the actual AST after all scheduling is done. Every schedule primitive in S-TIR is effectively a transformation from S-TIR to S-TIR, and the schedule class only provides necessary auxiliary data structures.
do you have any details of difference between new Tensor IR and former TIR
- S-TIR is scheduleable while NS-TIR is not;
- There is an IR construct called
Blockin S-TIR, which is introduced in full details in RFC, and is the core concept of how the system works; While in NS-TIR there is not.
3 Likes
@junrushao thank you for your reply, your explanation helps a lot. but I am a little confused about S-TIR & NS TIR.
can I understand like this:
- TF/ONNX/PyTorch => Relay => (with or without TE) => S-TIR => schedule => S-TIR => NS-TIR => code Gen
I still have two father more question about S-TIR:
- can S-TIR fully describe an op(e.g. conv), which can replace the TE’s work?
- is there any interaction process between S-TIR & NS-TIR?
I am also concerned about if Tensor IR development was finished or is still in process? I am really interensted in this idea, so i’d also like to read some code to understand it better.
S-TIR is a complete IR for the PrimFunc, while TE is just a DSL abstraction. So that S-TIR can fully describe ops which TE can describe. Here is a conv example in S-TIR (The syntax may change as we are discussing the final syntax
# TE
A = te.placeholder((16, 16, 14, 14), name="A")
W = te.placeholder((16, 3, 3, 32), name="W")
Apad = te.compute(
(batch, in_channel, size + 2, size + 2),
lambda nn, cc, yy, xx: tvm.tir.if_then_else(
tvm.tir.all(yy >= 1, yy - 1 < size, xx >= 1, xx - 1 < size),
A[nn, cc, yy - 1, xx - 1],
0.0,
),
name="Apad",
)
rc = te.reduce_axis((0, in_channel), name="rc")
ry = te.reduce_axis((0, kernel), name="ry")
rx = te.reduce_axis((0, kernel), name="rx")
B = te.compute(
(batch, out_channel, size, size),
lambda nn, ff, yy, xx: te.sum(
Apad[nn, rc, yy + ry, xx + rx] * W[rc, ry, rx, ff], axis=[rc, ry, rx]
),
name="B",
)
# S-TIR
@tvm.script.tir
def tir_conv2d(a: ty.handle, w: ty.handle, b: ty.handle) -> None:
A = tir.match_buffer(a, [16, 16, 14, 14])
W = tir.match_buffer(w, [16, 3, 3, 32])
B = tir.match_buffer(b, [16, 32, 14, 14])
Apad = tir.alloc_buffer([16, 16, 16, 16])
for n, c, y, x in tir.grid(16, 16, 16, 16):
with tir.block([16, 16, 16, 16], "Apad") as [nn, cc, yy, xx]:
Apad[nn, cc, yy, xx] = tir.if_then_else(
yy >= 1 and yy - 1 < 14 and xx >= 1 and xx - 1 < 14,
A[nn, cc, yy - 1, xx - 1],
0.0,
dtype="float32",
)
for n, f, y, x, kc, ky, kx in tir.grid(16, 32, 14, 14, 16, 3, 3):
with tir.block(
[16, 32, 14, 14, tir.reduce_axis(0, 16), tir.reduce_axis(0, 3), tir.reduce_axis(0, 3)], "B"
) as [nn, ff, yy, xx, rc, ry, rx]:
with tir.init():
B[nn, ff, yy, xx] = 0.0
B[nn, ff, yy, xx] += Apad[nn, rc, yy + ry, xx + rx] * W[rc, ry, rx, ff]
We have some passes to transform S-TIR to NS-TIR
It has nearly finished upstream. Please see https://github.com/apache/tvm/issues/7527. It would be an experimental feature in the v0.8 release(in one or two months)
2 Likes
Thanks for reply. I look forward to seeing your work soon. 
Amazing progress here! So great!
I have a short question, though. When trying out the Blitz tutorial, I didn’t get how to use constant to define e.g. shapes or types
@tvm.script.ir_module
class MyModule:
@T.prim_func
def main(a: T.handle, b: T.handle):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
A = T.match_buffer(a, (8,), dtype="float32")
B = T.match_buffer(b, (8,), dtype="float32")
...
I’d be interested how to do something like this:
@tvm.script.ir_module
class MyModule:
@T.prim_func
def main(a: T.handle, b: T.handle, x:int, dt: str):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
A = T.match_buffer(a, (x,), dtype=dt)
B = T.match_buffer(b, (x,), dtype=dt)
if x < cond1:
...
else:
....
...
I get an error message whenever using a construct similar to this.
Is this possible. If yes, it would be great if you could give a short example.
I’m not sure if there’s any easy way to represent string types, but integers can be represented as T.int32 or T.int64, etc. So in your example, you can probably use the below syntax:
@tvm.script.ir_module
class MyModule:
@T.prim_func
def main(a: T.handle, b: T.handle, x: T.int32):
T.func_attr({"global_symbol": "main", "tir.noalias": True})
A = T.match_buffer(a, (x,), dtype=dt)
B = T.match_buffer(b, (x,), dtype=dt)
...Note that the TVMScript is a represent for TensorIR (TVM IR) rather than part of runnable python code. That means we can only use TVM data structure (e.g. T.int32, T.handle) rather than python type (e.g. int, str).
Unfortunately, TVM does not have a type hint for type struct (str is not the type struct). So we can not use constants to define buffer types. We may support it in the future if it is needed.
Thanks for the replies! I’m not 100% sure if I understood everything correct.
The feature I’m looking for is to replace TE by TensorIR. Here an example (from the TVM repository) with lots of parameters defined as Python constants.
conv = te.compute(
(1, ofm_height, ofm_width, ofm_channels),
lambda nn, hh, ww, cc: te.sum(
dmaed_ifm(
nn, hh * stride_h + rh * dilation_h, ww * stride_w + rw * dilation_w, rc
).astype(ifm.dtype)
* weight[cc, rh, rw, rc].astype(ifm.dtype)
+ (scale_bias[cc, 0] * scale_bias[cc, 9]).astype(ifm.dtype),
axis=[rh, rw, rc],
),
name="ethosu_conv2d",
attrs=conv2d_attrs,
)
The way I interpret the figure from the Blitz tutorial is that I can replace TE by TVMScript(that represents Tensor IR). However, I don’t see the way to use (Python) variables to influence the schedule here.
You have indicated that it might not be possible yet using TVMScript. Is there another way of bringing the shapes (or type or other constants) of an operator into the TensorIR AST? At least from the tutorial it was not directly clear for me.
IIUC, TVMscript cannot handle such a case. This’s also my question. So I try to modify the TVMscript. The final python snippet is similar as the torchscript:
class ScriptModule(object):
def __init__(self, x: ty.int32):
self.x = x
def main(self, a:ty.handle, b:ty.handle, c:ty.handle)->None:
T.func_attr({"global_symbol": "main", "tir.noalias": True})
A = T.match_buffer(a, (self.x,), dtype=dt)
B = T.match_buffer(b, (self.x,), dtype=dt)
I recommend that you could read the code about torchscript.
I see your pain. Currently we cannot programmatically construct (meta-program) TVMScript from Python. See also the discussion in [RFC] Hybrid Script Support for TIR - #34 by masahi
1 Like
IMHO, we can do the same staff like torch script. But like you mentioned before, how to modulize the function may be the main obstacle. Cross function calles may affect codegen, split device and host, etc.