Tensor

@irdl_op_definition
class CastOp(IRDLOperation):
    """
    Tensor cast operation.

    Convert a tensor from one type to an equivalent type without changing any data elements.
    The source and destination types must both be tensor types with the same element type.
    If both are ranked, then the rank should be the same and static dimensions should match.
    The operation is invalid if converting to a mismatching constant dimension.

    https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorcast-tensorcastop
    """

    name = "tensor.cast"

    source = operand_def(
        base(TensorType[Attribute]) | base(UnrankedTensorType[Attribute])
    )
    dest = result_def(base(TensorType[Attribute]) | base(UnrankedTensorType[Attribute]))

    assembly_format = "$source attr-dict `:` type($source) `to` type($dest)"

    traits = traits_def(Pure())

    def __init__(self, source: SSAValue | Operation, dest: TensorType[Attribute]):
        super().__init__(operands=(source,), result_types=(dest,))

    def verify_(self):
        source_type = self.source.type
        dest_type = self.dest.type

        if isinstance(source_type, TensorType) and isinstance(dest_type, TensorType):
            # rank should be the same + constant shapes equal
            if len(source_type.get_shape()) != (len(dest_type.get_shape())):
                raise VerifyException("source and destination rank should be the same")
            for a, b in zip(source_type.get_shape(), dest_type.get_shape()):
                if a >= 0 and b >= 0 and a != b:
                    raise VerifyException(
                        "source and destination constant dimensions should match"
                    )

@irdl_op_definition
class ConcatOp(IRDLOperation):
    """
    Tensor concatenation operation.

    The “concat” operation constructs a tensor out of a variadic list of input tensors,
    concatenated along a static dimension number. All inputs and the result type must share the
    same rank.

    ``dim`` specifies the dimension along which to concatenate. The size of the concatenated
    dimension in the result must be equal to the sum of the sizes of the inputs along that
    dimension. All other dimensions in both the inputs and result must be the same size.

    https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorconcat-tensorconcatop
    """

    name = "tensor.concat"
    _TENSOR_ELEMENT: ClassVar[AttrConstraint] = VarConstraint(
        "Tensor Element", AnyAttr()
    )
    _RANK: ClassVar[IntConstraint] = IntVarConstraint("Rank", AtLeast(1))
    inputs = var_operand_def(
        RangeOf(
            TensorType.constr(
                _TENSOR_ELEMENT, ArrayOfConstraint(RangeOf(AnyAttr()).of_length(_RANK))
            )
        ).of_length(AtLeast(1))
    )
    dim = prop_def(IntegerAttr[I64])
    result = result_def(
        TensorType.constr(
            _TENSOR_ELEMENT, ArrayOfConstraint(RangeOf(AnyAttr()).of_length(_RANK))
        )
    )

    traits = traits_def(Pure())

    assembly_format = (
        " `dim` `(` $dim `)` $inputs attr-dict `:` functional-type(operands, results)"
    )

    def __init__(
        self,
        inputs: Sequence[SSAValue | Operation],
        dim: IntegerAttr | int,
        result_type: TensorType,
        attributes: Mapping[str, Attribute] | None = None,
    ):
        if isinstance(dim, int):
            dim = IntegerAttr(dim, i64)
        super().__init__(
            operands=(inputs,),
            result_types=(result_type,),
            properties={"dim": dim},
            attributes=attributes,
        )

    @staticmethod
    def _expected_concatenated_dim(dims: tuple[int]) -> int:
        """Return the expected length of concatenated dimension lengths for verifying the result type."""
        return DYNAMIC_INDEX if DYNAMIC_INDEX in dims else sum(dims)

    @staticmethod
    def _verify_and_get_non_concatenated_dim(dims: tuple[int], dim_index: int) -> int:
        """
        Raise a VerifyException if the lengths are inconsistent - ie. non-dynamic lengths are not all equal - else return the
        expected length of the non-concatenated dimension.
        """
        dim = DYNAMIC_INDEX
        for arg_dim in dims:
            if dim == DYNAMIC_INDEX:
                dim = arg_dim
            elif arg_dim != DYNAMIC_INDEX and dim != arg_dim:
                raise VerifyException(
                    f"static concatenation size mismatch along non-concatenated dimension {dim_index}"
                )
        return dim

    def _verify_result_shape(self, inferred_shape: tuple[int, ...]) -> None:
        """
        Verify the tensor shape of ``self.result.type`` by comparing it to the given inferred shape.

        Raises a VerifyException is any dimension that is not ``DYNAMIC_INDEX`` in either ``inferred_shape`` or the result type
        shape is unequal between the inferred size and the result type shape's dimension size.

        Expects that the length of inferred_shape is equal to the rank of the result type's shape, and that `self.result.type` is
        a TensorType."""
        for inferred_size, actual_size in zip(
            inferred_shape, self.result.type.shape, strict=True
        ):
            if (
                DYNAMIC_INDEX not in (inferred_size, actual_size.data)
                and inferred_size != actual_size.data
            ):
                raise VerifyException(
                    f"result type {self.result.type} does not match inferred shape {inferred_shape} static sizes"
                )

    def verify_(self) -> None:
        concat_dim = self.dim.value.data
        result_type = self.result.type
        if concat_dim >= result_type.get_num_dims():
            raise VerifyException("concatenation dim must be less than the tensor rank")

        transposed_shapes = zip(
            *(
                (int_attr.data for int_attr in cast(TensorType, arg.type).shape)
                for arg in self.inputs
            )
        )
        expected_result_shape = tuple(
            self._expected_concatenated_dim(grouped_dims)
            if i == concat_dim
            else self._verify_and_get_non_concatenated_dim(grouped_dims, i)
            for i, grouped_dims in enumerate(transposed_shapes)
        )
        self._verify_result_shape(expected_result_shape)

@irdl_op_definition
class DimOp(IRDLOperation):
    """
    Dimension index operation.

    The tensor.dim operation takes a tensor and a dimension operand of type index.
    It returns the size of the requested dimension of the given tensor.
    If the dimension index is out of bounds, the behavior is undefined.

    https://mlir.llvm.org/docs/Dialects/TensorOps/#tensordim-tensordimop
    """

    name = "tensor.dim"

    source = operand_def(
        base(TensorType[Attribute]) | base(UnrankedTensorType[Attribute])
    )
    index = operand_def(IndexType)
    result = result_def(IndexType)

    traits = traits_def(NoMemoryEffect())

    assembly_format = "attr-dict $source `,` $index `:` type($source)"

    def __init__(
        self,
        source: SSAValue | Operation,
        index: SSAValue | Operation,
        attributes: Mapping[str, Attribute] | None = None,
    ):
        super().__init__(
            operands=(source, index), result_types=(IndexType(),), attributes=attributes
        )

    def verify_(self):
        if isinstance((source_type := self.source.type), TensorType):
            if not len(source_type.get_shape()):
                raise VerifyException("cannot get dim of 0-rank tensor")

@irdl_op_definition
class EmptyOp(IRDLOperation):
    """
    Empty tensor operation.

    Defines a tensor of a particular shape which could be dynamic or static.
    The contents of the tensor are unspecified and the only purpose of the op
    result is to materialize the specified shape in IR and make it available
    to other transformations.

    https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorempty-tensoremptyop
    """

    name = "tensor.empty"

    dynamic_sizes = var_operand_def(IndexType)

    tensor = result_def(TensorType[Attribute])

    traits = traits_def(NoMemoryEffect())

    def __init__(self, dynamic_sizes: Sequence[SSAValue], tensor_type: Attribute):
        super().__init__(
            operands=(dynamic_sizes,),
            result_types=(tensor_type,),
        )

    def print(self, printer: Printer):
        if self.dynamic_sizes:
            printer.print_string("(")
            printer.print_list(self.dynamic_sizes, printer.print_ssa_value)
            printer.print_string(")")
        else:
            printer.print_string("(")
            printer.print_string(")")

        printer.print_string(" : ")
        printer.print_attribute(self.tensor.type)

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        pos = parser.pos
        parser.parse_punctuation("(")
        if parser.parse_optional_punctuation(")"):
            dynamic_sizes = ()
        else:
            unresolved_dynamic_sizes = parser.parse_comma_separated_list(
                Parser.Delimiter.NONE, parser.parse_unresolved_operand
            )
            unresolved_types = (IndexType(),) * len(unresolved_dynamic_sizes)
            parser.parse_punctuation(")")
            dynamic_sizes = parser.resolve_operands(
                unresolved_dynamic_sizes, unresolved_types, pos
            )
        parser.parse_punctuation(":")
        result_type = parser.parse_attribute()

        empty = cls(dynamic_sizes, result_type)

        return empty

@irdl_op_definition
class CollapseShapeOp(IRDLOperation):
    """
    Operation to produce a tensor with a smaller rank.

    The collapse_shape operation produces a new tensor of lower (or equal)
    rank whose dimension sizes are a reassociation of the original src dimensions.

    https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorcollapse_shape-tensorcollapseshapeop
    """

    name = "tensor.collapse_shape"

    src = operand_def(TensorType[Attribute])
    result = result_def(TensorType[Attribute])
    reassociation = prop_def(ContiguousArrayOfIntArray())
    assembly_format = (
        "$src $reassociation attr-dict `:` type($src) `into` type($result)"
    )

    traits = traits_def(Pure())

@irdl_op_definition
class ReshapeOp(IRDLOperation):
    """
    Tensor reshape operation.

    The reshape operation converts a tensor from one type to an equivalent
    type with a provided shape. The source and destination types are compatible
    if both have the same element type, same number of elements.

    https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorreshape-tensorreshapeop
    """

    name = "tensor.reshape"

    source = operand_def(TensorType[Attribute])
    shape = operand_def(TensorType[AnySignlessIntegerType | IndexType])
    result = result_def(TensorType[Attribute])
    assembly_format = "attr-dict $source `(` $shape `)` `:` `(` type($source) `,` type($shape) `)` `->` type($result)"

    traits = traits_def(Pure())

    def __init__(self, source: SSAValue, shape: SSAValue, result_type: Attribute):
        super().__init__(
            operands=(
                source,
                shape,
            ),
            result_types=(result_type,),
        )

    def verify_(self) -> None:
        if not isinstance(
            source_type := self.source.type, TensorType
        ) or not isinstance(shape_type := self.shape.type, TensorType):
            raise ValueError(
                "tensor elementwise operation operands and result must be of type TensorType"
            )

        source_type = cast(TensorType[Attribute], source_type)
        shape_type = cast(TensorType[Attribute], shape_type)
        res_type = self.result.type

        if source_type.element_type != res_type.element_type:
            raise VerifyException(
                "element types of source and result tensor types should be the same"
            )

        source_type = source_type.get_shape()
        shape_type = shape_type.get_shape()
        res_type = res_type.get_shape()

        if len(shape_type) != 1:
            raise VerifyException("shape tensor must have a rank one")

        # concerns the case of static reshaping
        if math.prod(source_type) != math.prod(res_type):
            raise VerifyException(
                "source and result tensor should have the same number of elements"
            )

        shape_size = shape_type[0]
        if shape_size != len(res_type):
            raise VerifyException(
                "length of shape operand differs from the result's tensor rank"
            )

@irdl_op_definition
class ExpandShapeOp(IRDLOperation):
    """
    Operation to produce a tensor with a higher rank.

    The tensor.expand_shape op produces a tensor of higher (or equal)
    rank than the operand src whose dimension sizes are a reassociation of src.

    https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorexpand_shape-tensorexpandshapeop
    """

    name = "tensor.expand_shape"

    src = operand_def(TensorType)
    dynamic_output_shape = var_operand_def(IndexType)

    reassociation = prop_def(ContiguousArrayOfIntArray())

    static_output_shape = prop_def(DenseArrayBase.constr(i64))

    result = result_def(TensorType[Attribute])

    traits = traits_def(Pure())

    def __init__(
        self,
        src: SSAValue | Operation,
        dynamic_output_shape: Sequence[SSAValue],
        reassociation: ArrayAttr[ArrayAttr[IntegerAttr]],
        static_output_shape: Sequence[int] | DenseArrayBase,
        result_type: TensorType[Attribute],
        attributes: dict[str, Attribute] | None = None,
    ):
        if not isinstance(static_output_shape, DenseArrayBase):
            static_output_shape = DenseArrayBase.from_list(i64, static_output_shape)

        super().__init__(
            operands=[src, dynamic_output_shape],
            result_types=[result_type],
            properties={
                "reassociation": reassociation,
                "static_output_shape": static_output_shape,
            },
            attributes=attributes,
        )

    def verify_(self):
        assert isinstance(self.src.type, ShapedType)
        assert isinstance(self.result.type, ShapedType)

        # make sure the static output shape matches the result type
        if len(self.static_output_shape) != len(self.result.type.get_shape()):
            raise VerifyException(
                "expected number of static shape dims to be equal to the output rank "
                f"({len(self.result.type.get_shape())}) but found {len(self.static_output_shape)} inputs instead"
            )

        verify_reshape_like_types(
            collapsed_type=self.src.type,
            expanded_type=self.result.type,
            reassociation=self.reassociation,
        )

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        src_operand = parser.parse_unresolved_operand()

        reassociation = parser.parse_attribute()
        parser.parse_characters("output_shape")
        index = IndexType()

        # Parse shape: mixture of ints and SSA values
        dyn_shape, static_shape = parse_dynamic_index_list_without_types(
            parser, dynamic_index=DYNAMIC_INDEX
        )

        dyn_shape = parser.resolve_operands(
            dyn_shape, (index,) * len(dyn_shape), parser.pos
        )

        attributes = parser.parse_optional_attr_dict()

        parser.parse_punctuation(":")
        src_type = parser.parse_type()
        parser.parse_characters("into")
        result_type = parser.parse_type()
        src = parser.resolve_operand(src_operand, src_type)

        shape_attr = DenseArrayBase.from_list(i64, static_shape)

        reassociation = cast(ArrayAttr[ArrayAttr[IntegerAttr]], reassociation)
        result_type = cast(TensorType[Attribute], result_type)

        return cls(src, dyn_shape, reassociation, shape_attr, result_type, attributes)

    def print(self, printer: Printer):
        printer.print_string(" ")
        printer.print_ssa_value(self.src)
        printer.print_string(" ")
        printer.print_attribute(self.reassociation)
        printer.print_string(" output_shape ")
        print_dynamic_index_list(
            printer,
            DYNAMIC_INDEX,
            self.dynamic_output_shape,
            self.static_output_shape.get_values(),
        )

        printer.print_op_attributes(attributes=self.attributes)

        printer.print_string(" : ")
        printer.print_attribute(self.src.type)
        printer.print_string(" into ")
        printer.print_attribute(self.result.type)

@irdl_op_definition
class ExtractSliceOp(IRDLOperation):
    """
    Extract slice operation.

    Extracts a tensor from another tensor as specified by the operation’s
    offsets, sizes and strides arguments.

    https://mlir.llvm.org/docs/Dialects/TensorOps/#tensorextract_slice-tensorextractsliceop
    """

    name = "tensor.extract_slice"

    source = operand_def(TensorType)
    offsets = var_operand_def(IndexType)
    sizes = var_operand_def(IndexType)
    strides = var_operand_def(IndexType)
    static_offsets = prop_def(DenseArrayBase.constr(i64))
    static_sizes = prop_def(DenseArrayBase.constr(i64))
    static_strides = prop_def(DenseArrayBase.constr(i64))
    result = result_def(TensorType)

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    traits = traits_def(Pure())

    assembly_format = (
        "$source ``"
        "custom<DynamicIndexList>($offsets, $static_offsets)"
        "custom<DynamicIndexList>($sizes, $static_sizes)"
        "custom<DynamicIndexList>($strides, $static_strides)"
        "attr-dict `:` type($source) `to` type($result)"
    )

    custom_directives = (DynamicIndexList,)

    def verify_(self) -> None:
        verify_dynamic_index_list(
            self.static_offsets.get_values(),
            self.offsets,
            DYNAMIC_INDEX,
            " in the offset arguments",
        )
        verify_dynamic_index_list(
            self.static_sizes.get_values(),
            self.sizes,
            DYNAMIC_INDEX,
            " in the size arguments",
        )
        verify_dynamic_index_list(
            self.static_strides.get_values(),
            self.strides,
            DYNAMIC_INDEX,
            " in the stride arguments",
        )

    def __init__(
        self,
        source: SSAValue | Operation,
        offsets: Sequence[SSAValue | int],
        sizes: Sequence[SSAValue | int],
        strides: Sequence[SSAValue | int],
        result_type: Attribute,
    ):
        """
        Build a tensor.extract_slice from offsets, sizes, and strides that may each
        be a static integer or a dynamic SSA value.
        """
        static_offsets, dyn_offsets = split_dynamic_index_list(offsets, DYNAMIC_INDEX)
        static_sizes, dyn_sizes = split_dynamic_index_list(sizes, DYNAMIC_INDEX)
        static_strides, dyn_strides = split_dynamic_index_list(strides, DYNAMIC_INDEX)

        super().__init__(
            operands=[source, dyn_offsets, dyn_sizes, dyn_strides],
            result_types=[result_type],
            properties={
                "static_offsets": DenseArrayBase.from_list(i64, static_offsets),
                "static_sizes": DenseArrayBase.from_list(i64, static_sizes),
                "static_strides": DenseArrayBase.from_list(i64, static_strides),
            },
        )

    @staticmethod
    def infer_result_type(
        source_type: TensorType[Attribute],
        sizes: Sequence[SSAValue | int],
    ) -> TensorType[Attribute]:
        """
        Infer the result type of a tensor.extract_slice from its source type and
        sizes. Offsets and strides do not affect the result shape.

        Raises:
            ValueError: If sizes do not match the source rank.
        """
        rank = source_type.get_num_dims()
        if len(sizes) != rank:
            raise ValueError("expected sizes to match source rank")

        result_shape = tuple(
            size if isinstance(size, int) else DYNAMIC_INDEX for size in sizes
        )

        return TensorType(
            source_type.get_element_type(),
            result_shape,
            source_type.encoding,
        )

    @staticmethod
    def from_static_parameters(
        source: SSAValue | Operation,
        offsets: Sequence[int],
        sizes: Sequence[int],
        strides: Sequence[int] | None = None,
        reduce_rank: bool = False,
    ) -> ExtractSliceOp:
        if strides is None:
            strides = [1] * len(offsets)
        source_v = SSAValue.get(source, type=TensorType)
        source_t = source_v.type

        if reduce_rank:
            result_sizes = list(s for s in sizes if s != 1)
        else:
            result_sizes = list(sizes)

        return_type = TensorType(source_t.get_element_type(), result_sizes)

        return ExtractSliceOp.build(
            operands=[source, [], [], []],
            result_types=[return_type],
            properties={
                "static_offsets": DenseArrayBase.from_list(i64, offsets),
                "static_sizes": DenseArrayBase.from_list(i64, result_sizes),
                "static_strides": DenseArrayBase.from_list(i64, strides),
            },
        )