[AMDGPU] Add .uniform_work_group_size metadata to v5

[vangthao]

Summary

Amdgpu kernel with function attribute "uniform-work-group-size"="true" requires
uniform work group size (i.e. each dimension of global size is a multiple of
corresponding dimension of work group size). hipExtModuleLaunchKernel allows to
launch HIP kernel with non-uniform workgroup size, which makes it necessary for
runtime to check and enforce uniform workgroup size if kernel requires it. To
let runtime be able to enforce that, this metadata is needed to indicate that
the kernel requires uniform workgroup size.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp

1087–1090

Don't need to query it twice, just use getFnAttribute once.

I thought there was a helper around for these awful bool string attributes

[vangthao]

Query function attribute once and use getValueAsBool helper function

[arsenm]

llvm/test/CodeGen/AMDGPU/hsa-metadata-uniform-workgroup-size-v5.ll

2

Don't need -verify-machineinstrs

[vangthao]

Remove -verify-machineinstrs from test

[kzhuravl]

lgtm

llvm/docs/AMDGPUUsage.rst

.. table:: AMDHSA Code Object V5 Kernel Metadata Map Additions
      ".uses_dynamic_stack"         boolean                  Indicates if the generated machine code
                                                             is using a dynamically sized stack.
      ".workgroup_processor_mode"   boolean                  (GFX10+) Controls ENABLE_WGP_MODE in
                                                             :ref:`amdgpu-amdhsa-kernel-descriptor-v3-table`.
      ============================= ============= ========== =======================================
 
 ..
 
+  .. table:: AMDHSA Code Object V5 Kernel Attribute Metadata Map
+     :name: amdgpu-amdhsa-code-object-kernel-attribute-metadata-map-v5-table
+
+     =========================== ============== ========= ==============================
+     String Key                  Value Type     Required? Description
+     =========================== ============== ========= ==============================
+     ".uniform_work_group_size"  integer                  Indicates if the kernel
+                                                          requires that each dimension
+                                                          of global size is a multiple
+                                                          of corresponding dimension of
+                                                          work-group size. Value of 1
+                                                          implies true and value of 0
+                                                          implies false. Metadata is
+                                                          only emitted when value is 1.
+     =========================== ============== ========= ==============================
+
+..
+
+..
+
   .. table:: AMDHSA Code Object V5 Kernel Argument Metadata Map Additions and Changes
      :name: amdgpu-amdhsa-code-object-kernel-argument-metadata-map-table-v5
 
      ====================== ============== ========= ================================
      String Key             Value Type     Required? Description
      ====================== ============== ========= ================================
      ".value_kind"          string         Required  Kernel argument kind that
                                                      specifies how to set up the

llvm/lib/BinaryFormat/AMDGPUMetadataVerifier.cpp

bool MetadataVerifier::verifyKernel(msgpack::DocNode &Node) {
   if (!verifyIntegerEntry(KernelMap, ".vgpr_count", true))
     return false;
   if (!verifyIntegerEntry(KernelMap, ".max_flat_workgroup_size", true))
     return false;
   if (!verifyIntegerEntry(KernelMap, ".sgpr_spill_count", false))
     return false;
   if (!verifyIntegerEntry(KernelMap, ".vgpr_spill_count", false))
     return false;
+  if (!verifyIntegerEntry(KernelMap, ".uniform_work_group_size", false))
+    return false;
+
 
   return true;
 }
 
 bool MetadataVerifier::verify(msgpack::DocNode &HSAMetadataRoot) {
   if (!HSAMetadataRoot.isMap())
     return false;
   auto &RootMap = HSAMetadataRoot.getMap();

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h

public:
 
   virtual void emitKernel(const MachineFunction &MF,
                           const SIProgramInfo &ProgramInfo) = 0;
 
 protected:
   virtual void emitVersion() = 0;
   virtual void emitHiddenKernelArgs(const MachineFunction &MF, unsigned &Offset,
                                     msgpack::ArrayDocNode Args) = 0;
+  virtual void emitKernelAttrs(const Function &Func,
+                               msgpack::MapDocNode Kern) = 0;
 };
 
 class MetadataStreamerMsgPackV3 : public MetadataStreamer {
 protected:
   std::unique_ptr<msgpack::Document> HSAMetadataDoc =
       std::make_unique<msgpack::Document>();
 
   void dump(StringRef HSAMetadataString) const;
msgpack::MapDocNode getHSAKernelProps(const MachineFunction &MF,
                                         const SIProgramInfo &ProgramInfo) const;
 
   void emitVersion() override;
 
   void emitPrintf(const Module &Mod);
 
   void emitKernelLanguage(const Function &Func, msgpack::MapDocNode Kern);
 
-  void emitKernelAttrs(const Function &Func, msgpack::MapDocNode Kern);
+  void emitKernelAttrs(const Function &Func, msgpack::MapDocNode Kern) override;
 
   void emitKernelArgs(const MachineFunction &MF, msgpack::MapDocNode Kern);
 
   void emitKernelArg(const Argument &Arg, unsigned &Offset,
                      msgpack::ArrayDocNode Args);
 
   void emitKernelArg(const DataLayout &DL, Type *Ty, Align Alignment,
                      StringRef ValueKind, unsigned &Offset,
void begin(const Module &Mod,
              const IsaInfo::AMDGPUTargetID &TargetID) override;
 };
 
 class MetadataStreamerMsgPackV5 final : public MetadataStreamerMsgPackV4 {
 protected:
   void emitVersion() override;
   void emitHiddenKernelArgs(const MachineFunction &MF, unsigned &Offset,
                             msgpack::ArrayDocNode Args) override;
+  void emitKernelAttrs(const Function &Func, msgpack::MapDocNode Kern) override;
 
 public:
   MetadataStreamerMsgPackV5() = default;
   ~MetadataStreamerMsgPackV5() = default;
 };
 
 // TODO: Rename MetadataStreamerV2 -> MetadataStreamerYamlV2.
 class MetadataStreamerYamlV2 final : public MetadataStreamer {
private:
   }
 
 protected:
   void emitVersion() override;
   void emitHiddenKernelArgs(const MachineFunction &MF, unsigned &Offset,
                             msgpack::ArrayDocNode Args) override {
     llvm_unreachable("Dummy override should not be invoked!");
   }
+  void emitKernelAttrs(const Function &Func,
+                       msgpack::MapDocNode Kern) override {
+    llvm_unreachable("Dummy override should not be invoked!");
+  }
 
 public:
   MetadataStreamerYamlV2() = default;
   ~MetadataStreamerYamlV2() = default;
 
   bool emitTo(AMDGPUTargetStreamer &TargetStreamer) override;
 
   void begin(const Module &Mod,

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp

void MetadataStreamerMsgPackV5::emitHiddenKernelArgs(
   } else {
     Offset += 8; // Skipped.
   }
 
   if (MFI.hasQueuePtr())
     emitKernelArg(DL, Int8PtrTy, Align(8), "hidden_queue_ptr", Offset, Args);
 }
 
+void MetadataStreamerMsgPackV5::emitKernelAttrs(const Function &Func,
+                                                msgpack::MapDocNode Kern) {
+  MetadataStreamerMsgPackV3::emitKernelAttrs(Func, Kern);
+
+  if (Func.getFnAttribute("uniform-work-group-size").getValueAsBool())
+    Kern[".uniform_work_group_size"] = Kern.getDocument()->getNode(1);
+}
+
+
 } // end namespace HSAMD
 } // end namespace AMDGPU
 } // end namespace llvm

llvm/test/CodeGen/AMDGPU/hsa-metadata-uniform-workgroup-size-v5.ll

+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 --amdhsa-code-object-version=5 -filetype=obj -o - < %s | llvm-readelf --notes - | FileCheck %s
+; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx900 -verify-machineinstrs --amdhsa-code-object-version=5 < %s | FileCheck %s
+
+; CHECK: ---
+; CHECK: amdhsa.kernels:
+; CHECK:  - .args:
+; CHECK-LABEL:     .name:           kernel_uniform_workgroup
+; CHECK:     .uniform_work_group_size: 1
+define amdgpu_kernel void @kernel_uniform_workgroup() #0 {
+bb:
+  ret void
+}
+
+; CHECK:  - .args:
+; CHECK-LABEL:     .name:           kernel_non_uniform_workgroup
+; CHECK-NOT:     .uniform_work_group_size:
+define amdgpu_kernel void @kernel_non_uniform_workgroup() #1 {
+bb:
+  ret void
+}
+
+; CHECK:  - .args:
+; CHECK-LABEL:     .name:           kernel_no_attr
+; CHECK-NOT:     .uniform_work_group_size:
+define amdgpu_kernel void @kernel_no_attr() {
+bb:
+  ret void
+}
+attributes #0 = { "uniform-work-group-size"="true" }
+attributes #1 = { "uniform-work-group-size"="false" }