From 6f58f472bb9ad53dee659e1b0d191335e89db1eb Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Mon, 27 Mar 2023 17:01:54 -0700 Subject: [PATCH] [Unity] Fix getting shapes for cutlass BYOC kernels --- python/tvm/contrib/cutlass/build.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py index 47bdcaa790b8..d9eefd34a38b 100644 --- a/python/tvm/contrib/cutlass/build.py +++ b/python/tvm/contrib/cutlass/build.py @@ -28,6 +28,7 @@ from tvm import relax, relay, runtime from tvm._ffi.registry import register_func from tvm.contrib.nvcc import get_cuda_version +from tvm.topi.utils import get_const_tuple from .gen_conv2d import CutlassConv2DProfiler from .gen_gemm import CutlassGemmProfiler @@ -545,7 +546,7 @@ def _extract_relax_function_signature(f): for i, arg in enumerate(f.params): sinfo = arg.struct_info - signature["arg%d_shape" % i] = list(sinfo.shape) + signature["arg%d_shape" % i] = get_const_tuple(sinfo.shape) signature["arg%d_dtype" % i] = sinfo.dtype ret_sinfo = f.ret_struct_info