You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2023/07/31 13:57:07 UTC
[tvm] branch unity updated: [Unity][Op] Implement basic `call_tir_inplace` operator (#15372)
This is an automated email from the ASF dual-hosted git repository.
tqchen 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 3727c048c6 [Unity][Op] Implement basic `call_tir_inplace` operator (#15372)
3727c048c6 is described below
commit 3727c048c63617a6e56bfb3597f781b76015400e
Author: Steven S. Lyubomirsky <sl...@octoml.ai>
AuthorDate: Mon Jul 31 09:57:00 2023 -0400
[Unity][Op] Implement basic `call_tir_inplace` operator (#15372)
---
include/tvm/relax/attrs/op.h | 14 ++
python/tvm/relax/op/base.py | 72 ++++++++++
python/tvm/script/ir_builder/relax/ir.py | 2 +
src/relax/op/op.cc | 145 ++++++++++++++++++++-
src/relax/transform/call_tir_rewrite.cc | 66 +++++++---
tests/python/relax/test_transform.py | 217 +++++++++++++++++++++++++++++++
tests/python/relax/test_vm_build.py | 101 ++++++++++++++
7 files changed, 599 insertions(+), 18 deletions(-)
diff --git a/include/tvm/relax/attrs/op.h b/include/tvm/relax/attrs/op.h
index d6caa70f52..8c0ee6abc6 100644
--- a/include/tvm/relax/attrs/op.h
+++ b/include/tvm/relax/attrs/op.h
@@ -43,6 +43,20 @@ struct CallTIRWithGradAttrs : public tvm::AttrsNode<CallTIRWithGradAttrs> {
}
}; // struct CallTIRAttrs
+/*! \brief Attributes used in call_tir_inplace */
+struct CallTIRInplaceAttrs : public tvm::AttrsNode<CallTIRInplaceAttrs> {
+ Array<Integer> inplace_indices;
+
+ TVM_DECLARE_ATTRS(CallTIRInplaceAttrs, "relax.attrs.CallTIRInplaceAttrs") {
+ TVM_ATTR_FIELD(inplace_indices)
+ .describe(
+ "Indices that describe which input corresponds to which output. If the `i`th member "
+ "has the value `k` >= 0, then that means that input `k` should be used to store the "
+ "`i`th output. If an element has the value -1, that means a new tensor should be "
+ "allocated for that output.");
+ }
+}; // struct CallTIRInplaceAttrs
+
} // namespace relax
} // namespace tvm
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index 4b71dd70c8..25c70e0493 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -169,6 +169,78 @@ def call_tir_with_grad(
)
+@args_converter.auto
+def call_tir_inplace(
+ gvar: GlobalVar,
+ args: Expr,
+ inplace_indices: Union[int, List[int]],
+ out_sinfo: Union[TensorStructInfo, List[TensorStructInfo]],
+ tir_vars: Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]] = None,
+) -> Call:
+ """
+ Call a TIR PrimFunc and return the result, doing the specified computations in-place
+ (based on the `inplace_indices` argument; outputs will alias the inputs
+ selected by in-place indices).
+
+ Warning: This operator is considered pure by the type system but actually mutates
+ the arguments specified by `inplace_indices`. This operator should not be used directly,
+ but rather should be inserted by passes that have checked whether it is safe to perform
+ operations in-place (i.e., none of the arguments specified as an output is aliased or is
+ live after calling call_tir_inplace).
+
+ Direct calls to this operator should be done for testing purposes only.
+
+ Parameters
+ ----------
+ gvar : GlobalVar
+ The GlobalVar referring to a TIR PrimFunc.
+
+ args : Expr
+ The input arguments.
+
+ input_indices : Union[int, List[int]]
+ Specify which arguments should be used for in-place computations.
+ If `input_indices` is a single integer, it will be made into a singleton list.
+ Suppose `input_indices[i] = j`, where `j >= 0`. Then the `i`th output
+ will be an alias of `args[j]`.
+ If `input_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
+ At least one member of `input_indices` must not be -1.
+
+ out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
+ The structure info of the call_tir_inplace output.
+ It should be a single `TensorStructInfo` or a list of `TensorStructInfo`.
+ Each one denotes the structure info of a returned tensor.
+ If a list of `TensorStructInfo` is given, the result will be a tuple of `TensorStructInfo`.
+
+ tir_vars : Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]]
+ ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
+
+ Returns
+ -------
+ ret: Call
+ A call node for the call_tir operator.
+ """
+ if isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore
+ args = RxTuple((args,))
+
+ if not isinstance(inplace_indices, list):
+ inplace_indices = [inplace_indices]
+
+ if not isinstance(out_sinfo, list):
+ out_sinfo = [out_sinfo]
+
+ if isinstance(tir_vars, (list, tuple)):
+ tir_vars = ShapeExpr(tir_vars)
+
+ return _ffi_api.call_tir_inplace( # type: ignore
+ gvar,
+ args,
+ inplace_indices,
+ out_sinfo,
+ tir_vars,
+ )
+
+
@args_converter.auto
def call_dps_packed(
func: Union[str, Expr],
diff --git a/python/tvm/script/ir_builder/relax/ir.py b/python/tvm/script/ir_builder/relax/ir.py
index 1ce19c2a9e..49d41765fd 100644
--- a/python/tvm/script/ir_builder/relax/ir.py
+++ b/python/tvm/script/ir_builder/relax/ir.py
@@ -54,6 +54,7 @@ from tvm.relax.op import (
call_builtin_with_ctx,
call_pure_packed,
call_tir,
+ call_tir_inplace,
call_tir_with_grad,
call_dps_packed,
ceil,
@@ -587,6 +588,7 @@ __all__ = [
"call_packed",
"call_pure_packed",
"call_tir",
+ "call_tir_inplace",
"call_tir_with_grad",
"call_dps_packed",
"call_builtin_with_ctx",
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index 549bf92278..af93b43dcf 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -126,7 +126,7 @@ TVM_REGISTER_GLOBAL("relax.op.call_pure_packed").set_body_typed(MakeCallPurePack
StructInfo InferStructInfoCallTIR(const Call& call, const BlockBuilder& ctx) {
if (call->sinfo_args.size() != 1) {
ctx->ReportFatal(Diagnostic::Error(call)
- << "sinfo_args should have exact 1 output struct info.");
+ << "sinfo_args should have exactly 1 output struct info.");
}
CHECK(call->args[0]->IsInstance<GlobalVarNode>())
<< "call_tir expects the first argument to be a GlobalVar referring to a TIR PrimFunc. "
@@ -223,6 +223,149 @@ Expr MakeCallTIRWithGrad(Expr func, Tuple args, Array<TensorStructInfo> out_sinf
TVM_REGISTER_GLOBAL("relax.op.call_tir_with_grad").set_body_typed(MakeCallTIRWithGrad);
+// call_tir_inplace
+
+StructInfo InferStructInfoCallTIRInplace(const Call& call, const BlockBuilder& ctx) {
+ if (call->sinfo_args.size() != 1) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "sinfo_args should have exactly 1 output struct info.");
+ }
+ CHECK(call->args[0]->IsInstance<GlobalVarNode>())
+ << "call_tir expects the first argument to be a GlobalVar referring to a TIR PrimFunc. "
+ << "However, gets " << call->args[0];
+ // there must be an inplace index for each output
+ const auto* attrs = call->attrs.as<CallTIRInplaceAttrs>();
+ size_t num_outputs = 1U;
+ if (auto* tup_info = call->sinfo_args[0].as<TupleStructInfoNode>()) {
+ num_outputs = tup_info->fields.size();
+ }
+ if (attrs->inplace_indices.size() != num_outputs) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "There must be an in-place index specified for each output");
+ }
+
+ // check the range for inplace indices, make sure at least one is not -1, ensure they're unique
+ size_t num_args = Downcast<Tuple>(call->args[1])->fields.size();
+ std::unordered_set<int> encountered;
+ for (size_t i = 0; i < attrs->inplace_indices.size(); i++) {
+ int index = attrs->inplace_indices[i].IntValue();
+ if (index < -1 || index >= static_cast<int>(num_args)) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "In-place index " << i << " is out of range (must be between -1 and "
+ << (num_args - 1) << ", inclusive, but is " << index << ")");
+ }
+ if (index != -1) {
+ if (encountered.count(index)) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "All in-place indices must be unique, but index " << index
+ << " appears more than once.");
+ }
+ encountered.insert(index);
+ }
+ }
+ if (encountered.empty()) {
+ ctx->ReportFatal(
+ Diagnostic::Error(call)
+ << "At least one index must have a value other than -1 (or else simply use call_tir)");
+ }
+
+ // for safety, we will make sure the output shape for each in-place argument exactly matches the
+ // input shape
+ // TODO(@slyubomirsky): eventually we will want to handle cases where that is not true
+ Tuple call_args = Downcast<Tuple>(call->args[1]);
+ if (attrs->inplace_indices.size() == 1) {
+ auto* out_sinfo = call->sinfo_args[0].as<TensorStructInfoNode>();
+ if (!out_sinfo) {
+ ctx->ReportFatal(Diagnostic::Error(call) << "The output struct info must be a tensor");
+ }
+ auto* input_sinfo = GetStructInfoAs<TensorStructInfoNode>(
+ call_args->fields[attrs->inplace_indices[0].IntValue()]);
+ if (!input_sinfo || !input_sinfo->shape.defined() ||
+ !CanProveShapeEqual(input_sinfo->shape.value(), out_sinfo->shape.value(),
+ ctx->GetAnalyzer())) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "The shape of output 0 must match input "
+ << attrs->inplace_indices[0].IntValue() << ", whereas we have "
+ << out_sinfo->shape.value() << " in output 0 versus "
+ << input_sinfo->shape.value() << " in input "
+ << attrs->inplace_indices[0].IntValue());
+ }
+ } else {
+ auto out_sinfos = call->sinfo_args[0].as<TupleStructInfoNode>()->fields;
+ for (size_t i = 0; i < attrs->inplace_indices.size(); i++) {
+ if (attrs->inplace_indices[i].IntValue() == -1) {
+ continue;
+ }
+ auto* out_sinfo = out_sinfos[i].as<TensorStructInfoNode>();
+ if (!out_sinfo) {
+ ctx->ReportFatal(Diagnostic::Error(call) << "The output struct info must be a tensor");
+ }
+ auto* input_sinfo = GetStructInfoAs<TensorStructInfoNode>(
+ call_args->fields[attrs->inplace_indices[i].IntValue()]);
+ if (!input_sinfo || !input_sinfo->shape.defined() ||
+ !CanProveShapeEqual(input_sinfo->shape.value(), out_sinfo->shape.value(),
+ ctx->GetAnalyzer())) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "The shape of output " << i << " must match that of input "
+ << attrs->inplace_indices[i].IntValue() << ", whereas we have "
+ << out_sinfo->shape.value() << " in output " << i << " versus "
+ << input_sinfo->shape.value() << " in input "
+ << attrs->inplace_indices[i].IntValue());
+ }
+ }
+ }
+
+ return call->sinfo_args[0];
+}
+
+TVM_REGISTER_NODE_TYPE(CallTIRInplaceAttrs);
+
+RELAY_REGISTER_OP("relax.call_tir_inplace")
+ .set_num_inputs(3)
+ .set_attrs_type<CallTIRInplaceAttrs>()
+ .add_argument("func", "Expr", "The destination-passing-style function.")
+ .add_argument("args", "Tuple", "The input arguments.")
+ .add_argument("packed_ints", "Expr",
+ "ShapeExpr representing a tuple of ints to unpack during runtime. Omitted from "
+ "args if unused")
+ .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoCallTIRInplace)
+ // Warning: considered pure, but it has the potential to create visible effects!
+ // This should only be used if it has been *checked* that it is safe (no aliases, in-place
+ // arguments will no longer be live)
+ .set_attr<Bool>("FPurity", Bool(true));
+
+Expr MakeCallTIRInplace(Expr func, Tuple args, Array<Integer> inplace_indices,
+ Array<TensorStructInfo> out_sinfo_list, Optional<Expr> packed_ints) {
+ for (const TensorStructInfo& sinfo : out_sinfo_list) {
+ const auto* shape = sinfo->shape.as<ShapeExprNode>();
+ CHECK(shape != nullptr) << "out_sinfo of call_tir should have defined ShapeExpr as shape. "
+ "However, one given structure info is "
+ << sinfo;
+ }
+
+ ObjectPtr<CallTIRInplaceAttrs> attrs = make_object<CallTIRInplaceAttrs>();
+ attrs->inplace_indices = Array<Integer>(inplace_indices.begin(), inplace_indices.end());
+
+ StructInfo out_sinfo{nullptr};
+ if (out_sinfo_list.size() == 1) {
+ out_sinfo = out_sinfo_list[0];
+ } else {
+ out_sinfo = TupleStructInfo({out_sinfo_list.begin(), out_sinfo_list.end()});
+ }
+
+ static const Op& op = Op::Get("relax.call_tir_inplace");
+ Call call;
+ if (!packed_ints) {
+ // don't use additional optional argument
+ call = Call(op, {func, args}, Attrs(attrs), {out_sinfo});
+ } else {
+ call = Call(op, {func, args, packed_ints.value()}, Attrs(attrs), {out_sinfo});
+ }
+ return call;
+}
+
+TVM_REGISTER_GLOBAL("relax.op.call_tir_inplace").set_body_typed(MakeCallTIRInplace);
+
// call_dps_packed
StructInfo InferStructInfoCallDPSPacked(const Call& call, const BlockBuilder& ctx) {
diff --git a/src/relax/transform/call_tir_rewrite.cc b/src/relax/transform/call_tir_rewrite.cc
index 6066ed8d2a..e040ccea14 100644
--- a/src/relax/transform/call_tir_rewrite.cc
+++ b/src/relax/transform/call_tir_rewrite.cc
@@ -20,6 +20,7 @@
* \file src/relax/transform/call_tir_rewrite.cc
* \brief Perform explicit tensor allocation for call_tir.
*/
+#include <tvm/relax/attrs/op.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/struct_info.h>
#include <tvm/relax/transform.h>
@@ -33,7 +34,7 @@ namespace relax {
// ==================
// CallTIRMutator
-// Perform explicit tensor allocation for call_tir or call_dps_packed.
+// Perform explicit tensor allocation for call_tir, call_tir_inplace, or call_dps_packed.
// Example:
// lv0: Tensor(n, m) = rx.call_tir(func, (x), (n, m), dtype="float32")
// -->
@@ -49,23 +50,36 @@ class CallTIRMutator : public ExprMutator {
call = expr.as<CallNode>();
static const Op& call_tir_op = Op::Get("relax.call_tir");
+ static const Op& call_tir_inplace_op = Op::Get("relax.call_tir_inplace");
static const Op& call_dps_packed_op = Op::Get("relax.call_dps_packed");
static const Op& alloc_tensor_op = Op::Get("relax.builtin.alloc_tensor");
static const Op& call_tir_dyn_op = Op::Get("relax.vm.call_tir_dyn");
- if (call->op == call_tir_op || call->op == call_dps_packed_op) {
+ if (call->op == call_tir_op || call->op == call_tir_inplace_op ||
+ call->op == call_dps_packed_op) {
+ bool is_inplace = (call->op == call_tir_inplace_op);
+ const auto* inplace_attrs = call->attrs.as<CallTIRInplaceAttrs>();
Array<Expr> outs;
if (const auto& _tensor_sinfo = MatchStructInfo<TensorStructInfo>(expr)) {
// single output case
const TensorStructInfo& tensor_sinfo = _tensor_sinfo.value();
ICHECK(tensor_sinfo->shape.defined())
<< "the TensorStructInfo shape of call_tir has not populated";
- outs.push_back(
- builder_->Emit(Call(alloc_tensor_op, //
- {Downcast<ShapeExpr>(tensor_sinfo->shape.value()),
- DataTypeImm(tensor_sinfo->dtype), PrimValue::Int64(0)}, //
- Attrs()),
- "alloc"));
+ if (!is_inplace) {
+ outs.push_back(
+ builder_->Emit(Call(alloc_tensor_op, //
+ {Downcast<ShapeExpr>(tensor_sinfo->shape.value()),
+ DataTypeImm(tensor_sinfo->dtype), PrimValue::Int64(0)}, //
+ Attrs()),
+ "alloc"));
+ } else {
+ // if there is only one output, it must be an in-place argument, but check anyway
+ ICHECK(inplace_attrs->inplace_indices[0].IntValue() != -1)
+ << "If calling call_tir_inplace and there is one output, its in-place index must not"
+ " be -1.";
+ outs.push_back(
+ Downcast<Tuple>(call->args[1])->fields[inplace_attrs->inplace_indices[0].IntValue()]);
+ }
} else if (const auto& _tuple_sinfo = MatchStructInfo<TupleStructInfo>(expr)) {
// multiple output case
const TupleStructInfo& tuple_sinfo = _tuple_sinfo.value();
@@ -79,12 +93,17 @@ class CallTIRMutator : public ExprMutator {
ICHECK(field_tensor->shape.defined())
<< "call_tir expects all TensorStructInfo has shape, but got " << field_tensor
<< " as an element of TupleStructInfo";
- outs.push_back(
- builder_->Emit(Call(alloc_tensor_op,
- {Downcast<ShapeExpr>(field_tensor->shape.value()),
- DataTypeImm(field_tensor->dtype), PrimValue::Int64(0)},
- Attrs()),
- "alloc"));
+ if (!is_inplace || inplace_attrs->inplace_indices[i].IntValue() == -1) {
+ outs.push_back(
+ builder_->Emit(Call(alloc_tensor_op,
+ {Downcast<ShapeExpr>(field_tensor->shape.value()),
+ DataTypeImm(field_tensor->dtype), PrimValue::Int64(0)},
+ Attrs()),
+ "alloc"));
+ } else {
+ outs.push_back(Downcast<Tuple>(call->args[1])
+ ->fields[inplace_attrs->inplace_indices[i].IntValue()]);
+ }
}
} else {
LOG(FATAL) << "TypeError: The struct info of call_tir expects to be TensorStructInfo or "
@@ -95,7 +114,16 @@ class CallTIRMutator : public ExprMutator {
Array<Expr> args;
if (call->args[1].as<TupleNode>()) {
args = Downcast<Tuple>(call->args[1])->fields;
- args.insert(args.end(), outs.begin(), outs.end());
+ // for call_tir_inplace, don't reinsert in-place args, only the newly allocated ones
+ if (!is_inplace) {
+ args.insert(args.end(), outs.begin(), outs.end());
+ } else {
+ for (size_t i = 0; i < outs.size(); i++) {
+ if (inplace_attrs->inplace_indices[i].IntValue() == -1) {
+ args.push_back(outs[i]);
+ }
+ }
+ }
if (call->args.size() == 2) {
builder_->Emit(Call(call->args[0], args), "_");
@@ -105,8 +133,12 @@ class CallTIRMutator : public ExprMutator {
builder_->Emit(Call(call_tir_dyn_op, {call->args[0], Tuple(args)}), "_");
}
} else {
- args = outs;
- args.insert(args.begin(), call->args[1]);
+ if (!is_inplace) {
+ args = outs;
+ args.insert(args.begin(), call->args[1]);
+ } else {
+ args.push_back(call->args[1]);
+ }
builder_->Emit(Call(call->args[0], args), "_");
}
diff --git a/tests/python/relax/test_transform.py b/tests/python/relax/test_transform.py
index 102f80b2b0..d690cc55e3 100644
--- a/tests/python/relax/test_transform.py
+++ b/tests/python/relax/test_transform.py
@@ -268,6 +268,223 @@ def test_call_dps_packed_rewrite():
assert s2.op.global_symbol == "test.op.identity"
+def test_call_tir_inplace_simple():
+ # simple case: one inplace argument
+ @tvm.script.ir_module
+ class Input:
+ @T.prim_func
+ def zeros(A: T.Buffer((2, 3), "int32")):
+ # just overwrites A with 0s
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_zeros"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.writes(A[ax0, ax1])
+ A[ax0, ax1] = T.int32(0)
+
+ @R.function
+ def foo(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "int32"):
+ # we expect RemovePurityChecking to have been used before this point
+ R.func_attr({"relax.force_pure": True})
+ gv0 = R.call_tir_inplace(Input.zeros, x, 0, R.Tensor((2, 3), dtype="int32"))
+ return gv0
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.prim_func
+ def zeros(A: T.Buffer((2, 3), "int32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_zeros"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.writes(A[ax0, ax1])
+ A[ax0, ax1] = T.int32(0)
+
+ @R.function
+ def foo(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "int32"):
+ R.func_attr({"relax.force_pure": True})
+ _ = Expected.zeros(x)
+ gv0 = x
+ return gv0
+
+ new_mod = relax.transform.CallTIRRewrite()(Input)
+ tvm.ir.assert_structural_equal(Expected["foo"], new_mod["foo"], map_free_vars=True)
+
+
+def test_call_tir_inplace_multiple_args():
+ @tvm.script.ir_module
+ class Input:
+ @T.prim_func
+ def copy(
+ A: T.Buffer((2, 3), "int32"), B: T.Buffer((2, 3), "int32"), C: T.Buffer((2, 3), "int32")
+ ):
+ # copies the contents of C into A and B
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_zeros"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(C[ax0, ax1])
+ T.writes(A[ax0, ax1], B[ax0, ax1])
+ A[ax0, ax1] = C[ax0, ax1]
+ B[ax0, ax1] = C[ax0, ax1]
+
+ @R.function
+ def foo(
+ x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
+ ) -> R.Tuple(R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")):
+ R.func_attr({"relax.force_pure": True})
+ gv0 = R.call_tir_inplace(
+ Input.copy,
+ (x, y, z),
+ [0, 1],
+ [R.Tensor((2, 3), dtype="int32"), R.Tensor((2, 3), dtype="int32")],
+ )
+ return gv0
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.prim_func
+ def copy(A: T.Buffer((2, 3), "int32"), B: T.Buffer((2, 3), "int32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_zeros"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(B[ax0, ax1])
+ T.writes(A[ax0, ax1])
+ A[ax0, ax1] = B[ax0, ax1]
+
+ @R.function
+ def foo(
+ x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
+ ) -> R.Tuple(R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")):
+ R.func_attr({"relax.force_pure": True})
+ _ = Expected.copy(x, y, z)
+ gv0 = (x, y)
+ return gv0
+
+ new_mod = relax.transform.CallTIRRewrite()(Input)
+ tvm.ir.assert_structural_equal(Expected["foo"], new_mod["foo"], map_free_vars=True)
+
+
+def test_call_tir_inplace_some_new():
+ @tvm.script.ir_module
+ class Input:
+ @T.prim_func
+ def copy(
+ A: T.Buffer((2, 3), "int32"),
+ B: T.Buffer((2, 3), "int32"),
+ C: T.Buffer((2, 3), "int32"),
+ out1: T.Buffer((2, 3), "int32"),
+ out2: T.Buffer((2, 3), "int32"),
+ ):
+ # copies the contents of C into A, out1, and out2
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_zeros"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(C[ax0, ax1])
+ T.writes(A[ax0, ax1], out1[ax0, ax1], out2[ax0, ax1])
+ A[ax0, ax1] = C[ax0, ax1]
+ out1[ax0, ax1] = C[ax0, ax1]
+ out2[ax0, ax1] = C[ax0, ax1]
+
+ @R.function
+ def foo(
+ x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
+ ) -> R.Tuple(
+ R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), dtype="int32")
+ ):
+ R.func_attr({"relax.force_pure": True})
+ gv0 = R.call_tir_inplace(
+ Input.copy,
+ (x, y, z),
+ [0, -1, -1],
+ [
+ R.Tensor((2, 3), dtype="int32"),
+ R.Tensor((2, 3), dtype="int32"),
+ R.Tensor((2, 3), dtype="int32"),
+ ],
+ )
+ return gv0
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.prim_func
+ def copy(
+ A: T.Buffer((2, 3), "int32"),
+ B: T.Buffer((2, 3), "int32"),
+ C: T.Buffer((2, 3), "int32"),
+ out1: T.Buffer((2, 3), "int32"),
+ out2: T.Buffer((2, 3), "int32"),
+ ):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_zeros"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(C[ax0, ax1])
+ T.writes(A[ax0, ax1], out1[ax0, ax1], out2[ax0, ax1])
+ A[ax0, ax1] = C[ax0, ax1]
+ out1[ax0, ax1] = C[ax0, ax1]
+ out2[ax0, ax1] = C[ax0, ax1]
+
+ @R.function
+ def foo(
+ x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
+ ) -> R.Tuple(
+ R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), dtype="int32")
+ ):
+ R.func_attr({"relax.force_pure": True})
+ gv0 = R.builtin.alloc_tensor(R.shape([2, 3]), "int32", R.prim_value(0))
+ gv1 = R.builtin.alloc_tensor(R.shape([2, 3]), "int32", R.prim_value(0))
+ _ = Expected.copy(x, y, z, gv0, gv1)
+ gv2 = (x, gv0, gv1)
+ return gv2
+
+ new_mod = relax.transform.CallTIRRewrite()(Input)
+ tvm.ir.assert_structural_equal(Expected["foo"], new_mod["foo"], map_free_vars=True)
+
+
+@pytest.mark.xfail()
+def test_call_tir_inplace_repeated_input():
+ @tvm.script.ir_module
+ class Input:
+ @T.prim_func
+ def func(
+ A: T.Buffer((2, 3), "int32"), B: T.Buffer((2, 3), "int32"), C: T.Buffer((2, 3), "int32")
+ ):
+ T.evaluate(0)
+
+ @R.function
+ def foo(
+ x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
+ ) -> R.Tuple(R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")):
+ R.func_attr({"relax.force_pure": True})
+ gv0 = R.call_tir_inplace(
+ Input.func,
+ (x, y, z),
+ # repeated 0 -> that's an error
+ [0, 0],
+ [R.Tensor((2, 3), dtype="int32"), R.Tensor((2, 3), dtype="int32")],
+ )
+ return gv0
+
+
+@pytest.mark.xfail()
+def test_call_tir_inplace_all_new():
+ @tvm.script.ir_module
+ class Input:
+ @T.prim_func
+ def func(A: T.Buffer((2, 3), "int32")):
+ T.evaluate(0)
+
+ @R.function
+ def foo(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "int32"):
+ R.func_attr({"relax.force_pure": True})
+ # cannot make the only output a fresh one
+ gv0 = R.call_tir_inplace(Input.func, x, -1, R.Tensor((2, 3), dtype="int32"))
+ return gv0
+
+
def test_vm_builtin_lower():
@tvm.script.ir_module
class TestVMBuiltinLower:
diff --git a/tests/python/relax/test_vm_build.py b/tests/python/relax/test_vm_build.py
index dfaf89e104..0b34d24540 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -204,6 +204,107 @@ def test_vm_compile_e2e_func_param_with_shape(exec_mode):
tvm.testing.assert_allclose(res.numpy(), expected, rtol=1e-6, atol=1e-6)
+@pytest.mark.parametrize("exec_mode", EXEC_MODE)
+def test_call_tir_inplace_e2e_simple(exec_mode):
+ @tvm.script.ir_module
+ class TestCallTIRInplaceE2ESimple:
+ @T.prim_func
+ def copy(
+ A: T.Buffer((2, 3), "int32"),
+ B: T.Buffer((2, 3), "int32"),
+ C: T.Buffer((2, 3), "int32"),
+ out1: T.Buffer((2, 3), "int32"),
+ ):
+ # copies the contents of C into A, B, and out1
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_zeros"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(C[ax0, ax1])
+ T.writes(A[ax0, ax1], B[ax0, ax1], out1[ax0, ax1])
+ A[ax0, ax1] = C[ax0, ax1]
+ B[ax0, ax1] = C[ax0, ax1]
+ out1[ax0, ax1] = C[ax0, ax1]
+
+ @R.function
+ def main(
+ x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
+ ) -> R.Tuple(
+ R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")
+ ):
+ res = R.call_tir_inplace(
+ TestCallTIRInplaceE2ESimple.copy,
+ (x, y, z),
+ [0, 1, -1],
+ [R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")],
+ )
+ return res
+
+ mod = TestCallTIRInplaceE2ESimple
+
+ target = tvm.target.Target("llvm", host="llvm")
+ ex = relax.build(mod, target, exec_mode=exec_mode)
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+
+ x = tvm.nd.array(np.zeros((2, 3)).astype(np.int32))
+ y = tvm.nd.array(np.zeros((2, 3)).astype(np.int32))
+ z = tvm.nd.array(np.ones((2, 3)).astype(np.int32))
+ vm.set_input("main", x, y, z)
+ vm.invoke_stateful("main")
+ outs = vm.get_outputs("main")
+ # check the expected aliasing (the last result is newly allocated)
+ assert x == outs[0]
+ assert y == outs[1]
+ assert x != y
+ assert x != outs[2]
+ assert y != outs[2]
+ tvm.testing.assert_allclose(x.numpy(), z.numpy(), rtol=1e-7, atol=1e-7)
+ tvm.testing.assert_allclose(y.numpy(), z.numpy(), rtol=1e-7, atol=1e-7)
+ tvm.testing.assert_allclose(outs[2].numpy(), z.numpy(), rtol=1e-7, atol=1e-7)
+
+
+@pytest.mark.parametrize("exec_mode", EXEC_MODE)
+def test_call_tir_inplace_e2e_rw(exec_mode):
+ # read and write from the same tensor
+ @tvm.script.ir_module
+ class TestCallTIRInplaceE2ERW:
+ @T.prim_func
+ def inplace_add(A: T.Buffer((2, 3), "int32"), B: T.Buffer((2, 3), "int32")):
+ # sums A and B, storing the result in A
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_add"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(A[ax0, ax1], B[ax0, ax1])
+ T.writes(A[ax0, ax1])
+ A[ax0, ax1] = A[ax0, ax1] + B[ax0, ax1]
+
+ @R.function
+ def main(
+ x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32")
+ ) -> R.Tensor((2, 3), "int32"):
+ res = R.call_tir_inplace(
+ TestCallTIRInplaceE2ERW.inplace_add, (x, y), [0], R.Tensor((2, 3), "int32")
+ )
+ return res
+
+ mod = TestCallTIRInplaceE2ERW
+
+ target = tvm.target.Target("llvm", host="llvm")
+ ex = relax.build(mod, target, exec_mode=exec_mode)
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+
+ x = tvm.nd.array(np.ones((2, 3)).astype(np.int32))
+ y = tvm.nd.array(np.ones((2, 3)).astype(np.int32))
+ vm.set_input("main", x, y)
+ vm.invoke_stateful("main")
+ out = vm.get_outputs("main")
+ expected = tvm.nd.array(np.full((2, 3), 2).astype(np.int32))
+
+ assert x == out
+ tvm.testing.assert_allclose(out.numpy(), expected.numpy(), rtol=1e-7, atol=1e-7)
+
+
@pytest.mark.parametrize("exec_mode", EXEC_MODE)
def test_vm_emit_te_extern(exec_mode):
if not tvm.get_global_func("tvm.contrib.cblas.matmul", True):