[Mlir-commits] [mlir] 58c1885 - [mlir][linalg] Fix `FoldInitTensorWithDimOp` if dim(init_tensor) is static.

[Alexander Belyaev]

Author: Alexander Belyaev
Date: 2021-09-06T10:47:26+02:00
New Revision: 58c188507f724feee181018bfb6a330be2c6ac9a

URL: https://github.com/llvm/llvm-project/commit/58c188507f724feee181018bfb6a330be2c6ac9a
DIFF: https://github.com/llvm/llvm-project/commit/58c188507f724feee181018bfb6a330be2c6ac9a.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index abb3c44158381..e2b9436f4a10e 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -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();
   }
 };