cmd/compile: look for newobject in register ABI for write barrier elision

If we are assigning a global address to an object that is
immediately returned from runtime.newobject, we omit the write
barrier because we know that both the source (static address) and
the destination (zeroed memory) do not need to be tracked by the
GC. Currently, the code that matches runtime.newobject's result
is specific to ABI0 layout. Update the code to work with register
ABI as well.

Change-Id: I7ab0833c6f745329271881ee4169956928a3a948
Reviewed-on: https://go-review.googlesource.com/c/go/+/308709
Trust: Cherry Zhang <cherryyz@google.com>
Run-TryBot: Cherry Zhang <cherryyz@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: David Chase <drchase@google.com>

1 files changed, 50 insertions, 31 deletions

diff --git a/src/cmd/compile/internal/ssa/writebarrier.go b/src/cmd/compile/internal/ssa/writebarrier.go
index d6af3578d0..419d91d0d3 100644
--- a/src/cmd/compile/internal/ssa/writebarrier.go
+++ b/src/cmd/compile/internal/ssa/writebarrier.go
@@ -38,9 +38,11 @@ func needwb(v *Value, zeroes map[ID]ZeroRegion) bool {
 	if IsStackAddr(v.Args[0]) {
 		return false // write on stack doesn't need write barrier
 	}
-	if v.Op == OpMove && IsReadOnlyGlobalAddr(v.Args[1]) && IsNewObject(v.Args[0], v.MemoryArg()) {
-		// Copying data from readonly memory into a fresh object doesn't need a write barrier.
-		return false
+	if v.Op == OpMove && IsReadOnlyGlobalAddr(v.Args[1]) {
+		if mem, ok := IsNewObject(v.Args[0]); ok && mem == v.MemoryArg() {
+			// Copying data from readonly memory into a fresh object doesn't need a write barrier.
+			return false
+		}
 	}
 	if v.Op == OpStore && IsGlobalAddr(v.Args[1]) {
 		// Storing pointers to non-heap locations into zeroed memory doesn't need a write barrier.
@@ -389,11 +391,7 @@ func (f *Func) computeZeroMap() map[ID]ZeroRegion {
 	// Find new objects.
 	for _, b := range f.Blocks {
 		for _, v := range b.Values {
-			if v.Op != OpLoad {
-				continue
-			}
-			mem := v.MemoryArg()
-			if IsNewObject(v, mem) {
+			if mem, ok := IsNewObject(v); ok {
 				nptr := v.Type.Elem().Size() / ptrSize
 				if nptr > 64 {
 					nptr = 64
@@ -578,39 +576,60 @@ func IsReadOnlyGlobalAddr(v *Value) bool {
 	return false
 }
 
-// IsNewObject reports whether v is a pointer to a freshly allocated & zeroed object at memory state mem.
-func IsNewObject(v *Value, mem *Value) bool {
-	// TODO this will need updating for register args; the OpLoad is wrong.
-	if v.Op != OpLoad {
-		return false
-	}
-	if v.MemoryArg() != mem {
-		return false
+// IsNewObject reports whether v is a pointer to a freshly allocated & zeroed object,
+// if so, also returns the memory state mem at which v is zero.
+func IsNewObject(v *Value) (mem *Value, ok bool) {
+	f := v.Block.Func
+	c := f.Config
+	if f.ABIDefault == f.ABI1 && len(c.intParamRegs) >= 1 {
+		if v.Op != OpSelectN || v.AuxInt != 0 {
+			return nil, false
+		}
+		// Find the memory
+		for _, w := range v.Block.Values {
+			if w.Op == OpSelectN && w.AuxInt == 1 && w.Args[0] == v.Args[0] {
+				mem = w
+				break
+			}
+		}
+		if mem == nil {
+			return nil, false
+		}
+	} else {
+		if v.Op != OpLoad {
+			return nil, false
+		}
+		mem = v.MemoryArg()
+		if mem.Op != OpSelectN {
+			return nil, false
+		}
+		if mem.Type != types.TypeMem {
+			return nil, false
+		} // assume it is the right selection if true
 	}
-	if mem.Op != OpSelectN {
-		return false
+	call := mem.Args[0]
+	if call.Op != OpStaticCall {
+		return nil, false
 	}
-	if mem.Type != types.TypeMem {
-		return false
-	} // assume it is the right selection if true
-	mem = mem.Args[0]
-	if mem.Op != OpStaticCall {
-		return false
+	if !isSameCall(call.Aux, "runtime.newobject") {
+		return nil, false
 	}
-	if !isSameCall(mem.Aux, "runtime.newobject") {
-		return false
+	if f.ABIDefault == f.ABI1 && len(c.intParamRegs) >= 1 {
+		if v.Args[0] == call {
+			return mem, true
+		}
+		return nil, false
 	}
 	if v.Args[0].Op != OpOffPtr {
-		return false
+		return nil, false
 	}
 	if v.Args[0].Args[0].Op != OpSP {
-		return false
+		return nil, false
 	}
-	c := v.Block.Func.Config
 	if v.Args[0].AuxInt != c.ctxt.FixedFrameSize()+c.RegSize { // offset of return value
-		return false
+		return nil, false
 	}
-	return true
+	return mem, true
 }
 
 // IsSanitizerSafeAddr reports whether v is known to be an address