[BUG FIX] Derive BatchRenderer texture descriptor capacity from scene textures

include/madrona/render/render_mgr.hpp

@@ -30,6 +30,8 @@ class RenderManager {
         uint32_t agentViewWidth;
         uint32_t agentViewHeight;
 
+        uint32_t maxTextures;
+
         uint32_t numWorlds;
         uint32_t maxViewsPerWorld;
         uint32_t maxLightsPerWorld;

src/bridge/bindings.cpp

@@ -51,7 +51,8 @@ NB_MODULE(_gs_madrona_batch_renderer, m) {
             nb::ndarray<const int32_t, nb::shape<-1>, nb::device::cpu> enabled_geom_groups,
             nb::ndarray<const int32_t, nb::shape<-1, -1>, nb::device::cpu> geom_env_mask,
             bool add_cam_debug_geo,
-            bool use_rt)
+            bool use_rt,
+            int64_t max_textures)
         {
             GSModelGeometry mesh_geo {
                 .vertices = (math::Vector3 *)mesh_vertices.data(),
@@ -112,6 +113,7 @@ NB_MODULE(_gs_madrona_batch_renderer, m) {
                     .batchRenderViewHeight = (uint32_t)batch_render_view_height,
                     .addCamDebugGeometry = add_cam_debug_geo,
                     .useRT = use_rt,
+                    .maxTextures = (uint32_t)max_textures,
                 },
                 gs_model,
                 Optional<VisualizerGPUHandles>::none()
@@ -151,6 +153,7 @@ NB_MODULE(_gs_madrona_batch_renderer, m) {
            nb::arg("geom_env_mask") = nb::none(),
            nb::arg("add_cam_debug_geo") = false,
            nb::arg("use_rt") = false,
+           nb::arg("max_textures") = 0,
            nb::keep_alive<1, 32>())
         .def("init", [](Manager &mgr,
             nb::ndarray<nb::pytorch, const float, nb::shape<-1, -1, 3>> geom_pos,

src/bridge/mgr.cpp

@@ -82,11 +82,23 @@ static inline Optional<render::RenderManager> initRenderManager(
         max_instances_per_world += gs_model.numCams;
     }
 
+    // Texture descriptor capacity: enough for every texture in the scene, at
+    // least one (Vulkan forbids a zero-sized descriptor array), and raised to
+    // the requested floor when a larger override is supplied.
+    uint32_t max_textures = gs_model.numTextures;
+    if (mgr_cfg.maxTextures > max_textures) {
+        max_textures = mgr_cfg.maxTextures;
+    }
+    if (max_textures == 0) {
+        max_textures = 1;
+    }
+
     return render::RenderManager(render_api, render_dev, {
         .enableBatchRenderer = true,
         .renderMode = render::RenderManager::Config::RenderMode::RGBD,
         .agentViewWidth = mgr_cfg.batchRenderViewWidth,
         .agentViewHeight = mgr_cfg.batchRenderViewHeight,
+        .maxTextures = max_textures,
         .numWorlds = mgr_cfg.numWorlds,
         .maxViewsPerWorld = gs_model.numCams,
         .maxLightsPerWorld = gs_model.numLights,

src/bridge/mgr.hpp

@@ -78,6 +78,9 @@ class Manager {
         uint32_t batchRenderViewHeight;
         bool addCamDebugGeometry = false;
         bool useRT = false;
+        // Texture descriptor capacity floor. 0 (the default) derives the
+        // capacity from the loaded scene; a larger value reserves headroom.
+        uint32_t maxTextures = 0;
     };
 
     MGR_EXPORT Manager(

src/render/batch_renderer.cpp

@@ -249,7 +249,8 @@ static HeapArray<LayeredTarget> makeLayeredTargets(uint32_t width,
 ////////////////////////////////////////////////////////////////////////////////
 // RENDER PIPELINE CREATION                                                   //
 ////////////////////////////////////////////////////////////////////////////////
-static vk::PipelineShaders makeDrawShaders(const vk::Device &dev, VkSampler repeat_sampler)
+static vk::PipelineShaders makeDrawShaders(
+    const vk::Device &dev, VkSampler repeat_sampler, uint32_t max_textures)
 {
     const char *py_root_env = getenv("MADRONA_ROOT_PATH");
     std::filesystem::path root_dir = py_root_env ? (std::string(py_root_env) + "/src/render") : STRINGIFY(MADRONA_RENDER_DATA_DIR);
@@ -267,11 +268,11 @@ static vk::PipelineShaders makeDrawShaders(const vk::Device &dev, VkSampler repe
         dev, tmp_alloc, shaders,
         Span<const vk::BindingOverride>({
             vk::BindingOverride {
-                3, 0, VK_NULL_HANDLE, InternalConfig::maxTextures,
+                3, 0, VK_NULL_HANDLE, max_textures,
                 VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT
             },
             vk::BindingOverride {
-                4, 0, VK_NULL_HANDLE, InternalConfig::maxTextures,
+                4, 0, VK_NULL_HANDLE, max_textures,
                 VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT
             },
             vk::BindingOverride {3, 1, repeat_sampler, 1, 0}
@@ -284,9 +285,10 @@ static PipelineMP<1> makeDrawPipeline(const vk::Device &dev,
                                     VkRenderPass render_pass,
                                     uint32_t num_frames,
                                     uint32_t num_pools,
-                                    VkSampler repeat_sampler)
+                                    VkSampler repeat_sampler,
+                                    uint32_t max_textures)
 {
-    auto shaders = makeDrawShaders(dev, repeat_sampler);
+    auto shaders = makeDrawShaders(dev, repeat_sampler, max_textures);
 
     VkPipelineVertexInputStateCreateInfo vert_info {};
     VkPipelineInputAssemblyStateCreateInfo input_assembly_info {};
@@ -1602,7 +1604,8 @@ BatchRenderer::Impl::Impl(const Config &cfg, RenderContext &rctx):
     batchDraw(
         makeDrawPipeline(
             dev, rctx.pipelineCache, VK_NULL_HANDLE, 
-            consts::numDrawCmdBuffers * cfg.numFrames, 5, rctx.repeatSampler)),
+            consts::numDrawCmdBuffers * cfg.numFrames, 5, rctx.repeatSampler,
+            cfg.maxTextures)),
     createVisualization(
         makeComputePipeline(
             dev, rctx.pipelineCache, 1, sizeof(uint32_t) * 2,

src/render/batch_renderer.hpp

@@ -177,6 +177,7 @@ struct BatchRenderer {
         uint32_t maxViewsPerWorld;
         uint32_t maxInstancesPerWorld;
         uint32_t maxLightsPerWorld;
+        uint32_t maxTextures;
         uint32_t numFrames;
     };
 

src/render/render_common.hpp

@@ -26,7 +26,6 @@ inline constexpr uint32_t initMaxTransforms = 100000;
 inline constexpr uint32_t initMaxMatIndices = 100000;
 inline constexpr uint32_t shadowMapSize = 2048;
 inline constexpr uint32_t maxLights = 10;
-inline constexpr uint32_t maxTextures = 128;
 inline constexpr uint32_t maxComponents = 4;
 inline constexpr VkFormat gbufferFormat = VK_FORMAT_R16G16B16A16_SFLOAT;
 inline constexpr VkFormat skyFormatHighp = VK_FORMAT_R32G32B32A32_SFLOAT;

src/render/render_ctx.cpp

@@ -267,7 +267,8 @@ static VkRenderPass makeShadowRenderPass(const Device &dev,
 
 
 static PipelineShaders makeDrawShaders(
-    const Device &dev, VkSampler repeat_sampler, VkSampler clamp_sampler)
+    const Device &dev, VkSampler repeat_sampler, VkSampler clamp_sampler,
+    uint32_t max_textures)
 {
     (void)repeat_sampler;
     (void)clamp_sampler;
@@ -308,7 +309,7 @@ static PipelineShaders makeDrawShaders(
                 2,
                 0,
                 VK_NULL_HANDLE,
-                InternalConfig::maxTextures,
+                max_textures,
                 VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT,
             },
             BindingOverride {
@@ -342,9 +343,11 @@ static Pipeline<1> makeDrawPipeline(const Device &dev,
                                     VkRenderPass render_pass,
                                     VkSampler repeat_sampler,
                                     VkSampler clamp_sampler,
-                                    uint32_t num_frames)
+                                    uint32_t num_frames,
+                                    uint32_t max_textures)
 {
-    auto shaders = makeDrawShaders(dev, repeat_sampler, clamp_sampler);
+    auto shaders = makeDrawShaders(
+        dev, repeat_sampler, clamp_sampler, max_textures);
     VkPipelineVertexInputStateCreateInfo vert_info {};
     VkPipelineInputAssemblyStateCreateInfo input_assembly_info {};
     VkPipelineViewportStateCreateInfo viewport_info {};
@@ -1245,6 +1248,49 @@ static Sky loadSky(const vk::Device &dev, MemoryAllocator &alloc, VkQueue queue)
     };
 }
 
+static uint32_t validateTextureCapacity(
+    const Backend &backend, const Device &dev, uint32_t max_textures)
+{
+    if (max_textures == 0) {
+        FATAL("Texture descriptor capacity must be greater than zero");
+    }
+
+    VkPhysicalDeviceProperties properties;
+    backend.dt.getPhysicalDeviceProperties(dev.phy, &properties);
+    const VkPhysicalDeviceLimits &limits = properties.limits;
+
+    // The BatchRenderer draw pipeline binds the widest set of sampled-image
+    // arrays of any pipeline: two arrays of this capacity in the fragment
+    // stage (base-color and emissive). That path therefore dominates the
+    // limit check, so every array must fit within both the per-stage and the
+    // whole-set sampled-image budgets. This is a conservative bound: it treats
+    // the two texture arrays as the only sampled images in the stage, leaving
+    // the remaining headroom for the engine's few auxiliary bindings (shadow
+    // maps, sky).
+    constexpr uint32_t texture_arrays_per_stage = 2;
+    const uint32_t per_stage_cap =
+        limits.maxPerStageDescriptorSampledImages / texture_arrays_per_stage;
+    const uint32_t per_set_cap =
+        limits.maxDescriptorSetSampledImages / texture_arrays_per_stage;
+    const uint32_t capacity_cap =
+        per_stage_cap < per_set_cap ? per_stage_cap : per_set_cap;
+    if (max_textures > capacity_cap) {
+        FATAL(
+            "Requested texture descriptor capacity %u exceeds the supported "
+            "maximum of %u for this device: BatchRenderer binds %u sampled-image "
+            "arrays of this capacity, and the Vulkan limits are "
+            "maxPerStageDescriptorSampledImages=%u, "
+            "maxDescriptorSetSampledImages=%u",
+            max_textures,
+            capacity_cap,
+            texture_arrays_per_stage,
+            limits.maxPerStageDescriptorSampledImages,
+            limits.maxDescriptorSetSampledImages);
+    }
+
+    return max_textures;
+}
+
 RenderContext::RenderContext(
     APIBackend *render_backend,
     GPUDevice *render_dev,
@@ -1255,6 +1301,7 @@ RenderContext::RenderContext(
     renderQueue(makeGFXQueue(dev, 0)),
     br_width_(cfg.agentViewWidth),
     br_height_(cfg.agentViewHeight),
+    max_textures_(validateTextureCapacity(backend, dev, cfg.maxTextures)),
     pipelineCache(getPipelineCache(dev)),
     repeatSampler(makeImmutableSampler(dev, VK_SAMPLER_ADDRESS_MODE_REPEAT)),
     clampSampler(makeImmutableSampler(dev, VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)),
@@ -1264,7 +1311,9 @@ RenderContext::RenderContext(
     shadowPass(makeShadowRenderPass(
         dev, InternalConfig::varianceFormat, InternalConfig::depthFormat)),
     instanceCull(makeCullPipeline(dev, pipelineCache, InternalConfig::numFrames)),
-    objectDraw(makeDrawPipeline(dev, pipelineCache, renderPass, repeatSampler, clampSampler, InternalConfig::numFrames)),
+    objectDraw(makeDrawPipeline(
+        dev, pipelineCache, renderPass, repeatSampler, clampSampler,
+        InternalConfig::numFrames, max_textures_)),
     asset_desc_pool_cull_(dev, instanceCull.shaders, 1, 1),
     asset_desc_pool_draw_(dev, objectDraw.shaders, 1, 1),
     asset_desc_pool_mat_tx_(dev, objectDraw.shaders, 2, 1),
@@ -1290,14 +1339,14 @@ RenderContext::RenderContext(
     {
         VkDescriptorPoolSize pool_sizes[] = {
             { VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 25 },
-            { VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, InternalConfig::maxTextures * 2 },
+            { VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, max_textures_ * 2 },
             { VK_DESCRIPTOR_TYPE_SAMPLER, 1 }
         };
 
         VkDescriptorPoolCreateInfo pool_info = {};
         pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
         pool_info.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;
-        pool_info.maxSets = 10 + InternalConfig::maxTextures + 1;
+        pool_info.maxSets = 10 + max_textures_ + 1;
         pool_info.poolSizeCount = 3;
         pool_info.pPoolSizes = pool_sizes;
         REQ_VK(dev.dt.createDescriptorPool(dev.hdl, &pool_info, nullptr, &asset_pool_));
@@ -1433,7 +1482,7 @@ RenderContext::RenderContext(
             {
                 .binding = 0,
                 .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
-                .descriptorCount = InternalConfig::maxTextures,
+                .descriptorCount = max_textures_,
                 .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT,
                 .pImmutableSamplers = nullptr,
             },
@@ -1489,6 +1538,7 @@ RenderContext::RenderContext(
         cfg.maxViewsPerWorld,
         cfg.maxInstancesPerWorld,
         cfg.maxLightsPerWorld,
+        max_textures_,
         1
     };
 
@@ -1880,6 +1930,15 @@ CountT RenderContext::loadObjects(Span<const imp::SourceObject> src_objs,
 
     assert(loaded_assets_.size() == 0);
 
+    if (textures.size() > max_textures_) {
+        FATAL(
+            "Render asset has %u textures, exceeding the configured texture "
+            "descriptor capacity of %u. The capacity is fixed when the renderer "
+            "is created; reserve more by passing a larger 'max_textures' at "
+            "construction, or reduce the scene's texture count",
+            (uint32_t)textures.size(), max_textures_);
+    }
+
     int64_t num_total_vertices = 0;
     int64_t num_total_indices = 0;
     int64_t num_total_meshes = 0;

src/render/render_ctx.hpp

@@ -30,6 +30,7 @@ struct RenderContext {
 
     uint32_t br_width_;
     uint32_t br_height_;
+    uint32_t max_textures_;
 
     VkPipelineCache pipelineCache;
     VkSampler repeatSampler;