Skip to content

Commit

Permalink
[TVMScript] Use op attribute to control whether to print dtype in TVM…
Browse files Browse the repository at this point in the history 
…Script (apache#14111)

This PR adds an op attribute `TScriptDtypePrintLocation`, and modifies the dtype printing logic of the builtin op to check this attribute. So that user defined operators can use it to specify how there dtype argument are printed by appending attributes instead of appending members to `dtype_first_arg`/`dtype_last_arg`.
  • Loading branch information
@liangW-intellif @yongwww
liangW-intellif authored and yongwww committed Feb 27, 2023
1 parent 17b0b50 commit 144fb49
Show file tree
Hide file tree
Showing 4 changed files with 82 additions and 41 deletions.
21 changes: 21 additions & 0 deletions include/tvm/tir/op_attr_types.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,27 @@ using FLegalize = runtime::TypedPackedFunc<PrimExpr(PrimExpr)>;
*/
using TScriptPrinterName = String;

/*!
* \brief Specifies that TVMScript printer prints the dtype as the first/last argument.
If not specified, dtype will not be printed.
*/
enum class ScriptDtypePrintLocation : int {
/*!
* \brief Do not print dtype as an argument.
*/
kNone = 0,
/*!
* \brief Print dtype as the first argument.
*/
kFirst = 1,
/*!
* \brief FPrint dtype as the last argument.
*/
kLast = 2,
};

using TScriptDtypePrintLocation = Integer;

/*!
* \brief The effect type of the call.
*/
Expand Down
31 changes: 9 additions & 22 deletions src/script/printer/tir/expr.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -222,33 +222,20 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<tir::Call>("", [](tir::Call call, ObjectPath call_p, IRDocsifier d) -> Doc {
static const OpAttrMap<tir::TScriptPrinterName>& op_names =
Op::GetAttrMap<tir::TScriptPrinterName>("TScriptPrinterName");
static const std::unordered_set<const Object*> dtype_first_arg = {
tir::builtin::reinterpret().get(),
tir::builtin::call_extern().get(),
tir::builtin::call_llvm_intrin().get(), //
tir::builtin::call_llvm_pure_intrin().get(), //
tir::builtin::call_pure_extern().get(), //
tir::builtin::ptx_mma().get(),
tir::builtin::ptx_mma_sp().get(),
tir::builtin::ptx_ldmatrix().get(),
tir::builtin::ptx_cp_async().get(),
tir::builtin::mma_store().get(),
tir::builtin::mma_fill().get(),
tir::builtin::vectorlow().get(),
tir::builtin::vectorhigh().get(),
tir::builtin::vectorcombine().get(),
Op::Get("tir.type_annotation").get(),
};
static const std::unordered_set<const Object*> dtype_last_arg = {
tir::builtin::tvm_struct_get().get(),
};
static const OpAttrMap<tir::TScriptDtypePrintLocation> dtype_locations =
Op::GetAttrMap<tir::TScriptDtypePrintLocation>("TScriptDtypePrintLocation");
tir::ScriptDtypePrintLocation dtype_print_location = tir::ScriptDtypePrintLocation::kNone;
ExprDoc prefix{nullptr};
if (const auto* op = call->op.as<OpNode>()) {
String name = op_names.get(GetRef<Op>(op), op->name);
if (op_names.count(GetRef<Op>(op)) == 0) {
LOG(WARNING) << "No TScriptPrinterName attribute for " << op->name;
}
prefix = TIR(d, name);
if (dtype_locations.count(GetRef<Op>(op))) {
dtype_print_location = static_cast<tir::ScriptDtypePrintLocation>(
dtype_locations[GetRef<Op>(op)].IntValue());
}
} else if (const auto* gv = call->op.as<GlobalVarNode>()) {
prefix = LiteralDoc::Str(gv->name_hint, call_p->Attr("op"));
} else {
Expand All @@ -257,13 +244,13 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
Array<ExprDoc> args;
int n_args = call->args.size();
args.reserve(n_args + 1);
if (dtype_first_arg.count(call->op.get())) {
if (dtype_print_location == tir::ScriptDtypePrintLocation::kFirst) {
args.push_back(LiteralDoc::DataType(call->dtype, call_p->Attr("dtype")));
}
for (int i = 0; i < n_args; ++i) {
args.push_back(d->AsDoc<ExprDoc>(call->args[i], call_p->Attr("args")->ArrayIndex(i)));
}
if (dtype_last_arg.count(call->op.get())) {
if (dtype_print_location == tir::ScriptDtypePrintLocation::kLast) {
args.push_back(LiteralDoc::DataType(call->dtype, call_p->Attr("dtype")));
}
return prefix->Call(args);
Expand Down
4 changes: 3 additions & 1 deletion src/tir/ir/stmt.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -692,7 +692,9 @@ PrimExpr TypeAnnotation(DataType dtype, Span span) {
}

TVM_TIR_REGISTER_OP("type_annotation")
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

} // namespace tir
} // namespace tvm
67 changes: 49 additions & 18 deletions src/tir/op/builtin.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ namespace builtin {

TIR_DEFINE_BUILTIN_FUNC(reinterpret)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst))
.set_num_inputs(1);

TIR_DEFINE_BUILTIN_FUNC(ret)
Expand Down Expand Up @@ -120,16 +122,24 @@ TIR_DEFINE_BUILTIN_FUNC(fma)
.set_attr<TVectorizable>("TVectorizable", true);

TIR_DEFINE_BUILTIN_FUNC(call_extern)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(call_pure_extern)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(call_llvm_intrin)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(call_llvm_pure_intrin)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(call_spirv_pure_glsl450)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
Expand All @@ -154,7 +164,9 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_tuple).set_attr<TCallEffectKind>("TCallEffectKind",

TIR_DEFINE_BUILTIN_FUNC(tvm_struct_get)
.set_num_inputs(3)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kReadState));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kReadState))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kLast));

TIR_DEFINE_BUILTIN_FUNC(tvm_struct_set)
.set_num_inputs(4)
Expand Down Expand Up @@ -249,40 +261,59 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_fill_fragment)
TIR_DEFINE_BUILTIN_FUNC(tvm_store_matrix_sync)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));

TIR_DEFINE_BUILTIN_FUNC(ptx_mma).set_attr<TCallEffectKind>("TCallEffectKind",
Integer(CallEffectKind::kOpaque));
TIR_DEFINE_BUILTIN_FUNC(ptx_mma)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(ptx_ldg32).set_num_inputs(4).set_attr<TCallEffectKind>(
"TCallEffectKind", Integer(CallEffectKind::kPure));

TIR_DEFINE_BUILTIN_FUNC(ptx_mma_sp)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(ptx_ldmatrix)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(ptx_cp_async)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(ptx_commit_group)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));

TIR_DEFINE_BUILTIN_FUNC(ptx_wait_group)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));

TIR_DEFINE_BUILTIN_FUNC(mma_store).set_attr<TCallEffectKind>("TCallEffectKind",
Integer(CallEffectKind::kOpaque));
TIR_DEFINE_BUILTIN_FUNC(mma_store)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(mma_fill).set_attr<TCallEffectKind>("TCallEffectKind",
Integer(CallEffectKind::kOpaque));
TIR_DEFINE_BUILTIN_FUNC(mma_fill)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(vectorhigh)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(vectorlow).set_attr<TCallEffectKind>("TCallEffectKind",
Integer(CallEffectKind::kPure));
TIR_DEFINE_BUILTIN_FUNC(vectorlow)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(vectorcombine)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure));
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));

TIR_DEFINE_BUILTIN_FUNC(atomic_add)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
Expand Down

0 comments on commit 144fb49

Please sign in to comment.