[OPTIONS] Remove dynamo_config['search_space'] (#342)

Continue with option cleaning. Remove `dynamo_config['search_space']`
hidet-org · Jul 22, 2024 · 4f738b7 · 4f738b7
1 parent 06ab1a0
commit 4f738b7
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 16 deletions.
diff --git a/python/hidet/cli/bench/model.py b/python/hidet/cli/bench/model.py
@@ -17,7 +17,6 @@
 
 
 class BenchModel:
-    search_space = 0
     dtype = None  # torch.float32
     tensor_core: bool = False
     disable_torch_cudnn_tf32 = False
@@ -150,9 +149,8 @@ def bench_inductor(self, mode: str) -> float:
         return self.bench_with_backend('inductor', mode=mode)
 
     def bench_hidet(self, mode, use_cuda_graph=True) -> float:
-        print('Benchmarking {} with backend {}...'.format(self, 'hidet(space={})'.format(self.search_space)))
+        print('Benchmarking {} with backend {}...'.format(self, 'hidet(mode={})'.format(mode)))
         config = hidet.torch.dynamo_config
-        config.search_space(self.search_space)
         config.use_cuda_graph(use_cuda_graph)
         config.use_tensor_core(self.tensor_core)
         return self.bench_with_backend('hidet', mode=mode)

diff --git a/python/hidet/graph/frontend/torch/dynamo_backends.py b/python/hidet/graph/frontend/torch/dynamo_backends.py
@@ -34,26 +34,21 @@
 def process_options(kwargs):
     # Default options for case mode is not passed to torch.compile()
     hidet.option.search_space(0)
-    hidet.torch.dynamo_config.search_space(0)
     hidet.torch.dynamo_config.use_cuda_graph(False)
 
     if 'mode' in kwargs:
         mode = kwargs['mode']
         if mode == 'max-autotune':
             hidet.option.search_space(2)
-            hidet.torch.dynamo_config.search_space(2)
             hidet.torch.dynamo_config.use_cuda_graph(True)
         elif mode == 'max-autotune-no-cudagraphs':
             hidet.option.search_space(2)
-            hidet.torch.dynamo_config.search_space(2)
             hidet.torch.dynamo_config.use_cuda_graph(False)
         elif mode == 'reduce-overhead':
             hidet.option.search_space(0)
-            hidet.torch.dynamo_config.search_space(0)
             hidet.torch.dynamo_config.use_cuda_graph(True)
         elif mode == 'default':
             hidet.option.search_space(0)
-            hidet.torch.dynamo_config.search_space(0)
             hidet.torch.dynamo_config.use_cuda_graph(False)
         else:
             raise ValueError(f'hidet_backend: unknown torch.compile mode={mode}')
@@ -107,7 +102,6 @@ def get_flow_graph(interpreter: Interpreter, example_inputs):
 
 
 def get_compiled_graph(flow_graph: FlowGraph):
-    search_space = dynamo_config['search_space']
     parallel_k = dynamo_config['parallel_k']
     tensor_core = dynamo_config['use_tensor_core']
     save_dir = dynamo_config['dump_graph_ir']
@@ -123,9 +117,9 @@ def get_compiled_graph(flow_graph: FlowGraph):
         graph_opt: FlowGraph = optimize(flow_graph)
         logger.info('finish optimizing the flow graph')
 
-    logger.info('schedule search space: %d', search_space)
+    logger.info('schedule search space: %d', hidet.option.get_search_space())
     logger.info('start to build the optimized computation graph')
-    cgraph: CompiledGraph = graph_opt.build(space=search_space)
+    cgraph: CompiledGraph = graph_opt.build(space=hidet.option.get_search_space())
     logger.info('finish building computation graph')
     return cgraph
 

diff --git a/python/hidet/graph/frontend/torch/dynamo_config.py b/python/hidet/graph/frontend/torch/dynamo_config.py
@@ -22,7 +22,6 @@ def dynamo_config_warning():
 
 class DynamoConfig:
     def __init__(self):
-        self._search_space: int = 0
         self._parallel_k: str = 'default'
         self._use_cuda_graph: bool = True
         self._use_tensor_core: bool = False
@@ -39,7 +38,6 @@ def reset(self):
         """
         Reset the configuration to the default values
         """
-        self._search_space: int = 0
         self._parallel_k: str = 'default'
         self._use_cuda_graph: bool = True
         self._use_tensor_core: bool = False
@@ -65,8 +63,7 @@ def search_space(self, level: int = 2):
         level: int
             The search space level.
         """
-        self._search_space = level
-        return self
+        dynamo_config_warning()
 
     def use_tensor_core(self, flag=True):
         """

diff --git a/scripts/regression/model_performance.py b/scripts/regression/model_performance.py
@@ -25,7 +25,7 @@ def enable_compiled_server():
     pass
 
 def setup_hidet_flags(dtype):
-    hidet.torch.dynamo_config.search_space(2)
+    hidet.option.search_space(2)
     hidet.torch.dynamo_config.use_tensor_core(True)
     hidet.torch.dynamo_config.use_cuda_graph(True)
     hidet.torch.dynamo_config.dump_graph_ir("./graph_ir")