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

* Add an index offset during codegen to allow arbitrary offsets when accessing buffers to avoid restrictive alignment constraints.
CodeGen now adds one int32 buffer offset param for each buffer after all other scalar args
The runtime packs these params into the uniform buffer for each storage buffer
Crop device now computes an index offset (instead of a byte offset
Copy to/from device recomputes a byte offset from this index offset

* Refactor cleanup to remove unused MemoryRegion.range.
Adjusted relative offsets to use indexing for buffer copies.
Added RegionAllocation and RegionIndexing to clarify mapping.
Updated all affected interfaces.

* Clang formatting pass

Author: derek-gerstmann

src/CodeGen_Vulkan_Dev.cpp

@@ -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);
@@ -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 {
@@ -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 << " "
@@ -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);
@@ -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) {
@@ -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();
 }
@@ -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) {
@@ -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) {
@@ -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");
@@ -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};
@@ -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};
@@ -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) {
@@ -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) {
 
@@ -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++;
         }
     }
 

src/runtime/internal/memory_resources.h

@@ -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