Scene Understanding: Leveraging Vulkan in the XR Pipeline

// 1. Ask OpenXR what it needs from Vulkan
XrGraphicsRequirementsVulkanKHR graphicsRequirements{XR_TYPE_GRAPHICS_REQUIREMENTS_VULKAN_KHR};
xrGetVulkanGraphicsRequirementsKHR(xrInstance, systemId, &graphicsRequirements);

// 2. Create the Vulkan Device with OpenXR-specific graphics binding
// This structure tells the OpenXR runtime which Vulkan device we are using
// so it can synchronize texture access between the OS and our app.
XrGraphicsBindingVulkanKHR graphicsBinding{XR_TYPE_GRAPHICS_BINDING_VULKAN_KHR};
graphicsBinding.instance = vkInstance;
graphicsBinding.physicalDevice = vkPhysicalDevice;
graphicsBinding.device = vkDevice;
graphicsBinding.queueFamilyIndex = queueFamily;
graphicsBinding.queueIndex = 0;

// 3. Link to XrSessionCreateInfo
XrSessionCreateInfo sessionCreateInfo{XR_TYPE_SESSION_CREATE_INFO};
sessionCreateInfo.next = &graphicsBinding;
sessionCreateInfo.systemId = systemId;
xrCreateSession(xrInstance, &sessionCreateInfo, &session);

// Managing Multiple Frames in Flight (XR Context)
struct FrameResources {
    vk::raii::Fence fence;
    vk::raii::Semaphore mlFinishedSemaphore; // Timeline Semaphore
    uint64_t mlValue;
};

void frameLoop() {
    xrWaitFrame(session, &waitInfo, &frameState);
    xrBeginFrame(session, &beginInfo);

    // 1. Acquire the Passthrough Texture
    XrSwapchainImageVulkanKHR xrImage;
    xrAcquireSwapchainImage(passthroughSwapchain, &acquireInfo, (XrSwapchainImageBaseHeader*)&xrImage);

    // 2. Submit ML Job to the COMPUTE Queue
    // We signal a Timeline Semaphore when the ML result is ready.
    uint64_t nextMlValue = ++globalMlCounter;

    vk::TimelineSemaphoreSubmitInfo timelineInfo;
    timelineInfo.setSignalSemaphoreValues(nextMlValue);

    vk::SubmitInfo submitInfo;
    submitInfo.setPNext(&timelineInfo);
    submitInfo.setSignalSemaphores(*mlTimelineSemaphore);

    computeQueue.submit(submitInfo);

    // 3. Render Virtual Content (GRAPHICS Queue)
    // We don't wait for the ML! We draw the virtual world immediately.
    // However, we use the ML result from 2-3 frames ago.
    uint64_t usableMlValue = std::max(0UL, nextMlValue - 2);
    renderVirtualWorld(frameState.predictedDisplayTime, usableMlValue);

    // 4. Compositor Submission
    xrEndFrame(session, &frameEndInfo);
}

#version 450
// Preprocessing both eyes in one go
layout(local_size_x = 16, local_size_y = 16, local_size_z = 2) in;

layout(set = 0, binding = 0) uniform sampler2DArray passthroughIn;
layout(set = 0, binding = 1) writeonly buffer Output { float data[]; } outTensor;

layout(push_constant) uniform Params {
    uint targetWidth;
    uint targetHeight;
} pc;

void main() {
    uint eyeIdx = gl_GlobalInvocationID.z; // 0 = Left, 1 = Right
    ivec2 pos = ivec2(gl_GlobalInvocationID.xy);

    if (pos.x >= pc.targetWidth || pos.y >= pc.targetHeight) return;

    // 1. Hardware-Accelerated Sample
    // uv coordinates are 0.0 to 1.0
    vec2 uv = (vec2(pos) + 0.5) / vec2(pc.targetWidth, pc.targetHeight);

    // Sample from the correct layer of the XR texture array
    vec3 rgb = texture(passthroughIn, vec3(uv, eyeIdx)).rgb;

    // 2. Normalization and Planar Transpose
    // Place into the NCHW tensor at the correct batch offset
    uint planeSize = pc.targetWidth * pc.targetHeight;
    uint batchStride = eyeIdx * (3 * planeSize);
    uint pixelIdx = pos.y * pc.targetWidth + pos.x;

    // Batching logic: ImageNet Normalization
    outTensor.data[batchStride + 0 * planeSize + pixelIdx] = (rgb.r - 0.485) / 0.229;
    outTensor.data[batchStride + 1 * planeSize + pixelIdx] = (rgb.g - 0.456) / 0.224;
    outTensor.data[batchStride + 2 * planeSize + pixelIdx] = (rgb.b - 0.406) / 0.225;
}

// SBS UV Logic
float splitX = (eyeIdx == 0) ? uv.x * 0.5 : 0.5 + uv.x * 0.5;
vec3 rgb = texture(passthroughSBS, vec2(splitX, uv.y)).rgb;

// Borrowing the texture for ML
void processML() {
    // 1. Wait for Compositor to release the image
    xrWaitSwapchainImage(passthroughSwapchain, &waitInfo);

    // 2. Transition layout for Compute
    vk::ImageMemoryBarrier barrier{
        .oldLayout = vk::ImageLayout::eUndefined,
        .newLayout = vk::ImageLayout::eShaderReadOnlyOptimal,
        .image = passthroughImage,
        // ...
    };

    // 3. Run ML Dispatch
    // ...

    // 4. Release back to Compositor
    xrReleaseSwapchainImage(passthroughSwapchain, &releaseInfo);
}

// Optimized Variable-Rate Blur Kernel
void main() {
    float mask = texture(mlMaskSampler, uv).r;
    vec3 color = texture(passthroughSampler, uv).rgb;

    // Adaptive Blur: Only compute the expensive blur if needed
    if (mask > 0.5) {
        vec3 blurred = vec3(0.0);
        float totalWeight = 0.0;
        float sigma = 10.0 * mask; // Blur intensity scales with mask confidence

        for(float x = -2.0; x <= 2.0; x += 1.0) {
            for(float y = -2.0; y <= 2.0; y += 1.0) {
                float weight = exp(-(x*x + y*y) / (2.0 * sigma * sigma));
                blurred += texture(passthroughSampler, uv + vec2(x, y) * texelSize).rgb * weight;
                totalWeight += weight;
            }
        }
        color = blurred / totalWeight;
    }
    outColor = vec4(color, 1.0);
}