[Mlir-commits] [mlir] [mlir][sparse] migrate datastructure tests to sparse_tensor.print (PR #83956)

Aart Bik llvmlistbot at llvm.org
Mon Mar 4 19:58:42 PST 2024


https://github.com/aartbik created https://github.com/llvm/llvm-project/pull/83956

Continuing the efforts started in llvm#83357

>From 8dde35cec43b5727c4e62e27af098c286e2764f2 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 4 Mar 2024 19:46:06 -0800
Subject: [PATCH] [mlir][sparse] migrate datastructure tests to
 sparse_tensor.print

Continuing the efforts started in llvm#83357
---
 .../SparseTensor/CPU/sparse_insert_1d.mlir    |  62 +++---
 .../SparseTensor/CPU/sparse_insert_2d.mlir    | 176 +++++-------------
 2 files changed, 72 insertions(+), 166 deletions(-)

diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir
index bc7ecb08ab2f49..61c68507ea5198 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir
@@ -10,7 +10,7 @@
 // DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
 // DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
 // DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
-// DEFINE: %{run_opts} = -e entry -entry-point-result=void
+// DEFINE: %{run_opts} = -e main -entry-point-result=void
 // DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
 // DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
 //
@@ -30,8 +30,6 @@
 // Do the same run, but now with direct IR generation and VLA vectorization.
 // RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %}
 
-// Insertion example using pure codegen (no sparse runtime support lib).
-
 #SparseVector = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
 
 #trait_mul_s = {
@@ -43,27 +41,7 @@
 }
 
 module {
-
-  // Dumps positions, indices, values for verification.
-  func.func @dump(%argx: tensor<1024xf32, #SparseVector>) {
-    %c0 = arith.constant 0 : index
-    %f0 = arith.constant 0.0 : f32
-    %p = sparse_tensor.positions %argx { level = 0 : index }
-       : tensor<1024xf32, #SparseVector> to memref<?xindex>
-    %i = sparse_tensor.coordinates %argx { level = 0 : index }
-       : tensor<1024xf32, #SparseVector> to memref<?xindex>
-    %v = sparse_tensor.values %argx
-       : tensor<1024xf32, #SparseVector> to memref<?xf32>
-    %vp = vector.transfer_read %p[%c0], %c0: memref<?xindex>, vector<2xindex>
-    %vi = vector.transfer_read %i[%c0], %c0: memref<?xindex>, vector<8xindex>
-    %vv = vector.transfer_read %v[%c0], %f0: memref<?xf32>,   vector<8xf32>
-    vector.print %vp : vector<2xindex>
-    vector.print %vi : vector<8xindex>
-    vector.print %vv : vector<8xf32>
-    return
-  }
-
-  func.func @entry() {
+  func.func @main() {
     %f1    = arith.constant 1.0 : f32
     %f2    = arith.constant 2.0 : f32
     %f3    = arith.constant 3.0 : f32
@@ -82,10 +60,17 @@ module {
     %4 = sparse_tensor.insert %f4 into %3[%c1023] : tensor<1024xf32, #SparseVector>
     %5 = sparse_tensor.load %4 hasInserts : tensor<1024xf32, #SparseVector>
 
-    // CHECK:      ( 0, 4 )
-    // CHECK-NEXT: ( 0, 1, 3, 1023
-    // CHECK-NEXT: ( 1, 2, 3, 4
-    call @dump(%5) : (tensor<1024xf32, #SparseVector>) -> ()
+    //
+    // CHECK:   ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 4
+    // CHECK-NEXT: dim = ( 1024 )
+    // CHECK-NEXT: lvl = ( 1024 )
+    // CHECK-NEXT: pos[0] : ( 0, 4,
+    // CHECK-NEXT: crd[0] : ( 0, 1, 3, 1023,
+    // CHECK-NEXT: values : ( 1, 2, 3, 4,
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %5 : tensor<1024xf32, #SparseVector>
 
     // Build another sparse vector in a loop.
     %6 = tensor.empty() : tensor<1024xf32, #SparseVector>
@@ -96,18 +81,17 @@ module {
     }
     %8 = sparse_tensor.load %7 hasInserts : tensor<1024xf32, #SparseVector>
 
-    // CHECK-NEXT: ( 0, 8 )
-    // CHECK-NEXT: ( 0, 3, 6, 9, 12, 15, 18, 21 )
-    // CHECK-NEXT: ( 1, 1, 1, 1, 1, 1, 1, 1 )
     //
-    call @dump(%8) : (tensor<1024xf32, #SparseVector>) -> ()
-
-    // CHECK-NEXT: 4
-    // CHECK-NEXT: 8
-    %noe1 = sparse_tensor.number_of_entries %5 : tensor<1024xf32, #SparseVector>
-    %noe2 = sparse_tensor.number_of_entries %8 : tensor<1024xf32, #SparseVector>
-    vector.print %noe1 : index
-    vector.print %noe2 : index
+    // CHECK-NEXT: ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 8
+    // CHECK-NEXT: dim = ( 1024 )
+    // CHECK-NEXT: lvl = ( 1024 )
+    // CHECK-NEXT: pos[0] : ( 0, 8,
+    // CHECK-NEXT: crd[0] : ( 0, 3, 6, 9, 12, 15, 18, 21,
+    // CHECK-NEXT: values : ( 1, 1, 1, 1, 1, 1, 1, 1,
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %8 : tensor<1024xf32, #SparseVector>
 
     // Free resources.
     bufferization.dealloc_tensor %5 : tensor<1024xf32, #SparseVector>
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_2d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_2d.mlir
index b8cc1997783aa5..d51b67792337d1 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_2d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_2d.mlir
@@ -10,7 +10,7 @@
 // DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
 // DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
 // DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
-// DEFINE: %{run_opts} = -e entry -entry-point-result=void
+// DEFINE: %{run_opts} = -e main -entry-point-result=void
 // DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
 // DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
 //
@@ -32,7 +32,7 @@
 }>
 
 #SortedCOO = #sparse_tensor.encoding<{
-  map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)
+  map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton(soa))
 }>
 
 #CSR = #sparse_tensor.encoding<{
@@ -48,94 +48,11 @@
 }>
 
 module {
-
-  func.func @dump_dense(%arg0: tensor<4x3xf64, #Dense>) {
-    %c0 = arith.constant 0 : index
-    %fu = arith.constant 99.0 : f64
-    %v = sparse_tensor.values %arg0 : tensor<4x3xf64, #Dense> to memref<?xf64>
-    %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<12xf64>
-    vector.print %vv : vector<12xf64>
-    return
-  }
-
-  func.func @dump_coo(%arg0: tensor<4x3xf64, #SortedCOO>) {
-    %c0 = arith.constant 0 : index
-    %cu = arith.constant -1 : index
-    %fu = arith.constant 99.0 : f64
-    %p0 = sparse_tensor.positions %arg0 { level = 0 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex>
-    %i0 = sparse_tensor.coordinates  %arg0 { level = 0 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
-    %i1 = sparse_tensor.coordinates  %arg0 { level = 1 : index } : tensor<4x3xf64, #SortedCOO> to memref<?xindex, strided<[?], offset: ?>>
-    %v = sparse_tensor.values %arg0 : tensor<4x3xf64, #SortedCOO> to memref<?xf64>
-    %vp0 = vector.transfer_read %p0[%c0], %cu: memref<?xindex>, vector<2xindex>
-    vector.print %vp0 : vector<2xindex>
-    %vi0 = vector.transfer_read %i0[%c0], %cu: memref<?xindex, strided<[?], offset: ?>>, vector<4xindex>
-    vector.print %vi0 : vector<4xindex>
-    %vi1 = vector.transfer_read %i1[%c0], %cu: memref<?xindex, strided<[?], offset: ?>>, vector<4xindex>
-    vector.print %vi1 : vector<4xindex>
-    %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<4xf64>
-    vector.print %vv : vector<4xf64>
-    return
-  }
-
-  func.func @dump_csr(%arg0: tensor<4x3xf64, #CSR>) {
-    %c0 = arith.constant 0 : index
-    %cu = arith.constant -1 : index
-    %fu = arith.constant 99.0 : f64
-    %p1 = sparse_tensor.positions %arg0 { level = 1 : index } : tensor<4x3xf64, #CSR> to memref<?xindex>
-    %i1 = sparse_tensor.coordinates  %arg0 { level = 1 : index } : tensor<4x3xf64, #CSR> to memref<?xindex>
-    %v = sparse_tensor.values %arg0 : tensor<4x3xf64, #CSR> to memref<?xf64>
-    %vp1 = vector.transfer_read %p1[%c0], %cu: memref<?xindex>, vector<5xindex>
-    vector.print %vp1 : vector<5xindex>
-    %vi1 = vector.transfer_read %i1[%c0], %cu: memref<?xindex>, vector<4xindex>
-    vector.print %vi1 : vector<4xindex>
-    %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<4xf64>
-    vector.print %vv : vector<4xf64>
-    return
-  }
-
-  func.func @dump_dcsr(%arg0: tensor<4x3xf64, #DCSR>) {
-    %c0 = arith.constant 0 : index
-    %cu = arith.constant -1 : index
-    %fu = arith.constant 99.0 : f64
-    %p0 = sparse_tensor.positions %arg0 { level = 0 : index } : tensor<4x3xf64, #DCSR> to memref<?xindex>
-    %i0 = sparse_tensor.coordinates  %arg0 { level = 0 : index } : tensor<4x3xf64, #DCSR> to memref<?xindex>
-    %p1 = sparse_tensor.positions %arg0 { level = 1 : index } : tensor<4x3xf64, #DCSR> to memref<?xindex>
-    %i1 = sparse_tensor.coordinates  %arg0 { level = 1 : index } : tensor<4x3xf64, #DCSR> to memref<?xindex>
-    %v = sparse_tensor.values %arg0 : tensor<4x3xf64, #DCSR> to memref<?xf64>
-    %vp0 = vector.transfer_read %p0[%c0], %cu: memref<?xindex>, vector<2xindex>
-    vector.print %vp0 : vector<2xindex>
-    %vi0 = vector.transfer_read %i0[%c0], %cu: memref<?xindex>, vector<3xindex>
-    vector.print %vi0 : vector<3xindex>
-    %vp1 = vector.transfer_read %p1[%c0], %cu: memref<?xindex>, vector<4xindex>
-    vector.print %vp1 : vector<4xindex>
-    %vi1 = vector.transfer_read %i1[%c0], %cu: memref<?xindex>, vector<4xindex>
-    vector.print %vi1 : vector<4xindex>
-    %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<4xf64>
-    vector.print %vv : vector<4xf64>
-    return
-  }
-
-  func.func @dump_row(%arg0: tensor<4x3xf64, #Row>) {
-    %c0 = arith.constant 0 : index
-    %cu = arith.constant -1 : index
-    %fu = arith.constant 99.0 : f64
-    %p0 = sparse_tensor.positions %arg0 { level = 0 : index } : tensor<4x3xf64, #Row> to memref<?xindex>
-    %i0 = sparse_tensor.coordinates  %arg0 { level = 0 : index } : tensor<4x3xf64, #Row> to memref<?xindex>
-    %v = sparse_tensor.values %arg0 : tensor<4x3xf64, #Row> to memref<?xf64>
-    %vp0 = vector.transfer_read %p0[%c0], %cu: memref<?xindex>, vector<2xindex>
-    vector.print %vp0 : vector<2xindex>
-    %vi0 = vector.transfer_read %i0[%c0], %cu: memref<?xindex>, vector<3xindex>
-    vector.print %vi0 : vector<3xindex>
-    %vv = vector.transfer_read %v[%c0], %fu: memref<?xf64>, vector<9xf64>
-    vector.print %vv : vector<9xf64>
-    return
-  }
-
   //
   // Main driver. We test the contents of various sparse tensor
   // schemes when they are still empty and after a few insertions.
   //
-  func.func @entry() {
+  func.func @main() {
     %c0 = arith.constant 0 : index
     %c2 = arith.constant 2 : index
     %c3 = arith.constant 3 : index
@@ -147,7 +64,12 @@ module {
     //
     // Dense case.
     //
-    // CHECK: ( 1, 0, 0, 0, 0, 0, 0, 0, 2, 3, 0, 4 )
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: dim = ( 4, 3 )
+    // CHECK-NEXT: lvl = ( 4, 3 )
+    // CHECK-NEXT: values : ( 1, 0, 0, 0, 0, 0, 0, 0, 2, 3, 0, 4,
+    // CHECK-NEXT: ----
     //
     %densea = tensor.empty() : tensor<4x3xf64, #Dense>
     %dense1 = sparse_tensor.insert %f1 into %densea[%c0, %c0] : tensor<4x3xf64, #Dense>
@@ -155,15 +77,20 @@ module {
     %dense3 = sparse_tensor.insert %f3 into %dense2[%c3, %c0] : tensor<4x3xf64, #Dense>
     %dense4 = sparse_tensor.insert %f4 into %dense3[%c3, %c2] : tensor<4x3xf64, #Dense>
     %densem = sparse_tensor.load %dense4 hasInserts : tensor<4x3xf64, #Dense>
-    call @dump_dense(%densem) : (tensor<4x3xf64, #Dense>) -> ()
+    sparse_tensor.print %densem : tensor<4x3xf64, #Dense>
 
     //
     // COO case.
     //
-    // CHECK-NEXT: ( 0, 4 )
-    // CHECK-NEXT: ( 0, 2, 3, 3 )
-    // CHECK-NEXT: ( 0, 2, 0, 2 )
-    // CHECK-NEXT: ( 1, 2, 3, 4 )
+    // CHECK-NEXT: ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 4
+    // CHECK-NEXT: dim = ( 4, 3 )
+    // CHECK-NEXT: lvl = ( 4, 3 )
+    // CHECK-NEXT: pos[0] : ( 0, 4,
+    // CHECK-NEXT: crd[0] : ( 0, 2, 3, 3,
+    // CHECK-NEXT: crd[1] : ( 0, 2, 0, 2,
+    // CHECK-NEXT: values : ( 1, 2, 3, 4,
+    // CHECK-NEXT: ----
     //
     %cooa = tensor.empty() : tensor<4x3xf64, #SortedCOO>
     %coo1 = sparse_tensor.insert %f1 into %cooa[%c0, %c0] : tensor<4x3xf64, #SortedCOO>
@@ -171,14 +98,19 @@ module {
     %coo3 = sparse_tensor.insert %f3 into %coo2[%c3, %c0] : tensor<4x3xf64, #SortedCOO>
     %coo4 = sparse_tensor.insert %f4 into %coo3[%c3, %c2] : tensor<4x3xf64, #SortedCOO>
     %coom = sparse_tensor.load %coo4 hasInserts : tensor<4x3xf64, #SortedCOO>
-    call @dump_coo(%coom) : (tensor<4x3xf64, #SortedCOO>) -> ()
+    sparse_tensor.print %coom : tensor<4x3xf64, #SortedCOO>
 
     //
     // CSR case.
     //
-    // CHECK-NEXT: ( 0, 1, 1, 2, 4 )
-    // CHECK-NEXT: ( 0, 2, 0, 2 )
-    // CHECK-NEXT: ( 1, 2, 3, 4 )
+    // CHECK-NEXT: ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 4
+    // CHECK-NEXT: dim = ( 4, 3 )
+    // CHECK-NEXT: lvl = ( 4, 3 )
+    // CHECK-NEXT: pos[1] : ( 0, 1, 1, 2, 4,
+    // CHECK-NEXT: crd[1] : ( 0, 2, 0, 2,
+    // CHECK-NEXT: values : ( 1, 2, 3, 4,
+    // CHECK-NEXT: ----
     //
     %csra = tensor.empty() : tensor<4x3xf64, #CSR>
     %csr1 = sparse_tensor.insert %f1 into %csra[%c0, %c0] : tensor<4x3xf64, #CSR>
@@ -186,16 +118,21 @@ module {
     %csr3 = sparse_tensor.insert %f3 into %csr2[%c3, %c0] : tensor<4x3xf64, #CSR>
     %csr4 = sparse_tensor.insert %f4 into %csr3[%c3, %c2] : tensor<4x3xf64, #CSR>
     %csrm = sparse_tensor.load %csr4 hasInserts : tensor<4x3xf64, #CSR>
-    call @dump_csr(%csrm) : (tensor<4x3xf64, #CSR>) -> ()
+    sparse_tensor.print %csrm : tensor<4x3xf64, #CSR>
 
     //
     // DCSR case.
     //
-    // CHECK-NEXT: ( 0, 3 )
-    // CHECK-NEXT: ( 0, 2, 3 )
-    // CHECK-NEXT: ( 0, 1, 2, 4 )
-    // CHECK-NEXT: ( 0, 2, 0, 2 )
-    // CHECK-NEXT: ( 1, 2, 3, 4 )
+    // CHECK-NEXT: ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 4
+    // CHECK-NEXT: dim = ( 4, 3 )
+    // CHECK-NEXT: lvl = ( 4, 3 )
+    // CHECK-NEXT: pos[0] : ( 0, 3,
+    // CHECK-NEXT: crd[0] : ( 0, 2, 3,
+    // CHECK-NEXT: pos[1] : ( 0, 1, 2, 4,
+    // CHECK-NEXT: crd[1] : ( 0, 2, 0, 2,
+    // CHECK-NEXT: values : ( 1, 2, 3, 4,
+    // CHECK-NEXT: ----
     //
     %dcsra = tensor.empty() : tensor<4x3xf64, #DCSR>
     %dcsr1 = sparse_tensor.insert %f1 into %dcsra[%c0, %c0] : tensor<4x3xf64, #DCSR>
@@ -203,14 +140,19 @@ module {
     %dcsr3 = sparse_tensor.insert %f3 into %dcsr2[%c3, %c0] : tensor<4x3xf64, #DCSR>
     %dcsr4 = sparse_tensor.insert %f4 into %dcsr3[%c3, %c2] : tensor<4x3xf64, #DCSR>
     %dcsrm = sparse_tensor.load %dcsr4 hasInserts : tensor<4x3xf64, #DCSR>
-    call @dump_dcsr(%dcsrm) : (tensor<4x3xf64, #DCSR>) -> ()
+    sparse_tensor.print %dcsrm : tensor<4x3xf64, #DCSR>
 
     //
     // Row case.
     //
-    // CHECK-NEXT: ( 0, 3 )
-    // CHECK-NEXT: ( 0, 2, 3 )
-    // CHECK-NEXT: ( 1, 0, 0, 0, 0, 2, 3, 0, 4 )
+    // CHECK-NEXT: ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 9
+    // CHECK-NEXT: dim = ( 4, 3 )
+    // CHECK-NEXT: lvl = ( 4, 3 )
+    // CHECK-NEXT: pos[0] : ( 0, 3,
+    // CHECK-NEXT: crd[0] : ( 0, 2, 3,
+    // CHECK-NEXT: values : ( 1, 0, 0, 0, 0, 2, 3, 0, 4,
+    // CHECK-NEXT: ----
     //
     %rowa = tensor.empty() : tensor<4x3xf64, #Row>
     %row1 = sparse_tensor.insert %f1 into %rowa[%c0, %c0] : tensor<4x3xf64, #Row>
@@ -218,27 +160,7 @@ module {
     %row3 = sparse_tensor.insert %f3 into %row2[%c3, %c0] : tensor<4x3xf64, #Row>
     %row4 = sparse_tensor.insert %f4 into %row3[%c3, %c2] : tensor<4x3xf64, #Row>
     %rowm = sparse_tensor.load %row4 hasInserts : tensor<4x3xf64, #Row>
-    call @dump_row(%rowm) : (tensor<4x3xf64, #Row>) -> ()
-
-    //
-    // NOE sanity check.
-    //
-    // CHECK-NEXT: 12
-    // CHECK-NEXT: 4
-    // CHECK-NEXT: 4
-    // CHECK-NEXT: 4
-    // CHECK-NEXT: 9
-    //
-    %noe1 = sparse_tensor.number_of_entries %densem : tensor<4x3xf64, #Dense>
-    %noe2 = sparse_tensor.number_of_entries %coom : tensor<4x3xf64, #SortedCOO>
-    %noe3 = sparse_tensor.number_of_entries %csrm : tensor<4x3xf64, #CSR>
-    %noe4 = sparse_tensor.number_of_entries %dcsrm : tensor<4x3xf64, #DCSR>
-    %noe5 = sparse_tensor.number_of_entries %rowm : tensor<4x3xf64, #Row>
-    vector.print %noe1 : index
-    vector.print %noe2 : index
-    vector.print %noe3 : index
-    vector.print %noe4 : index
-    vector.print %noe5 : index
+    sparse_tensor.print %rowm : tensor<4x3xf64, #Row>
 
     // Release resources.
     bufferization.dealloc_tensor %densem : tensor<4x3xf64, #Dense>



More information about the Mlir-commits mailing list