C#でLOUDS構造を用いたPatriciaTrieを実装する

ピン留めされたアイテム

kx-ras

(ピン止め用)リポジトリ

kx-ras

簡潔ビットベクトルの話とか
参考

更新前の実装

書いたコードがさすがにひどすぎたのでリファクタリング＆最適化とか
あと実装時のTipを残しておく

kx-ras

BitArrayを廃止

BitOperations.PopCountなどのCPU命令を活用したいのでBitArrayを廃止する
BitVector32構造体は16ビットのセクションを作成できないのでushortのリストでLBSを持つことにする

BitVector.cs

public class BitVector
{
    const int SMALL_BLOCK_SPLIT_SIZE = 16;

-   private BitArray bitVector;
+   private List<ushort> bitVectorList;
}

kx-ras

ビット反転

通常LOUDSでは子があったら1を置き、子がなくなったら0を置く
しかし実装しやすいので0と1を入れ替えることにする(理由)

BitVector.cs

const int MAX_INDEX = 15;

- private ushort bitVector = 0;
+ private ushort bitVector = UInt16.MaxValue;

- if (value) bitVector |= (ushort)(1 << MAX_INDEX - bitIndex);
+ if (!value) bitVector &= (ushort)~(1 << MAX_INDEX - bitIndex);

kx-ras

ビット追加

ListのCapacityを一括で調整
どうせ最終的に大ブロックの要素1つ分LBSを切り上げないといけないので、インデックスがCapacityを超えたら逐次追加する

BitVector.cs

    const int BIG_BLOCK_SPLIT_SIZE = 1024;
    const int SMALL_BLOCK_SPLIT_SIZE = 16;
    
    const int SMALL_BLOCK_SIZE = BIG_BLOCK_SPLIT_SIZE / SMALL_BLOCK_SPLIT_SIZE; //16

    private int capacity;

    public void Add(bool value)
    {
        
        // 追加操作(略)

        listIndex++;
        if (listIndex >= capacity)
        {
            capacity += SMALL_BLOCK_SIZE; //大ブロックの要素一つ分capacityを追加
            bitVectorList.Capacity = capacity;
        }
    }

Listの空要素を埋める
ビルド時にcapacity分MaxValueで埋める必要がある
1回しか呼ばれないし定数時間なので愚直な実装で妥協した

public void FillToCapacity()
{
    while (bitVectorList.Count < capacity)
    {
        bitVectorList.Add(UInt16.MaxValue);
    }
}

(速くしようと思えば速くできる)

public void HackedFillToCapacity()
{
    int index = bitVectorList.Count;
    if (index == capacity) return;

    CollectionsMarshal.SetCount(bitVectorList, capacity);
    var span = CollectionsMarshal.AsSpan(bitVectorList);

    span.Slice(index, capacity - index).Fill(UInt16.MaxValue);
}

kx-ras

LSB vs MSB

実装していたらビット演算が面倒なのでMSBではなくLSB基準のほうが良いことが分かった

BitVector.cs

// Add
- if (!value) bitVector &= (ushort)~(1 << MAX_INDEX - bitIndex);
+ if (!value) bitVector &= (ushort)~(1 << bitIndex);

// Get
var index = position / SMALL_BLOCK_SPLIT_SIZE;
var bitpos = position & SMALL_BLOCK_SPLIT_SIZE;

- return (_bitArray[index] & (1 << (MAX_INDEX - bitpos))) != 0;
+ return (_bitArray[index] & (1 << bitpos)) != 0;

今気づいたけど変数名に統一性がない...

デバッグ出力用メソッドにも変更を加える

BitVector.cs

- builder.Append(Convert.ToString(_bitArray[i], 2).PadLeft(MAX_INDEX + 1, '0'));
+ builder.Append(Convert.ToString(_bitArray[i], 2).PadLeft(MAX_INDEX + 1, '0').Reverse().ToArray());

kx-ras

Rank1の実装

BitVector.cs

private int Rank1(int position)
{
    if ((uint)position > _size) return -1;

    int count = 0;
    int bigBlockIndex = position / BIG_BLOCK_SPLIT_SIZE;
    count += _bigBlock[bigBlockIndex];
    int smallBlockIndex = position % BIG_BLOCK_SPLIT_SIZE / SMALL_BLOCK_SPLIT_SIZE;
    count += _smallBlock[bigBlockIndex, smallBlockIndex];
    uint mask = (uint)(1 << (position % SMALL_BLOCK_SPLIT_SIZE)) - 1;
    count += BitOperations.PopCount(_bitArray[position / SMALL_BLOCK_SPLIT_SIZE] & mask);
    return count;
}

ここが重要

BitVector.cs

uint mask = (uint)(1 << (position % SMALL_BLOCK_SPLIT_SIZE)) - 1;
count += BitOperations.PopCount(_bitArray[position / SMALL_BLOCK_SPLIT_SIZE] & mask);

positionまでのビットを立てたマスクを作る

BitVector.cs

uint mask = (uint)(1 << (position % SMALL_BLOCK_SPLIT_SIZE)) - 1;

AND演算を行い立っているビット数を数える

BitVector.cs

count += BitOperations.PopCount(_bitArray[position / SMALL_BLOCK_SPLIT_SIZE] & mask);

kx-ras

Select1の実装(前編)

長い

BitVector.cs

public int Select1(int count)
{
    if ((uint)count > _size) return -1;
    if (count == 0) return -1;

    int position = 0;
    int remain = count;

    int left = -1;
    int right = _bigBlock.Length;
    while (right - left > 1)
    {
        int mid = (left + right) / 2;

        if (_bigBlock[mid] < remain) left = mid;
        else right = mid;
    }
    int bigIndex = left;
    position += bigIndex * BIG_BLOCK_SPLIT_SIZE;
    remain -= _bigBlock[bigIndex];

    left = -1;
    right = _smallBlock.GetLength(1);
    while (right - left > 1)
    {
        int mid = (left + right) / 2;

        if (_smallBlock[bigIndex, mid] < remain) left = mid;
        else right = mid;
    }
    int smallIndex = left;
    position += smallIndex * SMALL_BLOCK_SPLIT_SIZE;
    remain -= _smallBlock[bigIndex, smallIndex];

    var mask = _bitArray[position / SMALL_BLOCK_SPLIT_SIZE];
    for (remain--; remain > 0; remain--)
    {
        mask &= (ushort)(mask - 1);
    }

    return position + BitOperations.TrailingZeroCount(mask) + 1;
}

半分以上二分探索なのでいいとして、最後のところが重要

BitVector.cs

var mask = _bitArray[position / SMALL_BLOCK_SPLIT_SIZE];
for (remain--; remain > 0; remain--)
{
    mask &= (ushort)(mask - 1);
}

return position + BitOperations.TrailingZeroCount(mask) + 1;

ushortのLBSを持ってくる

BitVector.cs

var mask = _bitArray[position / SMALL_BLOCK_SPLIT_SIZE];

残り回数-1回右端の1を消す

BitVector.cs

for (remain--; remain > 0; remain--)
{
    mask &= (ushort)(mask - 1);
}

右から数えて何番目に1が来るかTZCNTで数えて返す

BitVector.cs

return position + BitOperations.TrailingZeroCount(mask) + 1;

kx-ras

 Select1の実装(後編)実装で何が起こっているか示す。

例としてLBS0010010100001001のSelect1(3)を考える。
残り回数-1回右端の1を消す

ここでは残り回数は3回なので2回右端の1を消す

1を引いてAND演算すると右端の1を消すことができる。
1回目LBS:   0010010100001001
LBS-1: 0010010100001000
AND  : 0010010100001000
2回目

1が別の位にあっても同様に消せる。LBS:   0010010100001000
LBS-1: 0010010100000111
AND  : 0010010100000000
TrailingZeroCountで右端の0を数える

CPU命令TZCNTを活用して右端の0を数える。

ここでは0010010100000000なので8が得られる。
1を足す。

よってLBS0010010100001001のSelect1(3)は9となってこれは正しいことが分かる。

kx-ras

GetBitの実装

ushortのリストでLBSを持っているが、特定のビットにアクセスしたいときは？

BitVector.cs

public bool GetBit(int position)
{
    var index = position / SMALL_BLOCK_SPLIT_SIZE;
    var bitPos = position % SMALL_BLOCK_SPLIT_SIZE;
    return (_bitArray[index] & (1 << bitPos)) != 0;
}

ビットを立ててAND演算するだけ

kx-ras

String vs ReadOnlyMemory<char>

StartsWith()でspan.ToString()しているのでもったいない、ということでstring[]ではなくReadOnlyMemory<char>で持って比較の際にnodekey.Spanで切り出すことで最適化する。

LOUDSTrie.cs

- private string[] keys = null!;
+ private ReadOnlyMemory<char>[] keys = null!;

// ~略~

- else if (nodeKey.StartsWith(querySpan.ToString(), StringComparison.Ordinal))
+ else if (nodeKey.StartsWith(querySpan, StringComparison.Ordinal))

kx-ras

現段階でのベンチマーク

分かってはいたけど検索に関してはほとんど差がない。実際はもっと辞書サイズが大きいのでアロケーションが減ったのは良いと思う。ビルド時のパフォーマンスがかなり改善された。

Legacy

| Method    | Mean     | Error    | StdDev   | Gen0   | Allocated |
|---------- |---------:|---------:|---------:|-------:|----------:|
| BuildTrie | 15.72 us | 0.119 us | 0.111 us | 1.0376 |   8.69 KB |

| Method             | Mean      | Error    | StdDev   | Gen0   | Allocated |       
|------------------- |----------:|---------:|---------:|-------:|----------:|       
| ExactMatchSearch   |  22.59 ns | 0.069 ns | 0.054 ns |      - |         - |       
| CommonPrefixSearch | 148.21 ns | 0.617 ns | 0.547 ns | 0.0343 |     288 B |       
| PredictiveSearch   | 411.13 ns | 3.609 ns | 3.375 ns | 0.0820 |     688 B |

現在

| Method    | Mean     | Error     | StdDev    | Gen0   | Allocated |
|---------- |---------:|----------:|----------:|-------:|----------:|
| BuildTrie | 9.649 us | 0.0801 us | 0.0750 us | 0.6714 |   5.55 KB |

| Method             | Mean      | Error    | StdDev   | Gen0   | Allocated |       
|------------------- |----------:|---------:|---------:|-------:|----------:|       
| ExactMatchSearch   |  23.56 ns | 0.055 ns | 0.046 ns |      - |         - |       
| CommonPrefixSearch | 150.34 ns | 0.950 ns | 0.842 ns | 0.0257 |     216 B |       
| PredictiveSearch   | 339.12 ns | 2.198 ns | 1.835 ns | 0.0763 |     640 B |

kx-ras

分からない

どう考えてもC#のプロが必要、初心者の自分では限界がある

kx-ras

Save/Load

実装するの面倒。(理由：Mozc辞書などではTに対応するのがValueTuple<int, int, int, string>みたいな~~気持ち悪い~~複雑な型になるため)

ということで脳死でMemoryPackを使うことにします。あとで実際に辞書を使って検証が必要。

LOUDSTrieIO.cs

public static class LOUDSTrieIO
{
    public static async ValueTask SaveTrieAsync<T>(LOUDSTrie<T> trie, string filePath)
    {
        using (var fs = new FileStream(filePath, FileMode.Create, FileAccess.Write, FileShare.None, 4096, true))
        {
            await MemoryPackSerializer.SerializeAsync(fs, trie);
        }
    }

    public static async ValueTask<LOUDSTrie<T>> LoadTrieAsync<T>(string filePath)
    {
        using (var fs = new FileStream(filePath, FileMode.Open, FileAccess.Read, FileShare.Read, 4096, true))
        {
            var result = await MemoryPackSerializer.DeserializeAsync<LOUDSTrie<T>>(fs);
            if (result is null) throw new InvalidDataException("Failed to load LOUDSTrie.");
            return result;
        }
    }
}

kx-ras

Mozc辞書をビルドしてみる

だいたい4秒くらいでビルドできるみたい...遅すぎ(前使ってたライブラリに比べれば十分高速だが)。
やはり通常のトライ木に一度変換しているからアロケーションもめちゃめちゃ大きいし時間もかかる。

1-passでパトリシアトライを構築するのは難しいので、後々余裕があるときに実装することにする。

| Method | Mean    | Error    | StdDev   | Gen0        | Gen1       | Gen2      | Allocated |
|------- |--------:|---------:|---------:|------------:|-----------:|----------:|----------:|
| Build  | 3.884 s | 0.0767 s | 0.0942 s | 120000.0000 | 87000.0000 | 6000.0000 |   1.16 GB |

kx-ras

Mozc辞書の保存

ビルド＆保存のベンチマーク。やっぱりMemoryPackはすごかった。

| Method | Mean    | Error    | StdDev   | Gen0        | Gen1       | Gen2      | Allocated |
|------- |--------:|---------:|---------:|------------:|-----------:|----------:|----------:|
| Build  | 4.106 s | 0.0760 s | 0.0747 s | 120000.0000 | 87000.0000 | 6000.0000 |   1.22 GB |

一方ファイルサイズは52899KBという結果に。Darts-Cloneの性能表を見るとやや多いのでもう少し減らしたい。

kx-ras

Mozc辞書の読み出し

ビルドして保存したファイルの読み出し。実用上問題ない十分な速度が出ている。

| Method | Mean     | Error   | StdDev  | Gen0       | Gen1       | Gen2      | Allocated |
|------- |---------:|--------:|--------:|-----------:|-----------:|----------:|----------:|
| Load   | 266.0 ms | 5.17 ms | 6.90 ms | 16333.3333 | 16000.0000 | 2333.3333 | 134.51 MB |

kx-ras

UTF8 vs UTF16

実験を行った。Mozc辞書の保存ではUTF8方式では54890KBとなりデータ量が大きくなってしまった。表層形の部分に漢字が多く含まれているという点がデータ量増加の原因だと思われる。Ngramモデルでも漢字等が多く含まれているため、とりあえずUTF16でデータを扱うことにする。

kx-ras

Tail圧縮

現在まで、ノードに対応するキーはchar[][]で持っていた。TAIL圧縮ではすべての文字を連結してchar[]で持つ。今回はパトリシアトライなのでキーが何文字かを表現するためのBitVectorと共に使用する。

今まで

keys:
""
""
"an"
"i"
"o"
"f"
"ne"
"u"
"f"
"r"
"t"

TAIL圧縮後

tailKeys:  aniofneufrt
tailBits:0110111101111

tailBitsを用意することでtailKeys[Select1(index)..Select1(index+1)]としてノードに対応するキーを読むことができる。

Mozc辞書のファイルサイズは52899KB➡49999KBと約3MBほど圧縮できた。ただchar[][]からchar[]にするだけでこんなに容量が変化するとは思わなかったので驚いた。

kx-ras

outsの実装

今まで、通常のトライを作成してからパトリシアトライをビルドしていたので値を持つ配列のインデックスが変わってしまう問題があった。そこでとりあえず対応するインデックスをint[]で持って、ノードがleafでない場合は-1を入れていた。
パトリシアトライのビルド時にノードと値を対応する順番で保存するようにして、ノードがleafかどうかをBitVectorで管理する。これにより容量を削減できる。

今まで

public int[] indexes { get; private set; } = null!;

// Build()実装箇所
this.values = baseKeysets.ToArray();
indexList.Add(current.index);
indexes = indexList.ToArray();

// 検索部分
if (index >= 0) return values[index];

outs実装後

public BitVector indexBits { get; private set; } = null!;

// Build()実装箇所
if (current.index >= 0)
{
    indexBitsBuilder.Add(false);
    valueList.Add(baseKeysets[current.index]);
}
else
{
    indexBitsBuilder.Add(true);
}
indexBits = indexBitsBuilder.Build();
values = valueList.ToArray();

// 検索部分
if (!indexBits.GetBit(nodeNumber)) return values[indexBits.Rank0(nodeNumber)];

Mozc辞書のファイルサイズは49999KB➡46848KBとさらに3MB圧縮できた。ただ圧縮できた代わりにビルド時間は少し増えた。

kx-ras

DAWGまで実装できれば良いが、とりあえずいいところまでできたのでここでクローズ。
時間があったら記事にしておきたい

このスクラップは18日前にクローズされました