This is an automated email from the ASF dual-hosted git repository.
junrushao 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 0d462a012d [Unity] nn.Module Module and Effect (#15438)
0d462a012d is described below
commit 0d462a012dacfa7eb0809375bce2d88647fb5f69
Author: Yaxing Cai <[email protected]>
AuthorDate: Sat Jul 29 15:38:36 2023 -0700
[Unity] nn.Module Module and Effect (#15438)
This PR introduces the Module and Effect in the nn frontend. It adds the
built-in `Linear`,
`RMSNorm` and `Embedding` modules. And it adds a `KVCache` effect.
Also, this PR fix the interface of `R.nn.rms_norm` to support optional
`bias`.
Co-authored-by: Junru Shao <[email protected]>
---
include/tvm/topi/nn/rms_norm.h | 11 +-
python/tvm/relax/frontend/nn/__init__.py | 4 +-
python/tvm/relax/frontend/nn/core.py | 4 +-
python/tvm/relax/frontend/nn/modules.py | 279 ++++++++++++++++++++-
python/tvm/relax/frontend/nn/op.py | 12 +-
python/tvm/relax/op/nn/nn.py | 7 +-
python/tvm/relax/transform/legalize_ops/nn.py | 6 +-
python/tvm/topi/nn/rms_norm.py | 7 +-
src/relax/op/nn/nn.cc | 9 +-
src/relax/op/nn/nn.h | 2 +-
src/topi/nn.cc | 2 +-
tests/python/relax/test_frontend_nn_modules.py | 146 +++++++++++
tests/python/relax/test_frontend_nn_op.py | 17 +-
.../python/relax/test_transform_legalize_ops_nn.py | 118 +++++++--
14 files changed, 548 insertions(+), 76 deletions(-)
diff --git a/include/tvm/topi/nn/rms_norm.h b/include/tvm/topi/nn/rms_norm.h
index 44d38bae6d..55dac39b71 100644
--- a/include/tvm/topi/nn/rms_norm.h
+++ b/include/tvm/topi/nn/rms_norm.h
@@ -41,22 +41,18 @@ using namespace tvm::te;
* \param data N-D tensor with shape [d_0, d_1, ..., d_{N-1}]
* \param weight K-D tensor with shape [r_0, r_1, ..., r_{K-1}] where K ==
len(axis) and
* d_{axis_k} == r_k
- * \param bias Optional, K-D tensor with shape [r_0, r_1, ..., r_{K-1}] where
- * d_{axis_k} == r_k
* \param axis The axis to normalize over.
* \param epsilon The epsilon value to avoid division by zero.
* \param name The name of the operation.
* \param tag The tag to mark the operation.
* \return The normalized tensor, with the same shape as data.
*/
-inline Tensor rms_norm(const Tensor& data, const Tensor& weight, const Tensor&
bias,
- const Array<Integer>& axis, double epsilon, std::string
name = "T_rms_norm",
+inline Tensor rms_norm(const Tensor& data, const Tensor& weight, const
Array<Integer>& axis,
+ double epsilon, std::string name = "T_rms_norm",
std::string tag = kInjective) {
const auto& data_type = data->dtype;
const auto& weight_type = weight.defined() ? weight->dtype : data_type;
ICHECK(data_type == weight_type) << "rms_norm: data and weight must have the
same type";
- const auto& bias_type = bias.defined() ? bias->dtype : data_type;
- ICHECK(data_type == bias_type) << "rms_norm: data and bias must have the
same type";
auto square = multiply(data, data);
auto square_sum = sum(square, axis, /*keepdims=*/false, /*atleast1d=*/true);
@@ -80,9 +76,6 @@ inline Tensor rms_norm(const Tensor& data, const Tensor&
weight, const Tensor& b
auto output =
data(indices) * weight(reduce_indices) *
tvm::rsqrt(square_sum(non_reduce_indices) / reduce_extent +
make_const(data_type, epsilon));
- if (bias.defined()) {
- output += bias(reduce_indices);
- }
return output;
};
auto rms_norm = tvm::te::compute(data->shape, rms_norm_func, name, tag);
diff --git a/python/tvm/relax/frontend/nn/__init__.py
b/python/tvm/relax/frontend/nn/__init__.py
index 8375cb6618..cb8ba85295 100644
--- a/python/tvm/relax/frontend/nn/__init__.py
+++ b/python/tvm/relax/frontend/nn/__init__.py
@@ -15,5 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""A PyTorch-like API to build IRModules."""
-from . import spec
+from . import op, spec
from .core import Effect, Module, ModuleList, Parameter, Tensor
+from .modules import Embedding, IOEffect, KVCache, Linear, RMSNorm
+from .op import *
diff --git a/python/tvm/relax/frontend/nn/core.py
b/python/tvm/relax/frontend/nn/core.py
index 46197b7141..918c96d568 100644
--- a/python/tvm/relax/frontend/nn/core.py
+++ b/python/tvm/relax/frontend/nn/core.py
@@ -218,7 +218,7 @@ class Parameter(Tensor):
def to(self, dtype: Optional[str] = None) -> None: # pylint:
disable=invalid-name
"""Change the dtype of the parameter if it is not bound to any
concrete data"""
- if dtype is not None and self._data is not None:
+ if dtype is not None:
if self._data is not None:
raise ValueError(
"Changing the dtype of a Parameter that has been bound to
concrete "
@@ -267,7 +267,7 @@ class Module:
Yields
------
- (str, Parameter) – Tuple containing the name and parameter
+ (str, Parameter) - Tuple containing the name and parameter
"""
yield from _attribute_finder(
self, prefix, condition_yield=lambda x: isinstance(x, Parameter)
diff --git a/python/tvm/relax/frontend/nn/modules.py
b/python/tvm/relax/frontend/nn/modules.py
index 20f5a1d3c3..228f6bdd6e 100644
--- a/python/tvm/relax/frontend/nn/modules.py
+++ b/python/tvm/relax/frontend/nn/modules.py
@@ -15,11 +15,15 @@
# specific language governing permissions and limitations
# under the License.
"""Builtin Modules."""
-from typing import List, Optional
+from typing import List, Optional, Sequence, Union
from tvm import relax as rx
+from tvm import tir
+from tvm._ffi import register_func
+from tvm.runtime import NDArray
-from .core import Effect, Tensor
+from . import op
+from .core import Effect, Module, Parameter, Tensor, get_default_dtype
class IOEffect(Effect):
@@ -49,3 +53,274 @@ class IOEffect(Effect):
def print_(self, tensor: Tensor) -> None:
"""Encloses the side effect of NDArray printing"""
raise NotImplementedError
+
+
+@register_func("effect.print")
+def _print(_, array: NDArray) -> None:
+ print(f"effect.print: shape = {array.shape}, dtype = {array.dtype}, data
=\n{array}")
+
+
+class Linear(Module):
+ """
+ Module for linear layer.
+ """
+
+ def __init__( # pylint: disable=too-many-arguments
+ self,
+ in_features: int,
+ out_features: int,
+ bias: bool = True,
+ dtype: Optional[str] = None,
+ out_dtype: Optional[str] = None,
+ ):
+ super().__init__()
+ self.in_features = in_features
+ self.out_features = out_features
+ self.out_dtype = out_dtype
+ self.weight = Parameter((out_features, in_features), dtype)
+ if bias:
+ self.bias = Parameter((out_features,), dtype)
+ else:
+ self.bias = None
+
+ def forward(self, x: Tensor) -> Tensor: # pylint: disable=invalid-name
+ """
+ Forward method for linear layer.
+
+ Parameters
+ ----------
+ x : Tensor
+ The input tensor.
+
+ Returns
+ -------
+ ret : Tensor
+ The output tensor for the linear layer.
+ """
+ # x: [*B, in_features]
+ # w: [in_features, out_features]
+ w = op.permute_dims(self.weight) # pylint: disable=invalid-name
+ # x: [*B, out_features]
+ x = op.matmul(x, w, out_dtype=self.out_dtype)
+ if self.bias is not None:
+ x = x + self.bias
+ return x
+
+
+class RMSNorm(Module):
+ """
+ Module for rms norm layer.
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ axes: Union[int, List[int]],
+ epsilon: float = 1e-5,
+ bias: bool = True,
+ dtype: Optional[str] = None,
+ ):
+ super().__init__()
+ self.epsilon = epsilon
+ self.axes = axes
+ self.weight = Parameter((hidden_size,), dtype=dtype)
+ if bias:
+ self.bias = Parameter((hidden_size,), dtype=dtype)
+ else:
+ self.bias = None
+
+ # pylint: disable=invalid-name
+ def forward(self, x: Tensor):
+ """
+ Forward method for rms norm layer.
+
+ Parameters
+ ----------
+ x : Tensor
+ The input tensor.
+
+ Returns
+ -------
+ ret : Tensor
+ The output tensor for the rms norm layer.
+ """
+ out = op.rms_norm(x, weight=self.weight, axes=self.axes,
epsilon=self.epsilon)
+ if self.bias:
+ out = op.add(out, self.bias)
+ return out
+
+ # pylint: enable=invalid-name
+
+
+class KVCache(Effect):
+ """
+ Effect to implement KVCache.
+ """
+
+ init_seq_len: int
+ unit_shape: List[int]
+ dtype: str
+ cache: Optional[rx.Var]
+
+ def __init__(
+ self,
+ init_seq_len: int,
+ unit_shape: Sequence[int],
+ dtype: Optional[str] = None,
+ ):
+ if dtype is None:
+ dtype = get_default_dtype()
+ # Usually the shape is: [init_seq_len, num_heads, head_dim]
+ # and unit_shape = [num_heads, head_dim]
+ self.init_seq_len = init_seq_len
+ self.unit_shape = [int(i) for i in unit_shape]
+ self.dtype = dtype
+
+ def emit_init(self, name_hint: str, bb: rx.BlockBuilder): # pylint:
disable=arguments-renamed
+ """
+ Emit the initialization of the KVCache effect.
+
+ Parameters
+ ----------
+ name_hint : str
+ The name hint of the initialization binding Var.
+
+ bb : relax.BlockBuilder
+ The relax BlockBuilder to emit.
+ """
+ init_shape = rx.ShapeExpr([self.init_seq_len] + self.unit_shape)
+ return [
+ bb.emit(
+ rx.Call(
+ rx.extern("vm.builtin.attention_kv_cache_create"),
+ args=[rx.op.zeros(init_shape, self.dtype), init_shape,
rx.PrimValue(0)],
+ sinfo_args=[rx.ObjectStructInfo()],
+ ),
+ name_hint=name_hint,
+ )
+ ]
+
+ def create(self, name_hint: str) -> rx.Var:
+ """
+ Create the implicit inputs to a relax.Function that represents the
KVCache effect.
+
+ Parameters
+ ----------
+ name_hint : str
+ The name hint of the relax.Var.
+
+ Returns
+ -------
+ ret : relax.Var
+ The relax.Var for KVCache.
+ """
+ self.cache = rx.Var(name_hint, struct_info=rx.ObjectStructInfo())
+ return [self.cache]
+
+ def finalize(self) -> List[rx.Var]:
+ """
+ Finalize the KVCache effect as the implicit return value of a
relax.Function.
+
+ Returns
+ -------
+ ret : List[rx.Var]
+ The output relax.Var as KVCache.
+ """
+ result = self.cache
+ self.cache = None
+ return [result]
+
+ def to(self, dtype: Optional[str] = None) -> None:
+ """
+ Convert the KVCache effect to specific dtype.
+
+ Parameters
+ ----------
+ dtype : Optional[str]
+ The target data type to convert.
+ """
+ if dtype is not None:
+ self.dtype = dtype
+
+ def view(self, seq_len: tir.Var) -> Tensor:
+ """
+ View the last elements in KVCache.
+
+ Parameters
+ ----------
+ seq_len : tir.Var
+ The number of last elements to view.
+
+ Returns
+ -------
+ ret : Tensor
+ The last tensor to view.
+ """
+ shape = rx.ShapeExpr([seq_len] + self.unit_shape)
+ return Tensor(
+ _expr=rx.BlockBuilder.current().emit(
+ rx.Call(
+ rx.extern("vm.builtin.attention_kv_cache_view"),
+ args=[self.cache, shape],
+ sinfo_args=[rx.TensorStructInfo(shape, self.dtype)],
+ )
+ )
+ )
+
+ def append(self, new_element: Tensor) -> None:
+ """
+ Append a new element in KVCache.
+
+ Parameters
+ ----------
+ new_element : Tensor
+ The new tensor to append.
+ """
+ if new_element.dtype != self.dtype:
+ raise TypeError(
+ f'KVCache has been set to use dtype "{self.dtype}", '
+ f'but got "{new_element.dtype}"'
+ )
+ self.cache = rx.BlockBuilder.current().emit(
+ rx.Call(
+ rx.extern("vm.builtin.attention_kv_cache_append"),
+ args=[self.cache, new_element._expr], # pylint:
disable=protected-access
+ sinfo_args=[rx.ObjectStructInfo()],
+ )
+ )
+
+
+class Embedding(Module):
+ """
+ Module for embedding layer.
+ """
+
+ def __init__(self, num: int, dim: int, dtype: Optional[str] = None):
+ self.num = num
+ self.dim = dim
+ self.weight = Parameter((num, dim), dtype=dtype)
+
+ def forward(self, x: Tensor): # pylint: disable=invalid-name
+ """
+ Forward method for embedding layer.
+
+ Parameters
+ ----------
+ x : Tensor
+ The input tensor.
+
+ Returns
+ -------
+ ret : Tensor
+ The output tensor for the embedding layer.
+ """
+ if x.ndim == 1:
+ return op.take(self.weight, x, axis=0)
+ return op.reshape(
+ op.take(
+ self.weight,
+ op.reshape(x, shape=[-1]),
+ axis=0,
+ ),
+ shape=[*x.shape, self.dim], # TODO(@junrushao): revisit and
remove self.dim
+ )
diff --git a/python/tvm/relax/frontend/nn/op.py
b/python/tvm/relax/frontend/nn/op.py
index 74a337d8a3..9a98aa1adb 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -488,7 +488,6 @@ def softmax(x: Tensor, axis: int = -1, name: str =
"softmax") -> Tensor:
def rms_norm(
x: Tensor,
weight: Tensor,
- bias: Optional[Tensor],
axes: Union[int, List[int]],
epsilon: float = 1e-5,
name: str = "rms_norm",
@@ -501,7 +500,7 @@ def rms_norm(
.. math::
- out = \frac{data}{\sqrt{mean(data, axis)+\epsilon}} * weight + bias
+ out = \frac{data}{\sqrt{mean(data, axis)+\epsilon}} * weight
Parameters
----------
@@ -511,9 +510,6 @@ def rms_norm(
weight : Tensor
The scale factor.
- bias : Tensor
- Optional offset factor.
-
axes : Union[int, List[int]]
The axes that along which the normalization is applied.
@@ -528,11 +524,7 @@ def rms_norm(
result : Tensor
The computed result.
"""
- if bias is None:
- bias = _op.zeros(weight.shape, dtype=weight.dtype)
- else:
- bias = bias._expr
- return _wrap_nested(_op.nn.rms_norm(x._expr, weight._expr, bias, axes,
epsilon), name)
+ return _wrap_nested(_op.nn.rms_norm(x._expr, weight._expr, axes, epsilon),
name)
def triu(x: Tensor, diagonal: int = 0, name: str = "triu") -> Tensor:
diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py
index edba3c505f..aa5c7b18f6 100644
--- a/python/tvm/relax/op/nn/nn.py
+++ b/python/tvm/relax/op/nn/nn.py
@@ -20,8 +20,8 @@ from typing import List, Optional, Tuple, Union
from tvm import DataType
from tvm.tir import FloatImm
-from . import _ffi_api
from ...expr import Expr
+from . import _ffi_api
def conv1d(
@@ -928,8 +928,7 @@ def group_norm(
def rms_norm(
data: Expr,
weight: Expr,
- bias: Expr,
- axes: Union[int, List[int]],
+ axes: Union[int, List[int]] = -1,
epsilon: float = 1e-5,
) -> Expr:
r"""
@@ -966,7 +965,7 @@ def rms_norm(
"""
if isinstance(axes, int):
axes = [axes]
- return _ffi_api.rms_norm(data, weight, bias, axes, epsilon) # type: ignore
+ return _ffi_api.rms_norm(data, weight, axes, epsilon) # type: ignore
def dropout(data: Expr, rate: float = 0.5) -> Expr:
diff --git a/python/tvm/relax/transform/legalize_ops/nn.py
b/python/tvm/relax/transform/legalize_ops/nn.py
index b66b2d7b5d..5386fbf7cb 100644
--- a/python/tvm/relax/transform/legalize_ops/nn.py
+++ b/python/tvm/relax/transform/legalize_ops/nn.py
@@ -20,10 +20,11 @@ import logging
import math
from typing import Optional
-from tvm import topi, tir, te
+from tvm import te, tir, topi
+
from ...block_builder import BlockBuilder
from ...expr import Call, Expr
-from .common import register_legalize, _call_topi_without_attr
+from .common import _call_topi_without_attr, register_legalize
@register_legalize("relax.nn.conv1d")
@@ -340,7 +341,6 @@ def _nn_rms_norm(bb: BlockBuilder, call: Call) -> Expr:
topi.nn.rms_norm,
call.args[0],
call.args[1],
- call.args[2],
axis=call.attrs.axes,
epsilon=call.attrs.epsilon,
)
diff --git a/python/tvm/topi/nn/rms_norm.py b/python/tvm/topi/nn/rms_norm.py
index f2f5a7e674..9284517468 100644
--- a/python/tvm/topi/nn/rms_norm.py
+++ b/python/tvm/topi/nn/rms_norm.py
@@ -18,7 +18,7 @@
from .. import cpp
-def rms_norm(data, weight, bias, axis, epsilon=1e-5):
+def rms_norm(data, weight, axis, epsilon=1e-5):
"""Root mean square normalization operator. The output will have the same
data type as input.
Parameters
@@ -29,9 +29,6 @@ def rms_norm(data, weight, bias, axis, epsilon=1e-5):
weight: tvm.te.Tensor
K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and
d_{axis_k} == r_k
- bias: tvm.te.Tensor
- Optional, K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis)
and d_{axis_k} == r_k
-
axis : list of int
Axis over the normalization applied
@@ -43,4 +40,4 @@ def rms_norm(data, weight, bias, axis, epsilon=1e-5):
result : tvm.te.Tensor
N-D with shape (d_0, d_1, ..., d_{N-1})
"""
- return cpp.nn.rms_norm(data, weight, bias, axis, epsilon)
+ return cpp.nn.rms_norm(data, weight, axis, epsilon)
diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc
index 79ea8650dd..cbd71da848 100644
--- a/src/relax/op/nn/nn.cc
+++ b/src/relax/op/nn/nn.cc
@@ -440,13 +440,13 @@ TVM_REGISTER_OP("relax.nn.group_norm")
/* relax.nn.rms_norm */
TVM_REGISTER_NODE_TYPE(RMSNormAttrs);
-Expr rms_norm(Expr data, Expr weight, Expr bias, Array<Integer> axes, double
epsilon) {
+Expr rms_norm(Expr data, Expr weight, Array<Integer> axes, double epsilon) {
ObjectPtr<RMSNormAttrs> attrs = make_object<RMSNormAttrs>();
attrs->axes = std::move(axes);
attrs->epsilon = epsilon;
static const Op& op = Op::Get("relax.nn.rms_norm");
- return Call(op, {std::move(data), std::move(weight), std::move(bias)},
Attrs{attrs}, {});
+ return Call(op, {std::move(data), std::move(weight)}, Attrs{attrs}, {});
}
TVM_REGISTER_GLOBAL("relax.op.nn.rms_norm").set_body_typed(rms_norm);
@@ -466,7 +466,7 @@ InferLayoutOutput InferLayoutRMSNorm(const Call& call,
const VarLayoutMap& var_layout_map) {
ICHECK(NoDesiredLayout(call, desired_layouts));
std::vector<NLayout> initial_layouts;
- for (size_t i = 0; i < 3; ++i) {
+ for (size_t i = 0; i < 2; ++i) {
const auto* tensor_sinfo =
GetStructInfoAs<TensorStructInfoNode>(call->args[i]);
ICHECK(tensor_sinfo != nullptr) << "Invalid Call";
ICHECK(!tensor_sinfo->IsUnknownNdim()) << "Only support known ndim";
@@ -488,10 +488,9 @@ InferLayoutOutput InferLayoutRMSNorm(const Call& call,
TVM_REGISTER_OP("relax.nn.rms_norm")
.set_attrs_type<RMSNormAttrs>()
- .set_num_inputs(3)
+ .set_num_inputs(2)
.add_argument("data", "Tensor", "Input to which rms_norm will be applied.")
.add_argument("weight", "Tensor", "The scale factor.")
- .add_argument("bias", "Tensor", "The offset factor.")
.set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoRMSNorm)
.set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutRMSNorm)
.set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy",
MixedPrecisionPolicyKind::kFollow)
diff --git a/src/relax/op/nn/nn.h b/src/relax/op/nn/nn.h
index 624cfe9078..a3658fed54 100644
--- a/src/relax/op/nn/nn.h
+++ b/src/relax/op/nn/nn.h
@@ -79,7 +79,7 @@ Expr group_norm(Expr data, Expr gamma, Expr beta, int
num_groups, int channel_ax
Array<Integer> axes, double epsilon, bool center, bool scale);
/*! \brief Compute root mean square normalization. */
-Expr rms_norm(Expr data, Expr weight, Expr bias, Array<Integer> axes, double
epsilon);
+Expr rms_norm(Expr data, Expr weight, Array<Integer> axes, double epsilon);
/*!
* \brief Applies the dropout operation to the input tensor.
diff --git a/src/topi/nn.cc b/src/topi/nn.cc
index ba88f01c68..b6198b6ad4 100644
--- a/src/topi/nn.cc
+++ b/src/topi/nn.cc
@@ -179,7 +179,7 @@
TVM_REGISTER_GLOBAL("topi.nn.instance_norm").set_body([](TVMArgs args, TVMRetVal
/* Ops from nn/rms_norm.h */
TVM_REGISTER_GLOBAL("topi.nn.rms_norm").set_body([](TVMArgs args, TVMRetValue*
rv) {
- *rv = nn::rms_norm(args[0], args[1], args[2], args[3],
static_cast<double>(args[4]));
+ *rv = nn::rms_norm(args[0], args[1], args[2], args[3].operator double());
});
} // namespace topi
diff --git a/tests/python/relax/test_frontend_nn_modules.py
b/tests/python/relax/test_frontend_nn_modules.py
new file mode 100644
index 0000000000..27bbd68329
--- /dev/null
+++ b/tests/python/relax/test_frontend_nn_modules.py
@@ -0,0 +1,146 @@
+# 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.
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.ir import assert_structural_equal
+from tvm.relax.frontend.nn import core, modules, spec
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+def test_linear():
+ @R.function
+ def forward(
+ x: R.Tensor((1, 4), dtype="float32"),
+ weight: R.Tensor((8, 4), dtype="float32"),
+ bias: R.Tensor((8,), dtype="float32"),
+ _io: R.Object,
+ ) -> R.Tuple(R.Tensor((1, 8), dtype="float32"), R.Tuple(R.Object)):
+ with R.dataflow():
+ permute_dims: R.Tensor((4, 8), dtype="float32") =
R.permute_dims(weight, axes=None)
+ matmul: R.Tensor((1, 8), dtype="float32") = R.matmul(x,
permute_dims, out_dtype="void")
+ add: R.Tensor((1, 8), dtype="float32") = R.add(matmul, bias)
+ gv1: R.Tuple(R.Tensor((1, 8), dtype="float32"), R.Tuple(R.Object))
= add, (_io,)
+ R.output(gv1)
+ return gv1
+
+ mod = modules.Linear(4, 8)
+ tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((1, 4),
"float32")}})
+ assert_structural_equal(tvm_mod["forward"], forward, True)
+
+
+def test_rms_norm():
+ @R.function
+ def forward(
+ x: R.Tensor((2, 4, 8), dtype="float32"),
+ weight: R.Tensor((8,), dtype="float32"),
+ _io: R.Object,
+ ) -> R.Tuple(R.Tensor((2, 4, 8), dtype="float32"), R.Tuple(R.Object)):
+ with R.dataflow():
+ rms_norm: R.Tensor((2, 4, 8), dtype="float32") = R.nn.rms_norm(
+ x, weight, axes=[2], epsilon=1.0000000000000001e-05
+ )
+ gv1: R.Tuple(R.Tensor((2, 4, 8), dtype="float32"),
R.Tuple(R.Object)) = rms_norm, (_io,)
+ R.output(gv1)
+ return gv1
+
+ mod = modules.RMSNorm(8, [2], bias=False)
+ tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((2, 4, 8),
"float32")}})
+ tvm_mod.show()
+ assert_structural_equal(tvm_mod["forward"], forward, True)
+
+
+def test_embedding():
+ @R.function
+ def forward(
+ x: R.Tensor((1, 4), dtype="int32"),
+ weight: R.Tensor((4, 8), dtype="float32"),
+ _io: R.Object,
+ ) -> R.Tuple(R.Tensor((1, 4, 8), dtype="float32"), R.Tuple(R.Object)):
+ with R.dataflow():
+ reshape: R.Tensor((4,), dtype="int32") = R.reshape(x, R.shape([4]))
+ take: R.Tensor((4, 8), dtype="float32") = R.take(weight, reshape,
axis=0)
+ reshape1: R.Tensor((1, 4, 8), dtype="float32") = R.reshape(take,
R.shape([1, 4, 8]))
+ gv1: R.Tuple(R.Tensor((1, 4, 8), dtype="float32"),
R.Tuple(R.Object)) = reshape1, (_io,)
+ R.output(gv1)
+ return gv1
+
+ mod = modules.Embedding(4, 8, "float32")
+ tvm_mod, _ = mod.export_tvm(spec={"forward": {"x": spec.Tensor((1, 4),
"int32")}})
+ assert_structural_equal(tvm_mod["forward"], forward, True)
+
+
+def test_kv_cache():
+ @I.ir_module
+ class Module:
+ @R.function
+ def _initialize_effect() -> R.Tuple(R.Object, R.Object):
+ with R.dataflow():
+ _io: R.Object = R.null_value()
+ lv: R.Tensor((8, 2, 4), dtype="float32") = R.zeros(
+ R.shape([8, 2, 4]), dtype="float32"
+ )
+ cache: R.Object = R.call_packed(
+ "vm.builtin.attention_kv_cache_create",
+ lv,
+ R.shape([8, 2, 4]),
+ R.prim_value(0),
+ sinfo_args=(R.Object,),
+ )
+ lv1: R.Tuple(R.Object, R.Object) = _io, cache
+ gv: R.Tuple(R.Object, R.Object) = lv1
+ R.output(gv)
+ return gv
+
+ @R.function
+ def forward(
+ x: R.Tensor((2, 4), dtype="float32"), _io: R.Object, cache:
R.Object
+ ) -> R.Tuple(R.Tensor((4, 2, 4), dtype="float32"), R.Tuple(R.Object,
R.Object)):
+ with R.dataflow():
+ lv2: R.Object = R.call_packed(
+ "vm.builtin.attention_kv_cache_append", cache, x,
sinfo_args=(R.Object,)
+ )
+ lv3: R.Tensor((4, 2, 4), dtype="float32") = R.call_packed(
+ "vm.builtin.attention_kv_cache_view",
+ lv2,
+ R.shape([4, 2, 4]),
+ sinfo_args=(R.Tensor((4, 2, 4), dtype="float32"),),
+ )
+ gv1: R.Tuple(
+ R.Tensor((4, 2, 4), dtype="float32"), R.Tuple(R.Object,
R.Object)
+ ) = lv3, (_io, lv2)
+ R.output(gv1)
+ return gv1
+
+ class KVCacheTest(modules.Module):
+ def __init__(self) -> None:
+ self.cache = modules.KVCache(8, [2, 4])
+
+ def forward(self, x: core.Tensor) -> core.Tensor:
+ self.cache.append(x)
+ return self.cache.view(4)
+
+ tvm_mod, _ = KVCacheTest().export_tvm(spec={"forward": {"x":
spec.Tensor((2, 4), "float32")}})
+ assert_structural_equal(tvm_mod, Module, True)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_frontend_nn_op.py
b/tests/python/relax/test_frontend_nn_op.py
index 6668f02743..d2501cb0ef 100644
--- a/tests/python/relax/test_frontend_nn_op.py
+++ b/tests/python/relax/test_frontend_nn_op.py
@@ -15,13 +15,14 @@
# specific language governing permissions and limitations
# under the License.
import pytest
+
import tvm
import tvm.testing
from tvm import relax, te, tir
-from tvm.relax.frontend.nn import Tensor, Module, spec, op
+from tvm.relax.frontend.nn import Module, Tensor, op, spec
+from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tir as T
-from tvm.script import ir as I
def test_binary():
@@ -136,21 +137,21 @@ def test_datatype():
def test_nn():
class Model(Module):
- def test(self, x: Tensor, weight: Tensor, bias: Tensor):
+ def test(self, x: Tensor, weight: Tensor):
silu_out = op.silu(x)
softmax_out = op.softmax(x, axis=2)
- rms_norm_out = op.rms_norm(x, weight, bias, axes=[-2, -1])
- rms_norm_with_bias_out = op.rms_norm(x, weight, bias, axes=[-2,
-1])
+ rms_norm_out = op.rms_norm(x, weight, axes=[-2, -1])
+ rms_norm_with_bias_out = op.rms_norm(x, weight, axes=[-2, -1])
return x
# fmt: off
@R.function
- def test(x: R.Tensor((2, 3, 4, 5), dtype="float32"), weight: R.Tensor((4,
5), dtype="float32"), bias: R.Tensor((4, 5), dtype="float32"), _io: R.Object)
-> R.Tuple(R.Tensor((2, 3, 4, 5), dtype="float32"), R.Tuple(R.Object)):
+ def test(x: R.Tensor((2, 3, 4, 5), dtype="float32"), weight: R.Tensor((4,
5), dtype="float32"), _io: R.Object) -> R.Tuple(R.Tensor((2, 3, 4, 5),
dtype="float32"), R.Tuple(R.Object)):
with R.dataflow():
silu: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.silu(x)
softmax: R.Tensor((2, 3, 4, 5), dtype="float32") = R.nn.softmax(x,
axis=2)
- rms_norm: R.Tensor((2, 3, 4, 5), dtype="float32") =
R.nn.rms_norm(x, weight, bias, axes=[-2, -1], epsilon=1.0000000000000001e-05)
- rms_norm1: R.Tensor((2, 3, 4, 5), dtype="float32") =
R.nn.rms_norm(x, weight, bias, axes=[-2, -1], epsilon=1.0000000000000001e-05)
+ rms_norm: R.Tensor((2, 3, 4, 5), dtype="float32") =
R.nn.rms_norm(x, weight, axes=[-2, -1], epsilon=1.0000000000000001e-05)
+ rms_norm1: R.Tensor((2, 3, 4, 5), dtype="float32") =
R.nn.rms_norm(x, weight, axes=[-2, -1], epsilon=1.0000000000000001e-05)
gv1: R.Tuple(R.Tensor((2, 3, 4, 5), dtype="float32"),
R.Tuple(R.Object)) = x, (_io,)
R.output(gv1)
return gv1
diff --git a/tests/python/relax/test_transform_legalize_ops_nn.py
b/tests/python/relax/test_transform_legalize_ops_nn.py
index 66da2bae55..d750901b59 100644
--- a/tests/python/relax/test_transform_legalize_ops_nn.py
+++ b/tests/python/relax/test_transform_legalize_ops_nn.py
@@ -16,11 +16,13 @@
# under the License.
import pytest
+
import tvm
-from tvm.relax.transform import LegalizeOps
-from tvm.script import relax as R, tir as T, ir as I
import tvm.testing
-
+from tvm.relax.transform import LegalizeOps
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tir as T
##################### Neural network #####################
@@ -2747,8 +2749,8 @@ def test_rms_norm():
@tvm.script.ir_module
class RMSNorm:
@R.function
- def main(x: R.Tensor((2, 3, 4, 5), "float32"), weight: R.Tensor((4,
5), "float32"), bias: R.Tensor((4, 5), "float32")) -> R.Tensor((2, 3, 4, 5),
"float32"):
- gv: R.Tensor((2, 3, 4, 5), "float32") = R.nn.rms_norm(x, weight,
bias, axes=[-2, -1])
+ def main(x: R.Tensor((2, 3, 4, 5), "float32"), weight: R.Tensor((4,
5), "float32")) -> R.Tensor((2, 3, 4, 5), "float32"):
+ gv: R.Tensor((2, 3, 4, 5), "float32") = R.nn.rms_norm(x, weight,
axes=[-2, -1])
return gv
@tvm.script.ir_module
@@ -2757,7 +2759,6 @@ def test_rms_norm():
def rms_norm(
A: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)),
"float32"),
B: T.Buffer((T.int64(4), T.int64(5)), "float32"),
- C: T.Buffer((T.int64(4), T.int64(5)), "float32"),
T_rms_norm: T.Buffer(
(T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"
),
@@ -2795,7 +2796,6 @@ def test_rms_norm():
A[v_ax0, v_ax1, v_ax2, v_ax3],
B[v_ax2, v_ax3],
T_multiply_red[v_ax0, v_ax1],
- C[v_ax2, v_ax3],
)
T.writes(T_rms_norm[v_ax0, v_ax1, v_ax2, v_ax3])
T_rms_norm[v_ax0, v_ax1, v_ax2, v_ax3] = (
@@ -2805,19 +2805,17 @@ def test_rms_norm():
T_multiply_red[v_ax0, v_ax1] * T.float32(0.05)
+ T.float32(1e-5)
)
- + C[v_ax2, v_ax3]
)
@R.function
def main(
x: R.Tensor((2, 3, 4, 5), dtype="float32"),
weight: R.Tensor((4, 5), dtype="float32"),
- bias: R.Tensor((4, 5), dtype="float32"),
) -> R.Tensor((2, 3, 4, 5), dtype="float32"):
cls = Expected
gv = R.call_tir(
cls.rms_norm,
- (x, weight, bias),
+ (x, weight),
out_sinfo=R.Tensor((2, 3, 4, 5), dtype="float32"),
)
return gv
@@ -2831,8 +2829,8 @@ def test_rms_norm_fp16():
@tvm.script.ir_module
class RMSNorm:
@R.function
- def main(x: R.Tensor((2, 3, 4, 5), "float16"), weight: R.Tensor((4,
5), "float16"), bias: R.Tensor((4, 5), "float16")) -> R.Tensor((2, 3, 4, 5),
"float16"):
- gv: R.Tensor((2, 3, 4, 5), "float16") = R.nn.rms_norm(x, weight,
bias, axes=[-2, -1])
+ def main(x: R.Tensor((2, 3, 4, 5), "float16"), weight: R.Tensor((4,
5), "float16")) -> R.Tensor((2, 3, 4, 5), "float16"):
+ gv: R.Tensor((2, 3, 4, 5), "float16") = R.nn.rms_norm(x, weight,
axes=[-2, -1])
return gv
@tvm.script.ir_module
@@ -2841,7 +2839,6 @@ def test_rms_norm_fp16():
def rms_norm(
A: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)),
"float16"),
B: T.Buffer((T.int64(4), T.int64(5)), "float16"),
- C: T.Buffer((T.int64(4), T.int64(5)), "float16"),
T_rms_norm: T.Buffer(
(T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float16"
),
@@ -2881,7 +2878,6 @@ def test_rms_norm_fp16():
A[v_ax0, v_ax1, v_ax2, v_ax3],
B[v_ax2, v_ax3],
T_multiply_red[v_ax0, v_ax1],
- C[v_ax2, v_ax3],
)
T.writes(T_rms_norm[v_ax0, v_ax1, v_ax2, v_ax3])
T_rms_norm[v_ax0, v_ax1, v_ax2, v_ax3] = (
@@ -2891,19 +2887,17 @@ def test_rms_norm_fp16():
T_multiply_red[v_ax0, v_ax1] / (T.float16(4) *
T.float16(5))
+ T.float16(1e-5)
)
- + C[v_ax2, v_ax3]
)
@R.function
def main(
x: R.Tensor((2, 3, 4, 5), dtype="float16"),
weight: R.Tensor((4, 5), dtype="float16"),
- bias: R.Tensor((4, 5), dtype="float16"),
) -> R.Tensor((2, 3, 4, 5), dtype="float16"):
cls = Expected
gv = R.call_tir(
cls.rms_norm,
- (x, weight, bias),
+ (x, weight),
out_sinfo=R.Tensor((2, 3, 4, 5), dtype="float16"),
)
return gv
@@ -2917,24 +2911,23 @@ def test_rms_norm_symbolic():
@tvm.script.ir_module
class RMSNorm:
@R.function
- def main(x: R.Tensor(("n", "s", "f"), "float32"), weight:
R.Tensor(("s", "f"), "float32"), bias: R.Tensor(("s", "f"), "float32")) ->
R.Tensor(("n", "s", "f"), "float32"):
+ def main(x: R.Tensor(("n", "s", "f"), "float32"), weight:
R.Tensor(("s", "f"), "float32")) -> R.Tensor(("n", "s", "f"), "float32"):
n = T.int64()
s = T.int64()
f = T.int64()
- gv: R.Tensor((n, s, f), "float32") = R.nn.rms_norm(x, weight,
bias, axes=[1, 2])
+ gv: R.Tensor((n, s, f), "float32") = R.nn.rms_norm(x, weight,
axes=[1, 2])
return gv
@tvm.script.ir_module
class Expected:
@T.prim_func
def rms_norm(
- var_A: T.handle, var_B: T.handle, var_C: T.handle, var_T_rms_norm:
T.handle
+ var_A: T.handle, var_B: T.handle, var_T_rms_norm: T.handle
):
T.func_attr({"tir.noalias": T.bool(True)})
n, s, f = T.int64(), T.int64(), T.int64()
A = T.match_buffer(var_A, (n, s, f))
B = T.match_buffer(var_B, (s, f))
- C = T.match_buffer(var_C, (s, f))
T_rms_norm = T.match_buffer(var_T_rms_norm, (n, s, f))
# with T.block("root"):
T_multiply = T.alloc_buffer((n, s, f))
@@ -2964,7 +2957,6 @@ def test_rms_norm_symbolic():
A[v_ax0, v_ax1, v_ax2],
B[v_ax1, v_ax2],
T_multiply_red[v_ax0],
- C[v_ax1, v_ax2],
)
T.writes(T_rms_norm[v_ax0, v_ax1, v_ax2])
T_rms_norm[v_ax0, v_ax1, v_ax2] = (
@@ -2975,14 +2967,12 @@ def test_rms_norm_symbolic():
/ (T.Cast("float32", s) * T.Cast("float32", f))
+ T.float32(1e-5)
)
- + C[v_ax1, v_ax2]
)
@R.function
def main(
x: R.Tensor(("n", "s", "f"), dtype="float32"),
weight: R.Tensor(("s", "f"), dtype="float32"),
- bias: R.Tensor(("s", "f"), dtype="float32"),
) -> R.Tensor(("n", "s", "f"), dtype="float32"):
n = T.int64()
s = T.int64()
@@ -2990,7 +2980,7 @@ def test_rms_norm_symbolic():
cls = Expected
gv = R.call_tir(
cls.rms_norm,
- (x, weight, bias),
+ (x, weight),
out_sinfo=R.Tensor((n, s, f), dtype="float32"),
)
return gv
@@ -2999,6 +2989,84 @@ def test_rms_norm_symbolic():
tvm.ir.assert_structural_equal(mod, Expected)
+def test_rms_norm_no_bias():
+ # fmt: off
+ @tvm.script.ir_module
+ class RMSNorm:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32"), weight: R.Tensor((4,
5), "float32")) -> R.Tensor((2, 3, 4, 5), "float32"):
+ gv: R.Tensor((2, 3, 4, 5), "float32") = R.nn.rms_norm(x, weight,
axes=[-2, -1])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @T.prim_func
+ def rms_norm(
+ A: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)),
"float32"),
+ B: T.Buffer((T.int64(4), T.int64(5)), "float32"),
+ T_rms_norm: T.Buffer(
+ (T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"
+ ),
+ ):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ T_multiply = T.alloc_buffer((T.int64(2), T.int64(3), T.int64(4),
T.int64(5)))
+ T_multiply_red = T.alloc_buffer((T.int64(2), T.int64(3)))
+ for ax0, ax1, ax2, ax3 in T.grid(
+ T.int64(2), T.int64(3), T.int64(4), T.int64(5)
+ ):
+ with T.block("T_multiply"):
+ v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0,
ax1, ax2, ax3])
+ T.reads(A[v_ax0, v_ax1, v_ax2, v_ax3])
+ T.writes(T_multiply[v_ax0, v_ax1, v_ax2, v_ax3])
+ T_multiply[v_ax0, v_ax1, v_ax2, v_ax3] = (
+ A[v_ax0, v_ax1, v_ax2, v_ax3] * A[v_ax0, v_ax1, v_ax2,
v_ax3]
+ )
+ for ax0, ax1, k2, k3 in T.grid(T.int64(2), T.int64(3), T.int64(4),
T.int64(5)):
+ with T.block("T_multiply_red"):
+ v_ax0, v_ax1, v_k2, v_k3 = T.axis.remap("SSRR", [ax0, ax1,
k2, k3])
+ T.reads(T_multiply[v_ax0, v_ax1, v_k2, v_k3])
+ T.writes(T_multiply_red[v_ax0, v_ax1])
+ with T.init():
+ T_multiply_red[v_ax0, v_ax1] = T.float32(0)
+ T_multiply_red[v_ax0, v_ax1] = (
+ T_multiply_red[v_ax0, v_ax1] + T_multiply[v_ax0,
v_ax1, v_k2, v_k3]
+ )
+ for ax0, ax1, ax2, ax3 in T.grid(
+ T.int64(2), T.int64(3), T.int64(4), T.int64(5)
+ ):
+ with T.block("T_rms_norm"):
+ v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0,
ax1, ax2, ax3])
+ T.reads(
+ A[v_ax0, v_ax1, v_ax2, v_ax3],
+ B[v_ax2, v_ax3],
+ T_multiply_red[v_ax0, v_ax1],
+ )
+ T.writes(T_rms_norm[v_ax0, v_ax1, v_ax2, v_ax3])
+ T_rms_norm[v_ax0, v_ax1, v_ax2, v_ax3] = (
+ A[v_ax0, v_ax1, v_ax2, v_ax3]
+ * B[v_ax2, v_ax3]
+ * T.rsqrt(
+ T_multiply_red[v_ax0, v_ax1] * T.float32(0.05)
+ + T.float32(1e-05)
+ )
+ )
+
+ @R.function
+ def main(
+ x: R.Tensor((2, 3, 4, 5), dtype="float32"),
+ weight: R.Tensor((4, 5), dtype="float32"),
+ ) -> R.Tensor((2, 3, 4, 5), dtype="float32"):
+ cls = Expected
+ gv = R.call_tir(
+ cls.rms_norm, (x, weight), out_sinfo=R.Tensor((2, 3, 4, 5),
dtype="float32")
+ )
+ return gv
+ # fmt: on
+ mod = LegalizeOps()(RMSNorm)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
def test_attention():
# fmt: off
@tvm.script.ir_module
