Transform rank-local collective IR

Example source: examples/24_distributed_collective_ir.py

Example 24 constructs a rank-local Program containing collective operations and applies a caller-selected Compile pass to expose a generic reduction region. It compares two valid representations of ciphertext-add all-reduce without launching distributed execution.

Run the example

python examples/24_distributed_collective_ir.py

The script requires neither a process group nor multiple devices because it only builds and transforms IR.

Rank-local Program

The function receives:

• one rank-local CKKS ciphertext;
• one launch-bound fhelium_dist.group value.

It reads rank and group size, broadcasts the ciphertext from rank zero, and applies fhelium_dist.all_reduce_add_ciphertext.

Every operation describes the work observed by one rank. A later SPMD launch supplies device count, group rank, process-group lifecycle, and transport resources.

Specialized collective form

fhelium_dist.all_reduce_add_ciphertext states that the local ciphertext values are combined by ciphertext addition. A Backend with a whole-operation implementation may consume this representation directly.

The specialized operation provides one inspectable reduction form available to passes and execution owners.

Generic combine-region form

LowerSpecializedCollectivesPass can replace the specialized operation with:

fhelium_dist.all_reduce ... ({
  ^bb0(%left: !fhelium_ckks.ciphertext,
       %right: !fhelium_ckks.ciphertext):
    %sum = fhelium_ckks.add %left, %right
    fhelium_dist.yield %sum
})

The generic operation exposes the local binary combine computation as a region, allowing later passes to inspect or transform the CKKS addition.

The example compiles the same source Program twice:

LowerSpecializedCollectivesPass(lower_ciphertext_add=False)
LowerSpecializedCollectivesPass(lower_ciphertext_add=True)

The first request records a preserve decision. The second records a generic-combine-region decision and rewrites the IR. Both choices appear in pass reports.

Protocol checks beyond Program structure

Structural construction and transformation establish local IR validity. Distributed execution additionally requires checks for:

• that all ranks enter the same collective;
• that collectives occur in the same cross-rank order;
• that region control flow is rank-uniform;
• that the combine operation is associative under the selected arithmetic;
• that execution is free of deadlock;
• that process-group resources match the Program.

A caller may compose diagnostic passes for a particular protocol and run them before linking or execution.

Why preserve both forms

The two forms serve different transformation and execution choices:

• the specialized operation gives a provider one whole-operation unit;
• the generic operation exposes the local combine arithmetic;
• a distributed pass can choose according to a concrete workload and resource plan;
• the selected representation remains visible in Program text and pass decisions.

Neither choice becomes a permanent framework-wide lowering policy.

Source

#!/usr/bin/env python3

"""Build rank-local collective IR and expose its generic combine region.

A rank-local Program receives one ciphertext and one launch-bound process
group. The source Program preserves a specialized ciphertext-add all-reduce;
a caller-selected Compile pass may instead lower it to generic all-reduce with
a visible CKKS-add combine region. Neither representation proves cross-rank
ordering, uniform control flow, associativity, or deadlock freedom.
"""

from __future__ import annotations

from common import print_table
from xdsl.dialects.func import ReturnOp
from xdsl.ir import Block

from fhelium import compile as fh_compile
from fhelium import ir
from fhelium.compile.passes.distributed import (
    LowerSpecializedCollectivesPass,
)
from fhelium.ir.dialects import ckks, distributed


def rank_local_program() -> ir.Program:
    """Return one rank-local broadcast and ciphertext-reduction Program."""

    ciphertext_type = ckks.CiphertextType()
    group_type = distributed.GroupType()
    block = Block(arg_types=(ciphertext_type, group_type))
    rank = distributed.RankOp(block.args[1])
    group_size = distributed.GroupSizeOp(block.args[1])
    broadcast = distributed.BroadcastOp(
        block.args[0],
        block.args[1],
        root=0,
    )
    reduced = distributed.AllReduceAddCiphertextOp(
        broadcast.result,
        block.args[1],
    )
    block.add_ops(
        (
            rank,
            group_size,
            broadcast,
            reduced,
            ReturnOp(reduced.result),
        )
    )
    return ir.Program.from_function(block, (ciphertext_type,))


def _decision_rows(
    label: str,
    result: fh_compile.Compilation,
) -> list[list[object]]:
    """Return printable transformation choices from one Compile result."""

    return [
        [
            label,
            decision.subject,
            decision.selected,
            decision.candidates,
            decision.details,
        ]
        for report in result.reports
        for decision in report.decisions
    ]


def main() -> None:
    source = rank_local_program()
    preserved = fh_compile.Pipeline(
        (LowerSpecializedCollectivesPass(lower_ciphertext_add=False),)
    ).run(fh_compile.Compilation(source))
    lowered = fh_compile.Pipeline((LowerSpecializedCollectivesPass(),)).run(
        fh_compile.Compilation(source)
    )

    source_inventory = ir.inventory_program(source)
    preserved_inventory = ir.inventory_program(preserved.program)
    lowered_inventory = ir.inventory_program(lowered.program)
    rows = _decision_rows("preserve", preserved) + _decision_rows(
        "lower", lowered
    )

    print_table(
        ["Program", "operations", "collective form"],
        [
            [
                "source",
                sum(source_inventory.operation_counts.values()),
                "fhelium_dist.all_reduce_add_ciphertext",
            ],
            [
                "preserved",
                sum(preserved_inventory.operation_counts.values()),
                "fhelium_dist.all_reduce_add_ciphertext",
            ],
            [
                "lowered",
                sum(lowered_inventory.operation_counts.values()),
                "fhelium_dist.all_reduce + visible fhelium_ckks.add",
            ],
        ],
    )
    print()
    print_table(
        ["request", "subject", "selected", "candidates", "details"],
        rows,
    )
    print()
    print("--- specialized rank-local Program ---")
    print(ir.format_program(source))
    print()
    print("--- generic combine-region Program ---")
    print(ir.format_program(lowered.program))


if __name__ == "__main__":
    main()

Next steps

• SPMD execution model defines rank-local ownership and launch responsibilities.
• Communication semantics distinguishes typed ciphertext reduction from residue-row reconstruction.
• Neutral IR Programs explains regions and permissive mixed-level structure.