[llvm] [llvm][TypeSize] Consider TypeSize of '0' to be fixed/scalable-agnostic. (PR #72994)

[Sander de Smalen via llvm-commits]

https://github.com/sdesmalen-arm updated https://github.com/llvm/llvm-project/pull/72994

>From c342044da68a9090141e3618a2d51bb50124b930 Mon Sep 17 00:00:00 2001
From: Sander de Smalen <sander.desmalen at arm.com>
Date: Tue, 21 Nov 2023 13:51:09 +0000
Subject: [PATCH] [llvm][TypeSize] Consider TypeSize of '0' to be
 fixed/scalable-agnostic.

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.
---
 llvm/include/llvm/Support/TypeSize.h    | 20 +++++++++++++++++---
 llvm/unittests/Support/TypeSizeTest.cpp | 15 +++++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/Support/TypeSize.h b/llvm/include/llvm/Support/TypeSize.h
index ada98d809fc236f..988c62dae66dfb7 100644
--- a/llvm/include/llvm/Support/TypeSize.h
+++ b/llvm/include/llvm/Support/TypeSize.h
@@ -97,17 +97,29 @@ template <typename LeafTy, typename ValueTy> class FixedOrScalableQuantity {
 
   constexpr FixedOrScalableQuantity() = default;
   constexpr FixedOrScalableQuantity(ScalarTy Quantity, bool Scalable)
-      : Quantity(Quantity), Scalable(Scalable) {}
+      : Quantity(Quantity), Scalable(Quantity ? Scalable : false) {}
 
   friend constexpr LeafTy &operator+=(LeafTy &LHS, const LeafTy &RHS) {
-    assert(LHS.Scalable == RHS.Scalable && "Incompatible types");
+    assert((LHS.Quantity == 0 || RHS.Quantity == 0 ||
+            LHS.Scalable == RHS.Scalable) &&
+           "Incompatible types");
     LHS.Quantity += RHS.Quantity;
+    if (!LHS.Quantity)
+      LHS.Scalable = false;
+    else if (RHS.Quantity)
+      LHS.Scalable = RHS.Scalable;
     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.Scalable == RHS.Scalable) &&
+           "Incompatible types");
     LHS.Quantity -= RHS.Quantity;
+    if (!LHS.Quantity)
+      LHS.Scalable = false;
+    else if (RHS.Quantity)
+      LHS.Scalable = RHS.Scalable;
     return LHS;
   }
 
@@ -315,6 +327,8 @@ class TypeSize : public details::FixedOrScalableQuantity<TypeSize, uint64_t> {
       : FixedOrScalableQuantity(V) {}
 
 public:
+  constexpr TypeSize() : FixedOrScalableQuantity(0, false) {}
+
   constexpr TypeSize(ScalarTy Quantity, bool Scalable)
       : FixedOrScalableQuantity(Quantity, Scalable) {}
 
diff --git a/llvm/unittests/Support/TypeSizeTest.cpp b/llvm/unittests/Support/TypeSizeTest.cpp
index 33169a3d8b198c0..f1cdd8e22779e8c 100644
--- a/llvm/unittests/Support/TypeSizeTest.cpp
+++ b/llvm/unittests/Support/TypeSizeTest.cpp
@@ -81,6 +81,21 @@ static_assert(INT64_C(2) * TSFixed32 == TypeSize::getFixed(64));
 static_assert(UINT64_C(2) * TSFixed32 == TypeSize::getFixed(64));
 static_assert(alignTo(TypeSize::getFixed(7), 8) == TypeSize::getFixed(8));
 
+static_assert(TypeSize() == TypeSize::getFixed(0));
+static_assert(TypeSize() == TypeSize::getScalable(0));
+static_assert(TypeSize::getFixed(0) == TypeSize::getScalable(0));
+static_assert(TypeSize::getFixed(0) ==
+              (TypeSize::getScalable(4) - TypeSize::getScalable(4)));
+static_assert(!TypeSize().isScalable());
+static_assert(TypeSize::getFixed(0) + TypeSize::getScalable(8) ==
+              TypeSize::getScalable(8));
+static_assert(TypeSize::getScalable(8) + TypeSize::getFixed(0) ==
+              TypeSize::getScalable(8));
+static_assert(TypeSize::getFixed(8) + TypeSize::getScalable(0) ==
+              TypeSize::getFixed(8));
+static_assert(TypeSize::getScalable(0) + TypeSize::getFixed(8) ==
+              TypeSize::getFixed(8));
+
 TEST(TypeSize, FailIncompatibleTypes) {
   EXPECT_DEBUG_DEATH(TypeSize::getFixed(8) + TypeSize::getScalable(8),
                      "Incompatible types");