1264 lines · plain
1// RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops | FileCheck %s2 3func.func @generalize_matmul_buffer(%A : memref<16x8xf32>, %B: memref<8x32xf32>, %C: memref<16x32xf32>) {4 linalg.matmul ins(%A, %B: memref<16x8xf32>, memref<8x32xf32>)5 outs(%C: memref<16x32xf32>)6 return7}8 9 10// CHECK: #[[A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>11// CHECK: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>12// CHECK: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>13 14// CHECK: func @generalize_matmul_buffer15// CHECK-SAME: %[[A:.+]]: memref<16x8xf32>16// CHECK-SAME: %[[B:.+]]: memref<8x32xf32>17// CHECK-SAME: %[[C:.+]]: memref<16x32xf32>18 19// CHECK: linalg.generic20// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]]21// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]22// CHECK-SAME: ins(%[[A]], %[[B]]23// CHECK-SAME: outs(%[[C]]24 25// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)26// CHECK: %[[MUL:.+]] = arith.mulf %[[A_ARG]], %[[B_ARG]] : f3227// CHECK: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f3228// CHECK: linalg.yield %[[ADD]] : f3229 30// -----31 32func.func @matmul_bcast_a(%arg0: memref<5xf32>, %arg1: memref<5x7xf32>, %arg2: memref<3x7xf32>) {33 linalg.matmul indexing_maps = [34 affine_map<(d0, d1, d2) -> (d2)>,35 affine_map<(d0, d1, d2) -> (d2, d1)>,36 affine_map<(d0, d1, d2) -> (d0, d1)>37 ]38 ins(%arg0, %arg1 : memref<5xf32>, memref<5x7xf32>) outs(%arg2: memref<3x7xf32>)39 return40}41 42// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2) -> (d2)>43// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>44// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>45// CHECK-LABEL: func.func @matmul_bcast_a(46// CHECK-SAME: %[[VAL_0:.*]]: memref<5xf32>,47// CHECK-SAME: %[[VAL_1:.*]]: memref<5x7xf32>,48// CHECK-SAME: %[[VAL_2:.*]]: memref<3x7xf32>) {49// CHECK: linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_2]]], iterator_types = ["parallel", "parallel", "reduction"]} ins(%[[VAL_0]], %[[VAL_1]] : memref<5xf32>, memref<5x7xf32>) outs(%[[VAL_2]] : memref<3x7xf32>) {50// CHECK: ^bb0(%[[VAL_3:.*]]: f32, %[[VAL_4:.*]]: f32, %[[VAL_5:.*]]: f32):51// CHECK: %[[VAL_6:.*]] = arith.mulf %[[VAL_3]], %[[VAL_4]] : f3252// CHECK: %[[VAL_7:.*]] = arith.addf %[[VAL_5]], %[[VAL_6]] : f3253// CHECK: linalg.yield %[[VAL_7]] : f3254// CHECK: }55// CHECK: return56// CHECK: }57 58// -----59 60func.func @generalize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {61 %0 = linalg.matmul ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>)62 outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32>63 return %0: tensor<16x32xf32>64}65 66// CHECK: func @generalize_matmul_tensor67 68// CHECK: linalg.generic69// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xf32>, tensor<8x32xf32>)70// CHECK-SAME: outs(%{{.+}} : tensor<16x32xf32>)71 72// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)73// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_ARG]], %[[B_ARG]] : f3274// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f3275// CHECK-NEXT: linalg.yield %[[ADD]] : f3276// CHECK-NEXT: -> tensor<16x32xf32>77 78// -----79 80func.func @generalize_matmul_tensor_complex(%A : tensor<16x8xcomplex<f32>>,81 %B: tensor<8x32xcomplex<f32>>,82 %C: tensor<16x32xcomplex<f32>>)83 -> tensor<16x32xcomplex<f32>> {84 %0 = linalg.matmul ins(%A, %B: tensor<16x8xcomplex<f32>>, tensor<8x32xcomplex<f32>>)85 outs(%C: tensor<16x32xcomplex<f32>>) -> tensor<16x32xcomplex<f32>>86 return %0: tensor<16x32xcomplex<f32>>87}88 89// CHECK: func @generalize_matmul_tensor_complex90 91// CHECK: linalg.generic92// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xcomplex<f32>>, tensor<8x32xcomplex<f32>>)93// CHECK-SAME: outs(%{{.+}} : tensor<16x32xcomplex<f32>>)94 95// CHECK: ^{{.*}}(%[[A_ARG:.+]]: complex<f32>, %[[B_ARG:.+]]: complex<f32>, %[[C_ARG:.+]]: complex<f32>)96// CHECK-NEXT: %[[MUL:.+]] = complex.mul %[[A_ARG]], %[[B_ARG]] : complex<f32>97// CHECK-NEXT: %[[ADD:.+]] = complex.add %[[C_ARG]], %[[MUL]] : complex<f32>98// CHECK-NEXT: linalg.yield %[[ADD]] : complex<f32>99// CHECK-NEXT: -> tensor<16x32xcomplex<f32>>100 101// -----102 103func.func @depthwise_conv_2d_nhwc_hwcm(%input: memref<2x4x5x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<2x3x4x2x3xf32>) {104 linalg.depthwise_conv_2d_nhwc_hwcm105 { dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }106 ins(%input, %filter : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)107 outs(%output : memref<2x3x4x2x3xf32>)108 return109}110 111// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d5, d2 + d6, d3)>112// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d3, d4)>113// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4)>114 115// CHECK: func @depthwise_conv_2d_nhwc_hwcm116 117// CHECK: linalg.generic118// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]119// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}120// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)121// CHECK-SAME: outs(%{{.+}} : memref<2x3x4x2x3xf32>)122 123// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)124// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32125// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32126// CHECK-NEXT: linalg.yield %[[ADD]] : f32127 128// -----129 130func.func @depthwise_conv_2d_nhwc_hwcm(%input: memref<2x4x5x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<2x2x3x2x3xf32>) {131 linalg.depthwise_conv_2d_nhwc_hwcm132 { dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }133 ins(%input, %filter : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)134 outs(%output : memref<2x2x3x2x3xf32>)135 return136}137 138// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d5 * 2, d2 + d6 * 2, d3)>139// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d3, d4)>140// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4)>141 142// CHECK: func @depthwise_conv_2d_nhwc_hwcm143 144// CHECK: linalg.generic145// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]146// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}147// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)148// CHECK-SAME: outs(%{{.+}} : memref<2x2x3x2x3xf32>)149 150// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)151// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32152// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32153// CHECK-NEXT: linalg.yield %[[ADD]] : f32154 155// -----156 157func.func @depthwise_conv_2d_nhwc_hwc(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {158 linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}159 ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)160 outs(%output: memref<1x56x56x96xf32>)161 return162}163 164// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4, d2 * 2 + d5, d3)>165// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5, d3)>166// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>167 168// CHECK: func @depthwise_conv_2d_nhwc_hwc169 170// CHECK: linalg.generic171// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]172// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}173// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x113x113x96xf32>, memref<3x3x96xf32>)174// CHECK-SAME: outs(%{{.+}} : memref<1x56x56x96xf32>)175 176// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)177// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32178// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32179// CHECK-NEXT: linalg.yield %[[ADD]] : f32180 181// -----182 183func.func @conv_1d_nwc_wcf(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) {184 linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>,185 strides = dense<1> : tensor<1xi64>}186 ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>)187 outs (%output: memref<?x?x?xf32>)188 return189}190// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1 + d3, d4)>191// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4, d2)>192// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>193 194// CHECK: func @conv_1d_nwc_wcf195 196// CHECK: linalg.generic197// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]198// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]}199// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>)200// CHECK-SAME: outs(%{{.+}} : memref<?x?x?xf32>)201 202// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)203// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32204// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32205// CHECK-NEXT: linalg.yield %[[ADD]] : f32206 207// -----208 209func.func @conv_1d_ncw_fcw(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) {210 linalg.conv_1d_ncw_fcw {dilations = dense<1> : tensor<1xi64>,211 strides = dense<1> : tensor<1xi64>}212 ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>)213 outs (%output: memref<?x?x?xf32>)214 return215}216// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2 + d4)>217// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3, d4)>218// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>219 220// CHECK: func @conv_1d_ncw_fcw221 222// CHECK: linalg.generic223// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]224// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]}225// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>)226// CHECK-SAME: outs(%{{.+}} : memref<?x?x?xf32>)227 228// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)229// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32230// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32231// CHECK-NEXT: linalg.yield %[[ADD]] : f32232 233// -----234 235func.func @conv_2d_ngchw_gfchw_q(%input: memref<?x?x?x?x?xi8>, %filter: memref<?x?x?x?x?xi8>, %inputzp: i32, %filterzp: i32, %output: memref<?x?x?x?x?xi32>) {236 linalg.conv_2d_ngchw_gfchw_q {dilations = dense<1> : tensor<2xi64>,237 strides = dense<1> : tensor<2xi64>}238 ins (%input, %filter, %inputzp, %filterzp: memref<?x?x?x?x?xi8>, memref<?x?x?x?x?xi8>, i32, i32)239 outs (%output: memref<?x?x?x?x?xi32>)240 return241}242// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d5, d3 + d6, d4 + d7)>243// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d1, d2, d5, d6, d7)>244// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> ()>245// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4)>246 247// CHECK: func @conv_2d_ngchw_gfchw_q248 249// CHECK: linalg.generic250// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP2]], #[[MAP3]]]251// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]}252// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}} : memref<?x?x?x?x?xi8>, memref<?x?x?x?x?xi8>, i32, i32)253// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?x?xi32>)254 255// CHECK: ^{{.+}}(%[[BBARG0:.+]]: i8, %[[BBARG1:.+]]: i8, %[[BBARG2:.+]]: i32, %[[BBARG3:.+]]: i32, %[[BBARG4:.+]]: i32)256// CHECK-NEXT: %[[EXTSI0:.+]] = arith.extsi %[[BBARG0]] : i8 to i32257// CHECK-NEXT: %[[SUB0:.+]] = arith.subi %[[EXTSI0]], %[[BBARG2]] : i32258// CHECK-NEXT: %[[EXTSI1:.+]] = arith.extsi %[[BBARG1]] : i8 to i32259// CHECK-NEXT: %[[SUB1:.+]] = arith.subi %[[EXTSI1]], %[[BBARG3]] : i32260// CHECK-NEXT: %[[MUL:.+]] = arith.muli %[[SUB0]], %[[SUB1]] : i32261// CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[BBARG4]], %[[MUL]] : i32262// CHECK-NEXT: linalg.yield %[[ADD]] : i32263 264// -----265 266func.func @generalize_fill(%output: memref<?x?xf32>, %value : f32) {267 linalg.fill ins(%value : f32) outs(%output : memref<?x?xf32>)268 return269}270 271// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> ()>272// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>273 274// CHECK: func @generalize_fill275// CHECK-SAME: (%[[ARG0:.+]]: memref<?x?xf32>, %[[VAL:.+]]: f32)276 277// CHECK: linalg.generic278// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]279// CHECK-SAME: iterator_types = ["parallel", "parallel"]}280// CHECK-SAME: ins(%[[VAL]] : f32)281// CHECK-SAME: outs(%{{.+}} : memref<?x?xf32>)282 283// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)284// CHECK-NEXT: linalg.yield %[[BBARG0]] : f32285 286// -----287 288func.func @generalize_batch_matm_vec(%lhs : memref<?x?x?xi8>, %rhs: memref<?x?xi8>, %out: memref<?x?xf32>) {289 linalg.batch_matvec ins(%lhs, %rhs: memref<?x?x?xi8>, memref<?x?xi8>)290 outs(%out: memref<?x?xf32>)291 return292}293// CHECK: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>294// CHECK: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>295// CHECK: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>296 297// CHECK: @generalize_batch_matm_vec298 299// CHECK: linalg.generic300// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]301// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}302// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xi8>, memref<?x?xi8>)303// CHECK-SAME: outs(%{{.+}} : memref<?x?xf32>)304// CHECK: ^{{.+}}(%[[BBARG0:.+]]: i8, %[[BBARG1:.+]]: i8, %[[BBARG2:.+]]: f32)305// CHECK: %[[BBARG0_F32:.+]] = arith.sitofp %[[BBARG0]] : i8 to f32306// CHECK: %[[BBARG1_F32:.+]] = arith.sitofp %[[BBARG1]] : i8 to f32307// CHECK: %[[MUL:.+]] = arith.mulf %[[BBARG0_F32]], %[[BBARG1_F32]]308// CHECK: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]]309// CHECK: linalg.yield %[[ADD]] : f32310 311// -----312 313func.func @generalize_batch_vecmat(%lhs : memref<?x?xi8>, %rhs: memref<?x?x?xi8>, %out: memref<?x?xf32>) {314 linalg.batch_vecmat ins(%lhs, %rhs: memref<?x?xi8>, memref<?x?x?xi8>)315 outs(%out: memref<?x?xf32>)316 return317}318// CHECK: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>319// CHECK: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>320// CHECK: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>321 322// CHECK: @generalize_batch_vecmat323 324// CHECK: linalg.generic325// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]326// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}327// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?xi8>, memref<?x?x?xi8>)328// CHECK-SAME: outs(%{{.+}} : memref<?x?xf32>)329// CHECK: ^{{.+}}(%[[BBARG0:.+]]: i8, %[[BBARG1:.+]]: i8, %[[BBARG2:.+]]: f32)330// CHECK: %[[BBARG0_F32:.+]] = arith.sitofp %[[BBARG0]] : i8 to f32331// CHECK: %[[BBARG1_F32:.+]] = arith.sitofp %[[BBARG1]] : i8 to f32332// CHECK: %[[MUL:.+]] = arith.mulf %[[BBARG0_F32]], %[[BBARG1_F32]]333// CHECK: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]]334// CHECK: linalg.yield %[[ADD]] : f32335 336// -----337 338func.func @batch_reduce_gemm(%lhs: memref<7x8x9xf32>, %rhs: memref<7x9x8xf32>, %out: memref<8x8xf32>) {339 linalg.batch_reduce_matmul ins(%lhs, %rhs: memref<7x8x9xf32>, memref<7x9x8xf32>)340 outs(%out: memref<8x8xf32>)341 return342}343 344// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>345// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>346// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>347 348// CHECK: @batch_reduce_gemm349 350// CHECK: linalg.generic351// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]352// CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel", "reduction"]}353// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<7x8x9xf32>, memref<7x9x8xf32>)354// CHECK-SAME: outs(%{{.+}} : memref<8x8xf32>355// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)356// CHECK: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32357// CHECK: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32358// CHECK: linalg.yield %[[ADD]] : f32359 360// -----361 362func.func @generalize_batch_reduce_gemm_bf16(%lhs: memref<7x8x9xbf16>, %rhs: memref<7x9x8xbf16>, %out: memref<8x8xf32>) {363 linalg.batch_reduce_matmul ins(%lhs, %rhs: memref<7x8x9xbf16>, memref<7x9x8xbf16>)364 outs(%out: memref<8x8xf32>)365 return366}367 368// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>369// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>370// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>371 372// CHECK: @generalize_batch_reduce_gemm_bf16373 374// CHECK: linalg.generic375// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]376// CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel", "reduction"]}377// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<7x8x9xbf16>, memref<7x9x8xbf16>)378// CHECK-SAME: outs(%{{.+}} : memref<8x8xf32>379// CHECK: ^{{.+}}(%[[BBARG0:.+]]: bf16, %[[BBARG1:.+]]: bf16, %[[BBARG2:.+]]: f32)380// CHECK: %[[EXTBF16_0:.+]] = arith.extf %[[BBARG0]] : bf16 to f32381// CHECK: %[[EXTBF16_1:.+]] = arith.extf %[[BBARG1]] : bf16 to f32382// CHECK: %[[MUL:.+]] = arith.mulf %[[EXTBF16_0]], %[[EXTBF16_1]] : f32383// CHECK: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32384// CHECK: linalg.yield %[[ADD]] : f32385 386 387// -----388 389func.func @generalize_linalg_map(%arg0: memref<1x8x8x8xf32>, %arg1: memref<1x8x8x8xf32>, %arg2: memref<1x8x8x8xf32>) {390 %cst = arith.constant 0.000000e+00 : f32391 linalg.map {arith.addf} ins(%arg0, %arg1: memref<1x8x8x8xf32>, memref<1x8x8x8xf32>) outs(%arg2 : memref<1x8x8x8xf32>)392 return393}394 395// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>396 397// CHECK: @generalize_linalg_map398 399// CHECK: linalg.generic400// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP0]], #[[MAP0]]]401// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]}402// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x8x8x8xf32>, memref<1x8x8x8xf32>) outs(%{{.+}} : memref<1x8x8x8xf32>403// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)404// CHECK: %[[ADD:.+]] = arith.addf %[[BBARG0]], %[[BBARG1]] : f32405// CHECK: linalg.yield %[[ADD]] : f32406 407// -----408 409func.func @generalize_add(%lhs: memref<7x14x21xf32>, %rhs: memref<7x14x21xf32>,410 %out: memref<7x14x21xf32>) {411 linalg.add ins(%lhs, %rhs : memref<7x14x21xf32>, memref<7x14x21xf32>)412 outs(%out : memref<7x14x21xf32>)413 return414}415 416// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>417 418// CHECK: func @generalize_add419// CHECK-SAME: (%[[LHS:.+]]: memref<7x14x21xf32>, %[[RHS:.+]]: memref<7x14x21xf32>,420// CHECK-SAME: %[[OUT:.+]]: memref<7x14x21xf32>)421 422// CHECK: linalg.generic423// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]424// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}425// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<7x14x21xf32>, memref<7x14x21xf32>)426// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)427 428// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)429// CHECK-NEXT: %[[SUM:.+]] = arith.addf %[[BBARG0]], %[[BBARG1]] : f32430// CHECK-NEXT: linalg.yield %[[SUM]] : f32431 432// -----433 434func.func @generalize_sub(%lhs: memref<7x14x21xf32>, %rhs: memref<7x14x21xf32>,435 %out: memref<7x14x21xf32>) {436 linalg.sub ins(%lhs, %rhs : memref<7x14x21xf32>, memref<7x14x21xf32>)437 outs(%out : memref<7x14x21xf32>)438 return439}440 441// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>442 443// CHECK: func @generalize_sub444// CHECK-SAME: (%[[LHS:.+]]: memref<7x14x21xf32>, %[[RHS:.+]]: memref<7x14x21xf32>,445// CHECK-SAME: %[[OUT:.+]]: memref<7x14x21xf32>)446 447// CHECK: linalg.generic448// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]449// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}450// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<7x14x21xf32>, memref<7x14x21xf32>)451// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)452 453// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)454// CHECK-NEXT: %[[SUB:.+]] = arith.subf %[[BBARG0]], %[[BBARG1]] : f32455// CHECK-NEXT: linalg.yield %[[SUB]] : f32456 457// -----458 459func.func @generalize_mul(%lhs: memref<7x14x21xf32>, %rhs: memref<7x14x21xf32>,460 %out: memref<7x14x21xf32>) {461 linalg.mul ins(%lhs, %rhs : memref<7x14x21xf32>, memref<7x14x21xf32>)462 outs(%out : memref<7x14x21xf32>)463 return464}465 466// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>467 468// CHECK: func @generalize_mul469// CHECK-SAME: (%[[LHS:.+]]: memref<7x14x21xf32>, %[[RHS:.+]]: memref<7x14x21xf32>,470// CHECK-SAME: %[[OUT:.+]]: memref<7x14x21xf32>)471 472// CHECK: linalg.generic473// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]474// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}475// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<7x14x21xf32>, memref<7x14x21xf32>)476// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)477 478// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)479// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32480// CHECK-NEXT: linalg.yield %[[MUL]] : f32481 482// -----483 484func.func @generalize_div(%lhs: memref<7x14x21xf32>, %rhs: memref<7x14x21xf32>,485 %out: memref<7x14x21xf32>) {486 linalg.div ins(%lhs, %rhs : memref<7x14x21xf32>, memref<7x14x21xf32>)487 outs(%out : memref<7x14x21xf32>)488 return489}490 491// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>492 493// CHECK: func @generalize_div494// CHECK-SAME: (%[[LHS:.+]]: memref<7x14x21xf32>, %[[RHS:.+]]: memref<7x14x21xf32>,495// CHECK-SAME: %[[OUT:.+]]: memref<7x14x21xf32>)496 497// CHECK: linalg.generic498// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]499// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}500// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<7x14x21xf32>, memref<7x14x21xf32>)501// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)502 503// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)504// CHECK-NEXT: %[[DIV:.+]] = arith.divf %[[BBARG0]], %[[BBARG1]] : f32505// CHECK-NEXT: linalg.yield %[[DIV]] : f32506 507// -----508 509func.func @generalize_divu(%lhs: memref<7x14x21xi32>, %rhs: memref<7x14x21xi32>,510 %out: memref<7x14x21xi32>) {511 linalg.div_unsigned ins(%lhs, %rhs : memref<7x14x21xi32>, memref<7x14x21xi32>)512 outs(%out : memref<7x14x21xi32>)513 return514}515 516// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>517 518// CHECK: func @generalize_divu519// CHECK-SAME: (%[[LHS:.+]]: memref<7x14x21xi32>, %[[RHS:.+]]: memref<7x14x21xi32>,520// CHECK-SAME: %[[OUT:.+]]: memref<7x14x21xi32>)521 522// CHECK: linalg.generic523// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]524// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}525// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<7x14x21xi32>, memref<7x14x21xi32>)526// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xi32>)527 528// CHECK: ^{{.+}}(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: i32, %[[BBARG2:.+]]: i32)529// CHECK-NEXT: %[[DIVU:.+]] = arith.divui %[[BBARG0]], %[[BBARG1]] : i32530// CHECK-NEXT: linalg.yield %[[DIVU]] : i32531 532// -----533 534func.func @generalize_exp(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {535 linalg.exp ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)536 return537}538 539// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>540 541// CHECK: func @generalize_exp542// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)543 544// CHECK: linalg.generic545// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]546// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}547// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)548 549// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)550// CHECK-NEXT: %[[EXP:.+]] = math.exp %[[BBARG0]] : f32551// CHECK-NEXT: linalg.yield %[[EXP]] : f32552 553// -----554 555func.func @generalize_log(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {556 linalg.log ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)557 return558}559 560// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>561 562// CHECK: func @generalize_log563// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)564 565// CHECK: linalg.generic566// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]567// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}568// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)569 570// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)571// CHECK-NEXT: %[[log:.+]] = math.log %[[BBARG0]] : f32572// CHECK-NEXT: linalg.yield %[[log]] : f32573 574// -----575 576func.func @generalize_abs(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {577 linalg.abs ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)578 return579}580 581// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>582 583// CHECK: func @generalize_abs584// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)585 586// CHECK: linalg.generic587// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]588// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}589// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)590 591// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)592// CHECK-NEXT: %[[abs:.+]] = math.absf %[[BBARG0]] : f32593// CHECK-NEXT: linalg.yield %[[abs]] : f32594 595// -----596 597func.func @generalize_ceil(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {598 linalg.ceil ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)599 return600}601 602// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>603 604// CHECK: func @generalize_ceil605// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)606 607// CHECK: linalg.generic608// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]609// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}610// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)611 612// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)613// CHECK-NEXT: %[[ceil:.+]] = math.ceil %[[BBARG0]] : f32614// CHECK-NEXT: linalg.yield %[[ceil]] : f32615 616// -----617 618func.func @generalize_floor(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {619 linalg.floor ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)620 return621}622 623// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>624 625// CHECK: func @generalize_floor626// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)627 628// CHECK: linalg.generic629// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]630// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}631// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)632 633// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)634// CHECK-NEXT: %[[floor:.+]] = math.floor %[[BBARG0]] : f32635// CHECK-NEXT: linalg.yield %[[floor]] : f32636 637// -----638 639func.func @generalize_negf(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {640 linalg.negf ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)641 return642}643 644// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>645 646// CHECK: func @generalize_negf647// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)648 649// CHECK: linalg.generic650// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]651// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}652// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)653 654// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)655// CHECK-NEXT: %[[negf:.+]] = arith.negf %[[BBARG0]] : f32656// CHECK-NEXT: linalg.yield %[[negf]] : f32657 658// -----659 660func.func @generalize_reciprocal(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {661 linalg.reciprocal ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)662 return663}664 665// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>666 667// CHECK: func @generalize_reciprocal668// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)669 670// CHECK: %[[one:.+]] = arith.constant 1.000000e+00 : f32671 672// CHECK: linalg.generic673// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]674// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}675// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)676 677// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)678// CHECK-NEXT: %[[reciprocal:.+]] = arith.divf %[[one]], %[[BBARG0]] : f32679// CHECK-NEXT: linalg.yield %[[reciprocal]] : f32680 681// -----682 683func.func @generalize_round(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {684 linalg.round ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)685 return686}687 688// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>689 690// CHECK: func @generalize_round691// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)692 693// CHECK: linalg.generic694// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]695// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}696// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)697 698// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)699// CHECK-NEXT: %[[round:.+]] = math.round %[[BBARG0]] : f32700// CHECK-NEXT: linalg.yield %[[round]] : f32701 702// -----703 704func.func @generalize_sqrt(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {705 linalg.sqrt ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)706 return707}708 709// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>710 711// CHECK: func @generalize_sqrt712// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)713 714// CHECK: linalg.generic715// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]716// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}717// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)718 719// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)720// CHECK-NEXT: %[[sqrt:.+]] = math.sqrt %[[BBARG0]] : f32721// CHECK-NEXT: linalg.yield %[[sqrt]] : f32722 723// -----724 725func.func @generalize_rsqrt(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {726 linalg.rsqrt ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)727 return728}729 730// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>731 732// CHECK: func @generalize_rsqrt733// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)734 735// CHECK: linalg.generic736// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]737// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}738// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)739 740// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)741// CHECK-NEXT: %[[rsqrt:.+]] = math.rsqrt %[[BBARG0]] : f32742// CHECK-NEXT: linalg.yield %[[rsqrt]] : f32743 744// -----745 746func.func @generalize_square(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {747 linalg.square ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)748 return749}750 751// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>752 753// CHECK: func @generalize_square754// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)755 756// CHECK: linalg.generic757// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]758// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}759// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)760 761// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)762// CHECK-NEXT: %[[square:.+]] = arith.mulf %[[BBARG0]], %[[BBARG0]] : f32763// CHECK-NEXT: linalg.yield %[[square]] : f32764 765// -----766 767func.func @generalize_tanh(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {768 linalg.tanh ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)769 return770}771 772// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>773 774// CHECK: func @generalize_tanh775// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)776 777// CHECK: linalg.generic778// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]779// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}780// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)781 782// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)783// CHECK-NEXT: %[[tanh:.+]] = math.tanh %[[BBARG0]] : f32784// CHECK-NEXT: linalg.yield %[[tanh]] : f32785 786// -----787 788func.func @generalize_erf(%arg: memref<7x14x21xf32>, %out: memref<7x14x21xf32>) {789 linalg.erf ins(%arg : memref<7x14x21xf32>) outs(%out : memref<7x14x21xf32>)790 return791}792 793// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>794 795// CHECK: func @generalize_erf796// CHECK-SAME: (%[[ARG:.+]]: memref<7x14x21xf32>, %[[OUT:.+]]: memref<7x14x21xf32>)797 798// CHECK: linalg.generic799// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]800// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}801// CHECK-SAME: ins(%[[LHS]] : memref<7x14x21xf32>) outs(%[[OUT]] : memref<7x14x21xf32>)802 803// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)804// CHECK-NEXT: %[[erf:.+]] = math.erf %[[BBARG0]] : f32805// CHECK-NEXT: linalg.yield %[[erf]] : f32806 807// -----808 809func.func @generalize_max(%lhs: memref<7x14x21xf32>, %rhs: memref<7x14x21xf32>,810 %out: memref<7x14x21xf32>) {811 linalg.max ins(%lhs, %rhs : memref<7x14x21xf32>, memref<7x14x21xf32>)812 outs(%out : memref<7x14x21xf32>)813 return814}815 816// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>817 818// CHECK: func @generalize_max819// CHECK-SAME: (%[[LHS:.+]]: memref<7x14x21xf32>, %[[RHS:.+]]: memref<7x14x21xf32>,820// CHECK-SAME: %[[OUT:.+]]: memref<7x14x21xf32>)821 822// CHECK: linalg.generic823// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]824// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}825// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<7x14x21xf32>, memref<7x14x21xf32>)826// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)827 828// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)829// CHECK-NEXT: %[[max:.+]] = arith.maximumf %[[BBARG0]], %[[BBARG1]] : f32830// CHECK-NEXT: linalg.yield %[[max]] : f32831 832// -----833 834func.func @generalize_min(%lhs: memref<7x14x21xf32>, %rhs: memref<7x14x21xf32>,835 %out: memref<7x14x21xf32>) {836 linalg.min ins(%lhs, %rhs : memref<7x14x21xf32>, memref<7x14x21xf32>)837 outs(%out : memref<7x14x21xf32>)838 return839}840 841// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>842 843// CHECK: func @generalize_min844// CHECK-SAME: (%[[LHS:.+]]: memref<7x14x21xf32>, %[[RHS:.+]]: memref<7x14x21xf32>,845// CHECK-SAME: %[[OUT:.+]]: memref<7x14x21xf32>)846 847// CHECK: linalg.generic848// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]849// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}850// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<7x14x21xf32>, memref<7x14x21xf32>)851// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)852 853// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)854// CHECK-NEXT: %[[min:.+]] = arith.minimumf %[[BBARG0]], %[[BBARG1]] : f32855// CHECK-NEXT: linalg.yield %[[min]] : f32856 857 858// -----859 860func.func @generalize_powf(%lhs: memref<7x14x21xf32>, %rhs: memref<7x14x21xf32>,861 %out: memref<7x14x21xf32>) {862 linalg.powf ins(%lhs, %rhs : memref<7x14x21xf32>, memref<7x14x21xf32>)863 outs(%out : memref<7x14x21xf32>)864 return865}866 867// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>868 869// CHECK: func @generalize_powf870// CHECK-SAME: (%[[LHS:.+]]: memref<7x14x21xf32>, %[[RHS:.+]]: memref<7x14x21xf32>,871// CHECK-SAME: %[[OUT:.+]]: memref<7x14x21xf32>)872 873// CHECK: linalg.generic874// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]875// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}876// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<7x14x21xf32>, memref<7x14x21xf32>)877// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)878 879// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)880// CHECK-NEXT: %[[powf:.+]] = math.powf %[[BBARG0]], %[[BBARG1]] : f32881// CHECK-NEXT: linalg.yield %[[powf]] : f32882 883 884// -----885 886func.func @generalize_select(%cond: memref<7x14x21xi1>, %lhs: memref<7x14x21xf32>, %rhs: memref<7x14x21xf32>,887 %out: memref<7x14x21xf32>) {888 linalg.select ins(%cond, %lhs, %rhs: memref<7x14x21xi1>, memref<7x14x21xf32>, memref<7x14x21xf32>)889 outs(%out: memref<7x14x21xf32>)890 return891}892 893// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>894 895// CHECK: func @generalize_select896// CHECK-SAME: (%[[COND:.+]]: memref<7x14x21xi1>, %[[LHS:.+]]: memref<7x14x21xf32>, %[[RHS:.+]]: memref<7x14x21xf32>,897// CHECK-SAME: %[[OUT:.+]]: memref<7x14x21xf32>)898 899// CHECK: linalg.generic900// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]], #[[MAP]]]901// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}902// CHECK-SAME: ins(%[[COND]], %[[LHS]], %[[RHS]] : memref<7x14x21xi1>, memref<7x14x21xf32>, memref<7x14x21xf32>)903// CHECK-SAME: outs(%[[OUT]] : memref<7x14x21xf32>)904 905// CHECK: ^{{.+}}(%[[BBARG0:.+]]: i1, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32, %[[BBARG3:.+]]: f32)906// CHECK-NEXT: %[[select:.+]] = arith.select %[[BBARG0]], %[[BBARG1]], %[[BBARG2]] : f32907// CHECK-NEXT: linalg.yield %[[select]] : f32908 909 910// -----911 912// CHECK-LABEL: func @fill_tensor913func.func @fill_tensor(%f: f32, %v: vector<2x4xf32>) -> (tensor<f32>, tensor<vector<2x4xf32>>) {914 %e0 = tensor.empty() : tensor<f32>915 %0 = linalg.fill ins(%f : f32) outs(%e0 : tensor<f32>) -> tensor<f32>916// CHECK: linalg.generic917// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)918// CHECK-NEXT: linalg.yield %[[BBARG0]] : f32919 920 %e1 = tensor.empty() : tensor<vector<2x4xf32>>921 %1 = linalg.fill ins(%v : vector<2x4xf32>) outs(%e1 : tensor<vector<2x4xf32>>) -> tensor<vector<2x4xf32>>922// CHECK: linalg.generic923// CHECK: ^{{.+}}(%[[BBARG0:.+]]: vector<2x4xf32>, %[[BBARG1:.+]]: vector<2x4xf32>)924// CHECK-NEXT: linalg.yield %[[BBARG0]] : vector<2x4xf32>925 926 return %0, %1: tensor<f32>, tensor<vector<2x4xf32>>927}928 929// -----930 931// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>932// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>933// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>934 935// CHECK-LABEL: func.func @matmul_transpose_a_explicit(936// CHECK-SAME: %[[VAL_0:.*]]: memref<5x3xf32>,937// CHECK-SAME: %[[VAL_1:.*]]: memref<5x7xf32>,938// CHECK-SAME: %[[VAL_2:.*]]: memref<3x7xf32>) {939 940// CHECK: linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_2]]], iterator_types = ["parallel", "parallel", "reduction"]}941// CHECK: arith.mulf942// CHECK: arith.addf943 944func.func @matmul_transpose_a_explicit(%arg0: memref<5x3xf32>, %arg1: memref<5x7xf32>, %arg2: memref<3x7xf32>) {945 linalg.matmul indexing_maps = [946 affine_map<(d0, d1, d2) -> (d2, d0)>,947 affine_map<(d0, d1, d2) -> (d2, d1)>,948 affine_map<(d0, d1, d2) -> (d0, d1)>949 ]950 ins(%arg0, %arg1 : memref<5x3xf32>, memref<5x7xf32>)951 outs(%arg2: memref<3x7xf32>)952 return953}954 955// -----956 957// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>958// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>959// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>960// CHECK-LABEL: func.func @matmul_transpose_b_explicit(961// CHECK-SAME: %[[VAL_0:.*]]: memref<3x5xf32>,962// CHECK-SAME: %[[VAL_1:.*]]: memref<7x5xf32>,963// CHECK-SAME: %[[VAL_2:.*]]: memref<3x7xf32>) {964 965// CHECK: linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_2]]], iterator_types = ["parallel", "parallel", "reduction"]}966// CHECK: arith.mulf967// CHECK: arith.addf968 969func.func @matmul_transpose_b_explicit(%arg0: memref<3x5xf32>, %arg1: memref<7x5xf32>, %arg2: memref<3x7xf32>) {970 linalg.matmul indexing_maps = [971 affine_map<(d0, d1, d2) -> (d0, d2)>,972 affine_map<(d0, d1, d2) -> (d1, d2)>,973 affine_map<(d0, d1, d2) -> (d0, d1)>974 ]975 ins(%arg0, %arg1 : memref<3x5xf32>, memref<7x5xf32>)976 outs(%arg2: memref<3x7xf32>)977 return978}979 980// -----981 982// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>983// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>984// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>985 986// CHECK-LABEL: func.func @matmul_transpose_a_b_explicit(987// CHECK-SAME: %[[VAL_0:.*]]: memref<5x3xf32>,988// CHECK-SAME: %[[VAL_1:.*]]: memref<7x5xf32>,989// CHECK-SAME: %[[VAL_2:.*]]: memref<3x7xf32>) {990 991// CHECK: linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_2]]], iterator_types = ["parallel", "parallel", "reduction"]}992// CHECK: arith.mulf993// CHECK: arith.addf994 995func.func @matmul_transpose_a_b_explicit(%arg0: memref<5x3xf32>, %arg1: memref<7x5xf32>, %arg2: memref<3x7xf32>) {996 linalg.matmul indexing_maps = [997 affine_map<(d0, d1, d2) -> (d2, d0)>,998 affine_map<(d0, d1, d2) -> (d1, d2)>,999 affine_map<(d0, d1, d2) -> (d0, d1)>1000 ]1001 ins(%arg0, %arg1 : memref<5x3xf32>, memref<7x5xf32>)1002 outs(%arg2: memref<3x7xf32>)1003 return1004}1005 1006// -----1007 1008// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>1009// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>1010// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>1011 1012// CHECK-LABEL: func.func @batch_matmul(1013// CHECK-SAME: %[[VAL_0:.*]]: tensor<2x3x5xf32>,1014// CHECK-SAME: %[[VAL_1:.*]]: tensor<2x5x7xf32>,1015// CHECK-SAME: %[[VAL_2:.*]]: tensor<2x3x7xf32>) -> tensor<2x3x7xf32> {1016// CHECK: %[[VAL_3:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_0]], #[[$ATTR_1]], #[[$ATTR_2]]], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%[[VAL_0]], %[[VAL_1]] : tensor<2x3x5xf32>, tensor<2x5x7xf32>) outs(%[[VAL_2]] : tensor<2x3x7xf32>) {1017// CHECK: arith.mulf1018// CHECK: arith.addf1019 1020func.func @batch_matmul(%arg0: tensor<2x3x5xf32>, %arg1: tensor<2x5x7xf32>, %arg2: tensor<2x3x7xf32>) -> tensor<2x3x7xf32> {1021 %0 = linalg.batch_matmul indexing_maps = [1022 affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,1023 affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,1024 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>1025 ]1026 ins(%arg0, %arg1: tensor<2x3x5xf32>, tensor<2x5x7xf32>)1027 outs(%arg2: tensor<2x3x7xf32>) -> tensor<2x3x7xf32>1028 return %0 : tensor<2x3x7xf32>1029}1030 1031// -----1032 1033// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>1034// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>1035// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>1036 1037// CHECK-LABEL: func.func @batch_reduce_matmul(1038// CHECK-SAME: %[[A:.*]]: tensor<2x3x5xf32>,1039// CHECK-SAME: %[[B:.*]]: tensor<2x5x7xf32>,1040// CHECK-SAME: %[[C:.*]]: tensor<3x7xf32>) -> tensor<3x7xf32> {1041// CHECK: linalg.generic1042// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_B]], #[[$ACCESS_C]]],1043// CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel", "reduction"]}1044// CHECK: arith.mulf1045// CHECK: arith.addf1046// CHECK: linalg.yield1047 1048func.func @batch_reduce_matmul(%A: tensor<2x3x5xf32>, %B: tensor<2x5x7xf32>, %C: tensor<3x7xf32>) -> tensor<3x7xf32> {1049 %0 = linalg.batch_reduce_matmul indexing_maps = [1050 affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,1051 affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,1052 affine_map<(d0, d1, d2, d3) -> (d1, d2)>1053 ]1054 ins(%A, %B: tensor<2x3x5xf32>, tensor<2x5x7xf32>)1055 outs(%C: tensor<3x7xf32>) -> tensor<3x7xf32>1056 return %0 : tensor<3x7xf32>1057}1058 1059// -----1060 1061// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>1062// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>1063// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>1064 1065// CHECK-LABEL: func.func @contract_matmul(1066// CHECK-SAME: %[[A:.*]]: memref<3x5xf32>,1067// CHECK-SAME: %[[B:.*]]: memref<5x7xf32>,1068// CHECK-SAME: %[[C:.*]]: memref<3x7xf32>) {1069 1070// CHECK: linalg.generic1071// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_B]], #[[$ACCESS_C]]]1072// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]1073// CHECK-NEXT: ^{{.+}}(1074// CHECK-NEXT: arith.mulf1075// CHECK-NEXT: arith.addf1076// CHECK-NEXT: linalg.yield1077 1078func.func @contract_matmul(%A: memref<3x5xf32>, %B: memref<5x7xf32>, %C: memref<3x7xf32>) {1079 linalg.contract1080 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,1081 affine_map<(d0, d1, d2) -> (d2, d1)>,1082 affine_map<(d0, d1, d2) -> (d0, d1)>]1083 ins(%A, %B : memref<3x5xf32>, memref<5x7xf32>)1084 outs(%C: memref<3x7xf32>)1085 1086 return1087}1088 1089// -----1090 1091// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>1092// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>1093// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>1094 1095// CHECK-LABEL: func.func @contract_matmul_transpose_a_b(1096// CHECK-SAME: %[[A:.*]]: memref<5x3xf32>,1097// CHECK-SAME: %[[B:.*]]: memref<7x5xf32>,1098// CHECK-SAME: %[[C:.*]]: memref<3x7xf32>) {1099 1100// CHECK: linalg.generic1101// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_B]], #[[$ACCESS_C]]]1102// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]1103// CHECK-NEXT: ^{{.+}}(1104// CHECK-NEXT: arith.mulf1105// CHECK-NEXT: arith.addf1106// CHECK-NEXT: linalg.yield1107 1108func.func @contract_matmul_transpose_a_b(%A: memref<5x3xf32>, %B: memref<7x5xf32>, %C: memref<3x7xf32>) {1109 linalg.contract1110 indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d0)>,1111 affine_map<(d0, d1, d2) -> (d1, d2)>,1112 affine_map<(d0, d1, d2) -> (d0, d1)>]1113 ins(%A, %B : memref<5x3xf32>, memref<7x5xf32>)1114 outs(%C: memref<3x7xf32>)1115 return1116}1117 1118// -----1119 1120// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>1121// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>1122// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>1123 1124// CHECK-LABEL: func.func @contract_batch_matmul(1125// CHECK-SAME: %[[A:.*]]: memref<9x3x5xf32>,1126// CHECK-SAME: %[[B:.*]]: memref<9x5x7xf32>,1127// CHECK-SAME: %[[C:.*]]: memref<9x3x7xf32>) {1128 1129// CHECK: linalg.generic1130// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_B]], #[[$ACCESS_C]]]1131// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"]1132// CHECK-NEXT: ^{{.+}}(1133// CHECK-NEXT: arith.mulf1134// CHECK-NEXT: arith.addf1135// CHECK-NEXT: linalg.yield1136 1137func.func @contract_batch_matmul(%A: memref<9x3x5xf32>, %B: memref<9x5x7xf32>, %C: memref<9x3x7xf32>) {1138 linalg.contract1139 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,1140 affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>,1141 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]1142 ins(%A, %B : memref<9x3x5xf32>, memref<9x5x7xf32>)1143 outs(%C: memref<9x3x7xf32>)1144 return1145}1146 1147// -----1148 1149// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>1150// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>1151// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>1152 1153// CHECK-LABEL: func.func @contract_batch_reduce_matmul(1154// CHECK-SAME: %[[A:.*]]: memref<9x3x5xf32>,1155// CHECK-SAME: %[[B:.*]]: memref<9x5x7xf32>,1156// CHECK-SAME: %[[C:.*]]: memref<3x7xf32>) {1157 1158// CHECK: linalg.generic1159// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_B]], #[[$ACCESS_C]]]1160// CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel", "reduction"]1161// CHECK-NEXT: ^{{.+}}(1162// CHECK-NEXT: arith.mulf1163// CHECK-NEXT: arith.addf1164// CHECK-NEXT: linalg.yield1165 1166#accessA = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>1167#accessB = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>1168#accessC = affine_map<(d0, d1, d2, d3) -> (d1, d2)>1169func.func @contract_batch_reduce_matmul(1170 %A: memref<9x3x5xf32>, %B: memref<9x5x7xf32>, %C: memref<3x7xf32>) {1171 linalg.contract1172 indexing_maps = [#accessA, #accessB, #accessC]1173 ins(%A, %B : memref<9x3x5xf32>, memref<9x5x7xf32>)1174 outs(%C: memref<3x7xf32>)1175 return1176}1177 1178// -----1179 1180// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>1181// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>1182// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>1183 1184// CHECK-LABEL: func.func @contract_batch_reduce_matmul_permute_m_with_k_and_k_with_n(1185// CHECK-SAME: %[[A:.*]]: memref<9x5x3xf32>,1186// CHECK-SAME: %[[B:.*]]: memref<9x7x5xf32>,1187// CHECK-SAME: %[[C:.*]]: memref<3x7xf32>) {1188 1189// CHECK: linalg.generic1190// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_B]], #[[$ACCESS_C]]]1191// CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel", "reduction"]1192// CHECK-NEXT: ^{{.+}}(1193// CHECK-NEXT: arith.mulf1194// CHECK-NEXT: arith.addf1195// CHECK-NEXT: linalg.yield1196 1197#accessA = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>1198#accessB = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>1199#accessC = affine_map<(d0, d1, d2, d3) -> (d1, d2)>1200func.func @contract_batch_reduce_matmul_permute_m_with_k_and_k_with_n(1201 %A: memref<9x5x3xf32>, %B: memref<9x7x5xf32>, %C: memref<3x7xf32>) {1202 linalg.contract1203 indexing_maps = [#accessA, #accessB, #accessC]1204 ins(%A, %B : memref<9x5x3xf32>, memref<9x7x5xf32>)1205 outs(%C: memref<3x7xf32>)1206 return1207}1208 1209// -----1210 1211// CHECK: #[[$ACCESS_A_B:.+]] = affine_map<(d0) -> (d0)>1212// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0) -> ()>1213 1214// CHECK-LABEL: func.func @contract_dot(1215// CHECK-SAME: %[[A:.*]]: memref<9xf32>, %[[B:.*]]: memref<9xf32>,1216// CHECK-SAME: %[[C:.*]]: memref<f32>) {1217 1218// CHECK: linalg.generic1219// CHECK-SAME: indexing_maps = [#[[$ACCESS_A_B]], #[[$ACCESS_A_B]], #[[$ACCESS_C]]]1220// CHECK-SAME: iterator_types = ["reduction"]1221// CHECK-NEXT: ^{{.+}}(1222// CHECK-NEXT: arith.mulf1223// CHECK-NEXT: arith.addf1224// CHECK-NEXT: linalg.yield1225 1226#accessAB = affine_map<(d0) -> (d0)>1227#accessC = affine_map<(d0) -> ()>1228func.func @contract_dot(1229 %A: memref<9xf32>, %B: memref<9xf32>, %C: memref<f32>) {1230 linalg.contract1231 indexing_maps = [#accessAB, #accessAB, #accessC]1232 ins(%A, %B : memref<9xf32>, memref<9xf32>)1233 outs(%C: memref<f32>)1234 return1235}1236 1237// -----1238 1239// CHECK: #[[$ACCESS_A_B:.+]] = affine_map<(d0, d1, d2) -> (d2)>1240// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>1241 1242// CHECK-LABEL: func.func @contract_matmul_bcast_a_b(1243// CHECK-SAME: %[[A:.*]]: memref<5xf32>, %[[B:.*]]: memref<5xf32>,1244// CHECK-SAME: %[[C:.*]]: memref<3x7xf32>) {1245 1246// CHECK: linalg.generic1247// CHECK-SAME: indexing_maps = [#[[$ACCESS_A_B]], #[[$ACCESS_A_B]], #[[$ACCESS_C]]]1248// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]1249// CHECK-NEXT: ^{{.+}}(1250// CHECK-NEXT: arith.mulf1251// CHECK-NEXT: arith.addf1252// CHECK-NEXT: linalg.yield1253 1254#accessAB = affine_map<(d0, d1, d2) -> (d2)>1255#accessC = affine_map<(d0, d1, d2) -> (d0, d1)>1256func.func @contract_matmul_bcast_a_b(1257 %A: memref<5xf32>, %B: memref<5xf32>, %C: memref<3x7xf32>) {1258 linalg.contract1259 indexing_maps = [#accessAB, #accessAB, #accessC]1260 ins(%A, %B : memref<5xf32>, memref<5xf32>)1261 outs(%C: memref<3x7xf32>)1262 return1263}1264