[mlir] Add maximum SSA form transformation pass

[mortbopet]

Summary

This pass converts a program into maximal SSA form. In this form, any value referenced within a block is also defined within the block. This is done by adding block arguments to all basic block dominance chains which may lead to an operation that relied on referencing a Value based on basic block dominance.

This pass is useful in dataflow-style programming models since it renders all data flow within the program explicit (through block arguments) instead of implicit (through block dominance).

This pass only works on Standard-level IR, in that it expects all operations (and blocks) within a FuncOp to be within the same region. Furthermore, it is assumed that any value referenced by any operation is eligible to be passed around as a block argument.

[mortbopet]

incorrect namespace

[mortbopet]

missing }

[mortbopet]

argument and control flow mappings should be stored for any value

[rriddle]

mlir/include/mlir/Transforms/Passes.td

656

This option is giving me a bad code smell. Why is memref special cased here?

[rriddle]

mlir/lib/Transforms/MaxSSA.cpp

15

Do you actually reference anything from here?

16

What is being used from here? Realistically, passes in Transforms/ shouldn't depend on any dialects.

[mortbopet]

mlir/include/mlir/Transforms/Passes.td

656

Fully agree. I am a bit unsure on how to go about creating an option for steering the exclusion of user-requested types. I guess an alternative could be an optional list of strings, wherein these strings are string-compared to known types in the pass, and base exclusion on that. But that would still require handling for every possible type that could be requested.

ListOption<"ignoredTypes", "ignore-types", "std::string",...>
...
bool isIgnored(Value v) {
  return llvm::TypeSwitch<Type, bool>(v.getType)
    .Case<MemRefType>([&](auto){
      return llvm::find(ignoredTypes, "memref") != ignoredTypes.end()
    })
    .Default([&](auto){return false;});
}

Do you know a better way of going about this?

mlir/lib/Transforms/MaxSSA.cpp

15

Can be removed.

16

BranchOpInterface; will convert to #include "mlir/Interfaces/ControlFlowInterfaces.h"

[mortbopet]

update include

[mehdi_amini]

mlir/include/mlir/Transforms/Passes.td

656

What is the reason to not perform the transformations on memref values? Is this about the fact that it has pointer semantics? Would you want the same thing for the LLVM pointer type?

If so this is the kind of thing you can address with op-interfaces / type-interfaces.

[mortbopet]

mlir/include/mlir/Transforms/Passes.td

656

The specific use-case that i'm looking at is over in CIRCT where a pass is constrained to all memref's being statically determinable by either being function arguments or memref.alloc'ed in the entry block. This provides the motivation for somehow adding a hook into this pass to inform about certain types which should not be considered in SSA maximization - i think it would be best if this was done here, as to the alternative of maximizing all SSA values and then letting the user have to re-minimize things afterwards (which might be hard).

How do you see op/type interfaces being applicable to this problem, without unnecessarily polluting types with information specific to this single transformation?

[mehdi_amini]

mlir/include/mlir/Transforms/Passes.td

656

Interfaces (Op, Type, Attribute) are used to decouple the transformations from the entities (here a particular type).
Now it works only if there is an intrinsic property of the type that you could always use as discriminator to decide on whether to consider the type or not. But unfortunately it seems that your criteria is very much use-case specific, so interfaces won't be the right solution.

We're back to options: we'd need a serializable description of the configuration of the pass for that, a list of string could work even though it'll be limited in what it can describe.

[bondhugula]

mlir/test/Transforms/max_ssa.mlir

1 ↗

(On Diff #391405)

Please use hyphens instead of underscores for the test cases.

[bondhugula]

This would be a really useful utility/pass to have! It would be good to (now or later) factor the core utility into MLIRTransformsUtils and call this from the pass. Some initial mostly minor comments.

mlir/include/mlir/Transforms/Passes.h

133

-> ... is either defined with the same block or is one of its block arguments?

mlir/include/mlir/Transforms/Passes.td

632

Nit: convert-to-max-ssa or maximize-ssa to make it a verb?

645

-> ... standard control flow ...

646

to be within -> to be nested immediately within the same region?

[bondhugula]

mlir/include/mlir/Transforms/Passes.td

656

You can have the pass creation method take a lambda filterTypes -> bool to filter out certain types or have a string list of fully-specified type names. In your case, builtin.memref would be passed. I think it's a useful option in general but with specialized use cases.

[bondhugula]

mlir/lib/Transforms/MaxSSA.cpp

60

operand

144

Please document the return status and when/why it might fail at a high level.

153

Why do you need an std::set instead of a DenseSet here?

195

Use /*..=*/ form for nullptr argument.

[mehdi_amini]

mlir/include/mlir/Transforms/Passes.td

656

A lambda breaks reproducibility, so I'd really like to avoid this.

If you need to inject C++ code, then make it a utility instead of a pass, and then it's up to the user to wrap it in a pass that injects the lambda to the utility, but at least that keeps the pass mechanism as a "standalone" thing.

[bondhugula]

mlir/include/mlir/Transforms/Passes.td

656

Right - the lambda could be on the utility. Something like this will invariably be factored into a utility + "a pass calling that utility" like a lot of the other transforms infrastructure. Then, the pass will have to take a list of string type names?

[mortbopet]

address review comments

[mortbopet]

Rework the conversion pass into a transform utility. In doing so, also define a test pass for the filtered case. The transformation utility now takes a lambda to specify a value filter. The user-exposed pass convert-to-max-ssa no longer takes any arguments. For my own use-case, this is sufficient, since i'll be using the SSA maximization utility with a filtered type as a precondition for another conversion pass.

[rriddle]

Thanks, this works much better as a reusable utility.

mlir/include/mlir/Transforms/MaxSSAUtils.h

28

Does this need to be std::function? llvm::function_ref is generally a better choice if you only need a reference to a function (e.g. it doesn't allocate memory).

30–34

What aspects of this require being nested under FuncOp? This seems like a utility that could be written fairly generically over any root operation.

mlir/lib/Transforms/MaxSSA.cpp

29–30

36–40

mlir/lib/Transforms/Utils/MaxSSAUtils.cpp

118

Please use the global LogicalResult functions.

124

Same here and throughout.

188

Leftover debugging?

200

219–229

Please use explicit namespacing resolution:

https://llvm.org/docs/CodingStandards.html#use-namespace-qualifiers-to-implement-previously-declared-functions

[bondhugula]

mlir/include/mlir/Transforms/MaxSSAUtils.h

30

Nit: Start with Converts a ....

30–34

Many region holding operations don't allow block arguments to be arbitrarily added to them -- some have very specific meanings for those (scf.for, affine.for, gpu.launch) and some may not even allow block arguments (scf.execute_region). This can perhaps be generalized to FunctionLike ops.

[bondhugula]

mlir/lib/Transforms/Utils/MaxSSAUtils.cpp

20

Unnecessary include here.

73–74

You aren't using this variable's value later (it's being rewritten). Please use llvm::is_contained(...).

94

auto alreadyinBlockValues = ...;

[mortbopet]

mlir/lib/Transforms/MaxSSA.cpp

153

Changed to DenseSet.

mlir/lib/Transforms/Utils/MaxSSAUtils.cpp

73–74

llvm::is_contained(...) doesn't seem to apply for llvm::DenseMap. Will use inBlockValues[v].count(use) == 0 instead.

[mortbopet]

address review comments.

[rriddle]

mlir/include/mlir/Transforms/MaxSSAUtils.h

30–34

Right, the constraints here should likely be captured in a comment. In the future this would be better driven through interfaces/traits, rather than limiting to FunctionLike/FuncOp/etc.

[rriddle]

mlir/include/mlir/Transforms/MaxSSAUtils.h

63–78

This is missing the fact that the value is required to be defined by an operation within a FuncOp.

mlir/lib/Transforms/Utils/MaxSSAUtils.cpp

141–148

This seems more appropriate as an assert than an error.

mlir/test/Transforms/max-ssa-ignore.mlir

6–62

Can you trim this down to the smallest test you can create? This seems quite large for a focused test.

mlir/test/Transforms/max-ssa.mlir

113–196

What is this testing? It isn't clear that this is providing a focused test, and seemingly includes a lot of unnecessary operations.

[mortbopet]

add missing early_inc_range and simplify tests

[mortbopet]

address review comments

[bondhugula]

mlir/lib/Transforms/Utils/MaxSSAUtils.cpp

45

Please name this replaceAllUsesInBlockWith. Consider aligning the signature of this with mlir::replaceAllUsesInRegionWith. This can even be moved to RegionUtils.cpp if desired.

142

You can't use the result of getDefiningOp unchecked. You are missing an assert or this is a bug.

149

For private methods, I think it's a convention to have the doc comment on the function definition.

173

We need a comment here or below.

[mortbopet]

reduce number of times each block is visited

[rriddle]

mlir/lib/Transforms/MaxSSA.cpp

37–41

Use proper namespace resolution instead of using namespaces:

https://llvm.org/docs/CodingStandards.html#use-namespace-qualifiers-to-implement-previously-declared-functions

mlir/lib/Transforms/Utils/MaxSSAUtils.cpp

141

mlir/test/lib/Transforms/TestMaxSSAFiltered.cpp

2

nit: Please name this the same more generically as the pass itself, just in case we add more test passes related to the pass here.

[mortbopet]

Address review comments.

[bondhugula]

I'm deferring to River here given the more detailed review he has done. Just making sure this isn't waiting for me.

[mortbopet]

@rriddle do you have anything else you'd like too see done for the PR?

mlir/lib/Transforms/Utils/MaxSSAUtils.cpp

142

Thank you for catching that; rewriting it so it can be used with block arguments as well.

mlir/test/Transforms/max-ssa-ignore.mlir

6–62

Updated test cases (in max-ssa.mlir) to mostly contain control flow operators, and pass around values defined in the entry block.

[rriddle]

mlir/include/mlir/Transforms/MaxSSAUtils.h

29

nit: You can drop the llvm:: here.

65

mlir/lib/Transforms/Utils/MaxSSAUtils.cpp

35–38

Is this assuming that a branch only has unique successors? That isn't the case in MLIR, e.g. a conditional branch can go to the same block in both the true and false case.

91

Does this need to be inserted after? or could you fuse this into the if above? (i.e. call insert and check the return)

206

Only the class should be in the anonymous namespace, the functions should be at top-scope.

mlir/lib/Transforms/Utils/RegionUtils.cpp

34–36 ↗

(On Diff #393811)

This could be implemented using Value::replaceUsesWithIf.

mlir/include/mlir/Transforms/MaxSSAUtils.h

+//===- MaxSSAUtils.h - Maximal SSA form utilities ---------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines prototypes for various utilities for converting a
+// function using standard control flow into having maximal SSA form.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_MAX_SSA_UTILS_H
+#define MLIR_TRANSFORMS_MAX_SSA_UTILS_H
+
+#include "mlir/IR/Block.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
+#include <functional>
+
+namespace mlir {
+class FuncOp;
+
+/// A filter callback function that can be used to determine if a given value
+/// should be ignored during maximal SSA form conversion. If the function
+/// returns true, usages of the value will remain unmodified.
+using ValueFilterCallbackFn = std::function<bool(Value)>;
+
+/// This utility converts a FuncOp into maximal SSA form. In this form, any
+/// value referenced within a block is also defined within the block. This is
+/// done by adding block arguments to all basic block dominance chains which may
+/// lead to an operation that relied on referencing a Value based on basic block
+/// dominance.
+///
+/// This utility is useful in dataflow-style programming models since it renders
+/// all data flow within the program explicit (through block arguments) instead
+/// of implicit (through block dominance)
+///
+/// This utility only works on standard control flow, in that it expects all
+/// operations (and blocks) within a FuncOp to be nested immediately within the
+/// same region. Furthermore, it is assumed that any value referenced by any
+/// operation is eligible to be passed around as a block argument.
+///
+/// An optional ValueFilterCallbackFn may be provided to filter out values from
+/// being considered during SSA maximization.
+///
+/// Example:
+///   func @simple(%arg0 : index, %arg1 : index) -> (index, index, index) {
+///     br ^bb1(%arg0 : index)
+///   ^bb1(%0 : index):
+///     return %0, %arg0, %arg1 : index, index, index
+///   }
+/// -> convertToMaximalSSA(@simple)
+///   func @simple(%arg0 : index, %arg1 : index) -> (index, index, index) {
+///     br ^bb1(%arg0, %arg0, %arg1 : index, index, index)
+///   ^bb1(%0, %1, %2 : index, index):
+///     return %0, %1, %2 : index, index, index
+///   }
+LogicalResult convertToMaximalSSA(FuncOp func,
+                                  ValueFilterCallbackFn filterFn = nullptr);
+
+/// Like convertToMaximalSSA(FuncOp) but restricted to converting references to
+/// a single 'value' into maximal SSA form
+///
+/// Example:
+///   func @simple(%arg0 : index, %arg1 : index) -> (index, index, index) {
+///     br ^bb1(%arg0 : index)
+///   ^bb1(%0 : index):
+///     return %0, %arg0, %arg1 : index, index, index
+///   }
+/// -> convertToMaximalSSA(%arg1)
+///   func @simple(%arg0 : index, %arg1 : index) -> (index, index, index) {
+///     br ^bb1(%arg0, %arg1 : index, index)
+///   ^bb1(%0, %1 : index, index):
+///     return %0, %arg0, %1 : index, index, index
+///   }
+LogicalResult convertToMaximalSSA(Value value);
+
+} // end namespace mlir
+
+#endif // MLIR_TRANSFORMS_MAX_SSA_UTILS_H

mlir/include/mlir/Transforms/Passes.h

 /// Creates a pass which delete symbol operations that are unreachable. This
 /// pass may *only* be scheduled on an operation that defines a SymbolTable.
 std::unique_ptr<Pass> createSymbolDCEPass();
 
 /// Creates an interprocedural pass to normalize memrefs to have a trivial
 /// (identity) layout map.
 std::unique_ptr<OperationPass<ModuleOp>> createNormalizeMemRefsPass();
 
+/// Creates a pass which converts a program into maximal SSA form. In this form,
+/// any value referenced within a block is also defined within the same block.
+std::unique_ptr<mlir::Pass> createMaxSSAFormPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Transforms/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_PASSES_H

mlir/include/mlir/Transforms/Passes.td

let options = [
     Option<"printDataFlowEdges", "print-data-flow-edges", "bool",
            /*default=*/"true", "Print data flow edges">,
     Option<"printResultTypes", "print-result-types", "bool",
             /*default=*/"true", "Print result types of operations">
   ];
   let constructor = "mlir::createPrintOpGraphPass()";
 }
 
+def MaxSSAForm : FunctionPass<"convert-to-max-ssa"> {
+  let summary = "Converts a program into maximal SSA form";
+  let description = [{
+    This pass converts a program into maximal SSA form. In this form, any value
+    referenced within a block is also defined within the block. This is done by
+    adding block arguments to all basic block dominance chains which may lead
+    to an operation that relied on referencing a Value based on basic block
+    dominance.
+    
+    This pass is useful in dataflow-style programming models since it renders
+    all data flow within the program explicit (through block arguments) instead
+    of implicit (through block dominance).
+
+    This pass only works on standard control flow, in that it expects all operations
+    (and blocks) within a FuncOp to be nested immediately within the same region.
+    Furthermore, it is assumed that any value referenced by any operation is eligible
+    to be passed around as a block argument.
+  }];
+  let constructor = "mlir::createMaxSSAFormPass()";
+}
+
 #endif // MLIR_TRANSFORMS_PASSES

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   BufferOptimizations.cpp
   BufferResultsToOutParams.cpp
   BufferUtils.cpp
   Canonicalizer.cpp
   CSE.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopInvariantCodeMotion.cpp
+  MaxSSA.cpp
   NormalizeMemRefs.cpp
   OpStats.cpp
   ParallelLoopCollapsing.cpp
   PipelineDataTransfer.cpp
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
   ViewOpGraph.cpp

mlir/lib/Transforms/MaxSSA.cpp

+//===- MaxSSA.cpp - Maximal SSA form conversion ------------------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains the definitions of the maximal SSA form conversion pass.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Transforms/MaxSSAUtils.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+#include <set>
+
+using namespace mlir;
+
+namespace {
+
+struct MaxSSAFormPass : public MaxSSAFormBase<MaxSSAFormPass> {
+public:
+  void runOnFunction() override {
+    FuncOp func = getOperation();
+    if (convertToMaximalSSA(func).failed())
+      return signalPassFailure();
+  }
+};
+
+} // namespace
+
+namespace mlir {
+std::unique_ptr<mlir::Pass> createMaxSSAFormPass() {
+  return std::make_unique<MaxSSAFormPass>();
+}
+} // namespace mlir

mlir/lib/Transforms/Utils/CMakeLists.txt

 add_mlir_library(MLIRTransformUtils
   DialectConversion.cpp
   FoldUtils.cpp
   GreedyPatternRewriteDriver.cpp
   InliningUtils.cpp
   LoopFusionUtils.cpp
   LoopUtils.cpp
+  MaxSSAUtils.cpp
   RegionUtils.cpp
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS
   MLIRStandardOpsIncGen

mlir/lib/Transforms/Utils/MaxSSAUtils.cpp

+//===- MaxSSAUtils.cpp - Maximal SSA form utilities --------------*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains various utilities for converting a function using standard control
+// flow into having maximal SSA form.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/MaxSSAUtils.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+#include <set>
+
+using namespace mlir;
+
+/// Rewrites the terminator in 'from' to pass an additional argument 'v' when
+/// passing control flow to 'to'.
+static LogicalResult rewriteControlFlowToBlock(Block *from, Block *to,
+                                               Value v) {
+  auto termOp = from->getTerminator();
+  auto branchOp = dyn_cast<BranchOpInterface>(termOp);
+  if (!branchOp)
+    return termOp->emitOpError() << "expected terminator op within control "
+                                    "flow to be a branch-like op";
+
+  // Find 'to' successor in the branch op.
+  auto successors = branchOp->getSuccessors();
+  auto succIt = llvm::find(successors, to);
+  assert(succIt != successors.end());
+  unsigned succIdx = std::distance(successors.begin(), succIt);
+
+  // Add the value 'v' as a block argument.
+  branchOp.getMutableSuccessorOperands(succIdx)->append(v);
+  return success();
+}
+
+/// Rewrites uses of 'oldV' in 'b' to 'newV'.
+static void rewriteUsageInBlock(Block *b, Value oldV, Value newV) {
+  for (auto &use : llvm::make_early_inc_range(oldV.getUses()))
+    if (use.getOwner()->getBlock() == b)
+      use.set(newV);
+}
+
+/// Perform a depth-first search backwards through the CFG graph of a program,
+/// starting from 'use', add a new block argument of type(v) to the block and
+/// replaces all uses of 'v' with the new block argument.
+///
+/// Arguments: 'use': ablock where a value v flows through. 'succ': a successor
+/// block of the 'use' block. Notably this conversion backtracks through the BB
+/// CFG graph, so 'succ' will be a basic block that called backtrackAndConvert
+/// on 'use'. 'inBlockValues': A mapping containing the appended block argument
+/// when backtracking through a basic block. 'convertedControlFlow': A mapping
+/// containing, for a given block (key) which successor operand range in the
+/// terminator have been rewritten to the new block argument signature.
+static LogicalResult
+backtrackAndConvert(Block *use, Block *succ, Value v,
+                    DenseMap<Value, DenseMap<Block *, Value>> &inBlockValues,
+                    DenseMap<Value, DenseMap<Block *, DenseSet<Block *>>>
+                        &convertedControlFlow) {
+  // The base case is when we've backtracked to the Block which defines the
+  // value. In these cases, set the actual value as the converted value.
+  if (v.getParentBlock() == use)
+    inBlockValues[v][use] = v;
+
+  auto alreadyinBlockValues = inBlockValues[v].find(use);
+  if (alreadyinBlockValues == inBlockValues[v].end()) {
+    // Rewrite this blocks' block arguments to take in a new value of 'v' type.
+    use->addArgument(v.getType());
+    Value newBarg = use->getArguments().back();
+
+    // Register the converted block argument in case other branches in the CFG
+    // arrive here later.
+    inBlockValues[v][use] = newBarg;
+    rewriteUsageInBlock(use, v, newBarg);
+
+    // Recurse through the predecessors of this block.
+    for (auto pred : use->getPredecessors())
+      if (backtrackAndConvert(pred, use, v, inBlockValues, convertedControlFlow)
+              .failed())
+        return failure();
+  }
+
+  // Rewrite control flow to the 'succ' block through the terminator, if not
+  // already done.
+  if (succ && convertedControlFlow[v][use].count(succ) == 0) {
+    alreadyinBlockValues = inBlockValues[v].find(use);
+    assert(alreadyinBlockValues != inBlockValues[v].end());
+    if (rewriteControlFlowToBlock(use, succ, alreadyinBlockValues->second)
+            .failed())
+      return failure();
+    convertedControlFlow[v][use].insert(succ);
+  }
+
+  return success();
+}
+
+namespace {
+
+struct MaxSSAFormConverter {
+public:
+  /// An optional filterFn may be provided to dynamically filter out values
+  /// from being converted.
+  MaxSSAFormConverter(ValueFilterCallbackFn filterFn = nullptr)
+      : filterFn(filterFn) {}
+
+  LogicalResult convertFunction(FuncOp function) {
+    auto walkRes = function.walk([&](Operation *op) {
+      // Run on operation results.
+      if (llvm::any_of(op->getResults(),
+                       [&](Value res) { return runOnValue(res).failed(); })) {
+        return WalkResult::interrupt();
+      }
+      // Run on block arguments.
+      for (auto &block : function) {
+        if (llvm::any_of(block.getArguments(), [&](Value barg) {
+              return runOnValue(barg).failed();
+            })) {
+          return WalkResult::interrupt();
+        }
+      }
+      return WalkResult::advance();
+    });
+
+    if (walkRes.wasInterrupted())
+      return failure();
+
+    assert(
+        verifyFunction(function).succeeded() &&
+        "Some values were still referenced outside of their defining block!");
+
+    return success();
+  }
+
+  LogicalResult convertValue(Value v) {
+    auto *defOp = v.getDefiningOp();
+    auto funcOp = dyn_cast<FuncOp>(defOp->getParentOp());
+    if (!funcOp)
+      return defOp->emitOpError() << "Expected parent operation to be a "
+                                     "function, but got "
+                                  << defOp->getParentOp()->getName();
+
+    return runOnValue(v);
+  }
+
+private:
+  /// Verifies that all values which are not filtered indeed are only referenced
+  /// within their defining block.
+  LogicalResult verifyFunction(FuncOp f) const;
+
+  /// Driver which will run backtrackAndConvert on values referenced outside
+  /// their defining block. Returns failure in case the pass failed to apply.
+  /// This may happen when nested regions exist within the FuncOp which this
+  /// pass is applied to, or if non branch-like control flow is used.
+  LogicalResult runOnValue(Value v);
+
+  /// A mapping {original value : {block : replaced value}} representing
+  /// 'original value' has been replaced in 'block' with 'replaced value'".
+  DenseMap<Value, DenseMap<Block *, Value>> inBlockValues;
+
+  /// A mapping {original value : {block : succ block}} representing
+  /// 'original value' has already been passed from 'block' to 'succ block'
+  /// through the terminator of 'block'.
+  DenseMap<Value, DenseMap<Block *, DenseSet<Block *>>> convertedControlFlow;
+
+  /// An optional filter function to dynamically determine whether a value
+  /// should be considered for SSA maximization.
+  ValueFilterCallbackFn filterFn;
+};
+
+LogicalResult MaxSSAFormConverter::verifyFunction(FuncOp f) const {
+  for (auto &op : f.getOps()) {
+    auto isValid = [&](Value v) {
+      if (filterFn && filterFn(v))
+        return true;
+
+      return v.getParentBlock() == op.getBlock();
+    };
+
+    if (!llvm::all_of(op.getOperands(), isValid)) {
+      f.dump();
+      return op.emitOpError()
+             << "has operands that are not defined within its block";
+    }
+  }
+  return success();
+}
+
+LogicalResult MaxSSAFormConverter::runOnValue(Value v) {
+  if (filterFn && filterFn(v))
+    return success();
+  Block *definingBlock = v.getParentBlock();
+  for (auto user : v.getUsers()) {
+    Block *userBlock = user->getBlock();
+    if (definingBlock != userBlock) {
+      // This is a case of using an SSA value through basic block dominance.
+      if (userBlock->getParent() != definingBlock->getParent())
+        return user->emitOpError() << "can only convert SSA usage across "
+                                      "blocks in the same region.";
+
+      if (backtrackAndConvert(userBlock, /*succ=*/nullptr, v, inBlockValues,
+                              convertedControlFlow)
+              .failed())
+        return failure();
+    }
+  }
+  return success();
+}
+
+} // namespace
+
+namespace mlir {
+
+LogicalResult convertToMaximalSSA(FuncOp func, ValueFilterCallbackFn filterFn) {
+  return MaxSSAFormConverter(filterFn).convertFunction(func);
+}
+
+LogicalResult convertToMaximalSSA(Value value) {
+  return MaxSSAFormConverter().convertValue(value);
+}
+
+} // namespace mlir

mlir/test/Transforms/max-ssa-errors.mlir

+// RUN: mlir-opt -split-input-file --convert-to-max-ssa %s -verify-diagnostics
+
+func @scf(%i : index) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c0_i32 = arith.constant 0 : i32
+  %m = memref.alloc() : memref<4xi32>
+  scf.for %0 = %c0 to %c4 step %c1 {
+    // expected-error @+1 {{'memref.store' op can only convert SSA usage across blocks in the same region.}}
+    memref.store %c0_i32, %m[%i] : memref<4xi32>
+  }
+  return
+}

mlir/test/Transforms/max-ssa-ignore.mlir

+// RUN: mlir-opt -split-input-file --test-max-ssa-filtered %s | FileCheck %s
+
+// CHECK-LABEL:   func @memory_loop() -> i32 {
+// CHECK:           %[[VAL_0:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_2:.*]] = arith.constant 4 : index
+// CHECK:           %[[VAL_3:.*]] = arith.constant 5 : i32
+// CHECK:           %[[VAL_4:.*]] = memref.alloc() : memref<64xi32>
+// CHECK:           %[[VAL_5:.*]] = memref.alloc() : memref<64xi32>
+// CHECK:           br ^bb1(%[[VAL_0]], %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]] : index, index, index, index, i32)
+// CHECK:         ^bb1(%[[VAL_6:.*]]: index, %[[VAL_7:.*]]: index, %[[VAL_8:.*]]: index, %[[VAL_9:.*]]: index, %[[VAL_10:.*]]: i32):
+// CHECK:           %[[VAL_11:.*]] = arith.cmpi slt, %[[VAL_6]], %[[VAL_9]] : index
+// CHECK:           cond_br %[[VAL_11]], ^bb2(%[[VAL_7]], %[[VAL_6]], %[[VAL_8]], %[[VAL_9]], %[[VAL_10]] : index, index, index, index, i32), ^bb3(%[[VAL_7]], %[[VAL_8]], %[[VAL_9]] : index, index, index)
+// CHECK:         ^bb2(%[[VAL_12:.*]]: index, %[[VAL_13:.*]]: index, %[[VAL_14:.*]]: index, %[[VAL_15:.*]]: index, %[[VAL_16:.*]]: i32):
+// CHECK:           memref.store %[[VAL_16]], %[[VAL_4]]{{\[}}%[[VAL_13]]] : memref<64xi32>
+// CHECK:           %[[VAL_17:.*]] = arith.addi %[[VAL_13]], %[[VAL_14]] : index
+// CHECK:           br ^bb1(%[[VAL_17]], %[[VAL_12]], %[[VAL_14]], %[[VAL_15]], %[[VAL_16]] : index, index, index, index, i32)
+// CHECK:         ^bb3(%[[VAL_18:.*]]: index, %[[VAL_19:.*]]: index, %[[VAL_20:.*]]: index):
+// CHECK:           br ^bb4(%[[VAL_18]], %[[VAL_18]], %[[VAL_19]], %[[VAL_20]] : index, index, index, index)
+// CHECK:         ^bb4(%[[VAL_21:.*]]: index, %[[VAL_22:.*]]: index, %[[VAL_23:.*]]: index, %[[VAL_24:.*]]: index):
+// CHECK:           %[[VAL_25:.*]] = arith.cmpi slt, %[[VAL_21]], %[[VAL_24]] : index
+// CHECK:           cond_br %[[VAL_25]], ^bb5(%[[VAL_22]], %[[VAL_21]], %[[VAL_23]], %[[VAL_24]] : index, index, index, index), ^bb6(%[[VAL_22]] : index)
+// CHECK:         ^bb5(%[[VAL_26:.*]]: index, %[[VAL_27:.*]]: index, %[[VAL_28:.*]]: index, %[[VAL_29:.*]]: index):
+// CHECK:           %[[VAL_30:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_27]]] : memref<64xi32>
+// CHECK:           %[[VAL_31:.*]] = arith.addi %[[VAL_30]], %[[VAL_30]] : i32
+// CHECK:           memref.store %[[VAL_31]], %[[VAL_5]]{{\[}}%[[VAL_27]]] : memref<64xi32>
+// CHECK:           %[[VAL_32:.*]] = arith.addi %[[VAL_27]], %[[VAL_28]] : index
+// CHECK:           br ^bb4(%[[VAL_32]], %[[VAL_26]], %[[VAL_28]], %[[VAL_29]] : index, index, index, index)
+// CHECK:         ^bb6(%[[VAL_33:.*]]: index):
+// CHECK:           %[[VAL_34:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_33]]] : memref<64xi32>
+// CHECK:           return %[[VAL_34]] : i32
+// CHECK:         }
+func @memory_loop() -> i32 {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c5_i32 = arith.constant 5 : i32
+  %0 = memref.alloc() : memref<64xi32>
+  %1 = memref.alloc() : memref<64xi32>
+  br ^bb1(%c0 : index)
+^bb1(%2: index): 
+  %3 = arith.cmpi slt, %2, %c4 : index
+  cond_br %3, ^bb2, ^bb3
+^bb2:
+  memref.store %c5_i32, %0[%2] : memref<64xi32>
+  %4 = arith.addi %2, %c1 : index
+  br ^bb1(%4 : index)
+^bb3:
+  br ^bb4(%c0 : index)
+^bb4(%5: index):
+  %6 = arith.cmpi slt, %5, %c4 : index
+  cond_br %6, ^bb5, ^bb6
+^bb5:
+  %7 = memref.load %0[%5] : memref<64xi32>
+  %8 = arith.addi %7, %7 : i32
+  memref.store %8, %1[%5] : memref<64xi32>
+  %9 = arith.addi %5, %c1 : index
+  br ^bb4(%9 : index)
+^bb6:
+  %10 = memref.load %1[%c0] : memref<64xi32>
+  return %10 : i32
+}

mlir/test/Transforms/max-ssa.mlir

+// RUN: mlir-opt -split-input-file --convert-to-max-ssa %s | FileCheck %s
+
+// CHECK-LABEL:   func @simple(
+// CHECK-SAME:                 %[[VAL_0:.*]]: index,
+// CHECK-SAME:                 %[[VAL_1:.*]]: index) -> (index, index) {
+// CHECK:           br ^bb1(%[[VAL_0]], %[[VAL_1]] : index, index)
+// CHECK:         ^bb1(%[[VAL_2:.*]]: index, %[[VAL_3:.*]]: index):
+// CHECK:           return %[[VAL_2]], %[[VAL_3]] : index, index
+// CHECK:         }
+func @simple(%arg0 : index, %arg1 : index) -> (index, index) {
+  br ^bb1(%arg0 : index)
+^bb1(%0 : index):
+  return %0, %arg1 : index, index
+}
+
+// -----
+
+// CHECK-LABEL:   func @backedge() -> index {
+// CHECK:           %[[VAL_0:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_1:.*]] = arith.constant 42 : index
+// CHECK:           %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK:           br ^bb1(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]] : index, index, index)
+// CHECK:         ^bb1(%[[VAL_3:.*]]: index, %[[VAL_4:.*]]: index, %[[VAL_5:.*]]: index):
+// CHECK:           %[[VAL_6:.*]] = arith.cmpi slt, %[[VAL_3]], %[[VAL_4]] : index
+// CHECK:           cond_br %[[VAL_6]], ^bb2(%[[VAL_4]], %[[VAL_3]], %[[VAL_5]] : index, index, index), ^bb3(%[[VAL_3]] : index)
+// CHECK:         ^bb2(%[[VAL_7:.*]]: index, %[[VAL_8:.*]]: index, %[[VAL_9:.*]]: index):
+// CHECK:           %[[VAL_10:.*]] = arith.addi %[[VAL_8]], %[[VAL_9]] : index
+// CHECK:           br ^bb1(%[[VAL_10]], %[[VAL_7]], %[[VAL_9]] : index, index, index)
+// CHECK:         ^bb3(%[[VAL_11:.*]]: index):
+// CHECK:           return %[[VAL_11]] : index
+// CHECK:         }
+func @backedge() -> index {
+  %c1 = arith.constant 1 : index
+  %c42 = arith.constant 42 : index
+  %c1_0 = arith.constant 1 : index
+  br ^bb1(%c1 : index)
+^bb1(%0: index):
+  %1 = arith.cmpi slt, %0, %c42 : index
+  cond_br %1, ^bb2, ^bb3
+^bb2:
+  %2 = arith.addi %0, %c1_0 : index
+  br ^bb1(%2 : index)
+^bb3:
+  return %0 : index
+}
+
+// -----
+
+// CHECK-LABEL:   func @memory_loop() -> i32 {
+// CHECK:           %[[VAL_0:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_2:.*]] = arith.constant 4 : index
+// CHECK:           %[[VAL_3:.*]] = arith.constant 5 : i32
+// CHECK:           %[[VAL_4:.*]] = memref.alloc() : memref<64xi32>
+// CHECK:           %[[VAL_5:.*]] = memref.alloc() : memref<64xi32>
+// CHECK:           br ^bb1(%[[VAL_0]], %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]], %[[VAL_4]], %[[VAL_5]] : index, index, index, index, i32, memref<64xi32>, memref<64xi32>)
+// CHECK:         ^bb1(%[[VAL_6:.*]]: index, %[[VAL_7:.*]]: index, %[[VAL_8:.*]]: index, %[[VAL_9:.*]]: index, %[[VAL_10:.*]]: i32, %[[VAL_11:.*]]: memref<64xi32>, %[[VAL_12:.*]]: memref<64xi32>):
+// CHECK:           %[[VAL_13:.*]] = arith.cmpi slt, %[[VAL_6]], %[[VAL_9]] : index
+// CHECK:           cond_br %[[VAL_13]], ^bb2(%[[VAL_7]], %[[VAL_6]], %[[VAL_8]], %[[VAL_9]], %[[VAL_10]], %[[VAL_11]], %[[VAL_12]] : index, index, index, index, i32, memref<64xi32>, memref<64xi32>), ^bb3(%[[VAL_7]], %[[VAL_8]], %[[VAL_9]], %[[VAL_11]], %[[VAL_12]] : index, index, index, memref<64xi32>, memref<64xi32>)
+// CHECK:         ^bb2(%[[VAL_14:.*]]: index, %[[VAL_15:.*]]: index, %[[VAL_16:.*]]: index, %[[VAL_17:.*]]: index, %[[VAL_18:.*]]: i32, %[[VAL_19:.*]]: memref<64xi32>, %[[VAL_20:.*]]: memref<64xi32>):
+// CHECK:           memref.store %[[VAL_18]], %[[VAL_19]]{{\[}}%[[VAL_15]]] : memref<64xi32>
+// CHECK:           %[[VAL_21:.*]] = arith.addi %[[VAL_15]], %[[VAL_16]] : index
+// CHECK:           br ^bb1(%[[VAL_21]], %[[VAL_14]], %[[VAL_16]], %[[VAL_17]], %[[VAL_18]], %[[VAL_19]], %[[VAL_20]] : index, index, index, index, i32, memref<64xi32>, memref<64xi32>)
+// CHECK:         ^bb3(%[[VAL_22:.*]]: index, %[[VAL_23:.*]]: index, %[[VAL_24:.*]]: index, %[[VAL_25:.*]]: memref<64xi32>, %[[VAL_26:.*]]: memref<64xi32>):
+// CHECK:           br ^bb4(%[[VAL_22]], %[[VAL_22]], %[[VAL_23]], %[[VAL_24]], %[[VAL_25]], %[[VAL_26]] : index, index, index, index, memref<64xi32>, memref<64xi32>)
+// CHECK:         ^bb4(%[[VAL_27:.*]]: index, %[[VAL_28:.*]]: index, %[[VAL_29:.*]]: index, %[[VAL_30:.*]]: index, %[[VAL_31:.*]]: memref<64xi32>, %[[VAL_32:.*]]: memref<64xi32>):
+// CHECK:           %[[VAL_33:.*]] = arith.cmpi slt, %[[VAL_27]], %[[VAL_30]] : index
+// CHECK:           cond_br %[[VAL_33]], ^bb5(%[[VAL_28]], %[[VAL_27]], %[[VAL_29]], %[[VAL_30]], %[[VAL_31]], %[[VAL_32]] : index, index, index, index, memref<64xi32>, memref<64xi32>), ^bb6(%[[VAL_28]], %[[VAL_32]] : index, memref<64xi32>)
+// CHECK:         ^bb5(%[[VAL_34:.*]]: index, %[[VAL_35:.*]]: index, %[[VAL_36:.*]]: index, %[[VAL_37:.*]]: index, %[[VAL_38:.*]]: memref<64xi32>, %[[VAL_39:.*]]: memref<64xi32>):
+// CHECK:           %[[VAL_40:.*]] = memref.load %[[VAL_38]]{{\[}}%[[VAL_35]]] : memref<64xi32>
+// CHECK:           %[[VAL_41:.*]] = arith.addi %[[VAL_40]], %[[VAL_40]] : i32
+// CHECK:           memref.store %[[VAL_41]], %[[VAL_39]]{{\[}}%[[VAL_35]]] : memref<64xi32>
+// CHECK:           %[[VAL_42:.*]] = arith.addi %[[VAL_35]], %[[VAL_36]] : index
+// CHECK:           br ^bb4(%[[VAL_42]], %[[VAL_34]], %[[VAL_36]], %[[VAL_37]], %[[VAL_38]], %[[VAL_39]] : index, index, index, index, memref<64xi32>, memref<64xi32>)
+// CHECK:         ^bb6(%[[VAL_43:.*]]: index, %[[VAL_44:.*]]: memref<64xi32>):
+// CHECK:           %[[VAL_45:.*]] = memref.load %[[VAL_44]]{{\[}}%[[VAL_43]]] : memref<64xi32>
+// CHECK:           return %[[VAL_45]] : i32
+// CHECK:         }
+func @memory_loop() -> i32 {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c5_i32 = arith.constant 5 : i32
+  %0 = memref.alloc() : memref<64xi32>
+  %1 = memref.alloc() : memref<64xi32>
+  br ^bb1(%c0 : index)
+^bb1(%2: index): 
+  %3 = arith.cmpi slt, %2, %c4 : index
+  cond_br %3, ^bb2, ^bb3
+^bb2:
+  memref.store %c5_i32, %0[%2] : memref<64xi32>
+  %4 = arith.addi %2, %c1 : index
+  br ^bb1(%4 : index)
+^bb3:
+  br ^bb4(%c0 : index)
+^bb4(%5: index):
+  %6 = arith.cmpi slt, %5, %c4 : index
+  cond_br %6, ^bb5, ^bb6
+^bb5:
+  %7 = memref.load %0[%5] : memref<64xi32>
+  %8 = arith.addi %7, %7 : i32
+  memref.store %8, %1[%5] : memref<64xi32>
+  %9 = arith.addi %5, %c1 : index
+  br ^bb4(%9 : index)
+^bb6:
+  %10 = memref.load %1[%c0] : memref<64xi32>
+  return %10 : i32
+}
+
+// -----
+
+// CHECK-LABEL:   func @matrix_power(
+// CHECK-SAME:                       %[[VAL_0:.*]]: memref<400xi32>, %[[VAL_1:.*]]: memref<20xi32>, %[[VAL_2:.*]]: memref<20xi32>, %[[VAL_3:.*]]: memref<20xi32>) {
+// CHECK:           %[[VAL_4:.*]] = arith.constant 1 : i32
+// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_6:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_7:.*]] = arith.constant 20 : index
+// CHECK:           br ^bb1(%[[VAL_5]], %[[VAL_4]], %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]], %[[VAL_5]], %[[VAL_6]], %[[VAL_7]] : index, i32, memref<400xi32>, memref<20xi32>, memref<20xi32>, memref<20xi32>, index, index, index)
+// CHECK:         ^bb1(%[[VAL_8:.*]]: index, %[[VAL_9:.*]]: i32, %[[VAL_10:.*]]: memref<400xi32>, %[[VAL_11:.*]]: memref<20xi32>, %[[VAL_12:.*]]: memref<20xi32>, %[[VAL_13:.*]]: memref<20xi32>, %[[VAL_14:.*]]: index, %[[VAL_15:.*]]: index, %[[VAL_16:.*]]: index):
+// CHECK:           %[[VAL_17:.*]] = arith.cmpi slt, %[[VAL_8]], %[[VAL_16]] : index
+// CHECK:           cond_br %[[VAL_17]], ^bb2(%[[VAL_8]], %[[VAL_9]], %[[VAL_10]], %[[VAL_11]], %[[VAL_12]], %[[VAL_13]], %[[VAL_14]], %[[VAL_15]], %[[VAL_16]] : index, i32, memref<400xi32>, memref<20xi32>, memref<20xi32>, memref<20xi32>, index, index, index), ^bb5
+// CHECK:         ^bb2(%[[VAL_18:.*]]: index, %[[VAL_19:.*]]: i32, %[[VAL_20:.*]]: memref<400xi32>, %[[VAL_21:.*]]: memref<20xi32>, %[[VAL_22:.*]]: memref<20xi32>, %[[VAL_23:.*]]: memref<20xi32>, %[[VAL_24:.*]]: index, %[[VAL_25:.*]]: index, %[[VAL_26:.*]]: index):
+// CHECK:           %[[VAL_27:.*]] = arith.addi %[[VAL_18]], %[[VAL_24]] : index
+// CHECK:           %[[VAL_28:.*]] = arith.subi %[[VAL_18]], %[[VAL_24]] : index
+// CHECK:           %[[VAL_29:.*]] = arith.index_cast %[[VAL_28]] : index to i32
+// CHECK:           %[[VAL_30:.*]] = arith.addi %[[VAL_29]], %[[VAL_19]] : i32
+// CHECK:           %[[VAL_31:.*]] = arith.index_cast %[[VAL_30]] : i32 to index
+// CHECK:           br ^bb3(%[[VAL_25]], %[[VAL_19]], %[[VAL_20]], %[[VAL_21]], %[[VAL_22]], %[[VAL_23]], %[[VAL_24]], %[[VAL_25]], %[[VAL_26]], %[[VAL_27]], %[[VAL_31]] : index, i32, memref<400xi32>, memref<20xi32>, memref<20xi32>, memref<20xi32>, index, index, index, index, index)
+// CHECK:         ^bb3(%[[VAL_32:.*]]: index, %[[VAL_33:.*]]: i32, %[[VAL_34:.*]]: memref<400xi32>, %[[VAL_35:.*]]: memref<20xi32>, %[[VAL_36:.*]]: memref<20xi32>, %[[VAL_37:.*]]: memref<20xi32>, %[[VAL_38:.*]]: index, %[[VAL_39:.*]]: index, %[[VAL_40:.*]]: index, %[[VAL_41:.*]]: index, %[[VAL_42:.*]]: index):
+// CHECK:           %[[VAL_43:.*]] = arith.cmpi slt, %[[VAL_32]], %[[VAL_40]] : index
+// CHECK:           cond_br %[[VAL_43]], ^bb4(%[[VAL_32]], %[[VAL_33]], %[[VAL_34]], %[[VAL_35]], %[[VAL_36]], %[[VAL_37]], %[[VAL_38]], %[[VAL_39]], %[[VAL_40]], %[[VAL_41]], %[[VAL_42]] : index, i32, memref<400xi32>, memref<20xi32>, memref<20xi32>, memref<20xi32>, index, index, index, index, index), ^bb1(%[[VAL_41]], %[[VAL_33]], %[[VAL_34]], %[[VAL_35]], %[[VAL_36]], %[[VAL_37]], %[[VAL_38]], %[[VAL_39]], %[[VAL_40]] : index, i32, memref<400xi32>, memref<20xi32>, memref<20xi32>, memref<20xi32>, index, index, index)
+// CHECK:         ^bb4(%[[VAL_44:.*]]: index, %[[VAL_45:.*]]: i32, %[[VAL_46:.*]]: memref<400xi32>, %[[VAL_47:.*]]: memref<20xi32>, %[[VAL_48:.*]]: memref<20xi32>, %[[VAL_49:.*]]: memref<20xi32>, %[[VAL_50:.*]]: index, %[[VAL_51:.*]]: index, %[[VAL_52:.*]]: index, %[[VAL_53:.*]]: index, %[[VAL_54:.*]]: index):
+// CHECK:           %[[VAL_55:.*]] = arith.addi %[[VAL_44]], %[[VAL_50]] : index
+// CHECK:           %[[VAL_56:.*]] = memref.load %[[VAL_47]]{{\[}}%[[VAL_44]]] : memref<20xi32>
+// CHECK:           %[[VAL_57:.*]] = arith.index_cast %[[VAL_56]] : i32 to index
+// CHECK:           %[[VAL_58:.*]] = memref.load %[[VAL_49]]{{\[}}%[[VAL_44]]] : memref<20xi32>
+// CHECK:           %[[VAL_59:.*]] = arith.subi %[[VAL_54]], %[[VAL_50]] : index
+// CHECK:           %[[VAL_60:.*]] = memref.load %[[VAL_48]]{{\[}}%[[VAL_44]]] : memref<20xi32>
+// CHECK:           %[[VAL_61:.*]] = arith.index_cast %[[VAL_60]] : i32 to index
+// CHECK:           %[[VAL_62:.*]] = arith.muli %[[VAL_61]], %[[VAL_52]] : index
+// CHECK:           %[[VAL_63:.*]] = arith.addi %[[VAL_59]], %[[VAL_62]] : index
+// CHECK:           %[[VAL_64:.*]] = memref.load %[[VAL_46]]{{\[}}%[[VAL_63]]] : memref<400xi32>
+// CHECK:           %[[VAL_65:.*]] = arith.muli %[[VAL_58]], %[[VAL_64]] : i32
+// CHECK:           %[[VAL_66:.*]] = arith.muli %[[VAL_57]], %[[VAL_52]] : index
+// CHECK:           %[[VAL_67:.*]] = arith.addi %[[VAL_54]], %[[VAL_66]] : index
+// CHECK:           %[[VAL_68:.*]] = memref.load %[[VAL_46]]{{\[}}%[[VAL_67]]] : memref<400xi32>
+// CHECK:           %[[VAL_69:.*]] = arith.addi %[[VAL_68]], %[[VAL_65]] : i32
+// CHECK:           %[[VAL_70:.*]] = arith.muli %[[VAL_57]], %[[VAL_52]] : index
+// CHECK:           %[[VAL_71:.*]] = arith.addi %[[VAL_54]], %[[VAL_70]] : index
+// CHECK:           memref.store %[[VAL_69]], %[[VAL_46]]{{\[}}%[[VAL_71]]] : memref<400xi32>
+// CHECK:           br ^bb3(%[[VAL_55]], %[[VAL_45]], %[[VAL_46]], %[[VAL_47]], %[[VAL_48]], %[[VAL_49]], %[[VAL_50]], %[[VAL_51]], %[[VAL_52]], %[[VAL_53]], %[[VAL_54]] : index, i32, memref<400xi32>, memref<20xi32>, memref<20xi32>, memref<20xi32>, index, index, index, index, index)
+// CHECK:         ^bb5:
+// CHECK:           return
+// CHECK:         }
+func @matrix_power(%arg0: memref<400xi32>, %arg1: memref<20xi32>, %arg2: memref<20xi32>, %arg3: memref<20xi32>) {
+  %c1_i32 = arith.constant 1 : i32
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  %c20 = arith.constant 20 : index
+  br ^bb1(%c1 : index)
+^bb1(%0: index):  // 2 preds: ^bb0, ^bb3
+  %1 = arith.cmpi slt, %0, %c20 : index
+  cond_br %1, ^bb2(%0 : index), ^bb5
+^bb2(%2: index):  // pred: ^bb1
+  %3 = arith.addi %2, %c1 : index
+  %4 = arith.subi %2, %c1 : index
+  %5 = arith.index_cast %4 : index to i32
+  %6 = arith.addi %5, %c1_i32 : i32
+  %7 = arith.index_cast %6 : i32 to index
+  br ^bb3(%c0 : index)
+^bb3(%8: index):  // 2 preds: ^bb2, ^bb4
+  %9 = arith.cmpi slt, %8, %c20 : index
+  cond_br %9, ^bb4(%8 : index), ^bb1(%3 : index)
+^bb4(%10: index):  // pred: ^bb3
+  %11 = arith.addi %10, %c1 : index
+  %12 = memref.load %arg1[%10] : memref<20xi32>
+  %13 = arith.index_cast %12 : i32 to index
+  %14 = memref.load %arg3[%10] : memref<20xi32>
+  %15 = arith.subi %7, %c1 : index
+  %16 = memref.load %arg2[%10] : memref<20xi32>
+  %17 = arith.index_cast %16 : i32 to index
+  %18 = arith.muli %17, %c20 : index
+  %19 = arith.addi %15, %18 : index
+  %20 = memref.load %arg0[%19] : memref<400xi32>
+  %21 = arith.muli %14, %20 : i32
+  %22 = arith.muli %13, %c20 : index
+  %23 = arith.addi %7, %22 : index
+  %24 = memref.load %arg0[%23] : memref<400xi32>
+  %25 = arith.addi %24, %21 : i32
+  %26 = arith.muli %13, %c20 : index
+  %27 = arith.addi %7, %26 : index
+  memref.store %25, %arg0[%27] : memref<400xi32>
+  br ^bb3(%11 : index)
+^bb5:  // pred: ^bb1
+  return
+}

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestConstantFold.cpp
   TestInlining.cpp
   TestLoopFusion.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestLoopUnrolling.cpp
+  TestMaxSSAFiltered.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   LINK_LIBS PUBLIC
   MLIRTestDialect
   MLIRTransforms
   )
 
 target_include_directories(MLIRTestTransforms
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test
   ${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test
   )

mlir/test/lib/Transforms/TestMaxSSAFiltered.cpp

+//===- TestMaxSSAFiltered.cpp --- Max SSA with filtering test pass --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to test the max SSA utility with use of a
+// filtering function.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/MaxSSAUtils.h"
+#include "mlir/Transforms/Passes.h"
+
+#include "llvm/ADT/SetVector.h"
+
+using namespace mlir;
+
+namespace {
+class TestMaxSSAFilteredPass
+    : public PassWrapper<TestMaxSSAFilteredPass, FunctionPass> {
+public:
+  StringRef getArgument() const final { return "test-max-ssa-filtered"; }
+  StringRef getDescription() const final {
+    return "test the max SSA utility while filtering out certain values. The "
+           "test pass is configured to filter out memref values.";
+  }
+  explicit TestMaxSSAFilteredPass() {}
+
+  void getDependentDialects(DialectRegistry & /*registry*/) const override {}
+
+  void runOnFunction() override {
+    FuncOp func = getFunction();
+    auto filterFn = [](Value v) { return v.getType().isa<MemRefType>(); };
+
+    if (convertToMaximalSSA(func, filterFn).failed())
+      return signalPassFailure();
+  }
+};
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestMaxSSAFilteredPass() {
+  PassRegistration<TestMaxSSAFilteredPass>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestLinalgTransforms();
 void registerTestLivenessPass();
 void registerTestLoopFusion();
 void registerTestLoopMappingPass();
 void registerTestLoopUnrollingPass();
 void registerTestMatchReductionPass();
 void registerTestMathAlgebraicSimplificationPass();
 void registerTestMathPolynomialApproximationPass();
+void registerTestMaxSSAFilteredPass();
 void registerTestMemRefDependenceCheck();
 void registerTestMemRefStrideCalculation();
 void registerTestNumberOfBlockExecutionsPass();
 void registerTestNumberOfOperationExecutionsPass();
 void registerTestOpaqueLoc();
 void registerTestPDLByteCodePass();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
#endif
   mlir::test::registerTestLinalgTransforms();
   mlir::test::registerTestLivenessPass();
   mlir::test::registerTestLoopFusion();
   mlir::test::registerTestLoopMappingPass();
   mlir::test::registerTestLoopUnrollingPass();
   mlir::test::registerTestMatchReductionPass();
   mlir::test::registerTestMathAlgebraicSimplificationPass();
   mlir::test::registerTestMathPolynomialApproximationPass();
+  mlir::test::registerTestMaxSSAFilteredPass();
   mlir::test::registerTestMemRefDependenceCheck();
   mlir::test::registerTestMemRefStrideCalculation();
   mlir::test::registerTestNumberOfBlockExecutionsPass();
   mlir::test::registerTestNumberOfOperationExecutionsPass();
   mlir::test::registerTestOpaqueLoc();
   mlir::test::registerTestPDLByteCodePass();
   mlir::test::registerTestRecursiveTypesPass();
   mlir::test::registerTestSCFUtilsPass();