一种快速取模算法

2023-04-25 ⏳4.6分钟(1.8千字) g

最近业务系统使用map[uint64]float64保存某种业务数据,大约有两亿多条,但实际内存占用却将近10个G。团队最终设计了一种自定义数据结构1来解决内存消耗问题。在解决问题的过程各我们考查了 SwissTable 这种新哈希表2。虽然没能解决我们的问题,却给出了一种新的快速取模的算法,很有启发意义。本文就基于 Daniel 的文章3为大家分享这种取模算法。

传统的取模运算需要用到除法,所以会消耗很多指令周期。为了加快运算速度,常规的哈希表实现中通常把槽位的数量指定为二的幂次。如果N为二的幂次,那么x % N 就可以简化为:

x % N = x & (N - 1)

这样就能避开使用取模操作,提高运算速度。

但这也有一个很大的副作用,即哈希表的槽位数量最只能是二的幂次。这在某个情况下会浪费很多内存。比如假设有 1025 条数据。如果不考虑冲突,那么需要分配 2048 个槽位,因为 1025 大于 1024,下一个二的幂次只能取 2048 了。这就浪费了将近一半的内存。

可能有人会说这种浪费可以接受,因为后续插入新数据的时候就不需要扩容了。这在通常情况下是成立的。但我们的业务场景并非如此。我们的业务系统只会在启动时加载数据,加载完成后再也不会改动。后续有更新会生成新的 map 然后切换。所以 go 语言内置 map 的这种扩容行为非但帮不了我们,反而会拖后腿。

这就有点左右为难了。用二的幂次浪费内存,不用二的幂次浪费CPU。就没有两全其美的办法吗?还真就有了。但该算法也有自己的限制,我们后面会说。先讲算法本身。

假设x和N都是32位无符号整数,而且x在[0,2^32)范围内均匀分布,那么可以用如下函数实现快速「取模」运算4。

func fastMod(x, n uint32) uint32 {
  return uint32((uint64(x) * uint64(n)) >> 32)
}

写成数学形式为x*N÷232x*N \div 2^32。从指令上看只做了一次乘法和一次右移运算,比取模运算要快得多!但计算的结果均匀吗?确实是均匀的,至少不会比%效果差。现在给出证明。

由前提条件可知,x在[0,2^32)范围内均匀分布。所以x*N在[0,N*2^32)范围内均匀分布。我们把[0,N*2^32)分段并做如下映射:

范围 结果
[0*2^32,1*2^32) 0
[1*2^32,2*2^32) 1
[2*2^32,3*2^32) 2
… …
[(N-1)*2^32,N*2^32) N-1

也就是说对于落到每一段内的x,只需要右移32位,也就是除以2322^32,就能得到右边的结果。

现在的问题就是需要证明对于任意x,这种映射是均匀的。

对于每一段[k*2^32,(k+1)*2^32),我们需要将其分成多段,每一段的长度为N。有多少段,就代表本区间内有多少个x*N这样的数字。

分两种情况讨论。

如果k*2^32能被N整除,则意味着区间左边第一个值为某个x*N。整个区间内最多有 ⌈232÷N⌉\lceil 2^32 \div N \rceil 个小段。最后一段长度可以小于N。虽然该区间没有包含最后一个完整的小段,那该小段对应的x*N却落到了这个区间。所以这里使用了取天棚运算。也就是说,在这种情况下,落入本段的x*N最多有 ⌈232÷N⌉\lceil 2^32 \div N \rceil 个。

如果k*2^32不能被N整除,那么就意味着该区间两端各有一个不完整的小段。也就只能包含 ⌊232÷N⌋\lfloor 2^32 \div N \rfloor 个x*N。

综上所述,对于每个区间[k*2^32,(k+1)*2^32),包含形如x*N的数字的个数是 ⌈232÷N⌉\lceil 2^32 \div N \rceil 或者 ⌊232÷N⌋\lfloor 2^32 \div N \rfloor。所以是均匀的。

其实也可以简单理解。我们把x*N/232x*N/2^32改写为x/232*Nx/2^32*N。因为x在[0,2^32)内均匀分布,所以x/232x/2^32在[0,1)内均匀分布。所以[0,1)*N肯定在[0,N)内均匀分布。

当然了,实践是检验真理的唯一准则。我们最好还是通过代码检测一下新算法是否均匀。

检测思路是批量生成随机序列,然后选定某种哈希算法将其转换成uint64整数,再使用新的取模运算计算结果。最后给每个结果计数,然后跟传统取模运算作比较。

取模代码如下:

func fastMod(x, n uint32) uint32 {
  return uint32((uint64(x) * uint64(n)) >> 32)
}

func normalMod(x, n uint32) uint32 {
  return x % n
}

计数代码:

n := uint32(100000)

m1 := map[uint32]uint32{}
m2 := map[uint32]uint32{}

m := 1000000
var buf [100]byte
for i := 0; i < m; i++ {
   N, _ := rand.Read(buf[:])
   v := cityhash.CityHash64(buf[:N], uint32(N))
   a := fastMod(uint32(v), n)
   m1[a]++
   b := normalMod(uint32(v), n)
   m2[b]++
}

这里我用了 cityhash 计算哈希值,随机序列使用 crypto/rand 读取。

最后计算得分:

func score(m map[uint32]uint32, n uint32) float64 {
  var s uint32
  for _, v := range m {
    if v > n {
      v = v - n
    } else {
      v = n - v
    }
    s += v * v
  }

  return float64(s)
}

得分越高则证明分布越不均匀。

输出结果的时候以传统取模结果为准,计算差值百分比:

s1 := score(m1, uint32(m)/n)
s2 := score(m2, uint32(m)/n)
fmt.Printf("%.2f\n", (float64(s1)/float64(s2)-1)*100)

结果为正则说明新算法效果差,为负则说明效果好。值越大效果越明显。计算表明,新旧算法互有胜负,但差值一般都在 5% 以内。极端情况下可能会达到 10%。

最后还得验证它们的性能差异。这时需要添加二的幂次基准函数:

func fastMod2(x, n uint32) uint32 {
  return x & (n - 1)
}

然后是 benchmark 代码:

func BenchmarkMod(b *testing.B) {
  n := rand.Uint32() + 1000000
  u := rand.Uint32() + 1000000
  b.Run("fast2", func(b *testing.B) {
    for i := 0; i < b.N; i++ {
      fastMod2(u, n)
    }
  })
  b.Run("fast", func(b *testing.B) {
    for i := 0; i < b.N; i++ {
      fastMod(u, n)
    }
  })
  b.Run("normal", func(b *testing.B) {
    for i := 0; i < b.N; i++ {
      normalMod(u, n)
    }
  })
}

在我的电脑上运行结果如下:

goos: darwin
goarch: amd64
pkg: baz
cpu: Intel(R) Core(TM) i7-8557U CPU @ 1.70GHz
BenchmarkMod/fast2-8            1000000000               0.2354 ns/op
BenchmarkMod/fast-8             1000000000               0.2341 ns/op
BenchmarkMod/normal-8           1000000000               0.4591 ns/op
PASS
ok      baz     1.042s

新的取模算法跟用二的幂次快速取余速度不相上下,性能是传统取模运算的两倍。而 Daniel 说在他的设备上新算法是传统算法的四倍。

现在总结一下新取模算法的特点。

对于N,不能超过2322^32,大约四十亿。对于x,则可以只取低32位,一般的哈希函数都会保证高32位和低32位都均匀分布。

总得来说这是一种不错的方案。只要哈希函数够均匀,只要曹位不超过四十亿,就可以实现对任意槽位数的快速取模。