Allow variable number of operands in work-group metadata (intel#1726)

The reqd_work_group_size and work_group_size_hint metadata nodes are currently expected to have exactly 3 operands which are in turn translated to the 3 operands of their corresponding SPIR-V execution modes. This commit relaxes this requirement by allowing the metadata to have 1-3 operands, where the missing dimensions are padded with 1's in the corresponding execution mode. Note that the information about the dimensionality of the original metadata is lost after translation, so reverse translation will always result in the metadata having all three operands.

Motivation comes from intel#7450 where instead of doing the implicit padding during LLVM IR generation we want to simplify the metadata to only have the operands specified. Not only does it relax the requirements, it also lets us preserve information about dimensionality up until it is translated into SPIR-V.

Author: steffenlarsen

llvm-spirv/lib/SPIRV/OCLUtil.cpp

@@ -791,12 +791,15 @@ unsigned getOCLVersion(Module *M, bool AllowMulti) {
   return encodeOCLVer(Ver.first, Ver.second, 0);
 }
 
-void decodeMDNode(MDNode *N, unsigned &X, unsigned &Y, unsigned &Z) {
+SmallVector<unsigned, 3> decodeMDNode(MDNode *N) {
   if (N == NULL)
-    return;
-  X = getMDOperandAsInt(N, 0);
-  Y = getMDOperandAsInt(N, 1);
-  Z = getMDOperandAsInt(N, 2);
+    return {};
+  size_t NumOperands = N->getNumOperands();
+  SmallVector<unsigned, 3> ReadVals;
+  ReadVals.reserve(NumOperands);
+  for (unsigned I = 0; I < NumOperands; ++I)
+    ReadVals.push_back(getMDOperandAsInt(N, I));
+  return ReadVals;
 }
 
 /// Encode LLVM type by SPIR-V execution mode VecTypeHint

llvm-spirv/lib/SPIRV/OCLUtil.h

@@ -444,7 +444,7 @@ std::tuple<unsigned short, unsigned char, unsigned char>
 decodeOCLVer(unsigned Ver);
 
 /// Decode a MDNode assuming it contains three integer constants.
-void decodeMDNode(MDNode *N, unsigned &X, unsigned &Y, unsigned &Z);
+SmallVector<unsigned, 3> decodeMDNode(MDNode *N);
 
 /// Get full path from debug info metadata
 /// Return empty string if the path is not available.

llvm-spirv/lib/SPIRV/PreprocessMetadata.cpp

@@ -129,27 +129,30 @@ void PreprocessMetadataBase::visit(Module *M) {
 
     // !{void (i32 addrspace(1)*)* @kernel, i32 17, i32 X, i32 Y, i32 Z}
     if (MDNode *WGSize = Kernel.getMetadata(kSPIR2MD::WGSize)) {
-      unsigned X, Y, Z;
-      decodeMDNode(WGSize, X, Y, Z);
+      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(X)
-          .add(Y)
-          .add(Z)
+          .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)) {
-      unsigned X, Y, Z;
-      decodeMDNode(WGSizeHint, X, Y, Z);
+      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(X)
-          .add(Y)
-          .add(Z)
+          .add(DecodedVals[0])
+          .add(DecodedVals.size() >= 2 ? DecodedVals[1] : 1)
+          .add(DecodedVals.size() == 3 ? DecodedVals[2] : 1)
           .done();
     }
 
@@ -175,14 +178,16 @@ void PreprocessMetadataBase::visit(Module *M) {
     //         i32 Y, i32 Z}
     if (MDNode *MaxWorkgroupSizeINTEL =
             Kernel.getMetadata(kSPIR2MD::MaxWGSize)) {
-      unsigned X, Y, Z;
-      decodeMDNode(MaxWorkgroupSizeINTEL, X, Y, Z);
+      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(X)
-          .add(Y)
-          .add(Z)
+          .add(DecodedVals[0])
+          .add(DecodedVals[1])
+          .add(DecodedVals[2])
           .done();
     }
 

llvm-spirv/test/reqd_work_group_size_md.ll

@@ -0,0 +1,37 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc -o %t.spv
+; RUN: llvm-spirv -to-text %t.spv -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+;
+; 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.
+;
+; CHECK-SPIRV: EntryPoint 6 [[TEST1:[0-9]+]] "test1"
+; CHECK-SPIRV: EntryPoint 6 [[TEST2:[0-9]+]] "test2"
+; CHECK-SPIRV: EntryPoint 6 [[TEST3:[0-9]+]] "test3"
+; CHECK-SPIRV: ExecutionMode [[TEST1]] 17 1 2 3
+; CHECK-SPIRV: ExecutionMode [[TEST2]] 17 2 3 1
+; CHECK-SPIRV: ExecutionMode [[TEST3]] 17 3 1 1
+
+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-unknown"
+
+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-spirv/test/work_group_size_hint_md.ll

@@ -0,0 +1,37 @@
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-spirv %t.bc -o %t.spv
+; RUN: llvm-spirv -to-text %t.spv -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+;
+; The purpose of this test is to check that the work_group_size_hint metadata
+; is correctly converted to the LocalSizeHint execution mode for the kernels it
+; is applied to.
+;
+; CHECK-SPIRV: EntryPoint 6 [[TEST1:[0-9]+]] "test1"
+; CHECK-SPIRV: EntryPoint 6 [[TEST2:[0-9]+]] "test2"
+; CHECK-SPIRV: EntryPoint 6 [[TEST3:[0-9]+]] "test3"
+; CHECK-SPIRV: ExecutionMode [[TEST1]] 18 1 2 3
+; CHECK-SPIRV: ExecutionMode [[TEST2]] 18 2 3 1
+; CHECK-SPIRV: ExecutionMode [[TEST3]] 18 3 1 1
+
+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-unknown"
+
+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}