You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "Lunderberg (via GitHub)" <gi...@apache.org> on 2023/11/13 18:32:40 UTC
[PR] [Draft][Unity][Transform] Extract partial-tuple-usage from FuseTIR [tvm]
Lunderberg opened a new pull request, #16120:
URL: https://github.com/apache/tvm/pull/16120
Prior to this commit, the `FuseTIR` pass explicitly tracked usage of
tuple arguments, to minimize the set of arguments provided to each
kernel. The additional tgracking and handling of partially-used
tuples makes it difficult to follow the primary changes being made by
`FuseTIR`.
This commit implements the same functionality in terms of the
`ExpandTupleArguments` and `RemoveUnusedParameters` transforms,
introduced in https://github.com/apache/tvm/pull/16115 and
https://github.com/apache/tvm/pull/16116 respectively. By using these
passes before the main `FuseOps` changes, partial tuple usage is
already handled at that point.
This commit is intended to minimize any changes to user-facing
behavior, and so these pre-process passes are currently used
internally by `FuseTIR`. This may be avoided in the future by pulling
this internal delegation out into a lowering pipeline.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Extract partial-tuple-usage from FuseTIR [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on PR #16120:
URL: https://github.com/apache/tvm/pull/16120#issuecomment-1863304979
Rebased onto unity to avoid stale CI results.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Draft][Unity][Transform] Extract partial-tuple-usage from FuseTIR [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on PR #16120:
URL: https://github.com/apache/tvm/pull/16120#issuecomment-1808783752
This PR is marked as a draft, until the pre-requisite PRs #16115, #16116, #16117, and #16118 land.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Extract partial-tuple-usage from FuseTIR [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16120:
URL: https://github.com/apache/tvm/pull/16120#discussion_r1431893969
##########
src/relax/transform/fuse_tir.cc:
##########
@@ -719,64 +697,46 @@ class FusedTIRConstructor : public ExprVisitor {
}
/*!
- * \brief Create an TIR func params and buffers with specified relax type and shape
+ * \brief Collect TIR func params and buffers with specified relax type and shape
* \param struct_info The struct info
* \param name_hint The name hint for params and buffers
- * \param index The index used for unique name_hint if type is Tuple.
- * -1 means no need to add postfix since the relax param is not a Tuple.
- * \return The created TIR func params and buffers
+ * \param out The vector into which to collect the params/buffers
*/
- static std::pair<Array<tir::Var>, Array<tir::Buffer>> CreateParamsAndBuffers(
- StructInfo struct_info, const String& name_hint, int index = -1) {
- Array<tir::Var> params;
- Array<tir::Buffer> buffers;
- // The symbolic shape params must be defined at the end of the param list.
- bool symbolic_shape_param_started = false;
+ static void CollectPrimFuncParams(const Var& relax_param,
+ std::vector<Variant<tir::Var, tir::Buffer>>* out) {
+ auto struct_info = GetStructInfo(relax_param);
+
+ CHECK(!struct_info.as<TupleStructInfoNode>())
+ << "InternalError: "
+ << "All tuple parameters should be expanded before this point in FuseTIR. "
+ << "However, parameter " << relax_param << " has struct info " << struct_info;
+
+ auto name_hint = relax_param->name_hint();
+
if (const auto* tensor = struct_info.as<TensorStructInfoNode>()) {
- // Case 1. the relax param is a Tensor, we directly create a tir var and buffer
+ // Case 1. The relax param is a Tensor, we directly create a tir var and buffer
const auto* shape_expr = tensor->shape.as<ShapeExprNode>();
- ICHECK(shape_expr) << "FuseTIR expects all parameters are Tensors with symbolic shape.";
- CHECK(!symbolic_shape_param_started)
- << "The symbolic shape params must be defined at the end of the param "
- "list.";
- String name = index == -1 ? name_hint : name_hint + "_" + std::to_string(index);
+ ICHECK(shape_expr) << "FuseTIR expects all Tensor parameters have a known shape.";
DataType dtype = tensor->dtype;
- tir::Buffer buffer = tir::decl_buffer(shape_expr->values, dtype, name);
- // Differentiate buffer name and param name by adding prefix `v_` to param
- // Every symbol should be unique in TVMScript, and Buffer is used more than param
- // So we decide to make sure buffer names have better readability.
- tir::Var param = tir::Var("p_" + name, PrimType(DataType::Handle()));
- params.push_back(std::move(param));
- buffers.push_back(std::move(buffer));
- } else if (const auto* tuple = struct_info.as<TupleStructInfoNode>()) {
- // Case 2. the relax param is a Tuple, we recursively visit each field until it's a Tensor
- // Enable postfix
- CHECK(!symbolic_shape_param_started)
- << "The symbolic shape params must be defined at the end of the param "
- "list.";
- if (index == -1) index = 0;
- for (size_t i = 0; i < tuple->fields.size(); ++i) {
- auto [ret_params, ret_buffers] = CreateParamsAndBuffers(tuple->fields[i], name_hint, index);
- ICHECK_EQ(ret_params.size(), ret_buffers.size());
- // Adding tuple field results to the end of params and buffers.
- params.insert(params.end(), ret_params.begin(), ret_params.end());
- buffers.insert(buffers.end(), ret_buffers.begin(), ret_buffers.end());
- index += ret_params.size();
- }
+ tir::Buffer buffer = tir::decl_buffer(shape_expr->values, dtype, name_hint);
+ out->push_back(std::move(buffer));
+
+ } else if (const auto* prim_value = struct_info.as<PrimStructInfoNode>()) {
+ // Case 2. The relax param is a scalar, we directly create a tir var
+ ICHECK(prim_value->value->IsInstance<tir::VarNode>());
+ out->push_back(Downcast<tir::Var>(prim_value->value));
+
} else if (const auto* shape_expr = struct_info.as<ShapeStructInfoNode>()) {
- // Case 3. the relax param is a scalar, we directly create a tir var
- symbolic_shape_param_started = true;
- ICHECK(index == -1) << "TypeError: The ShapeExprNode should not be in a Tuple field.";
+ // Case 3. The relax param is a tuple of scalars, each represented as a tir var
for (const auto& var : shape_expr->values.value()) {
ICHECK(var->IsInstance<tir::VarNode>());
- params.push_back(Downcast<tir::Var>(var));
+ out->push_back(Downcast<tir::Var>(var));
}
} else {
ICHECK(false) << "TypeError: The param type of PrimFunc is expected to be Tensor, Tuple or "
Review Comment:
Nitpick: tuples are no longer accepted, so the message is inaccurate (though a tuple would have been caught above).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Extract partial-tuple-usage from FuseTIR [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16120:
URL: https://github.com/apache/tvm/pull/16120#discussion_r1432197350
##########
src/relax/transform/fuse_tir.cc:
##########
@@ -719,64 +697,46 @@ class FusedTIRConstructor : public ExprVisitor {
}
/*!
- * \brief Create an TIR func params and buffers with specified relax type and shape
+ * \brief Collect TIR func params and buffers with specified relax type and shape
* \param struct_info The struct info
* \param name_hint The name hint for params and buffers
- * \param index The index used for unique name_hint if type is Tuple.
- * -1 means no need to add postfix since the relax param is not a Tuple.
- * \return The created TIR func params and buffers
+ * \param out The vector into which to collect the params/buffers
*/
- static std::pair<Array<tir::Var>, Array<tir::Buffer>> CreateParamsAndBuffers(
- StructInfo struct_info, const String& name_hint, int index = -1) {
- Array<tir::Var> params;
- Array<tir::Buffer> buffers;
- // The symbolic shape params must be defined at the end of the param list.
- bool symbolic_shape_param_started = false;
+ static void CollectPrimFuncParams(const Var& relax_param,
+ std::vector<Variant<tir::Var, tir::Buffer>>* out) {
+ auto struct_info = GetStructInfo(relax_param);
+
+ CHECK(!struct_info.as<TupleStructInfoNode>())
+ << "InternalError: "
+ << "All tuple parameters should be expanded before this point in FuseTIR. "
+ << "However, parameter " << relax_param << " has struct info " << struct_info;
+
+ auto name_hint = relax_param->name_hint();
+
if (const auto* tensor = struct_info.as<TensorStructInfoNode>()) {
- // Case 1. the relax param is a Tensor, we directly create a tir var and buffer
+ // Case 1. The relax param is a Tensor, we directly create a tir var and buffer
const auto* shape_expr = tensor->shape.as<ShapeExprNode>();
- ICHECK(shape_expr) << "FuseTIR expects all parameters are Tensors with symbolic shape.";
- CHECK(!symbolic_shape_param_started)
- << "The symbolic shape params must be defined at the end of the param "
- "list.";
- String name = index == -1 ? name_hint : name_hint + "_" + std::to_string(index);
+ ICHECK(shape_expr) << "FuseTIR expects all Tensor parameters have a known shape.";
DataType dtype = tensor->dtype;
- tir::Buffer buffer = tir::decl_buffer(shape_expr->values, dtype, name);
- // Differentiate buffer name and param name by adding prefix `v_` to param
- // Every symbol should be unique in TVMScript, and Buffer is used more than param
- // So we decide to make sure buffer names have better readability.
- tir::Var param = tir::Var("p_" + name, PrimType(DataType::Handle()));
- params.push_back(std::move(param));
- buffers.push_back(std::move(buffer));
- } else if (const auto* tuple = struct_info.as<TupleStructInfoNode>()) {
- // Case 2. the relax param is a Tuple, we recursively visit each field until it's a Tensor
- // Enable postfix
- CHECK(!symbolic_shape_param_started)
- << "The symbolic shape params must be defined at the end of the param "
- "list.";
- if (index == -1) index = 0;
- for (size_t i = 0; i < tuple->fields.size(); ++i) {
- auto [ret_params, ret_buffers] = CreateParamsAndBuffers(tuple->fields[i], name_hint, index);
- ICHECK_EQ(ret_params.size(), ret_buffers.size());
- // Adding tuple field results to the end of params and buffers.
- params.insert(params.end(), ret_params.begin(), ret_params.end());
- buffers.insert(buffers.end(), ret_buffers.begin(), ret_buffers.end());
- index += ret_params.size();
- }
+ tir::Buffer buffer = tir::decl_buffer(shape_expr->values, dtype, name_hint);
+ out->push_back(std::move(buffer));
+
+ } else if (const auto* prim_value = struct_info.as<PrimStructInfoNode>()) {
+ // Case 2. The relax param is a scalar, we directly create a tir var
+ ICHECK(prim_value->value->IsInstance<tir::VarNode>());
+ out->push_back(Downcast<tir::Var>(prim_value->value));
+
} else if (const auto* shape_expr = struct_info.as<ShapeStructInfoNode>()) {
- // Case 3. the relax param is a scalar, we directly create a tir var
- symbolic_shape_param_started = true;
- ICHECK(index == -1) << "TypeError: The ShapeExprNode should not be in a Tuple field.";
+ // Case 3. The relax param is a tuple of scalars, each represented as a tir var
for (const auto& var : shape_expr->values.value()) {
ICHECK(var->IsInstance<tir::VarNode>());
- params.push_back(Downcast<tir::Var>(var));
+ out->push_back(Downcast<tir::Var>(var));
}
} else {
ICHECK(false) << "TypeError: The param type of PrimFunc is expected to be Tensor, Tuple or "
Review Comment:
Good catch, and updated to remove `Tuple` from the list, and add `PrimValue` as it is now handled.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Extract partial-tuple-usage from FuseTIR [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16120:
URL: https://github.com/apache/tvm/pull/16120#discussion_r1432200745
##########
src/relax/transform/fuse_tir.cc:
##########
@@ -385,58 +385,45 @@ class FusedTIRConstructor : public ExprVisitor {
: mod_(mod), func_name_(func_name) {}
void VisitExpr_(const FunctionNode* func) final {
- // Step 1. Create buffers for function params
-
- // Record which fields in a tuple passed as a parameter are actually accessed by the function.
- std::unordered_set<const Object*> tuple_param;
- for (auto param : func->params) {
- if (GetStructInfo(param)->IsInstance<TupleStructInfoNode>()) {
- tuple_param.insert(param.get());
- }
- }
-
- PostOrderVisit(func->body, [=, &tuple_param](Expr e) {
- if (auto tup_get = e.as<TupleGetItemNode>();
- tup_get && tuple_param.count(tup_get->tuple.get())) {
- func_info_.used_tuple_field_indices[tup_get->tuple.get()].insert(tup_get->index);
- }
- });
-
+ std::vector<Variant<tir::Var, tir::Buffer>> prim_func_params;
for (const Var& relax_param : func->params) {
- auto sinfo = GetStructInfo(relax_param);
- if (sinfo->IsInstance<ShapeStructInfoNode>()) {
- // It's a symbolic shape var, no need to alloc Buffers.
- continue;
- }
-
- auto [params, buffers] = [=]() {
- if (const auto* tuple = sinfo.as<TupleStructInfoNode>()) {
- // Add only those tuple fields which are actually used by the function body into the
- // function parameters.
- int index = 0;
- Array<tir::Var> params;
- Array<tir::Buffer> buffers;
- for (auto i : func_info_.used_tuple_field_indices[relax_param.get()]) {
- auto [ret_params, ret_buffers] =
- CreateParamsAndBuffers(tuple->fields[i], relax_param->name_hint(), index);
- ICHECK_EQ(ret_params.size(), ret_buffers.size());
- // Adding tuple field results to the end of params and buffers.
- params.insert(params.end(), ret_params.begin(), ret_params.end());
- buffers.insert(buffers.end(), ret_buffers.begin(), ret_buffers.end());
- index += ret_params.size();
+ size_t size_before = prim_func_params.size();
+ CollectPrimFuncParams(relax_param, &prim_func_params);
+
+ auto param_buffers = [&]() -> Array<tir::Buffer> {
+ Array<tir::Buffer> out;
+ for (size_t i = size_before; i < prim_func_params.size(); i++) {
+ if (auto buf = prim_func_params[i].as<tir::Buffer>()) {
+ out.push_back(buf.value());
}
- return std::make_pair(params, buffers);
- } else {
- return CreateParamsAndBuffers(sinfo, relax_param->name_hint());
}
+ return out;
}();
- ICHECK_EQ(params.size(), buffers.size());
- for (size_t i = 0; i < params.size(); ++i) {
- func_info_.buffer_map.Set(params[i], buffers[i]);
- func_info_.params.push_back(params[i]);
+ func_info_.expr2buffers.Set(relax_param, param_buffers);
+ }
+
+ // Move all scalar params after buffer params.
+ std::stable_sort(prim_func_params.begin(), prim_func_params.end(),
Review Comment:
Good call, and updated the comment. It's mainly to make sure that the order is consistent and predictable for unit tests, not for any correctness of the output.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Extract partial-tuple-usage from FuseTIR [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16120:
URL: https://github.com/apache/tvm/pull/16120#discussion_r1431894344
##########
src/relax/transform/fuse_tir.cc:
##########
@@ -385,58 +385,45 @@ class FusedTIRConstructor : public ExprVisitor {
: mod_(mod), func_name_(func_name) {}
void VisitExpr_(const FunctionNode* func) final {
- // Step 1. Create buffers for function params
-
- // Record which fields in a tuple passed as a parameter are actually accessed by the function.
- std::unordered_set<const Object*> tuple_param;
- for (auto param : func->params) {
- if (GetStructInfo(param)->IsInstance<TupleStructInfoNode>()) {
- tuple_param.insert(param.get());
- }
- }
-
- PostOrderVisit(func->body, [=, &tuple_param](Expr e) {
- if (auto tup_get = e.as<TupleGetItemNode>();
- tup_get && tuple_param.count(tup_get->tuple.get())) {
- func_info_.used_tuple_field_indices[tup_get->tuple.get()].insert(tup_get->index);
- }
- });
-
+ std::vector<Variant<tir::Var, tir::Buffer>> prim_func_params;
for (const Var& relax_param : func->params) {
- auto sinfo = GetStructInfo(relax_param);
- if (sinfo->IsInstance<ShapeStructInfoNode>()) {
- // It's a symbolic shape var, no need to alloc Buffers.
- continue;
- }
-
- auto [params, buffers] = [=]() {
- if (const auto* tuple = sinfo.as<TupleStructInfoNode>()) {
- // Add only those tuple fields which are actually used by the function body into the
- // function parameters.
- int index = 0;
- Array<tir::Var> params;
- Array<tir::Buffer> buffers;
- for (auto i : func_info_.used_tuple_field_indices[relax_param.get()]) {
- auto [ret_params, ret_buffers] =
- CreateParamsAndBuffers(tuple->fields[i], relax_param->name_hint(), index);
- ICHECK_EQ(ret_params.size(), ret_buffers.size());
- // Adding tuple field results to the end of params and buffers.
- params.insert(params.end(), ret_params.begin(), ret_params.end());
- buffers.insert(buffers.end(), ret_buffers.begin(), ret_buffers.end());
- index += ret_params.size();
+ size_t size_before = prim_func_params.size();
+ CollectPrimFuncParams(relax_param, &prim_func_params);
+
+ auto param_buffers = [&]() -> Array<tir::Buffer> {
+ Array<tir::Buffer> out;
+ for (size_t i = size_before; i < prim_func_params.size(); i++) {
+ if (auto buf = prim_func_params[i].as<tir::Buffer>()) {
+ out.push_back(buf.value());
}
- return std::make_pair(params, buffers);
- } else {
- return CreateParamsAndBuffers(sinfo, relax_param->name_hint());
}
+ return out;
}();
- ICHECK_EQ(params.size(), buffers.size());
- for (size_t i = 0; i < params.size(); ++i) {
- func_info_.buffer_map.Set(params[i], buffers[i]);
- func_info_.params.push_back(params[i]);
+ func_info_.expr2buffers.Set(relax_param, param_buffers);
+ }
+
+ // Move all scalar params after buffer params.
+ std::stable_sort(prim_func_params.begin(), prim_func_params.end(),
Review Comment:
I assume it's important for the relative ordering to be preserved (hence the stable sort), might be good to call that out.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Extract partial-tuple-usage from FuseTIR [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on PR #16120:
URL: https://github.com/apache/tvm/pull/16120#issuecomment-1836826990
All pre-requisite PRs have landed, and this is now ready for review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Extract partial-tuple-usage from FuseTIR [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg merged PR #16120:
URL: https://github.com/apache/tvm/pull/16120
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org