mbs-octoml commented on a change in pull request #8788:
URL: https://github.com/apache/tvm/pull/8788#discussion_r720593981
##########
File path: src/relay/analysis/dependency_graph.cc
##########
@@ -87,8 +83,16 @@ class DependencyGraph::Creator : private MixedModeVisitor {
void VisitExpr_(const CallNode* c) final {
DependencyGraph::Node* n = graph_.expr_node[GetRef<Expr>(c)];
Depend(n, c->op);
- for (const auto& a : c->args) {
- Depend(n, a);
+ if (c->op == on_device_op_) {
+ DependencyGraph::Node* body = NewNode(true);
Review comment:
Ended up reverting this -- it was a bit too subtle an approach. All the
"on_device" special casing is now handled explicitly by transforms which derive
from one of the three (!) device aware visitor classes.
##########
File path: src/relay/backend/graph_plan_memory.cc
##########
@@ -39,38 +42,45 @@ using backend::StaticMemoryPlan;
using backend::StorageInfo;
using IntegerArray = Array<Integer>;
+/*! A representation of a block of memory required at runtime on some device.
*/
struct StorageToken {
/*! \brief Reference counter */
int ref_counter{0};
/*! \brief number of bytes */
size_t max_bytes{0};
- /*! \brief The corresponding tensor type node. */
- const TensorTypeNode* ttype{nullptr};
- /*! \brief virtual device index that corresponds to the device_type in
- * DLDevice. */
- int device_type{0};
+ /*! \brief The corresponding tensor type. */
+ TensorType ttype{nullptr};
+ /*! \brief Device on which memory will reside. */
+ Device device{kInvalidDevice};
/*! \brief The storage id */
int64_t storage_id{-1};
+
+ bool is_valid() const { return device.device_type !=
kInvalidDevice.device_type; }
+
+ bool is_compatible(const StorageToken& that) const {
+ return device.device_type == that.device.device_type;
+ }
};
std::ostream& operator<<(std::ostream& os, StorageToken tok) {
return os << "StorageToken: " << std::endl
<< "ref_counter: " << tok.ref_counter << std::endl
<< "max_bytes: " << tok.max_bytes << std::endl
- << "tttype: " << tok.ttype
+ << "tttype: " << tok.ttype << std::endl
Review comment:
Zapped.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
