Fix conform request in the Vulkan memory allocator to avoid rounding up on subsequent passes.

src/runtime/internal/memory_resources.h

@@ -161,6 +161,30 @@ ALWAYS_INLINE size_t conform_size(size_t offset, size_t size, size_t alignment,
     }
 }
 
+// Updates request with conformed alignment, offset and size, computed from the given required size and alignment constraints,
+// rounding up to the nearest multiple if specified. The request is updated in-place.
+// -- Required alignment must be power of two!
+ALWAYS_INLINE void conform_memory_request(MemoryRequest *request, size_t required_size, size_t required_alignment, size_t nearest_multiple) {
+    size_t actual_alignment = conform_alignment(request->alignment, required_alignment);
+
+    // Ensure the request ends on an aligned address.
+    // NOTE: use the conformed alignment and not the input alignment. Otherwise, conform is not idempotent:
+    // A first pass that uses the original nearest_multiple writes a new alignment such that a subsequent pass
+    // sees a larger alignment and writes a larger nearest_multiple, which results in a larger size calculated
+    // in the second pass. The block allocator requires that conform is idempotent because it calls conform
+    // defensively several times at different layers.
+    if (actual_alignment > nearest_multiple) {
+        request->properties.nearest_multiple = actual_alignment;
+    }
+
+    size_t actual_offset = aligned_offset(request->offset, actual_alignment);
+    size_t actual_size = conform_size(actual_offset, required_size, actual_alignment, request->properties.nearest_multiple);
+
+    request->size = actual_size;
+    request->alignment = actual_alignment;
+    request->offset = actual_offset;
+}
+
 // Clamps the given value to be within the [min_value, max_value] range
 ALWAYS_INLINE size_t clamped_size(size_t value, size_t min_value, size_t max_value) {
     size_t result = (value < min_value) ? min_value : value;

src/runtime/internal/region_allocator.h

@@ -329,7 +329,7 @@ BlockRegion *RegionAllocator::find_block_region(void *user_context, const Memory
             debug(user_context) << "RegionAllocator: found suitable region ( "
                                 << "user_context=" << (void *)(user_context) << " "
                                 << "block_resource=" << (void *)block << " "
-                                << "block_size=" << (uint32_t)block->memory.allocation.size << " "
+                                << "block_size=" << (uint32_t)block->memory.size << " "
                                 << "block_reserved=" << (uint32_t)block->reserved << " "
                                 << "requested_size=" << (uint32_t)request.size << " "
                                 << "requested_is_dedicated=" << (request.dedicated ? "true" : "false") << " "

src/runtime/vulkan_memory.h

@@ -612,10 +612,24 @@ int VulkanMemoryAllocator::conform_block_request(void *instance_ptr, MemoryReque
                    << "uniform_buffer_offset_alignment=" << (uint32_t)instance->physical_device_limits.minUniformBufferOffsetAlignment << ", "
                    << "storage_buffer_offset_alignment=" << (uint32_t)instance->physical_device_limits.minStorageBufferOffsetAlignment << ", "
                    << "dedicated=" << (request->dedicated ? "true" : "false") << ")\n";
+
+    const MemoryRequest original_request = *request;
 #endif
 
-    request->size = memory_requirements.size;
     request->properties.alignment = memory_requirements.alignment;
+    conform_memory_request(request, memory_requirements.size, memory_requirements.alignment, instance->config.nearest_multiple);
+
+#if defined(HL_VK_DEBUG_MEM)
+    if ((request->size != original_request.size) || (request->alignment != original_request.alignment) || (request->offset != original_request.offset)) {
+        debug(nullptr) << "VulkanMemoryAllocator: Adjusting request to match requirements (\n"
+                       << "  size = " << (uint64_t)original_request.size << " => " << (uint64_t)request->size << ",\n"
+                       << "  alignment = " << (uint64_t)original_request.alignment << " => " << (uint64_t)request->alignment << ",\n"
+                       << "  offset = " << (uint64_t)original_request.offset << " => " << (uint64_t)request->offset << ",\n"
+                       << "  required.size = " << (uint64_t)memory_requirements.size << ",\n"
+                       << "  required.alignment = " << (uint64_t)memory_requirements.alignment << "\n)\n";
+    }
+#endif
+
     return halide_error_code_success;
 }
 
@@ -627,12 +641,12 @@ int VulkanMemoryAllocator::allocate_block(void *instance_ptr, MemoryBlock *block
 
     void *user_context = instance->owner_context;
     if ((instance->device == nullptr) || (instance->physical_device == nullptr)) {
-        error(user_context) << "VulkanBlockAllocator: Unable to deallocate block! Invalid device handle!\n";
+        error(user_context) << "VulkanBlockAllocator: Unable to allocate block! Invalid device handle!\n";
         return halide_error_code_internal_error;
     }
 
     if (block == nullptr) {
-        error(user_context) << "VulkanBlockAllocator: Unable to deallocate block! Invalid pointer!\n";
+        error(user_context) << "VulkanBlockAllocator: Unable to allocate block! Invalid pointer!\n";
         return halide_error_code_internal_error;
     }
 
@@ -1006,28 +1020,22 @@ int VulkanMemoryAllocator::conform(void *user_context, MemoryRequest *request) {
         }
     }
 
-    // Ensure the request ends on an aligned address
-    if (request->alignment > config.nearest_multiple) {
-        request->properties.nearest_multiple = request->alignment;
-    }
+#if defined(HL_VK_DEBUG_MEM)
+    const MemoryRequest original_request = *request;
+#endif
 
-    size_t actual_alignment = conform_alignment(request->alignment, memory_requirements.alignment);
-    size_t actual_offset = aligned_offset(request->offset, actual_alignment);
-    size_t actual_size = conform_size(actual_offset, memory_requirements.size, actual_alignment, request->properties.nearest_multiple);
+    conform_memory_request(request, memory_requirements.size, memory_requirements.alignment, config.nearest_multiple);
 
 #if defined(HL_VK_DEBUG_MEM)
-    if ((request->size != actual_size) || (request->alignment != actual_alignment) || (request->offset != actual_offset)) {
+    if ((request->size != original_request.size) || (request->alignment != original_request.alignment) || (request->offset != original_request.offset)) {
         debug(nullptr) << "VulkanMemoryAllocator: Adjusting request to match requirements (\n"
-                       << "  size = " << (uint64_t)request->size << " => " << (uint64_t)actual_size << ",\n"
-                       << "  alignment = " << (uint64_t)request->alignment << " => " << (uint64_t)actual_alignment << ",\n"
-                       << "  offset = " << (uint64_t)request->offset << " => " << (uint64_t)actual_offset << ",\n"
+                       << "  size = " << (uint64_t)original_request.size << " => " << (uint64_t)request->size << ",\n"
+                       << "  alignment = " << (uint64_t)original_request.alignment << " => " << (uint64_t)request->alignment << ",\n"
+                       << "  offset = " << (uint64_t)original_request.offset << " => " << (uint64_t)request->offset << ",\n"
                        << "  required.size = " << (uint64_t)memory_requirements.size << ",\n"
                        << "  required.alignment = " << (uint64_t)memory_requirements.alignment << "\n)\n";
     }
 #endif
-    request->size = actual_size;
-    request->alignment = actual_alignment;
-    request->offset = actual_offset;
 
     return halide_error_code_success;
 }

test/runtime/block_allocator.cpp

@@ -460,6 +460,70 @@ int main(int argc, char **argv) {
             }
         }
 
+        // test conform_memory_request
+        {
+            // conform_memory_request must be idempotent: the block allocator conforms
+            // requests defensively at several layers, feeding an already-conformed
+            // request back in. Restrict the sweep to request alignments that do not
+            // exceed the required alignment (both powers of two) so the conformed
+            // alignment stays a power of two and aligned_offset does not abort.
+            for (size_t required_alignment = 1; required_alignment <= 128; required_alignment *= 2) {
+                for (size_t nearest_multiple = 0; nearest_multiple <= 64; nearest_multiple = nearest_multiple ? nearest_multiple * 2 : 1) {
+                    for (size_t request_alignment = 0; request_alignment <= required_alignment; request_alignment = request_alignment ? request_alignment * 2 : 1) {
+                        for (size_t required_size = 1; required_size <= 100; ++required_size) {
+                            for (size_t offset = 0; offset <= 64; offset += 16) {
+                                MemoryRequest request = {0};
+                                request.offset = offset;
+                                request.alignment = request_alignment;
+
+                                conform_memory_request(&request, required_size, required_alignment, nearest_multiple);
+                                MemoryRequest conformed = request;
+
+                                conform_memory_request(&request, required_size, required_alignment, nearest_multiple);
+
+                                halide_abort_if_false(user_context, request.size == conformed.size);
+                                halide_abort_if_false(user_context, request.alignment == conformed.alignment);
+                                halide_abort_if_false(user_context, request.offset == conformed.offset);
+                                halide_abort_if_false(user_context, request.properties.nearest_multiple == conformed.properties.nearest_multiple);
+
+                                halide_abort_if_false(user_context, conformed.size >= required_size);
+                                halide_abort_if_false(user_context, conformed.size >= conformed.alignment);
+                                halide_abort_if_false(user_context, conformed.alignment == conform_alignment(request_alignment, required_alignment));
+                                halide_abort_if_false(user_context, conformed.offset == aligned_offset(offset, conformed.alignment));
+                                halide_abort_if_false(user_context, (conformed.offset % conformed.alignment) == 0);
+                                if (conformed.properties.nearest_multiple > 0) {
+                                    halide_abort_if_false(user_context, (conformed.size % conformed.properties.nearest_multiple) == 0);
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+
+            // Regression scenarios from the vulkan allocator fixes (Mesa lavapipe/RADV
+            // reporting 64-byte buffer alignment): a buffer whose size is a multiple of
+            // the config nearest_multiple but not of the device alignment must round up
+            // to the device alignment and stay stable across repeated conforms.
+            size_t regression_required_size[2] = {96, 9338976};
+            size_t regression_required_alignment[2] = {64, 64};
+            size_t regression_nearest_multiple[2] = {32, 32};
+            for (int i = 0; i < 2; ++i) {
+                MemoryRequest request = {0};
+
+                conform_memory_request(&request, regression_required_size[i], regression_required_alignment[i], regression_nearest_multiple[i]);
+                MemoryRequest conformed = request;
+
+                conform_memory_request(&request, regression_required_size[i], regression_required_alignment[i], regression_nearest_multiple[i]);
+
+                halide_abort_if_false(user_context, conformed.size >= regression_required_size[i]);
+                halide_abort_if_false(user_context, (conformed.size % regression_required_alignment[i]) == 0);
+                halide_abort_if_false(user_context, request.size == conformed.size);
+                halide_abort_if_false(user_context, request.alignment == conformed.alignment);
+                halide_abort_if_false(user_context, request.offset == conformed.offset);
+                halide_abort_if_false(user_context, request.properties.nearest_multiple == conformed.properties.nearest_multiple);
+            }
+        }
+
         BlockAllocator::destroy(user_context, instance);
         HALIDE_CHECK(user_context, get_allocated_system_memory() == 0);
     }