132 lines · plain
1//--------------------------------------------------------------------------------------------------2// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS.3//4// Set-up that's shared across all tests in this directory. In principle, this5// config could be moved to lit.local.cfg. However, there are downstream users that6// do not use these LIT config files. Hence why this is kept inline.7//8// DEFINE: %{sparsifier_opts} = enable-runtime-library=true9// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts}10// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"11// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"12// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils13// DEFINE: %{run_libs_sve} = -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils14// DEFINE: %{run_opts} = -e main -entry-point-result=void15// DEFINE: %{run} = mlir-runner %{run_opts} %{run_libs}16// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs_sve}17//18// DEFINE: %{env} =19//--------------------------------------------------------------------------------------------------20 21// REDEFINE: %{env} = TENSOR0="%mlir_src_dir/test/Integration/data/block.mtx"22// RUN: %{compile} | env %{env} %{run} | FileCheck %s23//24// Do the same run, but now with direct IR generation.25// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false26// RUN: %{compile} | env %{env} %{run} | FileCheck %s27//28// Do the same run, but now with direct IR generation and vectorization.29// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false enable-buffer-initialization=true vl=2 reassociate-fp-reductions=true enable-index-optimizations=true30// RUN: %{compile} | env %{env} %{run} | FileCheck %s31 32!Filename = !llvm.ptr33 34#BSR = #sparse_tensor.encoding<{35 map = (i, j) ->36 ( i floordiv 2 : dense37 , j floordiv 2 : compressed38 , i mod 2 : dense39 , j mod 2 : dense40 )41}>42 43#DSDD = #sparse_tensor.encoding<{44 map = (i, j, k, l) -> ( i : dense, j : compressed, k : dense, l : dense)45}>46 47#trait_scale_inplace = {48 indexing_maps = [49 affine_map<(i,j) -> (i,j)> // X (out)50 ],51 iterator_types = ["parallel", "parallel"]52}53 54//55// Example 2x2 block storage:56//57// +-----+-----+-----+ +-----+-----+-----+58// | 1 2 | . . | 4 . | | 1 2 | | 4 0 |59// | . 3 | . . | . 5 | | 0 3 | | 0 5 |60// +-----+-----+-----+ => +-----+-----+-----+61// | . . | 6 7 | . . | | | 6 7 | |62// | . . | 8 . | . . | | | 8 0 | |63// +-----+-----+-----+ +-----+-----+-----+64//65// Stored as:66//67// positions[1] : 0 2 368// coordinates[1] : 0 2 169// values : 1.000000 2.000000 0.000000 3.000000 4.000000 0.000000 0.000000 5.000000 6.000000 7.000000 8.000000 0.00000070//71module {72 73 func.func private @getTensorFilename(index) -> (!Filename)74 75 func.func @scale(%arg0: tensor<?x?xf64, #BSR>) -> tensor<?x?xf64, #BSR> {76 %c = arith.constant 3.0 : f6477 %0 = linalg.generic #trait_scale_inplace78 outs(%arg0: tensor<?x?xf64, #BSR>) {79 ^bb(%x: f64):80 %1 = arith.mulf %x, %c : f6481 linalg.yield %1 : f6482 } -> tensor<?x?xf64, #BSR>83 return %0 : tensor<?x?xf64, #BSR>84 }85 86 func.func @main() {87 %c0 = arith.constant 0 : index88 %f0 = arith.constant 0.0 : f6489 90 %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)91 %A = sparse_tensor.new %fileName : !Filename to tensor<?x?xf64, #BSR>92 93 // CHECK: ---- Sparse Tensor ----94 // CHECK-NEXT: nse = 1295 // CHECK-NEXT: dim = ( 4, 6 )96 // CHECK-NEXT: lvl = ( 2, 3, 2, 2 )97 // CHECK-NEXT: pos[1] : ( 0, 2, 3 )98 // CHECK-NEXT: crd[1] : ( 0, 2, 1 )99 // CHECK-NEXT: values : ( 1, 2, 0, 3, 4, 0, 0, 5, 6, 7, 8, 0 )100 // CHECK-NEXT: ----101 sparse_tensor.print %A : tensor<?x?xf64, #BSR>102 103 // CHECK-NEXT: ---- Sparse Tensor ----104 // CHECK-NEXT: nse = 12105 // CHECK-NEXT: dim = ( 2, 3, 2, 2 )106 // CHECK-NEXT: lvl = ( 2, 3, 2, 2 )107 // CHECK-NEXT: pos[1] : ( 0, 2, 3 )108 // CHECK-NEXT: crd[1] : ( 0, 2, 1 )109 // CHECK-NEXT: values : ( 1, 2, 0, 3, 4, 0, 0, 5, 6, 7, 8, 0 )110 // CHECK-NEXT: ----111 %t1 = sparse_tensor.reinterpret_map %A : tensor<?x?xf64, #BSR>112 to tensor<?x?x2x2xf64, #DSDD>113 sparse_tensor.print %t1 : tensor<?x?x2x2xf64, #DSDD>114 115 // CHECK-NEXT: ---- Sparse Tensor ----116 // CHECK-NEXT: nse = 12117 // CHECK-NEXT: dim = ( 4, 6 )118 // CHECK-NEXT: lvl = ( 2, 3, 2, 2 )119 // CHECK-NEXT: pos[1] : ( 0, 2, 3 )120 // CHECK-NEXT: crd[1] : ( 0, 2, 1 )121 // CHECK-NEXT: values : ( 3, 6, 0, 9, 12, 0, 0, 15, 18, 21, 24, 0 )122 // CHECK-NEXT: ----123 %As = call @scale(%A) : (tensor<?x?xf64, #BSR>) -> (tensor<?x?xf64, #BSR>)124 sparse_tensor.print %As : tensor<?x?xf64, #BSR>125 126 // Release the resources.127 bufferization.dealloc_tensor %A: tensor<?x?xf64, #BSR>128 129 return130 }131}132