From d11e3cf9c4fc2c7d8f430e0b45cea023f4209dec Mon Sep 17 00:00:00 2001 From: Florent Li <1508269885@qq.com> Date: Mon, 3 Aug 2026 14:54:48 +0800 Subject: [PATCH] add operator benchmark utility scripts --- scripts/README.md | 184 ++++++---- scripts/aggregate_results.py | 461 +++++++++++++++++++++++ scripts/generate_operator_inputs.py | 543 ++++++++++++++++++++++++++++ tests/test_operator_scripts.py | 78 ++++ 4 files changed, 1199 insertions(+), 67 deletions(-) create mode 100644 scripts/aggregate_results.py create mode 100644 scripts/generate_operator_inputs.py create mode 100644 tests/test_operator_scripts.py diff --git a/scripts/README.md b/scripts/README.md index 4c9b090d..8468d0e7 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -1,14 +1,15 @@ # Testing Scripts -Unified test execution scripts for InfiniMetrics. +Utilities for running InfiniMetrics tests, generating operator inputs, and +aggregating benchmark results. ## Quick Start ```bash -# Run tests with input file(s) +# Run one input file ./scripts/run_tests.sh test.json -# Run tests in a directory +# Run every input in a directory ./scripts/run_tests.sh test_dir/ # Run multiple input files @@ -17,99 +18,148 @@ Unified test execution scripts for InfiniMetrics. ## Structure -``` +```text scripts/ -├── run_tests.sh # Unified test execution script -└── common/ # Shared utilities - ├── install_deps.sh # Dependency management (check + install) - └── prepare_env.sh # Environment preparation functions +|-- run_tests.sh +|-- generate_operator_inputs.py +|-- aggregate_results.py +`-- common/ + |-- install_deps.sh + `-- prepare_env.sh ``` -## Script Organization - -### Main Script: `run_tests.sh` +## Test Runner -Unified test execution script with automatic dependency management. +`run_tests.sh` is the unified entry point with optional dependency checks. -**Usage (Direct Execution):** ```bash ./scripts/run_tests.sh [OPTIONS] ``` -**Usage (Source Mode - Environment Variables Persist):** +Options: + +- `--check ` checks comma-separated dependency groups: `hardware`, + `operator`, or `all`. +- `--no-check` skips dependency checks. +- `--help`, `-h` prints help. + +The script can also be sourced when environment changes must remain in the +current shell: + ```bash source scripts/run_tests.sh -run_tests [OPTIONS] +run_tests --check operator test.json ``` -**Options:** -```bash ---check Check specific dependencies before running (comma-separated) - Types: hardware, operator, all ---no-check Skip dependency checking ---help, -h Show help message -``` +## InfiniOps Adapter -**Input paths:** -- Can be JSON files or directories +Operator testcases whose framework segment is `InfiniOps` are dispatched to +the InfiniOps adapter. Cambricon and Ascend input configurations generated by +this directory use that framework automatically. -**Examples:** -```bash -# Direct execution (recommended for CI/automation) -./scripts/run_tests.sh test.json -./scripts/run_tests.sh test_dir/ -./scripts/run_tests.sh test1.json test2.json -./scripts/run_tests.sh --check hardware test.json +Runtime prerequisites depend on the selected device: -# Source mode (recommended for development) -source scripts/run_tests.sh -run_tests test.json -run_tests --check all test.json -``` +- InfiniOps Python bindings (`infini.ops`) +- PyTorch +- `torch_mlu` for Cambricon MLU +- `torch_npu` for Ascend NPU -### Common Functions (`common/`) +The adapter benchmarks the first implementation registered by InfiniOps, +including an ATen fallback when the runtime reports one. If the installed +InfiniOps build has no implementation for an operator and device, the test +fails with an actionable error instead of invoking an unregistered slot. It +reports latency, tensor accuracy, estimated TFLOPS, and estimated memory bandwidth. +An accuracy mismatch produces a nonzero `result_code`, so failed correctness +checks are also counted as failed runs by the result aggregator. -**`install_deps.sh`**: Unified dependency management (check + install) +## Operator Input Generator -Can be used standalone or sourced by other scripts. +`generate_operator_inputs.py` creates InfiniMetrics-compatible JSON files and +NumPy tensor data for small, medium, and large shape groups. -**Standalone usage:** ```bash -# Install specific component -export INFINICORE_PATH="/path/to/InfiniCore" -source scripts/common/install_deps.sh operator # Install InfiniCore -source scripts/common/install_deps.sh hardware # Build CUDA benchmark -source scripts/common/install_deps.sh all # Install everything +python scripts/generate_operator_inputs.py [OPTIONS] ``` -**Components:** -- `operator` - InfiniCore (operator testing) -- `hardware` - CUDA memory benchmark (hardware testing) +| Option | Default | Description | +|---|---|---| +| `--output`, `-o` | `./operator_test_inputs` | Output directory | +| `--operators` | `matmul add sub mul div` | Operators to generate | +| `--dtypes` | `float16 float32 bfloat16` | Tensor data types | +| `--scales` | `small medium large` | Shape groups | +| `--device` | `nvidia` | Target platform | +| `--seed` | `42` | Random seed | +| `--warmup` | `10` | Non-negative warmup iterations | +| `--measured` | `100` | Positive measured iterations | +| `--dry-run` | disabled | Print the plan without writing files | -**Checking functions** (available when sourced): -- `check_cuda` - Check NVIDIA CUDA toolkit -- `check_infinicore` - Check InfiniCore package +Examples: -**Installation functions** (available when sourced): -- `install_infinicore` - Install InfiniCore from source -- `install_hardware` - Build CUDA memory benchmark +```bash +# Preview Cambricon testcases without writing data +python scripts/generate_operator_inputs.py \ + --device cambricon --operators matmul add --dry-run + +# Generate selected operators and data types +python scripts/generate_operator_inputs.py \ + --output ./test_inputs \ + --device cambricon \ + --operators matmul add sub mul div \ + --dtypes float16 float32 + +# Run the generated configurations +python main.py ./test_inputs/configs/ +``` + +Generated layout: + +```text +operator_test_inputs/ +|-- configs/ +| `-- opbench.....json +|-- data/ +| `-- ____.npy +|-- all_test_inputs.json +`-- _generation_metadata.json +``` + +Supported generator operators are `matmul`, `mm`, `add`, `sub`, `mul`, `div`, +and `linear`. The adapter also accepts InfiniOps configurations for the common +`cast`, `cat`, and `gemm` operators. -**`prepare_env.sh`**: Environment preparation functions -- `log_test_start` - Log test start message with timestamp -- `log_test_end` - Log test completion with exit code -- `cleanup_on_error` - Error trap handler -- `get_timestamp` - Get current timestamp +## Result Aggregator -## Output +`aggregate_results.py` discovers `*_results.json` files recursively and builds +pass/fail and metric summaries by operator, scale, and dtype. -All test results are saved to: +```bash +python scripts/aggregate_results.py [OPTIONS] ``` -output/ + +| Option | Default | Description | +|---|---|---| +| `--input`, `-i` | `./output` | Result directory | +| `--output`, `-o` | `./aggregated_results.json` | Summary JSON path | +| `--print` | disabled | Print a console summary | +| `--filter-operator` | none | Include one operator | +| `--filter-scale` | none | Include one scale | +| `--filter-dtype` | none | Include one dtype | + +Examples: + +```bash +python scripts/aggregate_results.py --print +python scripts/aggregate_results.py -i ./output -o ./summary.json --print +python scripts/aggregate_results.py --filter-operator matmul --print +python scripts/aggregate_results.py \ + --filter-scale large --filter-dtype float16 --print ``` -## Requirements +Malformed JSON files are skipped with a warning. Metrics with nonnumeric values +remain in detailed records but are excluded from min/max/average calculations. + +## Common Helpers -- Python 3.10+ -- Bash 4.0+ -- CUDA toolkit (for CUDA hardware tests) -- InfiniCore source (for operator tests) +`common/install_deps.sh` provides dependency checks and installers for operator +and hardware tests. `common/prepare_env.sh` provides test logging, timestamps, +and cleanup helpers used by the runner. diff --git a/scripts/aggregate_results.py b/scripts/aggregate_results.py new file mode 100644 index 00000000..307125e3 --- /dev/null +++ b/scripts/aggregate_results.py @@ -0,0 +1,461 @@ +#!/usr/bin/env python3 +""" +Test Results Aggregator for InfiniMetrics + +Reads individual result JSON files from the output/ directory and produces a +structured summary with per-operator, per-scale, per-dtype breakdowns, including +metric values (latency, TFLOPS, bandwidth) and pass/fail statistics. + +Usage: + # Aggregate all results in default output directory + python scripts/aggregate_results.py + + # Specify input and output paths + python scripts/aggregate_results.py --input ./output --output ./summary.json + + # Print human-readable table to console + python scripts/aggregate_results.py --print + + # Filter by operator / scale / dtype + python scripts/aggregate_results.py --filter-operator matmul --filter-scale large +""" + +import argparse +import json +import logging +import numbers +import sys +from collections import defaultdict +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional + +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def discover_result_files(root: Path) -> List[Path]: + """Recursively find all *_results.json files under root.""" + return sorted(root.rglob("*_results.json")) + + +def parse_run_id(run_id: str) -> Dict[str, str]: + """ + Parse run_id like 'opbench.matmul.small.128x128.float16' into components. + Falls back to empty strings if the format doesn't match. + """ + parts = run_id.split(".") + if len(parts) >= 5 and parts[0] == "opbench": + return { + "operator": parts[1], + "scale": parts[2], + "shape": parts[3], + "dtype": parts[4], + } + return {"operator": "", "scale": "", "shape": "", "dtype": ""} + + +def load_result(path: Path) -> Optional[Dict[str, Any]]: + """Load a single result JSON file.""" + try: + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + except Exception as e: + logger.warning(f"Failed to load {path}: {e}") + return None + + +def extract_metrics(result: Dict[str, Any]) -> Dict[str, Any]: + """Extract metric values from a result dict.""" + metrics_out = {} + for m in result.get("metrics", []): + name = m.get("name", "") + value = m.get("value") + unit = m.get("unit", "") + metrics_out[name] = {"value": value, "unit": unit} + return metrics_out + + +# --------------------------------------------------------------------------- +# Aggregation logic +# --------------------------------------------------------------------------- + + +def aggregate( + files: List[Path], + filter_operator: Optional[str] = None, + filter_scale: Optional[str] = None, + filter_dtype: Optional[str] = None, +) -> Dict[str, Any]: + """Aggregate results from all files into a structured summary.""" + + all_records: List[Dict[str, Any]] = [] + + # Counters + total = 0 + passed = 0 + failed = 0 + + # Group-by containers + by_operator: Dict[str, Dict] = defaultdict( + lambda: {"passed": 0, "failed": 0, "records": []} + ) + by_scale: Dict[str, Dict] = defaultdict(lambda: {"passed": 0, "failed": 0}) + by_dtype: Dict[str, Dict] = defaultdict(lambda: {"passed": 0, "failed": 0}) + by_operator_dtype: Dict[str, Dict] = defaultdict( + lambda: {"passed": 0, "failed": 0, "records": []} + ) + + for f in files: + result = load_result(f) + if result is None: + continue + + run_id = result.get("run_id", f.stem) + info = parse_run_id(run_id) + config = result.get("config", {}) + inputs = config.get("inputs") or [{}] + operator = info["operator"] or config.get("operator", "") + dtype = info["dtype"] or inputs[0].get("dtype", "") + + # Apply filters + if filter_operator and operator != filter_operator: + continue + if filter_scale and info["scale"] != filter_scale: + continue + if filter_dtype and dtype != filter_dtype: + continue + + total += 1 + rc = result.get("result_code", -1) + is_pass = rc == 0 + if is_pass: + passed += 1 + else: + failed += 1 + + metrics = extract_metrics(result) + error_msg = result.get("error_msg") + + record = { + "run_id": run_id, + "testcase": result.get("testcase", ""), + "operator": operator, + "scale": info["scale"], + "shape": info["shape"], + "dtype": dtype, + "result_code": rc, + "passed": is_pass, + "error_msg": error_msg, + "time": result.get("time", ""), + "duration_sec": result.get("duration"), + "metrics": metrics, + "source_file": str(f), + } + all_records.append(record) + + op = record["operator"] + scale = record["scale"] + dtype = record["dtype"] + + by_operator[op]["passed" if is_pass else "failed"] += 1 + by_operator[op]["records"].append(record) + by_scale[scale]["passed" if is_pass else "failed"] += 1 + by_dtype[dtype]["passed" if is_pass else "failed"] += 1 + op_dt_key = f"{op}/{dtype}" + by_operator_dtype[op_dt_key]["passed" if is_pass else "failed"] += 1 + by_operator_dtype[op_dt_key]["records"].append(record) + + # Build per-operator metric summaries (only for passed tests with values) + operator_metrics_summary: Dict[str, Dict] = {} + for op, data in by_operator.items(): + latency_values = [] + flops_values = [] + bw_values = [] + for rec in data["records"]: + if not rec["passed"]: + continue + m = rec["metrics"] + lat = m.get("operator.latency", {}).get("value") + if isinstance(lat, numbers.Real) and not isinstance(lat, bool): + latency_values.append(lat) + fl = m.get("operator.flops", {}).get("value") + if isinstance(fl, numbers.Real) and not isinstance(fl, bool): + flops_values.append(fl) + bw = m.get("operator.bandwidth", {}).get("value") + if isinstance(bw, numbers.Real) and not isinstance(bw, bool): + bw_values.append(bw) + + summary: Dict[str, Any] = { + "passed": data["passed"], + "failed": data["failed"], + "total": data["passed"] + data["failed"], + } + if latency_values: + summary["latency_ms"] = { + "min": round(min(latency_values), 4), + "max": round(max(latency_values), 4), + "avg": round(sum(latency_values) / len(latency_values), 4), + "count": len(latency_values), + } + if flops_values: + summary["tflops"] = { + "min": round(min(flops_values), 4), + "max": round(max(flops_values), 4), + "avg": round(sum(flops_values) / len(flops_values), 4), + "count": len(flops_values), + } + if bw_values: + summary["bandwidth_gbs"] = { + "min": round(min(bw_values), 4), + "max": round(max(bw_values), 4), + "avg": round(sum(bw_values) / len(bw_values), 4), + "count": len(bw_values), + } + operator_metrics_summary[op] = summary + + # Build the detailed table (one row per test case) + detailed_table = [] + for rec in all_records: + row = { + "run_id": rec["run_id"], + "operator": rec["operator"], + "scale": rec["scale"], + "shape": rec["shape"], + "dtype": rec["dtype"], + "passed": rec["passed"], + "result_code": rec["result_code"], + "error_msg": rec["error_msg"], + } + for metric_name, metric_val in rec["metrics"].items(): + short = metric_name.replace("operator.", "") + row[short] = metric_val.get("value") + detailed_table.append(row) + + return { + "generated_at": datetime.now().isoformat(), + "total": total, + "passed": passed, + "failed": failed, + "pass_rate": f"{passed / total * 100:.1f}%" if total > 0 else "N/A", + "by_operator": operator_metrics_summary, + "by_scale": { + k: { + "passed": v["passed"], + "failed": v["failed"], + "total": v["passed"] + v["failed"], + } + for k, v in by_scale.items() + }, + "by_dtype": { + k: { + "passed": v["passed"], + "failed": v["failed"], + "total": v["passed"] + v["failed"], + } + for k, v in by_dtype.items() + }, + "by_operator_dtype": { + k: {"passed": v["passed"], "failed": v["failed"]} + for k, v in by_operator_dtype.items() + }, + "failed_details": [ + { + "run_id": r["run_id"], + "operator": r["operator"], + "scale": r["scale"], + "shape": r["shape"], + "dtype": r["dtype"], + "error_msg": r["error_msg"], + } + for r in all_records + if not r["passed"] + ], + "detailed_table": detailed_table, + } + + +# --------------------------------------------------------------------------- +# Console printing +# --------------------------------------------------------------------------- + + +def print_summary(agg: Dict[str, Any]) -> None: + """Print a human-readable summary to console.""" + + print("=" * 72) + print(" InfiniMetrics Test Results Summary") + print("=" * 72) + print(f" Total : {agg['total']}") + print(f" Passed : {agg['passed']}") + print(f" Failed : {agg['failed']}") + print(f" Rate : {agg['pass_rate']}") + print("=" * 72) + + # By operator + print("\n-- By Operator --") + print( + f" {'Operator':<12} {'Total':>5} {'Passed':>6} {'Failed':>6} {'Latency(ms)':>14} {'TFLOPS':>10} {'BW(GB/s)':>10}" + ) + print(" " + "-" * 68) + for op, data in sorted(agg["by_operator"].items()): + lat = data.get("latency_ms", {}) + lat_str = f"{lat['avg']:.4f}" if lat else "-" + fl = data.get("tflops", {}) + fl_str = f"{fl['avg']:.4f}" if fl else "-" + bw = data.get("bandwidth_gbs", {}) + bw_str = f"{bw['avg']:.4f}" if bw else "-" + print( + f" {op:<12} {data['total']:>5} {data['passed']:>6} {data['failed']:>6} " + f"{lat_str:>14} {fl_str:>10} {bw_str:>10}" + ) + + # By scale + print("\n-- By Scale --") + print(f" {'Scale':<10} {'Total':>5} {'Passed':>6} {'Failed':>6}") + print(" " + "-" * 30) + for scale, data in sorted(agg["by_scale"].items()): + print( + f" {scale:<10} {data['total']:>5} {data['passed']:>6} {data['failed']:>6}" + ) + + # By dtype + print("\n-- By Dtype --") + print(f" {'Dtype':<10} {'Total':>5} {'Passed':>6} {'Failed':>6}") + print(" " + "-" * 30) + for dtype, data in sorted(agg["by_dtype"].items()): + print( + f" {dtype:<10} {data['total']:>5} {data['passed']:>6} {data['failed']:>6}" + ) + + # Failed details + if agg["failed_details"]: + print(f"\n-- Failed Tests ({len(agg['failed_details'])}) --") + print(f" {'Run ID':<45} {'Error':<30}") + print(" " + "-" * 75) + for fd in agg["failed_details"]: + rid = fd["run_id"] + err = (fd["error_msg"] or "Unknown")[:30] + print(f" {rid:<45} {err:<30}") + + # Detailed table for passed tests with metrics + passed_rows = [r for r in agg["detailed_table"] if r["passed"]] + if passed_rows: + print(f"\n-- Passed Tests Detail ({len(passed_rows)}) --") + print( + f" {'Operator':<10} {'Scale':<8} {'Shape':<14} {'Dtype':<10} " + f"{'Latency(ms)':>12} {'TFLOPS':>10} {'BW(GB/s)':>10} {'Accuracy':>10}" + ) + print(" " + "-" * 90) + for r in passed_rows: + lat = r.get("latency") + lat_str = f"{lat:.4f}" if lat is not None else "-" + fl = r.get("flops") + fl_str = f"{fl:.4f}" if fl is not None else "-" + bw = r.get("bandwidth") + bw_str = f"{bw:.4f}" if bw is not None else "-" + acc = r.get("tensor_accuracy", "-") + print( + f" {r['operator']:<10} {r['scale']:<8} {r['shape']:<14} {r['dtype']:<10} " + f"{lat_str:>12} {fl_str:>10} {bw_str:>10} {str(acc):>10}" + ) + + print("\n" + "=" * 72) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Aggregate InfiniMetrics test results from output directory", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "Examples:\n" + " python scripts/aggregate_results.py\n" + " python scripts/aggregate_results.py --input ./output --print\n" + " python scripts/aggregate_results.py --filter-operator matmul --print\n" + ), + ) + parser.add_argument( + "--input", + "-i", + default="./output", + help="Input directory containing result JSON files (default: ./output)", + ) + parser.add_argument( + "--output", + "-o", + default="./aggregated_results.json", + help="Output JSON file path (default: ./aggregated_results.json)", + ) + parser.add_argument( + "--print", + dest="print_summary", + action="store_true", + help="Print human-readable summary table to console", + ) + parser.add_argument( + "--filter-operator", + default=None, + help="Only include results for this operator", + ) + parser.add_argument( + "--filter-scale", + default=None, + help="Only include results for this scale (small/medium/large)", + ) + parser.add_argument( + "--filter-dtype", + default=None, + help="Only include results for this dtype (float16/float32/bfloat16)", + ) + return parser.parse_args(argv) + + +def main(argv: Optional[List[str]] = None) -> int: + args = parse_args(argv) + + input_dir = Path(args.input) + if not input_dir.exists(): + print(f"Error: Input directory not found: {input_dir}") + return 1 + + files = discover_result_files(input_dir) + if not files: + print(f"No result files found in {input_dir}") + return 1 + + print(f"Found {len(files)} result file(s) in {input_dir}") + + agg = aggregate( + files, + filter_operator=args.filter_operator, + filter_scale=args.filter_scale, + filter_dtype=args.filter_dtype, + ) + + # Save to JSON + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w", encoding="utf-8") as f: + json.dump(agg, f, indent=2, ensure_ascii=False) + print(f"Aggregated results saved to {output_path}") + + # Print to console if requested + if args.print_summary: + print_summary(agg) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/generate_operator_inputs.py b/scripts/generate_operator_inputs.py new file mode 100644 index 00000000..d9f94ee5 --- /dev/null +++ b/scripts/generate_operator_inputs.py @@ -0,0 +1,543 @@ +#!/usr/bin/env python3 +""" +Operator Test Input Generator for InfiniMetrics + +Generates standardized test input data covering small, medium, and large tensor +scales with various shape and dtype combinations. Outputs InfiniMetrics-compatible +JSON configs and .npy data files ready for operator benchmarking. + +Usage: + # Generate inputs for all supported operators + python scripts/generate_operator_inputs.py --output ./test_inputs --seed 42 + + # Generate for specific operators and dtypes + python scripts/generate_operator_inputs.py --operators matmul add --dtypes float16 float32 + + # Generate only small-scale tests + python scripts/generate_operator_inputs.py --scales small medium + + # Dry run (print plan without generating files) + python scripts/generate_operator_inputs.py --dry-run +""" + +import argparse +import json +import logging +import sys +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Shape & dtype definitions +# --------------------------------------------------------------------------- + +# Shapes grouped by scale: (label, shapes_list) +# Each shape is a 2-element tuple for the "base shape" that operators interpret. +MATMUL_SHAPES = { + "small": [ + (64, 64), + (128, 128), + (256, 256), + ], + "medium": [ + (512, 512), + (768, 1024), + (1024, 768), + ], + "large": [ + (1024, 1024), + (2048, 2048), + (4096, 4096), + ], +} + +ELEMENTWISE_SHAPES = { + "small": [ + (64, 64), + (128, 256), + (256, 512), + ], + "medium": [ + (512, 1024), + (1024, 1024), + (2048, 512), + ], + "large": [ + (2048, 2048), + (4096, 4096), + (8192, 1024), + ], +} + +SUPPORTED_DTYPES = ["float16", "float32", "bfloat16"] + +DTYPE_BYTES = {"float16": 2, "float32": 4, "bfloat16": 2} + +# --------------------------------------------------------------------------- +# Operator specs +# --------------------------------------------------------------------------- + +OP_INPUT_NAMES = { + "matmul": ["a", "b"], + "mm": ["a", "b"], + "add": ["a", "b"], + "sub": ["a", "b"], + "mul": ["a", "b"], + "div": ["a", "b"], + "linear": ["input", "weight", "bias"], +} + +# Devices that use InfiniOps framework (ATen fallback) instead of InfiniCore +INFINIOPS_DEVICES = {"cambricon", "ascend"} + + +def non_negative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("must be a non-negative integer") + return parsed + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be a positive integer") + return parsed + + +def get_input_shapes(operator: str, base_shape: Tuple[int, ...]) -> List[List[int]]: + """Derive input shapes from a base shape for a given operator.""" + if operator in ("matmul", "mm"): + m, k = base_shape[0], base_shape[1] + n = k # square by default + return [[m, k], [k, n]] + elif operator == "linear": + m, k = base_shape[0], base_shape[1] + n = k + return [[m, k], [k, n], [n]] + else: + # element-wise: both inputs share the same shape + return [list(base_shape), list(base_shape)] + + +def get_output_shape(operator: str, input_shapes: List[List[int]]) -> List[int]: + """Calculate output shape from input shapes.""" + if operator in ("matmul", "mm"): + m = input_shapes[0][0] + n = input_shapes[1][1] + return [m, n] + elif operator == "linear": + m = input_shapes[0][0] + n = input_shapes[1][0] + return [m, n] + else: + return input_shapes[0].copy() + + +# --------------------------------------------------------------------------- +# Data generation +# --------------------------------------------------------------------------- + +NUMPY_DTYPE_MAP = { + "float16": np.float16, + "float32": np.float32, + "float64": np.float64, + "bfloat16": np.float32, # numpy lacks bfloat16; store as float32 + "int8": np.int8, + "int32": np.int32, +} + + +def generate_tensor( + shape: List[int], + dtype: str, + rng: np.random.Generator, + distribution: str = "uniform", +) -> np.ndarray: + """Generate a random tensor with the given distribution.""" + np_dtype = NUMPY_DTYPE_MAP[dtype] + if distribution == "uniform": + data = rng.uniform(-1.0, 1.0, shape).astype(np_dtype) + elif distribution == "normal": + data = rng.normal(0.0, 1.0, shape).astype(np_dtype) + else: + data = rng.uniform(-1.0, 1.0, shape).astype(np_dtype) + return data + + +# --------------------------------------------------------------------------- +# JSON config builder +# --------------------------------------------------------------------------- + + +def build_test_config( + operator: str, + device: str, + inputs: List[Dict[str, Any]], + output_shape: List[int], + dtype: str, + warmup: int = 10, + measured: int = 100, + atol: float = 1e-3, + rtol: float = 1e-3, +) -> Dict[str, Any]: + """Build an InfiniMetrics-compatible test input dict.""" + framework = "InfiniOps" if device.lower() in INFINIOPS_DEVICES else "InfiniCore" + if device.lower() == "cambricon": + atol = rtol = 0.01 + return { + "run_id": f"opbench.{operator}._", + "testcase": f"operator.{framework}.{operator.capitalize()}", + "config": { + "operator": operator, + "device": device, + "data_base_dir": "", + "inputs": inputs, + "outputs": [ + { + "name": "output", + "shape": output_shape, + "dtype": dtype, + } + ], + "warmup_iterations": warmup, + "measured_iterations": measured, + "tolerance": {"atol": atol, "rtol": rtol}, + }, + "metrics": [ + {"name": "operator.latency"}, + {"name": "operator.tensor_accuracy"}, + {"name": "operator.flops"}, + {"name": "operator.bandwidth"}, + ], + } + + +# --------------------------------------------------------------------------- +# Main generation logic +# --------------------------------------------------------------------------- + + +@dataclass +class TestCase: + """Represents one test case to generate.""" + + operator: str + scale: str + shape_label: str + base_shape: Tuple[int, ...] + dtype: str + index: int # unique within the whole batch + + +def enumerate_test_cases( + operators: List[str], + scales: List[str], + dtypes: List[str], +) -> List[TestCase]: + """Enumerate all test case combinations.""" + cases: List[TestCase] = [] + idx = 0 + for op in operators: + shape_table = ( + MATMUL_SHAPES if op in ("matmul", "mm", "linear") else ELEMENTWISE_SHAPES + ) + for scale in scales: + shapes = shape_table.get(scale, []) + for shape in shapes: + for dtype in dtypes: + cases.append( + TestCase( + operator=op, + scale=scale, + shape_label=f"{shape[0]}x{shape[1]}", + base_shape=shape, + dtype=dtype, + index=idx, + ) + ) + idx += 1 + return cases + + +def format_size_mb(shape: List[int], dtype: str) -> float: + """Calculate tensor size in MB.""" + elements = 1 + for d in shape: + elements *= d + return elements * DTYPE_BYTES.get(dtype, 4) / (1024 * 1024) + + +def generate_all( + cases: List[TestCase], + output_dir: Path, + device: str, + seed: int, + warmup: int, + measured: int, +) -> List[Path]: + """Generate data files and JSON configs for all test cases.""" + rng = np.random.default_rng(seed) + data_dir = output_dir / "data" + data_dir.mkdir(parents=True, exist_ok=True) + config_dir = output_dir / "configs" + config_dir.mkdir(parents=True, exist_ok=True) + + # Also produce a single combined JSON with all test cases + all_configs: List[Dict[str, Any]] = [] + + written: List[Path] = [] + total_data_mb = 0.0 + + for tc in cases: + input_shapes = get_input_shapes(tc.operator, tc.base_shape) + output_shape = get_output_shape(tc.operator, input_shapes) + input_names = OP_INPUT_NAMES.get(tc.operator, ["a", "b"]) + + # Restrict input_names to actual number of inputs + input_names = input_names[: len(input_shapes)] + + inputs_config: List[Dict[str, Any]] = [] + for name, shape in zip(input_names, input_shapes): + # Generate data + data = generate_tensor(shape, tc.dtype, rng) + shape_str = "x".join(str(d) for d in shape) + npy_name = f"{tc.operator}_{tc.scale}_{shape_str}_{tc.dtype}_{name}.npy" + npy_path = data_dir / npy_name + np.save(npy_path, data) + + total_data_mb += format_size_mb(shape, tc.dtype) + + inputs_config.append( + { + "name": name, + "shape": shape, + "dtype": tc.dtype, + "file_path": str(npy_path.resolve()), + "init_mode": "random", + } + ) + + # Build run_id + run_id = f"opbench.{tc.operator}.{tc.scale}.{tc.shape_label}.{tc.dtype}" + + config = build_test_config( + operator=tc.operator, + device=device, + inputs=inputs_config, + output_shape=output_shape, + dtype=tc.dtype, + warmup=warmup, + measured=measured, + ) + config["run_id"] = run_id + config["config"]["data_base_dir"] = str(data_dir.resolve()) + + # Write individual JSON + json_name = f"{run_id}.json" + json_path = config_dir / json_name + with open(json_path, "w", encoding="utf-8") as f: + json.dump(config, f, indent=2, ensure_ascii=False) + written.append(json_path) + + all_configs.append(config) + + logger.info( + f" [{tc.index + 1}/{len(cases)}] {run_id} " + f"(data ~{format_size_mb(output_shape, tc.dtype):.2f} MB output)" + ) + + # Write combined JSON (list of all configs) + combined_path = output_dir / "all_test_inputs.json" + with open(combined_path, "w", encoding="utf-8") as f: + json.dump(all_configs, f, indent=2, ensure_ascii=False) + written.append(combined_path) + + logger.info(f"Total data generated: {total_data_mb:.2f} MB") + return written + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate standardized operator test inputs for InfiniMetrics", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "Examples:\n" + " python scripts/generate_operator_inputs.py --output ./test_inputs --seed 42\n" + " python scripts/generate_operator_inputs.py --operators matmul --dtypes float16\n" + " python scripts/generate_operator_inputs.py --dry-run\n" + ), + ) + parser.add_argument( + "--output", + "-o", + default="./operator_test_inputs", + help="Output directory (default: ./operator_test_inputs)", + ) + parser.add_argument( + "--operators", + nargs="+", + default=["matmul", "add", "sub", "mul", "div"], + choices=["matmul", "mm", "add", "sub", "mul", "div", "linear"], + help="Operators to generate inputs for (default: matmul add sub mul div)", + ) + parser.add_argument( + "--dtypes", + nargs="+", + default=["float16", "float32", "bfloat16"], + choices=SUPPORTED_DTYPES, + help="Data types (default: float16 float32 bfloat16)", + ) + parser.add_argument( + "--scales", + nargs="+", + default=["small", "medium", "large"], + choices=["small", "medium", "large"], + help="Tensor scale categories (default: small medium large)", + ) + parser.add_argument( + "--device", + default="nvidia", + help="Target device (default: nvidia)", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed for reproducibility (default: 42)", + ) + parser.add_argument( + "--warmup", + type=non_negative_int, + default=10, + help="Warmup iterations (default: 10)", + ) + parser.add_argument( + "--measured", + type=positive_int, + default=100, + help="Measured iterations (default: 100)", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print plan without generating files", + ) + return parser.parse_args(argv) + + +def main(argv: Optional[List[str]] = None) -> int: + args = parse_args(argv) + + cases = enumerate_test_cases(args.operators, args.scales, args.dtypes) + + # Group by operator for summary + from collections import Counter + + op_counts = Counter(c.operator for c in cases) + scale_counts = Counter(c.scale for c in cases) + dtype_counts = Counter(c.dtype for c in cases) + + print("=" * 60) + print("Operator Test Input Generator") + print("=" * 60) + print(f" Total test cases : {len(cases)}") + print(f" Operators : {dict(op_counts)}") + print(f" Scales : {dict(scale_counts)}") + print(f" Dtypes : {dict(dtype_counts)}") + print(f" Device : {args.device}") + print(f" Seed : {args.seed}") + print(f" Output : {args.output}") + + # Estimate total data size + est_mb = 0.0 + for tc in cases: + input_shapes = get_input_shapes(tc.operator, tc.base_shape) + for s in input_shapes: + est_mb += format_size_mb(s, tc.dtype) + print(f" Est. data size : {est_mb:.1f} MB") + print("=" * 60) + + if args.dry_run: + print("\nDry run - test case details:") + print("-" * 60) + print(f"{'#':<4} {'Operator':<10} {'Scale':<8} {'Shape':<14} {'Dtype':<10}") + print("-" * 60) + for tc in cases: + input_shapes = get_input_shapes(tc.operator, tc.base_shape) + shapes_str = " @ ".join("x".join(str(d) for d in s) for s in input_shapes) + print( + f"{tc.index:<4} {tc.operator:<10} {tc.scale:<8} " + f"{shapes_str:<24} {tc.dtype:<10}" + ) + print("-" * 60) + print(f"Total: {len(cases)} test cases") + return 0 + + output_dir = Path(args.output) + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"\nGenerating {len(cases)} test cases...") + written = generate_all( + cases=cases, + output_dir=output_dir, + device=args.device, + seed=args.seed, + warmup=args.warmup, + measured=args.measured, + ) + + # Save generation metadata + metadata = { + "generated_at": datetime.now().isoformat(), + "seed": args.seed, + "device": args.device, + "operators": args.operators, + "scales": args.scales, + "dtypes": args.dtypes, + "warmup": args.warmup, + "measured": args.measured, + "total_cases": len(cases), + "total_files": len(written), + "estimated_data_mb": round(est_mb, 2), + "cases": [ + { + "operator": tc.operator, + "scale": tc.scale, + "base_shape": list(tc.base_shape), + "dtype": tc.dtype, + } + for tc in cases + ], + } + meta_path = output_dir / "_generation_metadata.json" + with open(meta_path, "w", encoding="utf-8") as f: + json.dump(metadata, f, indent=2, ensure_ascii=False) + + print(f"\nDone. Generated {len(written)} files in {output_dir}/") + print(f" configs/ - individual JSON test configs") + print(f" data/ - .npy tensor data files") + print(f" all_test_inputs.json - combined config (use with main.py)") + print(f" _generation_metadata.json - generation info") + print(f"\nUsage:") + print(f" python main.py {output_dir}/configs/") + print(f" python main.py {output_dir}/all_test_inputs.json") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_operator_scripts.py b/tests/test_operator_scripts.py new file mode 100644 index 00000000..b01fd79f --- /dev/null +++ b/tests/test_operator_scripts.py @@ -0,0 +1,78 @@ +import json + +import pytest + +from scripts.aggregate_results import aggregate, discover_result_files +from scripts.generate_operator_inputs import ( + build_test_config, + enumerate_test_cases, + get_input_shapes, + parse_args, +) + + +def test_generated_infiniops_config_matches_adapter_contract(): + config = build_test_config( + operator="sub", + device="cambricon", + inputs=[ + {"name": "a", "shape": [4, 8], "dtype": "float16"}, + {"name": "b", "shape": [4, 8], "dtype": "float16"}, + ], + output_shape=[4, 8], + dtype="float16", + ) + + assert config["testcase"] == "operator.InfiniOps.Sub" + assert config["config"]["device"] == "cambricon" + assert config["config"]["tolerance"] == {"atol": 0.01, "rtol": 0.01} + + +def test_case_enumeration_is_deterministic(): + cases = enumerate_test_cases(["mm", "add"], ["small"], ["float16"]) + + assert [case.index for case in cases] == list(range(len(cases))) + assert [(case.operator, case.shape_label) for case in cases[:2]] == [ + ("mm", "64x64"), + ("mm", "128x128"), + ] + + +def test_linear_generator_uses_matrix_compatible_weight_shape(): + assert get_input_shapes("linear", (4, 8)) == [[4, 8], [8, 8], [8]] + + +@pytest.mark.parametrize("args", [["--measured", "0"], ["--warmup", "-1"]]) +def test_cli_rejects_invalid_iteration_counts(args): + with pytest.raises(SystemExit): + parse_args(args) + + +def test_aggregate_falls_back_to_config_for_nonstandard_run_id(tmp_path): + result_path = tmp_path / "custom_results.json" + result_path.write_text( + json.dumps( + { + "run_id": "custom-run", + "testcase": "operator.InfiniOps.Add", + "result_code": 0, + "duration": 1.25, + "config": { + "operator": "add", + "inputs": [{"dtype": "float16"}], + }, + "metrics": [{"name": "operator.latency", "value": "unavailable"}], + } + ), + encoding="utf-8", + ) + + summary = aggregate( + discover_result_files(tmp_path), + filter_operator="add", + filter_dtype="float16", + ) + + assert summary["total"] == 1 + assert summary["by_operator"]["add"]["total"] == 1 + assert "latency_ms" not in summary["by_operator"]["add"]