[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