diff --git a/mlir/docs/Dialects/Affine.md b/mlir/docs/Dialects/Affine.md index f2d9fdabab51b5..7eab93ace42e5a 100644 --- a/mlir/docs/Dialects/Affine.md +++ b/mlir/docs/Dialects/Affine.md @@ -60,20 +60,26 @@ Example: ### Restrictions on Dimensions and Symbols The affine dialect imposes certain restrictions on dimension and symbolic -identifiers to enable powerful analysis and transformation. A symbolic -identifier can be bound to an SSA value that is either an argument to the -function, a value defined at the top level of that function (outside of all -loops and if operations), the result of a -[`constant` operation](Standard.md#constant-operation), or the result of an -[`affine.apply` operation](#affineapply-operation) that recursively takes as -arguments any symbolic identifiers, or the result of a [`dim` -operation](Standard.md#dim-operation) on either a memref that is a function -argument or a memref where the corresponding dimension is either static or a -dynamic one in turn bound to a symbolic identifier. Dimensions may be bound not -only to anything that a symbol is bound to, but also to induction variables of -enclosing [`affine.for`](#affinefor-affineforop) and -[`afffine.parallel`](#affineparallel-affineparallelop) operations, and the -result of an +identifiers to enable powerful analysis and transformation. An SSA value's use +can be bound to a symbolic identifier if that SSA value is either +1. a region argument for an op with trait `PolyhedralScope` (eg. `FuncOp`), +2. a value defined at the top level of a `PolyhedralScope` op (i.e., immediately +enclosed by the latter), +3. a value that dominates the `PolyhedralScope` op enclosing the value's use, +4. the result of a [`constant` operation](Standard.md#constant-operation), +5. the result of an [`affine.apply` +operation](#affineapply-operation) that recursively takes as arguments any valid +symbolic identifiers, or +6. the result of a [`dim` operation](Standard.md#dim-operation) on either a +memref that is an argument to a `PolyhedralScope` op or a memref where the +corresponding dimension is either static or a dynamic one in turn bound to a +valid symbol. + +Note that as a result of rule (3) above, symbol validity is sensitive to the +location of the SSA use. Dimensions may be bound not only to anything that a +symbol is bound to, but also to induction variables of enclosing +[`affine.for`](#affinefor-operation) and +[`affine.parallel`](#affineparallel-operation) operations, and the result of an [`affine.apply` operation](#affineapply-operation) (which recursively may use other dimensions and symbols). diff --git a/mlir/docs/Traits.md b/mlir/docs/Traits.md index 5931fd3f969877..0281185abf9137 100644 --- a/mlir/docs/Traits.md +++ b/mlir/docs/Traits.md @@ -219,6 +219,22 @@ foo.region_op { This trait is an important structural property of the IR, and enables operations to have [passes](PassManagement.md) scheduled under them. + +### PolyhedralScope + +* `OpTrait::PolyhedralScope` -- `PolyhedralScope` + +This trait is carried by region holding operations that define a new scope for +the purposes of polyhedral optimization and the affine dialect in particular. +Any SSA values of 'index' type that either dominate such operations, or are +defined at the top-level of such operations, or appear as region arguments for +such operations automatically become valid symbols for the polyhedral scope +defined by that operation. As a result, such SSA values could be used as the +operands or index operands of various affine dialect operations like affine.for, +affine.load, and affine.store. The polyhedral scope defined by an operation +with this trait includes all operations in its region excluding operations that +are nested inside of other operations that themselves have this trait. + ### Single Block with Implicit Terminator * `OpTrait::SingleBlockImplicitTerminator` : diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h index 6ce38bcddddc5e..e0741cfb3afac2 100644 --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h @@ -31,9 +31,10 @@ class AffineTerminatorOp; class FlatAffineConstraints; class OpBuilder; -/// A utility function to check if a value is defined at the top level of a -/// function. A value of index type defined at the top level is always a valid -/// symbol. +/// A utility function to check if a value is defined at the top level of an +/// op with trait `PolyhedralScope` or is a region argument for such an op. A +/// value of index type defined at the top level is always a valid symbol for +/// all its uses. bool isTopLevelValue(Value value); /// AffineDmaStartOp starts a non-blocking DMA operation that transfers data @@ -457,12 +458,22 @@ class AffineStoreOp : public Op &results); }; -/// Returns true if the given Value can be used as a dimension id. +/// Returns true if the given Value can be used as a dimension id in the region +/// of the closest surrounding op that has the trait `PolyhedralScope`. bool isValidDim(Value value); -/// Returns true if the given Value can be used as a symbol. +/// Returns true if the given Value can be used as a dimension id in `region`, +/// i.e., for all its uses in `region`. +bool isValidDim(Value value, Region *region); + +/// Returns true if the given value can be used as a symbol in the region of the +/// closest surrounding op that has the trait `PolyhedralScope`. bool isValidSymbol(Value value); +/// Returns true if the given Value can be used as a symbol for `region`, i.e., +/// for all its uses in `region`. +bool isValidSymbol(Value value, Region *region); + /// Modifies both `map` and `operands` in-place so as to: /// 1. drop duplicate operands /// 2. drop unused dims and symbols from map diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td index a49413fca27dec..ce93e103634ee2 100644 --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td @@ -83,12 +83,22 @@ def AffineApplyOp : Affine_Op<"apply", [NoSideEffect]> { /// Returns the affine value map computed from this operation. AffineValueMap getAffineValueMap(); - /// Returns true if the result of this operation can be used as dimension id. + /// Returns true if the result of this operation can be used as dimension id + /// in the region of the closest surrounding op with trait PolyhedralScope. bool isValidDim(); - /// Returns true if the result of this operation is a symbol. + /// Returns true if the result of this operation can be used as dimension id + /// within 'region', i.e., for all its uses with `region`. + bool isValidDim(Region *region); + + /// Returns true if the result of this operation is a symbol in the region + /// of the closest surrounding op that has the trait PolyhedralScope. bool isValidSymbol(); + /// Returns true if the result of this operation is a symbol for all its + /// uses in `region`. + bool isValidSymbol(Region *region); + operand_range getMapOperands() { return getOperands(); } }]; diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOpsBase.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOpsBase.td index 2883072d4aa98a..95232697789f96 100644 --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOpsBase.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOpsBase.td @@ -29,4 +29,11 @@ def AffineMapArrayAttr : TypedArrayAttrBase; + #endif // AFFINE_OPS_BASE diff --git a/mlir/include/mlir/Dialect/Affine/Traits.h b/mlir/include/mlir/Dialect/Affine/Traits.h new file mode 100644 index 00000000000000..6dc6f481d89d20 --- /dev/null +++ b/mlir/include/mlir/Dialect/Affine/Traits.h @@ -0,0 +1,41 @@ +//===- Traits.h - Traits for the affine dialect -----------------*- C++ -*-===// +// +// Part of the LLVM Project, 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 +// +//===----------------------------------------------------------------------===// +// +// This file declares traits that the affine dialect relies upon for analysis +// and transformation purposes, and that are also potentially used by other +// dialect entities not depending on the affine dialect. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_AFFINE_TRAITS +#define MLIR_DIALECT_AFFINE_TRAITS + +#include "mlir/IR/OpDefinition.h" + +namespace mlir { +namespace OpTrait { + +/// A trait of region holding operations that defines a new scope for polyhedral +/// optimization purposes. Any SSA values of 'index' type that either dominate +/// such an operation or are used at the top-level of such an operation +/// automatically become valid symbols for the polyhedral scope defined by that +/// operation. For more details, see `Traits.md#PolyhedralScope`. +template +class PolyhedralScope : public TraitBase { +public: + static LogicalResult verifyTrait(Operation *op) { + static_assert(!ConcreteType::template hasTrait(), + "expected operation to have one or more regions"); + return success(); + } +}; + +} // end namespace OpTrait +} // end namespace mlir + +#endif // MLIR_DIALECT_TRAITS diff --git a/mlir/include/mlir/IR/Function.h b/mlir/include/mlir/IR/Function.h index 0b725949576c9e..3e26f8f7399d5f 100644 --- a/mlir/include/mlir/IR/Function.h +++ b/mlir/include/mlir/IR/Function.h @@ -13,6 +13,7 @@ #ifndef MLIR_IR_FUNCTION_H #define MLIR_IR_FUNCTION_H +#include "mlir/Dialect/Affine/Traits.h" #include "mlir/IR/Block.h" #include "mlir/IR/FunctionSupport.h" #include "mlir/IR/OpDefinition.h" @@ -30,10 +31,11 @@ namespace mlir { /// implicitly capture global values, and all external references must use /// Function arguments or attributes that establish a symbolic connection(e.g. /// symbols referenced by name via a string attribute). -class FuncOp : public Op { +class FuncOp + : public Op { public: using Op::Op; using Op::print; diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index 5f490b0d3721b1..b64ce80bc0d093 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -8,6 +8,7 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Affine/IR/AffineValueMap.h" +#include "mlir/Dialect/Affine/Traits.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/IR/Function.h" #include "mlir/IR/IntegerSet.h" @@ -84,65 +85,110 @@ Operation *AffineDialect::materializeConstant(OpBuilder &builder, return builder.create(loc, type, value); } -/// A utility function to check if a given region is attached to a function. -static bool isFunctionRegion(Region *region) { - return llvm::isa(region->getParentOp()); +/// A utility function to check if a value is defined at the top level of an +/// op with trait `PolyhedralScope`. A value of index type defined at the top +/// level is always a valid symbol. +bool mlir::isTopLevelValue(Value value) { + if (auto arg = value.dyn_cast()) + return arg.getOwner()->getParentOp()->hasTrait(); + return value.getDefiningOp() + ->getParentOp() + ->hasTrait(); } -/// A utility function to check if a value is defined at the top level of a -/// function. A value of index type defined at the top level is always a valid -/// symbol. -bool mlir::isTopLevelValue(Value value) { +/// A utility function to check if a value is defined at the top level of +/// `region` or is an argument of `region`. A value of index type defined at the +/// top level of a `PolyhedralScope` region is always a valid symbol for all +/// uses in that region. +static bool isTopLevelValue(Value value, Region *region) { if (auto arg = value.dyn_cast()) - return isFunctionRegion(arg.getOwner()->getParent()); - return isFunctionRegion(value.getDefiningOp()->getParentRegion()); + return arg.getParentRegion() == region; + return value.getDefiningOp()->getParentOp() == region->getParentOp(); +} + +/// Returns the closest region enclosing `op` that is held by an operation with +/// trait `PolyhedralScope`. +// TODO: getAffineScope should be publicly exposed for affine passes/utilities. +static Region *getAffineScope(Operation *op) { + auto *curOp = op; + while (auto *parentOp = curOp->getParentOp()) { + if (parentOp->hasTrait()) + return curOp->getParentRegion(); + curOp = parentOp; + } + llvm_unreachable("op doesn't have an enclosing polyhedral scope"); } -// Value can be used as a dimension id if it is valid as a symbol, or -// it is an induction variable, or it is a result of affine apply operation -// with dimension id arguments. +// A Value can be used as a dimension id iff it meets one of the following +// conditions: +// *) It is valid as a symbol. +// *) It is an induction variable. +// *) It is the result of affine apply operation with dimension id arguments. bool mlir::isValidDim(Value value) { // The value must be an index type. if (!value.getType().isIndex()) return false; - if (auto *op = value.getDefiningOp()) { - // Top level operation or constant operation is ok. - if (isFunctionRegion(op->getParentRegion()) || isa(op)) - return true; - // Affine apply operation is ok if all of its operands are ok. - if (auto applyOp = dyn_cast(op)) - return applyOp.isValidDim(); - // The dim op is okay if its operand memref/tensor is defined at the top - // level. - if (auto dimOp = dyn_cast(op)) - return isTopLevelValue(dimOp.getOperand()); + if (auto *defOp = value.getDefiningOp()) + return isValidDim(value, getAffineScope(defOp)); + + // This value has to be a block argument for an op that has the + // `PolyhedralScope` trait or for an affine.for or affine.parallel. + auto *parentOp = value.cast().getOwner()->getParentOp(); + return parentOp->hasTrait() || + isa(parentOp) || isa(parentOp); +} + +// Value can be used as a dimension id iff it meets one of the following +// conditions: +// *) It is valid as a symbol. +// *) It is an induction variable. +// *) It is the result of an affine apply operation with dimension id operands. +bool mlir::isValidDim(Value value, Region *region) { + // The value must be an index type. + if (!value.getType().isIndex()) return false; + + // All valid symbols are okay. + if (isValidSymbol(value, region)) + return true; + + auto *op = value.getDefiningOp(); + if (!op) { + // This value has to be a block argument for an affine.for or an + // affine.parallel. + auto *parentOp = value.cast().getOwner()->getParentOp(); + return isa(parentOp) || isa(parentOp); } - // This value has to be a block argument of a FuncOp, an 'affine.for', or an - // 'affine.parallel'. - auto *parentOp = value.cast().getOwner()->getParentOp(); - return isa(parentOp) || isa(parentOp) || - isa(parentOp); + + // Affine apply operation is ok if all of its operands are ok. + if (auto applyOp = dyn_cast(op)) + return applyOp.isValidDim(region); + // The dim op is okay if its operand memref/tensor is defined at the top + // level. + if (auto dimOp = dyn_cast(op)) + return isTopLevelValue(dimOp.getOperand()); + return false; } /// Returns true if the 'index' dimension of the `memref` defined by -/// `memrefDefOp` is a statically shaped one or defined using a valid symbol. +/// `memrefDefOp` is a statically shaped one or defined using a valid symbol +/// for `region`. template -static bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp, - unsigned index) { +bool isMemRefSizeValidSymbol(AnyMemRefDefOp memrefDefOp, unsigned index, + Region *region) { auto memRefType = memrefDefOp.getType(); // Statically shaped. if (!memRefType.isDynamicDim(index)) return true; // Get the position of the dimension among dynamic dimensions; unsigned dynamicDimPos = memRefType.getDynamicDimIndex(index); - return isValidSymbol( - *(memrefDefOp.getDynamicSizes().begin() + dynamicDimPos)); + return isValidSymbol(*(memrefDefOp.getDynamicSizes().begin() + dynamicDimPos), + region); } -/// Returns true if the result of the dim op is a valid symbol. -static bool isDimOpValidSymbol(DimOp dimOp) { +/// Returns true if the result of the dim op is a valid symbol for `region`. +static bool isDimOpValidSymbol(DimOp dimOp, Region *region) { // The dim op is okay if its operand memref/tensor is defined at the top // level. if (isTopLevelValue(dimOp.getOperand())) @@ -152,43 +198,90 @@ static bool isDimOpValidSymbol(DimOp dimOp) { // whose corresponding size is a valid symbol. unsigned index = dimOp.getIndex(); if (auto viewOp = dyn_cast(dimOp.getOperand().getDefiningOp())) - return isMemRefSizeValidSymbol(viewOp, index); + return isMemRefSizeValidSymbol(viewOp, index, region); if (auto subViewOp = dyn_cast(dimOp.getOperand().getDefiningOp())) - return isMemRefSizeValidSymbol(subViewOp, index); + return isMemRefSizeValidSymbol(subViewOp, index, region); if (auto allocOp = dyn_cast(dimOp.getOperand().getDefiningOp())) - return isMemRefSizeValidSymbol(allocOp, index); + return isMemRefSizeValidSymbol(allocOp, index, region); return false; } -// Value can be used as a symbol if it is a constant, or it is defined at -// the top level, or it is a result of affine apply operation with symbol -// arguments, or a result of the dim op on a memref satisfying certain -// constraints. +// A value can be used as a symbol (at all its use sites) iff it meets one of +// the following conditions: +// *) It is a constant. +// *) Its defining op or block arg appearance is immediately enclosed by an op +// with `PolyhedralScope` trait. +// *) It is the result of an affine.apply operation with symbol operands. +// *) It is a result of the dim op on a memref whose corresponding size is a +// valid symbol. bool mlir::isValidSymbol(Value value) { // The value must be an index type. if (!value.getType().isIndex()) return false; - if (auto *op = value.getDefiningOp()) { - // Top level operation or constant operation is ok. - if (isFunctionRegion(op->getParentRegion()) || isa(op)) - return true; - // Affine apply operation is ok if all of its operands are ok. - if (auto applyOp = dyn_cast(op)) - return applyOp.isValidSymbol(); - if (auto dimOp = dyn_cast(op)) { - return isDimOpValidSymbol(dimOp); - } + // Check that the value is a top level value. + if (isTopLevelValue(value)) + return true; + + if (auto *defOp = value.getDefiningOp()) + return isValidSymbol(value, getAffineScope(defOp)); + + return false; +} + +// A value can be used as a symbol for `region` iff it meets onf of the the +// following conditions: +// *) It is a constant. +// *) It is defined at the top level of 'region' or is its argument. +// *) It dominates `region`'s parent op. +// *) It is the result of an affine apply operation with symbol arguments. +// *) It is a result of the dim op on a memref whose corresponding size is +// a valid symbol. +bool mlir::isValidSymbol(Value value, Region *region) { + // The value must be an index type. + if (!value.getType().isIndex()) + return false; + + // A top-level value is a valid symbol. + if (::isTopLevelValue(value, region)) + return true; + + auto *defOp = value.getDefiningOp(); + if (!defOp) { + // A block argument that is not a top-level value is a valid symbol if it + // dominates region's parent op. + if (!region->getParentOp()->isKnownIsolatedFromAbove()) + if (auto *parentOpRegion = region->getParentOp()->getParentRegion()) + return isValidSymbol(value, parentOpRegion); + return false; } - // Otherwise, check that the value is a top level value. - return isTopLevelValue(value); + + // Constant operation is ok. + Attribute operandCst; + if (matchPattern(defOp, m_Constant(&operandCst))) + return true; + + // Affine apply operation is ok if all of its operands are ok. + if (auto applyOp = dyn_cast(defOp)) + return applyOp.isValidSymbol(region); + + // Dim op results could be valid symbols at any level. + if (auto dimOp = dyn_cast(defOp)) + return isDimOpValidSymbol(dimOp, region); + + // Check for values dominating `region`'s parent op. + if (!region->getParentOp()->isKnownIsolatedFromAbove()) + if (auto *parentRegion = region->getParentOp()->getParentRegion()) + return isValidSymbol(value, parentRegion); + + return false; } // Returns true if 'value' is a valid index to an affine operation (e.g. -// affine.load, affine.store, affine.dma_start, affine.dma_wait). -// Returns false otherwise. -static bool isValidAffineIndexOperand(Value value) { - return isValidDim(value) || isValidSymbol(value); +// affine.load, affine.store, affine.dma_start, affine.dma_wait) where +// `region` provides the polyhedral symbol scope. Returns false otherwise. +static bool isValidAffineIndexOperand(Value value, Region *region) { + return isValidDim(value, region) || isValidSymbol(value, region); } /// Utility function to verify that a set of operands are valid dimension and @@ -203,9 +296,9 @@ verifyDimAndSymbolIdentifiers(OpTy &op, Operation::operand_range operands, unsigned opIt = 0; for (auto operand : operands) { if (opIt++ < numDims) { - if (!isValidDim(operand)) + if (!isValidDim(operand, getAffineScope(op))) return op.emitOpError("operand cannot be used as a dimension id"); - } else if (!isValidSymbol(operand)) { + } else if (!isValidSymbol(operand, getAffineScope(op))) { return op.emitOpError("operand cannot be used as a symbol"); } } @@ -273,6 +366,14 @@ bool AffineApplyOp::isValidDim() { [](Value op) { return mlir::isValidDim(op); }); } +// The result of the affine apply operation can be used as a dimension id if all +// its operands are valid dimension ids with the parent operation of `region` +// defining the polyhedral scope for symbols. +bool AffineApplyOp::isValidDim(Region *region) { + return llvm::all_of(getOperands(), + [&](Value op) { return ::isValidDim(op, region); }); +} + // The result of the affine apply operation can be used as a symbol if all its // operands are symbols. bool AffineApplyOp::isValidSymbol() { @@ -280,6 +381,14 @@ bool AffineApplyOp::isValidSymbol() { [](Value op) { return mlir::isValidSymbol(op); }); } +// The result of the affine apply operation can be used as a symbol in `region` +// if all its operands are symbols in `region`. +bool AffineApplyOp::isValidSymbol(Region *region) { + return llvm::all_of(getOperands(), [&](Value operand) { + return ::isValidSymbol(operand, region); + }); +} + OpFoldResult AffineApplyOp::fold(ArrayRef operands) { auto map = getAffineMap(); @@ -948,22 +1057,23 @@ LogicalResult AffineDmaStartOp::verify() { return emitOpError("incorrect number of operands"); } + Region *scope = getAffineScope(*this); for (auto idx : getSrcIndices()) { if (!idx.getType().isIndex()) return emitOpError("src index to dma_start must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) + if (!isValidAffineIndexOperand(idx, scope)) return emitOpError("src index must be a dimension or symbol identifier"); } for (auto idx : getDstIndices()) { if (!idx.getType().isIndex()) return emitOpError("dst index to dma_start must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) + if (!isValidAffineIndexOperand(idx, scope)) return emitOpError("dst index must be a dimension or symbol identifier"); } for (auto idx : getTagIndices()) { if (!idx.getType().isIndex()) return emitOpError("tag index to dma_start must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) + if (!isValidAffineIndexOperand(idx, scope)) return emitOpError("tag index must be a dimension or symbol identifier"); } return success(); @@ -1036,10 +1146,11 @@ ParseResult AffineDmaWaitOp::parse(OpAsmParser &parser, LogicalResult AffineDmaWaitOp::verify() { if (!getOperand(0).getType().isa()) return emitOpError("expected DMA tag to be of memref type"); + Region *scope = getAffineScope(*this); for (auto idx : getTagIndices()) { if (!idx.getType().isIndex()) return emitOpError("index to dma_wait must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) + if (!isValidAffineIndexOperand(idx, scope)) return emitOpError("index must be a dimension or symbol identifier"); } return success(); @@ -1811,10 +1922,11 @@ LogicalResult AffineLoadOp::verify() { "expects the number of subscripts to be equal to memref rank"); } + Region *scope = getAffineScope(*this); for (auto idx : getMapOperands()) { if (!idx.getType().isIndex()) return emitOpError("index to load must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) + if (!isValidAffineIndexOperand(idx, scope)) return emitOpError("index must be a dimension or symbol identifier"); } return success(); @@ -1909,10 +2021,11 @@ LogicalResult AffineStoreOp::verify() { "expects the number of subscripts to be equal to memref rank"); } + Region *scope = getAffineScope(*this); for (auto idx : getMapOperands()) { if (!idx.getType().isIndex()) return emitOpError("index to store must have 'index' type"); - if (!isValidAffineIndexOperand(idx)) + if (!isValidAffineIndexOperand(idx, scope)) return emitOpError("index must be a dimension or symbol identifier"); } return success(); @@ -2131,8 +2244,9 @@ static LogicalResult verify(AffinePrefetchOp op) { return op.emitOpError("too few operands"); } + Region *scope = getAffineScope(op); for (auto idx : op.getMapOperands()) { - if (!isValidAffineIndexOperand(idx)) + if (!isValidAffineIndexOperand(idx, scope)) return op.emitOpError("index must be a dimension or symbol identifier"); } return success(); diff --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir index 4ad8b7ed824706..c0855987ac3237 100644 --- a/mlir/test/Dialect/Affine/invalid.mlir +++ b/mlir/test/Dialect/Affine/invalid.mlir @@ -124,7 +124,7 @@ func @affine_if_invalid_dimop_dim(%arg0: index, %arg1: index, %arg2: index, %arg %0 = alloc(%arg0, %arg1, %arg2, %arg3) : memref %dim = dim %0, 0 : memref - // expected-error@+1 {{operand cannot be used as a dimension id}} + // expected-error@+1 {{operand cannot be used as a symbol}} affine.if #set0(%dim)[%n0] {} } return diff --git a/mlir/test/Dialect/Affine/ops.mlir b/mlir/test/Dialect/Affine/ops.mlir index cd42980f87c978..d6494e377d7bb1 100644 --- a/mlir/test/Dialect/Affine/ops.mlir +++ b/mlir/test/Dialect/Affine/ops.mlir @@ -115,6 +115,33 @@ func @valid_symbols(%arg0: index, %arg1: index, %arg2: index) { // ----- +// Test symbol constraints for ops with PolyhedralScope trait. + +// CHECK-LABEL: func @valid_symbol_polyhedral_scope +func @valid_symbol_polyhedral_scope(%n : index, %A : memref) { + test.polyhedral_scope { + %c1 = constant 1 : index + %l = subi %n, %c1 : index + // %l, %n are valid symbols since test.polyhedral_scope defines a new + // polyhedral scope. + affine.for %i = %l to %n { + %m = subi %l, %i : index + test.polyhedral_scope { + // %m and %n are valid symbols. + affine.for %j = %m to %n { + %v = affine.load %A[%n - 1] : memref + affine.store %v, %A[%n - 1] : memref + } + "terminate"() : () -> () + } + } + "terminate"() : () -> () + } + return +} + +// ----- + // CHECK-LABEL: @parallel // CHECK-SAME: (%[[N:.*]]: index) func @parallel(%N : index) { diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp index fa99d472676cd1..cbf53a5e3f224e 100644 --- a/mlir/test/lib/Dialect/Test/TestDialect.cpp +++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp @@ -201,6 +201,22 @@ static void print(OpAsmPrinter &p, IsolatedRegionOp op) { p.printRegion(op.region(), /*printEntryBlockArgs=*/false); } +//===----------------------------------------------------------------------===// +// Test PolyhedralScopeOp +//===----------------------------------------------------------------------===// + +static ParseResult parsePolyhedralScopeOp(OpAsmParser &parser, + OperationState &result) { + // Parse the body region, and reuse the operand info as the argument info. + Region *body = result.addRegion(); + return parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}); +} + +static void print(OpAsmPrinter &p, PolyhedralScopeOp op) { + p << "test.polyhedral_scope "; + p.printRegion(op.region(), /*printEntryBlockArgs=*/false); +} + //===----------------------------------------------------------------------===// // Test parser. //===----------------------------------------------------------------------===// diff --git a/mlir/test/lib/Dialect/Test/TestDialect.h b/mlir/test/lib/Dialect/Test/TestDialect.h index b4ca125cb3d615..405b6ee5d984fc 100644 --- a/mlir/test/lib/Dialect/Test/TestDialect.h +++ b/mlir/test/lib/Dialect/Test/TestDialect.h @@ -14,6 +14,7 @@ #ifndef MLIR_TESTDIALECT_H #define MLIR_TESTDIALECT_H +#include "mlir/Dialect/Affine/Traits.h" #include "mlir/Dialect/Traits.h" #include "mlir/IR/Dialect.h" #include "mlir/IR/OpDefinition.h" diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td index ad8c6fb99e6794..6bd5b1b16839e5 100644 --- a/mlir/test/lib/Dialect/Test/TestOps.td +++ b/mlir/test/lib/Dialect/Test/TestOps.td @@ -9,6 +9,7 @@ #ifndef TEST_OPS #define TEST_OPS +include "mlir/Dialect/Affine/IR/AffineOpsBase.td" include "mlir/IR/OpBase.td" include "mlir/IR/OpAsmInterface.td" include "mlir/IR/SymbolInterfaces.td" @@ -1129,6 +1130,17 @@ def IsolatedRegionOp : TEST_Op<"isolated_region", [IsolatedFromAbove]> { let printer = [{ return ::print(p, *this); }]; } +def PolyhedralScopeOp : TEST_Op<"polyhedral_scope", [PolyhedralScope]> { + let summary = "polyhedral scope operation"; + let description = [{ + Test op that defines a new polyhedral scope. + }]; + + let regions = (region SizedRegion<1>:$region); + let parser = [{ return ::parse$cppClass(parser, result); }]; + let printer = [{ return ::print(p, *this); }]; +} + def WrappingRegionOp : TEST_Op<"wrapping_region", [SingleBlockImplicitTerminator<"TestReturnOp">]> { let summary = "wrapping region operation";