[CIR][IR] Refactor for loops

This patch completes the deprecation of the generic `cir.loop` operation
by adding a new `cir.for` operation and removing the `cir.loop` op. The
new representation removes some bloat and places the regions in order of
execution.

ghstack-source-id: 886e0da
Pull Request resolved: #409

Author: sitio-couto

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

@@ -442,7 +442,7 @@ def StoreOp : CIR_Op<"store", [
 
 def ReturnOp : CIR_Op<"return", [ParentOneOf<["FuncOp", "ScopeOp", "IfOp",
                                               "SwitchOp", "DoWhileOp",
-                                              "WhileOp", "LoopOp"]>,
+                                              "WhileOp", "ForOp"]>,
                                  Terminator]> {
   let summary = "Return from function";
   let description = [{
@@ -635,7 +635,7 @@ def ConditionOp : CIR_Op<"condition", [
 //===----------------------------------------------------------------------===//
 
 def YieldOp : CIR_Op<"yield", [ReturnLike, Terminator,
-    ParentOneOf<["IfOp", "ScopeOp", "SwitchOp", "WhileOp", "LoopOp", "AwaitOp",
+    ParentOneOf<["IfOp", "ScopeOp", "SwitchOp", "WhileOp", "ForOp", "AwaitOp",
                  "TernaryOp", "GlobalOp", "DoWhileOp"]>]> {
   let summary = "Represents the default branching behaviour of a region";
   let description = [{
@@ -1159,106 +1159,6 @@ def BrCondOp : CIR_Op<"brcond",
   }];
 }
 
-//===----------------------------------------------------------------------===//
-// LoopOp
-//===----------------------------------------------------------------------===//
-
-def LoopOpKind_For : I32EnumAttrCase<"For", 1, "for">;
-
-def LoopOpKind : I32EnumAttr<
-    "LoopOpKind",
-    "Loop kind",
-    [LoopOpKind_For]> {
-  let cppNamespace = "::mlir::cir";
-}
-
-def LoopOp : CIR_Op<"loop",
-      [LoopOpInterface,
-       DeclareOpInterfaceMethods<LoopLikeOpInterface>,
-       DeclareOpInterfaceMethods<RegionBranchOpInterface>,
-       RecursivelySpeculatable, NoRegionArguments]> {
-  let summary = "Loop";
-  let description = [{
-    `cir.loop` represents C/C++ loop forms. It defines 3 blocks:
-    - `cond`: region can contain multiple blocks, terminated by regular
-    `cir.yield` when control should yield back to the parent, and
-    `cir.continue` when execution continues to the next region.
-    The region destination depends on the loop form specified.
-    - `step`: region with one block, containing code to compute the
-    loop step, must be terminated with `cir.yield`.
-    - `body`: region for the loop's body, can contain an arbitrary
-    number of blocks.
-
-    The loop form: `for`, `while` and `dowhile` must also be specified and
-    each implies the loop regions execution order.
-
-    ```mlir
-      // while (true) {
-      //  i = i + 1;
-      // }
-      cir.loop while(cond :  {
-        cir.continue
-      }, step :  {
-        cir.yield
-      })  {
-        %3 = cir.load %1 : cir.ptr <i32>, i32
-        %4 = cir.const(1 : i32) : i32
-        %5 = cir.binop(add, %3, %4) : i32
-        cir.store %5, %1 : i32, cir.ptr <i32>
-        cir.yield
-      }
-    ```
-  }];
-
-  let arguments = (ins Arg<LoopOpKind, "loop kind">:$kind);
-  let regions = (region AnyRegion:$cond, AnyRegion:$body,
-                        SizedRegion<1>:$step);
-
-  let assemblyFormat = [{
-    $kind
-    `(`
-    `cond` `:` $cond `,`
-    `step` `:` $step
-    `)`
-    $body
-    attr-dict
-  }];
-
-  let skipDefaultBuilders = 1;
-  let builders = [
-    OpBuilder<(ins
-      "cir::LoopOpKind":$kind,
-      CArg<"function_ref<void(OpBuilder &, Location)>",
-           "nullptr">:$condBuilder,
-      CArg<"function_ref<void(OpBuilder &, Location)>",
-           "nullptr">:$bodyBuilder,
-      CArg<"function_ref<void(OpBuilder &, Location)>",
-           "nullptr">:$stepBuilder
-      )>
-  ];
-
-  let hasVerifier = 1;
-
-  let extraClassDeclaration = [{
-    Region *maybeGetStep() {
-      if (getKind() == LoopOpKind::For)
-        return &getStep();
-      return nullptr;
-    }
-
-    llvm::SmallVector<Region *> getRegionsInExecutionOrder() {
-      switch(getKind()) {
-      case LoopOpKind::For:
-        return llvm::SmallVector<Region *, 3>{&getCond(), &getBody(), &getStep()};
-      // case LoopOpKind::While:
-      //   return llvm::SmallVector<Region *, 2>{&getCond(), &getBody()};
-      // case LoopOpKind::DoWhile:
-      //   return llvm::SmallVector<Region *, 2>{&getBody(), &getCond()};
-      }
-    }
-  }];
-}
-
 //===----------------------------------------------------------------------===//
 // While & DoWhileOp
 //===----------------------------------------------------------------------===//
@@ -1337,6 +1237,73 @@ def DoWhileOp : WhileOpBase<"do"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// ForOp
+//===----------------------------------------------------------------------===//
+
+def ForOp : CIR_Op<"for", [LoopOpInterface, NoRegionArguments]> {
+  let summary = "C/C++ for loop counterpart";
+  let description = [{
+    Represents a C/C++ for loop. It consists of three regions:
+
+     - `cond`: single block region with the loop's condition. Should be
+     terminated with a `cir.condition` operation.
+     - `body`: contains the loop body and an arbitrary number of blocks.
+     - `step`: single block region with the loop's step.
+
+    Example:
+
+    ```mlir
+    cir.for cond {
+      cir.condition(%val)
+    } body {
+      cir.break
+    ^bb2:
+      cir.yield
+    } step {
+      cir.yield
+    }
+    ```
+  }];
+
+  let regions = (region SizedRegion<1>:$cond,
+                        MinSizedRegion<1>:$body,
+                        SizedRegion<1>:$step);
+  let assemblyFormat = [{
+    `:` `cond` $cond
+    `body` $body
+    `step` $step
+    attr-dict
+  }];
+
+  let builders = [
+    OpBuilder<(ins "function_ref<void(OpBuilder &, Location)>":$condBuilder,
+                   "function_ref<void(OpBuilder &, Location)>":$bodyBuilder,
+                   "function_ref<void(OpBuilder &, Location)>":$stepBuilder), [{
+        OpBuilder::InsertionGuard guard($_builder);
+
+        // Build condition region.
+        $_builder.createBlock($_state.addRegion());
+        condBuilder($_builder, $_state.location);
+
+        // Build body region.
+        $_builder.createBlock($_state.addRegion());
+        bodyBuilder($_builder, $_state.location);
+
+        // Build step region.
+        $_builder.createBlock($_state.addRegion());
+        stepBuilder($_builder, $_state.location);
+      }]>
+  ];
+
+  let extraClassDeclaration = [{
+    Region *maybeGetStep() { return &getStep(); }
+    llvm::SmallVector<Region *> getRegionsInExecutionOrder() {
+      return llvm::SmallVector<Region *, 3>{&getCond(), &getBody(), &getStep()};
+    }
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // GlobalOp
 //===----------------------------------------------------------------------===//
@@ -2684,7 +2651,7 @@ def AllocException : CIR_Op<"alloc_exception", [
 
 def ThrowOp : CIR_Op<"throw", [
                      ParentOneOf<["FuncOp", "ScopeOp", "IfOp", "SwitchOp",
-                                  "DoWhileOp", "WhileOp", "LoopOp"]>,
+                                  "DoWhileOp", "WhileOp", "ForOp"]>,
                      Terminator]> {
   let summary = "(Re)Throws an exception";
   let description = [{

clang/lib/CIR/CodeGen/CIRGenBuilder.h

@@ -615,6 +615,15 @@ class CIRGenBuilderTy : public CIRBaseBuilderTy {
     return create<mlir::cir::WhileOp>(loc, condBuilder, bodyBuilder);
   }
 
+  /// Create a for operation.
+  mlir::cir::ForOp createFor(
+      mlir::Location loc,
+      llvm::function_ref<void(mlir::OpBuilder &, mlir::Location)> condBuilder,
+      llvm::function_ref<void(mlir::OpBuilder &, mlir::Location)> bodyBuilder,
+      llvm::function_ref<void(mlir::OpBuilder &, mlir::Location)> stepBuilder) {
+    return create<mlir::cir::ForOp>(loc, condBuilder, bodyBuilder, stepBuilder);
+  }
+
   mlir::cir::MemCpyOp createMemCpy(mlir::Location loc, mlir::Value dst,
                                    mlir::Value src, mlir::Value len) {
     return create<mlir::cir::MemCpyOp>(loc, dst, src, len);

clang/lib/CIR/CodeGen/CIRGenStmt.cpp

@@ -674,7 +674,7 @@ CIRGenFunction::buildDefaultStmt(const DefaultStmt &S, mlir::Type condType,
 mlir::LogicalResult
 CIRGenFunction::buildCXXForRangeStmt(const CXXForRangeStmt &S,
                                      ArrayRef<const Attr *> ForAttrs) {
-  mlir::cir::LoopOp loopOp;
+  mlir::cir::ForOp forOp;
 
   // TODO(cir): pass in array of attributes.
   auto forStmtBuilder = [&]() -> mlir::LogicalResult {
@@ -697,8 +697,8 @@ CIRGenFunction::buildCXXForRangeStmt(const CXXForRangeStmt &S,
     // sure we handle all cases.
     assert(!UnimplementedFeature::requiresCleanups());
 
-    loopOp = builder.create<LoopOp>(
-        getLoc(S.getSourceRange()), mlir::cir::LoopOpKind::For,
+    forOp = builder.createFor(
+        getLoc(S.getSourceRange()),
         /*condBuilder=*/
         [&](mlir::OpBuilder &b, mlir::Location loc) {
           assert(!UnimplementedFeature::createProfileWeightsForLoop());
@@ -743,12 +743,12 @@ CIRGenFunction::buildCXXForRangeStmt(const CXXForRangeStmt &S,
   if (res.failed())
     return res;
 
-  terminateBody(builder, loopOp.getBody(), getLoc(S.getEndLoc()));
+  terminateBody(builder, forOp.getBody(), getLoc(S.getEndLoc()));
   return mlir::success();
 }
 
 mlir::LogicalResult CIRGenFunction::buildForStmt(const ForStmt &S) {
-  mlir::cir::LoopOp loopOp;
+  mlir::cir::ForOp forOp;
 
   // TODO: pass in array of attributes.
   auto forStmtBuilder = [&]() -> mlir::LogicalResult {
@@ -764,8 +764,8 @@ mlir::LogicalResult CIRGenFunction::buildForStmt(const ForStmt &S) {
     // sure we handle all cases.
     assert(!UnimplementedFeature::requiresCleanups());
 
-    loopOp = builder.create<LoopOp>(
-        getLoc(S.getSourceRange()), mlir::cir::LoopOpKind::For,
+    forOp = builder.createFor(
+        getLoc(S.getSourceRange()),
         /*condBuilder=*/
         [&](mlir::OpBuilder &b, mlir::Location loc) {
           assert(!UnimplementedFeature::createProfileWeightsForLoop());
@@ -822,7 +822,7 @@ mlir::LogicalResult CIRGenFunction::buildForStmt(const ForStmt &S) {
   if (res.failed())
     return res;
 
-  terminateBody(builder, loopOp.getBody(), getLoc(S.getEndLoc()));
+  terminateBody(builder, forOp.getBody(), getLoc(S.getEndLoc()));
   return mlir::success();
 }
 

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

@@ -1238,42 +1238,6 @@ void CatchOp::build(
   catchBuilder(builder, result.location, result);
 }
 
-//===----------------------------------------------------------------------===//
-// LoopOp
-//===----------------------------------------------------------------------===//
-
-void LoopOp::build(OpBuilder &builder, OperationState &result,
-                   cir::LoopOpKind kind,
-                   function_ref<void(OpBuilder &, Location)> condBuilder,
-                   function_ref<void(OpBuilder &, Location)> bodyBuilder,
-                   function_ref<void(OpBuilder &, Location)> stepBuilder) {
-  OpBuilder::InsertionGuard guard(builder);
-  ::mlir::cir::LoopOpKindAttr kindAttr =
-      cir::LoopOpKindAttr::get(builder.getContext(), kind);
-  result.addAttribute(getKindAttrName(result.name), kindAttr);
-
-  Region *condRegion = result.addRegion();
-  builder.createBlock(condRegion);
-  condBuilder(builder, result.location);
-
-  Region *bodyRegion = result.addRegion();
-  builder.createBlock(bodyRegion);
-  bodyBuilder(builder, result.location);
-
-  Region *stepRegion = result.addRegion();
-  builder.createBlock(stepRegion);
-  stepBuilder(builder, result.location);
-}
-
-void LoopOp::getSuccessorRegions(mlir::RegionBranchPoint point,
-                                 SmallVectorImpl<RegionSuccessor> &regions) {
-  LoopOpInterface::getLoopOpSuccessorRegions(*this, point, regions);
-}
-
-llvm::SmallVector<Region *> LoopOp::getLoopRegions() { return {&getBody()}; }
-
-LogicalResult LoopOp::verify() { return success(); }
-
 //===----------------------------------------------------------------------===//
 // LoopOpInterface Methods
 //===----------------------------------------------------------------------===//
@@ -1296,6 +1260,14 @@ void WhileOp::getSuccessorRegions(
 
 ::llvm::SmallVector<Region *> WhileOp::getLoopRegions() { return {&getBody()}; }
 
+void ForOp::getSuccessorRegions(
+    ::mlir::RegionBranchPoint point,
+    ::llvm::SmallVectorImpl<::mlir::RegionSuccessor> &regions) {
+  LoopOpInterface::getLoopOpSuccessorRegions(*this, point, regions);
+}
+
+::llvm::SmallVector<Region *> ForOp::getLoopRegions() { return {&getBody()}; }
+
 //===----------------------------------------------------------------------===//
 // GlobalOp
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Interfaces/CIRLoopOpInterface.cpp

@@ -38,6 +38,8 @@ void LoopOpInterface::getLoopOpSuccessorRegions(
   // Branching from step: go to condition.
   else if (op.maybeGetStep() == point.getRegionOrNull()) {
     regions.emplace_back(&op.getCond(), op.getCond().getArguments());
+  } else {
+    llvm_unreachable("unexpected branch origin");
   }
 }
 

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

@@ -2261,7 +2261,6 @@ void ConvertCIRToLLVMPass::runOnOperation() {
                     // ,ConstantOp
                     // ,FuncOp
                     // ,LoadOp
-                    // ,LoopOp
                     // ,ReturnOp
                     // ,StoreOp
                     // ,YieldOp

clang/test/CIR/CodeGen/loop-scope.cpp

@@ -15,15 +15,15 @@ void l0(void) {
 // CPPSCOPE-NEXT:     %1 = cir.alloca !s32i, cir.ptr <!s32i>, ["j", init] {alignment = 4 : i64}
 // CPPSCOPE-NEXT:     %2 = cir.const(#cir.int<0> : !s32i) : !s32i
 // CPPSCOPE-NEXT:     cir.store %2, %0 : !s32i, cir.ptr <!s32i>
-// CPPSCOPE-NEXT:     cir.loop for(cond :  {
+// CPPSCOPE-NEXT:     cir.for : cond {
 
 // CSCOPE: cir.func @l0()
 // CSCOPE-NEXT: cir.scope {
 // CSCOPE-NEXT:   %0 = cir.alloca !s32i, cir.ptr <!s32i>, ["i", init] {alignment = 4 : i64}
 // CSCOPE-NEXT:   %1 = cir.const(#cir.int<0> : !s32i) : !s32i
 // CSCOPE-NEXT:   cir.store %1, %0 : !s32i, cir.ptr <!s32i>
-// CSCOPE-NEXT:   cir.loop for(cond :  {
+// CSCOPE-NEXT:   cir.for : cond {
 
-// CSCOPE:        })  {
+// CSCOPE:        } body {
 // CSCOPE-NEXT:     cir.scope {
 // CSCOPE-NEXT:       %2 = cir.alloca !s32i, cir.ptr <!s32i>, ["j", init] {alignment = 4 : i64}