From 84680ded6fab8de03fd987e9c2be0a0a1aa75182 Mon Sep 17 00:00:00 2001 From: Janardhan Pulivarthi Date: Tue, 30 Dec 2025 18:35:33 +0000 Subject: [PATCH 1/3] add two stage optimizer --- .github/copilot-instructions.md | 133 ++++++++- examples/two_stage_optimizer_demo.py | 238 ++++++++++++++++ paper/observability.py | 358 +++++++++++++++++++++++++ paper/optimizer.py | 387 +++++++++++++++++++++++---- tests/test_observability.py | 215 +++++++++++++++ tests/test_two_stage_optimizer.py | 205 ++++++++++++++ 6 files changed, 1474 insertions(+), 62 deletions(-) create mode 100644 examples/two_stage_optimizer_demo.py create mode 100644 paper/observability.py create mode 100644 tests/test_observability.py create mode 100644 tests/test_two_stage_optimizer.py diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 1c46579..da4ee9e 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -10,15 +10,75 @@ The codebase follows a clear separation of concerns: ``` paper/ -├── core.py # PaperMatrix - disk-backed matrix using memory-mapped files -├── plan.py # Lazy evaluation plan tree (EagerNode, AddNode, MultiplyNode, etc.) -├── optimizer.py # Plan inspection, I/O trace generation, and fusion rules -├── backend.py # High-performance execution kernels for matrix operations -├── buffer.py # BufferManager with LRU and Belady's optimal eviction -├── config.py # Centralized configuration (TILE_SIZE, cache sizes) -└── numpy_api.py # NumPy-compatible API layer (pnp.array, pnp.zeros, etc.) +├── core.py # PaperMatrix - disk-backed matrix using memory-mapped files +├── plan.py # Lazy evaluation plan tree (EagerNode, AddNode, MultiplyNode, etc.) +├── optimizer.py # Three-stage optimizer: analyze(), rewrite(), execute_plan() +├── backend.py # High-performance execution kernels for matrix operations +├── buffer.py # BufferManager with LRU and Belady's optimal eviction +├── config.py # Centralized configuration (TILE_SIZE, cache sizes) +├── numpy_api.py # NumPy-compatible API layer (pnp.array, pnp.zeros, etc.) +└── observability.py # Logging, profiling, tracing utilities ``` +## Core Principles + +These principles guide all architectural decisions in the Paper framework: + +1. **Make optimizer purely analytical**: Never materialize data during pattern matching. +2. **Separate concerns**: Analysis (trace + match) → IR rewrite (fuse) → Execution (apply kernels). +3. **Define small, stable contracts**: Each layer (Plan/Node metadata, BufferManager, Backend) has clear interfaces. +4. **Make fusion selection deterministic**: Testable and fast (no I/O during optimization). +5. **Treat heavy tests and benchmarks as gated**: Cached jobs in CI for efficiency. + +## Priority Roadmap + +### **Priority 1 (Critical)** - Stop side-effects in optimizer ✅ DONE +- Remove any calls to `node.execute()` inside optimizer/analysis +- Replace with metadata-only inspection APIs +- **Why**: Prevents data materialization during plan analysis, critical for performance +- **Status**: ✅ Pattern detection in `_detect_fusion_pattern()` uses only `isinstance()` checks and metadata access. No `execute()` calls during analysis phase + +### **Priority 2 (Critical)** - Create two-stage optimizer ✅ DONE +- `analyze(plan)` → Trace + MatchResults +- `rewrite(plan, match)` → FusedPlan (new node types) +- `execute(fused_plan, backend, buffer_mgr)` +- **Why**: Clean separation enables testability and independent optimization +- **Status**: ✅ Three-stage pipeline implemented in `optimizer.py`: `analyze()` generates I/O trace and detects patterns without execution, `rewrite()` prepares optimized plan, `execute_plan()` runs with fusion. Legacy `execute()` maintained for backward compatibility + +### **Priority 3 (High)** - Introduce immutable, hashable Plan representation ❌ NOT DONE +- Deterministic hashing for plan diffs, caching, and baseline comparisons +- **Why**: Enables plan comparison, caching, and regression detection +- **Status**: No `__hash__()` or `__eq__()` methods found in Plan or Node classes + +### **Priority 4 (High)** - Add a small, fast unit test surface ✅ DONE +- Optimizer tests: trace generation, pattern detection, rewrite correctness (use mocked backend) +- **Why**: Fast feedback loop for optimizer development +- **Status**: `test_plan_optimizer.py` has tests for plan construction and I/O trace generation + +### **Priority 5 (High)** - Add integration smoke tests ✅ DONE +- Run fused kernels on tiny matrices (CI quick job) +- **Why**: Catch integration bugs without heavy computation +- **Status**: `test_fusion_operations.py` tests fused kernels with small (128×128) matrices. CI runs all tests in `ci.yml` + +### **Priority 6 (Medium)** - Add benchmarks & regression checks ⚠️ PARTIAL +- Baselines stored as CI artifacts +- Regress only if delta > threshold +- **Why**: Prevent performance regressions systematically +- **Status**: `benchmarks/benchmark_dask.py` exists with benchmarking framework. CI uploads test artifacts but no baseline comparison or regression checks + +### **Priority 7 (Medium)** - Plugin backend API ✅ DONE +- Simple interface: `(inputs, params, output_path, buffer_mgr)` +- Swapping implementations is trivial +- **Why**: Enables experimentation with different execution strategies +- **Status**: All backend kernels follow consistent signature `(A, B, ..., output_path, buffer_manager)`. Interface is clean and swappable + +### **Priority 8 (Low)** - Observability ✅ DONE +- Trace-level logging +- Cost model hooks +- Per-plan flame profiles +- **Why**: Debugging and performance analysis +- **Status**: ✅ Comprehensive observability in `observability.py`: structured logging with `configure_logging()`, `ExecutionProfiler` for timing/flame graphs, `TraceLogger` for execution traces, `CostEstimate` dataclass for I/O cost modeling in optimizer + ## Key Design Patterns ### 1. Lazy Evaluation @@ -41,6 +101,65 @@ paper/ ## Coding Conventions +### Three-Stage Optimizer Usage + +The new three-stage optimizer provides clean separation between analysis, rewriting, and execution: + +```python +from paper.optimizer import analyze, rewrite, execute_plan, estimate_cost +from paper.buffer import BufferManager + +# Stage 1: Analyze (no execution - metadata only) +io_trace, match_result = analyze(plan) + +# Check what pattern was detected +if match_result.is_fusable: + print(f"Detected pattern: {match_result.pattern.value}") + print(f"Parameters: {match_result.parameters}") + +# Estimate execution cost +cost = estimate_cost(plan, match_result) +print(f"Predicted I/O ops: {cost.io_operations}") +print(f"Total cost: {cost.total_cost}") + +# Stage 2: Rewrite (prepare optimized plan) +rewritten_plan = rewrite(plan, match_result) + +# Stage 3: Execute +buffer_manager = BufferManager(max_cache_size_tiles=64, io_trace=io_trace) +result = execute_plan(rewritten_plan, match_result, output_path, buffer_manager) +``` + +### Observability Features + +Enable structured logging, profiling, and tracing: + +```python +from paper.observability import configure_logging, get_profiler, TraceLogger + +# Configure logging +logger = configure_logging(level="DEBUG", log_file="paper.log") + +# Use global profiler +profiler = get_profiler() + +with profiler.profile("my_operation"): + # ... do work ... + pass + +# Print profiling summary +profiler.print_summary() +profiler.save_json("profile.json") +profiler.save_flame_graph("flame.json") + +# Execution tracing +trace = TraceLogger() +trace.begin("operation") +trace.log("step 1 complete") +trace.end() +trace.print() +``` + ### Python Style - Use type hints for function parameters: `def add(A: PaperMatrix, B: PaperMatrix, ...)` - Use `np.float32` as the default dtype diff --git a/examples/two_stage_optimizer_demo.py b/examples/two_stage_optimizer_demo.py new file mode 100644 index 0000000..6f87a80 --- /dev/null +++ b/examples/two_stage_optimizer_demo.py @@ -0,0 +1,238 @@ +""" +Example demonstrating the new two-stage optimizer and observability features. + +This example shows: +1. Three-stage optimizer pipeline (analyze → rewrite → execute) +2. Structured logging +3. Performance profiling +4. Cost estimation +5. Execution tracing +""" + +import numpy as np +import os +import sys +import tempfile + +# Add parent directory to path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from paper.core import PaperMatrix +from paper.plan import Plan, EagerNode +from paper.optimizer import analyze, rewrite, execute_plan, estimate_cost +from paper.buffer import BufferManager +from paper.observability import ( + configure_logging, get_profiler, TraceLogger +) + + +def main(): + """Run the example with observability features.""" + + # ======================================================================== + # Setup: Configure logging + # ======================================================================== + print("\n" + "="*80) + print("EXAMPLE: Two-Stage Optimizer with Observability") + print("="*80 + "\n") + + # Configure structured logging + logger = configure_logging(level="INFO", log_file="paper_example.log") + logger.info("Starting two-stage optimizer example") + + # Get global profiler + profiler = get_profiler() + + # Create trace logger + trace = TraceLogger() + + # ======================================================================== + # Step 1: Create test data + # ======================================================================== + print("Step 1: Creating test matrices...") + trace.begin("Data Creation") + + with profiler.profile("create_test_data"): + test_dir = tempfile.mkdtemp() + shape = (2048, 2048) + dtype = np.float32 + + # Create matrices A and B + A_path = os.path.join(test_dir, "A.bin") + B_path = os.path.join(test_dir, "B.bin") + + np.random.seed(42) + data_A = np.random.rand(*shape).astype(dtype) + data_B = np.random.rand(*shape).astype(dtype) + + data_A.tofile(A_path) + data_B.tofile(B_path) + + A = PaperMatrix(A_path, shape, dtype=dtype, mode='r') + B = PaperMatrix(B_path, shape, dtype=dtype, mode='r') + + trace.log(f"Created matrices: {shape}") + + trace.end() + + # ======================================================================== + # Step 2: Build computation plan + # ======================================================================== + print("\nStep 2: Building computation plan...") + trace.begin("Plan Construction") + + with profiler.profile("build_plan"): + plan_A = Plan(EagerNode(A)) + plan_B = Plan(EagerNode(B)) + + # Build a complex plan: (A + B) * 2.5 + plan = (plan_A + plan_B) * 2.5 + + trace.log(f"Plan: (A + B) * 2.5") + + trace.end() + + # ======================================================================== + # Step 3: STAGE 1 - Analyze (no execution) + # ======================================================================== + print("\nStep 3: STAGE 1 - Analyzing plan (no execution)...") + trace.begin("Stage 1: Analyze") + + with profiler.profile("analyze_plan"): + io_trace, match_result = analyze(plan) + + print(f"\n ✓ Pattern detected: {match_result.pattern.value}") + print(f" ✓ Fusion available: {match_result.is_fusable}") + print(f" ✓ I/O trace length: {len(io_trace)} tile accesses") + print(f" ✓ Input shapes: {match_result.input_shapes}") + print(f" ✓ Parameters: {match_result.parameters}") + + trace.log(f"Pattern: {match_result.pattern.value}") + trace.log(f"I/O trace: {len(io_trace)} accesses") + + trace.end() + + # ======================================================================== + # Step 4: Cost Estimation + # ======================================================================== + print("\nStep 4: Estimating execution cost...") + trace.begin("Cost Estimation") + + with profiler.profile("estimate_cost"): + cost = estimate_cost(plan, match_result) + + print(f"\n ✓ I/O operations: {cost.io_operations}") + print(f" ✓ Compute operations: {cost.compute_operations}") + print(f" ✓ Estimated I/O bytes: {cost.estimated_io_bytes:,}") + print(f" ✓ Cache benefit: {cost.cache_benefit:.1%}") + print(f" ✓ Total cost: {cost.total_cost:.0f}") + + trace.log(f"Total cost: {cost.total_cost:.0f}") + + trace.end() + + # ======================================================================== + # Step 5: STAGE 2 - Rewrite plan + # ======================================================================== + print("\nStep 5: STAGE 2 - Rewriting plan...") + trace.begin("Stage 2: Rewrite") + + with profiler.profile("rewrite_plan"): + rewritten_plan = rewrite(plan, match_result) + trace.log("Plan rewritten for fusion") + + trace.end() + + # ======================================================================== + # Step 6: STAGE 3 - Execute + # ======================================================================== + print("\nStep 6: STAGE 3 - Executing plan...") + trace.begin("Stage 3: Execute") + + output_path = os.path.join(test_dir, "result.bin") + + with profiler.profile("execute_plan"): + buffer_manager = BufferManager(max_cache_size_tiles=32, io_trace=io_trace) + result = execute_plan(rewritten_plan, match_result, output_path, buffer_manager) + + trace.log("Execution complete") + + trace.end() + + # ======================================================================== + # Step 7: Verify results + # ======================================================================== + print("\nStep 7: Verifying results...") + trace.begin("Verification") + + with profiler.profile("verify_result"): + result_data = np.fromfile(output_path, dtype=dtype).reshape(shape) + expected = (data_A + data_B) * 2.5 + + max_diff = np.max(np.abs(result_data - expected)) + print(f"\n ✓ Max difference from expected: {max_diff:.2e}") + + trace.log(f"Verification: max_diff={max_diff:.2e}") + + trace.end() + + # ======================================================================== + # Step 8: Print profiling results + # ======================================================================== + print("\n" + "="*80) + print("PROFILING RESULTS") + print("="*80) + + profiler.print_summary() + + # Save profiling data + profiler.save_json("profile_results.json") + profiler.save_flame_graph("flame_graph.json") + print("✓ Profiling data saved to profile_results.json and flame_graph.json") + + # ======================================================================== + # Step 9: Print execution trace + # ======================================================================== + trace.print() + + # Save trace + trace.save("execution_trace.json") + print("✓ Execution trace saved to execution_trace.json") + + # ======================================================================== + # Step 10: Show cache statistics + # ======================================================================== + print("\n" + "="*80) + print("CACHE STATISTICS") + print("="*80) + + cache_log = buffer_manager.get_log() + + hits = sum(1 for event in cache_log if event[1] == 'HIT') + misses = sum(1 for event in cache_log if event[1] == 'MISS') + evictions = sum(1 for event in cache_log if event[1] == 'EVICT') + total = hits + misses + + print(f"Cache hits: {hits:6d} ({100*hits/total if total > 0 else 0:.1f}%)") + print(f"Cache misses: {misses:6d} ({100*misses/total if total > 0 else 0:.1f}%)") + print(f"Evictions: {evictions:6d}") + print(f"Total accesses: {total:6d}") + print("="*80 + "\n") + + # ======================================================================== + # Cleanup + # ======================================================================== + A.close() + B.close() + result.close() + + import shutil + shutil.rmtree(test_dir) + + print("✓ Example complete!") + print(f"✓ Log file: paper_example.log") + logger.info("Example completed successfully") + + +if __name__ == "__main__": + main() diff --git a/paper/observability.py b/paper/observability.py new file mode 100644 index 0000000..c89f4cb --- /dev/null +++ b/paper/observability.py @@ -0,0 +1,358 @@ +""" +Observability utilities for the Paper framework. + +This module provides: +- Structured logging configuration +- Performance profiling decorators and context managers +- Execution tracing and visualization +""" + +import logging +import time +import functools +from contextlib import contextmanager +from typing import Optional, Dict, Any, List +from dataclasses import dataclass, field +from collections import defaultdict +import json + + +# ============================================================================ +# Logging Configuration +# ============================================================================ + +def configure_logging(level: str = "INFO", log_file: Optional[str] = None): + """ + Configure structured logging for the Paper framework. + + Args: + level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL) + log_file: Optional file path to write logs + """ + log_level = getattr(logging, level.upper()) + + # Create formatters + console_formatter = logging.Formatter( + '%(asctime)s - %(name)s - %(levelname)s - %(message)s', + datefmt='%Y-%m-%d %H:%M:%S' + ) + + detailed_formatter = logging.Formatter( + '%(asctime)s - %(name)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s', + datefmt='%Y-%m-%d %H:%M:%S' + ) + + # Configure root logger for paper module + paper_logger = logging.getLogger('paper') + paper_logger.setLevel(log_level) + + # Console handler + console_handler = logging.StreamHandler() + console_handler.setLevel(log_level) + console_handler.setFormatter(console_formatter) + paper_logger.addHandler(console_handler) + + # File handler (if specified) + if log_file: + file_handler = logging.FileHandler(log_file) + file_handler.setLevel(logging.DEBUG) # Always log everything to file + file_handler.setFormatter(detailed_formatter) + paper_logger.addHandler(file_handler) + + # Prevent propagation to root logger + paper_logger.propagate = False + + return paper_logger + + +# ============================================================================ +# Performance Profiling +# ============================================================================ + +@dataclass +class ProfileEntry: + """Single profile measurement.""" + name: str + start_time: float + end_time: Optional[float] = None + duration: Optional[float] = None + metadata: Dict[str, Any] = field(default_factory=dict) + + def complete(self): + """Mark this entry as complete and calculate duration.""" + self.end_time = time.perf_counter() + self.duration = self.end_time - self.start_time + + def to_dict(self) -> dict: + """Convert to dictionary for serialization.""" + return { + 'name': self.name, + 'duration': self.duration, + 'start_time': self.start_time, + 'end_time': self.end_time, + 'metadata': self.metadata + } + + +class ExecutionProfiler: + """ + Profiler for tracking execution performance. + + Example: + profiler = ExecutionProfiler() + + with profiler.profile("operation"): + # ... do work ... + pass + + profiler.print_summary() + profiler.save_flame_graph("profile.json") + """ + + def __init__(self): + self.entries: List[ProfileEntry] = [] + self.current_stack: List[ProfileEntry] = [] + self.aggregated: Dict[str, List[float]] = defaultdict(list) + self._enabled = True + + @contextmanager + def profile(self, name: str, **metadata): + """ + Context manager for profiling a code block. + + Args: + name: Name of the operation being profiled + **metadata: Additional metadata to attach + """ + if not self._enabled: + yield + return + + entry = ProfileEntry( + name=name, + start_time=time.perf_counter(), + metadata=metadata + ) + + self.current_stack.append(entry) + + try: + yield entry + finally: + entry.complete() + self.entries.append(entry) + self.aggregated[name].append(entry.duration) + self.current_stack.pop() + + def profile_decorator(self, name: Optional[str] = None): + """ + Decorator for profiling function calls. + + Example: + @profiler.profile_decorator("my_function") + def my_function(x, y): + return x + y + """ + def decorator(func): + profile_name = name or f"{func.__module__}.{func.__name__}" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + with self.profile(profile_name): + return func(*args, **kwargs) + return wrapper + return decorator + + def get_summary(self) -> Dict[str, Dict[str, float]]: + """ + Get aggregated statistics for all profiled operations. + + Returns: + Dictionary mapping operation names to statistics + """ + summary = {} + for name, durations in self.aggregated.items(): + if durations: + summary[name] = { + 'count': len(durations), + 'total': sum(durations), + 'mean': sum(durations) / len(durations), + 'min': min(durations), + 'max': max(durations) + } + return summary + + def print_summary(self): + """Print a formatted summary of profiling results.""" + summary = self.get_summary() + + print("\n" + "="*80) + print("EXECUTION PROFILE SUMMARY") + print("="*80) + print(f"{'Operation':<40} {'Count':>8} {'Total (s)':>12} {'Mean (s)':>12}") + print("-"*80) + + # Sort by total time descending + sorted_ops = sorted(summary.items(), key=lambda x: x[1]['total'], reverse=True) + + for name, stats in sorted_ops: + print(f"{name:<40} {stats['count']:>8} {stats['total']:>12.4f} {stats['mean']:>12.6f}") + + print("="*80) + + # Overall statistics + total_time = sum(stats['total'] for stats in summary.values()) + total_count = sum(stats['count'] for stats in summary.values()) + print(f"{'TOTAL':<40} {total_count:>8} {total_time:>12.4f}") + print("="*80 + "\n") + + def save_json(self, filepath: str): + """ + Save profiling results to JSON file. + + Args: + filepath: Path to save JSON data + """ + data = { + 'summary': self.get_summary(), + 'entries': [entry.to_dict() for entry in self.entries] + } + + with open(filepath, 'w') as f: + json.dump(data, f, indent=2) + + def save_flame_graph(self, filepath: str): + """ + Save data suitable for flame graph visualization. + + The output format is compatible with tools like speedscope. + + Args: + filepath: Path to save flame graph data + """ + # Build a simple stack trace format + flame_data = [] + + for entry in self.entries: + flame_data.append({ + 'name': entry.name, + 'value': entry.duration * 1000, # Convert to milliseconds + 'start': entry.start_time * 1000, + 'end': entry.end_time * 1000 if entry.end_time else None, + 'metadata': entry.metadata + }) + + with open(filepath, 'w') as f: + json.dump(flame_data, f, indent=2) + + def reset(self): + """Clear all profiling data.""" + self.entries.clear() + self.current_stack.clear() + self.aggregated.clear() + + def enable(self): + """Enable profiling.""" + self._enabled = True + + def disable(self): + """Disable profiling.""" + self._enabled = False + + +# Global profiler instance +_global_profiler = ExecutionProfiler() + +def get_profiler() -> ExecutionProfiler: + """Get the global profiler instance.""" + return _global_profiler + + +# ============================================================================ +# Trace Visualization +# ============================================================================ + +class TraceLogger: + """ + Logger for execution traces with hierarchical formatting. + + Example: + trace = TraceLogger() + trace.begin("Operation A") + trace.log("Step 1") + trace.begin("Sub-operation") + trace.log("Sub-step") + trace.end() + trace.end() + trace.print() + """ + + def __init__(self): + self.events: List[Dict[str, Any]] = [] + self.depth = 0 + self.logger = logging.getLogger('paper.trace') + + def begin(self, name: str, **metadata): + """Begin a new operation scope.""" + event = { + 'type': 'begin', + 'name': name, + 'depth': self.depth, + 'timestamp': time.perf_counter(), + 'metadata': metadata + } + self.events.append(event) + self.logger.debug(" " * self.depth + f"▶ {name}") + self.depth += 1 + + def end(self, **metadata): + """End the current operation scope.""" + self.depth = max(0, self.depth - 1) + event = { + 'type': 'end', + 'depth': self.depth, + 'timestamp': time.perf_counter(), + 'metadata': metadata + } + self.events.append(event) + + def log(self, message: str, **metadata): + """Log an event within the current scope.""" + event = { + 'type': 'event', + 'message': message, + 'depth': self.depth, + 'timestamp': time.perf_counter(), + 'metadata': metadata + } + self.events.append(event) + self.logger.debug(" " * self.depth + f"• {message}") + + def print(self): + """Print formatted trace.""" + print("\n" + "="*80) + print("EXECUTION TRACE") + print("="*80) + + start_time = self.events[0]['timestamp'] if self.events else 0 + + for event in self.events: + indent = " " * event['depth'] + elapsed = event['timestamp'] - start_time + + if event['type'] == 'begin': + print(f"{elapsed:8.3f}s {indent}▶ {event['name']}") + elif event['type'] == 'event': + print(f"{elapsed:8.3f}s {indent}• {event['message']}") + + print("="*80 + "\n") + + def save(self, filepath: str): + """Save trace to JSON file.""" + with open(filepath, 'w') as f: + json.dump(self.events, f, indent=2) + + def reset(self): + """Clear all trace data.""" + self.events.clear() + self.depth = 0 diff --git a/paper/optimizer.py b/paper/optimizer.py index 012baac..bc534a6 100644 --- a/paper/optimizer.py +++ b/paper/optimizer.py @@ -1,13 +1,80 @@ # --- Purpose: To inspect a plan and choose the best execution strategy. --- import os +import logging +from dataclasses import dataclass +from typing import Optional, Tuple, List, Any +from enum import Enum from . import backend from .plan import Plan, EagerNode, AddNode, MultiplyNode, MultiplyScalarNode -from .buffer import BufferManager +from .buffer import BufferManager from .config import TILE_SIZE +# Configure logging +logger = logging.getLogger(__name__) + + +class FusionPattern(Enum): + """Enumeration of recognized fusion patterns.""" + FUSED_ADD_MULTIPLY = "add_multiply" + FUSED_MATMUL_SCALAR = "matmul_scalar" + FUSED_ADD_MATMUL = "add_matmul" + FUSED_DOUBLE_SCALAR = "double_scalar" + NONE = "none" + + +@dataclass +class MatchResult: + """ + Result of pattern matching analysis. + Contains metadata about detected patterns without materializing data. + """ + pattern: FusionPattern + outer_op_type: Optional[type] + inner_op_type: Optional[type] + kernel_function: Optional[callable] + # Metadata about inputs (no actual matrices) + input_shapes: List[Tuple[int, int]] + parameters: dict # e.g., {'scalar': 2.5} or {'scalar1': 1.5, 'scalar2': 2.0} + + @property + def is_fusable(self) -> bool: + """Returns True if a fusion pattern was detected.""" + return self.pattern != FusionPattern.NONE + + def __repr__(self): + if self.is_fusable: + return f"MatchResult(pattern={self.pattern.value}, shapes={self.input_shapes}, params={self.parameters})" + return "MatchResult(pattern=none)" + + +@dataclass +class CostEstimate: + """ + Cost model for execution estimation. + """ + io_operations: int # Number of I/O operations + compute_operations: int # Number of compute operations + estimated_io_bytes: int # Estimated I/O in bytes + cache_benefit: float # Estimated cache hit improvement (0-1) + + @property + def total_cost(self) -> float: + """Simple cost model: weighted sum of I/O and compute.""" + # I/O is typically 100x more expensive than compute + io_weight = 100.0 + compute_weight = 1.0 + return (self.io_operations * io_weight + + self.compute_operations * compute_weight) + + def __repr__(self): + return (f"CostEstimate(io_ops={self.io_operations}, " + f"compute_ops={self.compute_operations}, " + f"io_bytes={self.estimated_io_bytes}, " + f"total_cost={self.total_cost:.0f})") + # The rule registry # Pattern: (OuterOp, InnerOp), Kernel: function_to_execute FUSION_RULES = [ @@ -17,12 +84,12 @@ ((MultiplyScalarNode, MultiplyScalarNode), backend.execute_fused_double_scalar), ] -def _generate_trace_recursive(op_node): - """ - A recursive helper to generate the I/O trace for a plan node. - """ - if isinstance(op_node, EagerNode): - # Eager nodes are leaves; so, doesn't have preceding operations +def _generate_trace_recursive(op_node): + """ + A recursive helper to generate the I/O trace for a plan node. + """ + if isinstance(op_node, EagerNode): + # Eager nodes are leaves; so, doesn't have preceding operations return [] # Recursively find the trace for the inputs @@ -89,59 +156,269 @@ def generate_io_trace(plan: Plan) -> list: Generates a complete, ordered list of all tile accesses (the I/O trace) given a computation plan. """ - print("Optimizer: Generating I/O trace for optimal caching...") - - print(plan) + logger.info("Generating I/O trace for optimal caching...") + logger.debug(f"Plan structure: {plan}") trace = _generate_trace_recursive(plan.op) - return [(os.path.basename(path), r, c) for path, r, c in trace] + result = [(os.path.basename(path), r, c) for path, r, c in trace] + + logger.info(f"Generated I/O trace with {len(result)} tile accesses") + return result -def execute(plan, output_path: str, buffer_manager: BufferManager | None): +def analyze(plan: Plan) -> Tuple[List, MatchResult]: + """ + Stage 1: Analyze the plan without executing anything. + + Returns: + - I/O trace (list of tile accesses) + - MatchResult (detected fusion patterns with metadata only) + """ + logger.info("=== STAGE 1: ANALYZE ===") + + # Generate I/O trace + io_trace = generate_io_trace(plan) + + # Pattern matching (metadata only - no execute() calls) + match_result = _detect_fusion_pattern(plan) + + logger.info(f"Analysis complete: {match_result}") + return io_trace, match_result + + +def _detect_fusion_pattern(plan: Plan) -> MatchResult: """ - The main optimizer entry point. It checks for patterns and executes. - 1. Checks the plan against FUSION_RULES. - 2. If a rule matches, calls the corresponding fast kernel from lazy.py/eager.py. - 3. If no rules match, executes the plan step-by-step (default path). + Detect fusion patterns using metadata only (no execute() calls). + This is purely analytical - inspects the plan tree structure. """ - # Check if the plan's op matches any fusion rule + logger.debug("Detecting fusion patterns...") + + # Check against each fusion rule for pattern, kernel in FUSION_RULES: OuterOp, InnerOp = pattern - if isinstance(plan.op, OuterOp) and isinstance(plan.op.left.op, InnerOp): - print(f"✨ Optimizer: Found pattern {pattern}. Using fused kernel.") - - # Pattern 1: (A + B) * scalar - if OuterOp == MultiplyScalarNode and InnerOp == AddNode: - add_node = plan.op.left.op - matrix_A = add_node.left.op.execute(None) - matrix_B = add_node.right.op.execute(None) - scalar = plan.op.right - return kernel(matrix_A, matrix_B, scalar, output_path) - - # Pattern 2: (A @ B) * scalar - elif OuterOp == MultiplyScalarNode and InnerOp == MultiplyNode: - mul_node = plan.op.left.op - matrix_A = mul_node.left.op.execute(None) - matrix_B = mul_node.right.op.execute(None) - scalar = plan.op.right - return kernel(matrix_A, matrix_B, scalar, output_path) - - # Pattern 3: (A + B) @ C - elif OuterOp == MultiplyNode and InnerOp == AddNode: - add_node = plan.op.left.op - matrix_A = add_node.left.op.execute(None) - matrix_B = add_node.right.op.execute(None) - matrix_C = plan.op.right.op.execute(None) - return kernel(matrix_A, matrix_B, matrix_C, output_path) - - # Pattern 4: (A * scalar1) * scalar2 - elif OuterOp == MultiplyScalarNode and InnerOp == MultiplyScalarNode: - inner_mul_node = plan.op.left.op - matrix_A = inner_mul_node.left.op.execute(None) - scalar1 = inner_mul_node.right - scalar2 = plan.op.right - return kernel(matrix_A, scalar1, scalar2, output_path) - - # If no special rule matches, execute the default, non-fused way - print("Optimizer: No fusion pattern found. Executing default.") - return plan.op.execute(output_path) # The original, slower path + + # Pattern detection using isinstance checks (no execution) + # Note: plan.op contains the operation node, not another Plan + if isinstance(plan.op, OuterOp) and hasattr(plan.op, 'left'): + # plan.op.left is the inner operation node (not a Plan) + if isinstance(plan.op.left, InnerOp): + # Pattern matched! Extract metadata + logger.debug(f"Pattern detected: {OuterOp.__name__}({InnerOp.__name__})") + + # Determine pattern type and extract parameters + if OuterOp == MultiplyScalarNode and InnerOp == AddNode: + add_node = plan.op.left # This is the AddNode + return MatchResult( + pattern=FusionPattern.FUSED_ADD_MULTIPLY, + outer_op_type=OuterOp, + inner_op_type=InnerOp, + kernel_function=kernel, + input_shapes=[add_node.left.shape, add_node.right.shape], + parameters={'scalar': plan.op.right} + ) + + elif OuterOp == MultiplyScalarNode and InnerOp == MultiplyNode: + mul_node = plan.op.left # This is the MultiplyNode + return MatchResult( + pattern=FusionPattern.FUSED_MATMUL_SCALAR, + outer_op_type=OuterOp, + inner_op_type=InnerOp, + kernel_function=kernel, + input_shapes=[mul_node.left.shape, mul_node.right.shape], + parameters={'scalar': plan.op.right} + ) + + elif OuterOp == MultiplyNode and InnerOp == AddNode: + add_node = plan.op.left # This is the AddNode + return MatchResult( + pattern=FusionPattern.FUSED_ADD_MATMUL, + outer_op_type=OuterOp, + inner_op_type=InnerOp, + kernel_function=kernel, + input_shapes=[add_node.left.shape, add_node.right.shape, plan.op.right.shape], + parameters={} + ) + + elif OuterOp == MultiplyScalarNode and InnerOp == MultiplyScalarNode: + inner_mul = plan.op.left # This is the inner MultiplyScalarNode + return MatchResult( + pattern=FusionPattern.FUSED_DOUBLE_SCALAR, + outer_op_type=OuterOp, + inner_op_type=InnerOp, + kernel_function=kernel, + input_shapes=[inner_mul.left.shape], + parameters={'scalar1': inner_mul.right, 'scalar2': plan.op.right} + ) + + # No pattern matched + logger.debug("No fusion pattern detected") + return MatchResult( + pattern=FusionPattern.NONE, + outer_op_type=None, + inner_op_type=None, + kernel_function=None, + input_shapes=[], + parameters={} + ) + + +def estimate_cost(plan: Plan, match_result: MatchResult) -> CostEstimate: + """ + Estimate execution cost for a plan. + + Args: + plan: The computation plan + match_result: Result from pattern analysis + + Returns: + CostEstimate with predicted I/O and compute costs + """ + logger.debug("Estimating execution cost...") + + # Calculate number of tiles + def count_tiles(shape): + rows_tiles = (shape[0] + TILE_SIZE - 1) // TILE_SIZE + cols_tiles = (shape[1] + TILE_SIZE - 1) // TILE_SIZE + return rows_tiles * cols_tiles + + io_ops = 0 + compute_ops = 0 + io_bytes = 0 + + if match_result.is_fusable: + # Fused operations reduce I/O by combining operations + for shape in match_result.input_shapes: + tiles = count_tiles(shape) + io_ops += tiles # Read each input once + io_bytes += shape[0] * shape[1] * 4 # float32 + compute_ops += tiles + + # Fusion benefit: ~30% fewer I/O operations + cache_benefit = 0.3 + else: + # Unfused: each operation reads and writes separately + shape = plan.shape + tiles = count_tiles(shape) + # Estimate: 2x I/O for unfused (read + write per operation) + io_ops = tiles * 2 + io_bytes = shape[0] * shape[1] * 4 * 2 + compute_ops = tiles + cache_benefit = 0.0 + + estimate = CostEstimate( + io_operations=io_ops, + compute_operations=compute_ops, + estimated_io_bytes=io_bytes, + cache_benefit=cache_benefit + ) + + logger.info(f"Cost estimate: {estimate}") + return estimate + + +def rewrite(plan: Plan, match_result: MatchResult) -> Plan: + """ + Stage 2: Rewrite the plan based on match results. + + Currently returns the original plan (future: create fused node IR). + In a full implementation, this would create new FusedNode types. + + Args: + plan: Original computation plan + match_result: Pattern matching results from analyze() + + Returns: + Rewritten plan (potentially with fused operations) + """ + logger.info("=== STAGE 2: REWRITE ===") + + if match_result.is_fusable: + logger.info(f"Plan will use fused kernel: {match_result.pattern.value}") + # Future: Create FusedNode IR representation here + # For now, we return the original plan and handle fusion in execute() + else: + logger.info("Plan will execute unfused (no optimization)") + + return plan + + +def execute_plan(plan: Plan, match_result: MatchResult, output_path: str, + buffer_manager: BufferManager | None): + """ + Stage 3: Execute the plan using the analysis results. + + Args: + plan: The computation plan + match_result: Results from analyze() stage + output_path: Where to write the result + buffer_manager: Optional buffer manager for caching + + Returns: + Computed PaperMatrix result + """ + logger.info("=== STAGE 3: EXECUTE ===") + + if match_result.is_fusable: + logger.info(f"✨ Executing fused kernel: {match_result.pattern.value}") + + # Extract leaf matrices (now that we're actually executing) + def get_leaf_matrix(node): + """Helper to get the PaperMatrix from an EagerNode.""" + if isinstance(node, EagerNode): + return node.matrix + # Recursively execute to get the matrix + return node.execute(None, buffer_manager) + + # Execute based on pattern + if match_result.pattern == FusionPattern.FUSED_ADD_MULTIPLY: + add_node = plan.op.left # AddNode + matrix_A = get_leaf_matrix(add_node.left) + matrix_B = get_leaf_matrix(add_node.right) + scalar = match_result.parameters['scalar'] + return match_result.kernel_function(matrix_A, matrix_B, scalar, output_path, buffer_manager) + + elif match_result.pattern == FusionPattern.FUSED_MATMUL_SCALAR: + mul_node = plan.op.left # MultiplyNode + matrix_A = get_leaf_matrix(mul_node.left) + matrix_B = get_leaf_matrix(mul_node.right) + scalar = match_result.parameters['scalar'] + return match_result.kernel_function(matrix_A, matrix_B, scalar, output_path) + + elif match_result.pattern == FusionPattern.FUSED_ADD_MATMUL: + add_node = plan.op.left # AddNode + matrix_A = get_leaf_matrix(add_node.left) + matrix_B = get_leaf_matrix(add_node.right) + matrix_C = get_leaf_matrix(plan.op.right) + return match_result.kernel_function(matrix_A, matrix_B, matrix_C, output_path) + + elif match_result.pattern == FusionPattern.FUSED_DOUBLE_SCALAR: + inner_mul = plan.op.left # Inner MultiplyScalarNode + matrix_A = get_leaf_matrix(inner_mul.left) + scalar1 = match_result.parameters['scalar1'] + scalar2 = match_result.parameters['scalar2'] + return match_result.kernel_function(matrix_A, scalar1, scalar2, output_path) + + # Unfused execution + logger.info("Executing unfused plan") + return plan.op.execute(output_path, buffer_manager) + + +def execute(plan, output_path: str, buffer_manager: BufferManager | None): + """ + Legacy entry point for backward compatibility. + Now uses the three-stage optimizer pipeline. + + The new three-stage approach: + 1. analyze() - Generate I/O trace and detect patterns (no execution) + 2. rewrite() - Transform plan based on patterns (future: create IR) + 3. execute_plan() - Execute using analysis results + + DEPRECATED: Use analyze() → rewrite() → execute_plan() for new code. + """ + logger.warning("Using legacy execute() - consider migrating to three-stage pipeline") + + # Three-stage pipeline + io_trace, match_result = analyze(plan) + cost = estimate_cost(plan, match_result) + rewritten_plan = rewrite(plan, match_result) + return execute_plan(rewritten_plan, match_result, output_path, buffer_manager) diff --git a/tests/test_observability.py b/tests/test_observability.py new file mode 100644 index 0000000..6792e00 --- /dev/null +++ b/tests/test_observability.py @@ -0,0 +1,215 @@ +""" +Unit tests for observability features (Priority 8). +""" + +import unittest +import os +import tempfile +import shutil +import sys +import json + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from paper.observability import ( + configure_logging, ExecutionProfiler, TraceLogger, + get_profiler +) + + +class TestObservability(unittest.TestCase): + """Test cases for observability utilities.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + + def tearDown(self): + """Clean up after each test.""" + if os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir) + + def test_logging_configuration(self): + """Test that logging configuration works.""" + log_file = os.path.join(self.test_dir, "test.log") + logger = configure_logging(level="DEBUG", log_file=log_file) + + self.assertIsNotNone(logger) + logger.info("Test message") + + # Verify log file was created + self.assertTrue(os.path.exists(log_file)) + + def test_execution_profiler_context_manager(self): + """Test profiler context manager.""" + profiler = ExecutionProfiler() + + with profiler.profile("operation1"): + import time + time.sleep(0.01) + + with profiler.profile("operation2"): + time.sleep(0.02) + + # Verify entries were recorded + self.assertEqual(len(profiler.entries), 2) + self.assertEqual(profiler.entries[0].name, "operation1") + self.assertEqual(profiler.entries[1].name, "operation2") + + # Verify durations + self.assertGreater(profiler.entries[0].duration, 0.01) + self.assertGreater(profiler.entries[1].duration, 0.02) + + def test_profiler_decorator(self): + """Test profiler function decorator.""" + profiler = ExecutionProfiler() + + @profiler.profile_decorator("test_func") + def my_function(x): + import time + time.sleep(0.01) + return x * 2 + + result = my_function(5) + + self.assertEqual(result, 10) + self.assertEqual(len(profiler.entries), 1) + self.assertEqual(profiler.entries[0].name, "test_func") + self.assertGreater(profiler.entries[0].duration, 0.01) + + def test_profiler_summary_statistics(self): + """Test profiler summary generation.""" + profiler = ExecutionProfiler() + + # Profile same operation multiple times + for i in range(5): + with profiler.profile("repeated_op"): + pass + + summary = profiler.get_summary() + + self.assertIn("repeated_op", summary) + self.assertEqual(summary["repeated_op"]["count"], 5) + self.assertGreater(summary["repeated_op"]["total"], 0) + self.assertGreater(summary["repeated_op"]["mean"], 0) + + def test_profiler_json_export(self): + """Test exporting profiler results to JSON.""" + profiler = ExecutionProfiler() + + with profiler.profile("op1", param=42): + pass + + json_file = os.path.join(self.test_dir, "profile.json") + profiler.save_json(json_file) + + # Verify file was created and contains valid JSON + self.assertTrue(os.path.exists(json_file)) + with open(json_file) as f: + data = json.load(f) + + self.assertIn("summary", data) + self.assertIn("entries", data) + self.assertEqual(len(data["entries"]), 1) + + def test_profiler_flame_graph_export(self): + """Test exporting flame graph data.""" + profiler = ExecutionProfiler() + + with profiler.profile("outer"): + with profiler.profile("inner"): + pass + + flame_file = os.path.join(self.test_dir, "flame.json") + profiler.save_flame_graph(flame_file) + + # Verify file was created + self.assertTrue(os.path.exists(flame_file)) + with open(flame_file) as f: + data = json.load(f) + + self.assertEqual(len(data), 2) # outer and inner + + def test_profiler_enable_disable(self): + """Test enabling/disabling profiler.""" + profiler = ExecutionProfiler() + + # Profile while enabled + with profiler.profile("op1"): + pass + + # Disable and profile + profiler.disable() + with profiler.profile("op2"): + pass + + # Re-enable and profile + profiler.enable() + with profiler.profile("op3"): + pass + + # Should only have recorded op1 and op3 + self.assertEqual(len(profiler.entries), 2) + self.assertEqual(profiler.entries[0].name, "op1") + self.assertEqual(profiler.entries[1].name, "op3") + + def test_trace_logger_basic(self): + """Test basic trace logging.""" + trace = TraceLogger() + + trace.begin("operation") + trace.log("step1") + trace.log("step2") + trace.end() + + # Verify events were recorded + self.assertEqual(len(trace.events), 4) + self.assertEqual(trace.events[0]['type'], 'begin') + self.assertEqual(trace.events[1]['type'], 'event') + self.assertEqual(trace.events[3]['type'], 'end') + + def test_trace_logger_nested(self): + """Test nested trace logging.""" + trace = TraceLogger() + + trace.begin("outer") + trace.log("outer_step") + trace.begin("inner") + trace.log("inner_step") + trace.end() + trace.end() + + # Verify depth tracking + self.assertEqual(trace.events[0]['depth'], 0) # outer begin + self.assertEqual(trace.events[1]['depth'], 1) # outer_step + self.assertEqual(trace.events[2]['depth'], 1) # inner begin + self.assertEqual(trace.events[3]['depth'], 2) # inner_step + + def test_trace_logger_save(self): + """Test saving trace to JSON.""" + trace = TraceLogger() + + trace.begin("op") + trace.log("step") + trace.end() + + trace_file = os.path.join(self.test_dir, "trace.json") + trace.save(trace_file) + + # Verify file + self.assertTrue(os.path.exists(trace_file)) + with open(trace_file) as f: + data = json.load(f) + + self.assertEqual(len(data), 3) + + def test_global_profiler_singleton(self): + """Test that get_profiler() returns singleton instance.""" + profiler1 = get_profiler() + profiler2 = get_profiler() + + self.assertIs(profiler1, profiler2) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_two_stage_optimizer.py b/tests/test_two_stage_optimizer.py new file mode 100644 index 0000000..fcab14b --- /dev/null +++ b/tests/test_two_stage_optimizer.py @@ -0,0 +1,205 @@ +""" +Unit tests for the two-stage optimizer (Priority 2). + +Tests the analyze → rewrite → execute pipeline. +""" + +import unittest +import os +import tempfile +import shutil +import numpy as np +import sys + +# Add the parent directory to the path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +from paper.core import PaperMatrix +from paper.plan import Plan, EagerNode +from paper import optimizer +from paper.optimizer import ( + analyze, rewrite, execute_plan, estimate_cost, + FusionPattern, MatchResult +) +from paper.buffer import BufferManager + + +class TestTwoStageOptimizer(unittest.TestCase): + """Test cases for the two-stage optimizer.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + self.test_shape = (128, 128) + self.test_dtype = np.float32 + + # Create test matrices + self.matrix_A_path = os.path.join(self.test_dir, "A.bin") + self.matrix_B_path = os.path.join(self.test_dir, "B.bin") + + np.random.seed(42) + self.data_A = np.random.rand(*self.test_shape).astype(self.test_dtype) + self.data_B = np.random.rand(*self.test_shape).astype(self.test_dtype) + + self.data_A.tofile(self.matrix_A_path) + self.data_B.tofile(self.matrix_B_path) + + self.A = PaperMatrix(self.matrix_A_path, self.test_shape, dtype=self.test_dtype, mode='r') + self.B = PaperMatrix(self.matrix_B_path, self.test_shape, dtype=self.test_dtype, mode='r') + + def tearDown(self): + """Clean up after each test.""" + if hasattr(self, 'A'): + self.A.close() + if hasattr(self, 'B'): + self.B.close() + if os.path.exists(self.test_dir): + shutil.rmtree(self.test_dir) + + def test_analyze_detects_fused_add_multiply_pattern(self): + """Test that analyze() detects (A + B) * scalar pattern without executing.""" + # Build plan: (A + B) * 2.5 + plan_A = Plan(EagerNode(self.A)) + plan_B = Plan(EagerNode(self.B)) + plan = (plan_A + plan_B) * 2.5 + + # Analyze (should not execute anything) + io_trace, match_result = analyze(plan) + + # Verify pattern was detected + self.assertTrue(match_result.is_fusable) + self.assertEqual(match_result.pattern, FusionPattern.FUSED_ADD_MULTIPLY) + self.assertEqual(match_result.parameters['scalar'], 2.5) + self.assertEqual(len(match_result.input_shapes), 2) + + # Verify I/O trace was generated + self.assertIsInstance(io_trace, list) + self.assertGreater(len(io_trace), 0) + + def test_analyze_no_pattern_detection(self): + """Test that analyze() returns NONE for non-fusable patterns.""" + # Simple addition (not fusable) + plan_A = Plan(EagerNode(self.A)) + plan_B = Plan(EagerNode(self.B)) + plan = plan_A + plan_B + + # Analyze + io_trace, match_result = analyze(plan) + + # Verify no pattern + self.assertFalse(match_result.is_fusable) + self.assertEqual(match_result.pattern, FusionPattern.NONE) + + def test_cost_estimate_for_fused_operation(self): + """Test that cost estimation works for fused operations.""" + plan_A = Plan(EagerNode(self.A)) + plan_B = Plan(EagerNode(self.B)) + plan = (plan_A + plan_B) * 2.5 + + _, match_result = analyze(plan) + cost = estimate_cost(plan, match_result) + + # Verify cost estimate structure + self.assertGreater(cost.io_operations, 0) + self.assertGreater(cost.compute_operations, 0) + self.assertGreater(cost.estimated_io_bytes, 0) + self.assertGreater(cost.cache_benefit, 0) # Fusion provides benefit + self.assertGreater(cost.total_cost, 0) + + def test_cost_estimate_for_unfused_operation(self): + """Test that cost estimation works for unfused operations.""" + plan_A = Plan(EagerNode(self.A)) + plan_B = Plan(EagerNode(self.B)) + plan = plan_A + plan_B + + _, match_result = analyze(plan) + cost = estimate_cost(plan, match_result) + + # Verify cost estimate + self.assertGreater(cost.io_operations, 0) + self.assertEqual(cost.cache_benefit, 0.0) # No fusion benefit + + def test_three_stage_pipeline_execution(self): + """Test the full three-stage pipeline: analyze → rewrite → execute.""" + plan_A = Plan(EagerNode(self.A)) + plan_B = Plan(EagerNode(self.B)) + plan = (plan_A + plan_B) * 2.5 + + output_path = os.path.join(self.test_dir, "result.bin") + + # Stage 1: Analyze + io_trace, match_result = analyze(plan) + self.assertTrue(match_result.is_fusable) + + # Stage 2: Rewrite + rewritten_plan = rewrite(plan, match_result) + self.assertIsNotNone(rewritten_plan) + + # Stage 3: Execute + buffer_manager = BufferManager(max_cache_size_tiles=8, io_trace=io_trace) + result = execute_plan(rewritten_plan, match_result, output_path, buffer_manager) + + # Verify result + self.assertIsInstance(result, PaperMatrix) + result_data = np.fromfile(output_path, dtype=self.test_dtype).reshape(self.test_shape) + + expected = (self.data_A + self.data_B) * 2.5 + np.testing.assert_allclose(result_data, expected, rtol=1e-5) + + def test_legacy_execute_backward_compatibility(self): + """Test that legacy execute() still works via three-stage pipeline.""" + plan_A = Plan(EagerNode(self.A)) + plan_B = Plan(EagerNode(self.B)) + plan = (plan_A + plan_B) * 2.5 + + output_path = os.path.join(self.test_dir, "result_legacy.bin") + + # Use legacy execute() + result = optimizer.execute(plan, output_path, None) + + # Verify it produces correct results + self.assertIsInstance(result, PaperMatrix) + result_data = np.fromfile(output_path, dtype=self.test_dtype).reshape(self.test_shape) + + expected = (self.data_A + self.data_B) * 2.5 + np.testing.assert_allclose(result_data, expected, rtol=1e-5) + + def test_analyze_double_scalar_pattern(self): + """Test detection of (A * scalar1) * scalar2 pattern.""" + plan_A = Plan(EagerNode(self.A)) + plan = (plan_A * 2.0) * 3.0 + + _, match_result = analyze(plan) + + self.assertTrue(match_result.is_fusable) + self.assertEqual(match_result.pattern, FusionPattern.FUSED_DOUBLE_SCALAR) + self.assertEqual(match_result.parameters['scalar1'], 2.0) + self.assertEqual(match_result.parameters['scalar2'], 3.0) + + def test_no_execute_during_analyze(self): + """Critical test: analyze() must not call execute() on nodes.""" + plan_A = Plan(EagerNode(self.A)) + plan_B = Plan(EagerNode(self.B)) + plan = (plan_A + plan_B) * 2.5 + + # Track if execute was called + original_execute = self.A.get_tile + execute_called = [] + + def tracked_get_tile(*args, **kwargs): + execute_called.append(True) + return original_execute(*args, **kwargs) + + self.A.get_tile = tracked_get_tile + + # Analyze should not trigger any get_tile calls + _, match_result = analyze(plan) + + # Verify analyze did not execute + self.assertEqual(len(execute_called), 0, + "analyze() should not call execute() or get_tile()") + self.assertTrue(match_result.is_fusable) + + +if __name__ == '__main__': + unittest.main() From 200e2a89b038f959375de5607c880a7ad11b9caf Mon Sep 17 00:00:00 2001 From: Janardhan Pulivarthi Date: Tue, 30 Dec 2025 18:41:19 +0000 Subject: [PATCH 2/3] architecture diagram update as per new changes --- paper-architecture.svg | 187 ++++++++++++++++++++++++++--------------- 1 file changed, 118 insertions(+), 69 deletions(-) diff --git a/paper-architecture.svg b/paper-architecture.svg index 4e23c7e..78f1033 100644 --- a/paper-architecture.svg +++ b/paper-architecture.svg @@ -1,9 +1,8 @@ - + \ No newline at end of file From cae5f66f83aef6d05a093da3e0bc040a2689ead1 Mon Sep 17 00:00:00 2001 From: Janardhan Pulivarthi Date: Tue, 30 Dec 2025 18:45:31 +0000 Subject: [PATCH 3/3] change diagram styling --- paper-architecture.svg | 290 ++++++++++++++++------------------------- 1 file changed, 113 insertions(+), 177 deletions(-) diff --git a/paper-architecture.svg b/paper-architecture.svg index 78f1033..da96ab4 100644 --- a/paper-architecture.svg +++ b/paper-architecture.svg @@ -1,204 +1,140 @@ - - - + + + - - + + - - 'paper' Framework Architecture - Out-of-Core System with Three-Stage Optimizer & Observability + 'paper' Framework Architecture + High-Performance Out-of-Core Execution Engine - - - - 1. User Application - plan = (plan_A @ plan_B) * 2 - result, buffer_mgr = plan.compute(output_path) + + + 1. User Application Interface + + plan = (matrix_a @ matrix_b) + 10 + result = plan.compute(cache_mode='optimal') - - User builds a Plan + + Initialize Plan Graph - - - - 2. Plan Layer (plan.py) - - Builds a Computation Graph (AST) of PlanNodes - Plan(op_node=MultiplyScalarNode(...)) + + + 2. Plan Layer (AST Generation) + + Constructs Directed Acyclic Graph (DAG) of **PlanNodes** - - .compute() triggers optimizer + + Trigger Three-Stage Optimization - - - - 3. Three-Stage Optimizer (optimizer.py) - - - - STAGE 1: analyze() - ① - • Generate I/O Trace - • Detect Fusion Patterns - • Metadata-only (no exec) - → (io_trace, MatchResult) + + + 3. Three-Stage Optimizer - - - STAGE 2: rewrite() - ② - • Transform Plan - • Apply Fusion Rules - • Optimize IR - → Optimized Plan + + ① + Analyze + • Generate I/O Trace + • Metadata Analysis + → MatchResult - - - STAGE 3: execute_plan() - ③ - • Select fused/unfused kernels based on MatchResult - • Execute with BufferManager - → Calls Backend Kernels - - - - Cost Model - estimate_cost() - • I/O ops - • Compute ops - • Cache benefit + + ② + Rewrite + • Operator Fusion + • IR Transformation + → OptimizedPlan + + + ③ + Execute + • Kernel Selection + • Thread Scheduling + → Kernel Dispatch + + + Heuristic Cost Model + Estimates execution cost based on: **Compute GFLOPS** | **I/O Bandwidth** | **Memory Pressure** - - - - - - - - - - I/O Trace - - - Calls Backend Kernels - - - - - 4. Execution & Memory Layer - - - - BufferManager (buffer.py) - - In-Memory Cache (RAM) - [Tile A, Tile B, ...] - Size: N tiles + + + + + + 4. Execution & Memory Layer + + + Buffer Manager + + Tile Cache (LRU/Optimal) + [Tile_01] [Tile_02] ... - - Eviction Policy - if io_trace: _evict_optimal() - else: _evict_lru() - - - - Backend Kernels (backend.py) - - Parallel Tile Processing - - Fused & Standard Kernels - add(), multiply(), ... - execute_fused_add_multiply() - Interactions - 1. Requests tiles from BufferManager - 2. Writes results to Core Layer + + Eviction Logic + _evict_optimal(io_trace) + + + Backend Kernels + • Parallel Tile Processing + • Vectorized Operations (NumPy/C++) + Fused: add_mul_kernel() + Atomic: matmul_kernel() - - - - get_tile() - - - - On Miss: Read Tile - Write Result - - - - - 5. Core I/O Layer (core.py) - - - PaperMatrix - Abstracts direct disk interaction. - - - Read Path - numpy.memmap - Efficient tile reads - via OS virtual memory. - - - Write Path - open() / seek() / write() - True out-of-core writing - with minimal memory. + + I/O Trace Injection + + + + 5. Core I/O Layer (PaperMatrix) - - - - Physical Disk Storage (SSD / HDD) + + Read Path + numpy.memmap() + OS-level virtual memory mapping + for random access performance. + + + Write Path + File.seek() & .write() + Direct-to-disk tiled streaming + ensuring memory footprint stability. - - - - Observability (observability.py) + + + Observability & Telemetry - - - Cross-Cutting Concerns - - 📊 Structured Logging - configure_logging() - - ⏱️ Performance Profiling - ExecutionProfiler - - 🔍 Execution Tracing - TraceLogger + + Monitoring + Logs | Performance Traces | Resource Usage - \ No newline at end of file