Merge pull request #22 from software-mansion-labs/@ms/xnnpack-input-dim-order-1.4.1

XNNPACK: read inputs with a non-default dim order in their memory layout

Author: msluszniak

backends/xnnpack/test/test_xnnpack_partitioner.py

@@ -161,3 +161,32 @@ def example_inputs(self):
         fwd2_et = executorch_module.run_method("forward_2", example_inputs)
         self.assertTrue(torch.allclose(fwd1_eager, fwd1_et[0], 1e-3))
         self.assertTrue(torch.allclose(fwd2_eager, fwd2_et[0], 1e-3))
+
+    def test_input_with_non_contiguous_dim_order(self):
+        """
+        A delegate input whose dim order is neither contiguous nor channels-last,
+        here (1, 0, 2) because the example input is a transposed view, must be
+        read in its memory layout. Previously its sizes were serialized in
+        logical order while the runtime reshapes it in dim order, so shape
+        propagation failed with xnn_status_invalid_parameter.
+        """
+
+        class BatchMatmul(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.bmm(x, y)
+
+        x = torch.randn(4, 8, 16).transpose(0, 1)
+        y = torch.randn(8, 16, 4)
+        self.assertEqual(x.dim_order(), (1, 0, 2))
+        n = torch.export.Dim("n", min=2, max=8)
+        exported = export(BatchMatmul(), (x, y), dynamic_shapes=({1: n}, {2: n}))
+        executorch = to_edge_transform_and_lower(
+            exported, partitioner=[XnnpackPartitioner()]
+        ).to_executorch()
+
+        # The bindings copy an input's buffer as is, so hand over x's bytes in
+        # its memory order (dim order (1, 0, 2)) under its logical shape.
+        x_in_memory_order = x.transpose(0, 1).contiguous().view(x.shape)
+        executorch_module = _load_for_executorch_from_buffer(executorch.buffer)
+        out = executorch_module.run_method("forward", (x_in_memory_order, y))
+        self.assertTrue(torch.allclose(torch.bmm(x, y), out[0], atol=1e-5))

backends/xnnpack/xnnpack_preprocess.py

@@ -11,12 +11,20 @@
 import torch
 
 from executorch.backends.xnnpack._passes import XNNPACKPassManager
+from executorch.backends.xnnpack._passes.channels_last_tagged_reshape_pass import (
+    ChannelsLastTaggedReshapePass,
+)
 from executorch.backends.xnnpack._passes.convert_to_linear import ConvertToLinearPass
 from executorch.backends.xnnpack.operators.node_visitor import get_node_visitors
 
 from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import (
     ConstantDataOffset,
     XNNGraph,
+    XNNQuantizedTensorValue,
+    XNNStaticTranspose,
+    XNNTensorValue,
+    XNode,
+    XValue,
 )
 from executorch.backends.xnnpack.serialization.xnnpack_graph_serialize import (
     serialize_xnnpack_binary,
@@ -25,6 +33,7 @@
 from executorch.backends.xnnpack.utils.utils import is_param_node
 
 from executorch.backends.xnnpack.utils.xnnpack_constants import (
+    XNN_INVALID_VALUE_ID,
     XNN_VALUE_FLAG_EXTERNAL_INPUT,
     XNN_VALUE_FLAG_EXTERNAL_OUTPUT,
 )
@@ -99,6 +108,86 @@ def assert_default_dim_order(edge_graph_module: torch.fx.GraphModule) -> None:
                 )
 
 
+def define_inputs_in_memory_order(
+    xnn_graph: XNNGraph,
+    node_to_external_map: Dict[torch.fx.Node, ExternalMeta],
+    vals_to_ids: Dict[torch.fx.Node, int],
+) -> None:
+    """
+    Declares each graph input whose dim order is neither contiguous nor
+    channels-last in its memory order, followed by a static transpose back to
+    the logical layout the rest of the XNNPACK graph was built for.
+
+    At runtime XNNExecutor reshapes every external input with its sizes listed
+    in dim order, i.e. as the buffer is laid out in memory. The values here are
+    defined from the logical shape, so an input with, e.g., dim order (1, 0, 2)
+    would otherwise be reshaped to a transposed shape and either fail shape
+    propagation or be read with its axes swapped.
+    """
+    xvalues_by_id = {}
+    for xvalue in xnn_graph.xvalues:
+        union = xvalue.xvalue_union
+        tensor_value = (
+            union.tensor_value if isinstance(union, XNNQuantizedTensorValue) else union
+        )
+        xvalues_by_id[tensor_value.id_out] = xvalue
+
+    prologue = []
+    for node, meta in node_to_external_map.items():
+        if meta.io_type != XNN_VALUE_FLAG_EXTERNAL_INPUT or node not in vals_to_ids:
+            continue
+        val = node.meta.get("val", None)
+        if not isinstance(val, torch.Tensor):
+            continue
+        dim_order = tuple(val.dim_order())
+        if dim_order == tuple(range(val.dim())):
+            continue
+        if ChannelsLastTaggedReshapePass._is_nhwc_tensor(val):
+            # ChannelsLastTaggedReshapePass already bridges NHWC inputs.
+            continue
+
+        logical_id = vals_to_ids[node]
+        xvalue = xvalues_by_id[logical_id]
+        logical = xvalue.xvalue_union
+        if isinstance(logical, XNNQuantizedTensorValue):
+            raise RuntimeError(
+                f"XNNPACK does not support quantized input {node} with dim order "
+                f"{dim_order}; only contiguous or channels-last inputs are supported."
+            )
+
+        memory_id = len(xnn_graph.xvalues)
+        memory = XNNTensorValue(
+            datatype=logical.datatype,
+            num_dims=logical.num_dims,
+            dims=[logical.dims[d] for d in dim_order],
+            constant_buffer_idx=0,
+            external_id=logical.external_id,
+            flags=logical.flags,
+            id_out=memory_id,
+        )
+        xnn_graph.xvalues.append(XValue(xvalue_union=memory))
+        logical.external_id = XNN_INVALID_VALUE_ID
+        logical.flags = 0
+        xnn_graph.input_ids = [
+            memory_id if i == logical_id else i for i in xnn_graph.input_ids
+        ]
+
+        # Logical dim i is memory dim dim_order.index(i).
+        prologue.append(
+            XNode(
+                xnode_union=XNNStaticTranspose(
+                    num_dims=len(dim_order),
+                    perm=[dim_order.index(i) for i in range(len(dim_order))],
+                    input_id=memory_id,
+                    output_id=logical_id,
+                    flags=0,
+                ),
+                debug_handle=DEFAULT_DEBUG_HANDLE,
+            )
+        )
+    xnn_graph.xnodes[:0] = prologue
+
+
 @final
 class XnnpackBackend(BackendDetails):
     @staticmethod
@@ -185,6 +274,8 @@ def preprocess(
                 continue
             else:
                 raise RuntimeError(f"{node.op} is not supported in XNNPACK")
+
+        define_inputs_in_memory_order(xnnpack_graph, node_to_external_map, vals_to_ids)
         return PreprocessResult(
             processed_bytes=serialize_xnnpack_binary(
                 xnnpack_graph, constant_data_bytes