[llvm] [DataLayout] Change return type of `getStackAlignment` to `MaybeAlign` (PR #105478)

[Sergei Barannikov via llvm-commits]

https://github.com/s-barannikov created https://github.com/llvm/llvm-project/pull/105478

Currently, `getStackAlignment` asserts if the stack alignment wasn't specified. This makes it inconvenient to use and complicates testing.

This change also makes `exceedsNaturalStackAlignment` method redundant.

Depends on #105477.


>From ca9e2658c3c3f4d3abf525d6b8ce653a09e22eef Mon Sep 17 00:00:00 2001
From: Sergei Barannikov <barannikov88 at gmail.com>
Date: Wed, 21 Aug 2024 10:57:45 +0300
Subject: [PATCH] [DataLayout] Change return type of `getStackAlignment` to
 `MaybeAlign`

Currently, `getStackAlignment` asserts if the stack alignment wasn't
specified. This makes it inconvenient to use and complicates testing.

This change also makes `exceedsNaturalStackAlignment` method redundant.
---
 llvm/include/llvm/IR/DataLayout.h           | 12 +++---------
 llvm/lib/Transforms/IPO/ExpandVariadics.cpp |  6 +++---
 llvm/lib/Transforms/Utils/Local.cpp         |  3 ++-
 llvm/unittests/IR/DataLayoutTest.cpp        | 15 +++++++++++++++
 4 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/IR/DataLayout.h b/llvm/include/llvm/IR/DataLayout.h
index 2f06bda6c30a51..145f1a29c7dfb7 100644
--- a/llvm/include/llvm/IR/DataLayout.h
+++ b/llvm/include/llvm/IR/DataLayout.h
@@ -220,15 +220,9 @@ class DataLayout {
 
   bool isIllegalInteger(uint64_t Width) const { return !isLegalInteger(Width); }
 
-  /// Returns true if the given alignment exceeds the natural stack alignment.
-  bool exceedsNaturalStackAlignment(Align Alignment) const {
-    return StackNaturalAlign && (Alignment > *StackNaturalAlign);
-  }
-
-  Align getStackAlignment() const {
-    assert(StackNaturalAlign && "StackNaturalAlign must be defined");
-    return *StackNaturalAlign;
-  }
+  /// Returns the natural stack alignment, or MaybeAlign() if one wasn't
+  /// specified.
+  MaybeAlign getStackAlignment() const { return StackNaturalAlign; }
 
   unsigned getAllocaAddrSpace() const { return AllocaAddrSpace; }
 
diff --git a/llvm/lib/Transforms/IPO/ExpandVariadics.cpp b/llvm/lib/Transforms/IPO/ExpandVariadics.cpp
index 49bfec297bc173..a7a01ca1055dd3 100644
--- a/llvm/lib/Transforms/IPO/ExpandVariadics.cpp
+++ b/llvm/lib/Transforms/IPO/ExpandVariadics.cpp
@@ -748,10 +748,10 @@ bool ExpandVariadics::expandCall(Module &M, IRBuilder<> &Builder, CallBase *CB,
   // This is an awkward way to guess whether there is a known stack alignment
   // without hitting an assert in DL.getStackAlignment, 1024 is an arbitrary
   // number likely to be greater than the natural stack alignment.
-  // TODO: DL.getStackAlignment could return a MaybeAlign instead of assert
   Align AllocaAlign = MaxFieldAlign;
-  if (DL.exceedsNaturalStackAlignment(Align(1024)))
-    AllocaAlign = std::max(AllocaAlign, DL.getStackAlignment());
+  if (MaybeAlign StackAlign = DL.getStackAlignment();
+      StackAlign && *StackAlign > AllocaAlign)
+    AllocaAlign = *StackAlign;
 
   // Put the alloca to hold the variadic args in the entry basic block.
   Builder.SetInsertPointPastAllocas(CBF);
diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index efb02fdec56d7e..df0924dd8e4208 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -1506,7 +1506,8 @@ Align llvm::tryEnforceAlignment(Value *V, Align PrefAlign,
 
     // If the preferred alignment is greater than the natural stack alignment
     // then don't round up. This avoids dynamic stack realignment.
-    if (DL.exceedsNaturalStackAlignment(PrefAlign))
+    MaybeAlign StackAlign = DL.getStackAlignment();
+    if (StackAlign && PrefAlign > *StackAlign)
       return CurrentAlign;
     AI->setAlignment(PrefAlign);
     return PrefAlign;
diff --git a/llvm/unittests/IR/DataLayoutTest.cpp b/llvm/unittests/IR/DataLayoutTest.cpp
index 396d44af19f53f..16a603ff6416f4 100644
--- a/llvm/unittests/IR/DataLayoutTest.cpp
+++ b/llvm/unittests/IR/DataLayoutTest.cpp
@@ -444,6 +444,21 @@ TEST(DataLayout, ParseNonIntegralAddrSpace) {
         FailedWithMessage("address space 0 cannot be non-integral"));
 }
 
+TEST(DataLayout, GetStackAlignment) {
+  DataLayout Default;
+  EXPECT_FALSE(Default.getStackAlignment().has_value());
+
+  std::pair<StringRef, Align> Cases[] = {
+      {"S8", Align(1)},
+      {"S64", Align(8)},
+      {"S32768", Align(4096)},
+  };
+  for (auto [Layout, Val] : Cases) {
+    DataLayout DL = cantFail(DataLayout::parse(Layout));
+    EXPECT_EQ(DL.getStackAlignment(), Val) << Layout;
+  }
+}
+
 TEST(DataLayout, GetPointerSizeInBits) {
   std::tuple<StringRef, unsigned, unsigned, unsigned> Cases[] = {
       {"", 64, 64, 64},