Skip to content

Commit 1c09e5a

Browse files
maerhartmaerhart
committed
[RTG] Add simple linear scan register allocation pass
1 parent 44b5801 commit 1c09e5a

File tree

4 files changed

+248
-0
lines changed

4 files changed

+248
-0
lines changed

‎include/circt/Dialect/RTG/Transforms/RTGPasses.td

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,4 +64,21 @@ def EmitRTGISAAssemblyPass : Pass<"rtg-emit-isa-assembly", "mlir::ModuleOp"> {
6464
];
6565
}
6666

67+
def LinearScanRegisterAllocationPass : Pass<
68+
"rtg-linear-scan-register-allocation", "rtg::TestOp"> {
69+
70+
let summary = "simple linear scan register allocation for RTG";
71+
let description = [{
72+
Performs a simple version of the linear scan register allocation algorithm
73+
based on the 'rtg.virtual_reg' operations.
74+
75+
This pass is expected to be run after elaboration.
76+
}];
77+
78+
let statistics = [
79+
Statistic<"numRegistersSpilled", "num-registers-spilled",
80+
"Number of registers spilled to the stack.">,
81+
];
82+
}
83+
6784
#endif // CIRCT_DIALECT_RTG_TRANSFORMS_RTGPASSES_TD

‎lib/Dialect/RTG/Transforms/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
add_circt_dialect_library(CIRCTRTGTransforms
22
ElaborationPass.cpp
33
EmitRTGISAAssemblyPass.cpp
4+
LinearScanRegisterAllocationPass.cpp
45

56
DEPENDS
67
CIRCTRTGTransformsIncGen

‎lib/Dialect/RTG/Transforms/LinearScanRegisterAllocationPass.cpp

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
//===- LinearScanRegisterAllocationPass.cpp - Register Allocation ---------===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This pass allocates registers using a simple linear scan algorithm.
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
#include "circt/Dialect/RTG/IR/RTGISAAssemblyOpInterfaces.h"
14+
#include "circt/Dialect/RTG/IR/RTGOps.h"
15+
#include "circt/Dialect/RTG/Transforms/RTGPasses.h"
16+
#include "mlir/IR/PatternMatch.h"
17+
#include "llvm/Support/Debug.h"
18+
19+
namespace circt {
20+
namespace rtg {
21+
#define GEN_PASS_DEF_LINEARSCANREGISTERALLOCATIONPASS
22+
#include "circt/Dialect/RTG/Transforms/RTGPasses.h.inc"
23+
} // namespace rtg
24+
} // namespace circt
25+
26+
using namespace mlir;
27+
using namespace circt;
28+
29+
#define DEBUG_TYPE "rtg-linear-scan-register-allocation"
30+
31+
namespace {
32+
33+
/// Represents a register and its live range.
34+
struct RegisterLiveRange {
35+
rtg::RegisterAttrInterface fixedReg;
36+
rtg::VirtualRegisterOp regOp;
37+
unsigned start;
38+
unsigned end;
39+
};
40+
41+
class LinearScanRegisterAllocationPass
42+
: public circt::rtg::impl::LinearScanRegisterAllocationPassBase<
43+
LinearScanRegisterAllocationPass> {
44+
public:
45+
void runOnOperation() override;
46+
};
47+
48+
} // end namespace
49+
50+
static void expireOldInterval(SmallVector<RegisterLiveRange *> &active,
51+
RegisterLiveRange *reg) {
52+
// TODO: use a better datastructure for 'active'
53+
llvm::sort(active, [](auto *a, auto *b) { return a->end < b->end; });
54+
55+
for (auto *iter = active.begin(); iter != active.end(); ++iter) {
56+
auto *a = *iter;
57+
if (a->end >= reg->start)
58+
return;
59+
60+
active.erase(iter--);
61+
}
62+
}
63+
64+
void LinearScanRegisterAllocationPass::runOnOperation() {
65+
auto testOp = getOperation();
66+
67+
LLVM_DEBUG(llvm::dbgs() << "=== Processing test @" << testOp.getSymName()
68+
<< "\n\n");
69+
70+
DenseMap<Operation *, unsigned> opIndices;
71+
unsigned maxIdx;
72+
for (auto [i, op] : llvm::enumerate(*testOp.getBody())) {
73+
// TODO: ideally check that the IR is already fully elaborated
74+
opIndices[&op] = i;
75+
maxIdx = i;
76+
}
77+
78+
// Collect all the register intervals we have to consider.
79+
SmallVector<std::unique_ptr<RegisterLiveRange>> regRanges;
80+
SmallVector<RegisterLiveRange *> active;
81+
for (auto &op : *testOp.getBody()) {
82+
if (!isa<rtg::FixedRegisterOp, rtg::VirtualRegisterOp>(&op))
83+
continue;
84+
85+
RegisterLiveRange lr;
86+
lr.start = maxIdx;
87+
lr.end = 0;
88+
89+
if (auto regOp = dyn_cast<rtg::VirtualRegisterOp>(&op))
90+
lr.regOp = regOp;
91+
92+
if (auto regOp = dyn_cast<rtg::FixedRegisterOp>(&op))
93+
lr.fixedReg = regOp.getReg();
94+
95+
for (auto *user : op.getUsers()) {
96+
if (!isa<rtg::InstructionOpInterface>(user)) {
97+
user->emitError("only operations implementing 'InstructionOpInterface "
98+
"are allowed to use registers");
99+
return signalPassFailure();
100+
}
101+
102+
// TODO: support labels and control-flow loops (jumps in general)
103+
unsigned idx = opIndices.at(user);
104+
lr.start = std::min(lr.start, idx);
105+
lr.end = std::max(lr.end, idx);
106+
}
107+
108+
regRanges.emplace_back(std::make_unique<RegisterLiveRange>(lr));
109+
110+
// Reserve fixed registers from the start. It will be made available again
111+
// past the interval end. Not reserving it from the start can lead to the
112+
// same register being chosen for a virtual register that overlaps with the
113+
// fixed register interval.
114+
// TODO: don't overapproximate that much
115+
if (!lr.regOp)
116+
active.push_back(regRanges.back().get());
117+
}
118+
119+
// Sort such that we can process registers by increasing interval start.
120+
llvm::sort(regRanges, [](const auto &a, const auto &b) {
121+
return a->start < b->start || (a->start == b->start && !a->regOp);
122+
});
123+
124+
for (auto &lr : regRanges) {
125+
// Make registers out of live range available again.
126+
expireOldInterval(active, lr.get());
127+
128+
// Handle already fixed registers.
129+
if (!lr->regOp)
130+
continue;
131+
132+
// Handle virtual registers.
133+
rtg::RegisterAttrInterface availableReg;
134+
for (auto reg : lr->regOp.getAllowedRegs()) {
135+
if (llvm::none_of(active, [&](auto *r) { return r->fixedReg == reg; })) {
136+
availableReg = cast<rtg::RegisterAttrInterface>(reg);
137+
break;
138+
}
139+
}
140+
141+
if (!availableReg) {
142+
++numRegistersSpilled;
143+
lr->regOp->emitError(
144+
"need to spill this register, but not supported yet");
145+
return signalPassFailure();
146+
}
147+
148+
lr->fixedReg = availableReg;
149+
active.push_back(lr.get());
150+
}
151+
152+
LLVM_DEBUG({
153+
for (auto &regRange : regRanges) {
154+
llvm::dbgs() << "Start: " << regRange->start << ", End: " << regRange->end
155+
<< ", Selected: " << regRange->fixedReg << "\n";
156+
}
157+
llvm::dbgs() << "\n";
158+
});
159+
160+
for (auto &reg : regRanges) {
161+
// No need to fix already fixed registers.
162+
if (!reg->regOp)
163+
continue;
164+
165+
IRRewriter rewriter(reg->regOp);
166+
rewriter.replaceOpWithNewOp<rtg::FixedRegisterOp>(reg->regOp,
167+
reg->fixedReg);
168+
}
169+
}

‎test/Dialect/RTG/Transform/linear-scan-register-allocation.mlir

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
// RUN: circt-opt --rtg-linear-scan-register-allocation --split-input-file --verify-diagnostics %s | FileCheck %s
2+
3+
// CHECK-LABEL: @test0
4+
rtg.test @test0 : !rtg.dict<> {
5+
// CHECK: [[V0:%.+]] = rtg.fixed_reg #rtgtest.ra
6+
// CHECK: [[V1:%.+]] = rtg.fixed_reg #rtgtest.s1
7+
// CHECK: [[V2:%.+]] = rtg.fixed_reg #rtgtest.s0
8+
// CHECK: [[V3:%.+]] = rtg.fixed_reg #rtgtest.ra
9+
// CHECK: rtgtest.jalr [[V0]], [[V2]]
10+
// CHECK: rtgtest.jalr [[V1]], [[V0]]
11+
// CHECK: rtgtest.jalr [[V3]], [[V1]]
12+
// CHECK: rtgtest.jalr [[V2]], [[V3]]
13+
%0 = rtg.virtual_reg [#rtgtest.ra, #rtgtest.s0, #rtgtest.s1]
14+
%1 = rtg.virtual_reg [#rtgtest.ra, #rtgtest.s0, #rtgtest.s1]
15+
%2 = rtg.virtual_reg [#rtgtest.ra, #rtgtest.s0, #rtgtest.s1]
16+
%3 = rtg.virtual_reg [#rtgtest.ra, #rtgtest.s0, #rtgtest.s1]
17+
%imm = rtgtest.immediate #rtgtest.imm12<0>
18+
rtgtest.jalr %0, %2, %imm
19+
rtgtest.jalr %1, %0, %imm
20+
rtgtest.jalr %3, %1, %imm
21+
rtgtest.jalr %2, %3, %imm
22+
}
23+
24+
// CHECK-LABEL: @withFixedRegs
25+
rtg.test @withFixedRegs : !rtg.dict<> {
26+
// CHECK: [[V0:%.+]] = rtg.fixed_reg #rtgtest.ra
27+
// CHECK: [[V1:%.+]] = rtg.fixed_reg #rtgtest.s1
28+
// CHECK: [[V2:%.+]] = rtg.fixed_reg #rtgtest.s0
29+
// CHECK: [[V3:%.+]] = rtg.fixed_reg #rtgtest.ra
30+
// CHECK: rtgtest.jalr [[V0]], [[V2]]
31+
// CHECK: rtgtest.jalr [[V1]], [[V0]]
32+
// CHECK: rtgtest.jalr [[V3]], [[V1]]
33+
// CHECK: rtgtest.jalr [[V2]], [[V3]]
34+
%0 = rtg.fixed_reg #rtgtest.ra
35+
%1 = rtg.virtual_reg [#rtgtest.ra, #rtgtest.s0, #rtgtest.s1]
36+
%2 = rtg.fixed_reg #rtgtest.s0
37+
%3 = rtg.virtual_reg [#rtgtest.ra, #rtgtest.s0, #rtgtest.s1]
38+
%imm = rtgtest.immediate #rtgtest.imm12<0>
39+
rtgtest.jalr %0, %2, %imm
40+
rtgtest.jalr %1, %0, %imm
41+
rtgtest.jalr %3, %1, %imm
42+
rtgtest.jalr %2, %3, %imm
43+
}
44+
45+
// -----
46+
47+
rtg.test @spilling : !rtg.dict<> {
48+
%0 = rtg.virtual_reg [#rtgtest.ra]
49+
// expected-error @below {{need to spill this register, but not supported yet}}
50+
%1 = rtg.virtual_reg [#rtgtest.ra]
51+
%imm = rtgtest.immediate #rtgtest.imm12<0>
52+
rtgtest.jalr %0, %1, %imm
53+
}
54+
55+
// -----
56+
57+
rtg.test @unsupportedUser : !rtg.dict<> {
58+
%0 = rtg.virtual_reg [#rtgtest.ra]
59+
// expected-error @below {{only operations implementing 'InstructionOpInterface are allowed to use registers}}
60+
rtg.set_create %0 : !rtgtest.ireg
61+
}

0 commit comments

Comments
 (0)