Shape inference on modules returns ?

[denolf]

Hi,

Even for a module with just a single con2d layer, shape inference returns a set of question marks: func @forward(%arg0: !torch.vtensor<[1,1,64,64],f32>) -> !torch.vtensor<[?,?,?,?],f32>

while@annotate_args was used to specify the size of the input tensor.

Is there a pass available that can resolve this? Currently the pipelined torchscript-module-to-torch-backend-pipeline was used.

The python code for the small test is attached.

Additionally, what is the decomposition of the torchscript-module-to-torch-backend-pipeline? Using symbol-dce,torch-prepare-for-globalize-object-graph,torch-globalize-object-graph,symbol-dce,inline,torch-adjust-calling-conventions,torch-inline-global-slots does not produce the same output.

Thanks

Kristof cnv1TorchMLIR.py.txt

[silvasean]

I'm in the process of revamping our shape inference to be more precise. Historically, we only needed ranks, because linalg would do the per-dimension size inference. But now that we have more backends, a more generalized approach is needed.

You can see the definition of the pipeline here: https://github.com/llvm/torch-mlir/blob/1019ddf5a0edc388666bf88cd16d8cbbf9c66d9a/lib/Dialect/Torch/Transforms/Passes.cpp#L35

[denolf]

Hi Sean,

That sounds promising and interesting. Do you have any idea when that will be ready, I am eager to try it out.

Thanks

Kristof

[silvasean]

I was hoping by the end of this week, but I guess that has come. I still have a lot to do on it. Give me another week.

[sjarus]

@silvasean , I have what looks like a similar shape inference signature. For testing purposes, I added an ElementwiseNegModule op exactly on the lines of the Log one in e2e_testing/torchscript/elementwise.py , calling torch.neg instead of torch.log.

It fails with this signature:

    Exception:
    Lower Torch Backend IR -> TOSA Backend IR failed with the following diagnostics:
    error: unsupported by backend lowering: tensor with unknown rank or dtype
    note: see current operation: %0 = "torch.aten.neg"(%arg0) : (!torch.vtensor<[?,?],f32>) -> !torch.vtensor
    note: this is likely due to a missing case in RefineTypes

Looking at the generated MLIR:

module attributes {torch.debug_module_name = "ElementwiseNegModule"} { func @forward(%arg0: !torch.vtensor<[?,?],f32> loc(unknown)) -> !torch.vtensor { %0 = torch.aten.neg %arg0 : !torch.vtensor<[?,?],f32> -> !torch.vtensor loc(#loc1) return %0 : !torch.vtensor loc(#loc0) } loc(#loc0) } loc(#loc0)

This looks strange because the return type doesn't report torch.vtensor<[?,?], f32> as the Log test does. Manually adding the return types enables legalization to work.

Edit: this works if I just add AtenNegOp to RefineTypes as it suggests. I'll roll it in unless your work does it first.

[silvasean]

Yeah, I guess we didn't have AtenNegOp support on the linalg backend, so it required a case in RefineTypes (this is expected).

[denolf]

Could you give an update on when this is expected to be available for early testing? Thanks.

[silvasean]

The new system supports this. See https://github.com/llvm/torch-mlir/blob/main/docs/shape_lib.md