[Iluvatar][TLE] Sync iluvatar backend and extend TLE primitives support (#880)

* [Iluvatar] Update iluvatar backend to commit d10a6779d9a.

* [Iluvatar][TLE] Update for more tle support.
- Support async copy && tle.load, memory_space, cumsum and pipeline.
- Support tle.gpu.warp_specialize && tle.pipe only for ivcore11.
- Opt 01-fft && 03-topk.

* [CI/CD] Update iluvatar3.6 ci runner.

Author: huatuoli

.github/workflows/iluvatar3.6-build-and-test.yml

@@ -84,7 +84,7 @@ jobs:
 
           pip3 install "pytest==9.1.1"
 
-          for pkg in hypothesis absl-py scipy pytest-forked; do
+          for pkg in hypothesis absl-py scipy pytest-forked expecttest; do
             pip3 list "$pkg" | grep "$pkg" || pip3 install "$pkg"
           done
 
@@ -95,16 +95,77 @@ jobs:
             fi
           fi
 
-          # triton unit test
+          # triton language unit test
           UMD_CUDAMODULELOADING=0 pytest -v third_party/iluvatar/python/test/unit/language/test_core.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_annotations.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_block_pointer.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_compile_errors.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_conversions.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_decorator.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_frontend.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_libdevice.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_iluvatar_math_fp16_bf16.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_iluvatar_int8_upcast_dot_pipeline.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_matmul.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_mxfp.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_pipeliner.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_random.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_reproducer.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_standard.py
           pytest -v third_party/iluvatar/python/test/unit/language/test_subprocess.py
+          UMD_CUDAMODULELOADING=0 pytest -v third_party/iluvatar/python/test/unit/language/test_tensor_descriptor.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_tuple.py
+          pytest -v third_party/iluvatar/python/test/unit/language/test_warp_specialization.py
+
+          # triton operators unit test
           pytest -v third_party/iluvatar/python/test/unit/operators/test_blocksparse.py
           pytest -v third_party/iluvatar/python/test/unit/operators/test_cross_entropy.py
           pytest -v third_party/iluvatar/python/test/unit/operators/test_dot_trans.py
           pytest -v third_party/iluvatar/python/test/unit/operators/test_flash_attention.py
           pytest -v third_party/iluvatar/python/test/unit/operators/test_inductor.py
           pytest -v third_party/iluvatar/python/test/unit/operators/test_matmul.py
+          pytest -v third_party/iluvatar/python/test/unit/operators/test_sme.py
+
+          # triton runtime unit test
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_autotuner.py
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_bindings.py
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_build.py
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_cache.py
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_compilation_listener.py
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_driver.py
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_launch.py
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_specialize.py
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_subproc.py
+          pytest -v third_party/iluvatar/python/test/unit/runtime/test_iluvatar_loop_unroll_warning.py
+
+          # triton instrumentation unit test
+          pytest -v third_party/iluvatar/python/test/unit/instrumentation/test_gpuhello.py
+
+          # triton tools unit test
+          pytest -v third_party/iluvatar/python/test/unit/tools/test_irsource.py
+          pytest -v third_party/iluvatar/python/test/unit/tools/test_linear_layout.py
+
+          # triton misc unit test
+          pytest -v third_party/iluvatar/python/test/unit/test_debug_dump.py
           pytest -v third_party/iluvatar/python/test/unit/test_debug.py
+          pytest -v third_party/iluvatar/python/test/unit/test_debuginfo.py
+          pytest -v third_party/iluvatar/python/test/unit/test_filecheck.py
+          pytest -v third_party/iluvatar/python/test/unit/test_knobs.py
+          UMD_CUDAMODULELOADING=0 pytest -v third_party/iluvatar/python/test/unit/test_link.py
+          pytest -v third_party/iluvatar/python/test/unit/test_perf_warning.py
+
+          # triton regression test
+          pytest -v third_party/iluvatar/python/test/regression/test_cast_matmul.py
+          pytest -v third_party/iluvatar/python/test/regression/test_functional_regressions.py
+
+          # triton integrations unit test
+          pytest -v third_party/iluvatar/python/test/unit/integrations/vllm/chunk_o/test_chunk_fwd_kernel_o.py
+          pytest -v third_party/iluvatar/python/test/unit/integrations/vllm/wy_fast/test_recompute_w_u.py
+          pytest -v third_party/iluvatar/python/test/unit/integrations/fbgemm/test_jagged_flash_attention_bwd_basic_min.py
+          PUNICA_TEST_LEVEL=quick pytest -v third_party/iluvatar/python/test/unit/integrations/vllm/punica_lora/test_punica_ops.py
+          pytest -v third_party/iluvatar/python/test/unit/integrations/sglang/flash_mla/test_flash_mla_ut.py
+          pytest -v third_party/iluvatar/python/test/unit/integrations/inductor/test_bucketize_matmul.py
+          pytest -v third_party/iluvatar/python/test/unit/integrations/inductor/test_swfw3103_flex_attention_precision.py
 
           # TLE unit test
           pytest -v python/test/tle/integration/test_tle_local_store.py
@@ -113,6 +174,16 @@ jobs:
           pytest -v python/test/tle/unit/test_extract_tile_dynamic_index.py
           pytest -v python/test/tle/unit/test_insert_tile_static_index.py
           pytest -v python/test/tle/unit/test_insert_tile_dynamic_index.py
+          pytest -v python/test/tle/unit/test_tle.py
+          pytest -v third_party/iluvatar/python/test/unit/tle/test_tle_copy.py
+          pytest -v third_party/iluvatar/python/test/unit/tle/test_tle_async_load.py
+          pytest -v third_party/iluvatar/python/test/unit/tle/test_tle_memory_space.py
+          pytest -v third_party/iluvatar/python/test/unit/tle/test_tle_cumsum.py
+          pytest -v third_party/iluvatar/python/test/unit/tle/test_tle_pipeline.py
+          pytest -v third_party/iluvatar/python/test/unit/tle/test_tle_pipeline_e2e.py
+          pytest -v third_party/iluvatar/python/test/unit/tle/test_tle_warp_specialize.py
+          pytest -v third_party/iluvatar/python/test/unit/tle/test_tle_pipe.py
 
           # TLE tutorials unit tests
-          python3 python/tutorials/tle/01-fft.py --only_unit_test
+          python3 third_party/iluvatar/python/tutorials/tle/01-fft.py --only_unit_test
+          python3 third_party/iluvatar/python/tutorials/tle/03-topk.py --only_unit_test

python/setup_tools/setup_helper.py

@@ -563,6 +563,8 @@ def uninstall_triton():
 
 handle_flagtree_backend()
 
+register_backend_cache()
+
 # iluvatar
 cache.store(
     file="iluvatar-llvm22-x86_64",
@@ -572,13 +574,6 @@ def uninstall_triton():
     post_hook=set_llvm_env,
 )
 
-cache.store(
-    file="iluvatarTritonPlugin.so", condition=("iluvatar" == flagtree_backend) and (not configs.flagtree_plugin), url=
-    "https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/iluvatarTritonPlugin-cpython3.10-glibc2.30-glibcxx3.4.28-cxxabi1.3.12-ubuntu-x86_64_v0.3.0.tar.gz",
-    copy_dst_path=f"third_party/{flagtree_backend}", md5_digest="015b9af8")
-
-register_backend_cache()
-
 # mthreads
 cache.store(
     file="mthreads-llvm22",

python/triton/experimental/tle/language/gpu/core.py

@@ -954,10 +954,16 @@ def normcopy(
 
         try:
             if direction == CopyDirection.GM_TO_LOCAL:
-                # None fills the FlagTree hints slot; TLE copy has no hints to pass.
-                load_extra_args = () if (mthreads_enabled or iluvatar_enabled) else (None, )
-                tt_load = _semantic.load(src, mask, other, boundary_check, padding_option, cache_modifier,
-                                         eviction_policy, volatile, *load_extra_args)
+                if iluvatar_enabled:
+                    # Iluvatar's semantic.load carries an extra `stride` (SME) slot
+                    # right after `other`; TLE copy never uses the SME path.
+                    tt_load = _semantic.load(src, mask, other, None, boundary_check, padding_option, cache_modifier,
+                                             eviction_policy, volatile)
+                else:
+                    # None fills the FlagTree hints slot; TLE copy has no hints to pass.
+                    load_extra_args = () if mthreads_enabled else (None, )
+                    tt_load = _semantic.load(src, mask, other, boundary_check, padding_option, cache_modifier,
+                                             eviction_policy, volatile, *load_extra_args)
                 local_ptrs = local_ptr(dst, _make_full_indices(dst, _semantic), _semantic=_semantic)
                 _semantic.store(local_ptrs, tt_load, mask, boundary_check, cache_modifier, eviction_policy)
             else:

third_party/iluvatar/backend/compiler.py

@@ -16,6 +16,15 @@
 from pathlib import Path
 
 
+def has_tle_pass(pass_name: Optional[str] = None) -> bool:
+    tle_passes = getattr(iluvatar.passes, "tle", None)
+    if tle_passes is None:
+        return False
+    if pass_name is None:
+        return True
+    return hasattr(tle_passes, pass_name)
+
+
 def min_dot_size(target: GPUTarget):
 
     def check_dot_compatibility(lhs_type, rhs_type) -> Tuple[int, int, int]:  # [m, n, k]
@@ -106,7 +115,7 @@ def sm_arch_from_capability(capability: int):
 
 
 @dataclass(frozen=True)
-class CUDAOptions:
+class CorexOptions:
     num_warps: int = 4
     num_ctas: int = 1
     num_stages: int = 3
@@ -121,7 +130,9 @@ class CUDAOptions:
     enable_reflect_ftz: bool = True  # ftz in libdevice
     launch_cooperative_grid: bool = False
     launch_pdl: bool = False
-    supported_fp8_dtypes: Tuple[str] = ("fp8e5", "fp8e4b15")
+    # OCP only: fp8e5 always; fp8e4nv added for sm71 (ivcore11 SW) / sm89+.
+    # fp8e4b15 is NVIDIA PTX-asm software format — not supported on Iluvatar.
+    supported_fp8_dtypes: Tuple[str] = ("fp8e5", )
     deprecated_fp8_dot_operand_dtypes: Tuple[str] = ()
     default_dot_input_precision: str = "tf32"
     allowed_dot_input_precisions: Tuple[str] = ("tf32", "tf32x3", "ieee", 'bf16x3', 'bf16x6')
@@ -152,12 +163,12 @@ def hash(self):
         return hashlib.sha256(key.encode("utf-8")).hexdigest()
 
 
-class CUDABackend(BaseBackend):
+class CorexBackend(BaseBackend):
     instrumentation = None
 
     @staticmethod
     def supports_target(target: GPUTarget):
-        return target.backend == 'corex'
+        return target.backend in ('corex', 'cuda')
 
     def _parse_arch(self, arch):
         pattern = r"^sm(\d+)$"
@@ -171,6 +182,8 @@ def get_target_name(self, options) -> str:
         return f"cuda:{capability}"
 
     def __init__(self, target: GPUTarget) -> None:
+        if target.backend == 'cuda':
+            target = GPUTarget('corex', target.arch, target.warp_size)
         super().__init__(target)
         self.binary_ext = "cubin"
 
@@ -180,7 +193,7 @@ def parse_options(self, opts) -> Any:
             opts["debug"] = True
 
         args = {'arch': knobs.runtime.override_arch or f"sm{self.target.arch}"}
-        args.update({k: opts[k] for k in CUDAOptions.__dataclass_fields__.keys() if k in opts if opts[k] is not None})
+        args.update({k: opts[k] for k in CorexOptions.__dataclass_fields__.keys() if k in opts if opts[k] is not None})
         capability = int(self._parse_arch(args["arch"]))
 
         if args.get("num_ctas", 1) > 1 and capability < 90:
@@ -189,21 +202,18 @@ def parse_options(self, opts) -> Any:
                               f"Please set num_ctas=1 or target an SM90+ GPU."))
 
         if "supported_fp8_dtypes" not in args:
-            supported_fp8_dtypes = set(CUDAOptions.supported_fp8_dtypes)
-            if capability >= 89:
+            supported_fp8_dtypes = set(CorexOptions.supported_fp8_dtypes)
+            # ivcore11 (sm71): software FP8; ivcore30 will use native HW when mapped.
+            if capability >= 89 or capability == 71:
                 supported_fp8_dtypes.add("fp8e4nv")
             args["supported_fp8_dtypes"] = tuple(sorted(supported_fp8_dtypes))
 
-        if "deprecated_fp8_dot_operand_dtypes" not in args:
-            if capability >= 90:
-                args["deprecated_fp8_dot_operand_dtypes"] = ("fp8e4b15", )
-
         if "enable_fp_fusion" not in args:
             args["enable_fp_fusion"] = knobs.language.default_fp_fusion
 
         args["max_num_imprecise_acc_default"] = 2**30 if capability == 90 else 0
 
-        return CUDAOptions(**args)
+        return CorexOptions(**args)
 
     def pack_metadata(self, metadata):
         return (
@@ -228,8 +238,8 @@ def get_module_map(self) -> Dict[str, ModuleType]:
 
     def load_dialects(self, ctx):
         iluvatar.load_dialects(ctx)
-        if CUDABackend.instrumentation:
-            CUDABackend.instrumentation.load_dialects(ctx)
+        if CorexBackend.instrumentation:
+            CorexBackend.instrumentation.load_dialects(ctx)
 
     @staticmethod
     def make_ttir(mod, metadata, opt, capability):
@@ -263,15 +273,26 @@ def make_ttgir(mod, metadata, opt, capability):
         passes.ttgpuir.add_f32_dot_tc(pm, emuTF32)
         passes.ttgpuir.add_remove_layout_conversions(pm)
         passes.ttgpuir.add_optimize_thread_locality(pm)
-        if hasattr(iluvatar.passes, "tle"):
+        if has_tle_pass():
+            if has_tle_pass("add_optimize_local_pointer_async_stores"):
+                iluvatar.passes.tle.add_optimize_local_pointer_async_stores(pm)
+            if has_tle_pass("add_early_assign_memory_space"):
+                iluvatar.passes.tle.add_early_assign_memory_space(pm)
+            if has_tle_pass("add_optimize_exclusive_cumsum_layouts"):
+                iluvatar.passes.tle.add_optimize_exclusive_cumsum_layouts(pm)
+            if has_tle_pass("add_lower_exclusive_cumsum"):
+                iluvatar.passes.tle.add_lower_exclusive_cumsum(pm)
             iluvatar.passes.tle.add_insert_local_pointer_barriers(pm)
             iluvatar.passes.tle.add_optimize_local_pointer_loads(pm)
             iluvatar.passes.tle.add_optimize_local_pointer_stores(pm)
+            if has_tle_pass("add_lower_pipe_to_barriers"):
+                iluvatar.passes.tle.add_lower_pipe_to_barriers(pm)
         iluvatar.passes.ttgpuir.add_accelerate_matmul(pm, opt.use_sme)
         passes.ttgpuir.add_remove_layout_conversions(pm)
         iluvatar.passes.ttgpuir.add_mma_reduce_thread_locality(pm)
         iluvatar.passes.ttgpuir.add_optimize_epilogue(pm)
         passes.ttgpuir.add_optimize_dot_operands(pm, capability >= 71)
+        iluvatar.passes.ttgpuir.add_matmul_smeload(pm, capability)
         passes.ttir.add_loop_aware_cse(pm)
         if capability // 10 in [7, 8, 9]:
             passes.ttgpuir.add_fuse_nested_loops(pm)
@@ -281,6 +302,7 @@ def make_ttgir(mod, metadata, opt, capability):
             passes.ttgpuir.add_combine_tensor_select_and_if(pm)
             passes.ttgpuir.add_assign_latencies(pm, opt.num_stages)
             passes.ttgpuir.add_schedule_loops(pm)
+            passes.ttgpuir.add_pipeline(pm, opt.num_stages, dump_enabled)
         elif capability // 10 >= 10:
             passes.ttgpuir.add_fuse_nested_loops(pm)
             passes.common.add_canonicalizer(pm)
@@ -299,10 +321,11 @@ def make_ttgir(mod, metadata, opt, capability):
             passes.ttir.add_triton_licm(pm)
         passes.common.add_canonicalizer(pm)
         passes.ttir.add_loop_aware_cse(pm)
-        iluvatar.passes.ttgpuir.add_matmul_smeload(pm, capability)
         passes.ttgpuir.add_remove_layout_conversions(pm)
         passes.ttgpuir.add_prefetch(pm)
         passes.ttgpuir.add_optimize_dot_operands(pm, capability >= 71)
+        if has_tle_pass("add_lower_async_load"):
+            iluvatar.passes.tle.add_lower_async_load(pm)
         passes.ttgpuir.add_coalesce_async_copy(pm)
         passes.ttgpuir.add_remove_layout_conversions(pm)
         passes.ttgpuir.add_reduce_data_duplication(pm)
@@ -350,12 +373,15 @@ def make_llir(self, src, metadata, options, capability):
             # Call ConcurrencySanitizerPass here, before allocating global scratch memory but after allocating tensor and shared
             passes.ttgpuir.add_concurrency_sanitizer(pm)
         passes.ttgpuir.add_allocate_global_scratch_memory(pm)
-        if CUDABackend.instrumentation:
-            CUDABackend.instrumentation.patch("ttgpuir_to_llvmir", pm, mod.context)
+        if CorexBackend.instrumentation:
+            CorexBackend.instrumentation.patch("ttgpuir_to_llvmir", pm, mod.context)
         proc = sm_arch_from_capability(capability)
         iluvatar.passes.ttgpuir.add_to_llvmir(pm, proc, options.enable_reflect_ftz)
         passes.common.add_canonicalizer(pm)
         passes.common.add_cse(pm)
+        # [WA] On ivcore11 this relies on a shared-memory software barrier
+        # because the architecture lacks hardware named barriers and setmaxnreg.
+        iluvatar.passes.ttgpuir.add_warp_specialize_to_llvm(pm, proc)
         passes.common.add_canonicalizer(pm)
         passes.common.add_cse(pm)
         passes.common.add_symbol_dce(pm)
@@ -364,8 +390,8 @@ def make_llir(self, src, metadata, options, capability):
         if not knobs.compilation.disable_line_info and not knobs.compilation.dump_ir_extract_di_local_variables:
             passes.llvmir.add_di_scope(pm)
 
-        if CUDABackend.instrumentation:
-            CUDABackend.instrumentation.patch("llvmir_to_llvm", pm, mod.context)
+        if CorexBackend.instrumentation:
+            CorexBackend.instrumentation.patch("llvmir_to_llvm", pm, mod.context)
 
         pm.run(mod, 'make_llir')
 

third_party/iluvatar/backend/driver.py

@@ -51,11 +51,11 @@ def library_dirs():
 # ------------------------
 
 
-class CudaUtils(object):
+class CorexUtils(object):
 
     def __new__(cls):
         if not hasattr(cls, "instance"):
-            cls.instance = super(CudaUtils, cls).__new__(cls)
+            cls.instance = super(CorexUtils, cls).__new__(cls)
         return cls.instance
 
     def __init__(self):
@@ -674,7 +674,7 @@ def inner(*args):
     return inner
 
 
-class CudaLauncher(object):
+class CorexLauncher(object):
 
     def __init__(self, src, metadata):
         constants = src.constants if hasattr(src, "constants") else dict()
@@ -717,11 +717,11 @@ def allocate_scratch(size, align, allocator):
                     global_scratch, profile_scratch, *args)
 
 
-class CudaDriver(GPUDriver):
+class CorexDriver(GPUDriver):
 
     def __init__(self):
-        self.utils = CudaUtils()  # TODO: make static
-        self.launcher_cls = CudaLauncher
+        self.utils = CorexUtils()  # TODO: make static
+        self.launcher_cls = CorexLauncher
         super().__init__()
 
     def get_current_target(self):