diff --git a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-lib.mlir b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-lib.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-matvec-lib.mlir @@ -0,0 +1,98 @@ +// +// NOTE: this test requires gpu-sm80 +// +// with RT lib (SoA COO): +// +// RUN: mlir-opt %s \ +// RUN: --sparse-compiler="enable-runtime-library=true enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71" \ +// RUN: | mlir-cpu-runner \ +// RUN: --shared-libs=%mlir_cuda_runtime \ +// RUN: --shared-libs=%mlir_runner_utils \ +// RUN: --e main --entry-point-result=void \ +// RUN: | FileCheck %s +// +// TODO: without RT lib (AoS COO): + +#SortedCOO = #sparse_tensor.encoding<{ + dimLevelType = [ "compressed-nu", "singleton" ] +}> + +#CSR = #sparse_tensor.encoding<{ + dimLevelType = [ "dense", "compressed" ], + posWidth = 32, + crdWidth = 32 +}> + +module { + // Compute matrix vector y = Ax on COO with default index coordinates. + func.func @matvecCOO(%A: tensor, %x: tensor, %y_in: tensor) -> tensor { + %y_out = linalg.matvec + ins(%A, %x: tensor, tensor) + outs(%y_in: tensor) -> tensor + return %y_out : tensor + } + + // Compute matrix vector y = Ax on CSR with 32-bit positions and coordinates. + func.func @matvecCSR(%A: tensor, %x: tensor, %y_in: tensor) -> tensor { + %y_out = linalg.matvec + ins(%A, %x: tensor, tensor) + outs(%y_in: tensor) -> tensor + return %y_out : tensor + } + + func.func @main() { + %f0 = arith.constant 0.0 : f64 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + + // Stress test with a dense matrix DA. + %DA = tensor.generate { + ^bb0(%i: index, %j: index): + %k = arith.addi %i, %j : index + %l = arith.index_cast %k : index to i64 + %f = arith.uitofp %l : i64 to f64 + tensor.yield %f : f64 + } : tensor<1024x64xf64> + + // Convert to a "sparse" m x n matrix A. + %Acoo = sparse_tensor.convert %DA : tensor<1024x64xf64> to tensor + %Acsr = sparse_tensor.convert %DA : tensor<1024x64xf64> to tensor + + // Initialize dense vector with n elements: + // (1, 2, 3, 4, ..., n) + %d1 = tensor.dim %Acoo, %c1 : tensor + %x = tensor.generate %d1 { + ^bb0(%i : index): + %k = arith.addi %i, %c1 : index + %j = arith.index_cast %k : index to i64 + %f = arith.uitofp %j : i64 to f64 + tensor.yield %f : f64 + } : tensor + + // Initialize dense vector to m zeros. + %d0 = tensor.dim %Acoo, %c0 : tensor + %y = tensor.generate %d0 { + ^bb0(%i : index): + tensor.yield %f0 : f64 + } : tensor + + // Call the kernels. + %0 = call @matvecCOO(%Acoo, %x, %y) : (tensor, tensor, tensor) -> tensor + %1 = call @matvecCSR(%Acsr, %x, %y) : (tensor, tensor, tensor) -> tensor + + // + // Sanity check on results. + // + // CHECK-COUNT-2: ( 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 ) + // + %pb0 = vector.transfer_read %0[%c0], %f0 : tensor, vector<64xf64> + vector.print %pb0 : vector<64xf64> + %pb1 = vector.transfer_read %0[%c0], %f0 : tensor, vector<64xf64> + vector.print %pb1 : vector<64xf64> + + // Release the resources. + bufferization.dealloc_tensor %Acoo : tensor + bufferization.dealloc_tensor %Acsr : tensor + return + } +}