From 1b4c8446bcb3c6f49fa057c8f2fd443b14673c44 Mon Sep 17 00:00:00 2001
From: "Plyakhin, Yury" <yury.plyakhin@intel.com>
Date: Fri, 11 Oct 2024 01:20:20 +0200
Subject: [PATCH] draft

---
 include/LLVMSPIRVExtensions.inc       |  1 +
 lib/SPIRV/SPIRVUtil.cpp               | 24 +++++++++++++++++
 lib/SPIRV/libSPIRV/SPIRVInstruction.h | 39 +++++++++++++++++++++++++++
 3 files changed, 64 insertions(+)

diff --git a/include/LLVMSPIRVExtensions.inc b/include/LLVMSPIRVExtensions.inc
index f8edd290f..df33c7a32 100644
--- a/include/LLVMSPIRVExtensions.inc
+++ b/include/LLVMSPIRVExtensions.inc
@@ -74,3 +74,4 @@ EXT(SPV_INTEL_subgroup_requirements)
 EXT(SPV_INTEL_task_sequence)
 EXT(SPV_INTEL_maximum_registers)
 EXT(SPV_INTEL_bindless_images)
+EXT(SPV_INTEL_2d_block_io)
diff --git a/lib/SPIRV/SPIRVUtil.cpp b/lib/SPIRV/SPIRVUtil.cpp
index 4b0f721da..9f5d3579e 100644
--- a/lib/SPIRV/SPIRVUtil.cpp
+++ b/lib/SPIRV/SPIRVUtil.cpp
@@ -2518,6 +2518,30 @@ class SPIRVFriendlyIRMangleInfo : public BuiltinFuncMangleInfo {
     case internal::OpConvertHandleToSampledImageINTEL:
       addUnsignedArg(0);
       break;
+    case internal::OpSubgroup2DBlockLoadINTEL:
+    case internal::OpSubgroup2DBlockLoadTransposeINTEL:
+    case internal::OpSubgroup2DBlockLoadTransformINTEL:
+      addUnsignedArgs(0, 3);
+      setArgAttr(4, SPIR::ATTR_GLOBAL);
+      addVoidPtrArg(4);
+      addUnsignedArgs(5, 8);
+      setArgAttr(9, SPIR::ATTR_PRIVATE);
+      addVoidPtrArg(9);
+      break;
+    case internal::OpSubgroup2DBlockPrefetchINTEL:
+      addUnsignedArgs(0, 3);
+      setArgAttr(4, SPIR::ATTR_GLOBAL);
+      addVoidPtrArg(4);
+      addUnsignedArgs(5, 8);
+      break;
+    case internal::OpSubgroup2DBlockStoreINTEL:
+      addUnsignedArgs(0, 3);
+      setArgAttr(4, SPIR::ATTR_PRIVATE);
+      addVoidPtrArg(4);
+      setArgAttr(5, SPIR::ATTR_GLOBAL);
+      addVoidPtrArg(5);
+      addUnsignedArgs(6, 9);
+      break;
     default:;
       // No special handling is needed
     }
diff --git a/lib/SPIRV/libSPIRV/SPIRVInstruction.h b/lib/SPIRV/libSPIRV/SPIRVInstruction.h
index e1e9b98b5..2ffc199be 100644
--- a/lib/SPIRV/libSPIRV/SPIRVInstruction.h
+++ b/lib/SPIRV/libSPIRV/SPIRVInstruction.h
@@ -4195,5 +4195,44 @@ _SPIRV_OP(ConvertHandleToSamplerINTEL)
 _SPIRV_OP(ConvertHandleToSampledImageINTEL)
 #undef _SPIRV_OP
 
+class SPIRVSubgroup2DBlockIOINTELInstBase : public SPIRVInstTemplateBase {
+public:
+  std::optional<ExtensionID> getRequiredExtension() const override {
+    return ExtensionID::SPV_INTEL_2d_block_io;
+  }
+  SPIRVCapVec getRequiredCapability() const override {
+    return getVec(internal::Subgroup2DBlockIOINTEL);
+  }
+};
+
+class SPIRVSubgroup2DBlockLoadINTELInst : public SPIRVSubgroup2DBlockIOINTELInstBase {
+protected:
+  void validate() const override {
+    SPIRVInstruction::validate();
+    std::string InstName = "Subgroup2DBlockLoadINTEL";
+    SPIRVErrorLog &SPVErrLog = this->getModule()->getErrorLog();
+
+    SPVErrLog.checkError(
+        this->isOperandLiteral(0), SPIRVEC_InvalidInstruction,
+        InstName + "\nElement Size must be a constant instructions with scalar 32-bit integer type\n");
+    SPVErrLog.checkError(
+        this->isOperandLiteral(1), SPIRVEC_InvalidInstruction,
+        InstName + "\nBlock Width must be a constant instructions with scalar 32-bit integer type\n");
+    SPVErrLog.checkError(
+        this->isOperandLiteral(2), SPIRVEC_InvalidInstruction,
+        InstName + "\nBlock Height must be a constant instructions with scalar 32-bit integer type\n");
+    SPVErrLog.checkError(
+        this->isOperandLiteral(3), SPIRVEC_InvalidInstruction,
+        InstName + "\nElement Size must be a constant instructions with scalar 32-bit integer type\n");
+  }
+};
+
+#define _SPIRV_OP(x, ...)                                                      \
+  typedef SPIRVInstTemplate<SPIRVSubgroup2DBlockLoadINTELInst,                 \
+                            internal::Op##x##INTEL, __VA_ARGS__>               \
+      SPIRV##x##INTEL;
+_SPIRV_OP(Subgroup2DBlockLoad, true, 11)
+#undef _SPIRV_OP
+
 } // namespace SPIRV
 #endif // SPIRV_LIBSPIRV_SPIRVINSTRUCTION_H