Skip to content

[vulkan] Add a buffer index offset during codegen to allow arbitrary crop offsets #8954

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
derek-gerstmann merged 3 commits into from
Feb 20, 2026
Merged

[vulkan] Add a buffer index offset during codegen to allow arbitrary crop offsets #8954

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f9c904a
Add an index offset during codegen to allow arbitrary offsets when ac…
derek-gerstmann Feb 18, 2026
70830d4
Refactor cleanup to remove unused MemoryRegion.range.
derek-gerstmann Feb 19, 2026
bc65d21
Clang formatting pass
derek-gerstmann Feb 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 85 additions & 5 deletions src/CodeGen_Vulkan_Dev.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,8 @@ class CodeGen_Vulkan_Dev : public CodeGen_GPU_Dev {
void store_at_scalar_index(const Store *op, SpvId index_id, SpvId variable_id, Type value_type, Type storage_type, SpvStorageClass storage_class, SpvId value_id);
void store_at_vector_index(const Store *op, SpvId variable_id, Type value_type, Type storage_type, SpvStorageClass storage_class, SpvId value_id);

SpvId apply_storage_buffer_offset(SpvId variable_id, SpvId index_id);

SpvFactory::Components split_vector(Type type, SpvId value_id);
SpvId join_vector(Type type, const SpvFactory::Components &value_components);
SpvId fill_vector(Type type, SpvId value_id);
Expand Down Expand Up @@ -237,6 +239,9 @@ class CodeGen_Vulkan_Dev : public CodeGen_GPU_Dev {
using StorageAccessMap = std::unordered_map<SpvId, StorageAccess>;
StorageAccessMap storage_access_map;

using StorageBufferOffsetMap = std::unordered_map<SpvId, SpvId>;
StorageBufferOffsetMap storage_buffer_offset_map;

// Defines the binding information for a specialization constant
// that is exported by the module and can be overriden at runtime
struct SpecializationBinding {
Expand Down Expand Up @@ -1373,6 +1378,27 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::visit(const Select *op) {
builder.update_id(result_id);
}

SpvId CodeGen_Vulkan_Dev::SPIRV_Emitter::apply_storage_buffer_offset(SpvId variable_id, SpvId index_id) {
auto offset_map_it = storage_buffer_offset_map.find(variable_id);
if (offset_map_it == storage_buffer_offset_map.end()) {
return index_id;
}

SpvId offset_id = offset_map_it->second;
SpvId index_type_id = builder.declare_type(Int(32));
SpvId adjusted_index_id = builder.reserve_id(SpvResultId);

debug(2) << "CodeGen_Vulkan_Dev::SPIRV_Emitter::apply_storage_buffer_offset(): "
<< "variable_id=" << variable_id << " "
<< "index_type_id=" << index_type_id << " "
<< "index_id=" << index_id << " "
<< "offset_id=" << offset_id << " "
<< "adjusted_index_id=" << adjusted_index_id << "\n";

builder.append(SpvFactory::integer_add(index_type_id, adjusted_index_id, index_id, offset_id));
return adjusted_index_id;
}

void CodeGen_Vulkan_Dev::SPIRV_Emitter::load_from_scalar_index(const Load *op, SpvId index_id, SpvId variable_id, Type value_type, Type storage_type, SpvStorageClass storage_class) {
debug(2) << "CodeGen_Vulkan_Dev::SPIRV_Emitter::load_from_scalar_index(): "
<< "index_id=" << index_id << " "
Expand All @@ -1392,7 +1418,7 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::load_from_scalar_index(const Load *op, S

uint32_t zero = 0;
SpvId src_id = SpvInvalidId;
SpvId src_index_id = index_id;
SpvId src_index_id = apply_storage_buffer_offset(variable_id, index_id);
if (storage_class == SpvStorageClassUniform) {
if (builder.is_struct_type(base_type_id)) {
SpvId zero_id = builder.declare_constant(UInt(32), &zero);
Expand Down Expand Up @@ -1490,7 +1516,7 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::store_at_scalar_index(const Store *op, S

uint32_t zero = 0;
SpvId dst_id = SpvInvalidId;
SpvId dst_index_id = index_id;
SpvId dst_index_id = apply_storage_buffer_offset(variable_id, index_id);

SpvId ptr_type_id = builder.declare_pointer_type(storage_type, storage_class);
if (storage_class == SpvStorageClassUniform) {
Expand Down Expand Up @@ -2259,6 +2285,7 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::reset() {
SymbolScope empty;
symbol_table.swap(empty);
storage_access_map.clear();
storage_buffer_offset_map.clear();
descriptor_set_table.clear();
reset_workgroup_size();
}
Expand Down Expand Up @@ -2598,6 +2625,14 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::declare_entry_point(const Stmt &s, SpvId
builder.add_entry_point(kernel_func_id, SpvExecutionModelGLCompute, entry_point_variables);
}

namespace {

uint32_t align_offset(uint32_t offset, uint32_t alignment) {
return (offset + (alignment - 1)) & ~(alignment - 1);
}

} // namespace

void CodeGen_Vulkan_Dev::SPIRV_Emitter::declare_device_args(const Stmt &s, uint32_t entry_point_index,
const std::string &entry_point_name,
const std::vector<DeviceArgument> &args) {
Expand All @@ -2621,7 +2656,9 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::declare_device_args(const Stmt &s, uint3

// GLSL-style: each input buffer is a runtime array in a buffer struct
// All other params get passed in as a single uniform block
// First, need to count scalar parameters to construct the uniform struct
// First, need to count scalar parameters and buffer parameters to construct the uniform struct
uint32_t scalar_arg_count = 0;
uint32_t buffer_arg_count = 0;
SpvBuilder::StructMemberTypes param_struct_members;
for (const auto &arg : args) {
if (!arg.is_buffer) {
Expand All @@ -2634,11 +2671,21 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::declare_device_args(const Stmt &s, uint3

SpvId arg_type_id = builder.declare_type(arg.type);
param_struct_members.push_back(arg_type_id);
scalar_arg_count++;
} else {
buffer_arg_count++;
}
}

// Add a buffer offset parameter for each buffer (one Int32 per buffer)
// to support crops at arbitrary index offsets.
Type offset_type = Int(32);
SpvId offset_type_id = builder.declare_type(offset_type);
param_struct_members.insert(param_struct_members.end(), size_t(buffer_arg_count), offset_type_id);

// Add a binding for a uniform buffer packed with all scalar args
uint32_t binding_counter = 0;
SpvId param_pack_var_id = SpvInvalidId;
if (!param_struct_members.empty()) {

const std::string struct_name = std::string("k") + std::to_string(kernel_index) + std::string("_args_struct");
Expand All @@ -2647,6 +2694,8 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::declare_device_args(const Stmt &s, uint3
// Add a decoration describing the offset for each parameter struct member
uint32_t param_member_index = 0;
uint32_t param_member_offset = 0;

// First, add decorations for each scalar arg
for (const auto &arg : args) {
if (!arg.is_buffer) {
SpvBuilder::Literals param_offset_literals = {param_member_offset};
Expand All @@ -2656,13 +2705,24 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::declare_device_args(const Stmt &s, uint3
}
}

// Force alignment for the parameter offset (e.g. all Int32 members in Uniform blocks must be 4-byte aligned)
param_member_offset = align_offset(param_member_offset, offset_type.bytes());

// Next, add a decoration for the storage buffer offsets
for (uint32_t b = 0; b < buffer_arg_count; b++) {
SpvBuilder::Literals param_offset_literals = {param_member_offset};
builder.add_struct_annotation(param_struct_type_id, param_member_index, SpvDecorationOffset, param_offset_literals);
param_member_offset += offset_type.bytes();
param_member_index++;
}

// Add a Block decoration for the parameter pack itself
builder.add_annotation(param_struct_type_id, SpvDecorationBlock);

// Add a variable for the parameter pack
const std::string param_pack_var_name = std::string("k") + std::to_string(kernel_index) + std::string("_args_var");
SpvId param_pack_ptr_type_id = builder.declare_pointer_type(param_struct_type_id, SpvStorageClassUniform);
SpvId param_pack_var_id = builder.declare_global_variable(param_pack_var_name, param_pack_ptr_type_id, SpvStorageClassUniform);
param_pack_var_id = builder.declare_global_variable(param_pack_var_name, param_pack_ptr_type_id, SpvStorageClassUniform);

// We always pass in the parameter pack as the first binding
SpvBuilder::Literals binding_index = {0};
Expand All @@ -2672,7 +2732,7 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::declare_device_args(const Stmt &s, uint3
descriptor_set.uniform_buffer_count++;
binding_counter++;

// Declare all the args with appropriate offsets into the parameter struct
// Declare all the scalar args with appropriate offsets into the parameter struct
uint32_t scalar_index = 0;
for (const auto &arg : args) {
if (!arg.is_buffer) {
Expand All @@ -2692,6 +2752,8 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::declare_device_args(const Stmt &s, uint3
}

// Add bindings for all device buffers declared as GLSL-style buffer blocks in uniform storage
// and adjust the indices with the appropriate storage buffer offsets (to support arbitrary crops)
uint32_t buffer_index = 0;
for (const auto &arg : args) {
if (arg.is_buffer) {

Expand Down Expand Up @@ -2741,6 +2803,24 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::declare_device_args(const Stmt &s, uint3
access.storage_class = storage_class;
storage_access_map[buffer_block_var_id] = access;
descriptor_set.storage_buffer_count++;

// Load the storage buffer offset for this buffer from the uniform struct
// These offsets are stored *after* all scalar args in the uniform struct
if (param_pack_var_id != SpvInvalidId) {
uint32_t buffer_offset_index_param = scalar_arg_count + buffer_index;
SpvId buffer_offset_index_param_id = builder.declare_constant(UInt(32), &buffer_offset_index_param);
SpvId index_ptr_type_id = builder.declare_pointer_type(offset_type_id, SpvStorageClassUniform);
SpvFactory::Indices buffer_offset_index_access_indices = {buffer_offset_index_param_id};
SpvId buffer_offset_index_access_chain = builder.declare_access_chain(index_ptr_type_id, param_pack_var_id, buffer_offset_index_access_indices);

SpvId buffer_offset_index_id = builder.reserve_id(SpvResultId);
builder.append(SpvFactory::load(offset_type_id, buffer_offset_index_id, buffer_offset_index_access_chain));

// Store the mapping from the parameter defining the buffer index offset to the variable it should be applied to
storage_buffer_offset_map[buffer_block_var_id] = buffer_offset_index_id;
}

buffer_index++;
}
}

Expand Down
16 changes: 9 additions & 7 deletions src/runtime/internal/memory_resources.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,18 +69,20 @@ struct MemoryBlock {
MemoryProperties properties; //< properties for the allocated block
};

// Client-facing struct for specifying a range of a memory region (eg for crops)
struct MemoryRange {
size_t head_offset = 0; //< byte offset from start of region
size_t tail_offset = 0; //< byte offset from end of region
struct RegionAllocation {
size_t offset = 0; //< offset from base address in block (in bytes)
size_t size = 0; //< allocated size in block (in bytes)
};

struct RegionIndexing {
int32_t offset = 0; //< indexing offset from start of region (used to adjust indices in compute shader to avoid alignment constraints for arbitrary crops)
};

// Client-facing struct for exchanging memory region allocation requests
struct MemoryRegion {
void *handle = nullptr; //< client data storing native handle (managed by alloc_block_region/free_block_region) or a pointer to region owning allocation
size_t offset = 0; //< offset from base address in block (in bytes)
size_t size = 0; //< allocated size (in bytes)
MemoryRange range; //< optional range (e.g. for handling crops, etc)
RegionAllocation allocation; //< allocation in parent block for region
RegionIndexing indexing; //< indexing adjustments for controlling access
bool dedicated = false; //< flag indicating whether allocation is one dedicated resource (or split/shared into other resources)
bool is_owner = true; //< flag indicating whether allocation is owned by this region, in which case handle is a native handle. Otherwise handle points to owning region of alloction.
MemoryProperties properties; //< properties for the allocated region
Expand Down
Loading
Loading