From 25dbf8f3c66b625973f12ba81b99f8e797ab4e76 Mon Sep 17 00:00:00 2001
From: Alexander Belyaev <pifon@google.com>
Date: Tue, 21 Feb 2023 09:16:14 +0100
Subject: [PATCH] [mlir] Use foldDynamicIndexList instead of
 canonicalizeSubViewPart.

Differential Revision: https://reviews.llvm.org/D144452
---
 mlir/include/mlir/Dialect/Arith/Utils/Utils.h | 23 ++++++-----------------
 mlir/lib/Dialect/Arith/Utils/Utils.cpp        | 15 ---------------
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp      | 18 ++++++------------
 3 files changed, 12 insertions(+), 44 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Arith/Utils/Utils.h b/mlir/include/mlir/Dialect/Arith/Utils/Utils.h
index 4bed6e5..bfb4702 100644
--- a/mlir/include/mlir/Dialect/Arith/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Arith/Utils/Utils.h
@@ -26,12 +26,6 @@ namespace mlir {
 /// Matches a ConstantIndexOp.
 detail::op_matcher<arith::ConstantIndexOp> matchConstantIndex();
 
-/// Detects the `values` produced by a ConstantIndexOp and places the new
-/// constant in place of the corresponding sentinel value.
-/// TODO(pifon2a): Remove this function and use foldDynamicIndexList.
-void canonicalizeSubViewPart(SmallVectorImpl<OpFoldResult> &values,
-                             function_ref<bool(int64_t)> isDynamic);
-
 /// Returns `success` when any of the elements in `ofrs` was produced by
 /// arith::ConstantIndexOp. In that case the constant attribute replaces the
 /// Value. Returns `failure` when no folding happened.
@@ -50,20 +44,15 @@ public:
 
   LogicalResult matchAndRewrite(OpType op,
                                 PatternRewriter &rewriter) const override {
-    // No constant operand, just return;
-    if (llvm::none_of(op.getOperands(), [](Value operand) {
-          return matchPattern(operand, matchConstantIndex());
-        }))
-      return failure();
-
-    // At least one of offsets/sizes/strides is a new constant.
-    // Form the new list of operands and constant attributes from the existing.
     SmallVector<OpFoldResult> mixedOffsets(op.getMixedOffsets());
     SmallVector<OpFoldResult> mixedSizes(op.getMixedSizes());
     SmallVector<OpFoldResult> mixedStrides(op.getMixedStrides());
-    canonicalizeSubViewPart(mixedOffsets, ShapedType::isDynamic);
-    canonicalizeSubViewPart(mixedSizes, ShapedType::isDynamic);
-    canonicalizeSubViewPart(mixedStrides, ShapedType::isDynamic);
+
+    // No constant operands were folded, just return;
+    if (failed(foldDynamicIndexList(rewriter, mixedOffsets)) &&
+        failed(foldDynamicIndexList(rewriter, mixedSizes)) &&
+        failed(foldDynamicIndexList(rewriter, mixedStrides)))
+      return failure();
 
     // Create the new op in canonical form.
     ResultTypeFunc resultTypeFunc;
diff --git a/mlir/lib/Dialect/Arith/Utils/Utils.cpp b/mlir/lib/Dialect/Arith/Utils/Utils.cpp
index 4d8b5ad..45a4bf7 100644
--- a/mlir/lib/Dialect/Arith/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Arith/Utils/Utils.cpp
@@ -23,21 +23,6 @@ detail::op_matcher<arith::ConstantIndexOp> mlir::matchConstantIndex() {
   return detail::op_matcher<arith::ConstantIndexOp>();
 }
 
-// Detects the `values` produced by a ConstantIndexOp and places the new
-// constant in place of the corresponding sentinel value.
-void mlir::canonicalizeSubViewPart(
-    SmallVectorImpl<OpFoldResult> &values,
-    llvm::function_ref<bool(int64_t)> isDynamic) {
-  for (OpFoldResult &ofr : values) {
-    if (ofr.is<Attribute>())
-      continue;
-    // Newly static, move from Value to constant.
-    if (auto cstOp =
-            ofr.dyn_cast<Value>().getDefiningOp<arith::ConstantIndexOp>())
-      ofr = OpBuilder(cstOp).getIndexAttr(cstOp.value());
-  }
-}
-
 // Returns `success` when any of the elements in `ofrs` was produced by
 // arith::ConstantIndexOp. In that case the constant attribute replaces the
 // Value. Returns `failure` when no folding happened.
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index d567949..b6359a9 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2227,21 +2227,15 @@ public:
 
   LogicalResult matchAndRewrite(InsertOpTy insertSliceOp,
                                 PatternRewriter &rewriter) const override {
-    // No constant operand, just return.
-    if (llvm::none_of(insertSliceOp.getOperands(), [](Value operand) {
-          return matchPattern(operand, matchConstantIndex());
-        }))
-      return failure();
-
-    // At least one of offsets/sizes/strides is a new constant.
-    // Form the new list of operands and constant attributes from the
-    // existing.
     SmallVector<OpFoldResult> mixedOffsets(insertSliceOp.getMixedOffsets());
     SmallVector<OpFoldResult> mixedSizes(insertSliceOp.getMixedSizes());
     SmallVector<OpFoldResult> mixedStrides(insertSliceOp.getMixedStrides());
-    canonicalizeSubViewPart(mixedOffsets, ShapedType::isDynamic);
-    canonicalizeSubViewPart(mixedSizes, ShapedType::isDynamic);
-    canonicalizeSubViewPart(mixedStrides, ShapedType::isDynamic);
+
+    // No constant operands were folded, just return;
+    if (failed(foldDynamicIndexList(rewriter, mixedOffsets)) &&
+        failed(foldDynamicIndexList(rewriter, mixedSizes)) &&
+        failed(foldDynamicIndexList(rewriter, mixedStrides)))
+      return failure();
 
     // Create the new op in canonical form.
     auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType(
-- 
2.7.4