Unity:AndroidでComputeShaderを使ってみました

画像処理では、並列計算による高速化が求められます。（常に…）
Compute Shader を使った並列計算による高速化は実装が簡単で効果抜群なのですが、DirectX 依存なので（と思い込んでいた私は） PC 上での動作に限られていると思っていました。

しかし、ふと Unity の Compute Shader って Android 用にビルドした場合も動作するのかな？
だとしたら、Unity すばらしいなって思いまして、さっそく試してみました。

左は入力のカラー画像、右は出力のグレースケール画像（…イラストは私が描きました）

できた！Unityすばらしい！

使った Compute Shader のコードはこちら↓

// file head
RWStructuredBuffer<uint> _Histogram;
RWTexture2D<float4> _Result;

// Gathering pass
Texture2D<float4> _Source;

uint2 _SourceSize;

// groupshared uint gs_histogram[256]; // group shared memory can be used in a Windows PC, but it could not be used in an Android platform. why?

#pragma kernel KHistogramGather
[numthreads(32, 32, 1)]
void KHistogramGather(uint3 id : SV_DispatchThreadID, uint3 _group_thread_id : SV_GroupThreadID)
{
	const uint thread_id = _group_thread_id.y * 32 + _group_thread_id.x;

	//if (thread_id == 0)
	//{
	//	for (int i = 0; i < 256; i++)
	//	{
	//		gs_histogram[i] = 0;
	//	}
	//}

	// GroupMemoryBarrierWithGroupSync();

	if (id.x < _SourceSize.x && id.y < _SourceSize.y)
	{
		float3 color = saturate(_Source[id.xy].xyz);

		// Convert color to grayscale
		float luminance = dot(color.rgb, float3(0.2125, 0.7154, 0.0721));
		uint idx_l = (uint)round(luminance * 255.0);

		// InterlockedAdd(gs_histogram[idx_l], 1);
		InterlockedAdd(_Histogram[idx_l], 1); // it is an alternative method for using a group shared memory
		_Result[id.xy] = float4(luminance, luminance, luminance, 1);
	}

	// GroupMemoryBarrierWithGroupSync();

	//if (thread_id == 0)
	//{
	//	for (int i = 0; i < 256; i++)
	//	{
	//		InterlockedAdd(_Histogram[i], gs_histogram[i]);
	//	}
	//}
}

単にグレースケール化するだけじゃなくて、輝度ヒストグラムまで構築するスグレモノなのですが
残念なことに groupshared メモリを使った高速化がAndroid上だけ、できないことが分かったのでコメントアウトしています。

シェーダーの利用はスクリプト側で行いますので、利用スクリプトの方も以下に載せます。

using System.Collections;
using System.Collections.Generic;
using UnityEngine;

public class DeviceCameraBehaviour : MonoBehaviour {

    public GameObject resultCube;
    public ComputeShader RGBA2GRAY;
    public Texture2D Source;
    public RenderTexture Result;

    int _kernelIndex = 0;
    ComputeBuffer _HistogramBuffer = null;
    uint[] _emptyBuffer = new uint[256];
    uint[] _histogramBuffer = new uint[256];

    // Use this for initialization
    void Start () {

        Result = new RenderTexture(Source.width, Source.height, 0, RenderTextureFormat.ARGB32);
        Result.enableRandomWrite = true;
        Result.Create();

        _HistogramBuffer = new ComputeBuffer(256, sizeof(uint));
        for (int i = 0; i < _emptyBuffer.Length; i++)
        {
            _emptyBuffer[i] = 0;
        }
        _HistogramBuffer.SetData(_emptyBuffer);

        _kernelIndex = RGBA2GRAY.FindKernel("KHistogramGather");
        RGBA2GRAY.SetBuffer(_kernelIndex, "_Histogram", _HistogramBuffer);
    }
	
	// Update is called once per frame
	void Update () {
        if (Input.GetKeyDown(KeyCode.Space))
        {
            DoComputeShader();

            Debug.Log("Time = " + Time.time.ToString("000.0000"));
            WebCamDevice[] devices = WebCamTexture.devices;
            for (int i = 0; i < devices.Length; i++)
                Debug.Log(devices[i].name);
        }
        for (int i = 0; i < Input.touchCount; ++i)
        {
            if (Input.GetTouch(i).phase == TouchPhase.Began)
            {

                // Construct a ray from the current touch coordinates
                Ray ray = Camera.main.ScreenPointToRay(Input.GetTouch(i).position);
                // Create a particle if hit
                if (Physics.Raycast(ray))
                {
                    DoComputeShader();

                    Debug.Log("Time = " + Time.time.ToString("000.0000"));
                    WebCamDevice[] devices = WebCamTexture.devices;
                    for (int j = 0; j < devices.Length; j++)
                        Debug.Log(devices[j].name);
                }
            }
        }
    }

    private void DoComputeShader()
    {
        RGBA2GRAY.SetTexture(_kernelIndex, "_Source", Source);
        RGBA2GRAY.SetTexture(_kernelIndex, "_Result", Result);
        RGBA2GRAY.SetVector("_SourceSize", new Vector2(Source.width, Source.height));
        _HistogramBuffer.SetData(_emptyBuffer);
        RGBA2GRAY.Dispatch(_kernelIndex, Mathf.CeilToInt(Source.width / 32f), Mathf.CeilToInt(Source.height / 32f), 1);
        _HistogramBuffer.GetData(_histogramBuffer);
        uint total = 0;
        for (int i = 0; i < _histogramBuffer.Length; i++)
        {
            total += _histogramBuffer[i];
        }
        Debug.Log("histogram total = " + total);
        resultCube.GetComponent<Renderer>().material.mainTexture = Result;
    }
}

簡単にコードを解説してみますと

RGBA2GRAY というのが、Compute Shader で、先ほど示したコードで作ったシェーダーを Unity のエディタから紐づけます。
あとは、画像データを Source テクスチャとして用意して渡したら Disptach にて 32 分の 1 の幅と高さのスレッドグループ数を指定し実行しています。
なぜ 32 分の 1 にしなければいけないかというと、Compute Shader 側で[numthreads(32, 32, 1)]と、32 x 32のタイル領域ごとに処理をすると宣言しているからです。
Dispach で縦横のタイル数を指定してあげることで全画素についての処理が完成するというイメージを持っていただけると、理解しやすいかなって思います。

最後に Andorid のビルド設定ですが、以下のようにして、端末は2017年1月に購入した Phab 2 Pro で動作することを確認しました。（自身の携帯である Xperia Z3 では Compute Shader は動かなかった…）

今回の記事は以下のページを参考に作成しました。

まず Compute Shader を一度も触ったことがない人は、こちらの記事を参考にするとサクッとPC上で動作確認までできるようになると思います。
[Unity] UnityでComputeShaderを使う解説をしているページを訳してみた その2 - Qiita

次にテクスチャを入力に、Compute Shader で何かしら画素値を処理してテクスチャを出力する場合については次の記事を参考にします。
DirectCompute tutorial for Unity 3: Textures | Cheney Shen

ヒストグラムの計算については、次のプロジェクトの、このプルリクエストを参考にしました。（ヒストグラムなのに並列計算で高速化！？できないと思っていたので驚きましたよ、考えた人グッジョブです！）
Unity-Technologies / cinematic-image-effects / Pull request #31: [tcg] Histogram compute shader optimizations (~4x) — Bitbucket