[AllocaInst] Update `getAllocationSizeInBits` to return `TypeSize`.

[peterwaller-arm]

Tried to break it, wasn't able to. Looks straightforward.

Trivial nit.

llvm/unittests/IR/InstructionsTest.cpp

1381

Micro-nit: extra whitespace.

[sdesmalen]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

589

This should maintain the original operator> and instead GetAllocaSize should return RetSize->getFixedSize().

llvm/unittests/IR/InstructionsTest.cpp

502

nit: unnecessary change.

[fpetrogalli]

I reverted the return type of the lambda function to be integral type instead of TypeSize.

[fpetrogalli]

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

589

Good point! Done.

[sdesmalen]

Thanks for the changes @fpetrogalli, LGTM.

llvm/include/llvm/IR/Instructions.h

public:
 
   /// Overload to return most specific pointer type.
   PointerType *getType() const {
     return cast<PointerType>(Instruction::getType());
   }
 
   /// Get allocation size in bits. Returns None if size can't be determined,
   /// e.g. in case of a VLA.
-  Optional<uint64_t> getAllocationSizeInBits(const DataLayout &DL) const;
+  Optional<TypeSize> getAllocationSizeInBits(const DataLayout &DL) const;
 
   /// Return the type that is being allocated by the instruction.
   Type *getAllocatedType() const { return AllocatedType; }
   /// for use only in special circumstances that need to generically
   /// transform a whole instruction (eg: IR linking and vectorization).
   void setAllocatedType(Type *Ty) { AllocatedType = Ty; }
 
   /// Return the alignment of the memory that is being allocated by the

llvm/lib/IR/Instructions.cpp

 #include <vector>
 
 using namespace llvm;
 
 //===----------------------------------------------------------------------===//
 //                            AllocaInst Class
 //===----------------------------------------------------------------------===//
 
-Optional<uint64_t>
+Optional<TypeSize>
 AllocaInst::getAllocationSizeInBits(const DataLayout &DL) const {
-  uint64_t Size = DL.getTypeAllocSizeInBits(getAllocatedType());
+  TypeSize Size = DL.getTypeAllocSizeInBits(getAllocatedType());
   if (isArrayAllocation()) {
     auto *C = dyn_cast<ConstantInt>(getArraySize());
     if (!C)
       return None;
+    assert(!Size.isScalable() && "Array elements cannot have a scalable size");
     Size *= C->getZExtValue();
   }
   return Size;
 }
 
 //===----------------------------------------------------------------------===//
 //                              SelectInst Class
 //===----------------------------------------------------------------------===//

llvm/lib/Transforms/Coroutines/CoroFrame.cpp

void FrameTypeBuilder::addFieldForAllocas(const Function &F,
   StackLifetime StackLifetimeAnalyzer(F, ExtractAllocas(),
                                       StackLifetime::LivenessType::May);
   StackLifetimeAnalyzer.run();
   auto IsAllocaInferenre = [&](const AllocaInst *AI1, const AllocaInst *AI2) {
     return StackLifetimeAnalyzer.getLiveRange(AI1).overlaps(
         StackLifetimeAnalyzer.getLiveRange(AI2));
   };
   auto GetAllocaSize = [&](const AllocaInfo &A) {
-    Optional<uint64_t> RetSize = A.Alloca->getAllocationSizeInBits(DL);
-    assert(RetSize && "We can't handle scalable type now.\n");
+    Optional<TypeSize> RetSize = A.Alloca->getAllocationSizeInBits(DL);
+    assert(RetSize && "Variable Length Arrays (VLA) are not supported.\n");
+    assert(!RetSize->isScalable() && "Scalable vectors are not yet supported");
     return RetSize.getValue();
   };
   // Put larger allocas in the front. So the larger allocas have higher
   // priority to merge, which can save more space potentially. Also each
   // AllocaSet would be ordered. So we can get the largest Alloca in one
   // AllocaSet easily.
   sort(FrameData.Allocas, [&](const auto &Iter1, const auto &Iter2) {
-    return GetAllocaSize(Iter1) > GetAllocaSize(Iter2);
+    return TypeSize::isKnownGT(GetAllocaSize(Iter1), GetAllocaSize(Iter2));
   });
   for (const auto &A : FrameData.Allocas) {
     AllocaInst *Alloca = A.Alloca;
     bool Merged = false;
     // Try to find if the Alloca is not inferenced with any existing
     // NonOverlappedAllocaSet. If it is true, insert the alloca to that
     // NonOverlappedAllocaSet.
     for (auto &AllocaSet : NonOverlapedAllocas) {

llvm/unittests/IR/InstructionsTest.cpp

TEST(InstructionsTest, FPMathOperator) {
   Value *V1 = Builder.CreateFAdd(I, I, "", MD1);
   EXPECT_TRUE(isa<FPMathOperator>(V1));
   FPMathOperator *O1 = cast<FPMathOperator>(V1);
   EXPECT_EQ(O1->getFPAccuracy(), 1.0);
   V1->deleteValue();
   I->deleteValue();
 }
 
-
 TEST(InstructionsTest, isEliminableCastPair) {
   LLVMContext C;
 
   Type* Int16Ty = Type::getInt16Ty(C);
   Type* Int32Ty = Type::getInt32Ty(C);
   Type* Int64Ty = Type::getInt64Ty(C);
   Type* Int64PtrTy = Type::getInt64PtrTy(C);
 
ASSERT_TRUE(M);
     EXPECT_EQ(I2->getDebugLoc().getLine(), 2U);
     I2->dropLocation();
     EXPECT_EQ(I2->getDebugLoc().getLine(), 0U);
     EXPECT_EQ(I2->getDebugLoc().getScope(), Scope);
     EXPECT_EQ(I2->getDebugLoc().getInlinedAt(), nullptr);
   }
 }
 
+TEST(InstructionsTest, AllocaInst) {
+  LLVMContext Ctx;
+  std::unique_ptr<Module> M = parseIR(Ctx, R"(
+      %T = type { i64, [3 x i32]}
+      define void @f(i32 %n) {
+      entry:
+        %A = alloca i32, i32 1
+        %B = alloca i32, i32 4
+        %C = alloca i32, i32 %n
+        %D = alloca <8 x double>
+        %E = alloca <vscale x 8 x double>
+        %F = alloca [2 x half]
+        %G = alloca [2 x [3 x i128]]
+        %H = alloca %T
+        ret void
+      }
+    )");
+  const DataLayout &DL = M->getDataLayout();
+  ASSERT_TRUE(M);
+  Function *Fun = cast<Function>(M->getNamedValue("f"));
+  BasicBlock &BB = Fun->front();
+  auto It = BB.begin();
+  AllocaInst &A = cast<AllocaInst>(*It++);
+  AllocaInst &B = cast<AllocaInst>(*It++);
+  AllocaInst &C = cast<AllocaInst>(*It++);
+  AllocaInst &D = cast<AllocaInst>(*It++);
+  AllocaInst &E = cast<AllocaInst>(*It++);
+  AllocaInst &F = cast<AllocaInst>(*It++);
+  AllocaInst &G = cast<AllocaInst>(*It++);
+  AllocaInst &H = cast<AllocaInst>(*It++);
+  EXPECT_EQ(A.getAllocationSizeInBits(DL), TypeSize::getFixed(32));
+  EXPECT_EQ(B.getAllocationSizeInBits(DL), TypeSize::getFixed(128));
+  EXPECT_FALSE(C.getAllocationSizeInBits(DL));
+  EXPECT_EQ(D.getAllocationSizeInBits(DL), TypeSize::getFixed(512));
+  EXPECT_EQ(E.getAllocationSizeInBits(DL), TypeSize::getScalable(512));
+  EXPECT_EQ(F.getAllocationSizeInBits(DL), TypeSize::getFixed(32));
+  EXPECT_EQ(G.getAllocationSizeInBits(DL), TypeSize::getFixed(768));
+  EXPECT_EQ(H.getAllocationSizeInBits(DL), TypeSize::getFixed(160));
+}
+
 } // end anonymous namespace
 } // end namespace llvm