maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[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
This commit is contained in:
Aart Bik 2022-01-20 17:27:23 -08:00
parent 58ee14e29e
commit efa15f4178
7 changed files with 177 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
View File

@ -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

6
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
View File

@ -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.
//===----------------------------------------------------------------------===//

70
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
View File

@ -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<Value, 8> &params, 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<SparseTensorEncodingAttr::DimLevelType> 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<ShapedType>().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<NewOp> {
// inferred from the result type of the new operator.
SmallVector<Value, 4> sizes;
SmallVector<Value, 8> params;
sizesFromType(rewriter, sizes, op.getLoc(), resType.cast<ShapedType>());
ShapedType stp = resType.cast<ShapedType>();
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<InitOp> {
// Generate the call to construct empty tensor. The sizes are
// explicitly defined by the arguments to the init operator.
SmallVector<Value, 8> params;
newParams(rewriter, params, op, enc, Action::kEmpty, adaptor.getOperands());
ShapedType stp = resType.cast<ShapedType>();
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<ConvertOp> {
}
SmallVector<Value, 4> sizes;
SmallVector<Value, 8> params;
sizesFromPtr(rewriter, sizes, op, encSrc, srcType.cast<ShapedType>(),
src);
ShapedType stp = srcType.cast<ShapedType>();
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<ConvertOp> {
SmallVector<Value, 4> sizes;
SmallVector<Value, 8> 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<ConvertOp> {
SmallVector<Value, 4> sizes;
SmallVector<Value, 8> 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<OutOp> {
public:
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(OutOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
Location loc = op->getLoc();
ShapedType srcType = op.tensor().getType().cast<ShapedType>();
// Convert to default permuted COO.
Value src = adaptor.getOperands()[0];
auto encSrc = getSparseTensorEncoding(srcType);
SmallVector<Value, 4> sizes;
SmallVector<Value, 8> 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());
}

48
mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
View File

@ -26,6 +26,8 @@
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <limits>
#include <numeric>
#include <vector>
@ -713,6 +715,31 @@ static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t rank,
return tensor;
}
/// Writes the sparse tensor to extended FROSTT format.
template <typename V>
void outSparseTensor(const SparseTensorCOO<V> &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<SparseTensorCOO<V> *>(tensor); \
if (sort) \
coo->sort(); \
char *filename = static_cast<char *>(dest); \
outSparseTensor<V>(*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

24
mlir/test/Dialect/SparseTensor/conversion.mlir
View File

@ -468,3 +468,27 @@ func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>,
: tensor<8x8xf64, #SparseMatrix>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
return
}
// CHECK-LABEL: func @sparse_out1(
// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>,
// CHECK-SAME: %[[B:.*]]: !llvm.ptr<i8>)
// CHECK-DAG: %[[C:.*]] = arith.constant false
// CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
// CHECK: call @outSparseTensorF64(%[[T]], %[[B]], %[[C]]) : (!llvm.ptr<i8>, !llvm.ptr<i8>, i1) -> ()
// CHECK: return
func @sparse_out1(%arg0: tensor<?x?xf64, #SparseMatrix>, %arg1: !llvm.ptr<i8>) {
sparse_tensor.out %arg0, %arg1 : tensor<?x?xf64, #SparseMatrix>, !llvm.ptr<i8>
return
}
// CHECK-LABEL: func @sparse_out2(
// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>,
// CHECK-SAME: %[[B:.*]]: !llvm.ptr<i8>)
// CHECK-DAG: %[[C:.*]] = arith.constant true
// CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
// CHECK: call @outSparseTensorF32(%[[T]], %[[B]], %[[C]]) : (!llvm.ptr<i8>, !llvm.ptr<i8>, i1) -> ()
// CHECK: return
func @sparse_out2(%arg0: tensor<?x?x?xf32, #SparseTensor>, %arg1: !llvm.ptr<i8>) {
sparse_tensor.out %arg0, %arg1 : tensor<?x?x?xf32, #SparseTensor>, !llvm.ptr<i8>
return
}

8
mlir/test/Dialect/SparseTensor/invalid.mlir
View File

@ -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<i8>) {
// expected-error@+1 {{expected a sparse tensor for output}}
sparse_tensor.out %arg0, %arg1 : tensor<10xf64>, !llvm.ptr<i8>
return
}

14
mlir/test/Dialect/SparseTensor/roundtrip.mlir
View File

@ -179,3 +179,17 @@ func @sparse_compression(%arg0: tensor<8x8xf64, #SparseMatrix>,
: tensor<8x8xf64, #SparseMatrix>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
return
}
// -----
#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
// CHECK-LABEL: func @sparse_out(
// CHECK-SAME: %[[A:.*]]: tensor<?x?xf64, #sparse_tensor.encoding<{{.*}}>>,
// CHECK-SAME: %[[B:.*]]: !llvm.ptr<i8>)
// CHECK: sparse_tensor.out %[[A]], %[[B]] : tensor<?x?xf64, #sparse_tensor.encoding<{{.*}}>>, !llvm.ptr<i8>
// CHECK: return
func @sparse_out(%arg0: tensor<?x?xf64, #SparseMatrix>, %arg1: !llvm.ptr<i8>) {
sparse_tensor.out %arg0, %arg1 : tensor<?x?xf64, #SparseMatrix>, !llvm.ptr<i8>
return
}