From 0e704312619eb3bd9fb01224aaaf062683ee112e Mon Sep 17 00:00:00 2001
From: Sven van Haastregt <sven.vanhaastregt@arm.com>
Date: Wed, 11 Dec 2024 10:13:47 +0100
Subject: [PATCH] [Backport to 15] [SPIR-V 1.1/1.2] SPIRVReader: Add
 SubgroupsPerWorkgroup(Id) (#2916)

Add support for consuming the SubgroupsPerWorkgroup (SPIR-V 1.1) and
SubgroupsPerWorkgroupId (SPIR-V 1.2) execution modes.  Map both of
these to `spirv.ExecutionMode` named metadata.

(cherry picked from commit 98cd3e46bd62bd0e798a3c557f2ce2fe99bb51c9)
---
 lib/SPIRV/SPIRVReader.cpp           | 19 +++++++++++++++++++
 lib/SPIRV/libSPIRV/SPIRVEntry.cpp   |  2 ++
 test/SubgroupsPerWorkgroup.spvasm   | 24 ++++++++++++++++++++++++
 test/SubgroupsPerWorkgroupId.spvasm | 27 +++++++++++++++++++++++++++
 4 files changed, 72 insertions(+)
 create mode 100644 test/SubgroupsPerWorkgroup.spvasm
 create mode 100644 test/SubgroupsPerWorkgroupId.spvasm

diff --git a/lib/SPIRV/SPIRVReader.cpp b/lib/SPIRV/SPIRVReader.cpp
index c052f91ad..386ca93b9 100644
--- a/lib/SPIRV/SPIRVReader.cpp
+++ b/lib/SPIRV/SPIRVReader.cpp
@@ -4399,6 +4399,25 @@ bool SPIRVToLLVM::transMetadata() {
       auto SizeMD = ConstantAsMetadata::get(getUInt32(M, EM->getLiterals()[0]));
       F->setMetadata(kSPIR2MD::SubgroupSize, MDNode::get(*Context, SizeMD));
     }
+    // Generate metadata for SubgroupsPerWorkgroup/SubgroupsPerWorkgroupId.
+    auto EmitSubgroupsPerWorkgroupMD = [this, F](SPIRVExecutionModeKind EMK,
+                                                 uint64_t Value) {
+      NamedMDNode *ExecModeMD =
+          M->getOrInsertNamedMetadata(kSPIRVMD::ExecutionMode);
+      SmallVector<Metadata *, 2> OperandVec;
+      OperandVec.push_back(ConstantAsMetadata::get(F));
+      OperandVec.push_back(ConstantAsMetadata::get(getUInt32(M, EMK)));
+      OperandVec.push_back(ConstantAsMetadata::get(getUInt32(M, Value)));
+      ExecModeMD->addOperand(MDNode::get(*Context, OperandVec));
+    };
+    if (auto *EM = BF->getExecutionMode(ExecutionModeSubgroupsPerWorkgroup)) {
+      EmitSubgroupsPerWorkgroupMD(EM->getExecutionMode(), EM->getLiterals()[0]);
+    } else if (auto *EM = BF->getExecutionModeId(
+                   ExecutionModeSubgroupsPerWorkgroupId)) {
+      if (auto Val = transIdAsConstant(EM->getLiterals()[0])) {
+        EmitSubgroupsPerWorkgroupMD(EM->getExecutionMode(), *Val);
+      }
+    }
     // Generate metadata for max_work_group_size
     if (auto EM = BF->getExecutionMode(ExecutionModeMaxWorkgroupSizeINTEL)) {
       F->setMetadata(kSPIR2MD::MaxWGSize,
diff --git a/lib/SPIRV/libSPIRV/SPIRVEntry.cpp b/lib/SPIRV/libSPIRV/SPIRVEntry.cpp
index 57c2c2d29..5e5d492a6 100644
--- a/lib/SPIRV/libSPIRV/SPIRVEntry.cpp
+++ b/lib/SPIRV/libSPIRV/SPIRVEntry.cpp
@@ -636,6 +636,8 @@ void SPIRVExecutionMode::decode(std::istream &I) {
   case ExecutionModeSharedLocalMemorySizeINTEL:
   case ExecutionModeNamedBarrierCountINTEL:
   case ExecutionModeSubgroupSize:
+  case ExecutionModeSubgroupsPerWorkgroup:
+  case ExecutionModeSubgroupsPerWorkgroupId:
   case ExecutionModeMaxWorkDimINTEL:
   case ExecutionModeNumSIMDWorkitemsINTEL:
   case ExecutionModeSchedulerTargetFmaxMhzINTEL:
diff --git a/test/SubgroupsPerWorkgroup.spvasm b/test/SubgroupsPerWorkgroup.spvasm
new file mode 100644
index 000000000..94fa7ade6
--- /dev/null
+++ b/test/SubgroupsPerWorkgroup.spvasm
@@ -0,0 +1,24 @@
+; REQUIRES: spirv-as
+
+; RUN: spirv-as %s --target-env spv1.2 -o %t.spv
+; RUN: spirv-val %t.spv
+; RUN: llvm-spirv -r -o %t.rev.bc %t.spv
+; RUN: llvm-dis %t.rev.bc -o - | FileCheck %s
+
+               OpCapability Addresses
+               OpCapability Linkage
+               OpCapability Kernel
+               OpCapability SubgroupDispatch
+               OpMemoryModel Physical64 OpenCL
+               OpEntryPoint Kernel %fn "testSubgroupsPerWorkgroup"
+               OpExecutionMode %fn SubgroupsPerWorkgroup 8
+       %void = OpTypeVoid
+       %fnTy = OpTypeFunction %void
+
+; CHECK: !spirv.ExecutionMode = !{![[MD:[0-9]+]]}
+; CHECK: ![[MD]] = !{void ()* @testSubgroupsPerWorkgroup, i32 36, i32 8}
+
+         %fn = OpFunction %void None %fnTy
+      %entry = OpLabel
+               OpReturn
+               OpFunctionEnd
diff --git a/test/SubgroupsPerWorkgroupId.spvasm b/test/SubgroupsPerWorkgroupId.spvasm
new file mode 100644
index 000000000..09054da79
--- /dev/null
+++ b/test/SubgroupsPerWorkgroupId.spvasm
@@ -0,0 +1,27 @@
+; REQUIRES: spirv-as
+
+; RUN: spirv-as %s --target-env spv1.2 -o %t.spv
+; RUN: spirv-val %t.spv
+; RUN: llvm-spirv -r -o %t.rev.bc %t.spv
+; RUN: llvm-dis %t.rev.bc -o - | FileCheck %s
+
+               OpCapability Addresses
+               OpCapability Linkage
+               OpCapability Kernel
+               OpCapability SubgroupDispatch
+               OpMemoryModel Physical64 OpenCL
+               OpEntryPoint Kernel %fn "testSubgroupsPerWorkgroupId"
+               OpExecutionModeId %fn SubgroupsPerWorkgroupId %uint_8
+       %void = OpTypeVoid
+       %uint = OpTypeInt 32 0
+     %uint_4 = OpConstant %uint 4
+     %uint_8 = OpSpecConstantOp %uint IAdd %uint_4 %uint_4
+       %fnTy = OpTypeFunction %void
+
+; CHECK: !spirv.ExecutionMode = !{![[MD:[0-9]+]]}
+; CHECK: ![[MD]] = !{void ()* @testSubgroupsPerWorkgroupId, i32 37, i32 8}
+
+         %fn = OpFunction %void None %fnTy
+      %entry = OpLabel
+               OpReturn
+               OpFunctionEnd