Implement simple loop-invariant-code-motion based on dialect interfaces.

PiperOrigin-RevId: 275004258

mlir/include/mlir/Dialect/AffineOps/AffineOps.h

@@ -28,6 +28,7 @@
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
+#include "mlir/Transforms/LoopLikeInterface.h"
 
 namespace mlir {
 class AffineBound;

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

@@ -28,6 +28,11 @@
 include "mlir/IR/OpBase.td"
 #endif // OP_BASE
 
+#ifdef MLIR_LOOPLIKEINTERFACE
+#else
+include "mlir/Transforms/LoopLikeInterface.td"
+#endif
+
 include "mlir/Dialect/AffineOps/AffineOpsBase.td"
 
 def Affine_Dialect : Dialect {
@@ -53,7 +58,9 @@ class Affine_Op<string mnemonic, list<OpTrait> traits = []> :
 def ImplicitAffineTerminator
     : SingleBlockImplicitTerminator<"AffineTerminatorOp">;
 
-def AffineForOp : Affine_Op<"for", [ImplicitAffineTerminator]> {
+def AffineForOp : Affine_Op<"for",
+    [ImplicitAffineTerminator,
+     DeclareOpInterfaceMethods<LoopLikeOpInterface>]> {
   let summary = "for operation";
   let description = [{
     The "affine.for" operation represents an affine loop nest, defining an SSA

mlir/include/mlir/Dialect/LoopOps/LoopOps.h

@@ -26,6 +26,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Transforms/LoopLikeInterface.h"
 
 namespace mlir {
 namespace loop {

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

@@ -28,6 +28,11 @@
 include "mlir/IR/OpBase.td"
 #endif // OP_BASE
 
+#ifdef MLIR_LOOPLIKEINTERFACE
+#else
+include "mlir/Transforms/LoopLikeInterface.td"
+#endif
+
 def Loop_Dialect : Dialect {
   let name = "loop";
   let cppNamespace = "";
@@ -48,7 +53,8 @@ class Loop_Op<string mnemonic, list<OpTrait> traits = []> :
 }
 
 def ForOp : Loop_Op<"for",
-      [SingleBlockImplicitTerminator<"TerminatorOp">]> {
+      [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+       SingleBlockImplicitTerminator<"TerminatorOp">]> {
   let summary = "for operation";
   let description = [{
     The "loop.for" operation represents a loop nest taking 3 SSA value as

mlir/include/mlir/IR/Block.h

@@ -249,6 +249,15 @@ class Block : public IRObjectWithUseList,
     return op_filter_iterator<OpT>(end(), end());
   }
 
+  /// Return an iterator range over the operation within this block excluding
+  /// the terminator operation at the end.
+  llvm::iterator_range<iterator> without_terminator() {
+    if (begin() == end())
+      return {begin(), end()};
+    auto endIt = --end();
+    return {begin(), endIt};
+  }
+
   //===--------------------------------------------------------------------===//
   // Terminator management
   //===--------------------------------------------------------------------===//

mlir/include/mlir/Transforms/LoopLikeInterface.h

@@ -0,0 +1,35 @@
+//===- LoopLikeInterface.h - Loop-like operations interface ---------------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file implements the operation interface for loop like operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_LOOPLIKEINTERFACE_H_
+#define MLIR_TRANSFORMS_LOOPLIKEINTERFACE_H_
+
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/ArrayRef.h"
+
+namespace mlir {
+
+#include "mlir/Transforms/LoopLikeInterface.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_LOOPLIKEINTERFACE_H_

mlir/include/mlir/Transforms/LoopLikeInterface.td

@@ -0,0 +1,62 @@
+//===- LoopLikeInterface.td - LoopLike interface -----------*- tablegen -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// Defines the interface for loop-like operations as used by LICM.
+//
+//===----------------------------------------------------------------------===//
+
+#ifdef MLIR_LOOPLIKEINTERFACE
+#else
+#define MLIR_LOOPLIKEINTERFACE
+
+#ifdef OP_BASE
+#else
+include "mlir/IR/OpBase.td"
+#endif // OP_BASE
+
+def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
+  let description = [{
+    Encodes properties of a loop. Operations that implement this interface will
+    be considered by loop-invariant code motion.
+  }];
+
+  let methods = [
+    InterfaceMethod<[{
+        Returns true if the given value is defined outside of the loop.
+        A sensible implementation could be to check whether the value's defining
+        operation lies outside of the loops body region. If the loop uses
+        explicit capture of dependencies, an implementation could check whether
+        the value corresponds to a captured dependency.
+      }],
+      "bool", "isDefinedOutsideOfLoop", (ins "Value *":$value)
+    >,
+    InterfaceMethod<[{
+        Returns the region that makes up the body of the loop and should be
+        inspected for loop-invariant operations.
+      }],
+      "Region &", "getLoopBody"
+    >,
+    InterfaceMethod<[{
+        Moves the given vector of operations out of the loop. The vector is
+        sorted topologically.
+      }],
+      "LogicalResult", "moveOutOfLoop", (ins "ArrayRef<Operation *>":$ops)
+    >,
+  ];
+}
+
+#endif // MLIR_LOOPLIKEINTERFACE

mlir/include/mlir/Transforms/Passes.h

@@ -32,6 +32,7 @@ namespace mlir {
 class AffineForOp;
 class FuncOp;
 class ModuleOp;
+class Pass;
 template <typename T> class OpPassBase;
 
 /// Creates a constant folding pass. Note that this pass solely provides simple
@@ -90,7 +91,11 @@ createLoopFusionPass(unsigned fastMemorySpace = 0,
 
 /// Creates a loop invariant code motion pass that hoists loop invariant
 /// instructions out of the loop.
-std::unique_ptr<OpPassBase<FuncOp>> createLoopInvariantCodeMotionPass();
+std::unique_ptr<Pass> createLoopInvariantCodeMotionPass();
+
+/// Creates a loop invariant code motion pass that hoists loop invariant
+/// instructions out of affine loop.
+std::unique_ptr<OpPassBase<FuncOp>> createAffineLoopInvariantCodeMotionPass();
 
 /// Creates a pass to pipeline explicit movement of data across levels of the
 /// memory hierarchy.

mlir/include/mlir/Transforms/SideEffectsInterface.h

@@ -0,0 +1,73 @@
+//===- SideEffectsInterface.h - dialect interface modeling side effects ---===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file specifies a dialect interface to model side-effects.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_SIDEEFFECTSINTERFACE_H_
+#define MLIR_TRANSFORMS_SIDEEFFECTSINTERFACE_H_
+
+#include "mlir/IR/DialectInterface.h"
+#include "mlir/IR/Operation.h"
+
+namespace mlir {
+
+/// Specifies an interface for basic side-effect modelling that is used by the
+/// loop-invariant code motion pass.
+///
+/// TODO: This interface should be replaced by a more general solution.
+class SideEffectsDialectInterface
+    : public DialectInterface::Base<SideEffectsDialectInterface> {
+public:
+  SideEffectsDialectInterface(Dialect *dialect) : Base(dialect) {}
+
+  enum SideEffecting {
+    Never,     /* the operation has no side-effects */
+    Recursive, /* the operation has side-effects if a contained operation has */
+    Always     /* the operation has side-effects */
+  };
+
+  /// Checks whether the given operation has side-effects.
+  virtual SideEffecting isSideEffecting(Operation *op) const {
+    if (op->hasNoSideEffect())
+      return Never;
+    return Always;
+  };
+};
+
+class SideEffectsInterface
+    : public DialectInterfaceCollection<SideEffectsDialectInterface> {
+public:
+  using SideEffecting = SideEffectsDialectInterface::SideEffecting;
+  explicit SideEffectsInterface(MLIRContext *ctx)
+      : DialectInterfaceCollection<SideEffectsDialectInterface>(ctx) {}
+
+  SideEffecting isSideEffecting(Operation *op) const {
+    // First check generic trait.
+    if (op->hasNoSideEffect())
+      return SideEffecting::Never;
+    if (auto handler = getInterfaceFor(op))
+      return handler->isSideEffecting(op);
+
+    return SideEffecting::Always;
+  }
+};
+
+} // namespace mlir
+
+#endif // MLIR_TRANSFORMS_SIDEEFFECTSINTERFACE_H_

mlir/lib/Dialect/AffineOps/AffineOps.cpp

@@ -23,9 +23,11 @@
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/InliningUtils.h"
+#include "mlir/Transforms/SideEffectsInterface.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/Debug.h"
+
 using namespace mlir;
 using llvm::dbgs;
 
@@ -68,6 +70,19 @@ struct AffineInlinerInterface : public DialectInlinerInterface {
   /// Affine regions should be analyzed recursively.
   bool shouldAnalyzeRecursively(Operation *op) const final { return true; }
 };
+
+// TODO(mlir): Extend for other ops in this dialect.
+struct AffineSideEffectsInterface : public SideEffectsDialectInterface {
+  using SideEffectsDialectInterface::SideEffectsDialectInterface;
+
+  SideEffecting isSideEffecting(Operation *op) const override {
+    if (isa<AffineIfOp>(op)) {
+      return Recursive;
+    }
+    return SideEffectsDialectInterface::isSideEffecting(op);
+  };
+};
+
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
@@ -81,7 +96,7 @@ AffineOpsDialect::AffineOpsDialect(MLIRContext *context)
 #define GET_OP_LIST
 #include "mlir/Dialect/AffineOps/AffineOps.cpp.inc"
                 >();
-  addInterfaces<AffineInlinerInterface>();
+  addInterfaces<AffineInlinerInterface, AffineSideEffectsInterface>();
 }
 
 /// A utility function to check if a given region is attached to a function.
@@ -1530,6 +1545,18 @@ bool AffineForOp::matchingBoundOperandList() {
   return true;
 }
 
+Region &AffineForOp::getLoopBody() { return region(); }
+
+bool AffineForOp::isDefinedOutsideOfLoop(Value *value) {
+  return !region().isAncestor(value->getParentRegion());
+}
+
+LogicalResult AffineForOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
+  for (auto *op : ops)
+    op->moveBefore(*this);
+  return success();
+}
+
 /// Returns if the provided value is the induction variable of a AffineForOp.
 bool mlir::isForInductionVar(Value *val) {
   return getForInductionVarOwner(val) != AffineForOp();

mlir/lib/Dialect/LoopOps/LoopOps.cpp

@@ -29,10 +29,29 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Support/STLExtras.h"
+#include "mlir/Transforms/SideEffectsInterface.h"
 
 using namespace mlir;
 using namespace mlir::loop;
 
+//===----------------------------------------------------------------------===//
+// LoopOpsDialect Interfaces
+//===----------------------------------------------------------------------===//
+namespace {
+
+struct LoopSideEffectsInterface : public SideEffectsDialectInterface {
+  using SideEffectsDialectInterface::SideEffectsDialectInterface;
+
+  SideEffecting isSideEffecting(Operation *op) const override {
+    if (isa<IfOp>(op) || isa<ForOp>(op)) {
+      return Recursive;
+    }
+    return SideEffectsDialectInterface::isSideEffecting(op);
+  };
+};
+
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // LoopOpsDialect
 //===----------------------------------------------------------------------===//
@@ -43,6 +62,7 @@ LoopOpsDialect::LoopOpsDialect(MLIRContext *context)
 #define GET_OP_LIST
 #include "mlir/Dialect/LoopOps/LoopOps.cpp.inc"
       >();
+  addInterfaces<LoopSideEffectsInterface>();
 }
 
 //===----------------------------------------------------------------------===//
@@ -112,6 +132,18 @@ static ParseResult parseForOp(OpAsmParser &parser, OperationState &result) {
   return success();
 }
 
+Region &ForOp::getLoopBody() { return region(); }
+
+bool ForOp::isDefinedOutsideOfLoop(Value *value) {
+  return !region().isAncestor(value->getParentRegion());
+}
+
+LogicalResult ForOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
+  for (auto *op : ops)
+    op->moveBefore(this->getOperation());
+  return success();
+}
+
 ForOp mlir::loop::getForInductionVarOwner(Value *val) {
   auto *ivArg = dyn_cast<BlockArgument>(val);
   if (!ivArg)