IT博客汇
  • 首页
  • 精华
  • 技术
  • 设计
  • 资讯
  • 扯淡
  • 权利声明
    • 登录 注册

    128位整数的原子操作

    smallnest发表于 2024-06-16 03:45:28
    love 0

    我们已经知道，标准库中的 atomic 针对 int32/uint32、int64/uint64 提供了原子操作的方法和函数，但是如果针对 128 bit 的整数呢？

    当然使用128 bit 整数的原子操作的场景可能比较少，也不会有太多人有这个需求，但是如果我们需要对几个 32 bit、64 bit 变量进行原子操作吗， atomic128 可能就很有用。

    tmthrgd/atomic128 在几年前提供了 atomic 128 的实验性功能，最后放弃了，但是他提供了一个思路，可以使用 CMPXCHG16B 指令为 AMD 64 架构的CPU 提供 atomic 128 功能。

    CAFxX/atomic128 fork 了上面的项目，继续维护，还是使用 CMPXCHG16B 指令，只为 AMD 64 架构提供原子操作。

    首先我们看看它的功能然后再看一看它的实现，最后我们思路发散一下，看看使用 AVX 为 128 bit 甚至更多 bit 的整数提供原子操作是否可行。

    atomic128 的方法

    Package atomic128 实现了对 128 bit值的原子操作。在可能的情况下（例如，在支持 CMPXCHG16B 的 amd64 处理器上），它会自动使用 CPU 的原生特性来实现这些操作；否则，它会回退到基于互斥锁（mutexes）的方法。

    Go 的基本整数中不包含 int128/uint128，所以这个库先定义了一个 Int128 的类型：

    1
    2
    3
    4
    type Uint128 struct {
    	d [3]uint64
    	m sync.Mutex
    }

    然后类似标准库 atomic 中对各种整数的操作，它也提供了类似的方法：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    func AddUint128(ptr *Uint128, incr [2]uint64) [2]uint64
    func CompareAndSwapUint128(ptr *Uint128, old, new [2]uint64) bool
    func LoadUint128(ptr *Uint128) [2]uint64
    func StoreUint128(ptr *Uint128, new [2]uint64)
    func SwapUint128(ptr *Uint128, new [2]uint64) [2]uint64
    func OrUint128(ptr *Uint128, op [2]uint64) [2]uint64
    func AndUint128(ptr *Uint128, op [2]uint64) [2]uint64
    func XorUint128(ptr *Uint128, op [2]uint64) [2]uint64

    可以看到，除了正常的 Add、CAS、Load、Store、Swap 函数，还贴心的提供了 Or、And、Xor 三个位操作的函数。

    下面是一个简单的例子:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    n := &atomic128.Uint128{}
    v := atomic128.LoadUint128(n) // [2]uint64{0, 0}
    		
    atomic128.StoreUint128(n, [2]uint64{1, ^uint64(0)})
    v = atomic128.LoadUint128(n) // [2]uint64{1, ^uint64(0)}
    v = AddUint128(n, [2]uint64{2, 40})
    v = atomic128.LoadUint128(n) // [2]uint64{3, 40}
    v = atomic128.SwapUint128(n, [2]uint64{4, 50})
    v = atomic128.LoadUint128(n) // [2]uint64{4, 50}
    v = atomic128.CompareAndSwapUint128(n, [2]uint64{4, 50}, [2]uint64{5, 60})
    v = atomic128.LoadUint128(n) // [2]uint64{5, 60}
    v = atomic128.OrUint128(n, [2]uint64{0, 0})
    v = atomic128.LoadUint128(n) // [2]uint64{5, 60}

    atomic128 的实现

    聪明的你也许看到Uint128的定义的时候就会感觉有一点不对劲，为啥128bit的整数要用3个64bit的整数来表示呢？ 2个Uint64不就够了吗？

    这是为了保证128位对齐，类似的技术在Go 1.20之前的WaitGroup中也有使用。进一步了解可以查看：

    • https://go101.org/article/memory-layout.html
    • https://pkg.go.dev/sync/atomic#pkg-note-BUG

    通过包含三个Uint64元素的数组，我们总能通过下面的方法得到128位对齐的地址：

    1
    2
    3
    4
    5
    6
    func addr(ptr *Uint128) *[2]uint64 {
    	if (uintptr)((unsafe.Pointer)(&ptr.d[0]))%16 == 0 { // 指针已经128位对齐
    		return (*[2]uint64)((unsafe.Pointer)(&ptr.d[0]))
    	}
    	return (*[2]uint64)((unsafe.Pointer)(&ptr.d[1])) // 必然ptr.d[1]是128位对齐的 (AMD64架构)
    }

    通过变量useNativeAmd64判断CPU是否支持CMPXCHG16B指令：

    1
    2
    3
    func init() {
    	useNativeAmd64 = cpuid.CPU.Supports(cpuid.CX16)
    }

    如果不支持，回退到使用Mutex实现一个低效的atomic 128bit原子操作:

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    func CompareAndSwapUint128(ptr *Uint128, old, new [2]uint64) bool {
    	if runtime.GOARCH == "amd64" && useNativeAmd64 {
    		return compareAndSwapUint128amd64(addr(ptr), old, new)
    	}
        // 不支持CMPXCHG16B指令，使用Mutex
    	ptr.m.Lock()
    	v := load(ptr)
    	if v != old {
    		ptr.m.Unlock()
    		return false
    	}
    	store(ptr, new)
    	ptr.m.Unlock()
    	return true
    }

    如果支持CMPXCHG16B指令，直接调用compareAndSwapUint128amd64函数：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    TEXT ·compareAndSwapUint128amd64(SB),NOSPLIT,$0
    	MOVQ addr+0(FP), BP
    	MOVQ old+8(FP), AX
    	MOVQ old+16(FP), DX
    	MOVQ new+24(FP), BX
    	MOVQ new+32(FP), CX
    	LOCK
    	CMPXCHG16B (BP)
    	SETEQ swapped+40(FP)
    	RET

    主要依赖CMPXCHG16B实现。

    CMPXCHG16B是一条X86体系结构中的指令,全称为"Compare and Exchange 16 Bytes"。它用于原子地比较和交换16个字节(128位)的内存区域。
    这条指令的作用是:

    • 将要比较的16个字节的内存值加载到一个寄存器中。
    • 将要写入的16个字节的值加载到另一个寄存器中。
    • 比较内存中的值和第一个寄存器中的值是否相等。
    • 如果相等,则用第二个寄存器中的值覆盖内存中的值。
    • 根据比较结果,设置相应的标志位。

    思路发散

    当前很多号称性能优化的库，可能会使用SIMD指令集来提高性能，比如AVX、SSE等。那么，我们是否可以使用AVX指令集来实现对128位整数甚至256、512位整数的原子操作呢？

    有一篇很好的文章介绍了这方面的探索:Aligned AVX loads and stores are atomic。

    各家处理器手册中并没有为AVX指令集提供原子性的担保。The AMD64 Architecture Programmer’s Manual只是保证了内存操作最大8个字节，CMPXCHG16B是原子的。The Intel® 64 and IA-32 Architectures Software Developer’s Manual也做了类似的保证。此外，Intel手册明确指出AVX指令没有任何原子性保证。

    这篇文章的作者做了实验，得出下面的结论:

    尽管看起来对齐的 128 位操作室原子的，但是 CPU 提供商没有提供担保，我们还是使用 CMPXCHG16B 指令保险。



沪ICP备19023445号-2号
友情链接