fix: Fix arm64 assembly stack frame manipulation

arrow/compute/internal/kernels/cast_numeric_neon_arm64.s

@@ -10,13 +10,7 @@ TEXT ·_cast_type_numeric_neon(SB), $0-40
     MOVD len+32(FP), R4
 
 
-    // The Go ABI saves the frame pointer register one word below the 
-    // caller's frame. Make room so we don't overwrite it. Needs to stay 
-    // 16-byte aligned 
-    SUB $16, RSP
-    WORD $0xa9bf7bfd // stp    x29, x30, [sp, #-16]!
     WORD $0x7100181f // cmp    w0, #6
-    WORD $0x910003fd // mov    x29, sp
     BGT LBB0_17
 
     WORD $0x71000c1f // cmp    w0, #3
@@ -4450,9 +4444,6 @@ LBB0_892:
     WORD $0xb800452b // str    w11, [x9], #4
     BNE LBB0_892
 LBB0_893:
-    WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-    // Put the stack pointer back where it was 
-    ADD $16, RSP
     RET
 LBB0_894:
     WORD $0x927b6909 // and x9, x8, #0xffffffe0

arrow/math/int64_neon_arm64.s

@@ -10,13 +10,7 @@ TEXT ·_sum_int64_neon(SB), $0-24
 	MOVD	buf+0(FP), R0
 	MOVD	len+8(FP), R1
 	MOVD	res+16(FP), R2
-
-	// The Go ABI saves the frame pointer register one word below the 
-	// caller's frame. Make room so we don't overwrite it. Needs to stay 
-	// 16-byte aligned 
-	SUB $16, RSP
-	WORD $0xa9bf7bfd // stp x29, x30, [sp, #-16]!
-	WORD $0x910003fd // mov x29, sp
+
 	CBZ R1, LBB0_3
 	WORD $0xf1000c3f // cmp x1, #3
 	BHI LBB0_4
@@ -26,9 +20,6 @@ TEXT ·_sum_int64_neon(SB), $0-24
 LBB0_3:
 	WORD $0xaa1f03e9 // mov     x9, xzr
 	WORD $0xf9000049 // str     x9, [x2]
-	WORD $0xa8c17bfd // ldp     x29, x30, [sp], #16
-	// Put the stack pointer back where it was 
-	ADD $16, RSP
 	RET
 LBB0_4:
 	WORD $0x927ef428 // and     x8, x1, #0xfffffffffffffffc
@@ -59,8 +50,5 @@ LBB0_8:
 	BNE LBB0_8
 LBB0_9:
 	WORD $0xf9000049 // str     x9, [x2]
-	WORD $0xa8c17bfd // ldp     x29, x30, [sp], #16
-	// Put the stack pointer back where it was 
-	ADD $16, RSP
 	RET
 

arrow/math/uint64_neon_arm64.s

@@ -10,13 +10,7 @@ TEXT ·_sum_uint64_neon(SB), $0-24
 	MOVD	buf+0(FP), R0
 	MOVD	len+8(FP), R1
 	MOVD	res+16(FP), R2
-
-	// The Go ABI saves the frame pointer register one word below the 
-	// caller's frame. Make room so we don't overwrite it. Needs to stay 
-	// 16-byte aligned 
-	SUB $16, RSP
-	WORD $0xa9bf7bfd // stp x29, x30, [sp, #-16]!
-	WORD $0x910003fd // mov x29, sp
+
 	CBZ R1, LBB0_3
 	WORD $0xf1000c3f // cmp x1, #3
 	BHI LBB0_4
@@ -26,9 +20,6 @@ TEXT ·_sum_uint64_neon(SB), $0-24
 LBB0_3:
 	WORD $0xaa1f03e9 // mov     x9, xzr
 	WORD $0xf9000049 // str     x9, [x2]
-	WORD $0xa8c17bfd // ldp     x29, x30, [sp], #16
-	// Put the stack pointer back where it was 
-	ADD $16, RSP
 	RET
 LBB0_4:
 	WORD $0x927ef428 // and     x8, x1, #0xfffffffffffffffc
@@ -59,8 +50,5 @@ LBB0_8:
 	BNE LBB0_8
 LBB0_9:
 	WORD $0xf9000049 // str     x9, [x2]
-	WORD $0xa8c17bfd // ldp     x29, x30, [sp], #16
-	// Put the stack pointer back where it was 
-	ADD $16, RSP
 	RET
 

arrow/memory/memory_neon_arm64.s

@@ -11,14 +11,8 @@ TEXT ·_memset_neon(SB), $0-24
 	MOVD	len+8(FP), R1
 	MOVD	c+16(FP), R2
 
-	// The Go ABI saves the frame pointer register one word below the 
-	// caller's frame. Make room so we don't overwrite it. Needs to stay 
-	// 16-byte aligned 
-	SUB $16, RSP
-	WORD $0xa9bf7bfd // stp    x29, x30, [sp, #-16]!
 	WORD $0x8b010008 // add    x8, x0, x1
 	WORD $0xeb00011f // cmp    x8, x0
-	WORD $0x910003fd // mov    x29, sp
 	BLS	LBB0_7
 
 	WORD $0xf100803f // cmp    x1, #32
@@ -43,7 +37,4 @@ LBB0_6:
 	WORD $0xeb09011f // cmp    x8, x9
 	BNE	LBB0_6
 LBB0_7:
-	WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-	// Put the stack pointer back where it was 
-	ADD $16, RSP
 	RET

arrow/memory/memory_neon_arm64_test.go

@@ -0,0 +1,94 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file to you under the
+// Apache License, Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build arm64 && !noasm
+
+package memory_test
+
+import (
+	"os"
+	"os/exec"
+	"runtime/pprof"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+func TestNEONAssemblyTracebacks(t *testing.T) {
+	// This is a regression test for GH-983. _memset_neon used to modify the
+	// stack pointer in a way that the assembler didn't see, leading to
+	// broken tracebacks in CPU profiler samples when the function is
+	// running.
+
+	buf := make([]byte, 1<<20)
+	var traces []byte
+	duration := time.Second
+	d := t.TempDir()
+	for range 5 {
+		f, err := os.CreateTemp(d, "cpu.prof")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			f.Close()
+			t.Skipf("CPU profiling is already enabled: %v", err)
+		}
+
+		deadline := time.Now().Add(duration)
+		for time.Now().Before(deadline) {
+			memory.Set(buf, 0x1f)
+		}
+		pprof.StopCPUProfile()
+		if err := f.Close(); err != nil {
+			t.Fatal(err)
+		}
+
+		traces, err = exec.Command("go", "tool", "pprof", "-symbolize=none", "-traces", f.Name()).CombinedOutput()
+		if err != nil {
+			t.Fatalf("go tool pprof -traces: %v\n%s", err, traces)
+		}
+		if strings.Contains(string(traces), "memory._memset_neon") {
+			break
+		}
+		// Mitigate potential flakiness on CI runners
+		duration *= 2
+	}
+	if !strings.Contains(string(traces), "memory._memset_neon") {
+		t.Fatal("CPU profile contains no _memset_neon samples after 5 attempts")
+	}
+
+	want := strings.Join([]string{
+		"github.com/apache/arrow-go/v18/arrow/memory._memset_neon",
+		"github.com/apache/arrow-go/v18/arrow/memory.memory_memset_neon",
+		"github.com/apache/arrow-go/v18/arrow/memory.Set",
+		"github.com/apache/arrow-go/v18/arrow/memory_test.TestNEONAssemblyTracebacks",
+	}, "\n")
+
+	for trace := range strings.SplitSeq(string(traces), "-----------") {
+		if !strings.Contains(trace, "memory._memset_neon") {
+			continue
+		}
+
+		// pprof adds whitespace and marks inlined functions, neither of which is
+		// part of the call stack. Remove both before requiring adjacent frames.
+		trace = strings.ReplaceAll(trace, "(inline)", "")
+		trace = strings.Join(strings.Fields(trace), "\n")
+		if !strings.Contains(trace, want) {
+			t.Errorf("invalid _memset_neon traceback:\n%s", trace)
+		}
+	}
+}

internal/utils/min_max_neon_arm64.s

@@ -12,13 +12,7 @@ TEXT ·_int32_max_min_neon(SB), $0-32
 	MOVD    minout+16(FP), R2
 	MOVD    maxout+24(FP), R3
 
-	// The Go ABI saves the frame pointer register one word below the
-	// caller's frame. Make room so we don't overwrite it. Needs to stay
-	// 16-byte aligned
-	SUB $16, RSP
-	WORD $0xa9bf7bfd // stp    x29, x30, [sp, #-16]!
 	WORD $0x7100043f // cmp    w1, #1
-	WORD $0x910003fd // mov    x29, sp
 	BLT int32_early_exit
 
 	WORD $0x71001c3f // cmp    w1, #7
@@ -34,9 +28,6 @@ int32_early_exit:
 	WORD $0x52b0000b // mov    w11, #-2147483648
 	WORD $0xb900006b // str    w11, [x3]
 	WORD $0xb900004a // str    w10, [x2]
-	WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-	// Put the stack pointer back where it was
-	ADD $16, RSP
 	RET
 int32_neon:
 	WORD $0x927d7109 // and    x9, x8, #0xfffffff8
@@ -78,9 +69,6 @@ int32_scalar_loop:
 int32_done:
 	WORD $0xb900006b // str    w11, [x3]
 	WORD $0xb900004a // str    w10, [x2]
-	WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-	// Put the stack pointer back where it was
-	ADD $16, RSP
 	RET
 
 // func _uint32_max_min_neon(values unsafe.Pointer, length int, minout, maxout unsafe.Pointer)
@@ -91,13 +79,7 @@ TEXT ·_uint32_max_min_neon(SB), $0-32
 	MOVD    minout+16(FP), R2
 	MOVD    maxout+24(FP), R3
 
-	// The Go ABI saves the frame pointer register one word below the
-	// caller's frame. Make room so we don't overwrite it. Needs to stay
-	// 16-byte aligned
-	SUB $16, RSP
-	WORD $0xa9bf7bfd // stp    x29, x30, [sp, #-16]!
 	WORD $0x7100043f // cmp    w1, #1
-	WORD $0x910003fd // mov    x29, sp
 	BLT uint32_early_exit
 
 	WORD $0x71001c3f // cmp    w1, #7
@@ -113,9 +95,6 @@ uint32_early_exit:
 	WORD $0x1280000b // mov    w11, #-1
 	WORD $0xb900006a // str    w10, [x3]
 	WORD $0xb900004b // str    w11, [x2]
-	WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-	// Put the stack pointer back where it was
-	ADD $16, RSP
 	RET
 uint32_neon:
 	WORD $0x927d7109 // and    x9, x8, #0xfffffff8
@@ -157,9 +136,6 @@ uint32_scalar_loop:
 uint32_done:
 	WORD $0xb900006a // str    w10, [x3]
 	WORD $0xb900004b // str    w11, [x2]
-	WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-	// Put the stack pointer back where it was
-	ADD $16, RSP
 	RET
 
 // func _int64_max_min_neon(values unsafe.Pointer, length int, minout, maxout unsafe.Pointer)
@@ -170,13 +146,7 @@ TEXT ·_int64_max_min_neon(SB), $0-32
 	MOVD    minout+16(FP), R2
 	MOVD    maxout+24(FP), R3
 
-	// The Go ABI saves the frame pointer register one word below the
-	// caller's frame. Make room so we don't overwrite it. Needs to stay
-	// 16-byte aligned
-	SUB $16, RSP
-	WORD $0xa9bf7bfd // stp    x29, x30, [sp, #-16]!
 	WORD $0x7100043f // cmp    w1, #1
-	WORD $0x910003fd // mov    x29, sp
 	BLT int64_early_exit
 
 	WORD $0x2a0103e8 // mov    w8, w1
@@ -192,9 +162,6 @@ int64_early_exit:
 	WORD $0xd2f0000b // mov    x11, #-9223372036854775808
 	WORD $0xf900006b // str    x11, [x3]
 	WORD $0xf900004a // str    x10, [x2]
-	WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-	// Put the stack pointer back where it was
-	ADD $16, RSP
 	RET
 int64_neon:
 	WORD $0x927e7509 // and    x9, x8, #0xfffffffc
@@ -246,9 +213,6 @@ int64_scalar_loop:
 int64_done:
 	WORD $0xf900006b // str    x11, [x3]
 	WORD $0xf900004a // str    x10, [x2]
-	WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-	// Put the stack pointer back where it was
-	ADD $16, RSP
 	RET
 
 
@@ -260,13 +224,7 @@ TEXT ·_uint64_max_min_neon(SB), $0-32
 	MOVD    minout+16(FP), R2
 	MOVD    maxout+24(FP), R3
 
-	// The Go ABI saves the frame pointer register one word below the
-	// caller's frame. Make room so we don't overwrite it. Needs to stay
-	// 16-byte aligned
-	SUB $16, RSP
-	WORD $0xa9bf7bfd // stp    x29, x30, [sp, #-16]!
 	WORD $0x7100043f // cmp    w1, #1
-	WORD $0x910003fd // mov    x29, sp
 	BLT uint64_early_exit
 
 	WORD $0x71000c3f // cmp    w1, #3
@@ -282,9 +240,6 @@ uint64_early_exit:
 	WORD $0x9280000b // mov    x11, #-1
 	WORD $0xf900006a // str    x10, [x3]
 	WORD $0xf900004b // str    x11, [x2]
-	WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-	// Put the stack pointer back where it was
-	ADD $16, RSP
 	RET
 uint64_neon:
 	WORD $0x927e7509 // and    x9, x8, #0xfffffffc
@@ -336,7 +291,4 @@ uint64_scalar_loop:
 uint64_done:
 	WORD $0xf900006a // str    x10, [x3]
 	WORD $0xf900004b // str    x11, [x2]
-	WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-	// Put the stack pointer back where it was
-	ADD $16, RSP
 	RET

parquet/internal/bmi/bitmap_neon_arm64.s

@@ -8,13 +8,7 @@ TEXT ·_levels_to_bitmap_neon(SB), $0-32
     MOVD numLevels+8(FP), R1
     MOVD rhs+16(FP), R2
 
-    // The Go ABI saves the frame pointer register one word below the 
-    // caller's frame. Make room so we don't overwrite it. Needs to stay 
-    // 16-byte aligned 
-    SUB $16, RSP
-    WORD $0xa9bf7bfd // stp    x29, x30, [sp, #-16]!
     WORD $0x7100043f // cmp    w1, #1
-    WORD $0x910003fd // mov    x29, sp
     BLT LBB1_3
 
     WORD $0x71000c3f // cmp    w1, #3
@@ -82,9 +76,6 @@ LBB1_7:
     WORD $0xaa080168 // orr    x8, x11, x8
     BNE LBB1_7
 LBB1_8:
-    WORD $0xa8c17bfd // ldp    x29, x30, [sp], #16
-    // Put the stack pointer back where it was 
-    ADD $16, RSP
     MOVD R8, res+24(FP)
     RET