1// RUN: mlir-opt -linalg-bufferize -buffer-hoisting -buffer-deallocation -split-input-file %s | FileCheck %s 2 3#map0 = affine_map<(d0) -> (d0)> 4 5// CHECK-LABEL: func @multiple_results 6func @multiple_results(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { 7 %0, %1 = linalg.generic { 8 indexing_maps = [#map0, #map0, #map0], 9 iterator_types = ["parallel"] 10 } ins(%arg0 : tensor<4xf32>) { 11 ^bb0(%gen_arg1: f32): 12 %tmp1 = exp %gen_arg1 : f32 13 linalg.yield %tmp1, %tmp1 : f32, f32 14 } -> tensor<4xf32>, tensor<4xf32> 15 return %0, %1 : tensor<4xf32>, tensor<4xf32> 16} 17// CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]], %[[ARG1_RESULT:.*]]: [[TYPE]], %[[ARG2_RESULT:.*]]: [[TYPE]]) 18// CHECK: %[[FIRST_ALLOC:.*]] = alloc() : [[TYPE]] 19// CHECK: %[[SECOND_ALLOC:.*]] = alloc() : [[TYPE]] 20// CHECK: linalg.generic 21// CHECK-SAME: ins(%[[NEW_ARG0]] : [[TYPE]] 22// CHECK-SAME: outs(%[[FIRST_ALLOC]], %[[SECOND_ALLOC]] : [[TYPE]], [[TYPE]] 23// CHECK-NEXT: ^{{[a-z0-9_]*}} 24// CHECK-SAME: %{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32 25// CHECK-NEXT: %{{.*}} = exp 26// CHECK-NEXT: linalg.yield 27// CHECK: linalg.copy(%[[FIRST_ALLOC]], %[[ARG1_RESULT]]) 28// CHECK: dealloc %[[FIRST_ALLOC]] 29// CHECK: linalg.copy(%[[SECOND_ALLOC]], %[[ARG2_RESULT]]) 30// CHECK: dealloc %[[SECOND_ALLOC]] 31// CHECK: return 32 33// ----- 34 35#map0 = affine_map<(d0) -> (d0)> 36 37// CHECK-LABEL: func @chained_operations 38func @chained_operations(%arg0: tensor<4xf32>) -> tensor<4xf32> { 39 %0 = linalg.generic { 40 indexing_maps = [#map0, #map0], 41 iterator_types = ["parallel"] 42 } ins(%arg0 : tensor<4xf32>) { 43 ^bb0(%gen_arg1: f32): 44 %tmp1 = exp %gen_arg1 : f32 45 linalg.yield %tmp1 : f32 46 } -> tensor<4xf32> 47 48 %1 = linalg.generic { 49 indexing_maps = [#map0, #map0], 50 iterator_types = ["parallel"] 51 } ins(%0 : tensor<4xf32>) { 52 ^bb0(%gen_arg2: f32): 53 %tmp2 = exp %gen_arg2 : f32 54 linalg.yield %tmp2 : f32 55 } -> tensor<4xf32> 56 return %1 : tensor<4xf32> 57} 58// CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]], %[[ARG1_RESULT:.*]]: [[TYPE]]) 59// CHECK: %[[FIRST_ALLOC:.*]] = alloc() : [[TYPE]] 60// CHECK: linalg.generic 61// CHECK-SAME: ins(%[[NEW_ARG0]] : [[TYPE]] 62// CHECK-SAME: outs(%[[FIRST_ALLOC]] : [[TYPE]] 63// CHECK: ^{{[a-z0-9_]*}} 64// CHECK-SAME: %{{.*}}: f32, %{{.*}}: f32 65// CHECK: %[[SECOND_ALLOC:.*]] = alloc() : [[TYPE]] 66// CHECK: linalg.generic 67// CHECK-SAME: ins(%[[FIRST_ALLOC]] : [[TYPE]] 68// CHECK-SAME: outs(%[[SECOND_ALLOC]] : [[TYPE]] 69// CHECK: ^{{[a-z0-9_]*}} 70// CHECK-SAME: %{{.*}}: f32, %{{.*}}: f32 71// CHECK: dealloc %[[FIRST_ALLOC]] 72// CHECK: linalg.copy(%[[SECOND_ALLOC]], %[[ARG1_RESULT]]) 73// CHECK: dealloc %[[SECOND_ALLOC]] 74// CHECK: return 75 76// ----- 77 78// CHECK-LABEL: func @no_linalg_op 79func @no_linalg_op(%arg0: f32) -> (f32, f32) { 80 %0 = mulf %arg0, %arg0 : f32 81 return %0, %0 : f32, f32 82} 83// CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]]) -> ([[TYPE]], [[TYPE]]) 84// CHECK: %[[RESULT:.*]] = mulf %[[NEW_ARG0]], %[[NEW_ARG0]] : [[TYPE]] 85// CHECK: return %[[RESULT]], %[[RESULT]] : [[TYPE]], [[TYPE]] 86 87// ----- 88 89#map_2d = affine_map<(d0, d1) -> (d0, d1)> 90#map_2d_inv = affine_map<(d0, d1) -> (d1, d0)> 91 92func @dynamic_results(%arg0: tensor<?x?xf32>) 93 -> (tensor<?x?xf32>, tensor<?x?xf32>) { 94 %0, %1 = linalg.generic { 95 indexing_maps = [#map_2d, #map_2d, #map_2d_inv], 96 iterator_types = ["parallel", "parallel"] 97 } ins(%arg0 : tensor<?x?xf32>) { 98 ^bb0(%gen_arg1: f32): 99 %tmp1 = exp %gen_arg1 : f32 100 linalg.yield %tmp1, %tmp1 : f32, f32 101 } -> tensor<?x?xf32>, tensor<?x?xf32> 102 return %0, %1 : tensor<?x?xf32>, tensor<?x?xf32> 103} 104 105// CHECK: #map0 = affine_map<(d0, d1) -> (d0, d1)> 106// CHECK: #map1 = affine_map<(d0, d1) -> (d1, d0)> 107 108// CHECK-LABEL: func @dynamic_results 109// CHECK-SAME: (%[[INPUT:.*]]: [[TYPE:.*]], %[[OUT_1:.*]]: [[TYPE]], %[[OUT_2:.*]]: [[TYPE]]) { 110// CHECK: %[[C0:.*]] = constant 0 : index 111// CHECK: %[[DIM_0:.*]] = dim %[[INPUT]], %[[C0]] : [[TYPE]] 112// CHECK: %[[C1:.*]] = constant 1 : index 113// CHECK: %[[DIM_1:.*]] = dim %[[INPUT]], %[[C1]] : [[TYPE]] 114// CHECK: %[[OUT_BUF_1:.*]] = alloc(%[[DIM_0]], %[[DIM_1]]) : [[TYPE]] 115// CHECK: %[[OUT_BUF_2:.*]] = alloc(%[[DIM_1]], %[[DIM_0]]) : [[TYPE]] 116 117// CHECK: linalg.generic {indexing_maps = [#map0, #map0, #map1], {{.*}}} 118// CHECK-SAME: ins(%[[INPUT]] : [[TYPE]]) 119// CHECK-SAME: outs(%[[OUT_BUF_1]], %[[OUT_BUF_2]] : [[TYPE]], [[TYPE]]) { 120 121// CHECK: linalg.copy(%[[OUT_BUF_1]], %[[OUT_1]]) : [[TYPE]], [[TYPE]] 122// CHECK: dealloc %[[OUT_BUF_1]] : [[TYPE]] 123// CHECK: linalg.copy(%[[OUT_BUF_2]], %[[OUT_2]]) : [[TYPE]], [[TYPE]] 124// CHECK: dealloc %[[OUT_BUF_2]] : [[TYPE]] 125// CHECK: return 126 127// ----- 128 129func @foo() -> tensor<2x3xf32> { 130// CHECK-LABEL: func @foo( 131// CHECK-SAME: %[[A:[0-9a-z]*]]: memref<2x3xf32>) { 132 133 %0 = constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf32> 134// CHECK-NEXT: %[[ALLOC:.*]] = alloc() : memref<vector<6xf32>> 135// CHECK-NEXT: %[[CST:.*]] = constant dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00]> : vector<6xf32> 136// CHECK-NEXT: store %[[CST]], %[[ALLOC]][] : memref<vector<6xf32>> 137// CHECK-NEXT: %[[FLAT:.*]] = vector.type_cast %[[ALLOC]] : memref<vector<6xf32>> to memref<6xf32> 138// CHECK-NEXT: %[[RES:.*]] = linalg.reshape %[[FLAT]] {{.*}} : memref<6xf32> into memref<2x3xf32> 139 140 return %0 : tensor<2x3xf32> 141// CHECK-NEXT: linalg.copy(%[[RES]], %[[A]]) : memref<2x3xf32>, memref<2x3xf32> 142// CHECK-NEXT: dealloc %[[ALLOC]] : memref<vector<6xf32>> 143// CHECK-NEXT: return 144} 145 146func @bar() { 147// CHECK-LABEL: func @bar() { 148 149 %0 = call @foo() : () -> tensor<2x3xf32> 150// CHECK-NEXT: %[[ALLOC:.*]] = alloc() : memref<2x3xf32> 151// CHECK-NEXT: call @foo(%[[ALLOC]]) : (memref<2x3xf32>) -> () 152 153 // Instead of relying on tensor_store which introduces aliasing, we rely on 154 // the conversion of print_memref_f32(tensor<*xf32>) to 155 // print_memref_f32(memref<*xf32>). 156 // Note that this is skipping a step and we would need at least some function 157 // attribute to declare that this conversion is valid (e.g. when we statically 158 // know that things will play nicely at the C ABI boundary). 159 %unranked = tensor_cast %0 : tensor<2x3xf32> to tensor<*xf32> 160// CHECK-NEXT: %[[UNRANKED:.*]] = memref_cast %[[ALLOC]] : 161// CHECK-SAME: memref<2x3xf32> to memref<*xf32> 162 163 call @print_memref_f32(%unranked) : (tensor<*xf32>) -> () 164// CHECK-NEXT: call @print_memref_f32(%[[UNRANKED]]) : (memref<*xf32>) -> () 165 166 return 167// CHECK-NEXT: dealloc %[[ALLOC]] : memref<2x3xf32> 168// CHECK-NEXT: return 169} 170 171// This gets converted to a function operating on memref<*xf32>. 172// Note that this is skipping a step and we would need at least some function 173// attribute to declare that this conversion is valid (e.g. when we statically 174// know that things will play nicely at the C ABI boundary). 175func @print_memref_f32(%ptr : tensor<*xf32>) 176// CHECK-LABEL: func @print_memref_f32(memref<*xf32>) 177 178// ----- 179 180#accesses = [ 181 affine_map<(i, j, k) -> (j, i, k)>, 182 affine_map<(i, j, k) -> (i, j)> 183] 184 185#trait = { 186 indexing_maps = #accesses, 187 iterator_types = ["parallel", "parallel", "reduction"] 188} 189 190func @generic_with_init_tensor(%arg0: tensor<2x3x4xvector<3x4xi4>>, 191 %arg1: tensor<3x2xf32>) -> (tensor<3x2xf32>) { 192 193 %0 = linalg.generic #trait 194 ins(%arg0 : tensor<2x3x4xvector<3x4xi4>>) 195 init(%arg1 : tensor<3x2xf32>) { 196 ^bb(%v0: vector<3x4xi4>, %v1: f32) : 197 %f0 = constant 0.0 : f32 198 linalg.yield %f0 : f32 199 } -> tensor<3x2xf32> 200 201 return %0 : tensor<3x2xf32> 202} 203// CHECK-LABEL: func @generic_with_init_tensor 204// CHECK-SAME: (%[[ARG0:.*]]: memref<2x3x4xvector<3x4xi4>>, %[[ARG1:.*]]: memref<3x2xf32>, %[[RESULT0:.*]]: memref<3x2xf32>) { 205// CHECK-NEXT: linalg.generic 206// CHECK: linalg.copy(%[[ARG1]], %[[RESULT0]]) 207// CHECK-NEXT: return 208// CHECK-NOT: % 209 210// ----- 211 212#accesses = [ 213 affine_map<(i, j, k) -> (j, i, k)>, 214 affine_map<(i, j, k) -> (i, j)> 215] 216 217#trait = { 218 indexing_maps = #accesses, 219 iterator_types = ["parallel", "parallel", "reduction"] 220} 221 222func @init_tensor_with_2_uses(%arg0: tensor<2x3x4xvector<3x4xi4>>, 223 %arg1: tensor<3x2xf32>) -> (tensor<3x2xf32>, tensor<3x2xf32>) { 224 225 %0 = linalg.generic #trait 226 ins(%arg0 : tensor<2x3x4xvector<3x4xi4>>) 227 init(%arg1 : tensor<3x2xf32>) { 228 ^bb(%v0: vector<3x4xi4>, %v1: f32) : 229 %f0 = constant 0.0 : f32 230 linalg.yield %f0 : f32 231 } -> tensor<3x2xf32> 232 233 %1 = linalg.generic #trait 234 ins(%arg0 : tensor<2x3x4xvector<3x4xi4>>) 235 init(%arg1 : tensor<3x2xf32>) { 236 ^bb(%v0: vector<3x4xi4>, %v1: f32) : 237 %f0 = constant 0.0 : f32 238 linalg.yield %f0 : f32 239 } -> tensor<3x2xf32> 240 241 return %0, %1 : tensor<3x2xf32>, tensor<3x2xf32> 242} 243// CHECK-LABEL: func @init_tensor_with_2_uses 244// CHECK-SAME: (%[[ARG0:.*]]: memref<2x3x4xvector<3x4xi4>>, %[[ARG1:.*]]: memref<3x2xf32>, %[[RESULT0:.*]]: memref<3x2xf32>, %[[RESULT1:.*]]: memref<3x2xf32>) { 245// CHECK-NEXT: %[[ALLOC0:.*]] = alloc 246// CHECK-NEXT: linalg.copy(%[[ARG1]], %[[ALLOC0]]) 247// CHECK-NEXT: linalg.generic 248// CHECK-SAME: outs(%[[ALLOC0]] 249// CHECK-NEXT: ^bb 250// CHECK-NEXT: constant 251// CHECK-NEXT: yield 252// CHECK-NEXT: } 253// CHECK-NEXT: %[[ALLOC1:.*]] = alloc 254// CHECK-NEXT: linalg.copy(%[[ARG1]], %[[ALLOC1]]) 255// CHECK-NEXT: linalg.generic 256// CHECK-SAME: outs(%[[ALLOC1]] 257// CHECK-NEXT: ^bb 258// CHECK-NEXT: constant 259// CHECK-NEXT: yield 260// CHECK-NEXT: } 261// CHECK-NEXT: linalg.copy(%[[ALLOC0]], %[[RESULT0]]) 262// CHECK-NEXT: dealloc 263// CHECK-NEXT: linalg.copy(%[[ALLOC1]], %[[RESULT1]]) 264// CHECK-NEXT: dealloc 265// CHECK-NEXT: return 266// CHECK-NOT: % 267 268// ----- 269 270#accesses = [ 271 affine_map<(i, j, k) -> (j, i, k)>, 272 affine_map<(i, j, k) -> (i, j)> 273] 274 275#trait = { 276 indexing_maps = #accesses, 277 iterator_types = ["parallel", "parallel", "reduction"] 278} 279 280func @init_tensor_with_1_use_def_chain(%arg0: tensor<2x3x4xvector<3x4xi4>>, 281 %arg1: tensor<3x2xf32>) -> (tensor<3x2xf32>) { 282 283 %0 = linalg.generic #trait 284 ins(%arg0 : tensor<2x3x4xvector<3x4xi4>>) 285 init(%arg1 : tensor<3x2xf32>) { 286 ^bb(%v0: vector<3x4xi4>, %v1: f32) : 287 %f0 = constant 0.0 : f32 288 linalg.yield %f0 : f32 289 } -> tensor<3x2xf32> 290 291 %1 = linalg.generic #trait 292 ins(%arg0 : tensor<2x3x4xvector<3x4xi4>>) 293 init(%0 : tensor<3x2xf32>) { 294 ^bb(%v0: vector<3x4xi4>, %v1: f32) : 295 %f0 = constant 0.0 : f32 296 linalg.yield %f0 : f32 297 } -> tensor<3x2xf32> 298 299 return %1 : tensor<3x2xf32> 300} 301// CHECK-LABEL: func @init_tensor_with_1_use_def_chain 302// CHECK-SAME: (%[[ARG0:.*]]: memref<2x3x4xvector<3x4xi4>>, %[[ARG1:.*]]: memref<3x2xf32>, %[[RESULT0:.*]]: memref<3x2xf32>) { 303// CHECK-NEXT: linalg.generic 304// CHECK-NEXT: ^bb 305// CHECK-NEXT: constant 306// CHECK-NEXT: yield 307// CHECK-NEXT: } 308// CHECK-NEXT: linalg.generic 309// CHECK-NEXT: ^bb 310// CHECK-NEXT: constant 311// CHECK-NEXT: yield 312// CHECK-NEXT: } 313// CHECK-NEXT: linalg.copy(%[[ARG1]], %[[RESULT0]]) 314// CHECK-NEXT: return 315// CHECK-NOT: % 316