[NFCI] Add StackOffset class and base classes for ElementCount, TypeSize.

[sdesmalen]

Summary

This patch adds a linear polynomial base class, called LinearPolyBase, which
serves as a base class for StackOffset. It tries to represent a linear
polynomial like:

c0 * scale0 + c1 * scale1 + ... + cK * scaleK

where the scale is implicit, meaning that only the coefficients are
encoded.

This patch also adds a univariate linear polynomial, which serves as
a base class for ElementCount and TypeSize. This tries to represent a
linear polynomial where only one dimension can be set at any one time,
i.e. a TypeSize is either fixed-sized, or scalable-sized, but cannot be
a combination of the two.

class LinearPolyBase
   ^
   |
   +---- class StackOffset  (dimensions = 2 (fixed/scalable), type = int64_t)

class UnivariateLinearPolyBase
   |
   |
   +---- class LinearPolySize (dimensions = 2 (fixed/scalable))
                ^
                |
                +-------- class ElementCount  (type = unsigned)
                |
                |
                +-------- class TypeSize      (type = uint64_t)

[ctetreau]

This isn't really a true polynomial because it doesn't support multiplication of two terms: 2x * 2x = 4x^2

If what we have here is all we need, then I think it's fine (having not actually code reviewed the implementation at the time of this writing), but we should call it like it is, and ideally implement it in such a way that it would be straightforward to implement a true polynomial.

[sdesmalen]

In D88982#2317443, @ctetreau wrote:

This isn't really a true polynomial because it doesn't support multiplication of two terms: 2x * 2x = 4x^2

If what we have here is all we need, then I think it's fine (having not actually code reviewed the implementation at the time of this writing), but we should call it like it is, and ideally implement it in such a way that it would be straightforward to implement a true polynomial.

Are you suggesting rename PolyBase or just to remove the term 'polynomial' from the commit message and any of the comments?

[ctetreau]

In D88982#2317488, @sdesmalen wrote:

Are you suggesting rename PolyBase or just to remove the term 'polynomial' from the commit message and any of the comments?

I'm suggesting a rename. According to wikipedia, a polynomial of first degree is called a linear polynomial. So, I propose the following simplified hierarchy:

LinearPolyBase<Type, Dimensions, IsUnivariate>
   ^   
   |   
   +---- StackOffset : LinearPolyBase<int64_t, 2, false>
   ^
   |   
   +---- PolySize : LinearPolyBase<Type, 2, true>
                ^   
                |   
                +------------- ElementCount (typedef with Type=unsigned)
                |   
                +------------- TypeSize (subclass with Type=uint64_t)

If we do this, and don't define operator* with another linear polynomial on the right, then we should be good to go. I don't think MixedPolyBase2D bought us anything, all the new things you added were "scalable" specific anyway. Similarly, for SinglePolyBase2D: might as well just let it be a "scalable size" IMO.

[ctetreau]

llvm/include/llvm/Support/TypeSize.h

39

NIT: Nomenclature: suggest changing to IsUnivariate to match the terminology of the mathematical concept

41

NIT: might as well use std::array

47

NIT: it would be more immediately obvious what this is if you used std::numeric_limits<unsigned>::max()

85

I don't think the call to verify() is necessary in any of the arithmetic operators. Just assert that RHS and LHS have the same value set for ExclusiveDim.

You could even special case:

if (DimIsExclusive) {
   Coefficients[ExclusiveDim] += RHS.Coefficients[ExclusiveDim];
}
else {
   ...
}

[sdesmalen]

• Renamed PolyBase -> LinearPolyBase, DimIsExclusive -> IsUnivariate, etc.
• Removed verify() method.
• Replaced regular array with std::array.
• Changed name of new StackOffset class in TypeSize.h to NewStackOffset, so not to clash with the one defined in AArch64StackOffset.h. This is renamed to StackOffset in D88983 when it removes AArch64StackOffset.h entirely.
• Moved TestStackOffset.cpp from D88983 to this patch.
• Updated commit message.

[sdesmalen]

• Rebased patch.
• Added class LinearPolyBaseOperators which provides a set of overloaded operators for derived classes.
• Made ElementCount a deriving class, because it needs 'isScalar' and 'isVector' methods which cannot be part of its parent class.

[david-arm]

llvm/include/llvm/Support/TypeSize.h

34

Can you remove the need for ScalarTy by having Ty provide the scalar type? i.e. Ty::ScalarTy, where ScalarTy would be a typedef or something inside LinearPolyBase?

Instead of using multiple inheritance to get the operators, is it not possible to simply have the operators outside of any class? I think C++ allows you to do this provided the operators are friends of the relevant classes, i.e. just have

template <Ty>
LinearPolyBase<Ty> operator+(const LinearPolyBase<Ty> LHS, const LinearPolyBase<Ty> RHS) {

for (unsigned I = 0; I < Dimensions; ++I)
  LHS.Coefficients[I] -= RHS.Coefficients[I];
return LHS;

}

although I realise the example above only works if you have a simplified LinearPolyBase that assumes 2 dimensions, rather than the 3 template parameters you currently have for LinearPolyBase.

76

Do we actually need to support this right now? I just wonder if we can simplify things by just assuming it's always univariate for now. I guess I'm just worried that we're trying to provide theoretical support for something with no proof as yet that it's needed. We're always going to be targeting one backend, so when switching backends vscale will mean whatever is appropriate for the given backend. I think this code is trying to support a single backend having multiple vscales, unless I've misunderstood something?

Also, what happens when the user specifies 3 dimensions and IsUnivariate=true? I guess this is technically allowed, but looks a little odd that's all.

This is just a suggestion, but I wonder if the code might be made simpler by always assuming coefficient 0 is unscaled. Then 2 dimensions always implies both univariate and that coefficient 1 is scaled? By definition, 3 dimensions automatically implies IsUnivariate=false too I think? In this case you may actually only need two template parameters. Also, if we assume for now we're always dealing with exactly two dimensions then the enable_if stuff goes away too as then you only need one two-arg constructor.

[sdesmalen]

llvm/include/llvm/Support/TypeSize.h

34

Instead of using multiple inheritance to get the operators, is it not possible to simply have the operators outside of any class

That is kind of what this class tries to do. It defines the functions as friend functions, where it takes the result/operand types from the derived class, like TypeSize, ElementCount or StackOffset.

The problem with using inheritance for the operators is that it inherits the parent's operators, which return the parent class' type. That requires implicit conversion of parent -> derived class. By inheriting from this class like this:

class ElementCount : public LinearPolyBase<...>,
                     public LinearPolyBaseOperators<ElementCount, ...> {
};

it will automatically instantiate those operators for the derived class.

76

Do we actually need to support this right now? I just wonder if we can simplify things by just assuming it's always univariate for now.

We need this for StackOffset, which has IsUnivariate = false, because it needs to support a combination of fixed and scalable elements.

I think this code is trying to support a single backend having multiple vscales, unless I've misunderstood something?

I've made it generic enough in case this is needed in the future, mostly because it didn't really make things more complicated. I guess we could fix the Dimensions to 2, and remove this template operand from the class and explicitly name these dimensions 'Fixed' and 'Scalable', although I'm not sure if that would actually simplify these classes much.

Also, what happens when the user specifies 3 dimensions and IsUnivariate=true? I guess this is technically allowed, but looks a little odd that's all.
This is just a suggestion, but I wonder if the code might be made simpler by always assuming coefficient 0 is unscaled. Then 2 dimensions always implies both univariate and that coefficient 1 is scaled? By definition, 3 dimensions automatically implies IsUnivariate=false too I think? In this case you may actually only need two template parameters

The dimensions don't have a meaning in this class, it's just a representation for: c0 * scale0 + c1 * scale1 + ... + cK * scaleK. This abstract class could represent a K dimensionsal linear polynomial, where IsUnivariate would ensure that only one of {c0, c1, ..., ck} is non-zero.

Also, if we assume for now we're always dealing with exactly two dimensions then the enable_if stuff goes away too as then you only need one two-arg constructor

The enable_if is there to distinguish constructors for the Univariate and non-Univariate cases. i.e. If only one dimension is allowed to be set, you call it with (cK, #dimK), otherwise you call it with ({c0, c1, ..., cK}).

[ctetreau]

@sdesmalen I'm generally ok with this approach. This patch is pretty complicated, so I haven't had the time to fully digest it and decide if it's actually correct. I plan to try to find time to do this soon, but I'd feel more comfortable if others took a look as well.

llvm/include/llvm/Support/TypeSize.h

76

I like the template parameter for the number of different variables, and would prefer that we didn't fix it to 2. I think it doesn't really increase the complexity, and makes supporting more in a hypothetical future where some target has multiple vscale like things much easier. It is my personal opinion that N+1 template parameters aren't more complex than N template parameters unless N is zero.

[ctetreau]

llvm/include/llvm/Support/TypeSize.h

34

Is it ever valid for the operator to have a different scalar type than the polynomial?

If not, then I agree with David that LinearPolyBase should publicly expose a using ScalarType = T, and that this type should only 1 template parameter that is the polynomial type.

It might also be nice to do some type_traits magic to static_assert that the polynomial type actually inherits from LinearPolyBase so the user doesn't get 20 pages of template garbage if they mess it up.

76

NIT: add static_assert that Dimensions is not unsigned max

151

why "NewStackOffset"?

152

Shouldn't the first template parameter be NewStackOffset?

168

What's the point of this? NewStackOffset is already basically a pair of int64_t.

Also, the user must read the implementation to know which member of the pair is the fixed coefficient and which is the scalable one. At the very least, this should be documented.

183

If we name these dims, we can just use them instead of mapping bool -> unsigned

184–185

190–191

193

199

llvm/unittests/Support/StackOffsetTest.cpp

1 ↗

(On Diff #297619)

We should probably just write unit tests for LinearPolyBase, and get rid of tests for ElementCount, TypeSize, and StackOffset unless they test actual new features.

[sdesmalen]

• Removed need for LinearPolyBaseOperators. The code now passes the derived leaf type (i.e. ElementCount/TypeSize) directly to its parent classes and the operators in the base class are now defined to take/return the leaf type.
• The scalar type is no longer passed separately, but derived from the Leaf type using a type traits class.
• Added new tests based directly on LinearPolyBase.

[sdesmalen]

Thanks for the feedback @ctetreau!

llvm/include/llvm/Support/TypeSize.h

34

Is it ever valid for the operator to have a different scalar type than the polynomial?

It is not.

If not, then I agree with David that LinearPolyBase should publicly expose a using ScalarType = T, and that this type should only 1 template parameter that is the polynomial type.

It might also be nice to do some type_traits magic to static_assert that the polynomial type actually inherits from LinearPolyBase so the user doesn't get 20 pages of template garbage if they mess it up.

Okay I've changed things a bit in the latest update, by using a Traits class to get the ScalarType, and passing the Leaf type (i.e. StackOffset/TypeSize/ElementCount) directly to the base class. That also rendered the LinearPolyBaseOperators class unnecessary, as this can now move directly into the base class.

I've also tried adding some static_assert(std::is_base_of<LinearPolyBase, LeafTy>), but that falls over because LeafTy is incomplete at the point where this is being evaluated.

151

That was to avoid a name clash with StackOffset defined in AArch64StackOffset.h when both that header and TypeSize.h are included. This got renamed back in D88983. But it's better to put this in a separate namespace, and remove only that namespace in D88983.

152

Yes. Not sure how this happened.

168

The NewStackOffset is indeed a pair, but it's not a std::pair, so I thought it would be nice to have this as convenience function. But it's probably better to remove this interface for now, as I don't use it in any of the follow-up patches.

183

That's a great suggestion, thanks!

[ctetreau]

llvm/include/llvm/Support/TypeSize.h

41–52

Couldn't we just require that any type used as LeafTy have a ScalarTy typedef?

149–152

Does this need to be inside the NewStackOffset namespace? I believe that you are forward declaring llvm::StackOffset and then creating the type traits for it, not llvm::NewStackOffset::StackOffset

[sdesmalen]

• Fixed namespace for StackOffset traits.
• Noticed that StackOffset was still using int64_t instead of it's ScalarTy, so fixed it.
• Changed instance of std::all_of -> llvm::all_of

[sdesmalen]

llvm/include/llvm/Support/TypeSize.h

41–52

When I try that, I get:

error: invalid use of incomplete type 'class llvm::StackOffset'
   55 |   using SomeOtherTy = typename LeafTy::SomeTy;
      |         ^~~~~~~~~~~

149–152

Good catch! I'm not sure why this still built. I can't move the specialization into the NewStackOffset namespace, as this needs to happen in the same namespace as it's declaration, but I can move the forward declaration of StackOffset into a namespace, and specify the traits for NewStackOffset::StackOffset.

[georges]

llvm/include/llvm/Support/TypeSize.h

65–66

Not a suggestion, but just an observation that we're passing some arguments to the class via template parameters and others via the traits object. Seems inconsistent?

124–127

Is it worth specializing this struct template for IsUnivariate=true/false, since with the number of enable_if's scattered around they begin to look like quite different classes? It also seems a bit wasteful to be allocating Dimensions elements when all but one of them will be zero in the Univariate case?

178

nit: I think you can use the _v/_t versions of these functions here and elsewhere since we're in C++14 now, e.g. std::enable_if_t<std::is_signed_v<U>, LeafTy>.

[ctetreau]

llvm/include/llvm/Support/TypeSize.h

185–188

Out of bounds array access.

[ctetreau]

llvm/include/llvm/Support/TypeSize.h

41–52

I played around with it for a bit today, and I was unable to make it any nicer.

[sdesmalen]

llvm/include/llvm/Support/TypeSize.h

124–127

@ctetreau do you have any thoughts on this? I guess I could split this up and end up with a two base classes, as I agree with @georges that for >2 dimensions the code isn't optimal. By having two base classes, the univariate case can be optimized. The downside is some duplication, but it would avoid the need for enable_if and probably simplify the code.

185–188

Eek, good spot!

[ctetreau]

llvm/include/llvm/Support/TypeSize.h

124–127

I'm not so much worried about wasted space because Dimensions is likely to be low.

However, my general feeling is that this class getting way too fancy with the template metaprogramming. I wasn't going to complain about it, but since you're asking I think it would be more straightforward if they were separate classes.

I've also wondered on more than on occasion if we ought to just use a linear algebra lib for this. This class smells a lot like a vector to me...

[sdesmalen]

• Separated out the base classes, one specialized for StackOffset and another specialized for LinearPolySize.
• Moved dimensions to Traits class.
• Use enable_if_t instead of enable_if

[sdesmalen]

Okay, I have addressed all comments. Specializing the two base classes for Mixed and Univariate cases and moving the dimension to the Traits class improves the code quite a bit.
Let me know what you think!

llvm/include/llvm/Support/TypeSize.h

124–127

Using a linear algebra lib is probably a bit overkill here as the few supported operations are simple enough to implement. Unless there is already something in LLVM for this?

[georges]

llvm/include/llvm/Support/TypeSize.h

128

worth an assert that UniveriateDim is in range?

[georges]

Other than the nit, LGTM. I'll let others have a look and accept.

[ctetreau]

Looks good to me once you add that assert.

llvm/include/llvm/Support/TypeSize.h

124–127

Yeah, that was my thought too, which is why I didn't actually suggest it. I'm surprised we don't have a linear algebra lib already (either homegrown or not) though. Seems like it'd be useful for doing things like constexpr matrix multiplications.

128

I think this would be a useful assert

[sdesmalen]

Added assert.

[sdesmalen]

llvm/include/llvm/Support/TypeSize.h

128

Agreed, good catch!

[david-arm]

LGTM and a lot tidier now, provided you're also happy with this @ctetreau?

[ctetreau]

Yeah, I think this is good to go

[sdesmalen]

Thanks for the reviews and great feedback, I feel like the patch has much improved since its first revision!

llvm/include/llvm/Support/TypeSize.h

 // which may be scalable vectors. It provides convenience operators so that
 // it can be used in much the same way as a single scalar value.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_TYPESIZE_H
 #define LLVM_SUPPORT_TYPESIZE_H
 
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/WithColor.h"
 
-#include <cstdint>
 #include <cassert>
+#include <cstdint>
+#include <type_traits>
 
 namespace llvm {
 
 template <typename T> struct DenseMapInfo;
 
-// TODO: This class will be redesigned in a later patch that introduces full
-// polynomial behaviour, i.e. the ability to have composites made up of both
-// fixed and scalable sizes.
-template <typename T> class PolySize {
+/// PolyBase is a base class for ElementCount, TypeSize and StackOffset.
+/// It tries to represent a true polynomial, e.g.
+///  c0 * scale0 + c1 * scale1 + ... + cK * scaleK
+/// where the scale is implicit. Only the coefficients are encoded.
+///
+/// <typename T> is the type of the coefficients.
+/// <unsigned Dims> are the number of dimensions of the polynomial.
+/// <bool DimIsExclusive> lets PolyBase represent a polynomial with coefficients
+/// for only a single dimension when true, or multiple coefficients when false.
+template <typename T, unsigned Dims, bool DimIsExclusive> class PolyBase {
 protected:
-  T MinVal;        // The minimum value that it could be.
-  bool IsScalable; // If true, the total value is determined by multiplying
-                   // 'MinVal' by a runtime determinded quantity, 'vscale'.
+  T Coefficients[Dims];
+  unsigned ExclusiveDim;
 
-  constexpr PolySize(T MinVal, bool IsScalable)
-      : MinVal(MinVal), IsScalable(IsScalable) {}
+  // Create a PolyBase for multiple dimensions.
+  template <bool D = DimIsExclusive>
+  PolyBase(ArrayRef<typename std::enable_if<!D, T>::type> Values)
+      : ExclusiveDim(~0) {
+    assert(Values.size() == Dims && "Incorrect number of values");
+    for (unsigned I = 0; I < Dims; ++I)
+      Coefficients[I] = Values[I];
+    verify();
+  }
+
+  // Create a PolyBase for a single dimension.
+  template <bool D = DimIsExclusive>
+  PolyBase(const typename std::enable_if<(D), T>::type &Val,
+           unsigned ExclusiveDim)
+      : ExclusiveDim(ExclusiveDim) {
+    assert(ExclusiveDim < Dims && "Incorrect dimension");
+    for (unsigned I = 0; I < Dims; ++I)
+      Coefficients[I] = (I == ExclusiveDim ? Val : 0);
+    verify();
+  }
+
+  void verify() const {
+#ifndef NDEBUG
+    assert((!DimIsExclusive ||
+            llvm::count_if(ArrayRef<T>(Coefficients, Dims),
+                           [](const T &C) { return C > 0; }) <= 1) &&
+           "More than one dimensions set for Mutually Exclusive PolyBase");
+#endif
+  }
 
 public:
+  bool operator==(const PolyBase &RHS) const {
+    return std::equal(&Coefficients[0], &Coefficients[Dims],
+                      &RHS.Coefficients[0]);
+  }
+
+  bool operator!=(const PolyBase &RHS) const { return !(*this == RHS); }
 
-  static constexpr PolySize getFixed(T MinVal) { return {MinVal, false}; }
-  static constexpr PolySize getScalable(T MinVal) { return {MinVal, true}; }
-  static constexpr PolySize get(T MinVal, bool IsScalable) {
-    return {MinVal, IsScalable};
+  PolyBase &operator+=(const PolyBase &RHS) {
+    for (unsigned I = 0; I < Dims; ++I)
+      Coefficients[I] += RHS.Coefficients[I];
+    verify();
+    return *this;
   }
 
-  static constexpr PolySize getNull() { return {0, false}; }
+  PolyBase &operator-=(const PolyBase &RHS) {
+    for (unsigned I = 0; I < Dims; ++I)
+      Coefficients[I] -= RHS.Coefficients[I];
+    verify();
+    return *this;
+  }
 
-  /// Counting predicates.
-  ///
-  ///@{ No elements..
-  bool isZero() const { return MinVal == 0; }
-  /// At least one element.
-  bool isNonZero() const { return MinVal != 0; }
+  PolyBase &operator*=(T RHS) {
+    for (unsigned I = 0; I < Dims; ++I)
+      Coefficients[I] *= RHS;
+    return *this;
+  }
+
+  PolyBase operator+(const PolyBase &RHS) const {
+    PolyBase This = *this;
+    return This += RHS;
+  }
+
+  PolyBase operator-(const PolyBase &RHS) const {
+    PolyBase This = *this;
+    return This -= RHS;
+  }
+
+  PolyBase operator*(T RHS) const {
+    PolyBase This = *this;
+    return This *= RHS;
+  }
+
+  template <typename U = T>
+  typename std::enable_if<std::is_signed<U>::value, PolyBase>::type
+  operator-() const {
+    PolyBase This = *this;
+    return This *= (T)-1;
+  }
+
+  bool isZero() const {
+    return llvm::all_of(ArrayRef<T>(Coefficients, Dims),
+                        [](const T &C) { return C == 0; });
+  }
+
+  bool isNonZero() const { return !isZero(); }
+
+  explicit operator bool() const { return isNonZero(); }
+
+  T getValue(unsigned Dim) const { return Coefficients[Dim]; }
+
+  template <bool D = DimIsExclusive>
+  typename std::enable_if<(D), T>::type getExclusiveValue() const {
+    return Coefficients[ExclusiveDim];
+  }
+};
+
+
+template <typename T>
+using MixedPolyBase2DSuper_t =
+    PolyBase<T, /*Dims=*/2, /*DimIsExclusive=*/false>;
+
+/// MixedPolyBase2D is a base class to represent 2 dimensions, named
+/// fixed and scalable, respectively. This class allows a value for both
+/// dimensions to depict e.g. "8 bytes and 16 scalable bytes", which is needed
+/// to represent stack offsets.
+template <typename T> class MixedPolyBase2D : public MixedPolyBase2DSuper_t<T> {
+protected:
+  MixedPolyBase2D(T Fixed, T Scalable)
+      : MixedPolyBase2DSuper_t<T>({Fixed, Scalable}) {}
+
+public:
+  MixedPolyBase2D() : MixedPolyBase2DSuper_t<T>({0, 0}) {}
+  MixedPolyBase2D(const MixedPolyBase2DSuper_t<T> &Other)
+      : MixedPolyBase2DSuper_t<T>(Other) {}
+  static MixedPolyBase2D getFixed(T Fixed) { return {Fixed, 0}; }
+  static MixedPolyBase2D getScalable(T Scalable) { return {0, Scalable}; }
+  static MixedPolyBase2D get(T Fixed, T Scalable) { return {Fixed, Scalable}; }
+
+  T getFixed() const { return this->getValue(0); }
+  T getScalable() const { return this->getValue(1); }
+  std::pair<T, T> getMixed() const {
+    return {this->getValue(0), this->getValue(1)};
+  }
+};
+
+template <typename T>
+using SinglePolyBase2DSuper_t =
+    PolyBase<T, /*Dims=*/2, /*DimIsExclusive=*/true>;
+
+/// SinglePolyBase2D is a base class to represent 2 dimensions, where the
+/// base class can only represent 1 dimension exclusively at a time, i.e. it is
+/// either fixed-sized or it is scalable-sized, but it cannot be both.
+template <typename T>
+class SinglePolyBase2D : public SinglePolyBase2DSuper_t<T> {
+protected:
+  SinglePolyBase2D(T MinVal, bool IsScalable)
+      : SinglePolyBase2DSuper_t<T>(MinVal, !IsScalable ? 0 : 1) {}
+
+public:
+  SinglePolyBase2D(const SinglePolyBase2DSuper_t<T> &Other)
+      : SinglePolyBase2DSuper_t<T>(Other) {}
+  static SinglePolyBase2D getFixed(T MinVal) { return {MinVal, false}; }
+  static SinglePolyBase2D getScalable(T MinVal) { return {MinVal, true}; }
+  static SinglePolyBase2D get(T MinVal, bool Scalable) {
+    return {MinVal, Scalable};
+  }
+
+  /// Returns the minimum value this size can represent.
+  T getKnownMinValue() const { return this->getExclusiveValue(); }
+  /// Returns whether the size is scaled by a runtime quantity (vscale).
+  bool isScalable() const { return this->ExclusiveDim == 1; }
   /// A return value of true indicates we know at compile time that the number
   /// of elements (vscale * Min) is definitely even. However, returning false
   /// does not guarantee that the total number of elements is odd.
-  bool isKnownEven() const { return (MinVal & 0x1) == 0; }
-  ///@}
-
-  T getKnownMinValue() const { return MinVal; }
+  bool isKnownEven() const { return (getKnownMinValue() & 0x1) == 0; }
+  /// This function tells the caller whether the element count is known at
+  /// compile time to be a multiple of the scalar value RHS.
+  bool isKnownMultipleOf(T RHS) const { return getKnownMinValue() % RHS == 0; }
+  /// Exactly one element.
+  bool isScalar() const { return !isScalable() && getKnownMinValue() == 1; }
+  /// One or more elements.
+  bool isVector() const {
+    return (isScalable() && getKnownMinValue() != 0) || getKnownMinValue() > 1;
+  }
 
   // Return the minimum value with the assumption that the count is exact.
   // Use in places where a scalable count doesn't make sense (e.g. non-vector
   // types, or vectors in backends which don't support scalable vectors).
   T getFixedValue() const {
-    assert(!IsScalable &&
+    assert(!isScalable() &&
            "Request for a fixed element count on a scalable object");
-    return MinVal;
-  }
-
-  bool isScalable() const { return IsScalable; }
-
-  bool operator==(const PolySize &RHS) const {
-    return MinVal == RHS.MinVal && IsScalable == RHS.IsScalable;
+    return getKnownMinValue();
   }
 
-  bool operator!=(const PolySize &RHS) const { return !(*this == RHS); }
-
   // For some cases, size ordering between scalable and fixed size types cannot
   // be determined at compile time, so such comparisons aren't allowed.
   //
   // e.g. <vscale x 2 x i16> could be bigger than <4 x i32> with a runtime
   // vscale >= 5, equal sized with a vscale of 4, and smaller with
   // a vscale <= 3.
   //
   // All the functions below make use of the fact vscale is always >= 1, which
   // means that <vscale x 4 x i32> is guaranteed to be >= <4 x i32>, etc.
 
-  static bool isKnownLT(const PolySize &LHS, const PolySize &RHS) {
-    if (!LHS.IsScalable || RHS.IsScalable)
-      return LHS.MinVal < RHS.MinVal;
-
-    // LHS.IsScalable = true, RHS.IsScalable = false
+  static bool isKnownLT(const SinglePolyBase2D &LHS,
+                        const SinglePolyBase2D &RHS) {
+    if (!LHS.isScalable() || RHS.isScalable())
+      return LHS.getKnownMinValue() < RHS.getKnownMinValue();
     return false;
   }
 
-  static bool isKnownGT(const PolySize &LHS, const PolySize &RHS) {
-    if (LHS.IsScalable || !RHS.IsScalable)
-      return LHS.MinVal > RHS.MinVal;
-
-    // LHS.IsScalable = false, RHS.IsScalable = true
+  static bool isKnownGT(const SinglePolyBase2D &LHS,
+                        const SinglePolyBase2D &RHS) {
+    if (LHS.isScalable() || !RHS.isScalable())
+      return LHS.getKnownMinValue() > RHS.getKnownMinValue();
     return false;
   }
 
-  static bool isKnownLE(const PolySize &LHS, const PolySize &RHS) {
-    if (!LHS.IsScalable || RHS.IsScalable)
-      return LHS.MinVal <= RHS.MinVal;
-
-    // LHS.IsScalable = true, RHS.IsScalable = false
+  static bool isKnownLE(const SinglePolyBase2D &LHS,
+                        const SinglePolyBase2D &RHS) {
+    if (!LHS.isScalable() || RHS.isScalable())
+      return LHS.getKnownMinValue() <= RHS.getKnownMinValue();
     return false;
   }
 
-  static bool isKnownGE(const PolySize &LHS, const PolySize &RHS) {
-    if (LHS.IsScalable || !RHS.IsScalable)
-      return LHS.MinVal >= RHS.MinVal;
-
-    // LHS.IsScalable = false, RHS.IsScalable = true
+  static bool isKnownGE(const SinglePolyBase2D &LHS,
+                        const SinglePolyBase2D &RHS) {
+    if (LHS.isScalable() || !RHS.isScalable())
+      return LHS.getKnownMinValue() >= RHS.getKnownMinValue();
     return false;
   }
 
-  PolySize operator*(T RHS) { return {MinVal * RHS, IsScalable}; }
-
-  PolySize &operator*=(T RHS) {
-    MinVal *= RHS;
-    return *this;
-  }
-
-  friend PolySize operator-(const PolySize &LHS, const PolySize &RHS) {
-    assert(LHS.IsScalable == RHS.IsScalable &&
-           "Arithmetic using mixed scalable and fixed types");
-    return {LHS.MinVal - RHS.MinVal, LHS.IsScalable};
-  }
-
-  /// This function tells the caller whether the element count is known at
-  /// compile time to be a multiple of the scalar value RHS.
-  bool isKnownMultipleOf(T RHS) const { return MinVal % RHS == 0; }
-
   /// We do not provide the '/' operator here because division for polynomial
   /// types does not work in the same way as for normal integer types. We can
   /// only divide the minimum value (or coefficient) by RHS, which is not the
   /// same as
   ///   (Min * Vscale) / RHS
   /// The caller is recommended to use this function in combination with
   /// isKnownMultipleOf(RHS), which lets the caller know if it's possible to
   /// perform a lossless divide by RHS.
-  PolySize divideCoefficientBy(T RHS) const {
-    return PolySize(MinVal / RHS, IsScalable);
+  SinglePolyBase2D divideCoefficientBy(T RHS) const {
+    return SinglePolyBase2D(getKnownMinValue() / RHS, isScalable());
   }
 
-  PolySize coefficientNextPowerOf2() const {
-    return PolySize((T)llvm::NextPowerOf2(MinVal), IsScalable);
+  SinglePolyBase2D coefficientNextPowerOf2() const {
+    return SinglePolyBase2D((T)llvm::NextPowerOf2(getKnownMinValue()),
+                            isScalable());
   }
 
   /// Printing function.
   void print(raw_ostream &OS) const {
-    if (IsScalable)
+    if (isScalable())
       OS << "vscale x ";
-    OS << MinVal;
+    OS << getKnownMinValue();
   }
 };
 
-/// Stream operator function for `PolySize`.
+/// Stream operator function for `SinglePolyBase2D`.
 template <typename T>
-inline raw_ostream &operator<<(raw_ostream &OS, const PolySize<T> &PS) {
+inline raw_ostream &operator<<(raw_ostream &OS, const SinglePolyBase2D<T> &PS) {
   PS.print(OS);
   return OS;
 }
 
-class ElementCount : public PolySize<unsigned> {
-public:
+using ElementCount = SinglePolyBase2D<unsigned>;
+using StackOffset = MixedPolyBase2D<int64_t>;
 
-  constexpr ElementCount(PolySize<unsigned> V) : PolySize(V) {}
-
-  /// Counting predicates.
-  ///
-  /// Notice that MinVal = 1 and IsScalable = true is considered more than
-  /// one element.
-  ///
-  ///@{ No elements..
-  /// Exactly one element.
-  bool isScalar() const { return !IsScalable && MinVal == 1; }
-  /// One or more elements.
-  bool isVector() const { return (IsScalable && MinVal != 0) || MinVal > 1; }
-  ///@}
-};
-
-// This class is used to represent the size of types. If the type is of fixed
+// TODO: Most functionality in this class will gradually be phased out
+// so it will resemble SinglePolyBase2D<uint64_t> as much as possible.
+//
+// TypeSize is used to represent the size of types. If the type is of fixed
 // size, it will represent the exact size. If the type is a scalable vector,
 // it will represent the known minimum size.
-class TypeSize : public PolySize<uint64_t> {
+class TypeSize : public SinglePolyBase2D<uint64_t> {
 public:
-  constexpr TypeSize(PolySize<uint64_t> V) : PolySize(V) {}
+  TypeSize(const SinglePolyBase2D<uint64_t> &V) : SinglePolyBase2D(V) {}
 
-  constexpr TypeSize(uint64_t MinVal, bool IsScalable)
-      : PolySize(MinVal, IsScalable) {}
+  TypeSize(uint64_t MinVal, bool IsScalable)
+      : SinglePolyBase2D(MinVal, IsScalable) {}
 
-  static constexpr TypeSize Fixed(uint64_t MinVal) {
-    return TypeSize(MinVal, false);
-  }
-  static constexpr TypeSize Scalable(uint64_t MinVal) {
-    return TypeSize(MinVal, true);
-  }
+  static TypeSize Fixed(uint64_t MinVal) { return TypeSize(MinVal, false); }
+  static TypeSize Scalable(uint64_t MinVal) { return TypeSize(MinVal, true); }
 
   uint64_t getFixedSize() const { return getFixedValue(); }
   uint64_t getKnownMinSize() const { return getKnownMinValue(); }
 
+  // The comparison operators are in the process of being phased out
+  // in favour of isKnownLT/isKnownLE of its parent class.
+
   friend bool operator<(const TypeSize &LHS, const TypeSize &RHS) {
-    assert(LHS.IsScalable == RHS.IsScalable &&
+    assert(LHS.isScalable() == RHS.isScalable() &&
            "Ordering comparison of scalable and fixed types");
 
-    return LHS.MinVal < RHS.MinVal;
+    return LHS.getKnownMinValue() < RHS.getKnownMinValue();
   }
 
   friend bool operator>(const TypeSize &LHS, const TypeSize &RHS) {
     return RHS < LHS;
   }
 
   friend bool operator<=(const TypeSize &LHS, const TypeSize &RHS) {
     return !(RHS < LHS);
   }
 
-  friend bool operator>=(const TypeSize &LHS, const TypeSize& RHS) {
+  friend bool operator>=(const TypeSize &LHS, const TypeSize &RHS) {
     return !(LHS < RHS);
   }
 
-  TypeSize &operator-=(TypeSize RHS) {
-    assert(IsScalable == RHS.IsScalable &&
-           "Subtraction using mixed scalable and fixed types");
-    MinVal -= RHS.MinVal;
-    return *this;
-  }
-
-  TypeSize &operator+=(TypeSize RHS) {
-    assert(IsScalable == RHS.IsScalable &&
-           "Addition using mixed scalable and fixed types");
-    MinVal += RHS.MinVal;
-    return *this;
-  }
-
-  friend TypeSize operator-(const TypeSize &LHS, const TypeSize &RHS) {
-    assert(LHS.IsScalable == RHS.IsScalable &&
-           "Arithmetic using mixed scalable and fixed types");
-    return {LHS.MinVal - RHS.MinVal, LHS.IsScalable};
-  }
+  // All code for this class below this point is needed because of the
+  // temporary implicit conversion to uint64_t. The operator overloads are
+  // needed because otherwise the conversion of the parent class PolyBase ->
+  // TypeSize is ambiguous.
+  // TODO: Remove the implicit conversion.
 
   // Casts to a uint64_t if this is a fixed-width size.
   //
   // This interface is deprecated and will be removed in a future version
   // of LLVM in favour of upgrading uses that rely on this implicit conversion
   // to uint64_t. Calls to functions that return a TypeSize should use the
   // proper interfaces to TypeSize.
   // In practice this is mostly calls to MVT/EVT::getSizeInBits().
#else
     if (isScalable())
       WithColor::warning() << "Compiler has made implicit assumption that "
                               "TypeSize is not scalable. This may or may not "
                               "lead to broken code.\n";
     return getKnownMinValue();
 #endif
   }
 
-  // Convenience operators to obtain relative sizes independently of
-  // the scalable flag.
-  TypeSize operator*(unsigned RHS) const { return {MinVal * RHS, IsScalable}; }
-
-  friend TypeSize operator*(const unsigned LHS, const TypeSize &RHS) {
-    return {LHS * RHS.MinVal, RHS.IsScalable};
+  // Overload the parent operators.
+  TypeSize &operator+=(const TypeSize &RHS) {
+    return (TypeSize &)SinglePolyBase2D<uint64_t>::operator+=(RHS);
+  }
+  TypeSize &operator-=(const TypeSize &RHS) {
+    return (TypeSize &)SinglePolyBase2D<uint64_t>::operator-=(RHS);
+  }
+  TypeSize &operator*=(const uint64_t &RHS) {
+    return (TypeSize &)SinglePolyBase2D<uint64_t>::operator*=(RHS);
+  }
+  TypeSize operator+(const TypeSize &RHS) const {
+    return (TypeSize)SinglePolyBase2D<uint64_t>::operator+(RHS);
+  }
+  TypeSize operator-(const TypeSize &RHS) const {
+    return (TypeSize)SinglePolyBase2D<uint64_t>::operator-(RHS);
+  }
+  TypeSize operator*(const uint64_t &RHS) const {
+    return (TypeSize)SinglePolyBase2D<uint64_t>::operator*(RHS);
   }
 
-  // Additional convenience operators needed to avoid ambiguous parses.
-  // TODO: Make uint64_t the default operator?
-  TypeSize operator*(uint64_t RHS) const { return {MinVal * RHS, IsScalable}; }
-
-  TypeSize operator*(int RHS) const { return {MinVal * RHS, IsScalable}; }
-
-  TypeSize operator*(int64_t RHS) const { return {MinVal * RHS, IsScalable}; }
-
-  friend TypeSize operator*(const uint64_t LHS, const TypeSize &RHS) {
-    return {LHS * RHS.MinVal, RHS.IsScalable};
+  // Additional operators needed to avoid ambiguous parses
+  // because of the implicit conversion hack.
+  TypeSize operator*(const unsigned &RHS) const {
+    return (TypeSize)SinglePolyBase2D<uint64_t>::operator*(RHS);
+  }
+  TypeSize operator*(const int &RHS) const {
+    return (TypeSize)SinglePolyBase2D<uint64_t>::operator*(RHS);
+  }
+  TypeSize operator*(const int64_t &RHS) const {
+    return (TypeSize)SinglePolyBase2D<uint64_t>::operator*(RHS);
   }
 
   friend TypeSize operator*(const int LHS, const TypeSize &RHS) {
-    return {LHS * RHS.MinVal, RHS.IsScalable};
+    return RHS * LHS;
+  }
+  friend TypeSize operator*(const unsigned LHS, const TypeSize &RHS) {
+    return RHS * LHS;
   }
-
   friend TypeSize operator*(const int64_t LHS, const TypeSize &RHS) {
-    return {LHS * RHS.MinVal, RHS.IsScalable};
+    return RHS * LHS;
+  }
+  friend TypeSize operator*(const uint64_t LHS, const TypeSize &RHS) {
+    return RHS * LHS;
   }
 };
 
 /// Returns a TypeSize with a known minimum size that is the next integer
 /// (mod 2**64) that is greater than or equal to \p Value and is a multiple
 /// of \p Align. \p Align must be non-zero.
 ///
 /// Similar to the alignTo functions in MathExtras.h
 inline TypeSize alignTo(TypeSize Size, uint64_t Align) {
   assert(Align != 0u && "Align must be non-zero");
   return {(Size.getKnownMinValue() + Align - 1) / Align * Align,
           Size.isScalable()};
 }
 
 template <> struct DenseMapInfo<ElementCount> {
   static inline ElementCount getEmptyKey() {
     return ElementCount::getScalable(~0U);
   }
   static inline ElementCount getTombstoneKey() {
     return ElementCount::getFixed(~0U - 1);
   }
-  static unsigned getHashValue(const ElementCount& EltCnt) {
+  static unsigned getHashValue(const ElementCount &EltCnt) {
     unsigned HashVal = EltCnt.getKnownMinValue() * 37U;
     if (EltCnt.isScalable())
       return (HashVal - 1U);
 
     return HashVal;
   }
 
-  static bool isEqual(const ElementCount& LHS, const ElementCount& RHS) {
+  static bool isEqual(const ElementCount &LHS, const ElementCount &RHS) {
     return LHS == RHS;
   }
 };
 
 } // end namespace llvm
 
 #endif // LLVM_SUPPORT_TypeSize_H

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 #define MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 
 #include "mlir/IR/Types.h"
 
 namespace llvm {
-class ElementCount;
+template <typename T>
+class SinglePolyBase2D;
+using ElementCount = SinglePolyBase2D<unsigned>;
 class TypeSize;
 } // namespace llvm
 
 namespace mlir {
 
 class DialectAsmParser;
 class DialectAsmPrinter;