快捷方式

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, 可选) – 生成的函数的关键字参数

返回类型

可调用对象

示例

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 处于测试阶段,可能在将来的版本中发生更改。

警告

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源