1// RUN: mlir-opt %s -test-tensor-copy-insertion | FileCheck %s 2// RUN: mlir-opt %s -test-tensor-copy-insertion="bufferize-function-boundaries" | FileCheck %s --check-prefix=CHECK-FUNC 3 4#DCSR = #sparse_tensor.encoding<{ 5 map = (d0, d1) -> (d0 : compressed, d1 : compressed) 6}> 7 8// CHECK-LABEL: func @bufferization_alloc_tensor 9// CHECK-FUNC-LABEL: func @bufferization_alloc_tensor 10func.func @bufferization_alloc_tensor() -> tensor<20x40xf32, #DCSR> { 11 // CHECK: bufferization.alloc_tensor() 12 // CHECK-FUNC: bufferization.alloc_tensor() 13 %0 = bufferization.alloc_tensor() : tensor<20x40xf32, #DCSR> 14 %1 = sparse_tensor.load %0 : tensor<20x40xf32, #DCSR> 15 return %1 : tensor<20x40xf32, #DCSR> 16} 17 18!Filename = !llvm.ptr 19// CHECK-LABEL: func @sparse_tensor_new 20// CHECK-FUNC-LABEL: func @sparse_tensor_new 21func.func @sparse_tensor_new(%file: !Filename) -> tensor<20x40xf32, #DCSR> { 22 // CHECK: sparse_tensor.new {{.*}} 23 // CHECK-FUNC: sparse_tensor.new {{.*}} 24 %0 = sparse_tensor.new %file : !Filename to tensor<20x40xf32, #DCSR> 25 return %0 : tensor<20x40xf32, #DCSR> 26} 27 28// CHECK-LABEL: func @sparse_tensor_convert 29// CHECK-FUNC-LABEL: func @sparse_tensor_convert 30func.func @sparse_tensor_convert() -> tensor<20x40xf32> { 31 // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() 32 // CHECK-FUNC: %[[alloc:.*]] = bufferization.alloc_tensor() 33 %0 = bufferization.alloc_tensor() : tensor<20x40xf32, #DCSR> 34 // CHECK: %[[loaded:.*]] = sparse_tensor.load %[[alloc]] 35 // CHECK-FUNC: %[[loaded:.*]] = sparse_tensor.load %[[alloc]] 36 %1 = sparse_tensor.load %0 : tensor<20x40xf32, #DCSR> 37 // CHECK: sparse_tensor.convert %[[loaded]] 38 // CHECK-FUNC: sparse_tensor.convert %[[loaded]] 39 %2 = sparse_tensor.convert %1 : tensor<20x40xf32, #DCSR> to tensor<20x40xf32> 40 return %2 : tensor<20x40xf32> 41} 42 43#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }> 44 45#trait = { 46 indexing_maps = [ 47 affine_map<(i) -> (i)>, // A (in) 48 affine_map<(i) -> (i)> // X (out) 49 ], 50 iterator_types = ["parallel"] 51} 52 53// CHECK-LABEL: func @update_notinplace( 54// CHECK-SAME: %[[argb:.*]]: tensor<10xf32> 55// CHECK-FUNC-LABEL: func @update_notinplace( 56// CHECK-FUNC-SAME: %[[argb:.*]]: tensor<10xf32> 57func.func @update_notinplace(%argb: tensor<10xf32>, %arga: tensor<10xf32, #SV>) 58 -> (tensor<10xf32>, tensor<10xf32>) 59{ 60 // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[argb]]) : tensor<10xf32> 61 // CHECK: linalg.generic {{.*}} outs(%[[alloc]] 62 // CHECK-FUNC: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[argb]]) : tensor<10xf32> 63 // CHECK-FUNC: linalg.generic {{.*}} outs(%[[alloc]] 64 %0 = linalg.generic #trait 65 ins(%arga: tensor<10xf32, #SV>) 66 outs(%argb: tensor<10xf32>) { 67 ^bb(%a: f32, %x : f32): 68 %up = arith.addf %a, %x : f32 69 linalg.yield %up : f32 70 } -> tensor<10xf32> 71 return %0, %argb : tensor<10xf32>, tensor<10xf32> 72} 73 74#map = affine_map<(d0, d1) -> (d0, d1)> 75#map1 = affine_map<(d0, d1, d2) -> (d0, d2)> 76#map2 = affine_map<(d0, d1, d2) -> (d2, d1)> 77#map3 = affine_map<(d0, d1, d2) -> (d0, d1)> 78#sparse = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed), posWidth = 64, crdWidth = 64 }> 79 80// linalg.generic with sparse tensors does not necessarily bufferize to 81// element-wise access into the underlying sparse data structures. 82 83// CHECK-LABEL: func @sparse_non_elementwise( 84func.func @sparse_non_elementwise(%arg0: tensor<64x64xf32, #sparse>, %arg1: tensor<64x64xf32>, %arg2: tensor<64x64xf32>) -> tensor<64x64xf32> { 85 %cst = arith.constant 0.000000e+00 : f32 86 // CHECK: %[[alloc0:.*]] = bufferization.alloc_tensor() 87 // CHECK: %[[alloc1:.*]] = bufferization.alloc_tensor() 88 %0 = bufferization.alloc_tensor() : tensor<64x64xf32> 89 // CHECK: %[[generic0:.*]] = linalg.generic {{.*}} outs(%[[alloc1]] : {{.*}}) 90 %1 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel"]} outs(%0 : tensor<64x64xf32>) { 91 ^bb0(%out: f32): 92 linalg.yield %cst : f32 93 } -> tensor<64x64xf32> 94 // CHECK: linalg.generic {{.*}} outs(%[[generic0]] : {{.*}}) 95 %2 = linalg.generic {indexing_maps = [#map1, #map2, #map3], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg2, %arg2 : tensor<64x64xf32>, tensor<64x64xf32>) outs(%1 : tensor<64x64xf32>) { 96 ^bb0(%in: f32, %in_0: f32, %out: f32): 97 %4 = arith.mulf %in, %in_0 : f32 98 %5 = arith.addf %out, %4 : f32 99 linalg.yield %5 : f32 100 } -> tensor<64x64xf32> 101 // CHECK: linalg.generic {{.*}} outs(%[[alloc0]] : {{.*}}) 102 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %2 : tensor<64x64xf32, #sparse>, tensor<64x64xf32>) outs(%0 : tensor<64x64xf32>) attrs = {sorted = true} { 103 ^bb0(%in: f32, %in_0: f32, %out: f32): 104 %4 = arith.mulf %in, %in_0 : f32 105 linalg.yield %4 : f32 106 } -> tensor<64x64xf32> 107 return %3 : tensor<64x64xf32> 108} 109