Skip to content

Commit

Permalink
[mlir][sparse] migrate datastructure tests to sparse_tensor.print (#8…
Browse files Browse the repository at this point in the history 
…3956)

Continuing the efforts started in #83357
  • Loading branch information
@aartbik
aartbik authored Mar 5, 2024
1 parent d016712 commit 662d821
Show file tree
Hide file tree
Showing 2 changed files with 72 additions and 166 deletions.
62 changes: 23 additions & 39 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_1d.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}
//
Expand All @@ -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 = {
Expand All @@ -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
Expand All @@ -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>
Expand All @@ -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>
Expand Down
176 changes: 49 additions & 127 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_2d.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}
//
Expand All @@ -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<{
Expand All @@ -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
Expand All @@ -147,98 +64,103 @@ 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>
%dense2 = sparse_tensor.insert %f2 into %dense1[%c2, %c2] : tensor<4x3xf64, #Dense>
%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>
%coo2 = sparse_tensor.insert %f2 into %coo1[%c2, %c2] : tensor<4x3xf64, #SortedCOO>
%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>
%csr2 = sparse_tensor.insert %f2 into %csr1[%c2, %c2] : tensor<4x3xf64, #CSR>
%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>
%dcsr2 = sparse_tensor.insert %f2 into %dcsr1[%c2, %c2] : tensor<4x3xf64, #DCSR>
%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>
%row2 = sparse_tensor.insert %f2 into %row1[%c2, %c2] : tensor<4x3xf64, #Row>
%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>
Expand Down

0 comments on commit 662d821

Please sign in to comment.