300 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -cse -split-input-file | FileCheck %s2 3!type = tensor<2048x2048xf32>4func.func @fold_add_on_two_matmuls(%arg0: !type, %arg1: !type) -> !type {5 %0 = arith.constant dense<1.111111e+00> : !type6 %cst = arith.constant 0.000000e+00 : f327 %1 = tensor.empty() : !type8 %2 = linalg.fill ins(%cst : f32) outs(%1 : !type) -> !type9 %3 = linalg.matmul ins(%arg0, %0 : !type, !type) outs(%2 : !type) -> !type10 %4 = tensor.empty() : !type11 %5 = linalg.fill ins(%cst : f32) outs(%1 : !type) -> !type12 %6 = linalg.matmul ins(%arg1, %0 : !type, !type) outs(%5 : !type) -> !type13 %7 = linalg.add ins(%3, %6 : !type, !type) outs(%1 : !type) -> !type14 return %7 : !type15}16 17// CHECK-LABEL: func.func @fold_add_on_two_matmuls(18// CHECK-SAME: %[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})19// CHECK-NEXT: %[[DENSE:.*]] = arith.constant dense<1.1120// CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0.000000e+0021// CHECK-NEXT: %[[EMPTY:.*]] = tensor.empty()22// CHECK-NEXT: %[[FILLED:.*]] = linalg.fill ins(%[[ZERO]] : {{.*}}) outs(%[[EMPTY]] : {{.*}})23// CHECK-NEXT: %[[ACC:.+]] = linalg.matmul ins(%[[ARG0]], %[[DENSE]] : {{.*}}) outs(%[[FILLED]] : {{.*}})24// CHECK-NEXT: %[[RES:.+]] = linalg.matmul ins(%[[ARG1]], %[[DENSE]] : {{.*}}) outs(%[[ACC]] : {{.*}})25// CHECK-NOT: linalg.add26// CHECK-NEXT: return %[[RES]]27 28module attributes {transform.with_named_sequence} {29 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {30 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op31 transform.apply_patterns to %func {32 transform.apply_patterns.linalg.fold_add_into_dest33 } : !transform.any_op34 transform.yield35 }36}37// -----38 39!type = tensor<2048x2048xf32>40func.func @expect_no_fold_of_add_as_orig_dest_not_additive_zero(%arg0: !type, %arg1: !type) -> !type {41 %0 = arith.constant dense<1.111111e+00> : !type42 %cst = arith.constant 0.000000e+00 : f3243 %1 = tensor.empty() : !type44 %2 = linalg.fill ins(%cst : f32) outs(%1 : !type) -> !type45 %3 = linalg.matmul ins(%arg0, %0 : !type, !type) outs(%2 : !type) -> !type46 %4 = linalg.matmul ins(%arg1, %0 : !type, !type) outs(%0 : !type) -> !type47 %5 = linalg.add ins(%3, %4 : !type, !type) outs(%1 : !type) -> !type48 return %5 : !type49}50 51// CHECK-LABEL: func.func @expect_no_fold_of_add_as_orig_dest_not_additive_zero52// CHECK: linalg.fill53// CHECK-NEXT: linalg.matmul54// CHECK-NEXT: linalg.matmul55// CHECK-NEXT: linalg.add56// CHECK-NEXT: return57 58module attributes {transform.with_named_sequence} {59 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {60 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op61 transform.apply_patterns to %func {62 transform.apply_patterns.linalg.fold_add_into_dest63 } : !transform.any_op64 transform.yield65 }66}67 68// -----69 70!type = tensor<2048x2048xf32>71func.func @expect_no_fold_of_add_as_contraction_result_has_multiple_users(%arg0: !type, %arg1: !type) -> (!type, !type) {72 %0 = arith.constant dense<1.111111e+00> : !type73 %cst = arith.constant 0.000000e+00 : f3274 %1 = tensor.empty() : !type75 %2 = linalg.fill ins(%cst : f32) outs(%1 : !type) -> !type76 %3 = linalg.matmul ins(%arg0, %0 : !type, !type) outs(%2 : !type) -> !type77 %4 = linalg.matmul ins(%arg1, %0 : !type, !type) outs(%0 : !type) -> !type78 %5 = linalg.add ins(%3, %4 : !type, !type) outs(%1 : !type) -> !type79 %6 = linalg.mul ins(%4, %arg0 : !type, !type) outs(%1 : !type) -> !type80 return %5, %6 : !type, !type81}82 83// CHECK-LABEL: func.func @expect_no_fold_of_add_as_contraction_result_has_multiple_users84// CHECK: linalg.fill85// CHECK-NEXT: linalg.matmul86// CHECK-NEXT: linalg.matmul87// CHECK-NEXT: linalg.add88// CHECK-NEXT: linalg.mul89// CHECK-NEXT: return90 91module attributes {transform.with_named_sequence} {92 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {93 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op94 transform.apply_patterns to %func {95 transform.apply_patterns.linalg.fold_add_into_dest96 } : !transform.any_op97 transform.yield98 }99}100 101// -----102 103!type = tensor<2048x2048xf32>104func.func @fold_add_on_matmul_and_func_arg(%arg0: !type, %arg1: !type) -> !type {105 %0 = arith.constant dense<1.111111e+00> : !type106 %cst = arith.constant 0.000000e+00 : f32107 %1 = tensor.empty() : !type108 %2 = linalg.fill ins(%cst : f32) outs(%1 : !type) -> !type109 %3 = linalg.matmul ins(%arg0, %0 : !type, !type) outs(%2 : !type) -> !type110 %5 = linalg.add ins(%3, %arg1 : !type, !type) outs(%1 : !type) -> !type111 return %5 : !type112}113 114// CHECK-LABEL: func.func @fold_add_on_matmul_and_func_arg115// CHECK: %[[RES:.+]] = linalg.matmul116// CHECK-NOT: linalg.add117// CHECK-NEXT: return %[[RES]]118 119module attributes {transform.with_named_sequence} {120 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {121 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op122 transform.apply_patterns to %func {123 transform.apply_patterns.linalg.fold_add_into_dest124 } : !transform.any_op125 transform.yield126 }127}128 129// -----130 131!type = tensor<2048x2048xf32>132func.func @expect_no_fold_of_add_as_operands_do_not_dominate_each_other(%arg0: !type, %arg1: !type) -> !type {133 %0 = arith.constant dense<1.111111e+00> : !type134 %cst = arith.constant 0.000000e+00 : f32135 %1 = tensor.empty() : !type136 %2 = linalg.fill ins(%cst : f32) outs(%1 : !type) -> !type137 %3 = linalg.matmul ins(%arg0, %0 : !type, !type) outs(%2 : !type) -> !type138 %4 = linalg.add ins(%3, %3 : !type, !type) outs(%1 : !type) -> !type139 return %4 : !type140}141 142// CHECK-LABEL: func.func @expect_no_fold_of_add_as_operands_do_not_dominate_each_other143// CHECK: linalg.fill144// CHECK-NEXT: linalg.matmul145// CHECK-NEXT: linalg.add146// CHECK-NEXT: return147 148module attributes {transform.with_named_sequence} {149 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {150 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op151 transform.apply_patterns to %func {152 transform.apply_patterns.linalg.fold_add_into_dest153 } : !transform.any_op154 transform.yield155 }156}157 158// -----159 160!type = tensor<2048x2048xf32>161func.func @expect_no_fold_of_add_as_dominated_op_is_not_a_contraction(%arg0: !type, %arg1: !type) -> !type {162 %0 = arith.constant dense<1.111111e+00> : !type163 %cst = arith.constant 0.000000e+00 : f32164 %1 = tensor.empty() : !type165 %2 = linalg.fill ins(%cst : f32) outs(%1 : !type) -> !type166 %3 = linalg.matmul ins(%arg0, %0 : !type, !type) outs(%2 : !type) -> !type167 %4 = linalg.sub ins(%arg1, %0 : !type, !type) outs(%2 : !type) -> !type168 %5 = linalg.add ins(%3, %4 : !type, !type) outs(%1 : !type) -> !type169 return %5 : !type170}171 172// CHECK-LABEL: func.func @expect_no_fold_of_add_as_dominated_op_is_not_a_contraction173// CHECK: linalg.fill174// CHECK-NEXT: linalg.matmul175// CHECK-NEXT: linalg.sub176// CHECK-NEXT: linalg.add177// CHECK-NEXT: return178 179module attributes {transform.with_named_sequence} {180 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {181 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op182 transform.apply_patterns to %func {183 transform.apply_patterns.linalg.fold_add_into_dest184 } : !transform.any_op185 transform.yield186 }187}188 189// -----190 191#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>192#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>193#map2 = affine_map<(d0, d1, d2) -> (d1, d0)> // NB: not an ordered projection194 195!type = tensor<2048x2048xf32>196func.func @expect_no_fold_of_add_as_dest_accumulation_is_not_identity_mapped(%arg0: !type, %arg1: !type) -> !type {197 %0 = arith.constant dense<1.111111e+00> : !type198 %cst = arith.constant 0.000000e+00 : f32199 %1 = tensor.empty() : !type200 %2 = linalg.fill ins(%cst : f32) outs(%1 : !type) -> !type201 %3 = linalg.generic { indexing_maps = [#map0, #map1, #map2],202 iterator_types = ["parallel", "parallel", "reduction"] }203 ins(%arg0, %0: !type, !type) outs(%2: !type) {204 ^bb0(%a: f32, %b: f32, %c: f32):205 %5 = arith.mulf %a, %b : f32206 %6 = arith.addf %c, %5 : f32207 linalg.yield %6 : f32208 } -> !type209 %4 = linalg.add ins(%3, %arg1 : !type, !type) outs(%1 : !type) -> !type210 return %4 : !type211}212 213// CHECK-LABEL: func.func @expect_no_fold_of_add_as_dest_accumulation_is_not_identity_mapped214// CHECK: linalg.fill215// CHECK-NEXT: linalg.generic216// CHECK: linalg.add217// CHECK-NEXT: return218 219module attributes {transform.with_named_sequence} {220 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {221 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op222 transform.apply_patterns to %func {223 transform.apply_patterns.linalg.fold_add_into_dest224 } : !transform.any_op225 transform.yield226 }227}228 229// -----230 231#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>232#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>233#map2 = affine_map<(d0, d1, d2) -> (d0, d1)> // NB: is an ordered projection234 235!type = tensor<2048x2048xf32>236func.func @fold_add_on_a_generic_and_an_argument(%arg0: !type, %arg1: !type) -> !type {237 %0 = arith.constant dense<1.111111e+00> : !type238 %cst = arith.constant 0.000000e+00 : f32239 %1 = tensor.empty() : !type240 %2 = linalg.fill ins(%cst : f32) outs(%1 : !type) -> !type241 %3 = linalg.generic { indexing_maps = [#map0, #map1, #map2],242 iterator_types = ["parallel", "parallel", "reduction"] }243 ins(%arg0, %0: !type, !type) outs(%2: !type) {244 ^bb0(%a: f32, %b: f32, %c: f32):245 %5 = arith.mulf %a, %b : f32246 %6 = arith.addf %c, %5 : f32247 linalg.yield %6 : f32248 } -> !type249 %4 = linalg.add ins(%3, %arg1 : !type, !type) outs(%1 : !type) -> !type250 return %4 : !type251}252 253// CHECK-LABEL: func.func @fold_add_on_a_generic_and_an_argument254// CHECK: linalg.generic255// CHECK-NOT: linalg.add256// CHECK: return257 258module attributes {transform.with_named_sequence} {259 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {260 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op261 transform.apply_patterns to %func {262 transform.apply_patterns.linalg.fold_add_into_dest263 } : !transform.any_op264 transform.yield265 }266}267 268// -----269 270memref.global "private" constant @big_const : memref<2048x2048xf32> = dense<1.11111104> {alignment = 64 : i64}271func.func @expect_no_fold_due_to_no_memref_support(%arg0: memref<2048x2048xf32>, %arg1: memref<2048x2048xf32>) -> memref<2048x2048xf32> {272 %cst = arith.constant 0.000000e+00 : f32273 %0 = memref.get_global @big_const : memref<2048x2048xf32>274 %alloc = memref.alloc() {alignment = 64 : i64} : memref<2048x2048xf32>275 %alloc_0 = memref.alloc() {alignment = 64 : i64} : memref<2048x2048xf32>276 linalg.fill ins(%cst : f32) outs(%alloc_0 : memref<2048x2048xf32>)277 linalg.matmul ins(%arg0, %0 : memref<2048x2048xf32>, memref<2048x2048xf32>) outs(%alloc_0 : memref<2048x2048xf32>)278 linalg.fill ins(%cst : f32) outs(%alloc : memref<2048x2048xf32>)279 linalg.matmul ins(%arg1, %0 : memref<2048x2048xf32>, memref<2048x2048xf32>) outs(%alloc : memref<2048x2048xf32>)280 linalg.add ins(%alloc_0, %alloc : memref<2048x2048xf32>, memref<2048x2048xf32>) outs(%alloc : memref<2048x2048xf32>)281 memref.dealloc %alloc_0 : memref<2048x2048xf32>282 return %alloc : memref<2048x2048xf32>283}284 285// CHECK-LABEL: func.func @expect_no_fold_due_to_no_memref_support286// CHECK: linalg.matmul287// CHECK: linalg.matmul288// CHECK: linalg.add289// CHECK: return290 291module attributes {transform.with_named_sequence} {292 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {293 %func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op294 transform.apply_patterns to %func {295 transform.apply_patterns.linalg.fold_add_into_dest296 } : !transform.any_op297 transform.yield298 }299}300