#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
#  See https://llvm.org/LICENSE.txt for license information.
#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Model classes representing a tensor comprehension.

These classes model the language more at an AST level as evaluated. Reasoning
about it typically involves processing this form into config objects that
represent actual op definitions (i.e. YAML).
"""

from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple
from enum import Enum

from ..... import ir as _ir
from .affine import *
from .scalar_expr import *
from .types import *
from .yaml_helper import *

###############################################################################
# Tensor expression nodes.
###############################################################################


class TensorExpression:
    """An expression that can appear on the RHS of a comprehension."""

    def to_scalar_expression(self) -> ScalarExpression:
        raise NotImplementedError()

    def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
        """Visits all tensor expression reachable by the expression."""
        callback(self)

    def collect_dim_uses(self, uses: Set["DimDef"]):
        """Collects all DimDefs reachable through this expression."""

        def visit_dim_def(dim_def: AffineExprDef):
            if isinstance(dim_def, DimDef):
                uses.add(dim_def)

        def visit_affine_exprs(expr: "TensorExpression"):
            if isinstance(expr, TensorUse):
                for ind in expr.indices:
                    ind.visit_affine_exprs(visit_dim_def)
            if isinstance(expr, TensorReduceFn):
                for ind in expr.reduce_fn.reduce_dims:
                    ind.visit_affine_exprs(visit_dim_def)

        self.visit_tensor_exprs(visit_affine_exprs)

    def collect_tensor_uses(self, uses: Set["TensorUse"]):
        """Collects all TensorUses reachable through this expression."""

        def visit_tensor_use(expr: "TensorExpression"):
            if isinstance(expr, TensorUse):
                uses.add(expr)

        self.visit_tensor_exprs(visit_tensor_use)

    def collect_indices(self, indices: Set["index"]):
        """Collects all index accesses reachable through this expression."""

        def visit_index(expr: "TensorExpression"):
            if isinstance(expr, index):
                indices.add(expr)

        self.visit_tensor_exprs(visit_index)

    def collect_scalar_uses(self, uses: Set["ScalarDef"]):
        """Collects all ScalarDefs reachable through this expression."""

        def visit_scalar_def(expr: "TensorExpression"):
            if isinstance(expr, ScalarDef):
                uses.add(expr)

        self.visit_tensor_exprs(visit_scalar_def)

    def __add__(self, rhs: "TensorExpression") -> "TensorExpression":
        return BinaryFn.add(self, rhs)

    def __mul__(self, rhs) -> "TensorExpression":
        return BinaryFn.mul(self, rhs)

    def __sub__(self, rhs) -> "TensorExpression":
        return BinaryFn.sub(self, rhs)

    def __truediv__(self, rhs) -> "TensorExpression":
        return BinaryFn.div(self, rhs)

    def __hash__(self):
        return hash(id(self))


class TensorUse(TensorExpression):
    """A used tensor represented by its (tensor_name, indices).

    Note that forming a comprehension via direct assignment is performed through
    __setitem__ on the TensorDef level. However, performing a reduction with
    compound ops (+=, *=, etc) is done by doing a:
      TensorDef.__getitem__
      TensorUse.__iadd__
      TensorDef.__setitem__
    """

    def __init__(self, operand_def: "OperandDef", indices: Sequence[AffineExprDef]):
        self.operand_def = operand_def
        self.indices = tuple(indices)

    def to_scalar_expression(self) -> ScalarExpression:
        return ScalarArg(self.tensor_name).expr()

    @property
    def tensor_name(self) -> str:
        name = self.operand_def.name
        assert name is not None, "TensorDef not registered with an op"
        return name

    def _compute_reduce_dims(self, rhs: TensorExpression) -> Set[DimDef]:
        # Computes the reduction dims for implicit reductions. Assumes that the rhs
        # is the expression being reduced and self is being reduced into. Any
        # indices referenced on the rhs and not in self are considered reduction
        # dims and will be ordered as encountered on the rhs.
        rhs_dims = set()
        lhs_dims = set()
        rhs.collect_dim_uses(rhs_dims)
        self.collect_dim_uses(lhs_dims)
        return rhs_dims - lhs_dims

    def __iadd__(self, rhs: TensorExpression) -> "TensorReduceFn":
        return ReduceFnUse(BinaryFn.add, None, *self._compute_reduce_dims(rhs))(rhs)

    def __repr__(self):
        return (
            f"{self.operand_def.name}" f"[{', '.join([repr(i) for i in self.indices])}]"
        )


class TensorFn(TensorExpression):
    """Application of a tensor function."""

    def __init__(
        self,
        kind: "FunctionKind",
        name: Optional[str],
        operand_def: Optional["OperandDef"],
        type_var: Optional[TypeVar],
        args: Sequence[TensorExpression],
    ):
        if bool(name) + bool(operand_def) != 1:
            raise ValueError("One of 'name', 'operand_def' must be specified")
        self.name = name
        self.kind = kind
        self.operand_def = operand_def
        self.type_var = type_var
        self.args = args

    def to_scalar_expression(self) -> ScalarExpression:
        if self.operand_def:
            assert self.operand_def.name, "TensorFn not registered with an op"
        attr_name = self.operand_def.name if self.operand_def else None
        args = [arg.to_scalar_expression() for arg in self.args]
        return ScalarFn(self.kind, self.name, attr_name, self.type_var, args).expr()

    def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
        super().visit_tensor_exprs(callback)
        for arg in self.args:
            arg.visit_tensor_exprs(callback)

    def __repr__(self):
        name = self.operand_def.name if self.operand_def else self.name
        return (
            f"{self.kind.name}.{name}(type_var={self.type_var}, "
            f"args={', '.join(repr(a) for a in self.args)})"
        )


class TensorReduceFn(TensorExpression):
    """Application of a reduction function.

    This captures the lhs (initial value) separately from the rhs.
    """

    def __init__(self, reduce_use: "ReduceFnUse", args: Sequence[TensorExpression]):
        self.reduce_use = reduce_use
        self.lhs = None  # type: Optional[TensorUse]
        self.args = args

    def to_scalar_expression(self) -> ScalarExpression:
        if self.lhs is None:
            raise ValueError(
                f"Cannot scalarize a TensorReduceFn that has not been "
                f"bound to its lhs: {self}"
            )
        full_args = [self.lhs.to_scalar_expression()] + [
            arg.to_scalar_expression() for arg in self.args
        ]
        fn_name = None
        attr_name = None
        if self.reduce_use.binary_fn:
            fn_name = self.reduce_use.binary_fn.fn_name
        if self.reduce_use.binary_attr:
            attr_name = self.reduce_use.binary_attr.operand_def.name
        return ScalarFn(FunctionKind.BINARY, fn_name, attr_name, None, full_args).expr()

    def visit_tensor_exprs(self, callback: Callable[["TensorExpression"], None]):
        for arg in self.args:
            arg.visit_tensor_exprs(callback)

    def __repr__(self):
        return f"{repr(self.reduce_use)}({', '.join(repr(a) for a in self.args)})"


class const(TensorExpression):
    """Returns the given constant floating point or integer value."""

    def __init__(self, value: Any):
        with _ir.Context():
            if isinstance(value, float):
                self.value = str(_ir.FloatAttr.get_f64(float(value)))
            elif isinstance(value, int):
                self.value = str(
                    _ir.IntegerAttr.get(_ir.IntegerType.get_signless(64), int(value))
                )
            else:
                raise ValueError(f"const requires int or float but got {type(value)}")

    def to_scalar_expression(self) -> ScalarExpression:
        return ScalarConst(self.value).expr()

    def __repr__(self):
        return f"const({self.value})"


class index(TensorExpression):
    """Returns the iteration index for a given dimension name.

    Resolves the given dimension name to obtain its position in the iteration
    domain of the operation.
    """

    def __init__(self, dim: DimDef):
        self.dim_def = dim
        self.dim = -1

    def resolve_dimension_name(self, affine_state: AffineBuildState):
        self.dim = affine_state.get_dim(self.dim_def.dimname)

    def to_scalar_expression(self) -> ScalarExpression:
        assert self.dim != -1, "Dimension name not resolved"
        return ScalarIndex(self.dim).expr()

    def __repr__(self):
        return f"index({repr(self.dim)})"


###############################################################################
# Function types and function definitions.
###############################################################################


class FunctionKind(Enum):
    UNARY = 0
    BINARY = 1
    TERNARY = 2
    TYPE = 3


class UnaryFnType:
    """Unary function.

    A unary function takes one tensor expression and returns the
    function evaluation result.
    """

    def __init__(self, fn_name: str):
        self.fn_name = fn_name

    def __call__(self, arg: TensorExpression) -> "TensorFn":
        return TensorFn(FunctionKind.UNARY, self.fn_name, None, None, [arg])

    def __repr__(self):
        return f"{self.fn_name}"


class UnaryFn:
    """Unary function namespace."""

    exp = UnaryFnType("exp")
    log = UnaryFnType("log")
    abs = UnaryFnType("abs")
    ceil = UnaryFnType("ceil")
    floor = UnaryFnType("floor")
    negf = UnaryFnType("negf")
    reciprocal = UnaryFnType("reciprocal")
    round = UnaryFnType("round")
    sqrt = UnaryFnType("sqrt")
    rsqrt = UnaryFnType("rsqrt")
    square = UnaryFnType("square")
    tanh = UnaryFnType("tanh")
    erf = UnaryFnType("erf")


class BinaryFnType:
    """Binary function.

    A binary function takes two tensor expressions and returns the
    function evaluation result.
    """

    def __init__(self, fn_name: str):
        self.fn_name = fn_name

    def __call__(self, arg0: TensorExpression, arg1: TensorExpression) -> "TensorFn":
        return TensorFn(FunctionKind.BINARY, self.fn_name, None, None, [arg0, arg1])

    def __repr__(self):
        return f"{self.fn_name}"


class BinaryFn:
    """Binary function namespace.

    As the integer types are signless, signedness is implement by different
    functions that treat integers as signed or unsigned values.

    Examples:
    - max -> `arith.MaxSIOp`
    - max_unsigned -> `arith.MaxUIOp`
    """

    add = BinaryFnType("add")
    sub = BinaryFnType("sub")
    mul = BinaryFnType("mul")
    div = BinaryFnType("div")
    div_unsigned = BinaryFnType("div_unsigned")
    max_signed = BinaryFnType("max_signed")
    min_signed = BinaryFnType("min_signed")
    max_unsigned = BinaryFnType("max_unsigned")
    min_unsigned = BinaryFnType("min_unsigned")
    powf = BinaryFnType("powf")


class TernaryFnType:
    """Ternary function.

    A ternary function takes three tensor expressions and returns the
    function evaluation result.
    """

    def __init__(self, fn_name: str):
        self.fn_name = fn_name

    def __call__(
        self, arg0: TensorExpression, arg1: TensorExpression, arg2: TensorExpression
    ) -> "TensorFn":
        return TensorFn(
            FunctionKind.TERNARY, self.fn_name, None, None, [arg0, arg1, arg2]
        )

    def __repr__(self):
        return f"{self.fn_name}"


class TernaryFn:
    """Ternary function namespace."""

    select = TernaryFnType("select")


class TypeFnType:
    """Type conversion function.

    A type conversion function takes a target type and a tensor expression and
    returns the casted tensor expression.
    """

    def __init__(self, fn_name: str):
        self.fn_name = fn_name

    def __call__(self, type_var: TypeVar, arg: TensorExpression) -> "TensorFn":
        return TensorFn(FunctionKind.TYPE, self.fn_name, None, type_var, [arg])

    def __repr__(self):
        return f"{self.fn_name}"


class TypeFn:
    """Type conversion function namespace.

    As the integer types are signless, signedness is implement by different cast
    functions that treat integers as signed (`cast_signed`) or unsigned
    (`cast_unsigned`) values.

    Examples:
    - cast_signed(I32 -> I64) -> `arith.ExtSIOp`
    - cast_unsigned(I32 -> I64) -> `arith.ExtUIOp`
    """

    cast_signed = TypeFnType("cast_signed")
    cast_unsigned = TypeFnType("cast_unsigned")


class ReduceFnUse:
    """Reduction function use.

    A reduction use specifies the reduction function and dimensions.
    """

    def __init__(
        self,
        binary_fn: Optional[BinaryFnType],
        binary_attr: Optional["BinaryFnAttrDef"],
        *reduce_dims: DimDef,
    ):
        if bool(binary_fn) + bool(binary_attr) != 1:
            raise ValueError("One of 'binary_fn', 'binary_attr' must be specified")
        self.binary_fn = binary_fn
        self.binary_attr = binary_attr
        self.reduce_dims = reduce_dims

    def __call__(self, *args: TensorExpression) -> "TensorReduceFn":
        return TensorReduceFn(self, args)

    def __repr__(self):
        fn = self.binary_fn if self.binary_fn else self.binary_attr
        return f"reduce_{repr(fn)}({', '.join(repr(d) for d in self.reduce_dims)})"


class ReduceFnType:
    """Reduction function.

    A binary function that reduces its RHS into its LHS.
    """

    def __init__(self, binary_fn: BinaryFnType):
        if not isinstance(binary_fn, BinaryFnType):
            raise ValueError(f"Reduce expected a BinaryFnType but got {binary_fn}")
        self.binary_fn = binary_fn

    def __getitem__(self, reduce_dims: Tuple[DimDef]) -> ReduceFnUse:
        return ReduceFnUse(self.binary_fn, None, *reduce_dims)

    def __repr__(self):
        return f"reduce_{repr(self.binary_fn)}"


class ReduceFn:
    add = ReduceFnType(BinaryFn.add)
    mul = ReduceFnType(BinaryFn.mul)
    max_signed = ReduceFnType(BinaryFn.max_signed)
    min_signed = ReduceFnType(BinaryFn.min_signed)
    max_unsigned = ReduceFnType(BinaryFn.max_unsigned)
    min_unsigned = ReduceFnType(BinaryFn.min_unsigned)


###############################################################################
# Operand definitions.
###############################################################################


class OperandKind(Enum):
    INPUT_TENSOR = 0
    SCALAR = 1
    OUTPUT_TENSOR = 2
    INDEX_ATTR = 3
    UNARY_FN_ATTR = 4
    BINARY_FN_ATTR = 5
    TERNARY_FN_ATTR = 6
    TYPE_FN_ATTR = 7


class OperandDef:
    """Definition of an operand passed to an operation.

    Keep the meta information of Tensor, Scalar, and Attribute operands and
    provide the shared registration functionality.
    """

    def __init__(
        self,
        kind: OperandKind,
        type_var: Optional[TypeVar] = None,
        size_exprs: Optional[Sequence[AffineExprDef]] = None,
        index_dims: Optional[Sequence[DimDef]] = None,
        default_indices: Optional[Sequence[int]] = None,
        default_fn: Optional[str] = None,
    ):
        if type_var and not isinstance(type_var, TypeVar):
            raise ValueError(f"OperandDef requires a TypeVar but got {repr(type_var)}")
        self.owner = None  # type: Optional["LinalgOpDef"]
        self.type_var = type_var
        self.size_exprs = size_exprs
        self.index_dims = index_dims
        self.default_indices = default_indices
        self.default_fn = default_fn
        self.kind = kind
        self.name = None  # type: Optional[str]
        self.registered_index = -1  # type: int

    def attach(self, index: int, name: str, owner: "LinalgOpDef"):
        if self.owner:
            raise ValueError(f"OperandDef already registered with an op: {self}")
        self.registered_index = index
        self.name = name
        self.owner = owner

    def is_input(self) -> bool:
        return self.kind == OperandKind.SCALAR or self.kind == OperandKind.INPUT_TENSOR

    def is_tensor(self) -> bool:
        return (
            self.kind == OperandKind.INPUT_TENSOR
            or self.kind == OperandKind.OUTPUT_TENSOR
        )

    def is_attribute(self) -> bool:
        return (
            self.kind == OperandKind.INDEX_ATTR
            or self.kind == OperandKind.UNARY_FN_ATTR
            or self.kind == OperandKind.BINARY_FN_ATTR
            or self.kind == OperandKind.TERNARY_FN_ATTR
            or self.kind == OperandKind.TYPE_FN_ATTR
        )

    def __hash__(self):
        return hash(id(self))

    def __repr__(self):
        return (
            f"{self.name}:OperandDef(kind={self.kind.name}, "
            f"type={repr(self.type_var)}, size_exprs={self.size_exprs}, "
            f"index_dims={self.index_dims}, "
            f"default_indices={self.default_indices}, "
            f"default_fn={self.default_fn})"
        )


class TensorDef:
    """Tensor operand definition.

    Tensor operands are indexed using the associated indexing_map when forwarded
    to the body of the structured op. A unique name identifies the tensor operands
    and an index determines their position in the operation's parameter list. A
    tensor definition takes type, a shape, and an optional flag to mark output
    tensors. Additionally, a tuple of index dimensions may be used to map the
    tensor to the loop dimensions of the operation. This mapping is needed to
    compute the indexing map of shape-only tensors that have no uses.
    """

    def __init__(
        self,
        type_var: TypeVar,
        *shape: AffineExprDef,
        index_dims: Optional[Sequence[DimDef]] = None,
        output: bool = False,
    ):
        if index_dims and len(shape) != len(index_dims):
            raise ValueError(
                f"Expected the shape rank {len(shape)} to match the "
                f"number of index_dims {len(index_dims)}"
            )
        if index_dims and any(not isinstance(dim, DimDef) for dim in index_dims):
            raise ValueError(
                f"TensorDef requires index dims of type DimDef but " f"got {index_dims}"
            )
        kind = OperandKind.OUTPUT_TENSOR if output else OperandKind.INPUT_TENSOR
        self.operand_def = OperandDef(
            kind, type_var=type_var, size_exprs=shape, index_dims=index_dims
        )

    def __getitem__(self, dims: Sequence[AffineExprDef]) -> TensorUse:
        assert self.operand_def.owner, "TensorDef is not registered with an op"
        state = AffineBuildState(
            global_state=self.operand_def.owner._affine_state, allow_new_symbols=False
        )
        if not isinstance(dims, tuple):
            dims = (dims,)  # Handle single subscript case.
        # Special case: (None) is a 0d-scalar use.
        if dims == (None,):
            dims = ()

        exprs = []
        for expr_def in dims:
            if not isinstance(expr_def, AffineExprDef):
                raise KeyError(
                    "A TensorDef can only be subscripted by a tuple of affine dims"
                )
            exprs.append(expr_def)
        return TensorUse(self.operand_def, exprs)

    def __setitem__(self, dims: Sequence[AffineExprDef], value: TensorExpression):
        """Creates a new 1:1 comprehension by binding this tensor to an expression.

        Note that due to the way assignment works in Python, we have to capture
        direct assignment as a setitem on the TensorDef.
        """
        if not isinstance(value, TensorExpression):
            raise ValueError(
                f"Only TensorExpressions can be assigned to TensorDefs. "
                f"Got: {repr(value)}"
            )
        use = self[dims]
        comp = Comprehension((use, value))
        self.operand_def.owner.comprehensions.append(comp)


class ScalarDef(TensorExpression):
    """Scalar operand definition.

    Scalar operands are forwarded to the body of the structured op as they are.
    A unique name identifies the scalars and an index determines their position in
    the operation's parameter list.
    """

    def __init__(self, type_var: TypeVar):
        self.operand_def = OperandDef(OperandKind.SCALAR, type_var=type_var)

    @property
    def scalar_name(self) -> str:
        name = self.operand_def.name
        assert name is not None, "ScalarDef not registered with an op"
        return name

    def to_scalar_expression(self) -> ScalarExpression:
        return ScalarArg(self.scalar_name).expr()


class IndexAttrDef:
    """Index attribute definition.

    Index attributes provide a way to define and set symbols that can be used in
    indexing expressions. Every attribute specifies a tuple of symbols that at
    compile-time are replaced by integer values as well as their default values.
    """

    def __init__(self, *sizes: SymbolDef, default: Sequence[int]):
        if any(not isinstance(size, SymbolDef) for size in sizes):
            raise ValueError(
                f"IndexAttrDef requires sizes of type SymbolDef " f"but got {sizes}"
            )
        if any(not isinstance(default_val, int) for default_val in default):
            raise ValueError(
                f"IndexAttrDef requires default values of type int "
                f"but got {default}"
            )
        if len(sizes) != len(default):
            raise ValueError(
                f"IndexAttrDef expects {len(sizes)} default values "
                f"but got {len(default)}"
            )
        self.operand_def = OperandDef(
            OperandKind.INDEX_ATTR, size_exprs=sizes, default_indices=default
        )


class UnaryFnAttrDef:
    """Unary function attribute definition.

    Unary function attributes provide a way to make the arithmetic computation
    parametrizable. Every attribute specifies a default unary function
    that may be overwritten at operation instantiation time.
    """

    def __init__(self, default: "UnaryFnType"):
        if not isinstance(default, UnaryFnType):
            raise ValueError(
                f"UnaryFnAttrDef requires default of type UnaryFnType "
                f"but got {default}"
            )
        self.operand_def = OperandDef(
            OperandKind.UNARY_FN_ATTR, default_fn=default.fn_name
        )

    def __call__(self, arg: TensorExpression) -> TensorFn:
        return TensorFn(FunctionKind.UNARY, None, self.operand_def, None, [arg])


class BinaryFnAttrDef:
    """Binary function attribute definition.

    Binary function attributes provide a way to make the arithmetic computation
    parametrizable. Every attribute specifies a default binary function
    that may be overwritten at operation instantiation time.
    """

    def __init__(self, default: "BinaryFnType"):
        if not isinstance(default, BinaryFnType):
            raise ValueError(
                f"BinaryFnAttrDef requires default of type BinaryFnType "
                f"but got {default}"
            )
        self.operand_def = OperandDef(
            OperandKind.BINARY_FN_ATTR, default_fn=default.fn_name
        )

    def __call__(self, arg0: TensorExpression, arg1: TensorExpression) -> TensorFn:
        return TensorFn(FunctionKind.BINARY, None, self.operand_def, None, [arg0, arg1])

    def __getitem__(self, reduce_dims: Tuple[DimDef]) -> ReduceFnUse:
        return ReduceFnUse(None, self, *reduce_dims)


class TernaryFnAttrDef:
    """Ternary function attribute definition.

    Ternary function attributes provide a way to make the arithmetic computation
    parametrizable. Every attribute specifies a default Ternary function
    that may be overwritten at operation instantiation time.
    """

    def __init__(self, default: "TernaryFnType"):
        if not isinstance(default, TernaryFnType):
            raise ValueError(
                f"TernaryFnAttrDef requires default of type TernaryFnType "
                f"but got {default}"
            )
        self.operand_def = OperandDef(
            OperandKind.TERNARY_FN_ATTR, default_fn=default.fn_name
        )

    def __call__(self, arg0: TensorExpression, arg1: TensorExpression) -> TensorFn:
        return TensorFn(
            FunctionKind.TERNARY, None, self.operand_def, None, [arg0, arg1]
        )

    def __getitem__(self, reduce_dims: Tuple[DimDef]) -> ReduceFnUse:
        return ReduceFnUse(None, self, *reduce_dims)


class TypeFnAttrDef:
    """Type conversion function attribute definition.

    Type conversion function attributes provide a way to make type conversions
    parameterizable. Every attribute specifies a default type conversion function
    that may be overwritten at operation instantiation time.
    """

    def __init__(self, default: "TypeFnType"):
        if not isinstance(default, TypeFnType):
            raise ValueError(
                f"TypeFnAttrDef requires default of type TypeFnType "
                f"but got {default}"
            )
        self.operand_def = OperandDef(
            OperandKind.TYPE_FN_ATTR, default_fn=default.fn_name
        )

    def __call__(self, type_var: TypeVar, arg: TensorExpression) -> TensorFn:
        return TensorFn(FunctionKind.TYPE, None, self.operand_def, type_var, [arg])


###############################################################################
# Operation definition.
###############################################################################


class Comprehension:
    """Represents a single comprehension."""

    def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):
        self.definitions = list()  # List[TensorUse]
        self.values = list()  # List[TensorExpression]

        # Find the lhs to reduction rhs.
        for assign, value in bindings:
            if isinstance(value, TensorReduceFn):
                if value.lhs:
                    raise ValueError(f"Reduction expression already assigns: {value}")
                value.lhs = assign
            self.definitions.append(assign)
            self.values.append(value)

    @property
    def all_reduction_dims(self) -> Set[Tuple[DimDef, ...]]:
        """Gets the reduction dims for the comprehension or None."""
        result = set()
        for use in self.values:
            if isinstance(use, TensorReduceFn):
                result.add(use.reduce_use.reduce_dims)
            else:
                result.add(tuple())
        return result

    def __repr__(self):
        if len(self.definitions) > 1:
            defs_repr = f"({', '.join(repr(d) for d in self.definitions)})"
            values_repr = f"({', '.join(repr(v) for v in self.values)})"
        else:
            defs_repr = f"{repr(self.definitions[0])}"
            values_repr = f"{repr(self.values[0])}"

        return f"{defs_repr} = {values_repr}"


class OpInterfaceDef:
    """An interface that an op implements."""

    def __init__(self, cpp_name: str):
        self.cpp_name = cpp_name


ContractionOpInterface = OpInterfaceDef("LinalgContractionOpInterface")
ConvolutionOpInterface = OpInterfaceDef("LinalgConvolutionOpInterface")
FillOpInterface = OpInterfaceDef("LinalgFillOpInterface")


class OpDefinitionDef:
    """A method that an op implements."""

    def __init__(self, def_name: str):
        self.def_name = def_name


Canonicalizer = OpDefinitionDef("hasCanonicalizer")


class OpMetadataDef(YAMLObject):
    """Metadata about the op (generally not behavior impacting)."""

    yaml_tag = "!LinalgOpMetadata"

    def __init__(self, name: str, cpp_class_name: Optional[str], doc: Optional[str]):
        self.name = name
        self.cpp_class_name = cpp_class_name if cpp_class_name is not None else name
        self.doc = doc
        self.implements = []  # type: List[OpInterfaceDef]
        self.defines = []  # type: List[OpDefinitionsDef]

    def to_yaml_custom_dict(self):
        d = dict(
            name=self.name,
            cpp_class_name=self.cpp_class_name,
            doc=self.doc,
        )
        if self.implements:
            d["implements"] = [intr.cpp_name for intr in self.implements]
        if self.defines:
            d["defines"] = [defi.def_name for defi in self.defines]
        return d


class LinalgOpDef:
    """Definition of a linalg op."""

    def __init__(
        self, name: str, cpp_class_name: Optional[str] = None, doc: Optional[str] = None
    ):
        self.metadata = OpMetadataDef(name=name, cpp_class_name=cpp_class_name, doc=doc)
        self.registered_operands = dict()  # type: Dict[str, OperandDef]
        self.domain = list()  # type: List[DimDef]
        self.comprehensions = list()  # type: List[Comprehension]
        self._affine_state = AffineBuildState()

    def add_operand(self, name: str, operand: OperandDef):
        """Registers an operand."""
        if name in self.registered_operands:
            raise ValueError(
                f"The operand {name} is already registered "
                f"to {self.registered_operands['name']}"
            )
        structured_op_methods = [
            "inputs",
            "outputs",
            "result_tensors",
            "region",
            "iterator_types",
            "indexing_maps",
            "getRegionBuilder",
            "getLibraryCallName",
        ]
        if operand.is_attribute() and name in structured_op_methods:
            raise ValueError(
                f"The attribute name {name} conflicts with a structured "
                f"op method name"
            )
        # Ensure output tensors are registered after input tensors and scalars and
        # attributes are registered after all other operand types.
        if operand.is_input() and any(
            not op_def.is_input() for op_def in self.registered_operands.values()
        ):
            raise ValueError(f"Input {name} registered after an output or attribute")
        if operand.kind == OperandKind.OUTPUT_TENSOR and any(
            op_def.is_attribute() for op_def in self.registered_operands.values()
        ):
            raise ValueError(f"Output {name} registered after an attribute")
        operand.attach(len(self.registered_operands), name, self)
        self.registered_operands[name] = operand

    def __repr__(self):
        lines = [f"LinalgOpDef({self.metadata.name} -> {self.metadata.cpp_class_name},"]
        for name, operand in self.registered_operands.items():
            lines.append(f"  {operand}")
        if self.comprehensions:
            lines[-1] += " {"
            for comprehension in self.comprehensions:
                lines.append(f"    {comprehension}")
            lines.append("}")
        return "\n".join(lines)
