Page MenuHomePhabricator

[mlir] Define proper DenseMapInfo for Interfaces
ClosedPublic

Authored by zero9178 on Jul 2 2022, 5:21 AM.

Details

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.

Diff Detail

Event Timeline

zero9178 created this revision.Jul 2 2022, 5:21 AM
Herald added a project: Restricted Project. · View Herald Transcript
zero9178 requested review of this revision.Jul 2 2022, 5:21 AM
zero9178 updated this revision to Diff 441883.Jul 2 2022, 7:24 AM

Fix examples build

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 added inline comments.Jul 2 2022, 7:50 AM
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 updated this revision to Diff 441937.EditedJul 3 2022, 3:58 AM
  • 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 requested changes to this revision.Jul 4 2022, 10:04 PM
rriddle added inline comments.
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.

This revision now requires changes to proceed.Jul 4 2022, 10:04 PM
zero9178 updated this revision to Diff 442218.Jul 5 2022, 2:20 AM

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

zero9178 retitled this revision from [mlir][tblgen] Generate proper DenseMapInfo for Interfaces to [mlir] Define proper DenseMapInfo for Interfaces.Jul 5 2022, 2:23 AM
zero9178 edited the summary of this revision. (Show Details)
Mogball added inline comments.Jul 5 2022, 6:34 AM
mlir/include/mlir/Support/InterfaceSupport.h
288

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

zero9178 marked an inline comment as done.Jul 5 2022, 7:04 AM
zero9178 added inline comments.
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 added inline comments.Jul 5 2022, 7:34 AM
mlir/include/mlir/Support/InterfaceSupport.h
288

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

Mogball added inline comments.Jul 5 2022, 7:39 AM
mlir/include/mlir/Support/InterfaceSupport.h
288

Ah, right. Type and attribute interfaces.

zero9178 marked an inline comment as done.Jul 5 2022, 7:39 AM
zero9178 added inline comments.
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 updated this revision to Diff 442304.Jul 5 2022, 7:48 AM

Add tests for attribute and type interfaces as well.

Mogball accepted this revision.Jul 5 2022, 8:01 AM
Mogball added inline comments.
mlir/include/mlir/Support/InterfaceSupport.h
277

And below as well.

mlir/unittests/IR/InterfaceTest.cpp
13–19
47

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

zero9178 updated this revision to Diff 442314.Jul 5 2022, 8:17 AM

Address review comments

rriddle added inline comments.Jul 5 2022, 8:37 AM
mlir/include/mlir/Support/InterfaceSupport.h
287

Please separate out the namespace definition from the struct.

zero9178 updated this revision to Diff 442318.Jul 5 2022, 8:40 AM

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

Mogball added inline comments.Jul 5 2022, 3:00 PM
mlir/include/mlir/Support/InterfaceSupport.h
305
rriddle accepted this revision.Jul 6 2022, 1:43 AM
This revision is now accepted and ready to land.Jul 6 2022, 1:43 AM
This revision was automatically updated to reflect the committed changes.
zero9178 added inline comments.Jul 6 2022, 3:39 AM
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 added inline comments.Jul 6 2022, 8:17 AM
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 added inline comments.Jul 6 2022, 10:25 AM
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.