Merge pull request swiftlang#74304 from aschwaighofer/use_static_alloca_for_stackAlloc

IRGen: Use static alloca for builtin stackAlloc when possible

Author: aschwaighofer

lib/IRGen/GenOpaque.cpp

@@ -656,7 +656,14 @@ void IRGenFunction::emitDeallocateDynamicAlloca(StackAddress address,
     // NOTE: llvm does not support dynamic allocas in coroutines.
 
     auto allocToken = address.getExtraInfo();
-    assert(allocToken && "dynamic alloca in coroutine without alloc token?");
+    if (!allocToken) {
+#ifndef NDEBUG
+      auto *alloca = cast<llvm::AllocaInst>(address.getAddress().getAddress());
+      assert(isa<llvm::ConstantInt>(alloca->getArraySize()) &&
+             "Dynamic alloca without a token?!");
+#endif
+      return;
+    }
     Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_alloca_free, allocToken);
     return;
   }

lib/IRGen/IRGenSIL.cpp

@@ -3559,6 +3559,11 @@ static llvm::Value *getStackAllocationSize(IRGenSILFunction &IGF,
     result = IGF.Builder.CreateMul(capacity, stride);
   }
 
+  if (auto constResult = dyn_cast<llvm::ConstantInt>(result)) {
+    if (!constResult->getUniqueInteger().isZero())
+      return constResult;
+  }
+
   // If the caller requests a zero-byte allocation, allocate one byte instead
   // to ensure that the resulting pointer is valid and unique on the stack.
   return IGF.Builder.CreateIntrinsicCall(llvm::Intrinsic::umax,
@@ -3601,6 +3606,14 @@ static void emitBuiltinStackAlloc(IRGenSILFunction &IGF,
     IGF, i->getOperand(0), i->getOperand(1), loc);
   auto align = getStackAllocationAlignment(IGF, i->getOperand(2), loc);
 
+  // Emit a static alloca if the size is constant.
+  if (auto *constSize = dyn_cast<llvm::ConstantInt>(size)) {
+    auto stackAddress = IGF.createAlloca(IGF.IGM.Int8Ty, constSize, align,
+                                         "temp_alloc");
+    IGF.setLoweredStackAddress(i, {stackAddress});
+    return;
+  }
+
   auto stackAddress = IGF.emitDynamicAlloca(IGF.IGM.Int8Ty, size, align,
                                             false, "temp_alloc");
   IGF.setLoweredStackAddress(i, stackAddress);

test/IRGen/temporary_allocation/async.swift

@@ -9,5 +9,5 @@ func f() async {
     blackHole(buffer.baseAddress)
   }
 }
-// CHECK: alloca [123 x i8], align 1
+
 // CHECK-NOT: swift_task_alloc

test/IRGen/temporary_allocation/codegen.swift

@@ -10,10 +10,20 @@ do {
   blackHole(ptr)
   ptr.deallocate()
 }
-// CHECK: ptrtoint ptr {{.*}} to [[WORD:i[0-9]+]]
 
 // MARK: Trivial Cases
 
+// CHECK: [[ONE_BYTE_PTR_RAW:%temp_alloc[0-9]*]] = alloca i8, align 1
+// CHECK: [[FIVE_BYTE_PTR_RAW:%temp_alloc[0-9]*]] = alloca [5 x i8], align 1
+// CHECK: [[ONE_KB_PTR_RAW:%temp_alloc[0-9]*]] = alloca [1024 x i8], align 8
+// CHECK: [[ONE_KB_RAND_PTR_RAW:%temp_alloc[0-9]*]] = alloca [1024 x i8], align 16
+// CHECK: [[INT_PTR_RAW:%temp_alloc[0-9]*]] = alloca [16 x i8], align 4
+// CHECK: [[INT_PTR_RAW2:%temp_alloc[0-9]*]] = alloca [16 x i8], align 4
+// CHECK: [[VOID_PTR_RAW:%temp_alloc[0-9]*]] = alloca [2 x i8], align 1
+
+// CHECK: ptrtoint ptr {{.*}} to [[WORD:i[0-9]+]]
+
+
 withUnsafeTemporaryAllocation(byteCount: 0, alignment: 1) { buffer in
   blackHole(buffer.baseAddress)
 }
@@ -24,21 +34,18 @@ withUnsafeTemporaryAllocation(byteCount: 0, alignment: 1) { buffer in
 withUnsafeTemporaryAllocation(byteCount: 1, alignment: 1) { buffer in
   blackHole(buffer.baseAddress)
 }
-// CHECK: [[ONE_BYTE_PTR_RAW:%temp_alloc[0-9]*]] = alloca i8, align 1
 // CHECK: [[ONE_BYTE_PTR:%[0-9]+]] = ptrtoint ptr [[ONE_BYTE_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[ONE_BYTE_PTR]])
 
 withUnsafeTemporaryAllocation(byteCount: 5, alignment: 1) { buffer in
   blackHole(buffer.baseAddress)
 }
-// CHECK: [[FIVE_BYTE_PTR_RAW:%temp_alloc[0-9]*]] = alloca [5 x i8], align 1
 // CHECK: [[FIVE_BYTE_PTR:%[0-9]+]] = ptrtoint ptr [[FIVE_BYTE_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[FIVE_BYTE_PTR]])
 
 withUnsafeTemporaryAllocation(byteCount: 1024, alignment: 8) { buffer in
   blackHole(buffer.baseAddress)
 }
-// CHECK: [[ONE_KB_PTR_RAW:%temp_alloc[0-9]*]] = alloca [1024 x i8], align 8
 // CHECK: [[ONE_KB_PTR:%[0-9]+]] = ptrtoint ptr [[ONE_KB_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[ONE_KB_PTR]])
 
@@ -47,7 +54,6 @@ withUnsafeTemporaryAllocation(byteCount: 1024, alignment: 8) { buffer in
 withUnsafeTemporaryAllocation(byteCount: 1024, alignment: Int.random(in: 0 ..< 16)) { buffer in
   blackHole(buffer.baseAddress)
 }
-// CHECK: [[ONE_KB_RAND_PTR_RAW:%temp_alloc[0-9]*]] = alloca [1024 x i8], align 16
 // CHECK: [[ONE_KB_RAND_PTR:%[0-9]+]] = ptrtoint ptr [[ONE_KB_RAND_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[ONE_KB_RAND_PTR]])
 
@@ -56,21 +62,18 @@ withUnsafeTemporaryAllocation(byteCount: 1024, alignment: Int.random(in: 0 ..< 1
 withUnsafeTemporaryAllocation(of: Int32.self, capacity: 4) { buffer in
   blackHole(buffer.baseAddress)
 }
-// CHECK: [[INT_PTR_RAW:%temp_alloc[0-9]*]] = alloca [16 x i8], align 4
 // CHECK: [[INT_PTR:%[0-9]+]] = ptrtoint ptr [[INT_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[INT_PTR]])
 
 _withUnprotectedUnsafeTemporaryAllocation(of: Int32.self, capacity: 4) { buffer in
   blackHole(buffer.baseAddress)
 }
-// CHECK: [[INT_PTR_RAW:%temp_alloc[0-9]*]] = alloca [16 x i8], align 4
-// CHECK: [[INT_PTR:%[0-9]+]] = ptrtoint ptr [[INT_PTR_RAW]] to [[WORD]]
+// CHECK: [[INT_PTR:%[0-9]+]] = ptrtoint ptr [[INT_PTR_RAW2]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[INT_PTR]])
 
 withUnsafeTemporaryAllocation(of: Void.self, capacity: 2) { buffer in
   blackHole(buffer.baseAddress)
 }
-// CHECK: [[VOID_PTR_RAW:%temp_alloc[0-9]*]] = alloca [2 x i8], align 1
 // CHECK: [[VOID_PTR:%[0-9]+]] = ptrtoint ptr [[VOID_PTR_RAW]] to [[WORD]]
 // CHECK: call swiftcc void @blackHole([[WORD]] [[VOID_PTR]])
 

test/IRGen/temporary_allocation/codegen_very_large_allocation.swift

@@ -24,7 +24,7 @@ withUnsafeTemporaryAllocation(byteCount: 0x0FFF_FFFF, alignment: 1) { buffer in
 // CHECK-LARGE-HEAP-ALLOC-NEXT: {{(tail )?}}call void @swift_slowDealloc(ptr [[HEAP_PTR_RAW]], [[WORD]] -1, [[WORD]] -1)
 
 // CHECK-LARGE-STACK-ALLOC: [[STACK_PTR_RAW:%temp_alloc[0-9]*]] = alloca [268435455 x i8], align 1
-// CHECK-LARGE-STACK-ALLOC-NEXT: [[STACK_PTR:%[0-9]+]] = ptrtoint ptr [[STACK_PTR_RAW]] to [[WORD]]
+// CHECK-LARGE-STACK-ALLOC: [[STACK_PTR:%[0-9]+]] = ptrtoint ptr [[STACK_PTR_RAW]] to [[WORD]]
 // CHECK-LARGE-STACK-ALLOC-NEXT: call swiftcc void @blackHole([[WORD]] [[STACK_PTR]])
 
 // CHECK-LARGE-ALLOC-DAG: [[IS_SAFE:%[0-9]+]] = {{(tail )?}}call {{(zeroext )?}}i1 @swift_stdlib_isStackAllocationSafe([[WORD]] 268435455, [[WORD]] 1)