94 lines · plain
1// RUN: mlir-opt %s --stage-sparse-ops --lower-sparse-ops-to-foreach --canonicalize --cse | FileCheck %s2 3#SparseVector = #sparse_tensor.encoding<{4 map = (d0) -> (d0 : compressed)5}>6 7#CSR = #sparse_tensor.encoding<{8 map = (d0, d1) -> (d0 : dense, d1 : compressed)9}>10 11#CSC = #sparse_tensor.encoding<{12 map = (d0, d1) -> (d1 : dense, d0 : compressed)13}>14 15#SparseTensor = #sparse_tensor.encoding<{16 map = (d0, d1, d2) -> (d2 : dense, d0 : compressed, d1 : compressed)17}>18 19// CHECK-LABEL: func.func @sparse_convert_1d20// CHECK: sparse_tensor.foreach21// CHECK: scf.if22// CHECK: tensor.insert23// CHECK-NOT: sparse_tensor.reorder_coo24// CHECK: sparse_tensor.load25func.func @sparse_convert_1d(%arg0: tensor<?xi32>) -> tensor<?xi32, #SparseVector> {26 %0 = sparse_tensor.convert %arg0 : tensor<?xi32> to tensor<?xi32, #SparseVector>27 return %0 : tensor<?xi32, #SparseVector>28}29 30// CHECK-LABEL: func.func @sparse_convert_complex31// CHECK: sparse_tensor.foreach32// CHECK: scf.if33// CHECK: tensor.insert34// CHECK-NOT: sparse_tensor.reorder_coo35// CHECK: sparse_tensor.load36func.func @sparse_convert_complex(%arg0: tensor<100xcomplex<f64>>) -> tensor<100xcomplex<f64>, #SparseVector> {37 %0 = sparse_tensor.convert %arg0 : tensor<100xcomplex<f64>> to tensor<100xcomplex<f64>, #SparseVector>38 return %0 : tensor<100xcomplex<f64>, #SparseVector>39}40 41// CHECK-LABEL: func.func @sparse_convert_2d42// CHECK: sparse_tensor.foreach43// CHECK: scf.if44// CHECK: tensor.insert45// CHECK-NOT: sparse_tensor.reorder_coo46// CHECK: sparse_tensor.load47func.func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #CSR> {48 %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to tensor<2x4xf64, #CSR>49 return %0 : tensor<2x4xf64, #CSR>50}51 52// CHECK-LABEL: func.func @sparse_constant53// CHECK: sparse_tensor.foreach54// CHECK-NOT: scf.if55// CHECK: tensor.insert56// CHECK-NOT: sparse_tensor.reorder_coo57// CHECK: sparse_tensor.load58func.func @sparse_constant() -> tensor<8x7xf32, #CSR>{59 // Initialize a tensor.60 %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>61 // Convert the tensor to a sparse tensor.62 %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #CSR>63 return %1 : tensor<8x7xf32, #CSR>64}65 66// CHECK-LABEL: func.func @sparse_constant_csc67// CHECK: sparse_tensor.foreach68// CHECK-NOT: scf.if69// CHECK: tensor.insert70// CHECK-NOT: sparse_tensor.reorder_coo71// CHECK: sparse_tensor.load72func.func @sparse_constant_csc() -> tensor<8x7xf32, #CSC>{73 // Initialize a tensor.74 %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>75 // Convert the tensor to a sparse tensor.76 %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #CSC>77 return %1 : tensor<8x7xf32, #CSC>78}79 80// CHECK-LABEL: func.func @sparse_convert_3d81// CHECK: sparse_tensor.foreach82// CHECK: scf.if83// CHECK: tensor.insert84// CHECK: sparse_tensor.load85// CHECK: %[[TMP:.*]] = sparse_tensor.reorder_coo86// CHECK: sparse_tensor.foreach87// CHECK: tensor.insert88// CHECK: sparse_tensor.load89// CHECK: bufferization.dealloc_tensor %[[TMP]]90func.func @sparse_convert_3d(%arg0: tensor<?x?x?xf64>) -> tensor<?x?x?xf64, #SparseTensor> {91 %0 = sparse_tensor.convert %arg0 : tensor<?x?x?xf64> to tensor<?x?x?xf64, #SparseTensor>92 return %0 : tensor<?x?x?xf64, #SparseTensor>93}94