cmd/internal/obj: arm64 assembler support unshifted hi for large offsets

Extend splitImm24uScaled to support an unshifted hi value (hi <= 0xfff)
in addition to the shifted hi value (hi & ^0xfff000 == 0). This allows
load/store instructions to handle more offsets using ADD + load/store
sequences instead of falling back to the literal pool.

This will be used by a subsequent change to add FMOVQ support in SSA form.

Change-Id: I78490f5b1a60d49c1d42ad4daefb5d4e6021c965
Reviewed-on: https://go-review.googlesource.com/c/go/+/737320
Reviewed-by: Keith Randall <khr@golang.org>
Reviewed-by: Keith Randall <khr@google.com>
Auto-Submit: Keith Randall <khr@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: David Chase <drchase@google.com>

3 files changed, 182 insertions, 34 deletions

diff --git a/src/cmd/asm/internal/asm/testdata/arm64.s b/src/cmd/asm/internal/asm/testdata/arm64.s
index 6a44f0f5b0..edc6a8aafd 100644
--- a/src/cmd/asm/internal/asm/testdata/arm64.s
+++ b/src/cmd/asm/internal/asm/testdata/arm64.s
@@ -592,30 +592,30 @@ TEXT	foo(SB), DUPOK|NOSPLIT, $-8
 	MOVD R3, 4095(R17)                         // 3bfe3f91630300f9
 
 // large aligned offset, use two instructions(add+ldr/str).
-	MOVB	R1, 0x1001(R2) 		// MOVB		R1, 4097(R2)		// 5b04409161070039
+	MOVB	R1, 0x1001(R2) 		// MOVB		R1, 4097(R2)		// 5b08009161ff3f39
 	MOVB	R1, 0xffffff(R2)	// MOVB		R1, 16777215(R2)	// 5bfc7f9161ff3f39
-	MOVH	R1, 0x2002(R2)		// MOVH		R1, 8194(R2)		// 5b08409161070079
+	MOVH	R1, 0x2002(R2)		// MOVH		R1, 8194(R2)		// 5b10009161ff3f79
 	MOVH	R1, 0x1000ffe(R2)	// MOVH		R1, 16781310(R2)	// 5bfc7f9161ff3f79
-	MOVW	R1, 0x4004(R2)		// MOVW		R1, 16388(R2)		// 5b104091610700b9
+	MOVW	R1, 0x4004(R2)		// MOVW		R1, 16388(R2)		// 5b20009161ff3fb9
 	MOVW	R1, 0x1002ffc(R2)	// MOVW		R1, 16789500(R2)	// 5bfc7f9161ff3fb9
-	MOVD	R1, 0x8008(R2)		// MOVD		R1, 32776(R2)		// 5b204091610700f9
+	MOVD	R1, 0x8008(R2)		// MOVD		R1, 32776(R2)		// 5b40009161ff3ff9
 	MOVD	R1, 0x1006ff8(R2)	// MOVD		R1, 16805880(R2)	// 5bfc7f9161ff3ff9
-	FMOVS	F1, 0x4004(R2)		// FMOVS	F1, 16388(R2)		// 5b104091610700bd
+	FMOVS	F1, 0x4004(R2)		// FMOVS	F1, 16388(R2)		// 5b20009161ff3fbd
 	FMOVS	F1, 0x1002ffc(R2)	// FMOVS	F1, 16789500(R2)	// 5bfc7f9161ff3fbd
-	FMOVD	F1, 0x8008(R2)		// FMOVD	F1, 32776(R2)		// 5b204091610700fd
+	FMOVD	F1, 0x8008(R2)		// FMOVD	F1, 32776(R2)		// 5b40009161ff3ffd
 	FMOVD	F1, 0x1006ff8(R2)	// FMOVD	F1, 16805880(R2)	// 5bfc7f9161ff3ffd
 
-	MOVB	0x1001(R1), R2 		// MOVB		4097(R1), R2		// 3b04409162078039
+	MOVB	0x1001(R1), R2 		// MOVB		4097(R1), R2		// 3b08009162ffbf39
 	MOVB	0xffffff(R1), R2	// MOVB		16777215(R1), R2	// 3bfc7f9162ffbf39
-	MOVH	0x2002(R1), R2		// MOVH		8194(R1), R2		// 3b08409162078079
+	MOVH	0x2002(R1), R2		// MOVH		8194(R1), R2		// 3b10009162ffbf79
 	MOVH	0x1000ffe(R1), R2	// MOVH		16781310(R1), R2	// 3bfc7f9162ffbf79
-	MOVW	0x4004(R1), R2		// MOVW		16388(R1), R2		// 3b104091620780b9
+	MOVW	0x4004(R1), R2		// MOVW		16388(R1), R2		// 3b20009162ffbfb9
 	MOVW	0x1002ffc(R1), R2	// MOVW		16789500(R1), R2	// 3bfc7f9162ffbfb9
-	MOVD	0x8008(R1), R2		// MOVD		32776(R1), R2		// 3b204091620740f9
+	MOVD	0x8008(R1), R2		// MOVD		32776(R1), R2		// 3b40009162ff7ff9
 	MOVD	0x1006ff8(R1), R2	// MOVD		16805880(R1), R2	// 3bfc7f9162ff7ff9
-	FMOVS	0x4004(R1), F2		// FMOVS	16388(R1), F2		// 3b104091620740bd
+	FMOVS	0x4004(R1), F2		// FMOVS	16388(R1), F2		// 3b20009162ff7fbd
 	FMOVS	0x1002ffc(R1), F2	// FMOVS	16789500(R1), F2	// 3bfc7f9162ff7fbd
-	FMOVD	0x8008(R1), F2		// FMOVD	32776(R1), F2		// 3b204091620740fd
+	FMOVD	0x8008(R1), F2		// FMOVD	32776(R1), F2		// 3b40009162ff7ffd
 	FMOVD	0x1006ff8(R1), F2	// FMOVD	16805880(R1), F2	// 3bfc7f9162ff7ffd
 
 // very large or unaligned offset uses constant pool.
diff --git a/src/cmd/internal/obj/arm64/asm7.go b/src/cmd/internal/obj/arm64/asm7.go
index eedba28e3b..898977d62a 100644
--- a/src/cmd/internal/obj/arm64/asm7.go
+++ b/src/cmd/internal/obj/arm64/asm7.go
@@ -1392,10 +1392,8 @@ func roundUp(x, to uint32) uint32 {
 	return (x + to - 1) &^ (to - 1)
 }
 
-// splitImm24uScaled splits an immediate into a scaled 12 bit unsigned lo value
-// and an unscaled shifted 12 bit unsigned hi value. These are typically used
-// by adding or subtracting the hi value and using the lo value as the offset
-// for a load or store.
+// splitImm24uScaled returns hi, lo such that v == hi + lo<<shift.
+// Always 0 <= lo <= 0xfff, and hi is either 0 <= hi <= 0xfff, or (hi&0xfff == 0 && 0 <= hi <= 0xfff000).
 func splitImm24uScaled(v int32, shift int) (int32, int32, error) {
 	if v < 0 {
 		return 0, 0, fmt.Errorf("%d is not a 24 bit unsigned immediate", v)
@@ -1403,19 +1401,28 @@ func splitImm24uScaled(v int32, shift int) (int32, int32, error) {
 	if v > 0xfff000+0xfff<<shift {
 		return 0, 0, fmt.Errorf("%d is too large for a scaled 24 bit unsigned immediate", v)
 	}
-	if v&((1<<shift)-1) != 0 {
-		return 0, 0, fmt.Errorf("%d is not a multiple of %d", v, 1<<shift)
+
+	// Try hi <= 0xfff and lo <= 0xfff such that v = hi + (lo << shift).
+	hi := max(v-(0xfff<<shift), v&((1<<shift)-1))
+	if hi <= 0xfff {
+		lo := (v - hi) >> shift
+		if lo <= 0xfff {
+			return hi, lo, nil
+		}
 	}
+
+	// Try hi shifted left by 12 bits.
 	lo := (v >> shift) & 0xfff
-	hi := v - (lo << shift)
+	hi = v - (lo << shift)
 	if hi > 0xfff000 {
 		hi = 0xfff000
 		lo = (v - hi) >> shift
 	}
-	if hi & ^0xfff000 != 0 {
-		panic(fmt.Sprintf("bad split for %x with shift %v (%x, %x)", v, shift, hi, lo))
+	if hi&^0xfff000 == 0 && hi+lo<<shift == v {
+		return hi, lo, nil
 	}
-	return hi, lo, nil
+
+	return 0, 0, fmt.Errorf("%d cannot be split into valid hi/lo", v)
 }
 
 func (c *ctxt7) regoff(a *obj.Addr) int32 {
@@ -1969,28 +1976,28 @@ func (c *ctxt7) loadStoreClass(p *obj.Prog, lsc int, v int64) int {
 		if cmp(C_UAUTO8K, lsc) || cmp(C_UOREG8K, lsc) {
 			return lsc
 		}
-		if v >= 0 && v <= 0xfff000+0xfff<<1 && v&1 == 0 {
+		if v >= 0 && v <= 0xfff000+0xfff<<1 && (v&1 == 0 || v <= 0xfff+0xfff<<1) {
 			needsPool = false
 		}
 	case AMOVW, AMOVWU, AFMOVS:
 		if cmp(C_UAUTO16K, lsc) || cmp(C_UOREG16K, lsc) {
 			return lsc
 		}
-		if v >= 0 && v <= 0xfff000+0xfff<<2 && v&3 == 0 {
+		if v >= 0 && v <= 0xfff000+0xfff<<2 && (v&3 == 0 || v <= 0xfff+0xfff<<2) {
 			needsPool = false
 		}
 	case AMOVD, AFMOVD:
 		if cmp(C_UAUTO32K, lsc) || cmp(C_UOREG32K, lsc) {
 			return lsc
 		}
-		if v >= 0 && v <= 0xfff000+0xfff<<3 && v&7 == 0 {
+		if v >= 0 && v <= 0xfff000+0xfff<<3 && (v&7 == 0 || v <= 0xfff+0xfff<<3) {
 			needsPool = false
 		}
 	case AFMOVQ:
 		if cmp(C_UAUTO64K, lsc) || cmp(C_UOREG64K, lsc) {
 			return lsc
 		}
-		if v >= 0 && v <= 0xfff000+0xfff<<4 && v&15 == 0 {
+		if v >= 0 && v <= 0xfff000+0xfff<<4 && (v&15 == 0 || v <= 0xfff+0xfff<<4) {
 			needsPool = false
 		}
 	}
diff --git a/src/cmd/internal/obj/arm64/asm_test.go b/src/cmd/internal/obj/arm64/asm_test.go
index 6347bafc7e..bf00e4648b 100644
--- a/src/cmd/internal/obj/arm64/asm_test.go
+++ b/src/cmd/internal/obj/arm64/asm_test.go
@@ -46,8 +46,8 @@ func TestSplitImm24uScaled(t *testing.T) {
 		{
 			v:      0x1001,
 			shift:  0,
-			wantHi: 0x1000,
-			wantLo: 0x1,
+			wantHi: 0x2,
+			wantLo: 0xfff,
 		},
 		{
 			v:      0xffffff,
@@ -75,8 +75,8 @@ func TestSplitImm24uScaled(t *testing.T) {
 		{
 			v:      0x2002,
 			shift:  1,
-			wantHi: 0x2000,
-			wantLo: 0x1,
+			wantHi: 0x4,
+			wantLo: 0xfff,
 		},
 		{
 			v:      0xfffffe,
@@ -96,6 +96,11 @@ func TestSplitImm24uScaled(t *testing.T) {
 			wantErr: true,
 		},
 		{
+			v:       0x1000001,
+			shift:   1,
+			wantErr: true,
+		},
+		{
 			v:       0xfffffe,
 			shift:   2,
 			wantErr: true,
@@ -103,8 +108,8 @@ func TestSplitImm24uScaled(t *testing.T) {
 		{
 			v:      0x4004,
 			shift:  2,
-			wantHi: 0x4000,
-			wantLo: 0x1,
+			wantHi: 0x8,
+			wantLo: 0xfff,
 		},
 		{
 			v:      0xfffffc,
@@ -131,8 +136,8 @@ func TestSplitImm24uScaled(t *testing.T) {
 		{
 			v:      0x8008,
 			shift:  3,
-			wantHi: 0x8000,
-			wantLo: 0x1,
+			wantHi: 0x10,
+			wantLo: 0xfff,
 		},
 		{
 			v:      0xfffff8,
@@ -151,6 +156,127 @@ func TestSplitImm24uScaled(t *testing.T) {
 			shift:   3,
 			wantErr: true,
 		},
+		// Unshifted hi cases - hi <= 0xfff fits directly
+		{
+			v:      7,
+			shift:  3,
+			wantHi: 7,
+			wantLo: 0,
+		},
+		{
+			v:      0x8ff7,
+			shift:  3,
+			wantHi: 0xfff,
+			wantLo: 0xfff,
+		},
+		{
+			v:      0x7ff8,
+			shift:  3,
+			wantHi: 0,
+			wantLo: 0xfff,
+		},
+		{
+			v:      0xfff,
+			shift:  1,
+			wantHi: 1,
+			wantLo: 0x7ff,
+		},
+		{
+			v:      0xfff,
+			shift:  2,
+			wantHi: 3,
+			wantLo: 0x3ff,
+		},
+		{
+			v:      0xfff,
+			shift:  3,
+			wantHi: 7,
+			wantLo: 0x1ff,
+		},
+		{
+			v:      0x1ffe,
+			shift:  2,
+			wantHi: 2,
+			wantLo: 0x7ff,
+		},
+		{
+			v:      0x1ffe,
+			shift:  3,
+			wantHi: 6,
+			wantLo: 0x3ff,
+		},
+		{
+			v:      0x1fff,
+			shift:  1,
+			wantHi: 1,
+			wantLo: 0xfff,
+		},
+		{
+			v:      0x1fff,
+			shift:  2,
+			wantHi: 3,
+			wantLo: 0x7ff,
+		},
+		{
+			v:      0x1fff,
+			shift:  3,
+			wantHi: 7,
+			wantLo: 0x3ff,
+		},
+		{
+			v:      0x1001,
+			shift:  1,
+			wantHi: 1,
+			wantLo: 0x800,
+		},
+		{
+			v:      0x1001,
+			shift:  2,
+			wantHi: 1,
+			wantLo: 0x400,
+		},
+		{
+			v:      0x1001,
+			shift:  3,
+			wantHi: 1,
+			wantLo: 0x200,
+		},
+		{
+			v:      0x1000,
+			shift:  0,
+			wantHi: 0x1,
+			wantLo: 0xfff,
+		},
+		{
+			v:      0x8000,
+			shift:  3,
+			wantHi: 0x8,
+			wantLo: 0xfff,
+		},
+		{
+			v:      0xfff,
+			shift:  0,
+			wantHi: 0,
+			wantLo: 0xfff,
+		},
+		{
+			v:      0x1ffe,
+			shift:  1,
+			wantHi: 0,
+			wantLo: 0xfff,
+		},
+		{
+			v:      0x3ffc,
+			shift:  2,
+			wantHi: 0,
+			wantLo: 0xfff,
+		},
+		{
+			v:      0x10fef,
+			shift:  4,
+			wantHi: 0xfff,
+			wantLo: 0xfff,
+		},
 	}
 	for _, test := range tests {
 		hi, lo, err := splitImm24uScaled(test.v, test.shift)
@@ -179,6 +305,21 @@ func TestSplitImm24uScaled(t *testing.T) {
 			}
 		}
 	}
+
+	// Test the unshifted hi range specifically, including unaligned values.
+	// This exercises values where the unshifted path may be used.
+	for shift := 0; shift <= 3; shift++ {
+		maxUnshifted := int32(0xfff + 0xfff<<shift)
+		for v := int32(0); v <= maxUnshifted; v++ {
+			hi, lo, err := splitImm24uScaled(v, shift)
+			if err != nil {
+				t.Fatalf("splitImm24uScaled(%x, %x) failed: %v", v, shift, err)
+			}
+			if hi+lo<<shift != v {
+				t.Fatalf("splitImm24uScaled(%x, %x) = (%x, %x) is incorrect", v, shift, hi, lo)
+			}
+		}
+	}
 }
 
 // TestLarge generates a very large file to verify that large