CG×ML #3 DRを可能にするMitsuba3の中身 [Devlog #013]

Table of Contents

微分可能レンダリングを行うMitsuba3の仕組み

Dr.Jit（ジャストインタイムコンパイラ）

ここでは微分可能なレンダリングのためのジャストインタイムコンパイラについての論文1を理解する

Wenzel Jakob, Sébastien Speierer, Nicolas Roussel, and Delio Vicini. 2022. Dr.Jit: A Just-In-Time Compiler for Differentiable Rendering. In Transactions on Graphics (Proceedings of SIGGRAPH) 41(4).

Dr.Jitとは

物理ベース微分可能レンダリング（PBDR）のためのジャストインタイムコンパイラ

レンダリングアルゴリズムを実行するとグラフを生成

グラフとは

算術、ループ、レイトレーシングの操作、およびレンダリングアルゴリズムとシーンオブジェクト（形状、BSDF、テクスチャ、エミッターなど）間の情報を交換する多態的な呼び出し（例えば、BRDFのspecularの値を勾配方向に変えながらレイトレーシング操作を繰り返す）


Dr.Jitは提供されるシーンにグラフを特化させ、LLVMまたはOptiXを介してデータ並列カーネル（megakernel）にコンパイルする

megakernelアプローチとは

グラフィックスや物理ベースのレンダリングにおいて、多数の計算やタスクを一つの大きなカーネルに統合する手法

通常、GPU上で動作するプログラムは、多くの小さなカーネル（基本的な計算単位）に分割されて実行される

megakernelアプローチでは、これらの小さなタスクを一つの大きなカーネルに統合し、GPUの計算リソースをより効率的に活用することを目指す


通常のレンダリングのほか、微分シミュレーション（differential simulations）を動的にコンパイルする

大きな微分タスクを一連の漸近的なステップに分解するため、自動微分（automatic differentiation）の細かい実行をサポート

自動微分とは

自動微分（Automatic Differentiation ;AD）は、数値的に正確な微分値を計算する手段であり、関数の実行過程において、すべての計算ステップ（加算、乗算、関数呼び出しなど）の微分を追跡し、それらを組み合わせることで元の関数の導関数や勾配を導出する


画像空間における摂動（perturbation）を計算するForwardモードは、特定の入力パラメータに対する出力の変化を効率的に追跡する
（例えば、光源の位置や色の微小な変更がレンダリングされた画像にどのように影響するかを調査する）

計算の最後から逆順に微分を計算するReverseモードは、たくさんの入力パラメータを同時に最適化するためのパラメータ空間の導関数を提供する
（このアプローチは、多数の入力パラメータに対する少数の出力（例えば、誤差関数や目的関数の値）の勾配を求めるのに特に効率的である）

つまり、Forwardモードは特定のパラメータの影響を調べるのに、Reverseモードは多数のパラメータを最適化するのに利用される

パラメータ空間の例

シーン上での壁紙のテクセル（テクスチャ空間の基本単位）


どちらのモードも、関数の微分を計算する際には中間変数の導関数を逐次的に計算することで、微分の連鎖規則を適用する（Forwardモードは入力に近い部分から、Reverseモードは出力に近い部分から計算を始める）

目的

  1. Python等の高レベル言語で記述されたシミュレーションコードの実行を追跡し、どの部分のコードがどのように動作しているのか、どの部分が計算負荷が高いのかなどの情報をもとに、コードを効率的に最適化・変換し、高速なデータ並列カーネルを生成する

  2. 微分の過程をより効率的に、かつ細かい制御のもとで行うことをサポートし、微分可能なレンダリングアルゴリズムの開発を簡易化する（効率的な方法としてシミュレーションの微分を微分のシミュレーションに変える）

"シミュレーションの微分を微分のシミュレーションに変える"とは

Derivative of a simulation（シミュレーションの微分）：

シミュレーションが入力に対してどのように変化するかを示す

例えば、物理ベースのレンダリングの場合、光の強度や物質の性質などの入力パラメータに微小な変化を加えたときの出力画像の変化量がこれに該当する

Simulation of the derivative（微分のシミュレーション）：

微分そのものを新しいシミュレーションとして扱う考え方

つまり、入力の変化に対する出力の変化を直接シミュレートすることで、微分の効果を視覚化や解析するためのツールとして利用する


この変換はMitsubaにおいて重要。なぜなら、

多くのアルゴリズムや最適化手法は、微分の情報を使用して動作する

微分可能なレンダリングは、シーンのパラメータを最適化するための勾配情報を提供する能力が求められる

この勾配情報は、シーンの照明や材料の性質、カメラの位置などのパラメータを調整して、目的の画像や効果に近づけるために使用される


Dr.Jitは自動微分のプロセスを細かく制御してこの変換を支援する

また、Dr.Jitはデータの依存関係をグローバルに追跡し、最終的に勾配に影響を与えない計算（冗長なシミュレーション）を削除することで効率を上げる

Introduction

物理ベースの微分可能レンダリング（PBDR）

PBDRは、任意のシーンパラメータに基づいて光輸送シミュレーションを微分できる

勾配ベースの最適化と組み合わせ、多くの未知数（入力パラメータ）を持つ非線形の問題を解決する

この技術は、画像の逆解析が必要な多くの分野での利用が期待される

PBDRの課題

実用的な実装が難しい

例えば、パストレーシングのようなアルゴリズムの逆モード導関数は非常に複雑になる可能性がある

最適化は数千回の反復を行うかもしれないため、正確さだけでなく効率も重要となる

Dr.Jitの設計原理

PBDRのための実用的かつ効率的なプラットフォーム

アーキテクチャ内の決定は、DRを活用する開発の簡素化と効率の向上を目指す

手動の微分は面倒でエラーが発生しやすいため、開発を簡素化するために自動微分（AD）に基づいて構築する

自動微分を標準的に使うと、PBDRの要求を満たすのが難しく、その結果として得られる導関数は効率的でなく偏ってしまう

非効率性を解消するために、データを集める範囲や方法を再調整することで偏りを取り除く234

微分演算はランダムウォークに移行し、散乱がおこるたびに新しい偏微分項が導入されるが、微分は何らかの形で事前にコンパイルされなければならない（計算速度の問題で）

最適化がシーンのパラメータの一部のみを対象とする場合、計算された勾配に影響を与えることがないステップを削除する

このような問題には動的なコンパイルのアプローチが必要となるため、ジャストインタイム（JIT）コンパイルを中心としたアプローチを行う

JITとは？

ジャストインタイムコンパイル（Just-In-Time Compilation、略してJIT）は、プログラムの実行時にその一部をネイティブのマシンコードに変換する技術のことを指し、インタープリタベースの実行よりも高速にコードを実行することができる

JITコンパイルの基本的な考え方は次のとおり：

  • プログラムの実行が開始されると、最初はインタープリタのようにコードを逐次的に解釈しながら実行する
  • ある部分のコードが頻繁に実行されることを検知すると、その部分をネイティブのマシンコードにコンパイルする
  • 以降、その部分のコードが再度実行される際には、コンパイルされたネイティブのマシンコードを実行することで、インタープリタのオーバーヘッドを回避し、パフォーマンスを向上させる

モンテカルロ積分は、カーネル（megakernel）内で行われ、ほとんどのサンプル状態をレジスタに格納することで、メモリ使用量やカーネル間通信を削減することができる

Dr.Jitのアプローチ

Dr.Jitの微分レンダリングステップはループ、多態性（ポリモーフィズム）、レイの交差操作を中断することなく記録し、レンダリングアルゴリズムと、参照されるすべてのシーンオブジェクト（マテリアル、テクスチャ、ボリューム、光源など）の実装を包含する非常に大きなトレースの形で画像（未評価）を返す

トレースとは

トレースは、プログラムの実行中に行われる操作の記録または追跡を指す

具体的には、プログラムがどのような手順で実行され、どの関数や命令が順番に呼び出され、どのようなデータが処理されたかなどの情報が、時間の流れに沿って記録される

Dr.Jitの文脈でのトレースは、レンダリングのシミュレーションコードの実行中に行われるすべての操作を逐次的に記録することを指している

このトレースは、後でのコンパイルやターゲットデバイスでの実行のために用いられ、シミュレーションの各ステップや詳細を後から正確に再現または解析することが可能となる


生成されるトレースは、全体のレンダリングプロセスを表しており（グローバルな表現であり）、最適化の機会を明らかにする

配列プログラミング（Array programming）

（省略）

自動微分（Automatic Differentiation）

（省略）

ジャストインタイムコンパイル（Just-in-Time Compilation）

実行例（Running example）

（省略）

議論（Discussion）

（省略）

最適化（Oprimizations）

Dr.Jitはコード生成を改善するための最適化を行う

LLVMやOptiX自体が洗練された最適化コンパイラであるため、標準的な最適化手法はバックエンドのコンパイルステップに取り込まれるが、実行時にオブジェクトの型に基づいて特定のメソッドを呼び出す動的ディスパッチは、LLVM/OptiXにおける手続き間最適化を妨げる

シーンを表す大きなデータ構造において、最適化に使用されるフィールド（複合的なデータ構造の一部分を形成する各要素）はほんの一部であるため、グローバルな最適化によって定数の伝播、変数の参照カウントに従うデッドコードの削除、値のナンバリングに基づく重複の排除が行われる

論文の図4を参照

Dr.Jitが多態的なメソッド呼び出しに遭遇すると、到達可能なすべてのインスタンスの実装をトレースし、複数の最適化を実行する

図4で示されている例では、入力(𝑖1, 𝑖2, 𝑖3)を持つ呼び出しが出力(𝑜1, 𝑜2, 𝑜3)を返し、変数𝑖3は定数リテラルで、周囲のプログラムは出力𝑜2と𝑜3のみを参照する

(a). Dr.Jitは、キャプチャされたサブトレースに定数を伝播させると同時に、呼び出し境界を越えてデッドコードを排除する

(b). この例では、すべてのサブトレースが𝑜2を生成するための同じ計算を実行し、その計算はその後、デバーチャライズ（実仮想化解除）される（つまり、呼び出しから外される）

(c). 最後に、サブトレースのうち2つが同一であることが分かり、カーネルのアセンブリ中に単一の関数定義のみを生成する


また、同一のトレースは１つのサブルーチン（処理のまとまり）にまとめることができ、これはシーンが同じタイプのオブジェクトをたくさん含む可能性がある"レンダリング"において重要となる


しかし、型の等価性が、2つのインスタンスが同じ操作のシーケンスを実行することを保証するわけではない

それぞれのインスタンスで基礎となる複数のパラメータの異なるサブセットがアクティブになる可能性がある

最終的には、到達可能なすべてのインスタンスをトレースし、同一の関数をマージする方法が提案されている

その他の特徴（Additional features）

（省略）

結果（Results）

（省略）

微分法（Differentiation）

物理ベースの微分可能レンダリング（Physically Based Differentiable Rendering; PBDR）

パストレーシング

典型的なレンダリングアルゴリズムのパストレーシング5は、光源からセンサーへと光の放射を伝播させるため、センサーから始めて連続するパスの頂点を繰り返し処理する
What Is Path Tracing? - NVIDIA

シーンが散乱メディア（入射光がその中を通過する際に方向を変えるか、強度を変える物質や物体）を含む場合、多くの反復処理が必要となる可能性があり、このループが微分の際に問題となる

Reverseモードの導関数を計算するためには、生成された各反復ごとに中間の状態を再び組み立てる必要がある

パストレーシングの場合、何回の反復になるか予測できないため、モンテカルロ法のサンプルが数百万回ループを並列に実行することとなり、その結果として生じる中間の状態を保存するためには大量のメモリが必要となる

一般的に、ループの自動微分はほとんど満足のいくものではなく、同等かつより効率的なカスタムアジョイントを実装するために、ユーザには専門的な知識が必要になる

カスタムアジョイント（custom adjoint）とは

カスタムアジョイントは、微分の文脈、特にReverseモード自動微分（reverse-mode automatic differentiation）の中で使用される用語

自動微分技術は、コンピュータプログラムを通じて関数の導関数を効率的に計算するための手法の一つ

この技術の中で、「アジョイント」とは、ある関数の出力に関するある入力の導関数を意味する

通常、基本的な演算や関数については既知の導関数を持っているが、特定の領域や複雑な関数に関しては自動で効率的に微分するのが難しい場合がある

このような場合、専門的な知識を持つユーザがその関数の微分を効率的に計算するための「カスタムアジョイント」を提供することが求められることがある

例えば、複雑な物理シミュレーションや特定の数値手法など、特定の計算手順の導関数が一般的な方法での自動微分では非効率的または不正確になる可能性がある場合、カスタムアジョイントを使用して、その計算の微分をより効率的または正確に行うことができる


放射線逆伝播（Radiative Backpropagation; RB）

カスタムアジョイントの一つである放射線逆伝播6は、物理的な相互作用を利用して、シミュレーションの微分を「アジョイント放射線」という等価なシミュレーションに置き換えることで、ループの状態を保存する必要をなくす

画像の損失を微分すると、最初にアジョイント放射線が画像空間で得られ、レンダリングされた画像のピクセルがどのように変化すれば損失を減少させるかを示す

ビデオプロジェクターのように、カメラはこの符号付きの放射線をシーンに放出し、通常の光のように散乱させる

従来のレンダリングがシーンの表現を読み取るのに対して、放射線逆伝播は書き込みを行う

アジョイント放射線が微分可能なパラメータ（BSDFパラメータなど - 4-1図の𝛼）を持つ面に衝突するたびに、そのパラメータの勾配（どのように調整すれば最も画像の損失を減少させるかの情報）を保存する

微分可能なパラメータ（4-1図の𝛼）に対する微分も入射放射線の量に比例するため、アジョイント放射線と基本放射線のシミュレーションを同時に考慮しなければならず、その複雑な計算とその計算コストの高さが課題となる

パスリプレイ逆伝播（Path Replay Backpropagation; PRB）

上記の課題を解決するため、パスリプレイ逆伝播7は以下の手順を行う

  1. モンテカルロ法を用いてランダムなサンプルデータを生成し、保存する

  2. 保存されたデータに算術の反転性を利用してサンプルを再生成し、シーンの各頂点の光の強度や方向（入射放射線）を復元する
    ->これにより、望む勾配を蓄積することができる

  3. 異なるサンプル間での依存関係を持たない独立したレンダリング（非相関の基本レンダリング）を計算し、それを損失関数に提供する

パスリプレイ逆伝播におけるオクルージョンの問題

描画の過程で物体が他の物体の背後に隠れて見えなくなる、または明るさや色の突然の変化など、シーン内には非連続性の変化が起こる可能性がある

この非連続性を持つシーンを正確に計算・評価するのは難しく、特に積分の過程での計算は複雑になる

微分可能なパラメータを変更して物体の配置や光の動きなどを変化させた際に、もし上述の非連続性や可視性の変化が起こると、微分結果が予想よりも大きく変わる、つまり「偏る」可能性がある

したがって、この問題を解決するための特別な予防措置や修正手法が必要になる


従来のような基本的な積分を計算するために考えられたサンプリング戦略を、微分された関数（導関数）に対して使うと、その性能や精度が下がってしまう

この場合、微分された関数（導関数）の特性をよく考慮して、その特性に合わせた新しいサンプリング戦略を考える必要がある

目的（Objectives）

微分やモンテカルロ積分を効果的に扱うためのシステムの詳細とその課題

システムの中心的な原則

カーネル間の送信コストを避けるため、モンテカルロ積分は複数のカーネルに分割してはならない

ループの繰り返し処理を直列的に実行するような大きなカーネルの生成を避けるため、ループやポリモーフィズムを持った呼び出しは保存されるべき

微分に関する課題

パスリプレイ逆伝播（PRB）は、各散乱の相互作用ごとに微分可能な計算を実行し、その後、関連する導関数を後方または前方に伝播させるための自動微分（AD）を呼び出す

微分可能な計算には多くのポリモーフィズムを持った関数呼び出しが含まれている

導関数は、滑らかな法線やMIPマップの計算のような典型的な前処理ステップを通過するべきである

Dr.Jitは、微分タスクを複数のフェーズに分割し、情報の効率的な交換を行う必要がある

トレースに関する詳細

Dr.Jitは、トレースに基づくアプローチで問題に対処している

このトレーシングは、より高いアーキテクチャレベルで行われ、JITコンパイラの基盤に依存している

この組み合わせは調和しており、Dr.Jitは微分アルゴリズムをシーンや問題文に動的に特化させることができる

しかし、Dr.Jitは開発者が微分過程の微妙な詳細について考えることから解放するものではない

トレースは、微分伝播の操作をキャプチャする計算グラフ、または(Wengert)テープを指す

動的コンパイルによるトレース（Tracing with dynamic compilation）

カスタムアジョイント（Custom adjoints）

微分ポリモーフィズム（Differentiating polymorphism）

分離境界（Isolation boundaries）

再パラメータ化（Reparameterizations）

偏微分の選択的評価（Selective evaluation of partial derivatives）

AD tape surgery

結果（Results）

結論（Conclusion）


  1. Wenzel Jakob, Sébastien Speierer, Nicolas Roussel, and Delio Vicini. 2022. Dr.Jit: A Just-In-Time Compiler for Differentiable Rendering. In Transactions on Graphics (Proceedings of SIGGRAPH) 41(4). ↩︎

  2. Merlin Nimier-David, Sébastien Speierer, Benoît Ruiz, and Wenzel Jakob. 2020. Radiative Backpropagation: An Adjoint Method for Lightning-Fast Differentiable Rendering. ACM Trans. Graph. (SIGGRAPH) 39, 4 (July 2020). ↩︎

  3. Delio Vicini, Sébastien Speierer, and Wenzel Jakob. 2021. Path Replay Backpropagation: Differentiating Light Paths Using Constant Memory and Linear Time. ACM Trans. Graph. (SIGGRAPH) 40, 4 (Aug. 2021). ↩︎

  4. Guillaume Loubet, Nicolas Holzschuch, and Wenzel Jakob. 2019. Reparameterizing Discontinuous Integrands for Differentiable Rendering. ACM Trans. Graph. (SIGGRAPH Asia) 38, 6 (Nov. 2019). ↩︎

  5. James T. Kajiya. 1986. The Rendering Equation. SIGGRAPH Comput. Graph. 20, 4 (Aug. 1986). ↩︎

  6. Merlin Nimier-David, Sébastien Speierer, Benoît Ruiz, and Wenzel Jakob. 2020. Radiative Backpropagation: An Adjoint Method for Lightning-Fast Differentiable Rendering. ACM Trans. Graph. (SIGGRAPH) 39, 4 (July 2020). ↩︎

  7. Delio Vicini, Sébastien Speierer, and Wenzel Jakob. 2021. Path Replay Backpropagation: Differentiating Light Paths Using Constant Memory and Linear Time. ACM Trans. Graph. (SIGGRAPH) 40, 4 (Aug. 2021). ↩︎