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;