Moves the linalg_ext dialect to iree_linalg_ext under the iree-dialects project. (#7657)
Mostly NFC:
* Pre-factors iree-dialects into a consistent state wrt namespaces and directory layouts.
* Moves linalg_ext to a new dialect under iree-dialects. Does some adaptation to upstream style along the way.
* Redirects everything that was using it to use the new one.
* Enables tests for iree-dialects in the cmake CI (they were not enabled) and fixes some type prefixing that had drifted. (Will follow up with enabling them in the internal builds to guard against further regression).
Non-NFC:
*When tiling, the old pass was directly generating flow.dispatch.workgroup... ops to query the current workgroup id. We had been planning to add those to the input dialect, so I pulled part of that patch forward.
* Since these now newly lack lowerings, I was expecting to hit some test failures that would guide me to where to adapt, but this appears to be dead code (outside of integration tests?). We'll see what the CI says.
* Can keep coding on this patch to adapt whatever is needed.
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/PyDM/IR/PyDMOps.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/PyDM/IR/PyDMOps.cpp
new file mode 100644
index 0000000..7ef7bfc
--- /dev/null
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/PyDM/IR/PyDMOps.cpp
@@ -0,0 +1,810 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree-dialects/Dialect/PyDM/IR/PyDMOps.h"
+
+#include "iree-dialects/Dialect/PyDM/IR/PyDMDialect.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/FunctionImplementation.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/TypeUtilities.h"
+
+using namespace mlir;
+namespace PYDM = mlir::iree_compiler::IREE::PYDM;
+using namespace PYDM;
+
+using llvm::dbgs;
+
+using PyBoolType = PYDM::BoolType;
+using PyConstantOp = PYDM::ConstantOp;
+using PyIntegerType = PYDM::IntegerType;
+using PyRealType = PYDM::RealType;
+using PyCallOp = PYDM::CallOp;
+using PyFuncOp = PYDM::FuncOp;
+
+static LogicalResult verify(Operation *) { return success(); }
+
+//===----------------------------------------------------------------------===//
+// Utilities
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Generic pattern to unbox any operands that are a specific object
+/// type (i.e. object<integer>).
+struct UnboxOperands : public RewritePattern {
+ UnboxOperands(StringRef rootName, MLIRContext *context,
+ Optional<llvm::SmallSet<int, 4>> operandIndices = None)
+ : RewritePattern(rootName, 1, context), operandIndices(operandIndices) {}
+ LogicalResult matchAndRewrite(Operation *op,
+ PatternRewriter &rewriter) const override {
+ Location loc = op->getLoc();
+ bool changed = false;
+ SmallVector<Value> operands(op->getOperands());
+ auto excResultType = rewriter.getType<ExceptionResultType>();
+ for (int operandIndex = 0, e = operands.size(); operandIndex < e;
+ ++operandIndex) {
+ Value &operand = operands[operandIndex];
+ if (operandIndices && !operandIndices->contains(operandIndex)) continue;
+ if (auto objectType = operand.getType().dyn_cast<ObjectType>()) {
+ Type primitiveType = objectType.getPrimitiveType();
+ if (primitiveType) {
+ // Unbox.
+ auto unboxOp = rewriter.create<UnboxOp>(
+ loc, TypeRange{excResultType, primitiveType}, operand);
+ operand = unboxOp.primitive();
+ changed = true;
+ }
+ }
+ }
+
+ if (changed) {
+ rewriter.updateRootInPlace(op, [&]() { op->setOperands(operands); });
+ return success();
+ }
+
+ return failure();
+ }
+ Optional<llvm::SmallSet<int, 4>> operandIndices;
+};
+
+} // namespace
+
+static Value getNumericZeroConstant(Location loc, Type numericType,
+ OpBuilder &builder) {
+ return TypeSwitch<Type, Value>(numericType)
+ .Case([&](PyBoolType t) -> Value {
+ return builder.create<PyConstantOp>(loc, t, builder.getBoolAttr(false));
+ })
+ .Case([&](PyIntegerType t) -> Value {
+ return builder.create<PyConstantOp>(loc, t,
+ builder.getI64IntegerAttr(0));
+ })
+ .Case([&](PyRealType t) -> Value {
+ return builder.create<PyConstantOp>(loc, t,
+ builder.getF64FloatAttr(0.0));
+ });
+}
+
+static Value getBoolConstant(Location loc, bool pred, OpBuilder &builder) {
+ return builder.create<PyConstantOp>(loc, builder.getType<BoolType>(),
+ builder.getBoolAttr(pred));
+}
+
+//===----------------------------------------------------------------------===//
+// Constants
+//===----------------------------------------------------------------------===//
+
+OpFoldResult PyConstantOp::fold(ArrayRef<Attribute> operands) {
+ assert(operands.empty() && "constant has no operands");
+ return getValue();
+}
+
+OpFoldResult NoneOp::fold(ArrayRef<Attribute> operands) {
+ assert(operands.empty() && "constant has no operands");
+ return UnitAttr::get(getContext());
+}
+
+OpFoldResult SuccessOp::fold(ArrayRef<Attribute> operands) {
+ assert(operands.empty() && "constant has no operands");
+ return UnitAttr::get(getContext());
+}
+
+//===----------------------------------------------------------------------===//
+// Variables
+//===----------------------------------------------------------------------===//
+
+void AllocFreeVarOp::getAsmResultNames(OpAsmSetValueNameFn setNameFn) {
+ setNameFn(getResult(), name());
+}
+
+//===----------------------------------------------------------------------===//
+// ApplyBinaryOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct ApplyBinaryToSequenceClone : public OpRewritePattern<ApplyBinaryOp> {
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(ApplyBinaryOp op,
+ PatternRewriter &rewriter) const override {
+ if (op.dunder_name() != "mul") return failure();
+ Value listOperand;
+ Value countOperand;
+ if (isBuiltinSequence(op.left()) && isInteger(op.right())) {
+ listOperand = op.left();
+ countOperand = op.right();
+ } else if (isInteger(op.left()) && isBuiltinSequence(op.right())) {
+ countOperand = op.left();
+ listOperand = op.right();
+ } else {
+ return failure();
+ }
+ Type resultType = op.getResult().getType();
+ rewriter.replaceOpWithNewOp<SequenceCloneOp>(op, resultType, listOperand,
+ countOperand);
+ return success();
+ }
+
+ static bool isBuiltinSequence(Value operand) {
+ return operand.getType().isa<PYDM::ListType, PYDM::TupleType>();
+ }
+ static bool isInteger(Value operand) {
+ return operand.getType().isa<PYDM::IntegerType>();
+ }
+};
+} // namespace
+
+void ApplyBinaryOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+ MLIRContext *context) {
+ patterns.add<UnboxOperands>(getOperationName(), context);
+ patterns.add<ApplyBinaryToSequenceClone>(context);
+}
+
+bool ApplyBinaryOp::refineResultTypes() {
+ auto leftType = left().getType();
+ auto rightType = right().getType();
+ auto applyUpdates = [&](Type newResultType) -> bool {
+ if (newResultType != getResult().getType()) {
+ getResult().setType(newResultType);
+ return true;
+ }
+ return false;
+ };
+
+ // Both numeric types. It is only dynamically legal for statically known
+ // numeric types to be the same, in which case the result type must be the
+ // same as well.
+ auto ptLeft = leftType.dyn_cast<PythonTypeInterface>();
+ auto ptRight = rightType.dyn_cast<PythonTypeInterface>();
+ if (ptLeft && ptRight && ptLeft.getNumericPromotionOrder() &&
+ ptRight.getNumericPromotionOrder()) {
+ if (leftType == rightType) {
+ return applyUpdates(leftType);
+ }
+ }
+
+ // (list, integer) or (integer, list) refine to the list type.
+ if (dunder_name() == "mul") {
+ auto leftList = leftType.dyn_cast<ListType>();
+ auto rightList = rightType.dyn_cast<ListType>();
+ auto leftInteger = leftType.dyn_cast<IntegerType>();
+ auto rightInteger = rightType.dyn_cast<IntegerType>();
+ if (leftList && rightInteger) {
+ return applyUpdates(leftList);
+ } else if (leftInteger && rightList) {
+ return applyUpdates(rightList);
+ }
+ }
+
+ return false;
+}
+
+//===----------------------------------------------------------------------===//
+// ApplyCompareOp
+//===----------------------------------------------------------------------===//
+
+void ApplyCompareOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+ MLIRContext *context) {
+ patterns.add<UnboxOperands>(getOperationName(), context);
+}
+
+//===----------------------------------------------------------------------===//
+// AsBoolOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct FoldAsBoolFromBool : public OpRewritePattern<AsBoolOp> {
+ public:
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(AsBoolOp op,
+ PatternRewriter &rewriter) const override {
+ if (op.value().getType().isa<BoolType>()) {
+ rewriter.replaceOp(op, op.value());
+ return success();
+ }
+ return failure();
+ }
+};
+
+struct FoldAsBoolFromNumeric : public OpRewritePattern<AsBoolOp> {
+ public:
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(AsBoolOp op,
+ PatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto ptType = op.value().getType().dyn_cast<PythonTypeInterface>();
+ if (!ptType) return failure();
+ if (!ptType.getNumericPromotionOrder()) return failure();
+
+ auto boolType = rewriter.getType<BoolType>();
+ Value zeroValue =
+ getNumericZeroConstant(loc, op.value().getType(), rewriter);
+ Value trueValue = getBoolConstant(loc, true, rewriter);
+ Value falseValue = getBoolConstant(loc, false, rewriter);
+ Value cmpResult = rewriter.create<ApplyCompareOp>(
+ loc, boolType, rewriter.getStringAttr("eq"), op.value(), zeroValue);
+ rewriter.replaceOpWithNewOp<SelectOp>(op, boolType, cmpResult, falseValue,
+ trueValue);
+ return success();
+ }
+};
+
+} // namespace
+
+void AsBoolOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+ MLIRContext *context) {
+ patterns.add<FoldAsBoolFromBool, FoldAsBoolFromNumeric>(context);
+}
+
+OpFoldResult AsBoolOp::fold(ArrayRef<Attribute> operands) {
+ Builder b(getContext());
+ // Fold NoneType to False.
+ if (value().getType().isa<NoneType>()) {
+ return b.getBoolAttr(false);
+ }
+
+ return {};
+}
+
+//===----------------------------------------------------------------------===//
+// AssignSubscriptOp
+//===----------------------------------------------------------------------===//
+
+void AssignSubscriptOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+ MLIRContext *context) {
+ llvm::SmallSet<int, 4> unboxIndices;
+ unboxIndices.insert(0);
+ unboxIndices.insert(1);
+ patterns.add<UnboxOperands>(getOperationName(), context, unboxIndices);
+}
+
+//===----------------------------------------------------------------------===//
+// BoolToPredOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult BoolToPredOp::fold(ArrayRef<Attribute> operands) {
+ if (!operands[0]) return {};
+ // Since both BoolType and I1 share the attribute form (an IntegerAttr of I1),
+ // we can just return it.
+ return operands[0];
+}
+
+//===----------------------------------------------------------------------===//
+// BoxOp and UnboxOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult BoxOp::canonicalize(BoxOp op, PatternRewriter &rewriter) {
+ // Sometimes boxes are emitted when the input is an object. Just remove.
+ if (op.primitive().getType().isa<ObjectType>()) {
+ rewriter.replaceOp(op, op.primitive());
+ return success();
+ }
+
+ // Box to an appropriate type and static info cast.
+ ObjectType objectType = rewriter.getType<ObjectType>(nullptr);
+ if (op.object().getType() == objectType &&
+ !op.primitive().getType().isa<ObjectType>()) {
+ auto refinedBox = rewriter.create<BoxOp>(
+ op.getLoc(),
+ rewriter.getType<ObjectType>(
+ op.primitive().getType().cast<PrimitiveType>()),
+ op.primitive());
+ rewriter.replaceOpWithNewOp<StaticInfoCastOp>(op, op.object().getType(),
+ refinedBox);
+ return success();
+ }
+
+ return failure();
+}
+
+LogicalResult UnboxOp::canonicalize(UnboxOp unboxOp,
+ PatternRewriter &rewriter) {
+ auto loc = unboxOp.getLoc();
+
+ // Handle the case of an immediate BoxOp producer.
+ if (auto boxProducer =
+ dyn_cast_or_null<BoxOp>(unboxOp.object().getDefiningOp())) {
+ // If the producer is boxing to the same type we are unboxing, then
+ // just elide everything.
+ if (boxProducer.primitive().getType() == unboxOp.primitive().getType()) {
+ auto successValue = rewriter.create<SuccessOp>(
+ loc, rewriter.getType<ExceptionResultType>());
+ rewriter.replaceOp(unboxOp, {successValue, boxProducer.primitive()});
+ return success();
+ }
+ }
+ return failure();
+}
+
+//===----------------------------------------------------------------------===//
+// DynamicBinaryPromoteOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Resolves a DynamicBinaryPromote over numeric operands to either elide
+/// or insert specific PromoteNumeric ops.
+struct ResolveNumericDynamicBinaryPromote
+ : public OpRewritePattern<DynamicBinaryPromoteOp> {
+ public:
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(DynamicBinaryPromoteOp op,
+ PatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto leftType = op.left().getType();
+ auto rightType = op.right().getType();
+ auto leftResultType = op.getResultTypes()[0];
+ auto rightResultType = op.getResultTypes()[1];
+ auto leftPt = leftType.dyn_cast<PythonTypeInterface>();
+ auto rightPt = rightType.dyn_cast<PythonTypeInterface>();
+ if (!leftPt || !rightPt) return failure();
+
+ Optional<int> leftOrder = leftPt.getNumericPromotionOrder();
+ Optional<int> rightOrder = rightPt.getNumericPromotionOrder();
+ Value newLeft = op.left();
+ Value newRight = op.right();
+
+ if (leftOrder && rightOrder) {
+ // Both numeric.
+ if (*leftOrder > *rightOrder) {
+ newRight = rewriter.create<PromoteNumericOp>(loc, leftType, newRight);
+ }
+ if (*rightOrder > *leftOrder) {
+ newLeft = rewriter.create<PromoteNumericOp>(loc, rightType, newLeft);
+ }
+ } else {
+ return failure();
+ }
+
+ // Need to box back to the original type (which will always be a generic
+ // object).
+ newLeft = rewriter.create<BoxOp>(loc, leftResultType, newLeft);
+ newRight = rewriter.create<BoxOp>(loc, rightResultType, newRight);
+
+ rewriter.replaceOp(op, {newLeft, newRight});
+ return success();
+ }
+};
+
+/// If we statically determine one of the arguments to be a concrete, non
+/// numeric type, then the op has no meaning and is elided.
+struct ElideNonNumericDynamicBinaryPromote
+ : public OpRewritePattern<DynamicBinaryPromoteOp> {
+ public:
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(DynamicBinaryPromoteOp op,
+ PatternRewriter &rewriter) const override {
+ if ((!isConcreteNonNumericType(op.left().getType()) &&
+ !isConcreteNonNumericType(op.right().getType())))
+ return failure();
+
+ // Since DynamicBinaryPromote already returns object, and we only match
+ // non-object operands, box them back.
+ auto loc = op.getLoc();
+ auto leftResultType = op.getResultTypes()[0];
+ auto rightResultType = op.getResultTypes()[1];
+ Value newLeft = rewriter.create<BoxOp>(loc, leftResultType, op.left());
+ Value newRight = rewriter.create<BoxOp>(loc, rightResultType, op.right());
+ rewriter.replaceOp(op, {newLeft, newRight});
+ return success();
+ }
+
+ static bool isConcreteNonNumericType(Type t) {
+ if (t.isa<ObjectType>()) return false;
+ auto pt = t.dyn_cast<PythonTypeInterface>();
+ if (!pt || pt.getNumericPromotionOrder()) return false;
+ return true;
+ }
+};
+
+} // namespace
+
+void DynamicBinaryPromoteOp::getCanonicalizationPatterns(
+ RewritePatternSet &patterns, MLIRContext *context) {
+ patterns.add<ResolveNumericDynamicBinaryPromote>(context);
+ patterns.add<UnboxOperands>(getOperationName(), context);
+ patterns.add<ElideNonNumericDynamicBinaryPromote>(context);
+}
+
+//===----------------------------------------------------------------------===//
+// FunctionalIfOp
+//===----------------------------------------------------------------------===//
+
+::llvm::StringRef FunctionalIfOp::getDefaultDialect() { return "iree_pydm"; }
+
+static LogicalResult verify(FunctionalIfOp op) {
+ if (op.getNumResults() != 0 && op.elseRegion().empty())
+ return op.emitOpError("must have an else block if defining values");
+
+ return RegionBranchOpInterface::verifyTypes(op);
+}
+
+static ParseResult parseFunctionalIfOp(OpAsmParser &parser,
+ OperationState &result) {
+ // Create the regions for 'then'.
+ result.regions.reserve(2);
+ Region *thenRegion = result.addRegion();
+ Region *elseRegion = result.addRegion();
+
+ auto &builder = parser.getBuilder();
+ OpAsmParser::OperandType cond;
+ Type conditionType = builder.getType<PyBoolType>();
+ if (parser.parseOperand(cond) ||
+ parser.resolveOperand(cond, conditionType, result.operands))
+ return failure();
+ // Parse optional results type list.
+ if (parser.parseOptionalArrowTypeList(result.types)) return failure();
+ // Parse the 'then' region.
+ if (parser.parseRegion(*thenRegion, /*arguments=*/{}, /*argTypes=*/{}))
+ return failure();
+ // IfOp::ensureTerminator(*thenRegion, parser.getBuilder(), result.location);
+
+ // If we find an 'else' keyword then parse the 'else' region.
+ if (!parser.parseOptionalKeyword("else")) {
+ if (parser.parseRegion(*elseRegion, /*arguments=*/{}, /*argTypes=*/{}))
+ return failure();
+ // IfOp::ensureTerminator(*elseRegion, parser.getBuilder(),
+ // result.location);
+ }
+
+ // Parse the optional attribute list.
+ if (parser.parseOptionalAttrDict(result.attributes)) return failure();
+ return success();
+}
+
+static void print(OpAsmPrinter &p, FunctionalIfOp op) {
+ bool printBlockTerminators = false;
+
+ p << " " << op.condition();
+ if (!op.results().empty()) {
+ p << " -> (" << op.getResultTypes() << ")";
+ // Print yield explicitly if the op defines values.
+ printBlockTerminators = true;
+ }
+ p.printRegion(op.thenRegion(),
+ /*printEntryBlockArgs=*/false,
+ /*printBlockTerminators=*/printBlockTerminators);
+
+ // Print the 'else' regions if it exists and has a block.
+ auto &elseRegion = op.elseRegion();
+ if (!elseRegion.empty()) {
+ p << " else";
+ p.printRegion(elseRegion,
+ /*printEntryBlockArgs=*/false,
+ /*printBlockTerminators=*/printBlockTerminators);
+ }
+
+ p.printOptionalAttrDict(op->getAttrs());
+}
+
+/// Given the region at `index`, or the parent operation if `index` is None,
+/// return the successor regions. These are the regions that may be selected
+/// during the flow of control. `operands` is a set of optional attributes that
+/// correspond to a constant value for each operand, or null if that operand is
+/// not a constant.
+void FunctionalIfOp::getSuccessorRegions(
+ Optional<unsigned> index, ArrayRef<Attribute> operands,
+ SmallVectorImpl<RegionSuccessor> ®ions) {
+ // The `then` and the `else` region branch back to the parent operation.
+ if (index.hasValue()) {
+ regions.push_back(RegionSuccessor(getResults()));
+ return;
+ }
+
+ // Don't consider the else region if it is empty.
+ Region *elseRegion = &this->elseRegion();
+ if (elseRegion->empty()) elseRegion = nullptr;
+
+ // Otherwise, the successor is dependent on the condition.
+ if (auto condAttr = operands.front().dyn_cast_or_null<BoolAttr>()) {
+ bool condition = condAttr.getValue();
+ // Add the successor regions using the condition.
+ regions.push_back(RegionSuccessor(condition ? &thenRegion() : elseRegion));
+ } else {
+ // If the condition isn't constant, both regions may be executed.
+ regions.push_back(RegionSuccessor(&thenRegion()));
+ // If the else region does not exist, it is not a viable successor.
+ if (elseRegion) regions.push_back(RegionSuccessor(elseRegion));
+ }
+}
+
+//===----------------------------------------------------------------------===//
+// FuncOp
+//===----------------------------------------------------------------------===//
+
+::llvm::StringRef PyFuncOp::getDefaultDialect() { return "iree_pydm"; }
+
+LogicalResult PyFuncOp::verifyType() {
+ // TODO: Enforce arg/result invariants.
+ return success();
+}
+
+static ParseResult parseFuncOp(OpAsmParser &parser, OperationState &result) {
+ auto buildFuncType = [](Builder &builder, ArrayRef<Type> argTypes,
+ ArrayRef<Type> results,
+ function_like_impl::VariadicFlag, std::string &) {
+ return builder.getFunctionType(argTypes, results);
+ };
+
+ return function_like_impl::parseFunctionLikeOp(
+ parser, result, /*allowVariadic=*/false, buildFuncType);
+}
+
+static void print(PyFuncOp op, OpAsmPrinter &p) {
+ FunctionType fnType = op.getType();
+ function_like_impl::printFunctionLikeOp(
+ p, op, fnType.getInputs(), /*isVariadic=*/false, fnType.getResults());
+}
+
+static LogicalResult verify(PyFuncOp op) {
+ // TODO: Enforce invariants.
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// MakeListOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(MakeListOp op) {
+ auto listType = op.list().getType().cast<ListType>();
+ switch (listType.getStorageClass()) {
+ case CollectionStorageClass::Boxed:
+ for (auto element : op.elements()) {
+ if (!element.getType().isa<ObjectType>()) {
+ return op.emitOpError() << "making a list with boxed storage class "
+ "must have object elements. Got: "
+ << element.getType();
+ }
+ }
+ break;
+ case CollectionStorageClass::Unboxed:
+ for (auto element : op.elements()) {
+ if (element.getType().isa<ObjectType>()) {
+ return op.emitOpError() << "making a list with unboxed storage class "
+ "must not have object elements. Got: "
+ << element.getType();
+ }
+ }
+ break;
+ case CollectionStorageClass::Empty:
+ if (!op.elements().empty()) {
+ return op.emitOpError()
+ << "making a list with empty storage class must have zero "
+ "elements";
+ }
+ break;
+ }
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// NegOp
+//===----------------------------------------------------------------------===//
+
+void NegOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+ MLIRContext *context) {
+ patterns.add<UnboxOperands>(getOperationName(), context);
+}
+
+bool NegOp::refineResultTypes() {
+ if (value().getType() != getResult().getType()) {
+ getResult().setType(value().getType());
+ return true;
+ }
+ return false;
+}
+
+//===----------------------------------------------------------------------===//
+// PatternMatchCallOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult PatternMatchCallOp::verifySymbolUses(
+ SymbolTableCollection &symbolTable) {
+ auto verifySymbols = [&](ArrayAttr symbols) -> LogicalResult {
+ for (auto symbolAttr : symbols) {
+ auto symbol = symbolAttr.cast<FlatSymbolRefAttr>();
+ PyFuncOp fn =
+ symbolTable.lookupNearestSymbolFrom<PyFuncOp>(*this, symbol);
+ if (!fn)
+ return emitOpError() << "'" << symbol.getValue()
+ << "' does not reference a valid function";
+ }
+ return success();
+ };
+ auto genericsAttr = (*this)->getAttrOfType<ArrayAttr>("generic_match");
+ if (!genericsAttr)
+ return emitOpError(
+ "requires a 'generic_match' array of symbol reference attributes");
+ if (failed(verifySymbols(genericsAttr))) return failure();
+
+ auto specificsAttr = (*this)->getAttrOfType<ArrayAttr>("specific_match");
+ if (!specificsAttr)
+ return emitOpError(
+ "requires a 'specific_match' array of symbol reference attributes");
+ if (failed(verifySymbols(specificsAttr))) return failure();
+
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// PromoteNumericOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult PromoteNumericOp::fold(ArrayRef<Attribute> operands) {
+ if (!operands[0]) return {};
+
+ Builder b(getContext());
+ Attribute fromAttr = operands[0];
+ return TypeSwitch<Type, OpFoldResult>(getResult().getType())
+ .Case([&](PyIntegerType toType) -> OpFoldResult {
+ return TypeSwitch<Attribute, OpFoldResult>(fromAttr)
+ .Case([&](BoolAttr fromBool) -> OpFoldResult {
+ return b.getI64IntegerAttr(fromBool.getValue() ? 1 : 0);
+ })
+ .Default([](Attribute) -> OpFoldResult { return {}; });
+ })
+ .Case([&](PyRealType toType) -> OpFoldResult {
+ return TypeSwitch<Attribute, OpFoldResult>(fromAttr)
+ .Case([&](BoolAttr fromBool) -> OpFoldResult {
+ return b.getF64FloatAttr(fromBool.getValue() ? 1.0 : 0.0);
+ })
+ .Case([&](IntegerAttr fromInteger) -> OpFoldResult {
+ APInt value = fromInteger.getValue();
+ return b.getF64FloatAttr(value.getSExtValue());
+ })
+ .Default([](Attribute) -> OpFoldResult { return {}; });
+ })
+ .Default([](Type) -> OpFoldResult { return {}; });
+}
+
+LogicalResult PromoteNumericOp::canonicalize(PromoteNumericOp op,
+ PatternRewriter &rewriter) {
+ if (op.input().getType() == op.getResult().getType()) {
+ rewriter.replaceOp(op, op.input());
+ return success();
+ }
+ return failure();
+}
+
+//===----------------------------------------------------------------------===//
+// RaiseOnFailureOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult PYDM::RaiseOnFailureOp::fold(
+ ArrayRef<Attribute> operands, SmallVectorImpl<OpFoldResult> &results) {
+ assert(operands.size() == 1 && "expected one fold operand");
+ // Unit exception result is success. Just elide.
+ if (operands[0] && operands[0].isa<UnitAttr>()) {
+ erase();
+ return success();
+ }
+ return failure();
+}
+
+//===----------------------------------------------------------------------===//
+// SelectOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult SelectOp::fold(ArrayRef<Attribute> operands) {
+ if (!operands[0]) return {};
+
+ BoolAttr boolAttr = operands[0].cast<BoolAttr>();
+ if (boolAttr.getValue())
+ return true_value();
+ else
+ return false_value();
+}
+
+//===----------------------------------------------------------------------===//
+// SequenceCloneOp
+//===----------------------------------------------------------------------===//
+
+bool SequenceCloneOp::refineResultTypes() {
+ if (sequence().getType() != getResult().getType()) {
+ getResult().setType(sequence().getType());
+ return true;
+ }
+ return false;
+}
+
+//===----------------------------------------------------------------------===//
+// SubscriptOp
+//===----------------------------------------------------------------------===//
+
+void SubscriptOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+ MLIRContext *context) {
+ patterns.add<UnboxOperands>(getOperationName(), context);
+}
+
+//===----------------------------------------------------------------------===//
+// CallOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult PyCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+ // Check that the callee attribute was specified.
+ auto fnAttr = (*this)->getAttrOfType<FlatSymbolRefAttr>("callee");
+ if (!fnAttr)
+ return emitOpError("requires a 'callee' symbol reference attribute");
+ PyFuncOp fn = symbolTable.lookupNearestSymbolFrom<PyFuncOp>(*this, fnAttr);
+ if (!fn)
+ return emitOpError() << "'" << fnAttr.getValue()
+ << "' does not reference a valid function";
+
+ // Verify that the operand and result types match the callee.
+ auto fnType = fn.getType();
+ if (fnType.getNumInputs() != getNumOperands())
+ return emitOpError("incorrect number of operands for callee");
+
+ for (unsigned i = 0, e = fnType.getNumInputs(); i != e; ++i) {
+ if (getOperand(i).getType() != fnType.getInput(i)) {
+ return emitOpError("operand type mismatch: expected operand type ")
+ << fnType.getInput(i) << ", but provided "
+ << getOperand(i).getType() << " for operand number " << i;
+ }
+ }
+
+ if (fnType.getNumResults() != getNumResults())
+ return emitOpError("incorrect number of results for callee");
+
+ for (unsigned i = 0, e = fnType.getNumResults(); i != e; ++i) {
+ if (getResult(i).getType() != fnType.getResult(i)) {
+ auto diag = emitOpError("result type mismatch at index ") << i;
+ diag.attachNote() << " op result types: " << getResultTypes();
+ diag.attachNote() << "function result types: " << fnType.getResults();
+ return diag;
+ }
+ }
+
+ return success();
+}
+
+FunctionType PyCallOp::getCalleeType() {
+ return FunctionType::get(getContext(), getOperandTypes(), getResultTypes());
+}
+
+//===----------------------------------------------------------------------===//
+// DynamicCallOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult DynamicCallOp::verifySymbolUses(
+ SymbolTableCollection &symbolTable) {
+ // Check that the callee attribute was specified.
+ auto fnAttr = (*this)->getAttrOfType<FlatSymbolRefAttr>("callee");
+ if (!fnAttr)
+ return emitOpError("requires a 'callee' symbol reference attribute");
+ Operation *fn = symbolTable.lookupNearestSymbolFrom(*this, fnAttr);
+ if (!fn || !isa<PyFuncOp>(fn))
+ return emitOpError() << "'" << fnAttr.getValue()
+ << "' does not reference a valid function";
+ return success();
+}
+
+#define GET_OP_CLASSES
+#include "iree-dialects/Dialect/PyDM/IR/PyDMOps.cpp.inc"
