[mlir] Define proper DenseMapInfo for Interfaces

[zero9178]

Summary

Prior to this patch, using any kind of interface (op interface, attr interface, type interface) as the key of a llvm::DenseSet, llvm::DenseMap or other related containers, leads to invalid pointer dereferences, despite compiling.

The gist of the problem is that a llvm::DenseMapInfo specialization for the base type (aka one of Operation*, Type or Attribute) are selected when using an interface as a key, which uses getFromOpaquePointer with invalid pointer addresses to construct instances for the empty key and tombstone key values. The interface is then constructed with this invalid base type and an attempt is made to lookup the implementation in the interface map, which then dereferences the invalid pointer address. (For more details and the exact call chain involved see the GitHub issue below)

The current workaround is to use the more generic base type (eg. instead of DenseSet<FunctionOpInterface>, DenseSet<Operation*>), but this is strictly worse from a code perspective (doesn't enforce the invariant, code is less self documenting, having to insert casts etc).

This patch fixes that issue by defining a DenseMapInfo specialization of Interface subclasses which uses a new constructor to construct an instance without querying a concept. That allows getEmptyKey and getTombstoneKey to construct an interface with invalid pointer values.

Fixes https://github.com/llvm/llvm-project/issues/54908

---

Other bug fix approaches tried and tested that do not work:

• Specializing getFromOpaquePointer to mlir::detail::Interface: Doesn't work because the implementation can't differentiate between valid and invalid pointers and always skipping looking up the concept implementation would be incorrect.
• Specializing llvm::DenseMapInfo for mlir::detail::Interface subclasses: Doesn't work as the specialization would be ambiguous with existing specializations for OpState, Type and Attribute specializations. This is due to them only differing in the second std::enable_if_t parameter, which do not yield mutually exclusive results.

[zero9178]

Fix examples build

[jpienaar]

Thanks for digging into this one

mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h

20 ↗

(On Diff #441879)

What does the error message look like when one 1) doesn't set the dialect and/or 2) has it in namespace? My guess is not too intuitive. I wonder if we could improve that (beyond just sending out email and giving folks heads-up)

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

434 ↗

(On Diff #441879)

I wonder if we should make this a plain macro C++ side and just emit usage of it here. That way others can use when not using ODS (like the typeid construct). Also a bit easier to find if someone were to change DenseMapInfo at some point.

440 ↗

(On Diff #441879)

No space before to match template above?

[zero9178]

mlir/include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h

20 ↗

(On Diff #441879)

Not setting the cppNamespace is not problematic at all and would still work without errors.
The problem is that the llvm::DenseMapInfo would land within the namespace around the include and as you guessed they're not the most intuitive of errors. In the current state of the patch it leads to "no template named 'DenseMapInfo'; did you mean 'llvm::DenseMapInfo'?", which is very confusing for anyone just wanting to use interfaces. One could change the code a bit for better error messages but the gist of it not finding DenseMapInfo at the right place remains.

One could probably come up with a trick of a static_assert inserted into the TableGen generated code that only triggers if NOT included in the global namespace by abusing ADL a bit. Will have to think about how hacky such code would be :)

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

434 ↗

(On Diff #441879)

I actually like this idea a lot and it should work with minimal modifications. Let me give it a shot.

440 ↗

(On Diff #441879)

I am not 100% sure what you mean. It is indented as it is within the scope of the llvm::DenseMapInfo class. The template meanwhile is at namespace scope which we don't indent for. As precedent for this formatting I essentially copied the formatting from EnumsGen: https://github.com/llvm/llvm-project/blob/5c3b20520b5716d61833c8ce45d19faa48ce8db7/mlir/tools/mlir-tblgen/EnumsGen.cpp#L77

[zero9178]

• Use a macro to define the llvm::DefineMapInfo specialization. That way, handwritten interfaces can easily make use of it as well.
• Add some special code that checks that the interface declarations are being included in the global namespace.

Pros of this approach:

• Just a tiny bit of extra code
• A proper static_assert which can output a custom error that is as detailed as we want it to be
  
  Cons:
• The code is odd imho.
• Polluting the global namespace with MLIR_TEST_GLOBAL_NAMESPACE_INCLUDE (albeit it's named like a macro, hence not as big of a deal)

Demo of the error situation: https://godbolt.org/z/oWrToPvd5

[rriddle]

mlir/include/mlir/Support/InterfaceSupport.h

290

A lot of this complexity should not be necessary. DenseMapInfo supports SFINAE, so we should really only need a partial template specialization that checks if T is an interface.

[zero9178]

Change approach to using SFINAE to pick the correct template specialization. This sadly requires an intrusive change of the DenseMapInfo specializations of OpState, Type and Attribute to avoid an ambiguity of template specializations.

I have also chosen to make the Interface constructor with the NULL concept public for now, so that any subclasses of Interface work seamlessly with the template specialization without having to mark it as friend. Let me know if you prefer otherwise

[Mogball]

mlir/include/mlir/Support/InterfaceSupport.h

288

The intrusive change to DenseMapInfo for attribute, type, and op state shouldn't be necessary.

[zero9178]

mlir/include/mlir/Support/InterfaceSupport.h

288

Unless there is some other trick I am not aware of, I don't know how.

In an ideal world the specialization using mlir::detail::IsInterface<T> would be the most specialized implementation but that is sadly not how C++ ranks partial specialization as it does not consider SFINAE, but just the substituted types (I think C++20 concepts might be able to do this).

Instead when DenseMap & co. instantiates DenseMap<MemoryEffectOpInterface> it looks for any specializations of type DenseMapInfo<MemoryEffectOpInterface, void>. The problem now is that both the DenseMapInfo for OpState, as well as the one written here for any interfaces, are considered, as after substitution, both resolve to DenseMapInfo<MemoryEffectOpInterface, void> (since all OpInterface subclasses are also subclasses of OpState, and std::enable_if_t resolves to void).

Since they are both DenseMapInfo<MemoryEffectOpInterface, void>, none of them is strictly better/more specialized than the other and we get a compiler error:

llvm/ADT/DenseMap.h:458:12: error: ambiguous partial specializations of 'DenseMapInfo<mlir::MemoryEffectOpInterface, void>'
(...)
mlir/Support/InterfaceSupport.h:287:14: note: partial specialization matches [with T = mlir::MemoryEffectOpInterface]
struct llvm::DenseMapInfo<
             ^
mlir/IR/OpDefinition.h:1966:8: note: partial specialization matches [with T = mlir::MemoryEffectOpInterface]
struct DenseMapInfo<T,
       ^

Changing the type std::enable_if_t substitutes to also doesn't work because that just makes it not ever be selected, since DenseMap specifically requests a DenseMapInfo<mlir::MemoryEffectOpInterface, void> and not a DenseMapInfo<mlir::MemoryEffectOpInterface, mlir::MemoryEffectOpInterface>.

Hence this patch in the current state makes it so that both specializations are mutually exclusive to resolve the ambiguity.

[Mogball]

mlir/include/mlir/Support/InterfaceSupport.h

288

Why are the changes to the type and attribute DenseMapInfo necessary then?

[Mogball]

mlir/include/mlir/Support/InterfaceSupport.h

288

Ah, right. Type and attribute interfaces.

[zero9178]

mlir/include/mlir/Support/InterfaceSupport.h

288

Exactly the same reasons, just for Type and Attribute Interfaces, since they work exactly the same, just instead of OpState subclasses, it's Type and Attribute subclasses.
The Interface specialization works for all Type, Attribute and OpInterfaces.

I should probably add a test for them as well actually.

[zero9178]

Add tests for attribute and type interfaces as well.

[Mogball]

mlir/include/mlir/Support/InterfaceSupport.h

277

And below as well.

mlir/unittests/IR/InterfaceTest.cpp

12–18

46

Can you split up the op/type/attr interfaces tests?

[zero9178]

Address review comments

[rriddle]

mlir/include/mlir/Support/InterfaceSupport.h

287

Please separate out the namespace definition from the struct.

[zero9178]

Address review comments: Explicitly put DenseMapInfo specialization into the llvm namespace

[Mogball]

mlir/include/mlir/Support/InterfaceSupport.h

305

[zero9178]

mlir/include/mlir/Support/InterfaceSupport.h

305

This seems a bit odd to me and does not match the style in the rest of the file as well. Quick check:

C:\llvm-project\mlir>git grep "// namespace" | wc -l 
2160

C:\llvm-project\mlir>git grep "// end namespace" | wc -l 
29

seems to indicate the former being more common, so I pushed it without this suggestion for now.
Still open to fixing it post-commit of course.

[Mogball]

mlir/include/mlir/Support/InterfaceSupport.h

305

I see both and to be honest I don't know which one is preferred lol

[mehdi_amini]

mlir/include/mlir/Support/InterfaceSupport.h

305

The examples in the coding standard are the former form, and as mentioned above, a grep in the MLIR codebase shows 100 times more of it as well.

mlir/include/mlir/IR/Attributes.h

static unsigned getHashValue(mlir::Attribute val) {
     return mlir::hash_value(val);
   }
   static bool isEqual(mlir::Attribute LHS, mlir::Attribute RHS) {
     return LHS == RHS;
   }
 };
 template <typename T>
 struct DenseMapInfo<
-    T, std::enable_if_t<std::is_base_of<mlir::Attribute, T>::value>>
+    T, std::enable_if_t<std::is_base_of<mlir::Attribute, T>::value &&
+                        !mlir::detail::IsInterface<T>::value>>
     : public DenseMapInfo<mlir::Attribute> {
   static T getEmptyKey() {
     const void *pointer = llvm::DenseMapInfo<const void *>::getEmptyKey();
     return T::getFromOpaquePointer(pointer);
   }
   static T getTombstoneKey() {
     const void *pointer = llvm::DenseMapInfo<const void *>::getTombstoneKey();
     return T::getFromOpaquePointer(pointer);

mlir/include/mlir/IR/OpDefinition.h

 LogicalResult verifyCastInterfaceOp(
     Operation *op, function_ref<bool(TypeRange, TypeRange)> areCastCompatible);
 } // namespace impl
 } // namespace mlir
 
 namespace llvm {
 
 template <typename T>
-struct DenseMapInfo<
-    T, std::enable_if_t<std::is_base_of<mlir::OpState, T>::value>> {
+struct DenseMapInfo<T,
+                    std::enable_if_t<std::is_base_of<mlir::OpState, T>::value &&
+                                     !mlir::detail::IsInterface<T>::value>> {
   static inline T getEmptyKey() {
     auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
     return T::getFromOpaquePointer(pointer);
   }
   static inline T getTombstoneKey() {
     auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return T::getFromOpaquePointer(pointer);
   }

mlir/include/mlir/IR/Types.h

template <> struct DenseMapInfo<mlir::Type> {
   static mlir::Type getTombstoneKey() {
     auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return mlir::Type(static_cast<mlir::Type::ImplType *>(pointer));
   }
   static unsigned getHashValue(mlir::Type val) { return mlir::hash_value(val); }
   static bool isEqual(mlir::Type LHS, mlir::Type RHS) { return LHS == RHS; }
 };
 template <typename T>
-struct DenseMapInfo<T, std::enable_if_t<std::is_base_of<mlir::Type, T>::value>>
+struct DenseMapInfo<T, std::enable_if_t<std::is_base_of<mlir::Type, T>::value &&
+                                        !mlir::detail::IsInterface<T>::value>>
     : public DenseMapInfo<mlir::Type> {
   static T getEmptyKey() {
     const void *pointer = llvm::DenseMapInfo<const void *>::getEmptyKey();
     return T::getFromOpaquePointer(pointer);
   }
   static T getTombstoneKey() {
     const void *pointer = llvm::DenseMapInfo<const void *>::getTombstoneKey();
     return T::getFromOpaquePointer(pointer);

mlir/include/mlir/Support/InterfaceSupport.h

public:
   using Concept = typename Traits::Concept;
   template <typename T> using Model = typename Traits::template Model<T>;
   template <typename T>
   using FallbackModel = typename Traits::template FallbackModel<T>;
   using InterfaceBase =
       Interface<ConcreteType, ValueT, Traits, BaseType, BaseTrait>;
   template <typename T, typename U>
   using ExternalModel = typename Traits::template ExternalModel<T, U>;
+  using ValueType = ValueT;
 
   /// This is a special trait that registers a given interface with an object.
   template <typename ConcreteT>
   struct Trait : public BaseTrait<ConcreteT, Trait> {
     using ModelT = Model<ConcreteT>;
 
     /// Define an accessor for the ID of this interface.
     static TypeID getInterfaceID() { return TypeID::get<ConcreteType>(); }
public:
   /// interface's trait.
   template <typename T, typename std::enable_if_t<
                             std::is_base_of<Trait<T>, T>::value> * = nullptr>
   Interface(T t)
       : BaseType(t), impl(t ? ConcreteType::getInterfaceFor(t) : nullptr) {
     assert((!t || impl) && "expected value to provide interface instance");
   }
 
+  /// Constructor for DenseMapInfo's empty key and tombstone key.
+  Interface(ValueT t, std::nullptr_t) : BaseType(t), impl(nullptr) {}
+
   /// Support 'classof' by checking if the given object defines the concrete
   /// interface.
   static bool classof(ValueT t) { return ConcreteType::getInterfaceFor(t); }
 
   /// Define an accessor for the ID of this interface.
   static TypeID getInterfaceID() { return TypeID::get<ConcreteType>(); }
 
 protected:
const auto *it =
         });
     return (it != interfaces.end() && it->first == id) ? it->second : nullptr;
   }
 
   /// A list of interface instances, sorted by TypeID.
   SmallVector<std::pair<TypeID, void *>> interfaces;
 };
 
+template <typename ConcreteType, typename ValueT, typename Traits,
+          typename BaseType,
+          template <typename, template <typename> class> class BaseTrait>
+void isInterfaceImpl(
+    Interface<ConcreteType, ValueT, Traits, BaseType, BaseTrait> &);
+
+template <class T>
+using is_interface_t = decltype(isInterfaceImpl(std::declval<T&>()));
+
+template <class T>
+using IsInterface = llvm::is_detected<is_interface_t, T>;
+
 } // namespace detail
 } // namespace mlir
 
+template <class T>
+struct llvm::DenseMapInfo<
+    T, std::enable_if_t<mlir::detail::IsInterface<T>::value>> {
+  using ValueTypeInfo = llvm::DenseMapInfo<typename T::ValueType>;
+
+  static T getEmptyKey() { return T(ValueTypeInfo::getEmptyKey(), nullptr); }
+
+  static T getTombstoneKey() {
+    return T(ValueTypeInfo::getTombstoneKey(), nullptr);
+  }
+
+  static unsigned getHashValue(T val) {
+    return ValueTypeInfo::getHashValue(val);
+  }
+
+  static bool isEqual(T lhs, T rhs) { return ValueTypeInfo::isEqual(lhs, rhs); }
+};
+
 #endif

mlir/unittests/IR/CMakeLists.txt

 add_mlir_unittest(MLIRIRTests
   AttributeTest.cpp
   DialectTest.cpp
+  InterfaceTest.cpp
   InterfaceAttachmentTest.cpp
   OperationSupportTest.cpp
   PatternMatchTest.cpp
   ShapedTypeTest.cpp
   SubElementInterfaceTest.cpp
 
   DEPENDS
   MLIRTestInterfaceIncGen
 )
 target_include_directories(MLIRIRTests PRIVATE "${MLIR_BINARY_DIR}/test/lib/Dialect/Test")
 target_link_libraries(MLIRIRTests
   PRIVATE
   MLIRIR
   MLIRTestDialect)

mlir/unittests/IR/InterfaceTest.cpp

+//===- InterfaceTest.cpp - Test interfaces --------------------------------===//
 //
-// Created by Markus Böck on 02/07/2022.
+// 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/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "gtest/gtest.h"
+
+#include "../../test/lib/Dialect/Test/TestAttributes.h"
+#include "../../test/lib/Dialect/Test/TestDialect.h"
+#include "../../test/lib/Dialect/Test/TestTypes.h"
+#include "mlir/IR/OwningOpRef.h"
+
+using namespace mlir;
+using namespace test;
+
+TEST(InterfaceTest, DenseMapKey) {
+  DialectRegistry registry;
+  registry.insert<test::TestDialect>();
+
+  MLIRContext context(registry);
+  context.loadDialect<test::TestDialect>();
+
+  OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
+  OpBuilder builder(module->getBody(), module->getBody()->begin());
+  auto op1 = builder.create<test::SideEffectOp>(builder.getUnknownLoc(),
+                                                builder.getI32Type());
+  auto op2 = builder.create<test::SideEffectOp>(builder.getUnknownLoc(),
+                                                builder.getI32Type());
+  auto op3 = builder.create<test::SideEffectOp>(builder.getUnknownLoc(),
+                                                builder.getI32Type());
+  DenseSet<MemoryEffectOpInterface> set;
+  set.insert(op1);
+  set.insert(op2);
+  set.erase(op1);
+  EXPECT_FALSE(set.contains(op1));
+  EXPECT_TRUE(set.contains(op2));
+  EXPECT_FALSE(set.contains(op3));
+}