ガラスキューブの描画 反射と屈折

Ray tracingの環境が整ったのならば、一度はやってみたいガラスや水の透過、反射、屈折。
ということでvulkan ray tracingでガラスキューブで反射と屈折の描画をしてみたので、その間に調べたことや難しかった点などをまとめておきます。
(一番躓いた点は法線の計算がバグってたという、全く関係のない部分でした;;)
ソースコードの共有というよりはその周りの話がメインとなるので、Qiitaではなくこちらに投稿します。
結果はこんな感じ


f:id:Aqoole_Hateena:20210605201025p:plain
f:id:Aqoole_Hateena:20210605201101p:plain
f:id:Aqoole_Hateena:20210605201123p:plain
f:id:Aqoole_Hateena:20210605201150p:plain

オブジェクトは平面とガラスキューブだけという、非常にシンプルでray tracing学習者丸出しの内容です。

さて、このトレースではガラスキューブの入射光から、反射光と屈折光という2種類の計算を行っているのですが、何が難しいかと言うと
・シェーダ内で関数traceRayEXT()は再帰呼び出しできない(現状AMD GPUのみ?)
・GLSLでは一般の関数でも再帰呼び出しはできない(or 非推奨)
・しかしながら、ガラスに当たる光の計算は1種類の入射光から反射光と屈折光の2種類が生まれる

はい、そういうことです。
ガラス面に当たるたびに、ひとつの入射光から2つの光が生まれるので、2の指数の速さで光の本数が増えていき、コードの実装が大変です。例えばC++だとこのようなケースでは再帰関数で書いておいて、進んだり戻ったりを繰り返しトレースできそうなものですが、GLSLでかくシェーダになると再帰関数は使えないそうです。(ソースは若干古いです)

https://stackoverflow.com/questions/43601521/recursion-in-glsl-prohibited/43601658

GLSLで再帰関数が使えないのならtraceRayEXT関数を再帰的に使えばいいじゃない、という発想になります。つまりtraceRayEXT()でヒットした場合の呼び出し先のシェーダ内でさらに同じtraceRayEXT()を使うということですね。しかし結論的に、これはどうもできないようです。微妙な言い方になるのは、下記のNVIDIAの記事だと「遅くはなるが動く」と言っているようですが、私のAMDGPUだとVkPhysicalDeviceRayTracingPipelinePropertiesKHR内のmaxRayRecursionDepthという値を調べると1だったので、1回しか呼べないということになっています。つまり、再帰呼び出しできません。

https://github.com/nvpro-samples/vk_raytracing_tutorial_KHR/tree/master/ray_tracing_reflections

仕方がないので、再帰関数でスマートに書くことは諦めてトレースするレイの本数を直打ちで決めてトレースしました。最終的にはトレースする光線の数は2の指数の和になります。
この例ではdepth = 6, つまり6回光がガラス面にあたって反射/屈折した場合のトレースをしているので、1ピクセルあたり


\begin{eqnarray}
\sum_{k=1}^{6}2^k &=& 2^7 - 2 \\
&=& 126\
\end{eqnarray}

の光をトレースしています。
あらかじめ

vec3[126] colors;

に126本の光線のトレースをしてヒット先の色を格納した後、反射光の強度や全反射の有無にあわせて

/*
for(uint i = 62; i < 126; i++)
{
  uint reflectIndex = 2 * i + 2;
  uint refractIndex = 2 * i + 3;
  colors[i] = ColorBlend(reflection, colors[i], isPlane[i], isMiss[reflectIndex], colors[reflectIndex], isMiss[refractIndex], colors[refractIndex], isAllReflect[refractIndex]);
}
*/
for(uint i = 30; i < 62; i++)
{
  uint reflectIndex = 2 * i + 2;
  uint refractIndex = 2 * i + 3;
  colors[i] = ColorBlend(reflection, colors[i], isPlane[i], isMiss[reflectIndex], colors[reflectIndex], isMiss[refractIndex], colors[refractIndex], isAllReflect[refractIndex]);
}

for(uint i = 14; i < 30; i++)
{
  uint reflectIndex = 2 * i + 2;
  uint refractIndex = 2 * i + 3;
  colors[i] = ColorBlend(reflection, colors[i], isPlane[i], isMiss[reflectIndex], colors[reflectIndex], isMiss[refractIndex], colors[refractIndex], isAllReflect[refractIndex]);
}

for(uint i = 6; i < 14; i++)
{
  uint reflectIndex = 2 * i + 2;
  uint refractIndex = 2 * i + 3;
  colors[i] = ColorBlend(reflection, colors[i], isPlane[i], isMiss[reflectIndex], colors[reflectIndex], isMiss[refractIndex], colors[refractIndex], isAllReflect[refractIndex]);
}
for(uint i = 2; i < 6; i++)
{
  uint reflectIndex = 2 * i + 2;
  uint refractIndex = 2 * i + 3;
  colors[i] = ColorBlend(reflection, colors[i], isPlane[i], isMiss[reflectIndex], colors[reflectIndex], isMiss[refractIndex], colors[refractIndex], isAllReflect[refractIndex]);
}
for(uint i = 0; i < 2; i++)
{
  uint reflectIndex = 2 * i + 2;
  uint refractIndex = 2 * i + 3;
  colors[i] = ColorBlend(reflection, colors[i], isPlane[i], isMiss[reflectIndex], colors[reflectIndex], isMiss[refractIndex], colors[refractIndex], isAllReflect[refractIndex]);
}
return ColorBlend(reflection, surfaceColor, false, isMiss[0], colors[0], isMiss[1], colors[1], isAllReflect[1]);
vec3 ColorBlend(float reflection, vec3 ownColor, bool isPlane, bool reflectMiss, vec3 reflectColor, bool refractMiss, vec3 refractColor, bool isAllReflect)
{
  if(isPlane)
    return ownColor;
  if(reflectMiss && refractMiss)
    return vec3(0.0, 0.0, 1.0);
  if(reflectMiss)
  {
    if(isAllReflect)
      return ownColor;
    else 
      return mix(refractColor, ownColor, 0.1);
  }
  if(refractMiss)
    return mix(reflectColor, ownColor, 0.1);
  //both hit
  if(isAllReflect)
    return mix(reflectColor, ownColor, 0.1);
  else
  {
    vec3 traceColor = mix(refractColor,reflectColor,reflection);
    return mix(traceColor,ownColor,0.1);
  }
}

などとして色をブレンドしていきます。
この場合だとガラスから抜け出て平面にヒットした場合でも構わずトレースし続けているので改善の余地はありますが、予めトレースする光線の数を決めておくことに変わりはありません。GLSLではmalloc/newのように動的にメモリを確保する手段がないはずなので。
最初の/* */でコメントアウトしている部分はその次のdepthまでトレースした場合の処理ですが、ここまで行うと、どうやらvulkanの側でpipelineを作成する際にVK_ERROR_UNKNOWN = -13でエラーしていましたので、その一つ前のdepthで諦めています。fps的にはdepthが6の時点で400〜700fpsなので、もう少し余裕はあったかなと思ってます。

colors[126]にヒットした先の色が取ってこれたとして、次のColorBlend()の処理で色を合成しています。この処理はdepthが深いところの色を決めてから、その色でもって浅いところの色が決まっていきます。
ルールとしては
・平面にヒットした場合は、前の結果(よりdepthが大きい位置)にかかわらず平面の色を返す
・反射/屈折がミス(どこにもヒットしない)した場合、どちらかの色をガラス面とミックスして返す
・全反射が起きた場合、屈折光の色をガラス面とミックスして返す
・反射/屈折の両方がヒットしている場合は、反射率だけ反射光と屈折光の色を合成し、それをガラス面とミックスして返す
ガラス面とミックスする部分は、色の綺麗さなどに合わせてお好みかもしれません。
反射率は以下のサイトを参考に計算しています。
http://skomo.o.oo7.jp/f17/hp17_9.htm
https://ja.wikipedia.org/wiki/%E5%8F%8D%E5%B0%84%E7%8E%87

イデアとしてはたったこれだけなんですが、試行錯誤だったり何だったりで非常に時間がかかりました。とりあえず結果が得られたということで、満足しています。