[ADT] Allow specifying the size of resulting `SmallVector` in `map_to_vector`

[laszlokindrat]

Summary

This patch adds an overload for the map_to_vector helper template, exposing a parameter to control the size of the resulting SmallVector. A few call sites in mlir are updated to illustrate and test the change.

[laszlokindrat]

Making the formatter happy.

llvm/include/llvm/ADT/SmallVectorExtras.h

 #define LLVM_ADT_SMALLVECTOREXTRAS_H
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace llvm {
 
 /// Map a range to a SmallVector with element types deduced from the mapping.
+template <unsigned Size, class ContainerTy, class FuncTy>
+auto map_to_vector(ContainerTy &&C, FuncTy &&F) {
+  return to_vector<Size>(
+      map_range(std::forward<ContainerTy>(C), std::forward<FuncTy>(F)));
+}
 template <class ContainerTy, class FuncTy>
 auto map_to_vector(ContainerTy &&C, FuncTy &&F) {
   return to_vector(
       map_range(std::forward<ContainerTy>(C), std::forward<FuncTy>(F)));
 }
 
 } // namespace llvm
 
 #endif // LLVM_ADT_SMALLVECTOREXTRAS_H
 No newline at end of file

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/SmallBitVector.h"
+#include "llvm/ADT/SmallVectorExtras.h"
 #include <optional>
 
 namespace llvm {
 class SmallBitVector;
 } // namespace llvm
 
 namespace mlir {
 
public:
   AffineMap replace(const DenseMap<AffineExpr, AffineExpr> &map,
                     unsigned numResultDims, unsigned numResultSyms) const;
 
   /// Replace dims[offset ... numDims)
   /// by dims[offset + shift ... shift + numDims).
   AffineMap shiftDims(unsigned shift, unsigned offset = 0) const {
     assert(offset <= getNumDims());
     return AffineMap::get(getNumDims() + shift, getNumSymbols(),
-                          llvm::to_vector<4>(llvm::map_range(
+                          llvm::map_to_vector<4>(
                               getResults(),
                               [&](AffineExpr e) {
                                 return e.shiftDims(getNumDims(), shift, offset);
-                              })),
+                              }),
                           getContext());
   }
 
   /// Replace symbols[offset ... numSymbols)
   /// by symbols[offset + shift ... shift + numSymbols).
   AffineMap shiftSymbols(unsigned shift, unsigned offset = 0) const {
     return AffineMap::get(getNumDims(), getNumSymbols() + shift,
-                          llvm::to_vector<4>(llvm::map_range(
-                              getResults(),
+                          llvm::map_to_vector<4>(getResults(),
                                                  [&](AffineExpr e) {
-                                return e.shiftSymbols(getNumSymbols(), shift,
+                                                   return e.shiftSymbols(
+                                                       getNumSymbols(), shift,
                                                        offset);
-                              })),
+                                                 }),
                           getContext());
   }
 
   /// Returns a new AffineMap with the same number of dims and symbols and one
   /// less result at `pos`, dropped.
   AffineMap dropResult(int64_t pos) const {
     return dropResults(ArrayRef({pos}));
   }

mlir/lib/IR/Builders.cpp

 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/SymbolTable.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/ADT/SmallVectorExtras.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Locations.
 //===----------------------------------------------------------------------===//
 
 Location Builder::getUnknownLoc() { return UnknownLoc::get(context); }
StringAttr Builder::getStringAttr(const Twine &bytes) {
   return StringAttr::get(context, bytes);
 }
 
 ArrayAttr Builder::getArrayAttr(ArrayRef<Attribute> value) {
   return ArrayAttr::get(context, value);
 }
 
 ArrayAttr Builder::getBoolArrayAttr(ArrayRef<bool> values) {
-  auto attrs = llvm::to_vector<8>(llvm::map_range(
-      values, [this](bool v) -> Attribute { return getBoolAttr(v); }));
+  auto attrs = llvm::map_to_vector<8>(
+      values, [this](bool v) -> Attribute { return getBoolAttr(v); });
   return getArrayAttr(attrs);
 }
 
 ArrayAttr Builder::getI32ArrayAttr(ArrayRef<int32_t> values) {
-  auto attrs = llvm::to_vector<8>(llvm::map_range(
-      values, [this](int32_t v) -> Attribute { return getI32IntegerAttr(v); }));
+  auto attrs = llvm::map_to_vector<8>(
+      values, [this](int32_t v) -> Attribute { return getI32IntegerAttr(v); });
   return getArrayAttr(attrs);
 }
 ArrayAttr Builder::getI64ArrayAttr(ArrayRef<int64_t> values) {
-  auto attrs = llvm::to_vector<8>(llvm::map_range(
-      values, [this](int64_t v) -> Attribute { return getI64IntegerAttr(v); }));
+  auto attrs = llvm::map_to_vector<8>(
+      values, [this](int64_t v) -> Attribute { return getI64IntegerAttr(v); });
   return getArrayAttr(attrs);
 }
 
 ArrayAttr Builder::getIndexArrayAttr(ArrayRef<int64_t> values) {
-  auto attrs = llvm::to_vector<8>(
-      llvm::map_range(values, [this](int64_t v) -> Attribute {
-        return getIntegerAttr(IndexType::get(getContext()), v);
-      }));
+  auto attrs = llvm::map_to_vector<8>(values, [this](int64_t v) -> Attribute {
+  });
   return getArrayAttr(attrs);
 }
 
 ArrayAttr Builder::getF32ArrayAttr(ArrayRef<float> values) {
-  auto attrs = llvm::to_vector<8>(llvm::map_range(
-      values, [this](float v) -> Attribute { return getF32FloatAttr(v); }));
+  auto attrs = llvm::map_to_vector<8>(
+      values, [this](float v) -> Attribute { return getF32FloatAttr(v); });
   return getArrayAttr(attrs);
 }
 
 ArrayAttr Builder::getF64ArrayAttr(ArrayRef<double> values) {
-  auto attrs = llvm::to_vector<8>(llvm::map_range(
-      values, [this](double v) -> Attribute { return getF64FloatAttr(v); }));
+  auto attrs = llvm::map_to_vector<8>(
+      values, [this](double v) -> Attribute { return getF64FloatAttr(v); });
   return getArrayAttr(attrs);
 }
 
 ArrayAttr Builder::getStrArrayAttr(ArrayRef<StringRef> values) {
-  auto attrs = llvm::to_vector<8>(llvm::map_range(
-      values, [this](StringRef v) -> Attribute { return getStringAttr(v); }));
+  auto attrs = llvm::map_to_vector<8>(
+      values, [this](StringRef v) -> Attribute { return getStringAttr(v); });
   return getArrayAttr(attrs);
 }
 
 ArrayAttr Builder::getTypeArrayAttr(TypeRange values) {
-  auto attrs = llvm::to_vector<8>(llvm::map_range(
-      values, [](Type v) -> Attribute { return TypeAttr::get(v); }));
+  auto attrs = llvm::map_to_vector<8>(
+      values, [](Type v) -> Attribute { return TypeAttr::get(v); });
   return getArrayAttr(attrs);
 }
 
 ArrayAttr Builder::getAffineMapArrayAttr(ArrayRef<AffineMap> values) {
-  auto attrs = llvm::to_vector<8>(llvm::map_range(
-      values, [](AffineMap v) -> Attribute { return AffineMapAttr::get(v); }));
+  auto attrs = llvm::map_to_vector<8>(
+      values, [](AffineMap v) -> Attribute { return AffineMapAttr::get(v); });
   return getArrayAttr(attrs);
 }
 
 TypedAttr Builder::getZeroAttr(Type type) {
   if (llvm::isa<FloatType>(type))
     return getFloatAttr(type, 0.0);
   if (llvm::isa<IndexType>(type))
     return getIndexAttr(0);

mlir/lib/IR/TypeUtilities.cpp

 //===- TypeUtilities.cpp - Helper function for type queries ---------------===//
 //
 // 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 defines generic type utilities.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/TypeUtilities.h"
-
-#include <numeric>
-
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
+#include "llvm/ADT/SmallVectorExtras.h"
+#include <numeric>
 
 using namespace mlir;
 
 Type mlir::getElementTypeOrSelf(Type type) {
   if (auto st = llvm::dyn_cast<ShapedType>(type))
     return st.getElementType();
   return type;
 }
LogicalResult mlir::verifyCompatibleDims(ArrayRef<int64_t> dims) {
   }));
 }
 
 /// Returns success if all given types have compatible shapes. That is, they are
 /// all scalars (not shaped), or they are all shaped types and any ranked shapes
 /// have compatible dimensions. Dimensions are compatible if all non-dynamic
 /// dims are equal. The element type does not matter.
 LogicalResult mlir::verifyCompatibleShapes(TypeRange types) {
-  auto shapedTypes = llvm::to_vector<8>(llvm::map_range(
-      types, [](auto type) { return llvm::dyn_cast<ShapedType>(type); }));
+  auto shapedTypes = llvm::map_to_vector<8>(
+      types, [](auto type) { return llvm::dyn_cast<ShapedType>(type); });
   // Return failure if some, but not all are not shaped. Return early if none
   // are shaped also.
   if (llvm::none_of(shapedTypes, [](auto t) { return t; }))
     return success();
   if (!llvm::all_of(shapedTypes, [](auto t) { return t; }))
     return failure();
 
   // Return failure if some, but not all, are scalable vectors.
LogicalResult mlir::verifyCompatibleShapes(TypeRange types) {
   // All ranks should be equal
   auto firstRank = shapes.front().getRank();
   if (llvm::any_of(shapes,
                    [&](auto shape) { return firstRank != shape.getRank(); }))
     return failure();
 
   for (unsigned i = 0; i < firstRank; ++i) {
     // Retrieve all ranked dimensions
-    auto dims = llvm::to_vector<8>(llvm::map_range(
+    auto dims = llvm::map_to_vector<8>(
         llvm::make_filter_range(
             shapes, [&](auto shape) { return shape.getRank() >= i; }),
-        [&](auto shape) { return shape.getDimSize(i); }));
+        [&](auto shape) { return shape.getDimSize(i); });
     if (verifyCompatibleDims(dims).failed())
       return failure();
   }
 
   return success();
 }
 
 Type OperandElementTypeIterator::mapElement(Value value) const {
   return llvm::cast<ShapedType>(value.getType()).getElementType();
 }
 
 Type ResultElementTypeIterator::mapElement(Value value) const {
   return llvm::cast<ShapedType>(value.getType()).getElementType();
 }