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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Menooker (Menooker)

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp (+17-12) 
- (modified) mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir (+6-6) 


``````````diff
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index ac6dcf2513e99..03c5c2239ed7a 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,18 +382,10 @@ 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());
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) {

``````````

</details>


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