[CIR] Emit allocas into the proper lexical scope

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

@@ -22,6 +22,7 @@ namespace mlir {
 
 std::unique_ptr<Pass> createCIRCanonicalizePass();
 std::unique_ptr<Pass> createCIRFlattenCFGPass();
+std::unique_ptr<Pass> createHoistAllocasPass();
 
 void populateCIRPreLoweringPasses(mlir::OpPassManager &pm);
 

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

@@ -29,6 +29,16 @@ def CIRCanonicalize : Pass<"cir-canonicalize"> {
   let dependentDialects = ["cir::CIRDialect"];
 }
 
+def HoistAllocas : Pass<"cir-hoist-allocas"> {
+  let summary = "Hoist allocas to the entry of the function";
+  let description = [{
+    This pass hoist all non-dynamic allocas to the entry of the function.
+    This is helpful for later code generation.
+  }];
+  let constructor = "mlir::createHoistAllocasPass()";
+  let dependentDialects = ["cir::CIRDialect"];
+}
+
 def CIRFlattenCFG : Pass<"cir-flatten-cfg"> {
   let summary = "Produces flatten CFG";
   let description = [{

clang/include/clang/CIR/MissingFeatures.h

@@ -130,13 +130,16 @@ struct MissingFeatures {
   static bool continueOp() { return false; }
   static bool ifOp() { return false; }
   static bool labelOp() { return false; }
+  static bool ptrDiffOp() { return false; }
+  static bool ptrStrideOp() { return false; }
   static bool selectOp() { return false; }
   static bool switchOp() { return false; }
   static bool ternaryOp() { return false; }
   static bool tryOp() { return false; }
   static bool zextOp() { return false; }
-  static bool ptrStrideOp() { return false; }
-  static bool ptrDiffOp() { return false; }
+
+  // Future CIR attributes
+  static bool optInfoAttr() { return false; }
 };
 
 } // namespace cir

clang/lib/CIR/CodeGen/CIRGenDecl.cpp

@@ -49,7 +49,8 @@ CIRGenFunction::emitAutoVarAlloca(const VarDecl &d) {
   // A normal fixed sized variable becomes an alloca in the entry block,
   mlir::Type allocaTy = convertTypeForMem(ty);
   // Create the temp alloca and declare variable using it.
-  address = createTempAlloca(allocaTy, alignment, loc, d.getName());
+  address = createTempAlloca(allocaTy, alignment, loc, d.getName(),
+                             /*insertIntoFnEntryBlock=*/false);
   declare(address.getPointer(), &d, ty, getLoc(d.getSourceRange()), alignment);
 
   emission.Addr = address;

clang/lib/CIR/CodeGen/CIRGenExpr.cpp

@@ -317,10 +317,27 @@ void CIRGenFunction::emitIgnoredExpr(const Expr *e) {
 }
 
 mlir::Value CIRGenFunction::emitAlloca(StringRef name, mlir::Type ty,
-                                       mlir::Location loc,
-                                       CharUnits alignment) {
-  mlir::Block *entryBlock = getCurFunctionEntryBlock();
+                                       mlir::Location loc, CharUnits alignment,
+                                       bool insertIntoFnEntryBlock,
+                                       mlir::Value arraySize) {
+  mlir::Block *entryBlock = insertIntoFnEntryBlock
+                                ? getCurFunctionEntryBlock()
+                                : curLexScope->getEntryBlock();
+
+  // If this is an alloca in the entry basic block of a cir.try and there's
+  // a surrounding cir.scope, make sure the alloca ends up in the surrounding
+  // scope instead. This is necessary in order to guarantee all SSA values are
+  // reachable during cleanups.
+  assert(!cir::MissingFeatures::tryOp());
+
+  return emitAlloca(name, ty, loc, alignment,
+                    builder.getBestAllocaInsertPoint(entryBlock), arraySize);
+}
 
+mlir::Value CIRGenFunction::emitAlloca(StringRef name, mlir::Type ty,
+                                       mlir::Location loc, CharUnits alignment,
+                                       mlir::OpBuilder::InsertPoint ip,
+                                       mlir::Value arraySize) {
   // CIR uses its own alloca address space rather than follow the target data
   // layout like original CodeGen. The data layout awareness should be done in
   // the lowering pass instead.
@@ -331,7 +348,7 @@ mlir::Value CIRGenFunction::emitAlloca(StringRef name, mlir::Type ty,
   mlir::Value addr;
   {
     mlir::OpBuilder::InsertionGuard guard(builder);
-    builder.restoreInsertionPoint(builder.getBestAllocaInsertPoint(entryBlock));
+    builder.restoreInsertionPoint(ip);
     addr = builder.createAlloca(loc, /*addr type*/ localVarPtrTy,
                                 /*var type*/ ty, name, alignIntAttr);
     assert(!cir::MissingFeatures::astVarDeclInterface());
@@ -346,11 +363,13 @@ mlir::Value CIRGenFunction::createDummyValue(mlir::Location loc,
   return builder.createDummyValue(loc, t, alignment);
 }
 
-/// This creates an alloca and inserts it  at the current insertion point of the
-/// builder.
+/// This creates an alloca and inserts it into the entry block if
+/// \p insertIntoFnEntryBlock is true, otherwise it inserts it at the current
+/// insertion point of the builder.
 Address CIRGenFunction::createTempAlloca(mlir::Type ty, CharUnits align,
-                                         mlir::Location loc,
-                                         const Twine &name) {
-  mlir::Value alloca = emitAlloca(name.str(), ty, loc, align);
+                                         mlir::Location loc, const Twine &name,
+                                         bool insertIntoFnEntryBlock) {
+  mlir::Value alloca =
+      emitAlloca(name.str(), ty, loc, align, insertIntoFnEntryBlock);
   return Address(alloca, ty, align);
 }

clang/lib/CIR/CodeGen/CIRGenFunction.cpp

@@ -138,7 +138,8 @@ mlir::Location CIRGenFunction::getLoc(mlir::Location lhs, mlir::Location rhs) {
 void CIRGenFunction::emitAndUpdateRetAlloca(QualType type, mlir::Location loc,
                                             CharUnits alignment) {
   if (!type->isVoidType()) {
-    fnRetAlloca = emitAlloca("__retval", convertType(type), loc, alignment);
+    fnRetAlloca = emitAlloca("__retval", convertType(type), loc, alignment,
+                             /*insertIntoFnEntryBlock=*/false);
   }
 }
 
@@ -293,7 +294,8 @@ void CIRGenFunction::startFunction(GlobalDecl gd, QualType returnType,
 
     mlir::Value addrVal =
         emitAlloca(cast<NamedDecl>(paramVar)->getName(),
-                   convertType(paramVar->getType()), paramLoc, alignment);
+                   convertType(paramVar->getType()), paramLoc, alignment,
+                   /*insertIntoFnEntryBlock=*/true);
 
     declare(addrVal, paramVar, paramVar->getType(), paramLoc, alignment,
             /*isParam=*/true);

clang/lib/CIR/CodeGen/CIRGenFunction.h

@@ -109,7 +109,13 @@ class CIRGenFunction : public CIRGenTypeCache {
 
 public:
   mlir::Value emitAlloca(llvm::StringRef name, mlir::Type ty,
-                         mlir::Location loc, clang::CharUnits alignment);
+                         mlir::Location loc, clang::CharUnits alignment,
+                         bool insertIntoFnEntryBlock,
+                         mlir::Value arraySize = nullptr);
+  mlir::Value emitAlloca(llvm::StringRef name, mlir::Type ty,
+                         mlir::Location loc, clang::CharUnits alignment,
+                         mlir::OpBuilder::InsertPoint ip,
+                         mlir::Value arraySize = nullptr);
 
   mlir::Value createDummyValue(mlir::Location loc, clang::QualType qt);
 
@@ -483,7 +489,7 @@ class CIRGenFunction : public CIRGenTypeCache {
   LexicalScope *curLexScope = nullptr;
 
   Address createTempAlloca(mlir::Type ty, CharUnits align, mlir::Location loc,
-                           const Twine &name = "tmp");
+                           const Twine &name, bool insertIntoFnEntryBlock);
 };
 
 } // namespace clang::CIRGen

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

@@ -1,6 +1,7 @@
 add_clang_library(MLIRCIRTransforms
   CIRCanonicalize.cpp
   FlattenCFG.cpp
+  HoistAllocas.cpp
 
   DEPENDS
   MLIRCIRPassIncGen

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

@@ -0,0 +1,84 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 "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"
+#include "llvm/Support/TimeProfiler.h"
+
+using namespace mlir;
+using namespace cir;
+
+namespace {
+
+struct HoistAllocasPass : public HoistAllocasBase<HoistAllocasPass> {
+
+  HoistAllocasPass() = default;
+  void runOnOperation() override;
+};
+
+static void process(mlir::ModuleOp mod, cir::FuncOp func) {
+  if (func.getRegion().empty())
+    return;
+
+  // Hoist all static allocas to the entry block.
+  mlir::Block &entryBlock = func.getRegion().front();
+  mlir::Operation *insertPoint = &*entryBlock.begin();
+
+  // Post-order is the default, but the code below requires it, so
+  // let's not depend on the default staying that way.
+  func.getBody().walk<mlir::WalkOrder::PostOrder>([&](cir::AllocaOp alloca) {
+    if (alloca->getBlock() == &entryBlock)
+      return;
+    // Don't hoist allocas with dynamic alloca size.
+    assert(!cir::MissingFeatures::opAllocaDynAllocSize());
+
+    // Hoist allocas into the entry block.
+
+    // Preserving the `const` attribute on hoisted allocas can cause LLVM to
+    // incorrectly introduce invariant group metadata in some circumstances.
+    // The incubator performs some analysis to determine whether the attribute
+    // can be preserved, but it only runs this analysis when optimizations are
+    // enabled. Until we start tracking the optimization level, we can just
+    // always remove the `const` attribute.
+    assert(!cir::MissingFeatures::optInfoAttr());
+    if (alloca.getConstant())
+      alloca.setConstant(false);
+
+    alloca->moveBefore(insertPoint);
+  });
+}
+
+void HoistAllocasPass::runOnOperation() {
+  llvm::TimeTraceScope scope("Hoist Allocas");
+  llvm::SmallVector<Operation *, 16> ops;
+
+  Operation *op = getOperation();
+  auto mod = mlir::dyn_cast<mlir::ModuleOp>(op);
+  if (!mod)
+    mod = op->getParentOfType<mlir::ModuleOp>();
+
+  // If we ever introduce nested cir.function ops, we'll need to make this
+  // walk in post-order and recurse into nested functions.
+  getOperation()->walk<mlir::WalkOrder::PreOrder>([&](cir::FuncOp op) {
+    process(mod, op);
+    return mlir::WalkResult::skip();
+  });
+}
+
+} // namespace
+
+std::unique_ptr<Pass> mlir::createHoistAllocasPass() {
+  return std::make_unique<HoistAllocasPass>();
+}

clang/lib/CIR/Lowering/CIRPasses.cpp

@@ -37,6 +37,7 @@ mlir::LogicalResult runCIRToCIRPasses(mlir::ModuleOp theModule,
 namespace mlir {
 
 void populateCIRPreLoweringPasses(OpPassManager &pm) {
+  pm.addPass(createHoistAllocasPass());
   pm.addPass(createCIRFlattenCFGPass());
 }