Implement SPV_KHR_abort extension

include/LLVMSPIRVExtensions.inc

@@ -6,6 +6,7 @@ EXT(SPV_EXT_shader_atomic_float16_add)
 EXT(SPV_EXT_shader_atomic_float_add)
 EXT(SPV_EXT_shader_atomic_float_min_max)
 EXT(SPV_EXT_image_raw10_raw12)
+EXT(SPV_KHR_abort)
 EXT(SPV_KHR_no_integer_wrap_decoration)
 EXT(SPV_KHR_float_controls)
 EXT(SPV_KHR_linkonce_odr)

lib/SPIRV/SPIRVReader.cpp

@@ -1951,6 +1951,21 @@ Value *SPIRVToLLVM::transValueWithoutDecoration(SPIRVValue *BV, Function *F,
                                transValue(RV->getReturnValue(), F, BB), BB));
   }
 
+  case OpAbortKHR: {
+    // OpAbortKHR is a SPIR-V block terminator. In LLVM IR, model it as a call
+    // to the SPIR-V friendly builtin __spirv_AbortKHR followed by an
+    // 'unreachable' terminator.
+    auto *AbortInst =
+        transSPIRVBuiltinFromInst(static_cast<SPIRVAbortKHR *>(BV), BB);
+    if (auto *Call = dyn_cast<CallInst>(AbortInst)) {
+      Call->setDoesNotReturn();
+      if (auto *Callee = Call->getCalledFunction())
+        Callee->setDoesNotReturn();
+    }
+    new UnreachableInst(*Context, BB);
+    return mapValue(BV, AbortInst);
+  }
+
   case OpLifetimeStart: {
     SPIRVLifetimeStart *LTStart = static_cast<SPIRVLifetimeStart *>(BV);
     IRBuilder<> Builder(BB);

lib/SPIRV/SPIRVWriter.cpp

@@ -2343,7 +2343,7 @@ LLVMToSPIRVBase::transValueWithoutDecoration(Value *V, SPIRVBasicBlock *BB,
     return mapValue(V, BI);
   }
 
-  if (dyn_cast<UnreachableInst>(V))
+  if (isa<UnreachableInst>(V))
     return mapValue(V, BM->addUnreachableInst(BB));
 
   if (auto *RI = dyn_cast<ReturnInst>(V)) {
@@ -5141,9 +5141,10 @@ SPIRVValue *LLVMToSPIRVBase::transIntrinsicInst(IntrinsicInst *II,
   case Intrinsic::invariant_start:
   case Intrinsic::invariant_end:
   case Intrinsic::dbg_label:
-  // llvm.trap and llvm.debugtrap intrinsics are not implemented. But for now
-  // don't crash. This change is pending the trap/abort intrinsic
-  // implementation.
+  // TODO: lower llvm.trap / llvm.ubsantrap / llvm.debugtrap to OpAbortKHR
+  // (with a null/undefined Message) when the SPV_KHR_abort extension is
+  // enabled. For now we silently drop these intrinsics so that the translator
+  // doesn't crash on input that contains them.
   case Intrinsic::trap:
   case Intrinsic::ubsantrap:
   case Intrinsic::debugtrap:
@@ -6337,6 +6338,15 @@ void LLVMToSPIRVBase::transFunction(Function *I) {
     SPIRVBasicBlock *BB =
         static_cast<SPIRVBasicBlock *>(transValue(&FI, nullptr));
     for (auto &BI : FI) {
+      // SPV_KHR_abort: OpAbortKHR is itself a SPIR-V block terminator. Once a
+      // basic block has been terminated by OpAbortKHR, any subsequent LLVM IR
+      // instructions in the same block (typically lifetime intrinsics, a
+      // trailing `unreachable`, or a `ret`) must not be emitted, otherwise
+      // the resulting SPIR-V would have instructions after a block terminator
+      // and fail validation.
+      if (auto *Last = BB->getTerminateInstr();
+          Last && Last->getOpCode() == OpAbortKHR)
+        break;
       transValue(&BI, BB, false);
     }
   }
@@ -6910,6 +6920,16 @@ LLVMToSPIRVBase::transBuiltinToInstWithoutDecoration(Op OC, CallInst *CI,
     auto BArgs = transValue(getArguments(CI), BB);
     return BM->addControlBarrierInst(BArgs[0], BArgs[1], BArgs[2], BB);
   } break;
+  case OpAbortKHR: {
+    if (!BM->checkExtension(ExtensionID::SPV_KHR_abort,
+                            SPIRVEC_RequiresExtension, "SPV_KHR_abort\n"))
+      return nullptr;
+    BM->getErrorLog().checkError(
+        CI->arg_size() == 1, SPIRVEC_InvalidInstruction,
+        "__spirv_AbortKHR must be called with exactly one Message argument\n");
+    auto *Msg = transValue(CI->getArgOperand(0), BB);
+    return BM->addAbortKHRInst(Msg, BB);
+  } break;
   case OpGroupAsyncCopy: {
     auto BArgs = transValue(getArguments(CI), BB);
     return BM->addAsyncGroupCopy(BArgs[0], BArgs[1], BArgs[2], BArgs[3],
@@ -7600,7 +7620,8 @@ bool runSpirvBackend(Module *M, std::string &Result, std::string &ErrMsg,
       SPIRV::ExtensionID::SPV_INTEL_function_pointers,
       SPIRV::ExtensionID::SPV_KHR_shader_clock,
       SPIRV::ExtensionID::SPV_KHR_cooperative_matrix,
-      SPIRV::ExtensionID::SPV_KHR_non_semantic_info};
+      SPIRV::ExtensionID::SPV_KHR_non_semantic_info,
+      SPIRV::ExtensionID::SPV_KHR_abort};
   // The fallback for the Triple value.
   static const std::string DefaultTriple = "spirv64-unknown-unknown";
 

lib/SPIRV/libSPIRV/SPIRVEntry.cpp

@@ -171,6 +171,7 @@ bool SPIRVEntry::isEndOfBlock() const {
   case OpReturn:
   case OpReturnValue:
   case OpUnreachable:
+  case OpAbortKHR:
     return true;
   default:
     return false;

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -858,6 +858,50 @@ template <Op TheOpCode> class SPIRVInstNoOperand : public SPIRVInstruction {
 typedef SPIRVInstNoOperand<OpReturn> SPIRVReturn;
 typedef SPIRVInstNoOperand<OpUnreachable> SPIRVUnreachable;
 
+class SPIRVAbortKHR : public SPIRVInstruction {
+public:
+  static const Op OC = OpAbortKHR;
+  static const SPIRVWord FixedWordCount = 3;
+  // Complete constructor.
+  SPIRVAbortKHR(SPIRVValue *TheMessage, SPIRVBasicBlock *TheBB)
+      : SPIRVInstruction(FixedWordCount, OC, TheBB),
+        MessageTypeId(TheMessage->getType()->getId()),
+        MessageId(TheMessage->getId()) {
+    setAttr();
+    validate();
+    assert(TheBB && "Invalid BB");
+  }
+  // Incomplete constructor.
+  SPIRVAbortKHR()
+      : SPIRVInstruction(OC), MessageTypeId(SPIRVID_INVALID),
+        MessageId(SPIRVID_INVALID) {
+    setAttr();
+  }
+
+  SPIRVCapVec getRequiredCapability() const override {
+    return getVec(CapabilityAbortKHR);
+  }
+
+  std::optional<ExtensionID> getRequiredExtension() const override {
+    return ExtensionID::SPV_KHR_abort;
+  }
+
+  std::vector<SPIRVValue *> getOperands() override {
+    return std::vector<SPIRVValue *>(1, getValue(MessageId));
+  }
+
+  _SPIRV_DEF_ENCDEC2(MessageTypeId, MessageId)
+
+protected:
+  void setAttr() {
+    setHasNoId();
+    setHasNoType();
+  }
+  void validate() const override { SPIRVInstruction::validate(); }
+  SPIRVId MessageTypeId;
+  SPIRVId MessageId;
+};
+
 class SPIRVReturnValue : public SPIRVInstruction {
 public:
   static const Op OC = OpReturnValue;

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -549,6 +549,8 @@ class SPIRVModuleImpl : public SPIRVModule {
                                                   llvm::MDNode *MD) override;
   SPIRVInstruction *addAssumeTrueKHRInst(SPIRVValue *Condition,
                                          SPIRVBasicBlock *BB) override;
+  SPIRVInstruction *addAbortKHRInst(SPIRVValue *Message,
+                                    SPIRVBasicBlock *BB) override;
   SPIRVInstruction *addExpectKHRInst(SPIRVType *ResultTy, SPIRVValue *Value,
                                      SPIRVValue *ExpectedValue,
                                      SPIRVBasicBlock *BB) override;
@@ -2031,6 +2033,11 @@ SPIRVInstruction *SPIRVModuleImpl::addAssumeTrueKHRInst(SPIRVValue *Condition,
   return addInstruction(new SPIRVAssumeTrueKHR(Condition->getId(), BB), BB);
 }
 
+SPIRVInstruction *SPIRVModuleImpl::addAbortKHRInst(SPIRVValue *Message,
+                                                   SPIRVBasicBlock *BB) {
+  return addInstruction(new SPIRVAbortKHR(Message, BB), BB);
+}
+
 SPIRVInstruction *SPIRVModuleImpl::addExpectKHRInst(SPIRVType *ResultTy,
                                                     SPIRVValue *Value,
                                                     SPIRVValue *ExpectedValue,

lib/SPIRV/libSPIRV/SPIRVModule.h

@@ -529,6 +529,8 @@ class SPIRVModule {
                                       llvm::MDNode *MD) = 0;
   virtual SPIRVInstruction *addAssumeTrueKHRInst(SPIRVValue *Condition,
                                                  SPIRVBasicBlock *BB) = 0;
+  virtual SPIRVInstruction *addAbortKHRInst(SPIRVValue *Message,
+                                            SPIRVBasicBlock *BB) = 0;
   virtual SPIRVInstruction *addExpectKHRInst(SPIRVType *ResultTy,
                                              SPIRVValue *Value,
                                              SPIRVValue *ExpectedValue,

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -660,6 +660,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
   add(CapabilityTernaryBitwiseFunctionINTEL, "TernaryBitwiseFunctionINTEL");
   add(CapabilityFMAKHR, "FMAKHR");
   add(CapabilityPredicatedIOINTEL, "PredicatedIOINTEL");
+  add(CapabilityAbortKHR, "AbortKHR");
   // From spirv_internal.hpp
   add(internal::CapabilityTokenTypeINTEL, "TokenTypeINTEL");
   add(internal::CapabilityHWThreadQueryINTEL, "HWThreadQueryINTEL");

lib/SPIRV/libSPIRV/SPIRVOpCodeEnum.h

@@ -355,6 +355,7 @@ _SPIRV_OP(CooperativeMatrixStoreKHR, 4458)
 _SPIRV_OP(CooperativeMatrixMulAddKHR, 4459)
 _SPIRV_OP(CooperativeMatrixLengthKHR, 4460)
 _SPIRV_OP(ReadClockKHR, 5056)
+_SPIRV_OP(AbortKHR, 5121)
 _SPIRV_OP(SubgroupShuffleINTEL, 5571)
 _SPIRV_OP(SubgroupShuffleDownINTEL, 5572)
 _SPIRV_OP(SubgroupShuffleUpINTEL, 5573)

test/extensions/KHR/SPV_KHR_abort/abort-conditional.ll

@@ -0,0 +1,67 @@
+; Conditional abort — the assert() pattern where only some paths abort.
+; Verifies that:
+;   1. The abort BB gets OpAbortKHR (not OpUnreachable)
+;   2. The non-abort BB is unaffected (still has OpReturn)
+;   3. No double terminators in the abort block.
+
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_KHR_abort -o %t.spv
+; RUN: llvm-spirv %t.spv -to-text -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+
+; Round-trip
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis < %t.rev.bc | FileCheck %s --check-prefix=CHECK-LLVM
+
+; FIXME: enable the following run when the translator CI is updated to a new
+; version of the SPIR-V Tools that includes the support for the SPV_KHR_abort
+; extension.
+; RUN: not spirv-val %t.spv
+
+; ---- SPIR-V with extension ----
+; CHECK-SPIRV-DAG: Extension "SPV_KHR_abort"
+; CHECK-SPIRV-DAG: Capability AbortKHR
+; CHECK-SPIRV: Function
+;
+; Entry block: branch conditional
+; CHECK-SPIRV: BranchConditional
+;
+; OK block: normal return
+; CHECK-SPIRV: Return
+;
+; Trap block: abort only, no Unreachable / Return after it
+; CHECK-SPIRV: AbortKHR
+; CHECK-SPIRV-EMPTY:
+; CHECK-SPIRV-NEXT: FunctionEnd
+
+; ---- Round-trip LLVM IR ----
+; CHECK-LLVM: define spir_func void @assert_like
+; CHECK-LLVM: br i1
+; CHECK-LLVM: ret void
+; CHECK-LLVM: call spir_func void @{{.*__spirv_AbortKHR.*}}(i32 %{{.*}}){{.*}}#[[#ATTR:]]
+; CHECK-LLVM-NEXT: unreachable
+; CHECK-LLVM: attributes #[[#ATTR]] = {{{.*}}noreturn{{.*}}}
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spir64-unknown-unknown"
+
+define spir_func void @assert_like(i32 %gid, i32 %N, i32 %msg) {
+entry:
+  %cmp = icmp slt i32 %gid, %N
+  br i1 %cmp, label %ok, label %trap
+
+ok:
+  ret void
+
+trap:
+  call spir_func void @_Z16__spirv_AbortKHRIiEvT_(i32 %msg)
+  unreachable
+}
+
+declare spir_func void @_Z16__spirv_AbortKHRIiEvT_(i32)
+
+!opencl.spir.version = !{!0}
+!spirv.Source = !{!1}
+
+!0 = !{i32 1, i32 2}
+!1 = !{i32 4, i32 100000}

test/extensions/KHR/SPV_KHR_abort/abort-in-kernel.ll

@@ -0,0 +1,93 @@
+; Kernel function with assert-like abort pattern. Models the real-world HIP
+; assert() use case: kernel calls a helper which calls __spirv_AbortKHR.
+;
+; Verifies that OpAbortKHR works correctly inside callee functions invoked
+; from kernel entry points.
+;
+; Note: the kernel's assert.fail block has FunctionCall + Unreachable
+; (the unreachable is after the call, not after an abort — this is correct
+; because the abort is inside the callee, not the caller).
+
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_KHR_abort -o %t.spv
+; RUN: llvm-spirv %t.spv -to-text -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+
+; FIXME: enable the following run when the translator CI is updated to a new
+; version of the SPIR-V Tools that includes the support for the SPV_KHR_abort
+; extension.
+; RUN: not spirv-val %t.spv
+
+; Round-trip
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis < %t.rev.bc | FileCheck %s --check-prefix=CHECK-LLVM
+
+; ---- SPIR-V ----
+; CHECK-SPIRV-DAG: Capability AbortKHR
+; CHECK-SPIRV-DAG: Extension "SPV_KHR_abort"
+; CHECK-SPIRV-DAG: EntryPoint 6 [[#KernelId:]] "test_kernel"
+
+; __assert_fail_internal: contains the abort -> OpAbortKHR
+; CHECK-SPIRV: Function
+; CHECK-SPIRV: AbortKHR
+; CHECK-SPIRV: FunctionEnd
+
+; test_kernel: calls __assert_fail_internal, then unreachable (in caller)
+; CHECK-SPIRV: Function {{.*}} [[#KernelId]]
+; CHECK-SPIRV: BranchConditional
+; CHECK-SPIRV: Return
+; CHECK-SPIRV: FunctionCall
+; CHECK-SPIRV: Unreachable
+; CHECK-SPIRV: FunctionEnd
+
+; ---- Round-trip ----
+; CHECK-LLVM: define spir_func void @__assert_fail_internal
+; CHECK-LLVM: call spir_func void @{{.*__spirv_AbortKHR.*}}(i32 {{.*}})
+; CHECK-LLVM-NEXT: unreachable
+;
+; CHECK-LLVM: define spir_kernel void @test_kernel
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spir64-unknown-unknown"
+
+; Models __assert_fail from device libraries
+define spir_func void @__assert_fail_internal(i32 %msg) #0 {
+entry:
+  call spir_func void @_Z16__spirv_AbortKHRIiEvT_(i32 %msg)
+  unreachable
+}
+
+; Kernel entry point with conditional assert
+define spir_kernel void @test_kernel(ptr addrspace(1) %in, i32 %N) #1
+  !kernel_arg_addr_space !1 !kernel_arg_access_qual !2
+  !kernel_arg_type !3 !kernel_arg_base_type !3 !kernel_arg_type_qual !4 {
+entry:
+  %gid = call spir_func i64 @_Z13get_global_idj(i32 0)
+  %gid32 = trunc i64 %gid to i32
+  %cmp = icmp slt i32 %gid32, %N
+  br i1 %cmp, label %ok, label %assert.fail
+
+ok:
+  ret void
+
+assert.fail:
+  call spir_func void @__assert_fail_internal(i32 42)
+  unreachable
+}
+
+declare spir_func i64 @_Z13get_global_idj(i32) #2
+declare spir_func void @_Z16__spirv_AbortKHRIiEvT_(i32)
+
+attributes #0 = { noinline noreturn nounwind }
+attributes #1 = { nounwind }
+attributes #2 = { nounwind }
+
+!opencl.spir.version = !{!0}
+!spirv.Source = !{!5}
+
+!0 = !{i32 1, i32 2}
+!1 = !{i32 1, i32 0}
+!2 = !{!"none", !"none"}
+!3 = !{!"int*", !"int"}
+!4 = !{!"", !""}
+!5 = !{i32 4, i32 100000}

test/extensions/KHR/SPV_KHR_abort/abort-invalid-arg-count-no-args.ll

@@ -0,0 +1,18 @@
+; Negative test: OpAbortKHR takes exactly one Message operand; a call to
+; `__spirv_AbortKHR` with no arguments must be rejected by the translator.
+
+; RUN: llvm-as %s -o %t.bc
+; RUN: not llvm-spirv %t.bc --spirv-ext=+SPV_KHR_abort -o /dev/null 2>&1 | FileCheck %s
+
+; CHECK: InvalidInstruction
+; CHECK: __spirv_AbortKHR must be called with exactly one Message argument
+
+target triple = "spir64-unknown-unknown"
+
+define spir_func void @abort_no_args() {
+entry:
+  call spir_func void @_Z16__spirv_AbortKHRv()
+  ret void
+}
+
+declare spir_func void @_Z16__spirv_AbortKHRv()