502 lines · plain
1// RUN: mlir-opt --split-input-file --transform-interpreter %s | FileCheck %s2 3func.func @matmul_split(%A : tensor<16x256xf32>, %B: tensor<256x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {4 %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)5 outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>6 return %0: tensor<16x32xf32>7}8 9// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>10// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d1)>11// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>12// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>13// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>14// CHECK-LABEL: @matmul_split15// CHECK-DAG: %[[ID:.*]] = arith.constant 0.000000e+00 : f3216// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] output_shape [16, 4, 64] : tensor<16x256xf32> into tensor<16x4x64xf32>17// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] output_shape [4, 64, 32] : tensor<256x32xf32> into tensor<4x64x32xf32>18// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<16x32x4xf32>19// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<16x32x4xf32>) -> tensor<16x32x4xf32>20// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]21// CHECK-SAME: , iterator_types = ["parallel", "parallel", "parallel", "reduction"]}22// CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<16x4x64xf32>, tensor<4x64x32xf32>) outs(%[[F]] : tensor<16x32x4xf32>) {23// CHECK: arith.mulf24// CHECK: arith.addf25// CHECK: linalg.yield26// CHECK: } -> tensor<16x32x4xf32>27// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]],28// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]} ins(%[[G]] : tensor<16x32x4xf32>) outs(%{{.*}} : tensor<16x32xf32>) {29// CHECK: arith.addf30// CHECK: linalg.yield %{{.*}} : f3231// CHECK: } -> tensor<16x32xf32>32// CHECK: return %[[R]] : tensor<16x32xf32>33 34module attributes {transform.with_named_sequence} {35 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {36 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op37 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2}38 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)39 transform.yield40 }41}42 43// -----44 45func.func @generic_split_1d(%arg0: tensor<32xf32>, %arg1: tensor<f32>, %out: tensor<f32>) -> tensor<f32> {46 %red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,47 affine_map<(d0) -> ()>,48 affine_map<(d0) -> ()>],49 iterator_types = ["reduction"]}50 ins(%arg0, %arg1 : tensor<32xf32>, tensor<f32>)51 outs(%out : tensor<f32>) {52 ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):53 %40 = arith.subf %arg7, %arg8 : f3254 %41 = math.exp %40 : f3255 %42 = arith.mulf %41, %arg9 : f3256 linalg.yield %42 : f3257 } -> tensor<f32>58 return %red : tensor<f32>59}60 61// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>62// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> ()>63// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>64// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)>65// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> ()>66//CHECK-LABEL: @generic_split_1d67// CHECK-DAG: %[[ID:.*]] = arith.constant 1.000000e+00 : f3268// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] output_shape [4, 8] : tensor<32xf32> into tensor<4x8xf32>69// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<4xf32>70// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<4xf32>) -> tensor<4xf32>71// CHECK: %[[G:.*]] = linalg.generic72// CHECK: {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]],73// CHECK: iterator_types = ["parallel", "reduction"]} ins(%[[I1]], %{{.*}} : tensor<4x8xf32>, tensor<f32>) outs(%[[F]] : tensor<4xf32>) {74// CHECK: arith.subf75// CHECK: math.exp76// CHECK: arith.mulf77// CHECK: linalg.yield78// CHECK: } -> tensor<4xf32>79// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["reduction"]} ins(%[[G]] : tensor<4xf32>) outs(%{{.*}} : tensor<f32>) {80// CHECK: arith.mulf81// CHECK: linalg.yield82// CHECK: } -> tensor<f32>83// CHECK: return %[[R]] : tensor<f32>84 85module attributes {transform.with_named_sequence} {86 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {87 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op88 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 0}89 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)90 transform.yield91 }92}93 94// -----95 96func.func @generic_split_3d(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>)97 -> tensor<5x2xf32>98{99 %0 = linalg.generic {100 indexing_maps = [101 affine_map<(d0, d1, d2) -> (d1, d0)>,102 affine_map<(d0, d1, d2) -> (d2, d1)>,103 affine_map<(d0, d1, d2) -> (d2, d0)>104 ],105 iterator_types = ["parallel", "reduction", "parallel"]106 } ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) outs(%output : tensor<5x2xf32>) {107 ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):108 %3 = arith.addf %arg0, %arg1 : f32109 %4 = arith.maximumf %3, %arg2 : f32110 linalg.yield %4 : f32111 } -> tensor<5x2xf32>112 return %0 : tensor<5x2xf32>113}114 115// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d0)>116// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d2, d1)>117// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)>118// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>119// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>120// CHECK-LABEL: func @generic_split_3d121// CHECK-DAG: %[[ID:.*]] = arith.constant 0xFF800000 : f32122// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] output_shape [4, 8, 2] : tensor<32x2xf32> into tensor<4x8x2xf32>123// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] output_shape [5, 4, 8] : tensor<5x32xf32> into tensor<5x4x8xf32>124// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32>125// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32>126// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]}127// CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<4x8x2xf32>, tensor<5x4x8xf32>) outs(%[[F]] : tensor<5x2x4xf32>) {128// CHECK: arith.addf129// CHECK: arith.maximumf130// CHECK: linalg.yield131// CHECK: } -> tensor<5x2x4xf32>132// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}133// CHECK-SAME: ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {134// CHECK: arith.maximumf135// CHECK: linalg.yield136// CHECK: } -> tensor<5x2xf32>137// CHECK: return %[[R]] : tensor<5x2xf32>138 139module attributes {transform.with_named_sequence} {140 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {141 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op142 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2}143 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)144 transform.yield145 }146}147 148// -----149 150// Check that we don't use -inf as the neutral element for maxf when maxf has151// ninf. Instead check that we use the smallest finite floating point value.152// Also check that the fastmath flags are set on the created maxf153// instructions.154func.func @generic_split_3d_ninf(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>)155 -> tensor<5x2xf32>156{157 %0 = linalg.generic {158 indexing_maps = [159 affine_map<(d0, d1, d2) -> (d1, d0)>,160 affine_map<(d0, d1, d2) -> (d2, d1)>,161 affine_map<(d0, d1, d2) -> (d2, d0)>162 ],163 iterator_types = ["parallel", "reduction", "parallel"]164 } ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) outs(%output : tensor<5x2xf32>) {165 ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):166 %3 = arith.addf %arg0, %arg1 : f32167 %4 = arith.maximumf %3, %arg2 fastmath<nnan,ninf> : f32168 linalg.yield %4 : f32169 } -> tensor<5x2xf32>170 return %0 : tensor<5x2xf32>171}172 173// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d0)>174// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d2, d1)>175// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)>176// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>177// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>178// CHECK-LABEL: func @generic_split_3d_ninf179// CHECK-DAG: %[[ID:.*]] = arith.constant -3.40282347E+38 : f32180// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] output_shape [4, 8, 2] : tensor<32x2xf32> into tensor<4x8x2xf32>181// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] output_shape [5, 4, 8] : tensor<5x32xf32> into tensor<5x4x8xf32>182// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32>183// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32>184// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]}185// CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<4x8x2xf32>, tensor<5x4x8xf32>) outs(%[[F]] : tensor<5x2x4xf32>) {186// CHECK: arith.addf187// CHECK: arith.maximumf {{.*}} fastmath<nnan,ninf>188// CHECK: linalg.yield189// CHECK: } -> tensor<5x2x4xf32>190// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}191// CHECK-SAME: ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {192// CHECK: arith.maximumf {{.*}} fastmath<nnan,ninf>193// CHECK: linalg.yield194// CHECK: } -> tensor<5x2xf32>195// CHECK: return %[[R]] : tensor<5x2xf32>196 197module attributes {transform.with_named_sequence} {198 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {199 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op200 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2}201 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)202 transform.yield203 }204}205 206// -----207 208func.func @matmul_split(%A : tensor<16x256xf32>, %B: tensor<256x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {209 %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)210 outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>211 return %0: tensor<16x32xf32>212}213 214// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>215// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d1)>216// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>217// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>218// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>219// CHECK-LABEL: @matmul_split220// CHECK-DAG: %[[ID:.*]] = arith.constant 0.000000e+00 : f32221// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] output_shape [16, 64, 4] : tensor<16x256xf32> into tensor<16x64x4xf32>222// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] output_shape [64, 4, 32] : tensor<256x32xf32> into tensor<64x4x32xf32>223// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<16x32x4xf32>224// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<16x32x4xf32>) -> tensor<16x32x4xf32>225// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]226// CHECK-SAME: , iterator_types = ["parallel", "parallel", "reduction", "parallel"]}227// CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<16x64x4xf32>, tensor<64x4x32xf32>) outs(%[[F]] : tensor<16x32x4xf32>) {228// CHECK: arith.mulf229// CHECK: arith.addf230// CHECK: linalg.yield231// CHECK: } -> tensor<16x32x4xf32>232// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]],233// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]} ins(%[[G]] : tensor<16x32x4xf32>) outs(%{{.*}} : tensor<16x32xf32>) {234// CHECK: arith.addf235// CHECK: linalg.yield %{{.*}} : f32236// CHECK: } -> tensor<16x32xf32>237// CHECK: return %[[R]] : tensor<16x32xf32>238 239module attributes {transform.with_named_sequence} {240 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {241 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op242 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2, inner_parallel}243 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)244 transform.yield245 }246}247 248// -----249 250func.func @generic_split_1d(%arg0: tensor<32xf32>, %arg1: tensor<f32>, %out: tensor<f32>) -> tensor<f32> {251 %red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,252 affine_map<(d0) -> ()>,253 affine_map<(d0) -> ()>],254 iterator_types = ["reduction"]}255 ins(%arg0, %arg1 : tensor<32xf32>, tensor<f32>)256 outs(%out : tensor<f32>) {257 ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):258 %40 = arith.subf %arg7, %arg8 : f32259 %41 = math.exp %40 : f32260 %42 = arith.mulf %41, %arg9 : f32261 linalg.yield %42 : f32262 } -> tensor<f32>263 return %red : tensor<f32>264}265 266// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>267// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> ()>268// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d1)>269// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)>270// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> ()>271//CHECK-LABEL: @generic_split_1d272// CHECK-DAG: %[[ID:.*]] = arith.constant 1.000000e+00 : f32273// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] output_shape [8, 4] : tensor<32xf32> into tensor<8x4xf32>274// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<4xf32>275// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<4xf32>) -> tensor<4xf32>276// CHECK: %[[G:.*]] = linalg.generic277// CHECK: {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]],278// CHECK: iterator_types = ["reduction", "parallel"]} ins(%[[I1]], %{{.*}} : tensor<8x4xf32>, tensor<f32>) outs(%[[F]] : tensor<4xf32>) {279// CHECK: arith.subf280// CHECK: math.exp281// CHECK: arith.mulf282// CHECK: linalg.yield283// CHECK: } -> tensor<4xf32>284// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["reduction"]} ins(%[[G]] : tensor<4xf32>) outs(%{{.*}} : tensor<f32>) {285// CHECK: arith.mulf286// CHECK: linalg.yield287// CHECK: } -> tensor<f32>288// CHECK: return %[[R]] : tensor<f32>289 290module attributes {transform.with_named_sequence} {291 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {292 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op293 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 0, inner_parallel}294 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)295 transform.yield296 }297}298 299// -----300 301func.func @generic_split_3d(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>)302 -> tensor<5x2xf32>303{304 %0 = linalg.generic {305 indexing_maps = [306 affine_map<(d0, d1, d2) -> (d1, d0)>,307 affine_map<(d0, d1, d2) -> (d2, d1)>,308 affine_map<(d0, d1, d2) -> (d2, d0)>309 ],310 iterator_types = ["parallel", "reduction", "parallel"]311 } ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) outs(%output : tensor<5x2xf32>) {312 ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):313 %3 = arith.addf %arg0, %arg1 : f32314 %4 = arith.minimumf %3, %arg2 : f32315 linalg.yield %4 : f32316 } -> tensor<5x2xf32>317 return %0 : tensor<5x2xf32>318}319 320// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d0)>321// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)>322// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)>323// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>324// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>325// CHECK-LABEL: func @generic_split_3d326// CHECK-DAG: %[[ID:.*]] = arith.constant 0x7F800000 : f32327// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] output_shape [8, 4, 2] : tensor<32x2xf32> into tensor<8x4x2xf32>328// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] output_shape [5, 8, 4] : tensor<5x32xf32> into tensor<5x8x4xf32>329// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32>330// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32>331// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]}332// CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<8x4x2xf32>, tensor<5x8x4xf32>) outs(%[[F]] : tensor<5x2x4xf32>) {333// CHECK: arith.addf334// CHECK: arith.minimumf335// CHECK: linalg.yield336// CHECK: } -> tensor<5x2x4xf32>337// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}338// CHECK-SAME: ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {339// CHECK: arith.minimumf340// CHECK: linalg.yield341// CHECK: } -> tensor<5x2xf32>342// CHECK: return %[[R]] : tensor<5x2xf32>343 344module attributes {transform.with_named_sequence} {345 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {346 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op347 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2, inner_parallel}348 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)349 transform.yield350 }351}352 353// -----354 355// Check that we don't use +inf as the neutral element for minf when minf has356// ninf. Instead check that we use the largest finite floating point value.357// Also check that the fastmath flags are set on the created minf358// instructions.359func.func @generic_split_3d(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>)360 -> tensor<5x2xf32>361{362 %0 = linalg.generic {363 indexing_maps = [364 affine_map<(d0, d1, d2) -> (d1, d0)>,365 affine_map<(d0, d1, d2) -> (d2, d1)>,366 affine_map<(d0, d1, d2) -> (d2, d0)>367 ],368 iterator_types = ["parallel", "reduction", "parallel"]369 } ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) outs(%output : tensor<5x2xf32>) {370 ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):371 %3 = arith.addf %arg0, %arg1 : f32372 %4 = arith.minimumf %3, %arg2 fastmath<ninf> : f32373 linalg.yield %4 : f32374 } -> tensor<5x2xf32>375 return %0 : tensor<5x2xf32>376}377 378// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d0)>379// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)>380// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)>381// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>382// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>383// CHECK-LABEL: func @generic_split_3d384// CHECK-DAG: %[[ID:.*]] = arith.constant 3.40282347E+38 : f32385// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] output_shape [8, 4, 2] : tensor<32x2xf32> into tensor<8x4x2xf32>386// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] output_shape [5, 8, 4] : tensor<5x32xf32> into tensor<5x8x4xf32>387// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32>388// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32>389// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]}390// CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<8x4x2xf32>, tensor<5x8x4xf32>) outs(%[[F]] : tensor<5x2x4xf32>) {391// CHECK: arith.addf392// CHECK: arith.minimumf {{.*}} fastmath<ninf>393// CHECK: linalg.yield394// CHECK: } -> tensor<5x2x4xf32>395// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}396// CHECK-SAME: ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {397// CHECK: arith.minimumf {{.*}} fastmath<ninf>398// CHECK: linalg.yield399// CHECK: } -> tensor<5x2xf32>400// CHECK: return %[[R]] : tensor<5x2xf32>401 402module attributes {transform.with_named_sequence} {403 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {404 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op405 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2, inner_parallel}406 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)407 transform.yield408 }409}410 411// -----412// Checks we use nan as the neutral element for maxnumf op.413func.func @generic_split_maxnumf(%in: tensor<32xf32>, %out: tensor<f32>) -> tensor<f32> {414 %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,415 affine_map<(d0) -> ()>],416 iterator_types = ["reduction"]}417 ins(%in : tensor<32xf32>)418 outs(%out : tensor<f32>) {419 ^bb0(%arg1: f32, %arg2: f32):420 %y = arith.maxnumf %arg1, %arg2 : f32421 linalg.yield %y : f32422 } -> tensor<f32>423 return %r : tensor<f32>424}425 426// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>427// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>428// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0) -> (d0)>429// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> ()>430// CHECK-LABEL: func @generic_split_maxnumf431// The float value 0xFFC00000 that is filled into the init tensor represents negative NaN.432// CHECK-DAG: %[[ID:.*]] = arith.constant 0xFFC00000 : f32433// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] output_shape [8, 4] : tensor<32xf32> into tensor<8x4xf32>434// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<4xf32>435// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<4xf32>) -> tensor<4xf32>436// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]}437// CHECK-SAME: ins(%[[I1]] : tensor<8x4xf32>) outs(%[[F]] : tensor<4xf32>) {438// CHECK: arith.maxnumf439// CHECK: linalg.yield440// CHECK: } -> tensor<4xf32>441// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP3]]], iterator_types = ["reduction"]}442// CHECK-SAME: ins(%[[G]] : tensor<4xf32>) outs(%{{.*}} : tensor<f32>) {443// CHECK: arith.maxnumf {{.*}}444// CHECK: linalg.yield445// CHECK: } -> tensor<f32>446// CHECK: return %[[R]] : tensor<f32>447 448module attributes {transform.with_named_sequence} {449 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {450 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op451 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 0, inner_parallel}452 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)453 transform.yield454 }455}456 457// -----458// Checks we use nan as the neutral element for minnumf op.459func.func @generic_split_minnumf(%in: tensor<32xf32>, %out: tensor<f32>) -> tensor<f32> {460 %r = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,461 affine_map<(d0) -> ()>],462 iterator_types = ["reduction"]}463 ins(%in : tensor<32xf32>)464 outs(%out : tensor<f32>) {465 ^bb0(%arg1: f32, %arg2: f32):466 %y = arith.minnumf %arg1, %arg2 : f32467 linalg.yield %y : f32468 } -> tensor<f32>469 return %r : tensor<f32>470}471 472// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>473// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>474// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0) -> (d0)>475// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> ()>476// CHECK-LABEL: func @generic_split_minnumf477// The float value 0x7FC00000 that is filled into the init tensor represents positive NaN.478// CHECK-DAG: %[[ID:.*]] = arith.constant 0x7FC00000 : f32479// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] output_shape [8, 4] : tensor<32xf32> into tensor<8x4xf32>480// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<4xf32>481// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<4xf32>) -> tensor<4xf32>482// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]}483// CHECK-SAME: ins(%[[I1]] : tensor<8x4xf32>) outs(%[[F]] : tensor<4xf32>) {484// CHECK: arith.minnumf485// CHECK: linalg.yield486// CHECK: } -> tensor<4xf32>487// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP3]]], iterator_types = ["reduction"]}488// CHECK-SAME: ins(%[[G]] : tensor<4xf32>) outs(%{{.*}} : tensor<f32>) {489// CHECK: arith.minnumf {{.*}}490// CHECK: linalg.yield491// CHECK: } -> tensor<f32>492// CHECK: return %[[R]] : tensor<f32>493 494module attributes {transform.with_named_sequence} {495 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {496 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op497 %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 0, inner_parallel}498 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)499 transform.yield500 }501}502