[Mlir-commits] [mlir] [mlir][vector] Improve flattening vector.transfer_write ops. (PR #94051)

[Andrzej Warzyński]

================
@@ -510,20 +510,59 @@ static Value collapseInnerDims(PatternRewriter &rewriter, mlir::Location loc,
 /// the truncated indices where `firstDimToCollapse` is now the innermost dim.
 /// TODO: Extract the logic that writes to outIndices so that this method
 /// simply checks one pre-condition.
-static LogicalResult
-checkAndCollapseInnerZeroIndices(ValueRange indices, int64_t firstDimToCollapse,
-                                 SmallVector<Value> &outIndices) {
-  int64_t rank = indices.size();
-  if (firstDimToCollapse >= rank)
-    return failure();
-  for (int64_t i = firstDimToCollapse; i < rank; ++i) {
-    std::optional<int64_t> cst = getConstantIntValue(indices[i]);
-    if (!cst || cst.value() != 0)
-      return failure();
+static SmallVector<Value> getCollapsedIndices(RewriterBase &rewriter,
+                                              Location loc,
+                                              ArrayRef<int64_t> shape,
+                                              ValueRange indices,
+                                              int64_t firstDimToCollapse) {
+  assert(firstDimToCollapse < static_cast<int64_t>(indices.size()));
+
+  // If all the collapsed indices are zero then no extra logic is needed.
+  // Otherwise, a new offset/index has to be computed.
+  SmallVector<Value> collapsedIndices(indices.begin(),
+                                      indices.begin() + firstDimToCollapse);
----------------
banach-space wrote:

> How about renaming it to indicesAfterCollapsing?

👍🏻 

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