1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
|
// RUN: mlir-opt %s --linalg-generalize-named-ops \
// RUN: --pre-sparsification-rewrite \
// RUN: --sparsification="parallelization-strategy=dense-outer-loop" \
// RUN: --sparse-gpu-codegen | FileCheck %s
#CSR = #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>
//
// CHECK-LABEL: gpu.module @sparse_kernels
// CHECK: gpu.func @kernel1
// CHECK: gpu.func @kernel0
//
// CHECK-LABEL: func.func @matmuls
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: %[[T1:.*]] = gpu.launch_func async @sparse_kernels::@kernel1 blocks
// CHECK: gpu.memcpy async [%[[T1]]]
// CHECK: gpu.dealloc async
// CHECK: gpu.dealloc async
// CHECK: gpu.dealloc async
// CHECK: gpu.dealloc async
// CHECK: gpu.dealloc async
// CHECK: gpu.wait
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: gpu.alloc async
// CHECK: gpu.memcpy async
// CHECK: %[[T0:.*]] = gpu.launch_func async @sparse_kernels::@kernel0 blocks
// CHECK: gpu.memcpy async [%[[T0]]]
// CHECK: gpu.dealloc async
// CHECK: gpu.dealloc async
// CHECK: gpu.dealloc async
// CHECK: gpu.dealloc async
// CHECK: gpu.dealloc async
// CHECK: gpu.wait
//
func.func @matmuls(%A: tensor<1024x8xf64>,
%B: tensor<8x1024xf64, #CSR>,
%C: tensor<1024x1024xf64, #CSR>) -> tensor<1024x1024xf64> {
%Z = arith.constant dense<0.0> : tensor<1024x1024xf64>
%T = linalg.matmul
ins(%A, %B: tensor<1024x8xf64>, tensor<8x1024xf64, #CSR>)
outs(%Z: tensor<1024x1024xf64>) -> tensor<1024x1024xf64>
%D = linalg.matmul
ins(%T, %C: tensor<1024x1024xf64>, tensor<1024x1024xf64, #CSR>)
outs(%Z: tensor<1024x1024xf64>) -> tensor<1024x1024xf64>
return %D : tensor<1024x1024xf64>
}
|
