Loop unrolling pass update
- fix/complete forStmt cloning for unrolling to work for outer loops
- create IV const's only when needed
- test outer loop unrolling by creating a short trip count unroll pass for
loops with trip counts <= <parameter>
- add unrolling test cases for multiple op results, outer loop unrolling
- fix/clean up StmtWalker class while on this
- switch unroll loop iterator values from i32 to affineint
PiperOrigin-RevId: 207645967
diff --git a/tools/mlir-opt/mlir-opt.cpp b/tools/mlir-opt/mlir-opt.cpp
index c95b838..26417e2 100644
--- a/tools/mlir-opt/mlir-opt.cpp
+++ b/tools/mlir-opt/mlir-opt.cpp
@@ -54,6 +54,7 @@
enum Passes {
ConvertToCFG,
UnrollInnermostLoops,
+ UnrollShortLoops,
TFRaiseControlFlow,
};
@@ -63,6 +64,8 @@
"Convert all ML functions in the module to CFG ones"),
clEnumValN(UnrollInnermostLoops, "unroll-innermost-loops",
"Unroll innermost loops"),
+ clEnumValN(UnrollShortLoops, "unroll-short-loops",
+ "Unroll loops of trip count <= 2"),
clEnumValN(TFRaiseControlFlow, "tf-raise-control-flow",
"Dynamic TensorFlow Switch/Match nodes to a CFG")));
@@ -111,6 +114,9 @@
case UnrollInnermostLoops:
pass = createLoopUnrollPass();
break;
+ case UnrollShortLoops:
+ pass = createLoopUnrollPass(2);
+ break;
case TFRaiseControlFlow:
pass = createRaiseTFControlFlowPass();
break;