Update and generalize various passes to work on both CFG and ML functions,

simplifying them in minor ways.  The only significant cleanup here
is the constant folding pass.  All the other changes are simple and easy,
but this is still enough to shrink the compiler by 45LOC.

The one pass left to merge is the CSE pass, which will be move involved, so I'm
splitting it out to its own patch (which I'll tackle right after this).

This is step 28/n towards merging instructions and statements.

PiperOrigin-RevId: 227328115

1 files changed, 7 insertions, 5 deletions

diff --git a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
index 1a30e2b289d..49b33b0596b 100644
--- a/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
+++ b/mlir/lib/Transforms/MemRefDataFlowOpt.cpp
@@ -66,9 +66,7 @@ namespace {
 struct MemRefDataFlowOpt : public FunctionPass, InstWalker<MemRefDataFlowOpt> {
   explicit MemRefDataFlowOpt() : FunctionPass(&MemRefDataFlowOpt::passID) {}
 
-  // Not applicable to CFG functions.
-  PassResult runOnCFGFunction(Function *f) override { return success(); }
-  PassResult runOnMLFunction(Function *f) override;
+  PassResult runOnFunction(Function *f) override;
 
   void visitOperationInst(OperationInst *opInst);
 
@@ -210,7 +208,11 @@ void MemRefDataFlowOpt::visitOperationInst(OperationInst *opInst) {
   loadOpsToErase.push_back(loadOpInst);
 }
 
-PassResult MemRefDataFlowOpt::runOnMLFunction(Function *f) {
+PassResult MemRefDataFlowOpt::runOnFunction(Function *f) {
+  // Only supports single block functions at the moment.
+  if (f->getBlocks().size() != 1)
+    return success();
+
   DominanceInfo theDomInfo(f);
   domInfo = &theDomInfo;
   PostDominanceInfo thePostDomInfo(f);
@@ -233,7 +235,7 @@ PassResult MemRefDataFlowOpt::runOnMLFunction(Function *f) {
   for (auto *memref : memrefsToErase) {
     // If the memref hasn't been alloc'ed in this function, skip.
     OperationInst *defInst = memref->getDefiningInst();
-    if (!defInst || !cast<OperationInst>(defInst)->isa<AllocOp>())
+    if (!defInst || !defInst->isa<AllocOp>())
       // TODO(mlir-team): if the memref was returned by a 'call' instruction, we
       // could still erase it if the call has no side-effects.
       continue;