<commit_before>// Copyright 2019 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/hal/vulkan/direct_command_buffer.h"

#include <cstddef>
#include <cstdint>

#include "iree/base/api.h"
#include "iree/base/internal/inline_array.h"
#include "iree/base/internal/math.h"
#include "iree/base/logging.h"
#include "iree/base/status_cc.h"
#include "iree/base/tracing.h"
#include "iree/hal/utils/resource_set.h"
#include "iree/hal/vulkan/descriptor_set_arena.h"
#include "iree/hal/vulkan/dynamic_symbols.h"
#include "iree/hal/vulkan/native_descriptor_set.h"
#include "iree/hal/vulkan/native_event.h"
#include "iree/hal/vulkan/native_executable.h"
#include "iree/hal/vulkan/native_executable_layout.h"
#include "iree/hal/vulkan/status_util.h"
#include "iree/hal/vulkan/util/ref_ptr.h"
#include "iree/hal/vulkan/vma_buffer.h"

using namespace iree::hal::vulkan;

// Command buffer implementation that directly maps to VkCommandBuffer.
// This records the commands on the calling thread without additional threading
// indirection.
typedef struct iree_hal_vulkan_direct_command_buffer_t {
  iree_hal_command_buffer_t base;
  VkDeviceHandle* logical_device;
  iree_hal_vulkan_tracing_context_t* tracing_context;
  iree_arena_block_pool_t* block_pool;

  VkCommandPoolHandle* command_pool;
  VkCommandBuffer handle;

  DynamicSymbols* syms;

  // Maintains a reference to all resources used within the command buffer.
  // Reset on each begin.
  iree_hal_resource_set_t* resource_set;

  // TODO(benvanik): may grow large - should try to reclaim or reuse.
  DescriptorSetArena descriptor_set_arena;

  // The current descriptor set group in use by the command buffer, if any.
  // This must remain valid until all in-flight submissions of the command
  // buffer complete.
  DescriptorSetGroup descriptor_set_group;

  BuiltinExecutables* builtin_executables;

  // Shadow copy of push constants used during normal operation, for restoring
  // after builtin_executables uses vkCmdPushConstants. Size must be greater
  // than or equal to the push constant memory used by builtin_executables.
  // TODO(scotttodd): use [maxPushConstantsSize - 16, maxPushConstantsSize]
  //                  instead of [0, 16] to reduce frequency of updates
  uint8_t push_constants_storage[IREE_HAL_VULKAN_BUILTIN_PUSH_CONSTANT_COUNT];
} iree_hal_vulkan_direct_command_buffer_t;

namespace {
extern const iree_hal_command_buffer_vtable_t
    iree_hal_vulkan_direct_command_buffer_vtable;
}  // namespace

static iree_hal_vulkan_direct_command_buffer_t*
iree_hal_vulkan_direct_command_buffer_cast(
    iree_hal_command_buffer_t* base_value) {
  IREE_HAL_ASSERT_TYPE(base_value,
                       &iree_hal_vulkan_direct_command_buffer_vtable);
  return (iree_hal_vulkan_direct_command_buffer_t*)base_value;
}

iree_status_t iree_hal_vulkan_direct_command_buffer_allocate(
    iree_hal_device_t* device,
    iree::hal::vulkan::VkDeviceHandle* logical_device,
    iree::hal::vulkan::VkCommandPoolHandle* command_pool,
    iree_hal_command_buffer_mode_t mode,
    iree_hal_command_category_t command_categories,
    iree_hal_queue_affinity_t queue_affinity,
    iree_hal_vulkan_tracing_context_t* tracing_context,
    iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache,
    iree::hal::vulkan::BuiltinExecutables* builtin_executables,
    iree_arena_block_pool_t* block_pool,
    iree_hal_command_buffer_t** out_command_buffer) {
  IREE_ASSERT_ARGUMENT(logical_device);
  IREE_ASSERT_ARGUMENT(command_pool);
  IREE_ASSERT_ARGUMENT(descriptor_pool_cache);
  IREE_ASSERT_ARGUMENT(block_pool);
  IREE_ASSERT_ARGUMENT(out_command_buffer);
  IREE_TRACE_ZONE_BEGIN(z0);

  VkCommandBufferAllocateInfo allocate_info;
  allocate_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
  allocate_info.pNext = NULL;
  allocate_info.commandPool = *command_pool;
  allocate_info.commandBufferCount = 1;
  allocate_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;

  VkCommandBuffer handle = VK_NULL_HANDLE;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, command_pool->Allocate(&allocate_info, &handle));

  iree_hal_vulkan_direct_command_buffer_t* command_buffer = NULL;
  iree_status_t status =
      iree_allocator_malloc(logical_device->host_allocator(),
                            sizeof(*command_buffer), (void**)&command_buffer);
  if (iree_status_is_ok(status)) {
    iree_hal_command_buffer_initialize(
        device, mode, command_categories, queue_affinity,
        &iree_hal_vulkan_direct_command_buffer_vtable, &command_buffer->base);
    command_buffer->logical_device = logical_device;
    command_buffer->tracing_context = tracing_context;
    command_buffer->block_pool = block_pool;
    command_buffer->command_pool = command_pool;
    command_buffer->handle = handle;
    command_buffer->syms = logical_device->syms().get();

    new (&command_buffer->descriptor_set_arena)
        DescriptorSetArena(descriptor_pool_cache);
    new (&command_buffer->descriptor_set_group) DescriptorSetGroup();

    command_buffer->builtin_executables = builtin_executables;
    status = iree_hal_resource_set_allocate(block_pool,
                                            &command_buffer->resource_set);
  }

  if (iree_status_is_ok(status)) {
    *out_command_buffer = &command_buffer->base;
  } else {
    command_pool->Free(handle);
  }

  IREE_TRACE_ZONE_END(z0);
  return status;
}

static void iree_hal_vulkan_direct_command_buffer_reset(
    iree_hal_vulkan_direct_command_buffer_t* command_buffer) {
  // NOTE: we require that command buffers not be recorded while they are
  // in-flight so this is safe.
  IREE_IGNORE_ERROR(command_buffer->descriptor_set_group.Reset());
  iree_hal_resource_set_reset(command_buffer->resource_set);
}

bool iree_hal_vulkan_direct_command_buffer_isa(
    iree_hal_command_buffer_t* command_buffer) {
  return iree_hal_command_buffer_dyn_cast(
      command_buffer, &iree_hal_vulkan_direct_command_buffer_vtable);
}

static void* iree_hal_vulkan_direct_command_buffer_dyn_cast(
    iree_hal_command_buffer_t* command_buffer, const void* vtable) {
  if (vtable == &iree_hal_vulkan_direct_command_buffer_vtable) {
    IREE_HAL_ASSERT_TYPE(command_buffer, vtable);
    return command_buffer;
  }
  return NULL;
}

static void iree_hal_vulkan_direct_command_buffer_destroy(
    iree_hal_command_buffer_t* base_command_buffer) {
  iree_hal_vulkan_direct_command_buffer_t* command_buffer =
      iree_hal_vulkan_direct_command_buffer_cast(base_command_buffer);
  iree_allocator_t host_allocator =
      command_buffer->logical_device->host_allocator();
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_hal_vulkan_direct_command_buffer_reset(command_buffer);
  command_buffer->command_pool->Free(command_buffer->handle);

  command_buffer->descriptor_set_group.~DescriptorSetGroup();
  command_buffer->descriptor_set_arena.~DescriptorSetArena();

  iree_hal_resource_set_reset(command_buffer->resource_set);
  iree_allocator_free(host_allocator, command_buffer);

  IREE_TRACE_ZONE_END(z0);
}
<commit_msg>Fixing leak of Vulkan command buffer resource sets. (#8488)<commit_after>// Copyright 2019 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/hal/vulkan/direct_command_buffer.h"

#include <cstddef>
#include <cstdint>

#include "iree/base/api.h"
#include "iree/base/internal/inline_array.h"
#include "iree/base/internal/math.h"
#include "iree/base/logging.h"
#include "iree/base/status_cc.h"
#include "iree/base/tracing.h"
#include "iree/hal/utils/resource_set.h"
#include "iree/hal/vulkan/descriptor_set_arena.h"
#include "iree/hal/vulkan/dynamic_symbols.h"
#include "iree/hal/vulkan/native_descriptor_set.h"
#include "iree/hal/vulkan/native_event.h"
#include "iree/hal/vulkan/native_executable.h"
#include "iree/hal/vulkan/native_executable_layout.h"
#include "iree/hal/vulkan/status_util.h"
#include "iree/hal/vulkan/util/ref_ptr.h"
#include "iree/hal/vulkan/vma_buffer.h"

using namespace iree::hal::vulkan;

// Command buffer implementation that directly maps to VkCommandBuffer.
// This records the commands on the calling thread without additional threading
// indirection.
typedef struct iree_hal_vulkan_direct_command_buffer_t {
  iree_hal_command_buffer_t base;
  VkDeviceHandle* logical_device;
  iree_hal_vulkan_tracing_context_t* tracing_context;
  iree_arena_block_pool_t* block_pool;

  VkCommandPoolHandle* command_pool;
  VkCommandBuffer handle;

  DynamicSymbols* syms;

  // Maintains a reference to all resources used within the command buffer.
  // Reset on each begin.
  iree_hal_resource_set_t* resource_set;

  // TODO(benvanik): may grow large - should try to reclaim or reuse.
  DescriptorSetArena descriptor_set_arena;

  // The current descriptor set group in use by the command buffer, if any.
  // This must remain valid until all in-flight submissions of the command
  // buffer complete.
  DescriptorSetGroup descriptor_set_group;

  BuiltinExecutables* builtin_executables;

  // Shadow copy of push constants used during normal operation, for restoring
  // after builtin_executables uses vkCmdPushConstants. Size must be greater
  // than or equal to the push constant memory used by builtin_executables.
  // TODO(scotttodd): use [maxPushConstantsSize - 16, maxPushConstantsSize]
  //                  instead of [0, 16] to reduce frequency of updates
  uint8_t push_constants_storage[IREE_HAL_VULKAN_BUILTIN_PUSH_CONSTANT_COUNT];
} iree_hal_vulkan_direct_command_buffer_t;

namespace {
extern const iree_hal_command_buffer_vtable_t
    iree_hal_vulkan_direct_command_buffer_vtable;
}  // namespace

static iree_hal_vulkan_direct_command_buffer_t*
iree_hal_vulkan_direct_command_buffer_cast(
    iree_hal_command_buffer_t* base_value) {
  IREE_HAL_ASSERT_TYPE(base_value,
                       &iree_hal_vulkan_direct_command_buffer_vtable);
  return (iree_hal_vulkan_direct_command_buffer_t*)base_value;
}

iree_status_t iree_hal_vulkan_direct_command_buffer_allocate(
    iree_hal_device_t* device,
    iree::hal::vulkan::VkDeviceHandle* logical_device,
    iree::hal::vulkan::VkCommandPoolHandle* command_pool,
    iree_hal_command_buffer_mode_t mode,
    iree_hal_command_category_t command_categories,
    iree_hal_queue_affinity_t queue_affinity,
    iree_hal_vulkan_tracing_context_t* tracing_context,
    iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache,
    iree::hal::vulkan::BuiltinExecutables* builtin_executables,
    iree_arena_block_pool_t* block_pool,
    iree_hal_command_buffer_t** out_command_buffer) {
  IREE_ASSERT_ARGUMENT(logical_device);
  IREE_ASSERT_ARGUMENT(command_pool);
  IREE_ASSERT_ARGUMENT(descriptor_pool_cache);
  IREE_ASSERT_ARGUMENT(block_pool);
  IREE_ASSERT_ARGUMENT(out_command_buffer);
  IREE_TRACE_ZONE_BEGIN(z0);

  VkCommandBufferAllocateInfo allocate_info;
  allocate_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
  allocate_info.pNext = NULL;
  allocate_info.commandPool = *command_pool;
  allocate_info.commandBufferCount = 1;
  allocate_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;

  VkCommandBuffer handle = VK_NULL_HANDLE;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, command_pool->Allocate(&allocate_info, &handle));

  iree_hal_vulkan_direct_command_buffer_t* command_buffer = NULL;
  iree_status_t status =
      iree_allocator_malloc(logical_device->host_allocator(),
                            sizeof(*command_buffer), (void**)&command_buffer);
  if (iree_status_is_ok(status)) {
    iree_hal_command_buffer_initialize(
        device, mode, command_categories, queue_affinity,
        &iree_hal_vulkan_direct_command_buffer_vtable, &command_buffer->base);
    command_buffer->logical_device = logical_device;
    command_buffer->tracing_context = tracing_context;
    command_buffer->block_pool = block_pool;
    command_buffer->command_pool = command_pool;
    command_buffer->handle = handle;
    command_buffer->syms = logical_device->syms().get();

    new (&command_buffer->descriptor_set_arena)
        DescriptorSetArena(descriptor_pool_cache);
    new (&command_buffer->descriptor_set_group) DescriptorSetGroup();

    command_buffer->builtin_executables = builtin_executables;
    status = iree_hal_resource_set_allocate(block_pool,
                                            &command_buffer->resource_set);
  }

  if (iree_status_is_ok(status)) {
    *out_command_buffer = &command_buffer->base;
  } else {
    command_pool->Free(handle);
  }

  IREE_TRACE_ZONE_END(z0);
  return status;
}

static void iree_hal_vulkan_direct_command_buffer_reset(
    iree_hal_vulkan_direct_command_buffer_t* command_buffer) {
  // NOTE: we require that command buffers not be recorded while they are
  // in-flight so this is safe.
  IREE_IGNORE_ERROR(command_buffer->descriptor_set_group.Reset());
  iree_hal_resource_set_reset(command_buffer->resource_set);
}

bool iree_hal_vulkan_direct_command_buffer_isa(
    iree_hal_command_buffer_t* command_buffer) {
  return iree_hal_command_buffer_dyn_cast(
      command_buffer, &iree_hal_vulkan_direct_command_buffer_vtable);
}

static void* iree_hal_vulkan_direct_command_buffer_dyn_cast(
    iree_hal_command_buffer_t* command_buffer, const void* vtable) {
  if (vtable == &iree_hal_vulkan_direct_command_buffer_vtable) {
    IREE_HAL_ASSERT_TYPE(command_buffer, vtable);
    return command_buffer;
  }
  return NULL;
}

static void iree_hal_vulkan_direct_command_buffer_destroy(
    iree_hal_command_buffer_t* base_command_buffer) {
  iree_hal_vulkan_direct_command_buffer_t* command_buffer =
      iree_hal_vulkan_direct_command_buffer_cast(base_command_buffer);
  iree_allocator_t host_allocator =
      command_buffer->logical_device->host_allocator();
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_hal_vulkan_direct_command_buffer_reset(command_buffer);
  command_buffer->command_pool->Free(command_buffer->handle);

  command_buffer->descriptor_set_group.~DescriptorSetGroup();
  command_buffer->descriptor_set_arena.~DescriptorSetArena();

  iree_hal_resource_set_free(command_buffer->resource_set);
  iree_allocator_free(host_allocator, command_buffer);

  IREE_TRACE_ZONE_END(z0);
}

{'fraction_non_alphanumeric': 0.08139454128758941, 'fraction_numerical': 0.0030557677616501145, 'mean_word_length': 5.605504587155964, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 59, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6425945', 'n_tokens_mistral': 5043, 'n_tokens_neox': 4967, 'n_words': 788}