[microNPU] enable USMP

This commit enables USMP in the microNPU codegen
and tests.

Change-Id: Iafd7db8cd678f2b3cca8c06e5ea30e79a570faf9

include/tvm/tir/usmp/utils.h

@@ -251,24 +251,24 @@ struct AllocatedPoolInfoNode : public Object {
   PoolInfo pool_info;
   /*! \brief The allocated size into this pool */
   Integer allocated_size;
-  /*! \brief An optional associated pool Var*/
-  Optional<Var> pool_var;
+  /*! \brief An optional associated pool Var index of PrimFunc params*/
+  Optional<Integer> pool_var_idx;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("pool_info", &pool_info);
     v->Visit("allocated_size", &allocated_size);
-    v->Visit("pool_var", &pool_var);
+    v->Visit("pool_var_idx", &pool_var_idx);
   }
 
   bool SEqualReduce(const AllocatedPoolInfoNode* other, SEqualReducer equal) const {
     return equal(pool_info, other->pool_info) && equal(allocated_size, other->allocated_size) &&
-           equal(pool_var, other->pool_var);
+           equal(pool_var_idx, other->pool_var_idx);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
     hash_reduce(pool_info);
     hash_reduce(allocated_size);
-    hash_reduce(pool_var);
+    hash_reduce(pool_var_idx);
   }
 
   static constexpr const char* _type_key = "tir.usmp.AllocatedPoolInfo";
@@ -277,7 +277,8 @@ struct AllocatedPoolInfoNode : public Object {
 
 class AllocatedPoolInfo : public ObjectRef {
  public:
-  TVM_DLL AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size, Var pool_var = Var());
+  TVM_DLL AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size,
+                            Integer pool_var_idx = Integer());
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(AllocatedPoolInfo, ObjectRef, AllocatedPoolInfoNode);
 };
 

python/tvm/micro/model_library_format.py

@@ -177,14 +177,13 @@ def _build_function_memory_map(function_metadata):
     """
     device_max_workspace = dict()
     main_func_metadata = function_metadata[MAIN_FUNC_NAME_STR]
-    num_targets = len(main_func_metadata.workspace_sizes.items())
+    main_targets = dict(main_func_metadata.workspace_sizes).keys()
     from tvm.driver import tvmc  # pylint: disable=import-outside-toplevel
 
     external_codegens = tvmc.composite_target.get_codegen_names()
     func_entries = []
     target_local_entries = dict()
-    for i in range(num_targets):
-        main_target = main_func_metadata.workspace_sizes.items()[i][0]
+    for main_target in main_targets:
         device_max_workspace[main_target] = 0
         for func_name, finfo in function_metadata.items():
             if func_name == MAIN_FUNC_NAME_STR:
@@ -197,22 +196,18 @@ def _build_function_memory_map(function_metadata):
             # 2. BYOC operator implementations do not currently export useful FunctionInfo.
             if func_name == MAIN_FUNC_NAME_STR or not finfo.tir_primfuncs:
                 continue
-            assert (
-                len(finfo.constant_sizes.items()) == num_targets
-            ), f"{func_name}: found {finfo.constant_sizes!r} vs {num_targets}"
-            assert len(finfo.io_sizes.items()) == num_targets
-            target = finfo.workspace_sizes.items()[i][0]
-            workspace_size = finfo.workspace_sizes.items()[i][1]
-            target_entry = {
-                "device": int(target.kind.device_type),
-                "workspace_size_bytes": int(workspace_size),
-            }
-            target_local_entries[func_name].append(target_entry)
-            if workspace_size > device_max_workspace.get(target, 0):
-                device_max_workspace[target] = workspace_size
-            # TODO(Mousius) - Remove this massive hack when Targets are unified
-            if target.kind.name in external_codegens:
-                device_max_workspace[main_target] += int(workspace_size)
+            if main_target in finfo.workspace_sizes.keys():
+                workspace_size = finfo.workspace_sizes[main_target]
+                target_entry = {
+                    "device": int(main_target.kind.device_type),
+                    "workspace_size_bytes": int(workspace_size),
+                }
+                target_local_entries[func_name].append(target_entry)
+                if workspace_size > device_max_workspace.get(main_target, 0):
+                    device_max_workspace[main_target] = workspace_size
+                # TODO(Mousius) - Remove this massive hack when Targets are unified
+                if main_target.kind.name in external_codegens:
+                    device_max_workspace[main_target] += int(workspace_size)
 
     for func_name, target_entries_ in target_local_entries.items():
         func_entry = {
@@ -222,15 +217,22 @@ def _build_function_memory_map(function_metadata):
         func_entries.append(func_entry)
 
     target_main_entries = list()
-    for i in range(num_targets):
-        target = main_func_metadata.workspace_sizes.items()[i][0]
-        main_func_local_workspace = main_func_metadata.workspace_sizes.items()[i][1]
-        main_func_constants = main_func_metadata.constant_sizes.items()[i][1]
-        main_func_io = main_func_metadata.io_sizes.items()[i][1]
+    for main_target in main_targets:
+        main_func_local_workspace = main_func_metadata.workspace_sizes[main_target]
+        main_func_constants = (
+            main_func_metadata.constant_sizes[main_target]
+            if main_target in main_func_metadata.constant_sizes.keys()
+            else 0
+        )
+        main_func_io = (
+            main_func_metadata.io_sizes[main_target]
+            if main_target in main_func_metadata.io_sizes.keys()
+            else 0
+        )
         target_main_entries.append(
             {
-                "device": int(target.kind.device_type),
-                "workspace_size_bytes": int(device_max_workspace[target])
+                "device": int(main_target.kind.device_type),
+                "workspace_size_bytes": int(device_max_workspace[main_target])
                 + int(main_func_local_workspace),
                 "constants_size_bytes": int(main_func_constants),
                 "io_size_bytes": int(main_func_io),