[Fix][Kernel] Fix shared memory error in test_convolution.py

[3402956340]

Some test cases failed with 'Shared memory size error'.
Reduce num_stages to fit within shared memory limits.

[gemini-code-assist]

Code Review

This pull request reduces the default num_stages to 1 in two convolution kernel configurations within tileops/kernels/convolution.py. The reviewer points out that while this change prevents shared memory limit errors on lower-end devices, it causes performance regressions on high-end GPUs. To address this, the reviewer suggests dynamically determining the optimal num_stages based on the device's shared memory limit.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Hardcoding num_stages to 1 avoids shared memory limit errors on devices with smaller shared memory limits (such as MetaX MACA or SM86 GPUs), but it introduces a performance regression on high-end GPUs (like A100/H100) where num_stages = 2 fits comfortably and provides better pipelining performance.\n\nWe can dynamically determine the optimal num_stages by checking the device's shared memory limit using get_shared_memory_limit_bytes() and conv_shared_memory_bytes().

Suggested change

	"num_stages": 1,
	"num_stages": 2 if conv_shared_memory_bytes(64, 128, 128, 2, self.dtype) <= get_shared_memory_limit_bytes() else 1,

[gemini-code-assist]

Hardcoding num_stages to 1 avoids shared memory limit errors on devices with smaller shared memory limits (such as MetaX MACA), but it introduces a performance regression on standard NVIDIA GPUs (like RTX 3090/4090 or A100) where num_stages = 3 fits comfortably and provides better pipelining performance.\n\nWe can dynamically determine the optimal num_stages by checking the device's shared memory limit using get_shared_memory_limit_bytes() and conv_shared_memory_bytes().

Suggested change

	"num_stages": 1,
	"num_stages": 3 if conv_shared_memory_bytes(64, 256, 32, 3, self.dtype) <= get_shared_memory_limit_bytes() else (2 if conv_shared_memory_bytes(64, 256, 32, 2, self.dtype) <= get_shared_memory_limit_bytes() else 1),

[Five-HZ]

LGTM