From 230f110c2fa3f56b012c01cc0a296e901636fe2e Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Sun, 30 Nov 2025 04:34:48 +0000 Subject: [PATCH 1/5] [Bugfix] Prevent segfault when instantiating abstract SearchStrategy Add __init__ method to SearchStrategy class to prevent direct instantiation of the abstract class. This raises a TypeError with a helpful error message instead of causing a segmentation fault when SearchStrategy() is called directly or passed to TuneContext. Also add additional check in TuneContext.__init__ to ensure abstract SearchStrategy instances are not used. Fixes #18268 --- .../search_strategy/search_strategy.py | 19 ++++++++++ python/tvm/meta_schedule/tune_context.py | 8 ++++ .../test_meta_schedule_search_strategy.py | 37 +++++++++++++++++++ 3 files changed, 64 insertions(+) diff --git a/python/tvm/meta_schedule/search_strategy/search_strategy.py b/python/tvm/meta_schedule/search_strategy/search_strategy.py index 75b45cf424c3..5f1b6641cb5a 100644 --- a/python/tvm/meta_schedule/search_strategy/search_strategy.py +++ b/python/tvm/meta_schedule/search_strategy/search_strategy.py @@ -87,6 +87,25 @@ class SearchStrategy(Object): ], ] + def __init__(self, *args, **kwargs): + """Prevent direct instantiation of abstract SearchStrategy class. + + SearchStrategy is an abstract class and cannot be directly instantiated. + Use SearchStrategy.create() or a concrete subclass instead. + """ + # Check if this is a direct instantiation of SearchStrategy (not a subclass) + # Subclasses will have a different type + if type(self) is SearchStrategy: + raise TypeError( + "Cannot instantiate abstract class SearchStrategy. " + "Use SearchStrategy.create() with a valid strategy type " + "(e.g., 'evolutionary', 'replay-trace', 'replay-func') " + "or use a concrete subclass instead." + ) + + # Allow parent class initialization to proceed for proper subclasses + super().__init__(*args, **kwargs) + def _initialize_with_tune_context(self, context: "TuneContext") -> None: """Initialize the search strategy with tuning context. diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py index c3f496265a97..0cab52dba35a 100644 --- a/python/tvm/meta_schedule/tune_context.py +++ b/python/tvm/meta_schedule/tune_context.py @@ -123,6 +123,14 @@ def __init__( if search_strategy is not None: if not isinstance(search_strategy, SearchStrategy): search_strategy = SearchStrategy.create(search_strategy) + # Additional check: ensure it's not the abstract SearchStrategy class itself + elif type(search_strategy) is SearchStrategy: + raise TypeError( + "Cannot use abstract SearchStrategy class directly. " + "Use SearchStrategy.create() with a valid strategy type " + "(e.g., 'evolutionary', 'replay-trace', 'replay-func') " + "or use a concrete subclass instead." + ) if logger is None: logger = get_logger(__name__) if not isinstance(num_threads, int): diff --git a/tests/python/meta_schedule/test_meta_schedule_search_strategy.py b/tests/python/meta_schedule/test_meta_schedule_search_strategy.py index 29c20ced0488..e79a5dd69ac3 100644 --- a/tests/python/meta_schedule/test_meta_schedule_search_strategy.py +++ b/tests/python/meta_schedule/test_meta_schedule_search_strategy.py @@ -306,9 +306,46 @@ def __str__(self) -> str: assert candidates is None +def test_search_strategy_abstract_class_instantiation(): + """Test that directly instantiating abstract SearchStrategy raises TypeError instead of segfault.""" + from tvm.meta_schedule import SearchStrategy + from tvm.target import Target + from tvm.meta_schedule import TuneContext + from tvm.script import ir as I, relax as R + + # A minimal Relax IRModule (identity) + @I.ir_module + class M: + @R.function + def main(x: R.Tensor((1,), "float32")) -> R.Tensor((1,), "float32"): + return x + + # Test that direct instantiation raises TypeError + # This prevents segfault when SearchStrategy() is called directly + with pytest.raises(TypeError, match="Cannot instantiate abstract class SearchStrategy"): + SearchStrategy() + + # Test that TuneContext with SearchStrategy() raises TypeError + # The error should occur when trying to create SearchStrategy() instance + with pytest.raises(TypeError, match="Cannot instantiate abstract class SearchStrategy"): + TuneContext( + mod=M, + target=Target("llvm"), + search_strategy=SearchStrategy(), # This should fail before reaching TuneContext + ) + + # Test that SearchStrategy.create() works correctly + strategy = SearchStrategy.create("evolutionary") + assert strategy is not None + assert isinstance(strategy, SearchStrategy) + # Verify it's not the abstract class itself + assert type(strategy) is not SearchStrategy + + if __name__ == "__main__": test_meta_schedule_replay_func(ms.search_strategy.ReplayFunc) test_meta_schedule_replay_func(ms.search_strategy.ReplayTrace) test_meta_schedule_evolutionary_search() test_meta_schedule_evolutionary_search_early_stop() test_meta_schedule_evolutionary_search_fail_init_population() + test_search_strategy_abstract_class_instantiation() From 075fcb06f89a781208b91349e54d9561b8d21b0b Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Sun, 30 Nov 2025 04:34:48 +0000 Subject: [PATCH 2/5] Use __new__ instead of __init__ for abstract class prevention --- .../search_strategy/search_strategy.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/python/tvm/meta_schedule/search_strategy/search_strategy.py b/python/tvm/meta_schedule/search_strategy/search_strategy.py index 5f1b6641cb5a..ee068ff7b5b6 100644 --- a/python/tvm/meta_schedule/search_strategy/search_strategy.py +++ b/python/tvm/meta_schedule/search_strategy/search_strategy.py @@ -87,24 +87,23 @@ class SearchStrategy(Object): ], ] - def __init__(self, *args, **kwargs): + def __new__(cls, *args, **kwargs): """Prevent direct instantiation of abstract SearchStrategy class. SearchStrategy is an abstract class and cannot be directly instantiated. Use SearchStrategy.create() or a concrete subclass instead. """ - # Check if this is a direct instantiation of SearchStrategy (not a subclass) - # Subclasses will have a different type - if type(self) is SearchStrategy: + if cls is SearchStrategy: raise TypeError( "Cannot instantiate abstract class SearchStrategy. " "Use SearchStrategy.create() with a valid strategy type " "(e.g., 'evolutionary', 'replay-trace', 'replay-func') " "or use a concrete subclass instead." ) - - # Allow parent class initialization to proceed for proper subclasses - super().__init__(*args, **kwargs) + return super().__new__(cls) +======= + return super().__new__(cls) +>>>>>>> 2dba6e03d (Use __new__ instead of __init__ for abstract class prevention) def _initialize_with_tune_context(self, context: "TuneContext") -> None: """Initialize the search strategy with tuning context. From 286d7de8c00ff6e3f1f2aa2dda7e1fa9af628a39 Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Sun, 30 Nov 2025 04:51:52 +0000 Subject: [PATCH 3/5] Fix linting issues: remove merge conflict markers and trailing spaces --- python/tvm/meta_schedule/search_strategy/search_strategy.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/python/tvm/meta_schedule/search_strategy/search_strategy.py b/python/tvm/meta_schedule/search_strategy/search_strategy.py index ee068ff7b5b6..1b8302e55cbb 100644 --- a/python/tvm/meta_schedule/search_strategy/search_strategy.py +++ b/python/tvm/meta_schedule/search_strategy/search_strategy.py @@ -89,7 +89,7 @@ class SearchStrategy(Object): def __new__(cls, *args, **kwargs): """Prevent direct instantiation of abstract SearchStrategy class. - + SearchStrategy is an abstract class and cannot be directly instantiated. Use SearchStrategy.create() or a concrete subclass instead. """ @@ -101,9 +101,6 @@ def __new__(cls, *args, **kwargs): "or use a concrete subclass instead." ) return super().__new__(cls) -======= - return super().__new__(cls) ->>>>>>> 2dba6e03d (Use __new__ instead of __init__ for abstract class prevention) def _initialize_with_tune_context(self, context: "TuneContext") -> None: """Initialize the search strategy with tuning context. From a278d6171971510f5c434c03bd96c961605652c2 Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Sun, 30 Nov 2025 05:02:11 +0000 Subject: [PATCH 4/5] Fix pylint errors: add disable comments for necessary type checks --- python/tvm/meta_schedule/search_strategy/search_strategy.py | 4 ++-- python/tvm/meta_schedule/tune_context.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/python/tvm/meta_schedule/search_strategy/search_strategy.py b/python/tvm/meta_schedule/search_strategy/search_strategy.py index 1b8302e55cbb..cfb45dafdeb2 100644 --- a/python/tvm/meta_schedule/search_strategy/search_strategy.py +++ b/python/tvm/meta_schedule/search_strategy/search_strategy.py @@ -87,7 +87,7 @@ class SearchStrategy(Object): ], ] - def __new__(cls, *args, **kwargs): + def __new__(cls, *args, **kwargs): # pylint: disable=unused-argument """Prevent direct instantiation of abstract SearchStrategy class. SearchStrategy is an abstract class and cannot be directly instantiated. @@ -100,7 +100,7 @@ def __new__(cls, *args, **kwargs): "(e.g., 'evolutionary', 'replay-trace', 'replay-func') " "or use a concrete subclass instead." ) - return super().__new__(cls) + return super().__new__(cls) # pylint: disable=no-value-for-parameter def _initialize_with_tune_context(self, context: "TuneContext") -> None: """Initialize the search strategy with tuning context. diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py index 0cab52dba35a..35a8d468a75c 100644 --- a/python/tvm/meta_schedule/tune_context.py +++ b/python/tvm/meta_schedule/tune_context.py @@ -124,7 +124,8 @@ def __init__( if not isinstance(search_strategy, SearchStrategy): search_strategy = SearchStrategy.create(search_strategy) # Additional check: ensure it's not the abstract SearchStrategy class itself - elif type(search_strategy) is SearchStrategy: + # Use type() for exact type check (not isinstance which would match subclasses) + elif type(search_strategy) is SearchStrategy: # pylint: disable=unidiomatic-typecheck raise TypeError( "Cannot use abstract SearchStrategy class directly. " "Use SearchStrategy.create() with a valid strategy type " From 4a55738ba1b1d4c6d14e026d8c45bd0962a41266 Mon Sep 17 00:00:00 2001 From: Dayuxiaoshui <792179245@qq.com> Date: Sun, 30 Nov 2025 09:35:38 +0000 Subject: [PATCH 5/5] Use TIR module instead of Relax in abstract class instantiation test --- .../test_meta_schedule_search_strategy.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/tests/python/meta_schedule/test_meta_schedule_search_strategy.py b/tests/python/meta_schedule/test_meta_schedule_search_strategy.py index e79a5dd69ac3..04a6e187a6a7 100644 --- a/tests/python/meta_schedule/test_meta_schedule_search_strategy.py +++ b/tests/python/meta_schedule/test_meta_schedule_search_strategy.py @@ -311,14 +311,6 @@ def test_search_strategy_abstract_class_instantiation(): from tvm.meta_schedule import SearchStrategy from tvm.target import Target from tvm.meta_schedule import TuneContext - from tvm.script import ir as I, relax as R - - # A minimal Relax IRModule (identity) - @I.ir_module - class M: - @R.function - def main(x: R.Tensor((1,), "float32")) -> R.Tensor((1,), "float32"): - return x # Test that direct instantiation raises TypeError # This prevents segfault when SearchStrategy() is called directly @@ -326,12 +318,14 @@ def main(x: R.Tensor((1,), "float32")) -> R.Tensor((1,), "float32"): SearchStrategy() # Test that TuneContext with SearchStrategy() raises TypeError - # The error should occur when trying to create SearchStrategy() instance + # The error should occur when trying to create SearchStrategy() instance in the function call + # Since SearchStrategy() fails in __new__, it will fail before TuneContext.__init__ is called with pytest.raises(TypeError, match="Cannot instantiate abstract class SearchStrategy"): + # This will fail when evaluating SearchStrategy() as an argument TuneContext( - mod=M, + mod=Matmul, # Use the existing Matmul module from the test file target=Target("llvm"), - search_strategy=SearchStrategy(), # This should fail before reaching TuneContext + search_strategy=SearchStrategy(), # This should fail in __new__ before reaching TuneContext ) # Test that SearchStrategy.create() works correctly