[backend] Add register allocation
Bug: 295189858
Test: mm and berberis_host_tests
Change-Id: Iae5d53c45fad76fc0dbd3e93e50fd1de4c2e717d
Merged-In: Iae5d53c45fad76fc0dbd3e93e50fd1de4c2e717d
diff --git a/backend/Android.bp b/backend/Android.bp
index 078a2a6..e39f2c8 100644
--- a/backend/Android.bp
+++ b/backend/Android.bp
@@ -174,6 +174,7 @@
srcs: [
"common/lifetime_analysis.cc",
"common/machine_ir_debug.cc",
+ "common/reg_alloc.cc",
"x86_64/code.cc",
"x86_64/code_debug.cc",
"x86_64/code_emit.cc",
diff --git a/backend/common/reg_alloc.cc b/backend/common/reg_alloc.cc
new file mode 100644
index 0000000..3da5d06
--- /dev/null
+++ b/backend/common/reg_alloc.cc
@@ -0,0 +1,400 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// At the moment, this implements more-or-less traditional linear scan register
+// allocation.
+//
+// Input is virtual register lifetimes list, sorted by begin. Each lifetime is
+// a list of continuous live ranges with lifetime holes in between. Each live
+// range tracks insns that actually use the virtual register.
+//
+// Allocator walks sorted lifetime list and allocates lifetimes to hard
+// registers. When lifetimes do not interfere, so live ranges of one lifetime
+// fit into holes of another lifetime, both lifetimes can be allocated to the
+// same hard register.
+//
+// If there is no available hard register, allocator selects hard register to
+// free. All lifetimes allocated to that hard register that interfere with
+// lifetime being allocated are spilled.
+//
+// Lifetime being spilled is split into tiny lifetimes, each for one insn
+// where that virtual register is used. If register is read by insn, reload
+// insn is added before, and if register is written by insn, spill insn is
+// added after.
+//
+// Tiny lifetimes originated from spilling still need to be allocated. For tiny
+// lifetimes that end before lifetime in favor of which they were spilled,
+// previously allocated hard register is used. Tiny lifetimes that begin after
+// are merged into a list of not yet allocated lifetimes.
+//
+// The problematic case is when tiny lifetime overlaps with begin of lifetime
+// in favor of which it was spilled. In this case previously allocated hard
+// register can't be used (otherwise it doesn't become free) and tiny lifetime
+// can't be allocated later according to order by begin. In this case spill is
+// considered impossible.
+//
+// The above is the most significant difference from classic linear scan
+// register allocation algorithms, which usually solve tiny lifetimes
+// allocation either by backtracking or using reserved registers. Hopefully
+// this new approach works if there are more suitable hard registers than can
+// be used in one insn, so there is always a suitable register that is not used
+// at the point of spill.
+//
+// TODO(b/232598137): this might blow up when lifetimes compete for some single
+// specific register (ecx for x86 shift insn)! Can be solved by generating code
+// that minimizes lifetimes of such registers - moves in right before and moves
+// out right after the insn using such register.
+//
+// TODO(b/232598137): investigate how code quality is affected when there are few
+// available hard registers (x86_32)!
+
+#include "berberis/backend/common/reg_alloc.h"
+
+#include <iterator> // std::next()
+#include <string>
+
+#include "berberis/backend/common/lifetime.h"
+#include "berberis/backend/common/lifetime_analysis.h"
+#include "berberis/backend/common/machine_ir.h"
+#include "berberis/base/arena_alloc.h"
+#include "berberis/base/arena_list.h"
+#include "berberis/base/arena_vector.h"
+#include "berberis/base/logging.h"
+#include "berberis/runtime_primitives/config.h"
+
+// #define LOG_REG_ALLOC(...) ALOGE(__VA_ARGS__)
+#define LOG_REG_ALLOC(...) ((void)0)
+
+namespace berberis {
+
+namespace {
+
+// Lifetimes themselves are owned by lifetime list populated by lifetime
+// analysis. Use list of pointers to track lifetimes currently allocated
+// to particular hard register.
+using VRegLifetimePtrList = ArenaList<VRegLifetime*>;
+
+// How to spill one virtual register.
+struct VRegLifetimeSpill {
+ VRegLifetimePtrList::iterator lifetime;
+ SplitPos realloc_pos;
+
+ VRegLifetimeSpill(VRegLifetimePtrList::iterator i, const SplitPos& p)
+ : lifetime(i), realloc_pos(p) {}
+};
+
+// Every possible spill should have some smaller weight (CHECKed).
+const int kInfiniteSpillWeight = 99999;
+
+// Track what virtual registers are currently allocated to this particular
+// hard register, and how to spill them.
+class HardRegAllocation {
+ public:
+ explicit HardRegAllocation(Arena* arena)
+ : arena_(arena), lifetimes_(arena), new_lifetime_(nullptr), spills_(arena) {}
+
+ HardRegAllocation(const HardRegAllocation& other) = default;
+
+ // If new_lifetime doesn't interfere with lifetimes currently allocated
+ // to this hard register, allocate new_lifetime to this register as well.
+ bool TryAssign(VRegLifetime* new_lifetime);
+
+ // If TryAssign returned false:
+ // Check if it is possible to spill all lifetimes allocated to this hard
+ // register that interfere with new_lifetime, and return spill weight.
+ int ConsiderSpill(VRegLifetime* new_lifetime);
+
+ // If ConsiderSpill returned non-infinite weight:
+ // Given spill is possible, actually spill lifetimes that interfere with
+ // new_lifetime to spill_slot. Insert newly created tiny lifetimes into
+ // 'lifetimes' list, starting at position 'pos'.
+ void SpillAndAssign(VRegLifetime* new_lifetime,
+ int spill_slot,
+ VRegLifetimeList* lifetimes,
+ VRegLifetimeList::iterator pos);
+
+ private:
+ // Arena for allocations.
+ Arena* arena_;
+ // Lifetimes currently allocated to this hard register.
+ VRegLifetimePtrList lifetimes_;
+
+ // Last lifetime being allocated, for CHECKing.
+ // TODO(b/232598137): probably use this for ConsiderSpill and SpillAndAssign?
+ // This looks more natural...
+ VRegLifetime* new_lifetime_;
+
+ // How to free this register for last considered new lifetime.
+ // This is here for the following reasons:
+ // - it is highly coupled with lifetimes_
+ // - to avoid reallocating this for every spill consideration
+ ArenaVector<VRegLifetimeSpill> spills_;
+};
+
+bool HardRegAllocation::TryAssign(VRegLifetime* new_lifetime) {
+ // TODO(b/232598137): had to disable the check below! The problem is that when
+ // new_lifetime_ is split so that there remains no live ranges, we can't
+ // call begin() for it. Seems this place requires some rethinking, as such
+ // case means we can simply reorder lifetimes instead of actually splitting...
+ // Check lifetimes are processed in order by increasing begin.
+ // CHECK(!new_lifetime_ || new_lifetime_->begin() <= new_lifetime->begin());
+ new_lifetime_ = new_lifetime;
+
+ for (auto curr = lifetimes_.begin(); curr != lifetimes_.end();) {
+ VRegLifetime* curr_lifetime = *curr;
+
+ if (curr_lifetime->end() <= new_lifetime->begin()) {
+ // Curr lifetime ends before new lifetime starts, expire it.
+ curr = lifetimes_.erase(curr);
+ } else if (curr_lifetime->TestInterference(*new_lifetime)) {
+ // Lifetimes interfere, can't assign.
+ return false;
+ } else {
+ ++curr;
+ }
+ }
+
+ // No lifetimes interfere with new, can assign.
+ lifetimes_.push_back(new_lifetime);
+ return true;
+}
+
+int HardRegAllocation::ConsiderSpill(VRegLifetime* new_lifetime) {
+ CHECK_EQ(new_lifetime_, new_lifetime);
+
+ spills_.clear();
+ int weight = 0;
+
+ for (auto curr = lifetimes_.begin(); curr != lifetimes_.end(); ++curr) {
+ VRegLifetime* curr_lifetime = *curr;
+
+ if (!curr_lifetime->TestInterference(*new_lifetime)) {
+ // No interference, no need to spill.
+ continue;
+ }
+
+ SplitPos split_pos;
+ SplitKind split_kind = curr_lifetime->FindSplitPos(new_lifetime->begin(), &split_pos);
+ if (split_kind == SPLIT_IMPOSSIBLE) {
+ // Lifetimes interfere in such a way that spill is not possible.
+ return kInfiniteSpillWeight;
+ } else if (split_kind == SPLIT_CONFLICT) {
+ // A use within this lifetime conflicts with first use in 'new_lifetime'.
+ // If we spill it, it will compete with 'new_lifetime' at reallocation,
+ // and if it can only use register suitable for 'new_lifetime' as well,
+ // 'new_lifetime' can be evicted back, resulting in double spill.
+ if (curr_lifetime->GetRegClass()->IsSubsetOf(new_lifetime->GetRegClass())) {
+ return kInfiniteSpillWeight;
+ }
+ }
+
+ // Record spill.
+ spills_.push_back(VRegLifetimeSpill(curr, split_pos));
+ // Evicting tiny lifetime is free.
+ if (curr_lifetime->GetSpill() == -1) {
+ weight += curr_lifetime->spill_weight();
+ }
+ }
+
+ CHECK_LT(weight, kInfiniteSpillWeight);
+ return weight;
+}
+
+// Same as std::list::merge, but starting from a 'pos' position.
+void MergeVRegLifetimeList(VRegLifetimeList* dst,
+ VRegLifetimeList::iterator dst_pos,
+ VRegLifetimeList* src) {
+ while (!src->empty()) {
+ auto curr = src->begin();
+ for (; dst_pos != dst->end(); ++dst_pos) {
+ if (curr->begin() < dst_pos->begin()) {
+ break;
+ }
+ }
+ dst->splice(dst_pos, *src, curr);
+ }
+}
+
+void HardRegAllocation::SpillAndAssign(VRegLifetime* new_lifetime,
+ int spill_slot,
+ VRegLifetimeList* lifetimes,
+ VRegLifetimeList::iterator pos) {
+ CHECK_EQ(new_lifetime_, new_lifetime);
+ CHECK(!spills_.empty());
+
+ for (const auto& spill : spills_) {
+ VRegLifetime* spill_lifetime = *spill.lifetime;
+
+ // Assign spill slot.
+ // Lifetimes being spilled do not interfere, can share spill slot!
+ // TODO(b/232598137): evicted tiny lifetime have spill slot already.
+ // If we only evict tiny lifetimes, we might not need new spill_slot!
+ // Allocate spill slot here when needed.
+ if (spill_lifetime->GetSpill() == -1) {
+ spill_lifetime->SetSpill(spill_slot);
+ }
+
+ // Split spilled lifetime into tiny lifetimes, enqueue them for allocation.
+ VRegLifetimeList split(arena_);
+ spill_lifetime->Split(spill.realloc_pos, &split);
+ MergeVRegLifetimeList(lifetimes, pos, &split);
+
+ // Expire spilled lifetime.
+ lifetimes_.erase(spill.lifetime);
+ }
+
+ // Spilled all interfering lifetimes, can assign.
+ lifetimes_.push_back(new_lifetime);
+}
+
+// Simple register allocator.
+// Walk list of lifetimes sorted by begin and allocates in order.
+// Modifies lifetimes that have been spilled and adds tiny lifetimes split
+// from spilled lifetimes to the same list.
+class VRegLifetimeAllocator {
+ public:
+ VRegLifetimeAllocator(MachineIR* machine_ir, VRegLifetimeList* lifetimes)
+ : machine_ir_(machine_ir),
+ lifetimes_(lifetimes),
+ allocations_(config::kMaxHardRegs,
+ HardRegAllocation(machine_ir->arena()),
+ machine_ir->arena()) {}
+
+ void Allocate();
+
+ private:
+ void AllocateLifetime(VRegLifetimeList::iterator lifetime_it);
+
+ bool TryAssignHardReg(VRegLifetime* lifetime, MachineReg hard_reg);
+
+ int ConsiderSpillHardReg(MachineReg hard_reg, VRegLifetime* lifetime);
+
+ void SpillAndAssignHardReg(MachineReg hard_reg, VRegLifetimeList::iterator curr);
+
+ void RewriteAllocatedLifetimes();
+
+ MachineIR* machine_ir_;
+
+ VRegLifetimeList* lifetimes_;
+
+ ArenaVector<HardRegAllocation> allocations_;
+};
+
+int VRegLifetimeAllocator::ConsiderSpillHardReg(MachineReg hard_reg, VRegLifetime* lifetime) {
+ return allocations_[hard_reg.reg()].ConsiderSpill(lifetime);
+}
+
+bool VRegLifetimeAllocator::TryAssignHardReg(VRegLifetime* curr_lifetime, MachineReg hard_reg) {
+ if (allocations_[hard_reg.reg()].TryAssign(curr_lifetime)) {
+ curr_lifetime->set_hard_reg(hard_reg);
+ LOG_REG_ALLOC(".. to %s\n", GetMachineHardRegDebugName(hard_reg));
+ return true;
+ }
+ return false;
+}
+
+void VRegLifetimeAllocator::SpillAndAssignHardReg(MachineReg hard_reg,
+ VRegLifetimeList::iterator curr) {
+ auto next = std::next(curr);
+ allocations_[hard_reg.reg()].SpillAndAssign(&*curr, machine_ir_->AllocSpill(), lifetimes_, next);
+ curr->set_hard_reg(hard_reg);
+ LOG_REG_ALLOC(".. to %s (after spill)\n", GetMachineHardRegDebugName(hard_reg));
+}
+
+void VRegLifetimeAllocator::AllocateLifetime(VRegLifetimeList::iterator lifetime_it) {
+ VRegLifetime* lifetime = &*lifetime_it;
+ const MachineRegClass* reg_class = lifetime->GetRegClass();
+
+ LOG_REG_ALLOC(
+ "allocating lifetime %s:\n%s", reg_class->GetDebugName(), lifetime->GetDebugString().c_str());
+
+ // First try preferred register.
+ MachineReg pref_reg = lifetime->FindMoveHint()->hard_reg();
+ if (reg_class->HasReg(pref_reg) && TryAssignHardReg(lifetime, pref_reg)) {
+ return;
+ }
+
+ // Walk registers from reg class.
+ for (MachineReg hard_reg : *reg_class) {
+ if (hard_reg != pref_reg && TryAssignHardReg(lifetime, hard_reg)) {
+ return;
+ }
+ }
+
+ LOG_REG_ALLOC("... failed to find free hard reg, will try spilling");
+
+ // Walk registers again, consider each for spilling.
+ int best_spill_weight = kInfiniteSpillWeight;
+ MachineReg best_reg{0};
+ for (MachineReg hard_reg : *reg_class) {
+ int spill_weight = ConsiderSpillHardReg(hard_reg, lifetime);
+ LOG_REG_ALLOC(
+ "... consider spilling %s, weight %d", GetMachineHardRegDebugName(hard_reg), spill_weight);
+ if (best_spill_weight > spill_weight) {
+ best_spill_weight = spill_weight;
+ best_reg = hard_reg;
+ }
+ }
+
+ // Spill register with best spill weight.
+ CHECK_LT(best_spill_weight, kInfiniteSpillWeight);
+ SpillAndAssignHardReg(best_reg, lifetime_it);
+}
+
+void VRegLifetimeAllocator::RewriteAllocatedLifetimes() {
+ for (auto& lifetime : *lifetimes_) {
+ lifetime.Rewrite(machine_ir_);
+ }
+}
+
+void VRegLifetimeAllocator::Allocate() {
+ for (auto lifetime_it = lifetimes_->begin(); lifetime_it != lifetimes_->end(); ++lifetime_it) {
+ AllocateLifetime(lifetime_it);
+ }
+ RewriteAllocatedLifetimes();
+}
+
+void CollectLifetimes(const MachineIR* machine_ir, VRegLifetimeList* lifetimes) {
+ VRegLifetimeAnalysis lifetime_analysis(machine_ir->arena(), 2 * machine_ir->NumVReg(), lifetimes);
+
+ // Not 'const' because we need pointer to modifiable bb->insn_list().
+ for (auto* bb : machine_ir->bb_list()) {
+ for (const auto reg : bb->live_in()) {
+ lifetime_analysis.SetLiveIn(reg);
+ }
+
+ for (auto insn_it = bb->insn_list().begin(); insn_it != bb->insn_list().end(); ++insn_it) {
+ lifetime_analysis.AddInsn(MachineInsnListPosition(&(bb->insn_list()), insn_it));
+ }
+
+ for (const auto reg : bb->live_out()) {
+ lifetime_analysis.SetLiveOut(reg);
+ }
+ lifetime_analysis.EndBasicBlock();
+ }
+}
+
+} // namespace
+
+void AllocRegs(MachineIR* machine_ir) {
+ VRegLifetimeList lifetimes(machine_ir->arena());
+ CollectLifetimes(machine_ir, &lifetimes);
+
+ VRegLifetimeAllocator allocator(machine_ir, &lifetimes);
+ allocator.Allocate();
+}
+
+} // namespace berberis
diff --git a/backend/include/berberis/backend/common/reg_alloc.h b/backend/include/berberis/backend/common/reg_alloc.h
new file mode 100644
index 0000000..c390cf5
--- /dev/null
+++ b/backend/include/berberis/backend/common/reg_alloc.h
@@ -0,0 +1,28 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef BERBERIS_BACKEND_COMMON_REG_ALLOC_H_
+#define BERBERIS_BACKEND_COMMON_REG_ALLOC_H_
+
+namespace berberis {
+
+class MachineIR;
+
+void AllocRegs(MachineIR* machine_ir);
+
+} // namespace berberis
+
+#endif // BERBERIS_BACKEND_COMMON_REG_ALLOC_H_
diff --git a/runtime_primitives/include/berberis/runtime_primitives/config.h b/runtime_primitives/include/berberis/runtime_primitives/config.h
index 54fa647..26011dd 100644
--- a/runtime_primitives/include/berberis/runtime_primitives/config.h
+++ b/runtime_primitives/include/berberis/runtime_primitives/config.h
@@ -44,6 +44,8 @@
inline constexpr bool kLinkJumpsBetweenRegions = !kAllJumpsExitGeneratedCode;
// Guest page size. Always 4K for now.
inline constexpr size_t kGuestPageSize = 4096;
+// Number of hard registers assumed by the register allocator.
+inline constexpr uint32_t kMaxHardRegs = 64u;
} // namespace berberis::config
