[mlir] Use SmallBitVector instead of SmallDenseSet for AffineMap::compressSymbols

[bkramer]

Summary

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.

[rriddle]

Nice.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1739

Is this mlir:: necessary?

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

940

Is the mlir:: here necessary?

mlir/include/mlir/Dialect/Arithmetic/Utils/Utils.h

 /// Matches a ConstantIndexOp.
 detail::op_matcher<arith::ConstantIndexOp> matchConstantIndex();
 
 /// Detects the `values` produced by a ConstantIndexOp and places the new
 /// constant in place of the corresponding sentinel value.
 void canonicalizeSubViewPart(SmallVectorImpl<OpFoldResult> &values,
                              function_ref<bool(int64_t)> isDynamic);
 
-void getPositionsOfShapeOne(unsigned rank, ArrayRef<int64_t> shape,
-                            llvm::SmallDenseSet<unsigned> &dimsToProject);
+llvm::SmallBitVector getPositionsOfShapeOne(unsigned rank,
+                                            ArrayRef<int64_t> shape);
 
 /// Pattern to rewrite a subview op with constant arguments.
 template <typename OpType, typename ResultTypeFunc, typename CastOpFunc>
 class OpWithOffsetSizesAndStridesConstantArgumentFolder final
     : public OpRewritePattern<OpType> {
 public:
   using OpRewritePattern<OpType>::OpRewritePattern;
 

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let extraClassDeclaration = extraBaseClassDeclaration # [{
     }
 
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
 
     /// Return the dimensions of the source type that are dropped when
     /// the result is rank-reduced.
-    llvm::SmallDenseSet<unsigned> getDroppedDims();
+    llvm::SmallBitVector getDroppedDims();
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

let extraClassDeclaration = extraBaseClassDeclaration # [{
     }
 
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
 
     /// Return the dimensions of the source that are dropped in the
     /// result when the result is rank-reduced.
-    llvm::SmallDenseSet<unsigned> getDroppedDims();
-
+    llvm::SmallBitVector getDroppedDims();
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/AffineMap.h

 
 #ifndef MLIR_IR_AFFINEMAP_H
 #define MLIR_IR_AFFINEMAP_H
 
 #include "mlir/IR/AffineExpr.h"
 #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 {
 
 namespace detail {
 struct AffineMapStorage;
 } // namespace detail
 
 class Attribute;

 AffineMap compressUnusedDims(AffineMap map);
 
 /// Drop the dims that are not used by any of the individual maps in `maps`.
 /// Asserts that all maps in `maps` are normalized to the same number of
 /// dims and symbols.
 SmallVector<AffineMap> compressUnusedDims(ArrayRef<AffineMap> maps);
 
 /// Drop the dims that are not listed in `unusedDims`.
-AffineMap compressDims(AffineMap map,
-                       const llvm::SmallDenseSet<unsigned> &unusedDims);
+AffineMap compressDims(AffineMap map, const llvm::SmallBitVector &unusedDims);
 
 /// Drop the symbols that are not used.
 AffineMap compressUnusedSymbols(AffineMap map);
 
 /// Drop the symbols that are not used by any of the individual maps in `maps`.
 /// Asserts that all maps in `maps` are normalized to the same number of
 /// dims and symbols.
 SmallVector<AffineMap> compressUnusedSymbols(ArrayRef<AffineMap> maps);
 
 /// Drop the symbols that are not listed in `unusedSymbols`.
 AffineMap compressSymbols(AffineMap map,
-                          const llvm::SmallDenseSet<unsigned> &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`.
 AffineMap removeDuplicateExprs(AffineMap map);
 
 /// Returns a map of codomain to domain dimensions such that the first codomain
 /// dimension for a particular domain dimension is selected.
 /// Returns an empty map if the input map is empty.

 ///    projected_dimensions : {1}
 ///    result               : affine_map<(d0) -> (d0)>
 ///
 /// 3) map                  : affine_map<(d0, d1, d2) -> (d0, d1)>
 ///    projected_dimensions : {1}
 ///    result               : affine_map<(d0, d1) -> (d0, 0)>
 ///
 /// This function also compresses unused symbols away.
-AffineMap
-getProjectedMap(AffineMap map,
-                const llvm::SmallDenseSet<unsigned> &projectedDimensions);
+AffineMap getProjectedMap(AffineMap map,
+                          const llvm::SmallBitVector &projectedDimensions);
 
 /// Apply a permutation from `map` to `source` and return the result.
 template <typename T>
 SmallVector<T> applyPermutationMap(AffineMap map, llvm::ArrayRef<T> source) {
   assert(map.isProjectedPermutation());
   assert(map.getNumInputs() == source.size());
   SmallVector<T> result;
   result.reserve(map.getNumResults());

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

 #include "mlir/Conversion/LLVMCommon/Pattern.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Conversion/MemRefToLLVM/AllocLikeConversion.h"
 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #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;
 
 namespace {
 
 struct AllocOpLowering : public AllocLikeOpLLVMLowering {
   AllocOpLowering(LLVMTypeConverter &converter)
       : AllocLikeOpLLVMLowering(memref::AllocOp::getOperationName(),
if (!ShapedType::isDynamicStrideOrOffset(offset)) {
       targetMemRef.setOffset(rewriter, loc, baseOffset);
     }
 
     // Update sizes and strides.
     SmallVector<OpFoldResult> mixedSizes = subViewOp.getMixedSizes();
     SmallVector<OpFoldResult> mixedStrides = subViewOp.getMixedStrides();
     assert(mixedSizes.size() == mixedStrides.size() &&
            "expected sizes and strides of equal length");
-    llvm::SmallDenseSet<unsigned> 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.
       // In this case, the size is guaranteed to be interpreted as Dim and the
       // stride as 1.
       Value size, stride;
       if (static_cast<unsigned>(i) >= mixedSizes.size()) {
         // If the static size is available, use it directly. This is similar to

mlir/lib/Dialect/Arithmetic/Utils/Utils.cpp

 //===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities for the Linalg dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 
 /// Matches a ConstantIndexOp.
 /// TODO: This should probably just be a general matcher that uses matchConstant
 /// and checks the operation for an index type.
 detail::op_matcher<arith::ConstantIndexOp> mlir::matchConstantIndex() {
   return detail::op_matcher<arith::ConstantIndexOp>();
if (ofr.is<Attribute>())
       continue;
     // Newly static, move from Value to constant.
     if (auto cstOp =
             ofr.dyn_cast<Value>().getDefiningOp<arith::ConstantIndexOp>())
       ofr = OpBuilder(cstOp).getIndexAttr(cstOp.value());
   }
 }
 
-void mlir::getPositionsOfShapeOne(
-    unsigned rank, ArrayRef<int64_t> shape,
-    llvm::SmallDenseSet<unsigned> &dimsToProject) {
-  dimsToProject.reserve(rank);
+llvm::SmallBitVector mlir::getPositionsOfShapeOne(unsigned rank,
+                                                  ArrayRef<int64_t> 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,
                                             OpFoldResult ofr) {
   if (auto value = ofr.dyn_cast<Value>())
     return value;
   auto attr = ofr.dyn_cast<Attribute>().dyn_cast<IntegerAttr>();
   assert(attr && "expect the op fold result casts to an integer attribute");

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

if (consumerType != def.getType())
     def = b.create<tensor::CastOp>(fusedProducer.getLoc(), consumerType, def);
   consumerOpOperand.set(def);
   return FusionInfo{cast<LinalgOp>(producerOpResult.getOwner()), fusedProducer};
 }
 
 /// Prune all dimensions that are of reduction iterator type from `map`.
 static AffineMap pruneReductionDimsFromMap(ArrayRef<Attribute> iteratorTypes,
                                            AffineMap map) {
-  llvm::SmallDenseSet<unsigned> 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);
 }
 
 /// Returns the mapping from iterations in the consumer that write to the same
 /// location as the iterations in the producer. To do so use
 /// - indexing map of the fused view in the consumer : consumerIndexMap
 /// - indexing map of the fused view in the producer : producerIndexMap

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

static LogicalResult padOperandToSmallestStaticBoundingBox(
 
   // Cannot construct a static bounding box if the operand is not defined by an
   // ExtractSliceOp.
   auto sliceOp = opOperand->get().getDefiningOp<tensor::ExtractSliceOp>();
   if (!sliceOp)
     return failure(hasDynamicShape);
 
   // Compute the dropped dimensions if `sliceOp` is ranke-reducing.
-  llvm::SmallDenseSet<unsigned> droppedDims = sliceOp.getDroppedDims();
+  llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
 
   // Upper bound the `sliceOp` sizes to obtain a static bounding box.
   SmallVector<int64_t> staticSizes;
   staticSizes.reserve(shape.size());
   auto shapedOp = cast<OffsetSizeAndStrideOpInterface>(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<Attribute>()) {
       staticSizes.push_back(
           en.value().get<Attribute>().dyn_cast<IntegerAttr>().getInt());
       continue;
     }
     // Otherwise, try to compute a constant upper bound for the size value.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #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;
 
 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *MemRefDialect::materializeConstant(OpBuilder &builder,
                                               Attribute value, Type type,

 
 /// Given the `originalType` and a `candidateReducedType` whose shape is assumed
 /// to be a subset of `originalType` with some `1` entries erased, return the
 /// set of indices that specifies which of the entries of `originalShape` are
 /// dropped to obtain `reducedShape`.
 /// 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<llvm::SmallDenseSet<unsigned>>
+static llvm::Optional<llvm::SmallBitVector>
 computeMemRefRankReductionMask(MemRefType originalType, MemRefType reducedType,
                                ArrayRef<OpFoldResult> sizes) {
-  llvm::SmallDenseSet<unsigned> 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<Attribute>())
       if (attr.cast<IntegerAttr>().getInt() == 1)
-        unusedDims.insert(dim.index());
+        unusedDims.set(dim.index());
 
   SmallVector<int64_t> originalStrides, candidateStrides;
   int64_t originalOffset, candidateOffset;
   if (failed(
           getStridesAndOffset(originalType, originalStrides, originalOffset)) ||
       failed(
           getStridesAndOffset(reducedType, candidateStrides, candidateOffset)))
     return llvm::None;
 
   // For memrefs, a dimension is truly dropped if its corresponding stride is
   // also dropped. This is particularly important when more than one of the dims
   // is 1. Track the number of occurences of the strides in the original type
   // and the candidate type. For each unused dim that stride should not be
   // present in the candidate type. Note that there could be multiple dimensions
   // that have the same size. We dont need to exactly figure out which dim
   // corresponds to which stride, we just need to verify that the number of
   // reptitions of a stride in the original + number of unused dims with that
   // stride == number of repititions of a stride in the candidate.
   std::map<int64_t, unsigned> currUnaccountedStrides =
       getNumOccurences(originalStrides);
   std::map<int64_t, unsigned> candidateStridesNumOccurences =
       getNumOccurences(candidateStrides);
-  llvm::SmallDenseSet<unsigned> 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]) {
       // This dim can be treated as dropped.
       currUnaccountedStrides[originalStride]--;
       continue;
     }
     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] <
         candidateStridesNumOccurences[originalStride]) {
       // This should never happen. Cant have a stride in the reduced rank type
       // that wasnt in the original one.
       return llvm::None;
     }
   }
 
-  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<unsigned> SubViewOp::getDroppedDims() {
+llvm::SmallBitVector SubViewOp::getDroppedDims() {
   MemRefType sourceType = getSourceType();
   MemRefType resultType = getType();
-  llvm::Optional<llvm::SmallDenseSet<unsigned>> unusedDims =
+  llvm::Optional<llvm::SmallBitVector> unusedDims =
       computeMemRefRankReductionMask(sourceType, resultType, getMixedSizes());
   assert(unusedDims && "unable to find unused dims of subview");
   return *unusedDims;
 }
 
 OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
   // All forms of folding require a known index.
   auto index = operands[1].dyn_cast_or_null<IntegerAttr>();
if (auto alloca = dyn_cast_or_null<AllocaOp>(definingOp))
     return *(alloca.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto view = dyn_cast_or_null<ViewOp>(definingOp))
     return *(view.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto subview = dyn_cast_or_null<SubViewOp>(definingOp)) {
-    llvm::SmallDenseSet<unsigned> unusedDims = subview.getDroppedDims();
+    llvm::SmallBitVector unusedDims = subview.getDroppedDims();
     unsigned resultIndex = 0;
     unsigned sourceRank = subview.getSourceType().getRank();
     unsigned sourceIndex = 0;
     for (auto i : llvm::seq<unsigned>(0, sourceRank)) {
-      if (unusedDims.count(i))
+      if (unusedDims.test(i))
         continue;
       if (resultIndex == unsignedIndex) {
         sourceIndex = i;
         break;
       }
       resultIndex++;
     }
     assert(subview.isDynamicSize(sourceIndex) &&
Type SubViewOp::inferRankReducedResultType(unsigned resultRank,
                                            ArrayRef<int64_t> strides) {
   auto inferredType =
       inferResultType(sourceRankedTensorType, offsets, sizes, strides)
           .cast<MemRefType>();
   assert(inferredType.getRank() >= resultRank && "expected ");
   int rankDiff = inferredType.getRank() - resultRank;
   if (rankDiff > 0) {
     auto shape = inferredType.getShape();
-    llvm::SmallDenseSet<unsigned> dimsToProject;
-    mlir::getPositionsOfShapeOne(rankDiff, shape, dimsToProject);
+    llvm::SmallBitVector dimsToProject =
+        getPositionsOfShapeOne(rankDiff, shape);
     SmallVector<int64_t> 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();
     if (!map.isIdentity())
       map = getProjectedMap(map, dimsToProject);
     inferredType =
         MemRefType::get(projectedShape, inferredType.getElementType(), map,
                         inferredType.getMemorySpace());

 /// `currentResultType`.
 static MemRefType getCanonicalSubViewResultType(
     MemRefType currentResultType, MemRefType currentSourceType,
     MemRefType sourceType, ArrayRef<OpFoldResult> mixedOffsets,
     ArrayRef<OpFoldResult> mixedSizes, ArrayRef<OpFoldResult> mixedStrides) {
   auto nonRankReducedType = SubViewOp::inferResultType(sourceType, mixedOffsets,
                                                        mixedSizes, mixedStrides)
                                 .cast<MemRefType>();
-  llvm::Optional<llvm::SmallDenseSet<unsigned>> unusedDims =
+  llvm::Optional<llvm::SmallBitVector> unusedDims =
       computeMemRefRankReductionMask(currentSourceType, currentResultType,
                                      mixedSizes);
   // Return nullptr as failure mode.
   if (!unusedDims)
     return nullptr;
   SmallVector<int64_t> 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());
   }
   AffineMap layoutMap = nonRankReducedType.getLayout().getAffineMap();
   if (!layoutMap.isIdentity())
     layoutMap = getProjectedMap(layoutMap, unusedDims.getValue());
   return MemRefType::get(shape, nonRankReducedType.getElementType(), layoutMap,
                          nonRankReducedType.getMemorySpace());

mlir/lib/Dialect/MemRef/Transforms/FoldSubViewOps.cpp

 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #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;
 
 //===----------------------------------------------------------------------===//
 // Utility functions
 //===----------------------------------------------------------------------===//
 
 /// Given the 'indices' of an load/store operation where the memref is a result
resolveSourceIndices(Location loc, PatternRewriter &rewriter,
   SmallVector<OpFoldResult> mixedOffsets = subViewOp.getMixedOffsets();
   SmallVector<OpFoldResult> mixedSizes = subViewOp.getMixedSizes();
   SmallVector<OpFoldResult> mixedStrides = subViewOp.getMixedStrides();
 
   SmallVector<Value> useIndices;
   // Check if this is rank-reducing case. Then for every unit-dim size add a
   // zero to the indices.
   unsigned resultDim = 0;
-  llvm::SmallDenseSet<unsigned> unusedDims = subViewOp.getDroppedDims();
+  llvm::SmallBitVector unusedDims = subViewOp.getDroppedDims();
   for (auto dim : llvm::seq<unsigned>(0, subViewOp.getSourceType().getRank())) {
-    if (unusedDims.count(dim))
+    if (unusedDims.test(dim))
       useIndices.push_back(rewriter.create<arith::ConstantIndexOp>(loc, 0));
     else
       useIndices.push_back(indices[resultDim++]);
   }
   if (useIndices.size() != mixedOffsets.size())
     return failure();
   sourceIndices.resize(useIndices.size());
   for (auto index : llvm::seq<size_t>(0, mixedOffsets.size())) {

 }
 
 /// Given the permutation map of the original
 /// `vector.transfer_read`/`vector.transfer_write` operations compute the
 /// permutation map to use after the subview is folded with it.
 static AffineMapAttr getPermutationMapAttr(MLIRContext *context,
                                            memref::SubViewOp subViewOp,
                                            AffineMap currPermutationMap) {
-  llvm::SmallDenseSet<unsigned> unusedDims = subViewOp.getDroppedDims();
+  llvm::SmallBitVector unusedDims = subViewOp.getDroppedDims();
   SmallVector<AffineExpr> exprs;
   int64_t sourceRank = subViewOp.getSourceType().getRank();
   for (auto dim : llvm::seq<int64_t>(0, sourceRank)) {
-    if (unusedDims.count(dim))
+    if (unusedDims.test(dim))
       continue;
     exprs.push_back(getAffineDimExpr(dim, context));
   }
   auto resultDimToSourceDimMap = AffineMap::get(sourceRank, 0, exprs, context);
   return AffineMapAttr::get(
       currPermutationMap.compose(resultDimToSourceDimMap));
 }
 

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/Matchers.h"
 #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;
 
 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *TensorDialect::materializeConstant(OpBuilder &builder,
                                               Attribute value, Type type,
RankedTensorType ExtractSliceOp::inferRankReducedResultType(
     ArrayRef<int64_t> offsets, ArrayRef<int64_t> sizes,
     ArrayRef<int64_t> strides) {
   auto inferredType =
       inferResultType(sourceRankedTensorType, offsets, sizes, strides)
           .cast<RankedTensorType>();
   int rankDiff = inferredType.getRank() - resultRank;
   if (rankDiff > 0) {
     auto shape = inferredType.getShape();
-    llvm::SmallDenseSet<unsigned> dimsToProject;
-    mlir::getPositionsOfShapeOne(rankDiff, shape, dimsToProject);
+    llvm::SmallBitVector dimsToProject =
+        getPositionsOfShapeOne(rankDiff, shape);
     SmallVector<int64_t> 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());
   }
   return inferredType;
 }
 
 RankedTensorType ExtractSliceOp::inferRankReducedResultType(
getCanonicalSliceResultType(unsigned resultRank, RankedTensorType sourceType,
   if (resultType.getRank() != resultRank) {
     resultType = ExtractSliceOp::inferResultType(sourceType, mixedOffsets,
                                                  mixedSizes, mixedStrides)
                      .cast<RankedTensorType>();
   }
   return resultType;
 }
 
-llvm::SmallDenseSet<unsigned> ExtractSliceOp::getDroppedDims() {
-  llvm::SmallDenseSet<unsigned> droppedDims;
+llvm::SmallBitVector ExtractSliceOp::getDroppedDims() {
   ArrayRef<int64_t> resultShape = getType().getShape();
   SmallVector<OpFoldResult> mixedSizes = getMixedSizes();
+  llvm::SmallBitVector droppedDims(mixedSizes.size());
   unsigned shapePos = 0;
   for (const auto &size : enumerate(mixedSizes)) {
     Optional<int64_t> sizeVal = getConstantIntValue(size.value());
     // If the size is not 1, or if the current matched dimension of the result
     // is the same static shape as the size value (which is 1), then the
     // dimension is preserved.
     if (!sizeVal || sizeVal.getValue() != 1 ||
         (shapePos < resultShape.size() && resultShape[shapePos] == 1)) {
       shapePos++;
       continue;
     }
-    droppedDims.insert(size.index());
+    droppedDims.set(size.index());
   }
   return droppedDims;
 }
 
 LogicalResult ExtractSliceOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   reifiedReturnShapes.resize(1);
   reifiedReturnShapes[0].reserve(getType().getRank());
   SmallVector<OpFoldResult> mixedSizes = getMixedSizes();
-  llvm::SmallDenseSet<unsigned> 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<Attribute>()) {
       reifiedReturnShapes[0].push_back(builder.create<arith::ConstantIndexOp>(
           loc, attr.cast<IntegerAttr>().getInt()));
       continue;
     }
     reifiedReturnShapes[0].push_back(size.value().get<Value>());
   }

mlir/lib/IR/AffineMap.cpp

if (numResults == 0)
     return AffineMap();
   if (numResults > getNumResults())
     return *this;
   return getSubMap(llvm::to_vector<4>(
       llvm::seq<unsigned>(getNumResults() - numResults, getNumResults())));
 }
 
 AffineMap mlir::compressDims(AffineMap map,
-                             const llvm::SmallDenseSet<unsigned> &unusedDims) {
+                             const llvm::SmallBitVector &unusedDims) {
   unsigned numDims = 0;
   SmallVector<AffineExpr> 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));
   }
   SmallVector<AffineExpr> resultExprs;
   resultExprs.reserve(map.getNumResults());
   for (auto e : map.getResults())
     resultExprs.push_back(e.replaceDims(dimReplacements));
   return AffineMap::get(numDims, map.getNumSymbols(), resultExprs, context);
 }
 
 AffineMap mlir::compressUnusedDims(AffineMap map) {
-  llvm::SmallDenseSet<unsigned> usedDims;
+  llvm::SmallBitVector unusedDims(map.getNumDims(), true);
   map.walkExprs([&](AffineExpr expr) {
     if (auto dimExpr = expr.dyn_cast<AffineDimExpr>())
-      usedDims.insert(dimExpr.getPosition());
+      unusedDims.reset(dimExpr.getPosition());
   });
-  llvm::SmallDenseSet<unsigned> unusedDims;
-  for (unsigned d = 0, e = map.getNumDims(); d != e; ++d)
-    if (!usedDims.contains(d))
-      unusedDims.insert(d);
   return compressDims(map, unusedDims);
 }
 
 static SmallVector<AffineMap>
 compressUnusedImpl(ArrayRef<AffineMap> maps,
                    llvm::function_ref<AffineMap(AffineMap)> compressionFun) {
   if (maps.empty())
     return SmallVector<AffineMap>();
compressUnusedImpl(ArrayRef<AffineMap> maps,
   return res;
 }
 
 SmallVector<AffineMap> mlir::compressUnusedDims(ArrayRef<AffineMap> maps) {
   return compressUnusedImpl(maps,
                             [](AffineMap m) { return compressUnusedDims(m); });
 }
 
-AffineMap
-mlir::compressSymbols(AffineMap map,
-                      const llvm::SmallDenseSet<unsigned> &unusedSymbols) {
+AffineMap mlir::compressSymbols(AffineMap map,
+                                const llvm::SmallBitVector &unusedSymbols) {
   unsigned numSymbols = 0;
   SmallVector<AffineExpr> 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));
   }
   SmallVector<AffineExpr> resultExprs;
   resultExprs.reserve(map.getNumResults());
   for (auto e : map.getResults())
     resultExprs.push_back(e.replaceSymbols(symReplacements));
   return AffineMap::get(map.getNumDims(), numSymbols, resultExprs, context);
 }
 
 AffineMap mlir::compressUnusedSymbols(AffineMap map) {
-  llvm::SmallDenseSet<unsigned> usedSymbols;
+  llvm::SmallBitVector unusedSymbols(map.getNumSymbols(), true);
   map.walkExprs([&](AffineExpr expr) {
     if (auto symExpr = expr.dyn_cast<AffineSymbolExpr>())
-      usedSymbols.insert(symExpr.getPosition());
+      unusedSymbols.reset(symExpr.getPosition());
   });
-  llvm::SmallDenseSet<unsigned> unusedSymbols;
-  for (unsigned d = 0, e = map.getNumSymbols(); d != e; ++d)
-    if (!usedSymbols.contains(d))
-      unusedSymbols.insert(d);
   return compressSymbols(map, unusedSymbols);
 }
 
 SmallVector<AffineMap> mlir::compressUnusedSymbols(ArrayRef<AffineMap> maps) {
   return compressUnusedImpl(
       maps, [](AffineMap m) { return compressUnusedSymbols(m); });
 }
 
for (auto m : maps) {
 
     numSymbols += m.getNumSymbols();
     numDims = std::max(m.getNumDims(), numDims);
   }
   return AffineMap::get(numDims, numSymbols, results,
                         maps.front().getContext());
 }
 
-AffineMap
-mlir::getProjectedMap(AffineMap map,
-                      const llvm::SmallDenseSet<unsigned> &unusedDims) {
+AffineMap mlir::getProjectedMap(AffineMap map,
+                                const llvm::SmallBitVector &unusedDims) {
   return compressUnusedSymbols(compressDims(map, unusedDims));
 }
 
 //===----------------------------------------------------------------------===//
 // MutableAffineMap.
 //===----------------------------------------------------------------------===//
 
 MutableAffineMap::MutableAffineMap(AffineMap map)

mlir/test/Dialect/MemRef/fold-subview-ops.mlir

 // CHECK-SAME:   %[[ARG11:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG12:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG13:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG14:[a-zA-Z0-9_]+]]: index
 // CHECK-SAME:   %[[ARG15:[a-zA-Z0-9_]+]]: index
 // 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]]]
 //      CHECK:   memref.load %[[ARG0]][%[[I1]], %[[I2]], %[[I3]], %[[I4]], %[[I5]], %[[I6]]]
 
 // -----
 
 func @fold_vector_transfer_read_with_rank_reduced_subview(