transform.h

Go to the documentation of this file.

/*
 * 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.
 */

#ifndef TVM_RELAY_ATTRS_TRANSFORM_H_
#define TVM_RELAY_ATTRS_TRANSFORM_H_

#include <tvm/ir/attrs.h>
#include <tvm/relay/base.h>
#include <tvm/relay/expr.h>
#include <tvm/tir/index_map.h>

#include <string>

namespace tvm {
namespace relay {

struct SlidingWindowAttrs : public tvm::AttrsNode<SlidingWindowAttrs> {
  int axis;
  Array<Integer> window_shape;
  Array<Integer> strides;
  TVM_DECLARE_ATTRS(SlidingWindowAttrs, "relay.attrs.SlidingWindowAttrs") {
    TVM_ATTR_FIELD(axis).describe(
        "What axis the sliding window begin forming over."
        "Window will be slid over this axis and all following axes."
        "The axis value determines the window shape (and thus, the"
        "number of strides):"
        "window shape and strides must both be of length"
        "`data.ndim-axis`.");
    TVM_ATTR_FIELD(window_shape)
        .describe(
            "The window shape to form over the input."
            "Window shape must be of length `data.ndim-axis`.");
    TVM_ATTR_FIELD(strides).describe(
        "How to stride the window along each dimension."
        "Strides must be of length `data.ndim-axis`.");
  }
};  // struct SlidingWindowAttrs

struct CastAttrs : public tvm::AttrsNode<CastAttrs> {
  DataType dtype;

  TVM_DECLARE_ATTRS(CastAttrs, "relay.attrs.CastAttrs") {
    TVM_ATTR_FIELD(dtype).describe("Target data type");
  }
};  // struct CastAttrs.

struct ExpandDimsAttrs : public tvm::AttrsNode<ExpandDimsAttrs> {
  int axis;
  int num_newaxis;

  TVM_DECLARE_ATTRS(ExpandDimsAttrs, "relay.attrs.ExpandDimsAttrs") {
    TVM_ATTR_FIELD(axis).describe(
        "The axis at which the input array is expanded."
        "Should lie in range `[-data.ndim - 1, data.ndim]`."
        "If `axis < 0`, it is the first axis inserted;"
        "If `axis >= 0`, it is the last axis inserted in Python's negative indexing.");
    TVM_ATTR_FIELD(num_newaxis)
        .describe("Number of axes to be inserted. Should be >= 0.")
        .set_lower_bound(0)
        .set_default(1);
  }
};  // struct ExpandDimsAttrs

struct DynExpandDimsAttrs : public tvm::AttrsNode<DynExpandDimsAttrs> {
  int num_newaxis;

  TVM_DECLARE_ATTRS(DynExpandDimsAttrs, "relay.attrs.DynExpandDimsAttrs") {
    TVM_ATTR_FIELD(num_newaxis)
        .describe("Number of axes to be inserted. Should be >= 0.")
        .set_lower_bound(0)
        .set_default(1);
  }
};  // struct ExpandDimsAttrs

struct ConcatenateAttrs : public tvm::AttrsNode<ConcatenateAttrs> {
  int axis;
  TVM_DECLARE_ATTRS(ConcatenateAttrs, "relay.attrs.ConcatenateAttrs") {
    TVM_ATTR_FIELD(axis)
        .describe(
            "The axis at which the input arrays are concatenated."
            "Should lie in range `[-ndim, ndim)`.")
        .set_default(0);
  }
};  // struct ConcatenateAttrs

struct TransposeAttrs : public tvm::AttrsNode<TransposeAttrs> {
  Array<Integer> axes;
  TVM_DECLARE_ATTRS(TransposeAttrs, "relay.attrs.TransposeAttrs") {
    TVM_ATTR_FIELD(axes).describe("The target axes order, reverse order if not specified.");
  }
};  // struct TransposeAttrs

struct ReshapeAttrs : public tvm::AttrsNode<ReshapeAttrs> {
  Array<Integer> newshape;
  bool allowzero;
  TVM_DECLARE_ATTRS(ReshapeAttrs, "relay.attrs.ReshapeAttrs") {
    TVM_ATTR_FIELD(newshape).describe(
        "The new shape. Should be compatible with the original shape.");
    TVM_ATTR_FIELD(allowzero).set_default(0).describe(
        "Whether to honor the value of zero in newshape.");
  }
};  // struct ReshapeAttrs

struct ReshapeLikeAttrs : public tvm::AttrsNode<ReshapeLikeAttrs> {
  int lhs_begin;
  Integer lhs_end;  // can be None
  int rhs_begin;
  Integer rhs_end;  // can be None
  TVM_DECLARE_ATTRS(ReshapeLikeAttrs, "relay.attrs.ReshapeLikeAttrs") {
    TVM_ATTR_FIELD(lhs_begin).set_default(0).describe(
        "The axis of the input where reshaping should begin.");
    TVM_ATTR_FIELD(lhs_end)
        .set_default(NullValue<Integer>())
        .describe("The axis of the input where reshaping should end, exclusive.");
    TVM_ATTR_FIELD(rhs_begin).set_default(0).describe(
        "The axis of the shape_like tensor to begin taking dimensions from.");
    TVM_ATTR_FIELD(rhs_end)
        .set_default(NullValue<Integer>())
        .describe("The axis of the shape_like tensor to end taking dimensions from, exclusive.");
  }
};  // struct ReshapeLikeAttrs

struct ScatterAttrs : public tvm::AttrsNode<ScatterAttrs> {
  Integer axis;

  TVM_DECLARE_ATTRS(ScatterAttrs, "relay.attrs.ScatterAttrs") {
    TVM_ATTR_FIELD(axis).set_default(0).describe("The axis over which to select values.");
  }
};

struct ScatterAddAttrs : public tvm::AttrsNode<ScatterAddAttrs> {
  Integer axis;

  TVM_DECLARE_ATTRS(ScatterAddAttrs, "relay.attrs.ScatterAddAttrs") {
    TVM_ATTR_FIELD(axis).set_default(0).describe("The axis over which to select values.");
  }
};

struct ScatterNDAttrs : public tvm::AttrsNode<ScatterNDAttrs> {
  String mode;

  TVM_DECLARE_ATTRS(ScatterNDAttrs, "relay.attrs.ScatterNDAttrs") {
    TVM_ATTR_FIELD(mode).describe(
        "Accumulation mode of the scatter, either \"update\" or \"add\".");
  }
};

struct GatherAttrs : public tvm::AttrsNode<GatherAttrs> {
  Integer axis;

  TVM_DECLARE_ATTRS(GatherAttrs, "relay.attrs.GatherAttrs") {
    TVM_ATTR_FIELD(axis)
        .set_default(NullValue<Integer>())
        .describe("The axis over which to select values.");
  }
};

struct GatherNDAttrs : public tvm::AttrsNode<GatherNDAttrs> {
  Integer batch_dims;
  Optional<Integer> index_rank;

  TVM_DECLARE_ATTRS(GatherNDAttrs, "relay.attrs.GatherNDAttrs") {
    TVM_ATTR_FIELD(batch_dims).set_default(Integer(0)).describe("The number of batch dimensions.");
    TVM_ATTR_FIELD(index_rank)
        .set_default(NullValue<Integer>())
        .describe(
            "The size of an indexing tuple, which is a fixed value. Only needed when the number of "
            "indexting tuples is dynamic.");
  }
};

struct TakeAttrs : public tvm::AttrsNode<TakeAttrs> {
  Integer batch_dims;
  Integer axis;
  tvm::String mode;

  TVM_DECLARE_ATTRS(TakeAttrs, "relay.attrs.TakeAttrs") {
    TVM_ATTR_FIELD(batch_dims)
        .set_default(0)
        .describe("The batch_dims over which to select values.");
    TVM_ATTR_FIELD(axis)
        .set_default(NullValue<Integer>())
        .describe("The axis over which to select values.");
    TVM_ATTR_FIELD(mode).set_default("clip").describe(
        "Specify how out-of-bound indices will behave."
        "clip - clip to the range (default)"
        "wrap - wrap around the indices"
        "fast - no clip or wrap around (user must make sure indices are in-bound)");
  }
};

struct InitOpAttrs : public tvm::AttrsNode<InitOpAttrs> {
  Optional<Array<Integer>> shape;
  DataType dtype;

  TVM_DECLARE_ATTRS(InitOpAttrs, "relay.attrs.InitOpAttrs") {
    TVM_ATTR_FIELD(shape).describe("Target shape.");
    TVM_ATTR_FIELD(dtype).describe("Target data type.").set_default(NullValue<DataType>());
  }
};  // struct InitOpAttrs

struct ArangeAttrs : public tvm::AttrsNode<ArangeAttrs> {
  Expr start;
  Expr stop;
  Expr step;
  DataType dtype;

  TVM_DECLARE_ATTRS(ArangeAttrs, "relay.attrs.ArangeAttrs") {
    TVM_ATTR_FIELD(start).describe("Start of interval. The interval includes this value.");
    TVM_ATTR_FIELD(stop).describe("Stop of interval. The interval does not include this value.");
    TVM_ATTR_FIELD(step).describe("Spacing between values.");
    TVM_ATTR_FIELD(dtype).describe("Target data type.");
  }
};  // struct ArangeAttrs

struct MeshgridAttrs : public tvm::AttrsNode<MeshgridAttrs> {
  std::string indexing;

  TVM_DECLARE_ATTRS(MeshgridAttrs, "relay.attrs.MeshgridAttrs") {
    TVM_ATTR_FIELD(indexing)
        .describe(
            "Indexing mode, either \"ij\" for matrix or \"xy\" for cartesian in which first two"
            "dimensions are swapped.")
        .set_default("ij");
  }
};  // struct MeshgridAttrs

struct StackAttrs : public tvm::AttrsNode<StackAttrs> {
  Integer axis;
  TVM_DECLARE_ATTRS(StackAttrs, "relay.attrs.StackAttrs") {
    TVM_ATTR_FIELD(axis).set_default(0).describe(
        "The axis in the result array along which the input arrays are stacked.");
  }
};  // struct StackAttrs

struct RepeatAttrs : public tvm::AttrsNode<RepeatAttrs> {
  Integer repeats;
  Integer axis;
  TVM_DECLARE_ATTRS(RepeatAttrs, "relay.attrs.RepeatAttrs") {
    TVM_ATTR_FIELD(repeats).describe("The number of repetitions for each element.");
    TVM_ATTR_FIELD(axis)
        .set_default(NullValue<Integer>())
        .describe(" The axis along which to repeat values.");
  }
};  // struct RepeatAttrs

struct TileAttrs : public tvm::AttrsNode<TileAttrs> {
  Array<Integer> reps;
  TVM_DECLARE_ATTRS(TileAttrs, "relay.attrs.TileAttrs") {
    TVM_ATTR_FIELD(reps).describe(
        "The number of times for repeating the tensor a."
        "Each dim sizeof reps must be a positive integer.");
  }
};  // struct TileAttrs

struct ReverseAttrs : public tvm::AttrsNode<ReverseAttrs> {
  Integer axis;
  TVM_DECLARE_ATTRS(ReverseAttrs, "relay.attrs.ReverseAttrs") {
    TVM_ATTR_FIELD(axis)
        .set_default(NullValue<Integer>())
        .describe("The axis along which to reverse elements.");
  }
};  // struct ReverseAttrs

struct ReverseSequenceAttrs : public tvm::AttrsNode<ReverseSequenceAttrs> {
  Integer seq_axis;
  Integer batch_axis;

  TVM_DECLARE_ATTRS(ReverseSequenceAttrs, "relay.attrs.ReverseSequenceAttrs") {
    TVM_ATTR_FIELD(seq_axis).set_default(1).describe(
        "The seq axis along which to reverse elements.");
    TVM_ATTR_FIELD(batch_axis)
        .set_default(0)
        .describe("The batch axis along which to slice the tensor.");
  }
};  // struct ReverseSequenceAttrs

struct SqueezeAttrs : public tvm::AttrsNode<SqueezeAttrs> {
  // use axis to make the name numpy compatible.
  Array<Integer> axis;

  TVM_DECLARE_ATTRS(SqueezeAttrs, "relay.attrs.SqueezeAttrs") {
    TVM_ATTR_FIELD(axis)
        .describe(
            "The axis to squeeze in the input tensor."
            "If `axis = None`, all axis of dimension 1 get squeezed;"
            "Else, the dimension in axes get squeezed."
            "It is an error if an axis does not has dimension 1.")
        .set_default(NullValue<Array<Integer>>());
  }
};  // struct SqueezeAttrs

struct SplitAttrs : public tvm::AttrsNode<SplitAttrs> {
  ObjectRef indices_or_sections;
  int axis;

  TVM_DECLARE_ATTRS(SplitAttrs, "relay.attrs.SplitAttrs") {
    TVM_ATTR_FIELD(indices_or_sections)
        .describe(
            "Indices or sections to split into. Accepts an int or a tuple"
            "If indices_or_sections is an integer, the input will be divided equally"
            "along given axis. If such a split is not possible, an error is raised."
            "If indices_or_sections is a tuple of sorted integers,"
            "the entries indicate where along axis the array is split.");
    TVM_ATTR_FIELD(axis).set_default(0).describe("the axis to be splitted.");
  }
};

struct StridedSliceAttrs : public tvm::AttrsNode<StridedSliceAttrs> {
  Optional<Array<Integer>> begin;
  Optional<Array<Integer>> end;
  Optional<Array<Integer>> strides;
  tvm::String slice_mode;
  Optional<Array<Integer>> axes;

  TVM_DECLARE_ATTRS(StridedSliceAttrs, "relay.attrs.StridedSliceAttrs") {
    TVM_ATTR_FIELD(begin).describe("Indices for begin of slice, begin index is also inclusive");
    TVM_ATTR_FIELD(end).describe("Indices for end of slice, end index is exclusive");
    TVM_ATTR_FIELD(strides).describe(
        "Stride values of the slice, a stride can be negative, which causes a reverse slice.");
    TVM_ATTR_FIELD(slice_mode)
        .set_default("end")
        .describe(
            "The slice mode [end, size]."
            "end - The default slice mode, ending indices for the slice."
            "size - The input strides will be ignored, input end in this mode indicates the size"
            "of a slice starting at the location specified by begin. If end[i] is -1,"
            "all remaining elements in that dimension are included in the slice");
    TVM_ATTR_FIELD(axes).describe(
        "Axes along which slicing is applied. When it is specified, the length of begin, end, "
        "strides, and axes must be equal.");
  }
};

struct SliceLikeAttrs : public tvm::AttrsNode<SliceLikeAttrs> {
  Array<Integer> axes;

  TVM_DECLARE_ATTRS(SliceLikeAttrs, "relay.attrs.SliceLikeAttrs") {
    TVM_ATTR_FIELD(axes).describe(
        "List of axes on which input data will be sliced according to the "
        "corresponding size of the second input. By default will slice "
        "on all axes. Negative axes mean counting in reverse.");
  }
};

struct ClipAttrs : public tvm::AttrsNode<ClipAttrs> {
  double a_min;
  double a_max;

  TVM_DECLARE_ATTRS(ClipAttrs, "relay.attrs.ClipAttrs") {
    TVM_ATTR_FIELD(a_min).describe("The minimum clip value.");
    TVM_ATTR_FIELD(a_max).describe("The maximum clip value.");
  }
};

struct FixedPointMultiplyAttrs : public tvm::AttrsNode<FixedPointMultiplyAttrs> {
  int32_t multiplier;
  int32_t shift;

  TVM_DECLARE_ATTRS(FixedPointMultiplyAttrs, "relay.attrs.FixedPointMultiplyAttrs") {
    TVM_ATTR_FIELD(multiplier)
        .describe("Multiplier of a fixed floating point number described as multiplier*2^(shift)");
    TVM_ATTR_FIELD(shift).describe(
        "Shift of a fixed floating point number described as multiplier*2^(shift)");
  }
};

struct LayoutTransformAttrs : public tvm::AttrsNode<LayoutTransformAttrs> {
  std::string src_layout;
  std::string dst_layout;

  TVM_DECLARE_ATTRS(LayoutTransformAttrs, "relay.attrs.LayoutTransformAttrs") {
    TVM_ATTR_FIELD(src_layout).describe("The source layout of the tensor. (e.g. NCHW)");
    TVM_ATTR_FIELD(dst_layout).describe("The destination layout of the tensor. (e.g. NCHW16c)");
  }
};

struct AutoSchedulerLayoutTransformAttrs
    : public tvm::AttrsNode<AutoSchedulerLayoutTransformAttrs> {
  std::string src_layout;
  std::string dst_layout;

  TVM_DECLARE_ATTRS(AutoSchedulerLayoutTransformAttrs,
                    "relay.attrs.AutoSchedulerLayoutTransformAttrs") {
    TVM_ATTR_FIELD(src_layout).describe("The source layout of the tensor. (e.g. 1N32C112H112W)");
    TVM_ATTR_FIELD(dst_layout)
        .describe("The destination layout of the tensor. (e.g. 1N2C112H112W16c)");
  }
};

struct MetaScheduleLayoutTransformAttrs : public tvm::AttrsNode<MetaScheduleLayoutTransformAttrs> {
  tir::IndexMap index_map;

  TVM_DECLARE_ATTRS(MetaScheduleLayoutTransformAttrs,
                    "relay.attrs.MetaScheduleLayoutTransformAttrs") {
    TVM_ATTR_FIELD(index_map).describe(
        "The order of the extents, for example, "
        "let extents = [2, 3, 4], reorder = [0, 2, 1], and the shape of buffer A is (4, 6)"
        "then A[i, j] will be first rewritten to "
        "A[(6 * i + j) / 12, (6 * i + j) / 4 % 3 , (6 * i + j) % 4] according to the `extents`,"
        "and then reordered to A[(6 * i + j) / 12, (6 * i + j) % 4 , (6 * i + j) / 4 % 3]"
        "according to `reorder`");
  }
};

struct ShapeOfAttrs : public tvm::AttrsNode<ShapeOfAttrs> {
  DataType dtype;

  TVM_DECLARE_ATTRS(ShapeOfAttrs, "relay.attrs.ShapeOfAttrs") {
    TVM_ATTR_FIELD(dtype).describe("Target data type").set_default(NullValue<DataType>());
  }
};

struct SequenceMaskAttrs : public tvm::AttrsNode<SequenceMaskAttrs> {
  double mask_value;
  int axis;

  TVM_DECLARE_ATTRS(SequenceMaskAttrs, "relay.attrs.SequenceMaskAttrs") {
    TVM_ATTR_FIELD(mask_value).set_default(0).describe("The masking value.");
    TVM_ATTR_FIELD(axis).set_default(0).describe(
        "The axis of the length dimension. Can only be 0 or 1.");
  }
};  // struct SequenceMaskAttrs.

struct SparseToDenseAttrs : public tvm::AttrsNode<SparseToDenseAttrs> {
  Array<Integer> output_shape;

  TVM_DECLARE_ATTRS(SparseToDenseAttrs, "relay.attrs.SparseToDenseAttrs") {
    TVM_ATTR_FIELD(output_shape).describe("Shape of the dense output tensor");
  }
};  // struct SparseToDenseAttrs

struct NdarraySizeAttrs : public tvm::AttrsNode<NdarraySizeAttrs> {
  DataType dtype;

  TVM_DECLARE_ATTRS(NdarraySizeAttrs, "relay.attrs.NdarraySizeAttrs") {
    TVM_ATTR_FIELD(dtype).describe("Target data type").set_default(NullValue<DataType>());
  }
};

struct OneHotAttrs : public tvm::AttrsNode<OneHotAttrs> {
  int depth;
  int axis;
  DataType dtype;

  TVM_DECLARE_ATTRS(OneHotAttrs, "relay.attrs.OneHotAttrs") {
    TVM_ATTR_FIELD(depth).set_default(1).describe("Depth of the one hot dimension.");
    TVM_ATTR_FIELD(axis).set_default(-1).describe("Axis to fill.");
    TVM_ATTR_FIELD(dtype).set_default(NullValue<DataType>()).describe("Output data type.");
  }
};  // struct OneHotAttrs

struct MatrixSetDiagAttrs : public tvm::AttrsNode<MatrixSetDiagAttrs> {
  int k1;
  int k2;
  bool super_diag_right_align;
  bool sub_diag_right_align;

  TVM_DECLARE_ATTRS(MatrixSetDiagAttrs, "relay.attrs.MatrixSetDiagAttrs") {
    TVM_ATTR_FIELD(k1).set_default(0).describe("Lower limit (included) of the range of diagonals.");
    TVM_ATTR_FIELD(k2).set_default(0).describe("Upper limit (included) of the range of diagonals.");
    TVM_ATTR_FIELD(super_diag_right_align)
        .set_default(true)
        .describe("Bool, true iff super-diagonal is right aligned (left-padded).");
    TVM_ATTR_FIELD(sub_diag_right_align)
        .set_default(false)
        .describe("Bool, true iff sub-diagonal is right aligned (left-padded).");
  }
};  // struct MatrixSetDiagAttrs

struct ScanopAttrs : public tvm::AttrsNode<ScanopAttrs> {
  Integer axis;
  DataType dtype;
  Bool exclusive = Bool(false);
  TVM_DECLARE_ATTRS(ScanopAttrs, "relay.attrs.ScanopAttrs") {
    TVM_ATTR_FIELD(axis).describe("The axis to operate over").set_default(NullValue<Integer>());
    TVM_ATTR_FIELD(dtype).describe("Output data type").set_default(NullValue<DataType>());

    // Default is 0 which is "false"
    TVM_ATTR_FIELD(exclusive)
        .describe("The first element is not included")
        .set_default(Bool(false));
  }
};  // struct ScanopAttrs

struct UniqueAttrs : public tvm::AttrsNode<UniqueAttrs> {
  bool sorted;
  bool return_counts;
  TVM_DECLARE_ATTRS(UniqueAttrs, "relay.attrs.UniqueAttrs") {
    TVM_ATTR_FIELD(sorted).describe("Whether the unique elements are sorted").set_default(true);
    TVM_ATTR_FIELD(return_counts)
        .describe("Whether to return an additional tensor with counts of each unique elements")
        .set_default(false);
  }
};  // struct UniqueAttrs

struct EinsumAttrs : public tvm::AttrsNode<EinsumAttrs> {
  String equation;

  TVM_DECLARE_ATTRS(EinsumAttrs, "relay.attrs.EinsumAttrs") {
    TVM_ATTR_FIELD(equation).describe("The einsum expression string");
  }
};  // struct EinsumAttrs

struct StftAttrs : public tvm::AttrsNode<StftAttrs> {
  int n_fft;
  int hop_length;
  int win_length;
  bool normalized;
  bool onesided;

  TVM_DECLARE_ATTRS(StftAttrs, "relay.attrs.StftAttrs") {
    TVM_ATTR_FIELD(n_fft).set_default(-1).describe("The size of Fourier transform");
    TVM_ATTR_FIELD(hop_length)
        .set_default(-1)
        .describe("The distance between neighboring sliding window frames");
    TVM_ATTR_FIELD(win_length).set_default(-1).describe("The size of window frame and STFT filter");
    TVM_ATTR_FIELD(normalized)
        .set_default(false)
        .describe("Whether to return the normalized STFT results");
    TVM_ATTR_FIELD(onesided).set_default(true).describe(
        "Whether to return onesided result or fill with conjugate symmetry");
  }
};  // struct StftAttrs

}  // namespace relay
}  // namespace tvm
#endif  // TVM_RELAY_ATTRS_TRANSFORM_H_