SWDEV-523281 - [clr] Implementation of hipLaunchKernelExC and hipDrvLaunchKernelEx API with support for cooperative launch (#92)

hipamd/src/amdhip.def

@@ -489,4 +489,6 @@ hipLinkAddData
 hipLinkAddFile
 hipLinkComplete
 hipLinkCreate
-hipLinkDestroy
+hipLinkDestroy
+hipLaunchKernelExC
+hipDrvLaunchKernelEx

hipamd/src/hip_api_trace.cpp

@@ -823,6 +823,10 @@ hipError_t hipGraphBatchMemOpNodeSetParams(hipGraphNode_t hNode,
 hipError_t hipGraphExecBatchMemOpNodeSetParams(hipGraphExec_t hGraphExec, hipGraphNode_t hNode,
                                                const hipBatchMemOpNodeParams* nodeParams);
 hipError_t hipEventRecordWithFlags(hipEvent_t event, hipStream_t stream, unsigned flags);
+hipError_t hipLaunchKernelExC(const hipLaunchConfig_t* config, const void* fPtr, void** args);
+hipError_t hipDrvLaunchKernelEx(const HIP_LAUNCH_CONFIG* config, hipFunction_t f, void** params,
+                                void** extra);
+
 }  // namespace hip
 
 namespace hip {
@@ -1335,6 +1339,8 @@ void UpdateDispatchTable(HipDispatchTable* ptrDispatchTable) {
   ptrDispatchTable->hipGraphExecBatchMemOpNodeSetParams_fn =
       hip::hipGraphExecBatchMemOpNodeSetParams;
   ptrDispatchTable->hipEventRecordWithFlags_fn = hip::hipEventRecordWithFlags;
+  ptrDispatchTable->hipLaunchKernelExC_fn = hip::hipLaunchKernelExC;
+  ptrDispatchTable->hipDrvLaunchKernelEx_fn = hip::hipDrvLaunchKernelEx;
 }
 
 #if HIP_ROCPROFILER_REGISTER > 0
@@ -1972,15 +1978,18 @@ HIP_ENFORCE_ABI(HipDispatchTable, hipLinkAddFile_fn , 470)
 HIP_ENFORCE_ABI(HipDispatchTable, hipLinkComplete_fn , 471)
 HIP_ENFORCE_ABI(HipDispatchTable, hipLinkCreate_fn , 472)
 HIP_ENFORCE_ABI(HipDispatchTable, hipLinkDestroy_fn , 473)
+// HIP_RUNTIME_API_TABLE_STEP_VERSION == 11
+HIP_ENFORCE_ABI(HipDispatchTable, hipLaunchKernelExC_fn, 474);
+HIP_ENFORCE_ABI(HipDispatchTable, hipDrvLaunchKernelEx_fn, 475);
 // if HIP_ENFORCE_ABI entries are added for each new function pointer in the table, the number below
 // will be +1 of the number in the last HIP_ENFORCE_ABI line. E.g.:
 //
 //  HIP_ENFORCE_ABI(<table>, <functor>, 8)
 //
 //  HIP_ENFORCE_ABI_VERSIONING(<table>, 9) <- 8 + 1 = 9
-HIP_ENFORCE_ABI_VERSIONING(HipDispatchTable, 474)
+HIP_ENFORCE_ABI_VERSIONING(HipDispatchTable, 476)
 
-static_assert(HIP_RUNTIME_API_TABLE_MAJOR_VERSION == 0 && HIP_RUNTIME_API_TABLE_STEP_VERSION == 9,
+static_assert(HIP_RUNTIME_API_TABLE_MAJOR_VERSION == 0 && HIP_RUNTIME_API_TABLE_STEP_VERSION == 11,
               "If you get this error, add new HIP_ENFORCE_ABI(...) code for the new function "
               "pointers and then update this check so it is true");
 #endif

hipamd/src/hip_hcc.map.in

@@ -595,4 +595,12 @@ global:
     hipLinkDestroy;
 local:
     *;
-} hip_6.2;
+} hip_6.2;
+
+hip_6.5 {
+global:
+    hipLaunchKernelExC;
+    hipDrvLaunchKernelEx;
+local:
+    *;
+} hip_6.4;

hipamd/src/hip_module.cpp

@@ -1077,4 +1077,65 @@ hipError_t hipLinkDestroy(hipLinkState_t hip_link_state) {
   HIP_RETURN(hipSuccess);
 }
 
+hipError_t hipLaunchKernelExC(const hipLaunchConfig_t* config, const void* fPtr, void** args) {
+  HIP_INIT_API(hipLaunchKernelExC, config, fPtr, args);
+
+  if (fPtr == nullptr || config == nullptr || config->numAttrs == 0) {
+    HIP_RETURN(hipErrorInvalidConfiguration);
+  }
+
+  for (size_t attr_idx = 0; attr_idx < config->numAttrs; ++attr_idx) {
+    hipLaunchAttribute& attr = config->attrs[attr_idx];
+    switch (attr.id) {
+      case hipLaunchAttributeCooperative:
+        if (attr.val.cooperative != 0) {
+          HIP_RETURN_DURATION(
+              hipLaunchCooperativeKernel_common(fPtr, config->gridDim, config->blockDim, args,
+                                                config->dynamicSmemBytes, config->stream));
+        }
+        break;
+      default:
+        LogPrintfError("Attribute %u not supported", attr.id);
+        break;
+    }
+  }
+  HIP_RETURN(hipErrorInvalidConfiguration);
+}
+
+hipError_t hipDrvLaunchKernelEx(const HIP_LAUNCH_CONFIG* config, hipFunction_t f,
+                                void** kernelParams, void** extra) {
+  HIP_INIT_API(hipDrvLaunchKernelEx, config, f, kernelParams, extra);
+
+  if (f == nullptr || config == nullptr || config->numAttrs == 0) {
+    HIP_RETURN(hipErrorInvalidConfiguration);
+  }
+
+  size_t globalWorkSizeX = static_cast<size_t>(config->gridDimX) * config->blockDimX;
+  size_t globalWorkSizeY = static_cast<size_t>(config->gridDimY) * config->blockDimY;
+  size_t globalWorkSizeZ = static_cast<size_t>(config->gridDimZ) * config->blockDimZ;
+  if (globalWorkSizeX > std::numeric_limits<uint32_t>::max() ||
+      globalWorkSizeY > std::numeric_limits<uint32_t>::max() ||
+      globalWorkSizeZ > std::numeric_limits<uint32_t>::max()) {
+    HIP_RETURN(hipErrorInvalidConfiguration);
+  }
+
+  for (size_t attr_idx = 0; attr_idx < config->numAttrs; ++attr_idx) {
+    hipLaunchAttribute& attr = config->attrs[attr_idx];
+    switch (attr.id) {
+      case hipLaunchAttributeCooperative:
+        if (attr.value.cooperative != 0) {
+          HIP_RETURN(ihipModuleLaunchKernel(
+              f, static_cast<uint32_t>(globalWorkSizeX), static_cast<uint32_t>(globalWorkSizeY),
+              static_cast<uint32_t>(globalWorkSizeZ), config->blockDimX, config->blockDimY,
+              config->blockDimZ, config->sharedMemBytes, config->hStream, kernelParams, nullptr,
+              nullptr, nullptr, 0, amd::NDRangeKernelCommand::CooperativeGroups));
+        }
+        break;
+      default:
+        LogPrintfError("Attribute %u not supported", attr.id);
+        break;
+    }
+  }
+  HIP_RETURN(hipErrorInvalidConfiguration)
+}
 }  // namespace hip

hipamd/src/hip_table_interface.cpp

@@ -1868,4 +1868,13 @@ hipError_t hipGraphExecBatchMemOpNodeSetParams(hipGraphExec_t hGraphExec,
 }
 hipError_t hipEventRecordWithFlags(hipEvent_t event, hipStream_t stream, unsigned int flags) {
   return hip::GetHipDispatchTable()->hipEventRecordWithFlags_fn(event, stream, flags);
+}
+
+hipError_t hipLaunchKernelExC(const hipLaunchConfig_t* config, const void* fPtr, void** args) {
+  return hip::GetHipDispatchTable()->hipLaunchKernelExC_fn(config, fPtr, args);
+}
+
+hipError_t hipDrvLaunchKernelEx(const HIP_LAUNCH_CONFIG* config, hipFunction_t f, void** kernel,
+                                void** extra) {
+  return hip::GetHipDispatchTable()->hipDrvLaunchKernelEx_fn(config, f, kernel, extra);
 }