66 lines · plain
1//2// NOTE: this test requires gpu-sm803//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 %s11 12#CSR = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }>13 14module {15 // Compute matrix vector y = Ax16 func.func @matvec(%A: tensor<1024x64xf64, #CSR>, %x: tensor<64xf64>, %y_in: tensor<1024xf64>) -> tensor<1024xf64> {17 %y_out = linalg.matvec18 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 : f6427 %c0 = arith.constant 0 : index28 %c1 = arith.constant 1 : index29 30 // Stress test with a dense matrix DA.31 %DA = tensor.generate {32 ^bb0(%i: index, %j: index):33 %k = arith.addi %i, %j : index34 %l = arith.index_cast %k : index to i6435 %f = arith.uitofp %l : i64 to f6436 tensor.yield %f : f6437 } : 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 : f6446 } : 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 return64 }65}66