[Mlir-commits] [llvm] [mlir] [mlir][linalg/scf/transform] scalable tiling and fusion for pack/unpack ops (PR #204007)
Andrzej Warzyński
llvmlistbot at llvm.org
Wed Jul 1 03:29:34 PDT 2026
================
@@ -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]]]
+// CHECK: %[[SRC:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, %{{.*}}, 0, 0] [%[[OUTER0]], %[[OUTER1]], %[[C8_VSCALE]], %[[C4_VSCALE]]]
+// CHECK: %[[UNPACK:.*]] = linalg.unpack %[[SRC]]
+// CHECK-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+// CHECK-NOT: tensor.extract_slice %[[UNPACK]]
+// CHECK: tensor.insert_slice %[[UNPACK]]
+// CHECK: return %[[RES]]
+func.func @NCnc_to_NC_scalable_aligned(%source: tensor<4x8x?x?xf32>, %dest: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %c8 = arith.constant 8 : index
+ %c4 = arith.constant 4 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %0 = linalg.unpack %source inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %dest
+ : tensor<4x8x?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 `Multiple` hints are passed to hint the alignment between the loop tile sizes
+ // (16 * vscale, 8 * vscale) and the inner tile sizes (8 * vscale, 4 * vscale).
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [[16], [8]]
+ inner_tile_alignments = [Multiple, Multiple]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Unaligned scalable tiling - static tile sizes not aligned to scalable
+// inner tiles.
+
+// CHECK-LABEL: func.func @NCnc_to_NC_scalable_unaligned
+// CHECK: %[[RES:.*]] = scf.for
+// CHECK: scf.for
+// CHECK: %[[UNPACK:.*]] = linalg.unpack
+// CHECK-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice %[[UNPACK]]
+// CHECK: tensor.insert_slice %[[EXTRACT]]
+// CHECK: return %[[RES]]
+func.func @NCnc_to_NC_scalable_unaligned(%source: tensor<4x8x?x?xf32>, %dest: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %c8 = arith.constant 8 : index
+ %c4 = arith.constant 4 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %0 = linalg.unpack %source inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %dest
+ : tensor<4x8x?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
+ // No hint is passed: the static tile sizes (7, 5) are not aligned to the scalable
+ // inner tiles (8 * vscale, 4 * vscale), so the tiled unpack keeps a remainder slice.
+ %1, %loops:2 = transform.structured.tile_using_for %0 tile_sizes [7, 5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Producer fusion - linalg.unpack with scalable inner tiles fused as a producer
+// into an elementwise consumer that is tiled with scalable tile
+// sizes that are an integer multiple of the inner tiles (16*vscale of 8*vscale,
+// 8*vscale of 4*vscale). `Multiple` hint is passed accordingly.
+
+// CHECK: #[[$MAP_CEILDIV:.+]] = affine_map<(d0)[s0] -> (d0 ceildiv s0)>
+// CHECK-LABEL: func.func @unpack_elemwise_scalable_multiple
+// 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: scf.for
+// CHECK: %[[OUTER0:.*]] = affine.apply #[[$MAP_CEILDIV]](%{{.*}})[%[[C8_VSCALE]]]
+// CHECK: %[[OUTER1:.*]] = affine.apply #[[$MAP_CEILDIV]](%{{.*}})[%[[C4_VSCALE]]]
+// CHECK: %[[SRC:.*]] = tensor.extract_slice %{{.*}}[%{{.*}}, %{{.*}}, 0, 0] [%[[OUTER0]], %[[OUTER1]], %[[C8_VSCALE]], %[[C4_VSCALE]]]
+// CHECK: %[[UNPACK:.*]] = linalg.unpack %[[SRC]]
+// CHECK-SAME: tensor<?x?x?x?xf32> -> tensor<?x?xf32>
+// CHECK-NOT: tensor.extract_slice %[[UNPACK]]
+// CHECK: linalg.exp ins(%[[UNPACK]]
+// CHECK: return %[[RES]]
+func.func @unpack_elemwise_scalable_multiple(%arg0: tensor<4x8x?x?xf32>, %arg1: tensor<?x?xf32>, %arg2 : index, %arg3 : index) -> tensor<?x?xf32> {
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %c16 = arith.constant 16 : index
+ %vscale = vector.vscale
+ %c8_vscale = arith.muli %c8, %vscale : index
+ %c4_vscale = arith.muli %c4, %vscale : index
+ %t0 = arith.muli %c16, %vscale : index
+ %t1 = arith.muli %c8, %vscale : index
+ %0 = tensor.empty(%arg2, %arg3) : tensor<?x?xf32>
+ %1 = linalg.unpack %arg0 inner_dims_pos = [0, 1]
+ inner_tiles = [%c8_vscale, %c4_vscale] into %0
+ : tensor<4x8x?x?xf32> -> tensor<?x?xf32>
+ %2 = linalg.exp ins(%1: tensor<?x?xf32>)
+ outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %2 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %exp = transform.structured.match ops{["linalg.exp"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %mulis = transform.structured.match ops{["arith.muli"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %i0, %i1, %t0h, %t1h = transform.split_handle %mulis : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
----------------
banach-space wrote:
Is this required here? Normally you'd pass tile sizes as e.g. `[[16], [4]]`, no?
https://github.com/llvm/llvm-project/pull/204007
More information about the Mlir-commits
mailing list