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

[Ege Beysel via llvm-commits]

================
@@ -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]
+      inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, %c4_vscale]
+      into %dest : tensor<256x128xf32> -> tensor<?x?x?x?xf32>
+  return %pack : tensor<?x?x?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %pack = transform.structured.match ops{["linalg.pack"]} in %arg1
+        : (!transform.any_op) -> !transform.any_op
+    %loop = transform.structured.match ops{["scf.for"]} in %arg1
+        : (!transform.any_op) -> !transform.any_op
+    %a, %b = transform.test.fuse_consumer %pack into (%loop) {inner_tile_alignments = array<i64: 2, 0>}
+        : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
+
+// -----
+
+// Consumer fusion with a static producer step (64) and a scalable pack inner
+// tile (8*vscale), hinted `Multiple`. This is a caller assertion, the relationship
+// is not statically decidable, so fusion takes the aligned (non-equal) path and
+// the outer dim of the fused pack tile is dynamic (`64 ceildiv 8*vscale`).
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func.func @fuse_scalable_pack_consumer_aligned
+// CHECK-SAME:      %[[ARG0:.+]]: tensor<256x128xf32>, %[[ARG1:.+]]: tensor<256x128xf32>, %[[ARG2:.+]]: tensor<256x128xf32>, %[[DEST:.+]]: tensor<?x?x?x?xf32>
+//      CHECK:    %[[C64:.*]] = arith.constant 64 : index
+//      CHECK:    %[[RES:.*]]:2 = scf.for {{.*}} step %[[C64]]
+// CHECK-SAME:        iter_args(%{{.*}} = %[[ARG2]], %{{.*}} = %[[DEST]])
+//      CHECK:      %[[GENERIC:.*]] = linalg.generic
+//      CHECK:      %[[PACK:.*]] = linalg.pack %[[GENERIC]]
+// CHECK-SAME:          -> tensor<?x?x?x?xf32>
----------------
egebeysel wrote:

done

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