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

@@ -1939,6 +1939,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<SPIRVInstruction *>(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

@@ -2333,10 +2333,27 @@ LLVMToSPIRVBase::transValueWithoutDecoration(Value *V, SPIRVBasicBlock *BB,
     return mapValue(V, BI);
   }
 
-  if (dyn_cast<UnreachableInst>(V))
+  if (isa<UnreachableInst>(V)) {
+    // SPV_KHR_abort: OpAbortKHR is itself a block terminator. If the LLVM IR
+    // appends a trailing 'unreachable' after the abort call (the canonical
+    // pattern for noreturn calls), do not emit a second SPIR-V terminator.
+    if (auto *Last = BB->getTerminateInstr())
+      if (Last->getOpCode() == OpAbortKHR)
+        return mapValue(V, const_cast<SPIRVInstruction *>(Last));
     return mapValue(V, BM->addUnreachableInst(BB));
+  }
 
   if (auto *RI = dyn_cast<ReturnInst>(V)) {
+    // SPV_KHR_abort: OpAbortKHR is itself a SPIR-V block terminator. If the
+    // LLVM IR appends a trailing `ret void` after the abort call, do not emit
+    // a second SPIR-V terminator. We deliberately limit this suppression to
+    // `ret void`: if a non-void function ends with `ret <value>` after
+    // OpAbortKHR, fall through so that the value is still translated and the
+    // user sees a normal SPIR-V validation error rather than silently
+    // dropping the return value.
+    if (auto *Last = BB->getTerminateInstr();
+        !RI->getReturnValue() && Last && Last->getOpCode() == OpAbortKHR)
+      return mapValue(V, const_cast<SPIRVInstruction *>(Last));
     if (auto *RV = RI->getReturnValue()) {
       if (auto *II = dyn_cast<IntrinsicInst>(RV)) {
         if (II->getIntrinsicID() == Intrinsic::frexp) {
@@ -5131,9 +5148,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:
@@ -6902,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: {
+    BM->getErrorLog().checkError(
+        BM->isAllowedToUseExtension(ExtensionID::SPV_KHR_abort),
+        SPIRVEC_RequiresExtension, "SPV_KHR_abort\n");
+    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],
@@ -7592,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,49 @@ 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));
+  }
+
+protected:
+  void setAttr() {
+    setHasNoId();
+    setHasNoType();
+  }
+  _SPIRV_DEF_ENCDEC2(MessageTypeId, MessageId)
+  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

@@ -658,6 +658,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
       "SubgroupMatrixMultiplyAccumulateINTEL");
   add(CapabilityTernaryBitwiseFunctionINTEL, "TernaryBitwiseFunctionINTEL");
   add(CapabilityFMAKHR, "FMAKHR");
+  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)

spirv-headers-tag.conf

@@ -1 +1 @@
-9268f3057354a2cb65991ba5f38b16d81e803692
+ad9184e76a66b1001c29db9b0a3e87f646c64de0

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

@@ -0,0 +1,37 @@
+; Negative tests for the SPIR-V translator's handling of `__spirv_AbortKHR`.
+; OpAbortKHR takes exactly one Message operand; calls with any other arity must
+; be rejected by the translator with a clear diagnostic rather than producing
+; ill-formed SPIR-V.
+
+; RUN: split-file %s %t
+
+; RUN: llvm-as %t/no-args.ll -o %t/no-args.bc
+; RUN: not llvm-spirv %t/no-args.bc --spirv-ext=+SPV_KHR_abort -o /dev/null 2>&1 | FileCheck %s --check-prefix=CHECK-ARG-COUNT
+
+; RUN: llvm-as %t/two-args.ll -o %t/two-args.bc
+; RUN: not llvm-spirv %t/two-args.bc --spirv-ext=+SPV_KHR_abort -o /dev/null 2>&1 | FileCheck %s --check-prefix=CHECK-ARG-COUNT
+
+; CHECK-ARG-COUNT: InvalidInstruction
+; CHECK-ARG-COUNT: __spirv_AbortKHR must be called with exactly one Message argument
+
+;--- no-args.ll
+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()
+
+;--- two-args.ll
+target triple = "spir64-unknown-unknown"
+
+define spir_func void @abort_two_args(i32 %a, i32 %b) {
+entry:
+  call spir_func void @_Z16__spirv_AbortKHRii(i32 %a, i32 %b)
+  ret void
+}
+
+declare spir_func void @_Z16__spirv_AbortKHRii(i32, i32)

test/extensions/KHR/SPV_KHR_abort/abort.ll

@@ -0,0 +1,48 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc -o %t.spv --spirv-ext=+SPV_KHR_abort
+; RUN: llvm-spirv %t.spv -o %t.spt --to-text
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+; RUN: llvm-spirv %t.spv -o %t.rev.bc -r --spirv-target-env=SPV-IR
+; RUN: llvm-dis %t.rev.bc -o %t.rev.ll
+; RUN: FileCheck < %t.rev.ll %s --check-prefix=CHECK-LLVM
+; RUN: spirv-val %t.spv
+
+; RUN: not llvm-spirv %t.bc 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR
+; CHECK-ERROR: RequiresExtension: Feature requires the following SPIR-V extension:
+; CHECK-ERROR-NEXT: SPV_KHR_abort
+
+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"
+
+; CHECK-SPIRV: Capability AbortKHR
+; CHECK-SPIRV: Extension "SPV_KHR_abort"
+; CHECK-SPIRV: TypeInt [[#I32Ty:]] 32
+; CHECK-SPIRV: FunctionParameter [[#I32Ty]] [[#MsgId:]]
+; CHECK-SPIRV: AbortKHR [[#I32Ty]] [[#MsgId]]
+
+; CHECK-LLVM: call spir_func void @{{.*__spirv_AbortKHR.*}}(i32 %{{.*}}){{.*}}#[[#ATTR:]]
+; CHECK-LLVM-NEXT: unreachable
+; CHECK-LLVM: attributes #[[#ATTR]] = {{{.*}}noreturn{{.*}}}
+
+define spir_func void @test_abort(i32 %msg) {
+entry:
+  call spir_func void @_Z16__spirv_AbortKHRIiEvT_(i32 %msg)
+  ret void
+}
+
+; Same as @test_abort, but with an explicit `unreachable` terminator instead of
+; `ret void`. Both forms must lower to a single OpAbortKHR with no trailing
+; OpUnreachable / OpReturn.
+define spir_func void @test_abort_unreachable(i32 %msg) {
+entry:
+  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}