Improve python docs

Change-Id: I831137f8235665bc20ab4c060cc7049ffd48088a

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

@@ -25,7 +25,23 @@
 
 @tvm._ffi.register_object("contrib.ethosu.cascader.CascaderOptions")
 class CascaderOptions(Object):
-    """A class to hold configuration options for the cascader."""
+    """
+    A class to hold configuration options for the cascader.
+
+    Attributes
+    ----------
+    cascade_region : MemoryRegion
+        The MemoryRegion to place cascading buffers into.
+    max_proposals : int
+        The maximum number of Proposals to generate.
+    stripe_factors : int
+        How many striping factors to try per axis.
+    max_plan_size : int
+        The maximum number of Parts in a Plan.
+    always_copy_size : int
+        The maximum size of a Tensor that will always be copied into the cascade region.
+
+    """
 
     def __init__(
         self,

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

@@ -23,17 +23,17 @@
 from .plan import Plan
 
 
-def get_pareto_frontier(costs: List[List[float]]) -> List[bool]:
+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]:
+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]:
+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

@@ -27,7 +27,43 @@
 
 @tvm._ffi.register_object("contrib.ethosu.cascader.Plan")
 class Plan(Object):
-    """Plan class"""
+    """
+    A class which describes how to schedule a subgraph of Parts together.
+
+    A Plan takes the form of a subgraph of connected Parts (recorded in part_group) with
+    TensorConfigs for all of the required Tensors (recorded in tensor_configs). This information
+    can be used to produce a Tensor Expression schedule with inter-operator scheduling. A Plan is
+    necessarily single-output such that all non-output Parts are 'computed_at'ed the scope of the
+    output Part. This is what achieves the technique referred to as 'cascading'. A Plan also has
+    an interior memory region which specifies the region of memory into which all the Plans
+    intermediate buffers should be allocated.
+
+    Additionally, a Plan contains some other information used during the Plan generation and
+    selection algorithms. Both the memory and cycles required to run the Plan are accounted for so
+    that Plans can be ranked and Pareto-culled on these metrics. Furthermore, the TensorConfigs
+    which are 'open' is recorded indicating that these are valid points to merge with another Plan.
+    A Plan can only be turned into a schedule if it has no 'open' TensorConfigs - at which point
+    the Plan is said to be 'closed'.
+
+    Attributes
+    ----------
+    tensor_configs : Dict[Tensor, TensorConfig]
+        The TensorConfigs specified by the Plan.
+    open_configs : FrozenSet[TensorConfig]
+        The TensorConfigs which are 'open' meaning they are a Plan input/output but have
+        'interior' state.
+    output_config : TensorConfig
+        The TensorConfig of the Plan's output tensor.
+    part_group : FrozenSet[Part]
+        The Parts which are covered by the Plan.
+    interior_region : MemoryRegion
+        The MemoryRegion in which to store 'interior' Plan buffers.
+    memory_usage : int
+        The interior memory used by the Plan in bytes.
+    cycles : int
+        The cycles taken to execute the Plan.
+
+    """
 
     def __init__(
         self,
@@ -51,40 +87,72 @@ def __init__(
         )
 
     def merge(self, other):
-        return _ffi_api.PlanMerge(self, other)
+        """
+        Merge two Plans with share an 'open' TensorConfig.
 
-    def benchmark_merge(self, other, repeats):
-        return _ffi_api.PlanMergeBenchmark(self, other, repeats)
+        The current Plan is referred to as the 'upper Plan' and the other Plan as the 'lower
+        Plan'. The 'open' output config of the upper Plan must be an 'open' input config of the
+        lower Plan. The Tensor referenced by these configs is the Tensor on which the two Plans
+        will be merged. The merge process does the following:
+
+        The tensor config maps will be merged with TensorConfigs from the upper Plan taking
+        priority. The open configs will be merged with the TensorConfigs that are being merged
+        having been removed. The output config will be that of the lower Plan. The part groups
+        will be merged. The interior region is necessarily the same for both the upper and lower
+        Plan. The cycles and memory usage will be summed.
+
+        Parameters
+        ----------
+        other : Plan
+            The Plan to merge with.
+
+        Return
+        ------
+        Plan
+            The merged Plan.
+
+        """
+        return _ffi_api.PlanMerge(self, other)
 
     @property
     def tensor_configs(self):
+        """The TensorConfigs specified by the Plan."""
         tensor_configs = {}
         for config in self._tensor_configs:
             tensor_configs[config.tensor] = config
         return tensor_configs
 
     @property
     def open_configs(self):
+        """
+        The TensorConfigs which are 'open' meaning they are a Plan input/output but have
+        'interior' state.
+        """
         return frozenset(self._open_configs)
 
     @property
     def output_config(self):
+        """The TensorConfig of the Plan's output tensor."""
         return self._output_config
 
     @property
     def part_group(self):
+        """The Parts which are covered by the Plan."""
         return frozenset(self._part_group)
 
     @property
     def interior_region(self):
+        """The MemoryRegion in which to store 'interior' Plan buffers."""
         return self._interior_region
 
     @property
     def memory_usage(self):
+        """The interior memory used by the Plan in bytes."""
         return self._memory_usage
 
     @property
     def cycles(self):
+        """The cycles taken to execute the Plan."""
         return self._cycles
 
     def __repr__(self):

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

@@ -26,11 +26,11 @@
 from .graph import CascaderGraph, Part, Tensor
 
 
-def generate_output_stripe_configs(part: Part, stripe_factors: int) -> List[StripeConfig]:
+def _generate_output_stripe_configs(part: Part, stripe_factors: int) -> List[StripeConfig]:
     return list(_ffi_api.GenerateOutputStripeConfigs(part, stripe_factors))
 
 
-def generate_single_plans(
+def _generate_single_plans(
     part: Part,
     output_stripe_configs: List[StripeConfig],
     home_map: Dict[Tensor, List[MemoryRegion]],
@@ -39,7 +39,7 @@ def generate_single_plans(
     return list(_ffi_api.GenerateSinglePlans(part, output_stripe_configs, home_map, cascade_region))
 
 
-def generate_graph_plans(
+def _generate_graph_plans(
     graph: CascaderGraph,
     home_map: Dict[Tensor, List[MemoryRegion]],
     options: CascaderOptions,

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

@@ -28,13 +28,35 @@
 
 
 class TensorConfigState(IntEnum):
+    """
+    The 'state' of a TensorConfig as used in the Plan generation algorithm.
+
+    BOUNDARY - Should describe a Plan input/output Tensor.
+    INTERIOR - Should describe an intermediate Tensor in a 'closed' Plan.
+
+    """
+
     BOUNDARY = 0
     INTERIOR = 1
 
 
 @tvm._ffi.register_object("contrib.ethosu.cascader.MemoryRegion")
 class MemoryRegion(Object):
-    """MemoryRegion class"""
+    """
+    MemoryRegion class to store information about device memories.
+
+    Attributes
+    ----------
+    name : str
+        The name of the region.
+    size : int
+        The size of the region.
+    read_bandwidth : int
+        The read bandwidth of the region in bytes per cycle.
+    write_bandwidth : int
+        The write bandwidth of the region in bytes per cycle.
+
+    """
 
     def __init__(self, name: str, size: int, read_bandwidth: int, write_bandwidth: int):
         self.__init_handle_by_constructor__(
@@ -44,7 +66,55 @@ def __init__(self, name: str, size: int, read_bandwidth: int, write_bandwidth: i
 
 @tvm._ffi.register_object("contrib.ethosu.cascader.TensorConfig")
 class TensorConfig(Object):
-    """TensorConfig class"""
+    """
+    A class which describes how to realize a Tensor.
+
+    The TensorConfig describes both how a Tensor is scheduled (the order in which it's
+    produced/consumed) and how its allocated in memory (which region it should reside in
+    and whether it should be copied).
+
+    Attributes
+    ----------
+    tensor : Tensor
+        The Tensor the config applies to.
+    home_region : MemoryRegion
+        The region where the tensor is allocated.
+    state : TensorConfigState
+        The state of the TensorConfig.
+
+        The TensorConfigState is only used as part of the Plan generation algorithm. For a Plan
+        to be 'closed' (and therefore not subject to any further merging), all the TensorConfigs
+        that describe Plan input or output Tensors must be in the 'BOUNDARY' state with the rest
+        being 'INTERIOR'. If any of the input or output tensors are described by an 'INTERIOR'
+        TensorConfig, then the Plan is 'open' and should be merged with other 'open' Plans until
+        the result becomes 'closed'.
+    buffer_mode : BufferMode
+        The mode in which the buffer should be realized.
+
+        There are multiple buffering strategies by which a tensor may be realized (computed).
+        These affect the amount of recomputation necessary as well as the size of buffer required
+        to store the tensor. See 'BufferMode' for a description of the allowable buffering modes.
+    stripe_configs : List[StringConfig]
+       The StripeConfigs with which to compute the tensor.
+
+       The StripeConfigs determine the order in which the elements of the tensor should be
+       computed, including potentially computing them multiple times (recompute). Multiple
+       StripeConfigs are used over just a single StripeConfig for the case where the tensor is
+       consumed by two different Parts executing themselves with different StripeConfigs. In this
+       case, there is a StripeConfig per consumer of the tensor.
+    copy_tensor : bool, optional
+        Whether to copy the tensor.
+
+        While a tensor will originally reside in its home region, the TensorConfig may optionally
+        specify that the tensor should be copied (according to the StripeConfigs) into another
+        MemoryRegion. As an example for where this may be used, if a weights tensor initially
+        resides in slow Flash memory then necessarily the home region will be Flash. However, if
+        the weights values are used multiple times by a Part, it may be more performant to choose
+        to copy the weights into a faster memory like SRAM.
+    copy_region : Union[MemoryRegion, None], optional
+        The region to copy the tensor to.
+
+    """
 
     def __init__(
         self,
@@ -70,34 +140,52 @@ def __init__(
         )
 
     def get_buffer_size(self):
+        """
+        The size of the buffer needed for the TensorConfig.
+
+        The size of buffer necessary to store a tensor being produced using the TensorConfig is
+        not necessarily just the size of the tensor. In Plans, a tensor may be being produced and
+        consumed in 'stripes' which are smaller than the full tensor. Therefore, the buffer
+        necessary to store the tensor may only need to be as large as the stripe. The precise size
+        of the buffer will depend both on the BufferMode and StripeConfigs (as well as, of course,
+        the Tensor).
+
+        """
         return _ffi_api.TensorConfigGetBufferSize(self)
 
     @property
     def tensor(self):
+        """The Tensor the config applies to."""
         return self._tensor
 
     @property
     def home_region(self):
+        """The region where the tensor is allocated."""
         return self._home_region
 
     @property
     def state(self):
+        """The state of the TensorConfig."""
         return TensorConfigState(self._state)
 
     @property
     def buffer_mode(self):
+        """The mode in which the buffer should be realized."""
         return BufferMode(self._buffer_mode)
 
     @property
     def stripe_configs(self):
+        """The StripeConfigs with which to compute the tensor."""
         return list(self._stripe_configs)
 
     @property
     def copy_tensor(self):
+        """Whether to copy the tensor."""
         return bool(self._copy_tensor)
 
     @property
     def copy_region(self):
+        """The region to copy the tensor to."""
         return self._copy_region
 
     def __hash__(self):

tests/python/contrib/test_ethosu/cascader/test_pareto.py

@@ -15,7 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 from tvm.tir import IntImm
-from tvm.contrib.ethosu.cascader.pareto import get_pareto_frontier, thin_vector, pareto_cull_plans
+from tvm.contrib.ethosu.cascader.pareto import (
+    _get_pareto_frontier,
+    _thin_vector,
+    _pareto_cull_plans,
+)
 from tvm.contrib.ethosu.cascader import (
     Plan,
     StripeConfig,
@@ -78,7 +82,7 @@ def test_get_pareto_frontier(num_costs):
     for i in range(num_costs):
         costs.append(list(np.random.randint(cost_low, cost_high, size=(dims,))))
     reference = list(_ref_get_pareto_frontier(np.array(costs)))
-    result = get_pareto_frontier(costs)
+    result = _get_pareto_frontier(costs)
     assert result == reference
 
 
@@ -95,7 +99,7 @@ def _make_vector(length):
 
     vector = _make_vector(vec_length)
     reference = list(_ref_thin_vector(np.array(vector), max_size))
-    result = thin_vector(vector, max_size)
+    result = _thin_vector(vector, max_size)
     assert result == reference
 
 
@@ -137,7 +141,7 @@ def _make_plans(num):
 
     plans = _make_plans(num_plans)
     reference = list(_ref_pareto_cull_plans(plans, max_plans))
-    result = pareto_cull_plans(plans, max_plans)
+    result = _pareto_cull_plans(plans, max_plans)
     assert result == reference