[Lldb-commits] [lldb] [mlir] [llvm] [mlir] Introduce replaceWithZeroTripCheck in LoopLikeOpInterface (PR #80331)

[Jerry Wu via lldb-commits]

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/2] 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/2] 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);
+}