[MLIR] Add mapping based on ValueRange to BlockAndValueMapper.

[herhut]

Summary

It is often needed to map entire ranges rather than single values. To avoid
writing the same for loop every time, I have added an overload to the map
method.

[herhut]

Maybe this should even be a template to cover more cases like BlockArguments.

[merge_guards_bot]

Unit tests: pass. 62416 tests passed, 0 failed and 839 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/include/mlir/IR/BlockAndValueMapping.h

34

I would prefer that we use templates instead of specific classes. That would also remove the need for any additional includes, among other benefits.

[herhut]

Made it a template with help from @pifon2a.

[herhut]

Remove unneeded import.

[merge_guards_bot]

Unit tests: fail. 62437 tests passed, 1 failed and 845 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_timedmutex_requirements/thread_timedmutex_class/lock.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: pass. 62438 tests passed, 0 failed and 845 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/include/mlir/IR/BlockAndValueMapping.h

32

nit: I would just keep these the same for now. Is there a reason to templatize these yet?

49

nit: S&& and T&&

[herhut]

Less templates.

[herhut]

@rriddle Is this what you had in mind?

mlir/include/mlir/IR/BlockAndValueMapping.h

32

No, this was more me figuring out specifying the right template constraints. I have reverted that part.

[merge_guards_bot]

Unit tests: pass. 62438 tests passed, 0 failed and 845 were skipped.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

LGTM after fixing the format issue.

mlir/include/mlir/IR/BlockAndValueMapping.h

37

nit: Can you just use is_assignable for both?

[herhut]

Formatting.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[herhut]

use is_assignable

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/include/mlir/IR/BlockAndValueMapping.h

 // 'lookupOrNull' returns nullptr where as 'lookupOrDefault' will return the
 // lookup key.
 class BlockAndValueMapping {
 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<Block *, S>::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'.
   void erase(Block *from) { valueMap.erase(from); }
   void erase(Value from) { valueMap.erase(from.getAsOpaquePointer()); }
 
   /// Checks to see if a mapping for 'from' exists.
   bool contains(Block *from) const { return valueMap.count(from); }
   bool contains(Value from) const {
     return valueMap.count(from.getAsOpaquePointer());
   }
 
   /// Lookup a mapped value within the map. If a mapping for the provided value
   /// does not exist then return nullptr.
   Block *lookupOrNull(Block *from) const {
     return lookupOrValue(from, (Block *)nullptr);

mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp

if (funcOp) {
         std_store(callOp->getResult(i), genericOp.getOutputBuffer(i), indexing);
       }
       return;
     }
     // TODO(ntv): When a region inliner exists, use it.
     // 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.
     auto *yieldOp = cast<YieldOp>(block.back()).getOperation();
     assert(yieldOp->getNumOperands() == nOutputs);
     for (unsigned i = 0; i < nOutputs; ++i) {
       ValueHandleArray indexing(makeCanonicalAffineApplies(
           b, loc, genericOp.getOutputIndexingMap(i), allIvs));
if (auto funcOp = indexedGenericOp.getFunction()) {
                   indexing);
       }
       return;
     }
     // TODO(ntv): When a region inliner exists, use it.
     // 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.
     auto *yieldOp = cast<YieldOp>(block.back()).getOperation();
     assert(yieldOp->getNumOperands() == nOutputs);
     for (unsigned i = 0; i < nOutputs; ++i) {
       ValueHandleArray indexing(makeCanonicalAffineApplies(
           b, loc, indexedGenericOp.getOutputIndexingMap(i), allIvs));