Skip to content

Commit b6c7492

Browse files
author
Peiming Liu
committed
[mlir][sparse] support sparsifying 2:4 block sparsity
1 parent 30e4b09 commit b6c7492

File tree

5 files changed

+62
-26
lines changed

5 files changed

+62
-26
lines changed

‎mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -540,7 +540,8 @@ class Merger {
540540
bool isSparseLvlWithNonTrivialIdxExp(TensorLoopId b) const {
541541
if (isLvlWithNonTrivialIdxExp(b)) {
542542
auto dlt = getLoopDependentLevelType(b);
543-
return isCompressedDLT(dlt) || isSingletonDLT(dlt);
543+
return isCompressedDLT(dlt) || isSingletonDLT(dlt) ||
544+
isLooseCompressedDLT(dlt) || is2OutOf4DLT(dlt);
544545
}
545546
return false;
546547
}

‎mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -448,7 +448,7 @@ void LoopEmitter::initializeLoopEmit(
448448
positionsBuffers[t][l] = genToPositions(builder, loc, tensor, l);
449449
coordinatesBuffers[t][l] =
450450
genToCoordinates(builder, loc, tensor, l, cooStart);
451-
} else if (isSingletonDLT(lvlTp)) {
451+
} else if (isSingletonDLT(lvlTp) || is2OutOf4DLT(lvlTp)) {
452452
// Singleton level, fetch coordinates.
453453
coordinatesBuffers[t][l] =
454454
genToCoordinates(builder, loc, tensor, l, cooStart);
@@ -540,7 +540,8 @@ void LoopEmitter::categorizeLoopCondition(
540540
auto lvlType = lvlTypes[t][l];
541541
// Must be a recognizable DLT.
542542
assert(isDenseDLT(lvlType) || isCompressedDLT(lvlType) ||
543-
isLooseCompressedDLT(lvlType) || isSingletonDLT(lvlType));
543+
isLooseCompressedDLT(lvlType) || isSingletonDLT(lvlType) ||
544+
is2OutOf4DLT(lvlType));
544545

545546
bool isSparse = !isDenseDLT(lvlType);
546547
bool isSlice = isSparseSlices[t];
@@ -637,6 +638,7 @@ std::pair<Operation *, Value> LoopEmitter::emitForLoopOverTensorAtLvl(
637638
Value hi, MutableArrayRef<Value> reduc, bool isParallel) {
638639
bool isSparseCond = isCompressedDLT(lvlTypes[tid][lvl]) ||
639640
isLooseCompressedDLT(lvlTypes[tid][lvl]) ||
641+
is2OutOf4DLT(lvlTypes[tid][lvl]) ||
640642
isSingletonDLT(lvlTypes[tid][lvl]);
641643
// TODO: support dynamic slices.
642644
// Uses the first dimension here to build the loop bound (which is also the
@@ -1240,6 +1242,7 @@ void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
12401242

12411243
const Value c0 = C_IDX(0);
12421244
const Value c1 = C_IDX(1);
1245+
const Value c2 = C_IDX(2);
12431246
// Either the first level, or the previous level has been set.
12441247
/// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
12451248
assert(lvl == 0 || posits[tid][lvl - 1]);
@@ -1248,7 +1251,7 @@ void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
12481251

12491252
Value pLo = lvl == 0 ? c0 : posits[tid][lvl - 1];
12501253
if (isLooseCompressedDLT(lvlTp))
1251-
pLo = builder.create<arith::MulIOp>(loc, pLo, C_IDX(2));
1254+
pLo = builder.create<arith::MulIOp>(loc, pLo, c2);
12521255
posits[tid][lvl] = genIndexLoad(builder, loc, mem, pLo);
12531256

12541257
const Value pHi = ADDI(pLo, c1);
@@ -1271,7 +1274,13 @@ void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
12711274
: ADDI(pLo, c1);
12721275
return;
12731276
}
1274-
1277+
if (is2OutOf4DLT(lvlTp)) {
1278+
const Value pLo = lvl == 0 ? c0 : posits[tid][lvl - 1];
1279+
// Each 2:4 block has exactly two specified elements.
1280+
posits[tid][lvl] = MULI(pLo, c2);
1281+
highs[tid][lvl] = ADDI(posits[tid][lvl], c2);
1282+
return;
1283+
}
12751284
llvm_unreachable("Unrecognized level-type!");
12761285
}
12771286

‎mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -816,7 +816,7 @@ static bool computeIterationGraph(CodegenEnv &env, SortMask mask,
816816
for (LoopId i = 0; i < numLoops; i++) {
817817
const auto dltI = env.dlt(tid, i);
818818
if (isCompressedDLT(dltI) || isLooseCompressedDLT(dltI) ||
819-
isSingletonDLT(dltI)) {
819+
isSingletonDLT(dltI) || is2OutOf4DLT(dltI)) {
820820
for (LoopId j = 0; j < numLoops; j++)
821821
if (isUndefDLT(env.dlt(tid, j))) {
822822
addIterOrdering(i, j, adjM, inDegree);
@@ -1508,7 +1508,7 @@ static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId ldx,
15081508
assert(ldx == env.merger().loop(b));
15091509
Value clause;
15101510
if (isCompressedDLT(dlt) || isSingletonDLT(dlt) ||
1511-
isLooseCompressedDLT(dlt)) {
1511+
isLooseCompressedDLT(dlt) || is2OutOf4DLT(dlt)) {
15121512
assert(lvl.has_value());
15131513
const Value crd = env.emitter().getCoords()[tid][*lvl];
15141514
const Value lvar = env.getLoopVar(ldx);
@@ -1593,7 +1593,7 @@ static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
15931593
needsUniv = true;
15941594
}
15951595
if (isCompressedDLT(dlt) || isSingletonDLT(dlt) ||
1596-
isLooseCompressedDLT(dlt) || isIdxReduc) {
1596+
isLooseCompressedDLT(dlt) || is2OutOf4DLT(dlt) || isIdxReduc) {
15971597
// Only when this is a index reduction loop, can the dlt be undefined.
15981598
assert(!isUndefDLT(dlt) || isIdxReduc);
15991599
// sparse/singleton levels, or a dense/sparse index reduction loop.

‎mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -490,7 +490,7 @@ BitVector Merger::simplifyCond(LatSetId s0, LatPointId p0) {
490490
if (simple[b] && !isSparseLvlWithNonTrivialIdxExp(b)) {
491491
const auto dlt = getLvlType(b);
492492
if (!isCompressedDLT(dlt) && !isSingletonDLT(dlt) &&
493-
!isLooseCompressedDLT(dlt)) {
493+
!isLooseCompressedDLT(dlt) && !is2OutOf4DLT(dlt)) {
494494
if (reset)
495495
simple.reset(b);
496496
reset = true;
@@ -671,7 +671,7 @@ bool Merger::hasAnySparse(const BitVector &bits) const {
671671
for (TensorLoopId b : bits.set_bits()) {
672672
const auto dlt = getLvlType(b);
673673
if (isCompressedDLT(dlt) || isSingletonDLT(dlt) ||
674-
isLooseCompressedDLT(dlt))
674+
isLooseCompressedDLT(dlt) || is2OutOf4DLT(dlt))
675675
return true;
676676
}
677677
return hasSparseIdxReduction(bits);

‎mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_block_matmul.mlir

Lines changed: 42 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -47,19 +47,41 @@
4747
#BSR = #sparse_tensor.encoding<{
4848
map = ( i, j ) ->
4949
( i floordiv 2 : dense,
50-
j floordiv 3 : compressed,
50+
j floordiv 2 : compressed,
5151
i mod 2 : dense,
52-
j mod 3 : dense
52+
j mod 2 : dense
5353
)
5454
}>
5555

56+
#NV_24 = #sparse_tensor.encoding<{
57+
map = ( i, j ) ->
58+
( i : dense,
59+
j floordiv 4 : dense,
60+
j mod 4 : block2_4
61+
),
62+
}>
63+
5664
module {
5765

58-
func.func @mul(%arg0: tensor<4x6xf64>,
59-
%arg1: tensor<4x6xf64, #BSR>) -> tensor<4x4xf64> {
66+
func.func @mul(%arg0: tensor<4x8xf64>,
67+
%arg1: tensor<4x8xf64, #BSR>) -> tensor<4x4xf64> {
68+
%out = arith.constant dense<0.0> : tensor<4x4xf64>
69+
%0 = linalg.generic #trait_mul
70+
ins(%arg0, %arg1: tensor<4x8xf64>, tensor<4x8xf64, #BSR>)
71+
outs(%out: tensor<4x4xf64>) {
72+
^bb(%x: f64, %y : f64, %z : f64):
73+
%1 = arith.mulf %x, %y : f64
74+
%2 = arith.addf %1, %z : f64
75+
linalg.yield %2 : f64
76+
} -> tensor<4x4xf64>
77+
return %0 : tensor<4x4xf64>
78+
}
79+
80+
func.func @mul_24(%arg0: tensor<4x8xf64>,
81+
%arg1: tensor<4x8xf64, #NV_24>) -> tensor<4x4xf64> {
6082
%out = arith.constant dense<0.0> : tensor<4x4xf64>
6183
%0 = linalg.generic #trait_mul
62-
ins(%arg0, %arg1: tensor<4x6xf64>, tensor<4x6xf64, #BSR>)
84+
ins(%arg0, %arg1: tensor<4x8xf64>, tensor<4x8xf64, #NV_24>)
6385
outs(%out: tensor<4x4xf64>) {
6486
^bb(%x: f64, %y : f64, %z : f64):
6587
%1 = arith.mulf %x, %y : f64
@@ -69,11 +91,11 @@ func.func @mul(%arg0: tensor<4x6xf64>,
6991
return %0 : tensor<4x4xf64>
7092
}
7193

72-
func.func @mul_dense(%arg0: tensor<4x6xf64>,
73-
%arg1: tensor<4x6xf64>) -> tensor<4x4xf64> {
94+
func.func @mul_dense(%arg0: tensor<4x8xf64>,
95+
%arg1: tensor<4x8xf64>) -> tensor<4x4xf64> {
7496
%out = arith.constant dense<0.0> : tensor<4x4xf64>
7597
%0 = linalg.generic #trait_mul
76-
ins(%arg0, %arg1: tensor<4x6xf64>, tensor<4x6xf64>)
98+
ins(%arg0, %arg1: tensor<4x8xf64>, tensor<4x8xf64>)
7799
outs(%out: tensor<4x4xf64>) {
78100
^bb(%x: f64, %y : f64, %z : f64):
79101
%1 = arith.mulf %x, %y : f64
@@ -104,22 +126,26 @@ func.func @mul_dense(%arg0: tensor<4x6xf64>,
104126
%c2 = arith.constant 2 : index
105127

106128

107-
%td = arith.constant dense<[[ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
108-
[ 6.0, 7.0, 8.0, 9.0, 10.0, 11.0],
109-
[12.0, 13.0, 14.0, 15.0, 16.0, 17.0],
110-
[18.0, 19.0, 20.0, 21.0, 22.0, 23.0]]> : tensor<4x6xf64>
129+
%td = arith.constant dense<[[ 1.0, 2.0, 0.0, 0.0, 0.0, 0.0, 4.0, 5.0],
130+
[ 6.0, 7.0, 0.0, 0.0, 0.0, 0.0, 10.0, 11.0],
131+
[ 0.0, 0.0, 12.0, 13.0, 16.0, 17.0, 0.0, 0.0],
132+
[ 0.0, 0.0, 18.0, 19.0, 22.0, 23.0, 0.0, 0.0]]> : tensor<4x8xf64>
111133

112134

113-
%2 = sparse_tensor.convert %td : tensor<4x6xf64> to tensor<4x6xf64, #BSR>
135+
%2 = sparse_tensor.convert %td : tensor<4x8xf64> to tensor<4x8xf64, #BSR>
136+
%3 = sparse_tensor.convert %td : tensor<4x8xf64> to tensor<4x8xf64, #NV_24>
114137

115138
%d = call @mul_dense(%td, %td)
116-
: (tensor<4x6xf64>, tensor<4x6xf64>) -> tensor<4x4xf64>
139+
: (tensor<4x8xf64>, tensor<4x8xf64>) -> tensor<4x4xf64>
117140
%s = call @mul(%td, %2)
118-
: (tensor<4x6xf64>, tensor<4x6xf64, #BSR>) -> tensor<4x4xf64>
141+
: (tensor<4x8xf64>, tensor<4x8xf64, #BSR>) -> tensor<4x4xf64>
142+
%s24 = call @mul_24(%td, %3)
143+
: (tensor<4x8xf64>, tensor<4x8xf64, #NV_24>) -> tensor<4x4xf64>
119144

120-
// CHECK-COUNT-2: ( ( 55, 145, 235, 325 ), ( 145, 451, 757, 1063 ), ( 235, 757, 1279, 1801 ), ( 325, 1063, 1801, 2539 ) )
145+
// CHECK-COUNT-3: ( ( 46, 115, 0, 0 ), ( 115, 306, 0, 0 ), ( 0, 0, 858, 1206 ), ( 0, 0, 1206, 1698 ) )
121146
call @dumpf64(%d) : (tensor<4x4xf64>) -> ()
122147
call @dumpf64(%s) : (tensor<4x4xf64>) -> ()
148+
call @dumpf64(%s24) : (tensor<4x4xf64>) -> ()
123149

124150
return
125151
}

0 commit comments

Comments
 (0)