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

前書き

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

TL;DR;

PackageMaangerでインストール

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

使い方

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は、チェーンで積まれているときには各アイテムは評価されません。

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

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

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

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というラッパーで実装します。

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を入れてもよさそうですね。

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

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

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では無視します。

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などに重複した情報を残したくないですし、変なファイルが生まれるのもスマートではありません。

[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を使って構文を認識しながら解析するのがスマートです。しかし、今回は時間がなかったので、簡単な文字列操作でやります。

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#ファイルを抽出
2. AsGpuEnumerableという文字列を含む行を抽出
3. その行の下に続くSelect/Whereを正規表現でキャプチャ
4. ラムダ式をComputeShaderにコンバート
5. AddObjectToAsset

という流れです。

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

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を引くことができるものです。

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の前後で改行を入れてやる必要があります。

[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に有利になるかと思います。

まとめ

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

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

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

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