remove code comments and add type hints

Signed-off-by: Bangtian Liu <[email protected]>

tuner/tuner/candidate_gen_test.py

@@ -27,10 +27,7 @@ def tuner_ctx() -> Generator[common.TunerContext, None, None]:
     from logging import Logger
     from unittest.mock import MagicMock
 
-    # Mock the logger
     mock_logger = MagicMock(spec=Logger)
-
-    # Use TunerContext with the mocked logger
     with common.TunerContext(logger=mock_logger) as ctx:
         yield ctx
 

tuner/tuner/common.py

@@ -4,10 +4,10 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from __future__ import annotations
 import logging
 from dataclasses import astuple, dataclass, field
 from enum import Enum
+from types import TracebackType
 from typing import Optional
 from typing import Any
 from typing_extensions import Literal
@@ -46,13 +46,17 @@ def __init__(self, logger: Optional[logging.Logger] = None):
         )  # Default to "tune" logger
         self.type: CommonTypes = CommonTypes(self.mlir_ctx)
 
-    def __enter__(self) -> TunerContext:
+    def __enter__(self) -> "TunerContext":
         self.mlir_ctx.__enter__()
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback) -> Literal[False]:
-        self.mlir_ctx.__exit__(exc_type, exc_value, traceback)
-        return False
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        return self.mlir_ctx.__exit__(exc_type, exc_value, traceback)
 
 
 class DispatchKind(Enum):

tuner/tuner/common_test.py

@@ -23,10 +23,7 @@ def tuner_ctx() -> Generator[common.TunerContext, None, None]:
     from logging import Logger
     from unittest.mock import MagicMock
 
-    # Mock the logger
     mock_logger = MagicMock(spec=Logger)
-
-    # Use TunerContext with the mocked logger
     with common.TunerContext(logger=mock_logger) as ctx:
         yield ctx
 

tuner/tuner/dispatch_constraints_test.py

@@ -25,10 +25,7 @@ def tuner_ctx() -> Generator[common.TunerContext, None, None]:
     from logging import Logger
     from unittest.mock import MagicMock
 
-    # Mock the logger
     mock_logger = MagicMock(spec=Logger)
-
-    # Use TunerContext with the mocked logger
     with common.TunerContext(logger=mock_logger) as ctx:
         yield ctx
 

tuner/tuner/libtuner.py

@@ -642,15 +642,14 @@ def generate_candidate_specs(
         # source mlir.
         mlir_text = candidate_gen.strip_compilation_info(path_config.template_mlir)
         mlir_module = dispatch_parser.parse_mlir(mlir_text, tuning_client.tuner_context)
-        with tuning_client.tuner_context as tuner_context:
-            logging.debug("Captured messages from candidate_gen.py:")
-            config_specs: list[ir.Module] = candidate_gen.generate_configs_and_td_specs(
-                input_module=mlir_module,
-                tuner_context=tuner_context,
-                limit=args.num_candidates,
-                num_subgroups=args.num_subgroups,
-                codegen_pipeline=get_iree_codegen_pipeline(args.codegen_pipeline),
-            )
+        logging.debug("Captured messages from candidate_gen.py:")
+        config_specs: list[ir.Module] = candidate_gen.generate_configs_and_td_specs(
+            input_module=mlir_module,
+            tuner_context=tuning_client.tuner_context,
+            limit=args.num_candidates,
+            num_subgroups=args.num_subgroups,
+            codegen_pipeline=get_iree_codegen_pipeline(args.codegen_pipeline),
+        )
         logging.debug("candidate_gen.py ends")
         handle_error(
             condition=(len(config_specs) <= 1), msg="Failed to generate any candidates"