345 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -cse -split-input-file | FileCheck %s2 3!vecA = vector<1x1xf32>4!vecB = vector<1x64xf32>5!vecC = vector<1x64xf32>6#map = affine_map<(d0, d1, d2) -> (d0, d2)>7#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>8#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>9func.func @matmul_outer_product_to_fma(10 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC11{12 %0 = vector.contract {13 indexing_maps = [#map, #map1, #map2],14 iterator_types = ["parallel", "parallel", "reduction"],15 kind = #vector.kind<add>}16 %arg0, %arg1, %arg217 : !vecA, !vecB into !vecC18 return %0 : !vecC19}20 21// CHECK-LABEL: @matmul_outer_product_to_fma22// CHECK: vector.broadcast{{.*}}vector<1xf32> to vector<64xf32>23// CHECK: vector.fma{{.*}}vector<64xf32>24// CHECK: vector.shape_cast{{.*}}vector<64xf32> to vector<1x64xf32>25 26module attributes {transform.with_named_sequence} {27 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {28 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op29 transform.apply_patterns to %func {30 transform.apply_patterns.x86vector.vector_contract_to_fma31 } : !transform.any_op32 transform.yield33 }34}35 36// -----37 38!vecA = vector<64x1xf32>39!vecB = vector<1x1xf32>40!vecC = vector<64x1xf32>41#map = affine_map<(d0, d1, d2) -> (d0, d2)>42#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>43#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>44func.func @matmul_outer_product_to_fma_bcst_B(45 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC46{47 %0 = vector.contract {48 indexing_maps = [#map, #map1, #map2],49 iterator_types = ["parallel", "parallel", "reduction"],50 kind = #vector.kind<add>}51 %arg0, %arg1, %arg252 : !vecA, !vecB into !vecC53 return %0 : !vecC54}55 56// CHECK-LABEL: @matmul_outer_product_to_fma_bcst_B57// CHECK: vector.broadcast58// CHECK: vector.fma{{.*}}vector<64xf32>59 60module attributes {transform.with_named_sequence} {61 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {62 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op63 transform.apply_patterns to %func {64 transform.apply_patterns.x86vector.vector_contract_to_fma65 } : !transform.any_op66 transform.yield67 }68}69 70// -----71 72!vecA = vector<1x1x1xf32>73!vecB = vector<1x1x64xf32>74!vecC = vector<1x1x64xf32>75#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>76#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>77#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>78func.func @batch_matmul_to_fma(79 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC80{81 %0 = vector.contract {82 indexing_maps = [#map, #map1, #map2],83 iterator_types = ["parallel", "parallel", "parallel", "reduction"],84 kind = #vector.kind<add>}85 %arg0, %arg1, %arg286 : !vecA, !vecB into !vecC87 return %0 : !vecC88}89 90// CHECK-LABEL: @batch_matmul_to_fma91// CHECK: vector.broadcast92// CHECK: vector.fma{{.*}}vector<64xf32>93 94module attributes {transform.with_named_sequence} {95 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {96 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op97 transform.apply_patterns to %func {98 transform.apply_patterns.x86vector.vector_contract_to_fma99 } : !transform.any_op100 transform.yield101 }102}103 104// -----105 106!vecA = vector<1x64x1xf32>107!vecB = vector<1x1x1xf32>108!vecC = vector<1x64x1xf32>109#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>110#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>111#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>112func.func @batch_matmul_to_fma_bcst_B(113 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC114{115 %0 = vector.contract {116 indexing_maps = [#map, #map1, #map2],117 iterator_types = ["parallel", "parallel", "parallel", "reduction"],118 kind = #vector.kind<add>}119 %arg0, %arg1, %arg2120 : !vecA, !vecB into !vecC121 return %0 : !vecC122}123 124// CHECK-LABEL: @batch_matmul_to_fma_bcst_B125// CHECK: vector.broadcast126// CHECK: vector.fma{{.*}}vector<64xf32>127 128module attributes {transform.with_named_sequence} {129 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {130 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op131 transform.apply_patterns to %func {132 transform.apply_patterns.x86vector.vector_contract_to_fma133 } : !transform.any_op134 transform.yield135 }136}137 138// -----139 140!vecA = vector<1x1x1xf32>141!vecB = vector<1x1x64xf32>142!vecC = vector<1x64xf32>143#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>144#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>145#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>146func.func @brgemm_to_fma(147 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC148{149 %0 = vector.contract {150 indexing_maps = [#map, #map1, #map2],151 iterator_types = ["reduction", "parallel", "parallel", "reduction"],152 kind = #vector.kind<add>}153 %arg0, %arg1, %arg2154 : !vecA, !vecB into !vecC155 return %0 : !vecC156}157 158// CHECK-LABEL: @brgemm_to_fma159// CHECK: vector.broadcast160// CHECK: vector.fma{{.*}}vector<64xf32>161 162module attributes {transform.with_named_sequence} {163 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {164 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op165 transform.apply_patterns to %func {166 transform.apply_patterns.x86vector.vector_contract_to_fma167 } : !transform.any_op168 transform.yield169 }170}171 172// -----173 174!vecA = vector<1x64x1xf32>175!vecB = vector<1x1x1xf32>176!vecC = vector<64x1xf32>177#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>178#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>179#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>180func.func @brgemm_to_fma_bcst_B(181 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC182{183 %0 = vector.contract {184 indexing_maps = [#map, #map1, #map2],185 iterator_types = ["reduction", "parallel", "parallel", "reduction"],186 kind = #vector.kind<add>}187 %arg0, %arg1, %arg2188 : !vecA, !vecB into !vecC189 return %0 : !vecC190}191 192// CHECK-LABEL: @brgemm_to_fma_bcst_B193// CHECK: vector.broadcast194// CHECK: vector.fma{{.*}}vector<64xf32>195 196module attributes {transform.with_named_sequence} {197 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {198 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op199 transform.apply_patterns to %func {200 transform.apply_patterns.x86vector.vector_contract_to_fma201 } : !transform.any_op202 transform.yield203 }204}205 206// -----207 208!vecA = vector<3x1x1xf32>209!vecB = vector<3x1x64xf32>210!vecC = vector<3x1x64xf32>211#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>212#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>213#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>214func.func @negative_non_unit_batch_dim(215 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC216{217 %0 = vector.contract {218 indexing_maps = [#map, #map1, #map2],219 iterator_types = ["parallel", "parallel", "parallel", "reduction"],220 kind = #vector.kind<add>}221 %arg0, %arg1, %arg2222 : !vecA, !vecB into !vecC223 return %0 : !vecC224}225 226// Batch dimension should've been simplified earlier.227 228// CHECK-LABEL: @negative_non_unit_batch_dim229// CHECK-NOT: vector.fma230// CHECK: vector.contract231 232module attributes {transform.with_named_sequence} {233 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {234 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op235 transform.apply_patterns to %func {236 transform.apply_patterns.x86vector.vector_contract_to_fma237 } : !transform.any_op238 transform.yield239 }240}241 242// -----243 244!vecA = vector<3x1x1xf32>245!vecB = vector<3x1x64xf32>246!vecC = vector<1x64xf32>247#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>248#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>249#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>250func.func @negative_non_unit_batch_reduce_dim(251 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC252{253 %0 = vector.contract {254 indexing_maps = [#map, #map1, #map2],255 iterator_types = ["reduction", "parallel", "parallel", "reduction"],256 kind = #vector.kind<add>}257 %arg0, %arg1, %arg2258 : !vecA, !vecB into !vecC259 return %0 : !vecC260}261 262// Batch-reduce dimension should've been simplified earlier.263 264// CHECK-LABEL: @negative_non_unit_batch_reduce_dim265// CHECK-NOT: vector.fma266// CHECK: vector.contract267 268module attributes {transform.with_named_sequence} {269 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {270 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op271 transform.apply_patterns to %func {272 transform.apply_patterns.x86vector.vector_contract_to_fma273 } : !transform.any_op274 transform.yield275 }276}277 278// -----279 280!vecA = vector<1x1xf32>281!vecB = vector<1x64xf32>282!vecC = vector<1x64xf32>283#map = affine_map<(d0, d1, d2) -> (d0, d2)>284#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>285#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>286func.func @negative_invalid_kind(287 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC288{289 %0 = vector.contract {290 indexing_maps = [#map, #map1, #map2],291 iterator_types = ["parallel", "parallel", "reduction"],292 kind = #vector.kind<mul>}293 %arg0, %arg1, %arg2294 : !vecA, !vecB into !vecC295 return %0 : !vecC296}297 298// CHECK-LABEL: @negative_invalid_kind299// CHECK-NOT: vector.fma300// CHECK: vector.contract301 302module attributes {transform.with_named_sequence} {303 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {304 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op305 transform.apply_patterns to %func {306 transform.apply_patterns.x86vector.vector_contract_to_fma307 } : !transform.any_op308 transform.yield309 }310}311 312// -----313 314!vecA = vector<1x1x1xf32>315!vecB = vector<1x1x64xf32>316!vecC = vector<1x1x64xi32>317#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>318#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>319#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>320func.func @negative_accumulator_type(321 %arg0: !vecA, %arg1: !vecB, %arg2: !vecC) -> !vecC322{323 %0 = vector.contract {324 indexing_maps = [#map, #map1, #map2],325 iterator_types = ["parallel", "parallel", "parallel", "reduction"],326 kind = #vector.kind<add>}327 %arg0, %arg1, %arg2328 : !vecA, !vecB into !vecC329 return %0 : !vecC330}331 332// CHECK-LABEL: @negative_accumulator_type333// CHECK-NOT: vector.fma334// CHECK: vector.contract335 336module attributes {transform.with_named_sequence} {337 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {338 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op339 transform.apply_patterns to %func {340 transform.apply_patterns.x86vector.vector_contract_to_fma341 } : !transform.any_op342 transform.yield343 }344}345