From 24e5a72dac3f3facd365b8aa6a81b8c62c119608 Mon Sep 17 00:00:00 2001 From: Nicolas Vasilache Date: Thu, 17 Jan 2019 14:33:09 -0800 Subject: [PATCH] Fix AffineApply corner case This CL adds a test reported by andydavis@ and fixes the corner case that appears when operands do not come from an AffineApply and no Dim composition is needed. In such cases, we would need to create an empty map which is disallowed. The composition in such cases becomes trivial: there is no composition. This CL also updates the name AffineNormalizer to AffineApplyNormalizer. PiperOrigin-RevId: 229819234 --- mlir/include/mlir/IR/BuiltinOps.h | 7 ++- mlir/lib/Analysis/AffineAnalysis.cpp | 59 +++++++++++-------- mlir/test/Transforms/compose-affine-maps.mlir | 25 ++++++++ 3 files changed, 63 insertions(+), 28 deletions(-) diff --git a/mlir/include/mlir/IR/BuiltinOps.h b/mlir/include/mlir/IR/BuiltinOps.h index d9ed7cfd9c4c..189e3e248ea3 100644 --- a/mlir/include/mlir/IR/BuiltinOps.h +++ b/mlir/include/mlir/IR/BuiltinOps.h @@ -389,16 +389,19 @@ private: explicit ReturnOp(const OperationInst *state) : Op(state) {} }; -// Prints dimension and symbol list. +/// Prints dimension and symbol list. void printDimAndSymbolList(OperationInst::const_operand_iterator begin, OperationInst::const_operand_iterator end, unsigned numDims, OpAsmPrinter *p); -// Parses dimension and symbol list and returns true if parsing failed. +/// Parses dimension and symbol list and returns true if parsing failed. bool parseDimAndSymbolList(OpAsmParser *parser, SmallVector &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 *operands); diff --git a/mlir/lib/Analysis/AffineAnalysis.cpp b/mlir/lib/Analysis/AffineAnalysis.cpp index 19283b319b62..6f0b72e3f09f 100644 --- a/mlir/lib/Analysis/AffineAnalysis.cpp +++ b/mlir/lib/Analysis/AffineAnalysis.cpp @@ -1353,10 +1353,10 @@ bool mlir::checkMemrefAccessDependence( namespace { -/// An `AffineNormalizer` is a helper class that is not visible to the user and -/// supports renumbering operands of AffineApplyOp. -/// This acts as a reindexing map of Value* to positional dims or symbols and -/// allows simplifications such as: +/// An `AffineApplyNormalizer` is a helper class that is not visible to the user +/// and supports renumbering operands of AffineApplyOp. 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) @@ -1367,8 +1367,8 @@ namespace { /// ```mlir /// %1 = affine_apply () -> (0) /// ``` -struct AffineNormalizer { - AffineNormalizer(AffineMap map, ArrayRef operands); +struct AffineApplyNormalizer { + AffineApplyNormalizer(AffineMap map, ArrayRef operands); /// Returns the AffineMap resulting from normalization. AffineMap getAffineMap() { return affineMap; } @@ -1381,18 +1381,18 @@ struct AffineNormalizer { private: /// Helper function to insert `v` into the coordinate system of the current - /// AffineNormalizer. Returns the AffineDimExpr with the corresponding + /// AffineApplyNormalizer. Returns the AffineDimExpr with the corresponding /// renumbered position. AffineDimExpr applyOneDim(Value *v); /// Given an `other` normalizer, this rewrites `other.affineMap` in the - /// coordinate system of the current AffineNormalizer. + /// coordinate system of the current AffineApplyNormalizer. /// Returns the rewritten AffineMap and updates the dims and symbols of /// `this`. - AffineMap renumber(const AffineNormalizer &other); + AffineMap renumber(const AffineApplyNormalizer &other); /// Given an `app`, rewrites `app.getAffineMap()` in the coordinate system of - /// the current AffineNormalizer. + /// the current AffineApplyNormalizer. /// Returns the rewritten AffineMap and updates the dims and symbols of /// `this`. AffineMap renumber(const AffineApplyOp &app); @@ -1418,15 +1418,15 @@ private: } static constexpr unsigned kMaxAffineApplyDepth = 1; - AffineNormalizer() { affineApplyDepth()++; } + AffineApplyNormalizer() { affineApplyDepth()++; } public: - ~AffineNormalizer() { affineApplyDepth()--; } + ~AffineApplyNormalizer() { affineApplyDepth()--; } }; } // namespace -AffineDimExpr AffineNormalizer::applyOneDim(Value *v) { +AffineDimExpr AffineApplyNormalizer::applyOneDim(Value *v) { DenseMap::iterator iterPos; bool inserted = false; std::tie(iterPos, inserted) = @@ -1438,7 +1438,7 @@ AffineDimExpr AffineNormalizer::applyOneDim(Value *v) { .cast(); } -AffineMap AffineNormalizer::renumber(const AffineNormalizer &other) { +AffineMap AffineApplyNormalizer::renumber(const AffineApplyNormalizer &other) { SmallVector dimRemapping; for (auto *v : other.reorderedDims) { auto kvp = other.dimValueToPosition.find(v); @@ -1461,15 +1461,15 @@ AffineMap AffineNormalizer::renumber(const AffineNormalizer &other) { dimRemapping.size(), symRemapping.size()); } -AffineMap AffineNormalizer::renumber(const AffineApplyOp &app) { +AffineMap AffineApplyNormalizer::renumber(const AffineApplyOp &app) { assert(app.getAffineMap().getRangeSizes().empty() && "Non-empty range sizes"); - // Create the AffineNormalizer for the operands of this - // AffineApplyOp and combine it with the current AffineNormalizer. + // Create the AffineApplyNormalizer for the operands of this + // AffineApplyOp and combine it with the current AffineApplyNormalizer. SmallVector operands( const_cast(app).getOperands().begin(), const_cast(app).getOperands().end()); - AffineNormalizer normalizer(app.getAffineMap(), operands); + AffineApplyNormalizer normalizer(app.getAffineMap(), operands); return renumber(normalizer); } @@ -1484,16 +1484,13 @@ static unsigned getIndexOf(Value *v, const AffineApplyOp &op) { return static_cast(-1); } -AffineNormalizer::AffineNormalizer(AffineMap map, ArrayRef operands) - : AffineNormalizer() { +AffineApplyNormalizer::AffineApplyNormalizer(AffineMap map, + ArrayRef operands) + : AffineApplyNormalizer() { 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; - } - SmallVector exprs; for (auto en : llvm::enumerate(operands)) { auto *t = en.value(); @@ -1505,6 +1502,9 @@ AffineNormalizer::AffineNormalizer(AffineMap map, ArrayRef operands) if (en.index() < map.getNumDims()) { exprs.push_back(applyOneDim(t)); } else { + // Composition of mathematical symbols must occur by concatenation. + // A subsequent canonicalization will drop duplicates. Duplicates are + // not dropped here because it would just amount to code duplication. concatenatedSymbols.push_back(t); } } else { @@ -1516,14 +1516,21 @@ AffineNormalizer::AffineNormalizer(AffineMap map, ArrayRef operands) } } + // Map is already composed. + if (exprs.empty()) { + affineMap = map; + return; + } + auto numDims = dimValueToPosition.size(); auto numSymbols = concatenatedSymbols.size() - map.getNumSymbols(); auto exprsMap = AffineMap::get(numDims, numSymbols, exprs, {}); LLVM_DEBUG(map.print(dbgs() << "\nCompose map: ")); LLVM_DEBUG(exprsMap.print(dbgs() << "\nWith map: ")); + LLVM_DEBUG(map.compose(exprsMap).print(dbgs() << "\nResult: ")); affineMap = simplifyAffineMap(map.compose(exprsMap)); - LLVM_DEBUG(affineMap.print(dbgs() << "\nResult: ")); + LLVM_DEBUG(affineMap.print(dbgs() << "\nSimplified result: ")); } /// Implements `map` and `operands` composition and simplification to support @@ -1532,7 +1539,7 @@ AffineNormalizer::AffineNormalizer(AffineMap map, ArrayRef operands) /// immediately deleted. static void composeAffineMapAndOperands(AffineMap *map, SmallVectorImpl *operands) { - AffineNormalizer normalizer(*map, *operands); + AffineApplyNormalizer normalizer(*map, *operands); auto normalizedMap = normalizer.getAffineMap(); auto normalizedOperands = normalizer.getOperands(); canonicalizeMapAndOperands(&normalizedMap, &normalizedOperands); diff --git a/mlir/test/Transforms/compose-affine-maps.mlir b/mlir/test/Transforms/compose-affine-maps.mlir index 9ea9838aa19f..08ad50f5ed21 100644 --- a/mlir/test/Transforms/compose-affine-maps.mlir +++ b/mlir/test/Transforms/compose-affine-maps.mlir @@ -27,6 +27,12 @@ // Affine maps for test case: arg_used_as_dim_and_symbol // CHECK: [[MAP14:#map[0-9]+]] = (d0, d1, d2, d3)[s0, s1] -> (d2, d0 * -1 - d1 + d3 + s0 + s1) +// Affine maps for test case: zero_map +// CHECK: [[MAP15:#map[0-9]+]] = ()[s0] -> (s0) + +// Affine maps for test case: zero_map +// CHECK: [[MAP16:#map[0-9]+]] = () -> (0) + // CHECK-LABEL: func @compose_affine_maps_1dto2d_no_symbols() { func @compose_affine_maps_1dto2d_no_symbols() { %0 = alloc() : memref<4x4xf32> @@ -213,3 +219,22 @@ func @arg_used_as_dim_and_symbol(%arg0: memref<100x100xf32>, %arg1: index) { } return } + +// CHECK-LABEL: func @trivial_maps +func @trivial_maps() { + %0 = alloc() : memref<10xf32> + %c0 = constant 0 : index + %cst = constant 0.000000e+00 : f32 + for %i1 = 0 to 10 { + %1 = affine_apply ()[s0] -> (s0)()[%c0] + // CHECK: {{.*}} = affine_apply [[MAP15]]()[%c0] + store %cst, %0[%1] : memref<10xf32> + %2 = load %0[%c0] : memref<10xf32> + + %3 = affine_apply ()[] -> (0)()[] + // CHECK: {{.*}} = affine_apply [[MAP16]]() + store %cst, %0[%3] : memref<10xf32> + %4 = load %0[%c0] : memref<10xf32> + } + return +}