diff --git a/quadrants/codegen/spirv/detail/spirv_codegen.h b/quadrants/codegen/spirv/detail/spirv_codegen.h index 70add2f760..e7e3cc4685 100644 --- a/quadrants/codegen/spirv/detail/spirv_codegen.h +++ b/quadrants/codegen/spirv/detail/spirv_codegen.h @@ -17,6 +17,8 @@ #include #include +#include + namespace quadrants::lang { namespace spirv { namespace detail { @@ -154,6 +156,14 @@ class TaskCodegen : public IRVisitor { std::shared_ptr ir_; // spirv binary code builder std::unordered_map, spirv::Value, BufferInfoTypeTupleHasher> buffer_value_map_; std::unordered_map, uint32_t, BufferInfoTypeTupleHasher> buffer_binding_map_; + // All existing type views of each underlying storage buffer, in creation order. When a second or later + // view is minted in `get_buffer_value`, we decorate every entry here with `Aliased` so the driver is + // forbidden from assuming the views don't alias -- otherwise a plain load through one view is not + // ordered against an atomic through another view of the same memory, silently zeroing gradients on the + // load-and-clear reverse-mode pattern. See `get_buffer_value` for the decoration site and the commit + // message for the full failure matrix. + std::unordered_map, BufferInfoHasher> buffer_views_by_buffer_; + std::unordered_set aliased_decorated_buffer_ids_; std::vector shared_array_binds_; spirv::Value kernel_function_; spirv::Label kernel_return_label_; diff --git a/quadrants/codegen/spirv/spirv_codegen.cpp b/quadrants/codegen/spirv/spirv_codegen.cpp index 8439622967..f98281bd00 100644 --- a/quadrants/codegen/spirv/spirv_codegen.cpp +++ b/quadrants/codegen/spirv/spirv_codegen.cpp @@ -2096,16 +2096,44 @@ spirv::Value TaskCodegen::at_buffer(const Stmt *ptr, DataType dt) { return ret; } +// For primitive float types, access the storage buffer through the native type view instead of +// the uint-punned view. SPIR-V / Vulkan treats each (descriptor_set, binding) as a distinct +// variable; `at_buffer` creates a new binding per (buffer, element_type) pair, so the u32 view +// of a buffer and its f32 view are different variables pointing to the same memory. Without an +// `Aliased` decoration the driver / SPIRV-Tools is free to assume they do not alias, meaning a +// plain `OpLoad` through the u32 view is not ordered against a preceding `OpAtomicFAddEXT` on +// the f32 view at the same address. The reverse-mode pattern `m.grad[i][j,k] += loss.grad; +// tmp = m.grad[i][j,k]; m.grad[i][j,k] = 0; n.grad += tmp * factor` hits this: the load reads +// the stale zero initial value, `tmp = 0`, and the adjoint never propagates (`test_ad_dynamic_index.py:: +// test_matrix_non_constant_index[arch=vulkan]` asserts `0.0 == 1.0`). Using the native f32 view +// for plain load/store keeps them on the same binding as the atomic and removes the aliasing +// question entirely. Integer types (i32/u32/i16/u16/i8/u8) already route through their own uint +// view which matches the atomic path, so those stay as-is. `u1` stays on u8 because u1 has no +// native SPIR-V storage representation. +static DataType pick_buffer_access_type(DataType dt, const spirv::Value &ptr_val, spirv::IRBuilder &ir) { + if (dt->is_primitive(PrimitiveTypeID::u1)) { + return PrimitiveType::u8; + } + if (ptr_val.stype.dt == PrimitiveType::u64) { + return dt; + } + // Explicit whitelist of the real primitives we route natively, replacing the prior + // open-ended `is_real(dt)` predicate. Any future real-like primitive (e.g. a bfloat16, or an + // fp8 variant) would not have an audited SPIR-V storage-capability story yet -- rather than + // silently fall into the native-view branch, it must be added here deliberately after the + // storage-capability plumbing for its bit width is confirmed (see the + // `CapabilityStorageBuffer{8,16}BitAccess` emissions in `spirv_ir_builder.cpp`). + if (dt->is_primitive(PrimitiveTypeID::f16) || dt->is_primitive(PrimitiveTypeID::f32) || + dt->is_primitive(PrimitiveTypeID::f64)) { + return dt; + } + return ir.get_quadrants_uint_type(dt); +} + spirv::Value TaskCodegen::load_buffer(const Stmt *ptr, DataType dt) { spirv::Value ptr_val = ir_->query_value(ptr->raw_name()); - DataType ti_buffer_type = ir_->get_quadrants_uint_type(dt); - - if (dt->is_primitive(PrimitiveTypeID::u1)) { - ti_buffer_type = PrimitiveType::u8; - } else if (ptr_val.stype.dt == PrimitiveType::u64) { - ti_buffer_type = dt; - } + DataType ti_buffer_type = pick_buffer_access_type(dt, ptr_val, *ir_); auto buf_ptr = at_buffer(ptr, ti_buffer_type); auto val_bits = ir_->load_variable(buf_ptr, ir_->get_primitive_type(ti_buffer_type)); @@ -2117,18 +2145,31 @@ spirv::Value TaskCodegen::load_buffer(const Stmt *ptr, DataType dt) { void TaskCodegen::store_buffer(const Stmt *ptr, spirv::Value val) { spirv::Value ptr_val = ir_->query_value(ptr->raw_name()); - DataType ti_buffer_type = ir_->get_quadrants_uint_type(val.stype.dt); - + DataType ti_buffer_type = pick_buffer_access_type(val.stype.dt, ptr_val, *ir_); if (val.stype.dt->is_primitive(PrimitiveTypeID::u1)) { + // Stores go through i8 (matching the original path) so a signed i1 narrowing is preserved. ti_buffer_type = PrimitiveType::i8; - } else if (ptr_val.stype.dt == PrimitiveType::u64) { - ti_buffer_type = val.stype.dt; } auto buf_ptr = at_buffer(ptr, ti_buffer_type); - auto val_bits = val.stype.dt == ti_buffer_type - ? val - : ir_->make_value(spv::OpBitcast, ir_->get_primitive_type(ti_buffer_type), val); + spirv::Value val_bits; + if (val.stype.dt == ti_buffer_type) { + val_bits = val; + } else if (val.stype.dt->is_primitive(PrimitiveTypeID::u1)) { + // SPIR-V `OpBitcast` rejects bool operands (spec: operand must be numerical scalar / vector or + // pointer). Before this fix, a `u1` field / ndarray store emitted + // `OpBitcast %char %bool_val` and validated as + // `Expected input to be a pointer or int or float vector or scalar: Bitcast`. Most drivers + // ignore that and crash inside the pipeline compiler (observed on Mesa RADV: a hard SIGSEGV + // inside `libvulkan_radeon.so::create_compute_pipeline` the moment the offending kernel is + // registered). Route through `IRBuilder::cast`, which lowers `bool -> int` to `OpSelect` + // picking `1` or `0` of the target type -- that's the canonical spec-compliant way to widen a + // bool, matches what `load_buffer` already does on the reverse path, and keeps the + // "bool serialises as 0 / 1" behaviour every user of `to_numpy()` / `from_numpy()` depends on. + val_bits = ir_->cast(ir_->get_primitive_type(ti_buffer_type), val); + } else { + val_bits = ir_->make_value(spv::OpBitcast, ir_->get_primitive_type(ti_buffer_type), val); + } ir_->store_variable(buf_ptr, val_bits); } @@ -2166,6 +2207,37 @@ spirv::Value TaskCodegen::get_buffer_value(BufferInfo buffer, DataType dt) { ir_->decorate(spv::OpDecorate, buffer_value, spv::DecorationVolatile); } buffer_value_map_[key] = buffer_value; + + // Type-punned views of the same underlying `VkBuffer` need an explicit `Aliased` decoration. Every + // `(BufferInfo, element_type)` pair here gets its own `OpVariable` with a fresh `DescriptorSet` / + // `Binding`, so a field read through the `u32` view and a preceding `OpAtomicFAddEXT` through the + // `f32` view are separate SPIR-V variables pointing to the same memory. Without `Aliased` the + // driver is spec-free to assume they don't alias, and the plain load is not ordered against the + // atomic write -- the load reads the stale zero initial value and the reverse-mode adjoint silently + // drops to zero (`test_ad_dynamic_index.py::test_matrix_non_constant_index[arch=vulkan]` reproduces + // this on every device that exposes `shaderBufferFloat32AtomicAdd`). + // + // The aliasing hazard applies across every pairing of views on the same buffer, not just + // (native-float atomic, plain load): the CAS-emulation atomic path on devices without + // `shaderBufferFloat32AtomicAdd` routes float atomics through the `u32` view while other sites may + // emit atomics min / max / mul directly against the `u32` view -- each such pairing against a + // plain-load-through-any-view needs the same decoration. Decorating here rather than at first + // `at_buffer` call keeps the fix independent of which call site happens to introduce the second + // view. + // + // Decorate lazily: a single-view buffer stays un-decorated (no perf cost from disabled + // cross-variable scheduling optimizations), and only when a buffer gets its second distinct view do + // we retroactively decorate every view -- existing ones through the id-set guard below, and the new + // one unconditionally in the same sweep. + auto &views = buffer_views_by_buffer_[buffer]; + views.push_back(buffer_value); + if (views.size() >= 2) { + for (const auto &v : views) { + if (aliased_decorated_buffer_ids_.insert(v.id).second) { + ir_->decorate(spv::OpDecorate, v, spv::DecorationAliased); + } + } + } QD_TRACE("buffer name = {}, value = {}", buffer_instance_name(buffer), buffer_value.id); return buffer_value; diff --git a/quadrants/codegen/spirv/spirv_ir_builder.cpp b/quadrants/codegen/spirv/spirv_ir_builder.cpp index 0553d377cb..ba8420e432 100644 --- a/quadrants/codegen/spirv/spirv_ir_builder.cpp +++ b/quadrants/codegen/spirv/spirv_ir_builder.cpp @@ -54,6 +54,22 @@ void IRBuilder::init_header() { if (caps_->get(cap::spirv_has_int16)) { ib_.begin(spv::OpCapability).add(spv::CapabilityInt16).commit(&header_); } + // `CapabilityStorageBuffer{8,16}BitAccess` gate narrow-typed loads / stores through a + // descriptor-bound `StorageBuffer` pointer (e.g. `OpLoad %_ptr_StorageBuffer_ushort`). The + // existing codegen has been emitting these narrow-typed accesses via the uint-punning path in + // `load_buffer` / `store_buffer` for a while (`get_quadrants_uint_type(i16) = u16`, + // `get_quadrants_uint_type(i8) = u8`), which strict Vulkan validation requires these capabilities + // for -- so we emit them unconditionally whenever the queried feature is set, independent of + // whether the current kernel actually uses a narrow type. The cost is a single extra + // `OpCapability` word in the shader header; the upside is that every 16-bit / 8-bit field or + // ndarray access is spec-compliant on drivers that enforce the letter of + // `SPV_KHR_{8,16}bit_storage`. + if (caps_->get(cap::spirv_has_storage_buffer_8bit_access)) { + ib_.begin(spv::OpCapability).add(spv::CapabilityStorageBuffer8BitAccess).commit(&header_); + } + if (caps_->get(cap::spirv_has_storage_buffer_16bit_access)) { + ib_.begin(spv::OpCapability).add(spv::CapabilityStorageBuffer16BitAccess).commit(&header_); + } if (caps_->get(cap::spirv_has_int64)) { ib_.begin(spv::OpCapability).add(spv::CapabilityInt64).commit(&header_); } @@ -77,6 +93,17 @@ void IRBuilder::init_header() { ib_.begin(spv::OpExtension).add("SPV_KHR_storage_buffer_storage_class").commit(&header_); + // `SPV_KHR_{8,16}bit_storage` is paired with `CapabilityStorageBuffer{8,16}BitAccess` above. + // Both the capability and the extension are needed for narrow-typed `StorageBuffer` loads / + // stores to validate on Vulkan; declaring only the capability without the extension is + // ill-formed SPIR-V. + if (caps_->get(cap::spirv_has_storage_buffer_8bit_access)) { + ib_.begin(spv::OpExtension).add("SPV_KHR_8bit_storage").commit(&header_); + } + if (caps_->get(cap::spirv_has_storage_buffer_16bit_access)) { + ib_.begin(spv::OpExtension).add("SPV_KHR_16bit_storage").commit(&header_); + } + if (caps_->get(cap::spirv_has_no_integer_wrap_decoration)) { ib_.begin(spv::OpExtension).add("SPV_KHR_no_integer_wrap_decoration").commit(&header_); } @@ -366,6 +393,21 @@ SType IRBuilder::get_pointer_type(const SType &value_type, spv::StorageClass sto t.element_type_id = value_type.id; t.storage_class = storage_class; ib_.begin(spv::OpTypePointer).add_seq(t, storage_class, value_type).commit(&global_); + // An `OpTypePointer` in the `PhysicalStorageBuffer` storage class that points to a scalar or vector + // needs an explicit `ArrayStride` decoration for `OpPtrAccessChain`'s `Element` offset to be scaled + // correctly on Vulkan. Without the decoration, drivers that strictly follow SPV_KHR_physical_storage_buffer + // treat the stride as undefined and collapse every element index to the base address, which manifests as + // `arr[i]` reads returning `arr[0]` for all `i` across the whole kernel (and any indexed ndarray write + // landing on slot 0). Sized-pointees (structs, arrays) already carry explicit layout decorations so they + // don't need this; the fix is limited to scalars/vectors, where the natural stride is just the pointee + // byte size. Uniform / StorageBuffer / Input / Output / Workgroup pointers don't use PSB arithmetic, so + // the decoration is a no-op for them and is skipped. + if (storage_class == spv::StorageClassPhysicalStorageBuffer && value_type.flag == TypeKind::kPrimitive) { + size_t stride = get_primitive_type_size(value_type.dt); + if (stride > 0) { + this->decorate(spv::OpDecorate, t, spv::DecorationArrayStride, uint32_t(stride)); + } + } pointer_type_tbl_[key] = t; return t; } diff --git a/quadrants/inc/rhi_constants.inc.h b/quadrants/inc/rhi_constants.inc.h index 725171d5db..bdd3d1795c 100644 --- a/quadrants/inc/rhi_constants.inc.h +++ b/quadrants/inc/rhi_constants.inc.h @@ -14,6 +14,17 @@ PER_DEVICE_CAPABILITY(spirv_has_int16) PER_DEVICE_CAPABILITY(spirv_has_int64) PER_DEVICE_CAPABILITY(spirv_has_float16) PER_DEVICE_CAPABILITY(spirv_has_float64) +// Vulkan `VkPhysicalDevice8BitStorageFeatures::storageBuffer8BitAccess` / +// `VkPhysicalDevice16BitStorageFeatures::storageBuffer16BitAccess` queried flags, enabling +// `CapabilityStorageBuffer{8,16}BitAccess` emission in the SPIR-V header. Needed independently of +// `spirv_has_int{8,16}` / `spirv_has_float16`: `Int8` / `Int16` / `Float16` capabilities gate +// whether the type can be declared, whereas these gate whether it can be loaded / stored through +// a descriptor-bound `StorageBuffer` pointer. Emitting the storage caps without the underlying +// device feature makes the driver reject the shader at pipeline creation on strict +// implementations; skipping them when the device does support them is silent-corruption-UB on +// drivers that enforce them. +PER_DEVICE_CAPABILITY(spirv_has_storage_buffer_8bit_access) +PER_DEVICE_CAPABILITY(spirv_has_storage_buffer_16bit_access) PER_DEVICE_CAPABILITY(spirv_has_atomic_int64) PER_DEVICE_CAPABILITY(spirv_has_atomic_float16) // load, store, exchange PER_DEVICE_CAPABILITY(spirv_has_atomic_float16_add) diff --git a/quadrants/rhi/vulkan/vulkan_device_creator.cpp b/quadrants/rhi/vulkan/vulkan_device_creator.cpp index a958c5e73e..f8246afe71 100644 --- a/quadrants/rhi/vulkan/vulkan_device_creator.cpp +++ b/quadrants/rhi/vulkan/vulkan_device_creator.cpp @@ -304,6 +304,27 @@ void VulkanDeviceCreator::create_instance(uint32_t vk_api_version, bool manual_c } } + // Normalize `params_.enable_validation_layer` against the host's actual layer availability BEFORE + // the instance-reuse short-circuit below: downstream code (device-extension enumeration at + // `create_logical_device`, the `VK_KHR_SHADER_NON_SEMANTIC_INFO` cap gate, the device layer arrays + // at the tail of `create_logical_device`) reads this flag on every cycle, and `create_logical_device` + // runs on every re-init. If we let the request pass through unchanged on the reuse path, the cap is + // set to `true` every subsequent cycle even when the layer is not actually loaded on the cached + // instance (the only cycle where the layer-load flip happens is the very first, which is also the + // only cycle where a fresh `vkCreateInstance` runs). Users observed this as `test_overflow.py` + // passing on the first parametrization and failing on every subsequent one in the same pytest + // session: `DebugPrintf`-ext-imported shaders compile fine, but the validation layer that would + // route those messages to stdout was never loaded, so the overflow-detected strings never appear in + // `capfd`. Running the check here keeps the flag consistent across re-inits; re-running on every + // call is cheap (it enumerates instance layers, ~microseconds). + if (params_.enable_validation_layer && !check_validation_layer_support()) { + RHI_LOG_ERROR( + "Validation layers requested but not available, turning off... " + "Please make sure Vulkan SDK from https://vulkan.lunarg.com/sdk/home " + "is installed."); + params_.enable_validation_layer = false; + } + // Reuse the VkInstance from a previous init/reset cycle if available. // Repeated vkDestroyInstance/vkCreateInstance triggers an NVIDIA driver bug // that corrupts SubgroupLocalInvocationId after ~11 cycles. @@ -325,15 +346,8 @@ void VulkanDeviceCreator::create_instance(uint32_t vk_api_version, bool manual_c create_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO; create_info.pApplicationInfo = &app_info; - if (params_.enable_validation_layer) { - if (!check_validation_layer_support()) { - RHI_LOG_ERROR( - "Validation layers requested but not available, turning off... " - "Please make sure Vulkan SDK from https://vulkan.lunarg.com/sdk/home " - "is installed."); - params_.enable_validation_layer = false; - } - } + // `params_.enable_validation_layer` has already been normalized against the host's actual layer + // availability above, before the cached-instance short-circuit. No second check needed here. VkDebugUtilsMessengerCreateInfoEXT debug_create_info{}; @@ -789,6 +803,15 @@ void VulkanDeviceCreator::create_logical_device(bool manual_create) { features2.pNext = &shader_8bit_storage_feature; vkGetPhysicalDeviceFeatures2KHR(physical_device_, &features2); + // Gate the SPIR-V `CapabilityStorageBuffer8BitAccess` emission strictly on the queried + // feature bit. `VK_KHR_8bit_storage` promoted into Vulkan 1.2 core doesn't imply the feature + // is actually supported -- implementations can still expose `storageBuffer8BitAccess = FALSE` + // -- so the SPIR-V cap may only be emitted when this feature is true, otherwise strict + // validation rejects shaders that declare the capability. + if (shader_8bit_storage_feature.storageBuffer8BitAccess) { + caps.set(DeviceCapability::spirv_has_storage_buffer_8bit_access, true); + } + *pNextEnd = &shader_8bit_storage_feature; pNextEnd = &shader_8bit_storage_feature.pNext; } @@ -796,6 +819,10 @@ void VulkanDeviceCreator::create_logical_device(bool manual_create) { features2.pNext = &shader_16bit_storage_feature; vkGetPhysicalDeviceFeatures2KHR(physical_device_, &features2); + if (shader_16bit_storage_feature.storageBuffer16BitAccess) { + caps.set(DeviceCapability::spirv_has_storage_buffer_16bit_access, true); + } + *pNextEnd = &shader_16bit_storage_feature; pNextEnd = &shader_16bit_storage_feature.pNext; }