Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[RUNTIME] Add callback functions for external tools #738

Merged
merged 2 commits into from
Oct 5, 2022
Merged

[RUNTIME] Add callback functions for external tools #738

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e1d3e2a
Invoke callbacks in C; almost no overhead when hooks are not registered
Jokeren Oct 5, 2022
7ca2ea6
Merge remote-tracking branch 'origin/master' into keren/add-callback
Jokeren Oct 5, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 32 additions & 3 deletions python/triton/compiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -959,7 +959,7 @@ def format_of(ty):
"int64_t": "L",
}[ty]

format = "iiiiiKK" + ''.join([format_of(_extracted_type(ty)) for ty in signature.values()])
format = "iiiiiKKOOO" + ''.join([format_of(_extracted_type(ty)) for ty in signature.values()])

# generate glue code
src = f"""
Expand Down Expand Up @@ -1019,14 +1019,37 @@ def format_of(ty):
uint64_t _function;
int num_warps;
int shared_memory;
PyObject *launch_enter_hook = NULL;
PyObject *launch_exit_hook = NULL;
PyObject *compiled_kernel = NULL;
PyObject *hook_ret = NULL;
{' '.join([f"{_extracted_type(ty)} _arg{i}; " for i, ty in signature.items()])}
if(!PyArg_ParseTuple(args, \"{format}\", &gridX, &gridY, &gridZ, &num_warps, &shared_memory, &_stream, &_function, {', '.join(f"&_arg{i}" for i, ty in signature.items())})) {{
if(!PyArg_ParseTuple(args, \"{format}\", &gridX, &gridY, &gridZ, &num_warps, &shared_memory, &_stream, &_function, &launch_enter_hook, &launch_exit_hook, &compiled_kernel, {', '.join(f"&_arg{i}" for i, ty in signature.items())})) {{
return NULL;
}}

if (launch_enter_hook != Py_None) {{
PyObject *new_args = PyTuple_Pack(1, compiled_kernel);
hook_ret = PyObject_CallObject(launch_enter_hook, new_args);
Py_DECREF(new_args);
}}

_launch(gridX, gridY, gridZ, num_warps, shared_memory, (CUstream)_stream, (CUfunction)_function, {', '.join(f"getPointer(_arg{i},{i})" if ty[0]=="*" else f"_arg{i}"for i, ty in signature.items())});

if (launch_exit_hook != Py_None) {{
PyObject *new_args = NULL;
if (hook_ret) {{
new_args = PyTuple_Pack(2, compiled_kernel, hook_ret);
}} else {{
new_args = PyTuple_Pack(1, compiled_kernel);
}}
hook_ret = PyObject_CallObject(launch_exit_hook, new_args);
Py_DECREF(new_args);
}}

if (hook_ret) {{
Py_DECREF(hook_ret);
}}
if(PyErr_Occurred()) {{
return NULL;
}}
Expand Down Expand Up @@ -1242,6 +1265,10 @@ def compile(fn, signature: str, device: int = -1, constants=dict(), num_warps: i

class CompiledKernel:

# Hooks for external tools to monitor the execution of triton kernels
launch_enter_hook = None
launch_exit_hook = None

def __init__(self, fn_name, so_path, cache_dir, device):
# initialize launcher
import importlib.util
Expand All @@ -1267,6 +1294,7 @@ def __init__(self, fn_name, so_path, cache_dir, device):
self.asm["ttir"] = f.read()

mod, func, n_regs, n_spills = _triton.code_gen.load_binary(metadata["name"], self.asm["cubin"], self.shared, device)
self.fn_name = fn_name
self.cu_module = mod
self.cu_function = func
self.n_regs = n_regs
Expand All @@ -1276,7 +1304,8 @@ def __getitem__(self, grid):
def runner(*args, stream=None):
if stream is None:
stream = torch.cuda.current_stream().cuda_stream
self.c_wrapper(grid[0], grid[1], grid[2], self.num_warps, self.shared, stream, self.cu_function, *args)
self.c_wrapper(grid[0], grid[1], grid[2], self.num_warps, self.shared, stream, self.cu_function,
CompiledKernel.launch_enter_hook, CompiledKernel.launch_exit_hook, self, *args)
return runner

def get_sass(self, fun=None):
Expand Down
5 changes: 3 additions & 2 deletions python/triton/runtime/jit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ def launcher(*args, **kwargs):

class JITFunction(KernelInterface):

# Hook for inspecting compiled functions and modules
cache_hook = None
divisibility = 16

Expand Down Expand Up @@ -253,7 +254,7 @@ def {self.fn.__name__}({', '.join(self.arg_names)}, grid, num_warps=4, num_stage
try:
bin = cache[key]
if not warmup:
bin.c_wrapper(grid_0, grid_1, grid_2, bin.num_warps, bin.shared, stream, bin.cu_function, {args})
bin.c_wrapper(grid_0, grid_1, grid_2, bin.num_warps, bin.shared, stream, bin.cu_function, triton.compiler.CompiledKernel.launch_enter_hook, triton.compiler.CompiledKernel.launch_exit_hook, bin, {args})
return bin
# kernel not cached -- compile
except KeyError:
Expand All @@ -274,7 +275,7 @@ def {self.fn.__name__}({', '.join(self.arg_names)}, grid, num_warps=4, num_stage
if not self._call_hook(key, signature, device, constants, num_warps, num_stages, extern_libs, configs):
bin = triton.compile(self, signature, device, constants, num_warps, num_stages, extern_libs=extern_libs, configs=configs)
if not warmup:
bin.c_wrapper(grid_0, grid_1, grid_2, bin.num_warps, bin.shared, stream, bin.cu_function, *args)
bin.c_wrapper(grid_0, grid_1, grid_2, bin.num_warps, bin.shared, stream, bin.cu_function, triton.compiler.CompiledKernel.launch_enter_hook, triton.compiler.CompiledKernel.launch_exit_hook, bin, *args)
self.cache[key] = bin
return bin
return None
Expand Down