OSchip
/
llvm-project
mirror of https://github.com/THU-DSP-LAB/llvm-project.git
13
0
Fork 2
Code Issues Packages Projects Releases Wiki Activity

[mlir][sparse] first end-to-end matmul with codegen 

(1) also fixes memory leak in sparse2dense rewriting
(2) still needs fix in dense2sparse by skipping zeros

Reviewed By: wrengr

Differential Revision: https://reviews.llvm.org/D137736
This commit is contained in:
Aart Bik 2022-11-09 13:05:43 -08:00
parent 8bcf5df304
commit a61a9a700a
5 changed files with 28 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
View File

@ -947,6 +947,11 @@ Value mlir::sparse_tensor::allocDenseTensor(OpBuilder &builder, Location loc,
return mem;
}
void mlir::sparse_tensor::deallocDenseTensor(OpBuilder &builder, Location loc,
Value buffer) {
builder.create<memref::DeallocOp>(loc, buffer);
}
Value mlir::sparse_tensor::genValueForDense(OpBuilder &builder, Location loc,
Value tensor, ValueRange ivs) {
Value val = builder.create<tensor::ExtractOp>(loc, tensor, ivs);

3
mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h
View File

@ -145,6 +145,9 @@ Value genAllocaScalar(OpBuilder &builder, Location loc, Type tp);
Value allocDenseTensor(OpBuilder &builder, Location loc,
RankedTensorType tensorTp, ValueRange sizes);
/// Generates code to deallocate a dense buffer.
void deallocDenseTensor(OpBuilder &builder, Location loc, Value buffer);
/// Generates the code to read the value from tensor[ivs]. The generated code
/// looks like the following and the insertion point after this routine is
/// inside the if-then branch behind the assignment to ind.

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

@ -337,11 +337,6 @@ static Value genGetNextCall(OpBuilder &builder, Location loc, Value iter,
.getResult(0);
}
/// Generates code to deallocate a dense buffer.
static void deallocDenseTensor(OpBuilder &builder, Location loc, Value buffer) {
builder.create<memref::DeallocOp>(loc, buffer);
}
/// Converts a pointer to COO (from calls to iter->next()) into a vector of
/// indices, apply (optional) `offset` on `offsetDim`.
static SmallVector<Value, 4> loadIndices(OpBuilder &builder, Location loc,

9
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
View File

@ -564,7 +564,10 @@ private:
SmallVector<Value, 4> sizes;
sizesForTensor(rewriter, sizes, loc, srcTp, src);
Value dst = allocDenseTensor(rewriter, loc, dstTp, sizes);
Block *insertionBlock = rewriter.getInsertionBlock();
bool noEscape = bufferization::allocationDoesNotEscape(op->getOpResult(0));
rewriter.create<ForeachOp>(loc, src, llvm::None,
[&](OpBuilder &builder, Location loc,
@ -575,6 +578,12 @@ private:
});
rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstTp, dst);
// Deallocate the buffer.
if (noEscape) {
rewriter.setInsertionPoint(insertionBlock->getTerminator());
deallocDenseTensor(rewriter, loc, dst);
}
return success();
}

14
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
View File

@ -9,7 +9,13 @@
// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
// RUN: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
// RUN: FileCheck %s
//
// Do the same run, but now with direct IR generation.
//
// RUN: mlir-opt %s --sparse-compiler=enable-runtime-library=false | \
// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
// RUN: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
// RUN: FileCheck %s
#CSR = #sparse_tensor.encoding<{
dimLevelType = [ "dense", "compressed" ],
@ -221,8 +227,10 @@ module {
//
// Sanity check on nonzeros.
//
// CHECK: ( 30.5, 4.2, 4.6, 7, 8, -1, -1, -1 )
// CHECK: ( 30.5, 4.2, 4.6, 7, 8, -1, -1, -1 )
// FIXME: bring this back once dense2sparse skips zeros
//
// C_HECK: ( 30.5, 4.2, 4.6, 7, 8 )
// C_HECK: ( 30.5, 4.2, 4.6, 7, 8 )
//
%val7 = sparse_tensor.values %7 : tensor<4x4xf64, #CSR> to memref<?xf64>
%val8 = sparse_tensor.values %8 : tensor<4x4xf64, #DCSR> to memref<?xf64>