[Iluvatar][TLE] Sync iluvatar backend and extend TLE primitives support

[huatuoli]

[Iluvatar][TLE] Sync iluvatar backend and extend TLE primitives support

This PR syncs the backend to Iluvatar Triton commit d10a6779d9a and extends Iluvatar TLE lowering with async copy, pipeline, cumsum, memory_space, warp_specialize, and pipe support.

Included commits:

• 78ab6ee36c1 [Iluvatar] Update iluvatar backend to commit d10a6779d9a.
• acd391c7ff8 [Iluvatar][TLE] Update for more tle support.
  • Support async copy && tle.load, memory_space, cumsum and pipeline.
  • Support tle.gpu.warp_specialize && tle.pipe (currently supported on ivcore11).
  • Optimize TLE tutorial kernels in 01-fft.py and 03-topk.py.
• 8e31f05050d [CI/CD] Update iluvatar3.6 ci runner.

Main changes:

• Sync third_party/iluvatar backend to Iluvatar Triton d10a6779d9a: SME load path, load/store LLVM lowering, AccelerateMatmul, pipeline, FP8, integration tests (vLLM / SGLang / FBGEMM / Inductor), and expanded CI in test_triton.sh.
• Extend Iluvatar TLE with tle.load (async), tle.gpu.memory_space, tle.cumsum, async tle.gpu.copy, tle.gpu.pipeline, and tle.gpu.warp_specialize + tle.pipe (currently supported on ivcore11).
• Add 7 new TLE optimization/lowering passes, software mbarrier + warp-specialize LLVM lowering, and Iluvatar-specific TLE unit/integration tests.
• Add and optimize Iluvatar TLE tutorial benchmarks 01-fft.py and 03-topk.py (radix-select TopK) under third_party/iluvatar/python/tutorials/tle/, including CoreX-optimized kernel variants for performance evaluation.

Diff stats: 136 files changed, +17,687 / -1,493 (across both commits).

---

1. Overview

1.1 Backend sync (78ab6ee36c1)

Syncs third_party/iluvatar to Iluvatar Triton commit d10a6779d9a. Major areas:

Area	Changes
SME / matmul	add_matmul_smeload pass; SmeLoad.cpp rewrite; large test_sme.py suite.
Load/store LLVM	LoadStoreOpToLLVM.cpp expansion; new MembarUtility.
TTGIR transforms	AccelerateMatmul, pipeliner (AssignLatencies / LowerLoops), RemoveLayoutConversions, etc.
Integration tests	vLLM (Punica LoRA, chunk_o, wy_fast), SGLang FlashMLA, FBGEMM jagged FA, Inductor bucketize/flex-attention.
CI	test_triton.sh expanded with new language/operator/runtime/regression/integration suites.

1.2 TLE extension (acd391c7ff8)

Continues to reuse the shared TLE Python frontend under python/triton/experimental/tle/ and extends backend-specific builder bindings and lowering via iluvatar_tle.

• Python API: shared (no third_party/iluvatar/python/.../tle override).
• Builder overrides: third_party/iluvatar/tle/triton_iluvatar_tle.cc.
• MLIR dialect: iluvatar_tle (third_party/iluvatar/tle/include/IR/IluvatarTleOps.td).
• Optimization / lowering passes: 10 iluvatar-specific TLE passes (3 from [Iluvatar][TLE] Triton v3.6.x iluvatar backend and 5 TLE primitives support #724 + 7 new).
• LLVM lowering: TleToLLVM/ (extract/insert/local_pointers) + BarrierOpToLLVM.cpp + ConvertWarpSpecializeToLLVM.cpp.

---

2. Supported primitives

Covers the 5 primitives from #724 plus 6 new / extended primitives in this PR.

2.1–2.5 Primitives from #724

Primitive	Status
tle.gpu.alloc	Direct ttg.local_alloc; smem only.
tle.gpu.local_ptr	iluvatar_tle.local_pointers; integer indices; smem only.
tle.gpu.copy	Extended — GM→local now fuses to async copy (see §2.6).
tle.extract_tile	Register-shuffle or shared-memory relay.
tle.insert_tile	Register-merge or shared-memory relay.

2.6 tle.gpu.copy (async GM→local)

• Normal copy (tensor ↔ buffered_tensor): GM→local staging (tt.load + local_pointers tt.store) is recognized and rewritten to ttg.async_copy_global_to_local + commit/wait by triton-iluvatar-tle-optimize-local-pointer-async-stores.
• Local→GM: unchanged synchronous path via local_ptr → tl.load → tl.store.

2.7 tle.load (async)

• Frontend: tle.load(..., is_async=True) annotates tt.load with tt.load.async = true.
• Lowering: triton-iluvatar-tle-lower-async-load rewrites supported hints to ttg.async_copy_global_to_local + commit/wait + local_load.
• Coverage: pointer loads and block-pointer loads; fp32 / fp16 / bf16.

2.8 tle.gpu.memory_space

• Frontend: tle.gpu.memory_space(tensor, "shared_memory") marks tt.memory_space = "shared_memory".
• Lowering: triton-iluvatar-tle-early-assign-memory-space materializes before layout/MMA transforms:
  • Load producers: async global→shared copy path.
  • Other producers: initialized local_alloc + local_load.
• Limitations: only "shared_memory" / smem storage supported.

2.9 tle.cumsum

• MLIR: iluvatar_tle.exclusive_cumsum (create_exclusive_cumsum).
• Coverage: exclusive prefix sum + total; forward and reverse (reverse=True); 1D masked and 2D axis cases.
• Lowering: layout fold pass → tt.scan (add) + tt.reduce (total); exclusive = scan − src.
• Limitations: lowered through native Triton scan/reduce; dtype/shape constraints follow upstream Triton scan support.

2.10 tle.gpu.pipeline

• Frontend: for i in tle.gpu.pipeline(start, stop, step, num_stages=N) propagates tt.num_stages onto the generated scf.for.
• Coverage: loop iterator hint for software pipelining; numerically equivalent to tl.range.
• Integration: plugs into existing TTGIR add_pipeline pass when num_stages > 1.

2.11 tle.gpu.warp_specialize (currently supported on ivcore11)

• MLIR: ttg.warp_specialize via create_warp_specialize / create_warp_specialize_partitions / create_warp_return / create_warp_yield.
• Coverage: default + worker partitions; explicit captures; return values via warp_yield; concurrent partition execution.
• LLVM lowering: ConvertWarpSpecializeToLLVM.cpp — software named-barrier protocol.
• Limitations:
  • Currently validated on Iluvatar corex (ivcore11).
  • Worker partitions using block tensors should match default warp count for layout consistency.

2.12 tle.pipe (currently supported on ivcore11)

• MLIR ops: iluvatar_tle.pipe.create / writer_acquire|commit|close / reader_wait|release, plus init_barrier / wait_barrier / arrive_barrier.
• Coverage: CTA-scoped SPSC/SPMC local-store pipes used with ttg.warp_specialize; MVP producer/consumer e2e (tile accumulate).
• Lowering: triton-iluvatar-tle-lower-pipe-to-barriers → software mbarrier ring in shared memory → BarrierOpToLLVM.cpp.
• Limitations:
  • Currently supported on ivcore11.
  • Intended for warp-specialized producer/consumer patterns.

---

3. Lowering path

Key differences from the native Triton path:

Python TLE API
  alloc              -> ttg.local_alloc                          (direct)
  local_ptr          -> iluvatar_tle.local_pointers
  copy (GM->local)   -> tt.load + local_pointers store
                         -> [pass] async_copy_global_to_local    (NEW)
  copy (local->GM)   -> local_pointers load + tt.store           (sync)
  tle.load(async)    -> tt.load{async=true}
                         -> [pass] async_copy + local_load       (NEW)
  memory_space       -> tt.memory_space="shared_memory"
                         -> [pass] async copy or local_alloc     (NEW)
  cumsum             -> iluvatar_tle.exclusive_cumsum
                         -> [pass] tt.scan + tt.reduce           (NEW)
  pipeline           -> scf.for + tt.num_stages hint             (NEW)
  warp_specialize    -> ttg.warp_specialize                      (NEW, currently ivcore11)
  pipe               -> iluvatar_tle.pipe.*                      (NEW, currently ivcore11)
                         -> [pass] init/wait/arrive_barrier
  extract_tile       -> iluvatar_tle.extract_tile
  insert_tile        -> iluvatar_tle.insert_tile

TTGIR passes (iluvatar-specific, in pipeline order):
  triton-iluvatar-tle-optimize-local-pointer-async-stores   (NEW)
  triton-iluvatar-tle-early-assign-memory-space             (NEW)
  triton-iluvatar-tle-optimize-exclusive-cumsum-layouts     (NEW)
  triton-iluvatar-tle-lower-exclusive-cumsum                (NEW)
  triton-iluvatar-tle-insert-local-pointer-barriers         (#724)
  triton-iluvatar-tle-optimize-local-pointer-loads          (#724)
  triton-iluvatar-tle-optimize-local-pointer-stores         (#724)
  triton-iluvatar-tle-lower-pipe-to-barriers                (NEW)
  triton-iluvatar-tle-lower-async-load                      (NEW, post-TTGIR)

make_llir:
  populateTleToLLVMPatterns  (extract_tile / insert_tile / local_pointers)
  add_warp_specialize_to_llvm + BarrierOpToLLVM               (NEW, currently ivcore11)
  -> register shuffle / merge, shared-memory relay, or SW mbarrier -> LLVM

---

4. Tests

4.1 Shared TLE tests (FlagTree main tree)

Test
python/test/tle/integration/test_tle_local_store.py
python/test/tle/unit/test_tle_gpu_local_ptr.py
python/test/tle/unit/test_extract_tile_static_index.py
python/test/tle/unit/test_extract_tile_dynamic_index.py
python/test/tle/unit/test_insert_tile_static_index.py
python/test/tle/unit/test_insert_tile_dynamic_index.py
python/test/tle/unit/test_tle.py

4.2 Iluvatar-specific TLE tests (NEW in acd391c7ff8)

Test	Coverage
python/test/unit/tle/test_tle_copy.py	Async GM→local copy codegen + roundtrip
python/test/unit/tle/test_tle_async_load.py	tle.load(is_async=True) pointer / block_ptr
python/test/unit/tle/test_tle_memory_space.py	memory_space load vs non-load paths
python/test/unit/tle/test_tle_cumsum.py	Builder binding, TTIR/TTGIR lowering, runtime
python/test/unit/tle/test_tle_pipeline.py	tle.gpu.pipeline num_stages hint
python/test/unit/tle/test_tle_pipeline_e2e.py	Full alloc/copy/local_ptr/pipeline workflow
python/test/unit/tle/test_tle_warp_specialize.py	WS IR lowering + tle.gpu.warp_specialize e2e
python/test/unit/tle/test_tle_pipe.py	tle.pipe + warp_specialize producer/consumer e2e

---

5. Performance data

5.1 Measurement methodology

• Benchmark sources (Iluvatar-adapted TLE tutorials):
  • third_party/iluvatar/python/tutorials/tle/01-fft.py — complex FFT; compares Triton / TLE / TLE-CoreX / Torch.
  • third_party/iluvatar/python/tutorials/tle/03-topk.py — radix-select TopK; compares Triton-RadixSelect / RadixSelect-CoreX / Triton-TopK / Torch.

5.2 Environment

Field	Value
Hardware	Iluvatar Corex (ivcore11)
Driver / SDK	4.5.0
Torch	2.10.0
FlagTree	triton_v3.6.x_iluvatar branch

5.3 Representative results

5.3.1 FFT (01-fft.py, M=4096)

Correctness check passed (triton/tle).

	N	Triton (ms)	TLE (ms)	TLE-CoreX (ms)	Torch (ms)
0	64	0.075858	0.129763	0.049754	0.025894
1	128	0.092721	0.143077	0.059877	0.039996
2	256	0.146613	0.201750	0.117835	0.058927
3	512	0.310842	0.408877	0.253158	0.123965
4	1024	0.688510	0.972823	0.616052	0.224296

Speedup is computed as baseline_time / TLE_time:

Comparison	Mean
TLE FFT vs Triton FFT	0.69x
TLE-CoreX FFT vs Triton FFT	1.33x

5.3.2 TopK (03-topk.py)

Correctness check passed (radix + radix_corex + triton).

Note: Triton-RadixSelect and Triton-RadixSelect-CoreX use TLE primitives.

	M	N	K	Triton-RadixSelect (ms)	Triton-RadixSelect-CoreX (ms)	Triton-TopK (ms)	Torch-TopK (ms)
0	64	128	8	0.014785	0.012646	0.008335	0.019863
1	64	1024	32	0.036340	0.017712	0.012942	0.029604
2	64	8192	128	0.195965	0.054663	0.061383	0.136331
3	128	32768	256	1.503065	0.400435	0.572658	0.688038

Speedup is computed as baseline_time / CoreX_time, grouped by sequence length N:

N ≤ 1024 (rows 0–1: N = 128, 1024):

Comparison	Mean
RadixSelect-CoreX vs RadixSelect	1.61x
RadixSelect-CoreX vs Triton-TopK	0.69x
RadixSelect-CoreX vs Torch-TopK	1.62x

N > 1024 (rows 2–3: N = 8192, 32768):

Comparison	Mean
RadixSelect-CoreX vs RadixSelect	3.67x
RadixSelect-CoreX vs Triton-TopK	1.28x
RadixSelect-CoreX vs Torch-TopK	2.11x

CoreX optimizations show limited gain at small N (RadixSelect-CoreX is still slower than Triton-TopK when N ≤ 1024), but deliver strong speedups at large N (up to 3.75x vs baseline RadixSelect at N = 32768).

---

6. Status note

• 78ab6ee36c1 delivers a broad iluvatar backend sync (SME, matmul, pipeline, integrations, CI) aligned with Iluvatar Triton d10a6779d9a.
• acd391c7ff8 delivers functional support for 6 additional TLE primitives / extensions on top of [Iluvatar][TLE] Triton v3.6.x iluvatar backend and 5 TLE primitives support #724, with correctness validated by iluvatar-specific unit/integration tests.
• tle.gpu.warp_specialize and tle.pipe are currently supported and validated on ivcore11 (software mbarrier emulation); support on other Iluvatar architectures may follow in future commits.
• TLE-corex optimized paths (async copy, memory_space, pipeline) show measurable gains on representative workloads: FFT 1.33x vs native Triton; TopK RadixSelect-CoreX 3.67x vs baseline RadixSelect and 2.11x vs Torch when N > 1024. Base TLE paths without CoreX optimizations remain slower than native Triton on FFT; further tuning is planned for follow-up commits.

[CLAassistant]

All committers have signed the CLA.

[zhzhcookie]

LGTM