[CIR] Upstream initial support for CIR flattening (#130648)

The ClangIR CFG has to be flat before it can be lowered to LLVM IR. That
is, there can be no nested regions and all blocks in a region must
belong to the parent region. Currently only cir.scope operations violate
these rules, so the initial implementation of the cir-flatten-cfg pass
only has to transform scope operations.

Author: andykaylor

clang/include/clang/CIR/Dialect/CMakeLists.txt

@@ -1 +1,7 @@
 add_subdirectory(IR)
+
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name CIR)
+mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix CIR)
+mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix CIR)
+add_public_tablegen_target(MLIRCIRPassIncGen)

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

@@ -428,6 +428,46 @@ def ScopeOp : CIR_Op<"scope", [
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// BrOp
+//===----------------------------------------------------------------------===//
+
+def BrOp : CIR_Op<"br",
+      [DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>,
+     Pure, Terminator]> {
+  let summary = "Unconditional branch";
+  let description = [{
+    The `cir.br` branches unconditionally to a block. Used to represent C/C++
+    goto's and general block branching.
+
+    Note that for source level `goto`'s crossing scope boundaries, those are
+    usually represented with the "symbolic" `cir.goto` operation.
+
+    Example:
+
+    ```mlir
+      ...
+        cir.br ^bb3
+      ^bb3:
+        cir.return
+    ```
+  }];
+
+  let builders = [
+    OpBuilder<(ins "mlir::Block *":$dest,
+              CArg<"mlir::ValueRange", "{}">:$destOperands), [{
+      $_state.addSuccessors(dest);
+      $_state.addOperands(destOperands);
+    }]>
+  ];
+
+  let arguments = (ins Variadic<CIR_AnyType>:$destOperands);
+  let successors = (successor AnySuccessor:$dest);
+  let assemblyFormat = [{
+    $dest (`(` $destOperands^ `:` type($destOperands) `)`)? attr-dict
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // GlobalOp
 //===----------------------------------------------------------------------===//

clang/include/clang/CIR/Dialect/Passes.h

@@ -0,0 +1,39 @@
+//===- Passes.h - CIR pass entry points -------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines prototypes that expose pass constructors.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CLANG_CIR_DIALECT_PASSES_H
+#define CLANG_CIR_DIALECT_PASSES_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace clang {
+class ASTContext;
+}
+namespace mlir {
+
+std::unique_ptr<Pass> createCIRFlattenCFGPass();
+
+void populateCIRPreLoweringPasses(mlir::OpPassManager &pm);
+
+//===----------------------------------------------------------------------===//
+// Registration
+//===----------------------------------------------------------------------===//
+
+void registerCIRDialectTranslation(mlir::MLIRContext &context);
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "clang/CIR/Dialect/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // CLANG_CIR_DIALECT_PASSES_H

clang/include/clang/CIR/Dialect/Passes.td

@@ -0,0 +1,28 @@
+//===----------------------------------------------------------------------===//
+//
+// 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_DIALECT_PASSES_TD
+#define CLANG_CIR_DIALECT_PASSES_TD
+
+include "mlir/Pass/PassBase.td"
+
+def CIRFlattenCFG : Pass<"cir-flatten-cfg"> {
+  let summary = "Produces flatten CFG";
+  let description = [{
+    This pass transforms CIR by inlining all the nested regions. Thus,
+    the following conditions are true after the pass applied:
+    - there are no nested regions in any function body
+    - all the blocks in a function belong to the parent region
+    In other words, this pass removes such CIR operations like IfOp, LoopOp,
+    ScopeOp and etc. and produces a flat CIR.
+  }];
+  let constructor = "mlir::createCIRFlattenCFGPass()";
+  let dependentDialects = ["cir::CIRDialect"];
+}
+
+#endif // CLANG_CIR_DIALECT_PASSES_TD

clang/include/clang/CIR/MissingFeatures.h

@@ -82,6 +82,7 @@ struct MissingFeatures {
   static bool objCLifetime() { return false; }
   static bool emitNullabilityCheck() { return false; }
   static bool astVarDeclInterface() { return false; }
+  static bool stackSaveOp() { return false; }
 };
 
 } // namespace cir

clang/lib/CIR/Dialect/CMakeLists.txt

@@ -1 +1,2 @@
 add_subdirectory(IR)
+add_subdirectory(Transforms)

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

@@ -268,6 +268,19 @@ LogicalResult cir::ScopeOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// BrOp
+//===----------------------------------------------------------------------===//
+
+mlir::SuccessorOperands cir::BrOp::getSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return mlir::SuccessorOperands(getDestOperandsMutable());
+}
+
+Block *cir::BrOp::getSuccessorForOperands(ArrayRef<Attribute>) {
+  return getDest();
+}
+
 //===----------------------------------------------------------------------===//
 // GlobalOp
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Dialect/IR/CMakeLists.txt

@@ -11,6 +11,7 @@ add_clang_library(MLIRCIR
 
   LINK_LIBS PUBLIC
   MLIRIR
+  MLIRCIRInterfaces
   MLIRDLTIDialect
   MLIRDataLayoutInterfaces
   MLIRFuncDialect

clang/lib/CIR/Dialect/Transforms/CMakeLists.txt

@@ -0,0 +1,18 @@
+add_clang_library(MLIRCIRTransforms
+  FlattenCFG.cpp
+
+  DEPENDS
+  MLIRCIRPassIncGen
+
+  LINK_LIBS PUBLIC
+  clangAST
+  clangBasic
+
+  MLIRAnalysis
+  MLIRIR
+  MLIRPass
+  MLIRTransformUtils
+
+  MLIRCIR
+  MLIRCIRInterfaces
+)

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

@@ -0,0 +1,117 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 pass that inlines CIR operations regions into the parent
+// function region.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "clang/CIR/Dialect/IR/CIRDialect.h"
+#include "clang/CIR/Dialect/Passes.h"
+#include "clang/CIR/MissingFeatures.h"
+
+using namespace mlir;
+using namespace cir;
+
+namespace {
+
+struct CIRFlattenCFGPass : public CIRFlattenCFGBase<CIRFlattenCFGPass> {
+
+  CIRFlattenCFGPass() = default;
+  void runOnOperation() override;
+};
+
+class CIRScopeOpFlattening : public mlir::OpRewritePattern<cir::ScopeOp> {
+public:
+  using OpRewritePattern<cir::ScopeOp>::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(cir::ScopeOp scopeOp,
+                  mlir::PatternRewriter &rewriter) const override {
+    mlir::OpBuilder::InsertionGuard guard(rewriter);
+    mlir::Location loc = scopeOp.getLoc();
+
+    // Empty scope: just remove it.
+    // TODO: Remove this logic once CIR uses MLIR infrastructure to remove
+    // trivially dead operations. MLIR canonicalizer is too aggressive and we
+    // need to either (a) make sure all our ops model all side-effects and/or
+    // (b) have more options in the canonicalizer in MLIR to temper
+    // aggressiveness level.
+    if (scopeOp.isEmpty()) {
+      rewriter.eraseOp(scopeOp);
+      return mlir::success();
+    }
+
+    // Split the current block before the ScopeOp to create the inlining
+    // point.
+    mlir::Block *currentBlock = rewriter.getInsertionBlock();
+    mlir::Block *continueBlock =
+        rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint());
+    if (scopeOp.getNumResults() > 0)
+      continueBlock->addArguments(scopeOp.getResultTypes(), loc);
+
+    // Inline body region.
+    mlir::Block *beforeBody = &scopeOp.getScopeRegion().front();
+    mlir::Block *afterBody = &scopeOp.getScopeRegion().back();
+    rewriter.inlineRegionBefore(scopeOp.getScopeRegion(), continueBlock);
+
+    // Save stack and then branch into the body of the region.
+    rewriter.setInsertionPointToEnd(currentBlock);
+    assert(!cir::MissingFeatures::stackSaveOp());
+    rewriter.create<cir::BrOp>(loc, mlir::ValueRange(), beforeBody);
+
+    // Replace the scopeop return with a branch that jumps out of the body.
+    // Stack restore before leaving the body region.
+    rewriter.setInsertionPointToEnd(afterBody);
+    if (auto yieldOp = dyn_cast<cir::YieldOp>(afterBody->getTerminator())) {
+      rewriter.replaceOpWithNewOp<cir::BrOp>(yieldOp, yieldOp.getArgs(),
+                                             continueBlock);
+    }
+
+    // Replace the op with values return from the body region.
+    rewriter.replaceOp(scopeOp, continueBlock->getArguments());
+
+    return mlir::success();
+  }
+};
+
+void populateFlattenCFGPatterns(RewritePatternSet &patterns) {
+  patterns.add<CIRScopeOpFlattening>(patterns.getContext());
+}
+
+void CIRFlattenCFGPass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  populateFlattenCFGPatterns(patterns);
+
+  // Collect operations to apply patterns.
+  llvm::SmallVector<Operation *, 16> ops;
+  getOperation()->walk<mlir::WalkOrder::PostOrder>([&](Operation *op) {
+    if (isa<ScopeOp>(op))
+      ops.push_back(op);
+  });
+
+  // Apply patterns.
+  if (applyOpPatternsGreedily(ops, std::move(patterns)).failed())
+    signalPassFailure();
+}
+
+} // namespace
+
+namespace mlir {
+
+std::unique_ptr<Pass> createCIRFlattenCFGPass() {
+  return std::make_unique<CIRFlattenCFGPass>();
+}
+
+} // namespace mlir