934 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3from mlir.dialects import arith, func, linalg, tensor, memref, builtin4from mlir.dialects.linalg.opdsl.lang import *5from mlir.extras import types as T6from mlir.ir import *7 8 9def run(f):10 print("\nTEST:", f.__name__)11 f()12 return f13 14 15# CHECK-LABEL: TEST: testFill16@run17def testFill():18 with Context() as ctx, Location.unknown():19 module = Module.create()20 f32 = F32Type.get()21 with InsertionPoint(module.body):22 # CHECK-LABEL: func @fill_tensor23 # CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<12x?xf32>24 # CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f3225 # CHECK-NEXT: %[[RES:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[OUT]] : tensor<12x?xf32>) -> tensor<12x?xf32>26 # CHECK-NEXT: return %[[RES]] : tensor<12x?xf32>27 @func.FuncOp.from_py_func(28 RankedTensorType.get((12, ShapedType.get_dynamic_size()), f32)29 )30 def fill_tensor(out):31 zero = arith.ConstantOp(32 value=FloatAttr.get(f32, 0.0), result=f3233 ).result34 return linalg.fill(zero, outs=[out])35 36 # CHECK-LABEL: func @fill_buffer37 # CHECK-SAME: %[[OUT:[0-9a-z]+]]: memref<12x?xf32>38 # CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f3239 # CHECK-NEXT: linalg.fill ins(%[[CST]] : f32) outs(%[[OUT]] : memref<12x?xf32>)40 # CHECK-NEXT: return41 @func.FuncOp.from_py_func(42 MemRefType.get((12, ShapedType.get_dynamic_size()), f32)43 )44 def fill_buffer(out):45 zero = arith.ConstantOp(46 value=FloatAttr.get(f32, 0.0), result=f3247 ).result48 linalg.fill(zero, outs=[out])49 50 print(module)51 52 53# CHECK-LABEL: TEST: testIdentityRegionOps54@run55def testIdentityRegionOps():56 with Context(), Location.unknown():57 module = Module.create()58 f32 = F32Type.get()59 with InsertionPoint(module.body):60 # CHECK: %[[VAL_0:.*]] = tensor.empty() : tensor<1x13xf32>61 # CHECK: %[[VAL_1:.*]] = tensor.empty() : tensor<13x1xf32>62 op1 = tensor.EmptyOp([1, 13], f32)63 op2 = tensor.EmptyOp([13, 1], f32)64 # CHECK: %[[VAL_2:.*]] = linalg.transpose ins(%[[VAL_0]] : tensor<1x13xf32>) outs(%[[VAL_1]] : tensor<13x1xf32>) permutation = [1, 0]65 op3 = linalg.TransposeOp(66 result=[RankedTensorType.get((13, 1), f32)],67 input=op1,68 init=op2,69 permutation=[1, 0],70 )71 linalg.fill_builtin_region(op3.operation)72 73 # CHECK: %[[VAL_3:.*]] = linalg.transpose ins(%[[VAL_1]] : tensor<13x1xf32>) outs(%[[VAL_0]] : tensor<1x13xf32>) permutation = [1, 0]74 op4 = linalg.transpose(op2, outs=[op1], permutation=[1, 0])75 76 # CHECK: func.func @transpose_op(%[[VAL_4:.*]]: memref<1x13xf32>, %[[VAL_5:.*]]: memref<13x1xf32>)77 @func.FuncOp.from_py_func(78 MemRefType.get((1, 13), f32),79 MemRefType.get((13, 1), f32),80 )81 def transpose_op(op1, op2):82 # CHECK: linalg.transpose ins(%[[VAL_4]] : memref<1x13xf32>) outs(%[[VAL_5]] : memref<13x1xf32>) permutation = [1, 0]83 op3 = linalg.TransposeOp(84 result=[],85 input=op1,86 init=op2,87 permutation=[1, 0],88 )89 linalg.fill_builtin_region(op3.operation)90 # CHECK: linalg.transpose ins(%[[VAL_5]] : memref<13x1xf32>) outs(%[[VAL_4]] : memref<1x13xf32>) permutation = [1, 0]91 op4 = linalg.transpose(op2, outs=[op1], permutation=[1, 0])92 93 # CHECK: %[[VAL_6:.*]] = tensor.empty() : tensor<16xf32>94 # CHECK: %[[VAL_7:.*]] = tensor.empty() : tensor<16x64xf32>95 op1 = tensor.EmptyOp([16], f32)96 op2 = tensor.EmptyOp([16, 64], f32)97 # CHECK: %[[VAL_8:.*]] = linalg.broadcast ins(%[[VAL_6]] : tensor<16xf32>) outs(%[[VAL_7]] : tensor<16x64xf32>) dimensions = [1]98 op3 = linalg.BroadcastOp(99 result=[RankedTensorType.get((16, 64), f32)],100 input=op1,101 init=op2,102 dimensions=[1],103 )104 linalg.fill_builtin_region(op3.operation)105 106 # CHECK: %[[VAL_9:.*]] = tensor.empty() : tensor<64xf32>107 op4 = tensor.EmptyOp([64], f32)108 # CHECK: %[[VAL_10:.*]] = linalg.broadcast ins(%[[VAL_9]] : tensor<64xf32>) outs(%[[VAL_7]] : tensor<16x64xf32>) dimensions = [0]109 op5 = linalg.broadcast(op4, outs=[op2], dimensions=[0])110 111 # CHECK: func.func @broadcast_op(%[[VAL_11:.*]]: memref<16xf32>, %[[VAL_12:.*]]: memref<16x64xf32>, %[[VAL_13:.*]]: memref<64xf32>)112 @func.FuncOp.from_py_func(113 MemRefType.get((16,), f32),114 MemRefType.get((16, 64), f32),115 MemRefType.get((64,), f32),116 )117 def broadcast_op(op1, op2, op3):118 # CHECK: linalg.broadcast ins(%[[VAL_11]] : memref<16xf32>) outs(%[[VAL_12]] : memref<16x64xf32>) dimensions = [1]119 op4 = linalg.BroadcastOp(120 result=[],121 input=op1,122 init=op2,123 dimensions=[1],124 )125 linalg.fill_builtin_region(op4.operation)126 # CHECK: linalg.broadcast ins(%[[VAL_13]] : memref<64xf32>) outs(%[[VAL_12]] : memref<16x64xf32>) dimensions = [0]127 op5 = linalg.broadcast(op3, outs=[op2], dimensions=[0])128 129 print(module)130 131 132# CHECK-LABEL: TEST: testGenericOp133@run134def testGenericOp():135 with Context(), Location.unknown():136 module = Module.create()137 f32 = F32Type.get()138 memref_t = MemRefType.get([10, 10], f32)139 with InsertionPoint(module.body):140 id_map_1 = AffineMap.get_identity(2)141 # CHECK: %[[VAL_0:.*]] = tensor.empty() : tensor<16x16xf32>142 # CHECK: %[[VAL_1:.*]] = tensor.empty() : tensor<16x16xf32>143 x = tensor.empty((16, 16), f32)144 y = tensor.empty((16, 16), f32)145 146 # CHECK: %[[VAL_2:.*]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%[[VAL_0]] : tensor<16x16xf32>) outs(%[[VAL_1]] : tensor<16x16xf32>) {147 # CHECK: ^bb0(%in: f32, %out: f32):148 # CHECK: linalg.yield %in : f32149 # CHECK: } -> tensor<16x16xf32>150 @linalg.generic(151 [x],152 [y],153 [id_map_1, id_map_1],154 [linalg.IteratorType.parallel, linalg.IteratorType.parallel],155 )156 def f(a, b):157 assert isinstance(a, Value)158 assert isinstance(a.type, F32Type)159 assert isinstance(b, Value)160 assert isinstance(b.type, F32Type)161 return a162 163 assert isinstance(f, Value)164 assert isinstance(f.type, RankedTensorType)165 166 # CHECK: %[[VAL_3:.*]] = tensor.empty() : tensor<16x16x16xf32>167 z = tensor.empty((16, 16, 16), f32)168 169 minor_id = AffineMap.get_minor_identity(3, 2)170 id_map_2 = AffineMap.get_identity(3)171 172 # CHECK: %[[VAL_4:.+]]:2 = linalg.generic {indexing_maps = [#map1, #map2, #map2], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[VAL_0]] : tensor<16x16xf32>) outs(%[[VAL_3]], %[[VAL_3]] : tensor<16x16x16xf32>, tensor<16x16x16xf32>) {173 # CHECK: ^bb0(%in: f32, %out: f32, %out_1: f32):174 # CHECK: linalg.yield %in, %out : f32, f32175 # CHECK: } -> (tensor<16x16x16xf32>, tensor<16x16x16xf32>)176 @linalg.generic(177 [x],178 [z, z],179 [minor_id, id_map_2, id_map_2],180 [181 linalg.IteratorType.parallel,182 linalg.IteratorType.parallel,183 linalg.IteratorType.parallel,184 ],185 )186 def g(a, b, c):187 assert isinstance(a, Value)188 assert isinstance(a.type, F32Type)189 assert isinstance(b, Value)190 assert isinstance(b.type, F32Type)191 assert isinstance(c, Value)192 assert isinstance(c.type, F32Type)193 return a, b194 195 assert isinstance(g, OpResultList)196 assert len(g) == 2197 assert isinstance(g[0].type, RankedTensorType)198 assert isinstance(g[1].type, RankedTensorType)199 200 # CHECK: %[[VAL_5:.*]] = memref.alloc() : memref<10x10xf32>201 # CHECK: %[[VAL_6:.*]] = memref.alloc() : memref<10x10xf32>202 xx = memref.alloc(memref_t, [], [])203 yy = memref.alloc(memref_t, [], [])204 205 # CHECK: linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%[[VAL_5]] : memref<10x10xf32>) outs(%[[VAL_6]] : memref<10x10xf32>) {206 # CHECK: ^bb0(%in: f32, %out: f32):207 # CHECK: linalg.yield %in : f32208 # CHECK: }209 @linalg.generic(210 [xx],211 [yy],212 [id_map_1, id_map_1],213 [linalg.IteratorType.parallel, linalg.IteratorType.parallel],214 )215 def f(a, b):216 assert isinstance(a, Value)217 assert isinstance(a.type, F32Type)218 assert isinstance(b, Value)219 assert isinstance(b.type, F32Type)220 return a221 222 module.operation.verify()223 print(module)224 225 226# CHECK-LABEL: TEST: testMatmulOp227@run228def testMatmulOp():229 with Context(), Location.unknown():230 module = Module.create()231 f32 = F32Type.get()232 with InsertionPoint(module.body):233 a_shape = (4, 8)234 b_shape = (8, 12)235 b_transposed_shape = (12, 8)236 c_shape = (4, 12)237 238 dimM = ir.AffineDimExpr.get(0)239 dimN = ir.AffineDimExpr.get(1)240 dimK = ir.AffineDimExpr.get(2)241 242 # CHECK: #[[$A_MAP:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>243 # CHECK: #[[$BTrans_MAP:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>244 # CHECK: #[[$C_MAP:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>245 a_map = ir.AffineMap.get(3, 0, [dimM, dimK])246 b_map = ir.AffineMap.get(3, 0, [dimK, dimN])247 c_map = ir.AffineMap.get(3, 0, [dimM, dimN])248 b_transposed_map = ir.AffineMap.get(3, 0, [dimN, dimK])249 250 # CHECK: func.func @matmul_op(251 @func.FuncOp.from_py_func(252 # CHECK-SAME: %[[A:.*]]: tensor<4x8xf32>,253 RankedTensorType.get(a_shape, f32),254 # CHECK-SAME: %[[Amem:.*]]: memref<4x8xf32>,255 MemRefType.get(a_shape, f32),256 # CHECK-SAME: %[[B:.*]]: tensor<8x12xf32>,257 RankedTensorType.get(b_shape, f32),258 # CHECK-SAME: %[[Bmem:.*]]: memref<8x12xf32>,259 MemRefType.get(b_shape, f32),260 # CHECK-SAME: %[[BTrans:.*]]: tensor<12x8xf32>,261 RankedTensorType.get(b_transposed_shape, f32),262 # CHECK-SAME: %[[BTransmem:.*]]: memref<12x8xf32>,263 MemRefType.get(b_transposed_shape, f32),264 # CHECK-SAME: %[[C:.*]]: tensor<4x12xf32>,265 RankedTensorType.get(c_shape, f32),266 # CHECK-SAME: %[[Cmem:.*]]: memref<4x12xf32>)267 MemRefType.get(c_shape, f32),268 )269 def matmul_op(A, Amem, B, Bmem, Btransposed, Btransposedmem, C, Cmem):270 # CHECK: linalg.matmul ins(%[[A]], %[[B]] : tensor<4x8xf32>, tensor<8x12xf32>) outs(%[[C]] : tensor<4x12xf32>)271 res = linalg.MatmulOp(272 result_tensors=(C.type,),273 inputs=(A, B),274 outputs=(C,),275 )276 linalg.fill_builtin_region(res.operation)277 # CHECK: linalg.matmul ins(%[[A]], %[[B]] : tensor<4x8xf32>, tensor<8x12xf32>) outs(%[[C]] : tensor<4x12xf32>)278 res = linalg.matmul(A, B, outs=(C,))279 280 # CHECK: linalg.matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[BTrans]] : tensor<4x8xf32>, tensor<12x8xf32>) outs(%[[C]] : tensor<4x12xf32>)281 res = linalg.MatmulOp(282 result_tensors=(C.type,),283 inputs=(A, Btransposed),284 outputs=(C,),285 indexing_maps=[a_map, b_transposed_map, c_map],286 )287 linalg.fill_builtin_region(res.operation)288 # CHECK: linalg.matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[BTrans]] : tensor<4x8xf32>, tensor<12x8xf32>) outs(%[[C]] : tensor<4x12xf32>)289 res = linalg.matmul(290 A,291 Btransposed,292 outs=(C,),293 indexing_maps=[a_map, b_transposed_map, c_map],294 )295 296 # And now with memrefs...297 298 # CHECK: linalg.matmul ins(%[[Amem]], %[[Bmem]] : memref<4x8xf32>, memref<8x12xf32>) outs(%[[Cmem]] : memref<4x12xf32>)299 res = linalg.MatmulOp(300 result_tensors=[],301 inputs=(Amem, Bmem),302 outputs=(Cmem,),303 )304 linalg.fill_builtin_region(res.operation)305 # CHECK: linalg.matmul ins(%[[Amem]], %[[Bmem]] : memref<4x8xf32>, memref<8x12xf32>) outs(%[[Cmem]] : memref<4x12xf32>)306 linalg.matmul(Amem, Bmem, outs=(Cmem,))307 308 # CHECK: linalg.matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[BTransmem]] : memref<4x8xf32>, memref<12x8xf32>) outs(%[[Cmem]] : memref<4x12xf32>)309 res = linalg.MatmulOp(310 result_tensors=[],311 inputs=(Amem, Btransposedmem),312 outputs=(Cmem,),313 indexing_maps=[a_map, b_transposed_map, c_map],314 )315 linalg.fill_builtin_region(res.operation)316 # CHECK: linalg.matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[BTransmem]] : memref<4x8xf32>, memref<12x8xf32>) outs(%[[Cmem]] : memref<4x12xf32>)317 linalg.matmul(318 Amem,319 Btransposedmem,320 outs=(Cmem,),321 indexing_maps=[a_map, b_transposed_map, c_map],322 )323 324 print(module)325 326 327# CHECK-LABEL: TEST: testContractOp328@run329def testContractOp():330 with Context(), Location.unknown():331 module = Module.create()332 f32 = F32Type.get()333 with InsertionPoint(module.body):334 a_shape = (4, 8)335 b_shape = (8, 12)336 b_transposed_shape = (12, 8)337 c_shape = (4, 12)338 339 dimM = ir.AffineDimExpr.get(0)340 dimN = ir.AffineDimExpr.get(1)341 dimK = ir.AffineDimExpr.get(2)342 343 # CHECK: #[[$A_MAP:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>344 # CHECK: #[[$B_MAP:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>345 # CHECK: #[[$C_MAP:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>346 # CHECK: #[[$BTrans_MAP:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>347 a_map = ir.AffineMap.get(3, 0, [dimM, dimK])348 b_map = ir.AffineMap.get(3, 0, [dimK, dimN])349 c_map = ir.AffineMap.get(3, 0, [dimM, dimN])350 b_transposed_map = ir.AffineMap.get(3, 0, [dimN, dimK])351 352 # CHECK: func.func @matmul_as_contract_op(353 @func.FuncOp.from_py_func(354 # CHECK-SAME: %[[A:.*]]: tensor<4x8xf32>,355 RankedTensorType.get(a_shape, f32),356 # CHECK-SAME: %[[Amem:.*]]: memref<4x8xf32>,357 MemRefType.get(a_shape, f32),358 # CHECK-SAME: %[[B:.*]]: tensor<8x12xf32>,359 RankedTensorType.get(b_shape, f32),360 # CHECK-SAME: %[[Bmem:.*]]: memref<8x12xf32>,361 MemRefType.get(b_shape, f32),362 # CHECK-SAME: %[[BTrans:.*]]: tensor<12x8xf32>,363 RankedTensorType.get(b_transposed_shape, f32),364 # CHECK-SAME: %[[BTransmem:.*]]: memref<12x8xf32>,365 MemRefType.get(b_transposed_shape, f32),366 # CHECK-SAME: %[[C:.*]]: tensor<4x12xf32>,367 RankedTensorType.get(c_shape, f32),368 # CHECK-SAME: %[[Cmem:.*]]: memref<4x12xf32>)369 MemRefType.get(c_shape, f32),370 )371 def matmul_as_contract_op(372 A, Amem, B, Bmem, Btransposed, Btransposedmem, C, Cmem373 ):374 # CHECK: linalg.contract indexing_maps = [#[[$A_MAP]], #[[$B_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[B]] : tensor<4x8xf32>, tensor<8x12xf32>) outs(%[[C]] : tensor<4x12xf32>)375 op4 = linalg.ContractOp(376 result_tensors=(C.type,),377 inputs=(A, B),378 outputs=(C,),379 indexing_maps=[a_map, b_map, c_map],380 )381 linalg.fill_builtin_region(op4.operation)382 # CHECK: linalg.contract indexing_maps = [#[[$A_MAP]], #[[$B_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[B]] : tensor<4x8xf32>, tensor<8x12xf32>) outs(%[[C]] : tensor<4x12xf32>)383 op5 = linalg.contract(384 A, B, outs=(C,), indexing_maps=[a_map, b_map, c_map]385 )386 387 # CHECK: linalg.contract indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[BTrans]] : tensor<4x8xf32>, tensor<12x8xf32>) outs(%[[C]] : tensor<4x12xf32>)388 op4 = linalg.ContractOp(389 result_tensors=(C.type,),390 inputs=(A, Btransposed),391 outputs=(C,),392 indexing_maps=[a_map, b_transposed_map, c_map],393 )394 linalg.fill_builtin_region(op4.operation)395 # CHECK: linalg.contract indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[BTrans]] : tensor<4x8xf32>, tensor<12x8xf32>) outs(%[[C]] : tensor<4x12xf32>)396 op5 = linalg.contract(397 A,398 Btransposed,399 outs=(C,),400 indexing_maps=[a_map, b_transposed_map, c_map],401 )402 # And now with memrefs...403 404 # CHECK: linalg.contract indexing_maps = [#[[$A_MAP]], #[[$B_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[Bmem]] : memref<4x8xf32>, memref<8x12xf32>) outs(%[[Cmem]] : memref<4x12xf32>)405 op4 = linalg.ContractOp(406 result_tensors=[],407 inputs=(Amem, Bmem),408 outputs=(Cmem,),409 indexing_maps=[a_map, b_map, c_map],410 )411 linalg.fill_builtin_region(op4.operation)412 # CHECK: linalg.contract indexing_maps = [#[[$A_MAP]], #[[$B_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[Bmem]] : memref<4x8xf32>, memref<8x12xf32>) outs(%[[Cmem]] : memref<4x12xf32>)413 linalg.contract(414 Amem, Bmem, outs=(Cmem,), indexing_maps=[a_map, b_map, c_map]415 )416 417 # CHECK: linalg.contract indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[BTransmem]] : memref<4x8xf32>, memref<12x8xf32>) outs(%[[Cmem]] : memref<4x12xf32>)418 op4 = linalg.ContractOp(419 result_tensors=[],420 inputs=(Amem, Btransposedmem),421 outputs=(Cmem,),422 indexing_maps=[a_map, b_transposed_map, c_map],423 )424 linalg.fill_builtin_region(op4.operation)425 # CHECK: linalg.contract indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[BTransmem]] : memref<4x8xf32>, memref<12x8xf32>) outs(%[[Cmem]] : memref<4x12xf32>)426 linalg.contract(427 Amem,428 Btransposedmem,429 outs=(Cmem,),430 indexing_maps=[a_map, b_transposed_map, c_map],431 )432 433 print(module)434 435 436# CHECK-LABEL: TEST: testBatchMatmulOp437@run438def testBatchMatmulOp():439 with Context(), Location.unknown():440 module = Module.create()441 f32 = F32Type.get()442 with InsertionPoint(module.body):443 a_shape = (2, 4, 8)444 b_shape = (2, 8, 12)445 b_transposed_shape = (2, 12, 8)446 c_shape = (2, 4, 12)447 448 dimBatch = ir.AffineDimExpr.get(0)449 dimM = ir.AffineDimExpr.get(1)450 dimN = ir.AffineDimExpr.get(2)451 dimK = ir.AffineDimExpr.get(3)452 453 # CHECK: #[[$A_MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>454 # CHECK: #[[$BTrans_MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>455 # CHECK: #[[$C_MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>456 457 a_map = ir.AffineMap.get(4, 0, [dimBatch, dimM, dimK])458 b_transposed_map = ir.AffineMap.get(4, 0, [dimBatch, dimN, dimK])459 c_map = ir.AffineMap.get(4, 0, [dimBatch, dimM, dimN])460 461 # CHECK: func.func @batch_matmul_op(462 @func.FuncOp.from_py_func(463 # CHECK-SAME: %[[A:.*]]: tensor<2x4x8xf32>,464 RankedTensorType.get(a_shape, f32),465 # CHECK-SAME: %[[Amem:.*]]: memref<2x4x8xf32>,466 MemRefType.get(a_shape, f32),467 # CHECK-SAME: %[[B:.*]]: tensor<2x8x12xf32>,468 RankedTensorType.get(b_shape, f32),469 # CHECK-SAME: %[[Bmem:.*]]: memref<2x8x12xf32>,470 MemRefType.get(b_shape, f32),471 # CHECK-SAME: %[[BTrans:.*]]: tensor<2x12x8xf32>,472 RankedTensorType.get(b_transposed_shape, f32),473 # CHECK-SAME: %[[BTransmem:.*]]: memref<2x12x8xf32>,474 MemRefType.get(b_transposed_shape, f32),475 # CHECK-SAME: %[[C:.*]]: tensor<2x4x12xf32>,476 RankedTensorType.get(c_shape, f32),477 # CHECK-SAME: %[[Cmem:.*]]: memref<2x4x12xf32>)478 MemRefType.get(c_shape, f32),479 )480 def batch_matmul_op(A, Amem, B, Bmem, Btransposed, Btransposedmem, C, Cmem):481 # CHECK: linalg.batch_matmul ins(%[[A]], %[[B]] : tensor<2x4x8xf32>, tensor<2x8x12xf32>) outs(%[[C]] : tensor<2x4x12xf32>)482 res = linalg.BatchMatmulOp(483 result_tensors=(C.type,),484 inputs=(A, B),485 outputs=(C,),486 )487 linalg.fill_builtin_region(res.operation)488 # CHECK: linalg.batch_matmul ins(%[[A]], %[[B]] : tensor<2x4x8xf32>, tensor<2x8x12xf32>) outs(%[[C]] : tensor<2x4x12xf32>)489 res = linalg.batch_matmul(A, B, outs=(C,))490 491 # CHECK: linalg.batch_matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[BTrans]] : tensor<2x4x8xf32>, tensor<2x12x8xf32>) outs(%[[C]] : tensor<2x4x12xf32>)492 res = linalg.BatchMatmulOp(493 result_tensors=(C.type,),494 inputs=(A, Btransposed),495 outputs=(C,),496 indexing_maps=[a_map, b_transposed_map, c_map],497 )498 linalg.fill_builtin_region(res.operation)499 # CHECK: linalg.batch_matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[BTrans]] : tensor<2x4x8xf32>, tensor<2x12x8xf32>) outs(%[[C]] : tensor<2x4x12xf32>)500 res = linalg.batch_matmul(501 A,502 Btransposed,503 outs=(C,),504 indexing_maps=[a_map, b_transposed_map, c_map],505 )506 507 # CHECK: linalg.batch_matmul ins(%[[Amem]], %[[Bmem]] : memref<2x4x8xf32>, memref<2x8x12xf32>) outs(%[[Cmem]] : memref<2x4x12xf32>)508 res = linalg.BatchMatmulOp(509 result_tensors=[],510 inputs=(Amem, Bmem),511 outputs=(Cmem,),512 )513 linalg.fill_builtin_region(res.operation)514 # CHECK: linalg.batch_matmul ins(%[[Amem]], %[[Bmem]] : memref<2x4x8xf32>, memref<2x8x12xf32>) outs(%[[Cmem]] : memref<2x4x12xf32>)515 linalg.batch_matmul(Amem, Bmem, outs=(Cmem,))516 517 # CHECK: linalg.batch_matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[BTransmem]] : memref<2x4x8xf32>, memref<2x12x8xf32>) outs(%[[Cmem]] : memref<2x4x12xf32>)518 res = linalg.BatchMatmulOp(519 result_tensors=[],520 inputs=(Amem, Btransposedmem),521 outputs=(Cmem,),522 indexing_maps=[a_map, b_transposed_map, c_map],523 )524 linalg.fill_builtin_region(res.operation)525 # CHECK: linalg.batch_matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[BTransmem]] : memref<2x4x8xf32>, memref<2x12x8xf32>) outs(%[[Cmem]] : memref<2x4x12xf32>)526 linalg.batch_matmul(527 Amem,528 Btransposedmem,529 outs=(Cmem,),530 indexing_maps=[a_map, b_transposed_map, c_map],531 )532 533 print(module)534 535 536# CHECK-LABEL: TEST: testBatchReduceMatmulOp537@run538def testBatchReduceMatmulOp():539 with Context(), Location.unknown():540 module = Module.create()541 f32 = F32Type.get()542 with InsertionPoint(module.body):543 a_shape = (5, 4, 8)544 b_shape = (5, 8, 12)545 b_transposed_shape = (5, 12, 8)546 c_shape = (4, 12)547 548 dimBatch = ir.AffineDimExpr.get(0)549 dimM = ir.AffineDimExpr.get(1)550 dimN = ir.AffineDimExpr.get(2)551 dimK = ir.AffineDimExpr.get(3)552 553 # CHECK: #[[$A_MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>554 # CHECK: #[[$BTrans_MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>555 # CHECK: #[[$C_MAP:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>556 a_map = ir.AffineMap.get(4, 0, [dimBatch, dimM, dimK])557 b_transposed_map = ir.AffineMap.get(4, 0, [dimBatch, dimN, dimK])558 c_map = ir.AffineMap.get(4, 0, [dimM, dimN])559 560 # CHECK: func.func @batch_reduce_matmul_op(561 @func.FuncOp.from_py_func(562 # CHECK-SAME: %[[A:.*]]: tensor<5x4x8xf32>,563 RankedTensorType.get(a_shape, f32),564 # CHECK-SAME: %[[Amem:.*]]: memref<5x4x8xf32>,565 MemRefType.get(a_shape, f32),566 # CHECK-SAME: %[[B:.*]]: tensor<5x8x12xf32>,567 RankedTensorType.get(b_shape, f32),568 # CHECK-SAME: %[[Bmem:.*]]: memref<5x8x12xf32>,569 MemRefType.get(b_shape, f32),570 # CHECK-SAME: %[[BTrans:.*]]: tensor<5x12x8xf32>,571 RankedTensorType.get(b_transposed_shape, f32),572 # CHECK-SAME: %[[BTransmem:.*]]: memref<5x12x8xf32>,573 MemRefType.get(b_transposed_shape, f32),574 # CHECK-SAME: %[[C:.*]]: tensor<4x12xf32>,575 RankedTensorType.get(c_shape, f32),576 # CHECK-SAME: %[[Cmem:.*]]: memref<4x12xf32>)577 MemRefType.get(c_shape, f32),578 )579 def batch_reduce_matmul_op(580 A, Amem, B, Bmem, Btransposed, Btransposedmem, C, Cmem581 ):582 # CHECK: linalg.batch_reduce_matmul ins(%[[A]], %[[B]] : tensor<5x4x8xf32>, tensor<5x8x12xf32>) outs(%[[C]] : tensor<4x12xf32>)583 res = linalg.BatchReduceMatmulOp(584 result_tensors=(C.type,),585 inputs=(A, B),586 outputs=(C,),587 )588 linalg.fill_builtin_region(res.operation)589 # CHECK: linalg.batch_reduce_matmul ins(%[[A]], %[[B]] : tensor<5x4x8xf32>, tensor<5x8x12xf32>) outs(%[[C]] : tensor<4x12xf32>)590 res = linalg.batch_reduce_matmul(A, B, outs=(C,))591 592 # CHECK: linalg.batch_reduce_matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[BTrans]] : tensor<5x4x8xf32>, tensor<5x12x8xf32>) outs(%[[C]] : tensor<4x12xf32>)593 res = linalg.BatchReduceMatmulOp(594 result_tensors=(C.type,),595 inputs=(A, Btransposed),596 outputs=(C,),597 indexing_maps=[a_map, b_transposed_map, c_map],598 )599 linalg.fill_builtin_region(res.operation)600 # CHECK: linalg.batch_reduce_matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[A]], %[[BTrans]] : tensor<5x4x8xf32>, tensor<5x12x8xf32>) outs(%[[C]] : tensor<4x12xf32>)601 res = linalg.batch_reduce_matmul(602 A,603 Btransposed,604 outs=(C,),605 indexing_maps=[a_map, b_transposed_map, c_map],606 )607 608 # CHECK: linalg.batch_reduce_matmul ins(%[[Amem]], %[[Bmem]] : memref<5x4x8xf32>, memref<5x8x12xf32>) outs(%[[Cmem]] : memref<4x12xf32>)609 res = linalg.BatchReduceMatmulOp(610 result_tensors=[],611 inputs=(Amem, Bmem),612 outputs=(Cmem,),613 )614 linalg.fill_builtin_region(res.operation)615 # CHECK: linalg.batch_reduce_matmul ins(%[[Amem]], %[[Bmem]] : memref<5x4x8xf32>, memref<5x8x12xf32>) outs(%[[Cmem]] : memref<4x12xf32>)616 linalg.batch_reduce_matmul(Amem, Bmem, outs=(Cmem,))617 618 # CHECK: linalg.batch_reduce_matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[BTransmem]] : memref<5x4x8xf32>, memref<5x12x8xf32>) outs(%[[Cmem]] : memref<4x12xf32>)619 res = linalg.BatchReduceMatmulOp(620 result_tensors=[],621 inputs=(Amem, Btransposedmem),622 outputs=(Cmem,),623 indexing_maps=[a_map, b_transposed_map, c_map],624 )625 linalg.fill_builtin_region(res.operation)626 # CHECK: linalg.batch_reduce_matmul indexing_maps = [#[[$A_MAP]], #[[$BTrans_MAP]], #[[$C_MAP]]] ins(%[[Amem]], %[[BTransmem]] : memref<5x4x8xf32>, memref<5x12x8xf32>) outs(%[[Cmem]] : memref<4x12xf32>)627 linalg.batch_reduce_matmul(628 Amem,629 Btransposedmem,630 outs=(Cmem,),631 indexing_maps=[a_map, b_transposed_map, c_map],632 )633 634 print(module)635 636 637# CHECK-LABEL: TEST: testPackUnPackOp638@run639def testPackUnPackOp():640 with Context(), Location.unknown():641 module = Module.create()642 f32 = F32Type.get()643 with InsertionPoint(module.body):644 645 @func.FuncOp.from_py_func(646 RankedTensorType.get((128, 128), f32),647 RankedTensorType.get((16, 16, 8, 8), f32),648 )649 def tensor_pack(src, dst):650 packed = linalg.pack(651 src,652 dst,653 inner_dims_pos=[1, 0],654 inner_tiles=[8, 8],655 padding_value=arith.constant(f32, 0.0),656 )657 658 unpacked = linalg.unpack(659 packed,660 src,661 inner_dims_pos=[0, 1],662 inner_tiles=[8, 8],663 )664 665 return unpacked666 667 # CHECK-LABEL: func.func @tensor_pack(668 # CHECK-SAME: %[[VAL_0:.*]]: tensor<128x128xf32>, %[[VAL_1:.*]]: tensor<16x16x8x8xf32>) -> tensor<128x128xf32> {669 # CHECK: %[[VAL_2:.*]] = arith.constant 0.000000e+00 : f32670 # CHECK: %[[VAL_3:.*]] = linalg.pack %[[VAL_0]] padding_value(%[[VAL_2]] : f32) inner_dims_pos = [1, 0] inner_tiles = [8, 8] into %[[VAL_1]] : tensor<128x128xf32> -> tensor<16x16x8x8xf32>671 # CHECK: %[[VAL_4:.*]] = linalg.unpack %[[VAL_3]] inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %[[VAL_0]] : tensor<16x16x8x8xf32> -> tensor<128x128xf32>672 # CHECK: return %[[VAL_4]] : tensor<128x128xf32>673 # CHECK: }674 print(module)675 676 677# CHECK-LABEL: TEST: testElementwiseOp678@run679def testElementwiseOp():680 with Context(), Location.unknown():681 module = Module.create()682 f32 = F32Type.get()683 with InsertionPoint(module.body):684 rect_shape = (8, 16)685 vert_line_shape = (8,)686 hor_line_shape = (16,)687 transposed_rect_shape = (16, 8)688 689 # CHECK-DAG: #[[$IdentMap2D:.*]] = affine_map<(d0, d1) -> (d0, d1)>690 # CHECK-DAG: #[[$TransMap2D:.*]] = affine_map<(d0, d1) -> (d1, d0)>691 # CHECK-DAG: #[[$VertLineBCastMap:.*]] = affine_map<(d0, d1) -> (d0)>692 # CHECK-DAG: #[[$HorLineBCastMap:.*]] = affine_map<(d0, d1) -> (d1)>693 694 ident_map_2d = AffineMap.get_identity(2)695 transposed_map_2d = AffineMap.get_permutation((1, 0))696 vert_line_bcast_map = AffineMap.get(2, 0, [AffineDimExpr.get(0)])697 hor_line_bcast_map = AffineMap.get(2, 0, [AffineDimExpr.get(1)])698 699 # CHECK: func.func @elementwise_op(700 @func.FuncOp.from_py_func(701 # CHECK-SAME: %[[Rect:.*]]: tensor<8x16xf32>,702 RankedTensorType.get(rect_shape, f32),703 # CHECK-SAME: %[[RectMem:.*]]: memref<8x16xf32>,704 MemRefType.get(rect_shape, f32),705 # CHECK-SAME: %[[VertLine:.*]]: tensor<8xf32>,706 RankedTensorType.get(vert_line_shape, f32),707 # CHECK-SAME: %[[VertLineMem:.*]]: memref<8xf32>,708 MemRefType.get(vert_line_shape, f32),709 # CHECK-SAME: %[[HorLine:.*]]: tensor<16xf32>,710 RankedTensorType.get(hor_line_shape, f32),711 # CHECK-SAME: %[[HorLineMem:.*]]: memref<16xf32>,712 MemRefType.get(hor_line_shape, f32),713 # CHECK-SAME: %[[TransRect:.*]]: tensor<16x8xf32>,714 RankedTensorType.get(transposed_rect_shape, f32),715 # CHECK-SAME: %[[TransRectMem:.*]]: memref<16x8xf32>)716 MemRefType.get(transposed_rect_shape, f32),717 )718 def elementwise_op(719 rect,720 rect_mem,721 vert_line,722 vert_line_mem,723 hor_line,724 hor_line_mem,725 trans_rect,726 trans_rect_mem,727 ):728 # CHECK: %[[OutRect:.*]] = tensor.empty() : tensor<8x16xf32>729 out_rect = tensor.EmptyOp(rect_shape, f32)730 # CHECK: %[[OutRectMem:.*]] = memref.alloca() : memref<8x16xf32>731 out_rect_mem = memref.alloca(MemRefType.get(rect_shape, f32), [], [])732 733 if _inferred_affine_maps := True:734 # CHECK: linalg.elementwise735 # CHECK-SAME: kind=#linalg.elementwise_kind<exp>736 # CHECK-SAME: ins(%[[Rect]] : tensor<8x16xf32>)737 # CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>738 op1 = linalg.ElementwiseOp(739 result_tensors=(out_rect.result.type,),740 inputs=(rect,),741 outputs=(out_rect,),742 kind=linalg.ElementwiseKind.exp,743 )744 linalg.fill_builtin_region(op1.operation)745 746 # CHECK: linalg.elementwise747 # CHECK-SAME: kind=#linalg.elementwise_kind<exp>748 # CHECK-SAME: ins(%[[Rect]] : tensor<8x16xf32>)749 # CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>750 linalg.elementwise(751 rect,752 outs=(out_rect,),753 kind=linalg.ElementwiseKind.exp,754 )755 756 # CHECK: linalg.elementwise757 # CHECK-SAME: kind=#linalg.elementwise_kind<exp>758 # CHECK-SAME: ins(%[[RectMem]] : memref<8x16xf32>)759 # CHECK-SAME: outs(%[[OutRectMem]] : memref<8x16xf32>)760 linalg.elementwise(761 rect_mem,762 outs=(out_rect_mem,),763 kind=linalg.ElementwiseKind.exp,764 )765 766 if _explicit_ident_affine_maps := True:767 # Same as above but with default identity indexing_maps explicitly provided.768 # CHECK: linalg.elementwise769 # CHECK-SAME: kind=#linalg.elementwise_kind<exp>770 # CHECK-SAME: ins(%[[Rect]] : tensor<8x16xf32>)771 # CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>772 op3 = linalg.ElementwiseOp(773 result_tensors=(out_rect.result.type,),774 inputs=(rect,),775 outputs=(out_rect,),776 kind=linalg.ElementwiseKind.exp,777 indexing_maps=[ident_map_2d, ident_map_2d],778 )779 linalg.fill_builtin_region(op3.operation)780 781 # CHECK: linalg.elementwise782 # CHECK-SAME: kind=#linalg.elementwise_kind<exp>783 # CHECK-SAME: ins(%[[RectMem]] : memref<8x16xf32>)784 # CHECK-SAME: outs(%[[OutRectMem]] : memref<8x16xf32>)785 linalg.elementwise(786 rect_mem,787 outs=(out_rect_mem,),788 kind=linalg.ElementwiseKind.exp,789 indexing_maps=[ident_map_2d, ident_map_2d],790 )791 792 if _ops_with_non_ident_input_maps := True:793 # CHECK: linalg.elementwise kind=#linalg.elementwise_kind<exp>794 # CHECK-SAME: indexing_maps = [#[[$VertLineBCastMap]], #[[$IdentMap2D]]]795 # CHECK-SAME: ins(%[[VertLine]] : tensor<8xf32>)796 # CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>797 op4 = linalg.ElementwiseOp(798 result_tensors=(out_rect.result.type,),799 inputs=(vert_line,),800 outputs=(out_rect,),801 kind=linalg.ElementwiseKind.exp,802 indexing_maps=[vert_line_bcast_map, ident_map_2d],803 )804 linalg.fill_builtin_region(op4.operation)805 806 # CHECK: linalg.elementwise kind=#linalg.elementwise_kind<add>807 # CHECK-SAME: indexing_maps = [#[[$IdentMap2D]], #[[$VertLineBCastMap]], #[[$IdentMap2D]]]808 # CHECK-SAME: ins(%[[Rect]], %[[VertLine]] : tensor<8x16xf32>, tensor<8xf32>)809 # CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>810 op4 = linalg.ElementwiseOp(811 result_tensors=(out_rect.result.type,),812 inputs=(rect, vert_line),813 outputs=(out_rect,),814 kind=linalg.ElementwiseKind.add,815 indexing_maps=[ident_map_2d, vert_line_bcast_map, ident_map_2d],816 )817 linalg.fill_builtin_region(op4.operation)818 819 # CHECK: linalg.elementwise kind=#linalg.elementwise_kind<div>820 # CHECK-SAME: indexing_maps = [#[[$VertLineBCastMap]], #[[$HorLineBCastMap]], #[[$IdentMap2D]]]821 # CHECK-SAME: ins(%[[VertLine]], %[[HorLine]] : tensor<8xf32>, tensor<16xf32>)822 # CHECK-SAME: outs(%[[OutRect]] : tensor<8x16xf32>) -> tensor<8x16xf32>823 linalg.elementwise(824 vert_line,825 hor_line,826 outs=(out_rect,),827 kind=linalg.ElementwiseKind.div,828 indexing_maps=[829 vert_line_bcast_map,830 hor_line_bcast_map,831 ident_map_2d,832 ],833 )834 835 if _ops_with_non_ident_and_transposed_input_maps := True:836 # CHECK: %[[VertLineBoolsMem:.*]] = memref.alloca() : memref<8xi1>837 vert_line_bools_mem = memref.alloca(838 MemRefType.get(vert_line_shape, IntegerType.get_signless(1)),839 [],840 [],841 )842 # CHECK: linalg.elementwise kind=#linalg.elementwise_kind<select>843 # CHECK-SAME: indexing_maps = [#[[$VertLineBCastMap]], #[[$HorLineBCastMap]], #[[$TransMap2D]], #[[$IdentMap2D]]]844 # CHECK-SAME: ins(%[[VertLineBoolsMem]], %[[HorLineMem]], %[[TransRectMem]] : memref<8xi1>, memref<16xf32>, memref<16x8xf32>)845 # CHECK-SAME: outs(%[[OutRectMem]] : memref<8x16xf32>)846 linalg.elementwise(847 vert_line_bools_mem,848 hor_line_mem,849 trans_rect_mem,850 outs=(out_rect_mem,),851 kind=linalg.ElementwiseKind.select,852 indexing_maps=[853 vert_line_bcast_map,854 hor_line_bcast_map,855 transposed_map_2d,856 ident_map_2d,857 ],858 )859 860 print(module)861 862 863@run864def testReduceOp():865 with Context(), Location.unknown():866 f32 = T.f32()867 tensor_type = T.tensor(10, f32)868 869 @builtin.module870 def module():871 @func.func(tensor_type)872 def reduce_op(input):873 c1 = arith.constant(f32, 1.0)874 single_result = ir.RankedTensorType.get((), f32)875 dims = ir.DenseI64ArrayAttr.get([0])876 init = tensor.splat(single_result, c1, [])877 878 @linalg.reduce(879 result=[single_result],880 inputs=[input],881 inits=[init],882 dimensions=dims,883 )884 def reduced(element: f32, acc: f32):885 return arith.mulf(acc, element)886 887 return tensor.extract(reduced, [])888 889 print(module)890 891 892# CHECK-LABEL: func.func @reduce_op(893# CHECK-SAME: %[[ARG0:.*]]: tensor<10xf32>) -> f32 {894# CHECK: %[[CONSTANT_0:.*]] = arith.constant 1.000000e+00 : f32895# CHECK: %[[SPLAT_0:.*]] = tensor.splat %[[CONSTANT_0]] : tensor<f32>896# CHECK: %[[REDUCE_0:.*]] = linalg.reduce { arith.mulf } ins(%[[ARG0]] : tensor<10xf32>) outs(%[[SPLAT_0]] : tensor<f32>) dimensions = [0]897# CHECK: %[[EXTRACT_0:.*]] = tensor.extract %[[REDUCE_0]][] : tensor<f32>898# CHECK: return %[[EXTRACT_0]] : f32899# CHECK: }900 901 902@run903def testMapOp():904 with Context(), Location.unknown():905 f32 = T.f32()906 tensor_type = T.tensor(10, f32)907 908 @builtin.module909 def module():910 @func.func(tensor_type)911 def map_op(input):912 empty = tensor.empty(tensor_type.shape, f32)913 914 @linalg.map(915 result=[tensor_type],916 inputs=[input, input],917 init=empty,918 )919 def add(element: f32, acc: f32, init: f32):920 return arith.addf(element, acc)921 922 return add923 924 module.verify()925 print(module)926 927 928# CHECK-LABEL: func.func @map_op(929# CHECK-SAME: %[[ARG0:.*]]: tensor<10xf32>) -> tensor<10xf32> {930# CHECK: %[[EMPTY_0:.*]] = tensor.empty() : tensor<10xf32>931# CHECK: %[[MAP_0:.*]] = linalg.map { arith.addf } ins(%[[ARG0]], %[[ARG0]] : tensor<10xf32>, tensor<10xf32>) outs(%[[EMPTY_0]] : tensor<10xf32>)932# CHECK: return %[[MAP_0]] : tensor<10xf32>933# CHECK: }934