diff --git a/mlir/docs/Dialects/Affine.md b/mlir/docs/Dialects/Affine.md --- a/mlir/docs/Dialects/Affine.md +++ b/mlir/docs/Dialects/Affine.md @@ -60,20 +60,26 @@ ### 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 --- a/mlir/docs/Traits.md +++ b/mlir/docs/Traits.md @@ -219,6 +219,22 @@ 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 --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h @@ -31,9 +31,10 @@ 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 @@ SmallVectorImpl &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 --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td @@ -83,12 +83,22 @@ /// 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 --- a/mlir/include/mlir/IR/Function.h +++ b/mlir/include/mlir/IR/Function.h @@ -30,10 +30,11 @@ /// 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/include/mlir/IR/Module.h b/mlir/include/mlir/IR/Module.h --- a/mlir/include/mlir/IR/Module.h +++ b/mlir/include/mlir/IR/Module.h @@ -30,7 +30,8 @@ class ModuleOp : public Op< ModuleOp, OpTrait::ZeroOperands, OpTrait::ZeroResult, - OpTrait::IsIsolatedFromAbove, OpTrait::SymbolTable, + OpTrait::IsIsolatedFromAbove, OpTrait::PolyhedralScope, + OpTrait::SymbolTable, OpTrait::SingleBlockImplicitTerminator::Impl, SymbolOpInterface::Trait> { public: diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td --- a/mlir/include/mlir/IR/OpBase.td +++ b/mlir/include/mlir/IR/OpBase.td @@ -1648,6 +1648,8 @@ 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 --- a/mlir/include/mlir/IR/OpDefinition.h +++ b/mlir/include/mlir/IR/OpDefinition.h @@ -1034,6 +1034,21 @@ } }; +/// 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(); + } +}; + /// 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 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -84,65 +84,110 @@ 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 +197,90 @@ // 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 +295,9 @@ 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 @@ [](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 @@ [](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 operands) { auto map = getAffineMap(); @@ -948,22 +1056,23 @@ 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 @@ 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(); @@ -1814,10 +1924,11 @@ "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 @@ "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 @@ 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 --- a/mlir/test/Dialect/Affine/invalid.mlir +++ b/mlir/test/Dialect/Affine/invalid.mlir @@ -124,7 +124,7 @@ %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 --- a/mlir/test/Dialect/Affine/ops.mlir +++ b/mlir/test/Dialect/Affine/ops.mlir @@ -115,6 +115,33 @@ // ----- +// 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 --- a/mlir/test/lib/Dialect/Test/TestDialect.cpp +++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp @@ -201,6 +201,22 @@ 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 --- 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 @@ 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";