[mlir] Introduce wrapInZeroTripCheck in LoopLikeOpInterface

mlir/include/mlir/Interfaces/LoopLikeInterface.td

@@ -220,6 +220,30 @@ 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 same loop (moved) or a new loop (replaced) inside the
+        check. Returns "failure" if the loop doesn't support the 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. 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 condition out of the
+        loop for the zero-trip-check.
+      }],
+      /*retTy=*/"::mlir::FailureOr<::mlir::LoopLikeOpInterface>",
+      /*methodName=*/"wrapInZeroTripCheck",
+      /*args=*/(ins "::mlir::RewriterBase &":$rewriter),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::failure();
+      }]
     >
   ];
 

mlir/test/Dialect/SCF/wrap-in-zero-trip-check.mlir

@@ -0,0 +1,21 @@
+// RUN: mlir-opt %s -test-wrap-in-zero-trip-check -split-input-file  | FileCheck %s
+
+func.func @no_wrap_scf_while_in_zero_trip_check(%bound : i32) -> i32 {
+  %cst0 = arith.constant 0 : i32
+  %cst5 = arith.constant 5 : i32
+  %res:2 = scf.while (%iter = %cst0) : (i32) -> (i32, i32) {
+    %cond = arith.cmpi slt, %iter, %bound : i32
+    %inv = arith.addi %bound, %cst5 : i32
+    scf.condition(%cond) %iter, %inv : i32, i32
+  } do {
+  ^bb0(%arg1: i32, %arg2: i32):
+    %next = arith.addi %arg1, %arg2 : i32
+    scf.yield %next : i32
+  }
+  return %res#0 : i32
+}
+
+// TODO(pzread): Update the test once the wrapInZeroTripCheck is implemented.
+// CHECK-LABEL: func.func @no_wrap_scf_while_in_zero_trip_check
+// CHECK-NOT:     scf.if
+// CHECK:         scf.while

mlir/test/lib/Interfaces/LoopLikeInterface/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_mlir_library(MLIRLoopLikeInterfaceTestPasses
   TestBlockInLoop.cpp
+  TestWrapInZeroTripCheck.cpp
 
   EXCLUDE_FROM_LIBMLIR
 

mlir/test/lib/Interfaces/LoopLikeInterface/TestWrapInZeroTripCheck.cpp

@@ -0,0 +1,53 @@
+//===- TestWrapInZeroTripCheck.cpp -- Pass to test wrapInZeroTripCheck ----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the passes to test wrapInZeroTripCheck of loop ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+
+struct TestWrapInZeroTripCheck
+    : public PassWrapper<TestWrapInZeroTripCheck, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestWrapInZeroTripCheck)
+
+  StringRef getArgument() const final { return "test-wrap-in-zero-trip-check"; }
+  StringRef getDescription() const final {
+    return "test wrapInZeroTripCheck of loop ops";
+  }
+
+  void runOnOperation() override {
+    func::FuncOp func = getOperation();
+    MLIRContext *context = &getContext();
+    IRRewriter rewriter(context);
+    func.walk([&](LoopLikeOpInterface op) {
+      auto result = op.wrapInZeroTripCheck(rewriter);
+      if (failed(result)) {
+        // Ignore not implemented failure in tests. The expected output should
+        // catch problems (e.g. transformation doesn't happen).
+      }
+    });
+  }
+};
+
+} // namespace
+
+namespace mlir {
+namespace test {
+void registerTestWrapInZeroTripCheckPass() {
+  PassRegistration<TestWrapInZeroTripCheck>();
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

@@ -136,6 +136,7 @@ void registerTestTensorTransforms();
 void registerTestTopologicalSortAnalysisPass();
 void registerTestTransformDialectEraseSchedulePass();
 void registerTestTransformDialectInterpreterPass();
+void registerTestWrapInZeroTripCheckPass();
 void registerTestWrittenToPass();
 void registerTestVectorLowerings();
 void registerTestVectorReductionToSPIRVDotProd();
@@ -261,6 +262,7 @@ void registerTestPasses() {
   mlir::test::registerTestVectorLowerings();
   mlir::test::registerTestVectorReductionToSPIRVDotProd();
   mlir::test::registerTestNvgpuLowerings();
+  mlir::test::registerTestWrapInZeroTripCheckPass();
   mlir::test::registerTestWrittenToPass();
 #if MLIR_ENABLE_PDL_IN_PATTERNMATCH
   mlir::test::registerTestDialectConversionPasses();