apply code review suggestions

Author: cyx-6

include/tvm/topi/nn/rms_norm.h

@@ -41,25 +41,24 @@ 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 Array<Integer>& axis,
-                       double epsilon, std::string name = "T_rms_norm",
+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",
                        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";
-  ICHECK(data_type == DataType::Float(32) || data_type == DataType::Float(16))
-      << "rms_norm: only support float32 and float16 for now";
-  bool is_float16 = data_type == DataType::Float(16);
+  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 x = is_float16 ? cast(data, DataType::Float(32)) : data;
-  auto w = is_float16 ? cast(weight, DataType::Float(32)) : weight;
-  auto square = multiply(x, x);
+  auto square = multiply(data, data);
   auto square_sum = sum(square, axis, /*keepdims=*/false, /*atleast1d=*/true);
 
   auto ndim = data->shape.size();
@@ -79,12 +78,15 @@ inline Tensor rms_norm(const Tensor& data, const Tensor& weight, const Array<Int
       }
     }
     auto output =
-        x(indices) * w(reduce_indices) *
+        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);
-  return is_float16 ? cast(rms_norm, DataType::Float(16)) : rms_norm;
+  return rms_norm;
 }
 
 }  // namespace nn

python/tvm/topi/nn/rms_norm.py

@@ -18,10 +18,8 @@
 from .. import cpp
 
 
-def rms_norm(data, weight, axis, epsilon=1e-5):
-    """Root mean square normalization operator.
-    It accepts fp16 and fp32 as input data type. It will cast the input to fp32
-    to perform the computation. The output will have the same data type as input.
+def rms_norm(data, weight, bias, axis, epsilon=1e-5):
+    """Root mean square normalization operator. The output will have the same data type as input.
 
     Parameters
     ----------
@@ -31,6 +29,9 @@ def rms_norm(data, weight, 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
 
@@ -42,4 +43,4 @@ def rms_norm(data, weight, 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, axis, epsilon)
+    return cpp.nn.rms_norm(data, weight, bias, axis, epsilon)

python/tvm/topi/testing/rms_norm_python.py

@@ -19,7 +19,7 @@
 import numpy as np
 
 
-def rms_norm_python(data, weight, axis, epsilon=1e-5):
+def rms_norm_python(data, weight, bias, axis, epsilon=1e-5):
     """Root mean square normalization operator in Python.
 
     Parameters
@@ -30,6 +30,9 @@ def rms_norm_python(data, weight, axis, epsilon=1e-5):
     weight: numpy.ndarray
         K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
 
+    bias: numpy.ndarray
+        Optional, K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
+
     axis : int or tuple of ints
         Axis over the normalization applied
 
@@ -41,8 +44,8 @@ def rms_norm_python(data, weight, axis, epsilon=1e-5):
     result : np.ndarray
         N-D with shape (d_0, d_1, ..., d_{N-1})
     """
-    old_dtype = data.dtype
-    data = data.astype("float32")
     square_mean = np.mean(np.square(data), axis, keepdims=True)
     result = data * weight / np.sqrt(square_mean + epsilon)
-    return result.astype(old_dtype)
+    if bias is not None:
+        result += bias
+    return result

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], static_cast<double>(args[3]));
+  *rv = nn::rms_norm(args[0], args[1], args[2], args[3], static_cast<double>(args[4]));
 });
 
 }  // namespace topi

tests/python/topi/python/test_topi_rms_norm.py

@@ -35,24 +35,27 @@
 @tvm.testing.parametrize_targets("llvm")
 @pytest.mark.parametrize("shape,axis", [([4, 16], (1,)), ([4, 16, 16], (1, 2))])
 @pytest.mark.parametrize("dtype", ["float32", "float16"])
-def test_rms_norm(target, dev, shape, axis, dtype, episilon=1e-5, rtol=5e-4, atol=5e-4):
+def test_rms_norm(target, dev, shape, axis, dtype, episilon=1e-5, rtol=5e-3, atol=1e-4):
     data = te.placeholder(shape, dtype=dtype, name="data")
     scale_shape = [shape[dim] for dim in axis]
     weight = te.placeholder(scale_shape, dtype=dtype, name="weight")
-    B = topi.nn.rms_norm(data, weight, axis, episilon)
+    bias = te.placeholder(scale_shape, dtype=dtype, name="weight")
+    B = topi.nn.rms_norm(data, weight, bias, axis, episilon)
 
     data_np = np.random.uniform(size=shape).astype(dtype)
     weight_np = np.random.uniform(size=scale_shape).astype(dtype)
-    b_np = tvm.topi.testing.rms_norm_python(data_np, weight_np, axis, episilon)
+    bias_np = np.random.uniform(size=scale_shape).astype(dtype)
+    b_np = tvm.topi.testing.rms_norm_python(data_np, weight_np, bias_np, axis, episilon)
 
     with tvm.target.Target(target):
         s_func = tvm.topi.testing.dispatch(target, _rms_norm_schedule)
         s = s_func([B])
     data_tvm = tvm.nd.array(data_np, dev)
     weight_tvm = tvm.nd.array(weight_np, dev)
+    bias_tvm = tvm.nd.array(bias_np, dev)
     b_tvm = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=dtype), dev)
-    f = tvm.build(s, [data, weight, B], target)
-    f(data_tvm, weight_tvm, b_tvm)
+    f = tvm.build(s, [data, weight, bias, B], target)
+    f(data_tvm, weight_tvm, bias_tvm, b_tvm)
     tvm.testing.assert_allclose(b_tvm.numpy(), b_np, rtol=rtol, atol=atol)