schedule.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=unused-import
"""The computation schedule api of TVM."""
import collections
import inspect
from typing import Callable, List

import tvm._ffi
from tvm._ffi.base import string_types

from tvm.runtime import Object, convert
from tvm.ir import container as _container
from tvm.tir import IterVar, Buffer, Var, IndexMap

from . import tensor as _tensor
from . import _ffi_api


@tvm._ffi.register_object
class Split(Object):
    """Split operation on axis."""


@tvm._ffi.register_object
class Fuse(Object):
    """Fuse operation on axis."""


@tvm._ffi.register_object
class Singleton(Object):
    """Singleton axis."""


def create_schedule(ops):
    """Create a schedule for list of ops

    Parameters
    ----------
    ops : list of Operations
        The source expression.

    Returns
    -------
    sch : schedule.Schedule
        The created schedule.
    """
    if not isinstance(ops, (list, _container.Array)):
        ops = [ops]
    return _ffi_api.CreateSchedule(ops)


@tvm._ffi.register_object
class Schedule(Object):
    """Schedule for all the stages."""

    def __getitem__(self, k):
        if isinstance(k, _tensor.Tensor):
            k = k.op
        if not isinstance(k, _tensor.Operation):
            raise ValueError("Expect schedule key to be Tensor or Operation")
        if k not in self.stage_map:
            raise ValueError("Cannot find the operation %s in schedule" % (str(k)))
        return self.stage_map[k]

    def normalize(self):
        """Build a normalized schedule from the current schedule.

        Insert necessary rebase to make certain iter var to start from 0.
        This is needed before bound inference and followup step.

        Returns
        -------
        sch : Schedule
            The normalized schedule.
        """
        return _ffi_api.ScheduleNormalize(self)

    def create_group(self, outputs, inputs, include_inputs=False):
        """Create stage group by giving output and input boundary.

        The operators between outputs and inputs are placed as member of group.
        outputs are include in the group, while inputs are not included.

        Parameters
        ----------
        outputs : list of Tensors
            The outputs of the group.

        inputs : list of Tensors
            The inputs of the group.

        include_inputs : boolean, optional
            Whether include input operations in the group if they are used by outputs.

        Returns
        -------
        group : Stage
            A virtual stage represents the group, user can use compute_at to move
            the attachment point of the group.
        """
        if isinstance(outputs, _tensor.Tensor):
            outputs = [outputs]
        if isinstance(inputs, _tensor.Tensor):
            inputs = [inputs]
        return _ffi_api.ScheduleCreateGroup(self, outputs, inputs, include_inputs)

    def cache_read(self, tensor, scope, readers):
        """Create a cache read of original tensor for readers.

        This will mutate the body of the readers.
        A new cache stage will be created for the tensor.
        Call this before doing any split/fuse schedule.

        Parameters
        ----------
        tensor : Tensor
            The tensor to be cached.
        scope : str
            The scope of cached
        readers : list of Tensor or Operation
            The readers to read the cache.

        Returns
        -------
        cache : Tensor
            The created cache tensor.
        """
        if isinstance(readers, (_tensor.Tensor, _tensor.Operation)):
            readers = [readers]
        readers = [t.op if isinstance(t, _tensor.Tensor) else t for t in readers]
        return _ffi_api.ScheduleCacheRead(self, tensor, scope, readers)

    def cache_write(self, tensor, scope):
        """Create a cache write of original tensor, before storing into tensor.

        This will mutate the body of the tensor.
        A new cache stage will created before feed into the tensor.

        This function can be used to support data layout transformation.
        If there is a split/fuse/reorder on the data parallel axis of tensor
        before cache_write is called. The intermediate cache stores
        the data in the layout as the iteration order of leave axis.
        The data will be transformed back to the original layout in the original tensor.
        User can further call compute_inline to inline the original layout and keep
        the data stored in the transformed layout.

        Parameters
        ----------
        tensor : Tensor, list or tuple
            The tensors to be feed to. All the tensors must be produced by one computeOp
        scope : str
            The scope of cached

        Returns
        -------
        cache : Tensor
            The created cache tensor.
        """
        return _ffi_api.ScheduleCacheWrite(self, tensor, scope)

    def rfactor(self, tensor, axis, factor_axis=0):
        """Factor a reduction axis in tensor's schedule to be an explicit axis.

        This will create a new stage that generated the new tensor with axis
        as the first dimension. The tensor's body will be rewritten as a reduction
        over the factored tensor.

        Parameters
        ----------
        tensor : Tensor
            The tensor to be factored.
        axis : IterVar
            The reduction axis in the schedule to be factored.
        factor_axis : int
            The position where the new axis is placed.

        Returns
        -------
        tfactor : Tensor or Array of Tensor
            The created factored tensor.
        """
        factored = _ffi_api.ScheduleRFactor(self, tensor, axis, factor_axis)
        return factored[0] if len(factored) == 1 else factored


@tvm._ffi.register_object
class Stage(Object):
    """A Stage represents schedule for one operation."""

    def split(self, parent, factor=None, nparts=None):
        """Split the stage either by factor providing outer scope, or both

        Parameters
        ----------
        parent : IterVar
             The parent iter var.

        factor : Expr, optional
             The splitting factor

        nparts : Expr, optional
             The number of outer parts.

        Returns
        -------
        outer : IterVar
            The outer variable of iteration.

        inner : IterVar
            The inner variable of iteration.
        """
        if nparts is not None:
            if factor is not None:
                raise ValueError("Do not need to provide both outer and nparts")
            outer, inner = _ffi_api.StageSplitByNParts(self, parent, nparts)
        else:
            if factor is None:
                raise ValueError("Either nparts or factor need to be provided")
            outer, inner = _ffi_api.StageSplitByFactor(self, parent, factor)
        return outer, inner

    def fuse(self, *args):
        """Fuse multiple consecutive iteration variables into a single iteration variable.

        fused = fuse(...fuse(fuse(args[0], args[1]), args[2]),..., args[-1])
        The order is from outer to inner.

        Parameters
        ----------
        args : list of IterVars
            Itervars that proceeds each other

        Returns
        -------
        fused : IterVar
            The fused variable of iteration.
        """
        fused = _ffi_api.StageFuse(self, args)
        return fused

    def set_scope(self, scope):
        """Set the thread scope of this stage

        Parameters
        ----------
        scope : str
            The thread scope of this stage
        """
        return _ffi_api.StageSetScope(self, scope)

    def bind(self, ivar, thread_ivar):
        """Bind ivar to thread index thread_ivar

        Parameters
        ----------
        ivar : IterVar
            The iteration to be binded to thread.

        thread_ivar : IterVar
            The thread to be binded.
        """
        _ffi_api.StageBind(self, ivar, thread_ivar)

    def env_threads(self, threads):
        """Mark threads to be launched at the outer scope of composed op.

        Parameters
        ----------
        threads : list of threads
            The threads to be launched.
        """
        if isinstance(threads, IterVar):
            threads = [threads]
        _ffi_api.StageEnvThreads(self, threads)

    def set_store_predicate(self, predicate):
        """Set predicate under which store to the array can be performed.

        Use this when there are duplicated threads doing the same store and we only
        need one of them to do the store.

        Parameters
        ----------
        predicate : Expr
            The guard condition fo store.
        """
        _ffi_api.StageSetStorePredicate(self, predicate)

    def compute_at(self, parent, scope):
        """Attach the stage at parent's scope

        Parameters
        ----------
        parent : Stage
            The parent stage

        scope : IterVar
            The loop scope t be attached to.
        """
        _ffi_api.StageComputeAt(self, parent, scope)

    def compute_inline(self):
        """Mark stage as inline

        Parameters
        ----------
        parent : Stage
            The parent stage
        """
        _ffi_api.StageComputeInline(self)

    def compute_root(self):
        """Attach the stage at parent, and mark it as root

        Parameters
        ----------
        parent : Stage
            The parent stage
        """
        _ffi_api.StageComputeRoot(self)

    def reorder(self, *args):
        """reorder the arguments in the specified order.

        Parameters
        ----------
        args : list of IterVar
            The order to be ordered
        """
        _ffi_api.StageReorder(self, args)

    def tile(self, x_parent, y_parent, x_factor, y_factor):
        """Perform tiling on two dimensions

        The final loop order from outmost to inner most are
        [x_outer, y_outer, x_inner, y_inner]

        Parameters
        ----------
        x_parent : IterVar
            The original x dimension
        y_parent : IterVar
            The original y dimension
        x_factor : Expr
            The stride factor on x axis
        y_factor : Expr
            The stride factor on y axis

        Returns
        -------
        x_outer : IterVar
            Outer axis of x dimension
        y_outer : IterVar
            Outer axis of y dimension
        x_inner : IterVar
            Inner axis of x dimension
        p_y_inner : IterVar
            Inner axis of y dimension
        """
        x_outer, y_outer, x_inner, y_inner = _ffi_api.StageTile(
            self, x_parent, y_parent, x_factor, y_factor
        )
        return x_outer, y_outer, x_inner, y_inner

    def vectorize(self, var):
        """Vectorize the iteration.

        Parameters
        ----------
        var : IterVar
            The iteration to be vectorize
        """
        _ffi_api.StageVectorize(self, var)

    def tensorize(self, var, tensor_intrin):
        """Tensorize the computation enclosed by var with tensor_intrin

        Parameters
        ----------
        var : IterVar
            The iteration boundary of tensorization.

        tensor_intrin : TensorIntrin
            The tensor intrinsic used for computation.
        """
        _ffi_api.StageTensorize(self, var, tensor_intrin)

    def unroll(self, var):
        """Unroll the iteration.

        Parameters
        ----------
        var : IterVar
            The iteration to be unrolled.
        """
        _ffi_api.StageUnroll(self, var)

    def parallel(self, var):
        """Parallelize the iteration.

        Parameters
        ----------
        var : IterVar
            The iteration to be parallelized.
        """
        _ffi_api.StageParallel(self, var)

    def pragma(self, var, pragma_type, pragma_value=None):
        """Annotate the iteration with pragma

        This will translate to a pragma_scope surrounding
        the corresponding loop generated.
        Useful to support experimental features and extensions.

        Parameters
        ----------
        var : IterVar
            The iteration to be anotated

        pragma_type : str
             The pragma string to be annotated

        pragma_value : Expr, optional
             The pragma value to pass along the pragma

        Note
        ----
        Most pragmas are advanced/experimental features
        and may subject to change. List of supported pragmas:

        - **debug_skip_region**

          Force skip the region marked by the axis and turn it into no-op.
          This is useful for debug purposes.

        - **parallel_launch_point**

          Specify to launch parallel threads outside the
          specified iteration loop. By default the threads
          launch at the point of parallel construct.
          This pragma moves the launching point to even outer scope.
          The threads are launched once and reused across multiple
          parallel constructs as BSP style program.

        - **parallel_barrier_when_finish**

          Insert a synchronization barrier between working threads
          after the specified loop iteration finishes.

        - **parallel_stride_pattern**

          Hint parallel loop to execute in strided pattern.
          :code:`for (int i = task_id; i < end; i += num_task)`

        """
        if isinstance(pragma_value, string_types):
            pragma_value = convert(pragma_value)
        _ffi_api.StagePragma(self, var, pragma_type, pragma_value)

    def prefetch(self, tensor, var, offset):
        """Prefetch the specified variable

        Parameters
        ----------
        tensor : Tensor
            The tensor to be prefetched
        var : IterVar
            The loop point at which the prefetching is applied
        offset : Expr
            The number of iterations to be prefetched before actual execution
        """
        _ffi_api.StagePrefetch(self, tensor, var, offset)

    def storage_align(self, axis, factor, offset):
        """Set alignment requirement for specific axis

        This ensures that stride[axis] == k * factor + offset for some k.
        This is useful to set memory layout to for more friendly memory
        access pattern. For example, we can set alignment to be
        factor=2, offset=1 to avoid bank conflict for thread access on
        higher dimension in GPU shared memory.

        Parameters
        ----------
        axis : IterVar
            The axis dimension to be aligned.
        factor : int
            The factor in alignment specification.
        offset : int
            The offset in the alignment specification.
        """
        _ffi_api.StageStorageAlign(self, axis, factor, offset)

    def double_buffer(self):
        """Compute the current stage via double buffering.

        This can only be applied to intermediate stage.
        This will double the storage cost of the current stage.
        Can be useful to hide load latency.
        """
        _ffi_api.StageDoubleBuffer(self)

    def rolling_buffer(self):
        """Compute the current stage via rolling buffering.

        This can only be applied to intermediate stage.
        This will change the storage cost of the current stage.
        """
        _ffi_api.StageRollingBuffer(self)

    def transform_layout(self, mapping_function: Callable[..., List[tvm.tir.PrimExpr]]):
        """Defines the layout transformation for the current stage's tensor.

        The map from initial_indices to final_indices must be an
        invertible affine transformation.  This method may be called
        more than once for a given tensor, in which case each
        transformation is applied sequentially.

        If the stage is a ComputeOp, then the iteration order of the
        compute stage is rewritten to be a row-major traversal of the
        tensor, and the new loop iteration variables are returned.
        For all other stages, the loop iteration order is unmodified,
        and the return value is None.

        Parameters
        ----------
        mapping_function : Callable[..., List[tvm.tir.PrimExpr]]

            A callable that accepts N arguments of type tvm.tir.Var,
            and outputs a list of PrimExpr.  The input arguments
            represent the location of a value in the current stage's
            tensor, using the pre-transformation layout.  The return
            value of the function gives the location of that value in
            the current stage's tensor, using the post-transformation
            layout.

        Returns
        -------
        new_iter_vars : Optional[List[tvm.tir.IterVar]]

            If the stage is a ComputeOp, then the return will be the
            updated loop iteration variables over the data array, in
            the same order as the output values from the
            `mapping_function`.

            Otherwise, the return value is None.

        Examples
        --------
        .. code-block:: python

            # ``A`` is a tensor whose compute definition is in NHWC
            # format, and should be transformed into NCHWc format.

            s[A].transform_layout(
                lambda n,h,w,c: [n, c//4, h, w, c%4]
            )


        .. code-block:: python

            # ``A`` is a tensor whose compute definition is in an
            # arbitrary format, and should be transformed such that
            # the last index is split, with the slower-changing index
            # of the split placed at the slowest changing dimension.

            s[A].transform_layout(
                lambda *indices, i: [i//4, *indices, i%4]
            )

        .. code-block:: python

            # ``B`` is a tensor defined by te.compute to be a copy of
            # ``A`, and should be transformed such that ``B``'s layout
            # is a transpose of ``A``'s layout.  The loop iteration
            # that computes ``B`` will correspond to ``B``'s memory
            # layout.

            A = te.placeholder([n,m])
            B = te.compute(A.shape, lambda i,j: A[i,j])
            s = te.create_schedule(B.op)

            s[B].transform_layout(lambda i,j: [j,i])

        """

        ndim = len(self.op.output(0).shape)
        index_map, axis_separators = IndexMap.from_func_with_separators(mapping_function, ndim=ndim)

        new_iter_vars = _ffi_api.StageTransformLayout(
            self, index_map.initial_indices, index_map.final_indices
        )
        _ffi_api.StageSetAxisSeparators(self, axis_separators)

        return new_iter_vars or None


@tvm._ffi.register_object
class SpecializedCondition(Object):

    """Specialized condition to enable op specialization."""

    def __init__(self, conditions):
        """Create a specialized condition.

        .. note::
            Conditions are represented in conjunctive joint form (CNF).
            Each condition should be a simple expression, e.g., n > 16,
            m % 8 == 0, etc., where n, m are tvm.Var that represents a
            dimension in the tensor shape.

        Parameters
        ----------
        conditions : List of tvm.Expr
            List of conditions in conjunctive joint form (CNF).
        """
        if not isinstance(conditions, (list, _container.Array)):
            conditions = [conditions]
        self.__init_handle_by_constructor__(_ffi_api.CreateSpecializedCondition, conditions)

    @staticmethod
    def current():
        """Returns the current specialized condition"""
        return _ffi_api.GetCurrentSpecialization()

    def __enter__(self):
        _ffi_api.EnterSpecializationScope(self)
        return self

    def __exit__(self, ptype, value, trace):
        _ffi_api.ExitSpecializationScope(self)


# Sentinel value used to indicate which groups of pre-flattening axes
# should be used to post-flattening axes axes.  Moved from
# te.AXIS_SEPARATOR to tir.IndexMap.AXIS_SEPARATOR for general use,
# maintained here for backwards compatibility.
AXIS_SEPARATOR = IndexMap.AXIS_SEPARATOR


tvm._ffi._init_api("schedule", __name__)