[microNPU] Adding a option to disable striping

This commit adds a cascader option to disable striping. Change-Id: Ie32cef85814ce380a77daa4f8cf1f607891c9c69
apache · May 10, 2022 · 258715b · 258715b
1 parent 8d4f4dd
commit 258715b
Show file tree

Hide file tree

Showing 9 changed files with 78 additions and 14 deletions.
diff --git a/python/tvm/contrib/ethosu/cascader/cascader_options.py b/python/tvm/contrib/ethosu/cascader/cascader_options.py
@@ -40,6 +40,8 @@ class CascaderOptions(Object):
         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.
+    disable_striping : bool
+        A boolean option to disable striping
 
     """
 
@@ -50,6 +52,7 @@ def __init__(
         stripe_factors: int,
         max_plan_size: int,
         always_copy_size: int,
+        disable_striping: bool = False,
     ):
         self.__init_handle_by_constructor__(
             _ffi_api.CascaderOptions,
@@ -58,4 +61,5 @@ def __init__(
             stripe_factors,
             max_plan_size,
             always_copy_size,
+            disable_striping,
         )
diff --git a/python/tvm/contrib/ethosu/cascader/device_config.py b/python/tvm/contrib/ethosu/cascader/device_config.py
@@ -594,7 +594,7 @@ def _get_subkernel_propagator(
             if output_layout == "NHCWB16" and input_layout == "NHWC":
                 transform[3][-1] = depth
             elif output_layout == "NHCWB16" and input_layout == "NHCWB16":
-                transform[2][-1] = depth // 16
+                transform[2][-1] = 1 + ((depth - 1) // 16)
 
         return Propagator(transform, ifm_propagator.offset)
 

diff --git a/python/tvm/contrib/ethosu/cascader/plan_generator.py b/python/tvm/contrib/ethosu/cascader/plan_generator.py
@@ -26,8 +26,10 @@
 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_output_stripe_configs(
+    part: Part, stripe_factors: int, disable_striping: bool
+) -> List[StripeConfig]:
+    return list(_ffi_api.GenerateOutputStripeConfigs(part, stripe_factors, disable_striping))
 
 
 def _generate_single_plans(

diff --git a/src/contrib/ethosu/cascader/cascader_options.cc b/src/contrib/ethosu/cascader/cascader_options.cc
@@ -31,24 +31,27 @@ void CascaderOptionsNode::VisitAttrs(AttrVisitor* v) {
   v->Visit("stripe_factors", &stripe_factors);
   v->Visit("max_plan_size", &max_plan_size);
   v->Visit("always_copy_size", &always_copy_size);
+  v->Visit("disable_striping", &disable_striping);
 }
 
 CascaderOptions::CascaderOptions(const MemoryRegion& cascade_region, int max_proposals,
-                                 int stripe_factors, int max_plan_size, int always_copy_size) {
+                                 int stripe_factors, int max_plan_size, int always_copy_size,
+                                 bool disable_striping) {
   auto n = make_object<CascaderOptionsNode>();
   n->cascade_region = std::move(cascade_region);
   n->max_proposals = max_proposals;
   n->stripe_factors = stripe_factors;
   n->max_plan_size = max_plan_size;
   n->always_copy_size = always_copy_size;
+  n->disable_striping = disable_striping;
   data_ = std::move(n);
 }
 
 TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.CascaderOptions")
     .set_body_typed([](MemoryRegion cascade_region, int max_proposals, int stripe_factors,
-                       int max_plan_size, int always_copy_size) {
+                       int max_plan_size, int always_copy_size, bool disable_striping) {
       return CascaderOptions(cascade_region, max_proposals, stripe_factors, max_plan_size,
-                             always_copy_size);
+                             always_copy_size, disable_striping);
     });
 
 TVM_REGISTER_NODE_TYPE(CascaderOptionsNode);

diff --git a/src/contrib/ethosu/cascader/cascader_options.h b/src/contrib/ethosu/cascader/cascader_options.h
@@ -49,6 +49,8 @@ class CascaderOptionsNode : public Object {
   int max_plan_size;
   /*! \brief The maximum size of Tensor that will always be copied into the cascade region. */
   int always_copy_size;
+  /*! \brief A boolean option to disable striping. */
+  bool disable_striping;
 
   static constexpr const char* _type_key = "contrib.ethosu.cascader.CascaderOptions";
   TVM_DECLARE_FINAL_OBJECT_INFO(CascaderOptionsNode, Object)
@@ -58,7 +60,7 @@ class CascaderOptionsNode : public Object {
 class CascaderOptions : public ObjectRef {
  public:
   CascaderOptions(const MemoryRegion& cascade_region, int max_proposals, int stripe_factors,
-                  int max_plan_size, int always_copy_size);
+                  int max_plan_size, int always_copy_size, bool disable_striping = false);
 
   TVM_DEFINE_OBJECT_REF_METHODS(CascaderOptions, ObjectRef, CascaderOptionsNode);
 };

diff --git a/src/contrib/ethosu/cascader/plan_generator.cc b/src/contrib/ethosu/cascader/plan_generator.cc
@@ -105,7 +105,8 @@ std::vector<bool> GetCascadableAxes(const Part& part) {
   return cascadable_axes;
 }
 
-std::vector<StripeConfig> GenerateOutputStripeConfigs(const Part& part, int stripe_factors) {
+std::vector<StripeConfig> GenerateOutputStripeConfigs(const Part& part, int stripe_factors,
+                                                      bool disable_striping = false) {
   // If stripe_factors is <= 0, then we won't produce any StripeConfigs
   if (stripe_factors <= 0) {
     return std::vector<StripeConfig>();
@@ -134,7 +135,7 @@ std::vector<StripeConfig> GenerateOutputStripeConfigs(const Part& part, int stri
     auto axis = output_shape[i];
     auto axis_align = part->GetStripeAlignHint()[i];
     std::set<int> axis_splits;  // Note this is a set to remove duplicate splits
-    if (!cascadable_axes[i]) {
+    if (!cascadable_axes[i] || disable_striping) {
       axis_splits.insert(axis);
     } else {
       for (float factor : factors) {
@@ -436,7 +437,7 @@ std::unordered_map<std::vector<Part>, std::vector<Plan>> GenerateGraphPlans(
     // output of a Plan. The number generated is a function of stripe_factors and the number of
     // cascadable dimensions in the Part.
     std::vector<StripeConfig> stripe_configs =
-        GenerateOutputStripeConfigs(part, options->stripe_factors);
+        GenerateOutputStripeConfigs(part, options->stripe_factors, options->disable_striping);
     // Check to see if the output Tensor is part of any existing open Plans
     if (stripe_configs_by_tensor.find(part->GetOutputTensor()) != stripe_configs_by_tensor.end()) {
       // If there are other open Plans which have this Part's output Tensor as an input, then
@@ -514,11 +515,12 @@ std::unordered_map<std::vector<Part>, std::vector<Plan>> GenerateGraphPlans(
 }
 
 TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.GenerateOutputStripeConfigs")
-    .set_body_typed([](Part part, int stripe_factors) {
+    .set_body_typed([](Part part, int stripe_factors, bool disable_striping) {
       if (stripe_factors < 0) {
         return Array<StripeConfig>();
       }
-      return Array<StripeConfig>(GenerateOutputStripeConfigs(part, stripe_factors));
+      return Array<StripeConfig>(
+          GenerateOutputStripeConfigs(part, stripe_factors, disable_striping));
     });
 
 TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.GenerateSinglePlans")

diff --git a/tests/python/contrib/test_ethosu/cascader/infra.py b/tests/python/contrib/test_ethosu/cascader/infra.py
@@ -32,13 +32,15 @@ def make_options(
     stripe_factors: int = 1,
     max_plan_size: int = 1,
     always_copy_size: int = 1024,
+    disable_striping: bool = False,
 ):
     return cs.CascaderOptions(
         cascade_region=cascade_region,
         max_proposals=max_proposals,
         stripe_factors=stripe_factors,
         max_plan_size=max_plan_size,
         always_copy_size=always_copy_size,
+        disable_striping=disable_striping,
     )
 
 

diff --git a/tests/python/contrib/test_ethosu/cascader/test_plan_generator.py b/tests/python/contrib/test_ethosu/cascader/test_plan_generator.py
@@ -47,7 +47,33 @@ def test_generate_output_stripe_configs():
     tensor_1.add_consumer(part_1)
     tensor_2.add_producer(part_1)
 
-    assert len(_generate_output_stripe_configs(part_1, stripe_factors)) == expected_configs
+    assert (
+        len(_generate_output_stripe_configs(part_1, stripe_factors, disable_striping=False))
+        == expected_configs
+    )
+
+
+@pytest.mark.parametrize("stripe_factors", [3, 4, 8, 16, 10])
+def test_generate_output_stripe_configs_disable_striping(stripe_factors):
+    subgraph = cs.TESubgraph([], None)
+    part_1 = cs.InlinePart(
+        subgraph,
+        [
+            cs.Propagator(
+                [[2, 0, 0], [0, 2, 0], [0, 0, 1]],
+                [0, 0],
+            ),
+        ],
+    )
+    tensor_1 = cs.Tensor([800, 800], "uint8")
+    tensor_2 = cs.Tensor([400, 400], "uint8")
+
+    part_1.set_input(0, tensor_1)
+    part_1.set_output(tensor_2)
+    tensor_1.add_consumer(part_1)
+    tensor_2.add_producer(part_1)
+
+    assert len(_generate_output_stripe_configs(part_1, stripe_factors, disable_striping=True)) == 1
 
 
 def test_generate_single_plans(SRAM, DRAM):
@@ -74,7 +100,9 @@ def test_generate_single_plans(SRAM, DRAM):
         tensor_2: [SRAM],
     }
     options = make_options(cascade_region=SRAM, stripe_factors=1)
-    output_stripe_configs = _generate_output_stripe_configs(part_1, options.stripe_factors)
+    output_stripe_configs = _generate_output_stripe_configs(
+        part_1, options.stripe_factors, disable_striping=False
+    )
     plans = _generate_single_plans(part_1, output_stripe_configs, home_map, options)
     for plan in plans:
         assert plan.interior_region == SRAM

diff --git a/tests/python/contrib/test_ethosu/cascader/test_proposal_generator.py b/tests/python/contrib/test_ethosu/cascader/test_proposal_generator.py
@@ -134,6 +134,27 @@ def test_generate_proposals_mobilenetv2diamond(FLASH, SRAM, MobileNetv2DiamondGr
             assert min_sram < proposal.memory_usage < max_sram
             assert proposal.cycles > 0
 
+    def test_generate_proposals_mobilenetv1_disable_striping(FLASH, SRAM, MobileNetv1Graph):
+        graph = MobileNetv1Graph
+        home_map = make_simple_home_map(graph, SRAM, FLASH)
+        options = make_options(
+            cascade_region=SRAM,
+            max_proposals=32,
+            stripe_factors=5,
+            max_plan_size=10,
+            disable_striping=True,
+        )
+
+        proposals = generate_proposals(graph, home_map, options)
+        assert len(proposals) == 1
+        proposal = proposals[0]
+        for plan in proposal.plans:
+            for stripe_config in plan.output_config.stripe_configs:
+                for shape_dim, stride_dim in list(zip(stripe_config.shape, stripe_config.strides)):
+                    # The striding and shape sizes in each dimension should be the same
+                    # if striping is disabled
+                    assert int(shape_dim) == int(stride_dim)
+
 
 if __name__ == "__main__":
     pytest.main([__file__])