Implement SPV_INTEL_task_sequence extension (#2340)

Spec: KhronosGroup/SPIRV-Registry#192

Author: vmaksimo

include/LLVMSPIRVExtensions.inc

@@ -70,3 +70,4 @@ EXT(SPV_INTEL_fpga_latency_control)
 EXT(SPV_INTEL_fp_max_error)
 EXT(SPV_INTEL_cache_controls)
 EXT(SPV_INTEL_subgroup_requirements)
+EXT(SPV_INTEL_task_sequence)

lib/SPIRV/OCLUtil.cpp

@@ -908,6 +908,7 @@ SPIRAddressSpace getOCLOpaqueTypeAddrSpace(Op OpCode) {
   case internal::OpTypeJointMatrixINTEL:
   case internal::OpTypeJointMatrixINTELv2:
   case OpTypeCooperativeMatrixKHR:
+  case internal::OpTypeTaskSequenceINTEL:
     return SPIRAS_Global;
   default:
     if (isSubgroupAvcINTELTypeOpCode(OpCode))

lib/SPIRV/SPIRVInternal.h

@@ -938,6 +938,7 @@ template <> inline void SPIRVMap<std::string, Op, SPIRVOpaqueType>::init() {
   _SPIRV_OP(CooperativeMatrixKHR)
 #undef _SPIRV_OP
   add("JointMatrixINTEL", internal::OpTypeJointMatrixINTEL);
+  add("TaskSequenceINTEL", internal::OpTypeTaskSequenceINTEL);
 }
 
 // Check if the module contains llvm.loop.* metadata

lib/SPIRV/SPIRVReader.cpp

@@ -491,6 +491,9 @@ Type *SPIRVToLLVM::transType(SPIRVType *T, bool UseTPT) {
     return mapType(T, transType(static_cast<SPIRVType *>(
                           BM->getEntry(FP->getPointerId()))));
   }
+  case internal::OpTypeTaskSequenceINTEL:
+    return mapType(
+        T, llvm::TargetExtType::get(*Context, "spirv.TaskSequenceINTEL"));
 
   default: {
     auto OC = T->getOpCode();
@@ -2291,6 +2294,7 @@ Value *SPIRVToLLVM::transValueWithoutDecoration(SPIRVValue *BV, Function *F,
     }
     case internal::OpTypeJointMatrixINTEL:
     case OpTypeCooperativeMatrixKHR:
+    case internal::OpTypeTaskSequenceINTEL:
       return mapValue(BV, transSPIRVBuiltinFromInst(CC, BB));
     default:
       llvm_unreachable("Unhandled type!");
@@ -3392,6 +3396,7 @@ Instruction *SPIRVToLLVM::transSPIRVBuiltinFromInst(SPIRVInstruction *BI,
   case OpSUDotAccSatKHR:
   case internal::OpJointMatrixLoadINTEL:
   case OpCooperativeMatrixLoadKHR:
+  case internal::OpTaskSequenceCreateINTEL:
     AddRetTypePostfix = true;
     break;
   default: {

lib/SPIRV/SPIRVUtil.cpp

@@ -1252,7 +1252,8 @@ SPIR::TypePrimitiveEnum getOCLTypePrimitiveEnum(StringRef TyName) {
 /// \param Signed indicates integer type should be translated as signed.
 /// \param VoidPtr indicates i8* should be translated as void*.
 static SPIR::RefParamType transTypeDesc(Type *Ty,
-                                        const BuiltinArgTypeMangleInfo &Info) {
+                                        const BuiltinArgTypeMangleInfo &Info,
+                                        StringRef InstName = "") {
   bool Signed = Info.IsSigned;
   unsigned Attr = Info.Attr;
   bool VoidPtr = Info.IsVoidPtr;
@@ -1361,8 +1362,14 @@ static SPIR::RefParamType transTypeDesc(Type *Ty,
     auto *ET = TPT->getElementType();
     SPIR::ParamType *EPT = nullptr;
     if (isa<FunctionType>(ET)) {
-      assert(isVoidFuncTy(cast<FunctionType>(ET)) && "Not supported");
-      EPT = new SPIR::BlockType;
+      FunctionType *FT = cast<FunctionType>(ET);
+      if (InstName.consume_front(kSPIRVName::Prefix) &&
+          InstName.starts_with("TaskSequence")) {
+        EPT = new SPIR::PointerType(transTypeDesc(FT->getReturnType(), Info));
+      } else {
+        assert((isVoidFuncTy(FT)) && "Not supported");
+        EPT = new SPIR::BlockType;
+      }
     } else if (auto *StructTy = dyn_cast<StructType>(ET)) {
       LLVM_DEBUG(dbgs() << "ptr to struct: " << *Ty << '\n');
       auto TyName = StructTy->getStructName();
@@ -1690,7 +1697,7 @@ std::string mangleBuiltin(StringRef UniqName, ArrayRef<Type *> ArgTypes,
         T = MangleInfo.PointerTy;
       }
       FD.Parameters.emplace_back(
-          transTypeDesc(T, BtnInfo->getTypeMangleInfo(I)));
+          transTypeDesc(T, BtnInfo->getTypeMangleInfo(I), UniqName));
     }
   }
   // Ellipsis must be the last argument of any function

lib/SPIRV/SPIRVWriter.cpp

@@ -617,6 +617,8 @@ SPIRVType *LLVMToSPIRVBase::transType(Type *T) {
           Args.emplace_back(transConstant(getUInt32(M, Op)));
         return mapType(T, BM->addCooperativeMatrixKHRType(ElemTy, Args));
       }
+      case internal::OpTypeTaskSequenceINTEL:
+        return mapType(T, BM->addTaskSequenceINTELType());
       default:
         if (isSubgroupAvcINTELTypeOpCode(Opcode))
           return mapType(T, BM->addSubgroupAvcINTELType(Opcode));

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -226,7 +226,8 @@ class SPIRVInstTemplateBase : public SPIRVInstruction {
   virtual void init() {}
   virtual void initImpl(Op OC, bool HasId = true, SPIRVWord WC = 0,
                         bool VariWC = false, unsigned Lit1 = ~0U,
-                        unsigned Lit2 = ~0U, unsigned Lit3 = ~0U) {
+                        unsigned Lit2 = ~0U, unsigned Lit3 = ~0U,
+                        unsigned Lit4 = ~0U) {
     OpCode = OC;
     if (!HasId) {
       setHasNoId();
@@ -238,6 +239,7 @@ class SPIRVInstTemplateBase : public SPIRVInstruction {
     addLit(Lit1);
     addLit(Lit2);
     addLit(Lit3);
+    addLit(Lit4);
   }
   bool isOperandLiteral(unsigned I) const override { return Lit.count(I); }
   void addLit(unsigned L) {
@@ -364,14 +366,16 @@ class SPIRVInstTemplateBase : public SPIRVInstruction {
 
 template <typename BT = SPIRVInstTemplateBase, Op OC = OpNop, bool HasId = true,
           SPIRVWord WC = 0, bool HasVariableWC = false, unsigned Literal1 = ~0U,
-          unsigned Literal2 = ~0U, unsigned Literal3 = ~0U>
+          unsigned Literal2 = ~0U, unsigned Literal3 = ~0U,
+          unsigned Literal4 = ~0U>
 class SPIRVInstTemplate : public BT {
 public:
   typedef BT BaseTy;
   SPIRVInstTemplate() { init(); }
   ~SPIRVInstTemplate() override {}
   void init() override {
-    this->initImpl(OC, HasId, WC, HasVariableWC, Literal1, Literal2, Literal3);
+    this->initImpl(OC, HasId, WC, HasVariableWC, Literal1, Literal2, Literal3,
+                   Literal4);
   }
 };
 
@@ -3854,5 +3858,85 @@ template <Op OC> class SPIRVReadClockKHRInstBase : public SPIRVUnaryInst<OC> {
 _SPIRV_OP(ReadClockKHR)
 #undef _SPIRV_OP
 
+class SPIRVTaskSequenceINTELInstBase : public SPIRVInstTemplateBase {
+public:
+  std::optional<ExtensionID> getRequiredExtension() const override {
+    return ExtensionID::SPV_INTEL_task_sequence;
+  }
+};
+
+class SPIRVTaskSequenceINTELInst : public SPIRVTaskSequenceINTELInstBase {
+public:
+  SPIRVCapVec getRequiredCapability() const override {
+    return getVec(internal::CapabilityTaskSequenceINTEL);
+  }
+};
+
+class SPIRVTaskSequenceCreateINTELInst : public SPIRVTaskSequenceINTELInst {
+protected:
+  void validate() const override {
+    SPIRVInstruction::validate();
+    std::string InstName = "TaskSequenceCreateINTEL";
+    SPIRVErrorLog &SPVErrLog = this->getModule()->getErrorLog();
+
+    SPIRVType *ResTy = this->getType();
+    SPVErrLog.checkError(
+        ResTy->isTypeTaskSequenceINTEL(), SPIRVEC_InvalidInstruction,
+        InstName + "\nResult must be TaskSequenceINTEL type\n");
+
+    SPIRVValue *Func =
+        const_cast<SPIRVTaskSequenceCreateINTELInst *>(this)->getOperand(0);
+    SPVErrLog.checkError(
+        Func->getOpCode() == OpFunction, SPIRVEC_InvalidInstruction,
+        InstName + "\nFirst argument is expected to be a function.\n");
+
+    SPIRVConstant *PipelinedConst = static_cast<SPIRVConstant *>(
+        const_cast<SPIRVTaskSequenceCreateINTELInst *>(this)->getOperand(1));
+    const int Pipelined = PipelinedConst->getZExtIntValue();
+    SPVErrLog.checkError(Pipelined >= -1, SPIRVEC_InvalidInstruction,
+                         InstName + "\nPipeline must be a 32 bit integer with "
+                                    "the value bigger or equal to -1.\n");
+
+    const int ClusterMode =
+        static_cast<SPIRVConstant *>(
+            const_cast<SPIRVTaskSequenceCreateINTELInst *>(this)->getOperand(2))
+            ->getZExtIntValue();
+    SPVErrLog.checkError(
+        ClusterMode >= -1 && ClusterMode <= 1, SPIRVEC_InvalidInstruction,
+        InstName + "\nClusterMode valid values are -1, 0, 1.\n");
+
+    const uint32_t GetCapacity =
+        static_cast<SPIRVConstant *>(
+            const_cast<SPIRVTaskSequenceCreateINTELInst *>(this)->getOperand(3))
+            ->getZExtIntValue();
+    SPVErrLog.checkError(
+        GetCapacity, SPIRVEC_InvalidInstruction,
+        InstName + "\nGetCapacity must be unsigned 32 bit integer.\n");
+
+    const uint32_t AsyncCapacity =
+        static_cast<SPIRVConstant *>(
+            const_cast<SPIRVTaskSequenceCreateINTELInst *>(this)->getOperand(4))
+            ->getZExtIntValue();
+    SPVErrLog.checkError(
+        AsyncCapacity, SPIRVEC_InvalidInstruction,
+        InstName + "\nAsyncCapacity must be unsigned 32 bit integer.\n");
+  }
+};
+
+#define _SPIRV_OP(x, ...)                                                      \
+  typedef SPIRVInstTemplate<SPIRVTaskSequenceINTELInst,                        \
+                            internal::Op##x##INTEL, __VA_ARGS__>               \
+      SPIRV##x##INTEL;
+_SPIRV_OP(TaskSequenceAsync, false, 2, true)
+_SPIRV_OP(TaskSequenceGet, true, 4, false)
+_SPIRV_OP(TaskSequenceRelease, false, 2, false)
+#undef _SPIRV_OP
+#define _SPIRV_OP(x, ...)                                                      \
+  typedef SPIRVInstTemplate<SPIRVTaskSequenceCreateINTELInst,                  \
+                            internal::Op##x##INTEL, __VA_ARGS__>               \
+      SPIRV##x##INTEL;
+_SPIRV_OP(TaskSequenceCreate, true, 8, false, 1, 2, 3, 4)
+#undef _SPIRV_OP
+
 } // namespace SPIRV
 #endif // SPIRV_LIBSPIRV_SPIRVINSTRUCTION_H

lib/SPIRV/libSPIRV/SPIRVModule.cpp

@@ -247,6 +247,7 @@ class SPIRVModuleImpl : public SPIRVModule {
   addJointMatrixINTELType(SPIRVType *, std::vector<SPIRVValue *>) override;
   SPIRVTypeCooperativeMatrixKHR *
   addCooperativeMatrixKHRType(SPIRVType *, std::vector<SPIRVValue *>) override;
+  SPIRVTypeTaskSequenceINTEL *addTaskSequenceINTELType() override;
   SPIRVType *addOpaqueGenericType(Op) override;
   SPIRVTypeDeviceEvent *addDeviceEventType() override;
   SPIRVTypeQueue *addQueueType() override;
@@ -1021,6 +1022,10 @@ SPIRVModuleImpl::addCooperativeMatrixKHRType(SPIRVType *CompType,
       new SPIRVTypeCooperativeMatrixKHR(this, getId(), CompType, Args));
 }
 
+SPIRVTypeTaskSequenceINTEL *SPIRVModuleImpl::addTaskSequenceINTELType() {
+  return addType(new SPIRVTypeTaskSequenceINTEL(this, getId()));
+}
+
 SPIRVType *SPIRVModuleImpl::addOpaqueGenericType(Op TheOpCode) {
   return addType(new SPIRVTypeOpaqueGeneric(TheOpCode, this, getId()));
 }

lib/SPIRV/libSPIRV/SPIRVModule.h

@@ -97,6 +97,7 @@ class SPIRVTypeBufferSurfaceINTEL;
 class SPIRVTypeTokenINTEL;
 class SPIRVTypeJointMatrixINTEL;
 class SPIRVTypeCooperativeMatrixKHR;
+class SPIRVTypeTaskSequenceINTEL;
 
 typedef SPIRVBasicBlock SPIRVLabel;
 struct SPIRVTypeImageDescriptor;
@@ -264,6 +265,7 @@ class SPIRVModule {
   addJointMatrixINTELType(SPIRVType *, std::vector<SPIRVValue *>) = 0;
   virtual SPIRVTypeCooperativeMatrixKHR *
   addCooperativeMatrixKHRType(SPIRVType *, std::vector<SPIRVValue *>) = 0;
+  virtual SPIRVTypeTaskSequenceINTEL *addTaskSequenceINTELType() = 0;
   virtual SPIRVTypeVoid *addVoidType() = 0;
   virtual SPIRVType *addOpaqueGenericType(Op) = 0;
   virtual SPIRVTypeDeviceEvent *addDeviceEventType() = 0;

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -670,6 +670,7 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
       "CooperativeMatrixCheckedInstructionsINTEL");
   add(internal::CapabilitySubgroupRequirementsINTEL,
       "SubgroupRequirementsINTEL");
+  add(internal::CapabilityTaskSequenceINTEL, "TaskSequenceINTEL");
 }
 SPIRV_DEF_NAMEMAP(Capability, SPIRVCapabilityNameMap)