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<1024x64xf64, #CSR>, %x: tensor<64xf64>, %y_in: tensor<1024xf64>) -> tensor<1024xf64> { 17 %y_out = linalg.matvec 18 ins(%A, %x: tensor<1024x64xf64, #CSR>, tensor<64xf64>) 19 outs(%y_in: tensor<1024xf64>) -> tensor<1024xf64> 20 return %y_out : tensor<1024xf64> 21 } 22 23 memref.global "private" constant @__constant_64xf64 : memref<64xf64> = dense<1.000000e+00> {alignment = 64 : i64} 24 25 func.func @main() { 26 %f0 = arith.constant 0.0 : f64 27 %c0 = arith.constant 0 : index 28 %c1 = arith.constant 1 : index 29 30 // Stress test with a dense matrix DA. 31 %DA = tensor.generate { 32 ^bb0(%i: index, %j: index): 33 %k = arith.addi %i, %j : index 34 %l = arith.index_cast %k : index to i64 35 %f = arith.uitofp %l : i64 to f64 36 tensor.yield %f : f64 37 } : tensor<1024x64xf64> 38 39 // Convert to a "sparse" 1024 x 64 matrix A. 40 %A = sparse_tensor.convert %DA : tensor<1024x64xf64> to tensor<1024x64xf64, #CSR> 41 42 // Initialize dense vector to 1024 zeros. 43 %y = tensor.generate { 44 ^bb0(%i : index): 45 tensor.yield %f0 : f64 46 } : tensor<1024xf64> 47 48 // Call the kernel with an vector taken from global memory. 49 %xbuf = memref.get_global @__constant_64xf64 : memref<64xf64> 50 %x = bufferization.to_tensor %xbuf restrict : memref<64xf64> to tensor<64xf64> 51 %0 = call @matvec(%A, %x, %y) : (tensor<1024x64xf64, #CSR>, tensor<64xf64>, tensor<1024xf64>) -> tensor<1024xf64> 52 53 // 54 // Sanity check on results. 55 // 56 // CHECK: ( 2016, 2080, 2144, 2208, 2272, 2336, 2400, 2464, 2528, 2592, 2656, 2720, 2784, 2848, 2912, 2976, 3040, 3104, 3168, 3232, 3296, 3360, 3424, 3488, 3552, 3616, 3680, 3744, 3808, 3872, 3936, 4000, 4064, 4128, 4192, 4256, 4320, 4384, 4448, 4512, 4576, 4640, 4704, 4768, 4832, 4896, 4960, 5024, 5088, 5152, 5216, 5280, 5344, 5408, 5472, 5536, 5600, 5664, 5728, 5792, 5856, 5920, 5984, 6048 ) 57 // 58 %pb0 = vector.transfer_read %0[%c0], %f0 : tensor<1024xf64>, vector<64xf64> 59 vector.print %pb0 : vector<64xf64> 60 61 // Release the resources. 62 bufferization.dealloc_tensor %A : tensor<1024x64xf64, #CSR> 63 return 64 } 65} 66