From 6635c12ada0d2d4f2e3b579e818980d0a587135f Mon Sep 17 00:00:00 2001 From: Benjamin Kramer Date: Sun, 6 Feb 2022 14:06:34 +0100 Subject: [PATCH] [mlir] Use SmallBitVector instead of SmallDenseSet for AffineMap::compressSymbols This is both more efficient and more ergonomic to use, as inverting a bit vector is trivial while inverting a set is annoying. Sadly this leaks into a bunch of APIs downstream, so adapt them as well. This would be NFC, but there is an ordering dependency in MemRefOps's computeMemRefRankReductionMask. This is now deterministic, previously it was dependent on SmallDenseSet's unspecified iteration order. Differential Revision: https://reviews.llvm.org/D119076 --- .../mlir/Dialect/Arithmetic/Utils/Utils.h | 4 +- .../mlir/Dialect/MemRef/IR/MemRefOps.td | 2 +- .../mlir/Dialect/Tensor/IR/TensorOps.td | 3 +- mlir/include/mlir/IR/AffineMap.h | 15 ++++---- .../Conversion/MemRefToLLVM/MemRefToLLVM.cpp | 5 ++- mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp | 11 +++--- mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp | 4 +- .../Dialect/Linalg/Transforms/Transforms.cpp | 4 +- mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp | 37 ++++++++++--------- .../MemRef/Transforms/FoldSubViewOps.cpp | 9 +++-- mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 17 +++++---- mlir/lib/IR/AffineMap.cpp | 32 ++++++---------- .../test/Dialect/MemRef/fold-subview-ops.mlir | 4 +- 13 files changed, 71 insertions(+), 76 deletions(-) diff --git a/mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h b/mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h index a4dcf9622098..1d4353f176b1 100644 --- a/mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h +++ b/mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h @@ -31,8 +31,8 @@ detail::op_matcher matchConstantIndex(); void canonicalizeSubViewPart(SmallVectorImpl &values, function_ref isDynamic); -void getPositionsOfShapeOne(unsigned rank, ArrayRef shape, - llvm::SmallDenseSet &dimsToProject); +llvm::SmallBitVector getPositionsOfShapeOne(unsigned rank, + ArrayRef shape); /// Pattern to rewrite a subview op with constant arguments. template diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td index 6438e707f7c2..81839dfde9c5 100644 --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td @@ -1639,7 +1639,7 @@ def SubViewOp : MemRef_OpWithOffsetSizesAndStrides<"subview", [ /// Return the dimensions of the source type that are dropped when /// the result is rank-reduced. - llvm::SmallDenseSet getDroppedDims(); + llvm::SmallBitVector getDroppedDims(); }]; let hasCanonicalizer = 1; diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td index ec597e9f650a..a4dc2953b201 100644 --- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td +++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td @@ -316,8 +316,7 @@ def Tensor_ExtractSliceOp : Tensor_OpWithOffsetSizesAndStrides<"extract_slice", /// Return the dimensions of the source that are dropped in the /// result when the result is rank-reduced. - llvm::SmallDenseSet getDroppedDims(); - + llvm::SmallBitVector getDroppedDims(); }]; let hasCanonicalizer = 1; diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h index cd6fce22b614..82a33b065440 100644 --- a/mlir/include/mlir/IR/AffineMap.h +++ b/mlir/include/mlir/IR/AffineMap.h @@ -18,7 +18,10 @@ #include "mlir/Support/LLVM.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMapInfo.h" -#include "llvm/ADT/DenseSet.h" + +namespace llvm { +class SmallBitVector; +} // namespace llvm namespace mlir { @@ -372,8 +375,7 @@ AffineMap compressUnusedDims(AffineMap map); SmallVector compressUnusedDims(ArrayRef maps); /// Drop the dims that are not listed in `unusedDims`. -AffineMap compressDims(AffineMap map, - const llvm::SmallDenseSet &unusedDims); +AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims); /// Drop the symbols that are not used. AffineMap compressUnusedSymbols(AffineMap map); @@ -385,7 +387,7 @@ SmallVector compressUnusedSymbols(ArrayRef maps); /// Drop the symbols that are not listed in `unusedSymbols`. AffineMap compressSymbols(AffineMap map, - const llvm::SmallDenseSet &unusedSymbols); + const llvm::SmallBitVector &unusedSymbols); /// Returns a map with the same dimension and symbol count as `map`, but whose /// results are the unique affine expressions of `map`. @@ -521,9 +523,8 @@ AffineMap concatAffineMaps(ArrayRef maps); /// result : affine_map<(d0, d1) -> (d0, 0)> /// /// This function also compresses unused symbols away. -AffineMap -getProjectedMap(AffineMap map, - const llvm::SmallDenseSet &projectedDimensions); +AffineMap getProjectedMap(AffineMap map, + const llvm::SmallBitVector &projectedDimensions); /// Apply a permutation from `map` to `source` and return the result. template diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp index a4a82d66b0cb..d21d21e045fb 100644 --- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp +++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp @@ -18,6 +18,7 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/IR/AffineMap.h" #include "mlir/IR/BlockAndValueMapping.h" +#include "llvm/ADT/SmallBitVector.h" using namespace mlir; @@ -1503,10 +1504,10 @@ struct SubViewOpLowering : public ConvertOpToLLVMPattern { SmallVector mixedStrides = subViewOp.getMixedStrides(); assert(mixedSizes.size() == mixedStrides.size() && "expected sizes and strides of equal length"); - llvm::SmallDenseSet unusedDims = subViewOp.getDroppedDims(); + llvm::SmallBitVector unusedDims = subViewOp.getDroppedDims(); for (int i = inferredShapeRank - 1, j = resultShapeRank - 1; i >= 0 && j >= 0; --i) { - if (unusedDims.contains(i)) + if (unusedDims.test(i)) continue; // `i` may overflow subViewOp.getMixedSizes because of trailing semantics. diff --git a/mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp b/mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp index 74dd373f17cd..4e35c5b31924 100644 --- a/mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp @@ -12,6 +12,7 @@ #include "mlir/Dialect/Arithmetic/Utils/Utils.h" #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" +#include "llvm/ADT/SmallBitVector.h" using namespace mlir; @@ -37,16 +38,16 @@ void mlir::canonicalizeSubViewPart( } } -void mlir::getPositionsOfShapeOne( - unsigned rank, ArrayRef shape, - llvm::SmallDenseSet &dimsToProject) { - dimsToProject.reserve(rank); +llvm::SmallBitVector mlir::getPositionsOfShapeOne(unsigned rank, + ArrayRef shape) { + llvm::SmallBitVector dimsToProject(shape.size()); for (unsigned pos = 0, e = shape.size(); pos < e && rank > 0; ++pos) { if (shape[pos] == 1) { - dimsToProject.insert(pos); + dimsToProject.set(pos); --rank; } } + return dimsToProject; } Value mlir::getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, diff --git a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp index 9eda19ae3f4b..32d8ee098bce 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp @@ -520,10 +520,10 @@ mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult, /// Prune all dimensions that are of reduction iterator type from `map`. static AffineMap pruneReductionDimsFromMap(ArrayRef iteratorTypes, AffineMap map) { - llvm::SmallDenseSet projectedDims; + llvm::SmallBitVector projectedDims(iteratorTypes.size()); for (const auto &attr : llvm::enumerate(iteratorTypes)) { if (!isParallelIterator(attr.value())) - projectedDims.insert(attr.index()); + projectedDims.set(attr.index()); } return getProjectedMap(map, projectedDims); } diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp index d8e00ca5d4ed..65a0ed8ec68f 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @@ -187,7 +187,7 @@ static LogicalResult padOperandToSmallestStaticBoundingBox( return failure(hasDynamicShape); // Compute the dropped dimensions if `sliceOp` is ranke-reducing. - llvm::SmallDenseSet droppedDims = sliceOp.getDroppedDims(); + llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims(); // Upper bound the `sliceOp` sizes to obtain a static bounding box. SmallVector staticSizes; @@ -195,7 +195,7 @@ static LogicalResult padOperandToSmallestStaticBoundingBox( auto shapedOp = cast(sliceOp.getOperation()); for (const auto &en : enumerate(shapedOp.getMixedSizes())) { // Skip dropped dimensions. - if (droppedDims.contains(en.index())) + if (droppedDims.test(en.index())) continue; // If the size is an attribute add it directly to `staticSizes`. if (en.value().is()) { diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp index dad063ea30bc..21cb8d60f2d3 100644 --- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp +++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp @@ -20,6 +20,7 @@ #include "mlir/Interfaces/InferTypeOpInterface.h" #include "mlir/Interfaces/ViewLikeInterface.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallBitVector.h" using namespace mlir; using namespace mlir::memref; @@ -590,17 +591,17 @@ static std::map getNumOccurences(ArrayRef vals) { /// This accounts for cases where there are multiple unit-dims, but only a /// subset of those are dropped. For MemRefTypes these can be disambiguated /// using the strides. If a dimension is dropped the stride must be dropped too. -static llvm::Optional> +static llvm::Optional computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType, ArrayRef sizes) { - llvm::SmallDenseSet unusedDims; + llvm::SmallBitVector unusedDims(originalType.getRank()); if (originalType.getRank() == reducedType.getRank()) return unusedDims; for (const auto &dim : llvm::enumerate(sizes)) if (auto attr = dim.value().dyn_cast()) if (attr.cast().getInt() == 1) - unusedDims.insert(dim.index()); + unusedDims.set(dim.index()); SmallVector originalStrides, candidateStrides; int64_t originalOffset, candidateOffset; @@ -623,8 +624,9 @@ computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType, getNumOccurences(originalStrides); std::map candidateStridesNumOccurences = getNumOccurences(candidateStrides); - llvm::SmallDenseSet prunedUnusedDims; - for (unsigned dim : unusedDims) { + for (size_t dim = 0, e = unusedDims.size(); dim != e; ++dim) { + if (!unusedDims.test(dim)) + continue; int64_t originalStride = originalStrides[dim]; if (currUnaccountedStrides[originalStride] > candidateStridesNumOccurences[originalStride]) { @@ -635,7 +637,7 @@ computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType, if (currUnaccountedStrides[originalStride] == candidateStridesNumOccurences[originalStride]) { // The stride for this is not dropped. Keep as is. - prunedUnusedDims.insert(dim); + unusedDims.reset(dim); continue; } if (currUnaccountedStrides[originalStride] < @@ -646,17 +648,16 @@ computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType, } } - for (auto prunedDim : prunedUnusedDims) - unusedDims.erase(prunedDim); - if (unusedDims.size() + reducedType.getRank() != originalType.getRank()) + if ((int64_t)unusedDims.count() + reducedType.getRank() != + originalType.getRank()) return llvm::None; return unusedDims; } -llvm::SmallDenseSet SubViewOp::getDroppedDims() { +llvm::SmallBitVector SubViewOp::getDroppedDims() { MemRefType sourceType = getSourceType(); MemRefType resultType = getType(); - llvm::Optional> unusedDims = + llvm::Optional unusedDims = computeMemRefRankReductionMask(sourceType, resultType, getMixedSizes()); assert(unusedDims && "unable to find unused dims of subview"); return *unusedDims; @@ -698,12 +699,12 @@ OpFoldResult DimOp::fold(ArrayRef operands) { memrefType.getDynamicDimIndex(unsignedIndex)); if (auto subview = dyn_cast_or_null(definingOp)) { - llvm::SmallDenseSet unusedDims = subview.getDroppedDims(); + llvm::SmallBitVector unusedDims = subview.getDroppedDims(); unsigned resultIndex = 0; unsigned sourceRank = subview.getSourceType().getRank(); unsigned sourceIndex = 0; for (auto i : llvm::seq(0, sourceRank)) { - if (unusedDims.count(i)) + if (unusedDims.test(i)) continue; if (resultIndex == unsignedIndex) { sourceIndex = i; @@ -1734,11 +1735,11 @@ Type SubViewOp::inferRankReducedResultType(unsigned resultRank, int rankDiff = inferredType.getRank() - resultRank; if (rankDiff > 0) { auto shape = inferredType.getShape(); - llvm::SmallDenseSet dimsToProject; - mlir::getPositionsOfShapeOne(rankDiff, shape, dimsToProject); + llvm::SmallBitVector dimsToProject = + getPositionsOfShapeOne(rankDiff, shape); SmallVector projectedShape; for (unsigned pos = 0, e = shape.size(); pos < e; ++pos) - if (!dimsToProject.contains(pos)) + if (!dimsToProject.test(pos)) projectedShape.push_back(shape[pos]); AffineMap map = inferredType.getLayout().getAffineMap(); @@ -2015,7 +2016,7 @@ static MemRefType getCanonicalSubViewResultType( auto nonRankReducedType = SubViewOp::inferResultType(sourceType, mixedOffsets, mixedSizes, mixedStrides) .cast(); - llvm::Optional> unusedDims = + llvm::Optional unusedDims = computeMemRefRankReductionMask(currentSourceType, currentResultType, mixedSizes); // Return nullptr as failure mode. @@ -2023,7 +2024,7 @@ static MemRefType getCanonicalSubViewResultType( return nullptr; SmallVector shape; for (const auto &sizes : llvm::enumerate(nonRankReducedType.getShape())) { - if (unusedDims->count(sizes.index())) + if (unusedDims->test(sizes.index())) continue; shape.push_back(sizes.value()); } diff --git a/mlir/lib/Dialect/MemRef/Transforms/FoldSubViewOps.cpp b/mlir/lib/Dialect/MemRef/Transforms/FoldSubViewOps.cpp index 213b442998b0..5080a69688b9 100644 --- a/mlir/lib/Dialect/MemRef/Transforms/FoldSubViewOps.cpp +++ b/mlir/lib/Dialect/MemRef/Transforms/FoldSubViewOps.cpp @@ -18,6 +18,7 @@ #include "mlir/Dialect/Vector/IR/VectorOps.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/ADT/SmallBitVector.h" using namespace mlir; @@ -50,9 +51,9 @@ resolveSourceIndices(Location loc, PatternRewriter &rewriter, // Check if this is rank-reducing case. Then for every unit-dim size add a // zero to the indices. unsigned resultDim = 0; - llvm::SmallDenseSet unusedDims = subViewOp.getDroppedDims(); + llvm::SmallBitVector unusedDims = subViewOp.getDroppedDims(); for (auto dim : llvm::seq(0, subViewOp.getSourceType().getRank())) { - if (unusedDims.count(dim)) + if (unusedDims.test(dim)) useIndices.push_back(rewriter.create(loc, 0)); else useIndices.push_back(indices[resultDim++]); @@ -106,11 +107,11 @@ static Value getMemRefOperand(vector::TransferWriteOp op) { static AffineMapAttr getPermutationMapAttr(MLIRContext *context, memref::SubViewOp subViewOp, AffineMap currPermutationMap) { - llvm::SmallDenseSet unusedDims = subViewOp.getDroppedDims(); + llvm::SmallBitVector unusedDims = subViewOp.getDroppedDims(); SmallVector exprs; int64_t sourceRank = subViewOp.getSourceType().getRank(); for (auto dim : llvm::seq(0, sourceRank)) { - if (unusedDims.count(dim)) + if (unusedDims.test(dim)) continue; exprs.push_back(getAffineDimExpr(dim, context)); } diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp index 58233d6dd6ff..a26ad993b495 100644 --- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp +++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp @@ -19,6 +19,7 @@ #include "mlir/IR/PatternMatch.h" #include "mlir/IR/TypeUtilities.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallBitVector.h" using namespace mlir; using namespace mlir::tensor; @@ -935,11 +936,11 @@ RankedTensorType ExtractSliceOp::inferRankReducedResultType( int rankDiff = inferredType.getRank() - resultRank; if (rankDiff > 0) { auto shape = inferredType.getShape(); - llvm::SmallDenseSet dimsToProject; - mlir::getPositionsOfShapeOne(rankDiff, shape, dimsToProject); + llvm::SmallBitVector dimsToProject = + getPositionsOfShapeOne(rankDiff, shape); SmallVector projectedShape; for (unsigned pos = 0, e = shape.size(); pos < e; ++pos) - if (!dimsToProject.contains(pos)) + if (!dimsToProject.test(pos)) projectedShape.push_back(shape[pos]); inferredType = RankedTensorType::get(projectedShape, inferredType.getElementType()); @@ -1076,10 +1077,10 @@ getCanonicalSliceResultType(unsigned resultRank, RankedTensorType sourceType, return resultType; } -llvm::SmallDenseSet ExtractSliceOp::getDroppedDims() { - llvm::SmallDenseSet droppedDims; +llvm::SmallBitVector ExtractSliceOp::getDroppedDims() { ArrayRef resultShape = getType().getShape(); SmallVector mixedSizes = getMixedSizes(); + llvm::SmallBitVector droppedDims(mixedSizes.size()); unsigned shapePos = 0; for (const auto &size : enumerate(mixedSizes)) { Optional sizeVal = getConstantIntValue(size.value()); @@ -1091,7 +1092,7 @@ llvm::SmallDenseSet ExtractSliceOp::getDroppedDims() { shapePos++; continue; } - droppedDims.insert(size.index()); + droppedDims.set(size.index()); } return droppedDims; } @@ -1101,10 +1102,10 @@ LogicalResult ExtractSliceOp::reifyResultShapes( reifiedReturnShapes.resize(1); reifiedReturnShapes[0].reserve(getType().getRank()); SmallVector mixedSizes = getMixedSizes(); - llvm::SmallDenseSet droppedDims = getDroppedDims(); + llvm::SmallBitVector droppedDims = getDroppedDims(); Location loc = getLoc(); for (const auto &size : enumerate(mixedSizes)) { - if (droppedDims.count(size.index())) + if (droppedDims.test(size.index())) continue; if (auto attr = size.value().dyn_cast()) { reifiedReturnShapes[0].push_back(builder.create( diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp index ecdf8376b5fc..7fbe7488b97e 100644 --- a/mlir/lib/IR/AffineMap.cpp +++ b/mlir/lib/IR/AffineMap.cpp @@ -547,13 +547,13 @@ AffineMap AffineMap::getMinorSubMap(unsigned numResults) const { } AffineMap mlir::compressDims(AffineMap map, - const llvm::SmallDenseSet &unusedDims) { + const llvm::SmallBitVector &unusedDims) { unsigned numDims = 0; SmallVector dimReplacements; dimReplacements.reserve(map.getNumDims()); MLIRContext *context = map.getContext(); for (unsigned dim = 0, e = map.getNumDims(); dim < e; ++dim) { - if (unusedDims.contains(dim)) + if (unusedDims.test(dim)) dimReplacements.push_back(getAffineConstantExpr(0, context)); else dimReplacements.push_back(getAffineDimExpr(numDims++, context)); @@ -566,15 +566,11 @@ AffineMap mlir::compressDims(AffineMap map, } AffineMap mlir::compressUnusedDims(AffineMap map) { - llvm::SmallDenseSet usedDims; + llvm::SmallBitVector unusedDims(map.getNumDims(), true); map.walkExprs([&](AffineExpr expr) { if (auto dimExpr = expr.dyn_cast()) - usedDims.insert(dimExpr.getPosition()); + unusedDims.reset(dimExpr.getPosition()); }); - llvm::SmallDenseSet unusedDims; - for (unsigned d = 0, e = map.getNumDims(); d != e; ++d) - if (!usedDims.contains(d)) - unusedDims.insert(d); return compressDims(map, unusedDims); } @@ -613,15 +609,14 @@ SmallVector mlir::compressUnusedDims(ArrayRef maps) { [](AffineMap m) { return compressUnusedDims(m); }); } -AffineMap -mlir::compressSymbols(AffineMap map, - const llvm::SmallDenseSet &unusedSymbols) { +AffineMap mlir::compressSymbols(AffineMap map, + const llvm::SmallBitVector &unusedSymbols) { unsigned numSymbols = 0; SmallVector symReplacements; symReplacements.reserve(map.getNumSymbols()); MLIRContext *context = map.getContext(); for (unsigned sym = 0, e = map.getNumSymbols(); sym < e; ++sym) { - if (unusedSymbols.contains(sym)) + if (unusedSymbols.test(sym)) symReplacements.push_back(getAffineConstantExpr(0, context)); else symReplacements.push_back(getAffineSymbolExpr(numSymbols++, context)); @@ -634,15 +629,11 @@ mlir::compressSymbols(AffineMap map, } AffineMap mlir::compressUnusedSymbols(AffineMap map) { - llvm::SmallDenseSet usedSymbols; + llvm::SmallBitVector unusedSymbols(map.getNumSymbols(), true); map.walkExprs([&](AffineExpr expr) { if (auto symExpr = expr.dyn_cast()) - usedSymbols.insert(symExpr.getPosition()); + unusedSymbols.reset(symExpr.getPosition()); }); - llvm::SmallDenseSet unusedSymbols; - for (unsigned d = 0, e = map.getNumSymbols(); d != e; ++d) - if (!usedSymbols.contains(d)) - unusedSymbols.insert(d); return compressSymbols(map, unusedSymbols); } @@ -732,9 +723,8 @@ AffineMap mlir::concatAffineMaps(ArrayRef maps) { maps.front().getContext()); } -AffineMap -mlir::getProjectedMap(AffineMap map, - const llvm::SmallDenseSet &unusedDims) { +AffineMap mlir::getProjectedMap(AffineMap map, + const llvm::SmallBitVector &unusedDims) { return compressUnusedSymbols(compressDims(map, unusedDims)); } diff --git a/mlir/test/Dialect/MemRef/fold-subview-ops.mlir b/mlir/test/Dialect/MemRef/fold-subview-ops.mlir index fc06cd35cc8c..6f8e46fab444 100644 --- a/mlir/test/Dialect/MemRef/fold-subview-ops.mlir +++ b/mlir/test/Dialect/MemRef/fold-subview-ops.mlir @@ -154,8 +154,8 @@ func @fold_rank_reducing_subview_with_load // CHECK-SAME: %[[ARG16:[a-zA-Z0-9_]+]]: index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[I1:.+]] = affine.apply #[[MAP]](%[[ARG13]])[%[[ARG7]], %[[ARG1]]] -// CHECK-DAG: %[[I2:.+]] = affine.apply #[[MAP]](%[[ARG14]])[%[[ARG8]], %[[ARG2]]] -// CHECK-DAG: %[[I3:.+]] = affine.apply #[[MAP]](%[[C0]])[%[[ARG9]], %[[ARG3]]] +// CHECK-DAG: %[[I2:.+]] = affine.apply #[[MAP]](%[[C0]])[%[[ARG8]], %[[ARG2]]] +// CHECK-DAG: %[[I3:.+]] = affine.apply #[[MAP]](%[[ARG14]])[%[[ARG9]], %[[ARG3]]] // CHECK-DAG: %[[I4:.+]] = affine.apply #[[MAP]](%[[ARG15]])[%[[ARG10]], %[[ARG4]]] // CHECK-DAG: %[[I5:.+]] = affine.apply #[[MAP]](%[[ARG16]])[%[[ARG11]], %[[ARG5]]] // CHECK-DAG: %[[I6:.+]] = affine.apply #[[MAP]](%[[C0]])[%[[ARG12]], %[[ARG6]]]