add "extended_limits" flag support for metal: massive accels need this to trace properly.

Author: swfly

include/luisa/runtime/rhi/resource.h

@@ -155,6 +155,11 @@ struct ShaderOption {
     uint32_t max_registers{0};
     /// \brief Whether to measure time spent on each compilation phase.
     bool time_trace{false};
+    /// \brief Whether to enable extended acceleration structure limits.
+    /// \details If set to true, the shader will be compiled with support for
+    ///   massive instance counts (>2^24) in acceleration structures. Only has
+    ///   effect on the Metal backend; other backends ignore this option.
+    bool enable_extended_accel_limits{false};
     /// \brief A user-defined name for the shader.
     /// \details If provided, the shader will be read from or written to disk
     ///   via the `BinaryIO` object (passed to backends on device creation)
@@ -272,10 +277,12 @@ struct hash<compute::ShaderOption> {
         constexpr auto enable_fast_math_shift = 1u;
         constexpr auto enable_debug_info_shift = 2u;
         constexpr auto compile_only_shift = 3u;
+        constexpr auto enable_extended_accel_limits_shift = 4u;
         auto opt_hash = hash_value((static_cast<uint>(option.enable_cache) << enable_cache_shift) |
                                        (static_cast<uint>(option.enable_fast_math) << enable_fast_math_shift) |
                                        (static_cast<uint>(option.enable_debug_info) << enable_debug_info_shift) |
-                                       (static_cast<uint>(option.compile_only) << compile_only_shift),
+                                       (static_cast<uint>(option.compile_only) << compile_only_shift) |
+                                       (static_cast<uint>(option.enable_extended_accel_limits) << enable_extended_accel_limits_shift),
                                    seed);
         auto name_hash = hash_value(option.name, seed);
         return hash_combine({opt_hash, name_hash}, seed);

include/luisa/rust/api_types.h

@@ -521,6 +521,7 @@ typedef struct LCShaderOption {
     bool enable_debug_info;
     bool compile_only;
     bool time_trace;
+    bool enable_extended_accel_limits;
     uint32_t max_registers;
     const char *name;
     const char *native_include;

src/api/runtime.cpp

@@ -434,6 +434,7 @@ luisa_compute_shader_create(LCDevice device, LCKernelModule m, const LCShaderOpt
         .compile_only = option.compile_only,
         .max_registers = option.max_registers,
         .time_trace = option.time_trace,
+        .enable_extended_accel_limits = option.enable_extended_accel_limits,
         .name = luisa::string{option.name},
         .native_include = luisa::string{option.native_include},
     };

src/backends/metal/metal_accel.cpp

@@ -7,6 +7,10 @@
 
 namespace luisa::compute::metal {
 
+namespace {
+constexpr auto metal_accel_default_instance_limit = 1u << 24u;
+}// namespace
+
 MetalAccel::MetalAccel(MetalDevice *device, const AccelOption &option) noexcept
     : _update{device->builtin_update_accel_instances()},
       _option{option} { _resources.reserve(reserved_primitive_count); }
@@ -25,6 +29,7 @@ void MetalAccel::build(MetalCommandEncoder &encoder, AccelBuildCommand *command)
 
     auto device = encoder.device();
     auto instance_count = command->instance_count();
+    auto requires_extended_limits = instance_count > metal_accel_default_instance_limit;
     LUISA_ASSERT(instance_count > 0u, "Empty acceleration structure is not allowed.");
     if (auto size = instance_count * sizeof(MTL::AccelerationStructureInstanceDescriptor);
         _instance_buffer == nullptr || _instance_buffer->length() < size) {
@@ -91,8 +96,10 @@ void MetalAccel::build(MetalCommandEncoder &encoder, AccelBuildCommand *command)
     _requires_rebuild = _requires_rebuild /* pending rebuild */ ||
                         _descriptor == nullptr || old_instance_count != instance_count /* instance count has changed */ ||
                         _handle == nullptr /* not built before */ ||
+                        _requires_extended_limits != requires_extended_limits /* extended limits mode has changed */ ||
                         !_option.allow_update /* accel cannot be refitted */ ||
                         command->request() == AccelBuildRequest::FORCE_BUILD /* rebuild is forced */;
+    _requires_extended_limits = requires_extended_limits;
 
     // prepare the descriptor
     if (_requires_rebuild) {
@@ -114,6 +121,9 @@ void MetalAccel::build(MetalCommandEncoder &encoder, AccelBuildCommand *command)
             break;
     }
     if (_option.allow_update) { usage |= MTL::AccelerationStructureUsageRefit; }
+    if (_requires_extended_limits) {
+        usage |= MTL::AccelerationStructureUsageExtendedLimits;
+    }
     _descriptor->setUsage(usage);
 
     // update the descriptor

src/backends/metal/metal_accel.h

@@ -28,6 +28,7 @@ class MetalAccel {
     NS::String *_name{nullptr};
     AccelOption _option;
     bool _requires_rebuild{true};
+    bool _requires_extended_limits{false};
     spin_mutex _mutex;
 
 public:
@@ -48,11 +49,11 @@ class MetalAccel {
     [[nodiscard]] auto instance_buffer() const noexcept { return _instance_buffer; }
     [[nodiscard]] auto binding() const noexcept { return Binding{_handle->gpuResourceID(), _instance_buffer->gpuAddress()}; }
     [[nodiscard]] auto pointer_to_handle() const noexcept { return const_cast<void *>(static_cast<const void *>(&_handle)); }
+    [[nodiscard]] auto requires_extended_limits() const noexcept { return _requires_extended_limits; }
     void set_name(luisa::string_view name) noexcept;
     void mark_resource_usages(MetalCommandEncoder &encoder,
                               MTL::ComputeCommandEncoder *command_encoder,
                               MTL::ResourceUsage usage) noexcept;
 };
 
 }// namespace luisa::compute::metal
-

src/backends/metal/metal_builtin/metal_device_lib.metal

@@ -8,6 +8,14 @@ using namespace metal;
 #define lc_assume(...) __builtin_assume(__VA_ARGS__)
 #define lc_assert(...)// TODO: implement assert?
 
+#ifdef LUISA_ENABLE_EXTENDED_LIMITS
+#define LC_EXTENDED_LIMITS_TAG , extended_limits
+#define LC_INSTANCE_ACCELERATION_STRUCTURE instance_acceleration_structure<extended_limits>
+#else
+#define LC_EXTENDED_LIMITS_TAG
+#define LC_INSTANCE_ACCELERATION_STRUCTURE instance_acceleration_structure
+#endif
+
 template<typename Ptr>
 [[nodiscard]] inline auto lc_address_of_impl(Ptr ptr) {
     return reinterpret_cast<ulong>(ptr);
@@ -768,19 +776,19 @@ static_assert(sizeof(LCInstance) == 64u, "");
 
 struct LCAccel {
 #ifdef LUISA_ENABLE_MOTION_BLUR
-    acceleration_structure<instancing, primitive_motion, instance_motion> handle;
+    acceleration_structure<instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> handle;
 #else
-    instance_acceleration_structure handle;
+    LC_INSTANCE_ACCELERATION_STRUCTURE handle;
 #endif
     device LCInstance *__restrict__ instances;
 };
 
 [[nodiscard]] inline auto lc_intersector_base() {
 #ifdef LUISA_ENABLE_CURVE
 #ifdef LUISA_ENABLE_MOTION_BLUR
-    intersector<triangle_data, curve_data, instancing, primitive_motion, instance_motion> i;
+    intersector<triangle_data, curve_data, instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> i;
 #else
-    intersector<triangle_data, curve_data, instancing> i;
+    intersector<triangle_data, curve_data, instancing LC_EXTENDED_LIMITS_TAG> i;
 #endif
     i.assume_geometry_type(geometry_type::triangle | geometry_type::curve);
     i.assume_curve_type(curve_type::round);
@@ -807,9 +815,9 @@ struct LCAccel {
 #endif
 #else
 #ifdef LUISA_ENABLE_MOTION_BLUR
-    intersector<triangle_data, instancing, primitive_motion, instance_motion> i;
+    intersector<triangle_data, instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> i;
 #else
-    intersector<triangle_data, instancing> i;
+    intersector<triangle_data, instancing LC_EXTENDED_LIMITS_TAG> i;
 #endif
     i.assume_geometry_type(geometry_type::triangle);
 #endif
@@ -889,25 +897,25 @@ struct LCAccel {
 
 struct LCRayQuery {
 #ifdef LUISA_ENABLE_MOTION_BLUR
-    acceleration_structure<instancing, primitive_motion, instance_motion> accel;
+    acceleration_structure<instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> accel;
 #else
-    instance_acceleration_structure accel;
+    LC_INSTANCE_ACCELERATION_STRUCTURE accel;
 #endif
     ray ray;
     uint mask;
     float time;
     bool terminate_on_first_hit;
 #ifdef LUISA_ENABLE_CURVE
 #ifdef LUISA_ENABLE_MOTION_BLUR
-    thread intersection_query<triangle_data, curve_data, instancing, primitive_motion, instance_motion> *i;
+    thread intersection_query<triangle_data, curve_data, instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> *i;
 #else
-    thread intersection_query<triangle_data, curve_data, instancing> *i;
+    thread intersection_query<triangle_data, curve_data, instancing LC_EXTENDED_LIMITS_TAG> *i;
 #endif
 #else
 #ifdef LUISA_ENABLE_MOTION_BLUR
-    thread intersection_query<triangle_data, instancing, primitive_motion, instance_motion> *i;
+    thread intersection_query<triangle_data, instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> *i;
 #else
-    thread intersection_query<triangle_data, instancing> *i;
+    thread intersection_query<triangle_data, instancing LC_EXTENDED_LIMITS_TAG> *i;
 #endif
 #endif
 };
@@ -941,15 +949,15 @@ struct LCRayQuery {
 void ray_query_init(thread LCRayQuery &q,
 #ifdef LUISA_ENABLE_CURVE
 #ifdef LUISA_ENABLE_MOTION_BLUR
-                    thread intersection_query<triangle_data, curve_data, instancing, primitive_motion, instance_motion> &i,
+                    thread intersection_query<triangle_data, curve_data, instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> &i,
 #else
-                    thread intersection_query<triangle_data, curve_data, instancing> &i,
+                    thread intersection_query<triangle_data, curve_data, instancing LC_EXTENDED_LIMITS_TAG> &i,
 #endif
 #else
 #ifdef LUISA_ENABLE_MOTION_BLUR
-                    thread intersection_query<triangle_data, instancing, primitive_motion, instance_motion> &i,
+                    thread intersection_query<triangle_data, instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> &i,
 #else
-                    thread intersection_query<triangle_data, instancing> &i,
+                    thread intersection_query<triangle_data, instancing LC_EXTENDED_LIMITS_TAG> &i,
 #endif
 #endif
                     bool has_procedural_branch) {
@@ -993,18 +1001,18 @@ void ray_query_init(thread LCRayQuery &q,
 #ifdef LUISA_ENABLE_CURVE
 #ifdef LUISA_ENABLE_MOTION_BLUR
 #define LC_RAY_QUERY_SHADOW_VARIABLE(q) \
-    intersection_query<triangle_data, curve_data, instancing, primitive_motion, instance_motion> q##_i
+    intersection_query<triangle_data, curve_data, instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> q##_i
 #else
 #define LC_RAY_QUERY_SHADOW_VARIABLE(q) \
-    intersection_query<triangle_data, curve_data, instancing> q##_i
+    intersection_query<triangle_data, curve_data, instancing LC_EXTENDED_LIMITS_TAG> q##_i
 #endif
 #else
 #ifdef LUISA_ENABLE_MOTION_BLUR
 #define LC_RAY_QUERY_SHADOW_VARIABLE(q) \
-    intersection_query<triangle_data, instancing, primitive_motion, instance_motion> q##_i
+    intersection_query<triangle_data, instancing, primitive_motion, instance_motion LC_EXTENDED_LIMITS_TAG> q##_i
 #else
 #define LC_RAY_QUERY_SHADOW_VARIABLE(q) \
-    intersection_query<triangle_data, instancing> q##_i
+    intersection_query<triangle_data, instancing LC_EXTENDED_LIMITS_TAG> q##_i
 #endif
 #endif
 

src/backends/metal/metal_device.cpp

@@ -459,7 +459,13 @@ ShaderCreationInfo MetalDevice::create_shader(const ShaderOption &option, Functi
         }
 
         // create shader
-        auto pipeline = _compiler->compile(scratch.string_view(), option, metadata);
+        luisa::string source;
+        if (option.enable_extended_accel_limits) {
+            source = luisa::format("#define LUISA_ENABLE_EXTENDED_LIMITS\n{}", scratch.string_view());
+        } else {
+            source = luisa::string{scratch.string_view()};
+        }
+        auto pipeline = _compiler->compile(source, option, metadata);
         auto shader = luisa::new_with_allocator<MetalShader>(
             this, std::move(pipeline),
             std::move(metadata.argument_usages),

src/rust/luisa_compute_api_types/src/lib.rs

@@ -78,6 +78,7 @@ pub struct ShaderOption {
     pub enable_debug_info: bool,
     pub compile_only: bool,
     pub time_trace: bool,
+    pub enable_extended_accel_limits: bool,
     pub max_registers: u32,
     pub name: *const std::ffi::c_char,
     pub native_include: *const std::ffi::c_char,
@@ -95,6 +96,7 @@ impl Default for ShaderOption {
             enable_debug_info: false,
             compile_only: false,
             time_trace: false,
+            enable_extended_accel_limits: false,
             max_registers: 0,
             name: std::ptr::null(),
             native_include: std::ptr::null(),