VRMSpringBoneをC# Job System対応にする

2019/01/10修正

• コリジョンのローテーションが考慮されてないためあたり判定がおかしい場合があったのを修正

使用している各ソフトウェァのバージョン

この記事では

• Unity 2018.3
  • Burst 0.2.4-preview.37
  • Collections 0.0.9-preview.10
  • Jobs 0.0.7-preview.5
  • Mathematics 0.0.12-preview.19
• UniVRM 0.45

を使用している

VRMを使う際に困ったこと

VRMのキャラを複数読み込むと、CPU負荷が大きく上がる
Profilerを動かすと、揺れ物として使われているVRMSpringBoneの処理負荷が大きいことが分かる
特にVRoidで作られたキャラは髪の毛の揺れなどに結構ふんだんに使っているため顕著だ
そこで、VRMSpringBoneのC# Job System対応をしてみる

VRMSpringBoneをJob化したソース

オリジナルはとてもJob化しやすいコードなので、特に説明なくJob化したコードを以下に示す

VRMSpringBoneOptimizer.cs

using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using Unity.Mathematics;
using Unity.Jobs;
using Unity.Collections;
using static Unity.Mathematics.math;
using UnityEngine.Jobs;
using System.Linq;

namespace VRM
{
    public class VRMSpringBoneOptimizer : MonoBehaviour
    {
        // gにVRMSpringBoneOptimizerがまだついてないなら追加する
        public static void Attach(GameObject g)
        {
            if (!g.GetComponent<VRMSpringBoneOptimizer>())
            {
                g.AddComponent<VRMSpringBoneOptimizer>();
            }
        }

        private void Start()
        {
            // 全てのVRMSpringBoneを登録する
            foreach (var a in gameObject.GetComponentsInChildren<VRMSpringBone>())
            {
                VRMSpringBoneJobGroup.Register(gameObject, a);
            }
        }
    }

    public class VRMSpringBoneJobGroup : MonoBehaviour
    {
        bool initialized = false;

        List<VRMSpringBoneColliderGroup> groups;
        Transform m_center;
        List<VRMSpringBoneLogicInfo> logicInfoList = new List<VRMSpringBoneLogicInfo>();
        VRMSpringBoneLogic logic;
        JobHandle jobHandle;

        struct VRMSpringBoneLogicInfo
        {
            public VRMSpringBone bone;
            public Transform center;
            public Transform transform;
            public Vector3 pos;
        }

        struct SphereCollider
        {
            public float3 Position;
            public float Radius;
        }

        /// <summary>
        /// 
        /// original from
        /// 
        /// http://rocketjump.skr.jp/unity3d/109/
        /// 
        /// </summary>
        [Unity.Burst.BurstCompile]
        struct VRMSpringBoneLogic : System.IDisposable
        {
            GetRotationJob rotjob;
            CalcColliderPosJob coljob;
            CalcNextPosJob job;

            TransformAccessArray transformList;
            TransformAccessArray parentList;
            TransformAccessArray colliderGroupList;

            NativeArray<SphereCollider> precolliders;
            NativeArray<SphereCollider> colliders;
            NativeArray<VariableParam> vparamList;
            NativeArray<FixedParam> fparamList;
            NativeArray<ReadParam> paramList;


            // 変化がないパラメータ群
            public struct FixedParam
            {
                public float m_length;
                public quaternion m_localRotation;
                public float3 m_boneAxis;
                public float radius;
                public float stiffnessForce;
                public float dragForce;
                public float3 external;
            }

            // 今回と前回の頂点情報
            public struct VariableParam
            {
                public float3 m_currentTail;
                public float3 m_prevTail;
            }

            public struct ReadParam
            {
                public quaternion parentRotation;
            }

            public void InitializeColliderGroup(List<Transform> transforms, List<SphereCollider> spcolliders)
            {
                colliderGroupList = new TransformAccessArray(transforms.ToArray());
                precolliders = new NativeArray<SphereCollider>(spcolliders.ToArray(), Allocator.Persistent);
                colliders = new NativeArray<SphereCollider>(spcolliders.Count, Allocator.Persistent);
                coljob.inputList = precolliders;
                coljob.outputList = colliders;
                job.colliderList = colliders;
            }

            public void Initialize(List<VRMSpringBoneLogicInfo> alist)
            {
                transformList = new TransformAccessArray(alist.Select(x => x.transform).ToArray());
                parentList = new TransformAccessArray(alist.Select(x => x.transform.parent).ToArray());
                paramList = new NativeArray<ReadParam>(alist.Count, Allocator.Persistent);
                fparamList = new NativeArray<FixedParam>(alist.Count, Allocator.Persistent);
                vparamList = new NativeArray<VariableParam>(alist.Count, Allocator.Persistent);
                job.paramList = paramList;
                job.fparamList = fparamList;
                job.vparamList = vparamList;
                rotjob.paramList = paramList;

                for (int i = 0; i < alist.Count; ++i)
                {
                    var info = alist[i];

                    var worldChildPosition = info.transform.TransformPoint(info.pos);

                    var pp = info.center != null
                            ? info.center.InverseTransformPoint(worldChildPosition)
                            : worldChildPosition
                            ;

                    job.fparamList[i] = new FixedParam
                    {
                        m_localRotation = info.transform.localRotation,
                        m_boneAxis = info.pos.normalized,
                        m_length = info.pos.magnitude,
                        radius = info.bone.m_hitRadius,
                        stiffnessForce = info.bone.m_stiffnessForce,
                        dragForce = info.bone.m_dragForce,
                        external = info.bone.m_gravityDir * info.bone.m_gravityPower,
                    };

                    job.vparamList[i] = new VariableParam
                    {
                        m_currentTail = pp,
                        m_prevTail = pp,
                    };
                }
            }

            public void Dispose()
            {
                job.paramList.Dispose();
                job.fparamList.Dispose();
                job.vparamList.Dispose();
                transformList.Dispose();
                parentList.Dispose();
                colliderGroupList.Dispose();
                colliders.Dispose();
                precolliders.Dispose();
            }

            // トランスフォームからrotationを取り出すだけのJob
            [Unity.Burst.BurstCompile]
            struct GetRotationJob : IJobParallelForTransform
            {
                [WriteOnly] public NativeArray<ReadParam> paramList;
                void IJobParallelForTransform.Execute(int index, TransformAccess taccess)
                {
                    paramList[index] = new ReadParam
                    {
                        parentRotation = taccess.rotation
                    };
                }
            }

            // コライダーの位置を更新する
            [Unity.Burst.BurstCompile]
            struct CalcColliderPosJob : IJobParallelForTransform
            {
                [ReadOnly] public NativeArray<SphereCollider> inputList;
                [WriteOnly] public NativeArray<SphereCollider> outputList;
                void IJobParallelForTransform.Execute(int index, TransformAccess taccess)
                {
                    outputList[index] = new SphereCollider
                    {    // 2019/01/10修正
                         Position = (float3)taccess.position + mul(taccess.rotation, inputList[index].Position),
                         Radius = inputList[index].Radius,
                    };
                }
            }

            // 物理計算を行う
            [Unity.Burst.BurstCompile]
            struct CalcNextPosJob : IJobParallelForTransform
            {
                public NativeArray<VariableParam> vparamList;
                [ReadOnly] public NativeArray<SphereCollider> colliderList;
                [ReadOnly] public NativeArray<FixedParam> fparamList;
                [ReadOnly] public NativeArray<ReadParam> paramList;

                [ReadOnly] public float4x4 centerTrans;
                [ReadOnly] public float4x4 centerInvTrans;
                [ReadOnly] public float deltaTime;

                void IJobParallelForTransform.Execute(int index, TransformAccess taccess)
                {
                    float3 tposition = taccess.position;

                    var fparam = fparamList[index];
                    float3 currentTail = math.transform(centerTrans, vparamList[index].m_currentTail);
                    float3 prevTail = math.transform(centerTrans, vparamList[index].m_prevTail);

                    var param = paramList[index];

                    // verlet積分で次の位置を計算
                    var nextTail = currentTail
                        + (currentTail - prevTail) * (1.0f - fparam.dragForce) // 前フレームの移動を継続する(減衰もあるよ)
                        + mul(mul(param.parentRotation, fparam.m_localRotation), fparam.m_boneAxis) * (fparam.stiffnessForce * deltaTime) // 親の回転による子ボーンの移動目標
                        + (fparam.external * deltaTime) // 外力による移動量
                        ;

                    // 長さをboneLengthに強制
                    nextTail = tposition + normalize(nextTail - tposition) * fparam.m_length;

                    // Collisionで移動
                    for (int i = 0; i < colliderList.Length; ++i)
                    {
                        var collider = colliderList[i];

                        var r = fparam.radius + collider.Radius;
                        var dir = nextTail - collider.Position;
                        if (dot(dir, dir) <= (r * r))
                        {
                            // ヒット。Colliderの半径方向に押し出す
                            var normal = normalize(dir);
                            var posFromCollider = collider.Position + normal * (fparam.radius + collider.Radius);
                            // 長さをboneLengthに強制
                            nextTail = tposition + normalize(posFromCollider - tposition) * fparam.m_length;
                        }
                    }

                    var rotation = mul(param.parentRotation, fparam.m_localRotation);
                    taccess.rotation = Quaternion.FromToRotation(mul(rotation, fparam.m_boneAxis), nextTail - (float3)taccess.position) * rotation;

                    vparamList[index] = new VariableParam
                    {
                        m_prevTail = transform(centerInvTrans, currentTail),
                        m_currentTail = transform(centerInvTrans, nextTail),
                    };
                }
            }

            // Jobのための事前計算をし、Jobを発行する
            public JobHandle Process(Transform center)
            {
                job.deltaTime = Time.deltaTime;
                job.centerTrans = center != null ? center.localToWorldMatrix : Matrix4x4.identity;
                job.centerInvTrans = center != null ? center.worldToLocalMatrix : Matrix4x4.identity;

                return job.Schedule(transformList, JobHandle.CombineDependencies(rotjob.Schedule(parentList), coljob.Schedule(colliderGroupList)));
            }
        }

        // VRMSprintBoneを登録する
        public static void Register(GameObject target, VRMSpringBone bone)
        {
            VRMSpringBoneColliderGroup[] groups = bone.ColliderGroups;

            if (groups == null)
            {
                return;
            }

            bone.enabled = false;

            foreach (var item in target.GetComponents<VRMSpringBoneJobGroup>())
            {
                if (item.groups.Count != groups.Length || item.m_center != bone.m_center)
                {
                    continue;
                }

                bool same = true;
                for (int j = 0; j < groups.Length; ++j)
                {
                    if (groups[j] != item.groups[j])
                    {
                        same = false;
                        break;
                    }
                }

                if (!same)
                {
                    continue;
                }

                item.Setup(bone);
                return;
            }

            {
                var item = target.AddComponent<VRMSpringBoneJobGroup>();
                item.Setup(bone);
            }
        }

        void SetupColliders()
        {
            List<Transform> transforms = new List<Transform>();
            List<SphereCollider> spcolliders = new List<SphereCollider>();

            foreach (var group in groups)
            {
                if (group != null)
                {
                    foreach (var collider in group.Colliders)
                    {
                        transforms.Add(group.transform);

                        spcolliders.Add(new SphereCollider
                        {
                            Position = collider.Offset,
                            Radius = collider.Radius,
                        });
                    }
                }
            }

            logic.InitializeColliderGroup(transforms, spcolliders);
        }

        void Setup(VRMSpringBone bone)
        {
            groups = bone.ColliderGroups.ToList();
            m_center = bone.m_center;

            if (bone.RootBones != null)
            {
                foreach (var go in bone.RootBones)
                {
                    if (go != null)
                    {
                        SetupRecursive(bone, bone.m_center, go, logicInfoList);
                    }
                }
            }
        }

        void SetupRecursive(VRMSpringBone bone, Transform center, Transform parent, List<VRMSpringBoneLogicInfo> list)
        {
            if (parent.childCount == 0)
            {
                var delta = parent.position - parent.parent.position;
                var childPosition = parent.position + delta.normalized * 0.07f;
                list.Add(new VRMSpringBoneLogicInfo
                {
                    bone = bone,
                    center = center,
                    transform = parent,
                    pos = parent.worldToLocalMatrix.MultiplyPoint(childPosition)
                });
            }
            else
            {
                var firstChild = parent.childCount > 0 ? parent.GetChild(0) : null;
                var localPosition = firstChild.localPosition;
                var scale = firstChild.lossyScale;
                list.Add(new VRMSpringBoneLogicInfo
                {
                    bone = bone,
                    center = center,
                    transform = parent,
                    pos = new Vector3(
                        localPosition.x * scale.x,
                        localPosition.y * scale.y,
                        localPosition.z * scale.z
                        )
                });
            }

            foreach (Transform child in parent)
            {
                SetupRecursive(bone, center, child, list);
            }
        }

        private void LateUpdate()
        {
            if (!initialized)
            {
                SetupColliders();
                logic.Initialize(logicInfoList);
                initialized = true;
            }
            else
            {
                jobHandle.Complete();
            }

            jobHandle = logic.Process(m_center);

            JobHandle.ScheduleBatchedJobs();
        }

        private void OnDestroy()
        {
            if (initialized)
            {
                jobHandle.Complete();

                logic.Dispose();
            }
        }
    }
}

##使い方

上記ソースをVRMSpringBoneOptimizer.csという名前で保存し
VRMSpringBoneOptimizerコンポネントをVRMのルートのGameObjectに追加しておくと機能する

ランタイムで動的にVRMモデルを読み込んだ場合、VRMSpringBoneOptimizer.Attach(ルートのgameObject);と書けば機能する

工夫したところ

オリジナルのVRMSpringBoneの単位でJob化しても、一つの処理単位が小さすぎてあまり高速化されない

実際のVRMモデルの特徴を見てみると、ColliderGroupは同じものが割り当てられてるケースが多い
そこに着目し、同じColliderGroupを持つVRMSpringBoneを集めてそれを一つのJobの単位としている

Jobはトランスフォームからrotationを取り出すだけのJob
コライダーの位置を更新するJob
物理計算を行うJobの三つに分け、上の二つが終わったら下の一つが動くようになっている

性能

ひゃ～めっちゃ綺麗に並列化するようになったぞ。オラおでれぇたぞ pic.twitter.com/s6mxIuMKIl

— tatsunoru@DOFI (@tatsunoru) 2018年12月23日

オリジナルの実装から、C# Job Systemで書き直した揺れ物実装に途中で差し替えるデモ
10fpsが40fps以上になる

ただ、C# Job Systemに当てはめづらいパラメータ(Transform m_center)を消去し、gravityなどの値は動的に変更できないようになってるので、あくまで追い風参考値 pic.twitter.com/DoL7PxJtsD

— tatsunoru@DOFI (@tatsunoru) 2018年12月23日

(このツイートではm_centerを無視していると書いてるが、今は対応した)

Profilerを見るとコア数に応じて綺麗に並列化されていることがわかる
また、実際にWinのIL2CPPビルドでオンとオフを比べるとかなり高速化していることがわかる

制限

m_stiffnessForceなどのフィールドはオリジナルだと毎回参照しているが
今回のJob化したプログラムでは最初しか参照しないため、動的に変化させても動作に反映されない

Transformへの参照・反映のタイミングがオリジナルと違い、貫通が多く発生している可能性がある

初期化部分はかなり適当に書いているのでそこそこ重い処理となっている
今のままだとキャラが出た最初のフレームでかなり処理落ちするかも

複数体のモデルに大して適用しないとあまり並列化できないかも

本当は

もっと細かく説明を書きたいところだけど、どっからどう説明したらいいのかわからなくなったのでソースだけ示すことをお許しください

また、パッと見はそれっぽく動いてるだけで、そもそもちゃんと実装できてない可能性もあるので、何か問題点などありましたらご指摘ください

ライセンス

MIT LicenseかApache 2.0 Licenseのお好きな方で