[Mlir-commits] [mlir] 99fe3d2 - [mlir][sparse] 3-dimensional sparse tensor insertion test
Aart Bik
llvmlistbot at llvm.org
Tue Nov 8 14:49:39 PST 2022
Author: Aart Bik
Date: 2022-11-08T14:49:31-08:00
New Revision: 99fe3d2661085a14fa9f1f2861ab840704382e28
URL: https://github.com/llvm/llvm-project/commit/99fe3d2661085a14fa9f1f2861ab840704382e28
DIFF: https://github.com/llvm/llvm-project/commit/99fe3d2661085a14fa9f1f2861ab840704382e28.diff
LOG: [mlir][sparse] 3-dimensional sparse tensor insertion test
Reviewed By: Peiming
Differential Revision: https://reviews.llvm.org/D137668
Added:
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_3d.mlir
Modified:
Removed:
################################################################################
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_3d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_3d.mlir
new file mode 100644
index 000000000000..d05b3f901d25
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_3d.mlir
@@ -0,0 +1,122 @@
+// RUN: mlir-opt %s --sparse-compiler=enable-runtime-library=false | \
+// 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
+
+#TensorCSR = #sparse_tensor.encoding<{
+ dimLevelType = [ "compressed", "dense", "compressed" ]
+}>
+
+#TensorRow = #sparse_tensor.encoding<{
+ dimLevelType = [ "compressed", "compressed", "dense" ]
+}>
+
+module {
+
+ func.func @dump(%arg0: tensor<5x4x3xf64, #TensorCSR>) {
+ %c0 = arith.constant 0 : index
+ %fu = arith.constant 99.0 : f64
+ %p0 = sparse_tensor.pointers %arg0 { dimension = 0 : index } : tensor<5x4x3xf64, #TensorCSR> to memref<?xindex>
+ %i0 = sparse_tensor.indices %arg0 { dimension = 0 : index } : tensor<5x4x3xf64, #TensorCSR> to memref<?xindex>
+ %p2 = sparse_tensor.pointers %arg0 { dimension = 2 : index } : tensor<5x4x3xf64, #TensorCSR> to memref<?xindex>
+ %i2 = sparse_tensor.indices %arg0 { dimension = 2 : index } : tensor<5x4x3xf64, #TensorCSR> to memref<?xindex>
+ %v = sparse_tensor.values %arg0 : tensor<5x4x3xf64, #TensorCSR> to memref<?xf64>
+ %vp0 = vector.transfer_read %p0[%c0], %c0: memref<?xindex>, vector<2xindex>
+ vector.print %vp0 : vector<2xindex>
+ %vi0 = vector.transfer_read %i0[%c0], %c0: memref<?xindex>, vector<2xindex>
+ vector.print %vi0 : vector<2xindex>
+ %vp2 = vector.transfer_read %p2[%c0], %c0: memref<?xindex>, vector<9xindex>
+ vector.print %vp2 : vector<9xindex>
+ %vi2 = vector.transfer_read %i2[%c0], %c0: memref<?xindex>, vector<5xindex>
+ vector.print %vi2 : vector<5xindex>
+ %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<5xf64>
+ vector.print %vv : vector<5xf64>
+ return
+ }
+
+ func.func @dump_row(%arg0: tensor<5x4x3xf64, #TensorRow>) {
+ %c0 = arith.constant 0 : index
+ %fu = arith.constant 99.0 : f64
+ %p0 = sparse_tensor.pointers %arg0 { dimension = 0 : index } : tensor<5x4x3xf64, #TensorRow> to memref<?xindex>
+ %i0 = sparse_tensor.indices %arg0 { dimension = 0 : index } : tensor<5x4x3xf64, #TensorRow> to memref<?xindex>
+ %p1 = sparse_tensor.pointers %arg0 { dimension = 1 : index } : tensor<5x4x3xf64, #TensorRow> to memref<?xindex>
+ %i1 = sparse_tensor.indices %arg0 { dimension = 1 : index } : tensor<5x4x3xf64, #TensorRow> to memref<?xindex>
+ %v = sparse_tensor.values %arg0 : tensor<5x4x3xf64, #TensorRow> to memref<?xf64>
+ %vp0 = vector.transfer_read %p0[%c0], %c0: memref<?xindex>, vector<2xindex>
+ vector.print %vp0 : vector<2xindex>
+ %vi0 = vector.transfer_read %i0[%c0], %c0: memref<?xindex>, vector<2xindex>
+ vector.print %vi0 : vector<2xindex>
+ %vp1 = vector.transfer_read %p1[%c0], %c0: memref<?xindex>, vector<3xindex>
+ vector.print %vp1 : vector<3xindex>
+ %vi1 = vector.transfer_read %i1[%c0], %c0: memref<?xindex>, vector<4xindex>
+ vector.print %vi1 : vector<4xindex>
+ %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<12xf64>
+ vector.print %vv : vector<12xf64>
+ return
+ }
+
+ //
+ // Main driver.
+ //
+ func.func @entry() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %c4 = arith.constant 4 : index
+ %f1 = arith.constant 1.1 : f64
+ %f2 = arith.constant 2.2 : f64
+ %f3 = arith.constant 3.3 : f64
+ %f4 = arith.constant 4.4 : f64
+ %f5 = arith.constant 5.5 : f64
+
+ //
+ // CHECK: ( 0, 2 )
+ // CHECK-NEXT: ( 3, 4 )
+ // CHECK-NEXT: ( 0, 2, 2, 2, 3, 3, 3, 4, 5 )
+ // CHECK-NEXT: ( 1, 2, 1, 2, 2 )
+ // CHECK-NEXT: ( 1.1, 2.2, 3.3, 4.4, 5.5 )
+ //
+ %tensora = bufferization.alloc_tensor() : tensor<5x4x3xf64, #TensorCSR>
+ %tensor1 = sparse_tensor.insert %f1 into %tensora[%c3, %c0, %c1] : tensor<5x4x3xf64, #TensorCSR>
+ %tensor2 = sparse_tensor.insert %f2 into %tensor1[%c3, %c0, %c2] : tensor<5x4x3xf64, #TensorCSR>
+ %tensor3 = sparse_tensor.insert %f3 into %tensor2[%c3, %c3, %c1] : tensor<5x4x3xf64, #TensorCSR>
+ %tensor4 = sparse_tensor.insert %f4 into %tensor3[%c4, %c2, %c2] : tensor<5x4x3xf64, #TensorCSR>
+ %tensor5 = sparse_tensor.insert %f5 into %tensor4[%c4, %c3, %c2] : tensor<5x4x3xf64, #TensorCSR>
+ %tensorm = sparse_tensor.load %tensor5 hasInserts : tensor<5x4x3xf64, #TensorCSR>
+ call @dump(%tensorm) : (tensor<5x4x3xf64, #TensorCSR>) -> ()
+
+ //
+ // CHECK-NEXT: ( 0, 2 )
+ // CHECK-NEXT: ( 3, 4 )
+ // CHECK-NEXT: ( 0, 2, 4 )
+ // CHECK-NEXT: ( 0, 3, 2, 3 )
+ // CHECK-NEXT: ( 0, 1.1, 2.2, 0, 3.3, 0, 0, 0, 4.4, 0, 0, 5.5 )
+ //
+ %rowa = bufferization.alloc_tensor() : tensor<5x4x3xf64, #TensorRow>
+ %row1 = sparse_tensor.insert %f1 into %rowa[%c3, %c0, %c1] : tensor<5x4x3xf64, #TensorRow>
+ %row2 = sparse_tensor.insert %f2 into %row1[%c3, %c0, %c2] : tensor<5x4x3xf64, #TensorRow>
+ %row3 = sparse_tensor.insert %f3 into %row2[%c3, %c3, %c1] : tensor<5x4x3xf64, #TensorRow>
+ %row4 = sparse_tensor.insert %f4 into %row3[%c4, %c2, %c2] : tensor<5x4x3xf64, #TensorRow>
+ %row5 = sparse_tensor.insert %f5 into %row4[%c4, %c3, %c2] : tensor<5x4x3xf64, #TensorRow>
+ %rowm = sparse_tensor.load %row5 hasInserts : tensor<5x4x3xf64, #TensorRow>
+ call @dump_row(%rowm) : (tensor<5x4x3xf64, #TensorRow>) -> ()
+
+ // NOE sanity check.
+ //
+ // CHECK-NEXT: 5
+ // CHECK-NEXT: 12
+ //
+ %noe1 = sparse_tensor.number_of_entries %tensorm : tensor<5x4x3xf64, #TensorCSR>
+ vector.print %noe1 : index
+ %noe2 = sparse_tensor.number_of_entries %rowm : tensor<5x4x3xf64, #TensorRow>
+ vector.print %noe2 : index
+
+ // Release resources.
+ bufferization.dealloc_tensor %tensorm : tensor<5x4x3xf64, #TensorCSR>
+ bufferization.dealloc_tensor %rowm : tensor<5x4x3xf64, #TensorRow>
+
+ return
+ }
+}
More information about the Mlir-commits
mailing list