[microNPU][4] Add the cascader Proposal generator

The Proposal generator takes optimal Plans and combines
them to find optimal 'Proposals' - sets of disjoint
Plans that cover every Part in a CascaderGraph. It
ultimately produces a Pareto-frontier of 'optimal'
Proposals in terms of estimated cycles and memory usage.

Change-Id: Id42099819a596496a5769bae22f08eeb75ec69b6

Author: mbaret

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

@@ -21,6 +21,8 @@
 
 from . import _ffi_api
 from .plan import Plan
+from .proposal import Proposal
+from .tensor_config import MemoryRegion
 
 
 def get_pareto_frontier(costs: List[List[float]]) -> List[bool]:
@@ -37,3 +39,9 @@ def thin_vector(vec: List[Object], max_size: int) -> List[Object]:
 
 def pareto_cull_plans(plans: List[Plan], max_plans: int) -> List[Plan]:
     return list(_ffi_api.ParetoCullPlans(plans, max_plans))
+
+
+def pareto_cull_proposals(
+    proposals: List[Proposal], cascade_region: MemoryRegion, max_proposals: int
+) -> List[Proposal]:
+    return list(_ffi_api.ParetoCullProposals(proposals, cascade_region, max_proposals))

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

@@ -0,0 +1,72 @@
+# 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
+from .tensor_config import TensorConfig, MemoryRegion
+
+
+@tvm._ffi.register_object("contrib.ethosu.cascader.Proposal")
+class Proposal(Object):
+    """Proposal class"""
+
+    def __init__(
+        self,
+        part_group: FrozenSet[Part],
+        plans: List[Plan],
+        input_tensor_configs: Dict[Tensor, TensorConfig],
+        memory_usage: Dict[MemoryRegion, int],
+        cycles: int,
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.Proposal,
+            list(part_group),
+            plans,
+            input_tensor_configs,
+            memory_usage,
+            cycles,
+        )
+
+    @property
+    def graph(self):
+        return self._graph
+
+    @property
+    def part_group(self):
+        return frozenset(self._part_group)
+
+    @property
+    def plans(self):
+        return list(self._plans)
+
+    @property
+    def input_tensor_configs(self):
+        return dict(self._input_tensor_configs)
+
+    @property
+    def memory_usage(self):
+        return int(self._memory_usage)
+
+    @property
+    def cycles(self):
+        return int(self._cycles)

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

@@ -0,0 +1,38 @@
+# 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]:
+    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;
@@ -134,6 +161,12 @@ TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.ParetoCullPlans")
       return Array<Plan>(ParetoCullPlans(vplans, max_size));
     });
 
+TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.ParetoCullProposals")
+    .set_body_typed([](Array<Proposal> proposals, int max_size) {
+      std::vector<Proposal> vproposals(proposals.begin(), proposals.end());
+      return Array<Proposal>(ParetoCullProposals(vproposals, max_size));
+    });
+
 }  // namespace cascader
 }  // namespace ethosu
 }  // namespace contrib

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,81 @@
+/*
+ * 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 <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

src/contrib/ethosu/cascader/proposal.h

@@ -0,0 +1,105 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file src/contrib/ethosu/cascader/proposal.h
+ * \brief Proposal object for the NPU cascader
+ */
+#ifndef TVM_CONTRIB_ETHOSU_CASCADER_PROPOSAL_H_
+#define TVM_CONTRIB_ETHOSU_CASCADER_PROPOSAL_H_
+
+#include <tvm/node/reflection.h>
+#include <tvm/runtime/object.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "graph.h"
+#include "plan.h"
+#include "tensor_config.h"
+
+namespace tvm {
+namespace contrib {
+namespace ethosu {
+namespace cascader {
+
+using MemoryUsageMap = std::unordered_map<MemoryRegion, int, ObjectPtrHash, ObjectPtrEqual>;
+using TensorConfigMap = std::unordered_map<Tensor, TensorConfig, ObjectPtrHash, ObjectPtrEqual>;
+
+/*! \brief Node to represent a Proposal */
+class ProposalNode : public Object {
+ public:
+  void VisitAttrs(AttrVisitor* v);
+
+  /*! \return The CascaderGraph to which the Proposal applies */
+  const CascaderGraph GetGraph() const { return graph_; }
+  /*! \return The Parts which are covered by the Proposal */
+  const std::vector<Part> GetPartGroup() const { return part_group_; }
+  /*! \return The Plans used in the Proposal */
+  const std::vector<Plan> GetPlans() const { return plans_; }
+  /*! \return The TensorConfigs indexed by Tensor in the Proposal which aren't produced by a Plan */
+  const TensorConfigMap GetInputTensorConfigs() const { return input_tensor_configs_; }
+  /*! \return The MemoryRegion where cascading buffers should be homed */
+  const MemoryRegion GetCascadeRegion() const { return cascade_region_; }
+  /*! \return The memory required to execute the Proposal in the cascading MemoryRegion */
+  const int GetMemoryUsage() const { return memory_usage_; }
+  /*! \return The estimated cycles taken to execute the Proposal */
+  int GetCycles() const { return cycles_; }
+
+  static constexpr const char* _type_key = "contrib.ethosu.cascader.Proposal";
+  TVM_DECLARE_FINAL_OBJECT_INFO(ProposalNode, Object);
+
+ protected:
+  friend class Proposal;
+
+  /*! \brief The CascaderGraph to which the Proposal applies */
+  CascaderGraph graph_;
+  /*! \brief The Parts which are covered by the Proposal */
+  std::vector<Part> part_group_;
+  /*! \brief The Plans used in the Proposal */
+  std::vector<Plan> plans_;
+  /*! \brief The TensorConfigs indexed by Tensor in the Proposal which aren't produced by a Plan */
+  TensorConfigMap input_tensor_configs_;
+  /*! \brief The MemoryRegion where cascading buffers should be homed */
+  MemoryRegion cascade_region_;
+  /*! \brief The memory required to execute the Proposal in the cascading MemoryRegion */
+  int memory_usage_;
+  /*! \brief The estimated cycles taken to execute the Proposal */
+  int cycles_;
+};
+
+/*!
+ * \brief A class which describes how to schedule a CascaderGraph as a series of disjoint Plans.
+ */
+class Proposal : public ObjectRef {
+ public:
+  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);
+
+  TVM_DEFINE_OBJECT_REF_METHODS(Proposal, ObjectRef, ProposalNode);
+};
+
+}  // namespace cascader
+}  // namespace ethosu
+}  // namespace contrib
+}  // namespace tvm
+
+#endif  // TVM_CONTRIB_ETHOSU_CASCADER_PROPOSAL_H_