From 6699ebd34f46677a7db2859665c983ded27adbee Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Petr=20Balv=C3=ADn?= <opensource@petrbalvin.org>
Date: Sat, 29 Aug 2026 10:42:15 +0200
Subject: [PATCH] feat(asm): add prefetch hint encoding

---
 asm/assemble_test.go | 34 ++++++++++++++++++++++++++++++++++
 asm/encode.go        | 26 ++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/asm/assemble_test.go b/asm/assemble_test.go
index f4cff9c..e8cefcb 100644
--- a/asm/assemble_test.go
+++ b/asm/assemble_test.go
@@ -317,3 +317,37 @@ end:
 		t.Errorf("jump-folding mismatch:\n got:  %s\n want: %s", hexBytes(code), hexBytes(want))
 	}
 }
+
+func TestAssemblePrefetch(t *testing.T) {
+	fn := firstText(t, `
+#include "textflag.h"
+TEXT ·pf(SB), NOSPLIT, $0
+	PREFETCHNTA (AX)
+	PREFETCHT0 (BX)
+	PREFETCHT1 8(CX)
+	PREFETCHT2 -1(AX)(R12*1)
+	RET
+`)
+	code, _, err := Assemble(fn)
+	if err != nil {
+		t.Fatalf("Assemble: %v", err)
+	}
+	got := strings.Join(disasm(t, code), "\n")
+	want := strings.Join([]string{
+		"prefetchnta zmmword ptr [rax]",
+		"prefetcht0 zmmword ptr [rbx]",
+		"prefetcht1 zmmword ptr [rcx+0x8]",
+		"prefetcht2 zmmword ptr [rax+r12-0x1]",
+		"ret",
+	}, "\n")
+	if got != want {
+		t.Errorf("prefetch disassembly mismatch:\n got:\n%s\n want:\n%s", got, want)
+	}
+	// Byte-level expectations: 0F 18 with the variant in the reg field.
+	if hex := hexBytes(code[:3]); hex != "0f 18 00" {
+		t.Errorf("PREFETCHNTA bytes: got %s, want 0f 18 00", hex)
+	}
+	if hex := hexBytes(code[3:6]); hex != "0f 18 0b" {
+		t.Errorf("PREFETCHT0 bytes: got %s, want 0f 18 0b", hex)
+	}
+}
diff --git a/asm/encode.go b/asm/encode.go
index 333d87e..e5dba0d 100644
--- a/asm/encode.go
+++ b/asm/encode.go
@@ -114,6 +114,8 @@ func (e *enc) encode(mnem string, ops []Operand) error {
 		return e.encodeCount(base, ops, size)
 	case "BSWAP":
 		return e.encodeBswap(ops, size)
+	case "PREFETCHNTA", "PREFETCHT0", "PREFETCHT1", "PREFETCHT2":
+		return e.encodePrefetch(base, ops)
 	case "MOVBLZX", "MOVBQZX", "MOVWLZX", "MOVWQZX", "MOVWLSX", "MOVLQSX":
 		return e.encodeMovExtend(base, ops)
 	case "CVTSL2SD", "CVTSQ2SD":
@@ -124,6 +126,30 @@ func (e *enc) encode(mnem string, ops []Operand) error {
 	return fmt.Errorf("unsupported instruction %q", mnem)
 }
 
+// encodePrefetch emits the 0F 18 /r prefetch hints: the reg field selects
+// the locality (NTA=0, T0=1, T1=2, T2=3) and the single operand is memory.
+func (e *enc) encodePrefetch(base string, ops []Operand) error {
+	if len(ops) != 1 {
+		return fmt.Errorf("%s expects one memory operand", base)
+	}
+	m, ok := ops[0].(Mem)
+	if !ok {
+		return fmt.Errorf("%s requires a memory operand", base)
+	}
+	i := newInstr(0, []byte{0x0F, 0x18})
+	if err := setMem(i, prefetchVariant[base], m); err != nil {
+		return err
+	}
+	return e.emit(i)
+}
+
+var prefetchVariant = map[string]int{
+	"PREFETCHNTA": 0,
+	"PREFETCHT0":  1,
+	"PREFETCHT1":  2,
+	"PREFETCHT2":  3,
+}
+
 // splitSize separates a trailing B/W/L/Q size suffix from the mnemonic.
 func splitSize(upper string) (base string, size int) {
 	if upper == "" {