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. Similarly, the application can supply resources, or the context can allocate them based on the 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, bind any needed placeholder modules or data-graph binaries, then call Session::configure().

    • For a placeholder module, call Session::bindModule(workload.placeholderModule(index), implementation).

    • For an external data-graph binary, call Session::bindDataGraphBinary(binary). Both standalone and VGF graphs can specify a storage mode and pipeline selector. The final executable index defaults to 0; use it to select another graph in a VGF workload. Multiple executables can share one binary with different selectors or use separate binaries.

    A placeholder module and binary cannot target the same module. Unselected graph and compute executables continue to use their workload modules. A VGF graph selected for binary execution may omit its shader module.

  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.

By default, each Session owns its compiled pipelines and layouts. To share them across sessions for the same context and workload, pass a CompiledExecution to each session. Each session still has its own binding sets and command state.

Each session also has its own tiling selection. Building shared pipelines fixes graph tensor tiling and the tiling of its alias peers. Sessions created afterward start with those choices; other tiling choices remain independent. A session created earlier must select compatible graph tensor tiling before it can reuse the pipelines.

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));

utils::writeDeviceMemory(contextView.device, lhsBuffer.memory(), lhs);
utils::writeDeviceMemory(contextView.device, rhsBuffer.memory(), rhs);
utils::clearDeviceMemory(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:

RuntimeContextDeviceRequirements deviceRequirements;
deviceRequirements.capabilities = RuntimeDeviceCapability::Compute;
const utils::ApplicationContext applicationContext(deviceRequirements);
auto context = Context::wrap({applicationContext.vulkanInstance(), applicationContext.vulkanPhysicalDevice(),
                              applicationContext.vulkanDevice(), applicationContext.vulkanQueueFamilyIndex(),
                              applicationContext.vulkanQueue()});

Complete programs showing workload inspection, VGF and 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. 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, utils::bufferRequirements(byteSize)},
        {"rhs", 0, 1, ResourceAccess::Read, utils::bufferRequirements(byteSize)},
        {"output", 0, 2, ResourceAccess::Write, utils::bufferRequirements(byteSize)},
    };
    return description;
}

Standalone data-graph workloads use the same resource model and add graph pipeline metadata. Sample 5 leaves its module implementation unspecified so the Session can bind a module or binary later:

DataGraphDescription description;
description.module.codeKind = ModuleCodeKind::Missing;
description.entryPoint = "main";
description.resources = {
    {"input", 0, 0, ResourceAccess::Read, utils::tensorRequirements(vk::Format::eR8Sint, {1, 16, 16, 16})},
    {"output", 1, 1, ResourceAccess::Write, utils::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 execution sample encodes a small in-memory VGF with a placeholder module:

std::string addBuffersVgf(std::size_t elementCount) {
    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 = 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(lhs.size());
const auto workload = Workload::fromVGF(vgf.data(), vgf.size());

Workload inspection

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

void printResource(ResourceView resource) {
    const auto requirements = resource.requirements();
    std::cout << "  [" << resource.index() << "] " << resource.name() << ": "
              << utils::resourceKindName(requirements.kind()) << ", " << utils::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';
}

Similarly, executables expose their module and descriptor interface:

void printExecutable(ExecutableView executable) {
    const auto module = executable.module();
    std::cout << "  [" << executable.index() << "] " << executable.name() << ": "
              << utils::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=" << utils::accessName(binding.access) << '\n';
    }
}

Resource binding

The workload exposes its public resources in order through Workload::resources() or by index through Workload::resource(...).

Choose public tensor or image tiling with Session::setTensorTiling() or Session::setImageTiling() before calling configure(). Without a selection, tensors use Linear tiling and images use Optimal tiling. A tensor that requires image aliasing uses Optimal. Alias groups use a common default tiling: Optimal if any peer defaults to Optimal, unless a storage-buffer peer or explicit tensor strides require Linear. Image/tensor alias groups are rejected if a tensor peer’s innermost dimension exceeds four, regardless of tiling. Incompatible alias tiling requirements are rejected. For Optimal image/tensor aliases, the required tensor usage includes eImageAliasing.

The allocation overloads without a tiling argument create Linear tensors and Optimal images. ImageRequirementsView::usage() returns the usage for Optimal tiling.

To create and bind resources:

  1. Create each public resource from its kind and requirements, or use Context::createTensor(), Context::createBuffer(), or Context::createImage(). For tensors and images, use the tiling returned by session.tensorTiling(resource) or session.imageTiling(resource) when creating the Vulkan® object; pass it to the Context helper if using one. For a tensor created directly, query its required usage with resource.requirements().asTensor().usage(session.tensorTiling(resource)). For an image created directly, query its required usage with resource.requirements().asImage().usage(session.imageTiling(resource)).

  2. Bind every required public resource with bindTensor(), bindBuffer(), or bindImage() before preparing an execution.

BindingSet cannot inspect a caller-owned Vulkan® handle’s creation tiling.

Sample 3 binds 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, const Session &session,
                                      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, session.tensorTiling(resource));
            utils::clearDeviceMemory(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);
            utils::clearDeviceMemory(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, session.imageTiling(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.

Push constants

Call BindingSet::bindPushConstants(executable, data) before preparing an execution to supply push constants for an executable from the same workload. Each call copies the bytes and replaces that executable’s previous payload, without changing other executables’ payloads.

The payload size must equal the largest offset + size of the executable’s compiled push-constant ranges, including any leading gap. For example, a four-byte range at offset 16 requires a 20-byte payload. Empty data is accepted only when the executable requires zero bytes; it cannot clear a nonempty payload.

For a configured session whose first executable expects a single uint32_t at offset zero, using a binding set with its required resources already bound:

std::uint32_t value = 42;
const auto executable = workload.executable(0);

bindings.bindPushConstants(
    executable,
    {reinterpret_cast<const std::byte *>(&value), sizeof(value)});
auto execution = session.prepare(bindings);

The source bytes need not remain valid after bindPushConstants() returns. prepare() captures the bound payloads, so later binding changes do not affect an already prepared execution. bindings.pushConstants(executable) returns a copy of the currently bound payload.

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");
}

Pass RecordOptions to the record() overload to record caller-owned commands immediately before and after each executable dispatch. The hook receives the recording command buffer and an ExecutableView whose index, name, and kind identify the executable, together with an ExecutableRecordPhase indicating whether it runs before or after the dispatch. This supports timestamps and debug labels without transferring ownership of their Vulkan® objects to the library.

Pipeline options

SessionOptions accepts a pipeline cache and optional compute-pipeline robustness and neural statistics settings. The application must enable the corresponding Vulkan® extensions and device features before requesting them. Neural statistics settings are applied during graph pipeline and graph session creation. The pipeline cache is not used for binary-backed pipeline creation.

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.

  • CompiledExecution is tied to the Context and Workload supplied at creation. Sessions keep its compiled state alive if the CompiledExecution object is destroyed; the context and workload must still outlive those sessions.

  • 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.

  • The library borrows data-graph binary bytes when they are bound to a Session. The pointer must meet the device’s dataGraphPipelineBinaryDataAlignment requirement. With the default DataGraphBinaryStorageMode::ImplementationManaged mode, keep the bytes valid until Session::configure() returns. If pipeline creation fails and is retried, keep them alive or bind a replacement binary before retrying. The optional identifier is copied when bound.

  • With DataGraphBinaryStorageMode::ApplicationBacked, Vulkan® may use the supplied bytes directly. Keep them accessible and unchanged until every pipeline created from them is destroyed, including compiled state shared by Sessions.

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