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

lib/IR/AsmPrinter.cpp

diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index 8303f57..6bcaad4 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -1316,8 +1316,8 @@
   printOperand(stmt);
   os << " = " << *stmt->getLowerBound();
   os << " to " << *stmt->getUpperBound();
-  if (stmt->getStep()->getValue() != 1)
-    os << " step " << *stmt->getStep();
+  if (stmt->getStep() != 1)
+    os << " step " << stmt->getStep();
 
   os << " {\n";
   print(static_cast<const StmtBlock *>(stmt));

lib/IR/Builders.cpp

diff --git a/lib/IR/Builders.cpp b/lib/IR/Builders.cpp
index 5fce94c..b1dce25 100644
--- a/lib/IR/Builders.cpp
+++ b/lib/IR/Builders.cpp
@@ -103,8 +103,8 @@
   return ArrayAttr::get(value, context);
 }
 
-AffineMapAttr *Builder::getAffineMapAttr(AffineMap *value) {
-  return AffineMapAttr::get(value, context);
+AffineMapAttr *Builder::getAffineMapAttr(AffineMap *map) {
+  return AffineMapAttr::get(map, context);
 }
 
 TypeAttr *Builder::getTypeAttr(Type *type) {
@@ -207,9 +207,7 @@
 
 ForStmt *MLFuncBuilder::createFor(AffineConstantExpr *lowerBound,
                                   AffineConstantExpr *upperBound,
-                                  AffineConstantExpr *step) {
-  if (!step)
-    step = getConstantExpr(1);
+                                  int64_t step) {
   auto *stmt = new ForStmt(lowerBound, upperBound, step, context);
   block->getStatements().insert(insertPoint, stmt);
   return stmt;

lib/IR/StandardOps.cpp

diff --git a/lib/IR/StandardOps.cpp b/lib/IR/StandardOps.cpp
index 6d8c366..cec23e0 100644
--- a/lib/IR/StandardOps.cpp
+++ b/lib/IR/StandardOps.cpp
@@ -305,6 +305,22 @@
 }
 
 //===----------------------------------------------------------------------===//
+// AffineApplyOp
+//===----------------------------------------------------------------------===//
+
+OperationState AffineApplyOp::build(Builder *builder, AffineMap *map,
+                                    ArrayRef<SSAValue *> operands) {
+  SmallVector<Type *, 4> resultTypes(map->getNumResults(),
+                                     builder->getAffineIntType());
+
+  OperationState result(
+      builder->getIdentifier("affine_apply"), operands, resultTypes,
+      {{builder->getIdentifier("map"), builder->getAffineMapAttr(map)}});
+
+  return result;
+}
+
+//===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 

lib/IR/Statement.cpp

diff --git a/lib/IR/Statement.cpp b/lib/IR/Statement.cpp
index d20b45f..7da08c2 100644
--- a/lib/IR/Statement.cpp
+++ b/lib/IR/Statement.cpp
@@ -198,7 +198,7 @@
 //===----------------------------------------------------------------------===//
 
 ForStmt::ForStmt(AffineConstantExpr *lowerBound, AffineConstantExpr *upperBound,
-                 AffineConstantExpr *step, MLIRContext *context)
+                 int64_t step, MLIRContext *context)
     : Statement(Kind::For),
       MLValue(MLValueKind::ForStmt, Type::getAffineInt(context)),
       StmtBlock(StmtBlockKind::For), lowerBound(lowerBound),