[TOSA] Fold Up-Casts into MatMul when supported by tosa.matmul

[cmcgirr-amd]

In the newer version of Torch-MLIR using Torch 2.3 and Torch Dynamo + FxImporter we see that a GEMM with bias that operate on bf16 and f16 types decompose differently than the earlier versions using TorchScript.

  %0 = torch.aten.to.dtype %arg0, %int6, %false, %false, %none : !torch.vtensor<[4,8],bf16>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[4,8],f32>
  %1 = torch.aten.to.dtype %arg1, %int6, %false, %false, %none : !torch.vtensor<[8,16],bf16>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[8,16],f32>
  %2 = torch.aten.mm %0, %1 : !torch.vtensor<[4,8],f32>, !torch.vtensor<[8,16],f32> -> !torch.vtensor<[4,16],f32>

This decomposition contains a torch.aten.mm that operates on an accumulator type where the inputs are casted from the smaller representation to the accumulator type. E.g. bf16 -> f32. And since TOSA specifications give that the tosa.matmul can support these configurations, we can simply fold the cast into the operation signature rather than keeping the casting operations.

The resulting IR should look like this for the bf16 case.

%8 = tosa.matmul %6, %7 : (tensor<1x4x8xbf16>, tensor<1x8x16xbf16>) -> tensor<1x4x16xf32>

One caveat is the i16 -> i48 tosa.matmul case cannot be supported at the moment as there is no torch.int48 equivalent dtype to represent this special accumulator used in TOSA.

[roberteg16]

As you noted on your PR message, I think we should guard against i16 -> i48. Otherwise lgtm

[TinaAMD]

Nice idea to get rid of the casts this way! Implementationwise, I would have expected this to be a canonicalization on TOSA, i.e. you just lower naively from torch-mlir to tosa, and then canonicalize the tosa casts + matmul to a different matmul (then you could also easily cover the i48 case). WDYT?

[cmcgirr-amd]

Nice idea to get rid of the casts this way! Implementationwise, I would have expected this to be a canonicalization on TOSA, i.e. you just lower naively from torch-mlir to tosa, and then canonicalize the tosa casts + matmul to a different matmul (then you could also easily cover the i48 case). WDYT?

Kind of tricky because for some the accumulator types are not supported as input types so it would be invalid to have:

tosa.matmul %6, %7 : (tensor<1x4x8xi32>, tensor<1x8x16xi32>) -> tensor<1x4x16xi32>

Could be done for the float types, but integers may be difficult.

[cmcgirr-amd]

As you noted on your PR message, I think we should guard against i16 -> i48. Otherwise lgtm

In what sense do you mean guarded? Like don't convert in this case:

torch.aten.mm %0, %1 : !torch.vtensor<[4,8],si16>, !torch.vtensor<[8,16],si16> -> !torch.vtensor<[4,16],si48>

Or when there are casting ops? Because we cannot have a legal torch.aten.to_dtype op with si48 AFAIK

[TinaAMD]

Nice idea to get rid of the casts this way! Implementationwise, I would have expected this to be a canonicalization on TOSA, i.e. you just lower naively from torch-mlir to tosa, and then canonicalize the tosa casts + matmul to a different matmul (then you could also easily cover the i48 case). WDYT?

Kind of tricky because for some the accumulator types are not supported as input types so it would be invalid to have:

tosa.matmul %6, %7 : (tensor<1x4x8xi32>, tensor<1x8x16xi32>) -> tensor<1x4x16xi32>

Oh, I see, so you cannot actually lower the integer cases when the casts are not present, I didn't notice that this was disallowed. Makes sense to keep it here then, thanks for the explanation!

Could be done for the float types, but integers may be difficult.

[roberteg16]

As you noted on your PR message, I think we should guard against i16 -> i48. Otherwise lgtm

In what sense do you mean guarded? Like don't convert in this case:

torch.aten.mm %0, %1 : !torch.vtensor<[4,8],si16>, !torch.vtensor<[8,16],si16> -> !torch.vtensor<[4,16],si48>

Or when there are casting ops? Because we cannot have a legal torch.aten.to_dtype op with si48 AFAIK

I was referring to avoid converting a case like you described.

Because we cannot have a legal torch.aten.to_dtype op with si48 AFAIK

Had a look into: https://github.com/llvm/torch-mlir/blob/main/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td#L11547

I am no sure it is not allowed as you say

[cmcgirr-amd]

As you noted on your PR message, I think we should guard against i16 -> i48. Otherwise lgtm

In what sense do you mean guarded? Like don't convert in this case:

torch.aten.mm %0, %1 : !torch.vtensor<[4,8],si16>, !torch.vtensor<[8,16],si16> -> !torch.vtensor<[4,16],si48>

Or when there are casting ops? Because we cannot have a legal torch.aten.to_dtype op with si48 AFAIK

I was referring to avoid converting a case like you described.

Because we cannot have a legal torch.aten.to_dtype op with si48 AFAIK

Had a look into: https://github.com/llvm/torch-mlir/blob/main/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td#L11547

I am no sure it is not allowed as you say

That is true the definition allows for anything, but I would find it hard to find an integer value to describe the dtype if the importer does not have it defined: https://github.com/Xilinx/torch-mlir/blob/feature/backport_ea1_ops/python/torch_mlir/extras/fx_importer.py#L176-L193

[roberteg16]

As you noted on your PR message, I think we should guard against i16 -> i48. Otherwise lgtm

In what sense do you mean guarded? Like don't convert in this case:

torch.aten.mm %0, %1 : !torch.vtensor<[4,8],si16>, !torch.vtensor<[8,16],si16> -> !torch.vtensor<[4,16],si48>

Or when there are casting ops? Because we cannot have a legal torch.aten.to_dtype op with si48 AFAIK

I was referring to avoid converting a case like you described.

Because we cannot have a legal torch.aten.to_dtype op with si48 AFAIK

Had a look into: https://github.com/llvm/torch-mlir/blob/main/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td#L11547
I am no sure it is not allowed as you say

That is true the definition allows for anything, but I would find it hard to find an integer value to describe the dtype if the importer does not have it defined: https://github.com/Xilinx/torch-mlir/blob/feature/backport_ea1_ops/python/torch_mlir/extras/fx_importer.py#L176-L193

Oh I see, it seems that the imported does not allow it. Could we just harden it to make it safe for the future?

[TinaAMD]

LGTM