Ops

@irdl_op_definition
class GenericOp(LinalgStructuredOperation):
    name = "linalg.generic"

    # Trait attributes
    indexing_maps = prop_def(ArrayAttr[AffineMapAttr])
    iterator_types = prop_def(ArrayAttr[IteratorTypeAttr])
    doc = opt_prop_def(StringAttr)
    library_call = opt_prop_def(StringAttr)

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    def __init__(
        self,
        inputs: Sequence[SSAValue],
        outputs: Sequence[SSAValue],
        body: Region,
        indexing_maps: Sequence[AffineMapAttr] | ArrayAttr[AffineMapAttr],
        iterator_types: Sequence[Attribute] | ArrayAttr[Attribute],
        result_types: Sequence[Attribute] = (),
        doc: StringAttr | None = None,
        library_call: StringAttr | None = None,
    ) -> None:
        super().__init__(
            operands=[inputs, outputs],
            result_types=[result_types],
            properties={
                "indexing_maps": ArrayAttr(indexing_maps),
                "iterator_types": ArrayAttr(iterator_types),
                "doc": doc,
                "library_call": library_call,
            },
            regions=[body],
        )

    def get_indexing_maps(self) -> ArrayAttr[AffineMapAttr]:
        return self.indexing_maps

    def get_iterator_types(self) -> ArrayAttr[IteratorTypeAttr]:
        return self.iterator_types

    def print(self, printer: Printer):
        printer.print_string(" {indexing_maps = ")
        printer.print_attribute(self.indexing_maps)
        printer.print_string(", iterator_types = [")
        printer.print_list(
            self.iterator_types,
            lambda iterator_type: printer.print_string_literal(iterator_type.data),
        )
        printer.print_string("]")
        if self.doc:
            printer.print_string(", doc = ")
            printer.print_attribute(self.doc)
        if self.library_call:
            printer.print_string(", library_call = ")
            printer.print_attribute(self.library_call)
        printer.print_string("}")

        if self.inputs:
            printer.print_string(" ins(")
            printer.print_list(self.inputs, printer.print_ssa_value)
            printer.print_string(" : ")
            printer.print_list(self.inputs.types, printer.print_attribute)
            printer.print_string(")")

        if self.outputs:
            printer.print_string(" outs(")
            printer.print_list(self.outputs, printer.print_ssa_value)
            printer.print_string(" : ")
            printer.print_list(self.outputs.types, printer.print_attribute)
            printer.print_string(")")

        extra_attrs = self.attributes.copy()
        if "indexing_maps" in extra_attrs:
            del extra_attrs["indexing_maps"]
        if "iterator_types" in extra_attrs:
            del extra_attrs["iterator_types"]
        if "doc" in extra_attrs:
            del extra_attrs["doc"]
        if "library_call" in extra_attrs:
            del extra_attrs["library_call"]

        if extra_attrs:
            printer.print_string(" attrs = ")
            printer.print_op_attributes(extra_attrs)

        printer.print_string(" ")
        printer.print_region(self.body)

        if self.res:
            printer.print_string(" -> ")
            if len(self.res) == 1:
                printer.print_attribute(self.res[0].type)
            else:
                with printer.in_parens():
                    printer.print_list(
                        self.res, lambda res: printer.print_attribute(res.type)
                    )

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        attrs_start_pos = parser.pos
        attrs = parser.parse_optional_attr_dict()
        attrs_end_pos = parser.pos

        if "indexing_maps" in attrs:
            indexing_maps = attrs["indexing_maps"]
            assert isinstance(indexing_maps, ArrayAttr)
            indexing_maps = cast(ArrayAttr[AffineMapAttr], indexing_maps)
            del attrs["indexing_maps"]
        else:
            parser.raise_error(
                "Expected indexing_maps for linalg.generic",
                attrs_start_pos,
                attrs_end_pos,
            )

        if "iterator_types" in attrs:
            # Get iterator types and make sure they're an ArrayAttr
            parsed_iterator_types = attrs["iterator_types"]
            assert isinstance(parsed_iterator_types, ArrayAttr)
            parsed_iterator_types = cast(ArrayAttr[Attribute], parsed_iterator_types)
            del attrs["iterator_types"]

            # Make sure they're iterator types
            iterator_types: list[IteratorTypeAttr] = []
            for iterator_type in parsed_iterator_types:
                match iterator_type:
                    case IteratorTypeAttr():
                        iterator_types.append(iterator_type)
                    case StringAttr():
                        iterator_type = IteratorTypeAttr(
                            IteratorType(iterator_type.data)
                        )
                        iterator_types.append(iterator_type)
                    case _:
                        parser.raise_error(
                            f"Unknown iterator type {iterator_type}",
                            attrs_start_pos,
                            attrs_end_pos,
                        )
        else:
            parser.raise_error(
                "Expected iterator_types for linalg.generic",
                attrs_start_pos,
                attrs_end_pos,
            )

        if "doc" in attrs:
            doc = attrs["doc"]
            assert isinstance(doc, StringAttr)
            del attrs["doc"]
        else:
            doc = None

        if "library_call" in attrs:
            library_call = attrs["library_call"]
            assert isinstance(library_call, StringAttr)
            del attrs["library_call"]
        else:
            library_call = None

        pos = parser.pos
        if parser.parse_optional_characters("ins"):
            parser.parse_punctuation("(")
            unresolved_ins = parser.parse_comma_separated_list(
                Parser.Delimiter.NONE, parser.parse_unresolved_operand
            )
            parser.parse_punctuation(":")
            ins_types = parser.parse_comma_separated_list(
                Parser.Delimiter.NONE, parser.parse_type
            )
            parser.parse_punctuation(")")
            ins = parser.resolve_operands(unresolved_ins, ins_types, pos)
        else:
            ins = ()

        pos = parser.pos
        if parser.parse_optional_characters("outs"):
            parser.parse_punctuation("(")
            unresolved_outs = parser.parse_comma_separated_list(
                Parser.Delimiter.NONE, parser.parse_unresolved_operand
            )
            parser.parse_punctuation(":")
            outs_types = parser.parse_comma_separated_list(
                Parser.Delimiter.NONE, parser.parse_type
            )
            parser.parse_punctuation(")")
            outs = parser.resolve_operands(unresolved_outs, outs_types, pos)
        else:
            outs = ()

        if parser.parse_optional_keyword("attrs"):
            parser.parse_punctuation("=")
            extra_attrs = parser.expect(
                parser.parse_optional_attr_dict, "expect extra attributes"
            )
        else:
            extra_attrs = {}

        body = parser.parse_region()

        if parser.parse_optional_punctuation("->"):
            res_types = parser.parse_comma_separated_list(
                parser.Delimiter.NONE, parser.parse_attribute
            )
        else:
            res_types = ()

        generic = cls(
            ins,
            outs,
            body,
            indexing_maps,
            iterator_types,
            res_types,
            doc,
            library_call,
        )
        generic.attributes |= extra_attrs

        return generic

@irdl_op_definition
class YieldOp(AbstractYieldOperation[Attribute]):
    name = "linalg.yield"

    traits = traits_def(IsTerminator())

@irdl_op_definition
class IndexOp(IRDLOperation):
    name = "linalg.index"

    dim = prop_def(IntegerAttr[i64])

    result = result_def(IndexTypeConstr)

    traits = traits_def(HasParent(GenericOp))

    assembly_format = "$dim attr-dict `:` type($result)"

    def __init__(
        self,
        dim: int,
    ):
        dim_attr = IntegerAttr(dim, i64)
        super().__init__(properties={"dim": dim_attr}, result_types=[IndexType()])

@irdl_op_definition
class AddOp(ElementwiseOperation):
    """
    Adds two tensors elementwise.

    See external [documentation](https://mlir.llvm.org/docs/Dialects/Linalg/#linalgadd-linalgaddop).
    """

    name = "linalg.add"

    def __init__(
        self,
        inputs: Sequence[SSAValue],
        outputs: Sequence[SSAValue] = (),
        res: Sequence[Attribute] | None = None,
        attributes: dict[str, Attribute] | None = None,
    ):
        if res is None:
            result_types = tuple(output.type for output in outputs)
        else:
            result_types = res

        super().__init__(
            ins=inputs,
            outs=outputs,
            result_types=result_types,
            attributes=attributes,
            hidden_region=self.get_hidden_region(inputs, outputs),
        )

    @classmethod
    def get_hidden_region(
        cls, inputs: Sequence[SSAValue], outputs: Sequence[SSAValue]
    ) -> Region:
        arg_types = cls.body_arg_types((*inputs, *outputs))
        add = arith.AddfOp if isinstance(arg_types[-1], AnyFloat) else arith.AddiOp

        @Builder.implicit_region(arg_types)
        def hidden_region(args: tuple[BlockArgument, ...]) -> None:
            result = add(args[0], args[1])
            YieldOp(result)

        return hidden_region

@irdl_op_definition
class ExpOp(ElementwiseOperation, NamedOperation):
    """
    Applies exp(x) elementwise.

    See external [documentation](https://mlir.llvm.org/docs/Dialects/Linalg/#linalgexp-linalgexpop).
    """

    name = "linalg.exp"

    def __init__(
        self,
        inputs: Sequence[SSAValue],
        outputs: Sequence[SSAValue] = (),
        res: Sequence[Attribute] | None = None,
        attributes: dict[str, Attribute] | None = None,
    ):
        if res is None:
            result_types = tuple(output.type for output in outputs)
        else:
            result_types = res

        super().__init__(
            ins=inputs,
            outs=outputs,
            result_types=result_types,
            attributes=attributes,
            hidden_region=self.get_hidden_region(inputs, outputs),
        )

    @classmethod
    def get_hidden_region(
        cls, inputs: Sequence[SSAValue], outputs: Sequence[SSAValue]
    ) -> Region:
        arg_types = cls.body_arg_types((*inputs, *outputs))

        @Builder.implicit_region(arg_types)
        def hidden_region(args: tuple[BlockArgument, ...]) -> None:
            result = math.ExpOp(args[0])
            YieldOp(result)

        return hidden_region

@irdl_op_definition
class LogOp(ElementwiseOperation, NamedOperation):
    """
    Applies log(x) elementwise.

    See external [documentation](https://mlir.llvm.org/docs/Dialects/Linalg/#linalglog-linalglogop).
    """

    name = "linalg.log"

    def __init__(
        self,
        inputs: Sequence[SSAValue],
        outputs: Sequence[SSAValue] = (),
        res: Sequence[Attribute] | None = None,
        attributes: dict[str, Attribute] | None = None,
    ):
        if res is None:
            result_types = tuple(output.type for output in outputs)
        else:
            result_types = res

        super().__init__(
            ins=inputs,
            outs=outputs,
            result_types=result_types,
            attributes=attributes,
            hidden_region=self.get_hidden_region(inputs, outputs),
        )

    @classmethod
    def get_hidden_region(
        cls, inputs: Sequence[SSAValue], outputs: Sequence[SSAValue]
    ) -> Region:
        arg_types = cls.body_arg_types((*inputs, *outputs))

        @Builder.implicit_region(arg_types)
        def hidden_region(args: tuple[BlockArgument, ...]) -> None:
            result = math.LogOp(args[0])
            YieldOp(result)

        return hidden_region

@irdl_op_definition
class SubOp(ElementwiseOperation):
    """
    Subtracts two tensors elementwise.

    See external [documentation](https://mlir.llvm.org/docs/Dialects/Linalg/#linalgsub-linalgsubop).
    """

    name = "linalg.sub"

    def __init__(
        self,
        inputs: Sequence[SSAValue],
        outputs: Sequence[SSAValue] = (),
        res: Sequence[Attribute] | None = None,
        attributes: dict[str, Attribute] | None = None,
    ):
        if res is None:
            result_types = tuple(output.type for output in outputs)
        else:
            result_types = res

        self.body_arg_types((*inputs, *outputs))

        super().__init__(
            ins=inputs,
            outs=outputs,
            result_types=result_types,
            attributes=attributes,
            hidden_region=self.get_hidden_region(inputs, outputs),
        )

    @classmethod
    def get_hidden_region(
        cls, inputs: Sequence[SSAValue], outputs: Sequence[SSAValue]
    ) -> Region:
        arg_types = cls.body_arg_types((*inputs, *outputs))
        sub = arith.SubfOp if isinstance(arg_types[-1], AnyFloat) else arith.SubiOp

        @Builder.implicit_region(arg_types)
        def hidden_region(args: tuple[BlockArgument, ...]) -> None:
            result = sub(args[0], args[1])
            YieldOp(result)

        return hidden_region

@irdl_op_definition
class SqrtOp(ElementwiseOperation, NamedOperation):
    """
    Applies sqrt(x) elementwise.

    See external [documentation](https://mlir.llvm.org/docs/Dialects/Linalg/#linalgsqrt-linalgsqrtop).
    """

    name = "linalg.sqrt"

    def __init__(
        self,
        inputs: Sequence[SSAValue],
        outputs: Sequence[SSAValue] = (),
        res: Sequence[Attribute] | None = None,
        attributes: dict[str, Attribute] | None = None,
    ):
        if res is None:
            result_types = tuple(output.type for output in outputs)
        else:
            result_types = res

        super().__init__(
            ins=inputs,
            outs=outputs,
            result_types=result_types,
            attributes=attributes,
            hidden_region=self.get_hidden_region(inputs, outputs),
        )

    @classmethod
    def get_hidden_region(
        cls, inputs: Sequence[SSAValue], outputs: Sequence[SSAValue]
    ) -> Region:
        arg_types = cls.body_arg_types((*inputs, *outputs))

        @Builder.implicit_region(arg_types)
        def hidden_region(args: tuple[BlockArgument, ...]) -> None:
            result = math.SqrtOp(args[0])
            YieldOp(result)

        return hidden_region