electriclilies commented on a change in pull request #45: URL: https://github.com/apache/tvm-rfcs/pull/45#discussion_r762205103
########## File path: rfcs/1111-first-class-virtual-device.md ########## @@ -0,0 +1,146 @@ +- Feature Name: Add virtual device information as a first class field on Relay expressions. +- Start Date: 2021-11-20 +- RFC PR: [apache/tvm-rfcs#1111](https://github.com/apache/tvm-rfcs/pull/1111) +- GitHub Issue: [apache/tvm#1111](https://github.com/apache/tvm/issues/1111) + +# Summary +[summary]: #summary + +I propose adding a new field to Relay expressions, virtual_device_. This field will contain virtual device information [currently called SEScope]. + +# Motivation +[motivation]: #motivation + +Currently, the virtual device information (called SEScope today, but we will rename it soon) is stored in Function attributes and in on_device Relay ops. This op is a wrapper op that contains the virtual device information for an expression. + +Here's an example of how the virtual device information is stored in the program today (example from test_pass_plan_devices.py): + +(note that SEScope is just the virtual device, but we have not renamed it yet). + +``` +""" +#[version = "0.0.5"] +def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32], + %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32], + param_se_scopes=[meta[SEScope][0], meta[SEScope][0], meta[SEScope][1], meta[SEScope][1]], + result_se_scope=meta[SEScope][1]) { + %0 = add(%a, %b); + %1 = on_device(%0, se_scope=meta[SEScope][0], is_fixed=True); + %2 = device_copy(%1, src_se_scope=meta[SEScope][0], dst_se_scope=meta[SEScope][1]); + %3 = add(%c, %d); + subtract(%2, %3) +} +""" +``` + +Using this method to store the virtual device information has proven to be very fragile. + +Normal visitors that don't care about virtual devices need to peek inside on_device ops. + +Additionally, we need DeviceAware visitors to be able to know the virtual device of sub-expressions. Notice in the example above that on_device doesn't wrap every expression. Let's say we want to know the virtual device of %3 while visiting it. We can't look it up directly since the information is not stored on the node. So how do we get the information? Well, instead of a normal visitor, we need to use a DeviceAware visitor, which keeps track of the current virtual device when it visits sub-expressions. We can then get the virtual device from the DeviceAware visitor itself. + +Making virtual devices first class eliminates the need for this complexity, and will allow us to implement more features in device and memory planning in the future. + +# Guide-level explanation +[guide-level-explanation]: #guide-level-explanation + +Users can introduce new information about the virtual device through the on_device op. This behavior is the same as it was before. + +However, let's say you want to write a pass that uses the virtual devices after device planning. Now, you'll be able to use the virtual device directly in your pass, just like you can with the checked_type_ field. + +For example, in this visitor, we can just look at the virtual device directly. + +``` + Expr VisitExpr_(const LetNode* let_node) final { + Expr expr = GetRef<Expr>(let_node); + // Iterate through chained lets, provided they all agree on their device type. + SEScope scope = expr->virtual_device_; + ... + ``` + +# Reference-level explanation +[reference-level-explanation]: #reference-level-explanation + +The AST change will occur in include/tvm/ir/expr.h: + +``` +class RelayExprNode : public BaseExprNode { + public: + /*! + * \brief Stores the result of type inference(type checking). + * + * \note This can be undefined before type inference. + * This value is discarded during serialization. + */ + mutable Type checked_type_ = Type(nullptr); + /*! + * \return The checked_type + */ + inline const Type& checked_type() const; + /*! + * \brief Check if the inferred(checked) type of the Expr + * is backed by a TTypeNode and return it. + * + * \note This function will thrown an error if the node type + * of this Expr is not TTypeNode. + * + * \return The corresponding TTypeNode pointer. + * \tparam The specific TypeNode we look for. + */ + template <typename TTypeNode> + inline const TTypeNode* type_as() const; + + /*! + * \brief The virtual device (SEScope) for this node (the result of device planning). + * + * \note Unfortunately, the type of virtual_device_ needs to be ObjectRef to avoid a circular import. + * We can forward-declare the SEScope type for the getter function, but not for the field + * itself. + */ + mutable ObjectRef virtual_device_; + + /*! + * \return The virtual device (currently called SEScope, this will be changing soon.) + */ + SEScope virtual_device() const; + + static constexpr const char* _type_key = "RelayExpr"; + static constexpr const uint32_t _type_child_slots = 22; + TVM_DECLARE_BASE_OBJECT_INFO(RelayExprNode, BaseExprNode); +}; +``` + +Additionally, I will add virtual_device_ to the WithFields methods. + +# Drawbacks +[drawbacks]: #drawbacks + +One challenge with making virtual devices first class is that passes in TVM do not propagate all fields when they visit expressions. You can see this today with spans-- most visitors do not preserve spans. When we introduce the virtual device field, we will need to ensure that it is propagated correctly throughout the Relay program. To do this, I introduced WithFields (code: https://github.com/apache/tvm/blob/main/src/relay/ir/expr.cc#L79-L99), a COW constructor that copies extra fields, including spans. I will extend WithFields to also copy virtual devices. Then, we can use WithFields to ensure that the virtual device field is correctly propagated. Review comment: Thanks the feedback @comaniac! I outlined a few other options we considered in the rationales and alternatives section. I'll expand a bit on what I said there. The current approach device planning takes is to store the virtual device (SEScope) information in the Relay program in the on_device nodes. Passes will always preserve these nodes and provided that passes don't significantly change the lexical scope, they won't cause problems. If we store the virtual device information directly on nodes, we will need to preserve it, either manually or through the WithFields method. In general passes *should* be doing this anyways-- they shouldn't be throwing away information willy-nilly. In reality they don't, which why spans are pretty much useless right now. We need to balance adding overhead to passes with maintaining a global invariant. Based on our experience with the current implementation, we think that the overhead is necessary to maintain the invariant. Fortunately, the change is only difficult because we have to do it in so many places, but it will be a big PITA, which is why @mbs-octoml used on_device to represent this information in the first place. You can take a look at this PR for an example of what the changes would look like: https://github.com/apache/tvm/pull/9533 -- 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-unsubscr...@tvm.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
