Как сделать шаблонные вычислительные ядра в металле

Question

Я писал некоторые вычислительные ядра Metal.Итак, я написал ядро ​​со следующим объявлением:

kernel void
myKernel(const device uint32_t *inData [[buffer(MyKernelIn)]],
        device uint32_t *outData [[buffer(MyKernelOut)]],
        uint2                          gid       [[thread_position_in_grid]],
        uint2 thread_position_in_threadgroup         [[thread_position_in_threadgroup]],
        uint2       threads_per_threadgroup      [[threads_per_threadgroup]],
        uint2 threadgroup_position_in_grid       [[threadgroup_position_in_grid]]) 
{ }

Теперь я хочу написать вариант этого, который принимает inData типа uint8_t и float, как я могу это сделать?

Возможные способы сделать это:

1. Дублировать мои ядра с разными именами.(Не масштабируется)
2. Передайте некоторый флаг, на основании которого я могу добавить регистры переключения в свое ядро, которое я могу использовать всякий раз, читая / записывая любую область памяти в inData и outData.Это будет означать, что любые временные данные, которые я создаю, также будут преобразованы с использованием такой логики.(что опять-таки вызовет много косвенных ошибок в коде ядра, не знаю, как это повлияет на мою производительность)

Есть ли лучший способ сделать это?Я вижу, что Metal Performance Shaders работают на MTLTexture, которые указывают pixelFormat, и на основе этого pixelFormat, MPS, могут работать с большим диапазоном типов данных.Любые идеи о том, как это сделать?

Спасибо!

Answer 1

Один из подходов, который может сработать, заключается в следующем:

• Объявить inData как void*
• В теле шейдера ядра вызовите функцию шаблона, передаваяаргументы.Функция шаблона будет шаблонизироваться желаемым типом и получит inData в качестве указателя на этот тип.

Вы можете использовать входной параметр, чтобы динамически выбирать, какой вариантфункция шаблона для вызова.Но лучшим подходом, вероятно, является использование константы функции для выбора.Таким образом, выбор компилируется в.

Итак, что-то вроде:

constant int variant [[function_constant(0)]];

template<typename T> void
work(const device void *inData,
     device uint32_t *outData,
     uint2 gid,
     uint2 thread_position_in_threadgroup,
     uint2 threads_per_threadgroup,
     uint2 threadgroup_position_in_grid) 
{
    const device T *data = static_cast<const device T*>(inData);
    // ...
}

kernel void
myKernel(const device void *inData              [[buffer(MyKernelIn)]],
         device uint32_t *outData               [[buffer(MyKernelOut)]],
         uint2 gid                              [[thread_position_in_grid]],
         uint2 thread_position_in_threadgroup   [[thread_position_in_threadgroup]],
         uint2 threads_per_threadgroup          [[threads_per_threadgroup]],
         uint2 threadgroup_position_in_grid     [[threadgroup_position_in_grid]]) 
{
    if (variant == 0)
        work<uint32_t>(inData, outData, gid, thread_position_in_threadgroup,
                       threads_per_threadgroup, threadgroup_position_in_grid);
    else if (variant == 1)
        work<uint8_t>(inData, outData, gid, thread_position_in_threadgroup,
                      threads_per_threadgroup, threadgroup_position_in_grid);
    else
        work<float>(inData, outData, gid, thread_position_in_threadgroup,
                    threads_per_threadgroup, threadgroup_position_in_grid);
}