[Mlir-commits] [mlir] Folding extract_strided_metadata input into reinterpret_cast on constant layout (PR #134845)

[Matthias Springer]

================
@@ -2045,44 +2050,57 @@ struct ReinterpretCastOpExtractStridedMetadataFolder
         op.getSource().getDefiningOp<ExtractStridedMetadataOp>();
     if (!extractStridedMetadata)
       return failure();
+
     // Check if the reinterpret cast reconstructs a memref with the exact same
     // properties as the extract strided metadata.
-
-    // First, check that the strides are the same.
     SmallVector<OpFoldResult> extractStridesOfr =
         extractStridedMetadata.getConstifiedMixedStrides();
     SmallVector<OpFoldResult> reinterpretStridesOfr =
         op.getConstifiedMixedStrides();
-    if (extractStridesOfr.size() != reinterpretStridesOfr.size())
-      return failure();
+    auto isReinterpretCastNoop = [&]() -> bool {
+      // First, check that the strides are the same.
+      if (!llvm::equal(extractStridesOfr, reinterpretStridesOfr))
+        return false;
 
-    unsigned rank = op.getType().getRank();
-    for (unsigned i = 0; i < rank; ++i) {
-      if (extractStridesOfr[i] != reinterpretStridesOfr[i])
-        return failure();
-    }
+      // Second, check the sizes.
+      SmallVector<OpFoldResult> extractSizesOfr =
+          extractStridedMetadata.getConstifiedMixedSizes();
+      SmallVector<OpFoldResult> reinterpretSizesOfr =
+          op.getConstifiedMixedSizes();
+      if (!llvm::equal(extractSizesOfr, reinterpretSizesOfr))
----------------
matthias-springer wrote:

nit: I would get rid of all the variables here and directly use `op.getConstifiedMixedSizes()` inside of `llvm::equal`, etc.

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