Desymref

desymref

Desymrefier dataclass

Rewrites the program by removing all reads/writes from/to symbols and symbol definitions.

Source code in xdsl/transforms/desymref.py

@dataclass
class Desymrefier:
    """
    Rewrites the program by removing all reads/writes from/to symbols and symbol
    definitions.
    """

    def desymrefy(self, op: Operation):
        """
        Desymrefy an operation. This method guarantees that the operation does
        not have any symbols.
        """
        self.prepare_op(op)
        self.promote_op(op)

    def promote_op(self, op: Operation):
        """
        Promotes an operation. This method guarantees that the operation does
        not have any symbols.
        """
        # TODO: Add promoters in the next patch.
        pass

    def prepare_op(self, op: Operation):
        """
        Prepares an operation for promotion. This method guarantees that any
        symbol in any region of this operation is read at most once and written
        at most once.
        """

        # For operation with no regions we don't have to do any work.
        if len(op.regions) == 0:
            return

        # Otherwise, we have to prepare regions.
        for region in op.regions:
            self.prepare_region(region)

    def prepare_region(self, region: Region):
        """
        Prepares a region for promotion. This method guarantees that any symbol
        in the region is read at most once and written at most once.
        """
        num_blocks = len(region.blocks)
        if num_blocks == 1:
            # If there is only one block, preparing region is easier, so we
            # handle it separately.
            self.prepare_block(region.block)
        else:
            # TODO: Support regions with multiple blocks. This is not too
            # difficult but requires many more things:
            # 1. SSA-construction algorithm to prune symbol declarations. This
            #    also requires analyses passes (dominators).
            # 2. Insertion of entry/exit blocks to ensure dominance.
            raise FrontendProgramException(
                f"Running desymrefier on region with {num_blocks} > 1 blocks is "
                "not supported."
            )

    def prepare_block(self, block: Block):
        """Prepares a block for promotion."""

        # First, desymrefy nested regions.
        for op in block.ops:
            self.desymrefy(op)

        self.prune_definitions(block)
        self.prune_uses_without_definitions(block)

        # Ensure that each symbol is read/written at most once.
        symbols = get_symbols(block)
        for symbol in symbols:
            num_reads = sum(
                isinstance(op, symref.FetchOp)
                for op in block.ops
                if get_symbol(op) == symbol
            )
            num_writes = sum(
                isinstance(op, symref.UpdateOp)
                for op in block.ops
                if get_symbol(op) == symbol
            )
            if num_reads > 1 or num_writes > 1:
                raise FrontendProgramException(
                    f"Block {block} not ready for promotion: found {num_reads}"
                    f" reads and {num_writes} writes."
                )

    def prune_definitions(self, block: Block):
        """Removes all symbol definitions and their uses from the block."""

        # Find all symbol definitions in this block. If no definitions
        # found, terminate.
        while (
            len(
                definitions := [
                    op for op in block.ops if isinstance(op, symref.DeclareOp)
                ]
            )
            > 0
        ):
            # Otherwise, some definitions are still alive.
            for definition in definitions:
                symbol = get_symbol(definition)

                # Find all reads and writes for this symbol.
                reads = [
                    op
                    for op in block.ops
                    if isinstance(op, symref.FetchOp) and get_symbol(op) == symbol
                ]
                writes = [
                    op
                    for op in block.ops
                    if isinstance(op, symref.UpdateOp) and get_symbol(op) == symbol
                ]

                # Symbol is never read, so remove its definition and any writes.
                if len(reads) == 0:
                    for write in writes:
                        Rewriter.erase_op(write)
                    Rewriter.erase_op(definition)
                    continue

                # For symbols which are written once, the write dominates all
                # the uses and therefore can be trivially replaced.
                if len(writes) == 1:
                    write = writes[0]
                    for read in reads:
                        Rewriter.replace_op(read, [], [write.operands[0]])
                    Rewriter.erase_op(write)
                    Rewriter.erase_op(definition)
                    continue

                # If there are multiple reads and writes, replace every
                # read with the closest preceding write.
                for read in reads:
                    write = lower_positional_bound(writes, read)
                    if write is not None:
                        Rewriter.replace_op(read, [], [write.operands[0]])

    def _prune_unused_reads(self, block: Block):
        def is_unused_read(op: Operation) -> bool:
            return isinstance(op, symref.FetchOp) and not op.results[0].uses

        unused_reads = [op for op in block.ops if is_unused_read(op)]
        for read in unused_reads:
            Rewriter.erase_op(read)

    def prune_uses_without_definitions(self, block: Block):
        """Removes all possible symbol uses in a single block."""
        prepared_symbols: set[str] = set()

        while True:
            self._prune_unused_reads(block)

            # Find all symbols that are still in use in this block.
            symbol_worklist: set[str] = {
                symbol
                for symbol in get_symbols(block)
                if symbol not in prepared_symbols
            }
            if len(symbol_worklist) == 0:
                return

            for symbol in symbol_worklist:
                reads = [
                    op
                    for op in block.ops
                    if isinstance(op, symref.FetchOp) and get_symbol(op) == symbol
                ]
                writes = [
                    op
                    for op in block.ops
                    if isinstance(op, symref.UpdateOp) and get_symbol(op) == symbol
                ]
                assert len(reads) > 0 or len(writes) > 0

                # There are no reads, so we can only keep the last write to the
                # symbol.
                if len(reads) == 0:
                    for write in writes[:-1]:
                        Rewriter.erase_op(write)
                    prepared_symbols.add(symbol)
                    continue

                # There are no writes, so we can replace all reads with this
                # symbol.
                if len(writes) == 0:
                    for read in reads[1:]:
                        Rewriter.replace_op(read, [], [reads[0].results[0]])
                    prepared_symbols.add(symbol)
                    continue

                # sets of reads and writes are disjoint.
                last_read_idx = block.get_operation_index(reads[-1])
                first_write_idx = block.get_operation_index(writes[0])
                if last_read_idx < first_write_idx:
                    for read in reads[1:]:
                        Rewriter.replace_op(read, [], [reads[0].results[0]])
                    for write in writes[:-1]:
                        Rewriter.erase_op(write)
                    prepared_symbols.add(symbol)
                    continue

                # Otherwise, replace reads with the closest preceding write.
                for read in reads:
                    write = lower_positional_bound(writes, read)
                    if write is not None:
                        Rewriter.replace_op(read, [], [write.operands[0]])

__init__() -> None

desymrefy(op: Operation)

Desymrefy an operation. This method guarantees that the operation does not have any symbols.

Source code in xdsl/transforms/desymref.py

def desymrefy(self, op: Operation):
    """
    Desymrefy an operation. This method guarantees that the operation does
    not have any symbols.
    """
    self.prepare_op(op)
    self.promote_op(op)

promote_op(op: Operation)

Promotes an operation. This method guarantees that the operation does not have any symbols.

Source code in xdsl/transforms/desymref.py

def promote_op(self, op: Operation):
    """
    Promotes an operation. This method guarantees that the operation does
    not have any symbols.
    """
    # TODO: Add promoters in the next patch.
    pass

prepare_op(op: Operation)

Prepares an operation for promotion. This method guarantees that any symbol in any region of this operation is read at most once and written at most once.

Source code in xdsl/transforms/desymref.py

def prepare_op(self, op: Operation):
    """
    Prepares an operation for promotion. This method guarantees that any
    symbol in any region of this operation is read at most once and written
    at most once.
    """

    # For operation with no regions we don't have to do any work.
    if len(op.regions) == 0:
        return

    # Otherwise, we have to prepare regions.
    for region in op.regions:
        self.prepare_region(region)

prepare_region(region: Region)

Prepares a region for promotion. This method guarantees that any symbol in the region is read at most once and written at most once.

Source code in xdsl/transforms/desymref.py

def prepare_region(self, region: Region):
    """
    Prepares a region for promotion. This method guarantees that any symbol
    in the region is read at most once and written at most once.
    """
    num_blocks = len(region.blocks)
    if num_blocks == 1:
        # If there is only one block, preparing region is easier, so we
        # handle it separately.
        self.prepare_block(region.block)
    else:
        # TODO: Support regions with multiple blocks. This is not too
        # difficult but requires many more things:
        # 1. SSA-construction algorithm to prune symbol declarations. This
        #    also requires analyses passes (dominators).
        # 2. Insertion of entry/exit blocks to ensure dominance.
        raise FrontendProgramException(
            f"Running desymrefier on region with {num_blocks} > 1 blocks is "
            "not supported."
        )

prepare_block(block: Block)

Prepares a block for promotion.

Source code in xdsl/transforms/desymref.py

def prepare_block(self, block: Block):
    """Prepares a block for promotion."""

    # First, desymrefy nested regions.
    for op in block.ops:
        self.desymrefy(op)

    self.prune_definitions(block)
    self.prune_uses_without_definitions(block)

    # Ensure that each symbol is read/written at most once.
    symbols = get_symbols(block)
    for symbol in symbols:
        num_reads = sum(
            isinstance(op, symref.FetchOp)
            for op in block.ops
            if get_symbol(op) == symbol
        )
        num_writes = sum(
            isinstance(op, symref.UpdateOp)
            for op in block.ops
            if get_symbol(op) == symbol
        )
        if num_reads > 1 or num_writes > 1:
            raise FrontendProgramException(
                f"Block {block} not ready for promotion: found {num_reads}"
                f" reads and {num_writes} writes."
            )

prune_definitions(block: Block)

Removes all symbol definitions and their uses from the block.

Source code in xdsl/transforms/desymref.py

def prune_definitions(self, block: Block):
    """Removes all symbol definitions and their uses from the block."""

    # Find all symbol definitions in this block. If no definitions
    # found, terminate.
    while (
        len(
            definitions := [
                op for op in block.ops if isinstance(op, symref.DeclareOp)
            ]
        )
        > 0
    ):
        # Otherwise, some definitions are still alive.
        for definition in definitions:
            symbol = get_symbol(definition)

            # Find all reads and writes for this symbol.
            reads = [
                op
                for op in block.ops
                if isinstance(op, symref.FetchOp) and get_symbol(op) == symbol
            ]
            writes = [
                op
                for op in block.ops
                if isinstance(op, symref.UpdateOp) and get_symbol(op) == symbol
            ]

            # Symbol is never read, so remove its definition and any writes.
            if len(reads) == 0:
                for write in writes:
                    Rewriter.erase_op(write)
                Rewriter.erase_op(definition)
                continue

            # For symbols which are written once, the write dominates all
            # the uses and therefore can be trivially replaced.
            if len(writes) == 1:
                write = writes[0]
                for read in reads:
                    Rewriter.replace_op(read, [], [write.operands[0]])
                Rewriter.erase_op(write)
                Rewriter.erase_op(definition)
                continue

            # If there are multiple reads and writes, replace every
            # read with the closest preceding write.
            for read in reads:
                write = lower_positional_bound(writes, read)
                if write is not None:
                    Rewriter.replace_op(read, [], [write.operands[0]])

prune_uses_without_definitions(block: Block)

Removes all possible symbol uses in a single block.

Source code in xdsl/transforms/desymref.py

def prune_uses_without_definitions(self, block: Block):
    """Removes all possible symbol uses in a single block."""
    prepared_symbols: set[str] = set()

    while True:
        self._prune_unused_reads(block)

        # Find all symbols that are still in use in this block.
        symbol_worklist: set[str] = {
            symbol
            for symbol in get_symbols(block)
            if symbol not in prepared_symbols
        }
        if len(symbol_worklist) == 0:
            return

        for symbol in symbol_worklist:
            reads = [
                op
                for op in block.ops
                if isinstance(op, symref.FetchOp) and get_symbol(op) == symbol
            ]
            writes = [
                op
                for op in block.ops
                if isinstance(op, symref.UpdateOp) and get_symbol(op) == symbol
            ]
            assert len(reads) > 0 or len(writes) > 0

            # There are no reads, so we can only keep the last write to the
            # symbol.
            if len(reads) == 0:
                for write in writes[:-1]:
                    Rewriter.erase_op(write)
                prepared_symbols.add(symbol)
                continue

            # There are no writes, so we can replace all reads with this
            # symbol.
            if len(writes) == 0:
                for read in reads[1:]:
                    Rewriter.replace_op(read, [], [reads[0].results[0]])
                prepared_symbols.add(symbol)
                continue

            # sets of reads and writes are disjoint.
            last_read_idx = block.get_operation_index(reads[-1])
            first_write_idx = block.get_operation_index(writes[0])
            if last_read_idx < first_write_idx:
                for read in reads[1:]:
                    Rewriter.replace_op(read, [], [reads[0].results[0]])
                for write in writes[:-1]:
                    Rewriter.erase_op(write)
                prepared_symbols.add(symbol)
                continue

            # Otherwise, replace reads with the closest preceding write.
            for read in reads:
                write = lower_positional_bound(writes, read)
                if write is not None:
                    Rewriter.replace_op(read, [], [write.operands[0]])

FrontendDesymrefyPass dataclass

Bases: ModulePass

Source code in xdsl/transforms/desymref.py

class FrontendDesymrefyPass(ModulePass):
    name = "frontend-desymrefy"

    def apply(self, ctx: Context, op: builtin.ModuleOp):
        Desymrefier().desymrefy(op)

name = 'frontend-desymrefy' class-attribute instance-attribute

apply(ctx: Context, op: builtin.ModuleOp)

Source code in xdsl/transforms/desymref.py

def apply(self, ctx: Context, op: builtin.ModuleOp):
    Desymrefier().desymrefy(op)

has_symbol(op: Operation) -> bool

Source code in xdsl/transforms/desymref.py

def has_symbol(op: Operation) -> bool:
    return isinstance(op, symref.DeclareOp | symref.FetchOp | symref.UpdateOp)

get_symbol(op: Operation) -> str | None

Returns a symbol attribute for an operation. If operation has no symbol attributes, None is returned.

Source code in xdsl/transforms/desymref.py

def get_symbol(op: Operation) -> str | None:
    """
    Returns a symbol attribute for an operation. If operation has no symbol
    attributes, None is returned.
    """
    if isinstance(op, symref.DeclareOp):
        return op.sym_name.data
    elif isinstance(op, symref.FetchOp | symref.UpdateOp):
        return op.symbol.root_reference.data
    else:
        return None

get_symbols(block: Block) -> set[str]

Returns a set of all symbols defined/used in a basic block.

Source code in xdsl/transforms/desymref.py

def get_symbols(block: Block) -> set[str]:
    """Returns a set of all symbols defined/used in a basic block."""
    symbols: set[str] = set()
    for op in block.ops:
        if has_symbol(op):
            symbol = get_symbol(op)
            assert symbol is not None
            symbols.add(symbol)
    return symbols

lower_positional_bound(writes: list[symref.UpdateOp], read: symref.FetchOp) -> Operation | None

Returns a nearest write preceeding the read. If there is no such write, None is returned.

Pre-condition: list writes is sorted based on operation indices.

Source code in xdsl/transforms/desymref.py

def lower_positional_bound(
    writes: list[symref.UpdateOp], read: symref.FetchOp
) -> Operation | None:
    """
    Returns a nearest write preceeding the `read`. If there is no such write,
    `None` is returned.

    Pre-condition: list `writes` is sorted based on operation indices.
    """
    block = read.parent_block()
    assert block is not None

    idx = block.get_operation_index(read)
    low_idx = -1
    high_idx = len(writes) - 1

    # Binary search to find the right write.
    while low_idx < high_idx:
        mid_idx = (high_idx - low_idx + 1) // 2 + low_idx
        user_idx = block.get_operation_index(writes[mid_idx])

        if user_idx < idx:
            low_idx = mid_idx
        else:
            high_idx = mid_idx - 1

    if low_idx == -1:
        return None
    return writes[low_idx]