[llvm][TypeSize] Consider TypeSize of '0' to be fixed/scalable-agnostic.

[sdesmalen-arm]

This patch allows adding any quantity to a zero-initialized TypeSize, such
that e.g.:

TypeSize::Scalable(0) + TypeSize::Fixed(4) == TypeSize::Fixed(4)
TypeSize::Fixed(0) + TypeSize::Scalable(4) == TypeSize::Scalable(4)

This makes it easier to implement add-reductions using TypeSize where
the 'scalable' flag is not yet known before starting the reduction.

(this PR follows on from #72979)

[llvmbot]

@llvm/pr-subscribers-llvm-ir
@llvm/pr-subscribers-llvm-selectiondag
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-coroutines
@llvm/pr-subscribers-debuginfo
@llvm/pr-subscribers-backend-directx
@llvm/pr-subscribers-llvm-globalisel
@llvm/pr-subscribers-backend-webassembly
@llvm/pr-subscribers-backend-risc-v
@llvm/pr-subscribers-backend-aarch64
@llvm/pr-subscribers-backend-amdgpu
@llvm/pr-subscribers-clang
@llvm/pr-subscribers-backend-arm

@llvm/pr-subscribers-llvm-support

Author: Sander de Smalen (sdesmalen-arm)

Changes

This patch allows adding any quantity to a zero-initialized TypeSize, such
that e.g.:

TypeSize::Scalable(0) + TypeSize::Fixed(4) == TypeSize::Fixed(4)
TypeSize::Fixed(0) + TypeSize::Scalable(4) == TypeSize::Scalable(4)

This makes it easier to implement add-reductions using TypeSize where
the 'scalable' flag is not yet known before starting the reduction.

(this PR follows on from #72979)

---

Full diff: https://github.com/llvm/llvm-project/pull/72994.diff

2 Files Affected:

• (modified) llvm/include/llvm/Support/TypeSize.h (+25-17)
• (modified) llvm/unittests/Support/TypeSizeTest.cpp (+15)

diff --git a/llvm/include/llvm/Support/TypeSize.h b/llvm/include/llvm/Support/TypeSize.h
index 8e638f8278e828b..34c6c0ebfef25de 100644
--- a/llvm/include/llvm/Support/TypeSize.h
+++ b/llvm/include/llvm/Support/TypeSize.h
@@ -93,20 +93,26 @@ template <typename LeafTy, typename ValueTy> class FixedOrScalableQuantity {
 
 protected:
   ScalarTy Quantity = 0;
-  bool Scalable = false;
+  bool IsScalable = false;
 
   constexpr FixedOrScalableQuantity() = default;
-  constexpr FixedOrScalableQuantity(ScalarTy Quantity, bool Scalable)
-      : Quantity(Quantity), Scalable(Scalable) {}
+  constexpr FixedOrScalableQuantity(ScalarTy Quantity, bool IsScalable)
+      : Quantity(Quantity), IsScalable(IsScalable) {}
 
   friend constexpr LeafTy &operator+=(LeafTy &LHS, const LeafTy &RHS) {
-    assert(LHS.Scalable == RHS.Scalable && "Incompatible types");
+    assert((LHS.Quantity == 0 || RHS.Quantity == 0 ||
+            LHS.IsScalable == RHS.IsScalable) &&
+           "Incompatible types");
+    LHS.IsScalable = LHS.Quantity ? LHS.IsScalable : RHS.IsScalable;
     LHS.Quantity += RHS.Quantity;
     return LHS;
   }
 
   friend constexpr LeafTy &operator-=(LeafTy &LHS, const LeafTy &RHS) {
-    assert(LHS.Scalable == RHS.Scalable && "Incompatible types");
+    assert((LHS.Quantity == 0 || RHS.Quantity == 0 ||
+            LHS.IsScalable == RHS.IsScalable) &&
+           "Incompatible types");
+    LHS.IsScalable = LHS.Quantity ? LHS.IsScalable : RHS.IsScalable;
     LHS.Quantity -= RHS.Quantity;
     return LHS;
   }
@@ -140,11 +146,11 @@ template <typename LeafTy, typename ValueTy> class FixedOrScalableQuantity {
 
 public:
   constexpr bool operator==(const FixedOrScalableQuantity &RHS) const {
-    return Quantity == RHS.Quantity && Scalable == RHS.Scalable;
+    return Quantity == RHS.Quantity && IsScalable == RHS.IsScalable;
   }
 
   constexpr bool operator!=(const FixedOrScalableQuantity &RHS) const {
-    return Quantity != RHS.Quantity || Scalable != RHS.Scalable;
+    return Quantity != RHS.Quantity || IsScalable != RHS.IsScalable;
   }
 
   constexpr bool isZero() const { return Quantity == 0; }
@@ -155,14 +161,14 @@ template <typename LeafTy, typename ValueTy> class FixedOrScalableQuantity {
 
   /// Add \p RHS to the underlying quantity.
   constexpr LeafTy getWithIncrement(ScalarTy RHS) const {
-    return LeafTy::get(Quantity + RHS, Scalable);
+    return LeafTy::get(Quantity + RHS, IsScalable);
   }
 
   /// Returns the minimum value this quantity can represent.
   constexpr ScalarTy getKnownMinValue() const { return Quantity; }
 
   /// Returns whether the quantity is scaled by a runtime quantity (vscale).
-  constexpr bool isScalable() const { return Scalable; }
+  constexpr bool isScalable() const { return IsScalable; }
 
   /// A return value of true indicates we know at compile time that the number
   /// of elements (vscale * Min) is definitely even. However, returning false
@@ -277,8 +283,8 @@ template <typename LeafTy, typename ValueTy> class FixedOrScalableQuantity {
 //  - ElementCount::getScalable(4) : A scalable vector type holding 4 values.
 class ElementCount
     : public details::FixedOrScalableQuantity<ElementCount, unsigned> {
-  constexpr ElementCount(ScalarTy MinVal, bool Scalable)
-      : FixedOrScalableQuantity(MinVal, Scalable) {}
+  constexpr ElementCount(ScalarTy MinVal, bool IsScalable)
+      : FixedOrScalableQuantity(MinVal, IsScalable) {}
 
   constexpr ElementCount(
       const FixedOrScalableQuantity<ElementCount, unsigned> &V)
@@ -293,8 +299,8 @@ class ElementCount
   static constexpr ElementCount getScalable(ScalarTy MinVal) {
     return ElementCount(MinVal, true);
   }
-  static constexpr ElementCount get(ScalarTy MinVal, bool Scalable) {
-    return ElementCount(MinVal, Scalable);
+  static constexpr ElementCount get(ScalarTy MinVal, bool IsScalable) {
+    return ElementCount(MinVal, IsScalable);
   }
 
   /// Exactly one element.
@@ -315,11 +321,13 @@ class TypeSize : public details::FixedOrScalableQuantity<TypeSize, uint64_t> {
       : FixedOrScalableQuantity(V) {}
 
 public:
-  constexpr TypeSize(ScalarTy Quantity, bool Scalable)
-      : FixedOrScalableQuantity(Quantity, Scalable) {}
+  constexpr TypeSize() : FixedOrScalableQuantity(0, false) {}
 
-  static constexpr TypeSize get(ScalarTy Quantity, bool Scalable) {
-    return TypeSize(Quantity, Scalable);
+  constexpr TypeSize(ScalarTy Quantity, bool IsScalable)
+      : FixedOrScalableQuantity(Quantity, IsScalable) {}
+
+  static constexpr TypeSize get(ScalarTy Quantity, bool IsScalable) {
+    return TypeSize(Quantity, IsScalable);
   }
   static constexpr TypeSize Fixed(ScalarTy ExactSize) {
     return TypeSize(ExactSize, false);
diff --git a/llvm/unittests/Support/TypeSizeTest.cpp b/llvm/unittests/Support/TypeSizeTest.cpp
index 2850705d39f69f5..0331b4c00d5f1b1 100644
--- a/llvm/unittests/Support/TypeSizeTest.cpp
+++ b/llvm/unittests/Support/TypeSizeTest.cpp
@@ -81,4 +81,19 @@ static_assert(INT64_C(2) * TSFixed32 == TypeSize::Fixed(64));
 static_assert(UINT64_C(2) * TSFixed32 == TypeSize::Fixed(64));
 static_assert(alignTo(TypeSize::Fixed(7), 8) == TypeSize::Fixed(8));
 
+static_assert(TypeSize() == TypeSize::Fixed(0));
+static_assert(TypeSize() != TypeSize::Scalable(0));
+static_assert(!TypeSize().isScalable());
+static_assert(TypeSize::Fixed(0) + TypeSize::Scalable(8) ==
+              TypeSize::Scalable(8));
+static_assert(TypeSize::Scalable(8) + TypeSize::Fixed(0) ==
+              TypeSize::Scalable(8));
+static_assert(TypeSize::Fixed(8) + TypeSize::Scalable(0) == TypeSize::Fixed(8));
+static_assert(TypeSize::Scalable(0) + TypeSize::Fixed(8) == TypeSize::Fixed(8));
+
+TEST(TypeSize, FailIncompatibleTypes) {
+  EXPECT_DEBUG_DEATH(TypeSize::Fixed(8) + TypeSize::Scalable(8),
+                     "Incompatible types");
+}
+
 } // namespace

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[michaelmaitland]

TypeSize::Scalable(0) + TypeSize::Fixed(4) == TypeSize::Fixed(4)

From the docs

The number of elements is a constant integer value larger than 0; elementtype may be any integer, floating-point or pointer type. Vectors of size zero are not allowed. For scalable vectors, the total number of elements is a constant multiple (called vscale) of the specified number of elements; vscale is a positive integer that is unknown at compile time and the same hardware-dependent constant for all scalable vectors at run time. The size of a specific scalable vector type is thus constant within IR, even if the exact size in bytes cannot be determined until run time.

Creating a scalable vector of size 0 should not even be allowed.

[sdesmalen-arm]

Creating a scalable vector of size 0 should not even be allowed.

Correct, but that is no reason to disallow a value of '0' for TypeSize. This is just a class to represent a value (fixed or scalable), and it's up to the places where TypeSize is used on whether that use is valid. You could argue the same for an int, which is used to represent the size of a type.

FWIW, one of the uses for a zero-sized TypeSize value (and also the motivating reason for this patch) is e.g. to get the combined size of spills for N vectors, you could do:

TypeSize Size();
for (auto Obj : Objects)
  Size += Obj.getSize();

without having to determine before the loop whether Size needs to be Fixed or Scalable.

[michaelmaitland]

Correct, but that is no reason to disallow a value of '0' for TypeSize.

Fair enough!

[sdesmalen-arm]

(rebased patch after landing #72979)

[paulwalker-arm]

As discussed, this might be a little dangerous right now because it'll cause functions like isScalable to return false when given TypeSize::getScalable(0). There's also other operators like * which can produce zero without going via this constructor.

It might be better to ensure all the interfaces return canonical results for fixed and scalable zero rather than changing the way zero is stored.

Whilst canonicalisation is the correct way forward if you don't specifically need it at this time then it's best removed and done in isolation.

[sdesmalen-arm]

Thanks, I've removed the canonicalisation now.

[paulwalker-arm]

I think there's a hole in the testing but otherwise looks good.

[paulwalker-arm]

It's worth repeating this test for - to exercise the change to operator-=.