[ONNX] Disable type-unsafe Cast transpose fusion

[liquetxnx]

Summary

The canonicalize pass creates an invalid onnx.Cast.

The original Cast is valid:

%159 = "onnx.Transpose"(%158)
    <{perm = [0, 1, 3, 2]}>
    : (tensor<1x1x124x40xf16>) -> tensor<1x1x40x124xf16>

%160 = "onnx.Cast"(%159)
    <{saturate = 1 : si64, to = f32}>
    : (tensor<1x1x40x124xf16>) -> tensor<1x1x40x124xf32>

%161 = "onnx.Transpose"(%160)
    <{perm = [0, 2, 1, 3]}>
    : (tensor<1x1x40x124xf32>) -> tensor<1x40x1x124xf32>

After canonicalization, a rewritten Cast has to = f32 but an f16 result:

%cast = "onnx.Cast"(%value)
    <{saturate = 1 : si64, to = f32}>
    : (tensor<...xf16>) -> tensor<...xf16>

This causes:

'onnx.Cast' op element type does not match the 'to' attribute

Reproduction

onnx-mlir-opt input.onnx.mlir \
  --canonicalize \
  -o output.onnx.mlir

The invalid rewritten IR can be inspected with:

onnx-mlir-opt input.onnx.mlir \
  --verify-each=false \
  --canonicalize \
  -o output-unverified.onnx.mlir

Relevant data flow

The original data flow is:

Transpose
  -> Cast(f16 -> f32)
  -> Reshape
  -> Transpose

The Reshape input and result types are identical:

%178 = "onnx.Reshape"(%160, %shape)
    : (tensor<1x1x40x124xf32>, tensor<4xi64>)
      -> tensor<1x1x40x124xf32>

%179 = "onnx.Transpose"(%178)
    <{perm = [0, 2, 1, 3]}>
    : (tensor<1x1x40x124xf32>)
      -> tensor<1x40x1x124xf32>

Canonicalization can remove the identity Reshape, exposing:

Transpose -> Cast -> Transpose

Suspected cause

The likely source is the Transpose/elementwise fusion pattern in:

src/Dialect/ONNX/ONNXOps/Canonicalize.td

The relevant two-argument pattern reconstructs the elementwise operation using (about 680 lines):

(returnType $v)

This works for type-preserving operations such as Relu, but not for
ONNXCastOp.

When $v is tensor<...xf16>, the rewritten Cast receives an f16 result
type even though its to attribute is f32.

Expected behavior

The rewritten Cast should preserve the Cast result element type:

%cast = "onnx.Cast"(%value)
    <{saturate = 1 : si64, to = f32}>
    : (tensor<...xf16>) -> tensor<...xf32>

Suggested fix

• exclude ONNXCastOp from this generic type-preserving fusion pattern

[jenkins-droid]

Can one of the admins verify this patch?

[AlexandreEichenberger]

@jenkins-droid test this please

[AlexandreEichenberger]

@jenkin-droid test this please

[gongsu832]

@jenkins-droid test this please

[AlexandreEichenberger]

@jenkins-droid test this please