[llvm] MathExtras: template'ize alignToPowerOf2 (PR #97814)

[Ramkumar Ramachandra via llvm-commits]

https://github.com/artagnon created https://github.com/llvm/llvm-project/pull/97814

Follow up on 5627794 (MathExtras: avoid unnecessarily widening types) to change the overflow behavior of alignToPowerOf2 to only overflow if the result is not representable in the return type. This allows us to template'ize it, and avoid unnecessarily widening the types of arguments.

>From 8ef407dded3941f1037dbc4ff26d4718ffcfea0c Mon Sep 17 00:00:00 2001
From: Ramkumar Ramachandra <ramkumar.ramachandra at codasip.com>
Date: Fri, 5 Jul 2024 12:29:30 +0100
Subject: [PATCH] MathExtras: template'ize alignToPowerOf2

Follow up on 5627794 (MathExtras: avoid unnecessarily widening types) to
change the overflow behavior of alignToPowerOf2 to only overflow if the
result is not representable in the return type. This allows us to
template'ize it, and avoid unnecessarily widening the types of
arguments.
---
 llvm/include/llvm/Support/MathExtras.h    | 16 ++++++++++++----
 llvm/unittests/Support/MathExtrasTest.cpp |  7 +++++--
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/Support/MathExtras.h b/llvm/include/llvm/Support/MathExtras.h
index f6b1fdb6aba9e..81a9d16909cdc 100644
--- a/llvm/include/llvm/Support/MathExtras.h
+++ b/llvm/include/llvm/Support/MathExtras.h
@@ -488,13 +488,21 @@ constexpr uint64_t alignTo(uint64_t Value, uint64_t Align) {
   return CeilDiv * Align;
 }
 
+/// Will overflow only if result is not representable in T.
+template <typename U, typename V, typename T = common_uint<U, V>>
+constexpr T alignToPowerOf2(U Value, V Align) {
+  assert(Align != 0 && (Align & (Align - 1)) == 0 &&
+         "Align must be a power of 2");
+  T CeilDiv = divideCeil(Value, Align);
+  return CeilDiv * Align;
+}
+
+/// Fallback when arguments aren't integral.
 constexpr uint64_t alignToPowerOf2(uint64_t Value, uint64_t Align) {
   assert(Align != 0 && (Align & (Align - 1)) == 0 &&
          "Align must be a power of 2");
-  // Replace unary minus to avoid compilation error on Windows:
-  // "unary minus operator applied to unsigned type, result still unsigned"
-  uint64_t NegAlign = (~Align) + 1;
-  return (Value + Align - 1) & NegAlign;
+  uint64_t CeilDiv = divideCeil(Value, Align);
+  return CeilDiv * Align;
 }
 
 /// If non-zero \p Skew is specified, the return value will be a minimal integer
diff --git a/llvm/unittests/Support/MathExtrasTest.cpp b/llvm/unittests/Support/MathExtrasTest.cpp
index a557b61db9752..c0c7c301f1590 100644
--- a/llvm/unittests/Support/MathExtrasTest.cpp
+++ b/llvm/unittests/Support/MathExtrasTest.cpp
@@ -218,8 +218,11 @@ TEST(MathExtras, AlignToPowerOf2) {
   EXPECT_EQ(24u, alignToPowerOf2(17, 8));
   EXPECT_EQ(0u, alignToPowerOf2(~0LL, 8));
   EXPECT_EQ(240u, alignToPowerOf2(240, 16));
-  EXPECT_EQ(static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) + 1,
-            alignToPowerOf2(std::numeric_limits<uint32_t>::max(), 2));
+
+  // Overflow.
+  EXPECT_EQ(0u, alignToPowerOf2(static_cast<uint8_t>(200),
+                                static_cast<uint8_t>(128)));
+  EXPECT_EQ(0u, alignToPowerOf2(std::numeric_limits<uint32_t>::max(), 2));
 }
 
 TEST(MathExtras, AlignDown) {