Enable attaching LLVM loop annotations to scf.for

[xiaoleis-nv]

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.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: None (xiaoleis-nv)

Changes

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 MR allows the cf.cond_br to carry the loop annotation of scf.for after the convert-scf-to-cf pass.

---

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

2 Files Affected:

• (modified) mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp (+13-3)
• (modified) mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir (+23-1)

diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index 9eb8a289d7d658..82bbd373ce45ad 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/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/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -369,9 +370,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());
diff --git a/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir b/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
index caf17bc91ced23..9ea0093eff7868 100644
--- a/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
+++ b/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
+}
\ No newline at end of file

[xiaoleis-nv]

Hi @joker-eph , could you please review this PR? Thank you!

cc. @sherhut , @grypp , @durga4github

[xiaoleis-nv]

Can anyone help me merge this PR? I don't have permission to merge it. Thanks!