[microNPU][4] Add the cascader Proposal generator

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

@@ -0,0 +1,106 @@
+# 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.
+"""Proposal class to hold graph scheduling information."""
+from typing import Dict, FrozenSet, List
+import tvm._ffi
+from tvm.contrib.ethosu.cascader.plan import Plan
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+from .graph import Tensor, Part, CascaderGraph
+from .tensor_config import TensorConfig, MemoryRegion
+
+
+@tvm._ffi.register_object("contrib.ethosu.cascader.Proposal")
+class Proposal(Object):
+    """A class which describes how to schedule a CascaderGraph as a series of disjoint Plans.
+
+    Attributes
+    ----------
+    graph : CascaderGraph
+        The CascaderGraph to which the Proposal applies.
+    part_group : FrozenSet[Part]
+        The Parts which are covered by the Proposal.
+    plans : List[Plan]
+        The Plans used in the Proposal.
+    input_tensor_configs : Dict[Tensor, TensorConfig]
+        The TensorConfigs indexed by Tensor in the Proposal which aren't produced by a Plan.
+    cascade_region : MemoryRegion
+        The MemoryRegion where cascading buffers should be homed.
+    memory_usage : int
+        The memory required to execute the Proposal in the cascading MemoryRegion.
+    cycles : int
+        The estimated cycles taken to execute the Proposal.
+
+    """
+
+    def __init__(
+        self,
+        graph: CascaderGraph,
+        part_group: FrozenSet[Part],
+        plans: List[Plan],
+        input_tensor_configs: Dict[Tensor, TensorConfig],
+        cascade_region: MemoryRegion,
+        memory_usage: Dict[MemoryRegion, int],
+        cycles: int,
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.Proposal,
+            graph,
+            list(part_group),
+            plans,
+            input_tensor_configs,
+            cascade_region,
+            memory_usage,
+            cycles,
+        )
+
+    @property
+    def graph(self) -> CascaderGraph:
+        """The CascaderGraph to which the Proposal applies."""
+        return self._graph
+
+    @property
+    def part_group(self) -> FrozenSet[Part]:
+        """The Parts which are covered by the Proposal."""
+        return frozenset(self._part_group)
+
+    @property
+    def plans(self) -> List[Plan]:
+        """The Plans used in the Proposal."""
+        return list(self._plans)
+
+    @property
+    def input_tensor_configs(self) -> Dict[Tensor, TensorConfig]:
+        """The TensorConfigs indexed by Tensor in the Proposal which aren't produced by a Plan."""
+        return dict(self._input_tensor_configs)
+
+    @property
+    def cascade_region(self) -> MemoryRegion:
+        """The MemoryRegion where cascading buffers should be homed."""
+        return self._cascade_region
+
+    @property
+    def memory_usage(self) -> int:
+        """The memory required to execute the Proposal in the cascading MemoryRegion."""
+        return int(self._memory_usage)
+
+    @property
+    def cycles(self) -> int:
+        """The estimated cycles taken to execute the Proposal."""
+        return int(self._cycles)

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

@@ -0,0 +1,58 @@
+# 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 Proposals for a Graph."""
+from typing import List, Dict, FrozenSet
+
+from . import _ffi_api
+from .cascader_options import CascaderOptions
+from .plan import Plan
+from .proposal import Proposal
+from .graph import CascaderGraph, Part
+
+
+def generate_proposals(
+    graph: CascaderGraph,
+    home_map: Dict[FrozenSet[Part], List[Plan]],
+    options: CascaderOptions,
+) -> List[Proposal]:
+    """Generate Pareto optimal Proposals for a CascaderGraph.
+
+    This algorithm takes a top-down dynamic programming approach to determining how
+    to optimally combine Plans into Proposals.
+
+    Parameters
+    ----------
+    graph : CascaderGraph
+        The CascaderGraph to generate Proposals for.
+    home_map : Dict[FrozenSet[Part], List[Plan]]
+        The Tensor homing map defining valid memory homes for Tensors.
+    options : CascaderOptions
+        The configuration options with which to run the generator.
+
+    Returns
+    ------
+    List[Proposal]
+        A list of Pareto optimal Proposals.
+
+    """
+    return list(
+        _ffi_api.GenerateProposals(
+            graph,
+            home_map,
+            options,
+        )
+    )

src/contrib/ethosu/cascader/pareto.cc

@@ -29,6 +29,8 @@
 
 #include "common.h"
 #include "plan.h"
+#include "proposal.h"
+#include "tensor_config.h"
 
 namespace tvm {
 namespace contrib {
@@ -106,6 +108,31 @@ std::vector<Plan> ParetoCullPlans(std::vector<Plan> plans, size_t max_plans) {
   return ThinVector(optimal_plans, max_plans);
 }
 
+std::vector<Proposal> ParetoCullProposals(std::vector<Proposal> proposals, size_t max_proposals) {
+  std::sort(proposals.begin(), proposals.end(), [](const Proposal& a, const Proposal& b) -> bool {
+    return a->GetMemoryUsage() < b->GetMemoryUsage();
+  });
+  std::vector<std::array<float, 2>> costs;
+  for (const auto& proposal : proposals) {
+    std::array<float, 2> cost = {static_cast<float>(proposal->GetMemoryUsage()),
+                                 static_cast<float>(proposal->GetCycles())};
+    costs.emplace_back(cost);
+  }
+  std::vector<bool> is_optimal = GetParetoFrontier<2>(costs);
+  std::vector<Proposal> optimal_proposals;
+  size_t i = 0;
+  for (bool optimal : is_optimal) {
+    if (optimal) {
+      optimal_proposals.push_back(proposals[i]);
+    }
+    i++;
+  }
+  if (optimal_proposals.size() <= max_proposals) {
+    return optimal_proposals;
+  }
+  return ThinVector(optimal_proposals, max_proposals);
+}
+
 TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.GetParetoFrontier")
     .set_body_typed([](Array<Array<FloatImm>> tcosts) {
       std::vector<std::array<float, 2>> costs;

src/contrib/ethosu/cascader/pareto.h

@@ -37,6 +37,8 @@ namespace ethosu {
 namespace cascader {
 
 class Plan;
+class MemoryRegion;
+class Proposal;
 
 /*!
  * \brief Determine the Pareto optimal points.
@@ -65,6 +67,8 @@ std::vector<T> ThinVector(const std::vector<T>& vec, size_t max_size);
  */
 std::vector<Plan> ParetoCullPlans(std::vector<Plan> plans, size_t max_plans);
 
+std::vector<Proposal> ParetoCullProposals(std::vector<Proposal> proposals, size_t max_proposals);
+
 }  // namespace cascader
 }  // namespace ethosu
 }  // namespace contrib

src/contrib/ethosu/cascader/proposal.cc

@@ -0,0 +1,82 @@
+/*
+ * 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.
+ */
+#include "proposal.h"
+
+#include <tvm/runtime/container/array.h>
+#include <tvm/runtime/container/map.h>
+#include <tvm/runtime/object.h>
+#include <tvm/runtime/registry.h>
+
+#include <algorithm>
+#include <utility>
+#include <vector>
+
+#include "plan.h"
+
+namespace tvm {
+namespace contrib {
+namespace ethosu {
+namespace cascader {
+
+void ProposalNode::VisitAttrs(AttrVisitor* v) {
+  v->Visit("_graph", &graph_);
+  Array<Part> tmp_parts(part_group_.begin(), part_group_.end());
+  v->Visit("_part_group", &tmp_parts);
+  Array<Plan> tmp_plans(plans_.begin(), plans_.end());
+  v->Visit("_plans", &tmp_plans);
+  Map<Tensor, TensorConfig> tmp_tmap(input_tensor_configs_.begin(), input_tensor_configs_.end());
+  v->Visit("_input_tensor_configs", &tmp_tmap);
+  v->Visit("_cascade_region", &cascade_region_);
+  v->Visit("_memory_usage", &memory_usage_);
+  v->Visit("_cycles", &cycles_);
+}
+
+Proposal::Proposal(const CascaderGraph& graph, const std::vector<Part>& part_group,
+                   const std::vector<Plan>& plans, const TensorConfigMap& input_tensor_configs,
+                   const MemoryRegion& cascade_region, int memory_usage, int cycles) {
+  auto n = make_object<ProposalNode>();
+  n->graph_ = std::move(graph);
+  n->part_group_ = std::move(part_group);
+  std::sort(n->part_group_.begin(), n->part_group_.end());
+  n->plans_ = std::move(plans);
+  n->input_tensor_configs_ = std::move(input_tensor_configs);
+  n->cascade_region_ = std::move(cascade_region);
+  n->memory_usage_ = std::move(memory_usage);
+  n->cycles_ = cycles;
+  data_ = std::move(n);
+}
+
+TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.Proposal")
+    .set_body_typed([](CascaderGraph graph, Array<Part> part_group, Array<Plan> plans,
+                       Map<Tensor, TensorConfig> input_tensor_configs, MemoryRegion cascade_region,
+                       int memory_usage, int cycles) {
+      std::vector<Part> spart_group(part_group.begin(), part_group.end());
+      std::vector<Plan> vplans(plans.begin(), plans.end());
+      TensorConfigMap minput_tensor_configs(input_tensor_configs.begin(),
+                                            input_tensor_configs.end());
+      return Proposal(graph, spart_group, vplans, minput_tensor_configs, cascade_region,
+                      memory_usage, cycles);
+    });
+
+TVM_REGISTER_NODE_TYPE(ProposalNode);
+
+}  // namespace cascader
+}  // namespace ethosu
+}  // namespace contrib
+}  // namespace tvm