Skip to content

Air

air

Port of the AMD Xilinx AIR dialect for programming the AIEs on the AMD Xilinx Versal FPGA architecture. This is a higher-level dialect than the AIE dialect. It is used to program Versal cards over PCIe. AIE is a hardened systolic array present in the Versal devices. The dialect describes netlists of AIE components and it can be lowered to the processor's assembly using the vendor's compiler. A description of the original dialect can be found here https://xilinx.github.io/mlir-air/AIRDialect.html

AIR = Dialect('air', [AllocOp, ChannelOp, ChannelGetOp, ChannelPutOp, CustomOp, DeallocOp, DmaMemcpyNdOp, ExecuteOp, ExecuteTerminatorOp, HerdOp, HerdTerminatorOp, LaunchOp, LaunchTerminatorOp, HerdPipelineOp, PipelineGetOp, PipelinePutOp, PipelineStageOp, PipelineTerminatorOp, PipelineYieldOp, SegmentOp, SegmentTerminatorOp, WaitAllOp], [AsyncTokenAttr]) module-attribute

AsyncTokenAttr dataclass

Bases: ParametrizedAttribute, TypeAttribute

Source code in xdsl/dialects/experimental/air.py
@irdl_attr_definition
class AsyncTokenAttr(ParametrizedAttribute, TypeAttribute):
    name = "air.async.token"

name = 'air.async.token' class-attribute instance-attribute

AllocOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class AllocOp(IRDLOperation):
    name = "air.alloc"

    async_dependencies = var_operand_def(AsyncTokenAttr)

    async_token = result_def(AsyncTokenAttr)
    result = result_def(MemRefType)

    def __init__(
        self,
        async_dependencies: Operation,
        element_type: Attribute,
        shape: ArrayAttr[IntAttr],
    ):
        memref_type = MemRefType(element_type, shape)
        super().__init__(
            operands=[async_dependencies], result_types=[AsyncTokenAttr(), memref_type]
        )

name = 'air.alloc' class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr) class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr) class-attribute instance-attribute

result = result_def(MemRefType) class-attribute instance-attribute

__init__(async_dependencies: Operation, element_type: Attribute, shape: ArrayAttr[IntAttr])

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    async_dependencies: Operation,
    element_type: Attribute,
    shape: ArrayAttr[IntAttr],
):
    memref_type = MemRefType(element_type, shape)
    super().__init__(
        operands=[async_dependencies], result_types=[AsyncTokenAttr(), memref_type]
    )

ChannelOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class ChannelOp(IRDLOperation):
    name = "air.channel"

    sym_name = prop_def(SymbolRefAttr)
    size = prop_def(ArrayAttr)

    def __init__(
        self, sym_name: SymbolRefAttr, size: ArrayAttr[IntegerAttr]
    ):  # TODO: add verify to check 64-bit integer array attribute
        super().__init__(properties={"sym_name": sym_name, "size": size})

    assembly_format = "$sym_name $size attr-dict"

name = 'air.channel' class-attribute instance-attribute

sym_name = prop_def(SymbolRefAttr) class-attribute instance-attribute

size = prop_def(ArrayAttr) class-attribute instance-attribute

assembly_format = '$sym_name $size attr-dict' class-attribute instance-attribute

__init__(sym_name: SymbolRefAttr, size: ArrayAttr[IntegerAttr])

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self, sym_name: SymbolRefAttr, size: ArrayAttr[IntegerAttr]
):  # TODO: add verify to check 64-bit integer array attribute
    super().__init__(properties={"sym_name": sym_name, "size": size})

ChannelGetOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class ChannelGetOp(IRDLOperation):
    name = "air.channel.get"

    chan_name = attr_def(SymbolRefAttr)
    async_dependencies = var_operand_def(AsyncTokenAttr())
    indices = var_operand_def(AsyncTokenAttr())
    dst = operand_def(MemRefType)
    dst_offsets = var_operand_def(IndexType())
    dst_sizes = var_operand_def(IndexType())
    dst_strides = var_operand_def(IndexType())

    async_token = result_def(AsyncTokenAttr())

    irdl_options = (AttrSizedOperandSegments(),)

    def __init__(
        self,
        chan_name: SymbolRefAttr,
        async_dependencies: list[Operation | SSAValue],
        indices: list[Operation | SSAValue],
        dst: Operation | SSAValue,
        dst_offsets: list[Operation | SSAValue],
        dst_sizes: list[Operation | SSAValue],
        dst_strides: list[Operation | SSAValue],
    ):
        super().__init__(
            attributes={"chan_name": chan_name},
            operands=[
                async_dependencies,
                indices,
                dst,
                dst_offsets,
                dst_sizes,
                dst_strides,
            ],
        )

    assembly_format = "(`async` `[` $async_dependencies^ `]`)? $chan_name `[` $indices `]` `(` $dst `[` $dst_offsets `]``[` $dst_sizes `]``[` $dst_strides `]` `)` attr-dict `:` `(` type($dst) `)`"  # noqa: E501

name = 'air.channel.get' class-attribute instance-attribute

chan_name = attr_def(SymbolRefAttr) class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr()) class-attribute instance-attribute

indices = var_operand_def(AsyncTokenAttr()) class-attribute instance-attribute

dst = operand_def(MemRefType) class-attribute instance-attribute

dst_offsets = var_operand_def(IndexType()) class-attribute instance-attribute

dst_sizes = var_operand_def(IndexType()) class-attribute instance-attribute

dst_strides = var_operand_def(IndexType()) class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr()) class-attribute instance-attribute

irdl_options = (AttrSizedOperandSegments(),) class-attribute instance-attribute

assembly_format = '(`async` `[` $async_dependencies^ `]`)? $chan_name `[` $indices `]` `(` $dst `[` $dst_offsets `]``[` $dst_sizes `]``[` $dst_strides `]` `)` attr-dict `:` `(` type($dst) `)`' class-attribute instance-attribute

__init__(chan_name: SymbolRefAttr, async_dependencies: list[Operation | SSAValue], indices: list[Operation | SSAValue], dst: Operation | SSAValue, dst_offsets: list[Operation | SSAValue], dst_sizes: list[Operation | SSAValue], dst_strides: list[Operation | SSAValue])

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    chan_name: SymbolRefAttr,
    async_dependencies: list[Operation | SSAValue],
    indices: list[Operation | SSAValue],
    dst: Operation | SSAValue,
    dst_offsets: list[Operation | SSAValue],
    dst_sizes: list[Operation | SSAValue],
    dst_strides: list[Operation | SSAValue],
):
    super().__init__(
        attributes={"chan_name": chan_name},
        operands=[
            async_dependencies,
            indices,
            dst,
            dst_offsets,
            dst_sizes,
            dst_strides,
        ],
    )

ChannelPutOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class ChannelPutOp(IRDLOperation):
    name = "air.channel.put"

    chan_name = attr_def(SymbolRefAttr)

    async_dependencies = var_operand_def(AsyncTokenAttr())
    indices = var_operand_def(IndexType())
    src = operand_def(MemRefType)
    src_offsets = var_operand_def(IndexType())
    src_sizes = var_operand_def(IndexType())
    src_strides = var_operand_def(IndexType())

    async_token = result_def(AsyncTokenAttr())

    irdl_options = (AttrSizedOperandSegments(),)

    def __init__(
        self,
        chan_name: SymbolRefAttr,
        async_dependencies: list[Operation | SSAValue],
        indices: list[Operation | SSAValue],
        src: Operation | SSAValue,
        src_offsets: list[Operation | SSAValue],
        src_sizes: list[Operation | SSAValue],
        src_strides: list[Operation | SSAValue],
    ):
        super().__init__(
            properties={"chan_name": chan_name},
            operands=[
                async_dependencies,
                indices,
                src,
                src_offsets,
                src_sizes,
                src_strides,
            ],
            result_types=[AsyncTokenAttr()],
        )

    assembly_format = "(`async` `[` $async_dependencies^ `]`)? $chan_name `[` $indices `]` `(` $src `[` $src_offsets `]``[` $src_sizes `]``[` $src_strides `]` `)` attr-dict `:` `(` type($src) `)`"  # noqa: E501

name = 'air.channel.put' class-attribute instance-attribute

chan_name = attr_def(SymbolRefAttr) class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr()) class-attribute instance-attribute

indices = var_operand_def(IndexType()) class-attribute instance-attribute

src = operand_def(MemRefType) class-attribute instance-attribute

src_offsets = var_operand_def(IndexType()) class-attribute instance-attribute

src_sizes = var_operand_def(IndexType()) class-attribute instance-attribute

src_strides = var_operand_def(IndexType()) class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr()) class-attribute instance-attribute

irdl_options = (AttrSizedOperandSegments(),) class-attribute instance-attribute

assembly_format = '(`async` `[` $async_dependencies^ `]`)? $chan_name `[` $indices `]` `(` $src `[` $src_offsets `]``[` $src_sizes `]``[` $src_strides `]` `)` attr-dict `:` `(` type($src) `)`' class-attribute instance-attribute

__init__(chan_name: SymbolRefAttr, async_dependencies: list[Operation | SSAValue], indices: list[Operation | SSAValue], src: Operation | SSAValue, src_offsets: list[Operation | SSAValue], src_sizes: list[Operation | SSAValue], src_strides: list[Operation | SSAValue])

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    chan_name: SymbolRefAttr,
    async_dependencies: list[Operation | SSAValue],
    indices: list[Operation | SSAValue],
    src: Operation | SSAValue,
    src_offsets: list[Operation | SSAValue],
    src_sizes: list[Operation | SSAValue],
    src_strides: list[Operation | SSAValue],
):
    super().__init__(
        properties={"chan_name": chan_name},
        operands=[
            async_dependencies,
            indices,
            src,
            src_offsets,
            src_sizes,
            src_strides,
        ],
        result_types=[AsyncTokenAttr()],
    )

CustomOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class CustomOp(IRDLOperation):
    name = "air.custom"

    symbol = attr_def(SymbolRefAttr)
    async_dependencies = var_operand_def(AsyncTokenAttr)
    custom_operands = var_operand_def(Attribute)

    async_token = result_def(AsyncTokenAttr)

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    def __init__(
        self,
        symbol: SymbolRefAttr,
        async_dependencies: list[Operation | SSAValue],
        custom_operands: list[Operation | SSAValue],
    ):
        super().__init__(
            attributes={"symbol": symbol},
            operands=[async_dependencies, custom_operands],
            result_types=[AsyncTokenAttr()],
        )

name = 'air.custom' class-attribute instance-attribute

symbol = attr_def(SymbolRefAttr) class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr) class-attribute instance-attribute

custom_operands = var_operand_def(Attribute) class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr) class-attribute instance-attribute

irdl_options = (AttrSizedOperandSegments(as_property=True),) class-attribute instance-attribute

__init__(symbol: SymbolRefAttr, async_dependencies: list[Operation | SSAValue], custom_operands: list[Operation | SSAValue])

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    symbol: SymbolRefAttr,
    async_dependencies: list[Operation | SSAValue],
    custom_operands: list[Operation | SSAValue],
):
    super().__init__(
        attributes={"symbol": symbol},
        operands=[async_dependencies, custom_operands],
        result_types=[AsyncTokenAttr()],
    )

DeallocOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class DeallocOp(IRDLOperation):
    name = "air.dealloc"

    async_dependencies = var_operand_def(AsyncTokenAttr)
    memref = operand_def(MemRefType)

    async_token = result_def(AsyncTokenAttr)

    def __init__(
        self,
        async_dependencies: list[Operation | SSAValue],
        memref: Operation | SSAValue,
    ):
        super().__init__(
            operands=[async_dependencies, memref], result_types=[AsyncTokenAttr()]
        )

name = 'air.dealloc' class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr) class-attribute instance-attribute

memref = operand_def(MemRefType) class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr) class-attribute instance-attribute

__init__(async_dependencies: list[Operation | SSAValue], memref: Operation | SSAValue)

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    async_dependencies: list[Operation | SSAValue],
    memref: Operation | SSAValue,
):
    super().__init__(
        operands=[async_dependencies, memref], result_types=[AsyncTokenAttr()]
    )

DmaMemcpyNdOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class DmaMemcpyNdOp(IRDLOperation):
    name = "air.dma_memcpy_nd"

    async_dependencies = var_operand_def(AsyncTokenAttr())
    dst = operand_def(MemRefType)
    dst_offsets = var_operand_def(IndexType())
    dst_sizes = var_operand_def(IndexType())
    dst_strides = var_operand_def(IndexType())
    src = operand_def(MemRefType)
    src_offsets = var_operand_def(IndexType())
    src_sizes = var_operand_def(IndexType())
    src_strides = var_operand_def(IndexType())

    async_token = result_def(AsyncTokenAttr())

    irdl_options = (
        AttrSizedOperandSegments(as_property=True),
        ParsePropInAttrDict(),
    )

    def __init__(
        self,
        async_dependencies: list[Operation | SSAValue] | None,
        dst: Operation | SSAValue,
        dst_offsets: list[Operation | SSAValue],
        dst_sizes: list[Operation | SSAValue],
        dst_strides: list[Operation | SSAValue],
        src: Operation | SSAValue,
        src_offsets: list[Operation | SSAValue],
        src_sizes: list[Operation | SSAValue],
        src_strides: list[Operation | SSAValue],
    ):
        print(
            async_dependencies,
            dst,
            dst_offsets,
            dst_sizes,
            dst_strides,
            src,
            src_offsets,
            src_sizes,
            src_strides,
        )
        super().__init__(
            operands=[
                async_dependencies,
                dst,
                dst_offsets,
                dst_sizes,
                dst_strides,
                src,
                src_offsets,
                src_sizes,
                src_strides,
            ],
            result_types=[AsyncTokenAttr()],
        )

    assembly_format = "(`async` $async_dependencies^)? `(` $dst `[` $dst_offsets `]``[` $dst_sizes `]``[` $dst_strides `]` `,` $src `[` $src_offsets `]``[` $src_sizes `]``[` $src_strides `]` `)`  attr-dict `:` `(` type($dst) `,` type($src) `)`"  # noqa: E501

name = 'air.dma_memcpy_nd' class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr()) class-attribute instance-attribute

dst = operand_def(MemRefType) class-attribute instance-attribute

dst_offsets = var_operand_def(IndexType()) class-attribute instance-attribute

dst_sizes = var_operand_def(IndexType()) class-attribute instance-attribute

dst_strides = var_operand_def(IndexType()) class-attribute instance-attribute

src = operand_def(MemRefType) class-attribute instance-attribute

src_offsets = var_operand_def(IndexType()) class-attribute instance-attribute

src_sizes = var_operand_def(IndexType()) class-attribute instance-attribute

src_strides = var_operand_def(IndexType()) class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr()) class-attribute instance-attribute

irdl_options = (AttrSizedOperandSegments(as_property=True), ParsePropInAttrDict()) class-attribute instance-attribute

assembly_format = '(`async` $async_dependencies^)? `(` $dst `[` $dst_offsets `]``[` $dst_sizes `]``[` $dst_strides `]` `,` $src `[` $src_offsets `]``[` $src_sizes `]``[` $src_strides `]` `)` attr-dict `:` `(` type($dst) `,` type($src) `)`' class-attribute instance-attribute

__init__(async_dependencies: list[Operation | SSAValue] | None, dst: Operation | SSAValue, dst_offsets: list[Operation | SSAValue], dst_sizes: list[Operation | SSAValue], dst_strides: list[Operation | SSAValue], src: Operation | SSAValue, src_offsets: list[Operation | SSAValue], src_sizes: list[Operation | SSAValue], src_strides: list[Operation | SSAValue])

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    async_dependencies: list[Operation | SSAValue] | None,
    dst: Operation | SSAValue,
    dst_offsets: list[Operation | SSAValue],
    dst_sizes: list[Operation | SSAValue],
    dst_strides: list[Operation | SSAValue],
    src: Operation | SSAValue,
    src_offsets: list[Operation | SSAValue],
    src_sizes: list[Operation | SSAValue],
    src_strides: list[Operation | SSAValue],
):
    print(
        async_dependencies,
        dst,
        dst_offsets,
        dst_sizes,
        dst_strides,
        src,
        src_offsets,
        src_sizes,
        src_strides,
    )
    super().__init__(
        operands=[
            async_dependencies,
            dst,
            dst_offsets,
            dst_sizes,
            dst_strides,
            src,
            src_offsets,
            src_sizes,
            src_strides,
        ],
        result_types=[AsyncTokenAttr()],
    )

ExecuteOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class ExecuteOp(IRDLOperation):
    name = "air.execute"

    async_dependencies = var_operand_def(AsyncTokenAttr)
    async_token = result_def(AsyncTokenAttr)
    results_ = var_result_def(Attribute)
    body = region_def()

    traits = lazy_traits_def(
        lambda: (SingleBlockImplicitTerminator(ExecuteTerminatorOp),)
    )

    def __init__(
        self,
        async_dependencies: list[SSAValue] | None,
        result_types: list[Attribute] | None,
        body: Region,
    ):
        super().__init__(
            operands=[async_dependencies],
            result_types=[AsyncTokenAttr(), result_types],
            regions=[body],
        )

    @classmethod
    def parse(cls, parser: Parser) -> ExecuteOp:
        async_dependencies = parser.parse_optional_comma_separated_list(
            parser.Delimiter.SQUARE, parser.parse_operand
        )

        result_types: list[Attribute] | None = []
        if parser.parse_optional_characters("->"):
            result_types = parser.parse_optional_comma_separated_list(
                parser.Delimiter.PAREN, parser.parse_type
            )

        body = parser.parse_region()

        return ExecuteOp(async_dependencies, result_types, body)

name = 'air.execute' class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr) class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr) class-attribute instance-attribute

results_ = var_result_def(Attribute) class-attribute instance-attribute

body = region_def() class-attribute instance-attribute

traits = lazy_traits_def(lambda: (SingleBlockImplicitTerminator(ExecuteTerminatorOp),)) class-attribute instance-attribute

__init__(async_dependencies: list[SSAValue] | None, result_types: list[Attribute] | None, body: Region)

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    async_dependencies: list[SSAValue] | None,
    result_types: list[Attribute] | None,
    body: Region,
):
    super().__init__(
        operands=[async_dependencies],
        result_types=[AsyncTokenAttr(), result_types],
        regions=[body],
    )

parse(parser: Parser) -> ExecuteOp classmethod

Source code in xdsl/dialects/experimental/air.py
@classmethod
def parse(cls, parser: Parser) -> ExecuteOp:
    async_dependencies = parser.parse_optional_comma_separated_list(
        parser.Delimiter.SQUARE, parser.parse_operand
    )

    result_types: list[Attribute] | None = []
    if parser.parse_optional_characters("->"):
        result_types = parser.parse_optional_comma_separated_list(
            parser.Delimiter.PAREN, parser.parse_type
        )

    body = parser.parse_region()

    return ExecuteOp(async_dependencies, result_types, body)

ExecuteTerminatorOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class ExecuteTerminatorOp(IRDLOperation):
    name = "air.execute_terminator"

    # even though this is an operand they decided to name it "result" in the original specification
    results_op = var_operand_def()

    traits = traits_def(HasParent(ExecuteOp), IsTerminator())

    def __init__(self, results_op: list[Operation | SSAValue]):
        super().__init__(operands=[results_op])

    @classmethod
    def parse(cls, parser: Parser) -> ExecuteTerminatorOp:
        results_op: list[Operation | SSAValue] = []
        while not parser.parse_optional_characters(":"):
            results_op.append(parser.parse_operand())
        while parser.parse_optional_type():
            parser.parse_optional_characters(",")

        return ExecuteTerminatorOp(results_op)

name = 'air.execute_terminator' class-attribute instance-attribute

results_op = var_operand_def() class-attribute instance-attribute

traits = traits_def(HasParent(ExecuteOp), IsTerminator()) class-attribute instance-attribute

__init__(results_op: list[Operation | SSAValue])

Source code in xdsl/dialects/experimental/air.py
def __init__(self, results_op: list[Operation | SSAValue]):
    super().__init__(operands=[results_op])

parse(parser: Parser) -> ExecuteTerminatorOp classmethod

Source code in xdsl/dialects/experimental/air.py
@classmethod
def parse(cls, parser: Parser) -> ExecuteTerminatorOp:
    results_op: list[Operation | SSAValue] = []
    while not parser.parse_optional_characters(":"):
        results_op.append(parser.parse_operand())
    while parser.parse_optional_type():
        parser.parse_optional_characters(",")

    return ExecuteTerminatorOp(results_op)

HerdTerminatorOp dataclass

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class HerdTerminatorOp(IRDLOperation):
    name = "air.herd_terminator"

    traits = lazy_traits_def(lambda: (HasParent(HerdOp), IsTerminator()))

    assembly_format = "attr-dict"

name = 'air.herd_terminator' class-attribute instance-attribute

traits = lazy_traits_def(lambda: (HasParent(HerdOp), IsTerminator())) class-attribute instance-attribute

assembly_format = 'attr-dict' class-attribute instance-attribute

HerdOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class HerdOp(IRDLOperation):
    name = "air.herd"

    sym_name = opt_prop_def(StringAttr)
    async_dependencies = var_operand_def(AsyncTokenAttr())
    sizes = var_operand_def(IndexType())
    herd_operands = var_operand_def()
    async_token = opt_result_def(AsyncTokenAttr)
    region = opt_region_def()

    traits = traits_def(
        IsolatedFromAbove(), SingleBlockImplicitTerminator(HerdTerminatorOp)
    )

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    def __init__(
        self,
        sym_name: StringAttr | None,
        async_dependencies: list[SSAValue] | None,
        sizes: list[Operation | SSAValue],
        herd_operands: list[Operation | SSAValue] | None,
        region: Region | None,
    ):
        super().__init__(
            properties={"sym_name": sym_name},
            operands=[async_dependencies, sizes, herd_operands],
            result_types=[AsyncTokenAttr()],
            regions=[region],
        )

    def print(self, printer: Printer):
        printer.print_string(" tile")
        with printer.in_parens():
            if len(self.sizes) == 1:
                printer.print_string("%tx")
            if len(self.sizes) == 2:
                printer.print_string("%tx, %ty")
            if len(self.sizes) == 3:
                printer.print_string("%tx, %ty, %tz")
        printer.print_string(" in ")
        with printer.in_parens():
            if len(self.sizes) == 1:
                printer.print_string("%size_x = ")
                printer.print_ssa_value(self.sizes[0])
            if len(self.sizes) == 2:
                printer.print_string("%size_x = ")
                printer.print_ssa_value(self.sizes[0])
                printer.print_string(", ")
                printer.print_string("%size_y = ")
                printer.print_ssa_value(self.sizes[1])
            if len(self.sizes) == 3:
                printer.print_string("%size_x = ")
                printer.print_ssa_value(self.sizes[0])
                printer.print_string(", ")
                printer.print_string("%size_y = ")
                printer.print_ssa_value(self.sizes[1])
                printer.print_string(", ")
                printer.print_string("%size_z = ")
                printer.print_ssa_value(self.sizes[2])

        if self.herd_operands:
            printer.print_string(" args")
            with printer.in_parens():
                if len(self.herd_operands) == 1:
                    printer.print_string("%ext0 = ")
                    printer.print_ssa_value(self.herd_operands[0])
                if len(self.herd_operands) == 2:
                    printer.print_string("%ext0 = ")
                    printer.print_ssa_value(self.herd_operands[0])
                    printer.print_string(", ")
                    printer.print_string("%ext1 = ")
                    printer.print_ssa_value(self.herd_operands[1])
                if len(self.herd_operands) == 3:
                    printer.print_string("%ext0 = ")
                    printer.print_ssa_value(self.herd_operands[0])
                    printer.print_string(", ")
                    printer.print_string("%ext1 = ")
                    printer.print_ssa_value(self.herd_operands[1])
                    printer.print_string(", ")
                    printer.print_string("%ext2 = ")
                    printer.print_ssa_value(self.herd_operands[2])

                printer.print_string(" : ")
                printer.print_list(
                    self.herd_operands, lambda o: printer.print_attribute(o.type), ","
                )

        if self.region:
            printer.print_region(self.region)

    @classmethod
    def parse(cls, parser: Parser) -> HerdOp:
        sym_name = parser.parse_optional_symbol_name()
        async_dependencies = parser.parse_optional_comma_separated_list(
            parser.Delimiter.SQUARE, parser.parse_operand
        )
        parser.parse_keyword("tile")
        arg_list = parser.parse_op_args_list()
        parser.parse_keyword("in")

        parser.parse_characters("(")
        tile_size_lst: list[Operation | SSAValue] = []
        for n_arg in range(len(arg_list)):
            parser.parse_optional_argument(False)
            parser.parse_characters("=")
            tile_size = parser.parse_operand()
            tile_size_lst.append(tile_size)
            if n_arg < len(arg_list) - 1:
                parser.parse_characters(",")

        parser.parse_characters(")")
        arguments_lst: list[Parser.Argument] = []

        operands_lst: list[Operation | SSAValue] = []
        if parser.parse_optional_keyword("args"):
            parser.parse_characters("(")
            while True:
                argument = parser.parse_argument(expect_type=False)
                parser.parse_characters("=")
                operand = parser.parse_operand()

                # The type of the block argument is not known until the operand has been parsed, e.g. %ext0 = %arg0
                argument = argument.resolve(operand.type)
                arguments_lst.append(argument)
                operands_lst.append(operand)

                if not parser.parse_optional_characters(","):
                    break
            parser.parse_characters(")")

            parser.parse_characters(":")

            for n_arg in range(len(operands_lst)):
                parser.parse_type()
                parser.parse_optional_characters(",")

        parser.parse_keyword("attributes")
        parser.parse_optional_attr_dict()
        region = parser.parse_optional_region(arguments_lst)

        return HerdOp(sym_name, async_dependencies, tile_size_lst, operands_lst, region)

name = 'air.herd' class-attribute instance-attribute

sym_name = opt_prop_def(StringAttr) class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr()) class-attribute instance-attribute

sizes = var_operand_def(IndexType()) class-attribute instance-attribute

herd_operands = var_operand_def() class-attribute instance-attribute

async_token = opt_result_def(AsyncTokenAttr) class-attribute instance-attribute

region = opt_region_def() class-attribute instance-attribute

traits = traits_def(IsolatedFromAbove(), SingleBlockImplicitTerminator(HerdTerminatorOp)) class-attribute instance-attribute

irdl_options = (AttrSizedOperandSegments(as_property=True),) class-attribute instance-attribute

__init__(sym_name: StringAttr | None, async_dependencies: list[SSAValue] | None, sizes: list[Operation | SSAValue], herd_operands: list[Operation | SSAValue] | None, region: Region | None)

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    sym_name: StringAttr | None,
    async_dependencies: list[SSAValue] | None,
    sizes: list[Operation | SSAValue],
    herd_operands: list[Operation | SSAValue] | None,
    region: Region | None,
):
    super().__init__(
        properties={"sym_name": sym_name},
        operands=[async_dependencies, sizes, herd_operands],
        result_types=[AsyncTokenAttr()],
        regions=[region],
    )

print(printer: Printer)

Source code in xdsl/dialects/experimental/air.py
def print(self, printer: Printer):
    printer.print_string(" tile")
    with printer.in_parens():
        if len(self.sizes) == 1:
            printer.print_string("%tx")
        if len(self.sizes) == 2:
            printer.print_string("%tx, %ty")
        if len(self.sizes) == 3:
            printer.print_string("%tx, %ty, %tz")
    printer.print_string(" in ")
    with printer.in_parens():
        if len(self.sizes) == 1:
            printer.print_string("%size_x = ")
            printer.print_ssa_value(self.sizes[0])
        if len(self.sizes) == 2:
            printer.print_string("%size_x = ")
            printer.print_ssa_value(self.sizes[0])
            printer.print_string(", ")
            printer.print_string("%size_y = ")
            printer.print_ssa_value(self.sizes[1])
        if len(self.sizes) == 3:
            printer.print_string("%size_x = ")
            printer.print_ssa_value(self.sizes[0])
            printer.print_string(", ")
            printer.print_string("%size_y = ")
            printer.print_ssa_value(self.sizes[1])
            printer.print_string(", ")
            printer.print_string("%size_z = ")
            printer.print_ssa_value(self.sizes[2])

    if self.herd_operands:
        printer.print_string(" args")
        with printer.in_parens():
            if len(self.herd_operands) == 1:
                printer.print_string("%ext0 = ")
                printer.print_ssa_value(self.herd_operands[0])
            if len(self.herd_operands) == 2:
                printer.print_string("%ext0 = ")
                printer.print_ssa_value(self.herd_operands[0])
                printer.print_string(", ")
                printer.print_string("%ext1 = ")
                printer.print_ssa_value(self.herd_operands[1])
            if len(self.herd_operands) == 3:
                printer.print_string("%ext0 = ")
                printer.print_ssa_value(self.herd_operands[0])
                printer.print_string(", ")
                printer.print_string("%ext1 = ")
                printer.print_ssa_value(self.herd_operands[1])
                printer.print_string(", ")
                printer.print_string("%ext2 = ")
                printer.print_ssa_value(self.herd_operands[2])

            printer.print_string(" : ")
            printer.print_list(
                self.herd_operands, lambda o: printer.print_attribute(o.type), ","
            )

    if self.region:
        printer.print_region(self.region)

parse(parser: Parser) -> HerdOp classmethod

Source code in xdsl/dialects/experimental/air.py
@classmethod
def parse(cls, parser: Parser) -> HerdOp:
    sym_name = parser.parse_optional_symbol_name()
    async_dependencies = parser.parse_optional_comma_separated_list(
        parser.Delimiter.SQUARE, parser.parse_operand
    )
    parser.parse_keyword("tile")
    arg_list = parser.parse_op_args_list()
    parser.parse_keyword("in")

    parser.parse_characters("(")
    tile_size_lst: list[Operation | SSAValue] = []
    for n_arg in range(len(arg_list)):
        parser.parse_optional_argument(False)
        parser.parse_characters("=")
        tile_size = parser.parse_operand()
        tile_size_lst.append(tile_size)
        if n_arg < len(arg_list) - 1:
            parser.parse_characters(",")

    parser.parse_characters(")")
    arguments_lst: list[Parser.Argument] = []

    operands_lst: list[Operation | SSAValue] = []
    if parser.parse_optional_keyword("args"):
        parser.parse_characters("(")
        while True:
            argument = parser.parse_argument(expect_type=False)
            parser.parse_characters("=")
            operand = parser.parse_operand()

            # The type of the block argument is not known until the operand has been parsed, e.g. %ext0 = %arg0
            argument = argument.resolve(operand.type)
            arguments_lst.append(argument)
            operands_lst.append(operand)

            if not parser.parse_optional_characters(","):
                break
        parser.parse_characters(")")

        parser.parse_characters(":")

        for n_arg in range(len(operands_lst)):
            parser.parse_type()
            parser.parse_optional_characters(",")

    parser.parse_keyword("attributes")
    parser.parse_optional_attr_dict()
    region = parser.parse_optional_region(arguments_lst)

    return HerdOp(sym_name, async_dependencies, tile_size_lst, operands_lst, region)

LaunchTerminatorOp dataclass

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class LaunchTerminatorOp(IRDLOperation):
    name = "air.launch_terminator"

    traits = lazy_traits_def(
        lambda: (
            HasParent(LaunchOp),
            IsTerminator(),
        )
    )

name = 'air.launch_terminator' class-attribute instance-attribute

traits = lazy_traits_def(lambda: (HasParent(LaunchOp), IsTerminator())) class-attribute instance-attribute

LaunchOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class LaunchOp(IRDLOperation):
    name = "air.launch"

    sym_name = opt_prop_def(StringAttr)
    async_dependencies = var_operand_def(AsyncTokenAttr())
    sizes = var_operand_def(IndexType())
    launch_operands = var_operand_def()
    async_token = result_def(AsyncTokenAttr)
    body = opt_region_def()

    traits = traits_def(
        IsolatedFromAbove(), SingleBlockImplicitTerminator(LaunchTerminatorOp)
    )

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    def __init__(
        self,
        sym_name: StringAttr | None,
        async_dependencies: list[SSAValue] | None,
        sizes: list[Operation | SSAValue],
        launch_operands: list[Operation | SSAValue],
        body: Region | None,
    ):
        super().__init__(
            attributes={"sym_name": sym_name},
            operands=[async_dependencies, sizes, launch_operands],
            result_types=[AsyncTokenAttr()],
            regions=[body],
        )

    @classmethod
    def parse(cls, parser: Parser) -> LaunchOp:
        sym_name = parser.parse_optional_symbol_name()

        async_dependencies = parser.parse_optional_comma_separated_list(
            parser.Delimiter.SQUARE, parser.parse_operand
        )

        block_args_lst: list[Parser.Argument] = []
        if parser.parse_optional_characters("("):
            while not parser.parse_optional_characters(")"):
                b_arg = parser.parse_argument(expect_type=False)
                b_arg = b_arg.resolve(IndexType())
                block_args_lst.append(b_arg)
                parser.parse_optional_characters(",")

        parser.parse_keyword("in")
        parser.parse_characters("(")

        arguments_lst: list[Parser.Argument] = []

        sizes_operands_lst: list[Operation | SSAValue] = []
        while not parser.parse_optional_characters(")"):
            argument = parser.parse_argument(expect_type=False)
            parser.parse_characters("=")
            operand = parser.parse_operand()

            argument = argument.resolve(operand.type)
            arguments_lst.append(argument)
            sizes_operands_lst.append(operand)

            parser.parse_optional_characters(",")

        launch_operands_lst: list[Operation | SSAValue] = []
        attr_dict: dict[str, Attribute] = dict()
        if parser.parse_optional_keyword("args"):
            parser.parse_characters("(")
            while True:
                argument = parser.parse_argument(expect_type=False)
                parser.parse_characters("=")
                operand = parser.parse_operand()

                argument = argument.resolve(operand.type)
                block_args_lst.append(argument)
                launch_operands_lst.append(operand)

                if not parser.parse_optional_characters(","):
                    break
            parser.parse_characters(")")

            parser.parse_characters(":")
            for _ in range(len(launch_operands_lst)):
                parser.parse_type()
                parser.parse_optional_characters(",")

            if parser.parse_optional_keyword("attributes"):
                attr_dict = parser.parse_optional_attr_dict()

        body = parser.parse_optional_region(block_args_lst)

        launch_op = LaunchOp(
            sym_name, async_dependencies, sizes_operands_lst, launch_operands_lst, body
        )
        launch_op.attributes |= attr_dict

        return launch_op

name = 'air.launch' class-attribute instance-attribute

sym_name = opt_prop_def(StringAttr) class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr()) class-attribute instance-attribute

sizes = var_operand_def(IndexType()) class-attribute instance-attribute

launch_operands = var_operand_def() class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr) class-attribute instance-attribute

body = opt_region_def() class-attribute instance-attribute

traits = traits_def(IsolatedFromAbove(), SingleBlockImplicitTerminator(LaunchTerminatorOp)) class-attribute instance-attribute

irdl_options = (AttrSizedOperandSegments(as_property=True),) class-attribute instance-attribute

__init__(sym_name: StringAttr | None, async_dependencies: list[SSAValue] | None, sizes: list[Operation | SSAValue], launch_operands: list[Operation | SSAValue], body: Region | None)

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    sym_name: StringAttr | None,
    async_dependencies: list[SSAValue] | None,
    sizes: list[Operation | SSAValue],
    launch_operands: list[Operation | SSAValue],
    body: Region | None,
):
    super().__init__(
        attributes={"sym_name": sym_name},
        operands=[async_dependencies, sizes, launch_operands],
        result_types=[AsyncTokenAttr()],
        regions=[body],
    )

parse(parser: Parser) -> LaunchOp classmethod

Source code in xdsl/dialects/experimental/air.py
@classmethod
def parse(cls, parser: Parser) -> LaunchOp:
    sym_name = parser.parse_optional_symbol_name()

    async_dependencies = parser.parse_optional_comma_separated_list(
        parser.Delimiter.SQUARE, parser.parse_operand
    )

    block_args_lst: list[Parser.Argument] = []
    if parser.parse_optional_characters("("):
        while not parser.parse_optional_characters(")"):
            b_arg = parser.parse_argument(expect_type=False)
            b_arg = b_arg.resolve(IndexType())
            block_args_lst.append(b_arg)
            parser.parse_optional_characters(",")

    parser.parse_keyword("in")
    parser.parse_characters("(")

    arguments_lst: list[Parser.Argument] = []

    sizes_operands_lst: list[Operation | SSAValue] = []
    while not parser.parse_optional_characters(")"):
        argument = parser.parse_argument(expect_type=False)
        parser.parse_characters("=")
        operand = parser.parse_operand()

        argument = argument.resolve(operand.type)
        arguments_lst.append(argument)
        sizes_operands_lst.append(operand)

        parser.parse_optional_characters(",")

    launch_operands_lst: list[Operation | SSAValue] = []
    attr_dict: dict[str, Attribute] = dict()
    if parser.parse_optional_keyword("args"):
        parser.parse_characters("(")
        while True:
            argument = parser.parse_argument(expect_type=False)
            parser.parse_characters("=")
            operand = parser.parse_operand()

            argument = argument.resolve(operand.type)
            block_args_lst.append(argument)
            launch_operands_lst.append(operand)

            if not parser.parse_optional_characters(","):
                break
        parser.parse_characters(")")

        parser.parse_characters(":")
        for _ in range(len(launch_operands_lst)):
            parser.parse_type()
            parser.parse_optional_characters(",")

        if parser.parse_optional_keyword("attributes"):
            attr_dict = parser.parse_optional_attr_dict()

    body = parser.parse_optional_region(block_args_lst)

    launch_op = LaunchOp(
        sym_name, async_dependencies, sizes_operands_lst, launch_operands_lst, body
    )
    launch_op.attributes |= attr_dict

    return launch_op

HerdPipelineOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class HerdPipelineOp(IRDLOperation):
    name = "air.pipeline"

    body = opt_region_def()

    traits = traits_def(HasParent(HerdOp))

    def __init__(self, body: Region | None):
        super().__init__(regions=[body])

    @classmethod
    def parse(cls, parser: Parser) -> HerdPipelineOp:
        attr_dict: dict[str, Attribute] = dict()
        if parser.parse_optional_keyword("attributes"):
            attr_dict = parser.parse_optional_attr_dict()

        body = parser.parse_optional_region()

        herd_pipeline_op = HerdPipelineOp(body)
        herd_pipeline_op.attributes |= attr_dict

        return herd_pipeline_op

name = 'air.pipeline' class-attribute instance-attribute

body = opt_region_def() class-attribute instance-attribute

traits = traits_def(HasParent(HerdOp)) class-attribute instance-attribute

__init__(body: Region | None)

Source code in xdsl/dialects/experimental/air.py
def __init__(self, body: Region | None):
    super().__init__(regions=[body])

parse(parser: Parser) -> HerdPipelineOp classmethod

Source code in xdsl/dialects/experimental/air.py
@classmethod
def parse(cls, parser: Parser) -> HerdPipelineOp:
    attr_dict: dict[str, Attribute] = dict()
    if parser.parse_optional_keyword("attributes"):
        attr_dict = parser.parse_optional_attr_dict()

    body = parser.parse_optional_region()

    herd_pipeline_op = HerdPipelineOp(body)
    herd_pipeline_op.attributes |= attr_dict

    return herd_pipeline_op

PipelineGetOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class PipelineGetOp(IRDLOperation):
    name = "air.pipeline.get"

    src0 = operand_def(Attribute)
    src1 = operand_def(Attribute)
    results = var_result_def(Attribute)

    def __init__(
        self,
        src0: Operation | SSAValue,
        src1: Operation | SSAValue,
        result_types: list[Attribute],
    ):
        super().__init__(operands=[src0, src1], result_types=result_types)

name = 'air.pipeline.get' class-attribute instance-attribute

src0 = operand_def(Attribute) class-attribute instance-attribute

src1 = operand_def(Attribute) class-attribute instance-attribute

results = var_result_def(Attribute) class-attribute instance-attribute

__init__(src0: Operation | SSAValue, src1: Operation | SSAValue, result_types: list[Attribute])

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    src0: Operation | SSAValue,
    src1: Operation | SSAValue,
    result_types: list[Attribute],
):
    super().__init__(operands=[src0, src1], result_types=result_types)

PipelinePutOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class PipelinePutOp(IRDLOperation):
    name = "air.pipeline.put"

    dst0 = operand_def(Attribute)
    dst1 = operand_def(Attribute)
    opers = var_operand_def(Attribute)

    def __init__(
        self,
        dst0: Operation | SSAValue,
        dst1: Operation | SSAValue,
        opers: list[Operation | SSAValue],
    ):
        super().__init__(operands=[dst0, dst1, opers])

name = 'air.pipeline.put' class-attribute instance-attribute

dst0 = operand_def(Attribute) class-attribute instance-attribute

dst1 = operand_def(Attribute) class-attribute instance-attribute

opers = var_operand_def(Attribute) class-attribute instance-attribute

__init__(dst0: Operation | SSAValue, dst1: Operation | SSAValue, opers: list[Operation | SSAValue])

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    dst0: Operation | SSAValue,
    dst1: Operation | SSAValue,
    opers: list[Operation | SSAValue],
):
    super().__init__(operands=[dst0, dst1, opers])

PipelineStageOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class PipelineStageOp(IRDLOperation):
    name = "air.pipeline.stage"

    opers = var_operand_def(Attribute)
    result = var_result_def()

    body = opt_region_def()

    traits = traits_def(HasParent(HerdPipelineOp))

    def __init__(
        self,
        opers: list[Operation | SSAValue],
        result_types: Sequence[Attribute],
        body: Region | None,
    ):
        if not result_types:
            result_types = []
        super().__init__(operands=[opers], result_types=[result_types], regions=[body])

    @classmethod
    def parse(cls, parser: Parser) -> PipelineStageOp:
        kernel_ops_lst: list[Operation | SSAValue] = []
        result_types_lst: list[Attribute] = []

        arguments_lst: list[Parser.Argument] = []
        pipeline_operands_lst: list[Operation | SSAValue] = []
        if parser.parse_optional_keyword("args"):
            parser.parse_characters("(")
            while not parser.parse_optional_characters(")"):
                argument = parser.parse_argument(expect_type=False)
                parser.parse_characters("=")
                operand = parser.parse_operand()
                pipeline_operands_lst.append(operand)

                argument = argument.resolve(operand.type)
                arguments_lst.append(argument)

                parser.parse_optional_characters(",")

            if parser.parse_optional_characters(":"):
                for _ in range(len(arguments_lst)):
                    parser.parse_type()
                    parser.parse_optional_characters(",")

        body = parser.parse_optional_region(arguments=arguments_lst)
        if parser.parse_optional_characters(":"):
            result_types_lst.append(parser.parse_type())
            parser.parse_optional_characters(",")

        return PipelineStageOp(kernel_ops_lst, result_types_lst, body)

name = 'air.pipeline.stage' class-attribute instance-attribute

opers = var_operand_def(Attribute) class-attribute instance-attribute

result = var_result_def() class-attribute instance-attribute

body = opt_region_def() class-attribute instance-attribute

traits = traits_def(HasParent(HerdPipelineOp)) class-attribute instance-attribute

__init__(opers: list[Operation | SSAValue], result_types: Sequence[Attribute], body: Region | None)

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    opers: list[Operation | SSAValue],
    result_types: Sequence[Attribute],
    body: Region | None,
):
    if not result_types:
        result_types = []
    super().__init__(operands=[opers], result_types=[result_types], regions=[body])

parse(parser: Parser) -> PipelineStageOp classmethod

Source code in xdsl/dialects/experimental/air.py
@classmethod
def parse(cls, parser: Parser) -> PipelineStageOp:
    kernel_ops_lst: list[Operation | SSAValue] = []
    result_types_lst: list[Attribute] = []

    arguments_lst: list[Parser.Argument] = []
    pipeline_operands_lst: list[Operation | SSAValue] = []
    if parser.parse_optional_keyword("args"):
        parser.parse_characters("(")
        while not parser.parse_optional_characters(")"):
            argument = parser.parse_argument(expect_type=False)
            parser.parse_characters("=")
            operand = parser.parse_operand()
            pipeline_operands_lst.append(operand)

            argument = argument.resolve(operand.type)
            arguments_lst.append(argument)

            parser.parse_optional_characters(",")

        if parser.parse_optional_characters(":"):
            for _ in range(len(arguments_lst)):
                parser.parse_type()
                parser.parse_optional_characters(",")

    body = parser.parse_optional_region(arguments=arguments_lst)
    if parser.parse_optional_characters(":"):
        result_types_lst.append(parser.parse_type())
        parser.parse_optional_characters(",")

    return PipelineStageOp(kernel_ops_lst, result_types_lst, body)

PipelineTerminatorOp dataclass

Bases: AbstractYieldOperation[Attribute]

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class PipelineTerminatorOp(AbstractYieldOperation[Attribute]):
    name = "air.pipeline.terminator"

    traits = traits_def(HasParent(HerdPipelineOp), IsTerminator())

name = 'air.pipeline.terminator' class-attribute instance-attribute

traits = traits_def(HasParent(HerdPipelineOp), IsTerminator()) class-attribute instance-attribute

PipelineYieldOp dataclass

Bases: AbstractYieldOperation[Attribute]

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class PipelineYieldOp(AbstractYieldOperation[Attribute]):
    name = "air.pipeline.yield"

    traits = traits_def(HasParent(PipelineStageOp), IsTerminator())

name = 'air.pipeline.yield' class-attribute instance-attribute

traits = traits_def(HasParent(PipelineStageOp), IsTerminator()) class-attribute instance-attribute

SegmentTerminatorOp dataclass

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class SegmentTerminatorOp(IRDLOperation):
    name = "air.segment_terminator"

    traits = lazy_traits_def(lambda: (HasParent(SegmentOp), IsTerminator()))

name = 'air.segment_terminator' class-attribute instance-attribute

traits = lazy_traits_def(lambda: (HasParent(SegmentOp), IsTerminator())) class-attribute instance-attribute

SegmentOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class SegmentOp(IRDLOperation):
    name = "air.segment"

    sym_name = opt_prop_def(StringAttr)
    async_dependencies = var_operand_def(AsyncTokenAttr())
    sizes = var_operand_def(IndexType())
    segment_operands = var_operand_def()
    async_token = result_def(AsyncTokenAttr)

    body = opt_region_def()

    traits = traits_def(
        IsolatedFromAbove(), SingleBlockImplicitTerminator(SegmentTerminatorOp)
    )

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    def __init__(
        self,
        sym_name: StringAttr | None,
        async_dependencies: list[Operation | SSAValue],
        sizes: list[Operation | SSAValue],
        segment_operands: list[Operation | SSAValue],
        body: Region | None,
    ):
        super().__init__(
            attributes={"sym_name": sym_name},
            operands=[async_dependencies, sizes, segment_operands],
            result_types=[AsyncTokenAttr()],
            regions=[body],
        )

    @classmethod
    def parse(cls, parser: Parser) -> SegmentOp:
        sym_name = parser.parse_optional_symbol_name()
        async_dependencies: list[Operation | SSAValue] = []
        sizes: list[Operation | SSAValue] = []
        parser.parse_optional_keyword("async")
        # TODO: unclear from the tests how to parse async. Follow the C++ code for the original custom parser
        if parser.parse_optional_keyword("unroll"):
            pass  # TODO: unclear from the tests how to parse unroll. Follow the C++ code for the original custom parser
        arguments_lst: list[Parser.Argument] = []
        segment_operands_lst: list[Operation | SSAValue] = []
        if parser.parse_optional_keyword("args"):
            if parser.parse_optional_characters("("):
                while True:
                    argument = parser.parse_argument(expect_type=False)
                    parser.parse_characters("=")
                    operand = parser.parse_operand()

                    argument = argument.resolve(operand.type)
                    arguments_lst.append(argument)
                    segment_operands_lst.append(operand)

                    if not parser.parse_optional_characters(","):
                        break
                parser.parse_characters(")")

            parser.parse_characters(":")

            for _ in range(len(segment_operands_lst)):
                parser.parse_type()
                parser.parse_optional_characters(",")

        attr_dict: dict[str, Attribute] = dict()
        if parser.parse_optional_keyword("attributes"):
            attr_dict = parser.parse_optional_attr_dict()

        body = parser.parse_optional_region()

        segment_op = SegmentOp(
            sym_name, async_dependencies, sizes, segment_operands_lst, body
        )
        segment_op.attributes |= attr_dict
        return segment_op

name = 'air.segment' class-attribute instance-attribute

sym_name = opt_prop_def(StringAttr) class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr()) class-attribute instance-attribute

sizes = var_operand_def(IndexType()) class-attribute instance-attribute

segment_operands = var_operand_def() class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr) class-attribute instance-attribute

body = opt_region_def() class-attribute instance-attribute

traits = traits_def(IsolatedFromAbove(), SingleBlockImplicitTerminator(SegmentTerminatorOp)) class-attribute instance-attribute

irdl_options = (AttrSizedOperandSegments(as_property=True),) class-attribute instance-attribute

__init__(sym_name: StringAttr | None, async_dependencies: list[Operation | SSAValue], sizes: list[Operation | SSAValue], segment_operands: list[Operation | SSAValue], body: Region | None)

Source code in xdsl/dialects/experimental/air.py
def __init__(
    self,
    sym_name: StringAttr | None,
    async_dependencies: list[Operation | SSAValue],
    sizes: list[Operation | SSAValue],
    segment_operands: list[Operation | SSAValue],
    body: Region | None,
):
    super().__init__(
        attributes={"sym_name": sym_name},
        operands=[async_dependencies, sizes, segment_operands],
        result_types=[AsyncTokenAttr()],
        regions=[body],
    )

parse(parser: Parser) -> SegmentOp classmethod

Source code in xdsl/dialects/experimental/air.py
@classmethod
def parse(cls, parser: Parser) -> SegmentOp:
    sym_name = parser.parse_optional_symbol_name()
    async_dependencies: list[Operation | SSAValue] = []
    sizes: list[Operation | SSAValue] = []
    parser.parse_optional_keyword("async")
    # TODO: unclear from the tests how to parse async. Follow the C++ code for the original custom parser
    if parser.parse_optional_keyword("unroll"):
        pass  # TODO: unclear from the tests how to parse unroll. Follow the C++ code for the original custom parser
    arguments_lst: list[Parser.Argument] = []
    segment_operands_lst: list[Operation | SSAValue] = []
    if parser.parse_optional_keyword("args"):
        if parser.parse_optional_characters("("):
            while True:
                argument = parser.parse_argument(expect_type=False)
                parser.parse_characters("=")
                operand = parser.parse_operand()

                argument = argument.resolve(operand.type)
                arguments_lst.append(argument)
                segment_operands_lst.append(operand)

                if not parser.parse_optional_characters(","):
                    break
            parser.parse_characters(")")

        parser.parse_characters(":")

        for _ in range(len(segment_operands_lst)):
            parser.parse_type()
            parser.parse_optional_characters(",")

    attr_dict: dict[str, Attribute] = dict()
    if parser.parse_optional_keyword("attributes"):
        attr_dict = parser.parse_optional_attr_dict()

    body = parser.parse_optional_region()

    segment_op = SegmentOp(
        sym_name, async_dependencies, sizes, segment_operands_lst, body
    )
    segment_op.attributes |= attr_dict
    return segment_op

WaitAllOp

Bases: IRDLOperation

Source code in xdsl/dialects/experimental/air.py
@irdl_op_definition
class WaitAllOp(IRDLOperation):
    name = "air.wait_all"

    async_dependencies = var_operand_def(AsyncTokenAttr)
    async_token = result_def(AsyncTokenAttr)

    def __init__(self, async_dependencies: list[SSAValue] | None):
        super().__init__(operands=[async_dependencies], result_types=[AsyncTokenAttr()])

    def print(self, printer: Printer):
        printer.print_string(" async ")
        if self.async_dependencies:
            printer.print_list(self.async_dependencies, printer.print_ssa_value)

    @classmethod
    def parse(cls, parser: Parser) -> WaitAllOp:
        parser.parse_keyword("async")
        async_dependencies = parser.parse_optional_comma_separated_list(
            parser.Delimiter.SQUARE, parser.parse_operand
        )

        return WaitAllOp(async_dependencies)

name = 'air.wait_all' class-attribute instance-attribute

async_dependencies = var_operand_def(AsyncTokenAttr) class-attribute instance-attribute

async_token = result_def(AsyncTokenAttr) class-attribute instance-attribute

__init__(async_dependencies: list[SSAValue] | None)

Source code in xdsl/dialects/experimental/air.py
def __init__(self, async_dependencies: list[SSAValue] | None):
    super().__init__(operands=[async_dependencies], result_types=[AsyncTokenAttr()])

print(printer: Printer)

Source code in xdsl/dialects/experimental/air.py
def print(self, printer: Printer):
    printer.print_string(" async ")
    if self.async_dependencies:
        printer.print_list(self.async_dependencies, printer.print_ssa_value)

parse(parser: Parser) -> WaitAllOp classmethod

Source code in xdsl/dialects/experimental/air.py
@classmethod
def parse(cls, parser: Parser) -> WaitAllOp:
    parser.parse_keyword("async")
    async_dependencies = parser.parse_optional_comma_separated_list(
        parser.Delimiter.SQUARE, parser.parse_operand
    )

    return WaitAllOp(async_dependencies)