[CIR] Lower nested local constant alloca

Lancern · Lancern · commit 4e7443d9b6c4 · 2025-02-04T20:56:01.000+08:00
This patch adds support for lowering local constants in nested scopes, including
those in nested loops.

For those constant allocas in non-loop inner scopes, this patch keeps their
constant flags during alloca hoisting. LLVM lowering would correctly emit
necessary invariant metadata for those allocas.

For those constant allocas in a loop, this patch introduces a new operation
`cir.invariant_group` that marks the beginning of the lifetime of the constant
objects. This operation is put at the location of the alloca operation before
hoisting them. This patch updates LLVM lowering to emit the necessary invariant
metadata when loading and storing through such pointers.

This patch takes care of the special case where the constant alloca represents
a variable declared in the condition part of a while loop. In such a case, this
patch removes the constant flag on the alloca operation when hositing them.
diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td
@@ -3562,6 +3562,28 @@ def LLVMIntrinsicCallOp : CIR_Op<"llvm.intrinsic"> {
 
 }
 
+//===----------------------------------------------------------------------===//
+// InvariantGroupOp
+//===----------------------------------------------------------------------===//
+
+def InvariantGroupOp
+    : CIR_Op<"invariant_group", [Pure, SameOperandsAndResultType]> {
+  let summary = "Start an invariant group";
+  let description = [{
+    The `cir.invariant_group` operation takes a single pointer value as argument
+    and returns the same pointer value with a fresh invariant group. All loads
+    and stores that access the returned pointer value are presumed by the
+    optimizer to load or store the same value.
+  }];
+
+  let arguments = (ins CIR_PointerType:$ptr);
+  let results = (outs CIR_PointerType:$result);
+
+  let assemblyFormat = [{
+    $ptr `:` type($result) attr-dict
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // DeleteArrayOp
 //===----------------------------------------------------------------------===//
diff --git a/clang/lib/CIR/Dialect/Transforms/HoistAllocas.cpp b/clang/lib/CIR/Dialect/Transforms/HoistAllocas.cpp
@@ -28,6 +28,27 @@ struct HoistAllocasPass : public HoistAllocasBase<HoistAllocasPass> {
   void runOnOperation() override;
 };
 
+static bool isOpInLoop(mlir::Operation *op) {
+  return op->getParentOfType<cir::LoopOpInterface>();
+}
+
+static bool isWhileCondition(cir::AllocaOp alloca) {
+  for (mlir::Operation *user : alloca->getUsers()) {
+    if (!mlir::isa<cir::StoreOp>(user))
+      continue;
+
+    auto store = mlir::cast<cir::StoreOp>(user);
+    mlir::Operation *storeParentOp = store->getParentOp();
+    if (!mlir::isa<cir::WhileOp>(storeParentOp))
+      continue;
+
+    auto whileOp = mlir::cast<cir::WhileOp>(storeParentOp);
+    return &whileOp.getCond() == store->getParentRegion();
+  }
+
+  return false;
+}
+
 static void process(cir::FuncOp func) {
   if (func.getRegion().empty())
     return;
@@ -49,16 +70,27 @@ static void process(cir::FuncOp func) {
   mlir::Operation *insertPoint = &*entryBlock.begin();
 
   for (auto alloca : allocas) {
-    alloca->moveBefore(insertPoint);
     if (alloca.getConstant()) {
-      // Hoisted alloca may come from the body of a loop, in which case the
-      // stack slot is re-used by multiple objects alive in different iterations
-      // of the loop. In theory, each of these objects are still constant within
-      // their lifetimes, but currently we're not emitting metadata to further
-      // describe this. So for now let's behave conservatively and remove the
-      // const flag on nested allocas when hoisting them.
-      alloca.setConstant(false);
+      if (isOpInLoop(alloca)) {
+        mlir::OpBuilder builder(alloca);
+        auto invariantGroupOp =
+            builder.create<cir::InvariantGroupOp>(alloca.getLoc(), alloca);
+        alloca->replaceUsesWithIf(
+            invariantGroupOp,
+            [op = invariantGroupOp.getOperation()](mlir::OpOperand &use) {
+              return use.getOwner() != op;
+            });
+      } else if (isWhileCondition(alloca)) {
+        // The alloca represents a variable declared as the condition of a while
+        // loop. In CIR, the alloca would be emitted at a scope outside of the
+        // while loop. We have to remove the constant flag during hoisting,
+        // otherwise we would be telling the optimizer that the alloca-ed value
+        // is constant across all iterations of the while loop.
+        alloca.setConstant(false);
+      }
     }
+
+    alloca->moveBefore(insertPoint);
   }
 }
 
diff --git a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp
@@ -1649,6 +1649,15 @@ getLLVMMemOrder(std::optional<cir::MemOrder> &memorder) {
   llvm_unreachable("unknown memory order");
 }
 
+static bool isLoadOrStoreInvariant(mlir::Value addr) {
+  if (auto addrAllocaOp =
+          mlir::dyn_cast_if_present<cir::AllocaOp>(addr.getDefiningOp()))
+    return addrAllocaOp.getConstant();
+  if (mlir::isa_and_present<cir::InvariantGroupOp>(addr.getDefiningOp()))
+    return true;
+  return false;
+}
+
 mlir::LogicalResult CIRToLLVMLoadOpLowering::matchAndRewrite(
     cir::LoadOp op, OpAdaptor adaptor,
     mlir::ConversionPatternRewriter &rewriter) const {
@@ -1668,12 +1677,8 @@ mlir::LogicalResult CIRToLLVMLoadOpLowering::matchAndRewrite(
   auto invariant = false;
   // Under -O1 or higher optimization levels, add the invariant metadata if the
   // load operation loads from a constant object.
-  if (lowerMod &&
-      lowerMod->getContext().getCodeGenOpts().OptimizationLevel > 0) {
-    auto addrAllocaOp =
-        mlir::dyn_cast_if_present<cir::AllocaOp>(op.getAddr().getDefiningOp());
-    invariant = addrAllocaOp && addrAllocaOp.getConstant();
-  }
+  if (lowerMod && lowerMod->getContext().getCodeGenOpts().OptimizationLevel > 0)
+    invariant = isLoadOrStoreInvariant(op.getAddr());
 
   // TODO: nontemporal, syncscope.
   auto newLoad = rewriter.create<mlir::LLVM::LoadOp>(
@@ -1708,12 +1713,8 @@ mlir::LogicalResult CIRToLLVMStoreOpLowering::matchAndRewrite(
   auto invariant = false;
   // Under -O1 or higher optimization levels, add the invariant metadata if the
   // store operation stores to a constant object.
-  if (lowerMod &&
-      lowerMod->getContext().getCodeGenOpts().OptimizationLevel > 0) {
-    auto addrAllocaOp =
-        mlir::dyn_cast_if_present<cir::AllocaOp>(op.getAddr().getDefiningOp());
-    invariant = addrAllocaOp && addrAllocaOp.getConstant();
-  }
+  if (lowerMod && lowerMod->getContext().getCodeGenOpts().OptimizationLevel > 0)
+    invariant = isLoadOrStoreInvariant(op.getAddr());
 
   // Convert adapted value to its memory type if needed.
   mlir::Value value = emitToMemory(rewriter, dataLayout,
@@ -3700,6 +3701,20 @@ mlir::LogicalResult CIRToLLVMInlineAsmOpLowering::matchAndRewrite(
   return mlir::success();
 }
 
+mlir::LogicalResult CIRToLLVMInvariantGroupOpLowering::matchAndRewrite(
+    cir::InvariantGroupOp op, OpAdaptor adaptor,
+    mlir::ConversionPatternRewriter &rewriter) const {
+  if (!lowerMod ||
+      lowerMod->getContext().getCodeGenOpts().OptimizationLevel == 0) {
+    rewriter.replaceOp(op, adaptor.getPtr());
+    return mlir::success();
+  }
+
+  rewriter.replaceOpWithNewOp<mlir::LLVM::LaunderInvariantGroupOp>(
+      op, adaptor.getPtr());
+  return mlir::success();
+}
+
 mlir::LogicalResult CIRToLLVMPrefetchOpLowering::matchAndRewrite(
     cir::PrefetchOp op, OpAdaptor adaptor,
     mlir::ConversionPatternRewriter &rewriter) const {
@@ -4143,7 +4158,8 @@ void populateCIRToLLVMConversionPatterns(
       CIRToLLVMBaseDataMemberOpLowering,
       CIRToLLVMCmpOpLowering,
       CIRToLLVMDerivedDataMemberOpLowering,
-      CIRToLLVMGetRuntimeMemberOpLowering
+      CIRToLLVMGetRuntimeMemberOpLowering,
+      CIRToLLVMInvariantGroupOpLowering
       // clang-format on
       >(converter, patterns.getContext(), lowerModule);
   patterns.add<
diff --git a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.h b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.h
@@ -962,6 +962,21 @@ class CIRToLLVMInlineAsmOpLowering
                   mlir::ConversionPatternRewriter &) const override;
 };
 
+class CIRToLLVMInvariantGroupOpLowering
+    : public mlir::OpConversionPattern<cir::InvariantGroupOp> {
+  cir::LowerModule *lowerMod;
+
+public:
+  CIRToLLVMInvariantGroupOpLowering(const mlir::TypeConverter &typeConverter,
+                                    mlir::MLIRContext *context,
+                                    cir::LowerModule *lowerModule)
+      : OpConversionPattern(typeConverter, context), lowerMod(lowerModule) {}
+
+  mlir::LogicalResult
+  matchAndRewrite(cir::InvariantGroupOp op, OpAdaptor,
+                  mlir::ConversionPatternRewriter &) const override;
+};
+
 class CIRToLLVMPrefetchOpLowering
     : public mlir::OpConversionPattern<cir::PrefetchOp> {
 public:
diff --git a/clang/test/CIR/CodeGen/const-alloca.cpp b/clang/test/CIR/CodeGen/const-alloca.cpp
@@ -5,6 +5,7 @@
 
 int produce_int();
 void blackbox(const int &);
+void consume(int);
 
 void local_const_int() {
   const int x = produce_int();
@@ -85,3 +86,87 @@ int local_const_optimize() {
 // LLVM-NEXT:    call void @_Z8blackboxRKi(ptr nonnull %[[#slot]])
 // LLVM-NEXT:    ret i32 %[[#init]]
 // LLVM-NEXT:  }
+
+int local_scoped_const() {
+  {
+    const int x = produce_int();
+    blackbox(x);
+    return x;
+  }
+}
+
+// CIR-LABEL: @_Z18local_scoped_constv()
+//      CIR:    cir.scope {
+// CIR-NEXT:      %[[#x_slot:]] = cir.alloca !s32i, !cir.ptr<!s32i>, ["x", init, const]
+// CIR-NEXT:      %[[#init:]] = cir.call @_Z11produce_intv() : () -> !s32i
+// CIR-NEXT:      cir.store %[[#init]], %[[#x_slot]] : !s32i, !cir.ptr<!s32i> tbaa([#tbaa])
+// CIR-NEXT:      cir.call @_Z8blackboxRKi(%[[#x_slot]]) : (!cir.ptr<!s32i>) -> ()
+// CIR-NEXT:      %[[#x_reload:]] = cir.load %[[#x_slot]] : !cir.ptr<!s32i>, !s32i tbaa([#tbaa])
+// CIR-NEXT:      cir.store %[[#x_reload]], %[[#ret_slot:]] : !s32i, !cir.ptr<!s32i>
+// CIR-NEXT:      %[[#ret:]] = cir.load %[[#ret_slot]] : !cir.ptr<!s32i>, !s32i
+// CIR-NEXT:      cir.return %[[#ret]] : !s32i
+// CIR-NEXT:    }
+//      CIR:  }
+
+// LLVM-LABEL: @_Z18local_scoped_constv()
+// LLVM-NEXT:    %[[#x_slot:]] = alloca i32, align 4
+// LLVM-NEXT:    %[[#init:]] = tail call i32 @_Z11produce_intv()
+// LLVM-NEXT:    store i32 %[[#init]], ptr %[[#x_slot]], align 4, !invariant.group !{{.+}}
+// LLVM-NEXT:    call void @_Z8blackboxRKi(ptr nonnull %[[#x_slot]])
+// LLVM-NEXT:    ret i32 %[[#init]]
+// LLVM-NEXT:  }
+
+void local_const_in_loop() {
+  for (int i = 0; i < 10; ++i) {
+    const int x = produce_int();
+    blackbox(x);
+    consume(x);
+  }
+}
+
+// CIR-LABEL: @_Z19local_const_in_loopv
+//      CIR:    cir.scope {
+//      CIR:      cir.for : cond {
+//      CIR:      } body {
+// CIR-NEXT:        cir.scope {
+// CIR-NEXT:          %[[#x_slot:]] = cir.alloca !s32i, !cir.ptr<!s32i>, ["x", init, const]
+// CIR-NEXT:          %[[#init:]] = cir.call @_Z11produce_intv() : () -> !s32i
+// CIR-NEXT:          cir.store %[[#init]], %[[#x_slot]] : !s32i, !cir.ptr<!s32i> tbaa([#tbaa])
+// CIR-NEXT:          cir.call @_Z8blackboxRKi(%[[#x_slot]]) : (!cir.ptr<!s32i>) -> ()
+// CIR-NEXT:          %[[#x_reload:]] = cir.load %[[#x_slot]] : !cir.ptr<!s32i>, !s32i tbaa([#tbaa])
+// CIR-NEXT:          cir.call @_Z7consumei(%[[#x_reload]]) : (!s32i) -> ()
+// CIR-NEXT:        }
+// CIR-NEXT:        cir.yield
+// CIR-NEXT:      } step {
+//      CIR:      }
+// CIR-NEXT:    }
+// CIR-NEXT:    cir.return
+// CIR-NEXT:  }
+
+// LLVM-LABEL: @_Z19local_const_in_loopv()
+//      LLVM:    %[[#x_ptr:]] = call ptr @llvm.launder.invariant.group.p0(ptr nonnull %1)
+// LLVM-NEXT:    %[[#init:]] = call i32 @_Z11produce_intv()
+// LLVM-NEXT:    store i32 %[[#init]], ptr %[[#x_ptr]], align 4, !invariant.group !{{.+}}
+// LLVM-NEXT:    call void @_Z8blackboxRKi(ptr nonnull %[[#x_ptr]])
+// LLVM-NEXT:    call void @_Z7consumei(i32 %[[#init]])
+//      LLVM:  }
+
+void local_const_in_while_condition() {
+  while (const int x = produce_int()) {
+    blackbox(x);
+  }
+}
+
+// LLVM-LABEL: @_Z30local_const_in_while_conditionv()
+//      LLVM:    %[[#x_slot:]] = alloca i32, align 4
+// LLVM-NEXT:    %[[#init:]] = tail call i32 @_Z11produce_intv()
+// LLVM-NEXT:    store i32 %[[#init]], ptr %[[#x_slot]], align 4
+// LLVM-NEXT:    %[[loop_cond:.+]] = icmp eq i32 %[[#init]], 0
+// LLVM-NEXT:    br i1 %[[loop_cond]], label %{{.+}}, label %[[loop_body:.+]]
+//      LLVM:  [[loop_body]]:
+// LLVM-NEXT:    call void @_Z8blackboxRKi(ptr nonnull %[[#x_slot]])
+// LLVM-NEXT:    %[[#next:]] = call i32 @_Z11produce_intv()
+// LLVM-NEXT:    store i32 %[[#next]], ptr %[[#x_slot]], align 4
+// LLVM-NEXT:    %[[cond:.+]] = icmp eq i32 %[[#next]], 0
+// LLVM-NEXT:    br i1 %[[cond]], label %{{.+}}, label %[[loop_body]]
+//      LLVM:  }
