Google Git
Sign in
llvm / llvm-project / mlir / 0b3d5b8838288f5873894c53f30b007d6c83ee1d / . / test / Dialect / Linalg / transform-pack-greedily.mlir
blob: 100692426ef44ccc8b125f749c54dddbc134f99a [file] [log] [blame]
// RUN: mlir-opt %s -transform-interpreter --split-input-file | FileCheck %s
!A_mk = tensor<1023x255xf32>
!B_kn = tensor<255x127xf32>
!C_mn = tensor<1023x127xf32>
// Normalized dims are: ( k, m, n)(kk, mm, nn)
// CHECK-DAG: #[[$mk_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d0, d3, d4)>
// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>
// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d4, d5)>
// CHECK-LABEL: @matmul_mk_kn_mn(
func.func @matmul_mk_kn_mn(%A : !A_mk, %B : !B_kn, %C : !C_mn) -> !C_mn {
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]
// CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}
// CHECK-SAME: ins(%{{.*}} : tensor<128x8x32x8xf32>, tensor<8x8x32x16xf32>)
// CHECK-SAME: outs(%{{.*}} : tensor<128x8x8x16xf32>)
%0 = linalg.matmul ins(%A, %B : !A_mk, !B_kn) outs(%C : !C_mn) -> !C_mn
return %0 : !C_mn
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%matmul = transform.structured.match ops{["linalg.matmul"]} in %module_op
: (!transform.any_op) -> !transform.op<"linalg.matmul">
transform.structured.pack_greedily %matmul
matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
: (!transform.op<"linalg.matmul">) -> !transform.op<"linalg.generic">
transform.yield
}
}
// -----
!A_mk = tensor<1023x255xf32>
!B_nk = tensor<127x255xf32>
!C_nm = tensor<127x1023xf32>
#mkn_accesses = [
affine_map<(m, n, k) -> (m, k)>,
affine_map<(m, n, k) -> (n, k)>,
affine_map<(m, n, k) -> (n, m)>
]
#mkn_trait = {
indexing_maps = #mkn_accesses,
iterator_types = ["parallel", "parallel", "reduction"]
}
// Normalized dims are: ( k, m, n)(kk, mm, nn)
// CHECK-DAG: #[[$km_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d0, d3, d4)>
// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d0, d3, d5)>
// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d4, d5)>
// CHECK-LABEL: @matmul_mk_nk_nm(
func.func @matmul_mk_nk_nm(%A : !A_mk, %B : !B_nk, %C : !C_nm) -> !C_nm {
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]
// CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}
// CHECK-SAME: ins(%{{.*}} : tensor<128x8x32x8xf32>, tensor<8x8x32x16xf32>)
// CHECK-SAME: outs(%{{.*}} : tensor<8x128x8x16xf32>)
%0 = linalg.generic #mkn_trait ins(%A, %B : !A_mk, !B_nk) outs(%C : !C_nm) {
^bb0(%a: f32, %b: f32, %c: f32):
%d = arith.mulf %a, %b : f32
%e = arith.addf %c, %d : f32
linalg.yield %e : f32
} -> !C_nm
return %0 : !C_nm
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">
transform.structured.pack_greedily %generic
matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
: (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">
transform.yield
}
}
// -----
!A_mk = tensor<1023x255xf32>
!B_nk = tensor<127x255xf32>
!C_nm = tensor<127x1023xf32>
#mkn_accesses = [
affine_map<(k, m, n) -> (m, k)>,
affine_map<(k, m, n) -> (n, k)>,
affine_map<(k, m, n) -> (n, m)>
]
#mkn_trait = {
indexing_maps = #mkn_accesses,
iterator_types = ["reduction", "parallel", "parallel"]
}
// Normalized dims are: ( k, m, n)(kk, mm, nn)
// CHECK-DAG: #[[$mk_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d0, d3, d4)>
// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d0, d3, d5)>
// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d4, d5)>
// CHECK-LABEL: @matmul_mk_nk_nm_transposed(
func.func @matmul_mk_nk_nm_transposed(%A : !A_mk, %B : !B_nk, %C : !C_nm) -> !C_nm {
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]
// CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}
// CHECK-SAME: ins(%{{.*}} : tensor<128x8x32x8xf32>, tensor<8x8x32x16xf32>)
// CHECK-SAME: outs(%{{.*}} : tensor<8x128x8x16xf32>)
%0 = linalg.generic #mkn_trait ins(%A, %B : !A_mk, !B_nk) outs(%C : !C_nm) {
^bb0(%a: f32, %b: f32, %c: f32):
%d = arith.mulf %a, %b : f32
%e = arith.addf %c, %d : f32
linalg.yield %e : f32
} -> !C_nm
return %0 : !C_nm
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">
transform.structured.pack_greedily %generic
matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
: (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">
transform.yield
}
}
// -----
!A_bmkm2 = tensor<42x1023x255x33xf32>
!B_nkb = tensor<127x255x42xf32>
!C_nbm = tensor<127x42x1023xf32>
#mkn_accesses = [
affine_map<(k, m, n, b, m2) -> (b, m, k, m2)>,
affine_map<(k, m, n, b, m2) -> (n, k, b)>,
affine_map<(k, m, n, b, m2) -> (n, b, m)>
]
#mkn_trait = {
indexing_maps = #mkn_accesses,
iterator_types = ["reduction", "parallel", "parallel", "parallel", "parallel"]
}
// Normalized dims are: ( ?, ?, k, m, n)(kk, mm, nn)
// CHECK-DAG: #[[$bmkm2_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d3, d2, d1, d5, d6)>
// CHECK-DAG: #[[$nkb_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d4, d2, d0, d5, d7)>
// CHECK-DAG: #[[$nbm_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d4, d0, d3, d6, d7)>
// CHECK-LABEL: @contraction_bmkm2_nkb_nbm(
func.func @contraction_bmkm2_nkb_nbm(%A : !A_bmkm2, %B : !B_nkb, %C : !C_nbm) -> !C_nbm {
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$bmkm2_kkmm]], #[[$nkb_kknn]], #[[$nbm_mmnn]]]
// CHECK-SAME: ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}
// CHECK-SAME: ins(%{{.*}} : tensor<42x128x8x33x32x8xf32>, tensor<8x8x42x32x16xf32>)
// CHECK-SAME: outs(%{{.*}} : tensor<8x42x128x8x16xf32>)
%0 = linalg.generic #mkn_trait ins(%A, %B : !A_bmkm2, !B_nkb) outs(%C : !C_nbm) {
^bb0(%a: f32, %b: f32, %c: f32):
%d = arith.mulf %a, %b : f32
%e = arith.addf %c, %d : f32
linalg.yield %e : f32
} -> !C_nbm
return %0 : !C_nbm
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">
transform.structured.pack_greedily %generic
matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
: (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">
transform.yield
}
}
// -----
// Conv linguo: h w kh kw c n f cc nn ff
// Normalized dims are: ( ?, ?, ?, ?, k, m, n)(kk, mm, nn)
// n c h + kh w + kw cc nn
// CHECK-DAG: #[[$M1:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8, d9) -> (d5, d4, d0 + d2, d1 + d3, d7, d8)>
// f c kh kw cc ff
// CHECK-DAG: #[[$M2:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8, d9) -> (d6, d4, d2, d3, d7, d9)>
// n f h w nn ff
// CHECK-DAG: #[[$M3:.*]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8, d9) -> (d5, d6, d0, d1, d8, d9)>
// CHECK-LABEL: @conv_2d_nchw_fchw
func.func @conv_2d_nchw_fchw(%arg0: tensor<?x47x16x16xf32>, %arg2: tensor<?x16x14x14xf32>) -> tensor<?x16x14x14xf32> {
%c0 = arith.constant dense<0.1> : tensor<16x47x3x3xf32>
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$M1]], #[[$M2]], #[[$M3]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "reduction", "reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]
// CHECK-SAME: ins(%{{.*}} : tensor<?x2x16x16x32x8xf32>, tensor<1x2x3x3x32x16xf32>)
// CHECK-SAME: outs(%{{.*}} : tensor<?x1x14x14x8x16xf32>)
%0 = linalg.conv_2d_nchw_fchw
{dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
ins(%arg0, %c0: tensor<?x47x16x16xf32>, tensor<16x47x3x3xf32>)
outs(%arg2: tensor<?x16x14x14xf32>) -> tensor<?x16x14x14xf32>
return %0 : tensor<?x16x14x14xf32>
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%conv = transform.structured.match ops{["linalg.conv_2d_nchw_fchw"]} in %module_op
: (!transform.any_op) -> !transform.op<"linalg.conv_2d_nchw_fchw">
transform.structured.pack_greedily %conv
matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
: (!transform.op<"linalg.conv_2d_nchw_fchw">) -> !transform.op<"linalg.generic">
transform.yield
}
}
// -----
// These should fail to pack for now as they don't contain a contraction.
// CHECK-LABEL: @reduce_and_map
func.func @reduce_and_map(%arg0: tensor<10x100xf32>,
%arg1: tensor<10x100xf32>, %output: tensor<10xf32>) -> tensor<10xf32> {
%map_init = tensor.empty() : tensor<10x100xf32>
// CHECK: linalg.map
%mapped = linalg.map { arith.addf }
ins(%arg0, %arg1 : tensor<10x100xf32>, tensor<10x100xf32>)
outs(%map_init : tensor<10x100xf32>)
// CHECK: linalg.reduce
%res = linalg.reduce { arith.addf }
ins(%mapped: tensor<10x100xf32>)
outs(%output: tensor<10xf32>)
dimensions = [1]
return %res : tensor<10xf32>
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">
transform.structured.pack_greedily %generic
matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
: (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">
transform.yield
}
}
// -----
!A_mk = tensor<1023x255xf32>
!B_nk = tensor<127x255xf32>
!C_nm = tensor<127x1023xf32>
#mkn_accesses = [
affine_map<(m, n, k) -> (m, k)>,
affine_map<(m, n, k) -> (n, k)>,
affine_map<(m, n, k) -> (n, m)>
]
#mkn_trait = {
indexing_maps = #mkn_accesses,
iterator_types = ["parallel", "parallel", "reduction"]
}
// Normalized dims are: ( k, m, n)(kk, mm, nn)
// CHECK-DAG: #[[$km_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d0, d3, d4)>
// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d0, d3, d5)>
// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d1, d4, d5)>
// CHECK-LABEL: @matmul_mk_nk_nm(
func.func @matmul_mk_nk_nm(%A : !A_mk, %B : !B_nk, %C : !C_nm) -> !C_nm {
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]
// CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel", "parallel"]}
// CHECK-SAME: ins(%{{.*}} : tensor<128x8x32x8xf32>, tensor<1x8x32x130xf32>)
// CHECK-SAME: outs(%{{.*}} : tensor<1x128x8x130xf32>)
%0 = linalg.generic #mkn_trait ins(%A, %B : !A_mk, !B_nk) outs(%C : !C_nm) {
^bb0(%a: f32, %b: f32, %c: f32):
%d = arith.mulf %a, %b : f32
%e = arith.addf %c, %d : f32
linalg.yield %e : f32
} -> !C_nm
return %0 : !C_nm
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">
transform.structured.pack_greedily %generic
// In this spec, the "k" dimension is not packed but rather padded to the
// next multiple of 10 (i.e. 130).
matmul_packed_sizes = [8, 0, 32]
matmul_padded_sizes_next_multiple_of = [0, 10, 0]
matmul_inner_dims_order = [1, 2, 0]
: (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">
transform.yield
}
}
// -----
!A_mk = tensor<1023x255xf32>
!B_nk = tensor<127x255xf32>
!C_nm = tensor<127x1023xf32>
#mkn_accesses = [
affine_map<(m, n, k) -> (m, k)>,
affine_map<(m, n, k) -> (n, k)>,
affine_map<(m, n, k) -> (n, m)>
]
#mkn_trait = {
indexing_maps = #mkn_accesses,
iterator_types = ["parallel", "parallel", "reduction"]
}
// Normalized dims are: ( k, m, n)(kk, mm)
// CHECK-DAG: #[[$km_kkmm:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d0, d3)>
// CHECK-DAG: #[[$kn_kknn:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d0, d3, d4)>
// CHECK-DAG: #[[$mn_mmnn:.*]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d1, d4)>
// CHECK-LABEL: @matmul_mk_nk_nm(
func.func @matmul_mk_nk_nm(%A : !A_mk, %B : !B_nk, %C : !C_nm) -> !C_nm {
// CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$mk_kkmm]], #[[$kn_kknn]], #[[$mn_mmnn]]]
// CHECK-SAME: ["reduction", "parallel", "parallel", "reduction", "parallel"]}
// CHECK-SAME: ins(%{{.*}} : tensor<1023x8x32xf32>, tensor<1x8x32x130xf32>)
// CHECK-SAME: outs(%{{.*}} : tensor<1x1023x130xf32>)
%0 = linalg.generic #mkn_trait ins(%A, %B : !A_mk, !B_nk) outs(%C : !C_nm) {
^bb0(%a: f32, %b: f32, %c: f32):
%d = arith.mulf %a, %b : f32
%e = arith.addf %c, %d : f32
linalg.yield %e : f32
} -> !C_nm
return %0 : !C_nm
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%generic = transform.structured.match ops{["linalg.generic"]} in %module_op : (!transform.any_op) -> !transform.op<"linalg.generic">
transform.structured.pack_greedily %generic
// In this spec, the "n" dimension is neither packed not unpacked.
// We don't end up with an innermost matmul after packing but only with an
// innermost matvec.
matmul_packed_sizes = [0, 0, 32]
matmul_padded_sizes_next_multiple_of = [0, 10, 0]
matmul_inner_dims_order = [1, 2, 0]
: (!transform.op<"linalg.generic">) -> !transform.op<"linalg.generic">
transform.yield
}
}
// -----
!A = tensor<1023x255xf32>
!X = tensor<255xf32>
!Y = tensor<1023xf32>
// CHECK-LABEL: @matvec_fail(
func.func @matvec_fail(%A : !A, %x : !X, %y : !Y) -> !Y {
// CHECK: linalg.matvec
%0 = linalg.matvec ins(%A, %x : !A, !X) outs(%y : !Y) -> !Y
return %0 : !Y
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
%matmul = transform.structured.match ops{["linalg.matvec"]} in %module_op
: (!transform.any_op) -> !transform.op<"linalg.matvec">
transform.structured.pack_greedily %matmul
matmul_packed_sizes = [8, 16, 32] matmul_inner_dims_order = [1, 2, 0]
: (!transform.op<"linalg.matvec">) -> !transform.any_op
transform.yield
}
}
// -----
func.func @no_padding_on_packs(%A: tensor<32x32xf32>, %B: tensor<32x32xf32>, %C: tensor<32x32xf32>)
-> tensor<32x32xf32> {
%0 = linalg.matmul ins(%A, %B: tensor<32x32xf32>, tensor<32x32xf32>)
outs(%C: tensor<32x32xf32>)
-> tensor<32x32xf32>
return %0 : tensor<32x32xf32>
}
// CHECK-LABEL: no_padding_on_packs
// CHECK: tensor.pack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [8, 4]
// CHECK-SAME: into %{{.+}} : tensor<32x32xf32> -> tensor<4x8x8x4xf32>
// CHECK: tensor.pack %{{.+}} outer_dims_perm = [1, 0]
// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [4, 16] into %{{.+}} : tensor<32x32xf32> -> tensor<2x8x4x16xf32>
// CHECK: tensor.pack %{{.+}} inner_dims_pos = [0, 1] inner_tiles = [8, 16]
// CHECK-SAME: into %{{.+}} : tensor<32x32xf32> -> tensor<4x2x8x16xf32>
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.op<"linalg.matmul">
%1 = transform.structured.pack_greedily %0
matmul_packed_sizes = [8, 16, 4] matmul_inner_dims_order = [0, 1, 2]
: (!transform.op<"linalg.matmul">) -> !transform.op<"linalg.generic">
%pack = transform.get_producer_of_operand %1[1]
: (!transform.op<"linalg.generic">) -> (!transform.op<"tensor.pack">)
%2, %pack_2, %empty_unpack_2 =
transform.structured.pack_transpose %pack with_compute_op(%1)
outer_perm = [1, 0] inner_perm = [1, 0]
: (!transform.op<"tensor.pack">, !transform.op<"linalg.generic">)
-> (!transform.op<"linalg.generic">, !transform.op<"tensor.pack">, !transform.any_op)
transform.yield
}
}