[Vulkan ray tracing] sampleコードからの次の一歩

AMDのGPUでray tracingがしたいと思い、GWでvulkanのray tracingを学んでいた。探してみたがあまりQiitaにも記事がないので、hellow worldからの次の一歩をメモとして残しておこうと思う。
ray tracingでのhello world(hello triangle)のコードは公式のkhronosが公開しているsample codeを参考にすれば良いと思う。
https://github.com/KhronosGroup/Vulkan-Samples/tree/master/samples/extensions/raytracing_basic

ray tracingでの相違点

ある程度vulkanに馴染んでいることを前提としているが、うまくビルドできれば三角形が表示される。普通のレンダリングとray tracingの違いとして驚いたのは、shaderの出力結果がまったく異なるということ。通常のvert/fragのシェーダの出力として期待するのは各頂点ごとのベクトルであるが、ray tracingではraygen shaderが計算した結果はimageとして出力される。そのimageをディスプレイするimageと差し替えて(swapchain)、最終的に結果を描画する。

次ステップ

sample codeでは2次元の三角形を出力するだけなので、ライティングの効果はあまり実感できない。次ステップとして、次のサイトを参考にさせていただいた。
https://nvpro-samples.github.io/vk_raytracing_tutorial_KHR/#accelerationstructure/bottom-levelaccelerationstructure

ここではキューブの出力と、raygen/miss/rchitの各種shaderの応用が記述されている。
先のsample codeを参考にした場合に必要となる変更点は以下

bottom level acceleration structureの変更

primitive count

sample codeでは三角形を出力するだけなので、geometry情報が直打ちになっている。この箇所を修正する。

uint32_t primitiveCount = indicesCount / 3;
//build range info
VkAccelerationStructureBuildRangeInfoKHR mRangeInfo = {};
mRangeInfo.primitiveCount = primitiveCount;
mRangeInfo.primitiveOffset = 0;
mRangeInfo.firstVertex = 0;
mRangeInfo.transformOffset = 0;

update flag

sampleコードではtransform行列の変更によって画面を更新するということが考慮されていない。まずAS(Acceleration Structure)の情報の更新は初回のビルド時と同じく、ビルドコマンドで更新する。

pfnCmdBuildAccelerationStructuresKHR(
    commandBuffer, 1, &mBuildCommandGeometryInfo, rangeInfos.data());

pfnはファンクションポインタである。vulkanではextensionで提供される関数を用いる場合には、アプリに関数をポインタで渡してやる必要がある。

pfnCmdBuildAccelerationStructuresKHR = 
    reinterpret_cast<PFN_vkCmdBuildAccelerationStructuresKHR>(
        vkGetDeviceProcAddr(device,"vkCmdBuildAccelerationStructuresKHR"));

ASの情報を更新の上、この関数を再度用いることでASの更新を行うのだが、その際にASをビルドする時点で、フラグを設定しておく必要がある。

//build geometry info
mBuildCommandGeometryInfo = {};
mBuildCommandGeometryInfo.sType = 
    VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
mBuildCommandGeometryInfo.type = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR;
mBuildCommandGeometryInfo.flags =
    VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR |
    VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_UPDATE_BIT_KHR;
mBuildCommandGeometryInfo.mode = VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR;
mBuildCommandGeometryInfo.dstAccelerationStructure = mAS;
mBuildCommandGeometryInfo.geometryCount = 1;
mBuildCommandGeometryInfo.pGeometries = &mGeometry;
mBuildCommandGeometryInfo.scratchData.deviceAddress =
    scratchBufferDeviceAddress.deviceAddress;

gflagsのメンバについてVK_BUILD_ACCELERATION_STRUCTURE_ALLOW_UPDATE_BIT_KHRも設定しておかないと、更新後プログラムがハングし、PC自体が操作不能となったので注意が必要である。

push constants

lightの位置などの情報はどのshaderでも同じデータを扱うので、push constantsで提供する。

VkPushConstantRange pushConstant{};
pushConstant.stageFlags = VK_SHADER_STAGE_RAYGEN_BIT_KHR |
                          VK_SHADER_STAGE_MISS_BIT_KHR |
                          VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR;
pushConstant.offset = 0;
pushConstant.size = size;/*constantsのサイズ*/
std::vector<VkPushConstantRange> pushConstants = {pushConstant};
//pipelinelayoutinfoに追加
pipelineLayoutInfo.pushConstantRangeCount = pushConstants->size();
pipelineLayoutInfo.pPushConstantRanges = pushConstants->data();
if (vkCreatePipelineLayout(device, &pipelineLayoutInfo, nullptr, &mPipelineLayout) !=
    VK_SUCCESS)
        std::runtime_error("failed to create pipeline layout");

rchitシェーダの変更

sampleコードでは、頂点ごとの色の計算はrchitシェーダで行っている。これを次ステップの例に従って変更する。以下はmain関数だけを抜きだしたもの。

closesthit.rchit

void main()
{
  ivec3 ind = ivec3(indices[0].i[3 * gl_PrimitiveID + 0],   //
                    indices[0].i[3 * gl_PrimitiveID + 1],   //
                    indices[0].i[3 * gl_PrimitiveID + 2]);  //
  Vertex3D v0 = vertices[0].v[ind.x];
  Vertex3D v1 = vertices[0].v[ind.y];
  Vertex3D v2 = vertices[0].v[ind.z];　/*Vertrex3Dはユーザ定義しているstruct*/
  const vec3 barycentricCoords = vec3(1.0f - attribs.x - attribs.y, attribs.x, attribs.y);
  vec3 normal = v0.normal * barycentricCoords.x + 
                v1.normal * barycentricCoords.y +
                v2.normal * barycentricCoords.z;
  // camはuniform bufferで渡しているstruct
  vec4 normalP = normalize(cam.modelViewProj * vec4(normal, 0.0));
  vec3 worldPos = v0.pos * barycentricCoords.x +
                  v1.pos * barycentricCoords.y +
                  v2.pos * barycentricCoords.z;
  vec4 worldP = cam.modelViewProj * vec4(worldPos, 0.0);
  vec3 L;
  float lightIntensity = pushC.lightIntensity;
  float lightDistance = 10000.0;
  vec3 lDir = pushC.lightPosition - worldP.xyz;
  lightDistance = length(lDir);
  lightIntensity = lightIntensity / (lightDistance * lightDistance);
  L = normalize(lDir);
  float dotNL = min(lightIntensity * max(dot(normalP.xyz, L), 0.02), 1.0);
  hitValue = vec3(dotNL);
}

実行すると次のように、立方体(直方体にしか見えない)に光が当たっている例が得られる。

終わりに

この記事ではsampleコードからどのようにして次のステップへ進むのかを、vulkanのコードベースに記事にしてみた。ここではc++のコードの話しかしていないがvulkan APIでのray tracingの仕組みももちろん知っておいたほうが良いので、khronos公式のページは目を通しておいたほうが良いと思う。
https://www.khronos.org/blog/ray-tracing-in-vulkan
外に出る時間を極力減らす連休となってしまったが、この期間に得た知識を成果物として残しておきたかったので、何か出力できたものをここに残したいとの動機から今回記事を投稿することにした。欲を言えばもう少し難しいモデルにして、背景や光の当たり具合もきれいにしたかったが、仕組みもコードも思った以上に難しく、連休中にできたのはここまでとなった。同じ箇所で躓いている人にとって、役立つことを願っています。