Go to Qiita Advent Calendar 2024 Top
LoginSignup
18
8

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

UnityAdvent Calendar 2023

Day 25
@up-hash

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

Last updated at Posted at 2023-12-24

前書き

この記事は、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に有利になるかと思います。

image.png

まとめ

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

image.png

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

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

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

18
8
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
18
8

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?