[CIR] Add 'core-flat' as an option value for '-emit-mlir='

- Both 'func' and 'scf' have structure control flow constraints, e.g.,
  'func.return' needs to have 'func.func' parent op. The alternative
  approach to lower CIR into MLIR standard dialects is to lower CIR into
  'cf' dialect and then perform structurization on all or selected
  regions.
- Add 'core-flat' option to enable CFG flattening when lowering CIR into
  MLIR standard dialects.
- Enhance 'cir-flatten-cfg' pass to unify returns into branches to a
  dedicated return block.
- Fix 'cir.br' and 'cir.return' lowering to MLIR and allow function
  declarations.

Author: darkbuck

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

@@ -37,13 +37,15 @@ std::unique_ptr<Pass> createIdiomRecognizerPass(clang::ASTContext *astCtx);
 std::unique_ptr<Pass> createLibOptPass();
 std::unique_ptr<Pass> createLibOptPass(clang::ASTContext *astCtx);
 std::unique_ptr<Pass> createFlattenCFGPass();
+std::unique_ptr<Pass> createFlattenCFGPass(bool throughMLIR);
 std::unique_ptr<Pass> createHoistAllocasPass();
 std::unique_ptr<Pass> createGotoSolverPass();
 
 /// Create a pass to lower ABI-independent function definitions/calls.
 std::unique_ptr<Pass> createCallConvLoweringPass();
 
-void populateCIRPreLoweringPasses(mlir::OpPassManager &pm, bool useCCLowering);
+void populateCIRPreLoweringPasses(mlir::OpPassManager &pm, bool useCCLowering,
+                                  bool emitCore);
 
 //===----------------------------------------------------------------------===//
 // Registration

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

@@ -127,6 +127,10 @@ def FlattenCFG : Pass<"cir-flatten-cfg"> {
     In other words, this pass removes such CIR operations like IfOp, LoopOp,
     ScopeOp and etc. and produces a flat CIR.
   }];
+  let options = [
+    Option<"throughMLIR", "through-mlir", "bool", /*default=*/"false",
+        "Prepare the flat CIR for lowering through MLIR">,
+  ];
   let constructor = "mlir::createFlattenCFGPass()";
   let dependentDialects = ["cir::CIRDialect"];
 }

clang/include/clang/Driver/Options.td

@@ -3067,9 +3067,9 @@ def emit_mlir : Flag<["-"], "emit-mlir">, Visibility<[ClangOption]>, Group<Actio
 def emit_mlir_EQ : Joined<["-"], "emit-mlir=">, Visibility<[ClangOption, CC1Option]>, Group<Action_Group>,
   HelpText<"Build ASTs and then lower through ClangIR to the selected MLIR dialect, emit the .mlir file. "
   "Allowed values are `core` for MLIR standard dialects and `llvm` for the LLVM dialect.">,
-  Values<"core,llvm,cir,cir-flat">,
+  Values<"core,llvm,cir,cir-flat,core-flat">,
   NormalizedValuesScope<"frontend">,
-  NormalizedValues<["MLIR_CORE", "MLIR_LLVM", "MLIR_CIR", "MLIR_CIR_FLAT"]>,
+  NormalizedValues<["MLIR_CORE", "MLIR_LLVM", "MLIR_CIR", "MLIR_CIR_FLAT", "MLIR_CORE_FLAT"]>,
   MarshallingInfoEnum<FrontendOpts<"MLIRTargetDialect">, "MLIR_CORE">;
 def emit_cir : Flag<["-"], "emit-cir">, Visibility<[ClangOption, CC1Option]>,
   Group<Action_Group>, Alias<emit_mlir_EQ>, AliasArgs<["cir"]>,

clang/include/clang/Frontend/FrontendOptions.h

@@ -151,7 +151,13 @@ enum ActionKind {
   PrintDependencyDirectivesSourceMinimizerOutput
 };
 
-enum MLIRDialectKind { MLIR_CORE, MLIR_LLVM, MLIR_CIR, MLIR_CIR_FLAT };
+enum MLIRDialectKind {
+  MLIR_CORE,
+  MLIR_LLVM,
+  MLIR_CIR,
+  MLIR_CIR_FLAT,
+  MLIR_CORE_FLAT
+};
 
 } // namespace frontend
 

clang/lib/CIR/CodeGen/CIRPasses.cpp

@@ -76,7 +76,7 @@ mlir::LogicalResult runCIRToCIRPasses(
   pm.addPass(mlir::createLoweringPreparePass(&astContext));
 
   if (flattenCIR || enableMem2Reg)
-    mlir::populateCIRPreLoweringPasses(pm, enableCallConvLowering);
+    mlir::populateCIRPreLoweringPasses(pm, enableCallConvLowering, emitCore);
 
   if (enableMem2Reg)
     pm.addPass(mlir::createMem2Reg());
@@ -96,11 +96,12 @@ mlir::LogicalResult runCIRToCIRPasses(
 
 namespace mlir {
 
-void populateCIRPreLoweringPasses(OpPassManager &pm, bool useCCLowering) {
+void populateCIRPreLoweringPasses(OpPassManager &pm, bool useCCLowering,
+                                  bool emitCore) {
   if (useCCLowering)
     pm.addPass(createCallConvLoweringPass());
   pm.addPass(createHoistAllocasPass());
-  pm.addPass(createFlattenCFGPass());
+  pm.addPass(createFlattenCFGPass(/*throughMLIR=*/emitCore));
   pm.addPass(createGotoSolverPass());
 }
 

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

@@ -10,6 +10,7 @@
 // function region.
 //
 //===----------------------------------------------------------------------===//
+
 #include "PassDetail.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/PatternMatch.h"
@@ -50,6 +51,8 @@ struct FlattenCFGPass : public FlattenCFGBase<FlattenCFGPass> {
 
   FlattenCFGPass() = default;
   void runOnOperation() override;
+
+  void setThroughMLIR(bool v) { throughMLIR = v; }
 };
 
 struct CIRIfFlattening : public OpRewritePattern<IfOp> {
@@ -855,6 +858,7 @@ class CIRSwitchOpFlattening : public mlir::OpRewritePattern<cir::SwitchOp> {
     return mlir::success();
   }
 };
+
 class CIRTernaryOpFlattening : public mlir::OpRewritePattern<cir::TernaryOp> {
 public:
   using OpRewritePattern<cir::TernaryOp>::OpRewritePattern;
@@ -906,23 +910,87 @@ class CIRTernaryOpFlattening : public mlir::OpRewritePattern<cir::TernaryOp> {
   }
 };
 
-void populateFlattenCFGPatterns(RewritePatternSet &patterns) {
+/// Rewrite 'cir.return' within the specified 'cir.func' as a branch to a
+/// dedicated return block when the CIR needs lowering through MLIR, where
+/// 'func.return' must have 'func.func' as its parent. If the structured
+/// control flow is preferred, that implies a function could only have a single
+/// (unified) 'func.return' in that MLIR.
+class CIRReturnUnifying : public mlir::OpRewritePattern<cir::FuncOp> {
+  using OpRewritePattern<cir::FuncOp>::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(cir::FuncOp func,
+                  mlir::PatternRewriter &rewriter) const override {
+    mlir::OpBuilder::InsertionGuard guard(rewriter);
+
+    if (func.getRegion().empty())
+      return mlir::success();
+
+    // Collect operations to apply patterns.
+    llvm::SmallVector<cir::ReturnOp, 4> returnOps;
+    bool hasReturnOpWithFuncOpParentOnly = true;
+    func->walk([&](cir::ReturnOp ret) {
+      returnOps.push_back(ret);
+      // Check any 'cir.return' without 'cir.func' as the parent. Such
+      // 'cir.return' needs unifying even when there is just one.
+      if (!isa<cir::FuncOp>(ret->getParentOp()))
+        hasReturnOpWithFuncOpParentOnly = false;
+    });
+
+    // Skip unifying if there is only one 'cir.return' and it has 'cir.func' as
+    // the parent.
+    if (hasReturnOpWithFuncOpParentOnly && returnOps.size() < 2)
+      return mlir::success();
+
+    bool hasRetVals = (func.getNumResults() > 0);
+    auto *endBody = &func.getBody().back();
+    // Create a dedicated return block at the end of the function.
+    auto *retBlock = endBody->splitBlock(endBody->end());
+    if (hasRetVals)
+      retBlock->addArguments(func.getResultTypes(), func.getLoc());
+
+    // Rewrite all 'cir.return's as branches to that return block.
+    for (ReturnOp ret : returnOps) {
+      rewriter.setInsertionPoint(ret);
+      rewriter.replaceOpWithNewOp<cir::BrOp>(ret, retBlock, ret.getInput());
+    }
+
+    // Finally, add returnOp in that dedicated return block.
+    auto builder = OpBuilder::atBlockBegin(retBlock);
+    if (hasRetVals)
+      builder.create<cir::ReturnOp>(func.getLoc(), retBlock->getArguments());
+    else
+      builder.create<cir::ReturnOp>(func.getLoc());
+    return mlir::success();
+  }
+};
+
+void populateFlattenCFGPatterns(RewritePatternSet &patterns, bool throughMLIR) {
   patterns
       .add<CIRIfFlattening, CIRLoopOpInterfaceFlattening, CIRScopeOpFlattening,
            CIRSwitchOpFlattening, CIRTernaryOpFlattening, CIRTryOpFlattening>(
           patterns.getContext());
+  if (throughMLIR)
+    patterns.add<CIRReturnUnifying>(patterns.getContext());
 }
 
 void FlattenCFGPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
-  populateFlattenCFGPatterns(patterns);
+  populateFlattenCFGPatterns(patterns, throughMLIR);
 
   // Collect operations to apply patterns.
   llvm::SmallVector<Operation *, 16> ops;
   getOperation()->walk<mlir::WalkOrder::PostOrder>([&](Operation *op) {
     if (isa<IfOp, ScopeOp, SwitchOp, LoopOpInterface, TernaryOp, TryOp>(op))
       ops.push_back(op);
   });
+  // If CIR is to be lowered into MLIR, unify all 'cir.return's as
+  // 'func.return' needs to has 'func.func' as the parent.
+  if (throughMLIR)
+    getOperation()->walk<mlir::WalkOrder::PostOrder>([&](Operation *op) {
+      if (isa<FuncOp>(op))
+        ops.push_back(op);
+    });
 
   // Apply patterns.
   if (applyOpPatternsGreedily(ops, std::move(patterns)).failed())
@@ -937,4 +1005,10 @@ std::unique_ptr<Pass> createFlattenCFGPass() {
   return std::make_unique<FlattenCFGPass>();
 }
 
+std::unique_ptr<Pass> createFlattenCFGPass(bool throughMLIR) {
+  auto flatten = std::make_unique<FlattenCFGPass>();
+  flatten->setThroughMLIR(throughMLIR);
+  return std::move(flatten);
+}
+
 } // namespace mlir

clang/lib/CIR/FrontendAction/CIRGenAction.cpp

@@ -206,12 +206,16 @@ class CIRGenConsumer : public clang::ASTConsumer {
           feOptions.ClangIRCallConvLowering &&
           !(action == CIRGenAction::OutputType::EmitMLIR &&
             feOptions.MLIRTargetDialect == frontend::MLIR_CIR);
+
       bool flattenCIR =
           action == CIRGenAction::OutputType::EmitMLIR &&
-          feOptions.MLIRTargetDialect == clang::frontend::MLIR_CIR_FLAT;
+          (feOptions.MLIRTargetDialect == clang::frontend::MLIR_CORE_FLAT ||
+           feOptions.MLIRTargetDialect == clang::frontend::MLIR_CIR_FLAT);
 
-      bool emitCore = action == CIRGenAction::OutputType::EmitMLIR &&
-                      feOptions.MLIRTargetDialect == clang::frontend::MLIR_CORE;
+      bool emitCore =
+          action == CIRGenAction::OutputType::EmitMLIR &&
+          (feOptions.MLIRTargetDialect == clang::frontend::MLIR_CORE ||
+           feOptions.MLIRTargetDialect == clang::frontend::MLIR_CORE_FLAT);
 
       // Setup and run CIR pipeline.
       std::string passOptParsingFailure;
@@ -283,6 +287,7 @@ class CIRGenConsumer : public clang::ASTConsumer {
     case CIRGenAction::OutputType::EmitMLIR: {
       switch (feOptions.MLIRTargetDialect) {
       case clang::frontend::MLIR_CORE:
+      case clang::frontend::MLIR_CORE_FLAT:
         // case for direct lowering is already checked in compiler invocation
         // no need to check here
         emitMLIR(lowerFromCIRToMLIR(mlirMod, mlirCtx.get()), false);

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

@@ -4780,7 +4780,7 @@ std::unique_ptr<mlir::Pass> createConvertCIRToLLVMPass() {
 }
 
 void populateCIRToLLVMPasses(mlir::OpPassManager &pm, bool useCCLowering) {
-  populateCIRPreLoweringPasses(pm, useCCLowering);
+  populateCIRPreLoweringPasses(pm, useCCLowering, /*emitCore=*/false);
   pm.addPass(createConvertCIRToLLVMPass());
 }
 

clang/lib/CIR/Lowering/ThroughMLIR/LowerCIRToMLIR.cpp

@@ -69,6 +69,8 @@ class CIRReturnLowering : public mlir::OpConversionPattern<cir::ReturnOp> {
   mlir::LogicalResult
   matchAndRewrite(cir::ReturnOp op, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
+    assert(isa<mlir::FunctionOpInterface>(op->getParentOp()) &&
+           "'func.return' op expects parent op 'func.func'");
     rewriter.replaceOpWithNewOp<mlir::func::ReturnOp>(op,
                                                       adaptor.getOperands());
     return mlir::LogicalResult::success();
@@ -660,10 +662,13 @@ class CIRFuncOpLowering : public mlir::OpConversionPattern<cir::FuncOp> {
                                  resultType ? mlir::TypeRange(resultType)
                                             : mlir::TypeRange()));
 
-    if (failed(rewriter.convertRegionTypes(&op.getBody(), *typeConverter,
-                                           &signatureConversion)))
-      return mlir::failure();
-    rewriter.inlineRegionBefore(op.getBody(), fn.getBody(), fn.end());
+    if (!op.getBody().empty()) {
+      if (failed(rewriter.convertRegionTypes(&op.getBody(), *typeConverter,
+                                             &signatureConversion)))
+        return mlir::failure();
+      rewriter.inlineRegionBefore(op.getBody(), fn.getBody(), fn.end());
+    } else
+      fn.setPrivate();
 
     rewriter.eraseOp(op);
     return mlir::LogicalResult::success();
@@ -833,14 +838,15 @@ class CIRCmpOpLowering : public mlir::OpConversionPattern<cir::CmpOp> {
   }
 };
 
-class CIRBrOpLowering : public mlir::OpRewritePattern<cir::BrOp> {
+class CIRBrOpLowering : public mlir::OpConversionPattern<cir::BrOp> {
 public:
-  using OpRewritePattern<cir::BrOp>::OpRewritePattern;
+  using mlir::OpConversionPattern<cir::BrOp>::OpConversionPattern;
 
   mlir::LogicalResult
-  matchAndRewrite(cir::BrOp op,
-                  mlir::PatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<mlir::cf::BranchOp>(op, op.getDest());
+  matchAndRewrite(cir::BrOp op, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<mlir::cf::BranchOp>(op, op.getDest(),
+                                                    adaptor.getDestOperands());
     return mlir::LogicalResult::success();
   }
 };
@@ -1354,24 +1360,23 @@ class CIRPtrStrideOpLowering
 
 void populateCIRToMLIRConversionPatterns(mlir::RewritePatternSet &patterns,
                                          mlir::TypeConverter &converter) {
-  patterns.add<CIRReturnLowering, CIRBrOpLowering>(patterns.getContext());
-
-  patterns.add<
-      CIRCmpOpLowering, CIRCallOpLowering, CIRUnaryOpLowering, CIRBinOpLowering,
-      CIRLoadOpLowering, CIRConstantOpLowering, CIRStoreOpLowering,
-      CIRAllocaOpLowering, CIRFuncOpLowering, CIRScopeOpLowering,
-      CIRBrCondOpLowering, CIRTernaryOpLowering, CIRYieldOpLowering,
-      CIRCosOpLowering, CIRGlobalOpLowering, CIRGetGlobalOpLowering,
-      CIRCastOpLowering, CIRPtrStrideOpLowering, CIRSqrtOpLowering,
-      CIRCeilOpLowering, CIRExp2OpLowering, CIRExpOpLowering, CIRFAbsOpLowering,
-      CIRAbsOpLowering, CIRFloorOpLowering, CIRLog10OpLowering,
-      CIRLog2OpLowering, CIRLogOpLowering, CIRRoundOpLowering,
-      CIRPtrStrideOpLowering, CIRSinOpLowering, CIRShiftOpLowering,
-      CIRBitClzOpLowering, CIRBitCtzOpLowering, CIRBitPopcountOpLowering,
-      CIRBitClrsbOpLowering, CIRBitFfsOpLowering, CIRBitParityOpLowering,
-      CIRIfOpLowering, CIRVectorCreateLowering, CIRVectorInsertLowering,
-      CIRVectorExtractLowering, CIRVectorCmpOpLowering>(converter,
-                                                        patterns.getContext());
+  patterns.add<CIRBrOpLowering, CIRReturnLowering, CIRCmpOpLowering,
+               CIRCallOpLowering, CIRUnaryOpLowering, CIRBinOpLowering,
+               CIRLoadOpLowering, CIRConstantOpLowering, CIRStoreOpLowering,
+               CIRAllocaOpLowering, CIRFuncOpLowering, CIRScopeOpLowering,
+               CIRBrCondOpLowering, CIRTernaryOpLowering, CIRYieldOpLowering,
+               CIRCosOpLowering, CIRGlobalOpLowering, CIRGetGlobalOpLowering,
+               CIRCastOpLowering, CIRPtrStrideOpLowering, CIRSqrtOpLowering,
+               CIRCeilOpLowering, CIRExp2OpLowering, CIRExpOpLowering,
+               CIRFAbsOpLowering, CIRAbsOpLowering, CIRFloorOpLowering,
+               CIRLog10OpLowering, CIRLog2OpLowering, CIRLogOpLowering,
+               CIRRoundOpLowering, CIRPtrStrideOpLowering, CIRSinOpLowering,
+               CIRShiftOpLowering, CIRBitClzOpLowering, CIRBitCtzOpLowering,
+               CIRBitPopcountOpLowering, CIRBitClrsbOpLowering,
+               CIRBitFfsOpLowering, CIRBitParityOpLowering, CIRIfOpLowering,
+               CIRVectorCreateLowering, CIRVectorInsertLowering,
+               CIRVectorExtractLowering, CIRVectorCmpOpLowering>(
+      converter, patterns.getContext());
 }
 
 static mlir::TypeConverter prepareTypeConverter() {