1// RUN: mlir-opt %s --sparse-assembler -split-input-file | FileCheck %s 2 3// ----- 4 5// CHECK-LABEL: func.func @nop( 6// CHECK-SAME: %[[A:.*]]: tensor<100xf32>) -> tensor<100xf32> { 7// CHECK: return %[[A]] : tensor<100xf32> 8// CHECK: } 9func.func @nop(%arg0: tensor<100xf32>) -> tensor<100xf32> { 10 return %arg0 : tensor<100xf32> 11} 12 13// ----- 14 15// CHECK-LABEL: func.func @sparse_in( 16// CHECK-SAME: %[[B:.*0]]: tensor<?xindex>, 17// CHECK-SAME: %[[C:.*1]]: tensor<?xindex>, 18// CHECK-SAME: %[[A:.*]]: tensor<?xf32>) -> tensor<64x64xf32> { 19// CHECK: %[[I:.*]] = sparse_tensor.assemble (%[[B]], %[[C]]), %[[A]] 20// CHECK: %[[F:.*]] = call @_internal_sparse_in(%[[I]]) 21// CHECK: return %[[F]] : tensor<64x64xf32> 22// CHECK: } 23// CHECK: func.func private @_internal_sparse_in 24#sparse = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }> 25func.func @sparse_in(%arg0: tensor<64x64xf32, #sparse>) -> tensor<64x64xf32> { 26 %0 = sparse_tensor.convert %arg0 : tensor<64x64xf32, #sparse> to tensor<64x64xf32> 27 return %0 : tensor<64x64xf32> 28} 29 30// ----- 31 32// CHECK-LABEL: func.func @sparse_in2( 33// CHECK-SAME: %[[X:.*0]]: tensor<100xf32>, 34// CHECK-SAME: %[[B:.*1]]: tensor<?xindex>, 35// CHECK-SAME: %[[C:.*2]]: tensor<?xindex>, 36// CHECK-SAME: %[[A:.*3]]: tensor<?xf32>) -> tensor<64x64xf32> { 37// CHECK: %[[I:.*]] = sparse_tensor.assemble (%[[B]], %[[C]]), %[[A]] 38// CHECK: %[[F:.*]] = call @_internal_sparse_in2(%[[X]], %[[I]]) 39// CHECK: return %[[F]] : tensor<64x64xf32> 40// CHECK: } 41// CHECK: func.func private @_internal_sparse_in2 42#sparse = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }> 43func.func @sparse_in2(%arg0: tensor<100xf32>, %arg1: tensor<64x64xf32, #sparse>) -> tensor<64x64xf32> { 44 %0 = sparse_tensor.convert %arg1 : tensor<64x64xf32, #sparse> to tensor<64x64xf32> 45 return %0 : tensor<64x64xf32> 46} 47 48// ----- 49 50// CHECK-LABEL: func.func @sparse_out( 51// CHECK-SAME: %[[X:.*0]]: tensor<64x64xf32>, 52// CHECK-SAME: %[[B:.*1]]: tensor<?xindex>, 53// CHECK-SAME: %[[C:.*2]]: tensor<?xindex>, 54// CHECK-SAME: %[[A:.*3]]: tensor<?xf32>) 55// CHECK: %[[F:.*]] = call @_internal_sparse_out(%[[X]]) 56// CHECK: sparse_tensor.disassemble %[[F]] 57// CHECK: return 58// CHECK: } 59// CHECK: func.func private @_internal_sparse_out 60#sparse = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }> 61func.func @sparse_out(%arg0: tensor<64x64xf32>) -> tensor<64x64xf32, #sparse> { 62 %0 = sparse_tensor.convert %arg0 : tensor<64x64xf32> to tensor<64x64xf32, #sparse> 63 return %0 : tensor<64x64xf32, #sparse> 64} 65 66// ----- 67 68// CHECK-LABEL: func.func @sparse_out2( 69// CHECK-SAME: %[[X:.*0]]: tensor<64x64xf32>, 70// CHECK-SAME: %[[B:.*1]]: tensor<?xindex>, 71// CHECK-SAME: %[[C:.*2]]: tensor<?xindex>, 72// CHECK-SAME: %[[A:.*3]]: tensor<?xf32>) 73// CHECK: %[[F:.*]]:2 = call @_internal_sparse_out2(%[[X]]) 74// CHECK: sparse_tensor.disassemble %[[F]]#1 75// CHECK: return %[[F]]#0 76// CHECK: } 77// CHECK: func.func private @_internal_sparse_out2 78#sparse = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }> 79func.func @sparse_out2(%arg0: tensor<64x64xf32>) -> (tensor<64x64xf32>, tensor<64x64xf32, #sparse>) { 80 %0 = sparse_tensor.convert %arg0 : tensor<64x64xf32> to tensor<64x64xf32, #sparse> 81 return %arg0, %0 : tensor<64x64xf32>, tensor<64x64xf32, #sparse> 82} 83 84// ----- 85 86// CHECK-LABEL: func.func @sparse_inout( 87// CHECK-SAME: %[[B:.*0]]: tensor<?xindex>, 88// CHECK-SAME: %[[C:.*1]]: tensor<?xindex>, 89// CHECK-SAME: %[[A:.*2]]: tensor<?xf32>, 90// CHECK-SAME: %[[E:.*3]]: tensor<?xindex>, 91// CHECK-SAME: %[[F:.*4]]: tensor<?xindex>, 92// CHECK-SAME: %[[D:.*5]]: tensor<?xf32>) 93// CHECK: %[[I:.*]] = sparse_tensor.assemble (%[[B]], %[[C]]), %[[A]] 94// CHECK: %[[F:.*]] = call @_internal_sparse_inout(%[[I]]) 95// CHECK: sparse_tensor.disassemble %[[F]] 96// CHECK: return 97// CHECK: } 98// CHECK: func.func private @_internal_sparse_inout 99#sparse = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }> 100func.func @sparse_inout(%arg0: tensor<64x64xf32, #sparse>) -> tensor<64x64xf32, #sparse> { 101 return %arg0 : tensor<64x64xf32, #sparse> 102} 103 104// ----- 105 106// CHECK-LABEL: func.func @sparse_inout_coo_soa( 107// CHECK-SAME: %[[B:.*0]]: tensor<?xindex>, 108// CHECK-SAME: %[[C:.*1]]: tensor<?xindex>, 109// CHECK-SAME: %[[D:.*2]]: tensor<?xindex>, 110// CHECK-SAME: %[[A:.*3]]: tensor<?xf32>, 111// CHECK-SAME: %[[F:.*4]]: tensor<?xindex>, 112// CHECK-SAME: %[[G:.*5]]: tensor<?xindex>, 113// CHECK-SAME: %[[H:.*6]]: tensor<?xindex>, 114// CHECK-SAME: %[[E:.*7]]: tensor<?xf32>) 115// CHECK: %[[I:.*]] = sparse_tensor.assemble (%[[B]], %[[C]], %[[D]]), %[[A]] 116// CHECK: %[[F:.*]] = call @_internal_sparse_inout_coo_soa(%[[I]]) 117// CHECK: sparse_tensor.disassemble %[[F]] 118// CHECK: return 119// CHECK: } 120// CHECK: func.func private @_internal_sparse_inout 121#sparse = #sparse_tensor.encoding<{ 122 map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton(soa)) 123}> 124func.func @sparse_inout_coo_soa(%arg0: tensor<64x64xf32, #sparse>) -> tensor<64x64xf32, #sparse> { 125 return %arg0 : tensor<64x64xf32, #sparse> 126} 127