[Lldb-commits] [llvm] [mlir] [lldb] [mlir] Introduce replaceWithZeroTripCheck in LoopLikeOpInterface (PR #80331)
Jerry Wu via lldb-commits
lldb-commits at lists.llvm.org
Fri Feb 2 11:12:27 PST 2024
https://github.com/pzread updated https://github.com/llvm/llvm-project/pull/80331
>From 70f54b51bef87bde5e3f5ee067c0f2414d34e915 Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw at google.com>
Date: Thu, 1 Feb 2024 19:57:26 +0000
Subject: [PATCH 1/3] Add replaceWithZeroTripCheck to LoopLikeOpInterface
---
.../mlir/Interfaces/LoopLikeInterface.td | 22 +++++++++++++++++++
1 file changed, 22 insertions(+)
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index e2ac85a3f7725..77409cb3a8274 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -220,6 +220,28 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
/*defaultImplementation=*/[{
return ::mlir::failure();
}]
+ >,
+ InterfaceMethod<[{
+ Add a zero-trip-check around the loop to check if the loop body is ever
+ run and return the new loop inside the check. The loop body is moved
+ over to the new loop. Returns "failure" if the loop doesn't support
+ this transformation.
+
+ After the transformation, the ops inserted to the parent region of the
+ loop are guaranteed to be run only if the loop body runs at least one
+ iteration.
+
+ Note: Ops in the loop body might be rearranged because of loop rotating
+ to maintain the semantic. Terminators might be removed/added during this
+ transformation.
+ }],
+ /*retTy=*/"::mlir::FailureOr<::mlir::LoopLikeOpInterface>",
+ /*methodName=*/"replaceWithZeroTripCheck",
+ /*args=*/(ins "::mlir::RewriterBase &":$rewriter),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::failure();
+ }]
>
];
>From d6703ebbeb5ddc358929672b44994a9d05683523 Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw at google.com>
Date: Fri, 2 Feb 2024 18:59:03 +0000
Subject: [PATCH 2/3] Add tests
---
mlir/unittests/Interfaces/CMakeLists.txt | 3 +
.../Interfaces/LoopLikeInterfaceTest.cpp | 101 ++++++++++++++++++
2 files changed, 104 insertions(+)
create mode 100644 mlir/unittests/Interfaces/LoopLikeInterfaceTest.cpp
diff --git a/mlir/unittests/Interfaces/CMakeLists.txt b/mlir/unittests/Interfaces/CMakeLists.txt
index d192b2922d6b9..cab9503cf295b 100644
--- a/mlir/unittests/Interfaces/CMakeLists.txt
+++ b/mlir/unittests/Interfaces/CMakeLists.txt
@@ -3,6 +3,7 @@ add_mlir_unittest(MLIRInterfacesTests
DataLayoutInterfacesTest.cpp
InferIntRangeInterfaceTest.cpp
InferTypeOpInterfaceTest.cpp
+ LoopLikeInterfaceTest.cpp
)
target_link_libraries(MLIRInterfacesTests
@@ -12,7 +13,9 @@ target_link_libraries(MLIRInterfacesTests
MLIRDataLayoutInterfaces
MLIRDLTIDialect
MLIRFuncDialect
+ MLIRIR
MLIRInferIntRangeInterface
MLIRInferTypeOpInterface
+ MLIRLoopLikeInterface
MLIRParser
)
diff --git a/mlir/unittests/Interfaces/LoopLikeInterfaceTest.cpp b/mlir/unittests/Interfaces/LoopLikeInterfaceTest.cpp
new file mode 100644
index 0000000000000..b0b7680fed68e
--- /dev/null
+++ b/mlir/unittests/Interfaces/LoopLikeInterfaceTest.cpp
@@ -0,0 +1,101 @@
+//===- LoopLikeInterfaceTest.cpp - Unit tests for Loop Like Interfaces. ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Parser/Parser.h"
+
+#include <gtest/gtest.h>
+
+using namespace mlir;
+
+struct NoZeroTripCheckLoopOp
+ : public Op<NoZeroTripCheckLoopOp, LoopLikeOpInterface::Trait> {
+ using Op::Op;
+
+ static ArrayRef<StringRef> getAttributeNames() { return {}; }
+
+ static StringRef getOperationName() {
+ return "looptest.no_zero_trip_check_loop_op";
+ }
+
+ SmallVector<Region *> getLoopRegions() { return {}; }
+};
+
+struct ImplZeroTripCheckLoopOp
+ : public Op<ImplZeroTripCheckLoopOp, LoopLikeOpInterface::Trait> {
+ using Op::Op;
+
+ static ArrayRef<StringRef> getAttributeNames() { return {}; }
+
+ static StringRef getOperationName() {
+ return "looptest.impl_zero_trip_check_loop_op";
+ }
+
+ SmallVector<Region *> getLoopRegions() { return {}; }
+
+ FailureOr<LoopLikeOpInterface>
+ replaceWithZeroTripCheck(RewriterBase &rewriter) {
+ return cast<LoopLikeOpInterface>(this->getOperation());
+ }
+};
+
+/// A dialect putting all the above together.
+struct LoopTestDialect : Dialect {
+ explicit LoopTestDialect(MLIRContext *ctx)
+ : Dialect(getDialectNamespace(), ctx, TypeID::get<LoopTestDialect>()) {
+ addOperations<NoZeroTripCheckLoopOp, ImplZeroTripCheckLoopOp>();
+ }
+ static StringRef getDialectNamespace() { return "looptest"; }
+};
+
+TEST(LoopLikeOpInterface, NoReplaceWithZeroTripCheck) {
+ const char *ir = R"MLIR(
+ "looptest.no_zero_trip_check_loop_op"() : () -> ()
+ )MLIR";
+
+ DialectRegistry registry;
+ registry.insert<LoopTestDialect>();
+ MLIRContext ctx(registry);
+
+ OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
+ LoopLikeOpInterface testOp =
+ cast<LoopLikeOpInterface>(module->getBody()->getOperations().front());
+
+ IRRewriter rewriter(&ctx);
+ FailureOr<LoopLikeOpInterface> result =
+ testOp.replaceWithZeroTripCheck(rewriter);
+
+ EXPECT_TRUE(failed(result));
+}
+
+TEST(LoopLikeOpInterface, ImplReplaceWithZeroTripCheck) {
+ const char *ir = R"MLIR(
+ "looptest.impl_zero_trip_check_loop_op"() : () -> ()
+ )MLIR";
+
+ DialectRegistry registry;
+ registry.insert<LoopTestDialect>();
+ MLIRContext ctx(registry);
+
+ OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
+ LoopLikeOpInterface testOp =
+ cast<LoopLikeOpInterface>(module->getBody()->getOperations().front());
+
+ IRRewriter rewriter(&ctx);
+ FailureOr<LoopLikeOpInterface> result =
+ testOp.replaceWithZeroTripCheck(rewriter);
+
+ EXPECT_TRUE(succeeded(result));
+ EXPECT_EQ(*result, testOp);
+}
>From 43af5b4ff01b03cc8b69ba494d0f344ab645e3d1 Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw at google.com>
Date: Fri, 2 Feb 2024 19:12:14 +0000
Subject: [PATCH 3/3] Update comments
---
mlir/include/mlir/Interfaces/LoopLikeInterface.td | 5 ++++-
1 file changed, 4 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index 77409cb3a8274..81f202cf34186 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -233,7 +233,10 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
Note: Ops in the loop body might be rearranged because of loop rotating
to maintain the semantic. Terminators might be removed/added during this
- transformation.
+ transformation. Also callers are not required to check the side-effect
+ of loop condition, so the transformation needs to consider that to make
+ sure the loop behavior is unchanged when moving the condtion out of the
+ loop for the zero-trip-check.
}],
/*retTy=*/"::mlir::FailureOr<::mlir::LoopLikeOpInterface>",
/*methodName=*/"replaceWithZeroTripCheck",
More information about the lldb-commits
mailing list