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, 6 insertions, 2 deletions

diff --git a/mlir/lib/Transforms/MaterializeVectors.cpp b/mlir/lib/Transforms/MaterializeVectors.cpp
index 1ab1f6361d3..e95bb7307e3 100644
--- a/mlir/lib/Transforms/MaterializeVectors.cpp
+++ b/mlir/lib/Transforms/MaterializeVectors.cpp
@@ -197,7 +197,7 @@ struct MaterializationState {
 struct MaterializeVectorsPass : public FunctionPass {
   MaterializeVectorsPass() : FunctionPass(&MaterializeVectorsPass::passID) {}
 
-  PassResult runOnMLFunction(Function *f) override;
+  PassResult runOnFunction(Function *f) override;
 
   // Thread-safe RAII contexts local to pass, BumpPtrAllocator freed on exit.
   MLFunctionMatcherContext mlContext;
@@ -712,7 +712,11 @@ static bool materialize(Function *f,
   return false;
 }
 
-PassResult MaterializeVectorsPass::runOnMLFunction(Function *f) {
+PassResult MaterializeVectorsPass::runOnFunction(Function *f) {
+  // TODO(ntv): Check to see if this supports arbitrary top-level code.
+  if (f->getBlocks().size() != 1)
+    return success();
+
   using matcher::Op;
   LLVM_DEBUG(dbgs() << "\nMaterializeVectors on Function\n");
   LLVM_DEBUG(f->print(dbgs()));