Abstract ops

class LoopBoundSource(NamedTuple):
    """Source information for one loop upper bound."""

    operand: SSAValue
    """The shaped operand that provides the bound."""

    dim_index: int
    """The dimension index in the operand used for the bound."""

    dim_size: int
    """The size of that dimension, or DYNAMIC_INDEX if it is dynamic."""

class LinalgStructuredOperation(IRDLOperation, ABC):
    """
    Abstract base class for structured linalg operations, allowing them to be processed
    via a unified interface.
    """

    inputs = var_operand_def()
    """
    The operands that won't be mutated.
    """
    outputs = var_operand_def(ShapedType)
    """
    The operands that will be accumulated into.
    These inputs may be `memref`s, which will be mutated in-place, or `tensor`s, which will be returned as results.
    """

    res = var_result_def(TensorType)
    """
    The updated `outputs`, empty if the inputs are memrefs.
    """

    body = region_def("single_block")
    """
    The body implementing the combination of scalar elements of the inputs, and
    yielding the scalar elements of the outputs.
    """

    @abstractmethod
    def get_indexing_maps(self) -> ArrayAttr[AffineMapAttr]:
        """
        Get the indexing maps corresponding to this operation's operands, in order.
        """

    @abstractmethod
    def get_iterator_types(self) -> ArrayAttr[IteratorTypeAttr]:
        """
        Get the iterator types corresponding to this operation's loop, in order.
        """

    def get_num_loops(self) -> int:
        return self.get_indexing_maps().data[0].data.num_dims

    def get_loops_to_shapes_map(self) -> AffineMap:
        """
        Returns a map to answer the question: "given an iteration space over
        the codomain, what are the subshapes of the operands involved in the
        computation".
        The default behavior is to just concatenate all the indexing maps.
        """
        indexing_maps = tuple(attr.data for attr in self.get_indexing_maps())
        result_exprs = tuple(res for map in indexing_maps for res in map.results)

        dims = self.get_num_loops()

        # FIXME: Support symbols.
        for map in indexing_maps:
            if map.num_symbols != 0:
                raise NotImplementedError(
                    "Indexing maps with symbols not supported for now."
                )

        syms = 0
        return AffineMap(dims, syms, result_exprs)

    def get_shapes_to_loops_map(self) -> AffineMap:
        """
        Returns a map to answer the question: "Given a list of operand ranges,
        what is the subportion of the iteration space involved in the
        computation". This is the inverse problem of `get_loops_to_shapes_map`.
        Return the empty AffineMap when such an AffineMap cannot be
        constructed. The default behavior is based on a very simple inference
        procedure that only works with permutation affine maps. A more advanced
        Tensor-Comprehension like inference is possible but has proven to be
        ambiguous in unfavorable case. A safer and more robust alternative is
        to allow each op to define its own AffineMap.
        """
        loops_to_shapes = self.get_loops_to_shapes_map()
        inverse = loops_to_shapes.inverse_permutation()
        if not inverse:
            raise NotImplementedError(
                "Non-invertible maps need dynamic shapes, which are not implemented."
            )
        return inverse

    def get_loop_bound_sources(
        self,
    ) -> tuple[LoopBoundSource, ...]:
        """
        Return where each loop upper bound comes from.

        Each entry identifies the shaped operand, the dimension index, and the size value.
        """
        shapes_to_loops = self.get_shapes_to_loops_map()

        needed_positions = tuple(
            expr.position
            for expr in shapes_to_loops.results
            if isinstance(expr, AffineDimExpr)
        )
        assert len(shapes_to_loops.results) == len(needed_positions)

        flat_shape_dims = tuple(
            LoopBoundSource(operand, dim_index, dim_size)
            for operand in self.operands
            if isa(operand, SSAValue[ShapedType])
            for dim_index, dim_size in enumerate(operand.type.get_shape())
        )

        return tuple(flat_shape_dims[position] for position in needed_positions)

    def get_static_shapes(self) -> list[int]:
        return [
            dim
            for operand in self.operands
            if isinstance(operand.type, ShapedType)
            for dim in operand.type.get_shape()
        ]

    def get_static_loop_ranges(self) -> tuple[int, ...]:
        shapes_to_loops = self.get_shapes_to_loops_map()
        static_shapes = self.get_static_shapes()
        return shapes_to_loops.eval(static_shapes, [])

    def verify_(self) -> None:
        # Operands are all tensors or all memrefs. Each kind is written back a
        # different way, and MLIR requires one or the other too.
        shaped_types = [
            operand.type
            for operand in self.operands
            if isinstance(operand.type, ShapedType)
        ]
        if any(isinstance(t, TensorType) for t in shaped_types) and any(
            isinstance(t, MemRefType) for t in shaped_types
        ):
            raise VerifyException("expected to have pure tensor or buffer semantics")

        indexing_maps = tuple(attr.data for attr in self.get_indexing_maps())
        if len(indexing_maps) != len(self.operands):
            raise VerifyException(
                f"expected the number of indexing_map ({len(indexing_maps)}) to be "
                f"equal to the number of input/output operands ({len(self.operands)})"
            )

        # Every indexing map is over the same loops, one per iterator type, and
        # has one result per dimension of its operand, which says how the loops
        # reach that dimension. A scalar has none.
        num_loops = len(self.get_iterator_types())
        for index, (operand, indexing_map) in enumerate(
            zip(self.operands, indexing_maps, strict=True)
        ):
            if indexing_map.num_dims != num_loops:
                raise VerifyException(
                    f"expected indexing_map #{index} to have {num_loops} dim(s) to "
                    f"match the number of loops"
                )
            if indexing_map.num_symbols:
                raise VerifyException(f"unexpected symbols in indexing_map #{index}")
            rank = (
                len(operand.type.get_shape())
                if isinstance(operand.type, ShapedType)
                else 0
            )
            if len(indexing_map.results) != rank:
                raise VerifyException(
                    f"expected operand #{index} of rank {rank} to match the "
                    f"result count of indexing_map #{index}, "
                    f"{len(indexing_map.results)}"
                )

        # The body takes one argument per operand, of that operand's element
        # type, or of the operand's own type for a scalar.
        block_args = self.body.block.args
        if len(block_args) != len(self.operands):
            raise VerifyException(
                "expected as many non-induction variable region arguments as the "
                f"number of input/output operands, {len(self.operands)}, but got "
                f"{len(block_args)}"
            )
        for index, (operand, arg) in enumerate(
            zip(self.operands, block_args, strict=True)
        ):
            operand_type = operand.type
            element_type: Attribute = (
                operand_type.get_element_type()
                if isa(operand_type, MemRefType | TensorType)
                else operand_type
            )
            if arg.type != element_type:
                raise VerifyException(
                    f"expected type of bb argument #{index} ({arg.type}) to match "
                    f"element or self type of the corresponding operand "
                    f"({element_type})"
                )

        # With nothing to index there is nothing to reach into.
        if not indexing_maps:
            return

        # The loop ranges are read back off the operand shapes, which needs each
        # loop to appear on its own as some result of some map.
        loops_to_shapes = self.get_loops_to_shapes_map()
        shapes_to_loops = loops_to_shapes.inverse_permutation()
        if shapes_to_loops is None:
            raise VerifyException(
                f"invalid indexing maps are non-invertible: ({loops_to_shapes})"
            )

        # A shape only known at runtime gives a loop range only known then, and
        # nothing can be checked against it here.
        static_shapes = self.get_static_shapes()
        if any(dim < 0 for dim in static_shapes):
            return
        end_ranges = tuple(
            bound - 1 for bound in shapes_to_loops.eval(static_shapes, [])
        )
        start_ranges = (0,) * len(end_ranges)

        # Each operand is as large as the loops reach into it. A result that is
        # one loop dimension reaches every index up to that loop's range, so the
        # operand dimension is exactly that. One that is an expression over the
        # loops, `d0 + d1` or `d0 * 2`, reaches indices that may not fill the
        # dimension, so it only has to fit.
        for index, (operand, indexing_map) in enumerate(
            zip(self.operands, indexing_maps, strict=True)
        ):
            if not isinstance(operand.type, ShapedType):
                continue
            shape = operand.type.get_shape()
            starts = indexing_map.eval(start_ranges, [])
            ends = indexing_map.eval(end_ranges, [])
            for dim, (start, end, size) in enumerate(
                zip(starts, ends, shape, strict=True)
            ):
                if size == 0:
                    continue
                if min(start, end) < 0:
                    raise VerifyException(
                        f"unexpected result less than 0 at expression #{dim} in "
                        f"{indexing_map}"
                    )
                inferred = max(start, end) + 1
                if isinstance(indexing_map.results[dim], AffineDimExpr):
                    if inferred != size:
                        raise VerifyException(
                            f"inferred input/output operand #{index} has shape's "
                            f"dimension #{dim} to be {inferred}, but found {size}"
                        )
                elif inferred > size:
                    raise VerifyException(
                        f"inferred input/output operand #{index} has shape's "
                        f"dimension #{dim} to be greater than or equal to "
                        f"{inferred}, but found {size}"
                    )

class NamedOperation(LinalgStructuredOperation, ABC):
    """
    Abstract base class for named ops with hidden region.
    """

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

    PRINT_ATTRS_IN_FRONT: ClassVar[bool] = False

    def __init__(
        self,
        ins: Sequence[SSAValue],
        outs: Sequence[SSAValue],
        result_types: Sequence[Attribute | Sequence[Attribute] | None] | None = None,
        properties: Mapping[str, Attribute | None] | None = None,
        attributes: Mapping[str, Attribute | None] | None = None,
        hidden_region: Region | None = None,
    ):
        super().__init__(
            operands=[ins, outs],
            result_types=(
                result_types
                if result_types is not None and len(result_types) > 0
                else [[]]
            ),
            properties=properties,
            attributes=attributes,
            regions=[hidden_region],
        )

    @classmethod
    def parse(cls, parser: Parser):
        pos = parser.pos
        if cls.PRINT_ATTRS_IN_FRONT:
            attrs = parser.parse_optional_attr_dict()
        else:
            attrs = {}
        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 not cls.PRINT_ATTRS_IN_FRONT:
            if parser.parse_optional_keyword("attrs"):
                parser.parse_punctuation("=")
                attrs = parser.expect(
                    parser.parse_optional_attr_dict, "expect extra attributes"
                )
            else:
                attrs = {}

        if parser.parse_optional_punctuation("->"):
            res_types = parser.parse_optional_comma_separated_list(
                parser.Delimiter.PAREN, parser.parse_attribute
            )
            if res_types is None:
                res_types = [parser.parse_attribute()]
        else:
            res_types = ()

        prop_names = cls.get_irdl_definition().properties

        properties = {k: v for k, v in attrs.items() if k in prop_names}
        # Drop the values in properties from attrs
        for k in properties:
            if k in attrs:
                del attrs[k]

        try:
            return cls.build(
                operands=(ins, outs),
                result_types=(res_types,),
                properties=properties,
                attributes=attrs,
                regions=(cls.get_hidden_region(ins, outs),),
            )
        except ValueError:
            parser.raise_error("Could not build linalg op")

    def print(self, printer: Printer):
        extra_attrs = {**self.attributes, **self.properties}
        extra_attrs.pop("indexing_maps", None)
        extra_attrs.pop("linalg.memoized_indexing_maps", None)
        extra_attrs.pop("iterator_types", None)
        extra_attrs.pop("doc", None)
        extra_attrs.pop("library_call", None)
        extra_attrs.pop("operandSegmentSizes", None)

        if extra_attrs and self.PRINT_ATTRS_IN_FRONT:
            printer.print_op_attributes(extra_attrs)
        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(")")

        if extra_attrs and not self.PRINT_ATTRS_IN_FRONT:
            printer.print_string(" attrs = ")
            printer.print_op_attributes(extra_attrs)

        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)
                    )

    @staticmethod
    def body_arg_types(
        operands: Sequence[SSAValue],
    ) -> Sequence[AnyFloat | IntegerType]:
        """
        Return the element types of the arguments of the body of this operation
        """

        result: Sequence[AnyFloat | IntegerType] = []

        for op in operands:
            op_type = op.type
            if isa(op_type, MemRefType | TensorType):
                element_type = op_type.get_element_type()
            else:  # int or float
                element_type = op_type
            assert isa(element_type, AnyFloat | IntegerType)
            result.append(element_type)

        return result

    @classmethod
    @abstractmethod
    def get_hidden_region(
        cls, inputs: Sequence[SSAValue], outputs: Sequence[SSAValue]
    ) -> Region:
        """
        The hidden region for this linalg NamedOperation.
        """
        raise NotImplementedError

    @abstractmethod
    def get_default_indexing_maps(self) -> Sequence[AffineMap]:
        """
        Get the default indexing maps corresponding to this operation's operands, in order.
        """

    def get_indexing_maps(self) -> ArrayAttr[AffineMapAttr]:
        return ArrayAttr(
            AffineMapAttr(map_) for map_ in self.get_default_indexing_maps()
        )