Skip to content

Commit cd3e107

Browse files
authored
[Unity] Fix getting shapes for cutlass BYOC kernels (#14411)
1 parent 9a3ec23 commit cd3e107

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

python/tvm/contrib/cutlass/build.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from tvm import relax, relay, runtime
2929
from tvm._ffi.registry import register_func
3030
from tvm.contrib.nvcc import get_cuda_version
31+
from tvm.topi.utils import get_const_tuple
3132

3233
from .gen_conv2d import CutlassConv2DProfiler
3334
from .gen_gemm import CutlassGemmProfiler
@@ -545,7 +546,7 @@ def _extract_relax_function_signature(f):
545546

546547
for i, arg in enumerate(f.params):
547548
sinfo = arg.struct_info
548-
signature["arg%d_shape" % i] = list(sinfo.shape)
549+
signature["arg%d_shape" % i] = get_const_tuple(sinfo.shape)
549550
signature["arg%d_dtype" % i] = sinfo.dtype
550551

551552
ret_sinfo = f.ret_struct_info

0 commit comments

Comments
 (0)