Skip to content

Commit 96e1914

Browse files
author
Peiming Liu
committed
[mlir][sparse] fix crash when generating convolution kernel with sparse input in DCCD format.
Reviewed By: aartbik, anlunx Differential Revision: https://reviews.llvm.org/D159170
1 parent 26f230f commit 96e1914

File tree

3 files changed

+51
-4
lines changed

3 files changed

+51
-4
lines changed

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1800,20 +1800,20 @@ void LoopEmitter::genResolvedSliceBegin(OpBuilder &builder, Location loc,
18001800
pHi = genIndexLoad(builder, loc, positionsBuffers[tid][lvl],
18011801
ADDI(posits[tid][lvl - 1], c1));
18021802
}
1803-
// Fills out pIdxBuffer[tid][lvl][0] with [/*memSize =*/4, 0, 0, pHi]
1803+
// Fills out pIdxBuffer[tid][lvl][0] with [/*memSize =*/4, 0, pLo, pHi]
18041804
builder.create<memref::StoreOp>(loc, c4, sPtrBuf, c0); // memSize = 4
18051805
builder.create<memref::StoreOp>(loc, c0, sPtrBuf, c1); // index = 0
18061806
builder.create<memref::StoreOp>(loc, pLo, sPtrBuf, c2); // pLo
18071807
builder.create<memref::StoreOp>(loc, pHi, sPtrBuf, c3); // pHi
18081808

1809-
// This is an non empty tensor if 0 < pHi.
1810-
Value isNonEmpty = CMPI(ult, c0, pHi);
1809+
// This is an non empty tensor if pLo < pHi.
1810+
Value isNonEmpty = CMPI(ult, pLo, pHi);
18111811
// The minimal coord must be at the first on ordered level.
18121812
// FIXME: Technically we should load the coord only when the slice is
18131813
// nonempty. though we assume that even on empty sparse tensors, a non-empty
18141814
// ptr/idx buffer is allocated for each level so it would not cause OOB to
18151815
// avoid generating a ifOp here.
1816-
Value minCrd = genIndexLoad(builder, loc, coordinatesBuffers[tid][0], c0);
1816+
Value minCrd = genIndexLoad(builder, loc, coordinatesBuffers[tid][lvl], pLo);
18171817

18181818
// FIXME: We need the relative offset related to the base slice.
18191819
Value absOffset = offsetFromMinCoord(builder, loc, minCrd, size, isNonEmpty);

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
// DEFINE: %{env} =
1818
//--------------------------------------------------------------------------------------------------
1919

20+
// UNSUPPORTED: target={{.*}}
21+
2022
// RUN: %{compile} | %{env} %{run} | FileCheck %s
2123
//
2224
// Do the same run, but now with direct IR generation.

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,10 @@
3838
lvlTypes = [ "compressed", "dense", "compressed", "dense" ]
3939
}>
4040

41+
#DCCD = #sparse_tensor.encoding<{
42+
lvlTypes = [ "dense", "compressed", "compressed", "dense" ]
43+
}>
44+
4145
// Creates and returns 4-D buffer of size (%s1, %s2, %s3, %s4) filled with the value %f
4246
func.func @alloc_4d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %f : f32) -> tensor<?x?x?x?xf32> {
4347
%buf = bufferization.alloc_tensor(%s1, %s2, %s3, %s4) : tensor<?x?x?x?xf32>
@@ -77,6 +81,18 @@ func.func @conv_2d_nhwc_hwcf_CDCD(%arg0: tensor<?x?x?x?xf32, #CDCD>, %arg1: tens
7781
return %ret : tensor<?x?x?x?xf32, #CDCD>
7882
}
7983

84+
func.func @conv_2d_nhwc_hwcf_DCCD(%arg0: tensor<?x?x?x?xf32, #DCCD>, %arg1: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32, #DCCD> {
85+
%c1 = arith.constant 1 : index
86+
%c3 = arith.constant 3 : index
87+
%c6 = arith.constant 6 : index
88+
%s = bufferization.alloc_tensor(%c3, %c6, %c6, %c1) : tensor<?x?x?x?xf32, #DCCD>
89+
%ret = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
90+
strides = dense<1> : tensor<2xi64>}
91+
ins (%arg0, %arg1: tensor<?x?x?x?xf32, #DCCD>, tensor<?x?x?x?xf32>)
92+
outs (%s: tensor<?x?x?x?xf32, #DCCD>) -> tensor<?x?x?x?xf32, #DCCD>
93+
return %ret : tensor<?x?x?x?xf32, #DCCD>
94+
}
95+
8096
func.func @entry() {
8197
%c0 = arith.constant 0 : index
8298
%c1 = arith.constant 1 : index
@@ -96,10 +112,13 @@ func.func @entry() {
96112
: tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32, #CCCC>
97113
%in2D_nhwc_CDCD = sparse_tensor.convert %in2D_nhwc
98114
: tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32, #CDCD>
115+
%in2D_nhwc_DCCD = sparse_tensor.convert %in2D_nhwc
116+
: tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32, #DCCD>
99117

100118
%dense_ret = call @conv_2d_nhwc_hwcf(%in2D_nhwc, %filter2D_nhwc, %out2D_nhwc) : (tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) -> (tensor<?x?x?x?xf32>)
101119
%CCCC_ret = call @conv_2d_nhwc_hwcf_CCCC(%in2D_nhwc_CCCC, %filter2D_nhwc) : (tensor<?x?x?x?xf32, #CCCC>, tensor<?x?x?x?xf32>) -> (tensor<?x?x?x?xf32, #CCCC>)
102120
%CDCD_ret = call @conv_2d_nhwc_hwcf_CDCD(%in2D_nhwc_CDCD, %filter2D_nhwc) : (tensor<?x?x?x?xf32, #CDCD>, tensor<?x?x?x?xf32>) -> (tensor<?x?x?x?xf32, #CDCD>)
121+
%DCCD_ret = call @conv_2d_nhwc_hwcf_DCCD(%in2D_nhwc_DCCD, %filter2D_nhwc) : (tensor<?x?x?x?xf32, #DCCD>, tensor<?x?x?x?xf32>) -> (tensor<?x?x?x?xf32, #DCCD>)
103122

104123
// CHECK: ( ( ( ( 108 ), ( 124 ), ( 124 ), ( 124 ), ( 108 ), ( 108 ) ),
105124
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
@@ -171,16 +190,42 @@ func.func @entry() {
171190
: tensor<?x?x?x?xf32>, vector<3x6x6x1xf32>
172191
vector.print %v2 : vector<3x6x6x1xf32>
173192

193+
// CHECK: ( ( ( ( 108 ), ( 124 ), ( 124 ), ( 124 ), ( 108 ), ( 108 ) ),
194+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
195+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
196+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
197+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
198+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ) ),
199+
// CHECK-SAME: ( ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
200+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
201+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
202+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
203+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
204+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ) ),
205+
// CHECK-SAME: ( ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
206+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
207+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
208+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
209+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ),
210+
// CHECK-SAME: ( ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ), ( 108 ) ) ) )
211+
%3 = sparse_tensor.convert %DCCD_ret
212+
: tensor<?x?x?x?xf32, #DCCD> to tensor<?x?x?x?xf32>
213+
%v3 = vector.transfer_read %3[%c0, %c0, %c0, %c0], %zero
214+
: tensor<?x?x?x?xf32>, vector<3x6x6x1xf32>
215+
vector.print %v3 : vector<3x6x6x1xf32>
216+
174217
// Free the resources
175218
bufferization.dealloc_tensor %in2D_nhwc : tensor<?x?x?x?xf32>
176219
bufferization.dealloc_tensor %filter2D_nhwc : tensor<?x?x?x?xf32>
177220
bufferization.dealloc_tensor %out2D_nhwc : tensor<?x?x?x?xf32>
178221

179222
bufferization.dealloc_tensor %in2D_nhwc_CDCD : tensor<?x?x?x?xf32, #CDCD>
180223
bufferization.dealloc_tensor %in2D_nhwc_CCCC : tensor<?x?x?x?xf32, #CCCC>
224+
bufferization.dealloc_tensor %in2D_nhwc_DCCD : tensor<?x?x?x?xf32, #DCCD>
181225

182226
bufferization.dealloc_tensor %CCCC_ret : tensor<?x?x?x?xf32, #CCCC>
183227
bufferization.dealloc_tensor %CDCD_ret : tensor<?x?x?x?xf32, #CDCD>
228+
bufferization.dealloc_tensor %DCCD_ret : tensor<?x?x?x?xf32, #DCCD>
184229

185230
return
186231
}

0 commit comments

Comments
 (0)