Using the ML Workload Library for Vulkan®

Execution model

The library can create a runtime-owned Vulkan® context, or it can wrap Vulkan® objects supplied by the application without taking ownership. When wrapping objects, the application must enable the Vulkan® extensions and features required by the workload. Resources can likewise be supplied by the application or allocated through the context from workload resource requirements.

A workload is executed through the following sequence:

  1. Construct a Workload.

  2. Create a runtime-owned Context, or wrap the application’s Vulkan® instance, physical device, device, queue family, and queue.

  3. Create a Session, provide any missing module implementations, and configure it.

  4. Create a BindingSet and bind each public resource using caller-owned or runtime-owned allocations.

  5. Prepare a binding snapshot and either run it or record it into an application command buffer.

The following excerpt from the standalone compute sample shows the common flow with a runtime-owned context and buffers. The sample defines the GLSL workload description and mapped-memory helpers earlier in the source file.

const std::vector<int32_t> lhs = {1, 2, 3, 4};
const std::vector<int32_t> rhs = {10, 20, 30, 40};

const auto workload = Workload::fromComputeShader(addBuffersDescription(lhs.size()));
auto context = Context::create();
const auto contextView = context.contextView();

auto lhsBuffer = context.createBuffer(workload.resource(0));
auto rhsBuffer = context.createBuffer(workload.resource(1));
auto outputBuffer = context.createBuffer(workload.resource(2));

writeMemory(contextView.device, lhsBuffer.memory(), lhs);
writeMemory(contextView.device, rhsBuffer.memory(), rhs);
clearMemory(contextView.device, outputBuffer.memory());

Session session(context, workload);
session.configure();
auto bindings = session.createBindingSet();
bindings.bindBuffer(workload.resource(0), {lhsBuffer.handle(), lhsBuffer.memory()});
bindings.bindBuffer(workload.resource(1), {rhsBuffer.handle(), rhsBuffer.memory()});
bindings.bindBuffer(workload.resource(2), {outputBuffer.handle(), outputBuffer.memory()});

auto execution = session.prepare(bindings);
execution.run();

An application can instead retain ownership of its Vulkan® objects and wrap them in a library context:

const ApplicationVulkanContext applicationVulkan;
auto context = Context::wrap(applicationVulkan.view());

Complete programs showing VGF inspection and execution, standalone compute execution, and application-owned Vulkan® context wrapping are available in the samples directory.

Workload sources

Workload::fromVGF(...) loads a VGF file or decodes a caller-owned memory buffer. Workload::fromComputeShader(...) constructs a standalone compute workload, and Workload::fromDataGraph(...) constructs a standalone Vulkan® data-graph workload.

The standalone compute samples describe their module, dispatch, and resources directly:

inline ComputeShaderDescription addBuffersDescription(std::size_t elementCount) {
    ComputeShaderDescription description;
    description.module.codeKind = ModuleCodeKind::Glsl;
    description.module.source = addBuffersGlsl;
    description.dispatch = {static_cast<uint32_t>(elementCount), 1, 1};

    const auto byteSize = static_cast<vk::DeviceSize>(elementCount * sizeof(int32_t));
    description.resources = {
        {"lhs", 0, 0, ResourceAccess::Read, bufferRequirements(byteSize)},
        {"rhs", 0, 1, ResourceAccess::Read, bufferRequirements(byteSize)},
        {"output", 0, 2, ResourceAccess::Write, bufferRequirements(byteSize)},
    };
    return description;
}

Standalone data-graph workloads use the same resource model and add graph pipeline metadata. Sample 5 leaves the module implementation missing so it can be supplied to a session later:

DataGraphDescription description;
description.module.codeKind = ModuleCodeKind::Missing;
description.entryPoint = "main";
description.resources = {
    {"input", 0, 0, ResourceAccess::Read, tensorRequirements(vk::Format::eR8Sint, {1, 16, 16, 16})},
    {"output", 1, 1, ResourceAccess::Write, tensorRequirements(vk::Format::eR8Sint, {1, 8, 8, 16})},
};
description.pipeline.identifier = "maxpool";

const auto workload = Workload::fromDataGraph(std::move(description));

Modules can contain SPIR-V™, or GLSL and HLSL source when the matching optional backend is available. GLSL and HLSL modules are supported for compute executables. Check an optional backend with supports(Feature::GlslModules) or supports(Feature::HlslModules).

if (!supports(Feature::GlslModules)) {
    throw std::runtime_error("The sample requires GLSL module support");
}

A workload can also contain placeholder modules. Sample 3 binds a GLSL implementation before configuring its session:

ModuleImplementation implementation;
implementation.codeKind = ModuleCodeKind::Glsl;
implementation.source = addBuffersGlsl;
session.bindModule(workload.placeholderModule(0), std::move(implementation));
session.configure();

The VGF samples use a shared helper to encode an in-memory VGF:

inline std::string addBuffersVgf(std::size_t elementCount, bool embedImplementation = true) {
    if (elementCount > std::numeric_limits<uint32_t>::max()) {
        throw std::invalid_argument("The element count exceeds the VGF dispatch limit");
    }

    auto encoder = mlsdk::vgflib::CreateEncoder(VK_HEADER_VERSION);
    const auto module = embedImplementation
                            ? encoder->AddModule(mlsdk::vgflib::ModuleType::COMPUTE, "add_int32_buffers", "main",
                                                 mlsdk::vgflib::ShaderType::GLSL, std::string(addBuffersGlsl))
                            : encoder->AddModule(mlsdk::vgflib::ModuleType::COMPUTE, "add_int32_buffers", "main");

    const auto shape = std::vector<int64_t>{static_cast<int64_t>(elementCount)};
    const auto strides = std::vector<int64_t>{static_cast<int64_t>(sizeof(int32_t))};
    const auto lhs = encoder->AddInputResource(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_FORMAT_R32_SINT, shape, strides);
    const auto rhs = encoder->AddInputResource(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_FORMAT_R32_SINT, shape, strides);
    const auto output =
        encoder->AddOutputResource(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_FORMAT_R32_SINT, shape, strides);

    const auto lhsBinding = encoder->AddBindingSlot(0, lhs);
    const auto rhsBinding = encoder->AddBindingSlot(1, rhs);
    const auto outputBinding = encoder->AddBindingSlot(2, output);
    const auto descriptorSet = encoder->AddDescriptorSetInfo({lhsBinding, rhsBinding, outputBinding}, 0);
    const std::vector<mlsdk::vgflib::GraphConstantBindingRef> noGraphConstants;
    encoder->AddSegmentInfo(module, "add_int32_buffers_segment", {descriptorSet}, {lhsBinding, rhsBinding},
                            {outputBinding}, noGraphConstants, {static_cast<uint32_t>(elementCount), 1, 1});
    encoder->AddModelSequenceInputsOutputs({lhsBinding, rhsBinding}, {"lhs", "rhs"}, {outputBinding}, {"output"});
    encoder->Finish();

    std::stringstream stream;
    if (!encoder->WriteTo(stream)) {
        throw std::runtime_error("Failed to encode the sample VGF");
    }
    return stream.str();
}

The encoded bytes remain alive while Workload decodes their caller-owned memory:

const auto vgf = addBuffersVgf(4);
const auto workload = Workload::fromVGF(vgf.data(), vgf.size());

Workload inspection

A workload exposes non-owning views of its public resources. Sample 1 inspects each resource’s identity, access, kind, and kind-specific requirements:

void printResource(ResourceView resource) {
    const auto requirements = resource.requirements();
    std::cout << "  [" << resource.index() << "] " << resource.name() << ": " << resourceKindName(requirements.kind())
              << ", " << accessName(resource.access());

    switch (requirements.kind()) {
    case ResourceKind::Tensor: {
        std::cout << ", shape=[";
        const auto shape = requirements.asTensor().shape();
        for (std::size_t i = 0; i < shape.size(); ++i) {
            std::cout << (i == 0 ? "" : ", ") << shape[i];
        }
        std::cout << "]";
        break;
    }
    case ResourceKind::StorageBuffer:
        std::cout << ", bytes=" << requirements.byteSize();
        break;
    case ResourceKind::Image: {
        const auto extent = requirements.asImage().extent();
        std::cout << ", extent=" << extent.width << "x" << extent.height << "x" << extent.depth;
        break;
    }
    case ResourceKind::Unknown:
        break;
    }
    std::cout << '\n';
}

Executables similarly expose their module and descriptor interface:

void printExecutable(ExecutableView executable) {
    const auto module = executable.module();
    std::cout << "  [" << executable.index() << "] " << executable.name() << ": "
              << executableKindName(executable.type()) << ", module=" << module.name()
              << ", entry-point=" << module.entryPoint() << '\n';

    for (uint32_t i = 0; i < executable.interfaceDescriptorBindingCount(); ++i) {
        const auto binding = executable.interfaceDescriptorBinding(i);
        std::cout << "      set=" << binding.set << ", binding=" << binding.binding
                  << ", resource=" << binding.resourceIndex << ", access=" << accessName(binding.access) << '\n';
    }
}

Resource binding

The workload exposes its public resources in order through Workload::resources() or by index through Workload::resource(...). Inspect each resource’s kind and requirements before creating the corresponding Vulkan® object, or use Context::createTensor(), Context::createBuffer(), or Context::createImage() to create a compatible allocation. Then use bindTensor(), bindBuffer(), or bindImage(). Every required public resource must be bound before preparing an execution.

Sample 3 applies this process to every public resource according to its kind:

struct RuntimeResources {
    std::vector<TensorAllocation> tensors;
    std::vector<BufferAllocation> buffers;
    std::vector<ImageAllocation> images;
};

RuntimeResources bindRuntimeResources(Context &context, const Workload &workload, BindingSet &bindings) {
    RuntimeResources allocations;
    const auto contextView = context.contextView();

    for (const auto resource : workload.resources()) {
        const auto requirements = resource.requirements();
        if (requirements.participatesInAliasing()) {
            throw std::runtime_error("The sample cannot allocate aliased resource " + std::to_string(resource.index()));
        }

        switch (requirements.kind()) {
        case ResourceKind::Tensor: {
            auto allocation = context.createTensor(resource);
            clearMemory(contextView.device, allocation.memory());
            bindings.bindTensor(resource, {allocation.handle(), allocation.memory()});
            allocations.tensors.push_back(std::move(allocation));
            break;
        }
        case ResourceKind::StorageBuffer: {
            auto allocation = context.createBuffer(resource);
            clearMemory(contextView.device, allocation.memory());
            bindings.bindBuffer(resource, {allocation.handle(), allocation.memory()});
            allocations.buffers.push_back(std::move(allocation));
            break;
        }
        case ResourceKind::Image: {
            if (requirements.asImage().requiresSamplerBinding()) {
                throw std::runtime_error("The sample cannot provide a caller-owned sampler for image resource " +
                                         std::to_string(resource.index()));
            }
            auto allocation = context.createImage(resource);
            bindings.bindImage(resource, allocation.binding());
            allocations.images.push_back(std::move(allocation));
            break;
        }
        case ResourceKind::Unknown:
            throw std::runtime_error("Unsupported resource kind at index " + std::to_string(resource.index()));
        }
    }
    return allocations;
}

Supply BoundMemoryInfo when ResourceRequirementsView::requiresBoundMemoryInfo() is true. Image bindings must provide an image view and the required subresource range, and must provide a sampler only when the image requirements request one. An image layout is optional: when supplied, it describes the current layout and prepared execution transitions the image to the workload-required layout. When omitted, the application manages the image layout.

Running and recording

PreparedExecution::run() records, submits, and waits for completion using library-managed command and fence state.

PreparedExecution::record() records into a caller-owned command buffer. The application is responsible for the command-buffer lifecycle, submission, and synchronization.

const auto &device = contextView.device.get();
const auto &queue = contextView.queue.get();
const vk::raii::CommandPool commandPool(
    device, {vk::CommandPoolCreateFlagBits::eResetCommandBuffer, contextView.queueFamilyIndex});
auto commandBuffer =
    std::move(device.allocateCommandBuffers({*commandPool, vk::CommandBufferLevel::ePrimary, 1}).front());
const vk::raii::Fence fence(device, vk::FenceCreateInfo{});

commandBuffer.begin({vk::CommandBufferUsageFlagBits::eOneTimeSubmit});
execution.record(*commandBuffer);
commandBuffer.end();

const vk::SubmitInfo submitInfo({}, {}, *commandBuffer);
queue.submit(submitInfo, *fence);
if (device.waitForFences(*fence, true, std::numeric_limits<uint64_t>::max()) != vk::Result::eSuccess) {
    throw std::runtime_error("Timed out waiting for the recorded workload");
}

Ownership and lifetimes

  • The application retains ownership of Vulkan® objects and resources it supplies. Runtime allocations own the Vulkan® resource and backing memory they expose.

  • All bound resources, whether caller-owned or runtime-owned, must remain valid until submitted work has completed.

  • A runtime-owned Context must outlive its sessions, allocations, and any submitted work that uses them.

  • The Vulkan® objects wrapped by Context must outlive the context, its sessions, allocations created through it, and submitted work that uses them.

  • A Workload must outlive its non-owning views and every Session created from it. A Session must outlive its binding sets and prepared executions.

  • VGF memory passed to Workload::fromVGF(data, size) must remain valid for the workload lifetime. The same applies to caller-owned constant payloads in a standalone data-graph description.

  • Session::prepare() snapshots a binding set. Later changes to that binding set do not change the prepared execution.