Как реализовать дескрипторы для библиотеки API драйверов CUDA?

Note: The question has been updated to address the questions that have been raised in the comments, and to emphasize that the core of the question is about the interdependencies between the Runtime- and Driver API

Библиотеки времени выполнения CUDA (такие как CUBLAS или CUFFT) обычно используют концепцию "дескриптора", которая суммирует состояние и контекст такой библиотеки. Шаблон использования довольно прост:

// Create a handle
cublasHandle_t handle;
cublasCreate(&handle);

// Call some functions, always passing in the handle as the first argument
cublasSscal(handle, ...);

// When done, destroy the handle
cublasDestroy(handle);

Однако, есть много тонких деталей о том, как эти дескрипторы взаимодействуют с контекстами Driver- и среды выполнения, а также с несколькими потоками и устройствами. В документации перечислено несколько разрозненных деталей об обработке контекста:

Тем не менее, некоторая информация, кажется, не совсем обновлена (например, я думаю, что следует использовать cuCtxSetCurrent вместо cuCtxPushCurrent и cuCtxPopCurrent?), А некоторые, кажется, относятся ко времени, предшествующему "Первичному Обработка контекста была раскрыта через API драйвера, а некоторые части упрощены в том смысле, что они показывают только самые простые шаблоны использования, делают только расплывчатые или неполные утверждения о многопоточности или не могут применяться к понятию "дескрипторов", которое используется в библиотеки времени выполнения.

Моя цель состоит в том, чтобы реализовать библиотеку времени выполнения, которая предлагает свой собственный тип "дескриптора" и который позволяет использовать шаблоны, которые эквивалентны другим библиотекам времени выполнения с точки зрения обработки контекста и безопасности потока.

В случае, если библиотека может быть реализована внутри системы исключительно с использованием API времени выполнения, все может быть понятно: управление контекстом лежит исключительно на пользователе. Если он создает собственный контекст драйвера, будут применяться правила, изложенные в документации по Runtime- и управлению контекстом драйвера. В противном случае функции API времени выполнения позаботятся об обработке основных контекстов.

Однако может случиться так, что библиотеке придется использовать Driver API. Например, чтобы загрузить файлы PTX как объекты CUmodule и получить из них объекты CUfunction. И когда библиотека должна - для пользователя - вести себя как библиотека времени выполнения, но внутренне должна использовать API драйвера, возникают некоторые вопросы о том, как обработка контекста должна быть реализована "под капотом".

То, что я понял до сих пор, набросано здесь.

(Это "псевдокод" в том смысле, что он пропускает проверки ошибок и другие подробности, и... все это должно быть реализовано в Java, но это не должно быть здесь уместно)

1. "Дескриптор" - это, по сути, класс/структура, содержащая следующую информацию:

class Handle 
{
    CUcontext context;
    boolean usingPrimaryContext;
    CUdevice device;
}

2. При его создании необходимо рассмотреть два случая: его можно создать, когда контекст драйвера является текущим для вызывающего потока. В этом случае следует использовать этот контекст. В противном случае он должен использовать основной контекст текущего (во время выполнения) устройства:

Handle createHandle()
{
    cuInit(0);

    // Obtain the current context
    CUcontext context;
    cuCtxGetCurrent(&context);

    CUdevice device;

    // If there is no context, use the primary context
    boolean usingPrimaryContext = false;
    if (context == nullptr)
    {
        usingPrimaryContext = true;

        // Obtain the device that is currently selected via the runtime API
        int deviceIndex;
        cudaGetDevice(&deviceIndex);

        // Obtain the device and its primary context
        cuDeviceGet(&device, deviceIndex);
        cuDevicePrimaryCtxRetain(&context, device));
        cuCtxSetCurrent(context);
    }
    else
    {
        cuCtxGetDevice(device);
    }

    // Create the actual handle. This might internally allocate
    // memory or do other things that are specific for the context
    // for which the handle is created
    Handle handle = new Handle(device, context, usingPrimaryContext);
    return handle;
}

3. При вызове ядра библиотеки контекст соответствующего дескриптора обновляется для вызывающего потока:

void someLibraryFunction(Handle handle)
{
    cuCtxSetCurrent(handle.context);
    callMyKernel(...);
}

Здесь можно утверждать, что вызывающая сторона отвечает за то, чтобы требуемый контекст был текущим. Но если дескриптор был создан для основного контекста, этот контекст будет автоматически обновляться.

4. Когда дескриптор уничтожен, это означает, что должен вызываться cuDevicePrimaryCtxRelease, но только когда контекст является основным контекстом:

void destroyHandle(Handle handle)
{
    if (handle.usingPrimaryContext)
    {
        cuDevicePrimaryCtxRelease(handle.device);
    }
}

Из моих экспериментов до сих пор это, похоже, демонстрирует то же поведение, что и дескриптор CUBLAS, например. Но мои возможности для тщательного тестирования этого ограничены, потому что у меня есть только одно устройство, и, таким образом, я не могу проверить критические случаи, например иметь два контекста, по одному для каждого из двух устройств.

Итак, мои вопросы:

  • Существуют ли какие-либо установленные шаблоны для реализации такого "дескриптора"?
  • Существуют ли какие-либо шаблоны использования (например, с несколькими устройствами и одним контекстом на устройство), которые не могут быть охвачены описанным выше подходом, но будут охвачены реализациями "дескриптора" CUBLAS?
  • В более общем плане: есть ли рекомендации по улучшению текущей реализации "дескриптора"?
  • Риторика: доступен ли где-нибудь исходный код обработки дескриптора CUBLAS?

(Я также взглянул на обработку контекста в tenorflow, но я не уверен, можно ли извлечь из этого рекомендации о том, как реализовать дескрипторы для библиотеки времени выполнения...)

(An "Update" has been removed here, because it was added in response to the comments, and should no longer be relevant)

ОТВЕТЫ