You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "rutkoor (via GitHub)" <gi...@apache.org> on 2023/09/06 07:02:24 UTC
[GitHub] [tvm] rutkoor opened a new pull request, #15679: [Unity] Support Padding Reversal in Alter-Op pass
rutkoor opened a new pull request, #15679:
URL: https://github.com/apache/tvm/pull/15679
This patch adds support for non-bijective transformations in alter-op-impl pass. The idea is to introduce remove_pad operator, which copies non-padded buffer entries from padded_buffer. This additional copy operation will suffice as a placeholder for supporting non-bijective transformations. Further we can eliminate this copy operation in later stages of compilation 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
[GitHub] [tvm] masahi merged pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi merged PR #15679:
URL: https://github.com/apache/tvm/pull/15679
--
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
[GitHub] [tvm] rutkoor commented on pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on PR #15679:
URL: https://github.com/apache/tvm/pull/15679#issuecomment-1715399793
Gentle reminder to please review this PR.
--
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
[GitHub] [tvm] rutkoor commented on pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on PR #15679:
URL: https://github.com/apache/tvm/pull/15679#issuecomment-1711281848
@tvm-bot rerun
--
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
[GitHub] [tvm] rutkoor commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1338494963
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -182,7 +182,15 @@ def te_layout_transform(data, name):
)
index_map: tvm.tir.IndexMap = call.attrs.index_map
- pad_value = call.attrs.pad_value.value
+ pad_value = call.attrs.pad_value
+ if pad_value is not None:
+ pad_value = pad_value.value
+ else:
+ if "int" in call.args[0].struct_info.dtype:
+ pad_value = int(pad_value)
Review Comment:
Corrected.
--
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
[GitHub] [tvm] rutkoor commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1339485356
##########
src/relax/transform/alter_op_impl.cc:
##########
@@ -314,6 +377,8 @@ class AlterOpImplMutator : public ExprMutator {
Map<PrimFunc, GlobalVar> cache_;
/*! \brief Input IRModule */
const IRModule& mod_;
+ /*! \brief Map from shape_dim attribute to the remove_pad GlobalVar */
+ Map<String, GlobalVar> remove_pad_map_;
Review Comment:
Using `std::unordered_map<int, GlobalVar>` now.
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1339295756
##########
src/relax/transform/alter_op_impl.cc:
##########
@@ -314,6 +377,8 @@ class AlterOpImplMutator : public ExprMutator {
Map<PrimFunc, GlobalVar> cache_;
/*! \brief Input IRModule */
const IRModule& mod_;
+ /*! \brief Map from shape_dim attribute to the remove_pad GlobalVar */
+ Map<String, GlobalVar> remove_pad_map_;
Review Comment:
Why not `std::unordered_map<int, GlobalVar>'?
--
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
[GitHub] [tvm] rutkoor commented on pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on PR #15679:
URL: https://github.com/apache/tvm/pull/15679#issuecomment-1711288315
@tvm-bot rerun
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1323414157
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
How do you know that this removes the correct padded values?
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1336715785
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -182,7 +182,15 @@ def te_layout_transform(data, name):
)
index_map: tvm.tir.IndexMap = call.attrs.index_map
- pad_value = call.attrs.pad_value.value
+ pad_value = call.attrs.pad_value
+ if pad_value is not None:
+ pad_value = pad_value.value
+ else:
+ if "int" in call.args[0].struct_info.dtype:
+ pad_value = int(pad_value)
Review Comment:
pad_value is None here
##########
src/relax/transform/alter_op_impl.cc:
##########
@@ -176,16 +190,41 @@ class AlterOpImplMutator : public ExprMutator {
Expr TransformLayoutInverse(const Expr& expr, const IndexMap& index_map,
const TensorStructInfo& old_tensor_sinfo,
const Array<IntImm>& axis_separator) {
+ if (IsScalarConstant(expr) || index_map.get() == nullptr) {
+ return expr;
+ }
Array<PrimExpr> old_shape = GetShapeFromTensorStructInfo(old_tensor_sinfo);
Array<Range> initial_ranges = ConstructRangeFromShape(old_shape);
arith::Analyzer analyzer;
auto [inverse_index_map, padding_predicate] =
index_map.NonSurjectiveInverse(initial_ranges, &analyzer);
- ICHECK(tir::is_zero(padding_predicate))
- << "Only bijective transformations on input/output buffers are supported, but found "
- "padding predicate "
- << padding_predicate << " on initial range " << initial_ranges;
- return TransformLayout(expr, inverse_index_map, axis_separator);
+
+ if (tir::is_zero(padding_predicate)) {
+ return TransformLayout(expr, inverse_index_map, axis_separator);
+ } else {
+ auto padded_expr =
+ builder_->Normalize(TransformLayout(expr, inverse_index_map, axis_separator));
+ const auto& tensor_sinfo = Downcast<TensorStructInfo>(padded_expr->struct_info_);
+ Array<PrimExpr> padded_shape = GetShapeFromTensorStructInfo(tensor_sinfo);
+
+ te::Tensor placeholder_tensor = te::placeholder(padded_shape, tensor_sinfo->dtype, "input");
+ te::Tensor output_tensor = te::compute(
+ old_shape,
+ [&placeholder_tensor](const Array<tir::Var>& indices) {
+ return placeholder_tensor(indices);
+ },
+ "output", topi::kElementWise);
+
+ String op_name = "remove_pad";
+ PrimFunc remove_pad_with_frozen_layout =
+ WithAttr(CreatePrimFunc({placeholder_tensor, output_tensor}), kOperatorName, op_name);
+
+ GlobalVar gv_remove_pad = builder_->AddFunction(remove_pad_with_frozen_layout, op_name);
+ builder_->UpdateFunction(gv_remove_pad,
+ WithoutAttr(remove_pad_with_frozen_layout, "global_symbol"));
Review Comment:
Code between L210 - L224 is hard to read and it seems it doesn't need to be defined here. I think you can define `gv_remove_pad` with dynamic input shape once at the constructor and use that repeatedly here.
--
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
[GitHub] [tvm] rutkoor commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1338498012
##########
src/relax/transform/alter_op_impl.cc:
##########
@@ -176,16 +190,41 @@ class AlterOpImplMutator : public ExprMutator {
Expr TransformLayoutInverse(const Expr& expr, const IndexMap& index_map,
const TensorStructInfo& old_tensor_sinfo,
const Array<IntImm>& axis_separator) {
+ if (IsScalarConstant(expr) || index_map.get() == nullptr) {
+ return expr;
+ }
Array<PrimExpr> old_shape = GetShapeFromTensorStructInfo(old_tensor_sinfo);
Array<Range> initial_ranges = ConstructRangeFromShape(old_shape);
arith::Analyzer analyzer;
auto [inverse_index_map, padding_predicate] =
index_map.NonSurjectiveInverse(initial_ranges, &analyzer);
- ICHECK(tir::is_zero(padding_predicate))
- << "Only bijective transformations on input/output buffers are supported, but found "
- "padding predicate "
- << padding_predicate << " on initial range " << initial_ranges;
- return TransformLayout(expr, inverse_index_map, axis_separator);
+
+ if (tir::is_zero(padding_predicate)) {
+ return TransformLayout(expr, inverse_index_map, axis_separator);
+ } else {
+ auto padded_expr =
+ builder_->Normalize(TransformLayout(expr, inverse_index_map, axis_separator));
+ const auto& tensor_sinfo = Downcast<TensorStructInfo>(padded_expr->struct_info_);
+ Array<PrimExpr> padded_shape = GetShapeFromTensorStructInfo(tensor_sinfo);
+
+ te::Tensor placeholder_tensor = te::placeholder(padded_shape, tensor_sinfo->dtype, "input");
+ te::Tensor output_tensor = te::compute(
+ old_shape,
+ [&placeholder_tensor](const Array<tir::Var>& indices) {
+ return placeholder_tensor(indices);
+ },
+ "output", topi::kElementWise);
+
+ String op_name = "remove_pad";
+ PrimFunc remove_pad_with_frozen_layout =
+ WithAttr(CreatePrimFunc({placeholder_tensor, output_tensor}), kOperatorName, op_name);
+
+ GlobalVar gv_remove_pad = builder_->AddFunction(remove_pad_with_frozen_layout, op_name);
+ builder_->UpdateFunction(gv_remove_pad,
+ WithoutAttr(remove_pad_with_frozen_layout, "global_symbol"));
Review Comment:
I have introduced `GetOrCreateRemovePadOp` function which returns the `GlobalVar` of the `remove_pad PrimFunc` if it is present in `remove_pad_map_` map variable.
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1324919768
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
Better yet, is this actually needs to be an "op"? Since the op definition is so trivial, how about moving this compute definition to `alter_op_layout.cc` and use it there? You can remove all boilerplate and test cases required for a Relax op definition.
--
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
[GitHub] [tvm] rutkoor commented on pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on PR #15679:
URL: https://github.com/apache/tvm/pull/15679#issuecomment-1712420060
cc: @psrivas2 @masahi @junrushao
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1335797265
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
If all tests pass, I think it is fine to include such change in this PR.
--
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
[GitHub] [tvm] rutkoor commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1335492704
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
Hi @masahi, the error looks something like this,
`ValueError: StructuralEqual check failed, caused by lhs at <root>.functions[I.GlobalVar("remove_pad")].body.block.body.body.block.iter_vars[0].dom.min.dtype and rhs at <root>.functions[I.GlobalVar("remove_pad")].body.block.body.body.block.iter_vars[0].dom.min.dtype`
where the lhs dtype is coming out to be `int32` when running the `alter_op_impl `pass, and rhs dtype is `int64` which is coming from `Expected module`.
I have commited a change in `src/te/operation/compute_op.cc` which is resolving the issue.
My question is, does the change in `src/te/operation/compute_op.cc` should go as a separate PR or can I have it in this PR only?
--
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
[GitHub] [tvm] rutkoor commented on pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on PR #15679:
URL: https://github.com/apache/tvm/pull/15679#issuecomment-1711284797
@tvm-bot re-run
--
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
[GitHub] [tvm] rutkoor commented on pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on PR #15679:
URL: https://github.com/apache/tvm/pull/15679#issuecomment-1711285581
@tvm-bot re-build
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1324325599
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
> currently the elements are padded in the end of the original buffer
Isn't that a specific behavior of the layout transform pass? Here you are defining a general function that undoes padding, and padding can be anything along the begin / end of each axis.
--
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
[GitHub] [tvm] rutkoor commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1324346506
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
Technically the behavior of remove_pad can be extended to remove padding along the given axis and given offset. In this PR, the support is added for AlterOp pass explicitly, to remove the padding which is being introduced by layout transform pass. In the next PR, I plan to work on generalizing remove_pad functionality to be generic.
Will it be good idea to separate out this PR from generalizing remove_pad functionality?
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1339295756
##########
src/relax/transform/alter_op_impl.cc:
##########
@@ -314,6 +377,8 @@ class AlterOpImplMutator : public ExprMutator {
Map<PrimFunc, GlobalVar> cache_;
/*! \brief Input IRModule */
const IRModule& mod_;
+ /*! \brief Map from shape_dim attribute to the remove_pad GlobalVar */
+ Map<String, GlobalVar> remove_pad_map_;
Review Comment:
Why not `std::unordered_map<int, GlobalVar>`?
--
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
[GitHub] [tvm] rutkoor commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1334728620
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
Thanks @masahi for your helpful suggestions. In the latest changes I have defined remove_pad operator inside `alter_op_impl.cc`.
As of now, the test case in this PR fails because the `te::compute` introduces min Range value to `int32(0)`, whereas when reading `Expected` relax testcase the min Range value is being read as `int64(0)`.
Please let me know what I should do with min Range problem.
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1324923990
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
For example, https://github.com/mlc-ai/mlc-llm uses weight quantize / dequantize functions (called `encode / decode`), RMS norm etc implemented in TE without introducing Relax ops for them. They are represented by `call_tir`.
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1335452675
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
Can you show the error?
--
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
[GitHub] [tvm] rutkoor commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "rutkoor (via GitHub)" <gi...@apache.org>.
rutkoor commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1324286476
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
Hi @masahi, currently the elements are padded in the end of the original buffer. So, when we iterate over orig_shape(dimensions of non-padded buffer), we can be sure that we will be iterating over non-padded entries.
The one assumption that I have taken here is that padding will always be happening in the end of the buffer.
In future, if AlterOp introduces padding in-between original buffer then we should modify remove_pad operator to reflect the same.
--
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
[GitHub] [tvm] masahi commented on a diff in pull request #15679: [Unity] Support Padding Reversal in Alter-Op pass
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #15679:
URL: https://github.com/apache/tvm/pull/15679#discussion_r1324914791
##########
python/tvm/relax/transform/legalize_ops/manipulate.py:
##########
@@ -205,3 +213,16 @@ def te_layout_transform(data, name):
output_dtype = call_args[0].struct_info.dtype
output_sinfo = [TensorStructInfo(output_shape, output_dtype)]
return call_tir(gvar, call_args, output_sinfo, tir_vars)
+
+
+@register_legalize("relax.remove_pad")
+def _remove_pad(bb: BlockBuilder, call: Call) -> Expr:
+ orig_shape = call.attrs.orig_shape
+
+ def te_remove_pad(data):
+ """
+ Returns a new compute that restrict the original expression to the shape of orig_shape
+ """
+ return te.compute(orig_shape, data, name="te_remove_pad")
Review Comment:
We shouldn't introduce a half-baked op definition that is meant to be used only by a certain pass. The proper order is that a generic op definition should come first.
--
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