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 @@ -1150,10 +1150,11 @@ } /// Reads and parses the file's header. - void readHeader() { + void readHeader(bool checkValueType = false, + PrimaryType valTp = PrimaryType::kF64) { assert(file && "Attempt to readHeader() before openFile()"); if (strstr(filename, ".mtx")) - readMMEHeader(); + readMMEHeader(checkValueType, valTp); else if (strstr(filename, ".tns")) readExtFROSTTHeader(); else @@ -1209,7 +1210,7 @@ } private: - void readMMEHeader(); + void readMMEHeader(bool checkValueType, PrimaryType valTp); void readExtFROSTTHeader(); const char *filename; @@ -1222,7 +1223,7 @@ }; /// Read the MME header of a general sparse matrix of type real. -void SparseTensorFile::readMMEHeader() { +void SparseTensorFile::readMMEHeader(bool checkValueType, PrimaryType valTp) { char header[64]; char object[64]; char format[64]; @@ -1232,14 +1233,33 @@ if (fscanf(file, "%63s %63s %63s %63s %63s\n", header, object, format, field, symmetry) != 5) FATAL("Corrupt header in %s\n", filename); - // Set properties - isPattern_ = (strcmp(toLower(field), "pattern") == 0); + // Process `field`, which specify pattern or the data type of the values. + bool valueKindCompatible = true; + if (strcmp(toLower(field), "pattern") == 0) + isPattern_ = true; + else if (strcmp(toLower(field), "real") == 0) + valueKindCompatible = + (valTp >= PrimaryType::kF64 && valTp <= PrimaryType::kF32); + else if (strcmp(toLower(field), "integer") == 0) + valueKindCompatible = + (valTp >= PrimaryType::kI64 && valTp <= PrimaryType::kI8); + else if (strcmp(toLower(field), "complex") == 0) + valueKindCompatible = + (valTp >= PrimaryType::kC64 && valTp <= PrimaryType::kC32); + else + FATAL("Unexpected header field value in %s\n", filename); + + if (checkValueType && !valueKindCompatible) + FATAL( + "Header field value %s in %s not compatible with tensor element type\n", + field, filename); + + // Set properties. isSymmetric_ = (strcmp(toLower(symmetry), "symmetric") == 0); // Make sure this is a general sparse matrix. if (strcmp(toLower(header), "%%matrixmarket") || strcmp(toLower(object), "matrix") || strcmp(toLower(format), "coordinate") || - (strcmp(toLower(field), "real") && !isPattern_) || (strcmp(toLower(symmetry), "general") && !isSymmetric_)) FATAL("Cannot find a general sparse matrix in %s\n", filename); // Skip comments. @@ -1278,15 +1298,64 @@ 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 inline 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 inline void readCOOValue(SparseTensorCOO> *coo, + const std::vector indices, + char **linePtr, bool is_pattern, + bool add_symmetric_value) { + // Read two values to make a complex. The external formats always store + // 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. + 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 inline 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 -static SparseTensorCOO *openSparseTensorCOO(char *filename, uint64_t rank, - const uint64_t *shape, - const uint64_t *perm) { +static SparseTensorCOO * +openSparseTensorCOO(char *filename, uint64_t rank, const uint64_t *shape, + const uint64_t *perm, PrimaryType valTp) { SparseTensorFile stfile(filename); stfile.openFile(); - stfile.readHeader(); + stfile.readHeader(/* checkValueType = */ true, valTp); stfile.assertMatchesShape(rank, shape); // Prepare sparse tensor object with per-dimension sizes // and the number of nonzeros as initial capacity. @@ -1302,17 +1371,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(); @@ -1441,7 +1501,7 @@ if (action <= Action::kFromCOO) { \ if (action == Action::kFromFile) { \ char *filename = static_cast(ptr); \ - coo = openSparseTensorCOO(filename, rank, shape, perm); \ + coo = openSparseTensorCOO(filename, rank, shape, perm, valTp); \ } else if (action == Action::kFromCOO) { \ coo = static_cast *>(ptr); \ } else { \ @@ -1776,7 +1836,7 @@ assert(out && "Received nullptr for out-parameter"); SparseTensorFile stfile(filename); stfile.openFile(); - stfile.readHeader(); + stfile.readHeader(/* checkValueType = */ false); stfile.closeFile(); const uint64_t rank = stfile.getRank(); const uint64_t *dimSizes = stfile.getDimSizes(); diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt --- a/mlir/test/CMakeLists.txt +++ b/mlir/test/CMakeLists.txt @@ -48,8 +48,10 @@ 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 + ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/wide_integer.mtx DESTINATION ${MLIR_INTEGRATION_TEST_DIR}/data/) endif() diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir @@ -1,5 +1,5 @@ // RUN: mlir-opt %s --sparse-compiler | \ -// RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \ +// RUN: TENSOR0="%mlir_integration_test_dir/data/wide_integer.mtx" \ // RUN: mlir-cpu-runner \ // RUN: -e entry -entry-point-result=void \ // RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \ @@ -9,7 +9,7 @@ // // RUN: mlir-opt %s \ // RUN: --sparse-compiler="vectorization-strategy=2 vl=16 enable-simd-index32" | \ -// RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \ +// RUN: TENSOR0="%mlir_integration_test_dir/data/wide_integer.mtx" \ // RUN: mlir-cpu-runner \ // RUN: -e entry -entry-point-result=void \ // RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \ 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 complex 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 diff --git a/mlir/test/Integration/data/wide_integer.mtx b/mlir/test/Integration/data/wide_integer.mtx new file mode 100644 --- /dev/null +++ b/mlir/test/Integration/data/wide_integer.mtx @@ -0,0 +1,23 @@ +%%MatrixMarket matrix coordinate integer general +% +% This is a test sparse matrix in Matrix Market Exchange Format. +% see https://math.nist.gov/MatrixMarket +% +4 256 17 +1 1 -1 +1 127 2 +1 128 -3 +1 255 4 +2 2 -5 +2 254 6 +3 3 -7 +4 1 8 +4 2 -9 +4 4 10 +4 99 -11 +4 127 12 +4 128 -13 +4 129 14 +4 250 -15 +4 254 16 +4 256 -17