--- a/src/panfrost/vulkan/panvk_device.h +++ b/src/panfrost/vulkan/panvk_device.h @@ -6,6 +6,7 @@ #ifndef PANVK_DEVICE_H #define PANVK_DEVICE_H +#include #include #include "vk_debug_utils.h" @@ -25,6 +26,7 @@ #include "util/simple_mtx.h" #include "util/u_call_once.h" #include "util/u_printf.h" +#include "util/u_dynarray.h" #include "util/vma.h" /* On JM hardware, we need to allocate a buffer depending on vertex count. @@ -53,6 +55,8 @@ int queue_count; }; +struct panvk_priv_bo; + struct panvk_device { struct vk_device vk; @@ -63,8 +67,22 @@ struct util_vma_heap *priv_heap; bool split_heap; bool extended_range; + uint64_t priv_used; /* rocky-vulkan-llama: live priv-heap VA */ + uint64_t priv_peak; } as; + /* rocky-vulkan-llama: one device-wide TLS scratch BO shared across all + * command buffers on the serially-ordered compute queue (they never run + * concurrently), replacing the per-cmdbuf 304MB stack. Grow-only; superseded + * BOs are retired (kept alive while possibly in-flight) and freed at device + * destroy. Collapses N_inflight x stack -> 1 x stack. Gated by PANVK_SHARED_TLS. */ + struct { + simple_mtx_t lock; + struct panvk_priv_bo *bo; + uint64_t size; + struct util_dynarray retired; + } tls_scratch; + struct { struct pan_kmod_vm *vm; struct pan_kmod_dev *dev; @@ -190,6 +208,18 @@ { simple_mtx_lock(&device->as.lock); uint64_t address = util_vma_heap_alloc(heap, size, alignment); + if (heap == device->as.priv_heap) { + if (address) { + device->as.priv_used += size; + if (device->as.priv_used > device->as.priv_peak) + device->as.priv_peak = device->as.priv_used; + } + if (getenv("PANVK_AS_LOG")) + fprintf(stderr, "[panvk-as] priv alloc sz=%lu ok=%d used=%lu peak=%lu\n", + (unsigned long)size, address != 0, + (unsigned long)device->as.priv_used, + (unsigned long)device->as.priv_peak); + } simple_mtx_unlock(&device->as.lock); return address; } @@ -211,6 +241,8 @@ { simple_mtx_lock(&device->as.lock); util_vma_heap_free(heap, address, size); + if (heap == device->as.priv_heap) + device->as.priv_used -= size; simple_mtx_unlock(&device->as.lock); } --- a/src/panfrost/vulkan/panvk_vX_device.c +++ b/src/panfrost/vulkan/panvk_vX_device.c @@ -104,6 +104,9 @@ }; panvk_pool_init(&dev->mempools.exec, dev, NULL, NULL, &exec_pool_props); + + simple_mtx_init(&dev->tls_scratch.lock, mtx_plain); + util_dynarray_init(&dev->tls_scratch.retired, NULL); } static void @@ -112,6 +115,13 @@ panvk_pool_cleanup(&dev->mempools.rw); panvk_pool_cleanup(&dev->mempools.rw_nc); panvk_pool_cleanup(&dev->mempools.exec); + + util_dynarray_foreach(&dev->tls_scratch.retired, struct panvk_priv_bo *, b) + panvk_priv_bo_unref(*b); + util_dynarray_fini(&dev->tls_scratch.retired); + if (dev->tls_scratch.bo) + panvk_priv_bo_unref(dev->tls_scratch.bo); + simple_mtx_destroy(&dev->tls_scratch.lock); } static VkResult --- a/src/panfrost/vulkan/csf/panvk_vX_cmd_dispatch.c +++ b/src/panfrost/vulkan/csf/panvk_vX_cmd_dispatch.c @@ -13,6 +13,7 @@ #include "genxml/gen_macros.h" #include "panvk_buffer.h" +#include #include "panvk_cmd_alloc.h" #include "panvk_cmd_buffer.h" #include "panvk_cmd_desc_state.h" @@ -89,8 +90,14 @@ unsigned core_id_range; pan_query_core_count(&phys_dev->kmod.dev->props, &core_id_range); + /* rocky-vulkan-llama: size WLS with the grid-INDEPENDENT upper bound + * on instances (max_instances_per_core), not the per-dispatch grid. + * Makes wls_total_size constant per pipeline so the grow-only reuse + * below collapses to ~1 BO per distinct shader wls_size for the whole + * command buffer instead of one-per-new-high-water. Mirrors the + * indirect path and the JM/gallium single-batch-WLS model. */ tlsinfo.wls.instances = pan_calc_wls_instances( - &cs->cs.local_size, &phys_dev->kmod.dev->props, indirect ? NULL : dim); + &cs->cs.local_size, &phys_dev->kmod.dev->props, NULL); unsigned wls_total_size = pan_calc_total_wls_size( tlsinfo.wls.size, tlsinfo.wls.instances, core_id_range); @@ -100,10 +107,29 @@ * instance count) might differ significantly between dispatch commands, * rather than track a single maximum size, we might want to consider * multiple allocations for different size buckets. */ - tlsinfo.wls.ptr = - panvk_cmd_alloc_dev_mem(cmdbuf, tls, wls_total_size, 4096).gpu; - if (!tlsinfo.wls.ptr) - return 0; + /* rocky-vulkan-llama fix (addresses the TODO above): reuse a grow-only + * WLS buffer across dispatches. CSF allocated a fresh WLS BO per + * dispatch, exhausting the sub-4GB priv VA heap over a ~300-dispatch + * LLM prefill graph -> OOM at ub>64. Gated by PANVK_WLS_REUSE (A/B). */ + const bool wls_reuse = getenv("PANVK_WLS_REUSE") != NULL; + bool wls_did_reuse = false; + if (wls_reuse && wls_total_size <= cmdbuf->state.tls.wls_size) { + tlsinfo.wls.ptr = cmdbuf->state.tls.wls_ptr; + wls_did_reuse = true; + } else { + tlsinfo.wls.ptr = + panvk_cmd_alloc_dev_mem(cmdbuf, tls, wls_total_size, 4096).gpu; + if (!tlsinfo.wls.ptr) + return 0; + if (wls_reuse) { + cmdbuf->state.tls.wls_ptr = tlsinfo.wls.ptr; + cmdbuf->state.tls.wls_size = wls_total_size; + } + } + if (getenv("PANVK_WLS_LOG")) + fprintf(stderr, "[panvk-wls] total_size=%u %s prev_max=%u\n", + wls_total_size, wls_did_reuse ? "REUSE" : "ALLOC", + cmdbuf->state.tls.wls_size); } cmdbuf->state.tls.info.tls.size = --- a/src/panfrost/vulkan/csf/panvk_vX_cmd_buffer.c +++ b/src/panfrost/vulkan/csf/panvk_vX_cmd_buffer.c @@ -35,11 +35,36 @@ #include "pan_props.h" #include "pan_samples.h" +#include #include "util/bitscan.h" #include "vk_descriptor_update_template.h" #include "vk_format.h" #include "vk_synchronization.h" +static uint64_t +panvk_device_get_tls_scratch(struct panvk_device *dev, uint64_t size) +{ + simple_mtx_lock(&dev->tls_scratch.lock); + if (size > dev->tls_scratch.size) { + struct panvk_priv_bo *bo = NULL; + VkResult r = panvk_priv_bo_create( + dev, size, panvk_device_adjust_bo_flags(dev, PAN_KMOD_BO_FLAG_NO_MMAP), + VK_SYSTEM_ALLOCATION_SCOPE_DEVICE, &bo); + if (r != VK_SUCCESS) { + simple_mtx_unlock(&dev->tls_scratch.lock); + return 0; + } + if (dev->tls_scratch.bo) + util_dynarray_append(&dev->tls_scratch.retired, + dev->tls_scratch.bo); + dev->tls_scratch.bo = bo; + dev->tls_scratch.size = size; + } + uint64_t addr = dev->tls_scratch.bo->addr.dev; + simple_mtx_unlock(&dev->tls_scratch.lock); + return addr; +} + static void emit_tls(struct panvk_cmd_buffer *cmdbuf) { @@ -55,8 +80,15 @@ unsigned size = pan_get_total_stack_size(cmdbuf->state.tls.info.tls.size, thread_tls_alloc, core_id_range); - cmdbuf->state.tls.info.tls.ptr = - panvk_cmd_alloc_dev_mem(cmdbuf, tls, size, 4096).gpu; + if (getenv("PANVK_WLS_LOG")) + fprintf(stderr, "[panvk-tls] stack_size=%u (%.1f MB) tls_size=%u cores=%u\n", + size, size/1048576.0, cmdbuf->state.tls.info.tls.size, core_id_range); + if (getenv("PANVK_SHARED_TLS")) + cmdbuf->state.tls.info.tls.ptr = + panvk_device_get_tls_scratch(dev, ALIGN_POT(size, 4096)); + else + cmdbuf->state.tls.info.tls.ptr = + panvk_cmd_alloc_dev_mem(cmdbuf, tls, size, 4096).gpu; } assert(!cmdbuf->state.tls.info.wls.size); @@ -916,7 +948,11 @@ .owns_bos = true, .needs_locking = false, }; - panvk_pool_init(&cmdbuf->tls_pool, device, &pool->tls_bo_pool, &pool->tls_big_bo_pool, + /* rocky-vulkan-llama: pass NULL big_bo_pool so panvk_pool_reset FREES the tls + * pool's big BOs (WLS/TLS stack) on cmdbuf reset instead of recycling them into + * a free-list. On the sub-4GB CSF priv heap, the varying-size big BOs never match + * on recycle => VA grows monotonically => OOM. Freeing bounds it to live cmdbufs. */ + panvk_pool_init(&cmdbuf->tls_pool, device, &pool->tls_bo_pool, NULL, &tls_pool_props); for (uint32_t i = 0; i < ARRAY_SIZE(cmdbuf->utrace.uts); i++)