Implement operands for the 'if' statement.
This CL also includes two other minor changes:
- change the implemented syntax from 'if (cond)' to 'if cond', as specified by MLIR spec.
- a minor fix to the implementation of the ForStmt.
PiperOrigin-RevId: 210618122
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index 2cc20ac..21dd4c6 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -186,7 +186,7 @@
}
void ModuleState::visitIfStmt(const IfStmt *ifStmt) {
- recordIntegerSetReference(ifStmt->getCondition());
+ recordIntegerSetReference(ifStmt->getIntegerSet());
for (auto &childStmt : *ifStmt->getThen())
visitStatement(&childStmt);
if (ifStmt->hasElse())
@@ -1406,9 +1406,11 @@
}
void MLFunctionPrinter::print(const IfStmt *stmt) {
- os.indent(numSpaces) << "if (";
- printIntegerSetReference(stmt->getCondition());
- os << ") {\n";
+ os.indent(numSpaces) << "if ";
+ IntegerSet *set = stmt->getIntegerSet();
+ printIntegerSetReference(set);
+ printDimAndSymbolList(stmt->getStmtOperands(), set->getNumDims());
+ os << " {\n";
print(stmt->getThen());
os.indent(numSpaces) << "}";
if (stmt->hasElse()) {
diff --git a/lib/IR/Builders.cpp b/lib/IR/Builders.cpp
index 8d0f223..a03de2f 100644
--- a/lib/IR/Builders.cpp
+++ b/lib/IR/Builders.cpp
@@ -254,8 +254,9 @@
return stmt;
}
-IfStmt *MLFuncBuilder::createIf(Location *location, IntegerSet *condition) {
- auto *stmt = new IfStmt(location, condition);
+IfStmt *MLFuncBuilder::createIf(Location *location,
+ ArrayRef<MLValue *> operands, IntegerSet *set) {
+ auto *stmt = IfStmt::create(location, operands, set);
block->getStatements().insert(insertPoint, stmt);
return stmt;
}
diff --git a/lib/IR/Statement.cpp b/lib/IR/Statement.cpp
index 3f95a94..673aa43 100644
--- a/lib/IR/Statement.cpp
+++ b/lib/IR/Statement.cpp
@@ -15,7 +15,9 @@
// limitations under the License.
// =============================================================================
+#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/IntegerSet.h"
#include "mlir/IR/MLFunction.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/StandardOps.h"
@@ -115,8 +117,7 @@
case Kind::For:
return cast<ForStmt>(this)->getNumOperands();
case Kind::If:
- // TODO: query IfStmt once it has operands.
- return 0;
+ return cast<IfStmt>(this)->getNumOperands();
}
}
@@ -127,8 +128,7 @@
case Kind::For:
return cast<ForStmt>(this)->getStmtOperands();
case Kind::If:
- // TODO: query IfStmt once it has operands.
- return {};
+ return cast<IfStmt>(this)->getStmtOperands();
}
}
@@ -301,7 +301,9 @@
AffineMap *ubMap, int64_t step, MLIRContext *context)
: Statement(Kind::For, location),
MLValue(MLValueKind::ForStmt, Type::getAffineInt(context)),
- StmtBlock(StmtBlockKind::For), lbMap(lbMap), ubMap(ubMap), step(step) {}
+ StmtBlock(StmtBlockKind::For), lbMap(lbMap), ubMap(ubMap), step(step) {
+ operands.reserve(numOperands);
+}
const AffineBound ForStmt::getLowerBound() const {
return AffineBound(*this, 0, lbMap->getNumOperands(), lbMap);
@@ -357,18 +359,47 @@
// IfStmt
//===----------------------------------------------------------------------===//
-IfStmt::IfStmt(Location *location, IntegerSet *condition)
+IfStmt::IfStmt(Location *location, unsigned numOperands, IntegerSet *set)
: Statement(Kind::If, location), thenClause(new IfClause(this)),
- elseClause(nullptr), condition(condition) {}
+ elseClause(nullptr), set(set) {
+ operands.reserve(numOperands);
+}
IfStmt::~IfStmt() {
delete thenClause;
if (elseClause)
delete elseClause;
- // An IfStmt's IntegerSet 'condition' should not be deleted since it is
+ // An IfStmt's IntegerSet 'set' should not be deleted since it is
// allocated through MLIRContext's bump pointer allocator.
}
+IfStmt *IfStmt::create(Location *location, ArrayRef<MLValue *> operands,
+ IntegerSet *set) {
+ unsigned numOperands = operands.size();
+ assert(numOperands == set->getNumOperands() &&
+ "operand cound does not match the integer set operand count");
+
+ IfStmt *stmt = new IfStmt(location, numOperands, set);
+
+ for (auto *op : operands)
+ stmt->operands.emplace_back(StmtOperand(stmt, op));
+
+ return stmt;
+}
+
+const AffineCondition IfStmt::getCondition() const {
+ return AffineCondition(*this, set);
+}
+
+MLIRContext *IfStmt::getContext() const {
+ // Check for degenerate case of if statement with no operands.
+ // This is unlikely, but legal.
+ if (operands.empty())
+ return findFunction()->getContext();
+
+ return getOperand(0)->getType()->getContext();
+}
+
//===----------------------------------------------------------------------===//
// Statement Cloning
//===----------------------------------------------------------------------===//
@@ -428,9 +459,7 @@
// Otherwise, we must have an If statement.
auto *ifStmt = cast<IfStmt>(this);
- auto *newIf = new IfStmt(getLoc(), ifStmt->getCondition());
-
- // TODO: remap operands with remapOperand when if statements have them.
+ auto *newIf = IfStmt::create(getLoc(), operands, ifStmt->getIntegerSet());
auto *resultThen = newIf->getThen();
for (auto &childStmt : *ifStmt->getThen())
diff --git a/lib/IR/Verifier.cpp b/lib/IR/Verifier.cpp
index deda932..48f3643 100644
--- a/lib/IR/Verifier.cpp
+++ b/lib/IR/Verifier.cpp
@@ -372,7 +372,7 @@
// TODO: check that operation is not a return statement unless it's
// the last one in the function.
- // TODO: check that loop bounds are properly formed.
+ // TODO: check that loop bounds and if conditions are properly formed.
if (verifyReturn())
return true;