[PATCH] D73894: [MLIR] Add mapping based on ValueRange to BlockAndValueMapper.

[Stephan Herhut via Phabricator via llvm-commits]

herhut updated this revision to Diff 242292.
herhut added a comment.

Remove unneeded import.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D73894/new/

https://reviews.llvm.org/D73894

Files:
  mlir/include/mlir/IR/BlockAndValueMapping.h
  mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp


Index: mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp
===================================================================
--- mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp
+++ mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp
@@ -273,13 +273,11 @@
     // 2. Inline region, currently only works for a single basic block.
     BlockAndValueMapping map;
     auto &block = genericOp.region().front();
-    for (auto it : llvm::zip(block.getArguments(), indexedValues))
-      map.map(std::get<0>(it), std::get<1>(it));
+    map.map(block.getArguments(), indexedValues);
     for (auto &op : block.without_terminator()) {
       assert(op.getNumRegions() == 0);
       auto *newOp = b.clone(op, map);
-      for (auto it : llvm::zip(op.getResults(), newOp->getResults()))
-        map.map(std::get<0>(it), std::get<1>(it));
+      map.map(op.getResults(), newOp->getResults());
     }
 
     // 3. Emit std_store.
@@ -377,13 +375,11 @@
     // 2. Inline region, currently only works for a single basic block.
     BlockAndValueMapping map;
     auto &block = indexedGenericOp.region().front();
-    for (auto it : llvm::zip(block.getArguments(), indexedValues))
-      map.map(std::get<0>(it), std::get<1>(it));
+    map.map(block.getArguments(), indexedValues);
     for (auto &op : block.without_terminator()) {
       assert(op.getNumRegions() == 0);
       auto *newOp = b.clone(op, map);
-      for (auto it : llvm::zip(op.getResults(), newOp->getResults()))
-        map.map(std::get<0>(it), std::get<1>(it));
+      map.map(op.getResults(), newOp->getResults());
     }
 
     // 3. Emit std_store.
Index: mlir/include/mlir/IR/BlockAndValueMapping.h
===================================================================
--- mlir/include/mlir/IR/BlockAndValueMapping.h
+++ mlir/include/mlir/IR/BlockAndValueMapping.h
@@ -27,9 +27,26 @@
 public:
   /// Inserts a new mapping for 'from' to 'to'. If there is an existing mapping,
   /// it is overwritten.
-  void map(Block *from, Block *to) { valueMap[from] = to; }
-  void map(Value from, Value to) {
-    valueMap[from.getAsOpaquePointer()] = to.getAsOpaquePointer();
+  template <typename S, typename T,
+            std::enable_if_t<std::is_assignable<Value, S>::value &&
+                             std::is_assignable<Value, T>::value> * = nullptr>
+  void map(S from, T to) {
+    valueMap[Value{from}.getAsOpaquePointer()] = Value{to}.getAsOpaquePointer();
+  }
+
+  template <typename S, typename T,
+            std::enable_if_t<std::is_same<S, Block *>::value &&
+                             std::is_same<T, Block *>::value> * = nullptr>
+  void map(S from, T to) {
+    valueMap[from] = to;
+  }
+
+  template <typename S, typename T,
+            std::enable_if_t<!std::is_assignable<Value, S>::value &&
+                             !std::is_same<S, Block *>::value> * = nullptr>
+  void map(S from, T to) {
+    for (auto pair : llvm::zip(from, to))
+      map(std::get<0>(pair), std::get<1>(pair));
   }
 
   /// Erases a mapping for 'from'.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D73894.242292.patch
Type: text/x-patch
Size: 3025 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200204/6f5ea4a7/attachment.bin>