[mlir] [scf-to-cf] attach the loop annotation to latch block (#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.

Author: Menooker

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());

mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir

@@ -678,12 +678,24 @@ 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-DAG: #[[LOOP_UNROLL_DISABLE:.*]] = #llvm.loop_unroll<disable = true>
+
+// CHECK-DAG: #[[FULL_UNROLL:.*]] = #llvm.loop_annotation<unroll = #[[LOOP_UNROLL]]>
+// CHECK-DAG: #[[NO_UNROLL:.*]] = #llvm.loop_annotation<unroll = #[[LOOP_UNROLL_DISABLE]]>
+// CHECK: func @simple_std_for_loops_annotation
+//      CHECK: ^[[bb1:.*]](%{{.*}}: index):
+//      CHECK:   cf.cond_br %{{.*}}, ^[[bb2:.*]], ^[[bb6:.*]]
+//      CHECK: ^[[bb2]]:
+//      CHECK:   cf.br ^[[bb3:.*]]({{.*}})
+//      CHECK: ^[[bb3]](%{{.*}}: index):
+//      CHECK:   cf.cond_br %{{.*}}, ^[[bb4:.*]], ^[[bb5:.*]]
+//      CHECK: ^[[bb4]]:
+//      CHECK:   cf.br ^[[bb3]]({{.*}}) {llvm.loop_annotation = #[[FULL_UNROLL]]}
+//      CHECK: ^[[bb5]]:
+//      CHECK:   cf.br ^[[bb1]]({{.*}}) {llvm.loop_annotation = #[[NO_UNROLL]]}
+//      CHECK: ^[[bb6]]:
+//      CHECK:   return
 #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) {