没有说中文的么,哈哈,看着英文有点小费劲。我也在跟踪relax的开发,期待有可跑的dynamic-shape demo
I am also interested in dynamic shape feature in TVM, but did not find a suitable example or api interface.
checkout https://github.com/mlc-ai/mlc-llm/, whichi should contains examples about dynamic shape via tvm unity
Any detail code file or link to find the dynamic shape example? Thanks~
def matmul_nn(
M,
N,
K,
in_dtype="float16",
out_dtype="float16",
accum_dtype="float16",
with_bias=False,
):
if not isinstance(M, int):
M = tvm.te.var("m")
A = te.placeholder((M, K), name="A", dtype=in_dtype)
B = te.placeholder((K, N), name="B", dtype=in_dtype)
Bias = te.placeholder((N,), name="Bias", dtype=in_dtype)
# Describe the matrix multiplication in TE
k = te.reduce_axis((0, K), name="k")
C = te.compute(
(M, N),
lambda i, j: te.sum(A[i, k].astype(accum_dtype) * B[k, j].astype(accum_dtype), axis=k),
name="C",
)
last_output = C
if accum_dtype != out_dtype:
D = te.compute((M, N), lambda i, j: C[i, j].astype(out_dtype), name="D")
last_output = D
if with_bias:
E = te.compute((M, N), lambda i, j: last_output[i, j] + Bias[j], name="E")
last_output = E
args = [A, B, Bias, last_output] if with_bias else [A, B, last_output]
func = te.create_prim_func(args)
return tvm.IRModule.from_expr(func)
Take te as an example, M = tvm.te.var("m")
enables dynamic for M dimension, which also works for tir script.