From 667011f10320918e4dcd47ac2b57fe49849e5440 Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Thu, 27 May 2021 16:28:57 +0900
Subject: [PATCH] Squashed commit of the following:

commit 95242d86ea5de96925430c0a74b6e91e299fb5ab
Author: Masahiro Masuda <masahi129@gmail.com>
Date:   Thu May 27 15:45:19 2021 +0900

    Add function attribute for shape func for profiling
---
 src/relay/backend/vm/compiler.cc | 3 +++
 src/runtime/vm/profiler/vm.cc    | 4 ++++
 2 files changed, 7 insertions(+)

diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
index ad23e132c486..9a1c61b708e4 100644
--- a/src/relay/backend/vm/compiler.cc
+++ b/src/relay/backend/vm/compiler.cc
@@ -490,6 +490,9 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
       argument_registers.push_back(reg->second);
     }
 
+    // Extract functions attrs
+    op_attrs[op_index] = func->attrs->dict;
+
     Emit(Instruction::InvokePacked(op_index, argument_registers.size(), outputs.size(),
                                    argument_registers));
   }
diff --git a/src/runtime/vm/profiler/vm.cc b/src/runtime/vm/profiler/vm.cc
index 99126b159143..0a7795d600fe 100644
--- a/src/runtime/vm/profiler/vm.cc
+++ b/src/runtime/vm/profiler/vm.cc
@@ -105,6 +105,10 @@ void VirtualMachineDebug::InvokePacked(Index packed_index, const PackedFunc& fun
     }
 
     std::unordered_map<std::string, ObjectRef> metrics;
+
+    ICHECK(exec_->op_attrs.find(packed_index) != exec_->op_attrs.end())
+        << packed_index_map_[packed_index] << " not found in op attrs";
+
     auto& op_attrs = exec_->op_attrs.at(packed_index);
     for (auto p : op_attrs) {
       if (std::string(p.first).find("layout") != std::string::npos) {