[RFC] Support scalable dimensions in ShapedTypes

MLIR
2

CC @javiersetoain, @nicolasvasilache, @ThomasRaoux, @MaheshRavishankar, @aartbik, @zhanghb97

3

It would be useful to have the type system differentiate between scalable dim and static ones, in the same fashion as we do for RankedTensorType and unranked ones (or MemRef).

That would make the APIs more clear and safe in general, while avoiding redundant checks for scalable dimensions everywhere.
A type that does not support scalable dimensions (like Tensor for example) should be able to return only static dimensions.

 // Will assert `outputH` has a fixed size.
  for (int64_t d = 0; d < outputH.fixedSize(); d++) {
    // ...
  }

I would rather assert if fixedSize is called without checking for isDynamic here (like llvm::Error is setup): we should catch API misuses even when the runtime value wouldn’t trigger the error.

1 Like
4

I’m also not convinced by this part of the plan: without an explicit roadmap with aggressive conversion and depreciation of the implicit conversion, it’s not clear to me that this will pan out.

2 Likes
5

Conceptually, “scalable sizes” fit better with VectorType. If you move them to ShapedType, it gives the impression that “scalable sizes” is a property that any shaped type can/may have, but VectorType is actually the only such type.

This makes the VectorType somewhat tricky to work with. Whenever manipulating a shape, you have to remember to make the same change to the scalable dims. It is easy to accidentally lose your scalable dims, or end up with the wrong dims marked scalable.

How can this happen? Vector types are created through a VectorType constructor (VectorType::get(...)). ShapedType is just an interface, and there is no ShapedType::get(...). Maybe the real problem is that we use ShapedType in places where VectorType should be used? (E.g., helper functions of vector ops returning a ShapedType when it actually could be a VectorType.)

2 Likes
6

I agree with the sentiment here. This looks more like API misuse (ShapeType instead of VectorType) than a need to carry that information above vectors.

1 Like
7

Conceptually, “scalable sizes” fit better with VectorType . If you move them to ShapedType , it gives the impression that “scalable sizes” is a property that any shaped type can/may have, but VectorType is actually the only such type.

Though not done here, scalable sizes are pretty much a dynamic size with some more constraints. They could still be interpreted by other ShapedTypes, for example the shape for vector<2x[4]xf32> when applied unchanged to a memref be treated as memref<2x?xf32>.

How can this happen? Vector types are created through a VectorType constructor (VectorType::get(...) ). ShapedType is just an interface, and there is no ShapedType::get(...) . Maybe the real problem is that we use ShapedType in places where VectorType should be used? (E.g., helper functions of vector ops returning a ShapedType when it actually could be a VectorType .)

The issue I’m trying to get at is the size/element count of a dimension is currently split across two places which makes things harder to work with than they otherwise could be.

To take one example, trimming the leading 1 dims of a vector type.
Currently, it requires maintaining two lists and carefully checking scalability:

static VectorType trimLeadingOneDims(VectorType oldType) {
  ArrayRef<int64_t> oldShape = oldType.getShape();
  ArrayRef<int64_t> newShape = oldShape;

  ArrayRef<bool> oldScalableDims = oldType.getScalableDims();
  ArrayRef<bool> newScalableDims = oldScalableDims;

  while (!newShape.empty() && newShape.front() == 1 &&
         !newScalableDims.front()) {
    newShape = newShape.drop_front(1);
    newScalableDims = newScalableDims.drop_front(1);
  }

  // Make sure we have at least 1 dimension per vector type requirements.
  if (newShape.empty()) {
    newShape = oldShape.take_back();
    newScalableDims = oldType.getScalableDims().take_back();
  }
  return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
}

If that information was part of the shape itself the same transformation is a lot simpler:

static VectorType trimLeadingOneDims(VectorType oldType) {
  ArrayRef<ShapeDim> oldShape = oldType.getShape();
  ArrayRef<ShapeDim> newShape =
      oldShape.drop_while([](ShapeDim dim) { return dim == 1; });

  // Make sure we have at least 1 dimension per vector type requirements.
  if (newShape.empty())
    newShape = oldShape.take_back();

  return VectorType::get(newShape, oldType.getElementType());
}

Note that the above code is almost identical to how trimLeadingOneDims() was before it was fixed up for scalable vectors.

8

Thank you so much for working on this! :pray:t2:

The changes proposed here would allow us to remove a great deal of footguns in the context of scalable vectors. This would be really good for the overall health of the project.

However, this change would mean 2 things:

  1. Great improvement from the VectorType point of view (I don’t see any downsides, great!),
  2. Other types deriving from ShapedType would also be allowed to have “scalable” sizes - Is this desirable?

Looks like folks are quite unsure about 2 (quite understandable - it would be a huge shift).

One possible alternative would be to disconnect VectorType from ShapedType. Given that ShapedType defines kDynamic and VectorType does not use it, that would also make sense to me. What do people think?

-Andrzej

2 Likes
9

For those if us who have not been using VectorType in this way, could you define “scalable dim”?

From context, I take it to mean that it is “dynamic” in the ShapedType sense but with a static maximum size?

If that’s the case, we’ve wanted that for tensor basically forever, and if we could harmonize the type system, that would be nice.

I’m -1 on leaking a “stored as negative” approach to the public API and think this could be represented better.

ForTensorType, I expect this would be a new BoundedTensorType or something, and that there would be new ShapedType methods like getDimUpperBound(i) that could vary when isDimDynamic(i) (or something like that).

(Edit: I expect the same reasoning applies to MemRefType but I have not thought through the implications as that type is more complicated)

10

Not exactly. The size is a fixed runtime multiple of a static size. For instance, a compile-time vector<[4]xi32> can be, at runtime, equivalent to a vector<1x4xi32>, or vector<2x4xi32>, or vector<4x4xi32>, and so on. The hardware implementation of the vector unit is what defines the multiplicity, and ISAs may define its constrains, like a maximum.

11

From context, I take it to mean that it is “dynamic” in the ShapedType sense but with a static maximum size?

For a ShapedType it’s the opposite, it has a static minimum size, which at runtime will be scaled by a constant (that’s unknown at compile time), typically called vscale.

I’m -1 on leaking a “stored as negative” approach to the public API and think this could be represented better.

Me too, it’s not something I’d like to keep around (or ever use directly). It’s just a means to allow incremental change. I want to switch to using APIs that return a mlir::ShapeDim, with the underlying storage abstracted away.

In my current proof of concept PR, the only direct access to the negative values is by indexing into .getShape()[idx]. The .getDimSize(idx) method has been updated to return the minimum size of a dimension (or kDynamic). But the preferred way to access a dimension is via a new .getDim(idx) method, which returns a mlir::ShapeDim, which you can then query.

12

Ah, thanks. That sounds pretty specific to the VectorType and may just be part of its API surface.

ShapedType can represent less information than any specific concrete type. There is no rule that says it has to model every detail so long as it represents a subset of the information.

1 Like
13

One possible alternative would be to disconnect VectorType from ShapedType . Given that ShapedType defines kDynamic and VectorType does not use it, that would also make sense to me. What do people think?

Makes sense to me, I think it only makes sense for something to implement the ShapedType interface if shapes of that type are representable, and that does not seem to be the case for VectorType:

Supports kDynamic All shapes representable as ShapedTypes? Supports scalabe dimensions?
MemRef Yes Yes No
RankedTensor Yes Yes No
VectorType No No (no scalable dims) Yes
UnrankedMemRef N/A N/A (no shape) N/A
UnrankedTensor N/A N/A (no shape) N/A
1 Like
14

This makes sense to me.

Just a note of caution, ShapedType is used quite a lot (and I’m not sure what would be the consequences of removing it from VectorType). I tried something similar a while ago (splitting ShapedType into UnrankedShapedType and RankedShapedType; [RFC] ShapedType Type Interface Cleanup) and quickly ran into issues because so much code depends on the current API. (That change still has not landed.) It may not be as difficult for VectorType though…

1 Like
15

I’d like to argue for having scalable dimensions in memref and tensor, and I’ve had that “stored as negative” idea bouncing around our kernel generator’s todo list internally.

Specifically, the reason you want scalable tensor dimensions is being able to express things like “this dimension is of an unknown size that’s divisible by 4” or “this dimension’s a multiple of 32” which allows you to, for example, determine how best to vectorize loads/stores or to determine if you need to actually emit code to pad tiles out to an internal block size.

16

Notice that a scalable dimension in a vector is a stronger constraint than “it’s an unknown multiple N”, it means “it’s a hardware-defined multiple of N”.

An important implication of this difference, for instance, is that the scale factor (vscale) is constant, and the same for all vectors throughout the lifetime of the program. Also, the runtime size of the vector can be computed at runtime, always. It’s not a static value, but it’s a fixed, well-defined value. I don’t think that’s what you want for your memrefs.

What you propose would probably fit better as a memref attribute, not unlike alignment, offsets, etc.

1 Like
17

@nicolasvasilache and I have been intermittently discussing this for years. This is a useful thing to have, but more as a generic “property” of the memref/tensor type, potentially as part of the layout attribute. Divisibility is not the only piece of information that we may want to encode, we may also want constant bounds (runtime size unknonwn, but is guaranteed to be greater/less than X) and relations between dimensions (runtime size unknown, but guaranteed to be square, aka NxN rather than ?x?; or even Nx(2*N-3) kind of shapes that appear in dilated convolutions).

2 Likes
18

I agree that a more general constraints attribute could be useful, though that might be an independent discussion from cleaning up the ShapedType wart?

And I don’t want to see this restricted to memref - tensors should be able to carry those constraints as well.

19

Introducing such dimension constraints, in ShapedType or elsewhere, would indeed justify a separate RFC. However, propagating the idea of scalability and abusing it to represent a subset of what seems to be actually desired (i.e. constraints) sounds like the opposite of cleaning up ShapedType to me.

2 Likes
20

Please, don’t close the dynamic vector dimension door just yet :slight_smile:. Dynamic vector dimensions have been discussed in the past and we are currently working on a proposal for RISC-V dynamic vectors where vector<?xtype> is one of the representation options :slight_smile:

(Sorry, I haven’t had enough time to dig into all the details of this discussion yet)

21

Thanks all for your comments.

I think the conclusion is that having a VectorType where the shape matches the actual shape of the vector would require disconnecting VectorType from ShapedType, as ShapedTypes won’t support scalable dimensions.

For this new VectorType dimensions could be represented by llvm::ElementCount (or some new MLIR type to allow extending it to cover other size types, e.g. dynamic vectors).

This is just a possibility, that I may try sometime in the future (but no solid plans to do so). Also, I’d be happy to collaborate on this (or review patches) if anyone picks this up in the meantime.

1 Like