Develop Dim Order Representation to Accurately Reflect Tensor Memory Format in ExecuTorch

[Gasoonjia]

Context:

Dim order is a tensor-level memory format representation that describes the layout of a dense tensor in memory. It serves as the source of truth for understanding the memory format of input tensors across various components, particularly delegates and operators, in Executorch. The correctness of dim order is crucial to ensure the correct execution of tensor operations. Details can be found in the PyTorch core documentation and ExecuTorch documentation.

State of Dim Order Today:

• Enabling Dim Order as the Edge dialect default
  At this stage, support for dim order in the edge dialect has matured. During the exir phase, we have developed passes that can replace operators requiring memory format as input, such as to_copy, with our own equally performant functions, such as to_dim_order_copy, thus enabling support for multiple dim orders within a model graph. Users can also easily add customized export passes to change the dim order of specific parts of the model. Here is a simple example. Similarly, we have an evening verification mechanism to verify whether the current graph legally supports dim order.

  The current implementation of dim order has shape ambiguity issues, causing some tensors' dim orders to not accurately reflect their target memory formats that prevent us from making it the default state in our edge dialect. This is primarily because the current dim order calculation method in PyTorch Eager (AoT) relies on other tensor attributes (such as stride). This ambiguity leads to inconsistent dim orders for certain tensors blocking us from turning this on by default

• Portable operators and support for Dim Order ops at runtime
  During runtime, we can now use dim order as the basis for determining the memory format of input tensors. For all portable operators, we have added corresponding sanity checks to ensure that the memory format of input tensors remains consistent with what the operator expects. We have also provided special support for the contiguous and channels last dim orders on some operators. For operators that take memory format as input, we have also created corresponding dim order versions, except for empty.memory_format, which is a work-in-progress

• Delegates supporting Dim Order ops AoT and at runtime
  Support from delegates is crucial for dim order. Please refer to this call-for-action for delegate authors here Call for Action: Supporting Dim Order in Backend Delegates #4873, which includes details on delegate support and our next steps.

Next Steps and Timelines:

• Making Dim Order default in the Edge dialect
  We plan to modify the dim order generation approach from calculating based on existing attributes to computing based on the input memory format at tensor creation time to solve the issue, but may need to have a broader discussion with the core PyTorch team. We expected to have one to two weeks to figure out the best way to solve it, and another two to three weeks for implementation.

• Portable operators and support for Dim Order ops
  We estimate that it will take about one to two weeks to complete support for the empty op.

• Delegates supporting Dim Order ops
  XNNPACK and Arm ExecuTorch delegates already support this. We expect other delegates to add at least the minimal support before it becomes a blocking issue for the dim-order workstream

Timeline

By November, we hope to have a comprehensive solution for ambiguous dim orders
By December: support for the dim order version of empty.memory_format.
By next January: resolve the dim order ambiguity issue and make dim order the default in the edge dialect.

@digantdesai @larryliu0820 @mergennachin

cc @JacobSzwejbka @angelayi

[digantdesai]

@Gasoonjia - let's close this?