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

[aartbik]

Summary

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

[aartbik]

rebased with main

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

952

Should this first assert that the buffer has an appropriate type to be the argument to DeallocOp? I'm thinking mainly about ensuring nice error messages for debugging (since I'm guessing the error would only be caught rather further down down the compilation pipeline).

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

149

The type of buffer and the actual implementation are a lot more general than this function name would indicate. Perhaps rename to deallocMemref or deallocBuffer?

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

952

agreed, also pushed to the next revision (so this one is just moving it out)

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

149

agreed, but since I just migrated existing code out, and we have allocDenseTensor, let's do that as a generic follow-up on all these methods

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

 
 /// Generates code to allocate a buffer of the given type, and zero
 /// initialize it.  If the buffer type has any dynamic sizes, then the
 /// `sizes` parameter should be as filled by sizesFromPtr(); that way
 /// we can reuse the genDimSizeCall() results generated by sizesFromPtr().
 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.
 ///    if (tensor[ivs] != 0)
 ///      insert_point
 Value genValueForDense(OpBuilder &builder, Location loc, Value tensor,
                        ValueRange ivs);
 

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

if (shape[i] == ShapedType::kDynamicSize)
       dynamicSizes.push_back(sizes[i]);
   }
   Value mem = builder.create<memref::AllocOp>(loc, memTp, dynamicSizes);
   Value zero = constantZero(builder, loc, elemTp);
   builder.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{mem});
   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);
   Value cond = genIsNonzero(builder, loc, val);
   scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else*/ false);
   builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   return val;
 }

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

static Value genGetNextCall(OpBuilder &builder, Location loc, Value iter,
   Type elemTp = elemPtr.getType().cast<ShapedType>().getElementType();
   SmallString<10> name{"getNext", primaryTypeFunctionSuffix(elemTp)};
   SmallVector<Value, 3> params{iter, ind, elemPtr};
   Type i1 = builder.getI1Type();
   return createFuncCall(builder, loc, name, i1, params, EmitCInterface::On)
       .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,
                                          unsigned rank, Value ind,
                                          unsigned offsetDim = 0,
                                          Value offset = Value()) {
   SmallVector<Value, 4> ivs;
   ivs.reserve(rank);

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

LogicalResult sparse2DenseRewrite(ConvertOp op,
                                     PatternRewriter &rewriter) const {
     Location loc = op->getLoc();
     RankedTensorType dstTp = op.getType().cast<RankedTensorType>();
     Value src = op.getSource();
     RankedTensorType srcTp = src.getType().cast<RankedTensorType>();
 
     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,
                                    ValueRange args, Value v, ValueRange reduc) {
                                  builder.create<memref::StoreOp>(loc, v, dst,
                                                                  args);
                                  builder.create<sparse_tensor::YieldOp>(loc);
                                });
 
     rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstTp, dst);
+
+    // Deallocate the buffer.
+    if (noEscape) {
+      rewriter.setInsertionPoint(insertionBlock->getTerminator());
+      deallocDenseTensor(rewriter, loc, dst);
+    }
     return success();
   }
 
   // Handles sparse tensor to sparse tensor conversion as follows:
   //   if src is not COO
   //       construct a COO to represent the src
   //   sort the src COO
   //   foreach elemment in the sorted src COO

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // 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 parallelization.
 //
 // RUN: mlir-opt %s --sparse-compiler="parallelization-strategy=any-storage-any-loop" | \
 // 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" ],
   dimOrdering = affine_map<(i,j) -> (i,j)>
 }>
 
 #DCSR = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ],
func.func @entry() {
     //
     %c8 = sparse_tensor.convert %8 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
     %v8 = vector.transfer_read %c8[%c0, %c0], %d1 : tensor<4x4xf64>, vector<4x4xf64>
     vector.print %v8 : vector<4x4xf64>
 
     //
     // 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>
     %nz7 = vector.transfer_read %val7[%c0], %d1 : memref<?xf64>, vector<8xf64>
     %nz8 = vector.transfer_read %val8[%c0], %d1 : memref<?xf64>, vector<8xf64>
     vector.print %nz7 : vector<8xf64>
     vector.print %nz8 : vector<8xf64>