torch.cuda.jiterator._create_jit_fn

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

Create a jiterator-generated cuda kernel for an elementwise op.

The code string has to be a valid CUDA function that describes the computation for a single element. The code string has to follow the c++ template pattern, as shown in the example below. This function will be inlined into elementwise kernel template, and compiled on the fly. Compiled kernel will be cached in memory, as well as local temp dir.

Jiterator-generated kernels accepts noncontiguous tensors, and supports boardcasting and type promotion.

Parameters

• code_string (string) – CUDA code string to be compiled by jiterator.

• kwargs (Dict, optional) – Keyword arguments for generated function

Example

>>> 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)

Jiterator can be used together with python registration to override an operator’s cuda kernel

Following example is overriding gelu’s cuda kernel with relu:

>>> 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))

Warning

This API is in beta and may change in future releases.

Warning

Jiterator only supports up to 8 tensor inputs

Warning

All input tensors must live in CUDA device