[KernelGen][MThreads] Add im2col Moore Threads specialized operator

[Yukun-Cui]

[KernelGen][MThreads] Add im2col Moore Threads specialized operator

Summary

Add a Moore Threads (MUSA) specialized Triton kernel for im2col, overriding the generic implementation via runtime.replace_customized_ops(). The kernel gathers sliding local blocks into columns with a 2D tiled grid (tiling over the C*kH*kW rows and the outH*outW columns) and autotuned BLOCK_M/BLOCK_N. All index arithmetic is done in int64 to stay correct on large tensors.

Testing

• Reused the existing upstream accuracy tests tests/test_im2col.py (-m im2col)
• All 48 cases pass on the MUSA device; specialization confirmed active via the GEMS_MTHREADS IM2COL debug log
• Falls back to the generic implementation for unsupported dtype/device
• Test skip guard fix: tests/test_im2col.py hard-coded @pytest.mark.skipif(not torch.cuda.is_available()), which vacuously skips all cases on non-CUDA vendors (MUSA). Replaced it with the device-agnostic flag_gems.runtime.device.device_count == 0 guard, so the tests actually run on Moore Threads while remaining unchanged for CUDA.

Performance

Compared against the generic FlagGems implementation on Moore Threads (MUSA).

im2col

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup
float16	1, 3, 16, 16	0.023520	0.012880	1.826x
float16	1, 3, 16, 16	0.015560	0.014960	1.040x
float16	1, 3, 16, 16	0.020600	0.024320	0.847x
float16	1, 3, 16, 16	0.009160	0.005120	1.789x
float16	1, 3, 32, 32	0.016280	0.012640	1.288x
float16	1, 3, 32, 32	0.016240	0.025240	0.643x
float16	1, 3, 32, 32	0.021920	0.024200	0.906x
float16	1, 3, 32, 32	0.009680	0.005160	1.876x
float16	2, 16, 64, 64	0.029920	0.012360	2.421x
float16	2, 16, 64, 64	0.032040	0.016280	1.968x
float16	2, 16, 64, 64	0.042120	0.029200	1.442x
float16	2, 16, 64, 64	0.018680	0.003880	4.814x
float16	4, 32, 128, 128	0.317600	0.131240	2.420x
float16	4, 32, 128, 128	0.111560	0.080400	1.388x
float16	4, 32, 128, 128	0.248840	0.331880	0.750x
float16	4, 32, 128, 128	0.216840	0.011800	18.376x
float32	1, 3, 16, 16	0.015640	0.013600	1.150x
float32	1, 3, 16, 16	0.016320	0.015000	1.088x
float32	1, 3, 16, 16	0.021400	0.024520	0.873x
float32	1, 3, 16, 16	0.009240	0.004920	1.878x
float32	1, 3, 32, 32	0.016560	0.013240	1.251x
float32	1, 3, 32, 32	0.016600	0.025400	0.654x
float32	1, 3, 32, 32	0.022100	0.024440	0.904x
float32	1, 3, 32, 32	0.009720	0.005040	1.929x
float32	2, 16, 64, 64	0.029640	0.013360	2.219x
float32	2, 16, 64, 64	0.030680	0.016720	1.835x
float32	2, 16, 64, 64	0.041640	0.029360	1.418x
float32	2, 16, 64, 64	0.019120	0.004440	4.306x
float32	4, 32, 128, 128	0.328400	0.146960	2.235x
float32	4, 32, 128, 128	0.123720	0.081140	1.525x
float32	4, 32, 128, 128	0.259840	0.331400	0.784x
float32	4, 32, 128, 128	0.218240	0.020440	10.677x
bfloat16	1, 3, 16, 16	0.015400	0.012360	1.246x
bfloat16	1, 3, 16, 16	0.015640	0.014800	1.057x
bfloat16	1, 3, 16, 16	0.020400	0.024640	0.828x
bfloat16	1, 3, 16, 16	0.009000	0.005000	1.800x
bfloat16	1, 3, 32, 32	0.016160	0.012400	1.303x
bfloat16	1, 3, 32, 32	0.016080	0.025600	0.628x
bfloat16	1, 3, 32, 32	0.021720	0.024200	0.898x
bfloat16	1, 3, 32, 32	0.009560	0.005040	1.897x
bfloat16	2, 16, 64, 64	0.029840	0.012320	2.422x
bfloat16	2, 16, 64, 64	0.031920	0.016280	1.961x
bfloat16	2, 16, 64, 64	0.042000	0.029240	1.436x
bfloat16	2, 16, 64, 64	0.018720	0.003920	4.776x
bfloat16	4, 32, 128, 128	0.318240	0.131120	2.427x
bfloat16	4, 32, 128, 128	0.111760	0.080360	1.391x
bfloat16	4, 32, 128, 128	0.250480	0.331960	0.755x
bfloat16	4, 32, 128, 128	0.217440	0.011800	18.427x

Operator	Arithmetic Mean Speedup
im2col	2.54x

Files Changed

• src/flag_gems/runtime/backend/_mthreads/ops/im2col.py: Moore Threads Triton kernel + fallback
• src/flag_gems/runtime/backend/_mthreads/ops/__init__.py: Register import and __all__
• tests/test_im2col.py: device-agnostic skip guard so the tests run on non-CUDA vendors

[Dongxu-H]

/test | im2col:mthreads

[github-actions]

❌ On-demand test failed

Operator: im2col
Runner: mthreads
Backend: mthreads-musa520

The test failed to complete. Check the workflow run for details.

[Dongxu-H]

LGTM

[Dongxu-H]

LGTM

[Dongxu-H]

LGTM

[Dongxu-H]

/test | im2col:mthreads

[Dongxu-H]

PR contains large infrastructure changes (removing CODEOWNERS, modifying workflows, adding image-builder and sync-to-kernelgen workflows) that should not be mixed with operator implementation. Infrastructure changes should be submitted as separate PR. Only operator-specific changes should remain: src/flag_gems/runtime/backend/_mthreads/ops/im2col.py and tests/test_im2col.py.

[Dongxu-H]

/test | im2col:mthreads

[Dongxu-H]

/test | im2col:mthreads

[Dongxu-H]

LGTM

[Dongxu-H]

CI failures detected: backend-tests (mthreads-musa520) and python-op failed. Please verify the test results on MThreads device match the PR description (48 cases pass).