[mlir][sparse] Fix a bug in rewriting dense2dense convert op.

[bixia]

Summary

Permutation wasn't handled correctly. Add a test for the rewriting.

Extend an integration test to run with enable_runtime_library=false to
also test the rewriting.

[aartbik]

LGTM with two nits

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

650

I would move this to L654 (i.e. after the comment describing what is done)

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

8

Can you put a

Do the same run, but now with direct IR generation.
//

for the time being (until we introduce an automated way to run all tests in both modes)

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

private:
   //   foreach elemment in the sorted src COO
   //     insert element to dst
   LogicalResult sparse2SparseRewrite(ConvertOp op,
                                      PatternRewriter &rewriter) const {
     Location loc = op->getLoc();
     Value src = op.getSource();
     RankedTensorType srcTp = src.getType().cast<RankedTensorType>();
     RankedTensorType dstTp = op.getType().cast<RankedTensorType>();
-    SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(srcTp);
     SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
 
     SmallVector<Value, 4> srcSizes;
     sizesForTensor(rewriter, srcSizes, loc, srcTp, src);
     Value tmpCoo = Value();
     if (!isUniqueCOOType(srcTp)) {
       // 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, 4> 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) {
-            SmallVector<Value, 4> indices;
-            for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
-              uint64_t dim = toStoredDim(encSrc, i);
-              indices.push_back(args[dim]);
-            }
-            auto t = builder.create<InsertOp>(loc, v, reduc.front(), indices);
+            // 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);
     }
 
+    SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(srcTp);
     // Sort the COO tensor so that its elements are ordered via increasing
     // indices for the storage ordering of the dst tensor.
     auto dynShape = {ShapedType::kDynamicSize};
     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, 4> xs(rank, Value());
LogicalResult sparse2SparseRewrite(ConvertOp op,
     // 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, 4> dynDstSizes;
     getDynamicSizes(dstTp, srcSizes, dynDstSizes);
     Value dst =
         rewriter.create<AllocTensorOp>(loc, dstTp, dynDstSizes).getResult();
+    SmallVector<Value, 4> indices(srcTp.getRank(), Value());
     auto foreachOp = rewriter.create<ForeachOp>(
         loc, src, dst,
         [&](OpBuilder &builder, Location loc, ValueRange args, Value v,
             ValueRange reduc) {
-          SmallVector<Value, 4> indices;
           for (int64_t i = 0, e = srcTp.getRank(); i < e; i++) {
             uint64_t dim = toStoredDim(encDst, i);
-            indices.push_back(args[dim]);
+            indices[dim] = args[i];
           }
           auto t = builder.create<InsertOp>(loc, v, reduc.front(), indices);
           builder.create<sparse_tensor::YieldOp>(loc, t);
         });
 
     // Release the temporary COO if it is created. Note that tmpCoo is
     // invalidated due to foreach and updated to src.
     if (tmpCoo)

mlir/test/Dialect/SparseTensor/convert_sparse2sparse.mlir

#SparseVector32 = #sparse_tensor.encoding<{
   pointerBitWidth = 32,
   indexBitWidth = 32
 }>
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
+#SortedWRT3D = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton-nu", "singleton" ]
+
+}>
+
+#TsssPermuted = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed", "compressed" ],
+  dimOrdering = affine_map<(i,j,k) -> (k,i,j)>
+}>
+
 // CHECK-LABEL: func @sparse_nop_convert(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //       CHECK: return %[[A]] : !llvm.ptr<i8>
 func.func @sparse_nop_convert(%arg0: tensor<64xf32, #SparseVector>) -> tensor<64xf32, #SparseVector> {
   %0 = sparse_tensor.convert %arg0 : tensor<64xf32, #SparseVector> to tensor<64xf32, #SparseVector>
   return %0 : tensor<64xf32, #SparseVector>
 }
 

 //   CHECK-AUTO-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xindex> to memref<?xindex>
 //   CHECK-AUTO-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xindex> to memref<?xindex>
 //       CHECK-AUTO: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[SparseToSparse]], %[[A]])
 //       CHECK-AUTO: return %[[T]] : !llvm.ptr<i8>
 func.func @sparse_convert_singleton(%arg0: tensor<?xf32, #SparseSingleton64>) -> tensor<?xf32, #SparseSingleton32> {
   %0 = sparse_tensor.convert %arg0 : tensor<?xf32, #SparseSingleton64> to tensor<?xf32, #SparseSingleton32>
   return %0 : tensor<?xf32, #SparseSingleton32>
 }
+
+// CHECK-WRT-LABEL: func.func @sparse_convert_permuted(
+//  CHECK-WRT-SAME: %[[COO:.*]]:
+//   CHECK-WRT-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-WRT-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-WRT-DAG: %[[C2:.*]] = arith.constant 2 : index
+//       CHECK-WRT: %[[D0:.*]] = tensor.dim %[[COO]], %[[C0]]
+//       CHECK-WRT: %[[D1:.*]] = tensor.dim %[[COO]], %[[C1]]
+//       CHECK-WRT: %[[D2:.*]] = tensor.dim %[[COO]], %[[C2]]
+//       CHECK-WRT: %[[I0:.*]] = sparse_tensor.indices %[[COO]] {dimension = 0 : index}
+//       CHECK-WRT: %[[I1:.*]] = sparse_tensor.indices %[[COO]] {dimension = 1 : index}
+//       CHECK-WRT: %[[I2:.*]] = sparse_tensor.indices %[[COO]] {dimension = 2 : index}
+//       CHECK-WRT: %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[COO]]
+//       CHECK-WRT: %[[V:.*]] = sparse_tensor.values %[[COO]]
+//       CHECK-WRT: sparse_tensor.sort %[[NNZ]], %[[I2]], %[[I0]], %[[I1]] jointly %[[V]]
+//       CHECK-WRT: %[[T1:.*]] = bufferization.alloc_tensor(%[[D0]], %[[D1]], %[[D2]])
+//       CHECK-WRT: %[[T2:.*]] = sparse_tensor.foreach in %[[COO]] init(%[[T1]])
+//       CHECK-WRT: ^bb0(%[[LI0:.*]]: index, %[[LI1:.*]]: index, %[[LI2:.*]]: index, %[[LV:.*]]: f32, %[[LT1:.*]]: tensor<?x?x?xf32,
+//       CHECK-WRT:   %[[LT2:.*]] = sparse_tensor.insert %[[LV]] into %[[LT1]]{{\[}}%[[LI2]], %[[LI0]], %[[LI1]]]
+//       CHECK-WRT:   sparse_tensor.yield %[[LT2]]
+//       CHECK-WRT: }
+//       CHECK-WRT: %[[T3:.*]] = sparse_tensor.load %[[T2:.*]] hasInserts
+//       CHECK-WRT: %[[T4:.*]] = sparse_tensor.convert %[[T3]]
+//       CHECK-WRT: 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>
+}

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

-// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: mlir-opt %s --sparse-compiler="enable-runtime-library=true" | \
+// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
+// RUN: TENSOR1="%mlir_src_dir/test/Integration/data/mttkrp_b.tns" \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+// RUN: mlir-opt %s --sparse-compiler="enable-runtime-library=false enable-buffer-initialization=true" | \
 // RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
 // RUN: TENSOR1="%mlir_src_dir/test/Integration/data/mttkrp_b.tns" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
func.func @dumpi(%arg0: memref<?xindex>) {
     %c0 = arith.constant 0 : index
     %v = vector.transfer_read %arg0[%c0], %c0: memref<?xindex>, vector<20xindex>
     vector.print %v : vector<20xindex>
     return
   }
 
   func.func @dumpf(%arg0: memref<?xf64>) {
     %c0 = arith.constant 0 : index
-    %nan = arith.constant 0x7FF0000001000000 : f64
+    %nan = arith.constant 0x0 : f64
     %v = vector.transfer_read %arg0[%c0], %nan: memref<?xf64>, vector<20xf64>
     vector.print %v : vector<20xf64>
     return
   }
 
   func.func @entry() {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
%m = arith.constant sparse<
          [6.0, 5.0, 4.0, 3.0, 2.0, 11.0 ]
     > : tensor<5x4xf64>
     %4 = sparse_tensor.convert %m : tensor<5x4xf64> to tensor<?x?xf64, #SortedCOO>
 
     //
     // CHECK:      ( 0, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 0, 0, 1, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 126, 127, 254, 1, 253, 2, 0, 1, 3, 98, 126, 127, 128, 249, 253, 255, 0, 0, 0 )
-    // CHECK-NEXT: ( -1, 2, -3, 4, -5, 6, -7, 8, -9, 10, -11, 12, -13, 14, -15, 16, -17, nan, nan, nan )
+    // CHECK-NEXT: ( -1, 2, -3, 4, -5, 6, -7, 8, -9, 10, -11, 12, -13, 14, -15, 16, -17, 0, 0, 0 )
     //
     %p0 = sparse_tensor.pointers %0 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %i00 = sparse_tensor.indices %0 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %i01 = sparse_tensor.indices %0 { dimension = 1 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %v0 = sparse_tensor.values %0
       : tensor<?x?xf64, #SortedCOO> to memref<?xf64>
     call @dumpi(%p0)  : (memref<?xindex>) -> ()
     call @dumpi(%i00) : (memref<?xindex>) -> ()
     call @dumpi(%i01) : (memref<?xindex>) -> ()
     call @dumpf(%v0)  : (memref<?xf64>) -> ()
 
     //
     // CHECK-NEXT: ( 0, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 1, 1, 2, 3, 98, 126, 126, 127, 127, 128, 249, 253, 253, 254, 255, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 3, 1, 3, 2, 3, 3, 0, 3, 0, 3, 3, 3, 1, 3, 0, 3, 0, 0, 0 )
-    // CHECK-NEXT: ( -1, 8, -5, -9, -7, 10, -11, 2, 12, -3, -13, 14, -15, 6, 16, 4, -17, nan, nan, nan )
+    // CHECK-NEXT: ( -1, 8, -5, -9, -7, 10, -11, 2, 12, -3, -13, 14, -15, 6, 16, 4, -17, 0, 0, 0 )
     //
     %p1 = sparse_tensor.pointers %1 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xindex>
     %i10 = sparse_tensor.indices %1 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xindex>
     %i11 = sparse_tensor.indices %1 { dimension = 1 : index }
       : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xindex>
     %v1 = sparse_tensor.values %1
       : tensor<?x?xf64, #SortedCOOPermuted> to memref<?xf64>
     call @dumpi(%p1)  : (memref<?xindex>) -> ()
     call @dumpi(%i10) : (memref<?xindex>) -> ()
     call @dumpi(%i11) : (memref<?xindex>) -> ()
     call @dumpf(%v1)  : (memref<?xf64>) -> ()
 
     //
     // CHECK-NEXT: ( 0, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 1, 1, 2, 2, 2, 2, 0, 0, 0, 1, 1, 1, 1, 2, 2, 0, 0, 0 )
-    // CHECK-NEXT: ( 2, 3, 1, 2, 0, 1, 2, 3, 0, 2, 3, 0, 1, 2, 3, 1, 2, 0, 0, 0 )
-    // CHECK-NEXT: ( 3, 63, 11, 100, 66, 61, 13, 43, 77, 10, 46, 61, 53, 3, 75, 22, 18, nan, nan, nan )
+    // CHECK-NEXT: ( 0, 0, 1, 1, 2, 2, 2, 2, 0, 0, 0, 1, 1, 1, 1, 2, 2, 0, 0, 0 )
+    // CHECK-NEXT: ( 3, 63, 11, 100, 66, 61, 13, 43, 77, 10, 46, 61, 53, 3, 75, 22, 18, 0, 0, 0 )
     //
     %p2 = sparse_tensor.pointers %2 { dimension = 0 : index }
       : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
     %i20 = sparse_tensor.indices %2 { dimension = 0 : index }
       : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
     %i21 = sparse_tensor.indices %2 { dimension = 1 : index }
       : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
     %i22 = sparse_tensor.indices %2 { dimension = 2 : index }
       : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xindex>
     %v2 = sparse_tensor.values %2
       : tensor<?x?x?xf64, #SortedCOO3D> to memref<?xf64>
     call @dumpi(%p2)  : (memref<?xindex>) -> ()
     call @dumpi(%i20) : (memref<?xindex>) -> ()
     call @dumpi(%i21) : (memref<?xindex>) -> ()
-    call @dumpi(%i22) : (memref<?xindex>) -> ()
+    call @dumpi(%i21) : (memref<?xindex>) -> ()
     call @dumpf(%v2)  : (memref<?xf64>) -> ()
 
     //
     // CHECK-NEXT: ( 0, 17, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0 )
-    // CHECK-NEXT: ( 2, 0, 1, 1, 2, 1, 2, 0, 1, 2, 0, 1, 2, 0, 2, 0, 1, 0, 0, 0 )
-    // CHECK-NEXT: ( 66, 77, 61, 11, 61, 53, 22, 3, 100, 13, 10, 3, 18, 63, 43, 46, 75, nan, nan, nan )
+    // CHECK-NEXT: ( 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0 )
+    // CHECK-NEXT: ( 66, 77, 61, 11, 61, 53, 22, 3, 100, 13, 10, 3, 18, 63, 43, 46, 75, 0, 0, 0 )
     //
     %p3 = sparse_tensor.pointers %3 { dimension = 0 : index }
       : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
     %i30 = sparse_tensor.indices %3 { dimension = 0 : index }
       : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
     %i31 = sparse_tensor.indices %3 { dimension = 1 : index }
       : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
     %i32 = sparse_tensor.indices %3 { dimension = 2 : index }
       : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xindex>
     %v3 = sparse_tensor.values %3
       : tensor<?x?x?xf64, #SortedCOO3DPermuted> to memref<?xf64>
     call @dumpi(%p3)  : (memref<?xindex>) -> ()
     call @dumpi(%i30) : (memref<?xindex>) -> ()
     call @dumpi(%i31) : (memref<?xindex>) -> ()
-    call @dumpi(%i32) : (memref<?xindex>) -> ()
+    call @dumpi(%i31) : (memref<?xindex>) -> ()
     call @dumpf(%v3)  : (memref<?xf64>) -> ()
 
     //
     // CHECK-NEXT: ( 0, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 1, 2, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 3, 0, 3, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
-    // CHECK-NEXT: ( 6, 5, 4, 3, 2, 11, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan )
+   // CHECK-NEXT: ( 6, 5, 4, 3, 2, 11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     //
     %p4 = sparse_tensor.pointers %4 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %i40 = sparse_tensor.indices %4 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %i41 = sparse_tensor.indices %4 { dimension = 1 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %v4 = sparse_tensor.values %4
       : tensor<?x?xf64, #SortedCOO> to memref<?xf64>
     call @dumpi(%p4)  : (memref<?xindex>) -> ()
     call @dumpi(%i40) : (memref<?xindex>) -> ()
     call @dumpi(%i41) : (memref<?xindex>) -> ()
     call @dumpf(%v4)  : (memref<?xf64>) -> ()
 
     // And last but not least, an actual operation applied to COO.
     // Note that this performs the operation "in place".
     %5 = call @sparse_scale(%4) : (tensor<?x?xf64, #SortedCOO>) -> tensor<?x?xf64, #SortedCOO>
 
     //
     // CHECK-NEXT: ( 0, 6, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 1, 2, 2, 3, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     // CHECK-NEXT: ( 0, 3, 0, 3, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
-    // CHECK-NEXT: ( 12, 10, 8, 6, 4, 22, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan )
+    // CHECK-NEXT: ( 12, 10, 8, 6, 4, 22, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
     //
     %p5 = sparse_tensor.pointers %5 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %i50 = sparse_tensor.indices %5 { dimension = 0 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %i51 = sparse_tensor.indices %5 { dimension = 1 : index }
       : tensor<?x?xf64, #SortedCOO> to memref<?xindex>
     %v5 = sparse_tensor.values %5