[SPIR-V] reqd_work_group_size and work_group_size_hint metadata are correctly converted to the LocalSize and LocalSizeHint execution mode (#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.

reqd_work_group_size and work_group_size_hint require 3 operands (see
https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#Execution_Mode),
if metadata contains less operands, just add a default value (1).

Author: VyacheslavLevytskyy

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);
@@ -422,12 +423,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);
 }
 
@@ -473,17 +481,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);

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}

llvm/test/CodeGen/SPIRV/execution-mode-work_group_size_hint.ll

@@ -0,0 +1,34 @@
+; From Khronos Translator's test case: test/reqd_work_group_size_md.ll
+
+; The purpose of this test is to check that the work_group_size_hint metadata
+; is correctly converted to the LocalSizeHint execution mode.
+
+; 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]] LocalSizeHint 1 2 3
+; CHECK-DAG: OpExecutionMode %[[#ENTRY2]] LocalSizeHint 2 3 1
+; CHECK-DAG: OpExecutionMode %[[#ENTRY3]] LocalSizeHint 3 1 1
+
+define spir_kernel void @test1() !work_group_size_hint !1 {
+entry:
+  ret void
+}
+
+define spir_kernel void @test2() !work_group_size_hint !2 {
+entry:
+  ret void
+}
+
+define spir_kernel void @test3() !work_group_size_hint !3 {
+entry:
+  ret void
+}
+
+!1 = !{i32 1, i32 2, i32 3}
+!2 = !{i32 2, i32 3}
+!3 = !{i32 3}