[Mlir-commits] [mlir] [mlir][linalg/scf/transform] scalable tiling and fusion for pack/unpack ops (PR #204007)

[Andrzej Warzyński]

================
@@ -0,0 +1,457 @@
+// RUN: mlir-opt %s -transform-interpreter -canonicalize -cse -split-input-file --verify-diagnostics | FileCheck %s
+
+// Consumer fusion - linalg.pack with scalable inner tiles.
+// Producer step (8*vscale) equals the pack inner tile size
+// (8*vscale) on the tiled source dimension, so the outer
+// dim of the fused pack tile is statically 1.
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func.func @fuse_scalable_pack_consumer_equal
+// CHECK-SAME:      %[[ARG0:.+]]: tensor<256x128xf32>, %[[ARG1:.+]]: tensor<256x128xf32>, %[[ARG2:.+]]: tensor<256x128xf32>, %[[DEST:.+]]: tensor<?x?x?x?xf32>
+//      CHECK:    %[[C8:.*]] = arith.constant 8 : index
+//      CHECK:    %[[VSCALE:.*]] = vector.vscale
+//      CHECK:    %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+//      CHECK:    %[[RES:.*]]:2 = scf.for {{.*}} step %[[C8_VSCALE]]
+// CHECK-SAME:        iter_args(%{{.*}} = %[[ARG2]], %{{.*}} = %[[DEST]])
+//      CHECK:      %[[GENERIC:.*]] = linalg.generic
+//      CHECK:      %[[PACK:.*]] = linalg.pack %[[GENERIC]]
+// CHECK-SAME:          inner_tiles = [%[[C8_VSCALE]], %{{.*}}]
+// CHECK-SAME:          -> tensor<1x?x?x?xf32>
+//      CHECK:      scf.yield {{.*}}, %{{.*}} :
+//      CHECK:    return %[[RES]]#1
+func.func @fuse_scalable_pack_consumer_equal(
+    %arg0: tensor<256x128xf32>, %arg1: tensor<256x128xf32>,
+    %arg2: tensor<256x128xf32>, %dest: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
+  %c0 = arith.constant 0 : index
+  %c4 = arith.constant 4 : index
+  %c8 = arith.constant 8 : index
+  %c256 = arith.constant 256 : index
+  %vscale = vector.vscale
+  %c4_vscale = arith.muli %c4, %vscale : index
+  %c8_vscale = arith.muli %c8, %vscale : index
+
+  %0 = scf.for %iv = %c0 to %c256 step %c8_vscale iter_args(%out = %arg2) -> (tensor<256x128xf32>) {
+    %ext_out = tensor.extract_slice %out[%iv, 0] [%c8_vscale, 128] [1, 1]
+        : tensor<256x128xf32> to tensor<?x128xf32>
+    %ext_a = tensor.extract_slice %arg0[%iv, 0] [%c8_vscale, 128] [1, 1]
+        : tensor<256x128xf32> to tensor<?x128xf32>
+    %ext_b = tensor.extract_slice %arg1[%iv, 0] [%c8_vscale, 128] [1, 1]
+        : tensor<256x128xf32> to tensor<?x128xf32>
+    %computed = linalg.generic {
+        indexing_maps = [#map, #map, #map],
+        iterator_types = ["parallel", "parallel"]}
+        ins(%ext_a, %ext_b : tensor<?x128xf32>, tensor<?x128xf32>)
+        outs(%ext_out : tensor<?x128xf32>) {
+      ^bb0(%in0: f32, %in1: f32, %out_elem: f32):
+        %mul = arith.mulf %in0, %in1 : f32
+        linalg.yield %mul : f32
+    } -> tensor<?x128xf32>
+    %inserted = tensor.insert_slice %computed into %out[%iv, 0] [%c8_vscale, 128] [1, 1]
+        : tensor<?x128xf32> into tensor<256x128xf32>
+    scf.yield %inserted : tensor<256x128xf32>
+  }
+
+  %pack = linalg.pack %0 outer_dims_perm = [0, 1]
----------------
banach-space wrote:

[nit] Add comments to high-light that the loop tile and `linalg.pack` inner tile are identical. Otherwise it's hard to spot.

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