diff --git a/.gitignore b/.gitignore index 4e70a18..03554fc 100644 --- a/.gitignore +++ b/.gitignore @@ -51,6 +51,11 @@ experiments/ nccl_thread_sweep_*.log nccl_thread_sweep_summary_*.txt +# Test data and outputs (tests/gemm_analysis/) +expected_outputs/ +testdata/ +actual_outputs/ + # IDE/project-specific folders .vscode/ .idea/ diff --git a/docs/comprehensive_report.html b/docs/comprehensive_report.html index 803cb60..445aa0f 100644 --- a/docs/comprehensive_report.html +++ b/docs/comprehensive_report.html @@ -201,7 +201,7 @@

GEMM RCCL Overlap Comprehensive Performance Report

Complete Analysis: rocm-7.0.8-meta (Baseline) vs rocm-7.0.10-meta (Test)
Generated: 2025-12-08 14:59:23
- +

Quick Navigation

@@ -216,13 +216,13 @@

Quick Navigation

512t/70ch
- +

Executive Summary

This comprehensive report contains all performance analysis plots and metrics for RCCL comparisons across multiple configurations.

- +

Test Configuration

- +
- +

Configuration: 256 Threads, 28 Channels

- +
Base: rocm-7.0.8-meta
Test: rocm-7.0.10-meta
- +

Overview Plots

- +
Percentage Change Overview
Percentage Change Overview
- +
Absolute Time Comparison
Absolute Time Comparison
- +
Performance Heatmap
Performance Heatmap
- +
Total Execution Time by Rank
Total Execution Time by Rank
- +
- +

Detailed Metrics

- +
Computation Time Across Ranks
Computation Time Across Ranks
- +
Communication Time Across Ranks
Communication Time Across Ranks
- +
Idle Time Across Ranks
Idle Time Across Ranks
- +
Percentage Difference All Metrics
Percentage Difference All Metrics
- +
- +

NCCL Analysis

- +
NCCL Latency Analysis
NCCL Latency Analysis
- +
NCCL Summary Analysis
NCCL Summary Analysis
- +
- +
Back to Summary
- +

Configuration: 256 Threads, 42 Channels

- +
Base: rocm-7.0.8-meta
Test: rocm-7.0.10-meta
- +

Overview Plots

- +
Percentage Change Overview
Percentage Change Overview
- +
Absolute Time Comparison
Absolute Time Comparison
- +
Performance Heatmap
Performance Heatmap
- +
Total Execution Time by Rank
Total Execution Time by Rank
- +
- +

Detailed Metrics

- +
Computation Time Across Ranks
Computation Time Across Ranks
- +
Communication Time Across Ranks
Communication Time Across Ranks
- +
Idle Time Across Ranks
Idle Time Across Ranks
- +
Percentage Difference All Metrics
Percentage Difference All Metrics
- +
- +

NCCL Analysis

- +
NCCL Latency Analysis
NCCL Latency Analysis
- +
NCCL Summary Analysis
NCCL Summary Analysis
- +
- +
Back to Summary
- +

Configuration: 256 Threads, 56 Channels

- +
Base: rocm-7.0.8-meta
Test: rocm-7.0.10-meta
- +

Overview Plots

- +
Percentage Change Overview
Percentage Change Overview
- +
Absolute Time Comparison
Absolute Time Comparison
- +
Performance Heatmap
Performance Heatmap
- +
Total Execution Time by Rank
Total Execution Time by Rank
- +
- +

Detailed Metrics

- +
Computation Time Across Ranks
Computation Time Across Ranks
- +
Communication Time Across Ranks
Communication Time Across Ranks
- +
Idle Time Across Ranks
Idle Time Across Ranks
- +
Percentage Difference All Metrics
Percentage Difference All Metrics
- +
- +

NCCL Analysis

- +
NCCL Latency Analysis
NCCL Latency Analysis
- +
NCCL Summary Analysis
NCCL Summary Analysis
- +
- +
Back to Summary
- +

Configuration: 256 Threads, 70 Channels

- +
Base: rocm-7.0.8-meta
Test: rocm-7.0.10-meta
- +

Overview Plots

- +
Percentage Change Overview
Percentage Change Overview
- +
Absolute Time Comparison
Absolute Time Comparison
- +
Performance Heatmap
Performance Heatmap
- +
Total Execution Time by Rank
Total Execution Time by Rank
- +
- +

Detailed Metrics

- +
Computation Time Across Ranks
Computation Time Across Ranks
- +
Communication Time Across Ranks
Communication Time Across Ranks
- +
Idle Time Across Ranks
Idle Time Across Ranks
- +
Percentage Difference All Metrics
Percentage Difference All Metrics
- +
- +

NCCL Analysis

- +
NCCL Latency Analysis
NCCL Latency Analysis
- +
NCCL Summary Analysis
NCCL Summary Analysis
- +
- +
Back to Summary
- +

Configuration: 512 Threads, 28 Channels

- +
Base: rocm-7.0.8-meta
Test: rocm-7.0.10-meta
- +

Overview Plots

- +
Percentage Change Overview
Percentage Change Overview
- +
Absolute Time Comparison
Absolute Time Comparison
- +
Performance Heatmap
Performance Heatmap
- +
Total Execution Time by Rank
Total Execution Time by Rank
- +
- +

Detailed Metrics

- +
Computation Time Across Ranks
Computation Time Across Ranks
- +
Communication Time Across Ranks
Communication Time Across Ranks
- +
Idle Time Across Ranks
Idle Time Across Ranks
- +
Percentage Difference All Metrics
Percentage Difference All Metrics
- +
- +

NCCL Analysis

- +
NCCL Latency Analysis
NCCL Latency Analysis
- +
NCCL Summary Analysis
NCCL Summary Analysis
- +
- +
Back to Summary
- +

Configuration: 512 Threads, 42 Channels

- +
Base: rocm-7.0.8-meta
Test: rocm-7.0.10-meta
- +

Overview Plots

- +
Percentage Change Overview
Percentage Change Overview
- +
Absolute Time Comparison
Absolute Time Comparison
- +
Performance Heatmap
Performance Heatmap
- +
Total Execution Time by Rank
Total Execution Time by Rank
- +
- +

Detailed Metrics

- +
Computation Time Across Ranks
Computation Time Across Ranks
- +
Communication Time Across Ranks
Communication Time Across Ranks
- +
Idle Time Across Ranks
Idle Time Across Ranks
- +
Percentage Difference All Metrics
Percentage Difference All Metrics
- +
- +

NCCL Analysis

- +
NCCL Latency Analysis
NCCL Latency Analysis
- +
NCCL Summary Analysis
NCCL Summary Analysis
- +
- +
Back to Summary
- +

Configuration: 512 Threads, 56 Channels

- +
Base: rocm-7.0.8-meta
Test: rocm-7.0.10-meta
- +

Overview Plots

- +
Percentage Change Overview
Percentage Change Overview
- +
Absolute Time Comparison
Absolute Time Comparison
- +
Performance Heatmap
Performance Heatmap
- +
Total Execution Time by Rank
Total Execution Time by Rank
- +
- +

Detailed Metrics

- +
Computation Time Across Ranks
Computation Time Across Ranks
- +
Communication Time Across Ranks
Communication Time Across Ranks
- +
Idle Time Across Ranks
Idle Time Across Ranks
- +
Percentage Difference All Metrics
Percentage Difference All Metrics
- +
- +

NCCL Analysis

- +
NCCL Latency Analysis
NCCL Latency Analysis
- +
NCCL Summary Analysis
NCCL Summary Analysis
- +
- +
Back to Summary
- +

Configuration: 512 Threads, 70 Channels

- +
Base: rocm-7.0.8-meta
Test: rocm-7.0.10-meta
- +

Overview Plots

- +
Percentage Change Overview
Percentage Change Overview
- +
Absolute Time Comparison
Absolute Time Comparison
- +
Performance Heatmap
Performance Heatmap
- +
Total Execution Time by Rank
Total Execution Time by Rank
- +
- +

Detailed Metrics

- +
Computation Time Across Ranks
Computation Time Across Ranks
- +
Communication Time Across Ranks
Communication Time Across Ranks
- +
Idle Time Across Ranks
Idle Time Across Ranks
- +
Percentage Difference All Metrics
Percentage Difference All Metrics
- +
- +

NCCL Analysis

- +
NCCL Latency Analysis
NCCL Latency Analysis
- +
NCCL Summary Analysis
NCCL Summary Analysis
- +
- +
Back to Summary
- +

RCCL Performance Analysis Report

Generated with TraceLens Analysis Suite

diff --git a/scripts/gemm_analysis/run_tracelens_analysis.sh b/scripts/gemm_analysis/run_tracelens_analysis.sh index faf017f..f89d691 100755 --- a/scripts/gemm_analysis/run_tracelens_analysis.sh +++ b/scripts/gemm_analysis/run_tracelens_analysis.sh @@ -45,6 +45,10 @@ fi # Create output directory OUTPUT_DIR="${SWEEP_DIR}/tracelens_analysis" +# Use patched TraceLens entrypoint with GEMM recognition. +# Assume caller activates desired venv; fall back to python3 on PATH. +PYTHON_BIN="${PYTHON_BIN:-python3}" +TL_WRAPPER="scripts/gemm_analysis/tracelense_with_gemm_path.py" if ! mkdir -p "$OUTPUT_DIR" 2>/dev/null; then echo "Error: Cannot create output directory: $OUTPUT_DIR" echo "" @@ -134,14 +138,14 @@ for thread in "${THREAD_CONFIGS[@]}"; do OUTPUT="$OUTPUT_DIR/$thread/individual_reports/perf_${ch}ch_rank${rank}.xlsx" echo " Rank ${rank}..." - TraceLens_generate_perf_report_pytorch \ + "$PYTHON_BIN" "$TL_WRAPPER" generate_perf_report \ --profile_json_path "$TRACE" \ --output_xlsx_path "$OUTPUT" \ --include_unlinked_kernels \ --short_kernel_study \ --short_kernel_threshold_us 50 \ --topk_ops 100 \ - --enable_kernel_summary \ + --enable_kernel_summary \ --topk_roofline_ops 100 echo " [OK] $OUTPUT" @@ -181,16 +185,12 @@ for thread in "${THREAD_CONFIGS[@]}"; do echo "Processing $thread/${ch}ch (all 8 ranks)..." # Use trace_pattern instead of trace_dir for better subdirectory support - # It is not guaranteed that trace files will have the exact same name in all the ranks. - # To avoid file not found errors with `--trace_pattern` flag in TraceLens, we first - # create a directory called `trace` in all rank folders and then mv the respective - # trace file in the rank folder to the canonical `trace/pt.trace.json` path. - # This will satisfy TraceLens's requirement of only one `*` being present in the trace pattern - # while also avoiding FileNotFoundErrors due to different filenames. - find $TRACE_DIR/rank* -name "*.json" -exec sh -c 'mkdir -p "$(dirname "$0")/trace" && mv "$0" "$(dirname "$0")/trace/pt.trace.json"' {} \; - - TraceLens_generate_multi_rank_collective_report_pytorch \ - --trace_pattern "$TRACE_DIR/rank*/trace/pt.trace.json" \ + # Find the trace filename from rank0 + SAMPLE_TRACE=$(find "$TRACE_DIR/rank0" -name "*.json" | head -1) + TRACE_FILENAME=$(basename "$SAMPLE_TRACE") + + "$PYTHON_BIN" "$TL_WRAPPER" generate_multi_rank_collective \ + --trace_pattern "$TRACE_DIR/rank*/$TRACE_FILENAME" \ --world_size 8 \ --output_xlsx_path "$OUTPUT" \ --detailed_analysis \ diff --git a/scripts/gemm_analysis/tracelense_with_gemm_path.py b/scripts/gemm_analysis/tracelense_with_gemm_path.py new file mode 100644 index 0000000..d4d2aac --- /dev/null +++ b/scripts/gemm_analysis/tracelense_with_gemm_path.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +""" +TraceLens with GEMM Recognition Patches + +This script applies GEMM recognition patches and runs TraceLens commands. + +Usage: + python tracelens_with_gemm_patch.py generate_perf_report [args...] + python tracelens_with_gemm_patch.py generate_multi_rank_collective [args...] + python tracelens_with_gemm_patch.py compare_perf_reports [args...] +""" + +import re +import sys + + +def apply_gemm_patches(): + """Apply all GEMM recognition patches to TraceLens.""" + + print("Applying TraceLens GEMM recognition patches...") + + # Patch kernel_name_parser for enhanced ROCm GEMM recognition + try: + from TraceLens.PerfModel import kernel_name_parser + + def patched_is_rocm_gemm(kernel_name): + """ + Enhanced ROCm GEMM pattern matching for Tensile kernels. + Recognizes: Cijk_Alik_Bljk_... and variants with arbitrary prefixes. + """ + pattern = r"^.*C[a-z]{3}_A[a-z]{3}_B[a-z]{3}.*$" + return bool(re.match(pattern, kernel_name)) + + def patched_parse_rocm_gemm(kernel_name): + """Parse ROCm GEMM kernel details.""" + # Parse transpose flags + trans_a, trans_b = None, None + if "_Ailk_" in kernel_name: + trans_a = False + elif "_Alik_" in kernel_name: + trans_a = True + if "_Bljk_" in kernel_name: + trans_b = False + elif "_Bjlk_" in kernel_name: + trans_b = True + + # Parse macro tile size (MT64x16x64) + macro_tile_match = re.search(r"MT(\d+)x(\d+)x(\d+)", kernel_name) + if macro_tile_match: + mt_m = int(macro_tile_match.group(1)) + mt_n = int(macro_tile_match.group(2)) + depth_u = int(macro_tile_match.group(3)) + else: + mt_m, mt_n, depth_u = None, None, None + + return { + "transpose": (trans_a, trans_b), + "mt_m": mt_m, + "mt_n": mt_n, + "depth_u": depth_u, + } + + def patched_gemm_name_parser(kernel_name): + """Enhanced GEMM name parser with better ROCm support.""" + if patched_is_rocm_gemm(kernel_name): + return patched_parse_rocm_gemm(kernel_name) + elif kernel_name_parser.is_cuda_gemm(kernel_name): + return kernel_name_parser.parse_cuda_gemm(kernel_name) + return None + + kernel_name_parser.is_rocm_gemm = patched_is_rocm_gemm + kernel_name_parser.parse_rocm_gemm = patched_parse_rocm_gemm + kernel_name_parser.gemm_name_parser = patched_gemm_name_parser + + print(" [OK] Patched kernel_name_parser (ROCm GEMM recognition)") + except ImportError as e: + print(f" [WARN] Could not patch kernel_name_parser: {e}") + + # Patch Trace2Tree util for is_gemm_kernel function + try: + from TraceLens.Trace2Tree import util as trace_util + + def patched_is_gemm_kernel(kernel_event: dict) -> bool: + """Enhanced GEMM kernel detection.""" + assert kernel_event["cat"] == "kernel" + kernel_name = kernel_event["name"] + + # ROCm Tensile GEMM pattern: C[xyz]_A[xyz]_B[xyz] + pattern = r"^.*C[a-z]{3}_A[a-z]{3}_B[a-z]{3}.*$" + is_rocm_gemm = bool(re.match(pattern, kernel_name)) + + # CUDA GEMM pattern + is_cuda_gemm = kernel_name.startswith("nvjet") or "cublasLt" in kernel_name + + return is_rocm_gemm or is_cuda_gemm + + trace_util.is_gemm_kernel = patched_is_gemm_kernel + print(" [OK] Patched Trace2Tree.util (is_gemm_kernel)") + except ImportError as e: + print(f" [WARN] Could not patch Trace2Tree.util: {e}") + + # Patch TraceEventUtils to enhance GEMM keys + try: + from TraceLens import util as tracelens_util + + if hasattr(tracelens_util, 'TraceEventUtils'): + if hasattr(tracelens_util.TraceEventUtils, 'JaxOpKeys'): + original_gemm_keys = tracelens_util.TraceEventUtils.JaxOpKeys.GemmKeys + enhanced_gemm_keys = ["Cijk", "gemm", "nvjet", "cublasLt", "C[a-z]{3}_A[a-z]{3}_B[a-z]{3}"] + + all_keys = list(set(original_gemm_keys + enhanced_gemm_keys)) + tracelens_util.TraceEventUtils.JaxOpKeys.GemmKeys = all_keys + + print(" [OK] Patched TraceEventUtils.JaxOpKeys (GEMM keys enhanced)") + except (ImportError, AttributeError) as e: + print(f" [WARN] Could not patch TraceEventUtils: {e}") + + # Patch torch_op_mapping for better categorization + try: + from TraceLens.PerfModel import torch_op_mapping + + original_categorize = torch_op_mapping.categorize_torch_op + + def patched_categorize_torch_op(row): + """Enhanced categorization with better GEMM detection.""" + result = original_categorize(row) + + # If result is 'other', check for GEMM patterns in kernel names + if result == "other" and "kernel_details" in row and len(row["kernel_details"]) > 0: + kernel_name = row["kernel_details"][0]["name"] + pattern = r"^.*C[a-z]{3}_A[a-z]{3}_B[a-z]{3}.*$" + if re.match(pattern, kernel_name): + return "GEMM" + + return result + + torch_op_mapping.categorize_torch_op = patched_categorize_torch_op + print(" [OK] Patched torch_op_mapping (categorize_torch_op)") + except ImportError as e: + print(f" [WARN] Could not patch torch_op_mapping: {e}") + + print("[OK] All GEMM patches applied successfully!\n") + + +def main(): + if len(sys.argv) < 2: + print("Usage: tracelens_with_gemm_patch.py [args...]") + print("") + print("Commands:") + print(" generate_perf_report - Generate individual performance report") + print(" generate_multi_rank_collective - Generate multi-rank collective report") + print(" compare_perf_reports - Compare performance reports") + sys.exit(1) + + # Apply patches FIRST, before any TraceLens imports + apply_gemm_patches() + + # Import TraceLens modules - they will use the patched functions + from TraceLens.Reporting.generate_perf_report_pytorch import main as generate_perf_report_main + from TraceLens.Reporting.generate_multi_rank_collective_report_pytorch import main as generate_multi_rank_collective_report_main + from TraceLens.Reporting.compare_perf_reports_pytorch import main as compare_perf_reports_main + + # Verify patches are still in effect after imports + from TraceLens.PerfModel import torch_op_mapping + if hasattr(torch_op_mapping.categorize_torch_op, '__name__'): + print(f"[VERIFY] categorize_torch_op function name: {torch_op_mapping.categorize_torch_op.__name__}") + if hasattr(torch_op_mapping.categorize_torch_op, '_debug_counts'): + print("[VERIFY] Debug counters found - patch confirmed active!") + else: + print("[VERIFY] No debug counters - patch may not be active!") + + command = sys.argv[1] + + # Remove the command from argv so TraceLens sees only its args + sys.argv = [sys.argv[0]] + sys.argv[2:] + + if command == "generate_perf_report": + generate_perf_report_main() + elif command == "generate_multi_rank_collective": + generate_multi_rank_collective_report_main() + elif command == "compare_perf_reports": + compare_perf_reports_main() + else: + print(f"Error: Unknown command '{command}'") + print("") + print("Available commands:") + print(" generate_perf_report") + print(" generate_multi_rank_collective") + print(" compare_perf_reports") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/conftest.py b/tests/conftest.py index 775e953..1e88def 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,30 @@ import pytest +def pytest_configure(config): + """Configure pytest with custom markers for all tests.""" + config.addinivalue_line( + "markers", "integration: mark test as integration test (deselect with '-m \"not integration\"')" + ) + config.addinivalue_line( + "markers", "slow: mark test as slow running" + ) + config.addinivalue_line( + "markers", "gpu: mark test as requiring GPU" + ) + + +def pytest_addoption(parser): + """Add custom command-line options.""" + # GEMM regression test options + parser.addoption( + "--generate-baseline", + action="store_true", + default=False, + help="Generate baseline expected outputs for GEMM regression tests" + ) + + @pytest.fixture def sample_trace_event(): """Create a sample trace event for testing.""" diff --git a/tests/gemm_analysis/README.md b/tests/gemm_analysis/README.md new file mode 100644 index 0000000..063ef22 --- /dev/null +++ b/tests/gemm_analysis/README.md @@ -0,0 +1,272 @@ +# GEMM Analysis Regression Test Suite + +Automated regression testing for `scripts/gemm_analysis` pipeline. + +## Directory Structure + +``` +tests/gemm_analysis/ +├── testdata/ # Synthetic test data (not in git) +│ └── test_sweep/ # Fixed test data directory +├── expected_outputs/ # Baseline outputs (not in git) +├── actual_outputs/ # Current branch outputs (temp) +├── test_gemm_regression.py # Pytest test suite +├── generate_synthetic_data.py # Synthetic data generator +├── compare_outputs.py # Output comparison logic +└── README.md # This file +``` + +## Test Architecture + +### PipelineRunner Class +Encapsulates pipeline execution with methods: +- `run_tracelens()` - Run TraceLens analysis +- `run_analyze_gemm()` - Analyze GEMM reports +- `run_plot_variance()` - Generate variance plots +- `run_enhancement_scripts()` - Run additional analysis scripts +- `run_full_pipeline()` - Execute complete pipeline + +### Test Classes + +1. **TestFullPipeline** - Full end-to-end regression test +2. **TestIndividualSteps** - Test individual pipeline components +3. **TestCustomConfigurations** - Test with different thread/channel configurations +4. **TestErrorHandling** - Test error conditions and edge cases + +### Test Coverage + +Scripts tested: +- run_tracelens_analysis.sh +- analyze_gemm_reports.py +- plot_gemm_variance.py +- enhance_gemm_variance_with_timestamps.py +- gemm_report_with_collective_overlap.py +- process_gpu_timeline.py + +Validation: +- TraceLens analysis execution (skips if not installed) +- GEMM report analysis across configurations +- Variance plot generation +- Enhancement script outputs +- GPU timeline processing +- Numeric output consistency (tolerance: 1e-6) +- Script integration and data flow + +### Test Data + +- Synthetic PyTorch profiler traces (not real traces) +- Configurations: 2 threads (256, 512) × 2 channels (28, 56) × 8 ranks +- 2 batches (ProfilerSteps) per trace +- Based on MI350X trace structure +- Generated via `generate_synthetic_data.py` + +### Baseline Comparison + +- Baseline in `expected_outputs/` (generated from origin/main) +- Current branch outputs in `actual_outputs/` +- Numeric values: tolerance 1e-6 +- Ignores: whitespace, column ordering, timestamps + +## Setup + +```bash +source ~/venvs/aorta/bin/activate + +# Generate test data (once) +python tests/gemm_analysis/generate_synthetic_data.py --output-dir tests/gemm_analysis/testdata + +# Generate baseline (on origin/main) +pytest tests/gemm_analysis/test_gemm_regression.py --generate-baseline + +# If TraceLens not installed, run manually: +bash scripts/gemm_analysis/run_tracelens_analysis.sh tests/gemm_analysis/testdata/test_sweep +``` + +## Running Tests + +```bash +# All tests +pytest tests/gemm_analysis/ -v + +# Specific test class +pytest tests/gemm_analysis/test_gemm_regression.py::TestFullPipeline -v +pytest tests/gemm_analysis/test_gemm_regression.py::TestIndividualSteps -v +pytest tests/gemm_analysis/test_gemm_regression.py::TestCustomConfigurations -v +pytest tests/gemm_analysis/test_gemm_regression.py::TestErrorHandling -v + +# Integration tests only +pytest tests/gemm_analysis/ -m integration -v + +# Skip integration tests +pytest tests/gemm_analysis/ -m "not integration" -v +``` + +## Adding New Tests + +### Add Test Method to Existing Class + +```python +class TestIndividualSteps: + def test_new_feature(self, runner, test_data_dir, temp_output_dir): + """Test description.""" + # Setup + input_file = test_data_dir / "input.csv" + output_file = temp_output_dir / "output.csv" + + # Execute via runner + result = runner.run_new_feature(input_file, output_file) + + # Assert + assert result.exists() + assert result.stat().st_size > 0 +``` + +### Add New Test Class + +```python +class TestNewScript: + @pytest.fixture(scope="class") + def runner(self, scripts_dir): + return PipelineRunner(scripts_dir) + + def test_new_script_execution(self, runner, test_data_dir): + """Test new script execution.""" + # Add runner method first in PipelineRunner class + result = runner.run_new_script(test_data_dir) + assert result.returncode == 0 +``` + +### Add Runner Method + +In `PipelineRunner` class: + +```python +def run_new_script(self, input_path: Path, output_path: Path) -> Path: + """ + Run new_script.py step. + + Args: + input_path: Input file/directory + output_path: Output file/directory + + Returns: + Path to output + """ + print("\nN. Running new script...") + + cmd = [ + sys.executable, + str(self.scripts_dir / "new_script.py"), + "--input", str(input_path), + "--output", str(output_path) + ] + + result = subprocess.run( + cmd, capture_output=True, text=True, timeout=self.timeouts.get('new_script', 300) + ) + + if result.returncode != 0: + print(f"Error: {result.stderr}") + pytest.fail("new_script.py failed") + + print(f" Created: {output_path}") + return output_path +``` + +### Add to Full Pipeline + +In `run_full_pipeline()` method: + +```python +def run_full_pipeline(self, test_data_dir: Path, output_dir: Path) -> Dict[str, Path]: + # ... existing steps ... + + # Step N: New script + new_output = self.run_new_script(input_file, output_file) + + outputs['new_output'] = new_output + return outputs +``` + +### Update Baseline Comparison + +If script produces output that should be compared: + +1. Regenerate baseline: +```bash +pytest tests/gemm_analysis/test_gemm_regression.py --generate-baseline +``` + +2. Update `compare_outputs.py` if new comparison logic needed + +## Output Comparison Rules + +### Strict Requirements +- Core numeric metrics match within tolerance (1e-6) +- Required columns/fields present +- Per-GEMM aggregates consistent + +### Allowed Variations +- Extra columns (additive changes) +- Column reordering +- HTML styling changes +- Timestamps and build IDs +- File paths (structure maintained) +- Additional rows (core data present) + +## Modifying Test Data + +1. Edit `generate_synthetic_data.py` +2. Regenerate: +```bash +python tests/gemm_analysis/generate_synthetic_data.py --output-dir tests/gemm_analysis/testdata +bash scripts/gemm_analysis/run_tracelens_analysis.sh tests/gemm_analysis/testdata/test_sweep +pytest tests/gemm_analysis/test_gemm_regression.py --generate-baseline +``` + +## Troubleshooting + +**testdata not found** +```bash +python tests/gemm_analysis/generate_synthetic_data.py --output-dir tests/gemm_analysis/testdata +``` + +**Baseline outdated** +```bash +pytest tests/gemm_analysis/test_gemm_regression.py --generate-baseline +``` + +**TraceLens outputs missing** +```bash +bash scripts/gemm_analysis/run_tracelens_analysis.sh tests/gemm_analysis/testdata/test_sweep +``` + +**Dependencies missing** +```bash +pip install pytest pandas numpy openpyxl matplotlib +``` + +**Virtual environment not activated** +```bash +source ~/venvs/aorta/bin/activate +``` + +**Tests fail after generating synthetic data** +```bash +# Run TraceLens analysis on synthetic data +bash scripts/gemm_analysis/run_tracelens_analysis.sh tests/gemm_analysis/testdata/test_sweep +``` + +## CI/CD Integration + +```bash +source ~/venvs/aorta/bin/activate +python tests/gemm_analysis/generate_synthetic_data.py --output-dir tests/gemm_analysis/testdata +bash scripts/gemm_analysis/run_tracelens_analysis.sh tests/gemm_analysis/testdata/test_sweep +pytest tests/gemm_analysis/test_gemm_regression.py --generate-baseline # If baseline needed +pytest tests/gemm_analysis/ -v +``` + +Exit codes: +- 0: All tests passed +- 1: Test failures diff --git a/tests/gemm_analysis/compare_outputs.py b/tests/gemm_analysis/compare_outputs.py new file mode 100644 index 0000000..4f1cae8 --- /dev/null +++ b/tests/gemm_analysis/compare_outputs.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python3 +""" +Compare actual outputs against expected baseline outputs. +Allows tolerance for numeric values and ignores cosmetic differences. +""" + +import argparse +import json +import sys +from pathlib import Path +from typing import Dict, List, Tuple, Any +import pandas as pd +import numpy as np +from datetime import datetime + + +class OutputComparator: + """Compare analysis outputs with tolerance for acceptable differences.""" + + def __init__(self, expected_dir: Path, actual_dir: Path, verbose: bool = False): + self.expected_dir = Path(expected_dir) + self.actual_dir = Path(actual_dir) + self.verbose = verbose + self.numeric_tolerance = 1e-6 # Tolerance for numeric comparisons + self.errors = [] + self.warnings = [] + + def log_error(self, msg: str): + """Log an error message.""" + self.errors.append(msg) + if self.verbose: + print(f"ERROR: {msg}") + + def log_warning(self, msg: str): + """Log a warning message.""" + self.warnings.append(msg) + if self.verbose: + print(f"WARNING: {msg}") + + def log_success(self, msg: str): + """Log a success message.""" + if self.verbose: + print(f"OK: {msg}") + + def compare_csv_files(self, expected_file: Path, actual_file: Path) -> bool: + """Compare two CSV files with tolerance for numeric values.""" + try: + expected_df = pd.read_csv(expected_file) + actual_df = pd.read_csv(actual_file) + except Exception as e: + self.log_error(f"Failed to read CSV files: {e}") + return False + + # Check if required columns are present (actual may have more) + expected_cols = set(expected_df.columns) + actual_cols = set(actual_df.columns) + + missing_cols = expected_cols - actual_cols + if missing_cols: + self.log_error(f"Missing required columns in {actual_file.name}: {missing_cols}") + return False + + extra_cols = actual_cols - expected_cols + if extra_cols: + self.log_warning(f"Extra columns in {actual_file.name} (allowed): {extra_cols}") + + # Check row count (should be similar) + if len(expected_df) != len(actual_df): + self.log_warning(f"Different row counts in {actual_file.name}: " + f"expected {len(expected_df)}, got {len(actual_df)}") + + # Compare numeric columns with tolerance + for col in expected_cols: + if col not in actual_df.columns: + continue + + if pd.api.types.is_numeric_dtype(expected_df[col]): + # Compare numeric columns with tolerance + expected_vals = expected_df[col].fillna(0).values + actual_vals = actual_df[col].fillna(0).values[:len(expected_vals)] + + if not np.allclose(expected_vals, actual_vals, + rtol=self.numeric_tolerance, + atol=self.numeric_tolerance): + max_diff = np.max(np.abs(expected_vals - actual_vals)) + self.log_error(f"Numeric mismatch in column '{col}' of {actual_file.name}: " + f"max difference = {max_diff}") + return False + else: + # For non-numeric columns, check if core values are present + # (order may differ, some values may be added) + expected_values = set(expected_df[col].dropna().unique()) + actual_values = set(actual_df[col].dropna().unique()) + + missing_values = expected_values - actual_values + if missing_values and col in ['kernel_name', 'operation']: + # Critical columns should have all values + self.log_error(f"Missing values in column '{col}' of {actual_file.name}: " + f"{missing_values}") + return False + + self.log_success(f"CSV comparison passed: {actual_file.name}") + return True + + def compare_json_files(self, expected_file: Path, actual_file: Path) -> bool: + """Compare two JSON files allowing for structural flexibility.""" + try: + with open(expected_file, 'r') as f: + expected_data = json.load(f) + with open(actual_file, 'r') as f: + actual_data = json.load(f) + except Exception as e: + self.log_error(f"Failed to read JSON files: {e}") + return False + + return self.compare_json_objects(expected_data, actual_data, actual_file.name) + + def compare_json_objects(self, expected: Any, actual: Any, context: str) -> bool: + """Recursively compare JSON objects.""" + if type(expected) != type(actual): + self.log_error(f"Type mismatch in {context}: " + f"expected {type(expected)}, got {type(actual)}") + return False + + if isinstance(expected, dict): + # Check required keys are present + for key in expected: + if key not in actual: + self.log_error(f"Missing required key '{key}' in {context}") + return False + if not self.compare_json_objects(expected[key], actual[key], + f"{context}.{key}"): + return False + + # Extra keys are allowed + extra_keys = set(actual.keys()) - set(expected.keys()) + if extra_keys: + self.log_warning(f"Extra keys in {context} (allowed): {extra_keys}") + + elif isinstance(expected, list): + # Lists should have similar length + if len(expected) != len(actual): + self.log_warning(f"Different list lengths in {context}: " + f"expected {len(expected)}, got {len(actual)}") + + # Compare elements up to minimum length + min_len = min(len(expected), len(actual)) + for i in range(min_len): + if not self.compare_json_objects(expected[i], actual[i], + f"{context}[{i}]"): + return False + + elif isinstance(expected, (int, float)): + # Numeric comparison with tolerance + if not np.isclose(expected, actual, rtol=self.numeric_tolerance): + self.log_error(f"Numeric mismatch in {context}: " + f"expected {expected}, got {actual}") + return False + + elif isinstance(expected, str): + # String comparison (ignore timestamps and build IDs) + if self.is_volatile_string(expected) or self.is_volatile_string(actual): + # Skip comparison for volatile strings + pass + elif expected != actual: + self.log_error(f"String mismatch in {context}: " + f"expected '{expected}', got '{actual}'") + return False + + else: + # Direct comparison for other types + if expected != actual: + self.log_error(f"Value mismatch in {context}: " + f"expected {expected}, got {actual}") + return False + + return True + + def is_volatile_string(self, s: str) -> bool: + """Check if a string contains volatile content like timestamps or IDs.""" + volatile_patterns = [ + '2024', '2025', # Years in timestamps + 'build_', 'run_', # Build/run IDs + 'tmp/', # Temporary paths + '.log', # Log files + ] + return any(pattern in s for pattern in volatile_patterns) + + def compare_excel_files(self, expected_file: Path, actual_file: Path) -> bool: + """Compare Excel files by checking sheets and data.""" + try: + expected_xlsx = pd.ExcelFile(expected_file) + actual_xlsx = pd.ExcelFile(actual_file) + except Exception as e: + self.log_error(f"Failed to read Excel files: {e}") + return False + + # Check if required sheets are present + expected_sheets = set(expected_xlsx.sheet_names) + actual_sheets = set(actual_xlsx.sheet_names) + + missing_sheets = expected_sheets - actual_sheets + if missing_sheets: + self.log_error(f"Missing required sheets in {actual_file.name}: {missing_sheets}") + return False + + # Compare data in each expected sheet + for sheet_name in expected_sheets: + expected_df = pd.read_excel(expected_file, sheet_name=sheet_name) + actual_df = pd.read_excel(actual_file, sheet_name=sheet_name) + + # Similar logic to CSV comparison + if len(expected_df.columns) > len(actual_df.columns): + self.log_error(f"Missing columns in sheet '{sheet_name}' of {actual_file.name}") + return False + + self.log_success(f"Excel comparison passed: {actual_file.name}") + return True + + def compare_html_files(self, expected_file: Path, actual_file: Path) -> bool: + """Compare HTML files (check if key elements are present).""" + try: + with open(expected_file, 'r') as f: + expected_content = f.read() + with open(actual_file, 'r') as f: + actual_content = f.read() + except Exception as e: + self.log_error(f"Failed to read HTML files: {e}") + return False + + # Check for key markers/elements (not exact match due to styling changes) + key_markers = [ + ' bool: + """Compare all files in the output directories.""" + all_passed = True + + # Get all expected files + expected_files = list(self.expected_dir.rglob('*')) + expected_files = [f for f in expected_files if f.is_file()] + + print(f"\nComparing {len(expected_files)} files...") + print("=" * 50) + + for expected_file in expected_files: + # Get relative path + rel_path = expected_file.relative_to(self.expected_dir) + actual_file = self.actual_dir / rel_path + + if not actual_file.exists(): + self.log_error(f"Missing output file: {rel_path}") + all_passed = False + continue + + # Compare based on file type + if expected_file.suffix == '.csv': + if not self.compare_csv_files(expected_file, actual_file): + all_passed = False + elif expected_file.suffix == '.json': + if not self.compare_json_files(expected_file, actual_file): + all_passed = False + elif expected_file.suffix in ['.xlsx', '.xls']: + if not self.compare_excel_files(expected_file, actual_file): + all_passed = False + elif expected_file.suffix in ['.html', '.htm']: + if not self.compare_html_files(expected_file, actual_file): + all_passed = False + else: + # For other files, just check existence + self.log_success(f"File exists: {rel_path}") + + return all_passed + + def print_summary(self): + """Print summary of comparison results.""" + print("\n" + "=" * 50) + print("COMPARISON SUMMARY") + print("=" * 50) + + if self.errors: + print(f"\nERRORS ({len(self.errors)}):") + for error in self.errors[:10]: # Show first 10 errors + print(f" - {error}") + if len(self.errors) > 10: + print(f" ... and {len(self.errors) - 10} more errors") + + if self.warnings: + print(f"\nWARNINGS ({len(self.warnings)}):") + for warning in self.warnings[:5]: # Show first 5 warnings + print(f" - {warning}") + if len(self.warnings) > 5: + print(f" ... and {len(self.warnings) - 5} more warnings") + + if not self.errors: + print("\nAll regression tests PASSED!") + print(f" {len(self.warnings)} warnings (cosmetic differences allowed)") + else: + print(f"\nRegression tests FAILED!") + print(f" {len(self.errors)} errors, {len(self.warnings)} warnings") + + +def main(): + parser = argparse.ArgumentParser(description="Compare GEMM analysis outputs") + parser.add_argument("--expected", required=True, help="Path to expected outputs directory") + parser.add_argument("--actual", required=True, help="Path to actual outputs directory") + parser.add_argument("--verbose", action="store_true", help="Enable verbose output") + args = parser.parse_args() + + comparator = OutputComparator(args.expected, args.actual, args.verbose) + + # Run comparison + success = comparator.compare_directories() + + # Print summary + comparator.print_summary() + + # Exit with appropriate code + sys.exit(0 if success else 1) + + +if __name__ == "__main__": + main() diff --git a/tests/gemm_analysis/generate_synthetic_data.py b/tests/gemm_analysis/generate_synthetic_data.py new file mode 100644 index 0000000..73a6216 --- /dev/null +++ b/tests/gemm_analysis/generate_synthetic_data.py @@ -0,0 +1,641 @@ +#!/usr/bin/env python3 +""" +Generate synthetic test data for GEMM analysis regression tests. +Creates realistic PyTorch profiler traces based on actual MI350X traces +for two channel configurations, two threads, and 7 ranks. + +Key requirements for TraceLens GEMM sheet generation: +1. Use TraceLens-recognized GEMM operation names (aten::mm, aten::addmm, aten::bmm) +2. Provide proper 2D matrix Input Dims for GEMM ops: + - aten::mm: [[M, K], [K, N]] for two input matrices + - aten::addmm: [[N], [M, K], [K, N], ...] for bias + two matrices + - aten::bmm: [[B, M, K], [B, K, N]] for batched matrix multiplication +3. Include kernel_details on CPU ops with actual GEMM kernel names matching pattern C[xyz]_A[xyz]_B[xyz] +""" + +import json +import os +import sys +from pathlib import Path +from typing import Dict, List +import random +import argparse +from datetime import datetime + +# Set random seed for reproducibility +random.seed(42) + + +class SyntheticDataGenerator: + """Generate synthetic test data for regression testing.""" + + def __init__( + self, + output_dir: Path, + thread_configs: List[int] | None = None, + channel_configs: List[int] | None = None, + num_ranks: int = 8, + num_batches: int = 2, + ): + self.output_dir = Path(output_dir) + self.sweep_dir = self.output_dir / "test_sweep" # Fixed name without timestamp + + # Fixed configurations for testing (can be overridden) + self.thread_configs = thread_configs or [256, 512] + # include 28 by default so synthetic sweeps cover test expectations + self.channel_configs = channel_configs or [28, 56] + self.num_ranks = num_ranks + self.num_batches = num_batches # Generate batches (ProfilerSteps) per trace + + # Realistic GEMM kernel names from actual AMD MI350X traces + self.gemm_kernels = [ + # Simple names crafted to match TraceLens GEMM regex (Cxxx_Axxx_Bxxx) + "Cabc_Adef_Bghi_gemm0", + "Cxyz_Auvw_Brst_gemm1", + "Cfoo_Abar_Bbaz_gemm2", + # Most common GEMM kernels from actual traces + "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x240x64_MI16x16x1_SN_LDSB0_AFC0_AFEM1_AFEM1_ASEM1_CLR0_CADS0_DTLA1_DTLB1_DTVA0_DTVB0_EPS0_FDSI0_GRPM1_GRVWA8_GRVWB2_GSU0_GSUAMB_GLS0_ISA950_IU1_K1_LDSTI0_LBSPPA1024_LBSPPB256_LBSPPM0_LPA16_LPB16_LPM0_LRVW8_LWPMn1_MIAV0_MIWT4_15_MO40_NTn1_NTA0_NTB0_NTC6_NTD4_NTM0_NEPBS0_NLCA1_NLCB1_ONLL1_PGR2_PLR1_PKA1_SIA3_SS1_SPO1_SRVW0_SSO1_SVW4_SK3_SKFTR0_SKXCCM0_TLDS1_ULSGRO0_USL1_UIOFGRO0_USFGRO0_VSn1_VWA4_VWB1_WSGRA0_WSGRB0_WS64_WG64_4_1", + "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0_AFC0_AFEM1_AFEM1_ASEM1_CLR1_CADS0_DTLA1_DTLB1_DTVA0_DTVB0_EPS0_FDSI0_GRPM1_GRVWA8_GRVWB8_GSU0_GSUAMB_GLS0_ISA950_IU1_K1_LDSTI0_LBSPPA1024_LBSPPB1024_LBSPPM0_LPA16_LPB16_LPM0_LRVW8_LWPMn1_MIAV0_MIWT8_8_MO40_NTn1_NTA0_NTB0_NTC0_NTD4_NTM0_NEPBS0_NLCA1_NLCB1_ONLL1_PGR2_PLR1_PKA1_SIA3_SS1_SPO0_SRVW0_SSO0_SVW8_SK3_SKFTR0_SKXCCM0_TLDS1_ULSGRO0_USL1_UIOFGRO0_USFGRO0_VSn1_VWA8_VWB8_WSGRA0_WSGRB0_WS64_WG32_8_1", + # Additional GEMM variants + "Cijk_Alik_Bljk_SB_MT128x128x32_MI32x32x1x2", + "Cijk_Alik_Bljk_SB_MT256x256x64_MI64x64x2x2", + "Cijk_Alik_Bljk_SB_MT64x64x16_MI16x16x1x1", + ] + + # Other kernel types seen in real traces + self.other_kernels = [ + "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor >, std::array >(int, at::native::AUnaryFunctor >, std::array)", + "void at::native::unrolled_elementwise_kernel >, std::array, 4, TrivialOffsetCalculator<1, unsigned int>, TrivialOffsetCalculator<1, unsigned int>, at::native::memory::LoadWithoutCast, at::native::memory::StoreWithoutCast>(int, at::native::AUnaryFunctor >, std::array, TrivialOffsetCalculator<1, unsigned int>, TrivialOffsetCalculator<1, unsigned int>, at::native::memory::LoadWithoutCast, at::native::memory::StoreWithoutCast)", + "ncclDevKernel_Generic_1(ncclDevKernelArgsStorage<4096ul>)", + ] + + # NCCL operations from real traces + self.nccl_ops = [ + "nccl:all_reduce", + "nccl:broadcast", + "nccl:all_gather", + "nccl:reduce_scatter" + ] + + def create_pytorch_trace(self, rank: int, thread: int, channel: int) -> Dict: + """Create a realistic PyTorch profiler trace matching actual MI350X format. + + Generates trace events for multiple batches (ProfilerSteps). + Each ProfilerStep represents one training batch/iteration. + """ + + # Device properties for MI350X GPUs (exact format from real traces) + device_properties = [] + for gpu_id in range(8): + device_properties.append({ + "id": gpu_id, + "name": "AMD Instinct MI350X", + "totalGlobalMem": 309220868096, + "computeMajor": 9, + "computeMinor": 5, + "maxThreadsPerBlock": 1024, + "maxThreadsPerMultiprocessor": 2048, + "regsPerBlock": 65536, + "warpSize": 64, + "sharedMemPerBlock": 163840, + "numSms": 256, + "maxSharedMemoryPerMultiProcessor": 163840 + }) + + # Metadata matching real traces + metadata = { + "schemaVersion": 1, + "deviceProperties": device_properties, + "roctracer_version": 4.1, + "with_flops": 1, + "distributedInfo": { + "backend": "nccl", + "rank": rank, + "world_size": 8, + "pg_count": 1, + "pg_config": [{ + "pg_name": "0", + "pg_desc": "default_pg", + "backend_config": "cuda:nccl", + "pg_size": 8, + "ranks": list(range(8)) + }], + "nccl_version": "2.26.6" + }, + "record_shapes": 1, + "hip_runtime_version": 70051831, + "profile_memory": 1, + "hip_driver_version": 70051831, + "with_stack": 1, + "trace_id": f"F063B8B18C2C424DB348840DD69E22BB", + "displayTimeUnit": "ms", + "baseTimeNanoseconds": 1759300074000000000 + } + + # Generate trace events for multiple batches + events = [] + base_ts = 5917770837674.380 # Base timestamp in microseconds + pid = 258 # CPU Process ID + gpu_pid = 2 # GPU Process ID + tid_main = thread # Main thread ID + tid_nccl = 683 # NCCL thread ID (from real traces) + external_id_counter = 1 + correlation_counter = 3900 + + # Generate events for each batch (2 batches as requested) + for batch in range(self.num_batches): + batch_start_ts = base_ts + batch * 600000000 # ~600ms per batch (realistic timing) + current_ts = batch_start_ts + + # ProfilerStep event (marks each batch/iteration) + step_num = batch + 2 # Start from ProfilerStep#2 like real traces + step_duration = 536019.180 + random.uniform(-10000, 10000) # Slight variation per batch + + events.append({ + "ph": "X", + "cat": "user_annotation", + "name": f"ProfilerStep#{step_num}", + "pid": pid, + "tid": tid_main, + "ts": current_ts, + "dur": step_duration, + "args": { + "External id": external_id_counter, + "Record function id": 0, + "Ev Idx": batch * 1000 + } + }) + external_id_counter += 1 + + # DataLoader event (loading next batch of data) + events.append({ + "ph": "X", + "cat": "user_annotation", + "name": "enumerate(DataLoader)#_MultiProcessingDataLoaderIter.__next__", + "pid": pid, + "tid": tid_main, + "ts": current_ts + 50, + "dur": 141.768, + "args": { + "External id": external_id_counter, + "Record function id": 0, + "Ev Idx": batch * 1000 + 1 + } + }) + external_id_counter += 1 + + # Add CPU operations that TraceLens expects (generic ops, not GEMM) + # These are placeholder ops without kernels + cpu_op_names = ["aten::to", "aten::_to_copy", "aten::mul", "aten::add", "aten::copy_"] + for cpu_idx in range(5): + events.append({ + "ph": "X", + "cat": "cpu_op", + "name": cpu_op_names[cpu_idx % len(cpu_op_names)], + "pid": pid, + "tid": tid_main, + "ts": current_ts + 200 + cpu_idx * 100, + "dur": 100 + random.uniform(-20, 20), + "args": { + "External id": external_id_counter, + "Sequence number": 2351 + cpu_idx, + "Fwd thread id": 0, + "Record function id": 0, + "Concrete Inputs": ["", "6", "0", "", "", "True", "False", ""], + "Input type": ["float", "Scalar", "Scalar", "", "", "Scalar", "Scalar", ""], + "Input Strides": [[1441792, 8192, 256, 1], [], [], [], [], [], [], []], + "Input Dims": [[512, 176, 32, 256], [], [], [], [], [], [], []], + "Ev Idx": batch * 1000 + 2 + cpu_idx + } + }) + external_id_counter += 1 + + current_ts += 1000 + + # Generate NCCL broadcast at start of batch (parameter sync) + events.append({ + "ph": "X", + "cat": "user_annotation", + "name": "nccl:broadcast", + "pid": pid, + "tid": tid_main, + "ts": current_ts, + "dur": 5000 + random.uniform(-500, 500), + "args": { + "External id": external_id_counter, + "Record function id": 0, + "Ev Idx": batch * 1000 + 10 + } + }) + external_id_counter += 1 + current_ts += 6000 + + # Generate mix of kernels for this batch + num_kernels_per_batch = 30 + random.randint(-5, 5) + + for kernel_idx in range(num_kernels_per_batch): + # Mix GEMM kernels with other kernels (70% GEMM, 30% other) + is_gemm_kernel = random.random() < 0.7 + + if is_gemm_kernel: + # GEMM kernel + kernel_name_idx = (rank + kernel_idx + batch) % len(self.gemm_kernels) + kernel_name = self.gemm_kernels[kernel_name_idx] + + # Realistic timing based on kernel type and configuration + if "MT256x240" in kernel_name: + base_dur = 45000 + random.uniform(-5000, 5000) + grid_size = [385024, 1, 1] + elif "MT256x256" in kernel_name: + base_dur = 50000 + random.uniform(-5000, 5000) + grid_size = [65536, 1, 1] + else: + base_dur = 30000 + random.uniform(-3000, 3000) + grid_size = [32768, 1, 1] + + # Configuration-based variance + dur_variance = base_dur * (1 + (thread-256)/2000 + (channel-28)/200 + rank * 0.005) + else: + # Other kernel (elementwise, etc.) + kernel_name = random.choice(self.other_kernels) + dur_variance = 5000 + random.uniform(-1000, 1000) + grid_size = [16384, 1, 1] + + # Create a CPU launcher op that envelops the runtime + kernel so TraceLens can assign parents + cpu_launcher_ts = current_ts - 40 + cpu_launcher_dur = dur_variance + 80 + # Use GEMM operation names for GEMM kernels (TraceLens recognizes these) + # Use other names for non-GEMM kernels + if is_gemm_kernel: + # Use TraceLens-recognized GEMM operation names + gemm_op_names = ["aten::mm", "aten::addmm", "aten::bmm"] + launcher_name = gemm_op_names[kernel_idx % len(gemm_op_names)] + else: # Non-GEMM kernel + launcher_name = cpu_op_names[kernel_idx % len(cpu_op_names)] + + # NOTE: TraceLens needs kernel_details to categorize ops properly + # even though real PyTorch traces may not have them initially + kernel_detail_entry = { + "name": kernel_name, + "dur": dur_variance, + "stream": 0, + "grid": grid_size, + "block": [256, 1, 1], + } + + # Generate proper Input Dims based on operation type + # GEMM operations need 2D matrix dimensions for M, N, K extraction + if is_gemm_kernel: + # Realistic GEMM matrix dimensions (vary by operation type) + if launcher_name == "aten::mm": + # mm(A, B) where A is MxK, B is KxN + M, K, N = 8192 + rank * 1024, 4096, 2048 + rank * 512 + input_dims = [[M, K], [K, N]] + input_strides = [[K, 1], [N, 1]] + input_types = ["float", "float"] + elif launcher_name == "aten::addmm": + # addmm(bias, A, B) where bias is N, A is MxK, B is KxN + M, K, N = 8192 + rank * 1024, 4096, 2048 + rank * 512 + input_dims = [[N], [M, K], [K, N], [], []] + input_strides = [[1], [K, 1], [N, 1], [], []] + input_types = ["float", "float", "float", "", ""] + else: # aten::bmm + # bmm(A, B) where A is BxMxK, B is BxKxN + B, M, K, N = 512, 176, 256, 128 + input_dims = [[B, M, K], [B, K, N]] + input_strides = [[M*K, K, 1], [K*N, N, 1]] + input_types = ["float", "float"] + else: + # Non-GEMM operations can use simpler dims + input_dims = [[512, 176, 32, 256], [], [], [], [], [], [], []] + input_strides = [[1441792, 8192, 256, 1], [], [], [], [], [], [], []] + input_types = ["float", "Scalar", "Scalar", "", "", "Scalar", "Scalar", ""] + + events.append({ + "ph": "X", + "cat": "cpu_op", + "name": launcher_name, + "pid": pid, + "tid": tid_main, + "ts": cpu_launcher_ts, + "dur": cpu_launcher_dur, + "args": { + "External id": external_id_counter, + "Sequence number": 3000 + kernel_idx, + "Fwd thread id": 0, + "Record function id": 0, + "Concrete Inputs": ["", "6", "0", "", "", "True", "False", ""], + "Input type": input_types, + "Input Strides": input_strides, + "Input Dims": input_dims, + "Ev Idx": batch * 1000 + 600 + kernel_idx + }, + "kernel_details": [kernel_detail_entry], + "kernel_names": [kernel_name], + }) + external_id_counter += 1 + + # Create runtime launch and ac2g bridge to link kernels for TraceLens + corr_id = correlation_counter + runtime_ts = current_ts - 20 # slightly before kernel start + runtime_dur = 10 + + # cuda runtime launch event + events.append({ + "ph": "X", + "cat": "cuda_runtime", + "name": "hipLaunchKernel", + "pid": pid, + "tid": tid_main, + "ts": runtime_ts, + "dur": runtime_dur, + "args": { + "External id": external_id_counter, + "correlation": corr_id, + "kernel": kernel_name, + "grid": grid_size, + "block": [256, 1, 1], + } + }) + external_id_counter += 1 + + # ac2g start/end events to allow TraceLens to map runtime->GPU + events.append({ + "ph": "s", + "id": corr_id, + "pid": pid, + "tid": tid_main, + "ts": runtime_ts, + "cat": "ac2g", + "name": "ac2g" + }) + events.append({ + "ph": "f", + "id": corr_id, + "pid": gpu_pid, + "tid": 0, + "ts": current_ts, + "cat": "ac2g", + "name": "ac2g", + "bp": "e" + }) + + # GPU kernel event + events.append({ + "ph": "X", + "cat": "kernel", + "name": kernel_name, + "pid": gpu_pid, + "tid": 0, # GPU stream 0 + "ts": current_ts, + "dur": dur_variance, + "args": { + "External id": external_id_counter, + "device": rank % 8, # Device ID based on rank + "stream": 0, # Stream ID + "correlation": corr_id, + "kind": "Dispatch Kernel", # Required by TraceLens + "grid": grid_size, + "block": [256, 1, 1], + "kernel": kernel_name # Keep kernel name in args too + } + }) + + external_id_counter += 1 + correlation_counter += 1 + current_ts += dur_variance + random.uniform(100, 500) + + # Add periodic NCCL all_reduce operations (every 5 kernels for gradient sync) + if kernel_idx % 5 == 4: + events.append({ + "ph": "X", + "cat": "user_annotation", + "name": "nccl:all_reduce", + "pid": pid, + "tid": tid_nccl, # NCCL operations on different thread + "ts": current_ts, + "dur": 15000 + random.uniform(-2000, 2000), + "args": { + "External id": external_id_counter, + "Record function id": 0, + "Ev Idx": batch * 1000 + 100 + kernel_idx + } + }) + external_id_counter += 1 + + # Corresponding NCCL runtime launch and kernel + nccl_corr_id = correlation_counter + runtime_ts = current_ts + 50 + runtime_dur = 8 + + # CPU launcher for NCCL kernel (gives parent for short-kernel stats) + events.append({ + "ph": "X", + "cat": "cpu_op", + "name": "nccl::launch", + "pid": pid, + "tid": tid_nccl, + "ts": runtime_ts - 20, + "dur": 16000, + "args": { + "External id": external_id_counter, + "Sequence number": 5000 + kernel_idx, + "Fwd thread id": 0, + "Record function id": 0, + "Concrete Inputs": ["nccl_all_reduce"], + "Input type": ["string"], + "Input Strides": [[]], + "Input Dims": [[]], + "Ev Idx": batch * 1000 + 700 + kernel_idx + }, + "kernel_details": [{ + "name": "ncclDevKernel_Generic_1(ncclDevKernelArgsStorage<4096ul>)", + "dur": 14000, + "stream": 4, + "grid": [1, 1, 1], + "block": [1024, 1, 1], + }], + "kernel_names": ["ncclDevKernel_Generic_1(ncclDevKernelArgsStorage<4096ul>)"], + }) + external_id_counter += 1 + + events.append({ + "ph": "X", + "cat": "cuda_runtime", + "name": "hipLaunchKernel", + "pid": pid, + "tid": tid_nccl, + "ts": runtime_ts, + "dur": runtime_dur, + "args": { + "External id": external_id_counter, + "correlation": nccl_corr_id, + "kernel": "ncclDevKernel_Generic_1(ncclDevKernelArgsStorage<4096ul>)", + "grid": [1, 1, 1], + "block": [1024, 1, 1], + } + }) + external_id_counter += 1 + + events.append({ + "ph": "s", + "id": nccl_corr_id, + "pid": pid, + "tid": tid_nccl, + "ts": runtime_ts, + "cat": "ac2g", + "name": "ac2g" + }) + + events.append({ + "ph": "X", + "cat": "kernel", + "name": "ncclDevKernel_Generic_1(ncclDevKernelArgsStorage<4096ul>)", + "pid": gpu_pid, + "tid": 4, # Different GPU stream for NCCL + "ts": current_ts + 100, + "dur": 14000 + random.uniform(-2000, 2000), + "args": { + "External id": external_id_counter, + "device": rank % 8, # Device ID based on rank + "stream": 4, # NCCL stream + "correlation": nccl_corr_id, + "kind": "Dispatch Kernel", # Required by TraceLens + "grid": [1, 1, 1], + "block": [1024, 1, 1], + # Metadata required by TraceLens NcclAnalyser + "Collective name": "allreduce", + "In msg nelems": 1048576, + "Out msg nelems": 1048576, + "Group size": 8, + "Process Group Name": "0", + "Process Group Ranks": list(range(8)), + "dtype": "Float", + "Process Group Description": "default_pg", + "Process Group Ranks": list(range(8)), + } + }) + external_id_counter += 1 + events.append({ + "ph": "f", + "id": nccl_corr_id, + "pid": gpu_pid, + "tid": 4, + "ts": current_ts + 100, + "cat": "ac2g", + "name": "ac2g", + "bp": "e" + }) + + correlation_counter += 1 + current_ts += 16000 + + # Add GPU annotation events (match the ProfilerStep on GPU side) + for gpu_tid in [0, 5, 6]: # Multiple GPU threads like real traces + events.append({ + "ph": "X", + "cat": "gpu_user_annotation", + "name": f"ProfilerStep#{step_num}", + "pid": gpu_pid, + "tid": gpu_tid, + "ts": batch_start_ts, + "dur": step_duration, + "args": { + "External id": external_id_counter, + "Record function id": 0, + "Ev Idx": batch * 1000 + 500 + gpu_tid + } + }) + external_id_counter += 1 + + # Sort events by timestamp (important for trace validity) + events.sort(key=lambda x: x["ts"]) + + # Combine metadata and events + trace_data = metadata.copy() + trace_data["traceEvents"] = events + + return trace_data + + def generate_all_data(self) -> None: + """Generate all synthetic test data matching real trace structure.""" + print(f"Generating realistic synthetic test data in: {self.sweep_dir}") + print(f" Based on actual MI350X PyTorch profiler traces") + print(f" Generating {self.num_batches} batches (ProfilerSteps) per trace") + + for thread in self.thread_configs: + for channel in self.channel_configs: + config_name = f"{thread}thread/nccl_{channel}channels" + print(f"\n Generating traces for {config_name}...") + + # Create PyTorch traces for each rank + for rank in range(self.num_ranks): + trace_dir = self.sweep_dir / f"{config_name}/torch_profiler/rank{rank}" + trace_dir.mkdir(parents=True, exist_ok=True) + + # Generate realistic trace data + trace_data = self.create_pytorch_trace(rank, thread, channel) + + # Write to customer_trace_step10.json (matching real traces) + trace_file = trace_dir / "customer_trace_step10.json" + + with open(trace_file, 'w') as f: + json.dump(trace_data, f, indent=2) + + # Verify event counts + num_events = len(trace_data.get("traceEvents", [])) + print(f" Rank {rank}: Generated {num_events} events") + + # Create a metadata file for the sweep + metadata = { + "thread_configs": self.thread_configs, + "channel_configs": self.channel_configs, + "num_ranks": self.num_ranks, + "num_batches": self.num_batches, + "generated_at": datetime.now().isoformat(), + "based_on": "AMD MI350X actual traces", + "trace_format": "PyTorch Profiler JSON", + "profiler_steps": [f"ProfilerStep#{i+2}" for i in range(self.num_batches)] + } + + with open(self.sweep_dir / "metadata.json", 'w') as f: + json.dump(metadata, f, indent=2) + + print(f"\nSynthetic test data generated successfully!") + print(f" Thread configs: {self.thread_configs}") + print(f" Channel configs: {self.channel_configs}") + print(f" Number of ranks: {self.num_ranks}") + print(f" Batches per trace: {self.num_batches}") + print(f" ProfilerSteps: {metadata['profiler_steps']}") + + +def main(): + parser = argparse.ArgumentParser(description="Generate synthetic test data for GEMM analysis regression tests") + parser.add_argument("--output-dir", type=str, default="testdata", + help="Output directory for synthetic data (default: testdata)") + parser.add_argument("--threads", type=str, default="256,512", + help="Comma separated thread configs (default: 256,512)") + parser.add_argument("--channels", type=str, default="24,28,56", + help="Comma separated channel configs (default: 24,28,56)") + parser.add_argument("--ranks", type=int, default=8, + help="Number of ranks to generate (default: 8)") + parser.add_argument("--batches", type=int, default=2, + help="Profiler steps per trace (default: 2)") + args = parser.parse_args() + + output_dir = args.output_dir + thread_configs = [int(x) for x in args.threads.split(",") if x] + channel_configs = [int(x) for x in args.channels.split(",") if x] + + generator = SyntheticDataGenerator( + output_dir=output_dir, + thread_configs=thread_configs, + channel_configs=channel_configs, + num_ranks=args.ranks, + num_batches=args.batches, + ) + generator.generate_all_data() + + +if __name__ == "__main__": + main() diff --git a/tests/gemm_analysis/test_gemm_regression.py b/tests/gemm_analysis/test_gemm_regression.py new file mode 100644 index 0000000..3d283c1 --- /dev/null +++ b/tests/gemm_analysis/test_gemm_regression.py @@ -0,0 +1,586 @@ +#!/usr/bin/env python3 +""" +Modular regression tests for GEMM analysis pipeline. +Supports both full pipeline and individual component testing. + +Usage: + # Generate baseline (on origin/main) + pytest tests/gemm_analysis/test_gemm_regression.py --generate-baseline + + # Run all regression tests + pytest tests/gemm_analysis/test_gemm_regression.py + + # Run only full pipeline tests + pytest tests/gemm_analysis/test_gemm_regression.py::TestFullPipeline + + # Run only individual step tests + pytest tests/gemm_analysis/test_gemm_regression.py::TestIndividualSteps +""" + +import pytest +import subprocess +import sys +from pathlib import Path +import shutil +from typing import Dict, List, Optional, Tuple + +# Add parent directories to path +TEST_DIR = Path(__file__).parent # tests/gemm_analysis +REPO_ROOT = TEST_DIR.parent.parent # repository root +sys.path.insert(0, str(REPO_ROOT / 'scripts' / 'gemm_analysis')) + + +# ============================================================================ +# SHARED FIXTURES +# ============================================================================ + +@pytest.fixture(scope="session") +def test_data_dir(): + """Path to the test data directory.""" + return TEST_DIR / "testdata" / "test_sweep" + + +@pytest.fixture(scope="session") +def scripts_dir(): + """Path to the GEMM analysis scripts directory.""" + return REPO_ROOT / "scripts" / "gemm_analysis" + + +@pytest.fixture(scope="session") +def expected_outputs_dir(): + """Path to the expected outputs directory.""" + return TEST_DIR / "expected_outputs" + + +@pytest.fixture(scope="session") +def generate_baseline(request): + """Check if we're generating baseline outputs.""" + return request.config.getoption("--generate-baseline") + + +# ============================================================================ +# PIPELINE RUNNER CLASS +# ============================================================================ + +class PipelineRunner: + """Encapsulates pipeline execution logic for reuse across tests.""" + + def __init__(self, scripts_dir: Path, timeout_seconds: Dict[str, int] = None): + """ + Initialize the pipeline runner. + + Args: + scripts_dir: Path to GEMM analysis scripts + timeout_seconds: Optional timeout overrides for each step + """ + self.scripts_dir = scripts_dir + self.timeouts = timeout_seconds or { + 'tracelens': 600, + 'analyze': 300, + 'plot': 300, + 'enhance': 300, + 'overlap': 300, + 'timeline': 300 + } + + def setup_output_dirs(self, output_dir: Path) -> Tuple[Path, Path]: + """ + Create output directory structure. + + Returns: + Tuple of (output_dir, reports_dir) + """ + if output_dir.exists(): + shutil.rmtree(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + reports_dir = output_dir / "reports" + reports_dir.mkdir(parents=True, exist_ok=True) + return output_dir, reports_dir + + def run_tracelens(self, test_data_dir: Path, skip_on_missing: bool = True) -> Path: + """ + Run TraceLens analysis step. + + Args: + test_data_dir: Input data directory + skip_on_missing: Skip if TraceLens not installed (vs fail) + + Returns: + Path to tracelens output directory + """ + print("\n1. Running TraceLens analysis...") + tracelens_script = self.scripts_dir / "run_tracelens_analysis.sh" + + if tracelens_script.exists(): + result = subprocess.run( + ["bash", str(tracelens_script), str(test_data_dir)], + capture_output=True, text=True, timeout=self.timeouts['tracelens'] + ) + if result.returncode != 0: + if any(err in str(result.stderr) + str(result.stdout) + for err in ["TraceLens", "tracelens", "command not found"]): + if skip_on_missing: + print(" Warning: TraceLens not installed - skipping") + else: + pytest.fail("TraceLens not installed") + else: + pytest.fail(f"TraceLens analysis failed: {result.stderr}") + else: + print(" TraceLens analysis completed") + + # Check if outputs exist and contain actual Excel reports + tracelens_dir = test_data_dir / "tracelens_analysis" + if not tracelens_dir.exists(): + if skip_on_missing: + pytest.skip( + "TraceLens outputs not found. Either:\n" + " 1. Install TraceLens, or\n" + " 2. Run manually: bash scripts/gemm_analysis/run_tracelens_analysis.sh " + f"{test_data_dir}" + ) + else: + pytest.fail("TraceLens outputs not found") + + # Check if TraceLens generated actual Excel reports + excel_files = list(tracelens_dir.glob("**/individual_reports/*.xlsx")) + if not excel_files: + if skip_on_missing: + pytest.skip( + "TraceLens directory exists but contains no Excel reports.\n" + "TraceLens is either not installed or failed to run.\n" + "To fix, run: bash scripts/gemm_analysis/run_tracelens_analysis.sh " + f"{test_data_dir}" + ) + else: + pytest.fail("TraceLens outputs directory exists but contains no Excel reports") + + return tracelens_dir + + def run_analyze_gemm(self, + tracelens_dir: Path, + output_file: Path, + threads: List[int] = None, + channels: List[int] = None, + ranks: List[int] = None, + top_k: int = 10) -> Path: + """ + Run analyze_gemm_reports.py step. + + Args: + tracelens_dir: TraceLens output directory + output_file: Output CSV path + threads: Thread configurations (default: [256, 512]) + channels: Channel configurations (default: [28, 56]) + ranks: Rank list (default: [0-6]) + top_k: Number of top kernels to extract + + Returns: + Path to output file + """ + print("\n2. Analyzing GEMM reports...") + + # Default values + threads = threads or [256, 512] + channels = channels or [28, 56] + ranks = ranks or list(range(7)) + + # Require TraceLens Excel reports; otherwise skip gracefully + excel_files = list(tracelens_dir.glob("**/individual_reports/*.xlsx")) + if not excel_files: + pytest.skip( + "TraceLens outputs (individual_reports/*.xlsx) are required for analyze_gemm" + ) + + cmd = [ + sys.executable, + str(self.scripts_dir / "analyze_gemm_reports.py"), + "--base-path", str(tracelens_dir), + "--threads"] + [str(t) for t in threads] + [ + "--channels"] + [str(c) for c in channels] + [ + "--ranks"] + [str(r) for r in ranks] + [ + "--top-k", str(top_k), + "--output-file", str(output_file) + ] + + result = subprocess.run( + cmd, capture_output=True, text=True, timeout=self.timeouts['analyze'] + ) + + if result.returncode != 0: + print(f"Error: {result.stderr}") + pytest.fail(f"analyze_gemm_reports.py failed") + + print(f" Created: {output_file}") + return output_file + + def run_plot_variance(self, csv_file: Path, output_dir: Path, + skip_on_missing_deps: bool = True) -> Optional[Path]: + """ + Run plot_gemm_variance.py step. + + Args: + csv_file: Input CSV from analyze_gemm + output_dir: Output directory for plots + skip_on_missing_deps: Skip if matplotlib missing (vs fail) + + Returns: + Path to plots directory or None if skipped + """ + print("\n3. Generating variance plots...") + plots_dir = output_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + result = subprocess.run([ + sys.executable, + str(self.scripts_dir / "plot_gemm_variance.py"), + "--csv-path", str(csv_file), + "--output-dir", str(plots_dir) + ], capture_output=True, text=True, timeout=self.timeouts['plot']) + + if result.returncode != 0: + if "ModuleNotFoundError" in result.stderr or "matplotlib" in result.stderr: + if skip_on_missing_deps: + print(" Warning: Matplotlib not installed - skipping plots") + return None + else: + pytest.fail("Matplotlib not installed") + else: + print(f" Warning: plot_gemm_variance.py failed: {result.stderr}") + return None + + print(f" Plots created in: {plots_dir}") + return plots_dir + + def run_enhancement_scripts(self, + gemm_csv: Path, + test_data_dir: Path, + tracelens_dir: Path, + reports_dir: Path) -> Dict[str, Path]: + """ + Run all enhancement/additional analysis scripts. + + Args: + gemm_csv: GEMM kernels CSV from analyze step + test_data_dir: Test data directory + tracelens_dir: TraceLens output directory + reports_dir: Reports output directory + + Returns: + Dictionary of output file paths + """ + print("\n4. Running additional analysis scripts...") + outputs = {} + + # enhance_gemm_variance_with_timestamps.py + enhanced_file = reports_dir / "gemm_kernels_enhanced.csv" + result = subprocess.run([ + sys.executable, + str(self.scripts_dir / "enhance_gemm_variance_with_timestamps.py"), + "--input-csv", str(gemm_csv), + "--base-path", str(test_data_dir), + "--output-csv", str(enhanced_file) + ], capture_output=True, text=True, timeout=self.timeouts['enhance']) + + if result.returncode == 0: + print(f" Created: {enhanced_file}") + outputs['enhanced'] = enhanced_file + else: + print(f" Warning: Enhancement failed: {result.stderr}") + + # gemm_report_with_collective_overlap.py + overlap_file = reports_dir / "gemm_with_overlap.csv" + result = subprocess.run([ + sys.executable, + str(self.scripts_dir / "gemm_report_with_collective_overlap.py"), + "--input-csv", str(gemm_csv), + "--tracelens-path", str(tracelens_dir), + "--output-csv", str(overlap_file) + ], capture_output=True, text=True, timeout=self.timeouts['overlap']) + + if result.returncode == 0: + print(f" Created: {overlap_file}") + outputs['overlap'] = overlap_file + else: + print(f" Warning: Overlap analysis failed: {result.stderr}") + + # process_gpu_timeline.py + result = subprocess.run([ + sys.executable, + str(self.scripts_dir / "process_gpu_timeline.py"), + "--sweep-dir", str(test_data_dir) + ], capture_output=True, text=True, timeout=self.timeouts['timeline']) + + if result.returncode == 0: + print(" GPU timeline processed") + outputs['timeline'] = test_data_dir / "gpu_timeline_all_configs_mean.xlsx" + else: + print(f" Warning: GPU timeline processing failed: {result.stderr}") + + print("\n Additional analysis completed") + return outputs + + def run_full_pipeline(self, test_data_dir: Path, output_dir: Path) -> Dict[str, Path]: + """ + Run the complete analysis pipeline. + + Args: + test_data_dir: Input test data directory + output_dir: Output directory + + Returns: + Dictionary of all output file paths + """ + # Setup directories + output_dir, reports_dir = self.setup_output_dirs(output_dir) + + # Step 1: TraceLens + print('1') + tracelens_dir = self.run_tracelens(test_data_dir) + print('2') + + # Step 2: Analyze GEMM + gemm_csv = reports_dir / "gemm_kernels.csv" + self.run_analyze_gemm(tracelens_dir, gemm_csv) + + # Step 3: Plot variance + plots_dir = self.run_plot_variance(gemm_csv, output_dir) + + # Step 4: Enhancement scripts + enhancement_outputs = self.run_enhancement_scripts( + gemm_csv, test_data_dir, tracelens_dir, reports_dir + ) + + # Collect all outputs + outputs = { + 'gemm_csv': gemm_csv, + 'plots_dir': plots_dir, + **enhancement_outputs + } + + return outputs + + +# ============================================================================ +# TEST CLASSES +# ============================================================================ + +class TestFullPipeline: + """Full end-to-end regression tests.""" + + @pytest.mark.integration + @pytest.mark.slow + def test_regression(self, test_data_dir, expected_outputs_dir, scripts_dir, generate_baseline): + """ + Main regression test: either generate baseline or run regression test. + + Usage: + pytest test_gemm_regression.py --generate-baseline # Create baseline + pytest test_gemm_regression.py # Run regression test + """ + # Check prerequisites + if not test_data_dir.exists(): + pytest.skip( + "Synthetic test data not found\n" + "Run: python gemm_analysis/generate_synthetic_data.py" + ) + + # Initialize pipeline runner + runner = PipelineRunner(scripts_dir) + + # Determine mode and output directory + if generate_baseline: + # GENERATE BASELINE MODE + print(f"\n{'='*60}") + print("GENERATING BASELINE EXPECTED OUTPUTS") + print(f"{'='*60}") + + runner.run_full_pipeline(test_data_dir, expected_outputs_dir) + + print(f"\n{'='*60}") + print("BASELINE GENERATION COMPLETE") + print(f"Outputs saved to: {expected_outputs_dir}") + print(f"{'='*60}\n") + + else: + # REGRESSION TEST MODE + print(f"\n{'='*60}") + print("RUNNING REGRESSION TEST") + print(f"{'='*60}") + + # Check baseline exists + if not expected_outputs_dir.exists(): + pytest.skip( + "Baseline outputs not found\n" + "Run on origin/main: pytest test_gemm_regression.py --generate-baseline" + ) + + # Run pipeline to actual_outputs + actual_outputs_dir = TEST_DIR / "actual_outputs" + runner.run_full_pipeline(test_data_dir, actual_outputs_dir) + + # Compare outputs to baseline + print(f"\n5. Comparing outputs to baseline...") + print(f"{'='*60}") + + compare_script = TEST_DIR / "compare_outputs.py" + result = subprocess.run([ + sys.executable, + str(compare_script), + "--expected", str(expected_outputs_dir), + "--actual", str(actual_outputs_dir), + "--verbose" + ], capture_output=True, text=True, timeout=60) + + # Print comparison output + print(result.stdout) + if result.stderr: + print(result.stderr) + + # Fail if comparison found differences + if result.returncode != 0: + pytest.fail( + "REGRESSION DETECTED! Current branch outputs differ from baseline.\n" + "If changes are intentional, regenerate baseline:\n" + " pytest test_gemm_regression.py --generate-baseline" + ) + + print(f"\n{'='*60}") + print("ALL REGRESSION TESTS PASSED!") + print(f"{'='*60}\n") + + +class TestIndividualSteps: + """Test individual pipeline components in isolation.""" + + @pytest.fixture(scope="class") + def runner(self, scripts_dir): + """Create a PipelineRunner instance.""" + return PipelineRunner(scripts_dir) + + @pytest.fixture(scope="class") + def temp_output_dir(self, tmp_path_factory): + """Create a temporary output directory for tests.""" + return tmp_path_factory.mktemp("test_outputs") + + def test_analyze_gemm_step(self, runner, test_data_dir, temp_output_dir): + """Test analyze_gemm_reports.py independently.""" + tracelens_dir = test_data_dir / "tracelens_analysis" + if not tracelens_dir.exists(): + pytest.skip("TraceLens outputs required for this test") + + output_file = temp_output_dir / "gemm_test.csv" + + # Test with custom configuration + result = runner.run_analyze_gemm( + tracelens_dir=tracelens_dir, + output_file=output_file, + threads=[256], # Test with single thread config + channels=[28], # Test with single channel config + ranks=[0, 1], # Test with fewer ranks + top_k=5 # Test with different top_k + ) + + assert result.exists() + assert result.stat().st_size > 0 + + def test_plot_generation(self, runner, test_data_dir, temp_output_dir): + """Test plot_gemm_variance.py independently.""" + # First generate a CSV to plot + tracelens_dir = test_data_dir / "tracelens_analysis" + if not tracelens_dir.exists(): + pytest.skip("TraceLens outputs required for this test") + + csv_file = temp_output_dir / "gemm_for_plots.csv" + runner.run_analyze_gemm(tracelens_dir, csv_file) + + # Test plotting + plots_dir = runner.run_plot_variance(csv_file, temp_output_dir) + + if plots_dir is not None: + assert plots_dir.exists() + # Check for expected plot files + expected_plots = [ + "variance_by_threads_boxplot.png", + "variance_by_channels_boxplot.png" + ] + for plot_name in expected_plots: + plot_file = plots_dir / plot_name + if plot_file.exists(): + assert plot_file.stat().st_size > 0 + + +class TestCustomConfigurations: + """Test pipeline with various custom configurations.""" + + @pytest.fixture(scope="class") + def runner(self, scripts_dir): + """Create a PipelineRunner with custom timeouts.""" + custom_timeouts = { + 'tracelens': 300, # Faster timeout for testing + 'analyze': 120, + 'plot': 120, + 'enhance': 120, + 'overlap': 120, + 'timeline': 120 + } + return PipelineRunner(scripts_dir, timeout_seconds=custom_timeouts) + + @pytest.mark.parametrize("threads,channels", [ + ([256], [28]), # Minimal config + ([512], [56]), # Single alternative config + ([256, 512], [28]), # Mixed config + ]) + def test_different_configurations(self, runner, test_data_dir, threads, channels, tmp_path): + """Test pipeline with different thread/channel configurations.""" + tracelens_dir = test_data_dir / "tracelens_analysis" + if not tracelens_dir.exists(): + pytest.skip("TraceLens outputs required for this test") + + output_file = tmp_path / f"gemm_{threads}_{channels}.csv" + + result = runner.run_analyze_gemm( + tracelens_dir=tracelens_dir, + output_file=output_file, + threads=threads, + channels=channels, + ranks=list(range(3)) # Use fewer ranks for speed + ) + + assert result.exists() + # Could add more specific assertions about the output content + + +class TestErrorHandling: + """Test error handling and edge cases.""" + + @pytest.fixture(scope="class") + def runner(self, scripts_dir): + """Create a PipelineRunner instance.""" + return PipelineRunner(scripts_dir) + + def test_missing_tracelens_data(self, runner, tmp_path): + """Test graceful handling when TraceLens data is missing.""" + from _pytest.outcomes import Skipped + + fake_data_dir = tmp_path / "fake_data" + fake_data_dir.mkdir() + + with pytest.raises(Skipped): # pytest.skip raises Skipped exception + runner.run_tracelens(fake_data_dir, skip_on_missing=True) + + def test_invalid_output_path(self, runner, test_data_dir): + """Test handling of invalid output paths.""" + from _pytest.outcomes import Failed + + tracelens_dir = test_data_dir / "tracelens_analysis" + if not tracelens_dir.exists(): + pytest.skip("TraceLens outputs required for this test") + + invalid_path = Path("/invalid/path/that/does/not/exist/output.csv") + + with pytest.raises(Failed): # pytest.fail raises Failed exception + runner.run_analyze_gemm(tracelens_dir, invalid_path) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])