[Mlir-commits] [mlir] 0a7b8cc - [mlir][sparse] fully implement sparse tensor to sparse tensor conversions
Aart Bik
llvmlistbot at llvm.org
Fri Aug 27 15:08:26 PDT 2021
Author: Aart Bik
Date: 2021-08-27T15:08:18-07:00
New Revision: 0a7b8cc5dd8efadeef07bf6e4e93c7df2ebe7220
URL: https://github.com/llvm/llvm-project/commit/0a7b8cc5dd8efadeef07bf6e4e93c7df2ebe7220
DIFF: https://github.com/llvm/llvm-project/commit/0a7b8cc5dd8efadeef07bf6e4e93c7df2ebe7220.diff
LOG: [mlir][sparse] fully implement sparse tensor to sparse tensor conversions
with rigorous integration test
Reviewed By: bixia
Differential Revision: https://reviews.llvm.org/D108721
Added:
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion.mlir
Modified:
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
mlir/test/Dialect/SparseTensor/conversion.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index 3708345244f4e..e243e8e300d6a 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -283,11 +283,23 @@ class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
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
// ..
diff --git a/mlir/test/Dialect/SparseTensor/conversion.mlir b/mlir/test/Dialect/SparseTensor/conversion.mlir
index 33ddfa67543a3..f3ea0c0584f17 100644
--- a/mlir/test/Dialect/SparseTensor/conversion.mlir
+++ b/mlir/test/Dialect/SparseTensor/conversion.mlir
@@ -136,6 +136,16 @@ func @sparse_convert_1d(%arg0: tensor<?xi32>) -> 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
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion.mlir
new file mode 100644
index 0000000000000..8626220cf8767
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_conversion.mlir
@@ -0,0 +1,251 @@
+// 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
+ }
+}
+
More information about the Mlir-commits
mailing list