Skip to content

Commit 59e7eb9

Browse files
Lancernlanza
Lancern
authored and
lanza
committed
[CIR] Lower nested local constant alloca (#1261)
This PR adds support for lowering local constants in nested scopes, including those in nested loops. For those constant allocas in non-loop inner scopes, this PR 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 PR 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 PR updates LLVM lowering to emit the necessary invariant metadata when loading and storing through such pointers. This PR 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 PR removes the constant flag on the alloca operation when hositing them.
1 parent b9f0fb5 commit 59e7eb9

File tree

5 files changed

+341
-23
lines changed

5 files changed

+341
-23
lines changed

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

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3563,6 +3563,63 @@ def LLVMIntrinsicCallOp : CIR_Op<"llvm.intrinsic"> {
35633563

35643564
}
35653565

3566+
//===----------------------------------------------------------------------===//
3567+
// InvariantGroupOp
3568+
//===----------------------------------------------------------------------===//
3569+
3570+
def InvariantGroupOp
3571+
: CIR_Op<"invariant_group", [Pure, SameOperandsAndResultType]> {
3572+
let summary = "Start an invariant group";
3573+
let description = [{
3574+
The `cir.invariant_group` operation takes a single pointer value as argument
3575+
and returns the same pointer value with fresh [invariant group] information.
3576+
All loads and stores that access the returned pointer value are presumed by
3577+
the optimizer to load or store the same value.
3578+
3579+
[invariant group]: https://llvm.org/docs/LangRef.html#invariant-group-metadata
3580+
3581+
This operation is not emitted during CIRGen. Instead, it is created when
3582+
hoisting constant alloca operations to the entry block of a function. This
3583+
operation effectively marks the syntactic scope of the constant local
3584+
variable represented by the hosited alloca operation, and it allows for
3585+
better LLVMIR generation with potentially more optimizations.
3586+
3587+
For example, if we have the following CIR before alloca hoisting:
3588+
3589+
```mlir
3590+
cir.func @foo() {
3591+
cir.scope {
3592+
%0 = cir.alloca !s32i : !cir.ptr<!s32i>
3593+
use(%0)
3594+
}
3595+
}
3596+
```
3597+
3598+
After alloca hoisting:
3599+
3600+
```mlir
3601+
cir.func @foo() {
3602+
%0 = cir.alloca !s32i : !cir.ptr<!s32i>
3603+
cir.scope {
3604+
%1 = cir.invariant_group %0 : !cir.ptr<!s32i>
3605+
use(%1)
3606+
}
3607+
}
3608+
```
3609+
3610+
During LLVMIR lowering, load and store operations whose pointer operand
3611+
comes from `cir.invariant_group` are lowered to corresponding LLVM
3612+
instructions with invariant group metadata attached.
3613+
}];
3614+
3615+
let arguments = (ins CIR_PointerType:$ptr);
3616+
let results = (outs CIR_PointerType:$result);
3617+
3618+
let assemblyFormat = [{
3619+
$ptr `:` type($result) attr-dict
3620+
}];
3621+
}
3622+
35663623
//===----------------------------------------------------------------------===//
35673624
// DeleteArrayOp
35683625
//===----------------------------------------------------------------------===//

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

Lines changed: 155 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,142 @@ struct HoistAllocasPass : public HoistAllocasBase<HoistAllocasPass> {
2828
void runOnOperation() override;
2929
};
3030

31-
static void process(cir::FuncOp func) {
31+
static bool isOpInLoop(mlir::Operation *op) {
32+
return op->getParentOfType<cir::LoopOpInterface>();
33+
}
34+
35+
static bool hasStoreToAllocaInWhileCond(cir::AllocaOp alloca) {
36+
// This function determines whether the given alloca operation represents
37+
// a variable defined as a while loop's condition.
38+
//
39+
// Specifically, C/C++ allows the condition of a while loop be a variable
40+
// declaration:
41+
//
42+
// while (const int x = foo()) { /* body... */ }
43+
//
44+
// CIRGen would emit the following CIR for the above code:
45+
//
46+
// cir.scope {
47+
// %x.slot = cir.alloca !s32i [init, const]
48+
// cir.while {
49+
// %0 = cir.call @foo()
50+
// cir.store %0, %x
51+
// %1 = cir.load %x
52+
// %2 = cir.cast int_to_bool %1
53+
// cir.condition(%2)
54+
// } do {
55+
// // loop body goes here.
56+
// }
57+
// }
58+
//
59+
// Note that %x.slot is emitted outside the cir.while operation. Ideally, the
60+
// cir.while operation should cover this cir.alloca operation, but currently
61+
// CIR does not work this way. When hoisting such an alloca operation, one
62+
// must remove the "const" flag from it, otherwise LLVM lowering code will
63+
// mistakenly attach invariant group metadata to the load and store operations
64+
// in the while body, indicating that all loads and stores across all
65+
// iterations of the loop are constant.
66+
67+
for (mlir::Operation *user : alloca->getUsers()) {
68+
if (!mlir::isa<cir::StoreOp>(user))
69+
continue;
70+
71+
auto store = mlir::cast<cir::StoreOp>(user);
72+
mlir::Operation *storeParentOp = store->getParentOp();
73+
if (!mlir::isa<cir::WhileOp>(storeParentOp))
74+
continue;
75+
76+
auto whileOp = mlir::cast<cir::WhileOp>(storeParentOp);
77+
return &whileOp.getCond() == store->getParentRegion();
78+
}
79+
80+
return false;
81+
}
82+
83+
static void processConstAlloca(cir::AllocaOp alloca) {
84+
// When optimization is enabled, LLVM lowering would start emitting invariant
85+
// group metadata for loads and stores to alloca-ed objects with "const"
86+
// attribute. For example, the following CIR:
87+
//
88+
// %slot = cir.alloca !s32i [init, const]
89+
// cir.store %0, %slot
90+
// %1 = cir.load %slot
91+
//
92+
// would be lowered to the following LLVM IR:
93+
//
94+
// %slot = alloca i32, i64 1
95+
// store i32 %0, ptr %slot, !invariant.group !0
96+
// %1 = load i32, ptr %slot, !invariant.group !0
97+
//
98+
// The invariant group metadata would tell LLVM optimizer that the store and
99+
// load instruction would store and load the same value from %slot.
100+
//
101+
// So far so good. Things started to get tricky when such an alloca operation
102+
// appears in the body of a loop construct:
103+
//
104+
// cir.some_loop_construct {
105+
// %slot = cir.alloca !s32i [init, const]
106+
// cir.store %0, %slot
107+
// %1 = cir.load %slot
108+
// }
109+
//
110+
// After alloca hoisting, the CIR code above would be transformed into:
111+
//
112+
// %slot = cir.alloca !s32i [init, const]
113+
// cir.some_loop_construct {
114+
// cir.store %0, %slot
115+
// %1 = cir.load %slot
116+
// }
117+
//
118+
// Notice how alloca hoisting change the semantics of the program in such a
119+
// case. The transformed code now indicates the optimizer that the load and
120+
// store operations load and store the same value **across all iterations of
121+
// the loop**!
122+
//
123+
// To overcome this problem, we instead transform the program into this:
124+
//
125+
// %slot = cir.alloca !s32i [init, const]
126+
// cir.some_loop_construct {
127+
// %slot.inv = cir.invariant_group %slot
128+
// cir.store %0, %slot.inv
129+
// %1 = cir.load %slot.inv
130+
// }
131+
//
132+
// The cir.invariant_group operation attaches fresh invariant information to
133+
// the operand pointer and yields a pointer with the fresh invariant
134+
// information. Upon each loop iteration, the old invariant information is
135+
// disgarded, and a new invariant information is attached, thus the correct
136+
// program semantic retains. During LLVM lowering, the cir.invariant_group
137+
// operation would eventually become an intrinsic call to
138+
// @llvm.launder.invariant.group.
139+
140+
if (isOpInLoop(alloca)) {
141+
// Mark the alloca-ed pointer as invariant via the cir.invariant_group
142+
// operation.
143+
mlir::OpBuilder builder(alloca);
144+
auto invariantGroupOp =
145+
builder.create<cir::InvariantGroupOp>(alloca.getLoc(), alloca);
146+
147+
// And replace all uses of the original alloca-ed pointer with the marked
148+
// pointer (which carries invariant group information).
149+
alloca->replaceUsesWithIf(
150+
invariantGroupOp,
151+
[op = invariantGroupOp.getOperation()](mlir::OpOperand &use) {
152+
return use.getOwner() != op;
153+
});
154+
} else if (hasStoreToAllocaInWhileCond(alloca)) {
155+
// The alloca represents a variable declared as the condition of a while
156+
// loop. In CIR, the alloca would be emitted at a scope outside of the
157+
// while loop. We have to remove the constant flag during hoisting,
158+
// otherwise we would be telling the optimizer that the alloca-ed value
159+
// is constant across all iterations of the while loop.
160+
//
161+
// See the body of the isWhileCondition function for more details.
162+
alloca.setConstant(false);
163+
}
164+
}
165+
166+
static void process(mlir::ModuleOp mod, cir::FuncOp func) {
32167
if (func.getRegion().empty())
33168
return;
34169

@@ -47,25 +182,35 @@ static void process(cir::FuncOp func) {
47182
return;
48183

49184
mlir::Operation *insertPoint = &*entryBlock.begin();
185+
auto optInfoAttr = mlir::cast_if_present<cir::OptInfoAttr>(
186+
mod->getAttr(cir::CIRDialect::getOptInfoAttrName()));
187+
unsigned optLevel = optInfoAttr ? optInfoAttr.getLevel() : 0;
50188

51189
for (auto alloca : allocas) {
52-
alloca->moveBefore(insertPoint);
53190
if (alloca.getConstant()) {
54-
// Hoisted alloca may come from the body of a loop, in which case the
55-
// stack slot is re-used by multiple objects alive in different iterations
56-
// of the loop. In theory, each of these objects are still constant within
57-
// their lifetimes, but currently we're not emitting metadata to further
58-
// describe this. So for now let's behave conservatively and remove the
59-
// const flag on nested allocas when hoisting them.
60-
alloca.setConstant(false);
191+
if (optLevel == 0) {
192+
// Under non-optimized builds, just remove the constant flag.
193+
alloca.setConstant(false);
194+
continue;
195+
}
196+
197+
processConstAlloca(alloca);
61198
}
199+
200+
alloca->moveBefore(insertPoint);
62201
}
63202
}
64203

65204
void HoistAllocasPass::runOnOperation() {
66205
llvm::TimeTraceScope scope("Hoist Allocas");
67206
llvm::SmallVector<Operation *, 16> ops;
68-
getOperation()->walk([&](cir::FuncOp op) { process(op); });
207+
208+
Operation *op = getOperation();
209+
auto mod = mlir::dyn_cast<mlir::ModuleOp>(op);
210+
if (!mod)
211+
mod = op->getParentOfType<mlir::ModuleOp>();
212+
213+
getOperation()->walk([&](cir::FuncOp op) { process(mod, op); });
69214
}
70215

71216
} // namespace

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

Lines changed: 29 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1612,6 +1612,15 @@ getLLVMMemOrder(std::optional<cir::MemOrder> &memorder) {
16121612
llvm_unreachable("unknown memory order");
16131613
}
16141614

1615+
static bool isLoadOrStoreInvariant(mlir::Value addr) {
1616+
if (auto addrAllocaOp =
1617+
mlir::dyn_cast_if_present<cir::AllocaOp>(addr.getDefiningOp()))
1618+
return addrAllocaOp.getConstant();
1619+
if (mlir::isa_and_present<cir::InvariantGroupOp>(addr.getDefiningOp()))
1620+
return true;
1621+
return false;
1622+
}
1623+
16151624
mlir::LogicalResult CIRToLLVMLoadOpLowering::matchAndRewrite(
16161625
cir::LoadOp op, OpAdaptor adaptor,
16171626
mlir::ConversionPatternRewriter &rewriter) const {
@@ -1631,12 +1640,8 @@ mlir::LogicalResult CIRToLLVMLoadOpLowering::matchAndRewrite(
16311640
auto invariant = false;
16321641
// Under -O1 or higher optimization levels, add the invariant metadata if the
16331642
// load operation loads from a constant object.
1634-
if (lowerMod &&
1635-
lowerMod->getContext().getCodeGenOpts().OptimizationLevel > 0) {
1636-
auto addrAllocaOp =
1637-
mlir::dyn_cast_if_present<cir::AllocaOp>(op.getAddr().getDefiningOp());
1638-
invariant = addrAllocaOp && addrAllocaOp.getConstant();
1639-
}
1643+
if (lowerMod && lowerMod->getContext().getCodeGenOpts().OptimizationLevel > 0)
1644+
invariant = isLoadOrStoreInvariant(op.getAddr());
16401645

16411646
// TODO: nontemporal, syncscope.
16421647
auto newLoad = rewriter.create<mlir::LLVM::LoadOp>(
@@ -1674,12 +1679,8 @@ mlir::LogicalResult CIRToLLVMStoreOpLowering::matchAndRewrite(
16741679
auto invariant = false;
16751680
// Under -O1 or higher optimization levels, add the invariant metadata if the
16761681
// store operation stores to a constant object.
1677-
if (lowerMod &&
1678-
lowerMod->getContext().getCodeGenOpts().OptimizationLevel > 0) {
1679-
auto addrAllocaOp =
1680-
mlir::dyn_cast_if_present<cir::AllocaOp>(op.getAddr().getDefiningOp());
1681-
invariant = addrAllocaOp && addrAllocaOp.getConstant();
1682-
}
1682+
if (lowerMod && lowerMod->getContext().getCodeGenOpts().OptimizationLevel > 0)
1683+
invariant = isLoadOrStoreInvariant(op.getAddr());
16831684

16841685
// Convert adapted value to its memory type if needed.
16851686
mlir::Value value = emitToMemory(rewriter, dataLayout,
@@ -3666,6 +3667,20 @@ mlir::LogicalResult CIRToLLVMInlineAsmOpLowering::matchAndRewrite(
36663667
return mlir::success();
36673668
}
36683669

3670+
mlir::LogicalResult CIRToLLVMInvariantGroupOpLowering::matchAndRewrite(
3671+
cir::InvariantGroupOp op, OpAdaptor adaptor,
3672+
mlir::ConversionPatternRewriter &rewriter) const {
3673+
if (!lowerMod ||
3674+
lowerMod->getContext().getCodeGenOpts().OptimizationLevel == 0) {
3675+
rewriter.replaceOp(op, adaptor.getPtr());
3676+
return mlir::success();
3677+
}
3678+
3679+
rewriter.replaceOpWithNewOp<mlir::LLVM::LaunderInvariantGroupOp>(
3680+
op, adaptor.getPtr());
3681+
return mlir::success();
3682+
}
3683+
36693684
mlir::LogicalResult CIRToLLVMPrefetchOpLowering::matchAndRewrite(
36703685
cir::PrefetchOp op, OpAdaptor adaptor,
36713686
mlir::ConversionPatternRewriter &rewriter) const {
@@ -4107,7 +4122,8 @@ void populateCIRToLLVMConversionPatterns(
41074122
CIRToLLVMBaseDataMemberOpLowering,
41084123
CIRToLLVMCmpOpLowering,
41094124
CIRToLLVMDerivedDataMemberOpLowering,
4110-
CIRToLLVMGetRuntimeMemberOpLowering
4125+
CIRToLLVMGetRuntimeMemberOpLowering,
4126+
CIRToLLVMInvariantGroupOpLowering
41114127
// clang-format on
41124128
>(converter, patterns.getContext(), lowerModule);
41134129
patterns.add<

‎clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.h‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -975,6 +975,21 @@ class CIRToLLVMInlineAsmOpLowering
975975
mlir::ConversionPatternRewriter &) const override;
976976
};
977977

978+
class CIRToLLVMInvariantGroupOpLowering
979+
: public mlir::OpConversionPattern<cir::InvariantGroupOp> {
980+
cir::LowerModule *lowerMod;
981+
982+
public:
983+
CIRToLLVMInvariantGroupOpLowering(const mlir::TypeConverter &typeConverter,
984+
mlir::MLIRContext *context,
985+
cir::LowerModule *lowerModule)
986+
: OpConversionPattern(typeConverter, context), lowerMod(lowerModule) {}
987+
988+
mlir::LogicalResult
989+
matchAndRewrite(cir::InvariantGroupOp op, OpAdaptor,
990+
mlir::ConversionPatternRewriter &) const override;
991+
};
992+
978993
class CIRToLLVMPrefetchOpLowering
979994
: public mlir::OpConversionPattern<cir::PrefetchOp> {
980995
public:

0 commit comments

Comments
 (0)