[Relay] Handle memory scope during lowering from relay level#11874

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

Merged
junrushao merged 2 commits into from
Jun 29, 2022
Merged

[Relay] Handle memory scope during lowering from relay level #11874

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7a4e763
[Relay] Handle memory scope during lowering from relay level
elvin-n Jun 24, 2022
a46f8c4
Move BufferWithOffsetAlignment from driver_api to tir buffer
elvin-n Jun 29, 2022
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
1 change: 0 additions & 1 deletion include/tvm/driver/driver_api.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -165,7 +165,6 @@ TVM_DLL runtime::Module build(const Map<Target, IRModule>& input, const Target&
* \return The built module that contains code for different processors.
*/
TVM_DLL runtime::Module build(const Map<String, IRModule>& input, const Target& target_host);

} // namespace tvm

#endif // TVM_DRIVER_DRIVER_API_H_
3 changes: 2 additions & 1 deletion include/tvm/relay/expr.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -218,7 +218,8 @@ class VarNode : public ExprNode {

bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
equal->MarkGraphNode();
return equal(type_annotation, other->type_annotation) && equal(vid, other->vid);
return equal(type_annotation, other->type_annotation) && equal(vid, other->vid) &&
equal(virtual_device_, other->virtual_device_);
}

void SHashReduce(SHashReducer hash_reduce) const {
Expand Down
17 changes: 17 additions & 0 deletions include/tvm/tir/buffer.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -295,6 +295,23 @@ class DataProducer : public ObjectRef {
TVM_DEFINE_OBJECT_REF_METHODS(DataProducer, ObjectRef, DataProducerNode);
};

/*!
* \brief Creates TIR Buffer for provided parameters
* \param shape shape of the buffer
* \param dtype data type
* \param name buffer name
* \param data_alignment alignment requirement of data pointer in bytes
* \param offset_factor Factor of elem_offset field, elem_offset is guaranteed to be
* multiple of offset_factor
User can specify data_alignment and offset_factor to be 0
* A default value will be picked.
* \param compact If the statement has already bound to a compact buffer.
* \param memory_scope memory scope of the buffer
*/
TVM_DLL tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype,
std::string name, int data_alignment,
int offset_factor, bool compact,
std::string memory_scope = "");
} // namespace tir
} // namespace tvm
#endif // TVM_TIR_BUFFER_H_
30 changes: 2 additions & 28 deletions src/driver/driver_api.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -83,32 +83,6 @@ Target DefaultTargetHost(Target target) {
}
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
int data_alignment, int offset_factor, bool compact) {
DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
auto data = tir::Var(name, PointerType(PrimType(storage_dtype)));
bool has_any = false;
if (!compact) {
for (const auto& it : shape) {
if (it.as<tir::VarNode>()) {
has_any = true;
break;
}
}
}
tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

PrimExpr elem_offset;
if (offset_factor != 0) {
elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
} else {
elem_offset = PrimExpr();
}

return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
offset_factor, buffer_type);
}

void GetBinds(const Array<ObjectRef>& args, bool compact,
const std::unordered_map<te::Tensor, tir::Buffer>& binds,
Map<te::Tensor, tir::Buffer>* out_binds, Array<ObjectRef>* out_arg_list) {
Expand All@@ -118,8 +92,8 @@ void GetBinds(const Array<ObjectRef>& args, bool compact,
if (const te::TensorNode* tensor_node = x.as<te::TensorNode>()) {
te::Tensor x_ref = GetRef<te::Tensor>(tensor_node);
if (out_binds->find(x_ref) == out_binds->end()) {
tir::Buffer buf =
BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0, compact);
tir::Buffer buf = tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype,
x_ref->op->name, -1, 0, compact);
out_binds->Set(x_ref, buf);
out_arg_list->push_back(buf);
} else {
Expand Down
30 changes: 29 additions & 1 deletion src/relay/backend/te_compiler.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -414,6 +414,33 @@ class TECompilerImpl : public TECompilerNode {
}
// lower the function
std::unordered_map<te::Tensor, tir::Buffer> binds;

// If we have memory scopes, need to create tir::Buffer knowing this info
size_t i = 0; // for corresponding from tensor array
for (Var param : key->source_func->params) {
if (!param->virtual_device()->memory_scope.empty()) {
for (const auto& ttype : FlattenTupleType(param->checked_type())) {
te::Tensor x_ref = value->cached_func->inputs[i];
// verification if we have synced params and tensors
ICHECK(ttype->dtype == x_ref->dtype && ttype->shape.size() == x_ref->shape.size())
<< "function parameter does not correspond to prepared tensor";
binds[x_ref] =
tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
false, param->virtual_device()->memory_scope);
}
}
i++;
}
if (key->virtual_device != VirtualDevice::FullyUnconstrained() &&
!key->virtual_device->memory_scope.empty() &&
key->virtual_device->memory_scope != "global") {
ICHECK(value->cached_func->outputs.size() == 1)
<< "Expect only one output for defined memory scope";
te::Tensor x_ref = value->cached_func->outputs[0];
binds[x_ref] =
tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
false, key->virtual_device->memory_scope);
}
auto func_name = value->cached_func->prim_fn_var->name_hint;
VLOG(1) << "scheduling";
IRModule scheduled_module =
Expand DownExpand Up@@ -895,7 +922,8 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
} else {
// Cases 1 and 2: lower the primitive function for the desired target, possibly using external
// codegen.
CCacheKey key(Downcast<Function>(primitive_func), target);
CCacheKey key(Downcast<Function>(primitive_func), target,
GetVirtualDevice(GetRef<Call>(call_node)));
CachedFunc cfunc = compiler_->Lower(key, module_name_);
ICHECK(cfunc.defined());
return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
Expand Down
3 changes: 2 additions & 1 deletion src/relay/backend/te_compiler_cache.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,11 @@ LoweredOutput::LoweredOutput(tvm::Array<te::Tensor> outputs, OpImplementation im
data_ = std::move(n);
}

CCacheKey::CCacheKey(Function source_func, Target target) {
CCacheKey::CCacheKey(Function source_func, Target target, VirtualDevice vd) {
auto n = make_object<CCacheKeyNode>();
n->source_func = std::move(source_func);
n->target = std::move(target);
n->virtual_device = std::move(vd);
data_ = std::move(n);
}

Expand Down
7 changes: 6 additions & 1 deletion src/relay/backend/te_compiler_cache.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -82,10 +82,13 @@ class CCacheKeyNode : public Object {
Function source_func;
/*! \brief The hardware target.*/
Target target;
/*! \brief The virtual device constrains.*/
VirtualDevice virtual_device;

void VisitAttrs(tvm::AttrVisitor* v) {
v->Visit("source_func", &source_func);
v->Visit("target", &target);
v->Visit("virtual_device", &virtual_device);
}
/*! \return The hash value of CCacheKey. */
inline size_t Hash() const;
Expand DownExpand Up@@ -117,7 +120,8 @@ class CCacheKey : public ObjectRef {
* \param source_func The source function.
* \param target The target device.
*/
TVM_DLL CCacheKey(Function source_func, Target target);
TVM_DLL CCacheKey(Function source_func, Target target,
VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained());

const CCacheKeyNode* operator->() const { return static_cast<const CCacheKeyNode*>(get()); }
// comparator
Expand DownExpand Up@@ -244,6 +248,7 @@ inline size_t CCacheKeyNode::Hash() const {
inline bool CCacheKeyNode::Equal(const CCacheKeyNode* other) const {
if (Hash() != other->Hash()) return false;
return this->target->str() == other->target->str() &&
this->virtual_device == other->virtual_device &&
tvm::StructuralEqual()(this->source_func, other->source_func);
}

Expand Down
27 changes: 27 additions & 0 deletions src/tir/ir/buffer.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -585,6 +585,33 @@ Buffer::Buffer(Var data, DataType dtype, Array<PrimExpr> shape, Array<PrimExpr>
data_ = std::move(n);
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
int data_alignment, int offset_factor, bool compact,
std::string memory_scope) {
DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
auto data = tir::Var(name, PointerType(PrimType(storage_dtype), memory_scope));
bool has_any = false;
if (!compact) {
for (const auto& it : shape) {
if (it.as<tir::VarNode>()) {
has_any = true;
break;
}
}
}
tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

PrimExpr elem_offset;
if (offset_factor != 0) {
elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
} else {
elem_offset = PrimExpr();
}

return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
offset_factor, buffer_type);
}

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
.set_dispatch<BufferNode>([](const ObjectRef& node, ReprPrinter* p) {
auto* op = static_cast<const BufferNode*>(node.get());
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
[Relay] Handle memory scope during lowering from relay level#11874

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

 Merged

junrushao
merged 2 commits into

from

Jun 29, 2022
 Merged


[Relay] Handle memory scope during lowering from relay level
#11874
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
7a4e763

[Relay] Handle memory scope during lowering from relay level


elvin-n Jun 24, 2022
a46f8c4

Move BufferWithOffsetAlignment from driver_api to tir buffer


elvin-n Jun 29, 2022

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


1 change: 0 additions & 1 deletion
include/tvm/driver/driver_api.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -165,7 +165,6 @@ TVM_DLL runtime::Module build(const Map<Target, IRModule>& input, const Target&
 * \return The built module that contains code for different processors.
 */
TVM_DLL runtime::Module build(const Map<String, IRModule>& input, const Target& target_host);

} // namespace tvm

#endif // TVM_DRIVER_DRIVER_API_H_

3 changes: 2 additions & 1 deletion
include/tvm/relay/expr.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -218,7 +218,8 @@ class VarNode : public ExprNode {

 bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
 equal->MarkGraphNode();
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid);
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid) &&
 equal(virtual_device_, other->virtual_device_);
 }

 void SHashReduce(SHashReducer hash_reduce) const {
Expand Down

17 changes: 17 additions & 0 deletions
include/tvm/tir/buffer.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -295,6 +295,23 @@ class DataProducer : public ObjectRef {
 TVM_DEFINE_OBJECT_REF_METHODS(DataProducer, ObjectRef, DataProducerNode);
};

/*!
 * \brief Creates TIR Buffer for provided parameters
 * \param shape shape of the buffer
 * \param dtype data type
 * \param name buffer name
 * \param data_alignment alignment requirement of data pointer in bytes
 * \param offset_factor Factor of elem_offset field, elem_offset is guaranteed to be
 * multiple of offset_factor
 User can specify data_alignment and offset_factor to be 0
 * A default value will be picked.
 * \param compact If the statement has already bound to a compact buffer.
 * \param memory_scope memory scope of the buffer
 */
TVM_DLL tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype,
 std::string name, int data_alignment,
 int offset_factor, bool compact,
 std::string memory_scope = "");
} // namespace tir
} // namespace tvm
#endif // TVM_TIR_BUFFER_H_

30 changes: 2 additions & 28 deletions
src/driver/driver_api.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -83,32 +83,6 @@ Target DefaultTargetHost(Target target) {
 }
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype)));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

void GetBinds(const Array<ObjectRef>& args, bool compact,
 const std::unordered_map<te::Tensor, tir::Buffer>& binds,
 Map<te::Tensor, tir::Buffer>* out_binds, Array<ObjectRef>* out_arg_list) {
Expand All@@ -118,8 +92,8 @@ void GetBinds(const Array<ObjectRef>& args, bool compact,
 if (const te::TensorNode* tensor_node = x.as<te::TensorNode>()) {
 te::Tensor x_ref = GetRef<te::Tensor>(tensor_node);
 if (out_binds->find(x_ref) == out_binds->end()) {
 tir::Buffer buf =
BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0, compact);
 tir::Buffer buf = tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype,
 x_ref->op->name, -1, 0, compact);
 out_binds->Set(x_ref, buf);
 out_arg_list->push_back(buf);
 } else {
Expand Down

30 changes: 29 additions & 1 deletion
src/relay/backend/te_compiler.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -414,6 +414,33 @@ class TECompilerImpl : public TECompilerNode {
 }
 // lower the function
 std::unordered_map<te::Tensor, tir::Buffer> binds;

 // If we have memory scopes, need to create tir::Buffer knowing this info
 size_t i = 0; // for corresponding from tensor array
 for (Var param : key->source_func->params) {
 if (!param->virtual_device()->memory_scope.empty()) {
 for (const auto& ttype : FlattenTupleType(param->checked_type())) {
 te::Tensor x_ref = value->cached_func->inputs[i];
 // verification if we have synced params and tensors
 ICHECK(ttype->dtype == x_ref->dtype && ttype->shape.size() == x_ref->shape.size())
 << "function parameter does not correspond to prepared tensor";
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, param->virtual_device()->memory_scope);
 }
 }
 i++;
 }
 if (key->virtual_device != VirtualDevice::FullyUnconstrained() &&
 !key->virtual_device->memory_scope.empty() &&
 key->virtual_device->memory_scope != "global") {
 ICHECK(value->cached_func->outputs.size() == 1)
 << "Expect only one output for defined memory scope";
 te::Tensor x_ref = value->cached_func->outputs[0];
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, key->virtual_device->memory_scope);
 }
 auto func_name = value->cached_func->prim_fn_var->name_hint;
 VLOG(1) << "scheduling";
 IRModule scheduled_module =
Expand DownExpand Up@@ -895,7 +922,8 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
 } else {
 // Cases 1 and 2: lower the primitive function for the desired target, possibly using external
 // codegen.
 CCacheKey key(Downcast<Function>(primitive_func), target);
 CCacheKey key(Downcast<Function>(primitive_func), target,
 GetVirtualDevice(GetRef<Call>(call_node)));
 CachedFunc cfunc = compiler_->Lower(key, module_name_);
 ICHECK(cfunc.defined());
 return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
Expand Down

3 changes: 2 additions & 1 deletion
src/relay/backend/te_compiler_cache.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,11 @@ LoweredOutput::LoweredOutput(tvm::Array<te::Tensor> outputs, OpImplementation im
 data_ = std::move(n);
}

CCacheKey::CCacheKey(Function source_func, Target target) {
CCacheKey::CCacheKey(Function source_func, Target target, VirtualDevice vd) {
 auto n = make_object<CCacheKeyNode>();
 n->source_func = std::move(source_func);
 n->target = std::move(target);
 n->virtual_device = std::move(vd);
 data_ = std::move(n);
}

Expand Down

7 changes: 6 additions & 1 deletion
src/relay/backend/te_compiler_cache.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -82,10 +82,13 @@ class CCacheKeyNode : public Object {
 Function source_func;
 /*! \brief The hardware target.*/
 Target target;
 /*! \brief The virtual device constrains.*/
 VirtualDevice virtual_device;

 void VisitAttrs(tvm::AttrVisitor* v) {
 v->Visit("source_func", &source_func);
 v->Visit("target", &target);
 v->Visit("virtual_device", &virtual_device);
 }
 /*! \return The hash value of CCacheKey. */
 inline size_t Hash() const;
Expand DownExpand Up@@ -117,7 +120,8 @@ class CCacheKey : public ObjectRef {
 * \param source_func The source function.
 * \param target The target device.
 */
 TVM_DLL CCacheKey(Function source_func, Target target);
 TVM_DLL CCacheKey(Function source_func, Target target,
 VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained());

 const CCacheKeyNode* operator->() const { return static_cast<const CCacheKeyNode*>(get()); }
 // comparator
Expand DownExpand Up@@ -244,6 +248,7 @@ inline size_t CCacheKeyNode::Hash() const {
inline bool CCacheKeyNode::Equal(const CCacheKeyNode* other) const {
 if (Hash() != other->Hash()) return false;
 return this->target->str() == other->target->str() &&
 this->virtual_device == other->virtual_device &&
 tvm::StructuralEqual()(this->source_func, other->source_func);
}

Expand Down

27 changes: 27 additions & 0 deletions
src/tir/ir/buffer.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -585,6 +585,33 @@ Buffer::Buffer(Var data, DataType dtype, Array<PrimExpr> shape, Array<PrimExpr>
 data_ = std::move(n);
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact,
 std::string memory_scope) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype), memory_scope));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 .set_dispatch<BufferNode>([](const ObjectRef& node, ReprPrinter* p) {
 auto* op = static_cast<const BufferNode*>(node.get());
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[Relay] Handle memory scope during lowering from relay level#11874

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

 Merged

junrushao
merged 2 commits into

from

Jun 29, 2022
 Merged


[Relay] Handle memory scope during lowering from relay level
#11874
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
7a4e763

[Relay] Handle memory scope during lowering from relay level


elvin-n Jun 24, 2022
a46f8c4

Move BufferWithOffsetAlignment from driver_api to tir buffer


elvin-n Jun 29, 2022

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


1 change: 0 additions & 1 deletion
include/tvm/driver/driver_api.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -165,7 +165,6 @@ TVM_DLL runtime::Module build(const Map<Target, IRModule>& input, const Target&
 * \return The built module that contains code for different processors.
 */
TVM_DLL runtime::Module build(const Map<String, IRModule>& input, const Target& target_host);

} // namespace tvm

#endif // TVM_DRIVER_DRIVER_API_H_

3 changes: 2 additions & 1 deletion
include/tvm/relay/expr.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -218,7 +218,8 @@ class VarNode : public ExprNode {

 bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
 equal->MarkGraphNode();
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid);
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid) &&
 equal(virtual_device_, other->virtual_device_);
 }

 void SHashReduce(SHashReducer hash_reduce) const {
Expand Down

17 changes: 17 additions & 0 deletions
include/tvm/tir/buffer.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -295,6 +295,23 @@ class DataProducer : public ObjectRef {
 TVM_DEFINE_OBJECT_REF_METHODS(DataProducer, ObjectRef, DataProducerNode);
};

/*!
 * \brief Creates TIR Buffer for provided parameters
 * \param shape shape of the buffer
 * \param dtype data type
 * \param name buffer name
 * \param data_alignment alignment requirement of data pointer in bytes
 * \param offset_factor Factor of elem_offset field, elem_offset is guaranteed to be
 * multiple of offset_factor
 User can specify data_alignment and offset_factor to be 0
 * A default value will be picked.
 * \param compact If the statement has already bound to a compact buffer.
 * \param memory_scope memory scope of the buffer
 */
TVM_DLL tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype,
 std::string name, int data_alignment,
 int offset_factor, bool compact,
 std::string memory_scope = "");
} // namespace tir
} // namespace tvm
#endif // TVM_TIR_BUFFER_H_

30 changes: 2 additions & 28 deletions
src/driver/driver_api.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -83,32 +83,6 @@ Target DefaultTargetHost(Target target) {
 }
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype)));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

void GetBinds(const Array<ObjectRef>& args, bool compact,
 const std::unordered_map<te::Tensor, tir::Buffer>& binds,
 Map<te::Tensor, tir::Buffer>* out_binds, Array<ObjectRef>* out_arg_list) {
Expand All@@ -118,8 +92,8 @@ void GetBinds(const Array<ObjectRef>& args, bool compact,
 if (const te::TensorNode* tensor_node = x.as<te::TensorNode>()) {
 te::Tensor x_ref = GetRef<te::Tensor>(tensor_node);
 if (out_binds->find(x_ref) == out_binds->end()) {
 tir::Buffer buf =
BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0, compact);
 tir::Buffer buf = tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype,
 x_ref->op->name, -1, 0, compact);
 out_binds->Set(x_ref, buf);
 out_arg_list->push_back(buf);
 } else {
Expand Down

30 changes: 29 additions & 1 deletion
src/relay/backend/te_compiler.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -414,6 +414,33 @@ class TECompilerImpl : public TECompilerNode {
 }
 // lower the function
 std::unordered_map<te::Tensor, tir::Buffer> binds;

 // If we have memory scopes, need to create tir::Buffer knowing this info
 size_t i = 0; // for corresponding from tensor array
 for (Var param : key->source_func->params) {
 if (!param->virtual_device()->memory_scope.empty()) {
 for (const auto& ttype : FlattenTupleType(param->checked_type())) {
 te::Tensor x_ref = value->cached_func->inputs[i];
 // verification if we have synced params and tensors
 ICHECK(ttype->dtype == x_ref->dtype && ttype->shape.size() == x_ref->shape.size())
 << "function parameter does not correspond to prepared tensor";
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, param->virtual_device()->memory_scope);
 }
 }
 i++;
 }
 if (key->virtual_device != VirtualDevice::FullyUnconstrained() &&
 !key->virtual_device->memory_scope.empty() &&
 key->virtual_device->memory_scope != "global") {
 ICHECK(value->cached_func->outputs.size() == 1)
 << "Expect only one output for defined memory scope";
 te::Tensor x_ref = value->cached_func->outputs[0];
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, key->virtual_device->memory_scope);
 }
 auto func_name = value->cached_func->prim_fn_var->name_hint;
 VLOG(1) << "scheduling";
 IRModule scheduled_module =
Expand DownExpand Up@@ -895,7 +922,8 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
 } else {
 // Cases 1 and 2: lower the primitive function for the desired target, possibly using external
 // codegen.
 CCacheKey key(Downcast<Function>(primitive_func), target);
 CCacheKey key(Downcast<Function>(primitive_func), target,
 GetVirtualDevice(GetRef<Call>(call_node)));
 CachedFunc cfunc = compiler_->Lower(key, module_name_);
 ICHECK(cfunc.defined());
 return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
Expand Down

3 changes: 2 additions & 1 deletion
src/relay/backend/te_compiler_cache.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,11 @@ LoweredOutput::LoweredOutput(tvm::Array<te::Tensor> outputs, OpImplementation im
 data_ = std::move(n);
}

CCacheKey::CCacheKey(Function source_func, Target target) {
CCacheKey::CCacheKey(Function source_func, Target target, VirtualDevice vd) {
 auto n = make_object<CCacheKeyNode>();
 n->source_func = std::move(source_func);
 n->target = std::move(target);
 n->virtual_device = std::move(vd);
 data_ = std::move(n);
}

Expand Down

7 changes: 6 additions & 1 deletion
src/relay/backend/te_compiler_cache.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -82,10 +82,13 @@ class CCacheKeyNode : public Object {
 Function source_func;
 /*! \brief The hardware target.*/
 Target target;
 /*! \brief The virtual device constrains.*/
 VirtualDevice virtual_device;

 void VisitAttrs(tvm::AttrVisitor* v) {
 v->Visit("source_func", &source_func);
 v->Visit("target", &target);
 v->Visit("virtual_device", &virtual_device);
 }
 /*! \return The hash value of CCacheKey. */
 inline size_t Hash() const;
Expand DownExpand Up@@ -117,7 +120,8 @@ class CCacheKey : public ObjectRef {
 * \param source_func The source function.
 * \param target The target device.
 */
 TVM_DLL CCacheKey(Function source_func, Target target);
 TVM_DLL CCacheKey(Function source_func, Target target,
 VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained());

 const CCacheKeyNode* operator->() const { return static_cast<const CCacheKeyNode*>(get()); }
 // comparator
Expand DownExpand Up@@ -244,6 +248,7 @@ inline size_t CCacheKeyNode::Hash() const {
inline bool CCacheKeyNode::Equal(const CCacheKeyNode* other) const {
 if (Hash() != other->Hash()) return false;
 return this->target->str() == other->target->str() &&
 this->virtual_device == other->virtual_device &&
 tvm::StructuralEqual()(this->source_func, other->source_func);
}

Expand Down

27 changes: 27 additions & 0 deletions
src/tir/ir/buffer.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -585,6 +585,33 @@ Buffer::Buffer(Var data, DataType dtype, Array<PrimExpr> shape, Array<PrimExpr>
 data_ = std::move(n);
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact,
 std::string memory_scope) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype), memory_scope));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 .set_dispatch<BufferNode>([](const ObjectRef& node, ReprPrinter* p) {
 auto* op = static_cast<const BufferNode*>(node.get());
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[Relay] Handle memory scope during lowering from relay level#11874

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

 Merged

junrushao
merged 2 commits into

from

Jun 29, 2022
 Merged


[Relay] Handle memory scope during lowering from relay level
#11874
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
7a4e763

[Relay] Handle memory scope during lowering from relay level


elvin-n Jun 24, 2022
a46f8c4

Move BufferWithOffsetAlignment from driver_api to tir buffer


elvin-n Jun 29, 2022

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


1 change: 0 additions & 1 deletion
include/tvm/driver/driver_api.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -165,7 +165,6 @@ TVM_DLL runtime::Module build(const Map<Target, IRModule>& input, const Target&
 * \return The built module that contains code for different processors.
 */
TVM_DLL runtime::Module build(const Map<String, IRModule>& input, const Target& target_host);

} // namespace tvm

#endif // TVM_DRIVER_DRIVER_API_H_

3 changes: 2 additions & 1 deletion
include/tvm/relay/expr.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -218,7 +218,8 @@ class VarNode : public ExprNode {

 bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
 equal->MarkGraphNode();
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid);
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid) &&
 equal(virtual_device_, other->virtual_device_);
 }

 void SHashReduce(SHashReducer hash_reduce) const {
Expand Down

17 changes: 17 additions & 0 deletions
include/tvm/tir/buffer.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -295,6 +295,23 @@ class DataProducer : public ObjectRef {
 TVM_DEFINE_OBJECT_REF_METHODS(DataProducer, ObjectRef, DataProducerNode);
};

/*!
 * \brief Creates TIR Buffer for provided parameters
 * \param shape shape of the buffer
 * \param dtype data type
 * \param name buffer name
 * \param data_alignment alignment requirement of data pointer in bytes
 * \param offset_factor Factor of elem_offset field, elem_offset is guaranteed to be
 * multiple of offset_factor
 User can specify data_alignment and offset_factor to be 0
 * A default value will be picked.
 * \param compact If the statement has already bound to a compact buffer.
 * \param memory_scope memory scope of the buffer
 */
TVM_DLL tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype,
 std::string name, int data_alignment,
 int offset_factor, bool compact,
 std::string memory_scope = "");
} // namespace tir
} // namespace tvm
#endif // TVM_TIR_BUFFER_H_

30 changes: 2 additions & 28 deletions
src/driver/driver_api.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -83,32 +83,6 @@ Target DefaultTargetHost(Target target) {
 }
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype)));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

void GetBinds(const Array<ObjectRef>& args, bool compact,
 const std::unordered_map<te::Tensor, tir::Buffer>& binds,
 Map<te::Tensor, tir::Buffer>* out_binds, Array<ObjectRef>* out_arg_list) {
Expand All@@ -118,8 +92,8 @@ void GetBinds(const Array<ObjectRef>& args, bool compact,
 if (const te::TensorNode* tensor_node = x.as<te::TensorNode>()) {
 te::Tensor x_ref = GetRef<te::Tensor>(tensor_node);
 if (out_binds->find(x_ref) == out_binds->end()) {
 tir::Buffer buf =
BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0, compact);
 tir::Buffer buf = tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype,
 x_ref->op->name, -1, 0, compact);
 out_binds->Set(x_ref, buf);
 out_arg_list->push_back(buf);
 } else {
Expand Down

30 changes: 29 additions & 1 deletion
src/relay/backend/te_compiler.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -414,6 +414,33 @@ class TECompilerImpl : public TECompilerNode {
 }
 // lower the function
 std::unordered_map<te::Tensor, tir::Buffer> binds;

 // If we have memory scopes, need to create tir::Buffer knowing this info
 size_t i = 0; // for corresponding from tensor array
 for (Var param : key->source_func->params) {
 if (!param->virtual_device()->memory_scope.empty()) {
 for (const auto& ttype : FlattenTupleType(param->checked_type())) {
 te::Tensor x_ref = value->cached_func->inputs[i];
 // verification if we have synced params and tensors
 ICHECK(ttype->dtype == x_ref->dtype && ttype->shape.size() == x_ref->shape.size())
 << "function parameter does not correspond to prepared tensor";
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, param->virtual_device()->memory_scope);
 }
 }
 i++;
 }
 if (key->virtual_device != VirtualDevice::FullyUnconstrained() &&
 !key->virtual_device->memory_scope.empty() &&
 key->virtual_device->memory_scope != "global") {
 ICHECK(value->cached_func->outputs.size() == 1)
 << "Expect only one output for defined memory scope";
 te::Tensor x_ref = value->cached_func->outputs[0];
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, key->virtual_device->memory_scope);
 }
 auto func_name = value->cached_func->prim_fn_var->name_hint;
 VLOG(1) << "scheduling";
 IRModule scheduled_module =
Expand DownExpand Up@@ -895,7 +922,8 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
 } else {
 // Cases 1 and 2: lower the primitive function for the desired target, possibly using external
 // codegen.
 CCacheKey key(Downcast<Function>(primitive_func), target);
 CCacheKey key(Downcast<Function>(primitive_func), target,
 GetVirtualDevice(GetRef<Call>(call_node)));
 CachedFunc cfunc = compiler_->Lower(key, module_name_);
 ICHECK(cfunc.defined());
 return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
Expand Down

3 changes: 2 additions & 1 deletion
src/relay/backend/te_compiler_cache.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,11 @@ LoweredOutput::LoweredOutput(tvm::Array<te::Tensor> outputs, OpImplementation im
 data_ = std::move(n);
}

CCacheKey::CCacheKey(Function source_func, Target target) {
CCacheKey::CCacheKey(Function source_func, Target target, VirtualDevice vd) {
 auto n = make_object<CCacheKeyNode>();
 n->source_func = std::move(source_func);
 n->target = std::move(target);
 n->virtual_device = std::move(vd);
 data_ = std::move(n);
}

Expand Down

7 changes: 6 additions & 1 deletion
src/relay/backend/te_compiler_cache.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -82,10 +82,13 @@ class CCacheKeyNode : public Object {
 Function source_func;
 /*! \brief The hardware target.*/
 Target target;
 /*! \brief The virtual device constrains.*/
 VirtualDevice virtual_device;

 void VisitAttrs(tvm::AttrVisitor* v) {
 v->Visit("source_func", &source_func);
 v->Visit("target", &target);
 v->Visit("virtual_device", &virtual_device);
 }
 /*! \return The hash value of CCacheKey. */
 inline size_t Hash() const;
Expand DownExpand Up@@ -117,7 +120,8 @@ class CCacheKey : public ObjectRef {
 * \param source_func The source function.
 * \param target The target device.
 */
 TVM_DLL CCacheKey(Function source_func, Target target);
 TVM_DLL CCacheKey(Function source_func, Target target,
 VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained());

 const CCacheKeyNode* operator->() const { return static_cast<const CCacheKeyNode*>(get()); }
 // comparator
Expand DownExpand Up@@ -244,6 +248,7 @@ inline size_t CCacheKeyNode::Hash() const {
inline bool CCacheKeyNode::Equal(const CCacheKeyNode* other) const {
 if (Hash() != other->Hash()) return false;
 return this->target->str() == other->target->str() &&
 this->virtual_device == other->virtual_device &&
 tvm::StructuralEqual()(this->source_func, other->source_func);
}

Expand Down

27 changes: 27 additions & 0 deletions
src/tir/ir/buffer.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -585,6 +585,33 @@ Buffer::Buffer(Var data, DataType dtype, Array<PrimExpr> shape, Array<PrimExpr>
 data_ = std::move(n);
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact,
 std::string memory_scope) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype), memory_scope));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 .set_dispatch<BufferNode>([](const ObjectRef& node, ReprPrinter* p) {
 auto* op = static_cast<const BufferNode*>(node.get());
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
[Relay] Handle memory scope during lowering from relay level#11874

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

 Merged

junrushao
merged 2 commits into

from

Jun 29, 2022
 Merged


[Relay] Handle memory scope during lowering from relay level
#11874
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
7a4e763

[Relay] Handle memory scope during lowering from relay level


elvin-n Jun 24, 2022
a46f8c4

Move BufferWithOffsetAlignment from driver_api to tir buffer


elvin-n Jun 29, 2022

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


1 change: 0 additions & 1 deletion
include/tvm/driver/driver_api.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -165,7 +165,6 @@ TVM_DLL runtime::Module build(const Map<Target, IRModule>& input, const Target&
 * \return The built module that contains code for different processors.
 */
TVM_DLL runtime::Module build(const Map<String, IRModule>& input, const Target& target_host);

} // namespace tvm

#endif // TVM_DRIVER_DRIVER_API_H_

3 changes: 2 additions & 1 deletion
include/tvm/relay/expr.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -218,7 +218,8 @@ class VarNode : public ExprNode {

 bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
 equal->MarkGraphNode();
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid);
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid) &&
 equal(virtual_device_, other->virtual_device_);
 }

 void SHashReduce(SHashReducer hash_reduce) const {
Expand Down

17 changes: 17 additions & 0 deletions
include/tvm/tir/buffer.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -295,6 +295,23 @@ class DataProducer : public ObjectRef {
 TVM_DEFINE_OBJECT_REF_METHODS(DataProducer, ObjectRef, DataProducerNode);
};

/*!
 * \brief Creates TIR Buffer for provided parameters
 * \param shape shape of the buffer
 * \param dtype data type
 * \param name buffer name
 * \param data_alignment alignment requirement of data pointer in bytes
 * \param offset_factor Factor of elem_offset field, elem_offset is guaranteed to be
 * multiple of offset_factor
 User can specify data_alignment and offset_factor to be 0
 * A default value will be picked.
 * \param compact If the statement has already bound to a compact buffer.
 * \param memory_scope memory scope of the buffer
 */
TVM_DLL tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype,
 std::string name, int data_alignment,
 int offset_factor, bool compact,
 std::string memory_scope = "");
} // namespace tir
} // namespace tvm
#endif // TVM_TIR_BUFFER_H_

30 changes: 2 additions & 28 deletions
src/driver/driver_api.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -83,32 +83,6 @@ Target DefaultTargetHost(Target target) {
 }
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype)));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

void GetBinds(const Array<ObjectRef>& args, bool compact,
 const std::unordered_map<te::Tensor, tir::Buffer>& binds,
 Map<te::Tensor, tir::Buffer>* out_binds, Array<ObjectRef>* out_arg_list) {
Expand All@@ -118,8 +92,8 @@ void GetBinds(const Array<ObjectRef>& args, bool compact,
 if (const te::TensorNode* tensor_node = x.as<te::TensorNode>()) {
 te::Tensor x_ref = GetRef<te::Tensor>(tensor_node);
 if (out_binds->find(x_ref) == out_binds->end()) {
 tir::Buffer buf =
BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0, compact);
 tir::Buffer buf = tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype,
 x_ref->op->name, -1, 0, compact);
 out_binds->Set(x_ref, buf);
 out_arg_list->push_back(buf);
 } else {
Expand Down

30 changes: 29 additions & 1 deletion
src/relay/backend/te_compiler.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -414,6 +414,33 @@ class TECompilerImpl : public TECompilerNode {
 }
 // lower the function
 std::unordered_map<te::Tensor, tir::Buffer> binds;

 // If we have memory scopes, need to create tir::Buffer knowing this info
 size_t i = 0; // for corresponding from tensor array
 for (Var param : key->source_func->params) {
 if (!param->virtual_device()->memory_scope.empty()) {
 for (const auto& ttype : FlattenTupleType(param->checked_type())) {
 te::Tensor x_ref = value->cached_func->inputs[i];
 // verification if we have synced params and tensors
 ICHECK(ttype->dtype == x_ref->dtype && ttype->shape.size() == x_ref->shape.size())
 << "function parameter does not correspond to prepared tensor";
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, param->virtual_device()->memory_scope);
 }
 }
 i++;
 }
 if (key->virtual_device != VirtualDevice::FullyUnconstrained() &&
 !key->virtual_device->memory_scope.empty() &&
 key->virtual_device->memory_scope != "global") {
 ICHECK(value->cached_func->outputs.size() == 1)
 << "Expect only one output for defined memory scope";
 te::Tensor x_ref = value->cached_func->outputs[0];
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, key->virtual_device->memory_scope);
 }
 auto func_name = value->cached_func->prim_fn_var->name_hint;
 VLOG(1) << "scheduling";
 IRModule scheduled_module =
Expand DownExpand Up@@ -895,7 +922,8 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
 } else {
 // Cases 1 and 2: lower the primitive function for the desired target, possibly using external
 // codegen.
 CCacheKey key(Downcast<Function>(primitive_func), target);
 CCacheKey key(Downcast<Function>(primitive_func), target,
 GetVirtualDevice(GetRef<Call>(call_node)));
 CachedFunc cfunc = compiler_->Lower(key, module_name_);
 ICHECK(cfunc.defined());
 return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
Expand Down

3 changes: 2 additions & 1 deletion
src/relay/backend/te_compiler_cache.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,11 @@ LoweredOutput::LoweredOutput(tvm::Array<te::Tensor> outputs, OpImplementation im
 data_ = std::move(n);
}

CCacheKey::CCacheKey(Function source_func, Target target) {
CCacheKey::CCacheKey(Function source_func, Target target, VirtualDevice vd) {
 auto n = make_object<CCacheKeyNode>();
 n->source_func = std::move(source_func);
 n->target = std::move(target);
 n->virtual_device = std::move(vd);
 data_ = std::move(n);
}

Expand Down

7 changes: 6 additions & 1 deletion
src/relay/backend/te_compiler_cache.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -82,10 +82,13 @@ class CCacheKeyNode : public Object {
 Function source_func;
 /*! \brief The hardware target.*/
 Target target;
 /*! \brief The virtual device constrains.*/
 VirtualDevice virtual_device;

 void VisitAttrs(tvm::AttrVisitor* v) {
 v->Visit("source_func", &source_func);
 v->Visit("target", &target);
 v->Visit("virtual_device", &virtual_device);
 }
 /*! \return The hash value of CCacheKey. */
 inline size_t Hash() const;
Expand DownExpand Up@@ -117,7 +120,8 @@ class CCacheKey : public ObjectRef {
 * \param source_func The source function.
 * \param target The target device.
 */
 TVM_DLL CCacheKey(Function source_func, Target target);
 TVM_DLL CCacheKey(Function source_func, Target target,
 VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained());

 const CCacheKeyNode* operator->() const { return static_cast<const CCacheKeyNode*>(get()); }
 // comparator
Expand DownExpand Up@@ -244,6 +248,7 @@ inline size_t CCacheKeyNode::Hash() const {
inline bool CCacheKeyNode::Equal(const CCacheKeyNode* other) const {
 if (Hash() != other->Hash()) return false;
 return this->target->str() == other->target->str() &&
 this->virtual_device == other->virtual_device &&
 tvm::StructuralEqual()(this->source_func, other->source_func);
}

Expand Down

27 changes: 27 additions & 0 deletions
src/tir/ir/buffer.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -585,6 +585,33 @@ Buffer::Buffer(Var data, DataType dtype, Array<PrimExpr> shape, Array<PrimExpr>
 data_ = std::move(n);
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact,
 std::string memory_scope) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype), memory_scope));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 .set_dispatch<BufferNode>([](const ObjectRef& node, ReprPrinter* p) {
 auto* op = static_cast<const BufferNode*>(node.get());
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[Relay] Handle memory scope during lowering from relay level#11874

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

 Merged

junrushao
merged 2 commits into

from

Jun 29, 2022
 Merged


[Relay] Handle memory scope during lowering from relay level
#11874
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
7a4e763

[Relay] Handle memory scope during lowering from relay level


elvin-n Jun 24, 2022
a46f8c4

Move BufferWithOffsetAlignment from driver_api to tir buffer


elvin-n Jun 29, 2022

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


1 change: 0 additions & 1 deletion
include/tvm/driver/driver_api.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -165,7 +165,6 @@ TVM_DLL runtime::Module build(const Map<Target, IRModule>& input, const Target&
 * \return The built module that contains code for different processors.
 */
TVM_DLL runtime::Module build(const Map<String, IRModule>& input, const Target& target_host);

} // namespace tvm

#endif // TVM_DRIVER_DRIVER_API_H_

3 changes: 2 additions & 1 deletion
include/tvm/relay/expr.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -218,7 +218,8 @@ class VarNode : public ExprNode {

 bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
 equal->MarkGraphNode();
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid);
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid) &&
 equal(virtual_device_, other->virtual_device_);
 }

 void SHashReduce(SHashReducer hash_reduce) const {
Expand Down

17 changes: 17 additions & 0 deletions
include/tvm/tir/buffer.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -295,6 +295,23 @@ class DataProducer : public ObjectRef {
 TVM_DEFINE_OBJECT_REF_METHODS(DataProducer, ObjectRef, DataProducerNode);
};

/*!
 * \brief Creates TIR Buffer for provided parameters
 * \param shape shape of the buffer
 * \param dtype data type
 * \param name buffer name
 * \param data_alignment alignment requirement of data pointer in bytes
 * \param offset_factor Factor of elem_offset field, elem_offset is guaranteed to be
 * multiple of offset_factor
 User can specify data_alignment and offset_factor to be 0
 * A default value will be picked.
 * \param compact If the statement has already bound to a compact buffer.
 * \param memory_scope memory scope of the buffer
 */
TVM_DLL tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype,
 std::string name, int data_alignment,
 int offset_factor, bool compact,
 std::string memory_scope = "");
} // namespace tir
} // namespace tvm
#endif // TVM_TIR_BUFFER_H_

30 changes: 2 additions & 28 deletions
src/driver/driver_api.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -83,32 +83,6 @@ Target DefaultTargetHost(Target target) {
 }
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype)));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

void GetBinds(const Array<ObjectRef>& args, bool compact,
 const std::unordered_map<te::Tensor, tir::Buffer>& binds,
 Map<te::Tensor, tir::Buffer>* out_binds, Array<ObjectRef>* out_arg_list) {
Expand All@@ -118,8 +92,8 @@ void GetBinds(const Array<ObjectRef>& args, bool compact,
 if (const te::TensorNode* tensor_node = x.as<te::TensorNode>()) {
 te::Tensor x_ref = GetRef<te::Tensor>(tensor_node);
 if (out_binds->find(x_ref) == out_binds->end()) {
 tir::Buffer buf =
BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0, compact);
 tir::Buffer buf = tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype,
 x_ref->op->name, -1, 0, compact);
 out_binds->Set(x_ref, buf);
 out_arg_list->push_back(buf);
 } else {
Expand Down

30 changes: 29 additions & 1 deletion
src/relay/backend/te_compiler.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -414,6 +414,33 @@ class TECompilerImpl : public TECompilerNode {
 }
 // lower the function
 std::unordered_map<te::Tensor, tir::Buffer> binds;

 // If we have memory scopes, need to create tir::Buffer knowing this info
 size_t i = 0; // for corresponding from tensor array
 for (Var param : key->source_func->params) {
 if (!param->virtual_device()->memory_scope.empty()) {
 for (const auto& ttype : FlattenTupleType(param->checked_type())) {
 te::Tensor x_ref = value->cached_func->inputs[i];
 // verification if we have synced params and tensors
 ICHECK(ttype->dtype == x_ref->dtype && ttype->shape.size() == x_ref->shape.size())
 << "function parameter does not correspond to prepared tensor";
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, param->virtual_device()->memory_scope);
 }
 }
 i++;
 }
 if (key->virtual_device != VirtualDevice::FullyUnconstrained() &&
 !key->virtual_device->memory_scope.empty() &&
 key->virtual_device->memory_scope != "global") {
 ICHECK(value->cached_func->outputs.size() == 1)
 << "Expect only one output for defined memory scope";
 te::Tensor x_ref = value->cached_func->outputs[0];
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, key->virtual_device->memory_scope);
 }
 auto func_name = value->cached_func->prim_fn_var->name_hint;
 VLOG(1) << "scheduling";
 IRModule scheduled_module =
Expand DownExpand Up@@ -895,7 +922,8 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
 } else {
 // Cases 1 and 2: lower the primitive function for the desired target, possibly using external
 // codegen.
 CCacheKey key(Downcast<Function>(primitive_func), target);
 CCacheKey key(Downcast<Function>(primitive_func), target,
 GetVirtualDevice(GetRef<Call>(call_node)));
 CachedFunc cfunc = compiler_->Lower(key, module_name_);
 ICHECK(cfunc.defined());
 return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
Expand Down

3 changes: 2 additions & 1 deletion
src/relay/backend/te_compiler_cache.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,11 @@ LoweredOutput::LoweredOutput(tvm::Array<te::Tensor> outputs, OpImplementation im
 data_ = std::move(n);
}

CCacheKey::CCacheKey(Function source_func, Target target) {
CCacheKey::CCacheKey(Function source_func, Target target, VirtualDevice vd) {
 auto n = make_object<CCacheKeyNode>();
 n->source_func = std::move(source_func);
 n->target = std::move(target);
 n->virtual_device = std::move(vd);
 data_ = std::move(n);
}

Expand Down

7 changes: 6 additions & 1 deletion
src/relay/backend/te_compiler_cache.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -82,10 +82,13 @@ class CCacheKeyNode : public Object {
 Function source_func;
 /*! \brief The hardware target.*/
 Target target;
 /*! \brief The virtual device constrains.*/
 VirtualDevice virtual_device;

 void VisitAttrs(tvm::AttrVisitor* v) {
 v->Visit("source_func", &source_func);
 v->Visit("target", &target);
 v->Visit("virtual_device", &virtual_device);
 }
 /*! \return The hash value of CCacheKey. */
 inline size_t Hash() const;
Expand DownExpand Up@@ -117,7 +120,8 @@ class CCacheKey : public ObjectRef {
 * \param source_func The source function.
 * \param target The target device.
 */
 TVM_DLL CCacheKey(Function source_func, Target target);
 TVM_DLL CCacheKey(Function source_func, Target target,
 VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained());

 const CCacheKeyNode* operator->() const { return static_cast<const CCacheKeyNode*>(get()); }
 // comparator
Expand DownExpand Up@@ -244,6 +248,7 @@ inline size_t CCacheKeyNode::Hash() const {
inline bool CCacheKeyNode::Equal(const CCacheKeyNode* other) const {
 if (Hash() != other->Hash()) return false;
 return this->target->str() == other->target->str() &&
 this->virtual_device == other->virtual_device &&
 tvm::StructuralEqual()(this->source_func, other->source_func);
}

Expand Down

27 changes: 27 additions & 0 deletions
src/tir/ir/buffer.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -585,6 +585,33 @@ Buffer::Buffer(Var data, DataType dtype, Array<PrimExpr> shape, Array<PrimExpr>
 data_ = std::move(n);
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact,
 std::string memory_scope) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype), memory_scope));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 .set_dispatch<BufferNode>([](const ObjectRef& node, ReprPrinter* p) {
 auto* op = static_cast<const BufferNode*>(node.get());
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[Relay] Handle memory scope during lowering from relay level#11874

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

 Merged

junrushao
merged 2 commits into

from

Jun 29, 2022
 Merged


[Relay] Handle memory scope during lowering from relay level
#11874
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
7a4e763

[Relay] Handle memory scope during lowering from relay level


elvin-n Jun 24, 2022
a46f8c4

Move BufferWithOffsetAlignment from driver_api to tir buffer


elvin-n Jun 29, 2022

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


1 change: 0 additions & 1 deletion
include/tvm/driver/driver_api.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -165,7 +165,6 @@ TVM_DLL runtime::Module build(const Map<Target, IRModule>& input, const Target&
 * \return The built module that contains code for different processors.
 */
TVM_DLL runtime::Module build(const Map<String, IRModule>& input, const Target& target_host);

} // namespace tvm

#endif // TVM_DRIVER_DRIVER_API_H_

3 changes: 2 additions & 1 deletion
include/tvm/relay/expr.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -218,7 +218,8 @@ class VarNode : public ExprNode {

 bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
 equal->MarkGraphNode();
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid);
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid) &&
 equal(virtual_device_, other->virtual_device_);
 }

 void SHashReduce(SHashReducer hash_reduce) const {
Expand Down

17 changes: 17 additions & 0 deletions
include/tvm/tir/buffer.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -295,6 +295,23 @@ class DataProducer : public ObjectRef {
 TVM_DEFINE_OBJECT_REF_METHODS(DataProducer, ObjectRef, DataProducerNode);
};

/*!
 * \brief Creates TIR Buffer for provided parameters
 * \param shape shape of the buffer
 * \param dtype data type
 * \param name buffer name
 * \param data_alignment alignment requirement of data pointer in bytes
 * \param offset_factor Factor of elem_offset field, elem_offset is guaranteed to be
 * multiple of offset_factor
 User can specify data_alignment and offset_factor to be 0
 * A default value will be picked.
 * \param compact If the statement has already bound to a compact buffer.
 * \param memory_scope memory scope of the buffer
 */
TVM_DLL tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype,
 std::string name, int data_alignment,
 int offset_factor, bool compact,
 std::string memory_scope = "");
} // namespace tir
} // namespace tvm
#endif // TVM_TIR_BUFFER_H_

30 changes: 2 additions & 28 deletions
src/driver/driver_api.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -83,32 +83,6 @@ Target DefaultTargetHost(Target target) {
 }
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype)));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

void GetBinds(const Array<ObjectRef>& args, bool compact,
 const std::unordered_map<te::Tensor, tir::Buffer>& binds,
 Map<te::Tensor, tir::Buffer>* out_binds, Array<ObjectRef>* out_arg_list) {
Expand All@@ -118,8 +92,8 @@ void GetBinds(const Array<ObjectRef>& args, bool compact,
 if (const te::TensorNode* tensor_node = x.as<te::TensorNode>()) {
 te::Tensor x_ref = GetRef<te::Tensor>(tensor_node);
 if (out_binds->find(x_ref) == out_binds->end()) {
 tir::Buffer buf =
BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0, compact);
 tir::Buffer buf = tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype,
 x_ref->op->name, -1, 0, compact);
 out_binds->Set(x_ref, buf);
 out_arg_list->push_back(buf);
 } else {
Expand Down

30 changes: 29 additions & 1 deletion
src/relay/backend/te_compiler.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -414,6 +414,33 @@ class TECompilerImpl : public TECompilerNode {
 }
 // lower the function
 std::unordered_map<te::Tensor, tir::Buffer> binds;

 // If we have memory scopes, need to create tir::Buffer knowing this info
 size_t i = 0; // for corresponding from tensor array
 for (Var param : key->source_func->params) {
 if (!param->virtual_device()->memory_scope.empty()) {
 for (const auto& ttype : FlattenTupleType(param->checked_type())) {
 te::Tensor x_ref = value->cached_func->inputs[i];
 // verification if we have synced params and tensors
 ICHECK(ttype->dtype == x_ref->dtype && ttype->shape.size() == x_ref->shape.size())
 << "function parameter does not correspond to prepared tensor";
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, param->virtual_device()->memory_scope);
 }
 }
 i++;
 }
 if (key->virtual_device != VirtualDevice::FullyUnconstrained() &&
 !key->virtual_device->memory_scope.empty() &&
 key->virtual_device->memory_scope != "global") {
 ICHECK(value->cached_func->outputs.size() == 1)
 << "Expect only one output for defined memory scope";
 te::Tensor x_ref = value->cached_func->outputs[0];
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, key->virtual_device->memory_scope);
 }
 auto func_name = value->cached_func->prim_fn_var->name_hint;
 VLOG(1) << "scheduling";
 IRModule scheduled_module =
Expand DownExpand Up@@ -895,7 +922,8 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
 } else {
 // Cases 1 and 2: lower the primitive function for the desired target, possibly using external
 // codegen.
 CCacheKey key(Downcast<Function>(primitive_func), target);
 CCacheKey key(Downcast<Function>(primitive_func), target,
 GetVirtualDevice(GetRef<Call>(call_node)));
 CachedFunc cfunc = compiler_->Lower(key, module_name_);
 ICHECK(cfunc.defined());
 return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
Expand Down

3 changes: 2 additions & 1 deletion
src/relay/backend/te_compiler_cache.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,11 @@ LoweredOutput::LoweredOutput(tvm::Array<te::Tensor> outputs, OpImplementation im
 data_ = std::move(n);
}

CCacheKey::CCacheKey(Function source_func, Target target) {
CCacheKey::CCacheKey(Function source_func, Target target, VirtualDevice vd) {
 auto n = make_object<CCacheKeyNode>();
 n->source_func = std::move(source_func);
 n->target = std::move(target);
 n->virtual_device = std::move(vd);
 data_ = std::move(n);
}

Expand Down

7 changes: 6 additions & 1 deletion
src/relay/backend/te_compiler_cache.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -82,10 +82,13 @@ class CCacheKeyNode : public Object {
 Function source_func;
 /*! \brief The hardware target.*/
 Target target;
 /*! \brief The virtual device constrains.*/
 VirtualDevice virtual_device;

 void VisitAttrs(tvm::AttrVisitor* v) {
 v->Visit("source_func", &source_func);
 v->Visit("target", &target);
 v->Visit("virtual_device", &virtual_device);
 }
 /*! \return The hash value of CCacheKey. */
 inline size_t Hash() const;
Expand DownExpand Up@@ -117,7 +120,8 @@ class CCacheKey : public ObjectRef {
 * \param source_func The source function.
 * \param target The target device.
 */
 TVM_DLL CCacheKey(Function source_func, Target target);
 TVM_DLL CCacheKey(Function source_func, Target target,
 VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained());

 const CCacheKeyNode* operator->() const { return static_cast<const CCacheKeyNode*>(get()); }
 // comparator
Expand DownExpand Up@@ -244,6 +248,7 @@ inline size_t CCacheKeyNode::Hash() const {
inline bool CCacheKeyNode::Equal(const CCacheKeyNode* other) const {
 if (Hash() != other->Hash()) return false;
 return this->target->str() == other->target->str() &&
 this->virtual_device == other->virtual_device &&
 tvm::StructuralEqual()(this->source_func, other->source_func);
}

Expand Down

27 changes: 27 additions & 0 deletions
src/tir/ir/buffer.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -585,6 +585,33 @@ Buffer::Buffer(Var data, DataType dtype, Array<PrimExpr> shape, Array<PrimExpr>
 data_ = std::move(n);
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact,
 std::string memory_scope) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype), memory_scope));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 .set_dispatch<BufferNode>([](const ObjectRef& node, ReprPrinter* p) {
 auto* op = static_cast<const BufferNode*>(node.get());
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
[Relay] Handle memory scope during lowering from relay level#11874

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

 Merged

junrushao
merged 2 commits into

from

Jun 29, 2022
 Merged


[Relay] Handle memory scope during lowering from relay level
#11874
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


2 commits


Select commit
Hold shift + click to select a range
7a4e763

[Relay] Handle memory scope during lowering from relay level


elvin-n Jun 24, 2022
a46f8c4

Move BufferWithOffsetAlignment from driver_api to tir buffer


elvin-n Jun 29, 2022

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


1 change: 0 additions & 1 deletion
include/tvm/driver/driver_api.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -165,7 +165,6 @@ TVM_DLL runtime::Module build(const Map<Target, IRModule>& input, const Target&
 * \return The built module that contains code for different processors.
 */
TVM_DLL runtime::Module build(const Map<String, IRModule>& input, const Target& target_host);

} // namespace tvm

#endif // TVM_DRIVER_DRIVER_API_H_

3 changes: 2 additions & 1 deletion
include/tvm/relay/expr.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -218,7 +218,8 @@ class VarNode : public ExprNode {

 bool SEqualReduce(const VarNode* other, SEqualReducer equal) const {
 equal->MarkGraphNode();
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid);
 return equal(type_annotation, other->type_annotation) && equal(vid, other->vid) &&
 equal(virtual_device_, other->virtual_device_);
 }

 void SHashReduce(SHashReducer hash_reduce) const {
Expand Down

17 changes: 17 additions & 0 deletions
include/tvm/tir/buffer.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -295,6 +295,23 @@ class DataProducer : public ObjectRef {
 TVM_DEFINE_OBJECT_REF_METHODS(DataProducer, ObjectRef, DataProducerNode);
};

/*!
 * \brief Creates TIR Buffer for provided parameters
 * \param shape shape of the buffer
 * \param dtype data type
 * \param name buffer name
 * \param data_alignment alignment requirement of data pointer in bytes
 * \param offset_factor Factor of elem_offset field, elem_offset is guaranteed to be
 * multiple of offset_factor
 User can specify data_alignment and offset_factor to be 0
 * A default value will be picked.
 * \param compact If the statement has already bound to a compact buffer.
 * \param memory_scope memory scope of the buffer
 */
TVM_DLL tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype,
 std::string name, int data_alignment,
 int offset_factor, bool compact,
 std::string memory_scope = "");
} // namespace tir
} // namespace tvm
#endif // TVM_TIR_BUFFER_H_

30 changes: 2 additions & 28 deletions
src/driver/driver_api.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -83,32 +83,6 @@ Target DefaultTargetHost(Target target) {
 }
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype)));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

void GetBinds(const Array<ObjectRef>& args, bool compact,
 const std::unordered_map<te::Tensor, tir::Buffer>& binds,
 Map<te::Tensor, tir::Buffer>* out_binds, Array<ObjectRef>* out_arg_list) {
Expand All@@ -118,8 +92,8 @@ void GetBinds(const Array<ObjectRef>& args, bool compact,
 if (const te::TensorNode* tensor_node = x.as<te::TensorNode>()) {
 te::Tensor x_ref = GetRef<te::Tensor>(tensor_node);
 if (out_binds->find(x_ref) == out_binds->end()) {
 tir::Buffer buf =
BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0, compact);
 tir::Buffer buf = tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype,
 x_ref->op->name, -1, 0, compact);
 out_binds->Set(x_ref, buf);
 out_arg_list->push_back(buf);
 } else {
Expand Down

30 changes: 29 additions & 1 deletion
src/relay/backend/te_compiler.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -414,6 +414,33 @@ class TECompilerImpl : public TECompilerNode {
 }
 // lower the function
 std::unordered_map<te::Tensor, tir::Buffer> binds;

 // If we have memory scopes, need to create tir::Buffer knowing this info
 size_t i = 0; // for corresponding from tensor array
 for (Var param : key->source_func->params) {
 if (!param->virtual_device()->memory_scope.empty()) {
 for (const auto& ttype : FlattenTupleType(param->checked_type())) {
 te::Tensor x_ref = value->cached_func->inputs[i];
 // verification if we have synced params and tensors
 ICHECK(ttype->dtype == x_ref->dtype && ttype->shape.size() == x_ref->shape.size())
 << "function parameter does not correspond to prepared tensor";
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, param->virtual_device()->memory_scope);
 }
 }
 i++;
 }
 if (key->virtual_device != VirtualDevice::FullyUnconstrained() &&
 !key->virtual_device->memory_scope.empty() &&
 key->virtual_device->memory_scope != "global") {
 ICHECK(value->cached_func->outputs.size() == 1)
 << "Expect only one output for defined memory scope";
 te::Tensor x_ref = value->cached_func->outputs[0];
 binds[x_ref] =
 tir::BufferWithOffsetAlignment(x_ref->shape, x_ref->dtype, x_ref->op->name, -1, 0,
 false, key->virtual_device->memory_scope);
 }
 auto func_name = value->cached_func->prim_fn_var->name_hint;
 VLOG(1) << "scheduling";
 IRModule scheduled_module =
Expand DownExpand Up@@ -895,7 +922,8 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
 } else {
 // Cases 1 and 2: lower the primitive function for the desired target, possibly using external
 // codegen.
 CCacheKey key(Downcast<Function>(primitive_func), target);
 CCacheKey key(Downcast<Function>(primitive_func), target,
 GetVirtualDevice(GetRef<Call>(call_node)));
 CachedFunc cfunc = compiler_->Lower(key, module_name_);
 ICHECK(cfunc.defined());
 return MakeLoweredCall(primitive_func, cfunc->prim_fn_var, std::move(new_args),
Expand Down

3 changes: 2 additions & 1 deletion
src/relay/backend/te_compiler_cache.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -66,10 +66,11 @@ LoweredOutput::LoweredOutput(tvm::Array<te::Tensor> outputs, OpImplementation im
 data_ = std::move(n);
}

CCacheKey::CCacheKey(Function source_func, Target target) {
CCacheKey::CCacheKey(Function source_func, Target target, VirtualDevice vd) {
 auto n = make_object<CCacheKeyNode>();
 n->source_func = std::move(source_func);
 n->target = std::move(target);
 n->virtual_device = std::move(vd);
 data_ = std::move(n);
}

Expand Down

7 changes: 6 additions & 1 deletion
src/relay/backend/te_compiler_cache.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -82,10 +82,13 @@ class CCacheKeyNode : public Object {
 Function source_func;
 /*! \brief The hardware target.*/
 Target target;
 /*! \brief The virtual device constrains.*/
 VirtualDevice virtual_device;

 void VisitAttrs(tvm::AttrVisitor* v) {
 v->Visit("source_func", &source_func);
 v->Visit("target", &target);
 v->Visit("virtual_device", &virtual_device);
 }
 /*! \return The hash value of CCacheKey. */
 inline size_t Hash() const;
Expand DownExpand Up@@ -117,7 +120,8 @@ class CCacheKey : public ObjectRef {
 * \param source_func The source function.
 * \param target The target device.
 */
 TVM_DLL CCacheKey(Function source_func, Target target);
 TVM_DLL CCacheKey(Function source_func, Target target,
 VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained());

 const CCacheKeyNode* operator->() const { return static_cast<const CCacheKeyNode*>(get()); }
 // comparator
Expand DownExpand Up@@ -244,6 +248,7 @@ inline size_t CCacheKeyNode::Hash() const {
inline bool CCacheKeyNode::Equal(const CCacheKeyNode* other) const {
 if (Hash() != other->Hash()) return false;
 return this->target->str() == other->target->str() &&
 this->virtual_device == other->virtual_device &&
 tvm::StructuralEqual()(this->source_func, other->source_func);
}

Expand Down

27 changes: 27 additions & 0 deletions
src/tir/ir/buffer.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -585,6 +585,33 @@ Buffer::Buffer(Var data, DataType dtype, Array<PrimExpr> shape, Array<PrimExpr>
 data_ = std::move(n);
}

tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std::string name,
 int data_alignment, int offset_factor, bool compact,
 std::string memory_scope) {
 DataType storage_dtype = (dtype == DataType::Bool() ? DataType::Int(8) : dtype);
 auto data = tir::Var(name, PointerType(PrimType(storage_dtype), memory_scope));
 bool has_any = false;
 if (!compact) {
 for (const auto& it : shape) {
 if (it.as<tir::VarNode>()) {
 has_any = true;
 break;
 }
 }
 }
 tir::BufferType buffer_type = has_any ? tir::kAutoBroadcast : tir::kDefault;

 PrimExpr elem_offset;
 if (offset_factor != 0) {
 elem_offset = tir::Var(name + "_elem_offset", shape[0].dtype());
 } else {
 elem_offset = PrimExpr();
 }

 return tir::Buffer(data, dtype, shape, Array<PrimExpr>(), elem_offset, name, data_alignment,
 offset_factor, buffer_type);
}

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 .set_dispatch<BufferNode>([](const ObjectRef& node, ReprPrinter* p) {
 auto* op = static_cast<const BufferNode*>(node.get());
Expand Down