diff --git a/infinimetrics/common/constants.py b/infinimetrics/common/constants.py index 29991b0f..d1045342 100644 --- a/infinimetrics/common/constants.py +++ b/infinimetrics/common/constants.py @@ -125,6 +125,8 @@ class OperatorConfig: OUTPUTS = "outputs" ATTRIBUTES = "attributes" TOLERANCE = "tolerance" + WARMUP_ITERATIONS = "warmup_iterations" + MEASURED_ITERATIONS = "measured_iterations" INFINICORE_OP = "infinicore_op" TORCH_OP = "torch_op" @@ -142,6 +144,46 @@ class TensorSpec: INIT_MODE = "init_mode" +class AttributeSpec: + """Operator attribute field names.""" + + NAME = "name" + VALUE = "value" + + +class MetricSpec: + """Metric result field names.""" + + NAME = "name" + VALUE = "value" + TYPE = "type" + RAW_DATA_URL = "raw_data_url" + UNIT = "unit" + + +class MetricType: + """Metric value representations.""" + + SCALAR = "scalar" + + +class OperatorMetric: + """Metric names emitted by operator adapters.""" + + LATENCY = "operator.latency" + ACCURACY = "operator.tensor_accuracy" + FLOPS = "operator.flops" + BANDWIDTH = "operator.bandwidth" + + +class BandwidthField: + """Memory bandwidth calculation field names.""" + + READ_BYTES = "read_bytes" + WRITE_BYTES = "write_bytes" + TOTAL_BYTES = "total_bytes" + + class InfiniCoreResult: """InfiniCore test result field names""" @@ -165,6 +207,33 @@ class InfiniCoreResult: DEFAULT_TOLERANCE = {"atol": 1e-3, "rtol": 1e-3} +# InfiniOps runtime mappings. PyTorch dtype objects intentionally remain in +# the adapter so importing common constants does not import PyTorch. +INFINIOPS_PLATFORM_TO_TORCH_DEVICE = { + "nvidia": "cuda", + "metax": "cuda", + "iluvatar": "cuda", + "hygon": "cuda", + "moore": "musa", + "cambricon": "mlu", + "ascend": "npu", + "cpu": "cpu", +} + +INFINIOPS_DEVICE_PLUGIN_MODULES = { + "mlu": "torch_mlu", + "npu": "torch_npu", + "musa": "torch_musa", +} + +INFINIOPS_STREAM_ACCESSORS = { + "npu": ("npu", "npu_stream"), + "cuda": ("cuda", "cuda_stream"), + "mlu": ("mlu", "mlu_stream"), + "musa": ("musa", "musa_stream"), +} + + # ============================================================ # Hardware Test Adapter Constants # ============================================================ diff --git a/infinimetrics/dispatcher.py b/infinimetrics/dispatcher.py index 75856cc1..5b8f1cf8 100644 --- a/infinimetrics/dispatcher.py +++ b/infinimetrics/dispatcher.py @@ -17,6 +17,7 @@ # test_type must use TestCategory enum (not string literals) _ADAPTER_REGISTRY = { (TestCategory.OPERATOR, "infinicore"): lambda: _create_infinicore_adapter(), + (TestCategory.OPERATOR, "infiniops"): lambda: _create_infiniops_adapter(), (TestCategory.HARDWARE, "cudaunified"): lambda: _create_hardware_adapter(), (TestCategory.COMM, "nccltest"): lambda: _create_nccltests_adapter(), (TestCategory.INFER, "infinilm"): lambda: _create_inference_adapter(), @@ -40,6 +41,13 @@ def _create_infinicore_adapter(): return InfiniCoreAdapter() +def _create_infiniops_adapter(): + """Create InfiniOps adapter (lazy import).""" + from infinimetrics.operators.infiniops_adapter import InfiniOpsAdapter + + return InfiniOpsAdapter() + + def _create_nccltests_adapter(): """Create NCCL communication adapter.""" from infinimetrics.communication.nccl_adapter import NcclTestsAdapter diff --git a/infinimetrics/executor.py b/infinimetrics/executor.py index fe95b126..2537cbca 100644 --- a/infinimetrics/executor.py +++ b/infinimetrics/executor.py @@ -175,6 +175,13 @@ def execute(self) -> TestResult: test_result.result_file = result_file test_result.duration = time.time() - start_time + # Extract result_code and error_msg from adapter response + if isinstance(response, dict): + if "result_code" in response: + test_result.result_code = response["result_code"] + if "error_msg" in response: + test_result.error_msg = response["error_msg"] + logger.info( f"Executor: {self.testcase} completed in {test_result.duration:.2f}s" ) diff --git a/infinimetrics/operators/__init__.py b/infinimetrics/operators/__init__.py index 70e442c4..b1de08e7 100644 --- a/infinimetrics/operators/__init__.py +++ b/infinimetrics/operators/__init__.py @@ -5,10 +5,23 @@ FLOPSCalculator, calculate_bandwidth, ) -from infinimetrics.operators.infinicore_adapter import InfiniCoreAdapter __all__ = [ "FLOPSCalculator", "calculate_bandwidth", "InfiniCoreAdapter", + "InfiniOpsAdapter", ] + + +def __getattr__(name): + """Load adapters only when callers explicitly request them.""" + if name == "InfiniCoreAdapter": + from infinimetrics.operators.infinicore_adapter import InfiniCoreAdapter + + return InfiniCoreAdapter + if name == "InfiniOpsAdapter": + from infinimetrics.operators.infiniops_adapter import InfiniOpsAdapter + + return InfiniOpsAdapter + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/infinimetrics/operators/flops_calculator.py b/infinimetrics/operators/flops_calculator.py index 3ddb5097..f06f320b 100644 --- a/infinimetrics/operators/flops_calculator.py +++ b/infinimetrics/operators/flops_calculator.py @@ -8,7 +8,11 @@ from typing import Dict, List, Optional, Callable -from infinimetrics.common.constants import DTYPE_BYTES_MAP +from infinimetrics.common.constants import ( + BandwidthField, + DTYPE_BYTES_MAP, + TensorSpec, +) class FLOPSCalculator: @@ -103,7 +107,7 @@ def get_flops( @staticmethod def _get_tensor_size(tensor: Dict) -> int: """Get total number of elements in tensor""" - shape = tensor.get("shape", []) + shape = tensor.get(TensorSpec.SHAPE, []) size = 1 for dim in shape: size *= dim @@ -111,14 +115,14 @@ def _get_tensor_size(tensor: Dict) -> int: # Register matrix operations -@FLOPSCalculator.register(["matmul", "bmm", "batchmm"]) +@FLOPSCalculator.register(["matmul", "mm", "bmm", "batchmm"]) def _matmul_flops(inputs: List[Dict], outputs: List[Dict]) -> float: """Matrix Multiplication: C = A @ B (FLOPS = 2 * M * N * K)""" if len(inputs) < 2: return 0.0 - a_shape = inputs[0].get("shape", []) - b_shape = inputs[1].get("shape", []) + a_shape = inputs[0].get(TensorSpec.SHAPE, []) + b_shape = inputs[1].get(TensorSpec.SHAPE, []) if len(a_shape) == 2 and len(b_shape) == 2: m, k = a_shape @@ -135,14 +139,14 @@ def _matmul_flops(inputs: List[Dict], outputs: List[Dict]) -> float: return 0.0 -@FLOPSCalculator.register(["addmm", "linear"]) +@FLOPSCalculator.register(["addmm"]) def _addmm_flops(inputs: List[Dict], outputs: List[Dict]) -> float: """AddMM: C = beta * bias + alpha * (input @ weight)""" if len(inputs) < 3: return 0.0 - input_shape = inputs[1].get("shape", []) - weight_shape = inputs[2].get("shape", []) + input_shape = inputs[1].get(TensorSpec.SHAPE, []) + weight_shape = inputs[2].get(TensorSpec.SHAPE, []) if len(input_shape) >= 2 and len(weight_shape) >= 2: m, k = input_shape[-2], input_shape[-1] @@ -158,6 +162,28 @@ def _addmm_flops(inputs: List[Dict], outputs: List[Dict]) -> float: return 0.0 +@FLOPSCalculator.register(["linear"]) +def _linear_flops(inputs: List[Dict], outputs: List[Dict]) -> float: + """Linear: output = input @ weight + bias.""" + if len(inputs) < 2: + return 0.0 + + input_shape = inputs[0].get(TensorSpec.SHAPE, []) + weight_shape = inputs[1].get(TensorSpec.SHAPE, []) + output_shape = outputs[0].get(TensorSpec.SHAPE, []) if outputs else [] + if len(input_shape) < 2 or len(weight_shape) < 2 or not output_shape: + return 0.0 + + k = input_shape[-1] + n = output_shape[-1] + batch = 1 + for dim in input_shape[:-1]: + batch *= dim + + bias_flops = batch * n if len(inputs) >= 3 else 0 + return 2.0 * batch * n * k + bias_flops + + @FLOPSCalculator.register(["conv2d", "conv2d_backward"]) def _conv2d_flops(inputs: List[Dict], outputs: List[Dict]) -> float: """ @@ -170,11 +196,11 @@ def _conv2d_flops(inputs: List[Dict], outputs: List[Dict]) -> float: return 0.0 # Input: [N, C_in, H_in, W_in] - input_shape = inputs[0].get("shape", []) + input_shape = inputs[0].get(TensorSpec.SHAPE, []) # Weight: [C_out, C_in, K_h, K_w] - weight_shape = inputs[1].get("shape", []) + weight_shape = inputs[1].get(TensorSpec.SHAPE, []) # Output: [N, C_out, H_out, W_out] - output_shape = outputs[0].get("shape", []) + output_shape = outputs[0].get(TensorSpec.SHAPE, []) if len(input_shape) != 4 or len(weight_shape) != 4 or len(output_shape) != 4: return 0.0 @@ -214,7 +240,7 @@ def calculate_bandwidth( """ def get_tensor_bytes(tensor: Dict) -> int: - dtype = tensor.get("dtype", "float32").lower() + dtype = tensor.get(TensorSpec.DTYPE, "float32").lower() bytes_per_element = DTYPE_BYTES_MAP.get(dtype, 4) size = FLOPSCalculator._get_tensor_size(tensor) return size * bytes_per_element @@ -223,7 +249,7 @@ def get_tensor_bytes(tensor: Dict) -> int: write_bytes = sum(get_tensor_bytes(out) for out in outputs) return { - "read_bytes": read_bytes, - "write_bytes": write_bytes, - "total_bytes": read_bytes + write_bytes, + BandwidthField.READ_BYTES: read_bytes, + BandwidthField.WRITE_BYTES: write_bytes, + BandwidthField.TOTAL_BYTES: read_bytes + write_bytes, } diff --git a/infinimetrics/operators/infiniops_adapter.py b/infinimetrics/operators/infiniops_adapter.py new file mode 100644 index 00000000..8a1cb79e --- /dev/null +++ b/infinimetrics/operators/infiniops_adapter.py @@ -0,0 +1,693 @@ +#!/usr/bin/env python3 +"""InfiniOps operator performance adapter.""" + +import copy +import importlib +import logging +import time +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import torch + +import infini.ops + +from infinimetrics.adapter import BaseAdapter +from infinimetrics.common.constants import ( + AttributeSpec, + BandwidthField, + DEFAULT_MEASURED_ITERATIONS, + DEFAULT_TOLERANCE, + DEFAULT_WARMUP_ITERATIONS, + ErrorCode, + INFINIOPS_DEVICE_PLUGIN_MODULES, + INFINIOPS_PLATFORM_TO_TORCH_DEVICE, + INFINIOPS_STREAM_ACCESSORS, + InfiniMetricsJson, + MetricSpec, + MetricType, + OperatorConfig, + OperatorMetric, + TensorSpec, +) +from infinimetrics.operators.flops_calculator import ( + FLOPSCalculator, + calculate_bandwidth, +) + +logger = logging.getLogger(__name__) + + +_DTYPE_MAP = { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "int64": torch.int64, + "int32": torch.int32, + "int16": torch.int16, + "int8": torch.int8, +} + + +class OperatorFamily: + """Builder families used by the declarative operator registry.""" + + BINARY = "binary" + CAST = "cast" + CONCAT = "concat" + MATRIX = "matrix" + + +@dataclass(frozen=True) +class OperatorSpec: + """Description of a common operator and its builder family.""" + + name: str + family: str + torch_name: Optional[str] = None + scalar_args: Tuple[Any, ...] = () + + +@dataclass(frozen=True) +class BenchmarkCase: + """Prepared InfiniOps and reference calls for one benchmark.""" + + operation: Callable + reference: Callable + args: tuple + kwargs: dict = field(default_factory=dict) + + +OPERATOR_SPECS: Dict[str, OperatorSpec] = { + "add": OperatorSpec("add", OperatorFamily.BINARY, torch_name="add"), + "sub": OperatorSpec( + "sub", OperatorFamily.BINARY, torch_name="sub", scalar_args=(1.0,) + ), + "mul": OperatorSpec("mul", OperatorFamily.BINARY, torch_name="mul"), + "div": OperatorSpec("div", OperatorFamily.BINARY, torch_name="div"), + "cast": OperatorSpec("cast", OperatorFamily.CAST), + "cat": OperatorSpec("cat", OperatorFamily.CONCAT), + "gemm": OperatorSpec("gemm", OperatorFamily.MATRIX), + "matmul": OperatorSpec("matmul", OperatorFamily.MATRIX), + "mm": OperatorSpec("mm", OperatorFamily.MATRIX), + "linear": OperatorSpec("linear", OperatorFamily.MATRIX), +} + + +def _get_stream(device): + if isinstance(device, torch.device): + device = device.type + if isinstance(device, str) and ":" in device: + device = device.split(":")[0] + if device == "cpu": + return 0 + + mod_name, attr = INFINIOPS_STREAM_ACCESSORS.get(device, (None, None)) + if mod_name is None: + return 0 + mod = getattr(torch, mod_name, None) + if mod is None: + return 0 + stream = mod.current_stream() + return getattr(stream, attr, 0) + + +def _empty_strided(shape, strides, *, dtype=None, device=None): + if strides is None: + return torch.empty(shape, dtype=dtype, device=device) + return torch.empty_strided(shape, strides, dtype=dtype, device=device) + + +def _randn_strided(shape, strides, *, dtype=None, device=None): + out = _empty_strided(shape, strides, dtype=dtype, device=device) + out.as_strided( + (out.untyped_storage().size() // out.element_size(),), (1,) + ).normal_() + return out + + +def _clone_strided(inp): + out = _empty_strided(inp.size(), inp.stride(), dtype=inp.dtype, device=inp.device) + flat_args = (out.untyped_storage().size() // out.element_size(),), (1,) + out.as_strided(*flat_args).copy_(inp.as_strided(*flat_args)) + return out + + +def _clone(obj): + if isinstance(obj, torch.Tensor): + return _clone_strided(obj) + if isinstance(obj, tuple): + return tuple(_clone(arg) for arg in obj) + if isinstance(obj, list): + return [_clone(arg) for arg in obj] + if isinstance(obj, dict): + return {key: _clone(value) for key, value in obj.items()} + return obj + + +def _synchronize(device): + if device == "cpu": + return + mod = getattr(torch, device, None) + if mod is not None and hasattr(mod, "synchronize"): + mod.synchronize() + + +def _get_attributes(config: dict) -> dict: + """Convert the attribute list into a name-to-value mapping.""" + return { + attr[AttributeSpec.NAME]: attr[AttributeSpec.VALUE] + for attr in config.get(OperatorConfig.ATTRIBUTES, []) + } + + +def _load_device_plugin(device: str) -> None: + """Load the PyTorch extension that registers a vendor device.""" + module_name = INFINIOPS_DEVICE_PLUGIN_MODULES.get(device) + if module_name is None: + return + + try: + importlib.import_module(module_name) + except ImportError as exc: + raise RuntimeError( + f"PyTorch device plugin '{module_name}' is required for device " + f"'{device}'" + ) from exc + + +def _pick_slot(op_name: str, device: str) -> int: + """Return the first implementation slot registered by InfiniOps.""" + op_pascal = "".join(part.capitalize() for part in op_name.split("_")) + op_cls = getattr(infini.ops, op_pascal, None) + if op_cls is None: + raise ValueError(f"InfiniOps operator class not found: {op_pascal}") + if not hasattr(op_cls, "active_implementation_indices"): + raise RuntimeError( + f"InfiniOps operator '{op_name}' cannot report implementations" + ) + + indices = op_cls.active_implementation_indices(device) + if not indices: + raise RuntimeError( + f"InfiniOps operator '{op_name}' has no active implementation " + f"for device '{device}'" + ) + return indices[0] + + +def _tensor_specs(config: dict, field_name: str) -> List[Dict[str, Any]]: + specs = config.get(field_name, []) + if not isinstance(specs, list): + raise ValueError(f"{field_name} must be a list") + return specs + + +def _require_inputs(config: dict, count: int) -> List[Dict[str, Any]]: + inputs = _tensor_specs(config, OperatorConfig.INPUTS) + if len(inputs) < count: + raise ValueError(f"operator requires at least {count} input tensor(s)") + return inputs + + +def _output_shape(config: dict, fallback) -> Any: + outputs = _tensor_specs(config, OperatorConfig.OUTPUTS) + if outputs: + return outputs[0][TensorSpec.SHAPE] + return fallback + + +def _build_binary_case( + spec: OperatorSpec, torch_device: str, torch_dtype, config: dict +) -> BenchmarkCase: + inputs = _require_inputs(config, 2) + shape = inputs[0][TensorSpec.SHAPE] + out_shape = _output_shape(config, shape) + a = _randn_strided(shape, None, dtype=torch_dtype, device=torch_device) + b = _randn_strided( + inputs[1][TensorSpec.SHAPE], + None, + dtype=torch_dtype, + device=torch_device, + ) + out = _empty_strided(out_shape, None, dtype=torch_dtype, device=torch_device) + stream = _get_stream(torch_device) + slot = _pick_slot(spec.name, torch_device) + infini_op = getattr(infini.ops, spec.name) + torch_op = getattr(torch, spec.torch_name) + + def operation(a, b, out): + infini_op( + a, + b, + *spec.scalar_args, + out, + stream=stream, + implementation_index=slot, + ) + return out + + def reference(a, b, out): + torch_op(a, b, out=out) + return out + + return BenchmarkCase(operation, reference, (a, b, out)) + + +def _build_cast_case( + spec: OperatorSpec, torch_device: str, torch_dtype, config: dict +) -> BenchmarkCase: + inputs = _require_inputs(config, 1) + outputs = _tensor_specs(config, OperatorConfig.OUTPUTS) + shape = inputs[0][TensorSpec.SHAPE] + out_dtype_name = ( + outputs[0].get(TensorSpec.DTYPE, "float32") if outputs else "float32" + ) + if out_dtype_name not in _DTYPE_MAP: + raise ValueError(f"Unsupported output dtype: {out_dtype_name}") + + inp = _randn_strided(shape, None, dtype=torch_dtype, device=torch_device) + out = _empty_strided( + shape, + None, + dtype=_DTYPE_MAP[out_dtype_name], + device=torch_device, + ) + stream = _get_stream(torch_device) + slot = _pick_slot(spec.name, torch_device) + + def operation(inp, out): + infini.ops.cast(inp, out, stream=stream, implementation_index=slot) + return out + + def reference(inp, out): + out.copy_(inp.to(out.dtype)) + return out + + return BenchmarkCase(operation, reference, (inp, out)) + + +def _build_concat_case( + spec: OperatorSpec, torch_device: str, torch_dtype, config: dict +) -> BenchmarkCase: + inputs = _require_inputs(config, 1) + dim = _get_attributes(config).get("dim", 0) + tensors = tuple( + _randn_strided( + input_spec[TensorSpec.SHAPE], + None, + dtype=torch_dtype, + device=torch_device, + ) + for input_spec in inputs + ) + fallback_shape = list(tensors[0].shape) + fallback_shape[dim] = sum(tensor.shape[dim] for tensor in tensors) + out = _empty_strided( + _output_shape(config, fallback_shape), + None, + dtype=torch_dtype, + device=torch_device, + ) + stream = _get_stream(torch_device) + + def operation(*args): + inps = list(args[:-1]) + output = args[-1] + infini.ops.cat(inps[0], inps[1:], dim, output, stream=stream) + return output + + def reference(*args): + inps = list(args[:-1]) + output = args[-1] + output.copy_(torch.cat(inps, dim=dim)) + return output + + return BenchmarkCase(operation, reference, (*tensors, out)) + + +def _logical_matrix(tensor, transpose: bool): + return tensor.transpose(-2, -1) if transpose else tensor + + +def _build_matrix_case( + spec: OperatorSpec, torch_device: str, torch_dtype, config: dict +) -> BenchmarkCase: + inputs = _require_inputs(config, 2) + attrs = _get_attributes(config) + trans_a = bool(attrs.get("trans_a", False)) + trans_b = bool(attrs.get("trans_b", False)) + a_shape = inputs[0][TensorSpec.SHAPE] + b_shape = inputs[1][TensorSpec.SHAPE] + fallback_shape = [*a_shape[:-1], b_shape[-1]] + out_shape = _output_shape(config, fallback_shape) + + a = _randn_strided(a_shape, None, dtype=torch_dtype, device=torch_device) + b = _randn_strided(b_shape, None, dtype=torch_dtype, device=torch_device) + stream = _get_stream(torch_device) + slot = _pick_slot(spec.name, torch_device) + + if spec.name == "gemm": + alpha = attrs.get("alpha", 1.0) + beta = attrs.get("beta", 0.0) + out = _randn_strided(out_shape, None, dtype=torch_dtype, device=torch_device) + + def operation(a, b, alpha, beta, trans_a, trans_b, out): + infini.ops.gemm( + a, + b, + alpha, + beta, + trans_a, + trans_b, + out, + stream=stream, + implementation_index=slot, + ) + return out + + def reference(a, b, alpha, beta, trans_a, trans_b, out): + if alpha == 0: + out.mul_(beta) + return out + product = torch.matmul( + _logical_matrix(a, trans_a).float(), + _logical_matrix(b, trans_b).float(), + ) + out.copy_((alpha * product + beta * out.float()).to(out.dtype)) + return out + + args = (a, b, alpha, beta, trans_a, trans_b, out) + return BenchmarkCase(operation, reference, args) + + out = _empty_strided(out_shape, None, dtype=torch_dtype, device=torch_device) + + if spec.name == "linear": + has_bias = bool(attrs.get("has_bias", len(inputs) >= 3)) + bias = None + if has_bias: + bias_shape = ( + inputs[2][TensorSpec.SHAPE] if len(inputs) >= 3 else (out_shape[-1],) + ) + bias = _randn_strided( + bias_shape, None, dtype=torch_dtype, device=torch_device + ) + + def operation(a, b, bias, out): + infini.ops.linear( + a, + b, + bias, + trans_a, + trans_b, + out, + stream=stream, + implementation_index=slot, + ) + return out + + def reference(a, b, bias, out): + result = torch.matmul( + _logical_matrix(a, trans_a).float(), + _logical_matrix(b, trans_b).float(), + ) + if bias is not None: + result = result + bias.float() + out.copy_(result.to(out.dtype)) + return out + + return BenchmarkCase(operation, reference, (a, b, bias, out)) + + if spec.name == "matmul": + + def operation(a, b, out): + infini.ops.matmul( + a, + b, + out, + trans_a, + trans_b, + stream=stream, + implementation_index=slot, + ) + return out + + def reference(a, b, out): + result = torch.matmul( + _logical_matrix(a, trans_a).float(), + _logical_matrix(b, trans_b).float(), + ) + out.copy_(result.to(out.dtype)) + return out + + return BenchmarkCase(operation, reference, (a, b, out)) + + def operation(a, b, out): + infini.ops.mm(a, b, out, stream=stream, implementation_index=slot) + return out + + def reference(a, b, out): + out.copy_(torch.mm(a.float(), b.float()).to(out.dtype)) + return out + + return BenchmarkCase(operation, reference, (a, b, out)) + + +CASE_BUILDERS: Dict[str, Callable[..., BenchmarkCase]] = { + OperatorFamily.BINARY: _build_binary_case, + OperatorFamily.CAST: _build_cast_case, + OperatorFamily.CONCAT: _build_concat_case, + OperatorFamily.MATRIX: _build_matrix_case, +} + + +class InfiniOpsAdapter(BaseAdapter): + """Adapter for common InfiniOps operator performance tests.""" + + def __init__(self): + self._req_metrics_template = [] + + def process(self, test_input: Union[Dict[str, Any], Any]) -> Dict[str, Any]: + test_input = self._normalize_test_input(test_input) + if not test_input: + raise ValueError(f"Invalid test_input type: {type(test_input)}") + + testcase = test_input.get(InfiniMetricsJson.TESTCASE, "unknown") + logger.info(f"InfiniOpsAdapter: Processing {testcase}") + config = test_input.get(InfiniMetricsJson.CONFIG, {}) + self._req_metrics_template = test_input.get(InfiniMetricsJson.METRICS, []) + + try: + operator_name = str(config.get(OperatorConfig.OPERATOR, "")).lower() + operator_spec = OPERATOR_SPECS.get(operator_name) + if operator_spec is None: + return self._create_error_response( + f"Unsupported operator: {operator_name or ''}", + test_input, + ) + + platform = str(config.get(OperatorConfig.DEVICE, "cpu")).lower() + torch_device = INFINIOPS_PLATFORM_TO_TORCH_DEVICE.get(platform, platform) + _load_device_plugin(torch_device) + + inputs = _tensor_specs(config, OperatorConfig.INPUTS) + dtype_name = ( + str(inputs[0].get(TensorSpec.DTYPE, "float16")).lower() + if inputs + else "float16" + ) + if dtype_name not in _DTYPE_MAP: + raise ValueError(f"Unsupported dtype: {dtype_name}") + + warmup = config.get( + OperatorConfig.WARMUP_ITERATIONS, + DEFAULT_WARMUP_ITERATIONS, + ) + measured = config.get( + OperatorConfig.MEASURED_ITERATIONS, + DEFAULT_MEASURED_ITERATIONS, + ) + self._validate_iterations(warmup, measured) + self._check_op_available(operator_name) + + benchmark_case = self._build_case( + operator_spec, + torch_device, + _DTYPE_MAP[dtype_name], + config, + ) + tolerance = config.get(OperatorConfig.TOLERANCE, DEFAULT_TOLERANCE) + avg_latency_s, accuracy_pass = self._run_benchmark( + benchmark_case, + torch_device, + warmup, + measured, + tolerance, + ) + return self._create_response( + test_input, + config, + operator_name, + avg_latency_s, + accuracy_pass, + ) + except Exception as exc: + operator = config.get(OperatorConfig.OPERATOR, "unknown") + device = config.get(OperatorConfig.DEVICE, "unknown") + logger.error( + f"InfiniOpsAdapter: Operator test failed for {testcase}\n" + f" Operator: {operator}\n Device: {device}\n Error: {exc}", + exc_info=True, + ) + raise + + @staticmethod + def _validate_iterations(warmup: Any, measured: Any) -> None: + if not isinstance(warmup, int) or warmup < 0: + raise ValueError("warmup_iterations must be a non-negative integer") + if not isinstance(measured, int) or measured <= 0: + raise ValueError("measured_iterations must be a positive integer") + + @staticmethod + def _check_op_available(operator_name: str) -> None: + op_pascal = "".join(part.capitalize() for part in operator_name.split("_")) + if getattr(infini.ops, op_pascal, None) is None: + raise ValueError(f"InfiniOps operator class not found: {op_pascal}") + + @staticmethod + def _build_case( + spec: OperatorSpec, torch_device: str, torch_dtype, config: dict + ) -> BenchmarkCase: + builder = CASE_BUILDERS[spec.family] + return builder(spec, torch_device, torch_dtype, config) + + def _run_benchmark( + self, + benchmark_case: BenchmarkCase, + torch_device: str, + warmup: int, + measured: int, + tolerance: dict, + ) -> Tuple[float, bool]: + cloned_args = _clone(benchmark_case.args) + cloned_kwargs = _clone(benchmark_case.kwargs) + output = benchmark_case.operation(*benchmark_case.args, **benchmark_case.kwargs) + expected = benchmark_case.reference(*cloned_args, **cloned_kwargs) + accuracy_pass = self._outputs_match(output, expected, tolerance) + + for _ in range(warmup): + benchmark_case.operation(*benchmark_case.args, **benchmark_case.kwargs) + _synchronize(torch_device) + + total_time = 0.0 + for _ in range(measured): + _synchronize(torch_device) + start = time.perf_counter() + benchmark_case.operation(*benchmark_case.args, **benchmark_case.kwargs) + _synchronize(torch_device) + total_time += time.perf_counter() - start + + return total_time / measured, accuracy_pass + + @staticmethod + def _outputs_match(output: Any, expected: Any, tolerance: dict) -> bool: + atol = tolerance.get("atol", DEFAULT_TOLERANCE["atol"]) + rtol = tolerance.get("rtol", DEFAULT_TOLERANCE["rtol"]) + + if isinstance(output, tuple): + if not isinstance(expected, tuple) or len(output) != len(expected): + return False + return all( + InfiniOpsAdapter._tensor_matches(actual, reference, rtol, atol) + for actual, reference in zip(output, expected) + ) + if isinstance(output, torch.Tensor): + return InfiniOpsAdapter._tensor_matches(output, expected, rtol, atol) + return True + + @staticmethod + def _tensor_matches(actual, expected, rtol: float, atol: float) -> bool: + if actual.dtype.is_floating_point: + return torch.allclose( + actual, + expected, + rtol=rtol, + atol=atol, + equal_nan=True, + ) + return torch.equal(actual, expected) + + def _create_response( + self, + test_input: dict, + config: dict, + operator_name: str, + avg_latency_s: float, + accuracy_pass: bool, + ) -> Dict[str, Any]: + response = copy.deepcopy(test_input) + response[InfiniMetricsJson.TIME] = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + response[InfiniMetricsJson.RESULT_CODE] = ( + ErrorCode.SUCCESS if accuracy_pass else ErrorCode.INTERNAL + ) + if not accuracy_pass: + response[ + InfiniMetricsJson.ERROR_MSG + ] = f"Accuracy check failed for operator '{operator_name}'" + response[InfiniMetricsJson.METRICS] = self._compute_metrics( + config, avg_latency_s, accuracy_pass + ) + return response + + def _compute_metrics( + self, config: dict, avg_latency_s: float, accuracy_pass: bool + ) -> List[Dict]: + metrics = copy.deepcopy(self._req_metrics_template) + inputs = config.get(OperatorConfig.INPUTS, []) + outputs = config.get(OperatorConfig.OUTPUTS, []) + operator = config.get(OperatorConfig.OPERATOR, "").lower() + + for metric in metrics: + name = metric.get(MetricSpec.NAME, "") + if name == OperatorMetric.LATENCY: + self._set_scalar_metric(metric, round(avg_latency_s * 1000, 6), "ms") + elif name == OperatorMetric.ACCURACY: + metric.update( + { + MetricSpec.VALUE: "PASS" if accuracy_pass else "FAIL", + MetricSpec.UNIT: "", + } + ) + elif name == OperatorMetric.FLOPS: + value = 0.0 + if avg_latency_s > 0: + flops = FLOPSCalculator.get_flops(operator, inputs, outputs) + if flops > 0: + value = (flops / avg_latency_s) / 1e12 + if value >= 0.0001: + value = round(value, 4) + self._set_scalar_metric(metric, value, "TFLOPS") + elif name == OperatorMetric.BANDWIDTH: + value = 0.0 + if avg_latency_s > 0: + bandwidth = calculate_bandwidth(inputs, outputs) + total_bytes = bandwidth[BandwidthField.TOTAL_BYTES] + if total_bytes > 0: + value = (total_bytes / avg_latency_s) / 1e9 + if value >= 0.0001: + value = round(value, 4) + self._set_scalar_metric(metric, value, "GB/s") + + return metrics + + @staticmethod + def _set_scalar_metric(metric: dict, value: Any, unit: str) -> None: + metric.update( + { + MetricSpec.VALUE: value, + MetricSpec.TYPE: MetricType.SCALAR, + MetricSpec.RAW_DATA_URL: "", + MetricSpec.UNIT: unit, + } + ) diff --git a/tests/test_executor_adapter_response.py b/tests/test_executor_adapter_response.py new file mode 100644 index 00000000..4729a5c2 --- /dev/null +++ b/tests/test_executor_adapter_response.py @@ -0,0 +1,45 @@ +import json +from pathlib import Path + +from infinimetrics.adapter import BaseAdapter +from infinimetrics.common.constants import TestCategory as Category +from infinimetrics.dispatcher import Dispatcher, _ADAPTER_REGISTRY +from infinimetrics.executor import Executor + + +class ErrorResponseAdapter(BaseAdapter): + def process(self, test_input): + return { + "run_id": test_input["run_id"], + "testcase": test_input["testcase"], + "result_code": 2, + "error_msg": "accuracy failed", + "config": test_input["config"], + "metrics": [], + } + + +def test_dispatcher_registers_infiniops_framework(): + assert Dispatcher()._parse_testcase("operator.InfiniOps.Add") == ( + "operator", + "infiniops", + ) + assert (Category.OPERATOR, "infiniops") in _ADAPTER_REGISTRY + + +def test_executor_propagates_adapter_error_response(tmp_path, monkeypatch): + payload = { + "run_id": "adapter-error", + "testcase": "operator.InfiniOps.Add", + "config": {"output_dir": str(tmp_path)}, + "metrics": [], + } + executor = Executor(payload, ErrorResponseAdapter()) + monkeypatch.setattr(executor, "_enrich_environment", lambda response: response) + + result = executor.execute() + + assert result.result_code == 2 + assert result.error_msg == "accuracy failed" + saved = json.loads(Path(result.result_file).read_text(encoding="utf-8")) + assert saved["result_code"] == 2 diff --git a/tests/test_infiniops_adapter.py b/tests/test_infiniops_adapter.py new file mode 100644 index 00000000..67bed80c --- /dev/null +++ b/tests/test_infiniops_adapter.py @@ -0,0 +1,220 @@ +import importlib +import sys +import types + +import pytest + + +@pytest.fixture +def adapter_module(monkeypatch): + fake_torch = types.ModuleType("torch") + + class FakeDevice: + def __init__(self, device_type="cpu"): + self.type = device_type + + class FakeTensor: + pass + + fake_torch.device = FakeDevice + fake_torch.Tensor = FakeTensor + fake_torch.float32 = object() + fake_torch.float16 = object() + fake_torch.bfloat16 = object() + fake_torch.int64 = object() + fake_torch.int32 = object() + fake_torch.int16 = object() + fake_torch.int8 = object() + + fake_infini = types.ModuleType("infini") + fake_infini.__path__ = [] + fake_ops = types.ModuleType("infini.ops") + fake_infini.ops = fake_ops + + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "infini", fake_infini) + monkeypatch.setitem(sys.modules, "infini.ops", fake_ops) + monkeypatch.delitem( + sys.modules, "infinimetrics.operators.infiniops_adapter", raising=False + ) + + return importlib.import_module("infinimetrics.operators.infiniops_adapter") + + +def test_pick_slot_preserves_native_slot_zero(adapter_module): + class Add: + @staticmethod + def active_implementation_indices(device): + return [0, 8] + + adapter_module.infini.ops.Add = Add + + assert adapter_module._pick_slot("add", "mlu") == 0 + + +def test_pick_slot_preserves_registered_aten_fallback(adapter_module): + class Add: + @staticmethod + def active_implementation_indices(device): + return [8] + + adapter_module.infini.ops.Add = Add + + assert adapter_module._pick_slot("add", "mlu") == 8 + + +def test_pick_slot_rejects_unregistered_fallback(adapter_module): + class Add: + @staticmethod + def active_implementation_indices(device): + return [] + + adapter_module.infini.ops.Add = Add + + with pytest.raises(RuntimeError, match="no active implementation.*mlu"): + adapter_module._pick_slot("add", "mlu") + + +def test_only_common_operators_are_registered(adapter_module): + assert set(adapter_module.OPERATOR_SPECS) == { + "add", + "sub", + "mul", + "div", + "cast", + "cat", + "gemm", + "matmul", + "mm", + "linear", + } + + +def test_operator_specs_use_four_builder_families(adapter_module): + assert {spec.family for spec in adapter_module.OPERATOR_SPECS.values()} == { + adapter_module.OperatorFamily.BINARY, + adapter_module.OperatorFamily.CAST, + adapter_module.OperatorFamily.CONCAT, + adapter_module.OperatorFamily.MATRIX, + } + assert set(adapter_module.CASE_BUILDERS) == { + adapter_module.OperatorFamily.BINARY, + adapter_module.OperatorFamily.CAST, + adapter_module.OperatorFamily.CONCAT, + adapter_module.OperatorFamily.MATRIX, + } + + +def test_deferred_model_operators_are_not_registered(adapter_module): + deferred = { + "rms_norm", + "causal_softmax", + "swiglu", + "flash_attention", + "rotary_embedding", + "add_rms_norm", + "reshape_and_cache", + } + + assert deferred.isdisjoint(adapter_module.OPERATOR_SPECS) + + +def test_deferred_operator_returns_structured_error(adapter_module): + result = adapter_module.InfiniOpsAdapter().process( + { + "run_id": "deferred-op", + "testcase": "operator.InfiniOps.FlashAttention", + "config": {"operator": "flash_attention", "device": "cambricon"}, + "metrics": [], + } + ) + + assert result["result_code"] != 0 + assert result["error_msg"] == "Unsupported operator: flash_attention" + + +def test_runtime_mappings_come_from_common_constants(adapter_module): + assert adapter_module.INFINIOPS_PLATFORM_TO_TORCH_DEVICE["cambricon"] == "mlu" + assert adapter_module.INFINIOPS_DEVICE_PLUGIN_MODULES["mlu"] == "torch_mlu" + assert adapter_module.INFINIOPS_STREAM_ACCESSORS["mlu"] == ( + "mlu", + "mlu_stream", + ) + + +def test_missing_device_plugin_has_actionable_error(adapter_module, monkeypatch): + def missing_module(name): + raise ImportError(name) + + monkeypatch.setattr(adapter_module.importlib, "import_module", missing_module) + + with pytest.raises(RuntimeError, match="torch_mlu.*mlu"): + adapter_module._load_device_plugin("mlu") + + +def test_accuracy_failure_sets_nonzero_result_code(adapter_module, monkeypatch): + adapter = adapter_module.InfiniOpsAdapter() + monkeypatch.setitem( + adapter_module.OPERATOR_SPECS, + "test_op", + adapter_module.OperatorSpec("test_op", adapter_module.OperatorFamily.BINARY), + ) + monkeypatch.setattr(adapter, "_check_op_available", lambda *args: None) + monkeypatch.setattr( + adapter, + "_build_case", + lambda *args: adapter_module.BenchmarkCase(lambda: None, lambda: None, ()), + ) + monkeypatch.setattr(adapter, "_run_benchmark", lambda *args: (0.001, False)) + + result = adapter.process( + { + "run_id": "accuracy-failure", + "testcase": "operator.infiniops.TestOp", + "config": { + "operator": "test_op", + "device": "cpu", + "inputs": [{"shape": [2, 2], "dtype": "float16"}], + "outputs": [{"shape": [2, 2], "dtype": "float16"}], + "warmup_iterations": 0, + "measured_iterations": 1, + }, + "metrics": [{"name": "operator.tensor_accuracy"}], + } + ) + + assert result["result_code"] != 0 + assert result["error_msg"] == "Accuracy check failed for operator 'test_op'" + assert result["metrics"][0]["value"] == "FAIL" + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("warmup_iterations", -1, "non-negative integer"), + ("measured_iterations", 0, "positive integer"), + ], +) +def test_iteration_counts_are_validated( + adapter_module, monkeypatch, field, value, message +): + adapter = adapter_module.InfiniOpsAdapter() + monkeypatch.setattr(adapter_module, "_load_device_plugin", lambda device: None) + config = { + "operator": "add", + "device": "cpu", + "inputs": [{"shape": [1], "dtype": "float16"}], + "outputs": [{"shape": [1], "dtype": "float16"}], + "warmup_iterations": 0, + "measured_iterations": 1, + } + config[field] = value + + with pytest.raises(ValueError, match=message): + adapter.process( + { + "testcase": "operator.infiniops.Add", + "config": config, + "metrics": [], + } + ) diff --git a/tests/test_operator_flops.py b/tests/test_operator_flops.py new file mode 100644 index 00000000..239f59ac --- /dev/null +++ b/tests/test_operator_flops.py @@ -0,0 +1,20 @@ +from infinimetrics.operators.flops_calculator import FLOPSCalculator + + +def test_mm_flops_uses_matrix_formula(): + inputs = [ + {"shape": [2, 3], "dtype": "float16"}, + {"shape": [3, 4], "dtype": "float16"}, + ] + + assert FLOPSCalculator.get_flops("mm", inputs, [{"shape": [2, 4]}]) == 48 + + +def test_linear_flops_uses_input_weight_and_bias(): + inputs = [ + {"shape": [2, 3], "dtype": "float16"}, + {"shape": [3, 4], "dtype": "float16"}, + {"shape": [4], "dtype": "float16"}, + ] + + assert FLOPSCalculator.get_flops("linear", inputs, [{"shape": [2, 4]}]) == 56