AttributorAttributes: avoid a crashing on bad alignments

[durin42]

Summary

Prior to this change, LLVM would attempt to optimize an
aligned_alloc(33, ...) call to the stack. This flunked an assertion when
trying to emit the alloca, which crashed LLVM. Avoid that with extra
checks.

This was previously D119604, but when I landed that I missed a test
change and it was broken (and I was AFK, so a colleague rolled it back
for me.)

[durin42]

(going to reland under the old ID per a discussion, I'm new here)

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Assumptions.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Value.h"
 #include "llvm/IR/NoFolder.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/IPO/ArgumentPromotion.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <cassert>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "attributor"
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 (Value *Align = getAllocAlignment(AI.CB, TLI)) {
-      if (!getAPInt(A, *this, *Align)) {
+      Optional<APInt> APAlign = getAPInt(A, *this, *Align);
+      if (!APAlign) {
         // Can't generate an alloca which respects the required alignment
         // on the allocation.
         LLVM_DEBUG(dbgs() << "[H2S] Unknown allocation alignment: " << *AI.CB
                           << "\n");
         AI.Status = AllocationInfo::INVALID;
         Changed = ChangeStatus::CHANGED;
         continue;
+      } else {
+        if (APAlign->ugt(llvm::Value::MaximumAlignment) || !APAlign->isPowerOf2()) {
+          LLVM_DEBUG(dbgs() << "[H2S] Invalid allocation alignment: " << APAlign << "\n");
+          AI.Status = AllocationInfo::INVALID;
+          Changed = ChangeStatus::CHANGED;
+          continue;
+        }
       }
     }
 
     if (MaxHeapToStackSize != -1) {
       Optional<APInt> Size = getSize(A, *this, AI);
       if (!Size.hasValue() || Size.getValue().ugt(MaxHeapToStackSize)) {
         LLVM_DEBUG({
           if (!Size.hasValue())

llvm/test/Transforms/Attributor/heap_to_stack.ll

 ; CHECK-NEXT:    tail call void @free(i8* noalias nocapture [[TMP1]])
 ; CHECK-NEXT:    ret void
 ;
   %1 = tail call noalias i8* @aligned_alloc(i64 %alignment, i64 128)
   tail call void @free(i8* %1)
   ret void
 }
 
+; leave alone a constant-but-invalid alignment
+define void @test3d(i8* %p) {
+; CHECK-LABEL: define {{[^@]+}}@test3d
+; CHECK-SAME; (i8* nocapture [[P:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call noalias i8* @aligned_alloc(i64 noundef 33, i64 noundef 128)
+; CHECK:    tail call void @free(i8* noalias nocapture [[TMP1]])
+; CHECK-NEXT:    ret void
+;
+  %1 = tail call noalias i8* @aligned_alloc(i64 33, i64 128)
+  tail call void @nofree_arg_only(i8* %1, i8* %p)
+  tail call void @free(i8* %1)
+  ret void
+}
+
 declare noalias i8* @calloc(i64, i64)
 
 define void @test0() {
 ; IS________OPM-LABEL: define {{[^@]+}}@test0() {
 ; IS________OPM-NEXT:    [[TMP1:%.*]] = tail call noalias i8* @calloc(i64 noundef 2, i64 noundef 4)
 ; IS________OPM-NEXT:    tail call void @no_sync_func(i8* noalias nocapture nofree [[TMP1]])
 ; IS________OPM-NEXT:    tail call void @free(i8* noalias nocapture [[TMP1]])
 ; IS________OPM-NEXT:    ret void