-
Notifications
You must be signed in to change notification settings - Fork 49
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.
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
Implement __triton_launcher
as pure DLL
#3251
Conversation
Signed-off-by: Anatoly Myachev <[email protected]>
class TritonLauncher: | ||
|
||
def __init__(self, cache_path: str): | ||
self.shared_library = ctypes.PyDLL(cache_path) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Must use PyDLL
instead of CDLL
(as in d7d55b8) for python api to work correctly, otherwise I get segfaults.
|
||
def __call__(self, *args, **kwargs): | ||
# Serialize KernelArguments for SPIR-V Runner | ||
serialize_kernel_args = os.getenv('TRITON_XPU_DUMP_SPIRV_KERNEL_ARGS', None) | ||
if serialize_kernel_args: | ||
serialize_args(args, self.constants, self.signature) | ||
self.launch(*args, **kwargs) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If the kwargs
were not empty, the launcher would not start without issues, since according to its C signature it supports only args
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Previously, the code unpacked both args and kwargs into a single dictionary and passed that to the launcher - now it looks like we're passing the args
dict which seems incorrect as any values in kwargs
would be ignored. Could we create a new dictionary that contains both args
and kwargs
and pass that to maintain the previous behavior?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Previously, the code unpacked both args and kwargs into a single dictionary and passed that to the launcher - now it looks like we're passing the args dict which seems incorrect as any values in kwargs would be ignored.
I tried to set the dictionary and the launcher immediately crashes. Are you sure that kwargs were ever used before?
> self.launch(*args, **{"test": 2})
E TypeError: launch() takes no keyword arguments
Seems only tuple is expected:
if(!PyArg_ParseTuple(args, \"{format}\", &gridX, &gridY, &gridZ, &py_obj_stream, &py_kernel, |
Perhaps we need to make a cleanup in the triton code...
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It looks like the calling classes support them: https://github.com/intel/intel-xpu-backend-for-triton/blob/main/python/triton/compiler/compiler.py#L432
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It looks like the calling classes support them: https://github.com/intel/intel-xpu-backend-for-triton/blob/main/python/triton/compiler/compiler.py#L432
In this line, all parameters will be packed into one tuple, since parameters are not passed as ..., name_param=param, ...
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@alexbaden do you mind if I merge it?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ok, interesting - I am sure your Python knowledge is better than mine! :)
Do you think this is something we should propose changing upstream? Or is this specific to our backend?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do you think this is something we should propose changing upstream? Or is this specific to our backend?
Similar code exists only for AMD backend. I suggested removing unused code in triton-lang/triton#5694.
Just for reference how it works for NVIDIA backend:
self.launch(gridX, gridY, gridZ, stream, function, self.launch_cooperative_grid, global_scratch, *args) |
@@ -635,15 +649,14 @@ def __init__(self, src, metadata): | |||
self.constants = {arg_idx(idx): value for idx, value in constants.items()} | |||
self.signature = {idx: value for idx, value in src.signature.items()} | |||
src = make_launcher(self.constants, self.signature) | |||
mod = compile_module_from_src(src, "__triton_launcher") | |||
self.launch = mod.launch | |||
self.mod = compile_module_from_src(src, "__triton_launcher") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
To keep a reference to so/dll
and not call the destructor prematurely.
With unloaded DLL libraries, these changes are no longer necessary. However, two tests that hold a reference to the compiled kernel need to be adjusted - manually clear the cache (inside `JITFunction` object). CI: * https://github.com/intel/intel-xpu-backend-for-triton/actions/runs/12951798922 (passed) * https://github.com/intel/intel-xpu-backend-for-triton/actions/runs/12955820922 (check status) Blocked on #3251 Extra refs: * python/cpython#87319 --------- Signed-off-by: Anatoly Myachev <[email protected]>
Test CI:
https://github.com/intel/intel-xpu-backend-for-triton/actions/runs/12935201507(wrong branch)Benchmarks: