diff --git a/CMakeLists.txt b/CMakeLists.txt index f2c9b04193..b3a32872ba 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -113,9 +113,38 @@ include(AddMLIR) # Utilities function(add_triton_object name) cmake_parse_arguments(ARG "" "" "DEPENDS;LINK_LIBS" ${ARGN}) + + # FLAGTREE SPEC SOURCE OVERRIDE: + # For a per-backend whole-file replacement of a main-tree source, drop a copy + # (derived from the FlagTree source file + backend edits) at + # third_party/${FLAGTREE_BACKEND}/backend/spec// + # This lets the backend compile its own version of a shared .cpp WITHOUT + # editing the pristine main-tree file (mirrors set_flagtree_backend_td for .td). + # EXISTS-gated: no effect for backends / files without an override file. + set(_flagtree_srcs "") + foreach(_src ${ARG_UNPARSED_ARGUMENTS}) + set(_resolved "${_src}") + if(FLAGTREE_BACKEND) + if(IS_ABSOLUTE "${_src}") + set(_abs "${_src}") + else() + set(_abs "${CMAKE_CURRENT_SOURCE_DIR}/${_src}") + endif() + file(RELATIVE_PATH _rel "${PROJECT_SOURCE_DIR}" "${_abs}") + if(NOT _rel MATCHES "^\\.\\.") + set(_spec "${PROJECT_SOURCE_DIR}/third_party/${FLAGTREE_BACKEND}/backend/spec/${_rel}") + if(EXISTS "${_spec}") + set(_resolved "${_spec}") + message(STATUS "[FlagTree spec] source override: ${_rel} -> ${_spec}") + endif() + endif() + endif() + list(APPEND _flagtree_srcs "${_resolved}") + endforeach() + add_library(${name} OBJECT) target_sources(${name} - PRIVATE ${ARG_UNPARSED_ARGUMENTS} + PRIVATE ${_flagtree_srcs} INTERFACE $ ) diff --git a/python/setup_tools/utils/xpu.py b/python/setup_tools/utils/xpu.py index 1a17cb020b..cd8a02ef2f 100644 --- a/python/setup_tools/utils/xpu.py +++ b/python/setup_tools/utils/xpu.py @@ -392,8 +392,8 @@ def register_cache(cache, flagtree_backend, check_env, set_llvm_env): copy_src_path=f"{cache.dir_path}/{flagtree_backend}/xpu-device-libs", copy_dst_path=f"third_party/{flagtree_backend}/device") cache.store(file="xpu-sdnn-objects", condition=is_xpu, - url="https://klx-sdk-release-public.su.bcebos.com/XTriton/xpu-sdnn-objects_v0.3.6.6.0.tar.gz", - version="v0.3.6.6.0", post_hook=lambda path: install_sdnn_objects(path, cache.flagtree_dir)) + url="https://klx-sdk-release-public.su.bcebos.com/XTriton/xpu-sdnn-objects_v0.3.6.7.1.tar.gz", + version="v0.3.6.7.1", post_hook=lambda path: install_sdnn_objects(path, cache.flagtree_dir)) cache.store( files=("clang", "xpu-xxd", "xpu3-elfconv", "xpu3-elfconv-triton", "xpu-kernel.t", "ld.lld", "llvm-readelf", "llvm-objdump", "llvm-objcopy"), condition=is_xpu, diff --git a/third_party/xpu/CMakeLists.txt b/third_party/xpu/CMakeLists.txt index 33d41187f4..3d901361f8 100644 --- a/third_party/xpu/CMakeLists.txt +++ b/third_party/xpu/CMakeLists.txt @@ -1,3 +1,14 @@ +# XPU: conceal all IR text (MLIR/LLVM/ASM) in non-Debug builds, aligning with +# internal Triton 3.0 TRITON_CONCEAL_IR. XPU-only. Also cover +# TritonRelBuildWithAsserts, which FlagTree uses as the default RelWithAsserts +# type and does not inherit CMAKE_CXX_FLAGS_RELEASE. +set(CMAKE_C_FLAGS_RELEASE "${CMAKE_C_FLAGS_RELEASE} -DTRITON_CONCEAL_IR=1") +set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -DTRITON_CONCEAL_IR=1") +set(CMAKE_C_FLAGS_TRITONRELBUILDWITHASSERTS "${CMAKE_C_FLAGS_TRITONRELBUILDWITHASSERTS} -DTRITON_CONCEAL_IR=1") +set(CMAKE_CXX_FLAGS_TRITONRELBUILDWITHASSERTS "${CMAKE_CXX_FLAGS_TRITONRELBUILDWITHASSERTS} -DTRITON_CONCEAL_IR=1") +add_compile_definitions($<$>:TRITON_CONCEAL_IR=1>) +message(STATUS "[XPU] TRITON_CONCEAL_IR=1 enabled for non-Debug XPU builds") + # ============================================================================ # XPU SDNN Object File Validation # ============================================================================ @@ -109,8 +120,17 @@ add_subdirectory(include) # compiles, without changing main-tree CMake files. if(TARGET TritonGPUIR) add_dependencies(TritonGPUIR TritonXPUAttrDefsIncGen) + target_include_directories(TritonGPUIR BEFORE PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_BINARY_DIR}/include) endif() +# ir.cc / llvm.cc are compiled as part of the main `triton` target, which is +# created after this subdirectory. Defer so the conceal macro actually reaches +# those sources. +cmake_language(DEFER DIRECTORY "${CMAKE_SOURCE_DIR}" CALL + target_compile_definitions triton PRIVATE TRITON_CONCEAL_IR=1) + add_subdirectory(lib) # ==================== FLAGTREE XPU SYNC MARK ==================== diff --git a/third_party/xpu/backend/compiler.py b/third_party/xpu/backend/compiler.py index 8aff946f67..46ca250ed3 100644 --- a/third_party/xpu/backend/compiler.py +++ b/third_party/xpu/backend/compiler.py @@ -107,6 +107,7 @@ class XPUOptions: buffer_size_limit: int = int(os.environ.get("TRITONXPU_BUFFER_SIZE", 512)) groups_per_cluster: int = int(os.environ.get("TRITONXPU_GROUPS_PER_CLUSTER", 1)) unroll_num: int = int(os.environ.get("TRITONXPU_UNROLL_NUM", 2)) + vrf_budget: int = int(os.environ.get("TRITONXPU_VRF_BUDGET", 24)) is_use_mask_zero: bool = int(os.environ.get("TRITONXPU_IS_USE_MASK_ZERO", 0)) extern_libs: dict = None is_sdnn: bool = False @@ -143,6 +144,7 @@ class XPUOptions: # use_int4_w4a8 use_int4_w4a8: bool = False + load_tile_size: int = 131072 def __post_init__(self): default_libdir = Path(__file__).parent / f"xpu{self.arch}" @@ -243,6 +245,7 @@ def make_ttxir(mod, metadata, opt): elem_bytes = int(os.environ.get("TRITONXPU_ELEMBYTES", 0)) groups_per_cluster = metadata["groups_per_cluster"] unroll_num = metadata["unroll_num"] + vrf_budget = metadata.get("vrf_budget", 24) XPUBackend.buffer_len = xpu.get_buffer_len(mod, max_buffer_size, elem_bytes) # print(f"XPUBackend.buffer_len = {XPUBackend.buffer_len}") core_num = metadata["core_num"] @@ -278,31 +281,48 @@ def make_ttxir(mod, metadata, opt): xpu.passes.ttsdnnir.add_triton_convert_type_pass(pm, opt.arch, 2) else: xpu.passes.ttsdnnir.add_triton_convert_type_pass(pm, opt.arch, TTSDNN_F_MATMUL_FAST_MODE) - xpu.passes.ttsdnnir.add_convert_triton_to_tritonsdnn_pass(pm, opt.arch) + xpu.passes.ttsdnnir.add_convert_triton_to_tritonsdnn_pass(pm, opt.arch, opt.load_tile_size) passes.ttir.add_loop_aware_cse(pm) xpu.passes.ttsdnnir.add_linalg_to_tritonsdnn_pass(pm, opt.arch) passes.ttir.add_loop_aware_cse(pm) xpu.passes.ttsdnnir.add_tritonsdnn_legalize_pass(pm, opt.arch) + xpu.passes.ttsdnnir.add_tritonsdnn_merge_extern_ew_pass(pm) xpu.passes.ttsdnnir.add_tritonsdnn_combine_before_pass(pm, opt.arch) + xpu.passes.ttsdnnir.add_tritonsdnn_resolve_layout_conflict_pass(pm) + xpu.passes.ttsdnnir.add_tritonsdnn_hoist_ds_pass(pm) + xpu.passes.ttsdnnir.add_tritonsdnn_transpose_mma_pass(pm) + xpu.passes.ttsdnnir.add_tritonsdnn_transpose_ew_pass(pm) + xpu.passes.ttsdnnir.add_tritonsdnn_combine_before_pass(pm, opt.arch) + xpu.passes.ttsdnnir.add_tritonsdnn_eliminate_mma_acc_zero_pass(pm, opt.arch) + if opt.arch == 4: + xpu.passes.ttsdnnir.add_tritonsdnn_fuse_mma_vector_bias_pass(pm, opt.arch) + xpu.passes.ttsdnnir.add_tritonsdnn_optimize_rc_layout_pass(pm, opt.arch) + if opt.arch == 4: + xpu.passes.ttsdnnir.add_tritonsdnn_ewlite_scheduling_pass(pm, opt.arch) + if TTSDNN_F_DMA_MODE: + xpu.passes.ttsdnnir.add_tritonsdnn_remove_ds_op_pass(pm, opt.arch) + xpu.passes.ttsdnnir.add_tritonsdnn_dsa_copy_pass(pm) xpu.passes.ttsdnnir.add_tritonsdnn_bufferize_pass(pm, opt.arch) + xpu.passes.ttsdnnir.add_tritonsdnn_mx_scale_layout_pass(pm) xpu.passes.ttsdnnir.add_tritonsdnn_combine_pass(pm, opt.arch) + xpu.passes.ttsdnnir.add_tritonsdnn_fuse_relu_activation_pass(pm, opt.arch) if opt.exp_range != ":0": res = parse_floating_range_string(opt.exp_range) xpu.passes.ttsdnnir.add_tritonsdnn_ew_act_table_pass(pm, res) else: xpu.passes.ttsdnnir.add_tritonsdnn_ew_act_table_pass(pm, None) xpu.passes.ttsdnnir.add_tritonsdnn_loop_grid_pass(pm) - if opt.arch == 4 and TTSDNN_F_DMA_MODE: - xpu.passes.ttsdnnir.add_tritonsdnn_remove_ds_op_pass(pm, opt.arch) - if not TTSDNN_F_SINGLE_CORE_MODE: - xpu.passes.ttsdnnir.add_tritonsdnn_pipeline_pass(pm) + xpu.passes.ttsdnnir.add_tritonsdnn_hoist_loop_invariant_dma_pass(pm, opt.arch) + xpu.passes.ttsdnnir.add_tritonsdnn_pipeline_pass(pm, opt.arch) if opt.arch == 4 and TTSDNN_F_KILL_EW_FILL_MODE: xpu.passes.ttsdnnir.add_tritonsdnn_kloop_acc_elimination_pass(pm) xpu.passes.ttsdnnir.add_tritonsdnn_multi_buffer_pass(pm, opt.arch, opt.num_stages) + xpu.passes.ttsdnnir.add_tritonsdnn_lower_rc_subview_pass(pm, opt.arch) passes.common.add_symbol_dce(pm) passes.common.add_canonicalizer(pm) passes.common.add_cse(pm) else: + xpu.passes.ttxpuir.add_tritonxpu_legalize_extern_ew_pass(pm) xpu.passes.ttxpuir.add_convert_triton_to_tritonxpu_pass(pm, opt.arch, XPUBackend.buffer_len, core_num) xpu.passes.ttxpuir.add_tritonxpu_print_pass(pm) xpu.passes.ttxpuir.add_tritonxpu_gm2lm_pass(pm, opt.arch, TTXPU_O_ATOMIC_SIM, opt.isClusterOneCoreActOnly, @@ -313,6 +333,7 @@ def make_ttxir(mod, metadata, opt): passes.common.add_canonicalizer(pm) if TTXPU_F_DTYPE_CONVERT: xpu.passes.ttxpuir.add_tritonxpu_dtype_convert_pass(pm, opt.arch) + xpu.passes.ttxpuir.add_tritonxpu_vectorizability_analysis_pass(pm, True, True) if not metadata["isCloseCoreTiling"]: xpu.passes.ttxpuir.add_tritonxpu_core_tiling_pass( pm, 0, XPUBackend.buffer_len, core_num, groups_per_cluster, @@ -320,8 +341,12 @@ def make_ttxir(mod, metadata, opt): # xpu.passes.ttxpuir.add_tritonxpu_lm_to_sm_pass(pm) passes.common.add_cse(pm) if not metadata["isCloseOffsetAnalysis"]: + xpu.passes.ttxpuir.add_tritonxpu_scalar_analysis_pass( + pm, False) if not TTXPU_O_CLOSE_OPT else None xpu.passes.ttxpuir.add_tritonxpu_offset_state_pass( pm, 0, XPUBackend.buffer_len, is_use_mask_zero) if not TTXPU_O_CLOSE_OPT else None # dumpFlag=0 + xpu.passes.ttxpuir.add_tritonxpu_scalar_analysis_pass( + pm, True) if not TTXPU_O_CLOSE_OPT else None passes.common.add_canonicalizer(pm) xpu.passes.ttxpuir.add_tritonxpu_legalize_pass(pm, XPUBackend.buffer_len, core_num, groups_per_cluster, is_use_mask_zero) @@ -337,17 +362,21 @@ def make_ttxir(mod, metadata, opt): passes.common.add_canonicalizer(pm) if not metadata["isCloseVectorization"]: compareFusion = int(os.environ.get("TRITONXPU_COMPARE_FUSION", 0)) + xpu.passes.ttxpuir.add_tritonxpu_normalize_pass( + pm, 0, compareFusion) if not TTXPU_O_CLOSE_OPT else None # dumpFlag=0 + xpu.passes.ttxpuir.add_tritonxpu_vectorizability_analysis_pass(pm, True, False) xpu.passes.ttxpuir.add_tritonxpu_vectorize_pass( pm, 0, compareFusion) if not TTXPU_O_CLOSE_OPT else None # dumpFlag=0 passes.common.add_canonicalizer(pm) xpu.passes.ttxpuir.add_tritonxpu_alloca_pass(pm, XPUBackend.buffer_len, core_num) if not metadata["isCloseMemoryAsync"]: - xpu.passes.ttxpuir.add_tritonxpu_memory_async_pass(pm, - 0) if not TTXPU_O_CLOSE_OPT else None # dumpFlag=0 + xpu.passes.ttxpuir.add_tritonxpu_async_load_schedule_pass( + pm, 0) if not TTXPU_O_CLOSE_OPT else None # dumpFlag=0 if not metadata["isCloseUnrollControl"]: - xpu.passes.ttxpuir.add_tritonxpu_unroll_control_pass(pm, XPUBackend.buffer_len, core_num, - is_use_mask_zero, - unroll_num) if not TTXPU_O_CLOSE_OPT else None + xpu.passes.ttxpuir.add_tritonxpu_tile_analysis_pass(pm, vrf_budget) + xpu.passes.ttxpuir.add_tritonxpu_unroll_control_pass( + pm, XPUBackend.buffer_len, core_num, is_use_mask_zero, unroll_num, + vrf_budget, False, -1) if not TTXPU_O_CLOSE_OPT else None xpu.passes.ttxpuir.add_tritonxpu_store_control_pass(pm) if not TTXPU_O_CLOSE_OPT else None if not TTXPU_F_OHTER_VALUE_SIM: xpu.passes.ttxpuir.add_tritonxpu_other_sim_pass(pm, XPUBackend.buffer_len, core_num) @@ -356,6 +385,8 @@ def make_ttxir(mod, metadata, opt): if not metadata["isCloseClusterLoopGrid"]: xpu.passes.ttxpuir.add_tritonxpu_cf_to_scf_pass(pm) xpu.passes.ttxpuir.add_tritonxpu_loop_grid_pass(pm) + if int(os.environ.get("TRITONXPU_LOOP_INVARIANT_STAGING", 0)): + xpu.passes.ttxpuir.add_tritonxpu_loop_invariant_staging_pass(pm) passes.common.add_cse(pm) passes.common.add_licm(pm) passes.common.add_symbol_dce(pm) diff --git a/third_party/xpu/backend/spec/PROVENANCE.md b/third_party/xpu/backend/spec/PROVENANCE.md new file mode 100644 index 0000000000..cdb21d01fd --- /dev/null +++ b/third_party/xpu/backend/spec/PROVENANCE.md @@ -0,0 +1,31 @@ +# FlagTree XPU — main-tree source overrides (whole-file vendored replacement) + +> 机制：`add_triton_object`（top `CMakeLists.txt`）对每个源文件检查本目录下是否有同「项目相对路径」的覆盖文件；有则 XPU build 编本副本，非 XPU 编主树原版。EXISTS 门控，非 XPU / 无覆盖零影响。 +> 该机制用于保持共享主树源码不含 XPU 专属语义，同时允许后端维护完整的源文件替换。 + +## 为什么用整文件替换而非就地改主树 +主树为所有后端共享，禁止就地改（Q0b）。这些副本**以 FlagTree 主树文件为底**（保留 `flagtree_hints` 等本地适配）+ 叠加 XPU 专属改动；**不以 internal 版为底**（实测 `Ops.cpp` 与 internal 有 162 行冲突 gap，如 `LoadOp::build` 的 `flagtree_hints` vs `offsetState/syncMode`）。 + +## Provenance（派生基线，用于反 drift） +- 派生自 FlagTree `main` 提交：`40a57023ddbd075c732eefb6886d6d3a152a181e` +- 迁移目标 internal Triton：`triton_3.6` @ `d3c64dd65e401239608164d6be4d893b261b4869` +- 生成时间：2026-07-10；2026-07-16 对最新 `main` 复核 + +上述五个主树源文件在原始派生提交 `b26ec15c65bea71f4e011c7626f7de9f5146937e` +与当前 `main` 基线之间内容一致；本次更新提交号用于准确反映 PR 的实际基线。 + +## 覆盖文件清单（副本 = 上述 FlagTree 文件 + 下列改动；括号为相对 pristine 的变动行） +| 覆盖路径 | 叠加的 XPU 改动 | 变动行 | +|---|---|---| +| `lib/Dialect/Triton/IR/Ops.cpp` | DotOp::verify 放宽 i8×i4(w4a8) · ReshapeOp::fold 去 `!getAllowReorder()` · BitcastOp::verify vector↔vector | 46 | +| `lib/Dialect/Triton/IR/Traits.cpp` | verifyTensorSize：允许非 pow2 元素数；verifyTensorLayouts：SliceEncoding 递归解析到 triton_xpu parent | 38 | +| `lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp` | ArithConstantSplatOpConversion splat guard | 2 | +| `lib/Dialect/TritonGPU/IR/Dialect.cpp` | ceil-offset：`getTotalElemsPerThread`/`getElemsPerThread(Attribute,shape)` 对 XPU 层（ClusterLayoutAttr / XPU-backed SliceEncoding）走 ceil-based 派发（新增 `isXPUBackedLayout` + `#include TritonXPU/IR/Dialect.h` + 文件内前置声明 `getElemsPerThread(Attribute,shape)`）；绕开泛型 LinearEncoding 的 pow2 断言 | 59 | +| `include/triton/Dialect/Triton/IR/TritonTypes.td` | 将 `TT_Vector`/`TT_VectorTensor` 组成的 `TT_VectorLike` 加入 `TT_Type`，允许 XPU vectorize 后的 `tt.extern_elementwise` 等主 Triton op 接受 vector-like operand/result | 6 | + +> `Dialect.cpp` 关键点：internal 在主树 header `TritonGPU/IR/Dialect.h` 加了 `getElemsPerThread(Attribute,ArrayRef)` 声明；FlagTree 该 header **无**此声明（Q0b 不改共享 header），故 `getTotalElemsPerThread` 里 line~112 的 `getElemsPerThread(layout,shape)` 会误配到 `getElemsPerThread(Type)` 报 `Attribute→Type` 转换错。修法：在本 vendored 副本内 `namespace mlir::triton::gpu` 前置声明该 overload（不动主树 header）。已经 XTDK clang22 实测编译+链接通过。 + +> 注：`Dialect.cpp` 副本保留了 pristine 顶部的 `flagtree_spec.h` 原生守卫（`#if __has_include("flagtree_spec.h")` / `#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_IR_Dialect`）。XPU 未提供 `third_party/xpu/backend/spec/include/flagtree_spec.h`，故 `__has_include` 为假、宏未定义、整个 body 正常编译——守卫无副作用。新增的 `#include "triton/Dialect/TritonXPU/IR/Dialect.h"` 经 XPU 后端 include dir（`third_party/xpu/include`）解析，仅存在于本副本。 + +## 维护须知（drift） +FlagTree 主树每次升级上述任一文件，**必须**用新版主树文件重做副本（以新主树为底重叠 XPU 改动），并更新本文件的派生提交号。否则 XPU 编译的是过期主树逻辑。 diff --git a/third_party/xpu/backend/spec/include/triton/Dialect/Triton/IR/TritonTypes.td b/third_party/xpu/backend/spec/include/triton/Dialect/Triton/IR/TritonTypes.td new file mode 100644 index 0000000000..c67b231ea5 --- /dev/null +++ b/third_party/xpu/backend/spec/include/triton/Dialect/Triton/IR/TritonTypes.td @@ -0,0 +1,133 @@ +#ifndef TRITON_TYPES +#define TRITON_TYPES + +include "mlir/IR/AttrTypeBase.td" +include "mlir/IR/BuiltinTypeInterfaces.td" +include "triton/Dialect/Triton/IR/TritonDialect.td" + +// +// Types +// +class TritonTypeDef traits = []> + : TypeDef { + // Used by printer/parser + let mnemonic = _mnemonic; +} + +// Floating-point Type +def TT_Float : AnyTypeOf<[F8E4M3FN, F8E4M3FNUZ, F8E5M2, F8E5M2FNUZ, F16, BF16, F32, F64], "floating-point">; +def TT_FloatTensor : RankedTensorOf<[TT_Float]>; +def TT_FloatLike : AnyTypeOf<[TT_Float, TT_FloatTensor]>; + +// Boolean Type +// TT_Bool -> I1 +def TT_BoolTensor : RankedTensorOf<[I1]>; +def TT_BoolLike : AnyTypeOf<[I1, TT_BoolTensor]>; + +// Integer Type +def I4 : I<4>; +def TT_Int : AnyTypeOf<[I1, I4, I8, I16, I32, I64], "integer">; +def TT_IntTensor : RankedTensorOf<[TT_Int]>; +def TT_IntLike : AnyTypeOf<[TT_Int, TT_IntTensor]>; + +// I32 Type +// TT_I32 -> I32 +// TT_I32Tensor -> I32Tensor +def TT_I32Like : AnyTypeOf<[I32, I32Tensor]>; + +// I64 Type +// TT_I64 -> I64 +// TT_I64Tensor -> I64Tensor +def TT_I64Like : AnyTypeOf<[I64, I64Tensor]>; + +// Pointer Type in TableGen +class TT_PtrOf pointeeTypes> : + DialectType($_self)">, + Concat<"[](::mlir::Type pointeeType) { return ", + SubstLeaves<"$_self", "pointeeType", AnyTypeOf.predicate>, + "; }(::mlir::cast<::mlir::triton::PointerType>($_self).getPointeeType())">]>, + "ptr", "::mlir::triton::PointerType">; + +// Pointer Type in C++ (corresponding to `TT_PtrOf`) +def TT_PtrType : TritonTypeDef<"Pointer", "ptr"> { + let summary = "Pointer type (`::mlir::triton::PointerType`) in Triton IR type system"; + + let description = [{ + Pointer type in Triton IR type system, which could be pointing to scalars or tensors. + }]; + + let parameters = (ins "Type":$pointeeType, "int":$addressSpace); + + let builders = [ + TypeBuilderWithInferredContext<(ins + "Type":$pointeeType, + "int":$addressSpace + ), [{ + return $_get(pointeeType.getContext(), pointeeType, addressSpace); + }]> + ]; + + let hasCustomAssemblyFormat = 1; + + let skipDefaultBuilders = 1; +} + +// Scalar Pointer Type: `ptr<>` +def TT_Ptr : TT_PtrOf<[AnyType]>; + +// Tensor of Pointer Type: `tensor>` +def TT_PtrTensor : RankedTensorOf<[TT_Ptr]>; + +// Tensor of Pointer Type or Pointer type: `tensor>` or `ptr<>` +def TT_PtrLike : AnyTypeOf<[TT_Ptr, TT_PtrTensor]>; + +// Tensor Type +def TT_FpIntTensor : RankedTensorOf<[TT_Float, TT_Int]>; +def TT_Tensor : RankedTensorOf<[TT_Float, TT_Int, TT_Ptr]>; + +// Pointer Type to Tensor Type: `ptr>` +def TT_TensorPtr : TT_PtrOf<[TT_Tensor]>; + +def TT_Vector : FixedVectorOfNonZeroRankOf<[TT_FloatLike, TT_IntLike]>; +def TT_VectorTensor : TensorOf<[TT_Vector]>; +def TT_VectorLike : AnyTypeOf<[TT_Vector, TT_VectorTensor]>; + +// Any Type in Triton IR +def TT_Type : AnyTypeOf<[TT_FloatLike, TT_IntLike, TT_PtrLike, TT_TensorPtr, TT_VectorLike]>; + +// Result type of MakeTensorDescriptor +def TT_TensorDescType : TritonTypeDef<"TensorDesc", "tensordesc", []> { + let summary = "Tensor descriptor type (`::mlir::triton::TensorDescType`) in Triton IR type system"; + + let description = [{ + A portable abstraction for nvidia-TMA descriptors. + }]; + + let parameters = (ins "RankedTensorType":$blockType); + let assemblyFormat = "`<` $blockType `>`"; + + let builders = [ + TypeBuilder<(ins "RankedTensorType":$blockType, "bool":$isSigned), [{ + if (auto intTy = llvm::dyn_cast(blockType.getElementType())) { + auto sem = isSigned ? IntegerType::Signed : IntegerType::Unsigned; + auto elemTy = IntegerType::get($_ctxt, intTy.getWidth(), sem); + blockType = blockType.clone(elemTy); + } + return Base::get($_ctxt, blockType); + }]>, + ]; + let extraClassDeclaration = [{ + RankedTensorType getSignlessBlockType() const { + auto resTy = getBlockType(); + if (auto intTy = llvm::dyn_cast(resTy.getElementType())) { + auto width = resTy.getElementTypeBitWidth(); + auto signlessTy = IntegerType::get(getContext(), width); + resTy = resTy.clone(signlessTy); + } + return resTy; + } + }]; +} + +#endif diff --git a/third_party/xpu/backend/spec/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp b/third_party/xpu/backend/spec/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp new file mode 100644 index 0000000000..5acb0ef85e --- /dev/null +++ b/third_party/xpu/backend/spec/lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp @@ -0,0 +1,605 @@ +#include "mlir/Support/LLVM.h" +#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h" +#include "triton/Conversion/TritonGPUToLLVM/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Attributes.h" +#include "triton/Dialect/TritonGPU/IR/Types.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#include "triton/Tools/LayoutUtils.h" + +using namespace mlir; +using namespace mlir::triton; +using namespace mlir::triton::gpu; +using ::mlir::LLVM::getSharedMemoryObjectFromStruct; +namespace { + +Value bitOrPtrCast(Value val, Type type, TritonLLVMOpBuilder &b) { + if (isa(val.getType()) && + !isa(type)) { + return b.ptrtoint(type, val); + } else { + return b.bitcast(val, type); + } +} + +struct SplatOpConversion : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + // Convert SplatOp or arith::ConstantOp with SplatElementsAttr to a + // LLVM::StructType value. + // + // @elemType: the element type in operand. + // @resType: the return type of the Splat-like op. + // @constVal: a LLVM::ConstantOp or other scalar value. + static Value convertSplatLikeOp(Type elemType, Type resType, Value constVal, + const LLVMTypeConverter *typeConverter, + ConversionPatternRewriter &rewriter, + Location loc) { + auto b = TritonLLVMOpBuilder(loc, rewriter); + auto tensorTy = cast(resType); + // Check the converted type for the tensor as depending on the encoding the + // converter may pick different element types. + auto srcType = typeConverter->convertType(tensorTy); + if (auto structTy = dyn_cast(srcType)) + srcType = structTy.getBody()[0]; + // If the type sizes don't match we need to pack constants. + if (srcType.isIntOrFloat() && constVal.getType().getIntOrFloatBitWidth() != + srcType.getIntOrFloatBitWidth()) { + unsigned cstBitWidth = constVal.getType().getIntOrFloatBitWidth(); + unsigned srcBitWidth = srcType.getIntOrFloatBitWidth(); + assert(cstBitWidth <= srcBitWidth && srcBitWidth % cstBitWidth == 0); + unsigned ratio = srcBitWidth / cstBitWidth; + Type intTy = IntegerType::get(elemType.getContext(), cstBitWidth); + VectorType vecType = VectorType::get(ratio, intTy); + Value intCst = bitOrPtrCast(constVal, intTy, b); + Value vec = b.undef(vecType); + for (unsigned i = 0; i < ratio; ++i) + vec = b.insert_element(vecType, vec, intCst, b.int_val(32, i)); + constVal = vec; + } + Value llSrc = bitOrPtrCast(constVal, srcType, b); + size_t elemsPerThread = getTotalElemsPerThread(tensorTy); + llvm::SmallVector elems(elemsPerThread, llSrc); + return packLLElements(loc, typeConverter, elems, rewriter, resType); + } + LogicalResult matchAndRewrite(triton::SplatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op->getLoc(); + auto src = adaptor.getSrc(); + auto typeConverter = getTypeConverter(); + auto llStruct = convertSplatLikeOp(src.getType(), op.getType(), src, + typeConverter, rewriter, loc); + rewriter.replaceOp(op, {llStruct}); + return success(); + } +}; + +struct UnsplatOpConversion : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + LogicalResult matchAndRewrite(triton::UnsplatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op->getLoc(); + auto scrVals = unpackLLElements(loc, adaptor.getSrc(), rewriter); + rewriter.replaceOp(op, scrVals[0]); + return success(); + } +}; + +// This pattern helps to convert arith::ConstantOp(with SplatElementsAttr), +// the logic is the same as triton::SplatOp, so the underlying implementation +// is reused. +struct ArithConstantSplatOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + LogicalResult + matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto value = op.getValue(); + if (!mlir::dyn_cast(value)) + return failure(); + if (!mlir::isa(op.getType())) + return failure(); + auto loc = op->getLoc(); + LLVM::ConstantOp arithConstantOp; + auto values = mlir::dyn_cast(op.getValue()); + auto elemType = values.getElementType(); + Attribute val; + if (type::isFloat(elemType)) { + val = values.getValues()[0]; + } else if (type::isInt(elemType)) { + val = values.getValues()[0]; + } else { + llvm::errs() << "ArithConstantSplatOpConversion get unsupported type: " + << value.getType() << "\n"; + return failure(); + } + // Lower FP8 constant to int8 constant since FP8 types are not supported on + // LLVM IR. + if (type::isFloat8(elemType)) + elemType = rewriter.getIntegerType(8); + auto constOp = LLVM::ConstantOp::create(rewriter, loc, elemType, val); + auto typeConverter = getTypeConverter(); + auto llStruct = SplatOpConversion::convertSplatLikeOp( + elemType, op.getType(), constOp, typeConverter, rewriter, loc); + rewriter.replaceOp(op, llStruct); + return success(); + } +}; + +// Convert arith::ConstantOp with an array DenseElementsAttr to a +// LLVM::StructType value. +struct ArithConstantArrayOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + LogicalResult + matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto value = op.getValue(); + if (!mlir::dyn_cast(value)) + return failure(); + if (mlir::isa(value)) + return failure(); + auto tensorTy = cast(op.getType()); + auto loc = op->getLoc(); + auto values = mlir::dyn_cast(op.getValue()); + auto elemType = values.getElementType(); + SmallVector llVals; + for (auto v : values.getValues()) { + auto ll = LLVM::ConstantOp::create(rewriter, loc, elemType, v); + llVals.push_back(ll); + } + size_t elemsPerThread = getTotalElemsPerThread(tensorTy); + + if (elemsPerThread != llVals.size()) { + op->emitError( + "Right now we only support constant arrays with the same number of " + "elements as the number of threads per warp"); + return failure(); + } + auto llStruct = + packLLElements(loc, getTypeConverter(), llVals, rewriter, op.getType()); + rewriter.replaceOp(op, {llStruct}); + return success(); + } +}; + +struct CatOpConversion : public ConvertOpToLLVMPattern { + using OpAdaptor = typename CatOp::Adaptor; + explicit CatOpConversion(LLVMTypeConverter &typeConverter, + PatternBenefit benefit = patternBenefitDefault) + : ConvertOpToLLVMPattern(typeConverter, benefit) {} + LogicalResult + matchAndRewrite(CatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); + auto resultTy = cast(op.getType()); + unsigned elems = getTotalElemsPerThread(resultTy); + auto typeConverter = getTypeConverter(); + Type elemTy = typeConverter->convertType(resultTy.getElementType()); + SmallVector types(elems, elemTy); + // unpack input values + auto lhsVals = unpackLLElements(loc, adaptor.getLhs(), rewriter); + auto rhsVals = unpackLLElements(loc, adaptor.getRhs(), rewriter); + // concatenate (and potentially reorder) values + SmallVector retVals; + for (Value v : lhsVals) + retVals.push_back(v); + for (Value v : rhsVals) + retVals.push_back(v); + // pack and replace + Value ret = packLLElements(loc, typeConverter, retVals, rewriter, resultTy); + rewriter.replaceOp(op, ret); + return success(); + } +}; +struct JoinOpConversion : public ConvertOpToLLVMPattern { + using OpAdaptor = typename JoinOp::Adaptor; + explicit JoinOpConversion(LLVMTypeConverter &typeConverter, + PatternBenefit benefit = patternBenefitDefault) + : ConvertOpToLLVMPattern(typeConverter, benefit) {} + LogicalResult + matchAndRewrite(JoinOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + // We rely on the following invariants of this op (which are checked by its + // verifier): + // + // - The last dimension (the one we're joining) is also the most minor + // dimension. + // - The input and output encodings are the same, except the output has + // 2 elements per thread in the last dim. + // + // With these invariants, join is trivial: We can count how many contiguous + // registers belong to the same chunk then we merge the registers between + // two different chunks. + Location loc = op->getLoc(); + RankedTensorType dstTy = op.getType(); + auto ll = toLinearLayout(dstTy); + int splitDim = dstTy.getRank() - 1; + auto kReg = mlir::StringAttr::get(dstTy.getContext(), "register"); + const auto &bases = ll.getBases(); + const auto ®s = bases.find(kReg)->second; + int numContiguousValues = 1; + bool found = false; + for (const auto ® : regs) { + if (reg[splitDim] == 1) { + found = true; + break; + } + numContiguousValues *= 2; + } + assert(found && "Join dimension is not distributed along registers."); + SmallVector lhsVals = + unpackLLElements(loc, adaptor.getLhs(), rewriter); + SmallVector rhsVals = + unpackLLElements(loc, adaptor.getRhs(), rewriter); + assert(lhsVals.size() == rhsVals.size()); + SmallVector joinedVals; + joinedVals.resize(lhsVals.size() * 2); + for (int i = 0; i < lhsVals.size(); i += numContiguousValues) { + for (int j = 0; j < numContiguousValues; j++) { + joinedVals[2 * i + j] = lhsVals[i + j]; + joinedVals[2 * i + numContiguousValues + j] = rhsVals[i + j]; + } + } + auto typeConverter = getTypeConverter(); + Value ret = packLLElements(loc, typeConverter, joinedVals, rewriter, dstTy); + rewriter.replaceOp(op, ret); + return success(); + } +}; +struct SplitOpConversion : public ConvertOpToLLVMPattern { + using OpAdaptor = typename SplitOp::Adaptor; + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + LogicalResult + matchAndRewrite(SplitOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + // We rely on the following invariants of this op (which are checked by its + // verifier): + // + // - The layout distribute the last dimension along registers + // - The last dimension (the one we're splitting) has sizePerThread=2, + // threadPerWarp=1 and warpPerBlock=1. + // + // With these invariants, split is trivial: We can count how many contiguous + // registers belong to the same chunk then we separate the registers between + // two different chunks. + auto srcTy = cast(op.getSrc().getType()); + auto ll = toLinearLayout(srcTy); + int splitDim = srcTy.getRank() - 1; + auto kReg = mlir::StringAttr::get(srcTy.getContext(), "register"); + const auto &bases = ll.getBases(); + const auto ®s = bases.find(kReg)->second; + int numContiguousValues = 1; + bool found = false; + for (const auto ® : regs) { + if (reg[splitDim] == 1) { + found = true; + break; + } + numContiguousValues *= 2; + } + assert(found && "Split dimension is not distributed along registers."); + Location loc = op->getLoc(); + auto typeConverter = getTypeConverter(); + SmallVector srcVals = + unpackLLElements(loc, adaptor.getSrc(), rewriter); + assert(srcVals.size() % 2 == 0); + SmallVector outLhsVals; + SmallVector outRhsVals; + for (int i = 0; i < srcVals.size(); i += 2 * numContiguousValues) { + for (int j = 0; j < numContiguousValues; j++) { + outLhsVals.push_back(srcVals[i + j]); + outRhsVals.push_back(srcVals[i + numContiguousValues + j]); + } + } + auto resultTy = cast(op.getResult(0).getType()); + Value retLhs = + packLLElements(loc, typeConverter, outLhsVals, rewriter, resultTy); + Value retRhs = + packLLElements(loc, typeConverter, outRhsVals, rewriter, resultTy); + rewriter.replaceOp(op, {retLhs, retRhs}); + return success(); + } +}; +struct ReshapeOpConversion : public ConvertOpToLLVMPattern { + using OpAdaptor = typename ReshapeOp::Adaptor; + explicit ReshapeOpConversion(LLVMTypeConverter &typeConverter, + PatternBenefit benefit = patternBenefitDefault) + : ConvertOpToLLVMPattern(typeConverter, benefit) {} + LogicalResult + matchAndRewrite(ReshapeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); + if (triton::gpu::isExpensiveView(op.getSrc().getType(), op.getType())) { + return emitOptionalError(loc, + "expensive view not supported on reshape op"); + } + auto resultTy = cast(op.getType()); + auto srcTy = cast(op.getSrc().getType()); + auto typeConverter = getTypeConverter(); + auto vals = unpackLLElements(loc, adaptor.getSrc(), rewriter); + Value ret = packLLElements(loc, typeConverter, vals, rewriter, resultTy); + rewriter.replaceOp(op, ret); + return success(); + } +}; +struct ExpandDimsOpConversion : public ConvertOpToLLVMPattern { + using OpAdaptor = typename ExpandDimsOp::Adaptor; + explicit ExpandDimsOpConversion( + LLVMTypeConverter &typeConverter, + PatternBenefit benefit = patternBenefitDefault) + : ConvertOpToLLVMPattern(typeConverter, benefit) {} + LogicalResult + matchAndRewrite(ExpandDimsOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); + auto typeConverter = getTypeConverter(); + auto srcVals = unpackLLElements(loc, adaptor.getSrc(), rewriter); + auto srcTy = cast(op.getSrc().getType()); + auto resultTy = cast(op.getType()); + auto srcLayout = dyn_cast(srcTy.getEncoding()); + if (!srcLayout) { + return emitOptionalError( + loc, "ExpandDimsOp only supports SliceEncodingAttr as its input"); + } + auto resultLayout = resultTy.getEncoding(); + auto srcOffsets = emitOffsetForLayout(srcLayout, srcTy); + auto resultOffsets = emitOffsetForLayout(resultLayout, resultTy); + std::map, Value> srcValues; + for (size_t i = 0; i < srcOffsets.size(); i++) { + srcValues[srcOffsets[i]] = srcVals[i]; + } + SmallVector resultVals; + for (size_t i = 0; i < resultOffsets.size(); i++) { + auto offset = resultOffsets[i]; + offset.erase(offset.begin() + srcLayout.getDim()); + resultVals.push_back(srcValues.at(offset)); + } + Value ret = + packLLElements(loc, typeConverter, resultVals, rewriter, resultTy); + rewriter.replaceOp(op, ret); + return success(); + } +}; +struct MemDescTransOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + LogicalResult + matchAndRewrite(MemDescTransOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); + auto resultTy = cast(op.getType()); + auto llvmElemTy = + getTypeConverter()->convertType(resultTy.getElementType()); + auto srcSmemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(), + llvmElemTy, rewriter); + auto dstSmemObj = SharedMemoryObject( + srcSmemObj.getBase(), srcSmemObj.getBaseElemType(), + /*offsets=*/applyPermutation(srcSmemObj.getOffsets(), op.getOrder())); + auto retVal = getStructFromSharedMemoryObject(loc, dstSmemObj, rewriter); + rewriter.replaceOp(op, retVal); + return success(); + } +}; + +struct MemDescReshapeOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + LogicalResult + matchAndRewrite(MemDescReshapeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); + auto resultTy = cast(op.getType()); + auto llvmElemTy = + getTypeConverter()->convertType(resultTy.getElementType()); + auto srcSmemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(), + llvmElemTy, rewriter); + SmallVector offsets = srcSmemObj.getOffsets(); + // FIXME: This should be done by composing a linear layout with its + // reshaped counterpart. + SmallVector srcShape; + for (int64_t d : op.getSrc().getType().getShape()) + srcShape.push_back(d); + SmallVector dstShape; + for (int64_t d : op.getType().getShape()) + dstShape.push_back(d); + Value linearOffset = LLVM::linearize(rewriter, loc, offsets, srcShape); + SmallVector delinearizedOffset = + LLVM::delinearize(rewriter, loc, linearOffset, dstShape); + auto b = TritonLLVMOpBuilder(loc, rewriter); + auto dstSmemObj = SharedMemoryObject( + srcSmemObj.getBase(), srcSmemObj.getBaseElemType(), delinearizedOffset); + auto retVal = getStructFromSharedMemoryObject(loc, dstSmemObj, rewriter); + rewriter.replaceOp(op, retVal); + return success(); + } +}; + +struct TransOpConversion : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + LogicalResult + matchAndRewrite(TransOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + // By construction, TransOp::inferReturnTypes ensures that the src encoding + // is the same as the dst encoding so that this op is a no-op. + rewriter.replaceOp(op, adaptor.getSrc()); + return success(); + } +}; + +struct BroadcastOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + LogicalResult + matchAndRewrite(triton::BroadcastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + // Following the order of indices in the legacy code, a broadcast of: + // [s(0), s(1) ... s(k-1), 1, s(k+1), s(k+2) ... s(n-1)] + // => + // [s(0), s(1) ... s(k-1), s(k), s(k+1), s(k+2) ... s(n-1)] + // + // logically maps to a broadcast within a thread's scope: + // [cta(0)..cta(k-1), 1,cta(k+1)..cta(n-1),spt(0)..spt(k-1), + // 1,spt(k+1)..spt(n-1)] + // => + // [cta(0)..cta(k-1),cta(k),cta(k+1)..cta(n-1),spt(0)..spt(k-1),spt(k),spt(k+1)..spt(n-1)] + // + // regardless of the order of the layout + // + Location loc = op->getLoc(); + Value src = adaptor.getSrc(); + Value result = op.getResult(); + auto srcTy = cast(op.getSrc().getType()); + auto resultTy = cast(result.getType()); + auto srcLayout = srcTy.getEncoding(); + auto resultLayout = resultTy.getEncoding(); + auto srcShape = srcTy.getShape(); + auto resultShape = resultTy.getShape(); + unsigned rank = srcTy.getRank(); + auto typeConverter = getTypeConverter(); + assert(rank == resultTy.getRank()); + auto srcOffsets = emitOffsetForLayout(srcLayout, srcTy); + auto resultOffsets = emitOffsetForLayout(resultLayout, resultTy); + SmallVector srcVals = unpackLLElements(loc, src, rewriter); + std::map, Value> srcValues; + for (size_t i = 0; i < srcOffsets.size(); i++) { + srcValues[srcOffsets[i]] = srcVals[i]; + } + SmallVector resultVals; + for (size_t i = 0; i < resultOffsets.size(); i++) { + auto offset = resultOffsets[i]; + for (size_t j = 0; j < srcShape.size(); j++) + if (srcShape[j] == 1) + offset[j] = 0; + resultVals.push_back(srcValues.at(offset)); + } + Value resultStruct = + packLLElements(loc, typeConverter, resultVals, rewriter, resultTy); + rewriter.replaceOp(op, {resultStruct}); + return success(); + } +}; + +struct MemDescIndexOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern< + triton::gpu::MemDescIndexOp>::ConvertOpToLLVMPattern; + + LogicalResult + matchAndRewrite(triton::gpu::MemDescIndexOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); + auto *ctx = op->getContext(); + auto b = TritonLLVMOpBuilder(loc, rewriter); + auto srcTy = op.getSrc().getType(); + auto dstTy = op.getResult().getType(); + auto llvmElemTy = getTypeConverter()->convertType(srcTy.getElementType()); + + // getAllocationShapePerCTA returns the correct number fp4 elements that we + // need to skip when we have fp4Padded=True. getShapePerCTA does not account + // for this + auto stride = product( + getAllocationShapePerCTA(dstTy.getEncoding(), dstTy.getShape())); + Value offset = b.mul(op.getIndex(), b.i32_val(stride)); + auto smemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(), + llvmElemTy, rewriter); + auto base = smemObj.getBase(); + auto elemPtrTy = base.getType(); + auto prevOffsets = smemObj.getOffsets(); + SmallVector offsetVals(prevOffsets.end() - dstTy.getRank(), + prevOffsets.end()); + + // Apply padding based on the amount we move the base ptr + if (auto padEnc = dyn_cast(dstTy.getEncoding())) { + auto bitwidth = dstTy.getElementTypeBitWidth(); + Value padOffset = emitPadding(loc, rewriter, padEnc, bitwidth, offset, + /*offsetInBytes=*/false); + offset = b.add(offset, padOffset); + } + + // Advance the pointer and keep the opOffsets as the new shape + smemObj = SharedMemoryObject(b.gep(elemPtrTy, llvmElemTy, base, offset), + llvmElemTy, offsetVals); + auto retVal = getStructFromSharedMemoryObject(loc, smemObj, rewriter); + rewriter.replaceOp(op, retVal); + return success(); + } +}; + +struct MemDescSubsliceOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern< + triton::gpu::MemDescSubsliceOp>::ConvertOpToLLVMPattern; + + LogicalResult + matchAndRewrite(triton::gpu::MemDescSubsliceOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op->getLoc(); + auto *ctx = op->getContext(); + auto b = TritonLLVMOpBuilder(loc, rewriter); + auto srcTy = op.getSrc().getType(); + auto destTy = op.getResult().getType(); + auto llvmElemTy = getTypeConverter()->convertType(srcTy.getElementType()); + auto layoutOrder = getOrder(srcTy); + auto enc = srcTy.getEncoding(); + + auto smemObj = getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(), + llvmElemTy, rewriter); + auto opOffsetVals = op.getOffsets(); + + auto base = smemObj.getBase(); + auto elemPtrTy = base.getType(); + // Accumulate the logical offsets + SmallVector offsetVals; + for (auto [oldOffVal, opOff] : + llvm::zip(smemObj.getOffsets(), opOffsetVals)) { + offsetVals.push_back(b.add(oldOffVal, b.i32_val(opOff))); + } + smemObj = SharedMemoryObject(base, llvmElemTy, offsetVals); + auto retVal = getStructFromSharedMemoryObject(loc, smemObj, rewriter); + rewriter.replaceOp(op, retVal); + return success(); + } +}; + +struct MemDescReinterpretOpConversion + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + + LogicalResult matchAndRewrite(MemDescReinterpretOp op, OpAdaptor adaptor, + ConversionPatternRewriter &b) const override { + Location loc = op.getLoc(); + MemDescType srcTy = op.getSrc().getType(); + MemDescType dstTy = op.getType(); + Type srcElemTy = getTypeConverter()->convertType(srcTy.getElementType()); + Type dstElemTy = getTypeConverter()->convertType(dstTy.getElementType()); + + auto smemObj = + getSharedMemoryObjectFromStruct(loc, adaptor.getSrc(), srcElemTy, b); + Value newBase = smemObj.getShmemAffineBase(loc, b, srcTy); + SharedMemoryObject newObj(newBase, dstElemTy, dstTy.getRank(), loc, b); + b.replaceOp(op, getStructFromSharedMemoryObject(loc, newObj, b)); + return success(); + } +}; + +} // namespace + +void mlir::triton::populateViewOpToLLVMPatterns( + LLVMTypeConverter &typeConverter, RewritePatternSet &patterns, + PatternBenefit benefit) { + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add( + typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add(typeConverter, benefit); + patterns.add( + typeConverter, benefit); + patterns.add(typeConverter, benefit); +} diff --git a/third_party/xpu/backend/spec/lib/Dialect/Triton/IR/Ops.cpp b/third_party/xpu/backend/spec/lib/Dialect/Triton/IR/Ops.cpp new file mode 100644 index 0000000000..a3fd0bc8d7 --- /dev/null +++ b/third_party/xpu/backend/spec/lib/Dialect/Triton/IR/Ops.cpp @@ -0,0 +1,1514 @@ +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/Interfaces/FunctionImplementation.h" +#include "mlir/Interfaces/FunctionInterfaces.h" +#include "mlir/Support/LLVM.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/MathExtras.h" + +namespace mlir { +namespace triton { + +void LoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getPtrMutable(), + GlobalMemory::get()); + if (getIsVolatile()) + effects.emplace_back(MemoryEffects::Write::get()); +} + +} // namespace triton +} // namespace mlir + +#define GET_OP_CLASSES +#include "triton/Dialect/Triton/IR/Ops.cpp.inc" + +// enum attribute definitions +#include "triton/Dialect/Triton/IR/OpsEnums.cpp.inc" + +#include "TritonCanonicalize.inc" + +namespace mlir { +namespace triton { + +//-- LoadOp -- +void LoadOp::build(OpBuilder &builder, OperationState &state, Value ptr, + CacheModifier cache, EvictionPolicy evict, bool isVolatile, + mlir::StringAttr flagtree_hints) { + LoadOp::build(builder, state, ptr, /*mask=*/{}, /*other=*/{}, + /*boundaryCheck=*/ArrayRef{}, /*padding=*/std::nullopt, + cache, evict, isVolatile, flagtree_hints); +} + +void LoadOp::build(OpBuilder &builder, OperationState &state, Value ptr, + ArrayRef boundaryCheck, + std::optional padding, CacheModifier cache, + EvictionPolicy evict, bool isVolatile, + mlir::StringAttr flagtree_hints) { + LoadOp::build(builder, state, ptr, /*mask=*/{}, /*other=*/{}, boundaryCheck, + padding, cache, evict, isVolatile, flagtree_hints); +} + +void LoadOp::build(OpBuilder &builder, OperationState &state, Value ptr, + Value mask, CacheModifier cache, EvictionPolicy evict, + bool isVolatile, mlir::StringAttr flagtree_hints) { + LoadOp::build(builder, state, ptr, mask, /*other=*/{}, + /*boundaryCheck=*/ArrayRef{}, + /*padding=*/std::nullopt, cache, evict, isVolatile, + flagtree_hints); +} + +void LoadOp::build(OpBuilder &builder, OperationState &state, Value ptr, + Value mask, Value other, CacheModifier cache, + EvictionPolicy evict, bool isVolatile, + mlir::StringAttr flagtree_hints) { + LoadOp::build(builder, state, ptr, mask, other, + /*boundaryCheck=*/ArrayRef{}, + /*padding=*/std::nullopt, cache, evict, isVolatile, + flagtree_hints); +} + +void LoadOp::build(OpBuilder &builder, OperationState &state, Value ptr, + Value mask, Value other, ArrayRef boundaryCheck, + std::optional padding, CacheModifier cache, + EvictionPolicy evict, bool isVolatile, + mlir::StringAttr flagtree_hints) { + auto paddingAttr = + padding.has_value() + ? PaddingOptionAttr::get(builder.getContext(), padding.value()) + : PaddingOptionAttr(); + if (!flagtree_hints) // flagtree hints: if not provided, use empty string + flagtree_hints = builder.getStringAttr(""); + LoadOp::build(builder, state, ptr, mask, other, + builder.getDenseI32ArrayAttr(boundaryCheck), paddingAttr, cache, + evict, isVolatile, flagtree_hints); +} + +// load(ptr, splat(1), ...) -> load(ptr, ...) +// load(ptr, splat(0), other, ...) -> other +struct CanonicalizeMaskedLoadPattern : public OpRewritePattern { + CanonicalizeMaskedLoadPattern(MLIRContext *context) + : OpRewritePattern(context, 1) {} + + LogicalResult matchAndRewrite(LoadOp loadOp, + PatternRewriter &rewriter) const override { + auto mask = loadOp.getMask(); + if (!mask) + return failure(); + + auto constantMask = mask.getDefiningOp(); + if (!constantMask) + return failure(); + + auto splatMask = mlir::dyn_cast(constantMask.getValue()); + if (!splatMask) + return failure(); + + if (splatMask.getSplatValue().getValue() == true) { + // mask = splat(1) + rewriter.replaceOpWithNewOp( + loadOp, loadOp.getType(), loadOp.getPtr(), Value(), Value(), + loadOp.getBoundaryCheckAttr(), loadOp.getPaddingAttr(), + loadOp.getCache(), loadOp.getEvict(), loadOp.getIsVolatile()); + } else { + // mask = splat(0) + + // If there's no "other", the value is "undef". Perhaps we want to + // optimize it in the future.x + auto otherVal = loadOp.getOther(); + if (!otherVal) + return failure(); + rewriter.replaceOp(loadOp, otherVal); + } + return success(); + } +}; + +void LoadOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.add(context); +} + +//-- StoreOp -- +void StoreOp::build(OpBuilder &builder, OperationState &state, Value ptr, + Value value, CacheModifier cache, EvictionPolicy evict) { + return StoreOp::build(builder, state, ptr, value, /*mask=*/{}, + /*boundaryCheck=*/{}, cache, evict); +} + +void StoreOp::build(OpBuilder &builder, OperationState &state, Value ptr, + Value value, Value mask, CacheModifier cache, + EvictionPolicy evict) { + return StoreOp::build(builder, state, ptr, value, mask, /*boundaryCheck=*/{}, + cache, evict); +} + +void StoreOp::build(OpBuilder &builder, OperationState &state, Value ptr, + Value value, ArrayRef boundaryCheck, + CacheModifier cache, EvictionPolicy evict) { + return StoreOp::build(builder, state, ptr, value, /*mask=*/{}, + builder.getDenseI32ArrayAttr(boundaryCheck), cache, + evict); +} + +// store(ptr, value, splat(1), ...) -> store(ptr, value, ...) +// store(ptr, value, splat(0), ...) -> [none] +struct CanonicalizeMaskedStorePattern : public OpRewritePattern { + CanonicalizeMaskedStorePattern(MLIRContext *context) + : OpRewritePattern(context, 1) {} + + LogicalResult matchAndRewrite(StoreOp storeOp, + PatternRewriter &rewriter) const override { + auto mask = storeOp.getMask(); + if (!mask) + return failure(); + + auto constantMask = mask.getDefiningOp(); + if (!constantMask) + return failure(); + + auto splatMask = mlir::dyn_cast(constantMask.getValue()); + if (!splatMask) + return failure(); + + if (splatMask.getSplatValue().getValue() == true) { + // mask = splat(1) + rewriter.replaceOpWithNewOp( + storeOp, storeOp.getPtr(), storeOp.getValue(), storeOp.getCache(), + storeOp.getEvict()); + } else { + // mask = splat(0) + rewriter.eraseOp(storeOp); + } + return success(); + } +}; + +void StoreOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.add(context); +} + +//-- TransOp -- +OpFoldResult TransOp::fold(FoldAdaptor adaptor) { + // transpose(x, order=[0, 1, ...]) -> x + if (isIota(getOrder())) { + // If the source and result types are the same, we can return the source + // If their layout is different (even if structurally equivalent), we need + // to insert a convert_layout in between as otherwise ::fold complains + // We do this in CanonicalizeConvertFromTranspose + if (getSrc().getType() == getType()) { + return getSrc(); + } + } + + // transpose(transpose(x)) -> transpose(x) + if (auto innerTrans = getSrc().getDefiningOp()) { + setOrder(applyPermutation(innerTrans.getOrder(), getOrder())); + setOperand(innerTrans.getSrc()); + return getResult(); + } + + // Eliminate splat constant transpose ops. + if (auto attr = + llvm::dyn_cast_if_present(adaptor.getSrc())) + return attr.reshape(getType()); + + return {}; +} + +LogicalResult TransOp::verify() { + auto order = getOrder(); + auto srcTy = cast(getSrc().getType()); + if (order.size() != srcTy.getShape().size()) { + return emitError("order must have the same size as the source tensor"); + } + if (!isPermutationOfIota(order)) { + return emitError("order must be a permutation of 0..n-1"); + } + SmallVector retShape = applyPermutation(srcTy.getShape(), order); + if (retShape != getType().getShape()) { + return emitError( + "result shape must match the permutation of the source shape"); + } + return success(); +} + +LogicalResult +TransOp::inferReturnTypes(MLIRContext *context, std::optional loc, + TransOp::Adaptor adaptor, + SmallVectorImpl &inferredReturnTypes) { + + // type is the same as the input + auto argTy = cast(adaptor.getSrc().getType()); + auto shape = argTy.getShape(); + auto order = adaptor.getOrder(); + SmallVector retShape = applyPermutation(shape, order); + + auto retEltTy = argTy.getElementType(); + Attribute argEncoding = argTy.getEncoding(); + Attribute retEncoding; + if (argEncoding) { + Dialect &dialect = argEncoding.getDialect(); + auto inferLayoutInterface = cast(&dialect); + if (failed(inferLayoutInterface->inferTransOpEncoding( + argEncoding, shape, order, retEncoding, loc))) { + return failure(); + } + } + inferredReturnTypes.push_back( + RankedTensorType::get(retShape, retEltTy, retEncoding)); + return success(); +} + +//-- DotOp -- +LogicalResult +DotOp::inferReturnTypes(MLIRContext *context, std::optional location, + ValueRange operands, DictionaryAttr attributes, + OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + // type is the same as the accumulator + auto accTy = cast(operands[2].getType()); + inferredReturnTypes.push_back(accTy); + + // verify encodings + auto aEnc = cast(operands[0].getType()).getEncoding(); + auto bEnc = cast(operands[1].getType()).getEncoding(); + auto retEnc = accTy.getEncoding(); + if (aEnc) { + assert(bEnc && retEnc); + Dialect &dialect = retEnc.getDialect(); + auto interface = cast(&dialect); + if (interface->inferDotOpEncoding(aEnc, 0, retEnc, location).failed()) + return failure(); + if (interface->inferDotOpEncoding(bEnc, 1, retEnc, location).failed()) + return failure(); + } + return success(); +} + +LogicalResult DotOp::verify() { + auto aTy = getA().getType(); + auto bTy = getB().getType(); + auto aBitWidth = aTy.getElementType().getIntOrFloatBitWidth(); + auto bBitWidth = bTy.getElementType().getIntOrFloatBitWidth(); + if (aBitWidth != bBitWidth) { + // Allow i8 * i4 mixed-precision dot for w4a8 (int4 weight, int8 activation) + if (!(aBitWidth == 8 && bBitWidth == 4)) { + return emitError( + "element types of operands A and B must have same bit width"); + } + } + auto aEncoding = aTy.getEncoding(); + auto bEncoding = bTy.getEncoding(); + if (!aEncoding && !bEncoding) + return success(); + // Verify that the encodings are valid. + if (!aEncoding || !bEncoding) + return emitError("mismatching encoding between A and B operands"); + auto accTy = getC().getType(); + auto retEnc = accTy.getEncoding(); + if (!retEnc) + return emitError("miss encoding of C operand"); + Dialect &dialect = retEnc.getDialect(); + auto interface = cast(&dialect); + return interface->verifyDotOpEncodingCompatibility(getOperation(), aEncoding, + bEncoding); +} + +bool DotOp::verifyDims() { + auto aShape = this->getA().getType().getShape(); + auto bShape = this->getB().getType().getShape(); + + return aShape[aShape.size() - 1] == bShape[aShape.size() - 2]; +} + +//-- DotScaledOp -- +bool DotScaledOp::verifyDims() { + auto aShape = this->getA().getType().getShape(); + auto bShape = this->getB().getType().getShape(); + + auto aKdim = aShape[aShape.size() - 1]; + auto bKdim = bShape[aShape.size() - 2]; + if (this->getAElemType() == ScaleDotElemType::E2M1) { + if (this->getLhsKPack()) + aKdim *= 2; + } + if (this->getBElemType() == ScaleDotElemType::E2M1) { + if (this->getRhsKPack()) + bKdim *= 2; + } + + return aKdim == bKdim; +} + +bool DotScaledOp::verifyOutputDims() { + auto cShape = this->getC().getType().getShape(); + auto oMdim = cShape[cShape.size() - 2]; + auto oNdim = cShape[cShape.size() - 1]; + auto aShape = this->getA().getType().getShape(); + auto bShape = this->getB().getType().getShape(); + auto adim = aShape[aShape.size() - 2]; + auto bdim = bShape[bShape.size() - 1]; + if (this->getAElemType() == ScaleDotElemType::E2M1) { + if (!this->getLhsKPack()) + adim *= 2; + } + if (this->getBElemType() == ScaleDotElemType::E2M1) { + if (!this->getRhsKPack()) + bdim *= 2; + } + if (adim != oMdim || bdim != oNdim) + return false; + return true; +} + +LogicalResult DotScaledOp::verify() { + auto aShape = this->getA().getType().getShape(); + int64_t rank = aShape.size(); + + auto k = aShape[rank - 1]; + if (this->getAElemType() == ScaleDotElemType::E2M1) { + if (this->getLhsKPack()) + k *= 2; + } + auto cShape = this->getC().getType().getShape(); + int64_t mDim = cShape[cShape.size() - 2]; + int64_t nDim = cShape[cShape.size() - 1]; + + if (getAScale()) { + auto aScaleShape = getAScale().getType().getShape(); + if (aScaleShape[rank - 2] != mDim) + return this->emitError( + "scales M dimension must match the operand M dimension"); + int scale_factor = + isa(getAScale().getType().getElementType()) ? 16 : 32; + if (aScaleShape[rank - 1] != k / scale_factor) + return this->emitError("scales K dimension must match the operand K " + "divided by the scale factor"); + } + if (getBScale()) { + auto bScaleShape = getBScale().getType().getShape(); + if (bScaleShape[rank - 2] != nDim) + return this->emitError( + "scales N dimension must match the operand N dimension"); + int scale_factor = + isa(getBScale().getType().getElementType()) ? 16 : 32; + if (bScaleShape[rank - 1] != k / scale_factor) + return this->emitError("scales K dimension must match the operand K " + "divided by the scale factor"); + } + return success(); +} + +//-- MakeRangeOp -- +// XPU: disable folding of `make_range(start, start+1)` to a constant splat. +// On XPU, MakeRangeOp lowering (see +// third_party/xpu/lib/Conversion/TritonXPUToLLVM/MakeRangeOpToLLVM.cpp) is +// what injects the `core_id` based offset so that each of the 64 cores in a +// cluster computes a distinct index. If we fold size-1 ranges away upstream +// (in tt.make_range), the XPU pattern never runs, no `core_id` is emitted, +// and all 64 cores in a cluster end up computing the same address. Mirrors +// the 3.0 fork where this fold is intentionally disabled. +OpFoldResult MakeRangeOp::fold(FoldAdaptor adaptor) { return {}; } + +LogicalResult MakeRangeOp::verify() { + int64_t start = getStartAttr().getInt(); + int64_t end = getEndAttr().getInt(); + if (start >= end) { + return this->emitOpError() << "start must be less than end"; + } + auto ty = getType(); + if (ty.getShape().size() != 1) { + return this->emitOpError() << "return type must be a 1D tensor"; + } + if (end - start != ty.getShape()[0]) { + return this->emitOpError() + << "number of elements in returned tensor, " << ty.getShape()[0] + << ", must match size of range [" << start << ", " << end + << "), which has " << end - start << " elements"; + } + if (!ty.getElementType().isInteger(32)) { + return this->emitOpError() << "returned tensor must have i32 elements"; + } + return success(); +} + +//-- ReduceOp -- +static LogicalResult +inferReduceReturnShape(std::optional loc, RankedTensorType argTy, + Type retEltTy, int axis, + SmallVectorImpl &inferredReturnTypes) { + auto retShape = argTy.getShape().vec(); + retShape.erase(retShape.begin() + axis); + if (retShape.empty()) { + // 0d-tensor -> scalar + inferredReturnTypes.push_back(retEltTy); + } else { + // nd-tensor where n >= 1 + // infer encoding + Attribute argEncoding = argTy.getEncoding(); + Attribute retEncoding; + if (argEncoding) { + Dialect &dialect = argEncoding.getDialect(); + auto inferLayoutInterface = cast(&dialect); + if (failed(inferLayoutInterface->inferReduceOpEncoding( + argEncoding, axis, retEncoding, loc))) { + return failure(); + } + } + // create type + inferredReturnTypes.push_back( + RankedTensorType::get(retShape, retEltTy, retEncoding)); + } + return success(); +} + +LogicalResult +ReduceOp::inferReturnTypes(MLIRContext *context, std::optional loc, + ValueRange operands, DictionaryAttr attributes, + OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + Properties *prop = properties.as(); + int axis = prop->axis.getInt(); + for (auto arg : operands) { + auto argTy = cast(arg.getType()); + auto retEltTy = argTy.getElementType(); + if (failed(inferReduceReturnShape(loc, argTy, retEltTy, axis, + inferredReturnTypes))) { + return failure(); + } + } + return success(); +} + +// Helpers for Reductions and Scans +template LogicalResult verifyReduceScan(Op &op) { + if (op.getOperands().empty()) { + return op.emitOpError() << "must have at least 1 operand"; + } + if (op.getNumOperands() != op.getNumResults()) { + return op.emitOpError() << "must have the same number of inputs as outputs"; + } + + for (auto [opElemTy, resTy] : + llvm::zip(op.getElementTypes(), op.getResultTypes())) { + if (opElemTy != getElementTypeOrSelf(resTy)) { + return op.emitOpError() << "operand types and result types must agree"; + } + } + return success(); +} + +template +static LogicalResult verifyRegionsImpl(Op &op) { + auto argElementTypes = op.getElementTypes(); + const auto &operands = op.getOperands(); + const auto numArgs = 2 * operands.size(); + auto &block = *op.getBody(); + if (block.getNumArguments() != numArgs) { + return op.emitOpError() << "nested block must take " << numArgs + << " arguments, but given block with " + << block.getNumArguments() << " arguments"; + } + const auto &blockArgTypes = block.getArgumentTypes(); + for (unsigned i = 0; i < numArgs; ++i) { + const auto &blockArgTy = blockArgTypes[i]; + const auto &argElemTy = argElementTypes[i % operands.size()]; + if (blockArgTy != argElemTy) { + return op.emitOpError() + << "type mismatch on combine operation. Expected argument " << i + << " to have type " << argElemTy << " but got " << blockArgTy; + } + } + + auto terminator = dyn_cast(block.getTerminator()); + if (!terminator) { + return op.emitOpError() + << "combine operation must be terminated " + << "with a ReduceReturnOp but got " << block.getTerminator(); + } + const auto &combineResults = terminator->getOperands(); + if (combineResults.size() != operands.size()) { + return op.emitOpError() + << "expected combine operation to return " << operands.size() + << " values but got " << combineResults.size(); + } + for (unsigned i = 0; i < combineResults.size(); ++i) { + const auto &resultTy = combineResults[i].getType(); + const auto &argElemTy = argElementTypes[i]; + if (resultTy != argElemTy) { + return op.emitOpError() + << "type mismatch on combine operation. Expected argument " << i + << " to have type " << argElemTy << " but got " << resultTy; + } + } + return success(); +} + +static llvm::SmallVector +getInputTypesImpl(const Operation::operand_range &operands) { + llvm::SmallVector srcTys; + srcTys.reserve(operands.size()); + for (const auto &ty : operands.getTypes()) { + srcTys.push_back(cast(ty)); + } + return srcTys; +} + +template +static llvm::SmallVector getElementTypesImpl(const ValueRange &operands) { + llvm::SmallVector srcElemTys; + srcElemTys.reserve(operands.size()); + for (const auto &op : operands) { + srcElemTys.push_back(cast(op.getType()).getElementType()); + } + return srcElemTys; +} + +LogicalResult ReduceOp::verify() { return verifyReduceScan(*this); } + +LogicalResult ReduceOp::verifyRegions() { + return verifyRegionsImpl(*this); +} + +llvm::SmallVector ReduceOp::getInputTypes() { + return getInputTypesImpl(this->getOperands()); +} + +llvm::SmallVector ReduceOp::getElementTypes() { + return getElementTypesImpl(this->getOperands()); +} + +::mlir::Operation *ReduceOp::getSingleCombiner() { + if (getNumOperands() != 1 || getNumResults() != 1) + return nullptr; + Block *block = &(*getCombineOp().begin()); + Operation *yield = block->getTerminator(); + Operation *reduceOp = yield->getOperand(0).getDefiningOp(); + if (!reduceOp || reduceOp->getNumOperands() != 2 || + reduceOp->getNumResults() != 1) + return nullptr; + if (reduceOp->getOperand(0) != block->getArgument(0) || + reduceOp->getOperand(1) != block->getArgument(1)) + return nullptr; + + return reduceOp; +} + +unsigned ReduceOp::getNumOperands() { return this->getOperands().size(); } + +//-- ScanOp -- +void ScanOp::build(OpBuilder &builder, OperationState &state, + ValueRange operands, int axis, bool reverse) { + SmallVector inferredReturnTypes; + for (auto arg : operands) + inferredReturnTypes.push_back(arg.getType()); + ScanOp::build(builder, state, inferredReturnTypes, operands, axis, reverse); +} + +LogicalResult +ScanOp::inferReturnTypes(MLIRContext *context, std::optional location, + ValueRange operands, DictionaryAttr attributes, + OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + for (auto arg : operands) + inferredReturnTypes.push_back(arg.getType()); + return success(); +} + +LogicalResult ScanOp::verify() { return verifyReduceScan(*this); } + +LogicalResult ScanOp::verifyRegions() { + return verifyRegionsImpl(*this); +} + +llvm::SmallVector ScanOp::getInputTypes() { + return getInputTypesImpl(this->getOperands()); +} + +llvm::SmallVector ScanOp::getElementTypes() { + return getElementTypesImpl(this->getOperands()); +} + +unsigned ScanOp::getNumOperands() { return this->getOperands().size(); } + +//-- MapElementwiseOp +LogicalResult MapElementwiseOp::verify() { + if (getOperands().empty()) { + return emitOpError() << "MapElementwiseOp must have at least 1 operand"; + } + if (!llvm::isPowerOf2_32(getPack())) { + return emitOpError() << "Pack must be a power of 2"; + } + return success(); +} + +template +SmallVector repeatInterleave(const SmallVectorImpl &vs, int nRepeat) { + SmallVector result; + result.reserve(vs.size() * nRepeat); + for (auto v : vs) + for (auto _ : llvm::seq(nRepeat)) + result.push_back(v); + return result; +} + +LogicalResult MapElementwiseOp::verifyRegions() { + // Verify signature + auto *firstBlock = &getRegion().getBlocks().front(); + if (firstBlock->getNumArguments() != getNumOperands() * getPack()) { + return emitOpError() << "region has wrong number of arguments"; + } + + auto expectedArgTypes = + repeatInterleave(getElementTypesImpl(getOperands()), getPack()); + if (firstBlock->getArgumentTypes() != expectedArgTypes) { + return emitError() << "argument types did not match"; + } + auto expectedReturnTypes = + repeatInterleave(getElementTypesImpl(getResults()), getPack()); + auto walkRes = getRegion().walk([&](Operation *op) -> WalkResult { + auto memEffects = dyn_cast(op); + // Ban stores as we won't get the redundant masking correct by treating it + // as a scalar. + if (memEffects && memEffects.hasEffect()) { + return op->emitOpError() + << "Stores are not supported inside map_elementwise"; + } + if (isa(op) && + op->getOperandTypes() != expectedReturnTypes) { + return op->emitError() + << "region return does not match map_elementwise result"; + } + return WalkResult::advance(); + }); + return success(!walkRes.wasInterrupted()); +} + +//-- SplatOp -- +OpFoldResult SplatOp::fold(FoldAdaptor adaptor) { + auto value = adaptor.getSrc(); + if (!value) + return {}; + if (!isa(value)) + return {}; + auto shapedType = cast(getType()); + auto ret = SplatElementsAttr::get(shapedType, ArrayRef(value)); + return ret; +} + +//-- UnsplatOp -- +LogicalResult UnsplatOp::verify() { + auto srcShape = getSrc().getType().getShape(); + if (product(srcShape) != 1) { + return emitError("source tensor must have exactly one element"); + } + return success(); +} + +LogicalResult UnsplatOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + auto dstTy = cast(operands[0].getType()).getElementType(); + inferredReturnTypes.push_back(dstTy); + return success(); +} + +//-- ExpandDimsOp -- +LogicalResult ExpandDimsOp::inferReturnTypes( + MLIRContext *context, std::optional loc, ValueRange operands, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + // infer shape + auto arg = operands[0]; + auto argTy = cast(arg.getType()); + auto retShape = argTy.getShape().vec(); + Properties *prop = properties.as(); + int axis = prop->axis.getInt(); + retShape.insert(retShape.begin() + axis, 1); + // infer encoding + Attribute argEncoding = argTy.getEncoding(); + Attribute retEncoding; + if (argEncoding) { + Dialect &dialect = argEncoding.getDialect(); + auto inferLayoutInterface = cast(&dialect); + if (failed(inferLayoutInterface->inferExpandDimsOpEncoding( + argEncoding, axis, retEncoding, loc))) + return emitOptionalError(loc, "failed to infer layout for ExpandDimsOp"); + } + // create type + auto argEltTy = argTy.getElementType(); + inferredReturnTypes.push_back( + RankedTensorType::get(retShape, argEltTy, retEncoding)); + return success(); +} + +LogicalResult ExpandDimsOp::canonicalize(ExpandDimsOp op, + PatternRewriter &rewriter) { + auto definingOp = op.getSrc().getDefiningOp(); + if (!definingOp) { + return failure(); + } + // expand_dims(splat) -> splat + if (auto splat = dyn_cast(definingOp)) { + rewriter.replaceOpWithNewOp(op, op.getType(), splat.getSrc()); + return success(); + } + // expand_dims(broadcast(x)) -> broadcast(expand_dims(x)) + // + // On its own this doesn't do much, but consider + // broadcast(expand_dims(broadcast)) + // -> broadcast(broadcast(expand_dims)) + // -> broadcast(expand_dims) + if (auto broadcast = dyn_cast(definingOp)) { + auto src = broadcast.getSrc(); + auto srcTy = src.getType(); + SmallVector newExpandShape(srcTy.getShape()); + newExpandShape.insert(newExpandShape.begin() + op.getAxis(), 1); + + // Infer the encoding of the new expand op, if encodings are present. + Attribute newExpandEnc; + if (auto srcEnc = srcTy.getEncoding()) { + Dialect &dialect = srcEnc.getDialect(); + auto inferLayoutInterface = cast(&dialect); + if (failed(inferLayoutInterface->inferExpandDimsOpEncoding( + srcEnc, op.getAxis(), newExpandEnc, op.getLoc()))) { + return emitOptionalError(op.getLoc(), + "failed to infer layout for ExpandDimsOp"); + } + } + + auto newExpandTy = RankedTensorType::get( + newExpandShape, srcTy.getElementType(), newExpandEnc); + auto newExpand = ExpandDimsOp::create(rewriter, op.getLoc(), newExpandTy, + src, op.getAxis()); + auto newBroadcast = BroadcastOp::create( + rewriter, broadcast.getLoc(), op.getType(), newExpand.getResult()); + rewriter.replaceOp(op, {newBroadcast.getResult()}); + return success(); + } + + return failure(); +} + +template +static OpFoldResult foldViewLikeOp(ViewLikeOp op, Attribute value) { + if (!value) + return {}; + + auto shapedType = cast(op.getType()); + if (auto denseElemsAttr = dyn_cast(value)) { + if (denseElemsAttr.isSplat()) { + return denseElemsAttr.resizeSplat(shapedType); + } else { + return denseElemsAttr.reshape(shapedType); + } + } + return {}; +} + +OpFoldResult ExpandDimsOp::fold(FoldAdaptor adaptor) { + return foldViewLikeOp(*this, adaptor.getSrc()); +} + +//-- ReshapeOp -- + +void ReshapeOp::build(OpBuilder &builder, OperationState &state, + ArrayRef shape, Value src, bool allowReorder) { + auto srcTy = cast(src.getType()); + auto srcEnc = srcTy.getEncoding(); + Attribute dstEnc; + if (srcEnc) { + auto result = cast(&srcEnc.getDialect()) + ->inferReshapeOpEncoding(srcTy.getShape(), srcEnc, shape, + dstEnc, state.location); + assert(succeeded(result)); + } + auto dstTy = RankedTensorType::get(shape, srcTy.getElementType(), dstEnc); + build(builder, state, dstTy, src, allowReorder); +} + +LogicalResult ReshapeOp::canonicalize(ReshapeOp op, PatternRewriter &rewriter) { + if (op.getEfficientLayout()) + return failure(); + + auto definingOp = op.getSrc().getDefiningOp(); + if (!definingOp) { + return failure(); + } + + // reshape(reshape) -> reshape + if (auto parentReshape = dyn_cast(definingOp)) { + // Allow reorder if either reshape allowed it + const bool allowReorder = + (op.getAllowReorder() || parentReshape.getAllowReorder()); + rewriter.replaceOpWithNewOp(op, op.getType(), + parentReshape.getSrc(), allowReorder, + op.getEfficientLayout()); + return success(); + } + + // reshape(splat) -> splat + if (auto splat = dyn_cast(definingOp)) { + rewriter.replaceOpWithNewOp(op, op.getType(), splat.getSrc()); + return success(); + } + + return failure(); +} + +OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) { + //===-------------------- For Triton XPU -----------------------===// + // Restore the 3.0 fold behavior. Upstream 3.6 (commit 868561dd) added the + // `&& !getAllowReorder()` guard, which prevents folding an identity reshape + // (src type == dst type) that carries allow_reorder=true. An identity reshape + // between two *identical* types has nothing to reorder, so folding it is + // always safe. Without the fold, tl.sum(axis=None)'s identity reshape + // survives into the XPU pipeline; CoreTiling then assigns it a different + // ClusterLayout than its source, and isExpensiveView rejects the + // elems-per-thread mismatch ("expensive view not supported on reshape op"). + // Non-identity reshapes still use foldViewLikeOp below, so this only affects + // the degenerate src-type == dst-type case (safe for all backends). + if (getType() == getSrc().getType()) { + // no-op + return getSrc(); + } + //===-----------------------------------------------------------===// + + return foldViewLikeOp(*this, adaptor.getSrc()); +} + +LogicalResult ReshapeOp::verify() { + auto dstTy = getType(); + auto srcTy = getSrc().getType(); + if (getType().getNumElements() != srcTy.getNumElements()) { + return emitError( + "number of src and dst elements of reshape must be the same"); + } + + Attribute srcEnc = srcTy.getEncoding(); + Attribute dstEnc = dstTy.getEncoding(); + if (!!srcEnc != !!dstEnc) { + return emitError("Op requires that either (a) src and dst both have " + "encodings, or (b) neither does."); + } + + if (!srcEnc || getAllowReorder()) { + return success(); + } + + // Check that we can infer the dst encoding from the src encoding + // and that the inferred dst encoding is the same as the given dst encoding + Attribute inferredDstEnc; + auto layoutInterface = + cast(&srcEnc.getDialect()); + auto result = layoutInterface->inferReshapeOpEncoding( + srcTy.getShape(), srcEnc, dstTy.getShape(), inferredDstEnc, getLoc()); + if (failed(result)) + return failure(); + return layoutInterface->verifyLayoutsAreEqual( + dstTy.getShape(), inferredDstEnc, dstEnc, getLoc()); +} + +//-- FpToFpOp -- + +// Fold FpToFpOp when the input operand is a constant zero. +OpFoldResult FpToFpOp::fold(FoldAdaptor adaptor) { + auto srcVal = getSrc(); + auto dstTy = getType(); + // Fold trivial cast + if (srcVal.getType() == dstTy) { + return srcVal; + } + + auto resElemType = cast(getElementTypeOrSelf(getType())); + const llvm::fltSemantics &semantic = resElemType.getFloatSemantics(); + + if (matchPattern(srcVal, m_PosZeroFloat())) { + llvm::APFloat posZero = + llvm::APFloat::getZero(semantic, /*negative=*/false); + if (auto tensorTy = dyn_cast(dstTy)) + return DenseElementsAttr::get(tensorTy, posZero); + return Builder(getContext()).getFloatAttr(resElemType, posZero); + } + + if (matchPattern(srcVal, m_NegZeroFloat())) { + llvm::APFloat negZero = llvm::APFloat::getZero(semantic, /*negative=*/true); + if (auto tensorTy = dyn_cast(dstTy)) + return DenseElementsAttr::get(tensorTy, negZero); + return Builder(getContext()).getFloatAttr(resElemType, negZero); + } + + return {}; +} + +LogicalResult FpToFpOp::verify() { + auto dstType = getType(); + auto srcType = getSrc().getType(); + if (auto dstTensorType = dyn_cast(dstType)) + dstType = dstTensorType.getElementType(); + if (auto srcTensorType = dyn_cast(srcType)) + srcType = srcTensorType.getElementType(); + if ((dstType.getIntOrFloatBitWidth() < srcType.getIntOrFloatBitWidth()) && + (!getRounding().has_value())) { + return emitError("Rounding mode is required for FP downcast"); + } + return success(); +} + +//-- BitcastOp -- +LogicalResult BitcastOp::verify() { + // Bitcast only allows conversion between types with the same bit width. + Type dstType = getType(); + Type srcType = getSrc().getType(); + // Strip tensor shapes; SameOperandsAndResultShape guarantees shapes match. + if (auto dstTensorType = dyn_cast(dstType)) + dstType = dstTensorType.getElementType(); + if (auto srcTensorType = dyn_cast(srcType)) + srcType = srcTensorType.getElementType(); + // XPU backend lowers bitwise ops via bitcast on tensor> + // (nested vector element). For such vector-to-vector bitcast the scalar + // element types may differ (e.g. vector<32xi16> <-> vector<16xi32>); only + // the total bit width must match, so handle vectors before the scalar + // getIntOrFloatBitWidth() path below. + if (auto dstVecType = dyn_cast(dstType)) { + if (auto srcVecType = dyn_cast(srcType)) { + unsigned dstBits = dstVecType.getNumElements() * + dstVecType.getElementType().getIntOrFloatBitWidth(); + unsigned srcBits = srcVecType.getNumElements() * + srcVecType.getElementType().getIntOrFloatBitWidth(); + if (dstBits != srcBits) { + return emitError("Cannot bitcast data-type of size ") + << srcBits << " to data-type of size " << dstBits; + } + return success(); + } + return emitError("Cannot bitcast vector to non-vector type"); + } + bool dstIsPtr = isa(dstType); + bool srcIsPtr = isa(srcType); + if (dstIsPtr || srcIsPtr) { + // Bitcast supports pointer-to-pointer conversions but not + // pointer-to-scalar. + if (dstIsPtr && srcIsPtr) { + if (triton::getAddressSpace(dstType) != triton::getAddressSpace(srcType)) + return emitError( + "Cannot bitcast pointer between different address spaces"); + return success(); + } + return emitError("Cannot bitcast pointer to non-pointer type"); + } + unsigned dstBits = dstType.getIntOrFloatBitWidth(); + unsigned srcBits = srcType.getIntOrFloatBitWidth(); + if (dstBits != srcBits) { + return emitError("Cannot bitcast data-type of size ") + << srcBits << " to data-type of size " << dstBits; + } + return success(); +} + +//-- BroadcastOp -- +void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.add(context); +} + +OpFoldResult BroadcastOp::fold(FoldAdaptor adaptor) { + if (getType() == getSrc().getType()) { + // no-op + return getSrc(); + } + + auto value = adaptor.getSrc(); + if (!value) + return {}; + + if (auto denseElemsAttr = dyn_cast(value)) { + auto shapedType = cast(getType()); + return denseElemsAttr.resizeSplat(shapedType); + } + return {}; +} + +LogicalResult BroadcastOp::verify() { + auto src = getSrc(); + auto srcTensorType = cast(src.getType()); + auto srcShape = srcTensorType.getShape(); + auto result = getResult(); + auto resultTensorType = cast(result.getType()); + auto resultShape = resultTensorType.getShape(); + if (srcShape.size() != resultShape.size()) { + return emitError("rank of source must be same as rank of result"); + } + for (size_t i = 0; i < srcShape.size(); i++) { + if (srcShape[i] != 1 && srcShape[i] != resultShape[i]) { + return emitError("Different dimensions at index ") + << i << " between source and result. " + << "Broadcast requires the source dimension to be 1."; + } + } + return success(); +} + +//-- MakeTensorPtrOp -- +void MakeTensorPtrOp::build(OpBuilder &builder, OperationState &state, + Value base, ValueRange shape, ValueRange strides, + ValueRange offsets, ArrayRef tensorShape, + ArrayRef order) { + // Get pointer type from `base` + auto pointerType = cast(base.getType()); + assert(pointerType != nullptr); + + // Build type `tt.ptr>` + auto tensorType = RankedTensorType::get( + SmallVector(tensorShape.begin(), tensorShape.end()), + pointerType.getPointeeType()); + auto result = PointerType::get(tensorType, pointerType.getAddressSpace()); + + return build(builder, state, result, base, shape, strides, offsets, + builder.getDenseI32ArrayAttr(order)); +} + +//-- AddPtrOp -- +OpFoldResult AddPtrOp::fold(FoldAdaptor adaptor) { + // addptr(ptr, 0) -> ptr + if (matchPattern(adaptor.getOffset(), m_Zero())) { + return getPtr(); + } + return {}; +} + +//-- AdvanceOp -- +OpFoldResult AdvanceOp::fold(FoldAdaptor adaptor) { + // advance(ptr, 0, 0) -> ptr + SmallVector rawOffsets = getOffsets(); + auto offsets = getConstantIntValues(rawOffsets); + if (!offsets.has_value()) + return {}; + for (int64_t offset : offsets.value()) + if (offset != 0) + return {}; + return getPtr(); +} + +//-- MakeTensorDescOp -- +void MakeTensorDescOp::build(OpBuilder &builder, OperationState &state, + Value base, ValueRange shape, ValueRange strides, + ArrayRef blockShape, bool isSignedInteger, + triton::PaddingOption padding) { + auto ptrTy = dyn_cast(base.getType()); + if (!ptrTy) { + llvm::report_fatal_error("Expected pointer type"); + } + auto elemTy = ptrTy.getPointeeType(); + SmallVector blockShape64(blockShape); + auto blockTy = RankedTensorType::get(blockShape64, elemTy); + auto descTy = + TensorDescType::get(builder.getContext(), blockTy, isSignedInteger); + auto paddingAttr = PaddingOptionAttr::get(builder.getContext(), padding); + return build(builder, state, descTy, base, shape, strides, paddingAttr); +} + +// The following ops, including `call`, `func`, and `return` are copied and +// modified from +// https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Func/IR/FuncOps.cpp +// We could revert it back once MLIR has a better inliner interface. +//-- FuncOp -- +void FuncOp::build(OpBuilder &builder, OperationState &state, StringRef name, + FunctionType type, ArrayRef attrs, + ArrayRef argAttrs) { + state.addAttribute(SymbolTable::getSymbolAttrName(), + builder.getStringAttr(name)); + state.addAttribute(getFunctionTypeAttrName(state.name), TypeAttr::get(type)); + state.attributes.append(attrs.begin(), attrs.end()); + state.addRegion(); + + if (argAttrs.empty()) + return; + assert(type.getNumInputs() == argAttrs.size()); + call_interface_impl::addArgAndResultAttrs( + builder, state, argAttrs, /*resultAttrs=*/{}, + getArgAttrsAttrName(state.name), getResAttrsAttrName(state.name)); +} + +ParseResult FuncOp::parse(OpAsmParser &parser, OperationState &result) { + auto buildFuncType = + [](Builder &builder, ArrayRef argTypes, ArrayRef results, + function_interface_impl::VariadicFlag, + std::string &) { return builder.getFunctionType(argTypes, results); }; + + return function_interface_impl::parseFunctionOp( + parser, result, /*allowVariadic=*/false, + getFunctionTypeAttrName(result.name), buildFuncType, + getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name)); +} + +void FuncOp::print(OpAsmPrinter &printer) { + function_interface_impl::printFunctionOp( + printer, *this, /*isVariadic=*/false, getFunctionTypeAttrName(), + getArgAttrsAttrName(), getResAttrsAttrName()); +} + +// -- CallOp -- +LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) { + // Check that the callee attribute was specified. + auto fnAttr = (*this).getProperties().callee; + if (!fnAttr) + return emitOpError("requires a 'callee' symbol reference attribute"); + FuncOp fn = symbolTable.lookupNearestSymbolFrom(*this, fnAttr); + if (!fn) + return emitOpError() << "'" << fnAttr.getValue() + << "' does not reference a valid function"; + + // Verify that the operand and result types match the callee. + auto fnType = fn.getFunctionType(); + if (fnType.getNumInputs() != getNumOperands()) + return emitOpError("incorrect number of operands for callee"); + + for (unsigned i = 0, e = fnType.getNumInputs(); i != e; ++i) + if (getOperand(i).getType() != fnType.getInput(i)) + return emitOpError("operand type mismatch: expected operand type ") + << fnType.getInput(i) << ", but provided " + << getOperand(i).getType() << " for operand number " << i; + + if (fnType.getNumResults() != getNumResults()) + return emitOpError("incorrect number of results for callee"); + + for (unsigned i = 0, e = fnType.getNumResults(); i != e; ++i) + if (getResult(i).getType() != fnType.getResult(i)) { + auto diag = emitOpError("result type mismatch at index ") << i; + diag.attachNote() << " op result types: " << getResultTypes(); + diag.attachNote() << "function result types: " << fnType.getResults(); + return diag; + } + + return success(); +} + +// -- ReturnOp -- +LogicalResult ReturnOp::verify() { + auto function = cast((*this)->getParentOp()); + + // The operand number and types must match the function signature. + const auto &results = function.getFunctionType().getResults(); + if (getNumOperands() != results.size()) + return emitOpError("has ") + << getNumOperands() << " operands, but enclosing function (@" + << function.getName() << ") returns " << results.size(); + + for (unsigned i = 0, e = results.size(); i != e; ++i) + if (getOperand(i).getType() != results[i]) + return emitError() << "type of return operand " << i << " (" + << getOperand(i).getType() + << ") doesn't match function result type (" + << results[i] << ")" + << " in function @" << function.getName(); + + return success(); +} + +// -- JoinOp -- + +void JoinOp::build(OpBuilder &builder, OperationState &state, Value lhs, + Value rhs) { + auto lhsTy = cast(lhs.getType()); + SmallVector retShape(lhsTy.getShape()); + retShape.push_back(2); + + Attribute srcEnc = lhsTy.getEncoding(); + Attribute retEnc; + if (srcEnc) { + if (failed(cast(&srcEnc.getDialect()) + ->inferDefaultJoinOpEncoding( + srcEnc, retEnc, lhsTy.getShape(), state.location))) { + llvm_unreachable("failed to infer join encoding"); + } + } + auto retTy = RankedTensorType::get(retShape, lhsTy.getElementType(), retEnc); + JoinOp::build(builder, state, retTy, lhs, rhs); +} + +LogicalResult JoinOp::verify() { + RankedTensorType srcTy = getLhs().getType(); + SmallVector retShape(srcTy.getShape()); + retShape.push_back(2); + + RankedTensorType retTy = getType(); + if (SmallVector(retTy.getShape()) != retShape) { + return emitOpError("result shape must be (") + << retShape << "), but got " << retTy.getShape(); + } + if (retTy.getElementType() != srcTy.getElementType()) { + return emitOpError("result element type must match the input element type"); + } + Attribute retEnc = retTy.getEncoding(); + if (!retEnc) { + if (srcTy.getEncoding()) { + return emitOpError("result encoding must be specified"); + } + return success(); + } + // There are multiple correct destination layout for a given source layout but + // there is only one correct source layout for a given destination layout. So + // we verify that the source layout match the destination layout. + Attribute srcEnc; + Location location = getLoc(); + if (cast(&retEnc.getDialect()) + ->inferSplitOpEncoding(retEnc, srcEnc, retShape, location) + .failed()) { + return failure(); + } + + if (cast(&srcEnc.getDialect()) + ->verifyLayoutsAreEqual(srcTy.getShape(), srcEnc, srcTy.getEncoding(), + {}) + .failed()) { + return emitOpError("incompatible join layout"); + } + return success(); +} + +// -- SplitOp -- +LogicalResult SplitOp::inferReturnTypes( + MLIRContext *context, std::optional location, + SplitOp::Adaptor adaptor, SmallVectorImpl &inferredReturnTypes) { + auto srcTy = cast(adaptor.getSrc().getType()); + auto srcShape = srcTy.getShape(); + + if (srcShape.empty() || srcShape.back() != 2) { + return emitOptionalError(location, + "last dimension of input tensor must be 2"); + } + ArrayRef retShape(srcShape.begin(), srcShape.end() - 1); + + Attribute srcEnc = srcTy.getEncoding(); + Attribute retEnc; + if (srcEnc) { + if (cast(&srcEnc.getDialect()) + ->inferSplitOpEncoding(srcEnc, retEnc, srcTy.getShape(), location) + .failed()) { + return failure(); + } + } + auto retTy = RankedTensorType::get(retShape, srcTy.getElementType(), retEnc); + inferredReturnTypes.push_back(retTy); + inferredReturnTypes.push_back(retTy); + return success(); +} + +// -- ElementwiseInlineAsmOp -- +void ElementwiseInlineAsmOp::getEffects( + SmallVectorImpl> + &effects) { + if (getPure()) + return; + effects.emplace_back(MemoryEffects::Write::get()); + effects.emplace_back(MemoryEffects::Read::get()); +} + +Speculation::Speculatability ElementwiseInlineAsmOp::getSpeculatability() { + if (getPure()) + return Speculation::Speculatable; + return Speculation::NotSpeculatable; +} + +LogicalResult ElementwiseInlineAsmOp::verify() { + if (getNumOperands() >= 1) { + auto tensorType = dyn_cast(getOperand(0).getType()); + size_t numInputElems = tensorType ? tensorType.getNumElements() : 0; + if (numInputElems % this->getPackedElement() != 0) { + return emitError("number of input elements ") + << numInputElems + << " must be a multiple of the op's packed_element attribute, " + << getPackedElement(); + } + } + return success(); +} + +// -- ExternElementwiseOp -- +void ExternElementwiseOp::getEffects( + SmallVectorImpl> + &effects) { + if (getPure()) + return; + effects.emplace_back(MemoryEffects::Write::get()); + effects.emplace_back(MemoryEffects::Read::get()); +} + +Speculation::Speculatability ExternElementwiseOp::getSpeculatability() { + if (getPure()) + return Speculation::Speculatable; + return Speculation::NotSpeculatable; +} + +// -- GatherOp -- +LogicalResult GatherOp::verify() { + RankedTensorType indicesTy = getIndices().getType(); + RankedTensorType srcTy = getSrc().getType(); + RankedTensorType resTy = getResult().getType(); + + if (indicesTy.getShape() != resTy.getShape()) { + return emitOpError("indices and output shapes must match"); + } + if (indicesTy.getEncoding() != resTy.getEncoding()) { + return emitOpError("indices and output encodings must match"); + } + if (srcTy.getElementType() != resTy.getElementType()) { + return emitOpError("input and output element types must match"); + } + if (srcTy.getRank() != indicesTy.getRank()) { + return emitOpError("input and indices ranks must match"); + } + if (getAxis() >= srcTy.getRank()) { + return emitOpError("gather dimension must be less than the input rank"); + } + for (uint32_t dim = 0; dim < indicesTy.getRank(); ++dim) { + if (dim == getAxis()) + continue; + if (indicesTy.getShape()[dim] != srcTy.getShape()[dim]) { + return emitOpError("indices dimension ") + << dim << " must match the corresponding input dimension"; + } + } + + return success(); +} + +LogicalResult GatherOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + GatherOpAdaptor adaptor(operands, attributes, properties, regions); + auto indicesType = cast(adaptor.getIndices().getType()); + auto srcType = cast(adaptor.getSrc().getType()); + + // Shape and encoding of the indices with the element type of the src. + inferredReturnTypes.push_back(indicesType.clone(srcType.getElementType())); + return success(); +} + +// -- DescriptorGatherOp +LogicalResult +DescriptorGatherOp::verifyResultType(Operation *op, ShapedType resultType, + RankedTensorType indicesType) { + if (indicesType.getRank() != 1) + return op->emitOpError("x offsets must be a 1D tensor, but got ") + << indicesType; + if (resultType.getRank() != 2) + return op->emitOpError("result must be a 2D tensor, but got ") + << resultType; + + // The swizzling of TMA accesses matches that of the MMAv3 shared memory + // layouts. However, these have minimum size requirements. + // TODO: We can support smaller gather sizes by padding the `local_alloc` this + // lowers to to the nearest minimum tile size. + if (unsigned rows = resultType.getShape()[0]; rows < 8) { + return op->emitOpError("gather must have at least 8 rows, but got ") + << rows; + } + + Type dtype = resultType.getElementType(); + if (dtype.getIntOrFloatBitWidth() > 32) + return op->emitOpError("TMA dtype cannot be greater than 32 bits"); + + unsigned minCols = 32 / dtype.getIntOrFloatBitWidth() * 8; + if (unsigned cols = resultType.getShape()[1]; cols < minCols) { + return op->emitOpError("gather of ") + << dtype << " must have at least " << minCols << " columns, but got " + << cols; + } + + if (resultType.getShape()[0] != indicesType.getShape()[0]) { + return op->emitOpError("result tensor must have as many rows as indices (") + << indicesType.getShape()[0] << "), but got " << resultType; + } + + return success(); +} + +static LogicalResult verifyGatherScatterOp(Operation *op, + RankedTensorType blockType, + RankedTensorType resultType, + RankedTensorType indicesType) { + // Gather from `!tt.tensordesc>`. + if (blockType.getRank() != 2) { + return op->emitOpError("block must be a 2D tensor, but got ") << blockType; + } + if (blockType.getShape()[0] != 1) { + return op->emitOpError("block must have exactly 1 row, but got ") + << blockType; + } + + // With x offsets `tensor` into `tensor`. + if (failed(DescriptorGatherOp::verifyResultType(op, resultType, indicesType))) + return failure(); + + if (resultType.getShape()[1] != blockType.getShape()[1]) { + return op->emitOpError("result tensor number of columns must match block (") + << blockType.getShape()[1] << "), but got " << resultType; + } + if (resultType.getElementType() != blockType.getElementType()) { + return op->emitOpError("result tensor element type must match block (") + << blockType.getElementType() << "), but got " << resultType; + } + + return success(); +} + +LogicalResult DescriptorGatherOp::verify() { + return verifyGatherScatterOp(*this, + getDesc().getType().getSignlessBlockType(), + getResult().getType(), getXOffsets().getType()); +} + +// -- DescriptorScatterOp -- +LogicalResult DescriptorScatterOp::verify() { + return verifyGatherScatterOp(*this, + getDesc().getType().getSignlessBlockType(), + getSrc().getType(), getXOffsets().getType()); +} + +// -- DescriptorLoadOp -- +static LogicalResult verifyDescriptorLoadStoreType(Operation *op, + TensorDescType desc, + RankedTensorType tensor) { + RankedTensorType block = desc.getSignlessBlockType(); + ArrayRef blockShape = block.getShape(); + ArrayRef tensorShape = tensor.getShape(); + if (blockShape.size() > tensorShape.size()) { + // Allow ranked reduced load if the leading dimensions are all 1s. + for (int i = 0; i < blockShape.size() - tensorShape.size(); ++i) { + if (blockShape[i] != 1) + return op->emitOpError( + "ranked reduce load only allowed for unit dimension leading dim."); + } + blockShape = blockShape.take_back(tensorShape.size()); + } + + if (blockShape == tensorShape && + block.getElementType() == tensor.getElementType()) + return success(); + return op->emitOpError("tensor descriptor block and tensor types must match"); +} + +LogicalResult DescriptorLoadOp::verify() { + return verifyDescriptorLoadStoreType(*this, getDesc().getType(), getType()); +} + +// -- DescriptorStoreOp -- +LogicalResult DescriptorStoreOp::verify() { + return verifyDescriptorLoadStoreType(*this, getDesc().getType(), + getSrc().getType()); +} + +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/backend/spec/lib/Dialect/Triton/IR/Traits.cpp b/third_party/xpu/backend/spec/lib/Dialect/Triton/IR/Traits.cpp new file mode 100644 index 0000000000..8b0d3ba0a9 --- /dev/null +++ b/third_party/xpu/backend/spec/lib/Dialect/Triton/IR/Traits.cpp @@ -0,0 +1,268 @@ +#include "triton/Dialect/Triton/IR/Traits.h" + +#include + +#include "mlir/IR/TypeUtilities.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Types.h" +#include "llvm/Support/ErrorHandling.h" + +using namespace mlir; +using namespace mlir::triton::gpu; + +LogicalResult OpTrait::impl::verifyEquivalentType(Type typeA, Type typeB) { + auto memdescA = dyn_cast(typeA); + auto memdescB = dyn_cast(typeB); + if (memdescA || memdescB) { + if (!memdescA || !memdescB) + return failure(); + if (memdescA.getShape() != memdescB.getShape()) + return failure(); + if (memdescA.getAllocShape() != memdescB.getAllocShape()) + return failure(); + if (memdescA.getElementType() != memdescB.getElementType()) + return failure(); + if (memdescA.getMemorySpace() != memdescB.getMemorySpace()) + return failure(); + if (memdescA.getMutableMemory() != memdescB.getMutableMemory()) + return failure(); + + Attribute encodingA = memdescA.getEncoding(); + Attribute encodingB = memdescB.getEncoding(); + if (encodingA == encodingB) + return success(); + if (static_cast(encodingA) != static_cast(encodingB)) + return failure(); + + auto layoutInterface = + cast(&encodingA.getDialect()); + return layoutInterface->verifyLayoutsAreEqual(memdescA.getShape(), + encodingA, encodingB, {}); + } + auto tensorTypeA = dyn_cast(typeA); + auto tensorTypeB = dyn_cast(typeB); + if (!(bool(tensorTypeA) && bool(tensorTypeB))) + return typeA == typeB ? success() : failure(); + auto encodingA = tensorTypeA.getEncoding(); + auto encodingB = tensorTypeB.getEncoding(); + auto shapeA = tensorTypeA.getShape(); + auto shapeB = tensorTypeB.getShape(); + if (shapeA != shapeB) + return failure(); + if (tensorTypeA.getElementType() != tensorTypeB.getElementType()) + return failure(); + // If there's no encoding or the encodings are the same + if (encodingA == encodingB) + return success(); + if (bool(encodingA) != bool(encodingB)) + return failure(); + + return cast(&encodingA.getDialect()) + ->verifyLayoutsAreEqual(shapeA, encodingA, encodingB, {}); +} + +static LogicalResult verifySameEncoding(Type typeA, Type typeB, + bool allowTensorPointerType) { + // TODO(Keren): the allowTensorPointerType argument is a hack to allow. + // The type checking code is kind of a mess with the current design. + auto getEncoding = [=](Type type) -> Attribute { + Attribute ret; + if (auto tensorType = dyn_cast(type)) { + ret = tensorType.getEncoding(); + } + if (!allowTensorPointerType) { + assert(!triton::isTensorPointerType(type)); + } + return ret; + }; + auto encodingA = getEncoding(typeA); + auto encodingB = getEncoding(typeB); + if (!encodingA || !encodingB) + return success(); + return encodingA == encodingB ? success() : failure(); +} + +LogicalResult +OpTrait::impl::verifySameOperandsEncoding(Operation *op, + bool allowTensorPointerType) { + if (failed(verifyAtLeastNOperands(op, 1))) + return failure(); + + auto type = op->getOperand(0).getType(); + for (auto opType : llvm::drop_begin(op->getOperandTypes(), 1)) + if (failed(verifySameEncoding(opType, type, allowTensorPointerType))) + return op->emitOpError() << "requires the same encoding for all operands"; + + return success(); +} + +LogicalResult OpTrait::impl::verifySameOperandsAndResultEncoding( + Operation *op, bool allowTensorPointerType) { + if (op->getNumOperands() == 0) + return success(); + + if (failed(verifyAtLeastNOperands(op, 1)) || + failed(verifyAtLeastNResults(op, 1))) + return failure(); + + auto type = op->getOperand(0).getType(); + for (auto resultType : op->getResultTypes()) + if (failed(verifySameEncoding(resultType, type, allowTensorPointerType))) + return op->emitOpError() + << "requires the same encoding for all operands and results"; + + return verifySameOperandsEncoding(op, allowTensorPointerType); +} + +LogicalResult OpTrait::impl::verifyTensorSize(Operation *op) { + for (auto opType : op->getOperandTypes()) { + if (auto tensorType = dyn_cast(opType)) { + int64_t numElements = 1; + for (int64_t s : tensorType.getShape()) + numElements *= s; + if (numElements > maxTensorNumElements) + return op->emitError("Maximum allowed number of elements is ") + << maxTensorNumElements << ", but " << *op + << " has more than that"; + //===-------------------- For Triton XPU -----------------------===// + // [internal] Triton XPU does not require the power-of-two limitation + // (matches the 3.0 fork). MLIR-level shape verifier disabled so kernels + // with e.g. BLOCK_M=480 / non-pow2 block_size_candidates tiles compile. + // if ((numElements & (numElements - 1)) != 0) + // return op->emitError("Number of elements must be power-of-two, but ") + // << *op << " doesn't follow the rule (" << numElements << ")" + // << " elements"; + //===-----------------------------------------------------------===// + } + } + for (auto opType : op->getResultTypes()) { + if (auto tensorType = dyn_cast(opType)) { + int64_t numElements = 1; + for (int64_t s : tensorType.getShape()) + numElements *= s; + if (numElements > maxTensorNumElements) + return op->emitError("Maximum allowed number of elements is ") + << maxTensorNumElements << ", but " << *op + << " has more than that"; + //===-------------------- For Triton XPU -----------------------===// + // [internal] Triton XPU does not require the power-of-two limitation + // (matches the 3.0 fork). MLIR-level shape verifier disabled so kernels + // with e.g. BLOCK_M=480 / non-pow2 block_size_candidates tiles compile. + // if ((numElements & (numElements - 1)) != 0) + // return op->emitError("Number of elements must be power-of-two, but ") + // << *op << " doesn't follow the rule (" << numElements << ")" + // << " elements"; + //===-----------------------------------------------------------===// + } + } + return success(); +} + +// Check that the Triton layouts on op's operands and return types are valid. +// For example, we check that the number of warps per block in a Triton GPU +// blocked layout matches that of its module. +// +// It's a little weird to check these properties of a layout only when the +// layout is used in an op, since most of the properties don't actually depend +// on the op. They do depend on the *module*, though, and a layout is attached +// to a module only by virtue of being used in one of the module's ops. +LogicalResult OpTrait::impl::verifyTensorLayouts(Operation *op) { + auto checkLayout = [&](Value val, auto makeErr) -> LogicalResult { + // Only ranked tensors can have layouts. + auto rankedTy = dyn_cast(val.getType()); + if (!rankedTy) + return success(); + + mlir::Attribute layout = rankedTy.getEncoding(); + if (!layout) + return success(); + + Dialect *dialect = &layout.getDialect(); + if (auto sliceLayout = dyn_cast(layout)) { + Attribute parent = sliceLayout.getParent(); + while (auto parentSlice = + dyn_cast(parent)) + parent = parentSlice.getParent(); + if (parent.getDialect().getNamespace() == "triton_xpu") + dialect = &parent.getDialect(); + } + + auto verifyLayoutInterface = + dyn_cast(dialect); + if (verifyLayoutInterface) { + return verifyLayoutInterface->verifyTensorLayout(layout, rankedTy, op, + makeErr); + } + + return success(); + }; + + for (size_t i = 0; i < op->getNumOperands(); i++) { + auto operand = op->getOperand(i); + auto err = checkLayout(operand, [&]() { + // Stringify the operand using `printAsOperand`. This prints e.g. "%42" + // rather than the full definition. + std::string operandStr; + llvm::raw_string_ostream os(operandStr); + // If we don't assume verified, dump() will recursively call this + // function! + operand.printAsOperand(os, OpPrintingFlags().assumeVerified()); + + return op->emitError("Operand ") + << i << " (" << operand << ") has an invalid layout: "; + }); + if (!err.succeeded()) + return err; + } + + for (size_t i = 0; i < op->getNumResults(); i++) { + auto result = op->getResult(i); + auto err = checkLayout(result, [&]() { + if (op->getNumResults() == 1) { + return op->emitError("Result has an invalid layout: "); + } else { + return op->emitError("Result ") << i << " has an invalid layout: "; + } + }); + if (!err.succeeded()) + return err; + } + + return success(); +} + +static ArrayRef getTypeShape(Type type) { + auto rankedType = dyn_cast(type); + if (auto ptrType = dyn_cast(type)) + rankedType = dyn_cast(ptrType.getPointeeType()); + return rankedType ? rankedType.getShape() : ArrayRef(); +} + +LogicalResult OpTrait::impl::verifySameLoadStoreOperandsShape(Operation *op) { + if (failed(verifyAtLeastNOperands(op, 1))) + return failure(); + + auto firstOperandShape = getTypeShape(op->getOperand(0).getType()); + for (auto type : llvm::drop_begin(op->getOperandTypes(), 1)) + if (failed(verifyCompatibleShape(getTypeShape(type), firstOperandShape))) + return op->emitOpError() << "requires the same shape for all operands"; + + return success(); +} + +LogicalResult +OpTrait::impl::verifySameLoadStoreOperandsAndResultShape(Operation *op) { + if (failed(verifyAtLeastNOperands(op, 1)) || + failed(verifyAtLeastNResults(op, 1))) + return failure(); + + auto firstOperandShape = getTypeShape(op->getOperand(0).getType()); + for (auto type : op->getResultTypes()) + if (failed(verifyCompatibleShape(getTypeShape(type), firstOperandShape))) + return op->emitOpError() + << "requires the same shape for all operands and results"; + + return verifySameLoadStoreOperandsShape(op); +} diff --git a/third_party/xpu/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp b/third_party/xpu/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp new file mode 100644 index 0000000000..5ab801eaee --- /dev/null +++ b/third_party/xpu/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp @@ -0,0 +1,4080 @@ +#if __has_include("flagtree_spec.h") +#include "flagtree_spec.h" +#endif + +#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_IR_Dialect + +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include +#include +#include + +#include "mlir/Dialect/UB/IR/UBOps.h" +#include "mlir/IR/DialectImplementation.h" +#include "mlir/IR/OpImplementation.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/Support/LLVM.h" +#include "triton/Analysis/Utility.h" +#include "triton/Dialect/Triton/IR/Interfaces.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Attributes.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h" +#include "triton/Dialect/TritonGPU/IR/TritonGPUInterfaces.h" +#include "triton/Dialect/TritonGPU/IR/Types.h" +#include "triton/Dialect/TritonGPU/Transforms/Utility.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +// FlagTree XPU spec override: needed for triton::xpu::TritonXPU_AttrTrait +// dispatch in getTotalElemsPerThread / getElemsPerThread below. Resolved via +// the XPU backend include dir (third_party/xpu/include). This include is only +// present in the XPU vendored copy of this file, never in the pristine tree. +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "triton/Tools/LayoutUtils.h" +#include "triton/Tools/LinearLayout.h" +#include "triton/Tools/StrUtil.h" +#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/MathExtras.h" + +// Include TableGen'erated code +#include "triton/Dialect/TritonGPU/IR/Dialect.cpp.inc" +#include "triton/Dialect/TritonGPU/IR/OpInterfaces.cpp.inc" +#include "triton/Dialect/TritonGPU/IR/TypeInterfaces.cpp.inc" + +using namespace mlir; +using namespace mlir::triton; +using namespace mlir::triton::gpu; + +static SmallVector +basesPerDimImpl(const LinearLayout::BasesT &namedBases, StringAttr dimName, + size_t rank, bool skipBroadcast = true); + +// Utility +namespace mlir { +namespace triton { +namespace gpu { + +LinearEncodingAttr TritonGPUDialect::toLinearEncoding(ArrayRef shape, + Attribute layout) { + // LinearEncoding is a DistributedLayout + std::vector allocationShape; + CacheKey key{std::vector(shape.begin(), shape.end()), layout}; + if (auto result = leCache.get(key)) { + return *result; + } + auto linearLayout = toLinearLayout(shape, layout); + auto linearEncoding = + LinearEncodingAttr::get(layout.getContext(), std::move(linearLayout)); + leCache.set(key, linearEncoding); + return linearEncoding; +} + +LinearEncodingAttr toLinearEncoding(DistributedEncodingTrait layout, + ArrayRef shape) { + auto *ctx = layout.getContext(); + return ctx->getLoadedDialect()->toLinearEncoding(shape, + layout); +} + +LinearEncodingAttr toLinearEncoding(RankedTensorType type) { + auto *ctx = type.getContext(); + return ctx->getLoadedDialect()->toLinearEncoding( + type.getShape(), type.getEncoding()); +} + +//===-------------------- For Triton XPU -----------------------===// +// FlagTree XPU spec override: forward-declare the free getElemsPerThread +// (Attribute, shape) overload so the XPU dispatch in getTotalElemsPerThread +// below can call it before its definition. Upstream/internal declares this in +// TritonGPU/IR/Dialect.h; under Q0b we keep the declaration local to this +// vendored copy instead of editing the shared main-tree header. +SmallVector getElemsPerThread(Attribute layout, + ArrayRef shape); +// FlagTree XPU spec override: a ttg.slice whose parent chain bottoms out at an +// XPU ClusterLayoutAttr. Used to keep XPU-backed slices off the generic +// LinearEncoding path (which asserts on non-power-of-two shapes). For non-XPU +// builds this file is not used; the pristine tree returns false / stays on the +// generic path unchanged. +static bool isXPUBackedLayout(Attribute layout) { + while (auto slice = mlir::dyn_cast(layout)) + layout = slice.getParent(); + return mlir::isa(layout); +} +//===-----------------------------------------------------------===// + +unsigned getTotalElemsPerThread(Attribute layout, ArrayRef shape) { + //===-------------------- For Triton XPU -----------------------===// + // Match the 3.0 fork: dispatch XPU layouts (ClusterLayoutAttr) directly to + // their own ceil-based arity instead of the generic LinearEncoding path. + // The generic path re-derives a LinearEncodingAttr whose toLinearLayout + // re-runs ensureLayoutNotLargerThan against the raw (possibly non-power-of- + // two) shape, which asserts on XPU shapes like 100. XPU allows non-pow2. + if (auto tritonXPUAttr = + mlir::dyn_cast(layout)) + return tritonXPUAttr.getTotalElemsPerThread(shape, /*eltTy=*/Type()); + // A ttg.slice whose parent is an XPU cluster: its own dialect is triton_gpu + // but it must follow XPU ceil arity, not the generic LinearEncoding path. + if (isXPUBackedLayout(layout)) + return product(getElemsPerThread(layout, shape)); + //===-----------------------------------------------------------===// + return toLinearEncoding(cast(layout), shape) + .getTotalElemsPerThread(shape); +} + +SmallVector getElemsPerThread(Attribute layout, + ArrayRef shape) { + //===-------------------- For Triton XPU -----------------------===// + // See getTotalElemsPerThread above: XPU layouts use their own ceil-based + // per-dim arity, matching the 3.0 fork and avoiding the generic + // LinearEncoding pow2 assert. + if (auto tritonXPUAttr = + mlir::dyn_cast(layout)) + return tritonXPUAttr.getElemsPerThread(shape, /*eltTy=*/Type()); + // XPU-backed ttg.slice: recurse into the parent XPU layout with the padded + // parent shape, then erase the sliced dim. Mirrors the 3.0 fork's + // SliceEncodingAttr::getElemsPerThread recursion. + if (auto sliceLayout = mlir::dyn_cast(layout)) { + if (isXPUBackedLayout(sliceLayout)) { + auto parentElemsPerThread = getElemsPerThread( + sliceLayout.getParent(), sliceLayout.paddedShape(shape)); + parentElemsPerThread.erase(parentElemsPerThread.begin() + + sliceLayout.getDim()); + return parentElemsPerThread; + } + } + //===-----------------------------------------------------------===// + return toLinearEncoding(cast(layout), shape) + .getElemsPerThread(shape); +} + +SmallVector getElemsPerThread(Type type) { + if (type.isIntOrIndexOrFloat() || isa(type)) + return SmallVector(1, 1); + auto tensorType = cast(type); + return getElemsPerThread(tensorType.getEncoding(), tensorType.getShape()); +} + +unsigned getTotalElemsPerThread(Type type) { + if (type.isIntOrIndexOrFloat() || isa(type)) + return 1; + auto tensorType = cast(type); + return getTotalElemsPerThread(tensorType.getEncoding(), + tensorType.getShape()); +} + +SmallVector getThreadsPerWarp(Attribute layout, + ArrayRef shape) { + return toLinearEncoding(cast(layout), shape) + .getThreadsPerWarp(); +} + +SmallVector getWarpsPerCTA(Attribute layout, + ArrayRef shape) { + return toLinearEncoding(cast(layout), shape) + .getWarpsPerCTA(); +} + +SmallVector getContigPerThread(RankedTensorType type) { + return toLinearEncoding(type).getContigPerThread(); +} + +bool isExpensiveView(Type srcType, Type dstType) { + auto tensorSrcType = cast(srcType); + auto tensorDstType = cast(dstType); + auto llSrc = toLinearLayout(tensorSrcType); + auto llDst = toLinearLayout(tensorDstType); + // In case there are replicated value we need to make sure the new and old + // layout have matching masks. + for (auto [srcMask, dstMask] : + llvm::zip(llSrc.getFreeVariableMasks(), llDst.getFreeVariableMasks())) { + assert(srcMask.first == dstMask.first); + if (srcMask.second != dstMask.second) + return true; + } + return getTotalElemsPerThread(srcType) != getTotalElemsPerThread(dstType); +} + +/* Utility function used by get.*Order methods of SliceEncodingAttr. + * Erase dim and decrease all values larger than dim by 1. + * Example: order = [0, 2, 4, 3, 1], dim = 2 + * resOrder = [0, 3, 2, 1] + */ +static SmallVector eraseOrder(ArrayRef order, + unsigned dim) { + unsigned rank = order.size(); + assert(dim < rank && "Invalid dim to erase"); + SmallVector resOrder; + for (unsigned i : order) + if (i < dim) + resOrder.push_back(i); + else if (i > dim) + resOrder.push_back(i - 1); + return resOrder; +} + +SmallVector getMatrixOrder(unsigned rank, bool rowMajor) { + // Return the order that represents that the batch is in row-major or + // column-major order for a batch of matrices of shape [*, m, n] with + // len(shape) == rank. + SmallVector order(rank); + if (rank < 2) { + return order; + } + std::iota(order.rbegin(), order.rend(), 0); + if (!rowMajor) { + std::swap(order[0], order[1]); + } + return order; +} + +SmallVector getOrderForDotOperand(unsigned opIdx, unsigned rank, + bool kContig) { + // kContig: if true, the matrix is fastest-running on k, + // otherwise it is on m (resp. n) + // opIdx=0: [*batch, m, k] + // opIdx=1: [*batch, k, n] + assert(opIdx == 0 || opIdx == 1); + auto rowMajor = bool(opIdx) != kContig; + return getMatrixOrder(rank, rowMajor); +} + +SmallVector getRepOrder(RankedTensorType type) { + auto layout = type.getEncoding(); + if (auto distributedLayout = mlir::dyn_cast(layout)) + return distributedLayout.getRepOrder(); + else + llvm::report_fatal_error("Unimplemented usage of getRepOrder"); + return {}; +} + +// Legacy impl for now +// This one's not terribly bad as we don't broadcast ShareEncodings +SmallVector getOrder(SharedEncodingTrait layout, + ArrayRef shape) { + if (auto swizzledLayout = dyn_cast(layout)) { + return llvm::to_vector(swizzledLayout.getOrder()); + } + if (auto paddedEnc = dyn_cast(layout)) { + return paddedEnc.getOrder(); + } + if (auto linearEnc = dyn_cast(layout)) { + return linearEnc.getOrder(); + } + if (auto sharedLayout = dyn_cast(layout)) { + if (shape.size() == 1) { + return {0}; + } + return getMatrixOrder(shape.size(), !sharedLayout.getTransposed()); + } + if (auto sharedLayout = dyn_cast(layout)) { + return llvm::to_vector(sharedLayout.getOrder()); + } + llvm::report_fatal_error("Unimplemented usage of getOrder for MemDescType"); + return {}; +} + +SmallVector getOrder(DistributedEncodingTrait layout, + ArrayRef shape) { + return toLinearEncoding(layout, shape).getOrder(); +} + +SmallVector getOrderForMemory(DistributedEncodingTrait layout, + ArrayRef shape) { + auto linear = toLinearEncoding(layout, shape); + auto order = linear.getOrder(); + auto threadOrder = linear.getThreadOrder(); + if (order == threadOrder) { + return order; + } + // Heuristic: + // If the element contiguity does not align with the thread order + // because the thread order dimension has contiguity of 1---meaning that + // the order position of this dimension is irrelevant---we prefer + // to use the thread order for the memory layout + auto contig = linear.getElemsPerThread(shape); + if (contig[threadOrder[0]] == 1) { + return threadOrder; + } + return order; +} + +SmallVector getThreadOrder(DistributedEncodingTrait layout, + ArrayRef shape) { + return toLinearEncoding(layout, shape).getThreadOrder(); +} + +SmallVector getWarpOrder(DistributedEncodingTrait layout, + ArrayRef shape) { + return toLinearEncoding(layout, shape).getWarpOrder(); +} + +CTAEncodingAttr getCTALayout(Attribute layout) { + if (auto ttgLayout = mlir::dyn_cast(layout)) + return ttgLayout.getCTALayout(); + llvm::report_fatal_error("Unimplemented usage of getCTALayout"); + return {}; +} + +SmallVector getCTAsPerCGA(Attribute layout) { + if (auto ttgLayout = mlir::dyn_cast(layout)) + return ttgLayout.getCTALayout().getCTAsPerCGA(); + llvm::report_fatal_error("Unimplemented usage of getCTAsPerCGA"); +} + +SmallVector getCTASplitNum(Attribute layout) { + SmallVector res; + if (auto ttgLayout = mlir::dyn_cast(layout)) { + return ttgLayout.getCTALayout().getCTASplitNum(); + } else if (auto tmemLayout = + mlir::dyn_cast( + layout)) { + res.resize(2); + res[0] = tmemLayout.getCTASplitM(); + res[1] = tmemLayout.getCTASplitN(); + } else if (auto tmemScaleLayout = mlir::dyn_cast< + triton::nvidia_gpu::TensorMemoryScalesEncodingAttr>(layout)) { + res.resize(2); + res[0] = tmemScaleLayout.getCTASplitM(); + res[1] = tmemScaleLayout.getCTASplitN(); + } else { + assert(false && "Unimplemented usage of getCTASplitNum"); + } + return res; +} + +SmallVector getCTAOrder(Attribute layout) { + SmallVector res; + if (auto ttgLayout = mlir::dyn_cast(layout)) { + res = ttgLayout.getCTALayout().getCTAOrder(); + } else { + llvm::report_fatal_error("Unimplemented usage of getCTAOrder"); + } + return res; +} + +SmallVector getShapePerCTA(ArrayRef CTASplitNum, + ArrayRef shape) { + unsigned rank = shape.size(); + auto splitNum = llvm::to_vector(CTASplitNum); + if (splitNum.size() <= rank) { // pipelining + splitNum.insert(splitNum.begin(), rank - splitNum.size(), 1); + } else { // memory slicing + splitNum = + llvm::to_vector(llvm::drop_begin(splitNum, splitNum.size() - rank)); + } + SmallVector shapePerCTA(rank); + for (unsigned i = 0; i < rank; ++i) { + shapePerCTA[i] = shape[i] / std::min(shape[i], splitNum[i]); + } + return shapePerCTA; +} + +SmallVector getShapePerCTA(Attribute layout, ArrayRef shape) { + return getShapePerCTA(getCTASplitNum(layout), shape); +} + +SmallVector getAllocationShapePerCTA(Attribute layout, + ArrayRef shapeLogical) { + SmallVector shape(shapeLogical); + if (auto sharedMMALayout = dyn_cast(layout)) { + if (sharedMMALayout.getFp4Padded()) { + auto packedAxis = getOrder(sharedMMALayout, shapeLogical)[0]; + shape[packedAxis] *= 2; + } + } + return getShapePerCTA(layout, shape); +} + +SmallVector getShapePerCTA(Type type) { + auto tensorType = cast(type); + return getShapePerCTA(tensorType.getEncoding(), tensorType.getShape()); +} + +SmallVector getAllocationShapePerCTA(Type type) { + auto tensorType = cast(type); + return getAllocationShapePerCTA(tensorType.getEncoding(), + tensorType.getShape()); +} + +unsigned getNumCTAs(Attribute layout) { + return product(getCTAsPerCGA(layout)); +} + +SmallVector orderPerDimImpl(const LinearLayout &ll, + StringAttr dimName, + ArrayRef defaultOrder) { + assert(ll.getBases().contains(dimName)); + const auto &bases = ll.getBases().find(dimName)->second; + llvm::SetVector order; + auto nonZero = [](auto val) { return val != 0; }; + for (const auto &basis : bases) { + // Bases can have one or zero non-zero elements + // Skip a basis if it's broadcasting (all zeros) + // e.g. warps for DotOperandEncodingAttr (see ampereDotToLinearLayout) + auto it = std::find_if(basis.begin(), basis.end(), nonZero); + if (it != basis.end()) { + auto i = it - basis.begin(); + order.insert(i); + } + } + // If any dim is missing, we add them in the defaultOrder + for (auto i : defaultOrder) { + order.insert(i); + } + return order.takeVector(); +} + +bool isExpensiveCat(CatOp cat, Attribute targetEncoding) { + // If the new elements per thread is less than the old one, we will need to + // do convert encoding that goes through shared memory anyway. So we + // consider it as expensive. + RankedTensorType tensorTy = cat.getType(); + auto totalElemsPerThread = gpu::getTotalElemsPerThread(tensorTy); + auto shape = tensorTy.getShape(); + auto newTotalElemsPerThread = + gpu::getTotalElemsPerThread(targetEncoding, shape); + return newTotalElemsPerThread < totalElemsPerThread; +} + +static LogicalResult +verifyLayoutOrder(function_ref emitError, + ArrayRef order) { + if (!isPermutationOfIota(order)) { + return emitError() + << "order must be a permutation of 0..(rank-1), but was [" << order + << "]"; + } + return success(); +} + +LogicalResult +CTAEncodingAttr::verify(function_ref emitError, + LinearLayout linearLayout) { + if (linearLayout.getNumInDims() != 1) { + return emitError() << "CTA encoding must have exactly one input dimension " + "named 'block'."; + } + auto dim = *linearLayout.getInDimNames().begin(); + auto ctx = dim.getContext(); + if (dim != StringAttr::get(ctx, "block")) { + return emitError() << "CTA encoding must have exactly one input dimension " + "named 'block'."; + } + + auto outDimNames = linearLayout.getOutDimNames(); + auto expected = standardOutDimNames(ctx, linearLayout.getNumOutDims()); + if (!llvm::equal(outDimNames, expected)) { + return emitError() << "CTA encoding output dims must be [dim0, dim1, ...], " + "but got [" + << outDimNames << "]."; + } + + return success(); +} + +CTAEncodingAttr CTAEncodingAttr::getDefault(MLIRContext *ctx, int rank) { + auto kBlock = StringAttr::get(ctx, "block"); + LinearLayout::BasesT bases; + bases[kBlock] = {}; + auto dims = standardOutDimNames(ctx, rank); + return get(ctx, LinearLayout(bases, dims)); +} + +CTAEncodingAttr CTAEncodingAttr::fromSplitParams(MLIRContext *ctx, + ArrayRef CTAsPerCGA, + ArrayRef CTASplitNum, + ArrayRef CTAOrder) { + int rank = CTAOrder.size(); + auto outDimNames = standardOutDimNames(ctx, rank); + StringAttr kBlock = StringAttr::get(ctx, "block"); + + LinearLayout layout = LinearLayout::empty(); + SmallVector splitNums(CTASplitNum.begin(), CTASplitNum.end()); + SmallVector ctas(CTAsPerCGA.begin(), CTAsPerCGA.end()); + + for (int i = 0; i < rank; ++i) { + int dim = CTAOrder[i]; + unsigned split = splitNums[dim]; + unsigned total = ctas[dim]; + assert(total % split == 0 && "invalid CTA encoding parameters"); + layout *= LinearLayout::identity1D(split, kBlock, outDimNames[dim]) * + LinearLayout::zeros1D(total / split, kBlock, outDimNames[dim]); + } + + layout = layout.transposeOuts(outDimNames); + return CTAEncodingAttr::get(ctx, layout); +} + +SmallVector CTAEncodingAttr::getCTAsPerCGA() const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), StringAttr::get(getContext(), "block"), + rank, /*skipBroadcast=*/false); +} + +SmallVector CTAEncodingAttr::getCTASplitNum() const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), StringAttr::get(getContext(), "block"), + rank); +} + +SmallVector CTAEncodingAttr::getCTAOrder() const { + auto rank = getRank(); + SmallVector defaultOrder(rank); + std::iota(defaultOrder.begin(), defaultOrder.end(), 0); + return orderPerDimImpl(getLinearLayout(), + StringAttr::get(getContext(), "block"), defaultOrder); +} + +LogicalResult BlockedEncodingAttr::verify( + function_ref emitError, + ArrayRef sizePerThread, ArrayRef threadsPerWarp, + ArrayRef warpsPerCTA, ArrayRef order, + CTAEncodingAttr CTALayout) { + if (!llvm::all_equal({sizePerThread.size(), threadsPerWarp.size(), + warpsPerCTA.size(), order.size()})) { + return emitError() << "sizePerThread, threadsPerWarp, warpsPerCTA, and " + "order must all have the same rank."; + } + if (llvm::any_of(sizePerThread, + [](unsigned x) { return !llvm::isPowerOf2_64(x); })) { + return emitError() + << "Every element in sizePerThread must be a power of two."; + } + if (llvm::any_of(threadsPerWarp, + [](unsigned x) { return !llvm::isPowerOf2_64(x); })) { + return emitError() + << "Every element in threadsPerWarp must be a power of two."; + } + if (llvm::any_of(warpsPerCTA, + [](unsigned x) { return !llvm::isPowerOf2_64(x); })) { + return emitError() + << "Every element in warpsPerCTA must be a power of two."; + } + + // Empty CTALayout is allowed, but if it's present its rank must match the + // BlockedEncodingAttr's rank. + if (order.size() != CTALayout.getRank()) { + return emitError() << "BlockedEncodingAttr and CTALayout's fields must " + "have the same rank."; + } + return verifyLayoutOrder(emitError, order); +} + +// 1 element per thread +// order = reverse(arange(rank)) +triton::gpu::BlockedEncodingAttr +getDefaultBlockedEncoding(MLIRContext *context, ArrayRef shape, + int numWarps, int threadsPerWarp, int numCTAs) { + int rank = shape.size(); + llvm::SmallVector order(rank); + std::iota(order.begin(), order.end(), 0); + std::reverse(order.begin(), order.end()); + llvm::SmallVector sizePerThread(rank, 1); + triton::gpu::BlockedEncodingAttr encoding = + triton::gpu::BlockedEncodingAttr::get(context, shape, sizePerThread, + order, numWarps, threadsPerWarp, + numCTAs); + return encoding; +} + +LogicalResult tryJoinOnAxis(MLIRContext *ctx, const LinearLayout &inLl, + LinearLayout &outLl, bool fwdInference, int axis, + std::optional loc) { + auto kRegister = StringAttr::get(ctx, "register"); + auto outDims = llvm::to_vector(inLl.getOutDimNames()); + if (fwdInference) { + auto split = LinearLayout::identity1D(2, kRegister, outDims[axis]); + outLl = split * inLl; + } else { + // Assert that there is a dimension with size 2 in the axis + // that has contiguous elements + // Note that this is more general than the fwdInference case in that + // - It allows the dimension not to be the fastest running + // - It allows broadcasting + // In general, this allows us to split along any axis as long as + // the basis (0, 0, ..., 0, 1, 0, ..., 0) is in the registers. + bool found = false; + LinearLayout::BasesT newBases; + for (const auto &basesDim : inLl.getBases()) { + std::vector> newBasesDim; + for (auto base : basesDim.second) { + if (base[axis] == 1 && basesDim.first == kRegister) { + found = true; + continue; + } + base[axis] /= 2; + newBasesDim.push_back(std::move(base)); + } + newBases.insert({basesDim.first, std::move(newBasesDim)}); + } + if (!found) + return emitOptionalError(loc, + "Fp4ToFpOp/SplitOp requires at least 2 elements " + "per thread in the axis/last dimension"); + outLl = LinearLayout(std::move(newBases), std::move(outDims)); + } + return success(); +} + +} // namespace gpu +} // namespace triton +} // namespace mlir + +static LogicalResult parseIntAttrValue(AsmParser &parser, Attribute attr, + unsigned &value, StringRef desc) { + auto intAttr = mlir::dyn_cast(attr); + if (!intAttr) { + parser.emitError(parser.getNameLoc(), "expected an integer type in ") + << desc; + return failure(); + } + if (intAttr.getType().isSignedInteger()) { + int64_t attrVal = intAttr.getSInt(); + if (attrVal < 0) { + parser.emitError(parser.getNameLoc(), + "expected an unsigned integer value in ") + << desc; + return failure(); + } + value = attrVal; + } else if (intAttr.getType().isSignlessInteger()) { + int64_t attrVal = intAttr.getInt(); + if (attrVal < 0) { + parser.emitError(parser.getNameLoc(), + "expected an unsigned integer value in ") + << desc; + return failure(); + } + value = attrVal; + } else { + value = intAttr.getUInt(); + } + return success(); +} + +static LogicalResult parseBoolAttrValue(AsmParser &parser, Attribute attr, + bool &value, StringRef desc) { + auto boolAttr = mlir::dyn_cast(attr); + if (!boolAttr) { + parser.emitError(parser.getNameLoc(), "expected a bool type in ") << desc; + return failure(); + } + value = boolAttr.getValue(); + return success(); +} + +// parse an array of integers +static LogicalResult parseIntArrayAttr(AsmParser &parser, + const NamedAttribute &attr, + SmallVector &res, + StringRef desc) { + auto arrayAttr = mlir::dyn_cast(attr.getValue()); + if (!arrayAttr) { + parser.emitError(parser.getNameLoc(), "expected an array for ") << desc; + return failure(); + } + for (Attribute i : arrayAttr) { + unsigned value; + if (parseIntAttrValue(parser, i, value, desc).failed()) + return failure(); + res.push_back(value); + } + return success(); +}; + +static LogicalResult parseUInt(AsmParser &parser, const NamedAttribute &attr, + unsigned &value, StringRef desc) { + return parseIntAttrValue(parser, attr.getValue(), value, desc); +}; + +static LogicalResult parseBool(AsmParser &parser, const NamedAttribute &attr, + bool &value, StringRef desc) { + return parseBoolAttrValue(parser, attr.getValue(), value, desc); +}; + +static LogicalResult parseType(AsmParser &parser, const NamedAttribute &attr, + Type &value, StringRef desc) { + auto typeAttr = mlir::dyn_cast(attr.getValue()); + if (!typeAttr) { + parser.emitError(parser.getNameLoc(), "expected a Type in ") << desc; + return failure(); + } + value = typeAttr.getValue(); + return success(); +} + +std::optional +parseLinearLayout(const DictionaryAttr &dict, AsmParser &parser, + ArrayRef inDimNames) { + LinearLayout::BasesT bases; + + // Parse the basis names in order (the order is relevant) + for (const auto &inDimNameStr : inDimNames) { + auto inDimName = StringAttr::get(parser.getContext(), inDimNameStr); + Attribute value = dict.get(inDimName); + if (!value) { + parser.emitError(parser.getCurrentLocation(), "Expected basis of '") + << inDimName.getValue() << "' not found"; + return {}; + } + // Expecting an array of arrays + auto arrayOfArraysAttr = mlir::dyn_cast(value); + if (!arrayOfArraysAttr) { + parser.emitError(parser.getCurrentLocation(), + "Expected array of arrays for basis of '") + << inDimName.getValue() << "'"; + return {}; + } + + std::vector> inDimBases; + for (Attribute arrayAttr : arrayOfArraysAttr) { + auto intArrayAttr = mlir::dyn_cast(arrayAttr); + if (!intArrayAttr) { + parser.emitError(parser.getCurrentLocation(), + "Expected array of integers in basis for '") + << inDimName.getValue() << "'"; + return {}; + } + std::vector basis; + for (Attribute intAttr : intArrayAttr) { + auto intValueAttr = mlir::dyn_cast(intAttr); + if (!intValueAttr) { + parser.emitError(parser.getCurrentLocation(), + "Expected integer in basis for '") + << inDimName.getValue() << "'"; + return {}; + } + basis.push_back(intValueAttr.getInt()); + } + inDimBases.push_back(std::move(basis)); + } + bases[inDimName] = std::move(inDimBases); + } + size_t rank = 0; + for (const auto &basesDim : llvm::make_second_range(bases)) { + if (!basesDim.empty()) { + rank = basesDim[0].size(); + break; + } + } + + // To implement this we'd need to serialise the rank as well. + // We can do this if we ever need it + if (rank == 0) { + parser.emitError(parser.getCurrentLocation(), "Empty Layout not supported"); + return {}; + } + + // Generate standared outDimNames (dim0, dim1, ...) + SmallVector outDimNames; + for (int i = 0; i < rank; ++i) { + outDimNames.push_back( + StringAttr::get(parser.getContext(), "dim" + llvm::Twine(i))); + } + + // Create LinearLayout + return LinearLayout(std::move(bases), std::move(outDimNames)); +} + +// We don't use the default implementation as it's a bit too verbose +// This prints in the following format that is shape agnostic, in the sense +// that we don't print explicitly the outShape of the LL +// We always assume LLs to be surjective +// <{register = [[0, 1], [8, 0], [0, 8], [64, 0]], +// lane = [[0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], +// warp = [[16, 0], [32, 0]], +// block = []}> +static void printLinearLayout(AsmPrinter &printer, const LinearLayout &ll) { + printer << join(ll.getBases(), ", ", [](const auto &base) { + return base.first.str() + " = " + "[" + + join(base.second, ", ", + [](const std::vector &vec) { + return "[" + join(vec, ", ") + "]"; + }) + + "]"; + }); +} + +// Print the CTA encoding as `CGALayout = [[...]]` when the layout is +// non-trivial. +static void maybePrintCTALayout(mlir::MLIRContext *context, + mlir::AsmPrinter &printer, + CTAEncodingAttr layout, unsigned rank) { + if (layout == CTAEncodingAttr::getDefault(context, rank)) + return; + + auto kBlock = StringAttr::get(context, "block"); + const auto &basesMap = layout.getLinearLayout().getBases(); + auto it = basesMap.find(kBlock); + assert(it != basesMap.end()); + const auto &bases = it->second; + // This is the default layout + assert(!bases.empty()); + + printer << ", CGALayout = ["; + llvm::interleaveComma(bases, printer, [&](const std::vector &vec) { + printer << "["; + llvm::interleaveComma(vec, printer); + printer << "]"; + }); + printer << "]"; +} + +//===----------------------------------------------------------------------===// +// Attribute methods +//===----------------------------------------------------------------------===// + +#include "triton/Dialect/TritonGPU/IR/AttrInterfaces.cpp.inc" + +#define GET_ATTRDEF_CLASSES +#include "triton/Dialect/TritonGPU/IR/AttrDefs.cpp.inc" +#undef GET_ATTRDEF_CLASSES + +//===----------------------------------------------------------------------===// +// Blocked Encoding +//===----------------------------------------------------------------------===// + +std::optional parseCTAAttr(AsmParser &parser, Attribute attr, + unsigned rank) { + if (!attr) + return CTAEncodingAttr::getDefault(parser.getContext(), rank); + + auto array = llvm::dyn_cast(attr); + if (!array) { + parser.emitError(parser.getNameLoc(), + "expected array value for 'CGALayout'"); + return {}; + } + + auto ctx = parser.getContext(); + auto cgaName = StringAttr::get(ctx, "CGALayout"); + std::vector> bases; + bases.reserve(array.size()); + for (Attribute vecAttr : array) { + SmallVector basisValues; + NamedAttribute basisAttr(cgaName, vecAttr); + if (parseIntArrayAttr(parser, basisAttr, basisValues, "CGALayout entry") + .failed()) + return {}; + if (basisValues.size() != rank) { + parser.emitError(parser.getNameLoc()) + << "'CGALayout' entry length does not match rank " << rank; + return {}; + } + std::vector basis; + basis.reserve(basisValues.size()); + for (unsigned value : basisValues) + basis.push_back(static_cast(value)); + bases.push_back(std::move(basis)); + } + + LinearLayout::BasesT namedBases; + namedBases.insert( + std::make_pair(StringAttr::get(ctx, "block"), std::move(bases))); + LinearLayout ll(namedBases, standardOutDimNames(ctx, rank)); + return CTAEncodingAttr::get(ctx, std::move(ll)); +} + +Attribute BlockedEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + // Parse the data as a dictionary + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + SmallVector sizePerThread; + SmallVector threadsPerWarp; + SmallVector warpsPerCTA; + SmallVector order; + Attribute ctaAttr = nullptr; + + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "sizePerThread") { + if (parseIntArrayAttr(parser, attr, sizePerThread, + "number of elements per thread") + .failed()) + return {}; + } else if (attr.getName() == "threadsPerWarp") { + if (parseIntArrayAttr(parser, attr, threadsPerWarp, + "number of threads per warp") + .failed()) + return {}; + } else if (attr.getName() == "warpsPerCTA") { + if (parseIntArrayAttr(parser, attr, warpsPerCTA, + "number of warps per CTA") + .failed()) + return {}; + } else if (attr.getName() == "order") { + if (parseIntArrayAttr(parser, attr, order, "order").failed()) + return {}; + } else if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + } else { + parser.emitError(parser.getNameLoc(), "unexpected key: ") + << attr.getName().strref(); + return {}; + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, /*rank=*/sizePerThread.size()); + if (!CTALayout.has_value()) + return {}; + + return parser.getChecked(parser.getContext(), + sizePerThread, threadsPerWarp, + warpsPerCTA, order, *CTALayout); +} + +void BlockedEncodingAttr::print(mlir::AsmPrinter &printer) const { + printer << "<{" + << "sizePerThread = [" << ArrayRef(getSizePerThread()) << "]" + << ", threadsPerWarp = [" << ArrayRef(getThreadsPerWarp()) << "]" + << ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]" + << ", order = [" << getOrder() << "]"; + + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getSizePerThread().size()); + + printer << "}>"; +} + +// FIXME Can we take the LinearLayout by const&? +LogicalResult +LinearEncodingAttr::verify(function_ref emitError, + LinearLayout linearLayout) { + // Example of LinearEncodingAttr + // <{register = [[0, 1], [8, 0], [0, 8], [64, 0]], + // lane = [[0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], + // warp = [[16, 0], [32, 0]], + // block = []}> + // The input dims must be {register, lane, warp, block} + // The output dims of the linear layout should be dim0..dim[rank-1] + + static const auto expectedInDims = + SmallVector({"register", "lane", "warp", "block"}); + for (const auto &[i, dims] : llvm::enumerate( + llvm::zip(linearLayout.getInDimNames(), expectedInDims))) { + const auto &[dim, expectedDimStr] = dims; + if (dim.str() != expectedDimStr) { + return emitError() << "Expected input dimension " << i << " to be '" + << expectedDimStr << "'. Got " << dim; + } + } + + // outDims are ['dim0', 'dim1', ...] + for (auto [i, dim] : llvm::enumerate(linearLayout.getOutDimNames())) { + if (dim.str() != ("dim" + llvm::Twine(i)).str()) { + return emitError() + << "Expected output dimensions to be ['dim0', 'dim1', ...]. Got " + << dim << " at position " << i; + } + } + + const auto &bases = linearLayout.getBases(); + auto nonZero = [](auto val) { return val != 0; }; + for (const auto &dimBases : llvm::make_second_range(bases)) { + if (!llvm::all_of(dimBases, [&](const auto &basis) { + return std::count_if(basis.begin(), basis.end(), nonZero) <= 1; + })) { + return emitError() + << "In a distributed layout, each base must move in at most one " + "dimension."; + } + } + + return success(); +} + +// If we only had BlockedEncodingAttr, we could simply return ArrayRefs here. +// But we need to have a consistent interface with e.g. SliceEncodingAttr, which +// computes some of these fields. +SmallVector BlockedEncodingAttr::getRepOrder() const { + return SmallVector(getOrder()); +} + +//===----------------------------------------------------------------------===// +// Linear Encoding +//===----------------------------------------------------------------------===// + +void LinearEncodingAttr::print(mlir::AsmPrinter &printer) const { + printer << "<{"; + printLinearLayout(printer, getLinearLayout()); + printer << "}>"; +} + +Attribute LinearEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + + if (parser.parseGreater().failed()) + return {}; + + std::vector inDimNames = {"register", "lane", "warp", "block"}; + auto maybeLL = parseLinearLayout(dict, parser, inDimNames); + if (!maybeLL.has_value()) + return {}; + + // Create and return the LinearEncodingAttr + return parser.getChecked(parser.getContext(), + std::move(*maybeLL)); +} + +static SmallVector +basesPerDimImpl(const LinearLayout::BasesT &namedBases, StringAttr dimName, + size_t rank, bool skipBroadcast) { + const auto &bases = namedBases.find(dimName)->second; + + if (bases.empty()) { + return SmallVector(rank, 1); + } + + SmallVector ret(rank, 1); + auto nonZero = [](auto val) { return val != 0; }; + int nonZeroIdx = 0; + for (const auto &basis : bases) { + auto it = std::find_if(basis.begin(), basis.end(), nonZero); + // Bases can have one or zero non-zero elements + // Skip a basis if it's broadcasting (all zeros) + // e.g. warps for DotOperandEncodingAttr (see ampereDotToLinearLayout) + if (it != basis.end()) { + nonZeroIdx = it - basis.begin(); + ret[nonZeroIdx] *= 2; + } else if (!skipBroadcast) { + // If we've seen a non-zero basis, we double the size of the previous dim + // This is just needed to count the CTAsPerCGA + ret[nonZeroIdx] *= 2; + } + } + return ret; +} + +SmallVector +LinearEncodingAttr::basesPerDim(StringAttr dimName, bool skipBroadcast) const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), dimName, rank, skipBroadcast); +} + +CTAEncodingAttr linearToCTAEncodingAttr(const LinearLayout &ll, + ArrayRef cgaLogicalShape) { + // Compute the shapePerCTA + auto shape = ll.getOutDims(); + for (int i = 0; i < shape.size(); ++i) { + shape[i].second /= cgaLogicalShape[i]; + } + auto inDims = to_vector(ll.getInDimNames()); + auto kBlock = inDims.back(); + assert(kBlock.str() == "block"); + inDims.pop_back(); + auto outDims = to_vector(ll.getOutDimNames()); + auto subLl = ll.sublayout(inDims, outDims); + // sublayout returns the same output size. We trim it to the + // real size + subLl = LinearLayout(subLl.getBases(), shape, false); + // The ctaLayout is what we get after dividing on the left by + // the layout in a single CTA + auto maybeCtaLayout = divideLeft(ll, subLl); + assert(maybeCtaLayout.has_value()); + auto *ctx = inDims[0].getContext(); + auto ctaLayout = maybeCtaLayout->sublayout({kBlock}, outDims); + return CTAEncodingAttr::get(ctx, std::move(ctaLayout)); +} + +SmallVector +LinearEncodingAttr::orderPerDim(StringAttr dimName, + ArrayRef defaultOrder) const { + return orderPerDimImpl(getLinearLayout(), dimName, defaultOrder); +} + +// [Note. Divergence of methods wrt. legacy layouts] +// For smaller shapes where the CTATile is larger than the output +// tensor, some methods return different values than the legacy layouts. I think +// this is benign tho. An example: what is the vector of `warpsPerCTA` if +// all the warps hold the same data? I think it should be [1, 1], even if we +// have 4 warps. But perhaps for this we have to add some masking in some +// places... We'll see +SmallVector LinearEncodingAttr::getRepOrder() const { + // This is not correct, but: + // - It happens to agree in most places with the legacy layout + // - getRepOrder does not make sense for LinearEncodingAttr as it already has + // the same shape as the tensor that uses it + return getOrder(); +} + +CTAEncodingAttr LinearEncodingAttr::getCTALayout() const { + auto splitNum = basesPerDim(StringAttr::get(getContext(), "block")); + return linearToCTAEncodingAttr(getLinearLayout(), splitNum); +} +SmallVector LinearEncodingAttr::getWarpsPerCTA() const { + return basesPerDim(StringAttr::get(getContext(), "warp")); +} +SmallVector LinearEncodingAttr::getWarpOrder() const { + return orderPerDim(StringAttr::get(getContext(), "warp"), getOrder()); +} +SmallVector LinearEncodingAttr::getThreadsPerWarp() const { + return basesPerDim(StringAttr::get(getContext(), "lane")); +} +SmallVector LinearEncodingAttr::getThreadOrder() const { + return orderPerDim(StringAttr::get(getContext(), "lane"), getOrder()); +} + +SmallVector LinearEncodingAttr::getSizePerThread() const { + auto rank = getOrder().size(); + auto ll = getLinearLayout(); + auto ctx = getContext(); + auto kRegister = StringAttr::get(ctx, "register"); + auto splitNum = getCTALayout().getCTASplitNum(); + + // We canonicalize on the spot, as if we use CGAs the regs are not in + // canonical form The order is [reg, lane, warp, rep, block], so we first + // remove the blocks + llvm::SmallVector ctaShape; + for (auto [shape, cgaNum] : llvm::zip(ll.getOutDimSizes(), splitNum)) { + ctaShape.push_back(shape / cgaNum); + } + LinearLayout::BasesT bases = ll.getBases(); + + llvm::SetVector reverseRepOrder; + auto nonZero = [](auto val) { return val != 0; }; + auto ®isters = bases[kRegister]; + while (!registers.empty()) { + auto &basis = registers.back(); + auto it = std::find_if(basis.begin(), basis.end(), nonZero); + // If there's broadcasting (base == zeros) there are no more reps + if (it == basis.end()) { + break; + } + auto dim = it - basis.begin(); + reverseRepOrder.insert(dim); + // As soon as we stop finding reps, we stop + if (dim != reverseRepOrder.back() || 2 * basis[dim] != ctaShape[dim]) { + break; + } + ctaShape[dim] /= 2; + registers.pop_back(); + } + return basesPerDimImpl(bases, kRegister, rank); +} + +SmallVector LinearEncodingAttr::getOrder() const { + auto rank = getLinearLayout().getNumOutDims(); + SmallVector order(rank); + // Choose [rank-1, rank-2, ... 0] as the default order in case + // there are dims that do not move in the register + // This order is as good as any really + std::iota(order.rbegin(), order.rend(), 0); + + return orderPerDim(StringAttr::get(getContext(), "register"), order); +} + +LinearLayout LinearEncodingAttr::toLinearLayout(ArrayRef shape) const { + auto ll = getLinearLayout(); + auto canonicalDims = llvm::to_vector(ll.getOutDimNames()); + llvm::SmallDenseMap namedShape; + llvm::SmallVector permutedDims; + for (auto dim : getRepOrder()) { + permutedDims.push_back(canonicalDims[dim]); + namedShape[canonicalDims[dim]] = shape[dim]; + } + ll = ll.transposeOuts(permutedDims); + ll = ensureLayoutNotSmallerThan(ll, namedShape); + ll = ensureLayoutNotLargerThan(ll, namedShape, /*broadcastRegisters=*/false); + ll = ll.transposeOuts(canonicalDims); + return ll; +} + +SmallVector +LinearEncodingAttr::getElemsPerThread(ArrayRef shape) const { + // When broadcasting the layout the shape changes, otherwise the shape is + // the same as the shape of the tensor + // We can either have BroadcastOp with SameOperandsAndResultEncoding, or keep + // the invariant that the shape of the LL is that of the tensor + // We choose the former for BC + auto scaledLayout = get(getContext(), toLinearLayout(shape)); + auto kRegister = StringAttr::get(getContext(), "register"); + return scaledLayout.basesPerDim(kRegister, /*skipBroadcast=*/false); +} + +SmallVector +LinearEncodingAttr::getContig(const char *inDim, + SmallVector lowerContig) const { + auto ll = getLinearLayout(); + const auto &bases = + ll.getBases().find(StringAttr::get(getContext(), inDim))->second; + auto order = getOrder(); + auto rank = order.size(); + + SmallVector contig(lowerContig); + auto basisIt = bases.begin(); + for (unsigned dim : order) { + std::vector basis(rank, 0); + basis[dim] = contig[dim]; + + while (basisIt != bases.end() && *basisIt == basis) { + contig[dim] *= 2; + basis[dim] *= 2; + ++basisIt; + } + } + return contig; +} + +SmallVector LinearEncodingAttr::getContigPerThread() const { + SmallVector contig(getOrder().size(), 1); + return getContig("register", contig); +} + +SmallVector LinearEncodingAttr::getContigPerWarp() const { + return getContig("lane", getContigPerThread()); +} + +unsigned +LinearEncodingAttr::getTotalElemsPerThread(ArrayRef shape) const { + return product(getElemsPerThread(shape)); +} + +//===----------------------------------------------------------------------===// +// MMA encoding +//===----------------------------------------------------------------------===// + +Attribute NvidiaMmaEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned versionMajor = 0; + unsigned versionMinor = 0; + SmallVector warpsPerCTA; + SmallVector instrShape; + Attribute ctaAttr = nullptr; + + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "versionMajor") { + if (parseUInt(parser, attr, versionMajor, "versionMajor").failed()) + return {}; + } + if (attr.getName() == "versionMinor") { + if (parseUInt(parser, attr, versionMinor, "versionMinor").failed()) + return {}; + } + if (attr.getName() == "warpsPerCTA") { + if (parseIntArrayAttr(parser, attr, warpsPerCTA, "warpsPerCTA").failed()) + return {}; + } + if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + continue; + } + if (attr.getName() == "instrShape") { + if (parseIntArrayAttr(parser, attr, instrShape, "instrShape").failed()) { + return {}; + } + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, /*rank=*/warpsPerCTA.size()); + if (!CTALayout.has_value()) + return {}; + + return parser.getChecked( + parser.getContext(), versionMajor, versionMinor, warpsPerCTA, *CTALayout, + instrShape); +} + +void NvidiaMmaEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "versionMajor = " << getVersionMajor() + << ", versionMinor = " << getVersionMinor() // + << ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]"; + + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getRank()); + + printer << ", instrShape = [" << getInstrShape() << "]}>"; +} + +//===----------------------------------------------------------------------===// +// MFMA encoding +//===----------------------------------------------------------------------===// + +Attribute AMDMfmaEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned version = 0; + SmallVector warpsPerCTA; + SmallVector instrShape; + bool isTransposed; + SmallVector tilesPerWarp = {}; + unsigned elementBitWidth = 32; + Attribute ctaAttr = nullptr; + + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "version") { + if (parseUInt(parser, attr, version, "version").failed()) + return {}; + } + if (attr.getName() == "warpsPerCTA") { + if (parseIntArrayAttr(parser, attr, warpsPerCTA, "warpsPerCTA").failed()) + return {}; + } + if (attr.getName() == "instrShape") { + if (parseIntArrayAttr(parser, attr, instrShape, "instrShape").failed()) + return {}; + } + if (attr.getName() == "isTransposed") { + if (parseBool(parser, attr, isTransposed, "isTransposed").failed()) + return {}; + } + if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + continue; + } + if (attr.getName() == "tilesPerWarp") { + if (parseIntArrayAttr(parser, attr, tilesPerWarp, "tilesPerWarp") + .failed()) + return {}; + } + if (attr.getName() == "elementBitWidth") { + if (parseUInt(parser, attr, elementBitWidth, "elementBitWidth").failed()) + return {}; + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, /*rank=*/warpsPerCTA.size()); + if (!CTALayout.has_value()) + return {}; + + if (tilesPerWarp.empty()) + tilesPerWarp = SmallVector(instrShape.size(), 1); + + return parser.getChecked( + parser.getContext(), version, warpsPerCTA, instrShape, isTransposed, + *CTALayout, tilesPerWarp, elementBitWidth); +} + +void AMDMfmaEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "version = " << getVersion() // + << ", warpsPerCTA = [" << getWarpsPerCTA() << "]" // + << ", instrShape = [" << getInstrShape() << "]"; + + printer << ", isTransposed = " << getIsTransposed(); + + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getRank()); + + auto tilesPerWarp = getTilesPerWarp(); + if (!hasUnitTilesPerWarp()) + printer << ", tilesPerWarp = [" << getTilesPerWarp() << "]"; + + auto elementBitWidth = getElementBitWidth(); + if (elementBitWidth != 32) + printer << ", elementBitWidth = " << elementBitWidth; + + printer << "}>"; +} + +LogicalResult AMDMfmaEncodingAttr::verify( + function_ref emitError, unsigned version, + llvm::ArrayRef warpsPerCTA, + llvm::ArrayRef instrShape, bool isTransposed, + mlir::triton::gpu::CTAEncodingAttr, + llvm::ArrayRef tilesPerWarp, unsigned elementBitWidth) { + if (!(version >= 0 && version <= 4)) { + return emitError() << "version must be in the [0, 4] range"; + } + + auto mDim = instrShape[0]; + auto nDim = instrShape[1]; + const std::array, 4> validDims = { + {{32, 32}, {16, 16}, {64, 4}, {4, 64}}}; + if (!llvm::is_contained(validDims, std::make_pair(mDim, nDim))) { + return emitError() << "invalid (mDim, nDim) combination: (" << mDim << ", " + << nDim << ")"; + } + + if (!(elementBitWidth == 32 || elementBitWidth == 64)) + return emitError() << "elementBitWidth must be 32 or 64"; + + return success(); +} + +//===----------------------------------------------------------------------===// +// WMMA encoding +//===----------------------------------------------------------------------===// +bool AMDWmmaEncodingAttr::hasUnitTilesPerWarp() const { + return llvm::all_of(getTilesPerWarp(), [](int x) { return x == 1; }); +} + +Attribute AMDWmmaEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned version = 0; + bool isTransposed = false; + SmallVector warpsPerCTA; + SmallVector tilesPerWarp = {}; + SmallVector instrShape = getDefaultInstrShape(); + Attribute ctaAttr = nullptr; + + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "version") { + if (parseUInt(parser, attr, version, "version").failed()) + return {}; + } + if (attr.getName() == "isTranspose") { + if (parseBool(parser, attr, isTransposed, "isTranspose").failed()) + return {}; + } + if (attr.getName() == "warpsPerCTA") { + if (parseIntArrayAttr(parser, attr, warpsPerCTA, "warpsPerCTA").failed()) + return {}; + } + if (attr.getName() == "tilesPerWarp") { + if (parseIntArrayAttr(parser, attr, tilesPerWarp, "tilesPerWarp") + .failed()) + return {}; + } + if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + continue; + } + if (attr.getName() == "instrShape") { + instrShape.clear(); + if (parseIntArrayAttr(parser, attr, instrShape, "instrShape").failed()) { + return {}; + } + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, /*rank=*/warpsPerCTA.size()); + if (!CTALayout.has_value()) + return {}; + + if (tilesPerWarp.empty()) + tilesPerWarp = SmallVector(instrShape.size(), 1); + + return parser.getChecked( + parser.getContext(), version, isTransposed, warpsPerCTA, tilesPerWarp, + *CTALayout, instrShape); +} + +void AMDWmmaEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "version = " << getVersion() + << ", isTranspose = " << getIsTransposed() // + << ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]"; + + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getWarpsPerCTA().size()); + + auto tilesPerWarp = getTilesPerWarp(); + if (!hasUnitTilesPerWarp()) + printer << ", tilesPerWarp = [" << getTilesPerWarp() << "]"; + + if (getInstrShape() != ArrayRef(getDefaultInstrShape())) { + printer << ", instrShape = [" << getInstrShape() << "]"; + } + printer << "}>"; +} + +LogicalResult AMDWmmaEncodingAttr::verify( + function_ref emitError, unsigned version, + bool isTransposed, llvm::ArrayRef warpsPerCTA, + llvm::ArrayRef tilesPerWarp, CTAEncodingAttr ctaLayout, + llvm::ArrayRef instrShape) { + if (!(version >= 1 && version <= 3)) + return emitError() << "WMMA version must be in the [1, 3] range"; + + auto shape = SmallVector(instrShape); + auto validShapesV1 = std::vector>{{16, 16, 16}}; + if (version == 1 && !llvm::is_contained(validShapesV1, shape)) + return emitError() << "invalid WMMA v1 instruction shape"; + + auto validShapesV2 = + std::vector>{{16, 16, 16}, {16, 16, 32}}; + if (version == 2 && !llvm::is_contained(validShapesV2, shape)) + return emitError() << "invalid WMMA v2 instruction shape"; + + auto validShapesV3 = std::vector>{ + {16, 16, 4}, {16, 16, 32}, {16, 16, 64}, {16, 16, 128}}; + if (version == 3 && !llvm::is_contained(validShapesV3, shape)) + return emitError() << "invalid WMMA v3 instruction shape"; + + return success(); +} + +//===----------------------------------------------------------------------===// +// Sliced Encoding +//===----------------------------------------------------------------------===// + +Attribute SliceEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + NamedAttrList attrs; + if (parser.parseOptionalAttrDict(attrs).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + unsigned dim = mlir::cast(attrs.get("dim")).getInt(); + auto parent = mlir::dyn_cast(attrs.get("parent")); + if (!parent) { + parser.emitError(parser.getNameLoc(), + "expected a distributed encoding trait"); + return {}; + } + return parser.getChecked(parser.getContext(), dim, parent); +} + +void SliceEncodingAttr::print(mlir::AsmPrinter &printer) const { + printer << "<{" + << "dim = " << getDim() << ", " + << "parent = " << getParent() << "}>"; +} + +LogicalResult +SliceEncodingAttr::verify(function_ref emitError, + unsigned dim, DistributedEncodingTrait parent) { + unsigned rank = ::getCTALayout(parent).getRank(); + if (rank <= 1) + return emitError() << "parent layout must have at least rank >= 2"; + if (dim >= rank) { + return emitError() << "slice dim=" << dim + << " must be less than the parent rank=" << rank; + } + return success(); +} + +SmallVector SliceEncodingAttr::getRepOrder() const { + auto parentRepOrder = getParent().getRepOrder(); + return eraseOrder(parentRepOrder, getDim()); +} + +CTAEncodingAttr SliceEncodingAttr::getCTALayout() const { + auto layout = ::getCTALayout(getParent()).getLinearLayout(); + layout = removeStandardDim(layout, getDim()); + return CTAEncodingAttr::get(getContext(), layout); +} + +template +SmallVector SliceEncodingAttr::paddedShape(ArrayRef shape) const { + size_t rank = shape.size(); + unsigned dim = getDim(); + SmallVector retShape(rank + 1); + for (unsigned d = 0; d < rank + 1; ++d) { + if (d < dim) + retShape[d] = shape[d]; + else if (d == dim) + retShape[d] = 1; + else + retShape[d] = shape[d - 1]; + } + return retShape; +} +template SmallVector +SliceEncodingAttr::paddedShape(ArrayRef shape) const; +template SmallVector +SliceEncodingAttr::paddedShape(ArrayRef shape) const; + +template +Attribute parseSwizzledEncoding(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + // Parse the data as a dictionary + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned vec = 0; + unsigned perPhase = 0; + unsigned maxPhase = 0; + SmallVector order; + Attribute ctaAttr = nullptr; + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "vec") { + if (parseUInt(parser, attr, vec, "vec").failed()) + return {}; + } else if (attr.getName() == "perPhase") { + if (parseUInt(parser, attr, perPhase, "perPhase").failed()) + return {}; + } else if (attr.getName() == "maxPhase") { + if (parseUInt(parser, attr, maxPhase, "maxPhase").failed()) + return {}; + } else if (attr.getName() == "order") { + if (parseIntArrayAttr(parser, attr, order, "order").failed()) + return {}; + } else { + if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + } else { + parser.emitError(parser.getNameLoc(), "unexpected key: ") + << attr.getName().strref(); + return {}; + } + } + } + + if (auto CTALayout = parseCTAAttr(parser, ctaAttr, order.size())) + return parser.getChecked( + parser.getContext(), vec, perPhase, maxPhase, order, *CTALayout); + return {}; +} + +//===----------------------------------------------------------------------===// +// SwizzledShared encoding +//===----------------------------------------------------------------------===// + +LogicalResult +SwizzledSharedEncodingAttr::verify(function_ref emitError, + unsigned vec, unsigned perPhase, + unsigned maxPhase, ArrayRef order, + CTAEncodingAttr ctaLayout) { + if (order.size() != ctaLayout.getRank()) { + return emitError() << "order size (" << order.size() + << ") must match CTALayout rank (" << ctaLayout.getRank() + << ")"; + } + return verifyLayoutOrder(emitError, order); +} + +Attribute SwizzledSharedEncodingAttr::parse(AsmParser &parser, Type type) { + return parseSwizzledEncoding(parser, type); +} + +void SwizzledSharedEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "vec = " << getVec() // + << ", perPhase = " << getPerPhase() + << ", maxPhase = " << getMaxPhase() // + << ", order = [" << getOrder() << "]"; + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getOrder().size()); + printer << "}>"; +} + +//===----------------------------------------------------------------------===// +// SharedLinear encoding +//===----------------------------------------------------------------------===// + +LogicalResult +SharedLinearEncodingAttr::verify(function_ref emitError, + LinearLayout linearLayout, + unsigned layoutAlignment) { + if (layoutAlignment == 0 || !llvm::isPowerOf2_32(layoutAlignment)) { + return emitError() << "alignment must be a positive power of two"; + } + static const auto expectedInDims = + SmallVector({"offset", "block"}); + for (const auto &[index, dims] : llvm::enumerate( + llvm::zip(linearLayout.getInDimNames(), expectedInDims))) { + const auto &[dim, expected] = dims; + if (dim.str() != expected) { + return emitError() << "Expected input dimension " << index << " to be '" + << expected << "'. Got " << dim; + } + } + + for (auto [i, dim] : llvm::enumerate(linearLayout.getOutDimNames())) { + if (dim.str() != ("dim" + llvm::Twine(i)).str()) { + return emitError() + << "Expected output dimensions to be ['dim0', 'dim1', ...]. Got " + << dim << " at position " << i; + } + } + + SmallVector outDimNames = + llvm::to_vector(linearLayout.getOutDimNames()); + if (outDimNames.empty()) { + return emitError() + << "SharedLinearEncodingAttr requires at least one output" + " dimension."; + } + + auto *ctx = outDimNames.front().getContext(); + auto kOffset = StringAttr::get(ctx, "offset"); + auto kBlock = StringAttr::get(ctx, "block"); + + if (!linearLayout.isSurjective()) { + return emitError() << "The layout must be surjective"; + } + + LinearLayout withoutBroadcast = + linearLayout.removeZeroBasesAlongDim(kOffset).removeZeroBasesAlongDim( + kBlock); + if (!withoutBroadcast.isInvertible()) { + return emitError() + << "After removing the zero bases the layout must be bijective"; + } + + return success(); +} + +void SharedLinearEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{"; + auto layout = getLinearLayout(); + auto kBlock = StringAttr::get(getContext(), "block"); + auto kOffset = StringAttr::get(getContext(), "offset"); + if (layout.getBases().lookup(kBlock).empty()) { + layout = + layout.sublayout({kOffset}, llvm::to_vector(layout.getOutDimNames())); + } + printLinearLayout(printer, layout); + printer << "}, alignment = " << getAlignment() << ">"; +} + +Attribute SharedLinearEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + + DictionaryAttr layoutDictRaw; + if (parser.parseAttribute(layoutDictRaw).failed()) + return {}; + + if (layoutDictRaw.get("alignment")) { + parser.emitError(parser.getCurrentLocation()) + << "alignment must be specified outside of the linear layout braces"; + return {}; + } + + NamedAttrList layoutAttrList(layoutDictRaw.getValue()); + auto *ctx = parser.getContext(); + auto kBlock = StringAttr::get(ctx, "block"); + if (!layoutAttrList.get(kBlock)) { + layoutAttrList.push_back({kBlock, ArrayAttr::get(ctx, {})}); + } + + DictionaryAttr layoutDict = layoutAttrList.getDictionary(ctx); + + // Parse alignment + unsigned layoutAlignment; + if (parser.parseComma().failed()) + return {}; + if (parser.parseKeyword("alignment").failed() || parser.parseEqual().failed()) + return {}; + if (parser.parseInteger(layoutAlignment).failed()) + return {}; + + if (parser.parseGreater().failed()) + return {}; + + std::vector inDimNames = {"offset", "block"}; + auto maybeLL = parseLinearLayout(layoutDict, parser, inDimNames); + if (!maybeLL.has_value()) + return {}; + + // Special case for cleaner errors + if (layoutDict.get("alignment")) { + parser.emitError(parser.getCurrentLocation()) + << "alignment must be specified outside of the linear layout braces"; + return {}; + } + + if (layoutDict.size() != 2) { + parser.emitError(parser.getCurrentLocation()) + << "SharedLinearEncodingAttr must have exactly two attributes: offset " + "and block"; + return {}; + } + + return parser.getChecked( + parser.getContext(), std::move(*maybeLL), layoutAlignment); +} + +SmallVector +SharedLinearEncodingAttr::basesPerDim(StringAttr dimName, + bool skipBroadcast) const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), dimName, rank, skipBroadcast); +} + +SmallVector +SharedLinearEncodingAttr::orderPerDim(StringAttr dimName, + ArrayRef defaultOrder) const { + return orderPerDimImpl(getLinearLayout(), dimName, defaultOrder); +} + +SmallVector SharedLinearEncodingAttr::getOrder() const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + SmallVector defaultOrder(rank); + std::iota(defaultOrder.rbegin(), defaultOrder.rend(), 0); + return orderPerDim(StringAttr::get(getContext(), "offset"), defaultOrder); +} + +CTAEncodingAttr SharedLinearEncodingAttr::getCTALayout() const { + auto splitNum = basesPerDim(StringAttr::get(getContext(), "block")); + return linearToCTAEncodingAttr(getLinearLayout(), splitNum); +} +LinearLayout +SharedLinearEncodingAttr::toLinearLayout(ArrayRef shape) const { + auto ll = getLinearLayout(); + auto outDimNames = llvm::to_vector(ll.getOutDimNames()); + assert(shape.size() == outDimNames.size()); + // We don't support automatic broadcasting for shared linear layouts + for (auto [size, llSize] : llvm::zip(shape, ll.getOutDimSizes())) { + assert(size == llSize); + } + return ll; +} + +//===----------------------------------------------------------------------===// +// PaddedShared encoding +//===----------------------------------------------------------------------===// + +Attribute PaddedSharedEncodingAttr::parse(AsmParser &parser, Type type) { + // <[ + if (failed(parser.parseLess()) || failed(parser.parseLSquare())) + return {}; + + // :+ + SmallVector intervals, paddings; + auto parseIntervalPaddingPair = [&]() { + unsigned interval = 0, padding = 0; + if (failed(parser.parseInteger(interval)) || failed(parser.parseColon()) || + failed(parser.parsePlus()) || failed(parser.parseInteger(padding))) + return failure(); + intervals.push_back(interval); + paddings.push_back(padding); + return success(); + }; + // ] + if (failed(parser.parseCommaSeparatedList(parseIntervalPaddingPair)) || + failed(parser.parseRSquare())) + return {}; + + // {} + auto attrList = DictionaryAttr::get(parser.getContext()); + if (failed(parser.parseAttribute(attrList))) + return {}; + + // We have 2 possible formats for the attr-dict: + // 1) offset=[..], block=[..] handled by parseLinearLayout + // 2) order=[..], shape=[..] which creates an identity mapping + + std::optional maybeLL; + // Assume it's the first variant if offset or block is defined + if (attrList.contains("offset") || attrList.contains("block")) { + std::vector inDimNames = {"offset", "block"}; + // Error out on additional attribute names + for (const NamedAttribute &attr : attrList) { + if (!llvm::is_contained(inDimNames, attr.getName())) { + parser.emitError(parser.getCurrentLocation(), "Unexpected attribute ") + << attr.getName() << " found"; + } + } + maybeLL = parseLinearLayout(attrList, parser, inDimNames); + } else { + // Parse the second form + SmallVector order; + SmallVector shape; + for (const NamedAttribute &attr : attrList) { + if (attr.getName() == "order") { + if (parseIntArrayAttr(parser, attr, order, "order").failed()) + return {}; + } else if (attr.getName() == "shape") { + if (parseIntArrayAttr(parser, attr, shape, "shape").failed()) + return {}; + } else { + parser.emitError(parser.getCurrentLocation(), "Unexpected attribute ") + << attr.getName() << " found"; + return {}; + } + } + + if (order.size() != shape.size()) { + parser.emitError(parser.getCurrentLocation(), + "Mismatch of shape and order ranks in padded layout"); + return {}; + } + + // Create identity mapping based on shape and order + auto kOffset = StringAttr::get(parser.getContext(), "offset"); + maybeLL = identityStandardND(kOffset, shape, order); + maybeLL = combineCtaCgaWithShape( + *maybeLL, + CTAEncodingAttr::getDefault(parser.getContext(), shape.size()), + SmallVector(ArrayRef(shape))); + } + + if (!maybeLL.has_value()) + return {}; + + // > + if (parser.parseGreater().failed()) + return {}; + + return parser.getChecked( + parser.getContext(), intervals, paddings, *maybeLL); +} + +void PaddedSharedEncodingAttr::print(AsmPrinter &printer) const { + + auto *ctx = getContext(); + const auto &ll = getLinearComponent(); + + printer << "<["; + llvm::interleaveComma(llvm::zip(getIntervals(), getPaddings()), printer, + [&](std::tuple intervalPad) { + printer << std::get<0>(intervalPad) << ":+" + << std::get<1>(intervalPad); + }); + printer << "] {"; + + // We have a short hand form if linearComponent: + // 1) does have an empty CTA layout (empty block dim) + // 2) offsets are an identity mapping + auto kOffset = StringAttr::get(ctx, "offset"); + auto kBlock = StringAttr::get(ctx, "block"); + auto shape = SmallVector(ll.getOutDimSizes()); + + bool hasEmptyBlock = ll.getInDimSizeLog2(kBlock) == 0; + + LinearLayout identity = identityStandardND(kOffset, shape, getOrder()) + .transposeOuts(to_vector(ll.getOutDimNames())); + auto offsetLayout = ll.sublayout({kOffset}, to_vector(ll.getOutDimNames())); + + if (hasEmptyBlock && offsetLayout == identity) { + printer << "order = [" << ArrayRef(getOrder()) << "], shape = [" + << ArrayRef(shape) << "]"; + } else { + printLinearLayout(printer, getLinearComponent()); + } + + printer << "}>"; +} + +LogicalResult PaddedSharedEncodingAttr::verify( + function_ref emitError, ArrayRef intervals, + ArrayRef paddings, LinearLayout linearComponent) { + if (intervals.size() != paddings.size()) + return emitError() << "intervals size (" << intervals.size() + << ") must match paddings size (" << paddings.size() + << ")"; + + if (intervals.empty()) + return emitError() << "must have at least one interval-padding pair"; + + if (!llvm::all_of(intervals, llvm::isPowerOf2_32)) + return emitError() << "interval values must all be power of two"; + if (!llvm::all_of(paddings, llvm::isPowerOf2_32)) + return emitError() << "padding values must all be power of two"; + + llvm::SmallSet intervalValues(intervals.begin(), + intervals.end()); + if (intervalValues.size() != intervals.size()) + return emitError() << "interval values cannot have duplicates"; + + const auto &ll = linearComponent; + // The linear layout should map from [offset, block] to [dim0..dimN). All + // bases should be 0 or power of twos and move in a single direction without + // broadcasting + + if (ll == LinearLayout::empty()) + return emitError() << "linearComponent cannot be empty"; + + assert(!ll.getInDimNames().empty()); + auto *ctx = ll.getInDimNames().begin()->getContext(); + + if (!llvm::equal(ll.getInDimNames(), + std::array{StringAttr::get(ctx, "offset"), + StringAttr::get(ctx, "block")})) { + return emitError() + << "linearComponent must have [offset, block] as input dims"; + } + + if (!llvm::equal(ll.getOutDimNames(), + standardOutDimNames(ctx, ll.getNumOutDims()))) { + return emitError() + << "Expected output dimensions to be ['dim0', 'dim1', ...]."; + } + + const auto &bases = ll.getBases(); + + // Check that we are not broadcasting or having repeated bases + if (!ll.isInvertible()) { + return emitError() << "Broadcasting is not supported."; + } + + auto nonZero = [](auto val) { return val != 0; }; + for (const auto &dimBases : llvm::make_second_range(bases)) { + if (!llvm::all_of(dimBases, [&](const auto &basis) { + return llvm::count_if(basis, nonZero) <= 1; + })) { + return emitError() + << "Each offset basis must move in at most one dimension."; + } + // Ensure all non zero elements are a power of 2. Combined with the + // broadcast check above this prevents per element swizzling. The intent of + // the linear component is to rearrange whole rows or cache-line sized + // chunks of rows. + if (!llvm::all_of(dimBases, [&](const auto &basis) { + return llvm::all_of( + basis, [](auto v) { return v == 0 || llvm::isPowerOf2_32(v); }); + })) { + return emitError() << "Each offset basis must be 0 or a power of two."; + } + } + + return success(); +} + +PaddedSharedEncodingAttr PaddedSharedEncodingAttr::get( + MLIRContext *context, ArrayRef> intervalPads, + ArrayRef order, ArrayRef shape, + CTAEncodingAttr ctaLayout) { + auto outDimNames = standardOutDimNames(context, shape.size()); + StringAttr kOffset = StringAttr::get(context, "offset"); + + // Create identity mapping based on shape and order + LinearLayout linearComponent = + identityStandardND(kOffset, SmallVector(shape), order); + linearComponent = combineCtaCgaWithShape(linearComponent, ctaLayout, shape); + + return get(context, intervalPads, linearComponent); +} + +PaddedSharedEncodingAttr PaddedSharedEncodingAttr::get( + MLIRContext *context, ArrayRef> intervalPads, + LinearLayout linearComponent) { + SmallVector intervals, paddings; + intervals.reserve(intervalPads.size()); + paddings.reserve(intervalPads.size()); + for (auto [interval, padding] : intervalPads) { + intervals.push_back(interval); + paddings.push_back(padding); + } + return get(context, intervals, paddings, linearComponent); +} + +SmallVector +PaddedSharedEncodingAttr::basesPerDim(StringAttr dimName, + bool skipBroadcast) const { + const auto &ll = getLinearComponent(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), dimName, rank, skipBroadcast); +} + +int64_t PaddedSharedEncodingAttr::getPaddedSize(ArrayRef shape) const { + int64_t unpaddedSize = product(shape); + int64_t paddingSize = 0; + for (auto [interval, padding] : + llvm::zip_equal(getIntervals(), getPaddings())) { + paddingSize += (unpaddedSize >> llvm::Log2_32(interval)) + << llvm::Log2_32(padding); + // There is no need for padding after the last element + if (unpaddedSize % interval == 0) + paddingSize -= padding; + } + return unpaddedSize + paddingSize; +} + +SmallVector +PaddedSharedEncodingAttr::orderPerDim(StringAttr dimName, + ArrayRef defaultOrder) const { + return orderPerDimImpl(getLinearComponent(), dimName, defaultOrder); +} + +SmallVector PaddedSharedEncodingAttr::getOrder() const { + auto rank = getLinearComponent().getNumOutDims(); + SmallVector order(rank); + // Choose [rank-1, rank-2, ... 0] as the default order in case + // there are dims that do not move in the offsets + std::iota(order.rbegin(), order.rend(), 0); + + return orderPerDim(StringAttr::get(getContext(), "offset"), order); +} + +CTAEncodingAttr PaddedSharedEncodingAttr::getCTALayout() const { + auto splitNum = basesPerDim(StringAttr::get(getContext(), "block")); + return linearToCTAEncodingAttr(getLinearComponent(), splitNum); +} +//===----------------------------------------------------------------------===// +// NVMMAShared encoding +//===----------------------------------------------------------------------===// + +Attribute NVMMASharedEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + // Parse the data as a dictionary + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned swizzlingByteWidth; + bool transposed = false; + bool fp4Padded = false; + unsigned elementBitWidth; + unsigned layoutRank = 2; + Attribute ctaAttr = nullptr; + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "swizzlingByteWidth") { + if (parseUInt(parser, attr, swizzlingByteWidth, "swizzlingByteWidth") + .failed()) + return {}; + } else if (attr.getName() == "transposed") { + if (parseBool(parser, attr, transposed, "transposed").failed()) + return {}; + } else if (attr.getName() == "elementBitWidth") { + if (parseUInt(parser, attr, elementBitWidth, "elementBitWidth").failed()) + return {}; + } else if (attr.getName() == "fp4Padded") { + if (parseBool(parser, attr, fp4Padded, "fp4Padded").failed()) + return {}; + } else if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + } else if (attr.getName() == "rank") { + if (parseUInt(parser, attr, layoutRank, "rank").failed()) + return {}; + } else { + parser.emitError(parser.getNameLoc(), "unexpected key: ") + << attr.getName().strref(); + return {}; + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, layoutRank); + if (!CTALayout.has_value()) + return {}; + + return parser.getChecked( + parser.getContext(), swizzlingByteWidth, transposed, elementBitWidth, + fp4Padded, *CTALayout); +} + +void NVMMASharedEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "swizzlingByteWidth = " << getSwizzlingByteWidth() // + << ", transposed = " << getTransposed() // + << ", elementBitWidth = " << getElementBitWidth(); + if (getFp4Padded()) { + // Print only in this case to reduce the noise for the more common case. + printer << ", fp4Padded = true"; + } + unsigned rank = getCTALayout().getCTAOrder().size(); + auto *ctx = getContext(); + auto defaultLayout = CTAEncodingAttr::getDefault(ctx, rank); + if (getCTALayout() == defaultLayout && rank != 2) { + printer << ", rank = " << rank; + } else { + maybePrintCTALayout(ctx, printer, getCTALayout(), rank); + } + printer << "}>"; +} + +int NVMMASharedEncodingAttr::getVec() const { + if (getSwizzlingByteWidth() == 0) + return 1; + return 128 / getElementBitWidth(); +} + +int NVMMASharedEncodingAttr::getPerPhase() const { + if (getSwizzlingByteWidth() == 0) + return 1; + return 128 / getSwizzlingByteWidth(); +} + +int NVMMASharedEncodingAttr::getMaxPhase() const { + if (getSwizzlingByteWidth() == 0) + return 1; + return getSwizzlingByteWidth() / 16; +} + +int32_t NVMMASharedEncodingAttr::getAlignment() const { + return 128 * getMaxPhase(); +} + +//===----------------------------------------------------------------------===// +// AMDRotatingShared encoding +//===----------------------------------------------------------------------===// + +Attribute AMDRotatingSharedEncodingAttr::parse(AsmParser &parser, Type type) { + return parseSwizzledEncoding(parser, type); +} + +void AMDRotatingSharedEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "vec = " << getVec() // + << ", perPhase = " << getPerPhase() + << ", maxPhase = " << getMaxPhase() // + << ", order = [" << getOrder() << "]"; + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getOrder().size()); + printer << "}>"; +} + +//===----------------------------------------------------------------------===// +// Mfma encoding +//===----------------------------------------------------------------------===// +// TODO: there is a lot of common code with MmaEncoding here + +bool AMDMfmaEncodingAttr::hasUnitTilesPerWarp() const { + return llvm::all_of(getTilesPerWarp(), [](int x) { return x == 1; }); +} + +SmallVector +AMDMfmaEncodingAttr::getInstrShapeForOperand(int kWidth, int opIdx) const { + auto mnkDim = getInstrShape(); + unsigned mDim = mnkDim[0]; + unsigned nDim = mnkDim[1]; + assert((mDim == nDim) && (mDim == 32 || mDim == 16 || mDim == 4) || + (mDim == 64 && nDim == 4) || (mDim == 4 && nDim == 64)); + + constexpr int warpSize = 64; // MFMA is always based on the 64-wide warps. + int kGroups = warpSize / std::min(mDim, nDim); // for 64x4 and 4x64, + // kGroups = 16 + int64_t kDim = kWidth * kGroups; + + if (opIdx == 0) + return {mDim, kDim}; + else + assert(opIdx == 1); + return {kDim, nDim}; +} + +SmallVector AMDMfmaEncodingAttr::getRepOrder() const { + return getMatrixOrder(getRank(), /*rowMajor*/ true); +} + +SmallVector +AMDMfmaEncodingAttr::getRepOrderForOperand(int opIdx) const { + return getOrderForDotOperand(opIdx, getRank(), /*kContig*/ true); +} + +SmallVector +AMDMfmaEncodingAttr::getRepForOperand(ArrayRef operandShape, + int kWidth, int opIdx) const { + auto operandTileShape = getInstrShapeForOperand(kWidth, opIdx); + auto rank = operandShape.size(); + auto warpsPerCTA = getWarpsPerCTA(); + auto tilesPerWarp = getTilesPerWarp(); + + int numRepBatch = + rank == 3 ? std::max(1, operandShape[0] / warpsPerCTA[0]) : 1; + if (opIdx == 0) + return { + numRepBatch, + std::max(1, operandShape[rank - 2] / + (operandTileShape[0] * tilesPerWarp[rank - 2] * + warpsPerCTA[rank - 2])) * + tilesPerWarp[rank - 2], + std::max(1, operandShape[rank - 1] / operandTileShape[1])}; + else { + assert(opIdx == 1); + return { + numRepBatch, + std::max(1, operandShape[rank - 2] / operandTileShape[0]), + std::max(1, operandShape[rank - 1] / + (operandTileShape[1] * tilesPerWarp[rank - 1] * + warpsPerCTA[rank - 1])) * + tilesPerWarp[rank - 1]}; + } +} + +SwizzledSharedEncodingAttr AMDMfmaEncodingAttr::composeSharedLayoutForOperand( + CTAEncodingAttr ctaLayout, int operandIdx, ArrayRef operandShape, + ArrayRef sharedOrder, unsigned vectorSize, unsigned elemBitWidth, + bool needTrans) const { + int kDimIndex = operandIdx == 0 ? 1 : 0; + + // Disable swizzling for scales + if (operandIdx >= 2) { + return SwizzledSharedEncodingAttr::get(getContext(), 1, 1, 1, sharedOrder, + ctaLayout); + } + + if (needTrans) + kDimIndex = 1 - kDimIndex; + + bool isKContig = sharedOrder[0] == kDimIndex; + // GFX950 supports LDS transpose load instructions, so we need swizzling even + // when K dimension is not the contiguous dimension. + bool isGFX950 = getVersion() == 4; + bool swizzleNonKContig = + isGFX950 && (elemBitWidth == 8 || elemBitWidth == 16); + + if (!isKContig && !swizzleNonKContig) { + // Do not swizzle. In this case accesses will go in different banks even + // without swizzling. + return SwizzledSharedEncodingAttr::get(getContext(), 1, 1, 1, sharedOrder, + ctaLayout); + } + + const unsigned numBanks = isGFX950 ? 64 : 32; + const unsigned bankBitWidth = 32; + const unsigned simdWidth = 16; + + // Number of inner dimension rows per one pattern repeat + int innerDimLength = operandShape[sharedOrder[0]]; + int elemsPerOneBanksRow = (numBanks * bankBitWidth) / elemBitWidth; + + int perPhase = std::max(1, elemsPerOneBanksRow / innerDimLength); + int maxPhase = + std::max(std::min(simdWidth / perPhase, innerDimLength / vectorSize), 1u); + + // TODO (zhanglx): figure out better parameters for mfma4 + if (getInstrShape()[0] == 4) + maxPhase = 4; + + return SwizzledSharedEncodingAttr::get(getContext(), vectorSize, perPhase, + maxPhase, sharedOrder, ctaLayout); +} + +//===----------------------------------------------------------------------===// +// Wmma encoding +//===----------------------------------------------------------------------===// + +SmallVector AMDWmmaEncodingAttr::getRepOrder() const { + return getMatrixOrder(getRank(), /*rowMajor*/ true); +} + +SmallVector +AMDWmmaEncodingAttr::getRepOrderForOperand(int opIdx) const { + return getOrderForDotOperand(opIdx, getRank(), /*kContig*/ true); +} + +SmallVector +AMDWmmaEncodingAttr::getRepForOperand(ArrayRef operandShape, int kDim, + int opIdx) const { + auto mnkDim = getInstrShape(); + SmallVector operandTileShape{opIdx == 0 ? mnkDim[0] : kDim, + opIdx == 0 ? kDim : mnkDim[1]}; + + assert(operandTileShape.size() == 2); + auto warpsPerCTA = getWarpsPerCTA(); + auto tilesPerWarp = getTilesPerWarp(); + + auto rank = operandShape.size(); + assert(rank == 2 || rank == 3); + int numRepBatch = + rank == 3 ? std::max(1, operandShape[0] / warpsPerCTA[0]) : 1; + if (opIdx == 0) + return { + numRepBatch, + std::max(1, operandShape[rank - 2] / + (operandTileShape[0] * tilesPerWarp[rank - 2] * + warpsPerCTA[rank - 2])) * + tilesPerWarp[rank - 2], + std::max(1, operandShape[rank - 1] / operandTileShape[1])}; + else { + assert(opIdx == 1); + return { + numRepBatch, + std::max(1, operandShape[rank - 2] / operandTileShape[0]), + std::max(1, operandShape[rank - 1] / + (operandTileShape[1] * tilesPerWarp[rank - 1] * + warpsPerCTA[rank - 1])) * + tilesPerWarp[rank - 1]}; + } +} + +SwizzledSharedEncodingAttr AMDWmmaEncodingAttr::composeSharedLayoutForOperand( + CTAEncodingAttr ctaLayout, int operandIdx, ArrayRef operandShape, + ArrayRef sharedOrder, unsigned kWidth, unsigned elemBitWidth, + bool needTrans) const { + int kDimIndex = operandIdx == 0 ? 1 : 0; + bool isKContig = sharedOrder[0] == kDimIndex; + + if (!isKContig) { + // Do not swizzle. In this case accesses will go in different banks even + // without swizzling. + return SwizzledSharedEncodingAttr::get(getContext(), 1, 1, 1, sharedOrder, + ctaLayout); + } + + // max vectorization size for ds_load is 128 bits + int vectorSize = std::min(kWidth * elemBitWidth, 128u) / elemBitWidth; + + const int numBanks = 32; + const int bankBitWidth = 32; + + // Number of inner dimension rows per one pattern repeat + int innerDimLength = operandShape[sharedOrder[0]]; + int elemsPerOneBanksRow = (numBanks * bankBitWidth) / elemBitWidth; + + int perPhase = std::max(1, elemsPerOneBanksRow / innerDimLength); + // for both RDNA3 and RDNA4, the M/N dimension of wmma is 16 + // This represents the max number of rows that can be accessed + // at the same time + int mDim = getInstrShape()[0]; + int maxPhase = + std::max(std::min(mDim / perPhase, innerDimLength / vectorSize), 1); + + return SwizzledSharedEncodingAttr::get(getContext(), vectorSize, perPhase, + maxPhase, sharedOrder, ctaLayout); +} + +//===----------------------------------------------------------------------===// +// Mma encoding +//===----------------------------------------------------------------------===// + +bool NvidiaMmaEncodingAttr::isVolta() const { return getVersionMajor() == 1; } + +bool NvidiaMmaEncodingAttr::isTuring() const { + return getVersionMajor() == 2 && getVersionMinor() == 1; +} + +bool NvidiaMmaEncodingAttr::isAmpere() const { return getVersionMajor() == 2; } + +bool NvidiaMmaEncodingAttr::isHopper() const { return getVersionMajor() == 3; } + +SmallVector NvidiaMmaEncodingAttr::getRepOrder() const { + return getMatrixOrder(getRank(), /*rowMajor*/ true); +} + +SmallVector +NvidiaMmaEncodingAttr::getRepOrderForOperand(int opIdx) const { + return getOrderForDotOperand(opIdx, getRank(), /*kContig*/ true); +} + +SmallVector +NvidiaMmaEncodingAttr::getRepForOperand(ArrayRef shape, int bitwidth, + int kWidth, int opIdx) const { + assert(kWidth >= std::max(32 / bitwidth, 1) && + "kWidth must be >= max(32 / bitwidth, 1) for this function to be " + "well-defined"); + auto rank = shape.size(); + // Broadcast long K + auto warpsPerCTA = to_vector(getWarpsPerCTA()); + auto kDim = opIdx == 0 ? rank - 1 : rank - 2; + warpsPerCTA[kDim] = 1; + + SmallVector tileSize; + if (rank == 3) { + tileSize.push_back(1); + } + // warpSizeK * (warpRepK * VecBitWidth) + auto tileBitWidthK = (isAmpere() && bitwidth == 64) ? (4 * 256) : (4 * 64); + if (opIdx == 0) { + // m x k + tileSize.push_back(16); + tileSize.push_back(tileBitWidthK / bitwidth); + } else { + // k x n + // Hopper path never uses the n value, since this method is only invoked + // for in-RF (dotOpEnc) operands, but WGMMA only supports in A to be in RF + // so it's fine if the n is incorrect here + tileSize.push_back(tileBitWidthK / bitwidth); + tileSize.push_back(8); + } + + SmallVector numRep; + // Lezcano: This is odd. Why do we always return a vector of size 3? + if (rank != 3) { + numRep.push_back(1); + } + for (auto [s, size, warp] : llvm::zip(shape, tileSize, warpsPerCTA)) { + numRep.push_back(std::max(1, s / (size * warp))); + } + return numRep; +} + +//===----------------------------------------------------------------------===// +// DotOperand Encoding +//===----------------------------------------------------------------------===// + +SmallVector DotOperandEncodingAttr::getRepOrder() const { + if (auto mma = mlir::dyn_cast(getParent())) { + return mma.getRepOrderForOperand(getOpIdx()); + } else if (auto blocked = mlir::dyn_cast(getParent())) { + return to_vector(blocked.getOrder()); + } + llvm::report_fatal_error( + "getRepOrder not implemented for DotOperandEncodingAttr"); + return {}; +} + +CTAEncodingAttr DotOperandEncodingAttr::getCTALayout() const { + auto layout = ::getCTALayout(getParent()).getLinearLayout(); + auto bases = layout.getBases(); + auto kBlock = StringAttr::get(getContext(), "block"); + auto &blockBases = bases[kBlock]; + auto rank = layout.getNumOutDims(); + auto kDim = getOpIdx() == 0 ? rank - 1 : rank - 2; + for (auto &basis : blockBases) { + basis[kDim] = 0; + } + auto dims = layout.getOutDims(); + dims[kDim].second = 1; + return CTAEncodingAttr::get(getContext(), LinearLayout(bases, dims, true)); +} +LogicalResult DotOperandEncodingAttr::verify( + ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError, + unsigned opIdx, Attribute parent, unsigned kWidth) { + if (opIdx != 0 && opIdx != 1) { + return emitError() << "ttg.dot_op opIdx parameter can be 0 or 1, got: " + << opIdx; + } + if (!parent) { + return emitError() << "ttg.dot_op parent parameter cannot be null"; + } + if (auto parentAttr = mlir::dyn_cast(parent)) { + if (kWidth != 0 && !(parentAttr.isAmpere() || parentAttr.isHopper())) + return emitError() << "ttg.dot_op kWidth parameter can only be " + "non-zero for Ampere or Hopper MMA parent"; + if (kWidth == 0 && (parentAttr.isAmpere() || parentAttr.isHopper())) + return emitError() << "ttg.dot_op kWidth parameter is mandatory for " + "Ampere or Hopper MMA parent"; + if (opIdx != 0 && parentAttr.isHopper()) + return emitError() + << "ttg.dot_op opIdx parameter must be 0 for " + "Hopper MMA parent, since Hopper WGMMA only allows first " + "operand to be in registers"; + return success(); + } + + if (auto parentAttr = mlir::dyn_cast(parent)) { + if (parentAttr.getVersion() == 1 && (kWidth != 8 && kWidth != 16)) + return emitError() + << "ttg.dot_op kWidth parameter must be 8/16 for WMMA v1 " + "(including packed cases for `scaled_dot`)"; + if (parentAttr.getVersion() == 2 && !llvm::is_contained({4, 8, 16}, kWidth)) + return emitError() + << "ttg.dot_op kWidth parameter must be 4/8/16 for WMMA v2 " + "(including packed cases for `scaled_dot`)"; + if (parentAttr.getVersion() == 3 && !llvm::is_contained({2, 8, 16}, kWidth)) + return emitError() + << "ttg.dot_op kWidth parameter must be 2/8/16 for WMMA v3"; + return success(); + } + + if (auto parentAttr = mlir::dyn_cast(parent)) { + if (kWidth == 0) + return emitError() << "ttg.dot_op kWidth parameter is mandatory for " + "MFMA parent"; + return success(); + } + + if (auto parentAttr = mlir::dyn_cast(parent)) { + if (kWidth != 0) + return emitError() << "ttg.dot_op kWidth parameter is not supported " + "when the parent is a blocked layout"; + return success(); + } + + return emitError() << "ttg.dot_op unexpected parent layout: " << parent; +} + +//===----------------------------------------------------------------------===// +// ASM Interface (i.e.: alias) +//===----------------------------------------------------------------------===// + +class TritonGPUOpAsmInterface : public OpAsmDialectInterface { +public: + using OpAsmDialectInterface::OpAsmDialectInterface; + + AliasResult getAlias(Attribute attr, raw_ostream &os) const override { + // Encoding attributes + if (auto mmaAttr = mlir::dyn_cast(attr)) { + os << "mma"; + return AliasResult::FinalAlias; + } else if (auto sharedAttr = mlir::dyn_cast(attr)) { + os << "shared"; + return AliasResult::FinalAlias; + } else if (auto blockedAttr = mlir::dyn_cast(attr)) { + os << "blocked"; + return AliasResult::FinalAlias; + } else if (auto linearAttr = mlir::dyn_cast(attr)) { + os << "linear"; + return AliasResult::FinalAlias; + } /* else if (auto sliceAttr = dyn_cast(attr)) { + os << "slice"; + return AliasResult::FinalAlias; + } */ + // Memory space attributes + if (auto smem = mlir::dyn_cast(attr)) { + os << "smem"; + return AliasResult::FinalAlias; + } + return OpAsmDialectInterface::getAlias(attr, os); + } +}; + +struct TritonGPUInferLayoutInterface + : public triton::DialectInferLayoutInterface { + using DialectInferLayoutInterface::DialectInferLayoutInterface; + + LogicalResult + inferReduceOpEncoding(Attribute operandEncoding, unsigned axis, + Attribute &resultEncoding, + std::optional loc) const override { + resultEncoding = + SliceEncodingAttr::get(getDialect()->getContext(), axis, + cast(operandEncoding)); + return success(); + } + + // Infer the encoding of a tt.trans(x) given the encoding of x. + // + // Our goal is to choose an encoding so that the trans is a "nop". For + // example, in a blocked encoding, the same GPU threads hold the same + // elements, they're just "renamed" -- what was element [i,j] of the tensor is + // now element [j,i], but that element is held by the same GPU thread. + // + // For most properties of the encoding, we let + // outputEnc.prop = inputEnc.prop * trans.order, + // where `x * y` means we apply permutation y to x. + // + // This works because prop[i] tells you something about the i'th dimension of + // the tensor. (For example, sizePerThread[2] == 4 means that one GPU thread + // contains 4 elements along dim 2 of the tensor.) The transpose reorders the + // dimensions according to the perm trans.order, so we achieve our goal of + // having a "nop" transpose by reordering the values in the prop the same way. + // + // The big exception to this is the encoding's `order`. + // + // An encoding's order is a list of dimensions, from fastest moving (most + // minor) to slowest moving. Thus enc.order[i] does not tell you something + // about the i'th dimension of the tensor, and it would be disasterously + // incorrect to do enc.order * trans.order. + // + // But! If we invert enc.order, it *does* meet this criterion. For example, + // if enc.order = [2,0,1], inverse(enc.order) = [1,2,0]. If you stare at it, + // you'll see that inverse(enc.order)[i] == j means that dimension i is the + // j'th most minor. Therefore we can safely permute *this* by trans.order. + // + // Thus we have + // + // outputEnc.order = inverse(inverse(inputEnc.order) * trans.order) + // = inverse(trans.order) * inputEnc.order. + // + LogicalResult + inferTransOpEncoding(Attribute operandEncoding, ArrayRef shape, + ArrayRef order, Attribute &resultEncoding, + std::optional loc) const override { + // Note: inferFooOpEncoding should not crash if given invalid inputs, which + // happens when someone creates invalid IR. If we return failure() on + // error, then MLIR will generate a helpful error message. + if (isIota(order)) { + resultEncoding = operandEncoding; + return success(); + } + if (shape.size() != order.size()) { + return emitOptionalError(loc, "shape and order rank do not match: ", + shape.size(), " vs ", order.size()); + } + auto checkRank = [&](unsigned rank) { + if (rank != order.size()) { + return emitOptionalError(loc, "rank of encoding does not match order: ", + rank, " vs ", order.size()); + } + return success(); + }; + auto *ctx = getDialect()->getContext(); + + auto permuteCTALayout = [ctx](CTAEncodingAttr layout, + ArrayRef order) { + auto ll = transposeLinearLayout(layout.getLinearLayout(), order); + return CTAEncodingAttr::get(ctx, std::move(ll)); + }; + + auto invOrder = inversePermutation(order); + SmallVector invOrderUnsigned(invOrder.begin(), invOrder.end()); + + if (auto enc = dyn_cast(operandEncoding)) { + if (failed(checkRank(enc.getCTALayout().getRank()))) + return failure(); + + CTAEncodingAttr ctaLayout = permuteCTALayout(enc.getCTALayout(), order); + resultEncoding = SwizzledSharedEncodingAttr::get( + ctx, enc.getVec(), enc.getPerPhase(), enc.getMaxPhase(), + applyPermutation(invOrderUnsigned, enc.getOrder()), ctaLayout); + return success(); + } + + if (auto enc = dyn_cast(operandEncoding)) { + if (order == ArrayRef({1, 0})) { + if (failed(checkRank(enc.getCTALayout().getRank()))) + return failure(); + + CTAEncodingAttr ctaLayout = permuteCTALayout(enc.getCTALayout(), order); + resultEncoding = NVMMASharedEncodingAttr::get( + ctx, enc.getSwizzlingByteWidth(), !enc.getTransposed(), + enc.getElementBitWidth(), enc.getFp4Padded(), ctaLayout); + return success(); + } + } + + if (auto enc = dyn_cast(operandEncoding)) { + if (failed(checkRank(enc.getCTALayout().getRank()))) + return failure(); + + CTAEncodingAttr ctaLayout = permuteCTALayout(enc.getCTALayout(), order); + resultEncoding = BlockedEncodingAttr::get( + ctx, applyPermutation(enc.getSizePerThread(), order), + applyPermutation(enc.getThreadsPerWarp(), order), + applyPermutation(enc.getWarpsPerCTA(), order), + applyPermutation(invOrderUnsigned, enc.getOrder()), ctaLayout); + return success(); + } + // Generic case + auto padded = dyn_cast(operandEncoding); + + auto ll = padded ? padded.getLinearComponent() + : toLinearLayout(shape, operandEncoding); + if (failed(checkRank(ll.getNumOutDims()))) + return failure(); + auto transposedLl = transposeLinearLayout(ll, order); + if (isa(operandEncoding)) { + resultEncoding = LinearEncodingAttr::get(ctx, std::move(transposedLl)); + } else if (padded) { + resultEncoding = PaddedSharedEncodingAttr::get(ctx, padded.getIntervals(), + padded.getPaddings(), + std::move(transposedLl)); + } else { + auto shared = cast(operandEncoding); + resultEncoding = SharedLinearEncodingAttr::get( + ctx, std::move(transposedLl), shared.getAlignment()); + } + return success(); + } + + LogicalResult + inferExpandDimsOpEncoding(Attribute operandEncoding, unsigned axis, + Attribute &resultEncoding, + std::optional location) const override { + auto sliceEncoding = mlir::dyn_cast(operandEncoding); + if (!sliceEncoding) + return emitOptionalError( + location, "ExpandDimsOp operand encoding must be SliceEncodingAttr"); + if (sliceEncoding.getDim() != axis) + return emitOptionalError( + location, "Incompatible slice dimension for ExpandDimsOp operand"); + resultEncoding = sliceEncoding.getParent(); + return success(); + } + + LogicalResult + inferDotOpEncoding(Attribute operandEncoding, unsigned opIdx, + Attribute retEncoding, + std::optional location) const override { + auto mmaRetEncoding = mlir::dyn_cast(retEncoding); + if (mmaRetEncoding && mmaRetEncoding.isHopper()) { + auto dotOpEnc = mlir::dyn_cast(operandEncoding); + if (!mlir::isa( + operandEncoding) && + !(opIdx == 0 && dotOpEnc && dotOpEnc.getOpIdx() == 0 && + mlir::isa(dotOpEnc.getParent()))) { + return emitOptionalError( + location, "unexpected operand layout for NvidiaMmaEncodingAttr v3"); + } + } else if (auto dotOpEnc = + mlir::dyn_cast(operandEncoding)) { + if (opIdx != dotOpEnc.getOpIdx()) + return emitOptionalError(location, "Wrong opIdx"); + if (retEncoding != dotOpEnc.getParent()) + return emitOptionalError(location, "Incompatible parent encoding"); + } else + return emitOptionalError( + location, "Dot's a/b's encoding should be of DotOperandEncodingAttr"); + return success(); + } + + LogicalResult + verifyDotOpEncodingCompatibility(Operation *op, Attribute operandEncodingA, + Attribute operandEncodingB) const override { + auto aEncoding = + mlir::dyn_cast(operandEncodingA); + auto bEncoding = + mlir::dyn_cast(operandEncodingB); + if (!aEncoding && !bEncoding) + return mlir::success(); + if (!aEncoding || !bEncoding) + return op->emitError("mismatching encoding between A and B operands"); + // Verify that the encodings are valid. + if (aEncoding.getKWidth() != bEncoding.getKWidth()) + return op->emitError("mismatching kWidth between A and B operands"); + + // Check if we have already selected an MMA version for Nvidia. If so, + // validate that the encodings are correct and compatible. + auto mmaAEncoding = + dyn_cast_or_null(aEncoding.getParent()); + auto mmaBEncoding = + dyn_cast_or_null(bEncoding.getParent()); + auto dotOp = cast(op); + auto resEnc = dotOp.getResult().getType().getEncoding(); + auto mmaResEncoding = dyn_cast(resEnc); + if (mmaAEncoding || mmaBEncoding || mmaResEncoding) { + // Check that they are all set and have the same version. + if (!mmaAEncoding || !mmaBEncoding || !mmaResEncoding) + return op->emitError("mismatching MMA encoding"); + auto mmaBEncoding = cast(bEncoding.getParent()); + if (mmaAEncoding.getVersionMajor() != mmaBEncoding.getVersionMajor() || + mmaAEncoding.getVersionMajor() != mmaResEncoding.getVersionMajor()) { + return op->emitError("mismatched MMA version."); + } + // Verify that the operands are supported on the selected MMA version. + if (!supportMMA(dotOp, mmaResEncoding.getVersionMajor())) + return op->emitError("unsupported MMA version"); + } + return success(); + } + + // Given a src shape + encoding and a dst shape, our goal is to compute a dst + // encoding that makes the reshape a "nop". That is, if GPU thread [x,y,z] + // contains elements [a,b,c,d] before the reshape, it contains those same + // elements after the reshape, they're just "renamed". + // + // Using legacy layouts, a dst encoding that satisfies this property may not + // exist. Here are some positive and negative examples. + // + // - NOT OK: 4x4 order=[0,1] -> 16. Reshape merges elements so + // dim 1 is the fastest-changing in the dst, but the src has the opposite + // order. + // - OK: 2x2x32 order=[1,0,2] -> 4x32. We choose dst order [0,1]. + // What's important is that the 2x2 dimensions appear in major-to-minor + // order. + // - NOT OK: 32x32 sizePerThread=[2,2] -> 1024. Thread 0 in the src + // contains elements [(0,0), (0,1), (1,0), and (1,1)]. We cannot express + // this with an encoding based on the dst shape. + // - OK: 32x4 sizePerThread=[4,4] -> 128. dst with sizePerThread=[16] will + // contain the same elements as before. + // + // With linear layouts, we can always find a dst encoding that satisfies + // this property. See inferReshapeOpEncoding. + // + // Users of this function require that it is symmetrical: if + // (srcShape,srcEnc,dstShape) => dstEnc, then (dstShape,dstEnc,srcShape) => + // srcEnc. + LogicalResult inferReshapeOpLegacyEncoding(ArrayRef srcShape, + Attribute srcEnc, + ArrayRef dstShape, + Attribute &dstEnc) const { + auto src = mlir::dyn_cast(srcEnc); + if (!src) { + return failure(); + } + + // Nop reshape; we can always infer an encoding. + if (srcShape == dstShape) { + dstEnc = srcEnc; + return success(); + } + + // default -> default encoding is always a nop. + auto context = srcEnc.getContext(); + int32_t numWarps = product(src.getWarpsPerCTA()); + int32_t threadsPerWarp = product(src.getThreadsPerWarp()); + int32_t numCTAs = product(src.getCTALayout().getCTAsPerCGA()); + if (srcEnc == getDefaultBlockedEncoding(context, srcShape, numWarps, + threadsPerWarp, numCTAs)) { + dstEnc = getDefaultBlockedEncoding(context, dstShape, numWarps, + threadsPerWarp, numCTAs); + return success(); + } + + // Cowardly refuse to handle encodings with multiple CTAs. CTAsPerCGA + // should be like the other fields in blocked encoding, but I'm not sure how + // to handle CTASplitNum. + auto srcCTALayout = src.getCTALayout(); + if (!all_of(srcCTALayout.getCTAsPerCGA(), + [](int32_t x) { return x == 1; }) || + !all_of(srcCTALayout.getCTASplitNum(), + [](int32_t x) { return x == 1; })) { + return failure(); + } + + // Cowardly refuse to handle encodings where shape[dim] is not divisible by + // sizePerThread[dim], threadsPerWarp[dim], and warpsPerCTA[dim]. (We make + // an exception if the block is larger than the shape.) + auto checkDivisibility = [&](StringRef name, ArrayRef subblock) { + for (int dim = 0; dim < srcShape.size(); dim++) { + if (srcShape[dim] >= subblock[dim] && + srcShape[dim] % subblock[dim] != 0) { + return failure(); + } + } + return success(); + }; + if (!succeeded( + checkDivisibility("sizePerThread", src.getSizePerThread())) || + !succeeded( + checkDivisibility("threadsPerWarp", src.getThreadsPerWarp())) || + !succeeded(checkDivisibility("warpsPerCTA", src.getWarpsPerCTA()))) { + return failure(); + } + + SmallVector, SmallVector>> decomp = + getReshapeDecomposition(srcShape, dstShape); + + // enc.order[i] == j means that dimension j is the enc.order[i]'th most + // minor. But what we usually want is the inverse: inverse(enc.order)[i] = j + // means that dimension i is the j'th most minor (larger means more major). + auto srcInvOrder = inversePermutation(src.getOrder()); + + // If src dims [a,b,c] are to be merged, then they must be consecutive in + // physical order, with `a` being the most major. + for (const auto &[srcDims, dstDims] : decomp) { + if (!isConsecutive(to_vector(reverse(gather(srcInvOrder, srcDims))))) { + return failure(); + } + } + + // If src dims [a,b,c] are to be merged, then `c` must fill up sizePerThread + // / threadsPerWarp / blocksPerCTA before `b` can have any non-1 values. + // Examples: + // + // - NOT OK: shape=[4,4,4], sizePerThread=[1,2,2]. + // The total sizePerThread for dim 2 is 2, which is less than dim 2's + // size of 4. Therefore dim 1 cannot have non-1 sizePerThread. + // + // - OK: shape=[4,4,4], sizePerThread=[1,2,4]. + // Dim 2's sizePerThread covers its whole size, so dim 1 is allowed to + // have non-1 sizePerThread. + // + // - NOT OK: shape=[4,4,4], sizePerThread=[2,1,4]. + // Dim 1's sizePerThread does not cover its whole size, so dim 0 is not + // allowed to have non-1 sizePerThread. + // + // - NOT OK: shape=[4,4,4], sizePerThread=[1,1,2], + // threadsPerWarp=[1,2,1]. + // Dim 2 has 2 elems per thread and 1 thread per warp. 2*1 is less than + // dim 2's size. Therefore dim 1 must have threadsPerWarp=1. + // + // In addition, the encoding's block can be larger than the shape, but only + // in the most-major dimension of each decomposed chunk, and only after + // we've "used up" the more minor dims. Examples: + // + // - OK: shape=[4,4,4], sizePerThread=[1,2,4], threadsPerWarp=[16,2,1], + // warpsPerCTA=[4,1,1]. + // The whole size of dims 0 and 1 are covered by sizePerThread * + // threadsPerWarp. Therefore dim 2 is allowed to have threadsPerWarp and + // warpsPerCTA larger than its size. + for (const auto &[srcDims, dstDims] : decomp) { + auto shapeRemaining = gather(srcShape, srcDims); + auto checkSubblock = [&, srcDims = srcDims](ArrayRef subblock) { + // Iterate minor-to-major (i==0 is most major). + for (int i = srcDims.size() - 1; i >= 0; i--) { + int dim = srcDims[i]; + if (subblock[dim] == 1) { + continue; + } + + // Check that more-minor dims all have 1 in shapeRemaining. + for (int j = i + 1; j < srcDims.size(); j++) { + if (shapeRemaining[j] != 1) { + return failure(); + } + } + + if (shapeRemaining[i] >= subblock[dim]) { + assert(shapeRemaining[i] % subblock[dim] == 0); // checked earlier + shapeRemaining[i] /= subblock[dim]; + } else { + shapeRemaining[i] = 0; + } + + // Is the block larger than the shape in this dimension? This is OK + // only if we're the most-major dimension of the chunk and in all + // future chunks, only this most-major dim has a non-1 size. + if (shapeRemaining[i] == 0 && i != 0) { + return failure(); + } + } + return success(); + }; + if (!succeeded(checkSubblock(src.getSizePerThread())) || + !succeeded(checkSubblock(src.getThreadsPerWarp())) || + !succeeded(checkSubblock(src.getWarpsPerCTA()))) { + return failure(); + } + } + + // Given e.g. src.getSizePerThread(), computeSubblockSize computes e.g. + // dst.getSizePerThread(). This should be called for each of sizePerThread, + // threadsPerWarp, and warpsPerCTA, in that order. + SmallVector dstShapeRemaining(dstShape); + auto computeSubblockSize = [&](ArrayRef srcSubblock, + SmallVector &dstSubblock, + StringRef fieldName) -> LogicalResult { + // The dst subblock is "filled up" greedily starting with the most minor + // dim. When we're done, we are left with a smaller shape, of size + // dstShape / dstSubblock, which we store in dstShapeRemaining and use for + // the next call to computeSubblockSize. + dstSubblock.resize(dstShape.size()); + for (const auto &[srcDims, dstDims] : decomp) { + int64_t subblockRemaining = product(gather(srcSubblock, srcDims)); + for (int i = dstDims.size() - 1; i >= 0; i--) { + auto &val = dstSubblock[dstDims[i]]; + auto &shapeRemaining = dstShapeRemaining[dstDims[i]]; + val = std::min(subblockRemaining, shapeRemaining); + + assert(shapeRemaining % val == 0); // Checked earlier. + subblockRemaining /= val; + shapeRemaining /= val; + } + + // If there are any elems remaining in the subblock, it must be because + // the block is larger than the shape. This excess goes into the + // most-major dim of the subblock. + dstSubblock[dstDims[0]] *= subblockRemaining; + } + return success(); + }; + + SmallVector dstSizePerThread; + SmallVector dstThreadsPerWarp; + SmallVector dstWarpsPerCTA; + if (!succeeded(computeSubblockSize(src.getSizePerThread(), dstSizePerThread, + "sizePerThread")) || + !succeeded(computeSubblockSize(src.getThreadsPerWarp(), + dstThreadsPerWarp, "threadsPerWarp")) || + !succeeded(computeSubblockSize(src.getWarpsPerCTA(), dstWarpsPerCTA, + "warpsPerCTA"))) { + return failure(); + } + + // Since we know that each set of srcDims is consecutive, we can + // meaningfully sort decomp by the physical order of the src dimensions, + // major-to-minor. This will also be the order of the dst dimensions. + llvm::sort(decomp, [&](const auto &a, const auto &b) { + const auto &[srcDimsA, dstDimsA] = a; + const auto &[srcDimsB, dstDimsB] = b; + return srcInvOrder[srcDimsA.front()] < srcInvOrder[srcDimsB.front()]; + }); + + // Compute the dst order. Make the dimensions appear in the same order as + // their corresponding src dimensions. + SmallVector dstInvOrder(dstShape.size()); + int i = 0; + for (const auto &[srcDims, dstDims] : decomp) { + for (auto dim : reverse(dstDims)) { + dstInvOrder[dim] = i++; + } + } + auto dstOrder = inversePermutation(dstInvOrder); + + // CTALayout can be all 1's because we bailed on multi-CTA layouts above. + auto CTALayout = + CTAEncodingAttr::getDefault(src.getContext(), dstShape.size()); + + dstEnc = BlockedEncodingAttr::get(src.getContext(), dstSizePerThread, + dstThreadsPerWarp, dstWarpsPerCTA, + dstOrder, CTALayout); + + return success(); + } + + LogicalResult + verifyLayoutsAreEqual(ArrayRef shape, Attribute expected, + Attribute got, + std::optional loc) const override { + if (expected == got) { + return success(); + } + if (!expected || !got) + return failure(); + + // Check whether the encodings are structurally the same. + if (!areLayoutsEquivalent(shape, cast(expected), + cast(got))) { + return emitOptionalError(loc, "Expected result encoding ", expected, + " but was ", got); + } + return success(); + } + + LogicalResult + inferReshapeOpEncoding(ArrayRef srcShape, Attribute srcEnc, + ArrayRef dstShape, Attribute &dstEnc, + std::optional loc) const override { + if (product(srcShape) != product(dstShape)) { + return emitOptionalError(loc, "numel of dst shape does not match " + "numel of src shape"); + } + auto result = + inferReshapeOpLegacyEncoding(srcShape, srcEnc, dstShape, dstEnc); + if (succeeded(result)) { + return result; + } + if (!isa(srcEnc)) { + return emitOptionalError(loc, + "Failed MemDescReshapeOp encoding inference"); + } + // If the legacy encoding failed use LinearLayouts. + // Once LinearLayouts are more widely used, we can remove + // inferReshapeOpLegacyEncoding and simply use LLs. + + // HACK: We create a dummy tensor type to pass to inferReshapeLinearLayout. + auto ctx = srcEnc.getContext(); + auto fp32Type = IntegerType::get(ctx, 32, IntegerType::Unsigned); + auto srcTy = RankedTensorType::get(srcShape, fp32Type, srcEnc); + LinearLayout ll = + inferReshapeLinearLayout(cast(srcTy), dstShape); + + dstEnc = LinearEncodingAttr::get(srcEnc.getContext(), ll); + return success(); + } + + LogicalResult + inferDefaultJoinOpEncoding(Attribute srcEnc, Attribute &dstEnc, + ArrayRef shape, + std::optional loc) const override { + auto ctx = getContext(); + if (auto enc = mlir::dyn_cast(srcEnc); + enc && enc.getDim() == shape.size()) { + SmallVector joinedShape(shape); + joinedShape.push_back(2); + auto parent = enc.getParent(); + auto parentLL = toLinearLayout(joinedShape, parent); + + Attribute splitEnc; + auto result = inferSplitOpEncoding(parent, splitEnc, joinedShape, loc); + if (succeeded(result) && + areLayoutsEquivalent(shape, cast(splitEnc), + cast(srcEnc))) { + dstEnc = parent; + return success(); + } + } else if (auto enc = mlir::dyn_cast(srcEnc)) { + // JoinOp takes two tensors of shape AxBxC and generates a tensor of shape + // AxBxCx2. The encoding is the same as the input, but with 2 elems per + // thread in the new dimension. The new dimension is the fastest running + // dimension. + auto append = [](ArrayRef vals, int val) { + SmallVector ret(vals); + ret.push_back(val); + return ret; + }; + auto appendMajorDim = [](ArrayRef order) { + SmallVector ret(order); + ret.insert(ret.begin(), ret.size()); + return ret; + }; + auto ctall = enc.getCTALayout().getLinearLayout(); + auto kBlock = StringAttr::get(enc.getContext(), "block"); + auto newDim = standardOutDimNames( + enc.getContext(), ctall.getNumOutDims() + 1)[ctall.getNumOutDims()]; + ctall *= LinearLayout::identity1D(1, kBlock, newDim); + dstEnc = BlockedEncodingAttr::get( + enc.getContext(), append(enc.getSizePerThread(), 2), + append(enc.getThreadsPerWarp(), 1), append(enc.getWarpsPerCTA(), 1), + appendMajorDim(enc.getOrder()), + CTAEncodingAttr::get(enc.getContext(), ctall)); + return success(); + } + + // Append dim to shape + auto ll = toLinearLayout(shape, srcEnc); + SmallVector dstShape(shape.begin(), shape.end()); + dstShape.push_back(1); + ll = ll.reshapeOuts(standardOutDimPairs(ctx, dstShape)); + + // Try join on last dim + auto axis = dstShape.size() - 1; + auto newLl = LinearLayout::empty(); + auto result = + tryJoinOnAxis(ctx, ll, newLl, /*fwdInference=*/true, axis, loc); + + assert(result.succeeded()); + dstEnc = LinearEncodingAttr::get(ctx, newLl); + return success(); + } + + LogicalResult + inferSplitOpEncoding(Attribute srcEnc, Attribute &dstEnc, + ArrayRef shape, + std::optional loc) const override { + // SplitOp takes a tensor of shape AxBxCx2 and generates two tensors of + // shape AxBxC. The input must have 2 elements per thread in the last + // dimension, which must be the fastest running dimension. The result + // encoding is the same as the input, but with the last dimension removed. + auto enc = mlir::dyn_cast(srcEnc); + bool isSimpleSplit = (enc && (enc.getSizePerThread().back() == 2) && + (enc.getThreadsPerWarp().back() == 1) && + (enc.getWarpsPerCTA().back() == 1) && + (enc.getCTALayout().getCTAsPerCGA().back() == 1)); + if (isSimpleSplit) { + SmallVector newOrder(enc.getOrder()); + auto ctall = enc.getCTALayout().getLinearLayout(); + int splitDim = newOrder.size() - 1; + // Remove splitDim from order. + newOrder.erase(std::remove(newOrder.begin(), newOrder.end(), splitDim), + newOrder.end()); + // Remove last dimension from ctall. + ctall = ctall.unsqueezeOut(to_vector(ctall.getOutDimNames()).back()); + dstEnc = BlockedEncodingAttr::get( + enc.getContext(), // + ArrayRef(enc.getSizePerThread()).drop_back(1), + ArrayRef(enc.getThreadsPerWarp()).drop_back(1), + ArrayRef(enc.getWarpsPerCTA()).drop_back(1), ArrayRef(newOrder), + CTAEncodingAttr::get(enc.getContext(), ctall)); + return success(); + } + + auto axis = shape.size() - 1; + if (shape[axis] != 2) { + return emitOptionalError( + loc, "SplitOp input shape should have 2 in the last dim"); + } + + auto ctx = getContext(); + + // Split on last dim + auto ll = toLinearLayout(shape, srcEnc); + auto newLl = LinearLayout::empty(); + auto result = + tryJoinOnAxis(ctx, ll, newLl, /*fwdInference=*/false, axis, loc); + if (!result.succeeded()) { + return failure(); + } + // Remove last dim from newLl (which should be 1) + SmallVector dstShape(shape.begin(), shape.end()); + dstShape.pop_back(); + newLl = newLl.reshapeOuts(standardOutDimPairs(ctx, dstShape)); + dstEnc = LinearEncodingAttr::get(ctx, newLl); + return success(); + } + + LogicalResult + inferFp4ToFpOpEncoding(ArrayRef shape, int axis, Attribute inEnc, + Attribute &outEnc, bool fwdInference, + std::optional loc) const override { + // We implement two legacy layout propagations + // Once we fully migrate to LinearLayouts, we can remove these. + auto *ctx = getContext(); + // The output encoding will only be a legacy encoding if the axis is the + // fastest running dimension. + // FIXME: We should make sure that there are enough elements along the axis + // axis whenever fwdInference is false + if (getOrder(cast(inEnc), shape)[axis] == 0) { + // Dot operand: double kWidth if kDim == axis. + if (auto dotEnc = mlir::dyn_cast(inEnc)) { + auto kWidth = dotEnc.getKWidth(); + if (fwdInference) { + kWidth *= 2; + } else { + if (kWidth > 1) { + // bwd inference + kWidth /= 2; + } else { + return emitOptionalError(loc, + "Fp4ToFpOp requires at least 2 elements " + "per thread in the axis dimension"); + } + } + outEnc = DotOperandEncodingAttr::get(ctx, dotEnc.getOpIdx(), + dotEnc.getParent(), kWidth); + return success(); + } + + // Blocked layout: double elemsPerThread[axis]. + if (auto blockedEnc = mlir::dyn_cast(inEnc)) { + auto sizePerThread = llvm::to_vector(blockedEnc.getSizePerThread()); + if (fwdInference) { + sizePerThread[axis] *= 2; + } else { + if (sizePerThread[axis] > 1) { + sizePerThread[axis] /= 2; + } else { + return emitOptionalError( + loc, "Fp4ToFpOp requires at least 2 elements per " + "thread in the axis dimension"); + } + } + outEnc = BlockedEncodingAttr::get( + ctx, sizePerThread, blockedEnc.getThreadsPerWarp(), + blockedEnc.getWarpsPerCTA(), blockedEnc.getOrder(), + blockedEnc.getCTALayout()); + return success(); + } + } + + auto ll = toLinearLayout(shape, inEnc); + auto newLl = LinearLayout::empty(); + auto result = tryJoinOnAxis(ctx, ll, newLl, fwdInference, axis, loc); + if (!result.succeeded()) + return result; + outEnc = LinearEncodingAttr::get(ctx, newLl); + return success(); + } +}; + +struct TritonGPUVerifyTensorLayoutInterface + : public triton::DialectVerifyTensorLayoutInterface { + using DialectVerifyTensorLayoutInterface::DialectVerifyTensorLayoutInterface; + + LogicalResult verifyTensorLayout( + Attribute layout, RankedTensorType rankedTy, Operation *op, + function_ref makeErr) const override { + auto distr = dyn_cast(layout); + if (!distr) + return makeErr() + << "Non-distributed layout is not allowed in tensor type."; + auto rank = distr.getRepOrder().size(); + if (rank != rankedTy.getRank()) + return makeErr() << "Layout has rank " << rank + << ", but the tensor it's attached to has rank " + << rankedTy.getRank() << "."; + if (llvm::any_of(rankedTy.getShape(), + [](int64_t i) { return !llvm::isPowerOf2_64(i); })) { + return makeErr() << "Layout has shape " << rankedTy.getShape() + << ", but the tensor it's attached to has shape " + << rankedTy.getShape() + << " which is not a power of two."; + } + auto ll = toLinearLayout(rankedTy); + ModuleOp module = op->getParentOfType(); + + // Number of threads per warp. + auto kLane = StringAttr::get(module.getContext(), "lane"); + int moduleThreadsPerWarp = TritonGPUDialect::getThreadsPerWarp(module); + if (ll.getInDimSize(kLane) != moduleThreadsPerWarp) { + return makeErr() << layout << ".\nLayout has " << ll.getInDimSize(kLane) + << " threads per warp, but the module specifies " + << moduleThreadsPerWarp << " threads per warp."; + } + + // Number of warps per CTA. + std::optional moduleWarpsPerCTA = maybeLookupNumWarps(op); + if (!moduleWarpsPerCTA) { + return makeErr() + << "Could not determine the number of warps per CTA. Operation " + "is not in a context with `ttg.num-warps`."; + } + auto kWarp = StringAttr::get(module.getContext(), "warp"); + if (ll.getInDimSize(kWarp) != *moduleWarpsPerCTA) { + return makeErr() << layout << ".\nLayout has " << ll.getInDimSize(kWarp) + << " warps per CTA, but the context requires " + << *moduleWarpsPerCTA << " warps per CTA."; + } + + // Number of CTAs per CGA. + auto kBlock = StringAttr::get(module.getContext(), "block"); + int moduleCTAsPerCGA = TritonGPUDialect::getNumCTAs(module); + if (ll.getInDimSize(kBlock) != moduleCTAsPerCGA) { + return makeErr() << layout << ".\nLayout has " << ll.getInDimSize(kBlock) + << " CTAs per CGA, but the context requires " + << moduleCTAsPerCGA << " CTAs per CGA."; + } + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// Layout debug printing +//===----------------------------------------------------------------------===// + +// Return N-D delinearized indices from a linear index. +static SmallVector delinearizeIndex(int64_t idx, + ArrayRef shape) { + SmallVector ret(shape.size()); + for (int i = shape.size() - 1; i >= 0; i--) { + ret[i] = idx % shape[i]; + idx /= shape[i]; + } + return ret; +} + +// Returns how many padding characters are needed for the string representation +// of value to be the same as max. +static int numCharacterPadding(int value, int max) { + return std::to_string(max).size() - std::to_string(value).size(); +} + +// return the string padded to have the same length as max. +static std::string paddedString(int value, int max) { + int nbChar = numCharacterPadding(value, max); + std::string str; + for (int i = 0; i < nbChar; i++) + str += " "; + str += std::to_string(value); + return str; +} + +std::string mlir::triton::gpu::getSharedLayoutStr(LinearLayout &ll, + bool useHWPointOfView) { + // This RankedTensorType is a MemDescType (?!) + auto outDimNames = llvm::to_vector(ll.getOutDimNames()); + auto shape = convertType(llvm::to_vector(ll.getOutDimSizes())); + auto *ctx = outDimNames[0].getContext(); + + StringAttr kOffset = StringAttr::get(ctx, "offset"); + StringAttr kBlock = StringAttr::get(ctx, "block"); + int64_t tensorSize = product(shape); + unsigned numBlocks = ll.getInDimSize(kBlock); + int32_t blockSize = tensorSize / numBlocks; + + // elementMapping is for the non-hw layout, offsetMapping for hw-layout + std::vector elementMapping(tensorSize); + std::vector offsetMapping; + + // Shared layouts are a mapping of (block, offset) --> (...) + + // We can just use a single int to index into elementMapping because + // the 'swizzle' operation rearranges the indices---and we want to keep it + // that way + int32_t idx = 0; + // Enumerate all the offsets for each block + for (int32_t block = 0; block < numBlocks; block++) { + for (int32_t offset = 0; offset < blockSize; offset++) { + SmallVector> inputs = { + {kBlock, block}, + {kOffset, offset}, + }; + + SmallVector> outputs = ll.apply(inputs); + + std::string sharedInfo = "("; + std::string &value = elementMapping[idx]; + + if (!value.empty()) + value += "|"; + + value += "("; + // We can build up both strings (for hw/non-hw layouts) concurrently + for (int i = 0; i < outputs.size(); i++) { + // Based on the formatting from LinearLayout::toString, the format for + // the hw layout is slightly different. HW layouts use "," vs ":". + if (i > 0) { + sharedInfo += ","; + value += ":"; + } + auto index = paddedString(outputs[i].second, shape[i]); + sharedInfo += index; + value += index; + } + value += ")"; + sharedInfo += ")"; + + offsetMapping.push_back(sharedInfo); + + idx++; + } + } + + std::string layoutStr; + + if (!useHWPointOfView) { + int rank = shape.size(); + bool newLine = true; + for (int i = 0; i < tensorSize; i++) { + auto indices = delinearizeIndex(i, shape); + int numOpenBracket = 0; + for (int j = rank - 1; j >= 0; j--) { + if (indices[j] % shape[j] != 0) + break; + layoutStr += "["; + numOpenBracket++; + } + if (newLine) { + for (int j = 0; j < rank - numOpenBracket; j++) + layoutStr += " "; + newLine = false; + } + + layoutStr += elementMapping[i]; + auto nextIndices = delinearizeIndex(i + 1, shape); + for (int j = rank - 1; j >= 0; j--) { + if (nextIndices[j] % shape[j] != 0) + break; + layoutStr += "]"; + } + if (nextIndices.back() % shape.back() == 0) { + layoutStr += "\n"; + newLine = true; + } else { + layoutStr += ","; + } + } + } else { + // For the HW view here, print the (block, offset) --> (r,c) mapping + uint32_t idx = 0; + for (int32_t block = 0; block < numBlocks; block++) { + layoutStr += "Block: " + std::to_string(block) + ":\n"; + for (int32_t offset = 0; offset < (tensorSize / numBlocks); offset++) { + layoutStr += "Offset: " + std::to_string(offset) + " -> "; + layoutStr += offsetMapping[idx]; + layoutStr += "\n"; + idx++; + } + } + } + + return layoutStr; +} + +std::string mlir::triton::gpu::getDistributedLayoutStr(LinearLayout &ll, + bool useHWPointOfView) { + auto inDimNames = llvm::to_vector(ll.getInDimNames()); + auto *ctx = inDimNames[0].getContext(); + StringAttr kRegister = StringAttr::get(ctx, "register"); + StringAttr kLane = StringAttr::get(ctx, "lane"); + StringAttr kWarp = StringAttr::get(ctx, "warp"); + StringAttr kBlock = StringAttr::get(ctx, "block"); + + int64_t tensorSize = ll.getTotalOutDimSize(); + std::vector elementMapping(tensorSize); + std::vector threadMapping; + auto shape = convertType(llvm::to_vector(ll.getOutDimSizes())); + unsigned threadsPerWarp = ll.getInDimSize(kLane); + unsigned numWarpsPerCTA = ll.getInDimSize(kWarp); + unsigned numBlocks = ll.getInDimSize(kBlock); + int numElementsPerThreads = ll.getInDimSize(kRegister); + for (int blockId = 0; blockId < numBlocks; ++blockId) { + for (int warpId = 0; warpId < numWarpsPerCTA; warpId++) { + for (int tid = 0; tid < threadsPerWarp; ++tid) { + for (int idx = 0; idx < numElementsPerThreads; ++idx) { + SmallVector> inputs = { + {kBlock, blockId}, + {kWarp, warpId}, + {kLane, tid}, + {kRegister, idx}}; + SmallVector> outputs = + ll.apply(inputs); + int32_t linearizedIdx = 0; + int stride = 1; + for (int i = outputs.size() - 1; i >= 0; i--) { + linearizedIdx += outputs[i].second * stride; + stride *= shape[i]; + } + std::string &value = elementMapping[linearizedIdx]; + if (!value.empty()) + value += "|"; + int padding = numCharacterPadding(blockId, numBlocks) + + numCharacterPadding(tid + warpId * threadsPerWarp, + numWarpsPerCTA * threadsPerWarp) + + numCharacterPadding(idx, numElementsPerThreads); + for (int i = 0; i < padding; i++) + value += " "; + if (numBlocks > 1) + value += "B" + std::to_string(blockId) + ":"; + value += "T" + std::to_string(tid + warpId * threadsPerWarp) + ":" + + std::to_string(idx); + // Now also compute the thread mapping. + std::string threadInfo = "("; + for (int i = 0; i < outputs.size(); i++) { + if (i > 0) + threadInfo += ","; + threadInfo += paddedString(outputs[i].second, shape[i]); + } + threadInfo += ")"; + threadMapping.push_back(threadInfo); + } + } + } + } + std::string layoutStr; + if (!useHWPointOfView) { + // Printing the threads containing each elements of the tensor. + int rank = ll.getNumOutDims(); + bool newLine = true; + for (int i = 0; i < tensorSize; i++) { + auto indices = delinearizeIndex(i, shape); + int numOpenBracket = 0; + for (int j = rank - 1; j >= 0; j--) { + if (indices[j] % shape[j] != 0) + break; + layoutStr += "["; + numOpenBracket++; + } + if (newLine) { + for (int j = 0; j < rank - numOpenBracket; j++) + layoutStr += " "; + newLine = false; + } + + layoutStr += elementMapping[i]; + auto nextIndices = delinearizeIndex(i + 1, shape); + for (int j = rank - 1; j >= 0; j--) { + if (nextIndices[j] % shape[j] != 0) + break; + layoutStr += "]"; + } + if (nextIndices.back() % shape.back() == 0) { + layoutStr += "\n"; + newLine = true; + } else { + layoutStr += ", "; + } + } + } else { + // Printing the elements in each physical reg/warps/threads. + for (int blockId = 0; blockId < numBlocks; blockId++) { + if (numBlocks > 1) + layoutStr += "Block" + std::to_string(blockId) + ":\n"; + for (int warpId = 0; warpId < numWarpsPerCTA; warpId++) { + layoutStr += "Warp" + std::to_string(warpId) + ":\n"; + for (int idx = 0; idx < numElementsPerThreads; ++idx) { + for (int tid = 0; tid < threadsPerWarp; ++tid) { + int linearizedIdx = + blockId * numWarpsPerCTA * threadsPerWarp * + numElementsPerThreads + + warpId * threadsPerWarp * numElementsPerThreads + + tid * numElementsPerThreads + idx; + layoutStr += threadMapping[linearizedIdx]; + if (tid < threadsPerWarp - 1) + layoutStr += ", "; + } + layoutStr += "\n"; + } + } + } + } + return layoutStr; +} + +template +llvm::SmallVector +mlir::triton::gpu::expandMatrixShapeWithBatch(llvm::ArrayRef s) { + auto rank = s.size(); + assert(rank == 2 || rank == 3); + if (rank == 3) + return llvm::SmallVector(s); + return {1, s[0], s[1]}; +} + +template llvm::SmallVector +mlir::triton::gpu::expandMatrixShapeWithBatch( + llvm::ArrayRef s); + +template llvm::SmallVector +mlir::triton::gpu::expandMatrixShapeWithBatch( + llvm::ArrayRef s); + +llvm::SmallVector +mlir::triton::gpu::expandMatrixOrderWithBatch(llvm::ArrayRef o) { + int rank = o.size(); + assert(rank == 2 || rank == 3); + if (rank == 3) + return llvm::SmallVector(o); + llvm::SmallVector expanded(3, 0); + for (int i = 0; i < rank; ++i) + expanded[i] += o[i] + 1; + return expanded; +} + +std::string mlir::triton::gpu::getLayoutStr(RankedTensorType tensorType, + bool useHWPointOfView) { + auto layout = tensorType.getEncoding(); + LinearLayout ll = triton::gpu::toLinearLayout(tensorType.getShape(), layout); + + // tensorType is needed later on (e.g., getDimSize(j)), so we still have to + // pass it as a param + // TODO: Pass TensorOrMemDesc instead of RankedTensorType in + // triton-tensor-layout.cpp + if (mlir::isa(layout)) { + return getSharedLayoutStr(ll, useHWPointOfView); + } else if (mlir::isa(layout)) { + return getDistributedLayoutStr(ll, useHWPointOfView); + } + + // else unimplemented, return error + llvm::report_fatal_error("Unimplemented usage of getLayoutStr"); + return ""; +} + +void mlir::triton::gpu::dumpLayout(RankedTensorType tensorType) { + llvm::errs() << getLayoutStr(tensorType, /*useHWPointOfView=*/false); +} + +void mlir::triton::gpu::dumpHWLayout(RankedTensorType tensorType) { + llvm::errs() << getLayoutStr(tensorType, /*useHWPointOfView=*/true); +} + +namespace { +struct TensorModel + : public triton::gpu::TensorOrMemDesc::ExternalModel { + Type getElementType(Type pointer) const { + return cast(pointer).getElementType(); + } + Attribute getEncoding(Type pointer) const { + return cast(pointer).getEncoding(); + } + ArrayRef getShape(Type pointer) const { + return cast(pointer).getShape(); + } + int64_t getRank(Type pointer) const { + return cast(pointer).getRank(); + } + int64_t getElementTypeBitWidth(Type pointer) const { + return cast(pointer).getElementTypeBitWidth(); + } +}; + +struct MemDescModel + : public triton::gpu::TensorOrMemDesc::ExternalModel { + Type getElementType(Type pointer) const { + return cast(pointer).getElementType(); + } + Attribute getEncoding(Type pointer) const { + return cast(pointer).getEncoding(); + } + ArrayRef getShape(Type pointer) const { + return cast(pointer).getShape(); + } + int64_t getRank(Type pointer) const { + return cast(pointer).getShape().size(); + } + int64_t getElementTypeBitWidth(Type pointer) const { + return cast(pointer).getElementType().getIntOrFloatBitWidth(); + } +}; +} // namespace + +void TritonGPUDialect::initialize() { + registerTypes(); + + addAttributes< +#define GET_ATTRDEF_LIST +#include "triton/Dialect/TritonGPU/IR/AttrDefs.cpp.inc" + >(); + addOperations< +#define GET_OP_LIST +#include "triton/Dialect/TritonGPU/IR/Ops.cpp.inc" +#include "triton/Dialect/TritonGPU/IR/OpsEnums.cpp.inc" + >(); + addInterfaces(); + addInterfaces(); + addInterfaces(); + addInterfaces(); + + RankedTensorType::attachInterface(*getContext()); + MemDescType::attachInterface(*getContext()); +} + +LogicalResult TritonGPUDialect::verifyOperationAttribute(Operation *op, + NamedAttribute attr) { + // Verify that dialect attributes are attached to the right ops. + if (llvm::is_contained( + {AttrNumCTAsName, AttrTargetName, AttrNumThreadsPerWarp}, + attr.getName()) && + !isa(op)) { + return op->emitOpError("has unexpected attribute ") + << attr.getName() << " which is expected only on `module` ops"; + } + if (attr.getName() == AttrNumWarpsName && !isa(op)) { + return op->emitOpError("has unexpected attribute ") + << attr.getName() + << " which is expected only on `module` or `tt.func` ops"; + } + + // Verify that all ops in a tt.warp_specialize op have partition ids + if (attr.getName() == "tt.warp_specialize") { + if (!isa(op)) { + return op->emitOpError("has unexpected attribute ") + << attr.getName() << " which is expected only on `scf.for` ops"; + } + Operation *failedOp = nullptr; + op->walk([&](Operation *childOp) { + if (!childOp->hasAttr(kPartitionAttrName)) { + failedOp = childOp; + WalkResult::interrupt(); + } + }); + if (failedOp) { + return failedOp->emitOpError("does not have expected attribute ") + << kPartitionAttrName + << " which is expected on all child ops of an op with " + "attribute `tt.warp_specialize`"; + } + } + + // Verify that partition id lists are non-empty, sorted and have no duplicates + auto verifyPartitionIds = + [&](const ArrayRef &partitionIds) -> LogicalResult { + SetVector idSet; + for (auto id : partitionIds) { + if (idSet.contains(id)) + return op->emitOpError("has duplicated partition ids in attribute ") + << attr.getName(); + idSet.insert(id); + } + if (idSet.empty()) + return op->emitOpError("has no partition ids in attribute ") + << attr.getName(); + auto ids = idSet.takeVector(); + SmallVector sortedIds(ids.begin(), ids.end()); + std::sort(sortedIds.begin(), sortedIds.end()); + if (ids != sortedIds) + return op->emitOpError("partition ids not in sorted order in attribute ") + << attr.getName(); + return success(); + }; + + if (attr.getName() == kPartitionAttrName) { + auto result = verifyPartitionIds( + cast(attr.getValue()).asArrayRef()); + if (failed(result)) + return result; + } + if (attr.getName() == kPartitionOutputsAttrName) { + auto arrayAttr = cast(attr.getValue()); + for (auto idx = 0; idx < arrayAttr.size(); idx++) { + auto result = verifyPartitionIds( + cast(arrayAttr[idx]).asArrayRef()); + if (failed(result)) + return result; + } + } + + // Verify that op partitions include partitions of all child ops + if (attr.getName() == kPartitionAttrName && op->getNumRegions() != 0) { + SetVector expectedIds; + for (auto ®ion : op->getRegions()) { + for (auto &block : region.getBlocks()) { + for (auto &childOp : block.getOperations()) { + if (isa(childOp)) { + // yield ops and ub.poison do not need partition ids + continue; + } + if (!childOp.hasAttr(kPartitionAttrName)) + return childOp.emitOpError("does not have expected attribute ") + << kPartitionAttrName + << " which is expected for ops whose parent has partitions"; + auto ids = getPartitionIds(&childOp); + expectedIds.insert(ids.begin(), ids.end()); + } + } + } + auto partitionIds = getPartitionIds(op); + for (auto id : expectedIds) { + if (!partitionIds.contains(id)) { + return op->emitOpError("partition ids in attr ") + << attr.getName() + << " does not contain partition ids of all child ops"; + } + } + } + + if (attr.getName() == kPartitionOutputsAttrName) { + if (!isa(op)) + return op->emitOpError("has unexpected attribute ") << attr.getName(); + + // Verify that number of output partitions matches number of For/If results + size_t numResults = 0; + if (isa(op)) { + numResults = cast(op).getResults().size(); + } else if (isa(op)) { + numResults = cast(op).getResults().size(); + } else { + numResults = cast(op).getResults().size(); + } + + if (cast(attr.getValue()).size() != numResults) { + return op->emitOpError("does not have expected number of output " + "partition sets in attr ") + << attr.getName() << "; should match number of results"; + } + + // Verify that union of op output partitions is a subset of op partitions + if (!op->hasAttr(kPartitionAttrName)) + return op->emitOpError("does not have expected attribute ") + << kPartitionAttrName << " which is expected for ops with attr " + << kPartitionOutputsAttrName; + auto partitionIds = getPartitionIds(op); + + SetVector outputPartitionIdsUnion; + for (auto outputPartitionIds : getPartitionOutputs(op)) { + outputPartitionIdsUnion.insert(outputPartitionIds.begin(), + outputPartitionIds.end()); + } + if (!std::all_of(outputPartitionIdsUnion.begin(), + outputPartitionIdsUnion.end(), + [&](int id) { return partitionIds.contains(id); })) { + return op->emitOpError("partition ids in attr ") + << kPartitionAttrName + << " must be the union of all partition ids in " << attr.getName(); + } + } + + return success(); +} + +int TritonGPUDialect::getNumCTAs(ModuleOp module) { + if (auto attr = module->getAttrOfType(AttrNumCTAsName)) + return attr.getInt(); + return 1; +} + +int TritonGPUDialect::getThreadsPerWarp(ModuleOp module) { + if (auto attr = module->getAttrOfType(AttrNumThreadsPerWarp)) + return attr.getInt(); + return 32; +} + +std::optional triton::gpu::maybeLookupNumWarps(Operation *op) { + if (isa(op)) { + if (auto attr = op->getAttrOfType(AttrNumWarpsName)) + return attr.getInt(); + } else if (auto partitions = + dyn_cast(op->getParentOp())) { + unsigned idx = op->getParentRegion()->getRegionNumber(); + return partitions.getParentOp().getPartitionNumWarps()[idx]; + } + if (Operation *parent = op->getParentOp()) + return maybeLookupNumWarps(parent); + return {}; +} + +int triton::gpu::lookupNumWarps(Operation *op) { + std::optional numWarps = maybeLookupNumWarps(op); + if (!numWarps) { + op->emitOpError( + "is not contained within a context that specifies the number of warps"); + llvm::report_fatal_error("failed to lookup the number of warps, the " + "surrounding module should contain a " + + Twine(AttrNumWarpsName) + " attribute"); + } + return *numWarps; +} + +int triton::gpu::lookupNumWarps(Region *region) { + if (auto partitions = + dyn_cast(region->getParentOp())) { + unsigned idx = region->getRegionNumber(); + return partitions.getParentOp().getPartitionNumWarps()[idx]; + } + return lookupNumWarps(region->getParentOp()); +} + +int triton::gpu::lookupThreadsPerWarp(OpBuilder &rewriter) { + assert(rewriter.getInsertionBlock() && "expected an insertion point"); + Operation *op = + rewriter.getInsertionBlock()->getParentOp()->getParentOfType(); + assert(op && "cannot check threads per warp outside of module"); + return triton::gpu::TritonGPUDialect::getThreadsPerWarp(cast(op)); +} + +int triton::gpu::lookupNumCTAs(Operation *op) { + auto mod = op->getParentOfType(); + if (!mod) { + op->emitOpError( + "is not contained within a module, cannot lookup number of CTAs"); + llvm::report_fatal_error( + "failed to lookup the number of CTAs, the surrounding module should " + "contain a ModuleOp"); + } + return triton::gpu::TritonGPUDialect::getNumCTAs(mod); +} + +int triton::gpu::lookupNumCTAs(OpBuilder &rewriter) { + assert(rewriter.getInsertionBlock() && "expected an insertion point"); + Operation *op = + rewriter.getInsertionBlock()->getParentOp()->getParentOfType(); + assert(op && "cannot check number of CTAs outside of module"); + return triton::gpu::TritonGPUDialect::getNumCTAs(cast(op)); +} + +bool triton::gpu::areLayoutsEquivalent(ArrayRef shape, + LayoutEncodingTrait lhs, + LayoutEncodingTrait rhs) { + auto lhsLL = triton::gpu::toLinearLayout(shape, lhs); + auto rhsLL = triton::gpu::toLinearLayout(shape, rhs); + return lhsLL == rhsLL; +} + +bool triton::gpu::isInnermostContiguous(MemDescType type, unsigned numElems) { + ArrayRef shape = type.getShape(); + Attribute enc = type.getEncoding(); + MLIRContext *ctx = enc.getContext(); + + LinearLayout actual = toLinearLayout(type); + StringAttr fastestIn = *actual.getInDimNames().begin(); + + // Flatten actual outs in reverse order to produce a row-major flattening + // of the layout + auto outNames = actual.getOutDimNames(); + SmallVector revOut(outNames.begin(), outNames.end()); + std::reverse(revOut.begin(), revOut.end()); + actual = actual.transposeOuts(revOut).flattenOuts(); + + return actual.getNumConsecutiveInOut() >= numElems; +} + +LinearLayout triton::gpu::inferReshapeLinearLayout(TensorOrMemDesc srcTy, + ArrayRef dstShape) { + auto *ctx = srcTy.getContext(); + auto src = toLinearLayout(srcTy); + assert(product(srcTy.getShape()) == product(dstShape)); + auto dst = reshapeLayout(ctx, src, dstShape); + return dst; +} + +SetVector triton::gpu::getPartitionIds(Operation *op) { + auto attrs = op->getAttr(kPartitionAttrName); + SmallVector partitionIds; + for (auto id : cast(attrs).asArrayRef()) { + partitionIds.push_back(id); + } + std::sort(partitionIds.begin(), partitionIds.end()); + return SetVector(partitionIds.begin(), partitionIds.end()); +} + +SmallVector, 4> triton::gpu::getPartitionOutputs(Operation *op) { + SmallVector, 4> partitionOutputsIds; + if (op->getNumResults() == 0) { + return partitionOutputsIds; + } + auto arrayAttr = cast(op->getAttr(kPartitionOutputsAttrName)); + for (auto attr : arrayAttr) { + auto ids = cast(attr).asArrayRef(); + partitionOutputsIds.push_back(SetVector(ids.begin(), ids.end())); + } + return partitionOutputsIds; +} + +SetVector triton::gpu::getPartitionIds(OpOperand *use) { + auto owner = use->getOwner(); + if (isa(owner)) { + return getPartitionOutputs(owner->getParentOp())[use->getOperandNumber()]; + } else if (scf::ForOp forOp = dyn_cast(owner)) { + int idx = use->getOperandNumber() - forOp.getNumControlOperands(); + return idx >= 0 ? getPartitionOutputs(owner)[idx] : getPartitionIds(forOp); + } else { + return getPartitionIds(owner); + } +} + +bool triton::gpu::hasPartition(Operation *op) { + return op && op->hasAttr(kPartitionAttrName); +} + +bool triton::gpu::hasWarpSpecializeTag(Operation *op) { + return op && op->hasAttr(kWarpSpecializeTagAttrName); +} + +std::optional triton::gpu::getWarpSpecializeTag(Operation *op) { + if (hasWarpSpecializeTag(op)) { + return cast(op->getAttr(kWarpSpecializeTagAttrName)).getInt(); + } + return std::nullopt; +} + +#endif diff --git a/third_party/xpu/include/triton/Analysis/NewAnalysis/Utility.h b/third_party/xpu/include/triton/Analysis/NewAnalysis/Utility.h index c103707975..3b70bab9ce 100644 --- a/third_party/xpu/include/triton/Analysis/NewAnalysis/Utility.h +++ b/third_party/xpu/include/triton/Analysis/NewAnalysis/Utility.h @@ -67,7 +67,8 @@ enum class OffsetState { DiscreteSame = 0, Continuous = 1, Discrete = 2, - LocallyContinuous = 3 + LocallyContinuous = 3, + LocallyScalar = 4 }; enum class ElemState { diff --git a/third_party/xpu/include/triton/Analysis/ScalarAnalysis.h b/third_party/xpu/include/triton/Analysis/ScalarAnalysis.h new file mode 100644 index 0000000000..5940e9f838 --- /dev/null +++ b/third_party/xpu/include/triton/Analysis/ScalarAnalysis.h @@ -0,0 +1,240 @@ +#ifndef TRITON_XPU_ANALYSIS_SCALAR_ANALYSIS_H +#define TRITON_XPU_ANALYSIS_SCALAR_ANALYSIS_H + +#include "mlir/Analysis/DataFlow/SparseAnalysis.h" +#include "mlir/Pass/Pass.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include + +namespace mlir { +namespace triton { +namespace xpu { + +/// Lattice element for ScalarAnalysis. Classifies an SSA Value's per-lane +/// pattern starting from `tt.make_range` / `tt.splat` along the use-def chain. +/// +/// Lattice ordering: +/// Bottom (Unknown) pessimistic / uninitialized +/// |- Scalar uniform across lanes +/// |- VectorContig(s) per-lane affine: base + s * lane +/// Top (VectorOther) known not-uniform / not-affine; can't be sharpened +/// +/// In addition to the kind / stride, we track a **base alignment** `baseAlign`: +/// - `baseAlign = 0` : the value is provably exactly zero (i.e. divisible +/// by anything; identity for gcd merges). +/// - `baseAlign = N>=1`: the per-lane base is a known multiple of `N`. +/// +/// Tracking `baseAlign` lets the transfer functions for `arith.divsi` and +/// `arith.remsi` sharpen results when the divisor is compatible with the +/// known alignment (e.g. `(pid*64 + arange(64)) / 64` collapses to a Scalar). +/// +/// `join` is used by the dataflow solver to merge states arriving from +/// different control-flow paths; arithmetic transfer functions live in the +/// op visitor. +struct ScalarValueState { + enum class Kind : uint8_t { + Unknown = 0, // bottom + Scalar = 1, // lane-uniform + VectorContig = 2, // base + stride * lane (stride known) + BlockScalar = 3, // uniform within each rowLen-sized block; varies between + // blocks (e.g. xindex / 64 when xindex = arange contig) + BlockContig = 4, // stride-1 contig within each rowLen-sized block; base + // varies between blocks (e.g. xindex % 64) + VectorOther = 5, // top: lane-varying and non-affine + }; + + ScalarValueState() = default; + ScalarValueState(Kind k, int64_t s = 0, uint64_t ba = 1, int64_t rl = 0, + int64_t bs = 0, bool bsKnown = false, bool fromLoad = false) + : kind(k), stride(s), baseAlign(ba), rowLen(rl), blockStride(bs), + blockStrideKnown(bsKnown), blockFromLoad(fromLoad) {} + + static ScalarValueState scalar(uint64_t ba = 1) { + return {Kind::Scalar, 0, ba, 0, 0, false, false}; + } + static ScalarValueState contig(int64_t stride, uint64_t ba = 1) { + return {Kind::VectorContig, stride, ba, 0, 0, false, false}; + } + static ScalarValueState blockScalar(int64_t rowLen, uint64_t ba = 1, + int64_t bs = 0, bool bsKnown = false, + bool fromLoad = false) { + return {Kind::BlockScalar, 0, ba, rowLen, bs, bsKnown, fromLoad}; + } + static ScalarValueState blockContig(int64_t rowLen, uint64_t ba = 1, + int64_t bs = 0, bool bsKnown = false, + bool fromLoad = false) { + return {Kind::BlockContig, 1, ba, rowLen, bs, bsKnown, fromLoad}; + } + static ScalarValueState other() { + return {Kind::VectorOther, 0, 1, 0, 0, false, false}; + } + static ScalarValueState unknown() { + return {Kind::Unknown, 0, 1, 0, 0, false, false}; + } + + bool isUnknown() const { return kind == Kind::Unknown; } + bool isScalar() const { return kind == Kind::Scalar; } + bool isContig() const { return kind == Kind::VectorContig; } + bool isBlockScalar() const { return kind == Kind::BlockScalar; } + bool isBlockContig() const { return kind == Kind::BlockContig; } + bool isOther() const { return kind == Kind::VectorOther; } + + bool operator==(const ScalarValueState &rhs) const { + if (kind != rhs.kind) + return false; + if (kind == Kind::VectorContig && stride != rhs.stride) + return false; + if ((kind == Kind::BlockScalar || kind == Kind::BlockContig) && + rowLen != rhs.rowLen) + return false; + if ((kind == Kind::BlockScalar || kind == Kind::BlockContig) && + (blockStrideKnown != rhs.blockStrideKnown || + (blockStrideKnown && blockStride != rhs.blockStride))) + return false; + if ((kind == Kind::BlockScalar || kind == Kind::BlockContig) && + blockFromLoad != rhs.blockFromLoad) + return false; + if ((kind == Kind::Scalar || kind == Kind::VectorContig || + kind == Kind::BlockScalar || kind == Kind::BlockContig) && + baseAlign != rhs.baseAlign) + return false; + return true; + } + + /// Merge two alignment values. `0` is treated as "infinite alignment" + /// (i.e. the value is exactly zero) and acts as the identity for gcd. + static uint64_t mergeAlign(uint64_t a, uint64_t b) { + if (a == 0) + return b; + if (b == 0) + return a; + return std::gcd(a, b); + } + + /// Pessimistic value used by `setToEntryState` for entry/external Values + /// (block arguments, function arguments). + static ScalarValueState getPessimisticValueState(Value v) { + return ScalarValueState(); + } + + /// Lattice meet: merge two states reaching the same SSA value through + /// distinct control-flow paths. + static ScalarValueState join(const ScalarValueState &a, + const ScalarValueState &b) { + if (a.kind == Kind::Unknown) + return b; + if (b.kind == Kind::Unknown) + return a; + if (a.kind == Kind::VectorOther || b.kind == Kind::VectorOther) + return other(); + if (a.kind == Kind::Scalar && b.kind == Kind::Scalar) + return scalar(mergeAlign(a.baseAlign, b.baseAlign)); + if (a.kind == Kind::VectorContig && b.kind == Kind::VectorContig) + return a.stride == b.stride + ? contig(a.stride, mergeAlign(a.baseAlign, b.baseAlign)) + : other(); + if (a.kind == Kind::BlockScalar && b.kind == Kind::BlockScalar) { + if (a.rowLen != b.rowLen) + return other(); + bool bsk = a.blockStrideKnown && b.blockStrideKnown && + a.blockStride == b.blockStride; + return blockScalar(a.rowLen, mergeAlign(a.baseAlign, b.baseAlign), + bsk ? a.blockStride : 0, bsk, + a.blockFromLoad && b.blockFromLoad); + } + if (a.kind == Kind::BlockContig && b.kind == Kind::BlockContig) { + if (a.rowLen != b.rowLen) + return other(); + bool bsk = a.blockStrideKnown && b.blockStrideKnown && + a.blockStride == b.blockStride; + return blockContig(a.rowLen, mergeAlign(a.baseAlign, b.baseAlign), + bsk ? a.blockStride : 0, bsk, + a.blockFromLoad && b.blockFromLoad); + } + // Mixed kinds along the same SSA value cannot be reconciled. + return other(); + } + + void print(raw_ostream &os) const { + switch (kind) { + case Kind::Unknown: + os << "Unknown"; + break; + case Kind::Scalar: + os << "Scalar(baseAlign=" << baseAlign << ")"; + break; + case Kind::VectorContig: + os << "VectorContig(stride=" << stride << ", baseAlign=" << baseAlign + << ")"; + break; + case Kind::BlockScalar: + os << "BlockScalar(rowLen=" << rowLen << ", baseAlign=" << baseAlign + << ")"; + break; + case Kind::BlockContig: + os << "BlockContig(rowLen=" << rowLen << ", baseAlign=" << baseAlign + << ")"; + break; + case Kind::VectorOther: + os << "VectorOther"; + break; + } + } + + Kind kind = Kind::Unknown; + int64_t stride = 0; + uint64_t baseAlign = 1; + int64_t rowLen = 0; + // For BlockScalar / BlockContig kinds, the inter-row stride (i.e. how the + // block-index `idx / rowLen` is scaled before being added to the per-block + // pattern). Only meaningful when `blockStrideKnown == true`. The canonical + // BlockContig pattern produced by `arange % R` has blockStride == 1 (after + // the implicit `(idx / R) * 1` from the original `arange`); patterns like + // `(idx / R) * S + (idx % R)` carry blockStride = S. + int64_t blockStride = 0; + bool blockStrideKnown = false; + // True when this Block* pattern's inter-row base is produced by a runtime + // `triton_xpu.load` (gather), i.e. the genuine "k89" embedding-gather case + // that LocallyContinuous with rowStride=-1 was designed for. False for + // Block* patterns built purely from arithmetic on `arange` (e.g. the cat + // strided-copy `(idx/R)*S + (idx%R)`), where the inter-row stride S is a + // compile-time constant != R and must NOT be treated as row-by-row + // contiguous DMA — those are left to OffsetAnalysis. + bool blockFromLoad = false; +}; + +/// Forward sparse dataflow analysis that propagates `ScalarValueState` +/// through arithmetic / triton ops along the use-def chain. +class ScalarAnalysis : public dataflow::SparseForwardDataFlowAnalysis< + dataflow::Lattice> { +public: + using SparseForwardDataFlowAnalysis< + dataflow::Lattice>::SparseForwardDataFlowAnalysis; + + // LLVM 22 changed the transfer function to return LogicalResult (failure is + // reserved for hard errors; "cannot refine" is still success). + LogicalResult visitOperation( + Operation *op, + ArrayRef *> operands, + ArrayRef *> results) override; + + /// Entry / external Values default to Scalar for non-tensor (a single SSA + /// value is trivially lane-uniform) and VectorOther for tensor types + /// (without a defining op we cannot say anything). + void setToEntryState( + dataflow::Lattice *lattice) override { + Value v = lattice->getAnchor(); + ScalarValueState init = isa(v.getType()) + ? ScalarValueState::other() + : ScalarValueState::scalar(1); + propagateIfChanged(lattice, lattice->join(init)); + } +}; + +} // namespace xpu +} // namespace triton +} // namespace mlir + +#endif // TRITON_XPU_ANALYSIS_SCALAR_ANALYSIS_H diff --git a/third_party/xpu/include/triton/Analysis/TileAnalysis.h b/third_party/xpu/include/triton/Analysis/TileAnalysis.h new file mode 100644 index 0000000000..6f28b9719d --- /dev/null +++ b/third_party/xpu/include/triton/Analysis/TileAnalysis.h @@ -0,0 +1,160 @@ +#ifndef TRITONXPU_ANALYSIS_TILEANALYSIS_H +#define TRITONXPU_ANALYSIS_TILEANALYSIS_H + +#include "triton/Analysis/VectorizabilityAnalysis.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "llvm/ADT/SetVector.h" + +//===----------------------------------------------------------------------===// +// Tile analysis: the measurements a tiling decision needs, kept separate from +// the passes that act on them. +// +// A tile of an XPU kernel is described by three factors whose product is the +// per-core element count fixed by CoreTiling: +// +// E = vectorWidth * vregsPerIter * iterNum +// +// `vectorWidth` is chosen by Vectorize, `iterNum` by UnrollControl, and +// `vregsPerIter` is the residual that decides whether the tile spills. The +// passes that pick those factors are transforms; what they need in order to +// pick is measurement, and that is what lives here so a single planner can +// eventually consume all of it. +//===----------------------------------------------------------------------===// + +namespace mlir { +namespace triton { +namespace xpu { + +// Cluster layout of a tensor type, looking through the slice encoding that a +// reduction leaves on its result. +ClusterLayoutAttr getClusterLayout(RankedTensorType tensorTy); + +// Per-core register footprint of one SSA value, counted in registers of its own +// file. `isVector`, when given, reports which file that is. +// +// Vector values need no conversion: Vectorize divides sizePerCore.back() by the +// vector width, so product(sizePerCore) already is the number of vector +// registers held. +// +// Scalar-element values are counted too, one scalar register per element. +// Vectorize bails out whenever numElems < vectorWidth or +// numElems % vectorWidth != 0 (Vectorize.cpp:335), so a partially vectorized +// tree keeps scalar tensors live across the whole segment and used to be +// invisible here. The two files are kept apart on purpose: a vector register is +// 512 bits and a scalar one 32, they spill independently, and folding both into +// one budget makes the target meaningless (measured: it collapses layernorm to +// iterNum=1 and 23 vector spills). +int64_t getNumRegs(Type type, bool *isVector = nullptr); + +struct RegPressure { + int64_t vecPeak = 0; // 512-bit registers simultaneously live + int64_t scalarPeak = 0; // 32-bit registers simultaneously live + // Whole footprint the segment defines, which is what the un-tiled full width + // has to keep alive. + int64_t vecTotal = 0; + int64_t scalarTotal = 0; + // Widest per-core last dim among *vector* values only. Widths of the two + // files are in different units (registers vs elements) and must never be + // maxed together, or a width conversion against this divides the target down + // to 1. + int64_t maxVecWidth = 1; + // Narrowest per-core last dim among *vector* values, in vector slots, or 0 + // when the segment holds none. This is a legality bound, not a heuristic: the + // retyping slices sizePerCore with ceil(), so a trip count larger than a + // vector value's slot count saturates that value at one slot per iteration + // while the scalar values in the same tree keep dividing. The two sides then + // disagree on how many lanes one iteration covers, which is how a vselect + // ends up reading its condition past the end + // (VectorizedOpToLLVM.cpp:539-600 assumes condElems == numElems * vecSize). + int64_t minVecWidth = 0; +}; + +// Peak simultaneously-live registers across `opTree` as it stands, per file. +// `m` supplies the program order the liveness walk needs. +void getRegPressure(ModuleOp m, const llvm::SetVector &opTree, + RegPressure &p); + +// Same measurement over the whole block containing `insertPt`. Trees of +// different widths sharing a block must be tiled to the same per-iteration +// width rather than the same trip count, so the block-wide maxVecWidth matters +// as much as its peak. +void getBlockRegPressure(ModuleOp m, Operation *insertPt, RegPressure &p); + +//===----------------------------------------------------------------------===// +// E-dependent gates. +// +// These two read sizePerCore, so they cannot be answered before CoreTiling has +// fixed E. They used to sit in VectorizabilityAnalysis next to the E-independent +// op-kind and element-type whitelists; they live here instead so the state half +// of the vectorizability question carries no dependency on the numeric half and +// can be moved ahead of CoreTiling on its own (redesign-v2.md §4.2 steps 1.4 and +// 1.5). +// +// Both are phrased as "does a whole number of 512-bit vectors fit", which is the +// only thing either of them tests. Neither had any state, which is what makes +// them free functions here. +//===----------------------------------------------------------------------===// + +// Whether a root tensor type holds enough elements per core, in a whole number +// of vectors, to be worth vectorizing at all. `rowsPerCore` is factored out +// first, so a core holding several rows is judged on one row's width. +bool vectorFitsRoot(Type rootOpTy); + +// Same question for one operand of a reduce, measured along the reduced axis. +// Note this reads the *tensor shape* on that axis, not sizePerCore, so it is +// E-dependent only through the layout CoreTiling leaves behind. +bool vectorFitsReduceOperand(triton::xpu::ReduceOp redOp, Type operandTy); + +// The real fit oracle handed to `VectorizabilityAnalysis` (step 1.5b): answers +// the three footprint questions the closure walk used to answer inline by +// calling `getTotalElemsPerThread` itself. E-dependent, hence here; never +// returns `Fit::Unknown`, which is what makes the walk byte-identical to before +// on this side of CoreTiling. +Fit vectorFitsValue(Value value, FitQuery query, unsigned wantWidth); + +//===----------------------------------------------------------------------===// +// M6 -- the tile plan carrier (redesign-v2.md §3.7, step 1.6). +// +// The analysis that decides vectorizability runs before `tritonxpu-vectorize`; +// the transforms that consume the decision run around `tritonxpu-unroll-control`, +// with vectorize / canonicalize / alloca / memory-async in between. Something has +// to survive that gap and still identify *which root* each conclusion belongs to. +// +// §3.7 named two candidate keys and marked both unverified. This is the probe +// that decides between them by measuring, not by argument: the producer writes +// both keys for every root, the consumer looks both up again and reports two hit +// rates. +// +// loc -- the root's location, printed. Free, but a pass that rebuilds an op +// has to carry the loc across for this to work. +// plan_id -- an integer the producer stamps on the root op itself. Stable +// against renaming, but only survives if the op survives *and* its +// discardable attributes are copied. +// +// The payload is deliberately uninteresting (whatever the producing site already +// computed). Step 1.6 measures identification, not content. +// +// Everything here is off unless TRITONXPU_TILE_PLAN=1, and `tilePlanCheck` +// erases both keys unconditionally -- an entry reaching the emitted IR would +// break C1, which is the other half of this step's exit gate. +//===----------------------------------------------------------------------===// + +constexpr llvm::StringLiteral kTilePlanAttrName = "triton_xpu.tile_plan"; +constexpr llvm::StringLiteral kTilePlanIdAttrName = "triton_xpu.plan_id"; + +bool tilePlanProbeEnabled(); + +// Append one entry for `root`, stamping it with the next id. No-op when the +// probe is off. +void tilePlanRecord(ModuleOp mod, Operation *root, StringRef site, bool eligible, + int64_t closure); + +// Resolve every entry against the IR as it now stands, report per-entry how many +// live ops each key found, and erase everything the probe wrote. +void tilePlanCheck(ModuleOp mod); + +} // namespace xpu +} // namespace triton +} // namespace mlir + +#endif // TRITONXPU_ANALYSIS_TILEANALYSIS_H diff --git a/third_party/xpu/include/triton/Analysis/TileDecision.h b/third_party/xpu/include/triton/Analysis/TileDecision.h new file mode 100644 index 0000000000..8cea2042ea --- /dev/null +++ b/third_party/xpu/include/triton/Analysis/TileDecision.h @@ -0,0 +1,125 @@ +#ifndef TRITONXPU_ANALYSIS_TILEDECISION_H +#define TRITONXPU_ANALYSIS_TILEDECISION_H + +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" +#include +#include +#include + +//===----------------------------------------------------------------------===// +// Tile decision: picks `iterNum` out of the trip counts UnrollControl can +// express, by narrowing the candidate set one tier at a time. +// +// The order is a dictionary order, not a weighted sum: a lower tier is never +// traded against a higher one, so no unit conversion between "registers" and +// "instructions" has to be invented. Tier 1 and 2 are hard -- they filter -- +// and tier 3 is soft: it ranks whatever survived. Adding a consideration means +// pushing one more object into a tier, or adding a tier; nothing registers +// itself and there is no weight to tune. +// +// Every criterion has to end up in the remark (see `Decision::perTierTrace`). +// A criterion that cannot be observed cannot be verified, and the one time this +// pass shipped a silent fallback -- a mis-set budget making `budget-unreachable` +// the normal case -- the symptom was read as the design for weeks while every +// probe ran 8..11% slow. +//===----------------------------------------------------------------------===// + +namespace mlir { +namespace triton { +namespace xpu { + +// One expressible trip count. The boundary set and the segment id land here +// once M4/M5 exist; until then the candidate *is* the trip count. +struct TileCandidate { + int64_t iterNum = 1; +}; + +// What the criteria are allowed to read. Filled once per decision point so the +// criteria stay free of pass state and are testable on their own. +struct TileContext { + int64_t numCol = 1; + int64_t widthPerCore = 1; + // Peak simultaneously-live vector registers over the block the tile loop + // will live in, and the widest vector row in it. Both come from + // getRegPressure/getBlockRegPressure. + int64_t peakVRegs = 0; + int64_t maxVecWidth = 1; + int64_t scalarPeak = -1; // reported only, never drives a tier + // Narrowest vector row in slots, or 0 when the segment holds no vector + // value. A legality bound, so it is already folded into the candidate set; + // it is carried here only to name which ceiling a fallback hit. + int64_t vecRow = 0; + int64_t vrfBudget = 0; + // iter_args the tile loop will carry: 0 at a pointwise store segment + // (UnrollControl.cpp:1263 passes an empty range), one per reduce data + // operand at a reduce segment (:1908). + int64_t loopResults = 0; +}; + +// What one criterion said about the candidate set, for the remark. +struct CriterionTrace { + llvm::StringRef name; + unsigned tier = 0; + int64_t candidatesIn = 0; + int64_t candidatesOut = 0; + std::optional chosenCost; // set by soft criteria only + std::string why; // non-empty when the tier vetoed +}; + +struct Decision { + int64_t iterNum = 1; + std::string why; + llvm::SmallVector perTierTrace; +}; + +class TileCriterion { +public: + virtual ~TileCriterion() = default; + virtual llvm::StringRef name() const = 0; + // 1 = hard (VRF pressure proxy), 2 = hard (LM capacity), + // 3 = soft (tie-break). + virtual unsigned tier() const = 0; + + // Hard criteria implement this; fill `why` when rejecting. + virtual bool isFeasible(const TileCandidate &, const TileContext &, + std::string &why) const { + return true; + } + // Soft criteria implement this. Smaller is better; std::nullopt means the + // criterion has no opinion on this candidate. + virtual std::optional cost(const TileCandidate &, + const TileContext &) const { + return std::nullopt; + } +}; + +// The trip count a tree of the widest row needs so that the block's vector +// pressure lands inside the calibrated budget. +// +// NOT a zero-spill criterion, and it must not be described as one: the budget +// is a pressure ceiling calibrated against measured time, and on the welford +// reduce segment the time-optimal point spills 14 accumulators while the +// spill-free point is 15.6% slower. Once the pack/unpack unit price is +// calibrated this whole criterion moves to tier 3, so nothing about it may +// assume it sits on a hard tier. +int64_t vrfBudgetTarget(const TileContext &ctx); + +// Assembles the first-version criteria and runs the narrowing. `candidates` +// must already be the expressible trip counts, ascending. +class TileDecider { +public: + TileDecider(); + Decision decide(llvm::ArrayRef candidates, + const TileContext &ctx) const; + +private: + llvm::SmallVector> criteria; +}; + +} // namespace xpu +} // namespace triton +} // namespace mlir + +#endif // TRITONXPU_ANALYSIS_TILEDECISION_H diff --git a/third_party/xpu/include/triton/Analysis/VectorizabilityAnalysis.h b/third_party/xpu/include/triton/Analysis/VectorizabilityAnalysis.h new file mode 100644 index 0000000000..af7e7449d9 --- /dev/null +++ b/third_party/xpu/include/triton/Analysis/VectorizabilityAnalysis.h @@ -0,0 +1,382 @@ +#ifndef TRITONXPU_ANALYSIS_VECTORIZABILITYANALYSIS_H +#define TRITONXPU_ANALYSIS_VECTORIZABILITYANALYSIS_H + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "triton/Analysis/NewAnalysis/Utility.h" +#include "triton/Tools/Sys/GetEnv.hpp" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" + +#include + +//===----------------------------------------------------------------------===// +// Vectorizability analysis: which ops can take a `vector` element type, +// and how wide that vector is. +// +// This answers the question the Vectorize pass used to answer inline while it +// was already rewriting types. Separating the two matters because the tile +// factorization +// +// E = vectorWidth * vregsPerIter * iterNum +// +// couples Vectorize to UnrollControl: whether a tree vectorizes decides whether +// its values sit in 512-bit registers or in scalar ones, which decides the +// register pressure UnrollControl has to tile against. +// +// The predicates split into two groups, and step 1.4 of the redesign has made +// that split structural: the E-dependent half now lives in TileAnalysis.h +// (`vectorFitsRoot`, `vectorFitsReduceOperand`) and this unit only reaches it +// through the callback handed to the constructor. Nothing here includes +// TileAnalysis.h, which is what lets step 1.5 move this unit ahead of CoreTiling. +// +// * E-independent, so answerable before CoreTiling fixes sizePerCore: the +// op-kind whitelist in `getVectorizableClosure`, the element type whitelist +// in `vectorizedTyValid`, closure connectivity, the reduce combine-region +// check, and the ExternElementwise symbol whitelist. +// * E-dependent, so only answerable once sizePerCore is known: the root and +// reduce-operand fit tests, both of which divide an element count by the +// vector width. Now in TileAnalysis. +// +// Step 1.5b finished the split for the three cases whose residue sat inside the +// walk rather than at its edges: LoadOp's numElems % vectorWidth == 0, SplatOp's +// getTotalElemsPerThread(srcTy) == 1, and BroadcastOp's result element count all +// go through the `FitOracle` below now, so the walk itself no longer calls +// getTotalElemsPerThread. What each case keeps is its E-independent half (op +// kind, rank, shape relations, element width). Moving the walk ahead of +// CoreTiling is then a matter of handing it an all-`Unknown` oracle and +// re-testing the candidates in place, which is step 1.5c. +//===----------------------------------------------------------------------===// + +namespace mlir { +namespace triton { +namespace xpu { + +using OperationTree = llvm::SetVector; + +#define ARITH_BINARY_FLOAT_OP \ + arith::AddFOp, arith::SubFOp, arith::MulFOp, arith::DivFOp, \ + arith::MaximumFOp, arith::MinimumFOp, arith::MaxNumFOp, arith::MinNumFOp + +#define ARITH_BINARY_INT_OP \ + arith::SubIOp, arith::AndIOp, arith::OrIOp, arith::MulIOp, arith::AddIOp, \ + arith::XOrIOp + +#define MATH_UNARY_OP \ + math::ExpOp, math::SqrtOp, math::SinOp, math::CosOp, arith::ExtFOp, \ + arith::TruncFOp, math::AbsFOp, math::LogOp + +#define REDUCE_COMBINE_OP COMBINE_OP, triton::xpu::ReduceReturnOp + +// The scalar -> vector op table, single-sourced (redesign-v2.md §3.1, step 2.1). +// +// Vectorize.cpp expands this into the `VOp` specializations its rewrite +// builds from; `hasVectorForm` below expands the same list into a runtime +// predicate. Keeping the second consumer as a hand-maintained whitelist is what +// §3.1 forbids, and the cost of drift is not a missed optimization: the rewrite +// dispatch ends in `llvm_unreachable`, so a table the analysis believes in but +// the rewrite lacks is a crash. +// +// Membership here means only "this op kind has a 512-bit vector form". Element +// type (`vectorizedTyValid`), footprint (the FitOracle) and operand shape are +// separate questions asked elsewhere. +#define TTX_SCALAR_TO_VECTOR_OPS(FN) \ + FN(arith::AddFOp, triton::xpu::VvaddFOp) \ + FN(arith::SubFOp, triton::xpu::VvsubFOp) \ + FN(arith::MulFOp, triton::xpu::VvmulFOp) \ + FN(arith::DivFOp, triton::xpu::VvdivFOp) \ + FN(arith::MaximumFOp, triton::xpu::VvmaxFOp) \ + FN(arith::MinimumFOp, triton::xpu::VvminFOp) \ + FN(arith::MaxNumFOp, triton::xpu::VvmaxNumFOp) \ + FN(arith::MinNumFOp, triton::xpu::VvminNumFOp) \ + FN(arith::AddIOp, triton::xpu::VvaddIOp) \ + FN(arith::SubIOp, triton::xpu::VvsubIOp) \ + FN(arith::MulIOp, triton::xpu::VvmulIOp) \ + FN(arith::AndIOp, triton::xpu::VvandIOp) \ + FN(arith::XOrIOp, triton::xpu::VvxorIOp) \ + FN(arith::OrIOp, triton::xpu::VvorIOp) \ + FN(math::ExpOp, triton::xpu::VExpFOp) \ + FN(math::AbsFOp, triton::xpu::VAbsFOp) \ + FN(math::LogOp, triton::xpu::VLogFOp) \ + FN(math::SqrtOp, triton::xpu::VSqrtFOp) \ + FN(math::SinOp, triton::xpu::VSinFOp) \ + FN(math::CosOp, triton::xpu::VCosFOp) \ + FN(arith::ExtFOp, triton::xpu::VExtFOp) \ + FN(arith::TruncFOp, triton::xpu::VTruncFOp) \ + FN(arith::SIToFPOp, triton::xpu::VSIToFPOp) + +// Does this op kind have a vector form, per the table above? Op kind only. +bool hasVectorForm(Operation *op); + +// Is the region-interpreting reduce lowering in use? +// +// **Default off.** It was flipped on 2026-08-05 and flipped back on 2026-08-06: +// with it on, a *vectorized* welford reduce silently drops the lane-0 +// contribution of most cores, and which cores are affected changes from run to +// run. Minimal repro (FlagGems `var_mean_kernel_2`, one f32 vector per core): +// feed every partial `acc=1, average=0, count=1` so the exact fold must yield +// `nvar == BLOCK_NUM`; at BLOCK_NUM=1024 it yields 970 / 977 on successive runs, +// and sweeping the payload lane by lane shows the losses land exactly on the +// multiples of 16, i.e. the seed element `collapseVectorsJointly` extracts +// first. BLOCK_NUM <= 128 puts fewer than 16 elements on a core, skips the +// within-core fold, and stays exact. Observable as +// `test_accuracy_varmean[dtype1-*-*-{dim2,dim3}-shape1]` in +// third_party/xpu/test/FlagGems/tests/test_reduction_ops.py: those eight are the +// only regressions in that 3187-case suite, and TRITONXPU_REDUCE_REGION=0 fixes +// all eight while TRITONXPU_BUDGET_TILING=0 fixes none. +// +// The default is off rather than the admission gate in +// ReduceOpToLLVM::canInterpretCombine being narrowed, because the run-to-run +// variation means the real boundary is not known yet; a gate drawn around the +// one shape we happened to measure would be a made-up bound. +// +// `TRITONXPU_REDUCE_REGION=1` opts back in. The measured upside is real and +// waiting on that defect, all on xpu3 hardware: welford 801.4us vs 1319us all +// scalar = 1.65x (findings.md 1.17), pairmax 910.6us vs 1058.8us = 1.16x (1.22), +// and of the eleven golden probes only those two plus `bitred` change code at all +// -- the other eight are byte-identical either way, and `bitred` only differs in +// emission order (same instruction multiset). Combine-op coverage is closed in +// 1.24: twelve of the thirteen ops `isSupportedCombineOp` admits have a probe. +// +// Single-sourced here because three units have to agree on it: this analysis (op +// whitelist), Vectorize (region retyping), and ReduceOpToLLVM (which lowering to +// emit). A disagreement does not degrade, it hits llvm_unreachable. +inline bool reduceCombineRegionEnabled() { + static const bool enabled = + mlir::triton::tools::isEnvValueBool( + mlir::triton::tools::getStrEnv("TRITONXPU_REDUCE_REGION")) + .value_or(false); + return enabled; +} + +// The ops the region-interpreting lowering can also emit +// (ReduceOpToLLVM::emitCombineOp). Only reachable while the region lowering is +// on: with TRITONXPU_REDUCE_REGION=0 the lowering applies one op per output and +// these would be silently dropped. Constants are deliberately left scalar by the +// retyping and splatted at the use site. +// +// arith::NegFOp is deliberately absent: Triton's unary minus lowers to +// `subf(0.0, x)`, so no combine region can contain a NegFOp (findings.md 1.24). +#define REDUCE_COMBINE_REGION_OP \ + REDUCE_COMBINE_OP, arith::SubFOp, arith::DivFOp, arith::SelectOp, \ + arith::ConstantOp + +// Element types that have a 512-bit vector form on this target. Note i1 is +// absent, which is why a bool store never vectorizes. +bool vectorizedTyValid(Type elemTy); + +// Lanes a 512-bit vector register holds for `elemTy`. There is no freedom here +// today: the width is fully determined by the element width. +unsigned getVectorWidth(Type elemTy); + +// Whether every op in the reduce's combine region has a vector form, i.e. +// whether `Vectorize`'s wholesale retyping of that region can succeed. +// +// Single-sourced on purpose: the analysis and the rewrite both have to agree on +// REDUCE_COMBINE_OP, and when they disagree the rewrite does not fall back, it +// hits llvm_unreachable. +// +// E-independent: only looks at op kinds. +bool reduceCombineIsVectorizable(triton::xpu::ReduceOp redOp); + +// Root-level report, shared by `tritonxpu-vectorizability-analysis` and by +// Vectorize's own root enumeration. +// +// It exists because the two see different IR and the difference has to be +// diffable rather than assumed: the standalone pass runs before Vectorize, so +// before that pass' prologue (erf lowering, maximum/compare fusion, i1 logic to +// i8), and those rewrites change which roots form a closure. Emitting the same +// line format from both sides is what makes the gap step 1.5 has to close +// visible. Off unless TRITONXPU_VEC_REPORT=1. +bool vecReportEnabled(); +// `cands` is the number of footprint questions the walk left open (step 1.5c); +// -1 means "not applicable", which is every caller running the real oracle. +void reportVecRoot(const char *stage, const char *site, Operation *root, + Type rootOpTy, bool eligible, int64_t closureSize, + int64_t cands = -1); + +//===----------------------------------------------------------------------===// +// The fit oracle (redesign-v2.md §2.1.1, step 1.5b). +// +// Three cases inside the closure walk -- LoadOp, SplatOp, BroadcastOp -- ask +// about the per-core footprint of a value, which is exactly what CoreTiling +// fixes. Rather than invent a pre-tiling approximation for them, each case now +// keeps only its E-independent half and asks an injected oracle for the rest. +// +// `Unknown` is the pre-tiling answer: the walk must not veto on it, and records +// the query as a `FitCandidate` instead (step 1.5c) so the E-dependent half can +// be applied once sizePerCore exists. On the E-dependent side the real oracle +// answers Yes/No only, so the walk is byte-identical to the inline tests it +// replaces and records nothing. +// +// The width is not the oracle's business -- it follows from the element type +// alone, so the walk computes it and passes it in. That deliberately keeps +// inconsistency #3 of §2.1.1 (BroadcastOp hardcodes 16, right only for f32) +// visible at the call site instead of burying it in the oracle. +//===----------------------------------------------------------------------===// + +enum class Fit { Yes, No, Unknown }; + +// What is being asked of the value's per-core footprint: +// WholeVectors -- a nonzero whole number of `wantWidth`-lane vectors +// SingleElem -- exactly one element (`wantWidth` unused) +// AtLeastWidth -- at least `wantWidth` elements +enum class FitQuery { WholeVectors, SingleElem, AtLeastWidth }; + +using FitOracle = std::function; + +// One footprint question the walk could not answer, because the oracle it was +// handed returned `Unknown`. The walk does not veto on these -- it records them +// and keeps going, so the closure it reports is the E-independent answer and +// this list is exactly what still has to be checked against E. +// +// All-or-nothing on purpose, matching today's behavior: any candidate that +// later fails the real fit vetoes the whole tree, because that is what the +// inline test it replaces did. +struct FitCandidate { + Value value; + FitQuery query; + unsigned wantWidth; +}; + +// The pre-tiling oracle: knows nothing, so every query becomes a candidate. +Fit vectorFitUnknown(Value value, FitQuery query, unsigned wantWidth); + +class VectorizabilityAnalysis { +public: + // `reduceOperandFits` is the E-dependent gate the ReduceOp case needs, passed + // in rather than called directly so this unit keeps no compile-time dependency + // on TileAnalysis. Callers pass `vectorFitsReduceOperand`; it is owned, not + // borrowed, because every caller builds the analysis from a temporary. + VectorizabilityAnalysis( + bool reduceVec, bool dumpFlag, + std::function reduceOperandFits, + FitOracle fitOracle) + : ReduceVec(reduceVec), dumpFlag(dumpFlag), + reduceOperandFits(std::move(reduceOperandFits)), + fitOracle(std::move(fitOracle)) {} + + // Every op that can be retyped if `root` is, or an empty result when the + // closure cannot be formed. The closure is bidirectional (operands and users) + // and any single unsupported op vetoes the whole tree. + bool getVectorizableClosure(Operation *root, OperationTree &visited, + OperationTree &vectorizedOps) { + fitCandidates.clear(); + return vectorize(root, visited, vectorizedOps); + } + + // The footprint questions this walk left open, valid until the next + // `getVectorizableClosure` call. Empty with the real oracle. Only meaningful + // when the walk succeeded: a vetoed branch can leave candidates behind, and + // since a veto propagates to the root, a failed closure discards them. + const llvm::SmallVector &getFitCandidates() const { + return fitCandidates; + } + +private: + // Ask the oracle, recording rather than vetoing on `Unknown`. + Fit askFit(Value value, FitQuery query, unsigned wantWidth); + + Operation *getBlockArgumentOp(Value arg); + bool binLikeOpVectorize(Value lhs, Value rhs, OperationTree &visited, + OperationTree &vectorizedOps); + bool vectorize(Operation *op, OperationTree &visited, + OperationTree &vectorizedOps); + + bool ReduceVec; + bool dumpFlag; + std::function reduceOperandFits; + FitOracle fitOracle; + llvm::SmallVector fitCandidates; +}; + +//===----------------------------------------------------------------------===// +// Vector-Flow analysis (M1, redesign-v2.md §3.1, step 2.1). +// +// The closure walk above answers one question per root: "can this whole tree be +// retyped, yes or no". Anything it cannot retype vetoes the entire tree, which +// is why a single `arith.cmpi` in a store chain costs the chain its vector form. +// +// This unit answers a different question, per SSA value rather than per root: +// which state does this value *want*. Values that must agree are unioned into +// one class; a class holding both a Vector and a Scalar pin is not an error, it +// is a boundary -- somewhere on that class' edge a pack/unpack has to go. M2 +// decides where; this unit only reports. +// +// Termination is structural, not a fixed point: the propagation is union-find +// over equality edges, so every step strictly reduces the class count and there +// is nothing to iterate. That is a deliberate deviation from the "sparse +// DataFlowAnalysis" base §3.1 names -- the constraints it lists are all +// symmetric ("operands and result same state", "init <-> iter_arg <-> yield <-> +// result"), and a symmetric constraint system is a partition, not a lattice +// climb. It also sidesteps §4.2's convergence worry on control flow outright. +// +// Nothing consumes the result yet, so the exit gate is C1 byte equality. +//===----------------------------------------------------------------------===// + +// Unset <= {Vector, Scalar} <= Conflict, as in §3.1. +enum class VState { Unset, Vector, Scalar, Conflict }; + +const char *toString(VState state); + +struct VectorFlowStats { + int64_t values = 0; // SSA values reached + int64_t classes = 0; // equality classes after unioning + int64_t vectorClasses = 0; // ... of which pinned Vector only + int64_t scalarClasses = 0; // ... Scalar only + int64_t conflictClasses = 0; // ... both, i.e. a boundary + int64_t unsetClasses = 0; // ... neither, free to go either way + int64_t unions = 0; // equality edges that actually merged + int64_t vectorPins = 0; + int64_t scalarPins = 0; + int64_t externPins = 0; // extern_elementwise, pinned Scalar for now + int64_t unknownPins = 0; // op kind not modelled, pinned Scalar + // Step 2.2: the reduce entry as a boundary rather than a veto. + int64_t reduceOps = 0; // reduces with at least one data operand + int64_t reduceVectorEntries = 0; // ... data operands whose class came out Vector + int64_t reduceEntryUnpacks = 0; // ... of those, the ones needing a real unpack +}; + +class VectorFlowAnalysis { +public: + // The oracle is the same one the closure walk takes, and for the same reason: + // whether a load's footprint is a whole number of vectors is E-dependent, so + // it must not be answered here. `Unknown` seeds nothing. + explicit VectorFlowAnalysis(FitOracle fitOracle) + : fitOracle(std::move(fitOracle)) {} + + // Partition every value in `func`, then pin. Read-only on the IR. + void run(triton::FuncOp func); + + // Unset for values `run` never saw. + VState stateOf(Value value) const; + + // Whether `run` reached this value at all. `stateOf` folds two different + // answers into Unset -- "reached, and nothing constrains it" versus "never + // reached" -- and only the first is a usable answer. The per-root report + // prints both so an unmodelled value cannot pass as an unconstrained one. + bool isTracked(Value value) const { return ids.count(value) != 0; } + + const VectorFlowStats &getStats() const { return stats; } + +private: + unsigned idOf(Value value); + unsigned find(unsigned id); + void unite(Value a, Value b); + void pin(Value value, VState state); + void visit(Operation *op); + + FitOracle fitOracle; + llvm::DenseMap ids; + llvm::SmallVector parent; + llvm::SmallVector pins; // per class root, valid after `find` + VectorFlowStats stats; +}; + +// Off unless TRITONXPU_VFLOW_REPORT=1. +bool vflowReportEnabled(); + +} // namespace xpu +} // namespace triton +} // namespace mlir + +#endif // TRITONXPU_ANALYSIS_VECTORIZABILITYANALYSIS_H diff --git a/third_party/xpu/include/triton/Dialect/LLVMXPU/IR/LLVMXPUOps.td b/third_party/xpu/include/triton/Dialect/LLVMXPU/IR/LLVMXPUOps.td index 72827430b2..367a965ee1 100644 --- a/third_party/xpu/include/triton/Dialect/LLVMXPU/IR/LLVMXPUOps.td +++ b/third_party/xpu/include/triton/Dialect/LLVMXPU/IR/LLVMXPUOps.td @@ -143,6 +143,21 @@ def XPU_SM2GMOp_v3 : XPU_IntrOp<"sm2gm_v3", [], 0> { }]; } +def XPU_GM2SMOp_v3 : XPU_IntrOp<"gm2sm_v3", [], 0> { + let arguments = (ins LLVM_AnyPointer:$src, + LLVM_AnyPointer:$dst, + I32:$offset, + I32:$size); + string llvmBuilder = [{ + auto isrc = builder.CreatePtrToInt($src, builder.getInt64Ty()); + auto ioft = builder.CreateZExtOrTrunc($offset, builder.getInt64Ty()); + auto asrc = builder.CreateAdd(isrc, ioft); + auto fsrc = builder.CreateIntToPtr(asrc, $src->getType()); + auto zero = builder.getInt32(0); + createIntrinsicCall(builder, llvm::Intrinsic::xpu_gm2sm_v3, {$dst, fsrc, zero, $size}); + }]; +} + //===----------------------------------------------------------------------===// // XPU special register op definitions //===----------------------------------------------------------------------===// diff --git a/third_party/xpu/include/triton/Dialect/TritonXPU/IR/TritonXPUOps.td b/third_party/xpu/include/triton/Dialect/TritonXPU/IR/TritonXPUOps.td index 62ea268283..50f6fc8369 100644 --- a/third_party/xpu/include/triton/Dialect/TritonXPU/IR/TritonXPUOps.td +++ b/third_party/xpu/include/triton/Dialect/TritonXPU/IR/TritonXPUOps.td @@ -150,6 +150,35 @@ def TTX_LoadOp : TTX_Op<"load", [AttrSizedOperandSegments, let results = (outs TTX_Type:$result); } +def TTX_StageSMOp : TTX_Op<"stage_sm", + [MemoryEffects<[MemRead, MemWrite]>]> { + let summary = "stage a loop-invariant GM array into cluster shared memory " + "(SM) once, shared by all cores"; + let description = [{ + Copies a loop-invariant, read-only GM array into a cluster-shared SM + buffer once, outside the grid loop. Only core 0 issues the GM->SM DMA; + all cores then synchronize on a cluster barrier before reading. The + result is a scalar SM base pointer (addrspace 2) that subsequent + `triton_xpu.load_scalar_indexed` ops read from. + }]; + let arguments = (ins TT_Ptr:$ptr, + Optional:$len, + I32:$smOffset, + I32:$bufElems, + TT_SyncAttr:$syncMode); + let results = (outs TT_Ptr:$result); +} + +def TTX_LoadScalarIndexedOp : TTX_Op<"load_scalar_indexed", + [MemoryEffects<[MemRead]>]> { + let summary = "load one scalar from a staged buffer at a runtime element " + "index and broadcast it to the result tensor"; + let arguments = (ins TT_PtrLike:$ptr, + TT_Int:$index, + TT_SyncAttr:$syncMode); + let results = (outs TT_Type:$result); +} + def TTX_StoreOp : TTX_Op<"store", [AttrSizedOperandSegments, MemoryEffects<[MemWrite]>]> { let summary = "store"; diff --git a/third_party/xpu/include/triton/Dialect/TritonXPU/Transforms/Passes.h b/third_party/xpu/include/triton/Dialect/TritonXPU/Transforms/Passes.h index 686d9e7bde..f17f27c82c 100644 --- a/third_party/xpu/include/triton/Dialect/TritonXPU/Transforms/Passes.h +++ b/third_party/xpu/include/triton/Dialect/TritonXPU/Transforms/Passes.h @@ -12,6 +12,9 @@ namespace mlir { namespace triton { namespace xpu { +constexpr llvm::StringLiteral kBF16ToFP32VecOptOffAttrName = + "triton_xpu.bf16_to_fp32_vec_opt_off"; + // Generate the pass class declarations. #define GEN_PASS_DECL #include "triton/Dialect/TritonXPU/Transforms/Passes.h.inc" diff --git a/third_party/xpu/include/triton/Dialect/TritonXPU/Transforms/Passes.td b/third_party/xpu/include/triton/Dialect/TritonXPU/Transforms/Passes.td index fec79d3f29..d3b53bd077 100644 --- a/third_party/xpu/include/triton/Dialect/TritonXPU/Transforms/Passes.td +++ b/third_party/xpu/include/triton/Dialect/TritonXPU/Transforms/Passes.td @@ -33,6 +33,18 @@ def TritonXPUCreateGM2LM : Pass<"tritonxpu-create-gm2lm", "mlir::ModuleOp"> { ]; } +def TritonXPUTLELegalize : Pass<"tritonxpu-tle-legalize", "mlir::ModuleOp"> { + let summary = "Legalize for XPU TLE kernels (separate from tritonxpu-legalize)."; + + let description = [{ + TLE-specific legalize pass. Duplicated from tritonxpu-legalize so that TLE + kernels can be legalized independently without affecting the shared + (non-TLE) legalize behavior. + }]; + + let dependentDialects = ["mlir::triton::xpu::TritonXPUDialect"]; +} + def TritonXPULegalize : Pass<"tritonxpu-legalize", "mlir::ModuleOp"> { let summary = "Legalize for XPU."; @@ -113,6 +125,17 @@ def TritonXPUDtypeConvert : Pass<"tritonxpu-dtype-convert", "mlir::ModuleOp"> { ]; } +def TritonXPULoopInvariantStaging + : Pass<"tritonxpu-loop-invariant-staging", "mlir::ModuleOp"> { + let summary = "Stage a loop-invariant small-packet gm2lm out of the grid loop."; + + let description = [{ + Hoists a loop-invariant, read-only `gm2lm` out of the grid `scf.for`. + }]; + + let dependentDialects = ["mlir::triton::xpu::TritonXPUDialect"]; +} + def TritonXPULoopGrid : Pass<"tritonxpu-loop-grid", "mlir::ModuleOp"> { let summary = "Create loop on triton programs for out-of-bounds grid_size."; @@ -151,6 +174,15 @@ def TritonXPUUnrollControl : Pass<"tritonxpu-unroll-control", "mlir::ModuleOp"> Option<"unrollNum", "unroll-num", "uint32_t", /*default*/"2", "unroll num">, + Option<"vrfBudget", "vrf-budget", + "uint32_t", /*default*/"24", + "vector-register budget used when budget-tiling is on">, + Option<"budgetTiling", "budget-tiling", + "bool", /*default*/"0", + "Derive the tile factor from the vector register budget">, + Option<"pinUnrollNum", "pin-unroll-num", + "int32_t", /*default*/"-1", + "<0 keep pinned constants, 0 drop pins, >0 override them">, ]; } @@ -158,6 +190,22 @@ def TritonXPUUnrollControl : Pass<"tritonxpu-unroll-control", "mlir::ModuleOp"> // Optimization Pass //===----------------------------------------------------------------------===// +def TritonXPUScalarAnalysis : Pass<"tritonxpu-scalar-analysis", "mlir::ModuleOp"> { + let summary = "Scalar/Vector inference via tt.make_range + use-def chain."; + + let description = [{ + Classifies each tensor SSA value as Scalar or VectorContiguous and stamps + offsetState on gm2lm/lm2gm when the address is contig-1. + }]; + + let dependentDialects = ["mlir::triton::xpu::TritonXPUDialect"]; + + let options = [ + Option<"aggressive", "aggressive", "bool", /*default=*/"false", + "When true, skip blockFromLoad guard and only update ops with offsetState=-1"> + ]; +} + def TritonXPUOffsetAnalysis : Pass<"tritonxpu-offset-analysis", "mlir::ModuleOp"> { let summary = "Analysis Ptr's Offset State."; @@ -213,6 +261,26 @@ def TritonXPUCoreTiling : Pass<"tritonxpu-core-tiling", "mlir::ModuleOp"> { } +def TritonXPUNormalize : Pass<"tritonxpu-normalize", "mlir::ModuleOp"> { + let summary = "Pre-vectorization normalization rewrites."; + + let description = [{ + Scalar-to-scalar rewrites that put the IR into the shapes the vectorizer + matches. + }]; + + let dependentDialects = ["mlir::triton::xpu::TritonXPUDialect"]; + + let options = [ + Option<"dumpFlag", "dump-flag", + "bool", /*default*/"0", + "detail dump flag">, + Option<"compareFusion", "compare-fusion", + "bool", /*default*/"false", + "compare fusion"> + ]; +} + def TritonXPUVectorize : Pass<"tritonxpu-vectorize", "mlir::ModuleOp"> { let summary = "Vectorize Calculation."; @@ -231,6 +299,23 @@ def TritonXPUVectorize : Pass<"tritonxpu-vectorize", "mlir::ModuleOp"> { ]; } +def TritonXPUAsyncLoadSchedule + : Pass<"tritonxpu-async-load-schedule", "mlir::ModuleOp"> { + let summary = "Move XPU loads near first user and insert explicit mfence."; + + let dependentDialects = [ + "mlir::triton::xpu::TritonXPUDialect", + "mlir::arith::ArithDialect", + "mlir::LLVM::XPU::LLVMXPUDialect" + ]; + + let options = [ + Option<"dumpFlag", "dump-flag", + "bool", /*default*/"0", + "detail dump flag">, + ]; +} + def TritonXPUMemoryAsync : Pass<"tritonxpu-memory-async", "mlir::ModuleOp"> { let summary = "Memory Async Optimization."; @@ -335,4 +420,43 @@ def TritonXPUCFToSCF : Pass<"tritonxpu-cf-to-scf", "mlir::ModuleOp"> { let dependentDialects = ["mlir::triton::xpu::TritonXPUDialect"]; } +def TritonXPULegalizeExternEW : Pass<"tritonxpu-legalize-extern-ew", "mlir::ModuleOp"> { + let summary = "Legalize unsupported extern_elementwise ops for non-SDNN XPU."; + let description = [{ + Rewrites tt.extern_elementwise ops whose symbols are not present in + libdevice-xpu3.bc to equivalent native MLIR ops. + }]; + let dependentDialects = ["mlir::triton::xpu::TritonXPUDialect", + "mlir::math::MathDialect", + "mlir::arith::ArithDialect"]; +} + +def TritonXPUTileAnalysis : Pass<"tritonxpu-tile-analysis", "mlir::ModuleOp"> { + let summary = "Report register pressure and tile geometry per block."; + + let dependentDialects = ["mlir::triton::xpu::TritonXPUDialect"]; + + let options = [ + Option<"vrfBudget", "vrf-budget", + "uint32_t", /*default*/"24", + "same calibrated budget the tiling decision uses">, + ]; +} + +def TritonXPUVectorizabilityAnalysis + : Pass<"tritonxpu-vectorizability-analysis", "mlir::ModuleOp"> { + let summary = "Report which vectorization roots form a vectorizable closure."; + + let dependentDialects = ["mlir::triton::xpu::TritonXPUDialect"]; + + let options = [ + Option<"reduceVec", "reduce-vec", + "bool", /*default*/"1", + "match Vectorize's ReduceVec so the same roots are enumerated">, + Option<"preTiling", "pre-tiling", + "bool", /*default*/"0", + "this instance sits ahead of CoreTiling">, + ]; +} + #endif // TRITONXPU_PASSES diff --git a/third_party/xpu/lib/Analysis/NewAnalysis/CMakeLists.txt b/third_party/xpu/lib/Analysis/NewAnalysis/CMakeLists.txt index 632e24f378..cc12d6ded6 100644 --- a/third_party/xpu/lib/Analysis/NewAnalysis/CMakeLists.txt +++ b/third_party/xpu/lib/Analysis/NewAnalysis/CMakeLists.txt @@ -9,10 +9,15 @@ add_mlir_library(TritonXPUAnalysis PtrAnalysis.cpp UseAnalysis.cpp Utility.cpp + ../ScalarAnalysis.cpp + ../TileAnalysis.cpp + ../TileDecision.cpp + ../VectorizabilityAnalysis.cpp DEPENDS # Dialect.h includes the XPU attribute-interface declarations. TritonXPUAttrDefsIncGen + TritonXPUTableGen ) target_link_libraries(TritonXPUAnalysis PRIVATE TritonXPUAnalysisSDNN) diff --git a/third_party/xpu/lib/Analysis/NewAnalysis/Utility.cpp b/third_party/xpu/lib/Analysis/NewAnalysis/Utility.cpp index e47a6f0f08..4dad5ac121 100644 --- a/third_party/xpu/lib/Analysis/NewAnalysis/Utility.cpp +++ b/third_party/xpu/lib/Analysis/NewAnalysis/Utility.cpp @@ -88,6 +88,9 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const OffsetState &state) { case OffsetState::LocallyContinuous: os << "Locally Continuous"; break; + case OffsetState::LocallyScalar: + os << "Locally Scalar"; + break; default: os << "Invalid State"; break; diff --git a/third_party/xpu/lib/Analysis/ScalarAnalysis.cpp b/third_party/xpu/lib/Analysis/ScalarAnalysis.cpp new file mode 100644 index 0000000000..6d57ebfa74 --- /dev/null +++ b/third_party/xpu/lib/Analysis/ScalarAnalysis.cpp @@ -0,0 +1,615 @@ +#include "triton/Analysis/ScalarAnalysis.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" + +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/MathExtras.h" + +#include +#include + +using namespace mlir; +using namespace mlir::triton; +using namespace mlir::dataflow; + +#define DEBUG_TYPE "tritonxpu-scalar-analysis" + +namespace mlir { +namespace triton { +namespace xpu { + +namespace { + +using State = ScalarValueState; + +static bool isTensor(Value v) { return isa(v.getType()); } + +/// Default lattice value when an SSA value is produced by an op we don't +/// recognise — depend on whether the result is a tensor (lane-shaped) or a +/// plain scalar. +static State defaultFor(Value v) { + return isTensor(v) ? State::other() : State::scalar(1); +} + +/// Try to extract the (signed) integer value of a constant scalar / splat. +static std::optional getConstInt(Value v) { + if (!v) + return std::nullopt; + Operation *def = v.getDefiningOp(); + if (!def) + return std::nullopt; + auto cst = dyn_cast(def); + if (!cst) + return std::nullopt; + Attribute attr = cst.getValue(); + if (auto ia = dyn_cast(attr)) + return ia.getValue().getSExtValue(); + if (auto dense = dyn_cast(attr)) { + if (dense.isSplat()) + return dense.getSplatValue().getSExtValue(); + } + return std::nullopt; +} + +/// Alignment of a known integer constant (0 → exactly zero / "infinite"). +static uint64_t alignOfConst(int64_t c) { + if (c == 0) + return 0; + uint64_t u = c < 0 ? static_cast(-c) : static_cast(c); + return u; +} + +/// Saturating multiply of two alignment values, treating 0 as identity-zero +/// (i.e. exactly-zero values stay exactly-zero under multiplication). +static uint64_t mulAlign(uint64_t a, uint64_t b) { + if (a == 0 || b == 0) + return 0; + // Avoid overflow; cap at a large power of two. + if (a > (uint64_t(1) << 32) || b > (uint64_t(1) << 32)) + return 1; + return a * b; +} + +/// Lane count of a tensor result, or 0 if not a ranked tensor. +static int64_t laneCount(Value v) { + auto rt = dyn_cast(v.getType()); + if (!rt) + return 0; + int64_t n = 1; + for (int64_t d : rt.getShape()) + n *= d; + return n; +} + +/// `arith.addi` / `arith.subi` transfer. +static State combineAdd(const State &a, const State &b, bool isSub) { + if (a.isUnknown() || b.isUnknown()) + return State::unknown(); + if (a.isOther() || b.isOther()) + return State::other(); + uint64_t ba = State::mergeAlign(a.baseAlign, b.baseAlign); + if (a.isScalar() && b.isScalar()) + return State::scalar(ba); + if (a.isScalar() && b.isContig()) + return State::contig(isSub ? -b.stride : b.stride, ba); + if (b.isScalar() && a.isContig()) + return State::contig(a.stride, ba); + // Scalar + BlockX -> BlockX (Scalar is lane-uniform, doesn't disturb the + // per-block pattern; blockStride is preserved). + if (a.isScalar() && (b.isBlockScalar() || b.isBlockContig())) { + int64_t bs = b.blockStride; + if (isSub) bs = -bs; + if (b.isBlockScalar()) + return State::blockScalar(b.rowLen, ba, bs, b.blockStrideKnown, + b.blockFromLoad); + if (isSub) + return State::other(); + return State::blockContig(b.rowLen, ba, bs, b.blockStrideKnown, + b.blockFromLoad); + } + if (b.isScalar() && (a.isBlockScalar() || a.isBlockContig())) { + if (a.isBlockScalar()) + return State::blockScalar(a.rowLen, ba, a.blockStride, + a.blockStrideKnown, a.blockFromLoad); + return State::blockContig(a.rowLen, ba, a.blockStride, + a.blockStrideKnown, a.blockFromLoad); + } + // BlockScalar + BlockContig with same rowLen -> BlockContig. + // The resulting blockStride is the sum (or difference) of the two + // blockStrides. BlockContig's blockStride is "0" for the canonical + // `arange % R` pattern; BlockScalar's blockStride is "S" for + // `(arange / R) * S`. + // When subtracting a BlockContig, the intra-block stride flips sign + // (e.g. `S - (idx%R)` has stride=-1 within each R-block). + if (a.isBlockScalar() && b.isBlockContig() && a.rowLen == b.rowLen) { + if (isSub) + return State::other(); + bool bsk = a.blockStrideKnown && b.blockStrideKnown; + int64_t bs = bsk ? (a.blockStride + b.blockStride) : 0; + return State::blockContig(a.rowLen, ba, bs, bsk, + a.blockFromLoad || b.blockFromLoad); + } + if (b.isBlockScalar() && a.isBlockContig() && a.rowLen == b.rowLen) { + if (isSub) + return State::other(); + bool bsk = a.blockStrideKnown && b.blockStrideKnown; + int64_t bs = bsk ? (a.blockStride + b.blockStride) : 0; + return State::blockContig(a.rowLen, ba, bs, bsk, + a.blockFromLoad || b.blockFromLoad); + } + // BlockContig(R1) + BlockScalar(R2) where R1 | R2: BlockScalar is uniform + // within each R1-block (since R1 divides R2), so contig pattern preserved. + if (a.isBlockContig() && b.isBlockScalar() && a.rowLen != b.rowLen) { + int64_t R1 = a.rowLen, R2 = b.rowLen; + if (R1 > 0 && R2 > 0 && (R2 % R1) == 0) { + return State::blockContig(R1, ba, 0, false, + a.blockFromLoad || b.blockFromLoad); + } + } + if (b.isBlockContig() && a.isBlockScalar() && a.rowLen != b.rowLen) { + int64_t R1 = b.rowLen, R2 = a.rowLen; + if (R1 > 0 && R2 > 0 && (R2 % R1) == 0) { + if (isSub) + return State::other(); + return State::blockContig(R1, ba, 0, false, + a.blockFromLoad || b.blockFromLoad); + } + } + // BlockScalar + BlockScalar with same rowLen -> BlockScalar. + if (a.isBlockScalar() && b.isBlockScalar() && a.rowLen == b.rowLen) { + bool bsk = a.blockStrideKnown && b.blockStrideKnown; + int64_t bs = bsk ? (isSub ? a.blockStride - b.blockStride + : a.blockStride + b.blockStride) + : 0; + return State::blockScalar(a.rowLen, ba, bs, bsk, + a.blockFromLoad || b.blockFromLoad); + } + // BlockScalar(R1) + BlockScalar(R2) where R1 | R2 or R2 | R1: + // uniform within the smaller block is preserved. + if (a.isBlockScalar() && b.isBlockScalar() && a.rowLen != b.rowLen) { + int64_t R1 = a.rowLen, R2 = b.rowLen; + if (R1 > 0 && R2 > 0) { + int64_t Rmin = std::min(R1, R2); + int64_t Rmax = std::max(R1, R2); + if ((Rmax % Rmin) == 0) { + return State::blockScalar(Rmin, ba, 0, false, + a.blockFromLoad || b.blockFromLoad); + } + } + } + // BlockContig + BlockContig with same rowLen: stride doubles -> not contig-1 + // anymore; degrade. + if (a.isBlockContig() && b.isBlockContig()) + return State::other(); + if (a.isContig() && b.isContig()) { + int64_t s = isSub ? a.stride - b.stride : a.stride + b.stride; + return State::contig(s, ba); + } + // Mixed (e.g. VectorContig + BlockX) — conservative. + return State::other(); +} + +/// `arith.muli` transfer. Sharpens with a known constant operand when +/// possible: +/// Scalar(ba) * const(c) -> Scalar(ba * |c|) +/// VC(s, ba) * const(c) -> VC(s*c, ba * |c|) (tensor result) +static State combineMul(Value lhs, Value rhs, const State &a, const State &b, + bool resultIsTensor) { + if (a.isUnknown() || b.isUnknown()) + return State::unknown(); + if (a.isOther() || b.isOther()) + return resultIsTensor ? State::other() : State::scalar(1); + + std::optional cl = getConstInt(lhs); + std::optional cr = getConstInt(rhs); + + // Scalar * Scalar -> Scalar + if (a.isScalar() && b.isScalar()) { + uint64_t ba = mulAlign(a.baseAlign, b.baseAlign); + return State::scalar(ba); + } + // VC * Scalar with known constant on the scalar side -> sharpen stride. + if (a.isContig() && b.isScalar() && cr) { + int64_t c = *cr; + return State::contig(a.stride * c, mulAlign(a.baseAlign, alignOfConst(c))); + } + if (b.isContig() && a.isScalar() && cl) { + int64_t c = *cl; + return State::contig(b.stride * c, mulAlign(b.baseAlign, alignOfConst(c))); + } + // VC * Scalar without a known constant: stride unknown -> conservative. + // BlockScalar * Scalar (or const) -> BlockScalar (rowLen preserved). + // If the scalar side is a known constant, scale blockStride by that const. + if (a.isBlockScalar() && b.isScalar()) { + uint64_t ba = cr ? mulAlign(a.baseAlign, alignOfConst(*cr)) + : mulAlign(a.baseAlign, b.baseAlign); + if (cr && a.blockStrideKnown) { + return State::blockScalar(a.rowLen, ba, a.blockStride * (*cr), true, + a.blockFromLoad); + } + return State::blockScalar(a.rowLen, ba, 0, false, a.blockFromLoad); + } + if (b.isBlockScalar() && a.isScalar()) { + uint64_t ba = cl ? mulAlign(b.baseAlign, alignOfConst(*cl)) + : mulAlign(b.baseAlign, a.baseAlign); + if (cl && b.blockStrideKnown) { + return State::blockScalar(b.rowLen, ba, b.blockStride * (*cl), true, + b.blockFromLoad); + } + return State::blockScalar(b.rowLen, ba, 0, false, b.blockFromLoad); + } + // BlockContig * const: stride changes from 1 to c -> no longer contig-1. + // Degrade unless c == 1. + if (a.isBlockContig() && b.isScalar()) { + if (cr && *cr == 1) + return a; + return resultIsTensor ? State::other() : State::scalar(1); + } + if (b.isBlockContig() && a.isScalar()) { + if (cl && *cl == 1) + return b; + return resultIsTensor ? State::other() : State::scalar(1); + } + return resultIsTensor ? State::other() : State::scalar(1); +} + +/// `arith.remsi` transfer with a constant divisor `c`. +/// Scalar(ba) % c -> Scalar(gcd(ba, c)) +/// VC(s, ba) % c, where the lane span s*(n-1) < gcd(ba, c) +/// -> VC(s, gcd(ba, c)) +/// Otherwise conservative. +static State remsiByConst(const State &lhs, int64_t c, int64_t lanes, + bool resultIsTensor) { + if (c == 0) + return resultIsTensor ? State::other() : State::scalar(1); + uint64_t cu = static_cast(c < 0 ? -c : c); + if (lhs.isScalar()) { + uint64_t ba = State::mergeAlign(lhs.baseAlign, cu); + return State::scalar(ba); + } + if (lhs.isContig()) { + uint64_t ba = State::mergeAlign(lhs.baseAlign, cu); + int64_t span = lhs.stride * (lanes > 0 ? lanes - 1 : 0); + int64_t absSpan = span < 0 ? -span : span; + // No-wrap condition: every lane stays within a single c-block. + if (lhs.baseAlign != 0 && (cu % lhs.baseAlign) == 0 && + static_cast(absSpan) < lhs.baseAlign) { + return State::contig(lhs.stride, ba); + } + if (lhs.baseAlign == 0 && static_cast(absSpan) < cu) { + // Base is exactly zero; full span fits in [0, c). + return State::contig(lhs.stride, ba); + } + // Wraps: VC(1, ba) % c with span >= c. Within each c-lane chunk the + // pattern is contig-1 (possibly shifted by base%c). Promote to + // BlockContig(rowLen=c). Used by patterns like `xindex % 576` when + // followed by further `% 64` -- as long as 64 | 576 the wraps align + // with the 64-lane block boundaries. + // The canonical `arange % R` carries blockStride = 0 (no inter-row term). + if (lhs.stride == 1) { + return State::blockContig(c < 0 ? -c : c, ba, 0, true); + } + } + // BlockContig(rowLen=R, stride=1) % c with c | R -> BlockContig(rowLen=c). + if (lhs.isBlockContig() && lhs.stride == 1) { + int64_t R = lhs.rowLen; + if (R > 0 && c > 0) { + if ((R % c) == 0) { + uint64_t ba = State::mergeAlign(lhs.baseAlign, cu); + return State::blockContig(c, ba, 0, true); + } + // c >= R: within each R-block lanes are 0..R-1, all < c -> unchanged. + if ((c % R) == 0) { + uint64_t ba = State::mergeAlign(lhs.baseAlign, cu); + return State::blockContig(R, ba, lhs.blockStride, lhs.blockStrideKnown, + lhs.blockFromLoad); + } + } + } + // BlockScalar % c -> BlockScalar (uniform within block stays uniform). + if (lhs.isBlockScalar()) { + uint64_t ba = State::mergeAlign(lhs.baseAlign, cu); + return State::blockScalar(lhs.rowLen, ba, 0, false, lhs.blockFromLoad); + } + return resultIsTensor ? State::other() : State::scalar(1); +} + +/// `arith.divsi` transfer with a constant divisor `c`. +/// Scalar(ba) / c, where c | ba -> Scalar(ba / c) +/// VC(s, ba) / c, where c | ba and stride*(n-1) < c +/// -> Scalar(ba / c) (all lanes equal) +/// Otherwise conservative. +static State divsiByConst(const State &lhs, int64_t c, int64_t lanes, + bool resultIsTensor) { + if (c == 0) + return resultIsTensor ? State::other() : State::scalar(1); + uint64_t cu = static_cast(c < 0 ? -c : c); + if (lhs.isScalar()) { + if (lhs.baseAlign == 0) + return State::scalar(0); + if ((lhs.baseAlign % cu) == 0) + return State::scalar(lhs.baseAlign / cu); + return State::scalar(1); + } + if (lhs.isContig()) { + int64_t span = lhs.stride * (lanes > 0 ? lanes - 1 : 0); + int64_t absSpan = span < 0 ? -span : span; + uint64_t abs_u = static_cast(absSpan); + bool baseIsZero = (lhs.baseAlign == 0); + bool baseDividesByC = + (lhs.baseAlign != 0) && ((lhs.baseAlign % cu) == 0); // c | B + bool cDividesByBase = + (lhs.baseAlign != 0) && ((cu % lhs.baseAlign) == 0); // B | c + // Case 2a: v[0] on a c-boundary (v[0]=0 or c | B), and absSpan < c + // → all lanes fall in the same c-bucket → Scalar. + // Quotient k = v[0]/c = m·(B/c), so result baseAlign = B/c + // (or 0 when baseIsZero, i.e. quotient is exactly zero). + // + // Threshold must be c (not baseAlign): when B > c, a baseAlign-window + // straddles B/c c-buckets, so "no wrap across baseAlign" does NOT imply + // "no wrap across c". + // + // Requires stride >= 0: with negative stride the segment extends below + // v[0], and being on a c-boundary puts v[0] at the top of the previous + // bucket, so the segment immediately leaves it. + // + // Example: `(pid_y*8192 + arange(0, 64)) // 128` + // B=8192, c=128, k = m·(B/c) = m·64 + // → Scalar(baseAlign=64) + if ((baseIsZero || baseDividesByC) && lhs.stride >= 0 && abs_u < cu) { + uint64_t ba = baseIsZero ? 0 : (lhs.baseAlign / cu); + return State::scalar(ba); + } + // Case 2b: B | c (c is an integer multiple of B). v[0] lies on the + // B-grid but may sit at the tail of a c-bucket. When absSpan < B, the + // segment cannot reach the next B-grid point → stays in one c-bucket + // → Scalar. Quotient k = m ÷ (c/B) (integer division) can take any + // integer value → result baseAlign = 1. + // + // Threshold must be B (not c): v[0] may lie only B away from the next + // c-boundary, so absSpan reaching B risks crossing it. Dual to Case 2a + // where v[0] is guaranteed at a c-bucket start and can span a full c. + // + // Example: `(pid*1024 + arange(0, 1024)) // 409600` + // B=1024, c=400·B, k = m÷(c/B) = m÷400 (integer division) + // → Scalar(baseAlign=1) + if (!baseIsZero && cDividesByBase && lhs.stride >= 0 && + abs_u < lhs.baseAlign) { + return State::scalar(1); + } + // Wraps: VC(1, ba) / c with span >= c. Across c-lane chunks the + // quotient increases by 1; within each c-lane chunk it's uniform. + // -> BlockScalar(rowLen=c) with blockStride = 1. + if (lhs.stride == 1) { + uint64_t ba = baseIsZero ? 0 : 1; + return State::blockScalar(c < 0 ? -c : c, ba, 1, true); + } + } + // BlockScalar / c: uniform within block stays uniform -> BlockScalar. + if (lhs.isBlockScalar()) { + return State::blockScalar(lhs.rowLen, 1, 0, false, lhs.blockFromLoad); + } + // BlockContig(rowLen=R, stride=1) / c with c | R -> BlockScalar(rowLen=c). + // Within each R-block lanes are 0..R-1; dividing by c gives (R/c) groups of + // c equal quotients per R-block -> uniform within c-blocks. + if (lhs.isBlockContig() && lhs.stride == 1) { + int64_t R = lhs.rowLen; + if (R > 0 && c > 0 && (R % c) == 0) { + return State::blockScalar(c, 1); + } + // R | c: within each R-block span = R-1 < c, so all lanes in a block + // share the same quotient -> BlockScalar(rowLen=R). + if (R > 0 && c > 0 && (c % R) == 0) { + return State::blockScalar(R, 1, 0, false, lhs.blockFromLoad); + } + } + return resultIsTensor ? State::other() : State::scalar(1); +} + +} // namespace + +LogicalResult ScalarAnalysis::visitOperation( + Operation *op, + ArrayRef *> operands, + ArrayRef *> results) { + + auto setResult = [&](unsigned i, const State &st) { + propagateIfChanged(results[i], results[i]->join(st)); + }; + + // tt.make_range : VectorContig(stride = 1, baseAlign = |start|). + if (auto mr = dyn_cast(op)) { + int64_t start = static_cast(mr.getStart()); + setResult(0, State::contig(1, alignOfConst(start))); + return success(); + } + + // tt.splat : Scalar carrying the operand's alignment. + if (isa(op)) { + State a = operands[0]->getValue(); + if (a.isUnknown()) { + setResult(0, State::unknown()); + } else if (a.isScalar()) { + setResult(0, State::scalar(a.baseAlign)); + } else { + // Splatting a value that itself was tracked as a vector shouldn't + // happen in practice; be conservative. + setResult(0, State::scalar(1)); + } + return success(); + } + + // arith.constant : splat dense / scalar constant ⇒ Scalar(alignOfConst); + // non-splat dense tensor ⇒ VectorOther. + if (auto cst = dyn_cast(op)) { + Value res = cst.getResult(); + Attribute attr = cst.getValue(); + if (!isTensor(res)) { + if (auto ia = dyn_cast(attr)) { + setResult(0, State::scalar(alignOfConst(ia.getValue().getSExtValue()))); + } else { + setResult(0, State::scalar(1)); + } + } else if (auto dense = dyn_cast(attr)) { + if (dense.isSplat()) { + if (auto di = dyn_cast(dense)) { + setResult(0, State::scalar(alignOfConst( + di.getSplatValue().getSExtValue()))); + } else { + setResult(0, State::scalar(1)); + } + } else { + setResult(0, State::other()); + } + } else { + setResult(0, State::other()); + } + return success(); + } + + // arith.extsi / arith.extui / arith.trunci : integer width cast does NOT + // change the lane pattern (Scalar / Contig / BlockX all preserved). + // Without this rule, expressions like + // xindex = pid.to(i64) * XBLOCK + arange(0, XBLOCK).to(i64) + // would degrade to VectorOther at the extsi on `arange`, breaking the + // entire downstream BlockContig / Continuous detection. + if (isa(op)) { + setResult(0, operands[0]->getValue()); + return success(); + } + + // arith.addi / arith.subi. + if (isa(op)) { + State a = operands[0]->getValue(); + State b = operands[1]->getValue(); + setResult(0, combineAdd(a, b, isa(op))); + return success(); + } + + // arith.muli. + if (isa(op)) { + State a = operands[0]->getValue(); + State b = operands[1]->getValue(); + setResult(0, combineMul(op->getOperand(0), op->getOperand(1), a, b, + isTensor(op->getResult(0)))); + return success(); + } + + // arith.remsi by a known constant divisor. + if (isa(op)) { + State a = operands[0]->getValue(); + auto rhs = getConstInt(op->getOperand(1)); + if (!rhs) { + // Non-constant divisor: BlockScalar % anything is still BlockScalar + // (uniform within each rowLen-lane block stays uniform). All other + // kinds degrade conservatively. + if (a.isBlockScalar()) { + setResult(0, State::blockScalar(a.rowLen, 1, 0, false, + a.blockFromLoad)); + } else { + setResult(0, defaultFor(op->getResult(0))); + } + return success(); + } + setResult(0, remsiByConst(a, *rhs, laneCount(op->getResult(0)), + isTensor(op->getResult(0)))); + return success(); + } + + // arith.divsi / divui by a known constant divisor. + if (isa(op)) { + State a = operands[0]->getValue(); + auto rhs = getConstInt(op->getOperand(1)); + if (!rhs) { + // Non-constant divisor: BlockScalar / anything is still BlockScalar. + if (a.isBlockScalar()) { + setResult(0, State::blockScalar(a.rowLen, 1, 0, false, + a.blockFromLoad)); + } else { + setResult(0, defaultFor(op->getResult(0))); + } + return success(); + } + setResult(0, divsiByConst(a, *rhs, laneCount(op->getResult(0)), + isTensor(op->getResult(0)))); + return success(); + } + + // arith.select : if both branches share the same lattice kind, the + // selected value retains that kind (with merged alignments). Otherwise + // degrade conservatively. Note: this assumes the condition itself does + // not need to be lane-uniform — we only care about the value lattice. + if (auto sel = dyn_cast(op)) { + State t = operands[1]->getValue(); + State f = operands[2]->getValue(); + setResult(0, State::join(t, f)); + return success(); + } + + // tt.addptr(ptr_tensor, offset_tensor) -> ptr_tensor. + if (auto addPtr = dyn_cast(op)) { + State base = operands[0]->getValue(); + State off = operands[1]->getValue(); + if (base.isUnknown() || off.isUnknown()) { + setResult(0, State::unknown()); + } else if (base.isOther() || off.isOther()) { + setResult(0, State::other()); + } else if (base.isScalar() && off.isScalar()) { + setResult(0, State::scalar(1)); + } else if (base.isScalar()) { + // Pointer base lane-uniform; per-lane offset dictates the pattern. + setResult(0, off); + } else if (off.isScalar()) { + setResult(0, base); + } else if (base.isContig() && off.isContig()) { + // Both VectorContig. + setResult(0, State::contig(base.stride + off.stride, 1)); + } else { + // Mixed including BlockX / VectorContig — too complex to track precisely. + setResult(0, State::other()); + } + return success(); + } + + // triton_xpu.gm2lm : transparent w.r.t. the pointer's lane pattern. The + // output is a pointer-tensor pointing into LM that mirrors the GM input. + if (isa(op)) { + setResult(0, operands[0]->getValue()); + return success(); + } + + // triton_xpu.load : if the pointer is lane-uniform (Scalar), the loaded + // value is also lane-uniform across active lanes. We don't know the + // value's alignment, so use baseAlign = 1. + if (isa(op)) { + State ptr = operands[0]->getValue(); + if (ptr.isScalar()) { + setResult(0, State::scalar(1)); + return success(); + } + // BlockScalar pointer -> loaded value uniform within each block. + if (ptr.isBlockScalar()) { + // Mark `blockFromLoad`: the per-block value is produced by a runtime + // gather. This is the genuine k89 embedding-gather case that the + // LocallyContinuous(rowStride=-1) row-by-row DMA path was designed for. + setResult(0, State::blockScalar(ptr.rowLen, 1, 0, false, + /*fromLoad=*/true)); + return success(); + } + setResult(0, defaultFor(op->getResult(0))); + return success(); + } + + // Default: produce VectorOther for tensor results, Scalar for non-tensor. + for (unsigned i = 0; i < op->getNumResults(); ++i) { + setResult(i, defaultFor(op->getResult(i))); + } + return success(); +} + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Analysis/TileAnalysis.cpp b/third_party/xpu/lib/Analysis/TileAnalysis.cpp new file mode 100644 index 0000000000..f5b2116b7c --- /dev/null +++ b/third_party/xpu/lib/Analysis/TileAnalysis.cpp @@ -0,0 +1,316 @@ +#include "triton/Analysis/TileAnalysis.h" +#include "triton/Analysis/NewAnalysis/Utility.h" +// For getVectorWidth / vectorizedTyValid: the gates below ask how wide a vector +// of a given element type is, which is the state half's answer. The dependency +// is one-way on purpose -- the state unit must not come to depend on this one. +#include "triton/Analysis/VectorizabilityAnalysis.h" + +#include "triton/Tools/Sys/GetEnv.hpp" + +namespace mlir { +namespace triton { +namespace xpu { + +ClusterLayoutAttr getClusterLayout(RankedTensorType tensorTy) { + if (auto sliceEncoding = + dyn_cast(tensorTy.getEncoding())) + return dyn_cast(sliceEncoding.getParent()); + return dyn_cast(tensorTy.getEncoding()); +} + +int64_t getNumRegs(Type type, bool *isVector) { + if (isVector) + *isVector = false; + auto tensorTy = dyn_cast(type); + if (!tensorTy) + return 0; + auto clusterEncoding = getClusterLayout(tensorTy); + if (!clusterEncoding) + return 0; + int64_t units = 1; + for (auto sizePerCore : clusterEncoding.getSizePerCore()) + units *= sizePerCore; + units = std::max(units, 1); + auto elemTy = getElementTypeOrSelf(tensorTy); + if (isa(elemTy)) { + if (isVector) + *isVector = true; + return units; + } + // A tensor of pointers is an address computation, not a live data value: + // gm2lm/alloca own that footprint and it does not scale with the tile. + if (!elemTy.isIntOrFloat()) + return 0; + // sizePerCore counts elements here and a scalar register holds one of them. + return units; +} + +void getRegPressure(ModuleOp m, const llvm::SetVector &opTree, + RegPressure &p) { + DenseMap op2Line; + getOpLine(m, op2Line); + + SmallVector ordered(opTree.begin(), opTree.end()); + llvm::sort(ordered, [&](Operation *lhs, Operation *rhs) { + return op2Line[lhs] < op2Line[rhs]; + }); + + DenseMap vecRegs, scalarRegs; + DenseMap lastUse; + auto note = [&](Value val) { + bool isVector = false; + int64_t regs = getNumRegs(val.getType(), &isVector); + if (!regs) + return false; + if (!isVector) { + scalarRegs[val] = regs; + return true; + } + vecRegs[val] = regs; + if (auto layout = getClusterLayout(cast(val.getType()))) { + int64_t width = layout.getSizePerCore().back(); + p.maxVecWidth = std::max(p.maxVecWidth, width); + p.minVecWidth = + p.minVecWidth ? std::min(p.minVecWidth, width) : width; + } + return true; + }; + for (auto *op : ordered) { + unsigned line = op2Line[op]; + for (auto res : op->getResults()) + note(res); + for (auto operand : op->getOperands()) { + if (!note(operand)) + continue; + auto it = lastUse.find(operand); + if (it == lastUse.end() || it->second < line) + lastUse[operand] = line; + } + } + + // Values defined outside the tree are live on entry. Values with no use + // inside the tree are conservatively kept live to the end. + auto peakOf = [&](DenseMap ®s, int64_t &total) { + int64_t live = 0; + total = 0; + DenseMap deathAtLine; + for (auto &[val, n] : regs) { + total += n; + auto *defOp = val.getDefiningOp(); + if (!defOp || !opTree.contains(defOp)) + live += n; + if (auto it = lastUse.find(val); it != lastUse.end()) + deathAtLine[it->second] += n; + } + int64_t peak = live; + for (auto *op : ordered) { + for (auto res : op->getResults()) + live += regs.lookup(res); + peak = std::max(peak, live); + live -= deathAtLine.lookup(op2Line[op]); + } + return peak; + }; + p.vecPeak = peakOf(vecRegs, p.vecTotal); + p.scalarPeak = peakOf(scalarRegs, p.scalarTotal); +} + +void getBlockRegPressure(ModuleOp m, Operation *insertPt, RegPressure &p) { + Block *block = insertPt->getBlock(); + if (!block) + return; + llvm::SetVector blockOps; + for (auto &op : *block) + blockOps.insert(&op); + getRegPressure(m, blockOps, p); +} + +//===----------------------------------------------------------------------===// +// E-dependent gates, moved here from VectorizabilityAnalysis. +//===----------------------------------------------------------------------===// + +bool vectorFitsRoot(Type rootOpTy) { + auto rowsPerCore = 1; + if (auto rootOpTensorTy = mlir::dyn_cast(rootOpTy)) { + auto rank = rootOpTensorTy.getShape().size(); + if (rank > 1) { + rowsPerCore = mlir::cast( + rootOpTensorTy.getEncoding()) + .getSizePerCore()[0]; + } + } + + unsigned numElems = getTotalElemsPerThread(rootOpTy) / rowsPerCore; + Type vecTy = getElementTypeOrSelf(rootOpTy); + Type elemTy = getElementTypeOrSelf(vecTy); + auto vectorWidth = getVectorWidth(elemTy); + return numElems >= vectorWidth && numElems % vectorWidth == 0 && + vectorizedTyValid(elemTy); +} + +bool vectorFitsReduceOperand(triton::xpu::ReduceOp redOp, Type operandTy) { + unsigned numElems = 0; + auto axis = redOp.getAxis(); + + if (auto operandTensorTy = dyn_cast(operandTy)) { + auto operandShape = operandTensorTy.getShape(); + numElems = operandShape[axis]; + } + + Type vecTy = getElementTypeOrSelf(operandTy); + Type elemTy = getElementTypeOrSelf(vecTy); + auto elemWidth = elemTy.getIntOrFloatBitWidth(); + auto vectorWidth = 512 / elemWidth; + + if (numElems < vectorWidth || numElems % vectorWidth > 0 || + !vectorizedTyValid(elemTy)) + return false; + + return true; +} + +Fit vectorFitsValue(Value value, FitQuery query, unsigned wantWidth) { + unsigned numElems = getTotalElemsPerThread(value.getType()); + + switch (query) { + case FitQuery::WholeVectors: + // `wantWidth == 0` is unreachable today (it is 512 / elemWidth for an + // int-or-float element type), the guard only keeps the division defined. + return (wantWidth != 0 && numElems != 0 && numElems % wantWidth == 0) + ? Fit::Yes + : Fit::No; + case FitQuery::SingleElem: + return numElems == 1 ? Fit::Yes : Fit::No; + case FitQuery::AtLeastWidth: + return numElems >= wantWidth ? Fit::Yes : Fit::No; + } + llvm_unreachable("unhandled FitQuery"); +} + +//===----------------------------------------------------------------------===// +// M6 -- the tile plan carrier. +//===----------------------------------------------------------------------===// + +namespace { + +// Kernel plus printed location. The kernel name is part of the key because the +// probe suite compiles the same kernel twice at different shapes (softmax / +// shortrow), so `loc` alone is not even unique across a run. +std::string locKeyOf(Operation *op) { + std::string key; + llvm::raw_string_ostream os(key); + if (auto funcOp = op->getParentOfType()) + os << funcOp.getName() << "|"; + else + os << "|"; + op->getLoc().print(os); + return key; +} + +} // namespace + +bool tilePlanProbeEnabled() { + return mlir::triton::tools::getBoolEnv("TRITONXPU_TILE_PLAN"); +} + +void tilePlanRecord(ModuleOp mod, Operation *root, StringRef site, bool eligible, + int64_t closure) { + if (!tilePlanProbeEnabled() || !root) + return; + + MLIRContext *ctx = mod.getContext(); + SmallVector entries; + if (auto existing = mod->getAttrOfType(kTilePlanAttrName)) + entries.assign(existing.begin(), existing.end()); + + // The id is the entry's own index, so the two keys stay independent: nothing + // on the consumer side needs the array to be searched by id to find it. + int64_t id = entries.size(); + auto i64Ty = IntegerType::get(ctx, 64); + auto named = [&](StringRef name, Attribute value) { + return NamedAttribute(StringAttr::get(ctx, name), value); + }; + entries.push_back(DictionaryAttr::get( + ctx, {named("closure", IntegerAttr::get(i64Ty, closure)), + named("eligible", BoolAttr::get(ctx, eligible)), + named("id", IntegerAttr::get(i64Ty, id)), + named("loc", StringAttr::get(ctx, locKeyOf(root))), + named("root", StringAttr::get(ctx, root->getName().getStringRef())), + named("site", StringAttr::get(ctx, site))})); + + mod->setAttr(kTilePlanAttrName, ArrayAttr::get(ctx, entries)); + root->setAttr(kTilePlanIdAttrName, IntegerAttr::get(i64Ty, id)); +} + +void tilePlanCheck(ModuleOp mod) { + auto plan = mod->getAttrOfType(kTilePlanAttrName); + + // Unconditional, and before any early return: whatever the probe wrote must + // not reach the emitted IR, or C1 fails for a reason that has nothing to do + // with the decision being measured. + mod->removeAttr(kTilePlanAttrName); + llvm::DenseMap> byId; + llvm::StringMap> byLoc; + mod.walk([&](Operation *op) { + if (auto idAttr = op->getAttrOfType(kTilePlanIdAttrName)) { + byId[idAttr.getInt()].push_back(op); + op->removeAttr(kTilePlanIdAttrName); + } + byLoc[locKeyOf(op)].push_back(op); + }); + + if (!plan) + return; + + int64_t idHits = 0, locHits = 0, locKindHits = 0; + for (Attribute entry : plan) { + auto dict = dyn_cast(entry); + if (!dict) + continue; + int64_t id = cast(dict.get("id")).getInt(); + StringRef locKey = cast(dict.get("loc")).getValue(); + StringRef site = cast(dict.get("site")).getValue(); + StringRef rootName = cast(dict.get("root")).getValue(); + + auto idIt = byId.find(id); + size_t nById = idIt == byId.end() ? 0 : idIt->second.size(); + auto locIt = byLoc.find(locKey); + size_t nByLoc = locIt == byLoc.end() ? 0 : locIt->second.size(); + + // The same key narrowed by op kind, which a consumer legitimately knows: it + // enumerates stores, or reduces, not arbitrary ops. Reported so that "loc is + // ambiguous" cannot be answered with "then also match the op name" -- the + // number for that variant is right here. + size_t nByLocKind = 0; + if (locIt != byLoc.end()) + for (Operation *op : locIt->second) + if (op->getName().getStringRef() == rootName) + ++nByLocKind; + + // A hit is exactly one live op. Zero means the key did not survive; more + // than one means it does not identify a root, which is just as unusable -- + // counting either as a hit is the fallback masking the exit gate forbids. + if (nById == 1) + ++idHits; + if (nByLoc == 1) + ++locHits; + if (nByLocKind == 1) + ++locKindHits; + + // What the surviving op turned into, when the id found it: an op that was + // rebuilt under a different name is the interesting failure mode for `loc`. + StringRef nowName = nById == 1 ? idIt->second.front()->getName().getStringRef() + : StringRef("-"); + llvm::errs() << "[TilePlan] id=" << id << " site=" << site + << " root=" << rootName << " now=" << nowName + << " byId=" << nById << " byLoc=" << nByLoc + << " byLocKind=" << nByLocKind << " key=" << locKey << "\n"; + } + llvm::errs() << "[TilePlan] summary entries=" << plan.size() + << " idHits=" << idHits << " locHits=" << locHits + << " locKindHits=" << locKindHits << "\n"; +} + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Analysis/TileDecision.cpp b/third_party/xpu/lib/Analysis/TileDecision.cpp new file mode 100644 index 0000000000..10a5e9f82e --- /dev/null +++ b/third_party/xpu/lib/Analysis/TileDecision.cpp @@ -0,0 +1,165 @@ +#include "triton/Analysis/TileDecision.h" + +#include "llvm/Support/MathExtras.h" + +namespace mlir { +namespace triton { +namespace xpu { + +int64_t vrfBudgetTarget(const TileContext &ctx) { + // Trip count a tree of the widest row would need to fit the budget, ... + int64_t blockTarget = + ctx.vrfBudget > 0 ? llvm::divideCeil(ctx.peakVRegs, ctx.vrfBudget) : 1; + // ... converted to this tree's row: what is shared between the trees in one + // block is the per-iteration width, not the number of iterations. + int64_t target = llvm::divideCeil(blockTarget * ctx.widthPerCore, + std::max(ctx.maxVecWidth, 1)); + return std::max(target, 1); +} + +namespace { + +class VRFBudgetCriterion : public TileCriterion { +public: + llvm::StringRef name() const override { return "vrf-budget"; } + unsigned tier() const override { return 1; } + + bool isFeasible(const TileCandidate &cand, const TileContext &ctx, + std::string &why) const override { + // The smallest surviving trip count is the largest tile that fits the + // budget, so admitting everything at or above the target never over-tiles: + // tier 3 picks the smallest one back. + if (cand.iterNum >= vrfBudgetTarget(ctx)) + return true; + why = "budget-unreachable"; + return false; + } +}; + +// Placeholder for M3. Deciding nothing is deliberate: there is no calibrated +// capacity number yet (§3.3.1 -- legalize does not enforce an alloca ceiling +// either), and a criterion with a made-up bound would silently move decisions. +// It stays on the tier so the remark has a slot for the occupancy figure the +// moment M3 can produce one. +class LMCapacityCriterion : public TileCriterion { +public: + llvm::StringRef name() const override { return "lm-capacity"; } + unsigned tier() const override { return 2; } + + bool isFeasible(const TileCandidate &, const TileContext &, + std::string &) const override { + return true; + } +}; + +// Loop overhead, analytically: the index arithmetic plus a copy in and out for +// every value the loop carries, charged once per iteration. No calibration and +// no new measurement -- and being monotone in `iterNum`, its argmin is the +// smallest surviving trip count, which is exactly what this pass picked before +// the tiers existed. That is the point: the framework lands as a refactor with +// no observable change, and criteria that *do* move decisions get added one at +// a time behind their own evidence. +class LoopOverheadCriterion : public TileCriterion { +public: + llvm::StringRef name() const override { return "loop-overhead"; } + unsigned tier() const override { return 3; } + + std::optional cost(const TileCandidate &cand, + const TileContext &ctx) const override { + return cand.iterNum * (1 + 2 * std::max(ctx.loopResults, 0)); + } +}; + +} // namespace + +TileDecider::TileDecider() { + criteria.emplace_back(std::make_unique()); + criteria.emplace_back(std::make_unique()); + criteria.emplace_back(std::make_unique()); +} + +Decision TileDecider::decide(llvm::ArrayRef candidates, + const TileContext &ctx) const { + Decision d; + if (candidates.empty()) { + d.iterNum = 1; + d.why = "no-legal-trip-count"; + return d; + } + // The largest expressible trip count is the fallback: when a tier admits + // nothing, the most tiled legal point is the closest thing to satisfying it. + int64_t maxLegal = candidates.back(); + + llvm::SmallVector live(candidates.begin(), candidates.end()); + unsigned maxTier = 0; + for (auto &c : criteria) + maxTier = std::max(maxTier, c->tier()); + + for (unsigned tier = 1; tier <= maxTier; ++tier) { + for (auto &c : criteria) { + if (c->tier() != tier) + continue; + CriterionTrace trace; + trace.name = c->name(); + trace.tier = tier; + trace.candidatesIn = live.size(); + + llvm::SmallVector kept; + std::string why; + for (int64_t k : live) { + std::string thisWhy; + if (c->isFeasible(TileCandidate{k}, ctx, thisWhy)) + kept.emplace_back(k); + else if (why.empty()) + why = thisWhy; + } + if (kept.empty()) { + // Never silently: the fallback is reported, and it is the fallback the + // caller sees in `iterNum`. + trace.candidatesOut = 0; + // Which ceiling was hit is the useful part. The vector row is a + // correctness bound the model cannot trade away; the scalar row is one + // that tiling could in principle widen. + trace.why = ctx.vecRow && maxLegal == ctx.vecRow + ? "vector-row-bound" + : (why.empty() ? "infeasible" : why); + d.why = trace.why; + d.iterNum = maxLegal; + d.perTierTrace.emplace_back(std::move(trace)); + return d; + } + live = std::move(kept); + trace.candidatesOut = live.size(); + + // Soft side: rank what survived. Costs of one tier add up, and a + // criterion with no opinion contributes nothing. + std::optional best; + int64_t bestK = live.front(); + bool anyCost = false; + for (int64_t k : live) { + auto c0 = c->cost(TileCandidate{k}, ctx); + if (!c0) + continue; + anyCost = true; + if (!best || *c0 < *best) { + best = c0; + bestK = k; + } + } + if (anyCost) { + trace.chosenCost = best; + live.assign(1, bestK); + trace.candidatesOut = 1; + } + d.perTierTrace.emplace_back(std::move(trace)); + } + } + d.iterNum = live.front(); + if (d.why.empty()) + d.why = "budget"; + return d; +} + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Analysis/VectorizabilityAnalysis.cpp b/third_party/xpu/lib/Analysis/VectorizabilityAnalysis.cpp new file mode 100644 index 0000000000..4b3331ec97 --- /dev/null +++ b/third_party/xpu/lib/Analysis/VectorizabilityAnalysis.cpp @@ -0,0 +1,843 @@ +#include "triton/Tools/Sys/GetEnv.hpp" +#include "triton/Analysis/VectorizabilityAnalysis.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "tritonxpu-vectorizability" + +namespace mlir { +namespace triton { +namespace xpu { + +bool vectorizedTyValid(Type elemTy) { + if (elemTy.isF16() || elemTy.isF32() || elemTy.isBF16() || + elemTy.isInteger(8) || elemTy.isInteger(16) || elemTy.isInteger(32)) + return true; + return false; +} + +unsigned getVectorWidth(Type elemTy) { + return 512 / elemTy.getIntOrFloatBitWidth(); +} + +bool reduceCombineIsVectorizable(triton::xpu::ReduceOp redOp) { + // The region lowering emits the combine region op by op + // (ReduceOpToLLVM::interpretCombine) instead of applying the single defining op + // of each output, so a wider set of ops has a vector form. This predicate and + // Vectorize's retyping must widen together: that TypeSwitch ends in + // llvm_unreachable, it does not fall back. + bool region = reduceCombineRegionEnabled(); + for (Block &block : redOp.getCombineOp().getBlocks()) + for (auto &op : block) { + if (region ? !isa(op) + : !isa(op)) + return false; + // A value captured from outside the region is not retyped when Vectorize + // retypes the region, so it would leave a vector op with a scalar operand. + // Constants are the exception: Vectorize rematerializes those as in-region + // splats. + for (Value operand : op.getOperands()) + if (operand.getParentBlock() != &block && + !operand.getDefiningOp()) + return false; + } + return true; +} + +bool hasVectorForm(Operation *op) { + if (!op) + return false; +#define TTX_VF_HAS_CASE(SrcType, DstType) \ + if (isa(op)) \ + return true; + TTX_SCALAR_TO_VECTOR_OPS(TTX_VF_HAS_CASE) +#undef TTX_VF_HAS_CASE + return false; +} + +bool vecReportEnabled() { + return mlir::triton::tools::getBoolEnv("TRITONXPU_VEC_REPORT"); +} + +void reportVecRoot(const char *stage, const char *site, Operation *root, + Type rootOpTy, bool eligible, int64_t closureSize, + int64_t cands) { + StringRef kernel = ""; + if (auto funcOp = root->getParentOfType()) + kernel = funcOp.getName(); + + // Element count per core rather than the tensor type: it is the quantity + // `vectorFitsRoot` divides by the vector width, so a report that printed the + // type would hide why a root was rejected. + int64_t elemsPerCore = 0; + int64_t width = 0; + if (isa(rootOpTy)) { + elemsPerCore = getTotalElemsPerThread(rootOpTy); + Type elemTy = getElementTypeOrSelf(rootOpTy); + if (!isa(elemTy) && vectorizedTyValid(elemTy)) + width = getVectorWidth(elemTy); + } + + llvm::errs() << "[VecAnalysis] " << kernel << " stage=" << stage + << " site=" << site << " root=" << root->getName() + << " elemsPerCore=" << elemsPerCore << " vecWidth=" << width + << " eligible=" << eligible << " closure=" << closureSize; + if (cands >= 0) + llvm::errs() << " cands=" << cands; + llvm::errs() << " loc=" << root->getLoc() << "\n"; +} + +Fit vectorFitUnknown(Value value, FitQuery query, unsigned wantWidth) { + return Fit::Unknown; +} + +Fit VectorizabilityAnalysis::askFit(Value value, FitQuery query, + unsigned wantWidth) { + Fit fit = fitOracle(value, query, wantWidth); + if (fit == Fit::Unknown) + fitCandidates.push_back({value, query, wantWidth}); + return fit; +} + +Operation *VectorizabilityAnalysis::getBlockArgumentOp(Value arg) { + BlockArgument blockArg = mlir::dyn_cast(arg); + Block *block = blockArg.getOwner(); + unsigned argIndex = blockArg.getArgNumber(); + + if (auto forOp = dyn_cast(block->getParentOp())) { + // TODO[dyq]: check getIterOperands -> getInitArgs + Value initValue = + forOp.getInitArgs()[argIndex - forOp.getNumInductionVars()]; + return initValue.getDefiningOp(); + } + llvm_unreachable( + "[Vectorization]: Operand is Not a BlockArgument of scf::for."); + return nullptr; +} + +bool VectorizabilityAnalysis::binLikeOpVectorize(Value lhs, Value rhs, + OperationTree &visited, + OperationTree &vectorizedOps) { + bool isFP32Ty = getElementTypeOrSelf(lhs.getType()).isF32() && + getElementTypeOrSelf(rhs.getType()).isF32(); + bool isFP16Ty = getElementTypeOrSelf(lhs.getType()).isF16() && + getElementTypeOrSelf(rhs.getType()).isF16(); + bool isINT32Ty = getElementTypeOrSelf(lhs.getType()).isInteger(32) && + getElementTypeOrSelf(rhs.getType()).isInteger(32); + bool isINT16Ty = getElementTypeOrSelf(lhs.getType()).isInteger(16) && + getElementTypeOrSelf(rhs.getType()).isInteger(16); + bool isINT8Ty = getElementTypeOrSelf(lhs.getType()).isInteger(8) && + getElementTypeOrSelf(rhs.getType()).isInteger(8); + if (!isFP32Ty && !isFP16Ty && !isINT32Ty && !isINT16Ty && !isINT8Ty) { + return false; + } + + bool isVectorized = false; + + Operation *lhsOp = lhs.getDefiningOp(); + Operation *rhsOp = rhs.getDefiningOp(); + + Operation *lhsLoopInitOp = nullptr; + Operation *rhsLoopInitOp = nullptr; + + if (mlir::isa(lhs)) { + lhsLoopInitOp = getBlockArgumentOp(lhs); + } + + if (mlir::isa(rhs)) { + rhsLoopInitOp = getBlockArgumentOp(rhs); + } + + bool lhsVectorized = lhsOp ? vectorize(lhsOp, visited, vectorizedOps) + : vectorize(lhsLoopInitOp, visited, vectorizedOps); + bool rhsVectorized = rhsOp ? vectorize(rhsOp, visited, vectorizedOps) + : vectorize(rhsLoopInitOp, visited, vectorizedOps); + + isVectorized = lhsVectorized && rhsVectorized; + return isVectorized; +} + +bool VectorizabilityAnalysis::vectorize(Operation *op, OperationTree &visited, + OperationTree &vectorizedOps) { + if (!op) { + return false; + } + visited.insert(op); + + if (vectorizedOps.contains(op)) + return true; + + bool isVectorized = false; + TypeSwitch(op) + .Case([&](auto gm2lmOp) { isVectorized = true; }) + .Case( + [&](auto gm2lmmaskOp) { isVectorized = true; }) + .Case([&](auto lm2gmOp) { isVectorized = true; }) + .Case( + [&](auto lm2gmmaskOp) { isVectorized = true; }) + .Case( + [&](auto coreIdOp) { isVectorized = true; }) + .Case( + [&](auto programIdOp) { isVectorized = true; }) + .Case([&](auto constOp) { isVectorized = true; }) + .Case([&](auto unaryOp) { isVectorized = true; }) + .Case([&](auto loadOp) { + Type elemTy = getElementTypeOrSelf(loadOp.getType()); + // Pointers and an already-vectorized element type have no bitwidth. + if (!elemTy.isIntOrFloat()) + return; + auto vectorWidth = 512 / elemTy.getIntOrFloatBitWidth(); + // This case has no E-independent half (§2.1.1): the footprint question + // is all there is, so it goes to the oracle whole. Note it does not + // consult `vectorizedTyValid` and does not factor out `rowsPerCore`, + // unlike `vectorFitsRoot` -- inconsistencies #1 and #2, kept verbatim. + isVectorized = askFit(loadOp.getResult(), FitQuery::WholeVectors, + vectorWidth) != Fit::No; + }) + .Case([&](auto storeOp) { + isVectorized = vectorize(storeOp.getValue().getDefiningOp(), visited, + vectorizedOps); + }) + .Case([&](auto reduceOp) { + if (ReduceVec) { + isVectorized = true; + + // The trailing operand is the loop index, so `size() - 1` counts the + // inputs -- unsigned, hence the guard instead of a wrap. + if (reduceOp.getOperands().size() < 2) { + isVectorized = false; + return; + } + + for (int i = 0; i < reduceOp.getOperands().size() - 1; ++i) { + auto reduceOperand = reduceOp.getOperands()[i]; + auto reduceOperandTy = reduceOperand.getType(); + + if (!reduceOperandFits(reduceOp, reduceOperandTy)) { + isVectorized = false; + } + } + + if (!reduceCombineIsVectorizable(reduceOp)) + isVectorized = false; + } else { + isVectorized = false; + } + }) + .Case([&](auto extractOp) { + isVectorized = vectorize(extractOp.getTensor().getDefiningOp(), visited, + vectorizedOps); + }) + .Case([&](auto splatOp) { + auto defineOp = splatOp.getSrc().getDefiningOp(); + if (!defineOp) { // some splatOp deal in_ptr + isVectorized = true; + } else if (!mlir::isa(splatOp.getSrc().getType())) { + // A scalar source holds one element per thread by construction + // (ClusterLayoutAttr aside, Dialect.cpp:140-141), no distribution + // involved -- that is this case's E-independent half. + isVectorized = true; + } else { // some splatOp deal tensor + isVectorized = askFit(splatOp.getSrc(), FitQuery::SingleElem, + /*wantWidth=*/1) != Fit::No; + } + }) + .Case([&](auto broadCastOp) { + // Some BroadcastOp From ReduceOp + auto srcTy = + mlir::dyn_cast(broadCastOp.getSrc().getType()); + auto resTy = + mlir::dyn_cast(broadCastOp.getResult().getType()); + + // Unranked or non-tensor operands cannot be reasoned about here; the + // conservative answer is the `isVectorized = false` this case starts + // with. + if (!srcTy || !resTy) + return; + + auto srcShape = srcTy.getShape(); + auto resShape = resTy.getShape(); + + auto rank = srcTy.getRank(); + + // The rank and the two shape relations are the E-independent half; the + // element count is the oracle's. Testing them in this order only skips + // oracle calls, both predicates are pure. + if (rank == 2) { + // srcShape[0] > 32: Scalar Calculations Perform Better than Vector + // Calculations When The Data Size is Small. (Why > 32? Which Op?) + if ((srcShape[0] == resShape[0] && srcShape[1] == 1) || + (srcShape[0] == 1 && srcShape[1] == resShape[1])) { + // 16 regardless of element type: right for f32, wrong for f16/i8 + // (should be 32/64). Inconsistency #3 of §2.1.1, kept verbatim and + // deliberately left at the call site rather than in the oracle. + isVectorized = askFit(broadCastOp.getResult(), + FitQuery::AtLeastWidth, + /*wantWidth=*/16) != Fit::No; + } + } + }) + .Case([&](auto expandDimsOp) { + isVectorized = vectorize(expandDimsOp.getOperand().getDefiningOp(), + visited, vectorizedOps); + }) + .Case([&](auto addPtrOp) { + isVectorized = vectorize(addPtrOp.getPtr().getDefiningOp(), visited, + vectorizedOps) && + vectorize(addPtrOp.getOffset().getDefiningOp(), visited, + vectorizedOps); + }) + .Case([&](auto cvtOp) { + auto cvtResTy = + mlir::dyn_cast(cvtOp.getResult().getType()); + if (!cvtResTy) + return; + auto cvtOpResEncoding = cvtResTy.getEncoding(); + if (isa(cvtOpResEncoding)) { + isVectorized = vectorize(cvtOp.getOperand().getDefiningOp(), visited, + vectorizedOps); + } + }) + .Case([&](auto selectOp) { + auto tv = selectOp.getTrueValue(); + auto fv = selectOp.getFalseValue(); + auto tType = getElementTypeOrSelf(tv.getType()); + auto fType = getElementTypeOrSelf(fv.getType()); + isVectorized = (tType == fType && (tType.isF16() || tType.isF32()) && + binLikeOpVectorize(tv, fv, visited, vectorizedOps)); + }) + .Case([&](auto cmpIOp) { + isVectorized = false; + // TODO: Add vCmpIOp Support + // auto lhs = cmpIOp.getLhs(); + // auto rhs = cmpIOp.getRhs(); + // isVectorized = binLikeOpVectorize(lhs, rhs, visited, + // vectorizedOps); + }) + .Case([&](auto cmpFOp) { + auto lhs = cmpFOp.getLhs(); + auto rhs = cmpFOp.getRhs(); + isVectorized = binLikeOpVectorize(lhs, rhs, visited, vectorizedOps); + }) + .Case([&](auto truncIOp) { + isVectorized = false; + if (auto extElemwiseOp = dyn_cast_or_null( + truncIOp.getIn().getDefiningOp())) { + isVectorized = + vectorize(extElemwiseOp.getOperands().front().getDefiningOp(), + visited, vectorizedOps); + } + }) + .Case([&](auto cmpFOp) { + auto lhs = cmpFOp.getLhs(); + auto rhs = cmpFOp.getRhs(); + isVectorized = binLikeOpVectorize(lhs, rhs, visited, vectorizedOps); + }) + .Case([&](auto ifOp) { + // For then Region + Region &thenRegion = ifOp.getThenRegion(); + // getTerminator() asserts on an empty region or a block that does not + // end in a terminator, so a malformed region answers "not vectorizable" + // rather than tripping the assert. + if (thenRegion.empty() || !thenRegion.front().mightHaveTerminator()) + return; + Block &thenBlock = thenRegion.front(); + Operation *thenTerminator = thenBlock.getTerminator(); + isVectorized = true; + if (auto yieldOp = dyn_cast(thenTerminator)) { + for (int i = 0; i < yieldOp.getOperands().size(); ++i) { + if (auto yieldDef = yieldOp.getOperands()[i].getDefiningOp()) { + isVectorized &= vectorize(yieldDef, visited, vectorizedOps); + } + } + } + + // For Else Region + if (!ifOp.getElseRegion().empty() && + ifOp.getElseRegion().front().mightHaveTerminator()) { + Region &elseRegion = ifOp.getElseRegion(); + Block &elseBlock = elseRegion.front(); + Operation *elseTerminator = elseBlock.getTerminator(); + if (auto yieldOp = dyn_cast(elseTerminator)) { + for (int i = 0; i < yieldOp.getOperands().size(); ++i) { + if (auto yieldDef = yieldOp.getOperands()[i].getDefiningOp()) { + isVectorized &= vectorize(yieldDef, visited, vectorizedOps); + } + } + } + } + }) + .Case([&](auto forOp) { + // TODO[dyq]: check getIterOperands -> getInitArgs + auto iterArgsInitValues = forOp.getInitArgs(); + Region ®ion = forOp.getRegion(); + if (region.empty() || !region.front().mightHaveTerminator()) + return; + Block &block = region.front(); + Operation *terminator = block.getTerminator(); + isVectorized = true; + if (auto yieldOp = dyn_cast(terminator)) { + for (int i = 0; i < yieldOp.getOperands().size(); ++i) { + if (auto yieldDef = yieldOp.getOperands()[i].getDefiningOp()) { + isVectorized &= vectorize(yieldDef, visited, vectorizedOps); + } + } + } + }) + .Case([&](auto yieldOp) { + isVectorized = true; + for (int i = 0; i < yieldOp.getOperands().size(); ++i) { + if (auto yieldDef = yieldOp.getOperands()[i].getDefiningOp()) { + isVectorized &= vectorize(yieldDef, visited, vectorizedOps); + } + } + }) + .Case([&](auto extElemwiseOp) { + auto symbol = extElemwiseOp.getSymbol(); + // front() below reads the first operand, so check before, not after. + if (extElemwiseOp.getOperands().empty()) { + isVectorized = false; + return; + } + auto prevOp = extElemwiseOp.getOperands().front().getDefiningOp(); + if (symbol == "_ZN3xpu5tanhfEf") { + isVectorized = true; + for (auto operand : extElemwiseOp.getOperands()) { + isVectorized = + isVectorized && vectorize(prevOp, visited, vectorizedOps); + } + } else if (symbol == "_ZN3xpu4tanfEf") { + isVectorized = true; + for (auto operand : extElemwiseOp.getOperands()) { + isVectorized = + isVectorized && vectorize(prevOp, visited, vectorizedOps); + } + } else if (symbol == "_ZN3xpu3erfEf") { + isVectorized = true; + for (auto operand : extElemwiseOp.getOperands()) { + isVectorized = + isVectorized && vectorize(prevOp, visited, vectorizedOps); + } + } else if (symbol == "_ZN3xpu5atanfEf") { + isVectorized = true; + for (auto operand : extElemwiseOp.getOperands()) { + isVectorized = + isVectorized && vectorize(prevOp, visited, vectorizedOps); + } + } else if (symbol == "_ZN3xpu5isinfEf") { + isVectorized = false; + // TODO: check visinf logic + // isVectorized = true; + // for (auto operand : extElemwiseOp.getOperands()) { + // isVectorized = + // isVectorized && vectorize(prevOp, visited, vectorizedOps); + // } + } else if (symbol == "_ZN3xpu5isnanEf") { + isVectorized = true; + for (auto operand : extElemwiseOp.getOperands()) { + isVectorized = vectorize(prevOp, visited, vectorizedOps); + } + } else if (symbol == "_ZN3xpu6rsqrtfEf") { + auto outType = + getElementTypeOrSelf(extElemwiseOp.getResult().getType()); + for (auto operand : extElemwiseOp.getOperands()) { + isVectorized = + outType.isF32() && vectorize(prevOp, visited, vectorizedOps); + } + } else { + isVectorized = false; + LLVM_DEBUG(llvm::dbgs() + << "[Vectorization]: Unsupported LibDeviceOp Symbol" + << symbol << "\n"); + } + }) + .Case([&](arith::SIToFPOp unaryOp) { + auto inType = getElementTypeOrSelf(unaryOp.getIn().getType()); + isVectorized = inType.isInteger(32) && + vectorize(unaryOp.getOperand().getDefiningOp(), visited, + vectorizedOps); + }) + .Case([&](auto binOp) { + auto lhs = binOp.getLhs(); + auto rhs = binOp.getRhs(); + isVectorized = binLikeOpVectorize(lhs, rhs, visited, vectorizedOps); + }) + .Case([&](auto binOp) { + auto lhs = binOp.getLhs(); + auto rhs = binOp.getRhs(); + isVectorized = binLikeOpVectorize(lhs, rhs, visited, vectorizedOps); + }) + .Case([&](auto unaryOp) { + isVectorized = vectorize(unaryOp.getOperand().getDefiningOp(), visited, + vectorizedOps); + }); + + if (!isVectorized) { + if (dumpFlag) { + LLVM_DEBUG({ + op->dump(); + llvm_unreachable("[Vectorization]: Unsupported Operation"); + }); + } + return false; + } + + // Dont Need To Vectorize ReduceOp's Result + if (auto reduceOp = dyn_cast(op)) + return true; + + for (Operation *user : op->getUsers()) { + if (visited.contains(user)) + continue; + + // FIXME: We've omitted the `other` value of LoadOp when create GM2LMOp + // in the past. However, `other` value comes back as we are about to + // separate GM2LMOp and LoadOp, and it will lead to a user LoadOp be in + // the vectorization path. Actions should be taken to handle this case. + // Here we workaround to skip LoadOp's `other` value. + if (auto loadOp = dyn_cast(user)) { + if (op == loadOp.getOther().getDefiningOp()) { + continue; + } + } + + if (!vectorize(user, visited, vectorizedOps)) + return false; + } + + vectorizedOps.insert(op); + return true; +} + +//===----------------------------------------------------------------------===// +// Vector-Flow analysis (step 2.1). See the header for why this is a partition +// and not a lattice climb. +//===----------------------------------------------------------------------===// + +const char *toString(VState state) { + switch (state) { + case VState::Unset: + return "Unset"; + case VState::Vector: + return "Vector"; + case VState::Scalar: + return "Scalar"; + case VState::Conflict: + return "Conflict"; + } + return "?"; +} + +bool vflowReportEnabled() { + return mlir::triton::tools::getBoolEnv("TRITONXPU_VFLOW_REPORT"); +} + +namespace { +// Unset is the identity; two different pins are a boundary, not an error. +VState joinState(VState a, VState b) { + if (a == b) + return a; + if (a == VState::Unset) + return b; + if (b == VState::Unset) + return a; + return VState::Conflict; +} + +// Values that carry data whose representation is the question. Pointers, index +// and i1 plumbing are not: nothing retypes them, so putting them in classes +// would only inflate the counts. +bool isDataValue(Value value) { + auto tensorTy = dyn_cast(value.getType()); + if (!tensorTy) + return false; + Type elemTy = tensorTy.getElementType(); + return isa(elemTy) || elemTy.isIntOrFloat(); +} +} // namespace + +unsigned VectorFlowAnalysis::idOf(Value value) { + auto it = ids.find(value); + if (it != ids.end()) + return it->second; + unsigned id = parent.size(); + ids.insert({value, id}); + parent.push_back(id); + pins.push_back(VState::Unset); + ++stats.values; + return id; +} + +unsigned VectorFlowAnalysis::find(unsigned id) { + while (parent[id] != id) { + parent[id] = parent[parent[id]]; + id = parent[id]; + } + return id; +} + +void VectorFlowAnalysis::unite(Value a, Value b) { + if (!isDataValue(a) || !isDataValue(b)) + return; + unsigned ra = find(idOf(a)); + unsigned rb = find(idOf(b)); + if (ra == rb) + return; + parent[rb] = ra; + pins[ra] = joinState(pins[ra], pins[rb]); + ++stats.unions; +} + +void VectorFlowAnalysis::pin(Value value, VState state) { + if (!isDataValue(value)) + return; + unsigned root = find(idOf(value)); + pins[root] = joinState(pins[root], state); + if (state == VState::Vector) + ++stats.vectorPins; + else if (state == VState::Scalar) + ++stats.scalarPins; +} + +VState VectorFlowAnalysis::stateOf(Value value) const { + auto it = ids.find(value); + if (it == ids.end()) + return VState::Unset; + // `find` compresses, so do the walk read-only here. + unsigned id = it->second; + while (parent[id] != id) + id = parent[id]; + return pins[id]; +} + +void VectorFlowAnalysis::visit(Operation *op) { + // Element-wise, per the single-sourced table. `select` and the two compares + // ride along: they are retyped as one unit with their operands even though the + // compare's result element type differs. + if (hasVectorForm(op) || + isa(op)) { + if (op->getNumResults() != 1) + return; + Value result = op->getResult(0); + // A select's condition is not an equality edge. Vectorize builds VSelectOp + // with `selectOp.getCondition()` passed straight through, unretyped + // (Vectorize.cpp:332-335), so the mask follows its own producer: welford + // holds both shapes in one kernel, a vcmpf-fed vselect taking + // vector<16xi1> (welford.ttxir:110) and a mask-fed one taking a plain + // tensor<1x4096xi1> (:103). Uniting it dragged welford's whole f32 + // accumulator class Scalar through that i1 mask -- the two disagreeing + // roots step 2.1 measured. CmpFOp keeps the union: its result *is* retyped + // to vector<16xi1> together with its operands (Vectorize.cpp:339-348). + auto selectOp = dyn_cast(op); + for (Value operand : op->getOperands()) { + if (selectOp && operand == selectOp.getCondition()) + continue; + unite(result, operand); + } + // Op kind is the table's business; element type is not. i1 has no vector + // form, which is why a bool chain ends up Scalar here rather than Unset. + Type elemTy = getElementTypeOrSelf(result.getType()); + if (!isa(elemTy) && !vectorizedTyValid(elemTy)) + pin(result, VState::Scalar); + return; + } + + TypeSwitch(op) + .Case([&](auto loadOp) { + // Seed, but only the E-independent half is decided here: the footprint + // question goes to the oracle, and `Unknown` seeds nothing rather than + // guessing (this analysis may end up running before CoreTiling). + Value result = loadOp.getResult(); + Type elemTy = getElementTypeOrSelf(result.getType()); + if (!elemTy.isIntOrFloat() || !vectorizedTyValid(elemTy)) { + pin(result, VState::Scalar); + return; + } + unsigned width = getVectorWidth(elemTy); + Fit fit = fitOracle(result, FitQuery::WholeVectors, width); + if (fit == Fit::Yes) + pin(result, VState::Vector); + else if (fit == Fit::No) + pin(result, VState::Scalar); + }) + .Case( + [&](auto storeOp) { pin(storeOp.getValue(), VState::Vector); }) + .Case([&](auto reduceOp) { + // The transfer function of step 2.2 (§3.1): operands may be Vector, the + // entry is a boundary, the results are Scalar. + // + // Results, unconditionally: an XPU reduce yields one element per row + // (`tensor<1xf32>`, welford.ttxir:135) under either lowering, so there + // is no vector form for it to want. Note this is a pin on the *result* + // class alone -- no edge joins it to the operands, and that absence is + // what makes the entry a boundary instead of a Conflict. + for (Value result : reduceOp.getResults()) + pin(result, VState::Scalar); + + // No pin on the data operands, in either direction. + // + // Step 2.1 pinned them Scalar whenever the combine region could not be + // retyped, so that the partition would agree with the all-or-nothing + // closure walk. That is the pin 2.2 removes: it asserts "the producer + // chain must stay scalar", where the truth is "the chain may be Vector + // and the reduce unpacks at its entry" -- the boundary this analysis + // exists to report, and the one the closure walk structurally cannot + // express. Pinning Vector instead would be just as wrong: a scalar + // island producer (i1 chain, extern_elementwise) legitimately arrives + // Scalar. So each operand keeps its own class, and the boundaries are + // counted in `run` once the partition is final. + // + // Consequence to expect in the report, not to paper over: on a reduce + // whose combine cannot be retyped, the walk says 0 and this says + // Vector. That disagreement is the modelled boundary; VFlow's per-root + // line classifies it as `kind=boundary` rather than folding it into + // agree=1. + }) + .Case([&](auto bcOp) { + // A broadcast that replicates along the innermost (vectorized) axis is + // a Scalar -> Vector boundary, not an equality edge: processOpVecTy + // retypes only the *result* (Vectorize.cpp:383-387), leaving the source + // scalar, and SVOptimization_Cond recognises exactly this shape + // (`srcShape[1] == 1`, Vectorize.cpp:693-701). Uniting here is what made + // every reduce-consuming kernel report Conflict: the Scalar-pinned + // reduce result reached the Vector-pinned elementwise chain through the + // broadcast. + auto srcTy = mlir::dyn_cast(bcOp.getSrc().getType()); + auto resTy = + mlir::dyn_cast(bcOp.getResult().getType()); + if (srcTy && resTy && !srcTy.getShape().empty() && + !resTy.getShape().empty() && srcTy.getShape().back() == 1 && + resTy.getShape().back() > 1) + return; + unite(bcOp.getResult(), bcOp.getSrc()); + }) + .Case([&](Operation *passThrough) { + // Retyped in place, so source and result share a state. + if (passThrough->getNumResults() == 1 && + passThrough->getNumOperands() >= 1) + unite(passThrough->getResult(0), passThrough->getOperand(0)); + }) + .Case([&](Operation *materialized) { + // Rebuilt as VSplatOp / VConstOp with a scalar source, so the result is + // free and no edge crosses into the source. + }) + .Case([&](auto forOp) { + // init <-> iter_arg <-> yield <-> result, per index. These four are the + // edges the closure walk cannot express, and the reason a loop-carried + // value currently vetoes on the way in. + auto yieldOp = + dyn_cast(forOp.getRegion().front().getTerminator()); + for (unsigned i = 0, e = forOp.getInitArgs().size(); i < e; ++i) { + Value iterArg = forOp.getRegionIterArgs()[i]; + unite(iterArg, forOp.getInitArgs()[i]); + unite(iterArg, forOp.getResult(i)); + if (yieldOp && i < yieldOp.getOperands().size()) + unite(iterArg, yieldOp.getOperands()[i]); + } + }) + .Case([&](auto ifOp) { + for (Region *region : {&ifOp.getThenRegion(), &ifOp.getElseRegion()}) { + if (region->empty() || !region->front().mightHaveTerminator()) + continue; + auto yieldOp = + dyn_cast(region->front().getTerminator()); + if (!yieldOp) + continue; + for (unsigned i = 0, e = std::min(yieldOp.getOperands().size(), + ifOp.getResults().size()); + i < e; ++i) + unite(ifOp.getResult(i), yieldOp.getOperands()[i]); + } + }) + .Case([&](auto extOp) { + // Pinned Scalar on purpose. The symbol -> vector-symbol table lives + // inside processOpVecTy, and one of its entries (isnan) carries + // bitwidth-specific handling that a plain pair list cannot express, so + // single-sourcing it is its own step. Until then the conservative pin is + // the only answer that cannot drift; nothing consumes this yet. + for (Value result : extOp->getResults()) { + pin(result, VState::Scalar); + ++stats.externPins; + } + }) + .Case( + [&](Operation *handledElsewhere) {}) + .Default([&](Operation *other) { + // Not modelled: pin whatever data it produces Scalar. A local scalar + // island is the conservative answer and, unlike the closure walk's + // veto, it stays local instead of propagating to the root. + for (Value result : other->getResults()) { + if (!isDataValue(result)) + continue; + pin(result, VState::Scalar); + ++stats.unknownPins; + } + }); +} + +void VectorFlowAnalysis::run(triton::FuncOp func) { + // Two passes so that the edges are all in place before the pins are counted + // per class: pinning is order-independent only once the partition is final. + func.walk([&](Operation *op) { + if (hasVectorForm(op) || + isa(op)) + visit(op); + }); + func.walk([&](Operation *op) { + if (!(hasVectorForm(op) || + isa(op))) + visit(op); + }); + + // Step 2.2: how many reduce entries actually are boundaries. Counted here + // rather than in `visit` because a class is not final while the walk runs -- + // welford's accumulator operands only become Vector after the loop edges are + // in place. `combineVec` decides whether the boundary costs a real unpack: + // the region lowering consumes the vector operands inside the combine region, + // the legacy one-op-per-output lowering cannot and needs the unpack in front. + func.walk([&](triton::xpu::ReduceOp redOp) { + if (redOp.getOperands().size() < 2) + return; + ++stats.reduceOps; + bool combineVec = reduceCombineIsVectorizable(redOp); + for (unsigned i = 0, e = redOp.getOperands().size() - 1; i < e; ++i) { + if (stateOf(redOp.getOperands()[i]) != VState::Vector) + continue; + ++stats.reduceVectorEntries; + if (!combineVec) + ++stats.reduceEntryUnpacks; + } + }); + + llvm::DenseSet roots; + for (auto &entry : ids) + roots.insert(find(entry.second)); + stats.classes = roots.size(); + for (unsigned root : roots) { + switch (pins[root]) { + case VState::Vector: + ++stats.vectorClasses; + break; + case VState::Scalar: + ++stats.scalarClasses; + break; + case VState::Conflict: + ++stats.conflictClasses; + break; + case VState::Unset: + ++stats.unsetClasses; + break; + } + } +} + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Conversion/TritonSDNNToLLVM/CMakeLists.txt b/third_party/xpu/lib/Conversion/TritonSDNNToLLVM/CMakeLists.txt index 3ec3e3212c..248d4b97a4 100644 --- a/third_party/xpu/lib/Conversion/TritonSDNNToLLVM/CMakeLists.txt +++ b/third_party/xpu/lib/Conversion/TritonSDNNToLLVM/CMakeLists.txt @@ -11,4 +11,5 @@ add_xpu_sdnn_object(TritonSDNNToLLVM TritonSDNNIR TritonSDNNTransforms MLIRAffineToStandard + MLIRTargetLLVMIRImport ) diff --git a/third_party/xpu/lib/Conversion/TritonXPUToLLVM/LoadStoreOpToLLVM.cpp b/third_party/xpu/lib/Conversion/TritonXPUToLLVM/LoadStoreOpToLLVM.cpp index d2feafa7e6..db22321dbc 100644 --- a/third_party/xpu/lib/Conversion/TritonXPUToLLVM/LoadStoreOpToLLVM.cpp +++ b/third_party/xpu/lib/Conversion/TritonXPUToLLVM/LoadStoreOpToLLVM.cpp @@ -154,6 +154,12 @@ struct LoadStoreConversionBase { } } + void createGM2SMOp(ConversionPatternRewriter &rewriter, + mlir::MLIRContext *ctx, mlir::Location &loc, Value src, + Value dst, Value offset, Value size) const { + rewriter.create(loc, src, dst, offset, size); + } + void createMemOp(ConversionPatternRewriter &rewriter, mlir::MLIRContext *ctx, mlir::Location &loc, Value bufPtr, Value gmPtr, Value offset, Value size, MemCpyType memCpyType) const { @@ -173,9 +179,19 @@ struct LoadStoreConversionBase { } void createMfenceOp(ConversionPatternRewriter &rewriter, - mlir::Location &loc) const { - // The magic number 5(101) of MfenceOp means mfencing on LM and GM - rewriter.create(loc, i32_val(5)); + mlir::Location &loc, int32_t mfenceType = 5) const { + // Mfence mask bits: bit0=LM(1), bit1=SM(2), bit2=GM(4). 5 fences LM+GM. + rewriter.create(loc, i32_val(mfenceType)); + } + + void createMfenceLMOp(ConversionPatternRewriter &rewriter, + mlir::Location &loc) const { + createMfenceOp(rewriter, loc, 1); + } + + void createMfenceGMOp(ConversionPatternRewriter &rewriter, + mlir::Location &loc) const { + createMfenceOp(rewriter, loc, 4); } Value getStartPtr(ConversionPatternRewriter &rewriter, mlir::MLIRContext *ctx, @@ -1370,7 +1386,7 @@ struct XPULoadOpConversion : public ConvertOpToLLVMPattern, {fp16LM, fp32LM, i32_val(ptrNumElems * stride)}); mlir::LLVM::XPU::createDeviceCall("_ZN3xpu10fp16tofp32EPKNS_7float16EPfi", rewriter, op, singleOperandRange, loc); - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); ptrElemScalarTy = resElemScalarTy; } // bf16Tofp32 @@ -1556,6 +1572,76 @@ struct XPULoadOpConversion : public ConvertOpToLLVMPattern, } }; +struct XPULoadScalarIndexedOpConversion + : public ConvertOpToLLVMPattern, + public LoadStoreConversionBase { + XPULoadScalarIndexedOpConversion(LLVMTypeConverter &converter, + const xpu::TargetInfo &targetInfo, + ModuleAxisInfoAnalysis &axisAnalysisPass, + PatternBenefit benefit) + : ConvertOpToLLVMPattern(converter, + benefit), + LoadStoreConversionBase(targetInfo, axisAnalysisPass) {} + + LogicalResult + matchAndRewrite(triton::xpu::LoadScalarIndexedOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto loc = op->getLoc(); + MLIRContext *ctx = rewriter.getContext(); + auto typeConverter = getTypeConverter(); + + Value res = op.getResult(); + Type resTy = res.getType(); + Type resElemTy = typeConverter->convertType(getElementTypeOrSelf(resTy)); + Type resElemScalarTy = getElementTypeOrSelf(resElemTy); + resElemScalarTy = typeConverter->convertType(resElemScalarTy); + + unsigned addrSpace = 0; + Type ptrTy = op.getPtr().getType(); + if (auto ptrTensorTy = mlir::dyn_cast(ptrTy)) { + if (auto pt = + mlir::dyn_cast(ptrTensorTy.getElementType())) + addrSpace = pt.getAddressSpace(); + } else if (auto pt = mlir::dyn_cast(ptrTy)) { + addrSpace = pt.getAddressSpace(); + } + + auto llPtrs = unpackLLElements(loc, adaptor.getPtr(), rewriter); + Value basePtr = bitcast(llPtrs[0], ptr_ty(ctx, addrSpace)); + Value llIndex = adaptor.getIndex(); + Value elemPtr = + gep(ptr_ty(ctx, addrSpace), resElemScalarTy, basePtr, llIndex); + Value loaded = load(resElemScalarTy, elemPtr); + + unsigned resNumElems = getTotalElemsPerThread(resTy); + bool isVectorized = mlir::isa(resElemTy); + unsigned vecSize = 1u; + if (isVectorized) { + unsigned elemNbits = 0u; + getVectorInfo(resElemTy, vecSize, elemNbits); + } + + SmallVector loadedVals; + for (size_t elemIdx = 0; elemIdx < resNumElems; ++elemIdx) { + if (isVectorized) { + Value newVector = rewriter.create(loc, resElemTy); + for (size_t i = 0; i < vecSize; ++i) { + newVector = insert_element(resElemTy, newVector, loaded, i32_val(i)); + } + loadedVals.push_back(newVector); + } else { + loadedVals.push_back(loaded); + } + } + + Type llvmResultStructTy = typeConverter->convertType(resTy); + Value resultStruct = packLLElements(loc, typeConverter, loadedVals, + rewriter, llvmResultStructTy); + rewriter.replaceOp(op, {resultStruct}); + return success(); + } +}; + struct XPUStoreOpConversion : public ConvertOpToLLVMPattern, public LoadStoreConversionBase { @@ -1935,7 +2021,7 @@ struct XPUStoreOpConversion } } - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); // fp32 to fp16 if (fp32Tofp16) { @@ -1944,7 +2030,7 @@ struct XPUStoreOpConversion ValueRange singleOperandRange({fp32LM, fp16LM, i32_val(ptrNumElems)}); mlir::LLVM::XPU::createDeviceCall("_ZN3xpu10fp32tofp16Ef", rewriter, op, singleOperandRange, loc); - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); } rewriter.eraseOp(op); @@ -2170,7 +2256,7 @@ struct XPUGM2LMOpConversion createGM2LMOp(rewriter, ctx, loc, srcPtr, dstPtr, offsetBytes, readBytes); if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); } resultStruct = packLLElements(loc, typeConverter, llLMPtrs, rewriter, @@ -2196,6 +2282,30 @@ struct XPUGM2LMOpConversion SmallVector newLmBufPtrs(llLMPtrs.size(), llLMPtrs[0]); resultStruct = packLLElements(loc, typeConverter, newLmBufPtrs, rewriter, llLMPtr.getType()); + } else if (offsetState == OffsetState::LocallyScalar) { + int64_t rowLen = op.getRowLen(); + if (rowLen <= 0) + rowLen = static_cast(llLMPtrs.size()); + readBytes = elemBytes; + SmallVector newLmBufPtrs(llLMPtrs.size()); + for (size_t start = 0; start < llLMPtrs.size(); + start += static_cast(rowLen)) { + Value dstPtr = bitcast(llLMPtrs[start], ptr_ty(ctx, 0)); + Value srcPtr = bitcast(llGMPtrs[start], ptr_ty(ctx, 1)); + createGM2LMOp(rewriter, ctx, loc, srcPtr, dstPtr, offsetBytes, + readBytes); + size_t end = start + static_cast(rowLen); + if (end > llLMPtrs.size()) + end = llLMPtrs.size(); + for (size_t j = start; j < end; ++j) + newLmBufPtrs[j] = llLMPtrs[start]; + } + if (!async) + createMfenceLMOp(rewriter, loc); + resultStruct = packLLElements(loc, typeConverter, newLmBufPtrs, rewriter, + llvmResultStructTy); + rewriter.replaceOp(op, {resultStruct}); + return success(); } else if (offsetState == OffsetState::LocallyContinuous) { int64_t _rowLen = op.getRowLen(); int64_t _rowStride = op.getRowStride(); @@ -2230,7 +2340,7 @@ struct XPUGM2LMOpConversion } if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); resultStruct = packLLElements(loc, typeConverter, llLMPtrs, rewriter, llvmResultStructTy); @@ -2260,7 +2370,7 @@ struct XPUGM2LMOpConversion } if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); rewriter.replaceOp(op, {resultStruct}); return success(); @@ -2270,13 +2380,88 @@ struct XPUGM2LMOpConversion Value srcPtr = bitcast(llGMPtrs[0], ptr_ty(ctx, 1)); createGM2LMOp(rewriter, ctx, loc, srcPtr, dstPtr, offsetBytes, readBytes); if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); rewriter.replaceOp(op, {resultStruct}); return success(); } }; +struct XPUStageSMOpConversion + : public ConvertOpToLLVMPattern, + public LoadStoreConversionBase { + XPUStageSMOpConversion(LLVMTypeConverter &converter, + const xpu::TargetInfo &targetInfo, + ModuleAxisInfoAnalysis &axisAnalysisPass, + PatternBenefit benefit) + : ConvertOpToLLVMPattern(converter, benefit), + LoadStoreConversionBase(targetInfo, axisAnalysisPass) {} + + Value getGlobalSmemBase(Location loc, ConversionPatternRewriter &rewriter, + Operation *op) const { + ModuleOp mod = op->getParentOfType(); + LLVM::GlobalOp globalSmem; + mod.walk([&](LLVM::GlobalOp g) { + if (g.getSymName() == "global_smem") + globalSmem = g; + }); + assert(globalSmem && "global_smem not found; initSharedMemory must run " + "before StageSM lowering"); + Value addr = rewriter.create(loc, globalSmem); + return rewriter.create( + loc, LLVM::LLVMPointerType::get(rewriter.getContext(), 2), addr); + } + + LogicalResult + matchAndRewrite(triton::xpu::StageSMOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto loc = op->getLoc(); + MLIRContext *ctx = rewriter.getContext(); + auto typeConverter = getTypeConverter(); + + Value ptr = op.getPtr(); + Type ptrTy = ptr.getType(); + Type elemTy = mlir::cast(ptrTy).getPointeeType(); + elemTy = typeConverter->convertType(elemTy); + unsigned elemNbits = isa(elemTy) + ? 64u + : elemTy.getIntOrFloatBitWidth(); + + Value srcPtr = bitcast(adaptor.getPtr(), ptr_ty(ctx, 1)); + Value smBase = getGlobalSmemBase(loc, rewriter, op); + Value smDst = gep(ptr_ty(ctx, 2), i8_ty, smBase, adaptor.getSmOffset()); + + Value elemBytes = i32_val(elemNbits / 8u); + Value bufElems = adaptor.getBufElems(); + Value readLen = bufElems; + if (op.getLen()) { + Value reqLen = smax(adaptor.getLen(), i32_val(0)); + readLen = smin(reqLen, bufElems); + } + Value readBytes = mul(readLen, elemBytes); + + Value coreId = mlir::LLVM::XPU::getThreadId(rewriter, loc); + Value isCore0 = icmp_eq(coreId, i32_val(0)); + + Block *currentBlock = rewriter.getInsertionBlock(); + Block *afterBlock = + rewriter.splitBlock(currentBlock, rewriter.getInsertionPoint()); + Block *dmaBlock = rewriter.createBlock(afterBlock); + + rewriter.setInsertionPointToEnd(currentBlock); + rewriter.create(loc, isCore0, dmaBlock, afterBlock); + + rewriter.setInsertionPointToStart(dmaBlock); + createGM2SMOp(rewriter, ctx, loc, srcPtr, smDst, i32_val(0), readBytes); + rewriter.create(loc, afterBlock); + + rewriter.setInsertionPointToStart(afterBlock); + xpu_barrier(); + rewriter.replaceOp(op, {smDst}); + return success(); + } +}; + struct XPULM2GMOpConversion : public ConvertOpToLLVMPattern, public LoadStoreConversionBase { @@ -2425,7 +2610,7 @@ struct XPULM2GMOpConversion } } if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); rewriter.eraseOp(op); rewriter.create(loc, ValueRange{}, newBlock); return success(); @@ -2479,7 +2664,7 @@ struct XPULM2GMOpConversion break; } if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); rewriter.eraseOp(op); return success(); @@ -2618,7 +2803,7 @@ struct XPUGM2LMMaskOpConversion createGM2LMOp(rewriter, ctx, loc, srcPtr, dstPtr, offsetBytes, readBytes); if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); } else { // Unknown for (size_t i = 0; i < llGMPtrs.size(); ++i) { @@ -2636,7 +2821,7 @@ struct XPUGM2LMMaskOpConversion createGM2LMOp(rewriter, ctx, loc, srcPtr, dstPtr, offsetBytes, _readBytes); if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); } } resultStruct = packLLElements(loc, typeConverter, llLMPtrs, rewriter, @@ -2655,7 +2840,7 @@ struct XPUGM2LMMaskOpConversion readBytes = mask ? select(llMasks[0], readBytes, i32_val(0)) : readBytes; createGM2LMOp(rewriter, ctx, loc, srcPtr, dstPtr, offsetBytes, readBytes); if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); resultStruct = packLLElements(loc, typeConverter, newLmBufPtrs, rewriter, llvmResultStructTy); @@ -2664,7 +2849,7 @@ struct XPUGM2LMMaskOpConversion readBytes = mask ? select(llMasks[0], readBytes, i32_val(0)) : readBytes; createGM2LMOp(rewriter, ctx, loc, srcPtr, dstPtr, offsetBytes, readBytes); if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); SmallVector newLmBufPtrs(llLMPtrs.size(), llLMPtrs[0]); resultStruct = packLLElements(loc, typeConverter, newLmBufPtrs, rewriter, @@ -2695,7 +2880,7 @@ struct XPUGM2LMMaskOpConversion } } if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); resultStruct = packLLElements(loc, typeConverter, llLMPtrs, rewriter, llvmResultStructTy); rewriter.create(loc, ValueRange{}, newBlock); @@ -2737,7 +2922,7 @@ struct XPUGM2LMMaskOpConversion } if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); rewriter.replaceOp(op, {resultStruct}); return success(); @@ -2900,12 +3085,12 @@ struct XPULM2GMMaskOpConversion oldBlock, newBlock); } if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); rewriter.eraseOp(op); rewriter.create(loc, ValueRange{}, newBlock); return success(); } - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); rewriter.create(loc, ValueRange{}, newBlock); break; } @@ -2935,7 +3120,7 @@ struct XPULM2GMMaskOpConversion readBytes = mask ? select(_mask, elemBytes, i32_val(0)) : elemBytes; createLM2GMOp(rewriter, ctx, loc, srcPtr, dstPtr, offsetBytes, readBytes); - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); } break; } @@ -2945,7 +3130,7 @@ struct XPULM2GMMaskOpConversion } } if (!async) - createMfenceOp(rewriter, loc); + createMfenceLMOp(rewriter, loc); rewriter.eraseOp(op); return success(); @@ -3035,6 +3220,7 @@ void mlir::triton::xpu::populateLoadStoreOpToLLVMPatterns( RewritePatternSet &patterns, ModuleAxisInfoAnalysis &axisInfoAnalysis, PatternBenefit benefit) { patterns.add(typeConverter, targetInfo, diff --git a/third_party/xpu/lib/Conversion/TritonXPUToLLVM/ScanOpToLLVM.cpp b/third_party/xpu/lib/Conversion/TritonXPUToLLVM/ScanOpToLLVM.cpp index 02c05536d2..0370442292 100644 --- a/third_party/xpu/lib/Conversion/TritonXPUToLLVM/ScanOpToLLVM.cpp +++ b/third_party/xpu/lib/Conversion/TritonXPUToLLVM/ScanOpToLLVM.cpp @@ -1,8 +1,8 @@ -#include "PatternTritonXPUOpToLLVM.h" #include "Utility.h" #include "triton/Conversion/TritonGPUToLLVM/Utility.h" -#include "triton/Conversion/TritonXPUToLLVM/LegacyLLVMHelpers.h" // LLVM22 dragon-style macros for XPU only #include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "PatternTritonXPUOpToLLVM.h" +#include "triton/Conversion/TritonXPUToLLVM/LegacyLLVMHelpers.h" // LLVM22 dragon-style macros for XPU only using ::mlir::triton::gpu::getTotalElemsPerThread; @@ -136,7 +136,10 @@ static void applyLastElemScanOnSM(SmallVector> &srcValues, SmallVector smemBases, SmallVector accSmemBases, SmallVector smemTypes, Value groupId, - Value laneId, triton::xpu::ScanOp op) { + Value laneId, triton::xpu::ScanOp op, + ArrayRef carryIn = {}, + ArrayRef carryBases = {}, + Value isFirst = {}) { Location loc = helper.getXPULoc(); // Scan SM Last Elem Per Thread @@ -193,8 +196,19 @@ static void applyLastElemScanOnSM(SmallVector> &srcValues, for (unsigned srcIdx = 0; srcIdx < groupSizeInt; ++srcIdx) { if (srcIdx == 0) { // core0 dont't need to acc + // Base (no carry) exclusive prefix for this iteration. accs[srcIdx] = accumulate(helper, rewriter, accs[srcIdx], readValues[srcIdx]); + // When carrying across grid-stride iterations, fold the incoming carry + // into core0's prefix on every iteration except the first. Selecting on + // the *result* (rather than seeding with a hardcoded identity) keeps the + // first iteration exact for any combine op (min/max/prod/...). + if (!carryIn.empty()) { + SmallVector carry(carryIn.begin(), carryIn.end()); + auto withCarry = accumulate(helper, rewriter, carry, readValues[srcIdx]); + for (unsigned k = 0; k < accs[srcIdx].size(); ++k) + accs[srcIdx][k] = select(isFirst, accs[srcIdx][k], withCarry[k]); + } // ValueRange operandValueRange({accs[srcIdx][0], i32_val(srcIdx)}); // mlir::LLVM::XPU::createDeviceCall(calFunc, rewriter, op, @@ -227,10 +241,19 @@ static void applyLastElemScanOnSM(SmallVector> &srcValues, } } + // Persist this iteration's grand total (carryIn + sum of all core partials) + // to the carry slot so the next grid-stride iteration continues from it. + // carryOut = accumulate(accs[last], readValues[last]). + if (!carryBases.empty()) { + auto carryOut = accumulate(helper, rewriter, accs[groupSizeInt - 1], + readValues[groupSizeInt - 1]); + for (unsigned k = 0; k < carryBases.size(); ++k) + store_sm(carryOut[k], carryBases[k]); + } + // [dump] dump accs result if (/*dump_Out_SMValue*/ false) { auto operandIdx = 1; // 第一个输入 - for (int core_id_offset = 0; core_id_offset < 64; ++core_id_offset) { auto elemTy = accs[core_id_offset][operandIdx].getType(); @@ -305,7 +328,8 @@ static void applyCoreElemScanWithSMAcc( ConversionPatternRewriter &rewriter, const TargetInfoBase &targetInfo, ScanLoweringHelper &helper, SmallVector smemBases, SmallVector accSmemBases, SmallVector smemTypes, Value groupId, - Value laneId, triton::xpu::ScanOp op) { + Value laneId, triton::xpu::ScanOp op, ArrayRef carryIn = {}, + Value isFirst = {}) { Location loc = helper.getXPULoc(); Value threadId = mlir::LLVM::XPU::getThreadId(rewriter, loc); @@ -406,10 +430,23 @@ static void applyCoreElemScanWithSMAcc( } for (unsigned srcIndex = 0; srcIndex < srcValues.size(); srcIndex++) { + // core0 keeps its local prefix; when carrying across grid-stride + // iterations it must add the incoming carry (except on the first + // iteration). The combine region consumes all operands at once, so fold + // the carry over the full operand vector in a single accumulate rather + // than per operand (per-operand folding passes the wrong number of block + // arguments to the combine region and asserts for multi-operand scans). + SmallVector core0Vals = srcValues[srcIndex]; + if (!carryIn.empty()) { + SmallVector carry(carryIn.begin(), carryIn.end()); + auto withCarry = accumulate(helper, rewriter, carry, srcValues[srcIndex]); + for (unsigned k = 0; k < core0Vals.size(); ++k) + core0Vals[k] = select(isFirst, core0Vals[k], withCarry[k]); + } for (unsigned operandIdx = 0; operandIdx < (helper.getXPUNumOperands() - 1); ++operandIdx) { // skip loopIndex srcValues[srcIndex][operandIdx] = - select(coreIdInGroupZero, srcValues[srcIndex][operandIdx], + select(coreIdInGroupZero, core0Vals[operandIdx], newSrcValues[srcIndex][operandIdx]); } } @@ -434,10 +471,28 @@ struct XPUScanOpConversion return getTypeConverter()->convertType(ty); } + // Carry across grid-stride loop iterations is only meaningful when the loop + // tiles the scan axis itself (e.g. a 1-D scan split into tiles, as in the + // unique kernel's global cumsum). It must NOT fire when the loop iterates + // over independent scans (e.g. the per-row scan of a 2-D tensor along the + // last axis, where each row is a self-contained scan). Mirror + // ReduceOp::isNeedLoopCacheResult so scan and reduce agree. + bool isNeedLoopCarry(triton::xpu::ScanOp op) const { + if (!op.getLoopIndex()) + return false; + if (op.getAxis() == 1) + return false; + if (auto resultTy = + dyn_cast(op.getResult()[0].getType())) + return resultTy.getShape().size() == 1; + return true; + } + // Helper to compute the smem bases in both reductions and scans std::pair, SmallVector> getSmemBases(triton::xpu::ScanOp op, unsigned elems, - ConversionPatternRewriter &rewriter) const { + ConversionPatternRewriter &rewriter, + SmallVector *carryBases = nullptr) const { ScanLoweringHelper helper(op); SmallVector offsets; @@ -508,6 +563,30 @@ struct XPUScanOpConversion accCacheSmemBases[i] = indexToBaseForAccCache[i]; } + // Persistent carry region (1 slot per operand) used to accumulate a scan + // across grid-stride loop iterations. It lives just past the accCache + // region and is recorded in `offsets` so later scans/reduces don't reuse + // it. Only populated when the caller asks (scan inside a loop). + if (carryBases) { + carryBases->resize(op.getNumOperands() - 1); + std::map indexToBaseForCarry; + indexToBaseForCarry[0] = + gep(ptr_ty(rewriter.getContext(), 2), + getElementType(op, op.getNumOperands() - 2), + accCacheSmemBases.back(), i32_val(elems)); + offsets.push_back( + (getElementType(op, 0).getIntOrFloatBitWidth()) / 8); + for (unsigned i = 1; i < (op.getNumOperands() - 1); ++i) { + indexToBaseForCarry[i] = + gep(ptr_ty(rewriter.getContext(), 2), getElementType(op, i), + indexToBaseForCarry[i - 1], i32_val(1)); + offsets.push_back( + (getElementType(op, i).getIntOrFloatBitWidth()) / 8); + } + for (unsigned i = 0; i < (op.getNumOperands() - 1); ++i) + (*carryBases)[i] = indexToBaseForCarry[i]; + } + helper.setSMOffsets(helper.getScanId(), offsets); return {smemBases, accCacheSmemBases}; } @@ -614,7 +693,14 @@ struct XPUScanOpConversion // helper.dumpSMOffsets(); auto elems = helper.getScratchSizeInElemsXPU(); - auto [smemBases, accSmemBases] = getSmemBases(op, elems, rewriter); + // A scan whose axis is tiled across a grid-stride loop carries a running + // total across iterations via a persistent SM slot. Only allocate/use + // the carry machinery in that case (see isNeedLoopCarry); a per-row scan + // of a 2-D tensor must not carry between rows. + bool hasCarry = isNeedLoopCarry(op); + SmallVector carryBases; + auto [smemBases, accSmemBases] = + getSmemBases(op, elems, rewriter, hasCarry ? &carryBases : nullptr); // llvm::errs() << "\n [After getSmemBases]:\n" // << op->getParentOfType() << "\n"; @@ -632,6 +718,26 @@ struct XPUScanOpConversion storeGroupAccumulator(srcValues, rewriter, helper, laneId, groupId, smemBases, smemTypes); + // Read the carry produced by the previous grid-stride iteration. It is + // combined into core0's prefix by applyLastElemScanOnSM / + // applyCoreElemScanWithSMAcc, which select it away on the first + // iteration (isFirst) instead of relying on a hardcoded identity. Read + // it BEFORE applyLastElemScanOnSM overwrites the carry slot with this + // iteration's total. + SmallVector carryIn; + Value isFirst; + if (hasCarry) { + Value loopIdx = adaptor.getLoopIndex(); + Value zeroIdx = rewriter.create( + loc, loopIdx.getType(), + rewriter.getIntegerAttr(loopIdx.getType(), 0)); + isFirst = icmp_eq(loopIdx, zeroIdx); + for (unsigned k = 0; k < (op.getNumOperands() - 1); ++k) { + Type ety = smemTypes[k]; + carryIn.push_back(load_sm(ety, carryBases[k])); + } + } + // [dump][sm] operand-0: sm[0]-sm[63] or operand-1 sm[64]-sm[127] if (/*dump_In_SMValue*/ false) { auto operandIdx = 1; // 第一个输入 @@ -651,12 +757,9 @@ struct XPUScanOpConversion calFunc = "_ZN3xpu15printInt64_specEliii"; } - SmallVector operandValues; - operandValues.append({loadVal, /*cluster_id*/ i32_val(0), - /*core_id*/ i32_val(0), - /*custom_id*/ - i32_val(500 + core_id_offset)}); - ValueRange operandValueRange(operandValues); + ValueRange operandValueRange( + {loadVal, /*cluster_id*/ i32_val(0), /*core_id*/ i32_val(0), + /*custom_id*/ i32_val(500 + core_id_offset)}); mlir::LLVM::XPU::createDeviceCall(calFunc, rewriter, op, operandValueRange, loc); } @@ -672,7 +775,8 @@ struct XPUScanOpConversion // Read back the partial reduction of each warp and accumulate them // based on warpId. applyLastElemScanOnSM(srcValues, rewriter, targetInfo, helper, smemBases, - accSmemBases, smemTypes, groupId, laneId, op); + accSmemBases, smemTypes, groupId, laneId, op, + carryIn, carryBases, isFirst); // llvm::errs() << "\n After applyLastElemScanOnSM:\n" // << op->getParentOfType() << "\n"; @@ -699,12 +803,9 @@ struct XPUScanOpConversion calFunc = "_ZN3xpu15printInt64_specEliii"; } - SmallVector operandValues; - operandValues.append({loadVal, /*cluster_id*/ i32_val(0), - /*core_id*/ i32_val(0), - /*custom_id*/ - i32_val(600 + core_id_offset)}); - ValueRange operandValueRange(operandValues); + ValueRange operandValueRange( + {loadVal, /*cluster_id*/ i32_val(0), /*core_id*/ i32_val(0), + /*custom_id*/ i32_val(600 + core_id_offset)}); mlir::LLVM::XPU::createDeviceCall(calFunc, rewriter, op, operandValueRange, loc); } @@ -717,7 +818,7 @@ struct XPUScanOpConversion // adding the accumulated value from the previous lane. applyCoreElemScanWithSMAcc(srcValues, rewriter, targetInfo, helper, smemBases, accSmemBases, smemTypes, groupId, - laneId, op); + laneId, op, carryIn, isFirst); // llvm::errs() << "\n After applyCoreElemScanWithSMAcc:\n" // << op->getParentOfType() << "\n"; @@ -742,12 +843,9 @@ struct XPUScanOpConversion calFunc = "_ZN3xpu15printInt64_specEliii"; } - SmallVector operandValues; - operandValues.append({loadVal, /*cluster_id*/ i32_val(0), - /*core_id*/ i32_val(0), - /*custom_id*/ - i32_val(700 + core_id_offset)}); - ValueRange operandValueRange(operandValues); + ValueRange operandValueRange( + {loadVal, /*cluster_id*/ i32_val(0), /*core_id*/ i32_val(0), + /*custom_id*/ i32_val(700 + core_id_offset)}); mlir::LLVM::XPU::createDeviceCall(calFunc, rewriter, op, operandValueRange, loc); } diff --git a/third_party/xpu/lib/Dialect/TritonSDNN/Transforms/CMakeLists.txt b/third_party/xpu/lib/Dialect/TritonSDNN/Transforms/CMakeLists.txt index e367f25318..9387af5752 100644 --- a/third_party/xpu/lib/Dialect/TritonSDNN/Transforms/CMakeLists.txt +++ b/third_party/xpu/lib/Dialect/TritonSDNN/Transforms/CMakeLists.txt @@ -2,19 +2,37 @@ add_xpu_sdnn_object(TritonSDNNTransforms ${CMAKE_CURRENT_SOURCE_DIR}/TritonSDNNConversion.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/Bufferize.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/ConvertType.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/DSACopy.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/EWActTable.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/Legalize.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/MergeExternEW.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/Combine.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/CombinePatterns.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/CombineBefore.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/ResolveLayoutConflict.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/HoistDS.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/TransposeMMA.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/TransposeEW.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/LoopGrid.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/LoopUnroll.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/CoprocessorInterleaving.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/Pipeline.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/EWLiteScheduling.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/MultipleBuffer.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/HoistLoopInvariantDma.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/FuseReluActivation.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/EliminateMmaAccZero.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/FuseMmaVectorBias.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/OptimizeRcLayout.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/LowerRcSubview.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/KLoopAccModification.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/MaterializeDeferredRaw.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/Utility.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/SchedulePrintOp.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/RemoveDsOp.cpp.o ${CMAKE_CURRENT_SOURCE_DIR}/StripAllOp.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/InjectDMAFormat.cpp.o + ${CMAKE_CURRENT_SOURCE_DIR}/MxScaleLayout.cpp.o LINK_LIBS TritonIR diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/AsyncLoadSchedule.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/AsyncLoadSchedule.cpp new file mode 100644 index 0000000000..97773c35fb --- /dev/null +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/AsyncLoadSchedule.cpp @@ -0,0 +1,566 @@ +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Dominance.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/LLVMXPU/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/Transforms/Passes.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "tritonxpu-async-load-schedule" + +using namespace mlir; +using namespace mlir::triton; + +namespace mlir { +namespace triton { +namespace xpu { + +#define GEN_PASS_DEF_TRITONXPUASYNCLOADSCHEDULE +#include "triton/Dialect/TritonXPU/Transforms/Passes.h.inc" + +namespace { + +struct ForwardingChain { + SmallVector ops; + Operation *firstRealUser = nullptr; +}; + +class TritonXPUAsyncLoadSchedulePass + : public impl::TritonXPUAsyncLoadScheduleBase< + TritonXPUAsyncLoadSchedulePass> { +public: + using impl::TritonXPUAsyncLoadScheduleBase< + TritonXPUAsyncLoadSchedulePass>::TritonXPUAsyncLoadScheduleBase; + + TritonXPUAsyncLoadSchedulePass() = default; + TritonXPUAsyncLoadSchedulePass(bool dumpFlag) { this->dumpFlag = dumpFlag; } + + void runOnOperation() override { + ModuleOp mod = getOperation(); + + // Phase 0: Hoist gm2lm ops upward past non-dependent ops to enable + // earlier DMA issue and better overlap. + SmallVector gm2lmOps; + mod.walk([&](Operation *op) { + if (isa(op)) + gm2lmOps.push_back(op); + }); + for (auto *op : gm2lmOps) + tryHoistGM2LM(op); + + SmallVector loadOps; + mod.walk([&](triton::xpu::LoadOp loadOp) { loadOps.push_back(loadOp); }); + + llvm::DenseMap insertedMfences; + for (auto loadOp : loadOps) + trySchedule(loadOp, insertedMfences); + eraseDominatedMfences(insertedMfences); + } + +private: + // Check if op depends on (uses a result of) target. + bool dependsOn(Operation *op, Operation *target) const { + for (Value operand : op->getOperands()) { + if (Operation *def = operand.getDefiningOp()) { + if (def == target) + return true; + } + } + return false; + } + + // Check if it's safe to hoist gm2lm above op. + bool canHoistAcross(Operation *gm2lmOp, Operation *op) const { + // gm2lm cannot move above its own operand definitions. + for (Value operand : gm2lmOp->getOperands()) { + if (Operation *def = operand.getDefiningOp()) { + if (def == op) + return false; + } + } + // Don't cross terminators. + if (op->hasTrait()) + return false; + // Don't cross other gm2lm/store/lm2gm (memory ops). + if (isa(op)) + return false; + // Don't cross mfence. + if (isa(op)) + return false; + // Don't cross load that aliases with the gm2lm's LM buffer. + if (auto loadOp = dyn_cast(op)) { + Value gm2lmResult = gm2lmOp->getResult(0); + return !mayAliasLM(gm2lmResult, loadOp.getPtr()); + } + // Region ops (reduce/scan/if...) are safe to cross *iff* they have no + // memory effects — a pure reduction touches only SSA values, never the LM + // buffer our DMA lands in, so issuing the DMA before it just overlaps the + // transfer with the reduction. Ops the gm2lm actually depends on are held + // back separately (see collectLocalDeps `barriers`), so this only lets us + // cross *unrelated* region ops. Impure ops (side effects) stay barriers. + return isMemoryEffectFree(op); + } + + // Collect the set of same-block ops that must move together with `gm2lmOp` + // when it is hoisted to just before `insertBefore`: the transitive + // operand-closure ops that are defined at or after insertBefore and can be + // relocated (memory-effect-free, region-free address computation). Operands + // defined strictly before insertBefore already dominate and stay put. + // + // Any dependency that *cannot* move — an op with memory effects (e.g. a + // gather index load feeding the address) or a region op (reduce/scan) whose + // result the address depends on — is recorded in `barriers` instead. The + // gm2lm must never be hoisted above a barrier; the caller uses them as hard + // lower bounds for the insertion point (so we still hoist as far up as the + // nearest barrier rather than giving up entirely). `deps` is returned sorted + // in block order. + void collectLocalDeps(Operation *gm2lmOp, Operation *insertBefore, + SmallVectorImpl &deps, + DenseSet &barriers) const { + Block *block = gm2lmOp->getBlock(); + SmallVector worklist; + DenseSet visited; + worklist.push_back(gm2lmOp); + visited.insert(gm2lmOp); + while (!worklist.empty()) { + Operation *cur = worklist.pop_back_val(); + for (Value operand : cur->getOperands()) { + Operation *def = operand.getDefiningOp(); + // Block args / values from enclosing regions always dominate. + if (!def || def->getBlock() != block) + continue; + // Already dominates the insertion point: no need to move it. + if (def->isBeforeInBlock(insertBefore)) + continue; + if (!visited.insert(def).second) + continue; + // This def sits in (insertBefore, gm2lmOp]. Pure, region-free ops are + // address computation we carry along; anything else (memory effects or + // a nested region) cannot be relocated and becomes a hard lower bound. + if (isMemoryEffectFree(def) && def->getNumRegions() == 0) { + deps.push_back(def); + worklist.push_back(def); + } else { + barriers.insert(def); + } + } + } + + // Sort in original block order so we move them in topological order. + llvm::sort(deps, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + } + + // Verify that relocating every op in `moveSet` to just before `insertBefore` + // preserves SSA dominance. For each moving op, every same-block operand + // definition must either be part of the move set (relative order is kept) or + // already sit before insertBefore (so it still dominates after the move). + // Operands defined outside the block (block args / enclosing-region values) + // always dominate and are safe. + bool hoistPreservesDominance(ArrayRef moveSet, + Operation *insertBefore) const { + DenseSet moving(moveSet.begin(), moveSet.end()); + // insertBefore is the fixed anchor the whole move set is relocated in + // front of. If it is itself part of the move set, relocating the other + // ops "before insertBefore" reorders them around their own dependency + // (insertBefore) and produces a use-before-def. This happens when the + // backward scan stops on a dependency op (e.g. the topmost address + // computation, right below a region op we cannot cross). Bail out — the + // load simply stays synchronous. + if (moving.contains(insertBefore)) + return false; + Block *block = insertBefore->getBlock(); + for (Operation *op : moveSet) { + for (Value operand : op->getOperands()) { + Operation *def = operand.getDefiningOp(); + if (!def || def->getBlock() != block) + continue; + if (moving.contains(def)) + continue; + if (!def->isBeforeInBlock(insertBefore)) + return false; + } + } + return true; + } + + // Try to hoist a gm2lm op upward past non-dependent ops in the same block, + // bringing its local dependencies along. + void tryHoistGM2LM(Operation *gm2lmOp) const { + Block *block = gm2lmOp->getBlock(); + if (!block) + return; + + // First, collect the full set of local deps (address computation) that + // would need to move with gm2lm, plus the hard lower-bound barriers it + // depends on (impure/region ops that cannot be relocated). Use block begin + // as initial bound to discover everything. + SmallVector allDeps; + DenseSet barriers; + collectLocalDeps(gm2lmOp, &block->front(), allDeps, barriers); + + DenseSet depSet(allDeps.begin(), allDeps.end()); + depSet.insert(gm2lmOp); + + // Scan backwards from gm2lm to find the earliest safe insertion point, + // skipping over ops that are part of the dep set (they'll move too) and + // stopping at barriers (dependencies we cannot move above) or any op we + // cannot legally cross. + Operation *insertBefore = gm2lmOp; + for (Operation *op = gm2lmOp->getPrevNode(); op; op = op->getPrevNode()) { + if (depSet.contains(op)) { + insertBefore = op; + continue; + } + if (barriers.contains(op)) + break; + if (!canHoistAcross(gm2lmOp, op)) + break; + insertBefore = op; + } + if (insertBefore == gm2lmOp) + return; + // Check that insertBefore is actually before the earliest dep. If the + // backward scan stopped right on top of the earliest dependency, there is + // no non-dep op to anchor in front of — hoisting would place ops before + // their own operands. + if (!allDeps.empty() && + (insertBefore == allDeps.front() || depSet.contains(insertBefore))) + return; + + // Re-collect deps constrained to the actual insertion point. + SmallVector deps; + DenseSet unusedBarriers; + collectLocalDeps(gm2lmOp, insertBefore, deps, unusedBarriers); + + // Final dominance-safety gate: verify that moving the whole set to just + // before insertBefore keeps every operand definition dominating its uses. + // collectLocalDeps already guarantees the set is operand-closed, but this + // is a cheap, explicit correctness backstop against invalid IR + // ("operand does not dominate this use"). + SmallVector moveSet(deps.begin(), deps.end()); + moveSet.push_back(gm2lmOp); + if (!hoistPreservesDominance(moveSet, insertBefore)) + return; + + // Move deps first (in order), then gm2lm. + for (Operation *dep : deps) { + if (dep->isBeforeInBlock(insertBefore)) + continue; + dep->moveBefore(insertBefore); + } + if (!gm2lmOp->isBeforeInBlock(insertBefore)) + gm2lmOp->moveBefore(insertBefore); + LLVM_DEBUG(llvm::dbgs() + << "[AsyncLoadSchedule] hoisted gm2lm (with " << deps.size() + << " deps): " << *gm2lmOp << "\n"); + } + + bool isAsyncProducer(Operation *op) const { + if (auto gm2lmOp = dyn_cast(op)) + return gm2lmOp.getSyncMode() == MemorySyncMode::ASYNC; + if (auto gm2lmOp = dyn_cast(op)) + return gm2lmOp.getSyncMode() == MemorySyncMode::ASYNC; + return false; + } + + bool setAsyncProducer(Operation *op) const { + if (!isa(op)) + return false; + auto async = MemorySyncModeAttr::get(op->getContext(), MemorySyncMode::ASYNC); + op->setAttr("syncMode", async); + return true; + } + + bool hasSingleLoadUser(Operation *producer, + triton::xpu::LoadOp loadOp) const { + if (!producer || producer->getNumResults() != 1) + return false; + Value result = producer->getResult(0); + if (!result.hasOneUse()) + return false; + return *result.user_begin() == loadOp.getOperation(); + } + + Operation *getEarliestUserInBlock(Value value, Block *block) const { + Operation *earliest = nullptr; + for (Operation *user : value.getUsers()) { + if (user->getBlock() != block) + return nullptr; + if (!earliest || user->isBeforeInBlock(earliest)) + earliest = user; + } + return earliest; + } + + FailureOr getForwardingChain(triton::xpu::LoadOp loadOp) const { + ForwardingChain chain; + Block *block = loadOp->getBlock(); + Operation *firstUser = getEarliestUserInBlock(loadOp.getResult(), block); + if (!firstUser) + return failure(); + + if (auto broadcastOp = dyn_cast(firstUser)) { + if (!broadcastOp->getResult(0).hasOneUse()) + return failure(); + Operation *realUser = *broadcastOp->getResult(0).user_begin(); + if (realUser->getBlock() != block) + return failure(); + chain.ops.push_back(broadcastOp.getOperation()); + chain.firstRealUser = realUser; + return chain; + } + + chain.firstRealUser = firstUser; + return chain; + } + + Value getLMBase(Value ptr) const { + Operation *def = ptr.getDefiningOp(); + if (auto gm2lmOp = dyn_cast_or_null(def)) + return gm2lmOp.getBufPtr(); + if (auto gm2lmOp = dyn_cast_or_null(def)) + return gm2lmOp.getBufPtr(); + return ptr; + } + + bool mayAliasLM(Value lhs, Value rhs) const { + Value lhsBase = getLMBase(lhs); + Value rhsBase = getLMBase(rhs); + return !lhsBase || !rhsBase || lhsBase == rhsBase; + } + + bool canMoveAcross(Operation *op, Value movingPtr) const { + if (op->hasTrait() || op->getNumRegions() != 0) + return false; + if (auto loadOp = dyn_cast(op)) + return !mayAliasLM(movingPtr, loadOp.getPtr()); + if (auto gm2lmOp = dyn_cast(op)) + return !mayAliasLM(movingPtr, gm2lmOp.getResult()); + if (auto gm2lmOp = dyn_cast(op)) + return !mayAliasLM(movingPtr, gm2lmOp.getResult()); + if (isa(op)) + return false; + if (isa(op)) + return true; + return isMemoryEffectFree(op); + } + + bool canMoveTo(triton::xpu::LoadOp loadOp, + ArrayRef forwardingOps, + Operation *insertBefore) const { + if (!insertBefore || loadOp->getBlock() != insertBefore->getBlock()) + return false; + if (!loadOp->isBeforeInBlock(insertBefore)) + return false; + + DenseSet movingOps; + movingOps.insert(loadOp.getOperation()); + for (Operation *op : forwardingOps) + movingOps.insert(op); + + // Collect all values produced by moving ops. + DenseSet movingValues; + movingValues.insert(loadOp.getResult()); + for (Operation *op : forwardingOps) + for (Value result : op->getResults()) + movingValues.insert(result); + + Value movingPtr = loadOp.getPtr(); + for (Operation *op = loadOp->getNextNode(); op && op != insertBefore; + op = op->getNextNode()) { + if (movingOps.contains(op)) + continue; + if (!canMoveAcross(op, movingPtr)) + return false; + // Check if this op uses any result of the moving ops — if so, we'd + // break dominance by sinking past it. + for (Value operand : op->getOperands()) { + if (movingValues.contains(operand)) + return false; + } + } + return true; + } + + bool hasMfenceBefore(Operation *op) const { + Operation *prev = op->getPrevNode(); + return prev && isa(prev); + } + + void insertMfenceBefore( + Operation *op, Operation *producer, + llvm::DenseMap &insertedMfences) const { + if (hasMfenceBefore(op)) + return; + OpBuilder builder(op); + auto loc = op->getLoc(); + auto i32Ty = builder.getIntegerType(32); + // Fence LM only (mask bit0=1): the async producer is a GM2LM DMA whose + // data lands in LM, and the op we fence before consumes that LM buffer. + auto fenceValue = builder.create(loc, i32Ty, 1); + auto fenceOp = builder.create(loc, fenceValue); + insertedMfences[fenceOp.getOperation()] = producer; + } + + void eraseDominatedMfences( + llvm::DenseMap &insertedMfences) { + if (insertedMfences.empty()) + return; + DominanceInfo dominance(getOperation()); + + // Collect the fences we inserted, in program (walk) order. + SmallVector mfences; + getOperation()->walk([&](Operation *op) { + if (insertedMfences.count(op)) + mfences.push_back(op); + }); + + // An `mfence 1` is a full LM barrier: once it executes, every GM->LM DMA + // issued before it has completed. So a fence `m` guarding the load of + // producer `p_m` is redundant iff some *retained* fence `keep` is + // guaranteed to execute strictly between `p_m` (DMA issue) and `m` (the + // consuming load): dominates(p_m, keep) && dominates(keep, m). + // + // Crucially we only ever credit a fence that we KEEP. The previous + // implementation credited any dominating fence, which could itself be + // erased later against a fence sitting *before* p_m — collapsing the chain + // onto a fence that no longer covers p_m's buffer and letting the load read + // the buffer before its DMA landed (data race). Sweeping top-down and only + // trusting the most recent retained fence keeps the retained set a valid + // cover for every producer. + DenseSet eraseSet; + Operation *lastKept = nullptr; + for (Operation *m : mfences) { + Operation *producer = insertedMfences.lookup(m); + if (lastKept && producer && + dominance.dominates(producer, lastKept) && + dominance.dominates(lastKept, m)) { + eraseSet.insert(m); + } else { + lastKept = m; + } + } + + for (Operation *mfence : eraseSet) { + Value fenceValue = mfence->getOperand(0); + mfence->erase(); + Operation *constantOp = fenceValue.getDefiningOp(); + if (constantOp && constantOp->use_empty()) + constantOp->erase(); + } + } + + // Find the furthest op in [loadOp+1, insertBefore) that loadOp can move just + // before. Returns nullptr if no valid position exists beyond loadOp itself. + Operation *findFurthestMoveTarget(triton::xpu::LoadOp loadOp, + ArrayRef forwardingOps, + Operation *bound) const { + if (!bound || loadOp->getBlock() != bound->getBlock()) + return nullptr; + Operation *best = nullptr; + DenseSet movingOps; + movingOps.insert(loadOp.getOperation()); + for (Operation *op : forwardingOps) + movingOps.insert(op); + + // Collect all values produced by moving ops. + DenseSet movingValues; + movingValues.insert(loadOp.getResult()); + for (Operation *op : forwardingOps) + for (Value result : op->getResults()) + movingValues.insert(result); + + Value movingPtr = loadOp.getPtr(); + for (Operation *op = loadOp->getNextNode(); op && op != bound; + op = op->getNextNode()) { + if (movingOps.contains(op)) + continue; + if (!canMoveAcross(op, movingPtr)) + break; + // Check if this op uses any result of the moving ops. + bool usesMovingValue = false; + for (Value operand : op->getOperands()) { + if (movingValues.contains(operand)) { + usesMovingValue = true; + break; + } + } + if (usesMovingValue) + break; + best = op->getNextNode(); + } + return best; + } + + void trySchedule(triton::xpu::LoadOp loadOp, + llvm::DenseMap &insertedMfences) const { + Operation *producer = loadOp.getPtr().getDefiningOp(); + if (!isa_and_nonnull(producer)) + return; + if (!hasSingleLoadUser(producer, loadOp)) + return; + bool producerIsAsync = isAsyncProducer(producer); + + auto chainOr = getForwardingChain(loadOp); + if (failed(chainOr)) + return; + ForwardingChain chain = *chainOr; + if (!canMoveTo(loadOp, chain.ops, chain.firstRealUser)) { + // Fallback: move as far as possible short of firstRealUser. + Operation *fallback = + findFurthestMoveTarget(loadOp, chain.ops, chain.firstRealUser); + if (!fallback || fallback == loadOp->getNextNode()) { + if (producerIsAsync) + insertMfenceBefore(loadOp.getOperation(), producer, insertedMfences); + return; + } + loadOp->moveBefore(fallback); + Operation *insertAfter = loadOp.getOperation(); + for (Operation *forwardingOp : chain.ops) { + forwardingOp->moveAfter(insertAfter); + insertAfter = forwardingOp; + } + setAsyncProducer(producer); + insertMfenceBefore(loadOp.getOperation(), producer, insertedMfences); + if (dumpFlag) { + LLVM_DEBUG(llvm::dbgs() + << "[AsyncLoadSchedule] fallback move load: " << *loadOp + << "\n"); + } + return; + } + + Operation *insertBefore = chain.firstRealUser; + loadOp->moveBefore(insertBefore); + Operation *insertAfter = loadOp.getOperation(); + for (Operation *forwardingOp : chain.ops) { + forwardingOp->moveAfter(insertAfter); + insertAfter = forwardingOp; + } + + setAsyncProducer(producer); + insertMfenceBefore(loadOp.getOperation(), producer, insertedMfences); + + if (dumpFlag) { + LLVM_DEBUG(llvm::dbgs() + << "[AsyncLoadSchedule] move load before first user: " + << *loadOp << "\n"); + } + } +}; + +} // namespace + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/CMakeLists.txt b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/CMakeLists.txt index 5664e99749..220c8f1d3a 100644 --- a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/CMakeLists.txt +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/CMakeLists.txt @@ -1,13 +1,20 @@ add_triton_library(TritonXPUTransforms Alloca.cpp + AsyncLoadSchedule.cpp CoreTiling.cpp CreateGM2LM.cpp DtypeConvert.cpp Legalize.cpp + TLELegalize.cpp LoopGrid.cpp + LoopInvariantStaging.cpp LowerPrint.cpp Mask.cpp OffsetAnalysis.cpp + ScalarAnalysis.cpp + TileAnalysisPass.cpp + VectorizabilityAnalysisPass.cpp + Normalize.cpp Vectorize.cpp MemoryAsync.cpp UnrollControl.cpp @@ -18,6 +25,7 @@ add_triton_library(TritonXPUTransforms MemoryInplace.cpp FuncConvert.cpp CFToSCF.cpp + LegalizeExternEW.cpp DEPENDS TritonXPUTransformsIncGen diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/CoreTiling.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/CoreTiling.cpp index cdbe362020..9ec4f15820 100644 --- a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/CoreTiling.cpp +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/CoreTiling.cpp @@ -4,6 +4,7 @@ #include "triton/Dialect/TritonXPU/IR/Dialect.h" #include "triton/Dialect/TritonXPU/Transforms/Passes.h" +#include "llvm/ADT/TypeSwitch.h" #define DEBUG_TYPE "tritonxpu-core-tiling" @@ -283,14 +284,38 @@ struct TritonXPUCoreTilingPass } } + // 将 module 中的 op 按所属逻辑轴划分到两条 def-use 链： + // - innerChain：归属内轴（axis 0）的 op + // - outerChain：归属外轴（axis 1）的 op + // + // 划分以若干"锚点 op"（轴向明确）为起点，通过 getOpDefChainBwd + // 反向追溯它们的 producer 链： + // + // Walk 1（锚点分类）： + // * tt.expand_dims -> 直接由 getAxis() 决定 + // （axis 0 -> inner，axis 1 -> outer）， + // expand_dims 自身不计入链，只收其 producer。 + // * triton_xpu.broadcast -> 通过比较 src / result shape 推断： + // dim 0 被 broadcast 即视为内轴链，否则外轴链。 + // * triton_xpu.lm2gm[_mask] -> 若被存值由 tt.reduce 产生， + // 则按 reduce 的 axis 决定 inner / outer。 + // + // Walk 2（store 的回退分类）： + // 对 Walk 1 未能分类的 lm2gm[_mask]，向上查找 tt.make_range producer， + // 并继承其所在的链。用于处理不由 reduce 喂入的 store。 + // + // 得到的 innerChain / outerChain 会被后续 pass（如 recoverMakeRange） + // 用来为内 / 外轴分别分配不同的 encoding。 void getChain(ModuleOp &mod, SetVector &innerChain, SetVector &outerChain) { + // Walk 1: classify ops via unambiguous anchors. mod.walk([&](mlir::Operation *op) { - if (auto expandDimOp = dyn_cast(op)) { - auto src = expandDimOp.getSrc(); - auto result = expandDimOp.getResult(); - if (auto srcTy = mlir::dyn_cast(src.getType())) { - if (auto resTy = mlir::dyn_cast(result.getType())) { + llvm::TypeSwitch(op) + .Case([&](triton::ExpandDimsOp expandDimOp) { + auto srcTy = dyn_cast(expandDimOp.getSrc().getType()); + auto resTy = dyn_cast(expandDimOp.getResult().getType()); + if (!srcTy || !resTy) + return; if (expandDimOp.getAxis() == 0) { getOpDefChainBwd(innerChain, expandDimOp, expandDimOp); innerChain.remove(expandDimOp); @@ -300,76 +325,59 @@ struct TritonXPUCoreTilingPass } else { llvm_unreachable("expand dim axis must be 0 or 1"); } - dumpOpChain(dumpFlag, op, innerChain, outerChain); - } - } - } else if (auto broadcastOp = dyn_cast(op)) { - auto src = broadcastOp.getSrc(); - auto result = broadcastOp.getResult(); - if (auto srcTy = mlir::dyn_cast(src.getType())) { - if (auto resTy = mlir::dyn_cast(result.getType())) { - auto srcShape = srcTy.getShape(); - auto resShape = resTy.getShape(); - assert(srcShape.size() <= 2); - assert(resShape.size() <= 2); - assert(srcShape.size() == resShape.size()); - if (srcShape[0] != resShape[0]) { // unequal dim 0 shape means - // in the inner axis op chain - getOpDefChainBwd(innerChain, broadcastOp, broadcastOp); - innerChain.remove(broadcastOp); + dumpOpChain(dumpFlag, expandDimOp, innerChain, outerChain); + }) + .Case([&](triton::xpu::BroadcastOp broadcastOp) { + auto srcTy = dyn_cast(broadcastOp.getSrc().getType()); + auto resTy = dyn_cast(broadcastOp.getResult().getType()); + if (srcTy && resTy) { + auto srcShape = srcTy.getShape(); + auto resShape = resTy.getShape(); + assert(srcShape.size() <= 2); + assert(resShape.size() <= 2); + assert(srcShape.size() == resShape.size()); + if (srcShape[0] != resShape[0]) { + // Unequal dim 0 shape means in the inner axis op chain. + getOpDefChainBwd(innerChain, broadcastOp, broadcastOp); + innerChain.remove(broadcastOp); + } else { + getOpDefChainBwd(outerChain, broadcastOp, broadcastOp); + outerChain.remove(broadcastOp); + } + } + dumpOpChain(dumpFlag, broadcastOp, innerChain, outerChain); + }) + .Case([&](auto lm2gmOp) { + auto reduceOp = findDefOpBwd(lm2gmOp.getValue()); + if (!reduceOp) + return; + if (reduceOp.getAxis() == 0) { + getOpDefChainBwd(innerChain, lm2gmOp, lm2gmOp); + } else if (reduceOp.getAxis() == 1) { + getOpDefChainBwd(outerChain, lm2gmOp, lm2gmOp); } else { - getOpDefChainBwd(outerChain, broadcastOp, broadcastOp); - outerChain.remove(broadcastOp); + llvm_unreachable("reduce axis must be 0 or 1"); } - } - } - dumpOpChain(dumpFlag, op, innerChain, outerChain); - } else if (auto lm2gmOp = dyn_cast(op)) { - if (auto reduceOp = - findDefOpBwd(lm2gmOp.getValue())) { - if (reduceOp.getAxis() == 0) { - getOpDefChainBwd(innerChain, lm2gmOp, lm2gmOp); - } else if (reduceOp.getAxis() == 1) { - getOpDefChainBwd(outerChain, lm2gmOp, lm2gmOp); - } else { - llvm_unreachable("reduce axis must be 0 or 1"); - } - } - } else if (auto lm2gmOp = dyn_cast(op)) { - if (auto reduceOp = - findDefOpBwd(lm2gmOp.getValue())) { - if (reduceOp.getAxis() == 0) { - getOpDefChainBwd(innerChain, lm2gmOp, lm2gmOp); - } else if (reduceOp.getAxis() == 1) { - getOpDefChainBwd(outerChain, lm2gmOp, lm2gmOp); - } else { - llvm_unreachable("reduce axis must be 0 or 1"); - } - } - } + }); }); + // Walk 2: fallback for lm2gm[_mask] not classified by Walk 1; inherit + // chain membership from a backward MakeRangeOp producer. mod.walk([&](mlir::Operation *op) { - if (auto lm2gmOp = dyn_cast(op)) { - if (auto _rangeOp = - findDefOpBwd(lm2gmOp.getValue())) { - if (innerChain.contains(_rangeOp)) { - getOpDefChainBwd(innerChain, lm2gmOp, lm2gmOp); - } else if (outerChain.contains(_rangeOp)) { - getOpDefChainBwd(outerChain, lm2gmOp, lm2gmOp); - } - dumpOpChain(dumpFlag, op, innerChain, outerChain); - } - } else if (auto lm2gmOp = dyn_cast(op)) { - if (auto _rangeOp = - findDefOpBwd(lm2gmOp.getValue())) { - if (innerChain.contains(_rangeOp)) { - getOpDefChainBwd(innerChain, lm2gmOp, lm2gmOp); - } else if (outerChain.contains(_rangeOp)) { - getOpDefChainBwd(outerChain, lm2gmOp, lm2gmOp); - } - dumpOpChain(dumpFlag, op, innerChain, outerChain); - } - } + llvm::TypeSwitch(op) + .Case([&](auto lm2gmOp) { + // Skip if already classified by Walk 1 (e.g., via ReduceOp). + if (innerChain.contains(lm2gmOp) || outerChain.contains(lm2gmOp)) + return; + auto rangeOp = findDefOpBwd(lm2gmOp.getValue()); + if (!rangeOp) + return; + if (innerChain.contains(rangeOp)) { + getOpDefChainBwd(innerChain, lm2gmOp, lm2gmOp); + } else if (outerChain.contains(rangeOp)) { + getOpDefChainBwd(outerChain, lm2gmOp, lm2gmOp); + } + dumpOpChain(dumpFlag, lm2gmOp, innerChain, outerChain); + }); }); } @@ -377,31 +385,28 @@ struct TritonXPUCoreTilingPass // In this case, we need to create a new mrOp for innerChain. // The two mrOp will be modified with different [inner/outer] encodings. void recoverMakeRange(ModuleOp &mod) { - mod.walk([&](mlir::Operation *op) { - if (auto rangeOp = dyn_cast(op)) { - OpBuilder builder(rangeOp); - auto loc = builder.getUnknownLoc(); - // Get the Value (the result of the MakeRangeOp) - mlir::Value rangeValue = rangeOp.getResult(); - // Use a list to hold the operands to modify. Iterating over users - // while modifying is generally unsafe/tricky. - llvm::SmallVector usesToChange; - // Collect all uses (mlir::OpOperand*) except the first one (i=0) - int i = 0; - for (mlir::OpOperand &use : rangeValue.getUses()) { - if (i++ > 0) { - usesToChange.push_back(&use); - } - } - // Now, iterate over the collected OpOperands and perform the fix - for (mlir::OpOperand *operandToChange : usesToChange) { - // 1. Clone the operation - auto newRangeOp = builder.create( - loc, rangeOp.getType(), rangeOp.getStart(), rangeOp.getEnd()); - // 2. Set the operand to use the new operation's result - operandToChange->set(newRangeOp.getResult()); + mod.walk([&](triton::MakeRangeOp rangeOp) { + OpBuilder builder(rangeOp); + auto loc = rangeOp.getLoc(); + // Get the Value (the result of the MakeRangeOp) + mlir::Value rangeValue = rangeOp.getResult(); + // Use a list to hold the operands to modify. Iterating over users + // while modifying is generally unsafe/tricky. + llvm::SmallVector usesToChange; + // Collect all uses (mlir::OpOperand*) except the first one (i=0) + for (auto &&[i, use] : llvm::enumerate(rangeValue.getUses())) { + if (i > 0) { + usesToChange.push_back(&use); } } + // Now, iterate over the collected OpOperands and perform the fix + for (mlir::OpOperand *operandToChange : usesToChange) { + // 1. Clone the operation + auto newRangeOp = builder.create( + loc, rangeOp.getType(), rangeOp.getStart(), rangeOp.getEnd()); + // 2. Set the operand to use the new operation's result + operandToChange->set(newRangeOp.getResult()); + } }); } @@ -413,7 +418,7 @@ struct TritonXPUCoreTilingPass size_t groupsize = 64; bool isFirst = true; - auto getGroupInfo = [&](RankedTensorType &tensorType) { + auto getGroupInfo = [&](RankedTensorType tensorType) { if (auto globalEncoding = dyn_cast( tensorType.getEncoding())) { auto shape = tensorType.getShape(); @@ -448,27 +453,32 @@ struct TritonXPUCoreTilingPass "groups_per_cluster only could be 1, 2, 4, 8, 16, 32, 64"); groupsize = ceil(this->coreNum, ngroup); } else { - mod.walk([&](mlir::Operation *op) { - if (auto reduceOp = dyn_cast(op)) { - if (auto tensorType = - dyn_cast(reduceOp.getOperandTypes()[0])) { - if (tensorType.getShape().size() == 2) { - getGroupInfo(tensorType); - } else if (isAxisNone(reduceOp)) { - auto defOp = reduceOp.getSrcs()[0].getDefiningOp(); - if (auto reshapeOp = dyn_cast(defOp)) { - if (auto reshapeResTy = dyn_cast( - reshapeOp.getResult().getType())) { - if (reshapeResTy.getShape().size() == 1) { - auto reshapeSrcTy = cast( - reshapeOp.getOperand().getType()); - getGroupInfo(reshapeSrcTy); - } - } - } - } - } + mod.walk([&](triton::ReduceOp reduceOp) { + auto tensorType = + dyn_cast(reduceOp.getOperandTypes()[0]); + if (!tensorType) + return; + + if (tensorType.getShape().size() == 2) { + getGroupInfo(tensorType); + return; } + + if (!isAxisNone(reduceOp)) + return; + + auto reshapeOp = + reduceOp.getSrcs()[0].getDefiningOp(); + if (!reshapeOp) + return; + + auto reshapeResTy = + dyn_cast(reshapeOp.getResult().getType()); + if (!reshapeResTy || reshapeResTy.getShape().size() != 1) + return; + + getGroupInfo( + cast(reshapeOp.getOperand().getType())); }); } LLVM_DEBUG(llvm::dbgs() << "[Reduction SoftGroup]: " @@ -509,16 +519,21 @@ struct TritonXPUCoreTilingPass // must be globalEncoding if (auto parentEncoding = dyn_cast( sliceEncoding.getParent())) { + // Build a padded 2D type with the parent ClusterLayoutAttr so + // that getOptimizedGEncoding can process it (it requires + // ClusterLayoutAttr encoding, not SliceEncodingAttr). + auto paddedShape = sliceEncoding.paddedShape(shape); + auto paddedTy = RankedTensorType::get(paddedShape, elemTy, + parentEncoding); auto newParentEncoding = - getOptimizedGEncoding(context, resTy, innerChain, outerChain, - op, ngroup, groupsize); - // Mirror 3.0 behavior: getOptimizedGEncoding(resTy) returns null - // when resTy carries SliceEncoding (it only matches - // ClusterLayoutAttr). 3.0 builds a SliceEncodingAttr with a null - // Attribute parent and lets Step 2.2/2.4 overwrite the encoding - // later. 3.6 changed the parent to DistributedEncodingTrait, so a - // null cast would abort. Skip the rewrite here when the parent - // is unavailable; the follow-up walks fix it correctly. + getOptimizedGEncoding(context, paddedTy, innerChain, + outerChain, op, ngroup, groupsize); + // Pass paddedTy (not resTy): getOptimizedGEncoding only matches + // ClusterLayoutAttr, so handing it resTy (which carries + // SliceEncoding) returns null. The null guard stays as a + // safety net because 3.6 requires the parent to be a + // DistributedEncodingTrait, and casting null would abort; + // Step 2.2/2.4 overwrite the encoding later anyway. if (newParentEncoding) { newEncoding = triton::gpu::SliceEncodingAttr::get( context, sliceEncoding.getDim(), @@ -589,9 +604,7 @@ struct TritonXPUCoreTilingPass cast(cvtOpType.getEncoding()); auto newSliceEncoding = triton::gpu::SliceEncodingAttr::get( - context, sliceEncoding.getDim(), - cast( - static_cast(globalEncoding))); + context, sliceEncoding.getDim(), cast(static_cast(globalEncoding))); auto newResTy = RankedTensorType::get( cvtOpType.getShape(), cvtOpType.getElementType(), newSliceEncoding); @@ -657,9 +670,7 @@ struct TritonXPUCoreTilingPass cast(srcTy.getEncoding()); auto newEncoding = triton::gpu::SliceEncodingAttr::get( - context, resSliceEncoding.getDim(), - cast( - static_cast(srcGlobalEncoding))); + context, resSliceEncoding.getDim(), cast(static_cast(srcGlobalEncoding))); auto newResTy = RankedTensorType::get( resTy.getShape(), resTy.getElementType(), newEncoding); diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/LegalizeExternEW.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/LegalizeExternEW.cpp new file mode 100644 index 0000000000..82a11320f4 --- /dev/null +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/LegalizeExternEW.cpp @@ -0,0 +1,109 @@ +//===----------------------------------------------------------------------===// +// Legalize unsupported tt.extern_elementwise ops for the non-SDNN XPU path. +// +// libdevice.fast_expf and libdevice.fast_dividef are not present in +// libdevice-xpu3.bc. In the non-SDNN path these ops survive as unresolved +// external LLVM function calls and cause linker errors. This pass rewrites +// them to native MLIR ops that are fully supported by TritonXPUToLLVM: +// +// tt.extern_elementwise(%x) {symbol="libdevice.fast_expf"} +// -> math.exp(%x) +// -> LLVM::Exp2Op(%x) -> XPU ASM: exp.f.rn = e^x ✓ +// +// NOTE: Although the XPU backend maps math::ExpOp to LLVM::Exp2Op, the +// hardware instruction exp.f.rn computes the natural exponential e^x +// (not 2^x). No log2(e) pre-scaling is needed. +// +// tt.extern_elementwise(%a,%b) {symbol="libdevice.fast_dividef"} +// tt.extern_elementwise(%a,%b) {symbol="__triton_houyi_fast_fdividef"} +// -> arith.divf %a, %b +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/Transforms/Passes.h" + +#define DEBUG_TYPE "tritonxpu-legalize-extern-ew" + +namespace mlir { +namespace triton { +namespace xpu { + +#define GEN_PASS_DEF_TRITONXPULEGALIZEEXTERNEW +#include "triton/Dialect/TritonXPU/Transforms/Passes.h.inc" + +namespace { + +// fast_expf(x) = e^x +// +// The XPU backend maps math::ExpOp -> LLVM::Exp2Op -> XPU ASM exp.f.rn. +// The hardware instruction exp.f.rn computes the natural exponential e^x +// directly (not 2^x), so no log2(e) pre-scaling is required. +// +// tt.extern_elementwise(%x) {symbol="libdevice.fast_expf"} +// -> math.exp(%x) +struct FastExpfToMathExpPattern + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::ExternElementwiseOp op, + PatternRewriter &rewriter) const override { + if (op.getSymbol() != "libdevice.fast_expf") + return failure(); + auto srcs = op.getSrcs(); + if (srcs.size() != 1) + return failure(); + + rewriter.replaceOpWithNewOp(op, srcs[0]); + return success(); + } +}; + +// tt.extern_elementwise(%a, %b) {symbol="libdevice.fast_dividef"} +// tt.extern_elementwise(%a, %b) {symbol="__triton_houyi_fast_fdividef"} +// -> arith.divf %a, %b +struct FastDividefToArithDivfPattern + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::ExternElementwiseOp op, + PatternRewriter &rewriter) const override { + StringRef symbol = op.getSymbol(); + if (symbol != "libdevice.fast_dividef" && + symbol != "__triton_houyi_fast_fdividef") + return failure(); + auto srcs = op.getSrcs(); + if (srcs.size() != 2) + return failure(); + rewriter.replaceOpWithNewOp(op, srcs[0], srcs[1]); + return success(); + } +}; + +} // namespace + +struct TritonXPULegalizeExternEW + : public impl::TritonXPULegalizeExternEWBase { + + using impl::TritonXPULegalizeExternEWBase< + TritonXPULegalizeExternEW>::TritonXPULegalizeExternEWBase; + + void runOnOperation() override { + MLIRContext *context = &getContext(); + RewritePatternSet patterns(context); + patterns.add( + context); + if (failed(applyPatternsGreedily(getOperation(), + std::move(patterns)))) { + signalPassFailure(); + } + } +}; + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/LoopInvariantStaging.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/LoopInvariantStaging.cpp new file mode 100644 index 0000000000..81cdba1da4 --- /dev/null +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/LoopInvariantStaging.cpp @@ -0,0 +1,409 @@ +//===----------------------------------------------------------------------===// +// TritonXPULoopInvariantStaging +// +// Stage a loop-invariant, read-only, small-packet `gm2lm` out of the grid +// `scf.for` into a cluster-shared SM buffer once in the preheader, then read +// from that staged buffer inside the loop with a single indexed scalar load +// (`triton_xpu.load_scalar_indexed`). +// +// Target pattern (test 88/89): inside the grid loop a `DiscreteSame` gm2lm +// reads one scalar `idx[x1]` from a kernel-arg array `%arg0`, where `x1` is a +// per-iteration index uniform across all lanes. Instead of doing a per- +// iteration GM->LM small packet DMA + mfence, we stage the whole `idx` array +// once outside the loop into a cluster-shared SM buffer: core 0 issues a single +// GM->SM DMA (bounded at runtime by `n_idx = ceilDiv(xnumel, stride)`) followed +// by a cluster barrier, and all 64 cores then share that one SM copy. The +// staging buffer is sized dynamically from `n_idx` and bounded only by the +// physical SM ceiling (no external budget). +//===----------------------------------------------------------------------===// + +#include "triton/Analysis/NewAnalysis/Utility.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/Transforms/Passes.h" + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "tritonxpu-loop-invariant-staging" + +namespace mlir { +namespace triton { +namespace xpu { + +#define GEN_PASS_DEF_TRITONXPULOOPINVARIANTSTAGING +#include "triton/Dialect/TritonXPU/Transforms/Passes.h.inc" + +namespace { + +// Returns true if `v` is defined outside `loop` (i.e. loop-invariant w.r.t. the +// grid loop region). +static bool isDefinedOutside(Value v, scf::ForOp loop) { + if (auto *def = v.getDefiningOp()) + return !loop->isAncestor(def); + // Block argument: invariant iff its owning block is not inside the loop. + return !loop->isAncestor(v.getParentBlock()->getParentOp()); +} + +// Walk back through tt.splat / tt.addptr style chains to find whether the base +// pointer of `gm2lmOp` originates from a kernel pointer argument (FuncOp block +// argument). Returns that base Value if found. +static Value findInvariantBasePtr(triton::xpu::GM2LMOp gm2lmOp, + scf::ForOp loop) { + Value ptr = gm2lmOp.getPtr(); + while (auto addptr = ptr.getDefiningOp()) + ptr = addptr.getPtr(); + if (auto splat = ptr.getDefiningOp()) + ptr = splat.getSrc(); + if (isDefinedOutside(ptr, loop) && isa(ptr)) + return ptr; + return Value(); +} + +// Peel `gm2lmOp.getPtr()` to the index tensor feeding the addptr offset. +static Value getIndexTensor(triton::xpu::GM2LMOp gm2lmOp) { + Value ptr = gm2lmOp.getPtr(); + if (auto addptr = ptr.getDefiningOp()) + return addptr.getOffset(); + return Value(); +} + +// From the index tensor `%idx = arith.divsi(%dividend, %constStride)`, return +// the constant stride. Returns 0 if the pattern does not match. +static int64_t getConstStride(Value indexTensor) { + auto divsi = indexTensor.getDefiningOp(); + if (!divsi) + return 0; + auto cstOp = divsi.getRhs().getDefiningOp(); + if (!cstOp) + return 0; + auto denseAttr = dyn_cast(cstOp.getValue()); + if (!denseAttr || !denseAttr.isSplat()) + return 0; + return denseAttr.getSplatValue().getSExtValue(); +} + +// Staging the first `n_idx = ceilDiv(xnumel, stride)` elements of the array is +// only correct if the per-iteration index `x1 = dividend / stride` densely +// covers `0 .. n_idx-1`, i.e. `dividend` is derived from the grid loop's +// induction variable and the per-block `make_range(0..XBLOCK)` (a contiguous +// `xindex`). If `dividend` originates from some other source (a gathered pid, +// an index loaded from memory, a value-dependent selection), the +// contiguous-prefix assumption breaks and staging would read wrong / unstaged +// entries. +// +// A contiguous `xindex` may be reshaped into multi-dimensional grid +// coordinates by constant integer div/rem (e.g. `(pid / a) % b`). Such +// decompositions keep dense coverage: each derived coordinate still ranges +// over a contiguous grid, so they are allowed to sit on the chain. Only +// *non-constant* div/rem (data-dependent reshaping) and genuinely +// coverage-breaking ops (gather load, select, ...) are rejected. +// +// Returns true only when at least one contiguous source (induction var or +// make_range) is found and no unexpected non-invariant root is reached. +static bool isContiguousIndex(Value dividend, scf::ForOp loop) { + Value iv = loop.getInductionVar(); + SmallVector work{dividend}; + SmallPtrSet seen; + bool sawContiguousSource = false; + while (!work.empty()) { + Value v = work.pop_back_val(); + if (!seen.insert(v).second) + continue; + + if (v == iv) { + sawContiguousSource = true; + continue; + } + + Operation *def = v.getDefiningOp(); + if (!def) { + // Block argument other than the induction var: only safe if it is + // loop-invariant (e.g. a kernel scalar arg used as an additive offset). + if (isDefinedOutside(v, loop)) + continue; + return false; + } + + // A make_range is a contiguous 0..N source. + if (isa(def)) { + sawContiguousSource = true; + continue; + } + + // Constants are loop-invariant scalars/splats and never break contiguity of + // the dividend's variable part (they only scale/shift it). A constant whose + // defining op happens to sit inside the loop body would not be reported as + // "defined outside", so handle it explicitly here. + if (isa(def)) + continue; + + // Loop-invariant values (kernel args, anything defined before the loop) + // cannot break contiguity of the dividend's variable part. + if (isDefinedOutside(v, loop)) + continue; + + // Constant-divisor div/rem reshape a contiguous index into dense grid + // coordinates; keep walking only the dividend operand. A non-constant + // (loop-variant) divisor could reorder/sparsify -> bail out. + if (auto divsi = dyn_cast(def)) { + if (!isDefinedOutside(divsi.getRhs(), loop)) + return false; + work.push_back(divsi.getLhs()); + continue; + } + if (auto remsi = dyn_cast(def)) { + if (!isDefinedOutside(remsi.getRhs(), loop)) + return false; + work.push_back(remsi.getLhs()); + continue; + } + + // Linear / shape-only integer ops are allowed to sit between the + // contiguous source and the dividend. + if (isa(def)) { + for (Value operand : def->getOperands()) + work.push_back(operand); + continue; + } + + // Anything else (gather load, select, data-dependent index, ...) may + // destroy dense coverage -> bail out conservatively. + return false; + } + return sawContiguousSource; +} + + +// From the gm2lm `len` operand `%len = arith.subi(tt.splat(%xnumel), %idx)`, +// return the loop-invariant scalar `%xnumel`. The splat source may be wrapped +// in extension ops (e.g. `arith.extsi %arg3`); peel them to reach an invariant +// root scalar. Returns null if not matched. +static Value getXnumelFromLen(triton::xpu::GM2LMOp gm2lmOp, scf::ForOp loop) { + Value len = gm2lmOp.getLen(); + if (!len) + return Value(); + auto subi = len.getDefiningOp(); + if (!subi) + return Value(); + auto splat = subi.getLhs().getDefiningOp(); + if (!splat) + return Value(); + Value scalar = splat.getSrc(); + // Peel integer extension ops to reach the invariant root. + while (true) { + if (auto ext = scalar.getDefiningOp()) { + scalar = ext.getIn(); + continue; + } + if (auto ext = scalar.getDefiningOp()) { + scalar = ext.getIn(); + continue; + } + if (auto tr = scalar.getDefiningOp()) { + scalar = tr.getIn(); + continue; + } + break; + } + if (!isDefinedOutside(scalar, loop)) + return Value(); + return scalar; +} + +} // namespace + +struct TritonXPULoopInvariantStaging + : public impl::TritonXPULoopInvariantStagingBase< + TritonXPULoopInvariantStaging> { + + using impl::TritonXPULoopInvariantStagingBase< + TritonXPULoopInvariantStaging>::TritonXPULoopInvariantStagingBase; + + void runOnOperation() override { + ModuleOp m = getOperation(); + + SmallVector> candidates; + m.walk([&](scf::ForOp loop) { + // Only target the grid-dispatch loop produced by TritonXPULoopGrid: + // it lives directly in the FuncOp body. + if (!isa(loop->getParentOp())) + return; + + loop.getBody()->walk([&](triton::xpu::GM2LMOp gm2lmOp) { + auto offsetState = static_cast(gm2lmOp.getOffsetState()); + bool isScalarAttr = false; + if (auto a = gm2lmOp->getAttrOfType("isScalar")) + isScalarAttr = a.getValue(); + bool isSmallPacket = + isScalarAttr || offsetState == OffsetState::DiscreteSame; + if (!isSmallPacket) + return; + if (!findInvariantBasePtr(gm2lmOp, loop)) + return; + candidates.push_back({loop, gm2lmOp}); + }); + }); + + for (auto &[loop, gm2lmOp] : candidates) + tryStage(loop, gm2lmOp); + } + + void tryStage(scf::ForOp loop, triton::xpu::GM2LMOp gm2lmOp) { + Value base = findInvariantBasePtr(gm2lmOp, loop); + Value indexTensor = getIndexTensor(gm2lmOp); + if (!indexTensor) + return; + int64_t stride = getConstStride(indexTensor); + if (stride <= 0) + return; + // The staged buffer holds the contiguous prefix base[0 .. n_idx). This is + // only valid when the per-iteration index densely covers that prefix, i.e. + // the dividend feeding `x1 = dividend / stride` is built from the grid + // loop's induction variable / make_range. Reject non-contiguous pids. + auto divsi = indexTensor.getDefiningOp(); + if (!divsi || !isContiguousIndex(divsi.getLhs(), loop)) + return; + Value xnumel = getXnumelFromLen(gm2lmOp, loop); + if (!xnumel || !xnumel.getType().isIntOrIndex()) + return; + + // The gm2lm must feed a single load whose result we will replace. + triton::xpu::LoadOp loadOp; + for (auto *user : gm2lmOp.getResult().getUsers()) { + if (auto l = dyn_cast(user)) { + if (loadOp) + return; // more than one load: bail out conservatively. + loadOp = l; + } else { + return; // unexpected consumer. + } + } + if (!loadOp) + return; + + auto idxTensorTy = dyn_cast(indexTensor.getType()); + auto loadResTy = dyn_cast(loadOp.getResult().getType()); + if (!idxTensorTy || !loadResTy) + return; + + Type elemTy = idxTensorTy.getElementType(); // staged scalar element type. + + // --- Build the preheader staging ops, right before the grid loop. --- + OpBuilder builder(loop); + Location loc = gm2lmOp.getLoc(); + + auto i32Ty = builder.getI32Type(); + + // Runtime length n_idx = ceilDiv(xnumel, stride) = (xnumel + stride-1)/stride. + Type xnumelTy = xnumel.getType(); + auto sMinus1 = builder.create( + loc, xnumelTy, builder.getIntegerAttr(xnumelTy, stride - 1)); + auto strideC = builder.create( + loc, xnumelTy, builder.getIntegerAttr(xnumelTy, stride)); + Value sum = builder.create(loc, xnumel, sMinus1); + Value nIdx = builder.create(loc, sum, strideC); + + // n_idx as i32 (the DMA length / buffer-sizing arithmetic is byte-granular). + Value nIdxI32 = nIdx; + if (!xnumelTy.isInteger(32)) { + nIdxI32 = builder.create(loc, i32Ty, nIdx); + } + + // Staging capacity is bounded purely by the physical SM ceiling, not by an + // external budget. SM is 256KB shared cluster-wide; reduce/scan scratch + // grows from offset 0 upward, so the staging buffer is placed at the top of + // SM and may use at most `kStagingMaxBytes` (the rest is reserved for + // scratch). The element capacity `availElems = kStagingMaxBytes/elemBytes` + // is a compile-time constant derived from the staged element width. + unsigned elemBytes = elemTy.getIntOrFloatBitWidth() / 8u; + constexpr int32_t kSMTotalBytes = 256 * 1024; + // Reserve the lower half of SM for reduce/scan scratch; stage into the top. + constexpr int32_t kScratchReserveBytes = 128 * 1024; + constexpr int32_t kStagingMaxBytes = kSMTotalBytes - kScratchReserveBytes; + int32_t availElems = kStagingMaxBytes / static_cast(elemBytes); + + // bufElems = clamp(n_idx, 0, availElems) as i32 (runtime). The stage_sm DMA + // in the preheader runs unconditionally, so the buffer size must stay within + // the SM ceiling even when the runtime guard below selects the fallback. + auto zeroI32 = builder.create( + loc, i32Ty, builder.getI32IntegerAttr(0)); + auto availI32 = builder.create( + loc, i32Ty, builder.getI32IntegerAttr(availElems)); + Value bufElemsV = builder.create(loc, nIdxI32, zeroI32); + bufElemsV = builder.create(loc, bufElemsV, availI32); + + // smOffset = 256KB - bufElems*elemBytes (runtime, byte granularity). + auto elemBytesC = builder.create( + loc, i32Ty, builder.getI32IntegerAttr(static_cast(elemBytes))); + Value bufBytesV = builder.create(loc, bufElemsV, elemBytesC); + auto smTotalC = builder.create( + loc, i32Ty, builder.getI32IntegerAttr(kSMTotalBytes)); + Value smOffsetV = builder.create(loc, smTotalC, bufBytesV); + + // SM staging op: core0-only GM->SM DMA + cluster barrier. The result is a + // scalar SM base pointer (addrspace 2) shared by all cores. `$ptr` is the + // scalar GM base pointer and `$len` is the scalar n_idx; the DMA length is + // clamped to bufElems by the StageSM lowering. When n_idx > availElems this + // stages a truncated prefix; the runtime guard below makes sure the + // truncated buffer is never read. + auto smPtrTy = triton::PointerType::get(elemTy, /*addrSpace=*/2); + auto stageSM = builder.create( + loc, smPtrTy, /*ptr=*/base, /*len=*/nIdxI32, + /*smOffset=*/smOffsetV, /*bufElems=*/bufElemsV, gm2lmOp.getSyncMode()); + + // Runtime guard: the staged buffer covers at most availElems elements (the + // SM ceiling). When the live length n_idx exceeds it, staging is unsafe + // (truncated reads -> out-of-bounds / wrong results); fall back to the + // original per-iteration gm2lm+load. n_idx is uniform across cores, so this + // branch is taken identically by all lanes. + auto guardC = builder.create( + loc, xnumelTy, + builder.getIntegerAttr(xnumelTy, static_cast(availElems))); + Value useStaged = builder.create( + loc, arith::CmpIPredicate::sle, nIdx, guardC); + + // --- Inside the loop: replace the per-iteration gm2lm+load with a runtime + // branch selecting the staged SM read or the original gm2lm fallback. --- + OpBuilder inLoop(loadOp); + Location lloc = loadOp.getLoc(); + auto ifOp = inLoop.create(lloc, TypeRange{loadResTy}, useStaged, + /*withElseRegion=*/true); + + // then: read the staged scalar from the cluster-shared SM buffer. + { + OpBuilder::InsertionGuard g(inLoop); + inLoop.setInsertionPointToStart(ifOp.thenBlock()); + auto idxScalar = inLoop.create( + lloc, elemTy, inLoop.getI32IntegerAttr(0), indexTensor); + auto scalarLoad = inLoop.create( + lloc, loadResTy, stageSM.getResult(), idxScalar, + gm2lmOp.getSyncMode()); + inLoop.create(lloc, scalarLoad.getResult()); + } + + // else: original path. Move the per-iteration gm2lm + load into the else + // region and yield the loaded value. Their operands are defined before the + // grid loop body / before this point, so they still dominate the region. + scf::YieldOp elseYield; + { + OpBuilder::InsertionGuard g(inLoop); + gm2lmOp->moveBefore(ifOp.elseBlock(), ifOp.elseBlock()->end()); + loadOp->moveBefore(ifOp.elseBlock(), ifOp.elseBlock()->end()); + inLoop.setInsertionPointToEnd(ifOp.elseBlock()); + elseYield = inLoop.create(lloc, loadOp.getResult()); + } + + // Route every consumer outside the new region to the if-result; the only + // remaining direct use of loadOp is the else-region yield. + SmallPtrSet except; + except.insert(elseYield.getOperation()); + loadOp.getResult().replaceAllUsesExcept(ifOp.getResult(0), except); + } +}; + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/Normalize.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/Normalize.cpp new file mode 100644 index 0000000000..8b26d7380e --- /dev/null +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/Normalize.cpp @@ -0,0 +1,358 @@ +//===----------------------------------------------------------------------===// +// Pre-vectorization normalization. +// +// Split out of TritonXPUVectorize's prologue (redesign-v2 §4.2 step 1.5). None +// of this is vectorization: it rewrites scalar IR into the shapes the +// vectorizer knows how to match (libdevice calls instead of math dialect ops, +// fused max/min instead of [cmpf, cmpf, ori, select], compares that carry +// their i8 result in an i32 vector lane). Step 1.3 measured that these +// rewrites *create* vectorizability -- `boolfused` has closure=0 before them +// and closure=6 after -- so any analysis that predicts what Vectorize will do +// has to run after them. That is why they have to leave Vectorize before the +// state analysis can move ahead of it (step 1.5c). +// +// Position: immediately before tritonxpu-vectorize, i.e. exactly where the +// code used to run, with one difference -- Vectorize's own `vectorizeTLE` now +// runs *after* these rewrites instead of before them. That can only matter for +// a TLE kernel whose local-buffer store chain contains the NaN-aware max/min +// select pattern: `doMaximumFusion` folds it to arith.maximumf, which *is* in +// ARITH_BINARY_FLOAT_OP, so such a chain would newly vectorize. `vectorizeTLE` +// stays behind because it needs Vectorize.cpp's VOp table; no probe covers +// the TLE path (golden.py has none), so this is stated, not measured. +//===----------------------------------------------------------------------===// + +// clang-format off +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Math/IR/Math.h" +#include "triton/Analysis/Utility.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Analysis/VectorizabilityAnalysis.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/Transforms/Passes.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +// clang-format on + +#define DEBUG_TYPE "tritonxpu-normalize" + +namespace mlir { +namespace triton { +namespace xpu { + +#define GEN_PASS_DEF_TRITONXPUNORMALIZE +#include "triton/Dialect/TritonXPU/Transforms/Passes.h.inc" + +// ext(logic_op(cmp ne(a, 0))) -> ext(locic_op(a)) +template +struct ConvertI1LogicOpToI8 : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(arith::ExtUIOp extOp, + PatternRewriter &rewriter) const override { + auto i1TensorType = dyn_cast(extOp.getIn().getType()); + auto i8TensorType = dyn_cast(extOp.getType()); + if (!i1TensorType || !i8TensorType || + i1TensorType.getElementTypeBitWidth() != 1) + return failure(); + + auto logicOp = extOp.getIn().getDefiningOp(); + if (!logicOp) + return failure(); + + auto getI8Source = [&](Value v) -> Value { + auto cmpi = v.getDefiningOp(); + if (!cmpi || cmpi.getPredicate() != arith::CmpIPredicate::ne) + return nullptr; + + if (!getElementTypeOrSelf(cmpi.getLhs()).isInteger(8)) + return nullptr; + + if (isZeroConst(cmpi.getRhs())) + return cmpi.getLhs(); + + if (isZeroConst(cmpi.getLhs())) + return cmpi.getRhs(); + + return nullptr; + }; + + Value lhsSource = getI8Source(logicOp.getLhs()); + Value rhsSource = getI8Source(logicOp.getRhs()); + + if (lhsSource && rhsSource) { + rewriter.replaceOpWithNewOp(extOp, lhsSource, rhsSource); + return success(); + } + + return failure(); + } +}; + +// fold extui +struct BypassCmpIExtUI : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(arith::ExtUIOp extOp, + PatternRewriter &rewriter) const override { + if (extOp.getIn().getType() == extOp.getOut().getType()) { + rewriter.replaceOp(extOp, extOp.getIn()); + return success(); + } + return failure(); + } +}; + +struct TritonXPUNormalize + : public impl::TritonXPUNormalizeBase { + using impl::TritonXPUNormalizeBase::TritonXPUNormalizeBase; + + void doMaximumFusion(arith::SelectOp selectOp) { + if (auto orIOp = selectOp.getCondition().getDefiningOp()) { + if (orIOp.getResult().hasOneUse()) { + auto lhs = orIOp.getLhs().getDefiningOp(); + auto rhs = orIOp.getRhs().getDefiningOp(); + + // The null check comes before getPredicate() here; in Vectorize it came + // after, so an or of anything but two cmpf dereferenced null. Cannot + // change the emitted code: such a module crashed instead of compiling. + if (!lhs || !rhs || !lhs.getResult().hasOneUse() || + !rhs.getResult().hasOneUse()) + return; + + bool isMax = (lhs.getPredicate() == arith::CmpFPredicate::OGT && + rhs.getPredicate() == arith::CmpFPredicate::UNE) || + (lhs.getPredicate() == arith::CmpFPredicate::UNE && + rhs.getPredicate() == arith::CmpFPredicate::OGT); + bool isMin = (lhs.getPredicate() == arith::CmpFPredicate::OLT && + rhs.getPredicate() == arith::CmpFPredicate::UNE) || + (lhs.getPredicate() == arith::CmpFPredicate::UNE && + rhs.getPredicate() == arith::CmpFPredicate::OLT); + + OpBuilder builder(selectOp); + if (isMax) { + auto newMaxFOp = builder.create( + selectOp.getLoc(), selectOp.getType(), selectOp.getTrueValue(), + selectOp.getFalseValue()); + selectOp->replaceAllUsesWith(newMaxFOp); + selectOp->erase(); + orIOp->erase(); + lhs->erase(); + rhs->erase(); + LLVM_DEBUG(llvm::dbgs() + << "[Normalize]: Apply Maximum Fusion Optimization " + "For VVMax.\n"); + } else if (isMin) { + auto newMinFOp = builder.create( + selectOp.getLoc(), selectOp.getType(), selectOp.getTrueValue(), + selectOp.getFalseValue()); + selectOp->replaceAllUsesWith(newMinFOp); + selectOp->erase(); + orIOp->erase(); + lhs->erase(); + rhs->erase(); + LLVM_DEBUG(llvm::dbgs() + << "[Normalize]: Apply Minimum Fusion Optimization " + "For VVMin.\n"); + } + } + } + } + + void doCompareExtUI8Fusion(arith::ExtUIOp extUIOp) { + auto inTy = extUIOp.getIn().getType(); + auto outTy = extUIOp.getOut().getType(); + auto inElemTy = getElementTypeOrSelf(inTy); + auto outElemTy = getElementTypeOrSelf(outTy); + // Only Vectorize Do Fusion + auto rowsPerCore = 1; + if (auto outTensorTy = mlir::dyn_cast(outTy)) { + auto rank = outTensorTy.getShape().size(); + if (rank > 1) { + rowsPerCore = mlir::cast( + outTensorTy.getEncoding()) + .getSizePerCore()[0]; + } + } + unsigned numElems = getTotalElemsPerThread(outTy) / rowsPerCore; + Type elemTy = getElementTypeOrSelf(outTy); + auto elemWidth = elemTy.getIntOrFloatBitWidth(); + auto vectorWidth = 512 / elemWidth; + if (numElems < vectorWidth || numElems % vectorWidth > 0 || + !vectorizedTyValid(elemTy)) + return; + // Fuse CmpFOp(i1) + ExtUIOp(i8) + StoreOp = CmpFOp(i32) + StoreOp + if (auto cmpFOp = extUIOp.getIn().getDefiningOp()) { + if (inElemTy.isInteger(1) && outElemTy.isInteger(8)) { + for (auto user : extUIOp.getOut().getUsers()) { + if (auto storeOp = dyn_cast(user)) { + if (auto outTensorTy = dyn_cast(outTy)) { + auto lhsTy = cmpFOp.getLhs().getType(); + auto lhsElemTy = + getElementTypeOrSelf(getElementTypeOrSelf(lhsTy)); + auto context = storeOp.getContext(); + if (lhsElemTy.isF32()) { + auto dtype = DtypeAttr::get(context, Dtype::FP32); + storeOp->setAttr("dtype", dtype); + } else if (lhsElemTy.isF16()) { + auto dtype = DtypeAttr::get(context, Dtype::FP16); + storeOp->setAttr("dtype", dtype); + } else { + llvm_unreachable( + "CompareExtUI8Fusion only supports FP32 or FP16"); + } + OpBuilder builder(extUIOp); + auto newTensorTy = RankedTensorType::get( + outTensorTy.getShape(), builder.getIntegerType(32, false), + outTensorTy.getEncoding()); + auto newCmpFOp = builder.create( + extUIOp.getLoc(), newTensorTy, cmpFOp.getPredicate(), + cmpFOp.getLhs(), cmpFOp.getRhs()); + extUIOp.getOut().replaceAllUsesWith(newCmpFOp.getResult()); + extUIOp.erase(); + cmpFOp.erase(); + bf16ToFP32VecOptOff = true; + } + } + } + } + } + } + + void doCompareTruncI8Fusion(arith::TruncIOp truncIOp) { + auto inTy = truncIOp.getIn().getType(); + auto outTy = truncIOp.getOut().getType(); + auto inElemTy = getElementTypeOrSelf(inTy); + auto outElemTy = getElementTypeOrSelf(outTy); + // Only Vectorize Do Fusion + auto rowsPerCore = 1; + if (auto outTensorTy = mlir::dyn_cast(outTy)) { + auto rank = outTensorTy.getShape().size(); + if (rank > 1) { + rowsPerCore = mlir::cast( + outTensorTy.getEncoding()) + .getSizePerCore()[0]; + } + } + unsigned numElems = getTotalElemsPerThread(outTy) / rowsPerCore; + Type elemTy = getElementTypeOrSelf(outTy); + auto elemWidth = elemTy.getIntOrFloatBitWidth(); + auto vectorWidth = 512 / elemWidth; + if (numElems < vectorWidth || numElems % vectorWidth > 0 || + !vectorizedTyValid(elemTy)) + return; + // Fuse ExtElemwiseOp(i8) + TruncIOp(i8) + StoreOp = newExtElemwiseOp(i32) + + // StoreOp + if (auto extElemwiseOp = + truncIOp.getIn().getDefiningOp()) { + if (extElemwiseOp.getSymbol() == "_ZN3xpu5isnanEf" && + inElemTy.isInteger(32) && outElemTy.isInteger(8)) { + for (auto user : truncIOp.getOut().getUsers()) { + if (auto storeOp = dyn_cast(user)) { + if (auto outTensorTy = dyn_cast(outTy)) { + auto inTy = extElemwiseOp.getOperands().front().getType(); + auto inElemTy = getElementTypeOrSelf(getElementTypeOrSelf(inTy)); + auto context = storeOp.getContext(); + if (inElemTy.isF32()) { + auto dtype = DtypeAttr::get(context, Dtype::FP32); + storeOp->setAttr("dtype", dtype); + } else if (inElemTy.isF16()) { + auto dtype = DtypeAttr::get(context, Dtype::FP16); + storeOp->setAttr("dtype", dtype); + } else { + llvm_unreachable( + "CompareExtUI8Fusion only supports FP32 or FP16"); + } + OpBuilder builder(truncIOp); + auto newTensorTy = RankedTensorType::get( + outTensorTy.getShape(), builder.getIntegerType(32), + outTensorTy.getEncoding()); + auto newExtElemwiseOp = + builder.create( + truncIOp.getLoc(), newTensorTy, + extElemwiseOp.getOperands().front(), + extElemwiseOp.getLibname(), extElemwiseOp.getLibpath(), + extElemwiseOp.getSymbol(), extElemwiseOp.getPure()); + + truncIOp.getOut().replaceAllUsesWith( + newExtElemwiseOp.getResult()); + truncIOp.erase(); + extElemwiseOp.erase(); + bf16ToFP32VecOptOff = true; + } + } + } + } + } + } + + void runOnOperation() override { + context = &getContext(); + ModuleOp mod = getOperation(); + + LLVM_DEBUG(llvm::dbgs() << __FILE__ << " START\n" << mod << "\n"); + + // Lower math.erf to xpu libdevice erf (ExternElementwiseOp). The + // subsequent ExternElementwiseOp vectorization step will rewrite + // "_ZN3xpu3erfEf" to its vectorized counterpart "_ZN3xpu4verfEDv16_f". + // + // FIXME(XPUTC-7517): Generalize this ad-hoc walker into a proper + // PatternRewriter-based rewrite (e.g. OpRewritePattern) and + // table-drive the math-dialect-op -> xpu-libdevice-symbol mapping + // (math.erf -> _ZN3xpu3erfEf, math.tan -> _ZN3xpu4tanfEf, ...). This + // will let us cover other math dialect ops (atan, isinf, isnan, rsqrt, + // tanh, ...) without copy-pasting walker blocks here, and will compose + // cleanly with the existing ExternElementwiseOp vectorization Case. + mod.walk([&](math::ErfOp erfOp) { + OpBuilder builder(erfOp); + auto newExtElemwiseOp = builder.create( + erfOp.getLoc(), erfOp.getResult().getType(), + ValueRange{erfOp.getOperand()}, /*libname=*/"", /*libpath=*/"", + /*symbol=*/"_ZN3xpu3erfEf", /*pure=*/true); + erfOp.getResult().replaceAllUsesWith(newExtElemwiseOp.getResult()); + erfOp.erase(); + }); + + // Maximum Fusion Online + // [cmpf, cmpf, ori, select] -> [fmax] + if (maximumFusion) { + mod.walk([&](arith::SelectOp selectOp) { doMaximumFusion(selectOp); }); + } + + // Compare Fusion + // [cmpf, extui] -> [vcmpf(castToI8=True)] + if (this->compareFusion) { + mod.walk([&](arith::ExtUIOp extUIOp) { doCompareExtUI8Fusion(extUIOp); }); + } + mod.walk( + [&](arith::TruncIOp truncIOp) { doCompareTruncI8Fusion(truncIOp); }); + + { + RewritePatternSet patterns(context); + patterns.add, + ConvertI1LogicOpToI8, + ConvertI1LogicOpToI8, BypassCmpIExtUI>( + context); + if (failed(applyPatternsGreedily(mod, std::move(patterns)))) + signalPassFailure(); + } + + // Cross-pass signal, replacing the `BF16ToFP32VecOpt` member the two + // compare fusions used to clear inside Vectorize. Vectorize consumes and + // *erases* the marker, so it never reaches the emitted IR -- which is what + // keeps the `boolfused` probe (the only one where a compare fusion fires) + // byte-identical across the split. + if (bf16ToFP32VecOptOff) { + mod->setAttr(kBF16ToFP32VecOptOffAttrName, UnitAttr::get(context)); + } + + LLVM_DEBUG(llvm::dbgs() << __FILE__ << " END\n" << mod << "\n"); + } + +private: + MLIRContext *context; + bool maximumFusion = true; + bool bf16ToFP32VecOptOff = false; +}; + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/OffsetAnalysis.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/OffsetAnalysis.cpp index ed171ac976..13ea42ea68 100644 --- a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/OffsetAnalysis.cpp +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/OffsetAnalysis.cpp @@ -8,6 +8,8 @@ #include "triton/Dialect/TritonXPU/IR/Dialect.h" #include "triton/Dialect/TritonXPU/Transforms/Passes.h" +#include + #define DEBUG_TYPE "tritonxpu-offset-analysis" namespace mlir { @@ -214,6 +216,11 @@ struct TritonXPUOffsetAnalysisPass return mockVals; }; + auto getNumProgramsMockVals = []() { + SmallVector mockVals(1, 1); + return mockVals; + }; + auto getThreadIdMockVals = []() { SmallVector mockVals(/*core_num*/ 64); std::iota(mockVals.begin(), mockVals.end(), 0); @@ -274,6 +281,11 @@ struct TritonXPUOffsetAnalysisPass SmallVector mockVals = getProgramIdMockVals(); mockDataItems.emplace_back(MockData(getProgramIdOp, 0, mockVals)); }) + .Case([&](auto getNumProgramsOp) { + SmallVector mockVals = getNumProgramsMockVals(); + mockDataItems.emplace_back( + MockData(getNumProgramsOp, 0, mockVals)); + }) .Case([&](auto gm2lmOp) { SmallVector mockVals = getGM2LMOpMockVals(); mockDataItems.emplace_back(MockData(gm2lmOp, 0, mockVals)); @@ -500,12 +512,16 @@ struct TritonXPUOffsetAnalysisPass for (size_t i = 0; i < srcShape.size(); ++i) { if (srcShape[i] != resShape[i]) { - if (srcShape[i] == 1) { // [1x1xf32 -> 1xNxf32] + if (srcShape[i] == 1) { // [1xNxf32 -> MxNxf32] or [Mx1xf32 -> MxNxf32] unsigned numElems = resShape[resShape.size() - 1]; - if (i == srcShape[srcShape.size() - 1]) - return SmallVector(numElems, op2OffsetVal[operandOp][0]); - else + if (i == srcShape.size() - 1) { + SmallVector res; + for (int v : op2OffsetVal[operandOp]) + res.append(numElems, v); + return res; + } else { return op2OffsetVal[operandOp]; + } } else { // [1x2xf32 -> 1xNxf32] llvm_unreachable("[broadcastOpCalFunc] Only support broadcast 1->N"); } @@ -625,7 +641,7 @@ struct TritonXPUOffsetAnalysisPass auto hasDynamicInput = [](Operation *op) -> bool { for (auto operand : op->getOperands()) { if (mlir::isa(operand)) { - continue; + return true; } auto operandOp = operand.getDefiningOp(); if (!operandOp) { @@ -680,6 +696,10 @@ struct TritonXPUOffsetAnalysisPass op2OffsetVal[getProgramIdOp] = getProgramIdOpCalFunc(getProgramIdOp, op2OffsetVal, mockVal); }) + .Case([&](auto getNumProgramsOp) { + auto mockVal = op2MockVal[getNumProgramsOp]; + op2OffsetVal[getNumProgramsOp] = SmallVector(1, mockVal); + }) .Case([&](auto xpuGm2lmOp) { auto mockVal = op2MockVal[xpuGm2lmOp]; op2OffsetVal[xpuGm2lmOp] = @@ -709,7 +729,20 @@ struct TritonXPUOffsetAnalysisPass makeRangeOpCalFunc(makeRangeOp, op2OffsetVal); }) .Case([&](auto splatOp) { - if (hasDynamicInput(splatOp)) { + auto operand = splatOp.getOperand(); + if (mlir::isa(operand)) { + bool usedInMul = false; + for (auto user : splatOp.getResult().getUsers()) { + if (isa(user)) { + usedInMul = true; + break; + } + } + if (usedInMul) { + findUnsupportedOp = true; + return; + } + } else if (hasDynamicInput(splatOp)) { findUnsupportedOp = true; return; } @@ -821,7 +854,7 @@ struct TritonXPUOffsetAnalysisPass LLVM_DEBUG(llvm::dbgs() << "[OffsetState]: The 0th Address Is Not the Beginning " "of the Bank.\n"); - fixedStride = -1; + fixedStride = INT32_MIN; return OffsetState::Unknown; } } @@ -839,7 +872,7 @@ struct TritonXPUOffsetAnalysisPass } else { LLVM_DEBUG(llvm::dbgs() << "[OffsetState]: Addresses Are Not in the Same Bank.\n"); - fixedStride = -1; + fixedStride = INT32_MIN; return OffsetState::Unknown; } } @@ -947,13 +980,30 @@ struct TritonXPUOffsetAnalysisPass int64_t currRowLen = 2; int64_t currRowStride = 1; bool isFirst = true; + const int64_t base = res[0]; + auto gcd = [](int64_t a, int64_t b) { + while (b != 0) { + int64_t t = b; + b = a % b; + a = t; + } + return a; + }; for (int64_t i = 2; i < res.size(); i++) { + if (res[i] < base) { + rowLen = -1; + return false; + } if (res[i] - res[i - 1] == 1) { currRowLen++; } else { currRowStride = res[i] - res[i - 1] + currRowLen - 1; if (currRowStride < 0) { - return false; + rowStride = -1; + rowLen = gcd(rowLen, currRowLen); + currRowLen = 1; + currRowStride = 1; + continue; } if (isFirst) { @@ -968,15 +1018,6 @@ struct TritonXPUOffsetAnalysisPass if (currRowStride != rowStride) { rowStride = -1; } - - auto gcd = [](int64_t a, int64_t b) { - while (b != 0) { - int t = b; - b = a % b; - a = t; - } - return a; - }; rowLen = gcd(rowLen, currRowLen); } currRowLen = 1; @@ -1006,6 +1047,39 @@ struct TritonXPUOffsetAnalysisPass return OffsetState::Unknown; } + template ::value, bool> = true> + int64_t findForcedRowLen(T memoryOp) { + Value ptr = memoryOp.getPtr(); + Operation *ptrOp = ptr.getDefiningOp(); + while (ptrOp && + (isa(ptrOp) || isa(ptrOp))) + ptrOp = ptrOp->getOperand(0).getDefiningOp(); + auto addPtrOp = dyn_cast_or_null(ptrOp); + if (!addPtrOp) + return -1; + Operation *offsetDefineOp = addPtrOp.getOperand(1).getDefiningOp(); + if (!offsetDefineOp) + return -1; + llvm::SetVector opChain; + getOpChainBwdBFS(opChain, offsetDefineOp); + for (auto *op : opChain) { + auto remOp = dyn_cast(op); + if (!remOp) + continue; + auto constOp = remOp.getRhs().getDefiningOp(); + if (!constOp) + continue; + int64_t remConst = 0; + if (auto splatAttr = dyn_cast(constOp.getValue())) + remConst = splatAttr.getSplatValue().getSExtValue(); + else if (auto intAttr = dyn_cast(constOp.getValue())) + remConst = intAttr.getInt(); + if (remConst > 1) + return remConst; + } + return -1; + } + // -1 for Unknown // 0 for DiscreteSame // 1 for Continuous @@ -1128,6 +1202,28 @@ struct TritonXPUOffsetAnalysisPass memoryStateTransfer(memoryState, allOffsetStateResult[token]); } + if (memoryState == OffsetState::Continuous && + isa(memoryOp)) { + for (auto *op : opChain) { + if (auto remOp = dyn_cast(op)) { + if (auto constOp = remOp.getRhs().getDefiningOp()) { + int64_t remConst = 0; + if (auto splatAttr = dyn_cast(constOp.getValue())) + remConst = splatAttr.getSplatValue().getSExtValue(); + else if (auto intAttr = dyn_cast(constOp.getValue())) + remConst = intAttr.getInt(); + if (remConst > 0 && remConst > (int64_t)numElems) { + rowLen = remConst; + rowStride = -1; + fixedStride = 1; + memoryState = OffsetState::LocallyContinuous; + break; + } + } + } + } + } + bool canBeLrie = findUserOp(memoryOp) && !findUserOp(memoryOp); if (atomicSim && !analysisFlag && canBeLrie) { @@ -1166,49 +1262,91 @@ struct TritonXPUOffsetAnalysisPass mod.walk([&](mlir::Operation *op) { op2Line[op] = line++; }); + auto recoverForcedOffsetState = + [&](auto memoryOp, bool handwritten, OffsetState &offsetState, + int32_t &opFixedStride, int64_t &opRowLen, int64_t &opRowStride, + int32_t &opLrie) { + if (!handwritten) + return; + if (offsetState == OffsetState::Continuous) { + if (opFixedStride == INT32_MIN) + opFixedStride = 1; + if (opLrie <= 0) + opLrie = 1; + return; + } + if (offsetState != OffsetState::LocallyContinuous || opRowLen > 0) + return; + int64_t modLen = findForcedRowLen(memoryOp); + if (modLen > 1) { + opRowLen = modLen; + opRowStride = -1; + opFixedStride = 1; + return; + } + OffsetState inferred = getOffsetState(memoryOp); + opFixedStride = fixedStride; + opRowLen = rowLen; + opRowStride = rowStride; + opLrie = lrie; + if (inferred != OffsetState::LocallyContinuous) + offsetState = inferred; + fixedStride = INT32_MIN; + rowLen = -1; + rowStride = -1; + }; + mod.walk([&](triton::xpu::GM2LMOp gm2lmOp) { if (dumpFlag) LLVM_DEBUG(llvm::dbgs() << "\n=======================================\n"); + bool handwritten = gm2lmOp.getHandwrittenOffsetState(); OffsetState offsetState = - gm2lmOp.getHandwrittenOffsetState() - ? static_cast(gm2lmOp.getOffsetState()) - : getOffsetState(gm2lmOp); + handwritten ? static_cast(gm2lmOp.getOffsetState()) + : getOffsetState(gm2lmOp); + int32_t opFixedStride = handwritten ? gm2lmOp.getFixedStride() : fixedStride; + int64_t opRowLen = handwritten ? gm2lmOp.getRowLen() : rowLen; + int64_t opRowStride = handwritten ? gm2lmOp.getRowStride() : rowStride; + int32_t opLrie = handwritten ? gm2lmOp.getLrie() : lrie; + recoverForcedOffsetState(gm2lmOp, handwritten, offsetState, opFixedStride, + opRowLen, opRowStride, opLrie); if (dumpFlag) { LLVM_DEBUG(llvm::dbgs() << "\n" << gm2lmOp << "\n[OffsetState]: " << offsetState << "\n=======================================\n"); } - // In case `fixedStride` being modified by cluster(s) whose - // OffsetState is Continuous. - if (offsetState == OffsetState::Discrete) { - fixedStride = -1; - } else if (offsetState == OffsetState::Unknown && - (fixedStride == 1 | fixedStride == 0)) { - // Multi Memory State Like (Unknown & Continuous) - fixedStride = -1; + if (!handwritten) { + if (offsetState == OffsetState::Discrete) { + opFixedStride = INT32_MIN; + } else if (offsetState == OffsetState::Unknown && + (opFixedStride == 1 | opFixedStride == 0)) { + opFixedStride = INT32_MIN; + } } OpBuilder builder(gm2lmOp); int32_t offsetStateInt = static_cast(offsetState); gm2lmOp->setAttr("offsetState", builder.getSI32IntegerAttr(offsetStateInt)); - gm2lmOp->setAttr("fixedStride", builder.getSI32IntegerAttr(fixedStride)); + gm2lmOp->setAttr("fixedStride", builder.getSI32IntegerAttr(opFixedStride)); gm2lmOp->setAttr("rowLen", builder.getIntegerAttr( - builder.getIntegerType(64, true), rowLen)); + builder.getIntegerType(64, true), opRowLen)); gm2lmOp->setAttr( "rowStride", - builder.getIntegerAttr(builder.getIntegerType(64, true), rowStride)); - gm2lmOp->setAttr("lrie", builder.getSI32IntegerAttr(lrie)); + builder.getIntegerAttr(builder.getIntegerType(64, true), opRowStride)); + gm2lmOp->setAttr("lrie", builder.getSI32IntegerAttr(opLrie)); auto loadOp = cast(gm2lmOp->getNextNode()); loadOp->setOperand(0, gm2lmOp); - loadOp->setAttr("stride", builder.getSI32IntegerAttr(fixedStride)); - loadOp->setAttr("isDiscrete", builder.getBoolAttr(offsetState == - OffsetState::Discrete)); - fixedStride = -1; // reset - rowLen = -1; - rowStride = -1; + loadOp->setAttr("stride", builder.getSI32IntegerAttr(opFixedStride)); + loadOp->setAttr("isDiscrete", + builder.getBoolAttr(offsetState == OffsetState::Discrete || + offsetState == OffsetState::LocallyScalar)); + if (!handwritten) { + fixedStride = INT32_MIN; + rowLen = -1; + rowStride = -1; + } findUnsupportedOp = false; }); @@ -1216,15 +1354,16 @@ struct TritonXPUOffsetAnalysisPass if (dumpFlag) LLVM_DEBUG(llvm::dbgs() << "\n=======================================\n"); + bool handwritten = lm2gmOp.getHandwrittenOffsetState(); OffsetState offsetState = - lm2gmOp.getHandwrittenOffsetState() - ? static_cast(lm2gmOp.getOffsetState()) - : getOffsetState(lm2gmOp); - // Only able to handle continuous and unknown cases. + handwritten ? static_cast(lm2gmOp.getOffsetState()) + : getOffsetState(lm2gmOp); if (offsetState != OffsetState::Continuous && offsetState != OffsetState::LocallyContinuous) { offsetState = OffsetState::Unknown; } + int64_t opRowLen = handwritten ? lm2gmOp.getRowLen() : rowLen; + int64_t opRowStride = handwritten ? lm2gmOp.getRowStride() : rowStride; if (dumpFlag) { LLVM_DEBUG(llvm::dbgs() << "\n" @@ -1232,18 +1371,19 @@ struct TritonXPUOffsetAnalysisPass << "\n=======================================\n"); } OpBuilder builder(lm2gmOp); - // offsetState = OffsetState::Continuous; int32_t offsetStateInt = static_cast(offsetState); lm2gmOp->setAttr("offsetState", builder.getSI32IntegerAttr(offsetStateInt)); lm2gmOp->setAttr("rowLen", builder.getIntegerAttr( - builder.getIntegerType(64, true), rowLen)); + builder.getIntegerType(64, true), opRowLen)); lm2gmOp->setAttr( "rowStride", - builder.getIntegerAttr(builder.getIntegerType(64, true), rowStride)); - findUnsupportedOp = false; // reset - rowLen = -1; - rowStride = -1; + builder.getIntegerAttr(builder.getIntegerType(64, true), opRowStride)); + findUnsupportedOp = false; + if (!handwritten) { + rowLen = -1; + rowStride = -1; + } }); mod.walk([&](triton::xpu::SM2GMOp sm2gmOp) { @@ -1268,45 +1408,53 @@ struct TritonXPUOffsetAnalysisPass if (dumpFlag) LLVM_DEBUG(llvm::dbgs() << "\n=======================================\n"); + bool handwritten = gm2lmOp.getHandwrittenOffsetState(); OffsetState offsetState = - gm2lmOp.getHandwrittenOffsetState() - ? static_cast(gm2lmOp.getOffsetState()) - : getOffsetState(gm2lmOp); + handwritten ? static_cast(gm2lmOp.getOffsetState()) + : getOffsetState(gm2lmOp); + int32_t opFixedStride = handwritten ? gm2lmOp.getFixedStride() : fixedStride; + int64_t opRowLen = handwritten ? gm2lmOp.getRowLen() : rowLen; + int64_t opRowStride = handwritten ? gm2lmOp.getRowStride() : rowStride; + int32_t opLrie = handwritten ? gm2lmOp.getLrie() : lrie; + recoverForcedOffsetState(gm2lmOp, handwritten, offsetState, opFixedStride, + opRowLen, opRowStride, opLrie); if (dumpFlag) { LLVM_DEBUG(llvm::dbgs() << "\n" << gm2lmOp << "\n[OffsetState]: " << offsetState << "\n=======================================\n"); } - // In case `fixedStride` being modified by cluster(s) whose - // OffsetState is Continuous. - if (offsetState == OffsetState::Discrete) { - fixedStride = -1; - } else if (offsetState == OffsetState::Unknown && - (fixedStride == 1 | fixedStride == 0)) { - // Multi Memory State Like (Unknown & Continuous) - fixedStride = -1; + if (!handwritten) { + if (offsetState == OffsetState::Discrete) { + opFixedStride = INT32_MIN; + } else if (offsetState == OffsetState::Unknown && + (opFixedStride == 1 | opFixedStride == 0)) { + opFixedStride = INT32_MIN; + } } OpBuilder builder(gm2lmOp); int32_t offsetStateInt = static_cast(offsetState); gm2lmOp->setAttr("offsetState", builder.getSI32IntegerAttr(offsetStateInt)); - gm2lmOp->setAttr("fixedStride", builder.getSI32IntegerAttr(fixedStride)); + gm2lmOp->setAttr("fixedStride", builder.getSI32IntegerAttr(opFixedStride)); gm2lmOp->setAttr("rowLen", builder.getIntegerAttr( - builder.getIntegerType(64, true), rowLen)); + builder.getIntegerType(64, true), opRowLen)); gm2lmOp->setAttr( "rowStride", - builder.getIntegerAttr(builder.getIntegerType(64, true), rowStride)); - gm2lmOp->setAttr("lrie", builder.getSI32IntegerAttr(lrie)); + builder.getIntegerAttr(builder.getIntegerType(64, true), opRowStride)); + gm2lmOp->setAttr("lrie", builder.getSI32IntegerAttr(opLrie)); auto loadOp = cast(gm2lmOp->getNextNode()); loadOp->setOperand(0, gm2lmOp); - loadOp->setAttr("stride", builder.getSI32IntegerAttr(fixedStride)); - loadOp->setAttr("isDiscrete", builder.getBoolAttr(offsetState == - OffsetState::Discrete)); - fixedStride = -1; // reset - rowLen = -1; - rowStride = -1; + loadOp->setAttr("stride", builder.getSI32IntegerAttr(opFixedStride)); + loadOp->setAttr("isDiscrete", + builder.getBoolAttr(offsetState == OffsetState::Discrete || + offsetState == OffsetState::LocallyScalar)); + if (!handwritten) { + fixedStride = INT32_MIN; + rowLen = -1; + rowStride = -1; + } findUnsupportedOp = false; }); diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/ScalarAnalysis.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/ScalarAnalysis.cpp new file mode 100644 index 0000000000..73fef0776e --- /dev/null +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/ScalarAnalysis.cpp @@ -0,0 +1,310 @@ +//===----------------------------------------------------------------------===// +// TritonXPUScalarAnalysis pass +// +// Drives the dataflow `ScalarAnalysis` (defined in xpu/include/Analysis/ +// ScalarAnalysis.h) to classify every SSA value as Scalar / VectorContig / +// VectorOther starting from `tt.make_range` and `tt.splat`. Then, for every +// `triton_xpu.gm2lm` / `triton_xpu.lm2gm` whose pointer-tensor is +// `splat(scalar_base) + Contig(stride=1)`, mark `offsetState = Continuous` +// and `handwrittenOffsetState = true` so the downstream `OffsetAnalysis` +// preserves the result. +//===----------------------------------------------------------------------===// + +#include "triton/Analysis/NewAnalysis/Utility.h" +#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h" +#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h" +#include "mlir/Analysis/DataFlow/SparseAnalysis.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Analysis/ScalarAnalysis.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/Transforms/Passes.h" + +#include "llvm/Support/Debug.h" +#include +#include +#include + +#define DEBUG_TYPE "tritonxpu-scalar-analysis" + +namespace mlir { +namespace triton { +namespace xpu { + +#define GEN_PASS_DEF_TRITONXPUSCALARANALYSIS +#include "triton/Dialect/TritonXPU/Transforms/Passes.h.inc" + +namespace { + +struct TritonXPUScalarAnalysisPass + : public impl::TritonXPUScalarAnalysisBase { + + using impl::TritonXPUScalarAnalysisBase< + TritonXPUScalarAnalysisPass>::TritonXPUScalarAnalysisBase; + + void runOnOperation() override { + ModuleOp mod = getOperation(); + + // 1. Run the dataflow analysis over each function body. + DataFlowSolver solver; + solver.load(); + solver.load(); + solver.load(); + if (failed(solver.initializeAndRun(mod))) { + signalPassFailure(); + return; + } + + // 2. Walk gm2lm / lm2gm and stamp Continuous + handwritten when the + // pointer's classification is exactly VectorContig(stride = 1). + auto isContig1 = [&](Value ptrTensor) -> bool { + auto *lattice = solver.lookupState< + dataflow::Lattice>(ptrTensor); + if (!lattice) + return false; + const ScalarValueState &v = lattice->getValue(); + return v.isContig() && v.stride == 1; + }; + + // Returns rowLen (>= 2) when the pointer-tensor is BlockContig with + // stride=1 AND it is a *genuine runtime gather* (its inter-row base comes + // from a `triton_xpu.load`, the k89 embedding case) AND it is *not + // provably unsafe* to mark as LocallyContinuous. + // + // The downstream lowering treats LocallyContinuous(rowLen=R, rowStride=-1) + // as "iterate R-lane blocks in linear order, row base is data-dependent". + // That is correct only for true gathers, where each R-lane row is still + // walked in element order and the inter-row jump is a runtime value. + // + // It is NOT correct for Block* patterns built purely from arithmetic on + // `arange`, e.g. the cat strided-copy offset `(idx/R)*S + (idx%R)` where + // the inter-row stride S is a compile-time constant != R. Such patterns + // can surface as BlockContig with an *unknown* blockStride after the + // multi-dim index math is combined, which would slip past a "blockStride + // known && != R" check. Requiring `blockFromLoad` excludes them: they are + // left to OffsetAnalysis, which computes the correct fixed rowStride. + // + // blockFromLoad: true when the BlockContig's inter-row base originates + // from a `triton_xpu.load` result (i.e., the row base address is truly + // data-dependent at runtime, not derivable from arithmetic on arange). + // + // In aggressive mode (second run after OffsetAnalysis), the blockFromLoad + // guard is skipped: if OffsetAnalysis already failed (offsetState=-1), + // we accept arithmetic-derived BlockContig and mark LocallyContinuous as + // a best-effort fallback. + // + // We still reject when blockStride is statically known and != rowLen + // (e.g. flip's -R) as an extra guard. + auto getBlockContigRowLen = [&](Value ptrTensor) -> int64_t { + auto *lattice = solver.lookupState< + dataflow::Lattice>(ptrTensor); + if (!lattice) + return 0; + const ScalarValueState &v = lattice->getValue(); + if (!v.isBlockContig() || v.stride != 1 || v.rowLen < 2) + return 0; + if (!aggressive && !v.blockFromLoad) + return 0; // not a genuine runtime gather (e.g. cat strided copy) + if (v.blockStrideKnown && v.blockStride != v.rowLen) + return 0; // provably unsafe (e.g. flip's -R) + return v.rowLen; + }; + + auto markContinuous = [&](Operation *op) { + OpBuilder b(op); + op->setAttr("offsetState", + b.getSI32IntegerAttr( + static_cast(OffsetState::Continuous))); + op->setAttr("handwrittenOffsetState", b.getBoolAttr(true)); + // Stamp the auxiliary attributes that downstream OffsetAnalysis would + // otherwise compute via the (now skipped) inference path. For the + // VectorContig(stride=1) case these values are statically known: + // fixedStride = 1 (contiguous element stride) + // lrie = 1 (no runs of identical elements) + // rowLen / rowStride / tensorColSize remain -1 (only meaningful for + // LocallyContinuous / multi-row patterns), matching what + // `getOffsetState` produces for plain Continuous. + if (isa(op)) { + op->setAttr("fixedStride", b.getSI32IntegerAttr(1)); + op->setAttr("lrie", b.getSI32IntegerAttr(1)); + } + }; + + // Stamp LocallyContinuous with the detected rowLen. rowStride is left + // at -1 because the inter-row stride is data-dependent (e.g. comes from + // a gather index). The downstream LLVM lowering's LocallyContinuous path + // handles rowStride==-1 (row-by-row DMA) and even upgrades to Continuous + // when rowLen % numElems == 0. + auto markLocallyContinuous = [&](Operation *op, int64_t rowLen) { + OpBuilder b(op); + op->setAttr("offsetState", + b.getSI32IntegerAttr(static_cast( + OffsetState::LocallyContinuous))); + op->setAttr("handwrittenOffsetState", b.getBoolAttr(true)); + op->setAttr("rowLen", b.getIntegerAttr( + b.getIntegerType(64, /*isSigned=*/true), rowLen)); + op->setAttr("rowStride", b.getIntegerAttr( + b.getIntegerType(64, /*isSigned=*/true), -1)); + }; + + // Returns rowLen (>= sizePerCore) when the pointer-tensor is BlockScalar, + // meaning every `rowLen` lanes share the same address. Only mark when + // rowLen >= sizePerCore so that within each core the address is uniform. + auto getBlockScalarRowLen = [&](Value ptrTensor) -> int64_t { + auto *lattice = solver.lookupState< + dataflow::Lattice>(ptrTensor); + if (!lattice) + return 0; + const ScalarValueState &v = lattice->getValue(); + if (!v.isBlockScalar() || v.rowLen < 2) + return 0; + auto ty = dyn_cast(ptrTensor.getType()); + if (!ty) + return 0; + auto enc = dyn_cast(ty.getEncoding()); + if (!enc) + return 0; + int64_t sizePerCore = enc.getSizePerCore()[0]; + if (sizePerCore <= 0 || v.rowLen % sizePerCore != 0) + return 0; + return v.rowLen; + }; + + auto markDiscreteSame = [&](Operation *op, int64_t /*rowLen*/) { + OpBuilder b(op); + op->setAttr("offsetState", + b.getSI32IntegerAttr( + static_cast(OffsetState::DiscreteSame))); + op->setAttr("handwrittenOffsetState", b.getBoolAttr(true)); + op->setAttr("fixedStride", b.getSI32IntegerAttr(0)); + op->setAttr("rowLen", b.getIntegerAttr( + b.getIntegerType(64, /*isSigned=*/true), -1)); + op->setAttr("rowStride", b.getIntegerAttr( + b.getIntegerType(64, /*isSigned=*/true), -1)); + }; + + // Returns rowLen when the pointer-tensor is BlockScalar but each core spans + // *multiple whole blocks* (rowLen < sizePerCore && sizePerCore % rowLen == 0). + // DiscreteSame requires one uniform address per core (rowLen >= sizePerCore); + // this complementary "locally scalar" case keeps the per-block scalar address + // structure so the gm2lm can issue one scalar DMA per block instead of + // falling back to the per-element Unknown gather. + auto getLocalScalarRowLen = [&](Value ptrTensor) -> int64_t { + auto *lattice = solver.lookupState< + dataflow::Lattice>(ptrTensor); + if (!lattice) + return 0; + const ScalarValueState &v = lattice->getValue(); + if (!v.isBlockScalar() || v.rowLen < 2) + return 0; + auto ty = dyn_cast(ptrTensor.getType()); + if (!ty) + return 0; + auto enc = dyn_cast(ty.getEncoding()); + if (!enc) + return 0; + int64_t sizePerCore = enc.getSizePerCore()[0]; + if (sizePerCore <= 0) + return 0; + if (v.rowLen >= sizePerCore) + return 0; // DiscreteSame territory (handled by getBlockScalarRowLen) + if (sizePerCore % v.rowLen != 0) + return 0; // need whole blocks within a core + return v.rowLen; + }; + + auto markLocallyScalar = [&](Operation *op, int64_t rowLen) { + OpBuilder b(op); + op->setAttr("offsetState", + b.getSI32IntegerAttr( + static_cast(OffsetState::LocallyScalar))); + op->setAttr("handwrittenOffsetState", b.getBoolAttr(true)); + // INT32_MIN keeps the consuming load off the DiscreteSame(stride==0) path; + // OffsetAnalysis stamps isDiscrete=true so it uses the per-lane LM read. + op->setAttr("fixedStride", b.getSI32IntegerAttr(INT32_MIN)); + op->setAttr("rowLen", b.getIntegerAttr( + b.getIntegerType(64, /*isSigned=*/true), rowLen)); + op->setAttr("rowStride", b.getIntegerAttr( + b.getIntegerType(64, /*isSigned=*/true), -1)); + }; + + // In aggressive mode, only update ops where offsetState is still -1 + // (i.e., OffsetAnalysis failed to determine the state). + auto shouldSkipInAggressive = [&](Operation *op) -> bool { + if (!aggressive) + return false; + auto attr = op->getAttrOfType("offsetState"); + if (!attr) + return false; // no attr means not yet processed, ok to mark + return attr.getValue().getSExtValue() != -1; // already has valid state, skip + }; + + mod.walk([&](Operation *op) { + if (auto gm2lm = dyn_cast(op)) { + if (gm2lm.getHandwrittenOffsetState()) + return; + if (shouldSkipInAggressive(op)) + return; + if (isContig1(gm2lm.getPtr())) { + markContinuous(op); + return; + } + if (int64_t rl = getBlockContigRowLen(gm2lm.getPtr())) { + markLocallyContinuous(op, rl); + return; + } + if (int64_t rl = getBlockScalarRowLen(gm2lm.getPtr())) { + markDiscreteSame(op, rl); + return; + } + if (int64_t rl = getLocalScalarRowLen(gm2lm.getPtr())) { + markLocallyScalar(op, rl); + return; + } + return; + } + if (auto lm2gm = dyn_cast(op)) { + if (lm2gm.getHandwrittenOffsetState()) + return; + if (shouldSkipInAggressive(op)) + return; + if (isContig1(lm2gm.getPtr())) { + markContinuous(op); + return; + } + // dont mark locally continuous or DiscreteSame for lm2gm + // if (int64_t rl = getBlockContigRowLen(lm2gm.getPtr())) { + // markLocallyContinuous(op, rl); + // return; + // } + // if (int64_t rl = getBlockScalarRowLen(lm2gm.getPtr())) { + // markDiscreteSame(op, rl); + // } + return; + } + }); + + // 3. Debug mode: stamp every tensor result with its lattice state. + // Enable via env TRITONXPU_SCALAR_ANALYSIS_DEBUG=1 + if (std::getenv("TRITONXPU_SCALAR_ANALYSIS_DEBUG")) { + mod.walk([&](Operation *op) { + for (Value res : op->getResults()) { + auto *lattice = solver.lookupState< + dataflow::Lattice>(res); + if (!lattice) + continue; + const ScalarValueState &v = lattice->getValue(); + std::string desc; + llvm::raw_string_ostream os(desc); + v.print(os); + OpBuilder b(op); + op->setAttr("scalar_state", b.getStringAttr(desc)); + } + }); + } + } +}; + +} // namespace +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/TLELegalize.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/TLELegalize.cpp new file mode 100644 index 0000000000..5ad3d0fa76 --- /dev/null +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/TLELegalize.cpp @@ -0,0 +1,105 @@ +//===----------------------------------------------------------------------===// +// TODO[dyq]: Pass Description +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Dominance.h" +#include "mlir/Support/LLVM.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/Transforms/Passes.h" + +#include "mlir/Analysis/TopologicalSortUtils.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/IRMapping.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "tritonxpu-tle-legalize" + +namespace mlir { +namespace triton { +namespace xpu { + +#define GEN_PASS_DEF_TRITONXPUTLELEGALIZE +#include "triton/Dialect/TritonXPU/Transforms/Passes.h.inc" + +struct TritonXPUTLELegalizePass + : public impl::TritonXPUTLELegalizeBase { + + using impl::TritonXPUTLELegalizeBase< + TritonXPUTLELegalizePass>::TritonXPUTLELegalizeBase; + + TritonXPUTLELegalizePass() = default; + + void runOnOperation() override { + mlir::MLIRContext *context = &getContext(); + mlir::ModuleOp m = getOperation(); + + // --- TLE reduce lowering --------------------------------------------- + // The TLE pipeline does not run the normal tritonxpu-legalize pass, so + // tt.reduce is never turned into triton_xpu.reduce (that transform lives in + // Legalize.cpp and is entangled with core-tiling loops). For the current + // TLE model (fixed [1, BLOCK] tile, no core-tiling), a whole tile lives in + // one iteration, so we can convert tt.reduce -> triton_xpu.reduce directly + // with loopNum=1 and loopIndex=0. + SmallVector ttReduceOps; + m.walk([&](triton::ReduceOp op) { ttReduceOps.push_back(op); }); + for (auto reduceOp : ttReduceOps) { + OpBuilder builder(reduceOp); + auto loc = reduceOp->getLoc(); + Value loopIndex = builder.create( + loc, builder.getI32Type(), builder.getI32IntegerAttr(0)); + auto newReduceOp = builder.create( + loc, reduceOp->getResultTypes(), reduceOp.getSrcs(), + reduceOp.getAxis(), /*loopNum=*/1, loopIndex); + auto &newCombineOp = newReduceOp.getCombineOp(); + builder.cloneRegionBefore(reduceOp.getCombineOp(), newCombineOp, + newCombineOp.end()); + // tt.reduce.return -> triton_xpu.reduce.return inside the combine region. + for (auto &opInCombine : llvm::make_early_inc_range(newCombineOp.getOps())) { + if (auto redReturnOp = + dyn_cast(&opInCombine)) { + OpBuilder retBuilder(redReturnOp); + auto newRedReturnOp = retBuilder.create( + redReturnOp.getLoc(), redReturnOp.getOperands()); + redReturnOp->replaceAllUsesWith(newRedReturnOp->getResults()); + redReturnOp.erase(); + } + } + reduceOp->replaceAllUsesWith(newReduceOp->getResults()); + reduceOp->erase(); + } + + // Count reduce/scan ops for helper id assignment below. (TLE elementwise + // kernels have none; kept so ReduceOpHelper/ScanLoweringHelper stay valid.) + unsigned reduceId = 0; + unsigned reduceNum = 0; + unsigned scanId = 0; + unsigned scanNum = 0; + m.walk([&](triton::xpu::ReduceOp) { reduceNum++; }); + m.walk([&](triton::xpu::ScanOp) { scanNum++; }); + + // Set ReduceOpHelper + m.walk([&](triton::xpu::ReduceOp redOp) { + ReduceOpHelper helper(redOp); + helper.setReduceId(reduceId); + helper.setReduceNum(reduceNum); + reduceId++; + }); + + // Set ScanLoweringHelper + m.walk([&](triton::xpu::ScanOp scanOp) { + ScanLoweringHelper helper(scanOp); + helper.setScanId(scanId); + helper.setScanNum(scanNum); + scanId++; + }); + + } +}; + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/TileAnalysisPass.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/TileAnalysisPass.cpp new file mode 100644 index 0000000000..ebc8fd0242 --- /dev/null +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/TileAnalysisPass.cpp @@ -0,0 +1,137 @@ +//===----------------------------------------------------------------------===// +// tritonxpu-tile-analysis -- report-only pass over TileAnalysis. +// +// Gives the pressure measurement a pipeline slot of its own, so the transform +// that consumes it can eventually be handed the result instead of recomputing +// it (redesign-v2.md §4.2 step 1.3, contract in 1.6). Byte equivalence of the +// emitted code is its exit gate: the only thing it ever writes is the removal of +// the keys the M6 probe left behind, and the probe is off by default. +// +// What it can and cannot see: the geometry and the block-wide pressure are +// readable from the IR alone, but the *per-tree* pressure UnrollControl also +// uses comes from the op tree that pass builds while walking (getDAG / +// getOpChainBwd), which is not reproducible here. Reporting the block half is +// still the point of interest -- `decideIterNum` takes +// max(treeP.vecPeak, blockP.vecPeak), so the gap between the two numbers says +// which half dominates at each site. Measured, it goes both ways: on +// layernorm's reduce-for site the tree is 48 against a block of 24, while the +// pointwise sites match exactly. So the block figure is not an upper bound and +// must not be treated as one. +// +// One block-granularity trap the report makes visible: block-wide minVecWidth +// can be 1 (softmax) where the per-tree value is unconstrained, and since +// isLegalIterNum requires minVecWidth % iterNum == 0, that collapses the legal +// set to {1}. Widening the *target* to block scope is not the same as widening +// minVecWidth to block scope. +//===----------------------------------------------------------------------===// + +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Analysis/TileAnalysis.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/Transforms/Passes.h" + +#include "triton/Tools/Sys/GetEnv.hpp" + +#define DEBUG_TYPE "tritonxpu-tile-analysis" + +namespace mlir { +namespace triton { +namespace xpu { + +#define GEN_PASS_DEF_TRITONXPUTILEANALYSIS +#include "triton/Dialect/TritonXPU/Transforms/Passes.h.inc" + +namespace { + +// Same legality test as UnrollControl::isLegalIterNum. Duplicated rather than +// shared because sharing it means moving it into the analysis unit, which is a +// logic move and belongs to a later step -- this pass is only allowed to +// observe. When the predicate does move, this copy goes away. +bool isLegalIterNum(int64_t iterNum, int64_t numCol, int64_t widthPerCore, + int64_t minVecWidth) { + return iterNum >= 1 && numCol % iterNum == 0 && widthPerCore % iterNum == 0 && + (minVecWidth == 0 || minVecWidth % iterNum == 0); +} + +} // namespace + +struct TritonXPUTileAnalysisPass + : public impl::TritonXPUTileAnalysisBase { + +public: + using impl::TritonXPUTileAnalysisBase< + TritonXPUTileAnalysisPass>::TritonXPUTileAnalysisBase; + + TritonXPUTileAnalysisPass() = default; + TritonXPUTileAnalysisPass(unsigned vrfBudget) { this->vrfBudget = vrfBudget; } + + void report(const char *site, Operation *op, Type valTy) { + auto tensorTy = dyn_cast(valTy); + if (!tensorTy) + return; + int64_t numCol = tensorTy.getShape().back(); + int64_t widthPerCore = numCol, coresPerGroup = 1; + if (auto clusterEncoding = getClusterLayout(tensorTy)) { + coresPerGroup = clusterEncoding.getCoresPerGroup().back(); + widthPerCore = clusterEncoding.getSizePerCore().back(); + } + + RegPressure p; + getBlockRegPressure(getOperation(), op, p); + + int64_t target = + this->vrfBudget > 0 ? ceil(p.vecPeak, this->vrfBudget) : 1; + target = std::max(target, 1); + + // The legal set, printed in full: its sparseness is the reason step 3.6 + // exists, and a report that only gave the chosen factor would hide it. + std::string legal; + for (int64_t iterNum = 1; iterNum <= widthPerCore; ++iterNum) + if (isLegalIterNum(iterNum, numCol, widthPerCore, p.minVecWidth)) + legal += (legal.empty() ? "" : ",") + std::to_string(iterNum); + + StringRef kernel = ""; + if (auto funcOp = op->getParentOfType()) + kernel = funcOp.getName(); + + llvm::errs() << "[TileAnalysis] " << kernel << " site=" << site + << " numCol=" << numCol << " widthPerCore=" << widthPerCore + << " coresPerGroup=" << coresPerGroup + << " blockVecPeak=" << p.vecPeak + << " blockVecTotal=" << p.vecTotal + << " blockScalarPeak=" << p.scalarPeak + << " maxVecWidth=" << p.maxVecWidth + << " minVecWidth=" << p.minVecWidth + << " budget=" << this->vrfBudget << " blockTarget=" << target + << " legal={" << legal << "}\n"; + } + + void runOnOperation() override { + ModuleOp mod = getOperation(); + + // The consumer end of the M6 contract (step 1.6). This pass sits immediately + // before `tritonxpu-unroll-control`, which is where M4/M5 will read the plan, + // so it is the right place to ask whether either key still finds its root. + // Runs before the report gate and on its own switch, because it also erases + // what the probe wrote. (It is registered under `isCloseUnrollControl`, so + // with unroll control off the probe simply produces no report -- the writes + // are gated on the same env var, so nothing is left behind either way.) + tilePlanCheck(mod); + + // Cheap by default: the walk below is pure measurement, so it is only worth + // paying for when someone is reading the report. + if (!mlir::triton::tools::getBoolEnv("TRITONXPU_TILE_REPORT")) + return; + + mod.walk([&](triton::xpu::StoreOp storeOp) { + report("pointwise", storeOp, storeOp.getValue().getType()); + }); + mod.walk([&](triton::xpu::ReduceOp reduceOp) { + report("reduce", reduceOp, reduceOp.getInputTypes()[0]); + }); + } +}; + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/UnrollControl.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/UnrollControl.cpp index 760514f0ea..544f4dfa11 100644 --- a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/UnrollControl.cpp +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/UnrollControl.cpp @@ -1,5 +1,8 @@ #include "mlir/IR/IRMapping.h" #include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Tools/Sys/GetEnv.hpp" +#include "triton/Analysis/TileAnalysis.h" +#include "triton/Analysis/TileDecision.h" #include "triton/Dialect/TritonXPU/IR/Dialect.h" #include "triton/Dialect/TritonXPU/Transforms/Passes.h" @@ -15,11 +18,13 @@ namespace xpu { template struct COMOp; #define COMOP(SrcType, DstType) \ - template <> struct COMOp { \ - typedef DstType type; \ - }; + template <> struct COMOp { typedef DstType type; }; COMOP(arith::AddFOp, triton::xpu::VvaddFOp); +// subf/divf only reach a combine region on the region-interpreting path +// (TRITONXPU_REDUCE_REGION); COMBINE_BINARY_OP does not list them. +COMOP(arith::SubFOp, triton::xpu::VvsubFOp); +COMOP(arith::DivFOp, triton::xpu::VvdivFOp); COMOP(arith::MulFOp, triton::xpu::VvmulFOp); COMOP(arith::MaxNumFOp, triton::xpu::VvmaxNumFOp); COMOP(arith::MinNumFOp, triton::xpu::VvminNumFOp); @@ -35,13 +40,60 @@ struct TritonXPUUnrollControl TritonXPUUnrollControl>::TritonXPUUnrollControlBase; TritonXPUUnrollControl() = default; - TritonXPUUnrollControl(unsigned bufferSize, unsigned coreNum, - unsigned unrollNum) { + TritonXPUUnrollControl(unsigned bufferSize, unsigned coreNum, unsigned unrollNum) { this->bufferSize = bufferSize; this->coreNum = coreNum; this->unrollNum = unrollNum; } + // Enabled by TRITONXPU_UNROLL_DRYRUN: report the tile factor a + // register-budget model would pick, without changing the IR. + bool dryRun = false; + + // Set when a legality constraint dictates the factor, in which case the + // budget model must not override it. `pinReason` names the constraint. + static constexpr StringLiteral kUnrollLoopAttr = "triton_xpu.unroll_loop"; + int64_t pinnedUnrollNum = 0; + const char *pinReason = ""; + // Backs `pinReason` when the knob renames it; a plain literal otherwise. + std::string pinReasonStorage; + + // Step 3.5b: make the two pinned constants reachable from a knob, so they can + // be shown to be wrong. `pinUnrollNum` < 0 keeps the constant (the default, so + // artifacts stay byte-identical), 0 drops the pin and lets the pressure model + // decide, > 0 overrides it so its time curve can be swept. + // + // Both numbers move together on purpose: the pin's factor equals the legacy + // one only while `unrollNum == pinnedUnrollNum` (`findings.md` §1.30 proves + // the identity via `getNumUnroll`), so splitting them would make a sweep + // measure two things at once. + // + // Dropping a pin is not silent: the site then goes through the model, which + // remarks its own decision. `emitRemark` alone is not enough here -- nothing + // in this pipeline installs a handler for it, so those remarks never reach + // the log -- hence the `[PinKnob]` line, printed whenever the knob is doing + // something (`pinUnrollNum >= 0`). A default run prints nothing. + void applyPin(ModuleOp m, int64_t constant, const char *reason) { + if (this->pinUnrollNum == 0) { + m->emitRemark("[UnrollControl] pin " + std::string(reason) + + " suppressed by pin-unroll-num=0 (constant was " + + std::to_string(constant) + "); the model decides instead"); + llvm::errs() << "[PinKnob] pin=" << reason << " constant=" << constant + << " action=dropped(model-decides)\n"; + return; + } + int64_t value = this->pinUnrollNum > 0 ? this->pinUnrollNum : constant; + this->unrollNum = value; + pinnedUnrollNum = value; + pinReasonStorage = reason; + if (value != constant) + pinReasonStorage += ":pin-override=" + std::to_string(value); + pinReason = pinReasonStorage.c_str(); + if (this->pinUnrollNum > 0) + llvm::errs() << "[PinKnob] pin=" << reason << " constant=" << constant + << " action=override unrollNum=" << value << "\n"; + } + template static decltype(auto) createCombineVectorizedOp(T op) { OpBuilder builder(op); return builder.create::type>( @@ -51,7 +103,8 @@ struct TritonXPUUnrollControl void processOpVecTy(ModuleOp &m) { m.walk([&](Operation *op) { TypeSwitch(op) - .Case([&](auto combineBinaryOp) { + .Case( + [&](auto combineBinaryOp) { if (auto tensorTy = dyn_cast( combineBinaryOp.getResult().getType())) { if (isa(getElementTypeOrSelf(tensorTy))) { @@ -61,6 +114,19 @@ struct TritonXPUUnrollControl } } }) + .Case([&](auto selectOp) { + if (auto tensorTy = dyn_cast( + selectOp.getResult().getType())) { + if (isa(getElementTypeOrSelf(tensorTy))) { + OpBuilder builder(selectOp); + auto vecOp = builder.create( + selectOp.getLoc(), tensorTy, selectOp.getCondition(), + selectOp.getTrueValue(), selectOp.getFalseValue()); + selectOp.replaceAllUsesWith(vecOp.getResult()); + selectOp.erase(); + } + } + }) .Case([&](auto cmpFOp) { if (auto tensorTy = dyn_cast(cmpFOp.getResult().getType())) { @@ -326,6 +392,513 @@ struct TritonXPUUnrollControl return numUnroll; } + //===--------------------------------------------------------------------===// + // Dry-run register-pressure model. + // + // The current tile factor comes from `unrollNum`, a global knob with no + // notion of how many registers the segment actually needs. The code below + // computes what a budget-driven forward decision *would* pick and reports + // the delta. It does not touch the IR. + //===--------------------------------------------------------------------===// + + // Reports the gate that actually decides whether tiling happens at all: + // `numCol > numUnroll && numCol % numUnroll == 0`. A factor that does not + // divide the width silently disables tiling, which is why unrollNum + // 3/5/6/7/8 all produce identical code today. + void reportGate(const char *site, Operation *op, int64_t numCol, + int64_t numUnroll) { + StringRef kernel = ""; + if (auto funcOp = op->getParentOfType()) + kernel = funcOp.getName(); + const char *why = numCol <= numUnroll ? "skip:factor>=width" + : (numCol % numUnroll) ? "skip:width%factor!=0" + : "tile"; + llvm::errs() << "[UnrollControl][gate] " << kernel << " site=" << site + << " numCol=" << numCol << " numUnroll=" << numUnroll + << " unrollNum=" << this->unrollNum << " -> " << why << "\n"; + } + + void reportDryRun(const char *site, Operation *insertPt, + const SetVector &unrollOpTree, Type valTy, + int64_t numCol, int64_t numUnroll, int64_t iterNum) { + RegPressure p; + getRegPressure(getOperation(), unrollOpTree, p); + int64_t peakVRegs = p.vecPeak, totalVRegs = p.vecTotal; + + int64_t coresPerGroup = 1, widthPerCore = numCol; + if (auto tensorTy = dyn_cast(valTy)) { + if (auto clusterEncoding = getClusterLayout(tensorTy)) { + coresPerGroup = clusterEncoding.getCoresPerGroup().back(); + widthPerCore = clusterEncoding.getSizePerCore().back(); + } + } + + // Smallest trip count that fits the budget. + int64_t iterModel = + this->vrfBudget > 0 ? ceil(peakVRegs, this->vrfBudget) : 1; + iterModel = std::max(iterModel, 1); + // A trip count that does not divide the width makes setTensorType round + // the tile up, so the tiling is silently dropped (unrollNum 3/5/6/7/8 all + // emit byte-identical asm today). Snap up to the next legal divisor. + int64_t iterLegal = iterModel; + while (iterLegal < numCol && + (numCol % iterLegal || widthPerCore % iterLegal)) + ++iterLegal; + if (numCol % iterLegal || widthPerCore % iterLegal) + iterLegal = 1; + + StringRef kernel = ""; + if (auto funcOp = insertPt->getParentOfType()) + kernel = funcOp.getName(); + + llvm::errs() << "[UnrollControl][dry-run] " << kernel << " site=" << site + << " numCol=" << numCol << " widthPerCore=" << widthPerCore + << " coresPerGroup=" << coresPerGroup + << " treeOps=" << unrollOpTree.size() + << " peakVRegs=" << peakVRegs << " totalVRegs=" << totalVRegs + << " scalarPeak=" << p.scalarPeak + << " scalarTotal=" << p.scalarTotal + << " budget=" << this->vrfBudget + << " | now: unrollNum=" << this->unrollNum + << " numUnroll=" << numUnroll << " iterNum=" << iterNum + << " | model: iterNum=" << iterLegal << " (raw " << iterModel + << ") unrollNum=" + << ceil(numCol, iterLegal * coresPerGroup) + << (iterLegal == iterNum ? " SAME" : " DIFF") << "\n"; + } + + //===--------------------------------------------------------------------===// + // Forward tile decision. + // + // Instead of asking "does the global unrollNum happen to divide the width", + // decide the trip count from the register pressure of the segment and then + // intersect it with the trip counts the rewrite can actually express. + //===--------------------------------------------------------------------===// + + void getTileGeometry(Type valTy, int64_t numCol, int64_t &widthPerCore, + int64_t &coresPerGroup) { + widthPerCore = numCol; + coresPerGroup = 1; + if (auto tensorTy = dyn_cast(valTy)) { + if (auto clusterEncoding = getClusterLayout(tensorTy)) { + coresPerGroup = clusterEncoding.getCoresPerGroup().back(); + widthPerCore = clusterEncoding.getSizePerCore().back(); + } + } + } + + // setTensorType/createEncoding slice both the shape and sizePerCore with + // ceil(). A trip count that does not divide them evenly makes the last + // iteration run off the end, so only exact divisors are expressible. + // `minVecWidth` is the narrowest vector row in the tree, in slots, or 0 when + // the tree holds no vector value. It has to divide the trip count too: since + // M3.2 the type this geometry is read from can be the *scalar* side of a + // vector->scalar boundary (the unpack in front of a scalar reduce), whose row + // is vecSize times wider. Slicing by a factor the vector row cannot express + // saturates it at one slot while the scalar row keeps dividing, and the two + // then cover a different number of lanes per iteration. + bool isLegalIterNum(int64_t iterNum, int64_t numCol, int64_t widthPerCore, + int64_t minVecWidth = 0) { + return iterNum >= 1 && numCol % iterNum == 0 && + widthPerCore % iterNum == 0 && + (minVecWidth == 0 || minVecWidth % iterNum == 0); + } + + // Does the segment straddle a vector->scalar boundary? Only there do the two + // representations of one row -- V vector slots on one side, V * vecSize + // scalars on the other -- get sliced by the same factor while being counted + // in different units, which is what makes the vector row a legality bound + // rather than just the widest thing in the tree. + bool hasVecScalarBoundary(const SetVector &unrollOpTree) { + return false; + } + + // Is this the LM buffer tritonxpu-alloca attached to a vector<->scalar + // boundary? Recognised by its users rather than by a flag on the alloca, + // because Alloca.cpp creates one buffer per pack/unpack and nothing else + // ever consumes it. Conservative on purpose: a buffer shared with a store + // pointer is not treated as a boundary buffer. + bool isBoundaryBuffer(Operation *op) { + if (!isa(op)) + return false; + if (op->use_empty()) + return false; + return false; + } + + // Is there anything for the model to pick besides "do not tile"? Used at the + // points that decide whether to collect an unroll tree at all. + // True once the op sits inside a tile loop this pass created. + bool inUnrollLoop(Operation *op) { + for (auto *parent = op->getParentOp(); parent; + parent = parent->getParentOp()) + if (isa(parent) && parent->hasAttr(kUnrollLoopAttr)) + return true; + return false; + } + + bool canTile(Operation *op, Type valTy, int64_t numCol, int64_t numUnroll) { + if (this->budgetTiling && inUnrollLoop(op)) + return false; + if (!this->budgetTiling) + return numCol > numUnroll && numCol % numUnroll == 0; + int64_t widthPerCore = 1, coresPerGroup = 1; + getTileGeometry(valTy, numCol, widthPerCore, coresPerGroup); + for (int64_t iterNum = 2; iterNum <= widthPerCore; ++iterNum) + if (isLegalIterNum(iterNum, numCol, widthPerCore)) + return true; + return false; + } + + void remarkDecision(const char *site, Operation *insertPt, int64_t numCol, + int64_t widthPerCore, int64_t peakVRegs, int64_t target, + int64_t chosen, int64_t maxLegal, const char *why, + int64_t scalarPeak = -1, int64_t minVecWidth = -1, + const triton::xpu::Decision *decision = nullptr) { + std::string msg; + llvm::raw_string_ostream os(msg); + os << "[UnrollControl] site=" << site << " numCol=" << numCol + << " widthPerCore=" << widthPerCore << " minVecWidth=" << minVecWidth + << " peakVRegs=" << peakVRegs + << " scalarPeak=" << scalarPeak << " budget=" << this->vrfBudget + << " target=" << target << " maxLegal=" << maxLegal + << " -> iterNum=" << chosen << " (" << why << ")"; + // Every criterion that took part has to be visible, feasible or not: a + // criterion nobody can read is a criterion nobody can falsify. + if (decision) { + for (const auto &t : decision->perTierTrace) { + os << " [tier" << t.tier << ":" << t.name << " cands=" << t.candidatesIn + << "->" << t.candidatesOut; + if (t.chosenCost) + os << " cost=" << *t.chosenCost; + if (!t.why.empty()) + os << " veto=" << t.why; + os << "]"; + } + } + insertPt->emitRemark(msg); + if (dryRun) + llvm::errs() << "[UnrollControl][decide] " << msg << "\n"; + } + + // Step 3.4 groundwork, report-only. `peakVRegs` fed to the decision is + // max(treeP.vecPeak, blockP.vecPeak) and 3.4 replaces the block half with a + // per-segment peak. §1.7 measured the gap between the two halves going *both* + // ways (layernorm's reduce-for site is tree 48 against block 24, its + // pointwise sites are equal), so `blockP` is not an upper bound and the swap + // cannot be done blind -- the first step is to tell per site which half is + // in charge, and whether dropping the block term would move the factor at + // all. + // + // The replay uses the tree peak as the stand-in for a segment peak. That is + // the loosest end of the range 3.4 can land in: a segment is a superset of + // one op tree and a subset of the block, so the real per-segment peak sits + // between these two numbers. A site where the tree-only replay picks the same + // factor therefore cannot be moved by 3.4 in the loosening direction either, + // which is what makes this a usable filter rather than a guess. + void reportSegDominance(const char *site, Operation *insertPt, + const RegPressure &treeP, const RegPressure &blockP, + const triton::xpu::TileContext &ctx, + llvm::ArrayRef candidates, + const triton::xpu::Decision &decision) { + if (!std::getenv("TRITONXPU_TILE_REPORT")) + return; + triton::xpu::TileContext altCtx = ctx; + altCtx.peakVRegs = treeP.vecPeak; + altCtx.maxVecWidth = treeP.maxVecWidth; + triton::xpu::Decision alt = triton::xpu::TileDecider().decide(candidates, altCtx); + StringRef kernel = ""; + if (auto funcOp = insertPt->getParentOfType()) + kernel = funcOp.getName(); + llvm::errs() << "[SegDominance] " << kernel << " site=" << site + << " treePeak=" << treeP.vecPeak + << " blockPeak=" << blockP.vecPeak << " dominant=" + << (treeP.vecPeak > blockP.vecPeak + ? "tree" + : (treeP.vecPeak == blockP.vecPeak ? "equal" : "block")) + << " treeMaxVecWidth=" << treeP.maxVecWidth + << " blockMaxVecWidth=" << blockP.maxVecWidth + << " target=" << triton::xpu::vrfBudgetTarget(ctx) + << " treeOnlyTarget=" << triton::xpu::vrfBudgetTarget(altCtx) + << " iterNum=" << decision.iterNum + << " treeOnlyIterNum=" << alt.iterNum << " moved=" + << (alt.iterNum != decision.iterNum ? "yes" : "no") + << " why=" << decision.why << " treeOnlyWhy=" << alt.why + << "\n"; + } + + // One run of the pressure model, from the two measurements to the decider's + // verdict. Held together in a struct so the pin branch can ask what the model + // would have said without a second copy of the wiring: the one time this pass + // kept two copies of a derived number (`target`), they drifted, which is why + // 2.3 folded it into `vrfBudgetTarget`. + struct ModelRun { + triton::xpu::Decision decision; + RegPressure treeP, blockP; + triton::xpu::TileContext ctx; + SmallVector candidates; + int64_t maxLegal = 1; + // The tree holds no vector value, so the model has nothing to say and stops + // before `blockP` / `ctx` / `decision` are filled. The caller decides what + // to fall back to; this only reports that it happened. + bool noVectorValues = false; + }; + + void runModel(Operation *insertPt, const SetVector &unrollOpTree, + int64_t numCol, int64_t widthPerCore, int64_t loopResults, + ModelRun &run) { + getRegPressure(getOperation(), unrollOpTree, run.treeP); + // A tree that holds no vector value (a bool store, say) puts no pressure on + // the vector file, so this model has nothing to say about it: what tiling + // buys there is code size, which is not modelled. + if (run.treeP.vecPeak == 0) { + run.noVectorValues = true; + return; + } + getBlockRegPressure(getOperation(), insertPt, run.blockP); + // Everything the criteria may read, and nothing else: the target itself is + // computed by the tier-1 criterion out of these numbers (`vrfBudgetTarget`), + // so the pass no longer holds a second copy of it. + // + // Scalar pressure is measured and reported, but it does not drive the + // factor. Measured on the layernorm probe (bufSz=512): scalarPeak 389..770 + // against vecPeak 32..48, while the emitted scalar spill count stays + // at 7..9 no matter which trip count is picked. Those un-vectorized tensors + // are not register-resident at that scale, so charging them to a register + // budget only distorts the target; what they actually cost is instructions. + run.ctx.numCol = numCol; + run.ctx.widthPerCore = widthPerCore; + run.ctx.peakVRegs = std::max(run.treeP.vecPeak, run.blockP.vecPeak); + run.ctx.maxVecWidth = + std::max(run.treeP.maxVecWidth, run.blockP.maxVecWidth); + run.ctx.scalarPeak = std::max(run.treeP.scalarPeak, run.blockP.scalarPeak); + run.ctx.vecRow = + hasVecScalarBoundary(unrollOpTree) ? run.treeP.minVecWidth : 0; + run.ctx.vrfBudget = this->vrfBudget; + run.ctx.loopResults = loopResults; + + // A per-candidate pressure term for reduce segments that collapse an + // interpreted combine region (step 3.4b) lived here and was reverted on + // 2026-08-05: it drove welford's reduce segment to `vspill=0` at + // `iterNum=8` as designed, but the real-hardware sweep says that point is + // 15.6% *slower* than the `iterNum=2` the plain budget model picks, at all + // three occupancies measured (`findings.md` §1.17). Spilling 14 accumulator + // vregs is cheaper on this segment than running the collapse 4x more times, + // so zero spill is not the objective this model should be optimising. + // + // The candidate set is the expressible trip counts; which one to take is + // the decider's business, tier by tier. + for (int64_t iterNum = 1; iterNum <= widthPerCore; ++iterNum) { + if (!isLegalIterNum(iterNum, numCol, widthPerCore, run.ctx.vecRow)) + continue; + run.candidates.emplace_back(iterNum); + run.maxLegal = iterNum; + } + triton::xpu::TileDecider decider; + run.decision = decider.decide(run.candidates, run.ctx); + } + + // Step 3.5 groundwork, report-only. `pinnedUnrollNum` short-circuits the model + // before it ever runs, and the pin value is not reachable from any knob + // (`unroll_num` is read after it), so today neither magic number can be shown + // to be wrong: boolfused emits the same code at unroll_num 1/4/16. This puts + // the pin's factor next to what the model would have picked at the same site. + // + // The legacy factor is printed next to it because that is what the site would + // fall back to. The two turn out to be the *same* number by construction, not + // by coincidence: `getNumUnroll` already folds `unrollNum * coresPerGroup` + // into `numUnroll`, and both pin sites set `unrollNum` to the pinned value, so + // `ceil(numCol, pin * coresPerGroup) == ceil(numCol, numUnroll)`. Measured on + // both pin sites (`findings.md` §1.30). That is why 3.5b only has to stop the + // early return -- the pin's arithmetic contributes nothing of its own. + void reportPinShadow(const char *site, Operation *insertPt, + const SetVector &unrollOpTree, + int64_t numCol, int64_t widthPerCore, + int64_t coresPerGroup, int64_t legacyIterNum, + int64_t pinned, int64_t loopResults) { + if (!std::getenv("TRITONXPU_TILE_REPORT")) + return; + ModelRun run; + runModel(insertPt, unrollOpTree, numCol, widthPerCore, loopResults, run); + StringRef kernel = ""; + if (auto funcOp = insertPt->getParentOfType()) + kernel = funcOp.getName(); + llvm::errs() << "[PinShadow] " << kernel << " site=" << site + << " pin=" << pinReason + << " pinnedUnrollNum=" << pinnedUnrollNum + << " numCol=" << numCol << " widthPerCore=" << widthPerCore + << " coresPerGroup=" << coresPerGroup + << " pinIterNum=" << pinned + << " legacyIterNum=" << legacyIterNum; + // Without the pin this site would fall back to the legacy factor -- which + // the pin itself already overwrote (`this->unrollNum`), so the comparison is + // against the pinned unrollNum, not the user's. Worth seeing, not hiding. + if (run.noVectorValues) { + llvm::errs() << " modelWhy=no-vector-values:legacy treePeak=0" + << " treeScalarPeak=" << run.treeP.scalarPeak + << " modelIterNum=" << legacyIterNum << " moved=" + << (legacyIterNum != pinned ? "yes" : "no") << "\n"; + return; + } + // The candidate set in full: a pin value that is not even expressible, or a + // set collapsed to {1}, is the degradation path 3.5 must keep observable. + std::string legal; + for (int64_t c : run.candidates) + legal += (legal.empty() ? "" : ",") + std::to_string(c); + llvm::errs() << " treePeak=" << run.treeP.vecPeak + << " blockPeak=" << run.blockP.vecPeak + << " vecRow=" << run.ctx.vecRow + << " target=" << triton::xpu::vrfBudgetTarget(run.ctx) + << " maxLegal=" << run.maxLegal << " legal={" << legal << "}" + << " pinExpressible=" + << (llvm::is_contained(run.candidates, pinned) ? "yes" : "no") + << " modelIterNum=" << run.decision.iterNum + << " modelWhy=" << run.decision.why << " moved=" + << (run.decision.iterNum != pinned ? "yes" : "no") << "\n"; + } + + // The register file is shared by everything simultaneously live in the block, + // not by one op tree: sibling trees hold their values at the same time (the + // mean and var accumulators of a layernorm, say). Measuring per tree + // double-spends the file, so the pressure is taken over the whole block the + // tile loop will live in. Every vector value there scales with the same + // factor, so one target per block is the right granularity. + // Returns the whole decision, not just the factor: the trip count alone + // cannot say which criterion produced it, and the two call sites must not + // have to change again every time a criterion is added (§3.2.2). + // `loopResults` is the number of iter_args the tile loop will carry, which + // only the caller knows before the loop exists. + triton::xpu::Decision + decideIterNum(const char *site, Operation *insertPt, + const SetVector &unrollOpTree, Type valTy, + int64_t numCol, int64_t numUnroll, int64_t loopResults = 0) { + int64_t legacyIterNum = ceil(numCol, numUnroll); + if (dryRun) + reportDryRun(site, insertPt, unrollOpTree, valTy, numCol, numUnroll, + legacyIterNum); + int64_t widthPerCore = 1, coresPerGroup = 1; + getTileGeometry(valTy, numCol, widthPerCore, coresPerGroup); + + if (!this->budgetTiling) { + // The vector row is a legality bound, not part of the model, so it binds + // the legacy factor as well: ceil(numCol, unrollNum) is derived from the + // same possibly-scalar `valTy` and can exceed what the vector values in + // the tree can express. Fall back to the largest expressible factor that + // is no larger. + RegPressure p; + int64_t vecRow = 0; + if (hasVecScalarBoundary(unrollOpTree)) { + getRegPressure(getOperation(), unrollOpTree, p); + vecRow = p.minVecWidth; + } + if (vecRow && + !isLegalIterNum(legacyIterNum, numCol, widthPerCore, vecRow)) { + int64_t clamped = 1; + for (int64_t iterNum = 1; iterNum <= legacyIterNum; ++iterNum) + if (isLegalIterNum(iterNum, numCol, widthPerCore, vecRow)) + clamped = iterNum; + std::string msg; + llvm::raw_string_ostream os(msg); + os << "[UnrollControl] site=" << site << " numCol=" << numCol + << " widthPerCore=" << widthPerCore << " minVecWidth=" << vecRow + << " -> iterNum=" << clamped << " (vector-row-bound, legacy " + << legacyIterNum << ")"; + insertPt->emitRemark(msg); + legacyIterNum = clamped; + return {legacyIterNum, "vector-row-bound", {}}; + } + return {legacyIterNum, "legacy", {}}; + } + + // Collect the expressible trip counts once: the largest is the fallback + // when the budget cannot be reached at all. + int64_t maxLegal = 1; + + // A pinned factor claims to be a correctness constraint, not a heuristic, so + // the model gets no say here; it only reports which constraint took over. + // Whether that claim holds is now testable from the outside: `pin-unroll-num` + // (3.5b) can drop the pin or move its constant, and `[PinShadow]` says what + // the model would have picked at the same site. + if (pinnedUnrollNum > 0) { + int64_t pinned = ceil(numCol, pinnedUnrollNum * coresPerGroup); + pinned = std::max(pinned, 1); + remarkDecision(site, insertPt, numCol, widthPerCore, /*peakVRegs=*/-1, + /*target=*/-1, pinned, maxLegal, pinReason); + reportPinShadow(site, insertPt, unrollOpTree, numCol, widthPerCore, + coresPerGroup, legacyIterNum, pinned, loopResults); + return {pinned, pinReason, {}}; + } + + ModelRun run; + runModel(insertPt, unrollOpTree, numCol, widthPerCore, loopResults, run); + // Nothing for the model to say: keep the legacy factor. + if (run.noVectorValues) { + remarkDecision(site, insertPt, numCol, widthPerCore, /*peakVRegs=*/0, + /*target=*/-1, legacyIterNum, /*maxLegal=*/-1, + "no-vector-values:legacy", run.treeP.scalarPeak); + return {legacyIterNum, "no-vector-values:legacy", {}}; + } + remarkDecision(site, insertPt, numCol, widthPerCore, run.ctx.peakVRegs, + triton::xpu::vrfBudgetTarget(run.ctx), run.decision.iterNum, + run.maxLegal, run.decision.why.c_str(), run.ctx.scalarPeak, + run.ctx.vecRow, &run.decision); + reportSegDominance(site, insertPt, run.treeP, run.blockP, run.ctx, + run.candidates, run.decision); + return run.decision; + } + + // Earliest store of the tree plus the geometry it implies. Returns false when + // the stores of one tree live in different blocks, which the rewrite cannot + // handle yet. + bool getTreeUnrollInfo(const SetVector &unrollOpTree, + triton::xpu::StoreOp &insertPt, + SmallVector &allStoreOps, + int64_t &numCol, int64_t &numUnroll) { + for (auto op : unrollOpTree) { + auto storeOp = dyn_cast(op); + if (!storeOp) + continue; + auto type = storeOp.getValue().getType(); + numUnroll = numUnroll == 1 ? getNumUnroll(type) + : std::min(numUnroll, getNumCol(type)); + numCol = numCol == 1 ? getNumCol(type) : std::min(numCol, getNumCol(type)); + allStoreOps.emplace_back(storeOp); + //[TODO] To deal with the case that storeOps are in more than one block + if (insertPt && insertPt->getBlock() != storeOp->getBlock()) + return false; + if (!insertPt || storeOp->isBeforeInBlock(insertPt)) + insertPt = storeOp; + } + return true; + } + + // One decision per tree, taken before any IR is touched so that the discrete + // pointer rewrite can be kept in sync with it. + SmallVector + planIterNums(SmallVector> &unrollOpTrees, + const char *site) { + SmallVector plan; + for (auto &unrollOpTree : unrollOpTrees) { + triton::xpu::StoreOp insertPt; + SmallVector allStoreOps; + int64_t numCol = 1, numUnroll = 1; + if (!getTreeUnrollInfo(unrollOpTree, insertPt, allStoreOps, numCol, + numUnroll) || + !insertPt) { + plan.emplace_back(1); + continue; + } + // A pointwise store segment carries no iter_args (createFor is called + // with an empty range at :1263), so the loop-overhead criterion sees + // only the index arithmetic. + plan.emplace_back(decideIterNum(site, insertPt, unrollOpTree, + insertPt.getValue().getType(), numCol, + numUnroll, /*loopResults=*/0) + .iterNum); + } + return plan; + } + Type createPointerType(Type type, int64_t vecSize) { if (auto tensorType = dyn_cast(type)) { Type elemType = getElementTypeOrSelf(tensorType); @@ -631,6 +1204,10 @@ struct TritonXPUUnrollControl } else { forOp = builder.create(loc, lower, upper, step, iterArgs); } + // Later stages of this pass walk the module again and would otherwise tile + // an already tiled segment a second time, inserting the loop index twice. + // The legacy gate hid this because a tiled value has numCol == numUnroll. + forOp->setAttr(kUnrollLoopAttr, UnitAttr::get(forOp->getContext())); builder.setInsertionPointToStart(forOp.getBody()); idxVar = builder.create(loc, builder.getI32Type(), @@ -642,6 +1219,20 @@ struct TritonXPUUnrollControl SetVector &outerChain, arith::IndexCastOp &idxVar, IRMapping &mapping) { for (auto op : unrollOpTree) { + // A vector<->scalar boundary buffer must not be cloned into the tile + // loop. An alloca inside a loop body is never promoted, so every + // iteration grows the stack and the launch fails outright. The original + // alloca already sits before forOp (the tree precedes insertPt, and + // forOp was created at insertPt), so all this takes is not cloning it: + // the clones inside the body then reference the outside buffer, and + // eraseDAG leaves it alone because it is no longer use-empty. + // + // Keeping one full-width buffer for all iterations is correct because + // the boundary is written and read back within a single iteration, and + // leaving its type unsliced is harmless because lowering reads only + // element 0 of the bufPtr (LoadStoreOpToLLVM.cpp getBoundaryLMBase). + if (isBoundaryBuffer(op)) + continue; bool isOuter = inOpChain(outerChain, op); auto newOp = builder.clone(*op, mapping); setTensorType(context, newOp, iterNum, isOuter); @@ -651,7 +1242,7 @@ struct TritonXPUUnrollControl dyn_cast(loadOp.getPtr().getType())) { auto shape = tensorTy.getShape(); bool isOuter = (shape.size() == 2 && shape.back() == 1); - if (!isOuter && !loadOp.getSVOpt() && !loadOp.getIsDiscrete()) { + if (!isOuter && !loadOp.getSVOpt()) { insertIndex(newOp, idxVar); } } @@ -861,7 +1452,7 @@ struct TritonXPUUnrollControl void unrollControl(MLIRContext *context, SmallVector> &unrollOpTrees, - bool postReduce = false) { + ArrayRef plan, bool postReduce = false) { // Get outerChains SmallVector> outerChains; getOuterChains(unrollOpTrees, outerChains, postReduce); @@ -873,30 +1464,19 @@ struct TritonXPUUnrollControl int64_t numUnroll = 1; triton::xpu::StoreOp insertPt; SmallVector allStoreOps; - for (auto op : unrollOpTree) { - // 1.1 Get insertPt and tensor num - if (auto storeOp = dyn_cast(op)) { - auto type = storeOp.getValue().getType(); - numUnroll = numUnroll == 1 ? getNumUnroll(type) - : std::min(numUnroll, getNumCol(type)); - numCol = - numCol == 1 ? getNumCol(type) : std::min(numCol, getNumCol(type)); - allStoreOps.emplace_back(storeOp); - //[TODO] To deal with the case that storeOps are in more than one - // block - if (insertPt && insertPt->getBlock() != storeOp->getBlock()) { - return; - } - if (!insertPt || storeOp->isBeforeInBlock(insertPt)) { - insertPt = storeOp; - } - } - } + // 1.1 Get insertPt and tensor num + if (!getTreeUnrollInfo(unrollOpTree, insertPt, allStoreOps, numCol, + numUnroll)) + return; if (insertPt) { auto loc = insertPt.getLoc(); - int64_t iterNum = ceil(numCol, numUnroll); + // Decided in planIterNums, before any IR was touched. + int64_t iterNum = plan[i]; + // Skip this tree only: unlike the legacy gate, the budget model can + // legitimately answer "no loop" for one tree while the others still + // want one, so this must not abandon the remaining trees. if (iterNum <= 1) - return; + continue; LLVM_DEBUG(llvm::dbgs() << "[Unroll Control] Hit Unroll Control Pointwise\n"); // 2. Unroll control @@ -919,16 +1499,15 @@ struct TritonXPUUnrollControl } } + // `iterNum` is decided by the caller: it has to be known before + // findDiscretePtrChain() rewrites the pointer chain. void unrollControlReduce(MLIRContext *context, SetVector &unrollOpTree, Operation *insertPt, ValueRange &iterArgs, - ValueRange &returnOperands) { + ValueRange &returnOperands, int64_t iterNum) { SetVector outerChain; getOuterChain(unrollOpTree, outerChain); if (auto reduceOp = dyn_cast(insertPt)) { - int64_t numCol = 1, numUnroll = 1; - getUnrollInfoReduce(reduceOp, numCol, numUnroll); - int64_t iterNum = ceil(numCol, numUnroll); if (iterNum <= 1) return; OpBuilder builder(reduceOp); @@ -1063,7 +1642,8 @@ struct TritonXPUUnrollControl } void findDiscretePtrChain(SetVector &unrollOpTree, - SetVector &newUnrollOpTree) { + SetVector &newUnrollOpTree, + bool treeWillTile) { for (auto op : unrollOpTree) { if (auto loadOp = dyn_cast(op)) { bool isDiscrete = loadOp.getIsDiscrete(); @@ -1073,7 +1653,12 @@ struct TritonXPUUnrollControl auto resType = loadOp.getResult().getType(); int64_t numCol = getNumCol(resType); int64_t numUnroll = getNumUnroll(resType); - if (numCol > numUnroll && numCol % numUnroll == 0) { + // The rewrite only makes sense inside the tiling loop, so it must + // agree with the decision taken for the whole tree. + bool willTile = this->budgetTiling + ? treeWillTile + : (numCol > numUnroll && numCol % numUnroll == 0); + if (willTile) { auto lmPtr = loadOp.getPtr(); if (auto gm2lmOp = findDefOpBwd(lmPtr)) { auto gmPtrOp = findDefOpBwd(gm2lmOp.getPtr()); @@ -1118,9 +1703,10 @@ struct TritonXPUUnrollControl void findDiscretePtrChains(SmallVector> &unrollOpTrees, - SmallVector> &newUnrollOpTrees) { + SmallVector> &newUnrollOpTrees, + ArrayRef plan) { for (auto [i, unrollOpTree] : llvm::enumerate(unrollOpTrees)) { - findDiscretePtrChain(unrollOpTree, newUnrollOpTrees[i]); + findDiscretePtrChain(unrollOpTree, newUnrollOpTrees[i], plan[i] > 1); } } @@ -1137,6 +1723,13 @@ struct TritonXPUUnrollControl if (auto gm2lmOp = findDefOpBwd(lmPtr)) { auto gmPtrOp = findDefOpBwd(gm2lmOp.getPtr()); auto gmOffset = gmPtrOp.getOffset(); + // Nothing to rebase when the LM side has no addptr of its own: the + // one we found walking back *is* the gm2lm's GM addptr, which is what + // an untiled tree looks like. A discrete gm2lm already gathers into a + // 0-based LM buffer, so rewriting this offset would only strip the + // block's base off the GM address the gather reads from. + if (lmAddPtr == gmPtrOp) + return; auto extractOp = builder.create( loc, getElementTypeOrSelf(gmOffset), builder.getI32IntegerAttr(0), gmOffset); @@ -1152,6 +1745,13 @@ struct TritonXPUUnrollControl findDefOpBwd(lmPtr)) { auto gmPtrOp = findDefOpBwd(gm2lmOp.getPtr()); auto gmOffset = gmPtrOp.getOffset(); + // Nothing to rebase when the LM side has no addptr of its own: the + // one we found walking back *is* the gm2lm's GM addptr, which is what + // an untiled tree looks like. A discrete gm2lm already gathers into a + // 0-based LM buffer, so rewriting this offset would only strip the + // block's base off the GM address the gather reads from. + if (lmAddPtr == gmPtrOp) + return; auto extractOp = builder.create( loc, getElementTypeOrSelf(gmOffset), builder.getI32IntegerAttr(0), gmOffset); @@ -1181,7 +1781,9 @@ struct TritonXPUUnrollControl auto valType = storeOp.getValue().getType(); int64_t numCol = getNumCol(valType); int64_t numUnroll = getNumUnroll(valType); - if (numCol > numUnroll && numCol % numUnroll == 0) { + if (dryRun) + reportGate("pointwise", storeOp, numCol, numUnroll); + if (canTile(storeOp, valType, numCol, numUnroll)) { getDAG(storeOp, visitedOps, unrollOpTrees, excludeChainOps); } for (auto visitedOp : visitedOps) { @@ -1194,11 +1796,14 @@ struct TritonXPUUnrollControl return; // 1.3 Find ptr chain of discrete for moving to loop body + // The factor must be decided before this rewrite: it is only valid for + // trees that really end up inside a tiling loop. + SmallVector plan = planIterNums(unrollOpTrees, "pointwise"); SmallVector> newUnrollOpTrees(unrollOpTrees); - findDiscretePtrChains(unrollOpTrees, newUnrollOpTrees); + findDiscretePtrChains(unrollOpTrees, newUnrollOpTrees, plan); // 2. Deal with unroll opTrees - unrollControl(context, newUnrollOpTrees); + unrollControl(context, newUnrollOpTrees, plan); // 3. Calculate discrete offset in the runtime createDiscreteOffset(m); @@ -1297,7 +1902,9 @@ struct TritonXPUUnrollControl m.walk([&](triton::xpu::ReduceOp reduceOp) { int64_t numCol = 1, numUnroll = 1; getUnrollInfoReduce(reduceOp, numCol, numUnroll); - if (numCol > numUnroll && numCol % numUnroll == 0) { + if (dryRun) + reportGate("reduce", reduceOp, numCol, numUnroll); + if (canTile(reduceOp, reduceOp.getInputTypes()[0], numCol, numUnroll)) { llvm::SetVector reduceOpDefsBwd; getOpChainBwd(reduceOpDefsBwd, reduceOp); for (auto operand : reduceOpDefsBwd) { @@ -1327,10 +1934,13 @@ struct TritonXPUUnrollControl getDAG(storeOp, visitedOps, unrollOpTrees, excludeChainOps, true, true); // Find ptr chain of discrete for moving to loop body + SmallVector plan = + planIterNums(unrollOpTrees, "reduce-for"); SmallVector> newUnrollOpTrees( unrollOpTrees); - findDiscretePtrChains(unrollOpTrees, newUnrollOpTrees); - unrollControl(context, newUnrollOpTrees); + findDiscretePtrChains(unrollOpTrees, newUnrollOpTrees, + plan); + unrollControl(context, newUnrollOpTrees, plan); } hasIf = true; } @@ -1349,10 +1959,12 @@ struct TritonXPUUnrollControl getDAG(storeOp, visitedOps, unrollOpTrees, excludeChainOps, true, true); // Find ptr chain of discrete for moving to loop body + SmallVector plan = + planIterNums(unrollOpTrees, "reduce-for"); SmallVector> newUnrollOpTrees( unrollOpTrees); - findDiscretePtrChains(unrollOpTrees, newUnrollOpTrees); - unrollControl(context, newUnrollOpTrees); + findDiscretePtrChains(unrollOpTrees, newUnrollOpTrees, plan); + unrollControl(context, newUnrollOpTrees, plan); } } } @@ -1385,7 +1997,7 @@ struct TritonXPUUnrollControl SetVector unrollOpTree; int64_t numCol = 1, numUnroll = 1; getUnrollInfoReduce(reduceOp, numCol, numUnroll); - if (numCol > numUnroll && numCol % numUnroll == 0) { + if (canTile(reduceOp, reduceOp.getInputTypes()[0], numCol, numUnroll)) { LLVM_DEBUG(llvm::dbgs() << "[Unroll Control] Hit Unroll Control Reduction\n"); for (int i = 0; i < reduceOperandNum; ++i) { @@ -1394,6 +2006,24 @@ struct TritonXPUUnrollControl false); } } + // 0. Decide the factor up front: everything below mutates the IR + // (clones the operand chain, inlines the combine region) and is only + // valid if a loop is actually created afterwards. Deciding later + // would leave the inlined combine region orphaned. + SetVector probeOpTree; + for (auto ©OpTree : copyOpTrees) + for (auto *copyOp : copyOpTree) + probeOpTree.insert(copyOp); + // The reduce loop carries one accumulator per data operand (the + // iterArgs built at :1908), which is what the loop-overhead criterion + // charges for. + int64_t iterNum = + decideIterNum("reduce", reduceOp, probeOpTree, + reduceOp.getInputTypes()[0], numCol, numUnroll, + /*loopResults=*/reduceOperandNum) + .iterNum; + if (iterNum <= 1) + return; // 1. Copy Defined Op Chain of Reduce Operand for InitArgs IRMapping mapping; for (auto ©OpTree : copyOpTrees) { @@ -1412,61 +2042,88 @@ struct TritonXPUUnrollControl // Set Type for Cloned Ops auto tensorTy = reduceOp.getInputTypes()[0]; auto shape = tensorTy.getShape(); + // `tt.splat` cannot carry a vector element type. On the + // region-interpreting path (TRITONXPU_REDUCE_REGION) the combine + // region's constants are vector, so the broadcast has to be + // triton_xpu.vsplat instead -- the same choice Vectorize.cpp makes for + // splats it retypes. VSplatOpConversion broadcasts a *scalar* into every + // lane (insertelement into lane 0 + shuffle), so the constant is + // narrowed back to its splat value here rather than handed over as a + // vector. Ops are created at `anchor` because the cloned combine block + // sits ahead of the builder's insertion point. + auto createCombineSplat = [&](mlir::Type resTy, mlir::Value src, + mlir::Operation *anchor) -> mlir::Value { + OpBuilder b(anchor); + auto elemTy = + mlir::cast(resTy).getElementType(); + if (!mlir::isa(elemTy)) + return b.create(loc, resTy, src).getResult(); + auto cstOp = src.getDefiningOp(); + auto dense = mlir::cast(cstOp.getValue()); + assert(dense.isSplat() && "combine constant is not uniform"); + auto scalar = b.create( + cstOp.getLoc(), + mlir::cast(dense.getSplatValue())); + return b.create(loc, resTy, scalar) + .getResult(); + }; for (auto &op : newReduce) { if (isa(op) || isa(op)) { auto tensorTy0 = op.getOperand(0).getType(); auto tensorTy1 = op.getOperand(1).getType(); - int operandIndexNeedModify; + // The operand that is not a tensor is the one to broadcast. Asking + // "is it a Float or an Integer" is too narrow: on the vector path + // the combine region's constants are vector, and neither branch + // used to fire, leaving operandIndexNeedModify uninitialized and the + // assert below reading a garbage index. + int operandIndexNeedModify = -1; mlir::Type operandNeedReserved; if (tensorTy0 != tensorTy1) { - if ((mlir::isa(tensorTy0) || - mlir::isa(tensorTy0)) && + if (!mlir::isa(tensorTy0) && mlir::isa(tensorTy1)) { operandIndexNeedModify = 0; operandNeedReserved = tensorTy1; - } else if ((mlir::isa(tensorTy1) || - mlir::isa(tensorTy1)) && + } else if (!mlir::isa(tensorTy1) && mlir::isa(tensorTy0)) { operandIndexNeedModify = 1; operandNeedReserved = tensorTy0; } assert( + operandIndexNeedModify >= 0 && isa( op.getOperand(operandIndexNeedModify).getDefiningOp()) && "Unable to extract the non-constant operand."); - auto splatOp = builder.create( - loc, operandNeedReserved, - op.getOperand(operandIndexNeedModify)); - splatOp->moveBefore(&op); - op.setOperand(operandIndexNeedModify, splatOp.getResult()); + op.setOperand(operandIndexNeedModify, + createCombineSplat( + operandNeedReserved, + op.getOperand(operandIndexNeedModify), &op)); } } else if (auto selOp = dyn_cast(op)) { auto tensorTy1 = selOp.getODSOperands(1)[0].getType(); auto tensorTy2 = selOp.getODSOperands(2)[0].getType(); - int operandIndexNeedModify; + int operandIndexNeedModify = -1; mlir::Type operandNeedReserved; if (tensorTy1 != tensorTy2) { - if ((mlir::isa(tensorTy1) || - mlir::isa(tensorTy1)) && + if (!mlir::isa(tensorTy1) && mlir::isa(tensorTy2)) { operandIndexNeedModify = 1; operandNeedReserved = tensorTy2; - } else if ((mlir::isa(tensorTy2) || - mlir::isa(tensorTy2)) && + } else if (!mlir::isa(tensorTy2) && mlir::isa(tensorTy1)) { operandIndexNeedModify = 2; operandNeedReserved = tensorTy1; } - assert(isa( + assert(operandIndexNeedModify >= 0 && + isa( selOp.getOperand(operandIndexNeedModify) .getDefiningOp()) && "Unable to extract the non-constant operand."); - auto splatOp = builder.create( - loc, operandNeedReserved, - selOp.getOperand(operandIndexNeedModify)); - splatOp->moveBefore(&op); - selOp.setOperand(operandIndexNeedModify, splatOp.getResult()); + selOp.setOperand(operandIndexNeedModify, + createCombineSplat( + operandNeedReserved, + selOp.getOperand(operandIndexNeedModify), + &op)); } } for (auto [i, resTy] : llvm::enumerate(op.getResultTypes())) { @@ -1499,11 +2156,11 @@ struct TritonXPUUnrollControl } // Find ptr chain of discrete for moving to loop body SetVector newUnrollOpTree(unrollOpTree); - findDiscretePtrChain(unrollOpTree, newUnrollOpTree); + findDiscretePtrChain(unrollOpTree, newUnrollOpTree, iterNum > 1); // 3. Create Loop for ReduceWithinCore ValueRange iterArgsRange(iterArgs); unrollControlReduce(context, newUnrollOpTree, reduceOp, iterArgsRange, - returnOperands); + returnOperands, iterNum); // 4. For Vectorize: triton.addf->triton_xpu.vvaddf processOpVecTy(m); } @@ -1578,7 +2235,7 @@ struct TritonXPUUnrollControl int64_t numCol = getNumCol(valType); int64_t numUnroll = getNumUnroll(valType); bool _isPostReduceStore = isPostReduceStore(storeOp); - if (numCol > numUnroll && numCol % numUnroll == 0 && _isPostReduceStore) { + if (canTile(storeOp, valType, numCol, numUnroll) && _isPostReduceStore) { getPostReduceDAG(storeOp, visitedOps, unrollOpTrees, excludeChainOps); } }); @@ -1591,7 +2248,8 @@ struct TritonXPUUnrollControl // 3. Deal with unroll opTrees LLVM_DEBUG(llvm::dbgs() << "[Unroll Control] Hit Unroll Control Post Reduction\n"); - unrollControl(context, unrollOpTrees, true); + SmallVector plan = planIterNums(unrollOpTrees, "post-reduce"); + unrollControl(context, unrollOpTrees, plan, /*postReduce=*/true); } void reductionUnrollControl(ModuleOp &m, MLIRContext *context) { @@ -1611,6 +2269,8 @@ struct TritonXPUUnrollControl MLIRContext *context = &getContext(); ModuleOp m = getOperation(); + dryRun = std::getenv("TRITONXPU_UNROLL_DRYRUN") != nullptr; + bool isScan = false; m.walk([&](triton::xpu::ScanOp scanOp) { isScan = true; }); if (isScan) { @@ -1625,7 +2285,7 @@ struct TritonXPUUnrollControl auto ptrElemTy = getElementTypeOrSelf(getElementTypeOrSelf(ptrTy)); if (dtype == Dtype::FP32 && valElemTy.isInteger(32) && cast(ptrElemTy).getPointeeType().isInteger(8)) { - this->unrollNum = 4; + applyPin(m, /*constant=*/4, "bool-store-vectorize"); } }); @@ -1638,8 +2298,9 @@ struct TritonXPUUnrollControl auto layout = cast(operandType.getEncoding()); unsigned rowsPerCore = layout.getSizePerCore()[0]; - this->unrollNum = - (shape.size() == 2 && rowsPerCore > 1) ? 1 : this->unrollNum; + if (shape.size() == 2 && rowsPerCore > 1) { + applyPin(m, /*constant=*/1, "core-deal-multi-rows"); + } }); if (isReduce) { @@ -1647,6 +2308,13 @@ struct TritonXPUUnrollControl } else { pointwiseUnrollControl(m, context); } + + // The marker only exists to stop a later walk inside this pass from tiling + // an already tiled segment a second time; it must not survive into the + // emitted IR. With budget_tiling off the artifacts have to stay + // byte-identical to the legacy pipeline, and a leftover attribute is a + // visible difference (10/10 probes' .ttxir differed on just this line). + m.walk([&](scf::ForOp forOp) { forOp->removeAttr(kUnrollLoopAttr); }); } }; diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/VectorizabilityAnalysisPass.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/VectorizabilityAnalysisPass.cpp new file mode 100644 index 0000000000..afe907c85b --- /dev/null +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/VectorizabilityAnalysisPass.cpp @@ -0,0 +1,310 @@ +//===----------------------------------------------------------------------===// +// tritonxpu-vectorizability-analysis -- report-only pass over +// VectorizabilityAnalysis. +// +// Same purpose as TileAnalysisPass: give the measurement its own pipeline slot +// so the transform can later be handed the answer instead of recomputing it +// (redesign-v2.md §4.2 step 1.3). Never mutates the IR. +// +// It enumerates the roots Vectorize enumerates -- reduce operands under +// ReduceVec, then every store -- and reports, per root, whether the root is +// eligible and how big the closure is. +// +// Two instances of it are registered, and the point of the second one is to make +// a specific risk measurable instead of argued about. +// +// preTiling=false -- immediately before Vectorize. Since step 1.5 moved that +// pass' prologue into tritonxpu-normalize, this sees exactly the IR +// Vectorize rewrites, so `stage=pre-vectorize` is by construction the answer +// Vectorize will reach. +// +// preTiling=true -- ahead of CoreTiling. The three in-walk footprint tests +// (LoadOp, SplatOp, BroadcastOp) read sizePerCore, which does not exist yet, +// so this instance hands the walk an all-Unknown oracle: it never vetoes and +// records what it deferred (`stage=pre-tiling`, `cands=`). It then repeats +// the walk with the real oracle at the same position (`stage=pre-tiling-fit`). +// +// `pre-tiling-fit` vs `pre-vectorize` is the diff that matters: it holds the walk +// fixed and varies only the position, so a disagreement is CoreTiling or Legalize +// moving an E-dependent answer -- e.g. Legalize.cpp:245-247's +// `slicedShape[i] = max(shape[i]/iterCount[i], 1)`, whose saturation can flip +// BroadcastOp's `srcShape[1] == resShape[1]` from false to true. Agreement is +// what licenses moving the state half up (step 1.5c). +// +// Nothing is emitted unless TRITONXPU_VEC_REPORT=1, and neither instance mutates +// the IR. +//===----------------------------------------------------------------------===// + +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Analysis/TileAnalysis.h" +#include "triton/Analysis/VectorizabilityAnalysis.h" +#include "triton/Dialect/TritonXPU/IR/Dialect.h" +#include "triton/Dialect/TritonXPU/Transforms/Passes.h" + +#define DEBUG_TYPE "tritonxpu-vectorizability-analysis" + +namespace mlir { +namespace triton { +namespace xpu { + +#define GEN_PASS_DEF_TRITONXPUVECTORIZABILITYANALYSIS +#include "triton/Dialect/TritonXPU/Transforms/Passes.h.inc" + +struct TritonXPUVectorizabilityAnalysisPass + : public impl::TritonXPUVectorizabilityAnalysisBase< + TritonXPUVectorizabilityAnalysisPass> { + +public: + using impl::TritonXPUVectorizabilityAnalysisBase< + TritonXPUVectorizabilityAnalysisPass>:: + TritonXPUVectorizabilityAnalysisBase; + + TritonXPUVectorizabilityAnalysisPass() = default; + TritonXPUVectorizabilityAnalysisPass(bool reduceVec, bool preTiling) { + this->reduceVec = reduceVec; + this->preTiling = preTiling; + } + + // Today's answer, computed with the real oracle at whatever position this + // instance sits. A fresh analysis per root, matching + // vectorizeAndProcessOpVecTy: sharing `visited` across roots changes the + // answer, and this pass has to report what that function would see, not + // something tidier. + std::pair reportWithFit(const char *stage, const char *site, + Operation *root, Type rootOpTy) { + VectorizabilityAnalysis analysis(this->reduceVec, /*dumpFlag=*/false, + vectorFitsReduceOperand, vectorFitsValue); + bool eligible = vectorFitsRoot(rootOpTy); + int64_t closureSize = 0; + if (eligible) { + OperationTree visited, vectorizedOps; + if (analysis.getVectorizableClosure(root, visited, vectorizedOps)) + closureSize = vectorizedOps.size(); + } + if (vecReportEnabled()) + reportVecRoot(stage, site, root, rootOpTy, eligible, closureSize); + return {eligible, closureSize}; + } + + // What a state-only walk can conclude before CoreTiling: the all-Unknown + // oracle never vetoes, so `closure` here is the E-independent answer and + // `cands` is how many footprint questions were deferred. + void reportState(const char *site, Operation *root, Type rootOpTy) { + // Eligibility keeps only `vectorFitsRoot`'s E-independent conjunct. The + // other two (numElems >= width, numElems % width == 0) divide the per-core + // element count, so they belong to the deferred side -- but unlike the three + // in-walk cases they have no oracle to go through yet, which is why the root + // itself is not in `cands`. That is the hole step 2.1 has to close. + Type elemTy = getElementTypeOrSelf(getElementTypeOrSelf(rootOpTy)); + bool eligible = + isa(rootOpTy) && vectorizedTyValid(elemTy); + + // The reduce-operand gate is E-dependent too and its signature is + // (ReduceOp, Type) rather than a value footprint, so it cannot be a + // FitCandidate. Not vetoing here and counting it into `cands` keeps the + // "deferred, not decided" accounting honest. + int64_t redPending = 0; + auto reduceFitsPending = [&](triton::xpu::ReduceOp, Type) { + ++redPending; + return true; + }; + + VectorizabilityAnalysis analysis(this->reduceVec, /*dumpFlag=*/false, + reduceFitsPending, vectorFitUnknown); + int64_t closureSize = 0; + int64_t cands = 0; + if (eligible) { + OperationTree visited, vectorizedOps; + if (analysis.getVectorizableClosure(root, visited, vectorizedOps)) { + closureSize = vectorizedOps.size(); + cands = analysis.getFitCandidates().size() + redPending; + } + } + if (vecReportEnabled()) + reportVecRoot("pre-tiling", site, root, rootOpTy, eligible, closureSize, + cands); + } + + // Step 2.1: the same roots, answered by the Vector-Flow partition instead of + // by the closure walk. One analysis per function, built before any root is + // reported, because a per-value answer is only meaningful once the whole + // partition exists. + VectorFlowAnalysis *vflowFor(Operation *op) { + auto funcOp = op->getParentOfType(); + if (!funcOp) + return nullptr; + auto it = vflowByFunc.find(funcOp.getOperation()); + return it == vflowByFunc.end() ? nullptr : it->second; + } + + void buildVectorFlow(ModuleOp mod) { + mod.walk([&](triton::FuncOp funcOp) { + auto analysis = std::make_shared(vectorFitsValue); + analysis->run(funcOp); + const VectorFlowStats &st = analysis->getStats(); + llvm::errs() << "[VFlow] " << funcOp.getName() + << " summary values=" << st.values + << " classes=" << st.classes << " vector=" << st.vectorClasses + << " scalar=" << st.scalarClasses + << " conflict=" << st.conflictClasses + << " unset=" << st.unsetClasses << " unions=" << st.unions + << " pins{vec=" << st.vectorPins << ",scl=" << st.scalarPins + << ",extern=" << st.externPins + << ",unknown=" << st.unknownPins << "}" + << " reduce{ops=" << st.reduceOps + << ",vecEntry=" << st.reduceVectorEntries + << ",entryUnpack=" << st.reduceEntryUnpacks << "}\n"; + vflowByFunc[funcOp.getOperation()] = analysis.get(); + vflowOwned.push_back(std::move(analysis)); + }); + } + + // `keyValue` is the value whose state the closure verdict is about: the + // store's value operand, or the reduce operand. Agreement is the contract this + // step has to report -- a disagreement is either the partition seeing a + // boundary the walk had to veto on (expected, that is the point) or a bug. + void reportVFlowRoot(const char *site, Operation *root, Value keyValue, + int64_t closureSize) { + VectorFlowAnalysis *analysis = vflowFor(root); + if (!analysis || !keyValue) + return; + StringRef kernel = ""; + if (auto funcOp = root->getParentOfType()) + kernel = funcOp.getName(); + VState state = analysis->stateOf(keyValue); + bool walkSaysVector = closureSize > 0; + bool flowSaysVector = state == VState::Vector; + // Step 2.2 makes one class of disagreement expected, so it is named rather + // than counted as agreement. The walk is all-or-nothing and vetoes at a + // reduce whose combine region cannot be retyped; the partition answers + // Vector for the producer chain and puts a boundary at the reduce entry. + // That is the modelled case (`boundary`). Anything else that disagrees -- + // in particular the walk retyping a chain the partition calls Scalar -- is + // a real `MISMATCH`, greppable in upper case on purpose. + bool tracked = analysis->isTracked(keyValue); + const char *kind = "match"; + if (walkSaysVector != flowSaysVector) { + if (!walkSaysVector && flowSaysVector && + StringRef(site) == "reduce-combine-veto") + kind = "boundary"; + else if (walkSaysVector && state == VState::Unset) + // Unpinned, not contradicted: no seed in this class demands either + // representation. welford's `w` accumulator is the case -- it is built + // from constants and loop-carried values only, so the sole reason the + // walk retypes it is the reduce being a vector consumer, and *that* is + // the E-dependent fit question (`vectorFitsReduceOperand`) this analysis + // deliberately does not answer. Reported as its own kind rather than + // folded into agreement: a free class is a candidate, which is all M1 + // promises, but it is not the same statement as "Vector". + kind = tracked ? "free" : "UNTRACKED"; + else + kind = "MISMATCH"; + } + llvm::errs() << "[VFlow] " << kernel << " root site=" << site + << " root=" << root->getName() << " closure=" << closureSize + << " state=" << toString(state) << " tracked=" << tracked + << " agree=" << (walkSaysVector == flowSaysVector) + << " kind=" << kind << " loc=" << root->getLoc() << "\n"; + } + + void report(const char *site, Operation *root, Type rootOpTy, + Value siteValue = {}) { + if (!root) + return; + if (!this->preTiling) { + auto [eligible, closureSize] = + reportWithFit("pre-vectorize", site, root, rootOpTy); + if (vflowReportEnabled()) { + // The value whose representation the verdict is about. `siteValue` is + // supplied where the caller already holds it -- a reduce operand can be + // defined by a multi-result op (layernorm's rebuilt `scf.for`), and + // guessing "the root's single result" silently drops those roots. + Value keyValue = siteValue; + if (!keyValue) { + if (auto storeOp = dyn_cast(root)) + keyValue = storeOp.getValue(); + else if (root->getNumResults() == 1) + keyValue = root->getResult(0); + } + reportVFlowRoot(site, root, keyValue, closureSize); + } + // This instance sits where M1/M2 will produce, so it is the one that puts + // the conclusion into the plan (step 1.6). The pre-tiling instance must + // not: it would key roots that the walk there cannot even see, and its ids + // would collide with these. + tilePlanRecord(getOperation(), root, site, eligible, closureSize); + return; + } + reportState(site, root, rootOpTy); + // Same walk, real oracle, same position: today's answer as it would come + // out *here*. Diffing this against stage=pre-vectorize isolates one thing + // and nothing else -- whether CoreTiling and Legalize move the E-dependent + // answer between the two positions. If they do not, the state side above is + // free to move up; if they do, `cands` is where the difference has to be + // re-tested rather than assumed. + reportWithFit("pre-tiling-fit", site, root, rootOpTy); + } + + void runOnOperation() override { + // Three independent switches: the report, the plan probe, and the step 2.1 + // partition. Any one of them being on is reason enough to walk. + if (!vecReportEnabled() && !(tilePlanProbeEnabled() && !this->preTiling) && + !(vflowReportEnabled() && !this->preTiling)) + return; + + ModuleOp mod = getOperation(); + + // Before any root is reported: the per-root lines below read this. + if (vflowReportEnabled() && !this->preTiling) + buildVectorFlow(mod); + + if (this->reduceVec) { + llvm::SetVector reduceOps; + mod.walk([&](triton::xpu::ReduceOp redOp) { reduceOps.insert(redOp); }); + for (auto redOp : reduceOps) { + // Reported rather than skipped: an unvectorizable combine region is the + // reason a whole producer chain stays scalar, so it is the interesting + // case, and Vectorize's own `continue` here is measured to be redundant + // (the closure walk vetoes at the reduce anyway). + if (!reduceCombineIsVectorizable(redOp)) { + if (vecReportEnabled()) + reportVecRoot(this->preTiling ? "pre-tiling" : "pre-vectorize", + "reduce-combine-veto", redOp, + redOp.getInputTypes()[0], /*eligible=*/false, + /*closureSize=*/0); + if (!this->preTiling) { + tilePlanRecord(mod, redOp, "reduce-combine-veto", + /*eligible=*/false, /*closure=*/0); + // Vetoed roots are reported too, otherwise the partition's coverage + // is silently narrower than the walk's. + if (vflowReportEnabled() && !redOp.getOperands().empty()) + reportVFlowRoot("reduce-combine-veto", redOp, + redOp.getOperands()[0], /*closureSize=*/0); + } + continue; + } + for (int i = 0; i < redOp.getOperands().size() - 1; ++i) { + Value operand = redOp.getOperands()[i]; + report("reduce-operand", operand.getDefiningOp(), operand.getType(), + operand); + } + } + } + + mod.walk([&](triton::xpu::StoreOp storeOp) { + report("store", storeOp, storeOp.getValue().getType()); + }); + } + +private: + // Owned per pass run; `vflowByFunc` only borrows. Cleared implicitly when the + // pass instance dies, which is why nothing here outlives the report. + // shared_ptr rather than unique_ptr because MLIR's `clonePass()` copies the + // pass instance, and a unique_ptr member would make the pass non-copyable. + llvm::SmallVector> vflowOwned; + llvm::DenseMap vflowByFunc; +}; + +} // namespace xpu +} // namespace triton +} // namespace mlir diff --git a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/Vectorize.cpp b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/Vectorize.cpp index f01293c7d9..2697c27645 100644 --- a/third_party/xpu/lib/Dialect/TritonXPU/Transforms/Vectorize.cpp +++ b/third_party/xpu/lib/Dialect/TritonXPU/Transforms/Vectorize.cpp @@ -1795,6 +1795,11 @@ struct TritonXPUVectorizePass } } + if (mod->hasAttr(kBF16ToFP32VecOptOffAttrName)) { + BF16ToFP32VecOpt = false; + mod->removeAttr(kBF16ToFP32VecOptOffAttrName); + } + // bfloat16 -> float32 Vector Optimization if (BF16ToFP32VecOpt) { BF16ToFP32VecOptimize(mod); diff --git a/third_party/xpu/python/src/ir.cc b/third_party/xpu/python/src/ir.cc index 47489d7cb3..a581a8df25 100644 --- a/third_party/xpu/python/src/ir.cc +++ b/third_party/xpu/python/src/ir.cc @@ -506,7 +506,12 @@ void init_triton_ir(py::module &&m) { Location loc) { self.addArgument(ty, loc); }) .def("get_num_arguments", &Block::getNumArguments) .def("get_argument", &Block::getArgument) - .def("dump", &Block::dump) + .def("dump", + [](Block &self) { +#if !defined(TRITON_CONCEAL_IR) || (TRITON_CONCEAL_IR == 0) + self.dump(); +#endif + }) .def("move_before", [](Block &self, Block &dst) { self.moveBefore(&dst); }) .def("insert_before", &Block::insertBefore) @@ -591,20 +596,28 @@ void init_triton_ir(py::module &&m) { return self.getBody(idx); }, ret::reference) - .def("dump", [](OpState &self) { self->dump(); }) + .def("dump", [](OpState &self) { +#if !defined(TRITON_CONCEAL_IR) || (TRITON_CONCEAL_IR == 0) + self->dump(); +#endif + }) .def("__str__", [](OpState &self) -> std::string { std::string str; +#if !defined(TRITON_CONCEAL_IR) || (TRITON_CONCEAL_IR == 0) llvm::raw_string_ostream os(str); auto printingFlags = getOpPrintingFlags(); self->print(os, printingFlags); +#endif return str; }) .def("str_nodebug", [](OpState &self) -> std::string { std::string str; +#if !defined(TRITON_CONCEAL_IR) || (TRITON_CONCEAL_IR == 0) llvm::raw_string_ostream os(str); self->print(os); +#endif return str; }) .def("append_operand", @@ -674,13 +687,19 @@ void init_triton_ir(py::module &&m) { // module py::class_(m, "module", py::module_local(), py::dynamic_attr()) - .def("dump", &ModuleOp::dump) + .def("dump", [](ModuleOp &self) { +#if !defined(TRITON_CONCEAL_IR) || (TRITON_CONCEAL_IR == 0) + self.dump(); +#endif + }) .def("str", [](ModuleOp &self) -> std::string { std::string str; +#if !defined(TRITON_CONCEAL_IR) || (TRITON_CONCEAL_IR == 0) llvm::raw_string_ostream os(str); auto printingFlags = getOpPrintingFlags(); self.print(os, printingFlags); +#endif return str; }) .def("push_back", @@ -1964,6 +1983,7 @@ void init_triton_ir(py::module &&m) { .def("enable_debug", [](PassManager &self) -> bool { auto *context = self.getContext(); +#if !defined(TRITON_CONCEAL_IR) || (TRITON_CONCEAL_IR == 0) bool haveDump = ::triton::tools::getBoolEnv("MLIR_ENABLE_DUMP"); std::string funcToDump; if (!haveDump) { @@ -1973,6 +1993,10 @@ void init_triton_ir(py::module &&m) { if (!funcToDump.empty() && !isEnvValueBool) haveDump = true; } +#else + bool haveDump = false; + std::string funcToDump; +#endif if (haveDump) { context->disableMultithreading(); auto printingFlags = getOpPrintingFlags(); diff --git a/third_party/xpu/triton_xpu.cc b/third_party/xpu/triton_xpu.cc index 9016922db7..abae8aff30 100644 --- a/third_party/xpu/triton_xpu.cc +++ b/third_party/xpu/triton_xpu.cc @@ -91,6 +91,10 @@ void init_triton_xpu_passes_transform(py::module &&m) { {buffer_size, core_num, groups_per_cluster, isUseMaskZero})); }); + m.def("add_tritonxpu_tle_legalize_pass", [](mlir::PassManager &self) { + self.addPass(mlir::triton::xpu::createTritonXPUTLELegalize()); + }); + m.def("add_tritonxpu_mask_pass", [](mlir::PassManager &self, bool oneCoreActOnly, bool isUseMaskZero) { self.addPass(mlir::triton::xpu::createTritonXPUMask( @@ -112,15 +116,23 @@ void init_triton_xpu_passes_transform(py::module &&m) { self.addPass(mlir::triton::xpu::createTritonXPULoopGrid()); }); + m.def("add_tritonxpu_loop_invariant_staging_pass", + [](mlir::PassManager &self) { + self.addPass( + mlir::triton::xpu::createTritonXPULoopInvariantStaging()); + }); + m.def("add_tritonxpu_print_pass", [](mlir::PassManager &self) { self.addPass(mlir::triton::xpu::createTritonXPUPrint()); }); m.def("add_tritonxpu_unroll_control_pass", [](mlir::PassManager &self, uint32_t buffer_size, uint32_t core_num, - bool isUseMaskZero, uint32_t unroll_num) { + bool isUseMaskZero, uint32_t unroll_num, uint32_t vrf_budget = 24, + bool budget_tiling = false, int32_t pin_unroll_num = -1) { self.addPass(mlir::triton::xpu::createTritonXPUUnrollControl( - {buffer_size, core_num, isUseMaskZero, unroll_num})); + {buffer_size, core_num, isUseMaskZero, unroll_num, vrf_budget, + budget_tiling, pin_unroll_num})); }); m.def("add_tritonxpu_other_sim_pass", @@ -130,6 +142,12 @@ void init_triton_xpu_passes_transform(py::module &&m) { }); // Optimization Pass + m.def("add_tritonxpu_scalar_analysis_pass", + [](mlir::PassManager &self, bool aggressive) { + self.addPass( + mlir::triton::xpu::createTritonXPUScalarAnalysis({aggressive})); + }); + m.def("add_tritonxpu_offset_state_pass", [](mlir::PassManager &self, bool dump_flag, uint32_t buffer_size, bool isUseMaskZero) { @@ -146,17 +164,42 @@ void init_triton_xpu_passes_transform(py::module &&m) { triton_auto_core_tiling})); }); + m.def("add_tritonxpu_normalize_pass", + [](mlir::PassManager &self, bool dump_flag, bool compare_fusion) { + self.addPass(mlir::triton::xpu::createTritonXPUNormalize( + {dump_flag, compare_fusion})); + }); + m.def("add_tritonxpu_vectorize_pass", [](mlir::PassManager &self, bool dump_flag, bool compare_fusion) { self.addPass(mlir::triton::xpu::createTritonXPUVectorize( {dump_flag, compare_fusion})); }); + m.def("add_tritonxpu_tile_analysis_pass", + [](mlir::PassManager &self, uint32_t vrf_budget) { + self.addPass( + mlir::triton::xpu::createTritonXPUTileAnalysis({vrf_budget})); + }); + + m.def("add_tritonxpu_vectorizability_analysis_pass", + [](mlir::PassManager &self, bool reduce_vec, bool pre_tiling) { + self.addPass( + mlir::triton::xpu::createTritonXPUVectorizabilityAnalysis( + {reduce_vec, pre_tiling})); + }); + m.def("add_tritonxpu_memory_async_pass", [](mlir::PassManager &self, bool dump_flag) { self.addPass(mlir::triton::xpu::createTritonXPUMemoryAsync({dump_flag})); }); + m.def("add_tritonxpu_async_load_schedule_pass", + [](mlir::PassManager &self, bool dump_flag) { + self.addPass( + mlir::triton::xpu::createTritonXPUAsyncLoadSchedule({dump_flag})); + }); + m.def("add_tritonxpu_interleave_pass", [](mlir::PassManager &self) { self.addPass(mlir::triton::xpu::createTritonXPUInterleave()); }); @@ -185,6 +228,10 @@ void init_triton_xpu_passes_transform(py::module &&m) { m.def("add_tritonxpu_cf_to_scf_pass", [](mlir::PassManager &self) { self.addPass(mlir::triton::xpu::createTritonXPUCFToSCF()); }); + + m.def("add_tritonxpu_legalize_extern_ew_pass", [](mlir::PassManager &self) { + self.addPass(mlir::triton::xpu::createTritonXPULegalizeExternEW()); + }); } void init_triton_sdnn_passes_conversion(py::module &&m); @@ -391,6 +438,20 @@ void init_triton_xpu(py::module &&m) { // check if it is a sdnn kernel defineIsSDNNKernel(m); // is_sdnn_kernel + m.def("is_tle_kernel", [](mlir::ModuleOp &mod) { + bool hasTLEOp = false; + mod.walk([&](mlir::Operation *op) { + auto name = op->getName().getStringRef(); + if (name == "ttg.local_alloc" || + name == "triton_xpu.tle_copy_g2l" || + name == "triton_xpu.tle_copy_l2g" || + name == "triton_xpu.tle_local_ptr") { + hasTLEOp = true; + } + }); + return hasTLEOp; + }); + m.def("get_tensor_args", [](mlir::ModuleOp &mod, std::vector &tensorArgs) { mod.walk([&](mlir::triton::FuncOp funcOp) {