mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir

   1 // RUN: mlir-opt %s \
   2 // RUN:   --sparsification --sparse-tensor-conversion \
   3 // RUN:   --convert-vector-to-scf --convert-scf-to-std \
   4 // RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
   5 // RUN:   --std-bufferize --finalizing-bufferize --lower-affine \
   6 // RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm --reconcile-unrealized-casts | \
   7 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
   8 // RUN: mlir-cpu-runner \
   9 // RUN:  -e entry -entry-point-result=void  \
  10 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
  11 // RUN: FileCheck %s
  12 //
  13 // Do the same run, but now with SIMDization as well. This should not change the outcome.
  14 //
  15 // RUN: mlir-opt %s \
  16 // RUN:   --sparsification="vectorization-strategy=2 vl=4" --sparse-tensor-conversion \
  17 // RUN:   --convert-vector-to-scf --convert-scf-to-std \
  18 // RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
  19 // RUN:   --std-bufferize --finalizing-bufferize --lower-affine \
  20 // RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm --reconcile-unrealized-casts | \
  21 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
  22 // RUN: mlir-cpu-runner \
  23 // RUN:  -e entry -entry-point-result=void  \
  24 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
  25 // RUN: FileCheck %s
  26
  27 !Filename = type !llvm.ptr<i8>
  28
  29 #DCSR = #sparse_tensor.encoding<{
  30   dimLevelType = [ "compressed", "compressed" ],
  31   dimOrdering = affine_map<(i,j) -> (i,j)>
  32 }>
  33
  34 #eltwise_mult = {
  35   indexing_maps = [
  36     affine_map<(i,j) -> (i,j)>  // X (out)
  37   ],
  38   iterator_types = ["parallel", "parallel"],
  39   doc = "X(i,j) *= X(i,j)"
  40 }
  41
  42 //
  43 // Integration test that lowers a kernel annotated as sparse to
  44 // actual sparse code, initializes a matching sparse storage scheme
  45 // from file, and runs the resulting code with the JIT compiler.
  46 //
  47 module {
  48   //
  49   // A kernel that multiplies a sparse matrix A with itself
  50   // in an element-wise fashion. In this operation, we have
  51   // a sparse tensor as output, but although the values of the
  52   // sparse tensor change, its nonzero structure remains the same.
  53   //
  54   func @kernel_eltwise_mult(%argx: tensor<?x?xf64, #DCSR> {linalg.inplaceable = true})
  55     -> tensor<?x?xf64, #DCSR> {
  56     %0 = linalg.generic #eltwise_mult
  57       outs(%argx: tensor<?x?xf64, #DCSR>) {
  58       ^bb(%x: f64):
  59         %0 = arith.mulf %x, %x : f64
  60         linalg.yield %0 : f64
  61     } -> tensor<?x?xf64, #DCSR>
  62     return %0 : tensor<?x?xf64, #DCSR>
  63   }
  64
  65   func private @getTensorFilename(index) -> (!Filename)
  66
  67   //
  68   // Main driver that reads matrix from file and calls the sparse kernel.
  69   //
  70   func @entry() {
  71     %d0 = arith.constant 0.0 : f64
  72     %c0 = arith.constant 0 : index
  73
  74     // Read the sparse matrix from file, construct sparse storage.
  75     %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
  76     %x = sparse_tensor.new %fileName : !Filename to tensor<?x?xf64, #DCSR>
  77
  78     // Call kernel.
  79     %0 = call @kernel_eltwise_mult(%x) : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR>
  80
  81     // Print the result for verification.
  82     //
  83     // CHECK: ( 1, 1.96, 4, 6.25, 9, 16.81, 16, 27.04, 25 )
  84     //
  85     %m = sparse_tensor.values %0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
  86     %v = vector.transfer_read %m[%c0], %d0: memref<?xf64>, vector<9xf64>
  87     vector.print %v : vector<9xf64>
  88
  89     // Release the resources.
  90     sparse_tensor.release %x : tensor<?x?xf64, #DCSR>
  91
  92     return
  93   }
  94 }