You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by bo...@apache.org on 2023/04/24 18:52:05 UTC
[tvm] branch unity updated: [Unity] Reduce cast to fp32 for constant input in AMP (#14679)
This is an automated email from the ASF dual-hosted git repository.
bohan 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 129fd8f84f [Unity] Reduce cast to fp32 for constant input in AMP (#14679)
129fd8f84f is described below
commit 129fd8f84f07de081f5098ae5f228d901a15f2c5
Author: masahi <ma...@gmail.com>
AuthorDate: Tue Apr 25 03:51:59 2023 +0900
[Unity] Reduce cast to fp32 for constant input in AMP (#14679)
* Check the range of constant arg to reduce cast to fp32
* Support attention in AMP
* mic fix: constant fold R.full and update cutlass check
---
python/tvm/relax/backend/contrib/cutlass.py | 5 +-
src/relax/op/nn/attention.cc | 8 ++
src/relax/transform/fold_constant.cc | 3 +-
src/relax/transform/to_mixed_precision.cc | 79 ++++++++++++++----
.../relax/test_transform_to_mixed_precision.py | 96 +++++++++++++++++++++-
5 files changed, 174 insertions(+), 17 deletions(-)
diff --git a/python/tvm/relax/backend/contrib/cutlass.py b/python/tvm/relax/backend/contrib/cutlass.py
index 7d6dc6bf89..7c22ec3dda 100644
--- a/python/tvm/relax/backend/contrib/cutlass.py
+++ b/python/tvm/relax/backend/contrib/cutlass.py
@@ -86,6 +86,9 @@ def _check_residual(root_call: Call, context: PatternCheckContext) -> bool:
if "residual" in context.annotated_expr:
residual = context.annotated_expr["residual"]
if not isinstance(residual, Var):
+ if not residual in context.value_to_bound_var:
+ return False
+
residual = context.value_to_bound_var[residual]
root_var = context.value_to_bound_var[root_call]
@@ -141,7 +144,7 @@ def _check_matmul(context: PatternCheckContext) -> bool:
if not _is_supported_dtype(lhs_dtype, rhs_dtype):
return False
- if not _check_residual(lhs, context):
+ if not _check_residual(context.annotated_expr["root"], context):
return False
lhs_shape = lhs.struct_info.shape.values
diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc
index c27e8b68d0..7f0a1d3aea 100644
--- a/src/relax/op/nn/attention.cc
+++ b/src/relax/op/nn/attention.cc
@@ -108,12 +108,18 @@ StructInfo InferStructInfoAttention(const Call& call, const BlockBuilder& ctx) {
return TensorStructInfo(ShapeExpr(output_shape), q_sinfo->dtype);
}
+Call InferMixedPrecisionAttention(const Call& call, const DataType& out_dtype) {
+ return Downcast<Call>(attention(call->args[0], call->args[1], call->args[2], NullOpt, NullOpt));
+}
+
TVM_REGISTER_OP("relax.nn.attention")
.set_attrs_type<AttentionAttrs>()
.set_num_inputs(3)
.add_argument("query", "Tensor", "The input queries tensor.")
.add_argument("key", "Tensor", "The input keys tensor.")
.add_argument("value", "Tensor", "The input values tensor.")
+ .set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
+ .set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionAttention)
.set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoAttention);
TVM_REGISTER_OP("relax.nn.attention_bias")
@@ -123,6 +129,8 @@ TVM_REGISTER_OP("relax.nn.attention_bias")
.add_argument("key", "Tensor", "The input keys tensor.")
.add_argument("value", "Tensor", "The input values tensor.")
.add_argument("bias", "Tensor", "The input bias tensor.")
+ .set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
+ .set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionAttention)
.set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoAttention);
} // namespace relax
diff --git a/src/relax/transform/fold_constant.cc b/src/relax/transform/fold_constant.cc
index 539f20dd9d..a13b7f3d93 100644
--- a/src/relax/transform/fold_constant.cc
+++ b/src/relax/transform/fold_constant.cc
@@ -248,7 +248,8 @@ class ConstantFolder : public ExprMutator {
// Check if we can them to call_tir
if (legalize_map.count(op)) {
// Get the legalized expression
- Expr legalized_expr = builder_->Normalize(legalize_map[op](builder_, post_call));
+ Call post_call_normalized = Downcast<Call>(builder_->Normalize(post_call));
+ Expr legalized_expr = builder_->Normalize(legalize_map[op](builder_, post_call_normalized));
// If the legalized expression is call_tir, try to fold it.
const CallNode* call = legalized_expr.as<CallNode>();
if (call && call->op.same_as(call_tir_op)) {
diff --git a/src/relax/transform/to_mixed_precision.cc b/src/relax/transform/to_mixed_precision.cc
index e74ca7a4d8..a04d5dbd3a 100644
--- a/src/relax/transform/to_mixed_precision.cc
+++ b/src/relax/transform/to_mixed_precision.cc
@@ -26,6 +26,7 @@
#include <tvm/relax/transform.h>
#include <array>
+#include <cstdint>
#include "../op/nn/convolution.h"
#include "../op/tensor/datatype.h"
@@ -318,25 +319,75 @@ class ToMixedPrecisionRewriter : public ExprMutator {
return new_args;
}
- // Util function to check if any of the tensors in the args is fp32
- bool AnyArgIsFP32(const NType& cur_type) {
- bool result = false;
- auto fvisitleaf = [&](const String& dtype) {
- if (dtype == "float32") {
- result = true;
+ template <typename T>
+ bool CheckInFP16Range(const std::vector<uint8_t>& bytes, size_t num_elem) {
+ const T* ptr = reinterpret_cast<const T*>(bytes.data());
+ constexpr float fp16_max = 65504;
+ for (size_t i = 0; i < num_elem; ++i) {
+ if (std::abs(ptr[i]) > fp16_max) return false;
+ }
+ return true;
+ }
+
+ bool AllFP16Castable(const Array<Expr>& args) {
+ auto is_fp16 = [](StructInfo sinfo) {
+ if (auto tensor_sinfo = sinfo.as<TensorStructInfoNode>();
+ tensor_sinfo && tensor_sinfo->dtype == DataType::Float(16)) {
+ return true;
}
+ return false;
+ };
+
+ auto is_in_fp16_range = [this](const ConstantNode* constant) {
+ const auto& data = constant->data;
+ if (data->dtype.lanes > 1) {
+ // Skip vectorized types for now.
+ return false;
+ }
+
+ if (data.DataType() == DataType::Float(16)) {
+ return true;
+ }
+
+ auto shape = data.Shape();
+ size_t size_1d = 1;
+ for (size_t i = 0; i < shape.size(); ++i) {
+ size_1d *= shape[i];
+ }
+ auto elem_bytes = data->dtype.bits / 8;
+ std::vector<uint8_t> bytes(size_1d * elem_bytes);
+ data.CopyToBytes(bytes.data(), bytes.size());
+
+ if (data.DataType() == DataType::Float(32)) {
+ return CheckInFP16Range<float>(bytes, size_1d);
+ } else if (data.DataType() == DataType::Float(64)) {
+ return CheckInFP16Range<double>(bytes, size_1d);
+ } else if (data.DataType() == DataType::Int(8)) {
+ return CheckInFP16Range<std::int8_t>(bytes, size_1d);
+ } else if (data.DataType() == DataType::Int(16)) {
+ return CheckInFP16Range<std::int16_t>(bytes, size_1d);
+ } else if (data.DataType() == DataType::Int(32)) {
+ return CheckInFP16Range<std::int32_t>(bytes, size_1d);
+ } else if (data.DataType() == DataType::Int(64)) {
+ return CheckInFP16Range<std::int64_t>(bytes, size_1d);
+ }
+ return false;
};
- ForEachLeaf<String>(cur_type, fvisitleaf);
- return result;
- }
- bool AnyArgIsFP32(const Array<Expr>& args) {
for (const Expr& arg : args) {
- if (IsNestedTensor(arg)) {
- if (AnyArgIsFP32(NTypeFrom(arg))) return true;
+ auto sinfo = GetStructInfo(arg);
+ auto constant = arg.as<ConstantNode>();
+ auto tuple = arg.as<TupleNode>();
+
+ if (!IsNestedTensor(arg) || is_fp16(sinfo) || (constant && is_in_fp16_range(constant)) ||
+ (tuple && AllFP16Castable(tuple->fields))) {
+ continue;
+ } else {
+ return false;
}
}
- return false;
+
+ return true;
}
void CastIfFp16Only(const Var& var) {
@@ -421,7 +472,7 @@ class ToMixedPrecisionRewriter : public ExprMutator {
auto f = attr_map[op];
new_call = make_object<CallNode>(*(f(Call(new_call), output_dtype_).get()));
} else if (policy == kFollow) {
- to = AnyArgIsFP32(new_call->args) ? fp32_ : fp16_;
+ to = AllFP16Castable(new_call->args) ? fp16_ : fp32_;
} else if (policy == kNever) {
to = fp32_;
} else {
diff --git a/tests/python/relax/test_transform_to_mixed_precision.py b/tests/python/relax/test_transform_to_mixed_precision.py
index 721cbd9d58..6bae732927 100644
--- a/tests/python/relax/test_transform_to_mixed_precision.py
+++ b/tests/python/relax/test_transform_to_mixed_precision.py
@@ -30,7 +30,6 @@ def _assert_test(input, expected=None, expected2=None):
if expected2:
mod = ToMixedPrecision(out_dtype="float16")(input)
- print(mod.script())
tvm.ir.assert_structural_equal(mod, expected2)
@@ -896,5 +895,100 @@ def test_tuple_get():
_assert_test(Module, expected2=Expected)
+def test_conv2d_bias_fp32():
+ @tvm.script.ir_module
+ class Input:
+ @R.function
+ def main(
+ x: R.Tensor((1, 4, 64, 64), dtype="float32"),
+ w: R.Tensor((512, 4, 3, 3), dtype="float32"),
+ bias: R.Tensor((512,), dtype="float32"),
+ ) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
+ # block 0
+ with R.dataflow():
+ lv142: R.Tensor((1, 4, 64, 64), dtype="float32") = R.nn.conv2d(
+ x,
+ w,
+ strides=[1, 1],
+ padding=[0, 0, 0, 0],
+ out_dtype="float32",
+ )
+ lv143: R.Tensor((1, 4, 1, 1), dtype="float32") = R.reshape(bias, (1, 512, 1, 1))
+ lv144: R.Tensor((1, 4, 64, 64), dtype="float32") = R.add(lv142, lv143)
+ R.output(lv144)
+ return lv144
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((1, 4, 64, 64), dtype="float32"),
+ w: R.Tensor((512, 4, 3, 3), dtype="float32"),
+ bias: R.Tensor((512,), dtype="float32"),
+ ) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((1, 4, 64, 64), dtype="float16") = R.astype(x, dtype="float16")
+ lv1: R.Tensor((512, 4, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
+ lv142: R.Tensor((1, 512, 62, 62), dtype="float16") = R.nn.conv2d(
+ lv,
+ lv1,
+ strides=[1, 1],
+ padding=[0, 0, 0, 0],
+ out_dtype="float16",
+ )
+ lv2: R.Tensor((512,), dtype="float16") = R.astype(bias, dtype="float16")
+ lv143: R.Tensor((1, 512, 1, 1), dtype="float16") = R.reshape(
+ lv2, R.shape([1, 512, 1, 1])
+ )
+ lv3: R.Tensor((1, 512, 62, 62), dtype="float16") = R.add(lv142, lv143)
+ lv144: R.Tensor((1, 512, 62, 62), dtype="float32") = R.astype(lv3, dtype="float32")
+ R.output(lv144)
+ return lv144
+
+ @tvm.script.ir_module
+ class Expected_no_bias_cast:
+ @R.function
+ def main(
+ x: R.Tensor((1, 4, 64, 64), dtype="float32"),
+ w: R.Tensor((512, 4, 3, 3), dtype="float32"),
+ bias: R.Tensor((512,), dtype="float32"),
+ ) -> R.Tensor((1, 512, 64, 64), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((1, 4, 64, 64), dtype="float16") = R.astype(x, dtype="float16")
+ lv1: R.Tensor((512, 4, 3, 3), dtype="float16") = R.astype(w, dtype="float16")
+ lv142: R.Tensor((1, 512, 62, 62), dtype="float16") = R.nn.conv2d(
+ lv,
+ lv1,
+ strides=[1, 1],
+ padding=[0, 0, 0, 0],
+ out_dtype="float16",
+ )
+ lv143: R.Tensor((1, 512, 1, 1), dtype="float32") = R.reshape(
+ bias, R.shape([1, 512, 1, 1])
+ )
+ lv2: R.Tensor((1, 512, 62, 62), dtype="float32") = R.astype(lv142, dtype="float32")
+ lv144: R.Tensor((1, 512, 62, 62), dtype="float32") = R.add(lv2, lv143)
+ R.output(lv144)
+ return lv144
+
+ binding_np = {
+ "w": np.random.uniform(size=(512, 4, 3, 3)).astype("float32"),
+ "bias": np.random.uniform(size=(512,)).astype("float32"),
+ }
+ binding = {k: tvm.nd.array(v) for k, v in binding_np.items()}
+
+ Input_bound = relax.transform.BindParams("main", binding)(Input)
+ Expected = relax.transform.BindParams("main", binding)(Expected)
+
+ _assert_test(Input_bound, expected2=Expected)
+
+ binding_np["bias"][0] = 70000 # Out of fp16 range
+ binding = {k: tvm.nd.array(v) for k, v in binding_np.items()}
+ Input_bound = relax.transform.BindParams("main", binding)(Input)
+ Expected_no_bias_cast = relax.transform.BindParams("main", binding)(Expected_no_bias_cast)
+
+ _assert_test(Input_bound, expected2=Expected_no_bias_cast)
+
+
if __name__ == "__main__":
tvm.testing.main()