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/09 14:00:03 UTC
[tvm] branch unity updated: [Unity][Training] Registering te gradient (#15231)
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 399c5eaafe [Unity][Training] Registering te gradient (#15231)
399c5eaafe is described below
commit 399c5eaafe43f37ca630f9db71ff4ba2d0a5742b
Author: Yixin Dong <ub...@gmail.com>
AuthorDate: Sun Jul 9 21:59:57 2023 +0800
[Unity][Training] Registering te gradient (#15231)
---
include/tvm/relax/attrs/op.h | 49 ++++
python/tvm/relax/__init__.py | 2 +-
python/tvm/relax/block_builder.py | 59 ++++-
python/tvm/relax/op/base.py | 62 ++++-
python/tvm/relax/op/op_attrs.py | 5 +
python/tvm/relax/training/utils.py | 52 +++-
python/tvm/script/ir_builder/relax/ir.py | 2 +
src/relax/op/op.cc | 51 ++++
src/relax/transform/gradient.cc | 40 +--
src/script/printer/relax/call.cc | 25 +-
tests/python/relax/test_op_gradient_numeric.py | 20 +-
tests/python/relax/test_op_misc.py | 19 ++
.../relax/test_training_register_te_gradient.py | 280 +++++++++++++++++++++
tests/python/relax/test_tvmscript_parser.py | 28 +++
tests/python/relax/test_tvmscript_printer_relax.py | 21 ++
15 files changed, 691 insertions(+), 24 deletions(-)
diff --git a/include/tvm/relax/attrs/op.h b/include/tvm/relax/attrs/op.h
new file mode 100644
index 0000000000..d6caa70f52
--- /dev/null
+++ b/include/tvm/relax/attrs/op.h
@@ -0,0 +1,49 @@
+/*
+ * 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/op.h
+ * \brief Attributes for relax specific operators.
+ */
+#ifndef TVM_RELAX_ATTRS_OP_H_
+#define TVM_RELAX_ATTRS_OP_H_
+
+#include <tvm/relax/expr.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in call_tir_with_grad */
+struct CallTIRWithGradAttrs : public tvm::AttrsNode<CallTIRWithGradAttrs> {
+ String te_grad_name;
+ Map<String, ObjectRef> te_grad_kwargs;
+
+ TVM_DECLARE_ATTRS(CallTIRWithGradAttrs, "relax.attrs.CallTIRWithGradAttrs") {
+ TVM_ATTR_FIELD(te_grad_name)
+ .describe(
+ "The name of the te gradient function associated with this call_tir_with_grad node.");
+ TVM_ATTR_FIELD(te_grad_kwargs)
+ .describe("The keyword arguments passed to the te gradient function.");
+ }
+}; // struct CallTIRAttrs
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_ATTRS_OP_H_
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index e4d91c59ac..2945aee1a9 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -56,7 +56,7 @@ from .ty import Type, ObjectType, ShapeType, DynTensorType, TupleType, FuncType,
from .exec_builder import ExecBuilder
# Operator
-from .op.base import call_tir, call_pure_packed, call_dps_packed
+from .op.base import call_tir, call_pure_packed, call_dps_packed, call_tir_with_grad
# BlockBuilder
from .block_builder import BlockBuilder
diff --git a/python/tvm/relax/block_builder.py b/python/tvm/relax/block_builder.py
index 502073edf2..894d9324dd 100644
--- a/python/tvm/relax/block_builder.py
+++ b/python/tvm/relax/block_builder.py
@@ -32,7 +32,7 @@ from .expr import (
Binding,
)
from .struct_info import StructInfo
-from .op.base import call_tir
+from .op.base import call_tir, call_tir_with_grad
from . import _ffi_api
from .utils import gen_call_tir_inputs
@@ -331,6 +331,63 @@ class BlockBuilder(Object):
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+ def call_te_with_grad(
+ self,
+ func: Callable,
+ *args: Any,
+ te_grad_name: str,
+ te_grad_kwargs: Dict[str, Object] = None,
+ **kwargs: Any,
+ ) -> Expr:
+ """Generate a call node according to the te function.
+ This method will generate a call_tir_with_grad node, i.e. a call_tir node bound with a
+ te gradient function (refered by te_grad_name).
+
+ Parameters
+ ----------
+ func : Callable
+ A function that returns a te tensor or a list of te tensors.
+
+ args : Any, optional
+ arguments passed to the function.
+
+ te_grad_name : str
+ The registered name of the te gradient function associated with the call_tir_with_grad
+ node. Must be provided as a keyword argument.
+
+ te_grad_kwargs : Dict[str, Object], optional
+ The keyword arguments passed to the te gradient function.
+ Optionally provided as a keyword argument. Default: {}.
+
+ kwargs : Any, optional
+ The keyword arguments passed to the function.
+ Note that the following keyword args are reserved:
+
+ - 'primfunc_name_hint' for passing name hint to the PrimFunc
+ that gets generated.
+ - 'primfunc_attrs' is reserved for passing func attributes to
+ be added to the PrimFunc that gets created.
+
+ Returns
+ -------
+ ret : tvm.relax.Call
+ A newly created call node
+ """
+
+ primfunc_name = kwargs.pop("primfunc_name_hint", None)
+ tir_func, call_args, output_sinfo, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
+
+ if te_grad_kwargs is None:
+ te_grad_kwargs = {}
+
+ if not primfunc_name:
+ primfunc_name = func.__name__
+ gvar = self.add_func(tir_func, primfunc_name)
+
+ return call_tir_with_grad(
+ gvar, call_args, output_sinfo, te_grad_name, te_grad_kwargs, tir_vars
+ )
+
def emit_te(self, func: Callable, *args: Any, **kwargs: Any) -> Var:
"""Emit a call node according to the te function.
This function converts arguments from relax expression to te tensor,
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index a1a183003f..4b71dd70c8 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# pylint: disable=redefined-builtin
"""The base Relax operators."""
-from typing import Union, List, Tuple, Optional, Callable
+from typing import Dict, Union, List, Tuple, Optional, Callable
import tvm
@@ -109,6 +109,66 @@ def call_tir(
return _ffi_api.call_tir(gvar, args, out_sinfo, tir_vars) # type: ignore
+@args_converter.auto
+def call_tir_with_grad(
+ gvar: GlobalVar,
+ args: Expr,
+ out_sinfo: Union[TensorStructInfo, List[TensorStructInfo]],
+ te_grad_name: str,
+ te_grad_kwargs: Dict[str, Object] = None,
+ tir_vars: Optional[Union[ShapeExpr, Tuple[PrimExpr], List[PrimExpr]]] = None,
+) -> Call:
+ """
+ Call a tir.prim_func and return the output. This intrinsic will bind a te gradient function
+ (refered by te_grad_name) to the call_tir_with_grad node. The te gradient function will be
+ called by the Gradient pass.
+
+ Parameters
+ ----------
+ gvar : GlobalVar
+ The GlobalVar referring to a tir PrimFunc.
+
+ args : Expr
+ The input arguments.
+
+ out_sinfo : Union[TensorStructInfo, List[TensorStructInfo]]
+ The structure info of the call_tir_with_grad output.
+ It should be a single or a list of TensorStructInfo. Each one denotes the
+ structure info of a returned tensor.
+
+ te_grad_name : str
+ The registered name of the te gradient function associated with the call_tir_with_grad
+ node. Must be provided as a keyword argument.
+
+ te_grad_kwargs : Dict[str, Object], optional
+ The keyword arguments passed to the te gradient function.
+ Optionally provided as a keyword argument. Default: {}.
+
+ 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_with_grad operator.
+ """
+ if isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore
+ args = RxTuple((args,))
+
+ if not isinstance(out_sinfo, list):
+ out_sinfo = [out_sinfo]
+
+ if isinstance(tir_vars, (list, tuple)):
+ tir_vars = ShapeExpr(tir_vars)
+
+ if te_grad_kwargs is None:
+ te_grad_kwargs = {}
+
+ return _ffi_api.call_tir_with_grad( # type: ignore
+ gvar, args, out_sinfo, te_grad_name, te_grad_kwargs, tir_vars
+ )
+
+
@args_converter.auto
def call_dps_packed(
func: Union[str, Expr],
diff --git a/python/tvm/relax/op/op_attrs.py b/python/tvm/relax/op/op_attrs.py
index b08e853ef3..848c2094a0 100644
--- a/python/tvm/relax/op/op_attrs.py
+++ b/python/tvm/relax/op/op_attrs.py
@@ -19,6 +19,11 @@ from tvm.ir import Attrs
import tvm._ffi
+@tvm._ffi.register_object("relax.attrs.CallTIRWithGradAttrs")
+class CallTIRWithGradAttrs(Attrs):
+ """Attributes used in call_tir_with_grad operator"""
+
+
@tvm._ffi.register_object("relax.attrs.InitAttrs")
class InitAttrs(Attrs):
"""Attributes used in full/full_like, ones/ones_like, and zeros/zeros_like operator"""
diff --git a/python/tvm/relax/training/utils.py b/python/tvm/relax/training/utils.py
index 2944a203a9..69fa85043e 100644
--- a/python/tvm/relax/training/utils.py
+++ b/python/tvm/relax/training/utils.py
@@ -17,9 +17,13 @@
# pylint: disable=invalid-name
"""Utility functions for relax training."""
-from typing import Optional
+from typing import Optional, Callable
import tvm
+from tvm import relax
+from tvm._ffi.registry import register_func
+from tvm.relax.block_builder import BlockBuilder
+
from ..expr import Function
from . import _ffi_api
@@ -153,3 +157,49 @@ def AppendLoss(
num_backbone_outputs,
new_func_name,
)
+
+
+def register_te_gradient(te_grad_name: str, te_grad_func: Callable = None):
+ """Register a te gradient function bind with name te_grad_name. te_grad_name can be referenced
+ later in call_tir_with_grad nodes.
+
+ Parameters
+ ----------
+ te_grad_name : str
+ The registered name of the te gradient function. Should be align with the te_grad_name in
+ call_tir_with_grad nodes.
+
+ grad_func : Callable
+ The te grad function.
+ It must be a function taking (output_grad: Tensor, arg1: Tensor, arg2: Tensor, ...)
+ as inputs and returning a list of Tensor created by te.compute.
+
+ Returns
+ -------
+ mod : IRModule
+ The mod with corresponding attributes attached.
+ """
+
+ def register(func: Callable):
+ func_prefix = "tvm.relax.te_grad._register."
+
+ # The handler function is used to let the backend (cpp side) to emit_te.
+ # It's a wrapper of the te_grad_func.
+ # It takes the blockbuilder, the gradient var of the output and the forward call expr.
+ # It will return the emitted var.
+
+ def handler(
+ builder: BlockBuilder, output_grad_var: relax.Var, call_tir: relax.Call
+ ) -> relax.Expr:
+ return builder.emit_te(
+ func,
+ output_grad_var,
+ *call_tir.args[1],
+ **call_tir.attrs.te_grad_kwargs,
+ primfunc_name_hint=te_grad_name,
+ )
+
+ register_func(func_prefix + te_grad_name, handler)
+ return func
+
+ return register(te_grad_func) if te_grad_func else register
diff --git a/python/tvm/script/ir_builder/relax/ir.py b/python/tvm/script/ir_builder/relax/ir.py
index b06d9547ac..1ce19c2a9e 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_with_grad,
call_dps_packed,
ceil,
clip,
@@ -586,6 +587,7 @@ __all__ = [
"call_packed",
"call_pure_packed",
"call_tir",
+ "call_tir_with_grad",
"call_dps_packed",
"call_builtin_with_ctx",
"ceil",
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index 3e3cbd6b7d..549bf92278 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/relax/analysis.h>
+#include <tvm/relax/attrs/op.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/utils.h>
#include <tvm/relay/op.h>
@@ -172,6 +173,56 @@ Expr MakeCallTIR(Expr func, Tuple args, Array<TensorStructInfo> out_sinfo_list,
TVM_REGISTER_GLOBAL("relax.op.call_tir").set_body_typed(MakeCallTIR);
+// call_tir_with_grad
+
+TVM_REGISTER_NODE_TYPE(CallTIRWithGradAttrs);
+
+RELAY_REGISTER_OP("relax.call_tir_with_grad")
+ .set_num_inputs(3)
+ .set_attrs_type<CallTIRWithGradAttrs>()
+ .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", InferStructInfoCallTIR)
+ .set_attr<Bool>("FPurity", Bool(true));
+
+Expr MakeCallTIRWithGrad(Expr func, Tuple args, Array<TensorStructInfo> out_sinfo_list,
+ String te_grad_name, Map<String, ObjectRef> te_grad_kwargs,
+ 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_with_grad should have defined ShapeExpr as shape. "
+ "However, one given structure info is "
+ << sinfo;
+ }
+
+ 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()});
+ }
+
+ ObjectPtr<CallTIRWithGradAttrs> attrs = make_object<CallTIRWithGradAttrs>();
+ attrs->te_grad_name = te_grad_name;
+ attrs->te_grad_kwargs = te_grad_kwargs;
+
+ static const Op& op = Op::Get("relax.call_tir_with_grad");
+ 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_with_grad").set_body_typed(MakeCallTIRWithGrad);
+
// call_dps_packed
StructInfo InferStructInfoCallDPSPacked(const Call& call, const BlockBuilder& ctx) {
diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc
index 2cda7a972d..9a271a3ebe 100644
--- a/src/relax/transform/gradient.cc
+++ b/src/relax/transform/gradient.cc
@@ -25,6 +25,7 @@
* with respect to the only return value of the function, which needs to be scalar.
*/
+#include <tvm/relax/attrs/op.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/nested_msg.h>
#include <tvm/relax/op_attr_types.h>
@@ -103,27 +104,35 @@ class BackwardBindingGenerator : private ExprVisitor {
ExprVisitor::VisitBinding_(var_binding);
}
- // Handle the adjoint expr of the inputs of binding
+ // The following functions will handle the adjoint expr of the inputs of binding
// For call node, we would call the registered gradient functions
void VisitBinding_(const VarBindingNode* binding, const CallNode* call) final {
+ // Skip if it is not an Op
+ if (!call->op->IsInstance<OpNode>()) {
+ return;
+ }
+
static const OpAttrMap<FPrimalGradient>& gradient_op_map =
Op::GetAttrMap<FPrimalGradient>("FPrimalGradient");
Var adjoint_var = adjoint_var_map_[binding->var];
const Op& call_op = Downcast<Op>(call->op);
- const Array<Expr>& partials =
- gradient_op_map[call_op](binding->var, GetRef<Call>(call), adjoint_var, builder_);
- ICHECK(partials.size() == call->args.size()) << "partials number != inputs number";
-
- for (size_t i = 0; i < partials.size(); ++i) {
- Expr partial = partials[i];
- if (IsCallNoGrad(partial)) { // no grad: don't update
- continue;
- }
- if (!partial->struct_info_.defined()) {
- UpdateStructInfo(partial, GetStructInfo(call->args[i]));
+
+ if (call_op == Op::Get("relax.call_tir")) {
+ LOG(FATAL) << "Differentiation of call_tir op is not supported yet.";
+ } else {
+ Array<Expr> partials =
+ gradient_op_map[call_op](binding->var, GetRef<Call>(call), adjoint_var, builder_);
+ ICHECK(partials.size() == call->args.size()) << "partials number != inputs number";
+ for (size_t i = 0; i < partials.size(); ++i) {
+ if (IsCallNoGrad(partials[i])) { // no grad: don't update
+ continue;
+ }
+ if (!partials[i]->struct_info_.defined()) {
+ UpdateStructInfo(partials[i], GetStructInfo(call->args[i]));
+ }
+ UpdateAdjoint(call->args[i], partials[i]);
}
- UpdateAdjoint(call->args[i], partial);
}
}
@@ -339,9 +348,8 @@ class GradientMutator : private ExprMutator {
Function new_func_transformed = WithAttr(Downcast<Function>(mutator.VisitExpr(new_func)),
tvm::attr::kGlobalSymbol, new_name);
- IRModule new_module = GetRef<IRModule>(mod.CopyOnWrite());
- new_module->Add(GlobalVar(new_name), new_func_transformed);
- return new_module;
+ mutator.builder_->AddFunction(new_func_transformed, new_name);
+ return mutator.builder_->GetContextIRModule();
}
private:
diff --git a/src/script/printer/relax/call.cc b/src/script/printer/relax/call.cc
index 9bf9e50ee8..520058a8a5 100644
--- a/src/script/printer/relax/call.cc
+++ b/src/script/printer/relax/call.cc
@@ -16,6 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/relax/attrs/op.h>
+
#include "./utils.h"
namespace tvm {
@@ -95,7 +97,9 @@ Optional<ExprDoc> PrintCallTIRDPSPacked(const relax::Call& n, const ObjectPath&
const IRDocsifier& d) {
static const Op& call_tir_op = Op::Get("relax.call_tir");
static const Op& call_dps_packed_op = Op::Get("relax.call_dps_packed");
- if (!n->op.same_as(call_tir_op) && !n->op.same_as(call_dps_packed_op)) {
+ static const Op& call_tir_with_grad_op = Op::Get("relax.call_tir_with_grad");
+ if (!n->op.same_as(call_tir_op) && !n->op.same_as(call_dps_packed_op) &&
+ !n->op.same_as(call_tir_with_grad_op)) {
return NullOpt;
}
ICHECK(n->args.size() == 2 || n->args.size() == 3);
@@ -121,6 +125,23 @@ Optional<ExprDoc> PrintCallTIRDPSPacked(const relax::Call& n, const ObjectPath&
} else {
kwargs_values.push_back(d->AsDoc<ExprDoc>(o_sinfo, o_sinfo_p));
}
+
+ // start of specially handling call_tir_with_grad
+ if (const auto* call_tir_with_grad_attrs = n->attrs.as<relax::CallTIRWithGradAttrs>()) {
+ kwargs_keys.push_back("te_grad_name");
+ kwargs_values.push_back(LiteralDoc::Str(call_tir_with_grad_attrs->te_grad_name,
+ n_p->Attr("attrs")->Attr("te_grad_name")));
+ if (!call_tir_with_grad_attrs->te_grad_kwargs.empty()) {
+ kwargs_keys.push_back("te_grad_kwargs");
+ kwargs_values.push_back(d->AsDoc<ExprDoc>(call_tir_with_grad_attrs->te_grad_kwargs,
+ n_p->Attr("attrs")->Attr("te_grad_kwargs")));
+ }
+ }
+ if (n->op.same_as(call_tir_with_grad_op)) {
+ return Relax(d, "call_tir_with_grad")->Call(args, kwargs_keys, kwargs_values);
+ }
+ // end of specially handling call_tir_with_grad
+
if (n->op.same_as(call_dps_packed_op)) {
return Relax(d, "call_dps_packed")->Call(args, kwargs_keys, kwargs_values);
}
@@ -173,7 +194,7 @@ Optional<ExprDoc> PrintRelaxPrint(const relax::Call& n, const ObjectPath& n_p,
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<relax::Call>( //
"", [](relax::Call n, ObjectPath n_p, IRDocsifier d) -> Doc {
- // Special case: call_tir, call_dps_packed
+ // Special case: call_tir, call_dps_packed, call_tir_with_grad
if (Optional<ExprDoc> doc = PrintCallTIRDPSPacked(n, n_p, d)) {
return doc.value();
}
diff --git a/tests/python/relax/test_op_gradient_numeric.py b/tests/python/relax/test_op_gradient_numeric.py
index 574f40c753..0a7fa3672e 100644
--- a/tests/python/relax/test_op_gradient_numeric.py
+++ b/tests/python/relax/test_op_gradient_numeric.py
@@ -234,8 +234,6 @@ def test_unary(target, dev, unary_op_func, can_be_neg):
(relax.op.multiply,),
(relax.op.divide,),
(relax.op.power,),
- (relax.op.maximum,),
- (relax.op.minimum,),
)
@@ -246,6 +244,24 @@ def test_binary_arith(target, dev, binary_arith_op_func):
relax_check_gradients(binary_arith_op_func, [data1_numpy, data2_numpy], target, dev)
+(binary_minmax_op_func,) = tvm.testing.parameters(
+ (relax.op.maximum,),
+ (relax.op.minimum,),
+)
+
+
+@tvm.testing.parametrize_targets("llvm")
+def test_binary_minmax(target, dev, binary_minmax_op_func):
+ # Checking numerical gradient of min and max requires data1_numpy[i] != data2_numpy[i]
+ # for all possible i.
+ # If data1_numpy[i] == data2_numpy[i], the operator is not differentiable w.r.t. place i
+ data1_numpy = np.random.uniform(1, 1.1, (3, 3)).astype(np.float32)
+ delta = np.random.uniform(1, 1.1, (3, 3)).astype(np.float32)
+ sign = np.random.randint(0, 2, (3, 3)).astype(np.float32) * 2 - 1
+ data2_numpy = data1_numpy + delta * sign
+ relax_check_gradients(binary_minmax_op_func, [data1_numpy, data2_numpy], target, dev)
+
+
(binary_cmp_op_func,) = tvm.testing.parameters(
(relax.op.equal,),
(relax.op.greater,),
diff --git a/tests/python/relax/test_op_misc.py b/tests/python/relax/test_op_misc.py
index 2eeaa15208..2bbc18af6b 100644
--- a/tests/python/relax/test_op_misc.py
+++ b/tests/python/relax/test_op_misc.py
@@ -43,6 +43,25 @@ def test_call_tir() -> None:
v1 = rx.call_tir(identity_tir, [v0], R.Tensor((54, 96), "float32"))
+def test_call_tir_with_grad():
+ v0 = rx.Var("v0", R.Tensor([54, 96], "float32"))
+ v1 = rx.call_tir_with_grad(
+ identity_tir, (v0,), R.Tensor((54, 96), "float32"), te_grad_name="identity_grad"
+ )
+ assert v1.attrs.te_grad_name == "identity_grad"
+ v2 = rx.call_tir_with_grad(
+ identity_tir,
+ (v0,),
+ R.Tensor((54, 96), "float32"),
+ te_grad_name="identity_k_grad",
+ te_grad_kwargs={"k": 1.0},
+ )
+ assert v2.attrs.te_grad_name == "identity_k_grad"
+ assert isinstance(v2.attrs.te_grad_kwargs, tvm.ir.container.Map)
+ val = v2.attrs.te_grad_kwargs.items()[0]
+ assert val[0] == "k" and float(val[1]) == 1.0
+
+
def test_implicit_op():
m, n = tvm.tir.Var("m", "int64"), tvm.tir.Var("n", "int64")
x = rx.Var("x", R.Tensor([m, n], "float32"))
diff --git a/tests/python/relax/test_training_register_te_gradient.py b/tests/python/relax/test_training_register_te_gradient.py
new file mode 100644
index 0000000000..9803e4fab3
--- /dev/null
+++ b/tests/python/relax/test_training_register_te_gradient.py
@@ -0,0 +1,280 @@
+# 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.
+"""Unit tests for relax training utils."""
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.ir.base import assert_structural_equal
+from tvm.script.parser import relax as R, tir as T, ir as I
+
+from tvm.relax.training.utils import register_te_gradient
+from tvm.relax.transform import Gradient
+
+
+# Only run once in the whole test session
+@pytest.fixture(scope="module")
+def register_te_grads():
+ # register the gradient function
+ @register_te_gradient("f_mul_grad")
+ def f_mul_grad(output_grad, src1, src2):
+ def mul_grad_1(*idx):
+ return src2[idx] * output_grad[idx]
+
+ def mul_grad_2(*idx):
+ return src1[idx] * output_grad[idx]
+
+ return [
+ tvm.te.compute(src1.shape, mul_grad_1, name="f_mul_grad_1"),
+ tvm.te.compute(src2.shape, mul_grad_2, name="f_mul_grad_2"),
+ ]
+
+ # register the gradient function
+ @register_te_gradient("f_mulk_grad")
+ def f_mulk_grad(output_grad, src1, k):
+ def mulk_grad(*idx):
+ return output_grad[idx] * k
+
+ return [
+ tvm.te.compute(src1.shape, mulk_grad, name="f_mulk_grad"),
+ ]
+
+
+def get_expected_1():
+ # fmt: off
+ @I.ir_module
+ class Expected:
+ @T.prim_func
+ def f_mul(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), B: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul_1: T.Buffer((T.int64(5), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1 in T.grid(T.int64(5), T.int64(5)):
+ with T.block("f_mul"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(A[v_i0, v_i1], B[v_i0, v_i1])
+ T.writes(f_mul_1[v_i0, v_i1])
+ f_mul_1[v_i0, v_i1] = A[v_i0, v_i1] * B[v_i0, v_i1]
+
+ @T.prim_func
+ def f_mul_grad(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), B: T.Buffer((T.int64(5), T.int64(5)), "float32"), C: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul_grad_1: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul_grad_2: T.Buffer((T.int64(5), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1 in T.grid(T.int64(5), T.int64(5)):
+ with T.block("f_mul_grad_1"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(C[v_i0, v_i1], A[v_i0, v_i1])
+ T.writes(f_mul_grad_1[v_i0, v_i1])
+ f_mul_grad_1[v_i0, v_i1] = C[v_i0, v_i1] * A[v_i0, v_i1]
+ for i0, i1 in T.grid(T.int64(5), T.int64(5)):
+ with T.block("f_mul_grad_2"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(B[v_i0, v_i1], A[v_i0, v_i1])
+ T.writes(f_mul_grad_2[v_i0, v_i1])
+ f_mul_grad_2[v_i0, v_i1] = B[v_i0, v_i1] * A[v_i0, v_i1]
+
+ @R.function
+ def main_adjoint(a: R.Tensor((5, 5), dtype="float32"), b: R.Tensor((5, 5), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((5, 5), dtype="float32"), R.Tensor((5, 5), dtype="float32"))):
+ cls = Expected
+ with R.dataflow():
+ lv = R.call_tir(cls.f_mul, (a, b), out_sinfo=R.Tensor((5, 5), dtype="float32"), te_grad_name="f_mul_grad")
+ gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
+ gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
+ lv_adjoint: R.Tensor((5, 5), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([5, 5]))
+ lv_1 = R.call_tir(cls.f_mul_grad, (lv_adjoint, a, b), out_sinfo=[R.Tensor((5, 5), dtype="float32"), R.Tensor((5, 5), dtype="float32")])
+ a_adjoint: R.Tensor((5, 5), dtype="float32") = lv_1[0]
+ b_adjoint: R.Tensor((5, 5), dtype="float32") = lv_1[1]
+ R.output(gv, a_adjoint, b_adjoint)
+ return (gv, (a_adjoint, b_adjoint))
+
+ @R.function
+ def main(a: R.Tensor((5, 5), dtype="float32"), b: R.Tensor((5, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
+ cls = Expected
+ with R.dataflow():
+ lv = R.call_tir(cls.f_mul, (a, b), out_sinfo=R.Tensor((5, 5), dtype="float32"), te_grad_name="f_mul_grad")
+ gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
+ R.output(gv)
+ return gv
+ # fmt: on
+ return Expected
+
+
+@pytest.mark.skip("gradient will be refactored in later pull requests")
+def test_emit_te(register_te_grads):
+ # Build the target module using emit_te
+ def f_mul(src1, src2):
+ def mul(*idx):
+ return src1[idx] * src2[idx]
+
+ return tvm.te.compute(src1.shape, mul, name="f_mul")
+
+ a = relax.Var("a", relax.TensorStructInfo([5, 5], "float32"))
+ b = relax.Var("b", relax.TensorStructInfo([5, 5], "float32"))
+
+ bb = relax.BlockBuilder()
+ with bb.function("main", [a, b]):
+ with bb.dataflow():
+ d = bb.emit_te(f_mul, a, b, primfunc_name_hint="f_mul", te_grad_name="f_mul_grad")
+ out = bb.emit_output(R.sum(d))
+ bb.emit_func_output(out)
+
+ Before = bb.get()
+ After = Gradient("main")(Before)
+ assert_structural_equal(After, get_expected_1())
+
+
+@pytest.mark.skip("gradient will be refactored in later pull requests")
+def test_call_tir(register_te_grads):
+ # fmt: off
+ @I.ir_module
+ class Before:
+ @T.prim_func
+ def f_mul(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), B: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul_1: T.Buffer((T.int64(5), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1 in T.grid(T.int64(5), T.int64(5)):
+ with T.block("f_mul"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(A[v_i0, v_i1], B[v_i0, v_i1])
+ T.writes(f_mul_1[v_i0, v_i1])
+ f_mul_1[v_i0, v_i1] = A[v_i0, v_i1] * B[v_i0, v_i1]
+
+ @R.function
+ def main(a: R.Tensor((5, 5), dtype="float32"), b: R.Tensor((5, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
+ cls = Before
+ with R.dataflow():
+ lv = R.call_tir(cls.f_mul, (a, b), out_sinfo=R.Tensor((5, 5), dtype="float32"), te_grad_name="f_mul_grad")
+ gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
+ R.output(gv)
+ return gv
+ # fmt: off
+
+ After = Gradient("main")(Before)
+ assert_structural_equal(After, get_expected_1())
+
+
+def get_expected_2():
+ # fmt: off
+ @I.ir_module
+ class Expected:
+ @T.prim_func
+ def f_mul(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul2: T.Buffer((T.int64(5), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1 in T.grid(T.int64(5), T.int64(5)):
+ with T.block("f_mul2"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(A[v_i0, v_i1])
+ T.writes(f_mul2[v_i0, v_i1])
+ f_mul2[v_i0, v_i1] = A[v_i0, v_i1] * T.float32(2)
+
+ @T.prim_func
+ def f_mulk_grad(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), B: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mulk_grad_1: T.Buffer((T.int64(5), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1 in T.grid(T.int64(5), T.int64(5)):
+ with T.block("f_mulk_grad"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(A[v_i0, v_i1])
+ T.writes(f_mulk_grad_1[v_i0, v_i1])
+ f_mulk_grad_1[v_i0, v_i1] = A[v_i0, v_i1] * T.float32(2)
+
+ @R.function
+ def main_adjoint(a: R.Tensor((5, 5), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((5, 5), dtype="float32"))):
+ cls = Expected
+ with R.dataflow():
+ lv = R.call_tir(cls.f_mul, (a,), out_sinfo=R.Tensor((5, 5), dtype="float32"), te_grad_name="f_mulk_grad", te_grad_kwargs={"k": T.float32(2)})
+ gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
+ gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
+ lv_adjoint: R.Tensor((5, 5), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([5, 5]))
+ lv_1 = R.call_tir(cls.f_mulk_grad, (lv_adjoint, a), out_sinfo=R.Tensor((5, 5), dtype="float32"))
+ a_adjoint: R.Tensor((5, 5), dtype="float32") = lv_1
+ R.output(gv, a_adjoint)
+ return (gv, (a_adjoint,))
+
+ @R.function
+ def main(a: R.Tensor((5, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
+ cls = Expected
+ with R.dataflow():
+ lv = R.call_tir(cls.f_mul, (a,), out_sinfo=R.Tensor((5, 5), dtype="float32"), te_grad_name="f_mulk_grad", te_grad_kwargs={"k": T.float32(2)})
+ gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
+ R.output(gv)
+ return gv
+ # fmt: on
+ return Expected
+
+
+@pytest.mark.skip("gradient will be refactored in later pull requests")
+def test_emit_te_kwargs(register_te_grads):
+ # Build the target module using emit_te
+ def f_mul2(src):
+ return tvm.te.compute(src.shape, lambda *idx: src[idx] * T.float32(2), name="f_mul2")
+
+ a = relax.Var("a", relax.TensorStructInfo([5, 5], "float32"))
+
+ bb = relax.BlockBuilder()
+ with bb.function("main", [a]):
+ with bb.dataflow():
+ d = bb.emit_te(
+ f_mul2,
+ a,
+ primfunc_name_hint="f_mul",
+ te_grad_name="f_mulk_grad",
+ te_grad_kwargs={"k": T.float32(2)},
+ )
+ out = bb.emit_output(R.sum(d))
+ bb.emit_func_output(out)
+
+ Before = bb.get()
+ After = Gradient("main")(Before)
+ After.show(None, False)
+
+ assert_structural_equal(After, get_expected_2())
+
+
+@pytest.mark.skip("gradient will be refactored in later pull requests")
+def test_call_tir_kwargs(register_te_grads):
+ # fmt: off
+ @I.ir_module
+ class Before:
+ @T.prim_func
+ def f_mul(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul2: T.Buffer((T.int64(5), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ for i0, i1 in T.grid(T.int64(5), T.int64(5)):
+ with T.block("f_mul2"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ T.reads(A[v_i0, v_i1])
+ T.writes(f_mul2[v_i0, v_i1])
+ f_mul2[v_i0, v_i1] = A[v_i0, v_i1] * T.float32(2)
+
+ @R.function
+ def main(a: R.Tensor((5, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
+ cls = Before
+ with R.dataflow():
+ lv = R.call_tir(cls.f_mul, (a,), out_sinfo=R.Tensor((5, 5), dtype="float32"), te_grad_name="f_mulk_grad", te_grad_kwargs={"k": T.float32(2)})
+ gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
+ R.output(gv)
+ return gv
+ # fmt: on
+
+ After = Gradient("main")(Before)
+ assert_structural_equal(After, get_expected_2())
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_tvmscript_parser.py b/tests/python/relax/test_tvmscript_parser.py
index c9aa16b9b2..e917cc795c 100644
--- a/tests/python/relax/test_tvmscript_parser.py
+++ b/tests/python/relax/test_tvmscript_parser.py
@@ -865,6 +865,34 @@ def test_call_tir_with_tir_var():
_check(Module)
+def test_call_tir_with_grad():
+ @I.ir_module
+ class Module:
+ @T.prim_func
+ def identity_tir(a: T.handle, b: T.handle) -> None:
+ A = T.match_buffer(a, [54, 96])
+ B = T.match_buffer(b, [54, 96])
+
+ for i, j in T.grid(54, 96):
+ with T.block("compute"):
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = A[vi, vj]
+
+ @R.function
+ def main(v0: R.Tensor([54, 96], "float32")):
+ cls = Module
+ out = R.call_tir_with_grad(
+ cls.identity_tir,
+ (v0,),
+ R.Tensor((54, 96), "float32"),
+ te_grad_name="identity_k_grad",
+ te_grad_kwargs={"k": 1.0},
+ )
+ return out
+
+ _check(Module)
+
+
def test_local_function():
@R.function
def main(
diff --git a/tests/python/relax/test_tvmscript_printer_relax.py b/tests/python/relax/test_tvmscript_printer_relax.py
index c376943173..9f4ffd9acd 100644
--- a/tests/python/relax/test_tvmscript_printer_relax.py
+++ b/tests/python/relax/test_tvmscript_printer_relax.py
@@ -376,6 +376,27 @@ R.call_dps_packed("my_dps_func", (a,), out_sinfo=R.Tensor((1, x, 3), dtype="floa
)
+def test_call_tir_with_grad():
+ x = tir.Var("x", "int64")
+ v0 = relax.Var("v0", R.Tensor([54, 96], "float32"))
+ v1 = relax.call_tir_with_grad(
+ relax.GlobalVar("tir_func"),
+ (v0,),
+ R.Tensor((54, 96), "float32"),
+ te_grad_name="grad_func",
+ te_grad_kwargs={"k": 1.0, "x": x},
+ )
+
+ _assert_print(
+ v1,
+ """
+v0: R.Tensor((54, 96), dtype="float32")
+x = T.int64()
+R.call_tir_with_grad(tir_func, (v0,), out_sinfo=R.Tensor((54, 96), dtype="float32"), te_grad_name="grad_func", te_grad_kwargs={"k": T.float32(1), "x": x})
+""",
+ )
+
+
def test_seq_expr():
x = tir.Var("x", "int64")
a = relax.Var("a", relax.TensorStructInfo([1, x, 3], "float32"))