Add 3D (NCW) channels-last dim_order support to Neutron runtime (#19341)

[JakeStevens]

Summary:

Although 3dim dim order is not supported pytorch, it actually is properly supported within the NXP backend's implementation.

The only issue is within the runtime, which forces the check to four dimensions.

This PR relaxes and supports 3D.

Differential Revision: D102862305

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19341

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Run pull request jobs on OSDC runners in shadow mode

❌ 2 New Failures, 5 Unrelated Failures

As of commit 1ef67ca with merge base e4f5b38 ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / macos / macos-job (gh)
  exir/tests/test_joint_graph.py::TestJointGraph::test_joint_graph
• pull / unittest-nxp-neutron / linux-job (gh)
  RuntimeError: Command docker exec -t a78ebb27dc30e94be7547b4e08d875d0ee0586f08af44b9bde3752747b193cc6 /exec failed with exit code 1

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / unittest-editable / linux / linux-job (gh) (similar failure)
  exir/tests/test_memory_planning.py::TestMisc::test_multiple_pools_1

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / linux / linux-job (gh) (trunk failure)
  exir/tests/test_memory_planning.py::TestMisc::test_multiple_pools_1
• pull / unittest / windows / windows-job (gh) (trunk failure)
  backends/xnnpack/test/recipes/test_xnnpack_recipes.py::TestXnnpackRecipes::test_8a4w_recipe
• pull / unittest-editable / macos / macos-job (gh) (trunk failure)
  exir/tests/test_memory_planning.py::TestMisc::test_multiple_pools_1
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  backends/xnnpack/test/recipes/test_xnnpack_recipes.py::TestXnnpackRecipes::test_8a4w_recipe

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-codesync]

@JakeStevens has exported this pull request. If you are a Meta employee, you can view the originating Diff in D102862305.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[jirioc]

The implementation looks good. Note that sometimes the transpositions are done by the neutron instead of the backend (scheduled by AoT), so I you need to make sure that that part is also compatible with 3D.

[MartinPavella]

LGTM. But perhaps some test with a simple 3D channels last dim order model would be useful.

[JakeStevens]

LGTM. But perhaps some test with a simple 3D channels last dim order model would be useful.

Are there examples in the repo of cpp tests that use the backend cpp file?

[MartinPavella]

Are there examples in the repo of cpp tests that use the backend cpp file?

You can use a test like this:

class Conv1DModel(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv1d(2, 2, 3)

    def forward(self, x):
        return self.conv(x)


def test_3d_channels_last_dim_order__conv():
    model = Conv1DModel().eval().to(memory_format=torch.channels_last)

    # Use the `channels last` dim order.
    input_spec = [
        ModelInputSpec((1, 2, 5), dim_order=torch.channels_last)
    ]

    comparator = NumericalStatsOutputComparator(max_mse_error=8e-3)
    lower_run_compare(
        model,
        input_spec,
        dataset_creator=RandomDatasetCreator(),
        output_comparator=comparator,
        dlg_model_verifier=BaseGraphVerifier(1, []),
        reference_model=ReferenceModel.QUANTIZED_EDGE_PYTHON
    )

[MartinPavella]

According to your modifications to executorch_pipeline, the example input used to export this model will use the contiguous memory format. Also, you are not calling model.to(memory_format=torch.channels_last) here. If I understand correctly, this is because torchdoesn't support 3D channels last memory format. But how will thedim_orderentry in the nodes'meta` attribute be set? I.e. how will ExecuTorch know to use the 3D channels last dim_order?

To me, it looks like the test is just using regular contiguous dim_order, and the specified torch.channels_last in the input spec is being silently ignored.

Is my assumption correct, or am I missing something?