155 lines · plain
1// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py2// RUN: mlir-opt %s --sparse-reinterpret-map -sparsification | FileCheck %s3 4// Test to demonstrate the difference between non-annotated dense tensors5// and all-dense-annotated "sparse" tensors. The former class remains as6// two-dimensional tensors that are bufferized by subsequent passes. The7// latter class is linearized into one-dimensional buffers that are backed8// by the runtime support library.9 10#DenseMatrix = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : dense) }>11 12#trait_2d = {13 indexing_maps = [14 affine_map<(i,j) -> (i,j)>, // A15 affine_map<(i,j) -> (i,j)> // X (out)16 ],17 iterator_types = ["parallel", "parallel"],18 doc = "X(i,j) = A(i,j) + 1"19}20 21#trait_3d = {22 indexing_maps = [23 affine_map<(i,j,k) -> (i,j,k)>, // A24 affine_map<(i,j,k) -> (i,j)> // X (out)25 ],26 iterator_types = ["parallel", "parallel", "reduction"],27 doc = "X(i,j) += A(i,j,k)"28}29 30//31// Test with an all-dense-annotated "sparse" matrix as input and32// a non-annotated dense matrix as output.33//34// CHECK-LABEL: func @dense1(35// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse{{[0-9]*}}>,36// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {37// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1.000000e+00 : f3238// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index39// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index40// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index41// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index42// CHECK: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse{{[0-9]*}}> to memref<?xf32>43// CHECK: %[[VAL_8:.*]] = bufferization.to_buffer %[[VAL_1]] : tensor<32x16xf32> to memref<32x16xf32>44// CHECK: scf.for %[[VAL_9:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {45// CHECK: %[[VAL_11:.*]] = arith.muli %[[VAL_9]], %[[VAL_4]] : index46// CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {47// CHECK: %[[VAL_12:.*]] = arith.addi %[[VAL_10]], %[[VAL_11]] : index48// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_12]]] : memref<?xf32>49// CHECK: %[[VAL_14:.*]] = arith.addf %[[VAL_13]], %[[VAL_2]] : f3250// CHECK: memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_9]], %[[VAL_10]]] : memref<32x16xf32>51// CHECK: }52// CHECK: }53// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32x16xf32>54// CHECK: return %[[VAL_15]] : tensor<32x16xf32>55// CHECK: }56func.func @dense1(%arga: tensor<32x16xf32, #DenseMatrix>,57 %argx: tensor<32x16xf32>)58 -> tensor<32x16xf32> {59 %c = arith.constant 1.0 : f3260 %0 = linalg.generic #trait_2d61 ins(%arga: tensor<32x16xf32, #DenseMatrix>)62 outs(%argx: tensor<32x16xf32>) {63 ^bb(%a: f32, %x: f32):64 %1 = arith.addf %a, %c : f3265 linalg.yield %1 : f3266 } -> tensor<32x16xf32>67 return %0 : tensor<32x16xf32>68}69 70//71// Test with a non-annotated dense matrix as input and72// an all-dense annotated "sparse" matrix as output.73//74// CHECK-LABEL: func @dense2(75// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32>,76// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32, #sparse{{[0-9]*}}>) -> tensor<32x16xf32, #sparse{{[0-9]*}}> {77// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1.000000e+00 : f3278// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index79// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index80// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index81// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index82// CHECK: %[[VAL_7:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<32x16xf32> to memref<32x16xf32>83// CHECK: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse{{[0-9]*}}> to memref<?xf32>84// CHECK: scf.for %[[VAL_9:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {85// CHECK: %[[VAL_11:.*]] = arith.muli %[[VAL_9]], %[[VAL_4]] : index86// CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {87// CHECK: %[[VAL_12:.*]] = arith.addi %[[VAL_10]], %[[VAL_11]] : index88// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_9]], %[[VAL_10]]] : memref<32x16xf32>89// CHECK: %[[VAL_14:.*]] = arith.addf %[[VAL_13]], %[[VAL_2]] : f3290// CHECK: memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<?xf32>91// CHECK: }92// CHECK: }93// CHECK: %[[VAL_15:.*]] = sparse_tensor.load %[[VAL_1]] : tensor<32x16xf32, #sparse{{[0-9]*}}>94// CHECK: return %[[VAL_15]] : tensor<32x16xf32, #sparse{{[0-9]*}}>95// CHECK: }96func.func @dense2(%arga: tensor<32x16xf32>,97 %argx: tensor<32x16xf32, #DenseMatrix>)98 -> tensor<32x16xf32, #DenseMatrix> {99 %c = arith.constant 1.0 : f32100 %0 = linalg.generic #trait_2d101 ins(%arga: tensor<32x16xf32>)102 outs(%argx: tensor<32x16xf32, #DenseMatrix>) {103 ^bb(%a: f32, %x: f32):104 %1 = arith.addf %a, %c : f32105 linalg.yield %1 : f32106 } -> tensor<32x16xf32, #DenseMatrix>107 return %0 : tensor<32x16xf32, #DenseMatrix>108}109 110 111//112// Test with a non-annotated dense matrix as input and113// an all-dense annotated "sparse" matrix as output.114// The missing innermost "k" index (due to a reduction) is accounted115// for by scalarizing the reduction operation for the output tensor.116//117// CHECK-LABEL: func @dense3(118// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32>,119// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32, #sparse{{[0-9]*}}>) -> tensor<32x16xf32, #sparse{{[0-9]*}}> {120// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 8 : index121// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index122// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index123// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index124// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index125// CHECK: %[[VAL_7:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<32x16x8xf32> to memref<32x16x8xf32>126// CHECK: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse{{[0-9]*}}> to memref<?xf32>127// CHECK: scf.for %[[VAL_9:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {128// CHECK: %[[VAL_11:.*]] = arith.muli %[[VAL_9]], %[[VAL_4]] : index129// CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {130// CHECK: %[[VAL_12:.*]] = arith.addi %[[VAL_10]], %[[VAL_11]] : index131// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<?xf32>132// CHECK: %[[VAL_14:.*]] = scf.for %[[VAL_15:.*]] = %[[VAL_5]] to %[[VAL_2]] step %[[VAL_6]] iter_args(%[[VAL_16:.*]] = %[[VAL_13]]) -> (f32) {133// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_9]], %[[VAL_10]], %[[VAL_15]]] : memref<32x16x8xf32>134// CHECK: %[[VAL_18:.*]] = arith.addf %[[VAL_16]], %[[VAL_17]] : f32135// CHECK: scf.yield %[[VAL_18]] : f32136// CHECK: }137// CHECK: memref.store %[[VAL_19:.*]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<?xf32>138// CHECK: }139// CHECK: }140// CHECK: %[[VAL_20:.*]] = sparse_tensor.load %[[VAL_1]] : tensor<32x16xf32, #sparse{{[0-9]*}}>141// CHECK: return %[[VAL_20]] : tensor<32x16xf32, #sparse{{[0-9]*}}>142// CHECK: }143func.func @dense3(%arga: tensor<32x16x8xf32>,144 %argx: tensor<32x16xf32, #DenseMatrix>)145 -> tensor<32x16xf32, #DenseMatrix> {146 %0 = linalg.generic #trait_3d147 ins(%arga: tensor<32x16x8xf32>)148 outs(%argx: tensor<32x16xf32, #DenseMatrix>) {149 ^bb(%a: f32, %x: f32):150 %1 = arith.addf %x, %a : f32151 linalg.yield %1 : f32152 } -> tensor<32x16xf32, #DenseMatrix>153 return %0 : tensor<32x16xf32, #DenseMatrix>154}155