1 files changed, 151 insertions, 0 deletions
diff --git a/mlir/lib/AffineOps/AffineOps.cpp b/mlir/lib/AffineOps/AffineOps.cpp
new file mode 100644
index 00000000000..5b29467fc44
--- /dev/null
+++ b/mlir/lib/AffineOps/AffineOps.cpp
@@ -0,0 +1,151 @@
+//===- AffineOps.cpp - MLIR Affine Operations -----------------------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+
+#include "mlir/AffineOps/AffineOps.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/OpImplementation.h"
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// AffineOpsDialect
+//===----------------------------------------------------------------------===//
+
+AffineOpsDialect::AffineOpsDialect(MLIRContext *context)
+    : Dialect(/*namePrefix=*/"", context) {
+  addOperations<AffineIfOp>();
+}
+
+//===----------------------------------------------------------------------===//
+// AffineIfOp
+//===----------------------------------------------------------------------===//
+
+void AffineIfOp::build(Builder *builder, OperationState *result,
+                       IntegerSet condition,
+                       ArrayRef<Value *> conditionOperands) {
+  result->addAttribute(getConditionAttrName(), IntegerSetAttr::get(condition));
+  result->addOperands(conditionOperands);
+
+  // Reserve 2 block lists, one for the 'then' and one for the 'else' regions.
+  result->reserveBlockLists(2);
+}
+
+bool AffineIfOp::verify() const {
+  // Verify that we have a condition attribute.
+  auto conditionAttr = getAttrOfType<IntegerSetAttr>(getConditionAttrName());
+  if (!conditionAttr)
+    return emitOpError("requires an integer set attribute named 'condition'");
+
+  // Verify that the operands are valid dimension/symbols.
+  IntegerSet condition = conditionAttr.getValue();
+  for (unsigned i = 0, e = getNumOperands(); i != e; ++i) {
+    const Value *operand = getOperand(i);
+    if (i < condition.getNumDims() && !operand->isValidDim())
+      return emitOpError("operand cannot be used as a dimension id");
+    if (i >= condition.getNumDims() && !operand->isValidSymbol())
+      return emitOpError("operand cannot be used as a symbol");
+  }
+
+  // Verify that the entry of each child blocklist does not have arguments.
+  for (const auto &blockList : getInstruction()->getBlockLists()) {
+    if (blockList.empty())
+      continue;
+
+    // TODO(riverriddle) We currently do not allow multiple blocks in child
+    // block lists.
+    if (std::next(blockList.begin()) != blockList.end())
+      return emitOpError(
+          "expects only one block per 'if' or 'else' block list");
+    if (blockList.front().getTerminator())
+      return emitOpError("expects region block to not have a terminator");
+
+    for (const auto &b : blockList)
+      if (b.getNumArguments() != 0)
+        return emitOpError(
+            "requires that child entry blocks have no arguments");
+  }
+  return false;
+}
+
+bool AffineIfOp::parse(OpAsmParser *parser, OperationState *result) {
+  // Parse the condition attribute set.
+  IntegerSetAttr conditionAttr;
+  unsigned numDims;
+  if (parser->parseAttribute(conditionAttr, getConditionAttrName().data(),
+                             result->attributes) ||
+      parseDimAndSymbolList(parser, result->operands, numDims))
+    return true;
+
+  // Verify the condition operands.
+  auto set = conditionAttr.getValue();
+  if (set.getNumDims() != numDims)
+    return parser->emitError(
+        parser->getNameLoc(),
+        "dim operand count and integer set dim count must match");
+  if (numDims + set.getNumSymbols() != result->operands.size())
+    return parser->emitError(
+        parser->getNameLoc(),
+        "symbol operand count and integer set symbol count must match");
+
+  // Parse the 'then' block list.
+  if (parser->parseBlockList())
+    return true;
+
+  // If we find an 'else' keyword then parse the else block list.
+  if (!parser->parseOptionalKeyword("else")) {
+    if (parser->parseBlockList())
+      return true;
+  }
+
+  // Reserve 2 block lists, one for the 'then' and one for the 'else' regions.
+  result->reserveBlockLists(2);
+  return false;
+}
+
+void AffineIfOp::print(OpAsmPrinter *p) const {
+  auto conditionAttr = getAttrOfType<IntegerSetAttr>(getConditionAttrName());
+  *p << "if " << conditionAttr;
+  printDimAndSymbolList(operand_begin(), operand_end(),
+                        conditionAttr.getValue().getNumDims(), p);
+  p->printBlockList(getInstruction()->getBlockList(0));
+
+  // Print the 'else' block list if it has any blocks.
+  const auto &elseBlockList = getInstruction()->getBlockList(1);
+  if (!elseBlockList.empty()) {
+    *p << " else";
+    p->printBlockList(elseBlockList);
+  }
+}
+
+IntegerSet AffineIfOp::getIntegerSet() const {
+  return getAttrOfType<IntegerSetAttr>(getConditionAttrName()).getValue();
+}
+void AffineIfOp::setIntegerSet(IntegerSet newSet) {
+  setAttr(
+      Identifier::get(getConditionAttrName(), getInstruction()->getContext()),
+      IntegerSetAttr::get(newSet));
+}
+
+/// Returns the list of 'then' blocks.
+BlockList &AffineIfOp::getThenBlocks() {
+  return getInstruction()->getBlockList(0);
+}
+
+/// Returns the list of 'else' blocks.
+BlockList &AffineIfOp::getElseBlocks() {
+  return getInstruction()->getBlockList(1);
+}