[DRAFT] [RELAY] [DEVICES] Change the virtual device representation for Functions

include/tvm/ir/expr.h

@@ -39,6 +39,9 @@ namespace tvm {
 
 using tvm::runtime::String;
 
+// Forward-declare SEScope to avoid circular imports.
+class SEScope;
+
 /*!
  * \brief Base type of all the expressions.
  * \sa Expr
@@ -165,6 +168,29 @@ class RelayExprNode : public BaseExprNode {
   template <typename TTypeNode>
   inline const TTypeNode* type_as() const;
 
+  /*!
+   * \brief The virtual device (SEScope) for this node (the result of device planning).
+   * For first-order expressions (non functions), this describes where the result of evaluating the
+   * expression should be stored. Note that currently, all composite first-order values (tuples,
+   * references, ADTs) must be stored on the same virtual device. This means that it is not possible
+   * to store two tuple fields on different devices, so we only need one virtual device for these
+   * types.
+   *
+   * For expressions that have the function type, the virtual device describes where the result of
+   * the call to the function or closure is stored (instead of where the function itself is stored).
+   * The SEScope's Target field describes how the body of the function should be compiled.
+   *
+   * \note Unfortunately, the type of virtual_device_ needs to be ObjectRef to avoid a circular
+   * import.
+   */
+  mutable ObjectRef virtual_device_;
+
+  /*!
+   * \return The virtual device (SEScope).
+   * If the virtual device is not defined, returns SEScope::FullyUnconstrained().
+   */
+  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);

include/tvm/ir/function.h

@@ -190,26 +190,6 @@ constexpr const char* kTarget = "target";
  */
 constexpr const char* kGlobalSymbol = "global_symbol";
 
-/*!
- * \brief The SEScope which will hold each of the functions parameters.
- *
- * Only supported on Relay \p Functions. Generally added by the \p PlanDevices pass, but
- * may be included as an annotation on user programs.
- *
- * Type: Array<SEScope>
- */
-constexpr const char* kParamSEScopes = "param_se_scopes";
-
-/*!
- * \brief The SEScope which will hold the function result.
- *
- * Only supported on Relay \p Functions. Generally added by the \p PlanDevices pass, but
- * may be included as an annotation on user programs.
- *
- * Type: SEScope
- */
-constexpr const char* kResultSEScope = "result_se_scope";
-
 }  // namespace attr
 }  // namespace tvm
 #endif  // TVM_IR_FUNCTION_H_

include/tvm/relay/expr.h

@@ -28,6 +28,7 @@
 #include <tvm/ir/expr.h>
 #include <tvm/ir/module.h>
 #include <tvm/ir/op.h>
+#include <tvm/target/se_scope.h>
 
 #include <functional>
 #include <stack>
@@ -151,10 +152,14 @@ class Tuple : public Expr {
  * \param tuple The tuple to copy
  * \param opt_fields The (optional) fields for the copied tuple. If none, ret_tuple->fields =
  * tuple->fields.
- * \param opt_span The (optional) span for the copied tuple. If none, ret_tuple->span = tuple->span.
+ * \param opt_virtual_device The (optional) virtual_device for the copied tuple. If none,
+ * ret_tuple->virtual_device = tuple->virtual_device.
+ * \param opt_span The (optional) span for the copied tuple. If none,
+ * ret_tuple->span = tuple->span.
  */
 Tuple WithFields(Tuple tuple, Optional<Array<Expr>> opt_fields = Optional<Array<Expr>>(),
-                 Optional<Span> opt_span = Optional<Span>(nullptr));
+                 Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
+                 Optional<Span> opt_span = Optional<Span>());
 
 /*!
  * \brief Local variables used in the let expression.
@@ -240,14 +245,17 @@ class Var : public Expr {
  * \param opt_vid The (optional) vid for the copied var. If none, ret_var->vid = var->vid.
  * \param opt_type_annotation The (optional) type_annotation for the copied var. If none,
  * ret_var->type_annotation = var->type_annotation.
+ * \param opt_virtual_device The (optional) virtual_device for the copied tuple. If none,
+ * ret_tuple->virtual_device = tuple->virtual_device.
  * \param opt_span The (optional) span for the copied var. If none, ret_var->span = var->span.
- * \return If all properties are null or the same as the property in the input var
- * (i.e., opt_vid is null or opt_vid.value() == var->vid, etc.), then we return var. Otherwise,
- * we return a copy of call with the different fields overwritten. (i.e., if
- * opt_vid.value() != var->vid, then ret_var->vid = opt_.value()).
+ * \return If all properties are null or the same as the property in the input var (i.e., opt_vid is
+ * null or opt_vid.value() == var->vid, etc.), then we return var. Otherwise, we return a copy of
+ * call with the different fields overwritten. (i.e., if opt_vid.value() != var->vid, then
+ * ret_var->vid = opt_.value()).
  */
 Var WithFields(Var var, Optional<Id> opt_vid = Optional<Id>(),
                Optional<Type> opt_type_annotation = Optional<Type>(),
+               Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
                Optional<Span> opt_span = Optional<Span>());
 
 /*!
@@ -362,16 +370,19 @@ class Call : public Expr {
  * call->attrs.
  * \param opt_type_args The (optional) type args for the copied call. If none,
  * ret_call->type_args = call->type_args.
+ * \param opt_virtual_device The (optional) virtual_device for the copied call. If none,
+ * ret_call->virtual_device = call->virtual_device.
  * \param opt_span The (optional) span for the copied call. If none, ret_call->span = call->span.
- * \return If all properties are null or the same as the property in the input call
- * (i.e., opt_op is null or opt_op.value() == call->op, etc.), then we return call. Otherwise, we
- * return a copy of call with the different fields overwritten. (i.e., if opt_op.value() !=
- * call->op, then ret_call->op = opt_op.value()).
+ * \return If all properties are null or the same as the property in the input call (i.e., opt_op is
+ * null or opt_op.value() == call->op, etc.), then we return call. Otherwise, we return a copy of
+ * call with the different fields overwritten. (i.e., if opt_op.value() != call->op, then
+ * ret_call->op = opt_op.value()).
  */
 Call WithFields(Call call, Optional<Expr> opt_op = Optional<Expr>(),
                 Optional<Array<Expr>> opt_args = Optional<Array<Expr>>(),
                 Optional<Attrs> opt_attrs = Optional<Attrs>(),
                 Optional<Array<Type>> opt_type_args = Optional<Array<Type>>(),
+                Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
                 Optional<Span> opt_span = Optional<Span>());
 
 /*!
@@ -456,6 +467,8 @@ class Let : public Expr {
  * \param opt_var The (optional) var for the copied let. If none, ret_let->op = let->op.
  * \param opt_value The (optional) value for the copied let. If none, ret_let->args = let->args.
  * \param opt_body The (optional) body for the copied let. If none, ret_let->attrs = let->attrs.
+ * \param opt_virtual_device The (optional) virtual_device for the copied let. If none,
+ * ret_let->virtual_device = let->virtual_device.
  * \param opt_span The (optional) span for the copied let. If none, ret_let->span = let->span.
  * \return If all properties are null or the same as the property in the input let (i.e., opt_var is
  * null or opt_var.value() == let->var, etc.), then we return let. Otherwise, we return a copy of
@@ -465,6 +478,7 @@ class Let : public Expr {
 Let WithFields(Let let, Optional<Var> opt_var = Optional<Var>(),
                Optional<Expr> opt_value = Optional<Expr>(),
                Optional<Expr> opt_body = Optional<Expr>(),
+               Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
                Optional<Span> opt_span = Optional<Span>());
 
 /*!
@@ -539,17 +553,19 @@ class If : public Expr {
  * ret_if->true_branch = ret_if->false_branch.
  * \param opt_false_branch The (optional) false_branch
  * for the copied if_expr. If none, ret_if->false_branch = if_expr->false_branch.
- * \param opt_span
- * The (optional) span for the copied if_expr. If none, ret_if->span = if_expr->span.
- * \return If all
- * properties are null or the same as the property in the input if_expr (i.e., opt_cond is null or
- * opt_cond.value() == if_expr->cond, etc.), then we return if_expr. Otherwise, we return a copy of
- * if_expr with the different fields overwritten. (i.e., if opt_cond.value() != if_expr->cond, then
- * ret_if->cond = opt_cond.value()).
+ * \param opt_virtual_device The (optional) virtual_device for the copied if_expr. If none,
+ * ret_if->virtual_device = if_expr->virtual_device.
+ * \param opt_span The (optional) span for the copied if_expr. If none,
+ * ret_if->span = if_expr->span.
+ * \return If all properties are null or the same as the property in
+ * the input if_expr (i.e., opt_cond is null or opt_cond.value() == if_expr->cond, etc.), then we
+ * return if_expr. Otherwise, we return a copy of if_expr with the different fields overwritten.
+ * (i.e., if opt_cond.value() != if_expr->cond, then ret_if->cond = opt_cond.value()).
  */
 If WithFields(If if_expr, Optional<Expr> opt_cond = Optional<Expr>(),
               Optional<Expr> opt_true_branch = Optional<Expr>(),
               Optional<Expr> opt_false_branch = Optional<Expr>(),
+              Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
               Optional<Span> opt_span = Optional<Span>());
 
 /*! \brief Get index-th field out of a tuple. */
@@ -603,8 +619,9 @@ class TupleGetItem : public Expr {
  * ret_tuple_get_item->tuple = tuple_get_item->tuple.
  * \param opt_index The (optional) index for the copied tuple_get_item. If none,
  * ret_tuple_get_item->index = tuple_get_item->index.
- * \param
- * opt_span The (optional) span for the copied tuple_get_item. If none,
+ * \param opt_virtual_device The (optional) virtual_device for the copied tuple_get_item.
+ * If none, ret_tuple_get_item->virtual_device = tuple_get_item->virtual_device.
+ * \param opt_span The (optional) span for the copied tuple_get_item. If none,
  * ret_tuple_get_item->span = tuple_get_item->span.
  * \return If all properties are null or the same as the property in the input tuple_get_item
  * (i.e., opt_tuple is null or opt_tuple.value() == tuple_get_item->tuple, etc.), then we return
@@ -614,6 +631,7 @@ class TupleGetItem : public Expr {
  */
 TupleGetItem WithFields(TupleGetItem tuple_get_item, Optional<Expr> opt_tuple = Optional<Expr>(),
                         Optional<Integer> opt_index = Optional<Integer>(),
+                        Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
                         Optional<Span> opt_span = Optional<Span>());
 
 /*! \brief Create a new Reference out of initial value. */
@@ -663,6 +681,8 @@ class RefCreate : public Expr {
  * \param ref_create The ref_create to copy.
  * \param opt_value The (optional) value for the copied ref_create. If none,
  * ret_ref_create->value = ref_create->value.
+ * \param opt_virtual_device The (optional) virtual_device for the copied ref_create. If none,
+ * ret_ref_create->virtual_device = ref_create->virtual_device.
  * \param opt_span The (optional) span for the copied ref_create. If none,
  * ret_ref_create->span = ref_create->span.
  * \return If all properties are null or the same as the property in the input ref_create
@@ -672,6 +692,7 @@ class RefCreate : public Expr {
  * ret_ref_create->value = opt_value.value()).
  */
 RefCreate WithFields(RefCreate ref_create, Optional<Expr> opt_value = Optional<Expr>(),
+                     Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
                      Optional<Span> opt_span = Optional<Span>());
 
 /*! \brief Get value out of Reference. */
@@ -720,15 +741,18 @@ class RefRead : public Expr {
  * \param ref_read The ref_read to copy.
  * \param opt_ref The (optional) ref for the copied ref_read. If none, ret_ref_read->ref =
  * ref_read->ref.
- * \param opt_span
- * The (optional) span for the copied ref_read. If none, ret_ref_read->span = ref_read->span.
- * \return If all properties are null or the same as the property in the input ref_read
- * (i.e., opt_ref is null or opt_ref.value() == ref_read->ref, etc.), then we return ref_read.
- * Otherwise, we return a copy of ref_read with the different fields overwritten.
- * (i.e., if opt_ref.value() != ref_read->ref, then
- * ret_ref_read->ref = opt_ref.value()).
+ * \param opt_virtual_device
+ * The (optional) virtual_device for the copied ref_read. If none, ret_ref_read->virtual_device =
+ * ref_read->virtual_device.
+ * \param opt_span The (optional) span for the copied ref_read. If none, ret_ref_read->span =
+ * ref_read->span.
+ * \return If all properties are null or the same as the property in the input
+ * ref_read (i.e., opt_ref is null or opt_ref.value() == ref_read->ref, etc.), then we return
+ * ref_read. Otherwise, we return a copy of ref_read with the different fields overwritten. (i.e.,
+ * if opt_ref.value() != ref_read->ref, then ret_ref_read->ref = opt_ref.value()).
  */
 RefRead WithFields(RefRead ref_read, Optional<Expr> opt_ref = Optional<Expr>(),
+                   Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
                    Optional<Span> opt_span = Optional<Span>());
 
 /*! \brief Set value of Reference. The whole expression evaluates to an Empty Tuple. */
@@ -784,16 +808,19 @@ class RefWrite : public Expr {
  * ret_ref_write->ref = ref_write->ref.
  * \param opt_value The (optional) value for the copied ref_write. If none,
  * ret_ref_write->value = ref_write->value.
- * \param opt_span
- * The (optional) span for the copied ref_write. If none, ret_ref_write->span = ref_write->span.
- * \return If all properties are null or the same as the property in the input ref_write
- * (i.e., opt_ref is null or opt_ref.value() == ref_write->ref, etc.), then we return ref_write.
- * Otherwise, we return a copy of ref_write with the different fields overwritten.
- * (i.e., if ref_write.value() != ref_write->ref, then
- * ret_ref_write->ref = opt_ref.value()).
+ * \param opt_virtual_device
+ * The (optional) virtual_device for the copied ref_write. If none, ret_ref_write->virtual_device =
+ * ref_write->virtual_device.
+ * \param opt_span The (optional) span for the copied ref_write. If none, ret_ref_write->span =
+ * ref_write->span.
+ * \return If all properties are null or the same as the property in the input ref_write (i.e.,
+ * opt_ref is null or opt_ref.value() == ref_write->ref, etc.), then we return ref_write. Otherwise,
+ * we return a copy of ref_write with the different fields overwritten. (i.e., if ref_write.value()
+ * != ref_write->ref, then ret_ref_write->ref = opt_ref.value()).
  */
 RefWrite WithFields(RefWrite ref_write, Optional<Expr> opt_ref = Optional<Expr>(),
                     Optional<Expr> opt_value = Optional<Expr>(),
+                    Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
                     Optional<Span> opt_span = Optional<Span>());
 
 /*!

include/tvm/relay/function.h

@@ -134,6 +134,8 @@ class Function : public BaseFunc {
  * \param opt_attrs
  * The (optional) attributes for the copied function. If none,
  * ret_function->attrs = function->attrs.
+ * \param opt_virtual_device The (optional) virtual_device for the copied function. If none,
+ * ret_function->virtual_device = function->virtual_device.
  * \param opt_span The (optional) span for the copied function. If none,
  * ret_function->span = function->span.
  * \return If all properties are null or the same as the property in the input function
@@ -146,6 +148,7 @@ Function WithFields(Function function, Optional<Array<Var>> opt_params = Optiona
                     Optional<Type> opt_ret_type = Optional<Type>(),
                     Optional<Array<TypeVar>> opt_ty_params = Optional<Array<TypeVar>>(),
                     Optional<DictAttrs> opt_attrs = Optional<DictAttrs>(),
+                    Optional<SEScope> opt_virtual_device = Optional<SEScope>(),
                     Optional<Span> opt_span = Optional<Span>());
 
 /*

python/tvm/relay/op/annotation/annotation.py

@@ -56,30 +56,6 @@ def on_device(data, device, is_fixed=False):
     return _make.OnDevice(data, _make_se_scope(device), is_fixed)
 
 
-def function_on_device(function, param_devices, result_device):
-    """Annotates a Relay function with the device types on which its parameters and result should
-    be stored.
-
-    Parameters
-    ----------
-    function : tvm.relay.Function
-        The function to be annotated.
-
-    param_devices : Array[Union[:py:class:`Device`, str]]
-        The devices for each parameter.
-
-    result_device: Union[:py:class:`Device`, str]
-        The device for the function result.
-
-    Returns
-    -------
-    result : tvm.relay.Function
-        The annotated function.
-    """
-    return _make.FunctionOnDevice(
-        function, [_make_se_scope(d) for d in param_devices], _make_se_scope(result_device)
-    )
-
 
 def stop_fusion(data):
     """Annotate an expression to prevent it being fused with following expressions.

src/relay/analysis/extract_fused_functions.cc

@@ -52,15 +52,14 @@ class FusedFunctionExtractorWrapper : private ExprVisitor {
   // have the desired equals property
   Map<String, Function> functions;
 
-  void VisitExpr_(const FunctionNode* n) final {
-    if (n->HasNonzeroAttr(attr::kPrimitive)) {
+  void VisitExpr_(const FunctionNode* func_node) final {
+    if (func_node->HasNonzeroAttr(attr::kPrimitive)) {
       // Add function to functions, keyed by function hash string
-      Function func = Function(n->params, n->body, n->ret_type, n->type_params, n->attrs);
-      size_t hash_ = tvm::StructuralHash()(func);
-      this->functions.Set(std::to_string(hash_), func);
+      size_t hash_ = tvm::StructuralHash()(GetRef<Function>(func_node));
+      this->functions.Set(std::to_string(hash_), GetRef<Function>(func_node));
     }
 
-    ExprVisitor::VisitExpr_(n);
+    ExprVisitor::VisitExpr_(func_node);
   }
 };
 

src/relay/backend/contrib/cmsisnn/extract_constants.cc

@@ -46,15 +46,15 @@ class ExtractConstantsMutator : public MixedModeMutator {
  private:
   String gen_var_name() { return "tvm_var_extract_const_" + std::to_string(var_count_++); }
 
-  Expr VisitExpr_(const FunctionNode* function) final {
-    Function func = GetRef<Function>(function);
+  Expr VisitExpr_(const FunctionNode* func_node) final {
+    Function func = GetRef<Function>(func_node);
     function_to_constants_.Set(func, Array<Constant>{});
     functions_.push_back(func);
     auto new_body = VisitExpr(func->body);
     functions_.pop_back();
     if (function_to_constants_[func].size()) {
-      func = Function(FreeVars(new_body), new_body, func->ret_type, FreeTypeVars(new_body, mod_),
-                      func->attrs);
+      func = WithFields(std::move(func), FreeVars(new_body), new_body, func->ret_type,
+                        FreeTypeVars(new_body, mod_), func->attrs);
     }
     return func;
   }
@@ -145,8 +145,7 @@ IRModule ExtractConstants(const IRModule& mod) {
   auto new_main_body = extract_constants.VisitExpr(main_func->body);
   if (!new_main_body.same_as(main_func->body)) {
     auto main_var = mod->GetGlobalVar("main");
-    auto new_main_func = Function(main_func->params, new_main_body, main_func->ret_type,
-                                  main_func->type_params, main_func->attrs);
+    Function new_main_func = WithFields(std::move(main_func), main_func->params, new_main_body);
     mod->Update(main_var, new_main_func);
   }
   return mod;

src/relay/backend/contrib/cmsisnn/relay_to_tir.cc

@@ -44,13 +44,8 @@ class RelayToTIRVisitor : public MixedModeMutator {
 
   IRModule Mutate() {
     GlobalVar main_global_var = ir_module_->GetGlobalVar("main");
-    BaseFunc main = ir_module_->Lookup(main_global_var);
-    Function main_func = GetRef<Function>(main.as<FunctionNode>());
-
-    // Copy everything across and mutate the body
-    Function mutated_main =
-        Function(main_func->params, VisitExpr(main_func->body), main_func->ret_type,
-                 main_func->type_params, main_func->attrs, main_func->span);
+    Function main = Downcast<Function>(ir_module_->Lookup(main_global_var));
+    Function mutated_main = WithFields(std::move(main), main->params, VisitExpr(main->body));
 
     ir_module_->Update(main_global_var, mutated_main);