SROA optimizations depend on alloca type

[gbaraldi]

While debugging a regression in julia's codegen I saw this interesting behaviour. https://godbolt.org/z/KeEnhjc6G

I wasn't able to minimize it too much more than that sadly, it seems brittle, but basically SROA generates meaningfully different code depending on the type of the alloca, because SROA runs quite early for us (and in general) this ends up making a big difference in how LLVM optimizes this code, generating either

define swiftcc float @julia_f_3398(ptr nonnull swiftself "gcstack" %0, ptr nocapture noundef nonnull readonly align 4 dereferenceable(8) %1) local_unnamed_addr #0 {
  %3 = getelementptr inbounds nuw i8, ptr %0, i64 16
  %4 = load ptr, ptr %3, align 8, !tbaa !2
  %5 = getelementptr inbounds nuw i8, ptr %4, i64 16
  %6 = load atomic ptr, ptr %5 monotonic, align 8, !tbaa !6, !invariant.load !8
  fence syncscope("singlethread") seq_cst
  %7 = load volatile i64, ptr %6, align 8
  fence syncscope("singlethread") seq_cst
  %.sroa.06.0.copyload7 = load i64, ptr %1, align 4, !tbaa !9, !alias.scope !10, !noalias !14
  %.sroa.3.0.extract.shift12 = lshr i64 %.sroa.06.0.copyload7, 32
  %.sroa.3.0.extract.trunc13 = trunc nuw i64 %.sroa.3.0.extract.shift12 to i32
  %8 = bitcast i32 %.sroa.3.0.extract.trunc13 to float
  br label %9

9:                                                ; preds = %16, %2
  %10 = phi float [ %8, %2 ], [ %18, %16 ]
  %11 = phi i64 [ 1, %2 ], [ %17, %16 ]
  %.sroa.06.014 = phi i64 [ %.sroa.06.0.copyload7, %2 ], [ %.sroa.0.1, %16 ]
  %12 = and i64 %11, 1
  %.not = icmp eq i64 %12, 0
  br i1 %.not, label %19, label %16

13:                                               ; preds = %16
  %.sroa.010.0.extract.trunc = trunc i64 %.sroa.0.1 to i32
  %14 = bitcast i32 %.sroa.010.0.extract.trunc to float
  %15 = fadd float %14, %18
  ret float %15

16:                                               ; preds = %19, %9
  %.sroa.0.1 = phi i64 [ %.sroa.09.0.insert.insert, %19 ], [ %.sroa.06.014, %9 ]
  %17 = add nuw nsw i64 %11, 1
  %.sroa.3.0.extract.shift = lshr i64 %.sroa.0.1, 32
  %.sroa.3.0.extract.trunc = trunc nuw i64 %.sroa.3.0.extract.shift to i32
  %18 = bitcast i32 %.sroa.3.0.extract.trunc to float
  %exitcond = icmp eq i64 %17, 1001
  br i1 %exitcond, label %13, label %9

19:                                               ; preds = %9
  %20 = fadd float %10, 1.000000e+00
  %21 = bitcast float %20 to i32
  %.sroa.2.0.insert.ext = zext i32 %21 to i64
  %.sroa.2.0.insert.shift = shl nuw i64 %.sroa.2.0.insert.ext, 32
  %.sroa.09.0.insert.ext = and i64 %.sroa.06.014, 4294967295
  %.sroa.09.0.insert.insert = or disjoint i64 %.sroa.2.0.insert.shift, %.sroa.09.0.insert.ext
  br label %16
}

or

define swiftcc float @julia_f_4090(ptr nonnull swiftself %0, ptr nocapture noundef nonnull readonly align 4 dereferenceable(8) %1) local_unnamed_addr #0 {
  %3 = getelementptr inbounds nuw i8, ptr %0, i64 16
  %4 = load ptr, ptr %3, align 8, !tbaa !3
  %5 = getelementptr inbounds nuw i8, ptr %4, i64 16
  %6 = load ptr, ptr %5, align 8, !tbaa !7, !invariant.load !9
  fence syncscope("singlethread") seq_cst
  call void @julia.safepoint(ptr %6)
  fence syncscope("singlethread") seq_cst
  %.sroa.013.0.copyload = load float, ptr %1, align 4, !tbaa !10, !alias.scope !11, !noalias !15
  %.sroa.614.0..sroa_idx = getelementptr inbounds nuw i8, ptr %1, i64 4
  %.sroa.614.0.copyload = load float, ptr %.sroa.614.0..sroa_idx, align 4, !tbaa !10, !alias.scope !11, !noalias !15
  br label %7

7:                                                ; preds = %7, %2
  %8 = phi i64 [ 1, %2 ], [ %11, %7 ]
  %.sroa.614.015 = phi float [ %.sroa.614.0.copyload, %2 ], [ %spec.select, %7 ]
  %9 = and i64 %8, 1
  %.not = icmp eq i64 %9, 0
  %10 = fadd float %.sroa.614.015, 1.000000e+00
  %spec.select = select i1 %.not, float %10, float %.sroa.614.015
  %11 = add nuw nsw i64 %8, 1
  %exitcond = icmp eq i64 %11, 1001
  br i1 %exitcond, label %12, label %7

12:                                               ; preds = %7
  %13 = fadd float %.sroa.013.0.copyload, %10
  ret float %13
}

This second option is significantly faster to execute on the cpu this was targetting (apple-m3)

[gbaraldi]

Looking at the SROA source and debug output the difference seems to be
it allows merging the 2 x i32 alloca into an i64 but not the 2 x float

SROA debug output

Slices of alloca: %7 = alloca [2 x float], align 4
SROA alloca: %7 = alloca [2 x float], align 4
Rewriting FCA loads and stores...
Slices of alloca: %7 = alloca [2 x float], align 4
[0,8) slice #0 (splittable)
used by: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %6, ptr align 4 %7, i64 8, i1 false)
[0,8) slice #1 (splittable)
used by: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %5, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
[0,8) slice #2 (splittable)
used by: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %3, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
[0,8) slice #3 (splittable)
used by: call void @llvm.lifetime.end.p0(ptr %7)
[0,8) slice #4 (splittable)
used by: call void @llvm.lifetime.start.p0(ptr %7)
[0,8) slice #5 (splittable)
used by: call void @llvm.lifetime.start.p0(ptr %7)
Pre-splitting loads and stores
Searching for candidate loads and stores
Rewriting alloca partition [0,8) to: %7 = alloca [2 x float], align 4
rewriting [0,8) slice #0 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %6, ptr align 4 %7, i64 8, i1 false)
rewriting [0,8) slice #1 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %5, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
rewriting [0,8) slice #2 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %3, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
rewriting [0,8) slice #3 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.lifetime.end.p0(ptr %7)
to: call void @llvm.lifetime.end.p0(ptr %7)
rewriting [0,8) slice #4 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.lifetime.start.p0(ptr %7)
to: call void @llvm.lifetime.start.p0(ptr %7)
rewriting [0,8) slice #5 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.lifetime.start.p0(ptr %7)
to: call void @llvm.lifetime.start.p0(ptr %7)
Speculating PHIs
Rewriting Selects
Deleting dead instruction: call void @llvm.lifetime.start.p0(ptr %7)
Deleting dead instruction: call void @llvm.lifetime.start.p0(ptr %7)
Deleting dead instruction: call void @llvm.lifetime.end.p0(ptr %7)

vs

SROA alloca: %7 = alloca [2 x i32], align 4
Rewriting FCA loads and stores...
Slices of alloca: %7 = alloca [2 x i32], align 4
[0,8) slice #0 (splittable)
used by: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %6, ptr align 4 %7, i64 8, i1 false)
[0,8) slice #1 (splittable)
used by: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %5, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
[0,8) slice #2 (splittable)
used by: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %3, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
[0,8) slice #3 (splittable)
used by: call void @llvm.lifetime.end.p0(ptr %7)
[0,8) slice #4 (splittable)
used by: call void @llvm.lifetime.start.p0(ptr %7)
[0,8) slice #5 (splittable)
used by: call void @llvm.lifetime.start.p0(ptr %7)
Pre-splitting loads and stores
Searching for candidate loads and stores
Rewriting alloca partition [0,8) to: %.sroa.0 = alloca i64, align 8
rewriting [0,8) slice #0 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %6, ptr align 4 %7, i64 8, i1 false)
to: store i64 %.sroa.0.0.copyload, ptr %6, align 4
rewriting [0,8) slice #1 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %5, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
to: store i64 %.sroa.0.0.copyload1, ptr %.sroa.0, align 8, !tbaa !17, !alias.scope !19, !noalias !20
rewriting [0,8) slice #2 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %3, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
to: store i64 %.sroa.0.0.copyload2, ptr %.sroa.0, align 8, !tbaa !17, !alias.scope !19, !noalias !20
rewriting [0,8) slice #3 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.lifetime.end.p0(ptr %7)
to: call void @llvm.lifetime.end.p0(ptr %.sroa.0)
rewriting [0,8) slice #4 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.lifetime.start.p0(ptr %7)
to: call void @llvm.lifetime.start.p0(ptr %.sroa.0)
rewriting [0,8) slice #5 (splittable)
Begin:(0, 8) NewBegin:(0, 8) NewAllocaBegin:(0, 8)
original: call void @llvm.lifetime.start.p0(ptr %7)
to: call void @llvm.lifetime.start.p0(ptr %.sroa.0)
Speculating PHIs
Rewriting Selects
Deleting dead instruction: call void @llvm.lifetime.start.p0(ptr %7)
Deleting dead instruction: call void @llvm.lifetime.start.p0(ptr %7)
Deleting dead instruction: call void @llvm.lifetime.end.p0(ptr %7)
Deleting dead instruction: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %3, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
Deleting dead instruction: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %7, ptr align 4 %5, i64 8, i1 false), !tbaa !17, !alias.scope !19, !noalias !20
Deleting dead instruction: call void @llvm.memcpy.p0.p0.i64(ptr align 4 %6, ptr align 4 %7, i64 8, i1 false)
Deleting dead instruction: %7 = alloca [2 x i32], align 4

[gbaraldi]

llvm-project/llvm/lib/Transforms/Scalar/SROA.cpp

Lines 4885 to 4966 in e6b4a21

	Type *SliceTy = nullptr;
	VectorType *SliceVecTy = nullptr;
	const DataLayout &DL = AI.getDataLayout();
	unsigned VScale = AI.getFunction()->getVScaleValue();
	std::pair<Type *, IntegerType *> CommonUseTy =
	findCommonType(P.begin(), P.end(), P.endOffset());
	// Do all uses operate on the same type?
	if (CommonUseTy.first) {
	TypeSize CommonUseSize = DL.getTypeAllocSize(CommonUseTy.first);
	if (CommonUseSize.isFixed() && CommonUseSize.getFixedValue() >= P.size()) {
	SliceTy = CommonUseTy.first;
	SliceVecTy = dyn_cast<VectorType>(SliceTy);
	}
	}
	// If not, can we find an appropriate subtype in the original allocated type?
	if (!SliceTy)
	if (Type *TypePartitionTy = getTypePartition(DL, AI.getAllocatedType(),
	P.beginOffset(), P.size()))
	SliceTy = TypePartitionTy;
	// If still not, can we use the largest bitwidth integer type used?
	if (!SliceTy && CommonUseTy.second)
	if (DL.getTypeAllocSize(CommonUseTy.second).getFixedValue() >= P.size()) {
	SliceTy = CommonUseTy.second;
	SliceVecTy = dyn_cast<VectorType>(SliceTy);
	}
	if ((!SliceTy || (SliceTy->isArrayTy() &&
	SliceTy->getArrayElementType()->isIntegerTy())) &&
	DL.isLegalInteger(P.size() * 8)) {
	SliceTy = Type::getIntNTy(*C, P.size() * 8);
	}
	// If the common use types are not viable for promotion then attempt to find
	// another type that is viable.
	if (SliceVecTy && !checkVectorTypeForPromotion(P, SliceVecTy, DL, VScale))
	if (Type *TypePartitionTy = getTypePartition(DL, AI.getAllocatedType(),
	P.beginOffset(), P.size())) {
	VectorType *TypePartitionVecTy = dyn_cast<VectorType>(TypePartitionTy);
	if (TypePartitionVecTy &&
	checkVectorTypeForPromotion(P, TypePartitionVecTy, DL, VScale))
	SliceTy = TypePartitionTy;
	}
	if (!SliceTy)
	SliceTy = ArrayType::get(Type::getInt8Ty(*C), P.size());
	assert(DL.getTypeAllocSize(SliceTy).getFixedValue() >= P.size());
	bool IsIntegerPromotable = isIntegerWideningViable(P, SliceTy, DL);
	VectorType *VecTy =
	IsIntegerPromotable ? nullptr : isVectorPromotionViable(P, DL, VScale);
	if (VecTy)
	SliceTy = VecTy;
	// Check for the case where we're going to rewrite to a new alloca of the
	// exact same type as the original, and with the same access offsets. In that
	// case, re-use the existing alloca, but still run through the rewriter to
	// perform phi and select speculation.
	// P.beginOffset() can be non-zero even with the same type in a case with
	// out-of-bounds access (e.g. @PR35657 function in SROA/basictest.ll).
	AllocaInst *NewAI;
	if (SliceTy == AI.getAllocatedType() && P.beginOffset() == 0) {
	NewAI = &AI;
	// FIXME: We should be able to bail at this point with "nothing changed".
	// FIXME: We might want to defer PHI speculation until after here.
	// FIXME: return nullptr;
	} else {
	// Make sure the alignment is compatible with P.beginOffset().
	const Align Alignment = commonAlignment(AI.getAlign(), P.beginOffset());
	// If we will get at least this much alignment from the type alone, leave
	// the alloca's alignment unconstrained.
	const bool IsUnconstrained = Alignment <= DL.getABITypeAlign(SliceTy);
	NewAI = new AllocaInst(
	SliceTy, AI.getAddressSpace(), nullptr,
	IsUnconstrained ? DL.getPrefTypeAlign(SliceTy) : Alignment,
	AI.getName() + ".sroa." + Twine(P.begin() - AS.begin()),
	AI.getIterator());
	// Copy the old AI debug location over to the new one.
	NewAI->setDebugLoc(AI.getDebugLoc());
	++NumNewAllocas;
	}

is probably where the issue is