[Mlir-commits] [mlir] [mlir][x86] Add pattern to tile and unroll vector.contract (PR #229076)

[Julian Oppermann]

================
@@ -0,0 +1,1239 @@
+// RUN: mlir-opt %s -transform-interpreter -canonicalize -cse -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -transform-interpreter=entry-point=__transform_nano -canonicalize -cse -split-input-file \
+// RUN:   | FileCheck %s --check-prefix=NANO --implicit-check-not=vector.contract
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+func.func @amx_bf16_flat(%A: memref<?x?xbf16>, %B: memref<?x?xbf16>, %C: memref<?x?xf32>,
+                         %m: index, %n: index, %k_start: index, %k_end: index) {
+  %c32 = arith.constant 32 : index
+  %c0_0 = arith.constant 0.0 : bf16
+  %c0_0_f32 = arith.constant 0.0 : f32
+  %c = vector.transfer_read %C[%m, %n], %c0_0_f32 {in_bounds = [true, true]} : memref<?x?xf32>, vector<128x64xf32>
+  %res = scf.for %k = %k_start to %k_end step %c32 iter_args(%acc = %c) -> (vector<128x64xf32>) {
----------------
jopperm wrote:

> And if the k loop is not the innermost one?

It depends. 
- If the innermost loop is still a valid accumulation loop for the pattern, the transformation succeeds (most likely with the accumulator value being spilled in the surrounding loop). The AMX nanokernel handles this correctly, cf. the `@amx_bf16_flat_k_batch_order` I pushed with the last commit.
- Otherwise, the pattern bails out, see e.g. in the new `@negative_k_loop_not_innermost` test: Here, the IV of the loop wrapping the contraction is not used to index in the LHS operand load, so we don't touch it.

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