[mlir][sparse] compute allocation size_hint

This adds the hint to a number of tensor allocations in codegens,
shaving off quite some time from e.g. reading in sparse matrices
due to zero-reallocation scheme. Note that we can probably provide
hints on all allocations, and refine the heuristics that use them
for general tensors.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D143309

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

@@ -229,6 +229,10 @@ static void createAllocFields(OpBuilder &builder, Location loc, Type type,
       ptrHeuristic = constantIndex(builder, loc, 2);
       idxHeuristic = builder.create<arith::MulIOp>(
           loc, constantIndex(builder, loc, rank), sizeHint); // AOS
+    } else if (rank == 2 && isDenseDim(rtp, 0) && isCompressedDim(rtp, 1)) {
+      ptrHeuristic = builder.create<arith::AddIOp>(
+          loc, sizeHint, constantIndex(builder, loc, 1));
+      idxHeuristic = sizeHint;
     } else {
       ptrHeuristic = idxHeuristic = constantIndex(builder, loc, 16);
     }