[MLIR] Improve interaction of TypedValue with BlockAndValueMapping

[Tyker]

Summary

when BlockAndValueMapping::lookup is called with a TypedValue it will assert that the type of the return is correct. which might not always be true.
with this patch calls to BlockAndValueMapping::lookup will cast the argument to Value or Block*.

[rriddle]

mlir/include/mlir/IR/BlockAndValueMapping.h

81

mlir/unittests/IR/BlockAndValueMapping.cpp

18

Please drop all of the mlir::

[mscuttari]

Thanks for addressing the issue

[Tyker]

pushed as fa40fd40e055630c34890de84eb048796133fe03

mlir/include/mlir/IR/BlockAndValueMapping.h

public:
   /// does not exist then return the provided value.
   Block *lookupOrDefault(Block *from) const {
     return lookupOrValue(from, from);
   }
   Value lookupOrDefault(Value from) const { return lookupOrValue(from, from); }
 
   /// Lookup a mapped value within the map. This asserts the provided value
   /// exists within the map.
-  template <typename T>
-  T lookup(T from) const {
-    auto result = lookupOrNull(from);
-    assert(result && "expected 'from' to be contained within the map");
-    return result;
-  }
+  Block *lookup(Block *from) const { return lookupImpl(from); }
+  Value lookup(Value from) const { return lookupImpl(from); }
 
   /// Clears all mappings held by the mapper.
   void clear() { valueMap.clear(); }
 
   /// Return the held value mapping.
   const DenseMap<Value, Value> &getValueMap() const { return valueMap; }
 
   /// Return the held block mapping.
   const DenseMap<Block *, Block *> &getBlockMap() const { return blockMap; }
 
 private:
+  template<typename T>
+  T lookupImpl(T from) const {
+    auto result = lookupOrNull(from);
+    assert(result && "expected 'from' to be contained within the map");
+    return result;
+  }
+
   /// Utility lookupOrValue that looks up an existing key or returns the
   /// provided value.
   Block *lookupOrValue(Block *from, Block *value) const {
     auto it = blockMap.find(from);
     return it != blockMap.end() ? it->second : value;
   }
   Value lookupOrValue(Value from, Value value) const {
     auto it = valueMap.find(from);

mlir/unittests/IR/BlockAndValueMapping.cpp

+//===- BlockAndValueMapping.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Builders.h"
+#include "gtest/gtest.h"
+
+#include "../../test/lib/Dialect/Test/TestDialect.h"
+
+using namespace mlir;
+
+TEST(BlockAndValueMapping, TypedValue) {
+  mlir::MLIRContext context;
+
+  context.loadDialect<test::TestDialect>();
+
+  mlir::OpBuilder builder(&context);
+  mlir::Location loc = builder.getUnknownLoc();
+
+  mlir::Value i64Val = builder.create<test::TestOpConstant>(
+      loc, builder.getI64Type(), builder.getI64IntegerAttr(0));
+  mlir::Value f64Val = builder.create<test::TestOpConstant>(
+      loc, builder.getF64Type(), builder.getF64FloatAttr(0.0));
+
+  mlir::BlockAndValueMapping mapping;
+  mapping.map(i64Val, f64Val);
+  TypedValue<mlir::IntegerType> typedI64Val = i64Val;
+  mapping.lookup(typedI64Val);
+}

mlir/unittests/IR/CMakeLists.txt

 add_mlir_unittest(MLIRIRTests
   AttributeTest.cpp
+  BlockAndValueMapping.cpp
   DialectTest.cpp
   InterfaceTest.cpp
   InterfaceAttachmentTest.cpp
   OperationSupportTest.cpp
   PatternMatchTest.cpp
   ShapedTypeTest.cpp
   SubElementInterfaceTest.cpp
   TypeTest.cpp