[microNPU] Fix layout assignment in layout optimizer pass (#10143)

Fixes the layout optimizer incorrectly assigning layouts for graphs with
more complex topologies than previously considered. Specifically, this
commit now ensures that intermediate layouts match (e.g. parent output =
child input) and that all consumers are taken into account when altering
the output layout - something not done previously due to an incorrect
traversal order.

Previously, the input layout was always altered if the producer was an
NPU operation without regard to the output layout of that operation.
Additionally, is was possible for the output layout to be incorrectly
set due to a depth-first post-order of traversal of the graph, meaning
it was possible for not all consumers to be taken into account when
altering the layout.

Now the `AnalyzeConsumers` pass is run before `LayoutOptimization` which
determines a mapping from NPU operation to list of boolean values that
represent whether or not each consumer is an NPU operation. Since this
is completed before `LayoutOptimization`, all consumers are guaranteed
to be taken into account when altering the output layout. In turn, the
input layouts can correctly be determined by checking whether the output
of the producer will be altered.

Change-Id: I04e9605da65fa9f12801109dd50c5e3f08cbc73c

Author: lhutton1

python/tvm/relay/backend/contrib/ethosu/codegen.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Codegen for Arm(R) Ethos(TM)-U NPU"""
+from collections import defaultdict
 
 import tvm
 from tvm import relay
@@ -24,7 +25,7 @@
 from tvm.relay.backend.contrib.ethosu.legalize import LegalizeEthosU
 from tvm.relay.backend.contrib.ethosu import tir_to_cs_translator
 from tvm.relay.backend.contrib.ethosu import util
-from tvm.relay.expr_functor import ExprMutator
+from tvm.relay.expr_functor import ExprMutator, ExprVisitor
 
 # pylint: disable=unused-import
 from tvm.relay.backend.contrib.ethosu.op import op_attrs
@@ -138,38 +139,76 @@ def __call__(self, *args, **kwargs):
         pass
 
 
-class LayoutOptimization(ExprMutator):
-    """A pass to optimize the layout of NPU operations. If both the
-    producer and consumer of a tensor are NPU operators, then the
-    layout is converted from NHWC to NHCWB16.
+class AnalyzeConsumers(ExprVisitor):
+    """Traverses the graph to determine consumers that are NPU operations. The
+    result is maintained in `npu_consumers`.
 
     Attributes
     ----------
-    children : Dict[tvm.relay.expr.Call, List[tvm.relay.expr.Call]]
-        A map from current call to a list of calls that rely on the current
-        call. This allows the graph to be traversed backwards, which is useful
-        for checking whether the output layouts can be rewritten.
-    optimize_op : Dict[str, Callable]
-        A map from NPU op name to function that creates NPU op.
+    npu_consumers : Dict[tvm.relay.expr.Call, List[bool]]
+        Mapping from NPU operation to list of boolean values that represent
+        whether or not each consumer is an NPU operation.
+    optimize_ops : Dict[str, Callable]
+        A map from NPU operation name to function that creates NPU operation.
     """
 
-    def __init__(self):
-        self.children = {}
-        self.optimize_op = {
-            "contrib.ethosu.conv2d": op.ethosu_conv2d,
-            "contrib.ethosu.depthwise_conv2d": op.ethosu_depthwise_conv2d,
-            "contrib.ethosu.pooling": op.ethosu_pooling,
-            "contrib.ethosu.binary_elementwise": op.ethosu_binary_elementwise,
-            "contrib.ethosu.unary_elementwise": op.ethosu_unary_elementwise,
-        }
+    def __init__(self, optimize_ops):
+        self.npu_consumers = defaultdict(list)
+        self.optimize_ops = optimize_ops
+        super().__init__()
+
+    def visit_call(self, call: relay.Call):
+        is_npu_consumer = call.op.name in self.optimize_ops
+        args = []
 
+        # Expand tuples
+        for arg in call.args:
+            if isinstance(arg, relay.Tuple):
+                args.extend(arg.fields)
+            else:
+                args.append(arg)
+
+        for arg in args:
+            if isinstance(arg, relay.Call) and arg.op.name in self.optimize_ops:
+                self.npu_consumers[arg].append(is_npu_consumer)
+
+        super().visit_call(call)
+
+
+class LayoutOptimization(ExprMutator):
+    """A pass to optimize the layout of NPU operations by converting to brick format (NHCWB16).
+    This pass traverses the graph and attempts to alter the input/output layouts when an NPU
+    operation is visited. Whether or not the input/output layout can be altered for a given NPU
+    operation depends on the following:
+
+    Check alter input layout: For each argument, if the producer is also an NPU operation and
+        its output is altered to brick format, then the input layout with respect to the current
+        argument is altered to brick format.
+
+    Check alter output layout: If all consumers (child nodes) are an NPU operation, then the
+        output layout is altered to brick format.
+
+    Note
+    ----
+    In order for this pass to be run, the consumers of each NPU operation must first be analyzed
+    by the `AnalyzeConsumers` pass, since Relay doesn't keep a reference to child nodes.
+
+    Attributes
+    ----------
+    npu_consumers : Dict[tvm.relay.expr.Call, bool]
+        A map from current call to a list boolean values that state whether or not each consumer
+        is an NPU operation.
+    optimize_ops : Dict[str, Callable]
+        A map from NPU operation name to function that creates NPU operation.
+    """
+
+    def __init__(self, npu_consumers, optimize_ops):
+        self.npu_consumers = npu_consumers
+        self.optimize_ops = optimize_ops
         super().__init__()
 
     def alter_ethosu_op_layout(self, call: tvm.relay.expr.Call) -> tvm.relay.expr.Call:
-        """Alter the input and output layouts of an NPU operation if needed.
-        Input layout is only altered if the producing operation is an NPU
-        operation. Likewise, the output layout is only altered if the consuming
-        operation is an NPU operation.
+        """Alter the layouts of given NPU operation to brick format if possible.
 
         Parameters
         ----------
@@ -189,46 +228,26 @@ def alter_ethosu_op_layout(self, call: tvm.relay.expr.Call) -> tvm.relay.expr.Ca
         )
 
         new_attrs = dict(call.attrs)
-        parents = []
 
         # Check if we can rewrite the input layouts
         input_count = 0
         for arg in call.args:
             input_count += 1
-            if not isinstance(arg, tvm.relay.expr.Call):
+            if arg not in self.npu_consumers:
                 continue
-            if isinstance(arg.op, tvm.ir.op.Op) and arg.op.name in self.optimize_op:
+            consumers = self.npu_consumers[arg]
+            parent_has_brick_output = consumers and all(consumers)
+            if parent_has_brick_output:
                 layout_string = "ifm_layout" if input_count <= 1 else f"ifm{input_count}_layout"
                 new_attrs[layout_string] = "NHCWB16"
-            parents.append(arg)
 
         # Check if we can rewrite the output layouts
-        if call in self.children:
-            children = self.children[call]
-            if all(
-                isinstance(child, tvm.relay.expr.Call)
-                and isinstance(child.op, tvm.ir.op.Op)
-                and child.op.name in self.optimize_op
-                and child.attrs["ifm_layout"] == "NHCWB16"
-                for child in children
-            ):
-                new_attrs["ofm_layout"] = "NHCWB16"
+        consumers = self.npu_consumers[call]
+        if consumers and all(consumers):
+            new_attrs["ofm_layout"] = "NHCWB16"
 
         name = call.op.name
-        assert name in self.optimize_op, (
-            f"Could not create operator '{name}' as the creation function "
-            "is unknown. Please provide a mapping."
-        )
-        new_call = self.optimize_op[name](*call.args, **new_attrs)
-
-        # Update map of children
-        for input_arg in parents:
-            if input_arg in self.children:
-                self.children[input_arg].append(new_call)
-            else:
-                self.children[input_arg] = [new_call]
-
-        return super().visit_call(new_call)
+        return self.optimize_ops[name](*call.args, **new_attrs)
 
     def visit_call(self, call: tvm.relay.expr.Call) -> tvm.relay.expr.Call:
         """Recursively visit call nodes in the input graph and alter the
@@ -246,23 +265,33 @@ def visit_call(self, call: tvm.relay.expr.Call) -> tvm.relay.expr.Call:
             not refer to an Op. Else, a new call node with altered Op
             attributes.
         """
-        if isinstance(call.op, tvm.ir.op.Op) and call.op.name in self.optimize_op:
-            return self.alter_ethosu_op_layout(call)
+        if isinstance(call.op, tvm.ir.Op) and call.op.name in self.optimize_ops:
+            call = self.alter_ethosu_op_layout(call)
         return super().visit_call(call)
 
 
 @ir.transform.module_pass(opt_level=1, name="LayoutOptimizer")
 class LayoutOptimizer:
     """Register LayoutOptimizer as a Relay pass."""
 
+    OPTIMIZE_OPS = {
+        "contrib.ethosu.conv2d": op.ethosu_conv2d,
+        "contrib.ethosu.depthwise_conv2d": op.ethosu_depthwise_conv2d,
+        "contrib.ethosu.pooling": op.ethosu_pooling,
+        "contrib.ethosu.binary_elementwise": op.ethosu_binary_elementwise,
+        "contrib.ethosu.unary_elementwise": op.ethosu_unary_elementwise,
+    }
+
     def transform_module(self, mod: tvm.ir.IRModule, _) -> tvm.IRModule:
         """A pass to optimize the layout of NPU operations. If both the
         producer and consumer of a tensor are NPU operators, then the
         layout is converted from NHWC to NHCWB16 as this is the layout NPU
         uses internally."""
         assert len(mod.functions.items()) == 1, "Module can only contain one function."
         global_var, func = mod.functions.items()[0]
-        optimized_func = LayoutOptimization().visit(func)
+        analyze = AnalyzeConsumers(self.OPTIMIZE_OPS)
+        analyze.visit(func)
+        optimized_func = LayoutOptimization(analyze.npu_consumers, self.OPTIMIZE_OPS).visit(func)
         mod.update_func(global_var, optimized_func)
         return mod
 

tests/python/contrib/test_ethosu/test_layout_optimizer.py

@@ -538,6 +538,91 @@ def get_graph(get_expected=False):
     _assert_structural_equal(a, b)
 
 
+def test_op_without_ethosu_consumer():
+    """Test the layout optimization pass works as expected when
+    there is a case that the output layout should not be altered
+    since not all consumers are NPU operations (in this case conv).
+
+    depthwise
+        |
+      conv
+      /  \
+     |  pool
+     \   /
+    (concat)
+    """
+
+    def get_graph(get_expected=False):
+        exp_layout = "NHCWB16" if get_expected else "NHWC"
+
+        x = relay.var("x", shape=(1, 2, 2, 2), dtype="int8")
+        depthwise = infra.make_ethosu_depthwise_conv2d(
+            x, 2, (1, 1), (0, 0), (1, 1), (0, 0), ofm_layout=exp_layout
+        )
+        conv = infra.make_ethosu_conv2d(
+            depthwise,
+            2,
+            2,
+            (1, 1),
+            (0, 0),
+            (1, 1),
+            (0, 0),
+            ifm_layout=exp_layout,
+        )
+        pool = infra.make_ethosu_pooling(conv, "MAX", (1, 1), 2, (1, 1), (0, 0))
+        concat = relay.concatenate([conv, pool], axis=0)
+        return relay.Function(relay.analysis.free_vars(concat), concat)
+
+    a = _optimize(get_graph())
+    b = _optimize(get_graph(get_expected=True), optimize=False)
+    _assert_structural_equal(a, b)
+
+
+def test_diamond_graph():
+    """
+    Test the layout optimizer pass works as expected on a diamond graph
+    with a case where the operation dominating the output operation
+    cannot be altered, but operations within the diamond can.
+
+      pool_1
+        |
+      pool_2
+      /   \
+     |  pool_3
+     |     |
+     |  pool_4
+     |     |
+     |  pool_5
+     \    /
+    (concat)
+    """
+
+    def get_graph(get_expected=False):
+        exp_layout = "NHCWB16" if get_expected else "NHWC"
+        x = relay.var("x", shape=(1, 2, 2, 2), dtype="int8")
+        pool_1 = infra.make_ethosu_pooling(
+            x, "MAX", (1, 1), 2, (1, 1), (0, 0), ofm_layout=exp_layout
+        )
+        pool_2 = infra.make_ethosu_pooling(
+            pool_1, "MAX", (1, 1), 2, (1, 1), (0, 0), ifm_layout=exp_layout
+        )
+        pool_3 = infra.make_ethosu_pooling(
+            pool_2, "MAX", (1, 1), 2, (1, 1), (0, 0), ofm_layout=exp_layout
+        )
+        pool_4 = infra.make_ethosu_pooling(
+            pool_3, "MAX", (1, 1), 2, (1, 1), (0, 0), ifm_layout=exp_layout, ofm_layout=exp_layout
+        )
+        pool_5 = infra.make_ethosu_pooling(
+            pool_4, "MAX", (1, 1), 2, (1, 1), (0, 0), ifm_layout=exp_layout
+        )
+        concat = relay.concatenate([pool_2, pool_5], axis=0)
+        return relay.Function(relay.analysis.free_vars(concat), concat)
+
+    a = _optimize(get_graph())
+    b = _optimize(get_graph(get_expected=True), optimize=False)
+    _assert_structural_equal(a, b)
+
+
 def test_same_output_multiple_convolutions():
     """Test running the layout optimization pass with multiple convolutions
     gives same output as TFLite."""