How to use var to declare match_buffer and alloc_buffer in tvm.script?

Questions
#1

image

I can use const int to declare the match buffer. but I want to generate the match buffer by parameters such as shape_a how could I do？

#2

Hi,

If I understand your question correctly, you would like to declare match_buffers with variable shapes (please feel free to correct me if I’m mistaken). I think you can declare a variable using T.var(...) and then use that var in the match syntax.

For Example:

m = T.var("int32")
n = T.var("int32")
val = T.match_buffer(var_val, [m, n], dtype="float32")

This is taken from the test code written here.

In general, I think the test files in tests/python/ are probably the best resource after documentation to find good example use cases of the different features of TVM API (Just a hint since I took a long time to realize this :slight_smile:).

Thanks

#3

Since PR https://github.com/apache/tvm/pull/11097 is merged, we can use Python var in TVMScript.

from tvm.script import tir as T


def generator(shape=(128, 128)):
    @T.prim_func
    def func(a: T.handle):
        A = T.match_buffer(a, shape, "float32")
        T.evaluate(0)

    return func


print(generator().script())
1 Like