From 32ebcb8426defc94079cab4f0f017ee683261a58 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Wed, 3 Sep 2025 14:13:26 +0200
Subject: [PATCH] [DataLayout] Specialize the getTypeAllocSize() implementation

getTypeAllocSize() currently works by taking the type store size
and aligning it to the ABI alignment. However, this ends up doing
redundant work in various cases, for example arrays will
unnecessarily repeat the alignment step, and structs will fetch
the StructLayout multiple times.

As this code is rather hot (it is called every time we need to
calculate GEP offsets for example), specialize the implementation.
This repeats a small amount of logic from getAlignment(), but I
think that's worthwhile.
---
 llvm/include/llvm/IR/DataLayout.h             |  5 +--
 llvm/lib/IR/DataLayout.cpp                    | 38 +++++++++++++++++++
 llvm/test/CodeGen/AArch64/alloca-oversized.ll |  3 --
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/llvm/include/llvm/IR/DataLayout.h b/llvm/include/llvm/IR/DataLayout.h
index 2992484c47d06..2acae246c0b1e 100644
--- a/llvm/include/llvm/IR/DataLayout.h
+++ b/llvm/include/llvm/IR/DataLayout.h
@@ -501,10 +501,7 @@ class DataLayout {
   ///
   /// This is the amount that alloca reserves for this type. For example,
   /// returns 12 or 16 for x86_fp80, depending on alignment.
-  TypeSize getTypeAllocSize(Type *Ty) const {
-    // Round up to the next alignment boundary.
-    return alignTo(getTypeStoreSize(Ty), getABITypeAlign(Ty).value());
-  }
+  TypeSize getTypeAllocSize(Type *Ty) const;
 
   /// Returns the offset in bits between successive objects of the
   /// specified type, including alignment padding; always a multiple of 8.
diff --git a/llvm/lib/IR/DataLayout.cpp b/llvm/lib/IR/DataLayout.cpp
index ee43ad49d0df2..2cf96f8720a05 100644
--- a/llvm/lib/IR/DataLayout.cpp
+++ b/llvm/lib/IR/DataLayout.cpp
@@ -844,6 +844,44 @@ Align DataLayout::getAlignment(Type *Ty, bool abi_or_pref) const {
   }
 }
 
+TypeSize DataLayout::getTypeAllocSize(Type *Ty) const {
+  switch (Ty->getTypeID()) {
+  case Type::ArrayTyID: {
+    // The alignment of the array is the alignment of the element, so there
+    // is no need for further adjustment.
+    auto *ATy = cast<ArrayType>(Ty);
+    return ATy->getNumElements() * getTypeAllocSize(ATy->getElementType());
+  }
+  case Type::StructTyID: {
+    const StructLayout *Layout = getStructLayout(cast<StructType>(Ty));
+    TypeSize Size = Layout->getSizeInBytes();
+
+    if (cast<StructType>(Ty)->isPacked())
+      return Size;
+
+    Align A = std::max(StructABIAlignment, Layout->getAlignment());
+    return alignTo(Size, A.value());
+  }
+  case Type::IntegerTyID: {
+    unsigned BitWidth = Ty->getIntegerBitWidth();
+    TypeSize Size = TypeSize::getFixed(divideCeil(BitWidth, 8));
+    Align A = getIntegerAlignment(BitWidth, /*ABI=*/true);
+    return alignTo(Size, A.value());
+  }
+  case Type::PointerTyID: {
+    unsigned AS = Ty->getPointerAddressSpace();
+    TypeSize Size = TypeSize::getFixed(getPointerSize(AS));
+    return alignTo(Size, getPointerABIAlignment(AS).value());
+  }
+  case Type::TargetExtTyID: {
+    Type *LayoutTy = cast<TargetExtType>(Ty)->getLayoutType();
+    return getTypeAllocSize(LayoutTy);
+  }
+  default:
+    return alignTo(getTypeStoreSize(Ty), getABITypeAlign(Ty).value());
+  }
+}
+
 Align DataLayout::getABITypeAlign(Type *Ty) const {
   return getAlignment(Ty, true);
 }
diff --git a/llvm/test/CodeGen/AArch64/alloca-oversized.ll b/llvm/test/CodeGen/AArch64/alloca-oversized.ll
index e57bbcdf99800..81d301019e7b9 100644
--- a/llvm/test/CodeGen/AArch64/alloca-oversized.ll
+++ b/llvm/test/CodeGen/AArch64/alloca-oversized.ll
@@ -10,10 +10,7 @@ define void @test_oversized(ptr %dst, i32 %cond) {
 ; CHECK-NEXT:    .cfi_offset w30, -8
 ; CHECK-NEXT:    .cfi_offset w29, -16
 ; CHECK-NEXT:    mov x8, sp
-; CHECK-NEXT:    mov x9, #2305843009213693952 // =0x2000000000000000
-; CHECK-NEXT:    sub x8, x8, x9
 ; CHECK-NEXT:    sub x9, x29, #32
-; CHECK-NEXT:    mov sp, x8
 ; CHECK-NEXT:    cmp w1, #0
 ; CHECK-NEXT:    csel x8, x9, x8, eq
 ; CHECK-NEXT:    str x8, [x0]