maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][linalg] Fix `FoldInitTensorWithDimOp` if dim(init_tensor) is static. 

It looks like it was a typo. Instead of `*maybeConstantIndex`,
`initTensorOp.getStaticSize(*maybeConstantIndex)` should be used to access the
dim size of the tensor. There is a test for that in `canonicalize.mlir`, but it
was working correctly because `ReplaceStaticShapeDims` was canonicalizing DimOp
before `FoldInitTensorWithDimOp`. So, to make the patterns more "orthogonal",
this case is disabled.

Differential Revision: https://reviews.llvm.org/D109247
This commit is contained in:
Alexander Belyaev 2021-09-03 19:06:15 +02:00
parent 915a8bb5ee
commit 58c188507f
1 changed files with 3 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
View File

@ -977,12 +977,9 @@ struct FoldInitTensorWithDimOp : public OpRewritePattern<tensor::DimOp> {
auto initTensorOp = dimOp.source().getDefiningOp<linalg::InitTensorOp>();
if (!initTensorOp || !maybeConstantIndex)
return failure();
if (initTensorOp.isDynamicSize(*maybeConstantIndex)) {
rewriter.replaceOp(dimOp,
initTensorOp.getDynamicSize(*maybeConstantIndex));
return success();
}
rewriter.replaceOpWithNewOp<ConstantIndexOp>(dimOp, *maybeConstantIndex);
if (!initTensorOp.isDynamicSize(*maybeConstantIndex))
return failure();
rewriter.replaceOp(dimOp, initTensorOp.getDynamicSize(*maybeConstantIndex));
return success();
}
};