• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1// RUN: mlir-opt %s \
2// RUN:  -convert-scf-to-std -convert-vector-to-scf \
3// RUN:  -convert-linalg-to-llvm -convert-vector-to-llvm | \
4// RUN: MATRIX0="%mlir_integration_test_dir/data/test.mtx" \
5// RUN: mlir-cpu-runner \
6// RUN:  -e entry -entry-point-result=void  \
7// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
8// RUN: FileCheck %s
9
10module {
11  func private @openMatrix(!llvm.ptr<i8>, memref<index>, memref<index>, memref<index>) -> ()
12  func private @readMatrixItem(memref<index>, memref<index>, memref<f64>) -> ()
13  func private @closeMatrix() -> ()
14  func private @getMatrix(index) -> (!llvm.ptr<i8>)
15
16  func @entry() {
17    %d0  = constant 0.0 : f64
18    %c0  = constant 0 : index
19    %c1  = constant 1 : index
20    %c5  = constant 5 : index
21    %m   = alloc() : memref<index>
22    %n   = alloc() : memref<index>
23    %nnz = alloc() : memref<index>
24    %i   = alloc() : memref<index>
25    %j   = alloc() : memref<index>
26    %d   = alloc() : memref<f64>
27
28    //
29    // Read the header of a sparse matrix. This yields the
30    // size (m x n) and number of nonzero elements (nnz).
31    //
32    %file = call @getMatrix(%c0) : (index) -> (!llvm.ptr<i8>)
33    call @openMatrix(%file, %m, %n, %nnz)
34        : (!llvm.ptr<i8>, memref<index>,
35	                  memref<index>, memref<index>) -> ()
36    %M = load %m[]   : memref<index>
37    %N = load %n[]   : memref<index>
38    %Z = load %nnz[] : memref<index>
39
40    //
41    // At this point, code should prepare a proper sparse storage
42    // scheme for an m x n matrix with nnz nonzero elements. For
43    // simplicity, however, here we simply set up a dense matrix.
44    //
45    %a = alloc(%M, %N) : memref<?x?xf64>
46    scf.for %ii = %c0 to %M step %c1 {
47      scf.for %jj = %c0 to %N step %c1 {
48        store %d0, %a[%ii, %jj] : memref<?x?xf64>
49      }
50    }
51
52    //
53    // Now we are ready to read in the nonzero elements of the
54    // sparse matrix and insert these into a sparse storage
55    // scheme. In this example, we simply insert them in the
56    // dense matrix.
57    //
58    scf.for %k = %c0 to %Z step %c1 {
59      call @readMatrixItem(%i, %j, %d)
60          : (memref<index>, memref<index>, memref<f64>) -> ()
61      %I = load %i[] : memref<index>
62      %J = load %j[] : memref<index>
63      %D = load %d[] : memref<f64>
64      store %D, %a[%I, %J] : memref<?x?xf64>
65    }
66    call @closeMatrix() : () -> ()
67
68    //
69    // Verify that the results are as expected.
70    //
71    %A = vector.transfer_read %a[%c0, %c0], %d0 : memref<?x?xf64>, vector<5x5xf64>
72    vector.print %M : index
73    vector.print %N : index
74    vector.print %Z : index
75    vector.print %A : vector<5x5xf64>
76    //
77    // CHECK: 5
78    // CHECK: 5
79    // CHECK: 9
80    //
81    // CHECK: ( ( 1, 0, 0, 1.4, 0 ),
82    // CHECK-SAME: ( 0, 2, 0, 0, 2.5 ),
83    // CHECK-SAME: ( 0, 0, 3, 0, 0 ),
84    // CHECK-SAME: ( 4.1, 0, 0, 4, 0 ),
85    // CHECK-SAME: ( 0, 5.2, 0, 0, 5 ) )
86
87    //
88    // Free.
89    //
90    dealloc %m   : memref<index>
91    dealloc %n   : memref<index>
92    dealloc %nnz : memref<index>
93    dealloc %i   : memref<index>
94    dealloc %j   : memref<index>
95    dealloc %d   : memref<f64>
96    dealloc %a   : memref<?x?xf64>
97
98    return
99  }
100}
101