Replace loader handles with field at start of handle data

This replaces the handle logic in the loader from wrapped pointers
to a ddi table at the start of the handle struct itself.

Author: RossBrunton

unified-runtime/scripts/generate_code.py

@@ -258,25 +258,6 @@ def _mako_loader_cpp(path, namespace, tags, version, specs, meta):
         "make_loader_cpp path %s namespace %s version %s\n" % (path, namespace, version)
     )
     loc = 0
-    template = "ldrddi.hpp.mako"
-    fin = os.path.join(templates_dir, template)
-
-    name = "%s_ldrddi" % (namespace)
-    filename = "%s.hpp" % (name)
-    fout = os.path.join(path, filename)
-
-    print("Generating %s..." % fout)
-    loc += util.makoWrite(
-        fin,
-        fout,
-        name=name,
-        ver=version,
-        namespace=namespace,
-        tags=tags,
-        specs=specs,
-        meta=meta,
-    )
-
     template = "ldrddi.cpp.mako"
     fin = os.path.join(templates_dir, template)
 

unified-runtime/scripts/templates/helper.py

@@ -1510,20 +1510,6 @@ def get_initial_null_set(obj):
     return ""
 
 
-"""
-Public:
-    returns true if the function always wraps output pointers in loader handles
-"""
-
-
-def always_wrap_outputs(obj):
-    cname = obj_traits.class_name(obj)
-    return (cname, obj["name"]) in [
-        ("$xProgram", "Link"),
-        ("$xProgram", "LinkExp"),
-    ]
-
-
 """
 Private:
     returns the list of parameters, filtering based on desc tags

unified-runtime/scripts/templates/ldrddi.cpp.mako

@@ -20,6 +20,7 @@ from templates import helper as th
 //===----------------------------------------------------------------------===//
 #include <${n}_api.h>
 #include <${n}_ddi.h>
+#include <mutex>
 
 #include "ur_interface_loader.hpp"
 
@@ -68,22 +69,45 @@ ${X}_APIEXPORT ${x}_result_t ${X}_APICALL ${tbl['export']['name']}(
 } // extern "C"
 #endif
 
-#ifdef UR_STATIC_ADAPTER_${Adapter}
-namespace ur::${adapter} {
-ur_result_t urAdapterGetDdiTables(ur_dditable_t *ddi) {
-  if (ddi == nullptr) {
+namespace {
+ur_result_t populateDdiTable(ur_dditable_t *ddi) {
+    if (ddi == nullptr) {
     return UR_RESULT_ERROR_INVALID_NULL_POINTER;
   }
 
   ur_result_t result;
 
+#ifdef UR_STATIC_ADAPTER_${Adapter}
+#define NAMESPACE_ ::ur::${adapter}
+#else
+#define NAMESPACE_
+#endif
+
 %for tbl in th.get_pfntables(specs, meta, n, tags):
-  result = ${n}::${adapter}::${tbl['export']['name']}( ${X}_API_VERSION_CURRENT, &ddi->${tbl['name']} );
+  result = NAMESPACE_::${tbl['export']['name']}( ${X}_API_VERSION_CURRENT, &ddi->${tbl['name']} );
   if (result != UR_RESULT_SUCCESS)
     return result;
 %endfor
 
+#undef NAMESPACE_
+
   return result;
 }
 }
+
+
+namespace ur::${adapter} {
+const ${x}_dditable_t *ddi_getter::value() {
+  static std::once_flag flag;
+  static ${x}_dditable_t table;
+
+  std::call_once(flag, []() { populateDdiTable(&table); });
+  return &table;
+}
+
+#ifdef UR_STATIC_ADAPTER_${Adapter}
+ur_result_t urAdapterGetDdiTables(${x}_dditable_t *ddi) {
+  return populateDdiTable(ddi);
+}
 #endif
+}

unified-runtime/scripts/templates/ldrddi.hpp.mako

@@ -18,6 +18,8 @@ from templates import helper as th
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+#pragma once
+
 #include <${n}_api.h>
 #include <${n}_ddi.h>
 
@@ -36,4 +38,8 @@ ${x}_result_t ${th.make_func_name(n, tags, obj)}(
 #ifdef UR_STATIC_ADAPTER_LEVEL_ZERO
 ur_result_t urAdapterGetDdiTables(ur_dditable_t *ddi);
 #endif
+
+struct ddi_getter {
+  const static ${x}_dditable_t *value();
+};
 }

unified-runtime/scripts/templates/ur_interface_loader.cpp.mako

@@ -14,7 +14,7 @@
 #include "logger/ur_logger.hpp"
 #include "tracing.hpp"
 
-struct ur_adapter_handle_t_ {
+struct ur_adapter_handle_t_ : ur_handle_t_ {
   std::atomic<uint32_t> RefCount = 0;
   std::mutex Mutex;
   struct cuda_tracing_context_t_ *TracingCtx = nullptr;
@@ -41,7 +41,8 @@ class ur_legacy_sink : public logger::Sink {
 // through UR entry points.
 // https://github.com/oneapi-src/unified-runtime/issues/1330
 ur_adapter_handle_t_::ur_adapter_handle_t_()
-    : logger(logger::get_logger("cuda",
+    : ur_handle_t_(),
+      logger(logger::get_logger("cuda",
                                 /*default_log_level*/ logger::Level::ERR)) {
 
   if (std::getenv("UR_LOG_CUDA") != nullptr)

unified-runtime/scripts/templates/ur_interface_loader.hpp.mako

@@ -93,3 +93,8 @@ inline umf_result_t setCUMemoryProviderParams(
 }
 
 } // namespace umf
+
+struct cuda_ddi_getter {
+  const static ur_dditable_t *value();
+};
+using ur_handle_t_ = ur_handle_base_t_<cuda_ddi_getter>;

unified-runtime/source/adapters/cuda/adapter.cpp

@@ -108,7 +108,7 @@ CreateHostMemoryProviderPool(ur_device_handle_t_ *DeviceHandle,
   return UR_RESULT_SUCCESS;
 }
 
-struct ur_context_handle_t_ {
+struct ur_context_handle_t_ : ur_handle_t_ {
 
   struct deleter_data {
     ur_context_extended_deleter_t Function;
@@ -126,7 +126,7 @@ struct ur_context_handle_t_ {
   umf_memory_pool_handle_t MemoryPoolHost = nullptr;
 
   ur_context_handle_t_(const ur_device_handle_t *Devs, uint32_t NumDevices)
-      : Devices{Devs, Devs + NumDevices}, RefCount{1} {
+      : ur_handle_t_(), Devices{Devs, Devs + NumDevices}, RefCount{1} {
     for (auto &Dev : Devices) {
       urDeviceRetain(Dev);
     }

unified-runtime/source/adapters/cuda/common.hpp

@@ -16,7 +16,7 @@
 
 #include "common.hpp"
 
-struct ur_device_handle_t_ {
+struct ur_device_handle_t_ : ur_handle_t_ {
 private:
   using native_type = CUdevice;
 
@@ -40,8 +40,8 @@ struct ur_device_handle_t_ {
 public:
   ur_device_handle_t_(native_type cuDevice, CUcontext cuContext, CUevent evBase,
                       ur_platform_handle_t platform, uint32_t DevIndex)
-      : CuDevice(cuDevice), CuContext(cuContext), EvBase(evBase), RefCount{1},
-        Platform(platform), DeviceIndex{DevIndex} {
+      : ur_handle_t_(), CuDevice(cuDevice), CuContext(cuContext),
+        EvBase(evBase), RefCount{1}, Platform(platform), DeviceIndex{DevIndex} {
 
     UR_CHECK_ERROR(cuDeviceGetAttribute(
         &MaxRegsPerBlock, CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_BLOCK,

unified-runtime/source/adapters/cuda/context.hpp

@@ -24,7 +24,7 @@ ur_event_handle_t_::ur_event_handle_t_(ur_command_t Type,
                                        native_type EvEnd, native_type EvQueued,
                                        native_type EvStart, CUstream Stream,
                                        uint32_t StreamToken)
-    : CommandType{Type}, RefCount{1}, HasOwnership{true},
+    : ur_handle_t_(), CommandType{Type}, RefCount{1}, HasOwnership{true},
       HasBeenWaitedOn{false}, IsRecorded{false}, IsStarted{false},
       StreamToken{StreamToken}, EventID{0}, EvEnd{EvEnd}, EvStart{EvStart},
       EvQueued{EvQueued}, Queue{Queue}, Stream{Stream}, Context{Context} {
@@ -34,11 +34,12 @@ ur_event_handle_t_::ur_event_handle_t_(ur_command_t Type,
 
 ur_event_handle_t_::ur_event_handle_t_(ur_context_handle_t Context,
                                        CUevent EventNative)
-    : CommandType{UR_COMMAND_EVENTS_WAIT}, RefCount{1}, HasOwnership{false},
-      HasBeenWaitedOn{false}, IsRecorded{false}, IsStarted{false},
-      IsInterop{true}, StreamToken{std::numeric_limits<uint32_t>::max()},
-      EventID{0}, EvEnd{EventNative}, EvStart{nullptr}, EvQueued{nullptr},
-      Queue{nullptr}, Stream{nullptr}, Context{Context} {
+    : ur_handle_t_(), CommandType{UR_COMMAND_EVENTS_WAIT}, RefCount{1},
+      HasOwnership{false}, HasBeenWaitedOn{false}, IsRecorded{false},
+      IsStarted{false}, IsInterop{true},
+      StreamToken{std::numeric_limits<uint32_t>::max()}, EventID{0},
+      EvEnd{EventNative}, EvStart{nullptr}, EvQueued{nullptr}, Queue{nullptr},
+      Stream{nullptr}, Context{Context} {
   urContextRetain(Context);
 }
 

unified-runtime/source/adapters/cuda/device.hpp

@@ -17,7 +17,7 @@
 
 /// UR Event mapping to CUevent
 ///
-struct ur_event_handle_t_ {
+struct ur_event_handle_t_ : ur_handle_t_ {
 public:
   using native_type = CUevent;
 

unified-runtime/source/adapters/cuda/event.cpp

@@ -34,7 +34,7 @@
 /// A compiler pass converts the UR API local memory model into the
 /// CUDA shared model. This object simply calculates the total of
 /// shared memory, and the initial offsets of each parameter.
-struct ur_kernel_handle_t_ {
+struct ur_kernel_handle_t_ : ur_handle_t_ {
   using native_type = CUfunction;
 
   native_type Function;
@@ -250,8 +250,9 @@ struct ur_kernel_handle_t_ {
   ur_kernel_handle_t_(CUfunction Func, CUfunction FuncWithOffsetParam,
                       const char *Name, ur_program_handle_t Program,
                       ur_context_handle_t Context)
-      : Function{Func}, FunctionWithOffsetParam{FuncWithOffsetParam},
-        Name{Name}, Context{Context}, Program{Program}, RefCount{1} {
+      : ur_handle_t_(), Function{Func},
+        FunctionWithOffsetParam{FuncWithOffsetParam}, Name{Name},
+        Context{Context}, Program{Program}, RefCount{1} {
     urProgramRetain(Program);
     urContextRetain(Context);
     /// Note: this code assumes that there is only one device per context

unified-runtime/source/adapters/cuda/event.hpp

@@ -310,7 +310,7 @@ struct SurfaceMem {
 /// is on a different device, marked by
 /// LastQueueWritingToMemObj->getDevice()
 ///
-struct ur_mem_handle_t_ {
+struct ur_mem_handle_t_ : ur_handle_t_ {
   // Context where the memory object is accessible
   ur_context_handle_t Context;
 
@@ -345,17 +345,17 @@ struct ur_mem_handle_t_ {
   /// Constructs the UR mem handler for a non-typed allocation ("buffer")
   ur_mem_handle_t_(ur_context_handle_t Ctxt, ur_mem_flags_t MemFlags,
                    BufferMem::AllocMode Mode, void *HostPtr, size_t Size)
-      : Context{Ctxt}, RefCount{1}, MemFlags{MemFlags},
+      : ur_handle_t_(), Context{Ctxt}, RefCount{1}, MemFlags{MemFlags},
         HaveMigratedToDeviceSinceLastWrite(Context->Devices.size(), false),
         Mem{std::in_place_type<BufferMem>, Ctxt, this, Mode, HostPtr, Size} {
     urContextRetain(Context);
   };
 
   // Subbuffer constructor
   ur_mem_handle_t_(ur_mem_handle_t Parent, size_t SubBufferOffset)
-      : Context{Parent->Context}, RefCount{1}, MemFlags{Parent->MemFlags},
-        HaveMigratedToDeviceSinceLastWrite(Parent->Context->Devices.size(),
-                                           false),
+      : ur_handle_t_(), Context{Parent->Context}, RefCount{1},
+        MemFlags{Parent->MemFlags}, HaveMigratedToDeviceSinceLastWrite(
+                                        Parent->Context->Devices.size(), false),
         Mem{BufferMem{std::get<BufferMem>(Parent->Mem)}} {
     auto &SubBuffer = std::get<BufferMem>(Mem);
     SubBuffer.Parent = Parent;
@@ -376,7 +376,7 @@ struct ur_mem_handle_t_ {
   ur_mem_handle_t_(ur_context_handle_t Ctxt, ur_mem_flags_t MemFlags,
                    ur_image_format_t ImageFormat, ur_image_desc_t ImageDesc,
                    void *HostPtr)
-      : Context{Ctxt}, RefCount{1}, MemFlags{MemFlags},
+      : ur_handle_t_(), Context{Ctxt}, RefCount{1}, MemFlags{MemFlags},
         HaveMigratedToDeviceSinceLastWrite(Context->Devices.size(), false),
         Mem{std::in_place_type<SurfaceMem>,
             Ctxt,

unified-runtime/source/adapters/cuda/kernel.hpp

@@ -20,7 +20,7 @@
 /// UR queue mapping on physical memory allocations used in virtual memory
 /// management.
 ///
-struct ur_physical_mem_handle_t_ {
+struct ur_physical_mem_handle_t_ : ur_handle_t_ {
   using native_type = CUmemGenericAllocationHandle;
 
   std::atomic_uint32_t RefCount;
@@ -33,8 +33,8 @@ struct ur_physical_mem_handle_t_ {
   ur_physical_mem_handle_t_(native_type PhysMem, ur_context_handle_t_ *Ctx,
                             ur_device_handle_t Device, size_t Size,
                             ur_physical_mem_properties_t Properties)
-      : RefCount(1), PhysicalMem(PhysMem), Context(Ctx), Device(Device),
-        Size(Size), Properties(Properties) {
+      : ur_handle_t_(), RefCount(1), PhysicalMem(PhysMem), Context(Ctx),
+        Device(Device), Size(Size), Properties(Properties) {
     urContextRetain(Context);
     urDeviceRetain(Device);
   }

unified-runtime/source/adapters/cuda/memory.hpp

@@ -9,9 +9,10 @@
 //===----------------------------------------------------------------------===//
 #pragma once
 
+#include "common.hpp"
 #include <ur/ur.hpp>
 #include <vector>
 
-struct ur_platform_handle_t_ {
+struct ur_platform_handle_t_ : ur_handle_t_ {
   std::vector<std::unique_ptr<ur_device_handle_t_>> Devices;
 };