diff --git a/pass/alloc.go b/pass/alloc.go index fc7773a..39ada55 100644 --- a/pass/alloc.go +++ b/pass/alloc.go @@ -17,6 +17,7 @@ type edge struct { // Allocator is a graph-coloring register allocator. type Allocator struct { registers []reg.ID + priority map[reg.ID]int allocation reg.Allocation edges []*edge possible map[reg.ID][]reg.ID @@ -42,13 +43,16 @@ func NewAllocator(rs []reg.Physical) (*Allocator, error) { for id := range idset { ids = append(ids, id) } - sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) - return &Allocator{ + a := &Allocator{ registers: ids, + priority: map[reg.ID]int{}, allocation: reg.NewEmptyAllocation(), possible: map[reg.ID][]reg.ID{}, - }, nil + } + a.sortregisters() + + return a, nil } // NewAllocatorForKind builds an allocator for the given kind of registers. @@ -60,6 +64,25 @@ func NewAllocatorForKind(k reg.Kind) (*Allocator, error) { return NewAllocator(f.Registers()) } +// SetPriority sets the priority of the given regiser to p. Higher priority +// registers are preferred in allocations. By default all registers have 0 +// priority. Priority will only apply to subsequent Add() calls, therefore +// typically all SetPriority calls should happen at allocator initialization. +func (a *Allocator) SetPriority(id reg.ID, p int) { + a.priority[id] = p + a.sortregisters() +} + +// sortregisters sorts the list of available registers: higher priority first, +// falling back to sorting by ID. +func (a *Allocator) sortregisters() { + sort.Slice(a.registers, func(i, j int) bool { + ri, rj := a.registers[i], a.registers[j] + pi, pj := a.priority[ri], a.priority[rj] + return (pi > pj) || (pi == pj && ri < rj) + }) +} + // AddInterferenceSet records that r interferes with every register in s. Convenience wrapper around AddInterference. func (a *Allocator) AddInterferenceSet(r reg.Register, s reg.MaskSet) { for id, mask := range s { diff --git a/pass/alloc_test.go b/pass/alloc_test.go index d1935e7..cee9a9e 100644 --- a/pass/alloc_test.go +++ b/pass/alloc_test.go @@ -44,3 +44,55 @@ func TestAllocatorImpossible(t *testing.T) { t.Fatal("expected allocation error") } } + +func TestAllocatorPriority(t *testing.T) { + const n = 4 + + // Create an allocator with custom priorities. + a, err := NewAllocatorForKind(reg.KindVector) + if err != nil { + t.Fatal(err) + } + + a.SetPriority(reg.X0.ID(), -1) + a.SetPriority(reg.X7.ID(), 1) + a.SetPriority(reg.X13.ID(), 1) + a.SetPriority(reg.X3.ID(), 2) + + // The expected n highest priority registers. + expect := [n]reg.Physical{ + reg.X3, // priority 2, id 3 + reg.X7, // priority 1, id 7 + reg.X13, // priority 1, id 13 + reg.X1, // priority 0, id 1 (X0 has priority -1) + } + + // Setup allocation problem with n conflicting registers. + c := reg.NewCollection() + x := make([]reg.Virtual, n) + for i := range x { + x[i] = c.XMM() + } + + for i := range x { + a.Add(x[i].ID()) + } + + for i := 0; i < n; i++ { + for j := i + 1; j < n; j++ { + a.AddInterference(x[i].ID(), x[j].ID()) + } + } + + // Allocate and confirm expectation. + alloc, err := a.Allocate() + if err != nil { + t.Fatal(err) + } + + for i := range x { + if got := alloc.LookupRegister(x[i]); got != expect[i] { + t.Errorf("x[%d] allocated %s; expected %s", i, got.Asm(), expect[i].Asm()) + } + } +}