[Mlir-commits] [mlir] [mlir][x86] Support multi-level unroll of contractions with non-native K extent (PR #234593)

[Julian Oppermann]

================
@@ -1504,14 +1526,9 @@ module attributes {transform.with_named_sequence} {
 #map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
 #map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
 
-// Loop order K-batch, accumulating over both loops: the inner batch loop
-// carries the accumulator, its IV appears exactly once in each operand read and
-// its step is 1. It is a valid accumulation loop, indistinguishable from a
-// K-loop over a blocked layout, and is treated as such.
-//
-// Again, spilling the accumulator value to memory in the outer loop is correct,
-// but not ideal.
-func.func @amx_bf16_flat_k_batch_order(%A: memref<?x?x?xbf16>, %B: memref<?x?x?xbf16>, %C: memref<?x?xf32>,
+// Loop order K-batch, accumulating over both loops: the inner loops doesn't
+// satisfy the condition for block indexing (last two indices are non-zero).
+func.func @negative_k_batch_order(%A: memref<?x?x?xbf16>, %B: memref<?x?x?xbf16>, %C: memref<?x?xf32>,
----------------
jopperm wrote:

Yes, it is a regression. My rationale to "sacrifice" it for a simpler K-tiling support was that the VCMLU pattern doesn't understand nested accumulation loops yet (in contrast to the nanokernel), and this is one of the tests that I added to check that it at least doesn't break even when it is applied in a loop nest.

I see the following options to bring the test case back in its original form:
1. Introduce additional subviews with all the loop-independent indices before the accumulation loop.
2. Only enforce zero indices in the K position of the operand loads if we don't have to do K tiling. 
3. In the nanokernel, support that the K index is not the loop's IV, but derived from it.

Let me try (1). (2) is merely a workaround, and (3) would be nice to have in general, but IIUC would be a non-trivial change to the nanokernel pattern.

https://github.com/llvm/llvm-project/pull/234593