diff options
Diffstat (limited to 'mlir')
| -rw-r--r-- | mlir/include/mlir/AffineOps/AffineOps.h | 61 | ||||
| -rw-r--r-- | mlir/include/mlir/IR/BuiltinOps.h | 55 | ||||
| -rw-r--r-- | mlir/include/mlir/IR/Value.h | 9 | ||||
| -rw-r--r-- | mlir/include/mlir/Transforms/Utils.h | 1 | ||||
| -rw-r--r-- | mlir/lib/AffineOps/AffineOps.cpp | 242 | ||||
| -rw-r--r-- | mlir/lib/Analysis/AffineAnalysis.cpp | 6 | ||||
| -rw-r--r-- | mlir/lib/Analysis/AffineStructures.cpp | 2 | ||||
| -rw-r--r-- | mlir/lib/Analysis/Utils.cpp | 2 | ||||
| -rw-r--r-- | mlir/lib/IR/BuiltinOps.cpp | 208 | ||||
| -rw-r--r-- | mlir/lib/IR/Instruction.cpp | 36 | ||||
| -rw-r--r-- | mlir/lib/Transforms/ComposeAffineMaps.cpp | 1 | ||||
| -rw-r--r-- | mlir/lib/Transforms/Utils/Utils.cpp | 2 | ||||
| -rw-r--r-- | mlir/lib/Transforms/Vectorization/VectorizerTestPass.cpp | 1 |
13 files changed, 310 insertions, 316 deletions
diff --git a/mlir/include/mlir/AffineOps/AffineOps.h b/mlir/include/mlir/AffineOps/AffineOps.h index 0e71221871b..43ad0354b01 100644 --- a/mlir/include/mlir/AffineOps/AffineOps.h +++ b/mlir/include/mlir/AffineOps/AffineOps.h @@ -35,6 +35,55 @@ public: AffineOpsDialect(MLIRContext *context); }; +/// The "affine_apply" operation applies an affine map to a list of operands, +/// yielding a single result. The operand list must be the same size as the +/// number of arguments to the affine mapping. All operands and the result are +/// of type 'Index'. This operation requires a single affine map attribute named +/// "map". For example: +/// +/// %y = "affine_apply" (%x) { map: (d0) -> (d0 + 1) } : +/// (index) -> (index) +/// +/// equivalently: +/// +/// #map42 = (d0)->(d0+1) +/// %y = affine_apply #map42(%x) +/// +class AffineApplyOp : public Op<AffineApplyOp, OpTrait::VariadicOperands, + OpTrait::OneResult, OpTrait::HasNoSideEffect> { +public: + /// Builds an affine apply op with the specified map and operands. + static void build(Builder *builder, OperationState *result, AffineMap map, + ArrayRef<Value *> operands); + + /// Returns the affine map to be applied by this operation. + AffineMap getAffineMap() const { + return getAttrOfType<AffineMapAttr>("map").getValue(); + } + + /// Returns true if the result of this operation can be used as dimension id. + bool isValidDim() const; + + /// Returns true if the result of this operation is a symbol. + bool isValidSymbol() const; + + static StringRef getOperationName() { return "affine_apply"; } + + // Hooks to customize behavior of this op. + static bool parse(OpAsmParser *parser, OperationState *result); + void print(OpAsmPrinter *p) const; + bool verify() const; + Attribute constantFold(ArrayRef<Attribute> operands, + MLIRContext *context) const; + + static void getCanonicalizationPatterns(OwningRewritePatternList &results, + MLIRContext *context); + +private: + friend class Instruction; + explicit AffineApplyOp(const Instruction *state) : Op(state) {} +}; + /// The "for" instruction represents an affine loop nest, defining an SSA value /// for its induction variable. The induction variable is represented as a /// BlockArgument to the entry block of the body. The body and induction @@ -301,6 +350,18 @@ private: explicit AffineIfOp(const Instruction *state) : Op(state) {} }; +/// Returns true if the given Value can be used as a dimension id. +bool isValidDim(const Value *value); + +/// Returns true if the given Value can be used as a symbol. +bool isValidSymbol(const Value *value); + +/// Modifies both `map` and `operands` in-place so as to: +/// 1. drop duplicate operands +/// 2. drop unused dims and symbols from map +void canonicalizeMapAndOperands(AffineMap *map, + llvm::SmallVectorImpl<Value *> *operands); + } // end namespace mlir #endif diff --git a/mlir/include/mlir/IR/BuiltinOps.h b/mlir/include/mlir/IR/BuiltinOps.h index 8be2dbe7b79..02e3a39c09b 100644 --- a/mlir/include/mlir/IR/BuiltinOps.h +++ b/mlir/include/mlir/IR/BuiltinOps.h @@ -36,55 +36,6 @@ public: BuiltinDialect(MLIRContext *context); }; -/// The "affine_apply" operation applies an affine map to a list of operands, -/// yielding a single result. The operand list must be the same size as the -/// number of arguments to the affine mapping. All operands and the result are -/// of type 'Index'. This operation requires a single affine map attribute named -/// "map". For example: -/// -/// %y = "affine_apply" (%x) { map: (d0) -> (d0 + 1) } : -/// (index) -> (index) -/// -/// equivalently: -/// -/// #map42 = (d0)->(d0+1) -/// %y = affine_apply #map42(%x) -/// -class AffineApplyOp : public Op<AffineApplyOp, OpTrait::VariadicOperands, - OpTrait::OneResult, OpTrait::HasNoSideEffect> { -public: - /// Builds an affine apply op with the specified map and operands. - static void build(Builder *builder, OperationState *result, AffineMap map, - ArrayRef<Value *> operands); - - /// Returns the affine map to be applied by this operation. - AffineMap getAffineMap() const { - return getAttrOfType<AffineMapAttr>("map").getValue(); - } - - /// Returns true if the result of this operation can be used as dimension id. - bool isValidDim() const; - - /// Returns true if the result of this operation is a symbol. - bool isValidSymbol() const; - - static StringRef getOperationName() { return "affine_apply"; } - - // Hooks to customize behavior of this op. - static bool parse(OpAsmParser *parser, OperationState *result); - void print(OpAsmPrinter *p) const; - bool verify() const; - Attribute constantFold(ArrayRef<Attribute> operands, - MLIRContext *context) const; - - static void getCanonicalizationPatterns(OwningRewritePatternList &results, - MLIRContext *context); - -private: - friend class Instruction; - explicit AffineApplyOp(const Instruction *state) : Op(state) {} -}; - /// The "br" operation represents a branch instruction in a CFG function. /// The operation takes variable number of operands and produces no results. /// The operand number and types for each successor must match the @@ -397,12 +348,6 @@ bool parseDimAndSymbolList(OpAsmParser *parser, SmallVector<Value *, 4> &operands, unsigned &numDims); -/// Modifies both `map` and `operands` in-place so as to: -/// 1. drop duplicate operands -/// 2. drop unused dims and symbols from map -void canonicalizeMapAndOperands(AffineMap *map, - llvm::SmallVectorImpl<Value *> *operands); - } // end namespace mlir #endif diff --git a/mlir/include/mlir/IR/Value.h b/mlir/include/mlir/IR/Value.h index 842b432f0f5..acc39257e21 100644 --- a/mlir/include/mlir/IR/Value.h +++ b/mlir/include/mlir/IR/Value.h @@ -59,15 +59,6 @@ public: IRObjectWithUseList::replaceAllUsesWith(newValue); } - /// TODO: move isValidDim/isValidSymbol to a utility library specific to the - /// polyhedral operations. - - /// Returns true if the given Value can be used as a dimension id. - bool isValidDim() const; - - /// Returns true if the given Value can be used as a symbol. - bool isValidSymbol() const; - /// Return the function that this Value is defined in. Function *getFunction(); diff --git a/mlir/include/mlir/Transforms/Utils.h b/mlir/include/mlir/Transforms/Utils.h index 581c668a154..b4f1b83ba36 100644 --- a/mlir/include/mlir/Transforms/Utils.h +++ b/mlir/include/mlir/Transforms/Utils.h @@ -32,6 +32,7 @@ namespace mlir { +class AffineApplyOp; class AffineForOp; class FuncBuilder; class Location; diff --git a/mlir/lib/AffineOps/AffineOps.cpp b/mlir/lib/AffineOps/AffineOps.cpp index 2ef96aa3d14..39345d7fc7a 100644 --- a/mlir/lib/AffineOps/AffineOps.cpp +++ b/mlir/lib/AffineOps/AffineOps.cpp @@ -22,6 +22,8 @@ #include "mlir/IR/InstVisitor.h" #include "mlir/IR/IntegerSet.h" #include "mlir/IR/OpImplementation.h" +#include "mlir/IR/PatternMatch.h" +#include "llvm/ADT/SmallBitVector.h" using namespace mlir; //===----------------------------------------------------------------------===// @@ -30,7 +32,241 @@ using namespace mlir; AffineOpsDialect::AffineOpsDialect(MLIRContext *context) : Dialect(/*namePrefix=*/"", context) { - addOperations<AffineForOp, AffineIfOp>(); + addOperations<AffineApplyOp, AffineForOp, AffineIfOp>(); +} + +// 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. +bool mlir::isValidDim(const Value *value) { + if (auto *inst = value->getDefiningInst()) { + // Top level instruction or constant operation is ok. + if (inst->getParentInst() == nullptr || inst->isa<ConstantOp>()) + return true; + // Affine apply operation is ok if all of its operands are ok. + if (auto op = inst->dyn_cast<AffineApplyOp>()) + return op->isValidDim(); + return false; + } + // This value is a block argument. + return true; +} + +// 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. +bool mlir::isValidSymbol(const Value *value) { + if (auto *inst = value->getDefiningInst()) { + // Top level instruction or constant operation is ok. + if (inst->getParentInst() == nullptr || inst->isa<ConstantOp>()) + return true; + // Affine apply operation is ok if all of its operands are ok. + if (auto op = inst->dyn_cast<AffineApplyOp>()) + return op->isValidSymbol(); + return false; + } + // Otherwise, the only valid symbol is a non induction variable block + // argument. + return !isForInductionVar(value); +} + +//===----------------------------------------------------------------------===// +// AffineApplyOp +//===----------------------------------------------------------------------===// + +void AffineApplyOp::build(Builder *builder, OperationState *result, + AffineMap map, ArrayRef<Value *> operands) { + result->addOperands(operands); + result->types.append(map.getNumResults(), builder->getIndexType()); + result->addAttribute("map", builder->getAffineMapAttr(map)); +} + +bool AffineApplyOp::parse(OpAsmParser *parser, OperationState *result) { + auto &builder = parser->getBuilder(); + auto affineIntTy = builder.getIndexType(); + + AffineMapAttr mapAttr; + unsigned numDims; + if (parser->parseAttribute(mapAttr, "map", result->attributes) || + parseDimAndSymbolList(parser, result->operands, numDims) || + parser->parseOptionalAttributeDict(result->attributes)) + return true; + auto map = mapAttr.getValue(); + + if (map.getNumDims() != numDims || + numDims + map.getNumSymbols() != result->operands.size()) { + return parser->emitError(parser->getNameLoc(), + "dimension or symbol index mismatch"); + } + + result->types.append(map.getNumResults(), affineIntTy); + return false; +} + +void AffineApplyOp::print(OpAsmPrinter *p) const { + auto map = getAffineMap(); + *p << "affine_apply " << map; + printDimAndSymbolList(operand_begin(), operand_end(), map.getNumDims(), p); + p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/"map"); +} + +bool AffineApplyOp::verify() const { + // Check that affine map attribute was specified. + auto affineMapAttr = getAttrOfType<AffineMapAttr>("map"); + if (!affineMapAttr) + return emitOpError("requires an affine map"); + + // Check input and output dimensions match. + auto map = affineMapAttr.getValue(); + + // Verify that operand count matches affine map dimension and symbol count. + if (getNumOperands() != map.getNumDims() + map.getNumSymbols()) + return emitOpError( + "operand count and affine map dimension and symbol count must match"); + + // Verify that result count matches affine map result count. + if (map.getNumResults() != 1) + return emitOpError("mapping must produce one value"); + + return false; +} + +// The result of the affine apply operation can be used as a dimension id if it +// is a CFG value or if it is an Value, and all the operands are valid +// dimension ids. +bool AffineApplyOp::isValidDim() const { + return llvm::all_of(getOperands(), + [](const Value *op) { return mlir::isValidDim(op); }); +} + +// The result of the affine apply operation can be used as a symbol if it is +// a CFG value or if it is an Value, and all the operands are symbols. +bool AffineApplyOp::isValidSymbol() const { + return llvm::all_of(getOperands(), + [](const Value *op) { return mlir::isValidSymbol(op); }); +} + +Attribute AffineApplyOp::constantFold(ArrayRef<Attribute> operands, + MLIRContext *context) const { + auto map = getAffineMap(); + SmallVector<Attribute, 1> result; + if (map.constantFold(operands, result)) + return Attribute(); + return result[0]; +} + +namespace { +/// SimplifyAffineApply operations. +/// +struct SimplifyAffineApply : public RewritePattern { + SimplifyAffineApply(MLIRContext *context) + : RewritePattern(AffineApplyOp::getOperationName(), 1, context) {} + + PatternMatchResult match(Instruction *op) const override; + void rewrite(Instruction *op, std::unique_ptr<PatternState> state, + PatternRewriter &rewriter) const override; +}; +} // end anonymous namespace. + +namespace { +/// FIXME: this is massive overkill for simple obviously always matching +/// canonicalizations. Fix the pattern rewriter to make this easy. +struct SimplifyAffineApplyState : public PatternState { + AffineMap map; + SmallVector<Value *, 8> operands; + + SimplifyAffineApplyState(AffineMap map, + const SmallVector<Value *, 8> &operands) + : map(map), operands(operands) {} +}; + +} // end anonymous namespace. + +void mlir::canonicalizeMapAndOperands( + AffineMap *map, llvm::SmallVectorImpl<Value *> *operands) { + if (!map || operands->empty()) + return; + + assert(map->getNumInputs() == operands->size() && + "map inputs must match number of operands"); + + // Check to see what dims are used. + llvm::SmallBitVector usedDims(map->getNumDims()); + llvm::SmallBitVector usedSyms(map->getNumSymbols()); + map->walkExprs([&](AffineExpr expr) { + if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) + usedDims[dimExpr.getPosition()] = true; + else if (auto symExpr = expr.dyn_cast<AffineSymbolExpr>()) + usedSyms[symExpr.getPosition()] = true; + }); + + auto *context = map->getContext(); + + SmallVector<Value *, 8> resultOperands; + resultOperands.reserve(operands->size()); + + llvm::SmallDenseMap<Value *, AffineExpr, 8> seenDims; + SmallVector<AffineExpr, 8> dimRemapping(map->getNumDims()); + unsigned nextDim = 0; + for (unsigned i = 0, e = map->getNumDims(); i != e; ++i) { + if (usedDims[i]) { + auto it = seenDims.find((*operands)[i]); + if (it == seenDims.end()) { + dimRemapping[i] = getAffineDimExpr(nextDim++, context); + resultOperands.push_back((*operands)[i]); + seenDims.insert(std::make_pair((*operands)[i], dimRemapping[i])); + } else { + dimRemapping[i] = it->second; + } + } + } + llvm::SmallDenseMap<Value *, AffineExpr, 8> seenSymbols; + SmallVector<AffineExpr, 8> symRemapping(map->getNumSymbols()); + unsigned nextSym = 0; + for (unsigned i = 0, e = map->getNumSymbols(); i != e; ++i) { + if (usedSyms[i]) { + auto it = seenSymbols.find((*operands)[i + map->getNumDims()]); + if (it == seenSymbols.end()) { + symRemapping[i] = getAffineSymbolExpr(nextSym++, context); + resultOperands.push_back((*operands)[i + map->getNumDims()]); + seenSymbols.insert(std::make_pair((*operands)[i + map->getNumDims()], + symRemapping[i])); + } else { + symRemapping[i] = it->second; + } + } + } + *map = + map->replaceDimsAndSymbols(dimRemapping, symRemapping, nextDim, nextSym); + *operands = resultOperands; +} + +PatternMatchResult SimplifyAffineApply::match(Instruction *op) const { + auto apply = op->cast<AffineApplyOp>(); + auto map = apply->getAffineMap(); + + AffineMap oldMap = map; + SmallVector<Value *, 8> resultOperands(apply->getOperands().begin(), + apply->getOperands().end()); + canonicalizeMapAndOperands(&map, &resultOperands); + if (map != oldMap) + return matchSuccess( + std::make_unique<SimplifyAffineApplyState>(map, resultOperands)); + + return matchFailure(); +} + +void SimplifyAffineApply::rewrite(Instruction *op, + std::unique_ptr<PatternState> state, + PatternRewriter &rewriter) const { + auto *applyState = static_cast<SimplifyAffineApplyState *>(state.get()); + rewriter.replaceOpWithNewOp<AffineApplyOp>(op, applyState->map, + applyState->operands); +} + +void AffineApplyOp::getCanonicalizationPatterns( + OwningRewritePatternList &results, MLIRContext *context) { + results.push_back(std::make_unique<SimplifyAffineApply>(context)); } //===----------------------------------------------------------------------===// @@ -493,9 +729,9 @@ bool AffineIfOp::verify() const { IntegerSet condition = conditionAttr.getValue(); for (unsigned i = 0, e = getNumOperands(); i != e; ++i) { const Value *operand = getOperand(i); - if (i < condition.getNumDims() && !operand->isValidDim()) + if (i < condition.getNumDims() && !isValidDim(operand)) return emitOpError("operand cannot be used as a dimension id"); - if (i >= condition.getNumDims() && !operand->isValidSymbol()) + if (i >= condition.getNumDims() && !isValidSymbol(operand)) return emitOpError("operand cannot be used as a symbol"); } diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp index 0a2b283738c..2323cb3ef71 100644 --- a/mlir/lib/Analysis/AffineAnalysis.cpp +++ b/mlir/lib/Analysis/AffineAnalysis.cpp @@ -681,7 +681,7 @@ static void buildDimAndSymbolPositionMaps( for (unsigned i = 0, e = values.size(); i < e; ++i) { auto *value = values[i]; if (!isForInductionVar(values[i])) { - assert(values[i]->isValidSymbol() && + assert(isValidSymbol(values[i]) && "access operand has to be either a loop IV or a symbol"); valuePosMap->addSymbolValue(value); } else if (isSrc) { @@ -743,7 +743,7 @@ void initDependenceConstraints(const FlatAffineConstraints &srcDomain, auto setSymbolIds = [&](ArrayRef<Value *> values) { for (auto *value : values) { if (!isForInductionVar(value)) { - assert(value->isValidSymbol() && "expected symbol"); + assert(isValidSymbol(value) && "expected symbol"); dependenceConstraints->setIdValue(valuePosMap.getSymPos(value), value); } } @@ -913,7 +913,7 @@ addMemRefAccessConstraints(const AffineValueMap &srcAccessMap, if (isForInductionVar(operands[i])) continue; auto *symbol = operands[i]; - assert(symbol->isValidSymbol()); + assert(isValidSymbol(symbol)); // Check if the symbol is a constant. if (auto *opInst = symbol->getDefiningInst()) { if (auto constOp = opInst->dyn_cast<ConstantIndexOp>()) { diff --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp index 468e79b8545..b764d69c298 100644 --- a/mlir/lib/Analysis/AffineStructures.cpp +++ b/mlir/lib/Analysis/AffineStructures.cpp @@ -1268,7 +1268,7 @@ bool FlatAffineConstraints::addAffineForOpDomain( for (const auto &operand : operands) { unsigned loc; if (!findId(*operand, &loc)) { - if (operand->isValidSymbol()) { + if (isValidSymbol(operand)) { addSymbolId(getNumSymbolIds(), const_cast<Value *>(operand)); loc = getNumDimIds() + getNumSymbolIds() - 1; // Check if the symbol is a constant. diff --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp index d27715333c8..9f5fd65e774 100644 --- a/mlir/lib/Analysis/Utils.cpp +++ b/mlir/lib/Analysis/Utils.cpp @@ -168,7 +168,7 @@ bool MemRefRegion::compute(Instruction *inst, unsigned loopDepth) { } else { // Has to be a valid symbol. auto *symbol = accessValueMap.getOperand(i); - assert(symbol->isValidSymbol()); + assert(isValidSymbol(symbol)); // Check if the symbol is a constant. if (auto *inst = symbol->getDefiningInst()) { if (auto constOp = inst->dyn_cast<ConstantIndexOp>()) { diff --git a/mlir/lib/IR/BuiltinOps.cpp b/mlir/lib/IR/BuiltinOps.cpp index 1d225cd7a8a..bfa7701e500 100644 --- a/mlir/lib/IR/BuiltinOps.cpp +++ b/mlir/lib/IR/BuiltinOps.cpp @@ -27,7 +27,6 @@ #include "mlir/Support/MathExtras.h" #include "mlir/Support/STLExtras.h" #include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/SmallBitVector.h" #include "llvm/Support/raw_ostream.h" using namespace mlir; @@ -38,7 +37,7 @@ using namespace mlir; BuiltinDialect::BuiltinDialect(MLIRContext *context) : Dialect(/*namePrefix=*/"", context) { - addOperations<AffineApplyOp, BranchOp, CondBranchOp, ConstantOp, ReturnOp>(); + addOperations<BranchOp, CondBranchOp, ConstantOp, ReturnOp>(); addTypes<FunctionType, IndexType, UnknownType, FloatType, IntegerType, VectorType, RankedTensorType, UnrankedTensorType, MemRefType>(); } @@ -79,211 +78,6 @@ bool mlir::parseDimAndSymbolList(OpAsmParser *parser, } //===----------------------------------------------------------------------===// -// AffineApplyOp -//===----------------------------------------------------------------------===// - -void AffineApplyOp::build(Builder *builder, OperationState *result, - AffineMap map, ArrayRef<Value *> operands) { - result->addOperands(operands); - result->types.append(map.getNumResults(), builder->getIndexType()); - result->addAttribute("map", builder->getAffineMapAttr(map)); -} - -bool AffineApplyOp::parse(OpAsmParser *parser, OperationState *result) { - auto &builder = parser->getBuilder(); - auto affineIntTy = builder.getIndexType(); - - AffineMapAttr mapAttr; - unsigned numDims; - if (parser->parseAttribute(mapAttr, "map", result->attributes) || - parseDimAndSymbolList(parser, result->operands, numDims) || - parser->parseOptionalAttributeDict(result->attributes)) - return true; - auto map = mapAttr.getValue(); - - if (map.getNumDims() != numDims || - numDims + map.getNumSymbols() != result->operands.size()) { - return parser->emitError(parser->getNameLoc(), - "dimension or symbol index mismatch"); - } - - result->types.append(map.getNumResults(), affineIntTy); - return false; -} - -void AffineApplyOp::print(OpAsmPrinter *p) const { - auto map = getAffineMap(); - *p << "affine_apply " << map; - printDimAndSymbolList(operand_begin(), operand_end(), map.getNumDims(), p); - p->printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/"map"); -} - -bool AffineApplyOp::verify() const { - // Check that affine map attribute was specified. - auto affineMapAttr = getAttrOfType<AffineMapAttr>("map"); - if (!affineMapAttr) - return emitOpError("requires an affine map"); - - // Check input and output dimensions match. - auto map = affineMapAttr.getValue(); - - // Verify that operand count matches affine map dimension and symbol count. - if (getNumOperands() != map.getNumDims() + map.getNumSymbols()) - return emitOpError( - "operand count and affine map dimension and symbol count must match"); - - // Verify that result count matches affine map result count. - if (map.getNumResults() != 1) - return emitOpError("mapping must produce one value"); - - return false; -} - -// The result of the affine apply operation can be used as a dimension id if it -// is a CFG value or if it is an Value, and all the operands are valid -// dimension ids. -bool AffineApplyOp::isValidDim() const { - for (auto *op : getOperands()) { - if (!op->isValidDim()) - return false; - } - return true; -} - -// The result of the affine apply operation can be used as a symbol if it is -// a CFG value or if it is an Value, and all the operands are symbols. -bool AffineApplyOp::isValidSymbol() const { - for (auto *op : getOperands()) { - if (!op->isValidSymbol()) - return false; - } - return true; -} - -Attribute AffineApplyOp::constantFold(ArrayRef<Attribute> operands, - MLIRContext *context) const { - auto map = getAffineMap(); - SmallVector<Attribute, 1> result; - if (map.constantFold(operands, result)) - return Attribute(); - return result[0]; -} - -namespace { -/// SimplifyAffineApply operations. -/// -struct SimplifyAffineApply : public RewritePattern { - SimplifyAffineApply(MLIRContext *context) - : RewritePattern(AffineApplyOp::getOperationName(), 1, context) {} - - PatternMatchResult match(Instruction *op) const override; - void rewrite(Instruction *op, std::unique_ptr<PatternState> state, - PatternRewriter &rewriter) const override; -}; -} // end anonymous namespace. - -namespace { -/// FIXME: this is massive overkill for simple obviously always matching -/// canonicalizations. Fix the pattern rewriter to make this easy. -struct SimplifyAffineApplyState : public PatternState { - AffineMap map; - SmallVector<Value *, 8> operands; - - SimplifyAffineApplyState(AffineMap map, - const SmallVector<Value *, 8> &operands) - : map(map), operands(operands) {} -}; - -} // end anonymous namespace. - -void mlir::canonicalizeMapAndOperands( - AffineMap *map, llvm::SmallVectorImpl<Value *> *operands) { - if (!map || operands->empty()) - return; - - assert(map->getNumInputs() == operands->size() && - "map inputs must match number of operands"); - - // Check to see what dims are used. - llvm::SmallBitVector usedDims(map->getNumDims()); - llvm::SmallBitVector usedSyms(map->getNumSymbols()); - map->walkExprs([&](AffineExpr expr) { - if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) - usedDims[dimExpr.getPosition()] = true; - else if (auto symExpr = expr.dyn_cast<AffineSymbolExpr>()) - usedSyms[symExpr.getPosition()] = true; - }); - - auto *context = map->getContext(); - - SmallVector<Value *, 8> resultOperands; - resultOperands.reserve(operands->size()); - - llvm::SmallDenseMap<Value *, AffineExpr, 8> seenDims; - SmallVector<AffineExpr, 8> dimRemapping(map->getNumDims()); - unsigned nextDim = 0; - for (unsigned i = 0, e = map->getNumDims(); i != e; ++i) { - if (usedDims[i]) { - auto it = seenDims.find((*operands)[i]); - if (it == seenDims.end()) { - dimRemapping[i] = getAffineDimExpr(nextDim++, context); - resultOperands.push_back((*operands)[i]); - seenDims.insert(std::make_pair((*operands)[i], dimRemapping[i])); - } else { - dimRemapping[i] = it->second; - } - } - } - llvm::SmallDenseMap<Value *, AffineExpr, 8> seenSymbols; - SmallVector<AffineExpr, 8> symRemapping(map->getNumSymbols()); - unsigned nextSym = 0; - for (unsigned i = 0, e = map->getNumSymbols(); i != e; ++i) { - if (usedSyms[i]) { - auto it = seenSymbols.find((*operands)[i + map->getNumDims()]); - if (it == seenSymbols.end()) { - symRemapping[i] = getAffineSymbolExpr(nextSym++, context); - resultOperands.push_back((*operands)[i + map->getNumDims()]); - seenSymbols.insert(std::make_pair((*operands)[i + map->getNumDims()], - symRemapping[i])); - } else { - symRemapping[i] = it->second; - } - } - } - *map = - map->replaceDimsAndSymbols(dimRemapping, symRemapping, nextDim, nextSym); - *operands = resultOperands; -} - -PatternMatchResult SimplifyAffineApply::match(Instruction *op) const { - auto apply = op->cast<AffineApplyOp>(); - auto map = apply->getAffineMap(); - - AffineMap oldMap = map; - SmallVector<Value *, 8> resultOperands(apply->getOperands().begin(), - apply->getOperands().end()); - canonicalizeMapAndOperands(&map, &resultOperands); - if (map != oldMap) - return matchSuccess( - std::make_unique<SimplifyAffineApplyState>(map, resultOperands)); - - return matchFailure(); -} - -void SimplifyAffineApply::rewrite(Instruction *op, - std::unique_ptr<PatternState> state, - PatternRewriter &rewriter) const { - auto *applyState = static_cast<SimplifyAffineApplyState *>(state.get()); - rewriter.replaceOpWithNewOp<AffineApplyOp>(op, applyState->map, - applyState->operands); -} - -void AffineApplyOp::getCanonicalizationPatterns( - OwningRewritePatternList &results, MLIRContext *context) { - results.push_back(std::make_unique<SimplifyAffineApply>(context)); -} - -//===----------------------------------------------------------------------===// // BranchOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/IR/Instruction.cpp b/mlir/lib/IR/Instruction.cpp index b1653cb9e09..6720969ac0f 100644 --- a/mlir/lib/IR/Instruction.cpp +++ b/mlir/lib/IR/Instruction.cpp @@ -300,42 +300,6 @@ Function *Instruction::getFunction() const { return block ? block->getFunction() : nullptr; } -// 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. -bool Value::isValidDim() const { - if (auto *inst = getDefiningInst()) { - // Top level instruction or constant operation is ok. - if (inst->getParentInst() == nullptr || inst->isa<ConstantOp>()) - return true; - // Affine apply operation is ok if all of its operands are ok. - if (auto op = inst->dyn_cast<AffineApplyOp>()) - return op->isValidDim(); - return false; - } - // This value is either a function argument or an induction variable. Both - // are ok. - return true; -} - -// 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. -bool Value::isValidSymbol() const { - if (auto *inst = getDefiningInst()) { - // Top level instruction or constant operation is ok. - if (inst->getParentInst() == nullptr || inst->isa<ConstantOp>()) - return true; - // Affine apply operation is ok if all of its operands are ok. - if (auto op = inst->dyn_cast<AffineApplyOp>()) - return op->isValidSymbol(); - return false; - } - // Otherwise, the only valid symbol is a function argument. - auto *arg = dyn_cast<BlockArgument>(this); - return arg && arg->isFunctionArgument(); -} - /// Emit a note about this instruction, reporting up to any diagnostic /// handlers that may be listening. void Instruction::emitNote(const Twine &message) const { diff --git a/mlir/lib/Transforms/ComposeAffineMaps.cpp b/mlir/lib/Transforms/ComposeAffineMaps.cpp index 4a6430dc9be..289b00d3b51 100644 --- a/mlir/lib/Transforms/ComposeAffineMaps.cpp +++ b/mlir/lib/Transforms/ComposeAffineMaps.cpp @@ -21,6 +21,7 @@ // //===----------------------------------------------------------------------===// +#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/Attributes.h" diff --git a/mlir/lib/Transforms/Utils/Utils.cpp b/mlir/lib/Transforms/Utils/Utils.cpp index 524e8d542f5..1e9ad25fbec 100644 --- a/mlir/lib/Transforms/Utils/Utils.cpp +++ b/mlir/lib/Transforms/Utils/Utils.cpp @@ -124,7 +124,7 @@ bool mlir::replaceAllMemRefUsesWith(const Value *oldMemRef, Value *newMemRef, for (auto *extraIndex : extraIndices) { assert(extraIndex->getDefiningInst()->getNumResults() == 1 && "single result op's expected to generate these indices"); - assert((extraIndex->isValidDim() || extraIndex->isValidSymbol()) && + assert((isValidDim(extraIndex) || isValidSymbol(extraIndex)) && "invalid memory op index"); state.operands.push_back(extraIndex); } diff --git a/mlir/lib/Transforms/Vectorization/VectorizerTestPass.cpp b/mlir/lib/Transforms/Vectorization/VectorizerTestPass.cpp index 7d51637a6e1..3c706133a02 100644 --- a/mlir/lib/Transforms/Vectorization/VectorizerTestPass.cpp +++ b/mlir/lib/Transforms/Vectorization/VectorizerTestPass.cpp @@ -19,6 +19,7 @@ // //===----------------------------------------------------------------------===// +#include "mlir/AffineOps/AffineOps.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/NestedMatcher.h" #include "mlir/Analysis/SliceAnalysis.h" |

