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

[Yukun-Cui]

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

Summary

Add a Moore Threads (MUSA) specialized Triton kernel for arctan_ (in-place arctangent), overriding the generic implementation via runtime.replace_customized_ops(). The specialization computes atan in float32 via the cross-backend tl_extra_shim.atan shim (not the CUDA-only libdevice.atan) and writes back in the input dtype, with dtype-tuned vectorization (wider VEC for fp16) for L2-resident sizes and an out-of-place streaming kernel for the largest sizes.

Testing

• Reused the existing upstream accuracy tests tests/test_arctan_.py (-m arctan_)
• Validated against reference on the MUSA device; specialization confirmed active via the GEMS_MTHREADS ARCTAN_ debug log
• Falls back to the generic implementation for unsupported dtype/device/shape (fp64/int64 not supported on Moore Threads hardware)

Performance

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

arctan_

dtype	Size	Torch Latency (ms)	Gems Latency (ms)	Speedup
float16	1073741824	3.546000	3.545820	1.000x
float16	64, 64	0.003200	0.003840	0.833x
float16	4096, 4096	0.079920	0.079960	0.999x
float16	64, 512, 512	0.079800	0.079920	0.998x
float16	1024, 1024, 1024	3.547300	3.546500	1.000x
float32	1073741824	6.567620	6.489040	1.012x
float32	64, 64	0.002960	0.004200	0.705x
float32	4096, 4096	0.135960	0.132780	1.024x
float32	64, 512, 512	0.136520	0.133040	1.026x
float32	1024, 1024, 1024	6.559260	6.485000	1.011x
bfloat16	1073741824	3.547180	3.546620	1.000x
bfloat16	64, 64	0.003200	0.003920	0.816x
bfloat16	4096, 4096	0.079840	0.080080	0.997x
bfloat16	64, 512, 512	0.079840	0.079780	1.001x
bfloat16	1024, 1024, 1024	3.547240	3.546200	1.000x

Operator	Arithmetic Mean Speedup
arctan_	0.96x

Files Changed

• src/flag_gems/runtime/backend/_mthreads/ops/arctan_.py: Moore Threads Triton kernel + fallback
• src/flag_gems/runtime/backend/_mthreads/ops/__init__.py: Register import and __all__

[Dongxu-H]

/test | arctan_:mthreads

[Dongxu-H]

/test | arctan_:mthreads

[github-actions]

❌ On-demand test failed

Operator: | arctan_
Runner: mthreads
Backend: mthreads-musa520

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

[Dongxu-H]

/test | arctan_:mthreads

[Dongxu-H]

/test | arctan_:mthreads

[Dongxu-H]

/test | arctan_:mthreads

[Dongxu-H]

/test | arctan_:mthreads

[Dongxu-H]

LGTM

[Dongxu-H]

LGTM

[Dongxu-H]

LGTM