[KMCompiler] Add optimized index_copy and index_copy_ kernels for Ascend

[Onisen7]

PR Category

Operator

Type of Change

Performance Optimization

Description

This PR adds Ascend-specific Triton implementations of index_copy and
index_copy_.

The implementation is added in
src/flag_gems/runtime/backend/_ascend/ops/index_copy_.py and exported from
src/flag_gems/runtime/backend/_ascend/ops/__init__.py.

The Ascend implementation:

• uses a flat kernel for workloads with inner_size <= 4;
• uses a row-oriented 2D kernel for wider contiguous suffixes;
• uses 64-bit offset arithmetic for large tensors;
• keeps runtime dimensions out of Triton specialization to reduce unnecessary
  recompilation;
• selects the flat-kernel block size according to the workload size;
• skips kernel launches for empty index or inner dimensions;
• launches kernels under the input device context;
• masks invalid indices to prevent out-of-bounds stores; and
• supports both out-of-place index_copy and in-place index_copy_.

Issue

N/A

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue. (N/A: no linked issue)
• Change is fully covered by a UT.

Testing

Test environment:

• Device: Ascend 910B
• CANN: 8.5.0
• Python: 3.11.13
• flagtree: 0.6.0+ascend3.2

Functional test command:

python -m pytest -s tests/test_index_copy_.py

Result: 24 passed.

The test covers both index_copy and index_copy_ with the default dtype,
shape, and dimension parameterization in tests/test_index_copy_.py.

Performance

Benchmark command:

python -m pytest -s benchmark/test_index_copy.py

The default benchmark ran in kernel mode at the comprehensive level.
Speedup = PyTorch latency / FlagGems latency.

The generic implementation completed the functional test with 24 passed,
but did not complete the default Ascend benchmark. The new Ascend-specific
implementation completed all 48 benchmark rows: 2 APIs, 3 dtypes, and 8 input
configurations.

Operator	float16 mean	float32 mean	bfloat16 mean	Overall mean
index_copy	1.045x	0.872x	1.051x	0.989x
index_copy_	1.430x	1.345x	1.428x	1.401x
Combined	-	-	-	1.195x

The values above are arithmetic means of the configured benchmark rows.
Performance varies by shape, so the table does not imply that every individual
configuration is faster than the native implementation. Full per-shape logs
are available if additional review is needed.