Skip to content

Commit 67b14a2

Browse files
authored
[PASS] StorageRewrite Fold Inplace op storage when possible (#759)
* [PASS] StorageRewrite Fold Inplace op storage when possible * update comment to fix typos
1 parent d899f86 commit 67b14a2

File tree

5 files changed

+325
-75
lines changed

5 files changed

+325
-75
lines changed

include/tvm/ir.h

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,12 @@ constexpr const char* coproc_scope = "coproc_scope";
152152
constexpr const char* coproc_uop_scope = "coproc_uop_scope";
153153
/*! \brief Mark the scope as volatile access for certain handle. */
154154
constexpr const char* volatile_scope = "volatile_scope";
155+
/*!
156+
* \brief Mark the scope as generated by extern primitive.
157+
* such scope can contain arbitrary ir program and we need to be careful
158+
* when make certain assumptions about the structure of the program.
159+
*/
160+
constexpr const char* extern_scope = "extern_scope";
155161
/*!
156162
* \brief Mark the scope as when computation start to happen
157163
* This can hint some code generator to create a new function for compute.

src/op/extern_op.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ Stmt ExternOpNode::BuildProvide(
130130
const Stage& stage,
131131
const std::unordered_map<IterVar, Range>& dom_map) const {
132132
CHECK_EQ(stage->op.operator->(), this);
133-
Stmt ret = this->body;
133+
Stmt ret = AttrStmt::make(make_zero(Int(32)), attr::extern_scope, 0, this->body);
134134
auto f_push_bind = [&ret](Buffer buffer, Tensor tensor) {
135135
Array<NodeRef> bind_spec;
136136
Array<Expr> tuple;

0 commit comments

Comments
 (0)