Skip to content

Commit

Permalink
[ADT] Introduce map_to_vector helper
Browse files Browse the repository at this point in the history
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`. Only a handful of use cases have been updated to use the new helper.

Differential Revision: https://reviews.llvm.org/D145390
  • Loading branch information
Laszlo Kindrat committed May 4, 2023
1 parent 7c57195 commit 17faae9
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 10 deletions.
31 changes: 31 additions & 0 deletions llvm/include/llvm/ADT/SmallVectorExtras.h
@@ -0,0 +1,31 @@
//===- 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
5 changes: 2 additions & 3 deletions mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
Expand Up @@ -11,7 +11,7 @@
#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;
Expand Down Expand Up @@ -486,8 +486,7 @@ LogicalResult impl::scalarizeVectorOp(Operation *op, ValueRange operands,
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),
Expand Down
15 changes: 8 additions & 7 deletions mlir/lib/Dialect/Affine/IR/AffineOps.cpp
Expand Up @@ -20,6 +20,7 @@
#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>
Expand Down Expand Up @@ -1360,11 +1361,11 @@ 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,
Expand Down Expand Up @@ -4592,13 +4593,13 @@ void AffineDelinearizeIndexOp::build(OpBuilder &builder, OperationState &result,
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);
}

Expand Down

0 comments on commit 17faae9

Please sign in to comment.