You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ma...@apache.org on 2023/11/20 09:45:33 UTC
(tvm) branch unity updated: [Unity][MSC] Enable add attributes while fuse ops (#16128)
This is an automated email from the ASF dual-hosted git repository.
masahi 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 29450b927e [Unity][MSC] Enable add attributes while fuse ops (#16128)
29450b927e is described below
commit 29450b927e1b6379b849419705ae3931ce34605f
Author: Archermmt <ar...@126.com>
AuthorDate: Mon Nov 20 17:45:25 2023 +0800
[Unity][MSC] Enable add attributes while fuse ops (#16128)
* enable attrs_getter
* update test case for attrs_getter
* add name checker for test tensorrt
* remove pattern.py
* add unique name for tuple
* change attrs_getter argument
* change attrs getter signature
---
include/tvm/relax/transform.h | 14 ++-
python/tvm/contrib/msc/core/frontend/translate.py | 25 ++---
python/tvm/contrib/msc/core/transform/pattern.py | 115 +++++++++++++++------
python/tvm/contrib/msc/core/utils/namespace.py | 2 +
.../msc/framework/tensorrt/transform/pattern.py | 83 +++++++++++++--
python/tvm/relax/transform/transform.py | 8 +-
src/contrib/msc/core/transform/fuse_tuple.cc | 2 +
src/relax/transform/fuse_ops.cc | 37 ++++---
.../contrib/test_msc/test_translate_tensorrt.py | 36 +++++++
9 files changed, 247 insertions(+), 75 deletions(-)
diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h
index b4b9a3d78a..f99d74b7aa 100644
--- a/include/tvm/relax/transform.h
+++ b/include/tvm/relax/transform.h
@@ -322,11 +322,20 @@ class FusionPatternNode : public Object {
*/
Optional<PackedFunc> check;
+ /*!
+ * \brief The function to get attributes for fused function
+ *
+ * It should have signature
+ * Map<String, String>(const Map<String, Expr>& context)
+ */
+ Optional<PackedFunc> attrs_getter;
+
void VisitAttrs(tvm::AttrVisitor* v) {
v->Visit("name", &name);
v->Visit("pattern", &pattern);
v->Visit("annotation_patterns", &annotation_patterns);
v->Visit("check", &check);
+ v->Visit("attrs_getter", &attrs_getter);
}
static constexpr const char* _type_key = "relax.transform.FusionPattern";
@@ -336,9 +345,10 @@ class FusionPatternNode : public Object {
class FusionPattern : public ObjectRef {
public:
FusionPattern(String name, DFPattern pattern, Map<String, DFPattern> annotation_patterns,
- Optional<PackedFunc> check);
+ Optional<PackedFunc> check, Optional<PackedFunc> attrs_getter);
- FusionPattern(String name, DFPattern pattern) : FusionPattern(name, pattern, {}, NullOpt) {}
+ FusionPattern(String name, DFPattern pattern)
+ : FusionPattern(name, pattern, {}, NullOpt, NullOpt) {}
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(FusionPattern, ObjectRef, FusionPatternNode);
};
diff --git a/python/tvm/contrib/msc/core/frontend/translate.py b/python/tvm/contrib/msc/core/frontend/translate.py
index 1e7efe6489..4a7710f382 100644
--- a/python/tvm/contrib/msc/core/frontend/translate.py
+++ b/python/tvm/contrib/msc/core/frontend/translate.py
@@ -114,6 +114,7 @@ def from_relax(
)(mod)
patterns = get_patterns_with_prefix("msc.")
passes = [
+ msc_transform.SetExprName(),
tvm.relax.transform.FuseOpsByPattern(
patterns, bind_constants=False, annotate_codegen=False
),
@@ -306,20 +307,16 @@ def byoc_partition(
def _partition_mod(mod, as_msc=True):
patterns = get_patterns_with_prefix(target)
- if as_msc:
- passes = [tvm.relax.transform.FuseOpsByPattern(patterns, bind_constants=False)]
- else:
- passes = [tvm.relax.transform.FuseOpsByPattern(patterns, bind_constants=True)]
- passes.extend(
- [
- msc_transform.BindShape(),
- msc_transform.FuseTuple(target),
- tvm.relax.transform.MergeCompositeFunctions(),
- msc_transform.SetBYOCAttrs(target),
- msc_transform.SetExprName(target=target),
- msc_transform.SetExprLayout(trans_config.get("allow_layout_missing", True)),
- ]
- )
+ passes = [
+ msc_transform.SetExprName(),
+ tvm.relax.transform.FuseOpsByPattern(patterns, bind_constants=not as_msc),
+ msc_transform.BindShape(),
+ msc_transform.FuseTuple(target),
+ tvm.relax.transform.MergeCompositeFunctions(),
+ msc_transform.SetBYOCAttrs(target),
+ msc_transform.SetExprName(target=target),
+ msc_transform.SetExprLayout(trans_config.get("allow_layout_missing", True)),
+ ]
return tvm.transform.Sequential(passes)(mod)
def _is_target_func(func):
diff --git a/python/tvm/contrib/msc/core/transform/pattern.py b/python/tvm/contrib/msc/core/transform/pattern.py
index 08e054ac28..6d7ac5ba13 100644
--- a/python/tvm/contrib/msc/core/transform/pattern.py
+++ b/python/tvm/contrib/msc/core/transform/pattern.py
@@ -17,7 +17,8 @@
# pylint: disable=unused-argument
"""tvm.contrib.msc.core.transform.pattern"""
-from typing import Mapping, Tuple
+from typing import Mapping, Tuple, Dict
+from functools import partial
import tvm
from tvm.relax.dpl import pattern as relax_pattern
@@ -26,6 +27,36 @@ from tvm.relay import dataflow_pattern as relay_pattern
from tvm.relax.transform import PatternCheckContext
from tvm.relax.backend.pattern_registry import register_patterns
from tvm.relay.op.contrib.register import register_pattern_table
+from tvm.contrib.msc.core.utils.namespace import MSCMap, MSCKey
+from tvm.contrib.msc.core import utils as msc_utils
+
+
+def msc_attrs_getter(
+ annotated_expr: Dict[str, tvm.relax.Expr], anchor: str = "out"
+) -> Dict[str, str]:
+ """Get attributes for fused pattern
+
+ Parameters
+ ----------
+ annotated_expr: dict<str,Expr>
+ The annotated exprs during fus pattern
+ anchor: str
+ The anchor key of expr
+
+ Returns
+ -------
+ attrs: dict<str,str>
+ The extra attributes for msc.
+ """
+
+ fused_cnt = MSCMap.get(MSCKey.FUSED_CNT, 0)
+ unique_name = "msc_fused_" + str(fused_cnt)
+ if anchor in annotated_expr:
+ name = msc_utils.get_expr_name(annotated_expr[anchor])
+ if name:
+ unique_name = name
+ MSCMap.set(MSCKey.FUSED_CNT, fused_cnt + 1)
+ return {"unique_name": unique_name}
def make_relax_conv_bias_pattern(
@@ -56,7 +87,7 @@ def make_relax_conv_bias_pattern(
shape = relax_pattern.wildcard()
reshape = relax_pattern.is_op("relax.reshape")(bias, shape)
out = relax_pattern.is_op("relax.add")(conv, reshape)
- annotations = {"bias": bias, "reshape": reshape}
+ annotations = {"conv": conv, "bias": bias, "reshape": reshape, "out": out}
return out, annotations
@@ -75,9 +106,9 @@ def _check_relax_conv_bias(context: PatternCheckContext) -> bool:
return non_one_dims <= 1 and bias.struct_info.ndim == 1
-def make_relax_linear_pattern() -> Tuple[
- relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]
-]:
+def make_relax_linear_pattern() -> (
+ Tuple[relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]]
+):
"""A simple utility to create patterns for linear.
Returns
@@ -95,7 +126,7 @@ def make_relax_linear_pattern() -> Tuple[
weight = relax_pattern.is_const()
permute = relax_pattern.is_op("relax.permute_dims")(weight)
out = relax_pattern.is_op("relax.matmul")(data, permute)
- annotations = {"weight": weight, "permute": permute}
+ annotations = {"weight": weight, "permute": permute, "linear": out}
return out, annotations
@@ -113,9 +144,9 @@ def _check_relax_linear(context: PatternCheckContext) -> bool:
return weight.struct_info.ndim == 2 and not permute.attrs["axes"]
-def make_relax_linear_bias_pattern() -> Tuple[
- relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]
-]:
+def make_relax_linear_bias_pattern() -> (
+ Tuple[relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]]
+):
"""A simple utility to create patterns for linear with bias.
Returns
@@ -133,7 +164,7 @@ def make_relax_linear_bias_pattern() -> Tuple[
linear, annotations = make_relax_linear_pattern()
bias = relax_pattern.is_const()
out = relax_pattern.is_op("relax.add")(linear, bias)
- annotations.update({"bias": bias, "out": out})
+ annotations.update({"bias": bias, "bias_add": out})
return out, annotations
@@ -152,9 +183,9 @@ def _check_relax_linear_bias(context: PatternCheckContext) -> bool:
return bias.struct_info.ndim == 1
-def make_relax_embedding_pattern() -> Tuple[
- relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]
-]:
+def make_relax_embedding_pattern() -> (
+ Tuple[relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]]
+):
"""A simple utility to create patterns for embedding.
Returns
@@ -172,7 +203,7 @@ def make_relax_embedding_pattern() -> Tuple[
data = relax_pattern.wildcard()
astype = relax_pattern.is_op("relax.astype")(data)
out = relax_pattern.is_op("relax.take")(weight, astype)
- annotations = {"weight": weight, "astype": astype}
+ annotations = {"weight": weight, "astype": astype, "take": out}
return out, annotations
@@ -194,9 +225,9 @@ def _check_relax_embedding(context: PatternCheckContext) -> bool:
)
-def make_relax_reshape_embedding_pattern() -> Tuple[
- relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]
-]:
+def make_relax_reshape_embedding_pattern() -> (
+ Tuple[relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]]
+):
"""A simple utility to create patterns for reshaped embedding.
Returns
@@ -218,7 +249,13 @@ def make_relax_reshape_embedding_pattern() -> Tuple[
take = relax_pattern.is_op("relax.take")(weight, reduce_in)
expand_shape = relax_pattern.wildcard()
out = relax_pattern.is_op("relax.reshape")(take, expand_shape)
- annotations = {"weight": weight, "astype": astype, "reduce_in": reduce_in}
+ annotations = {
+ "weight": weight,
+ "astype": astype,
+ "reduce_in": reduce_in,
+ "take": take,
+ "out": out,
+ }
return out, annotations
@@ -241,9 +278,9 @@ def _check_relax_reshape_embedding(context: PatternCheckContext) -> bool:
return True
-def make_relax_attention_pattern() -> Tuple[
- relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]
-]:
+def make_relax_attention_pattern() -> (
+ Tuple[relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]]
+):
"""A simple utility to create patterns for attention.
Returns
@@ -264,7 +301,7 @@ def make_relax_attention_pattern() -> Tuple[
k_trans = relax_pattern.is_op("relax.permute_dims")(weight_k)
v_trans = relax_pattern.is_op("relax.permute_dims")(weight_v)
out = relax_pattern.is_op("relax.nn.attention")(q_trans, k_trans, v_trans)
- annotations = {"q_trans": q_trans, "k_trans": k_trans, "v_trans": v_trans}
+ annotations = {"q_trans": q_trans, "k_trans": k_trans, "v_trans": v_trans, "attention": out}
return out, annotations
@@ -280,9 +317,9 @@ def _check_relax_attention(context: PatternCheckContext) -> bool:
return True
-def make_relax_mask_attention_pattern() -> Tuple[
- relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]
-]:
+def make_relax_mask_attention_pattern() -> (
+ Tuple[relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]]
+):
"""A simple utility to create patterns for mask_attention.
Returns
@@ -304,7 +341,7 @@ def make_relax_mask_attention_pattern() -> Tuple[
k_trans = relax_pattern.is_op("relax.permute_dims")(weight_k)
v_trans = relax_pattern.is_op("relax.permute_dims")(weight_v)
out = relax_pattern.is_op("relax.nn.attention_bias")(q_trans, k_trans, v_trans, mask)
- annotations = {"q_trans": q_trans, "k_trans": k_trans, "v_trans": v_trans}
+ annotations = {"q_trans": q_trans, "k_trans": k_trans, "v_trans": v_trans, "attention": out}
return out, annotations
@@ -365,9 +402,9 @@ def _check_opt_relax_conv_bias(context: PatternCheckContext) -> bool:
return ndim_conv == ndim_bias and ndim_bias == ndim_out
-def make_opt_relax_linear_pattern() -> Tuple[
- relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]
-]:
+def make_opt_relax_linear_pattern() -> (
+ Tuple[relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]]
+):
"""Create patterns for an linear, for mod after optimize.
Returns
@@ -384,7 +421,7 @@ def make_opt_relax_linear_pattern() -> Tuple[
data = relax_pattern.wildcard()
weight = relax_pattern.is_const()
out = relax_pattern.is_op("relax.matmul")(data, weight)
- annotations = {"weight": weight}
+ annotations = {"weight": weight, "linear": out}
return out, annotations
@@ -401,9 +438,9 @@ def _check_opt_relax_linear(context: PatternCheckContext) -> bool:
return ndim_weight == 2
-def make_opt_relax_linear_bias_pattern() -> Tuple[
- relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]
-]:
+def make_opt_relax_linear_bias_pattern() -> (
+ Tuple[relax_pattern.DFPattern, Mapping[str, relax_pattern.DFPattern]]
+):
"""Create patterns for an linear_bias, for mod after optimize.
Returns
@@ -451,6 +488,7 @@ register_patterns(
"relax.nn.conv1d",
),
_check_opt_relax_conv_bias,
+ partial(msc_attrs_getter, anchor="conv"),
),
(
"msc.conv2d_bias",
@@ -458,16 +496,19 @@ register_patterns(
"relax.nn.conv2d",
),
_check_opt_relax_conv_bias,
+ partial(msc_attrs_getter, anchor="conv"),
),
(
"msc.linear",
*make_opt_relax_linear_pattern(),
_check_opt_relax_linear,
+ partial(msc_attrs_getter, anchor="linear"),
),
(
"msc.linear_bias",
*make_opt_relax_linear_bias_pattern(),
_check_opt_relax_linear_bias,
+ partial(msc_attrs_getter, anchor="linear"),
),
(
"msc.conv1d_bias",
@@ -475,6 +516,7 @@ register_patterns(
"relax.nn.conv1d",
),
_check_relax_conv_bias,
+ partial(msc_attrs_getter, anchor="conv"),
),
(
"msc.conv2d_bias",
@@ -482,36 +524,43 @@ register_patterns(
"relax.nn.conv2d",
),
_check_relax_conv_bias,
+ partial(msc_attrs_getter, anchor="conv"),
),
(
"msc.linear",
*make_relax_linear_pattern(),
_check_relax_linear,
+ partial(msc_attrs_getter, anchor="linear"),
),
(
"msc.linear_bias",
*make_relax_linear_bias_pattern(),
_check_relax_linear_bias,
+ partial(msc_attrs_getter, anchor="linear"),
),
(
"msc.embedding",
*make_relax_embedding_pattern(),
_check_relax_embedding,
+ partial(msc_attrs_getter, anchor="take"),
),
(
"msc.embedding",
*make_relax_reshape_embedding_pattern(),
_check_relax_reshape_embedding,
+ partial(msc_attrs_getter, anchor="take"),
),
(
"msc.attention",
*make_relax_attention_pattern(),
_check_relax_attention,
+ partial(msc_attrs_getter, anchor="attention"),
),
(
"msc.attention",
*make_relax_mask_attention_pattern(),
_check_relax_mask_attention,
+ partial(msc_attrs_getter, anchor="attention"),
),
]
)
diff --git a/python/tvm/contrib/msc/core/utils/namespace.py b/python/tvm/contrib/msc/core/utils/namespace.py
index dcf20bef1b..63a4365c54 100644
--- a/python/tvm/contrib/msc/core/utils/namespace.py
+++ b/python/tvm/contrib/msc/core/utils/namespace.py
@@ -60,6 +60,8 @@ class MSCKey:
MSC_STAGE = "msc_stage"
TIME_STAMPS = "time_stamps"
+ FUSED_CNT = "fused_cnt"
+
class MSCFramework:
"""Framework type for the MSC"""
diff --git a/python/tvm/contrib/msc/framework/tensorrt/transform/pattern.py b/python/tvm/contrib/msc/framework/tensorrt/transform/pattern.py
index fe42ee2ae7..9fa1cf2142 100644
--- a/python/tvm/contrib/msc/framework/tensorrt/transform/pattern.py
+++ b/python/tvm/contrib/msc/framework/tensorrt/transform/pattern.py
@@ -17,8 +17,8 @@
# pylint: disable=unused-argument
"""tvm.contrib.msc.framework.tensorrt.transform.pattern"""
-from typing import Mapping, Tuple, List, Union, Callable
-from functools import wraps
+from typing import Mapping, Tuple, List, Union, Callable, Dict
+from functools import wraps, partial
from tvm import relax
from tvm.relax.dpl import pattern
@@ -255,11 +255,13 @@ def wrap_basic_check(
CheckFunc = Callable[[Mapping[pattern.DFPattern, relax.Expr], relax.Expr], bool]
+GetterFunc = Callable[[Mapping[pattern.DFPattern, relax.Expr], relax.Expr], Dict[str, str]]
Pattern = Union[
FusionPattern,
Tuple[str, pattern.DFPattern],
Tuple[str, pattern.DFPattern, Mapping[str, pattern.DFPattern]],
Tuple[str, pattern.DFPattern, Mapping[str, pattern.DFPattern], CheckFunc],
+ Tuple[str, pattern.DFPattern, Mapping[str, pattern.DFPattern], CheckFunc, GetterFunc],
]
@@ -306,33 +308,91 @@ def get_patterns(target) -> List[Pattern]:
patterns = []
# basic ops
for op, in_types in basic_ops.items():
- patterns.append((target + "." + op, *basic_pattern("relax." + op, in_types), _basic_check))
+ patterns.append(
+ (
+ target + "." + op,
+ *basic_pattern("relax." + op, in_types),
+ _basic_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
+ )
+ )
# activation ops
for op in activation_ops:
- patterns.append((target + "." + op, *basic_pattern("relax." + op, ["input"]), _basic_check))
+ patterns.append(
+ (
+ target + "." + op,
+ *basic_pattern("relax." + op, ["input"]),
+ _basic_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
+ )
+ )
# reduce ops
for op in reduce_ops:
- patterns.append((target + "." + op, *basic_pattern("relax." + op, ["input"]), _basic_check))
+ patterns.append(
+ (
+ target + "." + op,
+ *basic_pattern("relax." + op, ["input"]),
+ _basic_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
+ )
+ )
# unary ops
for op in unary_ops:
- patterns.append((target + "." + op, *basic_pattern("relax." + op, ["input"]), _basic_check))
+ patterns.append(
+ (
+ target + "." + op,
+ *basic_pattern("relax." + op, ["input"]),
+ _basic_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
+ )
+ )
# elemwise ops
for op in elemwise_ops:
- patterns.append((target + "." + op, *elemwise_pattern("relax." + op), _elemwise_check))
+ patterns.append(
+ (
+ target + "." + op,
+ *elemwise_pattern("relax." + op),
+ _elemwise_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
+ )
+ )
# compare ops
for op in compare_ops:
- patterns.append((target + "." + op, *elemwise_pattern("relax." + op), _compare_check))
+ patterns.append(
+ (
+ target + "." + op,
+ *elemwise_pattern("relax." + op),
+ _compare_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
+ )
+ )
# special ops
patterns.extend(
[
- (target + ".take", *basic_pattern("relax.take", ["input", "input"]), _take_check),
- (target + ".argmax", *argmaxmin_pattern("relax.argmax"), _argmaxmin_check),
- (target + ".argmin", *argmaxmin_pattern("relax.argmin"), _argmaxmin_check),
+ (
+ target + ".take",
+ *basic_pattern("relax.take", ["input", "input"]),
+ _take_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
+ ),
+ (
+ target + ".argmax",
+ *argmaxmin_pattern("relax.argmax"),
+ _argmaxmin_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
+ ),
+ (
+ target + ".argmin",
+ *argmaxmin_pattern("relax.argmin"),
+ _argmaxmin_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
+ ),
(
target + ".reshape",
*basic_pattern("relax.reshape", ["input", "input"]),
_reshape_check,
+ partial(msc_pattern.msc_attrs_getter, anchor="out"),
),
]
)
@@ -343,6 +403,7 @@ def get_patterns(target) -> List[Pattern]:
target + ".msc.conv2d_bias",
*msc_pattern.make_opt_relax_conv_bias_pattern("relax.nn.conv2d"),
wrap_basic_check(msc_pattern._check_opt_relax_conv_bias),
+ partial(msc_pattern.msc_attrs_getter, anchor="conv"),
),
]
)
diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py
index bf37cd7526..5bdea37b50 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -647,19 +647,21 @@ class FusionPattern(Object):
pattern: DFPattern
annotation_patterns: Mapping[str, DFPattern]
check: Callable[[PatternCheckContext], bool]
+ attrs_getter: Callable[[Dict[str, Expr]], Dict[str, str]]
def __init__(
self,
name: str,
pattern: DFPattern,
annotation_patterns: Optional[Mapping[str, DFPattern]] = None,
- check: Optional[Callable[[Mapping[str, Expr]], bool]] = None,
+ check: Optional[Callable[[PatternCheckContext], bool]] = None,
+ attrs_getter: Optional[Callable[[Dict[str, Expr]], Dict[str, str]]] = None,
):
if annotation_patterns is None:
annotation_patterns = {}
self.__init_handle_by_constructor__(
- _ffi_api.FusionPattern, name, pattern, annotation_patterns, check # type: ignore
- )
+ _ffi_api.FusionPattern, name, pattern, annotation_patterns, check, attrs_getter
+ ) # type: ignore
def FuseOpsByPattern(
diff --git a/src/contrib/msc/core/transform/fuse_tuple.cc b/src/contrib/msc/core/transform/fuse_tuple.cc
index e18d2cc35f..37baff4fa7 100644
--- a/src/contrib/msc/core/transform/fuse_tuple.cc
+++ b/src/contrib/msc/core/transform/fuse_tuple.cc
@@ -153,6 +153,8 @@ class TupleFuser : public ExprMutator {
Map<String, ObjectRef> func_attrs;
func_attrs.Set(tvm::relax::attr::kComposite, target_ + func_name);
func_attrs.Set(tvm::relax::attr::kPrimitive, Integer(1));
+ func_attrs.Set("unique_name", SpanUtils::GetAttr(expr->span, "name"));
+
Function function = Function(/*params=*/params, //
/*body=*/body, //
/*ret_struct_info=*/NullOpt, //
diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc
index 5cabfc40ca..b0eeba399e 100644
--- a/src/relax/transform/fuse_ops.cc
+++ b/src/relax/transform/fuse_ops.cc
@@ -1033,23 +1033,26 @@ class PatternBasedPartitioner : ExprVisitor {
using PatternCheckContext = transform::PatternCheckContext;
using ExprVisitor::VisitExpr_;
using FCheckMatch = runtime::TypedPackedFunc<bool(const transform::PatternCheckContext&)>;
+ using FAttrsGetter = runtime::TypedPackedFunc<Map<String, String>(const Map<String, Expr>&)>;
static GroupMap Run(String pattern_name, DFPattern pattern,
Map<String, DFPattern> annotation_patterns, FCheckMatch check, Expr expr,
- support::Arena* arena) {
- PatternBasedPartitioner part(pattern_name, pattern, annotation_patterns, check, arena);
+ support::Arena* arena, FAttrsGetter attrs_getter) {
+ PatternBasedPartitioner part(pattern_name, pattern, annotation_patterns, check, arena,
+ attrs_getter);
part.VisitExpr(expr);
return part.group_map_;
}
PatternBasedPartitioner(String pattern_name, DFPattern pattern,
Map<String, DFPattern> annotation_patterns, FCheckMatch check,
- support::Arena* arena)
+ support::Arena* arena, FAttrsGetter attrs_getter)
: pat_name_(pattern_name),
pat_(pattern),
annotation_pat_(annotation_patterns),
check_(check),
- arena_(arena) {}
+ arena_(arena),
+ attrs_getter_(attrs_getter) {}
void VisitBindingBlock_(const DataflowBlockNode* block) final {
current_block_use_def_ = DataflowBlockUseDef(GetRef<DataflowBlock>(block));
@@ -1070,7 +1073,8 @@ class PatternBasedPartitioner : ExprVisitor {
void VisitBinding_(const VarBindingNode* binding, const CallNode* call) final {
VisitVarDef(binding->var);
if (auto matches_opt = ExtractMatchedExpr(pat_, GetRef<Call>(call), bindings_)) {
- if (check_ != nullptr && !check_(CreatePatternCheckContext(call, matches_opt.value()))) {
+ const auto& context = CreatePatternCheckContext(call, matches_opt.value());
+ if (check_ != nullptr && !check_(context)) {
return;
}
@@ -1102,6 +1106,12 @@ class PatternBasedPartitioner : ExprVisitor {
auto parent_group = GetGroupForBoundVar(binding->var);
ICHECK(parent_group);
parent_group->attrs.Set(attr::kComposite, pat_name_);
+ if (attrs_getter_ != nullptr) {
+ const auto& custom_attrs = attrs_getter_(context->annotated_expr);
+ for (const auto& pair : custom_attrs) {
+ parent_group->attrs.Set(pair.first, pair.second);
+ }
+ }
for (const auto& [pat, match] : matches_opt.value()) {
// Put all matching expressions into the parent group. But we need to be careful not to
// merge expressions matched by a wildcard pattern, since a wildcard can match an output of
@@ -1163,6 +1173,7 @@ class PatternBasedPartitioner : ExprVisitor {
Map<String, DFPattern> annotation_pat_;
FCheckMatch check_;
support::Arena* arena_;
+ FAttrsGetter attrs_getter_;
Map<Var, Expr> bindings_;
Map<Expr, Var> value_to_bound_var_;
Map<Var, Array<Var>> current_block_use_def_;
@@ -1255,9 +1266,10 @@ IRModule FuseOpsByPattern(const tvm::Array<transform::FusionPattern>& patterns,
if (entry.second->IsInstance<tir::PrimFuncNode>()) {
continue;
}
- auto map = PatternBasedPartitioner::Run(
- pattern->name, pattern->pattern, pattern->annotation_patterns,
- pattern->check.value_or(nullptr), entry.second, &arena);
+ auto map = PatternBasedPartitioner::Run(pattern->name, pattern->pattern,
+ pattern->annotation_patterns,
+ pattern->check.value_or(nullptr), entry.second,
+ &arena, pattern->attrs_getter.value_or(nullptr));
group_map.insert(map.begin(), map.end());
}
mod = MakeGroupedFunctions(mod, group_map, /*lift_constants*/ !bind_constants);
@@ -1271,21 +1283,22 @@ IRModule FuseOpsByPattern(const tvm::Array<transform::FusionPattern>& patterns,
namespace transform {
FusionPattern::FusionPattern(String name, DFPattern pattern,
- Map<String, DFPattern> annotation_patterns,
- Optional<PackedFunc> check) {
+ Map<String, DFPattern> annotation_patterns, Optional<PackedFunc> check,
+ Optional<PackedFunc> attrs_getter) {
ObjectPtr<FusionPatternNode> n = make_object<FusionPatternNode>();
n->name = std::move(name);
n->pattern = std::move(pattern);
n->annotation_patterns = std::move(annotation_patterns);
n->check = check;
+ n->attrs_getter = attrs_getter;
data_ = std::move(n);
}
TVM_REGISTER_NODE_TYPE(FusionPatternNode);
TVM_REGISTER_GLOBAL("relax.transform.FusionPattern")
.set_body_typed([](String name, DFPattern pattern, Map<String, DFPattern> annotation_patterns,
- Optional<PackedFunc> check) {
- return FusionPattern(name, pattern, annotation_patterns, check);
+ Optional<PackedFunc> check, Optional<PackedFunc> attrs_getter) {
+ return FusionPattern(name, pattern, annotation_patterns, check, attrs_getter);
});
PatternCheckContext::PatternCheckContext(Expr matched_expr, Map<String, Expr> annotated_expr,
diff --git a/tests/python/contrib/test_msc/test_translate_tensorrt.py b/tests/python/contrib/test_msc/test_translate_tensorrt.py
index 0c142c1d46..262cf40adc 100644
--- a/tests/python/contrib/test_msc/test_translate_tensorrt.py
+++ b/tests/python/contrib/test_msc/test_translate_tensorrt.py
@@ -25,6 +25,7 @@ from torch import fx
from torch.nn import Module
import tvm.testing
+from tvm.relax import PyExprVisitor
from tvm.relax.frontend.torch import from_fx
from tvm.contrib.msc.framework.tensorrt.frontend import translate
from tvm.contrib.msc.framework.tensorrt import codegen
@@ -52,6 +53,40 @@ def build_and_run(mod, inputs):
return [e.asnumpy() for e in res]
+def check_names(mod):
+ """Check the byoc name and unique_name"""
+
+ @tvm.relax.expr_functor.visitor
+ class NameChecker(PyExprVisitor):
+ """Checker to check if any non-target ops exist"""
+
+ def check(self, expr):
+ self._recorded_names = set()
+ if isinstance(expr, tvm.relax.Expr):
+ self.visit_expr(expr)
+ elif isinstance(expr, tvm.relax.BindingBlock):
+ self.visit_binding_block(expr)
+
+ def visit_function_(self, op: tvm.relax.Function) -> None:
+ if "Composite" in op.attrs:
+ assert "unique_name" in op.attrs, "Can not find unique_name for func " + str(op)
+ name = str(op.attrs["unique_name"])
+ assert name not in self._recorded_names, "Name {} is already in use".format(name)
+ self._recorded_names.add(name)
+ super().visit_function_(op)
+
+ def _is_target_func(func):
+ if "Codegen" not in func.attrs:
+ return False
+ return func.attrs["Codegen"] == "msc_tensorrt"
+
+ for _, func in mod.functions.items():
+ if not _is_target_func(func):
+ continue
+ assert "byoc_name" in func.attrs, "Can not find byoc_name from function attributes"
+ NameChecker().check(func)
+
+
def verify_model(torch_model, input_info, allow_incomplete=False):
"""Build model and verify results"""
@@ -68,6 +103,7 @@ def verify_model(torch_model, input_info, allow_incomplete=False):
mod, graph_infos = translate.partition_for_tensorrt(
mod, trans_config={"allow_incomplete": allow_incomplete}
)
+ check_names(mod)
output_folder = msc_utils.msc_dir()
# tranalte to tensorrt
mod = codegen.to_tensorrt(mod, graph_infos, output_folder=output_folder)