New Interfaces for LibTuner
#771
Conversation
src/flag_gems/utils/libentry.py
Outdated
| lambda idx0, idx1: self.strategy[idx0](args[idx1]), | ||
| enumerate(self.keys), | ||
| ) | ||
| ) + tuple(str(arg.dtype) for arg in args.values() if hasattr(arg, "dtype")) |
There was a problem hiding this comment.
append the data types for both status
| key = tuple(args[k] for k in self.keys if k in args) | ||
| else: | ||
| key = tuple( | ||
| starmap( | ||
| lambda idx0, idx1: self.strategy[idx0](args[idx1]), | ||
| enumerate(self.keys), | ||
| ) | ||
| ) | ||
| key += tuple(str(arg.dtype) for arg in args.values() if hasattr(arg, "dtype")) |
There was a problem hiding this comment.
Is it OK that self.keys is not filtered against args in the else branch?
There was a problem hiding this comment.
I think so. The previous implementation doesn't check args and works well. Do you think I should add it here?
There was a problem hiding this comment.
Not a big issue assuming self.key not including kwargs.
src/flag_gems/utils/libentry.py
Outdated
| """ | ||
|
|
||
| def decorator( | ||
| policy: Callable[ |
There was a problem hiding this comment.
What about renaming it as policy_impl to distinguish from AnonymousLibTunerImpl's policy method. It's bit confusing.
src/flag_gems/utils/libentry.py
Outdated
| args: Tuple[Any], | ||
| kwargs: Dict[str, Any], | ||
| ) -> Tuple[triton.Config, Dict[str, float]]: | ||
| return policy(fn, configs, args, kwargs) |
There was a problem hiding this comment.
Calling policy_impl instead.
| return strategy | ||
|
|
||
| return decorator | ||
|
|
||
| def run(self, *args, **kwargs): | ||
| self.nargs = dict(zip(self.arg_names, args)) |
There was a problem hiding this comment.
Please add a comment for clarification, noting that arg_names correspond to a JITFunction's parameter names in its signature.
There was a problem hiding this comment.
Take a look plz. I'm unsure whether it meets your requirements. Thanks!
| ), f"the length of strategy {len(strategy)} must match the length of keys {len(self.keys)}" | ||
| strategy: List[Callable[[Any], Any]] = [ | ||
| LibTuner.get_strategy(s) if isinstance(s, str) else s for s in strategy | ||
| ] | ||
| self.strategy = strategy | ||
| # Use table name with hash instead of hash in key | ||
| self.kernel_hash = None | ||
| self.table_name = f"{self.__name__}_{self.get_kernel_hash()}" |
There was a problem hiding this comment.
To deal with re-entrance is similar to what jit's _do_compile does in the async_mode, where the cache_key is voluntarily calculated with get_cache_key . I think it's more appropriate than just reusing cache key?
cache_key = get_cache_key(src, backend, options, env_vars)
PR Category
libtuner
Type of Change
New Feature
Description
This is a proposal to provide a new interface to extend the current
LibTuner. It redefinedLibTuneras a base class and allows extending it by overriding thepolicymethod. For easily usage, we also provide a subclass calledOfflineLibTunerso our current implementation, including the default one and new strategies, could be added to the newLibTunersystem bymakestatic method. For the future extension including online autotuner, it could be integrated by providing new subclasses. Totally our design principle is to allow different granularity to control the autotuner details.Except that, we also provide a dispatch mechanism so new policies and
LibTuners could be registered easily into the global dispatcher easily and use them with their name.If this work is approved, future work includes migrate #755, #763, and implement online autotuner #762 on it.
Issue
This is a competitive proposal against #763. Please have a detailed discussion on which design we should work on, or anything we can make better.
Progress
Performance