This is an archive of the discontinued LLVM Phabricator instance.

Add generic type attribute mapping infrastructure, use it in GpuToX
ClosedPublic

Authored by krzysz00 on Jan 19 2023, 2:09 PM.

Details

Summary

Remapping memory spaces is a function often needed in type
conversions, most often when going to LLVM or to/from SPIR-V (a future
commit), and it is possible that such remappings may become more
common in the future as dialects take advantage of the more generic
memory space infrastructure.

Currently, memory space remappings are handled by running a
special-purpose conversion pass before the main conversion that
changes the address space attributes. In this commit, this approach is
replaced by adding a notion of type attribute conversions
TypeConverter, which is then used to convert memory space attributes.

Then, we use this infrastructure throughout the *ToLLVM conversions.
This has the advantage of loosing the requirements on the inputs to
those passes from "all address spaces must be integers" to "all
memory spaces must be convertible to integer spaces", a looser
requirement that reduces the coupling between portions of MLIR.

ON top of that, this change leads to the removal of most of the calls
to getMemorySpaceAsInt(), bringing us closer to removing it.

(A rework of the SPIR-V conversions to use this new system will be in
a folowup commit.)

As a note, one long-term motivation for this change is that I would
eventually like to add an allocaMemorySpace key to MLIR data layouts
and then call getMemRefAddressSpace(allocaMemorySpace) in the
relevant *ToLLVM in order to ensure all alloca()s, whether incoming or
produces during the LLVM lowering, have the correct address space for
a given target.

I expect that the type attribute conversion system may be useful in
other contexts.

Diff Detail

Event Timeline

krzysz00 created this revision.Jan 19 2023, 2:09 PM
Herald added a project: Restricted Project. · View Herald Transcript
krzysz00 requested review of this revision.Jan 19 2023, 2:09 PM

The overall approach looks good to me. Two design comments:

  1. There is no reason to limit this to memory space attributes only and no good reason to introduce the notion of memory space to the top-level TypeConverter. Instead, this can be a facility to convert attributes. (Renaming TypeConverter to AttrAndTypeConverter is probably too much churn for one use case, so one can argue that it intends to support attributes present in types for now).
  2. It looks strange that the attribute conversion must always succeed. IMO, it should have an ability to fail, especially given the extension beyond memory spaces. When set up for conversions where there is a default, it is possible to add a low-priority catch-all conversion rule, e.g., for all attributes implementing the memory space interface. The code, inherited from type converters, already does a loop around conversions to support this.
mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h
135
mlir/include/mlir/Transforms/DialectConversion.h
163
167
mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp
320

We usually want a test for user-visible error messages.

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
689

Nit: I see it was there before this change, but could we avoid the magic number here?

ftynse requested changes to this revision.Jan 31 2023, 2:56 AM

(requesting changes so it shows up in my review thread again)

This revision now requires changes to proceed.Jan 31 2023, 2:56 AM

@ftynse I appreciate the design feedback!

I can agree that having a more general facility for attributes in types would be quite nice, and that having the conversion being able to fail would be good (and would make for a nice symmetry with the actual type converter).

My one concern is that attributes are used in different context, and the conversions could differ depending on those contexts. For example, we could imagine that, while memory spaces need to be converted to integers when lowering to LLVM (which makes IntegerAttr a trivial conversion), there could be some other type and attribute where you want integers to be converted to some kind of backend-specific enum, and so you'd need a different set of conversions for that context.

So the general signature might be Optional<Attribute> convertTypeAttribute(SomethingIDontKnowYet context, Attribute attribute);

One option for SomethingIDontKnowYet is the type the attribute came from.
Then, the signature of the conversion functions would be addTypeAttributeConversion([Type, Attribute] -> FailureOr<Attribute>) and you could, for example, add conversions for BaseMemRefType and so on. This also means attribute converters get some extra context to peek in to.

Another thought would be some sort of key, but then we have a stringly-typed hash table of conversion groups.

.... yeah, having written that out, convertTypeAttribute(Type type, Attribute attribute) is the better general signature.

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
689

I'd leave it there for backwards compatibility.

To amend, attribute conversions will need some way to return three states

  1. Failed, abort conversion attempt
  2. Not applicable
  3. Here's your result

which we can't do the way that type conversions can because nullptr is a meaningful attribute value for a lot of contexts.

Or we could not have the early-stop system.

krzysz00 updated this revision to Diff 495239.Feb 6 2023, 12:15 PM
krzysz00 retitled this revision from Add generic memory space mapping infrastructure, use it in GpuToX to Add generic type attribute mapping infrastructure, use it in GpuToX.
krzysz00 edited the summary of this revision. (Show Details)

Make the mapping system more generic and allow it to fail.

ftynse accepted this revision.Feb 7 2023, 7:32 AM

Originally, I was thinking that it was okay to just have context-free attribute conversion here. It's the same for types nested in other types: one may eventually need to convert types differently based on what they are nested in, but there is no support for context. Same for attributes, there's currently no need for context. But I also like the approach adopted here.

Regarding the early stopping and return values, it is possible to define a custom return type that steals two bits from the pointer to the attribute internal storage, using AttrConvresionResult = llvm::PointerIntPair<Attribute, 2> and interpret those bits as: 00 here's your result and it's null, 01 here's your result and it's a concrete attribute, 10 keep trying, 11 stop.

mlir/include/mlir/Transforms/DialectConversion.h
177–180

Nit: would putting this in a using declaration inside the function work instead?

258–259

Who determines the semantics of returning nullptr here, the caller?

397–399

This can grow quite fast... Not sure if caching with pair as index is worth it, especially given that the logic differs per type class (e.g, the same callback works for instances of MemRefType, but we will store for each instance).

This revision is now accepted and ready to land.Feb 7 2023, 7:32 AM

@ftynse Thanks for the feedback!

On the PointerIntPair method, would the solution to be making a wrapper around PointerIntPair<Attribute, 2>? Something along the lines of

enum TypeAttrConversionResult {
   Result(Attribute),
   NullAttribute,
    Continue,
   Abort
}

(that there being how I'd phrase it in Rust - the actual implementation would, I think, involve PointerIntPair and such)

The question I have is whether this struct would go inside TypeConverter or be free-standing?

mlir/include/mlir/Transforms/DialectConversion.h
177–180

We need it in the type of the vector of converters

258–259

Yeah. Good question. I should clarify.

397–399

Agreed, should probably remove.

On the PointerIntPair method, would the solution to be making a wrapper around PointerIntPair<Attribute, 2>? Something along the lines of <...>

SGTM, good for a separate patch.

krzysz00 updated this revision to Diff 495605.Feb 7 2023, 11:53 AM

Use AttributeConversionResult instead of std::optional<Attribute>

krzysz00 updated this revision to Diff 495945.Feb 8 2023, 2:12 PM

Rebase away a merge conflict

mehdi_amini added inline comments.Oct 26 2023, 1:33 AM
mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
97

You're dereferencing an optional here: this is crashing behavior, the error should be handled and propagated instead. Just hit this in https://github.com/llvm/llvm-project/issues/70160#issuecomment-1780644011

mlir/test/Conversion/GPUCommon/lower-memory-space-attrs.mlir