[mlir][sparse] Improve the implementation of sparse_tensor.new for the codegen path.

[bixia]

Summary

Rewrite a NewOp into a NewOp of a sorted COO tensor and a ConvertOp for
converting the sorted COO tensor to the destination tensor type.

Codegen a NewOp of a sorted COO tensor to use the new bulk reader API and sort
the elements only when the input is not sorted.

[Peiming]

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

1268

is it a dimRank or lvlRank? If you are reading in a COO, it should be lvlRank here? @wrengr

[aartbik]

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

1245

This needs a comment. All other new op situation are handled by the rewriting, and direct IR codegen only deals with any remaining COO new ops
(or something of the sort)

1256

indentation?

1299

getNextFuncName is a bit strange name for the ReaderRead call?

1324

This needs more documentation, since you have two tests ;-)

At compile time, you decide whether target needs ordering in the last level (and thus ordered COO).
So document that before L1324

Then, you have a run-time test on the notSorted boolean that comes out of getSparseTensorReaderRead(
Make that a bit more clear on L1325
(and yes, you have it at L1256 but good to repeat here.

[bixia]

Address review comments.

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

1268

I think it is dimRank. In line 1274, we use it to get the dimSizes of the tensor, for dynamic shape. In line 1305, we use it for the size of dim2lvl mapping.

1299

Renamed to getReadFuncName.

1324

Add comment.

[bixia]

Rebase.

[bixia]

Fix codegen.mlir

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

matchAndRewrite(UnpackOp op, OpAdaptor adaptor,
     Value nnz = toType(rewriter, loc, desc.getValMemSize(rewriter, loc),
                        op.getNnz().getType());
 
     rewriter.replaceOp(op, {data, indices, nnz});
     return success();
   }
 };
 
+struct SparseNewOpConverter : public OpConversionPattern<NewOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(NewOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    const auto dstTp = getSparseTensorType(op.getResult());
+    const auto encDst = dstTp.getEncoding();
+    // Creating COO with NewOp is handled by direct IR codegen. All other cases
+    // are handled by rewriting.
+    if (!dstTp.hasEncoding() || getCOOStart(encDst) != 0)
+      return failure();
+
+    // Implement the NewOp(filename) as follows:
+    //   reader = getSparseTensorReader(filename)
+    //   nse = getSparseTensorNNZ()
+    //   tmp = bufferization.alloc_tensor an ordered COO with
+    //          dst dim ordering, size_hint = nse
+    //   indices = to_indices_buffer(tmp)
+    //   values = to_values(tmp)
+    //   isSorted = getSparseTensorReaderRead(indices, values, dimOrdering)
+    //   if (!isSorted) sort_coo(nse, indices, values)
+    //   update storage specifier
+    //   dst = sparse_tensor.ConvertOp tmp
+
+    // Create a sparse tensor reader.
+    Value fileName = op.getSource();
+    Type opaqueTp = getOpaquePointerType(rewriter);
+    Value reader = createFuncCall(rewriter, loc, "createSparseTensorReader",
+                                  {opaqueTp}, {fileName}, EmitCInterface::Off)
+                       .getResult(0);
+
+    Type indexTp = rewriter.getIndexType();
+    const Dimension dimRank = dstTp.getDimRank();
+
+    // If the result tensor has dynamic dimensions, get the dynamic sizes from
+    // the sparse tensor reader.
+    SmallVector<Value> dynSizes;
+    if (dstTp.hasDynamicDimShape()) {
+      Value dimSizes = genAlloca(rewriter, loc, dimRank, indexTp);
+      createFuncCall(rewriter, loc, "copySparseTensorReaderDimSizes", {},
+                     {reader, dimSizes}, EmitCInterface::On)
+          .getResult(0);
+      ArrayRef<int64_t> dstShape = dstTp.getRankedTensorType().getShape();
+      for (auto &d : llvm::enumerate(dstShape)) {
+        if (d.value() == ShapedType::kDynamic) {
+          dynSizes.push_back(rewriter.create<memref::LoadOp>(
+              loc, dimSizes, constantIndex(rewriter, loc, d.index())));
+        }
+      }
+    }
+
+    Value nse = createFuncCall(rewriter, loc, "getSparseTensorReaderNNZ",
+                               {indexTp}, {reader}, EmitCInterface::Off)
+                    .getResult(0);
+    // Construct allocation for each field.
+    SmallVector<Value> fields;
+    createAllocFields(rewriter, loc, dstTp, dynSizes, /*enableInit=*/false,
+                      fields, nse);
+    MutSparseTensorDescriptor desc(dstTp, fields);
+
+    // Read the COO tensor data.
+    Type eltTp = dstTp.getElementType();
+    Type indBufEleTp = getIndexOverheadType(rewriter, encDst);
+    SmallString<32> getReadFuncName{"getSparseTensorReaderRead",
+                                    overheadTypeFunctionSuffix(indBufEleTp),
+                                    primaryTypeFunctionSuffix(eltTp)};
+
+    Value xs = desc.getAOSMemRef();
+    Value ys = desc.getValMemRef();
+    SmallVector<Value> dim2lvlValues(dimRank, Value());
+    if (auto dimOrder = encDst.getDimOrdering()) {
+      assert(dimOrder.isPermutation() && "Got non-permutation");
+      for (uint64_t l = 0; l < dimRank; l++) {
+        uint64_t d = dimOrder.getDimPosition(l);
+        dim2lvlValues[d] = constantIndex(rewriter, loc, l);
+      }
+    } else {
+      for (uint64_t l = 0; l < dimRank; l++)
+        dim2lvlValues[l] = constantIndex(rewriter, loc, l);
+    }
+    Value dim2lvl = allocaBuffer(rewriter, loc, dim2lvlValues);
+
+    Value f = constantI1(rewriter, loc, false);
+    Value isSorted =
+        createFuncCall(rewriter, loc, getReadFuncName, {f.getType()},
+                       {reader, dim2lvl, xs, ys}, EmitCInterface::On)
+            .getResult(0);
+
+    // If the destination tensor is a sorted COO, we need to sort the COO tensor
+    // data if the input elements aren't sorted yet.
+    if (encDst.isOrderedLvl(dimRank - 1)) {
+      Value notSorted = rewriter.create<arith::CmpIOp>(
+          loc, arith::CmpIPredicate::eq, isSorted, f);
+      scf::IfOp ifOp =
+          rewriter.create<scf::IfOp>(loc, notSorted, /*else*/ false);
+      rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
+      rewriter.create<SortCooOp>(
+          loc, nse, xs, ValueRange{ys}, rewriter.getIndexAttr(dimRank),
+          rewriter.getIndexAttr(0), SparseTensorSortKind::HybridQuickSort);
+      rewriter.setInsertionPointAfter(ifOp);
+    }
+
+    // Set PtrMemRef0[1] = nse.
+    Value c1 = constantIndex(rewriter, loc, 1);
+    Value ptrMemref0 = desc.getPtrMemRef(0);
+    Type ptrEleTy = getMemRefType(ptrMemref0).getElementType();
+    Value ptrNse =
+        ptrEleTy == nse.getType()
+            ? nse
+            : rewriter.create<arith::IndexCastOp>(loc, ptrEleTy, nse);
+    rewriter.create<memref::StoreOp>(loc, ptrNse, ptrMemref0, c1);
+
+    // Update storage specifier.
+    Value idxSize = rewriter.create<arith::MulIOp>(
+        loc, nse, constantIndex(rewriter, loc, dimRank));
+    desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::IdxMemSize, 0,
+                           idxSize);
+    desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::ValMemSize,
+                           std::nullopt, nse);
+
+    // Release the sparse tensor reader.
+    createFuncCall(rewriter, loc, "delSparseTensorReader", {}, {reader},
+                   EmitCInterface::Off);
+
+    // Replace operation with resulting memrefs.
+    rewriter.replaceOp(op, genTuple(rewriter, loc, dstTp, fields));
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     bool enableBufferInitialization) {
   patterns.add<SparsePackOpConverter, SparseUnpackOpConverter,
                SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
                SparseCastConverter, SparseTensorDeallocConverter,
                SparseExtractSliceCoverter, SparseTensorLoadConverter,
                SparseExpandConverter, SparseCompressConverter,
                SparseInsertConverter, SparseToPointersConverter,
                SparseToIndicesConverter, SparseToIndicesBufferConverter,
                SparseToValuesConverter, SparseConvertConverter,
-               SparseNumberOfEntriesConverter>(typeConverter,
-                                               patterns.getContext());
+               SparseNewOpConverter, SparseNumberOfEntriesConverter>(
+      typeConverter, patterns.getContext());
   patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
                                            enableBufferInitialization);
 }

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

 /// Sparse rewriting rule for the new operator.
 struct NewRewriter : public OpRewritePattern<NewOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(NewOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     const auto dstTp = getSparseTensorType(op.getResult());
     const auto encDst = dstTp.getEncoding();
-    if (!dstTp.hasEncoding())
+    if (!dstTp.hasEncoding() || getCOOStart(encDst) == 0)
       return failure();
 
-    // Create a sparse tensor reader.
-    Value fileName = op.getSource();
-    Type opaqueTp = getOpaquePointerType(rewriter);
-    Value reader = createFuncCall(rewriter, loc, "createSparseTensorReader",
-                                  {opaqueTp}, {fileName}, EmitCInterface::Off)
-                       .getResult(0);
-
-    // Allocate a temporary buffer for storing dimension sizes and indices.
-    Type indexTp = rewriter.getIndexType();
-    const Dimension dimRank = dstTp.getDimRank();
-    Value dimSizes = genAlloca(rewriter, loc, dimRank, indexTp);
-
-    // If the result tensor has dynamic dimensions, get the dynamic sizes from
-    // the sparse tensor reader.
-    SmallVector<Value> dynSizesArray;
-    if (dstTp.hasDynamicDimShape()) {
-      createFuncCall(rewriter, loc, "copySparseTensorReaderDimSizes", {},
-                     {reader, dimSizes}, EmitCInterface::On)
-          .getResult(0);
-      ArrayRef<int64_t> dstShape = dstTp.getRankedTensorType().getShape();
-      for (auto &d : llvm::enumerate(dstShape)) {
-        if (d.value() == ShapedType::kDynamic) {
-          dynSizesArray.push_back(rewriter.create<memref::LoadOp>(
-              loc, dimSizes, constantIndex(rewriter, loc, d.index())));
-        }
-      }
-    }
-
     // Implement the NewOp as follows:
-    //   %tmp = bufferization.alloc_tensor : an unordered COO with identity
-    //                                       storage ordering
-    //   for i = 0 to nnz
-    //     get the next element from the input file
-    //     insert the element to %tmp
-    //   %t = sparse_tensor.ConvertOp %tmp
-    Value c0 = constantIndex(rewriter, loc, 0);
-    Value c1 = constantIndex(rewriter, loc, 1);
-    Value nnz = createFuncCall(rewriter, loc, "getSparseTensorReaderNNZ",
-                               {indexTp}, {reader}, EmitCInterface::Off)
-                    .getResult(0);
+    //   %orderedCoo = sparse_tensor.new %filename
+    //   %t = sparse_tensor.ConvertOp %orderedCoo
     RankedTensorType cooTp =
-        getUnorderedCOOFromTypeWithOrdering(dstTp, dstTp.getDimToLvlMap());
-    Value cooBuffer =
-        rewriter
-            .create<AllocTensorOp>(loc, cooTp, dynSizesArray, Value(),
-                                   /*sizeHint=*/nnz, Attribute())
-            .getResult();
-
-    Type eltTp = dstTp.getElementType();
-    Value value = genAllocaScalar(rewriter, loc, eltTp);
-    scf::ForOp forOp = rewriter.create<scf::ForOp>(loc, c0, nnz, c1,
-                                                   ArrayRef<Value>(cooBuffer));
-    rewriter.setInsertionPointToStart(forOp.getBody());
-
-    SmallString<29> getNextFuncName{"getSparseTensorReaderNext",
-                                    primaryTypeFunctionSuffix(eltTp)};
-    Value indices = dimSizes; // Reuse the indices memref to store indices.
-    createFuncCall(rewriter, loc, getNextFuncName, {}, {reader, indices, value},
-                   EmitCInterface::On);
-    SmallVector<Value> indicesArray(dimRank);
-    for (Dimension d = 0; d < dimRank; d++) {
-      // FIXME: `toStoredDim` is deprecated
-      indicesArray[toStoredDim(encDst, d)] = rewriter.create<memref::LoadOp>(
-          loc, indices, constantIndex(rewriter, loc, d));
-    }
-    Value v = rewriter.create<memref::LoadOp>(loc, value);
-    Value t = rewriter.create<InsertOp>(loc, v, forOp.getRegionIterArg(0),
-                                        indicesArray);
-    rewriter.create<scf::YieldOp>(loc, ArrayRef<Value>(t));
-    rewriter.setInsertionPointAfter(forOp);
-    // Link SSA chain.
-    cooBuffer = forOp.getResult(0);
-
-    // Release the sparse tensor reader.
-    createFuncCall(rewriter, loc, "delSparseTensorReader", {}, {reader},
-                   EmitCInterface::Off);
-    cooBuffer = rewriter.create<LoadOp>(loc, cooBuffer, true);
-    Value newOp = rewriter.replaceOpWithNewOp<ConvertOp>(
-        op, dstTp.getRankedTensorType(), cooBuffer);
-
-    // Release the unordered COO tensor buffer.
-    rewriter.setInsertionPointAfterValue(newOp);
-    rewriter.create<DeallocTensorOp>(loc, cooBuffer);
+        getCOOFromTypeWithOrdering(dstTp, encDst.getDimOrdering(), true);
+    Value cooTensor = rewriter.create<NewOp>(loc, cooTp, op.getSource());
+    Value convert = rewriter.replaceOpWithNewOp<ConvertOp>(
+        op, dstTp.getRankedTensorType(), cooTensor);
+
+    // Release the ordered COO tensor.
+    rewriter.setInsertionPointAfterValue(convert);
+    rewriter.create<DeallocTensorOp>(loc, cooTensor);
 
     return success();
   }
 };
 
 struct OutRewriter : public OpRewritePattern<OutOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(OutOp op,

mlir/test/Dialect/SparseTensor/codegen.mlir

#Dense3D = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "dense", "dense" ],
   dimOrdering = affine_map<(i, j, k) -> (k, i, j)>
 }>
 
 #Coo = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed-nu", "singleton" ]
 }>
 
+#CooPNo = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton-no" ],
+  dimOrdering = affine_map<(i, j) -> (j, i)>
+}>
+
 #ccoo = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed-nu", "singleton" ]
 }>
 
 // CHECK-LABEL: func @sparse_nop(
 //  CHECK-SAME: %[[A0:.*]]: memref<?xi32>,
 //  CHECK-SAME: %[[A1:.*]]: memref<?xi64>,
 //  CHECK-SAME: %[[A2:.*]]: memref<?xf64>,

 //  CHECK-SAME: %[[A3:.*]]: memref<?xf32>,
 //  CHECK-SAME: %[[A4:.*]]: !sparse_tensor.storage_specifier
 //       CHECK: return  %[[A1]], %[[A2]], %[[A3]], %[[A4]] :
 //  CHECK-SAME:   memref<?xi32>, memref<?xi64>, memref<?xf32>, !sparse_tensor.storage_specifier
 func.func @sparse_nop_convert(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
   %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor<?xf32, #SparseVector>
   return %0 : tensor<?xf32, #SparseVector>
 }
+
+// CHECK-LABEL: func.func @sparse_new_coo(
+// CHECK-SAME:  %[[A0:.*]]: !llvm.ptr<i8>) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>>) {
+//   CHECK-DAG: %[[A1:.*]] = arith.constant false
+//   CHECK-DAG: %[[A2:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[A3:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[A4:.*]] = arith.constant 2 : index
+//       CHECK: %[[A5:.*]] = call @createSparseTensorReader(%[[A0]])
+//       CHECK: %[[A6:.*]] = memref.alloca() : memref<2xindex>
+//       CHECK: %[[A7:.*]] = memref.cast %[[A6]] : memref<2xindex> to memref<?xindex>
+//       CHECK: call @copySparseTensorReaderDimSizes(%[[A5]], %[[A7]]) : (!llvm.ptr<i8>, memref<?xindex>) -> ()
+//       CHECK: %[[A8:.*]] = memref.load %[[A6]]{{\[}}%[[A3]]] : memref<2xindex>
+//       CHECK: %[[A9:.*]] = memref.load %[[A6]]{{\[}}%[[A2]]] : memref<2xindex>
+//       CHECK: %[[A10:.*]] = call @getSparseTensorReaderNNZ(%[[A5]])
+//       CHECK: %[[A11:.*]] = arith.muli %[[A10]], %[[A4]] : index
+//       CHECK: %[[A12:.*]] = memref.alloc() : memref<2xindex>
+//       CHECK: %[[A13:.*]] = memref.cast %[[A12]] : memref<2xindex> to memref<?xindex>
+//       CHECK: %[[A14:.*]] = memref.alloc(%[[A11]]) : memref<?xindex>
+//       CHECK: %[[A15:.*]] = memref.alloc(%[[A10]]) : memref<?xf32>
+//       CHECK: %[[A16:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>>
+//       CHECK: %[[A17:.*]] = arith.index_cast %[[A8]] : index to i64
+//       CHECK: %[[A18:.*]] = sparse_tensor.storage_specifier.set %[[A16]]  dim_sz at 0 with %[[A17]]
+//       CHECK: %[[A19:.*]] = sparse_tensor.storage_specifier.get %[[A18]]  ptr_mem_sz at 0
+//       CHECK: %[[A20:.*]] = arith.index_cast %[[A19]] : i64 to index
+//       CHECK: %[[A21:.*]], %[[A22:.*]] = sparse_tensor.push_back %[[A20]], %[[A13]], %[[A3]]
+//       CHECK: %[[A23:.*]] = arith.index_cast %[[A22]] : index to i64
+//       CHECK: %[[A24:.*]] = sparse_tensor.storage_specifier.set %[[A18]]  ptr_mem_sz at 0 with %[[A23]]
+//       CHECK: %[[A25:.*]] = arith.index_cast %[[A9]] : index to i64
+//       CHECK: %[[A26:.*]] = sparse_tensor.storage_specifier.set %[[A24]]  dim_sz at 1 with %[[A25]]
+//       CHECK: %[[A27:.*]], %[[A28:.*]] = sparse_tensor.push_back %[[A22]], %[[A21]], %[[A3]], %[[A2]]
+//       CHECK: %[[A29:.*]] = arith.index_cast %[[A28]] : index to i64
+//       CHECK: %[[A30:.*]] = sparse_tensor.storage_specifier.set %[[A26]]  ptr_mem_sz at 0 with %[[A29]]
+//       CHECK: %[[A31:.*]] = memref.alloca() : memref<2xindex>
+//       CHECK: %[[A32:.*]] = memref.cast %[[A31]] : memref<2xindex> to memref<?xindex>
+//       CHECK: memref.store %[[A3]], %[[A31]]{{\[}}%[[A3]]] : memref<2xindex>
+//       CHECK: memref.store %[[A2]], %[[A31]]{{\[}}%[[A2]]] : memref<2xindex>
+//       CHECK: %[[A33:.*]] = call @getSparseTensorReaderRead0F32(%[[A5]], %[[A32]], %[[A14]], %[[A15]])
+//       CHECK: %[[A34:.*]] = arith.cmpi eq, %[[A33]], %[[A1]] : i1
+//       CHECK: scf.if %[[A34]] {
+//       CHECK:   sparse_tensor.sort_coo  hybrid_quick_sort %[[A10]], %[[A14]] jointly %[[A15]] {nx = 2 : index, ny = 0 : index} : memref<?xindex> jointly memref<?xf32>
+//       CHECK: }
+//       CHECK: memref.store %[[A10]], %[[A27]]{{\[}}%[[A2]]] : memref<?xindex>
+//       CHECK: %[[A35:.*]] = arith.index_cast %[[A11]] : index to i64
+//       CHECK: %[[A36:.*]] = sparse_tensor.storage_specifier.set %[[A30]]  idx_mem_sz at 0 with %[[A35]]
+//       CHECK: %[[A37:.*]] = arith.index_cast %[[A10]] : index to i64
+//       CHECK: %[[A38:.*]] = sparse_tensor.storage_specifier.set %[[A36]]  val_mem_sz with %[[A37]]
+//       CHECK: call @delSparseTensorReader(%[[A5]]) : (!llvm.ptr<i8>) -> ()
+//       CHECK: return %[[A27]], %[[A14]], %[[A15]], %[[A38]]
+func.func @sparse_new_coo(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #Coo> {
+  %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #Coo>
+  return %0 : tensor<?x?xf32, #Coo>
+}
+
+// CHECK-LABEL: func.func @sparse_new_coo_permute_no(
+//  CHECK-SAME: %[[A0:.*]]: !llvm.ptr<i8>) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>>) {
+//   CHECK-DAG: %[[A1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[A2:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[A3:.*]] = arith.constant 2 : index
+//       CHECK: %[[A4:.*]] = call @createSparseTensorReader(%[[A0]])
+//       CHECK: %[[A5:.*]] = memref.alloca() : memref<2xindex>
+//       CHECK: %[[A6:.*]] = memref.cast %[[A5]] : memref<2xindex> to memref<?xindex>
+//       CHECK: call @copySparseTensorReaderDimSizes(%[[A4]], %[[A6]])
+//       CHECK: %[[A7:.*]] = memref.load %[[A5]]{{\[}}%[[A2]]] : memref<2xindex>
+//       CHECK: %[[A8:.*]] = memref.load %[[A5]]{{\[}}%[[A1]]] : memref<2xindex>
+//       CHECK: %[[A9:.*]] = call @getSparseTensorReaderNNZ(%[[A4]])
+//       CHECK: %[[A10:.*]] = arith.muli %[[A9]], %[[A3]] : index
+//       CHECK: %[[A11:.*]] = memref.alloc() : memref<2xindex>
+//       CHECK: %[[A12:.*]] = memref.cast %[[A11]] : memref<2xindex> to memref<?xindex>
+//       CHECK: %[[A13:.*]] = memref.alloc(%[[A10]]) : memref<?xindex>
+//       CHECK: %[[A14:.*]] = memref.alloc(%[[A9]]) : memref<?xf32>
+//       CHECK: %[[A15:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ] }>>
+//       CHECK: %[[A16:.*]] = arith.index_cast %[[A8]] : index to i64
+//       CHECK: %[[A17:.*]] = sparse_tensor.storage_specifier.set %[[A15]]  dim_sz at 0 with %[[A16]]
+//       CHECK: %[[A18:.*]] = sparse_tensor.storage_specifier.get %[[A17]]  ptr_mem_sz at 0
+//       CHECK: %[[A19:.*]] = arith.index_cast %[[A18]] : i64 to index
+//       CHECK: %[[A20:.*]], %[[A21:.*]] = sparse_tensor.push_back %[[A19]], %[[A12]], %[[A2]]
+//       CHECK: %[[A22:.*]] = arith.index_cast %[[A21]] : index to i64
+//       CHECK: %[[A23:.*]] = sparse_tensor.storage_specifier.set %[[A17]]  ptr_mem_sz at 0 with %[[A22]]
+//       CHECK: %[[A24:.*]] = arith.index_cast %[[A7]] : index to i64
+//       CHECK: %[[A25:.*]] = sparse_tensor.storage_specifier.set %[[A23]]  dim_sz at 1 with %[[A24]]
+//       CHECK: %[[A26:.*]], %[[A27:.*]] = sparse_tensor.push_back %[[A21]], %[[A20]], %[[A2]], %[[A1]]
+//       CHECK: %[[A28:.*]] = arith.index_cast %[[A27]] : index to i64
+//       CHECK: %[[A29:.*]] = sparse_tensor.storage_specifier.set %[[A25]]  ptr_mem_sz at 0 with %[[A28]]
+//       CHECK: %[[A30:.*]] = memref.alloca() : memref<2xindex>
+//       CHECK: %[[A31:.*]] = memref.cast %[[A30]] : memref<2xindex> to memref<?xindex>
+//       CHECK: memref.store %[[A1]], %[[A30]]{{\[}}%[[A2]]] : memref<2xindex>
+//       CHECK: memref.store %[[A2]], %[[A30]]{{\[}}%[[A1]]] : memref<2xindex>
+//       CHECK: %[[A32:.*]] = call @getSparseTensorReaderRead0F32(%[[A4]], %[[A31]], %[[A13]], %[[A14]])
+//       CHECK: memref.store %[[A9]], %[[A26]]{{\[}}%[[A1]]] : memref<?xindex>
+//       CHECK: %[[A33:.*]] = arith.index_cast %[[A10]] : index to i64
+//       CHECK: %[[A34:.*]] = sparse_tensor.storage_specifier.set %[[A29]]  idx_mem_sz at 0 with %[[A33]]
+//       CHECK: %[[A35:.*]] = arith.index_cast %[[A9]] : index to i64
+//       CHECK: %[[A36:.*]] = sparse_tensor.storage_specifier.set %[[A34]]  val_mem_sz with %[[A35]]
+//       CHECK: call @delSparseTensorReader(%[[A4]]) : (!llvm.ptr<i8>) -> ()
+//       CHECK: return %[[A26]], %[[A13]], %[[A14]], %[[A36]]
+func.func @sparse_new_coo_permute_no(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CooPNo> {
+  %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #CooPNo>
+  return %0 : tensor<?x?xf32, #CooPNo>
+}
 No newline at end of file

mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir

 // RUN: mlir-opt %s -post-sparsification-rewrite="enable-runtime-library=false enable-convert=false" | \
 // RUN: FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]
 }>
 
 #CSC = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   dimOrdering = affine_map<(i, j) -> (j, i)>
 }>
 
+#COO = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ]
+}>
+
 // CHECK-LABEL:   func.func @sparse_new(
 // CHECK-SAME:    %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> {
-// CHECK-DAG:     %[[C2:.*]] = arith.constant 2 : index
-// CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
-// CHECK-DAG:     %[[C1:.*]] = arith.constant 1 : index
-// CHECK:         %[[R:.*]] = call @createSparseTensorReader(%[[A]])
-// CHECK:         %[[DS:.*]] = memref.alloca(%[[C2]]) : memref<?xindex>
-// CHECK:         call @copySparseTensorReaderDimSizes(%[[R]], %[[DS]])
-// CHECK:         %[[D0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]]
-// CHECK:         %[[D1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]]
-// CHECK:         %[[N:.*]] = call @getSparseTensorReaderNNZ(%[[R]])
-// CHECK:         %[[T:.*]] = bufferization.alloc_tensor(%[[D0]], %[[D1]]) size_hint=%[[N]]
-// CHECK:         %[[VB:.*]] = memref.alloca()
-// CHECK:         %[[T2:.*]] = scf.for %{{.*}} = %[[C0]] to %[[N]] step %[[C1]] iter_args(%[[A2:.*]] = %[[T]])
-// CHECK:           func.call @getSparseTensorReaderNextF32(%[[R]], %[[DS]], %[[VB]])
-// CHECK:           %[[E0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]]
-// CHECK:           %[[E1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]]
-// CHECK:           %[[V:.*]] = memref.load %[[VB]][]
-// CHECK:           %[[T1:.*]] = sparse_tensor.insert %[[V]] into %[[A2]]{{\[}}%[[E0]], %[[E1]]]
-// CHECK:           scf.yield %[[T1]]
-// CHECK:         }
-// CHECK:         call @delSparseTensorReader(%[[R]])
-// CHECK:         %[[T4:.*]] = sparse_tensor.load %[[T2]] hasInserts
-// CHECK:         %[[R:.*]] = sparse_tensor.convert %[[T4]]
-// CHECK:         bufferization.dealloc_tensor %[[T4]]
+// CHECK:         %[[COO:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr<i8> to tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>
+// CHECK:         %[[R:.*]] = sparse_tensor.convert %[[COO]]
+// CHECK:         bufferization.dealloc_tensor %[[COO]]
 // CHECK:         return %[[R]]
 func.func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CSR> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #CSR>
   return %0 : tensor<?x?xf32, #CSR>
 }
 
 // CHECK-LABEL:   func.func @sparse_new_csc(
 // CHECK-SAME:    %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>> {
-// CHECK-DAG:     %[[C2:.*]] = arith.constant 2 : index
-// CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
-// CHECK-DAG:     %[[C1:.*]] = arith.constant 1 : index
-// CHECK:         %[[R:.*]] = call @createSparseTensorReader(%[[A]])
-// CHECK:         %[[DS:.*]] = memref.alloca(%[[C2]]) : memref<?xindex>
-// CHECK:         call @copySparseTensorReaderDimSizes(%[[R]], %[[DS]])
-// CHECK:         %[[D0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]]
-// CHECK:         %[[D1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]]
-// CHECK:         %[[N:.*]] = call @getSparseTensorReaderNNZ(%[[R]])
-// CHECK:         %[[T:.*]] = bufferization.alloc_tensor(%[[D0]], %[[D1]]) size_hint=%[[N]]
-// CHECK:         %[[VB:.*]] = memref.alloca()
-// CHECK:         %[[T2:.*]] = scf.for %{{.*}} = %[[C0]] to %[[N]] step %[[C1]] iter_args(%[[A2:.*]] = %[[T]])
-// CHECK:           func.call @getSparseTensorReaderNextF32(%[[R]], %[[DS]], %[[VB]])
-// CHECK:           %[[E0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]]
-// CHECK:           %[[E1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]]
-// CHECK:           %[[V:.*]] = memref.load %[[VB]][]
-// CHECK:           %[[T1:.*]] = sparse_tensor.insert %[[V]] into %[[A2]]{{\[}}%[[E1]], %[[E0]]]
-// CHECK:           scf.yield %[[T1]]
-// CHECK:         }
-// CHECK:         call @delSparseTensorReader(%[[R]])
-// CHECK:         %[[T4:.*]] = sparse_tensor.load %[[T2]] hasInserts
-// CHECK:         %[[R:.*]] = sparse_tensor.convert %[[T4]]
-// CHECK:         bufferization.dealloc_tensor %[[T4]]
+// CHECK:         %[[COO:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr<i8> to tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>>
+// CHECK:         %[[R:.*]] = sparse_tensor.convert %[[COO]]
+// CHECK:         bufferization.dealloc_tensor %[[COO]]
 // CHECK:         return %[[R]]
 func.func @sparse_new_csc(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CSC> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #CSC>
   return %0 : tensor<?x?xf32, #CSC>
 }
 
+// CHECK-LABEL:   func.func @sparse_new_coo(
+// CHECK-SAME:    %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> {
+// CHECK:         %[[COO:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr<i8> to tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>
+// CHECK:         return %[[COO]]
+func.func @sparse_new_coo(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #COO> {
+  %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #COO>
+  return %0 : tensor<?x?xf32, #COO>
+}
+
 // CHECK-LABEL:   func.func @sparse_out(
 // CHECK-SAME:    %[[A:.*]]: tensor<10x20xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>,
 // CHECK-SAME:    %[[B:.*]]: !llvm.ptr<i8>) {
 // CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[C1:.*]] = arith.constant 1 : index
 // CHECK-DAG:     %[[C2:.*]] = arith.constant 2 : index
 // CHECK-DAG:     %[[C10:.*]] = arith.constant 10 : index
 // CHECK-DAG:     %[[C20:.*]] = arith.constant 20 : index