Motivation / The problems with VectorType
Currently, ShapedTypes can have dynamic dimensions or fixed/static dimensions. This leads to an awkward situation for VectorType. It is a ShapedType, but a large part of its shape, whether a dimension is scalable, is stored in a separate array of booleans unknown to ShapedType.
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. You also have to keep in mind the value of a dim has a different meaning depending on the corresponding scalable flag. vectorType.getDim(0) == 1 is only a unit dim if !vectorType.getScalableDims()[0].
Proposal
This RFC proposes to remove the scalableDims array of bools. Instead, scalability will be represented at the ShapedType level by using negative values for the dimension sizes.
| Current | Proposal | |
|---|---|---|
kDynamic |
std::numeric_limits<int64_t>::min() |
std::numeric_limits<int64_t>::min() |
Negative values (other than kDynamic) |
Invalid | Scalable dimensions |
| Positive values (including zero) | Fixed/static dimensions | Fixed/static dimensions |
Doing this while still treating dimensions as int64_t could prove error-prone, so a new class mlir::ShapeDim would be added to represent dimensions. This class just needs to wrap a single int64_t, so is no change to the sizeof() a ShapedType.
The mlir::ShapeDim class has methods to construct dims, and query if a dim is scalable, fixed, or dynamic, and safe ways to get the size (or min size) of a dimension. See the implementation below. This is similar to llvm::ElementCount, which is used to represent fixed or scalable vector quantities within LLVM. The main differences are the addition of the dynamic state, and that is it possible to round-trip ShapeDim → int64_t → ShapeDim without losing your scalable dimensions, which helps keep the initial change relatively small.
There is a lot of code that uses int64_t for dimensions, so currently mlir::ShapeDim can be implicitly converted to/from int64_t. It is hoped that code can transition to using ShapeDim and eventually the implicit conversions can be removed. This path is similar to llvm::ElementCount where implicit conversions are allowed unless compiled with STRICT_FIXED_SIZE_VECTORS, at which point implicit conversions become an error.
How does the affect existing ShapedTypes?
- No change in size (a dim is still just backed by an
int64_t) - If you don’t need scalable dims, things mostly stay the same, but there are
some extra assertions
Some code that never expects to work for scalable sizes:
// Note: This still works for the time being.
int64_t outputH = resultTy.getDimSize(0);
if (outputH == ShapedType::kDynamic) {
// ...
} else {
for (int64_t d = 0; d < outputH; d++) {
// ...
}
}
Would become:
ShapeDim outputH = resultTy.getDim(0);
if (outputH.isDynamic()) {
// ...
} else {
// Will assert `outputH` has a fixed size.
for (int64_t d = 0; d < outputH.fixedSize(); d++) {
// ...
}
}
Proof-of-concept implementation
Conclusion
Thanks for taking a look! We’d appreciate any feedback, or being told we’ve missed something obvious :). Note that this is being done as part of a larger effort to enable scalable vectorization (such as [RFC] Scalable Vectorisation in Linalg). This particular change is an attempt to alleviate some pain points we’ve noticed while doing that work.
mlir::ShapeDim class
struct ShapeDim {
/// Deprecated. Use ShapeDim::fixed(), ShapeDim::scalable(), or
/// ShapeDim::kDynamic() instead.
constexpr ShapeDim(int64_t size) : size(size) {}
/// Deprecated. Use an explicit conversion instead.
constexpr operator int64_t() const { return size; }
/// Construct a scalable dimension.
constexpr static ShapeDim scalable(int64_t size) {
assert(size > 0);
return ShapeDim{-size};
}
/// Construct a fixed dimension.
constexpr static ShapeDim fixed(int64_t size) {
assert(size >= 0);
return ShapeDim{size};
}
/// Construct a dynamic dimension.
constexpr static ShapeDim kDynamic() { return ShapeDim{}; }
/// Returns whether this is a dynamic dimension.
constexpr bool isDynamic() const { return size == kDynamic(); }
/// Returns whether this is a scalable dimension.
constexpr bool isScalable() const { return size < 0 && !isDynamic(); }
/// Returns whether this is a fixed dimension.
constexpr bool isFixed() const { return size >= 0; }
/// Asserts a dimension is fixed and returns its size.
constexpr int64_t fixedSize() const {
assert(isFixed());
return size;
};
/// Asserts a dimension is scalable and returns its size.
constexpr int64_t scalableSize() const {
assert(isScalable());
return -size;
}
/// Returns the minimum (runtime) size for this dimension.
constexpr int64_t minSize() const {
if (isScalable())
return scalableSize();
if (isFixed())
return fixedSize();
return 0;
}
private:
constexpr explicit ShapeDim()
: ShapeDim(std::numeric_limits<int64_t>::min()) {}
int64_t size;
};
. Dynamic vector dimensions have been discussed in the past and we are currently working on a proposal for RISC-V dynamic vectors where