You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/07/01 17:51:12 UTC
[GitHub] [tvm] vinx13 opened a new pull request #8384: [Arith] Inverse affine map
vinx13 opened a new pull request #8384:
URL: https://github.com/apache/tvm/pull/8384
This PR added a utility function to perform inverse transformation for the bijective affine map. This will be helpful for tensorization in the future.
@tqchen
--
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 #8384: [Arith] Inverse affine map
Posted by GitBox <gi...@apache.org>.
masahi merged pull request #8384:
URL: https://github.com/apache/tvm/pull/8384
--
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] vinx13 commented on a change in pull request #8384: [Arith] Inverse affine map
Posted by GitBox <gi...@apache.org>.
vinx13 commented on a change in pull request #8384:
URL: https://github.com/apache/tvm/pull/8384#discussion_r663136293
##########
File path: src/arith/iter_affine_map.cc
##########
@@ -1385,5 +1385,147 @@ TVM_REGISTER_GLOBAL("arith.SubspaceDivide")
return SubspaceDivide(bindings, root_iters, sub_iters, predicate, require_bijective, &ana);
});
+class InverseAffineIterMapTransformer {
+ public:
+ explicit InverseAffineIterMapTransformer(Analyzer* analyzer) : analyzer_(analyzer) {}
+
+ Map<Var, PrimExpr> operator()(const Array<IterSumExpr>& iter_map,
+ const Array<PrimExpr>& outputs) {
+ ICHECK(iter_map.size() == outputs.size());
+ std::vector<const IterMapExprNode*> post_dfs_order = ReverseTopologyOrder(iter_map);
+
+ // initialize back propagation accumulator
+ for (const IterMapExprNode* node : post_dfs_order) {
+ backprop_.Set(GetRef<IterMapExpr>(node), Integer(0));
+ }
+ for (size_t i = 0; i < iter_map.size(); i++) {
+ backprop_.Set(iter_map[i], outputs[i]);
+ }
+
+ // run back propagation
+ for (const IterMapExprNode* node : post_dfs_order) {
+ if (node->IsInstance<IterSumExprNode>()) {
+ Visit_(Downcast<IterSumExpr>(GetRef<IterMapExpr>(node)));
+ } else {
+ ICHECK(node->IsInstance<IterSplitExprNode>());
+ Visit_(Downcast<IterSplitExpr>(GetRef<IterMapExpr>(node)));
+ }
+ }
+ return std::move(inverse_);
+ }
+
+ private:
+ void Visit_(const IterSumExpr& iter_map_expr) {
+ PrimExpr input = backprop_.at(iter_map_expr) - iter_map_expr->base;
+
+ // Case 1: Propagate to the input node directly when the sum expression has only one components
+ if (iter_map_expr->args.size() == 1) {
+ const auto& source = iter_map_expr->args[0];
+ backprop_.Set(source, backprop_.at(source) + input);
+ return;
+ }
+
+ // Case 2: If the sum expression has multiple components, match the fuse pattern and then split
+ // the sum expression for each components.
+ // For example, consider the iterator i1[dom = (0, 16)], i2[dom = (0, 8)], fusing i1 and i2
+ // we will have i1_i2_fused[dom = (0, 64)]. During back propagation, we need to split the
+ // propagated value to get the corresponding components of i1 and i2, which are
+ // floordiv(i1_i2_fused, 8) and floormod(i1_i2_fused, 8), respectively.
+ Array<IterSplitExpr> splits = MatchFusePattern(iter_map_expr);
+ ICHECK(!splits.empty());
+
+ for (const IterSplitExpr& split : splits) {
+ backprop_.Set(split,
+ backprop_.at(split) + floormod(floordiv(input, split->scale), split->extent));
+ }
+ }
+
+ std::vector<const IterMapExprNode*> ReverseTopologyOrder(const Array<IterSumExpr>& iter_map) {
+ std::vector<const IterMapExprNode*> post_dfs_order;
+ std::unordered_map<IterMapExpr, bool, ObjectPtrHash, ObjectPtrEqual> visited;
+
+ std::function<void(const IterMapExpr&)> fvisit = [&](const IterMapExpr& expr) {
+ if (visited[expr]) {
+ return;
+ }
+ visited[expr] = true;
+ if (const auto* sum_expr = expr.as<IterSumExprNode>()) {
+ for (const IterSplitExpr& child : sum_expr->args) {
+ fvisit(child);
+ }
+ } else {
+ const auto* split_expr = expr.as<IterSplitExprNode>();
+ ICHECK(split_expr);
+ if (const auto* source = split_expr->source->source.as<IterMapExprNode>()) {
+ fvisit(GetRef<IterMapExpr>(source));
+ }
+ }
+ post_dfs_order.push_back(expr.get());
+ };
+ for (const IterSumExpr& expr : iter_map) {
+ fvisit(expr);
+ }
+ std::reverse(post_dfs_order.begin(), post_dfs_order.end());
+ return post_dfs_order;
+ }
+
+ void Visit_(const IterSplitExpr& iter_map_expr) {
+ PrimExpr input = backprop_.at(iter_map_expr) * iter_map_expr->lower_factor;
+ const IterMark& source = iter_map_expr->source;
+ if (source->source.as<IterSumExprNode>()) {
+ IterSumExpr source_expr = Downcast<IterSumExpr>(source->source);
+ backprop_.Set(source_expr, backprop_.at(source_expr) + input);
+ } else {
+ Var source_var = Downcast<Var>(source->source);
+ if (inverse_.count(source_var)) {
+ inverse_.Set(source_var, inverse_.at(source_var) + input);
+ } else {
+ inverse_.Set(source_var, input);
+ }
+ }
+ }
+
+ Array<IterSplitExpr> MatchFusePattern(const IterSumExpr sum_expr) {
+ IntImm base_scale(nullptr);
+ size_t base_index = 0;
+ for (size_t i = 0; i < sum_expr->args.size(); ++i) {
+ if (const auto* op = sum_expr->args[i]->scale.as<IntImmNode>()) {
+ if (!base_scale.defined() || op->value < base_scale->value) {
+ base_scale = GetRef<IntImm>(op);
+ base_index = i;
+ }
+ }
+ }
+ ICHECK(base_scale.defined());
+ std::vector<IterSplitExpr> iters;
+ std::vector<bool> visited(sum_expr->args.size(), false);
+ PrimExpr expected_scale = base_scale;
Review comment:
if the sum expression is produced by `DetectIterMap` this is true, can we can simply this whole function, I can add some assertions to check this assumption @tqchen
--
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] spectrometerHBH commented on a change in pull request #8384: [Arith] Inverse affine map
Posted by GitBox <gi...@apache.org>.
spectrometerHBH commented on a change in pull request #8384:
URL: https://github.com/apache/tvm/pull/8384#discussion_r662640827
##########
File path: src/arith/iter_affine_map.cc
##########
@@ -1385,5 +1385,147 @@ TVM_REGISTER_GLOBAL("arith.SubspaceDivide")
return SubspaceDivide(bindings, root_iters, sub_iters, predicate, require_bijective, &ana);
});
+class InverseAffineIterMapTransformer {
+ public:
+ explicit InverseAffineIterMapTransformer(Analyzer* analyzer) : analyzer_(analyzer) {}
+
+ Map<Var, PrimExpr> operator()(const Array<IterSumExpr>& iter_map,
+ const Array<PrimExpr>& outputs) {
+ ICHECK(iter_map.size() == outputs.size());
+ std::vector<const IterMapExprNode*> post_dfs_order = ReverseTopologyOrder(iter_map);
+
+ // initialize back propagation accumulator
+ for (const IterMapExprNode* node : post_dfs_order) {
+ backprop_.Set(GetRef<IterMapExpr>(node), Integer(0));
+ }
+ for (size_t i = 0; i < iter_map.size(); i++) {
+ backprop_.Set(iter_map[i], outputs[i]);
+ }
+
+ // run back propagation
+ for (const IterMapExprNode* node : post_dfs_order) {
+ if (node->IsInstance<IterSumExprNode>()) {
+ Visit_(Downcast<IterSumExpr>(GetRef<IterMapExpr>(node)));
+ } else {
+ ICHECK(node->IsInstance<IterSplitExprNode>());
+ Visit_(Downcast<IterSplitExpr>(GetRef<IterMapExpr>(node)));
+ }
+ }
+ return std::move(inverse_);
+ }
+
+ private:
+ void Visit_(const IterSumExpr& iter_map_expr) {
+ PrimExpr input = backprop_.at(iter_map_expr) - iter_map_expr->base;
+
+ // Case 1: Propagate to the input node directly when the sum expression has only one components
+ if (iter_map_expr->args.size() == 1) {
+ const auto& source = iter_map_expr->args[0];
+ backprop_.Set(source, backprop_.at(source) + input);
+ return;
+ }
+
+ // Case 2: If the sum expression has multiple components, match the fuse pattern and then split
+ // the sum expression for each components.
+ // For example, consider the iterator i1[dom = (0, 16)], i2[dom = (0, 8)], fusing i1 and i2
+ // we will have i1_i2_fused[dom = (0, 64)]. During back propagation, we need to split the
+ // propagated value to get the corresponding components of i1 and i2, which are
+ // floordiv(i1_i2_fused, 8) and floormod(i1_i2_fused, 8), respectively.
+ Array<IterSplitExpr> splits = MatchFusePattern(iter_map_expr);
+ ICHECK(!splits.empty());
+
+ for (const IterSplitExpr& split : splits) {
+ backprop_.Set(split,
+ backprop_.at(split) + floormod(floordiv(input, split->scale), split->extent));
+ }
+ }
+
+ std::vector<const IterMapExprNode*> ReverseTopologyOrder(const Array<IterSumExpr>& iter_map) {
+ std::vector<const IterMapExprNode*> post_dfs_order;
+ std::unordered_map<IterMapExpr, bool, ObjectPtrHash, ObjectPtrEqual> visited;
+
+ std::function<void(const IterMapExpr&)> fvisit = [&](const IterMapExpr& expr) {
+ if (visited[expr]) {
+ return;
+ }
+ visited[expr] = true;
+ if (const auto* sum_expr = expr.as<IterSumExprNode>()) {
+ for (const IterSplitExpr& child : sum_expr->args) {
+ fvisit(child);
+ }
+ } else {
+ const auto* split_expr = expr.as<IterSplitExprNode>();
+ ICHECK(split_expr);
+ if (const auto* source = split_expr->source->source.as<IterMapExprNode>()) {
+ fvisit(GetRef<IterMapExpr>(source));
+ }
+ }
+ post_dfs_order.push_back(expr.get());
+ };
+ for (const IterSumExpr& expr : iter_map) {
+ fvisit(expr);
+ }
+ std::reverse(post_dfs_order.begin(), post_dfs_order.end());
+ return post_dfs_order;
+ }
+
+ void Visit_(const IterSplitExpr& iter_map_expr) {
+ PrimExpr input = backprop_.at(iter_map_expr) * iter_map_expr->lower_factor;
+ const IterMark& source = iter_map_expr->source;
+ if (source->source.as<IterSumExprNode>()) {
+ IterSumExpr source_expr = Downcast<IterSumExpr>(source->source);
+ backprop_.Set(source_expr, backprop_.at(source_expr) + input);
+ } else {
+ Var source_var = Downcast<Var>(source->source);
+ if (inverse_.count(source_var)) {
+ inverse_.Set(source_var, inverse_.at(source_var) + input);
+ } else {
+ inverse_.Set(source_var, input);
+ }
+ }
+ }
+
+ Array<IterSplitExpr> MatchFusePattern(const IterSumExpr sum_expr) {
+ IntImm base_scale(nullptr);
+ size_t base_index = 0;
+ for (size_t i = 0; i < sum_expr->args.size(); ++i) {
+ if (const auto* op = sum_expr->args[i]->scale.as<IntImmNode>()) {
+ if (!base_scale.defined() || op->value < base_scale->value) {
+ base_scale = GetRef<IntImm>(op);
+ base_index = i;
+ }
+ }
+ }
+ ICHECK(base_scale.defined());
+ std::vector<IterSplitExpr> iters;
+ std::vector<bool> visited(sum_expr->args.size(), false);
+ PrimExpr expected_scale = base_scale;
Review comment:
I think we can expect the sum expression has its children in descending order, so to search for base scale is not necessary
##########
File path: tests/python/unittest/test_arith_iter_affine_map.py
##########
@@ -643,6 +643,67 @@ def test_normalize_iter_map_to_expr():
tvm.ir.assert_structural_equal(tvm.arith.normalize_iter_map_to_expr(res[1]), flm(x[0], 5))
+def test_inverse_affine_iter_map():
+ analyzer = tvm.arith.Analyzer()
+ l0 = create_iter("l0", 64)
+ l1 = create_iter("l1", 64)
+ l2 = create_iter("l2", 64)
+
+ # simple case
+ l0_0, l0_1 = isplit(l0, 16)
+ l1_0, l1_1 = isplit(l1, 4)
+ l0_1_l1_1_fused = ifuse([l0_1, l1_1])
+
+ iter_map = tvm.arith.detect_iter_map([l0_1_l1_1_fused[0], l0_0[0], l1_0[0]], var_dom([l0, l1]))
+ outputs = [tvm.tir.Var("output_{}".format(i), "int32") for i in range(len(iter_map))]
+ res = tvm.arith.inverse_affine_iter_map(iter_map, outputs)
+ print(res)
Review comment:
this line can be removed
--
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 change in pull request #8384: [Arith] Inverse affine map
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #8384:
URL: https://github.com/apache/tvm/pull/8384#discussion_r663428706
##########
File path: src/arith/iter_affine_map.cc
##########
@@ -1385,5 +1385,147 @@ TVM_REGISTER_GLOBAL("arith.SubspaceDivide")
return SubspaceDivide(bindings, root_iters, sub_iters, predicate, require_bijective, &ana);
});
+class InverseAffineIterMapTransformer {
+ public:
+ explicit InverseAffineIterMapTransformer(Analyzer* analyzer) : analyzer_(analyzer) {}
+
+ Map<Var, PrimExpr> operator()(const Array<IterSumExpr>& iter_map,
+ const Array<PrimExpr>& outputs) {
+ ICHECK(iter_map.size() == outputs.size());
+ std::vector<const IterMapExprNode*> post_dfs_order = ReverseTopologyOrder(iter_map);
+
+ // initialize back propagation accumulator
+ for (const IterMapExprNode* node : post_dfs_order) {
+ backprop_.Set(GetRef<IterMapExpr>(node), Integer(0));
+ }
+ for (size_t i = 0; i < iter_map.size(); i++) {
+ backprop_.Set(iter_map[i], outputs[i]);
+ }
+
+ // run back propagation
+ for (const IterMapExprNode* node : post_dfs_order) {
+ if (node->IsInstance<IterSumExprNode>()) {
+ Visit_(Downcast<IterSumExpr>(GetRef<IterMapExpr>(node)));
+ } else {
+ ICHECK(node->IsInstance<IterSplitExprNode>());
+ Visit_(Downcast<IterSplitExpr>(GetRef<IterMapExpr>(node)));
+ }
+ }
+ return std::move(inverse_);
+ }
+
+ private:
+ void Visit_(const IterSumExpr& iter_map_expr) {
+ PrimExpr input = backprop_.at(iter_map_expr) - iter_map_expr->base;
+
+ // Case 1: Propagate to the input node directly when the sum expression has only one components
+ if (iter_map_expr->args.size() == 1) {
+ const auto& source = iter_map_expr->args[0];
+ backprop_.Set(source, backprop_.at(source) + input);
+ return;
+ }
+
+ // Case 2: If the sum expression has multiple components, match the fuse pattern and then split
+ // the sum expression for each components.
+ // For example, consider the iterator i1[dom = (0, 16)], i2[dom = (0, 8)], fusing i1 and i2
+ // we will have i1_i2_fused[dom = (0, 64)]. During back propagation, we need to split the
+ // propagated value to get the corresponding components of i1 and i2, which are
+ // floordiv(i1_i2_fused, 8) and floormod(i1_i2_fused, 8), respectively.
+ Array<IterSplitExpr> splits = MatchFusePattern(iter_map_expr);
+ ICHECK(!splits.empty());
+
+ for (const IterSplitExpr& split : splits) {
+ backprop_.Set(split,
+ backprop_.at(split) + floormod(floordiv(input, split->scale), split->extent));
+ }
+ }
+
+ std::vector<const IterMapExprNode*> ReverseTopologyOrder(const Array<IterSumExpr>& iter_map) {
+ std::vector<const IterMapExprNode*> post_dfs_order;
+ std::unordered_map<IterMapExpr, bool, ObjectPtrHash, ObjectPtrEqual> visited;
+
+ std::function<void(const IterMapExpr&)> fvisit = [&](const IterMapExpr& expr) {
+ if (visited[expr]) {
+ return;
+ }
+ visited[expr] = true;
+ if (const auto* sum_expr = expr.as<IterSumExprNode>()) {
+ for (const IterSplitExpr& child : sum_expr->args) {
+ fvisit(child);
+ }
+ } else {
+ const auto* split_expr = expr.as<IterSplitExprNode>();
+ ICHECK(split_expr);
+ if (const auto* source = split_expr->source->source.as<IterMapExprNode>()) {
+ fvisit(GetRef<IterMapExpr>(source));
+ }
+ }
+ post_dfs_order.push_back(expr.get());
+ };
+ for (const IterSumExpr& expr : iter_map) {
+ fvisit(expr);
+ }
+ std::reverse(post_dfs_order.begin(), post_dfs_order.end());
+ return post_dfs_order;
+ }
+
+ void Visit_(const IterSplitExpr& iter_map_expr) {
+ PrimExpr input = backprop_.at(iter_map_expr) * iter_map_expr->lower_factor;
+ const IterMark& source = iter_map_expr->source;
+ if (source->source.as<IterSumExprNode>()) {
+ IterSumExpr source_expr = Downcast<IterSumExpr>(source->source);
+ backprop_.Set(source_expr, backprop_.at(source_expr) + input);
+ } else {
+ Var source_var = Downcast<Var>(source->source);
+ if (inverse_.count(source_var)) {
+ inverse_.Set(source_var, inverse_.at(source_var) + input);
+ } else {
+ inverse_.Set(source_var, input);
+ }
+ }
+ }
+
+ Array<IterSplitExpr> MatchFusePattern(const IterSumExpr sum_expr) {
+ IntImm base_scale(nullptr);
+ size_t base_index = 0;
+ for (size_t i = 0; i < sum_expr->args.size(); ++i) {
+ if (const auto* op = sum_expr->args[i]->scale.as<IntImmNode>()) {
+ if (!base_scale.defined() || op->value < base_scale->value) {
+ base_scale = GetRef<IntImm>(op);
+ base_index = i;
+ }
+ }
+ }
+ ICHECK(base_scale.defined());
+ std::vector<IterSplitExpr> iters;
+ std::vector<bool> visited(sum_expr->args.size(), false);
+ PrimExpr expected_scale = base_scale;
Review comment:
sorry I've just merged it...
--
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] tqchen commented on pull request #8384: [Arith] Inverse affine map
Posted by GitBox <gi...@apache.org>.
tqchen commented on pull request #8384:
URL: https://github.com/apache/tvm/pull/8384#issuecomment-873581831
@vinx13 let us followup with another PR to address the last comment @spectrometerHBH bought up
--
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