This is an automated email from the ASF dual-hosted git repository.
xiyou pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 34695afa5c [Unity][Op] introduce `ScatterElement` op (#14493)
34695afa5c is described below
commit 34695afa5c80b8795e32552933f5ff32d7245049
Author: Sunghyun Park <[email protected]>
AuthorDate: Thu Apr 6 14:54:52 2023 -0700
[Unity][Op] introduce `ScatterElement` op (#14493)
* feat: introduce scatter_element op
* fix whitespace
* fix typo
---
include/tvm/relax/attrs/manipulate.h | 12 ++
python/tvm/relax/op/manipulate.py | 70 ++++++++
.../tvm/relax/transform/legalize_ops/manipulate.py | 12 ++
python/tvm/script/ir_builder/relax/ir.py | 2 +
python/tvm/topi/scatter_elements.py | 2 +-
src/relax/op/tensor/manipulate.cc | 114 ++++++++++++
tests/python/relax/test_op_manipulate.py | 140 +++++++++++++++
.../test_transform_legalize_ops_manipulate.py | 198 +++++++++++++++++++++
8 files changed, 549 insertions(+), 1 deletion(-)
diff --git a/include/tvm/relax/attrs/manipulate.h
b/include/tvm/relax/attrs/manipulate.h
index 4aa51f2b73..be45dadc9b 100644
--- a/include/tvm/relax/attrs/manipulate.h
+++ b/include/tvm/relax/attrs/manipulate.h
@@ -140,6 +140,18 @@ struct CumsumAttrs : public tvm::AttrsNode<CumsumAttrs> {
}
}; // struct CumsumAttrs
+/*! \brief Attributes used in scatter_elements operators */
+struct ScatterElementsAttrs : public tvm::AttrsNode<ScatterElementsAttrs> {
+ Integer axis;
+ String reduction;
+
+ TVM_DECLARE_ATTRS(ScatterElementsAttrs, "relax.attrs.ScatterElementsAttrs") {
+ TVM_ATTR_FIELD(axis).set_default(0).describe("The axis over which to
select values.");
+ TVM_ATTR_FIELD(reduction).set_default("update").describe(
+ "Reduction mode of the scatter elements, "
+ "either \"update\", \"add\", \"mul\", \"mean\", \"min\" or \"max\".");
+ }
+}; // struct ScatterElementsAttrs
} // namespace relax
} // namespace tvm
diff --git a/python/tvm/relax/op/manipulate.py
b/python/tvm/relax/op/manipulate.py
index e9c3ce79d7..ce3df4499e 100644
--- a/python/tvm/relax/op/manipulate.py
+++ b/python/tvm/relax/op/manipulate.py
@@ -439,3 +439,73 @@ def cumsum(data: Expr, axis: Optional[int] = None, dtype:
Optional[Union[str, Da
-> [1, 1, 2, 2, 3, 4, 4]
"""
return _ffi_api.cumsum(data, axis, dtype) # type: ignore
+
+
+def scatter_elements(
+ data: Expr, indices: Expr, updates: Expr, axis: int = 0, reduction: str =
"update"
+):
+ """ONNX style scatter elements. This operation updates its value in `data`
to values
+ specified by `updates` at specific index positions specified by `indices`.
+ For example, in 2D tensor, the update corresponding to the [i][j] entry is
performed
+ as below:
+ output[indices[i][j]][j] = updates[i][j] if axis = 0
+ output[i][indices[i][j]] = updates[i][j] if axis = 1
+
+ When the `reduction` is set to some reduction function `f`, the update
corresponding to
+ [i][j] entry is performed as below:
+ output[indices[i][j]][j] += f(output[indices[i][j]][j], updates[i][j])
if axis = 0
+ output[i][indices[i][j]] += f(output[i][indices[i][j]], updates[i][j])
if axis = 1
+ Where `f` is update, add, mul, mean, max, min.
+
+ Parameters
+ ----------
+ data : relax.Expr
+ The input data to the operator.
+
+ indices: relax.Expr
+ The index positions to update in `data`.
+
+ updates: relax.Expr
+ Values to replace to.
+
+ axis: int
+ Axis to scatter on.
+
+ reduction: str
+ Type of reduction to apply: update, add, mul, mean, max, min.
+ It is "update" by default.
+
+ Returns
+ -------
+ result : relax.Expr
+ The result has the same size as data, and the same shape as data
+
+ Examples
+ --------
+ .. code-block:: python
+ # inputs
+ data = [
+ [0.0, 0.0, 0.0],
+ [0.0, 0.0, 0.0],
+ [0.0, 0.0, 0.0],
+ ]
+ indices = [
+ [1, 0, 2],
+ [0, 2, 1],
+ ]
+ updates = [
+ [1.0, 1.1, 1.2],
+ [2.0, 2.1, 2.2],
+ ]
+ axis = 0
+ reduction = "update"
+
+ # output P
+ output = [
+ [2.0, 1.1, 0.0]
+ [1.0, 0.0, 2.2]
+ [0.0, 2.1, 1.2]
+ ]
+
+ """
+ return _ffi_api.scatter_elements(data, indices, updates, axis, reduction)
# type: ignore
diff --git a/python/tvm/relax/transform/legalize_ops/manipulate.py
b/python/tvm/relax/transform/legalize_ops/manipulate.py
index 144ef04748..d226a48eed 100644
--- a/python/tvm/relax/transform/legalize_ops/manipulate.py
+++ b/python/tvm/relax/transform/legalize_ops/manipulate.py
@@ -151,3 +151,15 @@ def _tile(bb: BlockBuilder, call: Call) -> Expr:
@register_legalize("relax.cumsum")
def _cumsum(bb: BlockBuilder, call: Call) -> Expr:
return bb.call_te(topi.cumsum, call.args[0], call.attrs.axis,
call.attrs.dtype)
+
+
+@register_legalize("relax.scatter_elements")
+def _scatter_elements(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(
+ topi.scatter_elements,
+ call.args[0],
+ call.args[1],
+ call.args[2],
+ call.attrs.axis,
+ call.attrs.reduction,
+ )
diff --git a/python/tvm/script/ir_builder/relax/ir.py
b/python/tvm/script/ir_builder/relax/ir.py
index 9857853a80..4630a850bf 100644
--- a/python/tvm/script/ir_builder/relax/ir.py
+++ b/python/tvm/script/ir_builder/relax/ir.py
@@ -58,6 +58,7 @@ from tvm.relax.op import (
cos,
cosh,
cumsum,
+ scatter_elements,
divide,
equal,
ewise_fma,
@@ -567,6 +568,7 @@ __all__ = [
"cosh",
"const",
"cumsum",
+ "scatter_elements",
"dataflow",
"divide",
"dtype",
diff --git a/python/tvm/topi/scatter_elements.py
b/python/tvm/topi/scatter_elements.py
index 4c35578ffd..08fcb866b4 100644
--- a/python/tvm/topi/scatter_elements.py
+++ b/python/tvm/topi/scatter_elements.py
@@ -33,7 +33,7 @@ def scatter_elements(data, indices, updates, axis=0,
reduction="update"):
output[indices[i][j]][j] = f(output[indices[i][j]][j], updates[i][j])
if axis = 0
output[i][indices[i][j]] = f(output[i][indices[i][j]], updates[i][j])
if axis = 1
- where the update function f is determinted by the reduction.
+ where the update function f is determined by the reduction.
Five types of the function are supported: "update", "add", "mul", "min"
and "max" (see below)
Parameters
diff --git a/src/relax/op/tensor/manipulate.cc
b/src/relax/op/tensor/manipulate.cc
index faa5ee3bc0..c25ee94d38 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -1359,5 +1359,119 @@ TVM_REGISTER_OP("relax.cumsum")
.add_argument("data", "Tensor", "The input tensor.")
.set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoCumsum);
+/* relax.scatter_elements */
+TVM_REGISTER_NODE_TYPE(ScatterElementsAttrs);
+
+Expr scatter_elements(Expr data, Expr indices, Expr updates, int axis, String
reduction) {
+ auto attrs = make_object<ScatterElementsAttrs>();
+ attrs->axis = std::move(axis);
+ attrs->reduction = std::move(reduction);
+ static const Op& op = Op::Get("relax.scatter_elements");
+ return Call(op, {data, indices, updates}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.scatter_elements").set_body_typed(scatter_elements);
+
+StructInfo InferStructInfoScatterElements(const Call& call, const
BlockBuilder& ctx) {
+ arith::Analyzer* analyzer = ctx->GetAnalyzer();
+ const auto* data_sinfo =
GetStructInfoAs<TensorStructInfoNode>(call->args[0]);
+ const auto* indices_sinfo =
GetStructInfoAs<TensorStructInfoNode>(call->args[1]);
+ const auto* updates_sinfo =
GetStructInfoAs<TensorStructInfoNode>(call->args[2]);
+
+ auto diag_def = [&](const TensorStructInfoNode* sinfo, String name, String
type_key) {
+ if (sinfo == nullptr) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "ScatterElements requires the input " << name
+ << " to be a Tensor. However, the given one is " <<
type_key);
+ }
+ };
+
+ diag_def(data_sinfo, "data", call->args[0]->struct_info_->GetTypeKey());
+ diag_def(indices_sinfo, "indices",
call->args[1]->struct_info_->GetTypeKey());
+ diag_def(updates_sinfo, "updates",
call->args[2]->struct_info_->GetTypeKey());
+
+ if (data_sinfo->IsUnknownNdim()) {
+ // When `data` has unknown rank, assume rest of arguments are correct and
proceed.
+ // If the assumption turns out to be wrong, runtime error will be
triggered.
+ return TensorStructInfo(data_sinfo->dtype, kUnknownNDim);
+ }
+
+ if (!indices_sinfo->IsUnknownNdim() && !updates_sinfo->IsUnknownNdim()) {
+ if (data_sinfo->ndim != indices_sinfo->ndim) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "ScatterElements op requires the data tensor to have
the same rank with "
+ "indices tensor. However, the given dimensions are "
+ << "indices: " << indices_sinfo->ndim << ", data: " <<
data_sinfo->ndim);
+ }
+
+ if (indices_sinfo->ndim != updates_sinfo->ndim) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "ScatterElements op requires the indices tensor to have the same
rank with "
+ "updates tensor. However, the given dimensions are "
+ << "indices: " << indices_sinfo->ndim << ", updates: " <<
updates_sinfo->ndim);
+ }
+ }
+
+ if (data_sinfo->IsUnknownDtype() || updates_sinfo->IsUnknownDtype()) {
+ auto diag_dtype = [&](const TensorStructInfoNode* sinfo, String name) {
+ if (sinfo->IsUnknownDtype()) {
+ // TODO(tvm-team): Do we have an equivalent of `ctx->ReportFatal` for
warning?
+ LOG(WARNING) << "Data type of " << name
+ << " has not been specified. Assume it has an integer
type.";
+ }
+ };
+ diag_dtype(data_sinfo, "data");
+ diag_dtype(data_sinfo, "updates");
+ } else {
+ if (data_sinfo->dtype != updates_sinfo->dtype) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "ScatterElements op requires the input data to have
same type with "
+ "updates. However, the given types are "
+ << "data: " << data_sinfo->dtype << ", updates: " <<
updates_sinfo->dtype);
+ }
+ }
+
+ if (indices_sinfo->IsUnknownDtype()) {
+ // TODO(tvm-team): Do we have an equivalent of `ctx->ReportFatal` for
warning?
+ LOG(WARNING) << "Data type of indice has not been specified. Assume it has
an integer type.";
+ } else if (!(indices_sinfo->dtype.is_int() ||
indices_sinfo->dtype.is_uint())) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "ScatterElements op requires the input indices to have integer
dtype. However, the "
+ "given indices dtype is "
+ << indices_sinfo->dtype);
+ }
+
+ const auto* indices_shape = indices_sinfo->shape.as<ShapeExprNode>();
+ const auto* updates_shape = updates_sinfo->shape.as<ShapeExprNode>();
+ if (indices_shape && updates_shape) {
+ for (int i = 0; i < indices_sinfo->ndim; i++) {
+ if (analyzer->CanProve(indices_shape->values[i] !=
updates_shape->values[i])) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "ScatterElements op requires the indices tensor to have the
same shape with "
+ "updates tensor. However, the given shapes are "
+ << "indices: " << ShapeExpr(indices_shape->values)
+ << ", updates: " << ShapeExpr(updates_shape->values));
+ }
+ }
+ }
+ const auto* data_shape = data_sinfo->shape.as<ShapeExprNode>();
+ if (data_shape) {
+ return TensorStructInfo(ShapeExpr(data_shape->values), data_sinfo->dtype);
+ }
+ return TensorStructInfo(data_sinfo->dtype, data_sinfo->ndim);
+}
+
+// TODO(relax-team): implement FRelaxInferLayout for scatter_elements
+TVM_REGISTER_OP("relax.scatter_elements")
+ .set_attrs_type<ScatterElementsAttrs>()
+ .set_num_inputs(3)
+ .add_argument("data", "Tensor", "The input tensor.")
+ .add_argument("indices", "Tensor", "The indices tensor.")
+ .add_argument("updates", "Tensor", "The input tensor of updates.")
+ .set_attr<FInferStructInfo>("FInferStructInfo",
InferStructInfoScatterElements);
+
} // namespace relax
} // namespace tvm
diff --git a/tests/python/relax/test_op_manipulate.py
b/tests/python/relax/test_op_manipulate.py
index 3edf63764a..674287fcf3 100644
--- a/tests/python/relax/test_op_manipulate.py
+++ b/tests/python/relax/test_op_manipulate.py
@@ -43,6 +43,7 @@ def test_op_correctness():
y = relax.Var("x", R.Tensor((4, 5), "float32"))
assert relax.op.collapse_sum_like(x, y).op ==
Op.get("relax.collapse_sum_like")
assert relax.op.cumsum(x, axis=1, dtype="int32").op ==
Op.get("relax.cumsum")
+ assert relax.op.scatter_elements(x, x, x).op ==
Op.get("relax.scatter_elements")
def _check_inference(bb: relax.BlockBuilder, call: relax.Call, expected_sinfo:
relax.StructInfo):
@@ -3022,5 +3023,144 @@ def test_cumsum_infer_struct_info_wrong_input_type():
bb.normalize(relax.op.cumsum(x1, axis=1))
+def test_scatter_elements_infer_struct_info():
+ bb = relax.BlockBuilder()
+ d0 = relax.Var("data", R.Tensor((4, 4), "float32"))
+ d1 = relax.Var("data", R.Tensor(dtype="float32", ndim=2))
+ d2 = relax.Var("data", R.Tensor("float32"))
+ i0 = relax.Var("indices", R.Tensor((2, 2), "int64"))
+ i1 = relax.Var("indices", R.Tensor((2, 2)))
+ i2 = relax.Var("indices", R.Tensor(dtype="int64", ndim=2))
+ i3 = relax.Var("indices", R.Tensor(ndim=2))
+ u0 = relax.Var("updates", R.Tensor((2, 2), "float32"))
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d0, i0, u0, 0, "updates"),
+ relax.TensorStructInfo((4, 4), dtype="float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d1, i0, u0, 0, "updates"),
+ relax.TensorStructInfo(dtype="float32", ndim=2),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d2, i0, u0, 0, "updates"),
+ relax.TensorStructInfo(dtype="float32", ndim=-1),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d0, i1, u0, 0, "updates"),
+ relax.TensorStructInfo((4, 4), dtype="float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d1, i1, u0, 0, "updates"),
+ relax.TensorStructInfo(dtype="float32", ndim=2),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d2, i1, u0, 0, "updates"),
+ relax.TensorStructInfo(dtype="float32", ndim=-1),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d0, i2, u0, 0, "updates"),
+ relax.TensorStructInfo((4, 4), dtype="float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d1, i2, u0, 0, "updates"),
+ relax.TensorStructInfo(dtype="float32", ndim=2),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d2, i2, u0, 0, "updates"),
+ relax.TensorStructInfo(dtype="float32", ndim=-1),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d0, i3, u0, 0, "updates"),
+ relax.TensorStructInfo((4, 4), dtype="float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d1, i3, u0, 0, "updates"),
+ relax.TensorStructInfo(dtype="float32", ndim=2),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d2, i3, u0, 0, "updates"),
+ relax.TensorStructInfo(dtype="float32", ndim=-1),
+ )
+
+
+def test_scatter_elements_infer_struct_info_symbolic_shape():
+ bb = relax.BlockBuilder()
+ a = tir.Var("a", "int64")
+ b = tir.Var("b", "int64")
+ c = tir.Var("c", "int64")
+ d = tir.Var("d", "int64")
+ e = tir.Var("e", "int64")
+ f = tir.Var("f", "int64")
+
+ d0 = relax.Var("data", R.Tensor((a, b), "float32"))
+ i0 = relax.Var("indices", R.Tensor((c, d), "int64"))
+ u0 = relax.Var("updates", R.Tensor((c, d), "float32"))
+ u1 = relax.Var("updates", R.Tensor((e, f), "float32"))
+
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d0, i0, u0, 0, "updates"),
+ relax.TensorStructInfo((a, b), dtype="float32"),
+ )
+ _check_inference(
+ bb,
+ relax.op.scatter_elements(d0, i0, u1, 0, "updates"),
+ relax.TensorStructInfo((a, b), dtype="float32"),
+ )
+
+
+def test_scatter_elements_infer_struct_info_wrong_indices_type():
+ bb = relax.BlockBuilder()
+ d0 = relax.Var("data", R.Tensor((4, 4), "float32"))
+ i0 = relax.Var("indices", R.Tensor((2, 2), "float32"))
+ u0 = relax.Var("updates", R.Tensor((2, 2), "float32"))
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.scatter_elements(d0, i0, u0))
+
+
+def test_scatter_elements_infer_struct_info_rank_shape_mismatch():
+ a = tir.Var("a", "int64")
+ b = tir.Var("b", "int64")
+
+ bb = relax.BlockBuilder()
+ d0 = relax.Var("data", R.Tensor((4, 4), "float32"))
+ i0 = relax.Var("indices", R.Tensor((3, 3), "int64"))
+ i1 = relax.Var("indices", R.Tensor((3, 3, 3), "int64"))
+ i2 = relax.Var("indices", R.Tensor((a, b), "int64"))
+ u0 = relax.Var("updates", R.Tensor((3, 2), "float32"))
+ u1 = relax.Var("updates", R.Tensor((3, 2, 3), "float32"))
+ u2 = relax.Var("updates", R.Tensor((3, 3, 3), "float32"))
+ u3 = relax.Var("updates", R.Tensor((a + 1, b), "float32"))
+ u4 = relax.Var("updates", R.Tensor((3, 3), "float16"))
+
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.scatter_elements(d0, i0, u0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.scatter_elements(d0, i1, u0))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.scatter_elements(d0, i0, u1))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.scatter_elements(d0, i1, u1))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.scatter_elements(d0, i1, u2))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.scatter_elements(d0, i2, u3))
+ with pytest.raises(TVMError):
+ bb.normalize(relax.op.scatter_elements(d0, i0, u4))
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_manipulate.py
b/tests/python/relax/test_transform_legalize_ops_manipulate.py
index a1f8104cb4..88d45f075a 100644
--- a/tests/python/relax/test_transform_legalize_ops_manipulate.py
+++ b/tests/python/relax/test_transform_legalize_ops_manipulate.py
@@ -1353,5 +1353,203 @@ def test_cumsum_symbolic():
tvm.ir.assert_structural_equal(mod, Expected)
+def test_scatter_elements():
+ # fmt: off
+ @I.ir_module
+ class ScatterElements:
+ @R.function
+ def main(x: R.Tensor((4,4), "float32"), indices: R.Tensor((2,2),
"int64"), updates: R.Tensor((2,2), "float32")):
+ gv = R.scatter_elements(x, indices, updates, axis=1)
+ return gv
+ @I.ir_module
+ class Expected:
+ @T.prim_func
+ def scatter_elements(
+ var_rxplaceholder: T.handle,
+ var_rxplaceholder_1: T.handle,
+ var_rxplaceholder_2: T.handle,
+ out_buf: T.Buffer((T.int64(4), T.int64(4)), "float32"),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ rxplaceholder = T.match_buffer(
+ var_rxplaceholder, (T.int64(4), T.int64(4)), offset_factor=1
+ )
+ rxplaceholder_1 = T.match_buffer(
+ var_rxplaceholder_1, (T.int64(2), T.int64(2)), "int64",
offset_factor=1
+ )
+ rxplaceholder_2 = T.match_buffer(
+ var_rxplaceholder_2, (T.int64(2), T.int64(2)), offset_factor=1
+ )
+ with T.block("scatter_elements_generic"):
+ T.reads(
+ rxplaceholder[T.int64(0) : T.int64(4), T.int64(0) :
T.int64(4)],
+ rxplaceholder_1[T.int64(0) : T.int64(2), T.int64(0) :
T.int64(2)],
+ rxplaceholder_2[T.int64(0) : T.int64(2), T.int64(0) :
T.int64(2)],
+ )
+ T.writes(out_buf[T.int64(0) : T.int64(4), T.int64(0) :
T.int64(4)])
+ for i in T.parallel(T.int64(16)):
+ out_buf[i // T.int64(4), i % T.int64(4)] = rxplaceholder[
+ i // T.int64(4), i % T.int64(4)
+ ]
+ for fused in T.parallel(T.int64(2)):
+ for k in range(T.int64(2)):
+ out_buf[
+ (
+ fused * T.int64(4)
+ + (
+ rxplaceholder_1[
+ (fused * T.int64(2) + k) // T.int64(2),
+ (fused * T.int64(2) + k) % T.int64(2),
+ ]
+ + T.Cast(
+ "int64",
+ rxplaceholder_1[
+ (fused * T.int64(2) + k) //
T.int64(2),
+ (fused * T.int64(2) + k) %
T.int64(2),
+ ]
+ < T.int64(0),
+ )
+ * T.int64(4)
+ )
+ )
+ // T.int64(4),
+ (
+ fused * T.int64(4)
+ + (
+ rxplaceholder_1[
+ (fused * T.int64(2) + k) // T.int64(2),
+ (fused * T.int64(2) + k) % T.int64(2),
+ ]
+ + T.Cast(
+ "int64",
+ rxplaceholder_1[
+ (fused * T.int64(2) + k) //
T.int64(2),
+ (fused * T.int64(2) + k) %
T.int64(2),
+ ]
+ < T.int64(0),
+ )
+ * T.int64(4)
+ )
+ )
+ % T.int64(4),
+ ] = rxplaceholder_2[
+ (fused * T.int64(2) + k) // T.int64(2),
+ (fused * T.int64(2) + k) % T.int64(2),
+ ]
+
+ @R.function
+ def main(
+ x: R.Tensor((4, 4), dtype="float32"),
+ indices: R.Tensor((2, 2), dtype="int64"),
+ updates: R.Tensor((2, 2), dtype="float32"),
+ ) -> R.Tensor((4, 4), dtype="float32"):
+ gv = R.call_tir(
+ Expected.scatter_elements,
+ (x, indices, updates),
+ out_sinfo=R.Tensor((4, 4), dtype="float32"),
+ )
+ return gv
+
+ # fmt: on
+ mod = LegalizeOps()(ScatterElements)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_scatter_elements_symbolic():
+ # fmt: off
+ @I.ir_module
+ class ScatterElements:
+ @R.function
+ def main(x: R.Tensor(("a", "b"), "float32"), indices:R.Tensor(("m",
"n"), "int64"), updates:R.Tensor(("m","n"), "float32")):
+ gv = R.scatter_elements(x, indices, updates, axis=1)
+ return gv
+ @I.ir_module
+ class Expected:
+ @T.prim_func
+ def scatter_elements(
+ var_rxplaceholder: T.handle,
+ var_rxplaceholder_1: T.handle,
+ var_rxplaceholder_2: T.handle,
+ var_scatter_elements_generic: T.handle,
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ a, b = T.int64(), T.int64()
+ rxplaceholder = T.match_buffer(var_rxplaceholder, (a, b),
offset_factor=1)
+ m, n = T.int64(), T.int64()
+ rxplaceholder_1 = T.match_buffer(
+ var_rxplaceholder_1, (m, n), "int64", offset_factor=1
+ )
+ rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, (m, n),
offset_factor=1)
+ out_buf = T.match_buffer(var_scatter_elements_generic, (a, b))
+ with T.block("scatter_elements_generic"):
+ T.reads(
+ rxplaceholder[T.int64(0) : a, T.int64(0) : b],
+ rxplaceholder_1[T.int64(0) : m, T.int64(0) : n],
+ rxplaceholder_2[T.int64(0) : m, T.int64(0) : n],
+ )
+ T.writes(out_buf[T.int64(0) : a, T.int64(0) : b])
+ for i in T.parallel(a * b):
+ out_buf[i // b, i % b] = rxplaceholder[i // b, i % b]
+ for fused in T.parallel(m):
+ for k in range(n):
+ out_buf[
+ (
+ fused * b
+ + (
+ rxplaceholder_1[
+ (fused * n + k) // n, (fused * n + k)
% n
+ ]
+ + T.Cast(
+ "int64",
+ rxplaceholder_1[
+ (fused * n + k) // n, (fused * n +
k) % n
+ ]
+ < T.int64(0),
+ )
+ * b
+ )
+ )
+ // b,
+ (
+ fused * b
+ + (
+ rxplaceholder_1[
+ (fused * n + k) // n, (fused * n + k)
% n
+ ]
+ + T.Cast(
+ "int64",
+ rxplaceholder_1[
+ (fused * n + k) // n, (fused * n +
k) % n
+ ]
+ < T.int64(0),
+ )
+ * b
+ )
+ )
+ % b,
+ ] = rxplaceholder_2[(fused * n + k) // n, (fused * n +
k) % n]
+
+ @R.function
+ def main(
+ x: R.Tensor(("a", "b"), dtype="float32"),
+ indices: R.Tensor(("m", "n"), dtype="int64"),
+ updates: R.Tensor(("m", "n"), dtype="float32"),
+ ) -> R.Tensor(("a", "b"), dtype="float32"):
+ a = T.int64()
+ b = T.int64()
+ m = T.int64()
+ n = T.int64()
+ gv = R.call_tir(
+ Expected.scatter_elements,
+ (x, indices, updates),
+ out_sinfo=R.Tensor((a, b), dtype="float32"),
+ )
+ return gv
+ # fmt: on
+
+ mod = LegalizeOps()(ScatterElements)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
if __name__ == "__main__":
tvm.testing.main()
