[Mlir-commits] [mlir] [mlir][vector] Recompute load/store alignment when the index moves (PR #221385)

[Federico Bruzzone]

================
@@ -1284,21 +1284,36 @@ class ExtractOpFromLoad final : public OpRewritePattern<vector::ExtractOp> {
     rewriter.setInsertionPoint(loadOp);
     Location loc = loadOp.getLoc();
     ArithIndexingBuilder idxBuilderf(rewriter, loc);
+    SmallVector<int64_t> offsets(indices.size(), 0);
+    bool hasDynamicOffset = false;
     for (auto i : llvm::seq<int64_t>(rankOffset, indices.size() - finalRank)) {
       OpFoldResult pos = extractPos[i - rankOffset];
       if (isZeroInteger(pos))
         continue;
 
+      if (std::optional<int64_t> staticPos = getConstantIntValue(pos))
+        offsets[i] = *staticPos;
+      else
+        hasDynamicOffset = true;
       Value offset = getValueOrCreateConstantIndexOp(rewriter, loc, pos);
       indices[i] = idxBuilderf.add(indices[i], offset);
     }
 
+    // The new load is at an offset from the original one, so its alignment
+    // has to be recomputed, which needs a static offset.
+    llvm::MaybeAlign alignment;
+    if (!hasDynamicOffset)
+      alignment =
+          getAlignmentAfterOffset(loadOp.getMaybeAlign(), memType, offsets);
+    bool nontemporal = loadOp.getNontemporal();
----------------
FedericoBruzzone wrote:

[nit] I think we can fold this `nontemporal` local into the two use sites. Right?

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