Relax operand constraints on max_work_group_size metadata (#2711)

This allows `max_work_group_size` metadata to behave like
`reqd_work_group_size` and `work_group_size_hint` in that it may have
between 1 and 3 operands. Missing dimensions are filled in with 1s.

Original commit:
KhronosGroup/SPIRV-LLVM-Translator@c38378cd5d814ed

Author: frasercrmck

llvm-spirv/lib/SPIRV/PreprocessMetadata.cpp

@@ -129,32 +129,28 @@ void PreprocessMetadataBase::visit(Module *M) {
     // of ExecutionMode instructions.
 
     // !{void (i32 addrspace(1)*)* @kernel, i32 17, i32 X, i32 Y, i32 Z}
-    if (MDNode *WGSize = Kernel.getMetadata(kSPIR2MD::WGSize)) {
-      assert(WGSize->getNumOperands() >= 1 && WGSize->getNumOperands() <= 3 &&
-             "reqd_work_group_size does not have between 1 and 3 operands.");
-      SmallVector<unsigned, 3> DecodedVals = decodeMDNode(WGSize);
-      EM.addOp()
-          .add(&Kernel)
-          .add(spv::ExecutionModeLocalSize)
-          .add(DecodedVals[0])
-          .add(DecodedVals.size() >= 2 ? DecodedVals[1] : 1)
-          .add(DecodedVals.size() == 3 ? DecodedVals[2] : 1)
-          .done();
-    }
-
     // !{void (i32 addrspace(1)*)* @kernel, i32 18, i32 X, i32 Y, i32 Z}
-    if (MDNode *WGSizeHint = Kernel.getMetadata(kSPIR2MD::WGSizeHint)) {
-      assert(WGSizeHint->getNumOperands() >= 1 &&
-             WGSizeHint->getNumOperands() <= 3 &&
-             "work_group_size_hint does not have between 1 and 3 operands.");
-      SmallVector<unsigned, 3> DecodedVals = decodeMDNode(WGSizeHint);
-      EM.addOp()
-          .add(&Kernel)
-          .add(spv::ExecutionModeLocalSizeHint)
-          .add(DecodedVals[0])
-          .add(DecodedVals.size() >= 2 ? DecodedVals[1] : 1)
-          .add(DecodedVals.size() == 3 ? DecodedVals[2] : 1)
-          .done();
+    // !{void (i32 addrspace(1)*)* @kernel, i32 max_work_group_size, i32 X,
+    //         i32 Y, i32 Z}
+    std::pair<unsigned, const char *> WGSizeMDs[3] = {
+        {spv::ExecutionModeLocalSize, kSPIR2MD::WGSize},
+        {spv::ExecutionModeLocalSizeHint, kSPIR2MD::WGSizeHint},
+        {spv::ExecutionModeMaxWorkgroupSizeINTEL, kSPIR2MD::MaxWGSize},
+    };
+
+    for (auto &[ExMode, MDName] : WGSizeMDs) {
+      if (MDNode *WGMD = Kernel.getMetadata(MDName)) {
+        assert(WGMD->getNumOperands() >= 1 && WGMD->getNumOperands() <= 3 &&
+               "work-group metadata does not have between 1 and 3 operands.");
+        SmallVector<unsigned, 3> DecodedVals = decodeMDNode(WGMD);
+        EM.addOp()
+            .add(&Kernel)
+            .add(ExMode)
+            .add(DecodedVals[0])
+            .add(DecodedVals.size() >= 2 ? DecodedVals[1] : 1)
+            .add(DecodedVals.size() == 3 ? DecodedVals[2] : 1)
+            .done();
+      }
     }
 
     // !{void (i32 addrspace(1)*)* @kernel, i32 30, i32 hint}
@@ -184,23 +180,6 @@ void PreprocessMetadataBase::visit(Module *M) {
           .done();
     }
 
-    // !{void (i32 addrspace(1)*)* @kernel, i32 max_work_group_size, i32 X,
-    //         i32 Y, i32 Z}
-    if (MDNode *MaxWorkgroupSizeINTEL =
-            Kernel.getMetadata(kSPIR2MD::MaxWGSize)) {
-      assert(MaxWorkgroupSizeINTEL->getNumOperands() == 3 &&
-             "max_work_group_size does not have 3 operands.");
-      SmallVector<unsigned, 3> DecodedVals =
-          decodeMDNode(MaxWorkgroupSizeINTEL);
-      EM.addOp()
-          .add(&Kernel)
-          .add(spv::ExecutionModeMaxWorkgroupSizeINTEL)
-          .add(DecodedVals[0])
-          .add(DecodedVals[1])
-          .add(DecodedVals[2])
-          .done();
-    }
-
     // !{void (i32 addrspace(1)*)* @kernel, i32 no_global_work_offset}
     if (Kernel.getMetadata(kSPIR2MD::NoGlobalOffset)) {
       EM.addOp().add(&Kernel).add(spv::ExecutionModeNoGlobalOffsetINTEL).done();

llvm-spirv/test/extensions/INTEL/SPV_INTEL_kernel_attributes/max_work_group_size.ll

@@ -0,0 +1,54 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc --spirv-ext=+SPV_INTEL_kernel_attributes -o %t.spv
+; RUN: llvm-spirv %t.spv -to-text -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis < %t.rev.bc | FileCheck %s --check-prefix=CHECK-LLVM
+
+; RUN: llvm-spirv -spirv-text -r %t.spt -o %t.rev.bc
+; RUN: llvm-dis < %t.rev.bc | FileCheck %s --check-prefix=CHECK-LLVM
+
+; CHECK-SPIRV: Capability KernelAttributesINTEL
+; CHECK-SPIRV: Extension "SPV_INTEL_kernel_attributes"
+; CHECK-SPIRV: EntryPoint {{.*}} [[DIM1:[0-9]+]] "Dim1"
+; CHECK-SPIRV: EntryPoint {{.*}} [[DIM2:[0-9]+]] "Dim2"
+; CHECK-SPIRV: EntryPoint {{.*}} [[DIM3:[0-9]+]] "Dim3"
+; CHECK-SPIRV: ExecutionMode [[DIM1]] 5893 4 1 1
+; CHECK-SPIRV: ExecutionMode [[DIM2]] 5893 8 4 1
+; CHECK-SPIRV: ExecutionMode [[DIM3]] 5893 16 8 4
+; CHECK-SPIRV: Function {{.*}} [[DIM1]] {{.*}}
+; CHECK-SPIRV: Function {{.*}} [[DIM2]] {{.*}}
+; CHECK-SPIRV: Function {{.*}} [[DIM3]] {{.*}}
+
+; CHECK-LLVM: define spir_kernel void @Dim1()
+; CHECK-LLVM-SAME: !max_work_group_size ![[MAXWG1:[0-9]+]]
+
+; CHECK-LLVM: define spir_kernel void @Dim2()
+; CHECK-LLVM-SAME: !max_work_group_size ![[MAXWG2:[0-9]+]]
+
+; CHECK-LLVM: define spir_kernel void @Dim3()
+; CHECK-LLVM-SAME: !max_work_group_size ![[MAXWG3:[0-9]+]]
+
+; CHECK-LLVM: ![[MAXWG1]] = !{i32 4, i32 1, i32 1}
+; CHECK-LLVM: ![[MAXWG2]] = !{i32 8, i32 4, i32 1}
+; CHECK-LLVM: ![[MAXWG3]] = !{i32 16, i32 8, i32 4}
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
+target triple = "spir64-unknown-linux"
+
+define spir_kernel void @Dim1() !max_work_group_size !0 {
+  ret void
+}
+
+define spir_kernel void @Dim2() !max_work_group_size !1 {
+  ret void
+}
+
+define spir_kernel void @Dim3() !max_work_group_size !2 {
+  ret void
+}
+
+!0 = !{i32 4}
+!1 = !{i32 8, i32 4}
+!2 = !{i32 16, i32 8, i32 4}