From 00aac7015935a2da7a48c9ffc19e8057c7006411 Mon Sep 17 00:00:00 2001 From: Nicolas Vasilache Date: Mon, 7 Jan 2019 20:14:30 -0800 Subject: [PATCH] Move makeNormalizedAffineApply This CL is the 3rd on the path to simplifying AffineMap composition. This CL just moves `makeNormalizedAffineApply` from VectorAnalysis to AffineAnalysis where it more naturally belongs. PiperOrigin-RevId: 228277182 --- mlir/include/mlir/Analysis/AffineAnalysis.h | 29 ++- mlir/include/mlir/Analysis/VectorAnalysis.h | 19 -- mlir/lib/Analysis/AffineAnalysis.cpp | 204 +++++++++++++++++++ mlir/lib/Analysis/VectorAnalysis.cpp | 207 -------------------- 4 files changed, 229 insertions(+), 230 deletions(-) diff --git a/mlir/include/mlir/Analysis/AffineAnalysis.h b/mlir/include/mlir/Analysis/AffineAnalysis.h index ba6cb58111b4..0f165c967c0b 100644 --- a/mlir/include/mlir/Analysis/AffineAnalysis.h +++ b/mlir/include/mlir/Analysis/AffineAnalysis.h @@ -30,15 +30,19 @@ namespace mlir { +class AffineApplyOp; class AffineExpr; class AffineMap; class AffineValueMap; -class ForInst; -class MLIRContext; class FlatAffineConstraints; -class IntegerSet; -class OperationInst; +class ForInst; +class FuncBuilder; class Instruction; +class IntegerSet; +class Location; +class MLIRContext; +class OperationInst; +template class OpPointer; class Value; /// Simplify an affine expression by flattening and some amount of @@ -52,6 +56,23 @@ AffineExpr simplifyAffineExpr(AffineExpr expr, unsigned numDims, /// and sizes. AffineMap simplifyAffineMap(AffineMap map); +/// Creates an AffineApplyOp that is normalized for super-vectorization. That is +/// an AffineApplyOp with a single result and an unbounded AffineMap. The +/// operands of the AffineApplyOp are either dims, symbols or constants but can +/// never be obtained from other AffineApplyOp. +/// This is achieved by performing a composition at the single-result AffineMap +/// level. +/// +/// Prerequisite: +/// 1. `map` is a single result, unbounded, AffineMap; +/// 2. `operands` can involve at most a length-1 chain of AffineApplyOp. The +/// affine map for each of these AffineApplyOp is itself single result and +/// unbounded. Essentially, all ancestor AffineApplyOp must have been +/// constructed as single-result, unbounded, AffineMaps. +OpPointer +makeNormalizedAffineApply(FuncBuilder *b, Location loc, AffineMap map, + llvm::ArrayRef operands); + /// Returns the sequence of AffineApplyOp OperationInsts operation in /// 'affineApplyOps', which are reachable via a search starting from 'operands', /// and ending at operands which are not defined by AffineApplyOps. diff --git a/mlir/include/mlir/Analysis/VectorAnalysis.h b/mlir/include/mlir/Analysis/VectorAnalysis.h index 001bb4f99348..e34e3433f2fe 100644 --- a/mlir/include/mlir/Analysis/VectorAnalysis.h +++ b/mlir/include/mlir/Analysis/VectorAnalysis.h @@ -32,8 +32,6 @@ class Instruction; class Location; class MemRefType; class OperationInst; -template class OperandIterator; -template class OpPointer; class Value; class VectorType; @@ -128,23 +126,6 @@ AffineMap makePermutationMap(OperationInst *opInst, const llvm::DenseMap &loopToVectorDim); -/// Creates an AffineApplyOp that is normalized for super-vectorization. That is -/// an AffineApplyOp with a single result and an unbounded AffineMap. The -/// operands of the AffineApplyOp are either dims, symbols or constants but can -/// never be obtained from other AffineApplyOp. -/// This is achieved by performing a composition at the single-result AffineMap -/// level. -/// -/// Prerequisite: -/// 1. `map` is a single result, unbounded, AffineMap; -/// 2. `operands` can involve at most a length-1 chain of AffineApplyOp. The -/// affine map for each of these AffineApplyOp is itself single result and -/// unbounded. Essentially, all ancestor AffineApplyOp must have been -/// constructed as single-result, unbounded, AffineMaps. -OpPointer makeNormalizedAffineApply(FuncBuilder *b, Location loc, - AffineMap map, - ArrayRef operands); - namespace matcher { /// Matches vector_transfer_read, vector_transfer_write and ops that return a diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp index 7d31bdae387e..28551fc6f4a0 100644 --- a/mlir/lib/Analysis/AffineAnalysis.cpp +++ b/mlir/lib/Analysis/AffineAnalysis.cpp @@ -24,6 +24,7 @@ #include "mlir/Analysis/AffineStructures.h" #include "mlir/Analysis/Utils.h" #include "mlir/IR/AffineExprVisitor.h" +#include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Instructions.h" #include "mlir/StandardOps/StandardOps.h" @@ -36,6 +37,8 @@ using namespace mlir; +using llvm::dbgs; + /// Constructs an affine expression from a flat ArrayRef. If there are local /// identifiers (neither dimensional nor symbolic) that appear in the sum of /// products expression, 'localExprs' is expected to have the AffineExpr @@ -1352,3 +1355,204 @@ bool mlir::checkMemrefAccessDependence( LLVM_DEBUG(dependenceConstraints->dump()); return true; } + +namespace { + +/// A `SingleResultAffineNormalizer` is a helper class that is not visible to +/// the user and supports renumbering operands of single-result AffineApplyOp. +/// This operates on the assumption that only single-result unbounded AffineMap +/// are used for all operands. +/// This acts as a reindexing map of Value* to positional dims or symbols and +/// allows simplifications such as: +/// +/// ```mlir +/// %1 = affine_apply (d0, d1) -> (d0 - d1) (%0, %0) +/// ``` +/// +/// into: +/// +/// ```mlir +/// %1 = affine_apply () -> (0) +/// ``` +struct SingleResultAffineNormalizer { + SingleResultAffineNormalizer(AffineMap map, ArrayRef operands); + + /// Returns the single result, unbounded, AffineMap resulting from + /// normalization. + AffineMap getAffineMap() { + return AffineMap::get(reorderedDims.size(), reorderedSymbols.size(), {expr}, + {}); + } + + SmallVector getOperands() { + SmallVector res(reorderedDims); + res.append(reorderedSymbols.begin(), reorderedSymbols.end()); + return res; + } + +private: + /// Helper function to insert `v` into the coordinate system of the current + /// SingleResultAffineNormalizer (i.e. in the proper `xxxValueToPosition` and + /// the proper `reorderedXXX`). + /// Returns the AffineDimExpr or AffineSymbolExpr with the correponding + /// renumbered position. + template DimOrSymbol renumberOneIndex(Value *v); + + /// Given an `other` normalizer, this rewrites `other.expr` in the coordinate + /// system of the current SingleResultAffineNormalizer. + /// Returns the rewritten AffineExpr. + AffineExpr renumber(const SingleResultAffineNormalizer &other); + + /// Given an `app` with single result and unbounded AffineMap, this rewrites + /// the app's map single result AffineExpr in the coordinate system of the + /// current SingleResultAffineNormalizer. + /// Returns the rewritten AffineExpr. + AffineExpr renumber(AffineApplyOp *app); + + /// Maps of Value* to position in the `expr`. + DenseMap dimValueToPosition; + DenseMap symValueToPosition; + + /// Ordered dims and symbols matching positional dims and symbols in `expr`. + SmallVector reorderedDims; + SmallVector reorderedSymbols; + + AffineExpr expr; +}; + +} // namespace + +template +static DimOrSymbol make(unsigned position, MLIRContext *context); + +template <> AffineDimExpr make(unsigned position, MLIRContext *context) { + return getAffineDimExpr(position, context).cast(); +} + +template <> AffineSymbolExpr make(unsigned position, MLIRContext *context) { + return getAffineSymbolExpr(position, context).cast(); +} + +template +DimOrSymbol SingleResultAffineNormalizer::renumberOneIndex(Value *v) { + static_assert(std::is_same::value || + std::is_same::value, + "renumber(...) or renumber(...) " + "required"); + DenseMap &pos = + std::is_same::value ? symValueToPosition + : dimValueToPosition; + DenseMap::iterator iterPos; + bool inserted = false; + std::tie(iterPos, inserted) = pos.insert(std::make_pair(v, pos.size())); + if (inserted) { + std::is_same::value + ? reorderedDims.push_back(v) + : reorderedSymbols.push_back(v); + } + return make(iterPos->second, v->getFunction()->getContext()); +} + +AffineExpr SingleResultAffineNormalizer::renumber( + const SingleResultAffineNormalizer &other) { + SmallVector dimRemapping, symRemapping; + for (auto *v : other.reorderedDims) { + auto kvp = other.dimValueToPosition.find(v); + if (dimRemapping.size() <= kvp->second) + dimRemapping.resize(kvp->second + 1); + dimRemapping[kvp->second] = renumberOneIndex(kvp->first); + } + for (auto *v : other.reorderedSymbols) { + auto kvp = other.symValueToPosition.find(v); + if (symRemapping.size() <= kvp->second) + symRemapping.resize(kvp->second + 1); + symRemapping[kvp->second] = renumberOneIndex(kvp->first); + } + return other.expr.replaceDimsAndSymbols(dimRemapping, symRemapping); +} + +AffineExpr SingleResultAffineNormalizer::renumber(AffineApplyOp *app) { + // Sanity check, single result AffineApplyOp if one wants to use this. + assert(app->getNumResults() == 1 && "Not a single result AffineApplyOp"); + assert(app->getAffineMap().getRangeSizes().empty() && + "Non-empty range sizes"); + + // Create the SingleResultAffineNormalizer for the operands of this + // AffineApplyOp and combine it with the current SingleResultAffineNormalizer. + using ValueTy = decltype(*(app->getOperands().begin())); + SingleResultAffineNormalizer normalizer( + app->getAffineMap(), + functional::map([](ValueTy v) { return static_cast(v); }, + app->getOperands())); + + // We know this is a single result AffineMap, we need to append a + // renumbered AffineExpr. + return renumber(normalizer); +} + +SingleResultAffineNormalizer::SingleResultAffineNormalizer( + AffineMap map, ArrayRef operands) { + assert(map.getNumResults() == 1 && "Single-result map expected"); + assert(map.getRangeSizes().empty() && "Unbounded map expected"); + assert(map.getNumInputs() == operands.size() && + "number of operands does not match the number of map inputs"); + + if (operands.empty()) { + return; + } + + auto *context = operands[0]->getFunction()->getContext(); + SmallVector exprs; + for (auto en : llvm::enumerate(operands)) { + auto *t = en.value(); + assert(t->getType().isIndex()); + if (auto inst = t->getDefiningInst()) { + if (auto app = inst->dyn_cast()) { + // Sanity check, AffineApplyOp must always be composed by construction + // and there can only ever be a dependence chain of 1 AffineApply. So we + // can never get a second AffineApplyOp. + // This also guarantees we can build another + // SingleResultAffineNormalizer here that does not recurse a second + // time. + for (auto *pred : app->getOperands()) { + assert(!pred->getDefiningInst() || + !pred->getDefiningInst()->isa() && + "AffineApplyOp chain of length > 1"); + (void)pred; + } + exprs.push_back(renumber(app)); + } else if (auto constant = inst->dyn_cast()) { + // Constants remain constants. + auto affineConstant = inst->cast(); + exprs.push_back( + getAffineConstantExpr(affineConstant->getValue(), context)); + } else { + // DimOp, top of the function symbols are all symbols. + exprs.push_back(renumberOneIndex(t)); + } + } else if (en.index() < map.getNumDims()) { + assert(isa(t) && "ForInst expected for AffineDimExpr"); + exprs.push_back(renumberOneIndex(t)); + } else { + assert(!isa(t) && "unexpectd ForInst for a AffineSymbolExpr"); + exprs.push_back(renumberOneIndex(t)); + } + } + auto exprsMap = AffineMap::get(dimValueToPosition.size(), + symValueToPosition.size(), exprs, {}); + + expr = simplifyAffineExpr(map.getResult(0).compose(exprsMap), + exprsMap.getNumDims(), exprsMap.getNumSymbols()); + + LLVM_DEBUG(map.getResult(0).print(dbgs() << "\nCompose expr: ")); + LLVM_DEBUG(exprsMap.print(dbgs() << "\nWith map: ")); + LLVM_DEBUG(expr.print(dbgs() << "\nResult: ")); +} + +OpPointer +mlir::makeNormalizedAffineApply(FuncBuilder *b, Location loc, AffineMap map, + ArrayRef operands) { + SingleResultAffineNormalizer normalizer(map, operands); + return b->create(loc, normalizer.getAffineMap(), + normalizer.getOperands()); +} diff --git a/mlir/lib/Analysis/VectorAnalysis.cpp b/mlir/lib/Analysis/VectorAnalysis.cpp index cd74fc0c01c4..bc43e2ca5eb2 100644 --- a/mlir/lib/Analysis/VectorAnalysis.cpp +++ b/mlir/lib/Analysis/VectorAnalysis.cpp @@ -18,7 +18,6 @@ #include "mlir/Analysis/VectorAnalysis.h" #include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/LoopAnalysis.h" -#include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Instructions.h" #include "mlir/StandardOps/StandardOps.h" @@ -28,8 +27,6 @@ #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/SetVector.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/raw_ostream.h" /// /// Implements Analysis functions specific to vectors which support @@ -38,9 +35,6 @@ using namespace mlir; -#define DEBUG_TYPE "vector-analysis" - -using llvm::dbgs; using llvm::SetVector; Optional> mlir::shapeRatio(ArrayRef superShape, @@ -243,204 +237,3 @@ bool mlir::matcher::operatesOnSuperVectors(const OperationInst &opInst, return true; } - -namespace { - -/// A `SingleResultAffineNormalizer` is a helper class that is not visible to -/// the user and supports renumbering operands of single-result AffineApplyOp. -/// This operates on the assumption that only single-result unbounded AffineMap -/// are used for all operands. -/// This acts as a reindexing map of Value* to positional dims or symbols and -/// allows simplifications such as: -/// -/// ```mlir -/// %1 = affine_apply (d0, d1) -> (d0 - d1) (%0, %0) -/// ``` -/// -/// into: -/// -/// ```mlir -/// %1 = affine_apply () -> (0) -/// ``` -struct SingleResultAffineNormalizer { - SingleResultAffineNormalizer(AffineMap map, ArrayRef operands); - - /// Returns the single result, unbounded, AffineMap resulting from - /// normalization. - AffineMap getAffineMap() { - return AffineMap::get(reorderedDims.size(), reorderedSymbols.size(), {expr}, - {}); - } - - SmallVector getOperands() { - SmallVector res(reorderedDims); - res.append(reorderedSymbols.begin(), reorderedSymbols.end()); - return res; - } - -private: - /// Helper function to insert `v` into the coordinate system of the current - /// SingleResultAffineNormalizer (i.e. in the proper `xxxValueToPosition` and - /// the proper `reorderedXXX`). - /// Returns the AffineDimExpr or AffineSymbolExpr with the correponding - /// renumbered position. - template DimOrSymbol renumberOneIndex(Value *v); - - /// Given an `other` normalizer, this rewrites `other.expr` in the coordinate - /// system of the current SingleResultAffineNormalizer. - /// Returns the rewritten AffineExpr. - AffineExpr renumber(const SingleResultAffineNormalizer &other); - - /// Given an `app` with single result and unbounded AffineMap, this rewrites - /// the app's map single result AffineExpr in the coordinate system of the - /// current SingleResultAffineNormalizer. - /// Returns the rewritten AffineExpr. - AffineExpr renumber(AffineApplyOp *app); - - /// Maps of Value* to position in the `expr`. - DenseMap dimValueToPosition; - DenseMap symValueToPosition; - - /// Ordered dims and symbols matching positional dims and symbols in `expr`. - SmallVector reorderedDims; - SmallVector reorderedSymbols; - - AffineExpr expr; -}; - -} // namespace - -template -static DimOrSymbol make(unsigned position, MLIRContext *context); - -template <> AffineDimExpr make(unsigned position, MLIRContext *context) { - return getAffineDimExpr(position, context).cast(); -} - -template <> AffineSymbolExpr make(unsigned position, MLIRContext *context) { - return getAffineSymbolExpr(position, context).cast(); -} - -template -DimOrSymbol SingleResultAffineNormalizer::renumberOneIndex(Value *v) { - static_assert(std::is_same::value || - std::is_same::value, - "renumber(...) or renumber(...) " - "required"); - DenseMap &pos = - std::is_same::value ? symValueToPosition - : dimValueToPosition; - DenseMap::iterator iterPos; - bool inserted = false; - std::tie(iterPos, inserted) = pos.insert(std::make_pair(v, pos.size())); - if (inserted) { - std::is_same::value - ? reorderedDims.push_back(v) - : reorderedSymbols.push_back(v); - } - return make(iterPos->second, v->getFunction()->getContext()); -} - -AffineExpr SingleResultAffineNormalizer::renumber( - const SingleResultAffineNormalizer &other) { - SmallVector dimRemapping, symRemapping; - for (auto *v : other.reorderedDims) { - auto kvp = other.dimValueToPosition.find(v); - if (dimRemapping.size() <= kvp->second) - dimRemapping.resize(kvp->second + 1); - dimRemapping[kvp->second] = renumberOneIndex(kvp->first); - } - for (auto *v : other.reorderedSymbols) { - auto kvp = other.symValueToPosition.find(v); - if (symRemapping.size() <= kvp->second) - symRemapping.resize(kvp->second + 1); - symRemapping[kvp->second] = renumberOneIndex(kvp->first); - } - return other.expr.replaceDimsAndSymbols(dimRemapping, symRemapping); -} - -AffineExpr SingleResultAffineNormalizer::renumber(AffineApplyOp *app) { - // Sanity check, single result AffineApplyOp if one wants to use this. - assert(app->getNumResults() == 1 && "Not a single result AffineApplyOp"); - assert(app->getAffineMap().getRangeSizes().empty() && - "Non-empty range sizes"); - - // Create the SingleResultAffineNormalizer for the operands of this - // AffineApplyOp and combine it with the current SingleResultAffineNormalizer. - using ValueTy = decltype(*(app->getOperands().begin())); - SingleResultAffineNormalizer normalizer( - app->getAffineMap(), - functional::map([](ValueTy v) { return static_cast(v); }, - app->getOperands())); - - // We know this is a single result AffineMap, we need to append a - // renumbered AffineExpr. - return renumber(normalizer); -} - -SingleResultAffineNormalizer::SingleResultAffineNormalizer( - AffineMap map, ArrayRef operands) { - assert(map.getNumResults() == 1 && "Single-result map expected"); - assert(map.getRangeSizes().empty() && "Unbounded map expected"); - assert(map.getNumInputs() == operands.size() && - "number of operands does not match the number of map inputs"); - - if (operands.empty()) { - return; - } - - auto *context = operands[0]->getFunction()->getContext(); - SmallVector exprs; - for (auto en : llvm::enumerate(operands)) { - auto *t = en.value(); - assert(t->getType().isIndex()); - if (auto inst = t->getDefiningInst()) { - if (auto app = inst->dyn_cast()) { - // Sanity check, AffineApplyOp must always be composed by construction - // and there can only ever be a dependence chain of 1 AffineApply. So we - // can never get a second AffineApplyOp. - // This also guarantees we can build another - // SingleResultAffineNormalizer here that does not recurse a second - // time. - for (auto *pred : app->getOperands()) { - assert(!pred->getDefiningInst() || - !pred->getDefiningInst()->isa() && - "AffineApplyOp chain of length > 1"); - (void)pred; - } - exprs.push_back(renumber(app)); - } else if (auto constant = inst->dyn_cast()) { - // Constants remain constants. - auto affineConstant = inst->cast(); - exprs.push_back( - getAffineConstantExpr(affineConstant->getValue(), context)); - } else { - // DimOp, top of the function symbols are all symbols. - exprs.push_back(renumberOneIndex(t)); - } - } else if (en.index() < map.getNumDims()) { - assert(isa(t) && "ForInst expected for AffineDimExpr"); - exprs.push_back(renumberOneIndex(t)); - } else { - assert(!isa(t) && "unexpectd ForInst for a AffineSymbolExpr"); - exprs.push_back(renumberOneIndex(t)); - } - } - auto exprsMap = AffineMap::get(dimValueToPosition.size(), - symValueToPosition.size(), exprs, {}); - - expr = simplifyAffineExpr(map.getResult(0).compose(exprsMap), - exprsMap.getNumDims(), exprsMap.getNumSymbols()); - - LLVM_DEBUG(map.getResult(0).print(dbgs() << "\nCompose expr: ")); - LLVM_DEBUG(exprsMap.print(dbgs() << "\nWith map: ")); - LLVM_DEBUG(expr.print(dbgs() << "\nResult: ")); -} - -OpPointer -mlir::makeNormalizedAffineApply(FuncBuilder *b, Location loc, AffineMap map, - ArrayRef operands) { - SingleResultAffineNormalizer normalizer(map, operands); - return b->create(loc, normalizer.getAffineMap(), - normalizer.getOperands()); -}