You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "masahi (via GitHub)" <gi...@apache.org> on 2023/06/20 10:32:00 UTC
[GitHub] [tvm] masahi commented on a diff in pull request #15090: [Bugfix] [Relay] Insertion of "device_copy" CallNode to Resolve Device Conflict on Unconstrained Nodes
masahi commented on code in PR #15090:
URL: https://github.com/apache/tvm/pull/15090#discussion_r1235066382
##########
src/relay/transforms/device_planner.cc:
##########
@@ -404,6 +410,201 @@ class RewriteOnDevices : public ExprMutator {
/* =============== Phase 1 =============== */
+/*!
+ * \brief Add "device_copy" calls for these nodes that have conflicted virtual_devices.
+ *
+ * Eg Suppose an IRModule contains the following expr:
+ * \code
+ * %0 = add(%a, %b);
+ * %1 = on_device(%0, virtual_device=d1);
+ * %2 = add(%b, %c);
+ * %3 = on_device(%2, virtual_device=d2);
+ * \endcode
+ * Then \p %b has two possbile virtual_devices \p d1 and \p d2 from the above expr:
+ *
+ * - If \p d1 and \p d2 are joinable, replace \p d1 and \p d2 with the joined device \p d:
+ * \code
+ * %0 = add(%a, %b);
+ * %1 = on_device(%0, virtual_device=d);
+ * %2 = add(%b, %c);
+ * %3 = on_device(%2, virtual_device=d);
+ * \endcode
+ *
+ * - If \p d1 and \p d2 are unjoinable, insert a "device_copy" CallNode to copy \p %b to \p d2:
+ * \code
+ * %0 = add(%a, %b);
+ * %1 = on_device(%0, virtual_device=d);
+ * %2 = device_copy(%b, src_dev_type=d1, dst_dev_type=d2)
+ * %3 = add(%2, %c);
+ * %4 = on_device(%3, virtual_device=d);
+ * \endcode
+ */
+struct DeviceContext {
+ VirtualDevice VirtualDeviceFor(const ExprNode* expr) {
+ auto itr = expr_to_device.find(expr);
+ if (itr != expr_to_device.end()) {
+ return itr->second;
+ }
+ auto default_dev = VirtualDevice::FullyUnconstrained();
+ expr_to_device.emplace(expr, default_dev);
+ return default_dev;
+ }
+
+ bool Update(const ExprNode* expr, VirtualDevice dev) {
+ bool success = true;
+ auto pair = expr_to_device.emplace(expr, dev);
+ if (!pair.second) {
+ auto replaced_item = pair.first;
+ auto joined_dev = VirtualDevice::Join(replaced_item->second, dev);
+ if (joined_dev == nullptr) {
+ success = false;
+ } else {
+ replaced_item->second = joined_dev.value();
+ }
+ }
+ return success;
+ }
+
+ bool IsConflicted(const ExprNode* expr) {
+ auto itr = conflicted_nodes.find(expr);
+ return itr != conflicted_nodes.end();
+ }
+
+ std::unordered_set<const ExprNode*> conflicted_nodes;
+ std::unordered_map<const ExprNode*, VirtualDevice> expr_to_device;
+};
+
+/*!
+ * \brief Flows the device constraints over the module and find all the conflicted nodes. The
Review Comment:
Flow
##########
src/relay/transforms/device_planner.cc:
##########
@@ -77,6 +77,12 @@
*
* Phase 1
* -------
+ * We iterate process the programs to find those nodes with conflicting virtual devices. If the
Review Comment:
iteratively
##########
src/relay/transforms/device_planner.cc:
##########
@@ -404,6 +410,201 @@ class RewriteOnDevices : public ExprMutator {
/* =============== Phase 1 =============== */
+/*!
+ * \brief Add "device_copy" calls for these nodes that have conflicted virtual_devices.
+ *
+ * Eg Suppose an IRModule contains the following expr:
+ * \code
+ * %0 = add(%a, %b);
+ * %1 = on_device(%0, virtual_device=d1);
+ * %2 = add(%b, %c);
+ * %3 = on_device(%2, virtual_device=d2);
+ * \endcode
+ * Then \p %b has two possbile virtual_devices \p d1 and \p d2 from the above expr:
+ *
+ * - If \p d1 and \p d2 are joinable, replace \p d1 and \p d2 with the joined device \p d:
+ * \code
+ * %0 = add(%a, %b);
+ * %1 = on_device(%0, virtual_device=d);
+ * %2 = add(%b, %c);
+ * %3 = on_device(%2, virtual_device=d);
+ * \endcode
+ *
+ * - If \p d1 and \p d2 are unjoinable, insert a "device_copy" CallNode to copy \p %b to \p d2:
+ * \code
+ * %0 = add(%a, %b);
+ * %1 = on_device(%0, virtual_device=d);
+ * %2 = device_copy(%b, src_dev_type=d1, dst_dev_type=d2)
+ * %3 = add(%2, %c);
+ * %4 = on_device(%3, virtual_device=d);
+ * \endcode
+ */
+struct DeviceContext {
+ VirtualDevice VirtualDeviceFor(const ExprNode* expr) {
+ auto itr = expr_to_device.find(expr);
+ if (itr != expr_to_device.end()) {
+ return itr->second;
+ }
+ auto default_dev = VirtualDevice::FullyUnconstrained();
+ expr_to_device.emplace(expr, default_dev);
+ return default_dev;
+ }
+
+ bool Update(const ExprNode* expr, VirtualDevice dev) {
+ bool success = true;
+ auto pair = expr_to_device.emplace(expr, dev);
+ if (!pair.second) {
+ auto replaced_item = pair.first;
+ auto joined_dev = VirtualDevice::Join(replaced_item->second, dev);
+ if (joined_dev == nullptr) {
+ success = false;
+ } else {
+ replaced_item->second = joined_dev.value();
+ }
+ }
+ return success;
+ }
+
+ bool IsConflicted(const ExprNode* expr) {
+ auto itr = conflicted_nodes.find(expr);
+ return itr != conflicted_nodes.end();
+ }
+
+ std::unordered_set<const ExprNode*> conflicted_nodes;
+ std::unordered_map<const ExprNode*, VirtualDevice> expr_to_device;
+};
+
+/*!
+ * \brief Flows the device constraints over the module and find all the conflicted nodes. The
+ * conficted note only contains nodes that has no explicit constraints. Eg "on_device" nodes are not
+ * considerred as conflicted.
+ */
+class ConflictedNodeFinder : ExprVisitor {
+ public:
+ explicit ConflictedNodeFinder(IRModule mod)
+ : mod_(std::move(mod)), dev_ctx_(std::make_unique<DeviceContext>()) {}
+
+ std::unique_ptr<DeviceContext> Finder() {
+ VLOG_CONTEXT << "ConflictedNodeFinder";
+ for (const auto& kv : mod_->functions) {
+ if (const auto* function_node = AsOptimizableFunctionNode(kv.second)) {
+ VisitExpr(GetRef<Function>(function_node));
+ }
+ }
+ for (auto const node : dev_ctx_->conflicted_nodes) {
+ if (node->IsInstance<CallNode>()) {
+ auto call = Downcast<Call>(GetRef<Expr>(node));
+ // "DeviceCapturer" will insert "device_copy" for "on_device" calls.
+ // Therefore, "on_device" should not be considerred as conflicted.
+ if (call->op == OnDeviceOp()) {
+ dev_ctx_->conflicted_nodes.erase(node);
+ }
+ }
+ }
+ return std::move(dev_ctx_);
+ }
+
+ private:
+ void VisitExpr_(const CallNode* call_node) final {
+ VLOG(2) << "Initial call node: " << std::endl << PrettyPrint(GetRef<Call>(call_node));
+ auto call_dev = dev_ctx_->VirtualDeviceFor(call_node);
+ auto body_dev = call_dev;
+
+ auto on_dev_props = GetOnDeviceProps(call_node);
+ auto dev_cp_props = GetDeviceCopyProps(call_node);
+ if (call_node->op == OnDeviceOp()) {
+ if (on_dev_props.constrain_body) {
+ body_dev = on_dev_props.virtual_device;
+ }
+ if (on_dev_props.constrain_result) {
+ call_dev = on_dev_props.virtual_device;
+ }
+ } else if (call_node->op == DeviceCopyOp()) {
+ body_dev = dev_cp_props.src_virtual_device;
+ call_dev = dev_cp_props.dst_virtual_device;
+ }
+
+ if (!dev_ctx_->Update(call_node, call_dev) && call_node->op != OnDeviceOp()) {
+ LOG(FATAL) << "Mismatched device type after iterating args. Implied device: " << std::endl
+ << PrettyPrint(call_dev) << "and practial device:" << std::endl
+ << PrettyPrint(dev_ctx_->VirtualDeviceFor(call_node)) << std::endl
+ << "With CallNode: " << std::endl
+ << PrettyPrint(GetRef<Call>(call_node));
+ }
+
+ for (auto& arg : call_node->args) {
+ VLOG(3) << "Handle call node arg: " << std::endl << PrettyPrint(arg);
+ if (!dev_ctx_->Update(arg.get(), body_dev)) {
+ VLOG(2) << "Conflicted node found:" << std::endl
+ << PrettyPrint(GetRef<Expr>(arg.get())) << std::endl
+ << "With corresponding Callee:" << std::endl
+ << PrettyPrint(GetRef<Call>(call_node));
+ dev_ctx_->conflicted_nodes.emplace(arg.get());
+ }
+ }
+ for (auto& expr : call_node->args) {
+ VisitExpr(expr);
+ }
+ }
+
+ IRModule mod_;
+ std::unique_ptr<DeviceContext> dev_ctx_;
Review Comment:
Why need pointer here?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org