[llvm] [AMDGPU] Defaults for missing dimensions in SYCL required wg size (PR #72652)

[Jakub Chlanda via llvm-commits]

https://github.com/jchlanda created https://github.com/llvm/llvm-project/pull/72652

SYCL allows for required work group to be partially specified (i.e. not all 3 dimensions):
* https://registry.khronos.org/SYCL/specs/sycl-2020/html/sycl-2020.html#sec:kernel.attributes This fails AMDGPU's attribute verification. The patch aims to provide the default values for missing dimensions when dealing with SYCL kernels. Rather than modifying the module's metadata it uses internal data to padd missing values.

>From 1644043deeac8517b0fd61057a97d72129ac3440 Mon Sep 17 00:00:00 2001
From: Jakub Chlanda <jakub at codeplay.com>
Date: Fri, 17 Nov 2023 13:27:45 +0000
Subject: [PATCH] [AMDGPU] Defaults for missing dimensions in SYCL required wg
 size

SYCL allows for required work group to be partially specified (i.e. not
all 3 dimensions):
* https://registry.khronos.org/SYCL/specs/sycl-2020/html/sycl-2020.html#sec:kernel.attributes
This fails AMDGPU's attribute verification. The patch aims to provide
the default values for missing dimensions when dealing with SYCL
kernels. Rather than modifying the module's metadata it uses internal
data to padd missing values.
---
 llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp | 12 ++++++++----
 llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h   |  3 ++-
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp
index b51a876750b58b0..bff30bda357e2b9 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp
@@ -173,14 +173,18 @@ std::string MetadataStreamerMsgPackV4::getTypeName(Type *Ty,
 }
 
 msgpack::ArrayDocNode
-MetadataStreamerMsgPackV4::getWorkGroupDimensions(MDNode *Node) const {
+MetadataStreamerMsgPackV4::getWorkGroupDimensions(const Function &Func,
+                                                  MDNode *Node) const {
   auto Dims = HSAMetadataDoc->getArrayNode();
-  if (Node->getNumOperands() != 3)
+  if (Node->getNumOperands() != 3 && !Func.hasFnAttribute("sycl-module-id"))
     return Dims;
 
   for (auto &Op : Node->operands())
     Dims.push_back(Dims.getDocument()->getNode(
         uint64_t(mdconst::extract<ConstantInt>(Op)->getZExtValue())));
+  for (unsigned I = Dims.size(); I < 3; ++I)
+    Dims.push_back(Dims.getDocument()->getNode(1));
+
   return Dims;
 }
 
@@ -233,9 +237,9 @@ void MetadataStreamerMsgPackV4::emitKernelAttrs(const Function &Func,
                                                 msgpack::MapDocNode Kern) {
 
   if (auto Node = Func.getMetadata("reqd_work_group_size"))
-    Kern[".reqd_workgroup_size"] = getWorkGroupDimensions(Node);
+    Kern[".reqd_workgroup_size"] = getWorkGroupDimensions(Func, Node);
   if (auto Node = Func.getMetadata("work_group_size_hint"))
-    Kern[".workgroup_size_hint"] = getWorkGroupDimensions(Node);
+    Kern[".workgroup_size_hint"] = getWorkGroupDimensions(Func, Node);
   if (auto Node = Func.getMetadata("vec_type_hint")) {
     Kern[".vec_type_hint"] = Kern.getDocument()->getNode(
         getTypeName(
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h
index 18a7b5d7a9633e8..3214f096f27b9a9 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h
@@ -81,7 +81,8 @@ class MetadataStreamerMsgPackV4 : public MetadataStreamer {
 
   std::string getTypeName(Type *Ty, bool Signed) const;
 
-  msgpack::ArrayDocNode getWorkGroupDimensions(MDNode *Node) const;
+  msgpack::ArrayDocNode getWorkGroupDimensions(const Function &Func,
+                                               MDNode *Node) const;
 
   msgpack::MapDocNode getHSAKernelProps(const MachineFunction &MF,
                                         const SIProgramInfo &ProgramInfo,