From af3207bae943396476e141f2859719f63b5174d4 Mon Sep 17 00:00:00 2001 From: Michael Berg Date: Tue, 24 Jun 2025 12:32:18 -0700 Subject: [PATCH 1/2] [LoopDist] Add some runtime checks for cross partition loads Emit safety guards for ptr accesses when cross partition loads exist which have a corresponding store to the same address in a different partition. This will emit the necessary ptr checks for these accesses. --- llvm/lib/Transforms/Scalar/LoopDistribute.cpp | 15 ++ .../LoopDistribute/cross-partition-access.ll | 159 ++++++++++++++++++ 2 files changed, 174 insertions(+) create mode 100644 llvm/test/Transforms/LoopDistribute/cross-partition-access.ll diff --git a/llvm/lib/Transforms/Scalar/LoopDistribute.cpp b/llvm/lib/Transforms/Scalar/LoopDistribute.cpp index 0ac1a15981770..9800e8ba23486 100644 --- a/llvm/lib/Transforms/Scalar/LoopDistribute.cpp +++ b/llvm/lib/Transforms/Scalar/LoopDistribute.cpp @@ -521,6 +521,21 @@ class InstPartitionContainer { Partition = -1; } assert(Partition != -2 && "Pointer not belonging to any partition"); + // All the store context uses of our address were processed, + // Now make sure we don't have cross partition loads. + if (RtPtrCheck->Pointers[I].IsWritePtr) { + if (Ptr->hasOneUse() || Partition == -1) + continue; + + for (User *U : Ptr->users()) + if (auto *CurLoad = dyn_cast(U)) + if (L->contains(CurLoad->getParent())) + if (Partition != (int)this->InstToPartitionId[CurLoad]) { + // -1 means belonging to multiple partitions. + Partition = -1; + break; + } + } } return PtrToPartitions; diff --git a/llvm/test/Transforms/LoopDistribute/cross-partition-access.ll b/llvm/test/Transforms/LoopDistribute/cross-partition-access.ll new file mode 100644 index 0000000000000..6e1106c3277a7 --- /dev/null +++ b/llvm/test/Transforms/LoopDistribute/cross-partition-access.ll @@ -0,0 +1,159 @@ +; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5 +; RUN: opt -aa-pipeline=basic-aa -passes=loop-distribute -enable-loop-distribute -verify-loop-info -verify-dom-info -S %s | FileCheck %s + +; Test emit safety guards for ptr access of %a and %c on a cross parition load which +; has a corresponding store to the same address. This ensures that if %a and %c +; overlap in some scenarios, that we execute the original loop for safety reasons. + +define dso_local void @_Z13distribution3PiS_S_S_i(ptr nocapture noundef %a, ptr nocapture noundef readonly %b, ptr nocapture noundef %c, ptr nocapture noundef writeonly %d, i64 noundef signext %len) { +; CHECK-LABEL: define dso_local void @_Z13distribution3PiS_S_S_i( +; CHECK-SAME: ptr noundef captures(none) [[A:%.*]], ptr noundef readonly captures(none) [[B:%.*]], ptr noundef captures(none) [[C:%.*]], ptr noundef writeonly captures(none) [[D:%.*]], i64 noundef signext [[LEN:%.*]]) { +; CHECK-NEXT: [[ENTRY:.*:]] +; CHECK-NEXT: [[CMP:%.*]] = icmp sgt i64 [[LEN]], 0 +; CHECK-NEXT: br i1 [[CMP]], label %[[END:.*]], label %[[FOR_BODY_LVER_CHECK:.*]] +; CHECK: [[FOR_BODY_LVER_CHECK]]: +; CHECK-NEXT: [[TMP0:%.*]] = shl i64 [[LEN]], 2 +; CHECK-NEXT: [[SCEVGEP2:%.*]] = getelementptr i8, ptr [[A]], i64 [[TMP0]] +; CHECK-NEXT: [[SCEVGEP:%.*]] = getelementptr i8, ptr [[C]], i64 -4 +; CHECK-NEXT: [[SCEVGEP5:%.*]] = getelementptr i8, ptr [[C]], i64 [[TMP0]] +; CHECK-NEXT: [[SCEVGEP3:%.*]] = getelementptr i8, ptr [[D]], i64 [[TMP0]] +; CHECK-NEXT: [[SCEVGEP4:%.*]] = getelementptr i8, ptr [[B]], i64 [[TMP0]] +; CHECK-NEXT: [[BOUND1:%.*]] = icmp ult ptr [[A]], [[SCEVGEP5]] +; CHECK-NEXT: [[BOUND0:%.*]] = icmp ult ptr [[SCEVGEP]], [[SCEVGEP2]] +; CHECK-NEXT: [[FOUND_CONFLICT:%.*]] = and i1 [[BOUND1]], [[BOUND0]] +; CHECK-NEXT: [[BOUND07:%.*]] = icmp ult ptr [[A]], [[SCEVGEP3]] +; CHECK-NEXT: [[BOUND18:%.*]] = icmp ult ptr [[D]], [[SCEVGEP2]] +; CHECK-NEXT: [[FOUND_CONFLICT9:%.*]] = and i1 [[BOUND07]], [[BOUND18]] +; CHECK-NEXT: [[CONFLICT_RDX:%.*]] = or i1 [[FOUND_CONFLICT]], [[FOUND_CONFLICT9]] +; CHECK-NEXT: [[BOUND010:%.*]] = icmp ult ptr [[A]], [[SCEVGEP4]] +; CHECK-NEXT: [[BOUND111:%.*]] = icmp ult ptr [[B]], [[SCEVGEP2]] +; CHECK-NEXT: [[FOUND_CONFLICT12:%.*]] = and i1 [[BOUND010]], [[BOUND111]] +; CHECK-NEXT: [[CONFLICT_RDX13:%.*]] = or i1 [[CONFLICT_RDX]], [[FOUND_CONFLICT12]] +; CHECK-NEXT: [[BOUND014:%.*]] = icmp ult ptr [[SCEVGEP]], [[SCEVGEP3]] +; CHECK-NEXT: [[BOUND115:%.*]] = icmp ult ptr [[D]], [[SCEVGEP5]] +; CHECK-NEXT: [[FOUND_CONFLICT16:%.*]] = and i1 [[BOUND014]], [[BOUND115]] +; CHECK-NEXT: [[CONFLICT_RDX17:%.*]] = or i1 [[CONFLICT_RDX13]], [[FOUND_CONFLICT16]] +; CHECK-NEXT: [[BOUND018:%.*]] = icmp ult ptr [[D]], [[SCEVGEP4]] +; CHECK-NEXT: [[BOUND119:%.*]] = icmp ult ptr [[B]], [[SCEVGEP3]] +; CHECK-NEXT: [[FOUND_CONFLICT20:%.*]] = and i1 [[BOUND018]], [[BOUND119]] +; CHECK-NEXT: [[CONFLICT_RDX21:%.*]] = or i1 [[CONFLICT_RDX17]], [[FOUND_CONFLICT20]] +; CHECK-NEXT: br i1 [[CONFLICT_RDX21]], label %[[FOR_BODY_PH_LVER_ORIG:.*]], label %[[FOR_BODY_PH_LDIST1:.*]] +; CHECK: [[FOR_BODY_PH_LVER_ORIG]]: +; CHECK-NEXT: br label %[[FOR_BODY_LVER_ORIG:.*]] +; CHECK: [[FOR_BODY_LVER_ORIG]]: +; CHECK-NEXT: [[IDXPROM_LVER_ORIG:%.*]] = phi i64 [ 0, %[[FOR_BODY_PH_LVER_ORIG]] ], [ [[I6_LVER_ORIG:%.*]], %[[FOR_BODY_LVER_ORIG]] ] +; CHECK-NEXT: [[ARRAYIDX_LVER_ORIG:%.*]] = getelementptr inbounds i32, ptr [[B]], i64 [[IDXPROM_LVER_ORIG]] +; CHECK-NEXT: [[I2_LVER_ORIG:%.*]] = load i32, ptr [[ARRAYIDX_LVER_ORIG]], align 4, !tbaa [[TBAA0:![0-9]+]] +; CHECK-NEXT: [[ADD4_LVER_ORIG:%.*]] = add nsw i32 [[I2_LVER_ORIG]], 1 +; CHECK-NEXT: [[ARRAYIDX8_LVER_ORIG:%.*]] = getelementptr inbounds i32, ptr [[A]], i64 [[IDXPROM_LVER_ORIG]] +; CHECK-NEXT: store i32 [[ADD4_LVER_ORIG]], ptr [[ARRAYIDX8_LVER_ORIG]], align 4, !tbaa [[TBAA0]] +; CHECK-NEXT: [[I3_LVER_ORIG:%.*]] = getelementptr i32, ptr [[C]], i64 [[IDXPROM_LVER_ORIG]] +; CHECK-NEXT: [[ARRAYIDX17_LVER_ORIG:%.*]] = getelementptr i8, ptr [[I3_LVER_ORIG]], i64 -4 +; CHECK-NEXT: [[I4_LVER_ORIG:%.*]] = load i32, ptr [[ARRAYIDX17_LVER_ORIG]], align 4, !tbaa [[TBAA0]] +; CHECK-NEXT: [[SUB18_LVER_ORIG:%.*]] = sub nsw i32 [[ADD4_LVER_ORIG]], [[I4_LVER_ORIG]] +; CHECK-NEXT: store i32 [[SUB18_LVER_ORIG]], ptr [[I3_LVER_ORIG]], align 4, !tbaa [[TBAA0]] +; CHECK-NEXT: [[I5_LVER_ORIG:%.*]] = load i32, ptr [[ARRAYIDX8_LVER_ORIG]], align 4, !tbaa [[TBAA0]] +; CHECK-NEXT: [[ADD27_LVER_ORIG:%.*]] = add nsw i32 [[I5_LVER_ORIG]], 2 +; CHECK-NEXT: [[ARRAYIDX31_LVER_ORIG:%.*]] = getelementptr inbounds i32, ptr [[D]], i64 [[IDXPROM_LVER_ORIG]] +; CHECK-NEXT: store i32 [[ADD27_LVER_ORIG]], ptr [[ARRAYIDX31_LVER_ORIG]], align 4, !tbaa [[TBAA0]] +; CHECK-NEXT: [[I6_LVER_ORIG]] = add i64 [[IDXPROM_LVER_ORIG]], 1 +; CHECK-NEXT: [[CMP1_NOT_LVER_ORIG:%.*]] = icmp eq i64 [[I6_LVER_ORIG]], [[LEN]] +; CHECK-NEXT: br i1 [[CMP1_NOT_LVER_ORIG]], label %[[END_LOOPEXIT_LOOPEXIT:.*]], label %[[FOR_BODY_LVER_ORIG]], !llvm.loop [[LOOP4:![0-9]+]] +; CHECK: [[FOR_BODY_PH_LDIST1]]: +; CHECK-NEXT: br label %[[FOR_BODY_LDIST1:.*]] +; CHECK: [[FOR_BODY_LDIST1]]: +; CHECK-NEXT: [[IDXPROM_LDIST1:%.*]] = phi i64 [ 0, %[[FOR_BODY_PH_LDIST1]] ], [ [[I6_LDIST1:%.*]], %[[FOR_BODY_LDIST1]] ] +; CHECK-NEXT: [[ARRAYIDX_LDIST1:%.*]] = getelementptr inbounds i32, ptr [[B]], i64 [[IDXPROM_LDIST1]] +; CHECK-NEXT: [[I2_LDIST1:%.*]] = load i32, ptr [[ARRAYIDX_LDIST1]], align 4, !tbaa [[TBAA0]], !alias.scope [[META6:![0-9]+]] +; CHECK-NEXT: [[ADD4_LDIST1:%.*]] = add nsw i32 [[I2_LDIST1]], 1 +; CHECK-NEXT: [[ARRAYIDX8_LDIST1:%.*]] = getelementptr inbounds i32, ptr [[A]], i64 [[IDXPROM_LDIST1]] +; CHECK-NEXT: store i32 [[ADD4_LDIST1]], ptr [[ARRAYIDX8_LDIST1]], align 4, !tbaa [[TBAA0]], !alias.scope [[META9:![0-9]+]], !noalias [[META11:![0-9]+]] +; CHECK-NEXT: [[I3_LDIST1:%.*]] = getelementptr i32, ptr [[C]], i64 [[IDXPROM_LDIST1]] +; CHECK-NEXT: [[ARRAYIDX17_LDIST1:%.*]] = getelementptr i8, ptr [[I3_LDIST1]], i64 -4 +; CHECK-NEXT: [[I4_LDIST1:%.*]] = load i32, ptr [[ARRAYIDX17_LDIST1]], align 4, !tbaa [[TBAA0]], !alias.scope [[META14:![0-9]+]], !noalias [[META15:![0-9]+]] +; CHECK-NEXT: [[SUB18_LDIST1:%.*]] = sub nsw i32 [[ADD4_LDIST1]], [[I4_LDIST1]] +; CHECK-NEXT: store i32 [[SUB18_LDIST1]], ptr [[I3_LDIST1]], align 4, !tbaa [[TBAA0]], !alias.scope [[META14]], !noalias [[META15]] +; CHECK-NEXT: [[I6_LDIST1]] = add i64 [[IDXPROM_LDIST1]], 1 +; CHECK-NEXT: [[CMP1_NOT_LDIST1:%.*]] = icmp eq i64 [[I6_LDIST1]], [[LEN]] +; CHECK-NEXT: br i1 [[CMP1_NOT_LDIST1]], label %[[FOR_BODY_PH:.*]], label %[[FOR_BODY_LDIST1]], !llvm.loop [[LOOP16:![0-9]+]] +; CHECK: [[FOR_BODY_PH]]: +; CHECK-NEXT: br label %[[FOR_BODY:.*]] +; CHECK: [[FOR_BODY]]: +; CHECK-NEXT: [[IDXPROM:%.*]] = phi i64 [ 0, %[[FOR_BODY_PH]] ], [ [[I6:%.*]], %[[FOR_BODY]] ] +; CHECK-NEXT: [[ARRAYIDX8:%.*]] = getelementptr inbounds i32, ptr [[A]], i64 [[IDXPROM]] +; CHECK-NEXT: [[I5:%.*]] = load i32, ptr [[ARRAYIDX8]], align 4, !tbaa [[TBAA0]], !alias.scope [[META9]], !noalias [[META11]] +; CHECK-NEXT: [[ADD27:%.*]] = add nsw i32 [[I5]], 2 +; CHECK-NEXT: [[ARRAYIDX31:%.*]] = getelementptr inbounds i32, ptr [[D]], i64 [[IDXPROM]] +; CHECK-NEXT: store i32 [[ADD27]], ptr [[ARRAYIDX31]], align 4, !tbaa [[TBAA0]], !alias.scope [[META15]], !noalias [[META6]] +; CHECK-NEXT: [[I6]] = add i64 [[IDXPROM]], 1 +; CHECK-NEXT: [[CMP1_NOT:%.*]] = icmp eq i64 [[I6]], [[LEN]] +; CHECK-NEXT: br i1 [[CMP1_NOT]], label %[[END_LOOPEXIT_LOOPEXIT20:.*]], label %[[FOR_BODY]], !llvm.loop [[LOOP16]] +; CHECK: [[END_LOOPEXIT_LOOPEXIT]]: +; CHECK-NEXT: br label %[[END_LOOPEXIT:.*]] +; CHECK: [[END_LOOPEXIT_LOOPEXIT20]]: +; CHECK-NEXT: br label %[[END_LOOPEXIT]] +; CHECK: [[END_LOOPEXIT]]: +; CHECK-NEXT: br label %[[END]] +; CHECK: [[END]]: +; CHECK-NEXT: ret void +; +entry: + %cmp = icmp sgt i64 %len, 0 + br i1 %cmp, label %end, label %for.body.preheader + +for.body.preheader: ; preds = %entry + br label %for.body + +for.body: ; preds = %for.body, %for.body.preheader + %indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %for.body ] + %arrayidx = getelementptr inbounds i32, ptr %b, i64 %indvars.iv + %i2 = load i32, ptr %arrayidx, align 4, !tbaa !0 + %add4 = add nsw i32 %i2, 1 + %arrayidx8 = getelementptr inbounds i32, ptr %a, i64 %indvars.iv + store i32 %add4, ptr %arrayidx8, align 4, !tbaa !0 + %i3 = getelementptr i32, ptr %c, i64 %indvars.iv + %arrayidx17 = getelementptr i8, ptr %i3, i64 -4 + %i4 = load i32, ptr %arrayidx17, align 4, !tbaa !0 + %sub18 = sub nsw i32 %add4, %i4 + store i32 %sub18, ptr %i3, align 4, !tbaa !0 + %i5 = load i32, ptr %arrayidx8, align 4, !tbaa !0 + %add27 = add nsw i32 %i5, 2 + %arrayidx31 = getelementptr inbounds i32, ptr %d, i64 %indvars.iv + store i32 %add27, ptr %arrayidx31, align 4, !tbaa !0 + %indvars.iv.next = add i64 %indvars.iv, 1 + %cmp1.not = icmp eq i64 %indvars.iv.next, %len + br i1 %cmp1.not, label %end.loopexit, label %for.body, !llvm.loop !4 + +end.loopexit: ; preds = %for.body + br label %end + +end: ; preds = %end.loopexit, %entry + ret void +} + +!0 = !{!1, !1, i64 0} +!1 = !{!"int", !2, i64 0} +!2 = !{!"omnipotent char", !3, i64 0} +!3 = !{!"Simple C++ TBAA"} +!4 = distinct !{!4, !5} +!5 = !{!"llvm.loop.mustprogress"} + +;. +; CHECK: [[TBAA0]] = !{[[META1:![0-9]+]], [[META1]], i64 0} +; CHECK: [[META1]] = !{!"int", [[META2:![0-9]+]], i64 0} +; CHECK: [[META2]] = !{!"omnipotent char", [[META3:![0-9]+]], i64 0} +; CHECK: [[META3]] = !{!"Simple C++ TBAA"} +; CHECK: [[LOOP4]] = distinct !{[[LOOP4]], [[META5:![0-9]+]]} +; CHECK: [[META5]] = !{!"llvm.loop.mustprogress"} +; CHECK: [[META6]] = !{[[META7:![0-9]+]]} +; CHECK: [[META7]] = distinct !{[[META7]], [[META8:![0-9]+]]} +; CHECK: [[META8]] = distinct !{[[META8]], !"LVerDomain"} +; CHECK: [[META9]] = !{[[META10:![0-9]+]]} +; CHECK: [[META10]] = distinct !{[[META10]], [[META8]]} +; CHECK: [[META11]] = !{[[META12:![0-9]+]], [[META13:![0-9]+]], [[META7]]} +; CHECK: [[META12]] = distinct !{[[META12]], [[META8]]} +; CHECK: [[META13]] = distinct !{[[META13]], [[META8]]} +; CHECK: [[META14]] = !{[[META12]]} +; CHECK: [[META15]] = !{[[META13]]} +; CHECK: [[LOOP16]] = distinct !{[[LOOP16]], [[META5]]} +;. From 0dadf65b1a3b6b41d01ae289e3e744776e7450d8 Mon Sep 17 00:00:00 2001 From: Michael Berg Date: Mon, 4 Aug 2025 12:48:11 -0700 Subject: [PATCH 2/2] [LoopDist] Consider reads and writes together for runtime checks Add back in getInstructionsForAccess as query for Instructions modifying memory to mark cross partition accesses. --- llvm/lib/Analysis/LoopAccessAnalysis.cpp | 11 +++++----- llvm/lib/Transforms/Scalar/LoopDistribute.cpp | 21 ++++--------------- 2 files changed, 10 insertions(+), 22 deletions(-) diff --git a/llvm/lib/Analysis/LoopAccessAnalysis.cpp b/llvm/lib/Analysis/LoopAccessAnalysis.cpp index a5535339a714f..62baf9b632bc7 100644 --- a/llvm/lib/Analysis/LoopAccessAnalysis.cpp +++ b/llvm/lib/Analysis/LoopAccessAnalysis.cpp @@ -2404,12 +2404,13 @@ bool MemoryDepChecker::areDepsSafe(const DepCandidates &DepCands, SmallVector MemoryDepChecker::getInstructionsForAccess(Value *Ptr, bool IsWrite) const { MemAccessInfo Access(Ptr, IsWrite); - auto &IndexVector = Accesses.find(Access)->second; - + auto I = Accesses.find(Access); SmallVector Insts; - transform(IndexVector, - std::back_inserter(Insts), - [&](unsigned Idx) { return this->InstMap[Idx]; }); + if (I != Accesses.end()) { + transform(I->second, std::back_inserter(Insts), + [&](unsigned Idx) { return this->InstMap[Idx]; }); + } + return Insts; } diff --git a/llvm/lib/Transforms/Scalar/LoopDistribute.cpp b/llvm/lib/Transforms/Scalar/LoopDistribute.cpp index 9800e8ba23486..27d3004d81947 100644 --- a/llvm/lib/Transforms/Scalar/LoopDistribute.cpp +++ b/llvm/lib/Transforms/Scalar/LoopDistribute.cpp @@ -502,8 +502,10 @@ class InstPartitionContainer { SmallVector PtrToPartitions(N); for (unsigned I = 0; I < N; ++I) { Value *Ptr = RtPtrCheck->Pointers[I].PointerValue; - auto Instructions = - LAI.getInstructionsForAccess(Ptr, RtPtrCheck->Pointers[I].IsWritePtr); + auto Instructions = LAI.getInstructionsForAccess(Ptr, /* IsWrite */ true); + auto ReadInstructions = + LAI.getInstructionsForAccess(Ptr, /* IsWrite */ false); + Instructions.append(ReadInstructions.begin(), ReadInstructions.end()); int &Partition = PtrToPartitions[I]; // First set it to uninitialized. @@ -521,21 +523,6 @@ class InstPartitionContainer { Partition = -1; } assert(Partition != -2 && "Pointer not belonging to any partition"); - // All the store context uses of our address were processed, - // Now make sure we don't have cross partition loads. - if (RtPtrCheck->Pointers[I].IsWritePtr) { - if (Ptr->hasOneUse() || Partition == -1) - continue; - - for (User *U : Ptr->users()) - if (auto *CurLoad = dyn_cast(U)) - if (L->contains(CurLoad->getParent())) - if (Partition != (int)this->InstToPartitionId[CurLoad]) { - // -1 means belonging to multiple partitions. - Partition = -1; - break; - } - } } return PtrToPartitions;