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.
Must use PyDLL instead of CDLL (as in d7d55b8) for python api to work correctly, otherwise I get segfaults.
| 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.
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.
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.
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 argumentsSeems only tuple is expected:
Perhaps we need to make a cleanup in the triton code...
There was a problem hiding this comment.
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.
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.
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.
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:
| 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.
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]>
The same motivation as for #3251 and #3230 --------- Signed-off-by: Anatoly Myachev <[email protected]>
Test CI:
https://github.com/intel/intel-xpu-backend-for-triton/actions/runs/12935201507(wrong branch)Benchmarks: