[TIR][USMP] Added buffer info extraction pass

Updating the USMP utility tests to include tests
that test creation of PoolInfo and PoolAllocation
Objects.

Change-Id: I5d349d0ffcac6b0160072d832dd9d5418699228e

Author: manupak

python/tvm/tir/usmp/utils.py

@@ -17,7 +17,7 @@
 """USMP Utilities and Data Structures"""
 # pylint: disable=invalid-name
 
-from typing import Dict, Optional
+from typing import Dict, Optional, List
 
 from tvm._ffi import register_object
 from tvm.runtime import Object
@@ -27,6 +27,43 @@
 CANDIDATE_MEMORY_POOL_ATTR = "candidate_memory_pools"
 
 
+@register_object("tir.usmp.PoolInfo")
+class PoolInfo(Object):
+    """PoolInfo object holds information related to memory pools
+    where the statically sized allocate nodes will pooled into.
+
+    Parameters
+    ----------
+    pool_name : str
+        The name of the memory pool
+
+    target_access : Dict[Target, str]
+        A dictionary where keys describe which targets could
+        access the pool where value could take the values :
+        a) "rw" : read-write access
+        b) "ro" : write-only acesss
+
+    size_hint_bytes : Optional[int]
+        The expected size hint to be used by the allocator.
+        The default value would be -1 which means the pool
+        is not size restricted.
+
+    """
+
+    READ_WRITE_ACCESS = "rw"
+    READ_ONLY_ACCESS = "ro"
+
+    def __init__(
+        self, pool_name: str, target_access: Dict[Target, str], size_hint_bytes: Optional[int] = -1
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.PoolInfo,  # type: ignore # pylint: disable=no-member
+            pool_name,
+            target_access,
+            size_hint_bytes,
+        )
+
+
 @register_object("tir.usmp.BufferInfo")
 class BufferInfo(Object):
     """BufferInfo object holds information related to buffers
@@ -41,7 +78,10 @@ class BufferInfo(Object):
     size_bytes : int
         The size in bytes
 
-    alignment : int
+    pool_candidates : List[PoolInfo]
+        The list of candidates pools this buffer could be placed
+
+    alignment : Optional[int]
         The byte alignment required in the workspace memory
 
     """
@@ -50,12 +90,14 @@ def __init__(
         self,
         name_hint: str,
         size_bytes: int,
-        alignment: int = None,
+        pool_candidates: List[PoolInfo],
+        alignment: Optional[int] = None,
     ):
         self.__init_handle_by_constructor__(
             _ffi_api.BufferInfo,  # type: ignore # pylint: disable=no-member
             name_hint,
             size_bytes,
+            pool_candidates,
             alignment,
         )
 
@@ -72,43 +114,6 @@ def set_conflicts(self, conflicts: list):
         _ffi_api.BufferInfoSetConflicts(self, conflicts)
 
 
-@register_object("tir.usmp.PoolInfo")
-class PoolInfo(Object):
-    """PoolInfo object holds information related to memory pools
-    where the statically sized allocate nodes will pooled into.
-
-    Parameters
-    ----------
-    pool_name : str
-        The name of the memory pool
-
-    target_access : Dict[Target, str]
-        A dictionary where keys describe which targets could
-        access the pool where value could take the values :
-        a) "rw" : read-write access
-        b) "ro" : write-only acesss
-
-    size_hint_bytes : Optional[int]
-        The expected size hint to be used by the allocator.
-        The default value would be -1 which means the pool
-        is not size restricted.
-
-    """
-
-    READ_WRITE_ACCESS = "rw"
-    READ_ONLY_ACCESS = "ro"
-
-    def __init__(
-        self, pool_name: str, target_access: Dict[Target, str], size_hint_bytes: Optional[int] = -1
-    ):
-        self.__init_handle_by_constructor__(
-            _ffi_api.PoolInfo,  # type: ignore # pylint: disable=no-member
-            pool_name,
-            target_access,
-            size_hint_bytes,
-        )
-
-
 @register_object("tir.usmp.PoolAllocation")
 class PoolAllocation(Object):
     """PoolAllocation object holds information related to an allocation

tests/python/unittest/test_tir_usmp_analysis_extract_bufferinfo.py

@@ -53,7 +53,7 @@ def get_allocate(stmt):
     return allocates
 
 
-def assign_poolinfos_to_allocates_in_primfunc(primfunc, pool_infos):
+def _assign_poolinfos_to_allocates_in_primfunc(primfunc, pool_infos):
     """helper to assing poolinfos to allocate nodes in a tir.PrimFunc"""
 
     def set_poolinfos(stmt):
@@ -70,12 +70,12 @@ def set_poolinfos(stmt):
     return primfunc.with_body(stmt_functor.ir_transform(primfunc.body, None, set_poolinfos))
 
 
-def assign_poolinfos_to_allocates_in_irmodule(mod, pool_infos):
+def _assign_poolinfos_to_allocates_in_irmodule(mod, pool_infos):
     """helper to assing poolinfos to allocate nodes in a IRModule"""
     ret = tvm.IRModule()
     for global_var, basefunc in mod.functions.items():
         if isinstance(basefunc, tvm.tir.PrimFunc):
-            ret[global_var] = assign_poolinfos_to_allocates_in_primfunc(basefunc, pool_infos)
+            ret[global_var] = _assign_poolinfos_to_allocates_in_primfunc(basefunc, pool_infos)
     return ret
 
 
@@ -158,7 +158,7 @@ def test_linear():
         pool_name="slow_memory", target_access={Target("c"): usmp_utils.PoolInfo.READ_WRITE_ACCESS}
     )
     tir_mod = LinearStructure
-    tir_mod = assign_poolinfos_to_allocates_in_irmodule(
+    tir_mod = _assign_poolinfos_to_allocates_in_irmodule(
         tir_mod, [fast_memory_pool, slow_memory_pool]
     )
     buffer_info_map = tvm.tir.usmp.analysis.extract_buffer_info(
@@ -274,7 +274,7 @@ def test_parallel_serial_mixed_for_loops():
         target_access={Target("c"): usmp_utils.PoolInfo.READ_WRITE_ACCESS},
     )
     all_serial_tir_mod = AllSerialForLoops
-    all_serial_tir_mod = assign_poolinfos_to_allocates_in_irmodule(
+    all_serial_tir_mod = _assign_poolinfos_to_allocates_in_irmodule(
         all_serial_tir_mod, [global_ws_pool]
     )
     main_func = all_serial_tir_mod["tvmgen_default_run_model"]
@@ -287,7 +287,7 @@ def test_parallel_serial_mixed_for_loops():
         assert name in ["dummy_allocate", "Conv2dOutput_8", "PaddedInput_8"]
 
     parallel_serial_mixed_tir_mod = ParallelSerialMixedForLoops
-    parallel_serial_mixed_tir_mod = assign_poolinfos_to_allocates_in_irmodule(
+    parallel_serial_mixed_tir_mod = _assign_poolinfos_to_allocates_in_irmodule(
         parallel_serial_mixed_tir_mod, [global_ws_pool]
     )
     main_func = parallel_serial_mixed_tir_mod["tvmgen_default_run_model"]
@@ -634,7 +634,7 @@ def test_inception_structure():
         target_access={Target("c"): usmp_utils.PoolInfo.READ_WRITE_ACCESS},
     )
     tir_mod = InceptionStructure
-    tir_mod = assign_poolinfos_to_allocates_in_irmodule(tir_mod, [global_ws_pool])
+    tir_mod = _assign_poolinfos_to_allocates_in_irmodule(tir_mod, [global_ws_pool])
     main_func = tir_mod["tvmgen_default_run_model"]
     buffer_info_map = tvm.tir.usmp.analysis.extract_buffer_info(main_func, tir_mod)
     buffer_info_map = _replace_stmt_with_buf_var_names(buffer_info_map)