diff --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp --- a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp +++ b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp @@ -1278,6 +1278,49 @@ isValid = true; } +// Adds a value to a tensor in coordinate scheme. If is_symmetric_value is true, +// also adds the value to its symmetric location. +template +static void addValue(T *coo, V value, const std::vector indices, + bool is_symmetric_value) { + // TODO: + coo->add(indices, value); + // We currently chose to deal with symmetric matrices by fully constructing + // them. In the future, we may want to make symmetry implicit for storage + // reasons. + if (is_symmetric_value) + coo->add({indices[1], indices[0]}, value); +} + +// Reads an element of a complex type for the current indices in coordinate +// scheme. +template +static void readCOOValue(SparseTensorCOO> *coo, + const std::vector indices, char **linePtr, + bool is_pattern, bool add_symmetric_value) { + // For a pattern tensor, we arbitrarily pick the value 1 for all entries. + V re = is_pattern ? 1.0 : strtod(*linePtr, linePtr); + V im = is_pattern ? 1.0 : strtod(*linePtr, linePtr); + std::complex value = {re, im}; + addValue(coo, value, indices, add_symmetric_value); +} + +// Reads an element of a non-complex type for the current indices in coordinate +// scheme. +template , V>::value && + !std::is_same, V>::value>::type * = nullptr> +static void readCOOValue(SparseTensorCOO *coo, + const std::vector indices, char **linePtr, + bool is_pattern, bool is_symmetric_value) { + // The external formats always store these numerical values with the type + // double, but we cast these values to the sparse tensor object type. + // For a pattern tensor, we arbitrarily pick the value 1 for all entries. + double value = is_pattern ? 1.0 : strtod(*linePtr, linePtr); + addValue(coo, value, indices, is_symmetric_value); +} + /// Reads a sparse tensor with the given filename into a memory-resident /// sparse tensor in coordinate scheme. template @@ -1302,17 +1345,8 @@ // Add 0-based index. indices[perm[r]] = idx - 1; } - // The external formats always store the numerical values with the type - // double, but we cast these values to the sparse tensor object type. - // For a pattern tensor, we arbitrarily pick the value 1 for all entries. - double value = stfile.isPattern() ? 1.0 : strtod(linePtr, &linePtr); - // TODO: - coo->add(indices, value); - // We currently chose to deal with symmetric matrices by fully constructing - // them. In the future, we may want to make symmetry implicit for storage - // reasons. - if (stfile.isSymmetric() && indices[0] != indices[1]) - coo->add({indices[1], indices[0]}, value); + readCOOValue(coo, indices, &linePtr, stfile.isPattern(), + stfile.isSymmetric() && indices[0] != indices[1]); } // Close the file and return tensor. stfile.closeFile(); diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt --- a/mlir/test/CMakeLists.txt +++ b/mlir/test/CMakeLists.txt @@ -48,6 +48,7 @@ file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/mttkrp_b.tns ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/test.mtx ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/test_symmetric.mtx + ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/test_symmetric_complex.mtx ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/test.tns ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/wide.mtx DESTINATION ${MLIR_INTEGRATION_TEST_DIR}/data/) diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir @@ -0,0 +1,86 @@ +// RUN: mlir-opt %s --sparse-compiler | \ +// RUN: TENSOR0="%mlir_integration_test_dir/data/test_symmetric_complex.mtx" \ +// RUN: mlir-cpu-runner \ +// RUN: -e entry -entry-point-result=void \ +// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \ +// RUN: FileCheck %s + +!Filename = !llvm.ptr + +#SparseMatrix = #sparse_tensor.encoding<{ + dimLevelType = [ "compressed", "compressed" ] +}> + +#trait_sum_reduce = { + indexing_maps = [ + affine_map<(i,j) -> (i,j)>, // A + affine_map<(i,j) -> ()> // x (out) + ], + iterator_types = ["reduction", "reduction"], + doc = "x += A(i,j)" +} + +// +// Integration test that lowers a kernel annotated as sparse to +// actual sparse code, initializes a matching sparse storage scheme +// from file, and runs the resulting code with the JIT compiler. +// +module { + // + // A kernel that sum-reduces a matrix to a single scalar. + // + func.func @kernel_sum_reduce(%arga: tensor, #SparseMatrix>, + %argx: tensor> {linalg.inplaceable = true}) -> tensor> { + %0 = linalg.generic #trait_sum_reduce + ins(%arga: tensor, #SparseMatrix>) + outs(%argx: tensor>) { + ^bb(%a: complex, %x: complex): + %0 = complex.add %x, %a : complex + linalg.yield %0 : complex + } -> tensor> + return %0 : tensor> + } + + func.func private @getTensorFilename(index) -> (!Filename) + + // + // Main driver that reads matrix from file and calls the sparse kernel. + // + func.func @entry() { + //%d0 = arith.constant 0.0 : complex + %d0 = complex.constant [0.0 : f64, 0.0 : f64] : complex + %c0 = arith.constant 0 : index + + // Setup memory for a single reduction scalar, + // initialized to zero. + %xdata = memref.alloc() : memref> + memref.store %d0, %xdata[] : memref> + %x = bufferization.to_tensor %xdata : memref> + + // Read the sparse matrix from file, construct sparse storage. + %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename) + %a = sparse_tensor.new %fileName : !Filename to tensor, #SparseMatrix> + + // Call the kernel. + %0 = call @kernel_sum_reduce(%a, %x) + : (tensor, #SparseMatrix>, tensor>) -> tensor> + + // Print the result for verification. + // + // CHECK: 30.2 + // CHECK-NEXT: 22.2 + // + %m = bufferization.to_memref %0 : memref> + %v = memref.load %m[] : memref> + %real = complex.re %v : complex + %imag = complex.im %v : complex + vector.print %real : f64 + vector.print %imag : f64 + + // Release the resources. + memref.dealloc %xdata : memref> + sparse_tensor.release %a : tensor, #SparseMatrix> + + return + } +} diff --git a/mlir/test/Integration/data/test_symmetric_complex.mtx b/mlir/test/Integration/data/test_symmetric_complex.mtx new file mode 100644 --- /dev/null +++ b/mlir/test/Integration/data/test_symmetric_complex.mtx @@ -0,0 +1,13 @@ +%%MatrixMarket matrix coordinate real symmetric +% +% This is a test sparse matrix in Matrix Market Exchange Format. +% see https://math.nist.gov/MatrixMarket +% +5 5 7 +1 1 5.0 1.0 +1 3 4.1 2.1 +2 2 3.0 3.0 +2 4 2.0 4.0 +3 3 1.0 3.0 +4 4 4.0 2.0 +5 5 5.0 1.0