430 lines · plain
1// RUN: mlir-opt %s -test-vector-reduction-to-contract-patterns -split-input-file | FileCheck %s2 3// TODO: Separate tests for vector.multi_reduction -> vector.contract and4// * pre-op + vector.contract -> vector.contract,5// * vector.contract + post-op -> vector.contract.6 7// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>8// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>9 10// CHECK-LABEL: multidimreduction_contract11// CHECK-SAME: (%[[ARG0:.*]]: vector<8x32x16xf32>, %[[ARG1:.*]]: vector<8x32x16xf32>, %[[ARG2:.*]]: vector<8x16xf32>)12// CHECK-NEXT: %[[R:.+]] = vector.contract {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]],13// CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel"], kind = #vector.kind<add>}14// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x16xf32>15// CHECK-NEXT: return %[[R]] : vector<8x16xf32>16func.func @multidimreduction_contract(17 %arg0: vector<8x32x16xf32>,%arg1: vector<8x32x16xf32>, %acc: vector<8x16xf32>) -> vector<8x16xf32> {18 %0 = arith.mulf %arg0, %arg1 : vector<8x32x16xf32>19 %1 = vector.multi_reduction <add>, %0, %acc [1] : vector<8x32x16xf32> to vector<8x16xf32>20 return %1 : vector<8x16xf32> }21 22// -----23 24// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>25// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>26 27// CHECK-LABEL: multidimreduction_contract_int28// CHECK-SAME: (%[[ARG0:.*]]: vector<8x32x16xi32>, %[[ARG1:.*]]: vector<8x32x16xi32>, %[[ARG2:.*]]: vector<8x16xi32>)29// CHECK-NEXT: %[[R:.+]] = vector.contract {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP1]]],30// CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel"], kind = #vector.kind<add>}31// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<8x32x16xi32>, vector<8x32x16xi32> into vector<8x16xi32>32// CHECK-NEXT: return %[[R]] : vector<8x16xi32>33func.func @multidimreduction_contract_int(34 %arg0: vector<8x32x16xi32>,%arg1: vector<8x32x16xi32>, %acc: vector<8x16xi32>) -> vector<8x16xi32> {35 %0 = arith.muli %arg0, %arg1 : vector<8x32x16xi32>36 %1 = vector.multi_reduction <add>, %0, %acc [1] : vector<8x32x16xi32> to vector<8x16xi32>37 return %1 : vector<8x16xi32>38}39 40// -----41 42//-----------------------------------------------------------------------------43// [Pattern: CombineContractABTranspose]44//-----------------------------------------------------------------------------45 46#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>47#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>48 49// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)>50// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>51// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>52 53// CHECK-LABEL: contract_transpose54// CHECK-SAME: (%[[ARG0:.+]]: vector<32x16x8xf32>,55// CHECK-NEXT: %[[C0:.+]] = arith.constant dense<0.000000e+00> : vector<8x32xf32>56// CHECK-NEXT: %[[R:.+]] = vector.contract {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]],57// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}58// CHECK-SAME: %[[ARG0]], %{{.*}}, %[[C0]] : vector<32x16x8xf32>, vector<8x32x16xf32> into vector<8x32xf32>59// CHECK-NEXT: return %[[R]] : vector<8x32xf32>60func.func @contract_transpose(61 %arg0: vector<32x16x8xf32>, %arg1: vector<8x32x16xf32>) -> vector<8x32xf32> {62 %cst = arith.constant dense<0.000000e+00> : vector<8x32xf32>63 %0 = vector.transpose %arg0, [2, 0, 1] : vector<32x16x8xf32> to vector<8x32x16xf32>64 %1 = vector.contract {indexing_maps = [#map0, #map0, #map1],65 iterator_types = ["parallel", "parallel", "reduction"],66 kind = #vector.kind<add>} %0, %arg1, %cst : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>67 return %1 : vector<8x32xf32>68}69 70// -----71 72//-----------------------------------------------------------------------------73// [Pattern: CombineContractBroadcast]74//-----------------------------------------------------------------------------75 76#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>77#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>78 79// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>80// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>81// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>82 83// CHECK-LABEL: contract_broadcast84// CHECK-SAME: (%[[ARG0:.+]]: vector<32x16xf32>,85// CHECK-NEXT: %[[C0:.+]] = arith.constant dense<0.000000e+00> : vector<8x32xf32>86// CHECK-NEXT: %[[R:.+]] = vector.contract {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]],87// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}88// CHECK-SAME: %[[ARG0]], %{{.*}}, %[[C0]] : vector<32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>89// CHECK-NEXT: return %[[R]] : vector<8x32xf32>90func.func @contract_broadcast(91 %arg0: vector<32x16xf32>, %arg1: vector<8x32x16xf32>) -> vector<8x32xf32> {92 %cst = arith.constant dense<0.000000e+00> : vector<8x32xf32>93 %0 = vector.broadcast %arg0 : vector<32x16xf32> to vector<8x32x16xf32>94 %1 = vector.contract {indexing_maps = [#map0, #map0, #map1],95 iterator_types = ["parallel", "parallel", "reduction"],96 kind = #vector.kind<add>} %0, %arg1, %cst : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>97 return %1 : vector<8x32xf32>98}99 100// -----101 102// Same as above, but with a mask.103 104#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>105#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>106 107// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>108// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>109// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>110 111// CHECK-LABEL: contract_broadcast_masked112// CHECK-SAME: %[[ARG0:.*]]: vector<32x16xf32>,113// CHECK-SAME: %[[ARG1:.*]]: vector<8x32x16xf32>,114// CHECK-SAME: %[[MASK:.*]]: vector<8x32x16xi1>) -> vector<8x32xf32> {115// CHECK: %[[C0:.*]] = arith.constant dense<0.000000e+00> : vector<8x32xf32>116// CHECK: %[[R:.*]] = vector.mask %[[MASK]] {117// CHECK-SAME: vector.contract {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]],118// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"],119// CHECK-SAME: kind = #vector.kind<add>}120// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[C0]] : vector<32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>121// CHECK-SAME } : vector<8x32x16xi1> -> vector<8x32xf32>122// CHECK: return %[[R]] : vector<8x32xf32>123func.func @contract_broadcast_masked(124 %arg0: vector<32x16xf32>, %arg1: vector<8x32x16xf32>, %mask: vector<8x32x16xi1>) -> vector<8x32xf32> {125 %cst = arith.constant dense<0.000000e+00> : vector<8x32xf32>126 %0 = vector.broadcast %arg0 : vector<32x16xf32> to vector<8x32x16xf32>127 %1 = vector.mask %mask {128 vector.contract {indexing_maps = [#map0, #map0, #map1],129 iterator_types = ["parallel", "parallel", "reduction"],130 kind = #vector.kind<add>131 } %0, %arg1, %cst : vector<8x32x16xf32>, vector<8x32x16xf32> into vector<8x32xf32>132 } : vector<8x32x16xi1> -> vector<8x32xf32>133 return %1 : vector<8x32xf32>134}135 136// -----137 138// Same as above, but with a scalable dim.139 140#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>141#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>142 143// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>144// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>145// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>146 147// CHECK-LABEL: contract_broadcast_masked_scalable148// CHECK-SAME: %[[ARG0:.*]]: vector<[32]x16xf32>,149// CHECK-SAME: %[[ARG1:.*]]: vector<8x[32]x16xf32>,150// CHECK-SAME: %[[MASK:.*]]: vector<8x[32]x16xi1>) -> vector<8x32xf32> {151// CHECK: %[[C0:.*]] = arith.constant dense<0.000000e+00> : vector<8x32xf32>152// CHECK: %[[R:.*]] = vector.mask %[[MASK]] {153// CHECK-SAME: vector.contract {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]],154// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"],155// CHECK-SAME: kind = #vector.kind<add>}156// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[C0]] : vector<[32]x16xf32>, vector<8x[32]x16xf32> into vector<8x32xf32>157// CHECK-SAME } : vector<8x[32]x16xi1> -> vector<8x32xf32>158// CHECK: return %[[R]] : vector<8x32xf32>159func.func @contract_broadcast_masked_scalable(160 %arg0: vector<[32]x16xf32>, %arg1: vector<8x[32]x16xf32>, %mask: vector<8x[32]x16xi1>) -> vector<8x32xf32> {161 %cst = arith.constant dense<0.000000e+00> : vector<8x32xf32>162 %0 = vector.broadcast %arg0 : vector<[32]x16xf32> to vector<8x[32]x16xf32>163 %1 = vector.mask %mask {164 vector.contract {indexing_maps = [#map0, #map0, #map1],165 iterator_types = ["parallel", "parallel", "reduction"],166 kind = #vector.kind<add>167 } %0, %arg1, %cst : vector<8x[32]x16xf32>, vector<8x[32]x16xf32> into vector<8x32xf32>168 } : vector<8x[32]x16xi1> -> vector<8x32xf32>169 return %1 : vector<8x32xf32>170}171 172// -----173 174// Test that CombineContractBroadcast is able to combine a broadcast that175// creates a unit dim that is consumed by a reduction iterator, dropping that176// reduction iterator, as long as there is another reduction iterator left.177 178#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>179#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>180#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>181 182// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>183// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>184// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>185 186// CHECK-LABEL: contract_broadcast_unit_dim_reduction187// CHECK-SAME: (%[[ARG0:.+]]: vector<8x4xi32>, %[[ARG1:.+]]: vector<8x4xi32>, %[[ARG2:.+]]: vector<8x8xi32>)188// CHECK: vector.contract189// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]190// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]191// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<8x4xi32>, vector<8x4xi32> into vector<8x8xi32>192func.func @contract_broadcast_unit_dim_reduction(%arg0 : vector<8x4xi32>, %arg1 : vector<8x4xi32>, %arg2 : vector<8x8xi32>) -> vector<8x8xi32> {193 %0 = vector.broadcast %arg0 : vector<8x4xi32> to vector<1x8x4xi32>194 %1 = vector.broadcast %arg1 : vector<8x4xi32> to vector<1x8x4xi32>195 %result = vector.contract {196 indexing_maps = [#map0, #map1, #map2],197 iterator_types = ["reduction", "parallel", "parallel", "reduction"],198 kind = #vector.kind<add>199 } %0, %1, %arg2 : vector<1x8x4xi32>, vector<1x8x4xi32> into vector<8x8xi32>200 return %result : vector<8x8xi32>201}202 203// -----204 205// Same as above, but with a mask.206 207#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>208#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>209#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>210 211// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>212// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>213// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>214 215// CHECK-LABEL: contract_broadcast_unit_dim_reduction_masked216// CHECK-SAME: (%[[ARG0:.+]]: vector<8x4xi32>, %[[ARG1:.+]]: vector<8x4xi32>, %[[ARG2:.+]]: vector<8x8xi32>, %[[MASK:.+]]: vector<1x8x8x4xi1>)217// CHECK: %[[MASK_SC:.*]] = vector.shape_cast %[[MASK]] : vector<1x8x8x4xi1> to vector<8x8x4xi1>218// CHECK: %[[R:.*]] = vector.mask %[[MASK_SC]] {219// CHECK-SAME: vector.contract220// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]221// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]222// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<8x4xi32>, vector<8x4xi32> into vector<8x8xi32>223func.func @contract_broadcast_unit_dim_reduction_masked(%arg0 : vector<8x4xi32>, %arg1 : vector<8x4xi32>, %arg2 : vector<8x8xi32>, %mask: vector<1x8x8x4xi1>) -> vector<8x8xi32> {224 %0 = vector.broadcast %arg0 : vector<8x4xi32> to vector<1x8x4xi32>225 %1 = vector.broadcast %arg1 : vector<8x4xi32> to vector<1x8x4xi32>226 %result = vector.mask %mask {227 vector.contract {228 indexing_maps = [#map0, #map1, #map2],229 iterator_types = ["reduction", "parallel", "parallel", "reduction"],230 kind = #vector.kind<add>231 } %0, %1, %arg2 : vector<1x8x4xi32>, vector<1x8x4xi32> into vector<8x8xi32>232 } : vector<1x8x8x4xi1> -> vector<8x8xi32>233 return %result : vector<8x8xi32>234}235 236// -----237 238// Same as above, but with a scalable dim.239 240#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>241#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>242#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>243 244// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>245// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>246// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>247 248// CHECK-LABEL: contract_broadcast_unit_dim_reduction_masked_scalable249// CHECK-SAME: (%[[ARG0:.+]]: vector<8x4xi32>, %[[ARG1:.+]]: vector<[8]x4xi32>, %[[ARG2:.+]]: vector<8x[8]xi32>, %[[MASK:.+]]: vector<1x8x[8]x4xi1>)250// CHECK: %[[MASK_SC:.*]] = vector.shape_cast %[[MASK]] : vector<1x8x[8]x4xi1> to vector<8x[8]x4xi1>251// CHECK: %[[R:.*]] = vector.mask %[[MASK_SC]] {252// CHECK-SAME: vector.contract253// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]254// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]255// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<8x4xi32>, vector<[8]x4xi32> into vector<8x[8]xi32>256func.func @contract_broadcast_unit_dim_reduction_masked_scalable(%arg0 : vector<8x4xi32>, %arg1 : vector<[8]x4xi32>, %arg2 : vector<8x[8]xi32>, %mask: vector<1x8x[8]x4xi1>) -> vector<8x[8]xi32> {257 %0 = vector.broadcast %arg0 : vector<8x4xi32> to vector<1x8x4xi32>258 %1 = vector.broadcast %arg1 : vector<[8]x4xi32> to vector<1x[8]x4xi32>259 %result = vector.mask %mask {260 vector.contract {261 indexing_maps = [#map0, #map1, #map2],262 iterator_types = ["reduction", "parallel", "parallel", "reduction"],263 kind = #vector.kind<add>264 } %0, %1, %arg2 : vector<1x8x4xi32>, vector<1x[8]x4xi32> into vector<8x[8]xi32>265 } : vector<1x8x[8]x4xi1> -> vector<8x[8]xi32>266 return %result : vector<8x[8]xi32>267}268 269// -----270// Test that CombineContractBroadcast will not combine a broadcast that creates271// a non-unit dim that is consumed by a reduction iterator.272// Moreover, the affine_map's are permuting the position of that reduction273// iterator with the position of a parallel iterator to ensure that274// the logic guarding that case does not mix up dimensions here.275 276#map0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>277#map1 = affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>278#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d2)>279 280// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>281// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)>282// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2)>283 284// CHECK-LABEL: contract_broadcast_non_unit_dim_reduction_with_permutation285// CHECK-SAME: (%[[ARG0:.+]]: vector<8x4xi32>, %[[ARG1:.+]]: vector<8x4xi32>, %[[ARG2:.+]]: vector<8x8xi32>)286// CHECK: %[[BROADCAST0:.+]] = vector.broadcast %[[ARG0]] : vector<8x4xi32> to vector<2x8x4xi32>287// CHECK: %[[BROADCAST1:.+]] = vector.broadcast %[[ARG1]] : vector<8x4xi32> to vector<2x8x4xi32>288// CHECK: vector.contract289// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]290// CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel", "reduction"]291// CHECK-SAME: %[[BROADCAST0]], %[[BROADCAST1]], %[[ARG2]] : vector<2x8x4xi32>, vector<2x8x4xi32> into vector<8x8xi32>292func.func @contract_broadcast_non_unit_dim_reduction_with_permutation(%arg0 : vector<8x4xi32>, %arg1 : vector<8x4xi32>, %arg2 : vector<8x8xi32>) -> vector<8x8xi32> {293 %0 = vector.broadcast %arg0 : vector<8x4xi32> to vector<2x8x4xi32>294 %1 = vector.broadcast %arg1 : vector<8x4xi32> to vector<2x8x4xi32>295 %result = vector.contract {296 indexing_maps = [#map0, #map1, #map2],297 iterator_types = ["parallel", "reduction", "parallel", "reduction"],298 kind = #vector.kind<add>299 } %0, %1, %arg2 : vector<2x8x4xi32>, vector<2x8x4xi32> into vector<8x8xi32>300 return %result : vector<8x8xi32>301}302 303// -----304 305// Test that CombineContractBroadcast is not combining this case, as that would306// result in dropping this contract's only reduction iterator.307 308#map0 = affine_map<(d0, d1, d2) -> (d0, d1)>309#map1 = affine_map<(d0, d1, d2) -> (d0, d2)>310#map2 = affine_map<(d0, d1, d2) -> (d1, d2)>311 312// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>313// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>314// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>315 316// CHECK-LABEL: contract_broadcast_unit_dim_reduction_as_only_reduction317// CHECK-SAME: (%[[ARG0:.+]]: vector<8xi32>, %[[ARG1:.+]]: vector<8xi32>, %[[ARG2:.+]]: vector<8x8xi32>)318// CHECK: %[[BROADCAST0:.+]] = vector.broadcast %[[ARG0]] : vector<8xi32> to vector<1x8xi32>319// CHECK: %[[BROADCAST1:.+]] = vector.broadcast %[[ARG1]] : vector<8xi32> to vector<1x8xi32>320// CHECK: vector.contract321// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]322// CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel"]323// CHECK-SAME: %[[BROADCAST0]], %[[BROADCAST1]], %[[ARG2]] : vector<1x8xi32>, vector<1x8xi32> into vector<8x8xi32>324func.func @contract_broadcast_unit_dim_reduction_as_only_reduction(%arg0 : vector<8xi32>, %arg1 : vector<8xi32>, %arg2 : vector<8x8xi32>) -> vector<8x8xi32> {325 %0 = vector.broadcast %arg0 : vector<8xi32> to vector<1x8xi32>326 %1 = vector.broadcast %arg1 : vector<8xi32> to vector<1x8xi32>327 %result = vector.contract {328 indexing_maps = [#map0, #map1, #map2],329 iterator_types = ["reduction", "parallel", "parallel"],330 kind = #vector.kind<add>331 } %0, %1, %arg2 : vector<1x8xi32>, vector<1x8xi32> into vector<8x8xi32>332 return %result : vector<8x8xi32>333}334 335// -----336 337// Test that CombineContractBroadcast is not combining this case, as that would338// result in a dimension being unused in the LHS and RHS maps, which is illegal.339 340#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>341#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>342#map2 = affine_map<(d0, d1, d2) -> (d1)>343 344// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>345// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>346// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d1)>347 348// CHECK-LABEL: contract_broadcast_dimension_would_go_unused_in_lhs_rhs349// CHECK-SAME: (%[[ARG0:.+]]: vector<1x2xi32>, %[[ARG1:.+]]: vector<2xi32>, %[[ARG2:.+]]: vector<1xi32>)350// CHECK: %[[BROADCAST1:.+]] = vector.broadcast %[[ARG1]] : vector<2xi32> to vector<1x1x2xi32>351// CHECK: vector.contract352// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]353// CHECK-SAME: iterator_types = ["reduction", "parallel", "reduction"]354// CHECK-SAME: %[[ARG0]], %[[BROADCAST1]], %[[ARG2]] : vector<1x2xi32>, vector<1x1x2xi32> into vector<1xi32>355 356func.func @contract_broadcast_dimension_would_go_unused_in_lhs_rhs(%arg0 : vector<1x2xi32>, %arg1 : vector<2xi32>, %arg2 : vector<1xi32>) -> vector<1xi32> {357 %1 = vector.broadcast %arg1 : vector<2xi32> to vector<1x1x2xi32>358 %result = vector.contract {359 indexing_maps = [#map0, #map1, #map2],360 iterator_types = ["reduction", "parallel", "reduction"],361 kind = #vector.kind<add>362 } %arg0, %1, %arg2 : vector<1x2xi32>, vector<1x1x2xi32> into vector<1xi32>363 return %result : vector<1xi32>364}365 366// -----367 368// Test that CombineContractBroadcast is not combining this case, as that would369// result in a vector.contract without a reduction dimention pair, as the only370// reduction dimension would be used by only one side among LHS, RHS.371// This is arguably a convoluted edge case (the affine_maps here look weird!)372// but it is something that we actually ran into from linalg.matmul tests that373// were exercising 1x1 shapes, and using various drop-unit-dims patterns.374 375#map0 = affine_map<(d0, d1) -> (d1)>376#map1 = affine_map<(d0, d1) -> (d1, d0)>377#map2 = affine_map<(d0, d1) -> (d0)>378 379// CHECK-LABEL: contract_broadcast_would_have_no_reduction_dim_pair380// CHECK-SAME: (%[[ARG0:.+]]: vector<1xf32>, %[[ARG1:.+]]: vector<1xf32>, %[[ARG2:.+]]: vector<1xf32>)381// CHECK: %[[BROADCAST1:.+]] = vector.broadcast %[[ARG1]] : vector<1xf32> to vector<1x1xf32>382// CHECK: vector.contract383// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]384// CHECK-SAME: iterator_types = ["parallel", "reduction"]385// CHECK-SAME: %[[ARG0]], %[[BROADCAST1]], %[[ARG2]] : vector<1xf32>, vector<1x1xf32> into vector<1xf32>386 387func.func @contract_broadcast_would_have_no_reduction_dim_pair(%arg0 : vector<1xf32>, %arg1 : vector<1xf32>, %arg2 : vector<1xf32>) -> vector<1xf32> {388 %1 = vector.broadcast %arg1 : vector<1xf32> to vector<1x1xf32>389 %result = vector.contract {390 indexing_maps = [#map0, #map1, #map2],391 iterator_types = ["parallel", "reduction"],392 kind = #vector.kind<add>393 } %arg0, %1, %arg2 : vector<1xf32>, vector<1x1xf32> into vector<1xf32>394 return %result : vector<1xf32>395}396 397 398// -----399 400//-----------------------------------------------------------------------------401// [Pattern: CombineContractResultTranspose]402//-----------------------------------------------------------------------------403 404// CHECK-DAG: #[[$LHS_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>405// CHECK-DAG: #[[$RHS_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>406// CHECK-DAG: #[[$ACC_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1)>407 408// CHECK-LABEL: func.func @contract_result_transpose409// CHECK-SAME: (%[[LHS:.+]]: vector<2x4x4xf32>, %[[RHS:.+]]: vector<4x8xf32>, %[[ACC:.+]]: vector<2x8x4xf32>)410// CHECK: %[[CONTRACT:.+]] = vector.contract411// CHECK-SAME: indexing_maps = [#[[$LHS_MAP]], #[[$RHS_MAP]], #[[$ACC_MAP]]]412// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"]413// CHECK-SAME: kind = #vector.kind<add>414// CHECK-SAME: %[[LHS]], %[[RHS]], %[[ACC]]415// CHECK: return %[[CONTRACT]]416func.func @contract_result_transpose(%lhs : vector<2x4x4xf32>, %rhs: vector<4x8xf32>, %acc: vector<2x8x4xf32>) -> vector<2x8x4xf32> {417 %accT = vector.transpose %acc, [0, 2, 1] : vector<2x8x4xf32> to vector<2x4x8xf32>418 %contract = vector.contract {419 indexing_maps = [420 affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>,421 affine_map<(d0, d1, d2, d3) -> (d3, d2)>,422 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>423 ],424 iterator_types = ["parallel", "parallel", "parallel", "reduction"],425 kind = #vector.kind<add>426 } %lhs, %rhs, %accT : vector<2x4x4xf32>, vector<4x8xf32> into vector<2x4x8xf32>427 %resT = vector.transpose %contract, [0, 2, 1] : vector<2x4x8xf32> to vector<2x8x4xf32>428 return %resT : vector<2x8x4xf32>429}430