Skip to content

Commit 2de3d00

Browse files
authored
[mlir][tosa] Fix a crash in PadOp::fold (#114921)
This PR fixes a crash when padding of `tosa.pad` is not dense elements. Fixes #114762.
1 parent 8a1eba4 commit 2de3d00

File tree

2 files changed

+14
-2
lines changed

2 files changed

+14
-2
lines changed

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -878,8 +878,9 @@ OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
878878
OpFoldResult PadOp::fold(FoldAdaptor adaptor) {
879879
// If the pad is all zeros we can fold this operation away.
880880
if (adaptor.getPadding() && getInput1().getType() == getType()) {
881-
auto densePad = llvm::cast<DenseElementsAttr>(adaptor.getPadding());
882-
if (densePad.isSplat() && densePad.getSplatValue<APInt>().isZero()) {
881+
auto densePad = llvm::dyn_cast<DenseElementsAttr>(adaptor.getPadding());
882+
if (densePad && densePad.isSplat() &&
883+
densePad.getSplatValue<APInt>().isZero()) {
883884
return getInput1();
884885
}
885886
}

mlir/test/Dialect/Tosa/canonicalize.mlir

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,17 @@ func.func @pad_noop(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {
217217

218218
// -----
219219

220+
// CHECK-LABEL: @pad_noop_padding_mismatch_nofold
221+
func.func @pad_noop_padding_mismatch_nofold(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {
222+
// CHECK: %[[PAD:.+]] = tosa.pad
223+
// CHECK: return %[[PAD]]
224+
%0 = "tosa.const"() { value = dense_resource<__elided__> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
225+
%1 = tosa.pad %arg0, %0 : (tensor<?x?xf32>, tensor<2x2xi32>) -> tensor<?x?xf32>
226+
return %1 : tensor<?x?xf32>
227+
}
228+
229+
// -----
230+
220231
// CHECK-LABEL: @pad_noop_type_mismatch_nofold
221232
func.func @pad_noop_type_mismatch_nofold(%arg0: tensor<10xf32>) -> tensor<?xf32> {
222233
// CHECK: %[[PAD:.+]] = tosa.pad

0 commit comments

Comments
 (0)