torch.cuda.jiterator._create_multi_output_jit_fn

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

创建一个 jiterator 生成的 CUDA 核函数，用于支持返回一个或多个输出的元素级操作。

参数

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

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

• kwargs (Dict, optional) – 生成函数的关键字参数

返回类型

Callable

示例

code_string = "template <typename T> 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")
# invoke jitted function like a regular python function
result = jitted_fn(a, b, alpha=3.14)

警告

此 API 处于 Beta 版，未来版本中可能会更改。

警告

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