Fixes

Change-Id: I4f5f2a298bd3bb379c7c8d179150358923b0dd66

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

@@ -21,8 +21,6 @@
 
 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]:
@@ -39,9 +37,3 @@ 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

@@ -22,51 +22,85 @@
 from tvm.runtime import Object
 
 from . import _ffi_api
-from .graph import Tensor, Part
+from .graph import Tensor, Part, CascaderGraph
 from .tensor_config import TensorConfig, MemoryRegion
 
 
 @tvm._ffi.register_object("contrib.ethosu.cascader.Proposal")
 class Proposal(Object):
-    """Proposal class"""
+    """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):
+    def graph(self) -> CascaderGraph:
+        """The CascaderGraph to which the Proposal applies."""
         return self._graph
 
     @property
-    def part_group(self):
+    def part_group(self) -> FrozenSet[Part]:
+        """The Parts which are covered by the Proposal."""
         return frozenset(self._part_group)
 
     @property
-    def plans(self):
+    def plans(self) -> List[Plan]:
+        """The Plans used in the Proposal."""
         return list(self._plans)
 
     @property
-    def input_tensor_configs(self):
+    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 memory_usage(self):
+    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):
+    def cycles(self) -> int:
+        """The estimated cycles taken to execute the Proposal."""
         return int(self._cycles)

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

@@ -29,6 +29,26 @@ def generate_proposals(
     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,

src/contrib/ethosu/cascader/pareto.cc

@@ -161,12 +161,6 @@ 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/proposal.cc

@@ -23,6 +23,7 @@
 #include <tvm/runtime/object.h>
 #include <tvm/runtime/registry.h>
 
+#include <algorithm>
 #include <utility>
 #include <vector>
 

src/contrib/ethosu/cascader/proposal_generator.cc

@@ -51,13 +51,12 @@ std::unordered_set<TensorConfig> GetPlanBoundaryConfigs(const Plan& plan) {
   return boundary_configs;
 }
 
-bool IsPlanCompatible(const Proposal& proposal,
-                      const std::vector<Part>& plan_part_group,
+bool IsPlanCompatible(const Proposal& proposal, const std::vector<Part>& plan_part_group,
                       const std::unordered_set<TensorConfig>& plan_boundary_configs) {
   // Check the Plan Part group is disjoint with the Proposal Part group
-  for(const auto& plan_part : plan_part_group) {
-    for(const auto& proposal_part : proposal->GetPartGroup()) {
-      if(plan_part == proposal_part) {
+  for (const auto& plan_part : plan_part_group) {
+    for (const auto& proposal_part : proposal->GetPartGroup()) {
+      if (plan_part == proposal_part) {
         return false;
       }
     }
@@ -126,24 +125,25 @@ Proposal AddPlanToProposal(const Proposal& proposal, const Plan& plan,
   new_memory_usage = std::max(new_memory_usage, proposal->GetMemoryUsage());
   int new_cycles = proposal->GetCycles() + plan->GetCycles();
   std::vector<Part> new_part_group = proposal->GetPartGroup();
-  new_part_group.insert(new_part_group.end(), plan->GetPartGroup().begin(), plan->GetPartGroup().end());
+  new_part_group.insert(new_part_group.end(), plan->GetPartGroup().begin(),
+                        plan->GetPartGroup().end());
   std::sort(new_part_group.begin(), new_part_group.end());
   return Proposal(proposal->GetGraph(), new_part_group, new_plans, new_configs,
                   proposal->GetCascadeRegion(), new_memory_usage, new_cycles);
 }
 
-std::vector<Proposal> GeneratePartialProposals(const CascaderGraph& graph, const HomeMap& home_map,
-                                        const CascaderOptions options,
-                                        const std::unordered_map<Part, std::vector<Plan>, ObjectPtrHash, ObjectPtrEqual>& plans_by_part,
-                                        const std::vector<Part>& partial_proposal_group,
-                                        std::unordered_map<std::vector<Part>, std::vector<Proposal>>* proposals_by_group) {
+std::vector<Proposal> GeneratePartialProposals(
+    const CascaderGraph& graph, const HomeMap& home_map, const CascaderOptions options,
+    const std::unordered_map<Part, std::vector<Plan>, ObjectPtrHash, ObjectPtrEqual>& plans_by_part,
+    const std::vector<Part>& partial_proposal_group,
+    std::unordered_map<std::vector<Part>, std::vector<Proposal>>* proposals_by_group) {
   if (proposals_by_group->find(partial_proposal_group) != proposals_by_group->end()) {
     return proposals_by_group->at(partial_proposal_group);
   }
   if (partial_proposal_group.size() == 0) {
     (*proposals_by_group)[partial_proposal_group] =
-      std::vector<Proposal>{Proposal(graph, std::vector<Part>(), std::vector<Plan>(),
-                                     TensorConfigMap(), options->cascade_region, 0, 0)};
+        std::vector<Proposal>{Proposal(graph, std::vector<Part>(), std::vector<Plan>(),
+                                       TensorConfigMap(), options->cascade_region, 0, 0)};
   } else {
     Part part = partial_proposal_group.back();
     const auto& plans = plans_by_part.at(part);
@@ -158,26 +158,26 @@ std::vector<Proposal> GeneratePartialProposals(const CascaderGraph& graph, const
         // pick the current Plan.
         std::vector<Part> residual_proposal_group;
         std::copy_if(partial_proposal_group.begin(), partial_proposal_group.end(),
-                      std::back_inserter(residual_proposal_group), [&plan](Part value) {
-                        return std::find(plan->GetPartGroup().begin(),
-                                        plan->GetPartGroup().end(),
+                     std::back_inserter(residual_proposal_group), [&plan](Part value) {
+                       return std::find(plan->GetPartGroup().begin(), plan->GetPartGroup().end(),
                                         value) == plan->GetPartGroup().end();
-                      });
+                     });
         // std::sort(residual_proposal_group.begin(), residual_proposal_group.end());
-        const auto& residual_proposals = GeneratePartialProposals(graph, home_map, options, plans_by_part, residual_proposal_group, proposals_by_group);
+        const auto& residual_proposals = GeneratePartialProposals(
+            graph, home_map, options, plans_by_part, residual_proposal_group, proposals_by_group);
         auto plan_output_tensor = plan->GetOutputConfig()->GetTensor();
         ICHECK_LE(plan_output_tensor->GetProducers().size(), 1)
             << "All tensors must have at most one producer.";
         for (const auto& residual_proposal : residual_proposals) {
           if (IsPlanCompatible(residual_proposal, plan->GetPartGroup(), plan_boundary_configs)) {
-            (*proposals_by_group)[partial_proposal_group].push_back(AddPlanToProposal(
-                residual_proposal, plan, plan_boundary_configs));
+            (*proposals_by_group)[partial_proposal_group].push_back(
+                AddPlanToProposal(residual_proposal, plan, plan_boundary_configs));
           }
         }
       }
     }
-    (*proposals_by_group)[partial_proposal_group] = ParetoCullProposals(
-          proposals_by_group->at(partial_proposal_group), options->max_proposals);
+    (*proposals_by_group)[partial_proposal_group] =
+        ParetoCullProposals(proposals_by_group->at(partial_proposal_group), options->max_proposals);
   }
   return proposals_by_group->at(partial_proposal_group);
 }
@@ -194,7 +194,8 @@ std::vector<Proposal> GenerateProposals(const CascaderGraph& graph, const HomeMa
   std::vector<Part> partial_proposal_group = graph->GetPartOrder();
   // A map of Proposals indexed by the Part group they cover
   std::unordered_map<std::vector<Part>, std::vector<Proposal>> proposals_by_group;
-  return GeneratePartialProposals(graph, home_map, options, plans_by_part, partial_proposal_group, &proposals_by_group);
+  return GeneratePartialProposals(graph, home_map, options, plans_by_part, partial_proposal_group,
+                                  &proposals_by_group);
 }
 
 TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.GenerateProposals")