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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
|
// RUN: mlir-opt %s -sparsification | FileCheck %s
#trait = {
indexing_maps = [
affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1, 0)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1, 0)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1, 0)>,
affine_map<(d0, d1, d2, d3) -> (d3)>,
affine_map<(d0, d1, d2, d3) -> (d3)>,
affine_map<(d0, d1, d2, d3) -> (d2, d3)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]
}
#VEC = #sparse_tensor.encoding<{ lvlTypes = [ "compressed" ], posWidth = 32, crdWidth = 32 }>
#COO = #sparse_tensor.encoding<{ lvlTypes = [ "compressed-nu", "singleton" ], posWidth = 32, crdWidth = 32 }>
#CCC = #sparse_tensor.encoding<{ lvlTypes = [ "compressed", "compressed", "compressed" ], posWidth = 32, crdWidth = 32 }>
//
// This kernel can be sparsified as all unsparsifiable operations'
// operands are loaded from dense tensors.
//
// CHECK-LABEL: func @dense_op_without_sp_dep
// CHECK-NOT: linalg.generic {{.*}}
func.func @dense_op_without_sp_dep(%169: tensor<2x10x8xf32>,
%expanded_54: tensor<2x10x1xf32>,
%expanded_56: tensor<2x10x1xf32>,
%expanded_57: tensor<2x10x1xf32>,
%176: tensor<8xf32, #VEC>,
%177: tensor<8xf32, #VEC>,
%9: tensor<100x8xf32, #COO>) -> tensor<2x10x100xf32> {
%cst_13 = arith.constant -3.40282347E+38 : f32
%178 = tensor.empty() : tensor<2x10x100xf32>
%179 = linalg.generic #trait
ins(%169, %expanded_54, %expanded_56, %expanded_57, %176, %177, %9 :
tensor<2x10x8xf32>, tensor<2x10x1xf32>, tensor<2x10x1xf32>, tensor<2x10x1xf32>,
tensor<8xf32, #VEC>, tensor<8xf32, #VEC>, tensor<100x8xf32, #COO>)
outs(%178 : tensor<2x10x100xf32>) {
^bb0(%in: f32, %in_58: f32, %in_59: f32, %in_60: f32, %in_61: f32, %in_62: f32, %in_63: f32, %out: f32):
%180 = arith.mulf %in_60, %in_60 : f32
%181 = arith.mulf %in_59, %cst_13 : f32
%182 = arith.subf %181, %180 : f32
%183 = arith.maxf %182, %cst_13 : f32
%184 = arith.addf %183, %cst_13 : f32
%185 = math.rsqrt %184 : f32 // data dependent on sparse value.
%186 = arith.mulf %185, %in_61 : f32
%187 = arith.subf %in, %in_58 : f32
%188 = arith.mulf %187, %186 : f32
%189 = arith.addf %188, %in_62 : f32
%190 = arith.mulf %189, %in_63 : f32
%191 = arith.addf %out, %190 : f32
linalg.yield %191 : f32
} -> tensor<2x10x100xf32>
return %179 : tensor<2x10x100xf32>
}
//
// This kernel cannot be sparsified as some unsparsifiable operations'
// operands are loaded from sparse tensors.
//
// CHECK-LABEL: func @dense_op_with_sp_dep
// CHECK: linalg.generic {{.*}}
func.func @dense_op_with_sp_dep(%169: tensor<2x10x8xf32>,
%expanded_54: tensor<2x10x1xf32, #CCC>,
%expanded_56: tensor<2x10x1xf32, #CCC>,
%expanded_57: tensor<2x10x1xf32, #CCC>,
%176: tensor<8xf32, #VEC>,
%177: tensor<8xf32, #VEC>,
%9: tensor<100x8xf32, #COO>) -> tensor<2x10x100xf32> {
%cst_13 = arith.constant -3.40282347E+38 : f32
%178 = tensor.empty() : tensor<2x10x100xf32>
%179 = linalg.generic #trait
ins(%169, %expanded_54, %expanded_56, %expanded_57, %176, %177, %9 :
tensor<2x10x8xf32>, tensor<2x10x1xf32, #CCC>, tensor<2x10x1xf32, #CCC>, tensor<2x10x1xf32, #CCC>,
tensor<8xf32, #VEC>, tensor<8xf32, #VEC>, tensor<100x8xf32, #COO>)
outs(%178 : tensor<2x10x100xf32>) {
^bb0(%in: f32, %in_58: f32, %in_59: f32, %in_60: f32, %in_61: f32, %in_62: f32, %in_63: f32, %out: f32):
%180 = arith.mulf %in_60, %in_60 : f32
%181 = arith.mulf %in_59, %cst_13 : f32
%182 = arith.subf %181, %180 : f32
%183 = arith.maxf %182, %cst_13 : f32
%184 = arith.addf %183, %cst_13 : f32
%185 = math.rsqrt %184 : f32
%186 = arith.mulf %185, %in_61 : f32
%187 = arith.subf %in, %in_58 : f32
%188 = arith.mulf %187, %186 : f32
%189 = arith.addf %188, %in_62 : f32
%190 = arith.mulf %189, %in_63 : f32
%191 = arith.addf %out, %190 : f32
linalg.yield %191 : f32
} -> tensor<2x10x100xf32>
return %179 : tensor<2x10x100xf32>
}
|
