[flang] Add loop annotation attributes to the loop backedge (#126082)

Flang currently adds loop metadata to a conditional branch in the loop
preheader, while clang adds it to the loop latch's branch instruction.
Langref says:

> Currently, loop metadata is implemented as metadata attached to the
branch instruction in the loop latch block.
>
> https://llvm.org/docs/LangRef.html#llvm-loop

I misread langref a couple times, but I think this is the appropriate
branch op for the LoopAnnotationAttr. In a couple examples I found that
the metadata was lost entirely during canonicalization. This patch makes
the codegen look more like clang's and the annotations persist through
codegen.

* current clang: https://godbolt.org/z/8WhbcrnG3
* current flang: https://godbolt.org/z/TrPboqqcn
This commit is contained in:
Asher Mancinelli
2025-05-14 07:07:57 -07:00
committed by GitHub
parent 65a9a3d5b1
commit f486cc4417
5 changed files with 44 additions and 19 deletions

View File

@@ -123,23 +123,24 @@ public:
: terminator->operand_begin();
loopCarried.append(begin, terminator->operand_end());
loopCarried.push_back(itersMinusOne);
rewriter.create<mlir::cf::BranchOp>(loc, conditionalBlock, loopCarried);
auto backEdge =
rewriter.create<mlir::cf::BranchOp>(loc, conditionalBlock, loopCarried);
rewriter.eraseOp(terminator);
// Copy loop annotations from the do loop to the loop back edge.
if (auto ann = loop.getLoopAnnotation())
backEdge->setAttr("loop_annotation", *ann);
// Conditional block
rewriter.setInsertionPointToEnd(conditionalBlock);
auto zero = rewriter.create<mlir::arith::ConstantIndexOp>(loc, 0);
auto comparison = rewriter.create<mlir::arith::CmpIOp>(
loc, arith::CmpIPredicate::sgt, itersLeft, zero);
auto cond = rewriter.create<mlir::cf::CondBranchOp>(
rewriter.create<mlir::cf::CondBranchOp>(
loc, comparison, firstBlock, llvm::ArrayRef<mlir::Value>(), endBlock,
llvm::ArrayRef<mlir::Value>());
// Copy loop annotations from the do loop to the loop entry condition.
if (auto ann = loop.getLoopAnnotation())
cond->setAttr("loop_annotation", *ann);
// The result of the loop operation is the values of the condition block
// arguments except the induction variable on the last iteration.
auto args = loop.getFinalValue()

View File

@@ -13,7 +13,9 @@ func.func @_QPvector_always() -> i32 {
%c10_i32 = arith.constant 10 : i32
%c1_i32 = arith.constant 1 : i32
%c10 = arith.constant 10 : index
// CHECK: cf.cond_br %{{.*}}, ^{{.*}}, ^{{.*}} {loop_annotation = #[[ANNOTATION]]}
// CHECK: cf.cond_br
// CHECK-NOT: loop_annotation
// CHECK: cf.br ^{{.*}} {loop_annotation = #[[ANNOTATION]]}
%8:2 = fir.do_loop %arg0 = %c1 to %c10 step %c1 iter_args(%arg1 = %c1_i32) -> (index, i32) attributes {loopAnnotation = #loop_annotation} {
fir.result %c1, %c1_i32 : index, i32
}

View File

@@ -3,8 +3,10 @@
! CHECK-LABEL: unroll_dir
subroutine unroll_dir
integer :: a(10)
!dir$ unroll
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_ENABLE_FULL_ANNO:.*]]
!dir$ unroll
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[UNROLL_ENABLE_FULL_ANNO:.*]]
do i=1,10
a(i)=i
end do
@@ -14,7 +16,9 @@ end subroutine unroll_dir
subroutine unroll_dir_0
integer :: a(10)
!dir$ unroll 0
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_DISABLE_ANNO:.*]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[UNROLL_DISABLE_ANNO:.*]]
do i=1,10
a(i)=i
end do
@@ -24,7 +28,9 @@ end subroutine unroll_dir_0
subroutine unroll_dir_1
integer :: a(10)
!dir$ unroll 1
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_DISABLE_ANNO]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[UNROLL_DISABLE_ANNO]]
do i=1,10
a(i)=i
end do
@@ -34,7 +40,9 @@ end subroutine unroll_dir_1
subroutine unroll_dir_2
integer :: a(10)
!dir$ unroll 2
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_ENABLE_COUNT_2:.*]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[UNROLL_ENABLE_COUNT_2:.*]]
do i=1,10
a(i)=i
end do

View File

@@ -4,7 +4,9 @@
subroutine unroll_and_jam_dir
integer :: a(10)
!dir$ unroll_and_jam 4
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION:.*]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[ANNOTATION:.*]]
do i=1,10
a(i)=i
end do
@@ -14,7 +16,9 @@ end subroutine unroll_and_jam_dir
subroutine unroll_and_jam_dir_0
integer :: a(10)
!dir$ unroll_and_jam 0
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE:.*]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE:.*]]
do i=1,10
a(i)=i
end do
@@ -24,7 +28,9 @@ end subroutine unroll_and_jam_dir_0
subroutine unroll_and_jam_dir_1
integer :: a(10)
!dir$ unroll_and_jam 1
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE]]
do i=1,10
a(i)=i
end do
@@ -34,7 +40,9 @@ end subroutine unroll_and_jam_dir_1
subroutine nounroll_and_jam_dir
integer :: a(10)
!dir$ nounroll_and_jam
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE]]
do i=1,10
a(i)=i
end do
@@ -44,7 +52,9 @@ end subroutine nounroll_and_jam_dir
subroutine unroll_and_jam_dir_no_factor
integer :: a(10)
!dir$ unroll_and_jam
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_NO_FACTOR:.*]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[ANNOTATION_NO_FACTOR:.*]]
do i=1,10
a(i)=i
end do

View File

@@ -4,7 +4,9 @@
subroutine vector_always
integer :: a(10)
!dir$ vector always
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION:.*]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[ANNOTATION:.*]]
do i=1,10
a(i)=i
end do
@@ -14,7 +16,9 @@ end subroutine vector_always
subroutine no_vector
integer :: a(10)
!dir$ novector
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION2:.*]]
! CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}
! CHECK-NOT: !llvm.loop
! CHECK: br label {{.*}}, !llvm.loop ![[ANNOTATION2:.*]]
do i=1,10
a(i)=i
end do