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

[Andrzej Warzyński via llvm-commits]

================
@@ -0,0 +1,784 @@
+// RUN: mlir-opt %s -transform-interpreter -canonicalize -cse -split-input-file --verify-diagnostics | FileCheck %s
+
+// The transform ops below carry an optional `inner_tile_alignments` hint,
+// written as a per-dimension keyword list:
+//   - `Equal`:    the loop tile size equals the pack/unpack inner tile size.
+//   - `Multiple`: the loop tile size is an integer multiple of the inner tile.
+//   - `Unknown`:  the default; nothing is asserted for that dimension.
+
+// Perfect scalable tiling - scalable tile sizes equal scalable inner
+// tiles. Outer sizes of the tiled unpack should be 1's.
+
+// CHECK-LABEL: func.func @perfect_CKkc_to_KC_scalable
+// CHECK:         %[[RES:.*]] = scf.for
+// CHECK:           scf.for
+// CHECK:             %[[UNPACK:.*]] = linalg.unpack
+// CHECK-SAME:            tensor<1x1x?x?xf32> -> tensor<?x?xf32>
+// CHECK-NOT:         tensor.extract_slice %[[UNPACK]]
+// CHECK:             tensor.insert_slice %[[UNPACK]]
+// CHECK:         return %[[RES]]
+func.func @perfect_CKkc_to_KC_scalable(%source: tensor<32x4x?x?xf32>, %dest: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+  %vscale = vector.vscale
+  %c2_vscale = arith.muli %c2, %vscale : index
+  %c4_vscale = arith.muli %c4, %vscale : index
+  %0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
+      inner_tiles = [%c2_vscale, %c4_vscale] into %dest
+      : tensor<32x4x?x?xf32> -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+      %0 = transform.structured.match ops{["linalg.unpack"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+      // The `Equal` hints assert the loop tile sizes (2 * vscale, 4 * vscale) equal the
+      // inner tile sizes (2 * vscale, 4 * vscale).
+      %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [[2], [4]]
+          inner_tile_alignments = [Equal, Equal]
+          : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+      transform.yield
+  }
+}
+
+// -----
+
+// Aligned scalable tiling - scalable tile sizes ([16] = 16*vscale, [8] =
+// 8*vscale) that are integer multiples of the scalable inner tiles (2x 8*vscale
+// and 2x 4*vscale), so tiling is aligned and the tiled unpack needs no trailing
+// remainder slice.
+
+// CHECK:       #[[$MAP_CEILDIV:.+]] = affine_map<(d0)[s0] -> (d0 ceildiv s0)>
+// CHECK-LABEL: func.func @NCnc_to_NC_scalable_aligned
+//  CHECK-DAG:    %[[C8:.*]] = arith.constant 8 : index
+//  CHECK-DAG:    %[[C4:.*]] = arith.constant 4 : index
+//  CHECK-DAG:    %[[VSCALE:.*]] = vector.vscale
+//  CHECK-DAG:    %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+//  CHECK-DAG:    %[[C4_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index
+//      CHECK:    %[[RES:.*]] = scf.for
+//      CHECK:      %[[OUTER0:.*]] = affine.apply #[[$MAP_CEILDIV]](%{{.*}})[%[[C8_VSCALE]]]
+//      CHECK:      %[[OUTER1:.*]] = affine.apply #[[$MAP_CEILDIV]](%{{.*}})[%[[C4_VSCALE]]]
----------------
banach-space wrote:

[nit] `OUTER0` -> `OUTER_DIM0`

Similar comment for `OUTER1` and other tests.

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