[llvm] [SPIR-V] reqd_work_group_size and work_group_size_hint metadata are correctly converted to the LocalSize and LocalSizeHint execution mode (PR #92552)

[Vyacheslav Levytskyy via llvm-commits]

https://github.com/VyacheslavLevytskyy created https://github.com/llvm/llvm-project/pull/92552

The goal of this PR is to ensure that reqd_work_group_size and work_group_size_hint metadata are correctly converted to the LocalSize and LocalSizeHint execution mode.

>From 7f9ec9adb6b831518de454fc1fbd5d21615564bf Mon Sep 17 00:00:00 2001
From: "Levytskyy, Vyacheslav" <vyacheslav.levytskyy at intel.com>
Date: Fri, 17 May 2024 07:15:20 -0700
Subject: [PATCH] reqd_work_group_size metadata is correctly converted to the
 LocalSize execution mode

---
 llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp     | 20 +++++++----
 .../execution-mode-reqd_work_group_size.ll    | 35 +++++++++++++++++++
 2 files changed, 49 insertions(+), 6 deletions(-)
 create mode 100644 llvm/test/CodeGen/SPIRV/execution-mode-reqd_work_group_size.ll

diff --git a/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp b/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp
index 2ebe5bdc47715..23dd11cf83bed 100644
--- a/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp
@@ -69,7 +69,8 @@ class SPIRVAsmPrinter : public AsmPrinter {
   void outputOpFunctionEnd();
   void outputExtFuncDecls();
   void outputExecutionModeFromMDNode(Register Reg, MDNode *Node,
-                                     SPIRV::ExecutionMode::ExecutionMode EM);
+                                     SPIRV::ExecutionMode::ExecutionMode EM,
+                                     unsigned ExpectMDOps, int64_t DefVal);
   void outputExecutionModeFromNumthreadsAttribute(
       const Register &Reg, const Attribute &Attr,
       SPIRV::ExecutionMode::ExecutionMode EM);
@@ -425,12 +426,19 @@ static void addOpsFromMDNode(MDNode *MDN, MCInst &Inst,
 }
 
 void SPIRVAsmPrinter::outputExecutionModeFromMDNode(
-    Register Reg, MDNode *Node, SPIRV::ExecutionMode::ExecutionMode EM) {
+    Register Reg, MDNode *Node, SPIRV::ExecutionMode::ExecutionMode EM,
+    unsigned ExpectMDOps, int64_t DefVal) {
   MCInst Inst;
   Inst.setOpcode(SPIRV::OpExecutionMode);
   Inst.addOperand(MCOperand::createReg(Reg));
   Inst.addOperand(MCOperand::createImm(static_cast<unsigned>(EM)));
   addOpsFromMDNode(Node, Inst, MAI);
+  // reqd_work_group_size and work_group_size_hint require 3 operands,
+  // if metadata contains less operands, just add a default value
+  unsigned NodeSz = Node->getNumOperands();
+  if (ExpectMDOps > 0 && NodeSz < ExpectMDOps)
+    for (unsigned i = NodeSz; i < ExpectMDOps; ++i)
+      Inst.addOperand(MCOperand::createImm(DefVal));
   outputMCInst(Inst);
 }
 
@@ -476,17 +484,17 @@ void SPIRVAsmPrinter::outputExecutionMode(const Module &M) {
     Register FReg = MAI->getFuncReg(&F);
     assert(FReg.isValid());
     if (MDNode *Node = F.getMetadata("reqd_work_group_size"))
-      outputExecutionModeFromMDNode(FReg, Node,
-                                    SPIRV::ExecutionMode::LocalSize);
+      outputExecutionModeFromMDNode(FReg, Node, SPIRV::ExecutionMode::LocalSize,
+                                    3, 1);
     if (Attribute Attr = F.getFnAttribute("hlsl.numthreads"); Attr.isValid())
       outputExecutionModeFromNumthreadsAttribute(
           FReg, Attr, SPIRV::ExecutionMode::LocalSize);
     if (MDNode *Node = F.getMetadata("work_group_size_hint"))
       outputExecutionModeFromMDNode(FReg, Node,
-                                    SPIRV::ExecutionMode::LocalSizeHint);
+                                    SPIRV::ExecutionMode::LocalSizeHint, 3, 1);
     if (MDNode *Node = F.getMetadata("intel_reqd_sub_group_size"))
       outputExecutionModeFromMDNode(FReg, Node,
-                                    SPIRV::ExecutionMode::SubgroupSize);
+                                    SPIRV::ExecutionMode::SubgroupSize, 0, 0);
     if (MDNode *Node = F.getMetadata("vec_type_hint")) {
       MCInst Inst;
       Inst.setOpcode(SPIRV::OpExecutionMode);
diff --git a/llvm/test/CodeGen/SPIRV/execution-mode-reqd_work_group_size.ll b/llvm/test/CodeGen/SPIRV/execution-mode-reqd_work_group_size.ll
new file mode 100644
index 0000000000000..6e36b0bd5b9dc
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/execution-mode-reqd_work_group_size.ll
@@ -0,0 +1,35 @@
+; From Khronos Translator's test case: test/reqd_work_group_size_md.ll
+
+; The purpose of this test is to check that the reqd_work_group_size metadata
+; is correctly converted to the LocalSize execution mode for the kernels it is
+; applied to.
+
+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; CHECK: OpMemoryModel
+; CHECK-DAG: OpEntryPoint Kernel %[[#ENTRY1:]] "test1"
+; CHECK-DAG: OpEntryPoint Kernel %[[#ENTRY2:]] "test2"
+; CHECK-DAG: OpEntryPoint Kernel %[[#ENTRY3:]] "test3"
+; CHECK-DAG: OpExecutionMode %[[#ENTRY1]] LocalSize 1 2 3
+; CHECK-DAG: OpExecutionMode %[[#ENTRY2]] LocalSize 2 3 1
+; CHECK-DAG: OpExecutionMode %[[#ENTRY3]] LocalSize 3 1 1
+
+define spir_kernel void @test1() !reqd_work_group_size !1 {
+entry:
+  ret void
+}
+
+define spir_kernel void @test2() !reqd_work_group_size !2 {
+entry:
+  ret void
+}
+
+define spir_kernel void @test3() !reqd_work_group_size !3 {
+entry:
+  ret void
+}
+
+!1 = !{i32 1, i32 2, i32 3}
+!2 = !{i32 2, i32 3}
+!3 = !{i32 3}