[MemoryBuiltins] Add field for alignment argument

[Bryce-MW]

Summary

There are a few places where the alignment argument for AlignedAllocLike functions was previously hardcoded. This goes along with removing hardcoded arguments for the allocation size. It also allows alligned operator new to be given an alignment value in InstCombinerImpl::annotateAnyAllocSite and makes Heap To Stack add the proper alignment for aligned new. I'm sure this isn't the best way to do this but I am new to the LLVM codebase so I appreciate any suggestions on how to improve.

[Bryce-MW]

• Fix build issue

[Bryce-MW]

The AllocationFnData table was not properly formatted before my modifications. When running clang-format, it changed the formatting of the entire table. I formatted it back to how it looked before but fixed a few issues with it. I think this looks better and is easier to read but I am happy to set it back to the way clang-format likes it if that is desired. Let me know. I think that was the only issue with the build.

[xbolva00]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2594

Works fine currently but we should replace Call.getOperand(1) with some more general API like

getAllocSizeArg or so.

[reames]

llvm/include/llvm/Analysis/MemoryBuiltins.h

75

Move this down to the new property section please.

76

I'd suggest changing the return value from the operand index to the Value * of the operand or nullptr. This seems like it would cleanup the current client code a bit, and would also allow us to return constant exprs for any allocation function whose alignment is a (known) function of the constant allocated size. I'm not suggesting you do the last bit, just that you leave the potential open.

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

5957

I think this block would be cleaner as:
if (isAllocationfn(*AI.CB)) {

Value *Align = getAlignment(*AI.CB)
..

}

This lets you drop all the passing around indices.

6268–6270

Same.

[Bryce-MW]

• Return a Value and add a size call

[Bryce-MW]

• Remove the last index-based check

[Bryce-MW]

Thanks for the suggestions! I had a feeling that there would be a better way than passing indexes around. I think I've mostly implemented them correctly. I'm not entirely sure if my use of getInitialValueOfAllocation is correct. I am also not sure about the part where getAllocSize has the possibility of creating a multiplication. Maybe that should return a pair of Values and let the caller create a multiplication if it wants?

[reames]

Please move the getAllocSize parts to it's own patch. We were converging on an LGTM for the alignment bits, and I want to keep the discussion focused on that.

[reames]

I'd suggest leaving the Attributor change to it's own review. I'm not familiar enough to review it easily. Let's get the allocation accessor in, and then iterate.

llvm/lib/Analysis/MemoryBuiltins.cpp

278

Er, this does not work. You need to check that AlignmentParam is not -1, or this will crash.

[Bryce-MW]

• Rollback changes to Attributor and addition of size call

[Bryce-MW]

I've rolled back the addition of the size function and the changes to attributor that go with it (I'll put those together in a separate patch once this one is done). I've kept the changes to attributor that are related to alignment but I can remove those as well if you would like, I wasn't entirely sure if you meant for me to do that or not. Thanks!

llvm/lib/Analysis/MemoryBuiltins.cpp

278

Good point, for some reason I assumed getOperand would return null for that case. Do I also need to check if it's too large?

[Bryce-MW]

• Format (other than the table as mentioned)

[reames]

LGTM for everything outside of Attributor - those lines I'm unsure of whether there's a missing precondition on the argument being constant.

Do you have commit access? If not, I can land the LGTMed parts for you.

[Bryce-MW]

I am assuming that I don't have commit access since this is my first contribution so I would appreciate you landing the changes outside of Attributor.

I agree that someone experienced with Attributor should look at it to be sure. Looking at it again, I think that it probably does assume that it is constant and while I could put in a check, I'm not sure if there is a way to back out if it was not (maybe this was an issue even with the original aligned_alloc?). getAssumedConstant will return nullptr if it is sure that the value is not constant. getAPInt will return none if the result from getAssumedConstant is neither a known constant nor none (unsure if constant or not). The like after the line I changed then asserts that the result from getAPInt has a value.

[reames]

In D116851#3230237, @Bryce-MW wrote:

I am assuming that I don't have commit access since this is my first contribution so I would appreciate you landing the changes outside of Attributor.

Will do tomorrow morning.

So we're clear on the legal stuff, can you confirm that this is intended as a contribution to the LLVM project and that you agree to the license terms thereof?

[Bryce-MW]

I confirm that this is intended as a contribution to the LLVM project and that I agree to the license terms thereof.

[reames]

Landed. When I went to land it, I noticed a couple of small issues, so this ended up being a small mini series rather than one patch.

commit 7febd60a90960ff1c74027c05394e134dad655a1 (HEAD -> main, origin/main)
Author: Bryce Wilson <[e-mail]>
Date: Mon Jan 10 09:08:55 2022 -0800

[instcombine] Add align return attributes for operator new(..., align_val)

commit fb936595faa44ad9c8d8550b05ea95e7be5f4703
Author: Bryce Wilson <[e-mail]>
Date: Mon Jan 10 08:58:44 2022 -0800

[MemoryBuiltins] Add field for alignment argument [NFC]

commit 332642e693509947d33ec5dbcbdc3fcf149f4a87
Author: Philip Reames <[e-mail]>
Date: Mon Jan 10 08:38:40 2022 -0800

Add test coverage for D116851

commit f4c54683d684100fffd307ddf773078429b83442
Author: Philip Reames <[e-mail]>
Date: Mon Jan 10 08:25:20 2022 -0800

[instcombine] Infer alignment for aligned_alloc with potentially zero size

[xbolva00]

BTW, https://clang.llvm.org/docs/AttributeReference.html#id330 ... is it broken just for me or.. ?

void *b(void *v, int align) __attribute__((alloc_size(2), alloc_align(2))) ;

void *b(void *v, int align) {
    return v;
}

; ModuleID = '/app/example.cpp'
source_filename = "/app/example.cpp"
target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-unknown-linux-gnu"

; Function Attrs: mustprogress nofree norecurse nosync nounwind readnone uwtable willreturn allocsize(1)
define dso_local i8* @_Z1bPvi(i8* readnone returned %0, i32 %1) local_unnamed_addr #0 {
  ret i8* %0
}

attributes #0 = { mustprogress nofree norecurse nosync nounwind readnone uwtable willreturn allocsize(1) "frame-pointer"="none" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "tune-cpu"="generic" }

!llvm.module.flags = !{!0, !1}
!llvm.ident = !{!2}

!0 = !{i32 1, !"wchar_size", i32 4}
!1 = !{i32 7, !"uwtable", i32 1}
!2 = !{!"clang version 14.0.0 (https://github.com/llvm/llvm-project.git d23fa4f2f1319039b7d939a4bf68c493ba26b07a)"}

Clang seems to ignore alloc_align attribute?

[reames]

In D116851#3236691, @xbolva00 wrote:

BTW, https://clang.llvm.org/docs/AttributeReference.html#id330 ... is it broken just for me or.. ?
...
Clang seems to ignore alloc_align attribute?

If by "broken", you mean "unimplemented" yes. LLVM has no allocalign attribute.

(For the record, that's related, but not directly relevant to this patch. Your comment made it seem like there might be a regression here. I believe that not to be the case, do you agree?)

[xbolva00]

In D116851#3237552, @reames wrote:

In D116851#3236691, @xbolva00 wrote:

BTW, https://clang.llvm.org/docs/AttributeReference.html#id330 ... is it broken just for me or.. ?
...
Clang seems to ignore alloc_align attribute?

If by "broken", you mean "unimplemented" yes. LLVM has no allocalign attribute.

(For the record, that's related, but not directly relevant to this patch. Your comment made it seem like there might be a regression here. I believe that not to be the case, do you agree?)

Yes, not related to this patch, I was just curious about that (if you know something more about that, etc..). Thanks.

[Bryce-MW]

There probably should be an allocalign attribute in the future once allocation functions are less hardcoded

[xbolva00]

llvm/lib/Analysis/MemoryBuiltins.cpp

64

Thinking about this again..

Maybe this should go to Clang?

We should really use allocsize attribute and annotate LibFuncs instead of using this hardcoded list.

llvm/include/llvm/Analysis/MemoryBuiltins.h

 bool isMallocLikeFn(const Value *V,
                     function_ref<const TargetLibraryInfo &(Function &)> GetTLI);
 
 /// Tests if a value is a call or invoke to a library function that
 /// allocates uninitialized memory with alignment (such as aligned_alloc).
 bool isAlignedAllocLikeFn(const Value *V, const TargetLibraryInfo *TLI);
 bool isAlignedAllocLikeFn(
     const Value *V, function_ref<const TargetLibraryInfo &(Function &)> GetTLI);
+/// Gets the alignment argument for an aligned_alloc-like function
+int getAllocAlignmentArg(const Value *V, const TargetLibraryInfo *TLI);
 
 /// Tests if a value is a call or invoke to a library function that
 /// allocates zero-filled memory (such as calloc).
 bool isCallocLikeFn(const Value *V, const TargetLibraryInfo *TLI);
 
 /// Tests if a value is a call or invoke to a library function that
 /// allocates memory similar to malloc or calloc.
 bool isMallocOrCallocLikeFn(const Value *V, const TargetLibraryInfo *TLI);

llvm/lib/Analysis/MemoryBuiltins.cpp

enum AllocType : uint8_t {
   CallocLike         = 1<<3, // allocates + bzero
   ReallocLike        = 1<<4, // reallocates
   StrDupLike         = 1<<5,
   MallocOrCallocLike = MallocLike | CallocLike | AlignedAllocLike,
   AllocLike          = MallocOrCallocLike | StrDupLike,
   AnyAlloc           = AllocLike | ReallocLike
 };
 
 struct AllocFnsTy {
   AllocType AllocTy;
   unsigned NumParams;
   // First and Second size parameters (or -1 if unused)
   int FstParam, SndParam;
+  // Alignment parameter for aligned_alloc and aligned new
+  int AlignParam;
 };
 
 // FIXME: certain users need more information. E.g., SimplifyLibCalls needs to
 // know which functions are nounwind, noalias, nocapture parameters, etc.
 static const std::pair<LibFunc, AllocFnsTy> AllocationFnData[] = {
-  {LibFunc_malloc,              {MallocLike,  1, 0,  -1}},
-  {LibFunc_vec_malloc,          {MallocLike,  1, 0,  -1}},
-  {LibFunc_valloc,              {MallocLike,  1, 0,  -1}},
-  {LibFunc_Znwj,                {OpNewLike,   1, 0,  -1}}, // new(unsigned int)
-  {LibFunc_ZnwjRKSt9nothrow_t,  {MallocLike,  2, 0,  -1}}, // new(unsigned int, nothrow)
-  {LibFunc_ZnwjSt11align_val_t, {OpNewLike,   2, 0,  -1}}, // new(unsigned int, align_val_t)
-  {LibFunc_ZnwjSt11align_val_tRKSt9nothrow_t, // new(unsigned int, align_val_t, nothrow)
-                                {MallocLike,  3, 0,  -1}},
-  {LibFunc_Znwm,                {OpNewLike,   1, 0,  -1}}, // new(unsigned long)
-  {LibFunc_ZnwmRKSt9nothrow_t,  {MallocLike,  2, 0,  -1}}, // new(unsigned long, nothrow)
-  {LibFunc_ZnwmSt11align_val_t, {OpNewLike,   2, 0,  -1}}, // new(unsigned long, align_val_t)
-  {LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t, // new(unsigned long, align_val_t, nothrow)
-                                {MallocLike,  3, 0,  -1}},
-  {LibFunc_Znaj,                {OpNewLike,   1, 0,  -1}}, // new[](unsigned int)
-  {LibFunc_ZnajRKSt9nothrow_t,  {MallocLike,  2, 0,  -1}}, // new[](unsigned int, nothrow)
-  {LibFunc_ZnajSt11align_val_t, {OpNewLike,   2, 0,  -1}}, // new[](unsigned int, align_val_t)
-  {LibFunc_ZnajSt11align_val_tRKSt9nothrow_t, // new[](unsigned int, align_val_t, nothrow)
-                                {MallocLike,  3, 0,  -1}},
-  {LibFunc_Znam,                {OpNewLike,   1, 0,  -1}}, // new[](unsigned long)
-  {LibFunc_ZnamRKSt9nothrow_t,  {MallocLike,  2, 0,  -1}}, // new[](unsigned long, nothrow)
-  {LibFunc_ZnamSt11align_val_t, {OpNewLike,   2, 0,  -1}}, // new[](unsigned long, align_val_t)
-  {LibFunc_ZnamSt11align_val_tRKSt9nothrow_t, // new[](unsigned long, align_val_t, nothrow)
-                                 {MallocLike,  3, 0,  -1}},
-  {LibFunc_msvc_new_int,         {OpNewLike,   1, 0,  -1}}, // new(unsigned int)
-  {LibFunc_msvc_new_int_nothrow, {MallocLike,  2, 0,  -1}}, // new(unsigned int, nothrow)
-  {LibFunc_msvc_new_longlong,         {OpNewLike,   1, 0,  -1}}, // new(unsigned long long)
-  {LibFunc_msvc_new_longlong_nothrow, {MallocLike,  2, 0,  -1}}, // new(unsigned long long, nothrow)
-  {LibFunc_msvc_new_array_int,         {OpNewLike,   1, 0,  -1}}, // new[](unsigned int)
-  {LibFunc_msvc_new_array_int_nothrow, {MallocLike,  2, 0,  -1}}, // new[](unsigned int, nothrow)
-  {LibFunc_msvc_new_array_longlong,         {OpNewLike,   1, 0,  -1}}, // new[](unsigned long long)
-  {LibFunc_msvc_new_array_longlong_nothrow, {MallocLike,  2, 0,  -1}}, // new[](unsigned long long, nothrow)
-  {LibFunc_aligned_alloc,       {AlignedAllocLike, 2, 1,  -1}},
-  {LibFunc_memalign,            {AlignedAllocLike, 2, 1,  -1}},
-  {LibFunc_calloc,              {CallocLike,  2, 0,   1}},
-  {LibFunc_vec_calloc,          {CallocLike,  2, 0,   1}},
-  {LibFunc_realloc,             {ReallocLike, 2, 1,  -1}},
-  {LibFunc_vec_realloc,         {ReallocLike, 2, 1,  -1}},
-  {LibFunc_reallocf,            {ReallocLike, 2, 1,  -1}},
-  {LibFunc_strdup,              {StrDupLike,  1, -1, -1}},
-  {LibFunc_strndup,             {StrDupLike,  2, 1,  -1}},
-  {LibFunc___kmpc_alloc_shared, {MallocLike,  1, 0,  -1}},
+    {LibFunc_malloc,                            {MallocLike,       1,  0, -1, -1}},
+    {LibFunc_vec_malloc,                        {MallocLike,       1,  0, -1, -1}},
+    {LibFunc_valloc,                            {MallocLike,       1,  0, -1, -1}},
+    {LibFunc_Znwj,                              {OpNewLike,        1,  0, -1, -1}}, // new(unsigned int)
+    {LibFunc_ZnwjRKSt9nothrow_t,                {MallocLike,       2,  0, -1, -1}}, // new(unsigned int, nothrow)
+    {LibFunc_ZnwjSt11align_val_t,               {OpNewLike,        2,  0, -1,  1}}, // new(unsigned int, align_val_t)
+    {LibFunc_ZnwjSt11align_val_tRKSt9nothrow_t, {MallocLike,       3,  0, -1,  1}}, // new(unsigned int, align_val_t, nothrow)
+    {LibFunc_Znwm,                              {OpNewLike,        1,  0, -1, -1}}, // new(unsigned long)
+    {LibFunc_ZnwmRKSt9nothrow_t,                {MallocLike,       2,  0, -1, -1}}, // new(unsigned long, nothrow)
+    {LibFunc_ZnwmSt11align_val_t,               {OpNewLike,        2,  0, -1,  1}}, // new(unsigned long, align_val_t)
+    {LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t, {MallocLike,       3,  0, -1,  1}}, // new(unsigned long, align_val_t, nothrow)
+    {LibFunc_Znaj,                              {OpNewLike,        1,  0, -1, -1}}, // new[](unsigned int)
+    {LibFunc_ZnajRKSt9nothrow_t,                {MallocLike,       2,  0, -1, -1}}, // new[](unsigned int, nothrow)
+    {LibFunc_ZnajSt11align_val_t,               {OpNewLike,        2,  0, -1,  1}}, // new[](unsigned int, align_val_t)
+    {LibFunc_ZnajSt11align_val_tRKSt9nothrow_t, {MallocLike,       3,  0, -1,  1}}, // new[](unsigned int, align_val_t, nothrow)
+    {LibFunc_Znam,                              {OpNewLike,        1,  0, -1, -1}}, // new[](unsigned long)
+    {LibFunc_ZnamRKSt9nothrow_t,                {MallocLike,       2,  0, -1, -1}}, // new[](unsigned long, nothrow)
+    {LibFunc_ZnamSt11align_val_t,               {OpNewLike,        2,  0, -1,  1}}, // new[](unsigned long, align_val_t)
+    {LibFunc_ZnamSt11align_val_tRKSt9nothrow_t, {MallocLike,       3,  0, -1,  1}}, // new[](unsigned long, align_val_t, nothrow)
+    {LibFunc_msvc_new_int,                      {OpNewLike,        1,  0, -1, -1}}, // new(unsigned int)
+    {LibFunc_msvc_new_int_nothrow,              {MallocLike,       2,  0, -1, -1}}, // new(unsigned int, nothrow)
+    {LibFunc_msvc_new_longlong,                 {OpNewLike,        1,  0, -1, -1}}, // new(unsigned long long)
+    {LibFunc_msvc_new_longlong_nothrow,         {MallocLike,       2,  0, -1, -1}}, // new(unsigned long long, nothrow)
+    {LibFunc_msvc_new_array_int,                {OpNewLike,        1,  0, -1, -1}}, // new[](unsigned int)
+    {LibFunc_msvc_new_array_int_nothrow,        {MallocLike,       2,  0, -1, -1}}, // new[](unsigned int, nothrow)
+    {LibFunc_msvc_new_array_longlong,           {OpNewLike,        1,  0, -1, -1}}, // new[](unsigned long long)
+    {LibFunc_msvc_new_array_longlong_nothrow,   {MallocLike,       2,  0, -1, -1}}, // new[](unsigned long long, nothrow)
+    {LibFunc_aligned_alloc,                     {AlignedAllocLike, 2,  1, -1,  0}},
+    {LibFunc_memalign,                          {AlignedAllocLike, 2,  1, -1,  0}},
+    {LibFunc_calloc,                            {CallocLike,       2,  0,  1, -1}},
+    {LibFunc_vec_calloc,                        {CallocLike,       2,  0,  1, -1}},
+    {LibFunc_realloc,                           {ReallocLike,      2,  1, -1, -1}},
+    {LibFunc_vec_realloc,                       {ReallocLike,      2,  1, -1, -1}},
+    {LibFunc_reallocf,                          {ReallocLike,      2,  1, -1, -1}},
+    {LibFunc_strdup,                            {StrDupLike,       1, -1, -1, -1}},
+    {LibFunc_strndup,                           {StrDupLike,       2,  1, -1, -1}},
+    {LibFunc___kmpc_alloc_shared,               {MallocLike,       1,  0, -1, -1}},
     // TODO: Handle "int posix_memalign(void **, size_t, size_t)"
 };
 
 static const Function *getCalledFunction(const Value *V,
                                          bool &IsNoBuiltin) {
   // Don't care about intrinsics in this case.
   if (isa<IntrinsicInst>(V))
     return nullptr;
 
static Optional<AllocFnsTy> getAllocationSize(const Value *V,
 
   AllocFnsTy Result;
   // Because allocsize only tells us how many bytes are allocated, we're not
   // really allowed to assume anything, so we use MallocLike.
   Result.AllocTy = MallocLike;
   Result.NumParams = Callee->getNumOperands();
   Result.FstParam = Args.first;
   Result.SndParam = Args.second.getValueOr(-1);
+  // Allocsize has no way to specify an alignment argument
+  Result.AlignParam = -1;
   return Result;
 }
 
 static bool hasNoAliasAttr(const Value *V) {
   const auto *CB = dyn_cast<CallBase>(V);
   return CB && CB->hasRetAttr(Attribute::NoAlias);
 }
 
bool llvm::isAlignedAllocLikeFn(const Value *V, const TargetLibraryInfo *TLI) {
   return getAllocationData(V, AlignedAllocLike, TLI)
       .hasValue();
 }
 bool llvm::isAlignedAllocLikeFn(
     const Value *V, function_ref<const TargetLibraryInfo &(Function &)> GetTLI) {
   return getAllocationData(V, AlignedAllocLike, GetTLI)
       .hasValue();
 }
+/// Gets the alignment argument for an aligned_alloc-like function
+int getAllocAlignmentArg(const Value *V, const TargetLibraryInfo *TLI) {
+  const Optional<AllocFnsTy> FnData = getAllocationData(V, AnyAlloc, TLI);
+  if (!FnData.hasValue()) {
+    return -1;
+  }
+  return FnData->AlignParam;
+}
 
 /// Tests if a value is a call or invoke to a library function that
 /// allocates zero-filled memory (such as calloc).
 bool llvm::isCallocLikeFn(const Value *V, const TargetLibraryInfo *TLI) {
   return getAllocationData(V, CallocLike, TLI).hasValue();
 }
 
 /// Tests if a value is a call or invoke to a library function that

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

struct AllocationInfo {
 
     /// The kind of allocation.
     const enum class AllocationKind {
       MALLOC,
       CALLOC,
       ALIGNED_ALLOC,
     } Kind;
 
+    /// The position of the alignment argument (or -1)
+    const int AlignArgPos;
+
     /// The library function id for the allocation.
     LibFunc LibraryFunctionId = NotLibFunc;
 
     /// The status wrt. a rewrite.
     enum {
       STACK_DUE_TO_USE,
       STACK_DUE_TO_FREE,
       INVALID,
auto AllocationIdentifierCB = [&](Instruction &I) {
       bool IsCalloc =
           !IsMalloc && !IsAlignedAllocLike && isCallocLikeFn(CB, TLI);
       if (!IsMalloc && !IsAlignedAllocLike && !IsCalloc)
         return true;
       auto Kind =
           IsMalloc ? AllocationInfo::AllocationKind::MALLOC
                    : (IsCalloc ? AllocationInfo::AllocationKind::CALLOC
                                : AllocationInfo::AllocationKind::ALIGNED_ALLOC);
+      int AlignArgPos = getAllocAlignmentArg(CB, TLI);
 
-      AllocationInfo *AI = new (A.Allocator) AllocationInfo{CB, Kind};
+      AllocationInfo *AI =
+          new (A.Allocator) AllocationInfo{CB, Kind, AlignArgPos};
       AllocationInfos[CB] = AI;
       TLI->getLibFunc(*CB, AI->LibraryFunctionId);
       return true;
     };
 
     bool UsedAssumedInformation = false;
     bool Success = A.checkForAllCallLikeInstructions(
         AllocationIdentifierCB, *this, UsedAssumedInformation,
for (auto &It : AllocationInfos) {
         Size = AI.CB->getOperand(1);
       } else {
         Size = AI.CB->getOperand(0);
       }
 
       Align Alignment(1);
       if (MaybeAlign RetAlign = AI.CB->getRetAlign())
         Alignment = max(Alignment, RetAlign);
-      if (AI.Kind == AllocationInfo::AllocationKind::ALIGNED_ALLOC) {
+      if (AI.AlignArgPos != -1) {
         Optional<APInt> AlignmentAPI =
-            getAPInt(A, *this, *AI.CB->getArgOperand(0));
+            getAPInt(A, *this, *AI.CB->getArgOperand(AI.AlignArgPos));
         assert(AlignmentAPI.hasValue() &&
                "Expected an alignment during manifest!");
         Alignment =
             max(Alignment, MaybeAlign(AlignmentAPI.getValue().getZExtValue()));
       }
 
       unsigned AS = cast<PointerType>(AI.CB->getType())->getAddressSpace();
       Instruction *Alloca =
ChangeStatus AAHeapToStackFunction::updateImpl(Attributor &A) {
   // The actual update starts here. We look at all allocations and depending on
   // their status perform the appropriate check(s).
   for (auto &It : AllocationInfos) {
     AllocationInfo &AI = *It.second;
     if (AI.Status == AllocationInfo::INVALID)
       continue;
 
     if (MaxHeapToStackSize == -1) {
-      if (AI.Kind == AllocationInfo::AllocationKind::ALIGNED_ALLOC)
-        if (!getAPInt(A, *this, *AI.CB->getArgOperand(0)).hasValue()) {
+      if (AI.AlignArgPos != -1)
+        if (!getAPInt(A, *this, *AI.CB->getArgOperand(AI.AlignArgPos))
+                 .hasValue()) {
           LLVM_DEBUG(dbgs() << "[H2S] Unknown allocation alignment: " << *AI.CB
                             << "\n");
           AI.Status = AllocationInfo::INVALID;
           Changed = ChangeStatus::CHANGED;
           continue;
         }
     } else {
       Optional<APInt> Size = getSize(A, *this, AI);

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

if (isOpNewLikeFn(&Call, TLI))
       Call.addRetAttr(Attribute::getWithDereferenceableBytes(
           Call.getContext(), Size));
     else
       Call.addRetAttr(Attribute::getWithDereferenceableOrNullBytes(
           Call.getContext(), Size));
   }
 
   // Add alignment attribute if alignment is a power of two constant.
-  if (!isAlignedAllocLikeFn(&Call, TLI))
+  const int AlignmentArg = getAllocAlignmentArg(&Call, TLI);
+  if (AlignmentArg == -1)
     return;
 
-  ConstantInt *Op0C = dyn_cast<ConstantInt>(Call.getOperand(0));
-  if (Op0C && Op0C->getValue().ult(llvm::Value::MaximumAlignment) &&
+  ConstantInt *AlignOpC = dyn_cast<ConstantInt>(Call.getOperand(AlignmentArg));
+  if (AlignOpC && AlignOpC->getValue().ult(llvm::Value::MaximumAlignment) &&
       isKnownNonZero(Call.getOperand(1), DL, 0, &AC, &Call, &DT)) {
-    uint64_t AlignmentVal = Op0C->getZExtValue();
+    uint64_t AlignmentVal = AlignOpC->getZExtValue();
     if (llvm::isPowerOf2_64(AlignmentVal)) {
       Call.removeRetAttr(Attribute::Alignment);
       Call.addRetAttr(Attribute::getWithAlignment(Call.getContext(),
                                                   Align(AlignmentVal)));
     }
   }
 }