// NOTE: this test requires gpu-sm80
//
// DEFINE: %{compile} = mlir-opt %s \
// DEFINE: --sparsifier="enable-gpu-libgen gpu-triple=nvptx64-nvidia-cuda gpu-chip=sm_80 gpu-features=+ptx71 gpu-format=%gpu_compilation_format
// DEFINE: %{run} = \
// DEFINE: env TENSOR0="%mlir_src_dir/test/Integration/data/block.mtx" \
// DEFINE: mlir-cpu-runner \
// DEFINE: --shared-libs=%mlir_cuda_runtime \
// DEFINE: --shared-libs=%mlir_c_runner_utils \
// DEFINE: --e main --entry-point-result=void \
// DEFINE: | FileCheck %s
//
// with RT lib:
//
// RUN: %{compile} enable-runtime-library=true" | %{run}
//
// without RT lib:
//
// RUN: %{compile} enable-runtime-library=false" | %{run}
!Filename = !llvm.ptr
#CSR = #sparse_tensor.encoding<{
map = (d0, d1) -> (d0 : dense, d1 : compressed)
}>
#BSR = #sparse_tensor.encoding<{
map = (i, j) -> (
i floordiv 2 : dense,
j floordiv 2 : compressed,
i mod 2 : dense,
j mod 2 : dense)
}>
#trait_SDDMM = {
indexing_maps = [
affine_map<(i,j,k) -> (i,k)>, // A
affine_map<(i,j,k) -> (k,j)>, // B
affine_map<(i,j,k) -> (i,j)> // S (in/out)
],
iterator_types = ["parallel", "parallel", "reduction"],
doc = "S(i,j) += spy[S(i,j)] x SUM_k A(i,k) B(k,j)"
}
//
// Integration test that lowers a kernel annotated as sparse to
// actual sparse code, initializes sparse storage schemes, and
// runs the resulting code with the JIT compiler.
//
module {
llvm.func @mgpuCreateSparseEnv()
llvm.func @mgpuDestroySparseEnv()
//
// A kernel that computes a CSR sampled dense matrix matrix multiplication
// using a "spy" function and in-place update of the sampling sparse matrix.
//
func.func @SDDMM(%args: tensor<?x?xf32, #CSR>,
%arga: tensor<?x?xf32>,
%argb: tensor<?x?xf32>) -> tensor<?x?xf32, #CSR> {
%result = linalg.generic #trait_SDDMM
ins(%arga, %argb: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%args: tensor<?x?xf32, #CSR>) {
^bb(%a: f32, %b: f32, %s: f32):
%f0 = arith.constant 0.0 : f32
%u = sparse_tensor.unary %s : f32 to f32
present={
^bb0(%p: f32):
%mul = arith.mulf %a, %b : f32
sparse_tensor.yield %mul : f32
}
absent={}
%r = sparse_tensor.reduce %s, %u, %f0 : f32 {
^bb0(%p: f32, %q: f32):
%add = arith.addf %p, %q : f32
sparse_tensor.yield %add : f32
}
linalg.yield %r : f32
} -> tensor<?x?xf32, #CSR>
return %result : tensor<?x?xf32, #CSR>
}
//
// A kernel that computes a BSR sampled dense matrix matrix multiplication
// using a "spy" function and in-place update of the sampling sparse matrix.
//
func.func @SDDMM_block(%args: tensor<?x?xf32, #BSR>,
%arga: tensor<?x?xf32>,
%argb: tensor<?x?xf32>) -> tensor<?x?xf32, #BSR> {
%result = linalg.generic #trait_SDDMM
ins(%arga, %argb: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%args: tensor<?x?xf32, #BSR>) {
^bb(%a: f32, %b: f32, %s: f32):
%f0 = arith.constant 0.0 : f32
%u = sparse_tensor.unary %s : f32 to f32
present={
^bb0(%p: f32):
%mul = arith.mulf %a, %b : f32
sparse_tensor.yield %mul : f32
}
absent={}
%r = sparse_tensor.reduce %s, %u, %f0 : f32 {
^bb0(%p: f32, %q: f32):
%add = arith.addf %p, %q : f32
sparse_tensor.yield %add : f32
}
linalg.yield %r : f32
} -> tensor<?x?xf32, #BSR>
return %result : tensor<?x?xf32, #BSR>
}
func.func private @getTensorFilename(index) -> (!Filename)
//
// Main driver.
//
func.func @main() {
llvm.call @mgpuCreateSparseEnv() : () -> ()
%d0 = arith.constant 0.0 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c4 = arith.constant 4 : index
%c6 = arith.constant 6 : index
// Initialize dense matrices.
%a = tensor.generate %c4, %c4 {
^bb0(%i: index, %j: index):
%p = arith.addi %i, %c1 : index
%q = arith.index_cast %p : index to i32
%d = arith.sitofp %q : i32 to f32
tensor.yield %d : f32
} : tensor<?x?xf32>
%b = tensor.generate %c4, %c6 {
^bb0(%i: index, %j: index):
%p = arith.addi %j, %c1 : index
%q = arith.index_cast %p : index to i32
%d = arith.sitofp %q : i32 to f32
tensor.yield %d : f32
} : tensor<?x?xf32>
// Read the sparse matrix from file, construct sparse storage.
//
// +-----+-----+-----+
// | 1 2 | . . | 4 . |
// | . 3 | . . | . 5 |
// +-----+-----+-----+
// | . . | 6 7 | . . |
// | . . | 8 . | . . |
// +-----+-----+-----+
//
%fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
%m_csr = sparse_tensor.new %fileName : !Filename to tensor<?x?xf32, #CSR>
%m_bsr = sparse_tensor.new %fileName : !Filename to tensor<?x?xf32, #BSR>
// Call the kernel.
%0 = call @SDDMM(%m_csr, %a, %b)
: (tensor<?x?xf32, #CSR>,
tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32, #CSR>
%1 = call @SDDMM_block(%m_bsr, %a, %b)
: (tensor<?x?xf32, #BSR>,
tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32, #BSR>
//
// Print the result for verification. Note that the "spy" determines what
// dot products are sampled, but the original contents are added back to
// the result (which is why the block sparse version has actual results
// in the original zero positions).
//
// CHECK: ---- Sparse Tensor ----
// CHECK-NEXT: nse = 8
// CHECK-NEXT: dim = ( 4, 6 )
// CHECK-NEXT: lvl = ( 4, 6 )
// CHECK-NEXT: pos[1] : ( 0, 3, 5, 7, 8 )
// CHECK-NEXT: crd[1] : ( 0, 1, 4, 1, 5, 2, 3, 2 )
// CHECK-NEXT: values : ( 5, 10, 24, 19, 53, 42, 55, 56 )
// CHECK-NEXT: ----
//
// CHECK: ---- Sparse Tensor ----
// CHECK-NEXT: nse = 12
// CHECK-NEXT: dim = ( 4, 6 )
// CHECK-NEXT: lvl = ( 2, 3, 2, 2 )
// CHECK-NEXT: pos[1] : ( 0, 2, 3 )
// CHECK-NEXT: crd[1] : ( 0, 2, 1 )
// CHECK-NEXT: values : ( 5, 10, 8, 19, 24, 24, 40, 53, 42, 55, 56, 64 )
// CHECK-NEXT: ----
//
sparse_tensor.print %0 : tensor<?x?xf32, #CSR>
sparse_tensor.print %1 : tensor<?x?xf32, #BSR>
// Release the resources.
bufferization.dealloc_tensor %0 : tensor<?x?xf32, #CSR>
bufferization.dealloc_tensor %1 : tensor<?x?xf32, #BSR>
llvm.call @mgpuDestroySparseEnv() : () -> ()
return
}
}
Codebase Browser
llvm