diff --git a/mlir/include/mlir/Analysis/AffineStructures.h b/mlir/include/mlir/Analysis/AffineStructures.h --- a/mlir/include/mlir/Analysis/AffineStructures.h +++ b/mlir/include/mlir/Analysis/AffineStructures.h @@ -19,56 +19,16 @@ namespace mlir { -class AffineApplyOp; -class AffineBound; class AffineCondition; -class AffineMap; class AffineForOp; +class AffineMap; +class AffineValueMap; class IntegerSet; class MLIRContext; class Value; class HyperRectangularSet; class MemRefType; - -/// A mutable affine map. Its affine expressions are however unique. -struct MutableAffineMap { -public: - MutableAffineMap() {} - MutableAffineMap(AffineMap map); - - ArrayRef getResults() const { return results; } - AffineExpr getResult(unsigned idx) const { return results[idx]; } - void setResult(unsigned idx, AffineExpr result) { results[idx] = result; } - unsigned getNumResults() const { return results.size(); } - unsigned getNumDims() const { return numDims; } - void setNumDims(unsigned d) { numDims = d; } - unsigned getNumSymbols() const { return numSymbols; } - void setNumSymbols(unsigned d) { numSymbols = d; } - MLIRContext *getContext() const { return context; } - - /// Returns true if the idx'th result expression is a multiple of factor. - bool isMultipleOf(unsigned idx, int64_t factor) const; - - /// Resets this MutableAffineMap with 'map'. - void reset(AffineMap map); - - /// Simplify the (result) expressions in this map using analysis (used by - //-simplify-affine-expr pass). - void simplify(); - /// Get the AffineMap corresponding to this MutableAffineMap. Note that an - /// AffineMap will be uniqued and stored in context, while a mutable one - /// isn't. - AffineMap getAffineMap() const; - -private: - // Same meaning as AffineMap's fields. - SmallVector results; - unsigned numDims; - unsigned numSymbols; - /// A pointer to the IR's context to store all newly created - /// AffineExprStorage's. - MLIRContext *context; -}; +struct MutableAffineMap; /// A mutable integer set. Its affine expressions are however unique. struct MutableIntegerSet { @@ -96,78 +56,6 @@ SmallVector eqFlags; }; -/// An AffineValueMap is an affine map plus its ML value operands and -/// results for analysis purposes. The structure is still a tree form that is -/// same as that of an affine map or an AffineApplyOp. However, its operands, -/// results, and its map can themselves change as a result of -/// substitutions, simplifications, and other analysis. -// An affine value map can readily be constructed from an AffineApplyOp, or an -// AffineBound of a AffineForOp. It can be further transformed, substituted -// into, or simplified. Unlike AffineMap's, AffineValueMap's are created and -// destroyed during analysis. Only the AffineMap expressions that are pointed by -// them are unique'd. An affine value map, and the operations on it, maintain -// the invariant that operands are always positionally aligned with the -// AffineDimExpr and AffineSymbolExpr in the underlying AffineMap. -// TODO(bondhugula): Some of these classes could go into separate files. -class AffineValueMap { -public: - // Creates an empty AffineValueMap (users should call 'reset' to reset map - // and operands). - AffineValueMap() {} - AffineValueMap(AffineMap map, ArrayRef operands, - ArrayRef results = llvm::None); - - explicit AffineValueMap(AffineApplyOp applyOp); - explicit AffineValueMap(AffineBound bound); - - ~AffineValueMap(); - - // Resets this AffineValueMap with 'map', 'operands', and 'results'. - void reset(AffineMap map, ArrayRef operands, - ArrayRef results = llvm::None); - - /// Return the value map that is the difference of value maps 'a' and 'b', - /// represented as an affine map and its operands. The output map + operands - /// are canonicalized and simplified. - static void difference(const AffineValueMap &a, const AffineValueMap &b, - AffineValueMap *res); - - /// Return true if the idx^th result can be proved to be a multiple of - /// 'factor', false otherwise. - inline bool isMultipleOf(unsigned idx, int64_t factor) const; - - /// Return true if the idx^th result depends on 'value', false otherwise. - bool isFunctionOf(unsigned idx, Value value) const; - - /// Return true if the result at 'idx' is a constant, false - /// otherwise. - bool isConstant(unsigned idx) const; - - /// Return true if this is an identity map. - bool isIdentity() const; - - void setResult(unsigned i, AffineExpr e) { map.setResult(i, e); } - AffineExpr getResult(unsigned i) { return map.getResult(i); } - inline unsigned getNumOperands() const { return operands.size(); } - inline unsigned getNumDims() const { return map.getNumDims(); } - inline unsigned getNumSymbols() const { return map.getNumSymbols(); } - inline unsigned getNumResults() const { return map.getNumResults(); } - - Value getOperand(unsigned i) const; - ArrayRef getOperands() const; - AffineMap getAffineMap() const; - -private: - // A mutable affine map. - MutableAffineMap map; - - // TODO: make these trailing objects? - /// The SSA operands binding to the dim's and symbols of 'map'. - SmallVector operands; - /// The SSA results binding to the results of 'map'. - SmallVector results; -}; - /// An IntegerValueSet is an integer set plus its operands. // Both, the integer set being pointed to and the operands can change during // analysis, simplification, and transformation. diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineOps.h b/mlir/include/mlir/Dialect/AffineOps/AffineOps.h --- a/mlir/include/mlir/Dialect/AffineOps/AffineOps.h +++ b/mlir/include/mlir/Dialect/AffineOps/AffineOps.h @@ -73,6 +73,9 @@ return getAttrOfType("map").getValue(); } + /// 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. bool isValidDim(); @@ -528,6 +531,7 @@ /// 4. propagate constant operands and drop them void canonicalizeMapAndOperands(AffineMap *map, SmallVectorImpl *operands); + /// Canonicalizes an integer set the same way canonicalizeMapAndOperands does /// for affine maps. void canonicalizeSetAndOperands(IntegerSet *set, @@ -573,9 +577,6 @@ AffineForOp getAffineForOp() { return op; } AffineMap getMap() { return map; } - /// Returns an AffineValueMap representing this bound. - AffineValueMap getAsAffineValueMap(); - unsigned getNumOperands() { return opEnd - opStart; } Value getOperand(unsigned idx) { return op.getOperand(opStart + idx); } diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineOps.td b/mlir/include/mlir/Dialect/AffineOps/AffineOps.td --- a/mlir/include/mlir/Dialect/AffineOps/AffineOps.td +++ b/mlir/include/mlir/Dialect/AffineOps/AffineOps.td @@ -271,6 +271,76 @@ }]; } +def AffineParallelOp : Affine_Op<"parallel", [ImplicitAffineTerminator]> { + let summary = "multi-index parallel band operation"; + let description = [{ + The "affine.parallel" operation represents an affine parallel band, defining + multiple SSA values for its induction variables. It has one region capturing + the parallel band body. The induction variables are represented as arguments + of this region. These SSA values always have type index, which is the size + of the machine word. The strides, represented by steps, are positive + constant integers which defaults to "1" if not present. The lower and upper + bounds specify a half-open range: the range includes the lower bound but + does not include the upper bound. + + The body region must contain exactly one block that terminates with + "affine.terminator". Calling AffineParallelOp::build will create such + region and insert the terminator, so will the parsing even in cases if it is + absent from the custom format. + + The lower and upper bounds of a parallel operation are represented as an + application of an affine mapping to a list of SSA values passed to the map. + The same restrictions hold for these SSA values as for all bindings of SSA + values to dimensions and symbols. + + Example: + + affine.parallel (%i, %j) = [0, 0] to [10, 10] step (1, 1) { + ... + } + }]; + + let arguments = (ins + AffineMapAttr:$lowerBoundsMap, + AffineMapAttr:$upperBoundsMap, + I64ArrayAttr:$steps, + Variadic:$mapOperands); + let regions = (region SizedRegion<1>:$region); + + let builders = [ + OpBuilder<"Builder* builder, OperationState& result," + "ArrayRef ranges">, + OpBuilder<"Builder* builder, OperationState& result, AffineMap lbMap," + "ValueRange lbArgs, AffineMap ubMap, ValueRange ubArgs">, + OpBuilder<"Builder* builder, OperationState& result, AffineMap lbMap," + "ValueRange lbArgs, AffineMap ubMap, ValueRange ubArgs," + "ArrayRef steps"> + ]; + + let extraClassDeclaration = [{ + /// Get the number of dimensions. + unsigned getNumDims(); + + operand_range getLowerBoundsOperands(); + operand_range getUpperBoundsOperands(); + + AffineValueMap getLowerBoundsValueMap(); + AffineValueMap getUpperBoundsValueMap(); + AffineValueMap getRangesValueMap(); + + /// Get ranges as constants, may fail in dynamic case. + llvm::Optional> getConstantRanges(); + + mlir::Block *getBody(); + mlir::OpBuilder getBodyBuilder(); + void setSteps(ArrayRef newSteps); + + static StringRef getLowerBoundsMapAttrName() { return "lowerBoundsMap"; } + static StringRef getUpperBoundsMapAttrName() { return "upperBoundsMap"; } + static StringRef getStepsAttrName() { return "steps"; } + }]; +} + def AffinePrefetchOp : Affine_Op<"prefetch"> { let summary = "affine prefetch operation"; let description = [{ diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineValueMap.h b/mlir/include/mlir/Dialect/AffineOps/AffineValueMap.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/AffineOps/AffineValueMap.h @@ -0,0 +1,90 @@ +//===- AffineValueMap.h - MLIR Affine Value Map Class -----------*- 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 +// +//===----------------------------------------------------------------------===// +// +// An AffineValueMap is an affine map plus its ML value operands and results for +// analysis purposes. +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_AFFINEOPS_AFFINEVALUEMAP_H +#define MLIR_DIALECT_AFFINEOPS_AFFINEVALUEMAP_H + +#include "mlir/IR/AffineMap.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/IR/Value.h" + +namespace mlir { + +/// An AffineValueMap is an affine map plus its ML value operands and +/// results for analysis purposes. The structure is still a tree form that is +/// same as that of an affine map or an AffineApplyOp. However, its operands, +/// results, and its map can themselves change as a result of +/// substitutions, simplifications, and other analysis. +// An affine value map can readily be constructed from an AffineApplyOp, or an +// AffineBound of a AffineForOp. It can be further transformed, substituted +// into, or simplified. Unlike AffineMap's, AffineValueMap's are created and +// destroyed during analysis. Only the AffineMap expressions that are pointed by +// them are unique'd. An affine value map, and the operations on it, maintain +// the invariant that operands are always positionally aligned with the +// AffineDimExpr and AffineSymbolExpr in the underlying AffineMap. +class AffineValueMap { +public: + // Creates an empty AffineValueMap (users should call 'reset' to reset map + // and operands). + AffineValueMap() {} + AffineValueMap(AffineMap map, ValueRange operands, ValueRange results = {}); + + ~AffineValueMap(); + + // Resets this AffineValueMap with 'map', 'operands', and 'results'. + void reset(AffineMap map, ValueRange operands, ValueRange results = {}); + + /// Return the value map that is the difference of value maps 'a' and 'b', + /// represented as an affine map and its operands. The output map + operands + /// are canonicalized and simplified. + static void difference(const AffineValueMap &a, const AffineValueMap &b, + AffineValueMap *res); + + /// Return true if the idx^th result can be proved to be a multiple of + /// 'factor', false otherwise. + inline bool isMultipleOf(unsigned idx, int64_t factor) const; + + /// Return true if the idx^th result depends on 'value', false otherwise. + bool isFunctionOf(unsigned idx, Value value) const; + + /// Return true if the result at 'idx' is a constant, false + /// otherwise. + bool isConstant(unsigned idx) const; + + /// Return true if this is an identity map. + bool isIdentity() const; + + void setResult(unsigned i, AffineExpr e) { map.setResult(i, e); } + AffineExpr getResult(unsigned i) { return map.getResult(i); } + inline unsigned getNumOperands() const { return operands.size(); } + inline unsigned getNumDims() const { return map.getNumDims(); } + inline unsigned getNumSymbols() const { return map.getNumSymbols(); } + inline unsigned getNumResults() const { return map.getNumResults(); } + + Value getOperand(unsigned i) const; + ArrayRef getOperands() const; + AffineMap getAffineMap() const; + +private: + // A mutable affine map. + MutableAffineMap map; + + // TODO: make these trailing objects? + /// The SSA operands binding to the dim's and symbols of 'map'. + SmallVector operands; + /// The SSA results binding to the results of 'map'. + SmallVector results; +}; + +} // namespace mlir + +#endif // MLIR_DIALECT_AFFINEOPS_AFFINEVALUEMAP_H diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h --- a/mlir/include/mlir/IR/AffineMap.h +++ b/mlir/include/mlir/IR/AffineMap.h @@ -14,6 +14,7 @@ #ifndef MLIR_IR_AFFINE_MAP_H #define MLIR_IR_AFFINE_MAP_H +#include "mlir/IR/AffineExpr.h" #include "mlir/Support/LLVM.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMapInfo.h" @@ -24,7 +25,6 @@ struct AffineMapStorage; } // end namespace detail -class AffineExpr; class Attribute; struct LogicalResult; class MLIRContext; @@ -155,6 +155,46 @@ return ::llvm::hash_value(arg.map); } +/// A mutable affine map. Its affine expressions are however unique. +struct MutableAffineMap { +public: + MutableAffineMap() {} + MutableAffineMap(AffineMap map); + + ArrayRef getResults() const { return results; } + AffineExpr getResult(unsigned idx) const { return results[idx]; } + void setResult(unsigned idx, AffineExpr result) { results[idx] = result; } + unsigned getNumResults() const { return results.size(); } + unsigned getNumDims() const { return numDims; } + void setNumDims(unsigned d) { numDims = d; } + unsigned getNumSymbols() const { return numSymbols; } + void setNumSymbols(unsigned d) { numSymbols = d; } + MLIRContext *getContext() const { return context; } + + /// Returns true if the idx'th result expression is a multiple of factor. + bool isMultipleOf(unsigned idx, int64_t factor) const; + + /// Resets this MutableAffineMap with 'map'. + void reset(AffineMap map); + + /// Simplify the (result) expressions in this map using analysis (used by + //-simplify-affine-expr pass). + void simplify(); + /// Get the AffineMap corresponding to this MutableAffineMap. Note that an + /// AffineMap will be uniqued and stored in context, while a mutable one + /// isn't. + AffineMap getAffineMap() const; + +private: + // Same meaning as AffineMap's fields. + SmallVector results; + unsigned numDims; + unsigned numSymbols; + /// A pointer to the IR's context to store all newly created + /// AffineExprStorage's. + MLIRContext *context; +}; + /// Simplify an affine map by simplifying its underlying AffineExpr results. AffineMap simplifyAffineMap(AffineMap map); @@ -227,7 +267,8 @@ namespace llvm { // AffineExpr hash just like pointers -template <> struct DenseMapInfo { +template <> +struct DenseMapInfo { static mlir::AffineMap getEmptyKey() { auto pointer = llvm::DenseMapInfo::getEmptyKey(); return mlir::AffineMap(static_cast(pointer)); diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp --- a/mlir/lib/Analysis/AffineAnalysis.cpp +++ b/mlir/lib/Analysis/AffineAnalysis.cpp @@ -15,6 +15,7 @@ #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/Utils.h" #include "mlir/Dialect/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineValueMap.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/Builders.h" diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp --- a/mlir/lib/Analysis/AffineStructures.cpp +++ b/mlir/lib/Analysis/AffineStructures.cpp @@ -12,6 +12,7 @@ #include "mlir/Analysis/AffineStructures.h" #include "mlir/Dialect/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineValueMap.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/AffineExprVisitor.h" #include "mlir/IR/IntegerSet.h" @@ -136,51 +137,6 @@ localVarCst); } -//===----------------------------------------------------------------------===// -// MutableAffineMap. -//===----------------------------------------------------------------------===// - -MutableAffineMap::MutableAffineMap(AffineMap map) - : numDims(map.getNumDims()), numSymbols(map.getNumSymbols()), - // A map always has at least 1 result by construction - context(map.getResult(0).getContext()) { - for (auto result : map.getResults()) - results.push_back(result); -} - -void MutableAffineMap::reset(AffineMap map) { - results.clear(); - numDims = map.getNumDims(); - numSymbols = map.getNumSymbols(); - // A map always has at least 1 result by construction - context = map.getResult(0).getContext(); - for (auto result : map.getResults()) - results.push_back(result); -} - -bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const { - if (results[idx].isMultipleOf(factor)) - return true; - - // TODO(bondhugula): use simplifyAffineExpr and FlatAffineConstraints to - // complete this (for a more powerful analysis). - return false; -} - -// Simplifies the result affine expressions of this map. The expressions have to -// be pure for the simplification implemented. -void MutableAffineMap::simplify() { - // Simplify each of the results if possible. - // TODO(ntv): functional-style map - for (unsigned i = 0, e = getNumResults(); i < e; i++) { - results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols); - } -} - -AffineMap MutableAffineMap::getAffineMap() const { - return AffineMap::get(numDims, numSymbols, results); -} - MutableIntegerSet::MutableIntegerSet(IntegerSet set, MLIRContext *context) : numDims(set.getNumDims()), numSymbols(set.getNumSymbols()) { // TODO(bondhugula) @@ -191,110 +147,6 @@ MLIRContext *context) : numDims(numDims), numSymbols(numSymbols) {} -//===----------------------------------------------------------------------===// -// AffineValueMap. -//===----------------------------------------------------------------------===// - -AffineValueMap::AffineValueMap(AffineMap map, ArrayRef operands, - ArrayRef results) - : map(map), operands(operands.begin(), operands.end()), - results(results.begin(), results.end()) {} - -AffineValueMap::AffineValueMap(AffineApplyOp applyOp) - : map(applyOp.getAffineMap()), - operands(applyOp.operand_begin(), applyOp.operand_end()) { - results.push_back(applyOp.getResult()); -} - -AffineValueMap::AffineValueMap(AffineBound bound) - : map(bound.getMap()), - operands(bound.operand_begin(), bound.operand_end()) {} - -void AffineValueMap::reset(AffineMap map, ArrayRef operands, - ArrayRef results) { - this->map.reset(map); - this->operands.assign(operands.begin(), operands.end()); - this->results.assign(results.begin(), results.end()); -} - -void AffineValueMap::difference(const AffineValueMap &a, - const AffineValueMap &b, AffineValueMap *res) { - assert(a.getNumResults() == b.getNumResults() && "invalid inputs"); - - // Fully compose A's map + operands. - auto aMap = a.getAffineMap(); - SmallVector aOperands(a.getOperands().begin(), - a.getOperands().end()); - fullyComposeAffineMapAndOperands(&aMap, &aOperands); - - // Use the affine apply normalizer to get B's map into A's coordinate space. - AffineApplyNormalizer normalizer(aMap, aOperands); - SmallVector bOperands(b.getOperands().begin(), - b.getOperands().end()); - auto bMap = b.getAffineMap(); - normalizer.normalize(&bMap, &bOperands); - - assert(std::equal(bOperands.begin(), bOperands.end(), - normalizer.getOperands().begin()) && - "operands are expected to be the same after normalization"); - - // Construct the difference expressions. - SmallVector diffExprs; - diffExprs.reserve(a.getNumResults()); - for (unsigned i = 0, e = bMap.getNumResults(); i < e; ++i) - diffExprs.push_back(normalizer.getAffineMap().getResult(i) - - bMap.getResult(i)); - - auto diffMap = AffineMap::get(normalizer.getNumDims(), - normalizer.getNumSymbols(), diffExprs); - canonicalizeMapAndOperands(&diffMap, &bOperands); - diffMap = simplifyAffineMap(diffMap); - res->reset(diffMap, bOperands); -} - -// Returns true and sets 'indexOfMatch' if 'valueToMatch' is found in -// 'valuesToSearch' beginning at 'indexStart'. Returns false otherwise. -static bool findIndex(Value valueToMatch, ArrayRef valuesToSearch, - unsigned indexStart, unsigned *indexOfMatch) { - unsigned size = valuesToSearch.size(); - for (unsigned i = indexStart; i < size; ++i) { - if (valueToMatch == valuesToSearch[i]) { - *indexOfMatch = i; - return true; - } - } - return false; -} - -inline bool AffineValueMap::isMultipleOf(unsigned idx, int64_t factor) const { - return map.isMultipleOf(idx, factor); -} - -/// This method uses the invariant that operands are always positionally aligned -/// with the AffineDimExpr in the underlying AffineMap. -bool AffineValueMap::isFunctionOf(unsigned idx, Value value) const { - unsigned index; - if (!findIndex(value, operands, /*indexStart=*/0, &index)) { - return false; - } - auto expr = const_cast(this)->getAffineMap().getResult(idx); - // TODO(ntv): this is better implemented on a flattened representation. - // At least for now it is conservative. - return expr.isFunctionOfDim(index); -} - -Value AffineValueMap::getOperand(unsigned i) const { - return static_cast(operands[i]); -} - -ArrayRef AffineValueMap::getOperands() const { - return ArrayRef(operands); -} - -AffineMap AffineValueMap::getAffineMap() const { return map.getAffineMap(); } - -AffineValueMap::~AffineValueMap() {} - //===----------------------------------------------------------------------===// // FlatAffineConstraints. //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Analysis/LoopAnalysis.cpp b/mlir/lib/Analysis/LoopAnalysis.cpp --- a/mlir/lib/Analysis/LoopAnalysis.cpp +++ b/mlir/lib/Analysis/LoopAnalysis.cpp @@ -16,6 +16,7 @@ #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/NestedMatcher.h" #include "mlir/Dialect/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineValueMap.h" #include "mlir/Support/MathExtras.h" #include "llvm/ADT/DenseSet.h" @@ -185,7 +186,7 @@ auto composeOp = cast(affineApplyOps[0]); // We need yet another level of indirection because the `dim` index of the // access may not correspond to the `dim` index of composeOp. - return !(AffineValueMap(composeOp).isFunctionOf(0, iv)); + return !composeOp.getAffineValueMap().isFunctionOf(0, iv); } DenseSet mlir::getInvariantAccesses(Value iv, ArrayRef indices) { diff --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp --- a/mlir/lib/Analysis/Utils.cpp +++ b/mlir/lib/Analysis/Utils.cpp @@ -15,6 +15,7 @@ #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Dialect/AffineOps/AffineOps.h" +#include "mlir/Dialect/AffineOps/AffineValueMap.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/Support/Debug.h" diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp --- a/mlir/lib/Dialect/AffineOps/AffineOps.cpp +++ b/mlir/lib/Dialect/AffineOps/AffineOps.cpp @@ -7,6 +7,8 @@ //===----------------------------------------------------------------------===// #include "mlir/Dialect/AffineOps/AffineOps.h" +#include "mlir/Analysis/AffineStructures.h" +#include "mlir/Dialect/AffineOps/AffineValueMap.h" #include "mlir/Dialect/StandardOps/Ops.h" #include "mlir/IR/Function.h" #include "mlir/IR/IntegerSet.h" @@ -133,9 +135,11 @@ return isTopLevelValue(dimOp.getOperand()); return false; } - // This value has to be a block argument for a FuncOp or an affine.for. + // 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); + return isa(parentOp) || isa(parentOp) || + isa(parentOp); } /// Returns true if the 'index' dimension of the `memref` defined by @@ -235,6 +239,10 @@ result.addAttribute("map", AffineMapAttr::get(map)); } +AffineValueMap AffineApplyOp::getAffineValueMap() { + return AffineValueMap(getAffineMap(), getOperands(), getResult()); +} + ParseResult AffineApplyOp::parse(OpAsmParser &parser, OperationState &result) { auto &builder = parser.getBuilder(); auto indexTy = builder.getIndexType(); @@ -1332,7 +1340,7 @@ } static void print(OpAsmPrinter &p, AffineForOp op) { - p << "affine.for "; + p << op.getOperationName() << ' '; p.printOperand(op.getBody()->getArgument(0)); p << " = "; printBound(op.getLowerBoundMapAttr(), op.getLowerBoundOperands(), "max", p); @@ -2145,6 +2153,230 @@ return foldMemRefCast(*this); } +//===----------------------------------------------------------------------===// +// AffineParallelOp +//===----------------------------------------------------------------------===// + +void AffineParallelOp::build(Builder *builder, OperationState &result, + ArrayRef ranges) { + // Default initalize empty maps + auto lbMap = AffineMap::get(builder->getContext()); + auto ubMap = AffineMap::get(builder->getContext()); + // If ranges, set to [0, N) for each range + if (ranges.size()) { + SmallVector lbExprs; + SmallVector ubExprs; + // Make range expressions for each range + for (int64_t range : ranges) { + lbExprs.push_back(builder->getAffineConstantExpr(0)); + ubExprs.push_back(builder->getAffineConstantExpr(range)); + } + lbMap = AffineMap::get(0, 0, lbExprs); + ubMap = AffineMap::get(0, 0, ubExprs); + } + // Fall through + build(builder, result, lbMap, {}, ubMap, {}); +} + +void AffineParallelOp::build(Builder *builder, OperationState &result, + AffineMap lbMap, ValueRange lbArgs, + AffineMap ubMap, ValueRange ubArgs) { + // Verify sizes + auto dims = lbMap.getNumResults(); + assert(dims == ubMap.getNumResults()); + // Make default step sizes of 1 + SmallVector steps(dims, 1); + // Call through + build(builder, result, lbMap, lbArgs, ubMap, ubArgs, steps); +} + +void AffineParallelOp::build(Builder *builder, OperationState &result, + AffineMap lbMap, ValueRange lbArgs, + AffineMap ubMap, ValueRange ubArgs, + ArrayRef steps) { + // Verify sizes + auto dims = lbMap.getNumResults(); + assert(dims == ubMap.getNumResults()); + assert(dims == steps.size()); + // Set all of the attributes + result.addAttribute(getLowerBoundsMapAttrName(), AffineMapAttr::get(lbMap)); + result.addAttribute(getUpperBoundsMapAttrName(), AffineMapAttr::get(ubMap)); + result.addAttribute(getStepsAttrName(), builder->getI64ArrayAttr(steps)); + result.addOperands(lbArgs); + result.addOperands(ubArgs); + // Create a region and a block for the body. + auto bodyRegion = result.addRegion(); + auto body = new Block(); + // Add all the args + for (unsigned i = 0; i < dims; ++i) { + body->addArgument(IndexType::get(builder->getContext())); + } + bodyRegion->push_back(body); + // Terminate + ensureTerminator(*bodyRegion, *builder, result.location); +} + +AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() { + return {operand_begin(), operand_begin() + lowerBoundsMap().getNumInputs()}; +} + +AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() { + return {operand_begin() + lowerBoundsMap().getNumInputs(), operand_end()}; +} + +unsigned AffineParallelOp::getNumDims() { return steps().size(); } + +AffineValueMap AffineParallelOp::getLowerBoundsValueMap() { + return AffineValueMap(lowerBoundsMap(), getLowerBoundsOperands()); +} + +AffineValueMap AffineParallelOp::getUpperBoundsValueMap() { + return AffineValueMap(upperBoundsMap(), getUpperBoundsOperands()); +} + +AffineValueMap AffineParallelOp::getRangesValueMap() { + AffineValueMap out; + AffineValueMap::difference(getUpperBoundsValueMap(), getLowerBoundsValueMap(), + &out); + return out; +} + +Optional> AffineParallelOp::getConstantRanges() { + SmallVector out; + // Get the ranges + AffineValueMap rangesValueMap = getRangesValueMap(); + // Try to convert them to constants + for (unsigned i = 0, e = rangesValueMap.getNumResults(); i < e; ++i) { + auto expr = rangesValueMap.getResult(i); + if (auto cst = expr.dyn_cast()) { + out.push_back(cst.getValue()); + } else { + return llvm::None; + } + } + return out; +} + +mlir::Block *AffineParallelOp::getBody() { return ®ion().front(); } + +mlir::OpBuilder AffineParallelOp::getBodyBuilder() { + return mlir::OpBuilder(getBody(), std::prev(getBody()->end())); +} + +void AffineParallelOp::setSteps(ArrayRef newSteps) { + assert(newSteps.size() == getNumDims()); + setAttr("steps", getBodyBuilder().getI64ArrayAttr(newSteps)); +} + +static LogicalResult verify(AffineParallelOp op) { + auto dims = op.getNumDims(); + if (op.lowerBoundsMap().getNumResults() != dims || + op.upperBoundsMap().getNumResults() != dims || + op.steps().size() != dims || op.getBody()->getNumArguments() != dims) { + return op.emitOpError("region argument count and num results of upper " + "bounds, lower bounds, and steps must all match"); + } + return success(); +} + +static void print(OpAsmPrinter &p, AffineParallelOp op) { + p << op.getOperationName() << " ["; + p.printOperands(op.getBody()->getArguments()); + p << "] = ["; + p.printAffineMapOfSSAIds(op.lowerBoundsMapAttr(), + op.getLowerBoundsOperands()); + p << "] to ["; + p.printAffineMapOfSSAIds(op.upperBoundsMapAttr(), + op.getUpperBoundsOperands()); + p << ']'; + bool elideSteps = true; + for (auto attr : op.steps()) { + auto step = attr.cast().getValue(); + if (step != 1) { + elideSteps = false; + break; + } + } + if (!elideSteps) { + p << " step " << op.steps(); + } + p.printRegion(op.region(), /*printEntryBlockArgs=*/false, + /*printBlockTerminators=*/false); + p.printOptionalAttrDict( + op.getAttrs(), + /*elidedAttrs=*/{AffineParallelOp::getLowerBoundsMapAttrName(), + AffineParallelOp::getUpperBoundsMapAttrName(), + AffineParallelOp::getStepsAttrName()}); +} + +// +// operation ::= `affine.parallel` `[` ssa-ids `]` `=` `[` map-of-ssa-ids `]` +// `to` `[` map-of-ssa-ids `]` steps? region attr-dict? +// steps ::= `steps` `[` integer-literals `]` +// +static ParseResult parseAffineParallelOp(OpAsmParser &parser, + OperationState &result) { + auto &builder = parser.getBuilder(); + auto indexType = builder.getIndexType(); + AffineMapAttr lowerBoundsAttr, upperBoundsAttr; + SmallVector ivs, lowerBoundsMapOperands, + upperBoundsMapOperands; + if (parser.parseRegionArgumentList(ivs, /*requiredOperandCount=*/-1, + OpAsmParser::Delimiter::Square) || + parser.parseEqual() || + parser.parseAffineMapOfSSAIds( + lowerBoundsMapOperands, lowerBoundsAttr, + AffineParallelOp::getLowerBoundsMapAttrName(), result.attributes) || + parser.resolveOperands(lowerBoundsMapOperands, indexType, + result.operands) || + parser.parseKeyword("to") || + parser.parseAffineMapOfSSAIds( + upperBoundsMapOperands, upperBoundsAttr, + AffineParallelOp::getUpperBoundsMapAttrName(), result.attributes) || + parser.resolveOperands(upperBoundsMapOperands, indexType, + result.operands)) + return failure(); + + AffineMapAttr stepsMapAttr; + SmallVector stepsAttrs; + SmallVector stepsMapOperands; + if (parser.parseOptionalKeyword("step")) { + SmallVector steps(ivs.size(), 1); + result.addAttribute(AffineParallelOp::getStepsAttrName(), + builder.getI64ArrayAttr(steps)); + } else { + if (parser.parseAffineMapOfSSAIds(stepsMapOperands, stepsMapAttr, + AffineParallelOp::getStepsAttrName(), + stepsAttrs)) + return failure(); + + // Convert steps from an AffineMap into an I64ArrayAttr. + SmallVector steps; + auto stepsMap = stepsMapAttr.getValue(); + for (const auto &result : stepsMap.getResults()) { + auto constExpr = result.dyn_cast(); + if (!constExpr) + return parser.emitError(parser.getNameLoc(), + "steps must be constant integers"); + steps.push_back(constExpr.getValue()); + } + result.addAttribute(AffineParallelOp::getStepsAttrName(), + builder.getI64ArrayAttr(steps)); + } + + // Now parse the body. + Region *body = result.addRegion(); + SmallVector types(ivs.size(), indexType); + if (parser.parseRegion(*body, ivs, types) || + parser.parseOptionalAttrDict(result.attributes)) + return failure(); + + // Add a terminator if none was parsed. + AffineParallelOp::ensureTerminator(*body, builder, result.location); + + return success(); +} + //===----------------------------------------------------------------------===// // TableGen'd op method definitions //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/AffineOps/AffineValueMap.cpp b/mlir/lib/Dialect/AffineOps/AffineValueMap.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/AffineOps/AffineValueMap.cpp @@ -0,0 +1,102 @@ +//===- AffineValueMap.cpp - MLIR Affine Value Map Class -------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/AffineOps/AffineValueMap.h" +#include "mlir/Dialect/AffineOps/AffineOps.h" + +using namespace mlir; + +AffineValueMap::AffineValueMap(AffineMap map, ValueRange operands, + ValueRange results) + : map(map), operands(operands.begin(), operands.end()), + results(results.begin(), results.end()) {} + +void AffineValueMap::reset(AffineMap map, ValueRange operands, + ValueRange results) { + this->map.reset(map); + this->operands.assign(operands.begin(), operands.end()); + this->results.assign(results.begin(), results.end()); +} + +void AffineValueMap::difference(const AffineValueMap &a, + const AffineValueMap &b, AffineValueMap *res) { + assert(a.getNumResults() == b.getNumResults() && "invalid inputs"); + + // Fully compose A's map + operands. + auto aMap = a.getAffineMap(); + SmallVector aOperands(a.getOperands().begin(), + a.getOperands().end()); + fullyComposeAffineMapAndOperands(&aMap, &aOperands); + + // Use the affine apply normalizer to get B's map into A's coordinate space. + AffineApplyNormalizer normalizer(aMap, aOperands); + SmallVector bOperands(b.getOperands().begin(), + b.getOperands().end()); + auto bMap = b.getAffineMap(); + normalizer.normalize(&bMap, &bOperands); + + assert(std::equal(bOperands.begin(), bOperands.end(), + normalizer.getOperands().begin()) && + "operands are expected to be the same after normalization"); + + // Construct the difference expressions. + SmallVector diffExprs; + diffExprs.reserve(a.getNumResults()); + for (unsigned i = 0, e = bMap.getNumResults(); i < e; ++i) + diffExprs.push_back(normalizer.getAffineMap().getResult(i) - + bMap.getResult(i)); + + auto diffMap = AffineMap::get(normalizer.getNumDims(), + normalizer.getNumSymbols(), diffExprs); + canonicalizeMapAndOperands(&diffMap, &bOperands); + diffMap = simplifyAffineMap(diffMap); + res->reset(diffMap, bOperands); +} + +// Returns true and sets 'indexOfMatch' if 'valueToMatch' is found in +// 'valuesToSearch' beginning at 'indexStart'. Returns false otherwise. +static bool findIndex(Value valueToMatch, ArrayRef valuesToSearch, + unsigned indexStart, unsigned *indexOfMatch) { + unsigned size = valuesToSearch.size(); + for (unsigned i = indexStart; i < size; ++i) { + if (valueToMatch == valuesToSearch[i]) { + *indexOfMatch = i; + return true; + } + } + return false; +} + +inline bool AffineValueMap::isMultipleOf(unsigned idx, int64_t factor) const { + return map.isMultipleOf(idx, factor); +} + +/// This method uses the invariant that operands are always positionally aligned +/// with the AffineDimExpr in the underlying AffineMap. +bool AffineValueMap::isFunctionOf(unsigned idx, Value value) const { + unsigned index; + if (!findIndex(value, operands, /*indexStart=*/0, &index)) { + return false; + } + auto expr = const_cast(this)->getAffineMap().getResult(idx); + // TODO(ntv): this is better implemented on a flattened representation. + // At least for now it is conservative. + return expr.isFunctionOfDim(index); +} + +Value AffineValueMap::getOperand(unsigned i) const { + return static_cast(operands[i]); +} + +ArrayRef AffineValueMap::getOperands() const { + return ArrayRef(operands); +} + +AffineMap AffineValueMap::getAffineMap() const { return map.getAffineMap(); } + +AffineValueMap::~AffineValueMap() {} diff --git a/mlir/lib/Dialect/AffineOps/CMakeLists.txt b/mlir/lib/Dialect/AffineOps/CMakeLists.txt --- a/mlir/lib/Dialect/AffineOps/CMakeLists.txt +++ b/mlir/lib/Dialect/AffineOps/CMakeLists.txt @@ -1,5 +1,6 @@ add_llvm_library(MLIRAffineOps AffineOps.cpp + AffineValueMap.cpp DialectRegistration.cpp ADDITIONAL_HEADER_DIRS @@ -11,4 +12,3 @@ MLIRLoopLikeInterfaceIncGen MLIRStandardOps) target_link_libraries(MLIRAffineOps MLIRIR MLIRStandardOps) - diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp --- a/mlir/lib/IR/AffineMap.cpp +++ b/mlir/lib/IR/AffineMap.cpp @@ -326,3 +326,48 @@ } return numDims == 0 ? AffineMap() : AffineMap::get(numDims, 0, results); } + +//===----------------------------------------------------------------------===// +// MutableAffineMap. +//===----------------------------------------------------------------------===// + +MutableAffineMap::MutableAffineMap(AffineMap map) + : numDims(map.getNumDims()), numSymbols(map.getNumSymbols()), + // A map always has at least 1 result by construction + context(map.getResult(0).getContext()) { + for (auto result : map.getResults()) + results.push_back(result); +} + +void MutableAffineMap::reset(AffineMap map) { + results.clear(); + numDims = map.getNumDims(); + numSymbols = map.getNumSymbols(); + // A map always has at least 1 result by construction + context = map.getResult(0).getContext(); + for (auto result : map.getResults()) + results.push_back(result); +} + +bool MutableAffineMap::isMultipleOf(unsigned idx, int64_t factor) const { + if (results[idx].isMultipleOf(factor)) + return true; + + // TODO(bondhugula): use simplifyAffineExpr and FlatAffineConstraints to + // complete this (for a more powerful analysis). + return false; +} + +// Simplifies the result affine expressions of this map. The expressions have to +// be pure for the simplification implemented. +void MutableAffineMap::simplify() { + // Simplify each of the results if possible. + // TODO(ntv): functional-style map + for (unsigned i = 0, e = getNumResults(); i < e; i++) { + results[i] = simplifyAffineExpr(getResult(i), numDims, numSymbols); + } +} + +AffineMap MutableAffineMap::getAffineMap() const { + return AffineMap::get(numDims, numSymbols, results); +} diff --git a/mlir/test/AffineOps/invalid.mlir b/mlir/test/AffineOps/invalid.mlir --- a/mlir/test/AffineOps/invalid.mlir +++ b/mlir/test/AffineOps/invalid.mlir @@ -197,4 +197,40 @@ %0 = affine.max affine_map<(d0) -> (d0)> () return -} \ No newline at end of file +} + +// ----- + +// CHECK-LABEL: @affine_parallel +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}} + affine.parallel [%i0] = [0, 0] to [100, 100] step [10, 10] { + } +} + +// ----- + +// CHECK-LABEL: @affine_parallel +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}} + affine.parallel [%i0, %j0] = [0] to [100, 100] step [10, 10] { + } +} + +// ----- + +// CHECK-LABEL: @affine_parallel +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}} + affine.parallel [%i0, %j0] = [0, 0] to [100] step [10, 10] { + } +} + +// ----- + +// CHECK-LABEL: @affine_parallel +func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + // expected-error@+1 {{region argument count and num results of upper bounds, lower bounds, and steps must all match}} + affine.parallel [%i0, %j0] = [0, 0] to [100, 100] step [10] { + } +} diff --git a/mlir/test/AffineOps/ops.mlir b/mlir/test/AffineOps/ops.mlir --- a/mlir/test/AffineOps/ops.mlir +++ b/mlir/test/AffineOps/ops.mlir @@ -112,3 +112,17 @@ } return } + +// ----- + +// CHECK-LABEL: @parallel +// CHECK-SAME: (%[[N:.*]]: index) +func @parallel(%N : index) { + // CHECK: affine.parallel [%[[I0:.*]], %[[J0:.*]]] = [0, 0] to [symbol(%[[N]]), 100] step [10, 10] + affine.parallel [%i0, %j0] = [0, 0] to [symbol(%N), 100] step [10, 10] { + // CHECK-NEXT: affine.parallel [%{{.*}}, %{{.*}}] = [%[[I0]], %[[J0]]] to [%[[I0]] + 10, %[[J0]] + 10] + affine.parallel [%i1, %j1] = [%i0, %j0] to [%i0 + 10, %j0 + 10] { + } + } + return +}