parallel_inference.py

#!/usr/bin/env python3
"""
Parallel Inference via Linear Task Clustering
Based on: "Automatic Task Parallelization of Dataflow Graphs in ML/DL Models"
Algorithms cross-referenced with the Ramiel reference implementation.

Steps:
  1. Import model (GoogLeNet, ResNet-50, Inception-v3)
  2. Convert model to graph form via torch.fx
  3. Compute distance_to_end per node (FindDistanceToEnd / Algorithm 1)
  4. Linear clustering                 (Algorithm 2)
  5. Iterative cluster merging          (Algorithms 3 & 4)
     - Merge clusters whose spans do NOT overlap (confirmed from Ramiel source)
  6. Parallel execution via Python processes + mp.Queue
     - One process per merged cluster
     - Tensors sent as numpy arrays via queue.put() / queue.get(block=True)
  7. Measure speedup, scalability, execution efficiency vs sequential

Run:
  python3 parallel_inference.py
"""

import time
from collections import defaultdict, deque
import torch
import torch.nn as nn
import torchvision.models as models
import torch.fx as fx
import torch.multiprocessing as mp

mp.set_start_method("fork", force=True)


# ─────────────────────────────────────────────
# STEP 1: Import Models
# ─────────────────────────────────────────────

def load_model(name):
    if name == "googlenet":
        return models.googlenet(weights=None, aux_logits=False)
    elif name == "resnet50":
        return models.resnet50(weights=None)
    elif name == "inception_v3":
        return models.inception_v3(weights=None, aux_logits=False)
    elif name == "squeezenet":
        return models.squeezenet1_0(weights=None)
    elif name == "densenet121":
        return models.densenet121(weights=None)
    else:
        raise ValueError(f"Unknown model: {name}")


# ─────────────────────────────────────────────
# STEP 2: Build Graph via torch.fx
# ─────────────────────────────────────────────

def get_node_cost(node, model):
    """Base node cost per paper's operator cost table."""
    if node.op == "call_module":
        layer = model
        for part in node.target.split("."):
            layer = getattr(layer, part, None)
            if layer is None:
                return 1
        if isinstance(layer, nn.Linear):
            return 50
        if isinstance(layer, (nn.AdaptiveAvgPool2d, nn.AvgPool2d)):
            return 15
        if isinstance(layer, nn.Conv2d):
            w = layer.kernel_size[0] if isinstance(layer.kernel_size, tuple) else layer.kernel_size
            return 10 * w
        if isinstance(layer, nn.MaxPool2d):
            w = layer.kernel_size if isinstance(layer.kernel_size, int) else layer.kernel_size[0]
            return 5 * w
    if node.op == "call_function":
        if node.target in (torch.cat,) or "add" in str(node.target):
            return 15
    return 1


def build_graph(model, input_shape):
    """Trace model with torch.fx, capture tensor sizes via hooks."""
    model.eval()
    traced   = fx.symbolic_trace(model)
    fx_nodes = list(traced.graph.nodes)

    tensor_sizes = {}
    def make_hook(name):
        def hook(module, inp, output):
            if isinstance(output, torch.Tensor):
                tensor_sizes[name] = output.numel()
        return hook

    hooks = []
    for node in fx_nodes:
        if node.op == "call_module":
            submod, valid = model, True
            for part in node.target.split("."):
                submod = getattr(submod, part, None)
                if submod is None:
                    valid = False
                    break
            if valid:
                hooks.append(submod.register_forward_hook(make_hook(node.name)))

    with torch.no_grad():
        model(torch.zeros(1, *input_shape))
    for h in hooks:
        h.remove()

    placeholder_name = next(n.name for n in fx_nodes if n.op == "placeholder")

    graph = {}
    for node in fx_nodes:
        if node.op in ("placeholder", "output"):
            continue
        deps = []
        for arg in node.args:
            if isinstance(arg, fx.Node):
                deps.append(arg.name)
            elif isinstance(arg, (list, tuple)):
                for item in arg:
                    if isinstance(item, fx.Node):
                        deps.append(item.name)

        t_sz = tensor_sizes.get(node.name, 1)
        graph[node.name] = {
            "node":       node,
            "deps":       deps,
            "successors": [],
            "nodecost":   get_node_cost(node, model) * t_sz,
            "edgecost":   1 * t_sz,
        }

    for n in graph:
        for dep in graph[n]["deps"]:
            if dep in graph:
                graph[dep]["successors"].append(n)

    return graph, placeholder_name


# ─────────────────────────────────────────────
# STEP 3: FindDistanceToEnd
# ─────────────────────────────────────────────

def find_distance_to_end(graph):
    """
    Iterative distance-to-end using reverse topological order.
    Avoids recursion depth issues for deep models like ResNet-152.
    """
    # Topological sort
    in_degree = defaultdict(int)
    for n in graph:
        for dep in graph[n]["deps"]:
            if dep in graph:
                in_degree[n] += 1

    queue = deque([n for n in graph if in_degree[n] == 0])
    topo  = []
    while queue:
        node = queue.popleft()
        topo.append(node)
        for succ in graph[node]["successors"]:
            in_degree[succ] -= 1
            if in_degree[succ] == 0:
                queue.append(succ)

    # Process in reverse topological order (leaves first)
    dte = {n: 0 for n in graph}
    for n in reversed(topo):
        succs = [s for s in graph[n]["successors"] if s in graph]
        if succs:
            dte[n] = graph[n]["nodecost"] + max(dte[s] for s in succs) + graph[n]["edgecost"]

    return dte


# ─────────────────────────────────────────────
# STEP 4: Linear Clustering
# ─────────────────────────────────────────────

def linear_clustering(graph, dte):
    """
    Algorithm 2: Recursive Linear Clustering.
    Repeatedly extracts the critical path as a cluster, pruning edges
    as per the pseudocode until all nodes are clustered.
    """
    out_edges = {n: set(graph[n]["successors"]) for n in graph}
    in_edges  = {n: set(graph[n]["deps"]) & graph.keys() for n in graph}
    remaining = set(graph.keys())
    all_clusters = []

    while remaining:
        ready_l = [n for n in remaining if len(in_edges[n]) == 0] or list(remaining)
        c_node  = max(ready_l, key=lambda n: dte.get(n, 0))
        new_cluster = [c_node]
        remaining.remove(c_node)

        while out_edges[c_node]:
            s_node = max(out_edges[c_node], key=lambda n: dte.get(n, 0))
            new_cluster.append(s_node)
            if s_node in remaining:
                remaining.remove(s_node)
            for other in list(out_edges[c_node]):
                if other != s_node:
                    out_edges[c_node].discard(other)
                    in_edges[other].discard(c_node)
            for other in list(in_edges[s_node]):
                if other != c_node:
                    in_edges[s_node].discard(other)
                    out_edges[other].discard(s_node)
            out_edges[c_node].discard(s_node)
            in_edges[s_node].discard(c_node)
            c_node = s_node

        all_clusters.append(new_cluster)
    return all_clusters


# ─────────────────────────────────────────────
# STEP 5: Cluster Merging
# ─────────────────────────────────────────────

def clusters_overlap(cl1, cl2, dte):
    """
    Paper Algorithm 3: clusters overlap if
      sSpan(cl1) < eSpan(cl2) OR sSpan(cl2) < eSpan(cl1)
    sSpan = dte[entry node], eSpan = dte[exit node]
    """
    ss1 = dte.get(cl1[0], 0)
    es1 = dte.get(cl1[-1], 0)
    ss2 = dte.get(cl2[0], 0)
    es2 = dte.get(cl2[-1], 0)
    return (ss1 < es2) or (ss2 < es1)


def merge_clusters_pass(clusters, dte):
    """Algorithm 3: merge pairs of clusters that overlap."""
    merged, skip, done = [], set(), False
    for i, cl1 in enumerate(clusters):
        if i in skip:
            continue
        merged_flag = False
        for j, cl2 in enumerate(clusters):
            if i == j or j in skip:
                continue
            if clusters_overlap(cl1, cl2, dte):
                merged.append(cl1 + cl2)
                skip.update([i, j])
                done = True
                merged_flag = True
                break
        if not merged_flag and i not in skip:
            merged.append(cl1)
    return merged, done


def iterative_cluster_merging(clusters, dte):
    """Algorithm 4: repeat merge passes until stable."""
    current, done = clusters, True
    while done:
        current, done = merge_clusters_pass(current, dte)
    return current


# ─────────────────────────────────────────────
# STEP 6: Code Generation + Parallel Execution
# Generates a Python source file with one function per cluster,
# then forks processes to run them — matching Ramiel's approach.
# ─────────────────────────────────────────────

def profile_tensor_shapes(graph, model, node_to_cluster, placeholder_name, x, sends):
    """Run one sequential forward pass in topo order to get shapes of inter-cluster tensors."""
    # Topological sort
    in_degree = defaultdict(int)
    for n in graph:
        for dep in graph[n]["deps"]:
            if dep in graph:
                in_degree[n] += 1
    queue = deque([n for n in graph if in_degree[n] == 0])
    topo  = []
    while queue:
        node = queue.popleft()
        topo.append(node)
        for succ in graph[node]["successors"]:
            in_degree[succ] -= 1
            if in_degree[succ] == 0:
                queue.append(succ)

    outputs = {placeholder_name: x}
    shapes  = {}

    for node_name in topo:
        node = graph[node_name]["node"]

        def resolve(arg):
            if isinstance(arg, fx.Node):
                return outputs.get(arg.name)
            if isinstance(arg, (list, tuple)):
                resolved = [resolve(a) for a in arg]
                return type(arg)(resolved)
            return arg

        args   = tuple(resolve(a) for a in node.args)
        kwargs = {k: resolve(v) for k, v in node.kwargs.items()}

        # Skip if any required arg is None
        if any(a is None for a in args if isinstance(a, torch.Tensor)):
            continue

        try:
            with torch.no_grad():
                if node.op == "call_module":
                    submod = model
                    for part in node.target.split("."):
                        submod = getattr(submod, part)
                    out = submod(*args, **kwargs)
                elif node.op == "call_function":
                    out = node.target(*args, **kwargs)
                elif node.op == "call_method":
                    out = getattr(args[0], node.target)(*args[1:], **kwargs)
                else:
                    continue
        except Exception:
            continue

        outputs[node_name] = out
        if node_name in sends and isinstance(out, torch.Tensor):
            shapes[node_name] = (tuple(out.shape), out.dtype)

    return shapes


def generate_cluster_code(clusters, graph, node_to_cluster, placeholder_name):
    """
    Algorithm 7: generate Python source code for each cluster.
    Uses pre-allocated shared tensor buffers for zero-copy inter-cluster transfer.
    Only a signal integer is sent through the queue.
    """
    sends = defaultdict(list)
    recvs = {}

    for node_name in graph:
        ci = node_to_cluster[node_name]
        for dep in graph[node_name]["deps"]:
            cj = node_to_cluster.get(dep)
            if cj is not None and cj != ci:
                if dep not in recvs:
                    recvs[dep] = cj
                if ci not in sends[dep]:
                    sends[dep].append(ci)

    def arg_to_str(arg):
        if isinstance(arg, fx.Node):
            return f"local['{arg.name}']"
        if isinstance(arg, (list, tuple)):
            inner = ", ".join(arg_to_str(a) for a in arg)
            return f"[{inner}]"
        return repr(arg)

    lines = []
    lines.append("import torch")
    lines.append("import torch.nn.functional as F")
    lines.append("")

    for ci, cluster_nodes in enumerate(clusters):
        lines.append(f"def fparcluster_{ci}(model, x, sig_queues, shared_bufs, result_queue):")
        lines.append(f"    torch.set_num_threads(1)")
        lines.append(f"    model.eval()")
        lines.append(f"    local = {{'{placeholder_name}': x}}")
        lines.append(f"    with torch.no_grad():")

        for node_name in cluster_nodes:
            node = graph[node_name]["node"]

            # Recv: wait for signal, then read from shared buffer
            for dep in graph[node_name]["deps"]:
                if dep in recvs and node_to_cluster.get(dep) != ci:
                    src = recvs[dep]
                    lines.append(f"        if '{dep}' not in local:")
                    lines.append(f"            sig_queues[({src},{ci},'{dep}')].get(block=True)")
                    lines.append(f"            local['{dep}'] = shared_bufs['{dep}'].clone()")

            # Emit the op
            args_str   = ", ".join(arg_to_str(a) for a in node.args)
            kwargs_str = ", ".join(f"{k}={arg_to_str(v)}" for k, v in node.kwargs.items())
            all_args   = ", ".join(filter(None, [args_str, kwargs_str]))

            if node.op == "call_module":
                parts = node.target.split(".")
                path  = "model"
                for part in parts:
                    path = f"getattr({path}, '{part}')"
                lines.append(f"        local['{node_name}'] = {path}({all_args})")
            elif node.op == "call_function":
                fn = f"{node.target.__module__}.{node.target.__name__}"
                lines.append(f"        local['{node_name}'] = {fn}({all_args})")
            elif node.op == "call_method":
                obj  = arg_to_str(node.args[0])
                rest = ", ".join(arg_to_str(a) for a in node.args[1:])
                lines.append(f"        local['{node_name}'] = {obj}.{node.target}({rest})")

            # Send: copy into shared buffer, then signal
            if node_name in sends:
                for dst in sends[node_name]:
                    lines.append(f"        shared_bufs['{node_name}'].copy_(local['{node_name}'])")
                    lines.append(f"        sig_queues[({ci},{dst},'{node_name}')].put(1)")

        last = cluster_nodes[-1]
        lines.append(f"    result_queue.put(({ci}, local['{last}'].detach().contiguous().numpy()))")
        lines.append("")

    return "\n".join(lines), sends, recvs


def setup_parallel(clusters, graph, model, node_to_cluster, placeholder_name, x):
    src, sends, recvs = generate_cluster_code(
        clusters, graph, node_to_cluster, placeholder_name)

    import _operator
    import torchvision
    ns = {"torch": torch, "F": torch.nn.functional, "_operator": _operator, "torchvision": torchvision}
    exec(compile(src, "<generated_clusters>", "exec"), ns)

    # Profile tensor shapes for shared buffer allocation
    shapes = profile_tensor_shapes(
        graph, model, node_to_cluster, placeholder_name, x, sends)

    # Allocate shared memory tensors BEFORE forking
    shared_bufs = {}
    for node_name, dst_list in sends.items():
        if node_name in shapes:
            shape, dtype = shapes[node_name]
            buf = torch.zeros(shape, dtype=dtype).share_memory_()
            shared_bufs[node_name] = buf

    # Create signal queues (carry only int signals, not tensors)
    sig_queues = {}
    for node_name, dst_list in sends.items():
        src_ci = recvs[node_name]
        for dst_ci in dst_list:
            sig_queues[(src_ci, dst_ci, node_name)] = mp.Queue()

    result_queue  = mp.Queue()
    trigger_queues = [mp.Queue() for _ in range(len(clusters))]

    def run_cluster_persistent(ci, trigger_q):
        torch.set_num_threads(1)
        fn = ns[f"fparcluster_{ci}"]
        while True:
            msg = trigger_q.get()
            if msg == "STOP":
                break
            try:
                fn(model, msg, sig_queues, shared_bufs, result_queue)
            except Exception as e:
                import traceback
                traceback.print_exc()
                result_queue.put((ci, None))

    torch.set_num_threads(1)
    procs = [
        mp.Process(target=run_cluster_persistent, args=(i, trigger_queues[i]), daemon=True)
        for i in range(len(clusters))
    ]
    for p in procs:
        p.start()

    return procs, trigger_queues, result_queue, len(clusters)


def run_parallel(procs_info, x):
    """Run one inference using persistent processes."""
    procs, trigger_queues, result_queue, n_clusters = procs_info
    for tq in trigger_queues:
        tq.put(x)
    results = {}
    for _ in range(n_clusters):
        ci, arr = result_queue.get()
        if arr is not None:
            results[ci] = torch.from_numpy(arr)
    return results.get(n_clusters - 1)


def stop_parallel(procs_info):
    """Shut down persistent worker processes."""
    procs, trigger_queues, result_queue, _ = procs_info
    for tq in trigger_queues:
        tq.put("STOP")
    for p in procs:
        p.join()


def parallel_execute(clusters, graph, model, node_to_cluster, placeholder_name, x):
    """Single-run wrapper for timing compatibility."""
    info = setup_parallel(clusters, graph, model, node_to_cluster, placeholder_name, x)
    result = run_parallel(info, x)
    stop_parallel(info)
    return result


# ─────────────────────────────────────────────
# STEP 7: Measure & Compare
# ─────────────────────────────────────────────

def sequential_execute(model, x):
    start = time.perf_counter()
    with torch.no_grad():
        out = model(x)
    return out, time.perf_counter() - start


def benchmark(model_name, runs=10):
    print(f"\n{'='*55}")
    print(f"  Model: {model_name.upper()}")
    print(f"{'='*55}", flush=True)

    input_shape = (3, 224, 224)
    x           = torch.randn(1, *input_shape)
    model       = load_model(model_name)
    model.eval()

    graph, placeholder_name = build_graph(model, input_shape)
    dte             = find_distance_to_end(graph)
    clusters        = linear_clustering(graph, dte)
    pre_merge       = len(clusters)
    clusters        = iterative_cluster_merging(clusters, dte)
    node_to_cluster = {n: i for i, cl in enumerate(clusters) for n in cl}

    # Count inter-cluster transfers
    n_transfers = sum(
        1 for n in graph for dep in graph[n]["deps"]
        if node_to_cluster.get(dep) is not None and node_to_cluster.get(dep) != node_to_cluster[n]
    )
    print(f"  Inter-cluster transfers: {n_transfers}")

    total_work         = sum(graph[n]["nodecost"] for n in graph)
    critical_path_cost = max(dte.values()) if dte else 1
    theoretical_speedup = round(total_work / critical_path_cost, 4)

    print(f"  Nodes               : {len(graph)}")
    print(f"  Clusters (pre-merge): {pre_merge}")
    print(f"  Clusters (merged)   : {len(clusters)}")
    print(f"  Theoretical speedup : {theoretical_speedup:.3f}x")

    # Warm-up
    for _ in range(5):
        sequential_execute(model, x)

    # Sequential timing
    seq_times = []
    for _ in range(runs):
        _, t = sequential_execute(model, x)
        seq_times.append(t * 1000)  # ms

    seq_avg = sum(seq_times) / runs
    seq_std = (sum((t - seq_avg)**2 for t in seq_times) / runs) ** 0.5
    seq_min = min(seq_times)
    seq_max = max(seq_times)

    # Spawn worker processes ONCE before timing loop
    print("  Spawning workers...", flush=True)
    par_info = setup_parallel(clusters, graph, model, node_to_cluster, placeholder_name, x)
    for _ in range(5):
        run_parallel(par_info, x)
    par_times = []
    for _ in range(runs):
        t0 = time.perf_counter()
        run_parallel(par_info, x)
        par_times.append((time.perf_counter() - t0) * 1000)
    stop_parallel(par_info)

    par_avg = sum(par_times) / runs
    par_std = (sum((t - par_avg)**2 for t in par_times) / runs) ** 0.5
    par_min = min(par_times)
    par_max = max(par_times)

    speedup    = seq_avg / par_avg
    efficiency = speedup / len(clusters) * 100

    # Inter-cluster transfer count
    n_transfers = sum(
        1 for n in graph for dep in graph[n]["deps"]
        if node_to_cluster.get(dep) is not None and node_to_cluster.get(dep) != node_to_cluster[n]
    )

    print(f"\n  Sequential  : {seq_avg:.2f} ± {seq_std:.2f} ms  (min {seq_min:.2f}, max {seq_max:.2f})")
    print(f"  Parallel    : {par_avg:.2f} ± {par_std:.2f} ms  (min {par_min:.2f}, max {par_max:.2f})")
    print(f"  Speedup     : {speedup:.3f}x")
    print(f"  Theoretical : {theoretical_speedup:.3f}x")
    print(f"  Efficiency  : {efficiency:.1f}%")
    print(f"  Transfers   : {n_transfers}")
    print()

    return {
        "model":               model_name,
        "nodes":               len(graph),
        "pre_merge":           pre_merge,
        "clusters":            len(clusters),
        "n_transfers":         n_transfers,
        "seq_avg":             round(seq_avg, 3),
        "seq_std":             round(seq_std, 3),
        "seq_min":             round(seq_min, 3),
        "seq_max":             round(seq_max, 3),
        "par_avg":             round(par_avg, 3),
        "par_std":             round(par_std, 3),
        "par_min":             round(par_min, 3),
        "par_max":             round(par_max, 3),
        "speedup":             round(speedup, 4),
        "theoretical_speedup": theoretical_speedup,
        "efficiency_pct":      round(efficiency, 2),
        "seq_times":           [round(t, 3) for t in seq_times],
        "par_times":           [round(t, 3) for t in par_times],
    }


# ─────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────

if __name__ == "__main__":
    import json

    # Warmup run to prime Python imports, CPU cache, and process spawning
    print("Warming up system (throwaway run)...", flush=True)
    try:
        benchmark("resnet50", runs=5)
    except Exception:
        pass
    print("Warmup done.\n", flush=True)

    results = []
    for model_name in ["squeezenet", "googlenet", "resnet50", "densenet121", "inception_v3"]:
        try:
            r = benchmark(model_name, runs=30)
            results.append(r)
        except Exception as e:
            print(f"  ERROR benchmarking {model_name}: {e}")
            import traceback
            traceback.print_exc()

    # Save results for report generation
    with open("results.json", "w") as f:
        json.dump(results, f, indent=2)
    print("Results saved to results.json")

    print("\n-- Summary ------------------------------------------------------------------------------------------")
    print(f"{'Model':<14} {'Nodes':>6} {'Pre':>5} {'Post':>5} {'Seq(ms)':>9} {'Par(ms)':>9} {'Speedup':>9} {'Theor.':>8} {'Eff%':>7}")
    print("-" * 95)
    for r in results:
        print(f"{r['model']:<14} {r['nodes']:>6} {r['pre_merge']:>5} {r['clusters']:>5} "
              f"{r['seq_avg']:>9} {r['par_avg']:>9} {r['speedup']:>9} "
              f"{r['theoretical_speedup']:>8} {r['efficiency_pct']:>7}")