Change loop step to be a positive integral constant
Changing this per discussion on mlir-team. Spec updated.
PiperOrigin-RevId: 214868483
diff --git a/include/mlir/IR/Statements.h b/include/mlir/IR/Statements.h
index 46abe6f..4fe49bd 100644
--- a/include/mlir/IR/Statements.h
+++ b/include/mlir/IR/Statements.h
@@ -251,7 +251,10 @@
void setUpperBoundMap(AffineMap *map);
/// Set loop step.
- void setStep(int64_t step) { this->step = step; }
+ void setStep(int64_t step) {
+ assert(step > 0 && "step has to be a positive integer constant");
+ this->step = step;
+ }
/// Returns true if the lower bound is constant.
bool hasConstantLowerBound() const;
@@ -342,7 +345,8 @@
AffineMap *lbMap;
// Affine map for the upper bound.
AffineMap *ubMap;
- // Constant step.
+ // Positive constant step. Since affineint is int64_t, we restrict step to the
+ // set of positive integers that int64_t can represent.
int64_t step;
// Operands for the lower and upper bounds, with the former followed by the
// latter. Dimensional operands are followed by symbolic operands for each
diff --git a/lib/Analysis/LoopAnalysis.cpp b/lib/Analysis/LoopAnalysis.cpp
index fa283ae..f4e607e 100644
--- a/lib/Analysis/LoopAnalysis.cpp
+++ b/lib/Analysis/LoopAnalysis.cpp
@@ -66,13 +66,12 @@
auto *cExpr = dyn_cast<AffineConstantExpr>(loopSpanExpr);
if (!cExpr)
- return AffineBinaryOpExpr::getCeilDiv(loopSpanExpr, std::abs(step),
- context);
+ return AffineBinaryOpExpr::getCeilDiv(loopSpanExpr, step, context);
loopSpan = cExpr->getValue();
}
// 0 iteration loops.
- if ((loopSpan < 0 && step >= 1) || (loopSpan > 0 && step <= -1))
+ if (loopSpan < 0)
return 0;
return AffineConstantExpr::get(
diff --git a/lib/IR/Statement.cpp b/lib/IR/Statement.cpp
index c4eb5b8..7834d32 100644
--- a/lib/IR/Statement.cpp
+++ b/lib/IR/Statement.cpp
@@ -268,6 +268,7 @@
"lower bound operand count does not match the affine map");
assert(ubOperands.size() == ubMap->getNumInputs() &&
"upper bound operand count does not match the affine map");
+ assert(step > 0 && "step has to be a positive integer constant");
unsigned numOperands = lbOperands.size() + ubOperands.size();
ForStmt *stmt =
diff --git a/lib/Parser/Parser.cpp b/lib/Parser/Parser.cpp
index f9c801c..54ddb7a 100644
--- a/lib/Parser/Parser.cpp
+++ b/lib/Parser/Parser.cpp
@@ -2277,6 +2277,13 @@
if (consumeIf(Token::kw_step) && parseIntConstant(step))
return ParseFailure;
+ // The loop step is a positive integer constant. Since affineint is of int64_t
+ // type, we restrict step to be in the set of positive integers that int64_t
+ // can represent.
+ if (step < 1) {
+ return emitError("step has to be a positive integer");
+ }
+
// Create for statement.
ForStmt *forStmt =
builder.createFor(getEncodedSourceLocation(loc), lbOperands, lbMap,
diff --git a/test/IR/invalid.mlir b/test/IR/invalid.mlir
index 05557af..3e962f5 100644
--- a/test/IR/invalid.mlir
+++ b/test/IR/invalid.mlir
@@ -191,6 +191,12 @@
// -----
+mlfunc @for_negative_stride() {
+ for %i = 1 to 10 step -1
+} // expected-error {{step has to be a positive integer}}
+
+// -----
+
mlfunc @non_statement() {
asd // expected-error {{custom op 'asd' is unknown}}
}
diff --git a/test/IR/parser.mlir b/test/IR/parser.mlir
index c8c2237..49a7cad 100644
--- a/test/IR/parser.mlir
+++ b/test/IR/parser.mlir
@@ -214,8 +214,8 @@
%s = "foo"(%N) : (affineint) -> affineint
// CHECK: for %i0 = %0 to %arg0
for %i = %s to %N {
- // CHECK: for %i1 = #map1(%i0) to 0 step -1
- for %j = %i to 0 step -1 {
+ // CHECK: for %i1 = #map1(%i0) to 0
+ for %j = %i to 0 step 1 {
// CHECK: %1 = affine_apply #map{{.*}}(%i0, %i1)[%0]
%w = affine_apply(d0, d1)[s0] -> (d0+d1, s0+1) (%i, %j) [%s]
// CHECK: for %i2 = #map{{.*}}(%1#0, %i0)[%arg0] to #map{{.*}}(%1#1, %i1)[%0] {