1 files changed, 5 insertions, 5 deletions
diff --git a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
index 51e9debd0ad..4fc544074c0 100644
--- a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
+++ b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
@@ -212,13 +212,13 @@ void MemRefDataFlowOpt::forwardStoreToLoad(OpPointer<LoadOp> loadOp) {
 void MemRefDataFlowOpt::runOnFunction() {
   // Only supports single block functions at the moment.
   Function &f = getFunction();
-  if (f.getBlocks().size() != 1)
+  if (f.getBlocks().size() != 1) {
+    markAllAnalysesPreserved();
     return;
+  }
 
-  DominanceInfo theDomInfo(&f);
-  domInfo = &theDomInfo;
-  PostDominanceInfo thePostDomInfo(&f);
-  postDomInfo = &thePostDomInfo;
+  domInfo = &getAnalysisResult<DominanceInfo>();
+  postDomInfo = &getAnalysisResult<PostDominanceInfo>();
 
   loadOpsToErase.clear();
   memrefsToErase.clear();