Add a new form of load/store operations for cooperative matrices that accepts two separate arguments: the row index and the column index.

Author: YixingZhang007

lib/SPIRV/SPIRVReader.cpp

@@ -3540,6 +3540,7 @@ Instruction *SPIRVToLLVM::transSPIRVBuiltinFromInst(SPIRVInstruction *BI,
   case internal::OpJointMatrixLoadINTEL:
   case OpCooperativeMatrixLoadKHR:
   case internal::OpCooperativeMatrixLoadCheckedINTEL:
+  case internal::OpCooperativeMatrixLoadOffsetINTEL:
   case internal::OpTaskSequenceCreateINTEL:
   case internal::OpConvertHandleToImageINTEL:
   case internal::OpConvertHandleToSampledImageINTEL:

lib/SPIRV/libSPIRV/SPIRVEnum.h

@@ -221,6 +221,8 @@ template <> inline void SPIRVMap<SPIRVCapabilityKind, SPIRVCapVec>::init() {
                {CapabilityCooperativeMatrixKHR});
   ADD_VEC_INIT(internal::CapabilityCooperativeMatrixCheckedInstructionsINTEL,
                {CapabilityCooperativeMatrixKHR});
+  ADD_VEC_INIT(internal::CapabilityCooperativeMatrixOffsetInstructionsINTEL,
+               {CapabilityCooperativeMatrixKHR});
 }
 
 template <> inline void SPIRVMap<SPIRVExecutionModelKind, SPIRVCapVec>::init() {

lib/SPIRV/libSPIRV/SPIRVInstruction.h

@@ -3710,6 +3710,27 @@ _SPIRV_OP(CooperativeMatrixStoreChecked, false, 8, true, 8)
 _SPIRV_OP(CooperativeMatrixConstructChecked, true, 8)
 #undef _SPIRV_OP
 
+class SPIRVJointMatrixOffsetInstructionsINTELInstBase
+      : public SPIRVInstTemplateBase {
+protected:
+  std::optional<ExtensionID> getRequiredExtension() const override {
+    return ExtensionID::SPV_INTEL_joint_matrix;
+  }
+  SPIRVCapVec getRequiredCapability() const override {
+    return getVec(
+        internal::CapabilityCooperativeMatrixOffsetInstructionsINTEL);
+  }
+};
+
+#define _SPIRV_OP(x, ...)                                                      \
+  typedef SPIRVInstTemplate<                                                   \
+      SPIRVJointMatrixOffsetInstructionsINTELInstBase,                         \
+      internal::Op##x##INTEL, __VA_ARGS__>                                     \
+      SPIRV##x##INTEL;
+_SPIRV_OP(CooperativeMatrixLoadOffset, true, 8, true, 6)
+_SPIRV_OP(CooperativeMatrixStoreOffset, false, 7, true, 7)
+#undef _SPIRV_OP
+
 class SPIRVCooperativeMatrixInvocationInstructionsINTELInstBase
     : public SPIRVInstTemplateBase {
 protected:

lib/SPIRV/libSPIRV/SPIRVNameMapEnum.h

@@ -671,6 +671,8 @@ template <> inline void SPIRVMap<Capability, std::string>::init() {
       "CooperativeMatrixInvocationInstructionsINTEL");
   add(internal::CapabilityCooperativeMatrixCheckedInstructionsINTEL,
       "CooperativeMatrixCheckedInstructionsINTEL");
+  add(internal::CapabilityCooperativeMatrixOffsetInstructionsINTEL,
+      "CooperativeMatrixOffsetInstructionsINTEL");
   add(internal::CapabilitySubgroupRequirementsINTEL,
       "SubgroupRequirementsINTEL");
   add(internal::CapabilityTaskSequenceINTEL, "TaskSequenceINTEL");

lib/SPIRV/libSPIRV/SPIRVOpCodeEnumInternal.h

@@ -22,6 +22,10 @@ _SPIRV_OP_INTERNAL(CooperativeMatrixLoadCheckedINTEL,
                    internal::OpCooperativeMatrixLoadCheckedINTEL)
 _SPIRV_OP_INTERNAL(CooperativeMatrixStoreCheckedINTEL,
                    internal::OpCooperativeMatrixStoreCheckedINTEL)
+_SPIRV_OP_INTERNAL(CooperativeMatrixLoadOffsetINTEL,
+                   internal::OpCooperativeMatrixLoadOffsetINTEL)
+_SPIRV_OP_INTERNAL(CooperativeMatrixStoreOffsetINTEL,
+                   internal::OpCooperativeMatrixStoreOffsetINTEL)
 _SPIRV_OP_INTERNAL(CooperativeMatrixConstructCheckedINTEL,
                    internal::OpCooperativeMatrixConstructCheckedINTEL)
 _SPIRV_OP_INTERNAL(CooperativeMatrixApplyFunctionINTEL,

lib/SPIRV/libSPIRV/spirv_internal.hpp

@@ -77,6 +77,8 @@ enum InternalOp {
   IOpCooperativeMatrixLoadCheckedINTEL = 6193,
   IOpCooperativeMatrixStoreCheckedINTEL = 6194,
   IOpCooperativeMatrixConstructCheckedINTEL = 6195,
+  IOpCooperativeMatrixLoadOffsetINTEL = 6212,
+  IOpCooperativeMatrixStoreOffsetINTEL = 6213,
   IOpJointMatrixWorkItemLengthINTEL = 6410,
   IOpTypeTaskSequenceINTEL = 6199,
   IOpComplexFMulINTEL = 6415,
@@ -114,6 +116,7 @@ enum InternalCapability {
   ICapGlobalVariableDecorationsINTEL = 6146,
   ICapabilityTaskSequenceINTEL = 6162,
   ICapabilityCooperativeMatrixCheckedInstructionsINTEL = 6192,
+  ICapabilityCooperativeMatrixOffsetInstructionsINTEL = 6211,
   ICapabilityCooperativeMatrixPrefetchINTEL = 6411,
   ICapabilityComplexFloatMulDivINTEL = 6414,
   ICapabilityTensorFloat32RoundingINTEL = 6425,
@@ -187,6 +190,10 @@ _SPIRV_OP(Op, CooperativeMatrixLoadCheckedINTEL)
 _SPIRV_OP(Op, CooperativeMatrixStoreCheckedINTEL)
 _SPIRV_OP(Op, CooperativeMatrixConstructCheckedINTEL)
 
+_SPIRV_OP(Capability, CooperativeMatrixOffsetInstructionsINTEL)
+_SPIRV_OP(Op, CooperativeMatrixLoadOffsetINTEL)
+_SPIRV_OP(Op, CooperativeMatrixStoreOffsetINTEL)
+
 _SPIRV_OP(Capability, CooperativeMatrixInvocationInstructionsINTEL)
 _SPIRV_OP(Op, CooperativeMatrixApplyFunctionINTEL)