LINQのインターフェースでGPUを使いたい！（GPGPU）

# 前書き

この記事は、2023のUnityアドカレの12/24の記事です。
今年は、完走賞に挑戦してみたいと思います。Qiita君ぬい欲しい！

https://qiita.com/advent-calendar/2023/present-calendar#qiita_award

# TL;DR;

PackageMaangerでインストール

```
https://github.com/wallstudio/TechArticleSamples.git?path=GpuLinq/Packages/GpuLinq
```

使い方

```csharp
var gpu = source
    .AsGpuEnumerable() // switch GPGPU
    .Where(x => (uint)x % 2 == 0) // unorderd
    .Select(x => x * 2)
    .ToArray();
```

# はじめに

GPUのアーキテクチャは、多数の同じようなデータ列に対し、一律の処理をかけることに長けています。このタスク、C#でなじみのある機能に似ていませんか？

そう、LINQですね！

というわけで、今回はGPGPUでコレクションを超並列に処理するLINQ風のライブラリを作り、紹介したいと思います。

# LINQの構造

LINQは、チェーンで積まれているときには各アイテムは評価されません。

```csharp
var e = source.Select(x => x * 2);
// まだeの中身は評価されていない
// source, x => x * 2 のセットとして記録されている

e.ToArray(); // ここで評価される
```

イメージとしてはこんな感じです。`GetEnumerator`が呼び出されたときはじめてsourceへの読み取りを始めます。`source.GetEnumerator`を呼び出すということです。つまり、一つ前につながれたコレクションの評価を始めます。

```csharp
static IEnumerable<U> Select<T, U>(this IEnumerable<T> source, Func<T, U> filter)
{
    return new SelectEnumerable<T, U>(source, filter);
}

struct SelectEnumerable<T, U> : IEnumerable<U>
{
    IEnumerable<T> source;
    Func<T, U> filter;

    public SelectEnumerable(IEnumerable<T> source, Func<T, U> filter)
        => (this.source, this.filer) = (source, filter);

    public IEnumerator<U> GetEnumerator()
    {
        foreach(var item in source)
        {
            yield return filter(item);
        }
    }
}
```

これと同じような動作にしようと思います。

# 

GPUは並列なので、一つずつ評価することはできません。ひとつ目にアクセスした際に、全部評価するようにしましょう。メソッドチェーン呼び出し時には遅延評価のためのデリゲート（`Func<T[]>`）を作ります。そして評価時にデリゲートがInvokeされ評価される…というIEnumerableを`LazyGpuEnumerable`というラッパーで実装します。

```csharp
public interface IGpuEnumerable<T> : IEnumerable<T> where T : unmanaged {}

readonly struct LazyGpuEnumerable<T> : IGpuEnumerable<T> where T : unmanaged
{
    readonly Func<T[]> array;
    public LazyGpuEnumerable(Func<T[]> array) => this.array = array;
    public readonly IEnumerator<T> GetEnumerator() => array().AsEnumerable().GetEnumerator();
    readonly IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
}
```

`IGpuEnumerable`というのは、IEnumerableと同じですが、「ここからGPUで計算しますよ」という合図みたいなものです。PLINQの`IParallelEnumerable`というのがありますが、これを参考にしました。スレッド数などのConfigを入れてもよさそうですね。

```csharp
public static IGpuEnumerable<T> AsGpuEnumerable<T>(this IEnumerable<T> source)
    where T : unmanaged
    => new LazyGpuEnumerable<T>("AsGpuEnumerable", new (() => source.ToArray()));
```

GPUとの通信部分は`Dispatch`メソッドに同期的に実装するとしてWhere拡張メソッドを実装します。Lazyに詰め込んで遅延化します。

```csharp
public static IGpuEnumerable<T> Where<T>(
    this IGpuEnumerable<T> source, Func<T, bool> _,
    [CallerFilePath] string file = null, [CallerMemberName] string member = null, [CallerLineNumber] int line = 0)
    where T : unmanaged
{
    return new LazyGpuEnumerable<T>("Where", new (() => Dispatch<T, T>(source, useCounter: true, file, member, line)));
}
```

`Dispatch`の中身は普通にComputeShaderを実行するだけです。WhereはAppendStructuredBufferを使うので、カウント用のコードも入れておきます。これはSelectでは無視します。

```csharp
static unsafe U[] Dispatch<T, U>(
    IGpuEnumerable<T> source, bool useCounter,
    string file, string member, int line)
    where T : unmanaged where U : unmanaged
{
    var inputArr = source.ToArray();
    using var input = new GraphicsBuffer(
        GraphicsBuffer.Target.Structured,
        inputArr.Length, sizeof(T));
    input.SetData(inputArr);
    using var output = new GraphicsBuffer(
        GraphicsBuffer.Target.Structured | GraphicsBuffer.Target.Append,
        inputArr.Length, sizeof(U));
    using var immCounter = new GraphicsBuffer(
        GraphicsBuffer.Target.IndirectArguments,
        4, sizeof(int));
    using (var cmd = new CommandBuffer() { name = $"{file}:{member}:{line-1}" })
    {
        cmd.SetBufferCounterValue(output, 0);

        var (shader, code) = shaderLibrary.Value.Resolve(file, member, line-1);
        cmd.SetComputeBufferParam(shader, 0, "_Input", input);
        cmd.SetComputeBufferParam(shader, 0, "_Output", output);
        cmd.SetComputeIntParam(shader, "_InputLength", inputArr.Length);
        cmd.DispatchCompute(shader, 0, Mathf.CeilToInt(inputArr.Length / 64f), 1, 1);

        cmd.CopyCounterValue(output, immCounter, 0);
        Graphics.ExecuteCommandBuffer(cmd);
    }

    var countBuff = new int[1];
    immCounter.GetData(countBuff);
    var count = useCounter ? countBuff[0] : inputArr.Length;

    var outputArr = new U[count];
    output.GetData(outputArr);
    return outputArr;
}
```

# シェーダの構築

LINQの最大のメリットの一つは、インラインでかけることです。拡張メソッドの引数にComputeShaderのインスタンスをとってもよいのですが、それではあまりにも片手落ちですよね。C#からインラインのラムダ式を拾ってきてComputeShaderに変換します。

戦略としては、`__GpuLinq_ShaderLibrary__.gpulinqlib`というファイルをライブラリ内に作成し、そのScriptedImporterを実装します。このScriptedImporterによって、C#コードを解析し、ラムダ式を拾います。このようにしている理由は、ScriptedImporterで作成したアセットはメモリにしか乗らず、シリアライズされないという性質を使いたかったからです。Gitなどに重複した情報を残したくないですし、変なファイルが生まれるのもスマートではありません。

```csharp
[ScriptedImporter(1, "gpulinqlib")]
class GpuLinqShaderLibraryImporter : ScriptedImporter
{
    [InitializeOnLoadMethod]
    static void InitializeOnLoadMethod()
    {
        var lib = ShaderLibrary.Load();
        AssetDatabase.ImportAsset(AssetDatabase.GetAssetPath(lib));  
    }

    [SerializeField] bool enableDebugFlag = true;
    
    public override void OnImportAsset(AssetImportContext ctx)
    {
        var lib = ScriptableObject.CreateInstance<ShaderLibrary>();

        // C#を舐めて、対象のラムダ式ごとにComputeShaderを作成し、
        // ctx.AddObjectToAssetでアセットに追加する
    }
```

C#の解析は、Roslynを使って構文を認識しながら解析するのがスマートです。しかし、今回は時間がなかったので、簡単な文字列操作でやります。

```csharp
var lib = ScriptableObject.CreateInstance<ShaderLibrary>();
var sourceFiles = AssetDatabase.FindAssets("t:MonoScript")
    .Select(AssetDatabase.GUIDToAssetPath)
    .Select(AssetDatabase.LoadAssetAtPath<MonoScript>)
    .Where(x => x.name != "ShaderLibrary")
    .Where(x => x.text.Contains(".AsGpuEnumerable()"))
    .ToArray();
foreach (var file in sourceFiles)
{
    var line = file.text.Replace("\r\n", "\n").Split('\n');
    var funcs = line
        .Select((text, i) => (text, startIndex: i))
        .Where(startLine => startLine.text.Contains(".AsGpuEnumerable()"))
        .SelectMany(startLine => line
            .Select((text, i) => (text, targetIndex: i))
            .Skip(startLine.startIndex + 1)
            .Select(targetLine =>
            {
                var regex = new Regex(@"^\s*\.(?<ope>Where|Select)\((?<arg>.+)=>(?<func>.+)\);?\s*(\/\/.*)?\s*$");
                return (lineNo: targetLine.targetIndex, m: regex.Match(targetLine.text));
            })
            .TakeWhile(lno_m => lno_m.m.Success)
            .ToArray())
        .ToArray();
    foreach (var (lineNo, match) in funcs)
    {
        var ope = match.Groups["ope"].Value;
        var member = "?";
        var arg = match.Groups["arg"].Value;
        var func = match.Groups["func"].Value;
        var code = Convert(ope, arg, func);
        var text = new TextAsset(code){ name = $"{file.name}::{member}#{lineNo}" };
        ctx.AddObjectToAsset($"code_{text.name}", text);
        var shader = ShaderUtil.CreateComputeShaderAsset(ctx, code);
        shader.name = text.name;
        ctx.AddObjectToAsset($"shader_{shader.name}", shader);
        lib.shaders.Add(new ShaderLibrary.KeyValue
        {
            file = file.name,
            member = member,
            line = lineNo,
            shader = shader,
        });
    }
}
ctx.AddObjectToAsset("index", lib);
ctx.SetMainObject(lib);
```

1. `AsGpuEnumerable`という文字列を含むC#ファイルを抽出
1. `AsGpuEnumerable`という文字列を含む行を抽出
1. その行の下に続くSelect/Whereを正規表現でキャプチャ
1. ラムダ式をComputeShaderにコンバート
1. `AddObjectToAsset`

という流れです。

時間がなかったのでコンバータも非常にシンプルにしてあります。

```csharp
string Convert(string ope, string arg, string func)
{
    var sb = new StringBuilder();
    if (enableDebugFlag)
    {
        sb.AppendLine($"#pragma enable_d3d11_debug_symbols");
    }
    sb.AppendLine($"#pragma kernel CSMain");
    switch (ope)
    {
        case "Where":
            ConvertWhere(arg, func, sb);
            break;
        case "Select":
            ConvertSelect(arg, func, sb);
            break;
        default:
            throw new NotImplementedException($"Unknown operator: {ope}");
    }
    return sb.ToString();

    static void ConvertWhere(string arg, string func, StringBuilder sb)
    {
        sb.AppendLine($"StructuredBuffer<int> _Input;");
        sb.AppendLine($"AppendStructuredBuffer<int> _Output;");
        sb.AppendLine($"uint _InputLength;");
        sb.AppendLine();
        sb.AppendLine($"[numthreads(64, 1, 1)]");
        sb.AppendLine($"void CSMain(uint3 id : SV_DispatchThreadID)");
        sb.AppendLine($"{{");
        sb.AppendLine($"    if (id.x >= _InputLength) return;");
        sb.AppendLine($"");
        sb.AppendLine($"    int {arg} = _Input[id.x];");
        sb.AppendLine($"    if ({func})");
        sb.AppendLine($"    {{");
        sb.AppendLine($"        _Output.Append({arg});");
        sb.AppendLine($"    }}");
        sb.AppendLine($"}}");
    }

    static void ConvertSelect(string arg, string func, StringBuilder sb)
    {
        sb.AppendLine($"StructuredBuffer<int> _Input;");
        sb.AppendLine($"RWStructuredBuffer<int> _Output;");
        sb.AppendLine($"uint _InputLength;");
        sb.AppendLine();
        sb.AppendLine($"[numthreads(64, 1, 1)]");
        sb.AppendLine($"void CSMain(uint3 id : SV_DispatchThreadID)");
        sb.AppendLine($"{{");
        sb.AppendLine($"    if (id.x >= _InputLength) return;");
        sb.AppendLine($"");
        sb.AppendLine($"    int {arg} = _Input[id.x];");
        sb.AppendLine($"    _Output[id.x] = {func};");
        sb.AppendLine($"}}");
    }
}
```

作ったComputeShaderアセットをルーティングするためのScriptableObjectも実装します。GpuLinqの使用側からDictionaryで引きたいので、`ISerializationCallbackReceiver`でDictionaryに再構築しています。実体としては、ファイル名＋行番号をキーとしてComputeShaderを引くことができるものです。

```csharp
class ShaderLibrary : ScriptableObject, ISerializationCallbackReceiver
{
    public static ShaderLibrary Load() => Resources.LoadAll("__GpuLinq_ShaderLibrary__").OfType<ShaderLibrary>().First();

    [Serializable]
    public struct KeyValue
    {
        public string file;
        public string member;
        public int line;
        public ComputeShader shader;
        public TextAsset code;
    }

    [SerializeField] public List<KeyValue> shaders = new();

    Dictionary<(string, string, int), (ComputeShader cs, TextAsset code)> shaderMap;

    public (ComputeShader cs, TextAsset code) Resolve(string file, string member, int line)
        => shaderMap[(Path.GetFileNameWithoutExtension(file), "?", line)];

    public void OnBeforeSerialize() {}

    public void OnAfterDeserialize()
        => shaderMap = shaders.ToDictionary(x => (x.file, x.member, x.line), x => (x.shader, x.code));
}
```

# 実行

このように、メソッドチェーンでりようできます。C#の解析を端折っている都合、Where/Selectの前後で改行を入れてやる必要があります。

```csharp
[Test]
public void GpuLinqTestSimplePasses()
{
    var source = Enumerable.Range(1, 1_000_000).ToArray();
    
    var cpuSw = Stopwatch.StartNew();
    var cpu = source
        .Where(x => (uint)x % 2 == 0)
        .Select(x => x * 2)
        .ToArray();
    Debug.Log($"cpu: {cpuSw.ElapsedMilliseconds}ms");
    
    var gpuSw = Stopwatch.StartNew();
    var gpu = source
        .AsGpuEnumerable()
        .Where(x => (uint)x % 2 == 0) // unorderd
        .Select(x => x * 2)
        .ToArray();
    Debug.Log($"gpu: {gpuSw.ElapsedMilliseconds}ms");
    
    Assert.IsTrue(new HashSet<int>(cpu).SetEquals(gpu));
}
```

# パフォーマンス測定

上記のテストコードでは、GPUのオーバーヘッドの方が勝ってしまいました。ラムダ式の部分を複雑するとGPUに有利になるかと思います。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280975/8e36decf-767f-5b5a-088c-4ab1afd4b996.png)

# まとめ

以上のコードをUPMで公開しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/280975/5fa6dffc-c9f0-6a4a-1164-85e091ebff59.png)


```
https://github.com/wallstudio/TechArticleSamples.git?path=GpuLinq/Packages/GpuLinq
```

https://github.com/wallstudio/TechArticleSamples/tree/master/GpuLinq

ちょっと時間が足りなくて十分な実装ができませんでしたが、GPGPUを使ったLINQっぽいものは作ることができました。

満足😎
（また機会や需要があれば、ちゃんと実装しなおそうと思います）