Fix ValueError: numel (2097152) exceeds triton maximum tensor numel (1048576)

[mengluy0125]

Summary: We add block_size heuristics for autotune since it can have ValueError for tensors with large shape.

Differential Revision: D82742888

[facebook-github-bot]

@mengluy0125 has exported this pull request. If you are a Meta employee, you can view the originating diff in D82742888.

[facebook-github-bot]

@mengluy0125 has exported this pull request. If you are a Meta employee, you can view the originating diff in D82742888.

[jansel]

IMO capping the size on a per-dimension basis isn't ideal. Since if you have multiple dimensions you need the product of the block sizes to be less than the limit. So the cap is not dimension independent.

Maybe we could improve the heuristics in:

helion/helion/autotuner/config_generation.py

Lines 90 to 113 in 0a33998

	def shrink_config(
	self, flat_config: FlatConfig, max_elements_per_thread: int
	) -> None:
	"""
	Fully random configs tend to run out of resources and tile a long time to compile.
	Here we shrink the config to a reasonable size.
	Args:
	flat_config: config to mutate in place
	max_elements_per_thread: maximum number of elements per thread
	"""
	num_threads = warps_to_threads(cast("int", flat_config[self.num_warps_index]))
	max_elements = max_elements_per_thread * num_threads
	while self.block_numel(flat_config) > max_elements:
	changes = 0
	for i in self.block_size_indices:
	val = flat_config[i]
	assert isinstance(val, int)
	threshold = max(self.flat_spec[i].get_minimum(), self.min_block_size)
	if val // 2 >= threshold:
	flat_config[i] = val // 2
	changes += 1
	if changes == 0:
	break

to address this problem?

[mengluy0125]

IMO capping the size on a per-dimension basis isn't ideal. Since if you have multiple dimensions you need the product of the block sizes to be less than the limit. So the cap is not dimension independent.

Maybe we could improve the heuristics in:

helion/helion/autotuner/config_generation.py

Lines 90 to 113 in 0a33998

	def shrink_config(
	self, flat_config: FlatConfig, max_elements_per_thread: int
	) -> None:
	"""
	Fully random configs tend to run out of resources and tile a long time to compile.
	Here we shrink the config to a reasonable size.
	Args:
	flat_config: config to mutate in place
	max_elements_per_thread: maximum number of elements per thread
	"""
	num_threads = warps_to_threads(cast("int", flat_config[self.num_warps_index]))
	max_elements = max_elements_per_thread * num_threads
	while self.block_numel(flat_config) > max_elements:
	changes = 0
	for i in self.block_size_indices:
	val = flat_config[i]
	assert isinstance(val, int)
	threshold = max(self.flat_spec[i].get_minimum(), self.min_block_size)
	if val // 2 >= threshold:
	flat_config[i] = val // 2
	changes += 1
	if changes == 0:
	break

to address this problem?

Oh, make sense. Let me do the heuristics improvement inside the function.

[facebook-github-bot]

@mengluy0125 has exported this pull request. If you are a Meta employee, you can view the originating diff in D82742888.

[oulgen]

So I think this is overcounting, exact same problem I had on #485

self.block_numel(flat_config) > 2 ** 20

doesn't actually mean in the triton code we emit, we will end up with a tensor of numel 2 ** 20.

[jansel]

I agree, though maybe this is better than what we have now.

[jansel]

I agree, though maybe this is better than what we have now.

[yf225]

Merging it to fix benchmark CI issues.