How to make templated Compute Kernels in Metal

user_4106624 0
Viewed 400

I have been writing some Metal compute kernels. So, I wrote a kernel with the following declaration:

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]]) 
{ }

Now, I want to write a variant of this which takes inData of type uint8_t and float, how can I do that?

Possible ways I can think of doing this:

  1. Duplicate my Kernels, with different names. (Not scalable)
  2. Pass some flag based on which I can add switch cases to my kernel, which I can use whenever, reading/writing any memory location in inData and outData. This would mean any temporary data that I create to also be casted using such logic. (which would again induce a lot of indirections in Kernel code, not sure how it will impact my performance)

Is there any better way to do this? I see the Metal Performance Shaders working on MTLTexture, which specify pixelFormat, and based on that pixelFormat, MPS, can work on a large range of data types. Any insights on how that is done?

Thanks!

metal gpgpu gpu macos
2 Answers

One approach that might work is to:

  • Declare inData as a void*
  • In the body of the kernel shader, call a template function, passing along the arguments. The template function would be templated by the desired type and would receive inData as a pointer to that type.

You could use an input parameter to dynamically choose which variant of the template function to call. But a better approach is probably to use a function constant to pick. That way, the choice is compiled in.

So, something like:

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);
}
user_13121430

Section 5.1.11 from the Metal Spec. You can do this:

template<typename T> kernel void
myKernel(device const T *inData [[buffer(MyKernelIn)]],
         device 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]]) {
    ...
}

template [[ host_name("myKernel_float") ]] kernel void
ReduceKernel<float>(device const float*, 
                    device float*, 
                    uint2, 
                    uint2, 
                    uint2, 
                    uint2);

template [[ host_name("myKernel_uint32") ]] kernel void
ReduceKernel<uint32_t>(device const uint32_t*, 
                       device uint32_t*, 
                       uint2, 
                       uint2, 
                       uint2, 
                       uint2);
user_11701280
Related