pytorch / pytorch/executorch

ConvertToLinearPass is not sound when transposes are elided

Open
#10,499 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

module: xnnpack
Dominant language
Python
Stars
5k
Forks
1.2k
Avg merge
2d 10h
Merged PRs (30d)
581

Description

🐛 Describe the bug

The linear pattern has executorch_exir_dialects_edge__ops_aten_permute_copy_default and executorch_exir_dialects_edge__ops_aten_addmm_default:

input_1 = input
 aten_permute_copy_default: "f32[4, 8]" = executorch_exir_dialects_edge__ops_aten_permute_copy_default(p_weight, [1, 0]);  p_weight = None
aten_addmm_default: "f32[1, 8]" = executorch_exir_dialects_edge__ops_aten_addmm_default(p_bias, input_1, aten_permute_copy_default);  p_bias = input_1 = aten_permute_copy_default = None
return (aten_addmm_default,)

The ConvertToLinearPass tries to reconstruct linear from these ops. It correctly does this when permute is present, but incorrectly constructs linear from executorch_exir_dialects_edge__ops_aten_addmm_default when permute is not present. This can happen if you elide the transpose (e.g., with const propagation or RemoveRedundantTransposes).

Repro (modified from backends/xnnpack/test/passes/test_convert_to_linear):

import unittest

import torch
from executorch.backends.xnnpack._passes.convert_to_linear import ConvertToLinearPass
from executorch.backends.xnnpack.test.tester import RunPasses, Tester
from executorch.exir.passes.constant_prop_pass import constant_prop_pass


class TestConvertToLinear(unittest.TestCase):
    PassStage = RunPasses([ConvertToLinearPass])

    def setUp(self):
        torch._dynamo.reset()

    def test_fp32_convert_to_linear(self):
        in_sizes = [1, 4, 4]
        input_sizes = [4, 37, 17]
        output_sizes = [8, 17, 37]
        bias_vals = [True, True, False]

        for i, _ in enumerate(in_sizes):
            torch._dynamo.reset()
            in_size = int(in_sizes[i])
            input_size = int(input_sizes[i])
            output_size = int(output_sizes[i])
            linear = torch.nn.Linear(input_size, output_size, bias=bias_vals[i])
            inputs = (torch.randn(in_size, input_size),)

            to_edge_stage = Tester(linear, inputs).export().to_edge()
            constant_prop_pass(to_edge_stage.stages["ToEdge"].artifact.exported_program())
            (
                to_edge_stage
                .run_passes(self.PassStage)
                .check_count(
                    {"executorch_exir_dialects_edge__ops_aten_linear_default": 1}
                )
                .run_method_and_compare_outputs()
            )

This gives failure:

RuntimeError: a and b must have same reduction dim, but got [1, 4] X [8, 4].

I have seen this issue in the export llama script when linear ops are not delegated to XNNPACK.

Versions

NA

cc @digantdesai @mcr229 @cbilgin

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with executorch/backends/xnnpack/_passes/convert_to_linear.py and the tests under backends/xnnpack/test/passes/test_convert_to_linear. Run the provided constant-propagation reproduction and inspect how ConvertToLinearPass handles addmm when permute_copy is absent. Done means the pass no longer creates an invalid linear operation and the reproduced outputs compare successfully.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.