import torch
import torch.nn as nn
from torch.export import export, Dim

class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model=256, n_heads=4, d_ff=512):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-LN Self-Attention
        norm_x = self.ln1(x)
        attn_out, _ = self.attn(norm_x, norm_x, norm_x)
        x = x + attn_out
        # Pre-LN MLP Feed-Forward
        x = x + self.mlp(self.ln2(x))
        return x

def run_lab1():
    print("=" * 60)
    print("LAB 1: SOUND GRAPH CAPTURE WITH DYNAMIC SHAPES")
    print("=" * 60)
    
    model = TransformerEncoderBlock(d_model=256, n_heads=4, d_ff=512).eval()
    
    # 1. Define sample input: [batch=2, seq_len=64, d_model=256]
    example_x = torch.randn(2, 64, 256)
    
    # 2. Declare dynamic dimensions with strict symbolic inequalities
    batch_dim = Dim("batch", min=1, max=128)
    seq_dim = Dim("seq_len", min=8, max=1024)
    dynamic_shapes = {"x": (batch_dim, seq_dim, 256)}
    
    # 3. Sound Export via torch.export
    print("[*] Executing torch.export.export()...")
    exported_prog = export(model, args=(example_x,), dynamic_shapes=dynamic_shapes)
    
    print("
[+] Graph Capture Succeeded!")
    print(f"ExportedProgram Type: {type(exported_prog)}")
    print(f"Graph Signature Input Names: {exported_prog.graph_signature.user_inputs}")
    print(f"Graph Signature Output Names: {exported_prog.graph_signature.user_outputs}")
    
    # 4. Verify Symbolic Range Constraints
    print("
[*] Symbolic Range Constraints:")
    for sym_var, constraint in exported_prog.range_constraints.items():
        print(f"  Symbol {sym_var}: min={constraint.lower}, max={constraint.upper}")
        
    # 5. Execute with different variable shapes to prove zero-recompile execution
    test_x1 = torch.randn(8, 128, 256)
    test_x2 = torch.randn(32, 256, 256)
    
    out1 = exported_prog.module()(test_x1)
    out2 = exported_prog.module()(test_x2)
    print(f"
[+] Verified dynamic execution: Input [8, 128, 256] -> Output {list(out1.shape)}")
    print(f"[+] Verified dynamic execution: Input [32, 256, 256] -> Output {list(out2.shape)}")

if __name__ == "__main__":
    run_lab1()

import os
import torch
import torch.nn as nn
from torch.export import export, Dim

class FastClassifier(nn.Module):
    def __init__(self, in_features=512, hidden=256, num_classes=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, hidden),
            nn.LayerNorm(hidden),
            nn.ReLU(),
            nn.Linear(hidden, num_classes)
        )
    def forward(self, x):
        return self.net(x)

def run_lab2():
    print("=" * 60)
    print("LAB 2: AOTINDUCTOR PYTHON-FREE C++ COMPILATION")
    print("=" * 60)
    
    model = FastClassifier().eval()
    example_input = torch.randn(4, 512)
    batch = Dim("batch", min=1, max=64)
    
    # 1. Export graph
    print("[*] Exporting model graph...")
    exported = export(model, (example_input,), dynamic_shapes={"x": (batch, 512)})
    
    # 2. Ahead-of-Time Compile with Inductor
    output_package = "classifier_aoti.pt2"
    print(f"[*] Invoking AOTInductor compiler -> {output_package}...")
    
    try:
        from torch._inductor import aoti_compile_and_package
        compiled_package_path = aoti_compile_and_package(
            exported,
            package_path=output_package
        )
        print(f"[+] Successfully generated standalone package: {compiled_package_path}")
        print(f"[+] Package Size: {os.path.getsize(compiled_package_path) / 1024:.2f} KB")
        
        # 3. Load and execute the compiled package via Python C-Runner bindings
        from torch._inductor.aoti_runner import AOTIModelRunner
        runner = AOTIModelRunner(compiled_package_path)
        
        test_input = torch.randn(16, 512)
        out = runner(test_input)
        print(f"[+] Native AOTInductor Inference Successful! Output Shape: {list(out.shape)}")
        
    except Exception as e:
        print(f"[!] AOTInductor Note: Requires native C++ compiler toolchain (g++/clang/MSVC). Details: {e}")

if __name__ == "__main__":
    run_lab2()

import os
import time
import torch
import torch.nn as nn

class ConvFeatureExtractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.SiLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.SiLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(64, 128)
        )
    def forward(self, x):
        return self.features(x)

def run_lab3():
    print("=" * 60)
    print("LAB 3: MODERN ONNX DYNAMO EXPORT & ORT BENCHMARK")
    print("=" * 60)
    
    model = ConvFeatureExtractor().eval()
    dummy_input = torch.randn(1, 3, 64, 64)
    onnx_path = "conv_model_dynamo.onnx"
    
    # 1. Export with TorchDynamo Backend
    print(f"[*] Exporting model via torch.onnx.export(dynamo=True) -> {onnx_path}...")
    torch.onnx.export(
        model,
        args=(dummy_input,),
        f=onnx_path,
        opset_version=18,
        dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
        input_names=["input"],
        output_names=["output"]
    )
    print(f"[+] Model saved. File Size: {os.path.getsize(onnx_path) / 1024:.2f} KB")
    
    # 2. Benchmark with ONNX Runtime
    try:
        import onnxruntime as ort
        print("
[*] Initializing ONNX Runtime Session (CPU Execution Provider)...")
        session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
        
        # Test input with dynamic batch = 16
        input_numpy = torch.randn(16, 3, 64, 64).numpy()
        
        # Warmup
        for _ in range(10):
            _ = session.run(["output"], {"input": input_numpy})
            
        # Timed benchmark
        iterations = 50
        start = time.perf_counter()
        for _ in range(iterations):
            _ = session.run(["output"], {"input": input_numpy})
        elapsed = time.perf_counter() - start
        
        fps = (iterations * 16) / elapsed
        print(f"[+] ORT Benchmark Completed: {elapsed*1000/iterations:.2f} ms/batch (16 items)")
        print(f"[+] Throughput: {fps:.2f} images/second")
        
    except ImportError:
        print("[!] Note: Install 'onnxruntime' via pip to run the benchmarking suite.")

if __name__ == "__main__":
    run_lab3()

import torch
import torch.nn as nn
from torch.export import export

class MobileEdgeNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, 3, padding=1)
        self.act1 = nn.ReLU()
        self.conv2 = nn.Conv2d(16, 16, 3, padding=1)
        self.act2 = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 4)

    def forward(self, x):
        h1 = self.act1(self.conv1(x))
        h2 = self.act2(self.conv2(h1))
        return self.fc(self.pool(h2).flatten(1))

def run_lab4():
    print("=" * 60)
    print("LAB 4: EXECUTORCH EDGE LOWERING & MEMORY PLANNING")
    print("=" * 60)
    
    model = MobileEdgeNet().eval()
    example_x = torch.randn(1, 1, 28, 28)
    
    # 1. Step 1: Standard Sound Export
    print("[*] Stage 1: Capturing functional graph with torch.export...")
    exported_prog = export(model, (example_x,))
    
    # 2. Step 2: Simulate EXIR (Edge Dialect Lowering) & Liveness Memory Planning
    print("[*] Stage 2: Simulating ExecuTorch Memory Planning Pass...")
    
    tensor_lifespans = []
    # Calculate intermediate tensor sizes
    for node in exported_prog.graph_module.graph.nodes:
        if node.op == "call_function":
            # Estimate activation buffer size
            tensor_lifespans.append((node.name, 16 * 28 * 28 * 4)) # ~50 KB per activation
            
    print(f"[+] Total intermediate operators: {len(tensor_lifespans)}")
    unplanned_peak_mb = sum(size for _, size in tensor_lifespans) / (1024 * 1024)
    # With non-overlapping buffer reuse:
    planned_arena_mb = max(size for _, size in tensor_lifespans) * 2 / (1024 * 1024)
    
    print(f"[+] Unplanned Dynamic Heap Allocation: {unplanned_peak_mb:.2f} MB")
    print(f"[+] ExecuTorch Pre-Planned Static Arena: {planned_arena_mb:.2f} MB (Zero malloc at runtime!)")
    print("[+] Model ready for serialization into zero-copy .pte container.")

if __name__ == "__main__":
    run_lab4()

import time
import torch
import torch.nn as nn
import torch.ao.quantization as quantization

class DenseBenchmarkModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = quantization.QuantStub()
        self.fc1 = nn.Linear(1024, 512)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(512, 256)
        self.dequant = quantization.DeQuantStub()
        
    def forward(self, x):
        x = self.quant(x)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return self.dequant(x)

def run_lab5():
    print("=" * 60)
    print("LAB 5: POST-TRAINING INT8 QUANTIZATION (PTQ)")
    print("=" * 60)
    
    # 1. Instantiate baseline FP32 model
    model_fp32 = DenseBenchmarkModel().eval()
    
    # 2. Configure QConfig (X86 / ARM int8 engine)
    model_fp32.qconfig = quantization.get_default_qconfig("fbgemm")
    model_prepared = quantization.prepare(model_fp32)
    
    # 3. Calibration phase: Feed representative sample batches
    print("[*] Calibrating activation observers with representative dataset...")
    with torch.no_grad():
        for _ in range(50):
            calib_batch = torch.randn(32, 1024)
            model_prepared(calib_batch)
            
    # 4. Convert to INT8 Quantized Model
    print("[*] Converting prepared model to INT8 quantized representation...")
    model_int8 = quantization.convert(model_prepared)
    
    # 5. Measure model sizes in memory
    def get_model_size_kb(model):
        torch.save(model.state_dict(), "temp.p")
        size_kb = os.path.getsize("temp.p") / 1024
        if os.path.exists("temp.p"):
            os.remove("temp.p")
        return size_kb
        
    import os
    size_fp32 = get_model_size_kb(model_fp32)
    size_int8 = get_model_size_kb(model_int8)
    
    print(f"
[+] Baseline FP32 Model Size: {size_fp32:.2f} KB")
    print(f"[+] Quantized INT8 Model Size: {size_int8:.2f} KB")
    print(f"[+] Memory Footprint Compression: {size_fp32 / size_int8:.2f}x reduction!")
    
    # 6. Benchmark Inference Latency
    test_batch = torch.randn(64, 1024)
    
    # FP32 timing
    start = time.perf_counter()
    for _ in range(100):
        _ = model_fp32(test_batch)
    time_fp32 = (time.perf_counter() - start) * 10
    
    # INT8 timing
    start = time.perf_counter()
    for _ in range(100):
        _ = model_int8(test_batch)
    time_int8 = (time.perf_counter() - start) * 10
    
    print(f"[+] FP32 Latency (100 runs): {time_fp32:.2f} ms")
    print(f"[+] INT8 Latency (100 runs): {time_int8:.2f} ms")
    print(f"[+] Execution Speedup: {time_fp32 / time_int8:.2f}x faster on CPU!")

if __name__ == "__main__":
    run_lab5()