[Mlir-commits] [mlir] [mlir] [scf-to-cf] attach the loop annotation to latch block (PR #147462)

[llvmlistbot at llvm.org]

https://github.com/Menooker created https://github.com/llvm/llvm-project/pull/147462

As [required by LLVM](https://llvm.org/docs/LangRef.html#llvm-loop), the loop annotation (loop metadata) should be attached on the ["latch" block](https://llvm.org/docs/LoopTerminology.html). Otherwise, the annotation might be ignored by LLVM. This PR fixes this issue.

>From 8e5018f1ac7dbe7e6ccbf394de009cff041a4b82 Mon Sep 17 00:00:00 2001
From: Menooker <myjisgreat at live.cn>
Date: Mon, 7 Jul 2025 22:30:42 -0700
Subject: [PATCH 1/2] [mlir] [scf-to-cf] attach the loop annotation to latch
 block

---
 .../SCFToControlFlow/SCFToControlFlow.cpp     | 30 +++++++++++--------
 .../SCFToControlFlow/convert-to-cfg.mlir      | 12 ++++----
 2 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index ac6dcf2513e99..cd42465cd9e4d 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -347,7 +347,20 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
   SmallVector<Value, 8> loopCarried;
   loopCarried.push_back(stepped);
   loopCarried.append(terminator->operand_begin(), terminator->operand_end());
-  rewriter.create<cf::BranchOp>(loc, conditionBlock, loopCarried);
+  auto branchOp =
+      rewriter.create<cf::BranchOp>(loc, conditionBlock, loopCarried);
+
+  // Let the CondBranchOp carry the LLVM attributes from the ForOp, such as the
+  // llvm.loop_annotation attribute.
+  // LLVM requires the loop metadata to be attached on the "latch" block. Which
+  // is the back-edge to the header block (conditionBlock)
+  SmallVector<NamedAttribute> llvmAttrs;
+  llvm::copy_if(forOp->getAttrs(), std::back_inserter(llvmAttrs),
+                [](auto attr) {
+                  return isa<LLVM::LLVMDialect>(attr.getValue().getDialect());
+                });
+  branchOp->setDiscardableAttrs(llvmAttrs);
+
   rewriter.eraseOp(terminator);
 
   // Compute loop bounds before branching to the condition.
@@ -369,24 +382,15 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
   auto comparison = rewriter.create<arith::CmpIOp>(
       loc, arith::CmpIPredicate::slt, iv, upperBound);
 
-  auto condBranchOp = rewriter.create<cf::CondBranchOp>(
-      loc, comparison, firstBodyBlock, ArrayRef<Value>(), endBlock,
-      ArrayRef<Value>());
+  rewriter.create<cf::CondBranchOp>(loc, comparison, firstBodyBlock,
+                                    ArrayRef<Value>(), endBlock,
+                                    ArrayRef<Value>());
 
-  // Let the CondBranchOp carry the LLVM attributes from the ForOp, such as the
-  // llvm.loop_annotation attribute.
-  SmallVector<NamedAttribute> llvmAttrs;
-  llvm::copy_if(forOp->getAttrs(), std::back_inserter(llvmAttrs),
-                [](auto attr) {
-                  return isa<LLVM::LLVMDialect>(attr.getValue().getDialect());
-                });
-  condBranchOp->setDiscardableAttrs(llvmAttrs);
   // The result of the loop operation is the values of the condition block
   // arguments except the induction variable on the last iteration.
   rewriter.replaceOp(forOp, conditionBlock->getArguments().drop_front());
   return success();
 }
-
 LogicalResult IfLowering::matchAndRewrite(IfOp ifOp,
                                           PatternRewriter &rewriter) const {
   auto loc = ifOp.getLoc();
diff --git a/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir b/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
index 9ea0093eff786..209c1db4ceeb6 100644
--- a/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
+++ b/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
@@ -678,12 +678,12 @@ func.func @forall(%num_threads: index) {
 
 // -----
 
-//      CHECK: #loop_unroll = #llvm.loop_unroll<disable = true>
-// CHECK-NEXT: #loop_unroll1 = #llvm.loop_unroll<full = true>
-// CHECK-NEXT: #[[NO_UNROLL:.*]] = #llvm.loop_annotation<unroll = #loop_unroll>
-// CHECK-NEXT: #[[FULL_UNROLL:.*]] = #llvm.loop_annotation<unroll = #loop_unroll1>
-//      CHECK: cf.cond_br %{{.*}}, ^bb2, ^bb6 {llvm.loop_annotation = #[[NO_UNROLL]]}
-//      CHECK: cf.cond_br %{{.*}}, ^bb4, ^bb5 {llvm.loop_annotation = #[[FULL_UNROLL]]}
+//      CHECK: #[[LOOP_UNROLL:.*]] = #llvm.loop_unroll<full = true>
+// CHECK-NEXT: #[[LOOP_UNROLL_DISABLE:.*]] = #llvm.loop_unroll<disable = true>
+// CHECK-NEXT: #[[FULL_UNROLL:.*]] = #llvm.loop_annotation<unroll = #[[LOOP_UNROLL]]>
+// CHECK-NEXT: #[[NO_UNROLL:.*]] = #llvm.loop_annotation<unroll = #[[LOOP_UNROLL_DISABLE]]>
+//      CHECK: cf.br ^bb3({{.*}}) {llvm.loop_annotation = #[[FULL_UNROLL]]}
+//      CHECK: cf.br ^bb1({{.*}}) {llvm.loop_annotation = #[[NO_UNROLL]]}
 #no_unroll = #llvm.loop_annotation<unroll = <disable = true>>
 #full_unroll = #llvm.loop_annotation<unroll = <full = true>>
 func.func @simple_std_for_loops_annotation(%arg0 : index, %arg1 : index, %arg2 : index) {

>From b1ece0231ee5f2b7f74f188b313c28f3f0904972 Mon Sep 17 00:00:00 2001
From: Menooker <myjisgreat at live.cn>
Date: Mon, 7 Jul 2025 22:34:47 -0700
Subject: [PATCH 2/2] format

---
 mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index cd42465cd9e4d..03c5c2239ed7a 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -391,6 +391,7 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
   rewriter.replaceOp(forOp, conditionBlock->getArguments().drop_front());
   return success();
 }
+
 LogicalResult IfLowering::matchAndRewrite(IfOp ifOp,
                                           PatternRewriter &rewriter) const {
   auto loc = ifOp.getLoc();