Method __init__

models/custom_autotune.py:15–49

:param prune_configs_by: a dict of functions that are used to prune configs, fields: 'perf_model': performance model used to predicate running time with different configs, returns running time 'top_k': number of configs to bench 'prune_num_stages_by'(optional): a function used to prune n

(self, fn, arg_names, configs, key, reset_to_zero, prune_configs_by: Dict = None, nearest_power_of_two: bool = False)

class Autotuner(triton.KernelInterface):
	def __init__(self, fn, arg_names, configs, key, reset_to_zero, prune_configs_by: Dict = None, nearest_power_of_two: bool = False):
		'''
		:param prune_configs_by: a dict of functions that are used to prune configs, fields:
			'perf_model': performance model used to predicate running time with different configs, returns running time
			'top_k': number of configs to bench
			'prune_num_stages_by'(optional): a function used to prune num_stages. It take configs:List[Config] as its input, and returns pruned configs.
			'nearest_power_of_two'(optional): whether to round key arguments to the nearest power of two when caching tuning results
		'''
		if not configs:
			self.configs = [triton.Config({}, num_warps=4, num_stages=2)]
		else:
			self.configs = configs
		self.key_idx = [arg_names.index(k) for k in key]
		self.nearest_power_of_two = nearest_power_of_two
		self.cache = {}
		# hook to reset all required tensor to zeros before relaunching a kernel
		self.hook = lambda args: 0
		if reset_to_zero is not None:
			self.reset_idx = [arg_names.index(k) for k in reset_to_zero]

			def _hook(args):
				for i in self.reset_idx:
					args[i].zero_()
			self.hook = _hook
		self.arg_names = arg_names
		# prune configs
		if prune_configs_by:
			perf_model, top_k = prune_configs_by['perf_model'], prune_configs_by['top_k']
			if 'early_config_prune' in prune_configs_by:
				early_config_prune = prune_configs_by['early_config_prune']
		else:
			perf_model, top_k, early_config_prune = None, None, None
		self.perf_model, self.configs_top_k = perf_model, top_k
		self.early_config_prune = early_config_prune
		self.fn = fn

	def _bench(self, *args, config, **meta):
		# check for conflicts, i.e. meta-parameters both provided