Hi,
I recently begin to take to look for new dynamic shape feature supported by relax. While checking testcases as tests/python/relax/test_codegen_cutlass.py, I don’t see any dynamic shape infer has been made as a show case, or am I missing?
So what I want to know is for one simple conv2d forward, how could I test its dynamic shape feature with relax and use backend like cutlass?
Thx
I’m also new to relax. I think currently the best practice for dynamic shape in relax is MLC-LLM that uses it to handle the dynamic sequence length. As for your question, I find that there are some symbolic shapes like _vars[“a”] in tests/python/relax/test_codegen_cutlass.py. Relax uses it to represent the dynamic axis in the shape.
MLC-LLM is too complex to use as a start point to investigate the dynamic shape feature…
I try to modify test case in test_codegen_cutlass.py as below, and enable dynamic shape as below.
@tvm.script.ir_module
class Conv2dBiasReLU:
@R.function
def main(
data: R.Tensor((1, tvm.tir.expr.Var("a", "int64"), 400, 16), "float32"),
weight: R.Tensor((32, 3, 3, 16), "float32"),
bias: R.Tensor((1, 1, 1, 32), "float32"),
):
with R.dataflow():
conv1 = R.nn.relu(
R.nn.conv2d(data, weight, padding=(1, 1), data_layout="NHWC", kernel_layout="OHWI")
+ bias,
)
R.output(conv1)
return conv1
mod = Conv2dBiasReLU
mod = partition_for_cutlass(mod)
codegen_pass = relax.transform.RunCodegen({"cutlass": {"sm": 80, "find_first_valid": True}})
mod = codegen_pass(mod)
But error occurs:
op_name, op_def, _ = self.conv2d_profiler.profile(
File "/data/dev/tvm_dev/mlc-llm/3rdparty/tvm/python/tvm/contrib/cutlass/gen_conv2d.py", line 356, in profile
op = self.select_op(
File "/data/dev/tvm_dev/mlc-llm/3rdparty/tvm/python/tvm/contrib/cutlass/gen_conv2d.py", line 309, in select_op
args = (
TypeError: %d format: a real number is required, not Var
So it seems to me that build stage I still cannot make the data shape as dynamic?
You may learn dynamic shape support from my latest Unity branch (going to PR soon), which implements Llama from scratch in ~100 lines: https://github.com/junrushao/tvm/blob/b89e975eb994800f052809f59cf0178814bd7700/main.py#L109-L213
there is an official dynamic-shape example with dynamic_strided_slice() op in: tvm/tests/python/relax/test_e2e_op_dynamic.py at unity · apache/tvm (github.com)
The dynamic shape for Conv2D in cutlass BYOC hasn’t been supported yet, by checking the commit history. The test also doesn’t cover dynamic shape for Conv2D.
But matmul is supported now:
I see. I checked the matmul test, and yes it support dynamic shape well.