[CIR][Interfaces] Implement LoopOpInterface

clang/include/clang/CIR/Dialect/IR/CIRDialect.h

@@ -33,6 +33,7 @@
 
 #include "clang/CIR/Interfaces/ASTAttrInterfaces.h"
 #include "clang/CIR/Interfaces/CIROpInterfaces.h"
+#include "clang/CIR/Interfaces/LoopOpInterface.h"
 
 namespace mlir {
 namespace OpTrait {

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -20,6 +20,7 @@ include "clang/CIR/Dialect/IR/CIRAttrs.td"
 
 include "clang/CIR/Interfaces/ASTAttrInterfaces.td"
 include "clang/CIR/Interfaces/CIROpInterfaces.td"
+include "clang/CIR/Interfaces/LoopOpInterface.td"
 
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/FunctionInterfaces.td"
@@ -1147,7 +1148,8 @@ def LoopOpKind : I32EnumAttr<
 }
 
 def LoopOp : CIR_Op<"loop",
-      [DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+      [LoopOpInterface,
+       DeclareOpInterfaceMethods<LoopLikeOpInterface>,
        DeclareOpInterfaceMethods<RegionBranchOpInterface>,
        RecursivelySpeculatable, NoRegionArguments]> {
   let summary = "Loop";
@@ -1211,6 +1213,20 @@ def LoopOp : CIR_Op<"loop",
   ];
 
   let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+    Region *maybeGetStep() {
+      if (getKind() == LoopOpKind::For)
+        return &getStep();
+      return nullptr;
+    }
+
+    llvm::SmallVector<Region *> getRegionsInExecutionOrder() {
+      if (getKind() == LoopOpKind::For)
+        return llvm::SmallVector<Region *, 3>{&getCond(), &getBody(), &getStep()};
+      return llvm::SmallVector<Region *, 3>{&getCond(), &getBody()};
+    }
+  }];
 }
 
 //===----------------------------------------------------------------------===//

clang/include/clang/CIR/Interfaces/CMakeLists.txt

@@ -22,3 +22,4 @@ endfunction()
 
 add_clang_mlir_attr_interface(ASTAttrInterfaces)
 add_clang_mlir_op_interface(CIROpInterfaces)
+add_clang_mlir_op_interface(LoopOpInterface)

clang/include/clang/CIR/Interfaces/LoopOpInterface.h

@@ -0,0 +1,36 @@
+//===- LoopOpInterface.h - Interface for CIR loop-like ops -----*- C++ -*-===//
+//
+// 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
+//
+//===---------------------------------------------------------------------===//
+//
+// Defines the interface to generically handle CIR loop operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CLANG_INTERFACES_CIR_LOOPOPINTERFACE_H_
+#define CLANG_INTERFACES_CIR_LOOPOPINTERFACE_H_
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+
+namespace mlir {
+namespace cir {
+namespace detail {
+
+/// Verify invariants of the LoopOpInterface.
+::mlir::LogicalResult verifyLoopOpInterface(::mlir::Operation *op);
+
+} // namespace detail
+} // namespace cir
+} // namespace mlir
+
+/// Include the tablegen'd interface declarations.
+#include "clang/CIR/Interfaces/LoopOpInterface.h.inc"
+
+#endif // CLANG_INTERFACES_CIR_LOOPOPINTERFACE_H_

clang/include/clang/CIR/Interfaces/LoopOpInterface.td

@@ -0,0 +1,100 @@
+//===- LoopOpInterface.td - Interface for CIR loop-like ops ----*- C++ -*-===//
+//
+// 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
+//
+//===---------------------------------------------------------------------===//
+
+#ifndef CLANG_CIR_INTERFACES_LOOPOPINTERFACE
+#define CLANG_CIR_INTERFACES_LOOPOPINTERFACE
+
+include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/LoopLikeInterface.td"
+
+def LoopOpInterface : OpInterface<"LoopOpInterface", [
+  DeclareOpInterfaceMethods<RegionBranchOpInterface>,
+  DeclareOpInterfaceMethods<LoopLikeOpInterface>
+]> {
+  let description = [{
+    Contains helper functions to query properties and perform transformations
+    on a loop.
+  }];
+  let cppNamespace = "::mlir::cir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Returns the loop's conditional region.
+      }],
+      /*retTy=*/"mlir::Region &",
+      /*methodName=*/"getCond"
+    >,
+    InterfaceMethod<[{
+        Returns the loop's body region.
+      }],
+      /*retTy=*/"mlir::Region &",
+      /*methodName=*/"getBody"
+    >,
+    InterfaceMethod<[{
+        Returns a pointer to the loop's step region or nullptr.
+      }],
+      /*retTy=*/"mlir::Region *",
+      /*methodName=*/"maybeGetStep",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/"return nullptr;"
+    >,
+    InterfaceMethod<[{
+        Returns the first region to be executed in the loop.
+      }],
+      /*retTy=*/"mlir::Region &",
+      /*methodName=*/"getEntry",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/"return $_op.getCond();"
+    >,
+    InterfaceMethod<[{
+        Returns a list of regions in order of execution.
+      }],
+      /*retTy=*/"llvm::SmallVector<mlir::Region *>",
+      /*methodName=*/"getRegionsInExecutionOrder",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return llvm::SmallVector<mlir::Region *, 2>{&$_op.getRegion(0), &$_op.getRegion(1)};
+      }]
+    >,
+    InterfaceMethod<[{
+        Recursively walks the body of the loop in pre-order while skipping
+        nested loops and executing a callback on every other operation.
+      }],
+      /*retTy=*/"mlir::WalkResult",
+      /*methodName=*/"walkBodySkippingNestedLoops",
+      /*args=*/(ins "::llvm::function_ref<void (Operation *)>":$callback),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getBody().template walk<WalkOrder::PreOrder>([&](Operation *op) {
+          if (isa<LoopOpInterface>(op))
+            return mlir::WalkResult::skip();
+          callback(op);
+          return mlir::WalkResult::advance();
+        });
+      }]
+    >
+  ];
+
+  let extraClassDeclaration = [{
+    /// Generic method to retrieve the successors of a LoopOpInterface operation.
+    static void getLoopOpSuccessorRegions(
+        ::mlir::cir::LoopOpInterface op, ::mlir::RegionBranchPoint point,
+        ::mlir::SmallVectorImpl<::mlir::RegionSuccessor> &regions);
+  }];
+
+  let verify = [{
+    /// Verify invariants of the LoopOpInterface.
+    return detail::verifyLoopOpInterface($_op);
+  }];
+}
+
+#endif // CLANG_CIR_INTERFACES_LOOPOPINTERFACE

clang/lib/CIR/Dialect/IR/CIRDialect.cpp

@@ -14,6 +14,7 @@
 #include "clang/CIR/Dialect/IR/CIRAttrs.h"
 #include "clang/CIR/Dialect/IR/CIROpsEnums.h"
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
+#include "clang/CIR/Interfaces/LoopOpInterface.h"
 #include "llvm/Support/ErrorHandling.h"
 #include <optional>
 
@@ -232,7 +233,7 @@ void AllocaOp::build(::mlir::OpBuilder &odsBuilder,
 //===----------------------------------------------------------------------===//
 
 LogicalResult BreakOp::verify() {
-  if (!getOperation()->getParentOfType<LoopOp>() &&
+  if (!getOperation()->getParentOfType<LoopOpInterface>() &&
       !getOperation()->getParentOfType<SwitchOp>())
     return emitOpError("must be within a loop or switch");
   return success();
@@ -251,7 +252,7 @@ void ConditionOp::getSuccessorRegions(
   // down its list of possible successors.
 
   // Parent is a loop: condition may branch to the body or to the parent op.
-  if (auto loopOp = dyn_cast<LoopOp>(getOperation()->getParentOp())) {
+  if (auto loopOp = dyn_cast<LoopOpInterface>(getOperation()->getParentOp())) {
     regions.emplace_back(&loopOp.getBody(), loopOp.getBody().getArguments());
     regions.emplace_back(loopOp->getResults());
   }
@@ -269,7 +270,7 @@ ConditionOp::getMutableSuccessorOperands(RegionBranchPoint point) {
 }
 
 LogicalResult ConditionOp::verify() {
-  if (!isa<LoopOp, AwaitOp>(getOperation()->getParentOp()))
+  if (!isa<LoopOpInterface, AwaitOp>(getOperation()->getParentOp()))
     return emitOpError("condition must be within a conditional region");
   return success();
 }
@@ -365,7 +366,7 @@ OpFoldResult ConstantOp::fold(FoldAdaptor /*adaptor*/) { return getValue(); }
 //===----------------------------------------------------------------------===//
 
 LogicalResult ContinueOp::verify() {
-  if (!this->getOperation()->getParentOfType<LoopOp>())
+  if (!this->getOperation()->getParentOfType<LoopOpInterface>())
     return emitOpError("must be within a loop");
   return success();
 }
@@ -1264,38 +1265,14 @@ void LoopOp::build(OpBuilder &builder, OperationState &result,
   stepBuilder(builder, result.location);
 }
 
-/// Given the region at `index`, or the parent operation if `index` is None,
-/// return the successor regions. These are the regions that may be selected
-/// during the flow of control. `operands` is a set of optional attributes
-/// that correspond to a constant value for each operand, or null if that
-/// operand is not a constant.
 void LoopOp::getSuccessorRegions(mlir::RegionBranchPoint point,
                                  SmallVectorImpl<RegionSuccessor> &regions) {
-  // If any index all the underlying regions branch back to the parent
-  // operation.
-  if (!point.isParent()) {
-    regions.push_back(RegionSuccessor());
-    return;
-  }
-
-  // FIXME: we want to look at cond region for getting more accurate results
-  // if the other regions will get a chance to execute.
-  regions.push_back(RegionSuccessor(&this->getCond()));
-  regions.push_back(RegionSuccessor(&this->getBody()));
-  regions.push_back(RegionSuccessor(&this->getStep()));
+  LoopOpInterface::getLoopOpSuccessorRegions(*this, point, regions);
 }
 
 llvm::SmallVector<Region *> LoopOp::getLoopRegions() { return {&getBody()}; }
 
-LogicalResult LoopOp::verify() {
-  if (getCond().empty())
-    return emitOpError() << "cond region must not be empty";
-
-  if (!llvm::isa<ConditionOp>(getCond().back().getTerminator()))
-    return emitOpError() << "cond region terminate with 'cir.condition'";
-
-  return success();
-}
+LogicalResult LoopOp::verify() { return success(); }
 
 //===----------------------------------------------------------------------===//
 // GlobalOp

clang/lib/CIR/Dialect/Transforms/LifetimeCheck.cpp

@@ -674,27 +674,10 @@ void LifetimeCheckPass::checkLoop(LoopOp loopOp) {
   SmallVector<Region *, 4> regionsToCheck;
 
   auto setupLoopRegionsToCheck = [&](bool isSubsequentTaken = false) {
-    regionsToCheck.clear();
-    switch (loopOp.getKind()) {
-    case LoopOpKind::For: {
-      regionsToCheck.push_back(&loopOp.getCond());
-      regionsToCheck.push_back(&loopOp.getBody());
-      if (!isSubsequentTaken)
-        regionsToCheck.push_back(&loopOp.getStep());
-      break;
-    }
-    case LoopOpKind::While: {
-      regionsToCheck.push_back(&loopOp.getCond());
-      regionsToCheck.push_back(&loopOp.getBody());
-      break;
-    }
-    case LoopOpKind::DoWhile: {
-      // Note this is the reverse order from While above.
-      regionsToCheck.push_back(&loopOp.getBody());
-      regionsToCheck.push_back(&loopOp.getCond());
-      break;
-    }
-    }
+    regionsToCheck = loopOp.getRegionsInExecutionOrder();
+    // Drop step if it exists and we are not checking the subsequent taken.
+    if (loopOp.maybeGetStep() && !isSubsequentTaken)
+      regionsToCheck.pop_back();
   };
 
   // From 2.4.9 "Note":

clang/lib/CIR/Interfaces/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_clang_library(MLIRCIRInterfaces
   ASTAttrInterfaces.cpp
   CIROpInterfaces.cpp
+  LoopOpInterface.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces

clang/lib/CIR/Interfaces/LoopOpInterface.cpp

@@ -0,0 +1,52 @@
+//===- LoopOpInterface.cpp - Interface for CIR loop-like ops ---*- C++ -*-===//
+//
+// 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 "clang/CIR/Interfaces/LoopOpInterface.h"
+
+#include "clang/CIR/Dialect/IR/CIRDialect.h"
+#include "clang/CIR/Interfaces/LoopOpInterface.cpp.inc"
+
+namespace mlir {
+namespace cir {
+
+void LoopOpInterface::getLoopOpSuccessorRegions(
+    LoopOpInterface op, RegionBranchPoint point,
+    SmallVectorImpl<RegionSuccessor> &regions) {
+  assert(point.isParent() || point.getRegionOrNull());
+
+  // Branching to first region: go to condition or body (do-while).
+  if (point.isParent()) {
+    regions.emplace_back(&op.getEntry(), op.getEntry().getArguments());
+  }
+  // Branching from condition: go to body or exit.
+  else if (&op.getCond() == point.getRegionOrNull()) {
+    regions.emplace_back(RegionSuccessor(op->getResults()));
+    regions.emplace_back(&op.getBody(), op.getBody().getArguments());
+  }
+  // Branching from body: go to step (for) or condition (while).
+  else if (&op.getBody() == point.getRegionOrNull()) {
+    auto *afterBody = (op.maybeGetStep() ? op.maybeGetStep() : &op.getCond());
+    regions.emplace_back(afterBody, afterBody->getArguments());
+  }
+  // Branching from step: go to condition.
+  else if (op.maybeGetStep() == point.getRegionOrNull()) {
+    regions.emplace_back(&op.getCond(), op.getCond().getArguments());
+  }
+}
+
+/// Verify invariants of the LoopOpInterface.
+LogicalResult detail::verifyLoopOpInterface(Operation *op) {
+  auto loopOp = cast<LoopOpInterface>(op);
+  if (!isa<ConditionOp>(loopOp.getCond().back().getTerminator()))
+    return op->emitOpError(
+        "expected condition region to terminate with 'cir.condition'");
+  return success();
+}
+
+} // namespace cir
+} // namespace mlir

clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp

@@ -447,14 +447,14 @@ class CIRLoopOpLowering : public mlir::OpConversionPattern<mlir::cir::LoopOp> {
     // Lower continue statements.
     mlir::Block &dest =
         (kind != LoopKind::For ? condFrontBlock : stepFrontBlock);
-    walkRegionSkipping<mlir::cir::LoopOp>(
+    walkRegionSkipping<mlir::cir::LoopOpInterface>(
         loopOp.getBody(), [&](mlir::Operation *op) {
           if (isa<mlir::cir::ContinueOp>(op))
             lowerTerminator(op, &dest, rewriter);
         });
 
     // Lower break statements.
-    walkRegionSkipping<mlir::cir::LoopOp, mlir::cir::SwitchOp>(
+    walkRegionSkipping<mlir::cir::LoopOpInterface, mlir::cir::SwitchOp>(
         loopOp.getBody(), [&](mlir::Operation *op) {
           if (isa<mlir::cir::BreakOp>(op))
             lowerTerminator(op, continueBlock, rewriter);
@@ -1386,7 +1386,7 @@ class CIRSwitchOpLowering
       }
 
       // Handle break statements.
-      walkRegionSkipping<mlir::cir::LoopOp, mlir::cir::SwitchOp>(
+      walkRegionSkipping<mlir::cir::LoopOpInterface, mlir::cir::SwitchOp>(
           region, [&](mlir::Operation *op) {
             if (isa<mlir::cir::BreakOp>(op))
               lowerTerminator(op, exitBlock, rewriter);