Add optional code location entry point variants

callumfare · callumfare · commit 8efd8833412d · 2024-12-04T11:42:31.000Z
diff --git a/offload/new-api/API/Common.td b/offload/new-api/API/Common.td
@@ -89,3 +89,14 @@ def : Macro {
   let desc = "Success condition";
   let value = "NULL";
 }
+
+def : Struct {
+  let name = "offload_code_location_t";
+  let desc = "Code location information that can optionally be associated with an API call";
+  let members = [
+    StructMember<"const char*", "FunctionName", "Function name">,
+    StructMember<"const char*", "SourceFile", "Source code file">,
+    StructMember<"uint32_t", "LineNumber", "Source code line number">,
+    StructMember<"uint32_t", "ColumnNumber", "Source code column number">
+  ];
+}
diff --git a/offload/new-api/include/offload_api.h b/offload/new-api/include/offload_api.h
@@ -129,6 +129,16 @@ typedef const offload_error_struct_t *offload_result_t;
 #endif // !defined(OFFLOAD_SUCCESS)
 #endif // OFFLOAD_SUCCESS
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Code location information that can optionally be associated with an
+/// API call
+typedef struct offload_code_location_t {
+  const char *FunctionName; /// Function name
+  const char *SourceFile;   /// Source code file
+  uint32_t LineNumber;      /// Source code line number
+  uint32_t ColumnNumber;    /// Source code column number
+} offload_code_location_t;
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Retrieves all available platforms
 ///
@@ -397,6 +407,46 @@ typedef struct offload_device_get_info_params_t {
   size_t **ppPropSizeRet;
 } offload_device_get_info_params_t;
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Variant of offloadPlatformGet that also sets source code location
+/// information
+/// @details See also ::offloadPlatformGet
+OFFLOAD_APIEXPORT offload_result_t OFFLOAD_APICALL
+offloadPlatformGetWithCodeLoc(uint32_t NumEntries,
+                              offload_platform_handle_t *phPlatforms,
+                              uint32_t *pNumPlatforms,
+                              offload_code_location_t *pCodeLocation);
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Variant of offloadPlatformGetInfo that also sets source code location
+/// information
+/// @details See also ::offloadPlatformGetInfo
+OFFLOAD_APIEXPORT offload_result_t OFFLOAD_APICALL
+offloadPlatformGetInfoWithCodeLoc(offload_platform_handle_t hPlatform,
+                                  offload_platform_info_t propName,
+                                  size_t propSize, void *pPropValue,
+                                  size_t *pPropSizeRet,
+                                  offload_code_location_t *pCodeLocation);
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Variant of offloadDeviceGet that also sets source code location
+/// information
+/// @details See also ::offloadDeviceGet
+OFFLOAD_APIEXPORT offload_result_t OFFLOAD_APICALL offloadDeviceGetWithCodeLoc(
+    offload_platform_handle_t hPlatform, offload_device_type_t DeviceType,
+    uint32_t NumEntries, offload_device_handle_t *phDevices,
+    uint32_t *pNumDevices, offload_code_location_t *pCodeLocation);
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Variant of offloadDeviceGetInfo that also sets source code location
+/// information
+/// @details See also ::offloadDeviceGetInfo
+OFFLOAD_APIEXPORT offload_result_t OFFLOAD_APICALL
+offloadDeviceGetInfoWithCodeLoc(offload_device_handle_t hDevice,
+                                offload_device_info_t propName, size_t propSize,
+                                void *pPropValue, size_t *pPropSizeRet,
+                                offload_code_location_t *pCodeLocation);
+
 #if defined(__cplusplus)
 } // extern "C"
 #endif
diff --git a/offload/new-api/include/offload_entry_points.inc b/offload/new-api/include/offload_entry_points.inc
@@ -40,6 +40,16 @@ offloadPlatformGet(uint32_t NumEntries, offload_platform_handle_t *phPlatforms,
   }
   return result;
 }
+offload_result_t offloadPlatformGetWithCodeLoc(
+    uint32_t NumEntries, offload_platform_handle_t *phPlatforms,
+    uint32_t *pNumPlatforms, offload_code_location_t *pCodeLocation) {
+  CodeLocation() = pCodeLocation;
+  offload_result_t result =
+      offloadPlatformGet(NumEntries, phPlatforms, pNumPlatforms);
+
+  CodeLocation() = nullptr;
+  return result;
+}
 
 ///////////////////////////////////////////////////////////////////////////////
 offload_impl_result_t
@@ -88,6 +98,17 @@ OFFLOAD_APIEXPORT offload_result_t OFFLOAD_APICALL offloadPlatformGetInfo(
   }
   return result;
 }
+offload_result_t offloadPlatformGetInfoWithCodeLoc(
+    offload_platform_handle_t hPlatform, offload_platform_info_t propName,
+    size_t propSize, void *pPropValue, size_t *pPropSizeRet,
+    offload_code_location_t *pCodeLocation) {
+  CodeLocation() = pCodeLocation;
+  offload_result_t result = offloadPlatformGetInfo(
+      hPlatform, propName, propSize, pPropValue, pPropSizeRet);
+
+  CodeLocation() = nullptr;
+  return result;
+}
 
 ///////////////////////////////////////////////////////////////////////////////
 offload_impl_result_t offloadDeviceGet_val(offload_platform_handle_t hPlatform,
@@ -134,6 +155,17 @@ offloadDeviceGet(offload_platform_handle_t hPlatform,
   }
   return result;
 }
+offload_result_t offloadDeviceGetWithCodeLoc(
+    offload_platform_handle_t hPlatform, offload_device_type_t DeviceType,
+    uint32_t NumEntries, offload_device_handle_t *phDevices,
+    uint32_t *pNumDevices, offload_code_location_t *pCodeLocation) {
+  CodeLocation() = pCodeLocation;
+  offload_result_t result = offloadDeviceGet(hPlatform, DeviceType, NumEntries,
+                                             phDevices, pNumDevices);
+
+  CodeLocation() = nullptr;
+  return result;
+}
 
 ///////////////////////////////////////////////////////////////////////////////
 offload_impl_result_t offloadDeviceGetInfo_val(offload_device_handle_t hDevice,
@@ -183,3 +215,15 @@ OFFLOAD_APIEXPORT offload_result_t OFFLOAD_APICALL offloadDeviceGetInfo(
   }
   return result;
 }
+offload_result_t
+offloadDeviceGetInfoWithCodeLoc(offload_device_handle_t hDevice,
+                                offload_device_info_t propName, size_t propSize,
+                                void *pPropValue, size_t *pPropSizeRet,
+                                offload_code_location_t *pCodeLocation) {
+  CodeLocation() = pCodeLocation;
+  offload_result_t result = offloadDeviceGetInfo(hDevice, propName, propSize,
+                                                 pPropValue, pPropSizeRet);
+
+  CodeLocation() = nullptr;
+  return result;
+}
diff --git a/offload/new-api/src/offload_lib.cpp b/offload/new-api/src/offload_lib.cpp
@@ -26,6 +26,11 @@ ErrSetT &Errors() {
   return Errors;
 }
 
+offload_code_location_t *&CodeLocation() {
+  thread_local offload_code_location_t *CodeLoc = nullptr;
+  return CodeLoc;
+}
+
 // Pull in the declarations for the implementation funtions. The actual entry
 // points in this file wrap these.
 #include "offload_impl_func_decls.inc"
diff --git a/offload/tools/offload-tblgen/APIGen.cpp b/offload/tools/offload-tblgen/APIGen.cpp
@@ -168,6 +168,24 @@ typedef struct {1} {{
   OS << formatv("} {0};\n", Func.getParamStructName());
 }
 
+static void ProcessFuncWithCodeLocVariant(const FunctionRec &Func,
+                                          raw_ostream &OS) {
+
+  auto FuncWithCodeLocBegin = R"(
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Variant of {0} that also sets source code location information
+/// @details See also ::{0}
+OFFLOAD_APIEXPORT offload_result_t OFFLOAD_APICALL {0}WithCodeLoc(
+)";
+  OS << formatv(FuncWithCodeLocBegin, Func.getName());
+  auto Params = Func.getParams();
+  for (auto &Param : Params) {
+    OS << "  " << Param.getType() << " " << Param.getName();
+    OS << ",\n";
+  }
+  OS << "offload_code_location_t *pCodeLocation);\n\n";
+}
+
 void EmitOffloadAPI(RecordKeeper &Records, raw_ostream &OS) {
   OS << GenericHeader;
   OS << FileHeader;
@@ -193,5 +211,9 @@ void EmitOffloadAPI(RecordKeeper &Records, raw_ostream &OS) {
     ProcessFuncParamStruct(FunctionRec{R}, OS);
   }
 
+  for (auto *R : Records.getAllDerivedDefinitions("Function")) {
+    ProcessFuncWithCodeLocVariant(FunctionRec{R}, OS);
+  }
+
   OS << FileFooter;
 }
diff --git a/offload/tools/offload-tblgen/EntryPointGen.cpp b/offload/tools/offload-tblgen/EntryPointGen.cpp
@@ -102,10 +102,33 @@ static void EmitEntryPointFunc(const FunctionRec &F, raw_ostream &OS) {
   OS << "}\n";
 }
 
+static void EmitCodeLocWrapper(const FunctionRec &F, raw_ostream &OS) {
+  // Emit preamble
+  OS << formatv("{0}_result_t {1}WithCodeLoc(\n  ", PrefixLower, F.getName());
+  // Emit arguments
+  std::string ParamNameList = "";
+  for (auto &Param : F.getParams()) {
+    OS << Param.getType() << " " << Param.getName() << ", ";
+    ParamNameList += Param.getName().str();
+    if (Param != F.getParams().back()) {
+      ParamNameList += ", ";
+    }
+  }
+  OS << "offload_code_location_t *pCodeLocation";
+  OS << ") {\n";
+  OS << TAB_1 "CodeLocation() = pCodeLocation;\n";
+  OS << formatv(TAB_1 "{0}_result_t result = {1}({2});\n\n", PrefixLower,
+                F.getName(), ParamNameList);
+  OS << TAB_1 "CodeLocation() = nullptr;\n";
+  OS << TAB_1 "return result;\n";
+  OS << "}\n";
+}
+
 void EmitOffloadEntryPoints(RecordKeeper &Records, raw_ostream &OS) {
   OS << GenericHeader;
   for (auto *R : Records.getAllDerivedDefinitions("Function")) {
     EmitValidationFunc(FunctionRec{R}, OS);
     EmitEntryPointFunc(FunctionRec{R}, OS);
+    EmitCodeLocWrapper(FunctionRec{R}, OS);
   }
 }
