UR_KERNEL_SUB_GROUP_INFO_SUB_GROUP_SIZE_INTEL on Cuda and HIP (#17137)

For HIP the value of sub group size can either be 32 or 64, it can be
retrieved from `intel_reqd_sub_group_size` metadata node.

Cuda only supports 32, which is enforced in the compiler, see
[SemaSYCL::addIntelReqdSubGroupSizeAttr](https://github.com/intel/llvm/blob/sycl/clang/lib/Sema/SemaSYCLDeclAttr.cpp#L828).

---------

Co-authored-by: Nicolas Miller <[email protected]>

Author: 2 people

source/adapters/cuda/kernel.cpp

@@ -339,10 +339,15 @@ urKernelGetSubGroupInfo(ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
     return ReturnValue(0);
   }
   case UR_KERNEL_SUB_GROUP_INFO_SUB_GROUP_SIZE_INTEL: {
-    // Return value of 0 => unspecified or "auto" sub-group size
-    // Correct for now, since warp size may be read from special register
-    // TODO: Return warp size once default is primary sub-group size
-    // TODO: Revisit if we can recover [[sub_group_size]] attribute from PTX
+    const auto &KernelReqdSubGroupSizeMap =
+        hKernel->getProgram()->KernelReqdSubGroupSizeMD;
+    // If present, return the value of intel_reqd_sub_group_size metadata, if
+    // not: 0, which stands for unspecified or auto sub-group size.
+    if (auto KernelReqdSubGroupSize =
+            KernelReqdSubGroupSizeMap.find(hKernel->getName());
+        KernelReqdSubGroupSize != KernelReqdSubGroupSizeMap.end())
+      return ReturnValue(KernelReqdSubGroupSize->second);
+
     return ReturnValue(0);
   }
   default:

source/adapters/cuda/program.cpp

@@ -98,6 +98,9 @@ ur_program_handle_t_::setMetadata(const ur_program_metadata_t *Metadata,
     } else if (Tag ==
                __SYCL_UR_PROGRAM_METADATA_TAG_MAX_LINEAR_WORK_GROUP_SIZE) {
       KernelMaxLinearWorkGroupSizeMD[Prefix] = MetadataElement.value.data64;
+    } else if (Tag == __SYCL_UR_PROGRAM_METADATA_TAG_REQD_SUB_GROUP_SIZE) {
+      assert(MetadataElement.type == UR_PROGRAM_METADATA_TYPE_UINT32);
+      KernelReqdSubGroupSizeMD[Prefix] = MetadataElement.value.data32;
     }
   }
   return UR_RESULT_SUCCESS;

source/adapters/cuda/program.hpp

@@ -39,6 +39,7 @@ struct ur_program_handle_t_ {
   std::unordered_map<std::string, std::tuple<uint32_t, uint32_t, uint32_t>>
       KernelMaxWorkGroupSizeMD;
   std::unordered_map<std::string, uint64_t> KernelMaxLinearWorkGroupSizeMD;
+  std::unordered_map<std::string, uint32_t> KernelReqdSubGroupSizeMD;
 
   constexpr static size_t MaxLogSize = 8192u;
 
@@ -49,7 +50,8 @@ struct ur_program_handle_t_ {
   ur_program_handle_t_(ur_context_handle_t Context, ur_device_handle_t Device)
       : Module{nullptr}, Binary{}, BinarySizeInBytes{0}, RefCount{1},
         Context{Context}, Device{Device}, KernelReqdWorkGroupSizeMD{},
-        KernelMaxWorkGroupSizeMD{}, KernelMaxLinearWorkGroupSizeMD{} {
+        KernelMaxWorkGroupSizeMD{}, KernelMaxLinearWorkGroupSizeMD{},
+        KernelReqdSubGroupSizeMD{} {
     urContextRetain(Context);
     urDeviceRetain(Device);
   }

source/adapters/hip/kernel.cpp

@@ -274,10 +274,15 @@ urKernelGetSubGroupInfo(ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
     return ReturnValue(0);
   }
   case UR_KERNEL_SUB_GROUP_INFO_SUB_GROUP_SIZE_INTEL: {
-    // Return value of 0 => unspecified or "auto" sub-group size
-    // Correct for now, since warp size may be read from special register
-    // TODO: Return warp size once default is primary sub-group size
-    // TODO: Revisit if we can recover [[sub_group_size]] attribute from PTX
+    const auto &KernelReqdSubGroupSizeMap =
+        hKernel->getProgram()->KernelReqdSubGroupSizeMD;
+    // If present, return the value of intel_reqd_sub_group_size metadata, if
+    // not: 0, which stands for unspecified or auto sub-group size.
+    if (auto KernelReqdSubGroupSize =
+            KernelReqdSubGroupSizeMap.find(hKernel->getName());
+        KernelReqdSubGroupSize != KernelReqdSubGroupSizeMap.end())
+      return ReturnValue(KernelReqdSubGroupSize->second);
+
     return ReturnValue(0);
   }
   default:

source/adapters/hip/program.cpp

@@ -119,6 +119,9 @@ ur_program_handle_t_::setMetadata(const ur_program_metadata_t *Metadata,
       KernelReqdWorkGroupSizeMD[Prefix] =
           std::make_tuple(ReqdWorkGroupElements[0], ReqdWorkGroupElements[1],
                           ReqdWorkGroupElements[2]);
+    } else if (Tag == __SYCL_UR_PROGRAM_METADATA_TAG_REQD_SUB_GROUP_SIZE) {
+      assert(MetadataElement.type == UR_PROGRAM_METADATA_TYPE_UINT32);
+      KernelReqdSubGroupSizeMD[Prefix] = MetadataElement.value.data32;
     }
   }
 

source/adapters/hip/program.hpp

@@ -39,6 +39,7 @@ struct ur_program_handle_t_ {
   std::unordered_map<std::string, std::string> GlobalIDMD;
   std::unordered_map<std::string, std::tuple<uint32_t, uint32_t, uint32_t>>
       KernelReqdWorkGroupSizeMD;
+  std::unordered_map<std::string, uint32_t> KernelReqdSubGroupSizeMD;
 
   constexpr static size_t MAX_LOG_SIZE = 8192u;
 
@@ -48,7 +49,8 @@ struct ur_program_handle_t_ {
 
   ur_program_handle_t_(ur_context_handle_t Ctxt, ur_device_handle_t Device)
       : Module{nullptr}, Binary{}, BinarySizeInBytes{0}, RefCount{1},
-        Context{Ctxt}, Device{Device}, KernelReqdWorkGroupSizeMD{} {
+        Context{Ctxt}, Device{Device}, KernelReqdWorkGroupSizeMD{},
+        KernelReqdSubGroupSizeMD{} {
     urContextRetain(Context);
     urDeviceRetain(Device);
   }

source/ur/ur.hpp

@@ -57,6 +57,8 @@ const ur_command_t UR_EXT_COMMAND_TYPE_USER =
   "@max_work_group_size"
 #define __SYCL_UR_PROGRAM_METADATA_TAG_MAX_LINEAR_WORK_GROUP_SIZE              \
   "@max_linear_work_group_size"
+#define __SYCL_UR_PROGRAM_METADATA_TAG_REQD_SUB_GROUP_SIZE                     \
+  "@reqd_sub_group_size"
 #define __SYCL_UR_PROGRAM_METADATA_TAG_NEED_FINALIZATION "Requires finalization"
 
 // Terminates the process with a catastrophic error message.