reg,pass: refactor allocation of aliased registers (#121)

Issue #100 demonstrated that register allocation for aliased registers is
fundamentally broken. The root of the issue is that currently accesses to the
same virtual register with different masks are treated as different registers.
This PR takes a different approach:

* Liveness analysis is masked: we now properly consider which parts of a register are live
* Register allocation produces a mapping from virtual to physical ID, and aliasing is applied later

In addition, a new pass ZeroExtend32BitOutputs accounts for the fact that 32-bit writes in 64-bit mode should actually be treated as 64-bit writes (the result is zero-extended).

Closes #100

build/global.go

@@ -70,7 +70,13 @@ func Constraint(t buildtags.ConstraintConvertable) { ctx.Constraint(t) }
 // constraint comments.
 func ConstraintExpr(expr string) { ctx.ConstraintExpr(expr) }
 
-// GP8 allocates and returns a general-purpose 8-bit register.
+// GP8L allocates and returns a general-purpose 8-bit register (low byte).
+func GP8L() reg.GPVirtual { return ctx.GP8L() }
+
+// GP8H allocates and returns a general-purpose 8-bit register (high byte).
+func GP8H() reg.GPVirtual { return ctx.GP8H() }
+
+// GP8 allocates and returns a general-purpose 8-bit register (low byte).
 func GP8() reg.GPVirtual { return ctx.GP8() }
 
 // GP16 allocates and returns a general-purpose 16-bit register.

ir/ir.go

@@ -55,8 +55,8 @@ type Instruction struct {
 	Succ []*Instruction
 
 	// LiveIn/LiveOut are sets of live register IDs pre/post execution.
-	LiveIn  reg.Set
-	LiveOut reg.Set
+	LiveIn  reg.MaskSet
+	LiveOut reg.MaskSet
 }
 
 func (i *Instruction) node() {}

operand/types.go

@@ -141,10 +141,10 @@ func Registers(op Op) []reg.Register {
 func ApplyAllocation(op Op, a reg.Allocation) Op {
 	switch op := op.(type) {
 	case reg.Register:
-		return a.LookupDefault(op)
+		return a.LookupRegisterDefault(op)
 	case Mem:
-		op.Base = a.LookupDefault(op.Base)
-		op.Index = a.LookupDefault(op.Index)
+		op.Base = a.LookupRegisterDefault(op.Base)
+		op.Index = a.LookupRegisterDefault(op.Index)
 		return op
 	}
 	return op

pass/alloc.go

@@ -3,35 +3,51 @@ package pass
 import (
 	"errors"
 	"math"
+	"sort"
 
 	"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
+	X, Y reg.ID
 }
 
 // Allocator is a graph-coloring register allocator.
 type Allocator struct {
-	registers  []reg.Physical
+	registers  []reg.ID
 	allocation reg.Allocation
 	edges      []*edge
-	possible   map[reg.Virtual][]reg.Physical
-	vidtopid   map[reg.VID]reg.PID
+	possible   map[reg.ID][]reg.ID
 }
 
 // NewAllocator builds an allocator for the given physical registers.
 func NewAllocator(rs []reg.Physical) (*Allocator, error) {
-	if len(rs) == 0 {
-		return nil, errors.New("no registers")
+	// Set of IDs, excluding restricted registers.
+	idset := map[reg.ID]bool{}
+	for _, r := range rs {
+		if (r.Info() & reg.Restricted) != 0 {
+			continue
+		}
+		idset[r.ID()] = true
+	}
+
+	if len(idset) == 0 {
+		return nil, errors.New("no allocatable registers")
 	}
+
+	// Produce slice of unique register IDs.
+	var ids []reg.ID
+	for id := range idset {
+		ids = append(ids, id)
+	}
+	sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
+
 	return &Allocator{
-		registers:  rs,
+		registers:  ids,
 		allocation: reg.NewEmptyAllocation(),
-		possible:   map[reg.Virtual][]reg.Physical{},
-		vidtopid:   map[reg.VID]reg.PID{},
+		possible:   map[reg.ID][]reg.ID{},
 	}, nil
 }
 
@@ -45,23 +61,24 @@ func NewAllocatorForKind(k reg.Kind) (*Allocator, error) {
 }
 
 // AddInterferenceSet records that r interferes with every register in s. Convenience wrapper around AddInterference.
-func (a *Allocator) AddInterferenceSet(r reg.Register, s reg.Set) {
-	for y := range s {
-		a.AddInterference(r, y)
+func (a *Allocator) AddInterferenceSet(r reg.Register, s reg.MaskSet) {
+	for id, mask := range s {
+		if (r.Mask() & mask) != 0 {
+			a.AddInterference(r.ID(), id)
+		}
 	}
 }
 
 // AddInterference records that x and y must be assigned to non-conflicting physical registers.
-func (a *Allocator) AddInterference(x, y reg.Register) {
+func (a *Allocator) AddInterference(x, y reg.ID) {
 	a.Add(x)
 	a.Add(y)
 	a.edges = append(a.edges, &edge{X: x, Y: y})
 }
 
 // Add adds a register to be allocated. Does nothing if the register has already been added.
-func (a *Allocator) Add(r reg.Register) {
-	v, ok := r.(reg.Virtual)
-	if !ok {
+func (a *Allocator) Add(v reg.ID) {
+	if !v.IsVirtual() {
 		return
 	}
 	if _, found := a.possible[v]; found {
@@ -91,35 +108,22 @@ func (a *Allocator) Allocate() (reg.Allocation, error) {
 
 // update possible allocations based on edges.
 func (a *Allocator) update() error {
-	for v := range a.possible {
-		pid, found := a.vidtopid[v.VirtualID()]
-		if !found {
-			continue
-		}
-		a.possible[v] = filterregisters(a.possible[v], func(r reg.Physical) bool {
-			return r.PhysicalID() == pid
-		})
-	}
-
 	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)
-
+		x := a.allocation.LookupDefault(e.X)
+		y := a.allocation.LookupDefault(e.Y)
 		switch {
-		case vx != nil && vy != nil:
+		case x.IsVirtual() && y.IsVirtual():
 			rem = append(rem, e)
 			continue
-		case px != nil && py != nil:
-			if reg.AreConflicting(px, py) {
+		case x.IsPhysical() && y.IsPhysical():
+			if x == y {
 				return errors.New("impossible register allocation")
 			}
-		case px != nil && vy != nil:
-			a.discardconflicting(vy, px)
-		case vx != nil && py != nil:
-			a.discardconflicting(vx, py)
+		case x.IsPhysical() && y.IsVirtual():
+			a.discardconflicting(y, x)
+		case x.IsVirtual() && y.IsPhysical():
+			a.discardconflicting(x, y)
 		default:
 			panic("unreachable")
 		}
@@ -130,38 +134,36 @@ func (a *Allocator) update() error {
 }
 
 // mostrestricted returns the virtual register with the least possibilities.
-func (a *Allocator) mostrestricted() reg.Virtual {
+func (a *Allocator) mostrestricted() reg.ID {
 	n := int(math.MaxInt32)
-	var v reg.Virtual
-	for r, p := range a.possible {
-		if len(p) < n || (len(p) == n && v != nil && r.VirtualID() < v.VirtualID()) {
+	var v reg.ID
+	for w, p := range a.possible {
+		// On a tie, choose the smallest ID in numeric order. This avoids
+		// non-deterministic allocations due to map iteration order.
+		if len(p) < n || (len(p) == n && w < v) {
 			n = len(p)
-			v = r
+			v = w
 		}
 	}
 	return v
 }
 
 // discardconflicting removes registers from vs possible list that conflict with p.
-func (a *Allocator) discardconflicting(v reg.Virtual, p reg.Physical) {
-	a.possible[v] = filterregisters(a.possible[v], func(r reg.Physical) bool {
-		if pid, found := a.vidtopid[v.VirtualID()]; found && pid == p.PhysicalID() {
-			return true
-		}
-		return !reg.AreConflicting(r, p)
+func (a *Allocator) discardconflicting(v, p reg.ID) {
+	a.possible[v] = filterregisters(a.possible[v], func(r reg.ID) bool {
+		return r != p
 	})
 }
 
 // alloc attempts to allocate a register to v.
-func (a *Allocator) alloc(v reg.Virtual) error {
+func (a *Allocator) alloc(v reg.ID) error {
 	ps := a.possible[v]
 	if len(ps) == 0 {
 		return errors.New("failed to allocate registers")
 	}
 	p := ps[0]
 	a.allocation[v] = p
 	delete(a.possible, v)
-	a.vidtopid[v.VirtualID()] = p.PhysicalID()
 	return nil
 }
 
@@ -171,14 +173,14 @@ func (a *Allocator) remaining() int {
 }
 
 // possibleregisters returns all allocate-able registers for the given virtual.
-func (a *Allocator) possibleregisters(v reg.Virtual) []reg.Physical {
-	return filterregisters(a.registers, func(r reg.Physical) bool {
-		return v.SatisfiedBy(r) && (r.Info()&reg.Restricted) == 0
+func (a *Allocator) possibleregisters(v reg.ID) []reg.ID {
+	return filterregisters(a.registers, func(r reg.ID) bool {
+		return v.Kind() == r.Kind()
 	})
 }
 
-func filterregisters(in []reg.Physical, predicate func(reg.Physical) bool) []reg.Physical {
-	var rs []reg.Physical
+func filterregisters(in []reg.ID, predicate func(reg.ID) bool) []reg.ID {
+	var rs []reg.ID
 	for _, r := range in {
 		if predicate(r) {
 			rs = append(rs, r)

pass/alloc_test.go

@@ -15,9 +15,9 @@ func TestAllocatorSimple(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	a.Add(x)
-	a.Add(y)
-	a.AddInterference(x, y)
+	a.Add(x.ID())
+	a.Add(y.ID())
+	a.AddInterference(x.ID(), y.ID())
 
 	alloc, err := a.Allocate()
 	if err != nil {
@@ -26,7 +26,7 @@ func TestAllocatorSimple(t *testing.T) {
 
 	t.Log(alloc)
 
-	if alloc[x] != reg.X0 || alloc[y] != reg.Y1 {
+	if alloc.LookupRegister(x) != reg.X0 || alloc.LookupRegister(y) != reg.Y1 {
 		t.Fatalf("unexpected allocation")
 	}
 }
@@ -37,7 +37,7 @@ func TestAllocatorImpossible(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	a.AddInterference(reg.X7, reg.Z7)
+	a.AddInterference(reg.X7.ID(), reg.Z7.ID())
 
 	_, err = a.Allocate()
 	if err == nil {

pass/pass.go

@@ -16,6 +16,7 @@ var Compile = Concat(
 	FunctionPass(PruneDanglingLabels),
 	FunctionPass(LabelTarget),
 	FunctionPass(CFG),
+	InstructionPass(ZeroExtend32BitOutputs),
 	FunctionPass(Liveness),
 	FunctionPass(AllocateRegisters),
 	FunctionPass(BindRegisters),

pass/reg.go

@@ -8,6 +8,24 @@ import (
 	"github.com/mmcloughlin/avo/reg"
 )
 
+// ZeroExtend32BitOutputs applies the rule that "32-bit operands generate a
+// 32-bit result, zero-extended to a 64-bit result in the destination
+// general-purpose register" (Intel Software Developer’s Manual, Volume 1,
+// 3.4.1.1).
+func ZeroExtend32BitOutputs(i *ir.Instruction) error {
+	for j, op := range i.Outputs {
+		if !operand.IsR32(op) {
+			continue
+		}
+		r, ok := op.(reg.GP)
+		if !ok {
+			panic("r32 operand should satisfy reg.GP")
+		}
+		i.Outputs[j] = r.As64()
+	}
+	return nil
+}
+
 // Liveness computes register liveness.
 func Liveness(fn *ir.Function) error {
 	// Note this implementation is initially naive so as to be "obviously correct".
@@ -23,8 +41,8 @@ func Liveness(fn *ir.Function) error {
 
 	// Initialize.
 	for _, i := range is {
-		i.LiveIn = reg.NewSetFromSlice(i.InputRegisters())
-		i.LiveOut = reg.NewEmptySet()
+		i.LiveIn = reg.NewMaskSetFromRegisters(i.InputRegisters())
+		i.LiveOut = reg.NewEmptyMaskSet()
 	}
 
 	// Iterative dataflow analysis.
@@ -33,29 +51,16 @@ func Liveness(fn *ir.Function) error {
 
 		for _, i := range is {
 			// out[n] = UNION[s IN succ[n]] in[s]
-			nout := len(i.LiveOut)
 			for _, s := range i.Succ {
 				if s == nil {
 					continue
 				}
-				i.LiveOut.Update(s.LiveIn)
-			}
-			if len(i.LiveOut) != nout {
-				changes = true
+				changes = i.LiveOut.Update(s.LiveIn) || changes
 			}
 
 			// in[n] = use[n] UNION (out[n] - def[n])
-			nin := len(i.LiveIn)
-			def := reg.NewSetFromSlice(i.OutputRegisters())
-			i.LiveIn.Update(i.LiveOut.Difference(def))
-			for r := range i.LiveOut {
-				if _, found := def[r]; !found {
-					i.LiveIn.Add(r)
-				}
-			}
-			if len(i.LiveIn) != nin {
-				changes = true
-			}
+			def := reg.NewMaskSetFromRegisters(i.OutputRegisters())
+			changes = i.LiveIn.Update(i.LiveOut.Difference(def)) || changes
 		}
 
 		if !changes {
@@ -80,7 +85,7 @@ func AllocateRegisters(fn *ir.Function) error {
 				}
 				as[k] = a
 			}
-			as[k].Add(r)
+			as[k].Add(r.ID())
 		}
 	}
 
@@ -89,7 +94,7 @@ func AllocateRegisters(fn *ir.Function) error {
 		for _, d := range i.OutputRegisters() {
 			k := d.Kind()
 			out := i.LiveOut.OfKind(k)
-			out.Discard(d)
+			out.DiscardRegister(d)
 			as[k].AddInterferenceSet(d, out)
 		}
 	}