[microNPU][3] Plan generation for the cascader

The cascader creates 'Plans' which describe how
to schedule subgraphs. As part of the cascading
algorithm, it's necessary to explore a large
variety of Plans which are Pareto optimal (in
terms of memory usage and performance). This is
done by the Plan generation algorithm.

This commit adds the TensorConfig and Plan data
structures which hold information on how to schedule
the tensors/operators. Additionally, it includes
functions to calculate Pareto frontiers which are
used to cull sub-optimal Plans.

Change-Id: Ia358b2a1b29bd810df4441027752ced75812ad4e

Author: mbaret

python/tvm/contrib/ethosu/cascader/__init__.py

@@ -34,3 +34,6 @@
 )
 from .parts import InlinePart, EthosuPart
 from .device_config import EthosuDeviceConfig
+from .tensor_config import TensorConfigState, MemoryRegion, TensorConfig
+from .plan import Plan
+from .cascader_options import CascaderOptions

python/tvm/contrib/ethosu/cascader/cascader_options.py

@@ -0,0 +1,45 @@
+# 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.
+"""Object to hold options for the NPU cascader"""
+import tvm._ffi
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+from .tensor_config import MemoryRegion
+
+
+@tvm._ffi.register_object("contrib.ethosu.cascader.CascaderOptions")
+class CascaderOptions(Object):
+    """A class to hold configuration options for the cascader."""
+
+    def __init__(
+        self,
+        cascade_region: MemoryRegion,
+        max_proposals: int,
+        stripe_factors: int,
+        max_plan_size: int,
+        always_copy_size: int,
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.CascaderOptions,
+            cascade_region,
+            max_proposals,
+            stripe_factors,
+            max_plan_size,
+            always_copy_size,
+        )

python/tvm/contrib/ethosu/cascader/pareto.py

@@ -0,0 +1,39 @@
+# 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.
+"""Pareto optimisation functions for the NPU cascader."""
+from typing import List
+
+from tvm import Object
+
+from . import _ffi_api
+from .plan import Plan
+
+
+def get_pareto_frontier(costs: List[List[float]]) -> List[bool]:
+    for i, cost in enumerate(costs):
+        for j, value in enumerate(cost):
+            costs[i][j] = float(value)
+
+    return [bool(v) for v in _ffi_api.GetParetoFrontier(costs)]
+
+
+def thin_vector(vec: List[Object], max_size: int) -> List[Object]:
+    return list(_ffi_api.ThinVector(vec, max_size))
+
+
+def pareto_cull_plans(plans: List[Plan], max_plans: int) -> List[Plan]:
+    return list(_ffi_api.ParetoCullPlans(plans, max_plans))

python/tvm/contrib/ethosu/cascader/plan.py

@@ -0,0 +1,99 @@
+# 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.
+"""Plan class to hold subgraph scheduling information."""
+from typing import Dict, FrozenSet
+import tvm._ffi
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+from .graph import Tensor, Part
+from .tensor_config import TensorConfig, MemoryRegion
+
+
+@tvm._ffi.register_object("contrib.ethosu.cascader.Plan")
+class Plan(Object):
+    """Plan class"""
+
+    def __init__(
+        self,
+        tensor_configs: Dict[Tensor, TensorConfig],
+        open_configs: FrozenSet[TensorConfig],
+        output_config: TensorConfig,
+        part_group: FrozenSet[Part],
+        interior_region: MemoryRegion,
+        memory_usage: int,
+        cycles: int,
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.Plan,
+            list(tensor_configs.values()),
+            list(open_configs),
+            output_config,
+            list(part_group),
+            interior_region,
+            memory_usage,
+            cycles,
+        )
+
+    def merge(self, other):
+        return _ffi_api.PlanMerge(self, other)
+
+    def benchmark_merge(self, other, repeats):
+        return _ffi_api.PlanMergeBenchmark(self, other, repeats)
+
+    @property
+    def tensor_configs(self):
+        tensor_configs = {}
+        for config in self._tensor_configs:
+            tensor_configs[config.tensor] = config
+        return tensor_configs
+
+    @property
+    def open_configs(self):
+        return frozenset(self._open_configs)
+
+    @property
+    def output_config(self):
+        return self._output_config
+
+    @property
+    def part_group(self):
+        return frozenset(self._part_group)
+
+    @property
+    def interior_region(self):
+        return self._interior_region
+
+    @property
+    def memory_usage(self):
+        return self._memory_usage
+
+    @property
+    def cycles(self):
+        return self._cycles
+
+    def __repr__(self):
+        return (
+            f"Plan(tensor_configs={self.tensor_configs}, "
+            f"open_configs={self.open_configs}, "
+            f"output_config={self.output_config}, "
+            f"part_group={self.part_group}, "
+            f"interior_region={self.interior_region.name}, "
+            f"memory_usage={self.memory_usage}, "
+            f"cycles={self.cycles}, "
+        )

python/tvm/contrib/ethosu/cascader/plan_generator.py

@@ -0,0 +1,51 @@
+# 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.
+"""Algorithms to generate Plans for a CascaderGraph."""
+from typing import List, Dict
+
+from tvm.contrib.ethosu.cascader.tensor_config import MemoryRegion
+
+from . import _ffi_api
+from .cascader_options import CascaderOptions
+from .plan import Plan
+from .stripe_config import StripeConfig
+from .graph import CascaderGraph, Part, Tensor
+
+
+def generate_output_stripe_configs(part: Part, stripe_factors: int) -> List[StripeConfig]:
+    return list(_ffi_api.GenerateOutputStripeConfigs(part, stripe_factors))
+
+
+def generate_single_plans(
+    part: Part,
+    output_stripe_configs: List[StripeConfig],
+    home_map: Dict[Tensor, List[MemoryRegion]],
+    cascade_region: MemoryRegion,
+) -> List[Plan]:
+    return list(_ffi_api.GenerateSinglePlans(part, output_stripe_configs, home_map, cascade_region))
+
+
+def generate_graph_plans(
+    graph: CascaderGraph,
+    home_map: Dict[Tensor, List[MemoryRegion]],
+    options: CascaderOptions,
+):
+    return _ffi_api.GenerateGraphPlans(
+        graph,
+        home_map,
+        options,
+    )

python/tvm/contrib/ethosu/cascader/tensor_config.py

@@ -0,0 +1,118 @@
+# 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.
+"""Tensor config class to hold tensor scheduling information."""
+from typing import List, Union
+from enum import IntEnum
+import tvm._ffi
+from tvm.contrib.ethosu.cascader.stripe_config import StripeConfig
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+from .stripe_config import StripeConfig
+from .graph import Tensor, BufferMode
+
+
+class TensorConfigState(IntEnum):
+    BOUNDARY = 0
+    INTERIOR = 1
+
+
+@tvm._ffi.register_object("contrib.ethosu.cascader.MemoryRegion")
+class MemoryRegion(Object):
+    """MemoryRegion class"""
+
+    def __init__(self, name: str, size: int, read_bandwidth: int, write_bandwidth: int):
+        self.__init_handle_by_constructor__(
+            _ffi_api.MemoryRegion, name, size, read_bandwidth, write_bandwidth
+        )
+
+
+@tvm._ffi.register_object("contrib.ethosu.cascader.TensorConfig")
+class TensorConfig(Object):
+    """TensorConfig class"""
+
+    def __init__(
+        self,
+        tensor: Tensor,
+        home_region: MemoryRegion,
+        state: TensorConfigState,
+        buffer_mode: BufferMode,
+        stripe_configs: List[StripeConfig],
+        copy_tensor: bool = False,
+        copy_region: Union[MemoryRegion, None] = None,
+    ):
+        if copy_region is None:
+            copy_region = home_region
+        self.__init_handle_by_constructor__(
+            _ffi_api.TensorConfig,
+            tensor,
+            home_region,
+            state,
+            buffer_mode,
+            stripe_configs,
+            copy_tensor,
+            copy_region,
+        )
+
+    def get_buffer_size(self):
+        return _ffi_api.TensorConfigGetBufferSize(self)
+
+    @property
+    def tensor(self):
+        return self._tensor
+
+    @property
+    def home_region(self):
+        return self._home_region
+
+    @property
+    def state(self):
+        return TensorConfigState(self._state)
+
+    @property
+    def buffer_mode(self):
+        return BufferMode(self._buffer_mode)
+
+    @property
+    def stripe_configs(self):
+        return list(self._stripe_configs)
+
+    @property
+    def copy_tensor(self):
+        return bool(self._copy_tensor)
+
+    @property
+    def copy_region(self):
+        return self._copy_region
+
+    def __hash__(self):
+        return self._hash
+
+    def __eq__(self, other):
+        return _ffi_api.TensorConfigEqual(self, other)
+
+    def __repr__(self):
+        return (
+            f"TensorConfig(tensor={self.tensor}, "
+            f"home_region={self.home_region.name}, "
+            f"state={self.state.name}, "
+            f"buffer_mode={self.buffer_mode.name}, "
+            f"stripe_configs={self.stripe_configs}, "
+            f"copy_tensor={self.copy_tensor}, "
+            f"copy_region={self.copy_region.name}"
+        )