[mlir][sparse] fully implement sparse tensor to sparse tensor conversions

[aartbik]

Summary

with rigorous integration test

[aartbik]

lint issue

[aartbik]

rebase with main

[bixia]

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

118–168

I wonder where there is a better way to assert %10 == %20, %11 == %21, so that we don't have to repeat the dump output in this block, which is the same as line 110-115.
For example, we can just print a boolean (%10 == %20) the FileCheck the boolean, if we can't think of a better way?

[aartbik]

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

118–168

Yeah, fair enough. After working with FileCheck for close to two years now, I have a bit of the "when you have a hammer everything looks like a nail" approach to checking with output. I will rewrite this into a form that uses more intuitive comparisons, so this can become a simple pass/fail check. Thanks for the suggestion!

[aartbik]

exhaustive but silent checks

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

 class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ConvertOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     auto encDst = getSparseTensorEncoding(resType);
     auto encSrc = getSparseTensorEncoding(op.source().getType());
-    // TODO: implement sparse => sparse
-    //             and sparse => dense
-    if (!encDst || encSrc)
+    if (encDst && encSrc) {
+      // This is a sparse => sparse conversion, which is handled as follows:
+      //   t = src->asCOO();         ; src to COO in dst order
+      //   dst = newSparseTensor(t)
+      // Using the coordinate scheme as an intermediate does not always
+      // yield the fastest conversion but avoids the need for a full
+      // O(N^2) conversion matrix.
+      Value perm;
+      Value coo = genNewCall(rewriter, op, encDst, 3, perm, operands[0]);
+      rewriter.replaceOp(op, genNewCall(rewriter, op, encDst, 1, perm, coo));
+      return success();
+    }
+    if (!encDst || encSrc) {
+      // TODO: sparse => dense
       return failure();
-    // This is a dense => sparse conversion, that is handled as follows:
+    }
+    // This is a dense => sparse conversion, which is handled as follows:
     //   t = newSparseCOO()
     //   for i1 in dim1
     //    ..
     //     for ik in dimk
     //       val = a[i1,..,ik]
     //       if val != 0
     //         t->add(val, [i1,..,ik], [p1,..,pk])
     //   s = newSparseTensor(t)

mlir/test/Dialect/SparseTensor/conversion.mlir

 //       CHECK: }
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Y]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func @sparse_convert_1d(%arg0: tensor<?xi32>) -> tensor<?xi32, #SparseVector> {
   %0 = sparse_tensor.convert %arg0 : tensor<?xi32> to tensor<?xi32, #SparseVector>
   return %0 : tensor<?xi32, #SparseVector>
 }
 
+// CHECK-LABEL: func @sparse_convert_1d_ss(
+//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
+//       CHECK: %[[C:.*]] = call @newSparseTensor(%{{.}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%{{.}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
+//       CHECK: return %[[T]] : !llvm.ptr<i8>
+func @sparse_convert_1d_ss(%arg0: tensor<?xf32, #SparseVector64>) -> tensor<?xf32, #SparseVector32> {
+  %0 = sparse_tensor.convert %arg0 : tensor<?xf32, #SparseVector64> to tensor<?xf32, #SparseVector32>
+  return %0 : tensor<?xf32, #SparseVector32>
+}
+
 // CHECK-LABEL: func @sparse_convert_2d(
 //  CHECK-SAME: %[[A:.*]]: tensor<2x4xf64>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[C0:.*]] = constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = constant 1 : index
 //   CHECK-DAG: %[[U:.*]] = constant dense<[0, 1]> : tensor<2xi8>
 //   CHECK-DAG: %[[V:.*]] = constant dense<[2, 4]> : tensor<2xi64>
 //   CHECK-DAG: %[[W:.*]] = constant dense<[0, 1]> : tensor<2xi64>
 //   CHECK-DAG: %[[X:.*]] = tensor.cast %[[U]] : tensor<2xi8> to tensor<?xi8>

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

+// RUN: mlir-opt %s \
+// RUN:   --sparsification --sparse-tensor-conversion \
+// RUN:   --convert-vector-to-scf --convert-scf-to-std \
+// RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
+// RUN:   --std-bufferize --finalizing-bufferize  \
+// RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+#Tensor1  = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed", "compressed" ],
+  dimOrdering = affine_map<(i,j,k) -> (i,j,k)>
+}>
+
+#Tensor2  = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed", "compressed" ],
+  dimOrdering = affine_map<(i,j,k) -> (j,k,i)>
+}>
+
+#Tensor3  = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed", "compressed" ],
+  dimOrdering = affine_map<(i,j,k) -> (k,i,j)>
+}>
+
+//
+// Integration test that tests conversions between sparse tensors.
+//
+module {
+  func private @exit(index) -> ()
+
+  //
+  // Verify utilities.
+  //
+  func @checkf64(%arg0: memref<?xf64>, %arg1: memref<?xf64>) {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    // Same lengths?
+    %0 = memref.dim %arg0, %c0 : memref<?xf64>
+    %1 = memref.dim %arg1, %c0 : memref<?xf64>
+    %2 = cmpi ne, %0, %1 : index
+    scf.if %2 {
+      call @exit(%c1) : (index) -> ()
+    }
+    // Same content?
+    scf.for %i = %c0 to %0 step %c1 {
+      %a = memref.load %arg0[%i] : memref<?xf64>
+      %b = memref.load %arg1[%i] : memref<?xf64>
+      %c = cmpf une, %a, %b : f64
+      scf.if %c {
+        call @exit(%c1) : (index) -> ()
+      }
+    }
+    return
+  }
+  func @check(%arg0: memref<?xindex>, %arg1: memref<?xindex>) {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    // Same lengths?
+    %0 = memref.dim %arg0, %c0 : memref<?xindex>
+    %1 = memref.dim %arg1, %c0 : memref<?xindex>
+    %2 = cmpi ne, %0, %1 : index
+    scf.if %2 {
+      call @exit(%c1) : (index) -> ()
+    }
+    // Same content?
+    scf.for %i = %c0 to %0 step %c1 {
+      %a = memref.load %arg0[%i] : memref<?xindex>
+      %b = memref.load %arg1[%i] : memref<?xindex>
+      %c = cmpi ne, %a, %b : index
+      scf.if %c {
+        call @exit(%c1) : (index) -> ()
+      }
+    }
+    return
+  }
+
+  //
+  // Output utility.
+  //
+  func @dumpf64(%arg0: memref<?xf64>) {
+    %c0 = constant 0 : index
+    %d0 = constant 0.0 : f64
+    %0 = vector.transfer_read %arg0[%c0], %d0: memref<?xf64>, vector<24xf64>
+    vector.print %0 : vector<24xf64>
+    return
+  }
+
+  //
+  // Main driver.
+  //
+  func @entry() {
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+
+    //
+    // Initialize a 3-dim dense tensor.
+    //
+    %t = constant dense<[
+       [  [  1.0,  2.0,  3.0,  4.0 ],
+          [  5.0,  6.0,  7.0,  8.0 ],
+          [  9.0, 10.0, 11.0, 12.0 ] ],
+       [  [ 13.0, 14.0, 15.0, 16.0 ],
+          [ 17.0, 18.0, 19.0, 20.0 ],
+          [ 21.0, 22.0, 23.0, 24.0 ] ]
+    ]> : tensor<2x3x4xf64>
+
+    //
+    // Convert dense tensor directly to various sparse tensors.
+    //    tensor1: stored as 2x3x4
+    //    tensor2: stored as 3x4x2
+    //    tensor3: stored as 4x2x3
+    //
+    %1 = sparse_tensor.convert %t : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor1>
+    %2 = sparse_tensor.convert %t : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor2>
+    %3 = sparse_tensor.convert %t : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor3>
+
+    //
+    // Convert sparse tensor to various sparse tensors. Note that the result
+    // should always correspond to the direct conversion, since the sparse
+    // tensor formats have the ability to restore into the original ordering.
+    //
+    %a = sparse_tensor.convert %1 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf64, #Tensor1>
+    %b = sparse_tensor.convert %2 : tensor<2x3x4xf64, #Tensor2> to tensor<2x3x4xf64, #Tensor1>
+    %c = sparse_tensor.convert %3 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf64, #Tensor1>
+    %d = sparse_tensor.convert %1 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf64, #Tensor2>
+    %e = sparse_tensor.convert %2 : tensor<2x3x4xf64, #Tensor2> to tensor<2x3x4xf64, #Tensor2>
+    %f = sparse_tensor.convert %3 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf64, #Tensor2>
+    %g = sparse_tensor.convert %1 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf64, #Tensor3>
+    %h = sparse_tensor.convert %2 : tensor<2x3x4xf64, #Tensor2> to tensor<2x3x4xf64, #Tensor3>
+    %i = sparse_tensor.convert %3 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf64, #Tensor3>
+
+    //
+    // Check values equality.
+    //
+
+    %v1 = sparse_tensor.values %1 : tensor<2x3x4xf64, #Tensor1> to memref<?xf64>
+    %v2 = sparse_tensor.values %2 : tensor<2x3x4xf64, #Tensor2> to memref<?xf64>
+    %v3 = sparse_tensor.values %3 : tensor<2x3x4xf64, #Tensor3> to memref<?xf64>
+
+    %av = sparse_tensor.values %a : tensor<2x3x4xf64, #Tensor1> to memref<?xf64>
+    %bv = sparse_tensor.values %b : tensor<2x3x4xf64, #Tensor1> to memref<?xf64>
+    %cv = sparse_tensor.values %c : tensor<2x3x4xf64, #Tensor1> to memref<?xf64>
+    %dv = sparse_tensor.values %d : tensor<2x3x4xf64, #Tensor2> to memref<?xf64>
+    %ev = sparse_tensor.values %e : tensor<2x3x4xf64, #Tensor2> to memref<?xf64>
+    %fv = sparse_tensor.values %f : tensor<2x3x4xf64, #Tensor2> to memref<?xf64>
+    %gv = sparse_tensor.values %g : tensor<2x3x4xf64, #Tensor3> to memref<?xf64>
+    %hv = sparse_tensor.values %h : tensor<2x3x4xf64, #Tensor3> to memref<?xf64>
+    %iv = sparse_tensor.values %i : tensor<2x3x4xf64, #Tensor3> to memref<?xf64>
+
+    call @checkf64(%v1, %av) : (memref<?xf64>, memref<?xf64>) -> ()
+    call @checkf64(%v1, %bv) : (memref<?xf64>, memref<?xf64>) -> ()
+    call @checkf64(%v1, %cv) : (memref<?xf64>, memref<?xf64>) -> ()
+    call @checkf64(%v2, %dv) : (memref<?xf64>, memref<?xf64>) -> ()
+    call @checkf64(%v2, %ev) : (memref<?xf64>, memref<?xf64>) -> ()
+    call @checkf64(%v2, %fv) : (memref<?xf64>, memref<?xf64>) -> ()
+    call @checkf64(%v3, %gv) : (memref<?xf64>, memref<?xf64>) -> ()
+    call @checkf64(%v3, %hv) : (memref<?xf64>, memref<?xf64>) -> ()
+    call @checkf64(%v3, %iv) : (memref<?xf64>, memref<?xf64>) -> ()
+
+    //
+    // Check index equality.
+    //
+
+    %v10 = sparse_tensor.indices %1, %c0 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %v11 = sparse_tensor.indices %1, %c1 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %v12 = sparse_tensor.indices %1, %c2 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %v20 = sparse_tensor.indices %2, %c0 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %v21 = sparse_tensor.indices %2, %c1 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %v22 = sparse_tensor.indices %2, %c2 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %v30 = sparse_tensor.indices %3, %c0 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %v31 = sparse_tensor.indices %3, %c1 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %v32 = sparse_tensor.indices %3, %c2 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+
+    %a10 = sparse_tensor.indices %a, %c0 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %a11 = sparse_tensor.indices %a, %c1 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %a12 = sparse_tensor.indices %a, %c2 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %b10 = sparse_tensor.indices %b, %c0 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %b11 = sparse_tensor.indices %b, %c1 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %b12 = sparse_tensor.indices %b, %c2 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %c10 = sparse_tensor.indices %c, %c0 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %c11 = sparse_tensor.indices %c, %c1 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+    %c12 = sparse_tensor.indices %c, %c2 : tensor<2x3x4xf64, #Tensor1> to memref<?xindex>
+
+    %d10 = sparse_tensor.indices %d, %c0 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %d11 = sparse_tensor.indices %d, %c1 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %d12 = sparse_tensor.indices %d, %c2 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %e10 = sparse_tensor.indices %e, %c0 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %e11 = sparse_tensor.indices %e, %c1 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %e12 = sparse_tensor.indices %e, %c2 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %f10 = sparse_tensor.indices %f, %c0 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %f11 = sparse_tensor.indices %f, %c1 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+    %f12 = sparse_tensor.indices %f, %c2 : tensor<2x3x4xf64, #Tensor2> to memref<?xindex>
+
+    %g10 = sparse_tensor.indices %g, %c0 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %g11 = sparse_tensor.indices %g, %c1 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %g12 = sparse_tensor.indices %g, %c2 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %h10 = sparse_tensor.indices %h, %c0 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %h11 = sparse_tensor.indices %h, %c1 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %h12 = sparse_tensor.indices %h, %c2 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %i10 = sparse_tensor.indices %i, %c0 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %i11 = sparse_tensor.indices %i, %c1 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+    %i12 = sparse_tensor.indices %i, %c2 : tensor<2x3x4xf64, #Tensor3> to memref<?xindex>
+
+    call @check(%v10, %a10) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v11, %a11) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v12, %a12) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v10, %b10) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v11, %b11) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v12, %b12) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v10, %c10) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v11, %c11) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v12, %c12) : (memref<?xindex>, memref<?xindex>) -> ()
+
+    call @check(%v20, %d10) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v21, %d11) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v22, %d12) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v20, %e10) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v21, %e11) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v22, %e12) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v20, %f10) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v21, %f11) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v22, %f12) : (memref<?xindex>, memref<?xindex>) -> ()
+
+    call @check(%v30, %g10) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v31, %g11) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v32, %g12) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v30, %h10) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v31, %h11) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v32, %h12) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v30, %i10) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v31, %i11) : (memref<?xindex>, memref<?xindex>) -> ()
+    call @check(%v32, %i12) : (memref<?xindex>, memref<?xindex>) -> ()
+
+    //
+    // Sanity check direct results.
+    //
+    // CHECK:      ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24 )
+    // CHECK-NEXT: ( 1, 13, 2, 14, 3, 15, 4, 16, 5, 17, 6, 18, 7, 19, 8, 20, 9, 21, 10, 22, 11, 23, 12, 24 )
+    // CHECK-NEXT: ( 1, 5, 9, 13, 17, 21, 2, 6, 10, 14, 18, 22, 3, 7, 11, 15, 19, 23, 4, 8, 12, 16, 20, 24 )
+    //
+    call @dumpf64(%v1) : (memref<?xf64>) -> ()
+    call @dumpf64(%v2) : (memref<?xf64>) -> ()
+    call @dumpf64(%v3) : (memref<?xf64>) -> ()
+
+    return
+  }
+}
+