1// 2// NOTE: this test requires gpu-sm80 3// 4// RUN: mlir-opt %s \ 5// RUN: --sparsifier="enable-runtime-library=false parallelization-strategy=dense-outer-loop gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71 gpu-format=%gpu_compilation_format" \ 6// RUN: | mlir-runner \ 7// RUN: --shared-libs=%mlir_cuda_runtime \ 8// RUN: --shared-libs=%mlir_c_runner_utils \ 9// RUN: --e main --entry-point-result=void \ 10// RUN: | FileCheck %s 11 12#CSR = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }> 13 14module { 15 // Compute matrix vector y = Ax 16 func.func @matvec(%A: tensor<?x?xf64, #CSR>, %x: tensor<?xf64>, %y_in: tensor<?xf64>) -> tensor<?xf64> { 17 %y_out = linalg.matvec 18 ins(%A, %x: tensor<?x?xf64, #CSR>, tensor<?xf64>) 19 outs(%y_in: tensor<?xf64>) -> tensor<?xf64> 20 return %y_out : tensor<?xf64> 21 } 22 23 func.func @main() { 24 %f0 = arith.constant 0.0 : f64 25 %c0 = arith.constant 0 : index 26 %c1 = arith.constant 1 : index 27 28 // Stress test with a dense matrix DA. 29 %DA = tensor.generate { 30 ^bb0(%i: index, %j: index): 31 %k = arith.addi %i, %j : index 32 %l = arith.index_cast %k : index to i64 33 %f = arith.uitofp %l : i64 to f64 34 tensor.yield %f : f64 35 } : tensor<1024x64xf64> 36 37 // Convert to a "sparse" m x n matrix A. 38 %A = sparse_tensor.convert %DA : tensor<1024x64xf64> to tensor<?x?xf64, #CSR> 39 40 // Initialize dense vector with n elements: 41 // (1, 2, 3, 4, ..., n) 42 %d1 = tensor.dim %A, %c1 : tensor<?x?xf64, #CSR> 43 %x = tensor.generate %d1 { 44 ^bb0(%i : index): 45 %k = arith.addi %i, %c1 : index 46 %j = arith.index_cast %k : index to i64 47 %f = arith.uitofp %j : i64 to f64 48 tensor.yield %f : f64 49 } : tensor<?xf64> 50 51 // Initialize dense vector to m zeros. 52 %d0 = tensor.dim %A, %c0 : tensor<?x?xf64, #CSR> 53 %y = tensor.generate %d0 { 54 ^bb0(%i : index): 55 tensor.yield %f0 : f64 56 } : tensor<?xf64> 57 58 // Call the kernel. 59 %0 = call @matvec(%A, %x, %y) : (tensor<?x?xf64, #CSR>, tensor<?xf64>, tensor<?xf64>) -> tensor<?xf64> 60 61 // 62 // Sanity check on results. 63 // 64 // CHECK: ( 87360, 89440, 91520, 93600, 95680, 97760, 99840, 101920, 104000, 106080, 108160, 110240, 112320, 114400, 116480, 118560, 120640, 122720, 124800, 126880, 128960, 131040, 133120, 135200, 137280, 139360, 141440, 143520, 145600, 147680, 149760, 151840, 153920, 156000, 158080, 160160, 162240, 164320, 166400, 168480, 170560, 172640, 174720, 176800, 178880, 180960, 183040, 185120, 187200, 189280, 191360, 193440, 195520, 197600, 199680, 201760, 203840, 205920, 208000, 210080, 212160, 214240, 216320, 218400 ) 65 // 66 %pb0 = vector.transfer_read %0[%c0], %f0 : tensor<?xf64>, vector<64xf64> 67 vector.print %pb0 : vector<64xf64> 68 69 // Release the resources. 70 bufferization.dealloc_tensor %A : tensor<?x?xf64, #CSR> 71 return 72 } 73} 74