[torch] Add OnnxToTorch lowering for Onnx.ImageDecoder op

[vinayakdsci]

Implements a simplified OnnxToTorch lowering for the Onnx.ImageDecoder op.

The implementation assumes that the image in the respective format has already been loaded and converted to an appropriate tensor representation for simplicity, and therefore has different op semantics than the original Onnx definition.

[vinayakdsci]

cc @vivekkhandelwal1.

[ScottTodd]

The implementation assumes that the image in the respective format has already been loaded and converted to an appropriate tensor representation for simplicity, and therefore has different op semantics than the original Onnx definition.

This op takes an encoded stream of bytes (e.g. !torch.vtensor<[1058],ui8>) and decodes it. This is changing the top to take different inputs (an already decoded image, e.g. !torch.vtensor<[32,32,3],ui8>) and perform a different computation.

Here's an imported test case from the ONNX test suite using similar inputs: https://github.com/nod-ai/SHARK-TestSuite/blob/main/iree_tests/onnx/node/generated/test_image_decoder_decode_jpeg_bgr/model.mlir

module {
  func.func @test_image_decoder_decode_jpeg_bgr(%arg0: !torch.vtensor<[1058],ui8>) -> !torch.vtensor<[32,32,3],ui8> attributes {torch.onnx_meta.ir_version = 9 : si64, torch.onnx_meta.opset_version = 20 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
    %none = torch.constant.none
    %0 = torch.operator "onnx.ImageDecoder"(%arg0) {torch.onnx.pixel_format = "BGR"} : (!torch.vtensor<[1058],ui8>) -> !torch.vtensor<[32,32,3],ui8> 
    return %0 : !torch.vtensor<[32,32,3],ui8>
  }
}

Are there any other cases in torch-mlir where an op definition is changed like this? For this to work at all, input ONNX models and/or the ONNX importer would need to be changed to use this different op. I'm deeply skeptical about checking in code like this that uses the same name as the original op but with an entirely different implementation - that's a recipe for confusion and maintenance costs later on.

[vinayakdsci]

Hi @ScottTodd , I can totally understand your concern, but I am extremely limited by number of ways to overcome this issue of loading the image tensor, and I am very open to any tips you might have for this too.

However, all the steps that I follow after taking the input are logically correct and the code in the PR is closely modelled after the onnx reference implementation.

[ScottTodd]

and the code in the PR is closely modelled after the onnx reference implementation.

Which reference implementation are you looking at? The one I see is https://github.com/onnx/onnx/blob/main/onnx/reference/ops/op_image_decoder.py and that is calling

img = PIL.Image.open(io.BytesIO(encoded.tobytes()))

that's not something we can hand-wave away - it's a large chunk of code bundled into a complicated library, incompatible with this style of compiler / code generator.

Changing the op definition but using the same name does not count as "supporting" an op. An incorrect implementation is worse than no implementation. We could lower via a custom op somehow to backends that want to use their own implementation, but adding this style of lowering would prevent that.

[vinayakdsci]

that's not something we can hand-wave away - it's a large chunk of code bundled into a complicated library, incompatible with this style of compiler / code generator.

Exactly! But claiming support for the op appears to be a priority, and hence the only way that at the moment seems to get anywhere close to that, I implemented in this PR. I have no issues if we go ahead and decide to close this one as not feasible, as I agree with your opinions. But as I said, the use of PIL(and hence the large amount of bundled code) is an extremely limiting factor in terms of replication through compiler codegen.

So the decision is yours, whether the PR is reasonable, or not.