[SPIRV] Add support for SPV_KHR_abort extension

llvm/docs/SPIRVUsage.rst

@@ -213,6 +213,8 @@ Below is a list of supported SPIR-V extensions, sorted alphabetically by their e
      - Allows to allocate local arrays whose number of elements is unknown at compile time.
    * - ``SPV_INTEL_joint_matrix``
      - Adds few matrix capabilities on top of SPV_KHR_cooperative_matrix extension, such as matrix prefetch, get element coordinate and checked load/store/construct instructions, tensor float 32 and bfloat type interpretations for multiply-add instruction.
+   * - ``SPV_KHR_abort``
+     - Adds a function-termination instruction that signals an abnormal error to the client API. Currently supported in OpenCL environments only.
    * - ``SPV_KHR_bit_instructions``
      - Enables bit instructions to be used by SPIR-V modules without requiring the Shader capability.
    * - ``SPV_KHR_expect_assume``

llvm/include/llvm/IR/IntrinsicsSPIRV.td

@@ -39,6 +39,7 @@ let TargetPrefix = "spv" in {
   def int_spv_selection_merge : Intrinsic<[], [llvm_any_ty, llvm_i32_ty], [ImmArg<ArgIndex<1>>]>;
   def int_spv_cmpxchg : Intrinsic<[llvm_i32_ty], [llvm_any_ty, llvm_vararg_ty]>;
   def int_spv_unreachable : Intrinsic<[], []>;
+  def int_spv_abort : Intrinsic<[], [llvm_vararg_ty], [IntrNoReturn]>;
   def int_spv_alloca : Intrinsic<[llvm_any_ty], [llvm_i32_ty], [ImmArg<ArgIndex<0>>]>;
   def int_spv_alloca_array : Intrinsic<[llvm_any_ty], [llvm_anyint_ty, llvm_i32_ty], [ImmArg<ArgIndex<1>>]>;
   def int_spv_undef : Intrinsic<[llvm_i32_ty], []>;

llvm/lib/Target/SPIRV/SPIRVCommandLine.cpp

@@ -176,7 +176,8 @@ static const StringMap<SPIRV::Extension::Extension> SPIRVExtensionMap = {
     {"SPV_EXT_image_raw10_raw12",
      SPIRV::Extension::Extension::SPV_EXT_image_raw10_raw12},
     {"SPV_INTEL_unstructured_loop_controls",
-     SPIRV::Extension::Extension::SPV_INTEL_unstructured_loop_controls}};
+     SPIRV::Extension::Extension::SPV_INTEL_unstructured_loop_controls},
+    {"SPV_KHR_abort", SPIRV::Extension::Extension::SPV_KHR_abort}};
 
 bool SPIRVExtensionsParser::parse(cl::Option &O, StringRef ArgName,
                                   StringRef ArgValue, ExtensionSet &Vals) {

llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp

@@ -2253,6 +2253,29 @@ Instruction *SPIRVEmitIntrinsics::visitAtomicCmpXchgInst(AtomicCmpXchgInst &I) {
 Instruction *SPIRVEmitIntrinsics::visitUnreachableInst(UnreachableInst &I) {
   IRBuilder<> B(I.getParent());
   B.SetInsertPoint(&I);
+  // OpAbortKHR is itself a SPIR-V block terminator. If the immediately
+  // preceding instruction is a call to llvm.spv.abort, do not emit an
+  // additional OpUnreachable, which would leave the SPIR-V block with two
+  // terminators and produce invalid SPIR-V. The check is intentionally limited
+  // to the directly-preceding non-debug instruction: any real instruction
+  // sitting between `llvm.spv.abort` and `unreachable` would also be invalid
+  // SPIR-V (nothing can follow OpAbortKHR in the same block), so assert that
+  // shape if we see an `spv_abort` anywhere earlier in the block.
+  Instruction *Prev = I.getPrevNode();
+  while (Prev && Prev->isDebugOrPseudoInst())
+    Prev = Prev->getPrevNode();
+  if (auto *CI = dyn_cast_or_null<CallInst>(Prev);
+      CI && CI->getIntrinsicID() == Intrinsic::spv_abort)
+    return &I;
+#ifndef NDEBUG
+  for (Instruction *P = I.getPrevNode(); P; P = P->getPrevNode()) {
+    auto *CI = dyn_cast<CallInst>(P);
+    if (CI && CI->getIntrinsicID() == Intrinsic::spv_abort)
+      llvm_unreachable("llvm.spv.abort must be the last non-debug instruction "
+                       "before its block's `unreachable`; OpAbortKHR is itself "
+                       "a SPIR-V block terminator");
+  }
+#endif
   B.CreateIntrinsic(Intrinsic::spv_unreachable, {});
   return &I;
 }

llvm/lib/Target/SPIRV/SPIRVInstrInfo.td

@@ -641,6 +641,8 @@ let isReturn = 1, hasDelaySlot = 0, isBarrier = 0, isTerminator = 1, isNotDuplic
   def OpReturn: SimpleOp<"OpReturn", 253>;
   def OpReturnValue: Op<254, (outs), (ins ID:$ret), "OpReturnValue $ret">;
   def OpUnreachable: SimpleOp<"OpUnreachable", 255>;
+  def OpAbortKHR: Op<5121, (outs), (ins TYPE:$msg_type, ID:$msg),
+                  "OpAbortKHR $msg_type $msg">;
 }
 def OpLifetimeStart: Op<256, (outs), (ins ID:$ptr, i32imm:$sz), "OpLifetimeStart $ptr $sz">;
 def OpLifetimeStop: Op<257, (outs), (ins ID:$ptr, i32imm:$sz), "OpLifetimeStop $ptr $sz">;

llvm/lib/Target/SPIRV/SPIRVInstructionSelector.cpp

@@ -350,6 +350,7 @@ class SPIRVInstructionSelector : public InstructionSelector {
 
   bool diagnoseUnsupported(const MachineInstr &I, const Twine &Msg) const;
 
+  bool selectAbort(MachineInstr &I) const;
   bool selectFrameIndex(Register ResVReg, SPIRVTypeInst ResType,
                         MachineInstr &I) const;
   bool selectAllocaArray(Register ResVReg, SPIRVTypeInst ResType,
@@ -4635,6 +4636,8 @@ bool SPIRVInstructionSelector::selectIntrinsic(Register ResVReg,
     BuildMI(BB, I, I.getDebugLoc(), TII.get(SPIRV::OpUnreachable))
         .constrainAllUses(TII, TRI, RBI);
     return true;
+  case Intrinsic::spv_abort:
+    return selectAbort(I);
   case Intrinsic::spv_alloca:
     return selectFrameIndex(ResVReg, ResType, I);
   case Intrinsic::spv_alloca_array:
@@ -6250,6 +6253,63 @@ bool SPIRVInstructionSelector::selectAllocaArray(Register ResVReg,
   return true;
 }
 
+bool SPIRVInstructionSelector::selectAbort(MachineInstr &I) const {
+  if (!STI.canUseExtension(SPIRV::Extension::SPV_KHR_abort))
+    report_fatal_error("OpAbortKHR instruction requires the following "
+                       "SPIR-V extension: SPV_KHR_abort",
+                       false);
+  // The intrinsic is declared as variadic so it can carry composite message
+  // types (vectors and structs) without LLVM IR mangling restrictions, but
+  // OpAbortKHR takes exactly one Message operand.
+  if (I.getNumExplicitOperands() != 2)
+    report_fatal_error("llvm.spv.abort must be called with exactly one "
+                       "message argument",
+                       false);
+  Register MsgReg = I.getOperand(1).getReg();
+  SPIRVTypeInst MsgType = GR.getSPIRVTypeForVReg(MsgReg);
+  assert(MsgType && "Message argument of llvm.spv.abort has no SPIR-V type");
+  // SPV_KHR_abort requires Message Type to be a concrete type. Per the
+  // SPIR-V "Concrete Type" definition, that means a numerical scalar
+  // (int/float), a (physical) pointer, a vector, matrix, or any aggregate
+  // (array/struct) recursively containing only such types. OpTypeBool,
+  // OpTypeVoid, opaque handles and similar abstract/non-concrete types are
+  // rejected up front rather than emitting invalid SPIR-V. Validate
+  // recursively so that e.g. a struct containing a bool is also rejected.
+  SmallVector<SPIRVTypeInst, 4> Worklist{MsgType};
+  while (!Worklist.empty()) {
+    SPIRVTypeInst Ty = Worklist.pop_back_val();
+    switch (Ty->getOpcode()) {
+    case SPIRV::OpTypeInt:
+    case SPIRV::OpTypeFloat:
+    case SPIRV::OpTypePointer:
+      break;
+    case SPIRV::OpTypeVector:
+    case SPIRV::OpTypeMatrix:
+    case SPIRV::OpTypeArray:
+      // Operand 1 holds the element/component type id.
+      Worklist.push_back(GR.getSPIRVTypeForVReg(Ty->getOperand(1).getReg()));
+      break;
+    case SPIRV::OpTypeStruct:
+      // Operands 1..N hold the field type ids.
+      for (unsigned Idx = 1, E = Ty->getNumOperands(); Idx < E; ++Idx)
+        Worklist.push_back(
+            GR.getSPIRVTypeForVReg(Ty->getOperand(Idx).getReg()));
+      break;
+    default:
+      report_fatal_error("llvm.spv.abort message type must be a concrete "
+                         "SPIR-V type (numerical scalar, pointer, vector, "
+                         "matrix, or aggregate of such types)",
+                         false);
+    }
+  }
+  MachineBasicBlock &BB = *I.getParent();
+  BuildMI(BB, I, I.getDebugLoc(), TII.get(SPIRV::OpAbortKHR))
+      .addUse(GR.getSPIRVTypeID(MsgType))
+      .addUse(MsgReg)
+      .constrainAllUses(TII, TRI, RBI);
+  return true;
+}
+
 bool SPIRVInstructionSelector::selectFrameIndex(Register ResVReg,
                                                 SPIRVTypeInst ResType,
                                                 MachineInstr &I) const {

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

@@ -1945,6 +1945,14 @@ void addInstrRequirements(const MachineInstr &MI,
     Reqs.addExtension(SPIRV::Extension::SPV_KHR_shader_clock);
     Reqs.addCapability(SPIRV::Capability::ShaderClockKHR);
     break;
+  case SPIRV::OpAbortKHR:
+    if (!ST.canUseExtension(SPIRV::Extension::SPV_KHR_abort))
+      report_fatal_error("OpAbortKHR instruction requires the "
+                         "following SPIR-V extension: SPV_KHR_abort",
+                         false);
+    Reqs.addExtension(SPIRV::Extension::SPV_KHR_abort);
+    Reqs.addCapability(SPIRV::Capability::AbortKHR);
+    break;
   case SPIRV::OpFunctionPointerCallINTEL:
     if (ST.canUseExtension(SPIRV::Extension::SPV_INTEL_function_pointers)) {
       Reqs.addExtension(SPIRV::Extension::SPV_INTEL_function_pointers);

llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp

@@ -19,6 +19,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "SPIRV.h"
+#include "SPIRVBuiltins.h"
 #include "SPIRVSubtarget.h"
 #include "SPIRVTargetMachine.h"
 #include "SPIRVUtils.h"
@@ -33,6 +34,7 @@
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsSPIRV.h"
 #include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LowerMemIntrinsics.h"
 #include <regex>
 
@@ -43,6 +45,7 @@ namespace {
 class SPIRVPrepareFunctions : public ModulePass {
   const SPIRVTargetMachine &TM;
   bool substituteIntrinsicCalls(Function *F);
+  bool substituteAbortKHRCalls(Function *F);
   Function *removeAggregateTypesFromSignature(Function *F);
   bool removeAggregateTypesFromCalls(Function *F);
 
@@ -595,6 +598,62 @@ SPIRVPrepareFunctions::removeAggregateTypesFromSignature(Function *F) {
   return NewF;
 }
 
+// Replace OpenCL/SPIR-V style calls to `__spirv_AbortKHR(message)` with calls
+// to the `llvm.spv.abort` target intrinsic, so that they go through the same
+// instruction-selection path as the intrinsic and get lowered to OpAbortKHR.
+bool SPIRVPrepareFunctions::substituteAbortKHRCalls(Function *F) {
+  if (F->isDeclaration())
+    return false;
+
+  SmallVector<CallInst *> Calls;
+  for (Instruction &I : instructions(F)) {
+    auto *CI = dyn_cast<CallInst>(&I);
+    if (!CI)
+      continue;
+    Function *Callee = CI->getCalledFunction();
+    if (!Callee || Callee->isIntrinsic())
+      continue;
+    StringRef Demangled = Callee->getName();
+    std::string DemangledStr = getOclOrSpirvBuiltinDemangledName(Demangled);
+    if (DemangledStr.empty())
+      continue;
+    std::string BuiltinName = SPIRV::lookupBuiltinNameHelper(DemangledStr);
+    if (StringRef(BuiltinName) != "__spirv_AbortKHR")
+      continue;
+    if (CI->arg_size() != 1)
+      continue;
+    Calls.push_back(CI);
+  }
+
+  if (Calls.empty())
+    return false;
+
+  for (CallInst *CI : Calls) {
+    IRBuilder<> B(CI);
+    Value *Msg = CI->getArgOperand(0);
+    // The OpenCL C ABI may pass aggregate arguments by pointer (byval). In
+    // that case load the underlying value so that OpAbortKHR receives the
+    // composite itself, as required by the SPV_KHR_abort spec ("Message Type
+    // must be a concrete type").
+    if (CI->isByValArgument(0)) {
+      Type *AggTy = CI->getParamByValType(0);
+      Msg = B.CreateLoad(AggTy, Msg);
+    }
+    B.CreateIntrinsic(Intrinsic::spv_abort, {}, {Msg});
+    // OpAbortKHR is itself a SPIR-V function-termination instruction and must
+    // be the last instruction in its block. Drop the original call and
+    // everything that follows it (the OpenCL ABI typically appends stores into
+    // the return slot and a `ret`), and re-terminate the block with
+    // `unreachable`. We use changeToUnreachable so that any successor PHI
+    // nodes have the now-removed predecessor edge cleaned up; otherwise the
+    // IR verifier would reject mismatched PHI incoming entries. The matching
+    // suppression in SPIRVEmitIntrinsics::visitUnreachableInst ensures no
+    // extra OpUnreachable is emitted after OpAbortKHR.
+    changeToUnreachable(CI);
+  }
+  return true;
+}
+
 // Mutates indirect callsites iff if aggregate argument/return types are present
 // with the types replaced by i32 types. The change in types is noted in
 // 'spv.mutated_callsites' metadata for later restoration.
@@ -684,6 +743,7 @@ bool SPIRVPrepareFunctions::runOnModule(Module &M) {
 
   for (Function &F : M) {
     Changed |= substituteIntrinsicCalls(&F);
+    Changed |= substituteAbortKHRCalls(&F);
     Changed |= sortBlocks(F);
     Changed |= removeAggregateTypesFromCalls(&F);
   }

llvm/lib/Target/SPIRV/SPIRVSymbolicOperands.td

@@ -398,6 +398,10 @@ defm SPV_EXT_image_raw10_raw12 :ExtensionOperand<133, [EnvOpenCL, EnvVulkan]>;
 defm SPV_ALTERA_arbitrary_precision_floating_point: ExtensionOperand<134, [EnvOpenCL]>;
 defm SPV_KHR_fma : ExtensionOperand<135, [EnvVulkan, EnvOpenCL]>;
 defm SPV_INTEL_masked_gather_scatter : ExtensionOperand<136, [EnvOpenCL]>;
+// SPV_KHR_abort is currently only supported in the OpenCL environment.
+// Vulkan/Shader support requires emitting explicit layout decorations on the
+// Message Type, which is not yet implemented.
+defm SPV_KHR_abort : ExtensionOperand<137, [EnvOpenCL]>;
 
 //===----------------------------------------------------------------------===//
 // Multiclass used to define Capabilities enum values and at the same time
@@ -598,6 +602,7 @@ defm GlobalVariableFPGADecorationsINTEL : CapabilityOperand<6189, 0, 0, [SPV_INT
 defm CacheControlsINTEL : CapabilityOperand<6441, 0, 0, [SPV_INTEL_cache_controls], []>;
 defm CooperativeMatrixKHR : CapabilityOperand<6022, 0, 0, [SPV_KHR_cooperative_matrix], []>;
 defm ArithmeticFenceEXT : CapabilityOperand<6144, 0, 0, [SPV_EXT_arithmetic_fence], []>;
+defm AbortKHR : CapabilityOperand<5120, 0, 0, [SPV_KHR_abort], []>;
 defm SplitBarrierINTEL : CapabilityOperand<6141, 0, 0, [SPV_INTEL_split_barrier], []>;
 defm CooperativeMatrixCheckedInstructionsINTEL : CapabilityOperand<6192, 0, 0, [SPV_INTEL_joint_matrix], []>;
 defm CooperativeMatrixPrefetchINTEL : CapabilityOperand<6411, 0, 0, [SPV_INTEL_joint_matrix], []>;

llvm/test/CodeGen/SPIRV/extensions/SPV_KHR_abort/abort-composite-construct.ll

@@ -0,0 +1,30 @@
+; RUN: llc -O0 -mtriple=spirv32-unknown-unknown --spirv-ext=+SPV_KHR_abort %s -o - | FileCheck %s
+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown --spirv-ext=+SPV_KHR_abort %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown --spirv-ext=+SPV_KHR_abort %s -o - -filetype=obj | spirv-val %}
+
+;; Verify that OpAbortKHR can consume a composite produced by
+;; OpCompositeConstruct, mirroring the SPV_KHR_abort spec example which builds
+;; the message via OpCompositeConstruct.
+
+; CHECK-DAG: OpCapability AbortKHR
+; CHECK-DAG: OpExtension "SPV_KHR_abort"
+; CHECK-DAG: %[[#I32:]] = OpTypeInt 32 0
+; CHECK-DAG: %[[#V2:]] = OpTypeVector %[[#I32]] 2
+
+; CHECK:     %[[#CC:]] = OpCompositeConstruct %[[#V2]] %{{[0-9]+}} %{{[0-9]+}}
+; CHECK:     OpAbortKHR %[[#V2]] %[[#CC]]
+; CHECK-NOT: OpUnreachable
+
+declare void @llvm.spv.abort(...) #0
+
+define spir_kernel void @abort_composite_construct(i32 %a, i32 %b) {
+entry:
+  %va = insertelement <1 x i32> poison, i32 %a, i32 0
+  %vb = insertelement <1 x i32> poison, i32 %b, i32 0
+  %v = shufflevector <1 x i32> %va, <1 x i32> %vb,
+                     <2 x i32> <i32 0, i32 1>
+  call void (...) @llvm.spv.abort(<2 x i32> %v)
+  unreachable
+}
+
+attributes #0 = { noreturn }

llvm/test/CodeGen/SPIRV/extensions/SPV_KHR_abort/abort-composite.ll

@@ -0,0 +1,51 @@
+; RUN: llc -O0 -mtriple=spirv32-unknown-unknown --spirv-ext=+SPV_KHR_abort %s -o - | FileCheck %s
+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown --spirv-ext=+SPV_KHR_abort %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown --spirv-ext=+SPV_KHR_abort %s -o - -filetype=obj | spirv-val %}
+
+; Verify that OpAbortKHR can consume a composite (struct or vector) message,
+; built up via OpCompositeInsert, mirroring the example in the SPV_KHR_abort
+; specification.
+
+; CHECK-DAG: OpCapability AbortKHR
+; CHECK-DAG: OpExtension "SPV_KHR_abort"
+
+; CHECK-DAG: %[[#I32:]] = OpTypeInt 32 0
+; CHECK-DAG: %[[#F32:]] = OpTypeFloat 32
+; CHECK-DAG: %[[#STRUCT:]] = OpTypeStruct %[[#I32]] %[[#I32]] %[[#F32]]
+; CHECK-DAG: %[[#VEC:]] = OpTypeVector %[[#I32]] 4
+
+; CHECK: %[[#S0:]] = OpCompositeInsert %[[#STRUCT]] %{{[0-9]+}} %{{[0-9]+}} 0
+; CHECK: %[[#S1:]] = OpCompositeInsert %[[#STRUCT]] %{{[0-9]+}} %[[#S0]] 1
+; CHECK: %[[#S2:]] = OpCompositeInsert %[[#STRUCT]] %{{[0-9]+}} %[[#S1]] 2
+; CHECK: OpAbortKHR %[[#STRUCT]] %[[#S2]]
+
+; CHECK: %[[#V0:]] = OpCompositeInsert %[[#VEC]] %{{[0-9]+}} %{{[0-9]+}} 0
+; CHECK: %[[#V1:]] = OpCompositeInsert %[[#VEC]] %{{[0-9]+}} %[[#V0]] 1
+; CHECK: %[[#V2:]] = OpCompositeInsert %[[#VEC]] %{{[0-9]+}} %[[#V1]] 2
+; CHECK: %[[#V3:]] = OpCompositeInsert %[[#VEC]] %{{[0-9]+}} %[[#V2]] 3
+; CHECK: OpAbortKHR %[[#VEC]] %[[#V3]]
+
+%struct.Msg = type { i32, i32, float }
+
+declare void @llvm.spv.abort(...) #0
+
+define spir_kernel void @abort_with_struct(i32 %x, i32 %y, float %z) {
+entry:
+  %m0 = insertvalue %struct.Msg poison, i32 %x, 0
+  %m1 = insertvalue %struct.Msg %m0, i32 %y, 1
+  %m2 = insertvalue %struct.Msg %m1, float %z, 2
+  call void (...) @llvm.spv.abort(%struct.Msg %m2)
+  unreachable
+}
+
+define spir_kernel void @abort_with_vector(i32 %a, i32 %b, i32 %c, i32 %d) {
+entry:
+  %v0 = insertelement <4 x i32> poison, i32 %a, i32 0
+  %v1 = insertelement <4 x i32> %v0, i32 %b, i32 1
+  %v2 = insertelement <4 x i32> %v1, i32 %c, i32 2
+  %v3 = insertelement <4 x i32> %v2, i32 %d, i32 3
+  call void (...) @llvm.spv.abort(<4 x i32> %v3)
+  unreachable
+}
+
+attributes #0 = { noreturn }

llvm/test/CodeGen/SPIRV/extensions/SPV_KHR_abort/abort-invalid-message-type.ll

@@ -0,0 +1,26 @@
+; RUN: split-file %s %t
+; RUN: not llc -O0 -mtriple=spirv64-unknown-unknown --spirv-ext=+SPV_KHR_abort %t/top-level-bool.ll -o /dev/null 2>&1 | FileCheck %s
+; RUN: not llc -O0 -mtriple=spirv64-unknown-unknown --spirv-ext=+SPV_KHR_abort %t/nested-bool.ll -o /dev/null 2>&1 | FileCheck %s
+
+; CHECK: llvm.spv.abort message type must be a concrete SPIR-V type
+
+;--- top-level-bool.ll
+declare void @llvm.spv.abort(...) #0
+define void @abort_with_bool(i1 %b) {
+entry:
+  call void (...) @llvm.spv.abort(i1 %b)
+  unreachable
+}
+attributes #0 = { noreturn }
+
+;--- nested-bool.ll
+%B = type { i32, i1 }
+declare void @llvm.spv.abort(...) #0
+define void @abort_with_struct_of_bool(i1 %b) {
+entry:
+  %s0 = insertvalue %B poison, i32 0, 0
+  %s1 = insertvalue %B %s0, i1 %b, 1
+  call void (...) @llvm.spv.abort(%B %s1)
+  unreachable
+}
+attributes #0 = { noreturn }