How to use dynamic shape in relax?

Development unity
1

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

1 Like
2

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.

1 Like
3

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?

1 Like
4

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

2 Likes
5

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)

1 Like
6

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:

7

I see. I checked the matmul test, and yes it support dynamic shape well.

1 Like