Print auto_schedule Python schedule with topi op

comaniac · November 1, 2021, 4:17pm

This unfortunately cannot be controlled by the auto-scheduler. One good feature of auto-scheduler is that it schedules an entire fused compute graph together. On the other hand, it implies that we don’t differentiate ops once they are fused into a single TE compute. Specifically, the name of the TE tensors, such as PadInputs, T_reshape, and compute, are defined in the TOPI.

Reference:

github.com

apache/tvm/blob/813136401a11a49d6c15e6013c34dd822a5c4ff6/python/tvm/topi/nn/pad.py#L26


# under the License.
"""Pad the data by constant value """
from __future__ import absolute_import as _abs
import tvm
from tvm import te
from ..utils import equal_const_int
from .. import tag




@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
def pad(data, pad_before, pad_after=None, pad_value=0.0, name="PadInput"):
    """Pad Input with zeros.


    Parameters
    ----------
    data : tvm.te.Tensor
        n-D input, can be any layout.


    pad_before : list / tuple of n ints
        Pad width on each dimension to pad the before the axis begin.

github.com

apache/tvm/blob/5ad2f77403bed9a2bf356cc0d3d785ecc13e6c58/include/tvm/topi/transform.h#L235


/*!
 * \brief Reshape a tensor
 *
 * \param x The input tensor
 * \param newshape The new shape
 * \param name The name of the operation
 * \param tag The tag to mark the operation
 *
 * \return A Tensor whose op member is the reshape operation
 */
inline Tensor reshape(const Tensor& x, Array<PrimExpr> newshape, std::string name = "T_reshape",
                      std::string tag = kInjective) {
  auto x_shape = x->shape;
  Array<PrimExpr> target_shape;


  for (const auto& ele : newshape) {
    if (ele.as<IntImmNode>()) {
      target_shape.push_back(cast(DataType::Int(32), ele));
    } else {
      target_shape.push_back(ele);
    }

However, for your particular case, T_reshape seems come from the same op:

T_reshape_ax0, T_reshape_ax1, T_reshape_ax2, T_reshape_ax3 = \
  tuple(T_reshape.op.axis) + tuple(T_reshape.op.reduce_axis)

T_reshape_ax0, T_reshape_ax1, T_reshape_ax2, T_reshape_ax3 = \
  tuple(T_reshape.op.axis) + tuple(T_reshape.op.reduce_axis)

T_reshape_ax0, T_reshape_ax1, T_reshape_ax2, T_reshape_ax3 = \
  tuple(T_reshape.op.axis) + tuple(T_reshape.op.reduce_axis)

So IIUC, this is just the case that auto-scheduler printer wants to make sure the variable (e.g., T_reshape_ax0) it will be referring to is the right one, so it re-generates them using the same reshape op.