[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, _ = target_from_cli(args.target, additional_targets)
+
     compile_model(
         tvmc_model,
         args.target,
@@ -186,8 +193,9 @@ 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]),
     )
 
     return 0
@@ -212,6 +220,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.
 
@@ -313,6 +322,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 +338,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 +353,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 +392,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 +421,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,199 @@
+# 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.ir.memory_pools import PoolInfo, WorkspaceMemoryPools
+
+
+# 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-size-hint-bytes",
+        nargs="?",
+        help="""The expected size hint to be used by the allocator.
+                Example usage: --workspace-pools-size-hint-bytes=flash:8""",
+    )
+    parser.add_argument(
+        "--workspace-pools-target-access",
+        help="""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 access
+                Example usage: --workspace-pools-target-access=flash:accel:ro""",
+    )
+    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""",
+    )
+    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""",
+    )
+    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""",
+    )
+    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""",
+    )
+    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""",
+    )
+    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""",
+    )
+
+
+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, pool_name):
+    try:
+        return attr[pool_name]
+    except:
+        raise TVMCException(f'Workspace Pool "{pool_name}" using undefined Target.')
+
+
+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):
+        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 for target access "
+                "needs to be the same TVM target 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)) if attr else {}
+
+
+def workspace_pools_recombobulate(parsed, targets):
+    """Reconstructs the Workspace Pools args and returns a WorkspaceMemoryPool object"""
+    WORKSPACE_POOL_PARAMS = [
+        "workspace_pools_size_hint_bytes",
+        "workspace_pools_target_access",
+        "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_pools_target_burst_bytes",
+    ]
+    WORKSPACE_POOL_TARGET_PARAMS = [
+        "workspace_pools_target_access",
+        "workspace_pools_target_burst_bytes",
+    ]
+
+    workspace_pools = _split_pools(parsed.workspace_pools)
+    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_attributes = {
+        workspace_pool_param: _target_attributes_to_pools(
+            getattr(parsed, workspace_pool_param), targets
+        )
+        for workspace_pool_param in WORKSPACE_POOL_TARGET_PARAMS
+    }
+
+    return WorkspaceMemoryPools(
+        [
+            PoolInfo(
+                pool_name,
+                target_access=_parse_target_access(
+                    pool_to_target_attributes["workspace_pools_target_access"], pool_name
+                ),
+                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_attributes["workspace_pools_target_burst_bytes"], pool_name
+                ),
+            )
+            for pool_name in workspace_pools
+        ]
+    )

tests/python/driver/tvmc/test_command_line.py

@@ -21,6 +21,8 @@
 
 from pytest_lazyfixture import lazy_fixture
 from unittest import mock
+
+import tvm
 from tvm.driver.tvmc.main import _main
 from tvm.driver.tvmc.model import TVMCException
 from tvm.driver.tvmc import compiler
@@ -159,6 +161,29 @@ def test_tvmc_tune_file_check(capsys, invalid_input):
     assert captured.err == expected_err, on_assert_error
 
 
+@mock.patch("tvm.relay.build", side_effect=tvm.relay.build)
+@mock.patch("tvm.driver.tvmc.model.TVMCPackage.__init__", return_value=None)
+def test_tvmc_workspace_pools_check(mock_pkg, mock_relay, keras_simple, tmpdir_factory):
+    pytest.importorskip("tensorflow")
+    tmpdir = tmpdir_factory.mktemp("data")
+
+    # Test model compilation
+    package_path = os.path.join(tmpdir, "keras-tvm.tar")
+    compile_str = (
+        f"tvmc compile --target=llvm --workspace-pools=sram "
+        f"--workspace-pools-target-access=sram:llvm:rw "
+        f"--output={package_path} {keras_simple}"
+    )
+    compile_args = compile_str.split(" ")[1:]
+    _main(compile_args)
+    assert os.path.exists(package_path)
+    assert mock_relay.call_count == 1
+    assert mock_relay.call_args_list[0][1]["workspace_memory_pools"].pools[0].pool_name == "sram"
+    assert (
+        len(mock_relay.call_args_list[0][1]["workspace_memory_pools"].pools[0].target_access) == 1
+    )
+
+
 @pytest.fixture
 def paddle_model(paddle_resnet50):
     # If we can't import "paddle" module, skip testing paddle as the input model.

tests/python/driver/tvmc/test_compiler.py

@@ -24,6 +24,8 @@
 import pytest
 
 import tvm
+from tvm.ir.memory_pools import PoolInfo, WorkspaceMemoryPools
+from tvm.target import Target
 import tvm.testing
 from tvm.relay.op.contrib.ethosn import ethosn_available
 from tvm.relay.backend import Runtime, Executor
@@ -674,5 +676,27 @@ def test_compile_tflite_module_with_mod_name_and_ethosu(
             assert b"tvmgen_classify_ethos_u_main_" in content
 
 
+@mock.patch("tvm.relay.build")
+@mock.patch("tvm.driver.tvmc.load")
+@mock.patch("tvm.driver.tvmc.model.TVMCPackage.__init__", return_value=None)
+def test_compile_check_workspace_pools(mock_pkg, mock_fe, mock_relay):
+    mock_fe.return_value = mock.MagicMock()
+    mock_relay.return_value = mock.MagicMock()
+
+    memory_pools = WorkspaceMemoryPools(
+        [PoolInfo(pool_name="sram", target_access={Target("llvm"): "rw"})]
+    )
+
+    tvmc_model = tvmc.load("no_file_needed")
+    tvmc.compile(
+        tvmc_model,
+        target="llvm",
+        workspace_pools=memory_pools,
+    )
+
+    assert mock_relay.call_count == 1
+    assert mock_relay.call_args_list[0][1]["workspace_memory_pools"] == memory_pools
+
+
 if __name__ == "__main__":
     tvm.testing.main()