[AMDGPU] Surface QuadrantsAssertionError after in-kernel assert (fix barrier hang)

quadrants/rhi/amdgpu/amdgpu_driver.cpp

@@ -6,6 +6,41 @@
 
 namespace quadrants {
 namespace lang {
+namespace {
+
+AmdgpuLaunchFailureHook g_amdgpu_launch_failure_hook = nullptr;
+bool g_amdgpu_device_assert_already_surfaced = false;
+bool g_amdgpu_device_in_teardown = false;
+
+}  // namespace
+
+void set_amdgpu_launch_failure_hook(AmdgpuLaunchFailureHook hook) {
+  g_amdgpu_launch_failure_hook = hook;
+}
+
+AmdgpuLaunchFailureHook get_amdgpu_launch_failure_hook() {
+  return g_amdgpu_launch_failure_hook;
+}
+
+bool amdgpu_device_assert_already_surfaced() {
+  return g_amdgpu_device_assert_already_surfaced;
+}
+
+void amdgpu_reset_device_assert_surfaced_flag() {
+  g_amdgpu_device_assert_already_surfaced = false;
+}
+
+void amdgpu_mark_device_assert_surfaced() {
+  g_amdgpu_device_assert_already_surfaced = true;
+}
+
+bool amdgpu_device_in_teardown() {
+  return g_amdgpu_device_in_teardown;
+}
+
+void amdgpu_set_device_in_teardown(bool in_teardown) {
+  g_amdgpu_device_in_teardown = in_teardown;
+}
 
 std::string get_amdgpu_error_message(uint32 err) {
   auto err_name_ptr = AMDGPUDriver::get_instance_without_context().get_error_name(err);

quadrants/rhi/amdgpu/amdgpu_driver.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <exception>
 #include <mutex>
 
 #include "quadrants/common/dynamic_loader.h"
@@ -37,10 +38,33 @@ constexpr uint32 HIP_DEVICE_MINOR = 332 / 4;
 // offsetof(hipDeviceProp_t, minor) / 4
 constexpr uint32 HIP_DEVICE_MINOR_6 = 364 / 4;
 constexpr uint32 HIP_ERROR_ASSERT = 710;
+// hipErrorLaunchFailure - returned by the first hipStreamSynchronize after a device
+// __builtin_trap() (AMDGPU in-kernel assert path). Catchable; does not abort the process.
+constexpr uint32 HIP_ERROR_LAUNCH_FAILURE = 719;
 constexpr uint32 HIP_JIT_MAX_REGISTERS = 0;
 constexpr uint32 HIP_POINTER_ATTRIBUTE_MEMORY_TYPE = 2;
 constexpr uint32 HIP_SUCCESS = 0;
 constexpr uint32 HIP_MEMORYTYPE_DEVICE = 1;
+// hipHostMallocCoherent - required for host-visible reads of assert state after a trap.
+constexpr uint32 HIP_HOST_MALLOC_COHERENT = 0x40000000;
+
+// Optional hook invoked from AMDGPUFunction::operator() on HIP_ERROR_LAUNCH_FAILURE before the
+// generic QD_ERROR. LlvmRuntimeExecutor registers this in debug+amdgpu mode to surface
+// QuadrantsAssertionError from pinned assert state. May throw.
+using AmdgpuLaunchFailureHook = void (*)();
+void set_amdgpu_launch_failure_hook(AmdgpuLaunchFailureHook hook);
+AmdgpuLaunchFailureHook get_amdgpu_launch_failure_hook();
+// True after a debug-mode device assert has been surfaced as QuadrantsAssertionError.
+// The HIP context is dead afterward, so subsequent calls also return launch failure.
+bool amdgpu_device_assert_already_surfaced();
+void amdgpu_reset_device_assert_surfaced_flag();
+void amdgpu_mark_device_assert_surfaced();
+// True while a Program/executor is tearing down. Dead-context launch failures are only
+// suppressed inside this window (or while unwinding the just-thrown assertion) so destructors
+// do not std::terminate(); outside it, post-assert launch failures are surfaced as hard errors
+// instead of silently reported as success on a dead context. Set by LlvmRuntimeExecutor.
+bool amdgpu_device_in_teardown();
+void amdgpu_set_device_in_teardown(bool in_teardown);
 // `hipFuncAttributeMaxDynamicSharedMemorySize` from the `hipFuncAttribute` enum in ROCm/clr
 // hipamd/include/hip/hip_runtime_api.h. Used with `kernel_set_attribute` (`hipFuncSetAttribute`) to opt in to >48 KB
 // of dynamic shared memory for graph kernel nodes that request it.
@@ -87,6 +111,28 @@ class AMDGPUFunction {
 
   void operator()(Args... args) {
     auto err = call(args...);
+    // Intercept launch failure before the generic HIP error so a debug-mode in-kernel assert
+    // can raise QuadrantsAssertionError from pinned host memory (context is dead after trap).
+    if (err == HIP_ERROR_LAUNCH_FAILURE) {
+      if (auto hook = get_amdgpu_launch_failure_hook()) {
+        hook();  // may throw QuadrantsAssertionError on first surfacing
+      }
+      if (amdgpu_device_assert_already_surfaced()) {
+        // The HIP context is dead after the trap, so every later call also returns launch
+        // failure. Suppress only where throwing would std::terminate(): inside teardown
+        // destructors, or while unwinding the just-thrown QuadrantsAssertionError. Anywhere
+        // else (e.g. user code that caught the assertion and kept issuing GPU work) fail
+        // loudly instead of returning stale/uninitialized results as success.
+        if (amdgpu_device_in_teardown() || std::uncaught_exceptions() > 0) {
+          return;
+        }
+        QD_ERROR(
+            "AMDGPU device context is unusable after an in-kernel assertion failure; "
+            "re-initialize Quadrants in a fresh process before issuing further GPU work "
+            "(while calling {} ({}))",
+            name_, symbol_name_);
+      }
+    }
     QD_ERROR_IF(err, get_error_message(err));
   }
 

quadrants/runtime/llvm/llvm_context.cpp

@@ -717,6 +717,28 @@ std::unique_ptr<llvm::Module> QuadrantsLLVMContext::module_from_file(const std::
       patch_fence("block_mem_fence", "workgroup");
       patch_fence("grid_mem_fence", "agent");
 
+      // System-scope fence for the AMDGPU in-kernel assert path: publish pinned assert state to the
+      // host before `__builtin_trap()`. Default LLVM syncscope is system (includes host memory).
+      {
+        auto func = module->getFunction("amdgpu_system_mem_fence");
+        if (func) {
+          func->deleteBody();
+          auto bb = llvm::BasicBlock::Create(*ctx, "entry", func);
+          IRBuilder<> builder(*ctx);
+          builder.SetInsertPoint(bb);
+          builder.CreateFence(llvm::AtomicOrdering::SequentiallyConsistent);
+          builder.CreateRetVoid();
+          QuadrantsLLVMContext::mark_inline(func);
+        } else {
+          // If the symbol name ever drifts, the runtime.cpp fallback (a host-atomic stub) is a no-op on device,
+          // which would silently break the publish-before-trap ordering the in-kernel assert path depends on.
+          // Warn loudly rather than fail so non-debug builds (which never trap) still run.
+          QD_WARN(
+              "amdgpu_system_mem_fence not found while patching the AMDGPU runtime module; the in-kernel assert "
+              "publish-before-trap fence will be a device no-op. QuadrantsAssertionError may not surface correctly.");
+        }
+      }
+
       link_module_with_amdgpu_libdevice(module);
       patch_amdgpu_kernel_dim("block_dim", llvm::ConstantInt::get(llvm::Type::getInt32Ty(*ctx), 0));
       patch_amdgpu_kernel_dim("grid_dim", llvm::ConstantInt::get(llvm::Type::getInt32Ty(*ctx), 0));

quadrants/runtime/llvm/llvm_runtime_executor.cpp

@@ -1,11 +1,14 @@
 #include "quadrants/runtime/llvm/llvm_runtime_executor.h"
 #include "quadrants/program/adstack_size_expr_eval.h"
 
+#include <atomic>
+#include <cstddef>
 #include <cstdint>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <limits>
+#include <string>
 #include <vector>
 
 #include "quadrants/ir/stmt_op_types.h"
@@ -32,6 +35,10 @@
 #include "quadrants/rhi/amdgpu/amdgpu_context.h"
 #endif
 
+#include "quadrants/common/exceptions.h"
+#include "quadrants/inc/constants.h"
+#include "quadrants/util/str.h"
+
 namespace quadrants::lang {
 namespace {
 void assert_failed_host(const char *msg) {
@@ -42,6 +49,54 @@ void *host_allocate_aligned(HostMemoryPool *memory_pool, std::size_t size, std::
   return memory_pool->allocate(size, alignment);
 }
 
+#if defined(QD_WITH_AMDGPU)
+// Layout-compatible with `AmdgpuAssertErrorState` in llvm_runtime.h (must stay in sync). The host TU cannot include
+// the device-runtime header, so both structs are pinned to the same canonical layout via the shared constants below;
+// a change to either that is not mirrored breaks one of these static_asserts.
+struct AmdgpuAssertErrorStateHostView {
+  int64_t error_code;
+  char error_message_template[quadrants_error_message_max_length];
+  uint64_t error_message_arguments[quadrants_error_message_max_num_arguments];
+};
+
+static_assert(offsetof(AmdgpuAssertErrorStateHostView, error_code) == 0, "AmdgpuAssertErrorStateHostView layout drift");
+static_assert(offsetof(AmdgpuAssertErrorStateHostView, error_message_template) == sizeof(int64_t),
+              "AmdgpuAssertErrorStateHostView layout drift");
+static_assert(offsetof(AmdgpuAssertErrorStateHostView, error_message_arguments) ==
+                  sizeof(int64_t) + quadrants_error_message_max_length,
+              "AmdgpuAssertErrorStateHostView layout drift");
+static_assert(sizeof(AmdgpuAssertErrorStateHostView) ==
+                  sizeof(int64_t) + quadrants_error_message_max_length +
+                      quadrants_error_message_max_num_arguments * sizeof(uint64_t),
+              "AmdgpuAssertErrorStateHostView layout drift");
+
+// Host pointer published for the AMDGPU launch-failure hook (see amdgpu_driver.h). Only set while a
+// debug+amdgpu LlvmRuntimeExecutor owns a live pinned assert-error state.
+AmdgpuAssertErrorStateHostView *g_amdgpu_assert_error_state_host = nullptr;
+
+void amdgpu_launch_failure_assert_hook() {
+  auto *st = g_amdgpu_assert_error_state_host;
+  if (st == nullptr) {
+    return;
+  }
+  std::atomic_thread_fence(std::memory_order_acquire);
+  const int64_t code = __atomic_load_n(&st->error_code, __ATOMIC_SEQ_CST);
+  if (code != 1) {
+    return;
+  }
+  // Consume so Program teardown / subsequent HIP calls on the dead context do not re-throw.
+  __atomic_store_n(&st->error_code, (int64_t)0, __ATOMIC_SEQ_CST);
+  amdgpu_mark_device_assert_surfaced();
+  // Context is dead after the trap; format solely from the pinned buffer (no device retrieval).
+  std::string error_message_template(st->error_message_template);
+  const auto error_message_formatted =
+      format_error_message(error_message_template, [st](int argument_id) -> uint64 {
+        return st->error_message_arguments[argument_id];
+      });
+  throw QuadrantsAssertionError(error_message_formatted);
+}
+#endif
+
 }  // namespace
 
 LlvmRuntimeExecutor::LlvmRuntimeExecutor(CompileConfig &config, KernelProfilerBase *profiler, ProgramImpl *program_impl)
@@ -493,6 +548,15 @@ uint64_t *LlvmRuntimeExecutor::get_device_alloc_info_ptr(const DeviceAllocation
 
 void LlvmRuntimeExecutor::finalize() {
   profiler_ = nullptr;
+#if defined(QD_WITH_AMDGPU)
+  // Entering teardown: from here the driver may issue calls on a HIP context that a prior
+  // in-kernel assert left dead. Allow those launch failures to be swallowed (see
+  // AMDGPUFunction::operator()) so destructors do not std::terminate(); outside teardown a
+  // post-assert launch failure is still surfaced as a hard error.
+  if (config_.arch == Arch::amdgpu) {
+    amdgpu_set_device_in_teardown(true);
+  }
+#endif
   // Release the host-owned adstack heap before the device teardown below so its `DeviceAllocationGuard` destructor
   // runs while the RHI device is still valid. The destructor drops the allocation back to the driver memory pool
   // (or to the host allocator on CPU); deferring past `llvm_device()->clear()` would leak it.
@@ -589,6 +653,18 @@ void LlvmRuntimeExecutor::finalize() {
     adstack_overflow_task_id_host_ptr_ = nullptr;
     adstack_overflow_task_id_dev_ptr_ = nullptr;
   }
+#if defined(QD_WITH_AMDGPU)
+  if (assert_error_state_host_ptr_ != nullptr) {
+    if (g_amdgpu_assert_error_state_host == assert_error_state_host_ptr_) {
+      g_amdgpu_assert_error_state_host = nullptr;
+      set_amdgpu_launch_failure_hook(nullptr);
+    }
+    if (config_.arch == Arch::amdgpu) {
+      AMDGPUDriver::get_instance().mem_free_host(assert_error_state_host_ptr_);
+    }
+    assert_error_state_host_ptr_ = nullptr;
+  }
+#endif
   if (config_.arch == Arch::cuda || config_.arch == Arch::amdgpu) {
     preallocated_runtime_objects_allocs_.reset();
     preallocated_runtime_memory_allocs_.reset();
@@ -884,6 +960,28 @@ void LlvmRuntimeExecutor::materialize_runtime(KernelProfilerBase *profiler, uint
     runtime_jit->call<void *, void *>("runtime_set_adstack_overflow_task_id_dev_ptr", llvm_runtime_,
                                       adstack_overflow_task_id_dev_ptr_);
   }
+
+  // AMDGPU debug assert: allocate pinned coherent host memory for assert state so the host can
+  // format QuadrantsAssertionError after `__builtin_trap()` kills the dispatch (HIP context dead;
+  // device retrieval kernels cannot run). Gated on debug + amdgpu only.
+  if (config_.debug && config_.arch == Arch::amdgpu) {
+#if defined(QD_WITH_AMDGPU)
+    void *host_slot = nullptr;
+    AMDGPUDriver::get_instance().mem_alloc_host(&host_slot, sizeof(AmdgpuAssertErrorStateHostView),
+                                                HIP_HOST_MALLOC_COHERENT);
+    QD_ASSERT(host_slot != nullptr);
+    std::memset(host_slot, 0, sizeof(AmdgpuAssertErrorStateHostView));
+    assert_error_state_host_ptr_ = host_slot;
+    g_amdgpu_assert_error_state_host = static_cast<AmdgpuAssertErrorStateHostView *>(host_slot);
+    amdgpu_reset_device_assert_surfaced_flag();
+    amdgpu_set_device_in_teardown(false);
+    set_amdgpu_launch_failure_hook(amdgpu_launch_failure_assert_hook);
+    // UVA: host pointer is also a valid device pointer on GFX9+.
+    runtime_jit->call<void *, void *>("runtime_set_assert_error_state_dev_ptr", llvm_runtime_, host_slot);
+#else
+    QD_NOT_IMPLEMENTED;
+#endif
+  }
 }
 
 void LlvmRuntimeExecutor::destroy_snode_tree(SNodeTree *snode_tree) {

quadrants/runtime/llvm/llvm_runtime_executor.h

@@ -303,6 +303,10 @@ class LlvmRuntimeExecutor {
   int64_t *adstack_overflow_task_id_host_ptr_{nullptr};
   void *adstack_overflow_task_id_dev_ptr_{nullptr};
 
+  // AMDGPU debug-only: pinned host mirror of in-kernel assert state (`AmdgpuAssertErrorState` in
+  // llvm_runtime.h). Allocated in materialize_runtime when config_.debug && arch==amdgpu.
+  void *assert_error_state_host_ptr_{nullptr};
+
   std::unique_ptr<ThreadPool> thread_pool_{nullptr};
   std::shared_ptr<Device> device_{nullptr};
 

quadrants/runtime/llvm/runtime_module/adstack_runtime.cpp

@@ -787,6 +787,7 @@ void adstack_runtime_zero_init(LLVMRuntime *runtime) {
   runtime->adstack_bound_row_capacities_capacity = 0;
   runtime->adstack_overflow_flag_dev_ptr = nullptr;
   runtime->adstack_overflow_task_id_dev_ptr = nullptr;
+  runtime->assert_error_state_dev_ptr = nullptr;
 }
 
 extern "C" {  // local stack operations

quadrants/runtime/llvm/runtime_module/llvm_runtime.h

@@ -103,6 +103,26 @@ struct PreallocatedMemoryChunk {
   std::size_t preallocated_size = 0;
 };
 
+// AMDGPU debug-only pinned assert mirror. Host allocates via hipHostMalloc(Coherent); device publishes
+// here before `__builtin_trap()` so the host can format QuadrantsAssertionError after the HIP context dies.
+// See `LLVMRuntime::assert_error_state_dev_ptr`.
+struct AmdgpuAssertErrorState {
+  i64 error_code;
+  char error_message_template[quadrants_error_message_max_length];
+  uint64 error_message_arguments[quadrants_error_message_max_num_arguments];
+};
+
+// Keep in exact layout sync with `AmdgpuAssertErrorStateHostView` in llvm_runtime_executor.cpp, which is hand-mirrored
+// (the host TU cannot include this device-runtime header). Both structs are pinned to this canonical layout expressed
+// via the shared constants, so neither can drift without breaking its own static_assert.
+static_assert(offsetof(AmdgpuAssertErrorState, error_code) == 0, "AmdgpuAssertErrorState layout drift");
+static_assert(offsetof(AmdgpuAssertErrorState, error_message_template) == sizeof(i64), "AmdgpuAssertErrorState layout drift");
+static_assert(offsetof(AmdgpuAssertErrorState, error_message_arguments) == sizeof(i64) + quadrants_error_message_max_length,
+              "AmdgpuAssertErrorState layout drift");
+static_assert(sizeof(AmdgpuAssertErrorState) ==
+                  sizeof(i64) + quadrants_error_message_max_length + quadrants_error_message_max_num_arguments * sizeof(uint64),
+              "AmdgpuAssertErrorState layout drift");
+
 struct LLVMRuntime {
   PreallocatedMemoryChunk runtime_objects_chunk;
   PreallocatedMemoryChunk runtime_memory_chunk;
@@ -143,6 +163,7 @@ struct LLVMRuntime {
   uint64 error_message_arguments[quadrants_error_message_max_num_arguments];
   i32 error_message_lock = 0;
   i64 error_code = 0;
+
   // Dedicated overflow signal. Pointer to a 64-bit slot in pinned host memory (CUDA `cuMemAllocHost_v2`,
   // HIP `hipHostMalloc`; CPU plain malloc; on this struct stored as the device-mapped address obtained via
   // `cuMemHostGetDevicePointer` / HIP equivalent). The kernel-side `stack_push` writes via a system-wide
@@ -226,6 +247,13 @@ struct LLVMRuntime {
 
   i64 total_requested_memory;
 
+  // AMDGPU debug-only: device-mapped pointer to pinned `AmdgpuAssertErrorState` (see struct doc above).
+  // nullptr on non-AMDGPU / non-debug; nullptr-guarded in `quadrants_assert_format`. Appended at the end of
+  // the struct on purpose: the default-on offline cache keys on the numeric Quadrants version (not the runtime
+  // layout), so inserting a field mid-struct would shift every following field's offset and let an old cached
+  // kernel misread them. Keeping new fields at the tail preserves existing offsets for cached kernels.
+  AmdgpuAssertErrorState *assert_error_state_dev_ptr = nullptr;
+
   template <typename T>
   void set_result(std::size_t i, T t) {
     static_assert(sizeof(T) <= sizeof(uint64));