[Mlir-commits] [mlir] 17faae9 - [ADT] Introduce `map_to_vector` helper

[Laszlo Kindrat]

Author: Laszlo Kindrat
Date: 2023-05-04T06:25:25-05:00
New Revision: 17faae95d799e5b4cf8344589641efcb1c740454

URL: https://github.com/llvm/llvm-project/commit/17faae95d799e5b4cf8344589641efcb1c740454
DIFF: https://github.com/llvm/llvm-project/commit/17faae95d799e5b4cf8344589641efcb1c740454.diff

LOG: [ADT] Introduce `map_to_vector` helper

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

Added: 
    llvm/include/llvm/ADT/SmallVectorExtras.h

Modified: 
    mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/SmallVectorExtras.h b/llvm/include/llvm/ADT/SmallVectorExtras.h
new file mode 100644
index 0000000000000..8d5228025e0e7
--- /dev/null
+++ b/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
\ No newline at end of file

diff  --git a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
index dceb83c249eae..a13acc691ae35 100644
--- a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
+++ b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
@@ -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;
@@ -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),

diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 537ed8745ba19..6c239cf18c371 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -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>
@@ -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,
@@ -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);
 }