python - 如何在 pyopencl 中创建可变大小的 __local 内存？

Question

在我的 C OpenCL 代码中，我使用 clSetKernelArg 创建“可变大小”__local 内存以用于我的内核，这在 OpenCL 本身中不可用。看我的例子:

clSetKernelArg(clKernel, ArgCounter++, sizeof(cl_mem), (void *)&d_B);
...
clSetKernelArg(clKernel, ArgCounter++, sizeof(float)*block_size*block_size, NULL);
...
kernel="    
    matrixMul(__global float* C,
                   ...
              __local float* A_temp,
                  ...
             )"
{...

我现在的问题是，如何在 pyopencl 中做同样的事情？

我查看了 pyopencl 附带的示例，但我唯一能找到的是使用模板的方法，据我所知，这似乎是一种矫枉过正的做法。参见示例。

kernel = """
        __kernel void matrixMul(__global float* C,...){
                     ... 
            __local float A_temp[ %(mem_size) ];
                     ...
        }

你有什么建议？

Answer 1

它类似于 C。你将一个固定大小的数组作为本地传递给它。这是 Enja 基数排序的示例。注意最后一个参数是本地内存数组。

 def naive_scan(self, num):
    nhist = num/2/self.cta_size*16
    global_size = (nhist,)
    local_size = (nhist,)
    extra_space = nhist / 16 #NUM_BANKS defined as 16 in RadixSort.cpp
    shared_mem_size = self.uintsz * (nhist + extra_space)
    scan_args = (   self.mCountersSum,
                    self.mCounters,
                    np.uint32(nhist),
                    cl.LocalMemory(2*shared_mem_size)
                )
    self.radix_prg.scanNaive(self.queue, global_size, local_size, *(scan_args)).wait()