diff --git a/test/infinicore/framework/base.py b/test/infinicore/framework/base.py index e57e8ec11..74586a1ea 100644 --- a/test/infinicore/framework/base.py +++ b/test/infinicore/framework/base.py @@ -3,21 +3,23 @@ Contains TestConfig, TestRunner, and BaseOperatorTest classes. """ -import torch -import infinicore import traceback from abc import ABC, abstractmethod -from .results import CaseResult -from .datatypes import to_torch_dtype, to_infinicore_dtype +import torch + +import infinicore + +from .benchmark import BenchmarkUtils from .devices import InfiniDeviceNames, torch_device_map -from .tensor import TensorSpec, TensorInitializer +from .results import CaseResult +from .tensor import TensorSpec +from .utils.compare_utils import create_test_comparator from .utils.tensor_utils import ( clone_torch_tensor, infinicore_tensor_from_torch, + synchronize_device, ) -from .utils.compare_utils import create_test_comparator -from .benchmark import BenchmarkUtils class TestConfig: @@ -49,9 +51,7 @@ def __init__(self, test_cases, test_config): self.failed_tests = [] self.skipped_tests = [] # Track skipped tests (both operators not implemented) self.partial_tests = [] # Track partial tests (one operator not implemented) - self.passed_tests = ( - [] - ) # Track passed tests (both operators implemented and passed) + self.passed_tests = [] # Track passed tests (both operators implemented and passed) # Add benchmark timing statistics self.benchmark_times = { "torch_host_total": 0.0, @@ -76,9 +76,9 @@ def run_tests(self, devices, test_func, test_type="Test"): bool: True if no tests failed, False otherwise """ for device in devices: - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print(f"Testing {test_type} on {InfiniDeviceNames[device]}") - print(f"{'='*60}") + print(f"{'=' * 60}") # Keep InfiniCore's runtime aligned with the selected test device. infinicore.set_device(infinicore.device(torch_device_map[device], 0)) @@ -96,7 +96,7 @@ def run_tests(self, devices, test_func, test_type="Test"): self.passed_tests.append( f"{test_case} - {InfiniDeviceNames[device]}" ) - print(f"\033[92m✓\033[0m Passed") + print("\033[92m✓\033[0m Passed") elif test_result.return_code == -1: # Test failed - use the actual error message from test_result fail_msg = f"{test_case} - {InfiniDeviceNames[device]} - {test_result.error_message}" @@ -153,11 +153,11 @@ def print_summary(self): """ total_tests = len(self.test_cases) passed_count = len(self.passed_tests) - skipped_count = len(self.skipped_tests) + # skipped_count = len(self.skipped_tests) partial_count = len(self.partial_tests) failed_count = len(self.failed_tests) - print(f"\n{'='*60}") + print(f"\n{'=' * 60}") print("TEST SUMMARY") print(f"Total tests: {total_tests}") print(f"\033[92mPassed: {passed_count}\033[0m") @@ -179,10 +179,10 @@ def print_summary(self): # If there are skipped or partial tests, show appropriate message if self.skipped_tests or self.partial_tests: print( - f"\n\033[93mTests completed with some implementations missing\033[0m" + "\n\033[93mTests completed with some implementations missing\033[0m" ) else: - print(f"\n\033[92mAll tests passed!\033[0m") + print("\n\033[92mAll tests passed!\033[0m") # Print benchmark summary if benchmarking was enabled if self.config.bench and ( @@ -193,12 +193,12 @@ def print_summary(self): ): self._print_benchmark_summary() - print(f"{'='*60}") + print(f"{'=' * 60}") return result def _print_benchmark_summary(self): """Print benchmark timing summary""" - print(f"{'-'*60}") + print(f"{'-' * 60}") print("BENCHMARK SUMMARY") torch_host_total = self.benchmark_times["torch_host_total"] @@ -437,6 +437,7 @@ def run_test(self, device, test_case, config): try: torch_result = self.torch_operator(*inputs, **kwargs) + synchronize_device(device_str) if torch_result is None: torch_implemented = False except NotImplementedError as e: @@ -448,6 +449,7 @@ def run_test(self, device, test_case, config): try: infini_result = self.infinicore_operator(*infini_inputs, **infini_kwargs) + infinicore.sync_stream() if infini_result is None: infini_implemented = False except NotImplementedError as e: @@ -621,7 +623,7 @@ def run_test(self, device, test_case, config): is_valid = compare_fn(infini_comparison, torch_comparison) if not is_valid: - raise AssertionError(f"Result comparison failed.") + raise AssertionError("Result comparison failed.") # ========================================================================== # UNIFIED BENCHMARKING LOGIC @@ -653,15 +655,15 @@ def run_test(self, device, test_case, config): if hasattr(config, "_test_runner") and config._test_runner: # Accumulate total times config._test_runner.benchmark_times["torch_host_total"] += torch_host - config._test_runner.benchmark_times[ - "torch_device_total" - ] += torch_device - config._test_runner.benchmark_times[ - "infinicore_host_total" - ] += infini_host - config._test_runner.benchmark_times[ - "infinicore_device_total" - ] += infini_device + config._test_runner.benchmark_times["torch_device_total"] += ( + torch_device + ) + config._test_runner.benchmark_times["infinicore_host_total"] += ( + infini_host + ) + config._test_runner.benchmark_times["infinicore_device_total"] += ( + infini_device + ) # Test passed successfully test_result.success = True diff --git a/test/infinicore/framework/utils/tensor_utils.py b/test/infinicore/framework/utils/tensor_utils.py index 19467d238..1c14c585a 100644 --- a/test/infinicore/framework/utils/tensor_utils.py +++ b/test/infinicore/framework/utils/tensor_utils.py @@ -1,5 +1,7 @@ import torch + import infinicore + from ..datatypes import to_infinicore_dtype, to_torch_dtype # ================================================================= @@ -71,6 +73,7 @@ def convert_infinicore_to_torch(infini_result): ) temp_tensor = infinicore_tensor_from_torch(torch_result_from_infini) temp_tensor.copy_(infini_result) + infinicore.sync_stream() return torch_result_from_infini diff --git a/test/infinicore/ops/random_sample.py b/test/infinicore/ops/random_sample.py index 7d1bb7c48..225710766 100644 --- a/test/infinicore/ops/random_sample.py +++ b/test/infinicore/ops/random_sample.py @@ -1,14 +1,15 @@ -import sys import os +import sys sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -import infinicore import infinicore.nn.functional as F import torch -from framework import BaseOperatorTest, TensorSpec, TestCase, GenericTestRunner +from framework import BaseOperatorTest, GenericTestRunner, TensorSpec, TestCase from framework.tensor import TensorInitializer +import infinicore + # ============================================================================== # Operator-specific configuration # ============================================================================== @@ -65,7 +66,7 @@ def parse_test_cases(): output_spec=None, comparison_target=None, tolerance=tolerance, - description=f"RandomSample - OUT_OF_PLACE", + description="RandomSample - OUT_OF_PLACE", ) ) @@ -79,7 +80,7 @@ def parse_test_cases(): ), comparison_target="out", tolerance=tolerance, - description=f"RandomSample - INPLACE(out)", + description="RandomSample - INPLACE(out)", ) ) @@ -226,6 +227,7 @@ def run_test(self, device, test_case, config): from framework.utils.tensor_utils import ( convert_infinicore_to_torch, infinicore_tensor_from_torch, + synchronize_device, ) inputs, kwargs = self.prepare_pytorch_inputs_and_kwargs(test_case, device) @@ -249,7 +251,9 @@ def run_test(self, device, test_case, config): # Run both operators torch_result = self.torch_operator(*inputs, **kwargs) + synchronize_device(device) infini_result = self.infinicore_operator(*infini_inputs, **infini_kwargs) + infinicore.sync_stream() # Extract indices from results comparison_target = test_case.comparison_target