torch.cuda.jiterator._create_multi_output_jit_fn

torch.cuda.jiterator._create_multi_output_jit_fn(code_string, num_outputs, **kwargs)

创建一个由jiterator生成的cuda内核，用于支持返回一个或多个输出的元素操作。

Parameters

• code_string (str) – 要由jiterator编译的CUDA代码字符串。入口函数必须通过引用返回值。

• num_outputs (int) – 内核返回的输出数量

• kwargs (字典, 可选) – 生成函数的键值参数

Return type

可调用

示例：

code_string = "template  void my_kernel(T x, T y, T alpha, T& out) { out = -x + alpha * y; }"
jitted_fn = create_jit_fn(code_string, alpha=1.0)
a = torch.rand(3, device='cuda')
b = torch.rand(3, device='cuda')
# 像调用普通python函数一样调用jitted函数
result = jitted_fn(a, b, alpha=3.14)

警告

此API目前处于测试阶段，可能会在未来的版本中进行更改。

警告

此API仅支持最多8个输入和8个输出