364 lines · plain
1// RUN: mlir-opt %s -linalg-block-pack-matmul=block-factors=32,16,64 -canonicalize -split-input-file | FileCheck %s2 3func.func @block_matmul(4 %A: tensor<128x128xf32>, %B: tensor<128x128xf32>, %C: tensor<128x128xf32>) -> tensor<128x128xf32> {5 %0 = linalg.matmul ins(%A, %B : tensor<128x128xf32>, tensor<128x128xf32>)6 outs(%C : tensor<128x128xf32>) -> tensor<128x128xf32>7 return %0 : tensor<128x128xf32>8}9 10// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>11// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d4, d5)>12// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)>13 14// CHECK-LABEL: func @block_matmul(15// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<128x128xf32>, %[[B:[0-9a-z]+]]: tensor<128x128xf32>, %[[C:[0-9a-z]+]]: tensor<128x128xf32>16// CHECK: %[[PACK_DST_0:.+]] = tensor.empty() : tensor<4x2x32x64xf32>17// CHECK: %[[A_PACKED:.+]] = linalg.pack %[[A]]18// CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [32, 64]19// CHECK-SAME: into %[[PACK_DST_0]] : tensor<128x128xf32> -> tensor<4x2x32x64xf32>20// CHECK: %[[PACK_DST_1:.+]] = tensor.empty() : tensor<8x2x16x64xf32>21// CHECK: %[[B_PACKED:.+]] = linalg.pack %[[B]]22// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 64]23// CHECK-SAME: into %[[PACK_DST_1]] : tensor<128x128xf32> -> tensor<8x2x16x64xf32>24// CHECK: %[[PACK_DST_2:.+]] = tensor.empty() : tensor<4x8x32x16xf32>25// CHECK: %[[C_PACKED:.+]] = linalg.pack %[[C]]26// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]27// CHECK-SAME: into %[[PACK_DST_2]] : tensor<128x128xf32> -> tensor<4x8x32x16xf32>28// CHECK: %[[GEMM_RES_PACKED:.+]] = linalg.generic29// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]30// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]31// CHECK-SAME: ins(%[[A_PACKED]], %[[B_PACKED]] : tensor<4x2x32x64xf32>, tensor<8x2x16x64xf32>) outs(%[[C_PACKED]] : tensor<4x8x32x16xf32>)32// CHECK: %[[RES_UNPACKED:.+]] = linalg.unpack %[[GEMM_RES_PACKED]]33// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]34// CHECK-SAME: into %[[C]] : tensor<4x8x32x16xf32> -> tensor<128x128xf32>35// CHECK: return %[[RES_UNPACKED]] : tensor<128x128xf32>36 37// -----38 39func.func @block_matmul_dynamic(40 %A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {41 %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)42 outs(%C : tensor<?x?xf32>) -> tensor<?x?xf32>43 return %0 : tensor<?x?xf32>44}45 46// CHECK-DAG: #[[$MAP_M:.+]] = affine_map<()[s0] -> (s0 ceildiv 32)>47// CHECK-DAG: #[[$MAP_K:.+]] = affine_map<()[s0] -> (s0 ceildiv 64)>48// CHECK-DAG: #[[$MAP_N:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>49// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>50// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d4, d5)>51// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)>52 53// CHECK-LABEL: func @block_matmul_dynamic(54// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<?x?xf32>, %[[B:[0-9a-z]+]]: tensor<?x?xf32>, %[[C:[0-9a-z]+]]: tensor<?x?xf32>55// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index56// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index57// CHECK-DAG: %[[ZERO:.+]] = arith.constant 0.000000e+00 : f3258// CHECK-DAG: %[[A_M:.+]] = tensor.dim %[[A]], %[[C0]] : tensor<?x?xf32>59// CHECK-DAG: %[[A_K:.+]] = tensor.dim %[[A]], %[[C1]] : tensor<?x?xf32>60// CHECK-DAG: %[[A_OUTER_TILE_M:.+]] = affine.apply #[[$MAP_M]]()[%[[A_M]]]61// CHECK-DAG: %[[A_OUTER_TILE_K:.+]] = affine.apply #[[$MAP_K]]()[%[[A_K]]]62// CHECK: %[[PACK_DST_0:.+]] = tensor.empty(%[[A_OUTER_TILE_M]], %[[A_OUTER_TILE_K]]) : tensor<?x?x32x64xf32>63// CHECK: %[[A_PACKED:.+]] = linalg.pack %[[A]]64// CHECK-SAME: padding_value(%[[ZERO]] : f32)65// CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [32, 64]66// CHECK-SAME: into %[[PACK_DST_0]] : tensor<?x?xf32> -> tensor<?x?x32x64xf32>67// CHECK-DAG: %[[B_K:.+]] = tensor.dim %[[B]], %[[C0]] : tensor<?x?xf32>68// CHECK-DAG: %[[B_N:.+]] = tensor.dim %[[B]], %[[C1]] : tensor<?x?xf32>69// CHECK-DAG: %[[B_OUTER_TILE_K:.+]] = affine.apply #[[$MAP_K]]()[%[[B_K]]]70// CHECK-DAG: %[[B_OUTER_TILE_N:.+]] = affine.apply #[[$MAP_N]]()[%[[B_N]]]71// CHECK: %[[PACK_DST_1:.+]] = tensor.empty(%[[B_OUTER_TILE_N]], %[[B_OUTER_TILE_K]]) : tensor<?x?x16x64xf32>72// CHECK: %[[B_PACKED:.+]] = linalg.pack %[[B]]73// CHECK-SAME: padding_value(%[[ZERO]] : f32)74// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 64]75// CHECK-SAME: into %[[PACK_DST_1]] : tensor<?x?xf32> -> tensor<?x?x16x64xf32>76// CHECK-DAG: %[[C_M:.+]] = tensor.dim %[[C]], %[[C0]] : tensor<?x?xf32>77// CHECK-DAG: %[[C_N:.+]] = tensor.dim %[[C]], %[[C1]] : tensor<?x?xf32>78// CHECK-DAG: %[[C_OUTER_TILE_M:.+]] = affine.apply #[[$MAP_M]]()[%[[C_M]]]79// CHECK-DAG: %[[C_OUTER_TILE_N:.+]] = affine.apply #[[$MAP_N]]()[%[[C_N]]]80// CHECK: %[[PACK_DST_2:.+]] = tensor.empty(%[[C_OUTER_TILE_M]], %[[C_OUTER_TILE_N]]) : tensor<?x?x32x16xf32>81// CHECK: %[[C_PACKED:.+]] = linalg.pack %[[C]]82// CHECK-SAME: padding_value(%[[ZERO]] : f32)83// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]84// CHECK-SAME: into %[[PACK_DST_2]] : tensor<?x?xf32> -> tensor<?x?x32x16xf32>85// CHECK: %[[GEMM_RES_PACKED:.+]] = linalg.generic86// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]87// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]88// CHECK-SAME: ins(%[[A_PACKED]], %[[B_PACKED]] : tensor<?x?x32x64xf32>, tensor<?x?x16x64xf32>) outs(%[[C_PACKED]] : tensor<?x?x32x16xf32>)89// CHECK: %[[RES_UNPACKED:.+]] = linalg.unpack %[[GEMM_RES_PACKED]]90// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]91// CHECK-SAME: into %[[C]] : tensor<?x?x32x16xf32> -> tensor<?x?xf32>92// CHECK: return %[[RES_UNPACKED]] : tensor<?x?xf32>93 94// -----95 96func.func @block_matmul_with_constant(97 %A: tensor<128x128xf32>, %B: tensor<128x128xf32>) -> tensor<128x128xf32> {98 %cst_acc = arith.constant dense<0.0> : tensor<128x128xf32>99 %0 = linalg.matmul ins(%A, %B : tensor<128x128xf32>, tensor<128x128xf32>)100 outs(%cst_acc : tensor<128x128xf32>) -> tensor<128x128xf32>101 return %0 : tensor<128x128xf32>102}103 104// CHECK-LABEL: func @block_matmul_with_constant(105// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<128x128xf32>, %[[B:[0-9a-z]+]]: tensor<128x128xf32>106// CHECK-DAG: %[[CST_ACC_PACKED:.+]] = arith.constant dense<0.000000e+00> : tensor<4x8x32x16xf32>107// CHECK-DAG: %[[RES_DST:.+]] = arith.constant dense<0.000000e+00> : tensor<128x128xf32>108// CHECK: %[[GEMM_RES_PACKED:.+]] = linalg.generic109// CHECK-SAME: ins({{.*}} : tensor<4x2x32x64xf32>, tensor<8x2x16x64xf32>) outs(%[[CST_ACC_PACKED]] : tensor<4x8x32x16xf32>)110// CHECK: %[[RES_UNPACKED:.+]] = linalg.unpack %[[GEMM_RES_PACKED]]111// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]112// CHECK-SAME: into %[[RES_DST]] : tensor<4x8x32x16xf32> -> tensor<128x128xf32>113// CHECK: return %[[RES_UNPACKED]] : tensor<128x128xf32>114 115// -----116 117func.func @block_matmul_with_producer(118 %A: tensor<128x128xf32>, %B: tensor<128x128xf32>, %C: tensor<128x128xf32>) -> tensor<128x128xf32> {119 %cst = arith.constant 0.0 : f32120 %acc = linalg.fill ins(%cst : f32) outs(%C : tensor<128x128xf32>) -> tensor<128x128xf32>121 %1 = linalg.matmul ins(%A, %B : tensor<128x128xf32>, tensor<128x128xf32>)122 outs(%acc : tensor<128x128xf32>) -> tensor<128x128xf32>123 return %1 : tensor<128x128xf32>124}125 126// CHECK-LABEL: func @block_matmul_with_producer(127// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<128x128xf32>, %[[B:[0-9a-z]+]]: tensor<128x128xf32>, %[[C:[0-9a-z]+]]: tensor<128x128xf32>128// CHECK-DAG: %[[C0:.+]] = arith.constant 0.000000e+00 : f32129// CHECK: %[[FILL_DST_PACKED:.+]] = tensor.empty() : tensor<4x8x32x16xf32>130// CHECK: %[[ACC_PACKED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[FILL_DST_PACKED]] : tensor<4x8x32x16xf32>) -> tensor<4x8x32x16xf32>131// CHECK: %[[GEMM_RES_PACKED:.+]] = linalg.generic132// CHECK-SAME: ins({{.*}} : tensor<4x2x32x64xf32>, tensor<8x2x16x64xf32>) outs(%[[ACC_PACKED]] : tensor<4x8x32x16xf32>)133// CHECK: %[[RES_UNPACKED:.+]] = linalg.unpack %[[GEMM_RES_PACKED]]134// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]135// CHECK-SAME: into %[[C]] : tensor<4x8x32x16xf32> -> tensor<128x128xf32>136// CHECK: return %[[RES_UNPACKED]] : tensor<128x128xf32>137 138// -----139 140func.func @block_matmul_with_consumer(141 %A: tensor<128x128xf32>, %B: tensor<128x128xf32>, %C: tensor<128x128xf32>, %D: tensor<128x128xf32>) -> tensor<128x128xf32> {142 %0 = tensor.empty() : tensor<128x128xf32>143 %1 = linalg.matmul ins(%A, %B : tensor<128x128xf32>, tensor<128x128xf32>)144 outs(%C : tensor<128x128xf32>) -> tensor<128x128xf32>145 %2 = linalg.add ins(%1, %D : tensor<128x128xf32>, tensor<128x128xf32>)146 outs(%0 : tensor<128x128xf32>) -> tensor<128x128xf32>147 return %2 : tensor<128x128xf32>148}149 150// CHECK-LABEL: func @block_matmul_with_consumer(151// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<128x128xf32>, %[[B:[0-9a-z]+]]: tensor<128x128xf32>, %[[C:[0-9a-z]+]]: tensor<128x128xf32>, %[[D:[0-9a-z]+]]: tensor<128x128xf32>152// CHECK-DAG: %[[RES_DST:.+]] = tensor.empty() : tensor<128x128xf32>153// CHECK: %[[GEMM_RES_PACKED:.+]] = linalg.generic154// CHECK-SAME: outs({{.*}} : tensor<4x8x32x16xf32>)155// CHECK: %[[RES_UNPACKED:.+]] = linalg.unpack %[[GEMM_RES_PACKED]]156// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]157// CHECK-SAME: into %[[C]] : tensor<4x8x32x16xf32> -> tensor<128x128xf32>158// CHECK: %[[ADD_RES:.+]] = linalg.add159// CHECK-SAME: ins(%[[RES_UNPACKED]], %[[D]] : tensor<128x128xf32>, tensor<128x128xf32>) outs(%[[RES_DST]] : tensor<128x128xf32>)160// CHECK: return %[[ADD_RES]] : tensor<128x128xf32>161 162// -----163 164func.func @block_batch_matmul(165 %A: tensor<512x64x128xf32>, %B: tensor<512x128x64xf32>, %C: tensor<512x64x64xf32>) -> tensor<512x64x64xf32> {166 %0 = linalg.batch_matmul ins(%A, %B : tensor<512x64x128xf32>, tensor<512x128x64xf32>)167 outs(%C : tensor<512x64x64xf32>) -> tensor<512x64x64xf32>168 return %0 : tensor<512x64x64xf32>169}170 171// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d3, d4, d6)>172// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2, d3, d5, d6)>173// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d4, d5)>174 175// CHECK-LABEL: func @block_batch_matmul(176// CHECK-SAME: %[[A:.+]]: tensor<512x64x128xf32>, %[[B:.+]]: tensor<512x128x64xf32>, %[[C:.+]]: tensor<512x64x64xf32>177// CHECK: %[[PACK_DST_0:.+]] = tensor.empty() : tensor<512x2x2x32x64xf32>178// CHECK: %[[A_PACKED:.+]] = linalg.pack %[[A]]179// CHECK-SAME: outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [32, 64]180// CHECK-SAME: into %[[PACK_DST_0]] : tensor<512x64x128xf32> -> tensor<512x2x2x32x64xf32>181// CHECK: %[[PACK_DST_1:.+]] = tensor.empty() : tensor<512x4x2x16x64xf32>182// CHECK: %[[B_PACKED:.+]] = linalg.pack %[[B]]183// CHECK-SAME: outer_dims_perm = [0, 2, 1] inner_dims_pos = [2, 1] inner_tiles = [16, 64]184// CHECK-SAME: into %[[PACK_DST_1]] : tensor<512x128x64xf32> -> tensor<512x4x2x16x64xf32>185// CHECK: %[[PACK_DST_2:.+]] = tensor.empty() : tensor<512x2x4x32x16xf32>186// CHECK: %[[C_PACKED:.+]] = linalg.pack %[[C]]187// CHECK-SAME: inner_dims_pos = [1, 2] inner_tiles = [32, 16]188// CHECK-SAME: into %[[PACK_DST_2]] : tensor<512x64x64xf32> -> tensor<512x2x4x32x16xf32>189// CHECK: %[[GEMM_RES_PACKED:.+]] = linalg.generic190// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]191// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]192// CHECK-SAME: ins(%[[A_PACKED]], %[[B_PACKED]] : tensor<512x2x2x32x64xf32>, tensor<512x4x2x16x64xf32>) outs(%[[C_PACKED]] : tensor<512x2x4x32x16xf32>)193// CHECK: %[[RES_UNPACKED:.+]] = linalg.unpack %[[GEMM_RES_PACKED]]194// CHECK-SAME: inner_dims_pos = [1, 2] inner_tiles = [32, 16]195// CHECK-SAME: into %[[C]] : tensor<512x2x4x32x16xf32> -> tensor<512x64x64xf32>196// CHECK: return %[[RES_UNPACKED]] : tensor<512x64x64xf32>197 198// -----199 200#map = affine_map<(d0, d1, d2) -> (d0, d2)>201#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>202#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>203 204func.func @block_generic_matmul(205 %A: tensor<128x128xf32>, %B: tensor<128x128xf32>, %C: tensor<128x128xf32>) -> tensor<128x128xf32> {206 %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]}207 ins(%A, %B : tensor<128x128xf32>, tensor<128x128xf32>)208 outs(%C : tensor<128x128xf32>) {209 ^bb0(%in: f32, %in_0: f32, %out: f32):210 %1 = arith.mulf %in, %in_0 : f32211 %2 = arith.addf %out, %1 : f32212 linalg.yield %2 : f32213 } -> tensor<128x128xf32>214 return %0 : tensor<128x128xf32>215}216 217// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>218// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d4, d5)>219// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)>220 221// CHECK-LABEL: func @block_generic_matmul(222// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<128x128xf32>, %[[B:[0-9a-z]+]]: tensor<128x128xf32>, %[[C:[0-9a-z]+]]: tensor<128x128xf32>223// CHECK: %[[PACK_DST_0:.+]] = tensor.empty() : tensor<4x2x32x64xf32>224// CHECK: %[[A_PACKED:.+]] = linalg.pack %[[A]]225// CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [32, 64]226// CHECK-SAME: into %[[PACK_DST_0]] : tensor<128x128xf32> -> tensor<4x2x32x64xf32>227// CHECK: %[[PACK_DST_1:.+]] = tensor.empty() : tensor<8x2x16x64xf32>228// CHECK: %[[B_PACKED:.+]] = linalg.pack %[[B]]229// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 64]230// CHECK-SAME: into %[[PACK_DST_1]] : tensor<128x128xf32> -> tensor<8x2x16x64xf32>231// CHECK: %[[PACK_DST_2:.+]] = tensor.empty() : tensor<4x8x32x16xf32>232// CHECK: %[[C_PACKED:.+]] = linalg.pack %[[C]]233// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]234// CHECK-SAME: into %[[PACK_DST_2]] : tensor<128x128xf32> -> tensor<4x8x32x16xf32>235// CHECK: %[[GEMM_RES_PACKED:.+]] = linalg.generic236// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]237// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]238// CHECK-SAME: ins(%[[A_PACKED]], %[[B_PACKED]] : tensor<4x2x32x64xf32>, tensor<8x2x16x64xf32>) outs(%[[C_PACKED]] : tensor<4x8x32x16xf32>)239// CHECK: %[[RES_UNPACKED:.+]] = linalg.unpack %[[GEMM_RES_PACKED]]240// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]241// CHECK-SAME: into %[[C]] : tensor<4x8x32x16xf32> -> tensor<128x128xf32>242// CHECK: return %[[RES_UNPACKED]] : tensor<128x128xf32>243 244// -----245 246#map = affine_map<(d0, d1, d2) -> (d2, d0)>247#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>248#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>249 250func.func @block_generic_matmul_transpose_a(251 %A: tensor<128x64xf32>, %B: tensor<128x64xf32>, %C: tensor<64x64xf32>) -> tensor<64x64xf32> {252 %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]}253 ins(%A, %B : tensor<128x64xf32>, tensor<128x64xf32>)254 outs(%C : tensor<64x64xf32>) {255 ^bb0(%in: f32, %in_0: f32, %out: f32):256 %1 = arith.mulf %in, %in_0 : f32257 %2 = arith.addf %out, %1 : f32258 linalg.yield %2 : f32259 } -> tensor<64x64xf32>260 return %0 : tensor<64x64xf32>261}262 263// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>264// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d4, d5)>265// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)>266 267// CHECK-LABEL: func @block_generic_matmul_transpose_a(268// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<128x64xf32>, %[[B:[0-9a-z]+]]: tensor<128x64xf32>, %[[C:[0-9a-z]+]]: tensor<64x64xf32>269// CHECK: %[[PACK_DST_0:.+]] = tensor.empty() : tensor<2x2x32x64xf32>270// CHECK: %[[A_PACKED:.+]] = linalg.pack %[[A]]271// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [32, 64]272// CHECK-SAME: into %[[PACK_DST_0]] : tensor<128x64xf32> -> tensor<2x2x32x64xf32>273// CHECK: %[[PACK_DST_1:.+]] = tensor.empty() : tensor<4x2x16x64xf32>274// CHECK: %[[B_PACKED:.+]] = linalg.pack %[[B]]275// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 64]276// CHECK-SAME: into %[[PACK_DST_1]] : tensor<128x64xf32> -> tensor<4x2x16x64xf32>277// CHECK: %[[PACK_DST_2:.+]] = tensor.empty() : tensor<2x4x32x16xf32>278// CHECK: %[[C_PACKED:.+]] = linalg.pack %[[C]]279// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]280// CHECK-SAME: into %[[PACK_DST_2]] : tensor<64x64xf32> -> tensor<2x4x32x16xf32>281// CHECK: %[[GEMM_RES_PACKED:.+]] = linalg.generic282// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]283// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]284// CHECK-SAME: ins(%[[A_PACKED]], %[[B_PACKED]] : tensor<2x2x32x64xf32>, tensor<4x2x16x64xf32>) outs(%[[C_PACKED]] : tensor<2x4x32x16xf32>)285// CHECK: %[[RES_UNPACKED:.+]] = linalg.unpack %[[GEMM_RES_PACKED]]286// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]287// CHECK-SAME: into %[[C]] : tensor<2x4x32x16xf32> -> tensor<64x64xf32>288// CHECK: return %[[RES_UNPACKED]] : tensor<64x64xf32>289 290// -----291 292#map = affine_map<(d0, d1, d2) -> (d0, d2)>293#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>294#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>295 296func.func @block_generic_matmul_transpose_b(297 %A: tensor<64x128xf32>, %B: tensor<64x128xf32>, %C: tensor<64x64xf32>) -> tensor<64x64xf32> {298 %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]}299 ins(%A, %B : tensor<64x128xf32>, tensor<64x128xf32>)300 outs(%C : tensor<64x64xf32>) {301 ^bb0(%in: f32, %in_0: f32, %out: f32):302 %1 = arith.mulf %in, %in_0 : f32303 %2 = arith.addf %out, %1 : f32304 linalg.yield %2 : f32305 } -> tensor<64x64xf32>306 return %0 : tensor<64x64xf32>307}308 309// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)>310// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d4, d5)>311// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)>312 313// CHECK-LABEL: func @block_generic_matmul_transpose_b(314// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<64x128xf32>, %[[B:[0-9a-z]+]]: tensor<64x128xf32>, %[[C:[0-9a-z]+]]: tensor<64x64xf32>315// CHECK: %[[PACK_DST_0:.+]] = tensor.empty() : tensor<2x2x32x64xf32>316// CHECK: %[[A_PACKED:.+]] = linalg.pack %[[A]]317// CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [32, 64]318// CHECK-SAME: into %[[PACK_DST_0]] : tensor<64x128xf32> -> tensor<2x2x32x64xf32>319// CHECK: %[[PACK_DST_1:.+]] = tensor.empty() : tensor<4x2x16x64xf32>320// CHECK: %[[B_PACKED:.+]] = linalg.pack %[[B]]321// CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [16, 64]322// CHECK-SAME: into %[[PACK_DST_1]] : tensor<64x128xf32> -> tensor<4x2x16x64xf32>323// CHECK: %[[PACK_DST_2:.+]] = tensor.empty() : tensor<2x4x32x16xf32>324// CHECK: %[[C_PACKED:.+]] = linalg.pack %[[C]]325// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]326// CHECK-SAME: into %[[PACK_DST_2]] : tensor<64x64xf32> -> tensor<2x4x32x16xf32>327// CHECK: %[[GEMM_RES_PACKED:.+]] = linalg.generic328// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]329// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"]330// CHECK-SAME: ins(%[[A_PACKED]], %[[B_PACKED]] : tensor<2x2x32x64xf32>, tensor<4x2x16x64xf32>) outs(%[[C_PACKED]] : tensor<2x4x32x16xf32>)331// CHECK: %[[RES_UNPACKED:.+]] = linalg.unpack %[[GEMM_RES_PACKED]]332// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [32, 16]333// CHECK-SAME: into %[[C]] : tensor<2x4x32x16xf32> -> tensor<64x64xf32>334// CHECK: return %[[RES_UNPACKED]] : tensor<64x64xf32>335 336// -----337 338#map = affine_map<(d0, d1) -> (d0, d1)>339 340func.func @non_contraction_generic(341 %A: tensor<64x128xf32>) -> tensor<64x128xf32> {342 %c0 = arith.constant 0.000000e+00 : f32343 %0 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel"]}344 outs(%A : tensor<64x128xf32>) {345 ^bb0(%out: f32):346 %1 = arith.maximumf %out, %c0 : f32347 linalg.yield %1 : f32348 } -> tensor<64x128xf32>349 return %0 : tensor<64x128xf32>350}351 352// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1) -> (d0, d1)>353 354// CHECK-LABEL: func @non_contraction_generic(355// CHECK-SAME: %[[A:[0-9a-z]+]]: tensor<64x128xf32>356// CHECK-DAG: %[[C0:.+]] = arith.constant 0.000000e+00 : f32357// CHECK-NOT: linalg.pack358// CHECK: %[[GENERIC:.+]] = linalg.generic359// CHECK-SAME: indexing_maps = [#[[$MAP]]]360// CHECK-SAME: iterator_types = ["parallel", "parallel"]361// CHECK-SAME: outs(%[[A]] : tensor<64x128xf32>)362// CHECK-NOT: linalg.unpack363// CHECK: return %[[GENERIC]] : tensor<64x128xf32>364