diff --git a/CMakeLists.txt b/CMakeLists.txt index e3014d3ed2..42561233eb 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,15 +265,19 @@ 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) +# CommonIR is only part of builds that explicitly enable CommonIR. +if(FLAGTREE_COMMON_IR_ENABLED AND NOT TARGET CommonIRIR) + 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) - add_subdirectory(third_party/flir/include/mlir-ext/Dialect/TileIR) - add_subdirectory(third_party/flir/lib/Dialect/TileIR) + add_subdirectory(third_party/flir/include/mlir-ext/Dialect/CommonIR) + add_subdirectory(third_party/flir/lib/Dialect/CommonIR) endif() # link_directories(${LLVM_LIBRARY_DIR}) diff --git a/bin/RegisterTritonDialects.h b/bin/RegisterTritonDialects.h index 9ea88fe388..a6baffca5c 100644 --- a/bin/RegisterTritonDialects.h +++ b/bin/RegisterTritonDialects.h @@ -15,7 +15,10 @@ #include "tle/dialect/include/Transforms/Passes.h" #endif #ifdef __FLIR_TILEIR__ -#include "mlir-ext/Dialect/TileIR/IR/TileIRDialect.h" +#include "mlir-ext/Dialect/CommonIR/IR/CommonIRDialect.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" @@ -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(); @@ -161,7 +167,7 @@ inline void registerTritonDialects(mlir::DialectRegistry ®istry) { mlir::triton::tle::TleDialect, // flagtree tle raw #endif #ifdef __FLIR_TILEIR__ - mlir::triton::tile::TileIRDialect, + mlir::triton::tile::CommonIRDialect, #endif mlir::triton::gluon::GluonDialect>(); } diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp index bf509ffe52..a0ba2948a3 100644 --- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp +++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp @@ -606,6 +606,9 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter, GenericOpPattern, GenericOpPattern, GenericOpPattern, +#endif +#ifdef FLAGTREE_COMMON_IR + GenericOpPattern, #endif TritonExpandDimsPattern, TritonTransPattern, @@ -975,6 +978,10 @@ class ConvertTritonToTritonGPU TritonGPUTypeConverter typeConverter(context, numWarps, threadsPerWarp, numCTAs, enableSourceRemat); TritonGPUConversionTarget target(*context, typeConverter); +#ifdef FLAGTREE_COMMON_IR + target.addDynamicallyLegalOp( + [&](Operation *op) { return typeConverter.isLegal(op); }); +#endif // rewrite patterns RewritePatternSet patterns(context); // add rules diff --git a/python/setup_tools/utils/__init__.py b/python/setup_tools/utils/__init__.py index 4c618155b5..e273ebad1e 100644 --- a/python/setup_tools/utils/__init__.py +++ b/python/setup_tools/utils/__init__.py @@ -38,7 +38,11 @@ def _register_submodule(self, submodule): "name": "triton_shared", "url": "https://github.com/microsoft/triton-shared.git", "commit_id": "5842469a16b261e45a2c67fbfc308057622b03ee" }, - {"name": "flir", "url": "https://github.com/FlagTree/flir.git"}, + { + "name": "flir", + "url": "https://github.com/kateyijian/flir.git", + "commit_id": "a884fd3c608db2203f08ec5a281327f5dc0153b0", + }, {"name": "flagcx", "url": "https://github.com/flagos-ai/FlagCX.git", "relative_path": "tle/third_party/flagcx"}, )) diff --git a/python/triton/_common_ir.py b/python/triton/_common_ir.py new file mode 100644 index 0000000000..2d8f7645e6 --- /dev/null +++ b/python/triton/_common_ir.py @@ -0,0 +1,3 @@ +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 c8a6ec370b..2e466f2713 100644 --- a/python/triton/experimental/tle/language/gpu/core.py +++ b/python/triton/experimental/tle/language/gpu/core.py @@ -1,11 +1,10 @@ # flagtree tle -import os 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 ( @@ -18,22 +17,6 @@ SHARED_MEMORY_ADDRESS_SPACE = 3 -def _tileir_mode() -> bool: - return os.environ.get("TLE_GPU_TILEIR_MODE") == "1" - - -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 mode") - - -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. @@ -264,22 +247,9 @@ 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 _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 tle_semantic.COMMON_IR_ENABLED and storage != tle.smem: + raise ValueError("GPU CommonIR currently supports only shared-memory buffers") if layout is None: if storage == tle.smem: @@ -320,25 +290,25 @@ def alloc( # Use provided layout layout_handle = layout.to_ir(_semantic.builder) - 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 @@ -391,104 +361,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) @@ -502,23 +374,12 @@ 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: - 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)") + 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 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 @@ -526,22 +387,23 @@ 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 - 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") + elif tle_semantic.COMMON_IR_ENABLED and tl._unwrap_if_constexpr(offsets) is not None: + raise ValueError("GPU CommonIR normal copy does not support offsets; offset the global pointers instead") + + tle_semantic.copy(src, dst, shape, offsets if descriptor is not None else None, direction, _semantic) @tl.builtin @@ -557,50 +419,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 _tileir_mode(): - if not hasattr(_semantic.builder, "create_tile_to_tensor"): - raise RuntimeError("TLE GPU TileIR mode requires a TileIR-enabled TLE builder") - 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 _tileir_mode(): - if not hasattr(_semantic.builder, "create_tile_store_tensor"): - raise RuntimeError("TLE GPU TileIR mode requires a TileIR-enabled TLE builder") - _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) - _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 @@ -691,26 +517,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: @@ -728,8 +534,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] - 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..6c8fc4c52c 100644 --- a/python/triton/experimental/tle/language/gpu/semantic.py +++ b/python/triton/experimental/tle/language/gpu/semantic.py @@ -10,11 +10,160 @@ 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: + buffer = dst if isinstance(dst, tle.buffered_tensor) else src + pointers = src if buffer is dst else dst + requested_shape = tuple(int(tl._unwrap_if_constexpr(dim)) for dim in shape) + buffer_shape = tuple(buffer.type.shape) + if requested_shape != buffer_shape or tuple(pointers.shape) != buffer_shape: + raise ValueError("GPU CommonIR copy requires shape and pointer tensor shape to match the full buffer") + if not pointers.dtype.is_ptr() or pointers.dtype.element_ty != buffer.dtype: + raise ValueError("GPU CommonIR copy requires pointers with the buffer element type") + _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: + if tuple(result_type.shape) != tuple(buffer.type.shape): + raise ValueError("GPU CommonIR to_tensor requires the full buffer shape") + 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: + if tuple(value.shape) != tuple(dst.type.shape) or value.dtype != dst.dtype: + raise ValueError("GPU CommonIR store_tensor requires the full buffer shape and element type") + _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) + memdesc = get_memdesc(buffer, _semantic) + return _semantic.builder.create_local_pointers(result_ir, memdesc, *handles) + + +def get_memdesc(buffer, _semantic): + """Bridge a CommonIR 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 9dd3c5cc6d..426cb106e8 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,18 +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 _tileir_mode(): - if not hasattr(_semantic.builder, "create_tile_subview"): - raise RuntimeError("TLE GPU TileIR mode requires a TileIR-enabled TLE builder") - 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)) - 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) @@ -442,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 _tileir_mode() and 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), @@ -661,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()) @@ -788,10 +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) - 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, + _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) @@ -799,21 +785,16 @@ 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()) + _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 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, + _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()) @@ -830,24 +811,17 @@ 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 "", - self._reader_field_names()) + 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()) slot = self.pipe._make_slot(stage, _semantic=_semantic, field_names=self.field_names) return pipe_wait_result(slot, tl.tensor(is_closed, tl.int1)) @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, + _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 0f1c823362..bfd9115d14 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,10 +103,10 @@ 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) - return gpu_types.pipe_value(capacity, scope, name, fields, reader_names, one_shot=one_shot) + 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 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-async-copy.mlir b/test/CommonIR/gpu-tileir-async-copy.mlir new file mode 100644 index 0000000000..588c9750c9 --- /dev/null +++ b/test/CommonIR/gpu-tileir-async-copy.mlir @@ -0,0 +1,32 @@ +// REQUIRES: flagtree-common-ir +// RUN: triton-opt %s --convert-common-ir-to-ttgir='enable-async-copy=true' | FileCheck %s --check-prefix=ASYNC --implicit-check-not=tile. --implicit-check-not=tle. --implicit-check-not=builtin.unrealized_conversion_cast +// RUN: triton-opt %s --convert-common-ir-to-ttgir='enable-async-copy=true' --convert-triton-to-tritongpu='target=cuda:90 num-warps=4' | FileCheck %s --check-prefix=LAYOUT --implicit-check-not=builtin.unrealized_conversion_cast +// RUN: triton-opt %s --convert-common-ir-to-ttgir='enable-async-copy=false' | FileCheck %s --check-prefix=SYNC --implicit-check-not=ttg.async + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> + +module { + tt.func public @copy(%src: !tt.ptr, %dst: !tt.ptr) { + %index = tt.make_range {start = 0 : i32, end = 128 : i32} : tensor<128xi32> + %s = tt.splat %src : !tt.ptr -> tensor<128x!tt.ptr> + %d = tt.splat %dst : !tt.ptr -> tensor<128x!tt.ptr> + %sp = tt.addptr %s, %index : tensor<128x!tt.ptr>, tensor<128xi32> + %dp = tt.addptr %d, %index : tensor<128x!tt.ptr>, tensor<128xi32> + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[128], f32, shared> + tile.copy %sp -> %buf : tensor<128x!tt.ptr>, !tile.buf<[128], f32, shared> + tile.copy %buf -> %dp : !tile.buf<[128], f32, shared>, tensor<128x!tt.ptr> + tt.return + } +} + +// ASYNC: %[[BUF:.*]] = ttg.local_alloc +// ASYNC: %[[COPY:.*]] = ttg.async_copy_global_to_local %{{.*}}, %[[BUF]] +// ASYNC: %[[COMMIT:.*]] = ttg.async_commit_group tokens %[[COPY]] +// ASYNC: ttg.async_wait %[[COMMIT]] {num = 0 : i32} +// ASYNC: %[[VALUE:.*]] = ttg.local_load %[[BUF]] +// ASYNC: tt.store %{{.*}}, %[[VALUE]] +// LAYOUT: ttg.async_copy_global_to_local {{.*}} : tensor<128x!tt.ptr, #blocked> +// SYNC: tt.load +// SYNC: ttg.local_store +// SYNC: ttg.local_load +// SYNC: tt.store diff --git a/test/CommonIR/gpu-tileir-buffer-users.mlir b/test/CommonIR/gpu-tileir-buffer-users.mlir new file mode 100644 index 0000000000..3846fd7577 --- /dev/null +++ b/test/CommonIR/gpu-tileir-buffer-users.mlir @@ -0,0 +1,36 @@ +// REQUIRES: flagtree-common-ir +// RUN: triton-opt %s --convert-common-ir-to-ttgir | FileCheck %s --implicit-check-not=tile. --implicit-check-not=builtin.unrealized_conversion_cast + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory + +module { + // CHECK-LABEL: tt.func private @read + // CHECK-SAME: %[[ARG:.*]]: !ttg.memdesc<16xf32, #shared, #smem, mutable> + // CHECK: ttg.local_load %[[ARG]] + tt.func private @read(%buf: !tile.buf<[16], f32, shared>) -> tensor<16xf32> { + %desc = builtin.unrealized_conversion_cast %buf : !tile.buf<[16], f32, shared> to !ttg.memdesc<16xf32, #shared, #smem, mutable> + %value = ttg.local_load %desc : !ttg.memdesc<16xf32, #shared, #smem, mutable> -> tensor<16xf32> + tt.return %value : tensor<16xf32> + } + + // CHECK-LABEL: tt.func public @call_and_capture + // CHECK: %[[BUF:.*]] = ttg.local_alloc + // CHECK: tt.call @read(%[[BUF]]) : (!ttg.memdesc<16xf32, #shared, #smem, mutable>) + // CHECK: ttg.warp_specialize(%[[BUF]]) + // CHECK: partition0(%[[PART:.*]]: !ttg.memdesc<16xf32, #shared, #smem, mutable>) + // CHECK: ttg.local_load %[[PART]] + tt.func public @call_and_capture() { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + %value = tt.call @read(%buf) : (!tile.buf<[16], f32, shared>) -> tensor<16xf32> + ttg.warp_specialize(%buf) + default { + ttg.warp_yield + } + partition0(%arg: !tile.buf<[16], f32, shared>) num_warps(4) { + %tensor = tile.to_tensor %arg : <[16], f32, shared> -> tensor<16xf32> + ttg.warp_return + } : (!tile.buf<[16], f32, shared>) -> () + tt.return + } +} diff --git a/test/CommonIR/gpu-tileir-copy-groups.mlir b/test/CommonIR/gpu-tileir-copy-groups.mlir new file mode 100644 index 0000000000..d7ac1a7614 --- /dev/null +++ b/test/CommonIR/gpu-tileir-copy-groups.mlir @@ -0,0 +1,49 @@ +// REQUIRES: flagtree-common-ir +// RUN: triton-opt %s --convert-common-ir-to-ttgir='enable-async-copy=true' | FileCheck %s + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +module { + // CHECK-LABEL: @independent_copies + // CHECK: %[[A:.*]] = ttg.async_copy_global_to_local + // CHECK-NOT: ttg.async_wait + // CHECK: %[[B:.*]] = ttg.async_copy_global_to_local + // CHECK: %[[COMMIT:.*]] = ttg.async_commit_group tokens %[[A]], %[[B]] + // CHECK: ttg.async_wait %[[COMMIT]] + // CHECK-NOT: ttg.async_wait + tt.func @independent_copies(%src: tensor<16x!tt.ptr>) { + %a = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + %b = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + tile.copy %src -> %a : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + %c0 = arith.constant dense<0> : tensor<16xi32> + %ptr = tt.addptr %src, %c0 : tensor<16x!tt.ptr>, tensor<16xi32> + tile.copy %ptr -> %b : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + tt.return + } + + // CHECK-LABEL: @same_buffer + // CHECK: ttg.async_copy_global_to_local + // CHECK: ttg.async_wait + // CHECK: ttg.async_copy_global_to_local + // CHECK: ttg.async_wait + tt.func @same_buffer(%src: tensor<16x!tt.ptr>) { + %a = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + tile.copy %src -> %a : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + tile.copy %src -> %a : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + tt.return + } + + // CHECK-LABEL: @read_between_copies + // CHECK: ttg.async_copy_global_to_local + // CHECK: ttg.async_wait + // CHECK: ttg.local_load + // CHECK: ttg.async_copy_global_to_local + // CHECK: ttg.async_wait + tt.func @read_between_copies(%src: tensor<16x!tt.ptr>) -> tensor<16xf32> { + %a = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + %b = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + tile.copy %src -> %a : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + %value = tile.to_tensor %a : <[16], f32, shared> -> tensor<16xf32> + tile.copy %src -> %b : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + tt.return %value : tensor<16xf32> + } +} diff --git a/test/CommonIR/gpu-tileir-invalid.mlir b/test/CommonIR/gpu-tileir-invalid.mlir new file mode 100644 index 0000000000..bdc48e6d06 --- /dev/null +++ b/test/CommonIR/gpu-tileir-invalid.mlir @@ -0,0 +1,227 @@ +// REQUIRES: flagtree-common-ir +// RUN: triton-opt %s --split-input-file --convert-common-ir-to-ttgir --verify-diagnostics + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @partial_copy(%src: tensor<8x!tt.ptr>) { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + // expected-error@+1 {{requires a full-buffer tensor with matching shape and element type}} + tile.copy %src -> %buf : tensor<8x!tt.ptr>, !tile.buf<[16], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @copy_dtype(%src: tensor<16x!tt.ptr>) { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + // expected-error@+1 {{requires a full-buffer tensor with matching shape and element type}} + tile.copy %src -> %buf : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @copy_nonglobal(%src: tensor<16x!tt.ptr>) { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + // expected-error@+1 {{GPU tile.copy requires a tensor of global pointers}} + tile.copy %src -> %buf : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @copy_scalar(%src: !tt.ptr) { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + // expected-error@+1 {{GPU tile.copy requires a tensor of global pointers}} + tile.copy %src -> %buf : !tt.ptr, !tile.buf<[16], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @copy_transpose(%src: tensor<16x!tt.ptr>) { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + // expected-error@+1 {{GPU tile.copy supports only plain ND copies}} + tile.copy %src -> %buf {transpose} : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @copy_nz(%src: tensor<16x!tt.ptr>) { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + // expected-error@+1 {{GPU tile.copy supports only plain ND copies}} + tile.copy %src -> %buf {src_layout = 1 : i64} : tensor<16x!tt.ptr>, !tile.buf<[16], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @stride() { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + %c0 = arith.constant 0 : index + %c16 = arith.constant 16 : index + // expected-error@+1 {{GPU tile.subview requires unit strides}} + %view = tile.subview %buf[%c0] [[8]] [[2]] {tle.gpu_layout = #shared} : <[16], f32, shared> -> <[8], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @sizes() { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + %c0 = arith.constant 0 : index + %c16 = arith.constant 16 : index + // expected-error@+1 {{tile.subview sizes must match its result shape}} + %view = tile.subview %buf[%c0] [[4]] [[1]] {tle.gpu_layout = #shared} : <[16], f32, shared> -> <[8], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @offset_bounds() { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + %c0 = arith.constant 0 : index + %c16 = arith.constant 16 : index + // expected-error@+1 {{tile.subview offset is outside the source buffer}} + %view = tile.subview %buf[%c16] [[8]] [[1]] {tle.gpu_layout = #shared} : <[16], f32, shared> -> <[8], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @dynamic_same_rank(%idx: index) { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + %c0 = arith.constant 0 : index + %c16 = arith.constant 16 : index + // expected-error@+1 {{same-rank tile.subview currently requires static offsets}} + %view = tile.subview %buf[%idx] [[8]] [[1]] {tle.gpu_layout = #shared} : <[16], f32, shared> -> <[8], f32, shared> + tt.return + } +} + +// ----- + +#shared2 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}> +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @trailing_offset() { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared2} : <[2, 16], f32, shared> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + // expected-error@+1 {{rank-reducing tile.subview requires zero trailing offsets}} + %view = tile.subview %buf[%c0, %c1] [[16]] [[1]] {tle.gpu_layout = #shared} : <[2, 16], f32, shared> -> <[16], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @bridge_mismatch() { + %buf = tile.alloc {space = 7 : i64, tle.gpu_layout = #shared} : <[16], f32, shared> + // expected-error@+1 {{CommonIR buffer bridge descriptor type does not match the converted buffer}} + %desc = builtin.unrealized_conversion_cast %buf : !tile.buf<[16], f32, shared> to !ttg.memdesc<8xf32, #shared, #smem, mutable> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @unsupported_storage() { + // expected-error@+1 {{GPU TileIR conversion currently requires #tile.shared}} + %buf = tile.alloc {space = 8 : i64, tle.gpu_layout = #shared} : <[16], f32, local> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @missing_layout() { + // expected-error@+1 {{is missing the preserved tle.gpu_layout attribute}} + %buf = tile.alloc {space = 7 : i64} : <[16], f32, shared> + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @unsupported_set_flag() { + // expected-error@+1 {{GPU CommonIR conversion does not support tile engine/event synchronization}} + "tile.set_flag"() {producer = 1 : i64, consumer = 6 : i64, event = 0 : i64} : () -> () + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @unsupported_wait_flag() { + // expected-error@+1 {{GPU CommonIR conversion does not support tile engine/event synchronization}} + "tile.wait_flag"() {producer = 1 : i64, consumer = 6 : i64, event = 0 : i64} : () -> () + tt.return + } +} + +// ----- + +#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}> +#smem = #ttg.shared_memory +module { + tt.func public @unsupported_pipe_barrier() { + // expected-error@+1 {{GPU CommonIR conversion does not support tile engine/event synchronization}} + "tile.pipe_barrier"() {pipe = 1 : i64} : () -> () + tt.return + } +} diff --git a/test/CommonIR/gpu-tileir-to-ttgir.mlir b/test/CommonIR/gpu-tileir-to-ttgir.mlir new file mode 100644 index 0000000000..92f254fcd5 --- /dev/null +++ b/test/CommonIR/gpu-tileir-to-ttgir.mlir @@ -0,0 +1,68 @@ +// REQUIRES: flagtree-common-ir +// RUN: triton-opt %s --convert-common-ir-to-ttgir | FileCheck %s --implicit-check-not=tile. --implicit-check-not=builtin.unrealized_conversion_cast + +#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 = builtin.unrealized_conversion_cast %buf + : !tile.buf<[2, 16], f32, shared> to !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> + %slot_desc = builtin.unrealized_conversion_cast %slot + : !tile.buf<[16], f32, shared> to !ttg.memdesc<16xf32, #shared1, #smem, mutable> + %ptr = "tle.local_pointers"(%slot_desc) + : (!ttg.memdesc<16xf32, #shared1, #smem, mutable>) -> 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: ttg.local_store %{{.*}}, %[[SLOT]] +// CHECK: "tle.local_pointers"(%[[SLOT]]) +// CHECK-NOT: "tle.local_pointers" + +// CHECK-LABEL: tt.func public @copy_to_global +// CHECK: %[[LOCAL:.*]] = ttg.local_alloc +// CHECK: ttg.local_store %{{.*}}, %[[LOCAL]] +// CHECK: %[[VALUE:.*]] = ttg.local_load %[[LOCAL]] +// CHECK: tt.store %{{.*}}, %[[VALUE]] +// CHECK-NOT: "tle.local_pointers" +module { + tt.func public @copy_to_global(%dst: tensor<16x!tt.ptr>, %value: tensor<16xf32>) { + %buf = tile.alloc {layout = 0 : i64, space = 7 : i64, tle.gpu_layout = #shared1} + : <[16], f32, shared> + tile.store_tensor %value -> %buf : tensor<16xf32>, !tile.buf<[16], f32, shared> + tile.copy %buf -> %dst {src_layout = 0 : i64} + : !tile.buf<[16], f32, shared>, tensor<16x!tt.ptr> + tt.return + } +} + +// CHECK-LABEL: tt.func public @static_subview +// CHECK: %[[ALLOC:.*]] = ttg.local_alloc +// CHECK: %[[VIEW:.*]] = ttg.memdesc_subslice %[[ALLOC]][8] +// CHECK: ttg.local_load %[[VIEW]] +module { + tt.func public @static_subview() -> tensor<8xf32> { + %buf = tile.alloc {layout = 0 : i64, space = 7 : i64, tle.gpu_layout = #shared1} + : <[16], f32, shared> + %c8 = arith.constant 8 : index + %view = tile.subview %buf[%c8] [[8]] [[1]] {tle.gpu_layout = #shared1} + : <[16], f32, shared> -> <[8], f32, shared> + %value = tile.to_tensor %view : <[8], f32, shared> -> tensor<8xf32> + tt.return %value : tensor<8xf32> + } +} diff --git a/test/CommonIR/native_flash_gpu.py b/test/CommonIR/native_flash_gpu.py index 0f5b20072e..27abe27b7f 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, @@ -20,12 +16,9 @@ def _flash_gpu_kernel(q_ptr, k_ptr, v_ptr, o_ptr, stride_qm, stride_qd, stride_k offs_n = tl.arange(0, BLOCK_N) offs_d = tl.arange(0, D) - q_buf = tle.gpu.alloc([BLOCK_M, D], dtype=tl.float16, layout=None, scope=tle.gpu.smem, - nv_mma_shared_layout=False) - k_buf = tle.gpu.alloc([BLOCK_N, D], dtype=tl.float16, layout=None, scope=tle.gpu.smem, - nv_mma_shared_layout=False) - v_buf = tle.gpu.alloc([BLOCK_N, D], dtype=tl.float16, layout=None, scope=tle.gpu.smem, - nv_mma_shared_layout=False) + q_buf = tle.gpu.alloc([BLOCK_M, D], dtype=tl.float16, layout=None, scope=tle.gpu.smem, nv_mma_shared_layout=False) + k_buf = tle.gpu.alloc([BLOCK_N, D], dtype=tl.float16, layout=None, scope=tle.gpu.smem, nv_mma_shared_layout=False) + v_buf = tle.gpu.alloc([BLOCK_N, D], dtype=tl.float16, layout=None, scope=tle.gpu.smem, nv_mma_shared_layout=False) q_tile = q_ptr + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd tle.gpu.copy(q_tile, q_buf, [BLOCK_M, D]) @@ -52,8 +45,7 @@ def _flash_gpu_kernel(q_ptr, k_ptr, v_ptr, o_ptr, stride_qm, stride_qd, stride_k m_i = m_new out = acc / l_i[:, None] - out_buf = tle.gpu.alloc([BLOCK_M, D], dtype=tl.float32, layout=None, scope=tle.gpu.smem, - nv_mma_shared_layout=False) + out_buf = tle.gpu.alloc([BLOCK_M, D], dtype=tl.float32, layout=None, scope=tle.gpu.smem, nv_mma_shared_layout=False) tle.gpu.store_tensor(out, out_buf) out_vals = tle.gpu.to_tensor(out_buf, writable=False) o_tile = o_ptr + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od @@ -123,7 +115,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, @@ -150,10 +142,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) + nvidia.passes.commonir.add_to_ttgir(pm, _compile_options()[0].arch >= 80) + 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}") @@ -166,9 +164,8 @@ def run_check(n_ctx=32, d=32): out = torch.empty((n_ctx, d), device="cuda", dtype=torch.float32) sm_scale = d**-0.5 grid = (triton.cdiv(n_ctx, 16), ) - _flash_gpu_kernel[grid](q, k, v, out, q.stride(0), q.stride(1), k.stride(0), k.stride(1), v.stride(0), - v.stride(1), out.stride(0), out.stride(1), N_CTX=n_ctx, D=d, sm_scale=sm_scale, - BLOCK_M=16, BLOCK_N=16) + _flash_gpu_kernel[grid](q, k, v, out, q.stride(0), q.stride(1), k.stride(0), k.stride(1), v.stride(0), v.stride(1), + out.stride(0), out.stride(1), N_CTX=n_ctx, D=d, sm_scale=sm_scale, BLOCK_M=16, BLOCK_N=16) ref = torch.softmax(torch.matmul(q.float(), k.float().T) * sm_scale, dim=-1).matmul(v.float()) torch.testing.assert_close(out, ref, atol=5e-2, rtol=5e-2) print("[check] native_flash_gpu precision PASS") diff --git a/test/CommonIR/native_matmul_gpu.py b/test/CommonIR/native_matmul_gpu.py index ffdb380f00..b2d622eeb9 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): @@ -105,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, @@ -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) + nvidia.passes.commonir.add_to_ttgir(pm, _compile_options()[0].arch >= 80) + 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/test_gpu_semantics.py b/test/CommonIR/test_gpu_semantics.py new file mode 100644 index 0000000000..adc09ecde5 --- /dev/null +++ b/test/CommonIR/test_gpu_semantics.py @@ -0,0 +1,67 @@ +"""Full-buffer GPU copy contract and frontend diagnostics for CommonIR.""" +import pytest +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +from triton._common_ir import ENABLED + +pytestmark = pytest.mark.skipif(not ENABLED, reason="requires FLAGTREE_COMMON_IR build") + + +@triton.jit +def _roundtrip(src, dst, BLOCK: tl.constexpr): + idx = tl.arange(0, BLOCK) + buf = tle.gpu.alloc([BLOCK], tl.float32, scope=tle.gpu.smem, nv_mma_shared_layout=False) + tle.gpu.copy(src + idx, buf, [BLOCK]) + value = tle.gpu.to_tensor(buf, writable=False) + tle.gpu.store_tensor(value + 1, buf) + tle.gpu.copy(buf, dst + idx, [BLOCK]) + + +@pytest.mark.parametrize("block", [16, 128, 1024]) +def test_copy_roundtrip(block): + src = torch.arange(block, dtype=torch.float32, device="cuda") + dst = torch.empty_like(src) + kernel = _roundtrip[(1, )](src, dst, block) + torch.testing.assert_close(dst, src + 1, rtol=0, atol=0) + assert "tile." not in kernel.asm["ttgir"] + assert "unrealized_conversion_cast" not in kernel.asm["ttgir"] + + +@triton.jit +def _invalid(src, CASE: tl.constexpr): + idx = tl.arange(0, 16) + buf = tle.gpu.alloc([16], tl.float32, scope=tle.gpu.smem, nv_mma_shared_layout=False) + if CASE == 0: + tle.gpu.copy(src + idx, buf, [8]) + elif CASE == 1: + tle.gpu.copy(src + tl.arange(0, 8), buf, [16]) + elif CASE == 2: + tle.gpu.copy(src + idx, buf, [16], offsets=[1]) + elif CASE == 3: + tle.gpu.to_tensor(buf, target_shape=[8]) + elif CASE == 4: + tle.gpu.store_tensor(tl.full([8], 0, tl.float32), buf) + elif CASE == 5: + tle.gpu.store_tensor(tl.full([16], 0, tl.int32), buf) + elif CASE == 6: + tle.gpu.copy((src + idx).to(tl.pointer_type(tl.int32)), buf, [16]) + else: + tle.gpu.alloc([32, 32], tl.float32, scope=tle.gpu.tmem) + + +@pytest.mark.parametrize("case, message", [ + (0, "shape and pointer tensor shape"), + (1, "shape and pointer tensor shape"), + (2, "normal copy does not support offsets"), + (3, "to_tensor requires the full buffer shape"), + (4, "store_tensor requires the full buffer shape and element type"), + (5, "store_tensor requires the full buffer shape and element type"), + (6, "copy requires pointers with the buffer element type"), + (7, "currently supports only shared-memory buffers"), +]) +def test_invalid_buffer_semantics(case, message): + src = torch.empty(16, dtype=torch.float32, device="cuda") + with pytest.raises(triton.CompilationError, match=message): + _invalid[(1, )](src, case) diff --git a/test/CommonIR/tle_gpu_primitives.py b/test/CommonIR/tle_gpu_primitives.py index c6ba0b42b5..c32d94b639 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,14 +6,10 @@ 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) - storage = tle.gpu.alloc([2, BLOCK], dtype=tl.float32, layout=None, scope=tle.gpu.smem, - nv_mma_shared_layout=False) + storage = tle.gpu.alloc([2, BLOCK], dtype=tl.float32, layout=None, scope=tle.gpu.smem, nv_mma_shared_layout=False) pipe = tle.pipe(capacity=2, scope="cta", name="p", payload=storage) writer = pipe.writer() reader = pipe.reader() @@ -24,7 +19,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 +64,14 @@ def dump_tileir(path): "tile.alloc", "tile.copy", "tile.subview", - "tile.to_tensor", - "tile.set_flag", - "tile.wait_flag", - "tile.pipe_barrier", + "builtin.unrealized_conversion_cast", + "tle.local_pointers", + "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}") @@ -85,18 +79,22 @@ 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") - tle_ir.lower_gpu_tileir_to_ttir(module) + pm = ir.pass_manager(module.context) + nvidia.passes.commonir.add_to_ttgir(pm, _compile_options()[0].arch >= 80) + 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", "builtin.unrealized_conversion_cast"): + 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/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 144fac81d1..2cbb366610 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 @@ -264,11 +265,13 @@ 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() + # Lower CommonIR buffers before TTIR canonicalization. The frozen + # dialect does not model every GPU memory effect (e.g. to_tensor reads), + # so CSE must see the effect-aware TTGIR operations instead. + if COMMON_IR_ENABLED: + nvidia.passes.commonir.add_to_ttgir(pm, capability >= 80) passes.ttir.add_rewrite_tensor_pointer(pm) if capability // 10 < 9: passes.ttir.add_rewrite_tensor_descriptor_to_pointer(pm) diff --git a/third_party/nvidia/include/CMakeLists.txt b/third_party/nvidia/include/CMakeLists.txt index 2ef7aab106..dbae0d3cf8 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 CommonIRIR) + 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..121aaafd1b --- /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 options = [Option<"enableAsyncCopy", "enable-async-copy", "bool", "false", + "Use native async global-to-shared copies on SM80+">]; + let description = [{ + Converts the target-independent CommonIR produced by tle.gpu primitives into + the NVIDIA TritonGPU descriptor form. In particular, tile.alloc becomes + ttg.local_alloc with a ttg.memdesc result. Temporary CommonIR buffer bridges + are resolved to that descriptor before existing TLE operations are lowered. + }]; + let dependentDialects = ["mlir::arith::ArithDialect", + "mlir::triton::TritonDialect", + "mlir::triton::gpu::TritonGPUDialect", + "mlir::triton::tile::CommonIRDialect"]; +} + +#endif diff --git a/third_party/nvidia/lib/CMakeLists.txt b/third_party/nvidia/lib/CMakeLists.txt index 2ef7aab106..dbae0d3cf8 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 CommonIRIR) + 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..97ca313fc9 --- /dev/null +++ b/third_party/nvidia/lib/CommonIRToTTGIR/CMakeLists.txt @@ -0,0 +1,12 @@ +add_triton_library(CommonIRToTTGIR + CommonIRToTTGIR.cpp + + DEPENDS + CommonIRToTTGIRPassIncGen + + LINK_LIBS PUBLIC + MLIRArithDialect + CommonIRIR + TritonIR + TritonGPUIR +) diff --git a/third_party/nvidia/lib/CommonIRToTTGIR/CommonIRToTTGIR.cpp b/third_party/nvidia/lib/CommonIRToTTGIR/CommonIRToTTGIR.cpp new file mode 100644 index 0000000000..211ff2deae --- /dev/null +++ b/third_party/nvidia/lib/CommonIRToTTGIR/CommonIRToTTGIR.cpp @@ -0,0 +1,654 @@ +// MIT License +// +// Copyright (c) 2025 The FlagOS Contributors + +#include "nvidia/include/CommonIRToTTGIR/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" + +#include "mlir-ext/Dialect/CommonIR/IR/CommonIRDialect.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallPtrSet.h" +#include + +namespace mlir::triton { + +#define GEN_PASS_DEF_COMMONIRTOTTGIR +#include "nvidia/include/CommonIRToTTGIR/Passes.h.inc" + +namespace { + +namespace tile = mlir::triton::tile; +namespace ttg = mlir::triton::gpu; + +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(); + } + if (!isa(layout)) { + op->emitError("requires a preserved TritonGPU shared-memory layout"); + 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 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.getInputs().size() != 1 || + bridge.getInputs().front() != argument || bridge->getNumResults() != 1) + continue; + auto type = dyn_cast(bridge.getResult(0).getType()); + if (!type) + continue; + 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; + auto tensorType = dyn_cast(type); + return tensorType && isa(tensorType.getElementType()); +} + +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 integer.getInt(); +} + +static LogicalResult checkTensorType(Operation *op, Type type, + ttg::MemDescType bufferType) { + auto tensor = dyn_cast(type); + if (!tensor || tensor.getShape() != bufferType.getShape() || + tensor.getElementType() != bufferType.getElementType()) + return op->emitError( + "requires a full-buffer tensor with matching shape and element type"); + return success(); +} + +static LogicalResult checkCopy(tile::CopyOp op, Value pointers, + ttg::MemDescType bufferType) { + // The frozen CommonIR interface carries more copy modes than this backend + // supports. Reject their semantics rather than silently doing a plain copy. + if (op->hasAttr("engine") || op->hasAttr("transpose") || + op->hasAttr("dst_nz_layout") || + (op.getSrcLayout() && *op.getSrcLayout() != tile::Layout::ND)) + return op.emitError("GPU tile.copy supports only plain ND copies without " + "engine, transpose or NZ layout attributes"); + auto tensor = dyn_cast(pointers.getType()); + auto pointer = tensor ? dyn_cast(tensor.getElementType()) + : triton::PointerType(); + if (!pointer || pointer.getAddressSpace() != 1) + return op.emitError("GPU tile.copy requires a tensor of global pointers"); + return checkTensorType(op, getPointeeValueType(pointers.getType()), + bufferType); +} + +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 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 LogicalResult +convertBufferBridges(ModuleOp module, llvm::DenseMap &mapped, + ArrayRef argumentCasts, + SmallVectorImpl &eraseOps) { + WalkResult result = module.walk([&](UnrealizedConversionCastOp op) { + if (llvm::is_contained(argumentCasts, op) || op.getInputs().size() != 1 || + op->getNumResults() != 1 || + !isa(op.getInputs().front().getType()) || + !isa(op.getResult(0).getType())) + return WalkResult::advance(); + + auto source = lookupBuffer(op, op.getInputs().front(), mapped); + if (failed(source)) + return WalkResult::interrupt(); + if ((*source).getType() != op.getResult(0).getType()) { + op.emitError("CommonIR buffer bridge descriptor type does not match the " + "converted buffer"); + return WalkResult::interrupt(); + } + mapped[op.getResult(0)] = *source; + eraseOps.push_back(op); + return WalkResult::advance(); + }); + return failure(result.wasInterrupted()); +} + +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) { + 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()); + if (op.getOffsets().size() != sourceType.getRank() || + op.getSizes().size() != resultType->getRank() || + op.getStrides().size() != resultType->getRank()) { + op.emitError("GPU tile.subview requires source-rank offsets and " + "result-rank sizes and strides"); + return WalkResult::interrupt(); + } + for (auto [size, dim] : llvm::zip(op.getSizes(), resultType->getShape())) { + if (cast(size).getInt() != dim) { + op.emitError("tile.subview sizes must match its result shape"); + return WalkResult::interrupt(); + } + } + if (llvm::any_of(op.getStrides(), [](Attribute stride) { + return cast(stride).getInt() != 1; + })) { + op.emitError("GPU tile.subview requires unit strides"); + return WalkResult::interrupt(); + } + Value converted; + if (sourceType.getRank() == resultType->getRank() + 1) { + if (sourceType.getShape().drop_front() != resultType->getShape()) { + op.emitError("rank-reducing tile.subview must preserve trailing shape"); + return WalkResult::interrupt(); + } + for (Value offset : op.getOffsets().drop_front()) { + auto constant = getConstantOffset(offset); + if (failed(constant) || *constant != 0) { + op.emitError( + "rank-reducing tile.subview requires zero trailing offsets"); + return WalkResult::interrupt(); + } + } + auto leading = getConstantOffset(op.getOffsets().front()); + if (succeeded(leading) && + (*leading < 0 || *leading >= sourceType.getShape().front())) { + op.emitError("tile.subview offset is outside the source buffer"); + 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 (auto [axis, offset] : llvm::enumerate(op.getOffsets())) { + auto constant = getConstantOffset(offset); + if (failed(constant)) { + op.emitError( + "same-rank tile.subview currently requires static offsets"); + return WalkResult::interrupt(); + } + if (*constant < 0 || *constant > std::numeric_limits::max() || + *constant > + sourceType.getShape()[axis] - resultType->getShape()[axis]) { + op.emitError("tile.subview offset is outside the source buffer"); + return WalkResult::interrupt(); + } + offsets.push_back(static_cast(*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 enableAsyncCopy, + SmallVectorImpl &copies) { + 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.getSrc(), mapped); + if (failed(source)) { + failedConversion = true; + return; + } + if (failed( + checkTensorType(op, op.getResult().getType(), + cast((*source).getType())))) { + 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; + } + if (failed(checkTensorType( + op, op.getSrc().getType(), + cast((*destination).getType())))) { + 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 (isa(dst.getType())) { + if (failed(checkCopy(op, src, cast(dst.getType())))) { + failedConversion = true; + return; + } + if (enableAsyncCopy) { + auto copy = builder.create( + op.getLoc(), src, dst, Value(), Value(), + triton::CacheModifier::NONE, triton::EvictionPolicy::NORMAL, + false); + auto commit = builder.create( + op.getLoc(), ValueRange{copy.getToken()}); + copies.push_back(copy); + // tile.copy is synchronous to its users. Keep the completion token + // explicit so later GPU passes can reason about shared-memory reads. + builder.create(op.getLoc(), commit.getResult(), 0); + } else { + Value value = builder.create( + op.getLoc(), src, triton::CacheModifier::NONE, + triton::EvictionPolicy::NORMAL, false); + builder.create(op.getLoc(), value, dst); + } + } else if (isa(src.getType()) && + isPointerLike(dst.getType())) { + if (failed(checkCopy(op, dst, cast(src.getType())))) { + failedConversion = true; + return; + } + Type valueType = getPointeeValueType(dst.getType()); + if (!valueType) { + op.emitError("cannot determine tile.copy destination value type"); + failedConversion = true; + return; + } + Value value = + builder.create(op.getLoc(), valueType, src); + 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)) { + operation->emitError("GPU CommonIR conversion does not support tile " + "engine/event synchronization; use the existing " + "GPU synchronization path"); + failedConversion = true; + } + }); + + return failure(failedConversion); +} + +static Value getSharedAllocation(Value value) { + while (value) { + if (value.getDefiningOp()) + return value; + if (auto index = value.getDefiningOp()) + value = index.getSrc(); + else if (auto slice = value.getDefiningOp()) + value = slice.getSrc(); + else + return {}; + } + return {}; +} + +static void combineCopyWaits(ArrayRef copies) { + llvm::SmallPtrSet generated; + for (auto copy : copies) + generated.insert(copy); + for (auto copy : copies) { + auto commit = + dyn_cast_or_null(copy->getNextNode()); + if (!commit || commit.getInputTokens().size() != 1) + continue; + auto wait = dyn_cast_or_null(commit->getNextNode()); + Value allocation = getSharedAllocation(copy.getResult()); + if (!wait || !wait.getResult().use_empty() || !allocation) + continue; + SmallVector allocations{allocation}; + while (true) { + Operation *next = wait->getNextNode(); + // Only cross pure, region-free address/value computations. A read, + // write, barrier or control-flow boundary must observe the earlier copy. + while (next && !next->getNumRegions() && isMemoryEffectFree(next)) + next = next->getNextNode(); + if (!next || !generated.contains(next)) + break; + auto nextCopy = cast(next); + Value nextAllocation = getSharedAllocation(nextCopy.getResult()); + if (!nextAllocation || llvm::is_contained(allocations, nextAllocation)) + break; + auto nextCommit = + dyn_cast_or_null(next->getNextNode()); + auto nextWait = + nextCommit + ? dyn_cast_or_null(nextCommit->getNextNode()) + : ttg::AsyncWaitOp(); + if (!nextWait || !nextWait.getResult().use_empty()) + break; + SmallVector tokens(commit.getInputTokens()); + llvm::append_range(tokens, nextCommit.getInputTokens()); + nextCommit.getInputTokensMutable().assign(tokens); + wait.erase(); + commit.erase(); + commit = nextCommit; + wait = nextWait; + allocations.push_back(nextAllocation); + } + } +} + +class CommonIRToTTGIRPass + : public impl::CommonIRToTTGIRBase { +public: + using CommonIRToTTGIRBase::CommonIRToTTGIRBase; + + void runOnOperation() override { + ModuleOp module = getOperation(); + llvm::DenseMap mapped; + llvm::DenseMap convertedArgumentTypes; + SmallVector eraseOps; + SmallVector copies; + 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(convertBufferBridges(module, mapped, argumentCasts, eraseOps)) || + failed(convertTileUsers(module, mapped, eraseOps, enableAsyncCopy, + copies))) { + signalPassFailure(); + return; + } + + rewriteConvertedOperands(module, mapped); + + 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(); + } + + finalizeConvertedArguments(module, convertedArgumentTypes, argumentCasts); + combineCopyWaits(copies); + + bool hasRemainingGpuTileOps = false; + module.walk([&](Operation *op) { + if (op->getName().getDialectNamespace() == "tile") { + op->emitError("was not eliminated by GPU TileIR conversion"); + hasRemainingGpuTileOps = true; + return; + } + if (auto cast = dyn_cast(op); + cast && cast.getInputs().size() == 1 && + isa(cast.getInputs().front().getType())) { + cast.emitError("CommonIR buffer bridge was not eliminated"); + hasRemainingGpuTileOps = true; + } + }); + if (hasRemainingGpuTileOps) + signalPassFailure(); + } +}; + +} // namespace + +} // namespace mlir::triton diff --git a/third_party/nvidia/triton_nvidia.cc b/third_party/nvidia/triton_nvidia.cc index 70ffcb4707..2733cb8be8 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,13 @@ 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_OPTION_WRAPPER_1("add_to_ttgir", mlir::triton::createCommonIRToTTGIR, + bool); +} +#endif + void init_triton_hopper_passes(py::module &&m) { // Meta's autoWS ADD_PASS_OPTION_WRAPPER_2("add_hopper_warpspec", @@ -148,6 +158,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/CMakeLists.txt b/third_party/tle/CMakeLists.txt index 2e7c1a2109..f815554b13 100644 --- a/third_party/tle/CMakeLists.txt +++ b/third_party/tle/CMakeLists.txt @@ -8,11 +8,11 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}/utils/include) add_subdirectory(dialect) if(TRITON_BUILD_PYTHON_MODULE) set(TLE_TILEIR_LINK_LIBS "") - if(TARGET TileIRIR) - list(APPEND TLE_TILEIR_LINK_LIBS TileIRIR) + if(TARGET CommonIRIR) + list(APPEND TLE_TILEIR_LINK_LIBS CommonIRIR) endif() - if(TARGET TileIRTransforms) - list(APPEND TLE_TILEIR_LINK_LIBS TileIRTransforms) + if(TARGET CommonIRTransforms) + list(APPEND TLE_TILEIR_LINK_LIBS CommonIRTransforms) endif() add_triton_plugin(TritonTLE diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index 87764bde1a..06de10ff9c 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -31,6 +31,7 @@ #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinDialect.h" +#include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/MLIRContext.h" #include "mlir/IR/Value.h" @@ -48,18 +49,17 @@ #include "triton/Dialect/TritonGPU/Transforms/Utility.h" #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" #ifdef __FLIR_TILEIR__ -#include "mlir-ext/Dialect/TileIR/IR/TileIRDialect.h" +#include "mlir-ext/Dialect/CommonIR/IR/CommonIRDialect.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/ADT/SmallVectorExtras.h" #include "llvm/IRReader/IRReader.h" #include "llvm/Support/Casting.h" #include "llvm/Support/ErrorHandling.h" #include "llvm/Support/MemoryBuffer.h" -#include "llvm/Support/raw_ostream.h" #include "llvm/Support/SourceMgr.h" +#include "llvm/Support/raw_ostream.h" #include #include #include @@ -109,144 +109,9 @@ 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) { - // Get the existing builder class from the main ir module (TLX style) auto &builder_cls = *ir::getBuilderClass(); @@ -375,34 +240,48 @@ 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(), - /*layout=*/tile::LayoutAttr::get(self.getBuilder().getContext(), - tile::Layout::ND), + /*layout=*/ + tile::LayoutAttr::get(self.getBuilder().getContext(), + 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 { + bool interNoAlias) -> void { auto op = self.create( src, dst, /*engine=*/tile::EngineAttr(), - /*src_layout=*/tile::LayoutAttr::get( - self.getBuilder().getContext(), tile::Layout::ND), + /*src_layout=*/ + tile::LayoutAttr::get(self.getBuilder().getContext(), + tile::Layout::ND), /*dst_nz_layout=*/tile::NZLayoutAttr(), /*transpose=*/mlir::UnitAttr(), /*comment=*/mlir::StringAttr()); if (interNoAlias) - op->setAttr("inter_no_alias", self.getBuilder().getBoolAttr(true)); + op->setAttr("inter_no_alias", + self.getBuilder().getBoolAttr(true)); + }) + .def("create_tile_get_memdesc", + [](TritonOpBuilder &self, Type resultTy, Value source) -> Value { + return self + .create(TypeRange{resultTy}, + ValueRange{source}) + .getResult(0); }) .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(); @@ -413,18 +292,19 @@ void init_triton_tle_ir(py::module &&m) { } auto *ctx = builder.getContext(); auto srcBuf = mlir::cast(source.getType()); - auto resTy = tile::BufType::get(ctx, sizes, srcBuf.getElementType(), - srcBuf.getMemorySpace()); + auto resTy = tile::BufType::get( + ctx, sizes, srcBuf.getElementType(), srcBuf.getMemorySpace()); 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_to_tensor", [](TritonOpBuilder &self, Value &src, bool /*writable*/) -> Value { auto srcBuf = mlir::cast(src.getType()); - auto resTy = - RankedTensorType::get(srcBuf.getShape(), srcBuf.getElementType()); + auto resTy = RankedTensorType::get(srcBuf.getShape(), + srcBuf.getElementType()); auto op = self.create(resTy, src); return op.getResult(); }) @@ -831,11 +711,19 @@ 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; #ifdef __FLIR_TILEIR__ - registry.insert(); + registry.insert(); context.appendDialectRegistry(registry); #endif context.loadAllAvailableDialects(); @@ -843,12 +731,10 @@ void init_triton_tle(py::module &&m) { #ifdef __FLIR_TILEIR__ m.def("load_tile_dialects", [](mlir::MLIRContext &context) { mlir::DialectRegistry registry; - registry.insert(); + registry.insert(); 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"));