[flang] Add loop annotation attributes to the loop backedge instead of the loop header's conditional branch

[ashermancinelli]

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

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Asher Mancinelli (ashermancinelli)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/126082.diff

4 Files Affected:

• (modified) flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp (+6-6)
• (modified) flang/test/Fir/vector-always.fir (+3-1)
• (modified) flang/test/Integration/unroll.f90 (+3-1)
• (modified) flang/test/Integration/vector-always.f90 (+3-1)

diff --git a/flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp b/flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp index b09bbf6106dbbb1..0e03d574f40709e 100644 --- a/flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp +++ b/flang/lib/Optimizer/Transforms/ControlFlowConverter.cpp @@ -123,23 +123,23 @@ class CfgLoopConv : public mlir::OpRewritePattern<fir::DoLoopOp> { : 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() diff --git a/flang/test/Fir/vector-always.fir b/flang/test/Fir/vector-always.fir index 00eb0e7a756ee6f..ec06b94a3d0f8d2 100644 --- a/flang/test/Fir/vector-always.fir +++ b/flang/test/Fir/vector-always.fir @@ -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 } diff --git a/flang/test/Integration/unroll.f90 b/flang/test/Integration/unroll.f90 index 9d69605e10d1b3c..fbb41f72d324818 100644 --- a/flang/test/Integration/unroll.f90 +++ b/flang/test/Integration/unroll.f90 @@ -4,7 +4,9 @@ subroutine unroll_dir integer :: a(10) !dir$ unroll - ! 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 diff --git a/flang/test/Integration/vector-always.f90 b/flang/test/Integration/vector-always.f90 index 7216698f901c1fe..8c71f439027834f 100644 --- a/flang/test/Integration/vector-always.f90 +++ b/flang/test/Integration/vector-always.f90 @@ -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 

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[vzakhari]

Thank you, Asher!

This looks good to me and in line with LLVM LangRef. Please wait for comments from @DavidTruby

[jeanPerier]

Thanks! Looks in line with the doc you are linking, but I have little experience here.

[DavidTruby]

LGTM thanks for the fix!

[DavidTruby]

Should this be merged, since it has a number of approvals and seemingly no issues?

[ashermancinelli]

Sorry folks, I lost track of this. I believe it needs to be updated because of other changes to loop annotations. Let me rebase and run some more testing. Thanks for the ping!