[Mlir-commits] [mlir] [mlir] [scf-to-cf] attach the loop annotation to latch block (PR #147462)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jul 8 19:17:30 PDT 2025
https://github.com/Menooker updated https://github.com/llvm/llvm-project/pull/147462
>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/3] [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/3] 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();
>From 0fdd445e3f9f30b7e3bd96c681dc7f1575d9b8c0 Mon Sep 17 00:00:00 2001
From: Menooker <myjisgreat at live.cn>
Date: Tue, 8 Jul 2025 19:17:13 -0700
Subject: [PATCH 3/3] fix test
---
.../SCFToControlFlow/convert-to-cfg.mlir | 15 +++++++++++++--
1 file changed, 13 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir b/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
index 209c1db4ceeb6..19adefa5b4dc9 100644
--- a/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
+++ b/mlir/test/Conversion/SCFToControlFlow/convert-to-cfg.mlir
@@ -682,8 +682,19 @@ func.func @forall(%num_threads: index) {
// 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]]}
+// CHECK-LABEL: 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) {
More information about the Mlir-commits
mailing list