[Mlir-commits] [mlir] [mlir] Add Memref Normalization support for reinterpret_cast op (PR #133417)

[Uday Bondhugula]

================
@@ -353,16 +354,25 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(func::FuncOp funcOp,
   // Turn memrefs' non-identity layouts maps into ones with identity. Collect
   // alloc/alloca ops first and then process since normalizeMemRef
   // replaces/erases ops during memref rewriting.
-  SmallVector<memref::AllocOp, 4> allocOps;
-  funcOp.walk([&](memref::AllocOp op) { allocOps.push_back(op); });
-  for (memref::AllocOp allocOp : allocOps)
+  SmallVector<AllocOp, 4> allocOps;
+  funcOp.walk([&](AllocOp op) { allocOps.push_back(op); });
+  for (AllocOp allocOp : allocOps)
     (void)normalizeMemRef(allocOp);
 
-  SmallVector<memref::AllocaOp> allocaOps;
-  funcOp.walk([&](memref::AllocaOp op) { allocaOps.push_back(op); });
-  for (memref::AllocaOp allocaOp : allocaOps)
+  SmallVector<AllocaOp> allocaOps;
+  funcOp.walk([&](AllocaOp op) { allocaOps.push_back(op); });
+  for (AllocaOp allocaOp : allocaOps)
     (void)normalizeMemRef(allocaOp);
 
+  // Turn memrefs' non-identity layouts maps into ones with identity. Collect
+  // reinterpret_cast ops first and then process since normalizeMemRef
+  // replaces/erases ops during memref rewriting.
+  SmallVector<ReinterpretCastOp> reinterpretCastOps;
+  funcOp.walk(
+      [&](ReinterpretCastOp op) { reinterpretCastOps.push_back(op); });
----------------
bondhugula wrote:

Instead of walking the function thrice, can we walk once to collect all, or does the processing of one impact the other lists?

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