[Mlir-commits] [mlir] fa40fd4 - [MLIR] Improve interaction of TypedValue with BlockAndValueMapping

[llvmlistbot at llvm.org]

Author: Tyker
Date: 2022-09-09T08:55:05-07:00
New Revision: fa40fd40e055630c34890de84eb048796133fe03

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

LOG: [MLIR] Improve interaction of TypedValue with BlockAndValueMapping

Added: 
    mlir/unittests/IR/BlockAndValueMapping.cpp

Modified: 
    mlir/include/mlir/IR/BlockAndValueMapping.h
    mlir/unittests/IR/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/BlockAndValueMapping.h b/mlir/include/mlir/IR/BlockAndValueMapping.h
index 8134902bfd646..9776030282202 100644
--- a/mlir/include/mlir/IR/BlockAndValueMapping.h
+++ b/mlir/include/mlir/IR/BlockAndValueMapping.h
@@ -63,12 +63,8 @@ class BlockAndValueMapping {
 
   /// 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(); }
@@ -80,6 +76,13 @@ class BlockAndValueMapping {
   const DenseMap<Block *, Block *> &getBlockMap() const { return blockMap; }
 
 private:
+  template<typename T>
+  T lookupImpl(T from) const {
+    T 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 {

diff  --git a/mlir/unittests/IR/BlockAndValueMapping.cpp b/mlir/unittests/IR/BlockAndValueMapping.cpp
new file mode 100644
index 0000000000000..dad3c915c5b96
--- /dev/null
+++ b/mlir/unittests/IR/BlockAndValueMapping.cpp
@@ -0,0 +1,34 @@
+//===- 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) {
+  MLIRContext context;
+
+  context.loadDialect<test::TestDialect>();
+
+  OpBuilder builder(&context);
+  Location loc = builder.getUnknownLoc();
+
+  Value i64Val = builder.create<test::TestOpConstant>(
+      loc, builder.getI64Type(), builder.getI64IntegerAttr(0));
+  Value f64Val = builder.create<test::TestOpConstant>(
+      loc, builder.getF64Type(), builder.getF64FloatAttr(0.0));
+
+  BlockAndValueMapping mapping;
+  mapping.map(i64Val, f64Val);
+  TypedValue<IntegerType> typedI64Val = i64Val;
+  mapping.lookup(typedI64Val);
+}

diff  --git a/mlir/unittests/IR/CMakeLists.txt b/mlir/unittests/IR/CMakeLists.txt
index 51978aea6e880..904ea13447bb0 100644
--- a/mlir/unittests/IR/CMakeLists.txt
+++ b/mlir/unittests/IR/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_mlir_unittest(MLIRIRTests
   AttributeTest.cpp
+  BlockAndValueMapping.cpp
   DialectTest.cpp
   InterfaceTest.cpp
   InterfaceAttachmentTest.cpp