[Mlir-commits] [mlir] [mlir][vector]add foldConstantOp fold function and apply it to extractOp and insertOp. (PR #124399)

[Kunwar Grover]

================
@@ -1977,6 +1977,46 @@ static Value foldScalarExtractFromFromElements(ExtractOp extractOp) {
   return fromElementsOp.getElements()[flatIndex];
 }
 
+// If the dynamic operands of `extractOp` or `insertOp` is result of
+// `constantOp`, then fold it.
+template <typename T>
+static void foldConstantOp(T op, SmallVectorImpl<Value> &operands) {
+  auto staticPosition = op.getStaticPosition().vec();
+  OperandRange dynamicPosition = op.getDynamicPosition();
+
+  // If the dynamic operands is empty, it is returned directly.
+  if (!dynamicPosition.size())
+    return;
+  unsigned index = 0;
+
+  // `opChange` is a flog. If it is true, it means to update `op` in place.
+  bool opChange = false;
+  for (unsigned i = 0, e = staticPosition.size(); i < e; ++i) {
+    if (!ShapedType::isDynamic(staticPosition[i]))
+      continue;
+    Value position = dynamicPosition[index++];
+
+    // If it is a block parameter, proceed to the next iteration.
+    if (!position.getDefiningOp()) {
+      operands.push_back(position);
+      continue;
+    }
+
+    if (auto constantOp =
+            mlir::dyn_cast<arith::ConstantIndexOp>(position.getDefiningOp())) {
+      opChange = true;
+      staticPosition[i] = constantOp.value();
+      continue;
+    }
+    operands.push_back(position);
+  }
+
+  if (opChange) {
+    op.setStaticPosition(staticPosition);
+    op.getOperation()->setOperands(operands);
+  }
----------------
Groverkss wrote:

Can we actually do this? My understanding is that if there is a static position set, there should be no operand passed. This is creating an inconsistent state as can be seen in the test.

I don't think you can represent this "folding" with a fold pattern. Fold pattern cannot change the operation. Removing a operand changes the operation storage (no more storage for the operand), so you have to make this a canonicalization, instead of a fold. https://mlir.llvm.org/docs/Canonicalization/#canonicalizing-with-the-fold-method

https://github.com/llvm/llvm-project/pull/124399