asm: Add handling for atomic operations

As it stands, we only had proper handling of atomic add operations.
`lock *(u32 *)(r1 + 0x1) += w2` was the only instruction that was
properly handled. Atomic operations have the same opcode, but
and are differentiated by the imm value. So far we have not been looking
at the imm value, so all atomic operations look like atomic adds to us.

This commit adds decoding for all current atomic operations. We handle
them similarly to how we handle ISAv4 instructions which use the
offset to further specify the instruction.

In #1193 we expanded the opcode from a u8 to u16, which is bigger than
the actual size. This allowed us to still represent functionally
different opcodes in Go, even tough the kernel uses other bits of the
instruction. So our opcode in Go is not identical to the opcode in the
kernel. We translate during marshalling and unmarshaling.

So far, we have only needed a few additional bits, but the atomic ops
need 9 bits of imm to fully encode all possibilities. Since 9 + 8 > 16
we have to grow the opcode to 32 bits.

During unmarshaling, we simply take the lower 9 bits of the imm shift
it left by 16 bits and or it with the opcode. During marshaling this
process is reversed.

Signed-off-by: Dylan Reimerink <dylan.reimerink@isovalent.com>

Author: dylandreimerink

asm/instruction.go

@@ -62,6 +62,10 @@ func (ins *Instruction) Unmarshal(r io.Reader, bo binary.ByteOrder, platform str
 
 	ins.Offset = int16(bo.Uint16(data[2:4]))
 
+	// Convert to int32 before widening to int64
+	// to ensure the signed bit is carried over.
+	ins.Constant = int64(int32(bo.Uint32(data[4:8])))
+
 	if ins.IsBuiltinCall() {
 		fn, err := BuiltinFuncForPlatform(platform, uint32(ins.Constant))
 		if err != nil {
@@ -93,12 +97,13 @@ func (ins *Instruction) Unmarshal(r io.Reader, bo binary.ByteOrder, platform str
 				ins.Offset = 0
 			}
 		}
+	} else if ins.OpCode.Class() == StXClass &&
+		ins.OpCode.Mode() == AtomicMode {
+		// For atomic ops, consider the lower 16 bits of the constant as part of the opcode
+		// it tells us which sort of atomic operation we are doing.
+		ins.OpCode |= (OpCode((ins.Constant << 16)) & atomicMask)
 	}
 
-	// Convert to int32 before widening to int64
-	// to ensure the signed bit is carried over.
-	ins.Constant = int64(int32(bo.Uint32(data[4:8])))
-
 	if !ins.OpCode.IsDWordLoad() {
 		return nil
 	}
@@ -171,6 +176,9 @@ func (ins Instruction) Marshal(w io.Writer, bo binary.ByteOrder) (uint64, error)
 			return 0, fmt.Errorf("extended ALU opcodes should have an .Offset of 0: %s", ins)
 		}
 		ins.Offset = newOffset
+	} else if atomic := ins.OpCode.Atomic(); atomic != InvalidAtomic {
+		ins.OpCode = ins.OpCode &^ atomicMask
+		ins.Constant = int64(atomic >> 16)
 	}
 
 	op, err := ins.OpCode.bpfOpCode()
@@ -382,8 +390,8 @@ func (ins Instruction) Format(f fmt.State, c rune) {
 			fmt.Fprintf(f, "dst: %s src: %s imm: %d", ins.Dst, ins.Src, ins.Constant)
 		case MemMode, MemSXMode:
 			fmt.Fprintf(f, "dst: %s src: %s off: %d imm: %d", ins.Dst, ins.Src, ins.Offset, ins.Constant)
-		case XAddMode:
-			fmt.Fprintf(f, "dst: %s src: %s", ins.Dst, ins.Src)
+		case AtomicMode:
+			fmt.Fprintf(f, "dst: %s src: %s off: %d", ins.Dst, ins.Src, ins.Offset)
 		}
 
 	case cls.IsALU():

asm/instruction_test.go

@@ -328,6 +328,82 @@ func (t testFDer) FD() int {
 	return int(t)
 }
 
+func TestAtomics(t *testing.T) {
+	rawInsns := []byte{
+		0xc3, 0x21, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, // lock *(u32 *)(r1 + 0x1) += w2
+		0xc3, 0x21, 0x01, 0x00, 0x50, 0x00, 0x00, 0x00, // lock *(u32 *)(r1 + 0x1) &= w2
+		0xc3, 0x21, 0x01, 0x00, 0xa0, 0x00, 0x00, 0x00, // lock *(u32 *)(r1 + 0x1) ^= w2
+		0xc3, 0x21, 0x01, 0x00, 0x40, 0x00, 0x00, 0x00, // lock *(u32 *)(r1 + 0x1) |= w2
+
+		0xc3, 0x21, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, // lock *(u64 *)(r1 + 0x1) += r2
+		0xc3, 0x21, 0x01, 0x00, 0x50, 0x00, 0x00, 0x00, // lock *(u64 *)(r1 + 0x1) &= r2
+		0xc3, 0x21, 0x01, 0x00, 0xa0, 0x00, 0x00, 0x00, // lock *(u64 *)(r1 + 0x1) ^= r2
+		0xc3, 0x21, 0x01, 0x00, 0x40, 0x00, 0x00, 0x00, // lock *(u64 *)(r1 + 0x1) |= r2
+
+		0xc3, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, // w0 = atomic_fetch_add((u32 *)(r1 + 0x0), w0)
+		0xc3, 0x01, 0x00, 0x00, 0x51, 0x00, 0x00, 0x00, // w0 = atomic_fetch_and((u32 *)(r1 + 0x0), w0)
+		0xc3, 0x01, 0x00, 0x00, 0xa1, 0x00, 0x00, 0x00, // w0 = atomic_fetch_xor((u32 *)(r1 + 0x0), w0)
+		0xc3, 0x01, 0x00, 0x00, 0x41, 0x00, 0x00, 0x00, // w0 = atomic_fetch_or((u32 *)(r1 + 0x0), w0)
+
+		0xdb, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, // r0 = atomic_fetch_add((u64 *)(r1 + 0x0), r0)
+		0xdb, 0x01, 0x00, 0x00, 0x51, 0x00, 0x00, 0x00, // r0 = atomic_fetch_and((u64 *)(r1 + 0x0), r0)
+		0xdb, 0x01, 0x00, 0x00, 0xa1, 0x00, 0x00, 0x00, // r0 = atomic_fetch_xor((u64 *)(r1 + 0x0), r0)
+		0xdb, 0x01, 0x00, 0x00, 0x41, 0x00, 0x00, 0x00, // r0 = atomic_fetch_or((u64 *)(r1 + 0x0), r0)
+
+		0xc3, 0x01, 0x00, 0x00, 0xe1, 0x00, 0x00, 0x00, // w0 = xchg32_32(r1 + 0x0, w0)
+		0xdb, 0x01, 0x00, 0x00, 0xe1, 0x00, 0x00, 0x00, // r0 = xchg_64(r1 + 0x0, r0)
+
+		0xc3, 0x11, 0x00, 0x00, 0xf1, 0x00, 0x00, 0x00, // w0 = cmpxchg32_32(r1 + 0x0, w0, w1)
+		0xdb, 0x11, 0x00, 0x00, 0xf1, 0x00, 0x00, 0x00, // r0 = cmpxchg_64(r1 + 0x0, r0, r1)
+	}
+
+	insns, err := AppendInstructions(nil, bytes.NewReader(rawInsns), binary.LittleEndian, platform.Linux)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	lines := []string{
+		"StXAtomicAddW dst: r1 src: r2 off: 1",
+		"StXAtomicAndW dst: r1 src: r2 off: 1",
+		"StXAtomicXorW dst: r1 src: r2 off: 1",
+		"StXAtomicOrW dst: r1 src: r2 off: 1",
+		"StXAtomicAddDW dst: r1 src: r2 off: 1",
+		"StXAtomicAndDW dst: r1 src: r2 off: 1",
+		"StXAtomicXorDW dst: r1 src: r2 off: 1",
+		"StXAtomicOrDW dst: r1 src: r2 off: 1",
+		"StXAtomicFetchAddW dst: r1 src: r0 off: 0",
+		"StXAtomicFetchAndW dst: r1 src: r0 off: 0",
+		"StXAtomicFetchXorW dst: r1 src: r0 off: 0",
+		"StXAtomicFetchOrW dst: r1 src: r0 off: 0",
+		"StXAtomicFetchAddDW dst: r1 src: r0 off: 0",
+		"StXAtomicFetchAndDW dst: r1 src: r0 off: 0",
+		"StXAtomicFetchXorDW dst: r1 src: r0 off: 0",
+		"StXAtomicFetchOrDW dst: r1 src: r0 off: 0",
+		"StXAtomicXchgW dst: r1 src: r0 off: 0",
+		"StXAtomicXchgDW dst: r1 src: r0 off: 0",
+		"StXAtomicCmpXchgW dst: r1 src: r1 off: 0",
+		"StXAtomicCmpXchgDW dst: r1 src: r1 off: 0",
+	}
+
+	for i, ins := range insns {
+		if want, got := lines[i], fmt.Sprint(ins); want != got {
+			t.Errorf("Expected %q, got %q", want, got)
+		}
+	}
+
+	// Marshal and unmarshal again to make sure the instructions are
+	// still valid.
+	var buf bytes.Buffer
+	err = insns.Marshal(&buf, binary.LittleEndian)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !bytes.Equal(buf.Bytes(), rawInsns) {
+		t.Error("Expected instructions to be equal after marshalling")
+	}
+}
+
 func TestISAv4(t *testing.T) {
 	rawInsns := []byte{
 		0xd7, 0x01, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00, // r1 = bswap16 r1
@@ -355,6 +431,16 @@ func TestISAv4(t *testing.T) {
 
 		0x3c, 0x31, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, // w1 s/= w3
 		0x9c, 0x42, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, // w2 s%= w4
+
+		0xd3, 0x10, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, // w0 = load_acquire((u8 *)(r1 + 0x0))
+		0xcb, 0x10, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, // w0 = load_acquire((u16 *)(r1 + 0x0))
+		0xc3, 0x10, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, // w0 = load_acquire((u32 *)(r1 + 0x0))
+		0xdb, 0x10, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, // r0 = load_acquire((u64 *)(r1 + 0x0))
+
+		0xd3, 0x21, 0x00, 0x00, 0x10, 0x01, 0x00, 0x00, // store_release((u8 *)(r1 + 0x0), w2)
+		0xcb, 0x21, 0x00, 0x00, 0x10, 0x01, 0x00, 0x00, // store_release((u16 *)(r1 + 0x0), w2)
+		0xc3, 0x21, 0x00, 0x00, 0x10, 0x01, 0x00, 0x00, // store_release((u32 *)(r1 + 0x0), w2)
+		0xdb, 0x21, 0x00, 0x00, 0x10, 0x01, 0x00, 0x00, // store_release((u64 *)(r1 + 0x0), r2)
 	}
 
 	insns, err := AppendInstructions(nil, bytes.NewReader(rawInsns), binary.LittleEndian, platform.Linux)
@@ -381,6 +467,14 @@ func TestISAv4(t *testing.T) {
 		"SModReg dst: r2 src: r4",
 		"SDivReg32 dst: r1 src: r3",
 		"SModReg32 dst: r2 src: r4",
+		"StXAtomicLdAcqB dst: r0 src: r1 off: 0",
+		"StXAtomicLdAcqH dst: r0 src: r1 off: 0",
+		"StXAtomicLdAcqW dst: r0 src: r1 off: 0",
+		"StXAtomicLdAcqDW dst: r0 src: r1 off: 0",
+		"StXAtomicStRelB dst: r1 src: r2 off: 0",
+		"StXAtomicStRelH dst: r1 src: r2 off: 0",
+		"StXAtomicStRelW dst: r1 src: r2 off: 0",
+		"StXAtomicStRelDW dst: r1 src: r2 off: 0",
 	}
 
 	for i, ins := range insns {

asm/load_store.go

@@ -1,6 +1,6 @@
 package asm
 
-//go:generate go run golang.org/x/tools/cmd/stringer@latest -output load_store_string.go -type=Mode,Size
+//go:generate go run golang.org/x/tools/cmd/stringer@latest -output load_store_string.go -type=Mode,Size,Atomic
 
 // Mode for load and store operations
 //
@@ -26,8 +26,42 @@ const (
 	MemMode Mode = 0x60
 	// MemSXMode - load from memory, sign extension
 	MemSXMode Mode = 0x80
-	// XAddMode - add atomically across processors.
-	XAddMode Mode = 0xc0
+	// AtomicMode - add atomically across processors.
+	AtomicMode Mode = 0xc0
+)
+
+const atomicMask OpCode = 0x01ff_0000
+
+type Atomic uint32
+
+const (
+	InvalidAtomic Atomic = 0xffff_ffff
+	// AddAtomic - add src to memory address dst atomically
+	AddAtomic Atomic = Atomic(Add) << 16
+	// FetchAddAtomic - add src to memory address at dst atomically, store old value in src
+	FetchAddAtomic Atomic = AddAtomic | Fetch
+	// AndAtomic - bitwise AND src with memory address at dst atomically
+	AndAtomic Atomic = Atomic(And) << 16
+	// FetchAndAtomic - bitwise AND src with memory address at dst atomically, store old value in src
+	FetchAndAtomic Atomic = AndAtomic | Fetch
+	// OrAtomic - bitwise OR src with memory address at dst atomically
+	OrAtomic Atomic = Atomic(Or) << 16
+	// FetchOrAtomic - bitwise OR src with memory address at dst atomically, store old value in src
+	FetchOrAtomic Atomic = OrAtomic | Fetch
+	// XorAtomic - bitwise XOR src with memory address at dst atomically
+	XorAtomic Atomic = Atomic(Xor) << 16
+	// FetchXorAtomic - bitwise XOR src with memory address at dst atomically, store old value in src
+	FetchXorAtomic Atomic = XorAtomic | Fetch
+	// FetchAtomicOp - load the old value into the src register
+	Fetch Atomic = 0x0001_0000
+	// XchgAtomic - atomically exchange the old value with the new value
+	XchgAtomic Atomic = 0x00e0_0000 | Fetch
+	// CmpXchgAtomic - atomically compare and exchange the old value with the new value
+	CmpXchgAtomic Atomic = 0x00f0_0000 | Fetch
+	// LdAcqAtomic - atomically load with acquire semantics
+	LdAcqAtomic Atomic = 0x0100_0000
+	// StRelAtomic - atomically store with release semantics
+	StRelAtomic Atomic = 0x0110_0000
 )
 
 // Size of load and store operations
@@ -210,16 +244,109 @@ func StoreImm(dst Register, offset int16, value int64, size Size) Instruction {
 	}
 }
 
+func AtomicOp(size Size, atomic Atomic, fetch bool) OpCode {
+	switch atomic {
+	case AddAtomic, AndAtomic, OrAtomic, XorAtomic:
+		if fetch {
+			atomic |= Fetch
+		}
+
+		switch size {
+		case Byte, Half:
+			// 8-bit and 16-bit atomic copy-modify-write atomics are not supported
+			return InvalidOpCode
+		}
+	}
+
+	return OpCode(StXClass).SetMode(AtomicMode).SetSize(size).SetAtomic(atomic)
+}
+
 // StoreXAddOp returns the OpCode to atomically add a register to a value in memory.
 func StoreXAddOp(size Size) OpCode {
-	return OpCode(StXClass).SetMode(XAddMode).SetSize(size)
+	return AtomicOp(size, AddAtomic, false)
 }
 
-// StoreXAdd atomically adds src to *dst.
+// AtomicAnd atomically adds src to *(dst + offset), storing the old value in src if fetch is true.
+func AtomicAdd(dst, src Register, size Size, offset int16, fetch bool) Instruction {
+	return Instruction{
+		OpCode: AtomicOp(size, AddAtomic, fetch),
+		Dst:    dst,
+		Src:    src,
+		Offset: offset,
+	}
+}
+
+// Alias for AtomicAdd
 func StoreXAdd(dst, src Register, size Size) Instruction {
+	return AtomicAdd(dst, src, size, 0, false)
+}
+
+// AtomicAnd atomically bitwise ANDs src with *(dst + offset), storing the old value in src if fetch is true.
+func AtomicAnd(dst, src Register, size Size, offset int16, fetch bool) Instruction {
+	return Instruction{
+		OpCode: AtomicOp(size, AndAtomic, fetch),
+		Dst:    dst,
+		Src:    src,
+		Offset: offset,
+	}
+}
+
+// AtomicOr atomically bitwise ORs src with *(dst + offset), storing the old value in src if fetch is true.
+func AtomicOr(dst, src Register, size Size, offset int16, fetch bool) Instruction {
+	return Instruction{
+		OpCode: AtomicOp(size, OrAtomic, fetch),
+		Dst:    dst,
+		Src:    src,
+		Offset: offset,
+	}
+}
+
+// AtomicXor atomically bitwise XORs src with *(dst + offset), storing the old value in src if fetch is true.
+func AtomicXor(dst, src Register, size Size, offset int16, fetch bool) Instruction {
+	return Instruction{
+		OpCode: AtomicOp(size, XorAtomic, fetch),
+		Dst:    dst,
+		Src:    src,
+		Offset: offset,
+	}
+}
+
+// AtomicXchg atomically exchanges src with *(dst + offset), storing the old value in src if fetch is true.
+func AtomicXchg(dst, src Register, size Size, offset int16, fetch bool) Instruction {
+	return Instruction{
+		OpCode: AtomicOp(size, XchgAtomic, fetch),
+		Dst:    dst,
+		Src:    src,
+		Offset: offset,
+	}
+}
+
+// AtomicCmpXchg atomically compares src with *(dst + offset), storing the old value in src if fetch is true.
+func AtomicCmpXchg(dst, src Register, size Size, offset int16, fetch bool) Instruction {
 	return Instruction{
-		OpCode: StoreXAddOp(size),
+		OpCode: AtomicOp(size, CmpXchgAtomic, fetch),
 		Dst:    dst,
 		Src:    src,
+		Offset: offset,
+	}
+}
+
+// AtomicStRel atomically stores src in *(dst + offset) with release semantics.
+func AtomicStRel(dst, src Register, size Size, offset int16) Instruction {
+	return Instruction{
+		OpCode: AtomicOp(size, StRelAtomic, false),
+		Dst:    dst,
+		Src:    src,
+		Offset: offset,
+	}
+}
+
+// AtomicLdAcq atomically loads *(dst + offset) into src with acquire semantics.
+func AtomicLdAcq(dst, src Register, size Size, offset int16) Instruction {
+	return Instruction{
+		OpCode: AtomicOp(size, LdAcqAtomic, false),
+		Dst:    dst,
+		Src:    src,
+		Offset: offset,
 	}
 }