From 13c384c26176f0ca935e0a78224551f91835087b Mon Sep 17 00:00:00 2001 From: Joseph Huber Date: Fri, 29 Aug 2025 11:32:52 -0500 Subject: [PATCH 1/3] [InferAlignment] Increase alignment in masked load / store instrinsics if known Summary: The masked load / store LLVM intrinsics take an argument for the alignment. If the user is pessimistic about alignment they can provide a value of `1` for an unaligned load. This patch updates infer-alignment to increase the alignment value of the alignment argument if it is known greater than the provided one. Ignoring the gather / scatter versions for now since they contain many pointers. --- llvm/lib/Transforms/Scalar/InferAlignment.cpp | 37 ++++++++++++++++++- llvm/test/Transforms/InferAlignment/masked.ll | 34 +++++++++++++++++ 2 files changed, 69 insertions(+), 2 deletions(-) create mode 100644 llvm/test/Transforms/InferAlignment/masked.ll diff --git a/llvm/lib/Transforms/Scalar/InferAlignment.cpp b/llvm/lib/Transforms/Scalar/InferAlignment.cpp index e9bf59c6850a3..e7d4db401aa51 100644 --- a/llvm/lib/Transforms/Scalar/InferAlignment.cpp +++ b/llvm/lib/Transforms/Scalar/InferAlignment.cpp @@ -14,7 +14,9 @@ #include "llvm/Transforms/Scalar/InferAlignment.h" #include "llvm/Analysis/AssumptionCache.h" #include "llvm/Analysis/ValueTracking.h" +#include "llvm/IR/IRBuilder.h" #include "llvm/IR/Instructions.h" +#include "llvm/IR/IntrinsicInst.h" #include "llvm/Support/KnownBits.h" #include "llvm/Transforms/Scalar.h" #include "llvm/Transforms/Utils/Local.h" @@ -35,8 +37,39 @@ static bool tryToImproveAlign( return true; } } - // TODO: Also handle memory intrinsics. - return false; + + IntrinsicInst *II = dyn_cast(I); + if (!II) + return false; + + // TODO: Handle more memory intrinsics. + switch (II->getIntrinsicID()) { + case Intrinsic::masked_load: + case Intrinsic::masked_store: { + Value *PtrOp = II->getIntrinsicID() == Intrinsic::masked_load + ? II->getArgOperand(0) + : II->getArgOperand(1); + Value *AlignOp = II->getIntrinsicID() == Intrinsic::masked_load + ? II->getArgOperand(1) + : II->getArgOperand(2); + + Align OldAlign = cast(AlignOp)->getAlignValue(); + Align PrefAlign = getKnownAlignment(PtrOp, DL, II); + Align NewAlign = Fn(PtrOp, OldAlign, PrefAlign); + if (NewAlign <= OldAlign) + return false; + + Value *V = llvm::ConstantInt::get(llvm::Type::getInt32Ty(II->getContext()), + NewAlign.value()); + if (II->getIntrinsicID() == Intrinsic::masked_load) + II->setOperand(1, V); + else + II->setOperand(2, V); + return true; + } + default: + return false; + } } bool inferAlignment(Function &F, AssumptionCache &AC, DominatorTree &DT) { diff --git a/llvm/test/Transforms/InferAlignment/masked.ll b/llvm/test/Transforms/InferAlignment/masked.ll new file mode 100644 index 0000000000000..1b8d26417d75e --- /dev/null +++ b/llvm/test/Transforms/InferAlignment/masked.ll @@ -0,0 +1,34 @@ +; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5 +; RUN: opt < %s -passes=infer-alignment -S | FileCheck %s + +define <2 x i32> @load(<2 x i1> %mask, ptr %ptr) { +; CHECK-LABEL: define <2 x i32> @load( +; CHECK-SAME: <2 x i1> [[MASK:%.*]], ptr [[PTR:%.*]]) { +; CHECK-NEXT: [[ENTRY:.*:]] +; CHECK-NEXT: call void @llvm.assume(i1 true) [ "align"(ptr [[PTR]], i64 64) ] +; CHECK-NEXT: [[MASKED_LOAD:%.*]] = call <2 x i32> @llvm.masked.load.v2i32.p0(ptr [[PTR]], i32 64, <2 x i1> [[MASK]], <2 x i32> poison) +; CHECK-NEXT: ret <2 x i32> [[MASKED_LOAD]] +; +entry: + call void @llvm.assume(i1 true) [ "align"(ptr %ptr, i64 64) ] + %masked_load = call <2 x i32> @llvm.masked.load.v2i32.p0(ptr %ptr, i32 1, <2 x i1> %mask, <2 x i32> poison) + ret <2 x i32> %masked_load +} + +define void @store(<2 x i1> %mask, <2 x i32> %val, ptr %ptr) { +; CHECK-LABEL: define void @store( +; CHECK-SAME: <2 x i1> [[MASK:%.*]], <2 x i32> [[VAL:%.*]], ptr [[PTR:%.*]]) { +; CHECK-NEXT: [[ENTRY:.*:]] +; CHECK-NEXT: call void @llvm.assume(i1 true) [ "align"(ptr [[PTR]], i64 64) ] +; CHECK-NEXT: tail call void @llvm.masked.store.v2i32.p0(<2 x i32> [[VAL]], ptr [[PTR]], i32 64, <2 x i1> [[MASK]]) +; CHECK-NEXT: ret void +; +entry: + call void @llvm.assume(i1 true) [ "align"(ptr %ptr, i64 64) ] + tail call void @llvm.masked.store.v2i32.p0(<2 x i32> %val, ptr %ptr, i32 1, <2 x i1> %mask) + ret void +} + +declare void @llvm.assume(i1) +declare <2 x i32> @llvm.masked.load.v2i32.p0(ptr, i32, <2 x i1>, <2 x i32>) +declare void @llvm.masked.store.v2i32.p0(<2 x i32>, ptr, i32, <2 x i1>) From a74177b6d3cc2457faae87cbef16059ea58ea415 Mon Sep 17 00:00:00 2001 From: Joseph Huber Date: Sat, 30 Aug 2025 16:04:54 -0500 Subject: [PATCH 2/3] comments --- llvm/lib/Transforms/Scalar/InferAlignment.cpp | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/llvm/lib/Transforms/Scalar/InferAlignment.cpp b/llvm/lib/Transforms/Scalar/InferAlignment.cpp index e7d4db401aa51..9cdcfbc7cc788 100644 --- a/llvm/lib/Transforms/Scalar/InferAlignment.cpp +++ b/llvm/lib/Transforms/Scalar/InferAlignment.cpp @@ -14,7 +14,6 @@ #include "llvm/Transforms/Scalar/InferAlignment.h" #include "llvm/Analysis/AssumptionCache.h" #include "llvm/Analysis/ValueTracking.h" -#include "llvm/IR/IRBuilder.h" #include "llvm/IR/Instructions.h" #include "llvm/IR/IntrinsicInst.h" #include "llvm/Support/KnownBits.h" @@ -52,15 +51,18 @@ static bool tryToImproveAlign( Value *AlignOp = II->getIntrinsicID() == Intrinsic::masked_load ? II->getArgOperand(1) : II->getArgOperand(2); + Type *Type = II->getIntrinsicID() == Intrinsic::masked_load + ? II->getType() + : II->getArgOperand(0)->getType(); Align OldAlign = cast(AlignOp)->getAlignValue(); - Align PrefAlign = getKnownAlignment(PtrOp, DL, II); + Align PrefAlign = DL.getPrefTypeAlign(Type); Align NewAlign = Fn(PtrOp, OldAlign, PrefAlign); if (NewAlign <= OldAlign) return false; - Value *V = llvm::ConstantInt::get(llvm::Type::getInt32Ty(II->getContext()), - NewAlign.value()); + Value *V = + ConstantInt::get(Type::getInt32Ty(II->getContext()), NewAlign.value()); if (II->getIntrinsicID() == Intrinsic::masked_load) II->setOperand(1, V); else From b07bc2736a2eb0b944479b9b5394b646d49546cd Mon Sep 17 00:00:00 2001 From: Joseph Huber Date: Sat, 30 Aug 2025 20:08:15 -0500 Subject: [PATCH 3/3] comments --- llvm/lib/Transforms/Scalar/InferAlignment.cpp | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/llvm/lib/Transforms/Scalar/InferAlignment.cpp b/llvm/lib/Transforms/Scalar/InferAlignment.cpp index 9cdcfbc7cc788..b60b15b6c3a2b 100644 --- a/llvm/lib/Transforms/Scalar/InferAlignment.cpp +++ b/llvm/lib/Transforms/Scalar/InferAlignment.cpp @@ -45,17 +45,16 @@ static bool tryToImproveAlign( switch (II->getIntrinsicID()) { case Intrinsic::masked_load: case Intrinsic::masked_store: { + int AlignOpIdx = II->getIntrinsicID() == Intrinsic::masked_load ? 1 : 2; Value *PtrOp = II->getIntrinsicID() == Intrinsic::masked_load ? II->getArgOperand(0) : II->getArgOperand(1); - Value *AlignOp = II->getIntrinsicID() == Intrinsic::masked_load - ? II->getArgOperand(1) - : II->getArgOperand(2); Type *Type = II->getIntrinsicID() == Intrinsic::masked_load ? II->getType() : II->getArgOperand(0)->getType(); - Align OldAlign = cast(AlignOp)->getAlignValue(); + Align OldAlign = + cast(II->getArgOperand(AlignOpIdx))->getAlignValue(); Align PrefAlign = DL.getPrefTypeAlign(Type); Align NewAlign = Fn(PtrOp, OldAlign, PrefAlign); if (NewAlign <= OldAlign) @@ -63,10 +62,7 @@ static bool tryToImproveAlign( Value *V = ConstantInt::get(Type::getInt32Ty(II->getContext()), NewAlign.value()); - if (II->getIntrinsicID() == Intrinsic::masked_load) - II->setOperand(1, V); - else - II->setOperand(2, V); + II->setOperand(AlignOpIdx, V); return true; } default: