// RUN: mlir-opt %s --sparse-compiler | \ // RUN: mlir-cpu-runner -e entry -entry-point-result=void \ // RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \ // RUN: FileCheck %s #SparseVector = #sparse_tensor.encoding<{ dimLevelType = ["compressed"] }> #SparseMatrix = #sparse_tensor.encoding<{ dimLevelType = ["compressed", "compressed"] }> // // Test with various forms of the two most elementary reshape // operations: expand/collapse. // module { func.func @expand_dense(%arg0: tensor<12xf64>) -> tensor<3x4xf64> { %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64> return %0 : tensor<3x4xf64> } func.func @expand_from_sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64> { %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64> return %0 : tensor<3x4xf64> } func.func @expand_to_sparse(%arg0: tensor<12xf64>) -> tensor<3x4xf64, #SparseMatrix> { %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64, #SparseMatrix> return %0 : tensor<3x4xf64, #SparseMatrix> } func.func @expand_sparse2sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64, #SparseMatrix> { %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64, #SparseMatrix> return %0 : tensor<3x4xf64, #SparseMatrix> } func.func @collapse_dense(%arg0: tensor<3x4xf64>) -> tensor<12xf64> { %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64> into tensor<12xf64> return %0 : tensor<12xf64> } func.func @collapse_from_sparse(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64> { %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64, #SparseMatrix> into tensor<12xf64> return %0 : tensor<12xf64> } func.func @collapse_to_sparse(%arg0: tensor<3x4xf64>) -> tensor<12xf64, #SparseVector> { %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64> into tensor<12xf64, #SparseVector> return %0 : tensor<12xf64, #SparseVector> } func.func @collapse_sparse2sparse(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64, #SparseVector> { %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64, #SparseMatrix> into tensor<12xf64, #SparseVector> return %0 : tensor<12xf64, #SparseVector> } // // Main driver. // func.func @entry() { %c0 = arith.constant 0 : index %df = arith.constant -1.0 : f64 // Setup test vectors and matrices.. %v = arith.constant dense <[ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]> : tensor<12xf64> %m = arith.constant dense <[ [ 1.1, 1.2, 1.3, 1.4 ], [ 2.1, 2.2, 2.3, 2.4 ], [ 3.1, 3.2, 3.3, 3.4 ]]> : tensor<3x4xf64> %sv = sparse_tensor.convert %v : tensor<12xf64> to tensor<12xf64, #SparseVector> %sm = sparse_tensor.convert %m : tensor<3x4xf64> to tensor<3x4xf64, #SparseMatrix> // Call the kernels. %expand0 = call @expand_dense(%v) : (tensor<12xf64>) -> tensor<3x4xf64> %expand1 = call @expand_from_sparse(%sv) : (tensor<12xf64, #SparseVector>) -> tensor<3x4xf64> %expand2 = call @expand_to_sparse(%v) : (tensor<12xf64>) -> tensor<3x4xf64, #SparseMatrix> %expand3 = call @expand_sparse2sparse(%sv) : (tensor<12xf64, #SparseVector>) -> tensor<3x4xf64, #SparseMatrix> %collapse0 = call @collapse_dense(%m) : (tensor<3x4xf64>) -> tensor<12xf64> %collapse1 = call @collapse_from_sparse(%sm) : (tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64> %collapse2 = call @collapse_to_sparse(%m) : (tensor<3x4xf64>) -> tensor<12xf64, #SparseVector> %collapse3 = call @collapse_sparse2sparse(%sm) : (tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64, #SparseVector> // // Verify result. // // CHECK: ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) ) // CHECK-NEXT: ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) ) // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, -1, -1, -1, -1 ) // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, -1, -1, -1, -1 ) // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4 ) // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4 ) // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 ) // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 ) // %m0 = vector.transfer_read %expand0[%c0, %c0], %df: tensor<3x4xf64>, vector<3x4xf64> vector.print %m0 : vector<3x4xf64> %m1 = vector.transfer_read %expand1[%c0, %c0], %df: tensor<3x4xf64>, vector<3x4xf64> vector.print %m1 : vector<3x4xf64> %a2 = sparse_tensor.values %expand2 : tensor<3x4xf64, #SparseMatrix> to memref %m2 = vector.transfer_read %a2[%c0], %df: memref, vector<16xf64> vector.print %m2 : vector<16xf64> %a3 = sparse_tensor.values %expand3 : tensor<3x4xf64, #SparseMatrix> to memref %m3 = vector.transfer_read %a3[%c0], %df: memref, vector<16xf64> vector.print %m3 : vector<16xf64> %v0 = vector.transfer_read %collapse0[%c0], %df: tensor<12xf64>, vector<12xf64> vector.print %v0 : vector<12xf64> %v1 = vector.transfer_read %collapse1[%c0], %df: tensor<12xf64>, vector<12xf64> vector.print %v1 : vector<12xf64> %b2 = sparse_tensor.values %collapse2 : tensor<12xf64, #SparseVector> to memref %v2 = vector.transfer_read %b2[%c0], %df: memref, vector<16xf64> vector.print %v2 : vector<16xf64> %b3 = sparse_tensor.values %collapse3 : tensor<12xf64, #SparseVector> to memref %v3 = vector.transfer_read %b3[%c0], %df: memref, vector<16xf64> vector.print %v3 : vector<16xf64> // Release sparse resources. bufferization.dealloc_tensor %sv : tensor<12xf64, #SparseVector> bufferization.dealloc_tensor %sm : tensor<3x4xf64, #SparseMatrix> bufferization.dealloc_tensor %expand2 : tensor<3x4xf64, #SparseMatrix> bufferization.dealloc_tensor %expand3 : tensor<3x4xf64, #SparseMatrix> bufferization.dealloc_tensor %collapse2 : tensor<12xf64, #SparseVector> bufferization.dealloc_tensor %collapse3 : tensor<12xf64, #SparseVector> // Release dense resources. bufferization.dealloc_tensor %expand1 : tensor<3x4xf64> bufferization.dealloc_tensor %collapse1 : tensor<12xf64> return } }