[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>) {
+    %a = vector.transfer_read %A[%m, %k], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<128x32xbf16>
+    %b = vector.transfer_read %B[%k, %n], %c0_0 {in_bounds = [true, true]} : memref<?x?xbf16>, vector<32x64xbf16>
+    %d = vector.contract {
+      indexing_maps = [#map, #map1, #map2],
+      iterator_types = ["parallel", "parallel", "reduction"],
+      kind = #vector.kind<add>}
+      %a, %b, %acc : vector<128x32xbf16>, vector<32x64xbf16> into vector<128x64xf32>
+    scf.yield %d : vector<128x64xf32>
+  }
+  vector.transfer_write %res, %C[%m, %n] {in_bounds = [true, true]} : vector<128x64xf32>, memref<?x?xf32>
----------------
jopperm wrote:

Yes, the fallback is always storing the initial value to a stack buffer, and replacing the result value of the accumulation loop with a transfer_read of the stack buffer. In addition, in this PR, the pattern understands zero initialisation and pulls "simple" vector transfer ops into the loop nest.

There are a couple of low-hanging fruits I want to explore in follow-up PRs, such as pulling `arith.extf` and `arith.truncf` into the loop nest (and using a smaller buffer), and maybe even fusing a subset of common arith and math ops.

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