Enable attaching LLVM loop annotations to scf.for (#102562)

We recently discovered that the loop with a dynamic upper bound is
unexpectedly unrolled during the NVVM to PTX process. By attaching the
`llvm.loop_annotation`, we can control the unrolling behavior precisely.

This PR enables the `cf.cond_br` to retain the loop annotation of
`scf.for` after the `convert-scf-to-cf` pass. This change allows users
to have precise control over the loop behavior during backend
transformation.

---------

Co-authored-by: Xiaolei Shi <[email protected]>

Author: xiaoleis-nv

mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp

@@ -15,6 +15,7 @@
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/IR/Builders.h"
@@ -370,9 +371,18 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
   auto comparison = rewriter.create<arith::CmpIOp>(
       loc, arith::CmpIPredicate::slt, iv, upperBound);
 
-  rewriter.create<cf::CondBranchOp>(loc, comparison, firstBodyBlock,
-                                    ArrayRef<Value>(), endBlock,
-                                    ArrayRef<Value>());
+  auto condBranchOp = 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

@@ -1,4 +1,4 @@
-// RUN: mlir-opt -allow-unregistered-dialect -convert-scf-to-cf %s | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -convert-scf-to-cf -split-input-file %s | FileCheck %s
 
 // CHECK-LABEL: func @simple_std_for_loop(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
 //  CHECK-NEXT:  cf.br ^bb1(%{{.*}} : index)
@@ -675,3 +675,25 @@ func.func @forall(%num_threads: index) {
   }
   return
 }
+
+// -----
+
+//      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]]}
+#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) {
+  scf.for %i0 = %arg0 to %arg1 step %arg2 {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c4 = arith.constant 4 : index
+    scf.for %i1 = %c0 to %c4 step %c1 {
+      %c1_0 = arith.constant 1 : index
+    } {llvm.loop_annotation = #full_unroll}
+  } {llvm.loop_annotation = #no_unroll}
+  return
+}