[Mlir-commits] [mlir] [mlir][X86] Prepare contract interleave (PR #236750)

[Julian Oppermann]

================
@@ -9,71 +9,217 @@
 // CHECK-DAG: #[[$VNNI_RHS:[A-Za-z0-9_]+]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>
 // CHECK-DAG: #[[$VNNI_ACC:[A-Za-z0-9_]+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
 
-// fold VNNI unpacking of packed bf16 rhs read into contraction, keeping fastmath
-// CHECK-LABEL: func.func @direct_vnni_bf16(
-// CHECK-SAME: %[[A:[^,]+]]: memref<?x32x32xbf16>, %[[B:[^,]+]]: memref<?x16x32x2xbf16>
-// NANO-LABEL: func.func @direct_vnni_bf16(
-// NANO-NOT: vector.contract
+// Fold deinterleave/transpose/interleave/transpose unpacking of non-square
+// operands with strided memrefs, aligned nonzero inner read offsets, and
+// fastmath flags.
+// CHECK-LABEL: func.func @interleaved_vnni_bf16(
+// CHECK-SAME: %[[A:[^:]+]]: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>, %[[B:[^:]+]]: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>
+// NANO-LABEL: func.func @interleaved_vnni_bf16(
+// NANO: vector.contract {{.*}} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
 // NANO: return
-// CHECK: %[[A_PACKED:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x32x16x2xbf16>
-// CHECK: scf.for %[[K:.+]] = %{{.+}} to %{{.+}} step %{{.+}}
-// CHECK: %[[LHS:.+]] = vector.transfer_read %[[A_PACKED]][%[[K]], {{.*}} vector<32x16x2xbf16>
-// CHECK: %[[RHS:.+]] = vector.transfer_read %[[B]][%[[K]], {{.*}} vector<16x32x2xbf16>
-// CHECK: vector.contract {fastmath = #arith.fastmath<reassoc>, indexing_maps = [#[[$VNNI_LHS]], #[[$VNNI_RHS]], #[[$VNNI_ACC]]], iterator_types = ["parallel", "parallel", "reduction", "reduction"], kind = #vector.kind<add>} %[[LHS]], %[[RHS]], %{{.+}} : vector<32x16x2xbf16>, vector<16x32x2xbf16> into vector<32x32xf32>
-func.func @direct_vnni_bf16(%a: memref<?x32x32xbf16>,
-                            %b: memref<?x16x32x2xbf16>,
-                            %init: vector<32x32xf32>, %ub: index,
-                            %m: index) -> vector<32x32xf32> {
+// CHECK-DAG: %[[ONE:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[AP:.+]] = memref.expand_shape %[[A]] {{.*}} into memref<?x?x16x32x2xbf16, strided<[?, 1024, 64, 2, 1]>>
+// CHECK-DAG: %[[BP:.+]] = memref.expand_shape %[[B]] {{.*}} into memref<?x?x16x32x2xbf16, strided<[?, 1024, 64, 2, 1]>>
+// CHECK: scf.for %[[K:[^ ]+]] =
+// CHECK: %[[LHS:.+]] = vector.transfer_read %[[AP]][%{{.+}}, %[[K]], %{{.+}}, %[[ONE]], %{{.+}}]{{.*}} vector<16x16x2xbf16>
+// CHECK-NEXT: %[[RHS:.+]] = vector.transfer_read %[[BP]][%{{.+}}, %[[K]], %{{.+}}, %[[ONE]], %{{.+}}]{{.*}} vector<16x16x2xbf16>
+// CHECK-NEXT: %[[DOT:.+]] = vector.contract {fastmath = #arith.fastmath<reassoc>, indexing_maps = [#[[$VNNI_LHS]], #[[$VNNI_RHS]], #[[$VNNI_ACC]]], iterator_types = ["parallel", "parallel", "reduction", "reduction"], kind = #vector.kind<add>} %[[LHS]], %[[RHS]], %{{.+}} : vector<16x16x2xbf16>, vector<16x16x2xbf16> into vector<16x16xf32>
+// CHECK-NEXT: scf.yield %[[DOT]]
+func.func @interleaved_vnni_bf16(%a: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>,
+                                %b: memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>,
+                                %init: vector<16x16xf32>, %lb: index, %ub: index,
+                                %m: index, %n: index) -> vector<16x16xf32> {
   %zero = arith.constant 0 : index
-  %step = arith.constant 16 : index
+  %offset = arith.constant 2 : index
+  %step = arith.constant 1 : index
   %pad = arith.constant 0.0 : bf16
-  %result = scf.for %k = %zero to %ub step %step iter_args(%accum = %init) -> (vector<32x32xf32>) {
-    %lhs = vector.transfer_read %a[%k, %m, %zero], %pad {in_bounds = [true, true]} : memref<?x32x32xbf16>, vector<32x32xbf16>
-    %rhs = vector.transfer_read %b[%k, %zero, %zero, %zero], %pad {in_bounds = [true, true, true]} : memref<?x16x32x2xbf16>, vector<16x32x2xbf16>
-    %transpose = vector.transpose %rhs, [0, 2, 1] : vector<16x32x2xbf16> to vector<16x2x32xbf16>
-    %flat = vector.shape_cast %transpose : vector<16x2x32xbf16> to vector<32x32xbf16>
-    %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>, fastmath = #arith.fastmath<reassoc>} %lhs, %flat, %accum : vector<32x32xbf16>, vector<32x32xbf16> into vector<32x32xf32>
+  %result = scf.for %k = %lb to %ub step %step iter_args(%accum = %init) -> (vector<16x16xf32>) {
+    %lhs = vector.transfer_read %a[%m, %k, %zero, %offset], %pad {in_bounds = [true, true]} : memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>, vector<16x32xbf16>
+    %rhs = vector.transfer_read %b[%n, %k, %zero, %offset], %pad {in_bounds = [true, true]} : memref<?x?x16x64xbf16, strided<[?, 1024, 64, 1]>>, vector<16x32xbf16>
+    %even, %odd = vector.deinterleave %rhs : vector<16x32xbf16> -> vector<16x16xbf16>
+    %even_t = vector.transpose %even, [1, 0] : vector<16x16xbf16> to vector<16x16xbf16>
+    %odd_t = vector.transpose %odd, [1, 0] : vector<16x16xbf16> to vector<16x16xbf16>
+    %interleaved = vector.interleave %even_t, %odd_t : vector<16x16xbf16> -> vector<16x32xbf16>
+    %flat = vector.transpose %interleaved, [1, 0] : vector<16x32xbf16> to vector<32x16xbf16>
+    %dot = vector.contract {indexing_maps = [#lhs, #rhs, #acc], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>, fastmath = #arith.fastmath<reassoc>} %lhs, %flat, %accum : vector<16x32xbf16>, vector<32x16xbf16> into vector<16x16xf32>
----------------
jopperm wrote:

Should have asked more clearly: The test cases are about propagation of the flag when you create the replacement contraction, not influencing the pattern itself, correct?

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