[Mlir-commits] [mlir] 4803453 - [MLIR] Introduce op trait PolyhedralScope (revised)
Uday Bondhugula
llvmlistbot at llvm.org
Wed Apr 29 03:40:06 PDT 2020
Author: Uday Bondhugula
Date: 2020-04-29T16:08:23+05:30
New Revision: 480345381a3ec9b067bbccc509e2462c10cfc9e3
URL: https://github.com/llvm/llvm-project/commit/480345381a3ec9b067bbccc509e2462c10cfc9e3
DIFF: https://github.com/llvm/llvm-project/commit/480345381a3ec9b067bbccc509e2462c10cfc9e3.diff
LOG: [MLIR] Introduce op trait PolyhedralScope (revised)
(A previous version of this, dd2c639c3cd397dfef941186fb85c82e4e918425, was
reverted.)
Introduce op trait PolyhedralScope for ops to define a new scope for
polyhedral optimization / affine dialect purposes, thus generalizing
such scopes beyond FuncOp. Ops to which this trait is attached will
define a new scope for the consideration of SSA values as valid symbols
for the purposes of polyhedral analysis and optimization. Update methods
that check for dim/symbol validity to work based on this trait.
Differential Revision: https://reviews.llvm.org/D79060
Added:
Modified:
mlir/docs/Dialects/Affine.md
mlir/docs/Traits.md
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/include/mlir/IR/Function.h
mlir/include/mlir/IR/Module.h
mlir/include/mlir/IR/OpBase.td
mlir/include/mlir/IR/OpDefinition.h
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/test/Dialect/Affine/invalid.mlir
mlir/test/Dialect/Affine/ops.mlir
mlir/test/lib/Dialect/Test/TestDialect.cpp
mlir/test/lib/Dialect/Test/TestOps.td
Removed:
################################################################################
diff --git a/mlir/docs/Dialects/Affine.md b/mlir/docs/Dialects/Affine.md
index 246c944f66fd..7eab93ace42e 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
-[`affine.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 5931fd3f9698..0281185abf91 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<typename TerminatorOpType>` :
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index cbbc2d76c46e..b0bba011eb6b 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
@@ -316,12 +317,22 @@ class AffineDmaWaitOp : public Op<AffineDmaWaitOp, OpTrait::VariadicOperands,
SmallVectorImpl<OpFoldResult> &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 c8f56539fafe..e2b7716fc334 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/IR/Function.h b/mlir/include/mlir/IR/Function.h
index 9596b374a724..90b45910a2a3 100644
--- a/mlir/include/mlir/IR/Function.h
+++ b/mlir/include/mlir/IR/Function.h
@@ -30,10 +30,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<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
- OpTrait::IsIsolatedFromAbove, OpTrait::FunctionLike,
- OpTrait::AutomaticAllocationScope,
- CallableOpInterface::Trait, SymbolOpInterface::Trait> {
+class FuncOp
+ : public Op<FuncOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
+ OpTrait::IsIsolatedFromAbove, OpTrait::FunctionLike,
+ OpTrait::AutomaticAllocationScope, OpTrait::PolyhedralScope,
+ CallableOpInterface::Trait, SymbolOpInterface::Trait> {
public:
using Op::Op;
using Op::print;
diff --git a/mlir/include/mlir/IR/Module.h b/mlir/include/mlir/IR/Module.h
index 4f0bba773ee2..e724dd034030 100644
--- a/mlir/include/mlir/IR/Module.h
+++ b/mlir/include/mlir/IR/Module.h
@@ -30,7 +30,8 @@ class ModuleTerminatorOp;
class ModuleOp
: public Op<
ModuleOp, OpTrait::ZeroOperands, OpTrait::ZeroResult,
- OpTrait::IsIsolatedFromAbove, OpTrait::SymbolTable,
+ OpTrait::IsIsolatedFromAbove, OpTrait::PolyhedralScope,
+ OpTrait::SymbolTable,
OpTrait::SingleBlockImplicitTerminator<ModuleTerminatorOp>::Impl,
SymbolOpInterface::Trait> {
public:
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 3e55eadfd742..087aecb6b35d 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1648,6 +1648,8 @@ def ConstantLike : NativeOpTrait<"ConstantLike">;
def FunctionLike : NativeOpTrait<"FunctionLike">;
// Op is isolated from above.
def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
+// Op defines a polyhedral scope.
+def PolyhedralScope : NativeOpTrait<"PolyhedralScope">;
// Op results are float or vectors/tensors thereof.
def ResultsAreFloatLike : NativeOpTrait<"ResultsAreFloatLike">;
// Op has the same operand type.
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 107234a7158e..5dc481d1df2b 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1034,6 +1034,21 @@ class IsIsolatedFromAbove
}
};
+/// 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 <typename ConcreteType>
+class PolyhedralScope : public TraitBase<ConcreteType, PolyhedralScope> {
+public:
+ static LogicalResult verifyTrait(Operation *op) {
+ static_assert(!ConcreteType::template hasTrait<ZeroRegion>(),
+ "expected operation to have one or more regions");
+ return success();
+ }
+};
+
/// A trait of region holding operations that define a new scope for automatic
/// allocations, i.e., allocations that are freed when control is transferred
/// back from the operation's region. Any operations performing such allocations
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 0ea1f8dc84e7..8f533e9b2177 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -84,65 +84,110 @@ Operation *AffineDialect::materializeConstant(OpBuilder &builder,
return builder.create<ConstantOp>(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<FuncOp>(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<BlockArgument>())
+ return arg.getOwner()->getParentOp()->hasTrait<OpTrait::PolyhedralScope>();
+ return value.getDefiningOp()
+ ->getParentOp()
+ ->hasTrait<OpTrait::PolyhedralScope>();
}
-/// 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<BlockArgument>())
- 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<OpTrait::PolyhedralScope>())
+ 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<ConstantOp>(op))
- return true;
- // Affine apply operation is ok if all of its operands are ok.
- if (auto applyOp = dyn_cast<AffineApplyOp>(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<DimOp>(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<BlockArgument>().getOwner()->getParentOp();
+ return parentOp->hasTrait<OpTrait::PolyhedralScope>() ||
+ isa<AffineForOp>(parentOp) || isa<AffineParallelOp>(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<BlockArgument>().getOwner()->getParentOp();
+ return isa<AffineForOp>(parentOp) || isa<AffineParallelOp>(parentOp);
}
- // This value has to be a block argument of a FuncOp, an 'affine.for', or an
- // 'affine.parallel'.
- auto *parentOp = value.cast<BlockArgument>().getOwner()->getParentOp();
- return isa<FuncOp>(parentOp) || isa<AffineForOp>(parentOp) ||
- isa<AffineParallelOp>(parentOp);
+
+ // Affine apply operation is ok if all of its operands are ok.
+ if (auto applyOp = dyn_cast<AffineApplyOp>(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<DimOp>(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 <typename AnyMemRefDefOp>
-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 +197,90 @@ static bool isDimOpValidSymbol(DimOp dimOp) {
// whose corresponding size is a valid symbol.
unsigned index = dimOp.getIndex();
if (auto viewOp = dyn_cast<ViewOp>(dimOp.getOperand().getDefiningOp()))
- return isMemRefSizeValidSymbol<ViewOp>(viewOp, index);
+ return isMemRefSizeValidSymbol<ViewOp>(viewOp, index, region);
if (auto subViewOp = dyn_cast<SubViewOp>(dimOp.getOperand().getDefiningOp()))
- return isMemRefSizeValidSymbol<SubViewOp>(subViewOp, index);
+ return isMemRefSizeValidSymbol<SubViewOp>(subViewOp, index, region);
if (auto allocOp = dyn_cast<AllocOp>(dimOp.getOperand().getDefiningOp()))
- return isMemRefSizeValidSymbol<AllocOp>(allocOp, index);
+ return isMemRefSizeValidSymbol<AllocOp>(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<ConstantOp>(op))
- return true;
- // Affine apply operation is ok if all of its operands are ok.
- if (auto applyOp = dyn_cast<AffineApplyOp>(op))
- return applyOp.isValidSymbol();
- if (auto dimOp = dyn_cast<DimOp>(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<AffineApplyOp>(defOp))
+ return applyOp.isValidSymbol(region);
+
+ // Dim op results could be valid symbols at any level.
+ if (auto dimOp = dyn_cast<DimOp>(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 +295,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 +365,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 +380,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 mlir::isValidSymbol(operand, region);
+ });
+}
+
OpFoldResult AffineApplyOp::fold(ArrayRef<Attribute> operands) {
auto map = getAffineMap();
@@ -948,22 +1056,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 +1145,11 @@ ParseResult AffineDmaWaitOp::parse(OpAsmParser &parser,
LogicalResult AffineDmaWaitOp::verify() {
if (!getOperand(0).getType().isa<MemRefType>())
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();
@@ -1814,10 +1924,11 @@ LogicalResult verify(AffineLoadOp op) {
"expects the number of subscripts to be equal to memref rank");
}
+ Region *scope = getAffineScope(op);
for (auto idx : op.getMapOperands()) {
if (!idx.getType().isIndex())
return op.emitOpError("index to load must have 'index' type");
- if (!isValidAffineIndexOperand(idx))
+ if (!isValidAffineIndexOperand(idx, scope))
return op.emitOpError("index must be a dimension or symbol identifier");
}
return success();
@@ -1914,10 +2025,11 @@ LogicalResult verify(AffineStoreOp op) {
"expects the number of subscripts to be equal to memref rank");
}
+ Region *scope = getAffineScope(op);
for (auto idx : op.getMapOperands()) {
if (!idx.getType().isIndex())
return op.emitOpError("index to store must have 'index' type");
- if (!isValidAffineIndexOperand(idx))
+ if (!isValidAffineIndexOperand(idx, scope))
return op.emitOpError("index must be a dimension or symbol identifier");
}
return success();
@@ -2136,8 +2248,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 4ad8b7ed8247..c0855987ac32 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<?x?x?x?xf32>
%dim = dim %0, 0 : memref<?x?x?x?xf32>
- // expected-error at +1 {{operand cannot be used as a dimension id}}
+ // expected-error at +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 cd42980f87c9..d4b3a112dfc6 100644
--- a/mlir/test/Dialect/Affine/ops.mlir
+++ b/mlir/test/Dialect/Affine/ops.mlir
@@ -115,6 +115,44 @@ 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<?xf32>) {
+ 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<?xf32>
+ affine.store %v, %A[%n - 1] : memref<?xf32>
+ }
+ "terminate"() : () -> ()
+ }
+ }
+ "terminate"() : () -> ()
+ }
+ return
+}
+
+// -----
+
+// Test the fact that module op always provides a polyhedral scope.
+
+%idx = "test.foo"() : () -> (index)
+"test.func"() ({
+^bb0(%A : memref<?xf32>):
+ affine.load %A[%idx] : memref<?xf32>
+ "terminate"() : () -> ()
+}) : () -> ()
+
+// -----
+
// 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 048db3a2fcbd..4c67310e3705 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/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 8fb290a76e13..f9140f2e9bdc 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"
@@ -1138,6 +1139,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";
More information about the Mlir-commits
mailing list