commit 67701712d90d62da58fe72e58f83f81d98fbc986
author Uday Bondhugula <bondhugula@google.com> Tue Aug 21 16:01:23 2018 -0700
committer jpienaar <jpienaar@google.com> Fri Mar 29 13:03:38 2019 -0700
tree 92647ea626f46843a8712a4c5df31d2a65d267ae
parent 257339b552f534615a65b0fae9fb816d3ea2f4f2

Extend loop unrolling to unroll by a given factor; add builder for affine
apply op.

- add builder for AffineApplyOp (first one for an operation that has
  non-zero operands)
- add support for loop unrolling by a given factor; uses the affine apply op
  builder.

While on this, change 'step' of ForStmt to be 'unsigned' instead of
AffineConstantExpr *. Add setters for ForStmt lb, ub, step.

Sample Input:

// CHECK-LABEL: mlfunc @loop_nest_unroll_cleanup() {
mlfunc @loop_nest_unroll_cleanup() {
  for %i = 1 to 100 {
    for %j = 0 to 17 {
      %x = "addi32"(%j, %j) : (affineint, affineint) -> i32
      %y = "addi32"(%x, %x) : (i32, i32) -> i32
    }
  }
  return
}

Output:

$ mlir-opt -loop-unroll -unroll-factor=4 /tmp/single2.mlir
#map0 = (d0) -> (d0 + 1)
#map1 = (d0) -> (d0 + 2)
#map2 = (d0) -> (d0 + 3)
mlfunc @loop_nest_unroll_cleanup() {
  for %i0 = 1 to 100 {
    for %i1 = 0 to 17 step 4 {
      %0 = "addi32"(%i1, %i1) : (affineint, affineint) -> i32
      %1 = "addi32"(%0, %0) : (i32, i32) -> i32
      %2 = affine_apply #map0(%i1)
      %3 = "addi32"(%2, %2) : (affineint, affineint) -> i32
      %4 = affine_apply #map1(%i1)
      %5 = "addi32"(%4, %4) : (affineint, affineint) -> i32
      %6 = affine_apply #map2(%i1)
      %7 = "addi32"(%6, %6) : (affineint, affineint) -> i32
    }
    for %i2 = 16 to 17 {
      %8 = "addi32"(%i2, %i2) : (affineint, affineint) -> i32
      %9 = "addi32"(%8, %8) : (i32, i32) -> i32
    }
  }
  return
}

PiperOrigin-RevId: 209676220

tools/mlir-opt/mlir-opt.cpp

diff --git a/tools/mlir-opt/mlir-opt.cpp b/tools/mlir-opt/mlir-opt.cpp
index 2c6b76b..252e1f2 100644
--- a/tools/mlir-opt/mlir-opt.cpp
+++ b/tools/mlir-opt/mlir-opt.cpp
@@ -53,8 +53,7 @@
 
 enum Passes {
   ConvertToCFG,
-  UnrollInnermostLoops,
-  UnrollShortLoops,
+  LoopUnroll,
   TFRaiseControlFlow,
 };
 
@@ -62,10 +61,7 @@
     "", cl::desc("Compiler passes to run"),
     cl::values(clEnumValN(ConvertToCFG, "convert-to-cfg",
                           "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(LoopUnroll, "loop-unroll", "Unroll loops"),
                clEnumValN(TFRaiseControlFlow, "tf-raise-control-flow",
                           "Dynamic TensorFlow Switch/Match nodes to a CFG")));
 
@@ -112,11 +108,8 @@
     case ConvertToCFG:
       pass = createConvertToCFGPass();
       break;
-    case UnrollInnermostLoops:
-      pass = createLoopUnrollPass();
-      break;
-    case UnrollShortLoops:
-      pass = createLoopUnrollPass(2);
+    case LoopUnroll:
+      pass = createLoopUnrollPass(-1, -1);
       break;
     case TFRaiseControlFlow:
       pass = createRaiseTFControlFlowPass();