support qnn mean (dim=None)

[cccclai]

Summary: Address mean op lower failure. When dim is not specified, it will take mean across all axes. For QNN, we need to get axes based on input shape

Differential Revision: D83520776

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[facebook-github-bot]

@cccclai has exported this pull request. If you are a Meta employee, you can view the originating diff in D83520776.

[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.

[haowhsu-quic]

Do you need more configurations here? like Mean(dim=0), Mean(dim=0, keepdim=True).

[cccclai]

updated

[facebook-github-bot]

@cccclai has exported this pull request. If you are a Meta employee, you can view the originating diff in D83520776.

[facebook-github-bot]

@cccclai has exported this pull request. If you are a Meta employee, you can view the originating diff in D83520776.

[facebook-github-bot]

@cccclai has exported this pull request. If you are a Meta employee, you can view the originating diff in D83520776.

[cccclai]

Does it look good? I think it can fix 5 failing op tests

[facebook-github-bot]

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

[chenweng-quic]

Should it be torch.tensor(5.0) if you want to test scalar?

[cccclai]

Thanks, good catch

[cccclai]

Actually this is indeed a good test. ReduceMean doesn't support 0d tensor, I think we need to have a pass to convert 0d to 1d tensor if the input of mean is 0d. I have some sketch code here. What do you think? For now, I comment out the test case and think we can follow up on this

import torch
from executorch.exir.pass_base import ExportPass, PassResult

class Rank0ToRank1(ExportPass):
    """
    For selected ops and selected input positions, if the input is rank-0 (scalar),
    insert a reshape to [1] before the op.
    """

    def __init__(self, op_input_map=None) -> None:
        super().__init__()
        # key is the op, value is the input indices to be reshaped
        self.op_input_map = {
          torch.ops.aten.mean.dim: [0],
        }

    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
        graph = graph_module.graph
        changed = False

        for node in list(graph.nodes):
            if node.op == "call_function" and node.target in self.op_input_map:
                input_indices = self.op_input_map[node.target]

                new_args = list(node.args)
                for idx in input_indices:
                    if idx < len(new_args):
                        inp = new_args[idx]
                        if hasattr(inp, "meta"):
                            val = inp.meta.get("val", None)
                            if val is not None and hasattr(val, "shape") and val.shape == ():
                                # Insert reshape right before the op
                                with graph.inserting_before(node):
                                    reshape_node = graph.call_function(
                                        torch.ops.aten.reshape.default,
                                        args=(inp, (1,))
                                    )
                                    reshape_node.meta["val"] = val.reshape(1,)

                                # Replace arg idx with reshape_node
                                new_args[idx] = reshape_node
                                changed = True

                # update node args if modified
                node.args = tuple(new_args)

        if changed:
            graph_module.recompile()
        return PassResult(graph_module, changed)

[cccclai]

@pytorchbot cherry-pick --onto release/1.0 -c regression

[pytorchbot]

Cherry picking #14675

The cherry pick PR is at #14755 and it is recommended to link a regression cherry pick PR with an issue. The following tracker issues are updated:

• [v1.0.0] Release Tracker #14288 (comment)

Details for Dev Infra team

Raised by workflow job