/*
 * Copyright (c) 2015-2026 The Khronos Group Inc.
 * Copyright (C) 2025 Arm Limited.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 */

#include "layer_validation_tests.h"
#include "data_graph_objects.h"

class PositiveDataGraph : public DataGraphTest {};

void DataGraphTest::InitBasicDataGraph(bool optical_flow) {
    SetTargetApiVersion(VK_API_VERSION_1_4);
    AddRequiredExtensions(VK_ARM_TENSORS_EXTENSION_NAME);
    AddRequiredExtensions(VK_ARM_DATA_GRAPH_EXTENSION_NAME);
    AddRequiredFeature(vkt::Feature::tensors);
    AddRequiredFeature(vkt::Feature::dataGraph);
    AddRequiredFeature(vkt::Feature::dataGraphShaderModule);
    AddRequiredFeature(vkt::Feature::shaderTensorAccess);
    AddRequiredFeature(vkt::Feature::vulkanMemoryModel);
    AddRequiredFeature(vkt::Feature::shaderInt8);
    if (optical_flow) {
        AddRequiredExtensions(VK_ARM_DATA_GRAPH_OPTICAL_FLOW_EXTENSION_NAME);
        AddRequiredFeature(vkt::Feature::dataGraphOpticalFlow);
    }
    RETURN_IF_SKIP(Init());
}

const VkTensorDescriptionARM DataGraphTest::defaultConstantTensorDesc{DefaultConstantTensorDesc()};

std::vector<VkBindDataGraphPipelineSessionMemoryInfoARM> DataGraphTest::InitSessionBindInfo(
    const vkt::DataGraphPipelineSession& session, const std::vector<vkt::DeviceMemory>& device_mem) {
    const auto& bind_point_reqs = session.BindPointReqs();
    std::vector<VkBindDataGraphPipelineSessionMemoryInfoARM> session_bind_infos(session.MemReqs().size());
    uint32_t req_i = 0;
    for (uint32_t i = 0; i < bind_point_reqs.size(); i++) {
        if (bind_point_reqs[i].bindPointType != VK_DATA_GRAPH_PIPELINE_SESSION_BIND_POINT_TYPE_MEMORY_ARM) {
            continue;
        }

        for (uint32_t j = 0; j < bind_point_reqs[i].numObjects; j++) {
            session_bind_infos[req_i] = vku::InitStructHelper();
            session_bind_infos[req_i].session = session;
            session_bind_infos[req_i].memory = device_mem[req_i];
            session_bind_infos[req_i].bindPoint = bind_point_reqs[req_i].bindPoint;
            session_bind_infos[req_i].objectIndex = j;
            req_i++;
        }
    }
    return session_bind_infos;
}

// Trivial rank 1 tensor
VkTensorDescriptionARM DataGraphTest::DefaultDesc() {
    VkTensorDescriptionARM desc = vku::InitStructHelper();
    static const std::vector<int64_t> dimensions{2};
    desc.tiling = VK_TENSOR_TILING_LINEAR_ARM;
    desc.format = VK_FORMAT_R8_SINT;
    desc.dimensionCount = dimensions.size();
    desc.pDimensions = dimensions.data();
    desc.pStrides = nullptr;
    desc.usage = VK_TENSOR_USAGE_DATA_GRAPH_BIT_ARM;
    return desc;
}

// Tensor description for constant in GetSpirvModifyableDataGraph and GetSpirvMultiEntryTwoDataGraph
VkTensorDescriptionARM DataGraphTest::DefaultConstantTensorDesc() {
    VkTensorDescriptionARM desc = vku::InitStructHelper();
    static std::vector<int64_t> dimensions{1, 2, 4, 4};
    desc.tiling = VK_TENSOR_TILING_LINEAR_ARM;
    desc.format = VK_FORMAT_R32_SINT;
    desc.dimensionCount = dimensions.size();
    desc.pDimensions = dimensions.data();
    desc.usage = VK_TENSOR_USAGE_DATA_GRAPH_BIT_ARM;
    return desc;
}

// Assumes description defines a packed tensor with linear tiling.
inline VkDeviceSize GetLinearTensorBufferSizeBytes(const VkTensorDescriptionARM &description) {
    assert(description.dimensionCount > 0);
    assert(description.pDimensions);
    // NOTE: no assert for linear tiling: used for different tilings in negative tests

    VkDeviceSize size = vkuFormatTexelBlockSize(description.format);
    for (uint32_t i = 0; i < description.dimensionCount; ++i) {
        assert(description.pDimensions[i] > 0);
        size *= static_cast<VkDeviceSize>(description.pDimensions[i]);
    }
    return size;
}

// Constant for GetSpirvModifyableDataGraph and GetSpirvMultiEntryTwoDataGraph
VkDataGraphPipelineConstantARM DataGraphTest::GetConstant(const VkTensorDescriptionARM& desc) {
    VkDataGraphPipelineConstantARM constant = vku::InitStructHelper();
    constant.id = 0;

    VkDeviceSize buffer_size = GetLinearTensorBufferSizeBytes(desc);
    static std::vector<uint8_t> constant_data(buffer_size);
    constant.pConstantData = constant_data.data();
    constant.pNext = &desc;
    return constant;
}

TEST_F(PositiveDataGraph, ExecuteDataGraph) {
    TEST_DESCRIPTION("Create and execute a datagraph");
    RETURN_IF_SKIP(InitBasicDataGraph());

    vkt::dg::DataGraphPipelineHelper pipeline(*this);
    pipeline.CreateDataGraphPipeline();

    VkDataGraphPipelineSessionCreateInfoARM session_ci = vku::InitStructHelper();
    session_ci.dataGraphPipeline = pipeline;
    vkt::DataGraphPipelineSession session(*m_device, session_ci);

    auto& bind_point_reqs = session.BindPointReqs();
    std::vector<vkt::DeviceMemory> device_mem(bind_point_reqs.size());
    session.AllocSessionMem(device_mem);
    auto session_bind_infos = InitSessionBindInfo(session, device_mem);
    vk::BindDataGraphPipelineSessionMemoryARM(*m_device, session_bind_infos.size(), session_bind_infos.data());

    pipeline.descriptor_set_->WriteDescriptorTensorInfo(0, &pipeline.tensor_views_[0]->handle(), 0);
    pipeline.descriptor_set_->WriteDescriptorTensorInfo(1, &pipeline.tensor_views_[1]->handle(), 0);
    pipeline.descriptor_set_->UpdateDescriptorSets();

    m_command_buffer.Begin();
    vk::CmdBindPipeline(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, pipeline);
    vk::CmdBindDescriptorSets(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, pipeline.pipeline_layout_, 0, 1,
                              &pipeline.descriptor_set_.get()->set_, 0, nullptr);
    vk::CmdDispatchDataGraphARM(m_command_buffer, session, nullptr);
    m_command_buffer.End();

    m_default_queue->SubmitAndWait(m_command_buffer);
}

TEST_F(PositiveDataGraph, DISABLED_ProtectedMemoryDataGraph) {
    TEST_DESCRIPTION("Execute a datagraph with protected memory");
    SetTargetApiVersion(VK_API_VERSION_1_4);
    AddRequiredExtensions(VK_ARM_TENSORS_EXTENSION_NAME);
    AddRequiredExtensions(VK_ARM_DATA_GRAPH_EXTENSION_NAME);
    AddRequiredFeature(vkt::Feature::tensors);
    AddRequiredFeature(vkt::Feature::dataGraph);
    AddRequiredFeature(vkt::Feature::dataGraphShaderModule);
    AddRequiredFeature(vkt::Feature::shaderTensorAccess);
    AddRequiredFeature(vkt::Feature::vulkanMemoryModel);
    AddRequiredFeature(vkt::Feature::shaderInt8);
    AddRequiredFeature(vkt::Feature::protectedMemory);
    AddRequiredFeature(vkt::Feature::pipelineProtectedAccess);
    RETURN_IF_SKIP(InitFramework());
    RETURN_IF_SKIP(InitState(nullptr, nullptr, VK_COMMAND_POOL_CREATE_PROTECTED_BIT));

    vkt::dg::HelperParameters params;
    params.protected_tensors = true;
    vkt::dg::DataGraphPipelineHelper pipeline(*this, params);

    pipeline.pipeline_ci_.flags = VK_PIPELINE_CREATE_2_PROTECTED_ACCESS_ONLY_BIT_EXT;
    pipeline.CreateDataGraphPipeline();

    VkDataGraphPipelineSessionCreateInfoARM session_ci = vku::InitStructHelper();
    session_ci.dataGraphPipeline = pipeline;
    session_ci.flags = VK_DATA_GRAPH_PIPELINE_SESSION_CREATE_PROTECTED_BIT_ARM;
    vkt::DataGraphPipelineSession session(*m_device, session_ci);

    std::vector<vkt::DeviceMemory> device_mem(session.BindPointsCount());
    session.AllocSessionMem(device_mem, true);
    auto session_bind_infos = InitSessionBindInfo(session, device_mem);
    vk::BindDataGraphPipelineSessionMemoryARM(*m_device, session_bind_infos.size(), session_bind_infos.data());

    pipeline.descriptor_set_->WriteDescriptorTensorInfo(0, &pipeline.tensor_views_[0]->handle(), 0);
    pipeline.descriptor_set_->WriteDescriptorTensorInfo(1, &pipeline.tensor_views_[1]->handle(), 0);
    pipeline.descriptor_set_->UpdateDescriptorSets();

    m_command_buffer.Begin();
    vk::CmdBindPipeline(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, pipeline);
    vk::CmdBindDescriptorSets(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, pipeline.pipeline_layout_, 0, 1,
                              &pipeline.descriptor_set_.get()->set_, 0, nullptr);
    vk::CmdDispatchDataGraphARM(m_command_buffer, session, nullptr);
    m_command_buffer.End();

    m_default_queue->SubmitAndWait(m_command_buffer);
}

TEST_F(PositiveDataGraph, ShaderModuleInPNext) {
    TEST_DESCRIPTION(
        "Pass a VkShaderModuleCreateInfo in the pNext chain of pipeline info, not as "
        "VkDataGraphPipelineShaderModuleCreateInfoARM::module.");
    RETURN_IF_SKIP(InitBasicDataGraph());

    // create a ShaderModule to add in the pNext chain
    spvtools::SpirvTools tools{SPV_ENV_UNIVERSAL_1_6};
    const std::string& spirv_source = vkt::dg::DataGraphPipelineHelper::GetSpirvBasicDataGraph();
    std::vector<uint32_t> spirv_binary;
    if (!tools.Assemble(spirv_source, &spirv_binary)) {
        Monitor().SetError("Failed to compile SPIRV shader module");
        return;
    }
    VkShaderModuleCreateInfo shader_module_create_info = vku::InitStructHelper();
    shader_module_create_info.codeSize = spirv_binary.size() * sizeof(uint32_t);
    shader_module_create_info.pCode = spirv_binary.data();

    // 2 variants, adding the shader in different places in the pipeline's pNext chain, both must work

    {
        vkt::dg::DataGraphPipelineHelper pipeline(*this);
        // the helper constructor adds the shader module as VkDataGraphPipelineShaderModuleCreateInfoARM::module, get rid of it
        pipeline.shader_module_ci_.module = VK_NULL_HANDLE;

        // add the shader info in VkDataGraphPipelineShaderModuleCreateInfoARM::pNext
        pipeline.shader_module_ci_.pNext = &shader_module_create_info;
        pipeline.CreateDataGraphPipeline();
    }

    {
        vkt::dg::DataGraphPipelineHelper pipeline(*this);
        // the helper constructor adds the shader module as VkDataGraphPipelineShaderModuleCreateInfoARM::module, get rid of it
        pipeline.shader_module_ci_.module = VK_NULL_HANDLE;

        // add the shader info in VkDataGraphPipelineCreateInfoARM::pNext
        pipeline.pipeline_ci_.pNext = &shader_module_create_info;
        shader_module_create_info.pNext = &pipeline.shader_module_ci_;
        pipeline.shader_module_ci_.pNext = nullptr;
        pipeline.CreateDataGraphPipeline();
    }
}

TEST_F(PositiveDataGraph, DataGraphMultipleEntrypoints) {
    TEST_DESCRIPTION("Execute 2 different entrypoints in the datagraph's spirv.");
    RETURN_IF_SKIP(InitBasicDataGraph());

    // get spirv with 2 entrypoints
    const std::string two_entrypoint_spirv = vkt::dg::DataGraphPipelineHelper::GetSpirvMultiEntryTwoDataGraph();

    // create graph at entrypoint 1
    {
        // NOTE: even though there is an OpGraphConstantARM in the spirv, we don't need to initialize any
        // resource, as it is NOT used in this entrypoint
        vkt::dg::HelperParameters params;
        params.spirv_source = two_entrypoint_spirv.c_str();
        params.entrypoint = "entrypoint_1";
        vkt::dg::DataGraphPipelineHelper pipeline(*this, params);
        pipeline.CreateDataGraphPipeline();
    }

    // create graph at entrypoint 2
    {
        // NOTE: this entrypoint uses the OpGraphConstantARM in the spirv, so we have to provide a matching object
        vkt::dg::HelperParameters params;
        params.spirv_source = two_entrypoint_spirv.c_str();
        params.entrypoint = "entrypoint_2";
        vkt::dg::DataGraphPipelineHelper pipeline(*this, params);
        pipeline.CreateDataGraphPipeline();
    }
}

TEST_F(PositiveDataGraph, CmdDispatchDescriptorBuffer) {
    TEST_DESCRIPTION("Dispatch a datagraph using descriptor buffers.");
    SetTargetApiVersion(VK_API_VERSION_1_4);
    AddRequiredExtensions(VK_EXT_DESCRIPTOR_BUFFER_EXTENSION_NAME);
    AddRequiredExtensions(VK_KHR_BUFFER_DEVICE_ADDRESS_EXTENSION_NAME);
    AddRequiredFeature(vkt::Feature::descriptorBuffer);
    AddRequiredFeature(vkt::Feature::bufferDeviceAddress);
    AddRequiredFeature(vkt::Feature::dataGraphDescriptorBuffer);
    RETURN_IF_SKIP(InitBasicDataGraph());

    vkt::dg::DataGraphPipelineHelper pipeline(*this);

    // create a pipeline layout with the required flags
    VkDescriptorSetLayoutCreateInfo dsl_ci = vku::InitStructHelper();
    dsl_ci.flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_DESCRIPTOR_BUFFER_BIT_EXT;
    dsl_ci.bindingCount = pipeline.descriptor_set_layout_bindings_.size();
    dsl_ci.pBindings = pipeline.descriptor_set_layout_bindings_.data();

    vkt::DescriptorSetLayout dsl(*m_device, dsl_ci);
    vkt::PipelineLayout pipeline_layout(*m_device, {&dsl});
    ASSERT_TRUE(pipeline_layout.initialized());

    // set layout and flags for descriptor buffer
    pipeline.pipeline_ci_.flags |= VK_PIPELINE_CREATE_DESCRIPTOR_BUFFER_BIT_EXT;
    pipeline.pipeline_ci_.layout = pipeline_layout;
    pipeline.CreateDataGraphPipeline();

    VkDataGraphPipelineSessionCreateInfoARM session_ci = vku::InitStructHelper();
    session_ci.dataGraphPipeline = pipeline.Handle();
    vkt::DataGraphPipelineSession session(*m_device, session_ci);

    auto& bind_point_reqs = session.BindPointReqs();
    std::vector<vkt::DeviceMemory> device_mem(bind_point_reqs.size());
    session.AllocSessionMem(device_mem);
    auto session_bind_infos = InitSessionBindInfo(session, device_mem);
    vk::BindDataGraphPipelineSessionMemoryARM(*m_device, session_bind_infos.size(), session_bind_infos.data());

    pipeline.descriptor_set_->WriteDescriptorTensorInfo(0, &pipeline.tensor_views_[0]->handle(), 0);
    pipeline.descriptor_set_->WriteDescriptorTensorInfo(1, &pipeline.tensor_views_[1]->handle(), 0);
    pipeline.descriptor_set_->UpdateDescriptorSets();

    m_command_buffer.Begin();
    vk::CmdBindPipeline(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, pipeline.Handle());

    vkt::Buffer buffer(*m_device, 4096, VK_BUFFER_USAGE_RESOURCE_DESCRIPTOR_BUFFER_BIT_EXT, vkt::device_address);
    VkDescriptorBufferBindingInfoEXT dbbi = vku::InitStructHelper();
    dbbi.address = buffer.Address();
    dbbi.usage = VK_BUFFER_USAGE_RESOURCE_DESCRIPTOR_BUFFER_BIT_EXT;

    vk::CmdBindDescriptorBuffersEXT(m_command_buffer, 1, &dbbi);
    uint32_t index = 0;
    VkDeviceSize offset = 0;
    vk::CmdSetDescriptorBufferOffsetsEXT(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, pipeline_layout, 0, 1, &index,
                                         &offset);
    vk::CmdDispatchDataGraphARM(m_command_buffer, session, nullptr);
    m_errorMonitor->VerifyFound();
    m_command_buffer.End();
}

TEST_F(PositiveDataGraph, ResourceInfoImageLayoutsUnifiedImageLayouts) {
    AddRequiredFeature(vkt::Feature::unifiedImageLayouts);
    RETURN_IF_SKIP(InitBasicDataGraph(true));

    // Need to create a data graph with image resources instead of the usual tensors, so use optical flow.
    vkt::dg::OpticalFlowHelper optical_flow(*this);
    /* Fail to chain a VkDataGraphPipelineResourceInfoImageLayoutARM structure. Since the unifiedImageLayouts feature is enabled
     * this is valid usage. */
    optical_flow.dg_pipeline_.resources_[0].pNext = nullptr;
    optical_flow.CreateDataGraphPipeline();
    optical_flow.SetupImageDescriptors();

    VkDataGraphPipelineSessionCreateInfoARM session_ci = vku::InitStructHelper();
    session_ci.dataGraphPipeline = optical_flow.dg_pipeline_;

    vkt::DataGraphPipelineSession session(*m_device, session_ci);
    std::vector<vkt::DeviceMemory> device_mem(session.BindPointReqs().size());
    session.AllocSessionMem(device_mem);
    std::vector<VkBindDataGraphPipelineSessionMemoryInfoARM> session_bind_infos =
        DataGraphTest::InitSessionBindInfo(session, device_mem);
    vk::BindDataGraphPipelineSessionMemoryARM(*m_device, session_bind_infos.size(), session_bind_infos.data());

    VkDataGraphPipelineOpticalFlowDispatchInfoARM optical_flow_di = vku::InitStructHelper();
    VkDataGraphPipelineDispatchInfoARM pipeline_di = vku::InitStructHelper(&optical_flow_di);

    m_command_buffer.Begin();
    vk::CmdBindPipeline(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, optical_flow.dg_pipeline_);
    vk::CmdBindDescriptorSets(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, optical_flow.dg_pipeline_.pipeline_layout_,
                              0, 1, &optical_flow.dg_pipeline_.descriptor_set_.get()->set_, 0, nullptr);
    vk::CmdDispatchDataGraphARM(m_command_buffer, session, &pipeline_di);
    m_command_buffer.End();
}

TEST_F(PositiveDataGraph, OpticalFlow) {
    TEST_DESCRIPTION("Execute a datagraph with an optical flow node");
    RETURN_IF_SKIP(InitBasicDataGraph(true));

    vkt::dg::OpticalFlowHelper optical_flow(*this);
    optical_flow.CreateDataGraphPipeline();
    optical_flow.SetupImageDescriptors();

    VkDataGraphPipelineSessionCreateInfoARM session_ci = vku::InitStructHelper();
    session_ci.dataGraphPipeline = optical_flow.dg_pipeline_;
    vkt::DataGraphPipelineSession session(*m_device, session_ci);

    auto& bind_point_reqs = session.BindPointReqs();
    std::vector<vkt::DeviceMemory> device_mem(bind_point_reqs.size());
    session.AllocSessionMem(device_mem);
    auto session_bind_infos = InitSessionBindInfo(session, device_mem);
    vk::BindDataGraphPipelineSessionMemoryARM(*m_device, session_bind_infos.size(), session_bind_infos.data());

    VkDataGraphPipelineOpticalFlowDispatchInfoARM optical_flow_di = vku::InitStructHelper();
    optical_flow_di.meanFlowL1NormHint =
        optical_flow.optical_flow_ci_.height;  // must be less than or equal to optical_flow_ci.height/optical_flow_ci.width
    VkDataGraphPipelineDispatchInfoARM pipeline_di = vku::InitStructHelper(&optical_flow_di);

    m_command_buffer.Begin();
    vk::CmdBindPipeline(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, optical_flow.dg_pipeline_);
    vk::CmdBindDescriptorSets(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, optical_flow.PipelineLayout(), 0, 1,
                              optical_flow.DescriptorSet(), 0, nullptr);
    vk::CmdDispatchDataGraphARM(m_command_buffer, session, &pipeline_di);
    m_command_buffer.End();
}

TEST_F(PositiveDataGraph, OpticalFlowConnectionsImageLayoutsUnifiedImageLayouts) {
    TEST_DESCRIPTION(
        "Try to dispatch an optical flow pipeline where the connections specify layout VK_IMAGE_LAYOUT_GENERAL which is different "
        "to that of the corresponding pipeline resource info image layout. In this version of the test the unifiedImageLayouts "
        "feature is enabled.");
    AddRequiredFeature(vkt::Feature::unifiedImageLayouts);
    RETURN_IF_SKIP(InitBasicDataGraph(true));

    vkt::dg::OpticalFlowHelper optical_flow(*this);
    optical_flow.image_layouts_[0].layout = VK_IMAGE_LAYOUT_GENERAL;
    optical_flow.CreateDataGraphPipeline();
    optical_flow.SetupImageDescriptors();

    VkDataGraphPipelineSessionCreateInfoARM session_ci = vku::InitStructHelper();
    session_ci.dataGraphPipeline = optical_flow.dg_pipeline_;

    vkt::DataGraphPipelineSession session(*m_device, session_ci);
    std::vector<vkt::DeviceMemory> device_mem(session.BindPointReqs().size());
    session.AllocSessionMem(device_mem);
    std::vector<VkBindDataGraphPipelineSessionMemoryInfoARM> session_bind_infos =
        DataGraphTest::InitSessionBindInfo(session, device_mem);
    vk::BindDataGraphPipelineSessionMemoryARM(*m_device, session_bind_infos.size(), session_bind_infos.data());

    VkDataGraphPipelineOpticalFlowDispatchInfoARM optical_flow_di = vku::InitStructHelper();
    VkDataGraphPipelineDispatchInfoARM pipeline_di = vku::InitStructHelper(&optical_flow_di);

    m_command_buffer.Begin();
    vk::CmdBindPipeline(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, optical_flow.dg_pipeline_);
    vk::CmdBindDescriptorSets(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, optical_flow.dg_pipeline_.pipeline_layout_,
                              0, 1, &optical_flow.dg_pipeline_.descriptor_set_.get()->set_, 0, nullptr);
    vk::CmdDispatchDataGraphARM(m_command_buffer, session, &pipeline_di);
    m_command_buffer.End();
}

// unit test to verify the fix for bug in descriptor copying
TEST_F(PositiveDataGraph, CopyUpdateDescriptor) {
    TEST_DESCRIPTION("Dispatch a datagraph where a descriptor has been copy-updated after creation");
    RETURN_IF_SKIP(InitBasicDataGraph());

    vkt::dg::DataGraphPipelineHelper pipeline(*this);
    pipeline.CreateDataGraphPipeline();

    VkDataGraphPipelineSessionCreateInfoARM session_ci = vku::InitStructHelper();
    session_ci.dataGraphPipeline = pipeline;
    vkt::DataGraphPipelineSession session(*m_device, session_ci);

    auto& bind_point_reqs = session.BindPointReqs();
    std::vector<vkt::DeviceMemory> device_mem(bind_point_reqs.size());
    session.AllocSessionMem(device_mem);
    auto session_bind_infos = InitSessionBindInfo(session, device_mem);
    vk::BindDataGraphPipelineSessionMemoryARM(*m_device, session_bind_infos.size(), session_bind_infos.data());

    pipeline.descriptor_set_->WriteDescriptorTensorInfo(0, &pipeline.tensor_views_[0]->handle(), 0);
    pipeline.descriptor_set_->WriteDescriptorTensorInfo(1, &pipeline.tensor_views_[1]->handle(), 0);
    pipeline.descriptor_set_->UpdateDescriptorSets();

    // create a new tensor and copy it into the input tensor

    VkTensorDescriptionARM desc = pipeline.tensors_[0]->Description();
    desc.usage |= VK_TENSOR_USAGE_TRANSFER_SRC_BIT_ARM | VK_TENSOR_USAGE_TRANSFER_DST_BIT_ARM;
    VkTensorCreateInfoARM info = vku::InitStructHelper();
    info.pDescription = &desc;
    info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    vkt::Tensor tensor_src(*m_device, info);
    tensor_src.BindToMem();

    VkTensorViewCreateInfoARM tensor_view_create_info_src = vku::InitStructHelper();
    tensor_view_create_info_src.tensor = tensor_src;
    tensor_view_create_info_src.format = tensor_src.Format();
    vkt::TensorView view_src(*m_device, tensor_view_create_info_src);

    OneOffDescriptorSet descriptor_set_src(m_device, {{0, VK_DESCRIPTOR_TYPE_TENSOR_ARM, 1, VK_SHADER_STAGE_ALL, nullptr}});
    descriptor_set_src.WriteDescriptorTensorInfo(0, &view_src.handle());
    descriptor_set_src.UpdateDescriptorSets();

    VkCopyDescriptorSet copy_set = vku::InitStructHelper();
    copy_set.srcSet = descriptor_set_src.set_;
    copy_set.srcBinding = 0;
    copy_set.dstSet = pipeline.descriptor_set_->set_;
    copy_set.dstBinding = 0;
    copy_set.descriptorCount = 1;
    vk::UpdateDescriptorSets(*m_device, 0, nullptr, 1, &copy_set);

    // continue as normal to dispatch

    m_command_buffer.Begin();
    vk::CmdBindPipeline(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, pipeline);
    vk::CmdBindDescriptorSets(m_command_buffer, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, pipeline.pipeline_layout_, 0, 1,
                              &pipeline.descriptor_set_.get()->set_, 0, nullptr);
    vk::CmdDispatchDataGraphARM(m_command_buffer, session, nullptr);
    m_command_buffer.End();

    m_default_queue->SubmitAndWait(m_command_buffer);
}
