[mlir][sparse] avoid sorting when unnecessary when convert sparse tensors.

[Peiming]

minor fixes.

[aartbik]

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

52

can you please write this in the same style as above

if (auto enc = getSparseTensorEncoding(rtp))) {

...

}
return true

[Peiming]

address comments.

[Peiming]

rebase

[aartbik]

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

52

Yeah, reads much easier, thanks!

[Peiming]

remove redundant condition.

[Peiming]

minor fix.

[aartbik]

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

53

and it keeps getting better. ship it before no code is left! ;-)

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

 static bool isSparseTensor(OpOperand *op) {
   if (auto enc = getSparseTensorEncoding(op->get().getType())) {
     if (llvm::is_contained(enc.getDimLevelType(), DimLevelType::Compressed))
       return true;
   }
   return false;
 }
 
+static bool isAllDimOrdered(RankedTensorType rtp) {
+  if (auto enc = getSparseTensorEncoding(rtp))
+    return llvm::all_of(enc.getDimLevelType(), isOrderedDLT);
+
+  return true;
+}
+
+static bool hasSameDimOrdering(RankedTensorType rtp1, RankedTensorType rtp2) {
+  assert(rtp1.getRank() == rtp2.getRank());
+  AffineMap idMap =
+      AffineMap::getMultiDimIdentityMap(rtp1.getRank(), rtp1.getContext());
+
+  auto enc1 = getSparseTensorEncoding(rtp1);
+  auto enc2 = getSparseTensorEncoding(rtp2);
+
+  auto order1 = (enc1 && enc1.getDimOrdering()) ? enc1.getDimOrdering() : idMap;
+  auto order2 = (enc2 && enc2.getDimOrdering()) ? enc2.getDimOrdering() : idMap;
+
+  return order1 == order2;
+}
+
 // Helper method to find zero/uninitialized allocation.
 static bool isAlloc(OpOperand *op, bool isZero) {
   Value val = op->get();
   // Check allocation, with zero alloc when required.
   if (auto alloc = val.getDefiningOp<AllocTensorOp>()) {
     Value copy = alloc.getCopy();
     if (isZero)
       return copy && isZeroValue(copy);
LogicalResult sparse2SparseRewrite(ConvertOp op,
     Value src = op.getSource();
     RankedTensorType srcTp = src.getType().cast<RankedTensorType>();
     RankedTensorType dstTp = op.getType().cast<RankedTensorType>();
     SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
 
     SmallVector<Value> srcSizes;
     sizesForTensor(rewriter, srcSizes, loc, srcTp, src);
     Value tmpCoo = Value();
-    if (!isUniqueCOOType(srcTp)) {
+    // We need a tmp COO buffer if and only if
+    // 1. the src tensor is not a COO and
+    // 2. the src tensor is not ordered in the same way as the target
+    // tensor (e.g., src tensor is not ordered or src tensor haves a different
+    // dimOrdering).
+    if (!isUniqueCOOType(srcTp) &&
+        !(isAllDimOrdered(srcTp) && hasSameDimOrdering(srcTp, dstTp))) {
       // Construct a COO tensor from the src tensor.
       // TODO: there may be cases for which more efficiently without
       // going through an intermediate COO, such as cases that only change
       // the overhead types.
       SmallVector<Value> dynSrcSizes;
       getDynamicSizes(srcTp, srcSizes, dynSrcSizes);
       srcTp = getUnorderedCOOFromType(srcTp);
       tmpCoo =
           rewriter.create<AllocTensorOp>(loc, srcTp, dynSrcSizes).getResult();
       auto foreachOp = rewriter.create<ForeachOp>(
           loc, src, tmpCoo,
           [&](OpBuilder &builder, Location loc, ValueRange args, Value v,
               ValueRange reduc) {
             // The resulting COO tensor has identity ordering.
             auto t = builder.create<InsertOp>(loc, v, reduc.front(),
                                               args.slice(0, srcTp.getRank()));
             builder.create<sparse_tensor::YieldOp>(loc, t);
           });
       src = rewriter.create<LoadOp>(loc, foreachOp.getResult(0), true);
     }
 
+    // Only need to sort if the srcTp is not already sorted (we faithfully take
+    // the guarantee from the sparse tensor encoding).
+    if (!isAllDimOrdered(srcTp)) {
       // Sort the COO tensor so that its elements are ordered via increasing
       // indices for the storage ordering of the dst tensor.
       SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(srcTp);
       auto dynShape = {ShapedType::kDynamic};
       auto indTp =
           MemRefType::get(dynShape, getIndexOverheadType(rewriter, encSrc));
       uint64_t rank = dstTp.getRank();
       // Gather the indices-arrays in the dst tensor storage order.
       SmallVector<Value> xs(rank, Value());
       for (uint64_t i = 0; i < rank; i++) {
         uint64_t orgDim = toOrigDim(encSrc, i);
         xs[toStoredDim(encDst, orgDim)] = rewriter.create<ToIndicesOp>(
             loc, indTp, src, rewriter.getIndexAttr(i));
       }
 
       // Retrieve NNZ.
       Value nnz = rewriter.create<NumberOfEntriesOp>(loc, src);
-    nnz =
-        rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(), nnz);
+      nnz = rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(),
+                                                nnz);
 
       // Retrieve the values-array.
       auto valTp = MemRefType::get(dynShape, srcTp.getElementType());
       Value y = rewriter.create<ToValuesOp>(loc, valTp, src);
 
       // Sort the COO tensor.
       rewriter.create<SortOp>(loc, nnz, xs, ValueRange{y});
+    }
 
     // For each element in the COO tensor, insert the element to the dst tensor.
     SmallVector<Value> dynDstSizes;
     getDynamicSizes(dstTp, srcSizes, dynDstSizes);
     Value dst =
         rewriter.create<AllocTensorOp>(loc, dstTp, dynDstSizes).getResult();
     SmallVector<Value> indices(srcTp.getRank(), Value());
     auto foreachOp = rewriter.create<ForeachOp>(

mlir/test/Dialect/SparseTensor/convert_sparse2sparse.mlir

 //   CHECK-AUTO-DAG: %[[IotaP:.*]] = memref.cast %[[Iota]] : memref<1xindex> to memref<?xindex>
 //       CHECK-AUTO: %[[T:.*]] = call @newSparseTensor(%[[DimSizesP]], %[[LvlSizesP]], %[[LvlTypesP]], %[[IotaP]], %[[IotaP]], %{{.*}}, %{{.*}}, %{{.*}}, %[[SparseToSparse]], %[[A]])
 //       CHECK-AUTO: return %[[T]] : !llvm.ptr<i8>
 
 // CHECK-RWT-LABEL: func.func @sparse_convert(
 //  CHECK-RWT-SAME: %[[A:.*]]: tensor<?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ], pointerBitWidth = 64, indexBitWidth = 64 }>>)
 //  CHECK-RWT-DAG:  %[[C0:.*]] = arith.constant 0 : index
 //      CHECK-RWT:  %[[D:.*]] = tensor.dim %[[A]], %[[C0]]
-//      CHECK-RWT:  %[[I0:.*]] = sparse_tensor.indices %[[A]] {dimension = 0 : index}
-//      CHECK-RWT:  %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[A]]
-//      CHECK-RWT:  %[[V:.*]] = sparse_tensor.values %[[A]]
-//      CHECK-RWT:  sparse_tensor.sort %[[NNZ]], %[[I0]] jointly %[[V]]
 //      CHECK-RWT:  %[[DST:.*]] = bufferization.alloc_tensor(%[[D]])
 //      CHECK-RWT:  %[[RET:.*]] = sparse_tensor.foreach in %[[A]] init(%[[DST]])
 //      CHECK-RWT:  ^bb0(%[[FI2:.*]]: index, %[[FV2:.*]]: f32, %[[T:.*]]: tensor<?xf32,
 //      CHECK-RWT:    %[[I:.*]] = sparse_tensor.insert %[[FV2]] into %[[T]]{{\[}}%[[FI2]]]
 //      CHECK-RWT:    sparse_tensor.yield %[[I]]
 //      CHECK-RWT:  }
 //      CHECK-RWT:  %[[T:.*]] = sparse_tensor.load %[[RET]] hasInserts
 //      CHECK-RWT:  %[[R:.*]] = sparse_tensor.convert %[[T]]

 // CHECK-RWT-LABEL: func.func @sparse_convert_permuted(
 //  CHECK-RWT-SAME: %[[COO:.*]]:
 //   CHECK-RWT-DAG: %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-RWT-DAG: %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-RWT-DAG: %[[C2:.*]] = arith.constant 2 : index
 //       CHECK-RWT: %[[D0:.*]] = tensor.dim %[[COO]], %[[C0]]
 //       CHECK-RWT: %[[D1:.*]] = tensor.dim %[[COO]], %[[C1]]
 //       CHECK-RWT: %[[D2:.*]] = tensor.dim %[[COO]], %[[C2]]
-//       CHECK-RWT: %[[I0:.*]] = sparse_tensor.indices %[[COO]] {dimension = 0 : index}
-//       CHECK-RWT: %[[I1:.*]] = sparse_tensor.indices %[[COO]] {dimension = 1 : index}
-//       CHECK-RWT: %[[I2:.*]] = sparse_tensor.indices %[[COO]] {dimension = 2 : index}
-//       CHECK-RWT: %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[COO]]
-//       CHECK-RWT: %[[V:.*]] = sparse_tensor.values %[[COO]]
-//       CHECK-RWT: sparse_tensor.sort %[[NNZ]], %[[I2]], %[[I0]], %[[I1]] jointly %[[V]]
 //       CHECK-RWT: %[[T1:.*]] = bufferization.alloc_tensor(%[[D0]], %[[D1]], %[[D2]])
 //       CHECK-RWT: %[[T2:.*]] = sparse_tensor.foreach in %[[COO]] init(%[[T1]])
 //       CHECK-RWT: ^bb0(%[[LI0:.*]]: index, %[[LI1:.*]]: index, %[[LI2:.*]]: index, %[[LV:.*]]: f32, %[[LT1:.*]]: tensor<?x?x?xf32,
 //       CHECK-RWT:   %[[LT2:.*]] = sparse_tensor.insert %[[LV]] into %[[LT1]]{{\[}}%[[LI2]], %[[LI0]], %[[LI1]]]
 //       CHECK-RWT:   sparse_tensor.yield %[[LT2]]
 //       CHECK-RWT: }
 //       CHECK-RWT: %[[T3:.*]] = sparse_tensor.load %[[T2:.*]] hasInserts
 //       CHECK-RWT: %[[T4:.*]] = sparse_tensor.convert %[[T3]]
 //       CHECK-RWT: return %[[T4]]
 func.func @sparse_convert_permuted(%arg0: tensor<?x?x?xf32, #SortedCOO3D>) -> tensor<?x?x?xf32, #TsssPermuted> {
   %0 = sparse_tensor.convert %arg0 : tensor<?x?x?xf32, #SortedCOO3D> to tensor<?x?x?xf32, #TsssPermuted>
   return %0 : tensor<?x?x?xf32, #TsssPermuted>
 }