[ONNX] Systematically expand ONNX functions before conversion to Torch to avoid needing bespoke conversions

[andfau-amd]

Currently there are a large number of operators in ONNX that aren't supported by TorchOnnxToTorch. As a beginner, I was tasked with adding a conversion for one of those, namely MeanVarianceNormalization, in nod-ai/SHARK-ModelDev#697. But I noticed that this operator is actually a function, and operators that are functions are defined by a series of other ONNX operators. In this case, all those other operators are ones we already have conversions for. So I wondered why we don't expand these operator functions before converting, to avoid needing bespoke conversions for all of them.

I discussed a bit with some colleagues (@ScottTodd, @rsuderman) and it doesn't seem like there's a good reason not to do this, so I'm going to make an attempt at it. Here's some things we've learned that might be relevant:

• The ONNX API has a function called FunctionExpandHelper (in Python: function_expand_helper) that seems to be intended for exactly this kind of expansion. Here are some places it has been used:
  • ONNX uses a similar Python function for generating expanded versions of tests. For the example of MeanVarianceNormalization you can follow the flows of:
    • original Python -> (no expansion) -> generated ONNX -> MLIR
    • original Python -> (expansion) -> generated ONNX -> MLIR
  • ONNX also has a basic test of it
  • OpenVINO uses it in its ONNX frontend
• ONNX also has an inlining API but it's not clear to me if it's suitable for this use.
• Trivia: MeanVarianceNormalization seems to be the poster child for an ONNX operator that is a function! The guide on adding new ops gives it as an example of something that should be a function rather than a primitive operator; it predates functions actually being operators and then got promoted; and it is even used for testing function expansion. :)

Regarding actually implementing an expansion in our importer or converter, some things to keep track of:

• Once we have expansion working, there should be an assertion somewhere to catch cases where operators we expect to be expanded are somehow still present in the IR.
• In case we want to have a bespoke conversion for some particular operator, there should be some allowlist/denylist setup to prevent some operators being expanded, so they can be converted later.
• There should be some instrumentation on the expansion, so it's easy to see which operators are getting expanded, enabling decision-making about possible future special-casing.

If we do implement an expansion, then:

• Should we get rid of any existing bespoke lowerings that are now redundant? I know there's at least one. Maybe something for follow-up work.

[andfau-amd]

The ONNX API has a function called FunctionExpandHelper (in Python: function_expand_helper)

ONNX uses [the Python version of] it for generating expanded versions of tests

Oops, that's not the same function! I'd missed this:

https://github.com/onnx/onnx/blob/b90e252da11dea9bdc191d6b9b8d01511ef3e3bd/onnx/backend/test/case/node/__init__.py#L92-L95

So this Python variant of this function only exists internally to this test code. This might not be good for the idea of doing the expansion in our Python converter (python/torch_mlir/tools/import_onnx/__main__.py), but we'll see.

[stellaraccident]

I haven't looked at this super closely, but in my mind at the beginning, I thought the importer could just be taught to import such functions as private func.func and then emit a func.call to them when used. I think that would match semantics and leave most of the work in MLIR.

[andfau-amd]

Oh, that would actually suit me pretty well. Because the importer is written in Python, and there's no "proper" way to do single-step in-place expansion exposed by the public ONNX Python API right now, I'm currently attempting to do it in two steps: turn these ops into local functions first, then use the ONNX inlining API. But I could skip the inlining at the ONNX level and import them as MLIR local functions. That would give more readable MLIR actually (edit: and come to think of it, also avoid redundant work in the importer), it seems like the better idea.

[stellaraccident]

Yeah, this is how some other frontends do it: get it in mlir and let the infra there take over.

Also keeps the importer mechanical.

[andfau-amd]

It was more complicated than I expected, and I might have missed some things, but I managed to get something working!

I can get it to spit out an MLIR function for the relevant ONNX operator and call it:

python -m torch_mlir.tools.import_onnx mlir_venv/lib/python3.10/site-packages/onnx/backend/test/data/node/test_mvn/model.onnx

module {
  func.func @test_mvn(%arg0: !torch.vtensor<[3,3,3,1],f32>) -> !torch.vtensor<[3,3,3,1],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
    %none = torch.constant.none
    %0 = call @"('MeanVarianceNormalization', '', 18, [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 1\0A    }\0A  }\0A}\0A], [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 1\0A    }\0A  }\0A}\0A], [])"(%arg0) : (!torch.vtensor<[3,3,3,1],f32>) -> !torch.vtensor<[3,3,3,1],f32>
    return %0 : !torch.vtensor<[3,3,3,1],f32>
  }
  func.func @"('MeanVarianceNormalization', '', 18, [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 1\0A    }\0A  }\0A}\0A], [tensor_type {\0A  elem_type: 1\0A  shape {\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 3\0A    }\0A    dim {\0A      dim_value: 1\0A    }\0A  }\0A}\0A], [])"(%arg0: !torch.vtensor<[3,3,3,1],f32>) -> !torch.vtensor<[3,3,3,1],f32> attributes {torch.onnx_meta.ir_version = 0 : si64, torch.onnx_meta.opset_version = 18 : si64, torch.onnx_meta.producer_name = "", torch.onnx_meta.producer_version = ""} {
    %none = torch.constant.none
    %0 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<2.000000e+00> : tensor<f32>} : () -> !torch.vtensor<[],f32>
    %1 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<9.99999971E-10> : tensor<f32>} : () -> !torch.vtensor<[],f32>
    %2 = torch.operator "onnx.Constant"() {torch.onnx.value_ints = [0 : si64, 2 : si64, 3 : si64]} : () -> !torch.vtensor<[3],si64>
    %3 = torch.operator "onnx.ReduceMean"(%arg0, %2) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[3],si64>) -> !torch.vtensor<[1,3,1,1],f32>
    %4 = torch.operator "onnx.Pow"(%3, %0) : (!torch.vtensor<[1,3,1,1],f32>, !torch.vtensor<[],f32>) -> !torch.vtensor<[1,3,1,1],f32>
    %5 = torch.operator "onnx.Pow"(%arg0, %0) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[],f32>) -> !torch.vtensor<[3,3,3,1],f32>
    %6 = torch.operator "onnx.ReduceMean"(%5, %2) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[3],si64>) -> !torch.vtensor<[1,3,1,1],f32>
    %7 = torch.operator "onnx.Sub"(%6, %4) : (!torch.vtensor<[1,3,1,1],f32>, !torch.vtensor<[1,3,1,1],f32>) -> !torch.vtensor<[1,3,1,1],f32>
    %8 = torch.operator "onnx.Sqrt"(%7) : (!torch.vtensor<[1,3,1,1],f32>) -> !torch.vtensor<[1,3,1,1],f32>
    %9 = torch.operator "onnx.Sub"(%arg0, %3) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[1,3,1,1],f32>) -> !torch.vtensor<[3,3,3,1],f32>
    %10 = torch.operator "onnx.Add"(%8, %1) : (!torch.vtensor<[1,3,1,1],f32>, !torch.vtensor<[],f32>) -> !torch.vtensor<[1,3,1,1],f32>
    %11 = torch.operator "onnx.Div"(%9, %10) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[1,3,1,1],f32>) -> !torch.vtensor<[3,3,3,1],f32>
    return %11 : !torch.vtensor<[3,3,3,1],f32>
  }
}

I apologise for the very ugly function name. Hopefully a slightly more pleasant name mangling scheme is possible. :)

Then if I change the importer so that the main function is marked as public, and the operator's function is marked as private, MLIR can inline it and throw it away:

python -m torch_mlir.tools.import_onnx mlir_venv/lib/python3.10/site-packages/onnx/backend/test/data/node/test_mvn/model.onnx | build/bin/torch-mlir-opt --split-input-file --inline

module {
  func.func public @test_mvn(%arg0: !torch.vtensor<[3,3,3,1],f32>) -> !torch.vtensor<[3,3,3,1],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
    %0 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<2.000000e+00> : tensor<f32>} : () -> !torch.vtensor<[],f32>
    %1 = torch.operator "onnx.Constant"() {torch.onnx.value = dense<9.99999971E-10> : tensor<f32>} : () -> !torch.vtensor<[],f32>
    %2 = torch.operator "onnx.Constant"() {torch.onnx.value_ints = [0 : si64, 2 : si64, 3 : si64]} : () -> !torch.vtensor<[3],si64>
    %3 = torch.operator "onnx.ReduceMean"(%arg0, %2) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[3],si64>) -> !torch.vtensor<[1,3,1,1],f32>
    %4 = torch.operator "onnx.Pow"(%3, %0) : (!torch.vtensor<[1,3,1,1],f32>, !torch.vtensor<[],f32>) -> !torch.vtensor<[1,3,1,1],f32>
    %5 = torch.operator "onnx.Pow"(%arg0, %0) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[],f32>) -> !torch.vtensor<[3,3,3,1],f32>
    %6 = torch.operator "onnx.ReduceMean"(%5, %2) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[3],si64>) -> !torch.vtensor<[1,3,1,1],f32>
    %7 = torch.operator "onnx.Sub"(%6, %4) : (!torch.vtensor<[1,3,1,1],f32>, !torch.vtensor<[1,3,1,1],f32>) -> !torch.vtensor<[1,3,1,1],f32>
    %8 = torch.operator "onnx.Sqrt"(%7) : (!torch.vtensor<[1,3,1,1],f32>) -> !torch.vtensor<[1,3,1,1],f32>
    %9 = torch.operator "onnx.Sub"(%arg0, %3) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[1,3,1,1],f32>) -> !torch.vtensor<[3,3,3,1],f32>
    %10 = torch.operator "onnx.Add"(%8, %1) : (!torch.vtensor<[1,3,1,1],f32>, !torch.vtensor<[],f32>) -> !torch.vtensor<[1,3,1,1],f32>
    %11 = torch.operator "onnx.Div"(%9, %10) : (!torch.vtensor<[3,3,3,1],f32>, !torch.vtensor<[1,3,1,1],f32>) -> !torch.vtensor<[3,3,3,1],f32>
    return %11 : !torch.vtensor<[3,3,3,1],f32>
  }
}

And this looks fairly similar at first glance to the known-good expansion, though I'm not yet sure if it's perfectly equivalent.

You can see what I have so far at main...andfau-amd:torch-mlir:onnx-to-torch-function-expansion. The interesting stuff is in get_operator_function, the rest is mostly plumbing.

(Perhaps I should make a draft PR? I think it would be useful to get some early feedback.)

[andfau-amd]

Resolved my FIXMEs, now the MLIR perfectly matches for MeanVarianceNormalization once it's inlined and cleaned up:

$ python -m torch_mlir.tools.import_onnx mlir_venv/lib/python3.10/site-packages/onnx/backend/test/data/node/test_mvn/model.onnx | build/bin/torch-mlir-opt --split-input-file --inline --convert-torch-onnx-to-torch --cse --canonicalize > mvn-ours.txt
$ python -m torch_mlir.tools.import_onnx mlir_venv/lib/python3.10/site-packages/onnx/backend/test/data/node/test_mvn_expanded/model.onnx | build/bin/torch-mlir-opt --split-input-file --inline --convert-torch-onnx-to-torch --cse --canonicalize > mvn-theirs.txt
$ sed -e 's/_expanded//g' mvn-theirs.txt > mvn-theirs2.txt
$ diff mvn-ours.txt mvn-theirs2.txt
$

I'll try to add some tests and then I'll open a pull request.

[renxida]

This is awesome

[andfau-amd]

Thank you!

[andfau-amd]

I've now gotten to the point where cmake --build build --target check-torch-mlir no longer fails with my patch. Lots of edge cases had to be fixed that were causing the importer to crash in pre-existing regression tests. I think I know too much about ONNX now. :) I have actually had to denylist two operators to prevent them being expanded, for fairly good reasons. :(

With that out of the way I can actually start adding new tests.

[andfau-amd]

Created a PR to track things beyond this point and allow commenting on the code etc. It's not quite ready for merging yet though. #3409

[andfau-amd]

Marked the PR as ready for review. It still needs some more testing but that can be done in parallel I hope.

[andfau-amd]

Follow-up issue for expanding the allowlist: #3464