pass: first attempt at register allocation

pass/alloc.go

@@ -0,0 +1,153 @@
+package pass
+
+import (
+	"errors"
+	"math"
+
+	"github.com/mmcloughlin/avo/reg"
+)
+
+// edge is an edge of the interference graph, indicating that registers X and Y
+// must be in non-conflicting registers.
+type edge struct {
+	X, Y reg.Register
+}
+
+type Allocator struct {
+	registers  []reg.Physical
+	allocation reg.Allocation
+	edges      []*edge
+	possible   map[reg.Virtual][]reg.Physical
+}
+
+func NewAllocator(rs []reg.Physical) (*Allocator, error) {
+	if len(rs) == 0 {
+		return nil, errors.New("no registers")
+	}
+	return &Allocator{
+		registers:  rs,
+		allocation: reg.NewEmptyAllocation(),
+		possible:   map[reg.Virtual][]reg.Physical{},
+	}, nil
+}
+
+func NewAllocatorForKind(k reg.Kind) (*Allocator, error) {
+	f := reg.FamilyOfKind(k)
+	if f == nil {
+		return nil, errors.New("unknown register family")
+	}
+	return NewAllocator(f.Registers())
+}
+
+func (a *Allocator) AddInterferenceSet(r reg.Register, s reg.Set) {
+	for y := range s {
+		a.AddInterference(r, y)
+	}
+}
+
+func (a *Allocator) AddInterference(x, y reg.Register) {
+	a.add(x)
+	a.add(y)
+	a.edges = append(a.edges, &edge{X: x, Y: y})
+}
+
+func (a *Allocator) Allocate() (reg.Allocation, error) {
+	for a.remaining() > 0 {
+		a.update()
+
+		v := a.mostrestricted()
+		if err := a.alloc(v); err != nil {
+			return nil, err
+		}
+	}
+	return a.allocation, nil
+}
+
+// add adds a register.
+func (a *Allocator) add(r reg.Register) {
+	v, ok := r.(reg.Virtual)
+	if !ok {
+		return
+	}
+	a.possible[v] = a.registersofsize(v.Bytes())
+}
+
+// update possible allocations based on edges.
+func (a *Allocator) update() error {
+	var rem []*edge
+	for _, e := range a.edges {
+		e.X, e.Y = a.allocation.LookupDefault(e.X), a.allocation.LookupDefault(e.Y)
+
+		px, py := reg.ToPhysical(e.X), reg.ToPhysical(e.Y)
+		vx, vy := reg.ToVirtual(e.X), reg.ToVirtual(e.Y)
+
+		switch {
+		case vx != nil && vy != nil:
+			rem = append(rem, e)
+			continue
+		case px != nil && py != nil:
+			if reg.AreConflicting(px, py) {
+				return errors.New("impossible register allocation")
+			}
+		case px != nil && vy != nil:
+			a.discardconflicting(vy, px)
+		case vx != nil && py != nil:
+			a.discardconflicting(vx, py)
+		default:
+			panic("unreachable")
+		}
+	}
+	a.edges = rem
+	return nil
+}
+
+// mostrestricted returns the virtual register with the least possibilities.
+func (a *Allocator) mostrestricted() reg.Virtual {
+	n := int(math.MaxInt32)
+	var v reg.Virtual
+	for r, p := range a.possible {
+		if len(p) < n {
+			n = len(p)
+			v = r
+		}
+	}
+	return v
+}
+
+// discardconflicting removes registers from vs possible list that conflict with p.
+func (a *Allocator) discardconflicting(v reg.Virtual, p reg.Physical) {
+	var rs []reg.Physical
+	for _, r := range a.possible[v] {
+		if !reg.AreConflicting(r, p) {
+			rs = append(rs, r)
+		}
+	}
+	a.possible[v] = rs
+}
+
+// alloc attempts to allocate a register to v.
+func (a *Allocator) alloc(v reg.Virtual) error {
+	ps := a.possible[v]
+	if len(ps) == 0 {
+		return errors.New("failed to allocate registers")
+	}
+	a.allocation[v] = ps[0]
+	delete(a.possible, v)
+	return nil
+}
+
+// remaining returns the number of unallocated registers.
+func (a *Allocator) remaining() int {
+	return len(a.possible)
+}
+
+// registersofsize returns all registers of the given size.
+func (a *Allocator) registersofsize(n uint) []reg.Physical {
+	var rs []reg.Physical
+	for _, r := range a.registers {
+		if r.Bytes() == n {
+			rs = append(rs, r)
+		}
+	}
+	return rs
+}

pass/reg.go

@@ -1,12 +1,18 @@
 package pass
 
 import (
+	"errors"
+
 	"github.com/mmcloughlin/avo"
+	"github.com/mmcloughlin/avo/operand"
 	"github.com/mmcloughlin/avo/reg"
 )
 
 // Liveness computes register liveness.
 func Liveness(fn *avo.Function) error {
+	// Note this implementation is initially naive so as to be "obviously correct".
+	// There are a well-known optimizations we can apply if necessary.
+
 	is := fn.Instructions()
 
 	// Initialize to empty sets.
@@ -51,5 +57,58 @@ func Liveness(fn *avo.Function) error {
 }
 
 func AllocateRegisters(fn *avo.Function) error {
+	// Build one allocator per register kind and record register interferences.
+	as := map[reg.Kind]*Allocator{}
+	for _, i := range fn.Instructions() {
+		for _, d := range i.OutputRegisters() {
+			k := d.Kind()
+			if _, found := as[k]; !found {
+				a, err := NewAllocatorForKind(k)
+				if err != nil {
+					return err
+				}
+				as[k] = a
+			}
+
+			out := i.LiveOut.OfKind(k)
+			out.Discard(d)
+			as[k].AddInterferenceSet(d, out)
+		}
+	}
+
+	// Execute register allocation.
+	fn.Allocation = reg.NewEmptyAllocation()
+	for _, a := range as {
+		al, err := a.Allocate()
+		if err != nil {
+			return err
+		}
+		if err := fn.Allocation.Merge(al); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func BindRegisters(fn *avo.Function) error {
+	for _, i := range fn.Instructions() {
+		for idx := range i.Operands {
+			i.Operands[idx] = operand.ApplyAllocation(i.Operands[idx], fn.Allocation)
+		}
+	}
+	return nil
+}
+
+func VerifyAllocation(fn *avo.Function) error {
+	// All registers should be physical.
+	for _, i := range fn.Instructions() {
+		for _, r := range i.Registers() {
+			if reg.ToPhysical(r) == nil {
+				return errors.New("non physical register found")
+			}
+		}
+	}
+
 	return nil
 }