test/Dialect/Linalg/transform-op-multitile-sizes.mlir

// RUN: mlir-opt %s --transform-interpreter --split-input-file --verify-diagnostics | FileCheck %s

// CHECK-DAG: #[[$MAP13:.+]] = affine_map<() -> (13)>

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
      %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
      transform.structured.multitile_sizes %0 { target_size = 3, dimension = 0 } : (!transform.any_op) -> !transform.any_op
      transform.yield
  }
}

// CHECK-LABEL: @multitile_sizes_static
func.func @multitile_sizes_static(
  %arg0: tensor<13x34xf32>, %arg1: tensor<34x42xf32>, %arg2: tensor<13x42xf32>)
    -> tensor<13x42xf32> {
  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<13x34xf32>, tensor<34x42xf32>)
                     outs(%arg2: tensor<13x42xf32>)
    -> tensor<13x42xf32>
  // The first application computes the total size.
  // CHECK: %{{.*}} = affine.apply #[[$MAP13]]()
  // CHECK: %[[SIZE:.+]] = affine.apply #[[$MAP13]]()
  // CHECK: %[[COND:.+]] = arith.cmpi eq, %[[SIZE]], %{{.*}}
  // CHECK: cf.assert %[[COND]], "could not compute dynamic multi-size tile shapes"

  return %0 : tensor<13x42xf32>
}

// -----

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
      %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
      %low_tile, %high_tile, %split_point =
        transform.structured.multitile_sizes %0 { target_size = 3, dimension = 0 }
        : (!transform.any_op) -> !transform.param<i64>
      // expected-remark @below {{2 : i64}}
      transform.debug.emit_param_as_remark %low_tile : !transform.param<i64>
      // expected-remark @below {{3 : i64}}
      transform.debug.emit_param_as_remark %high_tile : !transform.param<i64>
      // expected-remark @below {{4 : i64}}
      transform.debug.emit_param_as_remark %split_point : !transform.param<i64>
      transform.yield
  }
}

// CHECK-LABEL: @multitile_sizes_static_gen
func.func @multitile_sizes_static_gen(
  %arg0: tensor<13x34xf32>, %arg1: tensor<34x42xf32>, %arg2: tensor<13x42xf32>)
    -> tensor<13x42xf32> {
  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<13x34xf32>, tensor<34x42xf32>)
                     outs(%arg2: tensor<13x42xf32>)
    -> tensor<13x42xf32>

  return %0 : tensor<13x42xf32>
}

// -----

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
      %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
      transform.structured.multitile_sizes %0 { target_size = 3, divisor = 2, dimension = 0 } : (!transform.any_op) -> !transform.any_op
      transform.yield
  }
}

// CHECK: #[[$MAP_A:.+]] = affine_map<()[s0] -> ([[A_IMPL:s0 floordiv 2]])>
// CHECK: #[[$MAP_T:.+]] = affine_map<() -> (2)>
// CHECK: #[[$MAP_D:.+]] = affine_map<()[s0] -> ([[D_IMPL:\(s0 floordiv 2 \+ 1\) floordiv 2]])>
// CHECK: #[[$MAP_S:.+]] = affine_map<()[s0] -> ((([[A_IMPL]]) floordiv ([[D_IMPL]])) * 2)>
// CHECK: #[[$MAP_V:.+]] = affine_map<()[s0] -> (([[A_IMPL]]) mod ([[D_IMPL]]))>
// CHECK: #[[$MAP_U:.+]] = affine_map<()[s0] -> ([[D_IMPL]] - ([[A_IMPL]]) mod ([[D_IMPL]]))>

// CHECK-LABEL: @multitile_sizes_dynamic
// CHECK-SAME: (%[[ARG0:.+]]: tensor<?x?xf32>, %{{.*}}: tensor<?x?xf32>, %{{.*}}: tensor<?x?xf32>)
func.func @multitile_sizes_dynamic(
  // For matmul, the extent of the first iteration space dimension is equal to
  // the size of the first dimension of the first tensor. The indexing map was
  // folded so there is no map application happening.
  //
  // CHECK: %[[C0:.+]] = arith.constant 0
  // CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]]
  //
  // The following are the maps as emitted by computeMultiTileSizes.
  // CHECK: affine.apply #[[$MAP_A]]()[%[[DIM]]]
  // CHECK: affine.apply #[[$MAP_T]]()
  // CHECK: affine.apply #[[$MAP_D]]()[%[[DIM]]]
  // CHECK: affine.apply #[[$MAP_S]]()[%[[DIM]]]
  // CHECK: affine.apply #[[$MAP_V]]()[%[[DIM]]]
  // CHECK: affine.apply #[[$MAP_U]]()[%[[DIM]]]
  %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
    -> tensor<?x?xf32> {
  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
                     outs(%arg2: tensor<?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.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
      // expected-error @below {{cannot compute parametric tile sizes for dynamically shaped payload op}}
      transform.structured.multitile_sizes %0 { target_size = 3, divisor = 2, dimension = 0 }
        : (!transform.any_op) -> !transform.param<i64>
        transform.yield
  }
}

func.func @multitile_sizes_dynamic_gen(
  %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
    -> tensor<?x?xf32> {
  // expected-note @below {{payload op}}
  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
                     outs(%arg2: tensor<?x?xf32>)
    -> tensor<?x?xf32>

  return %0 : tensor<?x?xf32>
}