Reference API Documentation

class Backend(ABC):
    """An abstract implementation of specific Graph implementation.

    A Backend is constructed from an input Graph and provides backend-specific
    implementations of the graph operations. It serves as a bridge between the abstract
    graph representation and concrete backend implementations (e.g., MLIR, TVM, JIR).

    The Implementer provides access to associated Scheduler and Compiler instances
    for applying transformations and generating executable code.
    """

    @abstractmethod
    def get_scheduler(self, **kwargs: Any) -> Scheduler:
        """Returns the scheduler associated with this implementation.

        Args:
            kwargs: scheduler configuration

        Returns:
            The scheduler for applying transformations
        """
        ...

    @abstractmethod
    def get_compiler(self, **kwargs: Any) -> Compiler:
        """Returns the compiler associated with this implementation.

        Args:
            kwargs: compiler configuration

        Returns:
            The compiler for generating executable code
        """
        ...

    @property
    @abstractmethod
    def graph(self) -> Graph:
        """Returns the graph being implemented.

        Returns:
            The source graph for this implementation
        """
        ...

class Compiler(ABC):
    """An abstract implementation of a compiler for a given backend and schedule.

    A Compiler takes a backend-specific implementation and schedule and generates
    an executable Module. It handles the final stage of converting the optimized
    intermediate representation into executable code for the target platform.
    """

    @abstractmethod
    def compile(self, schedule: Schedule) -> Module:
        """Compiles the implementation according to the given schedule.

        Args:
            schedule: The schedule specifying transformations and optimizations

        Returns:
            The compiled executable module
        """
        ...

    @property
    @abstractmethod
    def backend(self) -> "xtc.itf.back.Backend":
        """Returns the implementer associated with this compiler.

        Returns:
            The backend this compiler generates code for
        """
        ...

class Module(ABC):
    """An abstract representation of an executable module.

    A Module is the final output of the compilation process, representing
    compiled code that can be executed. It is produced by a Compiler after
    applying transformations specified by a Schedule to an Implementer's
    representation of a Graph.

    Modules can be exported as shared objects for direct execution and evaluation,
    or for usage in larger applications. They can be executed and evaluated using
    Executor and Evaluator classes to measure performance and validate correctness.
    """

    @property
    @abstractmethod
    def file_type(self) -> str:
        """The module type, can be target dependent.

        Available types are: "executable", "shlib"

        Returns:
            the type of the module
        """
        ...

    @property
    @abstractmethod
    def name(self) -> str:
        """The module name.

        The module name may be used to identify a module to
        execute.

        Returns:
            the name of the module
        """
        ...

    @property
    @abstractmethod
    def payload_name(self) -> str:
        """The payload name for the module.

        The name of the payload to execute for the module.
        Generally the entry point inside the module.

        Returns:
            the name of the executable payload inside the module
        """
        ...

    @property
    @abstractmethod
    def file_name(self) -> str:
        """The storage file name of the module.

        The file name extension should match the module file type.

        Returns:
            the path to the generated module file
        """
        ...

    @abstractmethod
    def export(self) -> None:
        """Exports the module to a format suitable for execution.

        This method handles the final step of making the compiled code
        available for execution, typically by writing it to a shared
        object file or similar executable format.
        """
        ...

    @abstractmethod
    def get_evaluator(self, **kwargs: Any) -> Evaluator:
        """Returns a suitable evaluator for the module.

        Args:
            kwargs: evaluator configuration

        Returns:
            The evaluator for executing the module
        """
        ...

    @abstractmethod
    def get_executor(self, **kwargs: Any) -> Executor:
        """Returns a suitable executor for the module.

        Args:
            kwargs: executor configuration

        Returns:
            The executor for executing the module
        """
        ...

class ConstantTensorType(TensorType):
    @property
    @abstractmethod
    @override
    def shape(self) -> ConstantShapeType:
        """Returns the tensor's constant shape as a tuple of dimension sizes.

        Returns:
            The size of each dimension in the tensor
        """
        ...

    @property
    @abstractmethod
    @override
    def dtype(self) -> ConstantDataType:
        """Returns the tensor's constant data type.

        Returns:
            The underlying data type of the tensor elements
        """
        ...

class Tensor(ABC):
    """An abstract representation of a multidimensional object.

    A Tensor is a fundamental input/output type in the dataflow graph,
    representing multidimensional data with associated type information.
    Tensors are used as inputs and outputs for Node operations in the Graph,
    and their dimensions and types can be used for inference throughout
    the compilation process.
    """

    @property
    @abstractmethod
    def type(self) -> TensorType:
        """Returns the tensor's type information.

        Returns:
            The type descriptor containing shape and dtype information
        """
        ...

    @property
    @abstractmethod
    def data(self) -> Any:
        """Returns the tensor's linearized data.

        Returns:
            any: The tensor's data
        """
        ...

    @abstractmethod
    def numpy(self) -> numpy.typing.NDArray:
        """Convert the tensor to a numpy array.

        Returns:
            The tensor's data as a numpy array
        """
        ...

class TensorType(ABC):
    """An abstract representation of a tensor's type information.

    TensorType defines the shape and data type characteristics of a tensor,
    providing the necessary information for type inference and validation
    during graph operations. This includes the tensor's dimensionality,
    size along each dimension, and the underlying data type.
    """

    @property
    @abstractmethod
    def shape(self) -> ShapeType:
        """Returns the tensor's shape as a tuple of dimension sizes.

        Returns:
            The size of each dimension in the tensor
        """
        ...

    @property
    @abstractmethod
    def dtype(self) -> DataType:
        """Returns the tensor's data type.

        Returns:
            The underlying data type of the tensor elements
        """
        ...

    @property
    @abstractmethod
    def ndim(self) -> int:
        """Returns the number of dimensions in the tensor.

        Returns:
            The tensor's dimensionality
        """
        ...

class Evaluator(ABC):
    """An abstract implementation of a Module performance evaluator.

    An Evaluator measures and validates the performance of compiled Modules.
    It works alongside Executors to provide performance metrics and correctness
    validation of the compiled code. This is crucial for assessing the
    effectiveness of different compilation strategies and optimizations.

    Evaluators can measure metrics like execution time, throughput, and
    validate output correctness against reference implementations.
    """

    @abstractmethod
    def evaluate(self) -> tuple[list[float], int, str]:
        """Evaluates the performance of the associated Module.

        Executes the Module multiple times to gather performance metrics,
        potentially validating correctness against reference implementations.

        Returns:
            List of performance measurements (typically execution times in seconds), error coe and error message
        """
        ...

    @property
    @abstractmethod
    def module(self) -> "xtc.itf.comp.Module":
        """Returns the Module being evaluated.

        Returns:
            The compiled Module this evaluator is measuring
        """
        ...

class Executor(ABC):
    """An abstract implementation of a Module executor.

    An Executor is responsible for running compiled Modules and validating their
    execution. It provides the runtime environment for executing the compiled code,
    typically in the form of shared objects generated by the compilation process.
    """

    @abstractmethod
    def execute(self) -> int:
        """Executes the associated Module.

        Runs the compiled code contained in the Module, providing the necessary
        runtime environment and handling any required setup/teardown.

        Returns:
            Status code indicating execution success (0) or failure (non-zero)
        """
        ...

    @property
    @abstractmethod
    def module(self) -> "xtc.itf.comp.Module":
        """Returns the Module being executed.

        Returns:
            The compiled Module this executor is responsible for running
        """
        ...

class Graph(ABC):
    """An abstract representation of a dataflow graph over Tensor types.

    A Graph is a directed acyclic graph (DAG) over Node objects with input Tensors
    and output Tensors. From given input Tensor dimensions and type, all node inputs,
    outputs and graph outputs Tensor dimensions and types can be inferred.

    A Graph can be evaluated by interpretation or compiled through various backends
    (mlir, tvm, jir) using corresponding Backend, Scheduler, and Compiler classes.

    Nodes in the graph are keyed by their uid which is globally unique.
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Returns the name of this graph. May be non-unique or empty.

        Returns:
            The graph's name
        """
        ...

    @property
    @abstractmethod
    def nodes(self) -> Mapping[str, Node]:
        """Returns a dictionary of all nodes in the graph, keyed by node uid.

        Returns:
            Dictionary mapping node names to Node objects
        """
        ...

    @property
    @abstractmethod
    def inputs(self) -> Sequence[str]:
        """Returns the list of input tensor uids for this graph.

        Returns:
            List of input tensor uids
        """
        ...

    @property
    @abstractmethod
    def outputs(self) -> Sequence[str]:
        """Returns the list of output tensor uids for the graph.

        Returns:
            List of output tensor uids
        """
        ...

    @property
    @abstractmethod
    def inputs_types(self) -> Sequence[TensorType] | None:
        """Returns the list of inputs tensor types

        Returns None when no input tensor type was given
        and forward_types was not called.

        Returns:
            List of input tensor types or None if undef
        """
        ...

    @property
    def outputs_types(self) -> Sequence[TensorType] | None:
        """Returns the list of outputs tensor types

        Returns None when no input tensor type was given
        and forward_types was not called.

        Returns:
            List of output tensor types or None if undef
        """
        ...

    @property
    @abstractmethod
    def inputs_nodes(self) -> Sequence[Node]:
        """Returns the list of inputs nodes.

        The list is such that:
        - nodes appear in the same order as outputs
        - nodes appear only once in the list, hence
          the size may differe from the inputs size

        Returns:
            List of input nodes
        """
        ...

    @property
    @abstractmethod
    def outputs_nodes(self) -> Sequence[Node]:
        """Returns the list of output nodes.

        The list is such that:
        - nodes appear in the same order as outputs
        - nodes appear only once in the list, hence
          the size may differe from the outputs size

        Returns:
            List of output nodes
        """
        ...

    @abstractmethod
    def forward_types(self, inputs_types: Sequence[TensorType]) -> Sequence[TensorType]:
        """Infers output tensor types from input tensor types.

        Args:
            inputs_types: List of input tensor types

        Returns:
            List of inferred output tensor types
        """
        ...

    @abstractmethod
    def forward(self, inputs: Sequence[Tensor]) -> Sequence[Tensor]:
        """Evaluate the graph with input tensors to produce output tensors.

        Args:
            inputs: List of input tensors

        Returns:
            List of output tensors
        """
        ...

class Node(ABC):
    """An abstract representation of a node in a dataflow graph.

    A Node represents a pure operation on input Tensor objects, resulting in output
    Tensor objects. Each node has a unique global uid, a set of input
    and output tensors, and an associated Operator that defines its semantic behavior.
    """

    @property
    @abstractmethod
    def uid(self) -> str:
        """Returns the globally unique id over all created nodes.

        Returns:
            The node's globally unique id
        """
        ...

    @property
    @abstractmethod
    def name(self) -> str:
        """Returns the name of this node. Can be non-unique or empty.

        Returns:
            The node's name
        """
        ...

    @property
    @abstractmethod
    def inputs(self) -> list[str]:
        """Returns the list of input tensors uids for this node.

        As of now nodes can only have one output, hence a
        tensor uid is the same as the producing node uid.

        Returns:
            List of input tensor names
        """
        ...

    @property
    @abstractmethod
    def outputs(self) -> list[str]:
        """Returns the list of output tensors uids for this node.

        As of now nodes can only have one output, hence the
        node outputs is the list containing the node uid itself.

        Returns:
            List of output tensor names
        """
        ...

    @property
    @abstractmethod
    def inputs_types(self) -> Sequence[TensorType] | None:
        """Returns the list of inputs tensor types

        Returns None when no input tensor type was given
        and forward_types was not called.

        Returns:
            List of input tensor types or None if undef
        """
        ...

    @property
    def outputs_types(self) -> Sequence[TensorType] | None:
        """Returns the list of outputs tensor types

        Returns None when no input tensor type was given
        and forward_types was not called.

        Returns:
            List of output tensor types or None if undef
        """
        ...

    @property
    @abstractmethod
    def preds(self) -> Sequence[str]:
        """Returns the list of predecessor nodes uids.

        The list is such that:
        - nodes appear in the same order of inputs
        - nodes appear only once, hence the number of preds
          may not be equal to the number of inputs

        Returns:
            List of predecessors nodes uids
        """
        ...

    @property
    @abstractmethod
    def preds_nodes(self) -> Sequence["Node"]:
        """Returns the list of predecessor nodes.

        The list is the same as for preds, but contains the nodes
        instead of the nodes' uids.

        Returns:
            List of predecessors nodes
        """
        ...

    @property
    @abstractmethod
    def operator(self) -> Operator:
        """Returns the operator that defines this node's behavior.

        Returns:
            The algebraic operator associated with this node
        """
        ...

    @property
    @abstractmethod
    def operation(self) -> Operation:
        """Returns the operation that defines this node's behavior.

        The operation specifies the operator behavior and the instanciated
        operator dimensions.

        Returns:
            The algebraic operation associated with this node
        """
        ...

    @abstractmethod
    def forward_types(self, inputs_types: Sequence[TensorType]) -> Sequence[TensorType]:
        """Infers output tensor types from input tensor types.

        Args:
            inputs_types: List of input tensor types

        Returns:
            List of inferred output tensor types
        """
        ...

    @abstractmethod
    def forward(self, inputs: Sequence[Tensor]) -> Sequence[Tensor]:
        """Evaluate the node with input tensors to produce output tensors.

        Args:
            inputs: List of input tensors

        Returns:
            List of output tensors
        """
        ...

class Operation(ABC):
    """An abstract representation of an Operation, itself a specialized Operator.

    An Operation represents the computation performed by a Node, i.e. an Operator
    specification and instanciated dimensions and types for the inputs and
    outputs.

    The Operation computation is currently internal,
    though the inputs and outputs accesses functions are available through the
    accesses_maps.
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Returns the unique name of this operations's operator.

        Returns:
            The node's operator unique name
        """
        ...

    @property
    @abstractmethod
    def attrs(self) -> OperationAttrs:
        """Returns the dict of attributes for this operation.

        Returns:
            Dict of attributes per name
        """
        ...

    @property
    @abstractmethod
    def inputs_types(self) -> Sequence[TensorType]:
        """Returns the list of input tensors types for this operation.

        Returns:
            List of input tensors types
        """
        ...

    @property
    @abstractmethod
    def outputs_types(self) -> Sequence[TensorType]:
        """Returns the list of output tensors types for this operation.

        Returns:
            List of output tensors types
        """
        ...

    @property
    @abstractmethod
    def dims(self) -> Mapping[DimSpec, DimSize]:
        """Returns the dict of dimensions size for this operation.

        A dimension size may be resolved (int) or symbolic (str).

        Returns:
            Dict of dim name, dim size
        """
        ...

    @abstractmethod
    def dims_kind(self, kind: str) -> Sequence[DimSpec]:
        """Returns the list of dimensions of the given kind.

        The kind argument is currently one of:
        - "P" for parallel dims
        - "R" for reduction axes

        Returns:
            List of dims names
        """
        ...

    @property
    @abstractmethod
    def accesses_maps(self) -> AccessesMaps:
        """Returns the accesses map for this operation.

        Accesses maps are a 3-tuple with:
        - operation dimensions names tuple,
        - tuple of inputs accesses tuples for each input,
        - tuple of outputs accesses tuples for each output,

        Returns:
            Accesses map for this operation
        """
        ...

class Operator(ABC):
    """An abstract representation of the algebraic operation for a node.

    An Operator defines the semantic behavior of operations in the graph, including
    how it transforms input tensor types and how it processes tensor data. It provides
    both type inference capabilities and concrete implementations of the operation.
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Returns the unique identifier for this operator type.

        Returns:
            The operator's name
        """
        ...

    @abstractmethod
    def forward_types(self, inputs_types: Sequence[TensorType]) -> Sequence[TensorType]:
        """Infers output tensor types from input tensor types.

        Args:
            inputs_types: List of input tensor types

        Returns:
            List of inferred output tensor types
        """
        ...

    @abstractmethod
    def forward(self, inputs: Sequence[Tensor]) -> Sequence[Tensor]:
        """Evaluate the operator with input tensors to produce output tensors.

        Args:
            inputs: List of input tensors

        Returns:
            List of output tensors
        """
        ...

class Schedule(ABC):
    """An abstract representation of the result of transformations from a scheduler.

    A Schedule captures all the transformations and optimizations that have been
    applied to an implementation by its associated Scheduler. It serves as an
    intermediate representation between scheduling operations and code generation,
    allowing Compilers to generate optimized executable code based on the
    specified transformations.

    Schedules are backend-specific and contain the necessary information for
    their associated Compiler to generate code optimized for the target
    platform and runtime. Common transformations captured in a Schedule include:
    - Tiling
    - Loop interchange
    - Vectorization
    - Parallelization
    - Loop unrolling
    """

    @property
    @abstractmethod
    def scheduler(self) -> "xtc.itf.schd.Scheduler":
        """Returns the scheduler that generated this schedule.

        Returns:
        """
        ...

class Scheduler(ABC):
    """An abstract implementation of the backend scheduler.

    A Scheduler is constructed from a given Backend and is responsible for
    applying primitive scheduling operations and transformations to the implementation.
    It generates a Schedule that captures these transformations, which can then be
    used by a Compiler to generate optimized executable code.

    Schedulers are backend-specific and work with their associated Backend
    to provide optimization capabilities appropriate for the target platform
    and runtime.
    """

    @abstractmethod
    def schedule(self) -> Schedule:
        """Creates a Schedule from the applied transformations.

        Returns a Schedule object that captures all the transformations and
        optimizations that have been applied to the implementation. This
        Schedule can then be used by a Compiler to generate executable code.

        Returns:
            Schedule: The resulting schedule containing all applied transformations
        """
        ...

    @property
    @abstractmethod
    def backend(self) -> "xtc.itf.back.Backend":
        """Returns the backend associated with this scheduler.

        Returns:
            Backend: The backend-specific implementation this scheduler
                     applies transformations to
        """
        ...

    @abstractmethod
    def set_dims(self, dims: list[str]) -> None:
        """Redefines dimensions names.

        Use provided abstract dimensions names for the scheduler
        transformantions instead of the default operation dimensions names.

        This should be set before applying the transformations

        Args:
            dims: list of dimensions names
        """
        ...

    @abstractmethod
    def split(
        self, dim: str, segments: dict[str, int], root: str = DEFAULT_ROOT
    ) -> None:
        """Split a dimension into `len(segments)` segments.

        Each segment is characterized by a starting/cutting point,
        which is also the endpoint of the previous segment, and by
        the name of the new axis created by the cut. The segments
        items must be provided in ascending order of the cut points
        on the axis.

        Args:
            dim: name of the dimension to split
            segments: ordered dict of new root name and segment
                      starting point
        """
        ...

    def strip_mine(
        self, dim: str, tiles: dict[str, int], root: str = DEFAULT_ROOT
    ) -> None:
        """Apply a multi level strip mining transformation on the given dimension.

        The strip mining can be seen as a multi level 1D tiling where the
        given tile sizes are interpreter outer to inner.
        After this transformation, the number of axis for the given initial
        dimension is `1 + len(tiles)` where the first axis inherits
        the name of the dimension, and the remaining axis names are
        given by the given tiles keys.
        Each 1D tile size must be greater or equal to the inner tile sizes.
        Some backend may not support non-divisible tile sizes, in which
        case an assertion is raised.

        Args:
            dim: name of the dimension to strip mine
            tiles: dict outer to inner of axis name and tile size
            root: the parent split (or the operator's absolute root)
        """
        self.tile(dim=dim, tiles=tiles, root=root)

    @abstractmethod
    def tile(self, dim: str, tiles: dict[str, int], root: str = DEFAULT_ROOT) -> None:
        """Apply a multi level tiling operation.

        As of now the interface is limited to a single dimension tiling,
        hence it is equivalent to strip mining the given dimension.

        In order to create multi dimensional tiles, strip mine each dimension
        with tile or stip_mine and use interchange to reorder generated axes
        accordingly.

        Args:
            dim: name of the dimension to tile
            tiles: dict outer to inner of axis name and tile size
            root: the parent split (or the operator's absolute root)
        """
        ...

    @abstractmethod
    def interchange(self, permutation: list[str], root: str = DEFAULT_ROOT) -> None:
        """Apply interchange over all axes.

        The given permutation of axes names is interpreted
        outer to inner and must have the same size as the
        number of axes after tiling.

        Args:
            permutation: outer to inner axes names permutation
            root: the parent split (or the operator's absolute root)
        """
        ...

    @abstractmethod
    def vectorize(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
        """Apply vectorizations on the given axes names.

        The axes names given must all be inner axes and parallel axes, full
        unrolling and vectorization of all given axes is implied.

        Args:
            axes: axes names to vectorize
            root: the parent split (or the operator's absolute root)
        """
        ...

    @abstractmethod
    def parallelize(self, axes: list[str], root: str = DEFAULT_ROOT) -> None:
        """Apply parallelization on the given axes names.

        The axes names must given must all be outer axes and parallel axes.

        Args:
            axes: axes names to parallelize
            root: the parent split (or the operator's absolute root)
        """
        ...

    @abstractmethod
    def unroll(self, unrolls: dict[str, int], root: str = DEFAULT_ROOT) -> None:
        """Apply unrolling on the given axes names.

        Each given axes name is unrolled with the specified unroll
        factor. The unroll factors must be greater or equal to 1.

        Args:
            unrolls: dict of axes names and unroll factor
            root: the parent split (or the operator's absolute root)
        """

    @abstractmethod
    def buffer_at(
        self, axis: str, mtype: str | None = None, root: str = DEFAULT_ROOT
    ) -> None:
        """Create a write buffer at a given level.

        A write buffer is created for the output under the given
        axis, The buffer memory type can be specified or defaults
        to the local memory at this level.

        Args:
            axis: localisation of the write buffer
            mtype: buffer memory type for the allocation
            root: the parent split (or the operator's absolute root)
        """
        ...

    @abstractmethod
    def pack_at(
        self,
        axis: str,
        input_idx: int,
        mtype: str | None = None,
        pad: bool = False,
        root: str = DEFAULT_ROOT,
    ) -> None:
        """Create a packed read buffer at a given level.

        A packed read buffer is created for the given input buffer index.
        The buffer memory type can be specified or defaults
        to the local memory at this level.
        When pad is true, a padding strategy is applied in order to reduce
        sets/banks conflicts.

        Args:
            axis: localisation of the write buffer
            input_idx: input buffer index for the scheduled computation
            mtype: buffer memory type for the allocation
            pad: whether to add padding or not
            root: the parent split (or the operator's absolute root)
        """
        ...

    @abstractmethod
    def fuse_producer_at(
        self, axis: str, input_idx: int, root: str = DEFAULT_ROOT
    ) -> None:
        """Fuse producer computation at the given consumer location.

        Given the input index identifying the producer of the input buffer,
        fuse the computation at the given scheduled consumer axis.
        The necessary input slices reads and computations will be inserted
        for computing the output tile at the given axis location.

        Args:
            axis: localisation of the fusion in the consumer
            input_idx: input index of the consumer
            root: the parent split (or the operator's absolute root)
        """
        ...

    @abstractmethod
    def define_memory_mesh(self, axes: dict[str, int]) -> None:
        """Define a memory mesh.

        Args:
            axes: dictionary where keys are axes names and values are the number of memories along each axis
        """
        ...

    @abstractmethod
    def define_processor_mesh(self, axes: dict[str, int]) -> None:
        """Define a processor mesh. It must be a superset of the memory mesh.

        Args:
            axes: dictionary where keys are axes names and values are the number of processors along each axis
        """
        ...

    @abstractmethod
    def distribute(
        self, axis: str, processor_axis: str, root: str = DEFAULT_ROOT
    ) -> None:
        """Distribute computation across processors along a given axis.

        This method distributes the computation of the specified axis across
        multiple processors or cores. The processor_axis parameter defines
        the axis that represents the processor dimension for this distribution.

        Args:
            axis: the axis to distribute across processors
            processor_axis: the axis representing the processor dimension
            root: the parent split (or the operator's absolute root)
        """
        ...

    @abstractmethod
    def distributed_buffer_at(
        self,
        axis: str,
        input_idx: int,
        memory_axes: list[str],
        root: str = DEFAULT_ROOT,
    ) -> None:
        """Create a distributed buffer at a given level across multiple memory axes.

        This method creates a distributed buffer for the given input buffer index
        at the specified axis level. The buffer is distributed across the provided
        memory axes, enabling distributed memory management and access patterns
        for improved performance in distributed computing environments.

        Args:
            axis: the axis level where the distributed buffer should be created
            input_idx: input buffer index for the scheduled computation
            memory_axes: list of memory axes across which to distribute the buffer
            root: the parent split (or the operator's absolute root)
        """
        ...