[TVMC] Workspace Pools Parameters

python/tvm/driver/tvmc/compiler.py

@@ -26,6 +26,7 @@
 from tvm import autotvm, auto_scheduler
 from tvm import relay
 from tvm.driver.tvmc.registry import generate_registry_args, reconstruct_registry_entity
+from tvm.ir.memory_pools import WorkspaceMemoryPools
 from tvm.target import Target
 from tvm.relay.backend import Executor, Runtime
 
@@ -37,6 +38,7 @@
 from .pass_list import parse_pass_list_str
 from .transform import convert_graph_layout
 from .shape_parser import parse_shape_string
+from .workspace_pools import generate_workspace_pools_args, workspace_pools_recombobulate
 
 # pylint: disable=invalid-name
 logger = logging.getLogger("TVMC")
@@ -142,10 +144,11 @@ def add_compile_parser(subparsers, _, json_params):
         default="default",
         help="The output module name. Defaults to 'default'.",
     )
-
     for one_entry in json_params:
         parser.set_defaults(**one_entry)
 
+    generate_workspace_pools_args(parser)
+
 
 def drive_compile(args):
     """Invoke tvmc.compiler module with command line arguments
@@ -161,6 +164,7 @@ def drive_compile(args):
         Zero if successfully completed
 
     """
+
     if not os.path.isfile(args.FILE):
         raise TVMCException(
             f"Input file '{args.FILE}' doesn't exist, is a broken symbolic link, or a directory."
@@ -170,6 +174,9 @@ def drive_compile(args):
 
     dump_code = [x.strip() for x in args.dump_code.split(",")] if args.dump_code else None
 
+    additional_targets = reconstruct_target_args(args)
+    workspace_pools_target, extra_targets = target_from_cli(args.target, additional_targets)
+
     compile_model(
         tvmc_model,
         args.target,
@@ -186,8 +193,11 @@ def drive_compile(args):
         desired_layout=args.desired_layout,
         disabled_pass=args.disabled_pass,
         pass_context_configs=args.pass_config,
-        additional_target_options=reconstruct_target_args(args),
         mod_name=args.module_name,
+        additional_target_options=additional_targets,
+        workspace_pools=(
+            workspace_pools_recombobulate(args, [workspace_pools_target], extra_targets)
+        ),
     )
 
     return 0
@@ -212,6 +222,7 @@ def compile_model(
     additional_target_options: Optional[Dict[str, Dict[str, Any]]] = None,
     use_vm: bool = False,
     mod_name: Optional[str] = "default",
+    workspace_pools: Optional[WorkspaceMemoryPools] = None,
 ):
     """Compile a model from a supported framework into a TVM module.
 
@@ -263,6 +274,9 @@ def compile_model(
         Whether to use the VM to compile the model as opposed to the graph executor
     mod_name: str, optional
         The module name
+    workspace_pools: WorkspaceMemoryPools, optional
+        Specification of WorkspacePoolInfo objects to be used as workspace memory in the
+        compilation.
 
     Returns
     -------
@@ -313,6 +327,7 @@ def compile_model(
                         params=params,
                         use_vm=use_vm,
                         mod_name=mod_name,
+                        workspace_pools=workspace_pools,
                     )
         else:
             with autotvm.apply_history_best(tuning_records):
@@ -328,6 +343,7 @@ def compile_model(
                         params=params,
                         use_vm=use_vm,
                         mod_name=mod_name,
+                        workspace_pools=workspace_pools,
                     )
     else:
         with tvm.transform.PassContext(
@@ -342,6 +358,7 @@ def compile_model(
                 params=params,
                 use_vm=use_vm,
                 mod_name=mod_name,
+                workspace_pools=workspace_pools,
             )
 
     # Generate output dump files with sources
@@ -380,6 +397,7 @@ def build(
     params: Dict[str, tvm.nd.NDArray],
     use_vm: bool,
     mod_name: str,
+    workspace_pools: Optional[WorkspaceMemoryPools],
 ):
     """
     Builds the model with the provided executor.
@@ -408,7 +426,13 @@ def build(
         return relay.vm.compile(mod, target=tvm_target, params=params)
     logger.debug("building with relay build")
     return relay.build(
-        mod, target=tvm_target, executor=executor, runtime=runtime, params=params, mod_name=mod_name
+        mod,
+        target=tvm_target,
+        executor=executor,
+        runtime=runtime,
+        params=params,
+        mod_name=mod_name,
+        workspace_memory_pools=workspace_pools,
     )
 
 

python/tvm/driver/tvmc/workspace_pools.py

@@ -0,0 +1,222 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Functions for processing dynamic workspace pool TVMC args
+"""
+
+
+import logging
+import re
+
+from tvm.driver.tvmc import TVMCException
+from tvm.target import Target
+from tvm.ir.memory_pools import PoolInfoProperties, WorkspaceMemoryPools, WorkspacePoolInfo
+
+
+# pylint: disable=invalid-name
+logger = logging.getLogger("TVMC")
+
+
+def generate_workspace_pools_args(parser):
+    """Generates arguments for each Workspace Pools's options"""
+    parser.add_argument(
+        "--workspace-pools",
+        help="""The name of the memory pool
+                Example usage: --workspace-pools=flash""",
+    )
+    parser.add_argument(
+        "--workspace-pools-targets",
+        help="""The name of the targets specified for the memory pool
+                Example usage: --workspace-pools-targets=flash:llvm""",
+        action="append",
+    )
+    parser.add_argument(
+        "--workspace-pools-size-hint-bytes",
+        nargs="?",
+        help="""The expected size hint to be used by the allocator.
+                Example usage: --workspace-pools-size-hint-bytes=flash:8""",
+        action="append",
+    )
+    parser.add_argument(
+        "--workspace-pools-clock-frequency-hz",
+        nargs="?",
+        help="""The clock frequency that the memory pool runs at in Hz.
+                Example usage: --workspace-pools-clock-frequency-hz=flash:70000000""",
+        action="append",
+    )
+    parser.add_argument(
+        "--workspace-pools-read-bandwidth-bytes-per-cycle",
+        nargs="?",
+        help="""The read bandwidth of the memory pool in bytes/cycle.
+                Example usage: --workspace-pools-read-bandwidth-bytes-per-cycle=flash:4""",
+        action="append",
+    )
+    parser.add_argument(
+        "--workspace-pools-write-bandwidth-bytes-per-cycle",
+        nargs="?",
+        help="""The write bandwidth of the memory pool in bytes/cycle.
+                Example usage: --workspace-pools-write-bandwidth-bytes-per-cycle=flash:8""",
+        action="append",
+    )
+    parser.add_argument(
+        "--workspace-pools-read-latency-cycles",
+        nargs="?",
+        help="""The read latency of the memory pool in cycles.
+                Example usage: --workspace-pools-read-latency-cycles=flash:4""",
+        action="append",
+    )
+    parser.add_argument(
+        "--workspace-pools-write-latency-cycles",
+        nargs="?",
+        help="""The write latency of the memory pool in cycles.
+                Example usage: --workspace-pools-write-latency-cycles=flash:8""",
+        action="append",
+    )
+    parser.add_argument(
+        "--workspace-pools-target-burst-bytes",
+        help="""The burst length of the memory pool in bytes per target.
+                Example usage: --workspace-pools-target-burst-bytes=flash:accel:1""",
+        action="append",
+    )
+
+
+def _parse_target_burst(attr, pool_name):
+    if pool_name not in attr:
+        return {}
+
+    return {target: int(attr[pool_name][target]) for target in attr[pool_name]}
+
+
+def _parse_target_access(attr, targets, pool_name):
+    if attr is None:
+        raise TVMCException(f'No target specified for Workspace Pool "{pool_name}"')
+    target_name = attr
+    matched_targets = [target for target in targets if target_name == target.kind.name]
+    if not matched_targets:
+        raise TVMCException(f'Workspace Pool "{pool_name}" using undefined Target "{target_name}"')
+
+    return matched_targets
+
+
+def _split_pools(attr):
+    return re.split(",", attr) if attr else []
+
+
+def _target_attributes_to_pools(attr, targets):
+    if not targets or attr is None:
+        return {}
+
+    target_attributes = {}
+    for pool_values in re.split(",", attr[0]):
+        pool_name, target_name, target_value = re.split(":", pool_values)
+        if pool_name not in target_attributes:
+            target_attributes[pool_name] = {}
+
+        matched_targets = [target for target in targets if target_name == target.kind.name]
+        if matched_targets:
+            target_attributes[pool_name][matched_targets[0]] = target_value
+        else:
+            raise TVMCException(
+                "The workspace pool target specification "
+                "needs to contain a subset of the same TVM "
+                "targets as when specifying targets to use."
+            )
+    return target_attributes
+
+
+def _attribute_to_pools(attr):
+    return dict(pool.split(":", maxsplit=1) for pool in re.split(",", attr[-1])) if attr else {}
+
+
+def workspace_pools_recombobulate(parsed, targets, extra_target):
+    """Reconstructs the Workspace Pools args and returns a WorkspaceMemoryPool object"""
+    WORKSPACE_POOL_PARAMS = [
+        "workspace_pools_size_hint_bytes",
+        "workspace_pools_targets",
+        "workspace_pools_clock_frequency_hz",
+        "workspace_pools_read_bandwidth_bytes_per_cycle",
+        "workspace_pools_write_bandwidth_bytes_per_cycle",
+        "workspace_pools_read_latency_cycles",
+        "workspace_pools_write_latency_cycles",
+    ]
+    WORKSPACE_POOL_TARGET_PARAMS = [
+        "workspace_pools_target_burst_bytes",
+    ]
+
+    # Load extra targets from CLI
+    additional_targets = []
+
+    for t in extra_target:
+        additional_targets.append(Target(t["raw"], host=targets[0].host or targets[0]))
+
+    target = targets + additional_targets
+    if targets[0].host:
+        target.append(targets[0].host)
+
+    workspace_pools = _split_pools(parsed.workspace_pools)
+    if not workspace_pools:
+        return None
+
+    pool_to_attributes = {
+        workspace_pool_param: _attribute_to_pools(getattr(parsed, workspace_pool_param))
+        for workspace_pool_param in WORKSPACE_POOL_PARAMS
+    }
+    pool_to_target_burst_bytes = {
+        workspace_pool_param: _target_attributes_to_pools(
+            getattr(parsed, workspace_pool_param), targets
+        )
+        for workspace_pool_param in WORKSPACE_POOL_TARGET_PARAMS
+    }
+
+    return WorkspaceMemoryPools(
+        [
+            WorkspacePoolInfo(
+                pool_name,
+                targets=_parse_target_access(
+                    pool_to_attributes["workspace_pools_targets"].get(pool_name), target, pool_name
+                ),
+                pool_info_properties=PoolInfoProperties(
+                    size_hint_bytes=int(
+                        pool_to_attributes["workspace_pools_size_hint_bytes"].get(pool_name, -1)
+                    ),
+                    clock_frequency_hz=int(
+                        pool_to_attributes["workspace_pools_clock_frequency_hz"].get(pool_name, -1)
+                    ),
+                    read_bandwidth_bytes_per_cycle=int(
+                        pool_to_attributes["workspace_pools_read_bandwidth_bytes_per_cycle"].get(
+                            pool_name, -1
+                        )
+                    ),
+                    write_bandwidth_bytes_per_cycle=int(
+                        pool_to_attributes["workspace_pools_write_bandwidth_bytes_per_cycle"].get(
+                            pool_name, -1
+                        )
+                    ),
+                    read_latency_cycles=int(
+                        pool_to_attributes["workspace_pools_read_latency_cycles"].get(pool_name, 0)
+                    ),
+                    write_latency_cycles=int(
+                        pool_to_attributes["workspace_pools_write_latency_cycles"].get(pool_name, 0)
+                    ),
+                    target_burst_bytes=_parse_target_burst(
+                        pool_to_target_burst_bytes["workspace_pools_target_burst_bytes"], pool_name
+                    ),
+                ),
+            )
+            for pool_name in workspace_pools
+        ]
+    )

python/tvm/ir/memory_pools.py

@@ -189,7 +189,7 @@ class WorkspaceMemoryPools(Object):
 
     def __init__(
         self,
-        pools: List[PoolInfo],
+        pools: List[WorkspacePoolInfo],
     ):
         self.__init_handle_by_constructor__(
             _ffi_api.WorkspaceMemoryPools, pools  # type: ignore # pylint: disable=no-member

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

@@ -32,7 +32,9 @@ runtime::Module TIRToRuntime(IRModule mod, Target target);
 
 TVM_REGISTER_TARGET_KIND("cmsis-nn", kDLCPU)
     .set_attr<FTVMRelayToTIR>(tvm::attr::kRelayToTIR, RelayToTIR())
-    .set_attr<FTVMTIRToRuntime>("TIRToRuntime", TIRToRuntime);
+    .set_attr<FTVMTIRToRuntime>("TIRToRuntime", TIRToRuntime)
+    .add_attr_option<Array<String>>("mattr")
+    .add_attr_option<String>("mcpu");
 
 }  // namespace cmsisnn
 }  // namespace contrib