SYCL] Support specialization constants in the L0 plugin (#1982)

Author: Jaime Arteaga <[email protected]>
Signed-off-by: Artur Gainullin <[email protected]>

Author: againull

sycl/plugins/level_zero/pi_level0.cpp

@@ -20,6 +20,7 @@
 #include <string>
 #include <thread>
 #include <utility>
+#include <vector>
 
 #include <level_zero/zet_api.h>
 
@@ -1824,9 +1825,29 @@ pi_result piProgramBuild(pi_program Program, pi_uint32 NumDevices,
   // Check that the program wasn't already built.
   assert(!Program->ZeModule);
 
+  // Translate collected specialization constants.
+  ze_module_constants_t ZeSpecConstants = {};
+  std::vector<uint32_t> ZeSpecContantsIds(Program->ZeSpecConstants.size());
+  std::vector<uint64_t> ZeSpecContantsValues(Program->ZeSpecConstants.size());
+  {
+    std::lock_guard<std::mutex> ZeSpecConstantsMutexGuard(
+        Program->ZeSpecConstantsMutex);
+    ZeSpecConstants.numConstants = Program->ZeSpecConstants.size();
+    auto ZeSpecContantsIdsIt = ZeSpecContantsIds.begin();
+    auto ZeSpecContantsValuesIt = ZeSpecContantsValues.begin();
+    for (auto &SpecConstant : Program->ZeSpecConstants) {
+      *ZeSpecContantsIdsIt = SpecConstant.first;
+      ZeSpecContantsIdsIt++;
+      *ZeSpecContantsValuesIt = SpecConstant.second;
+      ZeSpecContantsValuesIt++;
+    }
+    ZeSpecConstants.pConstantIds = ZeSpecContantsIds.data();
+    ZeSpecConstants.pConstantValues = ZeSpecContantsValues.data();
+  }
+
+  // Complete the module's descriptor
   Program->ZeModuleDesc.pBuildFlags = Options;
-  // TODO: set specialization constants here.
-  Program->ZeModuleDesc.pConstants = nullptr;
+  Program->ZeModuleDesc.pConstants = &ZeSpecConstants;
 
   ze_device_handle_t ZeDevice = Program->Context->Device->ZeDevice;
   ZE_CALL(zeModuleCreate(ZeDevice, &Program->ZeModuleDesc, &Program->ZeModule,
@@ -3713,12 +3734,20 @@ pi_result piKernelSetExecInfo(pi_kernel Kernel, pi_kernel_exec_info ParamName,
 }
 
 pi_result piextProgramSetSpecializationConstant(pi_program Prog,
-                                                pi_uint32 SpecID,
-                                                size_t SpecSize,
+                                                pi_uint32 SpecID, size_t,
                                                 const void *SpecValue) {
-  // TODO: implement
-  die("piextProgramSetSpecializationConstant: not implemented");
-  return {};
+  // Level Zero sets spec constants when creating modules,
+  // so save them for when program is built.
+  std::lock_guard<std::mutex> ZeSpecConstantsMutexGuard(
+      Prog->ZeSpecConstantsMutex);
+
+  // Pass SpecValue pointer. Spec constant value is retrieved
+  // by Level-Zero when creating the modul
+  //
+  // NOTE: SpecSize is unused in L0, the size is known from SPIR-V by SpecID.
+  Prog->ZeSpecConstants[SpecID] = reinterpret_cast<uint64_t>(SpecValue);
+
+  return PI_SUCCESS;
 }
 
 pi_result piPluginInit(pi_plugin *PluginInit) {

sycl/plugins/level_zero/pi_level0.hpp

@@ -316,6 +316,9 @@ struct _pi_program : _pi_object {
 
   // L0 module handle.
   ze_module_handle_t ZeModule;
+  // L0 module specialization constants
+  std::mutex ZeSpecConstantsMutex;
+  std::unordered_map<uint32_t, uint64_t> ZeSpecConstants;
 
   // L0 build log.
   ze_module_build_log_handle_t ZeBuildLog;