[mlir][sparse] add util for ToCoordinatesBuffer for COO AoS

[aartbik]

rebased with main

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

 
 /// Infers the result type and generates `ToCoordinatesOp`.  If the
 /// level is within a COO region, the result type is a memref with unknown
 /// stride and offset.  Otherwise, the result type is a memref without
 /// any specified layout.
 Value genToCoordinates(OpBuilder &builder, Location loc, Value tensor,
                        Level lvl, Level cooStart);
 
+/// Infers the result type and generates `ToCoordinatesBufferOp`.
+Value genToCoordinatesBuffer(OpBuilder &builder, Location loc, Value tensor);
+
 /// Infers the result type and generates `ToValuesOp`.
 Value genToValues(OpBuilder &builder, Location loc, Value tensor);
 
 /// Generates code to retrieve the values size for the sparse tensor.
 Value genValMemSize(OpBuilder &builder, Location loc, Value tensor);
 
 /// Generates code to retrieve the slice offset for the sparse tensor slice,
 /// return a constant if the offset is statically known.

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

Value sparse_tensor::genToCoordinates(OpBuilder &builder, Location loc,
                                       Value tensor, Level lvl, Level cooStart) {
   const auto srcTp = getSparseTensorType(tensor);
   const Type crdTp = srcTp.getEncoding().getCrdType();
   const Type memTp = get1DMemRefType(crdTp, /*withLayout=*/lvl >= cooStart);
   return builder.create<ToCoordinatesOp>(loc, memTp, tensor,
                                          builder.getIndexAttr(lvl));
 }
 
+Value sparse_tensor::genToCoordinatesBuffer(OpBuilder &builder, Location loc,
+                                            Value tensor) {
+  const auto srcTp = getSparseTensorType(tensor);
+  const Type crdTp = srcTp.getEncoding().getCrdType();
+  const Type memTp = get1DMemRefType(crdTp, /*withLayout=*/false);
+  return builder.create<ToCoordinatesBufferOp>(loc, memTp, tensor);
+}
+
 Value sparse_tensor::genToValues(OpBuilder &builder, Location loc,
                                  Value tensor) {
   RankedTensorType srcTp = getRankedTensorType(tensor);
   Type valTp = get1DMemRefType(srcTp.getElementType(),
                                /*withLayout=*/false);
   return builder.create<ToValuesOp>(loc, valTp, tensor);
 }
 

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

LogicalResult sparse2SparseRewrite(ConvertOp op,
     if (!srcTp.isAllOrdered()) {
       // Retrieve the values-array.
       Value y = genToValues(rewriter, loc, src);
       const auto encSrc = srcTp.getEncoding();
       // Sort the COO tensor so that its elements are ordered via increasing
       // coordinates for the storage ordering of the dst tensor.  Use SortCoo
       // if the COO tensor has the same ordering as the dst tensor.
       if (dimRank > 1 && srcTp.hasSameDimToLvlMap(dstTp)) {
-        MemRefType coordsTp =
-            get1DMemRefType(encSrc.getCrdType(), /*withLayout=*/false);
-        Value xs = rewriter.create<ToCoordinatesBufferOp>(loc, coordsTp, src);
+        Value xs = genToCoordinatesBuffer(rewriter, loc, src);
         rewriter.create<SortCooOp>(
             loc, nnz, xs, ValueRange{y}, rewriter.getIndexAttr(dimRank),
             rewriter.getIndexAttr(0), SparseTensorSortKind::HybridQuickSort);
       } else {
         // Gather the coordinates-arrays in the dst tensor storage order.
         SmallVector<Value> xs(dstLvlRank);
         const Level srcLvlRank = srcTp.getLvlRank();
         for (Level srcLvl = 0; srcLvl < srcLvlRank; srcLvl++) {