This is an archive of the discontinued LLVM Phabricator instance.

[mlir] Add a builtin distinct attribute
ClosedPublic

Authored by gysit on Jun 20 2023, 8:22 AM.

Details

Summary

A distinct attribute associates a referenced attribute to a unique
identifier. Every call to its create function allocates a new
distinct attribute instance. The address of the attribute instance
temporarily serves as its unique identifier. Similar to the names
of SSA values, the final unique identifiers are generated during
pretty printing.

Examples:
#distinct = distinct[0]<42.0 : f32>
#distinct1 = distinct[1]<42.0 : f32>
#distinct2 = distinct[2]<array<i32: 10, 42>>

This mechanism is meant to generate attributes with a unique
identifier, which can be used to mark groups of operations
that share a common properties such as if they are aliasing.

The design of the distinct attribute ensures minimal memory
footprint per distinct attribute since it only contains a reference
to another attribute. All distinct attributes are stored outside of
the storage uniquer in a thread local store that is part of the
context. It uses one bump pointer allocator per thread to ensure
distinct attributes can be created in-parallel.

Diff Detail

Event Timeline

gysit created this revision.Jun 20 2023, 8:22 AM
Herald added a project: Restricted Project. · View Herald Transcript
gysit requested review of this revision.Jun 20 2023, 8:22 AM
gysit added a comment.Jun 20 2023, 8:29 AM

This revision implements a custom numbering of distinct attributes during printing an parsing. This has the advantage that distinct attributes can be printing both as aliases or inline on the operation itself.

An alternative may be to use the alias printing mechanism to get a numbering for free. I have a WIP revision that implements this:

https://reviews.llvm.org/D153361

It is a bit more intrusive in terms of the necessary printer and parser changes though and I am unsure if we want to give up on the option of printing every attribute inline (not as an alias). If desired I can merge the WIP commit into this revision.

Added some small comment.

mlir/include/mlir/Support/ThreadLocalStore.h
45

Just to be sure: If I understand this correctly, each template instantiation will have its own bump allocator. Would it be more efficient to combine them into one bump allocator per thread? If so, this shouldn't be a static member function. Instead, it could just be a static function.

mlir/lib/AsmParser/AttributeParser.cpp
1235

Ultra NIT: I prefer "ID" over "Id", as it is an abbreviation.

mehdi_amini added inline comments.Jun 28 2023, 2:05 PM
mlir/include/mlir/Support/ThreadLocalStore.h
46

That is leaking memory for the entire process duration, it's not clear to me why this is needed?

mlir/lib/AsmParser/AttributeParser.cpp
1263

Can you print the existing definition?

gysit updated this revision to Diff 535713.Jun 29 2023, 4:20 AM
gysit marked 4 inline comments as done.

Address comments and rebase.

gysit added a comment.Jun 29 2023, 4:27 AM

@mehdi_amini Thanks for the review comments. ThreadLocalCache indeed already implements a thread local storage with a non-static lifetime. I changed my code to use it.

mehdi_amini added inline comments.Jun 29 2023, 7:16 AM
mlir/lib/IR/AttributeDetail.h
416

Why do we need all of this instead of just defining DistinctAttr in ODS as any other attribute?

mlir/test/IR/distinct-attr.mlir
3

Can you stick to using the test dialect and avoid -allow-unregistered-dialect ? This isn't really good practice to use this option: it disables some safety net.

gysit added inline comments.Jun 29 2023, 7:32 AM
mlir/lib/IR/AttributeDetail.h
416

The distinct attribute uses a different uniquer which cannot be generated using tablegen:

class DistinctAttr
    : public detail::StorageUserBase<DistinctAttr, Attribute,
                                     detail::DistinctAttrStorage,
                                     detail::DistinctAttributeUniquer> {

There is at least currently no way to set the uniquer to something else than the default storage uniquer. The distinct attribute uses a separate DistinctAttributeUniquer though. In theory ODS could be changed to support this but I don't think we want this since there will be only one builtin DistinctAttr and users are not supposed to create there own distinct attribute types.

I can try to put a comment somewhere that explains this design decision. Probably on the distinct attribute declaration.

mehdi_amini added inline comments.Jun 29 2023, 7:35 AM
mlir/lib/IR/AttributeDetail.h
416

Oh right I forgot about the uniquer at some point...

gysit updated this revision to Diff 535805.Jun 29 2023, 7:59 AM
gysit marked an inline comment as done.

Address comments.

Seems useful!

zero9178 accepted this revision.Jul 4 2023, 7:58 AM

LGTM

This revision is now accepted and ready to land.Jul 4 2023, 7:58 AM
gysit added a comment.Jul 4 2023, 8:02 AM

Are there any more comments? @rriddle @Mogball @mehdi_amini maybe?

Dinistro accepted this revision.Jul 4 2023, 9:13 AM

I did another pass over the revision and it's LGTM!
Not sure if the others have additional comments, though.

rriddle requested changes to this revision.Jul 4 2023, 12:08 PM

This patch looks like it's missing bytecode support, have you tested that case? Failure there will be a miscompile, not a hard error.

mlir/include/mlir/IR/BuiltinAttributes.h
1012–1041

All attributes should be defined in ODS, this should be moved over. The custom allocator could be solved other ways (e.g. a parameter that is a key that we have a custom allocator/counter/etc for inside of the context).

This revision now requires changes to proceed.Jul 4 2023, 12:08 PM
gysit added a comment.Jul 5 2023, 8:49 AM

This patch looks like it's missing bytecode support, have you tested that case? Failure there will be a miscompile, not a hard error.

The revision has two changes related to bytecode:

  • mlir/include/mlir/IR/BuiltinDialectBytecode.td
  • mlir/test/Dialect/Builtin/Bytecode/attrs.mlir

I would expect these roundtrip tests fail if the implementation is broken/incomplete?

mlir/include/mlir/IR/BuiltinAttributes.h
1012–1041

Let me rephrase to make sure I understand correctly. You would introduce a new tablegen attribute parameter type derived from AttrOrTypeParameter that uses a custom allocator and a custom printer and parser? A distinct attribute could then use a parameter of this new key type which for example contains the address of an allocation made using a custom allocator. A custom printer/parser then converts the address to a deterministic identifier. Does that make sense?

I think such a solution improves composability since a user attribute could contain such a key parameter and no boxing into a distinct attribute is needed.

There are also downsides:

  • The distinctness mechanism is less encapsulated since the get method will take a key parameter and there will be an accessor for the key parameter.
  • The storage usage may increase since there is an extra allocation for the key. An attribute that contains such a key will thus use 2 x 8 bytes while the proposed distinct attribute uses only 8 bytes.
rriddle accepted this revision.Jul 10 2023, 10:11 PM

I'm okay with landing this now, given we should only ever have a single distinct attribute type.

This revision is now accepted and ready to land.Jul 10 2023, 10:11 PM
gysit updated this revision to Diff 538922.Jul 10 2023, 11:04 PM

Rebase and improve comments.

This revision was automatically updated to reflect the committed changes.