[RFC] MemRef memory shape as Attribute

Hello!

I’d like to discuss one topic about MemRef type - the memory space property.

For now, it is represented as simple unsigned int value. This makes dumped IRs less readable, forcing reader to remember how to map this unsigned int value to actual memory hierarchy for particular Dialect. For example, you have to know that for GPU Dialect 0 means global memory, 3 means workgroup memory address, 5 means private memory address. For another Dialect those values might have different meaning.

My question - can we generalize the memory space property for MemRef type to make it more verbose regarding actual memory hierarchy and use cases? For example, use generic Attribute to store it. This will allows to use integer values, strings, enumerations or even custom Dialect attributes. For example, for GPU Dialect:

gpu.func @test(%arg0: memref<100xf32, "global">,
               %arg1: memref<100xf32, "workgroup">,
               %arg2: memref<100xf32, "private">)

One issue with this approach I see - some code use comparison operators for the memory space (DMA operations from Standard dialect, for example). This can be handled in 2 ways:

  1. Introduce separate ComparableAttrInterface interface to allow Attribute comparisons. In that case MemRef type can require from memory space attribute to implement this interface.
  2. Pass comparison callback as separate parameter for MemRef::get. This callback will be used to compare memory spaces for MemRef types.

What is your option about this?

+1 on removing the magic integer.

A couple of comments:

  • I am not convinced that we want an arbitrary, potentially foreign-dialect, attribute to specify the memory space. We would see crazy things like DenseElements or AffineMap. Wouldn’t a string attribute be enough?
  • Attributes are uniqued and directly comparable, so there is little interest in some additional mechanism unless you want to make two different attributes compare equal, e.g. IntegerAttr(5) == StringAttr(private). IMO, such implicit behavior is undesirable in an IR.

I agree with Alex that using pointer equality of the attribute is a good way to go. We should also align TensorType to use the same memory space as MemRefType IMO.

I don’t think there is any harm in allowing an arbitrary attribute here. Memory spaces are opaque in general - the only thing target-independent code can do is pointer-compare them. As such, it seems reasonable for targets to be able to define their own structures here. I agree that a string is going to be a very common use case, but I could imagine wanting to use structs with a few fields in some exotic cases.

1 Like

This sounds great to me - clearly, much more readable with strings or other structures if necessary. For the GPU dialect, we could even consider using “gpu.global”, “gpu.workgroup”, etc. How would printing of the memory space attribute be handled? (Would you get memref<..., 0 : i32> or would it be customized? What will the default / ellision rule be?

+1 on aligning tensor type with memref here. We’re already running into situations during bufferization where having a closer (possibly up to isomorphic) link between tensor and corresponding memref types seems to be the right answer.

It is a long TODO to make memory spaces strings instead of integer, I’d be very happy if someone take on this :slight_smile:
This is something I wanted to do in LLVM IR for a long time as well (but this requires a bit more effort).
Short of fixing LLVM IR, something that needs to be sorted out for MLIR will be the mapping of the strings to LLVM IR corresponding integers.

Thank you for your comments!

Some notes from my side.

By comparison interface I meant order comparison (greater / less), not only equality. There is code in std dialect (DMA operations). This code compares memory spaces with greater/less operator in the meaning faster/slower memory space. If we replace unsigned int with string/attribute this code will be broken.

As for LLVM IR mapping - it can be done as a part of lowering pass for particular Dialect. LLVM Dialect will require using IntergerAttr for memory space and lowering pass must convert internal memory space representation to one, supported by LLVM. But if the project doesn’t use lowering to LLVM is it free to use any exotic attributes for memory space as long as it can interpret it.

There is also another way, which is already discussed several times in different topics.

We need to convert ShapedType and MemRefBase from base classes to type interfaces. This will allow to Dialects define own types for tensors/memory buffers, but implement those interfaces.

This can be resolved with an AttributeInterface potentially.

I wonder if there is somebody who uses those ops and could share their use case. Given that nobody complained about them in over two years (for comparison, memref design and load/store model changed twice over the same period) and the description still refers to “affine context” and uses affine maps for layout, I am guessing nobody does. So I would rather change or remove the ops than complicate the design of a built-in type.

What do you expect to reuse through this interface? I assume virtually none of the code base will work seamlessly with memory buffers with the design significantly different from that of a memref. I am not saying we shouldn’t do this, but we should do it after weighing the pros and cons. We don’t have any built-in type implement an interface, so it is not clear whether the overhead of going through one is acceptable.

This will be the simplest solution. We can just remove the memory space comparison logic without removing the operations at start. But I’d like to be sure, that it is not used.

This interface will provide access to common properties (such as shape, element type, affine maps). But the custom type might add more information specific to its domain. One example is more complex memory space representation. Another example, is simplified layout information, when the affine maps will be built on top of it.

We do not need to rework all code base to the interfaces, only the parts that works with this common properties. If some parts requires more specifics for the types (like lowering to LLVM) it might require exact type usage (std.memref). In that case the Dialect must convert its own custom memory type to std.memref during lowering phase.

But, I agree with you, it is more radical change and it is need to be discussed more deeply.

A less intrusive change is for those ops to require IntegerAttr memory space.

Well, if a type has a shape in the same way MemRef does and the same affine map structure, how this type is different from a MemRef? :slight_smile: If it supports more or less complex memory addressing modes, or in particular modes that are not expressible with affine maps, than I don’t see a reason why this type should implement an interface that specifically promises there are affine maps. At this point, it’s just a ShapedType. At this point, ShapedType is mostly a way to share parsing/printing/verification + a couple of convenience functions. We can share the code without sharing concepts.

Terminology nit: it’s just memref, “standard” types do not belong to the standard dialect.

In principle, one could have my_memref which tracks random other data – whether it is read-only, upper bounds on the sizes of dimension, … That is, it makes sense for my_memref to be a type that carries a superset of the information carried by memref, and still satisfy MemrefTypeInterface.

1 Like

Submitted a patch with proposed change: https://reviews.llvm.org/D96145
Can we continue discussion there?

This change seems to have altered the meaning of the ‘default’ memory space. In the past, the value 0 was considered the default memory space. Now that the memory space is an Attribute, this no longer round-trips properly:

memref<2x4xf32, 0>

The code now seems to assume that the default memory space is a null Attribute (instead of an IntegerAttr with a value of 0).

I guess my question here is, what should the default value be? A null Attribute seems sensible, but then I would expect for memref<2x4xf32, 0> to roundtrip to memref<2x4xf32, 0> instead of memref<2x4xf32> as it does today.

The 0 value internally is converted to NULL Attribute to preserve its ‘default memory space’ meaning. I was trying to avoid extra changes when possible.

That a syntactic discrepancy, it isn’t clear to me why are we allowing integer in the syntax at all? If we expect an attribute, I think that is what the parser should get here (and what we should read).

Why are we allowing integer in the syntax at all?

Integer is a valid attribute (represented as IntegerAttr), that’s why it is allowed. I left integer memory space support for 2 reasons:

  • Compatibility with current code in LLVM/GPU/SPIRV dialects.
  • It might be used on simple cases like plain cache hierarchy (L1/L2, for example).

The 0 value is dropped as a ‘default’ memory space. This is similar to affine maps behavior, when identity maps are dropped (memref<1x2xf32, affine_map<(d0, d1) -> (d0, d1)>> is equivalent to memref<1x2xf32>).

Yes you’re right, I’m not used to see them without a type but it defaults to i64!