From fd7c81de3b2f1e023351b10478e647fdbf367acc Mon Sep 17 00:00:00 2001 From: Abhikrant Sharma Date: Tue, 9 Jul 2024 22:25:38 +0530 Subject: [PATCH] [TIR][Schedule] Remove `@type_check` for `set_axis_separator` (#17134) [TIR][Schedule] Remove @type_check decorator for set_axis_separator The decorator is not allowing types like Array to be passed to set_axis_separator directive. The FFI has a type checking implemented internally making this redundant. --- python/tvm/relax/transform/legalize_ops/manipulate.py | 1 - python/tvm/tir/schedule/schedule.py | 1 - 2 files changed, 2 deletions(-) diff --git a/python/tvm/relax/transform/legalize_ops/manipulate.py b/python/tvm/relax/transform/legalize_ops/manipulate.py index 4d30b97f6467..1efa78c069ad 100644 --- a/python/tvm/relax/transform/legalize_ops/manipulate.py +++ b/python/tvm/relax/transform/legalize_ops/manipulate.py @@ -213,7 +213,6 @@ def set_axis_sep(axis_sep: list, sch: tir.schedule, buffer_type: str): sch.transform_layout(primfunc_name, ("write", 0), index_map, pad_value) set_axis_sep(axis_separators, sch, "write") if input_axis_separators is not None: - input_axis_separators = [int(sep) for sep in input_axis_separators] set_axis_sep(input_axis_separators, sch, "read") gvar = bb.add_func(sch.mod["main"], primfunc_name) output_shape = index_map.map_shape(list(call_args[0].struct_info.shape)) diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index f477a0f11233..4127266da7e2 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -3490,7 +3490,6 @@ def after_transform_block_layout( self, block, index_map ) - @type_checked def set_axis_separator( self, block: Union[BlockRV, str],