pass: support priority in register allocator (#180)

Updates #156

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 {

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())
+		}
+	}
+}