[Unity] Relax op: index (#13987)

This PR is about the high-level tensor computation operators in Relax.

This PR includes the tensor indexing operators.

include/tvm/relax/attrs/index.h

@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/attrs/index.h
+ * \brief Attributes for indexing operators.
+ */
+#ifndef TVM_RELAX_ATTRS_INDEX_H_
+#define TVM_RELAX_ATTRS_INDEX_H_
+
+#include <tvm/relax/expr.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in take operator */
+struct TakeAttrs : public tvm::AttrsNode<TakeAttrs> {
+  Optional<Integer> axis;
+
+  TVM_DECLARE_ATTRS(TakeAttrs, "relax.attrs.TakeAttrs") {
+    TVM_ATTR_FIELD(axis).describe("The axis over which to select values.");
+  }
+};  // struct TakeAttrs
+
+/*! \brief Attributes used in strided_slice operator */
+struct StridedSliceAttrs : public tvm::AttrsNode<StridedSliceAttrs> {
+  Array<Integer> axes;
+  Array<PrimExpr> begin;
+  Array<PrimExpr> end;
+  Optional<Array<PrimExpr>> strides;
+
+  TVM_DECLARE_ATTRS(StridedSliceAttrs, "relax.attrs.StridedSliceAttrs") {
+    TVM_ATTR_FIELD(axes).describe("Axes along which slicing is applied.");
+    TVM_ATTR_FIELD(begin).describe("The indices to begin with in the slicing, inclusive.");
+    TVM_ATTR_FIELD(end).describe("The indices indicating end of the slice, exclusive.");
+    TVM_ATTR_FIELD(strides).describe(
+        "Specifies the stride values, it can be negative in that case, the input tensor will be "
+        "reversed in that particular axis. If not specified, it by default is an list of ones of "
+        "the same length as `axes`.");
+  }
+};  // struct StridedSliceAttrs
+
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_ATTRS_INDEX_H_

python/tvm/relax/op/__init__.py

@@ -20,6 +20,8 @@
 # Operators
 from .base import *
 from .binary import *
+from .index import *
 from .manipulate import *
+from .op_attrs import *
 from . import builtin
 from . import memory

python/tvm/relax/op/index.py

@@ -0,0 +1,90 @@
+# 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.
+"""Indexing operators."""
+from typing import List, Optional, Union
+
+from tvm.ir.expr import PrimExpr
+
+from . import _ffi_api
+from ..expr import Expr
+
+PrimExprLike = Union[int, PrimExpr]
+
+
+def take(x: Expr, indices: Expr, axis: Optional[int] = None) -> Expr:
+    """Take elements from a tensor along an axis.
+
+    Parameters
+    ----------
+    x : relax.Expr
+        The source tensor.
+
+    indices : relax.Expr
+        The indices of the values to extract.
+        It is required to be a one-dimensional tensor which has integer dtype.
+
+    axis : Optional[int]
+        The axis over which to select values.
+        If it is none, the input tensor is required to be one-dimensional.
+
+    Returns
+    -------
+    ret : relax.Expr
+        The taken result.
+    """
+    return _ffi_api.take(x, indices, axis)  # type: ignore
+
+
+def strided_slice(
+    x: Expr,
+    axes: List[int],
+    begin: List[PrimExprLike],
+    end: List[PrimExprLike],
+    strides: Optional[List[PrimExprLike]] = None,
+) -> Expr:
+    """Strided slice of a tensor.
+
+    Parameters
+    ----------
+    x : relax.Expr
+        The source tensor to be sliced.
+
+    axes : List[int]
+        Axes along which slicing is applied.
+
+    begin : List[PrimExprLike]
+        The indices to begin with in the slicing, inclusive.
+
+    end : List[PrimExprLike]
+        The indices indicating end of the slice, exclusive.
+
+    strides : Optional[List[PrimExprLike]]
+        Specifies the stride values, it can be negative in that case,
+        the input tensor will be reversed in that particular axis.
+        If not specified, it by default is an list of ones of the same length as `axes`.
+
+    Returns
+    -------
+    ret : relax.Expr
+        The sliced result.
+
+    Note
+    ----
+    strided_slice require the input `begin`, `end` and `strides` to have the
+    same length as `axes`.
+    """
+    return _ffi_api.strided_slice(x, axes, begin, end, strides)  # type: ignore

python/tvm/relax/op/op_attrs.py

@@ -0,0 +1,29 @@
+# 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.
+"""The attributes node used for Relax operators"""
+from tvm.ir import Attrs
+import tvm._ffi
+
+
+@tvm._ffi.register_object("relax.attrs.TakeAttrs")
+class TakeAttrs(Attrs):
+    """Attributes used in take operator"""
+
+
+@tvm._ffi.register_object("relax.attrs.StridedSliceAttrs")
+class StridedSliceAttrs(Attrs):
+    """Attributes used in strided_slice operator"""

python/tvm/script/ir_builder/relax/ir.py

@@ -42,6 +42,8 @@
     print,
     reshape,
     shape_of,
+    strided_slice,
+    take,
 )
 from tvm.relax.struct_info import StructInfo
 from tvm.relax.utils import args_converter
@@ -427,5 +429,7 @@ def dtype(value: Union[py_str, DataType]) -> Expr:
     "shape",
     "shape_of",
     "str",
+    "strided_slice",
+    "take",
     "tuple",
 ]