Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[SCHEDULE] New Reduction Mode for Tensorize #727

Merged
merged 5 commits into from
Dec 27, 2017
Merged

[SCHEDULE] New Reduction Mode for Tensorize #727

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e386176
when there is no intrin func, using body for initialization. For issu…
kun-zh Dec 25, 2017
1151fd6
Merge branch 'master' into master
kun-zh Dec 26, 2017
784a0fd
Refine code per review comments, and add a test case.
kun-zh Dec 26, 2017
5bbc941
Merge branch 'master' of https://github.com/kun-zh/tvm
kun-zh Dec 26, 2017
ab49e9f
Fix lint issues.
kun-zh Dec 27, 2017
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions src/op/op_util.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -208,5 +208,31 @@ Stmt Substitute(Stmt s,
return ir::Substitute(s, init);
}

Stmt TransformUpdate(const Stage& stage,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This function is specific to Tensorize, move it to tensorize.cc, no need to keep comment in header file, move comment document to tensorize.cc as well

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done, move the function into tensorize.cc

const std::unordered_map<IterVar, Range>& dom_map,
Stmt body,
Stmt update)
{
Expr condition = (make_zero(Int(32)) == make_zero(Int(32))); //will try Bool(1)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

build a Array<Expr> containing condition, use https://github.com/dmlc/tvm/blob/master/src/arithmetic/compute_expr.h#L42 to do reduction

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done, please check

for (size_t i = 0; i < stage->leaf_iter_vars.size(); ++i) {
IterVar iv = stage->leaf_iter_vars[i];
auto iit = stage->iter_var_attrs.find(iv);
if (iit != stage->iter_var_attrs.end()) {
const IterVarAttr& attr = (*iit).second;
if (attr->iter_type == kTensorized) {
break;
}
}
if (iv->iter_type == kCommReduce) {
auto vit = dom_map.find(iv);
CHECK(vit != dom_map.end());
const Range& vrange = vit->second;
Expr newcond = ( iv->var == vrange->min);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

use condition

if (likely(var1> min1) or likely(var2 > min2)) {
   update
} else {
   body
}

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

condition = condition && newcond;
}
}
return IfThenElse::make(condition, body, update);
}

} // namespace op
} // namespace tvm
14 changes: 14 additions & 0 deletions src/op/op_util.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,20 @@ Expr ReplaceTensor(Expr expr,
Stmt Substitute(Stmt stmt,
const std::unordered_map<IterVar, Expr>& value_map);


/*!
* \brief Transform the update part when there is no init func in tensorizing
* \param stage The stage for tensorizing.
* \param dom_map The range of each iter var.
* \param body The body func in tensorize intrin
* \param update The update func in tensorize intrin
* \return Transformed result.
*/
Stmt TransformUpdate(const Stage& stage,
const std::unordered_map<IterVar, Range>& dom_map,
Stmt body,
Stmt update);

} // namespace op
} // namespace tvm
#endif // TVM_OP_OP_UTIL_H_
50 changes: 33 additions & 17 deletions src/op/tensorize.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -416,32 +416,48 @@ Stmt MakeTensorize(const ComputeOpNode* self,
return MergeNest(nest, body);
} else {
// Need to split reduction
CHECK(intrin->reduce_init.defined())
<< "Reduction init op for intrin " << intrin << " is not defined";
// Comment out the following check for the case when there is no init func
//CHECK(intrin->reduce_init.defined())
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove lgtm, simply remove them

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

// << "Reduction init op for intrin " << intrin << " is not defined";
CHECK(intrin->reduce_update.defined())
<< "Reduction update op for intrin " << intrin << " is not defined";
// Need init and update steps
CHECK_NE(self->reduce_axis.size(), 0U);
std::vector<std::vector<Stmt> > common(
n.main_nest.begin(), n.main_nest.begin() + n.num_common_loop + 1);
// init nest
std::vector<std::vector<Stmt> > init_nest(
n.init_nest.begin(), n.init_nest.begin() + tloc + 1);
init_nest.emplace_back(op::MakeIfNest(n.init_predicates));
Stmt init = MergeNest(output_bind_nest, intrin->reduce_init);
init = Substitute(init, n.init_vmap);
init = MergeNest(init_nest, init);
// The update
std::vector<std::vector<Stmt> > update_nest(
n.main_nest.begin() + n.num_common_loop + 1, n.main_nest.begin() + tloc + 1);
update_nest.emplace_back(op::MakeIfNest(n.main_predicates));
Stmt update = MergeNest(output_bind_nest, intrin->reduce_update);
update = MergeNest(input_bind_nest, update);
update = Substitute(update, vmap);
update = MergeNest(binder.asserts(), update);
update = Substitute(update, n.main_vmap);
update = MergeNest(update_nest, update);
return MergeNest(common, Block::make(init, update));

if (intrin->reduce_init.defined()) {
// init nest
std::vector<std::vector<Stmt> > init_nest(
n.init_nest.begin(), n.init_nest.begin() + tloc + 1);
init_nest.emplace_back(op::MakeIfNest(n.init_predicates));
Stmt init = MergeNest(output_bind_nest, intrin->reduce_init);
init = Substitute(init, n.init_vmap);
init = MergeNest(init_nest, init);
// The update
Stmt update = MergeNest(output_bind_nest, intrin->reduce_update);
update = MergeNest(input_bind_nest, update);
update = Substitute(update, vmap);
update = MergeNest(binder.asserts(), update);
update = Substitute(update, n.main_vmap);
update = MergeNest(update_nest, update);
return MergeNest(common, Block::make(init, update));
} else {
// The update
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add a comment about what this logic is about. Something like

// when init intrinsic is not available, transform to reset in the first iter.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

Stmt update = TransformUpdate (stage, dom_map,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add CHECK(body.defined());

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

intrin->body,
intrin->reduce_update);
update = MergeNest(output_bind_nest, update);
update = MergeNest(input_bind_nest, update);
update = Substitute(update, vmap);
update = MergeNest(binder.asserts(), update);
update = Substitute(update, n.main_vmap);
update = MergeNest(update_nest, update);
return MergeNest(common, update);
}
}
}

Expand Down