Implementation of the CoopVec Inference and Training builtin intrinisics

include/dxc/DXIL/DxilConstants.h

@@ -162,24 +162,32 @@ const unsigned kDxilMaxOloadDims = 2;
 
 enum class ComponentType : uint32_t {
   Invalid = 0,
-  I1,
-  I16,
-  U16,
-  I32,
-  U32,
-  I64,
-  U64,
-  F16,
-  F32,
-  F64,
-  SNormF16,
-  UNormF16,
-  SNormF32,
-  UNormF32,
-  SNormF64,
-  UNormF64,
-  PackedS8x32,
-  PackedU8x32,
+  I1 = 1,
+  I16 = 2,
+  U16 = 3,
+  I32 = 4,
+  U32 = 5,
+  I64 = 6,
+  U64 = 7,
+  F16 = 8,
+  F32 = 9,
+  F64 = 10,
+  SNormF16 = 11,
+  UNormF16 = 12,
+  SNormF32 = 13,
+  UNormF32 = 14,
+  SNormF64 = 15,
+  UNormF64 = 16,
+  PackedS8x32 = 17,
+  PackedU8x32 = 18,
+
+  // BEGIN NEW FOR SM 6.9
+  U8 = 19,
+  I8 = 20,
+  F8_E4M3 = 21,
+  F8_E5M2 = 22,
+  // END
+
   LastEntry
 };
 
@@ -769,6 +777,19 @@ enum class OpCode : unsigned {
   CreateHandleForLib =
       160, // create resource handle from resource struct for library
 
+  // Linear Algebra Operations
+  MatVecMul =
+      305, // Multiplies a MxK dimension matrix and a K sized input vector
+  MatVecMulAdd = 306, // multiplies a MxK dimension matrix and a K sized input
+                      // vector and adds an M-sized bias vector
+  OuterProductAccumulate =
+      307, // Computes the outer product between column vectors and an MxN
+           // matrix is accumulated component-wise atomically (with device
+           // scope) in memory
+  VectorAccumulate = 308, // Accumulates the components of a vector
+                          // component-wise atomically (with device scope) to
+                          // the corresponding elements of an array in memory
+
   // Mesh shader instructions
   EmitIndices = 169, // emit a primitive's vertex indices in a mesh shader
   GetMeshPayload =
@@ -1047,7 +1068,7 @@ enum class OpCode : unsigned {
   NumOpCodes_Dxil_1_7 = 226,
   NumOpCodes_Dxil_1_8 = 258,
 
-  NumOpCodes = 305 // exclusive last value of enumeration
+  NumOpCodes = 309 // exclusive last value of enumeration
 };
 // OPCODE-ENUM:END
 
@@ -1188,6 +1209,12 @@ enum class OpCodeClass : unsigned {
   // Library create handle from resource struct (like HL intrinsic)
   CreateHandleForLib,
 
+  // Linear Algebra Operations
+  MatVecMul,
+  MatVecMulAdd,
+  OuterProductAccumulate,
+  VectorAccumulate,
+
   // Mesh shader instructions
   EmitIndices,
   GetMeshPayload,
@@ -1361,7 +1388,7 @@ enum class OpCodeClass : unsigned {
   NumOpClasses_Dxil_1_7 = 153,
   NumOpClasses_Dxil_1_8 = 174,
 
-  NumOpClasses = 179 // exclusive last value of enumeration
+  NumOpClasses = 183 // exclusive last value of enumeration
 };
 // OPCODECLASS-ENUM:END
 
@@ -1532,6 +1559,28 @@ const unsigned kMSStoreOutputColOpIdx = 3;
 const unsigned kMSStoreOutputVIdxOpIdx = 4;
 const unsigned kMSStoreOutputValOpIdx = 5;
 
+// MatVec Ops
+const unsigned kMatVecMulInputVectorIdx = 1;
+const unsigned kMatVecMulIsInputUnsignedIdx = 2;
+const unsigned kMatVecMulInputInterpretationIdx = 3;
+const unsigned kMatVecMulMatrixBufferIdx = 4;
+const unsigned kMatVecMulMatrixOffsetIdx = 5;
+const unsigned kMatVecMulMatrixInterpretationIdx = 6;
+const unsigned kMatVecMulMatrixMIdx = 7;
+const unsigned kMatVecMulMatrixKIdx = 8;
+const unsigned kMatVecMulMatrixLayoutIdx = 9;
+const unsigned kMatVecMulMatrixTransposeIdx = 10;
+const unsigned kMatVecMulMatrixStrideIdx = 11;
+const unsigned kMatVecMulIsOutputUnsignedIdx = 12;
+
+// MatVecAdd
+const unsigned kMatVecMulAddBiasInterpretation = 14;
+const unsigned kMatVecMulAddIsOutputUnsignedIdx = 15;
+
+// Outer Product Accumulate
+const unsigned kOuterProdAccMatrixInterpretation = 5;
+const unsigned kOuterProdAccMatrixLayout = 6;
+
 // TODO: add operand index for all the OpCodeClass.
 } // namespace OperandIndex
 
@@ -2101,6 +2150,13 @@ extern const char *kHostLayoutTypePrefix;
 
 extern const char *kWaveOpsIncludeHelperLanesString;
 
+enum class DXILMatrixLayout : uint32_t {
+  RowMajor = 0,
+  ColumnMajor = 1,
+  MulOptimal = 2,
+  OuterProductOptimal = 3,
+};
+
 } // namespace DXIL
 
 } // namespace hlsl

include/dxc/DXIL/DxilInstructions.h

@@ -9016,5 +9016,217 @@ struct DxilInst_RawBufferVectorStore {
                              llvm::APInt(32, (uint64_t)val)));
   }
 };
+
+/// This instruction Multiplies a MxK dimension matrix and a K sized input
+/// vector
+struct DxilInst_MatVecMul {
+  llvm::Instruction *Instr;
+  // Construction and identification
+  DxilInst_MatVecMul(llvm::Instruction *pInstr) : Instr(pInstr) {}
+  operator bool() const {
+    return hlsl::OP::IsDxilOpFuncCallInst(Instr, hlsl::OP::OpCode::MatVecMul);
+  }
+  // Validation support
+  bool isAllowed() const { return true; }
+  bool isArgumentListValid() const {
+    if (13 != llvm::dyn_cast<llvm::CallInst>(Instr)->getNumArgOperands())
+      return false;
+    return true;
+  }
+  // Metadata
+  bool requiresUniformInputs() const { return false; }
+  // Operand indexes
+  enum OperandIdx {
+    arg_inputVector = 1,
+    arg_isInputUnsigned = 2,
+    arg_inputInterpretation = 3,
+    arg_matrixBuffer = 4,
+    arg_matrixOffset = 5,
+    arg_matrixIntepretation = 6,
+    arg_matrixM = 7,
+    arg_matrixK = 8,
+    arg_matrixLayout = 9,
+    arg_matrixTranspose = 10,
+    arg_matrixStride = 11,
+    arg_isOutputUnsigned = 12,
+  };
+  // Accessors
+  llvm::Value *get_inputVector() const { return Instr->getOperand(1); }
+  void set_inputVector(llvm::Value *val) { Instr->setOperand(1, val); }
+  llvm::Value *get_isInputUnsigned() const { return Instr->getOperand(2); }
+  void set_isInputUnsigned(llvm::Value *val) { Instr->setOperand(2, val); }
+  llvm::Value *get_inputInterpretation() const { return Instr->getOperand(3); }
+  void set_inputInterpretation(llvm::Value *val) { Instr->setOperand(3, val); }
+  llvm::Value *get_matrixBuffer() const { return Instr->getOperand(4); }
+  void set_matrixBuffer(llvm::Value *val) { Instr->setOperand(4, val); }
+  llvm::Value *get_matrixOffset() const { return Instr->getOperand(5); }
+  void set_matrixOffset(llvm::Value *val) { Instr->setOperand(5, val); }
+  llvm::Value *get_matrixIntepretation() const { return Instr->getOperand(6); }
+  void set_matrixIntepretation(llvm::Value *val) { Instr->setOperand(6, val); }
+  llvm::Value *get_matrixM() const { return Instr->getOperand(7); }
+  void set_matrixM(llvm::Value *val) { Instr->setOperand(7, val); }
+  llvm::Value *get_matrixK() const { return Instr->getOperand(8); }
+  void set_matrixK(llvm::Value *val) { Instr->setOperand(8, val); }
+  llvm::Value *get_matrixLayout() const { return Instr->getOperand(9); }
+  void set_matrixLayout(llvm::Value *val) { Instr->setOperand(9, val); }
+  llvm::Value *get_matrixTranspose() const { return Instr->getOperand(10); }
+  void set_matrixTranspose(llvm::Value *val) { Instr->setOperand(10, val); }
+  llvm::Value *get_matrixStride() const { return Instr->getOperand(11); }
+  void set_matrixStride(llvm::Value *val) { Instr->setOperand(11, val); }
+  llvm::Value *get_isOutputUnsigned() const { return Instr->getOperand(12); }
+  void set_isOutputUnsigned(llvm::Value *val) { Instr->setOperand(12, val); }
+};
+
+/// This instruction multiplies a MxK dimension matrix and a K sized input
+/// vector and adds an M-sized bias vector
+struct DxilInst_MatVecMulAdd {
+  llvm::Instruction *Instr;
+  // Construction and identification
+  DxilInst_MatVecMulAdd(llvm::Instruction *pInstr) : Instr(pInstr) {}
+  operator bool() const {
+    return hlsl::OP::IsDxilOpFuncCallInst(Instr,
+                                          hlsl::OP::OpCode::MatVecMulAdd);
+  }
+  // Validation support
+  bool isAllowed() const { return true; }
+  bool isArgumentListValid() const {
+    if (16 != llvm::dyn_cast<llvm::CallInst>(Instr)->getNumArgOperands())
+      return false;
+    return true;
+  }
+  // Metadata
+  bool requiresUniformInputs() const { return false; }
+  // Operand indexes
+  enum OperandIdx {
+    arg_inputVector = 1,
+    arg_isInputUnsigned = 2,
+    arg_inputInterpretation = 3,
+    arg_matrixBuffer = 4,
+    arg_matrixOffset = 5,
+    arg_matrixIntepretation = 6,
+    arg_matrixM = 7,
+    arg_matrixK = 8,
+    arg_matrixLayout = 9,
+    arg_matrixTranspose = 10,
+    arg_matrixStride = 11,
+    arg_biasBuffer = 12,
+    arg_biasOffset = 13,
+    arg_biasIntepretation = 14,
+    arg_isOutputUnsigned = 15,
+  };
+  // Accessors
+  llvm::Value *get_inputVector() const { return Instr->getOperand(1); }
+  void set_inputVector(llvm::Value *val) { Instr->setOperand(1, val); }
+  llvm::Value *get_isInputUnsigned() const { return Instr->getOperand(2); }
+  void set_isInputUnsigned(llvm::Value *val) { Instr->setOperand(2, val); }
+  llvm::Value *get_inputInterpretation() const { return Instr->getOperand(3); }
+  void set_inputInterpretation(llvm::Value *val) { Instr->setOperand(3, val); }
+  llvm::Value *get_matrixBuffer() const { return Instr->getOperand(4); }
+  void set_matrixBuffer(llvm::Value *val) { Instr->setOperand(4, val); }
+  llvm::Value *get_matrixOffset() const { return Instr->getOperand(5); }
+  void set_matrixOffset(llvm::Value *val) { Instr->setOperand(5, val); }
+  llvm::Value *get_matrixIntepretation() const { return Instr->getOperand(6); }
+  void set_matrixIntepretation(llvm::Value *val) { Instr->setOperand(6, val); }
+  llvm::Value *get_matrixM() const { return Instr->getOperand(7); }
+  void set_matrixM(llvm::Value *val) { Instr->setOperand(7, val); }
+  llvm::Value *get_matrixK() const { return Instr->getOperand(8); }
+  void set_matrixK(llvm::Value *val) { Instr->setOperand(8, val); }
+  llvm::Value *get_matrixLayout() const { return Instr->getOperand(9); }
+  void set_matrixLayout(llvm::Value *val) { Instr->setOperand(9, val); }
+  llvm::Value *get_matrixTranspose() const { return Instr->getOperand(10); }
+  void set_matrixTranspose(llvm::Value *val) { Instr->setOperand(10, val); }
+  llvm::Value *get_matrixStride() const { return Instr->getOperand(11); }
+  void set_matrixStride(llvm::Value *val) { Instr->setOperand(11, val); }
+  llvm::Value *get_biasBuffer() const { return Instr->getOperand(12); }
+  void set_biasBuffer(llvm::Value *val) { Instr->setOperand(12, val); }
+  llvm::Value *get_biasOffset() const { return Instr->getOperand(13); }
+  void set_biasOffset(llvm::Value *val) { Instr->setOperand(13, val); }
+  llvm::Value *get_biasIntepretation() const { return Instr->getOperand(14); }
+  void set_biasIntepretation(llvm::Value *val) { Instr->setOperand(14, val); }
+  llvm::Value *get_isOutputUnsigned() const { return Instr->getOperand(15); }
+  void set_isOutputUnsigned(llvm::Value *val) { Instr->setOperand(15, val); }
+};
+
+/// This instruction Computes the outer product between column vectors and an
+/// MxN matrix is accumulated component-wise atomically (with device scope) in
+/// memory
+struct DxilInst_OuterProductAccumulate {
+  llvm::Instruction *Instr;
+  // Construction and identification
+  DxilInst_OuterProductAccumulate(llvm::Instruction *pInstr) : Instr(pInstr) {}
+  operator bool() const {
+    return hlsl::OP::IsDxilOpFuncCallInst(
+        Instr, hlsl::OP::OpCode::OuterProductAccumulate);
+  }
+  // Validation support
+  bool isAllowed() const { return true; }
+  bool isArgumentListValid() const {
+    if (8 != llvm::dyn_cast<llvm::CallInst>(Instr)->getNumArgOperands())
+      return false;
+    return true;
+  }
+  // Metadata
+  bool requiresUniformInputs() const { return false; }
+  // Operand indexes
+  enum OperandIdx {
+    arg_inputVector1 = 1,
+    arg_inputVector2 = 2,
+    arg_matrixBuffer = 3,
+    arg_matrixOffset = 4,
+    arg_matrixStride = 5,
+    arg_matrixIntepretation = 6,
+    arg_matrixLayout = 7,
+  };
+  // Accessors
+  llvm::Value *get_inputVector1() const { return Instr->getOperand(1); }
+  void set_inputVector1(llvm::Value *val) { Instr->setOperand(1, val); }
+  llvm::Value *get_inputVector2() const { return Instr->getOperand(2); }
+  void set_inputVector2(llvm::Value *val) { Instr->setOperand(2, val); }
+  llvm::Value *get_matrixBuffer() const { return Instr->getOperand(3); }
+  void set_matrixBuffer(llvm::Value *val) { Instr->setOperand(3, val); }
+  llvm::Value *get_matrixOffset() const { return Instr->getOperand(4); }
+  void set_matrixOffset(llvm::Value *val) { Instr->setOperand(4, val); }
+  llvm::Value *get_matrixStride() const { return Instr->getOperand(5); }
+  void set_matrixStride(llvm::Value *val) { Instr->setOperand(5, val); }
+  llvm::Value *get_matrixIntepretation() const { return Instr->getOperand(6); }
+  void set_matrixIntepretation(llvm::Value *val) { Instr->setOperand(6, val); }
+  llvm::Value *get_matrixLayout() const { return Instr->getOperand(7); }
+  void set_matrixLayout(llvm::Value *val) { Instr->setOperand(7, val); }
+};
+
+/// This instruction Accumulates the components of a vector component-wise
+/// atomically (with device scope) to the corresponding elements of an array in
+/// memory
+struct DxilInst_VectorAccumulate {
+  llvm::Instruction *Instr;
+  // Construction and identification
+  DxilInst_VectorAccumulate(llvm::Instruction *pInstr) : Instr(pInstr) {}
+  operator bool() const {
+    return hlsl::OP::IsDxilOpFuncCallInst(Instr,
+                                          hlsl::OP::OpCode::VectorAccumulate);
+  }
+  // Validation support
+  bool isAllowed() const { return true; }
+  bool isArgumentListValid() const {
+    if (4 != llvm::dyn_cast<llvm::CallInst>(Instr)->getNumArgOperands())
+      return false;
+    return true;
+  }
+  // Metadata
+  bool requiresUniformInputs() const { return false; }
+  // Operand indexes
+  enum OperandIdx {
+    arg_inputVector = 1,
+    arg_arrayBuffer = 2,
+    arg_arrayOffset = 3,
+  };
+  // Accessors
+  llvm::Value *get_inputVector() const { return Instr->getOperand(1); }
+  void set_inputVector(llvm::Value *val) { Instr->setOperand(1, val); }
+  llvm::Value *get_arrayBuffer() const { return Instr->getOperand(2); }
+  void set_arrayBuffer(llvm::Value *val) { Instr->setOperand(2, val); }
+  llvm::Value *get_arrayOffset() const { return Instr->getOperand(3); }
+  void set_arrayOffset(llvm::Value *val) { Instr->setOperand(3, val); }
+};
 // INSTR-HELPER:END
 } // namespace hlsl

include/dxc/DxilContainer/RDAT_LibraryTypes.inl

@@ -564,9 +564,13 @@ RDAT_DXIL_ENUM_START(hlsl::DXIL::ComponentType, uint32_t)
   RDAT_ENUM_VALUE_NODEF(UNormF64)
   RDAT_ENUM_VALUE_NODEF(PackedS8x32)
   RDAT_ENUM_VALUE_NODEF(PackedU8x32)
+  RDAT_ENUM_VALUE_NODEF(U8)
+  RDAT_ENUM_VALUE_NODEF(I8)
+  RDAT_ENUM_VALUE_NODEF(F8_E4M3)
+  RDAT_ENUM_VALUE_NODEF(F8_E5M2)
   RDAT_ENUM_VALUE_NODEF(LastEntry)
 #if DEF_RDAT_ENUMS == DEF_RDAT_DUMP_IMPL
-  static_assert((unsigned)hlsl::DXIL::ComponentType::LastEntry == 19,
+  static_assert((unsigned)hlsl::DXIL::ComponentType::LastEntry == 23,
                 "otherwise, RDAT_DXIL_ENUM definition needs updating");
 #endif
 RDAT_ENUM_END()