From efa15f417847439eb8884f1ef47e163fa450ac7c Mon Sep 17 00:00:00 2001 From: Aart Bik Date: Thu, 20 Jan 2022 17:27:23 -0800 Subject: [PATCH] [mlir][sparse] add ability for sparse tensor output Rationale: Although file I/O is a bit alien to MLIR itself, we provide two convenient ways for sparse tensor I/O. The input part was already there (behind the swiss army knife sparse_tensor.new). Now we have a sparse_tensor.out to write out data. As before, the ops are kept vague and may change in the future. For now this allows us to compare TACO vs MLIR very easily. Reviewed By: bixia Differential Revision: https://reviews.llvm.org/D117850 --- .../SparseTensor/IR/SparseTensorOps.td | 20 ++++++ .../SparseTensor/IR/SparseTensorDialect.cpp | 6 ++ .../Transforms/SparseTensorConversion.cpp | 70 +++++++++++++++---- .../lib/ExecutionEngine/SparseTensorUtils.cpp | 48 +++++++++++++ .../test/Dialect/SparseTensor/conversion.mlir | 24 +++++++ mlir/test/Dialect/SparseTensor/invalid.mlir | 8 +++ mlir/test/Dialect/SparseTensor/roundtrip.mlir | 14 ++++ 7 files changed, 177 insertions(+), 13 deletions(-) diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td index b7fce5b3137c..1209a70a72c2 100644 --- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td +++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td @@ -351,4 +351,24 @@ def SparseTensor_ReleaseOp : SparseTensor_Op<"release", []>, let assemblyFormat = "$tensor attr-dict `:` type($tensor)"; } +def SparseTensor_OutOp : SparseTensor_Op<"out", []>, + Arguments<(ins AnyType:$tensor, AnyType:$dest)> { + string summary = "Outputs a sparse tensor to the given destination"; + string description = [{ + Outputs the contents of a sparse tensor to the destination defined by an + opaque pointer provided by `dest`. For targets that have access to a file + system, for example, this pointer may specify a filename (or file) for output. + The form of the operation is kept deliberately very general to allow for + alternative implementations in the future, such as sending the contents to + a buffer defined by a pointer. + + Example: + + ```mlir + sparse_tensor.out %t, %dest : tensor<1024x1024xf64, #CSR>, !Dest + ``` + }]; + let assemblyFormat = "$tensor `,` $dest attr-dict `:` type($tensor) `,` type($dest)"; +} + #endif // SPARSETENSOR_OPS diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp index 8f44b5b1b9e2..8a7942c8d666 100644 --- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp +++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp @@ -329,6 +329,12 @@ static LogicalResult verify(ReleaseOp op) { return success(); } +static LogicalResult verify(OutOp op) { + if (!getSparseTensorEncoding(op.tensor().getType())) + return op.emitError("expected a sparse tensor for output"); + return success(); +} + //===----------------------------------------------------------------------===// // TensorDialect Methods. //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp index a28f9ac70b31..94e87b3b79b7 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp @@ -15,6 +15,7 @@ //===----------------------------------------------------------------------===// #include "CodegenUtils.h" + #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/Linalg/Utils/Utils.h" @@ -189,8 +190,8 @@ static Value genBuffer(ConversionPatternRewriter &rewriter, Location loc, /// computation. static void newParams(ConversionPatternRewriter &rewriter, SmallVector ¶ms, Operation *op, - SparseTensorEncodingAttr &enc, Action action, - ValueRange szs, Value ptr = Value()) { + ShapedType stp, SparseTensorEncodingAttr &enc, + Action action, ValueRange szs, Value ptr = Value()) { Location loc = op->getLoc(); ArrayRef dlt = enc.getDimLevelType(); unsigned sz = dlt.size(); @@ -218,7 +219,7 @@ static void newParams(ConversionPatternRewriter &rewriter, } params.push_back(genBuffer(rewriter, loc, rev)); // Secondary and primary types encoding. - Type elemTp = op->getResult(0).getType().cast().getElementType(); + Type elemTp = stp.getElementType(); params.push_back(constantPointerTypeEncoding(rewriter, loc, enc)); params.push_back(constantIndexTypeEncoding(rewriter, loc, enc)); params.push_back(constantPrimaryTypeEncoding(rewriter, loc, elemTp)); @@ -420,9 +421,10 @@ class SparseTensorNewConverter : public OpConversionPattern { // inferred from the result type of the new operator. SmallVector sizes; SmallVector params; - sizesFromType(rewriter, sizes, op.getLoc(), resType.cast()); + ShapedType stp = resType.cast(); + sizesFromType(rewriter, sizes, op.getLoc(), stp); Value ptr = adaptor.getOperands()[0]; - newParams(rewriter, params, op, enc, Action::kFromFile, sizes, ptr); + newParams(rewriter, params, op, stp, enc, Action::kFromFile, sizes, ptr); rewriter.replaceOp(op, genNewCall(rewriter, op, params)); return success(); } @@ -441,7 +443,9 @@ class SparseTensorInitConverter : public OpConversionPattern { // Generate the call to construct empty tensor. The sizes are // explicitly defined by the arguments to the init operator. SmallVector params; - newParams(rewriter, params, op, enc, Action::kEmpty, adaptor.getOperands()); + ShapedType stp = resType.cast(); + newParams(rewriter, params, op, stp, enc, Action::kEmpty, + adaptor.getOperands()); rewriter.replaceOp(op, genNewCall(rewriter, op, params)); return success(); } @@ -472,15 +476,15 @@ class SparseTensorConvertConverter : public OpConversionPattern { } SmallVector sizes; SmallVector params; - sizesFromPtr(rewriter, sizes, op, encSrc, srcType.cast(), - src); + ShapedType stp = srcType.cast(); + sizesFromPtr(rewriter, sizes, op, encSrc, stp, src); // Set up encoding with right mix of src and dst so that the two // method calls can share most parameters, while still providing // the correct sparsity information to either of them. auto enc = SparseTensorEncodingAttr::get( op->getContext(), encDst.getDimLevelType(), encDst.getDimOrdering(), encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth()); - newParams(rewriter, params, op, enc, Action::kToCOO, sizes, src); + newParams(rewriter, params, op, stp, enc, Action::kToCOO, sizes, src); Value coo = genNewCall(rewriter, op, params); params[3] = constantPointerTypeEncoding(rewriter, loc, encDst); params[4] = constantIndexTypeEncoding(rewriter, loc, encDst); @@ -512,7 +516,8 @@ class SparseTensorConvertConverter : public OpConversionPattern { SmallVector sizes; SmallVector params; sizesFromPtr(rewriter, sizes, op, encSrc, srcTensorTp, src); - newParams(rewriter, params, op, encDst, Action::kToIterator, sizes, src); + newParams(rewriter, params, op, dstTensorTp, encDst, Action::kToIterator, + sizes, src); Value iter = genNewCall(rewriter, op, params); Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType()); Value elemPtr = genAllocaScalar(rewriter, loc, elemTp); @@ -567,7 +572,7 @@ class SparseTensorConvertConverter : public OpConversionPattern { SmallVector sizes; SmallVector params; sizesFromSrc(rewriter, sizes, loc, src); - newParams(rewriter, params, op, encDst, Action::kEmptyCOO, sizes); + newParams(rewriter, params, op, stp, encDst, Action::kEmptyCOO, sizes); Value ptr = genNewCall(rewriter, op, params); Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType()); Value perm = params[2]; @@ -771,6 +776,45 @@ public: } }; +class SparseTensorOutConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(OutOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); + ShapedType srcType = op.tensor().getType().cast(); + // Convert to default permuted COO. + Value src = adaptor.getOperands()[0]; + auto encSrc = getSparseTensorEncoding(srcType); + SmallVector sizes; + SmallVector params; + sizesFromPtr(rewriter, sizes, op, encSrc, srcType, src); + auto enc = SparseTensorEncodingAttr::get( + op->getContext(), encSrc.getDimLevelType(), AffineMap(), + encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth()); + newParams(rewriter, params, op, srcType, enc, Action::kToCOO, sizes, src); + Value coo = genNewCall(rewriter, op, params); + // Then output the tensor to external file with indices in the externally + // visible lexicographic index order. A sort is required if the source was + // not in that order yet (note that the sort can be dropped altogether if + // external format does not care about the order at all, but here we assume + // it does). + bool sort = + encSrc.getDimOrdering() && !encSrc.getDimOrdering().isIdentity(); + params.clear(); + params.push_back(coo); + params.push_back(adaptor.getOperands()[1]); + params.push_back(constantI1(rewriter, loc, sort)); + Type eltType = srcType.getElementType(); + SmallString<18> name{"outSparseTensor", primaryTypeFunctionSuffix(eltType)}; + TypeRange noTp; + replaceOpWithFuncCall(rewriter, op, name, noTp, params, + EmitCInterface::Off); + return success(); + } +}; + } // namespace //===----------------------------------------------------------------------===// @@ -787,6 +831,6 @@ void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter, SparseTensorReleaseConverter, SparseTensorToPointersConverter, SparseTensorToIndicesConverter, SparseTensorToValuesConverter, SparseTensorLoadConverter, SparseTensorLexInsertConverter, - SparseTensorExpandConverter, SparseTensorCompressConverter>( - typeConverter, patterns.getContext()); + SparseTensorExpandConverter, SparseTensorCompressConverter, + SparseTensorOutConverter>(typeConverter, patterns.getContext()); } diff --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp index 20cd1b53d31b..605e17764773 100644 --- a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp +++ b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp @@ -26,6 +26,8 @@ #include #include #include +#include +#include #include #include #include @@ -713,6 +715,31 @@ static SparseTensorCOO *openSparseTensorCOO(char *filename, uint64_t rank, return tensor; } +/// Writes the sparse tensor to extended FROSTT format. +template +void outSparseTensor(const SparseTensorCOO &tensor, char *filename) { + auto &sizes = tensor.getSizes(); + auto &elements = tensor.getElements(); + uint64_t rank = tensor.getRank(); + uint64_t nnz = elements.size(); + std::fstream file; + file.open(filename, std::ios_base::out | std::ios_base::trunc); + assert(file.is_open()); + file << "; extended FROSTT format\n" << rank << " " << nnz << std::endl; + for (uint64_t r = 0; r < rank - 1; r++) + file << sizes[r] << " "; + file << sizes[rank - 1] << std::endl; + for (uint64_t i = 0; i < nnz; i++) { + auto &idx = elements[i].indices; + for (uint64_t r = 0; r < rank; r++) + file << (idx[r] + 1) << " "; + file << elements[i].value << std::endl; + } + file.flush(); + file.close(); + assert(file.good()); +} + } // namespace extern "C" { @@ -845,6 +872,17 @@ extern "C" { cursor, values, filled, added, count); \ } +#define IMPL_OUT(NAME, V) \ + void NAME(void *tensor, void *dest, bool sort) { \ + assert(tensor &&dest); \ + auto coo = static_cast *>(tensor); \ + if (sort) \ + coo->sort(); \ + char *filename = static_cast(dest); \ + outSparseTensor(*coo, filename); \ + delete coo; \ + } + // Assume index_t is in fact uint64_t, so that _mlir_ciface_newSparseTensor // can safely rewrite kIndex to kU64. We make this assertion to guarantee // that this file cannot get out of sync with its header. @@ -1026,6 +1064,14 @@ IMPL_EXPINSERT(expInsertI32, int32_t) IMPL_EXPINSERT(expInsertI16, int16_t) IMPL_EXPINSERT(expInsertI8, int8_t) +/// Helper to output a sparse tensor, one per value type. +IMPL_OUT(outSparseTensorF64, double) +IMPL_OUT(outSparseTensorF32, float) +IMPL_OUT(outSparseTensorI64, int64_t) +IMPL_OUT(outSparseTensorI32, int32_t) +IMPL_OUT(outSparseTensorI16, int16_t) +IMPL_OUT(outSparseTensorI8, int8_t) + #undef CASE #undef IMPL_SPARSEVALUES #undef IMPL_GETOVERHEAD @@ -1033,6 +1079,7 @@ IMPL_EXPINSERT(expInsertI8, int8_t) #undef IMPL_GETNEXT #undef IMPL_LEXINSERT #undef IMPL_EXPINSERT +#undef IMPL_OUT //===----------------------------------------------------------------------===// // @@ -1162,6 +1209,7 @@ void convertFromMLIRSparseTensor(void *tensor, uint64_t *pRank, uint64_t *pNse, *pValues = values; *pIndices = indices; } + } // extern "C" #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS diff --git a/mlir/test/Dialect/SparseTensor/conversion.mlir b/mlir/test/Dialect/SparseTensor/conversion.mlir index 89ee0d5b7c81..04c8a5181781 100644 --- a/mlir/test/Dialect/SparseTensor/conversion.mlir +++ b/mlir/test/Dialect/SparseTensor/conversion.mlir @@ -468,3 +468,27 @@ func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>, : tensor<8x8xf64, #SparseMatrix>, memref, memref, memref, memref, index return } + +// CHECK-LABEL: func @sparse_out1( +// CHECK-SAME: %[[A:.*]]: !llvm.ptr, +// CHECK-SAME: %[[B:.*]]: !llvm.ptr) +// CHECK-DAG: %[[C:.*]] = arith.constant false +// CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]]) +// CHECK: call @outSparseTensorF64(%[[T]], %[[B]], %[[C]]) : (!llvm.ptr, !llvm.ptr, i1) -> () +// CHECK: return +func @sparse_out1(%arg0: tensor, %arg1: !llvm.ptr) { + sparse_tensor.out %arg0, %arg1 : tensor, !llvm.ptr + return +} + +// CHECK-LABEL: func @sparse_out2( +// CHECK-SAME: %[[A:.*]]: !llvm.ptr, +// CHECK-SAME: %[[B:.*]]: !llvm.ptr) +// CHECK-DAG: %[[C:.*]] = arith.constant true +// CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]]) +// CHECK: call @outSparseTensorF32(%[[T]], %[[B]], %[[C]]) : (!llvm.ptr, !llvm.ptr, i1) -> () +// CHECK: return +func @sparse_out2(%arg0: tensor, %arg1: !llvm.ptr) { + sparse_tensor.out %arg0, %arg1 : tensor, !llvm.ptr + return +} diff --git a/mlir/test/Dialect/SparseTensor/invalid.mlir b/mlir/test/Dialect/SparseTensor/invalid.mlir index 06d662127174..84990221e4df 100644 --- a/mlir/test/Dialect/SparseTensor/invalid.mlir +++ b/mlir/test/Dialect/SparseTensor/invalid.mlir @@ -204,3 +204,11 @@ func @sparse_convert_dim_mismatch(%arg0: tensor<10x?xf32>) -> tensor<10x10xf32, %0 = sparse_tensor.convert %arg0 : tensor<10x?xf32> to tensor<10x10xf32, #CSR> return %0 : tensor<10x10xf32, #CSR> } + +// ----- + +func @invalid_out_dense(%arg0: tensor<10xf64>, %arg1: !llvm.ptr) { + // expected-error@+1 {{expected a sparse tensor for output}} + sparse_tensor.out %arg0, %arg1 : tensor<10xf64>, !llvm.ptr + return +} diff --git a/mlir/test/Dialect/SparseTensor/roundtrip.mlir b/mlir/test/Dialect/SparseTensor/roundtrip.mlir index 853befc1cdef..5457e55f57e6 100644 --- a/mlir/test/Dialect/SparseTensor/roundtrip.mlir +++ b/mlir/test/Dialect/SparseTensor/roundtrip.mlir @@ -179,3 +179,17 @@ func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>, : tensor<8x8xf64, #SparseMatrix>, memref, memref, memref, memref, index return } + +// ----- + +#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}> + +// CHECK-LABEL: func @sparse_out( +// CHECK-SAME: %[[A:.*]]: tensor>, +// CHECK-SAME: %[[B:.*]]: !llvm.ptr) +// CHECK: sparse_tensor.out %[[A]], %[[B]] : tensor>, !llvm.ptr +// CHECK: return +func @sparse_out(%arg0: tensor, %arg1: !llvm.ptr) { + sparse_tensor.out %arg0, %arg1 : tensor, !llvm.ptr + return +}