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
diff --git a/include/mlir/IR/Builders.h b/include/mlir/IR/Builders.h
index 62c0dcb..e8d3894 100644
--- a/include/mlir/IR/Builders.h
+++ b/include/mlir/IR/Builders.h
@@ -252,7 +252,8 @@
this->insertPoint = insertPoint;
}
- /// Set the insertion point to the specified operation.
+ /// Set the insertion point to the specified operation, which will cause
+ /// subsequent insertions to go right before it.
void setInsertionPoint(Statement *stmt) {
setInsertionPoint(stmt->getBlock(), StmtBlock::iterator(stmt));
}
@@ -298,8 +299,7 @@
// Creates for statement. When step is not specified, it is set to 1.
ForStmt *createFor(AffineConstantExpr *lowerBound,
- AffineConstantExpr *upperBound,
- AffineConstantExpr *step = nullptr);
+ AffineConstantExpr *upperBound, int64_t step = 1);
IfStmt *createIf(IntegerSet *condition) {
auto *stmt = new IfStmt(condition);
diff --git a/include/mlir/IR/Operation.h b/include/mlir/IR/Operation.h
index fed0d4e..af01a1d 100644
--- a/include/mlir/IR/Operation.h
+++ b/include/mlir/IR/Operation.h
@@ -47,6 +47,7 @@
struct OperationState {
Identifier name;
SmallVector<SSAValue *, 4> operands;
+ /// Types of the results of this operation.
SmallVector<Type *, 4> types;
SmallVector<NamedAttribute, 4> attributes;
diff --git a/include/mlir/IR/StandardOps.h b/include/mlir/IR/StandardOps.h
index 3b615cc..63921ca 100644
--- a/include/mlir/IR/StandardOps.h
+++ b/include/mlir/IR/StandardOps.h
@@ -77,6 +77,10 @@
class AffineApplyOp : public OpBase<AffineApplyOp, OpTrait::VariadicOperands,
OpTrait::VariadicResults> {
public:
+ /// Builds an affine apply op with the specified map and operands.
+ static OperationState build(Builder *builder, AffineMap *map,
+ ArrayRef<SSAValue *> operands);
+
// Returns the affine map to be applied by this operation.
AffineMap *getAffineMap() const {
return getAttrOfType<AffineMapAttr>("map")->getValue();
@@ -163,6 +167,7 @@
///
class ConstantFloatOp : public ConstantOp {
public:
+ /// Builds a constant float op producing a float of the specified type.
static OperationState build(Builder *builder, double value, FloatType *type);
double getValue() const {
diff --git a/include/mlir/IR/Statements.h b/include/mlir/IR/Statements.h
index f851d9f..1a68aab 100644
--- a/include/mlir/IR/Statements.h
+++ b/include/mlir/IR/Statements.h
@@ -199,7 +199,7 @@
// TODO: lower and upper bounds should be affine maps with
// dimension and symbol use lists.
explicit ForStmt(AffineConstantExpr *lowerBound,
- AffineConstantExpr *upperBound, AffineConstantExpr *step,
+ AffineConstantExpr *upperBound, int64_t step,
MLIRContext *context);
~ForStmt() {
@@ -216,7 +216,11 @@
AffineConstantExpr *getLowerBound() const { return lowerBound; }
AffineConstantExpr *getUpperBound() const { return upperBound; }
- AffineConstantExpr *getStep() const { return step; }
+ int64_t getStep() const { return step; }
+
+ void setLowerBound(AffineConstantExpr *lb) { lowerBound = lb; }
+ void setUpperBound(AffineConstantExpr *ub) { upperBound = ub; }
+ void setStep(unsigned s) { step = s; }
using Statement::dump;
using Statement::print;
@@ -242,7 +246,7 @@
// an affinemap and its operands as AffineBound.
AffineConstantExpr *lowerBound;
AffineConstantExpr *upperBound;
- AffineConstantExpr *step;
+ int64_t step;
};
/// An if clause represents statements contained within a then or an else clause