[Mlir-commits] [llvm] [mlir] Enable attaching LLVM loop annotations to scf.for (PR #102562)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Aug 8 19:23:19 PDT 2024
https://github.com/xiaoleis-nv created https://github.com/llvm/llvm-project/pull/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 MR allows the `cf.cond_br` to carry the loop annotation of `scf.for` after the `convert-scf-to-cf` pass.
>From c64eca0083aaa89ba4c92c139de51867f60eb9db Mon Sep 17 00:00:00 2001
From: Xiaolei Shi <xiaoleis at nvidia.com>
Date: Tue, 30 Apr 2024 10:35:15 +0800
Subject: [PATCH 1/3] remove redundant move in return statement
---
llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h | 2 +-
llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h | 8 ++++----
2 files changed, 5 insertions(+), 5 deletions(-)
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h
index f0263e0d78e75e..6dd99a76cf2afb 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h
@@ -77,7 +77,7 @@ class TrampolinePool {
std::lock_guard<std::mutex> Lock(TPMutex);
if (AvailableTrampolines.empty()) {
if (auto Err = grow())
- return std::move(Err);
+ return Err;
}
assert(!AvailableTrampolines.empty() && "Failed to grow trampoline pool");
auto TrampolineAddr = AvailableTrampolines.back();
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
index 3a71ddc88ce956..2608a95e40a36b 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
@@ -488,18 +488,18 @@ class LLJITBuilderSetters {
/// Create an instance of the JIT.
Expected<std::unique_ptr<JITType>> create() {
if (auto Err = impl().prepareForConstruction())
- return std::move(Err);
+ return Err;
Error Err = Error::success();
std::unique_ptr<JITType> J(new JITType(impl(), Err));
if (Err)
- return std::move(Err);
+ return Err;
if (impl().NotifyCreated)
if (Error Err = impl().NotifyCreated(*J))
- return std::move(Err);
+ return Err;
- return std::move(J);
+ return J;
}
protected:
>From 8cc0c082365d694ceb670d4c29781de92ab3390a Mon Sep 17 00:00:00 2001
From: Xiaolei Shi <xiaoleis at nvidia.com>
Date: Fri, 9 Aug 2024 10:11:40 +0800
Subject: [PATCH 2/3] Enable to carry the annotation of
---
.../SCFToControlFlow/SCFToControlFlow.cpp | 16 ++++++++++---
.../SCFToControlFlow/convert-to-cfg.mlir | 24 ++++++++++++++++++-
2 files changed, 36 insertions(+), 4 deletions(-)
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
>From ad9f1cbd51c5de94c9f6b02a84c74e4c37535c3e Mon Sep 17 00:00:00 2001
From: Xiaolei Shi <xiaoleis at nvidia.com>
Date: Fri, 9 Aug 2024 10:22:05 +0800
Subject: [PATCH 3/3] revert unwanted changes
---
llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h | 2 +-
llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h | 8 ++++----
2 files changed, 5 insertions(+), 5 deletions(-)
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h
index 6dd99a76cf2afb..f0263e0d78e75e 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/IndirectionUtils.h
@@ -77,7 +77,7 @@ class TrampolinePool {
std::lock_guard<std::mutex> Lock(TPMutex);
if (AvailableTrampolines.empty()) {
if (auto Err = grow())
- return Err;
+ return std::move(Err);
}
assert(!AvailableTrampolines.empty() && "Failed to grow trampoline pool");
auto TrampolineAddr = AvailableTrampolines.back();
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
index 2608a95e40a36b..3a71ddc88ce956 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
@@ -488,18 +488,18 @@ class LLJITBuilderSetters {
/// Create an instance of the JIT.
Expected<std::unique_ptr<JITType>> create() {
if (auto Err = impl().prepareForConstruction())
- return Err;
+ return std::move(Err);
Error Err = Error::success();
std::unique_ptr<JITType> J(new JITType(impl(), Err));
if (Err)
- return Err;
+ return std::move(Err);
if (impl().NotifyCreated)
if (Error Err = impl().NotifyCreated(*J))
- return Err;
+ return std::move(Err);
- return J;
+ return std::move(J);
}
protected:
More information about the Mlir-commits
mailing list