WIP

Author: jroesch

src/relay/backend/aot_executor_codegen.cc

@@ -330,7 +330,7 @@ class AOTExecutorCodegen : public ExprVisitor {
       fi_node->tir_primfuncs.Set(primfunc_target, primfunc);
       fi_node->relay_primfuncs.Set(primfunc_target, relay_func);
     }
-    function_metadata_.Set(cfunc->func_name, FunctionInfo(fi_node));
+    function_metadata_.Set(cfunc->prim_fn_var->name_hint, FunctionInfo(fi_node));
   }
 
   void VisitExpr_(const CallNode* op) override {
@@ -368,7 +368,7 @@ class AOTExecutorCodegen : public ExprVisitor {
       UpdateConstants(func, &params_);
 
       // Generate the TIR function call
-      CreateFuncCall(GetRef<Call>(op), ext_func->func_name);
+      CreateFuncCall(GetRef<Call>(op), ext_func->prim_fn_var->name_hint);
       return;
     }
 
@@ -403,7 +403,7 @@ class AOTExecutorCodegen : public ExprVisitor {
     UpdateFunctionMetadata(lowered_func, func, target);
 
     // Generate the TIR function call
-    CreateFuncCall(GetRef<Call>(op), lowered_func->func_name);
+    CreateFuncCall(GetRef<Call>(op), lowered_func->prim_fn_var->name_hint);
   }
 
   void VisitExpr_(const VarNode* op) override {

src/relay/backend/graph_executor_codegen.cc

@@ -189,7 +189,7 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
     targets_ = targets;
   }
 
-  /*!
+   /*!
    * \brief Update the "main" control function's metadata
    *
    * \param func The main function that contains calls to relay primitive functions
@@ -273,6 +273,9 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
     }
 
     function_metadata_.Set(String(runtime::symbol::tvm_module_main), FunctionInfo(fi_node));
+  }
+
+
   StorageInfo GetStorageInfo(const Expr& e) {
     size_t count = memory_plan_->expr_to_storage_info.count(e);
     ICHECK_GT(count, 0) << "Expr is not existing in storage plan";
@@ -466,16 +469,18 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
     return fields;
   }
 
-  std::vector<GraphNodeRef> GraphAddCallNode(const CallNode* op, const std::string& op_name,
-                                             const std::string& func_name, GraphAttrs attrs) {
+  std::vector<GraphNodeRef> GraphAddCallNode(const CallNode* op, const std::string& func_name, GraphAttrs op_attrs) {
     std::vector<GraphNodeRef> inputs;
     for (auto arg : op->args) {
       auto res = VisitExpr(arg);
       for (auto nr : res) {
         inputs.push_back(nr);
       }
     }
-    auto node = GraphOpNode::make_node_ptr(op_name, GraphAttrs(), func_name, inputs, attrs);
+
+    // Compute the operator name, because we used the get unique name when generating the kernel.
+    auto op_name =_GetUniqueName(func_name);
+    auto node = GraphOpNode::make_node_ptr(op_name, GraphAttrs(), func_name, inputs, op_attrs);
     return AddNode(node, GetRef<Expr>(op));
   }
 
@@ -484,8 +489,8 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
     if (auto global_node = call->op.as<GlobalVarNode>()) {
       auto prim_fn_name = global_node->name_hint;
 
-
-      return GraphAddCallNode(call_node, _GetUniqueName(prim_fn_name), prim_fn_name);
+      // TODO(@jroesch): attach attributes somehow
+      return GraphAddCallNode(call_node, prim_fn_name, GraphAttrs());
     } else {
       ICHECK(false) << "Non-primitive-call nodes should have been transformed away.\n"
                     << "The graph executor code generator expects all calls to have their callee "

src/relay/backend/graph_plan_memory.cc

@@ -34,10 +34,11 @@ namespace tvm {
 namespace relay {
 
 namespace backend {
-  StorageInfo::StorageInfo(std::vector<int64_t> storage_ids, std::vector<DLDeviceType> device_types) {
+  StorageInfo::StorageInfo(std::vector<int64_t> storage_ids, std::vector<DLDeviceType> device_types,  std::vector<int64_t> storage_sizes_in_bytes) {
     auto n = make_object<StorageInfoNode>();
     n->storage_ids = std::move(storage_ids);
     n->device_types = std::move(device_types);
+    n->storage_sizes_in_bytes = std::move(storage_sizes_in_bytes);
     data_ = std::move(n);
   }
 
@@ -233,7 +234,7 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
     for (const auto& kv : token_map_) {
       std::vector<int64_t> storage_ids;
       std::vector<DLDeviceType> device_types;
-      std::vector<Integer> sid_sizes_byte
+      std::vector<int64_t> sid_sizes_byte;
 
       for (StorageToken* tok : kv.second) {
         if (tok->device_type) {
@@ -243,7 +244,7 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
         storage_ids.push_back(tok->storage_id);
         device_types.push_back(static_cast<DLDeviceType>(tok->device_type));
       }
-      auto storage_info = backend::StorageInfo(storage_ids, device_types);
+      auto storage_info = backend::StorageInfo(storage_ids, device_types, sid_sizes_byte);
       smap.Set(GetRef<Expr>(kv.first), storage_info);
     }
     // Either all or none of the nodes should be annotated.

src/relay/backend/utils.h

@@ -99,6 +99,8 @@ class StorageInfoNode : public Object {
   std::vector<int64_t> storage_ids;
   /* \brief The type of "virtual devices" these expressions are stored on. */
   std::vector<DLDeviceType> device_types;
+  /* \brief The sizes of each storage element. */
+  std::vector<int64_t> storage_sizes_in_bytes;
 
   // TODO(@jroesch): expose the fields
   void VisitAttrs(AttrVisitor* v) {}
@@ -110,7 +112,7 @@ class StorageInfoNode : public Object {
 /*! \brief The storage information for a single expression. */
 class StorageInfo : public ObjectRef {
  public:
-  StorageInfo(std::vector<int64_t> storage_ids, std::vector<DLDeviceType> device_types);
+  StorageInfo(std::vector<int64_t> storage_ids, std::vector<DLDeviceType> device_types, std::vector<int64_t> storage_sizes_in_bytes);
   TVM_DEFINE_OBJECT_REF_METHODS(StorageInfo, ObjectRef, StorageInfoNode);
 };