[Linalg] linalg.generic and linalg.batch_matmul behaves differently with fold_type_extensions_into_contract

MLIR
1

Consider the following example.

#lhs = affine_map<(b, m, n, k) -> (b, m, k)>
#rhs = affine_map<(b, m, n, k) -> (b, k, n)>
#out = affine_map<(b, m, n, k) -> (b, m, n)>

func.func @mixed_bmm(%A: tensor<1x8x16xf32>, %B: tensor<1x16x8xf16>,
                     %C: tensor<1x8x8xf32>) -> tensor<1x8x8xf32> {
  %0 = linalg.batch_matmul ins(%A, %B : tensor<1x8x16xf32>, tensor<1x16x8xf16>)
                           outs(%C : tensor<1x8x8xf32>) -> tensor<1x8x8xf32>
  return %0 : tensor<1x8x8xf32>
}

func.func @mixed_generic(%A: tensor<1x8x16xf32>, %B: tensor<1x16x8xf16>,
                         %C: tensor<1x8x8xf32>) -> tensor<1x8x8xf32> {
  %0 = linalg.generic {indexing_maps = [#lhs, #rhs, #out],
                       iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
       ins(%A, %B : tensor<1x8x16xf32>, tensor<1x16x8xf16>)
       outs(%C : tensor<1x8x8xf32>) {
  ^bb0(%a: f32, %b: f16, %acc: f32):
    %narrow = arith.truncf %a : f32 to f16
    %ae = arith.extf %narrow : f16 to f32
    %be = arith.extf %b : f16 to f32
    %mul = arith.mulf %ae, %be : f32
    %add = arith.addf %acc, %mul : f32
    linalg.yield %add : f32
  } -> tensor<1x8x8xf32>
  return %0 : tensor<1x8x8xf32>
}

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
    %f = transform.structured.match ops{["func.func"]} in %root
        : (!transform.any_op) -> !transform.any_op
    %vf = transform.structured.vectorize_children_and_apply_patterns %f
        { fold_type_extensions_into_contract } : (!transform.any_op) -> !transform.any_op
    transform.yield
  }
}

the outputs for two cases are different with fold_type_extensions_into_contract enabled.

func.func @mixed_bmm(%arg0: tensor<1x8x16xf32>, %arg1: tensor<1x16x8xf16>, %arg2: tensor<1x8x8xf32>) -> tensor<1x8x8xf32> {
    %0 = ub.poison : f16
    %c0 = arith.constant 0 : index
    %1 = ub.poison : f32
    %2 = vector.transfer_read %arg0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} : tensor<1x8x16xf32>, vector<1x8x16xf32>
    %3 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} : tensor<1x16x8xf16>, vector<1x16x8xf16>
    %4 = vector.transfer_read %arg2[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} : tensor<1x8x8xf32>, vector<1x8x8xf32>
    %5 = arith.extf %3 : vector<1x16x8xf16> to vector<1x16x8xf32>
    %6 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %2, %5, %4 : vector<1x8x16xf32>, vector<1x16x8xf32> into vector<1x8x8xf32>
    %7 = vector.transfer_write %6, %arg2[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<1x8x8xf32>, tensor<1x8x8xf32>
    return %7 : tensor<1x8x8xf32>
  }
  func.func @mixed_generic(%arg0: tensor<1x8x16xf32>, %arg1: tensor<1x16x8xf16>, %arg2: tensor<1x8x8xf32>) -> tensor<1x8x8xf32> {
    %0 = ub.poison : f16
    %c0 = arith.constant 0 : index
    %1 = ub.poison : f32
    %2 = vector.transfer_read %arg0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} : tensor<1x8x16xf32>, vector<1x8x16xf32>
    %3 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} : tensor<1x16x8xf16>, vector<1x16x8xf16>
    %4 = vector.transfer_read %arg2[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} : tensor<1x8x8xf32>, vector<1x8x8xf32>
    %5 = arith.truncf %2 : vector<1x8x16xf32> to vector<1x8x16xf16>
    %6 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "reduction"], kind = #vector.kind<add>} %5, %3, %4 : vector<1x8x16xf16>, vector<1x16x8xf16> into vector<1x8x8xf32>
    %7 = vector.transfer_write %6, %arg2[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<1x8x8xf32>, tensor<1x8x8xf32>
    return %7 : tensor<1x8x8xf32>
  }

In BMM case we see upcast but for generic case we see the low-precision contract as expected. I understand that current implementation of fold_type_extensions_into_contract only fires if both A and B are coming from arith.extf, but shouldn’t these two show consistent behavior?

cc @rengolin

2

From the structured ops definition, the inputs types can be different in the contract, and their generalizations should only extend/truncate if necessary. I expected the linalg.contract to keep the operand element types as f32 and f16 respectively, even with the fold attribute, whereas the extension would be “logical” inside the contract and materialized only upon generalization.

Semantically speaking, both casts are wrong. Hardware architectures can have support for different combination of types for the inputs (element types, swizzles, shapes) and adding the cast to make both input types the same can stop further patterns to match the correct op/instruction without further cast fusion.

There doesn’t seem to be tests for mixed input types, so to me this looks like a bug in the code.

@mshahid @asiemien any ideas?

3

Definitely this is a reason for the rewrite to be rough around these edges.

The two linalg ops are not identical, the linalg.generic has different body compared to linalg.batch_matmul. That makes it harder to define if the outcome should be the same.
Both represent mixed precision batch matmul. However, during generic vectorization ops from the body will materialize. So the difference in IR becomes more explicit.

That’s another part of the equation. These casts tend to be loosely interpreted.
I think round tripping generic => named op => generic might even lose the extra trunc-ext pair. Whether that’s correct is a separate question.

Also, keep in mind that vectorize_children_and_apply_patterns doesn’t vectorize directly to vector.contract in either case. This transform is a bag of rewrite patterns going through multi_reduction, then reconstructing contract when possible, applying cleanups etc.

These two together result in slightly different vectorization results.
But again, not sure if consistency can or should be expected here.

4

agreed. it seems the vectorization makes arbitrary decisions on which types to cast to make A, B the same types. This is problematic because BMM itself support mixed types for A and B and vectorization does not respect them. It seems only way around is to have downstream target-specific contract legalization which adds necessary casts depending on what the target supports. I am not sure what the expected behavior for vectorization should be either?

agreed, but the issue here is orthogonal to that. why BMM vectorization decided to cast to f32, why not f16? problem is what logic decides that and what is the motivation. It is not quite clear to me at least.