[ADT] Introduce `map_to_vector` helper

[laszlokindrat]

Summary

The following pattern is common in the llvm codebase, as well as in downstream projects:

llvm::to_vector(llvm::map_range(container, lambda))

This patch introduces a shortcut for this called map_to_vector.

This template depends on both llvm/ADT/SmallVector.h and llvm/ADT/STLExtras.h, and since these are both relatively large and do not depend on each other, the map_to_vector helper is placed in a new header under llvm/ADT/SmallVectorExtras.h. I am happy to change this if reviewers can suggest a better location.

Only a handful of use cases have been updated to use the new helper, but I'm happy to update across the codebase (either in this patch, or in subsequent diffs) if this patch is accepted.

[bzcheeseman]

Nice, thank you.

[kuhar]

llvm/include/llvm/ADT/SmallVectorExtras.h

23–24

I'd expect this to optionally accept small size, just like to_vector<N>(...) does. Is there a specific reason why only the default size is allowed?

Also, do you plan on adding map_to_vector_of? IIRC to_vector_of was required for some common mlir functions like DenseElementAttribute::getValues<T>()

[laszlokindrat]

llvm/include/llvm/ADT/SmallVectorExtras.h

23–24

I'd expect this to optionally accept small size, just like to_vector<N>(...) does. Is there a specific reason why only the default size is allowed?

No reason other than I prefer to add complexity incrementally. But I agree that this would be nice, so I will add it in a subsequent patch.

Also, do you plan on adding map_to_vector_of?

I can't really find any use cases for this in the llvm monorepo. Do you have a downstream project that could benefit from it? I usually like to add this kind of sugar if/when there are lots of places that can benefit from it.

[kuhar]

llvm/include/llvm/ADT/SmallVectorExtras.h

23–24

I can't really find any use cases for this in the llvm monorepo. Do you have a downstream project that could benefit from it? I usually like to add this kind of sugar if/when there are lots of places that can benefit from it.

I think I wanted to use to_vector(map(...)) here https://github.com/llvm/llvm-project/blob/890aa28f1c9a502ad012d83ad4d6ad60f75efccb/mlir/lib/Dialect/Vector/IR/VectorOps.cpp#L2671C10-L2687 and https://github.com/llvm/llvm-project/blob/890aa28f1c9a502ad012d83ad4d6ad60f75efccb/mlir/lib/Dialect/Arith/Transforms/IntNarrowing.cpp#L201-L204 (in reduce or accumulate) but couldn't at the time. I think this can be a bit of a chicken and egg problem -- you may not be able to find many places that would immediately benefit form it because folks have already worked around it.

[laszlokindrat]

llvm/include/llvm/ADT/SmallVectorExtras.h

23–24

See https://reviews.llvm.org/D150601 for a patch that adds the ability to specify the size of the resulting vector.

[dblaikie]

(my 2c would be that I'm not sure the wrapper pulls its weight - writing the two functions calls out doesn't seem too bad, but no big deal either way)

llvm/include/llvm/ADT/SmallVectorExtras.h

+//===- llvm/ADT/SmallVectorExtras.h -----------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file defines less commonly used SmallVector utilities.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef 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 <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/lib/Conversion/GPUCommon/GPUOpsLowering.cpp

 //===- GPUOpsLowering.cpp - GPU FuncOp / ReturnOp lowering ----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "GPUOpsLowering.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
-#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 LogicalResult
 GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,
                                    ConversionPatternRewriter &rewriter) const {
   Location loc = gpuFuncOp.getLoc();
LogicalResult impl::scalarizeVectorOp(Operation *op, ValueRange operands,
 
   for (int64_t i = 0; i < vectorType.getNumElements(); ++i) {
     Value index = rewriter.create<LLVM::ConstantOp>(loc, indexType, i);
     auto extractElement = [&](Value operand) -> Value {
       if (!operand.getType().isa<VectorType>())
         return operand;
       return rewriter.create<LLVM::ExtractElementOp>(loc, operand, index);
     };
-    auto scalarOperands =
-        llvm::to_vector(llvm::map_range(operands, extractElement));
+    auto scalarOperands = llvm::map_to_vector(operands, extractElement);
     Operation *scalarOp =
         rewriter.create(loc, name, scalarOperands, elementType, op->getAttrs());
     rewriter.create<LLVM::InsertElementOp>(loc, result, scalarOp->getResult(0),
                                            index);
   }
 
   rewriter.replaceOp(op, result);
   return success();

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ShapedOpInterfaces.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallBitVector.h"
+#include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include <numeric>
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::affine;
 
return makeComposedFoldedAffineApply(
       b, loc, AffineMap::inferFromExprList(ArrayRef<AffineExpr>{expr}).front(),
       operands);
 }
 
 SmallVector<OpFoldResult>
 mlir::affine::makeComposedFoldedMultiResultAffineApply(
     OpBuilder &b, Location loc, AffineMap map,
     ArrayRef<OpFoldResult> operands) {
-  return llvm::to_vector(llvm::map_range(
-      llvm::seq<unsigned>(0, map.getNumResults()), [&](unsigned i) {
-        return makeComposedFoldedAffineApply(b, loc, map.getSubMap({i}),
-                                             operands);
-      }));
+  return llvm::map_to_vector(llvm::seq<unsigned>(0, map.getNumResults()),
+                             [&](unsigned i) {
+                               return makeComposedFoldedAffineApply(
+                                   b, loc, map.getSubMap({i}), operands);
+                             });
 }
 
 Value mlir::affine::makeComposedAffineMin(OpBuilder &b, Location loc,
                                           AffineMap map, ValueRange operands) {
   SmallVector<Value> allOperands = llvm::to_vector(operands);
   composeMultiResultAffineMap(map, allOperands);
   return b.createOrFold<AffineMinOp>(loc, b.getIndexType(), map, allOperands);
 }

 //===----------------------------------------------------------------------===//
 
 void AffineDelinearizeIndexOp::build(OpBuilder &builder, OperationState &result,
                                      Value linearIndex,
                                      ArrayRef<OpFoldResult> basis) {
   result.addTypes(SmallVector<Type>(basis.size(), builder.getIndexType()));
   result.addOperands(linearIndex);
   SmallVector<Value> basisValues =
-      llvm::to_vector(llvm::map_range(basis, [&](OpFoldResult ofr) -> Value {
+      llvm::map_to_vector(basis, [&](OpFoldResult ofr) -> Value {
         std::optional<int64_t> staticDim = getConstantIntValue(ofr);
         if (staticDim.has_value())
           return builder.create<arith::ConstantIndexOp>(result.location,
                                                         *staticDim);
         return ofr.dyn_cast<Value>();
-      }));
+      });
   result.addOperands(basisValues);
 }
 
 LogicalResult AffineDelinearizeIndexOp::verify() {
   if (getBasis().empty())
     return emitOpError("basis should not be empty");
   if (getNumResults() != getBasis().size())
     return emitOpError("should return an index for each basis element");