Adel Ejjeh 49250284cf
[AMDGPU][LoopUnroll] Enable allowexpensivetripcount for amdgpu when unroll pragmas are present (#181241)
This PR is intended as an AMDGPU-specific solution for #181267 while
discussions on changing the default behavior for all targets continue in
that PR.

Problem:
Loops with an explicit unroll pragma (#pragma unroll / #pragma clang
loop unroll(enable)) that have an expensive runtime trip count currently
don't get unrolled because UP.AllowExpensiveTripCount defaults to false.
The pragma is silently ignored. This is not the case when an unroll
factor is specified (PragmaCount > 0), where the pass sets
UP.AllowExpensiveTripCount = true.

Solution:
Set UP.AllowExpensiveTripCount to true for for loops that have an unroll
pramga in the AMDGPU TTI Implementation.
I've added a new lit test expensive-tripcount.ll that verifies
pragma-driven unrolling with expensive trip counts will work as
expected.

The change showed no meaningful regressions across a few different
workloads from Composable Kernels (CK) and llama.cpp as well as Pytorch
kernels on AMDGPU gfx950. Additionally, the change improves the
performance of PyTorch reduction loops on AMDGPU targets.
2026-03-13 09:56:20 -07:00

100 lines
6.4 KiB
LLVM

; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
; RUN: opt -mtriple=amdgcn-unknown-amdhsa < %s -S -passes=loop-unroll | FileCheck %s
; Checks that loops with expensive trip counts are unrolled when the loop.unroll.enable metadata is present.
; Function Attrs: nofree norecurse nosync nounwind memory(argmem: readwrite)
define dso_local void @complex_loop_unroll(i64 noundef %input_offset, i64 noundef %step, i64 noundef %n) {
; CHECK-LABEL: define dso_local void @complex_loop_unroll(
; CHECK-SAME: i64 noundef [[INPUT_OFFSET:%.*]], i64 noundef [[STEP:%.*]], i64 noundef [[N:%.*]]) {
; CHECK-NEXT: [[FOR_BODY_LR_PH:.*]]:
; CHECK-NEXT: [[TMP3:%.*]] = add nsw i64 [[STEP]], [[INPUT_OFFSET]]
; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 [[TMP3]])
; CHECK-NEXT: [[TMP4:%.*]] = sub i64 [[SMAX]], [[TMP3]]
; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP4]], i64 1)
; CHECK-NEXT: [[TMP5:%.*]] = sub i64 [[SMAX]], [[UMIN]]
; CHECK-NEXT: [[TMP6:%.*]] = sub i64 [[TMP5]], [[TMP3]]
; CHECK-NEXT: [[UMAX:%.*]] = call i64 @llvm.umax.i64(i64 [[STEP]], i64 1)
; CHECK-NEXT: [[TMP7:%.*]] = udiv i64 [[TMP6]], [[UMAX]]
; CHECK-NEXT: [[TMP8:%.*]] = add i64 [[UMIN]], [[TMP7]]
; CHECK-NEXT: [[TMP9:%.*]] = add i64 [[TMP8]], 1
; CHECK-NEXT: [[XTRAITER:%.*]] = and i64 [[TMP9]], 7
; CHECK-NEXT: [[TMP12:%.*]] = icmp ult i64 [[TMP8]], 7
; CHECK-NEXT: br i1 [[TMP12]], label %[[FOR_END14_LOOPEXIT1:.*]], label %[[FOR_BODY_LR_PH_NEW:.*]]
; CHECK: [[FOR_BODY_LR_PH_NEW]]:
; CHECK-NEXT: [[UNROLL_ITER:%.*]] = sub i64 [[TMP9]], [[XTRAITER]]
; CHECK-NEXT: br label %[[FOR_BODY1:.*]]
; CHECK: [[FOR_BODY1]]:
; CHECK-NEXT: [[TMP11:%.*]] = phi i64 [ 0, %[[FOR_BODY_LR_PH_NEW]] ], [ [[ADD_7:%.*]], %[[FOR_BODY1]] ]
; CHECK-NEXT: [[INDVARS_IV:%.*]] = phi i64 [ [[INPUT_OFFSET]], %[[FOR_BODY_LR_PH_NEW]] ], [ [[INDVARS_IV_NEXT_7:%.*]], %[[FOR_BODY1]] ]
; CHECK-NEXT: [[NITER:%.*]] = phi i64 [ 0, %[[FOR_BODY_LR_PH_NEW]] ], [ [[NITER_NEXT_7:%.*]], %[[FOR_BODY1]] ]
; CHECK-NEXT: [[ADD:%.*]] = add i64 [[TMP11]], [[INDVARS_IV]]
; CHECK-NEXT: [[INDVARS_IV_NEXT:%.*]] = add nsw i64 [[INDVARS_IV]], [[STEP]]
; CHECK-NEXT: [[ADD_1:%.*]] = add i64 [[ADD]], [[INDVARS_IV_NEXT]]
; CHECK-NEXT: [[INDVARS_IV_NEXT_1:%.*]] = add nsw i64 [[INDVARS_IV_NEXT]], [[STEP]]
; CHECK-NEXT: [[ADD_2:%.*]] = add i64 [[ADD_1]], [[INDVARS_IV_NEXT_1]]
; CHECK-NEXT: [[INDVARS_IV_NEXT_2:%.*]] = add nsw i64 [[INDVARS_IV_NEXT_1]], [[STEP]]
; CHECK-NEXT: [[ADD_3:%.*]] = add i64 [[ADD_2]], [[INDVARS_IV_NEXT_2]]
; CHECK-NEXT: [[INDVARS_IV_NEXT_3:%.*]] = add nsw i64 [[INDVARS_IV_NEXT_2]], [[STEP]]
; CHECK-NEXT: [[ADD_4:%.*]] = add i64 [[ADD_3]], [[INDVARS_IV_NEXT_3]]
; CHECK-NEXT: [[INDVARS_IV_NEXT_4:%.*]] = add nsw i64 [[INDVARS_IV_NEXT_3]], [[STEP]]
; CHECK-NEXT: [[ADD_5:%.*]] = add i64 [[ADD_4]], [[INDVARS_IV_NEXT_4]]
; CHECK-NEXT: [[INDVARS_IV_NEXT_5:%.*]] = add nsw i64 [[INDVARS_IV_NEXT_4]], [[STEP]]
; CHECK-NEXT: [[ADD_6:%.*]] = add i64 [[ADD_5]], [[INDVARS_IV_NEXT_5]]
; CHECK-NEXT: [[INDVARS_IV_NEXT_9:%.*]] = add nsw i64 [[INDVARS_IV_NEXT_5]], [[STEP]]
; CHECK-NEXT: [[ADD_7]] = add i64 [[ADD_6]], [[INDVARS_IV_NEXT_9]]
; CHECK-NEXT: [[INDVARS_IV_NEXT_7]] = add nsw i64 [[INDVARS_IV_NEXT_9]], [[STEP]]
; CHECK-NEXT: [[NITER_NEXT_7]] = add i64 [[NITER]], 8
; CHECK-NEXT: [[NITER_NCMP_7:%.*]] = icmp ne i64 [[NITER_NEXT_7]], [[UNROLL_ITER]]
; CHECK-NEXT: br i1 [[NITER_NCMP_7]], label %[[FOR_BODY1]], label %[[FOR_END14_LOOPEXIT_UNR_LCSSA1:.*]], !llvm.loop [[LOOP0:![0-9]+]]
; CHECK: [[FOR_END14_LOOPEXIT_UNR_LCSSA1]]:
; CHECK-NEXT: [[DOTUNR:%.*]] = phi i64 [ [[ADD_7]], %[[FOR_BODY1]] ]
; CHECK-NEXT: [[INDVARS_IV_UNR:%.*]] = phi i64 [ [[INDVARS_IV_NEXT_7]], %[[FOR_BODY1]] ]
; CHECK-NEXT: [[LCMP_MOD:%.*]] = icmp ne i64 [[XTRAITER]], 0
; CHECK-NEXT: br i1 [[LCMP_MOD]], label %[[FOR_END14_LOOPEXIT1]], label %[[FOR_END14_LOOPEXIT:.*]]
; CHECK: [[FOR_END14_LOOPEXIT1]]:
; CHECK-NEXT: [[DOTEPIL_INIT:%.*]] = phi i64 [ 0, %[[FOR_BODY_LR_PH]] ], [ [[DOTUNR]], %[[FOR_END14_LOOPEXIT_UNR_LCSSA1]] ]
; CHECK-NEXT: [[INDVARS_IV_EPIL_INIT:%.*]] = phi i64 [ [[INPUT_OFFSET]], %[[FOR_BODY_LR_PH]] ], [ [[INDVARS_IV_UNR]], %[[FOR_END14_LOOPEXIT_UNR_LCSSA1]] ]
; CHECK-NEXT: [[LCMP_MOD1:%.*]] = icmp ne i64 [[XTRAITER]], 0
; CHECK-NEXT: call void @llvm.assume(i1 [[LCMP_MOD1]])
; CHECK-NEXT: br label %[[FOR_BODY_EPIL:.*]]
; CHECK: [[FOR_BODY_EPIL]]:
; CHECK-NEXT: [[TMP13:%.*]] = phi i64 [ [[DOTEPIL_INIT]], %[[FOR_END14_LOOPEXIT1]] ], [ [[ADD_EPIL:%.*]], %[[FOR_BODY_EPIL]] ]
; CHECK-NEXT: [[INDVARS_IV_EPIL:%.*]] = phi i64 [ [[INDVARS_IV_EPIL_INIT]], %[[FOR_END14_LOOPEXIT1]] ], [ [[INDVARS_IV_NEXT_EPIL:%.*]], %[[FOR_BODY_EPIL]] ]
; CHECK-NEXT: [[EPIL_ITER:%.*]] = phi i64 [ 0, %[[FOR_END14_LOOPEXIT1]] ], [ [[EPIL_ITER_NEXT:%.*]], %[[FOR_BODY_EPIL]] ]
; CHECK-NEXT: [[ADD_EPIL]] = add i64 [[TMP13]], [[INDVARS_IV_EPIL]]
; CHECK-NEXT: [[INDVARS_IV_NEXT_EPIL]] = add nsw i64 [[INDVARS_IV_EPIL]], [[STEP]]
; CHECK-NEXT: [[CMP_EPIL:%.*]] = icmp slt i64 [[INDVARS_IV_NEXT_EPIL]], [[N]]
; CHECK-NEXT: [[EPIL_ITER_NEXT]] = add i64 [[EPIL_ITER]], 1
; CHECK-NEXT: [[EPIL_ITER_CMP:%.*]] = icmp ne i64 [[EPIL_ITER_NEXT]], [[XTRAITER]]
; CHECK-NEXT: br i1 [[EPIL_ITER_CMP]], label %[[FOR_BODY_EPIL]], label %[[FOR_END14_LOOPEXIT_UNR_LCSSA:.*]], !llvm.loop [[LOOP3:![0-9]+]]
; CHECK: [[FOR_END14_LOOPEXIT_UNR_LCSSA]]:
; CHECK-NEXT: br label %[[FOR_END14_LOOPEXIT]]
; CHECK: [[FOR_END14_LOOPEXIT]]:
; CHECK-NEXT: ret void
;
entry:
br label %for.body
for.body: ; preds = %for.body, %for.body.lr.ph
%3 = phi i64 [ 0, %entry ], [ %add, %for.body ]
%indvars.iv = phi i64 [ %input_offset, %entry ], [ %indvars.iv.next, %for.body ]
%add = add i64 %3, %indvars.iv
%indvars.iv.next = add nsw i64 %indvars.iv, %step
%cmp = icmp slt i64 %indvars.iv.next, %n
br i1 %cmp, label %for.body, label %for.end14, !llvm.loop !0
for.end14: ; preds = %for.body, %entry
ret void
}
!0 = distinct !{!0, !1, !2}
!1 = !{!"llvm.loop.mustprogress"}
!2 = !{!"llvm.loop.unroll.enable"}
;.
; CHECK: [[LOOP0]] = distinct !{[[LOOP0]], [[META1:![0-9]+]], [[META2:![0-9]+]]}
; CHECK: [[META1]] = !{!"llvm.loop.mustprogress"}
; CHECK: [[META2]] = !{!"llvm.loop.unroll.disable"}
; CHECK: [[LOOP3]] = distinct !{[[LOOP3]], [[META2]]}
;.