From e86732ca1d3b521effbf784cd73fc4eaca69b2fe Mon Sep 17 00:00:00 2001 From: macqueen09 Date: Tue, 11 Aug 2026 03:53:41 +0800 Subject: [PATCH 1/2] Route TLE GPU through TileIR conversion --- .../experimental/tle/language/gpu/core.py | 79 ++-- .../experimental/tle/language/gpu/types.py | 41 +- .../triton/experimental/tle/language/pipe.py | 11 +- test/CommonIR/native_flash_gpu.py | 16 +- test/CommonIR/native_matmul_gpu.py | 16 +- test/CommonIR/tle_gpu_primitives.py | 30 +- third_party/nvidia/backend/compiler.py | 7 +- third_party/tle/dialect/include/IR/TleOps.td | 12 +- .../tle/dialect/include/Transforms/Passes.td | 19 + third_party/tle/dialect/lib/IR/CMakeLists.txt | 4 + third_party/tle/dialect/lib/IR/Ops.cpp | 25 +- .../tle/dialect/lib/Transforms/CMakeLists.txt | 5 + .../Transforms/TleConvertGpuTileToTtgir.cpp | 411 ++++++++++++++++++ third_party/tle/triton_tle.cc | 163 +------ 14 files changed, 553 insertions(+), 286 deletions(-) create mode 100644 third_party/tle/dialect/lib/Transforms/TleConvertGpuTileToTtgir.cpp diff --git a/python/triton/experimental/tle/language/gpu/core.py b/python/triton/experimental/tle/language/gpu/core.py index c8a6ec370b..f5f22ecfed 100644 --- a/python/triton/experimental/tle/language/gpu/core.py +++ b/python/triton/experimental/tle/language/gpu/core.py @@ -1,5 +1,4 @@ # flagtree tle -import os import builtins import triton.language.core as tl from typing import Optional, Sequence @@ -18,8 +17,8 @@ SHARED_MEMORY_ADDRESS_SPACE = 3 -def _tileir_mode() -> bool: - return os.environ.get("TLE_GPU_TILEIR_MODE") == "1" +def _tileir_available(builder) -> bool: + return hasattr(builder, "create_tile_alloc") def _tile_storage_name(storage: tle.scope) -> str: @@ -27,7 +26,7 @@ def _tile_storage_name(storage: tle.scope) -> str: return "shared" if storage == tle.tmem: return "local" - raise ValueError(f"Storage type {storage} not supported by TileIR mode") + raise ValueError(f"Storage type {storage} not supported by TileIR") def _tile_shape(shape) -> list[int]: @@ -268,19 +267,6 @@ def alloc( dtype = tl._unwrap_if_constexpr(dtype) elem_type = dtype.to_ir(_semantic.builder) - if _tileir_mode(): - if not hasattr(_semantic.builder, "create_tile_alloc"): - raise RuntimeError("TLE GPU TileIR mode requires a TileIR-enabled TLE builder") - memory_space = _semantic.builder.tile_get_string_attr(_tile_storage_name(storage)) - tile_ty = _semantic.builder.tile_get_buffer_type(_tile_shape(unwrapped_shape), elem_type, memory_space) - tensor_handle = _semantic.builder.create_tile_alloc(tile_ty) - if layout is None: - if storage == tle.smem: - layout = tle.nv_mma_shared_layout.make_default(shape, dtype) - else: - layout = tle.tensor_memory_layout.make_default(shape) - return tle.buffered_tensor(tensor_handle, dtype, unwrapped_shape, storage, layout, _semantic) - if layout is None: if storage == tle.smem: if not nv_mma_shared_layout: @@ -320,6 +306,15 @@ def alloc( # Use provided layout layout_handle = layout.to_ir(_semantic.builder) + if _tileir_available(_semantic.builder): + memory_space = _semantic.builder.tile_get_string_attr(_tile_storage_name(storage)) + tile_ty = _semantic.builder.tile_get_buffer_type(_tile_shape(unwrapped_shape), elem_type, memory_space) + tensor_handle = _semantic.builder.create_tile_alloc(tile_ty, layout_handle) + result = tle.buffered_tensor(tensor_handle, dtype, unwrapped_shape, storage, layout, _semantic) + if init_value is not None: + _semantic.builder.create_tile_store_tensor(init_value.handle, tensor_handle) + return result + if storage == tle.smem: if init_value is not None: mutable_ty = _semantic.builder.get_memdesc_type(full_shape, elem_type, layout_handle, "smem") @@ -526,15 +521,13 @@ def tmacopy( shape = tuple(shape) else: raise ValueError(f"Shape parameter must be tuple or list, but got {type(shape)}") - if _tileir_mode(): - if not hasattr(_semantic.builder, "create_tile_copy"): - raise RuntimeError("TLE GPU TileIR mode requires a TileIR-enabled TLE builder") - src_handle = src.handle - dst_handle = dst.handle - _semantic.builder.create_tile_copy(src_handle, dst_handle, [], False) - if isinstance(dst, tle.buffered_tensor): - dst._tle_gpu_tile_src = src - dst._tle_gpu_last_copy_shape = shape + if hasattr(_semantic.builder, "create_tile_copy"): + tile_indices = [] + if not is_normcopy: + if offsets is None: + raise ValueError("TMA copy requires offsets") + tile_indices = _semantic._convert_to_ir_values(offsets, require_i64=False) + _semantic.builder.create_tile_copy(src.handle, dst.handle, tile_indices, False) return if is_normcopy: return normcopy(src, dst, shape, direction, _semantic) @@ -558,9 +551,7 @@ def to_tensor( shape = list(memref.shape if target_shape is None else target_shape) writable = tl._unwrap_if_constexpr(writable) - if _tileir_mode(): - if not hasattr(_semantic.builder, "create_tile_to_tensor"): - raise RuntimeError("TLE GPU TileIR mode requires a TileIR-enabled TLE builder") + if hasattr(_semantic.builder, "create_tile_to_tensor"): result = _semantic.builder.create_tile_to_tensor(memref.handle, bool(writable)) return tl.tensor(result, tl.block_type(memref.dtype, shape)) @@ -573,11 +564,8 @@ def store_tensor(tensor_value: tl.tensor, dst: tle.buffered_tensor, _semantic=No if not isinstance(dst, tle.buffered_tensor): raise ValueError(f"dst must be tle.gpu.buffered_tensor, got {type(dst).__name__}") - if _tileir_mode(): - if not hasattr(_semantic.builder, "create_tile_store_tensor"): - raise RuntimeError("TLE GPU TileIR mode requires a TileIR-enabled TLE builder") + if hasattr(_semantic.builder, "create_tile_store_tensor"): _semantic.builder.create_tile_store_tensor(tensor_value.handle, dst.handle) - dst._tle_gpu_tensor_value = tensor_value return ptrs = local_ptr(dst, _make_full_indices(dst, _semantic), _semantic=_semantic) @@ -691,26 +679,6 @@ def local_ptr( import warnings warnings.warn("TLE semantic analysis module not available, skipping validation", UserWarning) - if _tileir_mode(): - if not hasattr(_semantic.builder, "create_tile_to_tensor"): - raise RuntimeError("TLE GPU TileIR mode requires a TileIR-enabled TLE builder") - source = buffer.handle - result_shape = list(buffer_shape if view_shape is None else view_shape) - if not no_indices: - if not all_scalar_indices: - raise RuntimeError("TLE GPU TileIR local_ptr mode currently supports only full or scalar views") - offsets = [idx.handle for idx in idx_tensors] - sizes = [] - strides = [] - source = _semantic.builder.create_tile_subview(buffer.handle, offsets, sizes, strides) - result_shape = [] - value = _semantic.builder.create_tile_to_tensor(source, False) - if result_shape: - result_ty = tl.block_type(buffer.dtype, result_shape) - else: - result_ty = buffer.dtype - return tl.tensor(value, result_ty) - ptr_dtype = tl.pointer_type(buffer.type.element_ty, SHARED_MEMORY_ADDRESS_SPACE) insert_block = _semantic.builder.get_insertion_block() if insert_block is None: @@ -729,7 +697,10 @@ def local_ptr( result_ty = tl.block_type(ptr_dtype, list(view_shape)) result_ir = result_ty.to_ir(_semantic.builder) handles = [idx.handle for idx in idx_tensors] - local_ptr_op = _semantic.builder.create_local_pointers(result_ir, buffer.handle, *handles) + if hasattr(_semantic.builder, "create_tile_local_ptr"): + local_ptr_op = _semantic.builder.create_tile_local_ptr(result_ir, buffer.handle, *handles) + else: + local_ptr_op = _semantic.builder.create_local_pointers(result_ir, buffer.handle, *handles) result_tensor = tl.tensor(local_ptr_op.get_result(0), result_ty) diff --git a/python/triton/experimental/tle/language/gpu/types.py b/python/triton/experimental/tle/language/gpu/types.py index 9dd3c5cc6d..ff7fd0a3eb 100644 --- a/python/triton/experimental/tle/language/gpu/types.py +++ b/python/triton/experimental/tle/language/gpu/types.py @@ -1,6 +1,5 @@ # flagtree tle -import os import triton.language.core as tl from typing import Optional, List, Tuple from abc import abstractmethod @@ -43,21 +42,6 @@ def _storage_to_tile_space(storage: scope) -> str: raise ValueError(f"Unsupported TLE TileIR storage: {storage}") -def _tileir_mode() -> bool: - return os.environ.get("TLE_GPU_TILEIR_MODE") == "1" - - -_TILE_PIPE_V = 1 -_TILE_PIPE_S = 6 -_TILE_EVENT_READY = 0 -_TILE_EVENT_CLOSE = 1 - - -def _tile_sync_available(builder) -> bool: - return (_tileir_mode() and hasattr(builder, "create_tile_pipe_barrier") - and hasattr(builder, "create_tile_set_flag") and hasattr(builder, "create_tile_wait_flag")) - - class layout: def __init__(self): @@ -336,14 +320,13 @@ def slot(self, stage, _semantic=None): slot_layout = _make_slot_layout(self.type.layout, slot_shape) slot_ty = buffered_tensor_type(self.dtype, slot_shape, self.type.storage, slot_layout, _semantic, alloc_shape=slot_shape) - if _tileir_mode(): - if not hasattr(_semantic.builder, "create_tile_subview"): - raise RuntimeError("TLE GPU TileIR mode requires a TileIR-enabled TLE builder") + if hasattr(_semantic.builder, "create_tile_subview"): offsets = [stage_tensor.handle] for _ in range(len(self.shape) - 1): offsets.append(_semantic.to_tensor(0).handle) slot_handle = _semantic.builder.create_tile_subview(self.handle, offsets, slot_shape, - [1] * len(slot_shape)) + [1] * len(slot_shape), + slot_layout.to_ir(_semantic.builder)) return buffered_tensor(slot_handle, self.dtype, slot_shape, self.type.storage, slot_layout, _semantic, alloc_shape=slot_ty.alloc_shape) slot_handle = _semantic.builder.create_memdesc_index(slot_ty.to_ir(_semantic.builder), self.handle, @@ -442,7 +425,7 @@ def _flatten_ir_types(self, builder: ir.builder, out: List[ir.type]) -> None: def to_ir(self, builder: ir.builder) -> None: shape = self.shape builder = self.semantic.builder - if _tileir_mode() and hasattr(builder, "tile_get_buffer_type"): + if hasattr(builder, "tile_get_buffer_type"): memory_space = builder.tile_get_string_attr(_storage_to_tile_space(self.storage)) return builder.tile_get_buffer_type( [int(tl._unwrap_if_constexpr(dim)) for dim in shape], @@ -788,9 +771,6 @@ def __init__(self, pipe: pipe_value, reader_name=None, field_names=None): @tl.builtin def acquire(self, iter, _semantic=None): stage, phase = self.pipe._stage_phase(iter, _semantic=_semantic) - if _tile_sync_available(_semantic.builder): - _semantic.builder.create_tile_pipe_barrier(_TILE_PIPE_S) - return self.pipe._make_slot(stage, _semantic=_semantic) _semantic.builder.create_pipe_writer_acquire(self.pipe._field_handles(), stage.handle, phase.handle, self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names()) @@ -799,9 +779,6 @@ def acquire(self, iter, _semantic=None): @tl.builtin def commit(self, iter, _semantic=None): stage, _ = self.pipe._stage_phase(iter, _semantic=_semantic) - if _tile_sync_available(_semantic.builder): - _semantic.builder.create_tile_set_flag(_TILE_PIPE_S, _TILE_PIPE_V, _TILE_EVENT_READY) - return _semantic.builder.create_pipe_writer_commit(self.pipe._field_handles(), stage.handle, self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names()) @@ -810,9 +787,6 @@ def close(self, iter, _semantic=None): if self.pipe.one_shot: raise ValueError("tle.pipe one_shot pipes do not support close") stage, phase = self.pipe._stage_phase(iter, _semantic=_semantic) - if _tile_sync_available(_semantic.builder): - _semantic.builder.create_tile_set_flag(_TILE_PIPE_S, _TILE_PIPE_V, _TILE_EVENT_CLOSE) - return _semantic.builder.create_pipe_writer_close(self.pipe._field_handles(), stage.handle, phase.handle, self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names()) @@ -830,10 +804,6 @@ def _reader_field_names(self): @tl.builtin def wait(self, iter, _semantic=None): stage, phase = self.pipe._stage_phase(iter, _semantic=_semantic) - if _tile_sync_available(_semantic.builder): - _semantic.builder.create_tile_wait_flag(_TILE_PIPE_S, _TILE_PIPE_V, _TILE_EVENT_READY) - slot = self.pipe._make_slot(stage, _semantic=_semantic, field_names=self.field_names) - return pipe_wait_result(slot, _semantic.to_tensor(False)) is_closed = _semantic.builder.create_pipe_reader_wait(self.pipe._field_handles(), stage.handle, phase.handle, self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names(), self.reader_name or "", @@ -844,9 +814,6 @@ def wait(self, iter, _semantic=None): @tl.builtin def release(self, iter, _semantic=None): stage, _ = self.pipe._stage_phase(iter, _semantic=_semantic) - if _tile_sync_available(_semantic.builder): - _semantic.builder.create_tile_set_flag(_TILE_PIPE_V, _TILE_PIPE_S, _TILE_EVENT_READY) - return _semantic.builder.create_pipe_reader_release(self.pipe._field_handles(), stage.handle, self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names(), self.reader_name or "", diff --git a/python/triton/experimental/tle/language/pipe.py b/python/triton/experimental/tle/language/pipe.py index 0f1c823362..a40f455292 100644 --- a/python/triton/experimental/tle/language/pipe.py +++ b/python/triton/experimental/tle/language/pipe.py @@ -1,15 +1,9 @@ # flagtree tle -import os - import triton.language.core as tl from .gpu import types as gpu_types -def _tileir_mode() -> bool: - return os.environ.get("TLE_GPU_TILEIR_MODE") == "1" - - def _unwrap_pipe_constexpr(value): if isinstance(value, tl.constexpr): value = value.value @@ -109,9 +103,8 @@ def pipe( raise ValueError( f"tle.pipe field {field_name!r} leading dimension must equal capacity {capacity}, got {field.shape[0]}") - if not _tileir_mode(): - _semantic.builder.create_pipe_create([field.handle for field in fields.values()], capacity, scope, name or "", - list(fields.keys()), list(reader_names or ()), one_shot) + _semantic.builder.create_pipe_create([field.handle for field in fields.values()], capacity, scope, name or "", + list(fields.keys()), list(reader_names or ()), one_shot) return gpu_types.pipe_value(capacity, scope, name, fields, reader_names, one_shot=one_shot) diff --git a/test/CommonIR/native_flash_gpu.py b/test/CommonIR/native_flash_gpu.py index 0f5b20072e..e7894c650d 100644 --- a/test/CommonIR/native_flash_gpu.py +++ b/test/CommonIR/native_flash_gpu.py @@ -1,5 +1,4 @@ import argparse -import os from pathlib import Path import torch @@ -8,9 +7,6 @@ import triton.experimental.tle.language as tle -os.environ.setdefault("TLE_GPU_TILEIR_MODE", "1") - - @triton.jit def _flash_gpu_kernel(q_ptr, k_ptr, v_ptr, o_ptr, stride_qm, stride_qd, stride_km, stride_kd, stride_vm, stride_vd, stride_om, stride_od, N_CTX: tl.constexpr, D: tl.constexpr, sm_scale: tl.constexpr, @@ -150,10 +146,16 @@ def dump_ttir(path): pm = ir.pass_manager(module.context) passes.common.add_inliner(pm) pm.run(module, "native_flash_gpu.inliner") - tle_ir.lower_gpu_tileir_to_ttir(module) + pm = ir.pass_manager(module.context) + tle_ir.passes.add_convert_gpu_tile_to_ttgir(pm) + pm.run(module, "native_flash_gpu.tileir_to_ttgir") text = str(module) - if "tile." in text: - raise RuntimeError("TTIR dump still contains TileIR ops") + for needle in ("tile.alloc", "tile.copy", "tile.to_tensor", "tile.store_tensor"): + if needle in text: + raise RuntimeError(f"converted TTIR dump still contains {needle}") + for needle in ("ttg.local_alloc", "ttg.local_load", "ttg.local_store"): + if needle not in text: + raise RuntimeError(f"converted TTIR dump expected {needle}") Path(path).write_text(text, encoding="utf-8") print(f"[dump-ttir] wrote {path}") diff --git a/test/CommonIR/native_matmul_gpu.py b/test/CommonIR/native_matmul_gpu.py index ffdb380f00..b753c8ae98 100644 --- a/test/CommonIR/native_matmul_gpu.py +++ b/test/CommonIR/native_matmul_gpu.py @@ -1,5 +1,4 @@ import argparse -import os from pathlib import Path import torch @@ -8,9 +7,6 @@ import triton.experimental.tle.language as tle -os.environ.setdefault("TLE_GPU_TILEIR_MODE", "1") - - @triton.jit def _matmul_gpu_kernel(a_ptr, b_ptr, c_ptr, M, N, K, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr): @@ -130,10 +126,16 @@ def dump_ttir(path): pm = ir.pass_manager(module.context) passes.common.add_inliner(pm) pm.run(module, "native_matmul_gpu.inliner") - tle_ir.lower_gpu_tileir_to_ttir(module) + pm = ir.pass_manager(module.context) + tle_ir.passes.add_convert_gpu_tile_to_ttgir(pm) + pm.run(module, "native_matmul_gpu.tileir_to_ttgir") text = str(module) - if "tile." in text: - raise RuntimeError("TTIR dump still contains TileIR ops") + for needle in ("tile.alloc", "tile.copy", "tile.to_tensor", "tile.store_tensor"): + if needle in text: + raise RuntimeError(f"converted TTIR dump still contains {needle}") + for needle in ("ttg.local_alloc", "ttg.local_load", "ttg.local_store"): + if needle not in text: + raise RuntimeError(f"converted TTIR dump expected {needle}") Path(path).write_text(text, encoding="utf-8") print(f"[dump-ttir] wrote {path}") diff --git a/test/CommonIR/tle_gpu_primitives.py b/test/CommonIR/tle_gpu_primitives.py index c6ba0b42b5..9c880de8e0 100644 --- a/test/CommonIR/tle_gpu_primitives.py +++ b/test/CommonIR/tle_gpu_primitives.py @@ -1,5 +1,4 @@ import argparse -import os from pathlib import Path import triton @@ -7,9 +6,6 @@ import triton.experimental.tle.language as tle -os.environ.setdefault("TLE_GPU_TILEIR_MODE", "1") - - @triton.jit def _gpu_primitives_kernel(src, out, BLOCK: tl.constexpr): idx = tl.arange(0, BLOCK) @@ -24,7 +20,7 @@ def _gpu_primitives_kernel(src, out, BLOCK: tl.constexpr): writer.commit(0) wait_result = reader.wait(0) read_slot = wait_result.slot.payload - value = tle.gpu.local_ptr(write_slot) + value = tl.load(tle.gpu.local_ptr(read_slot)) reader.release(0) writer.close(0) tl.store(out + idx, value) @@ -69,15 +65,13 @@ def dump_tileir(path): "tile.alloc", "tile.copy", "tile.subview", - "tile.to_tensor", - "tile.set_flag", - "tile.wait_flag", - "tile.pipe_barrier", + "tile.local_ptr", + "tle.pipe.create", + "tle.pipe.writer_acquire", + "tle.pipe.reader_wait", ): if needle not in text: raise RuntimeError(f"expected {needle} in primitive TileIR dump") - if "tle.pipe" in text: - raise RuntimeError("primitive TileIR dump still contains raw tle.pipe ops") Path(path).write_text(text, encoding="utf-8") print(f"[dump-tileir] wrote {path}") @@ -91,12 +85,16 @@ def dump_ttir(path): pm = ir.pass_manager(module.context) passes.common.add_inliner(pm) pm.run(module, "tle_gpu_primitives.inliner") - tle_ir.lower_gpu_tileir_to_ttir(module) + pm = ir.pass_manager(module.context) + tle_ir.passes.add_convert_gpu_tile_to_ttgir(pm) + pm.run(module, "tle_gpu_primitives.tileir_to_ttgir") text = str(module) - if "tile." in text: - raise RuntimeError("primitive TTIR dump still contains TileIR ops") - if "ttg.local_barrier" not in text: - raise RuntimeError("primitive TTIR dump expected ttg.local_barrier") + for needle in ("tile.alloc", "tile.copy", "tile.subview", "tile.local_ptr"): + if needle in text: + raise RuntimeError(f"primitive converted TTIR still contains {needle}") + for needle in ("ttg.local_alloc", "ttg.memdesc_index", "tle.local_pointers", "tle.pipe.create"): + if needle not in text: + raise RuntimeError(f"primitive converted TTIR expected {needle}") Path(path).write_text(text, encoding="utf-8") print(f"[dump-ttir] wrote {path}") diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index 144fac81d1..f52ad1545d 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -264,9 +264,6 @@ def make_ttir(mod, metadata, opt, capability): passes.common.add_inliner(pm) pm.run(mod, 'make_ttir.inliner') - if hasattr(tle, "lower_gpu_tileir_to_ttir"): - tle.lower_gpu_tileir_to_ttir(mod) - pm = ir.pass_manager(mod.context) pm.enable_debug() passes.ttir.add_rewrite_tensor_pointer(pm) @@ -290,6 +287,10 @@ def make_ttgir(mod, metadata, opt, capability): pm = ir.pass_manager(mod.context) dump_enabled = pm.enable_debug() emuTF32 = (capability // 10 >= 8) + # TLE GPU always enters through TileIR when the dialect is available. + # Convert the target-independent buffers to TritonGPU descriptors + # before the standard TTIR-to-TTGIR conversion. + tle.passes.add_convert_gpu_tile_to_ttgir(pm) passes.ttir.add_convert_to_ttgpuir(pm, f"cuda:{capability}", opt.num_warps, 32, opt.num_ctas) # flagtree tle raw tle.raw_passes.add_tle_convert_arg_to_memdesc(pm) diff --git a/third_party/tle/dialect/include/IR/TleOps.td b/third_party/tle/dialect/include/IR/TleOps.td index ec89a213c3..147b3b9bfa 100644 --- a/third_party/tle/dialect/include/IR/TleOps.td +++ b/third_party/tle/dialect/include/IR/TleOps.td @@ -101,7 +101,7 @@ def Tle_MemDescWGMMAViewOp : Tle_Op<"memdesc_wgmma_view", [Pure, def Tle_PipeCreateOp : Tle_Op<"pipe.create"> { let summary = "create a typed TLE GPU pipe"; let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32Attr:$capacity, StrAttr:$scope, OptionalAttr:$pipe_name, @@ -115,7 +115,7 @@ def Tle_PipeCreateOp : Tle_Op<"pipe.create"> { def Tle_PipeWriterAcquireOp : Tle_Op<"pipe.writer_acquire"> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I1:$phase, I32Attr:$capacity, @@ -129,7 +129,7 @@ def Tle_PipeWriterAcquireOp : Tle_Op<"pipe.writer_acquire"> { def Tle_PipeWriterCommitOp : Tle_Op<"pipe.writer_commit"> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I32Attr:$capacity, StrAttr:$scope, @@ -142,7 +142,7 @@ def Tle_PipeWriterCommitOp : Tle_Op<"pipe.writer_commit"> { def Tle_PipeWriterCloseOp : Tle_Op<"pipe.writer_close"> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I1:$phase, I32Attr:$capacity, @@ -156,7 +156,7 @@ def Tle_PipeWriterCloseOp : Tle_Op<"pipe.writer_close"> { def Tle_PipeReaderWaitOp : Tle_Op<"pipe.reader_wait"> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I1:$phase, I32Attr:$capacity, @@ -174,7 +174,7 @@ def Tle_PipeReaderReleaseOp : Tle_Op<"pipe.reader_release", [ DeclareOpInterfaceMethods ]> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I32Attr:$capacity, StrAttr:$scope, diff --git a/third_party/tle/dialect/include/Transforms/Passes.td b/third_party/tle/dialect/include/Transforms/Passes.td index bef5701756..fb9916a959 100644 --- a/third_party/tle/dialect/include/Transforms/Passes.td +++ b/third_party/tle/dialect/include/Transforms/Passes.td @@ -4,6 +4,25 @@ include "mlir/Pass/PassBase.td" +def TritonTleConvertGpuTileToTtgir + : Pass<"triton-tle-convert-gpu-tile-to-ttgir", "mlir::ModuleOp"> { + let summary = "convert TLE GPU TileIR buffers to TritonGPU descriptors"; + + let description = [{ + Converts the target-independent TileIR produced by tle.gpu primitives into + the TritonGPU descriptor form expected by the NVIDIA lowering pipeline. + In particular, tile.alloc becomes ttg.local_alloc with a ttg.memdesc result, + and tile.local_ptr becomes tle.local_pointers. The conversion runs before + the standard TTIR-to-TTGIR pass and must eliminate all GPU buffer TileIR + operations. + }]; + + let dependentDialects = ["mlir::arith::ArithDialect", + "mlir::triton::TritonDialect", + "mlir::triton::gpu::TritonGPUDialect", + "mlir::triton::tle::TleDialect"]; +} + def TritonTleEarlyAssignMemorySpace : Pass<"triton-tle-early-assign-memory-space", "mlir::ModuleOp"> { let summary = "early assign memory space"; diff --git a/third_party/tle/dialect/lib/IR/CMakeLists.txt b/third_party/tle/dialect/lib/IR/CMakeLists.txt index 5e9e87a7cf..0ad0cd1546 100644 --- a/third_party/tle/dialect/lib/IR/CMakeLists.txt +++ b/third_party/tle/dialect/lib/IR/CMakeLists.txt @@ -13,3 +13,7 @@ add_triton_library(TleIR TritonIR TritonGPUIR ) + +if(TARGET TileIRIR) + target_link_libraries(TleIR PUBLIC TileIRIR) +endif() diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp index 3d26c14694..c011494aee 100644 --- a/third_party/tle/dialect/lib/IR/Ops.cpp +++ b/third_party/tle/dialect/lib/IR/Ops.cpp @@ -4,6 +4,9 @@ #include "triton/Dialect/Triton/IR/Types.h" #include "triton/Dialect/Triton/IR/Utility.h" #include "triton/Dialect/TritonGPU/IR/Dialect.h" +#ifdef __FLIR_TILEIR__ +#include "mlir-ext/Dialect/TileIR/IR/TileIRDialect.h" +#endif #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallSet.h" #include @@ -271,12 +274,24 @@ static LogicalResult verifyPipeAttrs(Operation *op, OperandRange fields) { if (fields.empty()) return op->emitOpError("expects at least one pipe field"); for (Value field : fields) { - auto type = cast(field.getType()); - if (!isa(type.getMemorySpace())) - return op->emitOpError("expects only shared-memory pipe fields"); - if (type.getRank() < 2) + ArrayRef shape; + if (auto type = dyn_cast(field.getType())) { + if (!isa(type.getMemorySpace())) + return op->emitOpError("expects only shared-memory pipe fields"); + shape = type.getShape(); +#ifdef __FLIR_TILEIR__ + } else if (auto type = dyn_cast(field.getType())) { + if (type.getMemorySpace() != triton::tile::MemorySpace::Shared) + return op->emitOpError("expects only shared-memory pipe fields"); + shape = type.getShape(); +#endif + } else { + return op->emitOpError( + "expects pipe fields to be TileIR buffers or ttg.memdesc values"); + } + if (shape.size() < 2) return op->emitOpError("expects pipe fields to have rank >= 2"); - if (type.getShape()[0] != capacity) + if (shape[0] != capacity) return op->emitOpError("expects field leading dimension to equal " "pipe capacity"); } diff --git a/third_party/tle/dialect/lib/Transforms/CMakeLists.txt b/third_party/tle/dialect/lib/Transforms/CMakeLists.txt index 7d2e6c7f7f..5647f87c22 100644 --- a/third_party/tle/dialect/lib/Transforms/CMakeLists.txt +++ b/third_party/tle/dialect/lib/Transforms/CMakeLists.txt @@ -1,5 +1,6 @@ # flagtree tle add_triton_library(TritonTLETransforms + TleConvertGpuTileToTtgir.cpp TleEarlyAssignMemorySpace.cpp TleSelectEncodings.cpp TleInsertLocalPointerBarriers.cpp @@ -37,6 +38,10 @@ add_triton_library(TritonTLETransforms NVWSIR ) +if(TARGET TileIRIR) + target_link_libraries(TritonTLETransforms PUBLIC TileIRIR) +endif() + target_include_directories(TritonTLETransforms PUBLIC ${CMAKE_SOURCE_DIR}/third_party/tle/dialect/include ${CMAKE_SOURCE_DIR}/include diff --git a/third_party/tle/dialect/lib/Transforms/TleConvertGpuTileToTtgir.cpp b/third_party/tle/dialect/lib/Transforms/TleConvertGpuTileToTtgir.cpp new file mode 100644 index 0000000000..ac5777f95e --- /dev/null +++ b/third_party/tle/dialect/lib/Transforms/TleConvertGpuTileToTtgir.cpp @@ -0,0 +1,411 @@ +// MIT License +// +// Copyright (c) 2025 The FlagOS Contributors + +#include "tle/dialect/include/Transforms/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinOps.h" +#include "tle/dialect/include/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" + +#ifdef __FLIR_TILEIR__ +#include "mlir-ext/Dialect/TileIR/IR/TileIRDialect.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" +#endif + +namespace mlir::triton::tle { + +#define GEN_PASS_DEF_TRITONTLECONVERTGPUTILETOTTGIR +#include "tle/dialect/include/Transforms/Passes.h.inc" + +namespace { + +#ifdef __FLIR_TILEIR__ +namespace tile = mlir::triton::tile; +namespace ttg = mlir::triton::gpu; + +constexpr int kSharedMemoryAddressSpace = 3; + +static FailureOr getMemDescType(Operation *op, + tile::BufType bufferType) { + if (bufferType.getMemorySpace() != tile::MemorySpace::Shared) { + op->emitError("GPU TileIR conversion currently requires #tile.shared"); + return failure(); + } + + Attribute layout = op->getAttr("tle.gpu_layout"); + if (!layout) { + op->emitError("is missing the preserved tle.gpu_layout attribute"); + return failure(); + } + + auto memorySpace = ttg::SharedMemorySpaceAttr::get(op->getContext()); + return ttg::MemDescType::get(bufferType.getShape(), + bufferType.getElementType(), layout, + memorySpace, /*mutableMemory=*/true); +} + +static FailureOr lookupBuffer(Operation *op, Value buffer, + const llvm::DenseMap &mapped) { + auto it = mapped.find(buffer); + if (it == mapped.end()) { + op->emitError("uses a TileIR buffer without a converted allocation"); + return failure(); + } + return it->second; +} + +static bool isPointerLike(Type type) { + if (isa(type)) + return true; + auto tensorType = dyn_cast(type); + return tensorType && isa(tensorType.getElementType()); +} + +static Type getLocalPointerType(Type valueType, Type elementType) { + Type ptrType = triton::PointerType::get(elementType, + kSharedMemoryAddressSpace); + if (auto tensorType = dyn_cast(valueType)) + return RankedTensorType::get(tensorType.getShape(), ptrType, + tensorType.getEncoding()); + return ptrType; +} + +static SmallVector createFullIndices(OpBuilder &builder, Location loc, + ArrayRef shape) { + SmallVector indices; + Type i32 = builder.getI32Type(); + for (auto [axis, dim] : llvm::enumerate(shape)) { + auto rangeType = RankedTensorType::get({dim}, i32); + Value index = builder.create(loc, rangeType, 0, dim); + SmallVector expandedShape{dim}; + + for (unsigned i = 0; i < axis; ++i) { + expandedShape.insert(expandedShape.begin(), 1); + auto expandedType = RankedTensorType::get(expandedShape, i32); + index = builder.create(loc, expandedType, index, 0); + } + for (unsigned i = axis + 1; i < shape.size(); ++i) { + unsigned expandAxis = expandedShape.size(); + expandedShape.push_back(1); + auto expandedType = RankedTensorType::get(expandedShape, i32); + index = builder.create(loc, expandedType, index, + expandAxis); + } + + auto resultType = RankedTensorType::get(shape, i32); + indices.push_back( + builder.create(loc, resultType, index)); + } + return indices; +} + +static Value createFullLocalPointers(OpBuilder &builder, Location loc, + Value memDesc, Type valueType) { + auto memDescType = cast(memDesc.getType()); + Type ptrType = + getLocalPointerType(valueType, memDescType.getElementType()); + SmallVector indices = + createFullIndices(builder, loc, memDescType.getShape()); + return builder.create(loc, ptrType, memDesc, indices); +} + +static Type getPointeeValueType(Type pointerType) { + if (auto ptr = dyn_cast(pointerType)) + return ptr.getPointeeType(); + auto tensorType = dyn_cast(pointerType); + if (!tensorType) + return {}; + auto ptr = dyn_cast(tensorType.getElementType()); + if (!ptr) + return {}; + return RankedTensorType::get(tensorType.getShape(), ptr.getPointeeType(), + tensorType.getEncoding()); +} + +static FailureOr castToI32(OpBuilder &builder, Location loc, + Value value) { + if (value.getType().isInteger(32)) + return value; + if (value.getType().isIndex()) + return builder.create(loc, builder.getI32Type(), value) + .getResult(); + emitError(loc) << "expected an index or i32 TileIR subview offset"; + return failure(); +} + +static FailureOr getConstantOffset(Value value) { + if (auto indexCast = value.getDefiningOp()) + value = indexCast.getIn(); + auto constant = value.getDefiningOp(); + if (!constant) + return failure(); + auto integer = dyn_cast(constant.getValue()); + if (!integer) + return failure(); + return static_cast(integer.getInt()); +} + +static LogicalResult convertAllocations( + ModuleOp module, llvm::DenseMap &mapped, + SmallVectorImpl &eraseOps) { + WalkResult result = module.walk([&](tile::AllocOp op) { + auto memDescType = + getMemDescType(op, cast(op.getResult().getType())); + if (failed(memDescType)) + return WalkResult::interrupt(); + OpBuilder builder(op); + mapped[op.getResult()] = + builder.create(op.getLoc(), *memDescType); + eraseOps.push_back(op); + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); +} + +static LogicalResult convertSubviews( + ModuleOp module, llvm::DenseMap &mapped, + SmallVectorImpl &eraseOps) { + WalkResult result = module.walk([&](tile::SubViewOp op) { + auto source = lookupBuffer(op, op.getSource(), mapped); + auto resultType = + getMemDescType(op, cast(op.getResult().getType())); + if (failed(source) || failed(resultType)) + return WalkResult::interrupt(); + + OpBuilder builder(op); + auto sourceType = cast((*source).getType()); + Value converted; + if (sourceType.getRank() == resultType->getRank() + 1) { + if (op.getOffsets().empty()) { + op.emitError("rank-reducing tile.subview requires a leading offset"); + return WalkResult::interrupt(); + } + auto index = castToI32(builder, op.getLoc(), op.getOffsets().front()); + if (failed(index)) + return WalkResult::interrupt(); + converted = builder.create( + op.getLoc(), *resultType, *source, *index); + } else if (sourceType.getRank() == resultType->getRank()) { + SmallVector offsets; + offsets.reserve(op.getOffsets().size()); + for (Value offset : op.getOffsets()) { + auto constant = getConstantOffset(offset); + if (failed(constant)) { + op.emitError("same-rank tile.subview currently requires static offsets"); + return WalkResult::interrupt(); + } + offsets.push_back(*constant); + } + converted = builder.create( + op.getLoc(), *resultType, *source, + builder.getDenseI32ArrayAttr(offsets)); + } else { + op.emitError("unsupported tile.subview rank change"); + return WalkResult::interrupt(); + } + + mapped[op.getResult()] = converted; + eraseOps.push_back(op); + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); +} + +static LogicalResult convertTileUsers( + ModuleOp module, llvm::DenseMap &mapped, + SmallVectorImpl &eraseOps) { + bool failedConversion = false; + OpBuilder builder(module.getContext()); + + module.walk([&](Operation *operation) { + if (failedConversion) + return; + + if (auto op = dyn_cast(operation)) { + auto source = lookupBuffer(op, op.getSource(), mapped); + if (failed(source)) { + failedConversion = true; + return; + } + builder.setInsertionPoint(op); + Value replacement = builder.create( + op.getLoc(), op.getResult().getType(), *source, op.getIndices()); + op.getResult().replaceAllUsesWith(replacement); + eraseOps.push_back(op); + return; + } + + if (auto op = dyn_cast(operation)) { + auto source = lookupBuffer(op, op.getSrc(), mapped); + if (failed(source)) { + failedConversion = true; + return; + } + builder.setInsertionPoint(op); + Value replacement = builder.create( + op.getLoc(), op.getResult().getType(), *source); + op.getResult().replaceAllUsesWith(replacement); + eraseOps.push_back(op); + return; + } + + if (auto op = dyn_cast(operation)) { + auto destination = lookupBuffer(op, op.getDst(), mapped); + if (failed(destination)) { + failedConversion = true; + return; + } + builder.setInsertionPoint(op); + builder.create(op.getLoc(), op.getSrc(), + *destination); + eraseOps.push_back(op); + return; + } + + if (auto op = dyn_cast(operation)) { + Value src = op.getSrc(); + Value dst = op.getDst(); + if (isa(src.getType())) { + auto converted = lookupBuffer(op, src, mapped); + if (failed(converted)) { + failedConversion = true; + return; + } + src = *converted; + } + if (isa(dst.getType())) { + auto converted = lookupBuffer(op, dst, mapped); + if (failed(converted)) { + failedConversion = true; + return; + } + dst = *converted; + } + + builder.setInsertionPoint(op); + if (!op.getIndices().empty()) { + SmallVector indices; + for (Value index : op.getIndices()) { + auto converted = castToI32(builder, op.getLoc(), index); + if (failed(converted)) { + failedConversion = true; + return; + } + indices.push_back(*converted); + } + builder.create(op.getLoc(), src, dst, indices); + } else if (isa(dst.getType())) { + Value value = src; + if (isPointerLike(src.getType())) + value = builder.create( + op.getLoc(), src, triton::CacheModifier::NONE, + triton::EvictionPolicy::NORMAL, false); + Value ptr = createFullLocalPointers(builder, op.getLoc(), dst, + value.getType()); + builder.create(op.getLoc(), ptr, value, + triton::CacheModifier::NONE, + triton::EvictionPolicy::NORMAL); + } else if (isa(src.getType()) && + isPointerLike(dst.getType())) { + Type valueType = getPointeeValueType(dst.getType()); + if (!valueType) { + op.emitError("cannot determine tile.copy destination value type"); + failedConversion = true; + return; + } + Value ptr = createFullLocalPointers(builder, op.getLoc(), src, + valueType); + Value value = builder.create( + op.getLoc(), ptr, triton::CacheModifier::NONE, + triton::EvictionPolicy::NORMAL, false); + builder.create(op.getLoc(), dst, value, + triton::CacheModifier::NONE, + triton::EvictionPolicy::NORMAL); + } else { + op.emitError("unsupported operands for GPU tile.copy conversion"); + failedConversion = true; + return; + } + eraseOps.push_back(op); + return; + } + + if (isa(operation)) { + builder.setInsertionPoint(operation); + builder.create(operation->getLoc()); + eraseOps.push_back(operation); + } + }); + + return failure(failedConversion); +} + +static LogicalResult rewritePipeOperands( + ModuleOp module, const llvm::DenseMap &mapped) { + bool failedConversion = false; + module.walk([&](Operation *op) { + if (!op->getName().getStringRef().starts_with("tle.pipe.")) + return; + for (OpOperand &operand : op->getOpOperands()) { + if (!isa(operand.get().getType())) + continue; + auto converted = lookupBuffer(op, operand.get(), mapped); + if (failed(converted)) { + failedConversion = true; + return; + } + operand.set(*converted); + } + }); + return failure(failedConversion); +} +#endif + +class ConvertGpuTileToTtgirPass + : public impl::TritonTleConvertGpuTileToTtgirBase< + ConvertGpuTileToTtgirPass> { + void runOnOperation() override { +#ifdef __FLIR_TILEIR__ + ModuleOp module = getOperation(); + llvm::DenseMap mapped; + SmallVector eraseOps; + + if (failed(convertAllocations(module, mapped, eraseOps)) || + failed(convertSubviews(module, mapped, eraseOps)) || + failed(convertTileUsers(module, mapped, eraseOps)) || + failed(rewritePipeOperands(module, mapped))) { + signalPassFailure(); + return; + } + + for (Operation *op : llvm::reverse(eraseOps)) { + if (llvm::any_of(op->getResults(), + [](Value result) { return !result.use_empty(); })) { + op->emitError("still has users after GPU TileIR conversion"); + signalPassFailure(); + return; + } + op->erase(); + } + + bool hasRemainingGpuTileOps = false; + module.walk([&](Operation *op) { + if (isa(op)) { + op->emitError("was not eliminated by GPU TileIR conversion"); + hasRemainingGpuTileOps = true; + } + }); + if (hasRemainingGpuTileOps) + signalPassFailure(); +#endif + } +}; + +} // namespace + +} // namespace mlir::triton::tle diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index 87764bde1a..b41c0b974f 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -51,7 +51,6 @@ #include "mlir-ext/Dialect/TileIR/IR/TileIRDialect.h" #include "triton/Dialect/Triton/IR/Dialect.h" #endif -#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallVectorExtras.h" #include "llvm/ADT/STLExtras.h" #include "llvm/IRReader/IRReader.h" @@ -109,140 +108,6 @@ static tile::MemorySpace attrToTileMemSpace(Attribute attr) { return tile::MemorySpace::Shared; } -static bool isTileBuffer(Value value) { - return value && isa(value.getType()); -} - -static bool isTTTensorOrPointer(Value value) { - if (!value) - return false; - auto type = value.getType(); - return isa(type); -} - -static bool isTTPointerLike(Value value) { - if (!value) - return false; - Type type = value.getType(); - if (isa(type)) - return true; - if (auto tensorType = dyn_cast(type)) - return isa(tensorType.getElementType()); - return false; -} - -static Value materializeTileBuffer(Value value, - llvm::DenseMap &bufferValues) { - while (isTileBuffer(value)) { - auto it = bufferValues.find(value); - if (it == bufferValues.end()) - return Value(); - value = it->second; - } - return value; -} - -static void eraseIfUnused(Operation *op) { - if (!op) - return; - if (llvm::all_of(op->getResults(), - [](Value result) { return result.use_empty(); })) - op->erase(); -} - -static void lowerGpuTileIRToTTIR(ModuleOp module) { - OpBuilder builder(module.getContext()); - llvm::DenseMap bufferValues; - SmallVector eraseOps; - - module.walk([&](Operation *op) { - if (auto copyOp = dyn_cast(op)) { - Value src = materializeTileBuffer(copyOp.getSrc(), bufferValues); - if (!src) - src = copyOp.getSrc(); - Value dst = copyOp.getDst(); - if (isTileBuffer(dst)) { - bufferValues[dst] = src; - eraseOps.push_back(op); - return; - } - if (isTTPointerLike(dst) && isTTTensorOrPointer(src)) { - builder.setInsertionPoint(op); - Value storeValue = src; - if (isTTPointerLike(storeValue)) { - auto load = builder.create( - op->getLoc(), storeValue, triton::CacheModifier::NONE, - triton::EvictionPolicy::NORMAL, false); - storeValue = load.getResult(); - } - builder.create( - op->getLoc(), dst, storeValue, triton::CacheModifier::NONE, - triton::EvictionPolicy::NORMAL); - eraseOps.push_back(op); - } - return; - } - - if (auto storeTensorOp = dyn_cast(op)) { - bufferValues[storeTensorOp.getDst()] = storeTensorOp.getSrc(); - eraseOps.push_back(op); - return; - } - - if (auto subviewOp = dyn_cast(op)) { - if (Value src = materializeTileBuffer(subviewOp.getSource(), bufferValues)) - bufferValues[subviewOp.getResult()] = src; - eraseOps.push_back(op); - return; - } - - if (auto gmOffsetOp = dyn_cast(op)) { - gmOffsetOp.getResult().replaceAllUsesWith(gmOffsetOp.getBase()); - eraseOps.push_back(op); - return; - } - - if (auto toTensorOp = dyn_cast(op)) { - Value value = materializeTileBuffer(toTensorOp.getSrc(), bufferValues); - if (!value) { - toTensorOp.emitError("cannot lower tile.to_tensor without a preceding " - "tile.copy or tile.store_tensor"); - return; - } - - builder.setInsertionPoint(op); - Value replacement = value; - if (isTTPointerLike(value)) { - auto load = builder.create( - op->getLoc(), value, triton::CacheModifier::NONE, - triton::EvictionPolicy::NORMAL, false); - replacement = load.getResult(); - } - toTensorOp.getResult().replaceAllUsesWith(replacement); - eraseOps.push_back(op); - } - - if (isa(op)) { - builder.setInsertionPoint(op); - builder.create(op->getLoc()); - eraseOps.push_back(op); - return; - } - }); - - for (Operation *op : llvm::reverse(eraseOps)) - eraseIfUnused(op); - - SmallVector cleanupOps; - module.walk([&](Operation *op) { - if (isa(op)) - cleanupOps.push_back(op); - }); - for (Operation *op : llvm::reverse(cleanupOps)) - eraseIfUnused(op); -} #endif void init_triton_tle_ir(py::module &&m) { @@ -375,9 +240,10 @@ void init_triton_tle_ir(py::module &&m) { return tile::BufType::get(ctx, shape, elementType, memSpace); }) .def("create_tile_alloc", - [](TritonOpBuilder &self, Type tileBufType) -> Value { + [](TritonOpBuilder &self, Type tileBufType, + Attribute targetLayout) -> Value { auto bufType = mlir::cast(tileBufType); - return self.create( + auto op = self.create( tileBufType, bufType.getMemorySpace(), /*shape=*/mlir::ArrayAttr(), /*dtype=*/mlir::TypeAttr(), /*policy=*/tile::PolicyAttr(), @@ -385,12 +251,14 @@ void init_triton_tle_ir(py::module &&m) { tile::Layout::ND), /*lifetime=*/tile::LifetimeAttr(), /*comment=*/mlir::StringAttr()); + op->setAttr("tle.gpu_layout", targetLayout); + return op.getResult(); }) .def("create_tile_copy", [](TritonOpBuilder &self, Value &src, Value &dst, - std::vector & /*shape*/, bool interNoAlias) -> void { + std::vector &indices, bool interNoAlias) -> void { auto op = self.create( - src, dst, /*engine=*/tile::EngineAttr(), + src, dst, indices, /*engine=*/tile::EngineAttr(), /*src_layout=*/tile::LayoutAttr::get( self.getBuilder().getContext(), tile::Layout::ND), /*dst_nz_layout=*/tile::NZLayoutAttr(), @@ -402,7 +270,8 @@ void init_triton_tle_ir(py::module &&m) { .def("create_tile_subview", [](TritonOpBuilder &self, Value source, std::vector &offsets, const std::vector &sizes, - const std::vector &strides) -> Value { + const std::vector &strides, + Attribute targetLayout) -> Value { SmallVector indexOffsets; auto &builder = self.getBuilder(); auto indexType = builder.getIndexType(); @@ -418,8 +287,18 @@ void init_triton_tle_ir(py::module &&m) { auto op = self.create( resTy, source, indexOffsets, builder.getI64ArrayAttr(sizes), builder.getI64ArrayAttr(strides)); + op->setAttr("tle.gpu_layout", targetLayout); return op.getResult(); }) + .def("create_tile_local_ptr", + [](TritonOpBuilder &self, Type resultTy, Value source, + py::args args) -> OpState { + llvm::SmallVector indices; + indices.reserve(args.size()); + for (const auto &arg : args) + indices.push_back(py::cast(arg)); + return self.create(resultTy, source, indices); + }) .def("create_tile_to_tensor", [](TritonOpBuilder &self, Value &src, bool /*writable*/) -> Value { auto srcBuf = mlir::cast(src.getType()); @@ -733,6 +612,8 @@ void init_triton_tle_ir(py::module &&m) { } void init_triton_tle_passes(py::module &&m) { + ADD_PASS_WRAPPER_0("add_convert_gpu_tile_to_ttgir", + tle::createTritonTleConvertGpuTileToTtgir); ADD_PASS_WRAPPER_0("add_early_assign_memory_space", tle::createTritonTleEarlyAssignMemorySpace); ADD_PASS_WRAPPER_0("add_select_encodings", @@ -847,8 +728,6 @@ void init_triton_tle(py::module &&m) { context.appendDialectRegistry(registry); context.loadAllAvailableDialects(); }); - m.def("lower_gpu_tileir_to_ttir", - [](ModuleOp &module) { lowerGpuTileIRToTTIR(module); }); #endif init_triton_tle_ir(m.def_submodule("ir")); From b6fe3db17a1b2ab9d7921c3b1085a5583bc00ef8 Mon Sep 17 00:00:00 2001 From: macqueen09 Date: Thu, 20 Aug 2026 23:14:54 +0800 Subject: [PATCH 2/2] Refine GPU CommonIR to TTGIR conversion --- CMakeLists.txt | 19 +- bin/RegisterTritonDialects.h | 6 + python/triton/_common_ir.py | 4 + .../experimental/tle/language/gpu/core.py | 247 +++--------------- .../experimental/tle/language/gpu/semantic.py | 143 +++++++++- .../experimental/tle/language/gpu/types.py | 48 ++-- .../triton/experimental/tle/language/pipe.py | 5 +- test/CMakeLists.txt | 1 + test/CommonIR/gpu-tileir-to-ttgir.mlir | 33 +++ test/CommonIR/native_flash_gpu.py | 4 +- test/CommonIR/native_matmul_gpu.py | 4 +- test/CommonIR/tle_gpu_primitives.py | 7 +- test/lit.cfg.py | 3 + test/lit.site.cfg.py.in | 1 + third_party/nvidia/CMakeLists.txt | 3 + third_party/nvidia/backend/compiler.py | 10 +- third_party/nvidia/include/CMakeLists.txt | 3 + .../include/CommonIRToTTGIR/CMakeLists.txt | 3 + .../nvidia/include/CommonIRToTTGIR/Passes.h | 16 ++ .../nvidia/include/CommonIRToTTGIR/Passes.td | 22 ++ third_party/nvidia/lib/CMakeLists.txt | 3 + .../nvidia/lib/CommonIRToTTGIR/CMakeLists.txt | 13 + .../lib/CommonIRToTTGIR/CommonIRToTTGIR.cpp} | 214 +++++++++++---- third_party/nvidia/triton_nvidia.cc | 12 + third_party/tle/dialect/include/IR/TleOps.td | 12 +- .../tle/dialect/include/Transforms/Passes.td | 19 -- third_party/tle/dialect/lib/IR/CMakeLists.txt | 4 - third_party/tle/dialect/lib/IR/Ops.cpp | 25 +- .../tle/dialect/lib/Transforms/CMakeLists.txt | 5 - third_party/tle/triton_tle.cc | 19 +- 30 files changed, 555 insertions(+), 353 deletions(-) create mode 100644 python/triton/_common_ir.py create mode 100644 test/CommonIR/gpu-tileir-to-ttgir.mlir create mode 100644 third_party/nvidia/include/CommonIRToTTGIR/CMakeLists.txt create mode 100644 third_party/nvidia/include/CommonIRToTTGIR/Passes.h create mode 100644 third_party/nvidia/include/CommonIRToTTGIR/Passes.td create mode 100644 third_party/nvidia/lib/CommonIRToTTGIR/CMakeLists.txt rename third_party/{tle/dialect/lib/Transforms/TleConvertGpuTileToTtgir.cpp => nvidia/lib/CommonIRToTTGIR/CommonIRToTTGIR.cpp} (69%) diff --git a/CMakeLists.txt b/CMakeLists.txt index e3014d3ed2..473960c790 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -16,6 +16,13 @@ if(FLAGTREE_BACKEND) add_definitions(-DFLAGTREE_BACKEND=\"${FLAGTREE_BACKEND}\") endif() +# CommonIR is an explicit build-time feature. Keep the public contract as the +# C/C++ FLAGTREE_COMMON_IR macro rather than adding a CMake cache option. +set(FLAGTREE_COMMON_IR_ENABLED "$ENV{FLAGTREE_COMMON_IR}") +if(FLAGTREE_COMMON_IR_ENABLED) + add_compile_definitions(FLAGTREE_COMMON_IR) +endif() + set(FLAGTREE_TLE ON) if(FLAGTREE_TLE) add_definitions(-D__TLE__) @@ -258,10 +265,14 @@ endif() include_directories(${PROJECT_SOURCE_DIR}/third_party) include_directories(${PROJECT_BINARY_DIR}/third_party) # Tablegen'd files -# TileIR dialect from FLIR is used by the TLE GPU -> Common IR POC path. -# Keep it optional at configure time so non-POC builds without FLIR are not -# forced to download the dependency. -if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/flir/CMakeLists.txt" AND NOT TARGET TileIRIR) +# TileIR is only part of builds that explicitly enable CommonIR. +if(FLAGTREE_COMMON_IR_ENABLED AND NOT TARGET TileIRIR) + if(FLAGTREE_BACKEND) + message(FATAL_ERROR "FLAGTREE_COMMON_IR on triton_v3.6.x requires the default NVIDIA backend") + endif() + if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/flir/CMakeLists.txt") + message(FATAL_ERROR "FLAGTREE_COMMON_IR requires FLIR at third_party/flir") + endif() add_definitions(-D__FLIR_TILEIR__) include_directories(${CMAKE_CURRENT_SOURCE_DIR}/third_party/flir/include) include_directories(${CMAKE_CURRENT_BINARY_DIR}/third_party/flir/include) diff --git a/bin/RegisterTritonDialects.h b/bin/RegisterTritonDialects.h index 9ea88fe388..e247e28761 100644 --- a/bin/RegisterTritonDialects.h +++ b/bin/RegisterTritonDialects.h @@ -17,6 +17,9 @@ #ifdef __FLIR_TILEIR__ #include "mlir-ext/Dialect/TileIR/IR/TileIRDialect.h" #endif +#ifdef FLAGTREE_COMMON_IR +#include "nvidia/include/CommonIRToTTGIR/Passes.h" +#endif #include "triton/Dialect/Gluon/Transforms/Passes.h" #include "triton/Dialect/Triton/IR/Dialect.h" #include "triton/Dialect/TritonGPU/IR/Dialect.h" @@ -77,6 +80,9 @@ inline void registerTritonDialects(mlir::DialectRegistry ®istry) { mlir::triton::gluon::registerGluonPasses(); #ifdef __TLE__ mlir::triton::tle::registerPasses(); // flagtree tle +#endif +#ifdef FLAGTREE_COMMON_IR + mlir::triton::registerCommonIRToTTGIRPasses(); #endif mlir::test::registerTestAliasPass(); mlir::test::registerTestAlignmentPass(); diff --git a/python/triton/_common_ir.py b/python/triton/_common_ir.py new file mode 100644 index 0000000000..542ffe49ce --- /dev/null +++ b/python/triton/_common_ir.py @@ -0,0 +1,4 @@ +from triton._C.libtriton import tle as tle_ir + + +ENABLED = tle_ir.is_common_ir_enabled() diff --git a/python/triton/experimental/tle/language/gpu/core.py b/python/triton/experimental/tle/language/gpu/core.py index f5f22ecfed..52d66a3d2d 100644 --- a/python/triton/experimental/tle/language/gpu/core.py +++ b/python/triton/experimental/tle/language/gpu/core.py @@ -1,10 +1,10 @@ # flagtree tle import builtins import triton.language.core as tl -from typing import Optional, Sequence from enum import Enum +from typing import Optional, Sequence from . import types as tle -from .mthreads import copy as mthreads_copy +from . import semantic as tle_semantic from triton.compiler.code_generator import flatten_values_to_ir, unflatten_ir_values from triton.language.core import ( @@ -17,22 +17,6 @@ SHARED_MEMORY_ADDRESS_SPACE = 3 -def _tileir_available(builder) -> bool: - return hasattr(builder, "create_tile_alloc") - - -def _tile_storage_name(storage: tle.scope) -> str: - if storage == tle.smem: - return "shared" - if storage == tle.tmem: - return "local" - raise ValueError(f"Storage type {storage} not supported by TileIR") - - -def _tile_shape(shape) -> list[int]: - return [int(tl._unwrap_if_constexpr(dim)) for dim in shape] - - class pipeline(range): """ Iterator that counts upward forever, with parallel execution semantics. @@ -263,9 +247,7 @@ def alloc( try: unwrapped_shape = [tl._unwrap_if_constexpr(dim) for dim in shape] - full_shape = unwrapped_shape dtype = tl._unwrap_if_constexpr(dtype) - elem_type = dtype.to_ir(_semantic.builder) if layout is None: if storage == tle.smem: @@ -306,34 +288,25 @@ def alloc( # Use provided layout layout_handle = layout.to_ir(_semantic.builder) - if _tileir_available(_semantic.builder): - memory_space = _semantic.builder.tile_get_string_attr(_tile_storage_name(storage)) - tile_ty = _semantic.builder.tile_get_buffer_type(_tile_shape(unwrapped_shape), elem_type, memory_space) - tensor_handle = _semantic.builder.create_tile_alloc(tile_ty, layout_handle) - result = tle.buffered_tensor(tensor_handle, dtype, unwrapped_shape, storage, layout, _semantic) - if init_value is not None: - _semantic.builder.create_tile_store_tensor(init_value.handle, tensor_handle) - return result - - if storage == tle.smem: - if init_value is not None: - mutable_ty = _semantic.builder.get_memdesc_type(full_shape, elem_type, layout_handle, "smem") - tensor_handle = _semantic.builder.create_local_alloc(mutable_ty, init_value.handle) - else: - tensor_handle = _semantic.builder.create_local_alloc(full_shape, elem_type, layout_handle) - else: - raise ValueError(f"Storage type {storage} not yet supported") - - return tle.buffered_tensor(tensor_handle, dtype, unwrapped_shape, storage, layout, _semantic) + return tle_semantic.alloc( + unwrapped_shape, + dtype, + storage, + layout, + layout_handle, + init_value, + _semantic, + ) except Exception as e: raise RuntimeError(f"Memory allocation failed: {str(e)}") from e class CopyDirection(Enum): - """Copy direction enum for data transfer operations""" - GM_TO_LOCAL = "GMTOLOCAL" # Global memory to local memory - LOCAL_TO_GM = "LOCALTOGM" # Local memory to global memory + """Copy direction for global/local transfers.""" + + GM_TO_LOCAL = "GMTOLOCAL" + LOCAL_TO_GM = "LOCALTOGM" @tl.builtin @@ -386,104 +359,6 @@ def copy( TMA copy with offsets: tle.copy(tma_desc, local_buf, [64, 64], [x_offset, y_offset]) """ - mthreads_enabled = mthreads_copy.enabled() - - def normcopy( - src: tl.tensor, - dst: tle.buffered_tensor, - shape: tuple, - direction, - _semantic=None, - ) -> None: - if mthreads_enabled: - mthreads_copy.validate_normal_copy(src, dst, shape, direction) - - # Semantic analysis - try: - from .semantic import TLESemantic - if isinstance(_semantic, TLESemantic): - _semantic.analyze_copy_operation(src, dst, shape) - except ImportError: - # If semantic analysis module is not available, continue with warning - import warnings - warnings.warn("TLE semantic analysis module not available, skipping validation", UserWarning) - - mask = None - other = None - boundary_check = () - padding_option = "" - cache_modifier = "" - eviction_policy = "" - volatile = False - - 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 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: - local_ptrs = local_ptr(src, _make_full_indices(src, _semantic), _semantic=_semantic) - load = tl.load(local_ptrs, _semantic=_semantic) - _semantic.store(dst, load, mask, boundary_check, cache_modifier, eviction_policy) - except Exception as e: - raise RuntimeError(f"copy operation failed: {str(e)}") from e - - # this api is use for tma copy - def tmacopy( - src: tle.buffered_tensor | tl.tensor_descriptor, - dst: tle.buffered_tensor | tl.tensor_descriptor, - direction, - shape: tuple, - offsets: Sequence[constexpr | tensor], - _semantic=None, - ) -> None: - # Parameter validation - valid_types = (tle.buffered_tensor, tl.tensor_descriptor) - - if not isinstance(src, valid_types): - raise ValueError( - f"Source parameter must be tle.buffered_tensor or tl.tensor_descriptor, but got {type(src).__name__}") - - if not isinstance(dst, valid_types): - raise ValueError( - f"Destination parameter must be tle.buffered_tensor or tl.tensor_descriptor, but got {type(dst).__name__}" - ) - - # Auto-determine copy direction based on operand types - if isinstance(src, tle.buffered_tensor) and isinstance(dst, tl.tensor_descriptor): - desc = dst - elif isinstance(src, tl.tensor_descriptor) and isinstance(dst, tle.buffered_tensor): - desc = src - else: - raise ValueError( - f"Invalid copy combination: src={type(src).__name__}, dst={type(dst).__name__}. " - "One operand must be tl.tensor_descriptor (global memory) and the other must be tle.buffered_tensor (local memory)" - ) - - if not isinstance(shape, (tuple, list)): - # Try to handle Triton tuple-like objects - if hasattr(shape, '__iter__'): - shape = tuple(shape) - else: - raise ValueError(f"Shape parameter must be tuple or list, but got {type(shape)}") - - if not isinstance(offsets, (tuple, list)): - # Try to handle Triton tuple-like objects - if hasattr(offsets, '__iter__'): - offsets = tuple(offsets) - else: - raise ValueError(f"Shape parameter must be tuple or list, but got {type(shape)}") - - # Note: Skip shape assertion at this level since it requires _semantic context - # assert desc.shape == shape, "Shape mismatch between descriptor and provided shape" - assert len(offsets) == len(desc.shape), "Offsets and shape must have the same length" - offsets = _semantic._convert_to_ir_values(offsets, require_i64=False) - _semantic.builder.create_tma_copy(src.handle, dst.handle, offsets) - return - # Parameter validation valid_types = (tl.tensor, tle.buffered_tensor, tl.tensor_descriptor) @@ -497,23 +372,15 @@ def tmacopy( f"Destination parameter must be tl.tensor or tle.buffered_tensor tl.tensor_descriptor, but got {type(dst).__name__}" ) - # Auto-determine copy direction based on operand types - if isinstance(src, tle.buffered_tensor) and isinstance(dst, tl.tensor): - direction = CopyDirection.LOCAL_TO_GM - is_normcopy = True - elif isinstance(src, tl.tensor) and isinstance(dst, tle.buffered_tensor): - direction = CopyDirection.GM_TO_LOCAL - is_normcopy = True - elif isinstance(src, tle.buffered_tensor) and isinstance(dst, tl.tensor_descriptor): - direction = CopyDirection.LOCAL_TO_GM - is_normcopy = False - elif isinstance(src, tl.tensor_descriptor) and isinstance(dst, tle.buffered_tensor): - direction = CopyDirection.GM_TO_LOCAL - is_normcopy = False - else: + local_to_global = isinstance(src, tle.buffered_tensor) and isinstance( + dst, (tl.tensor, tl.tensor_descriptor)) + global_to_local = isinstance(dst, tle.buffered_tensor) and isinstance( + src, (tl.tensor, tl.tensor_descriptor)) + if not (local_to_global or global_to_local): raise ValueError( f"Invalid copy combination: src={type(src).__name__}, dst={type(dst).__name__}. " - "One operand must be tl.tensor (global memory) and the other must be tle.buffered_tensor (local memory)") + "One operand must be a global tensor/descriptor and the other must be tle.buffered_tensor") + direction = CopyDirection.LOCAL_TO_GM if local_to_global else CopyDirection.GM_TO_LOCAL if not isinstance(shape, (tuple, list)): # Try to handle Triton tuple-like objects @@ -521,20 +388,21 @@ def tmacopy( shape = tuple(shape) else: raise ValueError(f"Shape parameter must be tuple or list, but got {type(shape)}") - if hasattr(_semantic.builder, "create_tile_copy"): - tile_indices = [] - if not is_normcopy: - if offsets is None: - raise ValueError("TMA copy requires offsets") - tile_indices = _semantic._convert_to_ir_values(offsets, require_i64=False) - _semantic.builder.create_tile_copy(src.handle, dst.handle, tile_indices, False) - return - if is_normcopy: - return normcopy(src, dst, shape, direction, _semantic) - if mthreads_enabled: - return mthreads_copy.tmacopy(src, dst, direction, shape, offsets, _semantic) - else: - return tmacopy(src, dst, direction, shape, offsets, _semantic) + + descriptor = src if isinstance(src, tl.tensor_descriptor) else dst if isinstance( + dst, tl.tensor_descriptor) else None + if descriptor is not None: + if offsets is None: + raise ValueError("TMA copy requires offsets") + if not isinstance(offsets, (tuple, list)): + if hasattr(offsets, '__iter__'): + offsets = tuple(offsets) + else: + raise ValueError(f"Offsets must be tuple or list, but got {type(offsets)}") + if len(offsets) != len(descriptor.shape): + raise ValueError("Offsets and descriptor shape must have the same rank") + + tle_semantic.copy(src, dst, shape, offsets if descriptor is not None else None, direction, _semantic) @tl.builtin @@ -550,45 +418,14 @@ def to_tensor( target_shape = tl._unwrap_if_constexpr(target_shape) shape = list(memref.shape if target_shape is None else target_shape) writable = tl._unwrap_if_constexpr(writable) - - if hasattr(_semantic.builder, "create_tile_to_tensor"): - result = _semantic.builder.create_tile_to_tensor(memref.handle, bool(writable)) - return tl.tensor(result, tl.block_type(memref.dtype, shape)) - - ptrs = local_ptr(memref, _make_full_indices(memref, _semantic), _semantic=_semantic) - return tl.load(ptrs, _semantic=_semantic) + return tle_semantic.to_tensor(memref, bool(writable), tl.block_type(memref.dtype, shape), _semantic) @tl.builtin def store_tensor(tensor_value: tl.tensor, dst: tle.buffered_tensor, _semantic=None) -> None: if not isinstance(dst, tle.buffered_tensor): raise ValueError(f"dst must be tle.gpu.buffered_tensor, got {type(dst).__name__}") - - if hasattr(_semantic.builder, "create_tile_store_tensor"): - _semantic.builder.create_tile_store_tensor(tensor_value.handle, dst.handle) - return - - ptrs = local_ptr(dst, _make_full_indices(dst, _semantic), _semantic=_semantic) - _semantic.store(ptrs, tensor_value, None, (), "", "") - - -def _expand_index_to_shape(index: tl.tensor, shape: Sequence[int], axis: int, _semantic) -> tl.tensor: - idx = index - for _ in builtins.range(axis): - idx = tl.expand_dims(idx, 0, _semantic=_semantic) - for _ in builtins.range(len(shape) - axis - 1): - idx = tl.expand_dims(idx, len(idx.shape), _semantic=_semantic) - return tl.broadcast_to(idx, *shape, _semantic=_semantic) - - -def _make_full_indices(buffer: tle.buffered_tensor, _semantic) -> tuple[tl.tensor, ...]: - shape = tuple(int(tl._unwrap_if_constexpr(dim)) for dim in buffer.type.shape) - indices = [] - for axis, dim in enumerate(shape): - idx = tl.arange(0, dim, _semantic=_semantic) - idx = _expand_index_to_shape(idx, shape, axis, _semantic) - indices.append(idx) - return tuple(indices) + tle_semantic.store_tensor(tensor_value, dst, _semantic) @tl.builtin @@ -696,11 +533,7 @@ def local_ptr( else: result_ty = tl.block_type(ptr_dtype, list(view_shape)) result_ir = result_ty.to_ir(_semantic.builder) - handles = [idx.handle for idx in idx_tensors] - if hasattr(_semantic.builder, "create_tile_local_ptr"): - local_ptr_op = _semantic.builder.create_tile_local_ptr(result_ir, buffer.handle, *handles) - else: - local_ptr_op = _semantic.builder.create_local_pointers(result_ir, buffer.handle, *handles) + local_ptr_op = tle_semantic.local_ptr(result_ir, buffer, idx_tensors, _semantic) result_tensor = tl.tensor(local_ptr_op.get_result(0), result_ty) diff --git a/python/triton/experimental/tle/language/gpu/semantic.py b/python/triton/experimental/tle/language/gpu/semantic.py index 06c531bd99..3581d5bb89 100644 --- a/python/triton/experimental/tle/language/gpu/semantic.py +++ b/python/triton/experimental/tle/language/gpu/semantic.py @@ -10,11 +10,152 @@ from typing import List, Optional, Sequence, Tuple, Union from triton._C.libtriton import ir -from triton import language as tl +import triton.language.core as tl +from triton._common_ir import ENABLED as COMMON_IR_ENABLED from . import types as tle +from .mthreads import copy as mthreads_copy from math import prod +def _expand_index_to_shape(index, shape, axis, _semantic): + result = index + for _ in range(axis): + result = tl.expand_dims(result, 0, _semantic=_semantic) + for _ in range(len(shape) - axis - 1): + result = tl.expand_dims(result, len(result.shape), _semantic=_semantic) + return tl.broadcast_to(result, *shape, _semantic=_semantic) + + +def _make_full_indices(buffer, _semantic): + shape = tuple(int(tl._unwrap_if_constexpr(dim)) for dim in buffer.type.shape) + return tuple( + _expand_index_to_shape(tl.arange(0, dim, _semantic=_semantic), shape, axis, _semantic) + for axis, dim in enumerate(shape) + ) + + +def alloc(shape, dtype, storage, layout, layout_handle, init_value, _semantic): + builder = _semantic.builder + shape = [int(tl._unwrap_if_constexpr(dim)) for dim in shape] + element_type = dtype.to_ir(builder) + if COMMON_IR_ENABLED: + memory_space = builder.tile_get_string_attr(tle._storage_to_tile_space(storage)) + tile_ty = builder.tile_get_buffer_type(shape, element_type, memory_space) + handle = builder.create_tile_alloc(tile_ty, layout_handle) + if init_value is not None: + builder.create_tile_store_tensor(init_value.handle, handle) + elif storage == tle.smem: + if init_value is None: + handle = builder.create_local_alloc(shape, element_type, layout_handle) + else: + result_type = builder.get_memdesc_type(shape, element_type, layout_handle, "smem") + handle = builder.create_local_alloc(result_type, init_value.handle) + else: + raise ValueError(f"Storage type {storage} not yet supported") + + result = tle.buffered_tensor(handle, dtype, list(shape), storage, layout, _semantic) + return result + + +def copy(src, dst, shape, offsets, direction, _semantic): + descriptor = src if isinstance(src, tl.tensor_descriptor) else dst if isinstance( + dst, tl.tensor_descriptor) else None + + if COMMON_IR_ENABLED: + if descriptor is None: + _semantic.builder.create_tile_copy(src.handle, dst.handle, False) + return + indices = _semantic._convert_to_ir_values(offsets, require_i64=False) + buffer = dst if isinstance(dst, tle.buffered_tensor) else src + memdesc = get_memdesc(buffer, _semantic) + src_handle = src.handle if src is descriptor else memdesc + dst_handle = dst.handle if dst is descriptor else memdesc + _semantic.builder.create_tma_copy(src_handle, dst_handle, indices) + return + + if descriptor is not None: + if mthreads_copy.enabled(): + mthreads_copy.tmacopy(src, dst, direction, shape, offsets, _semantic) + return + indices = _semantic._convert_to_ir_values(offsets, require_i64=False) + _semantic.builder.create_tma_copy(src.handle, dst.handle, indices) + return + + if mthreads_copy.enabled(): + mthreads_copy.validate_normal_copy(src, dst, shape, direction) + if isinstance(_semantic, TLESemantic): + _semantic.analyze_copy_operation(src, dst, shape) + + mask = None + boundary_check = () + cache_modifier = "" + eviction_policy = "" + if direction.name == "GM_TO_LOCAL": + load_extra_args = () if mthreads_copy.enabled() else (None, ) + value = _semantic.load(src, mask, None, boundary_check, "", cache_modifier, + eviction_policy, False, *load_extra_args) + from .core import local_ptr + ptrs = local_ptr(dst, _make_full_indices(dst, _semantic), _semantic=_semantic) + _semantic.store(ptrs, value, mask, boundary_check, cache_modifier, eviction_policy) + else: + from .core import local_ptr + ptrs = local_ptr(src, _make_full_indices(src, _semantic), _semantic=_semantic) + value = tl.load(ptrs, _semantic=_semantic) + _semantic.store(dst, value, mask, boundary_check, cache_modifier, eviction_policy) + + +def subview(src, offsets, shape, strides, layout, _semantic): + builder = _semantic.builder + if not COMMON_IR_ENABLED: + result_type = tle.buffered_tensor_type( + src.dtype, shape, src.type.storage, layout, _semantic, alloc_shape=shape).to_ir(builder) + return builder.create_memdesc_index(result_type, src.handle, offsets[0].handle) + return builder.create_tile_subview( + src.handle, + [offset.handle for offset in offsets], + shape, + strides, + layout.to_ir(builder), + ) + + +def to_tensor(buffer, writable, result_type, _semantic): + if COMMON_IR_ENABLED: + result = _semantic.builder.create_tile_to_tensor(buffer.handle, writable) + return tl.tensor(result, result_type) + from .core import local_ptr + ptrs = local_ptr(buffer, _make_full_indices(buffer, _semantic), _semantic=_semantic) + return tl.load(ptrs, _semantic=_semantic) + + +def store_tensor(value, dst, _semantic): + if COMMON_IR_ENABLED: + _semantic.builder.create_tile_store_tensor(value.handle, dst.handle) + return + from .core import local_ptr + ptrs = local_ptr(dst, _make_full_indices(dst, _semantic), _semantic=_semantic) + _semantic.store(ptrs, value, None, (), "", "") + + +def local_ptr(result_ir, buffer, indices, _semantic): + handles = [index.handle for index in indices] + if not COMMON_IR_ENABLED: + return _semantic.builder.create_local_pointers(result_ir, buffer.handle, *handles) + return _semantic.builder.create_tile_local_ptr( + result_ir, + buffer.handle, + *handles, + ) + + +def get_memdesc(buffer, _semantic): + """Bridge a TileIR buffer to the descriptor type consumed by TLE pipe ops.""" + if not COMMON_IR_ENABLED: + return buffer.handle + result_type = buffer.type.to_memdesc_ir(_semantic.builder) + return _semantic.builder.create_tile_get_memdesc(result_type, buffer.handle) + + class TLESemanticError(Exception): """TLE operation semantic error exception class""" diff --git a/python/triton/experimental/tle/language/gpu/types.py b/python/triton/experimental/tle/language/gpu/types.py index ff7fd0a3eb..58c0b1d436 100644 --- a/python/triton/experimental/tle/language/gpu/types.py +++ b/python/triton/experimental/tle/language/gpu/types.py @@ -320,17 +320,18 @@ def slot(self, stage, _semantic=None): slot_layout = _make_slot_layout(self.type.layout, slot_shape) slot_ty = buffered_tensor_type(self.dtype, slot_shape, self.type.storage, slot_layout, _semantic, alloc_shape=slot_shape) - if hasattr(_semantic.builder, "create_tile_subview"): - offsets = [stage_tensor.handle] - for _ in range(len(self.shape) - 1): - offsets.append(_semantic.to_tensor(0).handle) - slot_handle = _semantic.builder.create_tile_subview(self.handle, offsets, slot_shape, - [1] * len(slot_shape), - slot_layout.to_ir(_semantic.builder)) - return buffered_tensor(slot_handle, self.dtype, slot_shape, self.type.storage, slot_layout, _semantic, - alloc_shape=slot_ty.alloc_shape) - slot_handle = _semantic.builder.create_memdesc_index(slot_ty.to_ir(_semantic.builder), self.handle, - stage_tensor.handle) + offsets = [stage_tensor] + for _ in range(len(self.shape) - 1): + offsets.append(_semantic.to_tensor(0)) + from . import semantic as tle_semantic + slot_handle = tle_semantic.subview( + self, + offsets, + slot_shape, + [1] * len(slot_shape), + slot_layout, + _semantic, + ) return buffered_tensor(slot_handle, self.dtype, slot_shape, self.type.storage, slot_layout, _semantic, alloc_shape=slot_ty.alloc_shape) @@ -425,15 +426,19 @@ def _flatten_ir_types(self, builder: ir.builder, out: List[ir.type]) -> None: def to_ir(self, builder: ir.builder) -> None: shape = self.shape builder = self.semantic.builder - if hasattr(builder, "tile_get_buffer_type"): + from .semantic import COMMON_IR_ENABLED + if COMMON_IR_ENABLED: memory_space = builder.tile_get_string_attr(_storage_to_tile_space(self.storage)) return builder.tile_get_buffer_type( [int(tl._unwrap_if_constexpr(dim)) for dim in shape], self.element_ty.to_ir(builder), memory_space, ) + return self.to_memdesc_ir(builder) + + def to_memdesc_ir(self, builder: ir.builder): return builder.get_memdesc_type( - shape, + self.shape, self.element_ty.to_ir(builder), self.layout.to_ir(builder), _storage_to_memdesc_space(self.storage), @@ -644,8 +649,9 @@ def type(self): [(name, value.type) for name, value in self.fields.items()], self.readers, one_shot=self.one_shot) - def _field_handles(self): - return [field.handle for field in self.fields.values()] + def _field_handles(self, _semantic): + from . import semantic as tle_semantic + return [tle_semantic.get_memdesc(field, _semantic) for field in self.fields.values()] def _field_names(self): return list(self.fields.keys()) @@ -771,7 +777,7 @@ def __init__(self, pipe: pipe_value, reader_name=None, field_names=None): @tl.builtin def acquire(self, iter, _semantic=None): stage, phase = self.pipe._stage_phase(iter, _semantic=_semantic) - _semantic.builder.create_pipe_writer_acquire(self.pipe._field_handles(), stage.handle, phase.handle, + _semantic.builder.create_pipe_writer_acquire(self.pipe._field_handles(_semantic), stage.handle, phase.handle, self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names()) return self.pipe._make_slot(stage, _semantic=_semantic) @@ -779,7 +785,8 @@ def acquire(self, iter, _semantic=None): @tl.builtin def commit(self, iter, _semantic=None): stage, _ = self.pipe._stage_phase(iter, _semantic=_semantic) - _semantic.builder.create_pipe_writer_commit(self.pipe._field_handles(), stage.handle, self.pipe.capacity, + _semantic.builder.create_pipe_writer_commit(self.pipe._field_handles(_semantic), stage.handle, + self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names()) @tl.builtin @@ -787,7 +794,7 @@ def close(self, iter, _semantic=None): if self.pipe.one_shot: raise ValueError("tle.pipe one_shot pipes do not support close") stage, phase = self.pipe._stage_phase(iter, _semantic=_semantic) - _semantic.builder.create_pipe_writer_close(self.pipe._field_handles(), stage.handle, phase.handle, + _semantic.builder.create_pipe_writer_close(self.pipe._field_handles(_semantic), stage.handle, phase.handle, self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names()) @@ -804,7 +811,8 @@ def _reader_field_names(self): @tl.builtin def wait(self, iter, _semantic=None): stage, phase = self.pipe._stage_phase(iter, _semantic=_semantic) - is_closed = _semantic.builder.create_pipe_reader_wait(self.pipe._field_handles(), stage.handle, phase.handle, + is_closed = _semantic.builder.create_pipe_reader_wait(self.pipe._field_handles(_semantic), stage.handle, + phase.handle, self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names(), self.reader_name or "", self._reader_field_names()) @@ -814,7 +822,7 @@ def wait(self, iter, _semantic=None): @tl.builtin def release(self, iter, _semantic=None): stage, _ = self.pipe._stage_phase(iter, _semantic=_semantic) - _semantic.builder.create_pipe_reader_release(self.pipe._field_handles(), stage.handle, + _semantic.builder.create_pipe_reader_release(self.pipe._field_handles(_semantic), stage.handle, self.pipe.capacity, self.pipe.scope, self.pipe._ir_name(), self.pipe._field_names(), self.reader_name or "", self._reader_field_names()) diff --git a/python/triton/experimental/tle/language/pipe.py b/python/triton/experimental/tle/language/pipe.py index a40f455292..bfd9115d14 100644 --- a/python/triton/experimental/tle/language/pipe.py +++ b/python/triton/experimental/tle/language/pipe.py @@ -103,9 +103,10 @@ def pipe( raise ValueError( f"tle.pipe field {field_name!r} leading dimension must equal capacity {capacity}, got {field.shape[0]}") - _semantic.builder.create_pipe_create([field.handle for field in fields.values()], capacity, scope, name or "", + value = gpu_types.pipe_value(capacity, scope, name, fields, reader_names, one_shot=one_shot) + _semantic.builder.create_pipe_create(value._field_handles(_semantic), capacity, scope, name or "", list(fields.keys()), list(reader_names or ()), one_shot) - return gpu_types.pipe_value(capacity, scope, name, fields, reader_names, one_shot=one_shot) + return value pipe_slot = gpu_types.pipe_slot diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 03cd4a020b..39924f4924 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -2,6 +2,7 @@ add_subdirectory(lib) llvm_canonicalize_cmake_booleans( MLIR_ENABLE_BINDINGS_PYTHON + FLAGTREE_COMMON_IR_ENABLED ) configure_lit_site_cfg( diff --git a/test/CommonIR/gpu-tileir-to-ttgir.mlir b/test/CommonIR/gpu-tileir-to-ttgir.mlir new file mode 100644 index 0000000000..dab2873a98 --- /dev/null +++ b/test/CommonIR/gpu-tileir-to-ttgir.mlir @@ -0,0 +1,33 @@ +// REQUIRES: flagtree-common-ir +// RUN: triton-opt %s --convert-common-ir-to-ttgir | FileCheck %s + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#shared1 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module { + tt.func public @gpu_tileir_to_ttgir(%src: tensor<16x!tt.ptr>) { + %buf = tile.alloc {layout = 0 : i64, space = 7 : i64, tle.gpu_layout = #shared} + : <[2, 16], f32, shared> + %desc = tile.get_memdesc %buf + : !tile.buf<[2, 16], f32, shared> -> !ttg.memdesc<2x16xf32, #shared, #smem, mutable> + tle.pipe.create %desc {capacity = 2 : i32, field_names = ["payload"], pipe_name = "p", scope = "cta"} + : !ttg.memdesc<2x16xf32, #shared, #smem, mutable> + %c0 = arith.constant 0 : index + %slot = tile.subview %buf[%c0, %c0] [[16]] [[1]] {tle.gpu_layout = #shared1} + : <[2, 16], f32, shared> -> <[16], f32, shared> + tile.copy %src -> %slot {src_layout = 0 : i64} + : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + %ptr = "tile.local_ptr"(%slot) + : (!tile.buf<[16], f32, shared>) -> tensor<16x!tt.ptr> + tt.return + } +} + +// CHECK: %[[BUF:.*]] = ttg.local_alloc : () -> !ttg.memdesc<2x16xf32, #shared, #smem, mutable> +// CHECK: tle.pipe.create %[[BUF]] +// CHECK: %[[SLOT:.*]] = ttg.memdesc_index %[[BUF]][%{{.*}}] +// CHECK: tt.load %{{.*}} : tensor<16x!tt.ptr> +// CHECK: "tle.local_pointers"(%[[SLOT]], +// CHECK: "tle.local_pointers"(%[[SLOT]]) +// CHECK-NOT: tile. diff --git a/test/CommonIR/native_flash_gpu.py b/test/CommonIR/native_flash_gpu.py index e7894c650d..8af6c9b262 100644 --- a/test/CommonIR/native_flash_gpu.py +++ b/test/CommonIR/native_flash_gpu.py @@ -119,7 +119,7 @@ def dump_tileir(path): def dump_ttir(path): from triton._C.libtriton import ir from triton._C.libtriton import passes - from triton._C.libtriton import tle as tle_ir + from triton._C.libtriton import nvidia constants = { "N_CTX": 32, @@ -147,7 +147,7 @@ def dump_ttir(path): passes.common.add_inliner(pm) pm.run(module, "native_flash_gpu.inliner") pm = ir.pass_manager(module.context) - tle_ir.passes.add_convert_gpu_tile_to_ttgir(pm) + nvidia.passes.commonir.add_to_ttgir(pm) pm.run(module, "native_flash_gpu.tileir_to_ttgir") text = str(module) for needle in ("tile.alloc", "tile.copy", "tile.to_tensor", "tile.store_tensor"): diff --git a/test/CommonIR/native_matmul_gpu.py b/test/CommonIR/native_matmul_gpu.py index b753c8ae98..b971f9419c 100644 --- a/test/CommonIR/native_matmul_gpu.py +++ b/test/CommonIR/native_matmul_gpu.py @@ -101,7 +101,7 @@ def dump_tileir(path): def dump_ttir(path): from triton._C.libtriton import ir from triton._C.libtriton import passes - from triton._C.libtriton import tle as tle_ir + from triton._C.libtriton import nvidia constants = { "BLOCK_M": 16, @@ -127,7 +127,7 @@ def dump_ttir(path): passes.common.add_inliner(pm) pm.run(module, "native_matmul_gpu.inliner") pm = ir.pass_manager(module.context) - tle_ir.passes.add_convert_gpu_tile_to_ttgir(pm) + nvidia.passes.commonir.add_to_ttgir(pm) pm.run(module, "native_matmul_gpu.tileir_to_ttgir") text = str(module) for needle in ("tile.alloc", "tile.copy", "tile.to_tensor", "tile.store_tensor"): diff --git a/test/CommonIR/tle_gpu_primitives.py b/test/CommonIR/tle_gpu_primitives.py index 9c880de8e0..d5c3db6df9 100644 --- a/test/CommonIR/tle_gpu_primitives.py +++ b/test/CommonIR/tle_gpu_primitives.py @@ -66,6 +66,7 @@ def dump_tileir(path): "tile.copy", "tile.subview", "tile.local_ptr", + "tile.get_memdesc", "tle.pipe.create", "tle.pipe.writer_acquire", "tle.pipe.reader_wait", @@ -79,17 +80,17 @@ def dump_tileir(path): def dump_ttir(path): from triton._C.libtriton import ir from triton._C.libtriton import passes - from triton._C.libtriton import tle as tle_ir + from triton._C.libtriton import nvidia module = _frontend_module() pm = ir.pass_manager(module.context) passes.common.add_inliner(pm) pm.run(module, "tle_gpu_primitives.inliner") pm = ir.pass_manager(module.context) - tle_ir.passes.add_convert_gpu_tile_to_ttgir(pm) + nvidia.passes.commonir.add_to_ttgir(pm) pm.run(module, "tle_gpu_primitives.tileir_to_ttgir") text = str(module) - for needle in ("tile.alloc", "tile.copy", "tile.subview", "tile.local_ptr"): + for needle in ("tile.alloc", "tile.copy", "tile.subview", "tile.local_ptr", "tile.get_memdesc"): if needle in text: raise RuntimeError(f"primitive converted TTIR still contains {needle}") for needle in ("ttg.local_alloc", "ttg.memdesc_index", "tle.local_pointers", "tle.pipe.create"): diff --git a/test/lit.cfg.py b/test/lit.cfg.py index 0890e609b7..5e9960830c 100644 --- a/test/lit.cfg.py +++ b/test/lit.cfg.py @@ -19,6 +19,9 @@ # suffixes: A list of file extensions to treat as test files. config.suffixes = ['.mlir', '.ll'] +if getattr(config, "flagtree_common_ir_enabled", False): + config.available_features.add("flagtree-common-ir") + # test_source_root: The root path where tests are located. config.test_source_root = os.path.dirname(__file__) diff --git a/test/lit.site.cfg.py.in b/test/lit.site.cfg.py.in index fd1fb486dd..5eb2e9373e 100644 --- a/test/lit.site.cfg.py.in +++ b/test/lit.site.cfg.py.in @@ -14,6 +14,7 @@ config.lit_tools_dir = "@LLVM_LIT_TOOLS_DIR@" config.mlir_binary_dir = "@MLIR_BINARY_DIR@" config.python_executable = "@Python3_EXECUTABLE@" config.enable_bindings_python = @MLIR_ENABLE_BINDINGS_PYTHON@ +config.flagtree_common_ir_enabled = @FLAGTREE_COMMON_IR_ENABLED@ import lit.llvm diff --git a/third_party/nvidia/CMakeLists.txt b/third_party/nvidia/CMakeLists.txt index ff1ea248de..ed33015791 100644 --- a/third_party/nvidia/CMakeLists.txt +++ b/third_party/nvidia/CMakeLists.txt @@ -4,6 +4,9 @@ add_subdirectory(include) add_subdirectory(lib) if(TRITON_BUILD_PYTHON_MODULE) add_triton_plugin(TritonNVIDIA ${CMAKE_CURRENT_SOURCE_DIR}/triton_nvidia.cc LINK_LIBS TritonNVIDIAGPUToLLVM NVGPUToLLVM) + if(FLAGTREE_COMMON_IR_ENABLED AND TARGET CommonIRToTTGIR) + target_link_libraries(TritonNVIDIA PRIVATE CommonIRToTTGIR) + endif() target_link_libraries(TritonNVIDIA PRIVATE Python3::Module pybind11::headers) endif() if(TRITON_BUILD_UT) diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index f52ad1545d..7b5991246a 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -15,6 +15,7 @@ import os import subprocess from pathlib import Path +from triton._common_ir import ENABLED as COMMON_IR_ENABLED from .distributed import Distributed @@ -266,6 +267,11 @@ def make_ttir(mod, metadata, opt, capability): pm = ir.pass_manager(mod.context) pm.enable_debug() + # Lower TLE GPU CommonIR buffers before TTIR canonicalization. The + # upstream TileIR allocation op is pure, so leaving it until make_ttgir + # would allow CSE to merge distinct shared-memory allocations. + if COMMON_IR_ENABLED: + nvidia.passes.commonir.add_to_ttgir(pm) passes.ttir.add_rewrite_tensor_pointer(pm) if capability // 10 < 9: passes.ttir.add_rewrite_tensor_descriptor_to_pointer(pm) @@ -287,10 +293,6 @@ def make_ttgir(mod, metadata, opt, capability): pm = ir.pass_manager(mod.context) dump_enabled = pm.enable_debug() emuTF32 = (capability // 10 >= 8) - # TLE GPU always enters through TileIR when the dialect is available. - # Convert the target-independent buffers to TritonGPU descriptors - # before the standard TTIR-to-TTGIR conversion. - tle.passes.add_convert_gpu_tile_to_ttgir(pm) passes.ttir.add_convert_to_ttgpuir(pm, f"cuda:{capability}", opt.num_warps, 32, opt.num_ctas) # flagtree tle raw tle.raw_passes.add_tle_convert_arg_to_memdesc(pm) diff --git a/third_party/nvidia/include/CMakeLists.txt b/third_party/nvidia/include/CMakeLists.txt index 2ef7aab106..184864b537 100644 --- a/third_party/nvidia/include/CMakeLists.txt +++ b/third_party/nvidia/include/CMakeLists.txt @@ -1,3 +1,6 @@ add_subdirectory(Dialect) add_subdirectory(TritonNVIDIAGPUToLLVM) add_subdirectory(NVGPUToLLVM) +if(FLAGTREE_COMMON_IR_ENABLED AND TARGET TileIRIR) + add_subdirectory(CommonIRToTTGIR) +endif() diff --git a/third_party/nvidia/include/CommonIRToTTGIR/CMakeLists.txt b/third_party/nvidia/include/CommonIRToTTGIR/CMakeLists.txt new file mode 100644 index 0000000000..c338774421 --- /dev/null +++ b/third_party/nvidia/include/CommonIRToTTGIR/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name CommonIRToTTGIR) +add_public_tablegen_target(CommonIRToTTGIRPassIncGen) diff --git a/third_party/nvidia/include/CommonIRToTTGIR/Passes.h b/third_party/nvidia/include/CommonIRToTTGIR/Passes.h new file mode 100644 index 0000000000..0b5c60a4b2 --- /dev/null +++ b/third_party/nvidia/include/CommonIRToTTGIR/Passes.h @@ -0,0 +1,16 @@ +#ifndef TRITON_NVIDIA_COMMON_IR_TO_TTGIR_PASSES_H +#define TRITON_NVIDIA_COMMON_IR_TO_TTGIR_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir::triton { + +#define GEN_PASS_DECL +#include "nvidia/include/CommonIRToTTGIR/Passes.h.inc" + +#define GEN_PASS_REGISTRATION +#include "nvidia/include/CommonIRToTTGIR/Passes.h.inc" + +} // namespace mlir::triton + +#endif diff --git a/third_party/nvidia/include/CommonIRToTTGIR/Passes.td b/third_party/nvidia/include/CommonIRToTTGIR/Passes.td new file mode 100644 index 0000000000..9ce5ddef46 --- /dev/null +++ b/third_party/nvidia/include/CommonIRToTTGIR/Passes.td @@ -0,0 +1,22 @@ +#ifndef TRITON_NVIDIA_COMMON_IR_TO_TTGIR_PASSES +#define TRITON_NVIDIA_COMMON_IR_TO_TTGIR_PASSES + +include "mlir/Pass/PassBase.td" + +def CommonIRToTTGIR + : Pass<"convert-common-ir-to-ttgir", "mlir::ModuleOp"> { + let summary = "convert GPU CommonIR buffers to TritonGPU descriptors"; + let description = [{ + Converts the target-independent TileIR produced by tle.gpu primitives into + the NVIDIA TritonGPU descriptor form. In particular, tile.alloc becomes + ttg.local_alloc with a ttg.memdesc result and tile.local_ptr becomes + tle.local_pointers. + }]; + let dependentDialects = ["mlir::arith::ArithDialect", + "mlir::triton::TritonDialect", + "mlir::triton::gpu::TritonGPUDialect", + "mlir::triton::tle::TleDialect", + "mlir::triton::tile::TileIRDialect"]; +} + +#endif diff --git a/third_party/nvidia/lib/CMakeLists.txt b/third_party/nvidia/lib/CMakeLists.txt index 2ef7aab106..184864b537 100644 --- a/third_party/nvidia/lib/CMakeLists.txt +++ b/third_party/nvidia/lib/CMakeLists.txt @@ -1,3 +1,6 @@ add_subdirectory(Dialect) add_subdirectory(TritonNVIDIAGPUToLLVM) add_subdirectory(NVGPUToLLVM) +if(FLAGTREE_COMMON_IR_ENABLED AND TARGET TileIRIR) + add_subdirectory(CommonIRToTTGIR) +endif() diff --git a/third_party/nvidia/lib/CommonIRToTTGIR/CMakeLists.txt b/third_party/nvidia/lib/CommonIRToTTGIR/CMakeLists.txt new file mode 100644 index 0000000000..b404cc68f3 --- /dev/null +++ b/third_party/nvidia/lib/CommonIRToTTGIR/CMakeLists.txt @@ -0,0 +1,13 @@ +add_triton_library(CommonIRToTTGIR + CommonIRToTTGIR.cpp + + DEPENDS + CommonIRToTTGIRPassIncGen + + LINK_LIBS PUBLIC + MLIRArithDialect + TileIRIR + TleIR + TritonIR + TritonGPUIR +) diff --git a/third_party/tle/dialect/lib/Transforms/TleConvertGpuTileToTtgir.cpp b/third_party/nvidia/lib/CommonIRToTTGIR/CommonIRToTTGIR.cpp similarity index 69% rename from third_party/tle/dialect/lib/Transforms/TleConvertGpuTileToTtgir.cpp rename to third_party/nvidia/lib/CommonIRToTTGIR/CommonIRToTTGIR.cpp index ac5777f95e..e7ce1f3975 100644 --- a/third_party/tle/dialect/lib/Transforms/TleConvertGpuTileToTtgir.cpp +++ b/third_party/nvidia/lib/CommonIRToTTGIR/CommonIRToTTGIR.cpp @@ -2,7 +2,7 @@ // // Copyright (c) 2025 The FlagOS Contributors -#include "tle/dialect/include/Transforms/Passes.h" +#include "nvidia/include/CommonIRToTTGIR/Passes.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/IR/BuiltinOps.h" @@ -10,20 +10,17 @@ #include "triton/Dialect/Triton/IR/Dialect.h" #include "triton/Dialect/TritonGPU/IR/Dialect.h" -#ifdef __FLIR_TILEIR__ #include "mlir-ext/Dialect/TileIR/IR/TileIRDialect.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" -#endif -namespace mlir::triton::tle { +namespace mlir::triton { -#define GEN_PASS_DEF_TRITONTLECONVERTGPUTILETOTTGIR -#include "tle/dialect/include/Transforms/Passes.h.inc" +#define GEN_PASS_DEF_COMMONIRTOTTGIR +#include "nvidia/include/CommonIRToTTGIR/Passes.h.inc" namespace { -#ifdef __FLIR_TILEIR__ namespace tile = mlir::triton::tile; namespace ttg = mlir::triton::gpu; @@ -58,6 +55,101 @@ static FailureOr lookupBuffer(Operation *op, Value buffer, return it->second; } +static FailureOr inferMemDescType( + ModuleOp module, triton::FuncOp func, BlockArgument argument, + const llvm::DenseMap &mapped) { + ttg::MemDescType inferred; + for (Operation *user : argument.getUsers()) { + auto bridge = dyn_cast(user); + if (!bridge || bridge.getSource() != argument) + continue; + auto type = cast(bridge.getResult().getType()); + if (inferred && inferred != type) { + bridge.emitError("conflicting descriptor types for the same TileIR buffer"); + return failure(); + } + inferred = type; + } + if (!inferred) { + unsigned argumentNumber = argument.getArgNumber(); + module.walk([&](triton::CallOp call) { + if (inferred || module.lookupSymbol(call.getCallee()) != func || + argumentNumber >= call.getNumOperands()) + return; + Value operand = call.getOperand(argumentNumber); + auto it = mapped.find(operand); + if (it != mapped.end()) + operand = it->second; + inferred = dyn_cast(operand.getType()); + }); + } + if (!inferred) { + emitError(argument.getLoc()) + << "cannot infer a TritonGPU descriptor type for TileIR block argument"; + return failure(); + } + return inferred; +} + +static LogicalResult convertFunctionArguments( + ModuleOp module, llvm::DenseMap &mapped, + llvm::DenseMap &convertedArgumentTypes, + SmallVectorImpl &argumentCasts) { + for (auto func : module.getOps()) { + if (func.getBody().empty()) + continue; + Block &entry = func.getBody().front(); + for (BlockArgument argument : entry.getArguments()) { + if (!isa(argument.getType())) + continue; + auto type = inferMemDescType(module, func, argument, mapped); + if (failed(type)) + return failure(); + OpBuilder builder(func.getContext()); + builder.setInsertionPointToStart(&entry); + auto cast = builder.create( + argument.getLoc(), TypeRange{*type}, ValueRange{argument}); + mapped[argument] = cast.getResult(0); + convertedArgumentTypes[argument] = *type; + argumentCasts.push_back(cast); + } + } + return success(); +} + +static LogicalResult convertWarpSpecializeCaptures( + ModuleOp module, llvm::DenseMap &mapped, + llvm::DenseMap &convertedArgumentTypes, + SmallVectorImpl &argumentCasts) { + bool failedConversion = false; + module.walk([&](ttg::WarpSpecializeOp op) { + if (failedConversion) + return; + auto captures = op.getExplicitCaptures(); + for (Region *region : op.getPartitionRegions()) { + for (auto [index, capture] : llvm::enumerate(captures)) { + BlockArgument argument = region->getArgument(index); + if (!isa(argument.getType())) + continue; + auto converted = lookupBuffer(op, capture, mapped); + if (failed(converted)) { + failedConversion = true; + return; + } + OpBuilder builder(op.getContext()); + builder.setInsertionPointToStart(®ion->front()); + auto cast = builder.create( + argument.getLoc(), TypeRange{(*converted).getType()}, + ValueRange{argument}); + mapped[argument] = cast.getResult(0); + convertedArgumentTypes[argument] = (*converted).getType(); + argumentCasts.push_back(cast); + } + } + }); + return failure(failedConversion); +} + static bool isPointerLike(Type type) { if (isa(type)) return true; @@ -166,6 +258,42 @@ static LogicalResult convertAllocations( return failure(result.wasInterrupted()); } +static void rewriteConvertedOperands( + ModuleOp module, const llvm::DenseMap &mapped) { + module.walk([&](Operation *op) { + if (isa(op)) + return; + for (OpOperand &operand : op->getOpOperands()) { + auto it = mapped.find(operand.get()); + if (it != mapped.end()) + operand.set(it->second); + } + }); +} + +static void finalizeConvertedArguments( + ModuleOp module, const llvm::DenseMap &convertedArgumentTypes, + ArrayRef argumentCasts) { + for (auto [value, type] : convertedArgumentTypes) + cast(value).setType(type); + + for (UnrealizedConversionCastOp cast : argumentCasts) { + cast.getResult(0).replaceAllUsesWith(cast.getInputs().front()); + cast.erase(); + } + + for (auto func : module.getOps()) { + if (func.getBody().empty()) + continue; + SmallVector inputTypes(func.getBody().front().getArgumentTypes()); + SmallVector resultTypes(func.getResultTypes()); + func.setFunctionType( + FunctionType::get(func.getContext(), inputTypes, resultTypes)); + } +} + static LogicalResult convertSubviews( ModuleOp module, llvm::DenseMap &mapped, SmallVectorImpl &eraseOps) { @@ -239,6 +367,17 @@ static LogicalResult convertTileUsers( return; } + if (auto op = dyn_cast(operation)) { + auto source = lookupBuffer(op, op.getSource(), mapped); + if (failed(source)) { + failedConversion = true; + return; + } + op.getResult().replaceAllUsesWith(*source); + eraseOps.push_back(op); + return; + } + if (auto op = dyn_cast(operation)) { auto source = lookupBuffer(op, op.getSrc(), mapped); if (failed(source)) { @@ -287,18 +426,7 @@ static LogicalResult convertTileUsers( } builder.setInsertionPoint(op); - if (!op.getIndices().empty()) { - SmallVector indices; - for (Value index : op.getIndices()) { - auto converted = castToI32(builder, op.getLoc(), index); - if (failed(converted)) { - failedConversion = true; - return; - } - indices.push_back(*converted); - } - builder.create(op.getLoc(), src, dst, indices); - } else if (isa(dst.getType())) { + if (isa(dst.getType())) { Value value = src; if (isPointerLike(src.getType())) value = builder.create( @@ -344,44 +472,30 @@ static LogicalResult convertTileUsers( return failure(failedConversion); } -static LogicalResult rewritePipeOperands( - ModuleOp module, const llvm::DenseMap &mapped) { - bool failedConversion = false; - module.walk([&](Operation *op) { - if (!op->getName().getStringRef().starts_with("tle.pipe.")) - return; - for (OpOperand &operand : op->getOpOperands()) { - if (!isa(operand.get().getType())) - continue; - auto converted = lookupBuffer(op, operand.get(), mapped); - if (failed(converted)) { - failedConversion = true; - return; - } - operand.set(*converted); - } - }); - return failure(failedConversion); -} -#endif - -class ConvertGpuTileToTtgirPass - : public impl::TritonTleConvertGpuTileToTtgirBase< - ConvertGpuTileToTtgirPass> { +class CommonIRToTTGIRPass + : public impl::CommonIRToTTGIRBase { void runOnOperation() override { -#ifdef __FLIR_TILEIR__ ModuleOp module = getOperation(); llvm::DenseMap mapped; + llvm::DenseMap convertedArgumentTypes; SmallVector eraseOps; + SmallVector argumentCasts; if (failed(convertAllocations(module, mapped, eraseOps)) || + failed(convertWarpSpecializeCaptures(module, mapped, + convertedArgumentTypes, + argumentCasts)) || + failed(convertFunctionArguments(module, mapped, + convertedArgumentTypes, + argumentCasts)) || failed(convertSubviews(module, mapped, eraseOps)) || - failed(convertTileUsers(module, mapped, eraseOps)) || - failed(rewritePipeOperands(module, mapped))) { + failed(convertTileUsers(module, mapped, eraseOps))) { signalPassFailure(); return; } + rewriteConvertedOperands(module, mapped); + for (Operation *op : llvm::reverse(eraseOps)) { if (llvm::any_of(op->getResults(), [](Value result) { return !result.use_empty(); })) { @@ -392,20 +506,22 @@ class ConvertGpuTileToTtgirPass op->erase(); } + finalizeConvertedArguments(module, convertedArgumentTypes, argumentCasts); + bool hasRemainingGpuTileOps = false; module.walk([&](Operation *op) { if (isa(op)) { + tile::CopyOp, tile::ToTensorOp, tile::StoreTensorOp, + tile::GetMemDescOp>(op)) { op->emitError("was not eliminated by GPU TileIR conversion"); hasRemainingGpuTileOps = true; } }); if (hasRemainingGpuTileOps) signalPassFailure(); -#endif } }; } // namespace -} // namespace mlir::triton::tle +} // namespace mlir::triton diff --git a/third_party/nvidia/triton_nvidia.cc b/third_party/nvidia/triton_nvidia.cc index 70ffcb4707..42c9fe8eae 100644 --- a/third_party/nvidia/triton_nvidia.cc +++ b/third_party/nvidia/triton_nvidia.cc @@ -2,6 +2,9 @@ #include "Dialect/NVWS/IR/Dialect.h" #include "NVGPUToLLVM/Passes.h" #include "TritonNVIDIAGPUToLLVM/Passes.h" +#ifdef FLAGTREE_COMMON_IR +#include "CommonIRToTTGIR/Passes.h" +#endif #include "cublas_instance.h" #include "mlir/Pass/PassManager.h" #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h" @@ -88,6 +91,12 @@ void init_triton_nvidia_passes_nvws(py::module &&m) { mlir::triton::createNVWSInsertTmemAref); } +#ifdef FLAGTREE_COMMON_IR +void init_triton_nvidia_passes_commonir(py::module &&m) { + ADD_PASS_WRAPPER_0("add_to_ttgir", mlir::triton::createCommonIRToTTGIR); +} +#endif + void init_triton_hopper_passes(py::module &&m) { // Meta's autoWS ADD_PASS_OPTION_WRAPPER_2("add_hopper_warpspec", @@ -148,6 +157,9 @@ void init_triton_nvidia(py::module &&m) { init_triton_nvidia_passes_ttgpuir(passes.def_submodule("ttgpuir")); init_triton_nvidia_passes_ttnvgpuir(passes.def_submodule("ttnvgpuir")); init_triton_hopper_passes(passes.def_submodule("hopper")); +#ifdef FLAGTREE_COMMON_IR + init_triton_nvidia_passes_commonir(passes.def_submodule("commonir")); +#endif // load dialects m.def("load_dialects", [](mlir::MLIRContext &context) { diff --git a/third_party/tle/dialect/include/IR/TleOps.td b/third_party/tle/dialect/include/IR/TleOps.td index 147b3b9bfa..ec89a213c3 100644 --- a/third_party/tle/dialect/include/IR/TleOps.td +++ b/third_party/tle/dialect/include/IR/TleOps.td @@ -101,7 +101,7 @@ def Tle_MemDescWGMMAViewOp : Tle_Op<"memdesc_wgmma_view", [Pure, def Tle_PipeCreateOp : Tle_Op<"pipe.create"> { let summary = "create a typed TLE GPU pipe"; let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32Attr:$capacity, StrAttr:$scope, OptionalAttr:$pipe_name, @@ -115,7 +115,7 @@ def Tle_PipeCreateOp : Tle_Op<"pipe.create"> { def Tle_PipeWriterAcquireOp : Tle_Op<"pipe.writer_acquire"> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I1:$phase, I32Attr:$capacity, @@ -129,7 +129,7 @@ def Tle_PipeWriterAcquireOp : Tle_Op<"pipe.writer_acquire"> { def Tle_PipeWriterCommitOp : Tle_Op<"pipe.writer_commit"> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I32Attr:$capacity, StrAttr:$scope, @@ -142,7 +142,7 @@ def Tle_PipeWriterCommitOp : Tle_Op<"pipe.writer_commit"> { def Tle_PipeWriterCloseOp : Tle_Op<"pipe.writer_close"> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I1:$phase, I32Attr:$capacity, @@ -156,7 +156,7 @@ def Tle_PipeWriterCloseOp : Tle_Op<"pipe.writer_close"> { def Tle_PipeReaderWaitOp : Tle_Op<"pipe.reader_wait"> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I1:$phase, I32Attr:$capacity, @@ -174,7 +174,7 @@ def Tle_PipeReaderReleaseOp : Tle_Op<"pipe.reader_release", [ DeclareOpInterfaceMethods ]> { let arguments = (ins - Variadic:$fields, + Variadic:$fields, I32:$stage, I32Attr:$capacity, StrAttr:$scope, diff --git a/third_party/tle/dialect/include/Transforms/Passes.td b/third_party/tle/dialect/include/Transforms/Passes.td index fb9916a959..bef5701756 100644 --- a/third_party/tle/dialect/include/Transforms/Passes.td +++ b/third_party/tle/dialect/include/Transforms/Passes.td @@ -4,25 +4,6 @@ include "mlir/Pass/PassBase.td" -def TritonTleConvertGpuTileToTtgir - : Pass<"triton-tle-convert-gpu-tile-to-ttgir", "mlir::ModuleOp"> { - let summary = "convert TLE GPU TileIR buffers to TritonGPU descriptors"; - - let description = [{ - Converts the target-independent TileIR produced by tle.gpu primitives into - the TritonGPU descriptor form expected by the NVIDIA lowering pipeline. - In particular, tile.alloc becomes ttg.local_alloc with a ttg.memdesc result, - and tile.local_ptr becomes tle.local_pointers. The conversion runs before - the standard TTIR-to-TTGIR pass and must eliminate all GPU buffer TileIR - operations. - }]; - - let dependentDialects = ["mlir::arith::ArithDialect", - "mlir::triton::TritonDialect", - "mlir::triton::gpu::TritonGPUDialect", - "mlir::triton::tle::TleDialect"]; -} - def TritonTleEarlyAssignMemorySpace : Pass<"triton-tle-early-assign-memory-space", "mlir::ModuleOp"> { let summary = "early assign memory space"; diff --git a/third_party/tle/dialect/lib/IR/CMakeLists.txt b/third_party/tle/dialect/lib/IR/CMakeLists.txt index 0ad0cd1546..5e9e87a7cf 100644 --- a/third_party/tle/dialect/lib/IR/CMakeLists.txt +++ b/third_party/tle/dialect/lib/IR/CMakeLists.txt @@ -13,7 +13,3 @@ add_triton_library(TleIR TritonIR TritonGPUIR ) - -if(TARGET TileIRIR) - target_link_libraries(TleIR PUBLIC TileIRIR) -endif() diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp index c011494aee..3d26c14694 100644 --- a/third_party/tle/dialect/lib/IR/Ops.cpp +++ b/third_party/tle/dialect/lib/IR/Ops.cpp @@ -4,9 +4,6 @@ #include "triton/Dialect/Triton/IR/Types.h" #include "triton/Dialect/Triton/IR/Utility.h" #include "triton/Dialect/TritonGPU/IR/Dialect.h" -#ifdef __FLIR_TILEIR__ -#include "mlir-ext/Dialect/TileIR/IR/TileIRDialect.h" -#endif #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallSet.h" #include @@ -274,24 +271,12 @@ static LogicalResult verifyPipeAttrs(Operation *op, OperandRange fields) { if (fields.empty()) return op->emitOpError("expects at least one pipe field"); for (Value field : fields) { - ArrayRef shape; - if (auto type = dyn_cast(field.getType())) { - if (!isa(type.getMemorySpace())) - return op->emitOpError("expects only shared-memory pipe fields"); - shape = type.getShape(); -#ifdef __FLIR_TILEIR__ - } else if (auto type = dyn_cast(field.getType())) { - if (type.getMemorySpace() != triton::tile::MemorySpace::Shared) - return op->emitOpError("expects only shared-memory pipe fields"); - shape = type.getShape(); -#endif - } else { - return op->emitOpError( - "expects pipe fields to be TileIR buffers or ttg.memdesc values"); - } - if (shape.size() < 2) + auto type = cast(field.getType()); + if (!isa(type.getMemorySpace())) + return op->emitOpError("expects only shared-memory pipe fields"); + if (type.getRank() < 2) return op->emitOpError("expects pipe fields to have rank >= 2"); - if (shape[0] != capacity) + if (type.getShape()[0] != capacity) return op->emitOpError("expects field leading dimension to equal " "pipe capacity"); } diff --git a/third_party/tle/dialect/lib/Transforms/CMakeLists.txt b/third_party/tle/dialect/lib/Transforms/CMakeLists.txt index 5647f87c22..7d2e6c7f7f 100644 --- a/third_party/tle/dialect/lib/Transforms/CMakeLists.txt +++ b/third_party/tle/dialect/lib/Transforms/CMakeLists.txt @@ -1,6 +1,5 @@ # flagtree tle add_triton_library(TritonTLETransforms - TleConvertGpuTileToTtgir.cpp TleEarlyAssignMemorySpace.cpp TleSelectEncodings.cpp TleInsertLocalPointerBarriers.cpp @@ -38,10 +37,6 @@ add_triton_library(TritonTLETransforms NVWSIR ) -if(TARGET TileIRIR) - target_link_libraries(TritonTLETransforms PUBLIC TileIRIR) -endif() - target_include_directories(TritonTLETransforms PUBLIC ${CMAKE_SOURCE_DIR}/third_party/tle/dialect/include ${CMAKE_SOURCE_DIR}/include diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index b41c0b974f..69b5303a42 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -111,7 +111,6 @@ static tile::MemorySpace attrToTileMemSpace(Attribute attr) { #endif void init_triton_tle_ir(py::module &&m) { - // Get the existing builder class from the main ir module (TLX style) auto &builder_cls = *ir::getBuilderClass(); @@ -256,9 +255,9 @@ void init_triton_tle_ir(py::module &&m) { }) .def("create_tile_copy", [](TritonOpBuilder &self, Value &src, Value &dst, - std::vector &indices, bool interNoAlias) -> void { + bool interNoAlias) -> void { auto op = self.create( - src, dst, indices, /*engine=*/tile::EngineAttr(), + src, dst, /*engine=*/tile::EngineAttr(), /*src_layout=*/tile::LayoutAttr::get( self.getBuilder().getContext(), tile::Layout::ND), /*dst_nz_layout=*/tile::NZLayoutAttr(), @@ -267,6 +266,10 @@ void init_triton_tle_ir(py::module &&m) { if (interNoAlias) op->setAttr("inter_no_alias", self.getBuilder().getBoolAttr(true)); }) + .def("create_tile_get_memdesc", + [](TritonOpBuilder &self, Type resultTy, Value source) -> Value { + return self.create(resultTy, source); + }) .def("create_tile_subview", [](TritonOpBuilder &self, Value source, std::vector &offsets, const std::vector &sizes, @@ -612,8 +615,6 @@ void init_triton_tle_ir(py::module &&m) { } void init_triton_tle_passes(py::module &&m) { - ADD_PASS_WRAPPER_0("add_convert_gpu_tile_to_ttgir", - tle::createTritonTleConvertGpuTileToTtgir); ADD_PASS_WRAPPER_0("add_early_assign_memory_space", tle::createTritonTleEarlyAssignMemorySpace); ADD_PASS_WRAPPER_0("add_select_encodings", @@ -712,6 +713,14 @@ void init_llvm(py::module &&m) { } void init_triton_tle(py::module &&m) { + m.def("is_common_ir_enabled", []() { +#ifdef FLAGTREE_COMMON_IR + return true; +#else + return false; +#endif + }); + // load dialects m.def("load_dialects", [](mlir::MLIRContext &context) { mlir::DialectRegistry registry;