LICM now checks if parent loop has constant bounds/steps and isn't zero trip

Author: mbagherbeikTT

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

@@ -14,8 +14,9 @@
 #ifndef MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
 #define MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
 
-#include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Dominance.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
 
 namespace mlir {
 namespace SideEffects {
@@ -423,6 +424,15 @@ bool isOpTriviallyDead(Operation *op);
 /// Note: Terminators and symbols are never considered to be trivially dead.
 bool wouldOpBeTriviallyDead(Operation *op);
 
+/// Returns TRUE if the loop is dead/zero-trip,
+/// FALSE if loop has constant bounds/steps and has at least 1 iteration
+/// on every dimension, returns nullopt otherwise
+///
+/// Can only infer if loop is dead if it has constant loop bounds/steps.
+/// Otherwise we assume that it's dead to be conservative.
+///
+std::optional<bool> isZeroTrip(mlir::LoopLikeOpInterface loop);
+
 /// Returns true if the given operation is allowed to be moved under the
 /// memory effects interface.
 ///
@@ -449,13 +459,15 @@ bool isMemoryEffectFree(Operation *op);
 /// Returns true if the given operation has conflict-free write effects
 ///
 /// An operation is conflict free:
-/// (1) all of its memory effects are of type Write
-/// (2) there are no other ops with Alloc/Free/Write effects on the same
-/// resources within the ops parent region
-/// (3) all ops in the parent region with Read effects on the same resources
-/// are dominated by the operation
+/// (1) Parent is a loop with the LoopLikeOpInterface
+/// (2) Parent loop is not a zero trip loop and has constant bounds/steps
+/// (3) all of the op's memory effects are of type Write
+/// (4) there are no other ops with Alloc/Free/Write effects on the same
+/// resources within the op's parent loop region
+/// (5) all ops in the parent loop region with Read effects on the same
+/// resources are dominated by the operation
 ///
-/// If the operation meets all 3 criteria, then it is conflict free
+/// If the operation meets all criteria, then it is conflict free
 bool isMemoryEffectConflictFree(Operation *op);
 
 /// Returns true if op and/or any operations within its nested regions

mlir/lib/Interfaces/CMakeLists.txt

@@ -85,7 +85,21 @@ add_mlir_interface_library(MemorySlotInterfaces)
 add_mlir_interface_library(ParallelCombiningOpInterface)
 add_mlir_interface_library(RuntimeVerifiableOpInterface)
 add_mlir_interface_library(ShapedOpInterfaces)
-add_mlir_interface_library(SideEffectInterfaces)
+
+add_mlir_library(MLIRSideEffectInterfaces
+  SideEffectInterfaces.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces
+
+  DEPENDS
+  MLIRSideEffectInterfacesIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRDialectUtils
+  MLIRLoopLikeInterface
+)
 
 add_mlir_library(MLIRSubsetOpInterface
   SubsetOpInterface.cpp

mlir/lib/Interfaces/SideEffectInterfaces.cpp

@@ -8,10 +8,9 @@
 
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/SymbolTable.h"
-#include "llvm/ADT/SmallPtrSet.h"
-#include <unordered_set>
 #include <utility>
 
 using namespace mlir;
@@ -23,6 +22,13 @@ using namespace mlir;
 /// Include the definitions of the side effect interfaces.
 #include "mlir/Interfaces/SideEffectInterfaces.cpp.inc"
 
+// //===----------------------------------------------------------------------===//
+// // LoopLike Interfaces
+// //===----------------------------------------------------------------------===//
+
+// /// Include the definitions of the loop like interfaces.
+// #include "mlir/Interfaces/LoopLikeInterface.cpp.inc"
+
 //===----------------------------------------------------------------------===//
 // MemoryEffects
 //===----------------------------------------------------------------------===//
@@ -315,6 +321,33 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
   return wouldOpBeTriviallyDeadImpl(op);
 }
 
+std::optional<bool> mlir::isZeroTrip(mlir::LoopLikeOpInterface loop) {
+  auto lbs = loop.getLoopLowerBounds();
+  auto ubs = loop.getLoopUpperBounds();
+  auto steps = loop.getLoopSteps();
+
+  if (!lbs || !ubs || !steps)
+    return std::nullopt;
+
+  if (lbs->size() != ubs->size() || ubs->size() != steps->size())
+    return std::nullopt;
+
+  for (size_t i = 0; i < steps->size(); ++i) {
+    auto lb = getConstantIntValue((*lbs)[i]);
+    auto ub = getConstantIntValue((*ubs)[i]);
+    auto st = getConstantIntValue((*steps)[i]);
+
+    if (!lb || !ub || !st)
+      return std::nullopt; // non-constant -> unknown
+
+    if (*st >= 0 && *lb >= *ub)
+      return true;
+    if (*st < 0 && *lb <= *ub)
+      return true;
+  }
+  return false;
+}
+
 bool mlir::isMemoryEffectMovable(Operation *op) {
   if (isMemoryEffectFree(op))
     return true;
@@ -358,6 +391,19 @@ bool mlir::isMemoryEffectConflictFree(Operation *op) {
   // shouldn't be flagged as movable to be conservative
   if (!memInterface) return false;
 
+  // check parent loop to make sure it's not dead
+  Operation *parent = op->getParentOp();
+  if (!parent)
+    return false;
+
+  auto loopInterface = dyn_cast<LoopLikeOpInterface>(parent);
+  if (!loopInterface)
+    return false;
+
+  auto isDead = isZeroTrip(loopInterface);
+  if (!isDead.has_value() || isDead.value())
+    return false;
+
   // gather all effects on op
   llvm::SmallVector<MemoryEffects::EffectInstance> effects;
   memInterface.getEffects(effects);
@@ -376,8 +422,6 @@ bool mlir::isMemoryEffectConflictFree(Operation *op) {
     resourceCounts.try_emplace(effect.getResource()->getResourceID(), 0);
   }
 
-  // op itself is good, need to check rest of its parent region
-  Operation *parent = op->getParentOp();
   mlir::DominanceInfo dom(parent);
 
   for (Region &region : parent->getRegions())

mlir/test/Transforms/loop-invariant-code-motion.mlir

@@ -1480,8 +1480,8 @@ func.func @move_single_resource_write_dominant() attributes {} {
 
 // -----
 
-// CHECK-LABEL func.func @move_single_resource_read_dominant_negative
-func.func @move_single_resource_read_dominant_negative() attributes {} {
+// CHECK-LABEL func.func @move_single_resource_read_dominant
+func.func @move_single_resource_read_dominant() attributes {} {
   %c0_i32 = arith.constant 0 : i32
   %c1_i32 = arith.constant 10 : i32
   %c2_i32 = arith.constant 1 : i32
@@ -1530,8 +1530,8 @@ func.func @move_single_resource_basic_conflict() attributes {} {
 
 // -----
 
-// CHECK-LABEL func.func @move_single_resource_if_region_negative
-func.func @move_single_resource_if_region_negative() attributes {} {
+// CHECK-LABEL func.func @move_single_resource_if_region
+func.func @move_single_resource_if_region() attributes {} {
   %c0_i32 = arith.constant 0 : i32
   %c1_i32 = arith.constant 10 : i32
   %c2_i32 = arith.constant 1 : i32
@@ -1641,3 +1641,36 @@ func.func @move_multi_resource_comprehensive() attributes {} {
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL func.func @move_single_resource_dead_loops
+func.func @move_single_resource_dead_loops() attributes {} {
+  %c0_i32 = arith.constant 0 : i32
+  %c1_i32 = arith.constant 10 : i32
+  %c2_i32 = arith.constant 1 : i32
+  %c3_i32 = arith.constant -1 : i32
+  
+  scf.for %arg0 = %c0_i32 to %c1_i32 step %c3_i32  : i32 {
+    // CHECK: "test.test_effects_write_C"() : () -> ()
+    "test.test_effects_write_C"() : () -> ()
+
+    scf.for %arg1 = %c1_i32 to %c0_i32 step %c2_i32  : i32 {
+      // CHECK: "test.test_effects_write_B"() : () -> ()
+      "test.test_effects_write_B"() : () -> ()
+
+      scf.for %arg2 = %c0_i32 to %c0_i32 step %c0_i32  : i32 {
+        // CHECK: "test.test_effects_write_A"() : () -> ()
+        "test.test_effects_write_A"() : () -> ()
+
+        // CHECK: "test.test_effects_write_EF"() : () -> ()
+        scf.for %arg3 = %c0_i32 to %c1_i32 step %c2_i32  : i32 {
+          "test.test_effects_write_EF"() : () -> ()
+          // CHECK: "test.test_effects_read_EF"() : () -> ()
+          "test.test_effects_read_EF"() : () -> ()
+        }
+      }
+    }
+  }
+  return
+}