From 420f5dd41a21d5a7e5429a88e8c70182d6792e90 Mon Sep 17 00:00:00 2001
From: Sven van Haastregt <sven.vanhaastregt@arm.com>
Date: Thu, 5 Dec 2024 12:20:00 +0100
Subject: [PATCH] [SPIR-V 1.2] SPIRVReader: Support LocalSizeHintId (#2907)

If there is no `OpExecutionMode .. LocalSizeHint` in the input, see if
there is an `OpExecutionModeId .. LocalSizeHintId` and take the value
for the `work_group_size_hint` metadata from the referenced constants
instead.

Once `LocalSizeHintId` has been translated to LLVM IR, it is
indistinguishable from a (non-ID) `LocalSizeHint` execution mode.
---
 lib/SPIRV/SPIRVReader.cpp         | 10 ++++++++++
 lib/SPIRV/libSPIRV/SPIRVEntry.cpp |  1 +
 test/LocalSizeHintId.spvasm       | 27 +++++++++++++++++++++++++++
 3 files changed, 38 insertions(+)
 create mode 100644 test/LocalSizeHintId.spvasm

diff --git a/lib/SPIRV/SPIRVReader.cpp b/lib/SPIRV/SPIRVReader.cpp
index 584dd894c..8e8877519 100644
--- a/lib/SPIRV/SPIRVReader.cpp
+++ b/lib/SPIRV/SPIRVReader.cpp
@@ -4573,6 +4573,16 @@ bool SPIRVToLLVM::transMetadata() {
     if (auto *EM = BF->getExecutionMode(ExecutionModeLocalSizeHint)) {
       F->setMetadata(kSPIR2MD::WGSizeHint,
                      getMDNodeStringIntVec(Context, EM->getLiterals()));
+    } else if (auto *EM =
+                   BF->getExecutionModeId(ExecutionModeLocalSizeHintId)) {
+      std::vector<SPIRVWord> Values;
+      for (const auto Id : EM->getLiterals()) {
+        if (auto Val = transIdAsConstant(Id)) {
+          Values.emplace_back(static_cast<SPIRVWord>(*Val));
+        }
+      }
+      F->setMetadata(kSPIR2MD::WGSizeHint,
+                     getMDNodeStringIntVec(Context, Values));
     }
     // Generate metadata for vec_type_hint
     if (auto *EM = BF->getExecutionMode(ExecutionModeVecTypeHint)) {
diff --git a/lib/SPIRV/libSPIRV/SPIRVEntry.cpp b/lib/SPIRV/libSPIRV/SPIRVEntry.cpp
index 459109b4c..5a18cfd7e 100644
--- a/lib/SPIRV/libSPIRV/SPIRVEntry.cpp
+++ b/lib/SPIRV/libSPIRV/SPIRVEntry.cpp
@@ -658,6 +658,7 @@ void SPIRVExecutionMode::decode(std::istream &I) {
   case ExecutionModeLocalSize:
   case ExecutionModeLocalSizeId:
   case ExecutionModeLocalSizeHint:
+  case ExecutionModeLocalSizeHintId:
   case ExecutionModeMaxWorkgroupSizeINTEL:
     WordLiterals.resize(3);
     break;
diff --git a/test/LocalSizeHintId.spvasm b/test/LocalSizeHintId.spvasm
new file mode 100644
index 000000000..c6beebd74
--- /dev/null
+++ b/test/LocalSizeHintId.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
+               OpMemoryModel Physical64 OpenCL
+               OpEntryPoint Kernel %fn "testLocalSizeHintId"
+               OpExecutionModeId %fn LocalSizeHintId %uint_64 %uint_1 %uint_1sco
+       %void = OpTypeVoid
+       %uint = OpTypeInt 32 0
+     %uint_1 = OpConstant %uint 1
+    %uint_64 = OpConstant %uint 64
+  %uint_1sco = OpSpecConstantOp %uint UDiv %uint_64 %uint_64
+       %fnTy = OpTypeFunction %void
+
+; CHECK: define spir_kernel void @testLocalSizeHintId() {{.*}} !work_group_size_hint ![[MD:[0-9]+]]
+; CHECK: ![[MD]] = !{i32 64, i32 1, i32 1}
+
+         %fn = OpFunction %void None %fnTy
+      %entry = OpLabel
+               OpReturn
+               OpFunctionEnd