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

Defined a common base class for TensorComputeOp and ComputeOp #2587

Merged
merged 3 commits into from
Mar 1, 2019
Merged

Defined a common base class for TensorComputeOp and ComputeOp #2587

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b81a9fb
Defined a common base class for TensorComputeOp and ComputeOp
derisavi Feb 11, 2019
c618e5c
Made changes requested by @ZihengJiang
derisavi Feb 22, 2019
d1de651
added a testcase to assert that `tensorize` does not have any effect …
derisavi Feb 25, 2019
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
59 changes: 31 additions & 28 deletions include/tvm/operation.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,22 +184,45 @@ class PlaceholderOpNode : public OperationNode {

/*!
* \brief A Compute op that compute a tensor on certain domain.
Copy link
Contributor

Choose a reason for hiding this comment

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

stressing that it is the base class of ComputeOp and TensorComputeOp

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.

* This is the base class for ComputeOp (operating on a scalar at a time) and
* TensorComputeOp (operating on a TensorSlice at a time)
*/
class TVM_DLL ComputeOpNode : public OperationNode {
class TVM_DLL BaseComputeOpNode : public OperationNode {
public:
/*! \brief IterVar on each axis */
Array<IterVar> axis;
/*! \brief IterVar on each reduction axis, if the body is a Reduce */
Array<IterVar> reduce_axis;
// override functions
Array<IterVar> root_iter_vars() const final;
Array<Expr> output_shape(size_t idx) const final;
void GatherBound(
const Operation& self,
const std::unordered_map<Tensor, TensorDom>& tensor_dom,
std::unordered_map<IterVar, Range>* out_dom_map) const final;
Copy link
Contributor

Choose a reason for hiding this comment

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

Since ComputeOpNode is an exposed structure, can we keep it as the original, but we have another BasicComputeOpNode; ComputeOpNode and TensorComputeOpNode inherit it. in this way, we can avoid so many modifications, also we can save modifications for those in-house projects based on tvm.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

The way I see it is that TVM has defined a Node class, then derived OperationNode, then derived <X>OpNode, and so on. Therefore, before this PR, one can find out the class hierarchy only by looking at the names of the classes. I believe this was a good decision and I have seen it in other code bases. That's why I thought to keep it that way after this new class is introduced.

Regarding the amount of modifications needed both inside and outside TVM:

  • inside TVM: I agree that quite a few files have been affected. However, the changes are trivial, and have already been done fairly quickly using automatic refactoring tools so no additional effort is needed.
  • outside TVM: even if there are a lot of code out there using ComputeOpNode (which need to be renamed to ScalarComputeOpNode), I don't see why it's an inconvenient thing to do that. Maybe I missing something?
    Moreover, I think keeping the consistency of class naming in TVM (which I hope strives for many years to come) is more important than that level of inconvenience in external projects that depend on it.

That's my two cents.

Copy link
Contributor

@xqdan xqdan Feb 14, 2019 •
edited
Loading

Choose a reason for hiding this comment

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

The way I see it is that TVM has defined a Node class, then derived OperationNode, then derived <X>OpNode, and so on. Therefore, before this PR, one can find out the class hierarchy only by looking at the names of the classes. I believe this was a good decision and I have seen it in other code bases. That's why I thought to keep it that way after this new class is introduced.

agree. what i'm suggesting has no conflicts with this point.

Regarding the amount of modifications needed both inside and outside TVM:

  • inside TVM: I agree that quite a few files have been affected. However, the changes are trivial, and have already been done fairly quickly using automatic refactoring tools so no additional effort is needed.
  • outside TVM: even if there are a lot of code out there using ComputeOpNode (which need to be renamed to ScalarComputeOpNode), I don't see why it's an inconvenient thing to do that. Maybe I missing something?

We do have some custmized pass(C++ and python) and costmized schedule template(python) , in which we referenced ComputeOpNode and the code is not in public, if this is merged, we need to replace them with ScalarComputeOpNode when sync with upstream. I believe other inhouse projects may have the similar issue with this.

Moreover, I think keeping the consistency of class naming in TVM (which I hope strives for many years to come) is more important than that level of inconvenience in external projects that depend on it.

I don't see much differences if we change ComputeOpNode to ScalarComputeOpNode, since everybody who uses ComputeOpNode knew it's a scalar compute, for those who don't know, a comment on ComputeOpNode class will be enough.

That's my two cents.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thanks @xqdan. I followed the naming scheme you recommended. I used BaseComputeOpNode for the base class.

Stmt BuildRealize(
const Stage& stage,
const std::unordered_map<IterVar, Range>& realize_map,
const Stmt& body) const final;
virtual size_t num_schedulable_dims() const = 0;

static constexpr const char* _type_key = "BaseComputeOp";
TVM_DECLARE_BASE_NODE_INFO(BaseComputeOpNode, OperationNode);
};


/*!
* \brief A Compute op that compute a tensor on certain domain.
*/
class TVM_DLL ComputeOpNode : public BaseComputeOpNode {
public:
/*! \brief the compute expression */
Array<Expr> body;
/*! \brief constructor */
ComputeOpNode() {}
// override functions
int num_outputs() const final;
Array<IterVar> root_iter_vars() const final;
Type output_dtype(size_t i) const final;
Array<Expr> output_shape(size_t i) const final;
Array<Tensor> InputTensors() const final;
Operation ReplaceInputs(
const Operation& self,
Expand All @@ -208,18 +231,11 @@ class TVM_DLL ComputeOpNode : public OperationNode {
const Operation& self,
const std::unordered_map<const Variable*, IntSet>& dom_map,
std::unordered_map<Tensor, TensorDom>* out_dom_map) const final;
void GatherBound(
const Operation& self,
const std::unordered_map<Tensor, TensorDom>& tensor_dom,
std::unordered_map<IterVar, Range>* out_dom_map) const final;
Stmt BuildRealize(
const Stage& stage,
const std::unordered_map<IterVar, Range>& realize_map,
const Stmt& body) const final;
Stmt BuildProvide(
const Stage& stage,
const std::unordered_map<IterVar, Range>& dom_map,
bool debug_keep_trivial_loop) const final;
size_t num_schedulable_dims() const final;

void VisitAttrs(AttrVisitor* v) final {
v->Visit("name", &name);
Expand All @@ -236,18 +252,14 @@ class TVM_DLL ComputeOpNode : public OperationNode {
Array<Expr> body);

static constexpr const char* _type_key = "ComputeOp";
TVM_DECLARE_NODE_TYPE_INFO(ComputeOpNode, OperationNode);
TVM_DECLARE_NODE_TYPE_INFO(ComputeOpNode, BaseComputeOpNode);
};

/*!
* \brief A TenorCompute op that compute a tensor with an tensor intrinsic.
*/
class TensorComputeOpNode : public OperationNode {
class TensorComputeOpNode : public BaseComputeOpNode {
public:
/*! \brief IterVar on each axis */
Array<IterVar> axis;
/*! \brief IterVar on each reduction axis, if the intrin will use the reduce axis */
Array<IterVar> reduce_axis;
/*! \brief number of axes that can be scheduled */
int schedulable_ndim;
/*! \brief TensorIntrin used to compute */
Expand All @@ -260,9 +272,7 @@ class TensorComputeOpNode : public OperationNode {
TensorComputeOpNode() {}
// override functions
int num_outputs() const final;
Array<IterVar> root_iter_vars() const final;
Type output_dtype(size_t i) const final;
Array<Expr> output_shape(size_t i) const final;
Array<Tensor> InputTensors() const final;
Operation ReplaceInputs(
const Operation& self,
Expand All @@ -271,18 +281,11 @@ class TensorComputeOpNode : public OperationNode {
const Operation& self,
const std::unordered_map<const Variable*, IntSet>& dom_map,
std::unordered_map<Tensor, TensorDom>* out_dom_map) const final;
void GatherBound(
const Operation& self,
const std::unordered_map<Tensor, TensorDom>& tensor_dom,
std::unordered_map<IterVar, Range>* out_dom_map) const final;
Stmt BuildRealize(
const Stage& stage,
const std::unordered_map<IterVar, Range>& realize_map,
const Stmt& body) const final;
Stmt BuildProvide(
const Stage& stage,
const std::unordered_map<IterVar, Range>& dom_map,
bool debug_keep_trivial_loop) const final;
size_t num_schedulable_dims() const final;

void VisitAttrs(AttrVisitor* v) final {
v->Visit("name", &name);
Expand All @@ -304,7 +307,7 @@ class TensorComputeOpNode : public OperationNode {
Array<Region> regions);

static constexpr const char* _type_key = "TensorComputeOp";
TVM_DECLARE_NODE_TYPE_INFO(TensorComputeOpNode, OperationNode);
TVM_DECLARE_NODE_TYPE_INFO(TensorComputeOpNode, BaseComputeOpNode);
};

/*!
Expand Down
10 changes: 8 additions & 2 deletions python/tvm/tensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ class PlaceholderOp(Operation):


@register_node
class ComputeOp(Operation):
class BaseComputeOp(Operation):
"""Compute operation."""
@property
def axis(self):
Expand All @@ -160,7 +160,13 @@ def reduce_axis(self):


@register_node
class TensorComputeOp(Operation):
class ComputeOp(BaseComputeOp):
"""Scalar operation."""
pass


@register_node
class TensorComputeOp(BaseComputeOp):
"""Tensor operation."""


Expand Down
34 changes: 19 additions & 15 deletions src/op/compute_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ int ComputeOpNode::num_outputs() const {
return body.size();
}

Array<IterVar> ComputeOpNode::root_iter_vars() const {
Array<IterVar> BaseComputeOpNode::root_iter_vars() const {
if (reduce_axis.size() == 0) return axis;
Array<IterVar> ret = axis;
for (IterVar iv : reduce_axis) {
Expand All @@ -54,15 +54,15 @@ Type ComputeOpNode::output_dtype(size_t idx) const {
return body[idx].type();
}

Array<Expr> ComputeOpNode::output_shape(size_t idx) const {
Array<Expr> BaseComputeOpNode::output_shape(size_t idx) const {
CHECK_LT(idx, num_outputs());
// for now, all outputs of ComputeOp have the same shape
std::vector<Expr> shape;
for (size_t i = 0; i < axis.size(); ++i) {
const Range& r = axis[i]->dom;
// for now, all outputs of a BaseComputeOp have the same shape
Array<Expr> shape;
for (const auto& ivar : this->axis) {
const Range& r = ivar->dom;
shape.push_back(r->extent);
}
return Array<Expr>(shape);
return shape;
}

Tensor compute(Array<Expr> shape,
Expand Down Expand Up @@ -208,7 +208,7 @@ void ComputeOpNode::PropBoundToInputs(
for (auto& e : body) ir::PostOrderVisit(e, fvisit);
}

void ComputeOpNode::GatherBound(
void BaseComputeOpNode::GatherBound(
const Operation& self,
const std::unordered_map<Tensor, TensorDom>& tensor_dom,
std::unordered_map<IterVar, Range>* out_dom_map) const {
Expand All @@ -225,22 +225,22 @@ void ComputeOpNode::GatherBound(
}
}

Stmt ComputeOpNode::BuildRealize(
Stmt BaseComputeOpNode::BuildRealize(
const Stage& stage,
const std::unordered_map<IterVar, Range>& realize_map,
const Stmt& realize_body) const {
const Stmt& body) const {
CHECK_EQ(stage->op.get(), this);
HalideIR::Internal::Region bounds;
for (IterVar iv : this->axis) {
bounds.push_back(realize_map.at(iv));
}
Stmt realize = realize_body;
Stmt realize = body;
for (int i = this->num_outputs(); i > 0; --i) {
Tensor t = stage->op.output(i-1);
realize = ir::Realize::make(t->op, t->value_index,
t->dtype, bounds, const_true(), realize);
// alignment requirement, only useful for compute
for (size_t i = 0; i < this->axis.size(); ++i) {
for (size_t i = 0; i < num_schedulable_dims(); ++i) {
auto it = stage->iter_var_attrs.find(this->axis[i]);
if (it != stage->iter_var_attrs.end()) {
IterVarAttr attr = (*it).second;
Expand All @@ -259,6 +259,10 @@ Stmt ComputeOpNode::BuildRealize(
return realize;
}

size_t ComputeOpNode::num_schedulable_dims() const {
return axis.size();
}

// Build a reduction body.
void MakeReduction(const ComputeOpNode* op,
const Array<Tensor>& tensors,
Expand Down Expand Up @@ -414,7 +418,7 @@ Stmt ComputeOpNode::BuildProvide(
}

ComputeLoopNest ComputeLoopNest::make(
const ComputeOpNode* self,
const BaseComputeOpNode* self,
const Stage& stage,
const std::unordered_map<IterVar, Range>& dom_map,
bool debug_keep_trivial_loop) {
Expand All @@ -440,8 +444,8 @@ ComputeLoopNest ComputeLoopNest::make(
for (IterVar iv : self->reduce_axis) {
update_state[iv] = 2;
}
for (IterVar iv : self->axis) {
update_state[iv] = 1;
for (size_t i = 0; i < self->num_schedulable_dims(); ++i) {
update_state[self->axis[i]] = 1;
}
// find which iter var is related to reduction and which is related to axis.
schedule::PassDownBitMaskOr(stage, &update_state);
Expand Down
2 changes: 1 addition & 1 deletion src/op/compute_op.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ struct ComputeLoopNest {
* \return The constructed loop nest
*/
static ComputeLoopNest make(
const ComputeOpNode* self,
const BaseComputeOpNode* self,
const Stage& stage,
const std::unordered_map<IterVar, Range>& dom_map,
bool debug_keep_trivial_loop);
Expand Down
Loading