MP-SPDZ/Programs/Source/benchmark_priority_queue.mpc

from Compiler import library as lib, oram, path_oram
from Compiler.dijkstra import HeapQ
from Compiler.path_oblivious_heap import (
    POHToHeapQAdapter,
    POHVariant,
    path_oblivious_sort,
)
from Compiler.types import Array, sint
from Compiler.util import log2

DEBUG = True

from Compiler import path_oblivious_heap

path_oblivious_heap.DEBUG = False


def noop(*args, **kwargs):
    pass


# Only print if DEBUG is enabled
dprint = lib.print_ln if DEBUG else noop

# Benchmark types
INSERT = False
EXTRACT = False
SORTING = True

INSERT = INSERT or EXTRACT  # Always insert if we are going to extract

# Benchmark parameters
## General
KEY_SIZE = lambda n: 32
VALUE_SIZE = lambda n: log2(n)
N_THREADS = 1
N_PARALLEL = 1

## Insert / ExtractMin
RANGE = [2**i for i in range(1, 15)]
OPERATIONS_PER_STEP = 2
TIME_INIT = True
TREE_HEAP = False
TREE_PATH_HEAP = False
LINEAR_HEAP = False
OPTIMAL_TREE_HEAP = False
OPTIMAL_PATH_HEAP = False
POH_PATH = True
POH_PATH_CONSTANT_STASH = True

## Sorting
LENGTHS = [5, 10, 15, 20]
SORTING_BITS = 32
RADIX_SORT = True
POS = True
POS_CONSTANT_STASH = True

# Set module variables based on parameters
oram.n_threads = N_THREADS
oram.n_parallel = N_PARALLEL

# Timing with consecutive ids
timer_offset = 1000  # Hopefully run timers in an unused range


def start_fancy_timer(id: int | None = None) -> int:
    global timer_offset
    _id = id if id is not None else timer_offset
    lib.start_timer(_id)
    if id is None:
        timer_offset += 1
    return _id


def stop_fancy_timer(id):
    lib.stop_timer(id)


# BENCHMARK

if INSERT:

    def operation_round(q, apply_op, capacity, tag=""):
        global timer_offset
        id = None
        dprint(
            f"\n[{tag}] Running {OPERATIONS_PER_STEP} update{'s' if OPERATIONS_PER_STEP > 1 else ''} for capacity {capacity}"
        )
        id = timer_offset  # Handle id manually to avoid incrementing for each operation
        for i in range(OPERATIONS_PER_STEP):
            dprint(f"\n[{tag}] Update {i} for capacity {capacity}")
            start_fancy_timer(id)
            apply_op(q, i)
            stop_fancy_timer(id)
        timer_offset += 1

    def benchmark_operations(q_init, capacity, *args, tag="", **kwargs):
        apply_insert = lambda q, i: q.update(0, i)
        apply_extract = lambda q, _: q.pop()
        dprint(f"\n[{tag}] Initializing empty structure with capacity {capacity}")
        if TIME_INIT:
            id = start_fancy_timer()
        q = q_init(capacity, *args, **kwargs)
        if TIME_INIT:
            stop_fancy_timer(id)
        if INSERT:
            operation_round(q, apply_insert, capacity, tag=tag + " insert")
            if EXTRACT:
                operation_round(q, apply_extract, capacity, tag=tag + " extract_min")

    dprint(f"\n\nBENCHMARKING INSERT {'AND EXTRACT ' if EXTRACT else ''}TIME")
    for capacity in RANGE:
        entry_size = (KEY_SIZE(capacity), VALUE_SIZE(capacity))

        dprint(f"\nCAPACITY {capacity}")

        if TREE_HEAP:
            # Benchmark binary heap built on ORAM (Tree ORAM variant)
            benchmark_operations(
                HeapQ,
                capacity,
                oram_type=oram.RecursiveORAM,
                entry_size=entry_size,
                tag="ORAM Heap (Tree)",
            )

        if TREE_PATH_HEAP:
            # Benchmark binary heap built on ORAM (Path ORAM variant)
            benchmark_operations(
                HeapQ,
                capacity,
                oram_type=path_oram.RecursivePathORAM,
                entry_size=entry_size,
                tag="ORAM Heap (Path)",
            )

        if LINEAR_HEAP:
            # Benchmark binary heap built on ORAM (Linear ORAM variant)
            benchmark_operations(
                HeapQ,
                capacity,
                oram_type=oram.LinearORAM,
                entry_size=entry_size,
                tag="ORAM Heap (Linear)",
            )

        if OPTIMAL_TREE_HEAP:
            # Benchmark binary heap built on ORAM (OptimalORAM variant)
            benchmark_operations(
                HeapQ,
                capacity,
                oram_type=oram.OptimalORAM,
                entry_size=entry_size,
                tag="ORAM Heap (Optimal Tree)",
            )

        if OPTIMAL_PATH_HEAP:
            # Benchmark binary heap built on ORAM (OptimalORAM Path variant)
            benchmark_operations(
                HeapQ,
                capacity,
                oram_type=path_oram.OptimalORAM,
                entry_size=entry_size,
                tag="ORAM Heap (Optimal Path)",
            )

        if POH_PATH:
            # Benchmark Path Oblivious Heap (Path variant)
            benchmark_operations(
                POHToHeapQAdapter,
                capacity,
                bucket_size=2,
                stash_size=log2(capacity) ** 2,
                variant=POHVariant.PATH,
                entry_size=entry_size,
                tag="POH (Path (superlogarithmic stash size))",
            )

        if POH_PATH_CONSTANT_STASH:
            # Benchmark Path Oblivious Heap (Path variant with constant stash size)
            benchmark_operations(
                POHToHeapQAdapter,
                capacity,
                bucket_size=2,
                stash_size=20,  # based on empirical analysis by Keller and Scholl
                variant=POHVariant.PATH,
                entry_size=entry_size,
                tag="POH (Path (constant stash size))",
            )

if SORTING:
    dprint("\n\nBENCHMARKING SORTING TIME")
    for n in LENGTHS:
        dprint(f"\nLENGTH {n}")
        a = Array(n, sint)

        @lib.for_range(n)
        def _(i):
            a[i] = sint.get_random_int(SORTING_BITS)

        if RADIX_SORT:
            a_ = Array(n, sint).assign(a)
            lib.print_ln(
                "\n[Sorting (Radix)] Unsorted array of length %s: %s", n, a_.reveal()
            )
            lib.print_ln("[Sorting (Radix)] Sorting array...")
            id = start_fancy_timer()
            a_.sort()
            stop_fancy_timer(id)
            lib.print_ln("[Sorting (Radix)] Sorted array: %s", a_.reveal())

        if POS:
            a_ = Array(n, sint).assign(a)
            lib.print_ln(
                "\n[Sorting (POH) superlogarithmic stash size] Unsorted array of length %s: %s",
                n,
                a_.reveal(),
            )
            lib.print_ln("[Sorting (POH) superlogarithmic stash size] Sorting array...")
            id = start_fancy_timer()
            path_oblivious_sort(a_, a_, SORTING_BITS)
            stop_fancy_timer(id)
            lib.print_ln(
                "[Sorting (POH) superlogarithmic stash size] Sorted array: %s",
                a_.reveal(),
            )

        if POS_CONSTANT_STASH:
            a_ = Array(n, sint).assign(a)
            lib.print_ln(
                "\n[Sorting (POH) constant stash size] Unsorted array of length %s: %s",
                n,
                a_.reveal(),
            )
            lib.print_ln("[Sorting (POH) constant stash size] Sorting array...")
            id = start_fancy_timer()
            path_oblivious_sort(a_, a_, SORTING_BITS, stash_size=20)
            stop_fancy_timer(id)
            lib.print_ln(
                "[Sorting (POH) constant stash size] Sorted array: %s", a_.reveal()
            )