[SPIRV] Vulkan SPIR-V correctness: atomic-view aliasing, PSB stride, narrow storage caps, u1 cast, per-init layer recheck

quadrants/codegen/spirv/detail/spirv_codegen.h

@@ -17,6 +17,8 @@
 #include <spirv-tools/libspirv.hpp>
 #include <spirv-tools/optimizer.hpp>
 
+#include <unordered_set>
+
 namespace quadrants::lang {
 namespace spirv {
 namespace detail {
@@ -154,6 +156,14 @@ class TaskCodegen : public IRVisitor {
   std::shared_ptr<spirv::IRBuilder> ir_;  // spirv binary code builder
   std::unordered_map<std::pair<BufferInfo, int>, spirv::Value, BufferInfoTypeTupleHasher> buffer_value_map_;
   std::unordered_map<std::pair<BufferInfo, int>, 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<BufferInfo, std::vector<spirv::Value>, BufferInfoHasher> buffer_views_by_buffer_;
+  std::unordered_set<uint32_t> aliased_decorated_buffer_ids_;
   std::vector<spirv::Value> shared_array_binds_;
   spirv::Value kernel_function_;
   spirv::Label kernel_return_label_;

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;

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;
 }

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)

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,13 +803,26 @@ 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;
     }
     if (CHECK_VERSION(1, 1) || CHECK_EXTENSION(VK_KHR_16BIT_STORAGE_EXTENSION_NAME)) {
       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;
     }