Inconsistent lowering of 3D and 4D matmuls

[rengolin]

When lowering KernelBench level1/10_3D_tensor_matrix_multiplication.py and level1/11_4D_tensor_matrix_multiplication.py I get very different lowering styles.

Both kernels implement a similar einsum:

• 3D: Bmk,kn -> Bmn
• 4D: Bbmk,kn -> Bbmn

Reproducer

Clone Lighthouse

• Init uv as in README
• Run: $ uv run examples/KernelBench/test-kernel-bench.py --kernel level1/10_ --print-original-module
• Run: $ uv run examples/KernelBench/test-kernel-bench.py --kernel level1/11_ --print-original-module
• Compare the IRs

3D matmul

The 3D lowering materializes a broadcast of the B tensor to Bkn and then applies a batch_matmul:

  // Broadcast B into an empty
  %bcast = linalg.generic {...} ins(%arg1 : tensor<2048x768xf32>) outs(%empty : tensor<16x2048x768xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<16x2048x768xf32>

  // Batch matmul
  %mm = linalg.batch_matmul ins(%arg0, %bcast : tensor<16x1024x2048xf32>, tensor<16x2048x768xf32>) outs(%5 : tensor<16x1024x768xf32>) -> tensor<16x1024x768xf32>

While this is correct, I now have to elide the broadcast and convert the whole pattern to a contraction with broadcast maps. This could be trivially lowered to linalg.contract with #map2 being a "broadcast-B" map:

  linalg.contract { #map1, #map2, #map3 } ins(%arg0, %arg1 : tensor<16x1024x2048xf32>, tensor<2048x768xf32>) outs(%5 : tensor<16x1024x768xf32>) -> tensor<16x1024x768xf32>

4D matmul

The 4D lowering is even weirder. It collapses the 4D A shape to a 2D (as [0,1,2][3]) and then expands the shape of both B and the collapsed A to add a unit dimension, and then call a batch_matmul on that.

// Shape fiddling
%col_a = tensor.collapse_shape %0 [[0, 1, 2], [3]] : tensor<8x256x512x256xf32> into tensor<1048576x256xf32>
%ex_a = tensor.expand_shape %collapsed [[0, 1], [2]] output_shape [1, 1048576, 256] : tensor<1048576x256xf32> into tensor<1x1048576x256xf32>
%ex_b = tensor.expand_shape %1 [[0, 1], [2]] output_shape [1, 256, 768] : tensor<256x768xf32> into tensor<1x256x768xf32>

// Batch matmul
linalg.batch_matmul ins(%ex_a, %ex_b : tensor<1x1048576x256xf32>, tensor<1x256x768xf32>) outs(%empty : tensor<1x1048576x768xf32>) -> tensor<1x1048576x768xf32>

Again, this would be a trivial linalg.contract pattern, with #map2 being a "broadcast-B" map:

  linalg.contract { #map1, #map2, map3 } ins(%ex_a, %ex_b : tensor<8x256x512x256xf32>, tensor<256x768xf32>) outs(%empty : tensor<8x256x512x768xf32>) -> tensor<8x256x512x768xf32>

Linalg Contract

I understand this pattern has been created well before linalg.contract existed, so it may not be trivial or immediate to convert this lowering pattern. But at least have a common solution to ND (N>2) lowering to either collapse/expand shapes OR materialize broadcasts, not both.

The more patterns we have for this, the harder is the compiler's job to deal with different situations.

[hsqStephenZhang]

for 4d example, here is the IR before decomposition:

module {
  func.func @main(%arg0: !torch.vtensor<[8,256,512,256],f32>, %arg1: !torch.vtensor<[256,768],f32>) -> !torch.vtensor<[8,256,512,768],f32> {
    %str = torch.constant.str "Bbmk,kn->Bbmn"
    %0 = torch.prim.ListConstruct %arg0, %arg1 : (!torch.vtensor<[8,256,512,256],f32>, !torch.vtensor<[256,768],f32>) -> !torch.list<vtensor>
    %none = torch.constant.none
    %1 = torch.aten.einsum %str, %0, %none : !torch.str, !torch.list<vtensor>, !torch.none -> !torch.vtensor<[8,256,512,768],f32>
    return %1 : !torch.vtensor<[8,256,512,768],f32>
  }
}

after decomposition(run with -torch-decompose-complex-ops), it became a sequence of view + permuate + ListConstruct + bmm

module {
  func.func @main(%arg0: !torch.vtensor<[8,256,512,256],f32>, %arg1: !torch.vtensor<[256,768],f32>) -> !torch.vtensor<[8,256,512,768],f32> {
    %int1048576 = torch.constant.int 1048576
    %int8 = torch.constant.int 8
    %int256 = torch.constant.int 256
    %int512 = torch.constant.int 512
    %int768 = torch.constant.int 768
    %int0 = torch.constant.int 0
    %int1 = torch.constant.int 1
    %int2 = torch.constant.int 2
    %int3 = torch.constant.int 3
    %0 = torch.prim.ListConstruct %int0, %int1, %int2, %int3 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
    %1 = torch.aten.permute %arg0, %0 : !torch.vtensor<[8,256,512,256],f32>, !torch.list<int> -> !torch.vtensor<[8,256,512,256],f32>
    %2 = torch.prim.ListConstruct %int0, %int1 : (!torch.int, !torch.int) -> !torch.list<int>
    %3 = torch.aten.permute %arg1, %2 : !torch.vtensor<[256,768],f32>, !torch.list<int> -> !torch.vtensor<[256,768],f32>
    %4 = torch.prim.ListConstruct %int1, %int1048576, %int256 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
    %5 = torch.aten.view %1, %4 : !torch.vtensor<[8,256,512,256],f32>, !torch.list<int> -> !torch.vtensor<[1,1048576,256],f32>
    %6 = torch.prim.ListConstruct %int1, %int256, %int768 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int>
    %7 = torch.aten.view %3, %6 : !torch.vtensor<[256,768],f32>, !torch.list<int> -> !torch.vtensor<[1,256,768],f32>
    %8 = torch.aten.bmm %5, %7 : !torch.vtensor<[1,1048576,256],f32>, !torch.vtensor<[1,256,768],f32> -> !torch.vtensor<[1,1048576,768],f32>
    %9 = torch.prim.ListConstruct %int8, %int256, %int512, %int768 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
    %10 = torch.aten.view %8, %9 : !torch.vtensor<[1,1048576,768],f32>, !torch.list<int> -> !torch.vtensor<[8,256,512,768],f32>
    %11 = torch.prim.ListConstruct %int0, %int1, %int2, %int3 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
    %12 = torch.aten.permute %10, %11 : !torch.vtensor<[8,256,512,768],f32>, !torch.list<int> -> !torch.vtensor<[8,256,512,768],f32>
    return %12 : !torch.vtensor<[8,256,512,768],f32>
  }
}

after another run with -torch-backend-to-linalg-on-tensors-backend-pipeline, we can get the mlir that matches the issue's description above.

I think the problem is that the einsum decomposition was forced to match batch_matmul but not generic or contract, so the output looks not so natural.

[rengolin]

I think the problem is that the einsum decomposition was forced to match batch_matmul but not generic or contract, so the output looks not so natural.

That's my guess, too. At the very least it should be a generic, which can accept everything. One can use the Linalg morphisms (generic-to-contract or generic-to-named) afterwards.