pyopencl - 如何使用泛型类型?

问题描述

我使用 32 位浮点数和 32 位整数可互换地工作。我想要两个完全相同的内核,但一个用于整数,一个用于浮点数。一开始我以为我可以使用模板什么的,但似乎无法指定两个名称相同但参数类型不同的内核?

import pyopencl as cl
import numpy as np

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)

prg = cl.Program(ctx,"""
__kernel void arange(__global int *res_g)
{
  int gid = get_global_id(0);
  res_g[gid] = gid;
}

__kernel void arange(__global float *res_g)
{
  int gid = get_global_id(0);
  res_g[gid] = gid;
}
""").build()

错误

<kernel>:8:15: error: conflicting types for 'arange'
__kernel void arange(__global float *res_g)
              ^
<kernel>:2:15: note: prevIoUs deFinition is here
__kernel void arange(__global int *res_g)

最方便的方法是什么?

解决方法

#define 指令可用于:

code = """ 
__kernel void arange(__global TYPE *res_g)
{
  int gid = get_global_id(0);
  res_g[gid] = gid;
}
"""

prg_int = cl.Program(ctx,code).build("-DTYPE=int")
prg_float = cl.Program(ctx,code).build("-DTYPE=float")