torch.cuda.jiterator._create_jit_fn

torch.cuda.jiterator._create_jit_fn(code_string, **kwargs)

建立一個 jiterator 生成的針對逐元素運算的 cuda 核心。

code_string 必須是一個有效的 CUDA 函式，它描述了單個元素的計算。code_string 必須遵循 c++ 模板模式，如以下示例所示。此函式將被內聯到逐元素核心模板中，並即時編譯。編譯後的核心將快取在記憶體和本地臨時目錄中。

Jiterator 生成的核心支援非連續張量，並支援廣播和型別提升。

引數

• code_string (str) – 將由 jiterator 編譯的 CUDA 程式碼字串。入口函子必須透過值返回。

• kwargs (Dict, optional) – 為生成的函式提供的關鍵字引數

返回型別

Callable

示例

code_string = "template <typename T> T my_kernel(T x, T y, T alpha) { return -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)

code_string 也允許定義多個函式，最後一個函式將被視為入口函式。

示例

code_string = (
    "template <typename T> T util_fn(T x, T y) { return ::sin(x) + ::cos(y); }"
)
code_string += "template <typename T> T my_kernel(T x, T y, T val) { return ::min(val, util_fn(x, y)); }"
jitted_fn = create_jit_fn(code_string, val=0.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)  # using default val=0.0

Jiterator 可以與 Python 註冊一起使用，以覆蓋運算元的 cuda 核心。以下示例使用 relu 覆蓋了 gelu 的 cuda 核心。

示例

code_string = "template <typename T> T my_gelu(T a) { return a > 0 ? a : 0; }"
my_gelu = create_jit_fn(code_string)
my_lib = torch.library.Library("aten", "IMPL")
my_lib.impl("aten::gelu", my_gelu, "CUDA")
# torch.nn.GELU and torch.nn.function.gelu are now overridden
a = torch.rand(3, device="cuda")
torch.allclose(torch.nn.functional.gelu(a), torch.nn.functional.relu(a))

警告

此 API 處於 Beta 版，未來版本中可能會更改。

警告

此 API 最多僅支援 8 個輸入和 1 個輸出

警告

所有輸入張量必須位於 CUDA 裝置上